alloc 0.0.5__tar.gz → 0.0.6__tar.gz

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

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

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

@@ -2,7 +2,12 @@
 
 from __future__ import annotations
 
-__version__ = "0.0.5"
+import warnings as _warnings
+_warnings.filterwarnings("ignore", category=FutureWarning, module="pynvml")
+_warnings.filterwarnings("ignore", category=DeprecationWarning, module="pynvml")
+del _warnings
+
+__version__ = "0.0.6"
 
 from alloc.ghost import ghost, GhostReport
 from alloc.callbacks import AllocCallback as HuggingFaceCallback

{alloc-0.0.5 → alloc-0.0.6}/src/alloc/cli.py

@@ -16,8 +16,15 @@ from __future__ import annotations
 
 import os
 import sys
+import warnings
 from typing import List, Optional
 
+# Suppress noisy third-party warnings globally — pynvml deprecation and
+# urllib3 LibreSSL warnings clutter every CLI command on affected systems.
+warnings.filterwarnings("ignore", category=FutureWarning, module="pynvml")
+warnings.filterwarnings("ignore", category=DeprecationWarning, module="pynvml")
+warnings.filterwarnings("ignore", message=".*LibreSSL.*", module="urllib3")
+
 import typer
 from rich.console import Console
 
@@ -2107,7 +2114,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 +2159,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.6}/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.6}/src/alloc/model_extractor.py

@@ -99,6 +99,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],

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

@@ -83,6 +83,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 +180,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 +226,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 +275,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 +440,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 +464,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 +594,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.6}/src/alloc.egg-info/PKG-INFO

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