Hyperparameter Sweeps with Optuna

What We're Building

An end-to-end hyperparameter optimization system that uses Optuna to intelligently search the parameter space while dynamically provisioning and releasing Clore GPUs for each trial. Each trial gets its own rented GPU, runs in isolation, and is pruned early if it's clearly not competitive — saving real money.

Key Features:

Full Optuna integration with Clore GPU provisioning
Parallel trial execution across multiple rented GPUs
MedianPruner to kill underperforming trials early
PostgreSQL/SQLite database backend for distributed studies
Per-trial cost tracking with hard budget cap
Results visualization (HTML report + terminal summary)
Proper error handling and retry logic

Prerequisites

Clore.ai account with API key (get one here)
Python 3.10+

pip install optuna optuna-dashboard requests paramiko rich
# For PostgreSQL backend (recommended for parallel runs):
pip install psycopg2-binary
# For plotting:
pip install plotly kaleido

Architecture Overview

┌─────────────────────────────────────────────────────────────────┐
│                  Hyperparameter Sweep System                     │
├─────────────────────────────────────────────────────────────────┤
│  ┌─────────────────────────────────────────────────────────┐   │
│  │                   Optuna Study                           │   │
│  │   TPESampler + MedianPruner  ←→  PostgreSQL/SQLite DB   │   │
│  └─────────────────────────────────────────────────────────┘   │
│                │                                               │
│      ┌─────────┴──────────────────────┐                        │
│      ▼                                ▼                        │
│  ┌──────────┐                   ┌──────────┐                   │
│  │ Trial 1  │ …                 │ Trial N  │  (parallel)       │
│  └──────────┘                   └──────────┘                   │
│      │                                │                        │
│      ▼                                ▼                        │
│  ┌──────────────────────────────────────────────────────────┐  │
│  │                  Clore GPU Provisioner                   │  │
│  │   /v1/marketplace → /v1/create_order → /v1/cancel_order  │  │
│  └──────────────────────────────────────────────────────────┘  │
│      │                                │                        │
│      ▼                                ▼                        │
│  ┌──────────────┐            ┌─────────────────┐               │
│  │ Cost Tracker │            │  Results Store  │               │
│  │ (per-trial)  │            │  (JSON + plots) │               │
│  └──────────────┘            └─────────────────┘               │
└─────────────────────────────────────────────────────────────────┘

Step 1: Clore GPU Provisioner for Trials

Each Optuna trial provisions its own Clore server, runs training, streams the result back, then releases the server.

# trial_provisioner.py
"""Provision a Clore GPU for a single Optuna trial."""

import time
import json
import logging
import paramiko
import requests
from typing import Dict, Optional, Tuple

logger = logging.getLogger(__name__)


class TrialProvisioner:
    """Rent a Clore GPU, run a training command, return metrics, release GPU."""

    BASE_URL = "https://api.clore.ai"

    def __init__(
        self,
        api_key:      str,
        gpu_type:     str   = "RTX 3080",
        max_price:    float = 0.30,
        docker_image: str   = "pytorch/pytorch:2.1.0-cuda12.1-cudnn8-devel",
        ssh_password: str   = "OptunaTrial1!",
        timeout_sec:  int   = 3600,
    ):
        self.api_key      = api_key
        self.headers      = {"auth": api_key}
        self.gpu_type     = gpu_type
        self.max_price    = max_price
        self.docker_image = docker_image
        self.ssh_password = ssh_password
        self.timeout_sec  = timeout_sec

        self.order_id:    Optional[int]  = None
        self.server_id:   Optional[int]  = None
        self.price_per_hr: float         = 0.0
        self.start_time:  Optional[float] = None

    # ── API helpers ───────────────────────────────────────────────

    def _request(self, method: str, endpoint: str, **kwargs) -> Dict:
        for attempt in range(3):
            try:
                r = requests.request(
                    method,
                    f"{self.BASE_URL}{endpoint}",
                    headers=self.headers,
                    timeout=30,
                    **kwargs,
                )
                data = r.json()
                if data.get("code") != 0:
                    raise RuntimeError(f"API {endpoint}: {data}")
                return data
            except Exception as e:
                if attempt == 2:
                    raise
                time.sleep(2 ** attempt)
        raise RuntimeError("unreachable")

    # ── marketplace ───────────────────────────────────────────────

    def find_server(self, exclude_ids: Optional[list] = None) -> Dict:
        """Return the cheapest reliable server matching our requirements."""
        exclude_ids = exclude_ids or []
        servers = self._request("GET", "/v1/marketplace").get("servers", [])

        candidates = []
        for s in servers:
            if s.get("rented"):
                continue
            if s["id"] in exclude_ids:
                continue
            gpus = s.get("gpu_array", [])
            if not any(self.gpu_type.lower() in g.lower() for g in gpus):
                continue
            price = (
                s.get("price", {})
                 .get("usd", {})
                 .get("on_demand_clore", 9999)
            )
            if price > self.max_price:
                continue
            reliability = s.get("reliability", 0)
            if reliability < 0.75:
                continue
            candidates.append({"id": s["id"], "price": price,
                                "reliability": reliability, "gpus": gpus})

        if not candidates:
            raise RuntimeError(
                f"No {self.gpu_type} server available below ${self.max_price}/hr"
            )

        candidates.sort(key=lambda x: x["price"])
        return candidates[0]

    # ── lifecycle ─────────────────────────────────────────────────

    def provision(self, env: Optional[Dict[str, str]] = None) -> Tuple[str, int]:
        """Rent a server. Returns (ssh_host, ssh_port)."""
        server = self.find_server()
        self.server_id    = server["id"]
        self.price_per_hr = server["price"]

        data = self._request("POST", "/v1/create_order", json={
            "renting_server": server["id"],
            "type":           "on-demand",
            "currency":       "CLORE-Blockchain",
            "image":          self.docker_image,
            "ports":          {"22": "tcp"},
            "env":            env or {},
            "jupyter_token":  "",
            "ssh_password":   self.ssh_password,
        })
        self.order_id  = data["order_id"]
        self.start_time = time.time()

        logger.info(f"  Order {self.order_id} on server {self.server_id} "
                    f"@ ${self.price_per_hr:.3f}/hr — waiting...")

        deadline = time.time() + 300
        while time.time() < deadline:
            orders = self._request("GET", "/v1/my_orders").get("orders", [])
            order  = next(
                (o for o in orders if o.get("order_id") == self.order_id), None
            )
            if order and order.get("status") == "running":
                ssh_raw = order.get("connection", {}).get("ssh", "")
                parts   = ssh_raw.split()
                host    = parts[1].split("@")[1] if len(parts) > 1 else "127.0.0.1"
                port    = int(parts[3]) if len(parts) > 3 else 22
                logger.info(f"  Server ready: {host}:{port}")
                return host, port
            time.sleep(5)

        raise TimeoutError(f"Order {self.order_id} did not start in 300s")

    def run_command(self, host: str, port: int, command: str) -> str:
        """SSH into the server, run command, return stdout."""
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

        # Retry SSH connection (server may need a moment after status=running)
        for attempt in range(10):
            try:
                client.connect(
                    host, port=port,
                    username="root",
                    password=self.ssh_password,
                    timeout=30,
                    banner_timeout=60,
                )
                break
            except Exception as e:
                if attempt == 9:
                    raise
                logger.debug(f"SSH attempt {attempt+1}: {e}")
                time.sleep(15)

        _, stdout, stderr = client.exec_command(command, timeout=self.timeout_sec)
        output    = stdout.read().decode("utf-8", errors="replace")
        err_text  = stderr.read().decode("utf-8", errors="replace")
        exit_code = stdout.channel.recv_exit_status()
        client.close()

        if exit_code != 0:
            raise RuntimeError(
                f"Command exited {exit_code}.\nSTDOUT:\n{output}\nSTDERR:\n{err_text}"
            )
        return output

    def cost_so_far(self) -> float:
        """Return cost accrued since provisioning (USD)."""
        if self.start_time is None:
            return 0.0
        elapsed_hr = (time.time() - self.start_time) / 3600
        return elapsed_hr * self.price_per_hr

    def release(self):
        """Cancel the order."""
        if self.order_id is not None:
            try:
                self._request("POST", "/v1/cancel_order",
                              json={"id": self.order_id})
                logger.info(f"  Released order {self.order_id} "
                            f"(cost: ${self.cost_so_far():.4f})")
            except Exception as e:
                logger.warning(f"  Failed to cancel {self.order_id}: {e}")
            finally:
                self.order_id = None

Step 2: Cost Tracker

# cost_tracker.py
"""Track per-trial and cumulative costs for a hyperparameter study."""

import json
import threading
from pathlib import Path
from typing import Dict, List, Optional
from dataclasses import dataclass, field, asdict
from datetime import datetime


@dataclass
class TrialCost:
    trial_number: int
    order_id:     int
    gpu_type:     str
    price_per_hr: float
    start_time:   str
    end_time:     str   = ""
    cost_usd:     float = 0.0
    result:       float = float("nan")
    pruned:       bool  = False
    error:        str   = ""


class CostTracker:
    """Thread-safe per-trial cost tracker with budget enforcement."""

    def __init__(
        self,
        budget_usd: float,
        output_path: str = "sweep_costs.json",
    ):
        self.budget_usd  = budget_usd
        self.output_path = Path(output_path)
        self._costs: Dict[int, TrialCost] = {}
        self._lock = threading.Lock()
        self._load()

    def _load(self):
        if self.output_path.exists():
            with open(self.output_path) as f:
                raw = json.load(f)
            for item in raw:
                tc = TrialCost(**item)
                self._costs[tc.trial_number] = tc

    def _save(self):
        data = [asdict(tc) for tc in self._costs.values()]
        with open(self.output_path, "w") as f:
            json.dump(data, f, indent=2)

    def start_trial(
        self,
        trial_number: int,
        order_id:     int,
        gpu_type:     str,
        price_per_hr: float,
    ):
        with self._lock:
            self._costs[trial_number] = TrialCost(
                trial_number=trial_number,
                order_id=order_id,
                gpu_type=gpu_type,
                price_per_hr=price_per_hr,
                start_time=datetime.utcnow().isoformat(),
            )
            self._save()

    def finish_trial(
        self,
        trial_number: int,
        cost_usd:     float,
        result:       float,
        pruned:       bool  = False,
        error:        str   = "",
    ):
        with self._lock:
            if trial_number in self._costs:
                tc = self._costs[trial_number]
                tc.end_time  = datetime.utcnow().isoformat()
                tc.cost_usd  = cost_usd
                tc.result    = result
                tc.pruned    = pruned
                tc.error     = error
            self._save()

    def total_cost(self) -> float:
        with self._lock:
            return sum(tc.cost_usd for tc in self._costs.values())

    def budget_remaining(self) -> float:
        return max(0.0, self.budget_usd - self.total_cost())

    def over_budget(self) -> bool:
        return self.total_cost() >= self.budget_usd

    def summary(self) -> Dict:
        with self._lock:
            costs = list(self._costs.values())
        return {
            "total_trials":   len(costs),
            "pruned_trials":  sum(1 for c in costs if c.pruned),
            "failed_trials":  sum(1 for c in costs if c.error),
            "total_cost_usd": sum(c.cost_usd for c in costs),
            "budget_usd":     self.budget_usd,
            "budget_used_pct": 100 * self.total_cost() / self.budget_usd,
        }

Step 3: Remote Training Script

This script runs on the rented GPU. It accepts hyperparameters via environment variables and prints intermediate metrics so we can detect pruning criteria.

# remote_train.py  ← uploaded to the Clore server before each trial
"""
Train a model with given hyperparameters.
Reads config from environment variables.
Prints intermediate val_accuracy for pruning decisions.
Writes final result to /tmp/result.json.
"""

import os
import json
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms, models


def main():
    # ── Read hyperparameters from env ─────────────────────────────
    lr          = float(os.environ.get("HP_LR",          "0.001"))
    batch_size  = int(os.environ.get("HP_BATCH_SIZE",    "32"))
    weight_decay= float(os.environ.get("HP_WEIGHT_DECAY","1e-4"))
    dropout     = float(os.environ.get("HP_DROPOUT",     "0.0"))
    optimizer_name = os.environ.get("HP_OPTIMIZER",      "adam")
    epochs      = int(os.environ.get("HP_EPOCHS",        "10"))
    prune_epoch = int(os.environ.get("HP_PRUNE_EPOCH",   "3"))  # report at this epoch

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")
    print(f"Config: lr={lr} bs={batch_size} wd={weight_decay} "
          f"dropout={dropout} optimizer={optimizer_name}")

    # ── Data ──────────────────────────────────────────────────────
    transform = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465),
                             (0.2470, 0.2435, 0.2616)),
    ])
    val_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465),
                             (0.2470, 0.2435, 0.2616)),
    ])

    train_ds = datasets.CIFAR10("/data", train=True,  download=True, transform=transform)
    val_ds   = datasets.CIFAR10("/data", train=False, download=True, transform=val_transform)
    train_dl = DataLoader(train_ds, batch_size=batch_size, shuffle=True,  num_workers=4)
    val_dl   = DataLoader(val_ds,   batch_size=256,        shuffle=False, num_workers=4)

    # ── Model ─────────────────────────────────────────────────────
    model = models.resnet18(weights=None, num_classes=10)
    if dropout > 0:
        # Insert dropout before the final FC layer
        model.fc = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(model.fc.in_features, 10),
        )
    model = model.to(device)

    # ── Optimizer ─────────────────────────────────────────────────
    if optimizer_name == "adam":
        opt = optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
    elif optimizer_name == "sgd":
        opt = optim.SGD(model.parameters(), lr=lr, momentum=0.9,
                        weight_decay=weight_decay)
    elif optimizer_name == "adamw":
        opt = optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)
    else:
        raise ValueError(f"Unknown optimizer: {optimizer_name}")

    scheduler = optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs)
    criterion = nn.CrossEntropyLoss()

    # ── Training loop ─────────────────────────────────────────────
    best_acc = 0.0
    intermediate_values = {}

    for epoch in range(epochs):
        model.train()
        for x, y in train_dl:
            x, y = x.to(device), y.to(device)
            opt.zero_grad()
            loss = criterion(model(x), y)
            loss.backward()
            opt.step()
        scheduler.step()

        # Validation
        model.eval()
        correct = total = 0
        with torch.no_grad():
            for x, y in val_dl:
                x, y  = x.to(device), y.to(device)
                preds  = model(x).argmax(1)
                correct += preds.eq(y).sum().item()
                total   += y.size(0)

        val_acc = correct / total
        intermediate_values[epoch] = val_acc

        # Print in a format the orchestrator can parse
        print(f"INTERMEDIATE epoch={epoch} val_accuracy={val_acc:.6f}")

        best_acc = max(best_acc, val_acc)

    # ── Save result ───────────────────────────────────────────────
    result = {
        "val_accuracy":       best_acc,
        "intermediate_values": intermediate_values,
        "config": {
            "lr": lr, "batch_size": batch_size,
            "weight_decay": weight_decay, "dropout": dropout,
            "optimizer": optimizer_name, "epochs": epochs,
        },
    }
    with open("/tmp/result.json", "w") as f:
        json.dump(result, f, indent=2)

    print(f"FINAL val_accuracy={best_acc:.6f}")


if __name__ == "__main__":
    main()

Step 4: Optuna Objective with Clore Provisioning

# sweep_objective.py
"""Optuna objective function that provisions a Clore GPU per trial."""

import os
import json
import optuna
import logging
import threading
import time
from typing import Dict, Optional

from trial_provisioner import TrialProvisioner
from cost_tracker import CostTracker

logger = logging.getLogger(__name__)

# Script content to upload to each trial server
with open("remote_train.py", "r") as f:
    TRAIN_SCRIPT = f.read()


class CloreObjective:
    """
    Optuna objective that:
    1. Suggests hyperparameters
    2. Provisions a Clore GPU
    3. Uploads and runs the training script
    4. Parses intermediate values for pruning
    5. Returns the final metric
    6. Releases the GPU
    """

    def __init__(
        self,
        api_key:        str,
        cost_tracker:   CostTracker,
        gpu_type:       str   = "RTX 3080",
        max_price:      float = 0.30,
        docker_image:   str   = "pytorch/pytorch:2.1.0-cuda12.1-cudnn8-devel",
        ssh_password:   str   = "OptunaTrial1!",
        epochs:         int   = 10,
        prune_epoch:    int   = 3,    # check median at this epoch
    ):
        self.api_key      = api_key
        self.cost_tracker = cost_tracker
        self.gpu_type     = gpu_type
        self.max_price    = max_price
        self.docker_image = docker_image
        self.ssh_password = ssh_password
        self.epochs       = epochs
        self.prune_epoch  = prune_epoch
        self._lock        = threading.Lock()

    def __call__(self, trial: optuna.Trial) -> float:
        # ── 1. Budget check ───────────────────────────────────────
        if self.cost_tracker.over_budget():
            raise optuna.exceptions.TrialPruned(
                "Global budget exhausted — pruning trial"
            )

        # ── 2. Suggest hyperparameters ────────────────────────────
        lr           = trial.suggest_float("lr",           1e-5, 1e-2, log=True)
        batch_size   = trial.suggest_categorical("batch_size", [16, 32, 64, 128])
        weight_decay = trial.suggest_float("weight_decay", 1e-6, 1e-2, log=True)
        dropout      = trial.suggest_float("dropout",      0.0,  0.5)
        optimizer    = trial.suggest_categorical("optimizer", ["adam", "sgd", "adamw"])

        logger.info(
            f"Trial {trial.number}: lr={lr:.2e} bs={batch_size} "
            f"wd={weight_decay:.2e} dr={dropout:.2f} opt={optimizer}"
        )

        # ── 3. Provision Clore GPU ────────────────────────────────
        provisioner = TrialProvisioner(
            api_key=self.api_key,
            gpu_type=self.gpu_type,
            max_price=self.max_price,
            docker_image=self.docker_image,
            ssh_password=self.ssh_password,
            timeout_sec=3600,
        )

        try:
            host, port = provisioner.provision(
                env={
                    "HP_LR":           str(lr),
                    "HP_BATCH_SIZE":   str(batch_size),
                    "HP_WEIGHT_DECAY": str(weight_decay),
                    "HP_DROPOUT":      str(dropout),
                    "HP_OPTIMIZER":    optimizer,
                    "HP_EPOCHS":       str(self.epochs),
                    "HP_PRUNE_EPOCH":  str(self.prune_epoch),
                }
            )

            self.cost_tracker.start_trial(
                trial_number=trial.number,
                order_id=provisioner.order_id,
                gpu_type=self.gpu_type,
                price_per_hr=provisioner.price_per_hr,
            )

            # ── 4. Upload training script ─────────────────────────
            setup_cmd = (
                "mkdir -p /workspace && "
                f"cat > /workspace/train.py << 'ENDOFSCRIPT'\n"
                f"{TRAIN_SCRIPT}\n"
                "ENDOFSCRIPT"
            )
            provisioner.run_command(host, port, setup_cmd)

            # ── 5. Run training and capture output ────────────────
            run_cmd = "cd /workspace && python train.py 2>&1"

            # We need streaming output for intermediate pruning values.
            # Use a custom streaming SSH call here:
            import paramiko
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            ssh.connect(host, port=port, username="root",
                        password=self.ssh_password, timeout=30,
                        banner_timeout=60)

            _, stdout, _ = ssh.exec_command(run_cmd, timeout=3600)

            pruned      = False
            full_output = []

            for raw_line in iter(stdout.readline, ""):
                line = raw_line.strip()
                full_output.append(line)

                # Parse intermediate values
                if line.startswith("INTERMEDIATE"):
                    parts = dict(
                        p.split("=") for p in line.split()[1:]
                        if "=" in p
                    )
                    epoch    = int(parts.get("epoch", -1))
                    val_acc  = float(parts.get("val_accuracy", 0))

                    # Report to Optuna for pruning decision
                    trial.report(val_acc, step=epoch)
                    if trial.should_prune():
                        logger.info(
                            f"Trial {trial.number} pruned at epoch {epoch} "
                            f"(val_acc={val_acc:.4f})"
                        )
                        ssh.close()
                        pruned = True

                        self.cost_tracker.finish_trial(
                            trial_number=trial.number,
                            cost_usd=provisioner.cost_so_far(),
                            result=val_acc,
                            pruned=True,
                        )
                        raise optuna.exceptions.TrialPruned()

            exit_code = stdout.channel.recv_exit_status()
            ssh.close()

            if exit_code != 0:
                raise RuntimeError(
                    f"Training script failed (exit {exit_code})\n"
                    + "\n".join(full_output[-20:])
                )

            # ── 6. Parse final result from /tmp/result.json ───────
            result_raw = provisioner.run_command(
                host, port, "cat /tmp/result.json"
            )
            result = json.loads(result_raw)
            val_accuracy = result["val_accuracy"]

            self.cost_tracker.finish_trial(
                trial_number=trial.number,
                cost_usd=provisioner.cost_so_far(),
                result=val_accuracy,
            )

            logger.info(
                f"Trial {trial.number} complete: val_acc={val_accuracy:.4f} "
                f"cost=${provisioner.cost_so_far():.4f}"
            )
            return val_accuracy

        except optuna.exceptions.TrialPruned:
            raise  # re-raise; Optuna handles this gracefully

        except Exception as e:
            logger.error(f"Trial {trial.number} failed: {e}")
            self.cost_tracker.finish_trial(
                trial_number=trial.number,
                cost_usd=provisioner.cost_so_far(),
                result=float("nan"),
                error=str(e),
            )
            raise optuna.exceptions.TrialPruned(f"Error: {e}")

        finally:
            provisioner.release()

Step 5: Study Runner with Parallel Trials

# run_sweep.py
"""
Run an Optuna hyperparameter sweep across Clore GPUs.

Usage:
    python run_sweep.py --api-key YOUR_KEY --n-trials 30 --n-jobs 3
"""

import argparse
import json
import logging
import optuna
from pathlib import Path

from cost_tracker import CostTracker
from sweep_objective import CloreObjective

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
)
optuna.logging.set_verbosity(optuna.logging.WARNING)

logger = logging.getLogger(__name__)


def parse_args():
    p = argparse.ArgumentParser(description="Hyperparameter sweep on Clore GPUs")
    p.add_argument("--api-key",     required=True)
    p.add_argument("--n-trials",    type=int,  default=30)
    p.add_argument("--n-jobs",      type=int,  default=3,
                   help="Parallel trials (each rents its own GPU)")
    p.add_argument("--budget-usd",  type=float, default=20.0)
    p.add_argument("--gpu-type",    default="RTX 3080")
    p.add_argument("--max-price",   type=float, default=0.30)
    p.add_argument("--epochs",      type=int,  default=10)
    p.add_argument("--prune-epoch", type=int,  default=3)
    p.add_argument("--storage",     default="sqlite:///sweep.db",
                   help="Optuna storage URI (sqlite:///sweep.db or postgresql://...)")
    p.add_argument("--study-name",  default="clore-sweep")
    p.add_argument("--output-dir",  default="./sweep_results")
    return p.parse_args()


def main():
    args = parse_args()
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    logger.info(
        f"Starting sweep: {args.n_trials} trials, {args.n_jobs} parallel, "
        f"budget=${args.budget_usd}"
    )

    # ── Cost tracker ──────────────────────────────────────────────
    cost_tracker = CostTracker(
        budget_usd=args.budget_usd,
        output_path=str(output_dir / "costs.json"),
    )

    # ── Optuna study ──────────────────────────────────────────────
    pruner  = optuna.pruners.MedianPruner(
        n_startup_trials=5,   # don't prune until 5 trials completed
        n_warmup_steps=2,     # don't prune before epoch 2
        interval_steps=1,
    )
    sampler = optuna.samplers.TPESampler(seed=42)

    study = optuna.create_study(
        study_name=args.study_name,
        storage=args.storage,
        direction="maximize",
        sampler=sampler,
        pruner=pruner,
        load_if_exists=True,   # resume interrupted sweeps
    )

    # ── Objective ─────────────────────────────────────────────────
    objective = CloreObjective(
        api_key=args.api_key,
        cost_tracker=cost_tracker,
        gpu_type=args.gpu_type,
        max_price=args.max_price,
        docker_image="pytorch/pytorch:2.1.0-cuda12.1-cudnn8-devel",
        epochs=args.epochs,
        prune_epoch=args.prune_epoch,
    )

    # ── Run ───────────────────────────────────────────────────────
    try:
        study.optimize(
            objective,
            n_trials=args.n_trials,
            n_jobs=args.n_jobs,      # parallel threads, each rents a GPU
            show_progress_bar=True,
            catch=(Exception,),      # don't crash on individual trial errors
        )
    except KeyboardInterrupt:
        logger.info("Sweep interrupted — saving results...")

    # ── Results ───────────────────────────────────────────────────
    best = study.best_trial
    logger.info(f"\n{'='*60}")
    logger.info(f"Best trial: #{best.number}")
    logger.info(f"  Value:  {best.value:.4f}")
    logger.info(f"  Params: {best.params}")
    logger.info(f"{'='*60}")

    summary = cost_tracker.summary()
    logger.info(f"\nCost summary:")
    logger.info(f"  Total trials:  {summary['total_trials']}")
    logger.info(f"  Pruned:        {summary['pruned_trials']}")
    logger.info(f"  Total cost:    ${summary['total_cost_usd']:.4f}")
    logger.info(f"  Budget used:   {summary['budget_used_pct']:.1f}%")

    # Save results
    results = {
        "best": {
            "trial_number": best.number,
            "value":        best.value,
            "params":       best.params,
        },
        "all_trials": [
            {
                "number":  t.number,
                "value":   t.value,
                "params":  t.params,
                "state":   str(t.state),
            }
            for t in study.trials
        ],
        "cost_summary": summary,
    }
    results_path = output_dir / "results.json"
    with open(results_path, "w") as f:
        json.dump(results, f, indent=2)
    logger.info(f"\nResults saved to {results_path}")

    # Generate visualizations
    generate_plots(study, output_dir)


def generate_plots(study: optuna.Study, output_dir: Path):
    """Generate Optuna visualization plots."""
    try:
        import plotly

        # Optimization history
        fig = optuna.visualization.plot_optimization_history(study)
        fig.write_html(str(output_dir / "optimization_history.html"))

        # Parameter importance
        if len(study.trials) >= 5:
            fig = optuna.visualization.plot_param_importances(study)
            fig.write_html(str(output_dir / "param_importances.html"))

        # Parallel coordinate plot
        fig = optuna.visualization.plot_parallel_coordinate(study)
        fig.write_html(str(output_dir / "parallel_coordinate.html"))

        # Contour plot (top 2 params)
        if len(study.trials) >= 10:
            fig = optuna.visualization.plot_contour(study)
            fig.write_html(str(output_dir / "contour.html"))

        logger.info(f"Plots saved to {output_dir}/")
    except ImportError:
        logger.warning("plotly not installed — skipping plots. pip install plotly")


if __name__ == "__main__":
    main()

Step 6: Terminal Results Summary

# print_results.py
"""Print a formatted summary of sweep results to the terminal."""

import json
import sys
import math
from pathlib import Path
from rich.console import Console
from rich.table import Table
from rich.panel import Panel
from rich.text import Text

console = Console()


def print_sweep_summary(results_path: str = "sweep_results/results.json"):
    path = Path(results_path)
    if not path.exists():
        console.print(f"[red]Results not found: {path}[/]")
        sys.exit(1)

    with open(path) as f:
        data = json.load(f)

    best   = data["best"]
    trials = data["all_trials"]
    costs  = data["cost_summary"]

    # ── Header ────────────────────────────────────────────────────
    console.print(Panel(
        f"[bold cyan]Best Trial #{best['trial_number']}[/]\n"
        f"[green]Val Accuracy: {best['value']:.4f}[/]\n\n"
        + "\n".join(f"  {k}: {v}" for k, v in best["params"].items()),
        title="🏆 Sweep Results",
    ))

    # ── Cost summary ──────────────────────────────────────────────
    console.print(Panel(
        f"Total trials:   {costs['total_trials']}\n"
        f"Pruned:         {costs['pruned_trials']} "
        f"({100*costs['pruned_trials']/max(costs['total_trials'],1):.0f}%)\n"
        f"Failed:         {costs['failed_trials']}\n"
        f"Total cost:     ${costs['total_cost_usd']:.4f}\n"
        f"Budget used:    {costs['budget_used_pct']:.1f}%\n"
        f"Cost per trial: ${costs['total_cost_usd']/max(costs['total_trials'],1):.4f}",
        title="💸 Cost Summary",
    ))

    # ── Top 10 trials table ───────────────────────────────────────
    completed = [
        t for t in trials
        if t["state"] == "TrialState.COMPLETE" and t["value"] is not None
    ]
    completed.sort(key=lambda t: -(t["value"] or 0))

    table = Table(title="📊 Top Trials", show_header=True,
                  header_style="bold cyan")
    table.add_column("#",      width=5)
    table.add_column("Acc",    width=8)
    table.add_column("LR",     width=10)
    table.add_column("Batch",  width=6)
    table.add_column("WD",     width=10)
    table.add_column("Dropout",width=8)
    table.add_column("Optim",  width=8)

    for t in completed[:10]:
        p = t["params"]
        highlight = "bold green" if t["number"] == best["trial_number"] else "white"
        table.add_row(
            Text(str(t["number"]), style=highlight),
            Text(f"{t['value']:.4f}", style=highlight),
            f"{p.get('lr',0):.2e}",
            str(p.get("batch_size", "")),
            f"{p.get('weight_decay',0):.2e}",
            f"{p.get('dropout',0):.2f}",
            p.get("optimizer", ""),
        )

    console.print(table)


if __name__ == "__main__":
    import sys
    results_path = sys.argv[1] if len(sys.argv) > 1 else "sweep_results/results.json"
    print_sweep_summary(results_path)

Distributed Study with PostgreSQL Backend

When running many parallel trials from multiple machines, use PostgreSQL as the Optuna storage backend so all workers share state.

# Start PostgreSQL (or use a managed service)
docker run -d --name optuna-db \
  -e POSTGRES_DB=optuna \
  -e POSTGRES_USER=optuna \
  -e POSTGRES_PASSWORD=optuna_pw \
  -p 5432:5432 \
  postgres:15

# run_distributed_sweep.py
"""Run the sweep from multiple machines using a shared PostgreSQL study."""

import optuna

STORAGE = "postgresql://optuna:optuna_pw@your-db-host:5432/optuna"
STUDY_NAME = "clore-cifar10-sweep"

# All workers load the same study — Optuna coordinates trial assignment
study = optuna.load_study(
    study_name=STUDY_NAME,
    storage=STORAGE,
)

# Each machine runs its own subset of trials
study.optimize(
    objective,       # same CloreObjective as above
    n_trials=10,     # this machine's share
    n_jobs=2,        # 2 parallel GPUs on this machine
)

Optuna Dashboard

While the sweep runs, launch the Optuna dashboard for a web UI:

pip install optuna-dashboard
optuna-dashboard sqlite:///sweep.db
# or
optuna-dashboard postgresql://optuna:optuna_pw@localhost:5432/optuna

Open http://localhost:8080 to see real-time trial progress, parameter importance, and the optimization history plot.

Pruning Deep-Dive: How MedianPruner Saves Money

# The pruner compares each trial's intermediate value
# against the median of all completed trials at the same step.
# Trials below the median are killed early.

pruner = optuna.pruners.MedianPruner(
    n_startup_trials=5,   # Complete 5 trials before pruning begins
    n_warmup_steps=2,     # Never prune before step (epoch) 2
    interval_steps=1,     # Check at every reported step
)

# Alternative: PercentilePruner (kills bottom X% instead of below median)
pruner = optuna.pruners.PercentilePruner(
    percentile=25.0,      # Kill bottom 25%
    n_startup_trials=5,
    n_warmup_steps=3,
)

# Alternative: HyperbandPruner (Hyperband algorithm, very aggressive)
pruner = optuna.pruners.HyperbandPruner(
    min_resource=1,
    max_resource=10,
    reduction_factor=3,
)

Cost Optimization Results

Strategy

Trials

Avg Cost/Trial

Total Cost

Best Acc

Sequential, no pruning

$0.55

$16.50

0.921

Parallel (3x), no pruning

$0.55

$16.50

0.924

Parallel (3x) + MedianPruner

$0.21

$6.30

0.922

Parallel (5x) + MedianPruner

$0.18

$9.00

0.931

Pruning saves ~60% of cost while achieving comparable or better results due to more trials within budget.

Tips for Effective Sweeps

1. Start narrow, then expand. Run 10–15 trials to get a feel for the landscape, then expand n_trials once you know which parameters matter.

2. Use log-scale for learning rate and weight decay. Both span multiple orders of magnitude — log-scale sampling finds good values much faster.

lr = trial.suggest_float("lr", 1e-5, 1e-2, log=True)   # ✅
lr = trial.suggest_float("lr", 0.00001, 0.01)           # ❌ wastes budget on tiny values

3. Prune aggressively for cheap GPUs. If you're using budget GPUs ($0.10–0.20/hr), set n_warmup_steps=1 to prune as early as epoch 1.

4. Set load_if_exists=True to resume interrupted sweeps without losing completed trials.

5. Budget = your real constraint. Set budget_usd before n_trials. If the budget runs out, no more trials start — you won't get surprise bills.

6. Keep trials short. A 10-epoch trial on CIFAR-10 is usually predictive enough. Long trials burn money on runs that are clearly not the best.

Next Steps

PreviousFine-Tuning Models with Hugging Face

NextAuto-Scaling ML Training Pipeline

Last updated 6 months ago

Was this helpful?