MLflow Guide
pip install mlflow # or: uv add mlflow # Start UI: mlflow ui # Open: http://localhost:5000
MLflow tracks ML experiments — log parameters, metrics, artifacts (models, plots, data), and source code for every run. The UI lets you compare runs side-by-side and find the best hyperparameters. Use `mlflow.start_run()` as a context manager.
Beyond tracking, MLflow includes a Model Registry (version and stage models), a Deployment module (serve models as REST APIs), and an Evaluation module (benchmark LLMs and classical models). Models can be any framework: PyTorch, TensorFlow, sklearn, XGBoost, or custom.
MLflow supports autologging — `mlflow.autolog()` automatically logs params, metrics, and models for sklearn, PyTorch Lightning, XGBoost, and more. The MLflow UI is extensible with custom plugins. For production, use MLflow with a PostgreSQL backend and S3-compatible storage.
Tracking
mlflow ui # Default port 5000 mlflow ui --port 8080 mlflow ui --backend-store-uri sqlite:///mlflow.db # Open http://localhost:5000
import mlflow
with mlflow.start_run():
# Log parameters
mlflow.log_param('learning_rate', 0.01)
mlflow.log_param('batch_size', 32)
# Log metrics
for epoch in range(10):
accuracy = 0.8 + epoch * 0.02
mlflow.log_metric('accuracy', accuracy, step=epoch)
# Log artifacts
mlflow.log_artifact('confusion_matrix.png')
mlflow.log_artifact('model.pkl')
# Log model
mlflow.sklearn.log_model(model, 'model')# MLproject file:
# name: my_project
# conda_env: conda.yaml
# entry_points:
# main:
# parameters:
# data: {type: str, default: ./data}
# command: "python train.py {data}"
# Run project:
mlflow run . -P data=./dataset
# Remote execution:
mlflow run https://github.com/user/repo -P param=valueAutologging
import mlflow mlflow.autolog() # Enable autolog for all supported frameworks # Or specific framework: mlflow.sklearn.autolog() mlflow.pytorch.autolog() mlflow.xgboost.autolog() # Now any model training is auto-logged: from sklearn.ensemble import RandomForestClassifier model = RandomForestClassifier().fit(X_train, y_train) # All params, metrics, and model are logged automatically
Model Registry
# Register a model:
mlflow.register_model(
'runs:/<RUN_ID>/model',
'MyModel'
)
# CLI:
mlflow models -h # List registered models
# Load from registry:
import mlflow.pyfunc
model = mlflow.pyfunc.load_model(
model_uri='models:/MyModel/Production'
)
# Transition stage via UI or API:
from mlflow.tracking import MlflowClient
client = MlflowClient()
client.transition_model_version_stage(
name='MyModel',
version=3,
stage='Production'
)Serving
# Serve a registered model:
mlflow models serve -m models:/MyModel/Production -p 5001
# Test:
curl -X POST http://localhost:5001/invocations
-H 'Content-Type: application/json'
-d '{"inputs": {"feature1": 1.0, "feature2": 2.0}}'
# Serve from run:
mlflow models serve -m runs:/<RUN_ID>/model --port 5001Evaluation
import mlflow
from mlflow.metrics import rouge, flesch_kincaid
with mlflow.start_run():
eval_data = {
'inputs': ['What is MLflow?', 'Explain RAG'],
'ground_truth': [
'MLflow is an ML platform',
'RAG is retrieval-augmented generation'
]
}
results = mlflow.evaluate(
data=eval_data,
model='openai:/gpt-4o-mini',
model_type='text-summarization',
extra_metrics=[rouge(), flesch_kincaid()]
)
print(results.metrics)