Managing Models Using MLflow on Databricks

Table of Contents

1. Course Overview
2. Model Tracking with MLflow
   • 2.1 The Databricks Machine Learning Environment
   • 2.2 Introduction to MLflow Tracking
   • 2.3 Experiments and Runs
   • 2.4 Setting Up the Environment (Demo)
   • 2.5 Data Cleaning and Preprocessing (Demo)
   • 2.6 Creating an Experiment (Demo)
   • 2.7 Creating and Ending Runs (Demo)
   • 2.8 Tracking Metrics in a Run (Demo)
   • 2.9 Tracking Parameters in a Run (Demo)
   • 2.10 Visualizing and Sorting Runs (Demo)
   • 2.11 Comparing Runs (Demo)
   • 2.12 Autologging with MLflow (Demo)
   • 2.13 Programmatic Search and Sort of Runs (Demo)
3. Productionizing and Serving Models
   • 3.1 Model Registry and Model Serving
   • 3.2 Training Multiple Models (Demo)
   • 3.3 Registering a Model via the UI (Demo)
   • 3.4 Loading and Using a Model for Predictions (Demo)
   • 3.5 Lifecycle Management and Stages (Demo)
   • 3.6 Real-Time Serving via a REST Endpoint (Demo)
   • 3.7 Model Signature and Programmatic Management (Demo)
4. Custom Models and MLflow Projects
   • 4.1 Custom Models and Custom Flavors
   • 4.2 Creating a Custom Model (Demo)
   • 4.3 MLflow Projects
   • 4.4 Running an MLflow Project (Demo)
5. Summary and Additional Resources

1. Course Overview

This course covers the complete lifecycle management of machine learning models using MLflow on Databricks. The main topics covered are:

┌─────────────────────────────────────────────────────────────┐
│           MODEL MANAGEMENT WITH MLFLOW ON DATABRICKS        │
├─────────────────┬───────────────────────┬───────────────────┤
│  MODULE 1       │  MODULE 2             │  MODULE 3         │
│  Model          │  Productionizing      │  Custom Models    │
│  Tracking       │  and Serving          │  and Projects     │
├─────────────────┼───────────────────────┼───────────────────┤
│ • Experiments   │ • Model Registry      │ • Custom Models   │
│ • Runs          │ • Lifecycle stages    │ • Custom Flavors  │
│ • log_metric    │ • Batch inference     │ • MLflow Projects │
│ • log_param     │ • REST endpoint       │ • GitHub / DBFS   │
│ • autolog       │ • Model serving       │                   │
└─────────────────┴───────────────────────┴───────────────────┘

Prerequisites:

• Python programming
• Machine learning concepts (scikit-learn)
• Basic familiarity with the Databricks platform

2. Model Tracking with MLflow

2.1 The Databricks Machine Learning Environment

Databricks is a cloud-native platform for large-scale data processing, machine learning, and analytics. It is built on the Data Lakehouse architecture, which combines the best of data lakes and data warehouses.

The Three Personas (Environments) in a Databricks Workspace

graph TD
    WS[Databricks Workspace]
    WS --> SQL[Databricks SQL\nData Analysts\nSQL Queries / Dashboards]
    WS --> DSE[Data Science & Engineering\nEngineers / Data Scientists\nETL / Apache Spark]
    WS --> ML[Machine Learning\nData Scientists / ML Engineers\nModel Training and Deployment]

    style ML fill:#4a90d9,color:#fff
    style SQL fill:#7fb3d3
    style DSE fill:#7fb3d3

What the Databricks ML Runtime Provides

Full ML Model Lifecycle on Databricks

flowchart LR
    A[Data\nPreparation] --> B[Feature\nEngineering]
    B --> C[Model\nTraining]
    C --> D[MLflow\nTracking]
    D --> E[Model\nRegistry]
    E --> F[Deployment\n& Inference]

    C -->|AutoML| C
    C -->|HyperOpt| C

    style D fill:#ff7700,color:#fff
    style E fill:#ff7700,color:#fff

2.2 Introduction to MLflow Tracking

MLflow is an open-source platform created by the founders of Databricks to manage the machine learning lifecycle end-to-end.

Key Characteristics

• Library-agnostic: works with any ML library
• Multi-language: Python, R, Java, REST API, CLI
• Open source + managed version on Databricks

The Key Components of MLflow

graph LR
    MLflow --> T[Tracking\nExperiment tracking\nParameters and results]
    MLflow --> M[Models\nModel serialization\nto disk]
    MLflow --> P[Projects\nML code packaging\nReusable format]
    MLflow --> R[Model Registry\nCentralized management\nand versioning]
    MLflow --> S[Model Serving\nDeployment and\nREST inference]

    style T fill:#4a90d9,color:#fff
    style M fill:#27ae60,color:#fff
    style P fill:#8e44ad,color:#fff
    style R fill:#e67e22,color:#fff
    style S fill:#c0392b,color:#fff

What Model Tracking Can Capture

2.3 Experiments and Runs

Conceptual Hierarchy

MLflow Workspace
└── Experiment (loan_approval_prediction)
    ├── Run 1: exploratory_data_analysis
    │   ├── Artifacts: figure1.png, figure2.png, figure3.png
    │   └── Tags: ...
    ├── Run 2: RF_Default_params
    │   ├── Metrics: Test_accuracy, AUC_score, ...
    │   └── Tags: Classifier=RF-default_parameters
    ├── Run 3: RF_tuned_params_scenario1
    │   ├── Parameters: n_estimators=200, criterion=gini, ...
    │   ├── Metrics: Test_accuracy, AUC_score, ...
    │   └── Tags: ...
    └── Run N: ...

Types of Experiments

Run Lifecycle

stateDiagram-v2
    [*] --> RUNNING : mlflow.start_run()
    RUNNING --> FINISHED : mlflow.end_run() or end of with block
    RUNNING --> FAILED : Unhandled exception
    FINISHED --> [*]
    FAILED --> [*]
    FINISHED --> DELETED : mlflow.delete_run()
    DELETED --> FINISHED : Restore (30 days)

2.4 Setting Up the Environment (Demo)

Setup steps in the Databricks workspace:

1. Switch to the Machine Learning persona (and “pin” it to keep it as the default)
2. Create a single-node cluster with the Databricks ML Runtime (e.g. version 11.2, Spark 3.3.0, Scala 2.12)
3. Enable the DBFS File Browser: Admin Console → Workspace Settings → Advanced → DBFS File Browser
4. Upload data to DBFS → FileStore → datasets → customers.csv
5. Import the demo notebook to Workspace → [your user] → Import

Note: The ML Runtime comes pre-installed with scikit-learn, XGBoost, TensorFlow, PyTorch, and MLflow.

2.5 Data Cleaning and Preprocessing (Demo)

The customers.csv dataset contains 614 records of loan applications. The goal is to train a classification model to predict whether a customer gets loan approval (loan_approval_status).

Loading and Initial Exploration

import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

loan_data = pd.read_csv('/dbfs/FileStore/datasets/customers.csv', header='infer')
loan_data.head()

Issues identified in the data:

• num_of_dependents is of type object (due to the value "3+")
• spouse_income is of type object (due to malformed values like "9.857.999.878")
• Missing values (nulls) in several columns

Data Cleaning

# Check unique values
loan_data['num_of_dependents'].unique()
# → ['0', '1', '2', '3+']

# Drop rows with missing values
loan_data.dropna(axis=0, how='any', inplace=True)

# Remove rows with typos in spouse_income
loan_data = loan_data[loan_data.spouse_income != '9.857.999.878']
loan_data = loan_data[loan_data.spouse_income != '1.612.000.084']

# Convert to numeric type
loan_data['spouse_income'] = loan_data['spouse_income'].astype(float)

Exploratory Data Analysis (EDA)

# Boxplot: Income vs loan approval status
fig1, ax1 = plt.subplots()
plt.ylim(0, 20000)
sns.boxplot(x='loan_approval_status', y='income', data=loan_data)

# Barplot: Loan amount by gender and status
fig2, ax2 = plt.subplots()
sns.barplot(x='loan_approval_status', y='loan_amount', hue='Gender', data=loan_data)

# Barplot: Spouse income by credit history
fig3, ax3 = plt.subplots()
sns.barplot(x='loan_approval_status', y='spouse_income', hue='credit_history', data=loan_data)

Encoding Categorical Variables

from sklearn import preprocessing

# One-Hot Encoding for nominal variables with more than 2 categories
loan_data = pd.get_dummies(loan_data, columns=['Gender', 'property_type'], drop_first=True)

# Label Encoding for binary and ordinal variables
label_encoder = preprocessing.LabelEncoder()
for col in ['Married', 'education_level', 'working_status', 'loan_approval_status', 'num_of_dependents']:
    loan_data[col] = label_encoder.fit_transform(loan_data[col])

Saving and Splitting the Data

# Save preprocessed data
loan_data.drop(labels=['ID'], axis=1).to_csv(
    '/dbfs/FileStore/datasets/loan_data_processed.csv', index=False
)

# Train/test split (70%/30%)
X = loan_data.drop(labels=['ID', 'loan_approval_status'], axis=1)
y = loan_data['loan_approval_status']

from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=123)

2.6 Creating an Experiment (Demo)

Via the graphical interface:

1. Click on Experiments in the left navigation panel
2. Create a new Blank Experiment
3. Name the experiment: loan_approval_prediction
4. (Optional) Specify a location for artifacts

Via code:

import mlflow

mlflow.set_experiment(
    experiment_name='/Users/cloud.user@example.com/loan_approval_prediction'
)

Note: If set_experiment is not called, MLflow automatically creates a Notebook Experiment associated with the current notebook.

2.7 Creating and Ending Runs (Demo)

Method 1: start_run / end_run (explicit)

mlflow.start_run()

mlflow.log_figure(fig1, 'figure1.png')
mlflow.log_figure(fig2, 'figure2.png')
mlflow.log_figure(fig3, 'figure3.png')

mlflow.end_run()
# ⚠️ Forgetting end_run causes the error:
# "Run with UUID xxx is already active"

Method 2: start_run with run_name

mlflow.start_run(run_name='exploratory_data_analysis')

mlflow.log_figure(fig1, 'boxplot_approval_vs_income.png')
mlflow.log_figure(fig2, 'barplot_approval_amount_gender.png')
mlflow.log_figure(fig3, 'barplot_spouseincome_credithistory.png')

run = mlflow.active_run()
print('Active run_id: {}'.format(run.info.run_id))

mlflow.end_run()

Method 3: with Context Manager (recommended)

# The run is automatically ended at the end of the with block
with mlflow.start_run(run_name='eda_plots') as run1:
    mlflow.log_figure(fig1, 'boxplot_approval_vs_income.png')
    mlflow.log_figure(fig2, 'barplot_approval_amount_gender.png')
    mlflow.log_figure(fig3, 'barplot_spouseincome_credithistory.png')

2.8 Tracking Metrics in a Run (Demo)

from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score

with mlflow.start_run(run_name='RF_Default_params') as run2:

    model = RandomForestClassifier()
    model.fit(X_train, y_train)
    predictions = model.predict(X_test)
    predictions_proba = model.predict_proba(X_test)

    # Compute metrics
    test_accuracy        = accuracy_score(y_test, predictions)
    test_precision_score = precision_score(y_test, predictions)
    test_recall_score    = recall_score(y_test, predictions)
    test_f1_score        = f1_score(y_test, predictions)
    auc_score            = roc_auc_score(y_test, predictions_proba[:, 1])

    # Explicit metric logging
    mlflow.log_metric('Test_accuracy',        test_accuracy)
    mlflow.log_metric('Test_precision_score', test_precision_score)
    mlflow.log_metric('Test_recall_score',    test_recall_score)
    mlflow.log_metric('Test_f1_score',        test_f1_score)
    mlflow.log_metric('AUC_score',            auc_score)

    # Add a tag
    mlflow.set_tag('Classifier', 'RF-default_parameters')

Programmatic Access to Metrics via MlflowClient

client = mlflow.tracking.MlflowClient()

print('Model Parameters:', client.get_run(run2.info.run_id).data.params)
print('Metrics:',          client.get_run(run2.info.run_id).data.metrics)

2.9 Tracking Parameters in a Run (Demo)

Individual Logging with log_param

with mlflow.start_run(run_name='RF_tuned_params_scenario1') as run3:

    n_estimators      = 200
    criterion         = 'gini'
    min_samples_split = 5
    min_samples_leaf  = 2

    model = RandomForestClassifier(
        n_estimators=n_estimators,
        criterion=criterion,
        min_samples_split=min_samples_split,
        min_samples_leaf=min_samples_leaf
    )
    model.fit(X_train, y_train)
    predictions = model.predict(X_test)
    predictions_proba = model.predict_proba(X_test)

    # Log parameters (one by one)
    mlflow.log_param('No. of trees',       n_estimators)
    mlflow.log_param('Splitting criteria', criterion)
    mlflow.log_param('Min samples split',  min_samples_split)
    mlflow.log_param('Min samples leaf',   min_samples_leaf)

    # Log metrics (one by one)
    mlflow.log_metric('Test_accuracy',        accuracy_score(y_test, predictions))
    mlflow.log_metric('Test_precision_score', precision_score(y_test, predictions))
    mlflow.log_metric('Test_recall_score',    recall_score(y_test, predictions))
    mlflow.log_metric('Test_f1_score',        f1_score(y_test, predictions))
    mlflow.log_metric('AUC_score',            roc_auc_score(y_test, predictions_proba[:, 1]))

    mlflow.set_tag('Classifier', 'RF-tuned_params_sc1')

Dictionary Logging with log_params and log_metrics (recommended)

with mlflow.start_run(run_name='RF_tuned_params_scenario2') as run4:

    n_estimators      = 200
    criterion         = 'entropy'
    min_samples_split = 5
    min_samples_leaf  = 2

    model = RandomForestClassifier(
        n_estimators=n_estimators,
        criterion=criterion,
        min_samples_split=min_samples_split,
        min_samples_leaf=min_samples_leaf
    )
    model.fit(X_train, y_train)
    predictions = model.predict(X_test)
    predictions_proba = model.predict_proba(X_test)

    # Dictionary logging — more concise
    params = {
        'No. of trees':       n_estimators,
        'Splitting criteria': criterion,
        'Min samples split':  min_samples_split,
        'Min samples leaf':   min_samples_leaf
    }
    mlflow.log_params(params)

    metrics = {
        'Test_accuracy':        accuracy_score(y_test, predictions),
        'Test_precision_score': precision_score(y_test, predictions),
        'Test_recall_score':    recall_score(y_test, predictions),
        'Test_f1_score':        f1_score(y_test, predictions),
        'AUC_score':            roc_auc_score(y_test, predictions_proba[:, 1])
    }
    mlflow.log_metrics(metrics)

    mlflow.set_tag('Classifier', 'RF-tuned_params_sc2')

2.10 Visualizing and Sorting Runs (Demo)

# Scenario 3 with different parameters
with mlflow.start_run(run_name='RF_tuned_params_scenario3') as run5:

    n_estimators      = 500
    criterion         = 'gini'
    min_samples_split = 10
    min_samples_leaf  = 4

    model = RandomForestClassifier(
        n_estimators=n_estimators,
        criterion=criterion,
        min_samples_split=min_samples_split,
        min_samples_leaf=min_samples_leaf
    )
    model.fit(X_train, y_train)
    predictions = model.predict(X_test)
    predictions_proba = model.predict_proba(X_test)

    mlflow.log_params({
        'No. of trees':       n_estimators,
        'Splitting criteria': criterion,
        'Min samples split':  min_samples_split,
        'Min samples leaf':   min_samples_leaf
    })
    mlflow.log_metrics({
        'Test_accuracy':        accuracy_score(y_test, predictions),
        'Test_precision_score': precision_score(y_test, predictions),
        'Test_recall_score':    recall_score(y_test, predictions),
        'Test_f1_score':        f1_score(y_test, predictions),
        'AUC_score':            roc_auc_score(y_test, predictions_proba[:, 1])
    })
    mlflow.set_tag('Classifier', 'RF-tuned_params_sc3')

Experiments UI Features:

• Flask icon in the notebook → side panel with all runs
• Sort by date (ascending/descending)
• Sort by any metric or parameter
• Add custom metric/parameter columns via the + button

2.11 Comparing Runs (Demo)

Via the UI:

1. Select multiple runs (e.g. scenario1, 2 and 3)
2. Click the Compare button
3. The comparison page displays:
   • Run details: metadata for each run
   • Parameters: tabular comparison of hyperparameters
   • Metrics: tabular comparison of scores
   • Tags: label comparison
   • Parallel Coordinates Plot: visualization of parameter/metric relationships
   • Scatter Plot: e.g. Min_samples_leaf vs AUC_score
   • Box Plot: e.g. Splitting_criteria vs AUC_score

Example of the Parallel Coordinates Plot

Min samples leaf  |  Min samples split  |  No. of trees  |  AUC_score
─────────────────────────────────────────────────────────────────────
      2           │         5           │      200       │   0.87
      2           │         5           │      200       │   0.85
      4           │        10           │      500       │   0.86

2.12 Autologging with MLflow (Demo)

Autologging automatically logs parameters, metrics, and artifacts without any explicit logging instructions.

Enabling Autologging

# For scikit-learn only
mlflow.sklearn.autolog()

# For all frameworks (scikit-learn, XGBoost, TensorFlow, Keras, etc.)
mlflow.autolog()

Example with Autologging

mlflow.sklearn.autolog()

with mlflow.start_run(run_name='RF_tuned_params_scenario3_autolog') as run6:

    n_estimators      = 500
    criterion         = 'gini'
    min_samples_split = 10
    min_samples_leaf  = 4

    model = RandomForestClassifier(
        n_estimators=n_estimators,
        criterion=criterion,
        min_samples_split=min_samples_split,
        min_samples_leaf=min_samples_leaf
    )
    model.fit(X_train, y_train)

    # Only the tag is added manually — everything else is autologged
    mlflow.set_tag('Classifier', 'RF-tuned_params_sc3_autolog')

What Autolog Captures Automatically

Structure of an MLflow Model (standard format)

model/
├── MLmodel                 ← Flavor definitions and signature
├── model.pkl               ← Serialized model (scikit-learn)
├── conda.yaml              ← Conda environment
├── python_env.yaml         ← Python environment
├── requirements.txt        ← Dependencies
└── confusion_matrix.png    ← (auto-generated)
    precision_recall_curve.png
    training_roc_curve.png

Example MLmodel file:

flavors:
  python_function:
    env: conda.yaml
    loader_module: mlflow.sklearn
    model_path: model.pkl
    python_version: 3.8.x
  sklearn:
    pickled_model: model.pkl
    sklearn_version: 1.0.x
    serialization_format: cloudpickle
signature:
  inputs: '[{"name": "Married", "type": "long"}, ...]'
  outputs: '[{"type": "tensor", "tensor-spec": {...}}]'

Flavors: A convention allowing deployment tools to understand the model without having to integrate each tool with each library. The python_function flavor is supported by all MLflow deployment tools.

2.13 Programmatic Search and Sort of Runs (Demo)

# List all experiments in the workspace
experiments_list = client.list_experiments()

# Access an experiment by ID
experiment = mlflow.get_experiment(experiments_list[0].experiment_id)
print('Name:',             experiment.name)
print('Artifact Location:', experiment.artifact_location)
print('Tags:',             experiment.tags)
print('Lifecycle_stage:',  experiment.lifecycle_stage)

# Access an experiment by name
experiment = mlflow.get_experiment_by_name(
    '/Users/cloud.user@example.com/loan_approval_prediction'
)

# Get the tracking URI
mlflow.get_tracking_uri()

# List all runs of an experiment
mlflow.list_run_infos(experiment.experiment_id)

# Get the last active run
mlflow.last_active_run()

# Search runs with sorting
df_run_metrics = mlflow.search_runs(
    [experiment.experiment_id],
    order_by=['metrics.Test_accuracy DESC']
)
df_run_metrics

# Filter by condition (pandas)
best_runs_df = df_run_metrics[df_run_metrics['metrics.AUC_score'] > 0.7]

# Filter via MLflow filter_string
df_run_metrics = mlflow.search_runs(
    [experiment.experiment_id],
    filter_string="metrics.Test_accuracy > 0.8",
    order_by=['metrics.Test_accuracy DESC']
)

# Search with a condition on a parameter
# (via UI: params."Splitting criteria" = 'gini')

# Delete a run programmatically
random_run_id = df_run_metrics.loc[0, 'run_id']
mlflow.delete_run(random_run_id)
# Note: deleted runs can be restored for 30 days

Searches supported in the UI:

metrics.AUC_score > 0.8
params.`Splitting criteria` = 'gini'

3. Productionizing and Serving Models

3.1 Model Registry and Model Serving

Deployment Options with MLflow

graph TD
    M[Trained Model\nand Registered]
    M --> B[Batch Inference\nDelta Table → Spark DataFrame → UDF]
    M --> RT[Real-Time Serving]
    RT --> CL[Classic MLflow Serving\nREST Endpoint on a Cluster\nDedicated Single-Node Cluster]
    RT --> SV[Serverless Real-Time\nInferencing\nScalable Endpoint No Cluster\nManaged by Databricks]

    style B fill:#27ae60,color:#fff
    style CL fill:#e67e22,color:#fff
    style SV fill:#8e44ad,color:#fff

Classic MLflow Serving

• Hosts registered models from the Model Registry behind a REST endpoint
• Endpoints can be updated automatically based on new versions or stage changes
• Ideal for: predictable traffic, low throughput, non-critical applications

Serverless Real-Time Inferencing

• Scalable REST endpoints without cluster provisioning
• Compute fully managed by Databricks
• Built-in monitoring dashboard
• Integrated with the Databricks Feature Store
• ⚠️ At the time of course recording: in Public Preview

Model Registry

stateDiagram-v2
    [*] --> None : Initial Registration
    None --> Staging : transition_model_version_stage()
    Staging --> Production : transition_model_version_stage()
    Production --> Archived : transition_model_version_stage()
    Staging --> Archived : transition_model_version_stage()
    None --> Archived : transition_model_version_stage()

3.2 Training Multiple Models (Demo)

Training six different models with mlflow.autolog() to compare their performance and select the best one for deployment.

import mlflow
import numpy as np
import pandas as pd
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score
from sklearn.model_selection import train_test_split

# Load preprocessed data
loan_data = pd.read_csv('/dbfs/FileStore/datasets/loan_data_processed.csv', header='infer')

X = loan_data.drop(labels=['loan_approval_status'], axis=1)
y = loan_data['loan_approval_status']

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=123)

# Save test data for batch inference
X_test.to_csv('/dbfs/FileStore/datasets/processed_x_test.csv', index=False)

Model 1: Logistic Regression

from sklearn.linear_model import LogisticRegression

mlflow.autolog()

with mlflow.start_run(run_name='LR'):
    lr_model = LogisticRegression()
    lr_model.fit(X_train, y_train)

    predictions = lr_model.predict(X_test)
    predictions_predict_prob = lr_model.predict_proba(X_test)

    auc_score = roc_auc_score(y_test, predictions_predict_prob[:, 1])
    mlflow.log_metric('AUC_score', auc_score)  # AUC not autologged
    mlflow.set_tag('Classifier', 'LR-default parameters')

Model 2: Support Vector Classifier (SVC)

from sklearn.svm import SVC

mlflow.autolog()

with mlflow.start_run(run_name='SVC'):
    svc_model = SVC()
    svc_model.fit(X_train, y_train)
    predictions = svc_model.predict(X_test)

    auc_score = roc_auc_score(y_test, predictions_predict_prob[:, 1])
    mlflow.log_metric('AUC_score', auc_score)
    mlflow.set_tag('Classifier', 'SVC-default parameters')

Model 3: SGD Classifier

from sklearn.linear_model import SGDClassifier

mlflow.autolog()

with mlflow.start_run(run_name='SGD'):
    sgd_model = SGDClassifier()
    sgd_model.fit(X_train, y_train)
    predictions = sgd_model.predict(X_test)

    auc_score = roc_auc_score(y_test, predictions_predict_prob[:, 1])
    mlflow.log_metric('AUC_score', auc_score)
    mlflow.set_tag('Classifier', 'SGD-default parameters')

Model 4: Gaussian Naive Bayes (GNB)

from sklearn.naive_bayes import GaussianNB

mlflow.autolog()

with mlflow.start_run(run_name='GNB'):
    gnb_model = GaussianNB()
    gnb_model.fit(X_train, y_train)
    predictions = gnb_model.predict(X_test)
    predictions_predict_prob = gnb_model.predict_proba(X_test)

    auc_score = roc_auc_score(y_test, predictions_predict_prob[:, 1])
    mlflow.log_metric('AUC_score', auc_score)
    mlflow.set_tag('Classifier', 'GNB-default parameters')

Model 5: K-Nearest Neighbors (KNN)

from sklearn.neighbors import KNeighborsClassifier

mlflow.autolog()

with mlflow.start_run(run_name='KNN'):
    knn_model = KNeighborsClassifier()
    knn_model.fit(X_train, y_train)
    predictions = knn_model.predict(X_test)
    predictions_predict_prob = knn_model.predict_proba(X_test)

    auc_score = roc_auc_score(y_test, predictions_predict_prob[:, 1])
    mlflow.log_metric('AUC_score', auc_score)
    mlflow.set_tag('Classifier', 'KNN-default parameters')

Model 6: Random Forest (RF)

from sklearn.ensemble import RandomForestClassifier

mlflow.autolog()

with mlflow.start_run(run_name='RF'):
    rf_model = RandomForestClassifier()
    rf_model.fit(X_train, y_train)
    predictions = rf_model.predict(X_test)
    predictions_predict_prob = rf_model.predict_proba(X_test)

    auc_score = roc_auc_score(y_test, predictions_predict_prob[:, 1])
    mlflow.log_metric('AUC_score', auc_score)
    mlflow.set_tag('Classifier', 'RF-default parameters')

3.3 Registering a Model via the UI (Demo)

Via the graphical interface (for the GNB model):

1. Go to the GNB run page → Artifacts → model/ section
2. Click Register Model
3. Create a new model: loan_approval_gnb_model
4. In Models (left navigation) → the model appears as Version 1
5. Click Version 1 to view the input/output schema

Programmatic selection of the best model (by accuracy):

client = mlflow.tracking.MlflowClient()
experiments_list = client.list_experiments()

experiment = mlflow.get_experiment_by_name(experiments_list[0].name)

# Sort by test set accuracy (descending)
df_run_metrics = mlflow.search_runs(
    [experiment.experiment_id],
    order_by=['metrics.accuracy_score_X_test DESC']
)

df_run_metrics[['run_id', 'tags.Classifier', 'metrics.accuracy_score_X_test']]

best_run_id = df_run_metrics.loc[0, 'run_id']
print('Classifier:', df_run_metrics.loc[0, 'tags.Classifier'])
print('Run name:',   df_run_metrics.loc[0, 'tags.mlflow.runName'])

3.4 Loading and Using a Model for Predictions (Demo)

Loading a Model from a Run (without registration)

import mlflow.sklearn

# Build the model URI
model_uri = 'runs:/' + best_run_id + '/model'

# Load the model
best_model = mlflow.sklearn.load_model(model_uri=model_uri)

# Predictions
predictions_loaded   = best_model.predict(X_test)
predictions_original = lr_model.predict(X_test)

# Consistency check
assert(np.array_equal(predictions_loaded, predictions_original))

Programmatic Registration in the Model Registry

model_name = 'loan_approval_lr_model'

model_version = mlflow.register_model(
    f'runs:/{best_run_id}/model',
    model_name
)

3.5 Lifecycle Management and Stages (Demo)

Programmatic Stage Transitions

from mlflow.tracking import MlflowClient

client = MlflowClient()

# Transition to Production
client.transition_model_version_stage(
    name=model_name,
    version=model_version.version,
    stage='Production'
)

# Load from Production
loaded_model = mlflow.pyfunc.load_model(
    model_uri=f'models:/{model_name}/Production'
)
print(f'Accuracy_Test_set: {accuracy_score(y_test, loaded_model.predict(X_test))}')

Managing Multiple Versions and Models

# Select the best model by recall
best_run_recall = mlflow.search_runs(
    [experiment.experiment_id],
    order_by=['metrics.recall_score_X_test DESC']
)

print('Classifier:',  best_run_recall.loc[0, 'tags.Classifier'])
print('Best recall:', best_run_recall.loc[0, 'metrics.recall_score_X_test'])

# Register the new best model
best_recall_model_name = 'loan_approval_model'
best_recall_model = mlflow.register_model(
    f"runs:/{best_run_recall.loc[0, 'run_id']}/model",
    best_recall_model_name
)

# Archive the old model
client.transition_model_version_stage(
    name=model_name,
    version=model_version.version,
    stage='Archived'
)

# New model to Production
client.transition_model_version_stage(
    name=best_recall_model_name,
    version=best_recall_model.version,
    stage='Production'
)

Lifecycle Stage Summary

None ──► Staging ──► Production ──► Archived
         │                          ▲
         └──────────────────────────┘

3.6 Real-Time Serving via a REST Endpoint (Demo)

Steps in the UI:

1. Models → loan_approval_model → Serving → Enable Serving
2. Wait for the serving cluster to be created (visible in Compute → Job Clusters)
3. Generate an access token: Settings → User Settings → Generate New Token

Python Code to Call the REST Endpoint

import os
import requests
import json

os.environ['DATABRICKS_TOKEN'] = 'your_token_here'

def create_tf_serving_json(data):
    return {'inputs': {name: data[name].tolist() for name in data.keys()}}

def score_model(dataset):
    url = 'https://<your-workspace>.azuredatabricks.net/model/loan_approval_model/1/invocations'
    headers = {
        'Authorization': f'Bearer {os.environ.get("DATABRICKS_TOKEN")}',
        'Content-Type': 'application/json'
    }
    ds_dict = create_tf_serving_json(dataset)
    data_json = json.dumps(ds_dict, allow_nan=True)

    response = requests.request(method='POST', headers=headers, url=url, data=data_json)
    if response.status_code != 200:
        raise Exception(f'Request failed with status {response.status_code}, {response.text}')
    return response.json()

Comparing Predictions: Local Model vs Served Endpoint

model = mlflow.pyfunc.load_model(
    model_uri=f'models:/{best_recall_model_name}/{best_recall_model.version}'
)

num_predictions = 5
served_predictions = score_model(X_test[:num_predictions])
model_evaluations  = model.predict(X_test[:num_predictions])

pd.DataFrame({
    'Model Prediction':        model_evaluations,
    'Served Model Prediction': served_predictions['predictions']
})

Example JSON Payload for the Endpoint

{
  "inputs": {
    "Married":              [1, 1, 1, 0, 0],
    "num_of_dependents":    [1, 2, 0, 0, 0],
    "education_level":      [0, 0, 1, 0, 0],
    "working_status":       [1, 0, 0, 0, 0],
    "income":               [2895, 3283, 2333, 3676, 4300],
    "spouse_income":        [0.0, 2035.0, 1516.0, 4301.0, 0.0],
    "loan_amount":          [95.0, 148.0, 95.0, 172.0, 136.0],
    "monthly_installment":  [360.0, 360.0, 360.0, 360.0, 360.0],
    "credit_history":       [1.0, 1.0, 1.0, 1.0, 0.0],
    "Gender_Male":          [1, 1, 1, 1, 0],
    "Gender_Others":        [0, 0, 0, 0, 0],
    "property_type_Semiurban": [1, 0, 0, 0, 1],
    "property_type_Urban":  [0, 1, 1, 0, 0]
  }
}

3.7 Model Signature and Programmatic Management (Demo)

Logging a Model with Inferred Signature

from mlflow.models.signature import infer_signature

# Disable autologging for full control
mlflow.autolog(disable=True)

with mlflow.start_run(run_name='RF_with_inferred_signature'):
    model = RandomForestClassifier()
    model.fit(X_train, y_train)
    predictions = model.predict(X_test)
    predictions_predict_prob = model.predict_proba(X_test)

    metrics = {
        'Test_accuracy':        accuracy_score(y_test, predictions),
        'Test_precision_score': precision_score(y_test, predictions),
        'Test_recall_score':    recall_score(y_test, predictions),
        'Test_f1_score':        f1_score(y_test, predictions),
        'AUC_score':            roc_auc_score(y_test, predictions_predict_prob[:, 1])
    }
    mlflow.log_metrics(metrics)

    # Automatically infer the signature (input/output schema)
    signature = infer_signature(X_train, model.predict(X_train))

    # Log the model with signature in the Model Registry
    mlflow.sklearn.log_model(
        model,
        'RF_model_with_signature_example',
        registered_model_name='loan_approval_model',
        signature=signature
    )
    mlflow.set_tag('Classifier', 'RF-default_parameters-with_input_signature_example')

Programmatic Management of Registered Versions

from pprint import pprint

# Update version descriptions
client.update_model_version(
    name='loan_approval_model',
    version=1,
    description='This model version is a scikit-learn model which had the best recall score'
)

client.update_model_version(
    name='loan_approval_model',
    version=2,
    description='This model version is a scikit-learn random forest model with 100 decision trees'
)

# List all registered models
for rm in client.list_registered_models():
    pprint(dict(rm), indent=4)

# Search versions of a specific model
for mv in client.search_model_versions("name='loan_approval_model'"):
    pprint(dict(mv), indent=4)

4. Custom Models and MLflow Projects

4.1 Custom Models and Custom Flavors

MLflow natively supports many frameworks, but it is also possible to create custom Python models by subclassing mlflow.pyfunc.PythonModel.

When to Use Custom Models?

• Preprocessing logic to include in the model
• Models that don’t fit any standard framework
• Customization of how predictions are generated

Custom Model Architecture

mlflow.pyfunc.PythonModel (base class)
└── Your custom class
    ├── __init__(self, ...)      ← Initialization with parameters
    └── predict(self, context, model_input)  ← Custom prediction logic

4.2 Creating a Custom Model (Demo)

Goal: Create a custom model that standardizes numerical input columns before prediction.

import mlflow.pyfunc
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import StandardScaler

# Define columns
cat_cols = [
    'Gender_Male', 'Gender_Others', 'Married', 'num_of_dependents',
    'education_level', 'working_status', 'credit_history',
    'loan_approval_status', 'Property_type_Semiurban', 'Property_type_Urban'
]
numeric_cols = ['income', 'spouse_income', 'loan_amount', 'monthly_installment']

# Define the custom model
class DataScaler(mlflow.pyfunc.PythonModel):

    def __init__(self, numeric_cols):
        self.numeric_cols = numeric_cols

    def predict(self, context, model_input):
        ct = ColumnTransformer(
            [('stdscale', StandardScaler(), numeric_cols)],
            remainder='passthrough',
            verbose_feature_names_out=False
        )
        op_df_ct = pd.DataFrame(
            data=ct.fit_transform(model_input),
            columns=ct.get_feature_names_out()
        )
        return op_df_ct.reindex(columns=model_input.columns)

Saving and Logging the Custom Model

model_path = 'standard_scale_numeric_input_data'

with mlflow.start_run(run_name='custom_scale'):
    scale_model = DataScaler(numeric_cols)

    # Local save
    mlflow.pyfunc.save_model(path=model_path, python_model=scale_model)

    # Log to the experiment
    mlflow.pyfunc.log_model('std_scale', python_model=scale_model)

Loading and Using

loaded_model = mlflow.pyfunc.load_model(model_path)

model_input  = X_train
model_output = loaded_model.predict(model_input)
model_output

Accuracy Verification

from sklearn.preprocessing import StandardScaler

sc = StandardScaler()
X_train_num = sc.fit_transform(X_train[numeric_cols])

assert model_output[numeric_cols].equals(pd.DataFrame(X_train_num, columns=numeric_cols))
print('Dataframes are identical')

Custom Model Artifacts in MLflow

model/
├── MLmodel           ← python_function flavor
├── python_model.pkl  ← Serialized Python model (your class)
└── conda.yaml        ← Conda environment

4.3 MLflow Projects

An MLflow Project is a standard format for packaging ML code in a reusable and reproducible way.

Structure of an MLflow Project

my_project/
├── MLproject            ← Project definition file
├── conda.yaml           ← Execution environment
├── train_clf.py         ← Training script
├── loan_data_processed.csv
└── README.md

MLproject File

name: mlflow_rf

conda_env: conda.yaml

entry_points:
  main:
    parameters:
      n_estimators:      {type: int, default: 100}
      min_samples_split: {type: int, default: 2}
      min_samples_leaf:  {type: int, default: 2}
    command: "python train_clf.py {n_estimators} {min_samples_split} {min_samples_leaf}"

conda.yaml File

name: mlflow_rf
channels:
  - defaults
dependencies:
  - numpy>=1.14.3
  - pandas>=1.0.0
  - scikit-learn>=0.24.1
  - pip
  - pip:
    - mlflow

Training Script train_clf.py

import os
import warnings
import sys

import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score
from sklearn.model_selection import train_test_split

import mlflow
import mlflow.sklearn

if __name__ == '__main__':
    warnings.filterwarnings("ignore")
    np.random.seed(40)

    # Load data from the same directory
    loan_path = os.path.join(
        os.path.dirname(os.path.abspath(__file__)),
        'loan_data_processed.csv'
    )
    loan_data = pd.read_csv(loan_path)

    X = loan_data.drop(labels=['loan_approval_status'], axis=1)
    y = loan_data['loan_approval_status']

    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.3, stratify=y, random_state=123
    )

    # Read hyperparameters from command line
    n_estimators      = int(sys.argv[1]) if len(sys.argv) > 1 else 100
    min_samples_split = int(sys.argv[2]) if len(sys.argv) > 2 else 2
    min_samples_leaf  = int(sys.argv[3]) if len(sys.argv) > 3 else 2

    with mlflow.start_run():
        model = RandomForestClassifier(
            n_estimators=n_estimators,
            min_samples_split=min_samples_split,
            min_samples_leaf=min_samples_leaf
        )
        model.fit(X_train, y_train)
        predictions = model.predict(X_test)
        predictions_proba = model.predict_proba(X_test)

        metrics = {
            'Test_accuracy':        accuracy_score(y_test, predictions),
            'Test_precision_score': precision_score(y_test, predictions),
            'Test_recall_score':    recall_score(y_test, predictions),
            'Test_f1_score':        f1_score(y_test, predictions),
            'AUC_score':            roc_auc_score(y_test, predictions_proba[:, 1])
        }

        mlflow.log_metrics(metrics)
        mlflow.set_tag('Classifier', 'RF-tuned parameters')
        mlflow.sklearn.log_model(model, 'RF-tuned parameters')

4.4 Running an MLflow Project (Demo)

From GitHub

import mlflow

project_uri = 'https://github.com/loonyuser/mlflow-databricks'
params = {
    'n_estimators':      200,
    'min_samples_split': 3,
    'min_samples_leaf':  2
}

mlflow.run(project_uri, parameters=params)

From DBFS (Databricks File System)

# Upload project files to DBFS → FileStore → project/
project_uri = '/dbfs/FileStore/project'

params = {
    'n_estimators':      200,
    'min_samples_split': 3,
    'min_samples_leaf':  2
}

mlflow.run(project_uri, parameters=params)

Benefits of MLflow Projects

5. Summary and Additional Resources

Complete Workflow Recap

flowchart TD
    A[Raw Data\ncustomers.csv] --> B[Preprocessing\nEncoding, Split]
    B --> C[Experiments & Runs\nMLflow Tracking]

    C --> C1[log_param / log_params]
    C --> C2[log_metric / log_metrics]
    C --> C3[log_figure / log_artifact]
    C --> C4[mlflow.autolog]
    C --> C5[set_tag]

    C --> D[Run Comparison\nParallel Coordinates\nScatter / Box Plot]

    D --> E[Best Model Selection\nmlflow.search_runs]

    E --> F[Model Registry\nmlflow.register_model]

    F --> G1[Batch Inference\nDelta Table → Spark UDF]
    F --> G2[Classic Serving\nREST Endpoint on Cluster]
    F --> G3[Serverless Serving\nScalable No Cluster]

    F --> H[Lifecycle Management]
    H --> H1[None → Staging]
    H1 --> H2[Staging → Production]
    H2 --> H3[Production → Archived]

    style C fill:#ff7700,color:#fff
    style F fill:#27ae60,color:#fff
    style H fill:#8e44ad,color:#fff

Essential MLflow Functions Cheat Sheet

TRACKING
────────────────────────────────────────────────────────────────
mlflow.set_experiment(name)           → Set the active experiment
mlflow.start_run(run_name=...)        → Start a run
mlflow.end_run()                      → End a run
mlflow.active_run()                   → Get the active run
mlflow.last_active_run()              → Get the last run

LOGGING
────────────────────────────────────────────────────────────────
mlflow.log_param(key, value)          → Log a parameter
mlflow.log_params(dict)               → Log multiple parameters
mlflow.log_metric(key, value)         → Log a metric
mlflow.log_metrics(dict)              → Log multiple metrics
mlflow.log_figure(fig, filename)      → Log a matplotlib figure
mlflow.set_tag(key, value)            → Set a tag
mlflow.sklearn.autolog()              → sklearn autologging
mlflow.autolog()                      → All-framework autologging

SEARCH
────────────────────────────────────────────────────────────────
mlflow.search_runs([exp_id], ...)     → Search for runs
mlflow.get_experiment(exp_id)         → Get an experiment
mlflow.get_experiment_by_name(name)   → Get by name
mlflow.list_run_infos(exp_id)         → List runs
mlflow.delete_run(run_id)             → Delete a run

LOADING
────────────────────────────────────────────────────────────────
mlflow.sklearn.load_model(model_uri)  → Load sklearn model
mlflow.pyfunc.load_model(model_uri)   → Load pyfunc model
mlflow.pyfunc.save_model(path, ...)   → Save locally

REGISTRY
────────────────────────────────────────────────────────────────
mlflow.register_model(uri, name)      → Register a model
mlflow.sklearn.log_model(..., registered_model_name=...) → Log + Register
client.transition_model_version_stage(name, version, stage)
client.update_model_version(name, version, description)
client.list_registered_models()
client.search_model_versions("name='...'")

MLflow Model URI Syntax

runs:/<run_id>/model              → Model in a specific run
models:/<model_name>/<version>    → Specific registry version
models:/<model_name>/Production   → Production version
models:/<model_name>/Staging      → Staging version

Recommended Complementary Courses

Course by Janani Ravi — Loonycorn

Search Terms

managing · models · mlflow · databricks · azure · spark · data · engineering · analytics · model · custom · programmatic · runs · tracking · endpoint · lifecycle · loading · logging · run · serving · autologging · management · method · projects