alloc 0.0.1__py3-none-any.whl

alloc/__init__.py

@@ -0,0 +1,11 @@
+"""Alloc — GPU intelligence for ML training."""
+
+from __future__ import annotations
+
+__version__ = "0.0.1"
+
+from alloc.ghost import ghost, GhostReport
+from alloc.callbacks import AllocCallback as HuggingFaceCallback
+from alloc.callbacks import AllocLightningCallback as LightningCallback
+
+__all__ = ["ghost", "GhostReport", "HuggingFaceCallback", "LightningCallback", "__version__"]

alloc/artifact_writer.py

@@ -0,0 +1,67 @@
+"""Artifact Writer — write alloc_artifact.json.gz.
+
+Optionally uploads to W&B if wandb is active.
+"""
+
+from __future__ import annotations
+
+import gzip
+import json
+import os
+from datetime import datetime, timezone
+from typing import Optional
+
+
+def write_report(
+    ghost_report: Optional[dict] = None,
+    probe_result: Optional[dict] = None,
+    output_path: Optional[str] = None,
+    hardware_context: Optional[dict] = None,
+    context: Optional[dict] = None,
+) -> str:
+    """Write an artifact to disk.
+
+    Resolution order for output path:
+      1. Explicit output_path parameter
+      2. ALLOC_OUT env var
+      3. ./alloc_artifact.json.gz
+
+    Returns the path written to. Never raises.
+    """
+    try:
+        resolved_path = (
+            output_path
+            or os.environ.get("ALLOC_OUT", "")
+            or "alloc_artifact.json.gz"
+        )
+
+        report = {
+            "version": "0.0.1",
+            "timestamp": datetime.now(timezone.utc).isoformat(),
+            "ghost": ghost_report,
+            "probe": probe_result,
+            "hardware": hardware_context,
+            "context": context if context else None,
+        }
+
+        with gzip.open(resolved_path, "wt", encoding="utf-8") as f:
+            json.dump(report, f, indent=2)
+
+        _try_wandb_upload(resolved_path)
+        return resolved_path
+    except Exception:
+        return ""
+
+
+def _try_wandb_upload(path: str) -> None:
+    """Upload to W&B if wandb is active. Silent no-op otherwise."""
+    if not os.environ.get("WANDB_RUN_ID"):
+        return
+    try:
+        import wandb
+        if wandb.run is not None:
+            artifact = wandb.Artifact("alloc-profile", type="profile")
+            artifact.add_file(path)
+            wandb.run.log_artifact(artifact)
+    except Exception:
+        pass

alloc/callbacks.py

@@ -0,0 +1,342 @@
+"""Alloc Framework Callbacks — capture training step timing for artifact enrichment.
+
+Callbacks for popular ML frameworks. Write timing stats to a sidecar file
+(.alloc_callback.json) so the probe can compute throughput, step latency,
+and dataloader wait estimates.
+
+Usage (HuggingFace):
+    from alloc.callbacks import AllocCallback
+    trainer = Trainer(..., callbacks=[AllocCallback()])
+
+Usage (PyTorch Lightning):
+    from alloc.callbacks import AllocLightningCallback
+    trainer = Trainer(..., callbacks=[AllocLightningCallback()])
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import os
+import time
+from typing import Any, Dict, List, Optional
+
+
+# ── Helpers ──────────────────────────────────────────────────────────────
+
+_ROLLING_WINDOW = 200   # keep last N step times
+_WRITE_EVERY = 50       # flush sidecar every N steps
+
+
+def _compute_percentile(sorted_values, pct):
+    # type: (List[float], float) -> float
+    """Compute a percentile from an already-sorted list. No numpy."""
+    if not sorted_values:
+        return 0.0
+    n = len(sorted_values)
+    k = (pct / 100.0) * (n - 1)
+    f = math.floor(k)
+    c = math.ceil(k)
+    if f == c:
+        return sorted_values[int(k)]
+    return sorted_values[f] * (c - k) + sorted_values[c] * (k - f)
+
+
+def _compute_timing_stats(step_times_ms):
+    # type: (List[float]) -> Dict[str, float]
+    """Compute p50, p90, mean, std, cv from a list of step times (ms)."""
+    if not step_times_ms:
+        return {}
+    sorted_vals = sorted(step_times_ms)
+    n = len(sorted_vals)
+    mean = sum(sorted_vals) / n
+    variance = sum((x - mean) ** 2 for x in sorted_vals) / n
+    std = math.sqrt(variance)
+    cv = std / mean if mean > 0 else 0.0
+    return {
+        "p50": round(_compute_percentile(sorted_vals, 50), 2),
+        "p90": round(_compute_percentile(sorted_vals, 90), 2),
+        "mean": round(mean, 2),
+        "std": round(std, 2),
+        "cv": round(cv, 4),
+    }
+
+
+def _estimate_dataloader_wait(cv):
+    # type: (float) -> float
+    """Estimate dataloader wait % from step-time coefficient of variation.
+
+    Maps CV 0.1–0.5 linearly to 0–30%, capped at 30%.
+    Low CV = consistent steps = no data stalls.
+    High CV = irregular steps = likely waiting for data.
+
+    This is a heuristic estimate — direct measurement would require
+    monkey-patching the dataloader, which violates zero-config.
+    """
+    if cv <= 0.1:
+        return 0.0
+    if cv >= 0.5:
+        return 30.0
+    # Linear interpolation: 0.1 → 0%, 0.5 → 30%
+    return round((cv - 0.1) / 0.4 * 30.0, 1)
+
+
+def _detect_distributed():
+    # type: () -> tuple
+    """Detect if running inside a torch.distributed process group.
+
+    Returns (is_distributed, rank, world_size). Fail-safe: returns
+    (False, 0, 1) if torch.distributed is unavailable or not initialized.
+    """
+    try:
+        import torch.distributed as dist
+        if dist.is_initialized():
+            return True, dist.get_rank(), dist.get_world_size()
+    except Exception:
+        pass
+    return False, 0, 1
+
+
+def _estimate_comm_overhead(step_times_ms, dataloader_wait_pct=0.0):
+    # type: (List[float], float) -> Optional[float]
+    """Estimate communication overhead % for distributed training.
+
+    Uses the p90/p50 spread as a proxy for sync barrier delays.
+    Subtracts estimated dataloader contribution to avoid double-counting.
+    Returns None if insufficient data.
+    """
+    if len(step_times_ms) < 10:
+        return None
+    sorted_vals = sorted(step_times_ms)
+    p50 = _compute_percentile(sorted_vals, 50)
+    p90 = _compute_percentile(sorted_vals, 90)
+    if p50 <= 0:
+        return None
+    raw_pct = ((p90 - p50) / p50) * 100
+    comm_pct = max(0.0, raw_pct - dataloader_wait_pct)
+    return round(min(40.0, comm_pct), 1)
+
+
+def _write_callback_data(data):
+    # type: (Dict[str, Any]) -> None
+    """Write callback data to the alloc sidecar file.
+
+    Creates .alloc_callback.json next to the artifact.
+    Fail-safe: if write fails, training continues unaffected.
+    """
+    path = os.path.join(os.getcwd(), ".alloc_callback.json")
+    try:
+        with open(path, "w") as f:
+            json.dump(data, f)
+    except Exception:
+        pass
+
+
+def _build_sidecar(
+    framework,       # type: str
+    step_count,      # type: int
+    step_times_ms,   # type: List[float]
+    batch_size,      # type: Optional[int]
+    is_distributed=False,  # type: bool
+    rank=0,          # type: int
+    world_size=1,    # type: int
+):
+    # type: (...) -> Dict[str, Any]
+    """Build the sidecar dict from collected timing data."""
+    stats = _compute_timing_stats(step_times_ms)
+    cv = stats.get("cv", 0.0)
+    dataloader_wait_pct = _estimate_dataloader_wait(cv)
+
+    samples_per_sec = None  # type: Optional[float]
+    p50 = stats.get("p50")
+    if p50 and p50 > 0 and batch_size and batch_size > 0:
+        samples_per_sec = round(batch_size / (p50 / 1000.0), 2)
+
+    data = {
+        "framework": framework,
+        "step_count": step_count,
+        "step_time_ms_p50": stats.get("p50"),
+        "step_time_ms_p90": stats.get("p90"),
+        "step_time_ms_mean": stats.get("mean"),
+        "step_time_ms_std": stats.get("std"),
+        "samples_per_sec": samples_per_sec,
+        "batch_size": batch_size,
+        "dataloader_wait_pct": dataloader_wait_pct,
+    }
+
+    if is_distributed:
+        data["is_distributed"] = True
+        data["rank"] = rank
+        data["world_size"] = world_size
+        comm = _estimate_comm_overhead(step_times_ms, dataloader_wait_pct)
+        if comm is not None:
+            data["comm_overhead_pct"] = comm
+
+    return data
+
+
+# ── HuggingFace Callback ─────────────────────────────────────────────────
+
+try:
+    from transformers import TrainerCallback
+
+    class AllocCallback(TrainerCallback):
+        """HuggingFace Trainer callback that captures step timing for Alloc."""
+
+        def __init__(self):
+            # type: () -> None
+            self.step_count = 0          # type: int
+            self._step_times_ms = []     # type: List[float]
+            self._step_start = None      # type: Optional[float]
+            self._batch_size = None      # type: Optional[int]
+            self._last_write_step = 0    # type: int
+            self._dist_checked = False   # type: bool
+            self._is_distributed = False  # type: bool
+            self._rank = 0               # type: int
+            self._world_size = 1         # type: int
+
+        def on_step_begin(self, args, state, control, **kwargs):
+            self._step_start = time.monotonic()
+            # Detect distributed once after process group is initialized
+            if not self._dist_checked:
+                self._is_distributed, self._rank, self._world_size = _detect_distributed()
+                self._dist_checked = True
+
+        def on_step_end(self, args, state, control, **kwargs):
+            self.step_count = state.global_step
+
+            # Compute step duration
+            if self._step_start is not None:
+                elapsed_ms = (time.monotonic() - self._step_start) * 1000.0
+                self._step_times_ms.append(elapsed_ms)
+                # Rolling window
+                if len(self._step_times_ms) > _ROLLING_WINDOW:
+                    self._step_times_ms = self._step_times_ms[-_ROLLING_WINDOW:]
+                self._step_start = None
+
+            # Resolve batch size once
+            if self._batch_size is None:
+                try:
+                    bs = args.per_device_train_batch_size
+                    ga = getattr(args, "gradient_accumulation_steps", 1) or 1
+                    self._batch_size = bs * ga
+                except Exception:
+                    self._batch_size = None
+
+            # Periodic write
+            if self.step_count - self._last_write_step >= _WRITE_EVERY:
+                self._flush()
+                self._last_write_step = self.step_count
+
+        def on_train_end(self, args, state, control, **kwargs):
+            self.step_count = state.global_step
+            self._flush()
+
+        def _flush(self):
+            # type: () -> None
+            data = _build_sidecar(
+                framework="huggingface",
+                step_count=self.step_count,
+                step_times_ms=self._step_times_ms,
+                batch_size=self._batch_size,
+                is_distributed=self._is_distributed,
+                rank=self._rank,
+                world_size=self._world_size,
+            )
+            _write_callback_data(data)
+
+except ImportError:
+    # transformers not installed — provide a stub that raises a clear error
+    class AllocCallback:  # type: ignore[no-redef]
+        """Stub — install transformers to use AllocCallback."""
+
+        def __init__(self):
+            raise ImportError(
+                "AllocCallback requires the `transformers` package. "
+                "Install with: pip install transformers"
+            )
+
+
+# ── PyTorch Lightning Callback ───────────────────────────────────────────
+
+try:
+    from lightning.pytorch.callbacks import Callback as LightningBaseCallback
+
+    class AllocLightningCallback(LightningBaseCallback):
+        """PyTorch Lightning callback that captures step timing for Alloc."""
+
+        def __init__(self):
+            # type: () -> None
+            super().__init__()
+            self.step_count = 0          # type: int
+            self._step_times_ms = []     # type: List[float]
+            self._step_start = None      # type: Optional[float]
+            self._batch_size = None      # type: Optional[int]
+            self._last_write_step = 0    # type: int
+            self._dist_checked = False   # type: bool
+            self._is_distributed = False  # type: bool
+            self._rank = 0               # type: int
+            self._world_size = 1         # type: int
+
+        def on_train_batch_start(self, trainer, pl_module, batch, batch_idx):
+            self._step_start = time.monotonic()
+            if not self._dist_checked:
+                self._is_distributed, self._rank, self._world_size = _detect_distributed()
+                self._dist_checked = True
+
+        def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
+            self.step_count = trainer.global_step
+
+            # Compute step duration
+            if self._step_start is not None:
+                elapsed_ms = (time.monotonic() - self._step_start) * 1000.0
+                self._step_times_ms.append(elapsed_ms)
+                if len(self._step_times_ms) > _ROLLING_WINDOW:
+                    self._step_times_ms = self._step_times_ms[-_ROLLING_WINDOW:]
+                self._step_start = None
+
+            # Resolve batch size once
+            if self._batch_size is None:
+                try:
+                    if hasattr(batch, "__len__"):
+                        self._batch_size = len(batch)
+                    elif hasattr(batch, "shape"):
+                        self._batch_size = batch.shape[0]
+                    elif hasattr(trainer, "datamodule") and trainer.datamodule is not None:
+                        dm = trainer.datamodule
+                        if hasattr(dm, "batch_size"):
+                            self._batch_size = dm.batch_size
+                except Exception:
+                    self._batch_size = None
+
+            # Periodic write
+            if self.step_count - self._last_write_step >= _WRITE_EVERY:
+                self._flush()
+                self._last_write_step = self.step_count
+
+        def on_train_end(self, trainer, pl_module):
+            self.step_count = trainer.global_step
+            self._flush()
+
+        def _flush(self):
+            # type: () -> None
+            data = _build_sidecar(
+                framework="lightning",
+                step_count=self.step_count,
+                step_times_ms=self._step_times_ms,
+                batch_size=self._batch_size,
+                is_distributed=self._is_distributed,
+                rank=self._rank,
+                world_size=self._world_size,
+            )
+            _write_callback_data(data)
+
+except ImportError:
+    class AllocLightningCallback:  # type: ignore[no-redef]
+        """Stub — install lightning to use AllocLightningCallback."""
+
+        def __init__(self):
+            raise ImportError(
+                "AllocLightningCallback requires the `lightning` package. "
+                "Install with: pip install lightning"
+            )

alloc/catalog/__init__.py

@@ -0,0 +1,138 @@
+"""GPU catalog — offline hardware specs and pricing for CLI.
+
+Source of truth: apps/api/src/engine/catalog/gpus.v1.json
+This is a bundled copy for offline CLI use. Update when the API catalog changes.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Dict, List, Optional
+
+_CATALOG_DIR = Path(__file__).parent
+
+# Aliases for common shorthand names
+_ALIASES = {
+    "H100": "nvidia-h100-sxm-80gb",
+    "H100-80GB": "nvidia-h100-sxm-80gb",
+    "A100": "nvidia-a100-sxm-80gb",
+    "A100-80GB": "nvidia-a100-sxm-80gb",
+    "A100-40GB": "nvidia-a100-40gb",
+    "A10G": "nvidia-a10g-24gb",
+    "L40S": "nvidia-l40s-48gb",
+    "L4": "nvidia-l4-24gb",
+    "T4": "nvidia-t4-16gb",
+    "V100": "nvidia-v100-32gb",
+    "V100-32GB": "nvidia-v100-32gb",
+    "V100-16GB": "nvidia-v100-16gb",
+    "RTX-4090": "nvidia-rtx4090-24gb",
+    "RTX-3090": "nvidia-rtx3090-24gb",
+    "H200": "nvidia-h200-141gb",
+    "H100-NVL": "nvidia-h100-nvl-94gb",
+}
+
+
+def _load_catalog() -> dict:
+    """Load GPU catalog from bundled JSON."""
+    with open(_CATALOG_DIR / "gpus.v1.json") as f:
+        return json.load(f)
+
+
+def _load_rate_card() -> dict:
+    """Load default rate card from bundled JSON."""
+    with open(_CATALOG_DIR / "default_rate_card.json") as f:
+        return json.load(f)
+
+
+def list_gpus() -> List[dict]:
+    """Return all GPUs sorted by VRAM descending.
+
+    Each entry has: id, display_name, vendor, vram_gb, architecture,
+    bandwidth_gbps, bf16_tflops, tdp_watts, pricing.
+    """
+    catalog = _load_catalog()
+    rate_card = _load_rate_card()
+
+    result = []
+    for gpu_id, spec in catalog.get("gpus", {}).items():
+        pricing = rate_card.get("rates", {}).get(spec["display_name"], {})
+        result.append({
+            "id": gpu_id,
+            "display_name": spec["display_name"],
+            "vendor": spec.get("vendor", "nvidia"),
+            "vram_gb": spec["vram_gb"],
+            "architecture": spec.get("architecture", ""),
+            "bandwidth_gbps": spec.get("bandwidth_gbps", 0),
+            "bf16_tflops": spec.get("bf16_tflops", 0),
+            "fp16_tflops": spec.get("fp16_tflops", 0),
+            "fp32_tflops": spec.get("fp32_tflops", 0),
+            "tf32_tflops": spec.get("tf32_tflops", 0),
+            "tdp_watts": spec.get("tdp_watts", 0),
+            "interconnect": spec.get("interconnect"),
+            "pricing": pricing,
+        })
+
+    return sorted(result, key=lambda x: x["vram_gb"], reverse=True)
+
+
+def get_default_rate(gpu_name: str) -> Optional[float]:
+    """Look up the average default $/hr for a GPU by name or alias.
+
+    Tries to match the probe-reported GPU name against catalog display names.
+    Returns the average across clouds, or None if not found.
+    """
+    rate_card = _load_rate_card()
+    rates = rate_card.get("rates", {})
+
+    # Direct match by display name
+    for display_name, cloud_rates in rates.items():
+        if display_name.lower() in gpu_name.lower() or gpu_name.lower() in display_name.lower():
+            vals = [v for v in cloud_rates.values() if isinstance(v, (int, float))]
+            return sum(vals) / len(vals) if vals else None
+
+    # Try aliases → display name
+    for alias, stable_id in _ALIASES.items():
+        if alias.lower() in gpu_name.lower():
+            catalog = _load_catalog()
+            spec = catalog.get("gpus", {}).get(stable_id)
+            if spec:
+                dn = spec.get("display_name", "")
+                cloud_rates = rates.get(dn, {})
+                vals = [v for v in cloud_rates.values() if isinstance(v, (int, float))]
+                return sum(vals) / len(vals) if vals else None
+
+    return None
+
+
+def get_gpu(gpu_id: str) -> Optional[dict]:
+    """Look up a GPU by stable ID or alias.
+
+    Returns full spec dict or None if not found.
+    """
+    # Resolve aliases
+    resolved = _ALIASES.get(gpu_id, gpu_id)
+
+    catalog = _load_catalog()
+    rate_card = _load_rate_card()
+
+    spec = catalog.get("gpus", {}).get(resolved)
+    if not spec:
+        return None
+
+    pricing = rate_card.get("rates", {}).get(spec["display_name"], {})
+    return {
+        "id": resolved,
+        "display_name": spec["display_name"],
+        "vendor": spec.get("vendor", "nvidia"),
+        "vram_gb": spec["vram_gb"],
+        "architecture": spec.get("architecture", ""),
+        "bandwidth_gbps": spec.get("bandwidth_gbps", 0),
+        "bf16_tflops": spec.get("bf16_tflops", 0),
+        "fp16_tflops": spec.get("fp16_tflops", 0),
+        "fp32_tflops": spec.get("fp32_tflops", 0),
+        "tf32_tflops": spec.get("tf32_tflops", 0),
+        "tdp_watts": spec.get("tdp_watts", 0),
+        "interconnect": spec.get("interconnect"),
+        "pricing": pricing,
+    }

alloc/catalog/default_rate_card.json

@@ -0,0 +1,18 @@
+{
+  "version": "1",
+  "rates": {
+    "H200": { "aws": 5.50, "gcp": 5.30, "azure": 5.40 },
+    "H100-80GB": { "aws": 4.00, "gcp": 3.90, "azure": 3.85 },
+    "H100-NVL": { "aws": 4.50, "gcp": 4.40, "azure": 4.30 },
+    "A100-80GB": { "aws": 2.50, "gcp": 2.48, "azure": 2.55 },
+    "A100-40GB": { "aws": 2.00, "gcp": 1.95, "azure": 2.10 },
+    "A10G": { "aws": 0.75, "gcp": 0.70, "azure": 0.80 },
+    "L40S": { "aws": 1.50, "gcp": 1.45, "azure": 1.55 },
+    "L4": { "aws": 0.50, "gcp": 0.45, "azure": 0.55 },
+    "T4": { "aws": 0.35, "gcp": 0.30, "azure": 0.40 },
+    "V100-32GB": { "aws": 1.20, "gcp": 1.15, "azure": 1.25 },
+    "V100-16GB": { "aws": 0.90, "gcp": 0.85, "azure": 0.95 },
+    "RTX-4090": { "lambda": 0.70, "coreweave": 0.74 },
+    "RTX-3090": { "lambda": 0.50, "coreweave": 0.54 }
+  }
+}