fix(data): harden conveyor dedup locks and memory limits

Author: apstenku123

scripts/streaming_conveyor.py

@@ -121,6 +121,7 @@ def run_code_half(
     dedup_db: Path | None,
     dedup_near: bool,
     global_symbol_index: Path | None = None,
+    memory_limit_gb: float = 10.0,
 ) -> dict:
     """index+route+pack the repo's source via the EXISTING code stage, zstd-max.
 
@@ -132,7 +133,7 @@ def run_code_half(
     """
     info = sr.process_one_repo(
         repo, repo_dir, lengths_code, work_root, dedup_db, dedup_near,
-        global_symbol_index,
+        global_symbol_index, memory_limit_gb,
     )
     # zstd-max the per-length code parquet files this repo just wrote.
     for L in info.get("lengths", {}):
@@ -163,6 +164,7 @@ def run_commits_half(
     dedup_near: bool,
     pr_store: Path | None,
     repo_list: Path | None,
+    memory_limit_gb: float = 10.0,
 ) -> tuple[int, int]:
     """Extract commit records once, fan ranges to the pool. Returns (done, failed).
 
@@ -204,6 +206,7 @@ def run_commits_half(
             process_range, repo, repo_dir, records_jsonl,
             start, end, lengths_sorted, repo_work,
             dedup_db, dedup_near, pr_store, repo_list,
+            memory_limit_gb,
         )
         futures[fut] = (start, end)
 
@@ -259,6 +262,7 @@ def process_one_repo(
     pr_store: Path | None,
     repo_list: Path | None,
     global_symbol_index: Path | None = None,
+    memory_limit_gb: float = 10.0,
 ) -> dict:
     """Run BOTH halves for one already-extracted repo subtree, then delete it.
 
@@ -281,7 +285,7 @@ def process_one_repo(
             try:
                 cinfo = run_code_half(
                     repo, repo_dir, lengths_code, work_root, dedup_db, dedup_near,
-                    global_symbol_index,
+                    global_symbol_index, memory_limit_gb,
                 )
                 with manifest_lock:
                     manifest.mark_done(ck, cinfo)
@@ -303,7 +307,7 @@ def process_one_repo(
             done, failed = run_commits_half(
                 repo, repo_dir, repo_work, lengths_commits, range_size,
                 pool, manifest, manifest_lock, resume, cumulative,
-                dedup_db, dedup_near, pr_store, repo_list,
+                dedup_db, dedup_near, pr_store, repo_list, memory_limit_gb,
             )
             result["commits_done"] = done
             result["commits_failed"] = failed
@@ -370,6 +374,9 @@ def parse_args(argv: list[str]) -> argparse.Namespace:
                         "When set, the CODE half threads it into index_project so "
                         "unresolved base-lib callees are pulled as bounded depth-1 "
                         "deps tagged crosslib:<repo>. DEFAULT off (unchanged).")
+    p.add_argument("--memory-limit-gb", type=float, default=10.0,
+                   help="Per-stage fail-loud RSS limit passed to index/materialize/"
+                        "commit processors (default 10.0).")
     return p.parse_args(argv)
 
 
@@ -469,7 +476,7 @@ def should_process(repo: str) -> bool:
                 repo, repo_dir, lengths_code, lengths_commits, args.range_size,
                 work_root, pool, manifest, manifest_lock, resume, cumulative,
                 args.keep_temp, dedup_db, dedup_near, pr_store, repo_list,
-                global_symbol_index,
+                global_symbol_index, args.memory_limit_gb,
             )
             processed_repos += 1
             if isinstance(res.get("code"), dict):

scripts/streaming_reindex.py

@@ -293,6 +293,7 @@ def stage_index_source(
     dedup_db: Path | None = None,
     dedup_near: bool = True,
     global_symbol_index: Path | None = None,
+    memory_limit_gb: float = 10.0,
 ) -> Path:
     """index_project.py --enriched -> <repo>.enriched.jsonl.
 
@@ -313,6 +314,7 @@ def stage_index_source(
         "--max-tokens", str(TOKENIZE_BUDGET),
         "--exclude-dirs", SOURCE_EXTRA_EXCLUDE,
         "--tokenizer-path", TOKENIZER_PATH,
+        "--memory-limit-gb", str(memory_limit_gb),
     ]
     if dedup_db is not None:
         cmd += ["--dedup-db", str(dedup_db)]
@@ -336,7 +338,8 @@ def stage_index_commits(repo: str, commit_inputs: Sequence[Path], work: Path,
                         dedup_db: Path | None = None,
                         dedup_near: bool = True,
                         pr_store: Path | None = None,
-                        repo_list: Path | None = None) -> Path:
+                        repo_list: Path | None = None,
+                        memory_limit_gb: float = 10.0) -> Path:
     """process_commits.py -> <repo>.enriched.jsonl (commit edit-signal docs).
 
     A commit is an ATOMIC change-unit: process_commits dedups whole commit DOCS
@@ -352,6 +355,7 @@ def stage_index_commits(repo: str, commit_inputs: Sequence[Path], work: Path,
         "--max-tokens", str(TOKENIZE_BUDGET),
         "--tokenizer-path", TOKENIZER_PATH,
         "--format", "both",
+        "--memory-limit-gb", str(memory_limit_gb),
     ]
     if dedup_db is not None:
         cmd += ["--dedup-db", str(dedup_db)]
@@ -371,7 +375,12 @@ def stage_index_commits(repo: str, commit_inputs: Sequence[Path], work: Path,
     return enriched
 
 
-def stage_materialize(repo: str, enriched: Path, work: Path) -> Path:
+def stage_materialize(
+    repo: str,
+    enriched: Path,
+    work: Path,
+    memory_limit_gb: float = 10.0,
+) -> Path:
     """clang_enriched_to_parquet.py -> tokenized enriched parquet (single file)."""
     tok = work / f"{repo}.tok.parquet"
     run_checked(
@@ -385,6 +394,7 @@ def stage_materialize(repo: str, enriched: Path, work: Path) -> Path:
             "--materialize-tokenized-enriched",
             "--overflow-policy", "drop",
             "--size", _budget_size_label(TOKENIZE_BUDGET),
+            "--memory-limit-gb", str(memory_limit_gb),
         ],
         log_path=work / f"{repo}.materialize.log",
     )
@@ -469,6 +479,7 @@ def process_one_repo(
     dedup_db: Path | None = None,
     dedup_near: bool = True,
     global_symbol_index: Path | None = None,
+    memory_limit_gb: float = 10.0,
 ) -> dict:
     """Index a repo, then ROUTE each code doc to exactly ONE length bucket.
 
@@ -483,9 +494,11 @@ def process_one_repo(
     work = work_root / repo
     work.mkdir(parents=True, exist_ok=True)
     lengths_sorted = sorted(int(x) for x in target_lengths)
-    enriched = stage_index_source(repo, repo_dir, work, dedup_db, dedup_near,
-                                  global_symbol_index)
-    tok = stage_materialize(repo, enriched, work)
+    enriched = stage_index_source(
+        repo, repo_dir, work, dedup_db, dedup_near,
+        global_symbol_index, memory_limit_gb,
+    )
+    tok = stage_materialize(repo, enriched, work, memory_limit_gb)
 
     _bucket_for, route_by_fit = _route_by_fit_impl()
     route_dir = work / "routed"
@@ -509,13 +522,15 @@ def process_one_commit_source(
     repo_dir: Path | None,
     dedup_db: Path | None = None,
     dedup_near: bool = True,
+    memory_limit_gb: float = 10.0,
 ) -> dict:
     work = work_root / key
     work.mkdir(parents=True, exist_ok=True)
     lengths_sorted = sorted(int(x) for x in target_lengths)
     enriched = stage_index_commits(key, commit_inputs, work, repo_root, repo_dir,
-                                   dedup_db, dedup_near)
-    tok = stage_materialize(key, enriched, work)
+                                   dedup_db, dedup_near,
+                                   memory_limit_gb=memory_limit_gb)
+    tok = stage_materialize(key, enriched, work, memory_limit_gb)
 
     _bucket_for, route_by_fit = _route_by_fit_impl()
     route_dir = work / "routed"
@@ -571,6 +586,9 @@ def parse_args(argv: list[str]) -> argparse.Namespace:
                         "When set, the CODE stage threads it into index_project so "
                         "unresolved base-lib callees are pulled in as bounded "
                         "depth-1 deps tagged crosslib:<repo>. DEFAULT off.")
+    p.add_argument("--memory-limit-gb", type=float, default=10.0,
+                   help="Per-stage fail-loud RSS limit passed to index/materialize/"
+                        "commit processors (default 10.0).")
     return p.parse_args(argv)
 
 
@@ -650,7 +668,7 @@ def main(argv: list[str]) -> int:
                 key, files, target_lengths, work_root,
                 Path(args.commit_repo_root) if args.commit_repo_root else None,
                 Path(args.commit_repo_dir) if args.commit_repo_dir else None,
-                dedup_db, dedup_near,
+                dedup_db, dedup_near, args.memory_limit_gb,
             )
             manifest.mark_done(manifest_key, info)
             run_report[manifest_key] = info
@@ -681,7 +699,9 @@ def should_process(repo: str) -> bool:
             for repo, repo_dir in gen:
                 try:
                     info = process_one_repo(repo, repo_dir, target_lengths, work_root,
-                                            dedup_db, dedup_near, global_symbol_index)
+                                            dedup_db, dedup_near,
+                                            global_symbol_index,
+                                            args.memory_limit_gb)
                     manifest.mark_done(repo, info)
                     run_report[repo] = info
                     processed += 1

scripts/streaming_reindex_commits.py

@@ -363,6 +363,7 @@ def process_range(
     dedup_near: bool = True,
     pr_store: Path | None = None,
     repo_list: Path | None = None,
+    memory_limit_gb: float = 10.0,
 ) -> dict:
     """Full per-range pipeline. RAISES RepoFailure on any failure (no fallback)."""
     rkey = range_key(repo, start_idx)
@@ -375,8 +376,9 @@ def process_range(
         # process_commits needs the source tree for include resolution.
         enriched = stage_index_commits(rkey, [slice_jsonl], rwork, repo_dir, None,
                                        dedup_db, dedup_near,
-                                       pr_store=pr_store, repo_list=repo_list)
-        tok = stage_materialize(rkey, enriched, rwork)
+                                       pr_store=pr_store, repo_list=repo_list,
+                                       memory_limit_gb=memory_limit_gb)
+        tok = stage_materialize(rkey, enriched, rwork, memory_limit_gb)
 
         route_dir = rwork / "routed"
         routed = route_by_fit(tok, lengths_sorted, route_dir)
@@ -419,6 +421,7 @@ def process_one_repo(
     dedup_near: bool = True,
     pr_store: Path | None = None,
     repo_list: Path | None = None,
+    memory_limit_gb: float = 10.0,
 ) -> int:
     """Extract one repo, fan its ranges to the pool, wait for completion.
 
@@ -461,7 +464,7 @@ def process_one_repo(
             fut = pool.submit(
                 process_range, repo, repo_dir, records_jsonl,
                 start, end, lengths_sorted, repo_work,
-                dedup_db, dedup_near, pr_store, repo_list,
+                dedup_db, dedup_near, pr_store, repo_list, memory_limit_gb,
             )
             futures[fut] = (start, end)
 
@@ -537,6 +540,9 @@ def parse_args(argv: list[str]) -> argparse.Namespace:
     p.add_argument("--repo-list", default=None,
                    help="Path to outputs/pr_ingest/repo_list.json (bare-name -> "
                         "owner/repo map) for resolving the PR-store key.")
+    p.add_argument("--memory-limit-gb", type=float, default=10.0,
+                   help="Per-stage fail-loud RSS limit passed to process_commits/"
+                        "materializer (default 10.0).")
     return p.parse_args(argv)
 
 
@@ -636,6 +642,7 @@ def should_process(repo: str) -> bool:
                     work_root, pool, manifest, manifest_lock, resume,
                     args.token_budget, cumulative, args.keep_temp,
                     dedup_db, dedup_near, pr_store, repo_list,
+                    args.memory_limit_gb,
                 )
                 ranges_done += done
                 processed_repos += 1