diff --git a/ci/test-quality-policy.json b/ci/test-quality-policy.json index 7984660..c5a1f4d 100644 --- a/ci/test-quality-policy.json +++ b/ci/test-quality-policy.json @@ -12,6 +12,7 @@ "critical_cpu_paths": [ "obliteratus/runtime_contracts.py", "obliteratus/persistence_contracts.py", + "obliteratus/remote_contracts.py", "obliteratus/device.py", "obliteratus/models/loader.py", "obliteratus/architecture_profiles.py", diff --git a/ci/test-risk-map.json b/ci/test-risk-map.json index 06f6949..027df86 100644 --- a/ci/test-risk-map.json +++ b/ci/test-risk-map.json @@ -98,6 +98,7 @@ "obliteratus/local_ui.py", "obliteratus/models_client.py", "obliteratus/remote.py", + "obliteratus/remote_contracts.py", "obliteratus/ui_watchtower.py", "obliteratus/watchtower.py" ], @@ -105,6 +106,7 @@ "tests/test_cli.py", "tests/test_cli_boundaries.py", "tests/test_remote_boundaries.py", + "tests/test_remote_contracts.py", "tests/conditional/test_network_services.py", "tests/conditional/test_operator_ui.py", "tests/conditional/test_remote_runtime.py" @@ -461,6 +463,13 @@ "risk": "SSH discovery, quoting, execution, cancellation, and result synchronization", "required_tests": ["tests/test_remote_boundaries.py", "tests/conditional/test_remote_runtime.py"], "conditional_gates": ["remote-execution"] + }, + { + "path": "obliteratus/remote_contracts.py", + "risk_class": "cpu-contract", + "risk": "remote target validation, GPU selection normalization, and shell-safe command construction", + "required_tests": ["tests/test_remote_contracts.py", "tests/test_remote_boundaries.py"], + "conditional_gates": [] } ] } diff --git a/obliteratus/cli.py b/obliteratus/cli.py index 8b33af6..2c5ad85 100644 --- a/obliteratus/cli.py +++ b/obliteratus/cli.py @@ -35,6 +35,17 @@ def _positive_int(value: str) -> int: return parsed +def _ssh_port(value: str) -> int: + """Parse the public SSH port boundary.""" + try: + parsed = int(value) + except ValueError as exc: + raise argparse.ArgumentTypeError("must be an integer from 1 through 65535") from exc + if not 1 <= parsed <= 65535: + raise argparse.ArgumentTypeError("must be an integer from 1 through 65535") + return parsed + + def _add_gpu_args(parser): """Add --gpus flag for multi-GPU control.""" gpu_group = parser.add_argument_group("GPU selection") @@ -60,7 +71,7 @@ def _add_remote_args(parser): help="Path to SSH private key (default: use SSH agent or ~/.ssh/id_rsa)", ) remote_group.add_argument( - "--ssh-port", type=int, default=22, + "--ssh-port", type=_ssh_port, default=22, help="SSH port on remote host (default: 22)", ) remote_group.add_argument( @@ -707,10 +718,13 @@ def _cmd_run(args): user=config.remote.user, port=config.remote.port, ssh_key=config.remote.ssh_key, + known_hosts_file=config.remote.known_hosts_file, remote_dir=config.remote.remote_dir, + install_timeout=config.remote.install_timeout, python=config.remote.python, sync_results=config.remote.sync_results, gpus=config.remote.gpus, + install_source=config.remote.install_source, ) runner = RemoteRunner(rc) result_path = runner.run_config( diff --git a/obliteratus/config.py b/obliteratus/config.py index 96214e1..1c5ad95 100644 --- a/obliteratus/config.py +++ b/obliteratus/config.py @@ -44,10 +44,32 @@ class RemoteConfig: user: str = "root" port: int = 22 ssh_key: str | None = None + known_hosts_file: str | None = None remote_dir: str = "/tmp/obliteratus_run" + install_timeout: int = 600 python: str = "python3" sync_results: bool = True gpus: str | None = None # comma-separated GPU IDs or "all" + install_source: str = "git+https://github.com/elder-plinius/OBLITERATUS.git" + + def __post_init__(self) -> None: + from obliteratus.remote_contracts import validate_remote_settings + + self.gpus = validate_remote_settings( + host=self.host, + user=self.user, + port=self.port, + remote_dir=self.remote_dir, + python=self.python, + gpus=self.gpus, + install_source=self.install_source, + ) + if ( + isinstance(self.install_timeout, bool) + or not isinstance(self.install_timeout, int) + or self.install_timeout <= 0 + ): + raise ValueError("remote install timeout must be a positive integer") @dataclass @@ -114,7 +136,7 @@ class StudyConfig: ) def to_dict(self) -> dict: - return { + result = { "model": { "name": self.model.name, "task": self.model.task, @@ -138,3 +160,18 @@ class StudyConfig: "max_length": self.max_length, "output_dir": self.output_dir, } + if self.remote is not None: + result["remote"] = { + "host": self.remote.host, + "user": self.remote.user, + "port": self.remote.port, + "ssh_key": self.remote.ssh_key, + "known_hosts_file": self.remote.known_hosts_file, + "remote_dir": self.remote.remote_dir, + "install_timeout": self.remote.install_timeout, + "python": self.remote.python, + "sync_results": self.remote.sync_results, + "gpus": self.remote.gpus, + "install_source": self.remote.install_source, + } + return result diff --git a/obliteratus/remote.py b/obliteratus/remote.py index 5dd8c5f..55a0d4f 100644 --- a/obliteratus/remote.py +++ b/obliteratus/remote.py @@ -20,13 +20,26 @@ Usage (YAML config): from __future__ import annotations import os +import queue import shlex import subprocess +import tempfile +import threading +import time from dataclasses import dataclass from pathlib import Path from typing import Callable from rich.console import Console +import yaml + +from obliteratus import __version__ +from obliteratus.remote_contracts import ( + parse_remote_target, + remote_python_command, + remote_scp_spec, + validate_remote_settings, +) console = Console() @@ -45,6 +58,24 @@ class RemoteConfig: python: str = "python3" # remote python binary sync_results: bool = True gpus: str | None = None # comma-separated GPU IDs or "all" + install_source: str = "git+https://github.com/elder-plinius/OBLITERATUS.git" + + def __post_init__(self) -> None: + self.gpus = validate_remote_settings( + host=self.host, + user=self.user, + port=self.port, + remote_dir=self.remote_dir, + python=self.python, + gpus=self.gpus, + install_source=self.install_source, + ) + if ( + isinstance(self.install_timeout, bool) + or not isinstance(self.install_timeout, int) + or self.install_timeout <= 0 + ): + raise ValueError("remote install timeout must be a positive integer") @property def ssh_target(self) -> str: @@ -53,11 +84,7 @@ class RemoteConfig: @classmethod def from_cli_args(cls, remote_str: str, **kwargs) -> RemoteConfig: """Parse 'user@host' or just 'host' from CLI --remote flag.""" - if "@" in remote_str: - user, host = remote_str.rsplit("@", 1) - else: - user = "root" - host = remote_str + user, host = parse_remote_target(remote_str) return cls(host=host, user=user, **kwargs) @classmethod @@ -127,15 +154,51 @@ class RemoteRunner: text=True, bufsize=1, ) + if proc.stdout is None: + proc.kill() + proc.wait() + raise RuntimeError("remote process did not expose a stdout stream") + + output: queue.Queue[object] = queue.Queue() + finished = object() + + def read_output() -> None: + try: + for line in proc.stdout: + if not isinstance(line, str): + raise RuntimeError("remote process emitted non-text output") + output.put(line.rstrip("\n")) + except BaseException as exc: + output.put(exc) + finally: + output.put(finished) + + threading.Thread(target=read_output, daemon=True).start() + deadline = time.monotonic() + timeout if timeout is not None else None try: - for line in proc.stdout: - line = line.rstrip("\n") - self.on_log(line) - proc.wait(timeout=timeout) + while True: + remaining = None if deadline is None else max(0, deadline - time.monotonic()) + try: + item = output.get(timeout=remaining) + except queue.Empty as exc: + raise subprocess.TimeoutExpired(cmd, timeout) from exc + if item is finished: + break + if isinstance(item, BaseException): + raise item + self.on_log(item) + remaining = None if deadline is None else max(0, deadline - time.monotonic()) + proc.wait(timeout=remaining) except subprocess.TimeoutExpired: proc.kill() + proc.wait() self.on_log("[red]Remote command timed out[/]") - return 1 + return 124 + except KeyboardInterrupt: + proc.kill() + proc.wait() + self.on_log("[yellow]Remote command cancelled[/]") + raise return proc.returncode else: return subprocess.run( @@ -175,37 +238,43 @@ class RemoteRunner: self.on_log("[yellow]No GPUs detected on remote (nvidia-smi failed)[/]") return None - def _env_prefix(self) -> str: - """Build environment variable prefix for remote commands (e.g. CUDA_VISIBLE_DEVICES).""" - parts = [] - if self.config.gpus and self.config.gpus.lower() != "all": - parts.append(f"CUDA_VISIBLE_DEVICES={self.config.gpus}") - return " ".join(parts) + " " if parts else "" - def ensure_obliteratus(self) -> bool: """Install or update obliteratus on the remote if needed.""" - # Check if already installed + version_command = remote_python_command( + self.config.python, + ["-c", "import obliteratus; print(obliteratus.__version__)"], + ) check = self.run_ssh( - f"{self.config.python} -c \"import obliteratus; print(obliteratus.__version__)\"", + version_command, timeout=30, ) - if isinstance(check, subprocess.CompletedProcess) and check.returncode == 0: - version = check.stdout.strip() - self.on_log(f"Obliteratus {version} already installed on remote") + if ( + isinstance(check, subprocess.CompletedProcess) + and check.returncode == 0 + and check.stdout.strip() == __version__ + ): + self.on_log(f"Obliteratus {__version__} already installed on remote") return True - # Install from PyPI or git - self.on_log("Installing obliteratus on remote...") - install_cmd = ( - f"{self.config.python} -m pip install --quiet " - f"git+https://github.com/StellaAthena/OBLITERATUS.git" + self.on_log(f"Installing Obliteratus {__version__} on remote...") + install_cmd = remote_python_command( + self.config.python, + ["-m", "pip", "install", "--quiet", "--upgrade", self.config.install_source], ) rc = self.run_ssh(install_cmd, stream=True, timeout=self.config.install_timeout) if rc != 0: self.on_log("[red]Failed to install obliteratus on remote[/]") return False - self.on_log("Obliteratus installed successfully") + verify = self.run_ssh(version_command, timeout=30) + if not ( + isinstance(verify, subprocess.CompletedProcess) + and verify.returncode == 0 + and verify.stdout.strip() == __version__ + ): + self.on_log("[red]Installed Obliteratus version could not be verified[/]") + return False + self.on_log(f"Obliteratus {__version__} installed successfully") return True def sync_results_back(self, remote_output_dir: str, local_output_dir: str) -> bool: @@ -216,7 +285,7 @@ class RemoteRunner: self.on_log(f"Syncing results: {self.config.ssh_target}:{remote_output_dir} -> {local_output_dir}") cmd = self._scp_base_cmd() + [ - f"{self.config.ssh_target}:{remote_output_dir}/", + remote_scp_spec(self.config.ssh_target, remote_output_dir, directory=True), str(local_path), ] @@ -242,14 +311,24 @@ class RemoteRunner: refinement_passes: int | None = None, large_model: bool = False, verify_sample_size: int | None = None, + min_layer_fraction: float | None = None, + max_layer_fraction: float | None = None, + harmless_pc_count: int | None = None, + shield_concept_count: int | None = None, + shield_ridge: float | None = None, + shield_residualize: bool | None = None, + shield_layer_penalty: float | None = None, + projection_target: str | None = None, + projection_row_fraction: float | None = None, + refusal_max_tokens: int | None = None, ) -> str: """Build the remote obliteratus CLI command.""" remote_output = output_dir or f"{self.config.remote_dir}/output/{model.replace('/', '_')}" parts = [ - self._env_prefix() + self.config.python, "-m", "obliteratus", - "obliterate", shlex.quote(model), - "--output-dir", shlex.quote(remote_output), + "-m", "obliteratus", + "obliterate", model, + "--output-dir", remote_output, "--method", method, "--device", device, "--dtype", dtype, @@ -268,20 +347,36 @@ class RemoteRunner: parts.append("--large-model") if verify_sample_size is not None: parts.extend(["--verify-sample-size", str(verify_sample_size)]) + optional_values = ( + ("--min-layer-fraction", min_layer_fraction), + ("--max-layer-fraction", max_layer_fraction), + ("--harmless-pc-count", harmless_pc_count), + ("--shield-concept-count", shield_concept_count), + ("--shield-ridge", shield_ridge), + ("--shield-layer-penalty", shield_layer_penalty), + ("--projection-target", projection_target), + ("--projection-row-fraction", projection_row_fraction), + ("--refusal-max-tokens", refusal_max_tokens), + ) + for flag, value in optional_values: + if value is not None: + parts.extend([flag, str(value)]) + if shield_residualize: + parts.append("--shield-residualize") - return " ".join(parts) + return remote_python_command(self.config.python, parts, gpus=self.config.gpus) def build_run_command(self, remote_config_path: str, output_dir: str | None = None, preset: str | None = None) -> str: """Build remote 'obliteratus run' command.""" parts = [ - self._env_prefix() + self.config.python, "-m", "obliteratus", - "run", shlex.quote(remote_config_path), + "-m", "obliteratus", + "run", remote_config_path, ] if output_dir: - parts.extend(["--output-dir", shlex.quote(output_dir)]) + parts.extend(["--output-dir", output_dir]) if preset: parts.extend(["--preset", preset]) - return " ".join(parts) + return remote_python_command(self.config.python, parts, gpus=self.config.gpus) def build_tourney_command( self, @@ -299,9 +394,9 @@ class RemoteRunner: remote_output = output_dir or f"{self.config.remote_dir}/tourney/{model.replace('/', '_')}" parts = [ - self._env_prefix() + self.config.python, "-m", "obliteratus", - "tourney", shlex.quote(model), - "--output-dir", shlex.quote(remote_output), + "-m", "obliteratus", + "tourney", model, + "--output-dir", remote_output, "--device", device, "--dtype", dtype, "--dataset", dataset, @@ -314,18 +409,26 @@ class RemoteRunner: parts.extend(["--hub-repo", hub_repo]) if methods: parts.extend(["--methods"] + methods) - return " ".join(parts) + return remote_python_command(self.config.python, parts, gpus=self.config.gpus) def upload_config(self, local_config_path: str) -> str: - """Upload a YAML config file to the remote.""" + """Upload a YAML config without recursively redispatching remotely.""" remote_path = f"{self.config.remote_dir}/config.yaml" - self.run_ssh(f"mkdir -p {shlex.quote(self.config.remote_dir)}") + self.run_ssh(shlex.join(["mkdir", "-p", self.config.remote_dir])) - cmd = self._scp_base_cmd() - # scp uses -P not -p, already handled in _scp_base_cmd - cmd += [local_config_path, f"{self.config.ssh_target}:{remote_path}"] + source = Path(local_config_path) + payload = yaml.safe_load(source.read_text(encoding="utf-8")) + if not isinstance(payload, dict): + raise ValueError("remote study config must contain a YAML mapping") + payload.pop("remote", None) - result = subprocess.run(cmd, capture_output=True, text=True, timeout=60) + with tempfile.TemporaryDirectory(prefix="obliteratus-remote-config-") as temp_dir: + upload_path = Path(temp_dir) / "config.yaml" + upload_path.write_text(yaml.safe_dump(payload, sort_keys=False), encoding="utf-8") + cmd = self._scp_base_cmd() + # scp uses -P not -p, already handled in _scp_base_cmd + cmd += [str(upload_path), remote_scp_spec(self.config.ssh_target, remote_path)] + result = subprocess.run(cmd, capture_output=True, text=True, timeout=60) if result.returncode != 0: raise RuntimeError(f"Failed to upload config: {result.stderr}") self.on_log(f"Config uploaded to {remote_path}") @@ -353,7 +456,7 @@ class RemoteRunner: return None # 4. Create remote working directory - self.run_ssh(f"mkdir -p {shlex.quote(self.config.remote_dir)}") + self.run_ssh(shlex.join(["mkdir", "-p", self.config.remote_dir])) # 5. Build and run the command remote_output = f"{self.config.remote_dir}/output/{model.replace('/', '_')}" diff --git a/obliteratus/remote_contracts.py b/obliteratus/remote_contracts.py new file mode 100644 index 0000000..7b7a3b7 --- /dev/null +++ b/obliteratus/remote_contracts.py @@ -0,0 +1,99 @@ +"""Pure validation and command contracts for remote execution.""" + +from __future__ import annotations + +import re +import shlex + + +_REMOTE_USER = re.compile(r"[A-Za-z0-9_.-]+") + + +def _require_text(value: object, label: str) -> str: + if not isinstance(value, str) or not value: + raise ValueError(f"{label} must be a non-empty string") + if any(ord(character) < 32 or ord(character) == 127 for character in value): + raise ValueError(f"{label} may not contain control characters") + return value + + +def normalize_gpu_selection(value: str | None) -> str | None: + """Return a canonical CUDA device selection or reject malformed input.""" + if value is None: + return None + value = _require_text(value, "remote gpus").strip() + if value.lower() == "all": + return "all" + parts = value.split(",") + if any(not part.strip().isdigit() for part in parts): + raise ValueError("remote gpus must be 'all' or comma-separated non-negative integers") + return ",".join(str(int(part.strip())) for part in parts) + + +def _validate_remote_identity(host: str, user: str) -> None: + """Validate the host and user portions shared by CLI and YAML inputs.""" + host = _require_text(host, "remote host") + user = _require_text(user, "remote user") + if host.startswith("-") or "@" in host or any(character.isspace() for character in host): + raise ValueError("remote host must be a host name or address without user or options") + if user.startswith("-") or _REMOTE_USER.fullmatch(user) is None: + raise ValueError("remote user contains unsupported characters") + + +def validate_remote_settings( + *, + host: str, + user: str, + port: int, + remote_dir: str, + python: str, + gpus: str | None, + install_source: str, +) -> str | None: + """Validate public remote settings and return canonical GPU selection.""" + _validate_remote_identity(host, user) + remote_dir = _require_text(remote_dir, "remote directory") + _require_text(python, "remote Python") + _require_text(install_source, "remote install source") + + if isinstance(port, bool) or not isinstance(port, int) or not 1 <= port <= 65535: + raise ValueError("remote port must be an integer from 1 through 65535") + if not remote_dir.startswith("/"): + raise ValueError("remote directory must be an absolute POSIX path") + return normalize_gpu_selection(gpus) + + +def parse_remote_target(remote: str) -> tuple[str, str]: + """Parse and validate ``[USER@]HOST`` without accepting SSH options.""" + remote = _require_text(remote, "remote target").strip() + if remote.count("@") > 1: + raise ValueError("remote target may contain at most one user separator") + user, separator, host = remote.partition("@") + if not separator: + user, host = "root", user + _validate_remote_identity(host, user) + return user, host + + +def remote_python_command( + python: str, + arguments: list[str], + *, + gpus: str | None = None, +) -> str: + """Build one shell-safe command for the remote SSH shell.""" + tokens: list[str] = [] + normalized_gpus = normalize_gpu_selection(gpus) + if normalized_gpus not in (None, "all"): + tokens.extend(["env", f"CUDA_VISIBLE_DEVICES={normalized_gpus}"]) + tokens.extend([_require_text(python, "remote Python"), *arguments]) + return shlex.join(tokens) + + +def remote_scp_spec(target: str, path: str, *, directory: bool = False) -> str: + """Quote a remote SCP path while preserving the host/path separator.""" + target = _require_text(target, "SSH target") + path = _require_text(path, "remote SCP path") + if directory: + path = f"{path.rstrip('/')}/" + return f"{target}:{shlex.quote(path)}" diff --git a/pyproject.toml b/pyproject.toml index 073ef90..5a158d2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -120,6 +120,7 @@ only_mutate = [ "obliteratus/analysis/numerical_contracts.py", "obliteratus/runtime_contracts.py", "obliteratus/persistence_contracts.py", + "obliteratus/remote_contracts.py", "scripts/check_coverage_thresholds.py", ] pytest_add_cli_args = ["--no-cov", "-q"] @@ -129,6 +130,7 @@ pytest_add_cli_args_test_selection = [ "tests/test_coverage_thresholds.py", "tests/test_numerical_contracts.py", "tests/test_persistence_contracts.py", + "tests/test_remote_contracts.py", "tests/test_runtime_contracts.py", ] mutate_only_covered_lines = true diff --git a/scripts/run_repeat_gate.py b/scripts/run_repeat_gate.py index dfe77f0..c2d53c4 100644 --- a/scripts/run_repeat_gate.py +++ b/scripts/run_repeat_gate.py @@ -18,9 +18,13 @@ DEFAULT_TESTS = ( "tests/test_config.py", "tests/test_config_properties.py", "tests/test_coverage_thresholds.py", + "tests/test_numerical_contracts.py", + "tests/test_persistence_contracts.py", "tests/test_property_contracts.py", "tests/test_advanced_metrics.py", "tests/test_metrics.py", + "tests/test_remote_contracts.py", + "tests/test_runtime_contracts.py", ) HASH_SEEDS = ("0", "1", "8675309") diff --git a/tests/test_cli_boundaries.py b/tests/test_cli_boundaries.py index 406c1f2..811d936 100644 --- a/tests/test_cli_boundaries.py +++ b/tests/test_cli_boundaries.py @@ -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 diff --git a/tests/test_config.py b/tests/test_config.py index ef982f6..db38f25 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -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}) diff --git a/tests/test_remote_boundaries.py b/tests/test_remote_boundaries.py index 8ad20d1..1026170 100644 --- a/tests/test_remote_boundaries.py +++ b/tests/test_remote_boundaries.py @@ -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 diff --git a/tests/test_remote_contracts.py b/tests/test_remote_contracts.py new file mode 100644 index 0000000..20565d8 --- /dev/null +++ b/tests/test_remote_contracts.py @@ -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