From f23e38ad1787a42841aa9f29f148df8b1dd5c6ac Mon Sep 17 00:00:00 2001 From: Duc Hieu Ho <150573299+hohieuai@users.noreply.github.com> Date: Sat, 24 Oct 2026 11:48:23 +1100 Subject: [PATCH] fix: enforce --timeout as an idle limit without busy-waiting `--timeout` (default 30 s) is meant as an idle threshold, so that "long-running silent tasks" can raise it and code that keeps printing is never cut off. Runtime.execute_code passed it to jupyter-kernel-client's execute_interactive, where it is a deadline for the whole run. Past that deadline, 0.9.0 polls without waiting and holds a CPU core at ~100% until the run ends; 1.0.2 raises TimeoutError even while output flows. For streaming execution without stdin (run, exec, repl), the runtime now waits for kernel messages itself through the client's public API (execute, get_iopub_msg, get_shell_msg), in blocking one-second slices. It raises the same TimeoutError("Timeout waiting for output") once `timeout` seconds pass with no message from the execution, and RuntimeError if the connection drops. The default and the flag are unchanged; executions that use stdin keep the library path. Fixes #183 --- docs/02_execution_and_interactive.md | 2 + docs/05_run_command.md | 1 + src/colab_cli/runtime.py | 61 +++++++++++- tests/mock_kernel.py | 102 ++++++++++++++++++++ tests/test_runtime.py | 22 +++-- tests/test_streaming.py | 137 ++++++++++++++------------- 6 files changed, 250 insertions(+), 75 deletions(-) create mode 100644 tests/mock_kernel.py diff --git a/docs/02_execution_and_interactive.md b/docs/02_execution_and_interactive.md index ed1d782a..298900dd 100644 --- a/docs/02_execution_and_interactive.md +++ b/docs/02_execution_and_interactive.md @@ -5,6 +5,8 @@ log: 2026-05-07: Fixed `print_kitty` (used by `colab exec --output-image` and any image-producing exec) to no-op when `sys.stdout.isatty()` is false. The Kitty Graphics Protocol escape sequence is meaningless when stdout is a file or pipe and was visually corrupting captured output (a multi-KB base64 PNG blob would land in log files, grep targets, or showboat captures). Image bytes are still saved to disk via `handle_image`'s file-write path; only the inline-render attempt is suppressed. 2026-06-04: Bumped the default `--timeout` for `colab exec` from 10s to 30s (and the matching `colab run` default) so brief silent tasks are less likely to hit a premature `TimeoutError`. Explicit `--timeout` overrides are unaffected. + +2026-10-24: `--timeout` for `colab exec` (and `colab run`) is now an idle limit: execution fails only after that many seconds without output, so code that keeps printing is never cut off. `colab repl`, which passes no timeout, waits until each cell finishes. Neither busy-loops while waiting. See #183. --- # Design: Execution and Interactive Interaction (`repl`, `exec`, `console`) diff --git a/docs/05_run_command.md b/docs/05_run_command.md index 9a1d4214..64e21c9e 100644 --- a/docs/05_run_command.md +++ b/docs/05_run_command.md @@ -5,6 +5,7 @@ log: 2026-05-12: Initial design and implementation of `colab run [args...]`. Combines `colab new` + `colab exec` + `colab stop` into a single fire-and-forget invocation so a Python file can use `#!/usr/bin/env -S colab run` as a shebang line and execute on a freshly-allocated Colab VM. Adds `--keep` (skip auto-stop), `--gpu` / `--tpu` (passthrough to session creation), `-s/--session` (name the ephemeral session), and propagates the script's exit status (non-zero on any uncaught exception in the kernel). The script's `sys.argv` is re-set inside the kernel to mirror native `python script.py arg1 arg2` semantics, and `__name__` is set to `"__main__"`. 2026-05-12: Native CPython exit-code semantics for `sys.exit()` / `raise SystemExit(...)` from the script body. The Colab kernel reports a `SystemExit` as `output_type=='error'`, which under the previous logic would have (a) printed the IPython traceback (`An exception has occurred, use %tb...`) and (b) flagged the run as a failure regardless of the integer exit code. Now: `sys.exit()` / `sys.exit(0)` exit 0 silently; `sys.exit(N)` exits N; `sys.exit('msg')` exits 1 (matching CPython). The IPython "To exit: use 'exit', 'quit', or Ctrl-D." UserWarning is filtered via the prelude. Encoded after running `examples/gpu_hello.py` end-to-end and seeing the noisy `SystemExit: 0` traceback at the end of an otherwise-successful GPU run. 2026-06-04: Bumped the default value of the `--timeout` flag from 10.0s to 30.0s so short-but-silent tasks aren't prematurely killed out of the box. Mirrors the same change for `colab exec`. +2026-10-24: `--timeout` is now enforced as an idle limit: the run fails only after that many seconds without output, so a script that keeps printing can run as long as it needs. The CLI waits for kernel messages itself, in one-second blocking slices, instead of relying on `jupyter-kernel-client`'s `execute_interactive`, whose deadline covered the whole run and which, in 0.9.0, held a CPU core at ~100% once the deadline passed. See #183. --- # Design: `colab run` — Shebang-Compatible One-Shot Execution diff --git a/src/colab_cli/runtime.py b/src/colab_cli/runtime.py index 57e09fde..7dadff99 100644 --- a/src/colab_cli/runtime.py +++ b/src/colab_cli/runtime.py @@ -13,6 +13,7 @@ # limitations under the License. import logging +import queue import time from typing import Any, Callable, Dict, List, Optional @@ -20,6 +21,10 @@ import requests +# The longest single wait for a kernel message, so a lost connection is noticed. +_WAIT_SLICE_SEC = 1.0 + + class ColabRuntime: def __init__( self, @@ -234,9 +239,12 @@ def wrapped_output_hook(msg): if idx < len(outputs): output_hook(outputs[idx]) - reply = self.kernel_client.execute_interactive( - code, output_hook=wrapped_output_hook, **kwargs - ) + if allow_stdin: + reply = self.kernel_client.execute_interactive( + code, output_hook=wrapped_output_hook, **kwargs + ) + else: + reply = self._execute_streaming(code, wrapped_output_hook, timeout) # execute_interactive returns the raw reply message reply_content = reply["content"] if reply else {"status": "error"} else: @@ -261,6 +269,53 @@ def wrapped_output_hook(msg): return outputs + def _execute_streaming( + self, + code: str, + output_hook: Callable[[Dict[str, Any]], None], + timeout: Optional[float], + ) -> Dict[str, Any]: + """Run ``code``, pass each output message to ``output_hook``, return the reply. + + ``timeout`` is an idle limit: how long to wait, in seconds, for the next + message from this execution. A run that keeps printing can take as long + as it needs, and ``None`` waits without a limit. The deadline in + ``jupyter_kernel_client.execute_interactive`` covers the whole run + instead, and past it 0.9.0 polls without waiting, holding a CPU core + until the run ends. + """ + client = self.kernel_client._manager.client + msg_id = client.execute(code, allow_stdin=False) + last_message = time.monotonic() + + def next_message(get_msg): + nonlocal last_message + while True: + if not client.connection_ready.is_set(): + raise RuntimeError("Connection was lost.") + wait = _WAIT_SLICE_SEC + if timeout is not None: + left = last_message + timeout - time.monotonic() + if left <= 0: + raise TimeoutError("Timeout waiting for output") + wait = min(wait, left) + try: + msg = get_msg(timeout=wait) + except queue.Empty: + continue + if msg["parent_header"].get("msg_id") == msg_id: + last_message = time.monotonic() + return msg + + while True: + msg = next_message(client.get_iopub_msg) + output_hook(msg) + if ( + msg["header"]["msg_type"] == "status" + and msg["content"].get("execution_state") == "idle" + ): + return next_message(client.get_shell_msg) + def stop(self, shutdown_kernel: bool = False): if self._kernel_client: try: diff --git a/tests/mock_kernel.py b/tests/mock_kernel.py new file mode 100644 index 00000000..2639856d --- /dev/null +++ b/tests/mock_kernel.py @@ -0,0 +1,102 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""A mock of jupyter-kernel-client's websocket client for streaming tests. + +MockKernelClient replays a script of kernel messages on a simulated clock, so +tests can cover long runs and long silences without waiting. +""" + +import queue +import threading +import time +import types +from unittest.mock import MagicMock + +from colab_cli import runtime as runtime_module + +REQUEST = "req-1" + + +def _msg(msg_type, content, parent=REQUEST): + return { + "header": {"msg_type": msg_type}, + "parent_header": {"msg_id": parent}, + "content": content, + } + + +def _stdout(text, parent=REQUEST): + return _msg("stream", {"name": "stdout", "text": text}, parent) + + +def _idle(parent=REQUEST): + return _msg("status", {"execution_state": "idle"}, parent) + + +def _reply(status="ok", **content): + return _msg("execute_reply", {"status": status, **content}) + + +class MockKernelClient: + """The websocket client's public API, on a simulated clock. + + `iopub` and `shell` are lists of (seconds after the request, message). A + get_*_msg call that would wait past the next message advances the clock by + `timeout` and raises queue.Empty, as a quiet kernel would. + """ + + def __init__(self, iopub, shell=None): + self.now = 0.0 + self.iopub = list(iopub) + self.shell = list( + shell if shell is not None else [(self._last(iopub), _reply())] + ) + self.waits = [] + self.connection_ready = threading.Event() + self.connection_ready.set() + + @staticmethod + def _last(messages): + return messages[-1][0] if messages else 0.0 + + def execute(self, code, **kwargs): + self.code, self.execute_kwargs = code, kwargs + return REQUEST + + def _get(self, messages, timeout): + self.waits.append(timeout) + assert timeout is not None and timeout > 0, ( + "must wait in bounded, non-zero slices" + ) + if messages and messages[0][0] <= self.now + timeout: + at, msg = messages.pop(0) + self.now = max(self.now, at) + return msg + self.now += timeout + raise queue.Empty + + def get_iopub_msg(self, timeout=None): + return self._get(self.iopub, timeout) + + def get_shell_msg(self, timeout=None): + return self._get(self.shell, timeout) + + +def connect_runtime(runtime, mock_kc, monkeypatch): + """Make `runtime` talk to `mock_kc`, with the runtime's clock set to its clock.""" + clock = types.SimpleNamespace(monotonic=lambda: mock_kc.now, sleep=time.sleep) + monkeypatch.setattr(runtime_module, "time", clock) + runtime._kernel_client = MagicMock() + runtime._kernel_client._manager.client = mock_kc diff --git a/tests/test_runtime.py b/tests/test_runtime.py index 8d439d0a..da3a7fb4 100644 --- a/tests/test_runtime.py +++ b/tests/test_runtime.py @@ -14,6 +14,8 @@ from unittest.mock import MagicMock, patch +import pytest + import jupyter_kernel_client from colab_cli.runtime import ColabRuntime @@ -160,18 +162,22 @@ def test_colab_runtime_execute_code_with_timeout(): assert kwargs.get("timeout") == 600 -def test_colab_runtime_execute_interactive_with_timeout(): +def test_colab_runtime_execute_interactive_with_timeout(monkeypatch): """timeout must also be plumbed through the execute_interactive branch - (used when an output_hook is supplied).""" - runtime = ColabRuntime("http://url", "token123") - mock_kc = MagicMock() - runtime._kernel_client = mock_kc + (used when an output_hook is supplied). There it is an idle limit: the + longest silence allowed between messages from the execution.""" + from mock_kernel import MockKernelClient, _idle, _stdout, connect_runtime - mock_kc.execute_interactive.return_value = {"content": {"status": "ok"}} + runtime = ColabRuntime("http://url", "token123") + mock_kc = MockKernelClient([(599, _stdout("1\n")), (599.5, _idle())]) + connect_runtime(runtime, mock_kc, monkeypatch) runtime.execute_code("print(1)", output_hook=lambda o: None, timeout=600) - _, kwargs = mock_kc.execute_interactive.call_args - assert kwargs.get("timeout") == 600 + runtime = ColabRuntime("http://url", "token123") + mock_kc = MockKernelClient([(601, _stdout("1\n")), (601.5, _idle())]) + connect_runtime(runtime, mock_kc, monkeypatch) + with pytest.raises(TimeoutError): + runtime.execute_code("print(1)", output_hook=lambda o: None, timeout=600) def test_colab_runtime_stop(): diff --git a/tests/test_streaming.py b/tests/test_streaming.py index 50e9f4c2..e71d3090 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -12,87 +12,96 @@ # See the License for the specific language governing permissions and # limitations under the License. -from unittest.mock import MagicMock +import pytest +from mock_kernel import MockKernelClient, _idle, _reply, _stdout, connect_runtime + from colab_cli.runtime import ColabRuntime -def test_runtime_execute_code_streaming(): - runtime = ColabRuntime("http://url", "token") - mock_client = MagicMock() +@pytest.fixture +def run(monkeypatch): + """Run execute_code against a MockKernelClient; returns (outputs, streamed, mock_kc).""" - # Mock the return value of execute_interactive (the raw reply) - mock_client.execute_interactive.return_value = {"content": {"status": "ok"}} + def _run(mock_kc, code="work()", timeout=30.0): + runtime = ColabRuntime("http://url", "token") + connect_runtime(runtime, mock_kc, monkeypatch) + streamed = [] + outputs = runtime.execute_code( + code, output_hook=streamed.append, timeout=timeout + ) + return outputs, streamed, mock_kc - # Inject our mock client - runtime._kernel_client = mock_client + return _run - streamed_outputs = [] - def output_hook(o): - streamed_outputs.append(o) +def test_runtime_execute_code_streaming(run): + mock_kc = MockKernelClient( + [(0.1, _stdout("1\n")), (0.2, _stdout("2\n")), (0.3, _idle())] + ) + outputs, streamed, _ = run(mock_kc, code="print(1); print(2)") - code = "print(1); print(2)" + assert [o["text"] for o in outputs] == ["1\n", "2\n"] + assert [o["text"] for o in streamed] == ["1\n", "2\n"] + assert mock_kc.code == "print(1); print(2)" - # We need to simulate the execution where execute_interactive is called - # and it calls our wrapped_output_hook. - # Since wrapped_output_hook is defined inside execute_code, we have to - # intercept the call to execute_interactive to get a reference to it. - def side_effect(code, output_hook=None, **kwargs): - # Simulate messages arriving - msg1 = { - "header": {"msg_type": "stream"}, - "content": {"name": "stdout", "text": "1\n"}, - } - msg2 = { - "header": {"msg_type": "stream"}, - "content": {"name": "stdout", "text": "2\n"}, - } - if output_hook: - output_hook(msg1) - output_hook(msg2) - return {"content": {"status": "ok"}} +def test_runtime_execute_code_streaming_error_synthesis(run): + reply = _reply( + "error", ename="RuntimeError", evalue="something went wrong", traceback=["tb"] + ) + mock_kc = MockKernelClient([(0.1, _idle())], shell=[(0.1, reply)]) + outputs, streamed, _ = run(mock_kc, code="fail()") - mock_client.execute_interactive.side_effect = side_effect + # The synthesized error is added after the run, so it is returned but not streamed. + assert len(outputs) == 1 + assert outputs[0]["output_type"] == "error" + assert outputs[0]["ename"] == "RuntimeError" + assert streamed == [] - outputs = runtime.execute_code(code, output_hook=output_hook) - assert len(outputs) == 2 - assert outputs[0]["text"] == "1\n" - assert outputs[1]["text"] == "2\n" +def test_a_run_that_keeps_printing_is_not_cut_off(run): + """--timeout is an idle limit: a 10-minute run that prints every 20 s finishes.""" + lines = [(t, _stdout(f"step {t}\n")) for t in range(20, 601, 20)] + mock_kc = MockKernelClient([*lines, (601, _idle())]) + outputs, _, _ = run(mock_kc, timeout=30.0) - # Verify streaming hook was called - assert len(streamed_outputs) == 2 - assert streamed_outputs[0]["text"] == "1\n" - assert streamed_outputs[1]["text"] == "2\n" + assert "".join(o["text"] for o in outputs).count("step") == 30 -def test_runtime_execute_code_streaming_error_synthesis(): - runtime = ColabRuntime("http://url", "token") - mock_client = MagicMock() +def test_silence_longer_than_the_timeout_raises(run): + mock_kc = MockKernelClient([(5, _stdout("loading\n")), (100, _idle())]) + with pytest.raises(TimeoutError, match="Timeout waiting for output"): + run(mock_kc, timeout=30.0) + assert ( + 35 <= mock_kc.now <= 36 + ) # 30 s after the last message, not 30 s after the start - # Simulate an error reply but NO error output message - mock_client.execute_interactive.return_value = { - "content": { - "status": "error", - "ename": "RuntimeError", - "evalue": "something went wrong", - "traceback": ["tb line 1"], - } - } - runtime._kernel_client = mock_client - streamed_outputs = [] - outputs = runtime.execute_code( - "fail()", output_hook=lambda o: streamed_outputs.append(o) - ) +def test_no_timeout_waits_through_any_silence(run): + mock_kc = MockKernelClient([(3600, _stdout("done\n")), (3600.5, _idle())]) + outputs, _, _ = run(mock_kc, timeout=None) + assert outputs[0]["text"] == "done\n" + + +def test_waiting_never_spins(run): + """Each wait blocks for up to a second, so a long run costs a few calls a second.""" + lines = [(t, _stdout("tick\n")) for t in range(25, 301, 25)] + mock_kc = MockKernelClient([*lines, (301, _idle())]) + run(mock_kc, timeout=30.0) + assert mock_kc.waits and all(w > 0 for w in mock_kc.waits) + assert len(mock_kc.waits) < 330 # about one wait per second over 300 s + + +def test_messages_from_other_requests_do_not_count_as_output(run): + others = [(t, _stdout("other\n", parent="req-0")) for t in range(5, 100, 5)] + mock_kc = MockKernelClient([*others, (100, _idle())]) + with pytest.raises(TimeoutError): + run(mock_kc, timeout=30.0) + assert 30 <= mock_kc.now <= 31 - # Final outputs should include synthesized error - assert len(outputs) == 1 - assert outputs[0]["output_type"] == "error" - assert outputs[0]["ename"] == "RuntimeError" - # Note: synthesized error is added AFTER execute_interactive returns, - # so it won't be in streamed_outputs unless we specifically add logic for it. - # Currently it's only in the returned list. - assert len(streamed_outputs) == 0 +def test_a_lost_connection_stops_the_wait(run): + mock_kc = MockKernelClient([(5, _stdout("x\n")), (100, _idle())]) + mock_kc.connection_ready.clear() + with pytest.raises(RuntimeError, match="Connection was lost"): + run(mock_kc, timeout=None)