Merge unseen-pool-size: configurable pool size / slices

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

benchmarks/unseen/bisect_guard.py

@@ -50,8 +50,9 @@ def read_pair(public_path: Path, private_path: Path) -> tuple[dict, dict]:
             raise ValueError
         if set(hidden) != {r["id"] for r in rows if r.get("private")}:
             raise ValueError
+        # Slice IDs come from the draft; bisect may leave sparse slices.
         for row in rows:
-            if type(row["slice"]) is not int or row["slice"] not in (1, 2, 3):
+            if type(row["slice"]) is not int or row["slice"] < 1:
                 raise ValueError
             if row.get("private"):
                 full = hidden[row["id"]]

benchmarks/unseen/draft.py

@@ -37,8 +37,18 @@ def private_digest(row: dict) -> str:
     return hashlib.sha256(canonical({"url": row["url"], "commits": [[c["base"], c["head"]] for c in row["changes"]]})).hexdigest()
 
 
-def assign_slices(rows: list[dict], *, seed: int) -> list[dict]:
-    if len(rows) != 300 or len({r["repository"] for r in rows}) != 300:
+def validate_pool(pool_size: int, slices: int) -> None:
+    if pool_size < 1:
+        raise ValueError("--pool-size must be at least 1")
+    if slices < 1 or slices > pool_size:
+        raise ValueError("--slices must be between 1 and --pool-size")
+    if pool_size % slices:
+        raise ValueError("--pool-size must be divisible by --slices")
+
+
+def assign_slices(rows: list[dict], *, seed: int, pool_size: int = 300, slices: int = 3) -> list[dict]:
+    validate_pool(pool_size, slices)
+    if len(rows) != pool_size or len({r["repository"] for r in rows}) != pool_size:
         raise ValueError("pool_size")
     for row in rows:
         kinds = Counter(c["kind"] for c in row["changes"])
@@ -52,12 +62,12 @@ def assign_slices(rows: list[dict], *, seed: int) -> list[dict]:
     groups = defaultdict(list)
     for row in sorted(rows, key=lambda r: r["repository"]):
         groups[row["ecosystem"], row["post_cutoff"]].append(row)
-    sizes = Counter({1: 0, 2: 0, 3: 0})
+    sizes = Counter({i: 0 for i in range(1, slices + 1)})
     assigned = []
     for key in sorted(groups):
         group = groups[key]
         rng.shuffle(group)
-        quotient, remainder = divmod(len(group), 3)
+        quotient, remainder = divmod(len(group), slices)
         # Give each stratum its floor on all slices, then fill globally least-full slices.
         quotas = {i: quotient for i in sizes}
         order = list(sizes)
@@ -72,7 +82,7 @@ def assign_slices(rows: list[dict], *, seed: int) -> list[dict]:
                 assigned.append({**row, "id": "r-" + digest[:24], "slice": slice_id})
             sizes[slice_id] += quotas[slice_id]
             offset += quotas[slice_id]
-    if set(sizes.values()) != {100}:
+    if set(sizes.values()) != {pool_size // slices}:
         raise ValueError("slice_size")
     return sorted(assigned, key=lambda r: (r["slice"], r["id"]))
 
@@ -236,7 +246,7 @@ def repository_changes(git, candidate, fixes: set[str], *, seed: int):
     return [change, *ordinary], exclusions, proof
 
 
-def inputs(used, seed, known_empty, clones):
+def inputs(used, seed, known_empty, clones, *, pool_size=300, slices=3):
     from openultrasast.model.taxonomy import DEFAULT_FAMILIES_PATH
 
     parameters = {
@@ -245,6 +255,8 @@ def inputs(used, seed, known_empty, clones):
         "known_empty": sorted(known_empty),
         "clone_limit_bytes": getattr(clones, "limit_bytes", 500 * 1024 * 1024),
         "journal_version": 4,
+        "pool_size": pool_size,
+        "slices": slices,
     }
     # Hardcoded design thresholds and algorithms are covered as well as taxonomy data.
     for name in ("draft", "source", "extract", "eligibility", "journal"):
@@ -345,7 +357,7 @@ def finish_license(result, github):
 
 
 @contextmanager
-def ordered_extractions(rows, api, osv, used, clones, fixes, seed, checkpoint, entries, workers):
+def ordered_extractions(rows, api, osv, used, clones, fixes, seed, checkpoint, entries, workers, pool_size):
     """Bounded in-flight work, completion journaling, and seed-order acceptance.
 
     Only running jobs count against `workers`; completed results wait in the
@@ -390,16 +402,16 @@ def ordered_extractions(rows, api, osv, used, clones, fixes, seed, checkpoint, e
         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:
+        while (pending or not exhausted) and len(entries) < pool_size:
             # 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:
+                if len(entries) == pool_size:
                     return
-            while not exhausted and len(in_flight) < workers:
+            while not exhausted and len(in_flight) < min(workers, pool_size - len(entries)):
                 checkpoint.boundary()
                 row = next(ordered, None)
                 if row is None:
@@ -492,7 +504,10 @@ def run(
     max_hours=None,
     clock=None,
     extract_workers=None,
+    pool_size=300,
+    slices=3,
 ) -> dict:
+    validate_pool(pool_size, slices)
     workers = extract_workers if extract_workers is not None else getattr(clones, "extract_workers", 6)
     if workers < 1:
         raise ValueError("invalid_extract_workers")
@@ -512,7 +527,7 @@ def run(
         checkpoint = Journal(
             path, max_hours=max_hours, clock=clock, rate_limit=lambda: getattr(getattr(github, "transport", None), "rate_limit", {})
         )
-        with checkpoint.open(inputs(used, seed, known_empty, clones)):
+        with checkpoint.open(inputs(used, seed, known_empty, clones, pool_size=pool_size, slices=slices)):
             checkpoint.boundary()
             api = ResolvedGitHub(github, checkpoint)
             used = copy.deepcopy(used)  # canonicalization must not mutate the resume input
@@ -534,6 +549,7 @@ def run(
                 checkpoint,
                 entries,
                 workers,
+                pool_size,
             ) as outcomes:
                 for row, decision, result in outcomes:
                     counts[source.ecosystem(row)] += 1
@@ -556,7 +572,7 @@ def run(
                     # only zero accepted entries indicate a systemic failure.
                     if instrument_failures >= 25 and not entries:
                         raise extract.InstrumentFailure("systemic_extraction_failure")
-                    if len(entries) == 300:
+                    if len(entries) == pool_size:
                         break
             report.update(
                 {
@@ -567,11 +583,11 @@ def run(
                 }
             )
             checkpoint.boundary()
-            if len(entries) != 300:
+            if len(entries) != pool_size:
                 report["status"] = "insufficient_repositories"
                 checkpoint.progress(report["status"])
                 return report
-            assigned = assign_slices(entries, seed=seed)
+            assigned = assign_slices(entries, seed=seed, pool_size=pool_size, slices=slices)
             completion = {"manifest_sha256": digest(assigned), "repositories": len(entries)}
             if "all" not in checkpoint.rows["completion"]:
                 checkpoint.put("completion", "all", completion)
@@ -582,8 +598,8 @@ def run(
             report.update(
                 {
                     "status": "draft",
-                    "changes": 3900,
-                    "slices": 3,
+                    "changes": pool_size * 13,
+                    "slices": slices,
                     "license_split": dict(Counter(e["license_class"] for e in entries)),
                     "methods": dict(Counter(e["changes"][0]["method"] for e in entries)),
                     "post_cutoff": sum(e["post_cutoff"] for e in entries),
@@ -622,12 +638,18 @@ def main(argv: list[str] | None = None) -> int:
     parser.add_argument("--max-hours", type=float, help="Stop between units after this many hours per invocation")
     parser.add_argument("--extract-deadline", type=float, default=240.0, help="Extraction wall-clock budget per repository in seconds")
     parser.add_argument("--extract-workers", type=int, default=6, help="Maximum extraction processes, capped to CPU cores (default: 6)")
+    parser.add_argument("--pool-size", type=int, default=300, help="Repositories to accept (default: 300)")
+    parser.add_argument("--slices", type=int, default=3, help="Equal-sized slices (default: 3)")
     parser.add_argument("--seed", type=int, default=20260601)
     parser.add_argument("--local-store")
     parser.add_argument("--s3-store", default="s3://")
     parser.add_argument("--pointer-manifest", type=Path, action="append", default=[])
     parser.add_argument("--known-empty", choices=(*eligibility.SOURCES, "memory_local", "memory_s3"), action="append", default=[])
     args = parser.parse_args(argv)
+    try:
+        validate_pool(args.pool_size, args.slices)
+    except ValueError as exc:
+        parser.error(str(exc))
     if args.extract_workers < 1:
         parser.error("--extract-workers must be at least 1")
     try:
@@ -653,6 +675,8 @@ def main(argv: list[str] | None = None) -> int:
             clones=clones,
             used=used,
             seed=args.seed,
+            pool_size=args.pool_size,
+            slices=args.slices,
             known_empty=set(args.known_empty),
             journal_path=clones.cache / "journal",
             max_hours=args.max_hours,

tests/test_unseen_bisect.py

@@ -14,7 +14,7 @@ d = importlib.import_module("benchmarks.unseen.draft")
 b = importlib.import_module("benchmarks.unseen.bisect_guard")
 
 
-def draft_pair(tmp_path, count=9):
+def draft_pair(tmp_path, count=9, slices=3):
     root, cache = tmp_path / "tree", tmp_path / "cache"
     root.mkdir()
     cache.mkdir()
@@ -29,7 +29,7 @@ def draft_pair(tmp_path, count=9):
                 "url": "https://" + f"github.com/synthetic/candidate-{i}",
                 "license_class": "private" if i % 2 else "permissive",
                 "license_spdx": "GPL-3.0" if i % 2 else "MIT",
-                "slice": i % 3 + 1,
+                "slice": i % slices + 1,
                 "changes": changes,
             }
         )

tests/test_unseen_draft.py

@@ -451,7 +451,10 @@ def test_systemic_zero_source_is_instrument_failure_with_recorded_rejections(tmp
 @pytest.mark.parametrize(
     "deadline_args,deadline,workers", [([], 240.0, 6), (["--extract-deadline", "12.5", "--extract-workers", "2"], 12.5, 2)]
 )
-def test_cli_uses_environment_and_injected_stores(tmp_path, monkeypatch, capsys, deadline_args, deadline, workers):
+@pytest.mark.parametrize("pool_args,pool_size,slices", [([], 300, 3), (["--pool-size", "100", "--slices", "1"], 100, 1)])
+def test_cli_uses_environment_and_injected_stores(
+    tmp_path, monkeypatch, capsys, deadline_args, deadline, workers, pool_args, pool_size, slices
+):
     import json
 
     from openultrasast import config
@@ -464,10 +467,17 @@ def test_cli_uses_environment_and_injected_stores(tmp_path, monkeypatch, capsys,
     monkeypatch.setattr(d.eligibility, "build_used", lambda *a, **kw: e.UsedSet())
     monkeypatch.setattr(d.source, "GitHub", lambda token: seen.append(token) or FakeAPI(0))
     monkeypatch.setattr(d.source, "OSV", FakeOSV)
-    args = ["--root", str(tmp_path), "--cache", str(tmp_path.parent / "clone-cache")]
+    args = ["--root", str(tmp_path), "--cache", str(tmp_path.with_name(tmp_path.name + "-clone-cache"))]
     for label in e.SOURCES:
         args += ["--known-empty", label]
-    assert d.main(args + deadline_args) == 1
+    original_run = d.run
+
+    def run(*args, **kwargs):
+        assert kwargs["pool_size"] == pool_size and kwargs["slices"] == slices
+        return original_run(*args, **kwargs)
+
+    monkeypatch.setattr(d, "run", run)
+    assert d.main(args + deadline_args + pool_args) == 1
     output = capsys.readouterr().out
     assert json.loads(output)["extract_deadline_seconds"] == deadline
     assert json.loads(output)["extract_workers"] == min(workers, os.cpu_count() or 4)
@@ -789,7 +799,7 @@ def test_resume_refuses_changed_inputs_before_network(tmp_path, monkeypatch, par
         kwargs["clones"].limit_bytes = 12
     else:
         original = d.inputs
-        monkeypatch.setattr(d, "inputs", lambda *a: {**original(*a), "design_extract": "changed"})
+        monkeypatch.setattr(d, "inputs", lambda *a, **kw: {**original(*a, **kw), "design_extract": "changed"})
     report = d.run(**kwargs)
     assert report["status"] == "resume_refused"
     assert report["changed_parameters"] == [parameter]
@@ -1084,7 +1094,8 @@ def test_decide_http_failure_systemic_guard_aborts(tmp_path):
     assert report["rejections"].get("decide_http_failure") == 25
 
 
-def test_parallel_pool_is_byte_identical_and_fork_selection_is_seed_ordered(tmp_path, monkeypatch):
+@pytest.mark.parametrize("pool_size,slices", [(300, 3), (100, 1)])
+def test_parallel_pool_is_byte_identical_and_fork_selection_is_seed_ordered(tmp_path, monkeypatch, pool_size, slices):
     journal_module = importlib.import_module("benchmarks.unseen.journal")
     monkeypatch.setattr(journal_module.os, "fsync", lambda fd: None)
     coordinator = os.getpid()
@@ -1146,9 +1157,9 @@ def test_parallel_pool_is_byte_identical_and_fork_selection_is_seed_ordered(tmp_
 
     original_assign, original_put, original_boundary = d.assign_slices, d.Journal.put, d.Journal.boundary
 
-    def assign(rows, *, seed):
+    def assign(rows, *, seed, pool_size=300, slices=3):
         selected.append([row["repository"] for row in rows])
-        return original_assign(rows, seed=seed)
+        return original_assign(rows, seed=seed, pool_size=pool_size, slices=slices)
 
     def put(self, *args):
         assert os.getpid() == coordinator
@@ -1179,8 +1190,14 @@ def test_parallel_pool_is_byte_identical_and_fork_selection_is_seed_ordered(tmp_
             known_empty=set(used.sources),
             journal_path=path,
             extract_workers=workers,
+            pool_size=pool_size,
+            slices=slices,
         )
         assert report["status"] == "draft", report
+        assert report["repositories"] == pool_size
+        assert report["changes"] == pool_size * 13 and report["slices"] == slices
+        published = tomllib.loads((path / "draft-p1.toml").read_text())["repository"]
+        assert Counter(row["slice"] for row in published) == dict.fromkeys(range(1, slices + 1), pool_size // slices)
         assert report["rejections"] == {"fork_network": 2, "extract_git_exit_128": 2, "extraction_timeout": 1, "decide_http_failure": 1}
         assert report["extract_workers"] == min(workers, 4)
         assert 1 <= peak[0] <= min(workers, 4)
@@ -1190,7 +1207,7 @@ def test_parallel_pool_is_byte_identical_and_fork_selection_is_seed_ordered(tmp_
             (report["manifest_sha256"], *(path.joinpath(name).read_bytes() for name in ("draft-p1.toml", "private-draft-p1.toml")))
         )
     # Independent serial oracle: first successful repository wins each network.
-    expected = [name for i, name in enumerate(names) if i not in {0, 3, 5, 6, 7, 8}][:300]
+    expected = [name for i, name in enumerate(names) if i not in {0, 3, 5, 6, 7, 8}][:pool_size]
     assert selected == [expected, expected]
     assert snapshots[0] == snapshots[1]
 
@@ -1238,9 +1255,9 @@ def test_completion_journaling_refills_pool_while_seed_cursor_is_blocked(tmp_pat
                 assert set(range(1, release_rank + 1)) <= set(journaled)
                 release_first.set()
 
-    def assign(rows, *, seed):
+    def assign(rows, *, seed, pool_size=300, slices=3):
         accepted.append([ranks[row["repository"]] for row in rows])
-        return original_assign(rows, seed=seed)
+        return original_assign(rows, seed=seed, pool_size=pool_size, slices=slices)
 
     monkeypatch.setattr(d, "repository_changes", extract_changes)
     monkeypatch.setattr(d.Journal, "put", put)
@@ -1354,3 +1371,75 @@ def test_parallel_license_failure_is_a_resilient_skip(tmp_path, monkeypatch, fai
     )
     assert report["status"] == "insufficient_repositories"
     assert report["repositories"] == 1 and report["rejections"] == {reason: 1}
+
+
+@pytest.mark.parametrize("pool_size,slices", [(100, 1), (100, 4), (12, 12), (1, 1)])
+def test_configurable_slices_are_balanced_stratified_and_deterministic(pool_size, slices):
+    rows = entries()[:pool_size]
+    assigned = d.assign_slices(rows, seed=19, pool_size=pool_size, slices=slices)
+    assert assigned == d.assign_slices(list(reversed(rows)), seed=19, pool_size=pool_size, slices=slices)
+    assert Counter(row["slice"] for row in assigned) == dict.fromkeys(range(1, slices + 1), pool_size // slices)
+    assert assigned == sorted(assigned, key=lambda row: (row["slice"], row["id"]))
+    strata = Counter((row["ecosystem"], row["post_cutoff"]) for row in rows)
+    for slice_id in range(1, slices + 1):
+        counts = Counter((row["ecosystem"], row["post_cutoff"]) for row in assigned if row["slice"] == slice_id)
+        assert all(abs(counts[key] - total / slices) < 1 for key, total in strata.items())
+
+
+def test_default_assignment_matches_original_digest():
+    # Captured from HEAD's unparameterized algorithm on this synthetic fixture.
+    assert d.digest(d.assign_slices(entries(), seed=19)) == "811f0893c73ebb82a8d19948e6017b9a6f353f4de98c281f5861706d0f577b78"
+
+
+@pytest.mark.parametrize("count", [99, 101])
+def test_single_slice_requires_exact_pool_size(count):
+    with pytest.raises(ValueError, match="pool_size"):
+        d.assign_slices(entries()[:count], seed=19, pool_size=100, slices=1)
+
+
+def test_single_slice_rejects_duplicate_repositories():
+    rows = entries()[:100]
+    rows[-1] = rows[0]
+    with pytest.raises(ValueError, match="pool_size"):
+        d.assign_slices(rows, seed=19, pool_size=100, slices=1)
+
+
+@pytest.mark.parametrize("pool_size,slices", [(0, 1), (-1, 1), (100, 0), (100, -1), (100, 101), (100, 3)])
+def test_invalid_pool_configuration_rejected_before_io(tmp_path, monkeypatch, pool_size, slices):
+    from openultrasast import config
+
+    def unexpected(*args, **kwargs):
+        pytest.fail("invalid configuration performed I/O")
+
+    monkeypatch.setattr(config, "load_dotenv", unexpected)
+    with pytest.raises(SystemExit) as exc:
+        d.main(["--pool-size", str(pool_size), "--slices", str(slices)])
+    assert exc.value.code == 2
+    api, osv, clones = FakeAPI(), FakeOSV(), FakeClones()
+    with pytest.raises(ValueError):
+        d.run(tmp_path, github=api, osv=osv, clones=clones, used=e.UsedSet(), seed=7, pool_size=pool_size, slices=slices)
+    assert api.calls == osv.calls == 0 and not clones.counts
+    with pytest.raises(ValueError):
+        d.assign_slices([], seed=7, pool_size=pool_size, slices=slices)
+
+
+@pytest.mark.parametrize("changed", [{"pool_size": 200}, {"slices": 2}])
+def test_pool_configuration_change_refuses_resume_before_io(tmp_path, changed):
+    used, api, clones = e.UsedSet(), FakeAPI(), FakeClones()
+    path = tmp_path / "journal"
+    parameters = d.inputs(used, 7, set(used.sources), clones, pool_size=100, slices=1)
+    with d.Journal(path).open(parameters):
+        pass
+    report = d.run(
+        tmp_path / "repo",
+        github=api,
+        osv=FakeOSV(),
+        clones=clones,
+        used=used,
+        seed=7,
+        known_empty=set(used.sources),
+        journal_path=path,
+        **({"pool_size": 100, "slices": 1} | changed),
+    )
+    assert report["status"] == "resume_refused" and report["changed_parameters"] == list(changed)
+    assert api.calls == 0 and not clones.counts

tests/test_unseen_freeze.py

@@ -195,3 +195,30 @@ def test_freeze_cli_redacts_failure_and_reports_only_record(tmp_path, capsys):
     report.write_text(rows[0]["url"])
     assert f.main(["--root", str(root), "--used-report", str(report)]) == 1
     assert capsys.readouterr() == ('{"status": "freeze_refused"}\n', "")
+
+
+@pytest.mark.parametrize("slices", [1, 4])
+def test_hundred_repository_draft_bisects_and_freezes(tmp_path, slices):
+    root, draft, rows = draft_pair(tmp_path, count=100, slices=slices)
+    from test_unseen_bisect import staged
+
+    seen = []
+
+    def guard(root):
+        current = staged(root)
+        seen.append(len(current))
+        return 0
+
+    receipt = b.run(root, draft, guard_runner=guard)
+    assert receipt["input_count"] == 100 and receipt["dropped_count"] == 0
+    assert 100 in seen and staged(root) == rows
+    directory = root / "benchmarks/unseen"
+    (directory / "protocol-p1.md").write_bytes(Path("benchmarks/unseen/protocol-p1.md").read_bytes())
+    record = f.run(root, used_report=used())
+    public = tomllib.loads((directory / "pool-p1.toml").read_text())
+    private = tomllib.loads((directory / "private/pool-p1.toml").read_text())
+    assert len(public["repository"]) == 100 and len(private["repository"]) == 50
+    assert {row["slice"] for row in public["repository"]} == set(range(1, slices + 1))
+    assert record["freeze_digest"] == hashlib.sha256(d.canonical(public)).hexdigest()
+    assert sum(record["label_check_counts"].values()) == 100
+    assert f.run(root, used_report=used()) == record