|
| 1 | +"""Layer A vs Layer B memory-enforcement validation on the fixture repo (PRD section 4.4). |
| 2 | +
|
| 3 | +Runs three things under a 256 MB budget and prints each outcome in full: |
| 4 | + Control -- a dependency-free ``bytearray(400 MB)`` under a real RLIMIT_AS cap. This |
| 5 | + isolates RLIMIT_AS itself: if the control is not enforced, the bug is in |
| 6 | + linux_runner, independent of any model/library. |
| 7 | + Layer B -- the same fixture benchmark (bge-small + stub LLM) under the real RLIMIT_AS cap. |
| 8 | + Layer A -- the same benchmark under the portable app-level MemoryGovernor (peak-RSS cap). |
| 9 | +
|
| 10 | +For every RLIMIT_AS run it prints the kernel-confirmed limit (getrlimit *after* setrlimit), |
| 11 | +the capped child's PID, the exit code, and the failing exception's repr -- so a cap-induced |
| 12 | +failure is visible, not hidden behind a coarse "enforced" flag. No Ollama and no ianvs |
| 13 | +re-ingest. On macOS the RLIMIT_AS runs are reported as skipped (it cannot enforce them). |
| 14 | +""" |
| 15 | + |
| 16 | +from __future__ import annotations |
| 17 | + |
| 18 | +import os |
| 19 | +import tempfile |
| 20 | +from collections.abc import Callable |
| 21 | +from pathlib import Path |
| 22 | + |
| 23 | +from codecompass.benchmark.dataset import load_gold_questions |
| 24 | +from codecompass.benchmark.runner import BenchmarkRunner |
| 25 | +from codecompass.config import ChunkingConfig, EmbeddingConfig, RepoConfig |
| 26 | +from codecompass.constraints.linux_runner import ( |
| 27 | + EnforcementResult, |
| 28 | + UnsupportedPlatformError, |
| 29 | + run_under_memory_limit, |
| 30 | +) |
| 31 | +from codecompass.constraints.memory import peak_rss_bytes |
| 32 | +from codecompass.constraints.profiles import ConstraintProfile |
| 33 | +from codecompass.embedding.registry import build_embedder |
| 34 | +from codecompass.ingestion.github_source import build_file_documents |
| 35 | +from codecompass.pipeline import QueryPipeline, index_documents |
| 36 | +from codecompass.retrieval.retriever import Retriever |
| 37 | +from codecompass.store.milvus_lite import MilvusLiteStore |
| 38 | + |
| 39 | +_BUDGET_MB = int(os.environ.get("CC_EDGE_BUDGET_MB", "256")) |
| 40 | +_FIXTURE = Path("tests/fixtures/sample_repo") |
| 41 | +_GOLD_SET = Path("data/benchmarks/fixture_qa.yaml") |
| 42 | +_EMBEDDING = EmbeddingConfig(model_name="BAAI/bge-small-en-v1.5", dimension=384) |
| 43 | + |
| 44 | + |
| 45 | +class _StubProvider: |
| 46 | + """Offline LLM provider so the validation needs no Ollama.""" |
| 47 | + |
| 48 | + def generate(self, prompt: str) -> str: |
| 49 | + """Return a fixed answer.""" |
| 50 | + return "Validation stub answer." |
| 51 | + |
| 52 | + |
| 53 | +def _allocate_400mb() -> None: |
| 54 | + """Dependency-free 400 MB allocation -- isolates RLIMIT_AS from any model/library.""" |
| 55 | + block = bytearray(400 * 1024 * 1024) |
| 56 | + block[-1] = 1 |
| 57 | + |
| 58 | + |
| 59 | +def _build_runner(db_dir: Path) -> tuple[BenchmarkRunner, MilvusLiteStore]: |
| 60 | + """Ingest the fixture with bge-small and wire a stub-provider benchmark runner.""" |
| 61 | + embedder = build_embedder(_EMBEDDING) |
| 62 | + store = MilvusLiteStore(db_dir / "idx.db", dimension=embedder.dimension) |
| 63 | + documents = build_file_documents(_FIXTURE, RepoConfig()) |
| 64 | + index_documents(documents, embedder=embedder, store=store, chunking_config=ChunkingConfig()) |
| 65 | + pipeline = QueryPipeline(Retriever(embedder, store, top_k=5), _StubProvider()) |
| 66 | + return BenchmarkRunner(pipeline, top_k=5), store |
| 67 | + |
| 68 | + |
| 69 | +def _benchmark_workload() -> None: |
| 70 | + """Run the fixture benchmark once (bge load + index + score) -- the memory-heavy work.""" |
| 71 | + with tempfile.TemporaryDirectory() as tmp: |
| 72 | + runner, store = _build_runner(Path(tmp)) |
| 73 | + try: |
| 74 | + runner.score_once(load_gold_questions(_GOLD_SET)) |
| 75 | + finally: |
| 76 | + store.close() |
| 77 | + |
| 78 | + |
| 79 | +def _enforce(func: Callable[[], None]) -> EnforcementResult | None: |
| 80 | + """Run ``func`` under the real RLIMIT_AS cap; None when not on Linux.""" |
| 81 | + try: |
| 82 | + return run_under_memory_limit(func, memory_mb=_BUDGET_MB) |
| 83 | + except UnsupportedPlatformError: |
| 84 | + return None |
| 85 | + |
| 86 | + |
| 87 | +def _run_layer_a() -> tuple[int, int, int]: |
| 88 | + """Layer A: run the workload under the MemoryGovernor mem-256 cap. |
| 89 | +
|
| 90 | + Returns ``(failure_count, repeats, observed_peak_rss_mb)``. |
| 91 | + """ |
| 92 | + profile = ConstraintProfile(name=f"mem-{_BUDGET_MB}", memory_mb=_BUDGET_MB) |
| 93 | + with tempfile.TemporaryDirectory() as tmp: |
| 94 | + runner, store = _build_runner(Path(tmp)) |
| 95 | + try: |
| 96 | + result = runner.run(load_gold_questions(_GOLD_SET), profile, repeats=1) |
| 97 | + finally: |
| 98 | + store.close() |
| 99 | + return result.failure_count, result.repeats, peak_rss_bytes() // (1024 * 1024) |
| 100 | + |
| 101 | + |
| 102 | +def _print_enforcement(label: str, result: EnforcementResult | None) -> None: |
| 103 | + """Print a full RLIMIT_AS result: outcome, kernel-confirmed cap, PID, exit, detail.""" |
| 104 | + if result is None: |
| 105 | + print(f"{label} @ {_BUDGET_MB} MB: SKIPPED (not Linux; RLIMIT_AS unenforceable here).") |
| 106 | + return |
| 107 | + soft = result.applied_soft_bytes |
| 108 | + applied_mb = soft // (1024 * 1024) if soft is not None else None |
| 109 | + print( |
| 110 | + f"{label} @ {_BUDGET_MB} MB: outcome={result.outcome} " |
| 111 | + f"(completed={result.completed}, enforced={result.enforced}); " |
| 112 | + f"kernel RLIMIT_AS soft={applied_mb} MB; child_pid={result.child_pid}; " |
| 113 | + f"exit={result.exit_code}; detail={result.detail}" |
| 114 | + ) |
| 115 | + if result.traceback_text: |
| 116 | + print(f" ----- {label}: full exception chain (child PID {result.child_pid}) -----") |
| 117 | + for line in result.traceback_text.rstrip().splitlines(): |
| 118 | + print(f" {line}") |
| 119 | + print(" ----- end exception chain -----") |
| 120 | + |
| 121 | + |
| 122 | +def main() -> None: |
| 123 | + """Run Control + Layer B first (parent stays torch-free at fork), then Layer A; report.""" |
| 124 | + print(f"== Layer A vs Layer B memory enforcement @ {_BUDGET_MB} MB (fixture repo, stub LLM) ==") |
| 125 | + control = _enforce(_allocate_400mb) |
| 126 | + layer_b = _enforce(_benchmark_workload) |
| 127 | + failures, repeats, peak_mb = _run_layer_a() |
| 128 | + layer_a_completed = failures == 0 |
| 129 | + |
| 130 | + print("") |
| 131 | + _print_enforcement("Control (bytearray 400 MB, no deps)", control) |
| 132 | + _print_enforcement("Layer B (model workload)", layer_b) |
| 133 | + print( |
| 134 | + f"Layer A (MemoryGovernor peak-RSS cap): failures={failures}/{repeats}, " |
| 135 | + f"observed peak RSS={peak_mb} MB vs budget {_BUDGET_MB} MB -> " |
| 136 | + f"{'OK (ran)' if layer_a_completed else 'FAIL (rejected)'}" |
| 137 | + ) |
| 138 | + |
| 139 | + if control is not None and not control.completed and not control.enforced: |
| 140 | + # Only a control that neither completed (fit the budget) nor was enforced (a clean |
| 141 | + # MemoryError) is suspicious -- that means the pure-allocation control itself errored. |
| 142 | + print( |
| 143 | + f"WARNING: the dependency-free control errored (outcome={control.outcome}) -- " |
| 144 | + "RLIMIT_AS may be misbehaving in linux_runner; investigate the Layer-B result." |
| 145 | + ) |
| 146 | + if layer_b is None: |
| 147 | + print("Agreement: PENDING Layer B (not on Linux) -- run the edge-validation workflow.") |
| 148 | + return |
| 149 | + if layer_a_completed != layer_b.completed: |
| 150 | + print( |
| 151 | + f"Agreement: DISAGREE -- Layer A completed={layer_a_completed}, " |
| 152 | + f"Layer B completed={layer_b.completed}. Inspect the Layer-B detail above." |
| 153 | + ) |
| 154 | + elif layer_a_completed: |
| 155 | + print(f"Agreement: AGREE -- both layers ran the workload under {_BUDGET_MB} MB.") |
| 156 | + elif layer_b.enforced: |
| 157 | + print( |
| 158 | + f"Agreement: AGREE -- both rejected the workload under {_BUDGET_MB} MB " |
| 159 | + "(Layer B via a clean MemoryError/kill)." |
| 160 | + ) |
| 161 | + else: |
| 162 | + print( |
| 163 | + f"Agreement: LIKELY AGREE -- both failed to run under {_BUDGET_MB} MB, but Layer B's " |
| 164 | + "outcome=error is a non-MemoryError; confirm from the detail above that it is " |
| 165 | + "cap-induced (e.g. a failed mmap/import), not an unrelated failure." |
| 166 | + ) |
| 167 | + |
| 168 | + |
| 169 | +if __name__ == "__main__": |
| 170 | + main() |
0 commit comments