"""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-<id>", "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())