New --force flag for add. Appease PET (#197)

Author: WillEngler

rootstock/cli.py

@@ -243,6 +243,15 @@ def main():
         action="store_true",
         help="Skip the verify phase (download only). Login-node escape hatch.",
     )
+    add_parser.add_argument(
+        "--force",
+        action="store_true",
+        help=(
+            "Re-run the download even if the manifest records the checkpoint "
+            "as fetched — repairs a cache file that has gone missing. Cheap "
+            "when the cache is intact (cache hit, no transfer)."
+        ),
+    )
     add_parser.add_argument(
         "--root",
         default=os.environ.get(ROOTSTOCK_ROOT_ENV),

rootstock/commands/add.py

@@ -82,6 +82,7 @@ def cmd_add(args) -> int:
             device=args.device,
             verify=not args.no_verify,
             push=not args.no_push,
+            force=args.force,
             setup_kwargs=setup_kwargs,
             progress=print,
         )

rootstock/operations.py

@@ -1041,6 +1041,7 @@ def fetch_checkpoint(
     cache_root: Path | None = None,
     refresh: bool = True,
     push: bool = True,
+    force: bool = False,
     progress: Progress | None = None,
 ) -> FetchResult:
     """
@@ -1055,6 +1056,12 @@ def fetch_checkpoint(
     skips the trailing full manifest refresh (+ push): a batch driver running
     many fetches refreshes once at the end instead of once per checkpoint.
 
+    ``force=True`` re-runs the download even when the manifest records the
+    checkpoint as fetched — the repair path for cache files that have gone
+    missing behind the manifest's back (cleaned, or an interrupted download
+    that stamped anyway). The underlying download is cache-aware, so a
+    forced fetch of an intact checkpoint costs a cache hit, not a transfer.
+
     Raises CheckpointNotFoundError when no installed env declares the id,
     and OperationError when the download fails (also recorded in the
     manifest's ``last_error``).
@@ -1072,7 +1079,7 @@ def fetch_checkpoint(
     peek = load_manifest(root)
     peek_env = peek.environments.get(env_name) if peek else None
     peek_ckpt = peek_env.checkpoints.get(checkpoint) if peek_env else None
-    already_fetched = peek_ckpt is not None and peek_ckpt.fetched_at is not None
+    already_fetched = not force and peek_ckpt is not None and peek_ckpt.fetched_at is not None
     fetched_at = peek_ckpt.fetched_at if peek_ckpt else None
 
     # ---- Download (runs outside the lock) -------------------------------
@@ -1183,6 +1190,7 @@ def add_checkpoint(
     device: str = "cuda",
     verify: bool = True,
     push: bool = True,
+    force: bool = False,
     setup_kwargs: dict | None = None,
     cache_root: Path | None = None,
     progress: Progress | None = None,
@@ -1196,7 +1204,8 @@ def add_checkpoint(
     refresh at the end. The hosting env is resolved by matching the id
     against each installed env's CHECKPOINTS table. Downloads happen on CPU
     (the cache-aware path); verification runs on ``device`` unless
-    ``verify`` is False.
+    ``verify`` is False. ``force`` re-runs the download past the manifest's
+    fetched stamp (see :func:`fetch_checkpoint`).
 
     Raises CheckpointNotFoundError when no installed env declares the id,
     and OperationError for download/verify failures (the failure is also
@@ -1214,6 +1223,7 @@ def add_checkpoint(
         setup_kwargs=setup_kwargs,
         cache_root=cache_root,
         refresh=False,
+        force=force,
         progress=progress,
     )
 

sample_model_configurations/nvidia_configs/pet.py

@@ -24,7 +24,18 @@
 
 
 def setup(checkpoint: str, device: str = "cuda"):
+    from huggingface_hub import hf_hub_download
     from upet.calculator import UPETCalculator
 
+    # Passing model=/version= makes UPETCalculator resolve the name by listing
+    # the hub repo — an uncached API call that fails on workers, which run
+    # with HF_HUB_OFFLINE=1 (and on any node without internet). Fetch the
+    # pinned file ourselves — a cache hit needs no network even offline — and
+    # hand it over as checkpoint_path, which skips the resolve entirely.
     model, version = CHECKPOINTS[checkpoint].split("@", 1)
-    return UPETCalculator(model=model, version=version, device=device)
+    path = hf_hub_download(
+        repo_id="lab-cosmo/upet",
+        filename=f"{model}-v{version}.ckpt",
+        subfolder="models",
+    )
+    return UPETCalculator(checkpoint_path=path, device=device)

tests/cli/test_add.py

@@ -94,6 +94,7 @@ class _Args:
     args.kwarg = overrides.get("kwarg")
     args.device = overrides.get("device", "cuda")
     args.no_verify = overrides.get("no_verify", False)
+    args.force = overrides.get("force", False)
     args.root = str(root)
     args.no_push = overrides.get("no_push", True)
     return args
@@ -159,6 +160,22 @@ def fake_verify(*a, **kw):
     assert ckpt.last_error is None
 
 
+def test_add_force_redownloads(fake_root, monkeypatch):
+    """--force repairs a cache file gone missing behind the manifest's
+    fetched stamp."""
+    download_calls = []
+    monkeypatch.setattr(
+        operations,
+        "_run_download",
+        lambda *a, **kw: (download_calls.append(a), (True, None))[1],
+    )
+
+    assert cmd_add(_make_args(fake_root, no_verify=True)) == 0
+    assert cmd_add(_make_args(fake_root, no_verify=True, force=True)) == 0
+
+    assert len(download_calls) == 2, "--force must re-run the download"
+
+
 def test_add_records_download_failure_and_returns_1(fake_root, monkeypatch):
     monkeypatch.setattr(
         operations,

tests/commands/test_add_split.py

@@ -101,6 +101,22 @@ def test_fetch_is_idempotent(fake_root, refresh_calls, monkeypatch):
     assert second.fetched_at == first.fetched_at
 
 
+def test_fetch_force_redownloads_past_fetched_stamp(fake_root, refresh_calls, monkeypatch):
+    """The repair path for cache files gone missing behind the manifest."""
+    downloads = []
+    monkeypatch.setattr(
+        operations, "_run_download", lambda *a, **kw: (downloads.append(a), (True, None))[1]
+    )
+
+    first = fetch_checkpoint(fake_root, "mace-mp-0-medium")
+    second = fetch_checkpoint(fake_root, "mace-mp-0-medium", force=True)
+
+    assert len(downloads) == 2, "force must re-run the download"
+    assert not second.already_fetched
+    assert second.fetched_at is not None
+    assert second.fetched_at >= first.fetched_at
+
+
 def test_fetch_failure_records_last_error_and_raises(fake_root, refresh_calls, monkeypatch):
     monkeypatch.setattr(
         operations, "_run_download", lambda *a, **kw: (False, "ConnectionError: hub unreachable")

tests/sample_configs/test_pet_offline_load.py

@@ -0,0 +1,89 @@
+"""Tests for the pet env's hub-resolve bypass.
+
+UPETCalculator(model=..., version=...) resolves the name by listing the hub
+repo — an uncached API call that fails on workers (HF_HUB_OFFLINE=1) and on
+nodes without internet. The env fetches the pinned file itself via
+hf_hub_download (a cache hit needs no network) and passes checkpoint_path,
+which skips the resolve.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import sys
+import types
+from pathlib import Path
+
+import pytest
+
+_CONFIGS_DIR = (
+    Path(__file__).parent.parent.parent / "sample_model_configurations" / "nvidia_configs"
+)
+
+
+def _load_env_module():
+    spec = importlib.util.spec_from_file_location("pet_env", _CONFIGS_DIR / "pet.py")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+@pytest.fixture
+def stubbed_libs(monkeypatch):
+    """Stub huggingface_hub and upet; capture what setup() passes to each."""
+    captured = {}
+
+    hf = types.ModuleType("huggingface_hub")
+
+    def hf_hub_download(**kwargs):
+        captured["download"] = kwargs
+        return "/shared/cache/models/stub.ckpt"
+
+    hf.hf_hub_download = hf_hub_download
+
+    upet = types.ModuleType("upet")
+    upet_calculator = types.ModuleType("upet.calculator")
+
+    class UPETCalculator:
+        def __init__(self, **kwargs):
+            captured["calculator"] = kwargs
+
+    upet_calculator.UPETCalculator = UPETCalculator
+    upet.calculator = upet_calculator
+
+    monkeypatch.setitem(sys.modules, "huggingface_hub", hf)
+    monkeypatch.setitem(sys.modules, "upet", upet)
+    monkeypatch.setitem(sys.modules, "upet.calculator", upet_calculator)
+    return captured
+
+
+def test_setup_downloads_pinned_filename(stubbed_libs):
+    env = _load_env_module()
+    env.setup("pet-oam-xl", device="cuda")
+    assert stubbed_libs["download"] == {
+        "repo_id": "lab-cosmo/upet",
+        "filename": "pet-oam-xl-v1.0.0.ckpt",
+        "subfolder": "models",
+    }
+
+
+def test_setup_passes_checkpoint_path_not_model_name(stubbed_libs):
+    """model=/version= would trigger the hub-listing resolve — never pass them."""
+    env = _load_env_module()
+    env.setup("pet-omatpes-l", device="cpu")
+    calc_kwargs = stubbed_libs["calculator"]
+    assert calc_kwargs == {
+        "checkpoint_path": "/shared/cache/models/stub.ckpt",
+        "device": "cpu",
+    }
+
+
+def test_every_checkpoint_maps_to_parseable_filename():
+    """upet parses (model, size, version) out of the filename — every pinned
+    entry must render to the {model}-{size}-v{version}.ckpt shape."""
+    env = _load_env_module()
+    for upstream in env.CHECKPOINTS.values():
+        model, version = upstream.split("@", 1)
+        filename = f"{model}-v{version}.ckpt"
+        assert filename.endswith(f"-v{version}.ckpt")
+        assert "@" not in filename