alloc 0.0.3__tar.gz → 0.0.5__tar.gz

{alloc-0.0.3 → alloc-0.0.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alloc
-Version: 0.0.3
+Version: 0.0.5
 Summary: Engineer-first training calibration: estimate VRAM fit, profile short runs, and pick GPU configs under real budget constraints.
 Author-email: Alloc Labs <hello@alloclabs.com>
 License-Expression: Apache-2.0

{alloc-0.0.3 → alloc-0.0.5}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "alloc"
-version = "0.0.3"
+version = "0.0.5"
 description = "Engineer-first training calibration: estimate VRAM fit, profile short runs, and pick GPU configs under real budget constraints."
 readme = "README.md"
 license = "Apache-2.0"

{alloc-0.0.3 → alloc-0.0.5}/src/alloc/__init__.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-__version__ = "0.0.3"
+__version__ = "0.0.5"
 
 from alloc.ghost import ghost, GhostReport
 from alloc.callbacks import AllocCallback as HuggingFaceCallback

alloc-0.0.5/src/alloc/artifact_loader.py

@@ -0,0 +1,330 @@
+"""Artifact loader — parse alloc_artifact.json.gz for runtime-enhanced diagnosis.
+
+Loads the artifact created by `alloc run`, extracting GPU metrics, timing data,
+and per-rank distributed information for use by Phase 2 diagnosis rules.
+
+Never crashes. Returns None on any failure.
+"""
+
+from __future__ import annotations
+
+import glob
+import gzip
+import json
+import os
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional
+
+
+@dataclass
+class ArtifactData:
+    """Parsed runtime artifact — structured for rule consumption."""
+
+    # Hardware (from probe)
+    gpu_name: Optional[str] = None
+    gpu_count: int = 1
+    per_gpu_vram_total_mb: Optional[float] = None
+    per_gpu_vram_used_mb: Optional[List[float]] = None
+    gpu_utilization_pct: Optional[List[float]] = None
+    power_draw_w: Optional[List[float]] = None
+    sm_version: Optional[str] = None
+    interconnect: Optional[str] = None
+
+    # Timing (from callbacks — may be None)
+    step_times_ms: Optional[List[float]] = None
+    step_time_p50_ms: Optional[float] = None
+    step_time_p90_ms: Optional[float] = None
+    throughput_samples_per_sec: Optional[float] = None
+    dataloader_wait_pct: Optional[float] = None
+
+    # Phase timing (from callbacks with CUDA events)
+    phase_forward_ms_p50: Optional[float] = None
+    phase_forward_ms_p90: Optional[float] = None
+    phase_backward_ms_p50: Optional[float] = None
+    phase_backward_ms_p90: Optional[float] = None
+    phase_optimizer_ms_p50: Optional[float] = None
+    phase_optimizer_ms_p90: Optional[float] = None
+    phase_dataloader_ms_p50: Optional[float] = None
+    phase_dataloader_ms_p90: Optional[float] = None
+    has_phase_timing: bool = False
+
+    # Communication overhead (estimated from phase timing in distributed training)
+    comm_overhead_pct: Optional[float] = None
+
+    # Per-rank data (from distributed callbacks)
+    per_rank_peak_vram_mb: Optional[List[float]] = None
+    per_rank_step_times_ms: Optional[List[List[float]]] = None
+    straggler_ratio: Optional[float] = None
+
+    # Architecture metadata (from callback introspection)
+    architecture_type: Optional[str] = None
+    optimizer_type: Optional[str] = None
+    fine_tuning_method: Optional[str] = None
+    gradient_checkpointing: Optional[bool] = None
+    attention_type: Optional[str] = None
+    param_count: Optional[int] = None
+    trainable_param_count: Optional[int] = None
+
+    # Run metadata
+    exit_code: Optional[int] = None
+    duration_s: Optional[float] = None
+    command: Optional[str] = None
+    git_sha: Optional[str] = None
+    is_oom: bool = False
+
+    # Aggregate metrics (computed from samples)
+    avg_gpu_util: Optional[float] = None
+    peak_vram_mb: Optional[float] = None
+
+
+def load_artifact(path: str) -> Optional[ArtifactData]:
+    """Load and parse alloc_artifact.json.gz. Returns None on any failure."""
+    try:
+        if path.endswith(".gz"):
+            with gzip.open(path, "rt", encoding="utf-8") as f:
+                raw = json.load(f)
+        else:
+            with open(path, "r", encoding="utf-8") as f:
+                raw = json.load(f)
+    except Exception:
+        return None
+
+    return _parse_artifact(raw)
+
+
+def find_artifact(directory: str = ".") -> Optional[str]:
+    """Find most recent alloc_artifact*.json.gz in directory. Returns path or None."""
+    patterns = [
+        os.path.join(directory, "alloc_artifact*.json.gz"),
+        os.path.join(directory, "alloc_artifact*.json"),
+    ]
+    candidates = []  # type: List[str]
+    for pattern in patterns:
+        candidates.extend(glob.glob(pattern))
+
+    if not candidates:
+        return None
+
+    # Sort by modification time, return newest
+    return max(candidates, key=os.path.getmtime)
+
+
+def _parse_artifact(raw: dict) -> ArtifactData:
+    """Parse raw artifact JSON into ArtifactData."""
+    data = ArtifactData()
+
+    probe = raw.get("probe") or {}
+    hardware = raw.get("hardware") or {}
+    context = raw.get("context") or {}
+
+    # Hardware
+    data.gpu_name = hardware.get("gpu_name") or probe.get("gpu_name")
+    gpu_count = hardware.get("num_gpus_detected")
+    data.gpu_count = gpu_count if gpu_count is not None and gpu_count > 0 else 1
+    data.per_gpu_vram_total_mb = _float_or_none(
+        hardware.get("gpu_total_vram_mb") or probe.get("gpu_total_vram_mb")
+    )
+    data.sm_version = hardware.get("sm_version")
+    data.interconnect = probe.get("interconnect_type")
+
+    # Peak VRAM — from probe samples or direct field
+    peak = _float_or_none(probe.get("peak_vram_mb"))
+    data.peak_vram_mb = peak
+
+    # Per-GPU VRAM: use per_rank_peak_vram_mb if available, else single peak
+    per_rank = probe.get("per_rank_peak_vram_mb")
+    if isinstance(per_rank, list) and per_rank:
+        data.per_gpu_vram_used_mb = [x for x in (_float_or_none(v) for v in per_rank) if x is not None]
+    elif peak is not None:
+        data.per_gpu_vram_used_mb = [peak]
+
+    data.per_rank_peak_vram_mb = data.per_gpu_vram_used_mb
+
+    # GPU utilization from samples
+    samples = probe.get("samples") or []
+    if samples:
+        utils = [s.get("gpu_util_pct") for s in samples if s.get("gpu_util_pct") is not None]
+        if utils:
+            data.gpu_utilization_pct = [x for x in (_float_or_none(u) for u in utils) if x is not None]
+            data.avg_gpu_util = sum(data.gpu_utilization_pct) / len(data.gpu_utilization_pct)
+
+        powers = [s.get("power_w") for s in samples if s.get("power_w") is not None]
+        if powers:
+            data.power_draw_w = [x for x in (_float_or_none(p) for p in powers) if x is not None]
+
+    # Avg GPU util from probe aggregate
+    if data.avg_gpu_util is None:
+        data.avg_gpu_util = _float_or_none(probe.get("avg_gpu_util"))
+
+    # Timing (from callback sidecar data merged into probe)
+    data.step_time_p50_ms = _float_or_none(probe.get("step_time_ms_p50"))
+    data.step_time_p90_ms = _float_or_none(probe.get("step_time_ms_p90"))
+    data.throughput_samples_per_sec = _float_or_none(probe.get("samples_per_sec"))
+    data.dataloader_wait_pct = _float_or_none(probe.get("dataloader_wait_pct"))
+
+    # Phase timing (from CUDA events in callbacks)
+    data.phase_forward_ms_p50 = _float_or_none(probe.get("phase_forward_ms_p50"))
+    data.phase_forward_ms_p90 = _float_or_none(probe.get("phase_forward_ms_p90"))
+    data.phase_backward_ms_p50 = _float_or_none(probe.get("phase_backward_ms_p50"))
+    data.phase_backward_ms_p90 = _float_or_none(probe.get("phase_backward_ms_p90"))
+    data.phase_optimizer_ms_p50 = _float_or_none(probe.get("phase_optimizer_ms_p50"))
+    data.phase_optimizer_ms_p90 = _float_or_none(probe.get("phase_optimizer_ms_p90"))
+    data.phase_dataloader_ms_p50 = _float_or_none(probe.get("phase_dataloader_ms_p50"))
+    data.phase_dataloader_ms_p90 = _float_or_none(probe.get("phase_dataloader_ms_p90"))
+    data.has_phase_timing = bool(probe.get("has_phase_timing"))
+    data.comm_overhead_pct = _float_or_none(probe.get("comm_overhead_pct"))
+
+    # Architecture metadata (from callback introspection)
+    data.architecture_type = probe.get("architecture_type")
+    data.optimizer_type = probe.get("optimizer_type")
+    data.fine_tuning_method = probe.get("fine_tuning_method")
+    gc = probe.get("gradient_checkpointing")
+    data.gradient_checkpointing = bool(gc) if gc is not None else None
+    data.attention_type = probe.get("attention_type")
+    data.param_count = _int_or_none(probe.get("param_count"))
+    data.trainable_param_count = _int_or_none(probe.get("trainable_param_count"))
+
+    # Raw step times (for per-rank merge / straggler detection)
+    raw_times = probe.get("step_times_raw")
+    if isinstance(raw_times, list) and raw_times:
+        data.step_times_ms = [x for x in (_float_or_none(t) for t in raw_times) if x is not None]
+
+    # Run metadata
+    data.exit_code = probe.get("exit_code")
+    data.duration_s = _float_or_none(probe.get("duration_seconds"))
+    data.command = probe.get("command")
+
+    # Git SHA from context
+    git_ctx = context.get("git") or {}
+    data.git_sha = git_ctx.get("commit_sha")
+
+    # OOM detection: exit_code != 0 AND VRAM utilization > 95%
+    data.is_oom = _detect_oom(data)
+
+    return data
+
+
+def _detect_oom(data: ArtifactData) -> bool:
+    """Detect probable OOM from exit code and VRAM utilization."""
+    if data.exit_code is None or data.exit_code == 0:
+        return False
+
+    if data.per_gpu_vram_used_mb and data.per_gpu_vram_total_mb:
+        total = data.per_gpu_vram_total_mb
+        if total > 0:
+            max_used = max(data.per_gpu_vram_used_mb)
+            utilization = max_used / total
+            if utilization > 0.95:
+                return True
+
+    return False
+
+
+def _float_or_none(val) -> Optional[float]:
+    """Convert a value to float, returning None on failure."""
+    if val is None:
+        return None
+    try:
+        return float(val)
+    except (ValueError, TypeError):
+        return None
+
+
+def _int_or_none(val) -> Optional[int]:
+    """Convert a value to int, returning None on failure."""
+    if val is None:
+        return None
+    try:
+        return int(val)
+    except (ValueError, TypeError):
+        return None
+
+
+def find_rank_artifacts(directory: str = ".") -> List[str]:
+    """Find all alloc_artifact_rank*.json.gz files in directory."""
+    pattern = os.path.join(directory, "alloc_artifact_rank*.json.gz")
+    return sorted(glob.glob(pattern))
+
+
+def merge_artifacts(paths: List[str]) -> Optional[ArtifactData]:
+    """Load multiple rank artifacts and merge into a single ArtifactData.
+
+    Combines per-rank VRAM peaks, step times, and computes straggler metrics.
+    Returns None if no valid artifacts found.
+    """
+    artifacts = [load_artifact(p) for p in paths]
+    artifacts = [a for a in artifacts if a is not None]
+    if not artifacts:
+        return None
+
+    merged = ArtifactData()
+
+    # Use GPU info from rank 0
+    merged.gpu_name = artifacts[0].gpu_name
+    merged.gpu_count = len(artifacts)
+    merged.per_gpu_vram_total_mb = artifacts[0].per_gpu_vram_total_mb
+    merged.sm_version = artifacts[0].sm_version
+    merged.interconnect = artifacts[0].interconnect
+
+    # Per-rank peak VRAM
+    peaks = [a.peak_vram_mb for a in artifacts if a.peak_vram_mb is not None]
+    if peaks:
+        merged.per_rank_peak_vram_mb = peaks
+        merged.peak_vram_mb = max(peaks)
+        merged.per_gpu_vram_used_mb = peaks
+
+    # Per-rank step times (raw lists from each rank's rolling window)
+    rank_step_times = []
+    for a in artifacts:
+        if a.step_times_ms:
+            rank_step_times.append(a.step_times_ms)
+    if rank_step_times:
+        merged.per_rank_step_times_ms = rank_step_times
+
+    # Aggregate step timing from rank 0 (representative)
+    merged.step_time_p50_ms = artifacts[0].step_time_p50_ms
+    merged.step_time_p90_ms = artifacts[0].step_time_p90_ms
+    merged.throughput_samples_per_sec = artifacts[0].throughput_samples_per_sec
+
+    # Phase timing from rank 0 (representative)
+    merged.has_phase_timing = artifacts[0].has_phase_timing
+    merged.phase_forward_ms_p50 = artifacts[0].phase_forward_ms_p50
+    merged.phase_forward_ms_p90 = artifacts[0].phase_forward_ms_p90
+    merged.phase_backward_ms_p50 = artifacts[0].phase_backward_ms_p50
+    merged.phase_backward_ms_p90 = artifacts[0].phase_backward_ms_p90
+    merged.phase_optimizer_ms_p50 = artifacts[0].phase_optimizer_ms_p50
+    merged.phase_optimizer_ms_p90 = artifacts[0].phase_optimizer_ms_p90
+    merged.phase_dataloader_ms_p50 = artifacts[0].phase_dataloader_ms_p50
+    merged.phase_dataloader_ms_p90 = artifacts[0].phase_dataloader_ms_p90
+    merged.comm_overhead_pct = artifacts[0].comm_overhead_pct
+
+    # Straggler detection: ratio of slowest to fastest rank by step time p50
+    p50s = [a.step_time_p50_ms for a in artifacts if a.step_time_p50_ms is not None]
+    if p50s and len(p50s) > 1 and min(p50s) > 0:
+        merged.straggler_ratio = round(max(p50s) / min(p50s), 3)
+
+    # Aggregate GPU util from all ranks
+    all_utils = []
+    for a in artifacts:
+        if a.avg_gpu_util is not None:
+            all_utils.append(a.avg_gpu_util)
+    if all_utils:
+        merged.avg_gpu_util = sum(all_utils) / len(all_utils)
+
+    # Architecture metadata from rank 0
+    merged.architecture_type = artifacts[0].architecture_type
+    merged.optimizer_type = artifacts[0].optimizer_type
+    merged.fine_tuning_method = artifacts[0].fine_tuning_method
+    merged.gradient_checkpointing = artifacts[0].gradient_checkpointing
+    merged.attention_type = artifacts[0].attention_type
+    merged.param_count = artifacts[0].param_count
+    merged.trainable_param_count = artifacts[0].trainable_param_count
+
+    # Run metadata from rank 0
+    merged.exit_code = artifacts[0].exit_code
+    merged.duration_s = artifacts[0].duration_s
+    merged.command = artifacts[0].command
+    merged.git_sha = artifacts[0].git_sha
+    merged.is_oom = any(a.is_oom for a in artifacts)
+
+    return merged

{alloc-0.0.3 → alloc-0.0.5}/src/alloc/artifact_writer.py

@@ -44,6 +44,10 @@ def write_report(
             "context": context if context else None,
         }
 
+        if os.path.exists(resolved_path):
+            import sys
+            print(f"alloc: warning: overwriting existing artifact at {resolved_path}", file=sys.stderr)
+
         with gzip.open(resolved_path, "wt", encoding="utf-8") as f:
             json.dump(report, f, indent=2)
 

{alloc-0.0.3 → alloc-0.0.5}/src/alloc/browser_auth.py

@@ -11,6 +11,7 @@ from __future__ import annotations
 
 import base64
 import hashlib
+import html
 import secrets
 import socket
 import threading
@@ -40,7 +41,8 @@ def _find_open_port(start=17256, attempts=20):
     for port in range(start, start + attempts):
         try:
             with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-                s.bind(("127.0.0.1", port))
+                # Bind to all interfaces so both localhost and 127.0.0.1 work
+                s.bind(("0.0.0.0", port))
                 return port
         except OSError:
             continue
@@ -69,7 +71,7 @@ class _CallbackHandler(BaseHTTPRequestHandler):
             self._respond(
                 400,
                 "<html><body style='font-family:system-ui;text-align:center;padding:60px'>"
-                f"<h2>Login failed</h2><p>{error_desc}</p>"
+                f"<h2>Login failed</h2><p>{html.escape(error_desc)}</p>"
                 "</body></html>",
             )
         else:
@@ -108,7 +110,8 @@ def browser_login(
     verifier, challenge = _generate_pkce_pair()
     port = _find_open_port()
 
-    redirect_uri = f"http://localhost:{port}/callback"
+    # Use 127.0.0.1 (not localhost) — more reliable, avoids IPv6 resolution issues.
+    redirect_uri = f"http://127.0.0.1:{port}/callback"
 
     authorize_params = urlencode({
         "provider": provider,
@@ -118,7 +121,8 @@ def browser_login(
     })
     authorize_url = f"{supabase_url}/auth/v1/authorize?{authorize_params}"
 
-    server = HTTPServer(("127.0.0.1", port), _CallbackHandler)
+    # Bind to 0.0.0.0 so both localhost and 127.0.0.1 reach the server.
+    server = HTTPServer(("0.0.0.0", port), _CallbackHandler)
     server.auth_code = None  # type: ignore[attr-defined]
     server.auth_error = None  # type: ignore[attr-defined]
     server.timeout = 1  # poll interval for handle_request()
@@ -128,16 +132,21 @@ def browser_login(
     server_thread.daemon = True
     server_thread.start()
 
+    import sys
+
     # Open the browser (or print URL as fallback).
     try:
         opened = webbrowser.open(authorize_url)
     except Exception:
         opened = False
 
+    print(f"\nCallback server listening on http://127.0.0.1:{port}/callback", file=sys.stderr)
     if not opened:
         print(f"\nOpen this URL in your browser to log in:\n\n  {authorize_url}\n")
     else:
-        print("Opened browser for login. Waiting for callback...")
+        print("Opened browser for login. Waiting for callback...", file=sys.stderr)
+        print(f"If login completes but the terminal stays stuck, your Supabase", file=sys.stderr)
+        print(f"redirect allowlist may not include http://127.0.0.1:{port}/callback", file=sys.stderr)
 
     server_thread.join(timeout=timeout_seconds + 5)
     server.server_close()
@@ -146,7 +155,20 @@ def browser_login(
         raise RuntimeError(f"OAuth error: {server.auth_error}")
 
     if not server.auth_code:
-        raise RuntimeError("Login timed out — no callback received within 120 seconds.")
+        raise RuntimeError(
+            f"Login timed out — no callback received within {timeout_seconds} seconds.\n"
+            f"\n"
+            f"  The browser never reached http://127.0.0.1:{port}/callback.\n"
+            f"\n"
+            f"  Common causes:\n"
+            f"  1. Supabase redirect allowlist does not include http://127.0.0.1:{port}/**\n"
+            f"     (Check: Supabase Dashboard → Authentication → URL Configuration → Redirect URLs)\n"
+            f"  2. Browser redirected to your site URL instead of localhost\n"
+            f"  3. Firewall or antivirus blocked the local callback server\n"
+            f"\n"
+            f"  Workaround: alloc login --method token --token <paste-access-token>\n"
+            f"  (Copy token from browser DevTools → Application → Local Storage → sb-*-auth-token)"
+        )
 
     # Exchange auth code + verifier for tokens.
     with httpx.Client(timeout=15) as client: