How to use MLflow to track experiments and register models for classical ML

Summary

• MLflow tracks experiments by logging parameters, metrics, and artifacts, letting you compare model runs side-by-side.
• Use the Model Registry to version, stage, and promote models from Development to Production, with full lineage tracking.
• MLflow integrates with scikit-learn, XGBoost, and Spark ML for automatic logging of model hyperparameters and performance metrics.

How to use MLflow to track experiments and register models for classical ML

MLflow is the de facto standard for managing the classical ML lifecycle. It tracks experiments, versions models, and manages promotions from development to production, giving you reproducibility and auditability without extra infrastructure.

The MLflow workflow for classical ML

Step 1: Start an MLflow run and log experiments

import mlflow
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, roc_auc_score

    # Start an MLflow run
with mlflow.start_run():
    # Log hyperparameters
    mlflow.log_param("n_estimators", 100)
    mlflow.log_param("max_depth", 10)
    
    # Train model
    model = RandomForestClassifier(n_estimators=100, max_depth=10)
    model.fit(X_train, y_train)
    
    # Log metrics
    y_pred = model.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)
    auc = roc_auc_score(y_test, model.predict_proba(X_test)[:, 1])
    mlflow.log_metric("accuracy", accuracy)
    mlflow.log_metric("auc", auc)
    
    # Log the model
    mlflow.sklearn.log_model(model, "model")

MLflow auto-logs training time, framework version, and other metadata. Each run is stored with a unique ID and timestamp.

Step 2: Compare experiments in the MLflow UI

Navigate to the MLflow Experiments UI (built into Databricks notebooks). View all runs, sort by metrics, and compare hyperparameters across experiments.
Key columns:

• Parameters: Hyperparameters for each run
• Metrics: Accuracy, AUC, F1, or custom metrics
• Tags: Custom labels (e.g., "baseline", "production-candidate")
• Artifacts: Model files, plots, or any other outputs

Step 3: Register the best model

Once you identify a model that performs well, register it in the Model Registry.

from mlflow import MlflowClient

client = MlflowClient()

    # Register model from the run
model_uri = "runs:/" + run_id + "/model"
model_details = client.create_model_version(
    name="loan_classifier",
    source=model_uri,
    run_id=run_id
)

print(f"Registered model version {model_details.version}")

The Model Registry now tracks:

• Model name and version
• Stage (None → Staging → Production)
• Creation timestamp and author
• Source run and lineage

Step 4: Promote models through stages

Move your model from Staging (testing) to Production (serving).

client.transition_model_version_stage(
    name="loan_classifier",
    version=model_details.version,
    stage="Production"
)

Only one model version can be in the "Production" stage at a time. If you want to deploy a new version, the previous Production version moves to "Archived".

Best practices for classical ML experiment tracking

Practice	Why	How
Log hyperparameters early	Hyperparameters define the experiment; logging them first aids reproducibility	Use mlflow.log_param() before training
Use informative run names	Helps identify runs in the UI without reading parameters	Set run_name at start: mlflow.start_run(run_name="rf_max_depth_10")
Tag related runs	Group experiments by type (baseline, tuning, final) for easier filtering	Use mlflow.set_tag("type", "baseline")
Log data lineage	Enables reproducibility and debugging	Log dataset version/hash: mlflow.log_param("data_version", "v1.2.3")
Compare top-K models	Visualizes trade-offs between performance and complexity	Use the Experiments UI to sort by metrics and select top 3-5 runs
Version model artifacts	Enables rollback if a new model performs worse	MLflow auto-versions; old versions stay in Model Registry
Document model metadata	Helps others understand the model's purpose and limitations	Add description in Model Registry: client.update_model_version(..., description=...)

Example: tracking a hyperparameter tuning experiment

import mlflow
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score

experiment_name = "loan_rf_tuning"
mlflow.set_experiment(experiment_name)

    # Test different max_depth values
for max_depth in [4, 8, 12, 16]:
    with mlflow.start_run(run_name=f"rf_depth_{max_depth}"):
        # Log hyperparameters
        mlflow.log_param("max_depth", max_depth)
        mlflow.log_param("n_estimators", 100)
        mlflow.log_param("data_split", "80/20")
        
        # Train
        model = RandomForestClassifier(n_estimators=100, max_depth=max_depth)
        model.fit(X_train, y_train)
        
        # Log metrics
        accuracy = accuracy_score(y_test, model.predict(X_test))
        mlflow.log_metric("test_accuracy", accuracy)
        
        # Log model
        mlflow.sklearn.log_model(model, "model")

    # View all runs in the Experiments UI; pick the best

MLflow with Feature Store

When using a Feature Store, include feature versions in your MLflow logs.

with mlflow.start_run():
    mlflow.log_param("feature_set_version", "v1.3")
    mlflow.log_param("feature_store", "Databricks Feature Store")
    
    # Read features from Feature Store
    fs = feature_store.FeatureStoreClient()
    train_df = fs.score_batch("my_feature_table", df_with_keys)
    
    # Train model...

This enables full traceability: model → features → training data → business metrics.

FAQs

Do I need MLflow for classical ML or is it optional?

Optional but highly recommended. MLflow reduces boilerplate and provides reproducibility out of the box. For production systems, it's essential for version control and rollback.

Can I use MLflow to compare classical ML models with generative AI models?

Yes. Both can be logged in MLflow. The Model Registry stores both types. You can version and stage both classical models and LLM endpoints together.

How do I monitor model performance after deployment?

MLflow doesn't directly monitor production performance, but you can log actual outcomes back to MLflow or a monitoring system (Lakehouse Monitoring) to detect drift and trigger retraining.