feat: add lazy pi05 policy lifecycle

model_loader.py

@@ -0,0 +1,101 @@
+"""Thread-safe lazy lifecycle for the heavyweight π₀.₅ UR policy."""
+
+from __future__ import annotations
+
+import gc
+from pathlib import Path
+import sys
+import threading
+from collections.abc import Callable
+
+from artifacts import (
+    download_checkpoint,
+    normalize_checkpoint_path,
+    normalize_model_id,
+)
+
+
+class ModelUnavailableError(RuntimeError):
+    """Raised when the requested policy cannot be initialized."""
+
+
+def _release_gpu_memory() -> None:
+    gc.collect()
+    try:
+        import torch
+
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+    except ImportError:
+        pass
+
+
+class ModelManager:
+    def __init__(self, loader: Callable[[str, str], object] | None = None):
+        self._loader = loader or self._load_default
+        self._lock = threading.Lock()
+        self._value = None
+        self._active_key: tuple[str, str] | None = None
+        self._error: str | None = None
+
+    @property
+    def active_key(self) -> tuple[str, str] | None:
+        return self._active_key
+
+    @property
+    def health_message(self) -> str:
+        if self._value is not None and self._active_key is not None:
+            model_id, checkpoint_path = self._active_key
+            return f"Model ready: {model_id}/{checkpoint_path}."
+        if self._error:
+            return f"Model unavailable: {self._error}"
+        return "Model has not been loaded yet."
+
+    def get(self, model_id: str, checkpoint_path: str):
+        key = (
+            normalize_model_id(model_id),
+            normalize_checkpoint_path(checkpoint_path),
+        )
+        if self._value is not None and self._active_key == key:
+            return self._value
+        with self._lock:
+            if self._value is not None and self._active_key == key:
+                return self._value
+            if self._value is not None:
+                self._value = None
+                self._active_key = None
+                _release_gpu_memory()
+            self._error = None
+            try:
+                value = self._loader(*key)
+            except Exception as exc:
+                _release_gpu_memory()
+                detail = str(exc) or exc.__class__.__name__
+                self._error = f"{key[0]}/{key[1]}: {detail}"
+                raise ModelUnavailableError(self._error) from exc
+            self._value = value
+            self._active_key = key
+            return value
+
+    @staticmethod
+    def _load_default(model_id: str, checkpoint_path: str):
+        import torch
+
+        if not torch.cuda.is_available():
+            raise RuntimeError("CUDA GPU is required for π₀.₅ inference")
+        runtime = str(Path(__file__).resolve().parent / "openpi_runtime")
+        if runtime not in sys.path:
+            sys.path.insert(0, runtime)
+        from openpi.policies import policy_config
+        from openpi.training import config as openpi_config
+
+        paths = download_checkpoint(model_id, checkpoint_path)
+        config = openpi_config.get_config("pi05_ur_demo_state")
+        return policy_config.create_trained_policy(
+            config,
+            paths.checkpoint,
+            pytorch_device="cuda",
+        )
+
+
+MODEL_MANAGER = ModelManager()

tests/test_model_loader.py

@@ -0,0 +1,44 @@
+import unittest
+
+
+class ModelManagerTests(unittest.TestCase):
+    def test_manager_caches_same_key_and_replaces_changed_key(self):
+        from model_loader import ModelManager
+
+        calls = []
+        manager = ModelManager(
+            loader=lambda model, path: calls.append((model, path)) or object()
+        )
+        first = manager.get("owner/model", "a")
+        self.assertIs(manager.get("owner/model", "a"), first)
+        second = manager.get("owner/model", "b")
+        self.assertIsNot(second, first)
+        self.assertEqual(calls, [("owner/model", "a"), ("owner/model", "b")])
+
+    def test_failed_load_is_reported_and_can_retry(self):
+        from model_loader import ModelManager, ModelUnavailableError
+
+        attempts = 0
+
+        def loader(*_):
+            nonlocal attempts
+            attempts += 1
+            raise RuntimeError("bad checkpoint")
+
+        manager = ModelManager(loader=loader)
+        for _ in range(2):
+            with self.assertRaisesRegex(ModelUnavailableError, "bad checkpoint"):
+                manager.get("owner/model", "a")
+        self.assertEqual(attempts, 2)
+        self.assertIn("bad checkpoint", manager.health_message)
+
+    def test_invalid_identity_is_rejected_before_loader(self):
+        from model_loader import ModelManager
+
+        manager = ModelManager(loader=lambda *_: self.fail("loader was called"))
+        with self.assertRaises(ValueError):
+            manager.get("", "checkpoint")
+
+
+if __name__ == "__main__":
+    unittest.main()