diff --git a/src/hypertrain/challenge/app.py b/src/hypertrain/challenge/app.py index f32a80e..984e787 100644 --- a/src/hypertrain/challenge/app.py +++ b/src/hypertrain/challenge/app.py @@ -7,6 +7,7 @@ import hmac import json import os +import threading import time from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass @@ -278,6 +279,22 @@ async def read_model[M: BaseModel](request: Request, model: type[M]) -> M: raise ChallengeError(422, f"{where}: {first.get('msg', 'invalid')}") from None +def _lazy_role_launch(path: Path) -> role_dispatch.RoleLaunch: + """Load on first use: the supervisor canary boots with an empty /data and must answer.""" + lock, loaded = threading.Lock(), list[role_dispatch.RoleLaunch]() + + def launch(*args: Any) -> Any: + with lock: + if not loaded: + try: + loaded.append(role_dispatch.load(path)) + except (OSError, ValueError) as error: # pydantic ValidationError is a ValueError + raise ChallengeError(503, "ROLE_DISPATCH_UNAVAILABLE") from error + return loaded[0](*args) + + return launch + + def create_app( config: Config, *, @@ -310,7 +327,7 @@ def create_app( verify_beacon, objects or LocalFSStore(config.state_dir / "objects"), _role_launch=( - role_dispatch.load(config.role_dispatch_file) + _lazy_role_launch(config.role_dispatch_file) if config.role_dispatch_file is not None else None ), diff --git a/tests/challenge/test_role_dispatch_lazy.py b/tests/challenge/test_role_dispatch_lazy.py new file mode 100644 index 0000000..38c79d1 --- /dev/null +++ b/tests/challenge/test_role_dispatch_lazy.py @@ -0,0 +1,123 @@ +"""The supervisor canary boots the image with an empty /data: role dispatch loads on first use.""" + +from __future__ import annotations + +import json +import sys +import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from types import SimpleNamespace + +import pytest +from fastapi.testclient import TestClient + +from hypertrain.challenge import app as app_module +from hypertrain.challenge.store import ChallengeError +from hypertrain.gpu_ops import role_dispatch + +from .conftest import OWNER, make_config + +TREE = Path(__file__).parents[2] + + +def _app(tmp_path: Path, dispatch: Path) -> app_module.FastAPI: + # Canary: empty state dir, no secret files at all. + config = make_config(tmp_path / "state", tmp_path / "no-secrets") + config.role_dispatch_file = dispatch + return app_module.create_app(config) + + +def _write_valid(path: Path, tmp_path: Path) -> None: + path.write_text( + json.dumps( + { + "tree": str(TREE), + "run_dir": str(tmp_path / "lifecycle"), + "owner_key_file": str(tmp_path / "owner"), + "roles": {OWNER.ss58: "h1"}, + "run_id": "a" * 64, + } + ) + ) + + +def _foreign(launch: role_dispatch.RoleLaunch, tmp_path: Path) -> None: + # Foreign run: the real factory refuses before any I/O, proving the file was loaded. + with pytest.raises(RuntimeError, match="not the authorized run"): + launch("audit", SimpleNamespace(run_id="b" * 64), tmp_path, {}) # type: ignore[arg-type] + + +def test_boots_and_answers_version_without_dispatch_file(tmp_path: Path) -> None: + with TestClient(_app(tmp_path, tmp_path / "missing" / "dispatch.json")) as client: + response = client.get("/version") + assert response.status_code == 200 + assert response.json()["slug"] == "hypertrain" + + +def test_missing_then_fixed_file_without_restart(tmp_path: Path) -> None: + dispatch = tmp_path / "dispatch.json" + launch = _app(tmp_path, dispatch).state.store._role_launch + assert launch is not None + for broken in (None, "{not json", json.dumps({"tree": str(TREE), "extra": 1})): + if broken is not None: + dispatch.write_text(broken) + with pytest.raises(ChallengeError) as error: + launch("audit", SimpleNamespace(run_id="a" * 64), tmp_path, {}) + assert (error.value.status, error.value.detail) == (503, "ROLE_DISPATCH_UNAVAILABLE") + _write_valid(dispatch, tmp_path) + _foreign(launch, tmp_path) + + +def test_valid_file_loaded_once_same_behaviour( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + dispatch = tmp_path / "dispatch.json" + _write_valid(dispatch, tmp_path) + calls: list[Path] = [] + real = role_dispatch.load + + def spy(path: Path) -> role_dispatch.RoleLaunch: + calls.append(path) + return real(path) + + monkeypatch.setattr(role_dispatch, "load", spy) + monkeypatch.setattr(sys, "path", [p for p in sys.path if p != str(TREE / "scripts")]) + launch = _app(tmp_path, dispatch).state.store._role_launch + assert launch is not None + _foreign(launch, tmp_path) + assert sys.path[0] == str(TREE / "scripts") + _foreign(launch, tmp_path) + _foreign(real(dispatch), tmp_path) # same refusal as the eager loader + assert calls == [dispatch] + + +def test_concurrent_first_calls_load_once(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + dispatch = tmp_path / "dispatch.json" + _write_valid(dispatch, tmp_path) + calls: list[Path] = [] + real = role_dispatch.load + + def slow(path: Path) -> role_dispatch.RoleLaunch: + calls.append(path) + time.sleep(0.05) # widen the window: without the lock every thread loads + return real(path) + + monkeypatch.setattr(role_dispatch, "load", slow) + launch = _app(tmp_path, dispatch).state.store._role_launch + assert launch is not None + with ThreadPoolExecutor(16) as pool: + for future in [pool.submit(_foreign, launch, tmp_path) for _ in range(16)]: + future.result() + assert len(calls) == 1 + + +def test_unexpected_load_error_propagates(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + def broken(path: Path) -> role_dispatch.RoleLaunch: + raise RuntimeError("not a missing-file problem") + + monkeypatch.setattr(role_dispatch, "load", broken) + launch = _app(tmp_path, tmp_path / "dispatch.json").state.store._role_launch + assert launch is not None + with pytest.raises(RuntimeError, match="not a missing-file problem"): + launch("audit", SimpleNamespace(run_id="a" * 64), tmp_path, {})