Merge unseen-extract-ascompleted: saturate the extraction pool

Co-Authored-By: Claude Opus 4.8 <[email protected]>

benchmarks/unseen/draft.py

@@ -18,7 +18,7 @@ import subprocess
 import tempfile
 import time
 from collections import Counter, defaultdict, deque
-from concurrent.futures import ProcessPoolExecutor
+from concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, wait
 from concurrent.futures.process import BrokenProcessPool
 from contextlib import contextmanager
 from dataclasses import asdict
@@ -342,47 +342,60 @@ def finish_license(result, github):
 
 @contextmanager
 def ordered_extractions(rows, api, osv, used, clones, fixes, seed, checkpoint, entries, workers):
-    """Bounded speculation, ordered consumption, and coordinator-only durable writes.
+    """Bounded in-flight work, completion journaling, and seed-order acceptance.
 
-    At most `workers` candidates are ahead of the acceptance cursor. Counters are
-    private to each job, including clone cleanup, and merged exactly once here.
+    Only running jobs count against `workers`; completed results wait in the
+    journal while the acceptance cursor is blocked. Counters are private to each
+    job, including clone cleanup, and merged exactly once on completion here.
     Graceful STOP/guard/target exits drain started work into the journal for resume;
     a hard interrupt may lose unjournaled work but cannot change accepted order.
     """
-    pending = deque()
+    pending = deque()  # Seed-order acceptance cursor; contains no futures/results.
+    in_flight = {}  # Future -> extraction key, independent of the cursor.
     # Fork shares the coordinator's memory copy-on-write. Jobs receive only the
     # picklable extraction inputs, never the journal, REST clients or stores.
     context = multiprocessing.get_context("fork")
     pool = ProcessPoolExecutor(max_workers=workers, mp_context=context)
 
-    def collect(item):
-        row, decision, key, future = item
-        if future is not None:
-            try:
-                result = future.result()
-            except BrokenProcessPool:
-                # A killed worker also fails unrelated outstanding futures. Do not
-                # mistake these collateral failures for a broken git instrument.
-                result = {"rejected": "worker_crash", "ordinary_exclusions": {}, "counts": {}}
-            except Exception:
-                # Includes task/result serialization failures. Keep diagnostics
-                # identity-free; repeated task errors still reach the normal guard.
-                result = {"rejected": "extract_worker_failure", "ordinary_exclusions": {}, "counts": {}}
-            for name, value in result["counts"].items():
-                if name == "peak_clone_bytes":
-                    clones.counts[name] = max(clones.counts[name], value)
-                else:
-                    clones.counts[name] += value
-            result = finish_license(result, api)
-            checkpoint.put("extraction", key, result)
-        return row, decision, checkpoint.rows["extraction"].get(key)
+    def collect(future):
+        key = in_flight[future]
+        try:
+            result = future.result()
+        except BrokenProcessPool:
+            # A killed worker also fails unrelated outstanding futures. Do not
+            # mistake these collateral failures for a broken git instrument.
+            result = {"rejected": "worker_crash", "ordinary_exclusions": {}, "counts": {}}
+        except Exception:
+            # Includes task/result serialization failures. Keep diagnostics
+            # identity-free; repeated task errors still reach the normal guard.
+            result = {"rejected": "extract_worker_failure", "ordinary_exclusions": {}, "counts": {}}
+        for name, value in result["counts"].items():
+            if name == "peak_clone_bytes":
+                clones.counts[name] = max(clones.counts[name], value)
+            else:
+                clones.counts[name] += value
+        result = finish_license(result, api)
+        checkpoint.put("extraction", key, result)
+        del in_flight[future]
 
     def iterate():
         nonlocal pool
         ordered = iter(rows)
         exhausted = False
-        while pending or not exhausted:
-            while not exhausted and len(pending) < min(workers, 300 - len(entries)):
+
+        def ready():
+            return pending and (pending[0][2] is None or pending[0][2] in checkpoint.rows["extraction"])
+
+        while (pending or not exhausted) and len(entries) < 300:
+            # Only this cursor feeds acceptance, fork deduplication and the guard.
+            # Replayed results follow the same path and need no worker submission.
+            while ready():
+                checkpoint.boundary()
+                row, decision, key = pending.popleft()
+                yield row, decision, checkpoint.rows["extraction"].get(key)
+                if len(entries) == 300:
+                    return
+            while not exhausted and len(in_flight) < workers:
                 checkpoint.boundary()
                 row = next(ordered, None)
                 if row is None:
@@ -409,11 +422,21 @@ def ordered_extractions(rows, api, osv, used, clones, fixes, seed, checkpoint, e
                             pool.shutdown(wait=True, cancel_futures=True)
                             pool = ProcessPoolExecutor(max_workers=workers, mp_context=context)
                             future = pool.submit(extract_repository, *args)
-                pending.append((row, decision, extraction_key, future))
-            if not pending:
-                return
-            checkpoint.boundary()
-            yield collect(pending.popleft())
+                        in_flight[future] = extraction_key
+                pending.append((row, decision, extraction_key))
+                # Finalize cached decisions promptly (especially on resume), and
+                # service completed jobs even during a long run of decide skips.
+                if ready() or any(future.done() for future in in_flight):
+                    break
+            if ready():
+                continue
+            if in_flight:
+                checkpoint.boundary()
+                done, _ = wait(in_flight, return_when=FIRST_COMPLETED)
+                for future in done:
+                    collect(future)
+                # Freed slots are refilled on the next pass even if the oldest
+                # candidate is still running. No result() waits on that cursor.
 
     interrupted = False
     stopped = None
@@ -426,10 +449,9 @@ def ordered_extractions(rows, api, osv, used, clones, fixes, seed, checkpoint, e
     finally:
         pool.shutdown(wait=True, cancel_futures=True)
         if not interrupted:
-            for item in pending:
-                future = item[3]
-                if future is not None and not future.cancelled():
-                    collect(item)
+            for future in list(in_flight):
+                if not future.cancelled():
+                    collect(future)
             if stopped is not None:
                 checkpoint.progress("stopped", reason=stopped)
 

tests/test_unseen_draft.py

@@ -272,14 +272,16 @@ def test_systemic_extraction_guard_boundaries(tmp_path, monkeypatch, successes,
     expected_failures = min(failures, 201 if successes else 25)
     assert report["status"] == status
     assert report["rejections"] == {"extract_git_exit_128": expected_failures}
-    assert successes + expected_failures <= clones.counts["clones"] <= successes + min(failures, expected_failures + 5)
+    assert successes + expected_failures <= clones.counts["clones"] <= successes + failures
     assert not (journal_path / "draft-p1.toml").exists()
     if status == "instrument_failure":
         assert report["error_type"] == "InstrumentFailure"
-        # Replayed failures trip the same guard; only missing speculation may run.
+        # Completed speculation can run arbitrarily far ahead of a slow cursor.
+        # Replayed failures trip the same guard; cached results are not extracted again.
+        completed_clones = clones.counts["clones"]
         resumed = d.run(**kwargs)
         assert resumed["status"] == status and resumed["rejections"] == report["rejections"]
-        assert successes + expected_failures <= clones.counts["clones"] <= successes + min(failures, expected_failures + 5)
+        assert clones.counts["clones"] == completed_clones
 
 
 @pytest.mark.parametrize(
@@ -1123,6 +1125,75 @@ def test_parallel_pool_is_byte_identical_and_fork_selection_is_seed_ordered(tmp_
     assert snapshots[0] == snapshots[1]
 
 
+@pytest.mark.parametrize("failures", [0, 25])
+def test_completion_journaling_refills_pool_while_seed_cursor_is_blocked(tmp_path, monkeypatch, failures):
+    monkeypatch.setattr(d.os, "fsync", lambda fd: None)
+    monkeypatch.setattr(d.os, "cpu_count", lambda: 2)
+    context = multiprocessing.get_context("fork")
+    release_first, first_finished = context.Event(), context.Event()
+    ordered = s.ordered_advisories(list(FakeAPI(310 + failures).advisories()), 7)
+    names = [row["references"][0].split("/commit/")[0].removeprefix("https://github.com/") for row in ordered]
+    ranks = {name: i for i, name in enumerate(names)}
+    key_ranks = {d.digest({"repository": name, "advisory": ordered[i]["ghsa_id"], "fix": "a" * 40}): i for i, name in enumerate(names)}
+    release_rank = max(6, failures)
+    expected = [0, *range(failures + 1, failures + 300)]
+    changes = entries()[0]["changes"]
+    coordinator = os.getpid()
+    journaled, accepted = [], []
+    original_put, original_assign = d.Journal.put, d.assign_slices
+
+    def extract_changes(git, candidate, fixes, *, seed):
+        assert os.getpid() != coordinator
+        if ranks[candidate.repository] == 0:
+            # The release rank is beyond the old bounded cursor window.
+            assert release_first.wait(10), "pool waited on the oldest future"
+            first_finished.set()
+        if 1 <= ranks[candidate.repository] <= failures:
+            raise d.InstrumentFailure("git_exit_128")
+        return copy.deepcopy(changes), {}, {"bytes": 12, "files": 1}
+
+    def put(self, stage, key, value):
+        assert os.getpid() == coordinator
+        original_put(self, stage, key, value)
+        if stage == "extraction":
+            rank = key_ranks[key]
+            if 1 <= rank <= failures:
+                assert value["rejected"] == "extract_git_exit_128"
+            else:
+                assert "entry" in value, value
+            journaled.append(rank)
+            if rank == release_rank:
+                assert not first_finished.is_set()
+                assert 0 not in journaled
+                assert set(range(1, release_rank + 1)) <= set(journaled)
+                release_first.set()
+
+    def assign(rows, *, seed):
+        accepted.append([ranks[row["repository"]] for row in rows])
+        return original_assign(rows, seed=seed)
+
+    monkeypatch.setattr(d, "repository_changes", extract_changes)
+    monkeypatch.setattr(d.Journal, "put", put)
+    monkeypatch.setattr(d, "assign_slices", assign)
+    used = e.UsedSet()
+    path = tmp_path / "journal"
+    kwargs = dict(root=tmp_path / "repo", osv=FakeOSV(), used=used, seed=7, known_empty=set(used.sources), journal_path=path)
+    result = d.run(**kwargs, github=FakeAPI(310 + failures), clones=FakeClones(), extract_workers=2)
+    assert result["status"] == "draft", result
+    assert journaled.index(release_rank) < journaled.index(0)
+    assert accepted == [expected]
+    # Even 25 completed instrument failures must wait for the earlier success:
+    # counting them at completion would incorrectly trip the zero-entry guard.
+    assert result["rejections"] == ({"extract_git_exit_128": failures} if failures else {})
+    snapshots = [(path / name).read_bytes() for name in ("draft-p1.toml", "private-draft-p1.toml")]
+    before = (path / "extraction.jsonl").read_bytes()
+    resumed = d.run(**kwargs, github=FakeAPI(310 + failures), clones=FakeClones(), extract_workers=1)
+    assert resumed["status"] == "draft" and resumed["instrument"] == {}
+    assert accepted == [expected] * 2
+    assert (path / "extraction.jsonl").read_bytes() == before
+    assert snapshots == [(path / name).read_bytes() for name in ("draft-p1.toml", "private-draft-p1.toml")]
+
+
 @pytest.mark.parametrize("workers", [0, -1])
 def test_invalid_extract_workers_rejected_before_io(tmp_path, workers):
     used, api = e.UsedSet(), FakeAPI()