1  """
  2  CLI for runs
  3  """
  4  
  5  import json
  6  
  7  import click
  8  
  9  import mlflow
 10  from mlflow import MlflowClient
 11  from mlflow.entities import RunStatus, ViewType
 12  from mlflow.environment_variables import MLFLOW_EXPERIMENT_ID, MLFLOW_EXPERIMENT_NAME
 13  from mlflow.exceptions import MlflowException
 14  from mlflow.mcp.decorator import mlflow_mcp
 15  from mlflow.tracking import _get_store
 16  from mlflow.utils.string_utils import _create_table
 17  from mlflow.utils.time import conv_longdate_to_str
 18  
 19  RUN_ID = click.option("--run-id", type=click.STRING, required=True)
 20  
 21  
 22  @click.group("runs")
 23  def commands():
 24      """
 25      Manage runs. To manage runs of experiments associated with a tracking server, set the
 26      MLFLOW_TRACKING_URI environment variable to the URL of the desired server.
 27      """
 28  
 29  
 30  @commands.command("list")
 31  @mlflow_mcp(tool_name="list_runs")
 32  @click.option(
 33      "--experiment-id",
 34      envvar=MLFLOW_EXPERIMENT_ID.name,
 35      type=click.STRING,
 36      help="Specify the experiment ID for list of runs.",
 37      required=True,
 38  )
 39  @click.option(
 40      "--view",
 41      "-v",
 42      default="active_only",
 43      help="Select view type for list experiments. Valid view types are "
 44      "'active_only' (default), 'deleted_only', and 'all'.",
 45  )
 46  def list_run(experiment_id: str, view: str) -> None:
 47      """
 48      List all runs of the specified experiment in the configured tracking server.
 49      """
 50      store = _get_store()
 51      view_type = ViewType.from_string(view) if view else ViewType.ACTIVE_ONLY
 52      runs = store.search_runs([experiment_id], None, view_type)
 53      table = []
 54      for run in runs:
 55          run_name = run.info.run_name or ""
 56          table.append([conv_longdate_to_str(run.info.start_time), run_name, run.info.run_id])
 57      click.echo(_create_table(sorted(table, reverse=True), headers=["Date", "Name", "ID"]))
 58  
 59  
 60  @commands.command("delete")
 61  @mlflow_mcp(tool_name="delete_run")
 62  @RUN_ID
 63  def delete_run(run_id: str) -> None:
 64      """
 65      Mark a run for deletion. Return an error if the run does not exist or
 66      is already marked. You can restore a marked run with ``restore_run``,
 67      or permanently delete a run in the backend store.
 68      """
 69      store = _get_store()
 70      store.delete_run(run_id)
 71      click.echo(f"Run with ID {run_id} has been deleted.")
 72  
 73  
 74  @commands.command("restore")
 75  @mlflow_mcp(tool_name="restore_run")
 76  @RUN_ID
 77  def restore_run(run_id: str) -> None:
 78      """
 79      Restore a deleted run.
 80      Returns an error if the run is active or has been permanently deleted.
 81      """
 82      store = _get_store()
 83      store.restore_run(run_id)
 84      click.echo(f"Run with id {run_id} has been restored.")
 85  
 86  
 87  @commands.command("describe")
 88  @mlflow_mcp(tool_name="describe_run")
 89  @RUN_ID
 90  def describe_run(run_id: str) -> None:
 91      """
 92      All of run details will print to the stdout as JSON format.
 93      """
 94      store = _get_store()
 95      run = store.get_run(run_id)
 96      json_run = json.dumps(run.to_dictionary(), indent=4)
 97      click.echo(json_run)
 98  
 99  
100  @commands.command("create")
101  @mlflow_mcp(tool_name="create_run")
102  @click.option(
103      "--experiment-id",
104      envvar=MLFLOW_EXPERIMENT_ID.name,
105      type=click.STRING,
106      help="ID of the experiment under which to create the run. "
107      "Must specify either this or --experiment-name.",
108  )
109  @click.option(
110      "--experiment-name",
111      envvar=MLFLOW_EXPERIMENT_NAME.name,
112      type=click.STRING,
113      help="Name of the experiment under which to create the run. "
114      "Must specify either this or --experiment-id.",
115  )
116  @click.option(
117      "--run-name",
118      type=click.STRING,
119      help="Optional human-readable name for the run (e.g., 'baseline-model-v1').",
120  )
121  @click.option(
122      "--description",
123      type=click.STRING,
124      help="Optional longer description of what this run represents.",
125  )
126  @click.option(
127      "--tags",
128      "-t",
129      multiple=True,
130      help="Key-value pairs to categorize and filter runs. Use multiple times for "
131      "multiple tags. Format: key=value (e.g., env=prod, model=xgboost, version=1.0).",
132  )
133  @click.option(
134      "--status",
135      type=click.Choice(["FINISHED", "FAILED", "KILLED"], case_sensitive=False),
136      default="FINISHED",
137      help="Final status of the run. Options: FINISHED (default), FAILED, or KILLED.",
138  )
139  @click.option(
140      "--parent-run-id",
141      type=click.STRING,
142      help="Optional ID of a parent run to create a nested run under.",
143  )
144  def create_run(
145      experiment_id: str | None,
146      experiment_name: str | None,
147      run_name: str | None,
148      description: str | None,
149      tags: tuple[str, ...],
150      status: str,
151      parent_run_id: str | None,
152  ) -> None:
153      """
154      Create a new MLflow run and immediately end it with the specified status.
155  
156      This command is useful for creating runs programmatically for testing, scripting,
157      or recording completed experiments. The run will be created and immediately closed
158      with the specified status (FINISHED, FAILED, or KILLED).
159      """
160      # Validate that exactly one of experiment_id or experiment_name is provided
161      if (experiment_id is not None and experiment_name is not None) or (
162          experiment_id is None and experiment_name is None
163      ):
164          raise click.UsageError("Must specify exactly one of --experiment-id or --experiment-name.")
165  
166      # Parse tags from key=value format
167      tags_dict = {}
168      if tags:
169          for tag in tags:
170              match tag.split("=", 1):
171                  case [key, value]:
172                      if key in tags_dict:
173                          raise click.UsageError(f"Duplicate tag key: '{key}'")
174                      tags_dict[key] = value
175                  case _:
176                      raise click.UsageError(
177                          f"Invalid tag format: '{tag}'. Tags must be in key=value format."
178                      )
179  
180      # Set the experiment if using experiment_name
181      if experiment_name:
182          experiment = mlflow.set_experiment(experiment_name=experiment_name)
183          experiment_id = experiment.experiment_id
184  
185      # Start the run with the specified parameters
186      try:
187          # Start the run
188          active_run = mlflow.start_run(
189              experiment_id=experiment_id,
190              run_name=run_name,
191              nested=bool(parent_run_id),
192              parent_run_id=parent_run_id,
193              tags=tags_dict,
194              description=description,
195          )
196          run_id = active_run.info.run_id
197          actual_experiment_id = active_run.info.experiment_id
198  
199          # End the run with the specified status
200          mlflow.end_run(status=RunStatus.to_string(getattr(RunStatus, status.upper())))
201  
202          # Output the created run information
203          output = {
204              "run_id": run_id,
205              "experiment_id": actual_experiment_id,
206              "status": status.upper(),
207              "run_name": run_name,
208          }
209  
210          click.echo(json.dumps(output, indent=2))
211  
212      except MlflowException as e:
213          raise click.ClickException(f"Failed to create run: {e.message}")
214      except Exception as e:
215          raise click.ClickException(f"Unexpected error creating run: {e!s}")
216  
217  
218  @commands.command("link-traces")
219  @mlflow_mcp(tool_name="link_traces_to_run")
220  @click.option(
221      "--run-id",
222      type=click.STRING,
223      required=True,
224      help="ID of the run to link traces to.",
225  )
226  @click.option(
227      "trace_ids",
228      "--trace-id",
229      "-t",
230      multiple=True,
231      required=True,
232      help="Trace ID to link to the run. Can be specified multiple times (maximum 100 traces).",
233  )
234  def link_traces(run_id: str, trace_ids: tuple[str, ...]) -> None:
235      """
236      Link traces to a run.
237  
238      This command links one or more traces to an existing run. Traces can be
239      linked to runs to establish relationships between traces and runs.
240      Maximum 100 traces can be linked in a single command.
241      """
242      try:
243          client = MlflowClient()
244          client.link_traces_to_run(list(trace_ids), run_id)
245  
246          # Output success message with count
247          click.echo(f"Successfully linked {len(trace_ids)} trace(s) to run '{run_id}'")
248  
249      except MlflowException as e:
250          raise click.ClickException(f"Failed to link traces: {e.message}")
251      except Exception as e:
252          raise click.ClickException(f"Unexpected error linking traces: {e!s}")