MLflow Integration Guide

This guide provides comprehensive coverage of FlexLock’s MLflow integration, including the Shadow Run Pattern, best practices, and advanced usage patterns.

Table of Contents

Quick Start
The Shadow Run Pattern
API Reference
Common Workflows
Advanced Patterns
MLflow UI Filtering
Troubleshooting
Best Practices

Quick Start

Installation

pip install mlflow

# or with pixi
pixi add mlflow

Basic Usage

from flexlock import flexcli
from flexlock import mlflow_context
import mlflow

@flexcli
def train(cfg):
    # Your training code
    model = create_model(cfg)
    accuracy = train_model(model, cfg)

    # Log to MLflow
    with mlflow_context(cfg.save_dir, experiment_name="MyExperiment"):
        mlflow.log_metric("accuracy", accuracy)
        mlflow.log_artifact("model.pkl")

    return {"accuracy": accuracy}

That’s it! FlexLock automatically:

✅ Logs your config from run.lock
✅ Logs standard artifacts (experiment.log, etc.)
✅ Manages run lifecycle with Shadow Run Pattern
✅ Keeps your MLflow dashboard clean

The Shadow Run Pattern

The Problem: MLflow Run Management

Traditional MLflow usage creates a dilemma for iterative ML workflows:

Option 1: Resume Same Run

# Resume existing run
with mlflow.start_run(run_id="abc123"):
    mlflow.log_metric("new_metric", 0.95)

Problems:

🔴 Old and new metrics get mixed together
🔴 Hard to tell which values are current
🔴 Can’t easily “reset” diagnostic metrics

Option 2: Create New Run Every Time

# Create new run
with mlflow.start_run():
    mlflow.log_metric("accuracy", 0.95)

Problems:

🔴 After 50 diagnostic iterations, you have 50 runs
🔴 MLflow dashboard becomes cluttered
🔴 Hard to find the “current” state of an experiment

The Solution: Shadow Run Pattern

FlexLock implements “Always New + Deprecate Old” strategy:

Physical State (Disk):        MLflow Runs (Database):
┌──────────────────┐         ┌─────────────────────────┐
│ outputs/exp/r01/ │◄────────│ Run A (deprecated)      │
│                  │         │ - Config (snapshot)     │
│ • run.lock       │◄────┐   │ - Metrics v1            │
│ • model.pkl      │     │   └─────────────────────────┘
│ • experiment.log │     │
│ • results.json   │     │   ┌─────────────────────────┐
└──────────────────┘     └───│ Run B (ACTIVE) ⭐       │
  (Never moves)              │ - Config (inherited)    │
                             │ - Metrics v2 (latest)   │
                             │ - Model (inherited)     │
                             └─────────────────────────┘

Key Principles:

Physical Truth: Data lives in outputs/exp/run_01 (permanent)
MLflow Truth: We create a new Run ID every time
The Spotlight: New run tagged active, old run tagged deprecated
Artifact Inheritance: New run re-logs artifacts from disk

Result: Your MLflow dashboard shows exactly one active run per experiment, always with the latest metrics!

Shadow Run Lifecycle

Iteration 1:
┌─────────────────┐
│ Run A (active)  │  ← First run
│ tag: active     │
└─────────────────┘

Iteration 2:
┌─────────────────┐
│ Run A (depr.)   │  ← Automatically deprecated
│ tag: deprecated │
│ superseded_by:B │
└─────────────────┘
┌─────────────────┐
│ Run B (active)  │  ← New active run
│ tag: active     │
│ supersedes: A   │
└─────────────────┘

Iteration 3:
┌─────────────────┐
│ Run A (depr.)   │
└─────────────────┘
┌─────────────────┐
│ Run B (depr.)   │  ← Also deprecated now
│ superseded_by:C │
└─────────────────┘
┌─────────────────┐
│ Run C (active)  │  ← Latest active run
│ tag: active     │
│ supersedes: B   │
└─────────────────┘

Filter in MLflow UI: tags.flexlock.status = 'active'

You see only Run C (the latest), with all inherited artifacts and latest metrics!

API Reference

`mlflow_context()`

Context manager for MLflow logging with automatic run management.

from flexlock import mlflow_context

with mlflow_context(
    save_dir: str | Path,
    experiment_name: str = "Default",
    run_name: str = None,
    tags: Dict[str, str] = None,
    log_config: bool = True,
    log_artifacts: bool = True,
) as run:
    # Your MLflow logging code
    mlflow.log_metric("accuracy", 0.95)
    mlflow.log_artifact("plot.png")

Parameters

Parameter	Type	Default	Description
`save_dir`	`str \| Path`	required	Directory containing run data (e.g., `outputs/exp/run_01`)
`experiment_name`	`str`	`"Default"`	MLflow experiment name
`run_name`	`str`	`None`	Name for this MLflow run (defaults to directory name)
`tags`	`Dict[str, str]`	`None`	Additional custom tags
`log_config`	`bool`	`True`	If True, logs parameters from `run.lock` file
`log_artifacts`	`bool`	`True`	If True, logs standard artifacts

Automatic Behaviors

When log_config=True:

Reads save_dir/run.lock
Extracts configuration parameters
Flattens nested config (e.g., model.lr → model.lr)
Truncates long strings to 250 characters
Logs as MLflow parameters
Logs run.lock file as artifact

When log_artifacts=True: Automatically logs these files if they exist:

experiment.log - Main experiment log
stderr.log - Standard error output
stdout.log - Standard output

Tags Added Automatically:

flexlock.dir: Absolute path to save_dir
flexlock.status: "active" (new run) or "deprecated" (old run)
flexlock.supersedes: Run ID of previous active run (if exists)
flexlock.superseded_by: Run ID of next run (added to deprecated runs)

Returns

Yields mlflow.ActiveRun object, or None if MLflow not installed.

Example

from flexlock import mlflow_context
from pathlib import Path
import mlflow

save_dir = Path("outputs/my_exp/run_01")

with mlflow_context(
    save_dir=save_dir,
    experiment_name="ResNet_Training",
    run_name="resnet50_imagenet",
    tags={"model": "resnet50", "dataset": "imagenet"},
) as run:
    # Log training metrics
    mlflow.log_metric("train_loss", 0.32)
    mlflow.log_metric("val_accuracy", 0.95)

    # Log model
    mlflow.log_artifact("model.pkl")

    # Log plots
    mlflow.log_artifact(save_dir / "learning_curve.png")

    # Access run info
    print(f"Run ID: {run.info.run_id}")

Common Workflows

Workflow 1: Simple Training with Logging

Scenario: Train a model and log metrics to MLflow.

from flexlock import flexcli
from flexlock import mlflow_context
import mlflow

@flexcli
def train(cfg):
    """Train a model and log to MLflow."""
    # Training code
    model = build_model(cfg.model)
    history = model.fit(
        train_data,
        epochs=cfg.epochs,
        validation_data=val_data
    )

    # Evaluation
    test_accuracy = model.evaluate(test_data)

    # Save model to disk
    model.save(f"{cfg.save_dir}/model.h5")

    # Log to MLflow
    with mlflow_context(cfg.save_dir, experiment_name="Training"):
        # Metrics
        mlflow.log_metric("test_accuracy", test_accuracy)
        mlflow.log_metric("final_train_loss", history.history['loss'][-1])

        # Model
        mlflow.log_artifact(f"{cfg.save_dir}/model.h5")

    return {"test_accuracy": test_accuracy}

Run:

python train.py -c config.yaml

# View in MLflow UI
mlflow ui

Workflow 2: Separate Compute and Diagnostics

Scenario: Run expensive training once, then iterate on diagnostic plots and metrics.

Step 1: Training (Expensive, Run Once)

# train.py
from flexlock import flexcli

@flexcli
def train(cfg):
    """Expensive training - saves to disk, NO MLflow yet."""
    # Long training process
    model = train_for_24_hours(cfg)

    # Save results to disk
    save_results(cfg.save_dir, model, predictions)

    return {"status": "complete"}

# Run once
python train.py -c config.yaml
# Output: outputs/my_exp/2024-01-15_10-30-45/

Step 2: Diagnostics (Fast, Run Many Times)

# diagnose.py
from flexlock import flexcli
from flexlock import mlflow_context
import mlflow
import matplotlib.pyplot as plt

@flexcli
def diagnose(cfg):
    """Load results and create diagnostic plots."""
    # Load pre-computed results (fast!)
    results = load_results(cfg.save_dir)

    # Generate plots
    fig = create_confusion_matrix(results)
    fig.savefig(f"{cfg.save_dir}/confusion_matrix.png")

    fig2 = create_learning_curves(results)
    fig2.savefig(f"{cfg.save_dir}/learning_curves.png")

    # Calculate metrics
    f1 = compute_f1_score(results)
    precision = compute_precision(results)

    # Log to MLflow
    with mlflow_context(cfg.save_dir, experiment_name="Diagnostics"):
        mlflow.log_metric("f1_score", f1)
        mlflow.log_metric("precision", precision)
        mlflow.log_artifact(f"{cfg.save_dir}/confusion_matrix.png")
        mlflow.log_artifact(f"{cfg.save_dir}/learning_curves.png")

# Point to existing training directory
python diagnose.py save_dir=outputs/my_exp/2024-01-15_10-30-45/

# Tweak plotting code, re-run (creates NEW MLflow run)
python diagnose.py save_dir=outputs/my_exp/2024-01-15_10-30-45/

# Adjust metrics calculation, re-run again
python diagnose.py save_dir=outputs/my_exp/2024-01-15_10-30-45/

# ... iterate 50 times if needed!

MLflow Dashboard:

Filter: tags.flexlock.status = 'active'
See: Exactly 1 run with latest diagnostics
Contains: Original config + logs + latest metrics + latest plots

Workflow 3: Hyperparameter Sweeps

Scenario: Run parameter sweep and track all experiments.

from flexlock import flexcli
from flexlock import mlflow_context
import mlflow

@flexcli
def train_sweep(cfg):
    """Single training run in a sweep."""
    # Train with current hyperparameters
    model = train_model(cfg)
    val_acc = validate_model(model)

    # Log to MLflow
    with mlflow_context(
        save_dir=cfg.save_dir,
        experiment_name="HPO_Sweep",
        tags={
            "sweep_name": cfg.sweep_name,
            "optimizer": cfg.optimizer,
        }
    ):
        mlflow.log_metric("val_accuracy", val_acc)

        # Tag best run
        if val_acc > cfg.best_acc_threshold:
            mlflow.set_tag("is_best", "true")

    return {"val_acc": val_acc}

Config (sweep.yaml):

sweep_name: lr_optimizer_sweep
model:
  architecture: resnet50
training:
  epochs: 50
  batch_size: 32
best_acc_threshold: 0.90

# Sweep parameters
lr: ???  # To be overridden
optimizer: ???  # To be overridden

Run Sweep:

python train_sweep.py -c sweep.yaml \
  -o lr=0.001,0.01,0.1 \
     optimizer=adam,sgd,rmsprop

MLflow Dashboard:

# Filter: All runs in sweep
tags.sweep_name = 'lr_optimizer_sweep'

# Filter: Only best runs
tags.is_best = 'true'

# Filter: Specific optimizer
tags.optimizer = 'adam'

Columns to display:

params.lr
params.optimizer
metrics.val_accuracy
tags.is_best

Workflow 4: Multi-Stage Pipeline

Scenario: Pipeline with preprocess → train → evaluate, all logged to MLflow.

from flexlock import Project
from flexlock import mlflow_context
import mlflow

proj = Project("pipeline.yaml")

# Stage 1: Preprocess (no MLflow)
@proj.register("preprocess")
def preprocess(cfg):
    data = load_and_clean_data(cfg)
    save_data(cfg.save_dir, data)
    return {"num_samples": len(data)}

# Stage 2: Train with MLflow
@proj.register("train")
def train(cfg):
    data = load_data(cfg.save_dir)
    model = train_model(data, cfg)

    with mlflow_context(cfg.save_dir, experiment_name="Pipeline"):
        mlflow.log_metric("train_accuracy", model.train_acc)
        mlflow.log_artifact("model.pkl")

    return {"train_acc": model.train_acc}

# Stage 3: Evaluate with MLflow (same save_dir!)
@proj.register("evaluate")
def evaluate(cfg):
    model = load_model(cfg.save_dir)
    test_acc = model.evaluate(test_data)

    # Creates NEW MLflow run, deprecates training run
    # Inherits config and model.pkl, adds evaluation metrics
    with mlflow_context(cfg.save_dir, experiment_name="Pipeline"):
        mlflow.log_metric("test_accuracy", test_acc)
        mlflow.log_artifact("test_results.json")

    return {"test_acc": test_acc}

# Run pipeline
proj.submit(stage="preprocess", wait=True)
proj.submit(stage="train", wait=True)
proj.submit(stage="evaluate", wait=True)

Result in MLflow:

One active run per experiment
Contains:
- Config (from preprocess stage)
- Model (from train stage)
- Train accuracy (from train stage)
- Test accuracy (from evaluate stage)
- Test results (from evaluate stage)

Advanced Patterns

Pattern 1: Custom Run Names

Use descriptive run names for better organization:

from datetime import datetime

timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
run_name = f"{cfg.model_type}_{cfg.dataset}_{timestamp}"

with mlflow_context(
    cfg.save_dir,
    experiment_name="Experiments",
    run_name=run_name,
):
    mlflow.log_metric("accuracy", 0.95)

Pattern 2: Hierarchical Tags

Use dot-separated tags for hierarchical filtering:

with mlflow_context(
    cfg.save_dir,
    tags={
        "project": "image_classification",
        "model.architecture": "resnet50",
        "model.pretrained": "imagenet",
        "data.dataset": "cifar100",
        "data.augmentation": "heavy",
    }
):
    mlflow.log_metric("accuracy", 0.95)

Filter in MLflow UI:

tags.project = 'image_classification' AND tags.`model.architecture` = 'resnet50'

Pattern 3: Conditional Logging

Only log to MLflow in certain conditions:

from flexlock import mlflow_context
import mlflow

@flexcli
def train(cfg):
    model = train_model(cfg)
    accuracy = validate(model)

    # Only log if accuracy exceeds threshold
    if accuracy > cfg.log_threshold:
        with mlflow_context(
            cfg.save_dir,
            experiment_name="Good_Runs",
        ):
            mlflow.log_metric("val_accuracy", accuracy)
            mlflow.log_artifact("model.pkl")

    return {"accuracy": accuracy}

Pattern 4: Manual Parameter Logging

Add runtime parameters not in config:

import git
import platform

with mlflow_context(cfg.save_dir, experiment_name="Training"):
    # FlexLock automatically logs config from run.lock

    # Log additional runtime info
    repo = git.Repo(".")
    mlflow.log_param("git_commit", repo.head.commit.hexsha)
    mlflow.log_param("git_branch", repo.active_branch.name)
    mlflow.log_param("hostname", platform.node())

    # Log metrics
    mlflow.log_metric("accuracy", 0.95)

Pattern 5: Nested Experiments

Organize related experiments hierarchically:

# Parent experiment: Overall project
with mlflow_context(
    cfg.save_dir,
    experiment_name="ImageClassification/ResNet",
    tags={"project": "image_classification"},
):
    mlflow.log_metric("test_accuracy", 0.95)

# Child experiment: Specific variant
with mlflow_context(
    cfg.save_dir,
    experiment_name="ImageClassification/ResNet/Ablations",
    tags={"project": "image_classification", "study": "ablation"},
):
    mlflow.log_metric("ablation_accuracy", 0.92)

MLflow UI Filtering

Essential Filters

Show only active runs (recommended):

tags.flexlock.status = 'active'

Show deprecated runs:

tags.flexlock.status = 'deprecated'

Show all runs for specific directory:

tags.flexlock.dir = '/absolute/path/to/outputs/exp/run_01'

Advanced Filters

Runs from a specific sweep:

tags.flexlock.status = 'active' AND tags.sweep_name = 'my_sweep'

Best runs only:

tags.flexlock.status = 'active' AND tags.is_best = 'true'

Runs with specific model:

tags.flexlock.status = 'active' AND tags.model = 'resnet50'

Runs with accuracy above threshold:

tags.flexlock.status = 'active' AND metrics.val_accuracy > 0.90

Tracing Run History

Find what run A superseded:

tags.flexlock.superseded_by = '<run_A_id>'

Find what superseded run A:

tags.flexlock.supersedes = '<run_A_id>'

Complete lineage:

tags.flexlock.dir = '/path/to/experiment' ORDER BY start_time ASC

This shows the complete history: Run A → Run B → Run C → …

Troubleshooting

Problem: MLflow Not Installed

Error:

MLflow not installed. To use mlflow_context, install with: pip install mlflow

Solution:

pip install mlflow

# or with pixi
pixi add mlflow

Verify:

python -c "import mlflow; print(mlflow.__version__)"

Problem: Parameters Not Showing Up

Symptoms: No parameters logged in MLflow run

Possible Causes:

run.lock file doesn’t exist
```
ls outputs/my_exp/run_01/run.lock  # Should exist
```
Solution: Make sure you’re using @flexcli decorator or Project API

log_config=False

# Wrong
with mlflow_context(save_dir, log_config=False):  # Disables param logging

Solution: Use default log_config=True

Parameters too long (>250 chars)

FlexLock automatically truncates to 250 characters. Check logs:
```
MLflow config logging warning: ...
```

Problem: Duplicate Metrics from Multiple Runs

Symptoms: MLflow UI shows multiple runs for the same experiment

Cause: Creating new save_dir for each run instead of reusing

Wrong:

from datetime import datetime

@flexcli
def diagnose(cfg):
    # Creates NEW directory each time
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    save_dir = f"outputs/exp/{timestamp}"  # ❌ New dir every time

    with mlflow_context(save_dir):
        mlflow.log_metric("accuracy", 0.95)

Correct:

@flexcli
def diagnose(cfg):
    # Reuse EXISTING directory
    save_dir = "outputs/exp/2024-01-15_10-30-45"  # ✅ Same dir

    with mlflow_context(save_dir):
        mlflow.log_metric("accuracy", 0.95)

Result: Shadow Run Pattern activates, old run gets deprecated

Problem: “Can’t Find Previous Run”

Message in logs:

No previous run found or search failed: ...

This is normal! On the first run for a directory, there’s no previous run to deprecate.

Problem: Deprecated Runs Still Visible

Issue: MLflow UI shows both active and deprecated runs

Solution: Add filter in MLflow UI:

tags.flexlock.status = 'active'

Pro Tip: Save this filter as a view:

Apply filter
Click “Save View”
Name it “Active Runs Only”

Problem: Artifacts Not Logged

Symptoms: Missing experiment.log or other artifacts

Possible Causes:

Files don’t exist in save_dir

ls outputs/my_exp/run_01/experiment.log  # Should exist

log_artifacts=False

with mlflow_context(save_dir, log_artifacts=False):  # Disables artifacts

Solution: Use default log_artifacts=True

Custom artifact names

FlexLock auto-logs: experiment.log, stderr.log, stdout.log, run.lock

For custom files, log manually:
```
with mlflow_context(save_dir):
    mlflow.log_artifact("custom_results.json")
```

Best Practices

✅ DO

Use descriptive experiment names

# Good
with mlflow_context(save_dir, experiment_name="ResNet50_ImageNet"):

# Bad
with mlflow_context(save_dir, experiment_name="Exp1"):

Add meaningful tags

with mlflow_context(
    save_dir,
    tags={
        "model": "resnet50",
        "dataset": "imagenet",
        "optimizer": "adam",
    }
):

Reuse save_dir for diagnostics

# Training creates directory
save_dir = "outputs/exp/run_01"

# Diagnostics reuse same directory (enables Shadow Run Pattern)
python diagnose.py save_dir=outputs/exp/run_01

Filter by active status
```
tags.flexlock.status = 'active'
```
Keep deprecated runs for history

Don’t delete old runs - they provide audit trail!

❌ DON’T

Don’t manually manage run resumption

# Bad - manually managing runs
existing_run_id = find_run_id(...)
with mlflow.start_run(run_id=existing_run_id):
    ...

# Good - let mlflow_context handle it
with mlflow_context(save_dir):
    ...

Don’t log config parameters twice

# Bad - already logged automatically
with mlflow_context(save_dir):
    mlflow.log_param("learning_rate", cfg.lr)  # ❌ Duplicate

# Good - config auto-logged from run.lock
with mlflow_context(save_dir):
    mlflow.log_metric("accuracy", 0.95)  # ✅ Just log metrics

Don’t create new directories for diagnostic iterations

# Bad
save_dir = f"outputs/exp/diagnostic_{iteration}"  # ❌ New dir

# Good
save_dir = "outputs/exp/run_01"  # ✅ Reuse existing

Don’t use mlflow.start_run() directly

# Bad
with mlflow.start_run():
    ...

# Good
with mlflow_context(save_dir):
    ...

Don’t delete deprecated runs

Keep them for history and audit trails!

Summary

FlexLock’s MLflow integration provides:

Shadow Run Pattern - Clean dashboard with one active run per experiment
Automatic Logging - Config and artifacts logged from disk
Iterative Diagnostics - Run diagnostics 50 times, see only latest
Flexible Filtering - Powerful tags for organizing experiments

Key Takeaway: With mlflow_context(), you can separate expensive computation from iterative diagnostics while maintaining a clean, organized MLflow dashboard!

Next Steps

HPC Integration: Scale experiments with HPC Integration Guide
Debugging: Use interactive debugging with Debugging Guide
Python API: See Python API Reference for programmatic usage

Need Help?

Issues: Report bugs
Discussions: Ask questions
Docs: Full documentation

MLflow Integration Guide

Table of Contents

Quick Start

Installation

Basic Usage

The Shadow Run Pattern

The Problem: MLflow Run Management

The Solution: Shadow Run Pattern

Shadow Run Lifecycle

API Reference

mlflow_context()

Parameters

Automatic Behaviors

Returns

Example

Common Workflows

Workflow 1: Simple Training with Logging

Workflow 2: Separate Compute and Diagnostics

Step 1: Training (Expensive, Run Once)

Step 2: Diagnostics (Fast, Run Many Times)

Workflow 3: Hyperparameter Sweeps

Workflow 4: Multi-Stage Pipeline

Advanced Patterns

Pattern 1: Custom Run Names

Pattern 2: Hierarchical Tags

Pattern 3: Conditional Logging

Pattern 4: Manual Parameter Logging

Pattern 5: Nested Experiments

MLflow UI Filtering

Essential Filters

Advanced Filters

Tracing Run History

Troubleshooting

Problem: MLflow Not Installed

Problem: Parameters Not Showing Up

Problem: Duplicate Metrics from Multiple Runs

Problem: “Can’t Find Previous Run”

Problem: Deprecated Runs Still Visible

Problem: Artifacts Not Logged

Best Practices

✅ DO

❌ DON’T

Summary

Next Steps

Need Help?

`mlflow_context()`