Smarter download_checkpoints: skip download if /data/Lyra-2 already populated (direct model mount)

download_checkpoints.py

@@ -1,9 +1,10 @@
-"""Download Lyra-2 checkpoints from HuggingFace to the persistent /data volume.
+"""Verify or download Lyra-2.0 checkpoints.
 
-Run at container boot from entrypoint.sh. Idempotent: if the expected directory
-already exists and contains the model checkpoint, skip the download entirely.
-The ~50-80 GB download only happens on the first cold start after the paid
-Space's /data volume is provisioned — subsequent starts reuse the cached copy.
+Three cases handled, in priority order:
+1. /data/Lyra-2/checkpoints/model is already populated (read-only direct-mount
+   OR warm bucket): skip download, just symlink.
+2. /data is writable + empty: download from HF Hub.
+3. /data is unwritable + empty: fall back to ephemeral cache.
 """
 
 from __future__ import annotations
@@ -20,73 +21,83 @@ log = logging.getLogger("lyra2-space.checkpoints")
 logging.basicConfig(level=logging.INFO, format="[%(asctime)s] %(message)s")
 
 REPO_ID = "nvidia/Lyra-2.0"
-
-# /data is the HF Spaces persistent volume mount point on paid GPU Spaces.
-# Fall back to a local directory if /data is unavailable (e.g. local dev runs).
 DATA_ROOT = Path(os.environ.get("LYRA2_DATA_ROOT", "/data"))
-CHECKPOINTS_ROOT = DATA_ROOT / "Lyra-2"
-MODEL_MARKER = CHECKPOINTS_ROOT / "checkpoints" / "model"
-
 APP_ROOT = Path("/home/user/app")
-APP_LYRA2 = APP_ROOT / "Lyra-2"
-APP_CHECKPOINTS = APP_LYRA2 / "checkpoints"
+APP_CHECKPOINTS = APP_ROOT / "Lyra-2" / "checkpoints"
+
+
+def _already_populated(root: Path) -> bool:
+    """True if root/Lyra-2/checkpoints/model has content."""
+    marker = root / "Lyra-2" / "checkpoints" / "model"
+    return marker.is_dir() and any(marker.iterdir())
 
 
-def _ensure_writable_data_dir() -> Path:
+def _is_writable(root: Path) -> bool:
     try:
-        DATA_ROOT.mkdir(parents=True, exist_ok=True)
-        probe = DATA_ROOT / ".writable_probe"
+        root.mkdir(parents=True, exist_ok=True)
+        probe = root / ".writable_probe"
         probe.write_text("ok")
         probe.unlink()
-        return DATA_ROOT
-    except OSError as exc:
-        # /data isn't writable (local dev, wrong hardware tier, etc.). Fall back
-        # to an in-container path; note this means checkpoints re-download on
-        # every cold start. Good for local dev, bad for prod.
-        fallback = Path.home() / "app" / "_checkpoints_cache"
-        log.warning("Cannot use %s (%s); falling back to %s", DATA_ROOT, exc, fallback)
-        fallback.mkdir(parents=True, exist_ok=True)
-        return fallback
+        return True
+    except OSError:
+        return False
 
 
-def main() -> int:
-    data_root = _ensure_writable_data_dir()
-    lyra2_root = data_root / "Lyra-2"
-    lyra2_root.mkdir(parents=True, exist_ok=True)
+def _symlink_app_checkpoints(lyra2_root: Path) -> None:
+    """Point /home/user/app/Lyra-2/checkpoints at <lyra2_root>/checkpoints."""
+    src = lyra2_root / "checkpoints"
+    dst = APP_CHECKPOINTS
+    if dst.is_symlink() and dst.resolve() == src.resolve():
+        log.info("Checkpoints symlink already points to %s", src)
+        return
+    if dst.is_symlink() or dst.is_file():
+        dst.unlink()
+    elif dst.exists():
+        shutil.rmtree(dst)
+    dst.parent.mkdir(parents=True, exist_ok=True)
+    dst.symlink_to(src, target_is_directory=True)
+    log.info("Linked %s -> %s", dst, src)
+
 
-    model_marker = lyra2_root / "checkpoints" / "model"
-    if model_marker.exists() and any(model_marker.iterdir()):
-        log.info("Checkpoints already present at %s; skipping download.", model_marker)
-    else:
-        log.info("Downloading %s checkpoints/* to %s — this is ~50-80 GB and takes 20-40 min on cold start.",
+def main() -> int:
+    # Case 1: /data already has the checkpoints (direct-mount OR warm bucket)
+    if DATA_ROOT.is_dir() and _already_populated(DATA_ROOT):
+        log.info("Checkpoints present at %s/Lyra-2 — skipping download.", DATA_ROOT)
+        _symlink_app_checkpoints(DATA_ROOT / "Lyra-2")
+        return 0
+
+    # Case 2: /data writable but empty — download
+    if _is_writable(DATA_ROOT):
+        lyra2_root = DATA_ROOT / "Lyra-2"
+        lyra2_root.mkdir(parents=True, exist_ok=True)
+        log.info("Downloading %s checkpoints/* to %s (first cold start only).",
                  REPO_ID, lyra2_root)
         snapshot_download(
             repo_id=REPO_ID,
             allow_patterns=["checkpoints/*"],
             local_dir=str(lyra2_root),
-            local_dir_use_symlinks=False,
-            # Resume partial downloads across container restarts.
             max_workers=8,
         )
         log.info("Checkpoint download complete.")
-
-    # The Lyra-2 inference modules expect `checkpoints/` at Lyra-2/ repo root.
-    # Symlink the persistent-volume copy into the app's Lyra-2 tree so every
-    # existing `--checkpoint_dir checkpoints/model` CLI stays unchanged.
-    src = lyra2_root / "checkpoints"
-    dst = APP_CHECKPOINTS
-    if dst.is_symlink() or dst.exists():
-        if dst.is_symlink() and dst.resolve() == src.resolve():
-            log.info("Checkpoints symlink already points to %s", src)
-            return 0
-        log.info("Replacing stale %s with symlink to %s", dst, src)
-        if dst.is_symlink() or dst.is_file():
-            dst.unlink()
-        else:
-            shutil.rmtree(dst)
-    dst.parent.mkdir(parents=True, exist_ok=True)
-    dst.symlink_to(src, target_is_directory=True)
-    log.info("Linked %s -> %s", dst, src)
+        _symlink_app_checkpoints(lyra2_root)
+        return 0
+
+    # Case 3: /data unwritable + empty — ephemeral fallback
+    fallback = APP_ROOT / "_checkpoints_cache"
+    log.warning("Cannot use %s (read-only AND empty). Falling back to ephemeral %s",
+                DATA_ROOT, fallback)
+    fallback.mkdir(parents=True, exist_ok=True)
+    lyra2_root = fallback / "Lyra-2"
+    lyra2_root.mkdir(parents=True, exist_ok=True)
+    if not _already_populated(fallback):
+        log.info("Downloading %s checkpoints/* to %s.", REPO_ID, lyra2_root)
+        snapshot_download(
+            repo_id=REPO_ID,
+            allow_patterns=["checkpoints/*"],
+            local_dir=str(lyra2_root),
+            max_workers=8,
+        )
+    _symlink_app_checkpoints(lyra2_root)
     return 0