From bffd47ac5a0f3298e57fa0f2ab812c8299446145 Mon Sep 17 00:00:00 2001 From: Stanley Phoong Date: Sat, 8 Aug 2026 17:36:25 -0700 Subject: [PATCH] feat(swe-bench): content-addressed unit plan + durable mkdir-atomic work queue Adds the two foundations of the distributed SWE-bench harness: - units.py: shards an instance-id list into immutable, content-addressed units. The sha256 digest covers the ordered id list, so a plan cannot be silently reused across a different run, instance list, or ordering. - queue.py: a filesystem work queue whose claim is a bare os.mkdir (never makedirs(exist_ok=True), which hands a unit to every caller). available() is plan - claims - results, so deleting a result alone does NOT requeue a unit; requeue() is the only supported path and removes the result, the claim and the attempt records together. Env faults are ledgered separately from counted attempts, and abandoning a unit publishes a terminal result AND releases the claim so claims/ and results/ never disagree. --- .../swe_bench_distributed/__init__.py | 30 ++ .../evaluation/swe_bench_distributed/queue.py | 495 ++++++++++++++++++ .../evaluation/swe_bench_distributed/units.py | 187 +++++++ .../test_units_and_queue.py | 246 +++++++++ 4 files changed, 958 insertions(+) create mode 100644 src/inference_endpoint/evaluation/swe_bench_distributed/__init__.py create mode 100644 src/inference_endpoint/evaluation/swe_bench_distributed/queue.py create mode 100644 src/inference_endpoint/evaluation/swe_bench_distributed/units.py create mode 100644 tests/unit/evaluation/swe_bench_distributed/test_units_and_queue.py diff --git a/src/inference_endpoint/evaluation/swe_bench_distributed/__init__.py b/src/inference_endpoint/evaluation/swe_bench_distributed/__init__.py new file mode 100644 index 000000000..2042c3d55 --- /dev/null +++ b/src/inference_endpoint/evaluation/swe_bench_distributed/__init__.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Distributed SWE-bench execution across a fleet of SWE-bench services. + +The single-service :class:`~inference_endpoint.evaluation.swe_bench_scorer.SWEBenchScorer` +issues one run covering every instance. This package shards the instance list +into units, dispatches units across several services concurrently, classifies +infrastructure damage separately from genuine model failures, and refuses to +emit an accuracy number unless every planned instance id is accounted for +exactly once. +""" + +from .queue import ( + ClaimError, + UnitOutcome, + UnitResult, + WorkQueue, +) +from .units import Unit, UnitPlan, plan_units + +__all__ = [ + "ClaimError", + "Unit", + "UnitOutcome", + "UnitPlan", + "UnitResult", + "WorkQueue", + "plan_units", +] diff --git a/src/inference_endpoint/evaluation/swe_bench_distributed/queue.py b/src/inference_endpoint/evaluation/swe_bench_distributed/queue.py new file mode 100644 index 000000000..3feddae63 --- /dev/null +++ b/src/inference_endpoint/evaluation/swe_bench_distributed/queue.py @@ -0,0 +1,495 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Durable work queue for distributed SWE-bench units. + +The queue is a directory tree so that a client crash costs nothing but the +in-flight units, and so that the merge gate reads durable records rather than +process memory. + +Layout under ``root``:: + + units.json immutable plan (see units.py) + claims//owner json owner record, written temp+rename + claims//hb heartbeat, mtime only + results/.json terminal record (succeeded OR abandoned) + failed/..json one record per *counted* attempt + failed/env/.*.json environment faults, NOT counted + failed/artifacts/... evidence snapshot taken before a retry + +Two invariants are load-bearing and are enforced here rather than by +convention: + +1. ``claim()`` is ``os.mkdir`` and nothing else. ``mkdir`` on an existing + directory fails atomically with ``EEXIST`` on every filesystem we run on, + including Lustre, so exactly one of N racing callers wins. ``makedirs(..., + exist_ok=True)`` would hand the unit to every caller. +2. ``requeue()`` is the only way to make a terminal unit runnable again, and it + removes the result, the claim tombstone *and* the counted attempt records + together. Deleting a result file by hand does not requeue a unit -- the + claim tombstone still hides it -- and that misunderstanding has cost real + campaign time. +""" + +from __future__ import annotations + +import errno +import logging +import os +import shutil +import socket +import time +from dataclasses import dataclass, field +from enum import StrEnum +from pathlib import Path +from typing import Any + +import msgspec + +from .units import PLAN_FILENAME, UnitPlan, read_plan + +logger = logging.getLogger(__name__) + +_OWNER = "owner" +_HEARTBEAT = "hb" +_CLAIM_CONTENTS = frozenset({_OWNER, _HEARTBEAT}) + +# Small files that explain a failure. Snapshotted before a retry reuses the +# unit's run directory: a unit that fails and then succeeds otherwise leaves +# only the success's artifacts, and a post-mortem then reads the wrong run. +EVIDENCE_FILES = ( + "status.json", + "swe_bench_results.json", + "preds.json", + "swe_bench_service_status.json", +) +EVIDENCE_LOG_TAIL_BYTES = 200_000 +EVIDENCE_LOGS = ("swe_bench_agent.log", "swe_bench_eval.log") + + +class ClaimError(RuntimeError): + """A claim operation could not be performed.""" + + +class UnitOutcome(StrEnum): + """How an attempt at a unit ended. + + ``ENV_FAULT`` is deliberately separate from ``FAILED``: a broken service, an + unreachable endpoint or a refused gate is a property of the *worker*, not of + the unit. Charging it to the unit's attempt budget abandons perfectly good + units because they happened to land on a sick host. + """ + + SUCCEEDED = "succeeded" + INFRA = "infra" + FAILED = "failed" + ENV_FAULT = "env_fault" + + +#: Outcomes that consume one of the unit's ``max_attempts``. +COUNTED_OUTCOMES = frozenset({UnitOutcome.INFRA, UnitOutcome.FAILED}) + + +@dataclass(slots=True) +class UnitResult: + """A terminal or attempt record for one unit.""" + + unit_id: str + run_id: str + plan_digest: str + outcome: UnitOutcome + accounted_instance_ids: tuple[str, ...] = () + resolved_instance_ids: tuple[str, ...] = () + infra_error_count: int = 0 + genuine_error_count: int = 0 + error_kinds: dict[str, int] = field(default_factory=dict) + service_url: str | None = None + endpoint_fingerprint: str | None = None + service_run_id: str | None = None + attempt: int = 0 + abandoned: bool = False + duration_s: float = 0.0 + detail: str | None = None + finished_at: float = field(default_factory=time.time) + + def to_dict(self) -> dict[str, Any]: + return { + "unit_id": self.unit_id, + "run_id": self.run_id, + "plan_digest": self.plan_digest, + "outcome": self.outcome.value, + "accounted_instance_ids": list(self.accounted_instance_ids), + "resolved_instance_ids": list(self.resolved_instance_ids), + "infra_error_count": self.infra_error_count, + "genuine_error_count": self.genuine_error_count, + "error_kinds": dict(self.error_kinds), + "service_url": self.service_url, + "endpoint_fingerprint": self.endpoint_fingerprint, + "service_run_id": self.service_run_id, + "attempt": self.attempt, + "abandoned": self.abandoned, + "duration_s": self.duration_s, + "detail": self.detail, + "finished_at": self.finished_at, + } + + @classmethod + def from_dict(cls, raw: dict[str, Any]) -> UnitResult: + return cls( + unit_id=str(raw["unit_id"]), + run_id=str(raw["run_id"]), + plan_digest=str(raw["plan_digest"]), + outcome=UnitOutcome(str(raw["outcome"])), + accounted_instance_ids=tuple( + str(x) for x in raw.get("accounted_instance_ids") or () + ), + resolved_instance_ids=tuple( + str(x) for x in raw.get("resolved_instance_ids") or () + ), + infra_error_count=int(raw.get("infra_error_count") or 0), + genuine_error_count=int(raw.get("genuine_error_count") or 0), + error_kinds=dict(raw.get("error_kinds") or {}), + service_url=raw.get("service_url"), + endpoint_fingerprint=raw.get("endpoint_fingerprint"), + service_run_id=raw.get("service_run_id"), + attempt=int(raw.get("attempt") or 0), + abandoned=bool(raw.get("abandoned")), + duration_s=float(raw.get("duration_s") or 0.0), + detail=raw.get("detail"), + finished_at=float(raw.get("finished_at") or 0.0), + ) + + +@dataclass(slots=True) +class OwnerRecord: + unit_id: str + host: str + pid: int + boot_id: str + plan_digest: str + claimed_at: float + endpoint_fingerprint: str | None = None + slurm_job_id: str | None = None + slurm_step_id: str | None = None + + def to_dict(self) -> dict[str, Any]: + return { + "unit_id": self.unit_id, + "host": self.host, + "pid": self.pid, + "boot_id": self.boot_id, + "plan_digest": self.plan_digest, + "claimed_at": self.claimed_at, + "endpoint_fingerprint": self.endpoint_fingerprint, + "slurm_job_id": self.slurm_job_id, + "slurm_step_id": self.slurm_step_id, + } + + @classmethod + def from_dict(cls, raw: dict[str, Any]) -> OwnerRecord: + return cls( + unit_id=str(raw["unit_id"]), + host=str(raw.get("host") or ""), + pid=int(raw.get("pid") or 0), + boot_id=str(raw.get("boot_id") or ""), + plan_digest=str(raw.get("plan_digest") or ""), + claimed_at=float(raw.get("claimed_at") or 0.0), + endpoint_fingerprint=raw.get("endpoint_fingerprint"), + slurm_job_id=raw.get("slurm_job_id"), + slurm_step_id=raw.get("slurm_step_id"), + ) + + +def boot_id() -> str: + """Identify this boot of this host. + + A pid alone is not proof of liveness: after a reboot the same pid can belong + to something else entirely, and the reaper would then conclude a dead owner + is alive and leave its unit blocked forever. + """ + try: + return Path("/proc/sys/kernel/random/boot_id").read_text().strip() + except OSError: + try: + return str(int(time.time() - time.monotonic())) + except (OSError, ValueError): # pragma: no cover - defensive + return "unknown" + + +def _atomic_write_json(path: Path, payload: dict[str, Any]) -> None: + tmp = path.with_name(f".{path.name}.{os.getpid()}.tmp") + tmp.write_bytes(msgspec.json.encode(payload)) + tmp.replace(path) + + +class WorkQueue: + """Filesystem-backed queue over an immutable :class:`UnitPlan`.""" + + def __init__(self, root: os.PathLike[str] | str, plan: UnitPlan) -> None: + self.root = Path(root) + self.plan = plan + self.claims_dir = self.root / "claims" + self.results_dir = self.root / "results" + self.failed_dir = self.root / "failed" + self.env_failed_dir = self.failed_dir / "env" + self.artifacts_dir = self.failed_dir / "artifacts" + for directory in ( + self.root, + self.claims_dir, + self.results_dir, + self.failed_dir, + self.env_failed_dir, + self.artifacts_dir, + ): + directory.mkdir(parents=True, exist_ok=True) + self.plan.write(self.root) + self._boot_id = boot_id() + + # ------------------------------------------------------------------ open -- + + @classmethod + def open(cls, root: os.PathLike[str] | str) -> WorkQueue: + """Reopen an existing queue, reading its plan from disk.""" + root_path = Path(root) + return cls(root_path, read_plan(root_path / PLAN_FILENAME)) + + # ----------------------------------------------------------- inspection -- + + def claimed_unit_ids(self) -> set[str]: + try: + return {entry.name for entry in self.claims_dir.iterdir() if entry.is_dir()} + except FileNotFoundError: # pragma: no cover - created in __init__ + return set() + + def completed_unit_ids(self) -> set[str]: + return {path.stem for path in self.results_dir.glob("*.json")} + + def available_unit_ids(self) -> list[str]: + """Units that are neither claimed nor terminal, in plan order. + + Subtracting *both* claims and results is what makes a hand-deleted + result file a no-op: the claim tombstone still hides the unit. Use + :meth:`requeue`. + """ + taken = self.claimed_unit_ids() | self.completed_unit_ids() + return [unit_id for unit_id in self.plan.unit_ids if unit_id not in taken] + + def attempts(self, unit_id: str) -> int: + """Number of *counted* attempts recorded for a unit.""" + return len(list(self.failed_dir.glob(f"{unit_id}.*.json"))) + + def owner(self, unit_id: str) -> OwnerRecord | None: + path = self.claims_dir / unit_id / _OWNER + try: + raw = msgspec.json.decode(path.read_bytes(), type=dict) + except (OSError, msgspec.DecodeError): + return None + try: + return OwnerRecord.from_dict(raw) + except (KeyError, TypeError, ValueError): + return None + + def heartbeat_age(self, unit_id: str, *, now: float | None = None) -> float | None: + path = self.claims_dir / unit_id / _HEARTBEAT + try: + mtime = path.stat().st_mtime + except OSError: + return None + return (time.time() if now is None else now) - mtime + + def result(self, unit_id: str) -> UnitResult | None: + path = self.results_dir / f"{unit_id}.json" + try: + raw = msgspec.json.decode(path.read_bytes(), type=dict) + except (OSError, msgspec.DecodeError): + return None + try: + return UnitResult.from_dict(raw) + except (KeyError, TypeError, ValueError): + return None + + def results(self) -> dict[str, UnitResult]: + found: dict[str, UnitResult] = {} + for path in sorted(self.results_dir.glob("*.json")): + result = self.result(path.stem) + if result is not None: + found[path.stem] = result + return found + + # ---------------------------------------------------------------- claim -- + + def claim( + self, + unit_id: str, + *, + endpoint_fingerprint: str | None = None, + ) -> OwnerRecord | None: + """Take exclusive ownership of ``unit_id``; ``None`` if someone else has it.""" + if unit_id not in self.plan.unit_ids: + raise ClaimError( + f"{unit_id!r} is not in the plan for run {self.plan.run_id}" + ) + claim_dir = self.claims_dir / unit_id + try: + # THE RACE IS DECIDED HERE. mkdir, never makedirs/exist_ok. + os.mkdir(claim_dir) + except FileExistsError: + return None + except OSError as exc: + if exc.errno == errno.EEXIST: # pragma: no cover - platform variance + return None + raise ClaimError(f"could not claim {unit_id}: {exc}") from exc + + record = OwnerRecord( + unit_id=unit_id, + host=socket.gethostname(), + pid=os.getpid(), + boot_id=self._boot_id, + plan_digest=self.plan.digest, + claimed_at=time.time(), + endpoint_fingerprint=endpoint_fingerprint, + slurm_job_id=os.environ.get("SLURM_JOB_ID") or None, + slurm_step_id=os.environ.get("SLURM_STEP_ID") or None, + ) + # Sole owner from here, but still temp+rename so the reaper never reads + # a half-written owner record and calls it malformed. + _atomic_write_json(claim_dir / _OWNER, record.to_dict()) + (claim_dir / _HEARTBEAT).touch() + return record + + def beat(self, unit_id: str) -> None: + path = self.claims_dir / unit_id / _HEARTBEAT + try: + path.touch() + except OSError: + logger.debug("could not refresh heartbeat for %s", unit_id, exc_info=True) + + def release(self, unit_id: str) -> bool: + """Hand a claimed unit back to the queue. + + Removes the claim *directory*. Removing only the ``owner`` file leaves an + ownerless directory, which still hides the unit and merely relabels the + problem. + """ + claim_dir = self.claims_dir / unit_id + if not claim_dir.exists(): + return False + shutil.rmtree(claim_dir, ignore_errors=True) + return True + + def is_pure_bookkeeping(self, unit_id: str) -> bool: + """True when a claim directory holds only ``owner``/``hb``.""" + claim_dir = self.claims_dir / unit_id + try: + contents = {entry.name for entry in claim_dir.iterdir()} + except OSError: + return False + return not (contents - _CLAIM_CONTENTS) + + # -------------------------------------------------------------- publish -- + + def publish(self, result: UnitResult) -> None: + """Record a terminal result and release the claim. + + Releasing here is not optional. The abandon path once published a result + but kept the claim directory, so ``claims/`` and ``results/`` disagreed + for the rest of the campaign and every reaper pass had a phantom to + reason about. Releasing is safe because :meth:`available_unit_ids` + subtracts results as well as claims. + """ + self._check_digest(result) + _atomic_write_json( + self.results_dir / f"{result.unit_id}.json", result.to_dict() + ) + self.release(result.unit_id) + + def record_attempt(self, result: UnitResult) -> int: + """Record a non-terminal attempt. Returns the counted-attempt total. + + ``ENV_FAULT`` attempts are written to ``failed/env/`` and do not + increment the counter. + """ + self._check_digest(result) + if result.outcome is UnitOutcome.ENV_FAULT: + path = self.env_failed_dir / f"{result.unit_id}.{time.time_ns()}.json" + _atomic_write_json(path, result.to_dict()) + return self.attempts(result.unit_id) + count = self.attempts(result.unit_id) + 1 + result.attempt = count + _atomic_write_json( + self.failed_dir / f"{result.unit_id}.{count}.json", result.to_dict() + ) + return count + + def snapshot_evidence(self, unit_id: str, source_dir: Path, attempt: int) -> Path: + """Copy the small files that explain a failure before a retry overwrites them.""" + target = self.artifacts_dir / f"{unit_id}.attempt{attempt}" + target.mkdir(parents=True, exist_ok=True) + for name in EVIDENCE_FILES: + candidate = source_dir / name + if candidate.is_file(): + try: + shutil.copy2(candidate, target / name) + except OSError: + logger.debug("could not snapshot %s", candidate, exc_info=True) + for name in EVIDENCE_LOGS: + candidate = source_dir / name + if not candidate.is_file(): + continue + try: + with candidate.open("rb") as handle: + handle.seek(0, os.SEEK_END) + size = handle.tell() + handle.seek(max(0, size - EVIDENCE_LOG_TAIL_BYTES)) + (target / f"{name}.tail").write_bytes(handle.read()) + except OSError: + logger.debug("could not snapshot %s", candidate, exc_info=True) + return target + + def abandon(self, result: UnitResult) -> None: + """Publish a terminal, explicitly-abandoned result.""" + result.abandoned = True + self.publish(result) + + # -------------------------------------------------------------- requeue -- + + def requeue(self, unit_id: str) -> dict[str, list[str]]: + """Make a unit runnable again. The *only* supported way. + + Removes, together: the terminal result, the claim tombstone, and every + counted attempt record. Removing any subset leaves the unit invisible or + already out of attempts, which is how "I deleted the result, why is it + not rerunning?" happens. + """ + if unit_id not in self.plan.unit_ids: + raise ClaimError( + f"{unit_id!r} is not in the plan for run {self.plan.run_id}" + ) + removed: dict[str, list[str]] = {"results": [], "claims": [], "attempts": []} + result_path = self.results_dir / f"{unit_id}.json" + if result_path.exists(): + result_path.unlink() + removed["results"].append(str(result_path)) + claim_dir = self.claims_dir / unit_id + if claim_dir.exists(): + shutil.rmtree(claim_dir, ignore_errors=True) + removed["claims"].append(str(claim_dir)) + for path in sorted(self.failed_dir.glob(f"{unit_id}.*.json")): + path.unlink() + removed["attempts"].append(str(path)) + return removed + + # ---------------------------------------------------------------- utils -- + + def _check_digest(self, result: UnitResult) -> None: + if result.plan_digest != self.plan.digest: + raise ClaimError( + f"refusing to record {result.unit_id}: plan digest " + f"{result.plan_digest[:12]} does not match this queue's " + f"{self.plan.digest[:12]}" + ) + if result.run_id != self.plan.run_id: + raise ClaimError( + f"refusing to record {result.unit_id}: run id {result.run_id!r} " + f"does not match this queue's {self.plan.run_id!r}" + ) diff --git a/src/inference_endpoint/evaluation/swe_bench_distributed/units.py b/src/inference_endpoint/evaluation/swe_bench_distributed/units.py new file mode 100644 index 000000000..9a163a037 --- /dev/null +++ b/src/inference_endpoint/evaluation/swe_bench_distributed/units.py @@ -0,0 +1,187 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shard plan: the immutable binding between units and instance ids. + +The plan is content-addressed. Every unit result carries the plan digest, and +the merge gate refuses to combine results whose digest differs from the plan +being merged. That is what makes it impossible to accidentally merge results +from a different run, a different instance list, or a different ordering into +one accuracy number. +""" + +from __future__ import annotations + +import hashlib +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import msgspec + +PLAN_FILENAME = "units.json" + + +class PlanError(ValueError): + """The requested shard plan cannot be built, or a plan file is invalid.""" + + +@dataclass(frozen=True, slots=True) +class Unit: + """One dispatchable shard of a run.""" + + unit_id: str + run_id: str + shard: int + instance_ids: tuple[str, ...] + + def to_dict(self) -> dict[str, Any]: + return { + "unit_id": self.unit_id, + "run_id": self.run_id, + "shard": self.shard, + "instance_ids": list(self.instance_ids), + } + + +@dataclass(frozen=True, slots=True) +class UnitPlan: + """The full, immutable set of units for one run id.""" + + run_id: str + shard_size: int + digest: str + units: tuple[Unit, ...] + + @property + def instance_ids(self) -> tuple[str, ...]: + return tuple( + instance_id for unit in self.units for instance_id in unit.instance_ids + ) + + def unit(self, unit_id: str) -> Unit: + for candidate in self.units: + if candidate.unit_id == unit_id: + return candidate + raise KeyError(unit_id) + + @property + def unit_ids(self) -> tuple[str, ...]: + return tuple(unit.unit_id for unit in self.units) + + def to_dict(self) -> dict[str, Any]: + return { + "run_id": self.run_id, + "shard_size": self.shard_size, + "digest": self.digest, + "units": [unit.to_dict() for unit in self.units], + } + + def write(self, directory: Path) -> Path: + """Write the plan once. Rewriting an existing, different plan is an error.""" + directory.mkdir(parents=True, exist_ok=True) + path = directory / PLAN_FILENAME + if path.exists(): + existing = read_plan(path) + if existing.digest != self.digest or existing.run_id != self.run_id: + raise PlanError( + f"refusing to overwrite plan at {path}: existing run_id=" + f"{existing.run_id!r} digest={existing.digest[:12]} differs from " + f"new run_id={self.run_id!r} digest={self.digest[:12]}" + ) + return path + tmp = path.with_name(f".{path.name}.tmp") + tmp.write_bytes(msgspec.json.encode(self.to_dict())) + tmp.replace(path) + return path + + +def plan_digest(run_id: str, instance_ids: list[str] | tuple[str, ...]) -> str: + """Digest over the run id and the ordered instance list. + + Order is included deliberately: two plans over the same ids in a different + order produce different shards, so they are different plans. + """ + hasher = hashlib.sha256() + hasher.update(run_id.encode()) + hasher.update(b"\0") + for instance_id in instance_ids: + hasher.update(instance_id.encode()) + hasher.update(b"\n") + return hasher.hexdigest() + + +def plan_units( + run_id: str, + instance_ids: list[str] | tuple[str, ...], + *, + shard_size: int = 10, +) -> UnitPlan: + """Split ``instance_ids`` into fixed-size shards, in order. + + The final shard is short when the count is not a multiple of ``shard_size``; + it is never padded and never merged into its neighbour, because the merge + gate compares id sets and a padded shard would claim ids it never ran. + """ + if not run_id or "/" in run_id or run_id in {".", ".."}: + raise PlanError(f"invalid run_id: {run_id!r}") + if shard_size < 1: + raise PlanError(f"shard_size must be >= 1; got {shard_size}") + ordered = [str(instance_id) for instance_id in instance_ids] + if not ordered: + raise PlanError("cannot plan a run with no instance ids") + duplicates = sorted({x for x in ordered if ordered.count(x) > 1}) + if duplicates: + raise PlanError( + "instance ids must be unique; duplicated: " + ", ".join(duplicates[:10]) + ) + + digest = plan_digest(run_id, ordered) + units: list[Unit] = [] + for shard, start in enumerate(range(0, len(ordered), shard_size)): + chunk = tuple(ordered[start : start + shard_size]) + units.append( + Unit( + unit_id=f"{run_id}.s{shard:02d}", + run_id=run_id, + shard=shard, + instance_ids=chunk, + ) + ) + return UnitPlan( + run_id=run_id, shard_size=shard_size, digest=digest, units=tuple(units) + ) + + +def read_plan(path: Path) -> UnitPlan: + try: + raw = msgspec.json.decode(path.read_bytes(), type=dict) + except (OSError, msgspec.DecodeError) as exc: + raise PlanError(f"could not read unit plan at {path}") from exc + try: + units = tuple( + Unit( + unit_id=str(entry["unit_id"]), + run_id=str(entry["run_id"]), + shard=int(entry["shard"]), + instance_ids=tuple(str(x) for x in entry["instance_ids"]), + ) + for entry in raw["units"] + ) + plan = UnitPlan( + run_id=str(raw["run_id"]), + shard_size=int(raw["shard_size"]), + digest=str(raw["digest"]), + units=units, + ) + except (KeyError, TypeError, ValueError) as exc: + raise PlanError(f"malformed unit plan at {path}") from exc + + recomputed = plan_digest(plan.run_id, list(plan.instance_ids)) + if recomputed != plan.digest: + raise PlanError( + f"unit plan at {path} is inconsistent: recorded digest " + f"{plan.digest[:12]} does not match its own instance list " + f"({recomputed[:12]})" + ) + return plan diff --git a/tests/unit/evaluation/swe_bench_distributed/test_units_and_queue.py b/tests/unit/evaluation/swe_bench_distributed/test_units_and_queue.py new file mode 100644 index 000000000..6343b7536 --- /dev/null +++ b/tests/unit/evaluation/swe_bench_distributed/test_units_and_queue.py @@ -0,0 +1,246 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit plan and work-queue semantics.""" + +from __future__ import annotations + +import threading + +import pytest + +from inference_endpoint.evaluation.swe_bench_distributed.queue import ( + ClaimError, + UnitOutcome, + UnitResult, + WorkQueue, +) +from inference_endpoint.evaluation.swe_bench_distributed.units import ( + PlanError, + plan_units, + read_plan, +) + +pytestmark = pytest.mark.unit + + +def make_ids(n: int) -> list[str]: + return [f"repo__proj-{i:03d}" for i in range(n)] + + +@pytest.fixture +def queue(tmp_path): + plan = plan_units("run-a", make_ids(25), shard_size=10) + return WorkQueue(tmp_path / "wq", plan) + + +def result_for(queue: WorkQueue, unit_id: str, **overrides) -> UnitResult: + unit = queue.plan.unit(unit_id) + payload = { + "unit_id": unit_id, + "run_id": unit.run_id, + "plan_digest": queue.plan.digest, + "outcome": UnitOutcome.SUCCEEDED, + "accounted_instance_ids": unit.instance_ids, + "resolved_instance_ids": unit.instance_ids[:1], + } + payload.update(overrides) + return UnitResult(**payload) + + +class TestPlan: + def test_shards_in_order_with_a_short_tail(self): + plan = plan_units("run-a", make_ids(25), shard_size=10) + assert [len(unit.instance_ids) for unit in plan.units] == [10, 10, 5] + assert plan.unit_ids == ("run-a.s00", "run-a.s01", "run-a.s02") + # The short tail is never padded: a padded shard would claim ids it + # never ran and the merge gate compares ids, not counts. + assert plan.instance_ids == tuple(make_ids(25)) + + def test_digest_depends_on_order(self): + ids = make_ids(20) + assert ( + plan_units("r", ids).digest != plan_units("r", list(reversed(ids))).digest + ) + + def test_digest_depends_on_run_id(self): + ids = make_ids(20) + assert plan_units("r1", ids).digest != plan_units("r2", ids).digest + + def test_duplicate_instance_ids_are_refused(self): + with pytest.raises(PlanError, match="unique"): + plan_units("r", ["a", "b", "a"]) + + def test_empty_plan_is_refused(self): + with pytest.raises(PlanError): + plan_units("r", []) + + def test_plan_round_trips_and_self_verifies(self, tmp_path): + plan = plan_units("run-a", make_ids(12), shard_size=5) + path = plan.write(tmp_path) + assert read_plan(path).digest == plan.digest + + def test_rewriting_a_different_plan_is_refused(self, tmp_path): + plan_units("run-a", make_ids(10)).write(tmp_path) + with pytest.raises(PlanError, match="refusing to overwrite"): + plan_units("run-a", make_ids(11)).write(tmp_path) + + def test_tampered_plan_file_is_detected(self, tmp_path): + plan = plan_units("run-a", make_ids(10)) + path = plan.write(tmp_path) + raw = path.read_text().replace(plan.digest, "0" * 64) + path.write_text(raw) + with pytest.raises(PlanError, match="inconsistent"): + read_plan(path) + + +class TestClaims: + def test_a_second_claim_loses(self, queue): + assert queue.claim("run-a.s00") is not None + assert queue.claim("run-a.s00") is None + + def test_exactly_one_thread_wins_a_contested_claim(self, queue): + winners: list[object] = [] + barrier = threading.Barrier(8) + + def contend(): + barrier.wait() + if queue.claim("run-a.s01") is not None: + winners.append(object()) + + threads = [threading.Thread(target=contend) for _ in range(8)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + assert len(winners) == 1 + + def test_claiming_a_unit_outside_the_plan_is_an_error(self, queue): + with pytest.raises(ClaimError): + queue.claim("other-run.s00") + + def test_release_removes_the_whole_directory(self, queue): + queue.claim("run-a.s00") + assert queue.release("run-a.s00") + # Removing only `owner` would leave an ownerless directory that still + # hides the unit -- a relabelled problem, not a fix. + assert not (queue.claims_dir / "run-a.s00").exists() + assert "run-a.s00" in queue.available_unit_ids() + + def test_owner_record_carries_identity(self, queue): + record = queue.claim("run-a.s00") + stored = queue.owner("run-a.s00") + assert stored is not None + assert stored.pid == record.pid + assert stored.boot_id == record.boot_id + assert stored.plan_digest == queue.plan.digest + + +class TestAvailability: + def test_claims_and_results_both_hide_a_unit(self, queue): + queue.claim("run-a.s00") + queue.publish(result_for(queue, "run-a.s01")) + assert queue.available_unit_ids() == ["run-a.s02"] + + def test_publish_releases_the_claim(self, queue): + queue.claim("run-a.s00") + queue.publish(result_for(queue, "run-a.s00")) + # An abandoned or published unit that keeps its claim makes claims/ and + # results/ disagree for the rest of the run. + assert queue.claimed_unit_ids() == set() + + def test_abandon_publishes_and_releases(self, queue): + queue.claim("run-a.s00") + queue.abandon(result_for(queue, "run-a.s00", outcome=UnitOutcome.FAILED)) + stored = queue.result("run-a.s00") + assert stored is not None and stored.abandoned + assert queue.claimed_unit_ids() == set() + + def test_result_from_another_plan_is_refused(self, queue): + bad = result_for(queue, "run-a.s00", plan_digest="0" * 64) + with pytest.raises(ClaimError, match="plan digest"): + queue.publish(bad) + + +class TestRequeue: + def test_deleting_only_the_result_does_not_requeue(self, queue): + queue.claim("run-a.s00") + queue.publish(result_for(queue, "run-a.s00")) + # Re-claim so a tombstone exists, mimicking an interrupted retry. + queue.claim("run-a.s00") + (queue.results_dir / "run-a.s00.json").unlink() + assert "run-a.s00" not in queue.available_unit_ids() + + def test_deleting_only_the_claim_does_not_requeue(self, queue): + queue.claim("run-a.s00") + queue.publish(result_for(queue, "run-a.s00")) + queue.release("run-a.s00") + assert "run-a.s00" not in queue.available_unit_ids() + + def test_requeue_removes_result_claim_and_attempts(self, queue): + queue.claim("run-a.s00") + queue.record_attempt(result_for(queue, "run-a.s00", outcome=UnitOutcome.FAILED)) + queue.record_attempt(result_for(queue, "run-a.s00", outcome=UnitOutcome.INFRA)) + queue.publish(result_for(queue, "run-a.s00")) + queue.claim("run-a.s00") + + removed = queue.requeue("run-a.s00") + + assert len(removed["results"]) == 1 + assert len(removed["claims"]) == 1 + assert len(removed["attempts"]) == 2 + assert "run-a.s00" in queue.available_unit_ids() + assert queue.attempts("run-a.s00") == 0 + + def test_requeue_outside_the_plan_is_an_error(self, queue): + with pytest.raises(ClaimError): + queue.requeue("other-run.s00") + + +class TestAttemptLedger: + def test_environment_faults_do_not_consume_the_budget(self, queue): + for _ in range(5): + queue.record_attempt( + result_for(queue, "run-a.s00", outcome=UnitOutcome.ENV_FAULT) + ) + # A broken host is a property of the host, not of the unit. Charging it + # to the unit abandons good units for landing in the wrong place. + assert queue.attempts("run-a.s00") == 0 + + def test_counted_failures_increment(self, queue): + assert ( + queue.record_attempt( + result_for(queue, "run-a.s00", outcome=UnitOutcome.FAILED) + ) + == 1 + ) + assert ( + queue.record_attempt( + result_for(queue, "run-a.s00", outcome=UnitOutcome.INFRA) + ) + == 2 + ) + + def test_evidence_is_snapshotted_before_a_retry_overwrites_it( + self, queue, tmp_path + ): + source = tmp_path / "unit-run" + source.mkdir() + (source / "status.json").write_text('{"attempt": 1}') + (source / "swe_bench_agent.log").write_text("first attempt log") + + target = queue.snapshot_evidence("run-a.s00", source, attempt=1) + + # The retry reuses the run directory, so a unit that fails then succeeds + # would otherwise leave only the success's artifacts behind. + (source / "status.json").write_text('{"attempt": 2}') + assert (target / "status.json").read_text() == '{"attempt": 1}' + assert (target / "swe_bench_agent.log.tail").read_text() == "first attempt log" + + +class TestReopen: + def test_reopen_reads_the_plan_from_disk(self, queue): + queue.publish(result_for(queue, "run-a.s00")) + reopened = WorkQueue.open(queue.root) + assert reopened.plan.digest == queue.plan.digest + assert reopened.completed_unit_ids() == {"run-a.s00"}