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
+52
View File
@@ -144,6 +144,58 @@ class TestCLIDispatch:
main(["interactive"])
mock_cmd.assert_called_once()
def test_distributed_preflight_is_an_explicit_separate_command(self, tmp_path):
profile = tmp_path / "profile.json"
profile.write_text("{}", encoding="utf-8")
with patch("obliteratus.cli._cmd_distributed") as mock_cmd:
main(["distributed", "preflight", str(profile), "--json"])
args_passed = mock_cmd.call_args.args[0]
assert args_passed.command == "distributed"
assert args_passed.distributed_command == "preflight"
assert args_passed.profile == profile
def test_ordinary_command_never_infers_distributed_mode(self, monkeypatch):
monkeypatch.setenv("WORLD_SIZE", "2")
monkeypatch.setenv("RANK", "0")
monkeypatch.setenv("MASTER_ADDR", "10.10.0.10")
with (
patch("obliteratus.cli._cmd_abliterate") as ordinary,
patch("obliteratus.cli._cmd_distributed") as distributed,
):
main(["obliterate", "fake/model"])
ordinary.assert_called_once()
distributed.assert_not_called()
@pytest.mark.parametrize("option", ["--token", "--password", "--api-key", "--remote"])
def test_distributed_preflight_rejects_unknown_options_without_echoing_value(
self, option, capsys
):
secret = "hf_private_value_that_must_not_appear"
with pytest.raises(SystemExit) as error:
main(["distributed", "preflight", "profile.json", option, secret])
assert error.value.code == 2
captured = capsys.readouterr()
assert secret not in captured.out
assert secret not in captured.err
@pytest.mark.parametrize(
"argv",
[
["--token", "{secret}", "distributed", "preflight", "profile.json"],
["distributed", "preflight", "profile.json", "--token={secret}"],
],
)
def test_distributed_preflight_secret_prescan_cannot_be_bypassed(
self, argv, capsys
):
secret = "hf_private_value_that_must_not_appear"
with pytest.raises(SystemExit) as error:
main([item.replace("{secret}", secret) for item in argv])
assert error.value.code == 2
captured = capsys.readouterr()
assert secret not in captured.out
assert secret not in captured.err
# 9. --contribute and --contribute-notes are accepted on obliterate
def test_contribute_flags_on_obliterate(self):
"""Verify --contribute and --contribute-notes are accepted args."""