fix flaky crime download

pipeline/download/crime.py

@@ -15,6 +15,7 @@ import json
 import re
 import shutil
 import sys
+import time
 import zipfile
 from dataclasses import asdict, dataclass
 from datetime import UTC, datetime
@@ -35,6 +36,11 @@ ARCHIVE_LINK_RE = re.compile(
 )
 VALID_MD5_RE = re.compile(r"^[0-9a-fA-F]{32}$")
 MONTH_RE = re.compile(r"^\d{4}-\d{2}$")
+# Each archive is up to ~1.6 GB, so a transient failure (network blip, or the
+# work dir being removed underneath us by a concurrent run) must not abort the
+# whole multi-archive download. Retry a few times with a short backoff.
+DOWNLOAD_ATTEMPTS = 4
+RETRY_BACKOFF_SECONDS = 5
 STREET_CRIME_CSV_RE = re.compile(r"^\d{4}-\d{2}-.+-street\.csv$")
 CONTAINED_RANGE_RE = re.compile(
     r"Contains data from (?P<start_month>[A-Za-z]+) (?P<start_year>\d{4}) "
@@ -265,6 +271,10 @@ def file_md5(path: Path) -> str:
     return digest.hexdigest()
 
 
+class _RetryableDownloadError(Exception):
+    """A download attempt failed in a way that is worth retrying."""
+
+
 def download_archive(
     archive: CrimeArchive,
     archive_dir: Path,
@@ -273,7 +283,12 @@ def download_archive(
     force: bool,
     timeout: float,
 ) -> Path:
-    """Download one archive ZIP, resuming an existing .part file when possible."""
+    """Download one archive ZIP, resuming an existing .part file when possible.
+
+    Transient failures (network errors, or the work dir being removed
+    underneath us by a concurrent run) are retried with a short backoff so a
+    single hiccup does not abort a whole multi-archive download.
+    """
     dest = archive_dir / archive.filename
     partial = dest.with_suffix(dest.suffix + ".part")
 
@@ -297,6 +312,43 @@ def download_archive(
             print(f"{archive.filename}: already downloaded")
             return dest
 
+    last_error: Exception | None = None
+    for attempt in range(1, DOWNLOAD_ATTEMPTS + 1):
+        # The work dir can be removed underneath us (e.g. a concurrent run's
+        # startup cleanup), which is exactly what makes the final rename fail
+        # with ENOENT even after the download reached 100%. Recreate it so the
+        # (re)download has somewhere to land; a surviving .part is resumed.
+        archive_dir.mkdir(parents=True, exist_ok=True)
+        try:
+            _stream_to_partial(archive, partial, timeout=timeout)
+            partial.replace(dest)
+            if verify and archive.md5 is not None:
+                actual_md5 = file_md5(dest)
+                if actual_md5 != archive.md5:
+                    dest.unlink(missing_ok=True)
+                    raise _RetryableDownloadError(
+                        f"{archive.filename}: MD5 mismatch: "
+                        f"expected {archive.md5}, got {actual_md5}"
+                    )
+            return dest
+        except (httpx.HTTPError, OSError, _RetryableDownloadError) as error:
+            last_error = error
+            if attempt < DOWNLOAD_ATTEMPTS:
+                delay = RETRY_BACKOFF_SECONDS * attempt
+                print(
+                    f"{archive.filename}: download attempt {attempt} failed "
+                    f"({error}); retrying in {delay}s",
+                    file=sys.stderr,
+                )
+                time.sleep(delay)
+
+    raise RuntimeError(
+        f"{archive.filename}: download failed after {DOWNLOAD_ATTEMPTS} attempts"
+    ) from last_error
+
+
+def _stream_to_partial(archive: CrimeArchive, partial: Path, *, timeout: float) -> None:
+    """Stream one archive into its .part file, resuming a valid prefix."""
     resume_from = partial.stat().st_size if partial.exists() else 0
     headers = {"Range": f"bytes={resume_from}-"} if resume_from else {}
 
@@ -332,18 +384,6 @@ def download_archive(
                 output.write(chunk)
                 progress.update(len(chunk))
 
-    partial.replace(dest)
-
-    if verify and archive.md5 is not None:
-        actual_md5 = file_md5(dest)
-        if actual_md5 != archive.md5:
-            dest.unlink(missing_ok=True)
-            raise RuntimeError(
-                f"{archive.filename}: MD5 mismatch: expected {archive.md5}, got {actual_md5}"
-            )
-
-    return dest
-
 
 def _is_street_crime_csv(path: PurePosixPath | Path) -> bool:
     return STREET_CRIME_CSV_RE.fullmatch(path.name) is not None