import contextlib import json import logging import os import re import sys import warnings from datetime import timedelta from pathlib import Path import click from click import UsageError from click.core import ParameterSource from dotenv import load_dotenv import mlflow.db import mlflow.deployments.cli import mlflow.experiments import mlflow.runs import mlflow.store.artifact.cli from mlflow import ai_commands, projects, version from mlflow.entities import ViewType from mlflow.entities.lifecycle_stage import LifecycleStage from mlflow.environment_variables import ( MLFLOW_ENABLE_WORKSPACES, MLFLOW_EXPERIMENT_ID, MLFLOW_EXPERIMENT_NAME, MLFLOW_TRACE_ARCHIVAL_CONFIG, MLFLOW_WORKSPACE, MLFLOW_WORKSPACE_STORE_URI, ) from mlflow.exceptions import InvalidUrlException, MlflowException from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE, RESOURCE_DOES_NOT_EXIST, ErrorCode from mlflow.store.artifact.artifact_repository_registry import get_artifact_repository from mlflow.store.tracking import ( DEFAULT_ARTIFACTS_URI, DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH, ) from mlflow.store.workspace.utils import get_default_workspace_optional from mlflow.telemetry.events import TrackingServerStartEvent from mlflow.telemetry.track import _record_event from mlflow.tracing.trace_archival_config import load_trace_archival_server_config from mlflow.tracking import _get_store from mlflow.tracking._tracking_service.utils import ( _get_default_tracking_uri, is_tracking_uri_set, set_tracking_uri, ) from mlflow.tracking._workspace.registry import get_workspace_store from mlflow.utils import cli_args, workspace_context from mlflow.utils.logging_utils import eprint from mlflow.utils.os import is_windows from mlflow.utils.plugins import get_entry_points from mlflow.utils.process import ShellCommandException from mlflow.utils.server_cli_utils import ( artifacts_only_config_validation, assert_server_workspace_env_unset, resolve_default_artifact_root, ) from mlflow.utils.workspace_utils import resolve_workspace_store_uri _logger = logging.getLogger(__name__) class AliasedGroup(click.Group): def get_command(self, ctx, cmd_name): # `mlflow ui` is an alias for `mlflow server` cmd_name = "server" if cmd_name == "ui" else cmd_name return super().get_command(ctx, cmd_name) def _load_env_file(ctx: click.Context, param: click.Parameter, value: str | None) -> str | None: """ Click callback to load environment variables from a dotenv file. This function is designed to be used as an eager callback for the --env-file option, ensuring that environment variables are loaded before any command execution. """ if value is not None: env_path = Path(value) if not env_path.exists(): raise click.BadParameter(f"Environment file '{value}' does not exist.") # Load the environment file # override=False means existing environment variables take precedence load_dotenv(env_path, override=False) # Log that we've loaded the env file (using click.echo for CLI output) click.echo(f"Loaded environment variables from: {value}") return value @click.group(cls=AliasedGroup) @click.version_option(version=version.VERSION) @click.option( "--env-file", type=click.Path(exists=False), callback=_load_env_file, expose_value=True, is_eager=True, help="Load environment variables from a dotenv file before executing the command. " "Variables in the file will be loaded but won't override existing environment variables.", ) def cli(env_file): pass @cli.command() @click.argument("uri") @click.option( "--entry-point", "-e", metavar="NAME", default="main", help="Entry point within project. [default: main]. If the entry point is not found, " "attempts to run the project file with the specified name as a script, " "using 'python' to run .py files and the default shell (specified by " "environment variable $SHELL) to run .sh files", ) @click.option( "--version", "-v", metavar="VERSION", help="Version of the project to run, as a Git commit reference for Git projects.", ) @click.option( "--param-list", "-P", metavar="NAME=VALUE", multiple=True, help="A parameter for the run, of the form -P name=value. Provided parameters that " "are not in the list of parameters for an entry point will be passed to the " "corresponding entry point as command-line arguments in the form `--name value`", ) @click.option( "--docker-args", "-A", metavar="NAME=VALUE", multiple=True, help="A `docker run` argument or flag, of the form -A name=value (e.g. -A gpus=all) " "or -A name (e.g. -A t). The argument will then be passed as " "`docker run --name value` or `docker run --name` respectively. ", ) @click.option( "--experiment-name", envvar=MLFLOW_EXPERIMENT_NAME.name, help="Name of the experiment under which to launch the run. If not " "specified, 'experiment-id' option will be used to launch run.", ) @click.option( "--experiment-id", envvar=MLFLOW_EXPERIMENT_ID.name, type=click.STRING, help="ID of the experiment under which to launch the run.", ) # TODO: Add tracking server argument once we have it working. @click.option( "--backend", "-b", metavar="BACKEND", default="local", help="Execution backend to use for run. Supported values: 'local', 'databricks', " "kubernetes (experimental). Defaults to 'local'. If running against " "Databricks, will run against a Databricks workspace determined as follows: " "if a Databricks tracking URI of the form 'databricks://profile' has been set " "(e.g. by setting the MLFLOW_TRACKING_URI environment variable), will run " "against the workspace specified by . Otherwise, runs against the " "workspace specified by the default Databricks CLI profile. See " "https://github.com/databricks/databricks-cli for more info on configuring a " "Databricks CLI profile.", ) @click.option( "--backend-config", "-c", metavar="FILE", help="Path to JSON file (must end in '.json') or JSON string which will be passed " "as config to the backend. The exact content which should be " "provided is different for each execution backend and is documented " "at https://www.mlflow.org/docs/latest/projects.html.", ) @cli_args.ENV_MANAGER_PROJECTS @click.option( "--storage-dir", envvar="MLFLOW_TMP_DIR", help="Only valid when ``backend`` is local. " "MLflow downloads artifacts from distributed URIs passed to parameters of " "type 'path' to subdirectories of storage_dir.", ) @click.option( "--run-id", metavar="RUN_ID", help="If specified, the given run ID will be used instead of creating a new run. " "Note: this argument is used internally by the MLflow project APIs " "and should not be specified.", ) @click.option( "--run-name", metavar="RUN_NAME", help="The name to give the MLflow Run associated with the project execution. If not specified, " "the MLflow Run name is left unset.", ) @click.option( "--build-image", is_flag=True, default=False, show_default=True, help=( "Only valid for Docker projects. If specified, build a new Docker image that's based on " "the image specified by the `image` field in the MLproject file, and contains files in the " "project directory." ), ) def run( uri, entry_point, version, param_list, docker_args, experiment_name, experiment_id, backend, backend_config, env_manager, storage_dir, run_id, run_name, build_image, ): """ Run an MLflow project from the given URI. For local runs, the run will block until it completes. Otherwise, the project will run asynchronously. If running locally (the default), the URI can be either a Git repository URI or a local path. If running on Databricks, the URI must be a Git repository. By default, Git projects run in a new working directory with the given parameters, while local projects run from the project's root directory. """ if experiment_id is not None and experiment_name is not None: raise click.UsageError("Specify only one of 'experiment-name' or 'experiment-id' options.") param_dict = _user_args_to_dict(param_list) args_dict = _user_args_to_dict(docker_args, argument_type="A") if backend_config is not None and os.path.splitext(backend_config)[-1] != ".json": try: backend_config = json.loads(backend_config) except ValueError as e: raise click.UsageError(f"Invalid backend config JSON. Parse error: {e}") from e if backend == "kubernetes": if backend_config is None: raise click.UsageError("Specify 'backend_config' when using kubernetes mode.") try: projects.run( uri, entry_point, version, experiment_name=experiment_name, experiment_id=experiment_id, parameters=param_dict, docker_args=args_dict, backend=backend, backend_config=backend_config, env_manager=env_manager, storage_dir=storage_dir, synchronous=backend in ("local", "kubernetes") or backend is None, run_id=run_id, run_name=run_name, build_image=build_image, ) except projects.ExecutionException as e: _logger.error("=== %s ===", e) sys.exit(1) def _user_args_to_dict(arguments, argument_type="P"): user_dict = {} for arg in arguments: split = arg.split("=", maxsplit=1) # Docker arguments such as `t` don't require a value -> set to True if specified if len(split) == 1 and argument_type == "A": name = split[0] value = True elif len(split) == 2: name = split[0] value = split[1] else: raise click.UsageError( f"Invalid format for -{argument_type} parameter: '{arg}'. " f"Use -{argument_type} name=value." ) if name in user_dict: raise click.UsageError(f"Repeated parameter: '{name}'") user_dict[name] = value return user_dict def _validate_server_args( ctx=None, gunicorn_opts=None, workers=None, waitress_opts=None, uvicorn_opts=None, allowed_hosts=None, cors_allowed_origins=None, x_frame_options=None, disable_security_middleware=None, ): if sys.platform == "win32": if gunicorn_opts is not None: raise NotImplementedError( "gunicorn is not supported on Windows, cannot specify --gunicorn-opts" ) num_server_opts_specified = sum( 1 for opt in [gunicorn_opts, waitress_opts, uvicorn_opts] if opt is not None ) if num_server_opts_specified > 1: raise click.UsageError( "Cannot specify multiple server options. Choose one of: " "'--gunicorn-opts', '--waitress-opts', or '--uvicorn-opts'." ) using_flask_only = gunicorn_opts is not None or waitress_opts is not None # NB: Only check for security params that are explicitly passed via CLI (not env vars) # This allows Docker containers to set env vars while using gunicorn from click.core import ParameterSource security_params_specified = False if ctx: security_params_specified = any([ ctx.get_parameter_source("allowed_hosts") == ParameterSource.COMMANDLINE, ctx.get_parameter_source("cors_allowed_origins") == ParameterSource.COMMANDLINE, ( ctx.get_parameter_source("disable_security_middleware") == ParameterSource.COMMANDLINE ), ]) if using_flask_only and security_params_specified: raise click.UsageError( "Security middleware parameters (--allowed-hosts, --cors-allowed-origins, " "--disable-security-middleware) are only supported with " "the default uvicorn server. They cannot be used with --gunicorn-opts or " "--waitress-opts. To use security features, run without specifying a server " "option (uses uvicorn by default) or explicitly use --uvicorn-opts." ) def _validate_static_prefix(ctx, param, value): """ Validate that the static_prefix option starts with a "/" and does not end in a "/". Conforms to the callback interface of click documented at http://click.pocoo.org/5/options/#callbacks-for-validation. """ if value is not None: if not value.startswith("/"): raise UsageError("--static-prefix must begin with a '/'.") if value.endswith("/"): raise UsageError("--static-prefix should not end with a '/'.") return value @cli.command(short_help="Run the MLflow tracking server (UI + REST API).") @click.pass_context @click.option( "--backend-store-uri", envvar="MLFLOW_BACKEND_STORE_URI", metavar="PATH", default=None, help="URI to which to persist experiment and run data. Acceptable URIs are " "SQLAlchemy-compatible database connection strings " "(e.g. 'sqlite:///path/to/file.db') or local filesystem URIs " "(e.g. 'file:///absolute/path/to/directory'). By default, data is logged to a local " "SQLite database (sqlite:///mlflow.db), falling back to the ./mlruns directory when an " "existing ./mlruns store is present.", ) @click.option( "--read-replica-backend-store-uri", envvar="MLFLOW_READ_REPLICA_BACKEND_STORE_URI", metavar="URI", default=None, help="URI for a read-only database replica. When specified, read operations " "(e.g. search_runs, get_experiment) are routed to this URI while write operations " "use --backend-store-uri. Enables horizontal scaling via database read replicas. " "If not specified, all operations use --backend-store-uri. " "Note: there is no automatic failover to the primary if the replica becomes " "unavailable. Cloud-managed databases (Aurora, RDS) handle this at the DNS level. " "For self-hosted setups, use a connection proxy (PgBouncer, HAProxy) for failover.", ) @click.option( "--registry-store-uri", envvar="MLFLOW_REGISTRY_STORE_URI", metavar="URI", default=None, help="URI to which to persist registered models. Acceptable URIs are " "SQLAlchemy-compatible database connection strings (e.g. 'sqlite:///path/to/file.db'). " "If not specified, `backend-store-uri` is used.", ) @click.option( "--default-artifact-root", envvar="MLFLOW_DEFAULT_ARTIFACT_ROOT", metavar="URI", default=None, help="Directory in which to store artifacts for any new experiments created. For tracking " "server backends that rely on SQL, this option is required in order to store artifacts. " "Note that this flag does not impact already-created experiments with any previous " "configuration of an MLflow server instance. " f"By default, data will be logged to the {DEFAULT_ARTIFACTS_URI} uri proxy if " "the --serve-artifacts option is enabled. Otherwise, the default location will " f"be {DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH}.", ) @cli_args.SERVE_ARTIFACTS @click.option( "--artifacts-only", envvar="MLFLOW_ARTIFACTS_ONLY", is_flag=True, default=False, help="If specified, configures the mlflow server to be used only for proxied artifact serving. " "With this mode enabled, functionality of the mlflow tracking service (e.g. run creation, " "metric logging, and parameter logging) is disabled. The server will only expose " "endpoints for uploading, downloading, and listing artifacts. " "Default: False", ) @cli_args.ARTIFACTS_DESTINATION @cli_args.HOST @cli_args.PORT @cli_args.WORKERS @click.option( "--static-prefix", envvar="MLFLOW_STATIC_PREFIX", default=None, callback=_validate_static_prefix, help="A prefix which will be prepended to the path of all static paths.", ) @cli_args.ALLOWED_HOSTS @cli_args.CORS_ALLOWED_ORIGINS @cli_args.DISABLE_SECURITY_MIDDLEWARE @cli_args.X_FRAME_OPTIONS @click.option( "--gunicorn-opts", envvar="MLFLOW_GUNICORN_OPTS", default=None, help="Additional command line options forwarded to gunicorn processes.", ) @click.option( "--waitress-opts", default=None, help="Additional command line options for waitress-serve." ) @click.option( "--uvicorn-opts", envvar="MLFLOW_UVICORN_OPTS", default=None, help="Additional command line options forwarded to uvicorn processes (used by default).", ) @click.option( "--dev", is_flag=True, default=False, show_default=True, help=( "If enabled, run the server with debug logging and auto-reload. " "Should only be used for development purposes. " "Cannot be used with '--gunicorn-opts' or '--uvicorn-opts'. " "Unsupported on Windows." ), ) @click.option( "--expose-prometheus", envvar="MLFLOW_EXPOSE_PROMETHEUS", default=None, help="Path to the directory where metrics will be stored. If the directory " "doesn't exist, it will be created. " "Activate prometheus exporter to expose metrics on /metrics endpoint.", ) @click.option( "--app-name", default=None, type=click.Choice([e.name for e in get_entry_points("mlflow.app")]), show_default=True, help=( "Application name to be used for the tracking server. " "If not specified, 'mlflow.server:app' will be used." ), ) @click.option( "--trace-archival-config", envvar=MLFLOW_TRACE_ARCHIVAL_CONFIG.name, type=click.Path(exists=True, dir_okay=False, readable=True, path_type=Path), metavar="PATH", default=None, help=("Path to the YAML config file for server-owned trace archival."), ) @click.option( "--secrets-cache-ttl", type=click.IntRange(10, 300), default=60, show_default=True, help=( "Server-side secrets cache time-to-live in seconds. " "Controls how long decrypted secrets are cached in memory (encrypted with AES-GCM-256). " "Lower values (10-30s) are more secure but impact performance. " "Higher values (120-300s) improve performance but increase exposure window. " "Range: 10-300 seconds." ), ) @click.option( "--secrets-cache-max-size", type=click.IntRange(1, 10000), default=1000, show_default=True, help=( "Server-side secrets cache maximum entries. " "When exceeded, least recently used entries are evicted. " "Range: 1-10000 entries." ), ) @click.option( "--workspace-store-uri", envvar=MLFLOW_WORKSPACE_STORE_URI.name, metavar="URI", default=None, help=( "Workspace provider backend URI used for workspace CRUD APIs and request routing. " "When unspecified, defaults to the backend store URI. This only needs to be specified " "when using a workspace store plugin leveraging externally managed workspaces (e.g. " + "Kubernetes namespaces)." ), ) @click.option( "--enable-workspaces/--disable-workspaces", default=False, show_default=True, help="Enable backwards compatible workspaces mode for logical isolation of experiments, " + "registered models, and prompts.", ) def server( ctx, backend_store_uri, read_replica_backend_store_uri, registry_store_uri, default_artifact_root, serve_artifacts, artifacts_only, artifacts_destination, host, port, workers, allowed_hosts, cors_allowed_origins, disable_security_middleware, x_frame_options, static_prefix, gunicorn_opts, waitress_opts, expose_prometheus, app_name, trace_archival_config, dev, uvicorn_opts, secrets_cache_ttl, secrets_cache_max_size, workspace_store_uri, enable_workspaces, ): """Run the MLflow tracking server (UI + REST API). Whether you're doing LLMOps or training classic ML models, the server is the same. Pick your setup by how you're deploying: \b First time user? Try MLflow out locally: mlflow server (defaults to sqlite:///mlflow.db; reuses an existing ./mlruns file store if one is present) Team server (shared): mlflow server --backend-store-uri postgresql://... --host 0.0.0.0 Cloud artifact storage (proxied through server): ...add --artifacts-destination s3://my-bucket Exposed beyond LAN: ...add --allowed-hosts host.co --cors-allowed-origins https://host.co \b Options by topic: Storage --backend-store-uri, --registry-store-uri, --default-artifact-root, ... Network --host, --port, --workers, --static-prefix Security --allowed-hosts, --cors-allowed-origins, --x-frame-options, ... Advanced --app-name, --expose-prometheus, --enable-workspaces, ... \b Full guide: https://mlflow.org/docs/latest/self-hosting/architecture/tracking-server """ from mlflow.server import _run_server from mlflow.server.handlers import initialize_backend_stores # Get env_file from parent context env_file = ctx.parent.params.get("env_file") if ctx.parent else None if dev: if is_windows(): raise click.UsageError("'--dev' is not supported on Windows.") if gunicorn_opts: raise click.UsageError("'--dev' and '--gunicorn-opts' cannot be specified together.") if uvicorn_opts: raise click.UsageError("'--dev' and '--uvicorn-opts' cannot be specified together.") if app_name: raise click.UsageError( "'--dev' cannot be used with '--app-name'. Development mode with auto-reload " "is only supported for the default MLflow tracking server." ) uvicorn_opts = "--reload --log-level debug" _validate_server_args( ctx=ctx, gunicorn_opts=gunicorn_opts, workers=workers, waitress_opts=waitress_opts, uvicorn_opts=uvicorn_opts, allowed_hosts=allowed_hosts, cors_allowed_origins=cors_allowed_origins, x_frame_options=x_frame_options, disable_security_middleware=disable_security_middleware, ) # click treats any non-empty env var as "set" for flag options, which would interpret # MLFLOW_ENABLE_WORKSPACES="false" as True. If the flag wasn't set explicitly and # resolved to False, fall back to the env var parser to preserve "false"/"0". if ( ctx and not enable_workspaces and ctx.get_parameter_source("enable_workspaces") != ParameterSource.COMMANDLINE ): enable_workspaces = MLFLOW_ENABLE_WORKSPACES.get() assert_server_workspace_env_unset() if disable_security_middleware: os.environ["MLFLOW_SERVER_DISABLE_SECURITY_MIDDLEWARE"] = "true" else: if allowed_hosts: os.environ["MLFLOW_SERVER_ALLOWED_HOSTS"] = allowed_hosts if allowed_hosts == "*": click.echo( "WARNING: Accepting ALL hosts. " "This may leave the server vulnerable to DNS rebinding attacks." ) if cors_allowed_origins: os.environ["MLFLOW_SERVER_CORS_ALLOWED_ORIGINS"] = cors_allowed_origins if cors_allowed_origins == "*": click.echo( "WARNING: Allowing ALL origins for CORS. " "This allows ANY website to access your MLflow data. " "This configuration is only recommended for local development." ) if x_frame_options: os.environ["MLFLOW_SERVER_X_FRAME_OPTIONS"] = x_frame_options if not backend_store_uri: backend_store_uri = _get_default_tracking_uri() click.echo(f"Backend store URI not provided. Using {backend_store_uri}") if not registry_store_uri: registry_store_uri = backend_store_uri click.echo("Registry store URI not provided. Using backend store URI.") default_artifact_root = resolve_default_artifact_root( serve_artifacts, default_artifact_root, backend_store_uri ) artifacts_only_config_validation( artifacts_only, backend_store_uri, enable_workspaces, trace_archival_config_path=str(trace_archival_config) if trace_archival_config else None, ) if trace_archival_config is not None: try: load_trace_archival_server_config(trace_archival_config) except MlflowException as e: raise click.UsageError(e.message) from e # Keep environment flag in sync with the resolved boolean so server-side gating # (which reads MLFLOW_ENABLE_WORKSPACES.get()) has a single source of truth. os.environ[MLFLOW_ENABLE_WORKSPACES.name] = "true" if enable_workspaces else "false" if enable_workspaces and workspace_store_uri: os.environ[MLFLOW_WORKSPACE_STORE_URI.name] = workspace_store_uri elif workspace_store_uri: click.echo( "Ignoring --workspace-store-uri because workspaces are not enabled. " "Use --enable-workspaces to activate workspace mode.", err=True, ) if trace_archival_config is not None: os.environ[MLFLOW_TRACE_ARCHIVAL_CONFIG.name] = str(trace_archival_config) if not artifacts_only: try: initialize_backend_stores( backend_store_uri, registry_store_uri, default_artifact_root, workspace_store_uri=workspace_store_uri, read_replica_backend_store_uri=read_replica_backend_store_uri, ) except Exception as e: _logger.error("Error initializing backend store") _logger.exception(e) sys.exit(1) if disable_security_middleware: click.echo( "[MLflow] WARNING: Security middleware is DISABLED. " "Your MLflow server is vulnerable to various attacks.", err=True, ) elif not allowed_hosts and not cors_allowed_origins: click.echo( "[MLflow] Security middleware enabled with default settings (localhost-only). " "To allow connections from other hosts, use --host 0.0.0.0 and configure " "--allowed-hosts and --cors-allowed-origins.", err=True, ) else: parts = ["[MLflow] Security middleware enabled"] if allowed_hosts: hosts_list = allowed_hosts.split(",")[:3] if len(allowed_hosts.split(",")) > 3: hosts_list.append(f"and {len(allowed_hosts.split(',')) - 3} more") parts.append(f"Allowed hosts: {', '.join(hosts_list)}") if cors_allowed_origins: origins_list = cors_allowed_origins.split(",")[:3] if len(cors_allowed_origins.split(",")) > 3: origins_list.append(f"and {len(cors_allowed_origins.split(',')) - 3} more") parts.append(f"CORS origins: {', '.join(origins_list)}") click.echo(". ".join(parts) + ".", err=True) _record_event( TrackingServerStartEvent, TrackingServerStartEvent.parse({ "backend_store_uri": backend_store_uri, "serve_artifacts": serve_artifacts, "artifacts_only": artifacts_only, "expose_prometheus": expose_prometheus, "app_name": app_name, "enable_workspaces": enable_workspaces, "workers": workers, "dev": dev, }) or {}, ) try: _run_server( file_store_path=backend_store_uri, read_replica_backend_store_uri=read_replica_backend_store_uri, registry_store_uri=registry_store_uri, default_artifact_root=default_artifact_root, serve_artifacts=serve_artifacts, artifacts_only=artifacts_only, artifacts_destination=artifacts_destination, host=host, port=port, static_prefix=static_prefix, workers=workers, gunicorn_opts=gunicorn_opts, waitress_opts=waitress_opts, expose_prometheus=expose_prometheus, app_name=app_name, uvicorn_opts=uvicorn_opts, env_file=env_file, secrets_cache_ttl=secrets_cache_ttl, secrets_cache_max_size=secrets_cache_max_size, ) except ShellCommandException: eprint("Running the mlflow server failed. Please see the logs above for details.") sys.exit(1) def _gc_tracking_resources( backend_store, run_ids: list[str] | None, experiment_ids: list[str] | None, logged_model_ids: list[str] | None, older_than: str | None, time_delta: int, skip_experiments: bool, skip_logged_models: bool, ignore_not_found: bool = False, ): """ Perform garbage collection of tracking resources (runs, experiments, logged models). This is the core implementation of the gc command, extracted to support workspace iteration. Args: backend_store: The tracking store instance. run_ids: Optional list of specific run IDs to delete. experiment_ids: Optional list of specific experiment IDs to delete. logged_model_ids: Optional list of specific logged model IDs to delete. older_than: Original older_than string for error messages. time_delta: Time delta in milliseconds for age filtering. skip_experiments: Whether to skip experiment deletion. skip_logged_models: Whether to skip logged model deletion. ignore_not_found: If True, skip RESOURCE_DOES_NOT_EXIST errors for explicit IDs that may not exist (e.g., when iterating over multiple workspaces). """ from mlflow.utils.time import get_current_time_millis deleted_run_ids_older_than = backend_store._get_deleted_runs(older_than=time_delta) run_ids_to_delete = run_ids if run_ids is not None else list(deleted_run_ids_older_than) deleted_logged_model_ids = ( backend_store._get_deleted_logged_models() if not skip_logged_models else [] ) deleted_logged_model_ids_older_than = ( backend_store._get_deleted_logged_models(older_than=time_delta) if not skip_logged_models else [] ) logged_model_ids_to_delete = ( logged_model_ids if logged_model_ids is not None else list(deleted_logged_model_ids_older_than) ) time_threshold = get_current_time_millis() - time_delta experiment_ids_to_delete = [] if not skip_experiments: if experiment_ids: experiments = [] for exp_id in experiment_ids: try: experiments.append(backend_store.get_experiment(exp_id)) except MlflowException as exc: if ignore_not_found and exc.error_code == ErrorCode.Name( RESOURCE_DOES_NOT_EXIST ): continue raise # Ensure that the specified experiments are soft-deleted active_experiment_ids = [ e.experiment_id for e in experiments if e.lifecycle_stage != LifecycleStage.DELETED ] if active_experiment_ids: raise MlflowException( f"Experiments {active_experiment_ids} are not in the deleted lifecycle stage. " "Only experiments in the deleted lifecycle stage can be hard-deleted.", error_code=INVALID_PARAMETER_VALUE, ) # Ensure that the specified experiments are old enough if older_than: non_old_experiment_ids = [ e.experiment_id for e in experiments if e.last_update_time is None or e.last_update_time >= time_threshold ] if non_old_experiment_ids: raise MlflowException( f"Experiments {non_old_experiment_ids} are not older than the required " f"age. Only experiments older than {older_than} can be deleted.", error_code=INVALID_PARAMETER_VALUE, ) experiment_ids_to_delete = list(experiment_ids) else: filter_string = f"last_update_time < {time_threshold}" if older_than else None def fetch_experiments(token=None): page = backend_store.search_experiments( view_type=ViewType.DELETED_ONLY, filter_string=filter_string, page_token=token, ) return (page + fetch_experiments(page.token)) if page.token else page experiment_ids_to_delete = [exp.experiment_id for exp in fetch_experiments()] if experiment_ids_to_delete: def fetch_runs(token=None): page = backend_store.search_runs( experiment_ids=experiment_ids_to_delete, filter_string="", run_view_type=ViewType.DELETED_ONLY, page_token=token, ) return (page + fetch_runs(page.token)) if page.token else page run_ids_to_delete.extend([run.info.run_id for run in fetch_runs()]) for run_id in set(run_ids_to_delete): try: run = backend_store.get_run(run_id) except MlflowException as exc: if ignore_not_found and exc.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST): continue raise if run.info.lifecycle_stage != LifecycleStage.DELETED: raise MlflowException( f"Run {run_id} is not in `deleted` lifecycle stage. Only runs in" " `deleted` lifecycle stage can be deleted." ) # Raise MlflowException if run_id is newer than older_than parameter if older_than and run_id not in deleted_run_ids_older_than: raise MlflowException( f"Run {run_id} is not older than the required age. " f"Only runs older than {older_than} can be deleted.", error_code=INVALID_PARAMETER_VALUE, ) artifact_repo = get_artifact_repository(run.info.artifact_uri) try: artifact_repo.delete_artifacts() except InvalidUrlException as iue: click.echo( click.style( f"An exception {iue!r} was raised during the deletion of a model artifact", fg="yellow", ) ) click.echo( click.style( f"Unable to resolve the provided artifact URL: '{artifact_repo}'. " "The gc process will continue and bypass artifact deletion. " "Please ensure that the artifact exists " "and consider manually deleting any unused artifacts. ", fg="yellow", ), ) backend_store._hard_delete_run(run_id) click.echo(f"Run with ID {run_id} has been permanently deleted.") if not skip_logged_models: for model_id in set(logged_model_ids_to_delete): # First, check if the model exists (handles non-existent models correctly) try: logged_model = backend_store.get_logged_model(model_id, allow_deleted=True) except MlflowException as exc: if ignore_not_found and exc.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST): continue raise # Model exists - now check if it's in the deleted lifecycle stage # (never skip active models, even with ignore_not_found) if model_id not in deleted_logged_model_ids: raise MlflowException( f"Logged model {model_id} is not in `deleted` lifecycle stage. " "Only logged models in `deleted` lifecycle stage can be deleted." ) if older_than and model_id not in deleted_logged_model_ids_older_than: raise MlflowException( f"Logged model {model_id} is not older than the required age. " f"Only logged models older than {older_than} can be deleted.", error_code=INVALID_PARAMETER_VALUE, ) artifact_repo = get_artifact_repository(logged_model.artifact_location) try: artifact_repo.delete_artifacts() except InvalidUrlException as iue: click.echo( click.style( f"An exception {iue!r} was raised during the deletion of a model artifact", fg="yellow", ) ) click.echo( click.style( f"Unable to resolve the provided artifact URL: '{artifact_repo}'. " "The gc process will continue and bypass artifact deletion. " "Please ensure that the artifact exists " "and consider manually deleting any unused artifacts. ", fg="yellow", ), ) backend_store._hard_delete_logged_model(model_id) click.echo(f"Logged model with ID {model_id} has been permanently deleted.") if not skip_experiments: for experiment_id in experiment_ids_to_delete: backend_store._hard_delete_experiment(experiment_id) click.echo(f"Experiment with ID {experiment_id} has been permanently deleted.") def _resolve_gc_workspaces( backend_store, all_workspaces: bool, workspace: str | None, backend_store_uri: str | None, ) -> list[str | None]: """ Determine which workspaces to iterate over for garbage collection. Args: backend_store: The tracking store instance. all_workspaces: If True, return all workspaces from the workspace store. workspace: If provided, return a single-element list with this workspace. backend_store_uri: The backend store URI for resolving workspace store. Returns: List of workspace names to iterate over, or [None] for non-workspace mode. """ supports_workspaces = ( getattr(backend_store, "supports_workspaces", False) and MLFLOW_ENABLE_WORKSPACES.get() ) if not supports_workspaces: if all_workspaces or workspace: raise MlflowException.invalid_parameter_value( "Workspace selection flags are only supported when the tracking store " "supports workspaces." ) return [None] if all_workspaces: workspace_store = get_workspace_store( resolve_workspace_store_uri(tracking_uri=backend_store_uri) ) workspaces = [ws.name for ws in workspace_store.list_workspaces()] if not workspaces: raise MlflowException( "No workspaces found. Ensure the workspace provider is configured correctly.", error_code=INVALID_PARAMETER_VALUE, ) return workspaces if workspace: return [workspace] workspace_store = get_workspace_store( resolve_workspace_store_uri(tracking_uri=backend_store_uri) ) default_workspace, supports_default = get_default_workspace_optional(workspace_store) if supports_default and default_workspace is not None: return [default_workspace.name] raise MlflowException.invalid_parameter_value( "Active workspace is required. Configure a default workspace, set MLFLOW_WORKSPACE " "or use --workspace/--all-workspaces when workspaces are enabled." ) @cli.command(short_help="Permanently delete runs in the `deleted` lifecycle stage.") @click.option( "--older-than", default=None, help="Optional. Remove run(s) older than the specified time limit. " "Specify a string in #d#h#m#s format. Float values are also supported. " "For example: --older-than 1d2h3m4s, --older-than 1.2d3h4m5s", ) @click.option( "--backend-store-uri", metavar="PATH", default=None, help="URI of the backend store from which to delete runs. Acceptable URIs are " "SQLAlchemy-compatible database connection strings " "(e.g. 'sqlite:///path/to/file.db') or local filesystem URIs " "(e.g. 'file:///absolute/path/to/directory'). By default, data will be deleted " "from the ./mlruns directory.", ) @click.option( "--artifacts-destination", envvar="MLFLOW_ARTIFACTS_DESTINATION", metavar="URI", default=None, help=( "The base artifact location from which to resolve artifact upload/download/list requests " "(e.g. 's3://my-bucket'). This option only applies when the tracking server is configured " "to stream artifacts and the experiment's artifact root location is http or " "mlflow-artifacts URI. Otherwise, the default artifact location will be used." ), ) @click.option( "--run-ids", default=None, help="Optional comma separated list of runs to be permanently deleted. If run ids" " are not specified, data is removed for all runs in the `deleted`" " lifecycle stage.", ) @click.option( "--experiment-ids", default=None, help="Optional comma separated list of experiments to be permanently deleted including " "all of their associated runs. If experiment ids are not specified, data is removed for all " "experiments in the `deleted` lifecycle stage.", ) @click.option( "--logged-model-ids", default=None, help="Optional comma separated list of logged model IDs to be permanently deleted." " If logged model IDs are not specified, data is removed for all logged models in the `deleted`" " lifecycle stage.", ) @click.option( "--jobs", is_flag=True, default=False, help="Enable job cleanup. Without this flag, no jobs will be deleted." " When enabled, all jobs are deleted unless filtered by --older-than or --job-ids." " This option only works with database backends.", ) @click.option( "--job-ids", default=None, help="Optional comma separated list of job IDs to be permanently deleted." " Can be used with or without --jobs flag." " If --older-than is also specified, only jobs matching both filters are deleted.", ) @click.option( "--tracking-uri", default=os.environ.get("MLFLOW_TRACKING_URI"), help="Tracking URI to use for deleting 'deleted' runs e.g. http://127.0.0.1:8080", ) @click.option( "--workspace", envvar=MLFLOW_WORKSPACE.name, default=None, help=( "Target workspace for deletions when workspaces are enabled. Defaults to the active " "workspace (MLFLOW_WORKSPACE)." ), ) @click.option( "--all-workspaces", is_flag=True, default=False, help="Delete deleted resources across all workspaces (workspace mode only).", ) @click.pass_context def gc( ctx, older_than, backend_store_uri, artifacts_destination, run_ids, experiment_ids, logged_model_ids, jobs, job_ids, tracking_uri, workspace, all_workspaces, ): """ Permanently delete runs in the `deleted` lifecycle stage from the specified backend store. This command deletes all artifacts and metadata associated with the specified runs. If the provided artifact URL is invalid, the artifact deletion will be bypassed, and the gc process will continue. .. attention:: If you are running an MLflow tracking server with artifact proxying enabled, you **must** set the ``MLFLOW_TRACKING_URI`` environment variable before running this command. Otherwise, the ``gc`` command will not be able to resolve artifact URIs and will not be able to delete the associated artifacts. **What gets deleted:** This command permanently removes: - **Run metadata**: Parameters, metrics, tags, and all other run information from the backend store - **Artifacts**: All files stored in the run's artifact location (models, plots, data files, etc.) - **Experiment metadata**: When deleting experiments, removes the experiment record and all associated data - **Job records**: When using the --jobs flag, removes historical job records from the jobs table .. note:: This command only considers lifecycle stage and the specified deletion criteria. It does **not** check for pinned runs, registered models, or tags. Pinning is a UI-only feature that has no effect on garbage collection. Runs must be in the `deleted` lifecycle stage before they can be permanently deleted. **Examples:** .. code-block:: bash # Delete all runs that have been in the deleted state for more than 30 days mlflow gc --older-than 30d # Delete specific runs by ID (they must be in deleted state) mlflow gc --run-ids 'run1,run2,run3' # Delete all runs in specific experiments (experiments must be in deleted state) mlflow gc --experiment-ids 'exp1,exp2' # Combine criteria: delete runs older than 7 days in specific experiments mlflow gc --older-than 7d --experiment-ids 'exp1,exp2' # Delete deleted resources across all workspaces mlflow gc --all-workspaces --older-than 30d # Delete all finalized jobs older than 7 days (requires --jobs flag) mlflow gc --jobs --older-than 7d # Delete specific jobs by ID mlflow gc --job-ids 'job1,job2,job3' """ if (workspace or all_workspaces) and not MLFLOW_ENABLE_WORKSPACES.get(): os.environ[MLFLOW_ENABLE_WORKSPACES.name] = "true" backend_store = _get_store(backend_store_uri, artifacts_destination) # Only error if --workspace was explicitly provided on CLI (not from env var) workspace_from_cli = ctx.get_parameter_source("workspace") == ParameterSource.COMMANDLINE if workspace_from_cli and all_workspaces: raise UsageError("Cannot use --workspace and --all-workspaces together.") # If --all-workspaces is set, ignore workspace from env var if all_workspaces: workspace = None skip_experiments = False skip_logged_models = False if not hasattr(backend_store, "_hard_delete_run"): raise MlflowException( "This cli can only be used with a backend that allows hard-deleting runs" ) if not hasattr(backend_store, "_hard_delete_experiment"): warnings.warn( "The specified backend does not allow hard-deleting experiments. Experiments" " will be skipped.", FutureWarning, stacklevel=2, ) skip_experiments = True if not hasattr(backend_store, "_hard_delete_logged_model"): warnings.warn( "The specified backend does not allow hard-deleting logged models. Logged models" " will be skipped.", FutureWarning, stacklevel=2, ) skip_logged_models = True time_delta = 0 if older_than is not None: regex = re.compile( r"^((?P[\.\d]+?)d)?((?P[\.\d]+?)h)?((?P[\.\d]+?)m)" r"?((?P[\.\d]+?)s)?$" ) parts = regex.match(older_than) if parts is None: raise MlflowException( f"Could not parse any time information from '{older_than}'. " "Examples of valid strings: '8h', '2d8h5m20s', '2m4s'", error_code=INVALID_PARAMETER_VALUE, ) time_params = {name: float(param) for name, param in parts.groupdict().items() if param} time_delta = int(timedelta(**time_params).total_seconds() * 1000) if tracking_uri: set_tracking_uri(tracking_uri) if not is_tracking_uri_set(): raise MlflowException( "Tracking URL is not set. Please set MLFLOW_TRACKING_URI environment variable " "or provide --tracking-uri cli option." ) # Parse comma-separated IDs into lists run_ids_list = run_ids.split(",") if run_ids else None experiment_ids_list = experiment_ids.split(",") if experiment_ids else None logged_model_ids_list = logged_model_ids.split(",") if logged_model_ids else None # Prepare job cleanup if requested (database backends only) job_store = None job_ids_list = None if jobs or job_ids: from mlflow.utils.uri import extract_db_type_from_uri store_uri = backend_store_uri or os.environ.get("MLFLOW_BACKEND_STORE_URI") try: extract_db_type_from_uri(store_uri) except MlflowException: # Not a database backend - skip job cleanup silently pass else: if MLFLOW_ENABLE_WORKSPACES.get(): from mlflow.store.jobs.sqlalchemy_workspace_store import ( WorkspaceAwareSqlAlchemyJobStore, ) job_store = WorkspaceAwareSqlAlchemyJobStore(store_uri) else: from mlflow.store.jobs.sqlalchemy_store import SqlAlchemyJobStore job_store = SqlAlchemyJobStore(store_uri) job_ids_list = job_ids.split(",") if job_ids else None for workspace_name in _resolve_gc_workspaces( backend_store=backend_store, all_workspaces=all_workspaces, workspace=workspace, backend_store_uri=backend_store_uri, ): workspace_ctx = ( workspace_context.WorkspaceContext(workspace_name) if workspace_name else contextlib.nullcontext() ) with workspace_ctx: _gc_tracking_resources( backend_store=backend_store, run_ids=run_ids_list, experiment_ids=experiment_ids_list, logged_model_ids=logged_model_ids_list, older_than=older_than, time_delta=time_delta, skip_experiments=skip_experiments, skip_logged_models=skip_logged_models, ignore_not_found=all_workspaces, ) # Clean up jobs within the same workspace context if job_store is not None: deleted_job_ids = job_store.delete_jobs(older_than=time_delta, job_ids=job_ids_list) for job_id in deleted_job_ids: click.echo(f"Job with ID {job_id} has been permanently deleted.") @cli.command(short_help="Prints out useful information for debugging issues with MLflow.") @click.option( "--mask-envs", is_flag=True, help=( "If set (the default behavior without setting this flag is not to obfuscate information), " 'mask the MLflow environment variable values (e.g. `"MLFLOW_ENV_VAR": "***"`) ' "in the output to prevent leaking sensitive information." ), ) def doctor(mask_envs): mlflow.doctor(mask_envs) cli.add_command(mlflow.deployments.cli.commands) cli.add_command(mlflow.experiments.commands) cli.add_command(mlflow.store.artifact.cli.commands) cli.add_command(mlflow.runs.commands) cli.add_command(mlflow.db.commands) from mlflow.store.fs2db.cli import migrate_filestore cli.add_command(migrate_filestore) # Add traces CLI commands from mlflow.cli import traces cli.add_command(traces.commands) # Add scorers CLI commands from mlflow.cli import scorers cli.add_command(scorers.commands) # Add datasets CLI commands from mlflow.cli import datasets cli.add_command(datasets.commands) # Add demo CLI command from mlflow.cli.demo import demo cli.add_command(demo) # Add skills CLI command try: from mlflow.cli import skills cli.add_command(skills.commands) except ImportError: pass # Add AI commands CLI cli.add_command(ai_commands.commands) try: from mlflow.mcp.cli import cli as mcp_cli cli.add_command(mcp_cli) except ImportError: pass # Add Claude Code integration commands try: import mlflow.claude_code.cli cli.add_command(mlflow.claude_code.cli.commands) except ImportError: pass # Add Assistant CLI commands try: import mlflow.assistant.cli cli.add_command(mlflow.assistant.cli.commands) except ImportError: pass # Add Agent CLI commands try: import mlflow.agent.cli cli.add_command(mlflow.agent.cli.commands) except ImportError: pass # We are conditional loading these commands since the skinny client does # not support them due to the pandas and numpy dependencies of MLflow Models try: import mlflow.models.cli cli.add_command(mlflow.models.cli.commands) except ImportError: pass try: import mlflow.sagemaker.cli cli.add_command(mlflow.sagemaker.cli.commands) except ImportError: pass with contextlib.suppress(ImportError): import mlflow.gateway.cli cli.add_command(mlflow.gateway.cli.commands) # Add crypto CLI commands with contextlib.suppress(ImportError): from mlflow.cli import crypto cli.add_command(crypto.commands) if __name__ == "__main__": cli()