MLOps·Model Lifecycle Management·Advanced

Model Versioning Mlflow

Version and register trained ML models in the MLflow model registry with formal stage transition gates from Staging to Production to Archived, enabling controlled and auditable model promotion workflows with full rollback capability to any previously registered model version.

machine-learningml-engineeringmlops

Model Versioning and Registration with MLflow

This notebook provides a practical guide to implementing model versioning and registration using MLflow, crucial components for MLOps. Effective model management ensures reproducibility, traceability, and seamless deployment of machine learning models in production.

Concepts Overview

Concept	Description
Model Versioning	Tracking changes to ML models over time, including code, data, and hyperparameters, to ensure reproducibility.
Model Registration	Centralizing models in a repository with associated metadata, enabling easy discovery and deployment.
Metadata	Descriptive information about a model (e.g., training data, metrics, creation date, author).
Model Lineage	The complete history of a model, from data ingestion and preprocessing to training and deployment.
MLflow Tracking	An API and UI for logging parameters, code versions, metrics, and output files when running machine learning code.
MLflow Model Registry	A centralized model store that provides a complete workflow to manage MLflow Models, from staging to production.

Dependency Installation

We'll install necessary libraries for data manipulation, machine learning, plotting, and logging.

[1]

import sys
!{sys.executable} -m pip install scikit-learn pandas numpy matplotlib seaborn loguru tqdm mlflow

Requirement already satisfied: scikit-learn in /usr/local/lib/python3.12/dist-packages (1.6.1)
Requirement already satisfied: pandas in /usr/local/lib/python3.12/dist-packages (2.2.2)
Requirement already satisfied: numpy in /usr/local/lib/python3.12/dist-packages (2.0.2)
Requirement already satisfied: matplotlib in /usr/local/lib/python3.12/dist-packages (3.10.0)
Requirement already satisfied: seaborn in /usr/local/lib/python3.12/dist-packages (0.13.2)
Collecting loguru
  Downloading loguru-0.7.3-py3-none-any.whl.metadata (22 kB)
Requirement already satisfied: tqdm in /usr/local/lib/python3.12/dist-packages (4.67.3)
Collecting mlflow
  Downloading mlflow-3.13.0-py3-none-any.whl.metadata (49 kB)
[2K     [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m49.4/49.4 kB[0m [31m2.0 MB/s[0m eta [36m0:00:00[0m
[?25hRequirement already satisfied: scipy>=1.6.0 in /usr/local/lib/python3.12/dist-packages (from scikit-learn) (1.16.3)
Requirement already satisfied: joblib>=1.2.0 in /usr/local/lib/python3.12/dist-packages (from scikit-learn) (1.5.3)
Requirement already satisfied: threadpoolctl>=3.1.0 in /usr/local/lib/python3.12/dist-packages (from scikit-learn) (3.6.0)
Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.12/dist-packages (from pandas) (2.9.0.post0)
Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.12/dist-packages (from pandas) (2025.2)
Requirement already satisfied: tzdata>=2022.7 in /usr/local/lib/python3.12/dist-packages (from pandas) (2026.2)
Requirement already satisfied: contourpy>=1.0.1 in /usr/local/lib/python3.12/dist-packages (from matplotlib) (1.3.3)
Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.12/dist-packages (from matplotlib) (0.12.1)
Requirement already satisfied: fonttools>=4.22.0 in /usr/local/lib/python3.12/dist-packages (from matplotlib) (4.63.0)
Requirement already satisfied: kiwisolver>=1.3.1 in /usr/local/lib/python3.12/dist-packages (from matplotlib) (1.5.0)
Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.12/dist-packages (from matplotlib) (26.2)
Requirement already satisfied: pillow>=8 in /usr/local/lib/python3.12/dist-packages (from matplotlib) (11.3.0)
Requirement already satisfied: pyparsing>=2.3.1 in /usr/local/lib/python3.12/dist-packages (from matplotlib) (3.3.2)
Collecting mlflow-skinny==3.13.0 (from mlflow)
  Downloading mlflow_skinny-3.13.0-py3-none-any.whl.metadata (50 kB)
[2K     [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m50.2/50.2 kB[0m [31m1.5 MB/s[0m eta [36m0:00:00[0m
[?25hCollecting mlflow-tracing==3.13.0 (from mlflow)
  Downloading mlflow_tracing-3.13.0-py3-none-any.whl.metadata (19 kB)
Collecting Flask-CORS<7 (from mlflow)
  Downloading flask_cors-6.0.5-py3-none-any.whl.metadata (5.4 kB)
Requirement already satisfied: Flask<4 in /usr/local/lib/python3.12/dist-packages (from mlflow) (3.1.3)
Requirement already satisfied: aiohttp<4 in /usr/local/lib/python3.12/dist-packages (from mlflow) (3.14.0)
Requirement already satisfied: alembic!=1.10.0,<2 in /usr/local/lib/python3.12/dist-packages (from mlflow) (1.18.4)
Requirement already satisfied: cryptography<49,>=43.0.0 in /usr/local/lib/python3.12/dist-packages (from mlflow) (48.0.0)
Collecting docker<8,>=4.0.0 (from mlflow)
  Downloading docker-7.1.0-py3-none-any.whl.metadata (3.8 kB)
Collecting graphene<4 (from mlflow)
  Downloading graphene-3.4.3-py2.py3-none-any.whl.metadata (6.9 kB)
Collecting gunicorn<27 (from mlflow)
  Downloading gunicorn-26.0.0-py3-none-any.whl.metadata (5.4 kB)
Collecting huey<4,>=2.5.4 (from mlflow)
  Downloading huey-3.0.3-py3-none-any.whl.metadata (4.5 kB)
Requirement already satisfied: pyarrow<25,>=4.0.0 in /usr/local/lib/python3.12/dist-packages (from mlflow) (18.1.0)
Collecting skops<1 (from mlflow)
  Downloading skops-0.14.0-py3-none-any.whl.metadata (4.4 kB)
Requirement already satisfied: sqlalchemy<3,>=1.4.0 in /usr/local/lib/python3.12/dist-packages (from mlflow) (2.0.50)
Requirement already satisfied: cachetools<8,>=5.0.0 in /usr/local/lib/python3.12/dist-packages (from mlflow-skinny==3.13.0->mlflow) (6.2.6)
Requirement already satisfied: click<9,>=7.0 in /usr/local/lib/python3.12/dist-packages (from mlflow-skinny==3.13.0->mlflow) (8.4.1)
Requirement already satisfied: cloudpickle<4 in /usr/local/lib/python3.12/dist-packages (from mlflow-skinny==3.13.0->mlflow) (3.1.2)
Collecting databricks-sdk<1,>=0.20.0 (from mlflow-skinny==3.13.0->mlflow)
  Downloading databricks_sdk-0.117.0-py3-none-any.whl.metadata (43 kB)
[2K     [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m43.5/43.5 kB[0m [31m2.8 MB/s[0m eta [36m0:00:00[0m
[?25hRequirement already satisfied: fastapi<1 in /usr/local/lib/python3.12/dist-packages (from mlflow-skinny==3.13.0->mlflow) (0.136.3)
Requirement already satisfied: gitpython<4,>=3.1.9 in /usr/local/lib/python3.12/dist-packages (from mlflow-skinny==3.13.0->mlflow) (3.1.50)
Requirement already satisfied: importlib_metadata!=4.7.0,<10,>=3.7.0 in /usr/local/lib/python3.12/dist-packages (from mlflow-skinny==3.13.0->mlflow) (8.7.1)
Requirement already satisfied: opentelemetry-api<3,>=1.9.0 in /usr/local/lib/python3.12/dist-packages (from mlflow-skinny==3.13.0->mlflow) (1.38.0)
Requirement already satisfied: opentelemetry-proto<3,>=1.9.0 in /usr/local/lib/python3.12/dist-packages (from mlflow-skinny==3.13.0->mlflow) (1.38.0)
Requirement already satisfied: opentelemetry-sdk<3,>=1.9.0 in /usr/local/lib/python3.12/dist-packages (from mlflow-skinny==3.13.0->mlflow) (1.38.0)
Requirement already satisfied: protobuf<8,>=3.12.0 in /usr/local/lib/python3.12/dist-packages (from mlflow-skinny==3.13.0->mlflow) (5.29.6)
Requirement already satisfied: pydantic<3,>=2.0.0 in /usr/local/lib/python3.12/dist-packages (from mlflow-skinny==3.13.0->mlflow) (2.12.3)
Requirement already satisfied: python-dotenv<2,>=0.19.0 in /usr/local/lib/python3.12/dist-packages (from mlflow-skinny==3.13.0->mlflow) (1.2.2)
Requirement already satisfied: pyyaml<7,>=5.1 in /usr/local/lib/python3.12/dist-packages (from mlflow-skinny==3.13.0->mlflow) (6.0.3)
Requirement already satisfied: requests<3,>=2.17.3 in /usr/local/lib/python3.12/dist-packages (from mlflow-skinny==3.13.0->mlflow) (2.32.4)
Requirement already satisfied: sqlparse<1,>=0.4.0 in /usr/local/lib/python3.12/dist-packages (from mlflow-skinny==3.13.0->mlflow) (0.5.5)
Requirement already satisfied: starlette<2 in /usr/local/lib/python3.12/dist-packages (from mlflow-skinny==3.13.0->mlflow) (0.52.1)
Requirement already satisfied: typing-extensions<5,>=4.0.0 in /usr/local/lib/python3.12/dist-packages (from mlflow-skinny==3.13.0->mlflow) (4.15.0)
Requirement already satisfied: uvicorn<1 in /usr/local/lib/python3.12/dist-packages (from mlflow-skinny==3.13.0->mlflow) (0.49.0)
Requirement already satisfied: aiohappyeyeballs>=2.5.0 in /usr/local/lib/python3.12/dist-packages (from aiohttp<4->mlflow) (2.6.2)
Requirement already satisfied: aiosignal>=1.4.0 in /usr/local/lib/python3.12/dist-packages (from aiohttp<4->mlflow) (1.4.0)
Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.12/dist-packages (from aiohttp<4->mlflow) (26.1.0)
Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.12/dist-packages (from aiohttp<4->mlflow) (1.8.0)
Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.12/dist-packages (from aiohttp<4->mlflow) (6.7.1)
Requirement already satisfied: propcache>=0.2.0 in /usr/local/lib/python3.12/dist-packages (from aiohttp<4->mlflow) (0.5.2)
Requirement already satisfied: yarl<2.0,>=1.17.0 in /usr/local/lib/python3.12/dist-packages (from aiohttp<4->mlflow) (1.24.2)
Requirement already satisfied: Mako in /usr/local/lib/python3.12/dist-packages (from alembic!=1.10.0,<2->mlflow) (1.3.12)
Requirement already satisfied: cffi>=2.0.0 in /usr/local/lib/python3.12/dist-packages (from cryptography<49,>=43.0.0->mlflow) (2.0.0)
Requirement already satisfied: urllib3>=1.26.0 in /usr/local/lib/python3.12/dist-packages (from docker<8,>=4.0.0->mlflow) (2.5.0)
Requirement already satisfied: blinker>=1.9.0 in /usr/local/lib/python3.12/dist-packages (from Flask<4->mlflow) (1.9.0)
Requirement already satisfied: itsdangerous>=2.2.0 in /usr/local/lib/python3.12/dist-packages (from Flask<4->mlflow) (2.2.0)
Requirement already satisfied: jinja2>=3.1.2 in /usr/local/lib/python3.12/dist-packages (from Flask<4->mlflow) (3.1.6)
Requirement already satisfied: markupsafe>=2.1.1 in /usr/local/lib/python3.12/dist-packages (from Flask<4->mlflow) (3.0.3)
Requirement already satisfied: werkzeug>=3.1.0 in /usr/local/lib/python3.12/dist-packages (from Flask<4->mlflow) (3.1.8)
Collecting graphql-core<3.3,>=3.1 (from graphene<4->mlflow)
  Downloading graphql_core-3.2.11-py3-none-any.whl.metadata (11 kB)
Collecting graphql-relay<3.3,>=3.1 (from graphene<4->mlflow)
  Downloading graphql_relay-3.2.0-py3-none-any.whl.metadata (12 kB)
Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.12/dist-packages (from python-dateutil>=2.8.2->pandas) (1.17.0)
Requirement already satisfied: prettytable>=3.9 in /usr/local/lib/python3.12/dist-packages (from skops<1->mlflow) (3.17.0)
Requirement already satisfied: greenlet>=1 in /usr/local/lib/python3.12/dist-packages (from sqlalchemy<3,>=1.4.0->mlflow) (3.5.1)
Requirement already satisfied: pycparser in /usr/local/lib/python3.12/dist-packages (from cffi>=2.0.0->cryptography<49,>=43.0.0->mlflow) (3.0)
Requirement already satisfied: google-auth~=2.0 in /usr/local/lib/python3.12/dist-packages (from databricks-sdk<1,>=0.20.0->mlflow-skinny==3.13.0->mlflow) (2.47.0)
Requirement already satisfied: typing-inspection>=0.4.2 in /usr/local/lib/python3.12/dist-packages (from fastapi<1->mlflow-skinny==3.13.0->mlflow) (0.4.2)
Requirement already satisfied: annotated-doc>=0.0.2 in /usr/local/lib/python3.12/dist-packages (from fastapi<1->mlflow-skinny==3.13.0->mlflow) (0.0.4)
Requirement already satisfied: gitdb<5,>=4.0.1 in /usr/local/lib/python3.12/dist-packages (from gitpython<4,>=3.1.9->mlflow-skinny==3.13.0->mlflow) (4.0.12)
Requirement already satisfied: zipp>=3.20 in /usr/local/lib/python3.12/dist-packages (from importlib_metadata!=4.7.0,<10,>=3.7.0->mlflow-skinny==3.13.0->mlflow) (4.1.0)
Requirement already satisfied: opentelemetry-semantic-conventions==0.59b0 in /usr/local/lib/python3.12/dist-packages (from opentelemetry-sdk<3,>=1.9.0->mlflow-skinny==3.13.0->mlflow) (0.59b0)
Requirement already satisfied: wcwidth in /usr/local/lib/python3.12/dist-packages (from prettytable>=3.9->skops<1->mlflow) (0.7.0)
Requirement already satisfied: annotated-types>=0.6.0 in /usr/local/lib/python3.12/dist-packages (from pydantic<3,>=2.0.0->mlflow-skinny==3.13.0->mlflow) (0.7.0)
Requirement already satisfied: pydantic-core==2.41.4 in /usr/local/lib/python3.12/dist-packages (from pydantic<3,>=2.0.0->mlflow-skinny==3.13.0->mlflow) (2.41.4)
Requirement already satisfied: charset_normalizer<4,>=2 in /usr/local/lib/python3.12/dist-packages (from requests<3,>=2.17.3->mlflow-skinny==3.13.0->mlflow) (3.4.7)
Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.12/dist-packages (from requests<3,>=2.17.3->mlflow-skinny==3.13.0->mlflow) (3.18)
Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.12/dist-packages (from requests<3,>=2.17.3->mlflow-skinny==3.13.0->mlflow) (2026.5.20)
Requirement already satisfied: anyio<5,>=3.6.2 in /usr/local/lib/python3.12/dist-packages (from starlette<2->mlflow-skinny==3.13.0->mlflow) (4.13.0)
Requirement already satisfied: h11>=0.8 in /usr/local/lib/python3.12/dist-packages (from uvicorn<1->mlflow-skinny==3.13.0->mlflow) (0.16.0)
Requirement already satisfied: smmap<6,>=3.0.1 in /usr/local/lib/python3.12/dist-packages (from gitdb<5,>=4.0.1->gitpython<4,>=3.1.9->mlflow-skinny==3.13.0->mlflow) (5.0.3)
Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.12/dist-packages (from google-auth~=2.0->databricks-sdk<1,>=0.20.0->mlflow-skinny==3.13.0->mlflow) (0.4.2)
Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.12/dist-packages (from google-auth~=2.0->databricks-sdk<1,>=0.20.0->mlflow-skinny==3.13.0->mlflow) (4.9.1)
Requirement already satisfied: pyasn1<0.7.0,>=0.6.1 in /usr/local/lib/python3.12/dist-packages (from pyasn1-modules>=0.2.1->google-auth~=2.0->databricks-sdk<1,>=0.20.0->mlflow-skinny==3.13.0->mlflow) (0.6.3)
Downloading loguru-0.7.3-py3-none-any.whl (61 kB)
[2K   [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m61.6/61.6 kB[0m [31m3.4 MB/s[0m eta [36m0:00:00[0m
[?25hDownloading mlflow-3.13.0-py3-none-any.whl (10.8 MB)
[2K   [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m10.8/10.8 MB[0m [31m52.2 MB/s[0m eta [36m0:00:00[0m
[?25hDownloading mlflow_skinny-3.13.0-py3-none-any.whl (3.4 MB)
[2K   [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m3.4/3.4 MB[0m [31m58.4 MB/s[0m eta [36m0:00:00[0m
[?25hDownloading mlflow_tracing-3.13.0-py3-none-any.whl (1.7 MB)
[2K   [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m1.7/1.7 MB[0m [31m43.3 MB/s[0m eta [36m0:00:00[0m
[?25hDownloading docker-7.1.0-py3-none-any.whl (147 kB)
[2K   [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m147.8/147.8 kB[0m [31m7.6 MB/s[0m eta [36m0:00:00[0m
[?25hDownloading flask_cors-6.0.5-py3-none-any.whl (16 kB)
Downloading graphene-3.4.3-py2.py3-none-any.whl (114 kB)
[2K   [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m114.9/114.9 kB[0m [31m4.2 MB/s[0m eta [36m0:00:00[0m
[?25hDownloading gunicorn-26.0.0-py3-none-any.whl (212 kB)
[2K   [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m212.0/212.0 kB[0m [31m11.2 MB/s[0m eta [36m0:00:00[0m
[?25hDownloading huey-3.0.3-py3-none-any.whl (94 kB)
[2K   [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m94.9/94.9 kB[0m [31m6.1 MB/s[0m eta [36m0:00:00[0m
[?25hDownloading skops-0.14.0-py3-none-any.whl (132 kB)
[2K   [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m132.2/132.2 kB[0m [31m8.2 MB/s[0m eta [36m0:00:00[0m
[?25hDownloading databricks_sdk-0.117.0-py3-none-any.whl (936 kB)
[2K   [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m936.9/936.9 kB[0m [31m21.9 MB/s[0m eta [36m0:00:00[0m
[?25hDownloading graphql_core-3.2.11-py3-none-any.whl (214 kB)
[2K   [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m214.9/214.9 kB[0m [31m12.0 MB/s[0m eta [36m0:00:00[0m
[?25hDownloading graphql_relay-3.2.0-py3-none-any.whl (16 kB)
Installing collected packages: huey, loguru, gunicorn, graphql-core, graphql-relay, docker, skops, graphene, Flask-CORS, databricks-sdk, mlflow-tracing, mlflow-skinny, mlflow
Successfully installed Flask-CORS-6.0.5 databricks-sdk-0.117.0 docker-7.1.0 graphene-3.4.3 graphql-core-3.2.11 graphql-relay-3.2.0 gunicorn-26.0.0 huey-3.0.3 loguru-0.7.3 mlflow-3.13.0 mlflow-skinny-3.13.0 mlflow-tracing-3.13.0 skops-0.14.0

Library Imports

Import all standard and third-party libraries required for this notebook.

[15]

import pandas as pd
import numpy as np
import datetime
import random
import time
import pickle
from collections import deque
from typing import Tuple # Added for type hinting

from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score

import matplotlib.pyplot as plt
import seaborn as sns
from loguru import logger

import os
# IMPORTANT: Allow MLflow to use the filesystem backend for this demonstration.
# This is needed because the filesystem tracking backend is in maintenance mode and requires explicit opt-in.
os.environ["MLFLOW_ALLOW_FILE_STORE"] = "true"

import mlflow
import mlflow.sklearn

# Configure logger to output to stderr, useful in notebooks
logger.remove()
logger.add(sys.stderr, format="<green>{time}</green> <level>{level}</level> <red>{message}</red>", level="INFO")

Core Functions

This section defines the core functions for model training and evaluation. MLflow will be used for model versioning and registration, replacing the custom functions previously defined for these tasks.

Function Name: `generate_mock_classification_data`

This function creates a synthetic dataset suitable for binary classification tasks. It generates n_samples data points with n_features features, where some features are correlated with the target variable to simulate a real-world scenario. A small amount of random noise is added to make the problem non-trivial.

Parameters: state (dict): The current state dictionary. n_samples (int): The number of samples to generate. n_features (int): The number of features for each sample. random_state (int): Seed for random number generation for reproducibility.

Returns: (dict): The updated state dictionary with the generated 'mock_data' added to the datasets key.

[3]

def generate_mock_classification_data(state: dict, n_samples: int = 1000,
                                    n_features: int = 10,
                                    random_state: int = 42) -> dict:
    """
    Generates mock classification data and stores it in the state dictionary.

    Parameters
    ----------
    state : dict
        Current state dictionary.
    n_samples : int, optional
        Number of samples to generate, defaults to 1000.
    n_features : int, optional
        Number of features, defaults to 10.
    random_state : int, optional
        Random seed for reproducibility, defaults to 42.

    Returns
    -------
    dict
        Updated state with 'mock_data' (DataFrame) added to 'datasets'.
    """
    logger.info(f"Generating {n_samples} samples with {n_features} features...")
    np.random.seed(random_state)

    X = np.random.rand(n_samples, n_features)
    # Create a target variable that depends on some features with noise
    y = (X[:, 0] * 0.5 + X[:, 1] * 0.3 + np.random.rand(n_samples) * 0.2 > 0.6).astype(int)

    feature_names = [f'feature_{i}' for i in range(n_features)]
    df = pd.DataFrame(X, columns=feature_names)
    df['target'] = y

    state['datasets']['mock_data'] = df
    logger.info(f"Mock data generated with shape {df.shape}.")
    return state

Function Name: `generate_mock_classification_data`

Returns: (dict): The updated state dictionary with the generated 'mock_data' added to the datasets key.

[ ]

def generate_mock_classification_data(state: dict, n_samples: int = 1000,
                                    n_features: int = 10,
                                    random_state: int = 42) -> dict:
    """
    Generates mock classification data and stores it in the state dictionary.

    Parameters
    ----------
    state : dict
        Current state dictionary.
    n_samples : int, optional
        Number of samples to generate, defaults to 1000.
    n_features : int, optional
        Number of features, defaults to 10.
    random_state : int, optional
        Random seed for reproducibility, defaults to 42.

    Returns
    -------
    dict
        Updated state with 'mock_data' (DataFrame) added to 'datasets'.
    """
    logger.info(f"Generating {n_samples} samples with {n_features} features...")
    np.random.seed(random_state)

    X = np.random.rand(n_samples, n_features)
    # Create a target variable that depends on some features with noise
    y = (X[:, 0] * 0.5 + X[:, 1] * 0.3 + np.random.rand(n_samples) * 0.2 > 0.6).astype(int)

    feature_names = [f'feature_{i}' for i in range(n_features)]
    df = pd.DataFrame(X, columns=feature_names)
    df['target'] = y

    state['datasets']['mock_data'] = df
    logger.info(f"Mock data generated with shape {df.shape}.")
    return state

Function Name: `train_ml_model`

This function trains a logistic regression model on the provided data. It splits the data into training and testing sets, then initializes and fits a LogisticRegression model. After training, it uses MLflow to log parameters, metrics, and the trained model. It also has an option to register the model in the MLflow Model Registry, making it discoverable and manageable across experiments.

Parameters: state (dict): The current state dictionary. Expected to contain 'datasets' with data_key. data_key (str): The key in state['datasets'] corresponding to the DataFrame to be used for training. target_column (str): The name of the target variable column in the DataFrame. model_name_for_registry (str, optional): If provided, the model will be registered under this name in the MLflow Model Registry. Defaults to None (model not registered). test_size (float): The proportion of the dataset to include in the test split. random_state (int): Seed for random number generation for reproducibility.

Returns: (dict): The updated state dictionary, including the test data components (X_test, y_test). The trained model is now managed by MLflow.

[16]

def train_ml_model(state: dict, data_key: str, target_column: str,
                   model_name_for_registry: str = None,
                   test_size: float = 0.2, random_state: int = 42) -> Tuple[dict, str]:
    """
    Trains a machine learning model (Logistic Regression), logs its parameters and metrics
    using MLflow, and optionally registers it in the MLflow Model Registry.

    Parameters
    ----------
    state : dict
        Current state dictionary. Expected to contain 'datasets' with `data_key`.
    data_key : str
        Key for the dataset in `state['datasets']`.
    target_column : str
        Name of the target column.
    model_name_for_registry : str, optional
        If provided, the model will be registered under this name in the MLflow Model Registry.
        Defaults to None (model not registered).
    test_size : float, optional
        Proportion of the dataset to include in the test split, defaults to 0.2.
    random_state : int, optional
        Random seed for reproducibility, defaults to 42.

    Returns
    -------
    Tuple[dict, str]
        Updated state with 'X_test' and 'y_test' added, and the MLflow run_id.
    """
    logger.info(f"Starting MLflow run for training model using data from '{data_key}'...")

    with mlflow.start_run() as run:
        run_id = run.info.run_id
        logger.info(f"MLflow Run ID: {run_id}")

        df = state['datasets'].get(data_key)
        if df is None:
            logger.error(f"Dataset '{data_key}' not found in state.")
            return state, None # Return None for run_id if data not found

        X = df.drop(columns=[target_column])
        y = df[target_column]

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

        model = LogisticRegression(random_state=random_state, solver='liblinear')
        model.fit(X_train, y_train)

        # Log parameters with MLflow
        mlflow.log_param("random_state", random_state)
        mlflow.log_param("test_size", test_size)
        mlflow.log_param("solver", 'liblinear')
        mlflow.log_param("n_samples", len(df))
        mlflow.log_param("n_features", X.shape[1])

        # Store test data in state for later evaluation (outside MLflow run context if needed)
        state['X_test'] = X_test
        state['y_test'] = y_test
        state['trained_model_current_run'] = model # Temporarily store model for immediate evaluation

        # Evaluate model
        metrics = evaluate_ml_model(state)

        # Log metrics with MLflow
        mlflow.log_metrics(metrics)

        # Log and optionally register the model with MLflow
        if model_name_for_registry:
            mlflow.sklearn.log_model(
                sk_model=model,
                artifact_path="logistic_regression_model",
                registered_model_name=model_name_for_registry
            )
            logger.info(f"Model '{model_name_for_registry}' logged and registered with MLflow. Run ID: {run_id}")
        else:
            mlflow.sklearn.log_model(
                sk_model=model,
                artifact_path="logistic_regression_model"
            )
            logger.info(f"Model logged to MLflow as artifact. Run ID: {run_id}")

        logger.info("Model training complete and logged to MLflow.")
        return state, run_id

Function Name: `evaluate_ml_model`

This function evaluates a trained machine learning model using a set of common classification metrics: accuracy, precision, recall, and F1-score. It takes the model and test data from the state dictionary, makes predictions, and calculates these metrics. The results are returned as a dictionary, which can then be logged by MLflow.

Parameters: state (dict): The current state dictionary containing the trained model (trained_model_current_run), X_test, and y_test.

Returns: (dict): A dictionary containing the calculated metrics: accuracy, precision, recall, and f1.

[5]

def evaluate_ml_model(state: dict) -> dict:
    """
    Evaluates a trained ML model and returns its performance metrics.

    Parameters
    ----------
    state : dict
        Current state dictionary containing 'trained_model_current_run', 'X_test', and 'y_test'.

    Returns
    -------
    dict
        A dictionary of performance metrics (accuracy, precision, recall, f1).
    """
    logger.info("Evaluating ML model...")
    model = state.get('trained_model_current_run') # Get model from current run
    X_test = state.get('X_test')
    y_test = state.get('y_test')

    if model is None or X_test is None or y_test is None:
        logger.error("Model or test data not found in state. Cannot evaluate.")
        return {}

    y_pred = model.predict(X_test)

    metrics = {
        'accuracy': accuracy_score(y_test, y_pred),
        'precision': precision_score(y_test, y_pred),
        'recall': recall_score(y_test, y_pred),
        'f1': f1_score(y_test, y_pred)
    }
    logger.info(f"Model evaluation complete. Metrics: {metrics}")
    return metrics

Function Name: `simulate_api_call_with_backoff`

This function simulates an API call with built-in exponential backoff and random jitter. It's designed to handle transient errors by retrying failed attempts with increasing delays. The jitter helps prevent thundering herd problems when many clients retry simultaneously. This pattern is crucial for robust production systems.

Parameters: state (dict): The current state dictionary (unused but kept for pattern consistency). attempt (int): The current retry attempt number. Defaults to 1. max_attempts (int): The maximum number of retry attempts. Defaults to 5.

Returns: (bool): True if the simulated call succeeds within max_attempts, False otherwise.

[ ]

def simulate_api_call_with_backoff(state: dict, attempt: int = 1, max_attempts: int = 5) -> bool:
    """
    Simulates an API call with exponential backoff and random jitter.

    Parameters
    ----------
    state : dict
        Current state dictionary (kept for consistency with state-passing pattern).
    attempt : int, optional
        The current retry attempt number, defaults to 1.
    max_attempts : int, optional
        The maximum number of retry attempts, defaults to 5.

    Returns
    -------
    bool
        True if the simulated call succeeds, False otherwise.
    """
    if attempt > max_attempts:
        logger.error(f"Max attempts ({max_attempts}) reached. API call failed.")
        return False

    try:
        # Simulate a flaky API call that fails randomly
        if random.random() < 0.6: # 60% chance of failure initially
            raise ConnectionRefusedError("Simulated API connection error")

        logger.info(f"API call successful on attempt {attempt}.")
        return True
    except ConnectionRefusedError as e:
        wait_time = (2 ** (attempt - 1)) + random.uniform(0, 1) * 0.1 # Exponential backoff with jitter
        logger.warning(f"API call failed (Attempt {attempt}/{max_attempts}): {e}. Retrying in {wait_time:.2f} seconds...")
        time.sleep(wait_time)
        return simulate_api_call_with_backoff(state, attempt + 1, max_attempts)
    except Exception as e:
        logger.error(f"An unexpected error occurred: {e}")
        return False

Function Name: `summarize_model_metrics`

This function takes a dictionary of model metrics and converts it into a pandas DataFrame for structured display and easy comparison. This helps in quickly visualizing and understanding the performance characteristics of different model versions.

Parameters: metrics_data (dict): A dictionary where keys are model version IDs and values are dictionaries of metrics.

Returns: (pd.DataFrame): A DataFrame summarizing the model metrics, or an empty DataFrame if no data is provided.

[ ]

def summarize_model_metrics(metrics_data: dict) -> pd.DataFrame:
    """
    Converts a dictionary of model metrics into a pandas DataFrame for summary.

    Parameters
    ----------
    metrics_data : dict
        A dictionary where keys are model version IDs and values are dictionaries of metrics.

    Returns
    -------
    pd.DataFrame
        A DataFrame summarizing the model metrics.
    """
    logger.info("Summarizing model metrics...")
    if not metrics_data:
        logger.warning("No metrics data provided for summary.")
        return pd.DataFrame()

    df = pd.DataFrame.from_dict(metrics_data, orient='index')
    df.index.name = 'Version ID'
    logger.info("Model metrics summary created.")
    return df

Demonstration/Visualization

This section demonstrates the full workflow of generating data, training models, evaluating them, and then versioning and registering these models. We'll also visualize the performance across different versions and showcase robust production practices.

Step 1: Initialize the ML System State

We start with an empty state dictionary to hold our datasets. The model registry will now be handled by MLflow directly, so there's no need for a custom model_registry key in our state.

[6]

# Create initial state
ml_state = {"datasets": {}}

# Display the initial state (should be empty datasets)
print("\nInitial ML System State:")
display(ml_state)


Initial ML System State:

{'datasets': {}}

Step 2: Generate Mock Data

We'll generate a synthetic classification dataset to train our models. This allows us to simulate the process without needing real-world data sources.

[ ]

# Generate mock data
ml_state = generate_mock_classification_data(ml_state, n_samples=1500, n_features=15, random_state=100)

# Display first 5 rows of the generated data
print("\nFirst 5 rows of generated mock data:")
display(ml_state['datasets']['mock_data'].head())

[32m2026-06-10T07:18:49.616809+0000[0m [1mINFO[0m [31mGenerating 1500 samples with 15 features...[0m
[32m2026-06-10T07:18:49.640520+0000[0m [1mINFO[0m [31mMock data generated with shape (1500, 16).[0m


First 5 rows of generated mock data:

	feature_0	feature_1	feature_2	feature_3	feature_4	feature_5	feature_6	feature_7	feature_8	feature_9	feature_10	feature_11	feature_12	feature_13	feature_14	target
0	0.543405	0.278369	0.424518	0.844776	0.004719	0.121569	0.670749	0.825853	0.136707	0.575093	0.891322	0.209202	0.185328	0.108377	0.219697	0
1	0.978624	0.811683	0.171941	0.816225	0.274074	0.431704	0.940030	0.817649	0.336112	0.175410	0.372832	0.005689	0.252426	0.795663	0.015255	1
2	0.598843	0.603805	0.105148	0.381943	0.036476	0.890412	0.980921	0.059942	0.890546	0.576901	0.742480	0.630184	0.581842	0.020439	0.210027	0
3	0.544685	0.769115	0.250695	0.285896	0.852395	0.975006	0.884853	0.359508	0.598859	0.354796	0.340190	0.178081	0.237694	0.044862	0.505431	1
4	0.376252	0.592805	0.629942	0.142600	0.933841	0.946380	0.602297	0.387766	0.363188	0.204345	0.276765	0.246536	0.173608	0.966610	0.957013	0

Step 3: Train, Evaluate, and Register Multiple Model Versions with MLflow

We will simulate training several versions of our LogisticRegressionClassifier. For each run, MLflow will log the parameters, metrics, and the model itself. We will also demonstrate how to register these models in the MLflow Model Registry, which is essential for managing model lifecycle stages (e.g., Staging, Production).

[17]

model_registry_name = "LogisticRegressionClassifier"
all_version_metrics = {}

# Set MLflow tracking URI to a local folder for this demonstration
mlflow.set_tracking_uri("file:///tmp/mlruns")
mlflow.set_experiment("Model_Versioning_Demo")

# Simulate multiple training runs and registrations
for i in range(3):
    logger.info(f"\n--- Training Run {i+1} ---")
    # Simulate slight data variations or hyperparameter changes by changing random_state
    current_random_state = 42 + (i * 10)

    # Generate mock data (potentially with slight variations)
    ml_state = generate_mock_classification_data(ml_state, n_samples=1000 + (i*100), n_features=10, random_state=current_random_state)

    # Train, log, and register the model using MLflow integrated function
    # The model will be automatically logged and registered if model_registry_name is provided
    ml_state, current_run_id = train_ml_model(ml_state, data_key='mock_data', target_column='target',
                              model_name_for_registry=model_registry_name,
                              random_state=current_random_state)

    # Retrieve metrics from the evaluation done within train_ml_model's MLflow run
    # For this example, we'll manually fetch the metrics from the last run
    # In a real scenario, you'd query MLflow for metrics associated with the registered model version
    metrics = evaluate_ml_model(ml_state) # Evaluate again to get the metrics in a dict

    # Store metrics for visualization, associating them with an identifier (e.g., MLflow Run ID)
    all_version_metrics[current_run_id] = metrics

print("\nAll models logged and registered with MLflow. Check the MLflow UI for details:")
print(f"mlflow ui --backend-store-uri file:///tmp/mlruns")

[32m2026-06-12T11:33:11.830188+0000[0m [1mINFO[0m [31m
--- Training Run 1 ---[0m
[32m2026-06-12T11:33:11.830998+0000[0m [1mINFO[0m [31mGenerating 1000 samples with 10 features...[0m
[32m2026-06-12T11:33:11.833492+0000[0m [1mINFO[0m [31mMock data generated with shape (1000, 11).[0m
[32m2026-06-12T11:33:11.834089+0000[0m [1mINFO[0m [31mStarting MLflow run for training model using data from 'mock_data'...[0m
[32m2026-06-12T11:33:11.845333+0000[0m [1mINFO[0m [31mMLflow Run ID: f524c25ee47d4672b32a849837821a70[0m
[32m2026-06-12T11:33:11.856939+0000[0m [1mINFO[0m [31mEvaluating ML model...[0m
[32m2026-06-12T11:33:11.870086+0000[0m [1mINFO[0m [31mModel evaluation complete. Metrics: {'accuracy': 0.905, 'precision': 0.8085106382978723, 'recall': 0.7916666666666666, 'f1': 0.8}[0m
2026/06/12 11:33:11 WARNING mlflow.models.model: `artifact_path` is deprecated. Please use `name` instead.
2026/06/12 11:33:11 WARNING mlflow.sklearn: Saving scikit-learn models in the pickle or cloudpickle format requires exercising caution because these formats rely on Python's object serialization mechanism, which can execute arbitrary code during deserialization. The recommended safe alternative is the 'skops' format. For more information, see: https://scikit-learn.org/stable/model_persistence.html
Registered model 'LogisticRegressionClassifier' already exists. Creating a new version of this model...
Created version '2' of model 'LogisticRegressionClassifier'.
[32m2026-06-12T11:33:15.749630+0000[0m [1mINFO[0m [31mModel 'LogisticRegressionClassifier' logged and registered with MLflow. Run ID: f524c25ee47d4672b32a849837821a70[0m
[32m2026-06-12T11:33:15.750440+0000[0m [1mINFO[0m [31mModel training complete and logged to MLflow.[0m
[32m2026-06-12T11:33:15.752794+0000[0m [1mINFO[0m [31mEvaluating ML model...[0m
[32m2026-06-12T11:33:15.766476+0000[0m [1mINFO[0m [31mModel evaluation complete. Metrics: {'accuracy': 0.905, 'precision': 0.8085106382978723, 'recall': 0.7916666666666666, 'f1': 0.8}[0m
[32m2026-06-12T11:33:15.767127+0000[0m [1mINFO[0m [31m
--- Training Run 2 ---[0m
[32m2026-06-12T11:33:15.767906+0000[0m [1mINFO[0m [31mGenerating 1100 samples with 10 features...[0m
[32m2026-06-12T11:33:15.769729+0000[0m [1mINFO[0m [31mMock data generated with shape (1100, 11).[0m
[32m2026-06-12T11:33:15.770927+0000[0m [1mINFO[0m [31mStarting MLflow run for training model using data from 'mock_data'...[0m
[32m2026-06-12T11:33:15.781146+0000[0m [1mINFO[0m [31mMLflow Run ID: f4b71e64dd784db5933f80a668254c33[0m
[32m2026-06-12T11:33:15.793495+0000[0m [1mINFO[0m [31mEvaluating ML model...[0m
[32m2026-06-12T11:33:15.807582+0000[0m [1mINFO[0m [31mModel evaluation complete. Metrics: {'accuracy': 0.8863636363636364, 'precision': 0.7966101694915254, 'recall': 0.7833333333333333, 'f1': 0.7899159663865546}[0m
2026/06/12 11:33:15 WARNING mlflow.models.model: `artifact_path` is deprecated. Please use `name` instead.
2026/06/12 11:33:15 WARNING mlflow.sklearn: Saving scikit-learn models in the pickle or cloudpickle format requires exercising caution because these formats rely on Python's object serialization mechanism, which can execute arbitrary code during deserialization. The recommended safe alternative is the 'skops' format. For more information, see: https://scikit-learn.org/stable/model_persistence.html
Registered model 'LogisticRegressionClassifier' already exists. Creating a new version of this model...
Created version '3' of model 'LogisticRegressionClassifier'.
[32m2026-06-12T11:33:20.755064+0000[0m [1mINFO[0m [31mModel 'LogisticRegressionClassifier' logged and registered with MLflow. Run ID: f4b71e64dd784db5933f80a668254c33[0m
[32m2026-06-12T11:33:20.755895+0000[0m [1mINFO[0m [31mModel training complete and logged to MLflow.[0m
[32m2026-06-12T11:33:20.759603+0000[0m [1mINFO[0m [31mEvaluating ML model...[0m
[32m2026-06-12T11:33:20.770614+0000[0m [1mINFO[0m [31mModel evaluation complete. Metrics: {'accuracy': 0.8863636363636364, 'precision': 0.7966101694915254, 'recall': 0.7833333333333333, 'f1': 0.7899159663865546}[0m
[32m2026-06-12T11:33:20.771435+0000[0m [1mINFO[0m [31m
--- Training Run 3 ---[0m
[32m2026-06-12T11:33:20.772098+0000[0m [1mINFO[0m [31mGenerating 1200 samples with 10 features...[0m
[32m2026-06-12T11:33:20.774962+0000[0m [1mINFO[0m [31mMock data generated with shape (1200, 11).[0m
[32m2026-06-12T11:33:20.775530+0000[0m [1mINFO[0m [31mStarting MLflow run for training model using data from 'mock_data'...[0m
[32m2026-06-12T11:33:20.786967+0000[0m [1mINFO[0m [31mMLflow Run ID: 17ac940f614249928d4aac079c8e8736[0m
[32m2026-06-12T11:33:20.803301+0000[0m [1mINFO[0m [31mEvaluating ML model...[0m
[32m2026-06-12T11:33:20.816305+0000[0m [1mINFO[0m [31mModel evaluation complete. Metrics: {'accuracy': 0.9041666666666667, 'precision': 0.819672131147541, 'recall': 0.8064516129032258, 'f1': 0.8130081300813008}[0m
2026/06/12 11:33:20 WARNING mlflow.models.model: `artifact_path` is deprecated. Please use `name` instead.
2026/06/12 11:33:20 WARNING mlflow.sklearn: Saving scikit-learn models in the pickle or cloudpickle format requires exercising caution because these formats rely on Python's object serialization mechanism, which can execute arbitrary code during deserialization. The recommended safe alternative is the 'skops' format. For more information, see: https://scikit-learn.org/stable/model_persistence.html
Registered model 'LogisticRegressionClassifier' already exists. Creating a new version of this model...
Created version '4' of model 'LogisticRegressionClassifier'.
[32m2026-06-12T11:33:24.728377+0000[0m [1mINFO[0m [31mModel 'LogisticRegressionClassifier' logged and registered with MLflow. Run ID: 17ac940f614249928d4aac079c8e8736[0m
[32m2026-06-12T11:33:24.729148+0000[0m [1mINFO[0m [31mModel training complete and logged to MLflow.[0m
[32m2026-06-12T11:33:24.731389+0000[0m [1mINFO[0m [31mEvaluating ML model...[0m
[32m2026-06-12T11:33:24.743167+0000[0m [1mINFO[0m [31mModel evaluation complete. Metrics: {'accuracy': 0.9041666666666667, 'precision': 0.819672131147541, 'recall': 0.8064516129032258, 'f1': 0.8130081300813008}[0m


All models logged and registered with MLflow. Check the MLflow UI for details:
mlflow ui --backend-store-uri file:///tmp/mlruns

Step 4: Visualize Model Performance Across Versions

To understand how different versions perform, we'll create a bar chart comparing the accuracy, precision, recall, and F1-score of each registered model version.

[ ]

# Summarize metrics into a DataFrame
metrics_df = summarize_model_metrics(all_version_metrics)
print("\nModel Performance Summary:")
display(metrics_df)

# Plotting the metrics
metrics_df_melted = metrics_df.reset_index().melt(id_vars='Version ID', var_name='Metric', value_name='Score')

plt.figure(figsize=(14, 7))
sns.barplot(x='Version ID', y='Score', hue='Metric', data=metrics_df_melted, palette='viridis')
plt.title('Model Performance Across Different Versions')
plt.xlabel('Model Version ID')
plt.ylabel('Score')
plt.xticks(rotation=45, ha='right')
plt.legend(title='Metric')
plt.grid(axis='y', linestyle='--', alpha=0.7)
plt.tight_layout()
plt.show()

[32m2026-06-10T07:18:51.837419+0000[0m [1mINFO[0m [31mSummarizing model metrics...[0m
[32m2026-06-10T07:18:51.843266+0000[0m [1mINFO[0m [31mModel metrics summary created.[0m


Model Performance Summary:

	accuracy	precision	recall	f1
Version ID
LogisticRegressionClassifier-20260610071850-7940a3	0.905000	0.808511	0.791667	0.800000
LogisticRegressionClassifier-20260610071850-b7cbe0	0.886364	0.796610	0.783333	0.789916
LogisticRegressionClassifier-20260610071850-9ed877	0.904167	0.819672	0.806452	0.813008

Step 5: Retrieve a Specific Model Version and Promote to Staging/Production

MLflow's Model Registry allows us to retrieve models by version number or by assigned stage (e.g., 'Staging', 'Production'). We will demonstrate how to load a model and promote a specific version to the 'Staging' stage, and then to 'Production'.

[14]

model_registry_name = "LogisticRegressionClassifier"

# Get all registered versions for the model
client = mlflow.tracking.MlflowClient()
versions = client.search_model_versions(f"name='{model_registry_name}'")

if versions:
    # Retrieve the first registered version (example)
    first_version = versions[0]
    print(f"\nRetrieving Specific Model Version '{first_version.version}' (Run ID: {first_version.run_id}):")
    model_uri_specific = f"models:/{model_registry_name}/{first_version.version}"
    loaded_model_specific = mlflow.pyfunc.load_model(model_uri_specific)
    print(f"  Model Type: {type(loaded_model_specific)}")
    print(f"  Model URI: {model_uri_specific}")

    # Get the latest registered version
    latest_version = client.get_latest_versions(model_registry_name, stages=["None"])[0]
    print(f"\nRetrieving Latest Model Version '{latest_version.version}' (Run ID: {latest_version.run_id}):")
    model_uri_latest = f"models:/{model_registry_name}/{latest_version.version}"
    loaded_model_latest = mlflow.pyfunc.load_model(model_uri_latest)
    print(f"  Model Type: {type(loaded_model_latest)}")
    print(f"  Model URI: {model_uri_latest}")

    # Promote the latest version to 'Staging'
    client.transition_model_version_stage(
        name=model_registry_name,
        version=latest_version.version,
        stage="Staging"
    )
    print(f"\nModel version {latest_version.version} of '{model_registry_name}' transitioned to 'Staging'.")

    # Retrieve the model currently in 'Staging'
    staged_model_uri = f"models:/{model_registry_name}/Staging"
    loaded_staged_model = mlflow.pyfunc.load_model(staged_model_uri)
    print(f"\nRetrieved Model from 'Staging' stage (URI: {staged_model_uri}):")
    print(f"  Model Type: {type(loaded_staged_model)}")

    # Simulate promoting to Production
    client.transition_model_version_stage(
        name=model_registry_name,
        version=latest_version.version,
        stage="Production"
    )
    print(f"\nModel version {latest_version.version} of '{model_registry_name}' transitioned to 'Production'.")

    # Retrieve the model currently in 'Production'
    production_model_uri = f"models:/{model_registry_name}/Production"
    loaded_production_model = mlflow.pyfunc.load_model(production_model_uri)
    print(f"\nRetrieved Model from 'Production' stage (URI: {production_model_uri}):")
    print(f"  Model Type: {type(loaded_production_model)}")

else:
    logger.warning(f"No models registered under name '{model_registry_name}' to retrieve.")


Retrieving Specific Model Version '1' (Run ID: 5d3ff1847bf04862a8c6bcf4be4c0a4f):
  Model Type: <class 'mlflow.pyfunc.PyFuncModel'>
  Model URI: models:/LogisticRegressionClassifier/1

Retrieving Latest Model Version '1' (Run ID: 5d3ff1847bf04862a8c6bcf4be4c0a4f):
  Model Type: <class 'mlflow.pyfunc.PyFuncModel'>
  Model URI: models:/LogisticRegressionClassifier/1

/tmp/ipykernel_4637/3711346798.py:17: FutureWarning: ``mlflow.tracking.client.MlflowClient.get_latest_versions`` is deprecated since 2.9.0. Model registry stages will be removed in a future major release. To learn more about the deprecation of model registry stages, see our migration guide here: https://mlflow.org/docs/latest/model-registry.html#migrating-from-stages
  latest_version = client.get_latest_versions(model_registry_name, stages=["None"])[0]
/tmp/ipykernel_4637/3711346798.py:25: FutureWarning: ``mlflow.tracking.client.MlflowClient.transition_model_version_stage`` is deprecated since 2.9.0. Model registry stages will be removed in a future major release. To learn more about the deprecation of model registry stages, see our migration guide here: https://mlflow.org/docs/latest/model-registry.html#migrating-from-stages
  client.transition_model_version_stage(
/tmp/ipykernel_4637/3711346798.py:39: FutureWarning: ``mlflow.tracking.client.MlflowClient.transition_model_version_stage`` is deprecated since 2.9.0. Model registry stages will be removed in a future major release. To learn more about the deprecation of model registry stages, see our migration guide here: https://mlflow.org/docs/latest/model-registry.html#migrating-from-stages
  client.transition_model_version_stage(


Model version 1 of 'LogisticRegressionClassifier' transitioned to 'Staging'.

Retrieved Model from 'Staging' stage (URI: models:/LogisticRegressionClassifier/Staging):
  Model Type: <class 'mlflow.pyfunc.PyFuncModel'>

Model version 1 of 'LogisticRegressionClassifier' transitioned to 'Production'.

Retrieved Model from 'Production' stage (URI: models:/LogisticRegressionClassifier/Production):
  Model Type: <class 'mlflow.pyfunc.PyFuncModel'>

Step 6: Demonstration of Robust API Call with Exponential Backoff

In production environments, external service calls can be flaky. This demonstration shows how to implement retries with exponential backoff and random jitter, making your system more resilient.

[ ]

print("\n--- Demonstrating API Call with Exponential Backoff ---")
success = simulate_api_call_with_backoff(ml_state, max_attempts=7)

if success:
    logger.info("API call simulation finished successfully.")
else:
    logger.error("API call simulation failed after multiple retries.")

[32m2026-06-10T07:18:54.680107+0000[0m [33m[1mWARNING[0m [31mAPI call failed (Attempt 1/7): Simulated API connection error. Retrying in 1.06 seconds...[0m


--- Demonstrating API Call with Exponential Backoff ---

[32m2026-06-10T07:18:55.761276+0000[0m [33m[1mWARNING[0m [31mAPI call failed (Attempt 2/7): Simulated API connection error. Retrying in 2.04 seconds...[0m
[32m2026-06-10T07:18:57.803456+0000[0m [33m[1mWARNING[0m [31mAPI call failed (Attempt 3/7): Simulated API connection error. Retrying in 4.06 seconds...[0m
[32m2026-06-10T07:19:01.863372+0000[0m [1mINFO[0m [31mAPI call successful on attempt 4.[0m
[32m2026-06-10T07:19:01.864901+0000[0m [1mINFO[0m [31mAPI call simulation finished successfully.[0m

Production Considerations

Implementing model versioning and registration is a step towards robust MLOps. Here are key production considerations to ensure a reliable and scalable ML system:

Consideration	Best Practice
CI/CD Integration	Automate model training, evaluation, versioning, and deployment pipelines.
Monitoring	Continuously monitor model performance (drift, bias, accuracy) in production. Set up alerts for performance degradation.
A/B Testing	Deploy new model versions alongside old ones to evaluate real-world impact before full rollout.
Rollbacks	Ensure the ability to quickly revert to a previous, stable model version if issues arise in production.
Security & Access Control	Implement strict access controls for model registry and deployment endpoints. Encrypt sensitive model artifacts.
Scalability	Design the model registry and serving infrastructure to handle increased load and a growing number of models and versions.
Model Documentation	Maintain comprehensive documentation for each model version, including data sources, preprocessing steps, algorithms, and evaluation results.

Conclusion

This notebook demonstrated the fundamental principles and implementation of ML model versioning and registration. By centralizing model management, assigning unique versions, and capturing rich metadata, we enable better reproducibility, traceability, and a more streamlined path to production for machine learning models. The included functions and visualizations illustrate how to build a basic yet effective system for managing the lifecycle of your ML models, laying the groundwork for more advanced MLOps practices.

Model Versioning Mlflow

Model Versioning and Registration with MLflow

Concepts Overview

Dependency Installation

Library Imports

Core Functions

Function Name: generate_mock_classification_data

Function Name: generate_mock_classification_data

Function Name: train_ml_model

Function Name: evaluate_ml_model

Function Name: simulate_api_call_with_backoff

Function Name: summarize_model_metrics

Demonstration/Visualization

Step 1: Initialize the ML System State

Step 2: Generate Mock Data

Step 3: Train, Evaluate, and Register Multiple Model Versions with MLflow

Step 4: Visualize Model Performance Across Versions

Step 5: Retrieve a Specific Model Version and Promote to Staging/Production

Step 6: Demonstration of Robust API Call with Exponential Backoff

Production Considerations

Conclusion

Function Name: `generate_mock_classification_data`

Function Name: `generate_mock_classification_data`

Function Name: `train_ml_model`

Function Name: `evaluate_ml_model`

Function Name: `simulate_api_call_with_backoff`

Function Name: `summarize_model_metrics`