diff --git a/rootstock/operations.py b/rootstock/operations.py index 73def46..188723f 100644 --- a/rootstock/operations.py +++ b/rootstock/operations.py @@ -39,6 +39,7 @@ parse_checkpoints_dict, parse_clusters_list, parse_custom_checkpoint_ids, + verify_kwargs_for, ) from .exceptions import RootstockError from .layout import resolve_cache_root @@ -1211,12 +1212,17 @@ def fetch_checkpoint( manifest's ``last_error``). """ root = Path(root) - setup_kwargs = setup_kwargs or {} if cache_root is None: cache_root = resolve_cache_root(root) env_name = _resolve_built_env(root, checkpoint, cluster) + # No explicit kwargs → fall back to the env's declared VERIFY_KWARGS, + # exactly like the verify path: setup() runs the download here too, and a + # multi-head env (UMA's task, MACE-MH-1's head) raises without a head + # selection. Explicit kwargs always win. + setup_kwargs = setup_kwargs or verify_kwargs_for(root, env_name, checkpoint) + # Unlocked peek for idempotence; every write below loads fresh inside # its own transaction, so a racing writer costs at most a redundant # (idempotent) download, never a lost update. diff --git a/tests/commands/test_add_split.py b/tests/commands/test_add_split.py index ccd3c48..96ea03f 100644 --- a/tests/commands/test_add_split.py +++ b/tests/commands/test_add_split.py @@ -140,6 +140,71 @@ def test_fetch_refresh_knob(fake_root, refresh_calls, monkeypatch): assert len(refresh_calls) == 1 +_MULTIHEAD_ENV_SOURCE = '''\ +"""Multi-head env: setup() requires a head selection, no default.""" + +CHECKPOINTS = { + "uma-s-1p1": "uma-s-1p1", +} + +VERIFY_KWARGS = { + "uma-s-1p1": {"task": "omat"}, +} + + +def setup(checkpoint, device="cuda", task=None): + return None +''' + + +@pytest.fixture +def multihead_root(tmp_path: Path) -> Path: + """A root whose env declares VERIFY_KWARGS for its checkpoint.""" + root = tmp_path + env_dir = root / "envs" / "uma" + (env_dir / "bin").mkdir(parents=True) + (env_dir / "bin" / "python").touch() + (env_dir / "env_source.py").write_text(_MULTIHEAD_ENV_SOURCE) + + from rootstock.config import UserConfig + from rootstock.manifest import create_manifest, save_manifest + + cfg = UserConfig(name="t", email="t@t.t") + save_manifest(create_manifest(root, ["test"], cfg), root) + return root + + +def test_fetch_falls_back_to_verify_kwargs(multihead_root, refresh_calls, monkeypatch): + """No explicit kwargs → the env's VERIFY_KWARGS reach the download's + setup() call, same as the verify path (a multi-head setup() raises + without a head selection).""" + captured = {} + + def fake_download(root, env_name, checkpoint, setup_kwargs, **kw): + captured["setup_kwargs"] = setup_kwargs + return True, None + + monkeypatch.setattr(operations, "_run_download", fake_download) + + fetch_checkpoint(multihead_root, "uma-s-1p1") + + assert captured["setup_kwargs"] == {"task": "omat"} + + +def test_fetch_explicit_kwargs_beat_verify_kwargs(multihead_root, refresh_calls, monkeypatch): + captured = {} + + def fake_download(root, env_name, checkpoint, setup_kwargs, **kw): + captured["setup_kwargs"] = setup_kwargs + return True, None + + monkeypatch.setattr(operations, "_run_download", fake_download) + + fetch_checkpoint(multihead_root, "uma-s-1p1", setup_kwargs={"task": "omol"}) + + assert captured["setup_kwargs"] == {"task": "omol"} + + def test_fetch_fails_fast_when_env_not_built(tmp_path, refresh_calls): from rootstock.config import UserConfig from rootstock.manifest import create_manifest, save_manifest