From 862e43f46cd587231a6c7af36b7ff4400107fdc9 Mon Sep 17 00:00:00 2001 From: abhay-codes07 Date: Tue, 25 Aug 2026 00:06:14 +0530 Subject: [PATCH] feat(extract): add --memory-limit-mb so a budget overrun aborts instead of an OOM kill (#3011) `graphify extract` inside a memory-limited container could grow past the cgroup allowance and be OOM-killed: --max-workers bounds the AST pool, but the later JS/TS resolution passes retain source buffers and syntax trees for the whole corpus, and GRAPHIFY_REBUILD_MEMORY_LIMIT_MB only ever applied to hook/watch rebuilds. The kill left no graphify-specific failure, no stable exit status, and whatever had been written behind. `--memory-limit-mb N` / `GRAPHIFY_MEMORY_LIMIT_MB=N` on `extract`, `update` and the bare `graphify ` form: * caps the CLI process with setrlimit (RLIMIT_AS; RLIMIT_DATA on macOS) and, through a pool initializer, every extraction worker - workers start fresh under `spawn`, so they read the cap from the environment; * lets MemoryError propagate where the pipeline used to demote it: _safe_extract recorded it as a skipped file, and the pool's per-future handler warned and retried the file in-process, which would hit the same wall - either way a run could finish and publish a graph silently missing whatever came after; * reports the configured limit, the phase, and the observed peak, exits with status 3 (distinct from 1 = extraction failed, 2 = bad arguments) and writes no graph.json - the previous graph is left untouched; * refuses a malformed value (exit 2) rather than silently running with no budget, and on Windows says the budget cannot be enforced and continues, rather than pretending. The enforcement lives in a dependency-free graphify.memory_budget; watch._apply_resource_limits now delegates to it and also honours the general variable, with the hook-specific one keeping precedence. --- README.md | 2 + graphify/cli.py | 98 +++++++++- graphify/extract.py | 28 ++- graphify/memory_budget.py | 175 ++++++++++++++++++ graphify/watch.py | 18 +- tests/test_memory_budget.py | 346 ++++++++++++++++++++++++++++++++++++ 6 files changed, 654 insertions(+), 13 deletions(-) create mode 100644 graphify/memory_budget.py create mode 100644 tests/test_memory_budget.py diff --git a/README.md b/README.md index 0c14d207c..f06688426 100644 --- a/README.md +++ b/README.md @@ -521,6 +521,7 @@ These are only needed for **headless / CI extraction** (`graphify extract`). Whe | `AZURE_OPENAI_DEPLOYMENT` or `GRAPHIFY_AZURE_MODEL` | Azure deployment name | optional — default `gpt-4o` | | `AWS_*` / `~/.aws/credentials` | AWS Bedrock — standard credential chain | `--backend bedrock` (no API key, uses IAM) | | `GRAPHIFY_MAX_WORKERS` | AST parallelism thread count | optional — also `--max-workers` flag | +| `GRAPHIFY_MEMORY_LIMIT_MB` | Memory budget for `extract` / `update`: caps the CLI process and its extraction workers with `setrlimit`; exceeding it aborts with exit status 3 and no partial `graph.json` (Linux/macOS; reported as unenforceable on Windows) | optional — also `--memory-limit-mb` flag; set it below the container's cgroup limit | | `GRAPHIFY_MAX_OUTPUT_TOKENS` | Raise output cap for dense corpora | optional — e.g. `32768` for large files | | `GRAPHIFY_API_TIMEOUT` | Per-call timeout in seconds for HTTP, claude-cli, Anthropic SDK, and Bedrock backends (default: 600) | optional — also `--api-timeout` flag | | `GRAPHIFY_MAX_RETRIES` | How many times to retry a rate-limited (429) request before giving up (default: 6; honors `Retry-After`) | optional — raise for strict per-org limits (e.g. kimi); `0` disables | @@ -737,6 +738,7 @@ graphify extract ./docs --backend bedrock # AWS Bedrock via IAM - no API ke graphify extract ./docs --backend claude-cli # route through Claude Code CLI - no API key, uses your Claude subscription graphify extract ./docs --backend azure # Azure OpenAI (set AZURE_OPENAI_API_KEY + AZURE_OPENAI_ENDPOINT) graphify extract ./docs --max-workers 16 # AST parallelism (also GRAPHIFY_MAX_WORKERS) +graphify extract ./src --memory-limit-mb 6144 # abort (exit 3, no partial graph) instead of being OOM-killed; also GRAPHIFY_MEMORY_LIMIT_MB graphify extract --postgres "postgresql://user:pass@host/db" # introspect live PostgreSQL schema directly graphify extract ./my-workspace --cargo # introspect Rust Cargo workspace dependencies directly graphify extract ./docs --token-budget 30000 # smaller semantic chunks for local/small models diff --git a/graphify/cli.py b/graphify/cli.py index 5b7339726..8291939c5 100644 --- a/graphify/cli.py +++ b/graphify/cli.py @@ -594,6 +594,47 @@ def mark(self, stage: str) -> None: def total(self) -> None: if self.enabled: print(f"[graphify timing] total: {self._now() - self.start:.1f}s", file=sys.stderr) +def _arm_memory_budget(command: str, flag_value: int | None) -> None: + """Put this process under the memory budget (#3011) before extraction. + + ``flag_value`` (``--memory-limit-mb``) wins over ``GRAPHIFY_MEMORY_LIMIT_MB`` + and is written back to the environment so pool workers and nested + rebuilds see the same cap. A malformed env value is refused (exit 2) + rather than ignored - a budget that silently vanished is the failure this + exists to prevent. Where the platform cannot enforce a limit, say so once + and continue; the run is otherwise identical. + """ + from graphify.memory_budget import ( + ENV_VAR, + apply_memory_budget, + configured_limit_mb, + supports_enforcement, + ) + if flag_value is not None: + os.environ[ENV_VAR] = str(flag_value) + try: + limit = configured_limit_mb() + except ValueError as exc: + print(f"error: {exc}", file=sys.stderr) + sys.exit(2) + if limit is None: + return + if not supports_enforcement(): + print( + f"[{command}] warning: memory budget of {limit} MB cannot be enforced on " + f"this platform (no setrlimit); continuing without one", + file=sys.stderr, + ) + return + if apply_memory_budget(limit): + print(f"[{command}] memory budget: {limit} MB (applies to this process and its extraction workers)") + else: + print( + f"[{command}] warning: could not apply the {limit} MB memory budget; continuing without one", + file=sys.stderr, + ) + + def _enforce_graph_size_cap_or_exit(gp: Path) -> None: """Reject oversized graph files before parsing (CLI exit-on-fail flavor). @@ -2253,13 +2294,34 @@ def _clear_html_stale_marker() -> None: no_cluster = False args = sys.argv[2:] watch_arg: str | None = None + update_memory_limit_mb: int | None = None + _pending_flag: str | None = None + + def _parse_mem_flag(raw: str) -> int: + from graphify.memory_budget import parse_limit_mb + try: + return parse_limit_mb(raw) + except ValueError as exc: + print(f"error: --memory-limit-mb {exc}", file=sys.stderr) + sys.exit(2) + for a in args: + if _pending_flag == "--memory-limit-mb": + _pending_flag = None + update_memory_limit_mb = _parse_mem_flag(a) + continue if a == "--force": force = True continue if a == "--no-cluster": no_cluster = True continue + if a == "--memory-limit-mb": + _pending_flag = a + continue + if a.startswith("--memory-limit-mb="): + update_memory_limit_mb = _parse_mem_flag(a.split("=", 1)[1]) + continue if a.startswith("-"): print(f"error: unknown update option: {a}", file=sys.stderr) sys.exit(2) @@ -2277,16 +2339,29 @@ def _clear_html_stale_marker() -> None: watch_path = Path(saved.read_text(encoding="utf-8").strip()) else: watch_path = Path(".") + if _pending_flag is not None: + print(f"error: {_pending_flag} requires a value", file=sys.stderr) + sys.exit(2) if not watch_path.exists(): print(f"error: path not found: {watch_path}", file=sys.stderr) sys.exit(1) + _arm_memory_budget("graphify update", update_memory_limit_mb) from graphify.watch import _rebuild_code print(f"Re-extracting code files in {watch_path} (no LLM needed)...") # Interactive CLI: block on the per-repo lock rather than skip, so the # user sees their explicit `graphify update` complete instead of # exiting silently when a hook-driven rebuild happens to be running. - ok = _rebuild_code(watch_path, force=force, no_cluster=no_cluster, block_on_lock=True) + try: + ok = _rebuild_code(watch_path, force=force, no_cluster=no_cluster, block_on_lock=True) + except MemoryError as exc: + from graphify.memory_budget import ( + EXIT_MEMORY_BUDGET as _exit_mem, + budget_error as _budget_error, + report as _report_mem, + ) + _report_mem(_budget_error(exc, phase="code re-extraction")) + sys.exit(_exit_mem) if ok: print("Code graph updated. For doc/paper/image changes run /graphify --update in your AI assistant.") if not ( @@ -2995,7 +3070,7 @@ def _to_simple(g: "_nx.Graph") -> "_nx.Graph": "Usage: graphify extract [--backend gemini|kimi|claude|openai|deepseek|ollama] " "[--model M] [--mode deep] [--out DIR|--output DIR] [--google-workspace] [--no-cluster] " "[--no-gitignore] [--code-only] [--no-dedup] " - "[--max-workers N] [--token-budget N] [--max-concurrency N] " + "[--max-workers N] [--memory-limit-mb N] [--token-budget N] [--max-concurrency N] " "[--api-timeout S] [--postgres DSN] [--cargo] [--allow-partial] [--timing]", file=sys.stderr, ) @@ -3034,6 +3109,7 @@ def _to_simple(g: "_nx.Graph") -> "_nx.Graph": global_repo_tag: str | None = None # Performance/tuning knobs (issue #792). None means "use library default". cli_max_workers: int | None = None + cli_memory_limit_mb: int | None = None cli_token_budget: int | None = None cli_max_concurrency: int | None = None cli_api_timeout: float | None = None @@ -3111,6 +3187,10 @@ def _parse_float(name: str, raw: str) -> float: cli_max_workers = _parse_int("--max-workers", args[i + 1]); i += 2 elif a.startswith("--max-workers="): cli_max_workers = _parse_int("--max-workers", a.split("=", 1)[1]); i += 1 + elif a == "--memory-limit-mb" and i + 1 < len(args): + cli_memory_limit_mb = _parse_int("--memory-limit-mb", args[i + 1]); i += 2 + elif a.startswith("--memory-limit-mb="): + cli_memory_limit_mb = _parse_int("--memory-limit-mb", a.split("=", 1)[1]); i += 1 elif a == "--token-budget" and i + 1 < len(args): cli_token_budget = _parse_int("--token-budget", args[i + 1]); i += 2 elif a.startswith("--token-budget="): @@ -3183,6 +3263,10 @@ def _parse_float(name: str, raw: str) -> float: os.environ["GRAPHIFY_API_TIMEOUT"] = str(cli_api_timeout) if cli_max_workers is not None: os.environ["GRAPHIFY_MAX_WORKERS"] = str(cli_max_workers) + # Memory budget (#3011): the flag wins over the env var; either way it + # lands in the environment so the extraction workers (which start + # fresh under `spawn`) and any nested rebuild apply the same cap. + _arm_memory_budget("graphify extract", cli_memory_limit_mb) # Resolve output dir. The user-facing contract is "/graphify-out/" # so a fresh checkout writes graphify-out/ at the project root, matching @@ -3644,6 +3728,16 @@ def _ctx_identity(source_file) -> str | None: print(f"[graphify extract] AST extraction on {len(code_files)} code files...") try: ast_result = _ast_extract(code_files, **ast_kwargs) + except MemoryError as exc: + # The memory budget (#3011) was hit. Never a partial graph, never + # --allow-partial: the operator asked to be stopped here. + from graphify.memory_budget import ( + EXIT_MEMORY_BUDGET as _exit_mem, + budget_error as _budget_error, + report as _report_mem, + ) + _report_mem(_budget_error(exc, phase="AST extraction")) + sys.exit(_exit_mem) except Exception as exc: print(f"[graphify extract] AST extraction failed: {exc}", file=sys.stderr) # #2445: losing the whole AST pass is fatal by default. The diff --git a/graphify/extract.py b/graphify/extract.py index 89082af87..9ed442a1f 100644 --- a/graphify/extract.py +++ b/graphify/extract.py @@ -170,6 +170,11 @@ def _raise_recursion_limit() -> None: def _safe_extract(extractor: Callable, path: Path) -> dict: try: return extractor(path) + except MemoryError: + # Under a memory budget (#3011) this is the budget being hit, not a + # bad file. Recording it as a skipped file would let the run finish + # and publish a graph silently missing whatever came after. + raise except RecursionError: print(f" warning: skipped {path} (recursion limit exceeded)", file=sys.stderr, flush=True) return {"nodes": [], "edges": [], "error": "recursion_limit_exceeded"} @@ -5429,6 +5434,14 @@ def _safe_extract_with_xaml_root(extractor, path: Path, root: Path) -> dict: _XAML_ACTIVE_EXTRACT_ROOT = previous_root +def _pool_worker_init() -> None: + """Pool initializer: put each worker under the configured memory budget + (#3011). Under `fork` the parent's rlimit is inherited already; under + `spawn` the worker starts fresh and must apply it from the environment.""" + from graphify.memory_budget import apply_memory_budget + apply_memory_budget() + + def _extract_single_file(args: tuple) -> tuple[int, dict]: """Worker function for parallel extraction. Runs in a subprocess. @@ -5538,7 +5551,9 @@ def _extract_parallel( failed: list[int] = [] # positions into uncached_work whose future failed _PROGRESS_INTERVAL = 100 try: - with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as pool: + with concurrent.futures.ProcessPoolExecutor( + max_workers=max_workers, initializer=_pool_worker_init, + ) as pool: futures = { pool.submit(_extract_single_file, item): pos for pos, item in enumerate(work_items) @@ -5555,6 +5570,17 @@ def _extract_parallel( # swallowed here per-future — that left the remaining # per_file slots empty and silently dropped the files. raise + except MemoryError as exc: + # A worker hit the memory budget (#3011). This is not a + # per-file failure to warn about and retry in-process - + # the retry would hit the same wall in the parent, and a + # "skipped" file would leave the graph silently partial. + # Drop the queued work and abort the run. + from graphify.memory_budget import budget_error + pool.shutdown(wait=False, cancel_futures=True) + raise budget_error( + exc, phase=f"AST extraction of {work_items[futures[future]][1]}" + ) from exc except Exception as exc: pos = futures[future] print( diff --git a/graphify/memory_budget.py b/graphify/memory_budget.py new file mode 100644 index 000000000..c18588da5 --- /dev/null +++ b/graphify/memory_budget.py @@ -0,0 +1,175 @@ +"""A first-class memory budget for extraction (#3011). + +`graphify extract` inside a memory-limited container can grow past the +cgroup allowance — the AST pool is bounded by ``--max-workers``, but the +later JS/TS symbol-resolution passes retain source buffers and syntax trees +for the whole corpus — and the kernel then OOM-kills the pod with no +graphify-specific signal, leaving whatever it had written behind. + +``--memory-limit-mb N`` / ``GRAPHIFY_MEMORY_LIMIT_MB=N`` turns that into a +graphify failure: + +* the limit is applied with ``setrlimit`` to the CLI process and, through + the pool initializer, to every extraction worker; +* an allocation past it raises ``MemoryError`` where it happens, which the + pipeline lets propagate instead of recording as a skipped file; +* the CLI reports the configured limit, the observed peak, and the phase, + exits with :data:`EXIT_MEMORY_BUDGET`, and writes no ``graph.json`` — + the previous graph, if any, is untouched. + +``RLIMIT_AS`` bounds virtual address space, which over-approximates RSS +(the figure a cgroup accounts), so set the budget somewhat below the +container's limit. macOS uses ``RLIMIT_DATA`` because ``RLIMIT_AS`` is not +honoured by its allocator. Windows has neither; the CLI says so and runs +without a budget rather than pretending. + +The hooks' ``GRAPHIFY_REBUILD_MEMORY_LIMIT_MB`` predates this and keeps +working; it takes precedence on the hook/watch rebuild path only. +""" +from __future__ import annotations + +import os +import sys + +ENV_VAR = "GRAPHIFY_MEMORY_LIMIT_MB" +REBUILD_ENV_VAR = "GRAPHIFY_REBUILD_MEMORY_LIMIT_MB" +FLAG = "--memory-limit-mb" + +#: Exit status when the budget is exceeded. Distinct from 1 (extraction +#: failed) and 2 (bad arguments) so a wrapper can tell "give it more memory" +#: apart from "the corpus is broken" without parsing stderr. +EXIT_MEMORY_BUDGET = 3 + + +class MemoryBudgetExceeded(MemoryError): + """The configured memory budget was hit. + + A ``MemoryError`` subclass so code that already treats ``MemoryError`` + as fatal keeps doing so; carries what the CLI reports. + """ + + def __init__(self, limit_mb: int | None, *, phase: str = "extraction", + observed_mb: float | None = None, detail: str | None = None) -> None: + self.limit_mb = limit_mb + self.phase = phase + self.observed_mb = observed_mb + self.detail = detail + super().__init__(self.describe()) + + def describe(self) -> str: + limit = f"{self.limit_mb} MB" if self.limit_mb is not None else "the process limit" + msg = f"memory budget of {limit} exceeded during {self.phase}" + if self.observed_mb is not None: + msg += f" (peak observed in this process: ~{self.observed_mb:.0f} MB)" + if self.detail: + msg += f": {self.detail}" + return msg + + +def parse_limit_mb(raw: str) -> int: + """Validate a user-supplied budget. Raises ``ValueError`` with a reason.""" + try: + value = int(str(raw).strip()) + except ValueError: + raise ValueError(f"must be a positive integer number of megabytes (got {raw!r})") from None + if value <= 0: + raise ValueError(f"must be > 0 (got {value})") + return value + + +def configured_limit_mb(env: "os._Environ[str] | dict[str, str] | None" = None) -> int | None: + """The budget from :data:`ENV_VAR`, or ``None`` when unset. + + Raises ``ValueError`` on a malformed value so the caller can refuse the + run: a budget that silently became "no budget" is exactly the failure + this feature exists to prevent. + """ + env = os.environ if env is None else env + raw = (env.get(ENV_VAR) or "").strip() + if not raw: + return None + try: + return parse_limit_mb(raw) + except ValueError as exc: + raise ValueError(f"{ENV_VAR} {exc}") from None + + +def supports_enforcement() -> bool: + """True where ``setrlimit`` exists and is honoured for memory.""" + if sys.platform == "win32": + return False + try: + import resource # noqa: F401 + except ImportError: + return False + return True + + +def apply_memory_budget(limit_mb: int | None = None) -> bool: + """Cap this process's memory at ``limit_mb`` (default: the configured + budget). Returns True when a limit is now in force, False when there is + nothing to apply or the platform cannot enforce one. + + Never raises: a worker's initializer calls this, and a failure there + would take the whole pool down with a far less useful message. + """ + if limit_mb is None: + try: + limit_mb = configured_limit_mb() + except ValueError: + return False + if limit_mb is None or not supports_enforcement(): + return False + try: + import resource + which = resource.RLIMIT_DATA if sys.platform == "darwin" else resource.RLIMIT_AS + limit = int(limit_mb) * 1024 * 1024 + soft, hard = resource.getrlimit(which) + # Never raise a hard limit an operator (or a container runtime) + # already set lower than ours. + new_hard = hard if hard != resource.RLIM_INFINITY and hard < limit else limit + resource.setrlimit(which, (min(limit, new_hard), new_hard)) + return True + except (ImportError, ValueError, OSError): + return False + + +def peak_rss_mb() -> float | None: + """Peak resident size of this process in MB, when the platform reports it.""" + try: + import resource + ru = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + except (ImportError, AttributeError, OSError): + return None + # Linux reports kilobytes, macOS bytes. + return ru / 1024.0 if sys.platform != "darwin" else ru / (1024.0 * 1024.0) + + +def budget_error(exc: BaseException, *, phase: str) -> MemoryBudgetExceeded: + """Normalise any ``MemoryError`` raised under a budget into the typed form.""" + if isinstance(exc, MemoryBudgetExceeded): + return exc + try: + limit = configured_limit_mb() + except ValueError: + limit = None + detail = str(exc).strip() or None + return MemoryBudgetExceeded(limit, phase=phase, observed_mb=peak_rss_mb(), detail=detail) + + +def report(exc: MemoryBudgetExceeded, *, stream=None) -> None: + """Print the operator-facing account of a budget failure.""" + stream = sys.stderr if stream is None else stream + print(f"error: {exc.describe()}", file=stream) + if exc.limit_mb is not None: + print( + f" configured: {exc.limit_mb} MB ({FLAG} / {ENV_VAR}); " + f"the previous graph.json, if any, was left untouched.", + file=stream, + ) + print( + " Raise the budget, narrow the corpus (.graphifyignore, --exclude), " + "or lower --max-workers to reduce peak usage.", + file=stream, + ) + print(f" exit status {EXIT_MEMORY_BUDGET}", file=stream) diff --git a/graphify/watch.py b/graphify/watch.py index 8ad02c4df..9fc7ed875 100644 --- a/graphify/watch.py +++ b/graphify/watch.py @@ -226,21 +226,19 @@ def _apply_resource_limits() -> None: os.nice(10) except (OSError, AttributeError): pass - mb = os.environ.get("GRAPHIFY_REBUILD_MEMORY_LIMIT_MB", "").strip() + # The hook-specific variable keeps precedence; the general budget + # (#3011, `graphify extract --memory-limit-mb` / GRAPHIFY_MEMORY_LIMIT_MB) + # applies here too so one setting covers every rebuild path. + from graphify.memory_budget import ENV_VAR, REBUILD_ENV_VAR, apply_memory_budget + mb = os.environ.get(REBUILD_ENV_VAR, "").strip() or os.environ.get(ENV_VAR, "").strip() if not mb: return try: - limit = int(mb) * 1024 * 1024 + limit_mb = int(mb) except ValueError: return - try: - import resource - which = resource.RLIMIT_DATA if sys.platform == "darwin" else resource.RLIMIT_AS - soft, hard = resource.getrlimit(which) - new_hard = hard if hard != resource.RLIM_INFINITY and hard < limit else limit - resource.setrlimit(which, (limit, new_hard)) - except (ImportError, ValueError, OSError): - pass + if limit_mb > 0: + apply_memory_budget(limit_mb) def _git_head(cwd: Path | str | None = None) -> str | None: diff --git a/tests/test_memory_budget.py b/tests/test_memory_budget.py new file mode 100644 index 000000000..824db9b0d --- /dev/null +++ b/tests/test_memory_budget.py @@ -0,0 +1,346 @@ +"""`graphify extract --memory-limit-mb N`: a memory budget for extraction (#3011). + +Inside a memory-limited container graphify could grow past the cgroup +allowance (the AST pool is bounded by --max-workers, but the JS/TS +resolution passes retain source buffers and trees for the whole corpus) and +the kernel OOM-killed the pod: no graphify-specific failure, no stable exit +status, and whatever had been written was left behind. + +The budget must (1) reach the extraction workers, (2) turn an allocation +past it into an abort rather than a "skipped file" warning, and (3) exit +with a distinct status having written no graph.json. +""" +from __future__ import annotations + +import concurrent.futures +import io +import os +import subprocess +import sys +from contextlib import redirect_stdout +from pathlib import Path + +import pytest + +import graphify.__main__ as mainmod +import graphify.extract as extractmod +from graphify import memory_budget as mb +from graphify.extract import _safe_extract, extract +from graphify.memory_budget import ( + ENV_VAR, + EXIT_MEMORY_BUDGET, + REBUILD_ENV_VAR, + MemoryBudgetExceeded, + apply_memory_budget, + budget_error, + configured_limit_mb, + parse_limit_mb, + supports_enforcement, +) + +try: + from graphify.extract import _pool_worker_init +except ImportError: # pre-fix tree + _pool_worker_init = None + + +@pytest.fixture(autouse=True) +def _clean_env(monkeypatch): + monkeypatch.delenv(ENV_VAR, raising=False) + monkeypatch.delenv(REBUILD_ENV_VAR, raising=False) + monkeypatch.delenv("GRAPHIFY_MAX_WORKERS", raising=False) + + +# --------------------------------------------------------------------------- +# Parsing +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize("raw", ["0", "-1", "abc", "", "1.5"]) +def test_a_budget_must_be_a_positive_whole_number_of_megabytes(raw): + with pytest.raises(ValueError): + parse_limit_mb(raw) + + +def test_the_env_var_is_read_and_a_bad_value_is_refused_not_ignored(monkeypatch): + assert configured_limit_mb() is None + monkeypatch.setenv(ENV_VAR, " 6144 ") + assert configured_limit_mb() == 6144 + monkeypatch.setenv(ENV_VAR, "lots") + with pytest.raises(ValueError, match=ENV_VAR): + configured_limit_mb() + + +def test_the_typed_error_reports_limit_phase_and_peak(): + exc = MemoryBudgetExceeded(6144, phase="AST extraction", observed_mb=6200.4) + text = str(exc) + assert "6144 MB" in text and "AST extraction" in text and "6200 MB" in text + assert isinstance(exc, MemoryError) # existing MemoryError handling still applies + + +def test_a_plain_memory_error_is_normalised_with_the_configured_limit(monkeypatch): + monkeypatch.setenv(ENV_VAR, "512") + exc = budget_error(MemoryError("boom"), phase="resolution") + assert isinstance(exc, MemoryBudgetExceeded) + assert exc.limit_mb == 512 and exc.phase == "resolution" and "boom" in str(exc) + assert budget_error(exc, phase="x") is exc # already typed: passed through + + +# --------------------------------------------------------------------------- +# Enforcement +# --------------------------------------------------------------------------- + +def test_nothing_to_apply_without_a_budget(): + assert apply_memory_budget() is False + + +@pytest.mark.skipif(not supports_enforcement(), reason="no setrlimit on this platform") +def test_the_limit_is_really_applied_and_bites(tmp_path): + """Run in a subprocess so the test process itself is never capped.""" + code = ( + "import resource, sys\n" + "from graphify.memory_budget import apply_memory_budget\n" + "assert apply_memory_budget(256) is True\n" + "which = resource.RLIMIT_DATA if sys.platform == 'darwin' else resource.RLIMIT_AS\n" + "assert resource.getrlimit(which)[0] == 256 * 1024 * 1024\n" + "try:\n" + " blob = bytearray(2 * 1024 * 1024 * 1024)\n" + "except MemoryError:\n" + " print('MEMORY_ERROR_RAISED')\n" + "else:\n" + " print('ALLOCATED_PAST_LIMIT')\n" + ) + out = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True) + assert out.returncode == 0, out.stderr + assert "MEMORY_ERROR_RAISED" in out.stdout + + +@pytest.mark.skipif(not supports_enforcement(), reason="no setrlimit on this platform") +def test_a_lower_existing_hard_limit_is_never_raised(): + code = ( + "import resource, sys\n" + "from graphify.memory_budget import apply_memory_budget\n" + "which = resource.RLIMIT_DATA if sys.platform == 'darwin' else resource.RLIMIT_AS\n" + "resource.setrlimit(which, (512 * 2**20, 512 * 2**20))\n" + "apply_memory_budget(4096)\n" + "print(resource.getrlimit(which))\n" + ) + out = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True) + assert out.returncode == 0, out.stderr + soft, hard = eval(out.stdout.strip()) + assert hard == 512 * 2**20 and soft <= hard + + +def test_unsupported_platform_reports_false_instead_of_pretending(monkeypatch): + monkeypatch.setattr(mb, "supports_enforcement", lambda: False) + assert apply_memory_budget(1024) is False + + +@pytest.mark.skipif(_pool_worker_init is None, reason="pre-fix tree") +def test_the_pool_initializer_applies_the_budget_from_the_environment(monkeypatch): + """Under `spawn` a worker starts fresh: it must pick the cap up itself.""" + seen = [] + monkeypatch.setattr(mb, "apply_memory_budget", lambda limit_mb=None: seen.append(limit_mb) or True) + monkeypatch.setenv(ENV_VAR, "777") + _pool_worker_init() + assert seen == [None] # reads the env inside, not a stale argument + + +@pytest.mark.skipif(_pool_worker_init is None, reason="pre-fix tree") +def test_the_pool_is_constructed_with_the_initializer(tmp_path, monkeypatch): + seen = {} + + class RecordingPool(concurrent.futures.ThreadPoolExecutor): + def __init__(self, max_workers=None, initializer=None, initargs=(), **kw): + seen["initializer"] = initializer + super().__init__(max_workers=max_workers, initializer=initializer, initargs=initargs) + + monkeypatch.setattr(concurrent.futures, "ProcessPoolExecutor", RecordingPool) + _corpus(tmp_path, 25) + with redirect_stdout(io.StringIO()): + extract(sorted(tmp_path.glob("*.py")), cache_root=tmp_path, root=tmp_path, parallel=True) + assert seen["initializer"] is _pool_worker_init + + +def test_hook_rebuilds_honour_the_general_budget_too(monkeypatch): + """`watch._apply_resource_limits` predates this; the hook-specific variable + keeps precedence, the general one now applies when it is the only one.""" + from graphify import watch + seen = [] + monkeypatch.setattr(mb, "apply_memory_budget", lambda limit_mb=None: seen.append(limit_mb) or True) + monkeypatch.setattr(os, "nice", lambda n: None, raising=False) + watch._apply_resource_limits() + monkeypatch.setenv(ENV_VAR, "256") + watch._apply_resource_limits() + monkeypatch.setenv(REBUILD_ENV_VAR, "512") + watch._apply_resource_limits() + assert seen == [256, 512] + + +# --------------------------------------------------------------------------- +# An allocation past the budget aborts the run - it is not a skipped file +# --------------------------------------------------------------------------- + +def _corpus(root: Path, n: int) -> list[Path]: + files = [] + for i in range(n): + p = root / f"m{i}.py" + p.write_text(f"def f{i}():\n return {i}\n", encoding="utf-8") + files.append(p) + return files + + +def test_safe_extract_still_swallows_ordinary_failures(tmp_path): + def bad(path): + raise RuntimeError("parser exploded") + with redirect_stdout(io.StringIO()): + result = _safe_extract(bad, tmp_path / "x.py") + assert result["nodes"] == [] and "error" in result + + +def test_safe_extract_lets_a_memory_error_through(tmp_path): + def oom(path): + raise MemoryError() + with pytest.raises(MemoryError): + _safe_extract(oom, tmp_path / "x.py") + + +def test_sequential_extraction_aborts_instead_of_finishing_partial(tmp_path, monkeypatch): + files = _corpus(tmp_path, 3) + calls = [] + + def oom_on_second(path): + calls.append(path.name) + if len(calls) == 2: + raise MemoryError() + return {"nodes": [{"id": path.stem, "label": path.stem, "file_type": "code", + "source_file": str(path)}], "edges": []} + + monkeypatch.setattr(extractmod, "_get_extractor", lambda p: oom_on_second) + with pytest.raises(MemoryError), redirect_stdout(io.StringIO()): + extract(files, cache_root=tmp_path, root=tmp_path, parallel=False) + assert len(calls) == 2 # stopped there; the third file was never attempted + + +def test_a_worker_hitting_the_budget_aborts_the_pool_with_the_typed_error(tmp_path, monkeypatch): + """The per-future handler used to treat any exception as a per-file failure: + warn, then retry that file in-process - which would hit the same wall.""" + monkeypatch.setattr(concurrent.futures, "ProcessPoolExecutor", concurrent.futures.ThreadPoolExecutor) + monkeypatch.setenv(ENV_VAR, "2048") + files = _corpus(tmp_path, 25) + real = extractmod._extract_single_file + victim = files[7].name + + def worker(args): + if Path(args[1]).name == victim: + raise MemoryError() + return real(args) + + monkeypatch.setattr(extractmod, "_extract_single_file", worker) + out = io.StringIO() + with pytest.raises(MemoryBudgetExceeded) as info, redirect_stdout(out): + extract(files, cache_root=tmp_path, root=tmp_path, parallel=True) + assert info.value.limit_mb == 2048 + assert victim in str(info.value) + assert "worker failed" not in out.getvalue() # not demoted to a warning + + +# --------------------------------------------------------------------------- +# The CLI: distinct exit status, no partial graph, honest on Windows +# --------------------------------------------------------------------------- + +def _run_cli(monkeypatch, argv): + monkeypatch.setattr(mainmod, "_check_skill_version", lambda _: None) + monkeypatch.setattr(mainmod.sys, "argv", ["graphify", *argv]) + with pytest.raises(SystemExit) as info: + mainmod.main() + return info.value.code + + +@pytest.fixture +def corpus(tmp_path): + c = tmp_path / "corpus" + c.mkdir() + (c / "main.go").write_text("package main\nfunc main() {}\n", encoding="utf-8") + return c + + +def test_extract_exits_3_and_writes_no_graph_when_the_budget_is_hit(corpus, tmp_path, monkeypatch, capsys): + out_dir = tmp_path / "out" + + def oom(paths, **kw): + raise MemoryError() + + monkeypatch.setattr(extractmod, "extract", oom) + monkeypatch.setattr(mb, "supports_enforcement", lambda: True) + monkeypatch.setattr(mb, "apply_memory_budget", lambda limit_mb=None: True) + code = _run_cli(monkeypatch, ["extract", str(corpus), "--code-only", "--out", str(out_dir), + "--memory-limit-mb", "6144", "--allow-partial"]) + err = capsys.readouterr().err + assert code == EXIT_MEMORY_BUDGET == 3 + assert "memory budget of 6144 MB exceeded during AST extraction" in err + assert "--memory-limit-mb / GRAPHIFY_MEMORY_LIMIT_MB" in err + assert not (out_dir / "graphify-out" / "graph.json").exists() + assert os.environ.get(ENV_VAR) == "6144" # forwarded to workers via the environment + + +def test_the_flag_form_with_equals_and_the_env_var_both_work(corpus, tmp_path, monkeypatch, capsys): + monkeypatch.setattr(extractmod, "extract", lambda paths, **kw: (_ for _ in ()).throw(MemoryError())) + monkeypatch.setattr(mb, "supports_enforcement", lambda: True) + applied = [] + monkeypatch.setattr(mb, "apply_memory_budget", lambda limit_mb=None: applied.append(limit_mb) or True) + assert _run_cli(monkeypatch, ["extract", str(corpus), "--code-only", "--out", str(tmp_path / "a"), + "--memory-limit-mb=300"]) == 3 + monkeypatch.setenv(ENV_VAR, "400") + assert _run_cli(monkeypatch, ["extract", str(corpus), "--code-only", "--out", str(tmp_path / "b")]) == 3 + assert applied == [300, 400] + + +@pytest.mark.parametrize("value", ["0", "-5", "big"]) +def test_a_bad_flag_value_is_a_usage_error(corpus, monkeypatch, capsys, value): + assert _run_cli(monkeypatch, ["extract", str(corpus), "--memory-limit-mb", value]) == 2 + assert "--memory-limit-mb" in capsys.readouterr().err + + +def test_a_bad_env_value_is_refused_rather_than_silently_dropped(corpus, monkeypatch, capsys): + monkeypatch.setenv(ENV_VAR, "unlimited") + assert _run_cli(monkeypatch, ["extract", str(corpus), "--code-only"]) == 2 + assert ENV_VAR in capsys.readouterr().err + + +def test_an_unenforceable_platform_says_so_and_continues(corpus, tmp_path, monkeypatch, capsys): + monkeypatch.setattr(mb, "supports_enforcement", lambda: False) + monkeypatch.setattr(extractmod, "extract", lambda paths, **kw: (_ for _ in ()).throw(RuntimeError("stop here"))) + code = _run_cli(monkeypatch, ["extract", str(corpus), "--code-only", "--out", str(tmp_path / "o"), + "--memory-limit-mb", "1024"]) + err = capsys.readouterr().err + assert code == 1 # the ordinary failure path ran: the budget did not block the run + assert "cannot be enforced on this platform" in err + + +def test_the_ordinary_failure_path_is_unchanged(corpus, tmp_path, monkeypatch, capsys): + """A RuntimeError from the AST pass still exits 1 with the #2445 message.""" + monkeypatch.setattr(extractmod, "extract", lambda paths, **kw: (_ for _ in ()).throw(RuntimeError("worker pool failed"))) + code = _run_cli(monkeypatch, ["extract", str(corpus), "--code-only", "--out", str(tmp_path / "o")]) + assert code == 1 + assert "AST extraction failed: worker pool failed" in capsys.readouterr().err + + +def test_update_takes_the_flag_and_exits_3_on_the_budget(corpus, monkeypatch, capsys): + from graphify import watch + monkeypatch.chdir(corpus) + monkeypatch.setattr(mb, "supports_enforcement", lambda: True) + applied = [] + monkeypatch.setattr(mb, "apply_memory_budget", lambda limit_mb=None: applied.append(limit_mb) or True) + monkeypatch.setattr(watch, "_rebuild_code", lambda *a, **k: (_ for _ in ()).throw(MemoryError())) + code = _run_cli(monkeypatch, ["update", ".", "--memory-limit-mb", "2048"]) + err = capsys.readouterr().err + assert code == 3 + assert applied == [2048] + assert "exceeded during code re-extraction" in err + + +def test_update_rejects_a_dangling_or_bad_flag(corpus, monkeypatch, capsys): + monkeypatch.chdir(corpus) + assert _run_cli(monkeypatch, ["update", "--memory-limit-mb"]) == 2 + assert _run_cli(monkeypatch, ["update", "--memory-limit-mb=zero"]) == 2 + assert _run_cli(monkeypatch, ["update", "--bogus"]) == 2 # unknown options still refused