Add code index timeout guard

Author: apstenku123

scripts/streaming_conveyor.py

@@ -1199,6 +1199,7 @@ def run_code_half(
     global_symbol_index: Path | None = None,
     memory_limit_gb: float = 10.0,
     parse_workers: int = 2,
+    index_timeout_s: int | None = None,
 ) -> dict:
     """index+route+pack the repo's source via the EXISTING code stage, zstd-max.
 
@@ -1214,7 +1215,7 @@ def run_code_half(
     try:
         info = sr.process_one_repo(
             repo, repo_dir, lengths_code, work_root, dedup_db, dedup_near,
-            global_symbol_index, memory_limit_gb, parse_workers,
+            global_symbol_index, memory_limit_gb, parse_workers, index_timeout_s,
             promote_dedup_on_success=False,
         )
         # zstd-max the per-length code parquet files this repo just wrote.
@@ -1245,6 +1246,7 @@ def run_code_half(
         Path | None,
         float,
         int,
+        int | None,
     ],
     dict,
 ]
@@ -1309,6 +1311,7 @@ def run_code_half_adaptive(
     global_symbol_index: Path | None = None,
     memory_limit_gb: float = 10.0,
     parse_workers: int = 2,
+    index_timeout_s: int | None = None,
     *,
     runner: CodeRunner | None = None,
     isolate_dedup_on_retry: bool = True,
@@ -1331,6 +1334,7 @@ def run_code_half_adaptive(
             global_symbol_index,
             memory_limit_gb,
             parse_workers,
+            index_timeout_s,
         )
     except RepoFailure as exc:
         if parse_workers <= 1 or not is_index_project_memory_failure(exc):
@@ -1352,6 +1356,7 @@ def run_code_half_adaptive(
             global_symbol_index,
             memory_limit_gb,
             1,
+            index_timeout_s,
         )
 
 
@@ -2002,6 +2007,7 @@ def process_one_repo(
     *,
     code_memory_limit_gb: float | None = None,
     commit_memory_limit_gb: float | None = None,
+    code_index_timeout_s: int | None = None,
     reservations: UnitReservationLedger | None = None,
     range_submit_window: int = 1,
     analysis_cache_entries: int = 128,
@@ -2086,6 +2092,7 @@ def process_one_repo(
                             repo, repo_dir, lengths_code, work_root,
                             code_dedup_db, code_dedup_near,
                             global_symbol_index, code_limit, code_parse_workers,
+                            code_index_timeout_s,
                         )
                         with manifest_lock:
                             manifest.mark_done(ck, cinfo)
@@ -2289,6 +2296,9 @@ def parse_args(argv: list[str]) -> argparse.Namespace:
                    help="Parse workers passed to index_project for the CODE half "
                         "(default 2; avoid multiplying --repo-workers by the old "
                         "index_project default of 8 clang workers).")
+    p.add_argument("--code-index-timeout-s", type=int, default=0,
+                   help="Optional fail-loud timeout for each CODE index_project "
+                        "stage. 0 disables the timeout (default).")
     p.add_argument("--memory-budget-gb", type=float, default=None,
                    help="Global conveyor memory budget for heavy subprocesses. "
                         "Default is 55%% of physical RAM (or 48 GiB if RAM size "
@@ -2667,6 +2677,7 @@ def mark_no_git_repo(repo: str) -> None:
                     progress, checkpoint,
                     code_memory_limit_gb=code_memory_limit_gb,
                     commit_memory_limit_gb=commit_memory_limit_gb,
+                    code_index_timeout_s=args.code_index_timeout_s,
                     reservations=reservations,
                     range_submit_window=range_submit_window,
                     analysis_cache_entries=args.analysis_cache_entries,
@@ -2756,6 +2767,7 @@ def drain_one_or_more(block: bool = True) -> None:
                     progress, checkpoint,
                     code_memory_limit_gb=code_memory_limit_gb,
                     commit_memory_limit_gb=commit_memory_limit_gb,
+                    code_index_timeout_s=args.code_index_timeout_s,
                     reservations=reservations,
                     range_submit_window=range_submit_window,
                     analysis_cache_entries=args.analysis_cache_entries,

scripts/streaming_reindex.py

@@ -244,27 +244,51 @@ def run_checked(
     printable = " ".join(str(c) for c in cmd)
     print(f"  [{repo}] {stage}: {printable}", file=sys.stderr, flush=True)
     log_fh = open(log_path, "wb") if log_path else None
+    stdout_data: bytes | None = None
+    proc: subprocess.Popen[bytes] | None = None
     try:
-        proc = subprocess.run(
+        proc = subprocess.Popen(
             [str(c) for c in cmd],
             cwd=str(cwd) if cwd else None,
             env=_subprocess_env(),
             stdout=log_fh if log_fh else subprocess.PIPE,
             stderr=subprocess.STDOUT,
-            timeout=timeout,
+            start_new_session=bool(timeout),
         )
+        stdout_data, _ = proc.communicate(timeout=timeout)
     except subprocess.TimeoutExpired as exc:
+        if proc is not None:
+            try:
+                if timeout:
+                    os.killpg(proc.pid, signal.SIGTERM)
+                else:
+                    proc.terminate()
+            except ProcessLookupError:
+                pass
+            try:
+                proc.wait(timeout=5)
+            except subprocess.TimeoutExpired:
+                try:
+                    if timeout:
+                        os.killpg(proc.pid, signal.SIGKILL)
+                    else:
+                        proc.kill()
+                except ProcessLookupError:
+                    pass
+                proc.wait()
         raise RepoFailure(repo, stage, f"timed out after {timeout}s: {exc}") from exc
     finally:
         if log_fh:
             log_fh.close()
+    if proc is None:
+        raise RepoFailure(repo, stage, "subprocess did not start")
     if proc.returncode != 0:
         tail = ""
         if log_path and log_path.exists():
             data = log_path.read_bytes()[-4000:]
             tail = data.decode("utf-8", errors="replace")
-        elif proc.stdout:
-            tail = proc.stdout.decode("utf-8", errors="replace")[-4000:]
+        elif stdout_data:
+            tail = stdout_data.decode("utf-8", errors="replace")[-4000:]
         raise RepoFailure(
             repo, stage, f"exit code {proc.returncode}\n--- last output ---\n{tail}"
         )
@@ -423,6 +447,7 @@ def stage_index_source(
     global_symbol_index: Path | None = None,
     memory_limit_gb: float = 10.0,
     parse_workers: int = 2,
+    index_timeout_s: int | None = None,
 ) -> Path:
     """index_project.py --enriched -> <repo>.enriched.jsonl.
 
@@ -461,6 +486,7 @@ def stage_index_source(
         "index_project",
         cmd,
         log_path=work / f"{repo}.index.log",
+        timeout=index_timeout_s if index_timeout_s and index_timeout_s > 0 else None,
     )
     if not enriched.exists() or enriched.stat().st_size == 0:
         raise RepoFailure(repo, "index_project", f"empty enriched jsonl: {enriched}")
@@ -628,6 +654,7 @@ def process_one_repo(
     global_symbol_index: Path | None = None,
     memory_limit_gb: float = 10.0,
     parse_workers: int = 2,
+    index_timeout_s: int | None = None,
     *,
     promote_dedup_on_success: bool = True,
 ) -> dict:
@@ -655,7 +682,7 @@ def process_one_repo(
         enriched = stage_index_source(
             repo, repo_dir, work, dedup_db, dedup_near,
             stage_id, stage_db, global_symbol_index, memory_limit_gb,
-            parse_workers,
+            parse_workers, index_timeout_s,
         )
         timings["index_project_s"] = round(time.monotonic() - started, 6)
         started = time.monotonic()
@@ -795,6 +822,9 @@ def parse_args(argv: list[str]) -> argparse.Namespace:
     p.add_argument("--parse-workers", type=int, default=2,
                    help="Parse workers passed to index_project for the code stage "
                         "(default 2; keep this low when multiple repos run).")
+    p.add_argument("--code-index-timeout-s", type=int, default=0,
+                   help="Optional fail-loud timeout for each index_project code "
+                        "stage. 0 disables the timeout (default).")
     return p.parse_args(argv)
 
 
@@ -911,7 +941,8 @@ def should_process(repo: str) -> bool:
                                             dedup_db, dedup_near,
                                             global_symbol_index,
                                             args.memory_limit_gb,
-                                            args.parse_workers)
+                                            args.parse_workers,
+                                            args.code_index_timeout_s)
                     manifest.mark_done(repo, info)
                     run_report[repo] = info
                     processed += 1

tests/conveyor_ckpt_harness.py

@@ -165,7 +165,7 @@ def _lengths_info(lengths, valid):
 
 def fake_run_code_half(repo, repo_dir, lengths_code, work_root, dedup_db,
                        dedup_near, global_symbol_index=None, memory_limit_gb=10.0,
-                       parse_workers=2):
+                       parse_workers=2, index_timeout_s=None):
     return {"source": f"{repo}::code", "lengths": _lengths_info(lengths_code, 10)}
 
 

tests/test_streaming_conveyor_progress.py

@@ -284,8 +284,9 @@ def runner(
         global_symbol_index,
         memory_limit_gb,
         parse_workers,
+        index_timeout_s,
     ):
-        calls.append((parse_workers, dedup_db, dedup_near))
+        calls.append((parse_workers, dedup_db, dedup_near, index_timeout_s))
         if parse_workers > 1:
             raise streaming_conveyor.RepoFailure(
                 repo,
@@ -308,12 +309,13 @@ def runner(
         None,
         8.0,
         2,
+        7200,
         runner=runner,
     )
 
     assert calls == [
-        (2, tmp_path / "global.sqlite", True),
-        (1, None, False),
+        (2, tmp_path / "global.sqlite", True, 7200),
+        (1, None, False, 7200),
     ]
     assert info["lengths"]["1024"]["valid_tokens"] == 1024
 

tests/test_streaming_reindex_run_checked.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+import sys
+
+
+def test_run_checked_times_out_fail_loud(tmp_path) -> None:
+    import streaming_reindex
+
+    log_path = tmp_path / "sleep.log"
+    try:
+        streaming_reindex.run_checked(
+            "repo",
+            "index_project",
+            [
+                sys.executable,
+                "-c",
+                "import time; print('started', flush=True); time.sleep(30)",
+            ],
+            log_path=log_path,
+            timeout=1,
+        )
+    except streaming_reindex.RepoFailure as exc:
+        assert exc.repo == "repo"
+        assert exc.stage == "index_project"
+        assert "timed out after 1s" in exc.detail
+    else:  # pragma: no cover
+        raise AssertionError("expected run_checked to fail loud on timeout")
+
+    assert "started" in log_path.read_text(encoding="utf-8")