"""Hold a Colab runtime's tunnel open from a detached process. Why this exists: the runtime is kept alive by an *open tunnel*, not by polling the API. There is no keepalive endpoint -- `GET /tun/m/keepalive`, `/tun/m/heartbeat` and friends all 400, and the plugin's whole surface is connect/instances/launch/status/stop. The previous version of this ran as an asyncio task inside the kernel, which meant a kernel or daemon restart silently killed it; the runtime then aged out and took an in-progress training run with it. So: this is spawned with `start_new_session=True`, reparented to init, and outlives every kernel. It re-reads `~/.albedo/colab-keepalive.json` on each reconnect, so a refreshed runtime token from the session (written on any waking turn) is picked up without restarting the daemon itself. The tunnel is `ssh -N` with the same ProxyCommand `colab.connect()` uses: `python -m albedo_plugins.ws`, with the websocket url and the `X-Colab-Ssh-Pubkey` header in the environment so the token never lands on a command line. `-N` runs no remote command, which is what holds it open. python3 scripts/colab_keepalive.py # daemonise and supervise python3 scripts/colab_keepalive.py --once # one connection attempt, for testing python3 scripts/colab_keepalive.py --status # is it up, and what is it doing Credentials file (written by the session, never committed): {"host": "colab-", "ws_url": "wss://.../colab/ssh?...", "pubkey": "ssh-ed25519 AAAA...", "key": "/path/to/private/key-or-empty", "pythonpath": "/path/to/albedo/priv/python", "updated": 1792444490.0} """ import argparse import json import os import threading import urllib.parse import urllib.request import pathlib import shlex import subprocess import sys import time CREDS = pathlib.Path(os.path.expanduser("~/.albedo/colab-keepalive.json")) LOG = pathlib.Path(os.path.expanduser("~/.albedo/colab-keepalive.log")) PID = pathlib.Path(os.path.expanduser("~/.albedo/colab-keepalive.pid")) STALE_AFTER = 3 * 3600 # a creds file this old means the session stopped refreshing it RETRY_SECONDS = 30 def log(message): LOG.parent.mkdir(parents=True, exist_ok=True) with LOG.open("a") as fh: fh.write(f"{time.strftime('%Y-%m-%dT%H:%M:%S')} {message}\n") def poll_runtime(creds, stop): """Request the runtime through Colab's proxy, independently of the exclusive ssh slot. The ssh slot is one-per-runtime AND sticky: a closed session keeps answering 429 already-active-session for a long time, so a tunnel cannot be relied on as the only keepalive. This path is separate -- plain HTTPS to the runtime proxy -- so it works even while the slot is held by someone else. """ runtime = creds.get("runtime_url") if not runtime: log("poller: no runtime_url in creds") return while not stop.is_set(): try: with urllib.request.urlopen(urllib.request.Request(runtime, headers={ "X-Colab-Client-Agent": "colab-cli"}), timeout=25) as response: body = response.read(120).decode("utf-8", "replace") log(f"poll ok {response.status} {body[:80]}") except Exception as error: # noqa: BLE001 - the poller must never die log(f"poll failed {type(error).__name__}: {error}") stop.wait(60) def read_creds(): try: creds = json.loads(CREDS.read_text()) except (OSError, ValueError) as error: return None, f"cannot read {CREDS}: {error}" missing = [k for k in ("host", "ws_url", "pubkey") if not creds.get(k)] if missing: return None, f"creds missing {missing}" age = time.time() - float(creds.get("updated") or 0) if age > STALE_AFTER: return None, f"creds are {age / 3600:.1f}h old; the session stopped refreshing them" return creds, None def build(creds): proxy = f"{shlex.quote(sys.executable)} -m albedo_plugins.ws".replace("%", "%%") argv = ["ssh", "-N", "-T", "-o", f"ProxyCommand={proxy}", "-o", "StrictHostKeyChecking=no", "-o", "UserKnownHostsFile=/dev/null", "-o", "ServerAliveInterval=30", "-o", "ServerAliveCountMax=3", "-o", "ExitOnForwardFailure=yes"] if creds.get("key"): argv += ["-i", creds["key"], "-o", "IdentitiesOnly=yes"] argv += [f"root@{creds['host']}"] env = dict(os.environ) if creds.get("pythonpath"): env["PYTHONPATH"] = creds["pythonpath"] if creds.get("auth_sock"): env["SSH_AUTH_SOCK"] = creds["auth_sock"] env["ALBEDO_WS_URL"] = creds["ws_url"] env["ALBEDO_WS_HEADERS"] = json.dumps({"X-Colab-Ssh-Pubkey": creds["pubkey"]}) return argv, env def supervise(): PID.write_text(str(os.getpid())) log(f"supervisor started (pid {os.getpid()})") stop = threading.Event() initial, _ = read_creds() if initial: threading.Thread(target=poll_runtime, args=(initial, stop), daemon=True).start() else: log("no creds yet; poller not started") while True: creds, problem = read_creds() if creds is None: log(f"idle: {problem}") time.sleep(RETRY_SECONDS) continue argv, env = build(creds) log(f"connecting to {creds['host']}") started = time.time() try: with LOG.open("a") as fh: code = subprocess.call(argv, env=env, stdout=fh, stderr=fh) log(f"tunnel exited code={code} after {time.time() - started:.0f}s") except Exception as error: # noqa: BLE001 - never let the supervisor die log(f"tunnel raised {type(error).__name__}: {error}") time.sleep(RETRY_SECONDS) def is_running(): try: pid = int(PID.read_text().strip()) os.kill(pid, 0) return pid except (OSError, ValueError): return None def main(argv=None): ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) ap.add_argument("--once", action="store_true") ap.add_argument("--status", action="store_true") ap.add_argument("--stop", action="store_true") a = ap.parse_args(argv) if a.status: pid = is_running() creds, problem = read_creds() print(f"supervisor: {'pid ' + str(pid) if pid else 'not running'}") print(f"creds: {CREDS} ({'ok' if creds else problem})") if LOG.exists(): print("last log lines:") for line in LOG.read_text().splitlines()[-6:]: print(" " + line) return 0 if a.stop: pid = is_running() if pid: os.kill(pid, 15) print(f"stopped pid {pid}") else: print("not running") return 0 if a.once: creds, problem = read_creds() if creds is None: print(f"cannot connect: {problem}") return 1 argv, env = build(creds) print("would run:", " ".join(shlex.quote(x) for x in argv)) return 0 running = is_running() if running: print(f"already running as pid {running}") return 0 supervise() return 0 if __name__ == "__main__": sys.exit(main())