Improve runtime shape selection, resource reporting and allocation error handling in Colab CLI

[hcsolakoglu]

Hi Colab CLI team,

I have been using google-colab-cli==0.5.9 heavily from terminal-based coding agents on a Colab Pro+ account. The tool is very useful, but a few missing surfaces make it harder for agents and automation to
use safely and predictably.

Environment

• Package: google-colab-cli==0.5.9
• Install: uv tool install google-colab-cli
• Platform: Linux
• Account: Colab Pro+
• Commands used: colab run, colab new, colab sessions, colab status

Observations

1. High-memory / machine shape cannot be selected from the CLI

The Colab web UI exposes high-memory choices for some runtimes, for example CPU-only, T4, and A100.

In the installed CLI source, client.py defines:

• Shape.STANDARD
• Shape.HIGH_RAM

and listed assignments expose machineShape.

But colab new and colab run only expose:

• --gpu
• --tpu

and assignment creation appears to send only variant and accelerator, not machine shape.

This makes it impossible to request high-memory CPU/T4/A100 from the CLI, even though the web UI supports it.

Request:

• Add a shape option, for example:
  • --high-mem
  • or --shape standard|high-ram
• Show shape in colab status and colab sessions.

2. CU/hour rate is not exposed through the CLI

For automation, it is important to know the active runtime’s current compute-unit usage rate before starting long jobs.

Currently I could not find a CLI command that reports:

• current CU/hour for an active runtime
• remaining compute units / balance
• estimated runtime cost

The web UI shows approximate usage rates, but terminal agents cannot safely read that.

Request:

• Add current runtime CU/hour to colab status.
• Optionally add account balance / compute-unit info if available.
• If exact values are not available, expose the same approximate values shown in the Colab UI.

3. Allocation failures for scarce GPUs are raw and hard to automate

For scarce accelerators like G4, A100, and H100, allocation can return Service Unavailable.

Example pattern:

ColabRequestError: Failed to issue request POST
https://colab.research.google.com/tun/m/assign?...&variant=GPU&accelerator=H100: Service Unavailable

Observed behavior:

• T4 allocated on first attempt.
• L4 allocated on first attempt.
• G4 allocated after several retries.
• A100 allocated after several retries.
• H100 did not allocate after many retry attempts.

Request:

• Convert allocation failures into concise, machine-readable output.

• Consider an optional retry/backoff mode:

  • colab run --gpu H100 --retry 80 --retry-delay 20 ...
  • or colab new --gpu H100 --wait --timeout ...

• Return structured reason codes where possible: quota, entitlement, temporary_unavailable, capacity, etc.

4. CLI selector does not guarantee actual assigned hardware details

The selector is useful, but agents need to know what was actually assigned.

Examples from one Pro+ account:

• --gpu G4 assigned NVIDIA RTX PRO 6000 Blackwell Server Edition, ~98 GB VRAM.
• --gpu A100 assigned NVIDIA A100-SXM4-40GB.
• The web UI also showed an A100 high-memory option with NVIDIA A100-SXM4-80GB.

Request:

• Include actual assigned hardware details in colab status:
  • GPU model
  • GPU memory
  • driver / CUDA if available
  • TPU device kind
  • TPU HBM/device memory
  • CPU/RAM shape
  • disk size
  • machine shape

5. Better cleanup / interruption ergonomics for automation

For long retry loops or agent-driven workflows, a user interrupt can leave uncertainty about local wrapper processes or remote assignments.

Request:

• Provide a single cleanup command or mode that stops a named session and confirms server-side assignment state.
• Consider clearer documentation for interruption behavior and keep-alive behavior.

Why this matters

AI coding agents increasingly run workloads from terminals. For Colab CLI to be reliable in that setting, agents need to:

• request the right shape,
• estimate cost before long jobs,
• retry scarce accelerators safely,
• verify actual assigned resources,
• and clean up confidently after interruption.

These improvements would make the CLI much safer for automated ML workflows.

Thanks for building this tool.

[ddjukic]

Chiming in to strongly support section 2 (compute-unit rate / balance not exposed via the CLI) — we're independently driving google-colab-cli (0.5.11) from terminal coding agents on Colab Pro+, and this is the single biggest gap for safe automation. An agent can colab new --gpu A100 and run for minutes with no programmatic way to read remaining balance or current CU/hour, so cost control falls back to a human glancing at the web UI.

A colab usage (or colab quota) command exposing:

• remaining compute-unit balance
• current CU/hour per active runtime (+ a total)
• ideally the per-accelerator rate, so a job can estimate cost before colab new

would let agents gate long jobs on budget and pick the cheapest viable accelerator.

From the installed source, the CLI already authenticates to and calls colab.pa.googleapis.com (the same backend the web UI reads the balance from), but client.py only models session/assignment RPCs — there's no quota/units call. So this looks like surfacing an existing backend RPC rather than new infrastructure.

Willing to contribute a PR: yes. Happy to implement colab usage if a maintainer can point at the RuntimeService method (or confirm the endpoint) that returns CU balance/rate, since colab.pa.googleapis.com isn't publicly documented. Also glad to help with the machineShape surfacing in section 1 — Shape.STANDARD/Shape.HIGH_RAM are already defined in client.py.

Thanks for the tool — it's been genuinely useful for agent-driven GPU workflows.

[jthirty7]

Just adding that still can't connect to a high ram gpu through the CLI and this is necessary for my work. Basically makes CLI useless to me.