5.5. AI/ML Experiments

What is an AI/ML experiment?

An AI/ML experiment is a systematic and iterative process for building robust machine learning models. It involves testing different algorithms, tuning hyperparameters, and using various datasets to discover the optimal configuration for a specific predictive task. Each experiment is a structured trial designed to measure the impact of changes on model performance, such as accuracy, efficiency, and reliability.

Why is experiment tracking essential in AI/ML?

In MLOps, the complexity and often non-deterministic nature of model development require a disciplined approach. Experiment tracking provides the necessary structure, much like a lab notebook for a scientist. Key benefits include:

Ensuring Reproducibility: Tracking guarantees that every aspect of an experiment—code, data, environment, and parameters—is recorded. This allows you and your team to reliably replicate and verify results.
Optimizing Hyperparameters: It provides a systematic way to test and compare different hyperparameter configurations, helping you pinpoint the settings that maximize model performance.
Organizing Your Work: By logging experiments and using tags, you can categorize runs by model type, dataset, or objective. This organization is crucial for managing complex projects and quickly retrieving past results.
Monitoring Performance: Tracking metrics during a run offers real-time insight into how adjustments affect model behavior, enabling faster, data-driven decisions.
Seamless Framework Integration: Modern tracking tools integrate with popular AI/ML frameworks, creating a unified and streamlined workflow across your entire toolchain.

Which experiment tracking solution should you use?

Numerous solutions are available for tracking AI/ML experiments. Major cloud platforms like Google Cloud (Vertex AI), Azure (Azure ML), and AWS (SageMaker) offer powerful, integrated MLOps capabilities. Specialized commercial tools like Weights & Biases, Comet, and ClearML also provide excellent features.

A note on vendor risk: this chapter used to recommend Neptune.ai here. OpenAI acquired it in December 2025 and it wound down its external service, taking its documentation with it. That is worth internalising before you build a workflow on a hosted tracker: an experiment store is where your history lives, so prefer one whose format you can export, and treat a vendor's continued existence as an assumption rather than a given. It is one of the reasons this course teaches MLflow, which you can run yourself.

For those starting out or preferring an open-source, framework-agnostic solution, MLflow is an outstanding choice. It is versatile, robust, and integrates with a wide array of ML libraries.

To install MLflow, run:

uv add mlflow

To verify the installation and start the MLflow UI server locally:

uv run mlflow doctor
uv run mlflow server --backend-store-uri=sqlite:///mlflow.db --artifacts-destination=./mlruns

In the MLOps Python Package, that command is wrapped as a task, so you never have to retype the flags:

mise run mlflow:serve

For a more permanent setup using Docker, you can use a docker-compose.yml file to launch the MLflow server on the same store:

services:
  mlflow:
    image: ghcr.io/mlflow/mlflow:v3.15.1
    ports:
      - 5000:5000
    volumes:
      - ./mlflow.db:/mlflow/mlflow.db
      - ./mlruns:/mlflow/mlruns
    working_dir: /mlflow
    # SQLAlchemy backend, the same store the package writes to locally.
    command: mlflow server --host 0.0.0.0 --port 5000 --backend-store-uri sqlite:///mlflow.db --artifacts-destination ./mlruns

Run docker compose up to start the service. For information on production-grade deployments, refer to the MLflow documentation.

Where should MLflow store your experiments?

MLflow splits storage into two halves, and they do not have to live in the same place:

The backend store holds the metadata: experiments, runs, parameters, metrics, tags, and every entry in the model registry. It can be a local directory (the file store) or any SQLAlchemy-compatible database (SQLite, PostgreSQL, MySQL).
The artifact store holds the files: serialized models, plots, data samples, and anything else you log as an artifact. It is a filesystem path or an object store such as S3 or GCS.

This course uses SQLite for the backend store and the local ./mlruns directory for artifacts:

sqlite:///mlflow.db   # metadata: experiments, runs, metrics, registered models
./mlruns              # artifact files: models, plots, datasets

Three reasons make this the right default, even on a laptop:

It is the shape you will run in production. A production MLflow deployment uses a SQLAlchemy backend, almost always PostgreSQL. Developing against SQLite means the same store type, the same schema, and the same query semantics. Moving up to PostgreSQL becomes a change of connection string (MLFLOW_TRACKING_URI), not a rewrite of how your code talks to MLflow.
The model registry needs a database. The registry is designed around a relational store; the file store was never built to back it. You will see this the moment you try to register a model, assign an alias, or search versions. The next section builds directly on top of this choice.
You are agreeing with upstream, not working around it. Since MLflow 3.14, sqlite:///mlflow.db is MLflow's own default tracking URI (DEFAULT_TRACKING_URI in mlflow/store/tracking/__init__.py). MLflow keeps a backward-compatibility check that falls back to ./mlruns only when it detects an existing file store on disk. Setting the URI explicitly documents the intent and removes the ambiguity; it does not fight the library.

!!! warning "The file store is a legacy path" Earlier versions of this course configured ./mlruns as the backend store. If you have an existing ./mlruns directory holding file-store runs, MLflow will silently keep using it instead of SQLite. Delete or move that directory when you migrate, and set the URIs explicitly so the behavior no longer depends on what happens to be on disk.

How do you configure MLflow in a project?

To integrate MLflow, you first need to point it at the backend store. Set the tracking and registry URIs to your SQLite database, then define an experiment name to group related runs.

Enabling MLflow's autologging is highly recommended. It automatically captures metrics, parameters, and models from popular ML libraries without requiring explicit logging statements.

import mlflow

# Metadata goes to the SQLite database; artifact files land under ./mlruns
mlflow.set_tracking_uri("sqlite:///mlflow.db")
mlflow.set_registry_uri("sqlite:///mlflow.db")

# Set a name for the experiment
mlflow.set_experiment(experiment_name="Bike Sharing Demand Prediction")

# Enable autologging for automatic tracking
mlflow.autolog()

In the MLOps Python Package these URIs are fields of a configuration object rather than loose calls, so every job reads the same defaults and any of them can be overridden from a YAML file or an environment variable:

class MlflowService(Service):
    """Service for Mlflow tracking and registry."""

    # SQLAlchemy backends are the supported store in MLflow 3: SQLite gives the local
    # setup the same shape as a production database (Postgres, MySQL) with no server to
    # run, and it is the only local store the model registry is actually designed for.
    tracking_uri: str = "sqlite:///mlflow.db"
    registry_uri: str = "sqlite:///mlflow.db"
    experiment_name: str = "bikes"
    registry_name: str = "bikes"

Remember to add both mlflow.db and mlruns/ to your .gitignore: experiment metadata and model artifacts are outputs, not source code.

To start a new run, wrap your training code within an MLflow run context. This allows you to add descriptive metadata and enable system metric logging.

with mlflow.start_run(
    run_name="Forecast Model with Feature Engineering",
    description="Training a model with an enhanced feature set.",
    log_system_metrics=True,
) as run:
    # Your model training and evaluation code goes here
    print(f"MLflow Run ID: {run.info.run_id}")

What information can you track in an experiment?

While MLflow's autologging captures a wealth of information automatically, you can enhance it with manual logging for greater detail:

Parameters: Log key-value parameters with mlflow.log_param() or a dictionary of parameters with mlflow.log_params().
Metrics: Record single metrics over time (e.g., per epoch) with mlflow.log_metric() or a dictionary of metrics with mlflow.log_metrics().
Inputs: Log dataset details and context with mlflow.log_input(), including tags for better categorization.
Tags: Assign custom tags to a run for improved filtering and organization using mlflow.set_tag() or mlflow.set_tags().
Artifacts: Save output files, such as plots, model files, or data samples, with mlflow.log_artifact() for single files or mlflow.log_artifacts() for directories.

How can you compare experiments to find the best model?

Comparing experiments is essential for model selection. MLflow provides two powerful ways to do this: its web UI and its programmatic API.

Comparing via the MLflow Web UI

The MLflow UI offers an intuitive, visual way to compare runs.

Launch the MLflow Server: If it's not running, start it with mise run mlflow:serve.
Select Runs: Navigate to the experiment page, where all runs are listed. Use the checkboxes to select the runs you want to compare.
Click Compare: A "Compare" button will appear. Clicking it opens a detailed view that places the selected runs side-by-side.
Analyze Results: This view provides a comprehensive summary of parameters, metrics, and artifacts for each run. You can use it to identify which configurations yielded the best performance.

Comparing Programmatically

Programmatic comparison is ideal for automated analysis and custom reporting.

Query Runs: Use mlflow.search_runs() to fetch run data into a pandas DataFrame. You can filter by experiment, metrics, parameters, or tags.

import mlflow

# Fetch runs from specific experiments
experiment_ids = ["1", "2"]
runs_df = mlflow.search_runs(experiment_ids)

Filter and Sort: With the data in a DataFrame, you can use pandas to sort and filter the results to find the top-performing runs based on your criteria.

# Find the best run based on validation accuracy
best_run = runs_df.sort_values("metrics.validation_accuracy", ascending=False).iloc[0]
print(f"Best Run ID: {best_run.run_id}")

Visualize Comparisons: Use libraries like Matplotlib or Seaborn to create custom visualizations that make comparisons clear and intuitive.

import matplotlib.pyplot as plt

# Plot validation accuracy for the top 5 runs
top_5_runs = runs_df.sort_values("metrics.validation_accuracy", ascending=False).head(5)
plt.figure(figsize=(12, 7))
plt.bar(top_5_runs["run_id"].str[:7], top_5_runs["metrics.validation_accuracy"])
plt.title("Comparison of Validation Accuracy Across Top 5 Runs")
plt.xlabel("Run ID")
plt.ylabel("Validation Accuracy")
plt.show()

What are some best practices for experiment tracking?

To maximize the value of your experiments, adopt these practices:

Use Clear Naming Conventions: Give experiments and runs descriptive names to make them easily identifiable (e.g., PROD_Retraining_ResNet50 vs. test_run_1).
Align with Business Metrics: Ensure that the metrics you track are directly relevant to project goals and business outcomes.

Leverage Nested Runs: Use nested runs to organize complex experiments, such as hyperparameter tuning, where each child run explores a different parameter set.

with mlflow.start_run(run_name="Hyperparameter Search") as parent_run:
    params = [0.01, 0.02, 0.03]
    for p in params:
        with mlflow.start_run(nested=True, run_name=f"alpha_{p}") as child_run:
            mlflow.log_param("alpha", p)
            # ... training logic ...
            mlflow.log_metric("val_loss", val_loss)

Tag Extensively: Use tags to add metadata like the dataset version, model type, or evaluation status (e.g., dataset:v2, model:xgboost, status:production_candidate).
Track Progress Over Time: Log metrics at each step or epoch to visualize the learning process and diagnose issues like overfitting.
```
# Inside your training loop
mlflow.log_metric(key="train_loss", value=train_loss, step=epoch)
```
Register Promising Models: When a run produces a high-quality model, log it to the MLflow Model Registry to version it and prepare it for deployment.

How do you test code that tracks experiments?

A real backend store is not free, and your test suite is where you feel it. Every time MLflow opens a SQLite database that does not exist yet, it runs its Alembic schema migrations before the first row is written. That costs a few seconds — around seven in the MLOps Python Package — and it happens per database.

The naive fixture creates a fresh database for every test, so the suite pays that price on every single test. When the reference package moved from the file store to SQLite with a per-test database, its suite went from 34 seconds to 339 seconds: a tenfold regression caused entirely by re-running the same migrations hundreds of times.

The fix keeps the isolation and drops the cost. Build one migrated database in a session-scoped fixture, then copy the file for each test. A file copy takes about 0.07 seconds instead of seven, and each test still gets a private, empty store it can write to freely. With this fixture, the reference suite settles at 63 seconds — the honest price of testing against the store you actually ship.

@pytest.fixture(scope="session")
def mlflow_db_template(tmp_path_factory: pytest.TempPathFactory) -> str:
    """Return a migrated but empty MLflow database used as a template by every test.

    Creating an MLflow SQLite store runs its Alembic migrations, which costs seconds.
    Paying that once per session and copying the file per test keeps the isolation of a
    fresh database at the cost of a file copy.
    """
    path = tmp_path_factory.mktemp("mlflow") / "template.db"
    services.MlflowService(
        tracking_uri=f"sqlite:///{path}",
        registry_uri=f"sqlite:///{path}",
        experiment_name="Experiment-Template",
        registry_name="Registry-Template",
    ).start()
    return str(path)


@pytest.fixture(scope="function", autouse=True)
def mlflow_service(tmp_path: str, mlflow_db_template: str) -> T.Generator[services.MlflowService]:
    """Return and start the mlflow service."""
    # Each test gets its own SQLite file under tmp_path, so runs stay isolated the same
    # way the old per-test directories were, with the store the package actually ships.
    database = os.path.join(tmp_path, "mlflow.db")
    shutil.copyfile(mlflow_db_template, database)
    service = services.MlflowService(
        tracking_uri=f"sqlite:///{database}",
        registry_uri=f"sqlite:///{database}",
        experiment_name="Experiment-Testing",
        registry_name="Registry-Testing",
    )
    service.start()
    yield service
    service.stop()

The general lesson outlives MLflow: when a test fixture is slow, look for the expensive, deterministic setup step hiding inside it and pay for it once. A migrated schema, a compiled model, or a seeded dataset can almost always be built at session scope and cloned per test.

How does experiment tracking fit into the MLOps lifecycle?

Experiment tracking is a cornerstone of the MLOps lifecycle, bridging the gap between development and production.

Development: It provides the tools to systematically explore and refine models.
CI/CD Integration: The artifacts and metadata from experiments feed directly into continuous integration and deployment pipelines. For example, a CI pipeline can automatically trigger when a new model is registered, running tests and preparing it for deployment.
Production Monitoring: The metrics and parameters from training runs serve as a baseline for monitoring the model's performance in production. If performance degrades (a concept known as model drift), the tracked experiments provide a clear, reproducible history to inform retraining and debugging efforts.

By maintaining a detailed record of every experiment, you create a transparent, auditable, and efficient workflow that accelerates the entire MLOps cycle.