#!/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 dies
without closing cleanly, which leaves no way in short of destroying the VM. The kernel path has no
such limit: `jupyter_kernel_client` speaks the framing (it is the library google-colab-cli uses), so
this only has to hand it the runtime URL and the proxy token.
Usage: colab_kernel_exec.py --host HOST --token TOKEN --code "python"
"""
import argparse
import json
import 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()