feat: add safe distributed checkpoint intake and preflight

This commit is contained in:
Joseph Magly
2026-09-04 19:43:53 -04:00
parent 5cc43c6e52
commit 985c9e9363
108 changed files with 21726 additions and 92 deletions
+70
View File
@@ -9,6 +9,7 @@ from types import SimpleNamespace
from unittest.mock import MagicMock, Mock
import pytest
import torch
from obliteratus import cli
@@ -21,6 +22,7 @@ def ns(**values):
("argv", "target"),
[
(["gpu-calc", "--params", "1"], "_cmd_gpu_calc"),
(["checkpoint", "inspect", "local/checkpoint"], "_cmd_checkpoint"),
(["run", "config.yml"], "_cmd_run"),
(["interactive"], "_cmd_interactive"),
(["models"], "_cmd_models"),
@@ -144,6 +146,74 @@ def test_version_is_stable_and_does_not_dispatch(capsys):
assert capsys.readouterr().out.endswith(f"obliteratus {__version__}\n")
def test_checkpoint_inspect_json_is_machine_readable_without_banner(tmp_path, capsys):
from safetensors.torch import save_file
save_file({"weight": torch.ones(1)}, tmp_path / "model.safetensors")
cli.main(["checkpoint", "inspect", str(tmp_path), "--json"])
payload = json.loads(capsys.readouterr().out)
assert payload["schema_id"] == "obliteratus.checkpoint-descriptor"
assert payload["primary_format"] == "hf_safetensors"
assert payload["safety"]["inspection_level"] == "safe_structure"
def test_checkpoint_inspect_boundary_error_is_stable_json(tmp_path, capsys):
target = tmp_path / "target"
target.write_bytes(b"payload")
(tmp_path / "model.safetensors").symlink_to(target.name)
with pytest.raises(SystemExit) as caught:
cli.main(["checkpoint", "inspect", str(tmp_path), "--json"])
assert caught.value.code == 2
payload = json.loads(capsys.readouterr().out)
assert payload["code"] == "DCI_SOURCE_BOUNDARY_VIOLATION"
assert payload["detail"] == "source_symlink"
assert str(tmp_path) not in json.dumps(payload)
def test_checkpoint_inspect_human_output_honors_explicit_limits(tmp_path, capsys):
cli.main(
[
"checkpoint",
"inspect",
str(tmp_path),
"--max-files",
"10",
"--max-total-bytes",
"1024",
"--max-json-bytes",
"512",
"--max-header-bytes",
"256",
]
)
output = capsys.readouterr().out
assert "format" in output
assert "confidence" in output
assert "unknown" in output
assert "blocked" in output
assert "DCI_UNSUPPORTED_FORMAT_OR_VERSION" in output
def test_checkpoint_inspect_human_boundary_error_is_actionable(tmp_path, capsys):
target = tmp_path / "target"
target.write_bytes(b"payload")
(tmp_path / "model.safetensors").symlink_to(target.name)
with pytest.raises(SystemExit) as caught:
cli.main(["checkpoint", "inspect", str(tmp_path)])
assert caught.value.code == 2
output = capsys.readouterr().out
assert "DCI_SOURCE_BOUNDARY_VIOLATION" in output
assert "source_symlink" in output
assert "Repair the immutable local source boundary" in output
def test_gpu_selection_contract(monkeypatch):
monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
cli._apply_gpu_selection(ns(gpus=None, remote=None))