alloc 0.0.1__py3-none-any.whl

alloc/extractor_runner.py

@@ -0,0 +1,141 @@
+"""Standalone subprocess runner for model parameter extraction.
+
+This module is invoked by model_extractor.py as:
+    subprocess.run([python, "-m", "alloc.extractor_runner", sidecar_path, script_path])
+
+It imports the user's training script in an isolated process, finds
+nn.Module instances/classes, counts parameters, and writes results
+to a JSON sidecar file. Runs with CUDA_VISIBLE_DEVICES="" to prevent
+GPU allocation.
+
+Never imported directly by the rest of the CLI — only executed as __main__.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import json
+import os
+import sys
+
+
+def _patch_sys_exit():
+    """Prevent sys.exit() from killing the extractor."""
+    class _ExitCatch(SystemExit):
+        pass
+
+    def _fake_exit(code=0):
+        raise _ExitCatch(code)
+
+    sys.exit = _fake_exit
+    return _ExitCatch
+
+
+def _count_params(model):
+    """Count parameters in an nn.Module."""
+    total = 0
+    dtype_str = "float32"
+    try:
+        for i, (name, param) in enumerate(model.named_parameters()):
+            total += param.numel()
+            if i == 0:
+                dtype_str = str(param.dtype).replace("torch.", "")
+    except Exception:
+        pass
+    return total, dtype_str
+
+
+def main():
+    sidecar_path = sys.argv[1]
+    script_path = sys.argv[2]
+
+    result = {"status": "no_model"}
+
+    try:
+        import torch
+        import torch.nn as nn
+    except ImportError:
+        result = {"status": "error", "error": "torch not installed"}
+        with open(sidecar_path, "w") as f:
+            json.dump(result, f)
+        return
+
+    ExitCatch = _patch_sys_exit()
+
+    # Add script's directory to sys.path for relative imports
+    script_dir = os.path.dirname(os.path.abspath(script_path))
+    if script_dir not in sys.path:
+        sys.path.insert(0, script_dir)
+
+    # Import the user module
+    try:
+        spec = importlib.util.spec_from_file_location("__user_module__", script_path)
+        if spec is None or spec.loader is None:
+            result = {"status": "error", "error": "cannot load module"}
+            with open(sidecar_path, "w") as f:
+                json.dump(result, f)
+            return
+
+        module = importlib.util.module_from_spec(spec)
+        module.__name__ = "__user_module__"  # skip if __name__ == "__main__" guards
+        try:
+            spec.loader.exec_module(module)
+        except ExitCatch:
+            pass  # script called sys.exit(), that's fine
+        except SystemExit:
+            pass  # catch real SystemExit too
+    except Exception as e:
+        result = {"status": "error", "error": str(e)[:200]}
+        with open(sidecar_path, "w") as f:
+            json.dump(result, f)
+        return
+
+    # Search for nn.Module instances in module globals
+    models = []
+    for attr_name in dir(module):
+        try:
+            obj = getattr(module, attr_name)
+            if isinstance(obj, nn.Module):
+                count, dtype_str = _count_params(obj)
+                if count > 0:
+                    models.append((count, dtype_str, attr_name))
+        except Exception:
+            continue
+
+    # If no instances found, search for nn.Module subclasses and try instantiation
+    if not models:
+        for attr_name in dir(module):
+            try:
+                obj = getattr(module, attr_name)
+                if (isinstance(obj, type)
+                    and issubclass(obj, nn.Module)
+                    and obj is not nn.Module):
+                    try:
+                        instance = obj()
+                        count, dtype_str = _count_params(instance)
+                        if count > 0:
+                            models.append((count, dtype_str, attr_name))
+                    except Exception:
+                        continue
+            except Exception:
+                continue
+
+    if models:
+        # Pick largest model (most params = likely training target)
+        models.sort(key=lambda m: m[0], reverse=True)
+        best = models[0]
+        result = {
+            "status": "ok",
+            "param_count": best[0],
+            "dtype": best[1],
+            "model_name": best[2],
+        }
+    else:
+        result = {"status": "no_model"}
+
+    with open(sidecar_path, "w") as f:
+        json.dump(result, f)
+
+
+if __name__ == "__main__":
+    main()

alloc/ghost.py

@@ -0,0 +1,167 @@
+"""Ghost Scan — static VRAM analysis without executing the model.
+
+Two paths:
+  1. With torch: walk model.named_parameters() for exact sizing
+  2. Without torch: pure math from param_count (same formula as API engine)
+
+Never crashes user code. All exceptions are caught and logged.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, asdict
+from typing import Any, Optional
+
+# Canonical dtype → bytes mapping. Keep in sync with apps/api/src/engine/vram.py.
+BYTES_PER_DTYPE = {
+    "float32": 4,
+    "float16": 2,
+    "bfloat16": 2,
+    "int8": 1,
+    "int4": 0.5,
+    # Shorthand aliases (torch uses long names, API uses short names)
+    "fp32": 4,
+    "fp16": 2,
+    "bf16": 2,
+}
+
+# Adam: params (fp32) + momentum (fp32) + variance (fp32) = 3 states * 4 bytes
+# Keep in sync with OPTIMIZER_MULTIPLIER_FP16 in apps/api/src/engine/vram.py.
+OPTIMIZER_BYTES_PER_PARAM = 12
+
+# Buffer overhead: 10% of (weights + gradients + activations + optimizer)
+# for fragmentation and temp allocations.
+BUFFER_OVERHEAD_FACTOR = 0.1
+
+
+@dataclass
+class GhostReport:
+    """Result of a Ghost Scan — VRAM breakdown."""
+
+    param_count: int
+    param_count_b: float
+    dtype: str
+    weights_gb: float
+    gradients_gb: float
+    optimizer_gb: float
+    activations_gb: float
+    buffer_gb: float
+    total_gb: float
+    extraction_method: Optional[str] = None  # "execution", "ast", "manual"
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+
+def ghost(
+    model: Any = None,
+    *,
+    param_count: Optional[int] = None,
+    param_count_b: Optional[float] = None,
+    dtype: str = "float16",
+    batch_size: int = 32,
+    seq_length: int = 2048,
+    hidden_dim: int = 4096,
+) -> GhostReport:
+    """Run a Ghost Scan on a model or param count.
+
+    Usage:
+        # With a PyTorch model
+        report = alloc.ghost(model)
+
+        # Without torch — just param count
+        report = alloc.ghost(param_count_b=7.0)
+
+    Never raises. If something goes wrong, returns a best-effort report.
+    """
+    try:
+        return _ghost_impl(
+            model=model,
+            param_count=param_count,
+            param_count_b=param_count_b,
+            dtype=dtype,
+            batch_size=batch_size,
+            seq_length=seq_length,
+            hidden_dim=hidden_dim,
+        )
+    except Exception:
+        # Never crash user code
+        count = param_count or int((param_count_b or 0) * 1e9) or 0
+        return GhostReport(
+            param_count=count,
+            param_count_b=count / 1e9,
+            dtype=dtype,
+            weights_gb=0.0,
+            gradients_gb=0.0,
+            optimizer_gb=0.0,
+            activations_gb=0.0,
+            buffer_gb=0.0,
+            total_gb=0.0,
+        )
+
+
+def _ghost_impl(
+    model: Any,
+    param_count: Optional[int],
+    param_count_b: Optional[float],
+    dtype: str,
+    batch_size: int,
+    seq_length: int,
+    hidden_dim: int,
+) -> GhostReport:
+    """Core Ghost implementation."""
+    resolved_count = 0
+    resolved_dtype = dtype
+
+    if model is not None:
+        # Path 1: Walk torch model parameters
+        resolved_count, resolved_dtype = _count_from_model(model)
+    elif param_count is not None:
+        resolved_count = param_count
+    elif param_count_b is not None:
+        resolved_count = int(param_count_b * 1e9)
+
+    if resolved_count <= 0:
+        resolved_count = 0
+
+    bytes_per_param = BYTES_PER_DTYPE.get(resolved_dtype, 2)
+    to_gb = 1.0 / (1024 ** 3)
+
+    weights_bytes = resolved_count * bytes_per_param
+    gradients_bytes = resolved_count * bytes_per_param
+    optimizer_bytes = resolved_count * OPTIMIZER_BYTES_PER_PARAM
+    activations_bytes = batch_size * seq_length * hidden_dim * bytes_per_param
+
+    weights_gb = weights_bytes * to_gb
+    gradients_gb = gradients_bytes * to_gb
+    optimizer_gb = optimizer_bytes * to_gb
+    activations_gb = activations_bytes * to_gb
+    buffer_gb = BUFFER_OVERHEAD_FACTOR * (weights_gb + gradients_gb + optimizer_gb + activations_gb)
+    total_gb = weights_gb + gradients_gb + optimizer_gb + activations_gb + buffer_gb
+
+    return GhostReport(
+        param_count=resolved_count,
+        param_count_b=round(resolved_count / 1e9, 3),
+        dtype=resolved_dtype,
+        weights_gb=round(weights_gb, 2),
+        gradients_gb=round(gradients_gb, 2),
+        optimizer_gb=round(optimizer_gb, 2),
+        activations_gb=round(activations_gb, 2),
+        buffer_gb=round(buffer_gb, 2),
+        total_gb=round(total_gb, 2),
+    )
+
+
+def _count_from_model(model: Any) -> tuple:
+    """Extract param count and dtype from a torch model. Returns (count, dtype_str)."""
+    total = 0
+    dtype_str = "float16"
+    try:
+        for name, param in model.named_parameters():
+            total += param.numel()
+            # Use dtype of the first parameter
+            if total == param.numel():
+                dtype_str = str(param.dtype).replace("torch.", "")
+    except Exception:
+        pass
+    return total, dtype_str

alloc/model_extractor.py

@@ -0,0 +1,170 @@
+"""Model extraction — get param count from user scripts.
+
+Three extraction methods with fallback chain:
+  1. Manual override (--param-count-b flag) → instant, no file needed
+  2. Subprocess execution → import module, find nn.Module, count params
+  3. AST parsing → find from_pretrained() calls, match against known models
+
+Never crashes. Returns None when extraction fails.
+"""
+
+from __future__ import annotations
+
+import ast
+import json
+import os
+import subprocess
+import sys
+import tempfile
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class ModelInfo:
+    """Result of model extraction."""
+
+    param_count: int            # raw parameter count
+    dtype: str                  # "float32", "float16", etc.
+    model_name: Optional[str]   # class name if found
+    method: str                 # "execution" | "ast" | "manual"
+
+
+def extract_model_info(
+    script: str,
+    *,
+    timeout: int = 60,
+    param_count_b: Optional[float] = None,
+) -> Optional[ModelInfo]:
+    """Extract model info from a script with fallback chain.
+
+    Returns ModelInfo or None if extraction fails.
+    """
+    # 1. Manual override — no file needed
+    if param_count_b is not None:
+        return ModelInfo(
+            param_count=int(param_count_b * 1e9),
+            dtype="float16",
+            model_name=None,
+            method="manual",
+        )
+
+    # 2. File must exist for remaining methods
+    if not os.path.isfile(script):
+        return None
+
+    # 3. Try subprocess execution
+    result = _extract_via_subprocess(script, timeout=timeout)
+    if result is not None:
+        return result
+
+    # 4. Try AST parsing
+    result = _extract_via_ast(script)
+    if result is not None:
+        return result
+
+    return None
+
+
+# ---------------------------------------------------------------------------
+# Subprocess execution
+# ---------------------------------------------------------------------------
+
+def _extract_via_subprocess(
+    script: str,
+    *,
+    timeout: int = 60,
+) -> Optional[ModelInfo]:
+    """Execute user script in isolated subprocess, extract model info.
+
+    Invokes alloc.extractor_runner as a module in a subprocess with
+    CUDA_VISIBLE_DEVICES="" to prevent GPU allocation. Results are
+    communicated via a JSON sidecar file.
+    """
+    sidecar_fd = None
+    sidecar_path = None
+
+    try:
+        # Create sidecar file for IPC
+        sidecar_fd, sidecar_path = tempfile.mkstemp(suffix=".json", prefix="alloc_extract_")
+        os.close(sidecar_fd)
+        sidecar_fd = None
+
+        script_abs = os.path.abspath(script)
+
+        env = os.environ.copy()
+        env["CUDA_VISIBLE_DEVICES"] = ""  # prevent GPU allocation
+
+        subprocess.run(
+            [sys.executable, "-m", "alloc.extractor_runner", sidecar_path, script_abs],
+            timeout=timeout,
+            capture_output=True,
+            env=env,
+        )
+
+        # Read sidecar result
+        try:
+            with open(sidecar_path, "r") as f:
+                data = json.load(f)
+        except (json.JSONDecodeError, FileNotFoundError, OSError):
+            return None
+
+        if data.get("status") == "ok":
+            return ModelInfo(
+                param_count=data["param_count"],
+                dtype=data.get("dtype", "float32"),
+                model_name=data.get("model_name"),
+                method="execution",
+            )
+
+        return None
+
+    except subprocess.TimeoutExpired:
+        return None
+    except Exception:
+        return None
+    finally:
+        # Clean up sidecar temp file
+        if sidecar_path and os.path.exists(sidecar_path):
+            try:
+                os.unlink(sidecar_path)
+            except OSError:
+                pass
+
+
+# ---------------------------------------------------------------------------
+# AST-based extraction (static fallback)
+# ---------------------------------------------------------------------------
+
+from alloc.model_registry import KNOWN_MODELS
+
+
+def _extract_via_ast(script: str) -> Optional[ModelInfo]:
+    """Parse script AST, find from_pretrained() calls, match against known models."""
+    try:
+        with open(script, "r") as f:
+            source = f.read()
+        tree = ast.parse(source)
+    except (SyntaxError, OSError):
+        return None
+
+    # Walk AST looking for .from_pretrained("model_name") calls
+    for node in ast.walk(tree):
+        if not isinstance(node, ast.Call):
+            continue
+
+        # Match *.from_pretrained(...)
+        func = node.func
+        if isinstance(func, ast.Attribute) and func.attr == "from_pretrained":
+            if node.args and isinstance(node.args[0], ast.Constant) and isinstance(node.args[0].value, str):
+                model_id = node.args[0].value
+                param_count = KNOWN_MODELS.get(model_id)
+                if param_count is not None:
+                    return ModelInfo(
+                        param_count=param_count,
+                        dtype="float16",
+                        model_name=model_id,
+                        method="ast",
+                    )
+
+    return None

alloc/model_registry.py

@@ -0,0 +1,138 @@
+"""Canonical model registry — single source of truth for known model param counts.
+
+Used by:
+  - model_extractor.py (AST fallback: HuggingFace model IDs → raw param count)
+  - cli.py (alloc scan: short model names → param count in billions)
+
+Add new models here. Do NOT maintain separate lookup tables elsewhere.
+"""
+
+from __future__ import annotations
+
+from typing import Dict, Optional
+
+# HuggingFace model IDs → raw parameter count (int).
+# Used by AST extractor to match from_pretrained("model_id") calls.
+KNOWN_MODELS: Dict[str, int] = {
+    # GPT-2
+    "gpt2": int(124e6),
+    "gpt2-medium": int(355e6),
+    "gpt2-large": int(774e6),
+    "gpt2-xl": int(1.5e9),
+    # BERT
+    "bert-base-uncased": int(110e6),
+    "bert-large-uncased": int(340e6),
+    # Llama
+    "meta-llama/Llama-2-7b-hf": int(7e9),
+    "meta-llama/Llama-2-13b-hf": int(13e9),
+    "meta-llama/Llama-2-70b-hf": int(70e9),
+    # Mistral / Mixtral
+    "mistralai/Mistral-7B-v0.1": int(7.24e9),
+    "mistralai/Mixtral-8x7B-v0.1": int(46.7e9),
+    # Falcon
+    "tiiuae/falcon-7b": int(7e9),
+    "tiiuae/falcon-40b": int(40e9),
+    # Bloom
+    "bigscience/bloom-560m": int(560e6),
+    "bigscience/bloom-1b7": int(1.7e9),
+    "bigscience/bloom-7b1": int(7.1e9),
+    # GPT-Neo / GPT-J
+    "EleutherAI/gpt-neo-125M": int(125e6),
+    "EleutherAI/gpt-neo-1.3B": int(1.3e9),
+    "EleutherAI/gpt-neo-2.7B": int(2.7e9),
+    "EleutherAI/gpt-j-6B": int(6e9),
+    # Phi
+    "microsoft/phi-2": int(2.78e9),
+    # Gemma
+    "google/gemma-2b": int(2.51e9),
+    "google/gemma-7b": int(8.54e9),
+    # Qwen
+    "Qwen/Qwen-7B": int(7.72e9),
+    "Qwen/Qwen-14B": int(14.2e9),
+    "Qwen/Qwen-72B": int(72.7e9),
+    # DeepSeek
+    "deepseek-ai/deepseek-llm-7b-base": int(6.9e9),
+    "deepseek-ai/deepseek-llm-67b-base": int(67e9),
+    # T5
+    "t5-small": int(60e6),
+    "t5-base": int(220e6),
+    "t5-large": int(770e6),
+    "google/t5-xl-lm-adapt": int(3e9),
+    "google/t5-xxl-lm-adapt": int(11e9),
+    # Vision
+    "google/vit-base-patch16-224": int(86e6),
+    "google/vit-large-patch16-224": int(307e6),
+    # Whisper
+    "openai/whisper-small": int(244e6),
+    "openai/whisper-medium": int(769e6),
+    "openai/whisper-large-v3": int(1.55e9),
+}
+
+# Short CLI names → param count in billions.
+# Used by `alloc scan --model <name>` for quick lookups without a script.
+# Maps normalized lowercase names to billions (float).
+_CLI_MODEL_PARAMS: Dict[str, float] = {
+    # Llama
+    "llama-3-70b": 70.0,
+    "llama-3-8b": 8.03,
+    "llama-2-70b": 70.0,
+    "llama-2-13b": 13.0,
+    "llama-2-7b": 7.0,
+    # Mistral / Mixtral
+    "mistral-7b": 7.24,
+    "mixtral-8x7b": 46.7,
+    # GPT-2
+    "gpt2": 0.124,
+    "gpt2-medium": 0.355,
+    "gpt2-large": 0.774,
+    "gpt2-xl": 1.5,
+    # BERT
+    "bert-base": 0.110,
+    "bert-large": 0.340,
+    # T5
+    "t5-small": 0.060,
+    "t5-base": 0.220,
+    "t5-large": 0.770,
+    "t5-xl": 3.0,
+    "t5-xxl": 11.0,
+    # Falcon
+    "falcon-7b": 7.0,
+    "falcon-40b": 40.0,
+    # Phi
+    "phi-2": 2.78,
+    # Gemma
+    "gemma-2b": 2.51,
+    "gemma-7b": 8.54,
+    # Qwen
+    "qwen-7b": 7.72,
+    "qwen-14b": 14.2,
+    "qwen-72b": 72.7,
+    # DeepSeek
+    "deepseek-7b": 6.9,
+    "deepseek-67b": 67.0,
+    # Vision
+    "vit-base": 0.086,
+    "vit-large": 0.307,
+    # Whisper
+    "whisper-small": 0.244,
+    "whisper-medium": 0.769,
+    "whisper-large": 1.55,
+    # Bloom
+    "bloom-560m": 0.560,
+    "bloom-1b7": 1.7,
+    "bloom-7b1": 7.1,
+    # GPT-Neo / GPT-J
+    "gpt-neo-125m": 0.125,
+    "gpt-neo-1.3b": 1.3,
+    "gpt-neo-2.7b": 2.7,
+    "gpt-j-6b": 6.0,
+}
+
+
+def lookup_model_params(model: str) -> Optional[float]:
+    """Look up model param count (in billions) by short CLI name.
+
+    Returns None if model not found.
+    """
+    normalized = model.lower().strip()
+    return _CLI_MODEL_PARAMS.get(normalized)