Chapter 1: MLflow Introduction
What is MLflow?
MLflow is an open-source platform for managing the end-to-end machine learning lifecycle. It provides four main components:
Setting Up MLflow
Installation
# Basic installation
pip install mlflow
# With extras
pip install mlflow[extras]
# Start local tracking server
mlflow server --host 0.0.0.0 --port 5000
Production Setup
# docker-compose.yml
version: '3.8'
services:
mlflow:
image: ghcr.io/mlflow/mlflow:v2.10.0
ports:
- "5000:5000"
environment:
- MLFLOW_BACKEND_STORE_URI=postgresql://user:pass@postgres:5432/mlflow
- MLFLOW_DEFAULT_ARTIFACT_ROOT=s3://mlflow-artifacts/
- AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID}
- AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY}
command: >
mlflow server
--backend-store-uri postgresql://user:pass@postgres:5432/mlflow
--default-artifact-root s3://mlflow-artifacts/
--host 0.0.0.0
--port 5000
depends_on:
- postgres
postgres:
image: postgres:15
environment:
- POSTGRES_USER=user
- POSTGRES_PASSWORD=pass
- POSTGRES_DB=mlflow
volumes:
- postgres_data:/var/lib/postgresql/data
volumes:
postgres_data:
MLflow Tracking
Basic Experiment Tracking
import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, f1_score, roc_auc_score
# Set tracking URI
mlflow.set_tracking_uri("http://localhost:5000")
# Set experiment
mlflow.set_experiment("customer-churn-prediction")
# Start a run
with mlflow.start_run(run_name="random-forest-baseline"):
# Log parameters
n_estimators = 100
max_depth = 10
mlflow.log_param("n_estimators", n_estimators)
mlflow.log_param("max_depth", max_depth)
mlflow.log_param("model_type", "RandomForest")
# Train model
model = RandomForestClassifier(
n_estimators=n_estimators,
max_depth=max_depth,
random_state=42
)
model.fit(X_train, y_train)
# Make predictions
predictions = model.predict(X_test)
probabilities = model.predict_proba(X_test)[:, 1]
# Log metrics
mlflow.log_metric("accuracy", accuracy_score(y_test, predictions))
mlflow.log_metric("f1_score", f1_score(y_test, predictions))
mlflow.log_metric("auc_roc", roc_auc_score(y_test, probabilities))
# Log model
mlflow.sklearn.log_model(model, "model")
# Log artifacts (files)
mlflow.log_artifact("feature_importance.csv")
Autologging
# Enable autolog for sklearn
mlflow.sklearn.autolog()
# Now training automatically logs everything
model = RandomForestClassifier(n_estimators=100)
model.fit(X_train, y_train)
# Parameters, metrics, and model are logged automatically!
# Autolog for other frameworks
mlflow.pytorch.autolog()
mlflow.tensorflow.autolog()
mlflow.xgboost.autolog()
mlflow.lightgbm.autolog()
Nested Runs for Hyperparameter Tuning
from sklearn.model_selection import ParameterGrid
param_grid = {
'n_estimators': [100, 200, 300],
'max_depth': [5, 10, 15, None],
'min_samples_split': [2, 5, 10]
}
with mlflow.start_run(run_name="hyperparameter-search"):
best_auc = 0
best_params = None
for params in ParameterGrid(param_grid):
with mlflow.start_run(nested=True):
# Log parameters
mlflow.log_params(params)
# Train model
model = RandomForestClassifier(**params, random_state=42)
model.fit(X_train, y_train)
# Evaluate
auc = roc_auc_score(y_test, model.predict_proba(X_test)[:, 1])
mlflow.log_metric("auc_roc", auc)
# Track best
if auc > best_auc:
best_auc = auc
best_params = params
# Log best results to parent run
mlflow.log_params({f"best_{k}": v for k, v in best_params.items()})
mlflow.log_metric("best_auc_roc", best_auc)
MLflow Projects
Project Structure
ml-project/
├── MLproject # Project definition
├── conda.yaml # Environment specification
├── train.py # Training script
└── requirements.txt # Alternative to conda
MLproject File
# MLproject
name: customer-churn-model
conda_env: conda.yaml
entry_points:
main:
parameters:
n_estimators: {type: int, default: 100}
max_depth: {type: int, default: 10}
data_path: {type: string, default: "data/processed/"}
command: "python train.py --n-estimators {n_estimators} --max-depth {max_depth} --data-path {data_path}"
evaluate:
parameters:
model_path: string
test_data: string
command: "python evaluate.py --model-path {model_path} --test-data {test_data}"
Running Projects
# Run locally
mlflow run . -P n_estimators=200 -P max_depth=15
# Run from Git
mlflow run https://github.com/company/ml-project -P n_estimators=200
# Run with specific environment
mlflow run . --env-manager conda
MLflow Models
Model Signature and Input Example
from mlflow.models.signature import infer_signature
import pandas as pd
# Infer signature from data
signature = infer_signature(X_train, model.predict(X_train))
# Create input example
input_example = X_train.head(5)
# Log model with signature
mlflow.sklearn.log_model(
model,
"model",
signature=signature,
input_example=input_example,
registered_model_name="customer-churn-model"
)
Model Flavors
# Different ways to log models
# Sklearn
mlflow.sklearn.log_model(sklearn_model, "model")
# PyTorch
mlflow.pytorch.log_model(pytorch_model, "model")
# TensorFlow/Keras
mlflow.tensorflow.log_model(tf_model, "model")
# Custom models with pyfunc
class CustomModel(mlflow.pyfunc.PythonModel):
def load_context(self, context):
# Load artifacts
self.model = load_model(context.artifacts["model_path"])
def predict(self, context, model_input):
return self.model.predict(model_input)
mlflow.pyfunc.log_model(
"model",
python_model=CustomModel(),
artifacts={"model_path": "path/to/model"}
)
Key Takeaways
Summary
- MLflow Tracking logs parameters, metrics, and artifacts
- Autologging simplifies tracking for supported frameworks
- MLflow Projects ensure reproducible ML code
- MLflow Models provide standard packaging for deployment
- Model signatures define expected inputs and outputs