Merge PR #15: bound decompilation concurrency + auth preflight

fix: bound decompilation concurrency (#12) + add auth preflight (#14)

processors/ghidra_decompile/ghidra_decompile.py

@@ -19,6 +19,63 @@ from deepzero.engine.stage import (
 
 log = logging.getLogger("deepzero.processor.ghidra")
 
+# each ghidra headless worker is a separate JVM holding a full program
+# database in memory (roughly this many GiB during analysis). auto worker
+# counts are derived from this so `parallel: 0` cannot exhaust RAM.
+_GHIDRA_GB_PER_WORKER = 4.0
+_GHIDRA_MAX_AUTO_WORKERS = 16
+
+
+def _total_ram_gb() -> float | None:
+    """best-effort total physical RAM in GiB, or None if it can't be measured."""
+    try:
+        names = getattr(os, "sysconf_names", {})
+        if "SC_PHYS_PAGES" in names and "SC_PAGE_SIZE" in names:
+            pages = os.sysconf("SC_PHYS_PAGES")
+            page = os.sysconf("SC_PAGE_SIZE")
+            if pages > 0 and page > 0:
+                return pages * page / (1024**3)
+    except (ValueError, OSError, AttributeError):
+        pass
+    if sys.platform == "win32":
+        try:
+            import ctypes
+
+            class _MemStatus(ctypes.Structure):
+                _fields_ = [
+                    ("dwLength", ctypes.c_ulong),
+                    ("dwMemoryLoad", ctypes.c_ulong),
+                    ("ullTotalPhys", ctypes.c_ulonglong),
+                    ("ullAvailPhys", ctypes.c_ulonglong),
+                    ("ullTotalPageFile", ctypes.c_ulonglong),
+                    ("ullAvailPageFile", ctypes.c_ulonglong),
+                    ("ullTotalVirtual", ctypes.c_ulonglong),
+                    ("ullAvailVirtual", ctypes.c_ulonglong),
+                    ("ullAvailExtendedVirtual", ctypes.c_ulonglong),
+                ]
+
+            stat = _MemStatus()
+            stat.dwLength = ctypes.sizeof(_MemStatus)
+            if ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(stat)):
+                return stat.ullTotalPhys / (1024**3)
+        except (OSError, AttributeError, ValueError):
+            pass
+    return None
+
+
+def _auto_ghidra_workers(cpu: int | None = None, ram_gb: float | None = None) -> int:
+    """a memory-safe default worker count for ghidra headless.
+
+    bounded by cpu count, available RAM (~4 GiB/worker), and a hard cap. falls
+    back to a conservative cpu-only bound when RAM can't be measured.
+    """
+    cpu = cpu or os.cpu_count() or 4
+    ram = ram_gb if ram_gb is not None else _total_ram_gb()
+    if ram is None:
+        return max(1, min(cpu, 4))
+    by_ram = max(1, int(ram // _GHIDRA_GB_PER_WORKER))
+    return max(1, min(cpu, by_ram, _GHIDRA_MAX_AUTO_WORKERS))
+
 
 class GhidraDecompile(MapProcessor):
     description = (
@@ -34,6 +91,15 @@ class GhidraDecompile(MapProcessor):
         max_depth: int | None = None
         ghidra_install_dir: str = ""
         java_home: str = ""
+        # ceiling on auto (parallel: 0) concurrency. 0 = derive a memory-safe
+        # default (~4 GiB RAM per worker; each worker is a full JVM). raise it
+        # to use more of a large machine, e.g. max_parallel: 24.
+        max_parallel: int = 0
+
+    def max_parallelism(self) -> int | None:
+        if self.config.max_parallel and self.config.max_parallel > 0:
+            return self.config.max_parallel
+        return _auto_ghidra_workers()
 
     def validate(self, ctx: ProcessorContext) -> list[str]:
         if not self.config.ghidra_install_dir:

src/deepzero/cli.py

@@ -152,8 +152,19 @@ def main(ctx: click.Context):
 @click.option("--work-dir", "-w", default=None, help="work directory override")
 @click.option("--verbose", "-v", is_flag=True, help="verbose logging")
 @click.option("--clean", is_flag=True, help="permanently delete previous run data and start fresh")
+@click.option(
+    "--preflight",
+    is_flag=True,
+    help="verify the LLM backend is authenticated before running (makes one tiny call)",
+)
 def run(
-    target: str, pipeline: str, model: str | None, work_dir: str | None, verbose: bool, clean: bool
+    target: str,
+    pipeline: str,
+    model: str | None,
+    work_dir: str | None,
+    verbose: bool,
+    clean: bool,
+    preflight: bool,
 ):
     """run a pipeline against a target file or directory (resumes automatically)"""
     _setup_logging(verbose)
@@ -222,6 +233,16 @@ def run(
 
     runner, llm = _build_runner(pipeline_def, dashboard=dashboard)
 
+    if preflight and llm is not None:
+        console.print("[dim]preflight: checking LLM backend authentication...[/]")
+        problem = llm.check_auth()
+        if problem:
+            console.print(
+                f"[bold red]X ERROR[/]: LLM backend '{pipeline_def.model}' is not ready: {problem}"
+            )
+            raise SystemExit(1)
+        console.print("[green]\\[ok][/] LLM backend authenticated")
+
     if is_resume:
         run_state = existing_run
         run_state.status = RunStatus.RUNNING

src/deepzero/engine/backends/base.py

@@ -96,6 +96,12 @@ class LLMBackend(ABC):
         m = (model or "").strip()
         return m.split("/", 1)[1].strip() if "/" in m else ""
 
+    def check_auth(self) -> str | None:
+        """opt-in readiness probe. return None if the backend can authenticate,
+        else a human-readable reason. may make one minimal call. default assumes
+        ok - api-key presence is already checked by validate_model()."""
+        return None
+
     @property
     def provider_name(self) -> str:
         return self.scheme or "unknown"
@@ -290,3 +296,17 @@ class CLIAgentBackend(LLMBackend):
         if proc.returncode != 0:
             return f"exited {proc.returncode}: {(proc.stderr or '').strip()[:200]}"
         return None
+
+    def check_auth(self) -> str | None:
+        """definitive readiness probe: sends a minimal prompt and reports an
+        auth failure. spends a tiny amount of usage, so it is opt-in (e.g.
+        `deepzero run --preflight`). transient errors (rate limit, network) are
+        not treated as auth failures and do not block the run."""
+        try:
+            self.raw_complete([{"role": "user", "content": "ping"}], timeout=60)
+        except (BackendAuthError, BackendNotFoundError) as exc:
+            return str(exc)
+        except BackendError:
+            # transient (rate limit, timeout, network) - not an auth problem
+            return None
+        return None

src/deepzero/engine/llm.py

@@ -87,6 +87,11 @@ class LLMProvider:
 
         raise RuntimeError("exhausted retries without raising")
 
+    def check_auth(self) -> str | None:
+        """opt-in readiness probe; delegates to the backend. returns None if
+        ready, else a human-readable reason."""
+        return self.backend.check_auth()
+
     @property
     def provider_name(self) -> str:
         return self.backend.provider_name

src/deepzero/engine/runner.py

@@ -397,6 +397,24 @@ class PipelineRunner:
 
     # -- map execution --
 
+    def _resolve_parallelism(self, spec: StageSpec, processor: Processor) -> int:
+        # explicit `parallel: N` is honored as-is; `parallel: 0` auto-scales to
+        # cpu count but is clamped to the processor's declared ceiling (if any)
+        # so heavy stages don't spawn one worker per core.
+        parallelism = spec.parallel
+        if parallelism > 0:
+            return parallelism
+        parallelism = os.cpu_count() or 4
+        ceiling = processor.max_parallelism()
+        if ceiling is not None and 0 < ceiling < parallelism:
+            log.info(
+                "%s: capping auto concurrency to %d workers (processor limit)", spec.name, ceiling
+            )
+            parallelism = ceiling
+        else:
+            log.debug("%s: auto-scaled to %d workers", spec.name, parallelism)
+        return parallelism
+
     def _run_map(
         self,
         processor: MapProcessor,
@@ -416,10 +434,7 @@ class PipelineRunner:
         if not pending:
             return
 
-        parallelism = spec.parallel
-        if parallelism <= 0:
-            parallelism = os.cpu_count() or 4
-            log.debug("%s: auto-scaled to %d workers", spec.name, parallelism)
+        parallelism = self._resolve_parallelism(spec, processor)
 
         if parallelism <= 1:
             dirty: list[SampleState] = []

src/deepzero/engine/stage.py

@@ -265,6 +265,15 @@ class Processor(ABC):
         d.mkdir(parents=True, exist_ok=True)
         return d
 
+    def max_parallelism(self) -> int | None:
+        # optional hard ceiling on AUTO (parallel: 0) concurrency for this
+        # processor. return an int to stop auto-scaling from using every core;
+        # None means no ceiling. heavy processors that spawn external, memory-
+        # hungry workers (e.g. one JVM per sample) override this so a default
+        # `parallel: 0` cannot exhaust the machine. an explicit `parallel: N`
+        # in the pipeline is always honored as-is.
+        return None
+
     def setup(self, global_config: dict[str, Any]) -> None:
         # called once before pipeline execution begins
         pass

tests/test_claude_code_backend.py

@@ -308,3 +308,36 @@ class TestValidation:
         with patch.dict("os.environ", {}, clear=True):
             with patch(_FIND, return_value="/usr/bin/claude"):
                 assert self._validate("claude-code") == []
+
+
+class TestAuthPreflight:
+    def test_check_auth_ok_when_authenticated(self):
+        with patch("subprocess.run", return_value=_completed(_result_json())):
+            assert _backend().check_auth() is None
+
+    def test_check_auth_reports_401(self):
+        body = _result_json(
+            is_error=True, api_error_status=401, result="OAuth access token has been revoked"
+        )
+        with patch("subprocess.run", return_value=_completed(body)):
+            msg = _backend().check_auth()
+        assert msg and "not authenticated" in msg
+
+    def test_check_auth_does_not_block_on_transient(self):
+        # a rate limit during preflight is not an auth failure
+        body = _result_json(is_error=True, api_error_status=429, result="rate limit")
+        with patch("subprocess.run", return_value=_completed(body)):
+            assert _backend().check_auth() is None
+
+    def test_provider_check_auth_delegates(self):
+        with patch(_FIND, return_value="/usr/bin/claude"):
+            provider = LLMProvider("claude-code")
+        with patch("subprocess.run", return_value=_completed(_result_json())):
+            assert provider.check_auth() is None
+        with patch(
+            "subprocess.run",
+            return_value=_completed(
+                _result_json(is_error=True, api_error_status=401, result="revoked")
+            ),
+        ):
+            assert provider.check_auth() is not None

tests/test_ghidra_decompile.py

@@ -4,7 +4,37 @@ from unittest.mock import patch
 
 from deepzero.engine.stage import ProcessorContext, StageSpec
 from deepzero.engine.state import StageOutput
-from processors.ghidra_decompile.ghidra_decompile import GhidraDecompile
+from processors.ghidra_decompile.ghidra_decompile import (
+    _GHIDRA_MAX_AUTO_WORKERS,
+    GhidraDecompile,
+    _auto_ghidra_workers,
+)
+
+
+class TestDecompileParallelism:
+    def test_ram_caps_workers(self):
+        # each worker ~4 GiB; bounded by cpu, ram, and the hard cap
+        assert _auto_ghidra_workers(cpu=32, ram_gb=128) == _GHIDRA_MAX_AUTO_WORKERS  # hard cap
+        assert _auto_ghidra_workers(cpu=32, ram_gb=16) == 4  # 16 // 4
+        assert _auto_ghidra_workers(cpu=2, ram_gb=128) == 2  # cpu-bound
+        assert _auto_ghidra_workers(cpu=8, ram_gb=64) == 8  # cpu-bound
+
+    def test_unknown_ram_is_conservative(self):
+        with patch("processors.ghidra_decompile.ghidra_decompile._total_ram_gb", return_value=None):
+            assert _auto_ghidra_workers(cpu=64) == 4
+
+    def test_never_zero(self):
+        assert _auto_ghidra_workers(cpu=1, ram_gb=1) == 1
+
+    def test_config_max_parallel_overrides_auto(self):
+        proc = GhidraDecompile(
+            StageSpec(name="decompile", processor="ghidra_decompile", config={"max_parallel": 24})
+        )
+        assert proc.max_parallelism() == 24
+
+    def test_auto_ceiling_is_bounded_when_unset(self):
+        proc = GhidraDecompile(StageSpec(name="decompile", processor="ghidra_decompile", config={}))
+        assert 1 <= proc.max_parallelism() <= _GHIDRA_MAX_AUTO_WORKERS
 
 
 def _make_ctx(tmp_path, config=None, global_config=None):

tests/test_runner.py

@@ -302,3 +302,35 @@ class TestPipelineRunner:
         for s in samples:
             if "late_fail" in s.history:
                 assert s.history["late_fail"].status != StageStatus.FAILED
+
+
+class TestResolveParallelism:
+    # _resolve_parallelism uses no instance state, so call it unbound with None
+    class _CeilingProc(MapProcessor):
+        def max_parallelism(self):
+            return 3
+
+        def process(self, ctx, entry):
+            return ProcessorResult.ok()
+
+    class _NoCeilingProc(MapProcessor):
+        def process(self, ctx, entry):
+            return ProcessorResult.ok()
+
+    def test_auto_is_capped_to_processor_ceiling(self, monkeypatch):
+        monkeypatch.setattr("os.cpu_count", lambda: 32)
+        proc = self._CeilingProc(StageSpec(name="d", processor="x"))
+        spec = StageSpec(name="d", processor="x", parallel=0)
+        assert PipelineRunner._resolve_parallelism(None, spec, proc) == 3
+
+    def test_explicit_parallel_is_honored_over_ceiling(self, monkeypatch):
+        monkeypatch.setattr("os.cpu_count", lambda: 32)
+        proc = self._CeilingProc(StageSpec(name="d", processor="x"))
+        spec = StageSpec(name="d", processor="x", parallel=20)
+        assert PipelineRunner._resolve_parallelism(None, spec, proc) == 20
+
+    def test_auto_uses_cpu_when_no_ceiling(self, monkeypatch):
+        monkeypatch.setattr("os.cpu_count", lambda: 8)
+        proc = self._NoCeilingProc(StageSpec(name="d", processor="x"))
+        spec = StageSpec(name="d", processor="x", parallel=0)
+        assert PipelineRunner._resolve_parallelism(None, spec, proc) == 8