141 lines
4.9 KiB
Python
141 lines
4.9 KiB
Python
import os
|
|
|
|
import click
|
|
|
|
import mlflow
|
|
from mlflow.entities import ViewType
|
|
from mlflow.tracking import _get_store, fluent
|
|
from mlflow.utils.data_utils import is_uri
|
|
from mlflow.utils.string_utils import _create_table
|
|
|
|
EXPERIMENT_ID = click.option("--experiment-id", "-x", type=click.STRING, required=True)
|
|
|
|
|
|
@click.group("experiments")
|
|
def commands():
|
|
"""
|
|
Manage experiments. To manage experiments associated with a tracking server, set the
|
|
MLFLOW_TRACKING_URI environment variable to the URL of the desired server.
|
|
"""
|
|
|
|
|
|
@commands.command()
|
|
@click.option("--experiment-name", "-n", type=click.STRING, required=True)
|
|
@click.option(
|
|
"--artifact-location",
|
|
"-l",
|
|
help="Base location for runs to store artifact results. Artifacts will be stored "
|
|
"at $artifact_location/$run_id/artifacts. See "
|
|
"https://mlflow.org/docs/latest/tracking.html#where-runs-are-recorded for "
|
|
"more info on the properties of artifact location. "
|
|
"If no location is provided, the tracking server will pick a default.",
|
|
)
|
|
def create(experiment_name, artifact_location):
|
|
"""
|
|
Create an experiment.
|
|
|
|
All artifacts generated by runs related to this experiment will be stored under artifact
|
|
location, organized under specific run_id sub-directories.
|
|
|
|
Implementation of experiment and metadata store is dependent on backend storage. ``FileStore``
|
|
creates a folder for each experiment ID and stores metadata in ``meta.yaml``. Runs are stored
|
|
as subfolders.
|
|
"""
|
|
store = _get_store()
|
|
exp_id = store.create_experiment(experiment_name, artifact_location)
|
|
click.echo(f"Created experiment '{experiment_name}' with id {exp_id}")
|
|
|
|
|
|
@commands.command("search")
|
|
@click.option(
|
|
"--view",
|
|
"-v",
|
|
default="active_only",
|
|
help="Select view type for experiments. Valid view types are "
|
|
"'active_only' (default), 'deleted_only', and 'all'.",
|
|
)
|
|
def search_experiments(view):
|
|
"""
|
|
Search for experiments in the configured tracking server.
|
|
"""
|
|
view_type = ViewType.from_string(view) if view else ViewType.ACTIVE_ONLY
|
|
experiments = mlflow.search_experiments(view_type=view_type)
|
|
table = [
|
|
[
|
|
exp.experiment_id,
|
|
exp.name,
|
|
exp.artifact_location
|
|
if is_uri(exp.artifact_location)
|
|
else os.path.abspath(exp.artifact_location),
|
|
]
|
|
for exp in experiments
|
|
]
|
|
click.echo(_create_table(sorted(table), headers=["Experiment Id", "Name", "Artifact Location"]))
|
|
|
|
|
|
@commands.command("delete")
|
|
@EXPERIMENT_ID
|
|
def delete_experiment(experiment_id):
|
|
"""
|
|
Mark an active experiment for deletion. This also applies to experiment's metadata, runs and
|
|
associated data, and artifacts if they are store in default location. Use ``list`` command to
|
|
view artifact location. Command will throw an error if experiment is not found or already
|
|
marked for deletion.
|
|
|
|
Experiments marked for deletion can be restored using ``restore`` command, unless they are
|
|
permanently deleted.
|
|
|
|
Specific implementation of deletion is dependent on backend stores. ``FileStore`` moves
|
|
experiments marked for deletion under a ``.trash`` folder under the main folder used to
|
|
instantiate ``FileStore``. Experiments marked for deletion can be permanently deleted by
|
|
clearing the ``.trash`` folder. It is recommended to use a ``cron`` job or an alternate
|
|
workflow mechanism to clear ``.trash`` folder.
|
|
"""
|
|
store = _get_store()
|
|
store.delete_experiment(experiment_id)
|
|
click.echo(f"Experiment with ID {experiment_id} has been deleted.")
|
|
|
|
|
|
@commands.command("restore")
|
|
@EXPERIMENT_ID
|
|
def restore_experiment(experiment_id):
|
|
"""
|
|
Restore a deleted experiment. This also applies to experiment's metadata, runs and associated
|
|
data. The command throws an error if the experiment is already active, cannot be found, or
|
|
permanently deleted.
|
|
"""
|
|
store = _get_store()
|
|
store.restore_experiment(experiment_id)
|
|
click.echo(f"Experiment with id {experiment_id} has been restored.")
|
|
|
|
|
|
@commands.command("rename")
|
|
@EXPERIMENT_ID
|
|
@click.option("--new-name", type=click.STRING, required=True)
|
|
def rename_experiment(experiment_id, new_name):
|
|
"""
|
|
Renames an active experiment.
|
|
Returns an error if the experiment is inactive.
|
|
"""
|
|
store = _get_store()
|
|
store.rename_experiment(experiment_id, new_name)
|
|
click.echo(f"Experiment with id {experiment_id} has been renamed to '{new_name}'.")
|
|
|
|
|
|
@commands.command("csv")
|
|
@EXPERIMENT_ID
|
|
@click.option("--filename", "-o", type=click.STRING)
|
|
def generate_csv_with_runs(experiment_id, filename):
|
|
# type: (str, str) -> None
|
|
"""
|
|
Generate CSV with all runs for an experiment
|
|
"""
|
|
runs = fluent.search_runs(experiment_ids=experiment_id)
|
|
if filename:
|
|
runs.to_csv(filename, index=False)
|
|
click.echo(
|
|
f"Experiment with ID {experiment_id} has been exported as a CSV to file: {filename}."
|
|
)
|
|
else:
|
|
click.echo(runs.to_csv(index=False))
|