import os import shutil import subprocess import sys import tempfile import time from pathlib import Path from unittest import mock from urllib.parse import unquote, urlparse from urllib.request import url2pathname import click import numpy as np import pytest import requests from botocore.stub import Stubber from click.testing import CliRunner import mlflow from mlflow.cli import cli, doctor, gc, server from mlflow.data import numpy_dataset from mlflow.entities import Metric, ViewType from mlflow.entities.logged_model import LoggedModelParameter, LoggedModelTag from mlflow.environment_variables import ( MLFLOW_ENABLE_WORKSPACES, MLFLOW_TRACE_ARCHIVAL_CONFIG, MLFLOW_WORKSPACE_STORE_URI, ) from mlflow.exceptions import MlflowException from mlflow.server import handlers from mlflow.store.artifact.artifact_repository_registry import get_artifact_repository from mlflow.store.jobs.sqlalchemy_store import SqlAlchemyJobStore from mlflow.store.tracking.dbmodels.models import ( SqlLoggedModelMetric, SqlLoggedModelParam, SqlLoggedModelTag, ) from mlflow.store.tracking.file_store import FileStore from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore from mlflow.utils.os import is_windows from mlflow.utils.rest_utils import augmented_raise_for_status from mlflow.utils.time import get_current_time_millis from tests.helper_functions import ( get_safe_port, kill_process_tree, ) from tests.tracking.integration_test_utils import _await_server_up_or_die @pytest.mark.parametrize("command", ["server"]) def test_mlflow_server_command(command): port = get_safe_port() with subprocess.Popen([ sys.executable, "-m", "mlflow", command, "--port", str(port), ]) as process: try: _await_server_up_or_die(port) resp = requests.get(f"http://localhost:{port}/health") augmented_raise_for_status(resp) assert resp.text == "OK" finally: kill_process_tree(process.pid) def test_server_static_prefix_validation(): with mock.patch("mlflow.server._run_server") as run_server_mock: CliRunner().invoke(server) run_server_mock.assert_called_once() with mock.patch("mlflow.server._run_server") as run_server_mock: CliRunner().invoke(server, ["--static-prefix", "/mlflow"]) run_server_mock.assert_called_once() with mock.patch("mlflow.server._run_server") as run_server_mock: result = CliRunner().invoke(server, ["--static-prefix", "mlflow/"]) assert "--static-prefix must begin with a '/'." in result.output run_server_mock.assert_not_called() with mock.patch("mlflow.server._run_server") as run_server_mock: result = CliRunner().invoke(server, ["--static-prefix", "/mlflow/"]) assert "--static-prefix should not end with a '/'." in result.output run_server_mock.assert_not_called() def test_server_cli_fails_when_workspace_env_set(): env = os.environ.copy() # Trigger server-mode guard in mlflow.server.__init__ env["_MLFLOW_SERVER_SERVE_ARTIFACTS"] = "1" env["MLFLOW_WORKSPACE"] = "team-a" result = subprocess.run( [sys.executable, "-m", "mlflow", "server", "--port", "0"], capture_output=True, text=True, env=env, ) assert result.returncode != 0 assert "MLFLOW_WORKSPACE=team-a is client-only" in result.stderr def test_server_uvicorn_options(): with mock.patch("mlflow.server._run_server") as run_server_mock: # Test default behavior (uvicorn should be used when no server options specified) CliRunner().invoke(server) run_server_mock.assert_called_once_with( file_store_path=mock.ANY, read_replica_backend_store_uri=mock.ANY, registry_store_uri=mock.ANY, default_artifact_root=mock.ANY, serve_artifacts=mock.ANY, artifacts_only=mock.ANY, artifacts_destination=mock.ANY, host="127.0.0.1", port=5000, static_prefix=None, workers=None, gunicorn_opts=None, waitress_opts=None, expose_prometheus=None, app_name=None, uvicorn_opts=None, env_file=None, secrets_cache_ttl=60, secrets_cache_max_size=1000, ) with mock.patch("mlflow.server._run_server") as run_server_mock: # Test with uvicorn-opts - use different options than dev mode CliRunner().invoke(server, ["--uvicorn-opts", "--loop asyncio --limit-concurrency 100"]) run_server_mock.assert_called_once_with( file_store_path=mock.ANY, read_replica_backend_store_uri=mock.ANY, registry_store_uri=mock.ANY, default_artifact_root=mock.ANY, serve_artifacts=mock.ANY, artifacts_only=mock.ANY, artifacts_destination=mock.ANY, host="127.0.0.1", port=5000, static_prefix=None, workers=None, gunicorn_opts=None, waitress_opts=None, expose_prometheus=None, app_name=None, uvicorn_opts="--loop asyncio --limit-concurrency 100", env_file=None, secrets_cache_ttl=60, secrets_cache_max_size=1000, ) @pytest.mark.skipif(is_windows(), reason="--dev mode is not supported on Windows") def test_server_dev_mode(): with mock.patch("mlflow.server._run_server") as run_server_mock: # Test with --dev flag (should set uvicorn opts) CliRunner().invoke(server, ["--dev"]) run_server_mock.assert_called_once_with( file_store_path=mock.ANY, read_replica_backend_store_uri=mock.ANY, registry_store_uri=mock.ANY, default_artifact_root=mock.ANY, serve_artifacts=mock.ANY, artifacts_only=mock.ANY, artifacts_destination=mock.ANY, host="127.0.0.1", port=5000, static_prefix=None, workers=None, gunicorn_opts=None, waitress_opts=None, expose_prometheus=None, app_name=None, uvicorn_opts="--reload --log-level debug", env_file=None, secrets_cache_ttl=60, secrets_cache_max_size=1000, ) @pytest.mark.skipif(is_windows(), reason="Gunicorn is not supported on Windows") def test_server_gunicorn_options(): with mock.patch("mlflow.server._run_server") as run_server_mock: # Test that gunicorn-opts disables uvicorn CliRunner().invoke(server, ["--gunicorn-opts", "--timeout 120 --max-requests 1000"]) run_server_mock.assert_called_once_with( file_store_path=mock.ANY, read_replica_backend_store_uri=mock.ANY, registry_store_uri=mock.ANY, default_artifact_root=mock.ANY, serve_artifacts=mock.ANY, artifacts_only=mock.ANY, artifacts_destination=mock.ANY, host="127.0.0.1", port=5000, static_prefix=None, workers=None, gunicorn_opts="--timeout 120 --max-requests 1000", waitress_opts=None, expose_prometheus=None, app_name=None, uvicorn_opts=None, env_file=None, secrets_cache_ttl=60, secrets_cache_max_size=1000, ) # Test conflicting options result = CliRunner().invoke( server, ["--uvicorn-opts", "--reload", "--gunicorn-opts", "--log-level debug"] ) assert result.exit_code != 0 assert "Cannot specify multiple server options" in result.output def test_server_initializes_backend_store_when_tracking_enabled(): handlers._tracking_store = None handlers._model_registry_store = None runner = CliRunner() with runner.isolated_filesystem(): with ( mock.patch("mlflow.server.handlers.initialize_backend_stores") as init_backend_mock, mock.patch("mlflow.server._run_server") as run_server_mock, ): result = runner.invoke(server) assert result.exit_code == 0 init_backend_mock.assert_called_once_with( mock.ANY, mock.ANY, mock.ANY, workspace_store_uri=None, read_replica_backend_store_uri=None ) run_server_mock.assert_called_once() def test_server_skips_backend_store_init_in_artifacts_only_mode(): handlers._tracking_store = None handlers._model_registry_store = None runner = CliRunner() with runner.isolated_filesystem(): with ( mock.patch("mlflow.server.handlers.initialize_backend_stores") as init_backend_mock, mock.patch("mlflow.server._run_server") as run_server_mock, ): result = runner.invoke(server, ["--artifacts-only"]) assert result.exit_code == 0 init_backend_mock.assert_not_called() run_server_mock.assert_called_once() def test_server_mlflow_artifacts_options(): handlers._tracking_store = None handlers._model_registry_store = None with mock.patch( "mlflow.tracking._tracking_service.utils._has_existing_mlruns_data", return_value=False ): runner = CliRunner() with runner.isolated_filesystem(): with mock.patch("mlflow.server._run_server") as run_server_mock: runner.invoke(server, ["--artifacts-only"]) run_server_mock.assert_called_once() with mock.patch("mlflow.server._run_server") as run_server_mock: runner.invoke(server, ["--serve-artifacts"]) run_server_mock.assert_called_once() with mock.patch("mlflow.server._run_server") as run_server_mock: runner.invoke(server, ["--no-serve-artifacts"]) run_server_mock.assert_called_once() with mock.patch("mlflow.server._run_server") as run_server_mock: runner.invoke(server, ["--artifacts-only"]) run_server_mock.assert_called_once() def test_server_artifacts_only_conflicts_with_enable_workspaces(): with pytest.raises( click.UsageError, match="--enable-workspaces cannot be combined with --artifacts-only" ): CliRunner().invoke( server, ["--artifacts-only", "--enable-workspaces"], catch_exceptions=False, standalone_mode=False, ) def _write_trace_archival_config( tmp_path: Path, *, enabled: bool = True, location: str = "s3://archive/default", retention: str = "30d", long_retention_allowlist: list[str] | None = None, interval_seconds: int | None = None, max_traces_per_pass: int | None = None, ) -> Path: lines = [ "trace_archival:", f" enabled: {'true' if enabled else 'false'}", f" location: {location}", f" retention: {retention}", ] if long_retention_allowlist is not None: lines.append(" long_retention_allowlist:") lines.extend(f' - "{entry}"' for entry in long_retention_allowlist) if interval_seconds is not None: lines.append(f" interval_seconds: {interval_seconds}") if max_traces_per_pass is not None: lines.append(f" max_traces_per_pass: {max_traces_per_pass}") config_path = tmp_path / "trace-archival.yaml" config_path.parent.mkdir(parents=True, exist_ok=True) config_path.write_text("\n".join(lines) + "\n", encoding="utf-8") return config_path def test_server_artifacts_only_conflicts_with_trace_archival_config(tmp_path): config_path = _write_trace_archival_config(tmp_path) with pytest.raises( click.UsageError, match="--trace-archival-config cannot be combined with --artifacts-only", ): CliRunner().invoke( server, ["--artifacts-only", "--trace-archival-config", str(config_path)], catch_exceptions=False, standalone_mode=False, ) def test_server_artifacts_only_conflicts_with_trace_archival_config_env_var(tmp_path): config_path = _write_trace_archival_config(tmp_path) MLFLOW_TRACE_ARCHIVAL_CONFIG.set(str(config_path)) try: with pytest.raises( click.UsageError, match="--trace-archival-config cannot be combined with --artifacts-only", ): CliRunner().invoke( server, ["--artifacts-only"], catch_exceptions=False, standalone_mode=False, ) finally: MLFLOW_TRACE_ARCHIVAL_CONFIG.unset() def test_server_workspace_uri_sets_env_when_workspaces_enabled(tmp_path): handlers._tracking_store = None handlers._model_registry_store = None workspace_uri = f"sqlite:///{tmp_path / 'workspace.db'}" backend_uri = f"sqlite:///{tmp_path / 'backend.db'}" artifact_root_path = tmp_path / "artifacts" artifact_root_path.mkdir() artifact_root = artifact_root_path.as_uri() MLFLOW_WORKSPACE_STORE_URI.unset() MLFLOW_ENABLE_WORKSPACES.unset() MLFLOW_TRACE_ARCHIVAL_CONFIG.unset() try: with ( mock.patch("mlflow.server._run_server") as run_server_mock, mock.patch("mlflow.server.handlers.initialize_backend_stores") as init_backend, ): result = CliRunner().invoke( server, [ "--enable-workspaces", "--workspace-store-uri", workspace_uri, "--backend-store-uri", backend_uri, "--registry-store-uri", backend_uri, "--default-artifact-root", artifact_root, ], catch_exceptions=False, standalone_mode=False, ) assert result.exit_code == 0 run_server_mock.assert_called_once() init_backend.assert_called_once_with( backend_uri, backend_uri, artifact_root, workspace_store_uri=workspace_uri, read_replica_backend_store_uri=None, ) assert MLFLOW_WORKSPACE_STORE_URI.get() == workspace_uri assert MLFLOW_ENABLE_WORKSPACES.get() is True finally: MLFLOW_WORKSPACE_STORE_URI.unset() MLFLOW_ENABLE_WORKSPACES.unset() MLFLOW_TRACE_ARCHIVAL_CONFIG.unset() def test_server_trace_archival_config_sets_env(tmp_path): handlers._tracking_store = None handlers._model_registry_store = None backend_uri = f"sqlite:///{tmp_path / 'backend.db'}" artifact_root = (tmp_path / "artifacts").as_uri() config_path = _write_trace_archival_config( tmp_path, long_retention_allowlist=["1", "2"], interval_seconds=42, max_traces_per_pass=1000, ) MLFLOW_TRACE_ARCHIVAL_CONFIG.unset() try: with ( mock.patch("mlflow.server._run_server") as run_server_mock, mock.patch("mlflow.server.handlers.initialize_backend_stores") as init_backend, ): result = CliRunner().invoke( server, [ "--backend-store-uri", backend_uri, "--registry-store-uri", backend_uri, "--default-artifact-root", artifact_root, "--trace-archival-config", str(config_path), ], catch_exceptions=False, standalone_mode=False, ) assert result.exit_code == 0 run_server_mock.assert_called_once() init_backend.assert_called_once() assert MLFLOW_TRACE_ARCHIVAL_CONFIG.get() == str(config_path) finally: MLFLOW_TRACE_ARCHIVAL_CONFIG.unset() def test_server_trace_archival_config_env_var_sets_env(tmp_path): handlers._tracking_store = None handlers._model_registry_store = None backend_uri = f"sqlite:///{tmp_path / 'backend.db'}" artifact_root = (tmp_path / "artifacts").as_uri() config_path = _write_trace_archival_config(tmp_path, long_retention_allowlist=["1", "2", "3"]) MLFLOW_TRACE_ARCHIVAL_CONFIG.unset() try: MLFLOW_TRACE_ARCHIVAL_CONFIG.set(str(config_path)) with ( mock.patch("mlflow.server._run_server") as run_server_mock, mock.patch("mlflow.server.handlers.initialize_backend_stores") as init_backend, ): result = CliRunner().invoke( server, [ "--backend-store-uri", backend_uri, "--registry-store-uri", backend_uri, "--default-artifact-root", artifact_root, ], catch_exceptions=False, standalone_mode=False, ) assert result.exit_code == 0 run_server_mock.assert_called_once() init_backend.assert_called_once() assert MLFLOW_TRACE_ARCHIVAL_CONFIG.get() == str(config_path) finally: MLFLOW_TRACE_ARCHIVAL_CONFIG.unset() def test_server_trace_archival_config_cli_overrides_env_var(tmp_path): handlers._tracking_store = None handlers._model_registry_store = None backend_uri = f"sqlite:///{tmp_path / 'backend.db'}" artifact_root = (tmp_path / "artifacts").as_uri() env_config_path = _write_trace_archival_config(tmp_path / "env") cli_config_path = _write_trace_archival_config(tmp_path / "cli", interval_seconds=42) MLFLOW_TRACE_ARCHIVAL_CONFIG.unset() try: MLFLOW_TRACE_ARCHIVAL_CONFIG.set(str(env_config_path)) with ( mock.patch("mlflow.server._run_server") as run_server_mock, mock.patch("mlflow.server.handlers.initialize_backend_stores") as init_backend, ): result = CliRunner().invoke( server, [ "--backend-store-uri", backend_uri, "--registry-store-uri", backend_uri, "--default-artifact-root", artifact_root, "--trace-archival-config", str(cli_config_path), ], catch_exceptions=False, standalone_mode=False, ) assert result.exit_code == 0 run_server_mock.assert_called_once() init_backend.assert_called_once() assert MLFLOW_TRACE_ARCHIVAL_CONFIG.get() == str(cli_config_path) finally: MLFLOW_TRACE_ARCHIVAL_CONFIG.unset() def test_server_trace_archival_scheduler_defaults_disabled(tmp_path): handlers._tracking_store = None handlers._model_registry_store = None backend_uri = f"sqlite:///{tmp_path / 'backend.db'}" artifact_root = (tmp_path / "artifacts").as_uri() MLFLOW_TRACE_ARCHIVAL_CONFIG.unset() try: with ( mock.patch("mlflow.server._run_server") as run_server_mock, mock.patch("mlflow.server.handlers.initialize_backend_stores") as init_backend, ): result = CliRunner().invoke( server, [ "--backend-store-uri", backend_uri, "--registry-store-uri", backend_uri, "--default-artifact-root", artifact_root, ], catch_exceptions=False, standalone_mode=False, ) assert result.exit_code == 0 run_server_mock.assert_called_once() init_backend.assert_called_once() finally: MLFLOW_TRACE_ARCHIVAL_CONFIG.unset() def test_server_trace_archival_config_missing_file(tmp_path): result = CliRunner().invoke( server, ["--trace-archival-config", str(tmp_path / "missing.yaml")], ) assert result.exit_code != 0 assert "does not exist" in result.output def test_server_trace_archival_config_rejects_invalid_yaml(tmp_path): config_path = tmp_path / "trace-archival.yaml" config_path.write_text("trace_archival: [\n", encoding="utf-8") result = CliRunner().invoke(server, ["--trace-archival-config", str(config_path)]) assert result.exit_code != 0 assert "Invalid trace archival config file" in result.output assert "Failed to parse YAML" in result.output def test_server_trace_archival_config_rejects_invalid_experiment_id(tmp_path): config_path = _write_trace_archival_config( tmp_path, long_retention_allowlist=["1", "invalid id"], ) result = CliRunner().invoke(server, ["--trace-archival-config", str(config_path)]) assert result.exit_code != 0 assert "Invalid experiment ID: 'invalid id'" in result.output def test_server_trace_archival_config_rejects_invalid_interval(tmp_path): config_path = _write_trace_archival_config(tmp_path, interval_seconds=0) result = CliRunner().invoke(server, ["--trace-archival-config", str(config_path)]) assert result.exit_code != 0 assert "'trace_archival.interval_seconds' must be a positive integer" in result.output def test_server_trace_archival_config_rejects_invalid_max_traces_per_pass(tmp_path): config_path = _write_trace_archival_config(tmp_path, max_traces_per_pass=0) result = CliRunner().invoke(server, ["--trace-archival-config", str(config_path)]) assert result.exit_code != 0 assert "'trace_archival.max_traces_per_pass' must be a positive integer" in result.output def test_server_trace_archival_config_requires_archival_capable_tracking_store(tmp_path): handlers._tracking_store = None handlers._model_registry_store = None backend_uri = f"sqlite:///{tmp_path / 'backend.db'}" artifact_root = (tmp_path / "artifacts").as_uri() config_path = _write_trace_archival_config(tmp_path) mock_tracking_store = mock.Mock(supports_trace_archival=False) with ( mock.patch("mlflow.server._run_server") as run_server_mock, mock.patch("mlflow.server.handlers._get_tracking_store", return_value=mock_tracking_store), mock.patch("mlflow.server.handlers._get_model_registry_store", return_value=None), ): result = CliRunner().invoke( server, [ "--backend-store-uri", backend_uri, "--registry-store-uri", backend_uri, "--default-artifact-root", artifact_root, "--trace-archival-config", str(config_path), ], ) assert result.exit_code != 0 assert "Error initializing backend store" in result.output run_server_mock.assert_not_called() def test_server_trace_archival_config_allows_disabled_config_on_unsupported_tracking_store( tmp_path, ): handlers._tracking_store = None handlers._model_registry_store = None backend_uri = f"sqlite:///{tmp_path / 'backend.db'}" artifact_root = (tmp_path / "artifacts").as_uri() config_path = _write_trace_archival_config(tmp_path, enabled=False) mock_tracking_store = mock.Mock(supports_trace_archival=False) MLFLOW_TRACE_ARCHIVAL_CONFIG.unset() try: with ( mock.patch("mlflow.server._run_server") as run_server_mock, mock.patch( "mlflow.server.handlers._get_tracking_store", return_value=mock_tracking_store ), mock.patch("mlflow.server.handlers._get_model_registry_store", return_value=None), ): result = CliRunner().invoke( server, [ "--backend-store-uri", backend_uri, "--registry-store-uri", backend_uri, "--default-artifact-root", artifact_root, "--trace-archival-config", str(config_path), ], ) assert result.exit_code == 0 run_server_mock.assert_called_once() finally: MLFLOW_TRACE_ARCHIVAL_CONFIG.unset() @pytest.mark.parametrize("command", [server]) def test_tracking_uri_validation_failure(command): handlers._tracking_store = None with mock.patch("mlflow.server._run_server") as run_server_mock: # SQLAlchemy expects postgresql:// not postgres:// CliRunner().invoke( command, [ "--backend-store-uri", "postgres://user:pwd@host:5432/mydb", "--default-artifact-root", "./mlruns", ], ) run_server_mock.assert_not_called() @pytest.mark.parametrize("command", [server]) def test_tracking_uri_validation_sql_driver_uris(command): handlers._tracking_store = None handlers._model_registry_store = None with ( mock.patch("mlflow.server._run_server") as run_server_mock, mock.patch("mlflow.store.tracking.sqlalchemy_store.SqlAlchemyStore"), mock.patch("mlflow.store.model_registry.sqlalchemy_store.SqlAlchemyStore"), ): result = CliRunner().invoke( command, [ "--backend-store-uri", "mysql+pymysql://user:pwd@host:5432/mydb", "--default-artifact-root", "./mlruns", ], ) assert result.exit_code == 0 run_server_mock.assert_called() # Clean up the global variables set by the server mlflow.set_tracking_uri(None) mlflow.set_registry_uri(None) @pytest.mark.parametrize("command", [server]) def test_registry_store_uri_different_from_tracking_store(command): handlers._tracking_store = None handlers._model_registry_store = None from mlflow.server.handlers import ( ModelRegistryStoreRegistryWrapper, TrackingStoreRegistryWrapper, ) handlers._tracking_store_registry = TrackingStoreRegistryWrapper() handlers._model_registry_store_registry = ModelRegistryStoreRegistryWrapper() with ( mock.patch("mlflow.server._run_server") as run_server_mock, mock.patch("mlflow.store.tracking.file_store.FileStore") as tracking_store, mock.patch( "mlflow.store.model_registry.sqlalchemy_store.SqlAlchemyStore" ) as registry_store, ): result = CliRunner().invoke( command, [ "--backend-store-uri", "./mlruns", "--registry-store-uri", "mysql://user:pwd@host:5432/mydb", ], ) assert result.exit_code == 0 run_server_mock.assert_called() tracking_store.assert_called() registry_store.assert_called() # Clean up the global variables set by the server mlflow.set_tracking_uri(None) mlflow.set_registry_uri(None) @pytest.fixture def sqlite_store(db_uri: str, tmp_path: Path) -> tuple[SqlAlchemyStore, str]: artifact_uri = (tmp_path / "artifacts").as_uri() store = SqlAlchemyStore(db_uri, artifact_uri) return (store, db_uri) @pytest.fixture def file_store(): pytest.skip("FileStore is no longer supported.") ROOT_LOCATION = os.path.join(tempfile.gettempdir(), "test_mlflow_gc") file_store_uri = f"file:///{ROOT_LOCATION}" yield (FileStore(ROOT_LOCATION), file_store_uri) shutil.rmtree(ROOT_LOCATION) def _create_run_in_store(store, create_artifacts=True): config = { "experiment_id": "0", "user_id": "Anderson", "start_time": get_current_time_millis(), "tags": [], "run_name": "name", } run = store.create_run(**config) if create_artifacts: artifact_path = url2pathname(unquote(urlparse(run.info.artifact_uri).path)) if not os.path.exists(artifact_path): os.makedirs(artifact_path) return run @pytest.mark.parametrize("create_artifacts_in_run", [True, False]) def test_mlflow_gc_sqlite(sqlite_store, create_artifacts_in_run): store = sqlite_store[0] run = _create_run_in_store(store, create_artifacts=create_artifacts_in_run) store.delete_run(run.info.run_id) subprocess.check_call([ sys.executable, "-m", "mlflow", "gc", "--backend-store-uri", sqlite_store[1], ]) runs = store.search_runs(experiment_ids=["0"], filter_string="", run_view_type=ViewType.ALL) assert len(runs) == 0 with pytest.raises(MlflowException, match=r"Run .+ not found"): store.get_run(run.info.run_id) artifact_path = url2pathname(unquote(urlparse(run.info.artifact_uri).path)) assert not os.path.exists(artifact_path) def test_mlflow_gc_sqlite_older_than(sqlite_store): store = sqlite_store[0] run = _create_run_in_store(store) store.delete_run(run.info.run_id) with pytest.raises(subprocess.CalledProcessError, match=r".+") as exp: subprocess.run( [ sys.executable, "-m", "mlflow", "gc", "--backend-store-uri", sqlite_store[1], "--older-than", "10d10h10m10s", "--run-ids", run.info.run_id, ], check=True, capture_output=True, text=True, ) assert "is not older than the required age" in exp.value.stderr runs = store.search_runs(experiment_ids=["0"], filter_string="", run_view_type=ViewType.ALL) assert len(runs) == 1 time.sleep(1) subprocess.check_output([ sys.executable, "-m", "mlflow", "gc", "--backend-store-uri", sqlite_store[1], "--older-than", "1s", "--run-ids", run.info.run_id, ]) runs = store.search_runs(experiment_ids=["0"], filter_string="", run_view_type=ViewType.ALL) assert len(runs) == 0 @pytest.mark.parametrize("create_artifacts_in_run", [True, False]) def test_mlflow_gc_file_store(file_store, create_artifacts_in_run): pytest.skip("FileStore is no longer supported.") store = file_store[0] run = _create_run_in_store(store, create_artifacts=create_artifacts_in_run) store.delete_run(run.info.run_id) subprocess.check_output([ sys.executable, "-m", "mlflow", "gc", "--backend-store-uri", file_store[1], ]) runs = store.search_runs(experiment_ids=["0"], filter_string="", run_view_type=ViewType.ALL) assert len(runs) == 0 with pytest.raises(MlflowException, match=r"Run .+ not found"): store.get_run(run.info.run_id) artifact_path = url2pathname(unquote(urlparse(run.info.artifact_uri).path)) assert not os.path.exists(artifact_path) def test_mlflow_gc_file_store_passing_explicit_run_ids(file_store): pytest.skip("FileStore is no longer supported.") store = file_store[0] run = _create_run_in_store(store) store.delete_run(run.info.run_id) subprocess.check_output([ sys.executable, "-m", "mlflow", "gc", "--backend-store-uri", file_store[1], "--run-ids", run.info.run_id, ]) runs = store.search_runs(experiment_ids=["0"], filter_string="", run_view_type=ViewType.ALL) assert len(runs) == 0 with pytest.raises(MlflowException, match=r"Run .+ not found"): store.get_run(run.info.run_id) def test_mlflow_gc_not_deleted_run(file_store): pytest.skip("FileStore is no longer supported.") store = file_store[0] run = _create_run_in_store(store) with pytest.raises(subprocess.CalledProcessError, match=r".+"): subprocess.check_output([ sys.executable, "-m", "mlflow", "gc", "--backend-store-uri", file_store[1], "--run-ids", run.info.run_id, ]) runs = store.search_runs(experiment_ids=["0"], filter_string="", run_view_type=ViewType.ALL) assert len(runs) == 1 def test_mlflow_gc_file_store_older_than(file_store): pytest.skip("FileStore is no longer supported.") store = file_store[0] run = _create_run_in_store(store) store.delete_run(run.info.run_id) with pytest.raises(subprocess.CalledProcessError, match=r".+") as exp: subprocess.run( [ sys.executable, "-m", "mlflow", "gc", "--backend-store-uri", file_store[1], "--older-than", "10d10h10m10s", "--run-ids", run.info.run_id, ], check=True, capture_output=True, text=True, ) assert "is not older than the required age" in exp.value.stderr runs = store.search_runs(experiment_ids=["0"], filter_string="", run_view_type=ViewType.ALL) assert len(runs) == 1 time.sleep(1) subprocess.check_output([ sys.executable, "-m", "mlflow", "gc", "--backend-store-uri", file_store[1], "--older-than", "1s", "--run-ids", run.info.run_id, ]) runs = store.search_runs(experiment_ids=["0"], filter_string="", run_view_type=ViewType.ALL) assert len(runs) == 0 @pytest.mark.parametrize("get_store_details", ["file_store", "sqlite_store"]) def test_mlflow_gc_experiments(get_store_details, request): if get_store_details == "file_store": pytest.skip("FileStore is no longer supported.") def invoke_gc(*args): return CliRunner().invoke(gc, args, catch_exceptions=False) store, uri = request.getfixturevalue(get_store_details) exp_id_1 = store.create_experiment("1") run_id_1 = store.create_run(exp_id_1, user_id="user", start_time=0, tags=[], run_name="1") invoke_gc("--backend-store-uri", uri) experiments = store.search_experiments(view_type=ViewType.ALL) exp_ids = [e.experiment_id for e in experiments] runs = store.search_runs(experiment_ids=exp_ids, filter_string="", run_view_type=ViewType.ALL) assert sorted(exp_ids) == sorted([exp_id_1, store.DEFAULT_EXPERIMENT_ID]) assert [r.info.run_id for r in runs] == [run_id_1.info.run_id] store.delete_experiment(exp_id_1) invoke_gc("--backend-store-uri", uri) experiments = store.search_experiments(view_type=ViewType.ALL) runs = store.search_runs(experiment_ids=exp_ids, filter_string="", run_view_type=ViewType.ALL) assert [e.experiment_id for e in experiments] == [store.DEFAULT_EXPERIMENT_ID] assert runs == [] exp_id_2 = store.create_experiment("2") exp_id_3 = store.create_experiment("3") store.delete_experiment(exp_id_2) store.delete_experiment(exp_id_3) invoke_gc("--backend-store-uri", uri, "--experiment-ids", exp_id_2) experiments = store.search_experiments(view_type=ViewType.ALL) assert sorted([e.experiment_id for e in experiments]) == sorted([ exp_id_3, store.DEFAULT_EXPERIMENT_ID, ]) with mock.patch("time.time", return_value=0) as mock_time: exp_id_4 = store.create_experiment("4") store.delete_experiment(exp_id_4) mock_time.assert_called() invoke_gc("--backend-store-uri", uri, "--older-than", "1d") experiments = store.search_experiments(view_type=ViewType.ALL) assert sorted([e.experiment_id for e in experiments]) == sorted([ exp_id_3, store.DEFAULT_EXPERIMENT_ID, ]) invoke_gc("--backend-store-uri", uri, "--experiment-ids", exp_id_3, "--older-than", "0s") experiments = store.search_experiments(view_type=ViewType.ALL) assert [e.experiment_id for e in experiments] == [store.DEFAULT_EXPERIMENT_ID] exp_id_5 = store.create_experiment("5") store.delete_experiment(exp_id_5) with pytest.raises(MlflowException, match=r"Experiments .+ can be deleted."): invoke_gc( "--backend-store-uri", uri, "--experiment-ids", exp_id_5, "--older-than", "10d10h10m10s", ) experiments = store.search_experiments(view_type=ViewType.ALL) assert sorted([e.experiment_id for e in experiments]) == sorted([ exp_id_5, store.DEFAULT_EXPERIMENT_ID, ]) def test_mlflow_gc_experiment_with_logged_model_params_tags_and_metrics(sqlite_store): store, uri = sqlite_store exp_id = store.create_experiment("exp_with_logged_model") run = store.create_run(exp_id, user_id="user", start_time=0, tags=[], run_name="run") run_id = run.info.run_id model = store.create_logged_model(experiment_id=exp_id) store.log_logged_model_params( model_id=model.model_id, params=[ LoggedModelParameter(key="param1", value="value1"), LoggedModelParameter(key="param2", value="value2"), ], ) store.set_logged_model_tags( model_id=model.model_id, tags=[ LoggedModelTag(key="tag1", value="value1"), LoggedModelTag(key="tag2", value="value2"), ], ) store._log_model_metrics( run_id=run_id, metrics=[Metric(key="m1", value=1.0, timestamp=0, step=0, model_id=model.model_id)], experiment_id=exp_id, ) store.delete_experiment(exp_id) result = CliRunner().invoke(gc, ["--backend-store-uri", uri], catch_exceptions=False) assert result.exit_code == 0 experiments = store.search_experiments(view_type=ViewType.ALL) assert [e.experiment_id for e in experiments] == [store.DEFAULT_EXPERIMENT_ID] with pytest.raises(MlflowException, match=r".+ not found"): store.get_logged_model(model.model_id) with store.ManagedSessionMaker() as session: for table in [SqlLoggedModelParam, SqlLoggedModelTag, SqlLoggedModelMetric]: assert session.query(table).count() == 0 @pytest.fixture def sqlite_store_with_s3_artifact_repository( tmp_path: Path, ) -> tuple[SqlAlchemyStore, str, str]: db_path = tmp_path / "mlflow.db" db_uri = f"sqlite:///{db_path}" s3_uri = "s3://mlflow" store = SqlAlchemyStore(db_uri, s3_uri) return (store, db_uri, s3_uri) def test_mlflow_gc_sqlite_with_s3_artifact_repository( sqlite_store_with_s3_artifact_repository, ): store = sqlite_store_with_s3_artifact_repository[0] run = _create_run_in_store(store, create_artifacts=False) store.delete_run(run.info.run_id) artifact_repo = get_artifact_repository(run.info.artifact_uri) bucket, dest_path = artifact_repo.parse_s3_compliant_uri(run.info.artifact_uri) fake_artifact_path = os.path.join(dest_path, "fake_artifact.txt") with Stubber(artifact_repo._get_s3_client()) as s3_stubber: s3_stubber.add_response( "list_objects_v2", {"Contents": [{"Key": fake_artifact_path}]}, {"Bucket": bucket, "Prefix": dest_path}, ) s3_stubber.add_response( "delete_objects", {"Deleted": [{"Key": fake_artifact_path}]}, {"Bucket": bucket, "Delete": {"Objects": [{"Key": fake_artifact_path}]}}, ) CliRunner().invoke( gc, [ "--backend-store-uri", sqlite_store_with_s3_artifact_repository[1], "--artifacts-destination", sqlite_store_with_s3_artifact_repository[2], ], catch_exceptions=False, ) runs = store.search_runs(experiment_ids=["0"], filter_string="", run_view_type=ViewType.ALL) assert len(runs) == 0 with pytest.raises(MlflowException, match=r"Run .+ not found"): store.get_run(run.info.run_id) def test_mlflow_tracking_disabled_in_artifacts_only_mode(tmp_path: Path): port = get_safe_port() env = {**os.environ} env.pop(MLFLOW_TRACE_ARCHIVAL_CONFIG.name, None) with subprocess.Popen( [sys.executable, "-m", "mlflow", "server", "--port", str(port), "--artifacts-only"], cwd=tmp_path, env=env, ) as process: try: _await_server_up_or_die(port) resp = requests.get(f"http://localhost:{port}/api/2.0/mlflow/experiments/search") assert ( "Endpoint: /api/2.0/mlflow/experiments/search disabled due to the mlflow " "server running in `--artifacts-only` mode." in resp.text ) finally: kill_process_tree(process.pid) def test_mlflow_artifact_list_in_artifacts_only_mode(tmp_path: Path): port = get_safe_port() env = {**os.environ} env.pop(MLFLOW_TRACE_ARCHIVAL_CONFIG.name, None) with subprocess.Popen( [sys.executable, "-m", "mlflow", "server", "--port", str(port), "--artifacts-only"], cwd=tmp_path, env=env, ) as process: try: _await_server_up_or_die(port) resp = requests.get(f"http://localhost:{port}/api/2.0/mlflow-artifacts/artifacts") augmented_raise_for_status(resp) assert resp.status_code == 200 assert resp.text == "{}" finally: kill_process_tree(process.pid) def test_mlflow_artifact_service_unavailable_when_no_server_artifacts_is_specified(): port = get_safe_port() with subprocess.Popen([ sys.executable, "-m", "mlflow", "server", "--port", str(port), "--no-serve-artifacts", ]) as process: try: _await_server_up_or_die(port) endpoint = "/api/2.0/mlflow-artifacts/artifacts" resp = requests.get(f"http://localhost:{port}{endpoint}") assert ( f"Endpoint: {endpoint} disabled due to the mlflow server running with " "`--no-serve-artifacts`" in resp.text ) finally: kill_process_tree(process.pid) def test_mlflow_artifact_only_prints_warning_for_configs(): with ( mock.patch("mlflow.server._run_server") as run_server_mock, mock.patch("mlflow.store.tracking.sqlalchemy_store.SqlAlchemyStore"), mock.patch("mlflow.store.model_registry.sqlalchemy_store.SqlAlchemyStore"), ): result = CliRunner().invoke( server, ["--artifacts-only", "--backend-store-uri", "sqlite:///my.db"], catch_exceptions=False, ) msg = ( "Usage: server [OPTIONS]\nTry 'server --help' for help.\n\nError: You are starting a " "tracking server in `--artifacts-only` mode and have provided a value for " "`--backend_store_uri`" ) assert msg in result.output assert result.exit_code != 0 run_server_mock.assert_not_called() def test_mlflow_ui_is_alias_for_mlflow_server(): mlflow_ui_stdout = subprocess.check_output( [sys.executable, "-m", "mlflow", "ui", "--help"], text=True ) mlflow_server_stdout = subprocess.check_output( [sys.executable, "-m", "mlflow", "server", "--help"], text=True ) assert ( mlflow_ui_stdout.replace("Usage: python -m mlflow ui", "Usage: python -m mlflow server") == mlflow_server_stdout ) def test_server_help_has_quickstart(): out = subprocess.check_output([sys.executable, "-m", "mlflow", "server", "--help"], text=True) assert "Pick your setup" in out assert "Options by topic:" in out assert "sqlite:///mlflow.db" in out # Curated guidance is rendered before the exhaustive flat option list. assert out.index("Pick your setup") < out.index("\nOptions:") def test_cli_with_python_mod(): stdout = subprocess.check_output( [ sys.executable, "-m", "mlflow", "--version", ], text=True, ) assert stdout.rstrip().endswith(mlflow.__version__) stdout = subprocess.check_output( [ sys.executable, "-m", "mlflow", "server", "--help", ], text=True, ) assert "mlflow server" in stdout def test_doctor(): res = CliRunner().invoke(doctor, catch_exceptions=False) assert res.exit_code == 0 def test_env_file_loading(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: # Setup: Create an experiment using the Python SDK test_tracking_uri = f"sqlite:///{tmp_path / 'mlflow.db'}" test_experiment_name = "test_experiment_from_env" # Create experiment using SDK mlflow.set_tracking_uri(test_tracking_uri) mlflow.create_experiment(test_experiment_name) # Create a test .env file pointing to this tracking URI env_file_path = tmp_path / "test.env" env_file_path.write_text(f"MLFLOW_TRACKING_URI={test_tracking_uri}\n") runner = CliRunner() # Clear the MLflow environment variables to ensure we're testing the loading monkeypatch.delenv("MLFLOW_TRACKING_URI", raising=False) monkeypatch.delenv("MLFLOW_EXPERIMENT_NAME", raising=False) # Ensure variables are not set before running command assert "MLFLOW_TRACKING_URI" not in os.environ # Use the existing experiments search CLI command with --env-file result = runner.invoke( cli, ["--env-file", str(env_file_path), "experiments", "search"], catch_exceptions=False, ) # Check that the command executed successfully assert result.exit_code == 0 # Verify the experiment we created is found (proves env vars were loaded) assert test_experiment_name in result.output # Verify the loading message assert "Loaded environment variables from:" in result.output assert str(env_file_path) in result.output def test_env_file_loading_invalid_path() -> None: runner = CliRunner() # Test error handling for non-existent file result = runner.invoke( cli, ["--env-file", "nonexistent.env", "experiments", "search"], catch_exceptions=False, ) assert result.exit_code != 0 assert "Environment file 'nonexistent.env' does not exist" in result.output @pytest.mark.parametrize( ("args", "expected_max_results", "expected_view"), [ (["experiments", "search"], None, ViewType.ACTIVE_ONLY), (["experiments", "search", "--max-results", "50"], 50, ViewType.ACTIVE_ONLY), ( ["experiments", "search", "--view", "all", "--max-results", "100"], 100, ViewType.ALL, ), ], ) def test_experiments_search_max_results(args, expected_max_results, expected_view): with mock.patch("mlflow.search_experiments", return_value=[]) as mock_search: result = CliRunner().invoke(cli, args, catch_exceptions=False) assert result.exit_code == 0 mock_search.assert_called_once() call_kwargs = mock_search.call_args[1] assert call_kwargs.get("max_results") == expected_max_results assert call_kwargs.get("view_type") == expected_view def test_experiments_search_max_results_zero(): with mock.patch("mlflow.search_experiments", return_value=[]) as mock_search: result = CliRunner().invoke( cli, ["experiments", "search", "--max-results", "0"], catch_exceptions=False ) assert result.exit_code == 0 mock_search.assert_called_once() call_kwargs = mock_search.call_args[1] assert call_kwargs.get("max_results") == 0 def test_experiments_search_max_results_negative(): result = CliRunner().invoke( cli, ["experiments", "search", "--max-results", "-1"], catch_exceptions=False ) assert result.exit_code != 0 assert "max-results must be a non-negative integer" in result.output def test_server_with_env_file(tmp_path): env_file = tmp_path / ".env" env_file.write_text("TEST_VAR=test_value\n") with mock.patch("mlflow.server._run_server") as run_server_mock: result = CliRunner().invoke(cli, ["--env-file", str(env_file), "server"]) assert result.exit_code == 0 run_server_mock.assert_called_once() # Verify env_file parameter is passed assert run_server_mock.call_args.kwargs["env_file"] == str(env_file) def test_mlflow_gc_with_datasets(sqlite_store): store = sqlite_store[0] mlflow.set_tracking_uri(sqlite_store[1]) mlflow.set_experiment("dataset") dataset = numpy_dataset.from_numpy(np.array([1, 2, 3])) with mlflow.start_run() as run: experiment_id = run.info.experiment_id mlflow.log_input(dataset) experiments = store.search_experiments(view_type=ViewType.ALL) # default and datasets assert len(experiments) == 2 store.delete_experiment(experiment_id) # the new experiment is only marked as deleted, not removed experiments = store.search_experiments(view_type=ViewType.ALL) assert len(experiments) == 2 subprocess.check_call([ sys.executable, "-m", "mlflow", "gc", "--backend-store-uri", sqlite_store[1], ]) experiments = store.search_experiments(view_type=ViewType.ALL) # only default is left after GC assert len(experiments) == 1 assert experiments[0].experiment_id == "0" with pytest.raises(MlflowException, match=f"No Experiment with id={experiment_id} exists"): store.get_experiment(experiment_id) @pytest.mark.parametrize("get_store_details", ["file_store", "sqlite_store"]) def test_mlflow_gc_logged_model(get_store_details, request): if get_store_details == "file_store": pytest.skip("FileStore is no longer supported.") store, uri = request.getfixturevalue(get_store_details) exp_id = store.create_experiment("exp") model = store.create_logged_model(experiment_id=exp_id) store.delete_logged_model(model.model_id) subprocess.check_output([sys.executable, "-m", "mlflow", "gc", "--backend-store-uri", uri]) with pytest.raises(MlflowException, match=r".+ not found"): store.get_logged_model(model.model_id) @pytest.mark.parametrize("get_store_details", ["file_store", "sqlite_store"]) def test_mlflow_gc_logged_models_older_than(get_store_details, request): if get_store_details == "file_store": pytest.skip("FileStore is no longer supported.") store, uri = request.getfixturevalue(get_store_details) exp_id = store.create_experiment("exp") old_time = time.time() - (2 * 24 * 60 * 60) with mock.patch("time.time", return_value=old_time): model = store.create_logged_model(experiment_id=exp_id) store.delete_logged_model(model.model_id) subprocess.check_call([ sys.executable, "-m", "mlflow", "gc", "--backend-store-uri", uri, "--older-than", "1d", ]) retrieved_model = store.get_logged_model(model.model_id, allow_deleted=True) assert retrieved_model.model_id == model.model_id @pytest.mark.parametrize("get_store_details", ["file_store", "sqlite_store"]) def test_mlflow_gc_logged_models_deletes_when_older_than(get_store_details, request): if get_store_details == "file_store": pytest.skip("FileStore is no longer supported.") store, uri = request.getfixturevalue(get_store_details) exp_id = store.create_experiment("exp") model = store.create_logged_model(experiment_id=exp_id) old_deletion_ms = int((time.time() - (2 * 24 * 60 * 60)) * 1000) with mock.patch("mlflow.utils.time.get_current_time_millis", return_value=old_deletion_ms): store.delete_logged_model(model.model_id) subprocess.check_call([ sys.executable, "-m", "mlflow", "gc", "--backend-store-uri", uri, "--older-than", "1d", ]) with pytest.raises(MlflowException, match=r".+ not found"): store.get_logged_model(model.model_id) @pytest.mark.parametrize("get_store_details", ["file_store", "sqlite_store"]) def test_mlflow_gc_logged_models_mixed_time(get_store_details, request): if get_store_details == "file_store": pytest.skip("FileStore is no longer supported.") store, uri = request.getfixturevalue(get_store_details) exp_id = store.create_experiment("exp") old_model = store.create_logged_model(experiment_id=exp_id) recent_model = store.create_logged_model(experiment_id=exp_id) old_deletion_ms = int((time.time() - (3 * 24 * 60 * 60)) * 1000) with mock.patch("mlflow.utils.time.get_current_time_millis", return_value=old_deletion_ms): store.delete_logged_model(old_model.model_id) with mock.patch("mlflow.utils.time.get_current_time_millis") as current_time_mock: current_time_mock.return_value = int(time.time() * 1000) store.delete_logged_model(recent_model.model_id) subprocess.check_call([ sys.executable, "-m", "mlflow", "gc", "--backend-store-uri", uri, "--older-than", "1d", ]) with pytest.raises(MlflowException, match=r".+ not found"): store.get_logged_model(old_model.model_id) retrieved_model = store.get_logged_model(recent_model.model_id, allow_deleted=True) assert retrieved_model.model_id == recent_model.model_id @pytest.fixture def sqlite_store_with_jobs( db_uri: str, tmp_path: Path ) -> tuple[SqlAlchemyStore, SqlAlchemyJobStore, str]: artifact_uri = (tmp_path / "artifacts").as_uri() tracking_store = SqlAlchemyStore(db_uri, artifact_uri) job_store = SqlAlchemyJobStore(db_uri) return (tracking_store, job_store, db_uri) def _create_test_job(job_store, job_name="test_job", finalize=True): job = job_store.create_job(job_name=job_name, params="{}") if finalize: job_store.start_job(job.job_id) job_store.finish_job(job.job_id, result="done") return job_store.get_job(job.job_id) return job def _set_job_creation_time(job_store, job_id, creation_time_ms): from mlflow.store.tracking.dbmodels.models import SqlJob with job_store.ManagedSessionMaker(read_only=False) as session: job = session.query(SqlJob).filter(SqlJob.id == job_id).first() job.creation_time = creation_time_ms def test_mlflow_gc_no_jobs_deleted_without_jobs_flag(sqlite_store_with_jobs): _, job_store, db_uri = sqlite_store_with_jobs job1 = _create_test_job(job_store, "job1") job2 = _create_test_job(job_store, "job2") subprocess.check_call([ sys.executable, "-m", "mlflow", "gc", "--backend-store-uri", db_uri, ]) retrieved_job1 = job_store.get_job(job1.job_id) retrieved_job2 = job_store.get_job(job2.job_id) assert retrieved_job1.job_id == job1.job_id assert retrieved_job2.job_id == job2.job_id def test_mlflow_gc_all_jobs_deleted_with_jobs_flag(sqlite_store_with_jobs): _, job_store, db_uri = sqlite_store_with_jobs job1 = _create_test_job(job_store, "job1") job2 = _create_test_job(job_store, "job2") job3 = _create_test_job(job_store, "job3") subprocess.check_call([ sys.executable, "-m", "mlflow", "gc", "--backend-store-uri", db_uri, "--jobs", ]) with pytest.raises(MlflowException, match=r"Job .+ not found"): job_store.get_job(job1.job_id) with pytest.raises(MlflowException, match=r"Job .+ not found"): job_store.get_job(job2.job_id) with pytest.raises(MlflowException, match=r"Job .+ not found"): job_store.get_job(job3.job_id) def test_mlflow_gc_specific_jobs_deleted_with_job_ids(sqlite_store_with_jobs): _, job_store, db_uri = sqlite_store_with_jobs job1 = _create_test_job(job_store, "job1") job2 = _create_test_job(job_store, "job2") job3 = _create_test_job(job_store, "job3") subprocess.check_call([ sys.executable, "-m", "mlflow", "gc", "--backend-store-uri", db_uri, "--job-ids", f"{job1.job_id},{job2.job_id}", ]) with pytest.raises(MlflowException, match=r"Job .+ not found"): job_store.get_job(job1.job_id) with pytest.raises(MlflowException, match=r"Job .+ not found"): job_store.get_job(job2.job_id) retrieved_job3 = job_store.get_job(job3.job_id) assert retrieved_job3.job_id == job3.job_id def test_mlflow_gc_jobs_deleted_with_older_than_flag(sqlite_store_with_jobs): _, job_store, db_uri = sqlite_store_with_jobs old_job = _create_test_job(job_store, "old_job") new_job = _create_test_job(job_store, "new_job") one_hour_ago = get_current_time_millis() - (60 * 60 * 1000) _set_job_creation_time(job_store, old_job.job_id, one_hour_ago) subprocess.check_call([ sys.executable, "-m", "mlflow", "gc", "--backend-store-uri", db_uri, "--jobs", "--older-than", "30m", ]) with pytest.raises(MlflowException, match=r"Job .+ not found"): job_store.get_job(old_job.job_id) retrieved_new_job = job_store.get_job(new_job.job_id) assert retrieved_new_job.job_id == new_job.job_id def test_mlflow_gc_job_ids_with_older_than_filter(sqlite_store_with_jobs): _, job_store, db_uri = sqlite_store_with_jobs old_job = _create_test_job(job_store, "old_job") new_job = _create_test_job(job_store, "new_job") one_hour_ago = get_current_time_millis() - (60 * 60 * 1000) _set_job_creation_time(job_store, old_job.job_id, one_hour_ago) subprocess.check_call([ sys.executable, "-m", "mlflow", "gc", "--backend-store-uri", db_uri, "--job-ids", new_job.job_id, "--older-than", "30m", ]) retrieved_new_job = job_store.get_job(new_job.job_id) assert retrieved_new_job.job_id == new_job.job_id retrieved_old_job = job_store.get_job(old_job.job_id) assert retrieved_old_job.job_id == old_job.job_id def test_mlflow_gc_only_deletes_finalized_jobs(sqlite_store_with_jobs): from mlflow.entities._job_status import JobStatus _, job_store, db_uri = sqlite_store_with_jobs # Create jobs with different statuses pending_job = _create_test_job(job_store, "pending_job", finalize=False) assert pending_job.status == JobStatus.PENDING running_job = _create_test_job(job_store, "running_job", finalize=False) job_store.start_job(running_job.job_id) running_job = job_store.get_job(running_job.job_id) assert running_job.status == JobStatus.RUNNING succeeded_job = _create_test_job(job_store, "succeeded_job", finalize=True) assert succeeded_job.status == JobStatus.SUCCEEDED failed_job = _create_test_job(job_store, "failed_job", finalize=False) job_store.start_job(failed_job.job_id) job_store.fail_job(failed_job.job_id, "error") failed_job = job_store.get_job(failed_job.job_id) assert failed_job.status == JobStatus.FAILED timeout_job = _create_test_job(job_store, "timeout_job", finalize=False) job_store.start_job(timeout_job.job_id) job_store.mark_job_timed_out(timeout_job.job_id) timeout_job = job_store.get_job(timeout_job.job_id) assert timeout_job.status == JobStatus.TIMEOUT canceled_job = _create_test_job(job_store, "canceled_job", finalize=False) job_store.cancel_job(canceled_job.job_id) canceled_job = job_store.get_job(canceled_job.job_id) assert canceled_job.status == JobStatus.CANCELED subprocess.check_call([ sys.executable, "-m", "mlflow", "gc", "--backend-store-uri", db_uri, "--jobs", ]) retrieved_pending = job_store.get_job(pending_job.job_id) assert retrieved_pending.job_id == pending_job.job_id assert retrieved_pending.status == JobStatus.PENDING retrieved_running = job_store.get_job(running_job.job_id) assert retrieved_running.job_id == running_job.job_id assert retrieved_running.status == JobStatus.RUNNING with pytest.raises(MlflowException, match=r"Job .+ not found"): job_store.get_job(succeeded_job.job_id) with pytest.raises(MlflowException, match=r"Job .+ not found"): job_store.get_job(failed_job.job_id) with pytest.raises(MlflowException, match=r"Job .+ not found"): job_store.get_job(timeout_job.job_id) with pytest.raises(MlflowException, match=r"Job .+ not found"): job_store.get_job(canceled_job.job_id)