Resume TerminalBench Harbor runs from checkpoints (#768)

Checkpoint TerminalBench Harbor trials across retries

Co-authored-by: neubig <neubig@users.noreply.github.com>
Co-authored-by: openhands <openhands@all-hands.dev>

Author: neubig

benchmarks/terminalbench/run_infer.py

@@ -13,6 +13,8 @@
 import os
 import subprocess
 import sys
+import tempfile
+import threading
 from datetime import datetime, timezone
 from pathlib import Path
 
@@ -21,6 +23,7 @@
 from benchmarks.utils.harbor import (
     HarborCredentialMode,
     check_harbor_installed as _check_harbor_installed,
+    completed_harbor_task_ids,
     convert_harbor_to_eval_output,
     get_supported_task_filter_flag,
     run_harbor_evaluation as _run_harbor_evaluation,
@@ -33,6 +36,107 @@
 
 # Output filename for results
 OUTPUT_FILENAME = "output.jsonl"
+CHECKPOINT_INTERVAL_SECONDS = 60
+
+
+def _checkpoint_gcs_uri() -> str | None:
+    bucket = os.environ.get("RESULTS_BUCKET")
+    model_slug = os.environ.get("MODEL_SLUG")
+    run_id = os.environ.get("GITHUB_RUN_ID")
+    if not bucket or not model_slug or not run_id:
+        return None
+    return f"gs://{bucket}/terminalbench/{model_slug}/{run_id}/checkpoint.tar.gz"
+
+
+def restore_harbor_checkpoint(structured_output_dir: Path) -> bool:
+    """Restore incremental Harbor results after a Kubernetes pod retry."""
+    checkpoint_uri = _checkpoint_gcs_uri()
+    if checkpoint_uri is None:
+        return False
+    stat = subprocess.run(
+        ["gsutil", "-q", "stat", checkpoint_uri],
+        capture_output=True,
+        text=True,
+    )
+    if stat.returncode != 0:
+        return False
+
+    structured_output_dir.parent.mkdir(parents=True, exist_ok=True)
+    with tempfile.TemporaryDirectory(prefix="terminalbench-checkpoint-") as tmpdir:
+        archive = Path(tmpdir) / "checkpoint.tar.gz"
+        subprocess.run(["gsutil", "cp", checkpoint_uri, str(archive)], check=True)
+        subprocess.run(
+            [
+                "tar",
+                "-xzf",
+                str(archive),
+                "-C",
+                str(structured_output_dir.parent),
+            ],
+            check=True,
+        )
+    logger.info("Restored Harbor checkpoint from %s", checkpoint_uri)
+    return True
+
+
+def upload_harbor_checkpoint(
+    harbor_output_dir: Path,
+    output_path: Path,
+) -> bool:
+    """Convert and upload all complete Harbor trials seen so far."""
+    checkpoint_uri = _checkpoint_gcs_uri()
+    if checkpoint_uri is None or not harbor_output_dir.exists():
+        return False
+    completed_ids = completed_harbor_task_ids(harbor_output_dir)
+    if not completed_ids:
+        return False
+
+    convert_harbor_to_eval_output(
+        harbor_output_dir=harbor_output_dir,
+        eval_output_path=output_path,
+    )
+    with tempfile.TemporaryDirectory(prefix="terminalbench-checkpoint-") as tmpdir:
+        archive = Path(tmpdir) / "checkpoint.tar.gz"
+        subprocess.run(
+            [
+                "tar",
+                "-czf",
+                str(archive),
+                "-C",
+                str(output_path.parent.parent),
+                output_path.parent.name,
+            ],
+            check=True,
+        )
+        subprocess.run(
+            [
+                "gsutil",
+                "-h",
+                "Cache-Control:no-cache, no-store, must-revalidate",
+                "cp",
+                str(archive),
+                checkpoint_uri,
+            ],
+            check=True,
+        )
+    logger.info(
+        "Uploaded resumable Harbor checkpoint with %d tasks to %s",
+        len(completed_ids),
+        checkpoint_uri,
+    )
+    return True
+
+
+def _checkpoint_loop(
+    stop_event: threading.Event,
+    harbor_output_dir: Path,
+    output_path: Path,
+) -> None:
+    while not stop_event.wait(CHECKPOINT_INTERVAL_SECONDS):
+        try:
+            upload_harbor_checkpoint(harbor_output_dir, output_path)
+        except Exception:
+            logger.exception("Failed to upload Harbor checkpoint; will retry")
 
 
 def check_harbor_installed() -> bool:
@@ -193,6 +297,11 @@ def main() -> None:
     )
 
     logger.info(f"Output directory: {structured_output_dir}")
+    structured_output_path = Path(structured_output_dir)
+    try:
+        restore_harbor_checkpoint(structured_output_path)
+    except Exception:
+        logger.exception("Failed to restore Harbor checkpoint; starting without it")
     os.makedirs(structured_output_dir, exist_ok=True)
 
     # Save metadata
@@ -222,20 +331,51 @@ def main() -> None:
     if not args.skip_harbor:
         # Run harbor evaluation
         try:
-            harbor_output_dir = run_harbor_evaluation(
-                llm=llm,
-                dataset=args.dataset,
-                output_dir=structured_output_dir,
-                num_workers=args.num_workers,
-                task_ids=task_ids,
-                n_limit=args.n_limit,
-            )
+            harbor_output_dir = Path(structured_output_dir) / "harbor_output"
+            if task_ids and harbor_output_dir.exists():
+                completed_ids = completed_harbor_task_ids(harbor_output_dir)
+                if completed_ids:
+                    task_ids = [
+                        task_id for task_id in task_ids if task_id not in completed_ids
+                    ]
+                    logger.info(
+                        "Restored %d completed Harbor tasks; %d selected tasks remain",
+                        len(completed_ids),
+                        len(task_ids),
+                    )
+
+            if task_ids is None or task_ids:
+                stop_event = threading.Event()
+                checkpoint_thread = threading.Thread(
+                    target=_checkpoint_loop,
+                    args=(stop_event, harbor_output_dir, output_path),
+                    name="terminalbench-checkpoint",
+                    daemon=True,
+                )
+                checkpoint_thread.start()
+                try:
+                    harbor_output_dir = run_harbor_evaluation(
+                        llm=llm,
+                        dataset=args.dataset,
+                        output_dir=structured_output_dir,
+                        num_workers=args.num_workers,
+                        task_ids=task_ids,
+                        n_limit=args.n_limit,
+                    )
+                finally:
+                    stop_event.set()
+                    checkpoint_thread.join()
+            else:
+                logger.info(
+                    "All selected tasks were restored; skipping Harbor execution"
+                )
 
             # Convert harbor output to standard format
             convert_harbor_to_eval_output(
                 harbor_output_dir=harbor_output_dir,
                 eval_output_path=output_path,
             )
+            upload_harbor_checkpoint(harbor_output_dir, output_path)
 
         except Exception as e:
             logger.error(f"Evaluation failed: {e}")

benchmarks/utils/harbor.py

@@ -204,6 +204,27 @@ def _find_job_dir(harbor_output_dir: Path) -> Path:
     return sorted(candidates)[-1]
 
 
+def find_harbor_trial_result_files(harbor_output_dir: Path) -> list[Path]:
+    """Return trial results from every timestamped Harbor attempt."""
+    return sorted(harbor_output_dir.glob("*/*/result.json"))
+
+
+def completed_harbor_task_ids(harbor_output_dir: Path) -> set[str]:
+    """Return task IDs with a complete, parseable Harbor trial result."""
+    completed: set[str] = set()
+    for result_file in find_harbor_trial_result_files(harbor_output_dir):
+        try:
+            with result_file.open() as stream:
+                trial = json.load(stream)
+        except (json.JSONDecodeError, OSError):
+            continue
+        task_name = trial.get("task_name")
+        rewards = (trial.get("verifier_result") or {}).get("rewards") or {}
+        if isinstance(task_name, str) and task_name and "reward" in rewards:
+            completed.add(task_name)
+    return completed
+
+
 def convert_harbor_to_eval_output(
     harbor_output_dir: Path,
     eval_output_path: Path,
@@ -214,17 +235,34 @@ def convert_harbor_to_eval_output(
     logger.info(f"Converting harbor output from {harbor_output_dir}")
 
     canonicalize = canonicalize_instance_id or (lambda instance_id: instance_id)
-    job_dir = _find_job_dir(harbor_output_dir)
-    logger.info(f"Using harbor job directory: {job_dir}")
-
-    result_files = [f for f in job_dir.glob("*/result.json") if f.parent != job_dir]
+    result_files = find_harbor_trial_result_files(harbor_output_dir)
     if not result_files:
+        # Preserve the established distinction between no Harbor job at all
+        # and a completed job that contains no trial results.
+        job_dir = _find_job_dir(harbor_output_dir)
         raise RuntimeError(
             f"No trial result files found in {job_dir}. "
-            f"Expected result.json files in trial subdirectories."
+            "Expected result.json files in timestamp/trial subdirectories."
         )
 
-    logger.info(f"Found {len(result_files)} trial results in {job_dir}")
+    logger.info(
+        f"Found {len(result_files)} trial results across Harbor attempts in "
+        f"{harbor_output_dir}"
+    )
+
+    # A retry normally filters restored IDs, but a crash can occur between a
+    # result write and the next checkpoint. Prefer the later timestamped copy.
+    latest_result_by_instance: dict[str, Path] = {}
+    unreadable_result_files: list[Path] = []
+    for result_file in result_files:
+        try:
+            with result_file.open() as stream:
+                trial = json.load(stream)
+            instance_id = canonicalize(trial.get("task_name", result_file.parent.name))
+            latest_result_by_instance[instance_id] = result_file
+        except (json.JSONDecodeError, OSError):
+            unreadable_result_files.append(result_file)
+    result_files = list(latest_result_by_instance.values()) + unreadable_result_files
 
     results: list[dict] = []
     errors: list[dict] = []

tests/test_terminalbench.py

@@ -8,9 +8,12 @@
 from benchmarks.terminalbench.config import INFER_DEFAULTS
 from benchmarks.terminalbench.eval_infer import process_terminalbench_results
 from benchmarks.terminalbench.run_infer import (
+    _checkpoint_gcs_uri,
     convert_harbor_to_eval_output,
     run_harbor_evaluation,
+    upload_harbor_checkpoint,
 )
+from benchmarks.utils.harbor import completed_harbor_task_ids
 from openhands.sdk import LLM
 
 
@@ -522,3 +525,81 @@ def test_discovery_finds_all_trials(self, tmp_path: Path) -> None:
         assert len(entries) == 5
         instance_ids = {e["instance_id"] for e in entries}
         assert instance_ids == {f"task-{i}" for i in range(5)}
+
+    def test_completed_ids_require_a_primary_reward(self, tmp_path: Path) -> None:
+        harbor_dir = tmp_path / "harbor_output"
+        attempt = harbor_dir / "2026-01-01__00-00-00"
+        for task_name, payload in {
+            "passed": {"verifier_result": {"rewards": {"reward": 1.0}}},
+            "failed": {"verifier_result": {"rewards": {"reward": 0.0}}},
+            "incomplete": {"verifier_result": {"rewards": {}}},
+        }.items():
+            trial_dir = attempt / task_name
+            trial_dir.mkdir(parents=True)
+            (trial_dir / "result.json").write_text(
+                json.dumps({"task_name": task_name, **payload})
+            )
+        malformed_dir = attempt / "malformed"
+        malformed_dir.mkdir()
+        (malformed_dir / "result.json").write_text("{")
+
+        assert completed_harbor_task_ids(harbor_dir) == {"passed", "failed"}
+
+    def test_conversion_combines_attempts_and_prefers_latest(
+        self, tmp_path: Path
+    ) -> None:
+        harbor_dir = tmp_path / "harbor_output"
+        for attempt_name, task_name, trial_name, reward in [
+            ("2026-01-01__00-00-00", "task-a", "task-a__old", 0.0),
+            ("2026-01-02__00-00-00", "task-b", "task-b__new", 1.0),
+            ("2026-01-02__00-00-00", "task-a", "task-a__new", 1.0),
+        ]:
+            attempt = harbor_dir / attempt_name
+            attempt.mkdir(parents=True, exist_ok=True)
+            (attempt / "result.json").write_text(json.dumps({"id": attempt_name}))
+            trial_dir = attempt / trial_name
+            trial_dir.mkdir()
+            (trial_dir / "result.json").write_text(
+                json.dumps(
+                    {
+                        "task_name": task_name,
+                        "trial_name": trial_name,
+                        "agent_result": {},
+                        "verifier_result": {"rewards": {"reward": reward}},
+                        "exception_info": None,
+                    }
+                )
+            )
+
+        output_file = tmp_path / "output.jsonl"
+        convert_harbor_to_eval_output(harbor_dir, output_file)
+        entries = {
+            entry["instance_id"]: entry
+            for entry in map(json.loads, output_file.read_text().splitlines())
+        }
+
+        assert set(entries) == {"task-a", "task-b"}
+        assert entries["task-a"]["test_result"]["trial_name"] == "task-a__new"
+        assert entries["task-a"]["test_result"]["passed"] is True
+
+    def test_checkpointing_is_disabled_without_run_coordinates(
+        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        monkeypatch.delenv("RESULTS_BUCKET", raising=False)
+        monkeypatch.delenv("MODEL_SLUG", raising=False)
+        monkeypatch.delenv("GITHUB_RUN_ID", raising=False)
+
+        assert _checkpoint_gcs_uri() is None
+        assert upload_harbor_checkpoint(tmp_path, tmp_path / "output.jsonl") is False
+
+    def test_checkpoint_uri_uses_run_coordinates(
+        self, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        monkeypatch.setenv("RESULTS_BUCKET", "results")
+        monkeypatch.setenv("MODEL_SLUG", "model")
+        monkeypatch.setenv("GITHUB_RUN_ID", "123")
+
+        assert (
+            _checkpoint_gcs_uri()
+            == "gs://results/terminalbench/model/123/checkpoint.tar.gz"
+        )