Checkpoint experiments through SWE pilot v2

scripts/remote_docker.py

@@ -0,0 +1,96 @@
+"""Run a local command against the experiment's Docker daemon over SSH.
+
+The command itself runs locally. Only Docker CLI requests made by it are directed to
+the remote daemon. A server architecture check fails closed before the command starts.
+"""
+from __future__ import annotations
+
+import argparse
+import os
+import subprocess
+import sys
+from urllib.parse import urlsplit
+
+
+DEFAULT_DOCKER_HOST = "ssh://[email protected]"
+EXPECTED_SERVER = "linux/amd64"
+
+
+def docker_host(environ: dict[str, str]) -> str:
+    return environ.get("MBB_DOCKER_HOST", DEFAULT_DOCKER_HOST)
+
+
+def validate_host(host: str) -> None:
+    parsed = urlsplit(host)
+    if (
+        parsed.scheme != "ssh"
+        or not parsed.hostname
+        or parsed.password is not None
+        or parsed.query
+        or parsed.fragment
+        or parsed.path not in ("", "/")
+    ):
+        raise ValueError(
+            "MBB_DOCKER_HOST must be an ssh://[user@]host URI without a password, "
+            "path, query, or fragment"
+        )
+
+
+def check_daemon(host: str, environ: dict[str, str]) -> str:
+    env = dict(environ)
+    env["DOCKER_HOST"] = host
+    result = subprocess.run(
+        [
+            "docker",
+            "version",
+            "--format",
+            "{{.Server.Os}}/{{.Server.Arch}}",
+        ],
+        env=env,
+        check=True,
+        capture_output=True,
+        text=True,
+        timeout=30,
+    )
+    server = result.stdout.strip()
+    if server != EXPECTED_SERVER:
+        raise RuntimeError(
+            f"Refusing Docker server {server!r}; expected {EXPECTED_SERVER!r} at {host}"
+        )
+    return server
+
+
+def parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(description=__doc__)
+    p.add_argument(
+        "command",
+        nargs=argparse.REMAINDER,
+        help="local command to run; place it after -- (omit it for a connection check)",
+    )
+    return p
+
+
+def main(argv: list[str] | None = None, environ: dict[str, str] | None = None) -> int:
+    args = parser().parse_args(argv)
+    env = dict(os.environ if environ is None else environ)
+    host = docker_host(env)
+    try:
+        validate_host(host)
+        server = check_daemon(host, env)
+    except (ValueError, RuntimeError, subprocess.SubprocessError, OSError) as exc:
+        print(f"Remote Docker preflight failed: {exc}", file=sys.stderr)
+        return 2
+
+    print(f"Remote Docker ready: {host} ({server})", flush=True)
+    command = args.command[1:] if args.command[:1] == ["--"] else args.command
+    if not command:
+        return 0
+    env["DOCKER_HOST"] = host
+    try:
+        return subprocess.run(command, env=env).returncode
+    except KeyboardInterrupt:
+        return 130
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())