autoprover: re-fetch a completed cloud job on transport failure

composer/prover/cloud.py

@@ -7,6 +7,7 @@
 
 import asyncio
 import logging
+import shutil
 import tempfile
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
@@ -159,6 +160,42 @@ def _results_api() -> ProverOutputAPI:
     return ProverOutputAPI(enable_cache=False)
 
 
+#: A fetch failure here is a *transport* fault, not a bad job: by the time we download documents
+#: the job has already polled ``SUCCEEDED``. Re-fetching the same completed job is cheap; re-running
+#: the whole proof (what the caller does if this context raises) throws away a finished — often
+#: hours-long — verification. So retry the fetch itself a few times with exponential backoff before
+#: giving up. POU retries at the individual-request level; this covers a whole-fetch failure that
+#: outlives those (e.g. a mid-transfer reset that exhausts the request-level retries on one file).
+_FETCH_MAX_ATTEMPTS = 3
+_FETCH_BACKOFF_BASE_S = 2.0
+
+
+async def _fetch_results(job_id: str, dest: Path) -> None:
+    """Download a completed job's sources + tree view into ``dest``, retrying a transient failure.
+
+    Each attempt starts from an empty ``dest`` so a partially-written archive from a failed attempt
+    cannot leave a truncated file behind. Never re-runs the job — only re-reads it.
+    """
+    for attempt in range(1, _FETCH_MAX_ATTEMPTS + 1):
+        try:
+            await asyncio.to_thread(
+                _results_api().fetch_sources_and_treeview_files, job_id, dest
+            )
+            return
+        except Exception as exc:
+            if attempt == _FETCH_MAX_ATTEMPTS:
+                raise
+            backoff = _FETCH_BACKOFF_BASE_S * 2 ** (attempt - 1)
+            logger.warning(
+                "Fetching results for completed job %s failed (attempt %d/%d): %s. "
+                "Re-fetching the finished job in %.0fs (not re-proving).",
+                job_id[:8], attempt, _FETCH_MAX_ATTEMPTS, exc, backoff,
+            )
+            for child in dest.iterdir():
+                shutil.rmtree(child) if child.is_dir() else child.unlink()
+            await asyncio.sleep(backoff)
+
+
 @asynccontextmanager
 async def cloud_results(
     run_result_link: str,
@@ -202,7 +239,5 @@ async def on_status(status: str) -> None:
         # tens of gigabytes on real jobs and used to exhaust the disk; across the
         # jobs measured here these two subtrees are ~3% of the archive. POU writes
         # them in the same layout the archive had, so the parse is unchanged.
-        await asyncio.to_thread(
-            _results_api().fetch_sources_and_treeview_files, cloud_job.job_id, dest
-        )
+        await _fetch_results(cloud_job.job_id, dest)
         yield (dest, runtime_ms)

tests/test_cloud_fetch_retry.py

@@ -0,0 +1,52 @@
+"""Unit tests for cloud_results' fetch retry: re-fetch a *completed* job, never re-prove."""
+
+import asyncio
+from pathlib import Path
+
+import pytest
+
+import composer.prover.cloud as cloud
+
+
+class _FakeAPI:
+    """Stand-in for the POU results client whose fetch fails ``fail_times`` times, then succeeds."""
+
+    def __init__(self, fail_times: int, files: list[str]):
+        self.fail_times = fail_times
+        self.files = files
+        self.calls = 0
+
+    def fetch_sources_and_treeview_files(self, job_id: str, dest: Path) -> None:
+        self.calls += 1
+        if self.calls <= self.fail_times:
+            # a mid-transfer reset can leave a partial file behind
+            (Path(dest) / f"partial_{self.calls}.txt").write_text("x")
+            raise ConnectionResetError("Connection reset by peer")
+        for name in self.files:
+            (Path(dest) / name).write_text("ok")
+
+
+@pytest.fixture(autouse=True)
+def _no_backoff_wait(monkeypatch):
+    # keep the test instant — exercise the retry loop, not the clock
+    monkeypatch.setattr(cloud, "_FETCH_BACKOFF_BASE_S", 0.0)
+
+
+def _fetch(fake: _FakeAPI, dest: Path, monkeypatch) -> None:
+    monkeypatch.setattr(cloud, "_results_api", lambda: fake)
+    asyncio.run(cloud._fetch_results("deadbeefcafebabe", dest))
+
+
+def test_succeeds_after_transient_failures(tmp_path, monkeypatch):
+    fake = _FakeAPI(fail_times=2, files=["a.json", "b.json"])
+    _fetch(fake, tmp_path, monkeypatch)
+    assert fake.calls == 3  # 2 transient failures, then success — the finished job is only re-read
+    # each failed attempt's partial write was cleared before the next; only the good set remains
+    assert {p.name for p in tmp_path.iterdir()} == {"a.json", "b.json"}
+
+
+def test_gives_up_after_max_attempts(tmp_path, monkeypatch):
+    fake = _FakeAPI(fail_times=99, files=["a.json"])
+    with pytest.raises(ConnectionResetError):
+        _fetch(fake, tmp_path, monkeypatch)
+    assert fake.calls == cloud._FETCH_MAX_ATTEMPTS  # capped — it never loops forever