Operating MLflow at scale

MLflow · Expert · 14 min

A shared MLflow becomes critical infrastructure; outages or lost metadata block every ML team.

Running MLflow for an organization raises operational questions beyond mlflow server.

Architecture. Run stateless server replicas behind a load balancer, with a managed PostgreSQL or MySQL backend and an object store for artifacts. Consider separate artifact-serving instances (--artifacts-only) when large model uploads would otherwise slow API traffic. Kubernetes deployments typically use the community Helm charts.

Access control. Built-in basic auth (--app-name basic-auth) supports users and per-experiment and per-model permissions (READ, EDIT, MANAGE). Larger setups put an SSO proxy in front, or use workspaces to isolate teams on one server. Managed platforms (Databricks, SageMaker, Azure ML, Nebius) provide hosted MLflow with integrated auth.

Upgrades. New versions may change the database schema. Back up the database, then run mlflow db upgrade <backend-uri> before starting new servers. Old clients usually keep working against newer servers, but test before rolling out.

Cleanup. Deleting runs or experiments only marks them deleted. mlflow gc --backend-store-uri ... --older-than 30d permanently removes them and their artifacts. Set retention policies for checkpoints and large artifacts.

Performance. The metrics table grows fastest; log step metrics at sensible intervals, use log_metrics batches, and enable asynchronous logging (MLFLOW_ENABLE_ASYNC_LOGGING=true) so training is not blocked by network latency. Index-heavy queries such as broad search_runs over huge experiments should be paginated.

Example

# Shell (operations):
#   mlflow db upgrade postgresql://mlflow:pw@db:5432/mlflow      # after backup
#   mlflow gc --backend-store-uri postgresql://mlflow:pw@db:5432/mlflow --older-than 30d
#   mlflow server --app-name basic-auth --workers 4 \
#       --backend-store-uri postgresql://mlflow:pw@db:5432/mlflow \
#       --artifacts-destination s3://acme-ml-artifacts
import mlflow

mlflow.config.enable_async_logging(True)       # don't block training on logging I/O
with mlflow.start_run():
    for step in range(0, 1000, 100):            # log at intervals, not every step
        mlflow.log_metrics({"loss": 1 / (step + 1), "lr": 1e-3}, step=step)
mlflow.flush_async_logging()
print("logged 10 batched points per metric")

Remember

  • Stateless servers + managed SQL database + object store is the standard layout
  • Back up, then mlflow db upgrade before running a new server version
  • Deleted runs need mlflow gc to be removed; log metrics in batches and asynchronously

Watch out for

  • Upgrading the MLflow server package without running the database migration, which makes the server fail to start or behave inconsistently.

Check yourself

  1. Users deleted thousands of old runs but storage did not shrink. Why?

Sign in to answer, see the explanations, and get this lesson back as a quick review later.

Official docs: mlflow.org/docs/latest/self-hosting/