mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-17 16:37:30 +02:00
Harden remote execution contracts and tests
This commit is contained in:
@@ -168,7 +168,8 @@ def test_run_local_overrides_and_remote_config(monkeypatch, tmp_path):
|
||||
|
||||
config.remote = ns(
|
||||
host="host", user="user", port=2200, ssh_key="key", remote_dir="/work",
|
||||
python="python", sync_results=True, gpus="0",
|
||||
known_hosts_file="known_hosts", install_timeout=120, python="python",
|
||||
sync_results=True, gpus="0", install_source="obliteratus==0.1.2",
|
||||
)
|
||||
runner = MagicMock()
|
||||
runner.run_config.return_value = "/local/results"
|
||||
@@ -347,6 +348,15 @@ def test_remote_runner_factory_and_commands(monkeypatch):
|
||||
monkeypatch.setattr(obliteratus.remote, "RemoteRunner", Mock(return_value=runner))
|
||||
args = _remote_args()
|
||||
assert cli._make_remote_runner(args) is runner
|
||||
obliteratus.remote.RemoteConfig.from_cli_args.assert_called_once_with(
|
||||
"user@host",
|
||||
port=22,
|
||||
ssh_key=None,
|
||||
remote_dir="/work",
|
||||
python="python3",
|
||||
sync_results=True,
|
||||
gpus="0",
|
||||
)
|
||||
|
||||
monkeypatch.setattr(cli, "_make_remote_runner", lambda _args: runner)
|
||||
runner.run_obliterate.return_value = "results"
|
||||
@@ -359,6 +369,25 @@ def test_remote_runner_factory_and_commands(monkeypatch):
|
||||
cli._cmd_remote_abliterate(args)
|
||||
assert "refusal_max_tokens" not in runner.run_obliterate.call_args.kwargs
|
||||
|
||||
for name in (
|
||||
"quantization", "n_directions", "direction_method", "regularization",
|
||||
"refinement_passes", "min_layer_fraction", "max_layer_fraction",
|
||||
"harmless_pc_count", "shield_concept_count", "shield_ridge",
|
||||
"shield_residualize", "shield_layer_penalty", "projection_target",
|
||||
"projection_row_fraction", "verify_sample_size", "refusal_max_tokens",
|
||||
):
|
||||
setattr(args, name, None)
|
||||
args.large_model = False
|
||||
runner.run_obliterate.reset_mock()
|
||||
cli._cmd_remote_abliterate(args)
|
||||
assert runner.run_obliterate.call_args.kwargs == {
|
||||
"model": "model",
|
||||
"local_output_dir": "out",
|
||||
"method": "advanced",
|
||||
"device": "cuda",
|
||||
"dtype": "float16",
|
||||
}
|
||||
|
||||
runner.run_config.return_value = "results"
|
||||
cli._cmd_remote_run(args)
|
||||
runner.run_tourney.return_value = "results"
|
||||
@@ -375,6 +404,20 @@ def test_remote_runner_factory_and_commands(monkeypatch):
|
||||
cli._cmd_remote_tourney(args)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("port", ["0", "65536", "not-a-port"])
|
||||
def test_remote_cli_rejects_invalid_ssh_ports_before_dispatch(port):
|
||||
with pytest.raises(SystemExit) as exc:
|
||||
cli.main(["run", "config.yml", "--remote", "host", "--ssh-port", port])
|
||||
assert exc.value.code == 2
|
||||
|
||||
|
||||
def test_remote_cli_accepts_valid_ssh_port_and_dispatches(monkeypatch):
|
||||
dispatch = Mock()
|
||||
monkeypatch.setattr(cli, "_cmd_remote_run", dispatch)
|
||||
cli.main(["run", "config.yml", "--remote", "host", "--ssh-port", "2222"])
|
||||
assert dispatch.call_args.args[0].ssh_port == 2222
|
||||
|
||||
|
||||
def test_abliterate_pipeline_callbacks_residue_and_contribution(monkeypatch, tmp_path):
|
||||
import obliteratus.abliterate
|
||||
import obliteratus.community
|
||||
|
||||
@@ -4,6 +4,7 @@ from __future__ import annotations
|
||||
|
||||
|
||||
import yaml
|
||||
import pytest
|
||||
|
||||
from obliteratus.config import StudyConfig
|
||||
|
||||
@@ -71,3 +72,39 @@ class TestStudyConfig:
|
||||
assert config.model.quantization == "4bit"
|
||||
assert config.model.num_labels == 7
|
||||
assert StudyConfig.from_dict(config.to_dict()).model == config.model
|
||||
|
||||
def test_remote_config_roundtrip_preserves_security_and_execution_settings(self):
|
||||
raw = {
|
||||
**SAMPLE_CONFIG,
|
||||
"remote": {
|
||||
"host": "compute.example",
|
||||
"user": "runner",
|
||||
"port": 2222,
|
||||
"ssh_key": "/keys/id",
|
||||
"known_hosts_file": "/keys/known_hosts",
|
||||
"remote_dir": "/srv/obliteratus",
|
||||
"install_timeout": 120,
|
||||
"python": "/opt/python",
|
||||
"sync_results": False,
|
||||
"gpus": "02, 0",
|
||||
"install_source": "obliteratus==0.1.2",
|
||||
},
|
||||
}
|
||||
config = StudyConfig.from_dict(raw)
|
||||
assert config.remote is not None
|
||||
assert config.remote.gpus == "2,0"
|
||||
assert StudyConfig.from_dict(config.to_dict()) == config
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"remote",
|
||||
[
|
||||
{"host": ""},
|
||||
{"host": "host", "port": 0},
|
||||
{"host": "host", "remote_dir": "relative"},
|
||||
{"host": "host", "gpus": "0; injected"},
|
||||
{"host": "host", "install_timeout": 0},
|
||||
],
|
||||
)
|
||||
def test_remote_config_rejects_invalid_public_values(self, remote):
|
||||
with pytest.raises(ValueError):
|
||||
StudyConfig.from_dict({**SAMPLE_CONFIG, "remote": remote})
|
||||
|
||||
@@ -2,6 +2,15 @@
|
||||
|
||||
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
|
||||
|
||||
|
||||
@@ -34,3 +43,375 @@ def test_remote_config_accepts_versioned_known_hosts_setting():
|
||||
)
|
||||
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
|
||||
|
||||
@@ -0,0 +1,180 @@
|
||||
"""Mutation-safe contracts for remote input and shell construction."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import shlex
|
||||
|
||||
import pytest
|
||||
|
||||
from obliteratus.remote_contracts import (
|
||||
normalize_gpu_selection,
|
||||
parse_remote_target,
|
||||
remote_python_command,
|
||||
remote_scp_spec,
|
||||
validate_remote_settings,
|
||||
)
|
||||
|
||||
|
||||
VALID_REMOTE_SETTINGS = {
|
||||
"host": "compute.example",
|
||||
"user": "runner",
|
||||
"port": 22,
|
||||
"remote_dir": "/tmp/obliteratus",
|
||||
"python": "python3",
|
||||
"gpus": None,
|
||||
"install_source": "obliteratus",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("raw", "expected"),
|
||||
[(None, None), ("all", "all"), ("ALL", "all"), ("0", "0"), (" 02, 0,11 ", "2,0,11")],
|
||||
)
|
||||
def test_gpu_selection_normalizes_public_values(raw, expected):
|
||||
assert normalize_gpu_selection(raw) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("raw", "message"),
|
||||
[
|
||||
("", "remote gpus must be a non-empty string"),
|
||||
("\x7f", "remote gpus may not contain control characters"),
|
||||
("0,,1", "remote gpus must be 'all' or comma-separated non-negative integers"),
|
||||
("-1", "remote gpus must be 'all' or comma-separated non-negative integers"),
|
||||
(
|
||||
"0; touch /tmp/pwned",
|
||||
"remote gpus must be 'all' or comma-separated non-negative integers",
|
||||
),
|
||||
("gpu0", "remote gpus must be 'all' or comma-separated non-negative integers"),
|
||||
("0\n1", "remote gpus may not contain control characters"),
|
||||
],
|
||||
)
|
||||
def test_gpu_selection_rejects_non_device_input(raw, message):
|
||||
with pytest.raises(ValueError) as error:
|
||||
normalize_gpu_selection(raw)
|
||||
assert str(error.value) == message
|
||||
|
||||
|
||||
def test_remote_target_defaults_user_and_parses_single_at_sign():
|
||||
assert parse_remote_target("compute.example") == ("root", "compute.example")
|
||||
assert parse_remote_target("runner@compute.example") == ("runner", "compute.example")
|
||||
assert parse_remote_target(" runner@compute.example ") == ("runner", "compute.example")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("target", "message"),
|
||||
[
|
||||
("", "remote target must be a non-empty string"),
|
||||
("\x7f", "remote target may not contain control characters"),
|
||||
("runner@", "remote host must be a non-empty string"),
|
||||
("@host", "remote user must be a non-empty string"),
|
||||
("bad user@host", "remote user contains unsupported characters"),
|
||||
(
|
||||
"runner@-oProxyCommand=evil",
|
||||
"remote host must be a host name or address without user or options",
|
||||
),
|
||||
("runner@host name", "remote host must be a host name or address without user or options"),
|
||||
("a@b@host", "remote target may contain at most one user separator"),
|
||||
],
|
||||
)
|
||||
def test_remote_target_rejects_empty_or_option_like_components(target, message):
|
||||
with pytest.raises(ValueError) as error:
|
||||
parse_remote_target(target)
|
||||
assert str(error.value) == message
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("override", "message"),
|
||||
[
|
||||
({"host": ""}, "remote host must be a non-empty string"),
|
||||
(
|
||||
{"host": "bad@host"},
|
||||
"remote host must be a host name or address without user or options",
|
||||
),
|
||||
(
|
||||
{"host": "bad host"},
|
||||
"remote host must be a host name or address without user or options",
|
||||
),
|
||||
({"user": ""}, "remote user must be a non-empty string"),
|
||||
({"user": "-oProxy"}, "remote user contains unsupported characters"),
|
||||
({"user": "bad user"}, "remote user contains unsupported characters"),
|
||||
({"port": 0}, "remote port must be an integer from 1 through 65535"),
|
||||
({"port": 65536}, "remote port must be an integer from 1 through 65535"),
|
||||
({"port": True}, "remote port must be an integer from 1 through 65535"),
|
||||
({"port": "22"}, "remote port must be an integer from 1 through 65535"),
|
||||
({"remote_dir": ""}, "remote directory must be a non-empty string"),
|
||||
({"remote_dir": "relative"}, "remote directory must be an absolute POSIX path"),
|
||||
({"python": ""}, "remote Python must be a non-empty string"),
|
||||
({"install_source": ""}, "remote install source must be a non-empty string"),
|
||||
(
|
||||
{"install_source": "bad\nsource"},
|
||||
"remote install source may not contain control characters",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_remote_settings_reject_invalid_boundaries(override, message):
|
||||
values = dict(VALID_REMOTE_SETTINGS)
|
||||
values.update(override)
|
||||
with pytest.raises(ValueError) as error:
|
||||
validate_remote_settings(**values)
|
||||
assert str(error.value) == message
|
||||
|
||||
|
||||
@pytest.mark.parametrize("port", [1, 65535])
|
||||
def test_remote_settings_accept_port_boundaries(port):
|
||||
values = dict(VALID_REMOTE_SETTINGS, port=port, gpus=" 02,0 ")
|
||||
assert validate_remote_settings(**values) == "2,0"
|
||||
|
||||
|
||||
def test_remote_python_command_preserves_each_untrusted_value_as_one_argument():
|
||||
command = remote_python_command(
|
||||
"/opt/python builds/current/python",
|
||||
["-m", "obliteratus", "run", "/tmp/a config.yml", "--preset", "x; echo injected"],
|
||||
gpus="02,0",
|
||||
)
|
||||
assert shlex.split(command) == [
|
||||
"env",
|
||||
"CUDA_VISIBLE_DEVICES=2,0",
|
||||
"/opt/python builds/current/python",
|
||||
"-m",
|
||||
"obliteratus",
|
||||
"run",
|
||||
"/tmp/a config.yml",
|
||||
"--preset",
|
||||
"x; echo injected",
|
||||
]
|
||||
|
||||
|
||||
def test_remote_python_command_omits_environment_for_all_devices():
|
||||
assert shlex.split(remote_python_command("python3", ["-V"], gpus="all")) == ["python3", "-V"]
|
||||
assert shlex.split(remote_python_command("python3", ["-V"])) == ["python3", "-V"]
|
||||
|
||||
|
||||
def test_remote_python_command_reports_invalid_python_contract():
|
||||
with pytest.raises(ValueError) as error:
|
||||
remote_python_command("", ["-V"])
|
||||
assert str(error.value) == "remote Python must be a non-empty string"
|
||||
|
||||
|
||||
def test_remote_scp_spec_quotes_remote_paths_and_directory_suffix():
|
||||
assert remote_scp_spec("runner@host", "/tmp/result file", directory=True) == (
|
||||
"runner@host:'/tmp/result file/'"
|
||||
)
|
||||
assert remote_scp_spec("runner@host", "/tmp/results/", directory=True) == (
|
||||
"runner@host:/tmp/results/"
|
||||
)
|
||||
assert remote_scp_spec("runner@host", "/tmp/config.yml") == "runner@host:/tmp/config.yml"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("target", "path", "message"),
|
||||
[
|
||||
("", "/tmp/config.yml", "SSH target must be a non-empty string"),
|
||||
("runner@host", "", "remote SCP path must be a non-empty string"),
|
||||
("runner@host", "/tmp/bad\x7fpath", "remote SCP path may not contain control characters"),
|
||||
],
|
||||
)
|
||||
def test_remote_scp_spec_reports_invalid_input_contract(target, path, message):
|
||||
with pytest.raises(ValueError) as error:
|
||||
remote_scp_spec(target, path)
|
||||
assert str(error.value) == message
|
||||
Reference in New Issue
Block a user