from __future__ import annotations
import json
import os
import selectors
import signal
import subprocess
import time
from collections import deque
from contextlib import suppress
from pathlib import Path
CELL_REPEATS = 50
SHELL_REPEATS = 5
MIXED_CELLS = 35
OUTPUT_BYTES = 32 * 1024
FRAME_ROWS = 10_000
FRAME_COLUMNS = 8
BASH_EMPTY = "_r = await bash(':')\nassert _r.exit_code == 0 and _r.output == ''"
GIT_STATUS = "_r = await bash('git status --porcelain')\nassert _r.exit_code == 0 and _r.output == ''"
BASH_OUTPUT = (
f"_r = await bash(\"printf '%{OUTPUT_BYTES}s' ''\")\n"
f"assert _r.exit_code == 0 and len(_r.output) == {OUTPUT_BYTES}\nprint(_r.output, end='')"
)
class Kernel:
def __init__(self, command: list[str], cwd: Path, env: dict[str, str], transcript: Path):
self.transcript = transcript
self.stderr = transcript.with_suffix(".stderr").open("wb")
self.selector = selectors.DefaultSelector()
self.raw: list[bytes] = []
self.pending = b""
self.lines: deque[bytes] = deque()
self.bytes = 0
self.sequence = 0
self.started = time.perf_counter()
try:
self.process = subprocess.Popen(
command,
cwd=cwd,
env=env,
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=self.stderr,
bufsize=0,
start_new_session=True,
)
except Exception:
self.stderr.close()
self.selector.close()
raise
assert self.process.stdout is not None
self.selector.register(self.process.stdout, selectors.EVENT_READ)
def event(self, deadline: float) -> dict:
while not self.lines:
remaining = deadline - time.perf_counter()
if remaining <= 0 or not self.selector.select(remaining):
raise TimeoutError("Timed out waiting for the kernel protocol")
assert self.process.stdout is not None
chunk = os.read(self.process.stdout.fileno(), 65536)
if not chunk:
raise RuntimeError("Kernel exited before completing the request")
self.bytes += len(chunk)
if self.bytes > 8_000_000:
raise RuntimeError("Kernel transcript exceeds 8 MB")
self.raw.append(chunk)
chunks = (self.pending + chunk).split(b"\n")
self.pending = chunks.pop()
self.lines.extend(chunks)
event = json.loads(self.lines.popleft())
if not isinstance(event, dict) or not isinstance(event.get("event"), str):
raise ValueError("Invalid kernel protocol event")
return event
def ready(self) -> tuple[float, str]:
event = self.event(time.perf_counter() + 30)
if event.get("event") != "ready" or event.get("protocol") != 3:
raise ValueError("Expected kernel protocol 3 ready event")
return time.perf_counter() - self.started, str(event.get("python", "unknown"))
def send(self, request: dict) -> None:
assert self.process.stdin is not None
data = memoryview((json.dumps(request) + "\n").encode())
while data:
written = self.process.stdin.write(data)
if not written:
raise RuntimeError("Kernel protocol write failed")
data = data[written:]
def done(self, request_id: str, *, interrupted: bool = False) -> list[dict]:
deadline = time.perf_counter() + 30
events = []
while True:
event = self.event(deadline)
events.append(event)
if event.get("event") == "done" and event.get("id") == request_id:
if interrupted:
if event.get("status") != "error" or not any(
e.get("ename") == "KeyboardInterrupt" and e.get("id") == request_id for e in events
):
raise RuntimeError("Kernel did not acknowledge the interrupt")
elif event.get("status") != "ok":
raise RuntimeError(f"Kernel request failed: {str(events)[-400:]}")
return events
def request(self, kind: str, **fields: object) -> list[dict]:
self.sequence += 1
request_id = str(self.sequence)
self.send({"type": kind, "id": request_id, **fields})
return self.done(request_id)
def execute(self, code: str) -> list[dict]:
return self.request("execute", code=code)
def batch(self, code: str, repeats: int, *, output_bytes: int | None = None) -> float:
started = time.perf_counter()
for _ in range(repeats):
events = self.execute(code)
if output_bytes is not None:
output = "".join(e.get("text", "") for e in events if e["event"] == "stdout")
if len(output.encode()) != output_bytes:
raise RuntimeError("Kernel output was truncated or missing")
return (time.perf_counter() - started) / repeats
def mixed(self) -> float:
started = time.perf_counter()
for index in range(MIXED_CELLS):
self.execute(GIT_STATUS if index % 4 == 0 else "pass")
return time.perf_counter() - started
def interrupt(self) -> float:
self.sequence += 1
request_id = str(self.sequence)
self.send(
{
"type": "execute",
"id": request_id,
"code": "import asyncio\nprint('interrupt-ready', flush=True)\nawait asyncio.sleep(60)",
}
)
deadline = time.perf_counter() + 30
while True:
event = self.event(deadline)
if (
event.get("id") == request_id
and event["event"] == "stdout"
and "interrupt-ready" in event.get("text", "")
):
break
if event["event"] in ("error", "done"):
raise RuntimeError("Interrupt fixture did not reach the running state")
started = time.perf_counter()
self.send({"type": "interrupt", "id": request_id})
self.done(request_id, interrupted=True)
elapsed = time.perf_counter() - started
self.execute("assert 1 + 1 == 2")
return elapsed
def close(self) -> None:
try:
if self.process.poll() is None:
self.send({"type": "shutdown"})
self.process.wait(timeout=5)
except (OSError, subprocess.TimeoutExpired):
with suppress(ProcessLookupError):
os.killpg(self.process.pid, signal.SIGKILL)
self.process.wait(timeout=5)
finally:
for stream in (self.process.stdin, self.process.stdout):
if stream is not None:
stream.close()
self.selector.close()
self.stderr.close()
self.transcript.with_suffix(".jsonl").write_bytes(b"".join(self.raw))