Skip to content
Deadline to Die

MLflow Deep Dive

🧪 MLflow: The Definitive Deep Dive

MLflow is an open-source platform to manage the ML lifecycle. It consists of four primary components: Tracking, Projects, Models, and Registry.

🟢 Level 1: Tracking Metadata

Every time you train a model, you should record the “DNA” of that run.

  • Params: Key-value inputs (e.g., learning_rate=0.01).
  • Metrics: Numeric outputs (e.g., accuracy=0.92).
  • Tags: Descriptive labels (e.g., env=staging).
import mlflow

# Point to a central tracking server (Postgres + S3)
mlflow.set_tracking_uri("http://mlflow.internal.company.com")
mlflow.set_experiment("fraud-detection")

with mlflow.start_run(run_name="xgboost-baseline"):
    # ... training code ...
    mlflow.log_params({"max_depth": 5, "n_estimators": 100})
    mlflow.log_metric("auc", 0.88)

🟡 Level 2: Artifacts & Model Flavors

1. Artifact Store

Artifacts are files like trained models, plots, or datasets. They are stored in an Artifact URI (usually S3 or Azure Blob).

2. Model Flavors

MLflow supports multiple “flavors” (e.g., sklearn, pytorch, onnx). This allows the same model to be served in different environments.

# Log a model with a specific flavor
mlflow.sklearn.log_model(model, "model_dir")

🔴 Level 3: The Model Registry Pattern

The Registry is a central place to manage a model’s life from “Developer laptop” to “Production API.”

  1. Register: Add a logged model to the registry.
  2. Version: Automatically increments versions (v1, v2, v3).
  3. Stage Transition: Manually or automatically move a model through stages:
    • Staging: Ready for integration tests.
    • Production: Ready for real traffic.
    • Archived: Deprecated.