MLflow Experiment Tracking: From Single Run to Team Workflow

We'll run a tiny PyTorch training loop. The MLflow wrapper is 3 lines:

mlflow.set_tracking_uri('file:///tmp/mlruns')
with mlflow.start_run(run_name='lr=1e-3 bs=32') as run:
    mlflow.log_params({'lr': 1e-3, 'batch_size': 32})
    for step, loss in training_loop():
        mlflow.log_metric('train_loss', loss, step=step)

What gets recorded automatically:

• run_id — unique UUID for this run
• start_time, end_time, status
• experiment_id — a grouping label
• Everything you explicitly logged via log_params, log_metric, log_artifact
• If you turn on autolog(): the model, sample inputs, and framework version too

Set up MLflow — file backend + experiment

We point MLflow at a clean file backend under /tmp so every lab session starts fresh. In production you'd swap this for a tracking server URI (http://mlflow.company.internal:5000) with Postgres + S3 behind it — but the Python API is identical either way, so code written against the file backend ports unchanged.

import os, tempfile, shutil, time
import mlflow, torch, torch.nn as nn

# Use a clean file backend under /tmp so each run starts fresh
TRACKING_DIR = tempfile.mkdtemp(prefix='mlruns_')
mlflow.set_tracking_uri(f'file://{TRACKING_DIR}')
print(f"MLflow tracking URI: file://{TRACKING_DIR}")

experiment_name = 'preporato-lab-i13'
experiment_id = mlflow.set_experiment(experiment_name).experiment_id
print(f"Experiment: {experiment_name} (id={experiment_id})")

# A deliberately small model so the lab stays fast
device = 'cuda'
torch.manual_seed(42)
model = nn.Sequential(nn.Linear(64, 128), nn.GELU(), nn.Linear(128, 10)).to(device)
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()

n_steps = 50
logged_params_dict = {
    'learning_rate':  1e-3,
    'batch_size':     64,
    'optimizer':      'Adam',
    'model_arch':     'MLP-64-128-10',
    'n_steps':        n_steps,
    'device':         'cuda',
}

The training run — tags + params + per-step metrics

Everything inside the with mlflow.start_run(): block gets attached to this run. Three kinds of data go in:

• Tags — the run's identity (team, owner, git commit, env). Used for filtering.
• Params — the static config (lr, batch, arch). Used for comparison across runs.
• Metrics — time-series measurements (loss per step). Used for plots and model selection.

# The main MLflow-wrapped training run
loss_series = []
with mlflow.start_run(run_name='baseline-lr1e-3') as run:
    run_id = run.info.run_id
    mlflow.set_tags({'team': 'ml-platform', 'owner': 'preporato-labs', 'git_commit': 'abc123def', 'env': 'dev'})
    mlflow.log_params(logged_params_dict)
    for step in range(n_steps):
        x = torch.randn(64, 64, device=device)
        y = torch.randint(0, 10, (64,), device=device)
        opt.zero_grad()
        out = model(x)
        loss = loss_fn(out, y)
        loss.backward(); opt.step()
        mlflow.log_metric('train_loss', float(loss), step=step)
        loss_series.append((step, float(loss)))
    mlflow.log_metric('final_loss', float(loss))

Post-run view — query the run back from MLflow

The Run object on-disk now has everything the training job logged. In production you'd browse this in the MLflow UI (mlflow ui --host 0.0.0.0 --port 5000) or query programmatically via search_runs (we do that in step 3).

# Assemble the post-run view
finished = mlflow.get_run(run_id)
run_info = {
    'run_id':        finished.info.run_id,
    'experiment_id': finished.info.experiment_id,
    'status':        finished.info.status,
    'start_time':    finished.info.start_time,
    'end_time':      finished.info.end_time,
}
logged_params  = logged_params_dict
logged_metrics = {'train_loss': loss_series}

print(f"\nRun {run_id[:12]} finished: status={run_info['status']}")
print(f"  params logged: {list(logged_params.keys())}")
print(f"  train_loss: {logged_metrics['train_loss'][0][1]:.4f} → {logged_metrics['train_loss'][-1][1]:.4f} over {n_steps} steps")