MLflow Cheat Sheet
A reference for MLflow's experiment tracking, model registry, and CLI commands used to log, compare, and deploy machine learning models.
Experiment Tracking
Log params, metrics, and artifacts for a run.
import mlflowmlflow.set_tracking_uri('http://localhost:5000')mlflow.set_experiment('my-experiment')with mlflow.start_run(): mlflow.log_param('lr', 0.01) mlflow.log_metric('accuracy', 0.92) mlflow.log_artifact('model.pkl') mlflow.sklearn.log_model(model, 'model')
CLI Commands
Common MLflow command-line operations.
mlflow ui # Launch tracking UI (default :5000)mlflow run . -P alpha=0.5 # Run an MLproject entry pointmlflow models serve -m runs:/<run_id>/model -p 1234mlflow experiments listmlflow server --backend-store-uri sqlite:///mlflow.db --default-artifact-root ./mlruns
Model Registry
Register and promote model versions.
from mlflow import MlflowClient# Register a run's model as a new model versionresult = mlflow.register_model('runs:/<run_id>/model', 'my-model')client = MlflowClient()client.transition_model_version_stage( name='my-model', version=3, stage='Production')
Core Components
The four pillars of the MLflow platform.
- Tracking- Logs parameters, metrics, artifacts, and source code for every training run
- Projects- Packages code with an MLproject file for reproducible, shareable runs
- Models- Standard packaging format supporting many flavors (sklearn, pytorch, xgboost, etc.)
- Model Registry- Central store for versioning models and managing stage transitions (Staging/Production/Archived)
- Autologging- mlflow.autolog() auto-captures params and metrics for supported frameworks with no manual logging
Nested Runs for Hyperparameter Search
Group child trials under a parent run to keep sweep results organized in the UI.
import mlflowwith mlflow.start_run(run_name='hpo-parent') as parent: best_score = -1 for lr in [0.001, 0.01, 0.1]: with mlflow.start_run(run_name=f'trial-lr-{lr}', nested=True): mlflow.log_param('lr', lr) score = train_and_eval(lr) mlflow.log_metric('val_accuracy', score) if score > best_score: best_score = score mlflow.log_metric('best_val_accuracy', best_score)
Custom PyFunc Models
Wrap arbitrary logic (pre/post-processing, ensembles) in the generic pyfunc flavor for uniform serving.
import mlflow.pyfuncclass EnsembleModel(mlflow.pyfunc.PythonModel): def load_context(self, context): import joblib self.model_a = joblib.load(context.artifacts['model_a']) self.model_b = joblib.load(context.artifacts['model_b']) def predict(self, context, model_input, params=None): pred_a = self.model_a.predict(model_input) pred_b = self.model_b.predict(model_input) return (pred_a + pred_b) / 2with mlflow.start_run(): mlflow.pyfunc.log_model( artifact_path='ensemble', python_model=EnsembleModel(), artifacts={'model_a': 'model_a.pkl', 'model_b': 'model_b.pkl'}, )
Model Signatures & Input Examples
Attach an inferred schema so served models validate inputs and document expected shape.
from mlflow.models import infer_signaturepredictions = model.predict(X_train)signature = infer_signature(X_train, predictions)mlflow.sklearn.log_model( model, 'model', signature=signature, input_example=X_train.iloc[:5],)# Log the training dataset for lineage trackingdataset = mlflow.data.from_pandas(train_df, source='s3://bucket/train.csv', name='train')mlflow.log_input(dataset, context='training')
Querying & Comparing Runs Programmatically
Use the fluent search API to pull the best run across an experiment without opening the UI.
runs = mlflow.search_runs( experiment_names=['my-experiment'], filter_string="metrics.accuracy > 0.9 and params.lr = '0.01'", order_by=['metrics.accuracy DESC'], max_results=5,)best_run_id = runs.iloc[0]['run_id']# Load a logged model directly from a run for offline inferenceloaded = mlflow.pyfunc.load_model(f'runs:/{best_run_id}/model')preds = loaded.predict(X_test)
Registry & Client Lifecycle APIs
MlflowClient calls used to automate promotion, aliasing, and cleanup beyond the UI.
- set_registered_model_alias- Points a mutable alias like 'champion' at a specific model version, decoupling deploy targets from version numbers
- search_model_versions- Queries versions across models with a filter string, e.g. "name='my-model'"
- set_model_version_tag- Attaches metadata (e.g. 'validated_by') to a specific model version for audit trails
- delete_run / restore_run- Soft-deletes a run (recoverable within retention window) to prune noisy experiment history
- get_latest_versions- Returns the newest model version(s) filtered by stage, useful for rollback scripts
- log_batch- Logs many params/metrics/tags in a single API call to cut network overhead during large sweeps
Call mlflow.autolog() at the top of your training script before fitting a model — it automatically captures framework-specific parameters, metrics, and model artifacts for libraries like scikit-learn, XGBoost, and PyTorch Lightning without any manual log_param or log_metric calls.