MLflow is an open-source platform designed to manage the complete machine learning lifecycle, including experimentation, reproducibility, deployment, and a central model registry. Developed by Databricks, it provides a set of lightweight APIs and tools that integrate with any machine learning library, framework, or language. Key components include Tracking (for logging parameters, metrics, and artifacts), Projects (for packaging code in a reusable, reproducible format), Models (for deploying models in diverse serving environments), and a Model Registry (for collaborative model versioning and lifecycle management).
1.1 MLflow Tracking
1.1.1 Experiment and Run Management
MLflow Tracking organizes experiments, each representing a logical grouping of runs. A run corresponds to a single execution of a machine learning process, such as training a model or tuning hyperparameters. The API allows users to create experiments and start runs, associating each run with a unique ID, timestamps, and user-defined tags. The UI and client library provide mechanisms to query, filter, and compare runs across experiments.
1.1.2 Logging Parameters, Metrics, and Artifacts
During a run, users can log parameters (key-value pairs for input settings), metrics (numerical values such as accuracy or loss that can be updated over time), and artifacts (arbitrary files, including model binaries, plots, or data files). Parameters and metrics are stored as structured data in the backend store, while artifacts are stored in a configurable artifact store (local filesystem, cloud storage, or database). The logging functions (e.g., mlflow.log_param, mlflow.log_metric, mlflow.log_artifact) are idempotent and can be called repeatedly.
1.1.3 Comparison of Runs via UI or API
The MLflow UI displays runs in a tabular format with sortable columns and filterable fields. Users can select multiple runs to compare their parameters and metrics side by side. Programmatic comparison is supported via the mlflow.search_runs API, which returns a pandas DataFrame of run data. This enables automated analysis and visualization outside the UI.
1.2 MLflow Projects
1.2.1 Project Structure and Conda/Docker Environments
An MLflow Project is a directory containing an MLproject file that specifies the project's entry points, parameters, and environment. The environment can be defined using a Conda environment YAML file (conda.yaml) or a Dockerfile. The MLproject file also lists dependencies and default parameter values. This structure ensures reproducibility by explicitly capturing the software environment alongside the code.
1.2.2 Running Projects Remotely
Projects can be executed on remote platforms such as Databricks, Kubernetes, or any system that supports Docker containers. The mlflow run command accepts a URI pointing to the project directory and an optional backend configuration. When running remotely, MLflow packages the project code and environment, transfers them to the remote executor, and retrieves the results (parameters, metrics, artifacts) back to the tracking server.
1.2.3 Parameter Entry Points
Each MLflow Project defines one or more entry points – scripts that accept parameters. Parameters are declared in the MLproject file with types (string, float, int, path) and default values. At runtime, parameters can be overridden via command-line arguments or programmatically. This design allows users to parameterize experiments without modifying the underlying code.
1.3 MLflow Models
1.3.1 Model Flavors (Python Function, PyTorch, TensorFlow, scikit-learn, etc.)
MLflow Models abstract the underlying machine learning framework into “flavors.” Common flavors include python_function (a generic inference interface), pytorch, tensorflow, sklearn, keras, xgboost, and many more. Each flavor defines a standard format for saving and loading models, as well as a default prediction signature. The python_function flavor acts as a universal backend for deployment, wrapping all other flavors.
1.3.2 Model Serialization and Loading
Models are serialized using the flavor’s native format (e.g., PyTorch’s state_dict, TensorFlow’s SavedModel) and stored as an artifact within a run. The mlflow.<flavor>.log_model function creates an MLflow Model directory containing the serialized model, a conda environment, and metadata. Loading is performed by mlflow.<flavor>.load_model or mlflow.pyfunc.load_model, which restores the model to a Python object ready for inference.
1.3.3 Deployment to Serving Platforms
1.3.3.1 Local REST Server
The mlflow models serve command starts a local REST server that exposes a prediction endpoint (typically at http://localhost:5001/invocations). The server loads a specified model and accepts JSON or CSV input according to the model’s signature. This is useful for testing and small-scale deployment.
1.3.3.2 Cloud Platforms (Amazon SageMaker, Azure ML, Databricks)
MLflow integrates with major cloud ML platforms. For Amazon SageMaker, the mlflow.sagemaker.deploy function packages the model into a Docker container and deploys it to a SageMaker endpoint. Similar integrations exist for Azure Machine Learning and Databricks serving. These platforms handle scaling, monitoring, and load balancing automatically.
1.4 MLflow Model Registry
1.4.1 Model Versioning
The Model Registry stores models as registered entities with version numbers. Each version corresponds to a specific model artifact logged in a run. New versions can be created by logging a model with the same registered model name. The registry maintains a list of all versions, their URIs, and creation timestamps.
1.4.2 Stage Transitions (Staging, Production, Archived)
Each model version can be assigned a stage: None, Staging, Production, or Archived. Stage transitions are performed via the UI or API (mlflow.<model_version>.transition_stage). For example, a candidate model in Staging can be promoted to Production after validation. Archived versions are kept for historical reference but are no longer recommended for deployment.
1.4.3 Model Annotations and Descriptions
Users can attach descriptions, tags, and comments to registered models and individual versions. This facilitates communication within data science teams – for instance, noting a model’s performance on a hold-out set or linking to an A/B test report. The annotations are stored in the backend database and visible in the registry UI.
2.1 Installation via pip or Conda
MLflow is available as a Python package on PyPI and Conda. Installation is straightforward: pip install mlflow or conda install -c conda-forge mlflow. Additional dependencies for specific frameworks (e.g., mlflow[extras]) can be specified as needed. The installation includes the tracking server, CLI, and APIs.
2.2 Backend Store and Artifact Store
2.2.1 Local File System
By default, MLflow stores tracking data (experiments, runs, parameters, metrics) as files in a local directory (./mlruns). Artifacts are also stored under mlruns. This setup is sufficient for single‑user experimentation.
2.2.2 Database Backends (SQLite, MySQL, PostgreSQL)
For persistent, multi‑user setups, relational database backends are supported. The backend store URI can be set to sqlite:///path/to/db.sqlite, mysql://user:pass@host/db, or postgresql://user:pass@host/db. Migrations are handled automatically when the tracking server starts.
2.2.3 Cloud Storage (S3, GCS, Azure Blob)
Artifacts can be stored in cloud blob storage. Configuration is done via environment variables or a URI (e.g., s3://bucket/prefix, gs://bucket/path, wasbs://container@storage.blob.core.windows.net). The client uses the respective SDK or authentication credentials (e.g., AWS credentials, Google service account, Azure storage keys).
2.3 Tracking Server Deployment
2.3.1 Single-User Mode
In single-user mode, the tracking server is run locally with mlflow server and default settings. The UI is accessible at http://127.0.0.1:5000. No authentication is enforced, and all users have full access.
2.3.2 Multi-User and Authentication Options
For collaborative environments, MLflow supports reverse‑proxy authentication (e.g., with Nginx and basic auth) or integration with identity providers (LDAP, OAuth) via plugins. The server can be configured to require tokens (see Section 5.4.1). Production deployments often use a combination of database backends and cloud artifact stores.
2.4 Environment Variables and Configuration
Key environment variables include MLFLOW_TRACKING_URI, MLFLOW_S3_ENDPOINT_URL, MLFLOW_ARTIFACT_LOCATION, and MLFLOW_DEFAULT_ARTIFACT_ROOT. The tracking server itself accepts command‑line flags (e.g., --host, --port, --backend-store-uri, --default-artifact-root). Configuration can also be placed in a mlflow.yml file, though this is less common.
3.1 Single-User Experimentation
3.1.1 Creating an Experiment
A user creates an experiment using mlflow.create_experiment("experiment_name") or via the UI. The experiment is assigned a unique ID. Runs logged to this experiment are grouped together in the UI.
3.1.2 Logging Runs Programmatically
Within a Python script, the user starts a run with mlflow.start_run(), logs parameters and metrics, and ends the run. Example: mlflow.log_param("learning_rate", 0.001); mlflow.log_metric("accuracy", 0.95). The run is automatically associated with the current (or specified) experiment.
3.1.3 Viewing Results in the UI
The user navigates to the tracking server’s UI, selects the experiment, and sees all runs. They can sort by metrics, filter by parameters, and compare selected runs. The UI also provides a dashboard for visualizing metric histories.
3.2 Collaborative Development
3.2.1 Sharing Tracking Server
A team deploys a shared tracking server (see Section 2.3.2) and configures all team members’ clients to point to its URI. Each developer logs runs from their own environment. The team can see each other’s experiments in real time, avoiding duplication of work.
3.2.2 Using Model Registry for Team Reviews
Team members register candidate models to the Model Registry. A reviewer can examine a model version’s parameters, metrics, and artifacts in the UI. After validation, the reviewer transitions the version to Staging or Production. The registry supports adding comments for discussion.
3.3 Production Deployment
3.3.1 Exporting a Model as a Docker Container
The mlflow models build-docker command creates a Docker image from a saved model. The image contains the model artifact, conda environment, and a Python‑based prediction server. This image can be deployed to any container orchestration platform (Kubernetes, ECS, etc.).
3.3.2 Serving with MLflow’s Built-in REST API
For simpler deployments, use mlflow models serve -m <model_uri>. This starts a synchronous REST API with endpoints for health checks and prediction. The server is suitable for low‑latency, single‑instances.
3.3.3 Integrating with CI/CD Pipelines
MLflow actions can be embedded in CI/CD pipelines (e.g., GitHub Actions, Jenkins). A typical workflow: train a model in the pipeline, log it to the tracking server, register a new version, run tests against the version, and if tests pass, promote to production. The mlflow run command can trigger remote execution on a cluster.
4.1 Machine Learning Frameworks
4.1.1 TensorFlow and Keras
MLflow provides mlflow.tensorflow and mlflow.keras modules. They offer autologging capabilities that automatically capture model parameters, metrics, and graphs during training. The log_model method saves the model in TensorFlow’s SavedModel format.
4.1.2 PyTorch and Lightning
For PyTorch, mlflow.pytorch supports autologging via callbacks or explicit calls. PyTorch Lightning users can use MlflowLogger for logging metrics and mlflow.pytorch.log_model for checkpoint serialization.
4.1.3 scikit-learn and XGBoost
mlflow.sklearn and mlflow.xgboost allow logging and loading models with minimal overhead. Autologging for scikit‑learn captures estimator parameters, cross‑validation scores, and feature importances. XGBoost models are saved in its native .json or .ubj format.
4.2 Orchestration Platforms
4.2.1 Apache Airflow and Kubeflow
MLflow can be used within Airflow DAGs by calling the Python API to log runs. Kubeflow Pipelines can trigger MLflow runs as components, and the resulting metrics are sent to a shared tracking server. Integration is facilitated by the mlflow.run function.
4.2.2 Databricks
As the original developer, Databricks offers deep integration: MLflow is pre‑installed on Databricks clusters, and the tracking server is managed automatically. Users can log runs from notebooks, manage experiments via a native UI, and deploy models to Databricks serving endpoints.
4.3 Hyperparameter Tuning Libraries
4.3.1 Hyperopt
MLflow integrates with Hyperopt through the mlflow.hyperopt module. Each trial of Hyperopt is logged as a separate run with its parameters and objective value. The best model can be retrieved from the search results.
4.3.2 Optuna
Optuna’s integration uses a callback (mlflow.log_metric per trial) or a dedicated MLflowCallback. Optuna studies can be tracked in MLflow experiments, allowing visualization of the optimization history in the UI.
5.1 Custom Model Flavors
5.1.1 Implementing the PythonModel Interface
Users can create a custom flavor by implementing mlflow.pyfunc.PythonModel. This class must define predict(context, model_input). The context contains environment configuration (e.g., Spark session). The model is saved with mlflow.pyfunc.save_model and loaded like any other flavor.
5.1.2 Adding Custom Serialization
For frameworks not natively supported, developers can extend mlflow.pyfunc by writing custom save/load functions. The model is stored as a PythonModel and the custom serialization logic (e.g., pickle, ONNX) is executed inside the predict method.
5.2 Metadata and Lineage Tracking
5.2.1 Capturing Dataset Versions
Dataset versions can be logged as artifacts (e.g., CSV files) or referenced via tags. Some integrations (e.g., with DVC or Feast) allow linking a run to the exact dataset snapshot used. The mlflow.data module (in newer versions) provides helpers for dataset lineage.
5.2.2 Linking Artifacts to Runs
Each artifact is stored under the run’s artifact URI. The artifact’s path is recorded in the tracking database. Users can programmatically list artifacts with run.list_artifacts() and download them. This enables reproducibility by retrieving the exact data or model used in a run.
5.3 Performance and Scaling
5.3.1 Large-Scale Artifact Stores
When dealing with many or large artifacts (e.g., hundreds of GBs), cloud object stores are preferred over local filesystems. MLflow can leverage multipart uploads and streaming. For even larger artifacts, consider using a dedicated storage solution (e.g., HDFS) and setting the artifact location accordingly.
5.3.2 Asynchronous Logging
The default logging is synchronous, which can become a bottleneck in high‑throughput scenarios. MLflow supports asynchronous logging by using the mlflow.tracking.MlflowClient’s create_run, log_batch, etc., with a queue‑based approach. Alternatively, users can separate logging into a background thread.
5.4 Security and Access Control
5.4.1 Token-Based Authentication
The MLflow tracking server can be started with the --app-name=mlflow_auth option (using the MLflow Authentication plugin). Users authenticate with a token passed via the MLFLOW_TRACKING_TOKEN environment variable. The plugin also provides role‑based access control for experiments and registered models.
5.4.2 Encryption at Rest and In Transit
Encryption at rest is handled by the underlying storage systems (e.g., S3‑SSE, GCS‑CMEK). For transmission, the tracking server should be deployed behind HTTPS (using a reverse proxy like Nginx). Database connections can be encrypted (e.g., MySQL with SSL). Artifact downloads from cloud storage can use signed URLs or TLS‑enabled endpoints.
6.1 Official Documentation and Repository
The official documentation is hosted at mlflow.org and includes tutorials, API references, and deployment guides. The source code is maintained on GitHub under the mlflow/mlflow repository, where users can report issues, submit pull requests, and access release notes.
6.2 Plugins and Extensions
The MLflow community has developed a variety of plugins, such as mlflow-torchserve (deployment to TorchServe), mlflow-gluon (Apache MXNet), and integration with tracking backends like MLflow‑OpenMLDB. The plugin registry is maintained in the official documentation.
6.3 Migration from Similar Tools (e.g., Kubeflow, Neptune)
For teams migrating from Kubeflow, MLflow offers a lighter‑weight alternative for experiment tracking and model registry. Tools like mlflow-kubeflow bridge can migrate existing metadata. From Neptune, users can export runs to MLflow format using community scripts. The official docs contain a migration guide covering common patterns.