mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-17 16:37:30 +02:00
547 lines
20 KiB
Python
547 lines
20 KiB
Python
"""Remote execution support for Obliteratus.
|
|
|
|
Run abliteration pipelines on remote GPU nodes via SSH. The remote machine
|
|
must have CUDA-capable GPUs and a Python environment. Obliteratus will be
|
|
auto-installed if not present.
|
|
|
|
Usage (CLI):
|
|
obliteratus obliterate meta-llama/Llama-3.1-8B-Instruct \
|
|
--remote user@gpu-node \
|
|
--ssh-key ~/.ssh/id_rsa
|
|
|
|
Usage (YAML config):
|
|
remote:
|
|
host: gpu-node
|
|
user: root
|
|
ssh_key: ~/.ssh/id_rsa
|
|
remote_dir: /tmp/obliteratus_run
|
|
"""
|
|
|
|
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()
|
|
|
|
|
|
@dataclass
|
|
class RemoteConfig:
|
|
"""SSH connection and remote execution settings."""
|
|
|
|
host: str
|
|
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 # seconds
|
|
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:
|
|
return f"{self.user}@{self.host}"
|
|
|
|
@classmethod
|
|
def from_cli_args(cls, remote_str: str, **kwargs) -> RemoteConfig:
|
|
"""Parse 'user@host' or just 'host' from CLI --remote flag."""
|
|
user, host = parse_remote_target(remote_str)
|
|
return cls(host=host, user=user, **kwargs)
|
|
|
|
@classmethod
|
|
def from_dict(cls, d: dict) -> RemoteConfig:
|
|
return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__})
|
|
|
|
|
|
class RemoteRunner:
|
|
"""Execute Obliteratus commands on a remote machine via SSH."""
|
|
|
|
def __init__(
|
|
self,
|
|
config: RemoteConfig,
|
|
on_log: Callable[[str], None] | None = None,
|
|
):
|
|
self.config = config
|
|
self.on_log = on_log or (lambda msg: console.print(f"[dim][remote][/] {msg}"))
|
|
|
|
def _ssh_base_cmd(self) -> list[str]:
|
|
"""Build base SSH command with common options."""
|
|
cmd = [
|
|
"ssh",
|
|
"-o", "StrictHostKeyChecking=yes",
|
|
"-o", "BatchMode=yes",
|
|
"-o", "ConnectTimeout=30",
|
|
"-p", str(self.config.port),
|
|
]
|
|
if self.config.known_hosts_file:
|
|
known_hosts = os.path.expanduser(self.config.known_hosts_file)
|
|
cmd.extend(["-o", f"UserKnownHostsFile={known_hosts}"])
|
|
if self.config.ssh_key:
|
|
key_path = os.path.expanduser(self.config.ssh_key)
|
|
cmd.extend(["-i", key_path])
|
|
cmd.append(self.config.ssh_target)
|
|
return cmd
|
|
|
|
def _scp_base_cmd(self) -> list[str]:
|
|
"""Build base SCP command."""
|
|
cmd = [
|
|
"scp",
|
|
"-o", "StrictHostKeyChecking=yes",
|
|
"-o", "BatchMode=yes",
|
|
"-P", str(self.config.port),
|
|
"-r",
|
|
]
|
|
if self.config.known_hosts_file:
|
|
known_hosts = os.path.expanduser(self.config.known_hosts_file)
|
|
cmd.extend(["-o", f"UserKnownHostsFile={known_hosts}"])
|
|
if self.config.ssh_key:
|
|
key_path = os.path.expanduser(self.config.ssh_key)
|
|
cmd.extend(["-i", key_path])
|
|
return cmd
|
|
|
|
def run_ssh(self, remote_cmd: str, stream: bool = False, timeout: int | None = None) -> subprocess.CompletedProcess | int:
|
|
"""Run a command on the remote host.
|
|
|
|
If stream=True, streams stdout/stderr in real-time and returns the
|
|
exit code. Otherwise returns CompletedProcess.
|
|
"""
|
|
cmd = self._ssh_base_cmd() + [remote_cmd]
|
|
|
|
if stream:
|
|
proc = subprocess.Popen(
|
|
cmd,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
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:
|
|
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 124
|
|
except KeyboardInterrupt:
|
|
proc.kill()
|
|
proc.wait()
|
|
self.on_log("[yellow]Remote command cancelled[/]")
|
|
raise
|
|
return proc.returncode
|
|
else:
|
|
return subprocess.run(
|
|
cmd,
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=timeout,
|
|
)
|
|
|
|
def check_connection(self) -> bool:
|
|
"""Verify SSH connectivity."""
|
|
self.on_log(f"Testing SSH connection to {self.config.ssh_target}...")
|
|
result = self.run_ssh("echo ok", timeout=30)
|
|
if isinstance(result, subprocess.CompletedProcess) and result.returncode == 0:
|
|
self.on_log("SSH connection OK")
|
|
return True
|
|
self.on_log("[red]SSH connection failed[/]")
|
|
return False
|
|
|
|
def check_gpu(self) -> str | None:
|
|
"""Check for CUDA GPUs on remote. Returns nvidia-smi output or None."""
|
|
result = self.run_ssh(
|
|
"nvidia-smi --query-gpu=index,name,memory.total,memory.free --format=csv,noheader",
|
|
timeout=30,
|
|
)
|
|
if isinstance(result, subprocess.CompletedProcess) and result.returncode == 0:
|
|
gpu_info = result.stdout.strip()
|
|
lines = gpu_info.split("\n")
|
|
self.on_log(f"Remote GPUs ({len(lines)} detected):")
|
|
for line in lines:
|
|
self.on_log(f" {line.strip()}")
|
|
if self.config.gpus and self.config.gpus.lower() != "all":
|
|
self.on_log(f" Selected GPUs: {self.config.gpus}")
|
|
else:
|
|
self.on_log(f" Using: all {len(lines)} GPUs")
|
|
return gpu_info
|
|
self.on_log("[yellow]No GPUs detected on remote (nvidia-smi failed)[/]")
|
|
return None
|
|
|
|
def ensure_obliteratus(self) -> bool:
|
|
"""Install or update obliteratus on the remote if needed."""
|
|
version_command = remote_python_command(
|
|
self.config.python,
|
|
["-c", "import obliteratus; print(obliteratus.__version__)"],
|
|
)
|
|
check = self.run_ssh(
|
|
version_command,
|
|
timeout=30,
|
|
)
|
|
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
|
|
|
|
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
|
|
|
|
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:
|
|
"""Copy results from remote back to local machine via scp."""
|
|
local_path = Path(local_output_dir)
|
|
local_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
self.on_log(f"Syncing results: {self.config.ssh_target}:{remote_output_dir} -> {local_output_dir}")
|
|
|
|
cmd = self._scp_base_cmd() + [
|
|
remote_scp_spec(self.config.ssh_target, remote_output_dir, directory=True),
|
|
str(local_path),
|
|
]
|
|
|
|
result = subprocess.run(cmd, capture_output=True, text=True, timeout=3600)
|
|
if result.returncode == 0:
|
|
self.on_log(f"Results synced to {local_output_dir}")
|
|
return True
|
|
else:
|
|
self.on_log(f"[red]SCP failed: {result.stderr}[/]")
|
|
return False
|
|
|
|
def build_obliterate_command(
|
|
self,
|
|
model: str,
|
|
output_dir: str | None = None,
|
|
method: str = "advanced",
|
|
device: str = "auto",
|
|
dtype: str = "float16",
|
|
quantization: str | None = None,
|
|
gpu_memory_utilization: float | None = None,
|
|
n_directions: int | None = None,
|
|
direction_method: str | None = None,
|
|
regularization: float | None = None,
|
|
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 = [
|
|
"-m", "obliteratus",
|
|
"obliterate", model,
|
|
"--output-dir", remote_output,
|
|
"--method", method,
|
|
"--device", device,
|
|
"--dtype", dtype,
|
|
]
|
|
if quantization:
|
|
parts.extend(["--quantization", quantization])
|
|
if gpu_memory_utilization is not None:
|
|
parts.extend(["--gpu-memory-utilization", str(gpu_memory_utilization)])
|
|
if n_directions is not None:
|
|
parts.extend(["--n-directions", str(n_directions)])
|
|
if direction_method:
|
|
parts.extend(["--direction-method", direction_method])
|
|
if regularization is not None:
|
|
parts.extend(["--regularization", str(regularization)])
|
|
if refinement_passes is not None:
|
|
parts.extend(["--refinement-passes", str(refinement_passes)])
|
|
if large_model:
|
|
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 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 = [
|
|
"-m", "obliteratus",
|
|
"run", remote_config_path,
|
|
]
|
|
if output_dir:
|
|
parts.extend(["--output-dir", output_dir])
|
|
if preset:
|
|
parts.extend(["--preset", preset])
|
|
return remote_python_command(self.config.python, parts, gpus=self.config.gpus)
|
|
|
|
def build_tourney_command(
|
|
self,
|
|
model: str,
|
|
output_dir: str | None = None,
|
|
device: str = "auto",
|
|
dtype: str = "float16",
|
|
quantization: str | None = None,
|
|
methods: list[str] | None = None,
|
|
hub_org: str | None = None,
|
|
hub_repo: str | None = None,
|
|
dataset: str = "builtin",
|
|
) -> str:
|
|
"""Build remote 'obliteratus tourney' command."""
|
|
remote_output = output_dir or f"{self.config.remote_dir}/tourney/{model.replace('/', '_')}"
|
|
|
|
parts = [
|
|
"-m", "obliteratus",
|
|
"tourney", model,
|
|
"--output-dir", remote_output,
|
|
"--device", device,
|
|
"--dtype", dtype,
|
|
"--dataset", dataset,
|
|
]
|
|
if quantization:
|
|
parts.extend(["--quantization", quantization])
|
|
if hub_org:
|
|
parts.extend(["--hub-org", hub_org])
|
|
if hub_repo:
|
|
parts.extend(["--hub-repo", hub_repo])
|
|
if methods:
|
|
parts.extend(["--methods"] + methods)
|
|
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 without recursively redispatching remotely."""
|
|
remote_path = f"{self.config.remote_dir}/config.yaml"
|
|
self.run_ssh(shlex.join(["mkdir", "-p", self.config.remote_dir]))
|
|
|
|
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)
|
|
|
|
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}")
|
|
return remote_path
|
|
|
|
def run_obliterate(
|
|
self,
|
|
model: str,
|
|
local_output_dir: str | None = None,
|
|
**kwargs,
|
|
) -> str | None:
|
|
"""Full remote obliteration: setup, run, sync results.
|
|
|
|
Returns local path to results, or None on failure.
|
|
"""
|
|
# 1. Verify connection
|
|
if not self.check_connection():
|
|
return None
|
|
|
|
# 2. Check GPUs
|
|
self.check_gpu()
|
|
|
|
# 3. Ensure obliteratus is installed
|
|
if not self.ensure_obliteratus():
|
|
return None
|
|
|
|
# 4. Create remote working directory
|
|
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('/', '_')}"
|
|
cmd = self.build_obliterate_command(model, output_dir=remote_output, **kwargs)
|
|
self.on_log(f"Running: {cmd}")
|
|
|
|
rc = self.run_ssh(cmd, stream=True)
|
|
if rc != 0:
|
|
self.on_log(f"[red]Remote obliteration failed (exit code {rc})[/]")
|
|
return None
|
|
|
|
# 6. Sync results back
|
|
if self.config.sync_results:
|
|
local_output = local_output_dir or f"abliterated/{model.replace('/', '_')}"
|
|
if self.sync_results_back(remote_output, local_output):
|
|
return local_output
|
|
return None
|
|
|
|
self.on_log(f"Results on remote: {remote_output}")
|
|
return remote_output
|
|
|
|
def run_config(
|
|
self,
|
|
local_config_path: str,
|
|
local_output_dir: str | None = None,
|
|
preset: str | None = None,
|
|
) -> str | None:
|
|
"""Upload config, run study remotely, sync results."""
|
|
if not self.check_connection():
|
|
return None
|
|
self.check_gpu()
|
|
if not self.ensure_obliteratus():
|
|
return None
|
|
|
|
# Upload config
|
|
remote_config = self.upload_config(local_config_path)
|
|
|
|
# Determine remote output dir
|
|
remote_output = f"{self.config.remote_dir}/results"
|
|
cmd = self.build_run_command(remote_config, output_dir=remote_output, preset=preset)
|
|
self.on_log(f"Running: {cmd}")
|
|
|
|
rc = self.run_ssh(cmd, stream=True)
|
|
if rc != 0:
|
|
self.on_log(f"[red]Remote run failed (exit code {rc})[/]")
|
|
return None
|
|
|
|
if self.config.sync_results:
|
|
local_output = local_output_dir or "results"
|
|
if self.sync_results_back(remote_output, local_output):
|
|
return local_output
|
|
return None
|
|
|
|
return remote_output
|
|
|
|
def run_tourney(
|
|
self,
|
|
model: str,
|
|
local_output_dir: str | None = None,
|
|
**kwargs,
|
|
) -> str | None:
|
|
"""Run tournament remotely, sync results."""
|
|
if not self.check_connection():
|
|
return None
|
|
self.check_gpu()
|
|
if not self.ensure_obliteratus():
|
|
return None
|
|
|
|
remote_output = f"{self.config.remote_dir}/tourney/{model.replace('/', '_')}"
|
|
cmd = self.build_tourney_command(model, output_dir=remote_output, **kwargs)
|
|
self.on_log(f"Running: {cmd}")
|
|
|
|
rc = self.run_ssh(cmd, stream=True)
|
|
if rc != 0:
|
|
self.on_log(f"[red]Remote tourney failed (exit code {rc})[/]")
|
|
return None
|
|
|
|
if self.config.sync_results:
|
|
local_output = local_output_dir or f"/tmp/obliteratus_tourney/{model.replace('/', '_')}"
|
|
if self.sync_results_back(remote_output, local_output):
|
|
return local_output
|
|
return None
|
|
|
|
return remote_output
|