Make EasyCom adjudication resumable

adjudicate_easycom_eot_int_candidates.py

@@ -41,6 +41,16 @@ HARD_NEGATIVE_TYPES = {
     "INT_WEARER_HOLDS_FLOOR",
     "OTHER_HARD_NEGATIVE",
 }
+
+
+class BatchValidationError(RuntimeError):
+    """The provider responded, but the batch output violated the label schema."""
+
+
+class APIRequestError(RuntimeError):
+    """The provider could not return a usable response after transport retries."""
+
+
 def _progress_line(
     completed: int, total: int, *, label: str, elapsed_seconds: float, width: int = 26
 ) -> str:
@@ -274,7 +284,9 @@ def validate_decisions(
         if verdict != "POSITIVE" and event_type is not None:
             raise ValueError(f"Only POSITIVE may have event_type for {uid}")
         if verdict == "HARD_NEGATIVE" and subtype not in HARD_NEGATIVE_TYPES:
-            raise ValueError(f"HARD_NEGATIVE requires a valid subtype for {uid}")
+            raise ValueError(
+                f"HARD_NEGATIVE requires a valid subtype for {uid}; got {subtype!r}"
+            )
         if verdict != "HARD_NEGATIVE" and subtype is not None:
             raise ValueError(f"Only HARD_NEGATIVE may have subtype for {uid}")
         if verdict == "EXCLUDE" and exclude_reason not in EXCLUDE_REASONS:
@@ -357,6 +369,19 @@ def _post(
         try:
             with urllib.request.urlopen(request, timeout=timeout_seconds) as response:
                 envelope = json.loads(response.read().decode("utf-8"))
+        except (
+            urllib.error.HTTPError,
+            urllib.error.URLError,
+            TimeoutError,
+        ) as error:
+            last_error = error
+            if attempt >= max_retries:
+                raise APIRequestError(
+                    f"API request failed after {max_retries + 1} attempts: {last_error}"
+                ) from error
+            time.sleep(min(30.0, 2.0**attempt))
+            continue
+        try:
             content = envelope["choices"][0]["message"]["content"]
             parsed = json.loads(content)
             allowed_opportunities = {}
@@ -383,19 +408,93 @@ def _post(
                 decision["uid"] = original_by_transport[decision["uid"]]
             return decisions, envelope.get("usage", {})
         except (
-            urllib.error.HTTPError,
-            urllib.error.URLError,
-            TimeoutError,
             KeyError,
             TypeError,
             json.JSONDecodeError,
             ValueError,
         ) as error:
             last_error = error
+            # A malformed multi-item answer is immediately bisected. Retrying
+            # the identical large prompt is both expensive and unlikely to
+            # identify which item caused the schema violation. A singleton is
+            # retried before it is placed in the human queue.
+            if len(rows) > 1:
+                raise BatchValidationError(str(error)) from error
             if attempt >= max_retries:
-                break
+                raise BatchValidationError(
+                    f"singleton validation failed after {max_retries + 1} attempts: {error}"
+                ) from error
             time.sleep(min(30.0, 2.0**attempt))
-    raise RuntimeError(f"DeepSeek request failed after {max_retries + 1} attempts: {last_error}")
+    raise RuntimeError(f"Unreachable request state: {last_error}")
+
+
+def _atomic_write_json(path: Path, payload: dict[str, Any]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp")
+    try:
+        temporary.write_text(
+            json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
+        )
+        os.replace(temporary, path)
+    finally:
+        if temporary.exists():
+            temporary.unlink()
+
+
+def _atomic_write_jsonl(path: Path, rows: Sequence[dict[str, Any]]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp")
+    try:
+        with temporary.open("w", encoding="utf-8") as handle:
+            for row in rows:
+                handle.write(json.dumps(row, ensure_ascii=False) + "\n")
+        os.replace(temporary, path)
+    finally:
+        if temporary.exists():
+            temporary.unlink()
+
+
+def _run_fingerprint(
+    rows: Sequence[dict[str, Any]], *, model: str, base_url: str, passes: int, batch_size: int
+) -> str:
+    payload = {
+        "prompt_version": PROMPT_VERSION,
+        "model": model,
+        "base_url": base_url.rstrip("/"),
+        "passes": passes,
+        "batch_size": batch_size,
+        "candidates": [candidate_for_prompt(row) for row in rows],
+    }
+    return hashlib.sha256(
+        json.dumps(payload, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode(
+            "utf-8"
+        )
+    ).hexdigest()
+
+
+def _initial_checkpoint(
+    rows: Sequence[dict[str, Any]], *, model: str, base_url: str, passes: int,
+    batch_size: int,
+) -> dict[str, Any]:
+    return {
+        "schema_version": 1,
+        "status": "running",
+        "run_fingerprint": _run_fingerprint(
+            rows, model=model, base_url=base_url, passes=passes, batch_size=batch_size
+        ),
+        "config": {
+            "prompt_version": PROMPT_VERSION,
+            "model": model,
+            "base_url": base_url.rstrip("/"),
+            "passes": passes,
+            "batch_size": batch_size,
+        },
+        "candidate_uids": [row["uid"] for row in rows],
+        "decisions": {},
+        "prompt_hashes": {},
+        "pending": {},
+        "usage": {},
+    }
 
 
 def add_suggestions(
@@ -408,57 +507,151 @@ def add_suggestions(
     batch_size: int = 8,
     timeout_seconds: float,
     max_retries: int,
-) -> tuple[list[dict[str, Any]], Counter[str]]:
-    by_uid: dict[str, list[dict[str, Any]]] = {row["uid"]: [] for row in rows}
-    prompt_hashes: dict[str, list[str]] = {row["uid"]: [] for row in rows}
-    usage: Counter[str] = Counter()
+    checkpoint_path: Path | None = None,
+    pending_output_path: Path | None = None,
+    resume: bool = True,
+) -> tuple[list[dict[str, Any]], list[dict[str, Any]], Counter[str]]:
     if batch_size <= 0:
         raise ValueError("batch_size must be positive")
+    expected_state = _initial_checkpoint(
+        rows, model=model, base_url=base_url, passes=passes, batch_size=batch_size
+    )
+    if checkpoint_path is not None and resume and checkpoint_path.exists():
+        state = json.loads(checkpoint_path.read_text(encoding="utf-8"))
+        if state.get("run_fingerprint") != expected_state["run_fingerprint"]:
+            raise ValueError(
+                f"Checkpoint does not match this input/model/prompt: {checkpoint_path}. "
+                "Pass --overwrite to start a new run."
+            )
+    else:
+        state = expected_state
+        if checkpoint_path is not None:
+            _atomic_write_json(checkpoint_path, state)
+
+    decisions_by_uid: dict[str, dict[str, dict[str, Any]]] = state["decisions"]
+    hashes_by_uid: dict[str, dict[str, str]] = state["prompt_hashes"]
+    pending_by_uid: dict[str, dict[str, Any]] = state["pending"]
+    usage: Counter[str] = Counter(state.get("usage", {}))
+    source_by_uid = {row["uid"]: row for row in rows}
+
+    def materialize_pending() -> list[dict[str, Any]]:
+        result: list[dict[str, Any]] = []
+        for uid in state["candidate_uids"]:
+            if uid not in pending_by_uid:
+                continue
+            pending = dict(source_by_uid[uid])
+            pending.update(
+                human_status="UNREVIEWED",
+                human_verdict=None,
+                human_event_type=None,
+                human_subtype=None,
+                exclude_reason=None,
+            )
+            pending["llm_adjudication_failure"] = pending_by_uid[uid]
+            result.append(pending)
+        return result
+
+    def save_checkpoint() -> None:
+        state["usage"] = dict(usage)
+        if checkpoint_path is not None:
+            _atomic_write_json(checkpoint_path, state)
+        if pending_output_path is not None:
+            _atomic_write_jsonl(pending_output_path, materialize_pending())
+
     total_work = len(rows) * passes
-    batches_per_pass = (len(rows) + batch_size - 1) // batch_size
-    completed_work = 0
+    completed_work = len(pending_by_uid) * passes + sum(
+        len(pass_decisions)
+        for uid, pass_decisions in decisions_by_uid.items()
+        if uid not in pending_by_uid
+    )
     progress = _ProgressBar(total_work)
-    progress.update(0, f"pass 1/{passes} · waiting")
+    if completed_work:
+        print(
+            f"Resuming {checkpoint_path}: completed={completed_work}/{total_work} "
+            f"pending_human={len(pending_by_uid)}"
+        )
+    progress.update(completed_work, f"pass 1/{passes} · waiting")
+
+    def process_batch(batch: Sequence[dict[str, Any]], pass_index: int) -> None:
+        nonlocal completed_work
+        if not batch:
+            return
+        transported, _ = _transport_rows(batch)
+        prompt_hash = hashlib.sha256(
+            (SYSTEM_PROMPT + "\n" + build_user_prompt(transported, pass_index)).encode("utf-8")
+        ).hexdigest()
+        progress.update(completed_work, f"pass {pass_index + 1}/{passes} · request {len(batch)}")
+        try:
+            decisions, pass_usage = _post(
+                rows=batch,
+                pass_index=pass_index,
+                api_key=api_key,
+                model=model,
+                base_url=base_url,
+                timeout_seconds=timeout_seconds,
+                max_retries=max_retries,
+            )
+        except BatchValidationError as error:
+            if len(batch) > 1:
+                midpoint = len(batch) // 2
+                progress.update(
+                    completed_work,
+                    f"validation split {len(batch)}→{midpoint}+{len(batch) - midpoint}",
+                )
+                process_batch(batch[:midpoint], pass_index)
+                process_batch(batch[midpoint:], pass_index)
+                return
+            uid = batch[0]["uid"]
+            pending_by_uid[uid] = {
+                "failed_pass": pass_index + 1,
+                "error_type": type(error).__name__,
+                "error": str(error),
+                "model": model,
+                "base_url": base_url.rstrip("/"),
+                "prompt_version": PROMPT_VERSION,
+                "request_prompt_sha256": prompt_hash,
+            }
+            completed_work += passes - pass_index
+            save_checkpoint()
+            progress.update(completed_work, f"queued for human review: {uid}")
+            return
+
+        pass_key = str(pass_index)
+        for decision in decisions:
+            uid = decision["uid"]
+            decisions_by_uid.setdefault(uid, {})[pass_key] = decision
+            hashes_by_uid.setdefault(uid, {})[pass_key] = prompt_hash
+        for key, value in pass_usage.items():
+            if isinstance(value, int):
+                usage[key] += value
+        completed_work += len(batch)
+        save_checkpoint()
+        progress.update(completed_work, f"pass {pass_index + 1}/{passes} · saved {len(batch)}")
+
     try:
         for pass_index in range(passes):
-            for batch_index, start in enumerate(range(0, len(rows), batch_size), 1):
-                batch = rows[start:start + batch_size]
-                progress.update(
-                    completed_work,
-                    f"pass {pass_index + 1}/{passes} · batch {batch_index}/{batches_per_pass}",
-                )
-                transported, _ = _transport_rows(batch)
-                prompt_hash = hashlib.sha256(
-                    (SYSTEM_PROMPT + "\n" + build_user_prompt(transported, pass_index)).encode("utf-8")
-                ).hexdigest()
-                decisions, pass_usage = _post(
-                    rows=batch,
-                    pass_index=pass_index,
-                    api_key=api_key,
-                    model=model,
-                    base_url=base_url,
-                    timeout_seconds=timeout_seconds,
-                    max_retries=max_retries,
-                )
-                for decision in decisions:
-                    by_uid[decision["uid"]].append(decision)
-                    prompt_hashes[decision["uid"]].append(prompt_hash)
-                for key, value in pass_usage.items():
-                    if isinstance(value, int):
-                        usage[key] += value
-                completed_work += len(batch)
-                progress.update(
-                    completed_work,
-                    f"pass {pass_index + 1}/{passes} · batch {batch_index}/{batches_per_pass}",
-                )
+            eligible = [
+                row for row in rows
+                if row["uid"] not in pending_by_uid
+                and str(pass_index) not in decisions_by_uid.get(row["uid"], {})
+            ]
+            for start in range(0, len(eligible), batch_size):
+                process_batch(eligible[start:start + batch_size], pass_index)
     except Exception:
-        progress.update(completed_work, "failed", final=True)
+        progress.update(completed_work, "failed; checkpoint retained", final=True)
         raise
     progress.update(completed_work, "complete", final=True)
 
     output: list[dict[str, Any]] = []
     for source in rows:
-        decisions = by_uid[source["uid"]]
+        uid = source["uid"]
+        if uid in pending_by_uid:
+            continue
+        pass_decisions = decisions_by_uid.get(uid, {})
+        missing_passes = [index + 1 for index in range(passes) if str(index) not in pass_decisions]
+        if missing_passes:
+            raise RuntimeError(f"Internal checkpoint error: {uid} missing passes {missing_passes}")
+        decisions = [pass_decisions[str(index)] for index in range(passes)]
         labels = {
             (d["verdict"], d["event_type"], d["subtype"], d["exclude_reason"])
             for d in decisions
@@ -531,14 +724,20 @@ def add_suggestions(
         )
         enriched["llm_adjudication"] = {
             "model": model,
+            "base_url": base_url.rstrip("/"),
             "prompt_version": PROMPT_VERSION,
             "passes": passes,
-            "request_prompt_sha256": prompt_hashes[source["uid"]],
+            "request_prompt_sha256": [
+                hashes_by_uid[uid][str(index)] for index in range(passes)
+            ],
             "decisions": decisions,
             "human_review_required": True,
         }
         output.append(enriched)
-    return output, usage
+    state["status"] = "complete"
+    save_checkpoint()
+    pending_output = materialize_pending()
+    return output, pending_output, usage
 
 
 def parse_args() -> argparse.Namespace:
@@ -553,6 +752,16 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--selection-seed", type=int, default=2027)
     parser.add_argument("--timeout-seconds", type=float, default=120.0)
     parser.add_argument("--max-retries", type=int, default=3)
+    parser.add_argument(
+        "--checkpoint",
+        type=Path,
+        help="Resume state (default: OUTPUT.checkpoint.json).",
+    )
+    parser.add_argument(
+        "--pending-output",
+        type=Path,
+        help="Singleton validation failures (default: OUTPUT stem + _pending.jsonl).",
+    )
     parser.add_argument("--dry-run", action="store_true")
     parser.add_argument("--allow-external-transcript-upload", action="store_true")
     parser.add_argument("--overwrite", action="store_true")
@@ -591,7 +800,17 @@ def main() -> None:
     output = args.output.expanduser().resolve()
     if output.exists() and not args.overwrite:
         raise FileExistsError(f"Refusing to overwrite {output}; pass --overwrite")
-    enriched, usage = add_suggestions(
+    checkpoint = (
+        args.checkpoint.expanduser().resolve()
+        if args.checkpoint is not None
+        else output.with_suffix(output.suffix + ".checkpoint.json")
+    )
+    pending_output_path = (
+        args.pending_output.expanduser().resolve()
+        if args.pending_output is not None
+        else output.with_name(output.stem + "_pending.jsonl")
+    )
+    enriched, pending, usage = add_suggestions(
         rows,
         api_key=api_key,
         model=args.model,
@@ -600,12 +819,15 @@ def main() -> None:
         batch_size=args.batch_size,
         timeout_seconds=args.timeout_seconds,
         max_retries=args.max_retries,
+        checkpoint_path=checkpoint,
+        pending_output_path=pending_output_path,
+        resume=not args.overwrite,
     )
-    output.parent.mkdir(parents=True, exist_ok=True)
-    with output.open("w", encoding="utf-8") as handle:
-        for row in enriched:
-            handle.write(json.dumps(row, ensure_ascii=False) + "\n")
+    _atomic_write_jsonl(output, enriched)
+    _atomic_write_jsonl(pending_output_path, pending)
     print(f"Saved {len(enriched)} suggested candidates to {output}")
+    print(f"Saved {len(pending)} pending-human candidates to {pending_output_path}")
+    print(f"Checkpoint: {checkpoint}")
     print("Usage:", dict(usage))