| 1 | --- |
| 2 | log: |
| 3 | 2026-08-09: Added `--high-mem` flag (passthrough to session creation; sends `shape=hm` on assign when supported). |
| 4 | 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__"`. |
| 5 | 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. |
| 6 | 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`. |
| 7 | --- |
| 8 | |
| 9 | # Design: `colab run` — Shebang-Compatible One-Shot Execution |
| 10 | |
| 11 | ## Motivation |
| 12 | 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: |
| 13 | |
| 14 | 1. Allocates a Colab VM according to user-supplied flags (CPU / GPU / TPU). |
| 15 | 2. Executes the body of the file on that VM. |
| 16 | 3. Tears the VM down when execution finishes — UNLESS told otherwise. |
| 17 | |
| 18 | This is the natural ergonomic top-end of `colab-cli`: no boilerplate, no stale sessions, a single file is the unit of work. |
| 19 | |
| 20 | ## User Surface |
| 21 | |
| 22 | ``` |
| 23 | colab run [OPTIONS] SCRIPT [SCRIPT_ARGS]... |
| 24 | ``` |
| 25 | |
| 26 | | Flag | Type | Default | Purpose | |
| 27 | |---|---|---|---| |
| 28 | | `SCRIPT` | positional | — | Local path to a `.py` file. Required. | |
| 29 | | `SCRIPT_ARGS` | variadic | — | Extra args forwarded to the script as `sys.argv[1:]`. | |
| 30 | | `-s`, `--session` | str | auto | Name the ephemeral session (helpful with `--keep`). Auto-generated as `run-<6 hex>` if omitted. | |
| 31 | | `--gpu` | str | None | Same set as `colab new --gpu` (T4, L4, G4, H100, A100). | |
| 32 | | `--tpu` | str | None | Same set as `colab new --tpu` (v5e1, v6e1). | |
| 33 | | `--high-mem` | bool | False | Same as `colab new --high-mem` — request high-RAM when supported. | |
| 34 | | `--keep` | bool | False | Do **not** stop the session after the script finishes. | |
| 35 | | `--timeout` | float | 30.0 | Timeout in seconds for code execution to prevent hanging on silent tasks. | |
| 36 | |
| 37 | ### Shebang usage |
| 38 | With `--keep` and `--gpu` baked into the shebang line, an entire one-file workload becomes: |
| 39 | |
| 40 | ```python |
| 41 | #!/usr/bin/env -S colab run --gpu T4 |
| 42 | import torch |
| 43 | print(torch.cuda.get_device_name(0)) |
| 44 | ``` |
| 45 | |
| 46 | `chmod +x` and `./script.py` is then a single-step "rent a GPU, run, return". |
| 47 | |
| 48 | > 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. |
| 49 | |
| 50 | ## Behavior |
| 51 | |
| 52 | 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>`. |
| 53 | 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`: |
| 54 | ```python |
| 55 | import sys |
| 56 | sys.argv = ['<basename>', 'arg1', 'arg2', ...] |
| 57 | __name__ = '__main__' |
| 58 | ``` |
| 59 | Then executes the script body in the same kernel cell so any `if __name__ == "__main__":` guard fires. |
| 60 | 3. **Detect failure**: If the kernel returns any output of `output_type == "error"` (uncaught exception, syntax error, etc.) the CLI exits non-zero. |
| 61 | 4. **Tear down**: In a `finally` block, unless `--keep` was passed, the CLI: |
| 62 | - Sends `runtime.stop(shutdown_kernel=True)` (best-effort). |
| 63 | - Calls `state.client.unassign(endpoint)` to free the billable VM. |
| 64 | - Removes the session from `StateStore`. |
| 65 | - Kills the keep-alive daemon (`kill_process(s.keep_alive_pid)`). |
| 66 | - Logs `session_terminated` with `reason="run_completed"` (or `"run_failed"`). |
| 67 | |
| 68 | 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). |
| 69 | |
| 70 | ## AGENTS.md Constraints Honoured |
| 71 | - **Item 7 (no background threads)**: The keep-alive daemon is the existing detached process from `colab new`; this command introduces no new threads. |
| 72 | - **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. |
| 73 | - **Item 16 (daemon flag propagation)**: Reuses `spawn_keep_alive(...)` which already propagates `--auth` and `--config`. |
| 74 | - **Item 17 (persist-before-spawn)**: Uses the same persist-before-spawn pattern as `colab new`. |
| 75 | |
| 76 | ## Testing Strategy (TDD) |
| 77 | |
| 78 | ### Unit tests (`tests/test_run.py`) |
| 79 | 1. **`test_run_basic_flow`** — Happy path: create session, execute script, unassign on exit. Mocks `client.assign`, `client.unassign`, `ColabRuntime`. Asserts unassign is called. |
| 80 | 2. **`test_run_keep_skips_unassign`** — With `--keep`, `unassign` is NOT called and the session remains in the store. |
| 81 | 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']`. |
| 82 | 4. **`test_run_sets_dunder_main`** — The execute payload contains `__name__ = '__main__'`. |
| 83 | 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`. |
| 84 | 6. **`test_run_with_gpu_flag`** — `colab run --gpu T4 script.py` calls `client.assign(..., variant=GPU, accelerator=T4)`. |
| 85 | 7. **`test_run_missing_script_errors`** — `colab run` with no script path errors out (Typer-level). |
| 86 | 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. |
| 87 | 9. **`test_run_unassign_called_on_exception_during_execute`** — If `runtime.execute_code` raises, unassign is still called (try/finally guarantee). |
| 88 | |
| 89 | ### Integration test (`integration/repro_run_command/test.sh`) |
| 90 | - Write a tiny script that prints its argv and exits 0. |
| 91 | - Run `colab run /tmp/script.py hello world`. |
| 92 | - Assert stdout contains `argv=['script.py', 'hello', 'world']`. |
| 93 | - Assert `colab sessions` returns "No active sessions" afterward (cleanup happened). |
| 94 | - Repeat with `--keep`: assert the session shows up in `colab sessions`, then call `colab stop` to clean up. |