Something went wrong. Try again.
strata but optimized for m5 (and maybe more later)
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104"""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 ordaemon restart is enough), the runtime answers `429 {"error":"already-active-session"}` to every newtunnel and there is no local remedy: `pgrep albedo_plugins.ws` is empty, the token is still valid foran 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 anhour, so re-list the instance rather than caching them."""import argparseimport osimport sysimport timeimport urllib.errorimport 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())