| 1 | # Copyright 2026 Google LLC |
| 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 | HIGH_MEM_ONLY_ACCELERATORS, |
| 45 | PostAssignmentResponse, |
| 46 | Shape, |
| 47 | ) |
| 48 | from colab_cli.commands.execution import _build_env_prelude, _parse_env_vars |
| 49 | from colab_cli.commands.session import ( |
| 50 | _is_scope_error, |
| 51 | _scope_remediation_message, |
| 52 | resolve_runtime_options, |
| 53 | spawn_keep_alive, |
| 54 | ) |
| 55 | from colab_cli.runtime import ColabRuntime |
| 56 | from colab_cli.state import SessionState |
| 57 | from colab_cli.utils import get_status_code, is_terminal_error |
| 58 | |
| 59 | |
| 60 | def _build_script_payload( |
| 61 | script_path: str, script_args: List[str], env_vars: Optional[dict[str, str]] = None |
| 62 | ) -> str: |
| 63 | """Wrap the script body so it executes with native-`python`-like semantics. |
| 64 | |
| 65 | Specifically: |
| 66 | - `sys.argv = [<basename>, *script_args]` so `argparse` etc. work. |
| 67 | - `__name__ = '__main__'` so `if __name__ == "__main__":` guards fire. |
| 68 | - Requested `--env KEY=VALUE` pairs are written into `os.environ`. |
| 69 | - Suppress the IPython UserWarning "To exit: use 'exit', 'quit', or |
| 70 | Ctrl-D." which fires whenever the script calls `sys.exit(...)`. This |
| 71 | warning is meaningful in an interactive REPL, but for `colab run` it |
| 72 | is pure noise that doesn't appear when running `python script.py`. |
| 73 | |
| 74 | The script body is appended verbatim; the prelude is short so any |
| 75 | traceback line numbers from user code remain close to the original. |
| 76 | """ |
| 77 | basename = os.path.basename(script_path) |
| 78 | with open(script_path, "r", encoding="utf-8") as f: |
| 79 | body = f.read() |
| 80 | |
| 81 | # `repr()` produces a safe, round-trippable Python literal for arbitrary |
| 82 | # strings (handles quotes, backslashes, non-ASCII). |
| 83 | argv_literal = f"[{', '.join(repr(x) for x in [basename] + script_args)}]" |
| 84 | |
| 85 | return ( |
| 86 | "import sys, warnings\n" |
| 87 | f"sys.argv = {argv_literal}\n" |
| 88 | "__name__ = '__main__'\n" |
| 89 | "warnings.filterwarnings('ignore', message=\"To exit: use\")\n" |
| 90 | + _build_env_prelude(env_vars or {}) |
| 91 | + _strip_shebang(body) |
| 92 | ) |
| 93 | |
| 94 | |
| 95 | def _extract_env_args_from_script_args( |
| 96 | script_args: List[str], |
| 97 | ) -> tuple[List[str], List[str]]: |
| 98 | """Pull colab's --env options out of variadic script args. |
| 99 | |
| 100 | `colab run` intentionally forwards unknown options after the script path to |
| 101 | the user's script. Since `--env` is now a colab option, support both |
| 102 | `colab run --env KEY=VALUE script.py` (Typer parses this) and |
| 103 | `colab run script.py --env KEY=VALUE` (this helper parses it). |
| 104 | """ |
| 105 | forwarded_args = [] |
| 106 | env = [] |
| 107 | i = 0 |
| 108 | while i < len(script_args): |
| 109 | arg = script_args[i] |
| 110 | if arg == "--env": |
| 111 | if i + 1 >= len(script_args): |
| 112 | env.append("") |
| 113 | i += 1 |
| 114 | else: |
| 115 | env.append(script_args[i + 1]) |
| 116 | i += 2 |
| 117 | elif arg.startswith("--env="): |
| 118 | env.append(arg.split("=", 1)[1]) |
| 119 | i += 1 |
| 120 | else: |
| 121 | forwarded_args.append(arg) |
| 122 | i += 1 |
| 123 | return forwarded_args, env |
| 124 | |
| 125 | |
| 126 | def _strip_shebang(body: str) -> str: |
| 127 | """Remove a leading `#!...\\n` if present. The remote kernel doesn't need |
| 128 | or understand it (it's a contract between the local kernel and the file's |
| 129 | executable bit), and leaving it in just adds noise. |
| 130 | """ |
| 131 | if body.startswith("#!"): |
| 132 | nl = body.find("\n") |
| 133 | return body[nl + 1 :] if nl != -1 else "" |
| 134 | return body |
| 135 | |
| 136 | |
| 137 | def _is_systemexit(out) -> bool: |
| 138 | """True iff this output is a `raise SystemExit(...)` (a.k.a. `sys.exit`).""" |
| 139 | return out.get("output_type") == "error" and out.get("ename") == "SystemExit" |
| 140 | |
| 141 | |
| 142 | def _systemexit_code(out) -> int: |
| 143 | """Map a SystemExit kernel output back to a CPython-style integer exit code. |
| 144 | |
| 145 | CPython conventions (mirrored): |
| 146 | - `sys.exit()` / `sys.exit(None)` / `sys.exit(0)` -> 0 |
| 147 | - `sys.exit(<int>)` -> <int> |
| 148 | - `sys.exit('msg')` (any non-int) -> 1 |
| 149 | """ |
| 150 | evalue = (out.get("evalue") or "").strip() |
| 151 | if evalue in ("", "None", "0"): |
| 152 | return 0 |
| 153 | try: |
| 154 | return int(evalue) |
| 155 | except ValueError: |
| 156 | return 1 |
| 157 | |
| 158 | |
| 159 | def _exit_code_from_outputs(outputs) -> int: |
| 160 | """Derive the CLI's exit code from the kernel's outputs for a single cell. |
| 161 | |
| 162 | A `SystemExit` is treated like CPython would treat the same call from a |
| 163 | plain `python script.py` invocation. Any *other* error (uncaught |
| 164 | exception, NameError, etc.) is exit 1. |
| 165 | """ |
| 166 | code = 0 |
| 167 | for o in outputs: |
| 168 | if o.get("output_type") != "error": |
| 169 | continue |
| 170 | if _is_systemexit(o): |
| 171 | ec = _systemexit_code(o) |
| 172 | # Last SystemExit wins, matching the runtime — and any non-zero |
| 173 | # eclipses any prior zero. |
| 174 | code = ec if ec != 0 else code |
| 175 | else: |
| 176 | return 1 |
| 177 | return code |
| 178 | |
| 179 | |
| 180 | def _make_run_output_hook(output_image=None): |
| 181 | """Build an `output_hook` for `runtime.execute_code` that: |
| 182 | - Routes normal output to `display_output` (stream/image/error). |
| 183 | - Suppresses the `SystemExit` traceback so `sys.exit(0)` is silent (it |
| 184 | wouldn't print anything under `python script.py` either) and |
| 185 | `sys.exit(N)` doesn't dump a noisy IPython-styled traceback when the |
| 186 | intent is "shell exit code N". |
| 187 | |
| 188 | The kernel still RETURNS the SystemExit output to us (so we can derive the |
| 189 | exit code in `_exit_code_from_outputs`); we just don't render it. |
| 190 | """ |
| 191 | # Imported here to avoid a circular import via execution.py at module load. |
| 192 | from colab_cli.commands.execution import display_output |
| 193 | |
| 194 | def hook(out): |
| 195 | if _is_systemexit(out): |
| 196 | return |
| 197 | display_output(out, output_image) |
| 198 | |
| 199 | return hook |
| 200 | |
| 201 | |
| 202 | def run_command( |
| 203 | ctx: typer.Context, |
| 204 | script: Annotated[ |
| 205 | str, |
| 206 | typer.Argument( |
| 207 | help="Path to a local Python file to execute on a fresh Colab VM." |
| 208 | ), |
| 209 | ], |
| 210 | script_args: Annotated[ |
| 211 | Optional[List[str]], |
| 212 | typer.Argument( |
| 213 | help=( |
| 214 | "Arguments forwarded to the script as sys.argv[1:]. " |
| 215 | "Anything after the script path is passed through verbatim." |
| 216 | ), |
| 217 | ), |
| 218 | ] = None, |
| 219 | session: Annotated[ |
| 220 | Optional[str], |
| 221 | typer.Option( |
| 222 | "-s", |
| 223 | "--session", |
| 224 | help=( |
| 225 | "Name for the ephemeral session (auto-generated if omitted). " |
| 226 | "Useful with --keep so you can attach later via `colab exec -s <name>`." |
| 227 | ), |
| 228 | ), |
| 229 | ] = None, |
| 230 | tpu: Annotated[ |
| 231 | Optional[str], |
| 232 | typer.Option(help="TPU accelerator variant. Supported: v5e1, v6e1."), |
| 233 | ] = None, |
| 234 | gpu: Annotated[ |
| 235 | Optional[str], |
| 236 | typer.Option( |
| 237 | help=( |
| 238 | "GPU accelerator variant. Supported: T4, L4, G4, H100, A100. " |
| 239 | "If omitted (along with --tpu), a CPU runtime is created." |
| 240 | ), |
| 241 | ), |
| 242 | ] = None, |
| 243 | high_mem: Annotated[ |
| 244 | bool, |
| 245 | typer.Option( |
| 246 | "--high-mem", |
| 247 | help=( |
| 248 | "Request a high-RAM machine shape. Requires Colab Pro or Pro+ " |
| 249 | "entitlement. Ignored for L4 and TPU accelerators." |
| 250 | ), |
| 251 | ), |
| 252 | ] = False, |
| 253 | keep: Annotated[ |
| 254 | bool, |
| 255 | typer.Option( |
| 256 | "--keep", |
| 257 | help=( |
| 258 | "Do not stop the session after the script finishes. The session " |
| 259 | "remains in `colab sessions` until you run `colab stop`." |
| 260 | ), |
| 261 | ), |
| 262 | ] = False, |
| 263 | timeout: Annotated[ |
| 264 | Optional[float], |
| 265 | typer.Option("--timeout", help="Timeout in seconds for code execution"), |
| 266 | ] = 30.0, |
| 267 | env: Annotated[ |
| 268 | Optional[List[str]], |
| 269 | typer.Option( |
| 270 | "--env", |
| 271 | help=( |
| 272 | "Set an environment variable in the remote kernel as KEY=VALUE. " |
| 273 | "Repeat for multiple variables." |
| 274 | ), |
| 275 | ), |
| 276 | ] = None, |
| 277 | ): |
| 278 | """Run a Python script on a fresh Colab VM, then release the VM |
| 279 | |
| 280 | Designed to be used as a shebang interpreter, e.g. |
| 281 | |
| 282 | #!/usr/bin/env -S colab run --gpu T4 |
| 283 | |
| 284 | so a single executable .py file can rent a GPU, run, and clean up after |
| 285 | itself. |
| 286 | """ |
| 287 | from colab_cli.common import state |
| 288 | |
| 289 | script_args = script_args or [] |
| 290 | script_args, inline_env = _extract_env_args_from_script_args(script_args) |
| 291 | env_vars = _parse_env_vars([*(env or []), *inline_env]) |
| 292 | |
| 293 | # AGENTS.md item 10: validate locally BEFORE allocating a VM. A typo'd |
| 294 | # script path should not cost the user real compute. |
| 295 | if not os.path.isfile(script): |
| 296 | typer.echo(f"[colab] Script not found: {script}", err=True) |
| 297 | raise typer.Exit(2) |
| 298 | |
| 299 | name = session or f"run-{uuid.uuid4().hex[:6]}" |
| 300 | variant, accelerator, shape = resolve_runtime_options( |
| 301 | gpu, tpu, high_mem=high_mem |
| 302 | ) |
| 303 | |
| 304 | if high_mem and accelerator in HIGH_MEM_ONLY_ACCELERATORS: |
| 305 | typer.echo( |
| 306 | "[colab] --high-mem ignored: this accelerator only offers one " |
| 307 | "machine shape.", |
| 308 | err=True, |
| 309 | ) |
| 310 | |
| 311 | typer.echo(f"[colab] Creating session '{name}'...", err=True) |
| 312 | try: |
| 313 | res = state.client.assign( |
| 314 | uuid.uuid4(), variant=variant, accelerator=accelerator, shape=shape |
| 315 | ) |
| 316 | except ColabRequestError as e: |
| 317 | # Mirror `colab new`'s friendly accelerator-quota message. |
| 318 | if get_status_code(e) == 400 and accelerator != Accelerator.NONE: |
| 319 | typer.echo( |
| 320 | f"[colab] Backend rejected accelerator '{accelerator.value}'. " |
| 321 | "You may not have quota or entitlement for this accelerator on " |
| 322 | "your account. Try a different one (e.g. --gpu T4) or omit " |
| 323 | "--gpu/--tpu for a CPU runtime.", |
| 324 | err=True, |
| 325 | ) |
| 326 | raise typer.Exit(code=1) |
| 327 | raise |
| 328 | |
| 329 | if isinstance(res, PostAssignmentResponse): |
| 330 | token = res.runtime_proxy_info.token |
| 331 | url = res.runtime_proxy_info.url |
| 332 | endpoint = res.endpoint |
| 333 | else: |
| 334 | token = ( |
| 335 | res.runtime_proxy_info.token |
| 336 | if hasattr(res, "runtime_proxy_info") |
| 337 | else getattr(res, "runtime_proxy_token", "") |
| 338 | ) |
| 339 | url = res.runtime_proxy_info.url if hasattr(res, "runtime_proxy_info") else "" |
| 340 | endpoint = res.endpoint |
| 341 | |
| 342 | s = SessionState( |
| 343 | name=name, |
| 344 | token=token, |
| 345 | url=url, |
| 346 | endpoint=endpoint, |
| 347 | variant=variant.value, |
| 348 | accelerator=accelerator.value, |
| 349 | machine_shape=( |
| 350 | Shape.HIGH_RAM.name if shape == Shape.HIGH_RAM else Shape.STANDARD.name |
| 351 | ), |
| 352 | ) |
| 353 | |
| 354 | # Pre-flight keep-alive: same scope-detection dance as `colab new` so a |
| 355 | # missing OAuth scope doesn't leak a billable assignment. |
| 356 | try: |
| 357 | state.client.keep_alive_assignment(endpoint) |
| 358 | except ColabRequestError as e: |
| 359 | if get_status_code(e) == 403 and _is_scope_error(e): |
| 360 | typer.echo( |
| 361 | "[colab] Keep-alive pre-flight failed: your credentials " |
| 362 | "are missing an OAuth scope required by Colab.\n", |
| 363 | err=True, |
| 364 | ) |
| 365 | typer.echo(_scope_remediation_message(state.auth_provider), err=True) |
| 366 | try: |
| 367 | state.client.unassign(endpoint) |
| 368 | except Exception: |
| 369 | pass |
| 370 | raise typer.Exit(code=1) |
| 371 | # Other failures: don't block — the daemon will retry. |
| 372 | |
| 373 | # AGENTS.md item 17: persist BEFORE spawning the daemon so the daemon's |
| 374 | # initial state.store.get(name) doesn't race the parent. |
| 375 | state.store.add(s) |
| 376 | s.keep_alive_pid = spawn_keep_alive( |
| 377 | endpoint, |
| 378 | name, |
| 379 | auth_provider=state.auth_provider, |
| 380 | config_path=state.config_path, |
| 381 | ) |
| 382 | state.store.add(s) |
| 383 | state.history.log_event( |
| 384 | name, |
| 385 | "session_created", |
| 386 | { |
| 387 | "endpoint": endpoint, |
| 388 | "variant": variant.value, |
| 389 | "accelerator": accelerator.value, |
| 390 | "machine_shape": s.machine_shape, |
| 391 | "via": "run", |
| 392 | }, |
| 393 | ) |
| 394 | typer.echo(f"[colab] Session READY ({name}). Executing {script}...", err=True) |
| 395 | |
| 396 | # ----- Execute the script ------------------------------------------------- |
| 397 | exit_code = 0 |
| 398 | cleanup_reason = "run_completed" |
| 399 | |
| 400 | def on_started(kid): |
| 401 | s.kernel_id = kid |
| 402 | state.store.add(s) |
| 403 | |
| 404 | def on_sess_started(sid): |
| 405 | s.session_id = sid |
| 406 | state.store.add(s) |
| 407 | |
| 408 | runtime = ColabRuntime( |
| 409 | s.url, |
| 410 | s.token, |
| 411 | kernel_id=s.kernel_id, |
| 412 | session_id=s.session_id, |
| 413 | on_kernel_started=on_started, |
| 414 | on_session_started=on_sess_started, |
| 415 | ) |
| 416 | |
| 417 | try: |
| 418 | # Same /content prelude as `colab exec` for consistency. |
| 419 | try: |
| 420 | runtime.execute_code( |
| 421 | "import os; os.makedirs('/content', exist_ok=True); " |
| 422 | "os.chdir('/content')" |
| 423 | ) |
| 424 | except Exception as e: |
| 425 | if is_terminal_error(e): |
| 426 | typer.echo( |
| 427 | f"[colab] Session '{name}' appears to be lost (404/401).", |
| 428 | err=True, |
| 429 | ) |
| 430 | state.prune_session(name) |
| 431 | raise typer.Exit(1) |
| 432 | raise |
| 433 | |
| 434 | payload = _build_script_payload(script, script_args, env_vars) |
| 435 | s.running = f"run({os.path.basename(script)})" |
| 436 | s.last_execution = ( |
| 437 | script, |
| 438 | None, |
| 439 | datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), |
| 440 | ) |
| 441 | state.store.add(s) |
| 442 | |
| 443 | try: |
| 444 | outputs = runtime.execute_code( |
| 445 | payload, output_hook=_make_run_output_hook(), timeout=timeout |
| 446 | ) |
| 447 | except Exception: |
| 448 | # Genuine transport-level failure. Cleanup still happens via the |
| 449 | # outer finally; surface non-zero exit so callers/CI notice. |
| 450 | exit_code = 1 |
| 451 | cleanup_reason = "run_failed" |
| 452 | raise |
| 453 | else: |
| 454 | exit_code = _exit_code_from_outputs(outputs) |
| 455 | if exit_code != 0: |
| 456 | cleanup_reason = "run_failed" |
| 457 | state.history.log_event( |
| 458 | name, |
| 459 | "execution", |
| 460 | {"code": payload, "outputs": outputs, "via": "run"}, |
| 461 | ) |
| 462 | finally: |
| 463 | s.running = None |
| 464 | state.store.add(s) |
| 465 | # Best-effort runtime close (keeps remote kernel alive for --keep). |
| 466 | try: |
| 467 | runtime.stop() |
| 468 | except Exception: |
| 469 | pass |
| 470 | |
| 471 | if not keep: |
| 472 | _teardown(name, s, reason=cleanup_reason) |
| 473 | |
| 474 | if exit_code != 0: |
| 475 | raise typer.Exit(exit_code) |
| 476 | |
| 477 | |
| 478 | def _teardown(name: str, s: SessionState, *, reason: str) -> None: |
| 479 | """Best-effort full session teardown: kill the keep-alive daemon, ask the |
| 480 | remote kernel to shut down, unassign the VM, and remove local state. |
| 481 | |
| 482 | Mirrors `commands.session.stop` but with a richer history event reason and |
| 483 | swallowing all errors (we don't want a teardown failure to mask the user's |
| 484 | exit code). |
| 485 | """ |
| 486 | from colab_cli.common import kill_process, state |
| 487 | |
| 488 | typer.echo(f"[colab] Stopping session '{name}'...", err=True) |
| 489 | if s.keep_alive_pid: |
| 490 | try: |
| 491 | kill_process(s.keep_alive_pid) |
| 492 | except Exception: |
| 493 | pass |
| 494 | |
| 495 | try: |
| 496 | rt = ColabRuntime(s.url, s.token, kernel_id=s.kernel_id) |
| 497 | rt.stop(shutdown_kernel=True) |
| 498 | except Exception: |
| 499 | pass |
| 500 | |
| 501 | try: |
| 502 | state.client.unassign(s.endpoint) |
| 503 | except Exception: |
| 504 | pass |
| 505 | |
| 506 | try: |
| 507 | state.store.remove(name) |
| 508 | except Exception: |
| 509 | pass |
| 510 | |
| 511 | try: |
| 512 | state.history.log_event(name, "session_terminated", {"reason": reason}) |
| 513 | except Exception: |
| 514 | pass |
| 515 | typer.echo("[colab] Session terminated.", err=True) |
| 516 | |
| 517 | |
| 518 | def register(app: typer.Typer) -> None: |
| 519 | # `context_settings` lets unknown args after the script path flow through |
| 520 | # as positional `script_args` so users can pass `--flags-for-the-script` |
| 521 | # without Typer trying to consume them. |
| 522 | app.command( |
| 523 | name="run", |
| 524 | context_settings={ |
| 525 | "allow_extra_args": True, |
| 526 | "ignore_unknown_options": True, |
| 527 | }, |
| 528 | )(run_command) |