aigroup-workflow 2.2.1 → 2.2.2

package/skills/ai-ml/ml-pipeline/references/experiment-tracking.md

@@ -1,833 +1,833 @@
-# Experiment Tracking
-
----
-
-## Overview
-
-Experiment tracking enables reproducibility, comparison, and collaboration in ML development. It captures hyperparameters, metrics, artifacts, and model versions to ensure every experiment can be reproduced and compared.
-
-## When to Use This Reference
-
-- Setting up MLflow for experiment tracking
-- Implementing Weights & Biases integration
-- Creating model registries and versioning
-- Comparing experiments and selecting models
-- Building custom tracking solutions
-
-## When NOT to Use
-
-- Quick one-off experiments without reproducibility needs
-- Simple scripts without hyperparameters
-- Non-ML projects
-
----
-
-## MLflow Integration
-
-### Basic Experiment Tracking
-
-```python
-import mlflow
-from mlflow.tracking import MlflowClient
-from pathlib import Path
-import json
-
-class MLflowTracker:
-    """MLflow experiment tracking wrapper."""
-
-    def __init__(
-        self,
-        experiment_name: str,
-        tracking_uri: str = "http://localhost:5000",
-        artifact_location: str = None,
-    ):
-        mlflow.set_tracking_uri(tracking_uri)
-
-        # Create or get experiment
-        experiment = mlflow.get_experiment_by_name(experiment_name)
-        if experiment is None:
-            self.experiment_id = mlflow.create_experiment(
-                experiment_name,
-                artifact_location=artifact_location,
-            )
-        else:
-            self.experiment_id = experiment.experiment_id
-
-        mlflow.set_experiment(experiment_name)
-        self.client = MlflowClient()
-        self.run = None
-
-    def start_run(
-        self,
-        run_name: str = None,
-        tags: dict = None,
-        nested: bool = False,
-    ) -> str:
-        """Start a new MLflow run."""
-        self.run = mlflow.start_run(
-            run_name=run_name,
-            experiment_id=self.experiment_id,
-            nested=nested,
-        )
-
-        if tags:
-            mlflow.set_tags(tags)
-
-        return self.run.info.run_id
-
-    def end_run(self, status: str = "FINISHED") -> None:
-        """End the current run."""
-        mlflow.end_run(status=status)
-        self.run = None
-
-    def log_params(self, params: dict) -> None:
-        """Log hyperparameters."""
-        mlflow.log_params(params)
-
-    def log_metrics(self, metrics: dict, step: int = None) -> None:
-        """Log metrics with optional step."""
-        for key, value in metrics.items():
-            mlflow.log_metric(key, value, step=step)
-
-    def log_artifact(self, local_path: str, artifact_path: str = None) -> None:
-        """Log file or directory as artifact."""
-        mlflow.log_artifact(local_path, artifact_path)
-
-    def log_model(
-        self,
-        model,
-        artifact_path: str,
-        registered_model_name: str = None,
-        signature=None,
-        input_example=None,
-    ) -> str:
-        """Log model with optional registration."""
-        from mlflow.models import infer_signature
-
-        if signature is None and input_example is not None:
-            signature = infer_signature(input_example, model.predict(input_example))
-
-        model_info = mlflow.sklearn.log_model(
-            model,
-            artifact_path=artifact_path,
-            registered_model_name=registered_model_name,
-            signature=signature,
-            input_example=input_example,
-        )
-
-        return model_info.model_uri
-
-# Usage example
-def train_with_mlflow(
-    model,
-    X_train,
-    y_train,
-    X_val,
-    y_val,
-    params: dict,
-):
-    """Complete training run with MLflow tracking."""
-    tracker = MLflowTracker("my_experiment")
-
-    tracker.start_run(
-        run_name=f"run_{params['model_type']}",
-        tags={
-            "model_type": params["model_type"],
-            "dataset_version": "v1.0",
-            "author": "ml-team",
-        },
-    )
-
-    try:
-        # Log parameters
-        tracker.log_params(params)
-
-        # Train model
-        model.fit(X_train, y_train)
-
-        # Evaluate and log metrics
-        train_score = model.score(X_train, y_train)
-        val_score = model.score(X_val, y_val)
-
-        tracker.log_metrics({
-            "train_accuracy": train_score,
-            "val_accuracy": val_score,
-        })
-
-        # Log model
-        model_uri = tracker.log_model(
-            model,
-            artifact_path="model",
-            registered_model_name="my_model",
-            input_example=X_train[:5],
-        )
-
-        tracker.end_run()
-        return model_uri
-
-    except Exception as e:
-        tracker.end_run(status="FAILED")
-        raise
-```
-
-### PyTorch Model Logging
-
-```python
-import mlflow.pytorch
-import torch
-
-def log_pytorch_model(
-    model: torch.nn.Module,
-    artifact_path: str,
-    registered_model_name: str = None,
-    sample_input: torch.Tensor = None,
-) -> str:
-    """Log PyTorch model with signature inference."""
-    from mlflow.models import infer_signature
-
-    # Create signature from sample input
-    signature = None
-    if sample_input is not None:
-        model.eval()
-        with torch.no_grad():
-            sample_output = model(sample_input)
-
-        signature = infer_signature(
-            sample_input.numpy(),
-            sample_output.numpy(),
-        )
-
-    model_info = mlflow.pytorch.log_model(
-        model,
-        artifact_path=artifact_path,
-        registered_model_name=registered_model_name,
-        signature=signature,
-    )
-
-    return model_info.model_uri
-
-def load_pytorch_model(model_uri: str, device: str = "cpu") -> torch.nn.Module:
-    """Load PyTorch model from MLflow."""
-    model = mlflow.pytorch.load_model(model_uri, map_location=device)
-    return model
-```
-
-### Model Registry Operations
-
-```python
-from mlflow.tracking import MlflowClient
-from mlflow.entities.model_registry import ModelVersion
-
-class ModelRegistry:
-    """MLflow Model Registry wrapper."""
-
-    def __init__(self, tracking_uri: str = "http://localhost:5000"):
-        mlflow.set_tracking_uri(tracking_uri)
-        self.client = MlflowClient()
-
-    def register_model(
-        self,
-        model_uri: str,
-        name: str,
-        tags: dict = None,
-        description: str = None,
-    ) -> ModelVersion:
-        """Register a new model version."""
-        result = mlflow.register_model(model_uri, name)
-
-        if tags:
-            for key, value in tags.items():
-                self.client.set_model_version_tag(name, result.version, key, value)
-
-        if description:
-            self.client.update_model_version(
-                name,
-                result.version,
-                description=description,
-            )
-
-        return result
-
-    def transition_model_stage(
-        self,
-        name: str,
-        version: str,
-        stage: str,
-        archive_existing: bool = True,
-    ) -> ModelVersion:
-        """Transition model to new stage (Staging, Production, Archived)."""
-        return self.client.transition_model_version_stage(
-            name=name,
-            version=version,
-            stage=stage,
-            archive_existing_versions=archive_existing,
-        )
-
-    def get_latest_version(
-        self,
-        name: str,
-        stages: list[str] = None,
-    ) -> list[ModelVersion]:
-        """Get latest model versions by stage."""
-        return self.client.get_latest_versions(name, stages=stages)
-
-    def load_production_model(self, name: str) -> any:
-        """Load the production model."""
-        model_uri = f"models:/{name}/Production"
-        return mlflow.pyfunc.load_model(model_uri)
-
-    def compare_versions(
-        self,
-        name: str,
-        version_a: str,
-        version_b: str,
-    ) -> dict:
-        """Compare two model versions."""
-        v_a = self.client.get_model_version(name, version_a)
-        v_b = self.client.get_model_version(name, version_b)
-
-        run_a = self.client.get_run(v_a.run_id)
-        run_b = self.client.get_run(v_b.run_id)
-
-        return {
-            "version_a": {
-                "version": version_a,
-                "metrics": run_a.data.metrics,
-                "params": run_a.data.params,
-            },
-            "version_b": {
-                "version": version_b,
-                "metrics": run_b.data.metrics,
-                "params": run_b.data.params,
-            },
-        }
-```
-
----
-
-## Weights & Biases Integration
-
-### Basic W&B Tracking
-
-```python
-import wandb
-from pathlib import Path
-
-class WandbTracker:
-    """Weights & Biases experiment tracking wrapper."""
-
-    def __init__(
-        self,
-        project: str,
-        entity: str = None,
-        config: dict = None,
-    ):
-        self.project = project
-        self.entity = entity
-        self.config = config
-        self.run = None
-
-    def start_run(
-        self,
-        name: str = None,
-        tags: list[str] = None,
-        group: str = None,
-        job_type: str = "train",
-        resume: str = None,
-    ) -> wandb.Run:
-        """Initialize W&B run."""
-        self.run = wandb.init(
-            project=self.project,
-            entity=self.entity,
-            name=name,
-            config=self.config,
-            tags=tags,
-            group=group,
-            job_type=job_type,
-            resume=resume,
-        )
-        return self.run
-
-    def log(self, data: dict, step: int = None, commit: bool = True) -> None:
-        """Log metrics and data."""
-        wandb.log(data, step=step, commit=commit)
-
-    def log_artifact(
-        self,
-        name: str,
-        artifact_type: str,
-        path: str,
-        metadata: dict = None,
-    ) -> wandb.Artifact:
-        """Log artifact (model, dataset, etc.)."""
-        artifact = wandb.Artifact(
-            name=name,
-            type=artifact_type,
-            metadata=metadata,
-        )
-
-        if Path(path).is_dir():
-            artifact.add_dir(path)
-        else:
-            artifact.add_file(path)
-
-        self.run.log_artifact(artifact)
-        return artifact
-
-    def log_model(
-        self,
-        model_path: str,
-        name: str,
-        metadata: dict = None,
-        aliases: list[str] = None,
-    ) -> wandb.Artifact:
-        """Log model artifact with aliases."""
-        artifact = wandb.Artifact(
-            name=name,
-            type="model",
-            metadata=metadata,
-        )
-
-        if Path(model_path).is_dir():
-            artifact.add_dir(model_path)
-        else:
-            artifact.add_file(model_path)
-
-        self.run.log_artifact(artifact, aliases=aliases or ["latest"])
-        return artifact
-
-    def watch_model(
-        self,
-        model,
-        log: str = "all",
-        log_freq: int = 100,
-    ) -> None:
-        """Watch model for gradient and parameter logging."""
-        wandb.watch(model, log=log, log_freq=log_freq)
-
-    def finish(self, exit_code: int = 0) -> None:
-        """Finish the run."""
-        wandb.finish(exit_code=exit_code)
-
-# Usage with PyTorch
-def train_with_wandb(
-    model: torch.nn.Module,
-    train_loader,
-    val_loader,
-    config: dict,
-):
-    """Training with W&B tracking."""
-    tracker = WandbTracker(
-        project="my-project",
-        config=config,
-    )
-
-    tracker.start_run(
-        name=f"experiment_{config['model_type']}",
-        tags=["baseline", config["model_type"]],
-        group="hyperparameter_search",
-    )
-
-    # Watch model gradients
-    tracker.watch_model(model)
-
-    for epoch in range(config["epochs"]):
-        model.train()
-        for batch_idx, (data, target) in enumerate(train_loader):
-            # Training step
-            loss = train_step(model, data, target)
-
-            tracker.log({
-                "train/loss": loss,
-                "train/epoch": epoch,
-            })
-
-        # Validation
-        val_metrics = evaluate(model, val_loader)
-        tracker.log({
-            "val/loss": val_metrics["loss"],
-            "val/accuracy": val_metrics["accuracy"],
-            "epoch": epoch,
-        })
-
-    # Save and log model
-    torch.save(model.state_dict(), "model.pt")
-    tracker.log_model(
-        "model.pt",
-        name="trained_model",
-        metadata={"accuracy": val_metrics["accuracy"]},
-        aliases=["latest", "best"],
-    )
-
-    tracker.finish()
-```
-
-### W&B Sweeps for Hyperparameter Tuning
-
-```python
-import wandb
-
-sweep_config = {
-    "method": "bayes",  # bayes, grid, random
-    "metric": {
-        "name": "val/loss",
-        "goal": "minimize",
-    },
-    "parameters": {
-        "learning_rate": {
-            "distribution": "log_uniform_values",
-            "min": 1e-5,
-            "max": 1e-2,
-        },
-        "batch_size": {
-            "values": [16, 32, 64, 128],
-        },
-        "hidden_size": {
-            "values": [128, 256, 512],
-        },
-        "dropout": {
-            "distribution": "uniform",
-            "min": 0.1,
-            "max": 0.5,
-        },
-    },
-    "early_terminate": {
-        "type": "hyperband",
-        "min_iter": 3,
-    },
-}
-
-def sweep_train():
-    """Training function for sweep."""
-    with wandb.init() as run:
-        config = wandb.config
-
-        model = build_model(
-            hidden_size=config.hidden_size,
-            dropout=config.dropout,
-        )
-
-        optimizer = torch.optim.Adam(
-            model.parameters(),
-            lr=config.learning_rate,
-        )
-
-        train_loader = DataLoader(train_dataset, batch_size=config.batch_size)
-
-        for epoch in range(10):
-            loss = train_epoch(model, train_loader, optimizer)
-            val_loss = evaluate(model, val_loader)
-
-            wandb.log({
-                "train/loss": loss,
-                "val/loss": val_loss,
-                "epoch": epoch,
-            })
-
-# Run sweep
-sweep_id = wandb.sweep(sweep_config, project="my-project")
-wandb.agent(sweep_id, function=sweep_train, count=50)
-```
-
----
-
-## Custom Experiment Tracking
-
-### Lightweight Tracker
-
-```python
-import json
-from datetime import datetime
-from pathlib import Path
-from dataclasses import dataclass, field, asdict
-from typing import Optional
-import hashlib
-import uuid
-
-@dataclass
-class Experiment:
-    """Experiment metadata and results."""
-    experiment_id: str
-    name: str
-    params: dict
-    metrics: dict = field(default_factory=dict)
-    artifacts: list = field(default_factory=list)
-    tags: dict = field(default_factory=dict)
-    start_time: str = field(default_factory=lambda: datetime.utcnow().isoformat())
-    end_time: Optional[str] = None
-    status: str = "running"
-
-    def to_dict(self) -> dict:
-        return asdict(self)
-
-class SimpleTracker:
-    """Lightweight file-based experiment tracker."""
-
-    def __init__(self, experiments_dir: str = "./experiments"):
-        self.experiments_dir = Path(experiments_dir)
-        self.experiments_dir.mkdir(parents=True, exist_ok=True)
-        self.current_experiment: Optional[Experiment] = None
-
-    def start_experiment(
-        self,
-        name: str,
-        params: dict,
-        tags: dict = None,
-    ) -> Experiment:
-        """Start a new experiment."""
-        experiment_id = str(uuid.uuid4())[:8]
-
-        self.current_experiment = Experiment(
-            experiment_id=experiment_id,
-            name=name,
-            params=params,
-            tags=tags or {},
-        )
-
-        # Create experiment directory
-        exp_dir = self.experiments_dir / experiment_id
-        exp_dir.mkdir(exist_ok=True)
-
-        self._save_experiment()
-        return self.current_experiment
-
-    def log_metrics(self, metrics: dict, step: int = None) -> None:
-        """Log metrics to current experiment."""
-        if self.current_experiment is None:
-            raise ValueError("No active experiment")
-
-        for key, value in metrics.items():
-            if key not in self.current_experiment.metrics:
-                self.current_experiment.metrics[key] = []
-
-            self.current_experiment.metrics[key].append({
-                "value": value,
-                "step": step,
-                "timestamp": datetime.utcnow().isoformat(),
-            })
-
-        self._save_experiment()
-
-    def log_artifact(self, path: str, name: str = None) -> str:
-        """Copy artifact to experiment directory."""
-        if self.current_experiment is None:
-            raise ValueError("No active experiment")
-
-        import shutil
-
-        source = Path(path)
-        artifact_name = name or source.name
-        exp_dir = self.experiments_dir / self.current_experiment.experiment_id
-        dest = exp_dir / "artifacts" / artifact_name
-
-        dest.parent.mkdir(parents=True, exist_ok=True)
-
-        if source.is_dir():
-            shutil.copytree(source, dest)
-        else:
-            shutil.copy2(source, dest)
-
-        self.current_experiment.artifacts.append(str(dest))
-        self._save_experiment()
-
-        return str(dest)
-
-    def end_experiment(self, status: str = "completed") -> None:
-        """End current experiment."""
-        if self.current_experiment is None:
-            return
-
-        self.current_experiment.status = status
-        self.current_experiment.end_time = datetime.utcnow().isoformat()
-        self._save_experiment()
-        self.current_experiment = None
-
-    def _save_experiment(self) -> None:
-        """Save experiment to JSON file."""
-        if self.current_experiment is None:
-            return
-
-        exp_dir = self.experiments_dir / self.current_experiment.experiment_id
-        with open(exp_dir / "experiment.json", "w") as f:
-            json.dump(self.current_experiment.to_dict(), f, indent=2)
-
-    def load_experiment(self, experiment_id: str) -> Experiment:
-        """Load experiment by ID."""
-        exp_file = self.experiments_dir / experiment_id / "experiment.json"
-        with open(exp_file) as f:
-            data = json.load(f)
-        return Experiment(**data)
-
-    def list_experiments(self, tags: dict = None) -> list[Experiment]:
-        """List all experiments, optionally filtered by tags."""
-        experiments = []
-
-        for exp_dir in self.experiments_dir.iterdir():
-            if not exp_dir.is_dir():
-                continue
-
-            exp_file = exp_dir / "experiment.json"
-            if not exp_file.exists():
-                continue
-
-            exp = self.load_experiment(exp_dir.name)
-
-            if tags:
-                if not all(exp.tags.get(k) == v for k, v in tags.items()):
-                    continue
-
-            experiments.append(exp)
-
-        return sorted(experiments, key=lambda x: x.start_time, reverse=True)
-
-    def compare_experiments(self, experiment_ids: list[str]) -> dict:
-        """Compare metrics across experiments."""
-        comparison = {}
-
-        for exp_id in experiment_ids:
-            exp = self.load_experiment(exp_id)
-            comparison[exp_id] = {
-                "name": exp.name,
-                "params": exp.params,
-                "final_metrics": {
-                    k: v[-1]["value"] if v else None
-                    for k, v in exp.metrics.items()
-                },
-            }
-
-        return comparison
-```
-
----
-
-## Experiment Comparison and Analysis
-
-### Metrics Comparison
-
-```python
-import pandas as pd
-import matplotlib.pyplot as plt
-from mlflow.tracking import MlflowClient
-
-def compare_runs(
-    experiment_name: str,
-    metric_keys: list[str],
-    n_runs: int = 10,
-) -> pd.DataFrame:
-    """Compare recent runs in an experiment."""
-    client = MlflowClient()
-    experiment = client.get_experiment_by_name(experiment_name)
-
-    runs = client.search_runs(
-        experiment_ids=[experiment.experiment_id],
-        order_by=["start_time DESC"],
-        max_results=n_runs,
-    )
-
-    data = []
-    for run in runs:
-        row = {
-            "run_id": run.info.run_id,
-            "run_name": run.info.run_name,
-            "status": run.info.status,
-            "start_time": run.info.start_time,
-        }
-        row.update(run.data.params)
-        row.update({k: run.data.metrics.get(k) for k in metric_keys})
-        data.append(row)
-
-    return pd.DataFrame(data)
-
-def plot_metric_comparison(
-    runs_df: pd.DataFrame,
-    metric: str,
-    group_by: str = None,
-) -> plt.Figure:
-    """Plot metric comparison across runs."""
-    fig, ax = plt.subplots(figsize=(10, 6))
-
-    if group_by:
-        for group, group_df in runs_df.groupby(group_by):
-            ax.bar(group_df["run_name"], group_df[metric], label=str(group))
-        ax.legend(title=group_by)
-    else:
-        ax.bar(runs_df["run_name"], runs_df[metric])
-
-    ax.set_xlabel("Run")
-    ax.set_ylabel(metric)
-    ax.set_title(f"Comparison of {metric}")
-    plt.xticks(rotation=45, ha="right")
-    plt.tight_layout()
-
-    return fig
-```
-
----
-
-## Best Practices
-
-### What to Track
-
-```python
-# Always track:
-REQUIRED_PARAMS = [
-    "learning_rate",
-    "batch_size",
-    "epochs",
-    "model_architecture",
-    "optimizer",
-    "random_seed",
-    "dataset_version",
-]
-
-REQUIRED_METRICS = [
-    "train_loss",
-    "val_loss",
-    "train_accuracy",
-    "val_accuracy",
-]
-
-REQUIRED_ARTIFACTS = [
-    "model_checkpoint",
-    "training_config",
-    "requirements.txt",
-]
-
-# Recommended tags
-RECOMMENDED_TAGS = {
-    "author": "username",
-    "environment": "dev|staging|prod",
-    "model_type": "classification|regression|etc",
-    "dataset": "dataset_name",
-    "git_commit": "commit_hash",
-}
-```
-
-### Experiment Naming Conventions
-
-```python
-# Good naming patterns
-run_name = f"{model_type}_{dataset}_{timestamp}"
-run_name = f"exp_{experiment_number:03d}_{description}"
-run_name = f"{feature_flag}_{ablation_type}_{seed}"
-
-# Organize with groups and tags
-tags = {
-    "project": "recommendation_engine",
-    "sprint": "sprint_42",
-    "hypothesis": "larger_embedding_helps",
-}
-```
-
----
-
-## Related References
-
-- `training-pipelines.md` - Integrating tracking with training
-- `model-validation.md` - Validating tracked models
-- `pipeline-orchestration.md` - Tracking in automated pipelines
-
-## Cross-Reference Skills
-
-- **DevOps Engineer** - MLflow server deployment
-- **Data Engineer** - Artifact storage integration
+# Experiment Tracking
+
+---
+
+## Overview
+
+Experiment tracking enables reproducibility, comparison, and collaboration in ML development. It captures hyperparameters, metrics, artifacts, and model versions to ensure every experiment can be reproduced and compared.
+
+## When to Use This Reference
+
+- Setting up MLflow for experiment tracking
+- Implementing Weights & Biases integration
+- Creating model registries and versioning
+- Comparing experiments and selecting models
+- Building custom tracking solutions
+
+## When NOT to Use
+
+- Quick one-off experiments without reproducibility needs
+- Simple scripts without hyperparameters
+- Non-ML projects
+
+---
+
+## MLflow Integration
+
+### Basic Experiment Tracking
+
+```python
+import mlflow
+from mlflow.tracking import MlflowClient
+from pathlib import Path
+import json
+
+class MLflowTracker:
+    """MLflow experiment tracking wrapper."""
+
+    def __init__(
+        self,
+        experiment_name: str,
+        tracking_uri: str = "http://localhost:5000",
+        artifact_location: str = None,
+    ):
+        mlflow.set_tracking_uri(tracking_uri)
+
+        # Create or get experiment
+        experiment = mlflow.get_experiment_by_name(experiment_name)
+        if experiment is None:
+            self.experiment_id = mlflow.create_experiment(
+                experiment_name,
+                artifact_location=artifact_location,
+            )
+        else:
+            self.experiment_id = experiment.experiment_id
+
+        mlflow.set_experiment(experiment_name)
+        self.client = MlflowClient()
+        self.run = None
+
+    def start_run(
+        self,
+        run_name: str = None,
+        tags: dict = None,
+        nested: bool = False,
+    ) -> str:
+        """Start a new MLflow run."""
+        self.run = mlflow.start_run(
+            run_name=run_name,
+            experiment_id=self.experiment_id,
+            nested=nested,
+        )
+
+        if tags:
+            mlflow.set_tags(tags)
+
+        return self.run.info.run_id
+
+    def end_run(self, status: str = "FINISHED") -> None:
+        """End the current run."""
+        mlflow.end_run(status=status)
+        self.run = None
+
+    def log_params(self, params: dict) -> None:
+        """Log hyperparameters."""
+        mlflow.log_params(params)
+
+    def log_metrics(self, metrics: dict, step: int = None) -> None:
+        """Log metrics with optional step."""
+        for key, value in metrics.items():
+            mlflow.log_metric(key, value, step=step)
+
+    def log_artifact(self, local_path: str, artifact_path: str = None) -> None:
+        """Log file or directory as artifact."""
+        mlflow.log_artifact(local_path, artifact_path)
+
+    def log_model(
+        self,
+        model,
+        artifact_path: str,
+        registered_model_name: str = None,
+        signature=None,
+        input_example=None,
+    ) -> str:
+        """Log model with optional registration."""
+        from mlflow.models import infer_signature
+
+        if signature is None and input_example is not None:
+            signature = infer_signature(input_example, model.predict(input_example))
+
+        model_info = mlflow.sklearn.log_model(
+            model,
+            artifact_path=artifact_path,
+            registered_model_name=registered_model_name,
+            signature=signature,
+            input_example=input_example,
+        )
+
+        return model_info.model_uri
+
+# Usage example
+def train_with_mlflow(
+    model,
+    X_train,
+    y_train,
+    X_val,
+    y_val,
+    params: dict,
+):
+    """Complete training run with MLflow tracking."""
+    tracker = MLflowTracker("my_experiment")
+
+    tracker.start_run(
+        run_name=f"run_{params['model_type']}",
+        tags={
+            "model_type": params["model_type"],
+            "dataset_version": "v1.0",
+            "author": "ml-team",
+        },
+    )
+
+    try:
+        # Log parameters
+        tracker.log_params(params)
+
+        # Train model
+        model.fit(X_train, y_train)
+
+        # Evaluate and log metrics
+        train_score = model.score(X_train, y_train)
+        val_score = model.score(X_val, y_val)
+
+        tracker.log_metrics({
+            "train_accuracy": train_score,
+            "val_accuracy": val_score,
+        })
+
+        # Log model
+        model_uri = tracker.log_model(
+            model,
+            artifact_path="model",
+            registered_model_name="my_model",
+            input_example=X_train[:5],
+        )
+
+        tracker.end_run()
+        return model_uri
+
+    except Exception as e:
+        tracker.end_run(status="FAILED")
+        raise
+```
+
+### PyTorch Model Logging
+
+```python
+import mlflow.pytorch
+import torch
+
+def log_pytorch_model(
+    model: torch.nn.Module,
+    artifact_path: str,
+    registered_model_name: str = None,
+    sample_input: torch.Tensor = None,
+) -> str:
+    """Log PyTorch model with signature inference."""
+    from mlflow.models import infer_signature
+
+    # Create signature from sample input
+    signature = None
+    if sample_input is not None:
+        model.eval()
+        with torch.no_grad():
+            sample_output = model(sample_input)
+
+        signature = infer_signature(
+            sample_input.numpy(),
+            sample_output.numpy(),
+        )
+
+    model_info = mlflow.pytorch.log_model(
+        model,
+        artifact_path=artifact_path,
+        registered_model_name=registered_model_name,
+        signature=signature,
+    )
+
+    return model_info.model_uri
+
+def load_pytorch_model(model_uri: str, device: str = "cpu") -> torch.nn.Module:
+    """Load PyTorch model from MLflow."""
+    model = mlflow.pytorch.load_model(model_uri, map_location=device)
+    return model
+```
+
+### Model Registry Operations
+
+```python
+from mlflow.tracking import MlflowClient
+from mlflow.entities.model_registry import ModelVersion
+
+class ModelRegistry:
+    """MLflow Model Registry wrapper."""
+
+    def __init__(self, tracking_uri: str = "http://localhost:5000"):
+        mlflow.set_tracking_uri(tracking_uri)
+        self.client = MlflowClient()
+
+    def register_model(
+        self,
+        model_uri: str,
+        name: str,
+        tags: dict = None,
+        description: str = None,
+    ) -> ModelVersion:
+        """Register a new model version."""
+        result = mlflow.register_model(model_uri, name)
+
+        if tags:
+            for key, value in tags.items():
+                self.client.set_model_version_tag(name, result.version, key, value)
+
+        if description:
+            self.client.update_model_version(
+                name,
+                result.version,
+                description=description,
+            )
+
+        return result
+
+    def transition_model_stage(
+        self,
+        name: str,
+        version: str,
+        stage: str,
+        archive_existing: bool = True,
+    ) -> ModelVersion:
+        """Transition model to new stage (Staging, Production, Archived)."""
+        return self.client.transition_model_version_stage(
+            name=name,
+            version=version,
+            stage=stage,
+            archive_existing_versions=archive_existing,
+        )
+
+    def get_latest_version(
+        self,
+        name: str,
+        stages: list[str] = None,
+    ) -> list[ModelVersion]:
+        """Get latest model versions by stage."""
+        return self.client.get_latest_versions(name, stages=stages)
+
+    def load_production_model(self, name: str) -> any:
+        """Load the production model."""
+        model_uri = f"models:/{name}/Production"
+        return mlflow.pyfunc.load_model(model_uri)
+
+    def compare_versions(
+        self,
+        name: str,
+        version_a: str,
+        version_b: str,
+    ) -> dict:
+        """Compare two model versions."""
+        v_a = self.client.get_model_version(name, version_a)
+        v_b = self.client.get_model_version(name, version_b)
+
+        run_a = self.client.get_run(v_a.run_id)
+        run_b = self.client.get_run(v_b.run_id)
+
+        return {
+            "version_a": {
+                "version": version_a,
+                "metrics": run_a.data.metrics,
+                "params": run_a.data.params,
+            },
+            "version_b": {
+                "version": version_b,
+                "metrics": run_b.data.metrics,
+                "params": run_b.data.params,
+            },
+        }
+```
+
+---
+
+## Weights & Biases Integration
+
+### Basic W&B Tracking
+
+```python
+import wandb
+from pathlib import Path
+
+class WandbTracker:
+    """Weights & Biases experiment tracking wrapper."""
+
+    def __init__(
+        self,
+        project: str,
+        entity: str = None,
+        config: dict = None,
+    ):
+        self.project = project
+        self.entity = entity
+        self.config = config
+        self.run = None
+
+    def start_run(
+        self,
+        name: str = None,
+        tags: list[str] = None,
+        group: str = None,
+        job_type: str = "train",
+        resume: str = None,
+    ) -> wandb.Run:
+        """Initialize W&B run."""
+        self.run = wandb.init(
+            project=self.project,
+            entity=self.entity,
+            name=name,
+            config=self.config,
+            tags=tags,
+            group=group,
+            job_type=job_type,
+            resume=resume,
+        )
+        return self.run
+
+    def log(self, data: dict, step: int = None, commit: bool = True) -> None:
+        """Log metrics and data."""
+        wandb.log(data, step=step, commit=commit)
+
+    def log_artifact(
+        self,
+        name: str,
+        artifact_type: str,
+        path: str,
+        metadata: dict = None,
+    ) -> wandb.Artifact:
+        """Log artifact (model, dataset, etc.)."""
+        artifact = wandb.Artifact(
+            name=name,
+            type=artifact_type,
+            metadata=metadata,
+        )
+
+        if Path(path).is_dir():
+            artifact.add_dir(path)
+        else:
+            artifact.add_file(path)
+
+        self.run.log_artifact(artifact)
+        return artifact
+
+    def log_model(
+        self,
+        model_path: str,
+        name: str,
+        metadata: dict = None,
+        aliases: list[str] = None,
+    ) -> wandb.Artifact:
+        """Log model artifact with aliases."""
+        artifact = wandb.Artifact(
+            name=name,
+            type="model",
+            metadata=metadata,
+        )
+
+        if Path(model_path).is_dir():
+            artifact.add_dir(model_path)
+        else:
+            artifact.add_file(model_path)
+
+        self.run.log_artifact(artifact, aliases=aliases or ["latest"])
+        return artifact
+
+    def watch_model(
+        self,
+        model,
+        log: str = "all",
+        log_freq: int = 100,
+    ) -> None:
+        """Watch model for gradient and parameter logging."""
+        wandb.watch(model, log=log, log_freq=log_freq)
+
+    def finish(self, exit_code: int = 0) -> None:
+        """Finish the run."""
+        wandb.finish(exit_code=exit_code)
+
+# Usage with PyTorch
+def train_with_wandb(
+    model: torch.nn.Module,
+    train_loader,
+    val_loader,
+    config: dict,
+):
+    """Training with W&B tracking."""
+    tracker = WandbTracker(
+        project="my-project",
+        config=config,
+    )
+
+    tracker.start_run(
+        name=f"experiment_{config['model_type']}",
+        tags=["baseline", config["model_type"]],
+        group="hyperparameter_search",
+    )
+
+    # Watch model gradients
+    tracker.watch_model(model)
+
+    for epoch in range(config["epochs"]):
+        model.train()
+        for batch_idx, (data, target) in enumerate(train_loader):
+            # Training step
+            loss = train_step(model, data, target)
+
+            tracker.log({
+                "train/loss": loss,
+                "train/epoch": epoch,
+            })
+
+        # Validation
+        val_metrics = evaluate(model, val_loader)
+        tracker.log({
+            "val/loss": val_metrics["loss"],
+            "val/accuracy": val_metrics["accuracy"],
+            "epoch": epoch,
+        })
+
+    # Save and log model
+    torch.save(model.state_dict(), "model.pt")
+    tracker.log_model(
+        "model.pt",
+        name="trained_model",
+        metadata={"accuracy": val_metrics["accuracy"]},
+        aliases=["latest", "best"],
+    )
+
+    tracker.finish()
+```
+
+### W&B Sweeps for Hyperparameter Tuning
+
+```python
+import wandb
+
+sweep_config = {
+    "method": "bayes",  # bayes, grid, random
+    "metric": {
+        "name": "val/loss",
+        "goal": "minimize",
+    },
+    "parameters": {
+        "learning_rate": {
+            "distribution": "log_uniform_values",
+            "min": 1e-5,
+            "max": 1e-2,
+        },
+        "batch_size": {
+            "values": [16, 32, 64, 128],
+        },
+        "hidden_size": {
+            "values": [128, 256, 512],
+        },
+        "dropout": {
+            "distribution": "uniform",
+            "min": 0.1,
+            "max": 0.5,
+        },
+    },
+    "early_terminate": {
+        "type": "hyperband",
+        "min_iter": 3,
+    },
+}
+
+def sweep_train():
+    """Training function for sweep."""
+    with wandb.init() as run:
+        config = wandb.config
+
+        model = build_model(
+            hidden_size=config.hidden_size,
+            dropout=config.dropout,
+        )
+
+        optimizer = torch.optim.Adam(
+            model.parameters(),
+            lr=config.learning_rate,
+        )
+
+        train_loader = DataLoader(train_dataset, batch_size=config.batch_size)
+
+        for epoch in range(10):
+            loss = train_epoch(model, train_loader, optimizer)
+            val_loss = evaluate(model, val_loader)
+
+            wandb.log({
+                "train/loss": loss,
+                "val/loss": val_loss,
+                "epoch": epoch,
+            })
+
+# Run sweep
+sweep_id = wandb.sweep(sweep_config, project="my-project")
+wandb.agent(sweep_id, function=sweep_train, count=50)
+```
+
+---
+
+## Custom Experiment Tracking
+
+### Lightweight Tracker
+
+```python
+import json
+from datetime import datetime
+from pathlib import Path
+from dataclasses import dataclass, field, asdict
+from typing import Optional
+import hashlib
+import uuid
+
+@dataclass
+class Experiment:
+    """Experiment metadata and results."""
+    experiment_id: str
+    name: str
+    params: dict
+    metrics: dict = field(default_factory=dict)
+    artifacts: list = field(default_factory=list)
+    tags: dict = field(default_factory=dict)
+    start_time: str = field(default_factory=lambda: datetime.utcnow().isoformat())
+    end_time: Optional[str] = None
+    status: str = "running"
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+class SimpleTracker:
+    """Lightweight file-based experiment tracker."""
+
+    def __init__(self, experiments_dir: str = "./experiments"):
+        self.experiments_dir = Path(experiments_dir)
+        self.experiments_dir.mkdir(parents=True, exist_ok=True)
+        self.current_experiment: Optional[Experiment] = None
+
+    def start_experiment(
+        self,
+        name: str,
+        params: dict,
+        tags: dict = None,
+    ) -> Experiment:
+        """Start a new experiment."""
+        experiment_id = str(uuid.uuid4())[:8]
+
+        self.current_experiment = Experiment(
+            experiment_id=experiment_id,
+            name=name,
+            params=params,
+            tags=tags or {},
+        )
+
+        # Create experiment directory
+        exp_dir = self.experiments_dir / experiment_id
+        exp_dir.mkdir(exist_ok=True)
+
+        self._save_experiment()
+        return self.current_experiment
+
+    def log_metrics(self, metrics: dict, step: int = None) -> None:
+        """Log metrics to current experiment."""
+        if self.current_experiment is None:
+            raise ValueError("No active experiment")
+
+        for key, value in metrics.items():
+            if key not in self.current_experiment.metrics:
+                self.current_experiment.metrics[key] = []
+
+            self.current_experiment.metrics[key].append({
+                "value": value,
+                "step": step,
+                "timestamp": datetime.utcnow().isoformat(),
+            })
+
+        self._save_experiment()
+
+    def log_artifact(self, path: str, name: str = None) -> str:
+        """Copy artifact to experiment directory."""
+        if self.current_experiment is None:
+            raise ValueError("No active experiment")
+
+        import shutil
+
+        source = Path(path)
+        artifact_name = name or source.name
+        exp_dir = self.experiments_dir / self.current_experiment.experiment_id
+        dest = exp_dir / "artifacts" / artifact_name
+
+        dest.parent.mkdir(parents=True, exist_ok=True)
+
+        if source.is_dir():
+            shutil.copytree(source, dest)
+        else:
+            shutil.copy2(source, dest)
+
+        self.current_experiment.artifacts.append(str(dest))
+        self._save_experiment()
+
+        return str(dest)
+
+    def end_experiment(self, status: str = "completed") -> None:
+        """End current experiment."""
+        if self.current_experiment is None:
+            return
+
+        self.current_experiment.status = status
+        self.current_experiment.end_time = datetime.utcnow().isoformat()
+        self._save_experiment()
+        self.current_experiment = None
+
+    def _save_experiment(self) -> None:
+        """Save experiment to JSON file."""
+        if self.current_experiment is None:
+            return
+
+        exp_dir = self.experiments_dir / self.current_experiment.experiment_id
+        with open(exp_dir / "experiment.json", "w") as f:
+            json.dump(self.current_experiment.to_dict(), f, indent=2)
+
+    def load_experiment(self, experiment_id: str) -> Experiment:
+        """Load experiment by ID."""
+        exp_file = self.experiments_dir / experiment_id / "experiment.json"
+        with open(exp_file) as f:
+            data = json.load(f)
+        return Experiment(**data)
+
+    def list_experiments(self, tags: dict = None) -> list[Experiment]:
+        """List all experiments, optionally filtered by tags."""
+        experiments = []
+
+        for exp_dir in self.experiments_dir.iterdir():
+            if not exp_dir.is_dir():
+                continue
+
+            exp_file = exp_dir / "experiment.json"
+            if not exp_file.exists():
+                continue
+
+            exp = self.load_experiment(exp_dir.name)
+
+            if tags:
+                if not all(exp.tags.get(k) == v for k, v in tags.items()):
+                    continue
+
+            experiments.append(exp)
+
+        return sorted(experiments, key=lambda x: x.start_time, reverse=True)
+
+    def compare_experiments(self, experiment_ids: list[str]) -> dict:
+        """Compare metrics across experiments."""
+        comparison = {}
+
+        for exp_id in experiment_ids:
+            exp = self.load_experiment(exp_id)
+            comparison[exp_id] = {
+                "name": exp.name,
+                "params": exp.params,
+                "final_metrics": {
+                    k: v[-1]["value"] if v else None
+                    for k, v in exp.metrics.items()
+                },
+            }
+
+        return comparison
+```
+
+---
+
+## Experiment Comparison and Analysis
+
+### Metrics Comparison
+
+```python
+import pandas as pd
+import matplotlib.pyplot as plt
+from mlflow.tracking import MlflowClient
+
+def compare_runs(
+    experiment_name: str,
+    metric_keys: list[str],
+    n_runs: int = 10,
+) -> pd.DataFrame:
+    """Compare recent runs in an experiment."""
+    client = MlflowClient()
+    experiment = client.get_experiment_by_name(experiment_name)
+
+    runs = client.search_runs(
+        experiment_ids=[experiment.experiment_id],
+        order_by=["start_time DESC"],
+        max_results=n_runs,
+    )
+
+    data = []
+    for run in runs:
+        row = {
+            "run_id": run.info.run_id,
+            "run_name": run.info.run_name,
+            "status": run.info.status,
+            "start_time": run.info.start_time,
+        }
+        row.update(run.data.params)
+        row.update({k: run.data.metrics.get(k) for k in metric_keys})
+        data.append(row)
+
+    return pd.DataFrame(data)
+
+def plot_metric_comparison(
+    runs_df: pd.DataFrame,
+    metric: str,
+    group_by: str = None,
+) -> plt.Figure:
+    """Plot metric comparison across runs."""
+    fig, ax = plt.subplots(figsize=(10, 6))
+
+    if group_by:
+        for group, group_df in runs_df.groupby(group_by):
+            ax.bar(group_df["run_name"], group_df[metric], label=str(group))
+        ax.legend(title=group_by)
+    else:
+        ax.bar(runs_df["run_name"], runs_df[metric])
+
+    ax.set_xlabel("Run")
+    ax.set_ylabel(metric)
+    ax.set_title(f"Comparison of {metric}")
+    plt.xticks(rotation=45, ha="right")
+    plt.tight_layout()
+
+    return fig
+```
+
+---
+
+## Best Practices
+
+### What to Track
+
+```python
+# Always track:
+REQUIRED_PARAMS = [
+    "learning_rate",
+    "batch_size",
+    "epochs",
+    "model_architecture",
+    "optimizer",
+    "random_seed",
+    "dataset_version",
+]
+
+REQUIRED_METRICS = [
+    "train_loss",
+    "val_loss",
+    "train_accuracy",
+    "val_accuracy",
+]
+
+REQUIRED_ARTIFACTS = [
+    "model_checkpoint",
+    "training_config",
+    "requirements.txt",
+]
+
+# Recommended tags
+RECOMMENDED_TAGS = {
+    "author": "username",
+    "environment": "dev|staging|prod",
+    "model_type": "classification|regression|etc",
+    "dataset": "dataset_name",
+    "git_commit": "commit_hash",
+}
+```
+
+### Experiment Naming Conventions
+
+```python
+# Good naming patterns
+run_name = f"{model_type}_{dataset}_{timestamp}"
+run_name = f"exp_{experiment_number:03d}_{description}"
+run_name = f"{feature_flag}_{ablation_type}_{seed}"
+
+# Organize with groups and tags
+tags = {
+    "project": "recommendation_engine",
+    "sprint": "sprint_42",
+    "hypothesis": "larger_embedding_helps",
+}
+```
+
+---
+
+## Related References
+
+- `training-pipelines.md` - Integrating tracking with training
+- `model-validation.md` - Validating tracked models
+- `pipeline-orchestration.md` - Tracking in automated pipelines
+
+## Cross-Reference Skills
+
+- **DevOps Engineer** - MLflow server deployment
+- **Data Engineer** - Artifact storage integration