Resilient worker (transient CaptureContractError -> unscorable, systematic still fatal), resume-from-repo, eval worker base-cells download fix

eval_watch.py

@@ -75,6 +75,10 @@ def main():
         print("[watch] evaluating base model first", flush=True)
         run_eval(BASE_MODEL, "/tmp/base_cells.json")
         upload("/tmp/base_cells.json", "eval/base_cells.json")
+    elif not Path("/tmp/base_cells.json").exists():
+        # A previous worker already computed base cells: reuse them locally.
+        local = api.hf_hub_download(REPO, "eval/base_cells.json")
+        Path("/tmp/base_cells.json").write_text(Path(local).read_text())
 
     summary = {}
     if "eval/summary.json" in fs:

run_job.sh

@@ -126,11 +126,39 @@ for name, url in [("vllm", "http://127.0.0.1:8001/v1/models"),
         raise SystemExit(f"{name} server never became ready")
 EOF
 
+echo "=== [resume] latest pushed checkpoint, if any ==="
+# Optimizer shards are not uploaded, so resume restores the policy, scheduler,
+# task position (rollout_state.json) and RNG of the sampler — with a fresh
+# optimizer state. Close enough for RL.
+RESUME_FROM=""
+if [ -n "${ADAPTER_REPO:-}" ]; then
+  RESUME_FROM=$(python3 - <<'EOF'
+import os, re
+from huggingface_hub import HfApi, snapshot_download
+
+repo = os.environ.get("ADAPTER_REPO", "")
+try:
+    files = HfApi().list_repo_files(repo)
+except Exception:
+    files = []
+ckpts = sorted(
+    {f.split("/")[0] for f in files if f.startswith("checkpoint-") and f.endswith("adapter_config.json")},
+    key=lambda s: int(s.rsplit("-", 1)[1]),
+)
+if ckpts:
+    path = snapshot_download(repo, allow_patterns=[f"{ckpts[-1]}/*"])
+    print(os.path.join(path, ckpts[-1]))
+EOF
+)
+fi
+echo "RESUME_FROM=${RESUME_FROM:-<fresh start>}"
+
 echo "=== [train] GPU 0, MAX_STEPS=${MAX_STEPS:-1000} ==="
 CUDA_VISIBLE_DEVICES=0 \
 VLLM_URL=http://127.0.0.1:8001 \
 ENV_URL=http://127.0.0.1:8000 \
 TRAIN_SPLIT=/app/prepared/datasets/train \
+RESUME_FROM="$RESUME_FROM" \
     python3 train_lora.py
 
 if [ "${SKIP_EVAL:-0}" != "1" ] && [ -n "${ADAPTER_REPO:-}" ]; then

train_lora.py

@@ -27,7 +27,11 @@ from peft import LoraConfig
 from transformers import AutoTokenizer, TrainerCallback
 from trl.experimental.async_grpo.async_grpo_config import AsyncGRPOConfig
 from trl.experimental.async_grpo.async_grpo_trainer import AsyncGRPOTrainer
-from trl.experimental.async_grpo.openenv_harness import HarnessRolloutWorker
+from trl.experimental.async_grpo.openenv_harness import (
+    CaptureContractError,
+    HarnessRolloutWorker,
+    _HarnessRolloutLoop,
+)
 
 from harbor_env.harness import HarborSessionFactory  # via .pth to OpenEnv/envs
 
@@ -84,6 +88,43 @@ class UploadCheckpointsCallback(TrainerCallback):
         return control
 
 
+class ResilientHarnessRolloutLoop(_HarnessRolloutLoop):
+    """Score a contract-violating capture as unscorable instead of killing the worker.
+
+    TRL makes CaptureContractError fatal: never train on a corrupted capture.
+    Right in principle — but in practice the corruption source is transient infra
+    (a dropped websocket between worker and env server), and one bad capture used
+    to kill a 25h run at step 551. Dropping the rollout keeps bad data out of
+    training just the same. To keep the safety property for SYSTEMATIC corruption
+    (e.g. every capture broken, as in the 401 incident), re-raise after a full
+    two groups' worth of consecutive contract errors.
+    """
+
+    MAX_CONSECUTIVE_CONTRACT_ERRORS = 16
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._contract_errors = 0
+
+    def _run_session(self, prompt, group_id=0):
+        import logging
+
+        try:
+            result = super()._run_session(prompt, group_id)
+        except CaptureContractError as exc:
+            self._contract_errors += 1
+            if self._contract_errors >= self.MAX_CONSECUTIVE_CONTRACT_ERRORS:
+                raise
+            logging.warning("capture contract violation; scoring unscorable (%s)", exc)
+            return self._EMPTY_ROLLOUT, None
+        self._contract_errors = 0
+        return result
+
+
+class ResilientHarnessRolloutWorker(HarnessRolloutWorker):
+    _loop_cls = ResilientHarnessRolloutLoop
+
+
 class HeartbeatCallback(TrainerCallback):
     """Touch a file at every optimizer step; run_job.sh's watchdog kills the
     job if it goes stale. The first full run stalled silently for 13h when the
@@ -124,7 +165,7 @@ def make_factory():
 
 
 def make_rollout_worker(dataset, tokenizer):
-    return HarnessRolloutWorker(
+    return ResilientHarnessRolloutWorker(
         harness_session_factory=make_factory(),
         harness_adapter=None,     # let the harness own the loop (black-box capture)
         rollout_reward_fn=correctness_plus_efficiency,
@@ -205,7 +246,10 @@ def main():
         peft_config=peft_config,
         callbacks=[UploadCheckpointsCallback(ADAPTER_REPO), HeartbeatCallback()],
     )
-    trainer.train()
+    resume = os.environ.get("RESUME_FROM", "") or None
+    if resume:
+        print(f"resuming from {resume}")
+    trainer.train(resume_from_checkpoint=resume)
     trainer.save_model("./checkpoints/final")   # adapter_model.safetensors + adapter_config.json
     if ADAPTER_REPO:
         trainer.push_to_hub()