feat: add --high-mem for high-RAM machine shape selection

README.md

@@ -68,21 +68,21 @@ Run `colab <command> --help` to view specific options, defaults, and detailed he
 ### Session Management
 | Command | Description |
 | --- | --- |
-| `colab new [-s NAME] [--gpu GPU] [--tpu TPU]` | Allocate a new CPU, GPU, or TPU VM runtime |
+| `colab new [-s NAME] [--gpu GPU] [--tpu TPU] [--high-mem]` | Allocate a new CPU, GPU, or TPU VM runtime (optionally high-RAM) |
 | `colab sessions` | List all active sessions currently active on the backend |
-| `colab status [-s NAME]` | Display hardware, status, and local metadata for active sessions |
+| `colab status [-s NAME]` | Display hardware, machine shape, status, and local metadata for active sessions |
 | `colab restart-kernel [-s NAME]` | Restart the active session's Jupyter kernel |
 | `colab stop [-s NAME]` | Terminate a session VM and tear down its keep-alive daemon |
 | `colab url [-s NAME] [--open]` | Print or open a browser URL connecting to the active session |
 
 ### Execution
 | Command | Description |
 | --- | --- |
-| `colab run [--gpu GPU] [--tpu TPU] [--keep] SCRIPT [ARGS...]` | Run a local script on a fresh VM, forwarding arguments, then release it |
+| `colab run [--gpu GPU] [--tpu TPU] [--high-mem] [--keep] SCRIPT [ARGS...]` | Run a local script on a fresh VM, forwarding arguments, then release it |
 | `colab exec [-s NAME] [-f FILE] [--output-image PATH]` | Execute Python code from stdin, a local `.py` file, or a `.ipynb` notebook |
 | `colab repl [-s NAME] [--output-image PATH]` | Start an interactive Python REPL on the VM (exits cleanly on piped EOF) |
 | `colab console [-s NAME]` | Connect to a raw interactive TTY shell (tmux) on the remote VM |
-| `colab ssh [-s NAME] [--proxy-mode] [-i KEY]` | Open an SSH shell to the runtime over WebSocket, or act as an OpenSSH `ProxyCommand` bridge for IDE remote-dev |
+| `colab ssh [-s NAME] [--proxy-mode] [-i KEY] [--gpu GPU] [--tpu TPU] [--high-mem]` | Open an SSH shell to the runtime over WebSocket, or act as an OpenSSH `ProxyCommand` bridge for IDE remote-dev |
 
 ### File Operations
 | Command | Description |
@@ -142,6 +142,7 @@ colab stop -s analysis
 
 ## Usage Notes
 
+* **Machine shape:** Use `--high-mem` with `colab new`, `colab run`, or `colab ssh` (when auto-creating a runtime) to request a high-RAM machine shape. Requires Colab Pro or Pro+ entitlement for supported accelerators (CPU, T4, A100, etc.). L4 and TPU runtimes ignore this flag because they only offer one shape. Machine shape is shown in `colab sessions` and `colab status`.
 * **TTY Requirements:** The interactive commands `repl` and `console` require a local TTY. When running inside automated scripts or pipelines, make sure to pipe stdin (e.g., `echo "print(1)" | colab repl`) to trigger non-interactive execution modes.
 * **Transparent Code Execution:** When calling `colab exec -f file.py`, the CLI reads the file locally and transmits its content to the remote kernel. You do not need to manually upload files before execution.
 * **Storage & State Paths:** Session tokens and metadata are stored at `~/.config/colab-cli/sessions.json`. Global CLI settings are located at `~/.config/colab-cli/settings.json`. These can be customized or isolated via the global `--config` flag.

docs/01_session_management.md

@@ -1,5 +1,6 @@
 ---
 log:
+2026-08-09: Added `--high-mem` to `colab new`, `colab run`, and `colab ssh` (auto-create). Assign requests now send `shape=hm` when high-RAM is requested; `colab sessions` and `colab status` display machine shape.
 2026-06-15: Switched the keep-alive daemon from the `colab.pa.googleapis.com` `RuntimeService/KeepAliveAssignment` RPC to a Tunnel Frontend HTTP ping (`GET /tun/m/<endpoint>/keep-alive/` with `X-Colab-Tunnel: Google`) on `colab.research.google.com`. The RPC required `serviceusage` consumer access to Colab's internal project `1014160490159`, which ordinary user accounts lack, so every external user hit HTTP 403 `USER_PROJECT_DENIED` and their CLI sessions were idle-pruned within minutes (issue #14). Reproduced live with a third-party account; verified the tunnel ping is accepted by the same bearer-token credential that already works for `assign`. A `ReadTimeout` on the ping is treated as success (TFE records activity before forwarding to the often-non-responding VM). Generalized the pre-flight remediation messaging away from the now-irrelevant `colaboratory`/`pa.googleapis.com` framing, and removed the dead grpc-web client-registry/API-key code.
 2026-06-10: Replaced the POSIX-only `fcntl.flock` file locking in `_LockedFileStore` with the cross-platform `filelock` library (reported broken on Windows). Reads use `ReadWriteLock.read_lock()` (shared) and writes use `write_lock()` (exclusive), preserving the original `LOCK_SH`/`LOCK_EX` semantics. The lock is constructed with `is_singleton=False` so two `StateStore` instances for the same path in one process don't collapse into a single reentrant lock (which would raise `RuntimeError` on multi-threaded write contention). Added shared-read, cross-process exclusion, and multi-thread/multi-process regression tests.
 ---
@@ -32,11 +33,20 @@ Defines the specific hardware model.
     - `V5E1`: TPU v5e (1 core, optimized for inference/efficient training).
     - `V6E1`: TPU v6e (1 core, high performance).
 
-### 3. CLI Mapping
+### 3. Machine shape (`shape`)
+Defines the RAM profile for runtimes that support a choice (CPU, T4, A100, etc.).
+- `STANDARD` (default): omit the `shape` query param on assign.
+- `HIGH_RAM`: send `shape=hm` on assign (requires Colab Pro/Pro+ entitlement).
+
+Accelerators with only one shape (L4, v5e1, v6e1) ignore `--high-mem`.
+
+### 4. CLI Mapping
 The CLI maps user flags to these backend parameters:
 - `colab new my-session` -> `variant=DEFAULT`, `accelerator=NONE`
-- `colab new my-session -gpu=L4` -> `variant=GPU`, `accelerator=L4`
-- `colab new my-session -tpu=v5e1` -> `variant=TPU`, `accelerator=V5E1`
+- `colab new my-session --gpu=L4` -> `variant=GPU`, `accelerator=L4`
+- `colab new my-session --tpu=v5e1` -> `variant=TPU`, `accelerator=V5E1`
+- `colab new my-session --high-mem` -> adds `shape=hm` (when supported)
+- `colab new my-session --gpu A100 --high-mem` -> `variant=GPU`, `accelerator=A100`, `shape=hm`
 
 ## Approach
 

docs/05_run_command.md

@@ -1,5 +1,6 @@
 ---
 log:
+2026-08-09: Added `--high-mem` flag (passthrough to session creation; sends `shape=hm` on assign when supported).
 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__"`.
 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.
 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`.
@@ -29,6 +30,7 @@ colab run [OPTIONS] SCRIPT [SCRIPT_ARGS]...
 | `-s`, `--session` | str | auto | Name the ephemeral session (helpful with `--keep`). Auto-generated as `run-<6 hex>` if omitted. |
 | `--gpu` | str | None | Same set as `colab new --gpu` (T4, L4, G4, H100, A100). |
 | `--tpu` | str | None | Same set as `colab new --tpu` (v5e1, v6e1). |
+| `--high-mem` | bool | False | Same as `colab new --high-mem` — request high-RAM when supported. |
 | `--keep` | bool | False | Do **not** stop the session after the script finishes. |
 | `--timeout` | float | 30.0 | Timeout in seconds for code execution to prevent hanging on silent tasks. |
 

docs/06_ssh_access.md

@@ -1,5 +1,6 @@
 ---
 log:
+2026-08-09: Added `--high-mem` passthrough when `colab ssh` auto-creates a runtime (forwards to `colab new --high-mem`).
 2026-07-17: Initial design and implementation of `colab ssh` — client side of SSH-over-WebSocket runtime access. Adds three modes (interactive shell, `-s SESSION`, and `--proxy-mode` OpenSSH ProxyCommand bridge), `--identity/-i` key selection, and per-HTTP-status handshake error messages. Server side is out of scope for this repo; the subcommand is a no-op against runtimes that do not expose the `/colab/ssh` endpoint (surfaces an actionable HTTP 404 message).
 2026-07-22: Bare `colab ssh` now auto-creates a runtime (via `colab new`) when you have no active session, with `--gpu/--tpu` passthrough and `--rm` to stop an auto-created runtime on exit. Fixed two client bugs: the dead 403 branch (feature-off returns 404, not 403) and the RSA guidance (all `ssh-rsa` keys are server-rejected, so `id_rsa` is no longer auto-scanned and the 400 message no longer advertises `rsa-sha2`). Added `tests/test_ssh_wire_contract.py` (real loopback-server wire assertions) and `tests/test_ssh_autocreate.py`.
 2026-07-22: Interactive `colab ssh` now starts in `/content` (Colab's working dir) instead of `/root`, via a forced PTY (`-t`) plus a remote `cd /content 2>/dev/null; exec $SHELL -l`. A missing `/content` falls back to the login home. Added `tests/test_ssh_workdir.py`.
@@ -30,6 +31,7 @@ colab ssh [OPTIONS]
 | `-i`, `--identity` | str | auto | Private key for the public key sent to Colab. Default: first of `~/.ssh/id_ed25519`, `id_ecdsa`. |
 | `--gpu` | str | None | GPU accelerator for a runtime this command creates (T4, L4, G4, H100, A100). |
 | `--tpu` | str | None | TPU accelerator for a runtime this command creates (v5e1, v6e1). |
+| `--high-mem` | bool | False | Request high-RAM when this command auto-creates a runtime (ignored when connecting to an existing session). |
 | `--rm` | bool | False | Stop the runtime when the session ends. Interactive: only a runtime `colab ssh` auto-created (a reused session is never removed). `--proxy-mode`: the bridged session, on disconnect. |
 
 ### `~/.ssh/config` usage

src/colab_cli/client.py

@@ -91,6 +91,37 @@ class Shape(int, Enum):
     HIGH_RAM = 1
 
 
+# Accelerators that only exist in a single (high-memory) shape; the assign
+# endpoint ignores shape=hm for these (colab-vscode forces STANDARD).
+HIGH_MEM_ONLY_ACCELERATORS = frozenset(
+    {Accelerator.L4, Accelerator.V5E1, Accelerator.V6E1}
+)
+
+
+def resolve_assign_shape(
+    accelerator: Optional[Accelerator],
+    *,
+    high_mem: bool = False,
+) -> Optional[Shape]:
+    """Map CLI intent to the shape query param for /tun/m/assign.
+
+    Returns ``Shape.HIGH_RAM`` when high memory was requested and the
+    accelerator supports a choice; otherwise ``None`` (omit the URL param).
+    """
+    if not high_mem:
+        return None
+    if accelerator in HIGH_MEM_ONLY_ACCELERATORS:
+        return None
+    return Shape.HIGH_RAM
+
+
+def shape_display_label(shape: Union[Shape, str, int, None]) -> str:
+    """Human-friendly label for sessions/status output."""
+    if shape in (Shape.HIGH_RAM, "HIGH_RAM", 1):
+        return "High-RAM"
+    return "Standard"
+
+
 class RuntimeProxyInfo(BaseModel):
     token: str
     token_expires_in_seconds: int = Field(..., alias="tokenExpiresInSeconds")
@@ -228,14 +259,17 @@ def assign(
         notebook_hash: uuid.UUID,
         variant: Optional[Variant] = None,
         accelerator: Optional[Accelerator] = None,
+        shape: Optional[Shape] = None,
     ) -> Union[PostAssignmentResponse, Assignment]:
-        assignment = self._get_assignment(notebook_hash, variant, accelerator)
+        assignment = self._get_assignment(
+            notebook_hash, variant, accelerator, shape
+        )
         if isinstance(assignment, Assignment):
             return assignment
 
         try:
             res = self._post_assignment(
-                notebook_hash, assignment.token, variant, accelerator
+                notebook_hash, assignment.token, variant, accelerator, shape
             )
         except ColabRequestError as e:
             if get_status_code(e) == 412:
@@ -249,13 +283,16 @@ def _build_assign_url(
         notebook_hash: uuid.UUID,
         variant: Optional[Variant] = None,
         accelerator: Optional[Accelerator] = None,
+        shape: Optional[Shape] = None,
     ) -> str:
         url = urljoin(self.colab_domain, f"{TUN_ENDPOINT}/assign")
         params = {"nbh": uuid_to_web_safe_base64(notebook_hash)}
         if variant:
             params["variant"] = variant.value
         if accelerator:
             params["accelerator"] = accelerator.value
+        if shape == Shape.HIGH_RAM:
+            params["shape"] = "hm"
 
         req = requests.Request("GET", url, params=params)
         prep = req.prepare()
@@ -266,8 +303,9 @@ def _get_assignment(
         notebook_hash: uuid.UUID,
         variant: Optional[Variant] = None,
         accelerator: Optional[Accelerator] = None,
+        shape: Optional[Shape] = None,
     ) -> Union[GetAssignmentResponse, Assignment]:
-        url = self._build_assign_url(notebook_hash, variant, accelerator)
+        url = self._build_assign_url(notebook_hash, variant, accelerator, shape)
         return self._issue_request(url, schema=Union[GetAssignmentResponse, Assignment])
 
     def _post_assignment(
@@ -276,8 +314,9 @@ def _post_assignment(
         xsrf_token: str,
         variant: Optional[Variant] = None,
         accelerator: Optional[Accelerator] = None,
+        shape: Optional[Shape] = None,
     ) -> PostAssignmentResponse:
-        url = self._build_assign_url(notebook_hash, variant, accelerator)
+        url = self._build_assign_url(notebook_hash, variant, accelerator, shape)
         headers = {COLAB_XSRF_TOKEN_HEADER["key"]: xsrf_token}
         return self._issue_request(
             url, method="POST", headers=headers, schema=PostAssignmentResponse

src/colab_cli/commands/run.py

@@ -41,41 +41,22 @@
 from colab_cli.client import (
     Accelerator,
     ColabRequestError,
+    HIGH_MEM_ONLY_ACCELERATORS,
     PostAssignmentResponse,
-    Variant,
+    Shape,
 )
 from colab_cli.commands.execution import _build_env_prelude, _parse_env_vars
 from colab_cli.commands.session import (
     _is_scope_error,
     _scope_remediation_message,
+    resolve_runtime_options,
     spawn_keep_alive,
 )
 from colab_cli.runtime import ColabRuntime
 from colab_cli.state import SessionState
 from colab_cli.utils import get_status_code, is_terminal_error
 
 
-# TODO(sethtroisi): dedupe this logic with similar in session.py
-def _resolve_accelerator(gpu: Optional[str], tpu: Optional[str]):
-    """Mirror the mapping logic in `commands.session.new`. Centralised so the
-    two commands stay in lock-step on supported accelerator names.
-    """
-    if tpu:
-        variant = Variant.TPU
-        accelerator = Accelerator.V5E1 if tpu.lower() == "v5e1" else Accelerator.V6E1
-        return variant, accelerator
-    if gpu:
-        mapping = {
-            "a100": Accelerator.A100,
-            "h100": Accelerator.H100,
-            "l4": Accelerator.L4,
-            "t4": Accelerator.T4,
-            "g4": Accelerator.G4,
-        }
-        return Variant.GPU, mapping.get(gpu.lower(), Accelerator.A100)
-    return Variant.DEFAULT, Accelerator.NONE
-
-
 def _build_script_payload(
     script_path: str, script_args: List[str], env_vars: Optional[dict[str, str]] = None
 ) -> str:
@@ -259,6 +240,16 @@ def run_command(
             ),
         ),
     ] = None,
+    high_mem: Annotated[
+        bool,
+        typer.Option(
+            "--high-mem",
+            help=(
+                "Request a high-RAM machine shape. Requires Colab Pro or Pro+ "
+                "entitlement. Ignored for L4 and TPU accelerators."
+            ),
+        ),
+    ] = False,
     keep: Annotated[
         bool,
         typer.Option(
@@ -306,12 +297,21 @@ def run_command(
         raise typer.Exit(2)
 
     name = session or f"run-{uuid.uuid4().hex[:6]}"
-    variant, accelerator = _resolve_accelerator(gpu, tpu)
+    variant, accelerator, shape = resolve_runtime_options(
+        gpu, tpu, high_mem=high_mem
+    )
+
+    if high_mem and accelerator in HIGH_MEM_ONLY_ACCELERATORS:
+        typer.echo(
+            "[colab] --high-mem ignored: this accelerator only offers one "
+            "machine shape.",
+            err=True,
+        )
 
     typer.echo(f"[colab] Creating session '{name}'...", err=True)
     try:
         res = state.client.assign(
-            uuid.uuid4(), variant=variant, accelerator=accelerator
+            uuid.uuid4(), variant=variant, accelerator=accelerator, shape=shape
         )
     except ColabRequestError as e:
         # Mirror `colab new`'s friendly accelerator-quota message.
@@ -346,6 +346,9 @@ def run_command(
         endpoint=endpoint,
         variant=variant.value,
         accelerator=accelerator.value,
+        machine_shape=(
+            Shape.HIGH_RAM.name if shape == Shape.HIGH_RAM else Shape.STANDARD.name
+        ),
     )
 
     # Pre-flight keep-alive: same scope-detection dance as `colab new` so a
@@ -384,6 +387,7 @@ def run_command(
             "endpoint": endpoint,
             "variant": variant.value,
             "accelerator": accelerator.value,
+            "machine_shape": s.machine_shape,
             "via": "run",
         },
     )