mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-17 16:37:30 +02:00
418 lines
16 KiB
Python
418 lines
16 KiB
Python
"""Offline contracts for strict SSH command construction."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import shlex
|
|
import subprocess
|
|
from pathlib import Path
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
import yaml
|
|
|
|
from obliteratus import __version__
|
|
from obliteratus.remote import RemoteConfig, RemoteRunner
|
|
|
|
|
|
def test_remote_commands_require_strict_host_key_verification(tmp_path):
|
|
key = tmp_path / "key"
|
|
known_hosts = tmp_path / "known_hosts"
|
|
config = RemoteConfig(
|
|
host="compute.example",
|
|
user="runner",
|
|
port=2222,
|
|
ssh_key=str(key),
|
|
known_hosts_file=str(known_hosts),
|
|
)
|
|
runner = RemoteRunner(config)
|
|
for command in (runner._ssh_base_cmd(), runner._scp_base_cmd()):
|
|
assert "StrictHostKeyChecking=yes" in command
|
|
assert f"UserKnownHostsFile={known_hosts}" in command
|
|
assert "StrictHostKeyChecking=no" not in command
|
|
assert str(key) in command
|
|
|
|
|
|
def test_remote_config_accepts_versioned_known_hosts_setting():
|
|
config = RemoteConfig.from_dict(
|
|
{
|
|
"host": "compute.example",
|
|
"user": "runner",
|
|
"known_hosts_file": "/secure/known_hosts",
|
|
"unknown": "ignored",
|
|
}
|
|
)
|
|
assert config.known_hosts_file == "/secure/known_hosts"
|
|
assert config.ssh_target == "runner@compute.example"
|
|
|
|
|
|
def _completed(returncode=0, stdout="", stderr=""):
|
|
return subprocess.CompletedProcess([], returncode, stdout=stdout, stderr=stderr)
|
|
|
|
|
|
def test_remote_config_cli_defaults_and_optional_command_flags_are_shell_safe():
|
|
config = RemoteConfig.from_cli_args(
|
|
"runner@compute.example",
|
|
python="/opt/python builds/python",
|
|
remote_dir="/srv/remote work",
|
|
gpus="02, 0",
|
|
)
|
|
runner = RemoteRunner(config, on_log=lambda _message: None)
|
|
command = runner.build_obliterate_command(
|
|
"org/model; touch /tmp/injected",
|
|
method="advanced",
|
|
device="cuda",
|
|
dtype="float16",
|
|
quantization="4bit",
|
|
n_directions=3,
|
|
direction_method="svd",
|
|
regularization=0.2,
|
|
refinement_passes=2,
|
|
large_model=True,
|
|
verify_sample_size=7,
|
|
min_layer_fraction=0.1,
|
|
max_layer_fraction=0.9,
|
|
harmless_pc_count=4,
|
|
shield_concept_count=5,
|
|
shield_ridge=0.05,
|
|
shield_residualize=True,
|
|
shield_layer_penalty=0.3,
|
|
projection_target="attention",
|
|
projection_row_fraction=0.25,
|
|
refusal_max_tokens=256,
|
|
)
|
|
tokens = shlex.split(command)
|
|
assert tokens[:4] == ["env", "CUDA_VISIBLE_DEVICES=2,0", "/opt/python builds/python", "-m"]
|
|
assert tokens[4:7] == ["obliteratus", "obliterate", "org/model; touch /tmp/injected"]
|
|
expected = {
|
|
"--quantization": "4bit",
|
|
"--n-directions": "3",
|
|
"--direction-method": "svd",
|
|
"--regularization": "0.2",
|
|
"--refinement-passes": "2",
|
|
"--verify-sample-size": "7",
|
|
"--min-layer-fraction": "0.1",
|
|
"--max-layer-fraction": "0.9",
|
|
"--harmless-pc-count": "4",
|
|
"--shield-concept-count": "5",
|
|
"--shield-ridge": "0.05",
|
|
"--shield-layer-penalty": "0.3",
|
|
"--projection-target": "attention",
|
|
"--projection-row-fraction": "0.25",
|
|
"--refusal-max-tokens": "256",
|
|
}
|
|
for flag, value in expected.items():
|
|
assert tokens[tokens.index(flag) + 1] == value
|
|
assert "--large-model" in tokens
|
|
assert "--shield-residualize" in tokens
|
|
|
|
|
|
def test_remote_config_rejects_invalid_install_timeout():
|
|
with pytest.raises(ValueError, match="install timeout"):
|
|
RemoteConfig(host="host", install_timeout=0)
|
|
|
|
|
|
def test_run_and_tourney_commands_quote_all_public_values():
|
|
runner = RemoteRunner(
|
|
RemoteConfig(host="host", user="runner", gpus="all"),
|
|
on_log=lambda _message: None,
|
|
)
|
|
assert shlex.split(
|
|
runner.build_run_command("/tmp/a config.yml", output_dir="/tmp/out dir", preset="x; echo bad")
|
|
) == [
|
|
"python3", "-m", "obliteratus", "run", "/tmp/a config.yml",
|
|
"--output-dir", "/tmp/out dir", "--preset", "x; echo bad",
|
|
]
|
|
tokens = shlex.split(
|
|
runner.build_tourney_command(
|
|
"org/model", output_dir="/tmp/out dir", quantization="8bit",
|
|
hub_org="org; bad", hub_repo="org/repo bad", methods=["basic", "advanced"],
|
|
dataset="data; bad",
|
|
)
|
|
)
|
|
assert tokens[tokens.index("--hub-org") + 1] == "org; bad"
|
|
assert tokens[tokens.index("--hub-repo") + 1] == "org/repo bad"
|
|
assert tokens[tokens.index("--dataset") + 1] == "data; bad"
|
|
assert tokens[-3:] == ["--methods", "basic", "advanced"]
|
|
assert shlex.split(runner.build_run_command("config.yml", preset="quick")) == [
|
|
"python3", "-m", "obliteratus", "run", "config.yml", "--preset", "quick",
|
|
]
|
|
|
|
|
|
def test_run_ssh_non_stream_uses_argv_and_propagates_timeout(monkeypatch):
|
|
observed = {}
|
|
|
|
def fake_run(command, **kwargs):
|
|
observed.update(command=command, kwargs=kwargs)
|
|
return _completed(stdout="ok\n")
|
|
|
|
monkeypatch.setattr(subprocess, "run", fake_run)
|
|
runner = RemoteRunner(RemoteConfig(host="host"), on_log=lambda _message: None)
|
|
result = runner.run_ssh("printf '%s' 'safe value'", timeout=9)
|
|
assert result.stdout == "ok\n"
|
|
assert observed["command"][-1] == "printf '%s' 'safe value'"
|
|
assert observed["kwargs"] == {"capture_output": True, "text": True, "timeout": 9}
|
|
|
|
|
|
class _StreamProcess:
|
|
def __init__(self, stdout):
|
|
self.stdout = stdout
|
|
self.returncode = 0
|
|
self.killed = False
|
|
self.wait_calls = 0
|
|
|
|
def wait(self, timeout=None):
|
|
self.wait_calls += 1
|
|
return self.returncode
|
|
|
|
def kill(self):
|
|
self.killed = True
|
|
|
|
|
|
def test_run_ssh_streams_lines_and_returns_exit_status(monkeypatch):
|
|
process = _StreamProcess(["first\n", "second\n"])
|
|
monkeypatch.setattr(subprocess, "Popen", lambda *_args, **_kwargs: process)
|
|
logs = []
|
|
runner = RemoteRunner(RemoteConfig(host="host"), on_log=logs.append)
|
|
assert runner.run_ssh("command", stream=True, timeout=4) == 0
|
|
assert logs == ["first", "second"]
|
|
assert process.wait_calls == 1
|
|
|
|
|
|
def test_run_ssh_timeout_kills_and_reaps_process(monkeypatch):
|
|
process = _StreamProcess([])
|
|
|
|
def wait(timeout=None):
|
|
process.wait_calls += 1
|
|
if process.wait_calls == 1:
|
|
raise subprocess.TimeoutExpired("ssh", timeout)
|
|
return 0
|
|
|
|
process.wait = wait
|
|
monkeypatch.setattr(subprocess, "Popen", lambda *_args, **_kwargs: process)
|
|
logs = []
|
|
runner = RemoteRunner(RemoteConfig(host="host"), on_log=logs.append)
|
|
assert runner.run_ssh("command", stream=True, timeout=4) == 124
|
|
assert process.killed is True
|
|
assert process.wait_calls == 2
|
|
assert any("timed out" in line for line in logs)
|
|
|
|
|
|
class _CancelledOutput:
|
|
def __iter__(self):
|
|
return self
|
|
|
|
def __next__(self):
|
|
raise KeyboardInterrupt
|
|
|
|
|
|
def test_run_ssh_cancellation_kills_reaps_and_reraises(monkeypatch):
|
|
process = _StreamProcess(_CancelledOutput())
|
|
monkeypatch.setattr(subprocess, "Popen", lambda *_args, **_kwargs: process)
|
|
logs = []
|
|
runner = RemoteRunner(RemoteConfig(host="host"), on_log=logs.append)
|
|
with pytest.raises(KeyboardInterrupt):
|
|
runner.run_ssh("command", stream=True)
|
|
assert process.killed is True
|
|
assert process.wait_calls == 1
|
|
assert any("cancelled" in line for line in logs)
|
|
|
|
|
|
def test_run_ssh_rejects_malformed_process_without_stdout(monkeypatch):
|
|
process = _StreamProcess(None)
|
|
monkeypatch.setattr(subprocess, "Popen", lambda *_args, **_kwargs: process)
|
|
runner = RemoteRunner(RemoteConfig(host="host"), on_log=lambda _message: None)
|
|
with pytest.raises(RuntimeError, match="stdout"):
|
|
runner.run_ssh("command", stream=True)
|
|
assert process.killed is True
|
|
assert process.wait_calls == 1
|
|
|
|
|
|
def test_connection_and_gpu_probes_cover_success_and_malformed_responses():
|
|
logs = []
|
|
runner = RemoteRunner(RemoteConfig(host="host", gpus="1"), on_log=logs.append)
|
|
runner.run_ssh = Mock(return_value=_completed(stdout="ok\n"))
|
|
assert runner.check_connection() is True
|
|
runner.run_ssh.return_value = 0
|
|
assert runner.check_connection() is False
|
|
|
|
runner.run_ssh.return_value = _completed(stdout="0, A100, 80 GiB, 70 GiB\n1, A100, 80 GiB, 60 GiB\n")
|
|
assert runner.check_gpu().splitlines() == [
|
|
"0, A100, 80 GiB, 70 GiB",
|
|
"1, A100, 80 GiB, 60 GiB",
|
|
]
|
|
assert any("Selected GPUs: 1" in line for line in logs)
|
|
runner.run_ssh.return_value = _completed(returncode=1, stderr="missing")
|
|
assert runner.check_gpu() is None
|
|
|
|
|
|
def test_gpu_probe_reports_all_devices_when_unrestricted():
|
|
logs = []
|
|
runner = RemoteRunner(RemoteConfig(host="host", gpus="all"), on_log=logs.append)
|
|
runner.run_ssh = Mock(return_value=_completed(stdout="0, A100\n"))
|
|
assert runner.check_gpu() == "0, A100"
|
|
assert any("Using: all 1 GPUs" in line for line in logs)
|
|
|
|
|
|
def test_ensure_obliteratus_accepts_exact_version_or_installs_and_verifies():
|
|
logs = []
|
|
runner = RemoteRunner(RemoteConfig(host="host"), on_log=logs.append)
|
|
runner.run_ssh = Mock(return_value=_completed(stdout=f"{__version__}\n"))
|
|
assert runner.ensure_obliteratus() is True
|
|
assert runner.run_ssh.call_count == 1
|
|
|
|
runner.run_ssh = Mock(side_effect=[_completed(stdout="0.0.1\n"), 0, _completed(stdout=f"{__version__}\n")])
|
|
assert runner.ensure_obliteratus() is True
|
|
install_command = runner.run_ssh.call_args_list[1].args[0]
|
|
assert shlex.split(install_command)[-1] == "git+https://github.com/elder-plinius/OBLITERATUS.git"
|
|
|
|
|
|
def test_ensure_obliteratus_reports_install_and_verification_failures():
|
|
runner = RemoteRunner(RemoteConfig(host="host"), on_log=lambda _message: None)
|
|
runner.run_ssh = Mock(side_effect=[_completed(returncode=1), 1])
|
|
assert runner.ensure_obliteratus() is False
|
|
runner.run_ssh = Mock(side_effect=[_completed(stdout="old\n"), 0, _completed(stdout="still-old\n")])
|
|
assert runner.ensure_obliteratus() is False
|
|
|
|
|
|
def test_result_sync_creates_local_directory_and_quotes_remote_path(monkeypatch, tmp_path):
|
|
responses = iter([_completed(), _completed(returncode=1, stderr="denied")])
|
|
observed = []
|
|
|
|
def fake_run(command, **_kwargs):
|
|
observed.append(command)
|
|
return next(responses)
|
|
|
|
monkeypatch.setattr(subprocess, "run", fake_run)
|
|
runner = RemoteRunner(RemoteConfig(host="host", user="runner"), on_log=lambda _message: None)
|
|
local = tmp_path / "local results"
|
|
assert runner.sync_results_back("/tmp/remote results", str(local)) is True
|
|
assert local.is_dir()
|
|
assert observed[0][-2] == "runner@host:'/tmp/remote results/'"
|
|
assert runner.sync_results_back("/tmp/remote results", str(local)) is False
|
|
|
|
|
|
def test_upload_config_returns_remote_path_or_raises(monkeypatch, tmp_path):
|
|
config_path = tmp_path / "study config.yml"
|
|
config_path.write_text("model: fixture\nremote:\n host: compute.example\n")
|
|
responses = iter([_completed(), _completed(returncode=1, stderr="denied")])
|
|
observed = []
|
|
uploaded_payloads = []
|
|
|
|
def fake_run(command, **_kwargs):
|
|
observed.append(command)
|
|
uploaded_payloads.append(yaml.safe_load(Path(command[-2]).read_text()))
|
|
return next(responses)
|
|
|
|
monkeypatch.setattr(subprocess, "run", fake_run)
|
|
runner = RemoteRunner(
|
|
RemoteConfig(host="host", user="runner", remote_dir="/srv/remote work"),
|
|
on_log=lambda _message: None,
|
|
)
|
|
runner.run_ssh = Mock()
|
|
assert runner.upload_config(str(config_path)) == "/srv/remote work/config.yaml"
|
|
assert observed[0][-1] == "runner@host:'/srv/remote work/config.yaml'"
|
|
assert uploaded_payloads[0] == {"model": "fixture"}
|
|
with pytest.raises(RuntimeError, match="denied"):
|
|
runner.upload_config(str(config_path))
|
|
|
|
|
|
def test_upload_config_requires_yaml_mapping(monkeypatch, tmp_path):
|
|
config_path = tmp_path / "study.yml"
|
|
config_path.write_text("- not\n- a\n- mapping\n")
|
|
runner = RemoteRunner(RemoteConfig(host="host"), on_log=lambda _message: None)
|
|
runner.run_ssh = Mock()
|
|
monkeypatch.setattr(subprocess, "run", Mock())
|
|
with pytest.raises(ValueError, match="must contain a YAML mapping"):
|
|
runner.upload_config(str(config_path))
|
|
subprocess.run.assert_not_called()
|
|
|
|
|
|
def _prepared_runner(*, sync_results=True):
|
|
runner = RemoteRunner(
|
|
RemoteConfig(host="host", remote_dir="/srv/run", sync_results=sync_results),
|
|
on_log=lambda _message: None,
|
|
)
|
|
runner.check_connection = Mock(return_value=True)
|
|
runner.check_gpu = Mock(return_value="gpu")
|
|
runner.ensure_obliteratus = Mock(return_value=True)
|
|
return runner
|
|
|
|
|
|
def test_remote_obliterate_orchestration_success_failure_and_sync_paths():
|
|
runner = _prepared_runner(sync_results=False)
|
|
runner.run_ssh = Mock(side_effect=[_completed(), 0])
|
|
assert runner.run_obliterate("org/model") == "/srv/run/output/org_model"
|
|
|
|
runner = _prepared_runner()
|
|
runner.run_ssh = Mock(side_effect=[_completed(), 9])
|
|
assert runner.run_obliterate("org/model") is None
|
|
|
|
runner = _prepared_runner()
|
|
runner.run_ssh = Mock(side_effect=[_completed(), 0])
|
|
runner.sync_results_back = Mock(return_value=True)
|
|
assert runner.run_obliterate("org/model", local_output_dir="local") == "local"
|
|
runner.run_ssh = Mock(side_effect=[_completed(), 0])
|
|
runner.sync_results_back.return_value = False
|
|
assert runner.run_obliterate("org/model") is None
|
|
|
|
|
|
def test_remote_obliterate_stops_at_connection_or_install_failure():
|
|
runner = _prepared_runner()
|
|
runner.check_connection.return_value = False
|
|
assert runner.run_obliterate("model") is None
|
|
runner.check_connection.return_value = True
|
|
runner.ensure_obliteratus.return_value = False
|
|
assert runner.run_obliterate("model") is None
|
|
|
|
|
|
def test_remote_config_orchestration_success_failure_and_sync_paths():
|
|
runner = _prepared_runner(sync_results=False)
|
|
runner.upload_config = Mock(return_value="/srv/run/config.yaml")
|
|
runner.run_ssh = Mock(return_value=0)
|
|
assert runner.run_config("local.yml", preset="quick") == "/srv/run/results"
|
|
|
|
runner = _prepared_runner()
|
|
runner.upload_config = Mock(return_value="/srv/run/config.yaml")
|
|
runner.run_ssh = Mock(return_value=2)
|
|
assert runner.run_config("local.yml") is None
|
|
|
|
runner.run_ssh.return_value = 0
|
|
runner.sync_results_back = Mock(return_value=True)
|
|
assert runner.run_config("local.yml", local_output_dir="local") == "local"
|
|
runner.sync_results_back.return_value = False
|
|
assert runner.run_config("local.yml") is None
|
|
|
|
|
|
def test_remote_config_stops_at_connection_or_install_failure():
|
|
runner = _prepared_runner()
|
|
runner.check_connection.return_value = False
|
|
assert runner.run_config("local.yml") is None
|
|
runner.check_connection.return_value = True
|
|
runner.ensure_obliteratus.return_value = False
|
|
assert runner.run_config("local.yml") is None
|
|
|
|
|
|
def test_remote_tourney_orchestration_success_failure_and_sync_paths():
|
|
runner = _prepared_runner(sync_results=False)
|
|
runner.run_ssh = Mock(return_value=0)
|
|
assert runner.run_tourney("org/model") == "/srv/run/tourney/org_model"
|
|
|
|
runner = _prepared_runner()
|
|
runner.run_ssh = Mock(return_value=3)
|
|
assert runner.run_tourney("org/model") is None
|
|
|
|
runner.run_ssh.return_value = 0
|
|
runner.sync_results_back = Mock(return_value=True)
|
|
assert runner.run_tourney("org/model", local_output_dir="local") == "local"
|
|
runner.sync_results_back.return_value = False
|
|
assert runner.run_tourney("org/model") is None
|
|
|
|
|
|
def test_remote_tourney_stops_at_connection_or_install_failure():
|
|
runner = _prepared_runner()
|
|
runner.check_connection.return_value = False
|
|
assert runner.run_tourney("model") is None
|
|
runner.check_connection.return_value = True
|
|
runner.ensure_obliteratus.return_value = False
|
|
assert runner.run_tourney("model") is None
|