diff --git a/benchmarks/unseen/draft.py b/benchmarks/unseen/draft.py index f9765b6..f25dd07 100644 --- a/benchmarks/unseen/draft.py +++ b/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) diff --git a/tests/test_unseen_draft.py b/tests/test_unseen_draft.py index c0f451c..e210b3c 100644 --- a/tests/test_unseen_draft.py +++ b/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()