MLflow in Python: Track Experiments and Version Models

Compare three Random Forest configurations on Iris, save their parameters and validation metrics, and retrieve the selected model from MLflow. The example records a separate final test and registers a version under the champion alias. Its outputs are local files and registry records; serving predictions to users requires a separate application.

Local setup and storage

The program was checked with Python 3.13.5 on Windows and the following direct dependencies. Create a project folder and install them in a virtual environment:

python -m venv .venv
.\.venv\Scripts\python.exe -m pip install `
    mlflow==3.14.0 scikit-learn==1.9.0 pandas==2.3.3 `
    numpy==2.4.6 skops==0.14.0

The commands use the Windows environment's Python executable. On Linux or macOS, use .venv/bin/python in the same commands. The validation used an existing environment with these versions, rather than a fresh installation or a complete transitive lockfile.

The script explicitly selects a local SQLite database for tracking and registry metadata. In this setup, mlflow.db holds the records and mlruns/ holds model and CSV artifacts. Keep both when reopening the experiment. Current MLflow defaults to a database backend; the old file-only metadata store is a legacy option. Backend store documentation.

Copy the four Python sections below into train.py, in order. Run the file from the project folder:

.\.venv\Scripts\python.exe train.py

Record three comparable runs

The stratified split reserves 30 rows for final testing, then divides the remaining rows into 90 for training and 30 for validation. Each forest learns only from training rows. Validation metrics select a candidate; the final test remains outside that decision.

batch_id identifies the three runs created by this execution. It prevents the later query from selecting an older run from another experiment batch. Each run stores a feature-importance CSV and a logged model with an input example and inferred signature. Explicit model requirements record the tested direct dependencies and avoid scanning the environment to infer them.

from pathlib import Path
from uuid import uuid4

import mlflow
import mlflow.sklearn
import pandas as pd
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score
from sklearn.model_selection import train_test_split

REQUIREMENTS = [
    "mlflow==3.14.0",
    "scikit-learn==1.9.0",
    "pandas==2.3.3",
    "numpy==2.4.6",
    "skops==0.14.0",
]
tracking_uri = "sqlite:///" + Path("mlflow.db").resolve().as_posix()
mlflow.set_tracking_uri(tracking_uri)
mlflow.set_registry_uri(tracking_uri)
mlflow.set_experiment("iris-random-forest")
batch_id = uuid4().hex

iris = load_iris()
X = pd.DataFrame(iris.data, columns=iris.feature_names)
y = iris.target
X_dev, X_test, y_dev, y_test = train_test_split(
    X, y, test_size=0.2, stratify=y, random_state=42
)
X_train, X_validation, y_train, y_validation = train_test_split(
    X_dev, y_dev, test_size=0.25, stratify=y_dev, random_state=42
)

configs = [
    {"n_estimators": 50, "max_depth": 3, "min_samples_split": 2},
    {"n_estimators": 100, "max_depth": 5, "min_samples_split": 5},
    {"n_estimators": 200, "max_depth": 10, "min_samples_split": 2},
]
model_uris = {}
for config in configs:
    with mlflow.start_run(tags={"batch_id": batch_id}) as run:
        mlflow.log_params({
            **config,
            "random_state": 42,
            "training_rows": len(X_train),
            "validation_rows": len(X_validation),
            "test_rows": len(X_test),
        })
        model = RandomForestClassifier(**config, random_state=42, n_jobs=1)
        model.fit(X_train, y_train)
        prediction = model.predict(X_validation)
        mlflow.log_metrics({
            "validation_accuracy": accuracy_score(y_validation, prediction),
            "validation_f1": f1_score(y_validation, prediction, average="weighted"),
            "validation_precision": precision_score(
                y_validation, prediction, average="weighted"
            ),
            "validation_recall": recall_score(y_validation, prediction, average="weighted"),
        })
        importance = pd.DataFrame({
            "feature": iris.feature_names,
            "importance": model.feature_importances_,
        }).sort_values("importance", ascending=False)
        importance.to_csv("feature_importance.csv", index=False)
        mlflow.log_artifact("feature_importance.csv")
        info = mlflow.sklearn.log_model(
            model,
            name="random_forest_model",
            input_example=X_train.iloc[:3],
            pip_requirements=REQUIREMENTS,
            serialization_format="skops",
        )
        model_uris[run.info.run_id] = info.model_uri

The saved model uses the skops format, supported by the pinned MLflow version. This example loads artifacts it created itself. Review provenance and any required trusted types before loading an external model; keep training and loading dependencies compatible. A signature records the expected data structure, rather than establishing that future inputs are meaningful or safe for a production service. scikit-learn model flavor.

Select using validation metrics

Query the current batch and sort by validation accuracy. Ties use the most recent start time as an explicit tie-breaker; this does not establish that one tied model generalizes better. The selected URI comes from the actual ModelInfo returned during logging, instead of reconstructing an artifact path.

from mlflow.tracking import MlflowClient

client = MlflowClient()
experiment = client.get_experiment_by_name("iris-random-forest")
if experiment is None:
    raise RuntimeError("Run the training section first.")
runs = client.search_runs(
    experiment_ids=[experiment.experiment_id],
    filter_string=f"tags.batch_id = '{batch_id}'",
    order_by=["metrics.validation_accuracy DESC", "attributes.start_time DESC"],
    max_results=1,
)
if not runs:
    raise RuntimeError("No runs found for this training batch.")
best_run = runs[0]
model_uri = model_uris[best_run.info.run_id]
print("Selected parameters:", best_run.data.params)
print("Validation accuracy:", best_run.data.metrics["validation_accuracy"])

Validation scores are part of model selection. Treating the largest of them as an unbiased final score would reuse the same observations for selection and evaluation.

Load and evaluate once on the final test

Load the selected artifact, compute its test accuracy and attach that metric only to the selected run. The two example flowers show how to pass named columns to the recovered model; the program prints its actual species predictions.

loaded_model = mlflow.sklearn.load_model(model_uri)
test_prediction = loaded_model.predict(X_test)
test_accuracy = accuracy_score(y_test, test_prediction)
with mlflow.start_run(run_id=best_run.info.run_id):
    mlflow.log_metric("test_accuracy", test_accuracy)
print("Final test accuracy:", test_accuracy)

new_data = pd.DataFrame({
    "sepal length (cm)": [5.1, 6.7],
    "sepal width (cm)": [3.5, 3.0],
    "petal length (cm)": [1.4, 5.2],
    "petal width (cm)": [0.2, 2.3],
})
predictions = loaded_model.predict(new_data)
print("Predicted species:", iris.target_names[predictions].tolist())

This small dataset and one partition illustrate the workflow. They do not establish production quality or a reliable performance ranking. Changing the models after inspecting test results would turn those test observations into further selection data. Preserve an external evaluation set or use an appropriate outer evaluation when iterating.

Register a version and assign an alias

Registration creates a version in IrisClassifier; champion is a mutable reference to that version. A later assignment changes which model a subsequent alias lookup retrieves. Model stages have been deprecated since MLflow 2.9.0; aliases and tags are the recommended workflow, rather than a stage system first replaced in version 3. Registry workflow and migration.

registered = mlflow.register_model(model_uri=model_uri, name="IrisClassifier")
client.set_registered_model_alias(
    name="IrisClassifier", alias="champion", version=registered.version
)
champion_model = mlflow.sklearn.load_model("models:/IrisClassifier@champion")
print("Registered version:", registered.version)

The registry stores the version and its relationship to the source run. An application still needs to load that version and expose inference, with its own access controls and operational checks. Assigning champion does not start that application.

Inspect the local experiment

From the same project folder, start the UI against the database used by the script:

.\.venv\Scripts\python.exe -m mlflow server `
    --backend-store-uri sqlite:///mlflow.db `
    --host 127.0.0.1 --port 5000

Open http://127.0.0.1:5000, find iris-random-forest, and inspect the three runs: parameters, validation_* metrics, CSV artifacts and logged models. Only the selected run has test_accuracy. The model registry shows the version and alias. The server binds to your own machine; this is a local inspection setup.

The example records the inputs described above, not a complete reproducibility or governance system. Teams also need data provenance, source and environment versions, evaluation policies and reviewed deployment decisions. For the next steps, see scheduled pipelines with Prefect and actual inference with FastAPI and Docker.