"""Pull files off a Colab runtime over the HTTP runtime proxy, with no ssh.
The ssh tunnel is a websocket bridged through `albedo_plugins.ws`, and Colab allows exactly **one**
ssh session per runtime. If the process holding that bridge dies without releasing it (a kernel or
daemon restart is enough), the runtime answers `429 {"error":"already-active-session"}` to every new
tunnel and there is no local remedy: `pgrep albedo_plugins.ws` is empty, the token is still valid for
an hour, and the API's only escape hatch is `/tun/m/unassign`, which destroys the VM.
The runtime proxy itself is a different path and is not locked. The same `colab-runtime-proxy-token`
authenticates plain HTTP against it, so the Jupyter REST API stays reachable while ssh is stuck:
GET /api/status -> {"connections":..,"kernels":..,"last_activity":..} (is it alive)
GET /api/contents/<p> -> directory listing, or one file as text
GET /files/<p> -> the file; honours Range, so `curl -C -` style resume works
Usage:
python3 scripts/pull_colab_http.py --host <runtime.netloc> --token <tok> \
--remote content/swift15/checkpoints/x/latest.pt --local out/latest.pt
`--host` is the netloc of `instance.url` and `--token` is `instance.token`; both expire after about an
hour, so re-list the instance rather than caching them.
"""
import argparse
import os
import sys
import time
import urllib.error
import urllib.request
CHUNK = 1 << 20 # 1 MiB: stream to disk, never hold the file in memory
def size_and_status(url, headers):
req = urllib.request.Request(url, headers={**headers, "Range": "bytes=0-0"})
try:
with urllib.request.urlopen(req, timeout=60) as r:
cr = r.headers.get("Content-Range") or ""
total = int(cr.split("/")[-1]) if "/" in cr else int(r.headers.get("Content-Length") or 0)
return total, r.status
except urllib.error.HTTPError as e:
return -1, e.code
def pull(url, local, headers, resume=True, quiet=False):
os.makedirs(os.path.dirname(os.path.abspath(local)) or ".", exist_ok=True)
# A checkpoint written every optimizer step keeps identical byte sizes (same tensor shapes), so
# "same size" is NOT "same content": resume-by-size would silently never refresh latest.pt.
# Pass resume=False for anything that is rewritten in place.
have = os.path.getsize(local) if (resume and os.path.exists(local)) else 0
total, status = size_and_status(url, headers)
if status == 404:
raise SystemExit(f"{local}: remote 404")
if total > 0 and have == total:
print(f"{local}: already complete ({total:,} bytes)")
return total
h = dict(headers)
if have:
h["Range"] = f"bytes={have}-"
t0 = time.time()
try:
r = urllib.request.urlopen(urllib.request.Request(url, headers=h), timeout=120)
except urllib.error.HTTPError as e:
if e.code == 416:
print(f"{local}: already complete")
return have
raise
with r as resp, open(local, "ab" if have else "wb") as fh:
mode = "resumed" if have else "new"
while True:
chunk = resp.read(CHUNK)
if not chunk:
break
fh.write(chunk)
have += len(chunk)
if not quiet:
pct = (100.0 * have / total) if total > 0 else 0.0
sys.stderr.write(f"\r {local}: {have:,}/{total:,} ({pct:5.1f}%)")
sys.stderr.flush()
dt = time.time() - t0
rate = (have / dt / 1e6) if dt else 0.0
print(f"{local}: {mode}, {have:,} bytes in {dt:.1f}s ({rate:.1f} MB/s)")
return have
def main(argv=None):
ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
ap.add_argument("--host", required=True, help="netloc of colab instance.url")
ap.add_argument("--token", required=True, help="colab instance.token (query param)")
ap.add_argument("--remote", required=True, help="path under the runtime's /content root")
ap.add_argument("--local", required=True, help="destination path on this host")
ap.add_argument("--no-resume", action="store_true")
ap.add_argument("--quiet", action="store_true")
a = ap.parse_args(argv)
remote = a.remote.lstrip("/")
url = f"https://{a.host}/files/{remote}?colab-runtime-proxy-token={a.token}"
headers = {"X-Colab-Client-Agent": "colab-cli"}
pull(url, a.local, headers, resume=not a.no_resume, quiet=a.quiet)
return 0
if __name__ == "__main__":
sys.exit(main())