Files
2026-07-13 13:22:34 +08:00

1443 lines
52 KiB
Python

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 <profile>. 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<days>[\.\d]+?)d)?((?P<hours>[\.\d]+?)h)?((?P<minutes>[\.\d]+?)m)"
r"?((?P<seconds>[\.\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()