remote-bench: detach bench in tmux so SSH disconnects can't kill it

Long aggregations exceed the IAP tunnel idle timeout (~30 min in
practice). When the tunnel resets, gcloud's `compute ssh --command`
session dies, the bench process gets SIGHUP'd, and we lose the run.
Hit this today running c4-standard-{8,16,32}: VMs went idle ~30 min
in, no result files written.

Fix: split the remote work across two SSH calls.

  1. Synchronous setup (existing ssh_exec): apt + rustup + uv + clone +
     prereqs. Short, well-bounded; live output streams to local stdout.
  2. The actual bench is launched into a detached `tmux` session
     (`tmux new-session -d -s leanbench ./run-bench.sh`) so it's owned
     by tmux, not by the SSH session. The setup ssh_exec returns
     immediately after launching it.

The orchestrator then enters _poll_bench_completion, which every
--poll-interval-s seconds (default 30):

  - tail -c +N leanBench/bench.log    (stream new output to local stdout)
  - tmux has-session -t leanbench     (check if bench has finished)
  - on completion, read leanBench/bench.exit for the exit code

Tolerates transient SSH failures (up to 6 in a row) so a brief IAP
hiccup doesn't trigger the outer try/except → finally → destroy
cascade and abort a healthy bench. On final-give-up, prints the
gcloud command to manually `tmux attach -t leanbench` and recover.

remote_setup.sh apt-installs `tmux`, writes a small run-bench.sh that
captures the bench exit code to bench.exit, and launches it via tmux.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: unnawut

scripts/cloud/remote_bench.py

@@ -31,6 +31,7 @@
 import signal
 import subprocess
 import sys
+import time
 from pathlib import Path
 
 from .provisioners import Instance, InstanceSpec
@@ -92,6 +93,11 @@ def main():
                     help="Don't destroy the VM if the bench fails — useful for debugging")
     ap.add_argument("--out-dir", type=Path, default=Path("results"))
     ap.add_argument("--ssh-timeout-s", type=int, default=300)
+    ap.add_argument("--poll-interval-s", type=int, default=30,
+                    help="Seconds between bench-completion polls. The bench "
+                         "runs in a detached tmux session on the VM; the "
+                         "orchestrator tails its log and checks for exit on "
+                         "this cadence.")
     ap.add_argument("--signers-cache", type=Path, default=None,
                     help="Path to a local benchmark_signers_cache_<hash>.bin file to "
                          "pre-upload to each VM. Skips the ~few-minute lazy regen on "
@@ -272,9 +278,7 @@ def run_one_machine(
             prov.scp_to(inst, args.signers_cache,
                         f"leanBench-signers/{args.signers_cache.name}")
 
-        print(f"{prefix}==> running setup + bench")
-        if not prefix:
-            print("─" * 64)
+        print(f"{prefix}==> running setup + launching bench in detached tmux")
         bench_args = f"--label {machine_type} {args.bench_args}".strip()
         # shlex.quote wraps each value in shell-safe single quotes so a
         # malicious branch / args string can't break out into bash.
@@ -285,10 +289,19 @@ def run_one_machine(
         )
         cmd = env_exports + REMOTE_SETUP_SH.read_text()
         rc = prov.ssh_exec(inst, cmd, prefix=prefix)
+        if rc != 0:
+            raise RuntimeError(f"setup + launch failed with code {rc}")
+
+        # Bench is now running in a detached tmux session on the VM, so a
+        # dropped SSH/IAP tunnel can't kill it. Poll separately for
+        # progress + completion.
+        print(f"{prefix}==> polling bench (interval {args.poll_interval_s}s)")
+        if not prefix:
+            print("─" * 64)
+        _poll_bench_completion(prov, inst, prefix=prefix,
+                               poll_interval_s=args.poll_interval_s)
         if not prefix:
             print("─" * 64)
-        if rc != 0:
-            raise RuntimeError(f"benchmark exited with code {rc}")
 
         marker = prov.ssh_capture(
             inst,
@@ -331,6 +344,66 @@ def run_one_machine(
     return summary
 
 
+def _poll_bench_completion(prov, inst, prefix: str, poll_interval_s: int) -> None:
+    """Poll the VM until the detached bench tmux session exits.
+
+    Streams new bench.log bytes to stdout (with optional `prefix`) so
+    progress is visible. Raises RuntimeError if bench exited non-zero
+    or never wrote its exit-code file.
+
+    Tolerates transient SSH failures: a brief IAP tunnel outage shouldn't
+    tear down a healthy VM mid-bench. We only give up after MAX_FAILURES
+    consecutive bad polls.
+    """
+    MAX_FAILURES = 6
+    last_byte = 0
+    consecutive_failures = 0
+    while True:
+        try:
+            chunk = prov.ssh_capture(
+                inst,
+                f"tail -c +{last_byte + 1} leanBench/bench.log 2>/dev/null || true",
+            )
+            if chunk:
+                for line in chunk.splitlines():
+                    print(f"{prefix}{line}")
+                last_byte += len(chunk.encode("utf-8"))
+
+            status = prov.ssh_capture(
+                inst,
+                "tmux has-session -t leanbench 2>/dev/null && echo running || echo done",
+            )
+            consecutive_failures = 0
+            if status == "done":
+                break
+        except subprocess.CalledProcessError as e:
+            consecutive_failures += 1
+            print(
+                f"{prefix}    poll failed ({consecutive_failures}/{MAX_FAILURES}): "
+                f"{e}",
+                file=sys.stderr,
+            )
+            if consecutive_failures >= MAX_FAILURES:
+                raise RuntimeError(
+                    f"bench poll failed {MAX_FAILURES} times in a row; "
+                    f"VM may be unreachable. Bench may still be running on "
+                    f"the VM — re-attach with: "
+                    f"gcloud compute ssh {inst.name} --zone {inst.data.get('zone', '?')} "
+                    f"--tunnel-through-iap -- tmux attach -t leanbench"
+                ) from e
+        time.sleep(poll_interval_s)
+
+    exit_str = prov.ssh_capture(
+        inst, "cat leanBench/bench.exit 2>/dev/null || true",
+    ).strip()
+    if exit_str.startswith("EXIT="):
+        code = int(exit_str.split("=", 1)[1])
+        if code != 0:
+            raise RuntimeError(f"bench exited with code {code}")
+    else:
+        raise RuntimeError("bench did not write bench.exit; outcome unknown")
+
+
 def _confirm(prompt: str) -> bool:
     try:
         ans = input(f"{prompt} [y/N] ").strip().lower()

scripts/cloud/remote_setup.sh

@@ -19,7 +19,7 @@ sudo cloud-init status --wait >/dev/null 2>&1 || true
 echo '==> [remote] installing build prerequisites...'
 sudo DEBIAN_FRONTEND=noninteractive apt-get update -qq
 sudo DEBIAN_FRONTEND=noninteractive apt-get install -y -qq \
-    build-essential git curl ca-certificates pkg-config
+    build-essential git curl ca-certificates pkg-config tmux
 
 if ! command -v cargo >/dev/null 2>&1; then
     echo '==> [remote] installing rustup...'
@@ -50,11 +50,17 @@ git reset --hard --quiet "origin/$BRANCH"
 mkdir -p "$HOME/leanBench-signers"
 export SIGNERS_CACHE_DIR="$HOME/leanBench-signers"
 
-echo '==> [remote] running benchmark...'
-# Intentionally unquoted: BENCH_ARGS is multi-arg (e.g. "--label foo --samples 10").
+echo '==> [remote] launching bench in detached tmux session...'
+# Detach the long bench (cargo build + 17 workloads, can run > 30 min)
+# so a dropped SSH/IAP tunnel can't kill it. The orchestrator polls
+# bench.log + tmux session state separately.
+cat > run-bench.sh <<RUN
+#!/usr/bin/env bash
 # shellcheck disable=SC2086
-uv run bench $BENCH_ARGS
-
-# Echo a parseable marker so the orchestrator knows where the result
-# landed (independent of bench.py's free-form output).
-echo "RESULT_FILE=$(ls -t results/*.json 2>/dev/null | grep -v 'results/index.json' | head -1)"
+uv run bench $BENCH_ARGS > bench.log 2>&1
+echo "EXIT=\$?" > bench.exit
+RUN
+chmod +x run-bench.sh
+tmux kill-session -t leanbench 2>/dev/null || true
+tmux new-session -d -s leanbench ./run-bench.sh
+echo '==> [remote] tmux session leanbench started; orchestrator will poll'