Skip to content
Collimate

Cookbook: live Jupyter

02 / Live service

For data-science and notebook RL, the slow part of a rollout is often kernel startup: spinning up an IPython kernel, importing pandas / numpy / your analysis stack, and re-establishing the ZMQ sockets — paid before the first cell runs. Bake a live-service ready state: start a Jupyter kernel at bake time and leave it running, so every fork wakes with a live kernel and a rollout just executes cells against it.

The idea

Section titled “The idea”
bake:  boot guest → start kernel → block until kernel_info replies → SNAPSHOT → publish
fork:  restore snapshot → kernel is already alive → rollout executes cells

The warm command starts the kernel detached, then blocks until the kernel actually answers a kernel_info_request over ZMQ — not merely until the process exists. See the warm-command contract.

The warm script

Section titled “The warm script”

/usr/local/bin/collimate-warm launches an IPython kernel with a fixed connection file, then probes it for readiness before exiting:

#!/bin/sh
set -e


# 1. Bring loopback up. The kernel binds its ZMQ sockets to 127.0.0.1; at bake
#    time the guest's `lo` may not be up yet.
ip link set lo up 2>/dev/null || true


# 2. Launch DETACHED with a FIXED connection file, so both the readiness probe
#    and every rollout know where to find the kernel's ZMQ ports. setsid makes
#    it outlive this script and land in the snapshot.
export JUPYTER_CONNECTION_FILE=/tmp/kernel.json
setsid python3 -m ipykernel_launcher -f "$JUPYTER_CONNECTION_FILE" \
  >/var/log/kernel.log 2>&1 < /dev/null &


# 3. Block until the kernel actually answers kernel_info over ZMQ — the real
#    "serving" signal, not just "the connection file exists".
for i in $(seq 1 60); do
  if [ -f "$JUPYTER_CONNECTION_FILE" ] && python3 - <<'PY' 2>/dev/null
from jupyter_client import BlockingKernelClient
c = BlockingKernelClient(connection_file="/tmp/kernel.json")
c.load_connection_file(); c.start_channels()
c.kernel_info(); c.get_shell_msg(timeout=2)   # raises if the kernel isn't serving
PY
  then
    echo kernel-ready; exit 0
  fi
  sleep 1
done
echo kernel-timeout >&2; exit 1

Bake it, pointing the ready command at the script:

from collimate_rl import Ready, Template, connect


client = connect()   # a col_pro_* key via COLLIMATE_API_KEY


handle = Template.bake(
    client,
    name="jupyter",
    image="ghcr.io/you/jupyter@sha256:…",              # pin by digest
    ready=Ready(cmd="/usr/local/bin/collimate-warm"),  # launch detached → block → exit 0
    width=64,
)
handle.wait(timeout=600)

Run the warm script at build time as a self-test so a broken kernel is a red build, not a bad bake.

Using it from the SDK

Section titled “Using it from the SDK”

A rollout attaches to the live kernel using the baked connection file and runs cells against it — the imports and kernel spin-up are already paid:

from collimate_rl import connect


client = connect(api_key="col_...")


# The rollout body runs inside the sandbox. It connects to the kernel the
# ready-state snapshot is already running — no per-fork kernel launch.
ROLLOUT = r"""
from jupyter_client import BlockingKernelClient


kc = BlockingKernelClient(connection_file="/tmp/kernel.json")
kc.load_connection_file(); kc.start_channels()


def run_cell(src):
    msg_id = kc.execute(src)
    out = []
    while True:
        m = kc.get_iopub_msg(timeout=30)
        if m["parent_header"].get("msg_id") != msg_id:
            continue
        t = m["msg_type"]
        if t == "stream":
            out.append(m["content"]["text"])
        elif t == "execute_result":
            out.append(m["content"]["data"].get("text/plain", ""))
        elif t == "status" and m["content"]["execution_state"] == "idle":
            break
    return "".join(out)


# The kernel already has the analysis stack imported from bake — this is cheap.
print(run_cell("import pandas as pd; df = pd.DataFrame({'x':[1,2,3]}); df.x.sum()"))
"""


sb = client.create_sandbox("jupyter")
result = client.exec(sb["id"], commands=[["python3", "-c", ROLLOUT]], timeout_seconds=60)
print(result["stdout"])
client.delete_sandbox(sb["id"])

Fan a warmed parent out across many rollouts with fork — every child inherits the same live kernel with its state and imports intact:

# count is the per-call fan-out, capped by tier (Pro 32, Ent 64); call again for more.
children = client.fork(sb["id"], count=32)["children"]
for child in children:
    client.exec(child["id"], commands=[["python3", "-c", ROLLOUT]], timeout_seconds=60)

Reach it from outside

Section titled “Reach it from outside”

The kernel above binds its ZMQ sockets to 127.0.0.1, so it is reachable in-guest only — a connection arriving on the fork's NIC reaches nothing, even with the port baked. To drive a notebook from a client on your machine — the connect-from-outside access model — bake the Jupyter server (HTTP/WebSocket, not the bare kernel) bound to all interfaces:

  • In the warm script, launch the server on all interfaces with a token — e.g. jupyter lab --ip 0.0.0.0 --port 8888 --no-browser --ServerApp.token=<token> — detached, then block until it answers (the same launch-then-probe shape as the kernel above).
  • Bake it with ready=Ready(cmd="/usr/local/bin/collimate-warm", ports=[8888]) so the fork gets a guest NIC with 8888 exposed.

Then forward the port and point a Jupyter client or your browser at the local URL:

import urllib.request
from collimate_rl import connect


client = connect(api_key="col_...")


TOKEN = "…"   # the same token baked into the warm launch


with client.session("jupyter") as sb:
    with sb.port_forward(8888) as fwd:
        # fwd.url == http://127.0.0.1:<local>; open <fwd.url>/lab?token=<TOKEN> in a
        # browser, or point a Jupyter client at fwd.url with the token.
        req = urllib.request.Request(f"{fwd.url}/api", headers={"Authorization": f"token {TOKEN}"})
        print(urllib.request.urlopen(req, timeout=10).read())
# leaving the `with` deletes the sandbox

Next steps

Section titled “Next steps”