Add SSH polling to mesh GPU snapshots

Author: tsotchke

docs/ci_and_scripts.md

@@ -1089,7 +1089,8 @@ python3 scripts/mesh_worker_gpu_snapshot_selftest.py
 
 ### `mesh_worker_gpu_snapshot.py` and `mesh_worker_gpu_reconcile.py`
 
-`mesh_worker_gpu_snapshot.py` emits host-local `nvidia-smi` GPU/process state.
+`mesh_worker_gpu_snapshot.py` emits host-local `nvidia-smi` GPU/process state,
+or SSH-polled state with `--ssh-host` for the scheduler-VM v0 deployment model.
 `mesh_worker_gpu_reconcile.py` compares that worker snapshot with central
 arbiter lease state and reports a failed reconciliation when CUDA processes are
 running without an active lease for the resource. Feed reconciliation JSON into
@@ -1098,6 +1099,10 @@ GPU activity fail the control-plane audit.
 
 ```sh
 python3 scripts/mesh_worker_gpu_snapshot.py --resource cosbox:cuda3090 --json
+python3 scripts/mesh_worker_gpu_snapshot.py \
+  --resource cosbox:cuda3090 \
+  --ssh-host cosbox \
+  --json
 python3 scripts/mesh_worker_gpu_reconcile.py \
   --snapshot-json worker_gpu_snapshot.json \
   --arbiter-status-json arbiter_status.json \

scripts/mesh_worker_gpu_snapshot.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-"""Emit host-local GPU process state for scheduler reconciliation."""
+"""Emit local or SSH-polled GPU process state for scheduler reconciliation."""
 
 from __future__ import annotations
 
@@ -30,6 +30,12 @@ def run_capture(argv: list[str], *, timeout: float) -> subprocess.CompletedProce
     )
 
 
+def nvidia_smi_command(args: argparse.Namespace, nvidia_smi_args: list[str]) -> list[str]:
+    if args.ssh_host:
+        return ["ssh", args.ssh_host, args.nvidia_smi, *nvidia_smi_args]
+    return [args.nvidia_smi, *nvidia_smi_args]
+
+
 def parse_compute_apps(stdout: str) -> list[dict[str, Any]]:
     apps: list[dict[str, Any]] = []
     for row in csv.reader(io.StringIO(stdout)):
@@ -92,15 +98,17 @@ def as_int(value: str) -> int | None:
     return rows
 
 
-def query_compute_apps(nvidia_smi: str, timeout: float) -> dict[str, Any]:
+def query_compute_apps(args: argparse.Namespace) -> dict[str, Any]:
     try:
         proc = run_capture(
-            [
-                nvidia_smi,
-                "--query-compute-apps=pid,process_name,used_gpu_memory",
-                "--format=csv,noheader,nounits",
-            ],
-            timeout=timeout,
+            nvidia_smi_command(
+                args,
+                [
+                    "--query-compute-apps=pid,process_name,used_gpu_memory",
+                    "--format=csv,noheader,nounits",
+                ],
+            ),
+            timeout=args.timeout_sec,
         )
     except FileNotFoundError as exc:
         return {"ok": False, "reason": "nvidia_smi_not_found", "error": str(exc), "apps": []}
@@ -115,15 +123,17 @@ def query_compute_apps(nvidia_smi: str, timeout: float) -> dict[str, Any]:
     }
 
 
-def query_gpus(nvidia_smi: str, timeout: float) -> dict[str, Any]:
+def query_gpus(args: argparse.Namespace) -> dict[str, Any]:
     try:
         proc = run_capture(
-            [
-                nvidia_smi,
-                "--query-gpu=index,uuid,pci.bus_id,name,memory.total,memory.used,memory.free,utilization.gpu",
-                "--format=csv,noheader,nounits",
-            ],
-            timeout=timeout,
+            nvidia_smi_command(
+                args,
+                [
+                    "--query-gpu=index,uuid,pci.bus_id,name,memory.total,memory.used,memory.free,utilization.gpu",
+                    "--format=csv,noheader,nounits",
+                ],
+            ),
+            timeout=args.timeout_sec,
         )
     except FileNotFoundError as exc:
         return {"ok": False, "reason": "nvidia_smi_not_found", "error": str(exc), "gpus": []}
@@ -146,22 +156,24 @@ def proc_cmdline(pid: int) -> str:
         return ""
 
 
-def enrich_apps(apps: list[dict[str, Any]]) -> list[dict[str, Any]]:
+def enrich_apps(apps: list[dict[str, Any]], *, include_cmdline: bool) -> list[dict[str, Any]]:
     enriched = []
     for app in apps:
         row = dict(app)
         pid = row.get("pid")
-        if isinstance(pid, int):
+        if include_cmdline and isinstance(pid, int):
             row["cmdline"] = proc_cmdline(pid)
         enriched.append(row)
     return enriched
 
 
 def build_payload(args: argparse.Namespace) -> dict[str, Any]:
-    gpu_query = query_gpus(args.nvidia_smi, args.timeout_sec)
-    compute_query = query_compute_apps(args.nvidia_smi, args.timeout_sec)
-    apps = enrich_apps(compute_query.get("apps", []))
+    gpu_query = query_gpus(args)
+    compute_query = query_compute_apps(args)
+    apps = enrich_apps(compute_query.get("apps", []), include_cmdline=not args.ssh_host)
     ok = bool(gpu_query.get("ok")) and bool(compute_query.get("ok"))
+    poller_host = socket.gethostname()
+    worker_host = args.ssh_host or poller_host
     reasons = []
     if not gpu_query.get("ok"):
         reasons.append(str(gpu_query.get("reason") or "gpu_query_failed"))
@@ -172,7 +184,9 @@ def build_payload(args: argparse.Namespace) -> dict[str, Any]:
         "ok": ok,
         "reason": "ok" if ok else ",".join(reasons),
         "checked_at_unix": time.time(),
-        "worker_host": socket.gethostname(),
+        "worker_host": worker_host,
+        "poller_host": poller_host,
+        "ssh_host": args.ssh_host,
         "resource": args.resource,
         "nvidia_smi": args.nvidia_smi,
         "gpus": gpu_query.get("gpus", []),
@@ -188,6 +202,7 @@ def parse_args(argv: list[str]) -> argparse.Namespace:
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument("--resource", default="cuda")
     parser.add_argument("--nvidia-smi", default="nvidia-smi")
+    parser.add_argument("--ssh-host", default="")
     parser.add_argument("--timeout-sec", type=float, default=10.0)
     parser.add_argument("--json", action="store_true")
     args = parser.parse_args(argv)

scripts/mesh_worker_gpu_snapshot_selftest.py

@@ -5,7 +5,9 @@
 
 import importlib.machinery
 import importlib.util
+import argparse
 import pathlib
+import subprocess
 from types import ModuleType
 
 
@@ -32,6 +34,55 @@ def test_snapshot_parses_nvidia_smi_rows() -> None:
     assert gpus[0]["memory_total_mib"] == 24576
 
 
+def test_snapshot_supports_ssh_polling_command_prefix() -> None:
+    mod = load_script("mesh_worker_gpu_snapshot_under_test", ROOT / "scripts" / "mesh_worker_gpu_snapshot.py")
+    args = argparse.Namespace(ssh_host="cosbox", nvidia_smi="nvidia-smi")
+    argv = mod.nvidia_smi_command(args, ["--query-gpu=index"])
+    assert argv == ["ssh", "cosbox", "nvidia-smi", "--query-gpu=index"]
+    args.ssh_host = ""
+    assert mod.nvidia_smi_command(args, ["--query-gpu=index"]) == [
+        "nvidia-smi",
+        "--query-gpu=index",
+    ]
+
+
+def test_snapshot_records_worker_and_poller_hosts_for_ssh() -> None:
+    mod = load_script("mesh_worker_gpu_snapshot_under_test", ROOT / "scripts" / "mesh_worker_gpu_snapshot.py")
+
+    def fake_run_capture(argv: list[str], *, timeout: float) -> subprocess.CompletedProcess[str]:
+        del timeout
+        if "--query-gpu=index,uuid,pci.bus_id,name,memory.total,memory.used,memory.free,utilization.gpu" in argv:
+            return subprocess.CompletedProcess(
+                argv,
+                0,
+                stdout="0, GPU-test, 00000000:01:00.0, RTX 3090, 24576, 1024, 23552, 7\n",
+                stderr="",
+            )
+        return subprocess.CompletedProcess(argv, 0, stdout="1234, python, 8192\n", stderr="")
+
+    original_run_capture = mod.run_capture
+    original_hostname = mod.socket.gethostname
+    mod.run_capture = fake_run_capture
+    mod.socket.gethostname = lambda: "poller"
+    try:
+        payload = mod.build_payload(
+            argparse.Namespace(
+                resource="cosbox:cuda3090",
+                nvidia_smi="nvidia-smi",
+                ssh_host="cosbox",
+                timeout_sec=1.0,
+            )
+        )
+    finally:
+        mod.run_capture = original_run_capture
+        mod.socket.gethostname = original_hostname
+
+    assert payload["worker_host"] == "cosbox"
+    assert payload["poller_host"] == "poller"
+    assert payload["cuda_pids"] == [1234]
+    assert "cmdline" not in payload["cuda_apps"][0]
+
+
 def test_reconcile_drains_unleased_cuda() -> None:
     mod = load_script("mesh_worker_gpu_reconcile_under_test", ROOT / "scripts" / "mesh_worker_gpu_reconcile.py")
     snapshot = {
@@ -65,6 +116,8 @@ def test_reconcile_allows_leased_cuda() -> None:
 
 def main() -> int:
     test_snapshot_parses_nvidia_smi_rows()
+    test_snapshot_supports_ssh_polling_command_prefix()
+    test_snapshot_records_worker_and_poller_hosts_for_ssh()
     test_reconcile_drains_unleased_cuda()
     test_reconcile_allows_leased_cuda()
     print("mesh worker GPU snapshot selftest OK")