google-colab-cli

repository·main·Indexed 20 days ago

https://github.com/googlecolab/google-colab-cli

A command-line interface for Google Colab that enables developers to provision high-performance runtimes (CPU, GPU, TPU), execute code, manage remote files, and automate cloud pipelines from a terminal. It supports session management, interactive REPLs, SSH proxy modes, and a keep-alive mechanism to prevent idle timeouts. Currently supports Linux and macOS.

Tokens
19.2K
Snippets
96
Records
117
Agent score
71%

What's inside google-colab-cli

  1. Use interactive console and REPL via pipes

    main

    The colab console and colab repl commands accept piped stdin, making them useful for programmatic troubleshooting or one-shot commands.

    • colab console: Connects to a tmux-wrapped pty. Note that output may contain terminal-control bytes; use grep -a to process it safely.
    • colab repl: Provides a one-shot Python REPL.

    Example: Querying disk usage via a pipe:

    echo 'import shutil; print(shutil.disk_usage("/").free // 2**30, "GB free")' | uv run colab --auth=adc repl -s debug
    echo 'import shutil; print(shutil.disk_usage("/").free // 2**30, "GB free")' | uv run colab --auth=adc repl -s debug
  2. How the Colab CLI keep-alive mechanism works

    main

    To prevent Colab VMs from being pruned due to idle timeouts (typically ~90 minutes), the CLI automatically manages a background keep-alive process when you create a session.

    Mechanism

    • Daemon Process: When colab new is executed, a detached background process is spawned to run a keep-alive command.
    • Tunnel Ping: Every 60 seconds, the daemon sends an HTTP GET request to the Colab Tunnel Frontend (/tun/m/<endpoint>/keep-alive/) with the header X-Colab-Tunnel: Google. This refreshes the LastActiveTime on the backend.
    • Timeout Handling: A ReadTimeout on the ping is treated as a success, as the Tunnel Frontend often records activity before the VM itself responds.
    • Termination: The daemon is terminated explicitly via colab stop, implicitly if the session is pruned, or automatically after 24 hours as a safety fallback.

    Troubleshooting Keep-Alive

    If you encounter issues with sessions being pruned, use colab log to view structured diagnostic logs, which include events like keep_alive_started, keep_alive_error, and keep_alive_stopped.

  3. How `colab run` manages session lifecycles

    main

    The colab run command follows a strict lifecycle to ensure resources are used efficiently and costs are minimized:

    1. Allocation: A fresh session is created (mirroring colab new). If --gpu or --tpu is specified, the session is allocated with that specific accelerator.
    2. Execution: The script is read and executed in a kernel cell. A prelude is prepended to the execution to ensure sys.argv and __name__ match standard Python execution.
    3. Failure Detection: If the kernel returns an error (uncaught exception, syntax error), the CLI flags the run as a failure and exits with a non-zero status.
    4. Teardown:
      • Default Behavior: A finally block ensures that even if the script fails, the CLI attempts to stop the runtime, unassign the VM to free billable resources, and remove the session from the StateStore.
      • With --keep: If the --keep flag is provided, the teardown step is skipped. The session remains active and visible in colab sessions and colab status, allowing you to interact with it later via colab exec or colab repl.
  4. Authentication strategies for Colab CLI

    main

    The Colab CLI supports two authentication strategies to interact with the Colab backend, selectable via the global --auth=<provider> flag. This determines how the CLI obtains credentials for the colab.research.google.com host.

    1. oauth2 (Default)

    Uses a remote copy-paste flow designed for local, remote, headless, or container environments.

    • How it works: The CLI prints an authorization URL. You sign in via your browser, copy the authorization code displayed by Google, and paste it back into the CLI prompt.
    • Storage: The refresh token is cached at ~/.config/colab-cli/token.json.
    • Note: If you have old cached tokens from before the remote flow update, you must delete ~/.config/colab-cli/token.json to trigger a fresh consent flow.

    2. adc (Application Default Credentials)

    Uses the standard Google discovery chain via google.auth.default(). This is ideal if you are running the CLI in an environment that already has ambient Google credentials (e.g., via GOOGLE_APPLICATION_CREDENTIALS, gcloud auth application-default login, or GCE/GKE metadata servers).

    Important for adc users: If using gcloud auth application-default login, you must explicitly include the required scopes, otherwise the CLI will fail. Use the following command to re-authenticate with the necessary scopes:

    gcloud auth application-default login \
        --scopes=openid,\nhttps://www.googleapis.com/auth/cloud-platform,\nhttps://www.googleapis.com/auth/userinfo.email,\nhttps://www.googleapis.com/auth/colaboratory
  5. Execute code on a Colab session

    main

    Once a session is running, you can execute code using several methods. Note that kernel state (imports, variables, functions) persists across calls in the same session.

    Run a local script

    colab exec -s <name> -f <script.py> reads the local file and sends it to the remote kernel.

    Run piped code

    echo "print(1)" | colab exec -s <name>

    Run a Jupyter Notebook

    colab exec -s <name> -f nb.ipynb runs each cell and writes the results to <basename>_output.ipynb locally.

    Shell commands

    echo "cmd" | colab console -s <name> wraps bash in tmux. For faster execution without a full shell, use exec with piped commands.

    Handling Images/Plots

    Use --output-image <path> with exec or repl to save intercepted PNG/JPEG outputs to a specific location.

    colab exec -s my-session -f script.py
  6. Provision, execute, and stop a Colab session

    main

    The standard lifecycle for a Colab CLI workflow follows a three-step pattern:

    1. Provision: Create a new session with a specific name using colab new -s <session_name>.
    2. Execute: Run code or scripts on the session using colab exec -s <session_name>.
    3. Stop: Terminate the session and release VM resources using colab stop -s <session_name>.

    Example workflow:

    # 1. Provision
    uv run colab --auth=adc new -s research
    
    # 2. Execute (passing code via stdin)
    cat <<'EOF' | uv run colab --auth=adc exec -s research
    print('Hello from Colab!')
    EOF
    
    # 3. Stop
    uv run colab --auth=adc stop -s research
    uv run colab --auth=adc new -s research
    cat <<'EOF' | uv run colab --auth=adc exec -s research
    print('Hello from Colab!')
    EOF
    uv run colab --auth=adc stop -s research
  7. Execute scripts with `colab run`

    main

    The colab run command provides a "one-shot" execution model. It automates the entire lifecycle of a Colab session: allocating a VM (CPU, GPU, or TPU), executing a local Python script, and automatically tearing down the VM once the script finishes. This allows you to treat a single Python file as a self-contained, cloud-hosted workload.

    Key behaviors:

    • Automatic Cleanup: By default, the VM is unassigned and the session is destroyed immediately after the script completes or fails.
    • Native Python Semantics: The command re-sets sys.argv and __name__ inside the kernel so that your script behaves exactly like it was run via python script.py [args]. This ensures if __name__ == "__main__": blocks work correctly.
    • Exit Codes: The CLI propagates the script's exit status. If the script encounters an uncaught exception, the CLI exits with a non-zero status. It also supports native CPython exit semantics (e.g., sys.exit(0) exits silently, while sys.exit(N) exits with code N).
    colab run [OPTIONS] SCRIPT [SCRIPT_ARGS]...
  8. Use `colab run` for Ephemeral Accelerator Jobs

    main

    The colab run command is designed for automated tasks where you want to provision hardware, run a script, and immediately release the VM. It handles the entire lifecycle (provisioning $\rightarrow$ execution $\rightarrow$ teardown) in one command.

    Example: Running a training script on a T4 GPU

    colab run --gpu T4 train.py

    Shebang Support: You can make a Python script executable directly by using the colab run interpreter in the shebang line. Use the --keep flag if you want to prevent the VM from being torn down immediately after execution.

    #!/usr/bin/env -S colab run --gpu L4 --keep
    import torch
    
    print("L4 GPU Available:", torch.cuda.is_available())
    print("Device Name:", torch.cuda.get_device_name(0))
  9. Use `colab repl` for interactive Python sessions

    main

    The colab repl command provides an interactive Python REPL by communicating with the Jupyter kernel on the Colab VM via the Jupyter Kernel Messaging Protocol.

    Key Behaviors:

    • Interactive Mode: Uses a local interactive console for standard input/output.
    • Piping Support: If stdin is not a TTY (e.g., you pipe a file into it), the command detects this and sends the entire input as a single execution request rather than entering interactive mode.

    Example (Piping):

    cat script.py | colab repl -s my-session
    cat script.py | colab repl -s my-session
  10. Configure Authentication for Colab CLI

    main

    Authentication is the most common point of failure. The CLI uses a global --auth flag which must be placed before the subcommand. The default is adc (Application Default Credentials).

    For headless or agent use, you must re-mint your ADC with specific scopes to ensure the CLI can manage sessions and keep-alive RPCs. Use the following command:

    gcloud auth application-default login \
      --scopes=openid,\nhttps://www.googleapis.com/auth/cloud-platform,\nhttps://www.googleapis.com/auth/userinfo.email,\nhttps://www.googleapis.com/auth/colaboratory

    Required Scopes:

    • userinfo.email: Required for the session backend.
    • colaboratory: Required for the RuntimeService keep-alive.
    • openid + cloud-platform: Mandated by gcloud.

    OAuth2 Setup

    Using colab --auth=oauth2 <subcommand> triggers a browser consent flow. This requires a human to interact with the browser and is generally not suitable for automated agents.

    Troubleshooting Auth

    • Use colab whoami to debug active email, scopes, audience, and expiry. If you see 403 errors against colab.pa.googleapis.com, check if the colaboratory scope is present.
    • Note: colab auth is not for CLI authentication; it is used to inject GCP credentials inside the VM kernel for notebook code. Do not use it to fix CLI 401/403 errors.
    gcloud auth application-default login \
      --scopes=openid,\nhttps://www.googleapis.com/auth/cloud-platform,\nhttps://www.googleapis.com/auth/userinfo.email,\nhttps://www.googleapis.com/auth/colaboratory
  11. Quick Start: Provision, Execute, and Stop a Session

    main

    To run a basic workflow on a CPU-based VM, follow these three steps:

    1. Provision: Create a new session.
    2. Execute: Send code to the session via stdin.
    3. Stop: Terminate the session and release resources.

    Note: If only one session is active, you can omit the -s, --session option; the CLI will automatically target the active session.

    # 1. Provision a new session
    colab new
    
    # 2. Execute code from stdin
    echo "print('Hello from Google Colab!')" | colab exec
    
    # 3. Stop and release the VM resource
    colab stop