Something went wrong. Try again.
strata but optimized for m5 (and maybe more later)
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081#!/usr/bin/env python3"""Run code inside a Colab runtime through its Jupyter kernel, the way `colab exec` does.
Colab's ssh proxy serves one WebSocket client per runtime and wedges server-side when a client dieswithout closing cleanly, which leaves no way in short of destroying the VM. The kernel path has nosuch limit: `jupyter_kernel_client` speaks the framing (it is the library google-colab-cli uses), sothis only has to hand it the runtime URL and the proxy token.
Usage: colab_kernel_exec.py --host HOST --token TOKEN --code "python""""import argparseimport jsonimport urllib.request
import jupyter_kernel_client as jkc
def make_client(host, token, kernel_id=None): kwargs = { "subprotocol": jkc.JupyterSubprotocol.DEFAULT, "extra_params": {"colab-runtime-proxy-token": token}, } server_url = f"https://{host}" headers = {"X-Colab-Client-Agent": "albedo", "X-Colab-Runtime-Proxy-Token": token} if hasattr(jkc, "ColabKernelClient"): cls, tok = jkc.ColabKernelClient, "proxy_token" elif hasattr(jkc, "KernelClient"): cls, tok = jkc.KernelClient, "token" else: cls, tok = jkc.JupyterKernelClient, "token" client = cls(server_url=server_url, **{tok: token}, kernel_id=kernel_id, client_kwargs=kwargs, headers=headers) # Never let the library reap a kernel it did not create. client._own_kernel = False return client
def run(host, token, code, timeout, kernel_id=None): if kernel_id: client = make_client(host, token, kernel_id) client.start() return client, kernel_id req = urllib.request.Request( f"https://{host}/api/kernels?colab-runtime-proxy-token={token}", data=json.dumps({"name": "python3"}).encode(), method="POST", headers={"Content-Type": "application/json"}) with urllib.request.urlopen(req, timeout=30) as resp: kernel_id = json.loads(resp.read())["id"] client = make_client(host, token, kernel_id) client.start() return client, kernel_id
def main(): ap = argparse.ArgumentParser() ap.add_argument("--host", required=True) ap.add_argument("--token", required=True) ap.add_argument("--code", default="") ap.add_argument("--file", default="") ap.add_argument("--timeout", type=float, default=120.0) ap.add_argument("--kernel-id", default="") a = ap.parse_args()
code = Path(a.file).read_text() if a.file else a.code client, kernel_id = run(a.host, a.token, code, a.timeout, a.kernel_id or None) print(f"KERNEL_ID {kernel_id}", flush=True) result = client.execute(code, timeout=a.timeout) for out in result.get("outputs", []): kind = out.get("output_type") if kind == "stream": print(out.get("text", ""), end="", flush=True) elif kind in ("execute_result", "display_data"): print(out.get("data", {}).get("text/plain", ""), flush=True) elif kind == "error": print("\n".join(out.get("traceback", [])), flush=True) print(f"\nEXEC_STATUS {result.get('status')}", flush=True)
if __name__ == "__main__": main()