feat: Expose --timeout flag for exec and run commands (#38)
* feat: Expose --timeout flag for exec and run commands * fix: explicitly set default timeout to 10.0
Xiaoquan Kong committed
Jun 3, 2026 at 21:40 UTC
96ef983083d197455096988a1d2f7ec19be99892
7 files changed
+66
-6
docs/02_execution_and_interactive.md
+1
@@ -26,6 +26,7 @@ Execution involves sending Python code (or shell commands) to the Jupyter kernel
26
- If file path is local: Read content, send as code.
27
- If file path is remote: Execute `!python <path>`.
28
- **Multi-Modal Output**: Handle `display_data` messages (e.g., `image/png`, `text/html`). For the CLI, we'll save images to temporary files and print their paths, or if the terminal supports it (e.g., iTerm2), inline them.
29
+- **Timeout Configuration**: Exposes a `--timeout` flag (default 10s) to allow long-running silent tasks (like model compilation or data downloading) to execute without being prematurely killed.
30
31
### 3. Console (`colab console`)
32
- **Implementation**: Connects directly to the backend terminal endpoint (`/colab/tty`) via WebSockets using `websocket-client`.
docs/05_run_command.md
+1
@@ -29,6 +29,7 @@ colab run [OPTIONS] SCRIPT [SCRIPT_ARGS]...
29
| `--gpu` | str | None | Same set as `colab new --gpu` (T4, L4, G4, H100, A100). |
30
| `--tpu` | str | None | Same set as `colab new --tpu` (v5e1, v6e1). |
31
| `--keep` | bool | False | Do **not** stop the session after the script finishes. |
32
+| `--timeout` | float | 10.0 | Timeout in seconds for code execution to prevent hanging on silent tasks. |
33
34
### Shebang usage
35
With `--keep` and `--gpu` baked into the shebang line, an entire one-file workload becomes:
src/colab_cli/commands/execution.py
+6
-1
@@ -108,6 +108,9 @@ def exec_command(
108
output_image: Annotated[
109
Optional[str], typer.Option("--output-image", help="Path to save plot")
110
] = None,
111
+ timeout: Annotated[
112
+ Optional[float], typer.Option("--timeout", help="Timeout in seconds for code execution")
113
+ ] = 10.0,
114
):
115
"""Execute code in a session"""
116
from colab_cli.common import state
@@ -205,7 +208,9 @@ def exec_command(
208
state.store.add(s)
209
210
outputs = runtime.execute_code(
208
- code, output_hook=lambda o: display_output(o, output_image)
211
+ code,
212
+ output_hook=lambda o: display_output(o, output_image),
213
+ timeout=timeout
214
)
215
if "cell" in block:
216
save_output(outputs, block["cell"])
src/colab_cli/commands/run.py
+8
-1
@@ -233,6 +233,9 @@ def run_command(
233
),
234
),
235
] = False,
236
+ timeout: Annotated[
237
+ Optional[float], typer.Option("--timeout", help="Timeout in seconds for code execution")
238
+ ] = 10.0,
239
):
240
"""Run a Python script on a fresh Colab VM, then release the VM
241
@@ -386,7 +389,11 @@ def run_command(
389
state.store.add(s)
390
391
try:
389
- outputs = runtime.execute_code(payload, output_hook=_make_run_output_hook())
392
+ outputs = runtime.execute_code(
393
+ payload,
394
+ output_hook=_make_run_output_hook(),
395
+ timeout=timeout
396
+ )
397
except Exception:
398
# Genuine transport-level failure. Cleanup still happens via the
399
# outer finally; surface non-zero exit so callers/CI notice.
tests/test_exec.py
+23
-2
@@ -58,7 +58,7 @@ def test_cli_exec_file(mock_store, mock_runtime_class, mock_common_state, tmp_pa
58
mock_runtime.execute_code.assert_any_call(
59
"import os; os.makedirs('/content', exist_ok=True); os.chdir('/content')"
60
)
61
- mock_runtime.execute_code.assert_any_call("print('hello')", output_hook=ANY)
61
+ mock_runtime.execute_code.assert_any_call("print('hello')", output_hook=ANY, timeout=10.0)
62
63
64
def test_cli_exec_stdin(mock_store, mock_runtime_class, mock_common_state):
@@ -80,7 +80,7 @@ def test_cli_exec_stdin(mock_store, mock_runtime_class, mock_common_state):
80
assert mock_session.last_execution[1] is None
81
assert mock_session.last_execution[2] is not None
82
mock_store.add.assert_called_with(mock_session)
83
- mock_runtime.execute_code.assert_any_call("print(42)", output_hook=ANY)
83
+ mock_runtime.execute_code.assert_any_call("print(42)", output_hook=ANY, timeout=10.0)
84
85
86
def test_cli_exec_not_found(mock_common_state):
@@ -176,3 +176,24 @@ def test_cli_exec_lost_session_prunes(
176
assert result.exit_code == 1
177
assert "appears to be lost" in result.output
178
mock_common_state.prune_session.assert_called_once_with("lost-sess")
179
+
180
+
181
+def test_cli_exec_timeout(mock_store, mock_runtime_class, mock_common_state, tmp_path):
182
+ mock_session = MagicMock()
183
+ mock_session.url = "http://url"
184
+ mock_session.token = "token123"
185
+ mock_session.name = "s1"
186
+ mock_session.kernel_id = None
187
+ mock_session.session_id = None
188
+ mock_store.get.return_value = mock_session
189
+
190
+ mock_common_state.resolve_session.return_value = "s1"
191
+ mock_runtime = mock_runtime_class.return_value
192
+ mock_runtime.execute_code.return_value = [{"text": "hello\n"}]
193
+
194
+ script = tmp_path / "script.py"
195
+ script.write_text("print('hello')")
196
+
197
+ result = runner.invoke(app, ["exec", "-s", "s1", "-f", str(script), "--timeout", "3600"])
198
+ assert result.exit_code == 0
199
+ mock_runtime.execute_code.assert_any_call("print('hello')", output_hook=ANY, timeout=3600.0)
tests/test_ipynb_exec.py
+2
-2
@@ -98,10 +98,10 @@ class TestIpynbExec(unittest.TestCase):
98
"os.chdir", mock_runtime.execute_code.call_args_list[0].args[0]
99
)
100
mock_runtime.execute_code.assert_any_call(
101
- "print('cell 1')", output_hook=ANY
101
+ "print('cell 1')", output_hook=ANY, timeout=10.0
102
)
103
mock_runtime.execute_code.assert_any_call(
104
- "print('cell 2')", output_hook=ANY
104
+ "print('cell 2')", output_hook=ANY, timeout=10.0
105
)
106
107
@patch("colab_cli.commands.execution.ColabRuntime")
tests/test_run.py
+25
@@ -504,3 +504,28 @@ def test_run_prelude_suppresses_ipython_exit_warning(
504
# Look for the warnings filter targeting IPython's exit-warning text.
505
assert "warnings.filterwarnings" in body
506
assert "To exit: use" in body
507
+
508
+
509
+def test_run_with_timeout_flag(
510
+ mock_client,
511
+ mock_store,
512
+ mock_runtime_class,
513
+ mock_spawn_keep_alive,
514
+ assign_response,
515
+ script_path,
516
+):
517
+ """`colab run --timeout 3600 script.py` must pass timeout down to the runtime."""
518
+ mock_client.assign.return_value = assign_response
519
+ mock_runtime = mock_runtime_class.return_value
520
+ mock_runtime.execute_code.return_value = []
521
+
522
+ persisted = {}
523
+ mock_store.add.side_effect = lambda s: persisted.setdefault("s", s)
524
+ mock_store.get.side_effect = lambda name: persisted.get("s")
525
+
526
+ result = runner.invoke(app, ["run", "--timeout", "3600", str(script_path)])
527
+ assert result.exit_code == 0, result.output
528
+
529
+ code_calls = mock_runtime.execute_code.call_args_list
530
+ body_call = next(c for c in code_calls if "hello from script" in c.args[0])
531
+ assert body_call.kwargs.get("timeout") == 3600.0