Experiment Tracker in Machine Learning

import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, f1_score
from sklearn.model_selection import train_test_split

# Set the experiment name (creates if not exists)
mlflow.set_experiment("fraud-detection-v2")

# Load and split data
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# Start a tracked run
with mlflow.start_run(run_name="rf-baseline"):
    # Log hyperparameters
    params = {
        "n_estimators": 200,
        "max_depth": 10,
        "min_samples_split": 5,
        "class_weight": "balanced",
        "random_state": 42
    }
    mlflow.log_params(params)
    
    # Log data metadata
    mlflow.log_param("train_size", len(X_train))
    mlflow.log_param("test_size", len(X_test))
    mlflow.log_param("positive_ratio", float(y_train.mean()))
    
    # Train model
    model = RandomForestClassifier(**params)
    model.fit(X_train, y_train)
    
    # Log metrics
    y_pred = model.predict(X_test)
    mlflow.log_metric("accuracy", accuracy_score(y_test, y_pred))
    mlflow.log_metric("f1_score", f1_score(y_test, y_pred))
    mlflow.log_metric("precision", precision_score(y_test, y_pred))
    mlflow.log_metric("recall", recall_score(y_test, y_pred))
    
    # Log the model as an artifact
    mlflow.sklearn.log_model(model, "model")
    
    # Log additional artifacts
    mlflow.log_artifact("feature_importance.png")
    mlflow.log_artifact("confusion_matrix.png")
    
    print(f"Run ID: {mlflow.active_run().info.run_id}")

import wandb
import torch

# Initialize a W&B run with config
config = {
    "architecture": "transformer",
    "hidden_dim": 256,
    "num_heads": 8,
    "num_layers": 6,
    "learning_rate": 3e-4,
    "batch_size": 32,
    "epochs": 50,
    "optimizer": "adamw",
    "weight_decay": 0.01,
    "scheduler": "cosine_warmup",
    "warmup_steps": 1000
}

run = wandb.init(
    project="text-classification",
    name="transformer-v3-cosine",
    config=config,
    tags=["transformer", "production-candidate"],
    notes="Testing cosine warmup scheduler with AdamW"
)

# Training loop with step-level logging
for epoch in range(config["epochs"]):
    model.train()
    epoch_loss = 0
    for batch_idx, (inputs, labels) in enumerate(train_loader):
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        scheduler.step()
        
        # Log every step for detailed curves
        wandb.log({
            "train/batch_loss": loss.item(),
            "train/learning_rate": scheduler.get_last_lr()[0],
            "train/global_step": epoch * len(train_loader) + batch_idx
        })
        epoch_loss += loss.item()
    
    # Log epoch-level metrics
    val_loss, val_acc, val_f1 = evaluate(model, val_loader)
    wandb.log({
        "train/epoch_loss": epoch_loss / len(train_loader),
        "val/loss": val_loss,
        "val/accuracy": val_acc,
        "val/f1": val_f1,
        "epoch": epoch
    })
    
    # Log model checkpoint as artifact
    if val_f1 > best_f1:
        best_f1 = val_f1
        artifact = wandb.Artifact(
            f"model-checkpoint", type="model",
            metadata={"val_f1": val_f1, "epoch": epoch}
        )
        torch.save(model.state_dict(), "best_model.pt")
        artifact.add_file("best_model.pt")
        run.log_artifact(artifact)

# Log final summary metrics
wandb.summary["best_val_f1"] = best_f1
wandb.finish()

import neptune
import optuna
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.metrics import roc_auc_score

def objective(trial):
    # Create a Neptune run for each Optuna trial
    run = neptune.init_run(
        project="team/fraud-detection",
        tags=["optuna-sweep", "gbm"]
    )
    
    # Sample hyperparameters
    params = {
        "n_estimators": trial.suggest_int("n_estimators", 100, 1000),
        "max_depth": trial.suggest_int("max_depth", 3, 12),
        "learning_rate": trial.suggest_float("learning_rate", 1e-4, 0.3, log=True),
        "subsample": trial.suggest_float("subsample", 0.6, 1.0),
        "min_samples_leaf": trial.suggest_int("min_samples_leaf", 1, 50)
    }
    
    # Log all parameters
    run["parameters"] = params
    run["trial_number"] = trial.number
    
    # Train and evaluate
    model = GradientBoostingClassifier(**params)
    model.fit(X_train, y_train)
    
    y_prob = model.predict_proba(X_test)[:, 1]
    auc = roc_auc_score(y_test, y_prob)
    
    # Log metrics and feature importance
    run["metrics/auc"] = auc
    run["metrics/train_auc"] = roc_auc_score(
        y_train, model.predict_proba(X_train)[:, 1]
    )
    
    # Log feature importance as a series
    for i, imp in enumerate(model.feature_importances_):
        run[f"feature_importance/{feature_names[i]}"] = imp
    
    run.stop()
    return auc

# Run Optuna study with Neptune tracking
study = optuna.create_study(direction="maximize")
study.optimize(objective, n_trials=100)

print(f"Best AUC: {study.best_value:.4f}")
print(f"Best params: {study.best_params}")

import mlflow
from mlflow.tracking import MlflowClient

client = MlflowClient()

def find_best_run(experiment_name, metric="f1_score", min_runs=5):
    experiment = client.get_experiment_by_name(experiment_name)
    if not experiment:
        raise ValueError(f"Experiment '{experiment_name}' not found")
    
    runs = client.search_runs(
        experiment_ids=[experiment.experiment_id],
        filter_string="status = 'FINISHED'",
        order_by=[f"metrics.{metric} DESC"],
        max_results=100
    )
    
    if len(runs) < min_runs:
        raise ValueError(f"Only {len(runs)} finished runs; need {min_runs}")
    
    best = runs[0]
    print(f"Best run: {best.info.run_id}")
    print(f"  {metric}: {best.data.metrics[metric]:.4f}")
    print(f"  Parameters: {best.data.params}")
    return best

def promote_to_registry(run, model_name, stage="Staging"):
    model_uri = f"runs:/{run.info.run_id}/model"
    
    # Register the model (creates a new version)
    mv = mlflow.register_model(model_uri, model_name)
    
    # Transition to the target stage
    client.transition_model_version_stage(
        name=model_name,
        version=mv.version,
        stage=stage
    )
    
    # Tag with lineage info
    client.set_model_version_tag(
        name=model_name,
        version=mv.version,
        key="source_experiment",
        value=run.info.experiment_id
    )
    
    print(f"Registered {model_name} v{mv.version} -> {stage}")
    return mv

# Usage in CI/CD pipeline
best_run = find_best_run("fraud-detection-v2", metric="f1_score")
promote_to_registry(best_run, "fraud-detector", stage="Staging")

# MLflow tracking server configuration
# docker-compose.yml
version: '3.8'
services:
  mlflow-server:
    image: ghcr.io/mlflow/mlflow:v2.12.0
    command: >
      mlflow server
      --backend-store-uri postgresql://mlflow:password@postgres:5432/mlflow
      --default-artifact-root s3://mlflow-artifacts/
      --host 0.0.0.0
      --port 5000
      --workers 4
    ports:
      - "5000:5000"
    environment:
    depends_on:
      - postgres

  postgres:
    image: postgres:16
    environment:
      POSTGRES_DB: mlflow
      POSTGRES_USER: mlflow
      POSTGRES_PASSWORD: password
    volumes:
      - pgdata:/var/lib/postgresql/data

volumes:
  pgdata: