feat: resolve pi05 checkpoint artifacts

artifacts.py

@@ -0,0 +1,64 @@
+"""Resolve and validate Hugging Face π₀.₅ UR checkpoint artifacts."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+import os
+from pathlib import Path, PurePosixPath
+
+
+DEFAULT_CHECKPOINT_PATH = "checkpoint"
+
+
+@dataclass(frozen=True)
+class ArtifactPaths:
+    checkpoint: Path
+
+
+def snapshot_download(**kwargs) -> str:
+    """Import Hugging Face Hub lazily so validation tests stay lightweight."""
+    from huggingface_hub import snapshot_download as download
+
+    return download(**kwargs)
+
+
+def normalize_model_id(value: str) -> str:
+    if not isinstance(value, str) or not value.strip():
+        raise ValueError("Hugging Face model ID is required")
+    result = value.strip()
+    if any(character.isspace() for character in result):
+        raise ValueError("Hugging Face model ID cannot contain whitespace")
+    return result
+
+
+def normalize_checkpoint_path(value: str) -> str:
+    if not isinstance(value, str) or not value.strip():
+        raise ValueError("checkpoint path is required")
+    candidate = value.strip().replace("\\", "/")
+    path = PurePosixPath(candidate)
+    if path.is_absolute() or ".." in path.parts or path == PurePosixPath("."):
+        raise ValueError("checkpoint path must be a relative path without parent traversal")
+    return path.as_posix()
+
+
+def resolve_model_id() -> str:
+    return os.getenv("PI05_MODEL_ID", "")
+
+
+def resolve_checkpoint_path() -> str:
+    return os.getenv("PI05_CHECKPOINT_PATH", DEFAULT_CHECKPOINT_PATH)
+
+
+def download_checkpoint(model_id: str, checkpoint_path: str) -> ArtifactPaths:
+    model_id = normalize_model_id(model_id)
+    relative = normalize_checkpoint_path(checkpoint_path)
+    root = Path(snapshot_download(repo_id=model_id))
+    checkpoint = root.joinpath(*PurePosixPath(relative).parts)
+    if not (checkpoint / "params").is_dir() and not (checkpoint / "model.safetensors").is_file():
+        raise FileNotFoundError(
+            f"checkpoint has neither params/ nor model.safetensors: {checkpoint}"
+        )
+    statistics = checkpoint / "assets/ur_demo/norm_stats.json"
+    if not statistics.is_file():
+        raise FileNotFoundError(f"UR normalization statistics not found: {statistics}")
+    return ArtifactPaths(checkpoint=checkpoint)

tests/test_artifacts.py

@@ -0,0 +1,38 @@
+from pathlib import Path
+import tempfile
+import unittest
+from unittest import mock
+
+
+class ArtifactTests(unittest.TestCase):
+    def test_checkpoint_path_rejects_absolute_parent_and_empty_paths(self):
+        import artifacts
+
+        for value in ("/tmp/model", "../model", "a/../../model", ""):
+            with self.subTest(value=value), self.assertRaises(ValueError):
+                artifacts.normalize_checkpoint_path(value)
+
+    def test_download_checkpoint_validates_required_layout(self):
+        import artifacts
+
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            checkpoint = root / "checkpoints/30000"
+            (checkpoint / "params").mkdir(parents=True)
+            (checkpoint / "assets/ur_demo").mkdir(parents=True)
+            (checkpoint / "assets/ur_demo/norm_stats.json").write_text("{}")
+            with mock.patch.object(artifacts, "snapshot_download", return_value=str(root)):
+                paths = artifacts.download_checkpoint("owner/model", "checkpoints/30000")
+            self.assertEqual(paths.checkpoint, checkpoint)
+
+    def test_download_checkpoint_reports_missing_weights(self):
+        import artifacts
+
+        with tempfile.TemporaryDirectory() as directory:
+            with mock.patch.object(artifacts, "snapshot_download", return_value=directory):
+                with self.assertRaisesRegex(FileNotFoundError, "params/|model.safetensors"):
+                    artifacts.download_checkpoint("owner/model", "checkpoint")
+
+
+if __name__ == "__main__":
+    unittest.main()