alloc 0.0.5__tar.gz → 0.0.7__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alloc
-Version: 0.0.5
+Version: 0.0.7
 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.5 → alloc-0.0.7}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "alloc"
-version = "0.0.5"
+version = "0.0.7"
 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.7/src/alloc/__init__.py

@@ -0,0 +1,18 @@
+"""Alloc — GPU intelligence for ML training."""
+
+from __future__ import annotations
+
+import warnings as _warnings
+_warnings.filterwarnings("ignore", category=FutureWarning, module="pynvml")
+_warnings.filterwarnings("ignore", category=DeprecationWarning, module="pynvml")
+_warnings.filterwarnings("ignore", category=FutureWarning, module=r"torch\.cuda")
+_warnings.filterwarnings("ignore", category=DeprecationWarning, module=r"torch\.cuda")
+del _warnings
+
+__version__ = "0.0.7"
+
+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-0.0.5 → alloc-0.0.7}/src/alloc/cli.py

@@ -16,8 +16,17 @@ from __future__ import annotations
 
 import os
 import sys
+import warnings
 from typing import List, Optional
 
+# Suppress noisy third-party warnings globally — pynvml deprecation (emitted
+# from torch.cuda.__init__) and urllib3 LibreSSL warnings clutter CLI output.
+warnings.filterwarnings("ignore", category=FutureWarning, module="pynvml")
+warnings.filterwarnings("ignore", category=DeprecationWarning, module="pynvml")
+warnings.filterwarnings("ignore", category=FutureWarning, module=r"torch\.cuda")
+warnings.filterwarnings("ignore", category=DeprecationWarning, module=r"torch\.cuda")
+warnings.filterwarnings("ignore", message=".*LibreSSL.*", module="urllib3")
+
 import typer
 from rich.console import Console
 
@@ -68,6 +77,19 @@ def ghost(
             console.print(f"[dim]Tip: alloc ghost {script} --param-count-b 7.0[/dim]")
         raise typer.Exit(1)
 
+    if info.extraction_error:
+        if json_output:
+            _print_json({
+                "error": info.extraction_error,
+                "detail": info.extraction_detail,
+                "supported": False,
+            })
+        else:
+            console.print(f"[yellow]{info.extraction_detail}[/yellow]")
+            if info.extraction_error == "distributed_entrypoint":
+                console.print("[dim]Tip: alloc ghost model.py  (point to the file that defines your model)[/dim]")
+        raise typer.Exit(1)
+
     # Use dtype from execution if available, otherwise CLI flag
     resolved_dtype = info.dtype if info.method == "execution" else dtype
 
@@ -2092,12 +2114,32 @@ def scan(
 
     try:
         headers = {"Content-Type": "application/json"}
+        used_auth = bool(token)
+
         if token:
             headers["Authorization"] = f"Bearer {token}"
+            endpoint = "/scans"
+        else:
+            endpoint = "/scans/cli"
 
-        endpoint = "/scans" if token else "/scans/cli"
         with httpx.Client(timeout=30) as client:
             resp = client.post(f"{api_url}{endpoint}", json=payload, headers=headers)
+
+            # On 401 with a saved token: try refresh, then fall back to public endpoint
+            if resp.status_code == 401 and used_auth:
+                new_token = try_refresh_access_token()
+                if new_token:
+                    headers["Authorization"] = f"Bearer {new_token}"
+                    resp = client.post(f"{api_url}/scans", json=payload, headers=headers)
+                else:
+                    # Token refresh failed — fall back to unauthenticated scan
+                    console.print(
+                        "[yellow]Session expired — falling back to public scan "
+                        "(org fleet context unavailable). Run `alloc login` to restore.[/yellow]",
+                    )
+                    del headers["Authorization"]
+                    resp = client.post(f"{api_url}/scans/cli", json=payload, headers=headers)
+
             resp.raise_for_status()
             result = resp.json()
 
@@ -2107,7 +2149,12 @@ def scan(
             _print_scan_result(result, gpu, strategy)
     except httpx.HTTPStatusError as e:
         if json_output:
-            _print_json({"error": f"API error {e.response.status_code}"})
+            detail = ""
+            try:
+                detail = e.response.json().get("detail", "")
+            except Exception:
+                pass
+            _print_json({"error": f"API error {e.response.status_code}", "detail": detail})
         elif e.response.status_code == 403:
             console.print("[yellow]AI analysis requires a Pro or Enterprise plan.[/yellow]")
             console.print("[dim]The scan still works — just without AI-powered analysis.[/dim]")
@@ -2147,18 +2194,30 @@ def login(
     ),
 ):
     """Authenticate with Alloc dashboard."""
-    # Suppress noisy third-party warnings (urllib3 LibreSSL, pynvml deprecation)
-    # that clutter the auth flow output.
-    import warnings
-    warnings.filterwarnings("ignore", category=DeprecationWarning, module="pynvml")
-    warnings.filterwarnings("ignore", message=".*LibreSSL.*", module="urllib3")
-    warnings.filterwarnings("ignore", message=".*pynvml.*", category=FutureWarning)
-
     import httpx
     from alloc.config import get_supabase_url, get_supabase_anon_key, load_config, save_config
 
     # --- Browser OAuth flow ---
     if browser:
+        # Detect headless/SSH environments where browser login won't work
+        is_headless = (
+            not os.environ.get("DISPLAY")
+            and not os.environ.get("WAYLAND_DISPLAY")
+            and sys.platform != "darwin"
+            and sys.platform != "win32"
+        )
+        is_ssh = bool(os.environ.get("SSH_CLIENT") or os.environ.get("SSH_TTY"))
+        if is_headless or is_ssh:
+            console.print("[yellow]Headless/SSH environment detected — browser login won't work here.[/yellow]")
+            console.print()
+            console.print("Use token login instead:")
+            console.print("  1. Log in at [cyan]https://alloclabs.com[/cyan] in your local browser")
+            console.print("  2. Open DevTools → Application → Local Storage → copy your access_token")
+            console.print("  3. Run: [green]alloc login --method token --token <paste-token>[/green]")
+            console.print()
+            console.print("[dim]Or set ALLOC_TOKEN=<token> in your environment.[/dim]")
+            raise typer.Exit(1)
+
         provider = (provider or "").strip().lower()
         if provider not in ("google", "azure"):
             console.print("[red]Invalid --provider. Use: google or azure[/red]")

{alloc-0.0.5 → alloc-0.0.7}/src/alloc/extractor_runner.py

@@ -217,7 +217,9 @@ def main():
         try:
             obj = getattr(module, attr_name)
             if isinstance(obj, nn.Module):
-                count, dtype_str = _count_params(obj)
+                # Unwrap DDP/FSDP wrappers to get the underlying model
+                unwrapped = getattr(obj, "module", obj)
+                count, dtype_str = _count_params(unwrapped)
                 if count > 0:
                     models.append((count, dtype_str, attr_name))
         except Exception:

{alloc-0.0.5 → alloc-0.0.7}/src/alloc/model_extractor.py

@@ -33,6 +33,8 @@ class ModelInfo:
     seq_length: Optional[int] = None
     activation_memory_bytes: Optional[int] = None
     activation_method: Optional[str] = None  # "traced" | None
+    extraction_error: Optional[str] = None   # "distributed_entrypoint" | None
+    extraction_detail: Optional[str] = None  # human-readable explanation
 
 
 def extract_model_info(
@@ -99,6 +101,13 @@ def _extract_via_subprocess(
 
         env = os.environ.copy()
         env["CUDA_VISIBLE_DEVICES"] = ""  # prevent GPU allocation
+        # Set distributed env vars so DDP scripts don't crash on
+        # torch.distributed.init_process_group() during model extraction.
+        env.setdefault("RANK", "0")
+        env.setdefault("LOCAL_RANK", "0")
+        env.setdefault("WORLD_SIZE", "1")
+        env.setdefault("MASTER_ADDR", "127.0.0.1")
+        env.setdefault("MASTER_PORT", "29500")
 
         subprocess.run(
             [sys.executable, "-m", "alloc.extractor_runner", sidecar_path, script_abs],
@@ -127,6 +136,24 @@ def _extract_via_subprocess(
                 activation_method=data.get("activation_method"),
             )
 
+        # Structured degradation for distributed scripts
+        if data.get("status") == "error":
+            error_msg = data.get("error", "")
+            _dist_keywords = ("init_process_group", "NCCL", "gloo", "distributed",
+                              "MASTER_ADDR", "MASTER_PORT", "RendezvousError")
+            if any(kw.lower() in error_msg.lower() for kw in _dist_keywords):
+                return ModelInfo(
+                    param_count=0,
+                    dtype="float16",
+                    model_name=None,
+                    method="execution",
+                    extraction_error="distributed_entrypoint",
+                    extraction_detail=(
+                        "Script requires a distributed runtime (e.g. torchrun). "
+                        "Run ghost on the model definition file instead of the launcher script."
+                    ),
+                )
+
         return None
 
     except subprocess.TimeoutExpired:

{alloc-0.0.5 → alloc-0.0.7}/src/alloc/probe.py

@@ -18,6 +18,13 @@ from dataclasses import dataclass, field
 from enum import Enum
 from typing import List, Optional
 
+import warnings as _warnings
+_warnings.filterwarnings("ignore", category=FutureWarning, module="pynvml")
+_warnings.filterwarnings("ignore", category=DeprecationWarning, module="pynvml")
+_warnings.filterwarnings("ignore", category=FutureWarning, module=r"torch\.cuda")
+_warnings.filterwarnings("ignore", category=DeprecationWarning, module=r"torch\.cuda")
+del _warnings
+
 
 class StopReason(str, Enum):
     STABLE = "stable"
@@ -83,6 +90,85 @@ def _try_import_pynvml():
 
 
 
+def _parse_launcher_gpu_count(command):
+    # type: (list) -> Optional[int]
+    """Extract expected GPU count from launcher command line.
+
+    Recognizes: torchrun, torch.distributed.launch, accelerate launch,
+    deepspeed, mpirun/mpiexec.
+    """
+    args = [str(a) for a in command]
+    cmd_str = " ".join(args)
+
+    # torchrun --nproc_per_node=N or --nproc-per-node=N
+    for i, a in enumerate(args):
+        for flag in ("--nproc_per_node", "--nproc-per-node"):
+            if a.startswith(flag + "="):
+                try:
+                    return int(a.split("=", 1)[1])
+                except ValueError:
+                    pass
+            elif a == flag and i + 1 < len(args):
+                try:
+                    return int(args[i + 1])
+                except ValueError:
+                    pass
+
+    # accelerate launch --num_processes=N or --num-processes=N
+    for i, a in enumerate(args):
+        for flag in ("--num_processes", "--num-processes"):
+            if a.startswith(flag + "="):
+                try:
+                    return int(a.split("=", 1)[1])
+                except ValueError:
+                    pass
+            elif a == flag and i + 1 < len(args):
+                try:
+                    return int(args[i + 1])
+                except ValueError:
+                    pass
+
+    # deepspeed --num_gpus=N or --num-gpus=N
+    for i, a in enumerate(args):
+        for flag in ("--num_gpus", "--num-gpus"):
+            if a.startswith(flag + "="):
+                try:
+                    return int(a.split("=", 1)[1])
+                except ValueError:
+                    pass
+            elif a == flag and i + 1 < len(args):
+                try:
+                    return int(args[i + 1])
+                except ValueError:
+                    pass
+
+    # mpirun/mpiexec -np N or -n N
+    for i, a in enumerate(args):
+        if a in ("-np", "-n") and i + 1 < len(args):
+            if any(x in cmd_str for x in ("mpirun", "mpiexec")):
+                try:
+                    return int(args[i + 1])
+                except ValueError:
+                    pass
+
+    return None
+
+
+def _read_child_env(pid, var_name):
+    # type: (int, str) -> Optional[str]
+    """Read an environment variable from a child process via /proc (Linux only)."""
+    try:
+        env_path = "/proc/{}/environ".format(pid)
+        with open(env_path, "rb") as f:
+            env_data = f.read()
+        for entry in env_data.split(b"\x00"):
+            if entry.startswith(var_name.encode() + b"="):
+                return entry.split(b"=", 1)[1].decode("utf-8")
+    except Exception:
+        pass
+    return None
+
+
 def _get_child_pids(pid):
     # type: (int) -> List[int]
     """Get child PIDs of a process. Returns empty list on failure."""
@@ -101,11 +187,15 @@ def _get_child_pids(pid):
     return []
 
 
-def _discover_gpu_indices(proc_pid, pynvml, fallback_index=0):
-    # type: (int, ..., int) -> List[int]
+def _discover_gpu_indices(proc_pid, pynvml, fallback_index=0, expected_gpus=None):
+    # type: (int, ..., int, Optional[int]) -> List[int]
     """Discover which GPUs a process (and its children) are using.
 
-    Iterates all GPU devices and checks running compute processes.
+    Two strategies:
+    1. PID-matching: walks process tree, matches PIDs against NVML compute processes.
+    2. Active-GPU counting: counts GPUs with ANY compute processes. Used when PID
+       matching finds fewer GPUs than expected (common for DDP launchers).
+
     Falls back to [fallback_index] if discovery fails or finds nothing.
     """
     try:
@@ -143,11 +233,48 @@ def _discover_gpu_indices(proc_pid, pynvml, fallback_index=0):
             for ggchild in _get_child_pids(grandchild):
                 target_pids.add(ggchild)
 
+    # Also try reading WORLD_SIZE from child process environments (Linux)
+    child_world_size = None
+    for child in _get_child_pids(proc_pid):
+        ws = _read_child_env(child, "WORLD_SIZE")
+        if ws:
+            try:
+                child_world_size = int(ws)
+            except ValueError:
+                pass
+            break
+        for grandchild in _get_child_pids(child):
+            ws = _read_child_env(grandchild, "WORLD_SIZE")
+            if ws:
+                try:
+                    child_world_size = int(ws)
+                except ValueError:
+                    pass
+                break
+            for ggchild in _get_child_pids(grandchild):
+                ws = _read_child_env(ggchild, "WORLD_SIZE")
+                if ws:
+                    try:
+                        child_world_size = int(ws)
+                    except ValueError:
+                        pass
+                    break
+
+    # Determine how many GPUs we expect
+    effective_expected = expected_gpus
+    if child_world_size is not None:
+        if effective_expected is None or child_world_size > effective_expected:
+            effective_expected = child_world_size
+
+    # Strategy 1: PID-based matching
     found_indices = []
+    active_indices = []  # GPUs with ANY compute processes
     for idx in search_indices:
         try:
             handle = pynvml.nvmlDeviceGetHandleByIndex(idx)
             procs = pynvml.nvmlDeviceGetComputeRunningProcesses(handle)
+            if procs:
+                active_indices.append(idx)
             for p in procs:
                 if p.pid in target_pids:
                     found_indices.append(idx)
@@ -155,6 +282,16 @@ def _discover_gpu_indices(proc_pid, pynvml, fallback_index=0):
         except Exception:
             continue
 
+    # Strategy 2: If PID matching found fewer than expected, use active GPU count.
+    # This handles the common case where DDP workers are running but their PIDs
+    # weren't in our process tree (e.g., process group isolation, timing issues).
+    if effective_expected is not None and len(found_indices) < effective_expected:
+        if len(active_indices) >= effective_expected:
+            return active_indices[:effective_expected]
+        # Even if active < expected, active is better than found
+        if len(active_indices) > len(found_indices):
+            return active_indices
+
     return found_indices if found_indices else [fallback_index]
 
 
@@ -310,12 +447,17 @@ def probe_command(
             discovery_attempts = 0
             max_discovery_attempts = 3  # Retry at samples 5, 15, 30
 
-            # Determine expected GPU count from environment for retry logic
+            # Determine expected GPU count from command line + environment
             expected_gpus = 1
+            launcher_gpus = _parse_launcher_gpu_count(command)
+            if launcher_gpus is not None and launcher_gpus > 1:
+                expected_gpus = launcher_gpus
             ws = os.environ.get("WORLD_SIZE", "").strip()
             if ws:
                 try:
-                    expected_gpus = max(1, int(ws))
+                    ws_int = max(1, int(ws))
+                    if ws_int > expected_gpus:
+                        expected_gpus = ws_int
                 except ValueError:
                     pass
 
@@ -329,7 +471,10 @@ def probe_command(
                     and proc.pid):
                     discovery_attempts += 1
                     try:
-                        discovered = _discover_gpu_indices(proc.pid, pynvml, fallback_index=gpu_index)
+                        discovered = _discover_gpu_indices(
+                            proc.pid, pynvml, fallback_index=gpu_index,
+                            expected_gpus=expected_gpus if expected_gpus > 1 else None,
+                        )
                         if len(discovered) > 1:
                             handles = []
                             pmap = []
@@ -456,9 +601,12 @@ def probe_command(
     if calibration_time_ref[0] is not None:
         cal_duration = round(calibration_time_ref[0] - start_time, 2)
 
-    # Environment-based fallback: if NVML discovery found fewer GPUs than
-    # WORLD_SIZE indicates, trust the environment. The probe may miss GPUs
+    # Final fallback: if NVML discovery found fewer GPUs than expected,
+    # trust the command-line / environment signals. The probe may miss GPUs
     # due to DDP per-rank CVD isolation or timing races.
+    launcher_count = _parse_launcher_gpu_count(command)
+    if launcher_count is not None and launcher_count > num_gpus_ref[0]:
+        num_gpus_ref[0] = launcher_count
     env_world = os.environ.get("WORLD_SIZE", "").strip()
     if env_world:
         try:

{alloc-0.0.5 → alloc-0.0.7}/src/alloc.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alloc
-Version: 0.0.5
+Version: 0.0.7
 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.5 → alloc-0.0.7}/src/alloc.egg-info/SOURCES.txt

@@ -43,11 +43,13 @@ tests/test_diagnosis_engine.py
 tests/test_diagnosis_rules.py
 tests/test_extractor_activation.py
 tests/test_ghost.py
+tests/test_ghost_degradation.py
 tests/test_init_from_org.py
 tests/test_interconnect.py
 tests/test_model_extractor.py
 tests/test_probe_hw.py
 tests/test_probe_multi.py
+tests/test_scan_auth.py
 tests/test_stability.py
 tests/test_upload.py
 tests/test_verdict.py

{alloc-0.0.5 → alloc-0.0.7}/tests/test_auth.py

@@ -223,6 +223,7 @@ def test_browser_login_saves_tokens(tmp_path: Path):
     env = {
         "HOME": str(tmp_path),
         "ALLOC_API_URL": "https://api.example.com",
+        "DISPLAY": ":0",  # prevent headless detection in CI
     }
 
     with patch("alloc.browser_auth.browser_login", return_value=mock_result):
@@ -248,6 +249,7 @@ def test_browser_login_with_azure_provider(tmp_path: Path):
     env = {
         "HOME": str(tmp_path),
         "ALLOC_API_URL": "https://api.example.com",
+        "DISPLAY": ":0",
     }
 
     with patch("alloc.browser_auth.browser_login", return_value=mock_result) as mock_bl:
@@ -263,7 +265,7 @@ def test_browser_login_with_azure_provider(tmp_path: Path):
 
 
 def test_browser_login_invalid_provider(tmp_path: Path):
-    env = {"HOME": str(tmp_path)}
+    env = {"HOME": str(tmp_path), "DISPLAY": ":0"}
     result = runner.invoke(
         app, ["login", "--browser", "--provider", "facebook"], env=env
     )
@@ -271,10 +273,19 @@ def test_browser_login_invalid_provider(tmp_path: Path):
     assert "Invalid --provider" in result.output
 
 
+def test_browser_login_headless_detection(tmp_path: Path):
+    """In headless/SSH environment, browser login should fail with guidance."""
+    env = {"HOME": str(tmp_path), "SSH_CLIENT": "1.2.3.4 1234 22", "DISPLAY": ":0"}
+    result = runner.invoke(app, ["login", "--browser"], env=env)
+    assert result.exit_code != 0
+    assert "Headless" in result.output or "token" in result.output
+
+
 def test_browser_login_timeout(tmp_path: Path):
     env = {
         "HOME": str(tmp_path),
         "ALLOC_API_URL": "https://api.example.com",
+        "DISPLAY": ":0",
     }
 
     with patch(

alloc-0.0.7/tests/test_ghost_degradation.py

@@ -0,0 +1,145 @@
+"""Tests for ghost structured degradation on distributed scripts."""
+
+from __future__ import annotations
+
+import json
+import os
+import tempfile
+from pathlib import Path
+from unittest.mock import patch
+
+from typer.testing import CliRunner
+
+from alloc.cli import app
+from alloc.model_extractor import ModelInfo, extract_model_info
+
+runner = CliRunner()
+
+
+def test_distributed_error_returns_structured_modelinfo():
+    """When extractor subprocess fails with a distributed keyword, return structured ModelInfo."""
+    sidecar_data = json.dumps({
+        "status": "error",
+        "error": "RuntimeError: torch.distributed.init_process_group requires MASTER_ADDR",
+    })
+
+    def _fake_subprocess_run(*args, **kwargs):
+        # Write the sidecar file
+        sidecar_path = args[0][3]  # [python, -m, alloc.extractor_runner, sidecar_path, script_path]
+        with open(sidecar_path, "w") as f:
+            f.write(sidecar_data)
+
+    # Create a dummy script
+    fd, script_path = tempfile.mkstemp(suffix=".py", prefix="alloc_test_dist_")
+    os.write(fd, b"import torch\ntorch.distributed.init_process_group('nccl')\n")
+    os.close(fd)
+
+    try:
+        with patch("subprocess.run", side_effect=_fake_subprocess_run):
+            info = extract_model_info(script_path)
+
+        assert info is not None
+        assert info.extraction_error == "distributed_entrypoint"
+        assert "distributed runtime" in info.extraction_detail
+        assert info.param_count == 0
+    finally:
+        os.unlink(script_path)
+
+
+def test_distributed_error_nccl_keyword():
+    """NCCL errors should be caught as distributed failures."""
+    sidecar_data = json.dumps({
+        "status": "error",
+        "error": "NCCL error: unhandled system error",
+    })
+
+    fd, script_path = tempfile.mkstemp(suffix=".py", prefix="alloc_test_nccl_")
+    os.write(fd, b"pass\n")
+    os.close(fd)
+
+    try:
+        def _fake_run(*args, **kwargs):
+            sidecar_path = args[0][3]
+            with open(sidecar_path, "w") as f:
+                f.write(sidecar_data)
+
+        with patch("subprocess.run", side_effect=_fake_run):
+            info = extract_model_info(script_path)
+
+        assert info is not None
+        assert info.extraction_error == "distributed_entrypoint"
+    finally:
+        os.unlink(script_path)
+
+
+def test_non_distributed_error_returns_none():
+    """Non-distributed errors should still return None (fall through to AST)."""
+    sidecar_data = json.dumps({
+        "status": "error",
+        "error": "ImportError: No module named 'custom_lib'",
+    })
+
+    fd, script_path = tempfile.mkstemp(suffix=".py", prefix="alloc_test_other_")
+    # Script with no from_pretrained so AST also returns None
+    os.write(fd, b"import custom_lib\n")
+    os.close(fd)
+
+    try:
+        def _fake_run(*args, **kwargs):
+            sidecar_path = args[0][3]
+            with open(sidecar_path, "w") as f:
+                f.write(sidecar_data)
+
+        with patch("subprocess.run", side_effect=_fake_run):
+            info = extract_model_info(script_path)
+
+        # Should be None because error is not distributed and AST won't find a model either
+        assert info is None
+    finally:
+        os.unlink(script_path)
+
+
+def test_ghost_cli_distributed_error_json(tmp_path: Path):
+    """ghost --json shows structured error for distributed scripts."""
+    script_path = tmp_path / "train_ddp.py"
+    script_path.write_text("import torch\ntorch.distributed.init_process_group('nccl')\n")
+
+    dist_info = ModelInfo(
+        param_count=0,
+        dtype="float16",
+        model_name=None,
+        method="execution",
+        extraction_error="distributed_entrypoint",
+        extraction_detail="Script requires a distributed runtime (e.g. torchrun). Run ghost on the model definition file instead of the launcher script.",
+    )
+
+    with patch("alloc.model_extractor.extract_model_info", return_value=dist_info):
+        result = runner.invoke(app, ["ghost", str(script_path), "--json"])
+
+    assert result.exit_code != 0
+    data = json.loads(result.output)
+    assert data["error"] == "distributed_entrypoint"
+    assert data["supported"] is False
+    assert "distributed runtime" in data["detail"]
+
+
+def test_ghost_cli_distributed_error_human(tmp_path: Path):
+    """ghost shows human-readable message with tip for distributed scripts."""
+    script_path = tmp_path / "train_ddp.py"
+    script_path.write_text("import torch\n")
+
+    dist_info = ModelInfo(
+        param_count=0,
+        dtype="float16",
+        model_name=None,
+        method="execution",
+        extraction_error="distributed_entrypoint",
+        extraction_detail="Script requires a distributed runtime (e.g. torchrun). Run ghost on the model definition file instead of the launcher script.",
+    )
+
+    with patch("alloc.model_extractor.extract_model_info", return_value=dist_info):
+        result = runner.invoke(app, ["ghost", str(script_path)])
+
+    assert result.exit_code != 0
+    assert "distributed runtime" in result.output
+    assert "model.py" in result.output  # tip about pointing to model file

{alloc-0.0.5 → alloc-0.0.7}/tests/test_probe_multi.py

@@ -4,7 +4,12 @@ from __future__ import annotations
 
 from unittest.mock import MagicMock, patch
 
-from alloc.probe import _discover_gpu_indices, _get_child_pids, ProbeResult
+from alloc.probe import (
+    _discover_gpu_indices,
+    _get_child_pids,
+    _parse_launcher_gpu_count,
+    ProbeResult,
+)
 
 
 def _mock_pynvml_multi_gpu(proc_pid, gpu_process_map):
@@ -112,3 +117,76 @@ def test_probe_result_defaults_single_gpu():
     result = ProbeResult()
     assert result.num_gpus_detected == 1
     assert result.process_map is None
+
+
+# ── Launcher command-line parsing ──
+
+
+def test_parse_torchrun_equals():
+    assert _parse_launcher_gpu_count(["torchrun", "--nproc_per_node=2", "train.py"]) == 2
+
+
+def test_parse_torchrun_space():
+    assert _parse_launcher_gpu_count(["torchrun", "--nproc_per_node", "4", "train.py"]) == 4
+
+
+def test_parse_torchrun_hyphen():
+    assert _parse_launcher_gpu_count(["torchrun", "--nproc-per-node=8", "train.py"]) == 8
+
+
+def test_parse_accelerate_equals():
+    assert _parse_launcher_gpu_count(["accelerate", "launch", "--num_processes=4", "train.py"]) == 4
+
+
+def test_parse_accelerate_hyphen_space():
+    assert _parse_launcher_gpu_count(["accelerate", "launch", "--num-processes", "8", "train.py"]) == 8
+
+
+def test_parse_deepspeed_equals():
+    assert _parse_launcher_gpu_count(["deepspeed", "--num_gpus=4", "train.py"]) == 4
+
+
+def test_parse_deepspeed_hyphen_space():
+    assert _parse_launcher_gpu_count(["deepspeed", "--num-gpus", "2", "train.py"]) == 2
+
+
+def test_parse_mpirun():
+    assert _parse_launcher_gpu_count(["mpirun", "-np", "8", "python", "train.py"]) == 8
+
+
+def test_parse_plain_python():
+    assert _parse_launcher_gpu_count(["python", "train.py"]) is None
+
+
+def test_parse_torch_distributed_launch():
+    assert _parse_launcher_gpu_count([
+        "python", "-m", "torch.distributed.launch", "--nproc_per_node=2", "train.py"
+    ]) == 2
+
+
+# ── Active-GPU fallback discovery ──
+
+
+def test_active_gpu_fallback_when_pid_mismatch():
+    """When PID matching fails but GPUs have active compute processes,
+    and expected_gpus matches, use active GPUs."""
+    mock = _mock_pynvml_multi_gpu(
+        proc_pid=1000,
+        gpu_process_map={0: [9999], 1: [8888]},  # PIDs don't match 1000
+    )
+    with patch("alloc.probe._get_child_pids", return_value=[]):
+        with patch("alloc.probe._read_child_env", return_value=None):
+            result = _discover_gpu_indices(1000, mock, fallback_index=0, expected_gpus=2)
+    assert result == [0, 1]
+
+
+def test_active_gpu_fallback_not_used_without_expected():
+    """Without expected_gpus, active GPU fallback is not used."""
+    mock = _mock_pynvml_multi_gpu(
+        proc_pid=1000,
+        gpu_process_map={0: [9999], 1: [8888]},
+    )
+    with patch("alloc.probe._get_child_pids", return_value=[]):
+        with patch("alloc.probe._read_child_env", return_value=None):
+            result = _discover_gpu_indices(1000, mock, fallback_index=0)
+    assert result == [0]  # Falls back to default

alloc-0.0.7/tests/test_scan_auth.py

@@ -0,0 +1,142 @@
+"""Tests for scan command 401 retry + /scans/cli fallback."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import httpx
+from typer.testing import CliRunner
+
+from alloc.cli import app
+
+runner = CliRunner()
+
+
+def _make_resp(status_code: int, body: dict, url: str = "https://api.example.com/scans"):
+    req = httpx.Request("POST", url)
+    return httpx.Response(
+        status_code,
+        request=req,
+        content=json.dumps(body).encode(),
+        headers={"content-type": "application/json"},
+    )
+
+
+def test_scan_401_refresh_retry(tmp_path: Path):
+    """On 401, refresh token and retry on /scans."""
+    resp_401 = _make_resp(401, {"detail": "unauthorized"})
+    resp_ok = _make_resp(200, {"vram_gb": 16.0, "configs": []})
+
+    mock_client = MagicMock()
+    mock_client.__enter__.return_value = mock_client
+    mock_client.__exit__.return_value = False
+    mock_client.post.side_effect = [resp_401, resp_ok]
+
+    cfg_file = tmp_path / ".alloc" / "config.json"
+    cfg_file.parent.mkdir(parents=True)
+    cfg_file.write_text(json.dumps({"token": "old-tok", "refresh_token": "rt"}))
+
+    env = {
+        "HOME": str(tmp_path),
+        "ALLOC_API_URL": "https://api.example.com",
+    }
+
+    with (
+        patch("httpx.Client", return_value=mock_client),
+        patch("alloc.cli.try_refresh_access_token", return_value="new-tok"),
+    ):
+        result = runner.invoke(app, ["scan", "--model", "llama-3-8b", "--json"], env=env)
+
+    assert result.exit_code == 0
+    assert mock_client.post.call_count == 2
+    # Second call should use refreshed token
+    second_call = mock_client.post.call_args_list[1]
+    assert "Bearer new-tok" in str(second_call)
+
+
+def test_scan_401_refresh_fails_fallback_public(tmp_path: Path):
+    """On 401 + refresh failure, fall back to /scans/cli with warning."""
+    resp_401 = _make_resp(401, {"detail": "unauthorized"})
+    resp_ok = _make_resp(200, {"vram_gb": 16.0, "configs": []},
+                          url="https://api.example.com/scans/cli")
+
+    mock_client = MagicMock()
+    mock_client.__enter__.return_value = mock_client
+    mock_client.__exit__.return_value = False
+    mock_client.post.side_effect = [resp_401, resp_ok]
+
+    cfg_file = tmp_path / ".alloc" / "config.json"
+    cfg_file.parent.mkdir(parents=True)
+    cfg_file.write_text(json.dumps({"token": "old-tok", "refresh_token": "rt"}))
+
+    env = {
+        "HOME": str(tmp_path),
+        "ALLOC_API_URL": "https://api.example.com",
+    }
+
+    with (
+        patch("httpx.Client", return_value=mock_client),
+        patch("alloc.cli.try_refresh_access_token", return_value=None),
+    ):
+        result = runner.invoke(app, ["scan", "--model", "llama-3-8b", "--json"], env=env)
+
+    assert result.exit_code == 0
+    assert mock_client.post.call_count == 2
+    # Second call should hit /scans/cli
+    second_url = str(mock_client.post.call_args_list[1])
+    assert "/scans/cli" in second_url
+
+
+def test_scan_401_fallback_warns_about_dropped_features(tmp_path: Path):
+    """Fallback to public scan warns user about lost org context."""
+    resp_401 = _make_resp(401, {"detail": "unauthorized"})
+    resp_ok = _make_resp(200, {"vram_gb": 16.0, "configs": []})
+
+    mock_client = MagicMock()
+    mock_client.__enter__.return_value = mock_client
+    mock_client.__exit__.return_value = False
+    mock_client.post.side_effect = [resp_401, resp_ok]
+
+    cfg_file = tmp_path / ".alloc" / "config.json"
+    cfg_file.parent.mkdir(parents=True)
+    cfg_file.write_text(json.dumps({"token": "old-tok", "refresh_token": "rt"}))
+
+    env = {
+        "HOME": str(tmp_path),
+        "ALLOC_API_URL": "https://api.example.com",
+    }
+
+    with (
+        patch("httpx.Client", return_value=mock_client),
+        patch("alloc.cli.try_refresh_access_token", return_value=None),
+    ):
+        # Non-JSON mode to see the warning message
+        result = runner.invoke(app, ["scan", "--model", "llama-3-8b"], env=env)
+
+    assert result.exit_code == 0
+    assert "expired" in result.output.lower() or "falling back" in result.output.lower()
+
+
+def test_scan_no_token_uses_public_directly(tmp_path: Path):
+    """Without a token, scan goes directly to /scans/cli."""
+    resp_ok = _make_resp(200, {"vram_gb": 16.0, "configs": []})
+
+    mock_client = MagicMock()
+    mock_client.__enter__.return_value = mock_client
+    mock_client.__exit__.return_value = False
+    mock_client.post.return_value = resp_ok
+
+    env = {
+        "HOME": str(tmp_path),
+        "ALLOC_API_URL": "https://api.example.com",
+    }
+
+    with patch("httpx.Client", return_value=mock_client):
+        result = runner.invoke(app, ["scan", "--model", "llama-3-8b", "--json"], env=env)
+
+    assert result.exit_code == 0
+    assert mock_client.post.call_count == 1
+    call_url = str(mock_client.post.call_args_list[0])
+    assert "/scans/cli" in call_url

alloc-0.0.5/src/alloc/__init__.py

@@ -1,11 +0,0 @@
-"""Alloc — GPU intelligence for ML training."""
-
-from __future__ import annotations
-
-__version__ = "0.0.5"
-
-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__"]