diff --git a/composer/prover/cloud.py b/composer/prover/cloud.py index 33dba95b..1aa865f2 100644 --- a/composer/prover/cloud.py +++ b/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) diff --git a/tests/test_cloud_fetch_retry.py b/tests/test_cloud_fetch_retry.py new file mode 100644 index 00000000..7f8bf26e --- /dev/null +++ b/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