@samitouri / QOSami-colab-cli-1 / commits / 6c49d59

Run shebang command (#7)

* Add `colab run <script.py>` shebang-friendly one-shot execution Combines `colab new` + `colab exec` + `colab stop` so a user can write a single executable .py file with a shebang line that allocates a Colab VM, runs the script with native `sys.argv` / `__name__ == '__main__'` semantics, and releases the VM on exit. `--keep` opts out of the auto-teardown for debugging. `--gpu` / `--tpu` mirror `colab new`. The script's exit status is propagated (non-zero on any uncaught exception), and cleanup runs in a `finally` so a failure can't leak a billable VM. * Honor CPython sys.exit() semantics in `colab run` The Colab kernel reports `raise SystemExit(N)` as an output of `output_type=='error'` with `ename=='SystemExit'`. Previously this caused two user-hostile behaviors: 1. The IPython traceback (`An exception has occurred, use %tb to see the full traceback. SystemExit: 0`) was printed at the end of every script that called `sys.exit()` \xe2\x80\x94 i.e. every well-formed CLI script. 2. The CLI exited 1 regardless of the integer arg, so `sys.exit(0)` was indistinguishable from a real failure. This commit: - Suppresses the SystemExit traceback in the run-mode output hook. - Maps the SystemExit `evalue` back to a CPython-style exit code (`None`/`0` -> 0, `<int>` -> N, anything else -> 1) and propagates it as the CLI's exit status. - Filters the IPython 'To exit: use exit, quit, or Ctrl-D' UserWarning from the script body via a one-liner in the prelude. Adds four regression tests covering each SystemExit shape. * AGENTS.md: shebang invocations bypass uv run, need explicit `uv tool install` Adds a note to item 8 capturing why the SystemExit-suppression fix appeared broken when run via `./examples/hello_colab.py` even though the editable install was up to date \xe2\x80\x94 the shebang's `#!/usr/bin/env -S colab run` resolves `colab` through `$PATH` to the uv-tool global, which was pinned to the previous commit. Future shebang-related testing must `uv tool install --reinstall --force --from . colab` first and verify the SHA with `colab version`. * comments from review - fix comprehension, docstring, add todo

Tyler committed May 12, 2026 at 15:05 UTC 6c49d59e066ed81f5185692327d4da33f1422e27

7 files changed +1204 -2

AGENTS.md

+2 -1

@@ -50,7 +50,7 @@

50 5. **Mocking Interactivity**: When testing commands that branch on `stdin.isatty()`, use the `is_stdin_tty` helper in `execution.py` and mock it via `mocker.patch("colab_cli.commands.execution.is_stdin_tty", return_value=...)`. This ensures tests don't hang in CI/agent environments.

51 6. **State Isolation**: Always patch the `colab_cli.common.state` singleton in tests to control session persistence and client behavior. Refer to `tests/conftest.py` for the standard global fixture.

52 7. **Fire-and-Forget Architecture**: The Colab CLI is a "fire-and-forget" tool. Avoid using background threads for long-running tasks within the main command flows. For persistent needs such as keep-alive, utilize detached background daemon processes (with PID tracking in the session state).

53 -8. **Verify the Local Install**: A globally-installed `colab` may exist on `PATH` (e.g. at `~/.local/bin/colab`) and can shadow the project's editable install when `uv run` is invoked from outside the repo. ALWAYS run shell commands with the repo as the working directory (e.g. via the `workdir` parameter, never `cd && cmd`) so `uv run colab ...` resolves to `.venv/bin/colab`. Confirm with `which colab` and `uv run which colab` if a CLI test produces unexpected results (e.g. flag-not-recognized errors for flags you just added).

53 +8. **Verify the Local Install**: A globally-installed `colab` may exist on `PATH` (e.g. at `~/.local/bin/colab`) and can shadow the project's editable install when `uv run` is invoked from outside the repo. ALWAYS run shell commands with the repo as the working directory (e.g. via the `workdir` parameter, never `cd && cmd`) so `uv run colab ...` resolves to `.venv/bin/colab`. Confirm with `which colab` and `uv run which colab` if a CLI test produces unexpected results (e.g. flag-not-recognized errors for flags you just added). **Shebang invocations always resolve via `$PATH`**, so a script like `#!/usr/bin/env -S colab run ...` will pick up the stale global tool even when the editable install is current — when testing shebang-based behavior after a code change, always run `uv tool install --reinstall --force --from . colab` first, then verify with `colab version` (the version string includes the git short SHA). Encoded 2026-05-12 after the SystemExit-suppression fix appeared not to work in `examples/hello_colab.py` because the shebang resolved to a uv-tool install pinned to the prior commit.

54 9. **Isolate the Regression First**: When a user reports an error in code you just touched, do NOT assume your change caused it. First, reproduce the failure on `main` (or the branch point) to determine whether the bug is pre-existing. Only after confirming the regression is yours should you start debugging the new code. Encoded after spending a turn debugging "ADC broke `colab new`" only to discover `colab new --gpu A100` was already failing on `main` due to an A100-quota-vs-default issue unrelated to ADC.

55 10. **Live Probes Allocate Real Resources**: Probing the Colab API to debug an issue creates real, billable assignments — every successful POST `/tun/m/assign` reserves a VM. Prefer GET-only (read) probes whenever possible. For any state-mutating call, (a) record every endpoint you create as you go, and (b) clean up via `client.unassign(endpoint)` (or `colab stop`) before declaring the investigation done. Then verify with `colab sessions` that nothing was orphaned.

56 11. **Push Freshness**: The remote may have advanced during a session (other contributors land commits while you work). ALWAYS `git fetch <remote>` immediately before pushing or merging. If `git log main..<remote>/main` is non-empty, reset local `main` to the remote, rebase feature branches onto it, retest, then push. NEVER force-push `main` to recover from divergence.

@@ -65,6 +65,7 @@

65 20. **Suggest the branch-diff review command after committing**: The user reviews changes with `git diff main..<branch-name>` (full cumulative diff against `main`, not just the latest commit). After landing one or more commits on a feature branch, ALWAYS suggest the exact command — e.g. "Review with `git diff main..sort-help-commands`" — instead of `git show <sha>` (which only shows a single commit and misses context when a branch has multiple commits). Encoded 2026-05-05 after suggesting `git show 9d9c7da` for a branch the user wanted to review holistically.

66 21. **Verify research-tool claims with primary sources**: Research tools and AI assistants can be confidently wrong, especially about edge cases or features outside their training corpus. When such a tool says "X is not used / not parsed / doesn't exist", treat it as a hypothesis to verify, not a fact. Always cross-check against the primary source (the actual code or config) — and when a tool names files, check whether the indirection chain it describes actually exists. The cost of believing the tool when it's wrong is shipping a non-functional feature; the cost of double-checking is small. Encoded 2026-05-05 after a `colab url` first-cut shipped the wrong URL format because of unverified output.

67 22. **Clean up orphaned assignments before finishing live tests**: After running a live integration test, `colab sessions` may show server-side assignments that the local `colab stop` couldn't see (e.g. assignments leaked from earlier in the session, or from crashed prior runs). Always run `colab sessions` as the final cleanup step, and for any `[?]`-marked orphan, run `python -c "from colab_cli.common import state; state.client.unassign('<endpoint>')"` (the CLI doesn't expose a direct unassign-by-endpoint command). Re-verify with `colab sessions` returning "No active sessions". Encoded 2026-05-05 after the first `colab url` live test left an orphan from a prior conversation turn that would have idled-out and billed compute units.

68 +23. **Forward unknown args through Typer with `context_settings`**: Typer/Click consumes any token starting with `-`/`--` as a flag of the parent command unless told otherwise. For a subcommand like `colab run script.py --some-script-flag` (where `--some-script-flag` belongs to the user's script, not to `colab`), declare it with `app.command(name="run", context_settings={"allow_extra_args": True, "ignore_unknown_options": True})` and accept the positional with `Annotated[Optional[List[str]], typer.Argument(...)] = None`. Also use `repr()` (not f-string interpolation) when embedding those forwarded strings into kernel-side Python source — `repr()` produces a safe round-trippable literal regardless of the user's shell-passed quotes, backslashes, or non-ASCII bytes. Encoded 2026-05-12 while implementing `colab run`.

70 ## Agent Execution Limitations (What I Can vs Cannot Run)

71 As an AI agent operating via non-interactive shell tools (`run_shell_command`), there are strict limits on what I can test autonomously without human intervention:

docs/05_run_command.md new

+90

@@ -0,0 +1,90 @@

1 +---

2 +log:

3 +2026-05-12: Initial design and implementation of `colab run <script.py> [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__"`.

4 +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.

5 +---

6 +

7 +# Design: `colab run` — Shebang-Compatible One-Shot Execution

8 +

9 +## Motivation

10 +Inspired by the `llm` shebang pattern (https://til.simonwillison.net/llms/llm-shebang), users should be able to write a single self-contained Python file with a shebang line that:

11 +

12 +1. Allocates a Colab VM according to user-supplied flags (CPU / GPU / TPU).

13 +2. Executes the body of the file on that VM.

14 +3. Tears the VM down when execution finishes — UNLESS told otherwise.

15 +

16 +This is the natural ergonomic top-end of `colab-cli`: no boilerplate, no stale sessions, a single file is the unit of work.

17 +

18 +## User Surface

19 +

20 +```

21 +colab run [OPTIONS] SCRIPT [SCRIPT_ARGS]...

22 +```

23 +

25 +|---|---|---|---|

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). |

32 +

33 +### Shebang usage

34 +With `--keep` and `--gpu` baked into the shebang line, an entire one-file workload becomes:

35 +

36 +```python

37 +#!/usr/bin/env -S colab run --gpu T4

38 +import torch

39 +print(torch.cuda.get_device_name(0))

40 +```

41 +

42 +`chmod +x` and `./script.py` is then a single-step "rent a GPU, run, return".

43 +

44 +> The `-S` flag of `env` is necessary on Linux/macOS to allow multiple words after `colab run` in a shebang line; without it the kernel passes the whole tail as one argument.

45 +

46 +## Behavior

47 +

48 +1. **Allocate**: Creates a fresh session (mirrors `colab new` end-to-end: `assign` → keep-alive pre-flight → spawn keep-alive daemon → persist `SessionState`). Session name defaults to `run-<6 hex>`.

49 +2. **Execute**: Reads the script file. Prepends a deterministic prelude that re-sets `sys.argv` and `__name__` so the script body sees the same execution context as `python script.py arg1 arg2`:

50 + ```python

51 + import sys

52 + sys.argv = ['<basename>', 'arg1', 'arg2', ...]

53 + __name__ = '__main__'

54 + ```

55 + Then executes the script body in the same kernel cell so any `if __name__ == "__main__":` guard fires.

56 +3. **Detect failure**: If the kernel returns any output of `output_type == "error"` (uncaught exception, syntax error, etc.) the CLI exits non-zero.

57 +4. **Tear down**: In a `finally` block, unless `--keep` was passed, the CLI:

58 + - Sends `runtime.stop(shutdown_kernel=True)` (best-effort).

59 + - Calls `state.client.unassign(endpoint)` to free the billable VM.

60 + - Removes the session from `StateStore`.

61 + - Kills the keep-alive daemon (`kill_process(s.keep_alive_pid)`).

62 + - Logs `session_terminated` with `reason="run_completed"` (or `"run_failed"`).

63 +

64 +If `--keep` is set, the session remains visible in `colab sessions` and `colab status` and can be reused with `colab exec -s <name>`, `colab repl -s <name>`, etc., until the user runs `colab stop` (or the keep-alive daemon hits its 24h cap).

65 +

66 +## AGENTS.md Constraints Honoured

67 +- **Item 7 (no background threads)**: The keep-alive daemon is the existing detached process from `colab new`; this command introduces no new threads.

68 +- **Item 10 (live probes allocate real resources)**: The teardown is in a `try/finally` so an exception during execution still releases the VM. Tests assert `unassign` is called even when the script errors.

69 +- **Item 16 (daemon flag propagation)**: Reuses `spawn_keep_alive(...)` which already propagates `--auth` and `--config`.

70 +- **Item 17 (persist-before-spawn)**: Uses the same persist-before-spawn pattern as `colab new`.

71 +

72 +## Testing Strategy (TDD)

73 +

74 +### Unit tests (`tests/test_run.py`)

75 +1. **`test_run_basic_flow`** — Happy path: create session, execute script, unassign on exit. Mocks `client.assign`, `client.unassign`, `ColabRuntime`. Asserts unassign is called.

76 +2. **`test_run_keep_skips_unassign`** — With `--keep`, `unassign` is NOT called and the session remains in the store.

77 +3. **`test_run_passes_argv`** — `colab run script.py a b c` results in a kernel `execute_code` call whose payload contains `sys.argv = ['script.py', 'a', 'b', 'c']`.

78 +4. **`test_run_sets_dunder_main`** — The execute payload contains `__name__ = '__main__'`.

79 +5. **`test_run_propagates_error_exit_code`** — When `runtime.execute_code` returns an output of `output_type == "error"`, the CLI exits non-zero AND still calls `unassign`.

80 +6. **`test_run_with_gpu_flag`** — `colab run --gpu T4 script.py` calls `client.assign(..., variant=GPU, accelerator=T4)`.

81 +7. **`test_run_missing_script_errors`** — `colab run` with no script path errors out (Typer-level).

82 +8. **`test_run_nonexistent_script_errors_before_assign`** — `colab run does-not-exist.py` MUST exit non-zero **without** calling `client.assign` so users don't burn a VM on a typo.

83 +9. **`test_run_unassign_called_on_exception_during_execute`** — If `runtime.execute_code` raises, unassign is still called (try/finally guarantee).

84 +

85 +### Integration test (`integration/repro_run_command/test.sh`)

86 +- Write a tiny script that prints its argv and exits 0.

87 +- Run `colab run /tmp/script.py hello world`.

88 +- Assert stdout contains `argv=['script.py', 'hello', 'world']`.

89 +- Assert `colab sessions` returns "No active sessions" afterward (cleanup happened).

90 +- Repeat with `--keep`: assert the session shows up in `colab sessions`, then call `colab stop` to clean up.

integration/repro_run_command/test.sh new

+132

@@ -0,0 +1,132 @@

1 +#!/bin/bash

3 +#

4 +# Licensed under the Apache License, Version 2.0 (the "License");

5 +# you may not use this file except in compliance with the License.

6 +# You may obtain a copy of the License at

7 +#

8 +# http://www.apache.org/licenses/LICENSE-2.0

9 +#

10 +# Unless required by applicable law or agreed to in writing, software

11 +# distributed under the License is distributed on an "AS IS" BASIS,

12 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

13 +# See the License for the specific language governing permissions and

14 +# limitations under the License.

15 +

16 +# Integration Test: `colab run <script.py> [args...]`

17 +#

18 +# Verifies the shebang-friendly one-shot execution flow:

19 +# 1. `colab run` allocates a CPU VM, runs the script, releases the VM.

20 +# 2. `sys.argv` and `__name__ == "__main__"` are honored.

21 +# 3. After the run finishes, no orphan VMs remain.

22 +# 4. `colab run --keep` leaves the session alive; `colab stop` clears it.

23 +

24 +# Don't `set -e` so we can capture failures and clean up explicitly.

25 +

26 +# ---------- Auth detection (mirrors integration/repro_keep_alive/test.sh) ----

27 +if [ -f "$HOME/.config/colab-cli/token.json" ]; then

28 + AUTH_FLAGS="--auth=oauth2"

29 +elif command -v gcloud > /dev/null && gcloud auth application-default print-access-token > /dev/null 2>&1; then

30 + ADC_TOKEN=$(gcloud auth application-default print-access-token 2>/dev/null)

31 + ADC_SCOPES=$(curl -s "https://www.googleapis.com/oauth2/v3/tokeninfo?access_token=$ADC_TOKEN" | python3 -c "import json,sys; print(json.load(sys.stdin).get('scope',''))" 2>/dev/null)

32 + if echo "$ADC_SCOPES" | grep -q "colaboratory" && echo "$ADC_SCOPES" | grep -q "userinfo.email"; then

33 + AUTH_FLAGS="--auth=adc"

34 + else

35 + echo "Error: ADC token lacks the required scopes (colaboratory + userinfo.email)."

36 + echo "Re-issue ADC creds with all required scopes:"

37 + echo " gcloud auth application-default login \\"

38 + echo " --scopes=openid,\\"

39 + echo " https://www.googleapis.com/auth/cloud-platform,\\"

40 + echo " https://www.googleapis.com/auth/userinfo.email,\\"

41 + echo " https://www.googleapis.com/auth/colaboratory"

42 + exit 1

43 + fi

44 +else

45 + echo "Error: No usable auth provider found."

46 + exit 1

47 +fi

48 +echo "[*] Using $AUTH_FLAGS"

49 +

50 +# ---------- Isolated session state -------------------------------------------

51 +TMP_DIR=$(mktemp -d)

52 +SESSION_FILE="$TMP_DIR/sessions.json"

53 +SCRIPT_PATH="$TMP_DIR/script.py"

54 +KEEP_SESSION_NAME="repro-run-keep-$(date +%s)"

55 +

56 +cleanup() {

57 + echo "[*] Cleaning up..."

58 + uv run colab $AUTH_FLAGS --config "$SESSION_FILE" stop -s "$KEEP_SESSION_NAME" 2>/dev/null || true

59 + rm -rf "$TMP_DIR"

60 +}

61 +trap cleanup EXIT

62 +

63 +cat > "$SCRIPT_PATH" <<'PYEOF'

64 +import sys

65 +print(f"argv={sys.argv}")

66 +print(f"is_main={__name__ == '__main__'}")

67 +PYEOF

68 +SCRIPT_BASENAME=$(basename "$SCRIPT_PATH")

69 +

70 +# ---------- Phase 1: basic run + auto-cleanup --------------------------------

71 +echo "[*] Phase 1: colab run <script.py> hello world"

72 +OUTPUT=$(uv run colab $AUTH_FLAGS --config "$SESSION_FILE" run "$SCRIPT_PATH" hello world 2>&1)

73 +RC=$?

74 +echo "$OUTPUT"

75 +

76 +if [ $RC -ne 0 ]; then

77 + echo "[FAILURE] colab run exited $RC"

78 + exit 1

79 +fi

80 +if ! echo "$OUTPUT" | grep -q "argv=\['$SCRIPT_BASENAME', 'hello', 'world'\]"; then

81 + echo "[FAILURE] argv was not propagated as expected."

82 + echo " Wanted substring: argv=['$SCRIPT_BASENAME', 'hello', 'world']"

83 + exit 1

84 +fi

85 +if ! echo "$OUTPUT" | grep -q "is_main=True"; then

86 + echo "[FAILURE] __name__ was not set to '__main__'."

87 + exit 1

88 +fi

89 +

90 +# Verify cleanup actually happened — no orphan assignments remain.

91 +SESSIONS_OUT=$(uv run colab $AUTH_FLAGS --config "$SESSION_FILE" sessions 2>&1)

92 +echo "$SESSIONS_OUT"

93 +if ! echo "$SESSIONS_OUT" | grep -q "No active sessions found on server."; then

94 + echo "[FAILURE] After auto-cleanup, server still reports active sessions."

95 + echo " (Possible orphan VM — investigate.)"

96 + exit 1

97 +fi

98 +echo "[SUCCESS] Phase 1 passed: argv passthrough, __main__, and auto-cleanup."

99 +

100 +# ---------- Phase 2: --keep leaves the session alive -------------------------

101 +echo ""

102 +echo "[*] Phase 2: colab run --keep -s $KEEP_SESSION_NAME <script.py>"

103 +OUTPUT=$(uv run colab $AUTH_FLAGS --config "$SESSION_FILE" run --keep -s "$KEEP_SESSION_NAME" "$SCRIPT_PATH" keep_arg 2>&1)

104 +RC=$?

105 +echo "$OUTPUT"

106 +

107 +if [ $RC -ne 0 ]; then

108 + echo "[FAILURE] colab run --keep exited $RC"

109 + exit 1

110 +fi

111 +if ! echo "$OUTPUT" | grep -q "argv=\['$SCRIPT_BASENAME', 'keep_arg'\]"; then

112 + echo "[FAILURE] --keep run did not produce the expected argv output."

113 + exit 1

114 +fi

115 +

116 +SESSIONS_OUT=$(uv run colab $AUTH_FLAGS --config "$SESSION_FILE" sessions 2>&1)

117 +echo "$SESSIONS_OUT"

118 +if ! echo "$SESSIONS_OUT" | grep -q "\[$KEEP_SESSION_NAME\]"; then

119 + echo "[FAILURE] --keep session $KEEP_SESSION_NAME not found in colab sessions."

120 + exit 1

121 +fi

122 +

123 +uv run colab $AUTH_FLAGS --config "$SESSION_FILE" stop -s "$KEEP_SESSION_NAME"

124 +SESSIONS_OUT=$(uv run colab $AUTH_FLAGS --config "$SESSION_FILE" sessions 2>&1)

125 +if ! echo "$SESSIONS_OUT" | grep -q "No active sessions found on server."; then

126 + echo "[FAILURE] After manual stop of $KEEP_SESSION_NAME, sessions remain."

127 + exit 1

128 +fi

129 +

130 +echo "[SUCCESS] Phase 2 passed: --keep persists the session, manual stop clears it."

131 +echo "[SUCCESS] All phases passed."

132 +exit 0

src/colab_cli/cli.py

+2 -1

@@ -23,7 +23,7 @@ from typing_extensions import Annotated

23 from colab_cli import auto_update

24 from colab_cli.auth import AuthProvider

25 from colab_cli.common import state, setup_logging

26 -from colab_cli.commands import session, execution, files, automation, utility

26 +from colab_cli.commands import session, execution, files, automation, run, utility

29 class AlphabeticalGroup(TyperGroup):

@@ -139,6 +139,7 @@ session.register(app)

139 execution.register(app)

140 files.register(app)

141 automation.register(app)

142 +run.register(app)

143 utility.register(app)

144

145

src/colab_cli/commands/run.py new

+471

@@ -0,0 +1,471 @@

2 +#

3 +# Licensed under the Apache License, Version 2.0 (the "License");

4 +# you may not use this file except in compliance with the License.

5 +# You may obtain a copy of the License at

6 +#

7 +# http://www.apache.org/licenses/LICENSE-2.0

8 +#

9 +# Unless required by applicable law or agreed to in writing, software

10 +# distributed under the License is distributed on an "AS IS" BASIS,

11 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

12 +# See the License for the specific language governing permissions and

13 +# limitations under the License.

14 +

15 +"""

16 +`colab run <script.py> [args...]` — shebang-friendly one-shot execution.

17 +

18 +Combines `colab new` + `colab exec` + `colab stop` into a single fire-and-forget

19 +invocation. The Python script's body runs in a freshly-allocated Colab kernel

20 +with `sys.argv` set as if it had been invoked via `python script.py [args...]`,

21 +and the VM is automatically released when the script finishes (unless `--keep`

22 +is passed).

23 +

24 +Designed to support shebangs:

25 +

26 + #!/usr/bin/env -S colab run --gpu T4

27 + import torch

28 + print(torch.cuda.get_device_name(0))

29 +

30 +See docs/05_run_command.md for the full design.

31 +"""

32 +

33 +import datetime

34 +import os

35 +import uuid

36 +from typing import List, Optional

37 +

38 +import typer

39 +from typing_extensions import Annotated

40 +

41 +from colab_cli.client import (

42 + Accelerator,

43 + ColabRequestError,

44 + PostAssignmentResponse,

45 + Variant,

46 +)

47 +from colab_cli.commands.session import (

48 + _is_scope_error,

49 + _scope_remediation_message,

50 + spawn_keep_alive,

51 +)

52 +from colab_cli.runtime import ColabRuntime

53 +from colab_cli.state import SessionState

54 +from colab_cli.utils import get_status_code, is_terminal_error

55 +

56 +

57 +# TODO(sethtroisi): dedupe this logic with similar in session.py

58 +def _resolve_accelerator(gpu: Optional[str], tpu: Optional[str]):

59 + """Mirror the mapping logic in `commands.session.new`. Centralised so the

60 + two commands stay in lock-step on supported accelerator names.

61 + """

62 + if tpu:

63 + variant = Variant.TPU

64 + accelerator = Accelerator.V5E1 if tpu.lower() == "v5e1" else Accelerator.V6E1

65 + return variant, accelerator

66 + if gpu:

67 + mapping = {

68 + "a100": Accelerator.A100,

69 + "h100": Accelerator.H100,

70 + "l4": Accelerator.L4,

71 + "t4": Accelerator.T4,

72 + "g4": Accelerator.G4,

73 + }

74 + return Variant.GPU, mapping.get(gpu.lower(), Accelerator.A100)

75 + return Variant.DEFAULT, Accelerator.NONE

76 +

77 +

78 +def _build_script_payload(script_path: str, script_args: List[str]) -> str:

79 + """Wrap the script body so it executes with native-`python`-like semantics.

80 +

81 + Specifically:

82 + - `sys.argv = [<basename>, *script_args]` so `argparse` etc. work.

83 + - `__name__ = '__main__'` so `if __name__ == "__main__":` guards fire.

84 + - Suppress the IPython UserWarning "To exit: use 'exit', 'quit', or

85 + Ctrl-D." which fires whenever the script calls `sys.exit(...)`. This

86 + warning is meaningful in an interactive REPL, but for `colab run` it

87 + is pure noise that doesn't appear when running `python script.py`.

88 +

89 + The script body is appended verbatim; the prelude is short so any

90 + traceback line numbers from user code remain close to the original.

91 + """

92 + basename = os.path.basename(script_path)

93 + with open(script_path, "r", encoding="utf-8") as f:

94 + body = f.read()

95 +

96 + # `repr()` produces a safe, round-trippable Python literal for arbitrary

97 + # strings (handles quotes, backslashes, non-ASCII).

98 + argv_literal = f"[{', '.join(repr(x) for x in [basename] + script_args)}]"

99 +

100 + return (

101 + "import sys, warnings\n"

102 + f"sys.argv = {argv_literal}\n"

103 + "__name__ = '__main__'\n"

104 + "warnings.filterwarnings('ignore', message=\"To exit: use\")\n"

105 + + _strip_shebang(body)

106 + )

107 +

108 +

109 +def _strip_shebang(body: str) -> str:

110 + """Remove a leading `#!...\\n` if present. The remote kernel doesn't need

111 + or understand it (it's a contract between the local kernel and the file's

112 + executable bit), and leaving it in just adds noise.

113 + """

114 + if body.startswith("#!"):

115 + nl = body.find("\n")

116 + return body[nl + 1 :] if nl != -1 else ""

117 + return body

118 +

119 +

120 +def _is_systemexit(out) -> bool:

121 + """True iff this output is a `raise SystemExit(...)` (a.k.a. `sys.exit`)."""

122 + return out.get("output_type") == "error" and out.get("ename") == "SystemExit"

123 +

124 +

125 +def _systemexit_code(out) -> int:

126 + """Map a SystemExit kernel output back to a CPython-style integer exit code.

127 +

128 + CPython conventions (mirrored):

129 + - `sys.exit()` / `sys.exit(None)` / `sys.exit(0)` -> 0

130 + - `sys.exit(<int>)` -> <int>

131 + - `sys.exit('msg')` (any non-int) -> 1

132 + """

133 + evalue = (out.get("evalue") or "").strip()

134 + if evalue in ("", "None", "0"):

135 + return 0

136 + try:

137 + return int(evalue)

138 + except ValueError:

139 + return 1

140 +

141 +

142 +def _exit_code_from_outputs(outputs) -> int:

143 + """Derive the CLI's exit code from the kernel's outputs for a single cell.

144 +

145 + A `SystemExit` is treated like CPython would treat the same call from a

146 + plain `python script.py` invocation. Any *other* error (uncaught

147 + exception, NameError, etc.) is exit 1.

148 + """

149 + code = 0

150 + for o in outputs:

151 + if o.get("output_type") != "error":

152 + continue

153 + if _is_systemexit(o):

154 + ec = _systemexit_code(o)

155 + # Last SystemExit wins, matching the runtime — and any non-zero

156 + # eclipses any prior zero.

157 + code = ec if ec != 0 else code

158 + else:

159 + return 1

160 + return code

161 +

162 +

163 +def _make_run_output_hook(output_image=None):

164 + """Build an `output_hook` for `runtime.execute_code` that:

165 + - Routes normal output to `display_output` (stream/image/error).

166 + - Suppresses the `SystemExit` traceback so `sys.exit(0)` is silent (it

167 + wouldn't print anything under `python script.py` either) and

168 + `sys.exit(N)` doesn't dump a noisy IPython-styled traceback when the

169 + intent is "shell exit code N".

170 +

171 + The kernel still RETURNS the SystemExit output to us (so we can derive the

172 + exit code in `_exit_code_from_outputs`); we just don't render it.

173 + """

174 + # Imported here to avoid a circular import via execution.py at module load.

175 + from colab_cli.commands.execution import display_output

176 +

177 + def hook(out):

178 + if _is_systemexit(out):

179 + return

180 + display_output(out, output_image)

181 +

182 + return hook

183 +

184 +

185 +def run_command(

186 + ctx: typer.Context,

187 + script: Annotated[

188 + str,

189 + typer.Argument(

190 + help="Path to a local Python file to execute on a fresh Colab VM."

191 + ),

192 + ],

193 + script_args: Annotated[

194 + Optional[List[str]],

195 + typer.Argument(

196 + help=(

197 + "Arguments forwarded to the script as sys.argv[1:]. "

198 + "Anything after the script path is passed through verbatim."

199 + ),

200 + ),

201 + ] = None,

202 + session: Annotated[

203 + Optional[str],

204 + typer.Option(

205 + "-s",

206 + "--session",

207 + help=(

208 + "Name for the ephemeral session (auto-generated if omitted). "

209 + "Useful with --keep so you can attach later via `colab exec -s <name>`."

210 + ),

211 + ),

212 + ] = None,

213 + tpu: Annotated[

214 + Optional[str],

215 + typer.Option(help="TPU accelerator variant. Supported: v5e1, v6e1."),

216 + ] = None,

217 + gpu: Annotated[

218 + Optional[str],

219 + typer.Option(

220 + help=(

221 + "GPU accelerator variant. Supported: T4, L4, G4, H100, A100. "

222 + "If omitted (along with --tpu), a CPU runtime is created."

223 + ),

224 + ),

225 + ] = None,

226 + keep: Annotated[

227 + bool,

228 + typer.Option(

229 + "--keep",

230 + help=(

231 + "Do not stop the session after the script finishes. The session "

232 + "remains in `colab sessions` until you run `colab stop`."

233 + ),

234 + ),

235 + ] = False,

236 +):

237 + """Run a Python script on a fresh Colab VM, then release the VM.

238 +

239 + Designed to be used as a shebang interpreter, e.g.

240 +

241 + #!/usr/bin/env -S colab run --gpu T4

242 +

243 + so a single executable .py file can rent a GPU, run, and clean up after

244 + itself.

245 + """

246 + from colab_cli.common import state

247 +

248 + script_args = script_args or []

249 +

250 + # AGENTS.md item 10: validate locally BEFORE allocating a VM. A typo'd

251 + # script path should not cost the user real compute.

252 + if not os.path.isfile(script):

253 + typer.echo(f"[colab] Script not found: {script}", err=True)

254 + raise typer.Exit(2)

255 +

256 + name = session or f"run-{uuid.uuid4().hex[:6]}"

257 + variant, accelerator = _resolve_accelerator(gpu, tpu)

258 +

259 + typer.echo(f"[colab] Creating session '{name}'...", err=True)

260 + try:

261 + res = state.client.assign(

262 + uuid.uuid4(), variant=variant, accelerator=accelerator

263 + )

264 + except ColabRequestError as e:

265 + # Mirror `colab new`'s friendly accelerator-quota message.

266 + if get_status_code(e) == 400 and accelerator != Accelerator.NONE:

267 + typer.echo(

268 + f"[colab] Backend rejected accelerator '{accelerator.value}'. "

269 + "You may not have quota or entitlement for this accelerator on "

270 + "your account. Try a different one (e.g. --gpu T4) or omit "

271 + "--gpu/--tpu for a CPU runtime.",

272 + err=True,

273 + )

274 + raise typer.Exit(code=1)

275 + raise

276 +

277 + if isinstance(res, PostAssignmentResponse):

278 + token = res.runtime_proxy_info.token

279 + url = res.runtime_proxy_info.url

280 + endpoint = res.endpoint

281 + else:

282 + token = (

283 + res.runtime_proxy_info.token

284 + if hasattr(res, "runtime_proxy_info")

285 + else getattr(res, "runtime_proxy_token", "")

286 + )

287 + url = res.runtime_proxy_info.url if hasattr(res, "runtime_proxy_info") else ""

288 + endpoint = res.endpoint

289 +

290 + s = SessionState(

291 + name=name,

292 + token=token,

293 + url=url,

294 + endpoint=endpoint,

295 + variant=variant.value,

296 + accelerator=accelerator.value,

297 + )

298 +

299 + # Pre-flight keep-alive: same scope-detection dance as `colab new` so a

300 + # missing OAuth scope doesn't leak a billable assignment.

301 + try:

302 + state.client.keep_alive_assignment(endpoint)

303 + except ColabRequestError as e:

304 + if get_status_code(e) == 403 and _is_scope_error(e):

305 + typer.echo(

306 + "[colab] Keep-alive pre-flight failed: your OAuth "

307 + "credentials are missing the 'colaboratory' scope, which "

308 + "is required by the Colab RuntimeService.\n",

309 + err=True,

310 + )

311 + typer.echo(_scope_remediation_message(state.auth_provider), err=True)

312 + try:

313 + state.client.unassign(endpoint)

314 + except Exception:

315 + pass

316 + raise typer.Exit(code=1)

317 + # Other failures: don't block — the daemon will retry.

318 +

319 + # AGENTS.md item 17: persist BEFORE spawning the daemon so the daemon's

320 + # initial state.store.get(name) doesn't race the parent.

321 + state.store.add(s)

322 + s.keep_alive_pid = spawn_keep_alive(

323 + endpoint,

324 + name,

325 + auth_provider=state.auth_provider,

326 + config_path=state.config_path,

327 + )

328 + state.store.add(s)

329 + state.history.log_event(

330 + name,

331 + "session_created",

332 + {

333 + "endpoint": endpoint,

334 + "variant": variant.value,

335 + "accelerator": accelerator.value,

336 + "via": "run",

337 + },

338 + )

339 + typer.echo(f"[colab] Session READY ({name}). Executing {script}...", err=True)

340 +

341 + # ----- Execute the script -------------------------------------------------

342 + exit_code = 0

343 + cleanup_reason = "run_completed"

344 +

345 + def on_started(kid):

346 + s.kernel_id = kid

347 + state.store.add(s)

348 +

349 + def on_sess_started(sid):

350 + s.session_id = sid

351 + state.store.add(s)

352 +

353 + runtime = ColabRuntime(

354 + s.url,

355 + s.token,

356 + kernel_id=s.kernel_id,

357 + session_id=s.session_id,

358 + on_kernel_started=on_started,

359 + on_session_started=on_sess_started,

360 + )

361 +

362 + try:

363 + # Same /content prelude as `colab exec` for consistency.

364 + try:

365 + runtime.execute_code(

366 + "import os; os.makedirs('/content', exist_ok=True); "

367 + "os.chdir('/content')"

368 + )

369 + except Exception as e:

370 + if is_terminal_error(e):

371 + typer.echo(

372 + f"[colab] Session '{name}' appears to be lost (404/401).",

373 + err=True,

374 + )

375 + state.prune_session(name)

376 + raise typer.Exit(1)

377 + raise

378 +

379 + payload = _build_script_payload(script, script_args)

380 + s.running = f"run({os.path.basename(script)})"

381 + s.last_execution = (

382 + script,

383 + None,

384 + datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),

385 + )

386 + state.store.add(s)

387 +

388 + try:

389 + outputs = runtime.execute_code(payload, output_hook=_make_run_output_hook())

390 + except Exception:

391 + # Genuine transport-level failure. Cleanup still happens via the

392 + # outer finally; surface non-zero exit so callers/CI notice.

393 + exit_code = 1

394 + cleanup_reason = "run_failed"

395 + raise

396 + else:

397 + exit_code = _exit_code_from_outputs(outputs)

398 + if exit_code != 0:

399 + cleanup_reason = "run_failed"

400 + state.history.log_event(

401 + name,

402 + "execution",

403 + {"code": payload, "outputs": outputs, "via": "run"},

404 + )

405 + finally:

406 + s.running = None

407 + state.store.add(s)

408 + # Best-effort runtime close (keeps remote kernel alive for --keep).

409 + try:

410 + runtime.stop()

411 + except Exception:

412 + pass

413 +

414 + if not keep:

415 + _teardown(name, s, reason=cleanup_reason)

416 +

417 + if exit_code != 0:

418 + raise typer.Exit(exit_code)

419 +

420 +

421 +def _teardown(name: str, s: SessionState, *, reason: str) -> None:

422 + """Best-effort full session teardown: kill the keep-alive daemon, ask the

423 + remote kernel to shut down, unassign the VM, and remove local state.

424 +

425 + Mirrors `commands.session.stop` but with a richer history event reason and

426 + swallowing all errors (we don't want a teardown failure to mask the user's

427 + exit code).

428 + """

429 + from colab_cli.common import kill_process, state

430 +

431 + typer.echo(f"[colab] Stopping session '{name}'...", err=True)

432 + if s.keep_alive_pid:

433 + try:

434 + kill_process(s.keep_alive_pid)

435 + except Exception:

436 + pass

437 +

438 + try:

439 + rt = ColabRuntime(s.url, s.token, kernel_id=s.kernel_id)

440 + rt.stop(shutdown_kernel=True)

441 + except Exception:

442 + pass

443 +

444 + try:

445 + state.client.unassign(s.endpoint)

446 + except Exception:

447 + pass

448 +

449 + try:

450 + state.store.remove(name)

451 + except Exception:

452 + pass

453 +

454 + try:

455 + state.history.log_event(name, "session_terminated", {"reason": reason})

456 + except Exception:

457 + pass

458 + typer.echo("[colab] Session terminated.", err=True)

459 +

460 +

461 +def register(app: typer.Typer) -> None:

462 + # `context_settings` lets unknown args after the script path flow through

463 + # as positional `script_args` so users can pass `--flags-for-the-script`

464 + # without Typer trying to consume them.

465 + app.command(

466 + name="run",

467 + context_settings={

468 + "allow_extra_args": True,

469 + "ignore_unknown_options": True,

470 + },

471 + )(run_command)

tests/conftest.py

@@ -34,5 +34,6 @@ def mock_common_state(mocker):

34 mocker.patch("colab_cli.commands.session.ColabRuntime")

35 mocker.patch("colab_cli.commands.execution.ColabRuntime")

36 mocker.patch("colab_cli.commands.automation.ColabRuntime")

37 + mocker.patch("colab_cli.commands.run.ColabRuntime")

39 return mock_state

tests/test_run.py new

+506

@@ -0,0 +1,506 @@

2 +#

3 +# Licensed under the Apache License, Version 2.0 (the "License");

4 +# you may not use this file except in compliance with the License.

5 +# You may obtain a copy of the License at

6 +#

7 +# http://www.apache.org/licenses/LICENSE-2.0

8 +#

9 +# Unless required by applicable law or agreed to in writing, software

10 +# distributed under the License is distributed on an "AS IS" BASIS,

11 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

12 +# See the License for the specific language governing permissions and

13 +# limitations under the License.

14 +

15 +"""Tests for `colab run <script.py> [args...]` — shebang-friendly one-shot

16 +execution that bundles `colab new` + `colab exec` + `colab stop`.

17 +"""

18 +

19 +from unittest.mock import MagicMock

20 +

21 +import pytest

22 +from typer.testing import CliRunner

23 +

24 +from colab_cli.cli import app

25 +from colab_cli.client import (

26 + Accelerator,

27 + PostAssignmentResponse,

28 + Variant,

29 +)

30 +

31 +runner = CliRunner()

32 +

33 +

34 +@pytest.fixture

35 +def mock_client(mock_common_state):

36 + return mock_common_state.client

37 +

38 +

39 +@pytest.fixture

40 +def mock_store(mock_common_state):

41 + return mock_common_state.store

42 +

43 +

44 +@pytest.fixture

45 +def mock_runtime_class(mocker):

46 + """Patch ColabRuntime in the run module specifically."""

47 + return mocker.patch("colab_cli.commands.run.ColabRuntime")

48 +

49 +

50 +@pytest.fixture

51 +def mock_spawn_keep_alive(mocker):

52 + """Don't actually spawn a daemon during tests."""

53 + return mocker.patch("colab_cli.commands.run.spawn_keep_alive", return_value=12345)

54 +

55 +

56 +@pytest.fixture

57 +def assign_response():

58 + """A minimal PostAssignmentResponse-shaped mock for client.assign."""

59 + res = MagicMock()

60 + res.__class__ = PostAssignmentResponse

61 + res.runtime_proxy_info.token = "tok"

62 + res.runtime_proxy_info.url = "http://runtime"

63 + res.endpoint = "ep-123"

64 + return res

65 +

66 +

67 +@pytest.fixture

68 +def script_path(tmp_path):

69 + p = tmp_path / "script.py"

70 + p.write_text("print('hello from script')\n")

71 + return p

72 +

73 +

74 +# ---------------------------------------------------------------------------

75 +# Happy path

76 +# ---------------------------------------------------------------------------

77 +

78 +

79 +def test_run_basic_flow(

80 + mock_client,

81 + mock_store,

82 + mock_runtime_class,

83 + mock_spawn_keep_alive,

84 + assign_response,

85 + script_path,

86 +):

87 + """`colab run script.py` should: assign, exec, unassign."""

88 + mock_client.assign.return_value = assign_response

89 + mock_runtime = mock_runtime_class.return_value

90 + mock_runtime.execute_code.return_value = []

91 +

92 + # Simulate the persisted SessionState being readable by the run command.

93 + persisted = {}

94 +

95 + def store_add(s):

96 + persisted["s"] = s

97 +

98 + def store_get(name):

99 + return persisted.get("s")

100 +

101 + mock_store.add.side_effect = store_add

102 + mock_store.get.side_effect = store_get

103 +

104 + result = runner.invoke(app, ["run", str(script_path)])

105 +

106 + assert result.exit_code == 0, result.output

107 + # Allocation happened

108 + mock_client.assign.assert_called_once()

109 + # Script body was executed (the prelude + body is one execute_code call)

110 + code_calls = [c.args[0] for c in mock_runtime.execute_code.call_args_list]

111 + assert any("hello from script" in code for code in code_calls), (

112 + f"Script body never sent to runtime. Calls: {code_calls}"

113 + )

114 + # Cleanup happened

115 + mock_client.unassign.assert_called_once_with("ep-123")

116 +

117 +

118 +# ---------------------------------------------------------------------------

119 +# --keep flag

120 +# ---------------------------------------------------------------------------

121 +

122 +

123 +def test_run_keep_skips_unassign(

124 + mock_client,

125 + mock_store,

126 + mock_runtime_class,

127 + mock_spawn_keep_alive,

128 + assign_response,

129 + script_path,

130 +):

131 + """With `--keep`, the session must NOT be unassigned after the script

132 + finishes — the user wants to attach to it later."""

133 + mock_client.assign.return_value = assign_response

134 + mock_runtime = mock_runtime_class.return_value

135 + mock_runtime.execute_code.return_value = []

136 +

137 + persisted = {}

138 + mock_store.add.side_effect = lambda s: persisted.setdefault("s", s)

139 + mock_store.get.side_effect = lambda name: persisted.get("s")

140 +

141 + result = runner.invoke(app, ["run", "--keep", str(script_path)])

142 +

143 + assert result.exit_code == 0, result.output

144 + mock_client.assign.assert_called_once()

145 + mock_client.unassign.assert_not_called()

146 + mock_store.remove.assert_not_called()

147 +

148 +

149 +# ---------------------------------------------------------------------------

150 +# argv passthrough

151 +# ---------------------------------------------------------------------------

152 +

153 +

154 +def test_run_passes_argv(

155 + mock_client,

156 + mock_store,

157 + mock_runtime_class,

158 + mock_spawn_keep_alive,

159 + assign_response,

160 + script_path,

161 +):

162 + """Args after the script must be exposed as `sys.argv` inside the kernel."""

163 + mock_client.assign.return_value = assign_response

164 + mock_runtime = mock_runtime_class.return_value

165 + mock_runtime.execute_code.return_value = []

166 +

167 + persisted = {}

168 + mock_store.add.side_effect = lambda s: persisted.setdefault("s", s)

169 + mock_store.get.side_effect = lambda name: persisted.get("s")

170 +

171 + result = runner.invoke(

172 + app, ["run", str(script_path), "alpha", "beta", "--flag-for-script"]

173 + )

174 +

175 + assert result.exit_code == 0, result.output

176 + code_calls = [c.args[0] for c in mock_runtime.execute_code.call_args_list]

177 + # The execute_code call that contains the script body must also set

178 + # sys.argv to mirror native python invocation.

179 + body_calls = [c for c in code_calls if "hello from script" in c]

180 + assert body_calls, f"Body never executed. Calls: {code_calls}"

181 + body = body_calls[0]

182 + assert "sys.argv" in body

183 + assert "'script.py'" in body

184 + assert "'alpha'" in body

185 + assert "'beta'" in body

186 + assert "'--flag-for-script'" in body

187 +

188 +

189 +def test_run_sets_dunder_main(

190 + mock_client,

191 + mock_store,

192 + mock_runtime_class,

193 + mock_spawn_keep_alive,

194 + assign_response,

195 + script_path,

196 +):

197 + """The script must run with __name__ == '__main__'."""

198 + mock_client.assign.return_value = assign_response

199 + mock_runtime = mock_runtime_class.return_value

200 + mock_runtime.execute_code.return_value = []

201 +

202 + persisted = {}

203 + mock_store.add.side_effect = lambda s: persisted.setdefault("s", s)

204 + mock_store.get.side_effect = lambda name: persisted.get("s")

205 +

206 + result = runner.invoke(app, ["run", str(script_path)])

207 + assert result.exit_code == 0, result.output

208 + code_calls = [c.args[0] for c in mock_runtime.execute_code.call_args_list]

209 + body = next(c for c in code_calls if "hello from script" in c)

210 + assert "__name__" in body and "'__main__'" in body

211 +

212 +

213 +# ---------------------------------------------------------------------------

214 +# Error handling

215 +# ---------------------------------------------------------------------------

216 +

217 +

218 +def test_run_propagates_error_exit_code(

219 + mock_client,

220 + mock_store,

221 + mock_runtime_class,

222 + mock_spawn_keep_alive,

223 + assign_response,

224 + script_path,

225 +):

226 + """If the kernel reports an error, the CLI must exit non-zero AND still

227 + unassign the VM (try/finally guarantee — AGENTS.md item 10)."""

228 + mock_client.assign.return_value = assign_response

229 + mock_runtime = mock_runtime_class.return_value

230 +

231 + def execute_with_error(code, output_hook=None, **kwargs):

232 + outputs = [

233 + {

234 + "output_type": "error",

235 + "ename": "ValueError",

236 + "evalue": "boom",

237 + "traceback": ["Traceback...\n", "ValueError: boom\n"],

238 + }

239 + ]

240 + if output_hook:

241 + for o in outputs:

242 + output_hook(o)

243 + return outputs

244 +

245 + mock_runtime.execute_code.side_effect = execute_with_error

246 +

247 + persisted = {}

248 + mock_store.add.side_effect = lambda s: persisted.setdefault("s", s)

249 + mock_store.get.side_effect = lambda name: persisted.get("s")

250 +

251 + result = runner.invoke(app, ["run", str(script_path)])

252 + assert result.exit_code != 0

253 + # Cleanup MUST happen even on script failure.

254 + mock_client.unassign.assert_called_once_with("ep-123")

255 +

256 +

257 +def test_run_unassign_called_on_exception_during_execute(

258 + mock_client,

259 + mock_store,

260 + mock_runtime_class,

261 + mock_spawn_keep_alive,

262 + assign_response,

263 + script_path,

264 +):

265 + """Even if `runtime.execute_code` raises (e.g. websocket dies), the VM

266 + must be released."""

267 + mock_client.assign.return_value = assign_response

268 + mock_runtime = mock_runtime_class.return_value

269 + mock_runtime.execute_code.side_effect = RuntimeError("websocket closed")

270 +

271 + persisted = {}

272 + mock_store.add.side_effect = lambda s: persisted.setdefault("s", s)

273 + mock_store.get.side_effect = lambda name: persisted.get("s")

274 +

275 + result = runner.invoke(app, ["run", str(script_path)])

276 + assert result.exit_code != 0

277 + mock_client.unassign.assert_called_once_with("ep-123")

278 +

279 +

280 +# ---------------------------------------------------------------------------

281 +# Accelerator passthrough

282 +# ---------------------------------------------------------------------------

283 +

284 +

285 +def test_run_with_gpu_flag(

286 + mock_client,

287 + mock_store,

288 + mock_runtime_class,

289 + mock_spawn_keep_alive,

290 + assign_response,

291 + script_path,

292 +):

293 + """`colab run --gpu T4 script.py` must request a T4 GPU."""

294 + mock_client.assign.return_value = assign_response

295 + mock_runtime = mock_runtime_class.return_value

296 + mock_runtime.execute_code.return_value = []

297 +

298 + persisted = {}

299 + mock_store.add.side_effect = lambda s: persisted.setdefault("s", s)

300 + mock_store.get.side_effect = lambda name: persisted.get("s")

301 +

302 + result = runner.invoke(app, ["run", "--gpu", "T4", str(script_path)])

303 + assert result.exit_code == 0, result.output

304 +

305 + _, kwargs = mock_client.assign.call_args

306 + assert kwargs["variant"] is Variant.GPU

307 + assert kwargs["accelerator"] is Accelerator.T4

308 +

309 +

310 +def test_run_with_tpu_flag(

311 + mock_client,

312 + mock_store,

313 + mock_runtime_class,

314 + mock_spawn_keep_alive,

315 + assign_response,

316 + script_path,

317 +):

318 + """`colab run --tpu v5e1 script.py` must request a TPU."""

319 + mock_client.assign.return_value = assign_response

320 + mock_runtime = mock_runtime_class.return_value

321 + mock_runtime.execute_code.return_value = []

322 +

323 + persisted = {}

324 + mock_store.add.side_effect = lambda s: persisted.setdefault("s", s)

325 + mock_store.get.side_effect = lambda name: persisted.get("s")

326 +

327 + result = runner.invoke(app, ["run", "--tpu", "v5e1", str(script_path)])

328 + assert result.exit_code == 0, result.output

329 +

330 + _, kwargs = mock_client.assign.call_args

331 + assert kwargs["variant"] is Variant.TPU

332 + assert kwargs["accelerator"] is Accelerator.V5E1

333 +

334 +

335 +# ---------------------------------------------------------------------------

336 +# Argument validation — fail FAST, before allocating a VM

337 +# ---------------------------------------------------------------------------

338 +

339 +

340 +def test_run_missing_script_errors(mock_client):

341 + """Typer should reject the invocation if no script path is given."""

342 + result = runner.invoke(app, ["run"])

343 + assert result.exit_code != 0

344 + mock_client.assign.assert_not_called()

345 +

346 +

347 +def test_run_nonexistent_script_errors_before_assign(mock_client):

348 + """If the script doesn't exist locally, fail BEFORE allocating a VM —

349 + otherwise a typo would burn billable compute."""

350 + result = runner.invoke(app, ["run", "/no/such/file.py"])

351 + assert result.exit_code != 0

352 + mock_client.assign.assert_not_called()

353 +

354 +

355 +# ---------------------------------------------------------------------------

356 +# SystemExit handling — the kernel reports `sys.exit(N)` as an error output of

357 +# `ename=='SystemExit'`. We want native-`python`-like semantics: exit 0 for

358 +# `SystemExit(0)` (no traceback printed), and propagate the integer for

359 +# `SystemExit(N)`.

360 +# ---------------------------------------------------------------------------

361 +

362 +

363 +def _systemexit_output(evalue: str):

364 + """Shape of the kernel's error output for `raise SystemExit(<evalue>)`."""

365 + return {

366 + "output_type": "error",

367 + "ename": "SystemExit",

368 + "evalue": evalue,

369 + "traceback": [

370 + "An exception has occurred, use %tb to see the full traceback.\n",

371 + f"\x1b[0;31mSystemExit\x1b[0m\x1b[0;31m:\x1b[0m {evalue}\n",

372 + ],

373 + }

374 +

375 +

376 +def test_run_systemexit_zero_treated_as_success(

377 + mock_client,

378 + mock_store,

379 + mock_runtime_class,

380 + mock_spawn_keep_alive,

381 + assign_response,

382 + script_path,

383 + capfd,

384 +):

385 + """`raise SystemExit(0)` from the script body must NOT make the CLI exit

386 + non-zero, AND the SystemExit traceback must NOT be printed (it's noise

387 + that doesn't appear when running `python script.py`)."""

388 + mock_client.assign.return_value = assign_response

389 + mock_runtime = mock_runtime_class.return_value

390 +

391 + def execute_with_systemexit(code, output_hook=None, **kwargs):

392 + outputs = [_systemexit_output("0")]

393 + if output_hook:

394 + for o in outputs:

395 + output_hook(o)

396 + return outputs

397 +

398 + mock_runtime.execute_code.side_effect = execute_with_systemexit

399 +

400 + persisted = {}

401 + mock_store.add.side_effect = lambda s: persisted.setdefault("s", s)

402 + mock_store.get.side_effect = lambda name: persisted.get("s")

403 +

404 + result = runner.invoke(app, ["run", str(script_path)])

405 + captured = capfd.readouterr()

406 +

407 + assert result.exit_code == 0, result.output

408 + # The IPython "An exception has occurred..." traceback must be suppressed.

409 + assert "An exception has occurred" not in (

410 + result.output + result.stderr + captured.out + captured.err

411 + )

412 + # Cleanup still happened.

413 + mock_client.unassign.assert_called_once_with("ep-123")

414 +

415 +

416 +def test_run_systemexit_nonzero_propagates_code(

417 + mock_client,

418 + mock_store,

419 + mock_runtime_class,

420 + mock_spawn_keep_alive,

421 + assign_response,

422 + script_path,

423 +):

424 + """`raise SystemExit(7)` from the script must surface as exit code 7

425 + (matching `python script.py` semantics)."""

426 + mock_client.assign.return_value = assign_response

427 + mock_runtime = mock_runtime_class.return_value

428 +

429 + def execute_with_systemexit(code, output_hook=None, **kwargs):

430 + outputs = [_systemexit_output("7")]

431 + if output_hook:

432 + for o in outputs:

433 + output_hook(o)

434 + return outputs

435 +

436 + mock_runtime.execute_code.side_effect = execute_with_systemexit

437 +

438 + persisted = {}

439 + mock_store.add.side_effect = lambda s: persisted.setdefault("s", s)

440 + mock_store.get.side_effect = lambda name: persisted.get("s")

441 +

442 + result = runner.invoke(app, ["run", str(script_path)])

443 + assert result.exit_code == 7

444 + mock_client.unassign.assert_called_once_with("ep-123")

445 +

446 +

447 +def test_run_systemexit_string_message_exits_one(

448 + mock_client,

449 + mock_store,

450 + mock_runtime_class,

451 + mock_spawn_keep_alive,

452 + assign_response,

453 + script_path,

454 +):

455 + """`sys.exit('boom')` (string arg, like `python -c "import sys; sys.exit(\"x\")"`)

456 + must (a) exit non-zero (CPython uses 1) and (b) print the message so the

457 + user sees what went wrong."""

458 + mock_client.assign.return_value = assign_response

459 + mock_runtime = mock_runtime_class.return_value

460 +

461 + def execute_with_systemexit(code, output_hook=None, **kwargs):

462 + outputs = [_systemexit_output("boom")]

463 + if output_hook:

464 + for o in outputs:

465 + output_hook(o)

466 + return outputs

467 +

468 + mock_runtime.execute_code.side_effect = execute_with_systemexit

469 +

470 + persisted = {}

471 + mock_store.add.side_effect = lambda s: persisted.setdefault("s", s)

472 + mock_store.get.side_effect = lambda name: persisted.get("s")

473 +

474 + result = runner.invoke(app, ["run", str(script_path)])

475 + assert result.exit_code == 1

476 + mock_client.unassign.assert_called_once_with("ep-123")

477 +

478 +

479 +def test_run_prelude_suppresses_ipython_exit_warning(

480 + mock_client,

481 + mock_store,

482 + mock_runtime_class,

483 + mock_spawn_keep_alive,

484 + assign_response,

485 + script_path,

486 +):

487 + """The prelude must mute IPython's 'To exit: use exit, quit, or Ctrl-D'

488 + UserWarning, which fires whenever the user calls `sys.exit(...)` (i.e.

489 + every well-formed CLI script)."""

490 + mock_client.assign.return_value = assign_response

491 + mock_runtime = mock_runtime_class.return_value

492 + mock_runtime.execute_code.return_value = []

493 +

494 + persisted = {}

495 + mock_store.add.side_effect = lambda s: persisted.setdefault("s", s)

496 + mock_store.get.side_effect = lambda name: persisted.get("s")

497 +

498 + result = runner.invoke(app, ["run", str(script_path)])

499 + assert result.exit_code == 0, result.output

500 +

501 + # Find the body-bearing execute_code call.

502 + code_calls = [c.args[0] for c in mock_runtime.execute_code.call_args_list]

503 + body = next(c for c in code_calls if "hello from script" in c)

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