@samitouri / QOSami-colab-cli-2 / commits / b04e83a

feat: add --high-mem for high-RAM machine shape selection (#105)

Koji committed Aug 11, 2026 at 19:22 UTC b04e83a92aae3549f7bc734546d36d56a3146773

13 files changed +371 -73

README.md

+5 -4

@@ -68,9 +68,9 @@ Run `colab <command> --help` to view specific options, defaults, and detailed he

68 ### Session Management

69 | Command | Description |

70 | --- | --- |

71 -| `colab new [-s NAME] [--gpu GPU] [--tpu TPU]` | Allocate a new CPU, GPU, or TPU VM runtime |

71 +| `colab new [-s NAME] [--gpu GPU] [--tpu TPU] [--high-mem]` | Allocate a new CPU, GPU, or TPU VM runtime (optionally high-RAM) |

72 | `colab sessions` | List all active sessions currently active on the backend |

73 -| `colab status [-s NAME]` | Display hardware, status, and local metadata for active sessions |

73 +| `colab status [-s NAME]` | Display hardware, machine shape, status, and local metadata for active sessions |

74 | `colab restart-kernel [-s NAME]` | Restart the active session's Jupyter kernel |

75 | `colab stop [-s NAME]` | Terminate a session VM and tear down its keep-alive daemon |

76 | `colab url [-s NAME] [--open]` | Print or open a browser URL connecting to the active session |

@@ -78,11 +78,11 @@ Run `colab <command> --help` to view specific options, defaults, and detailed he

78 ### Execution

79 | Command | Description |

80 | --- | --- |

81 -| `colab run [--gpu GPU] [--tpu TPU] [--keep] SCRIPT [ARGS...]` | Run a local script on a fresh VM, forwarding arguments, then release it |

81 +| `colab run [--gpu GPU] [--tpu TPU] [--high-mem] [--keep] SCRIPT [ARGS...]` | Run a local script on a fresh VM, forwarding arguments, then release it |

82 | `colab exec [-s NAME] [-f FILE] [--output-image PATH]` | Execute Python code from stdin, a local `.py` file, or a `.ipynb` notebook |

83 | `colab repl [-s NAME] [--output-image PATH]` | Start an interactive Python REPL on the VM (exits cleanly on piped EOF) |

84 | `colab console [-s NAME]` | Connect to a raw interactive TTY shell (tmux) on the remote VM |

85 -| `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 |

85 +| `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 |

87 ### File Operations

88 | Command | Description |

@@ -142,6 +142,7 @@ colab stop -s analysis

142

143 ## Usage Notes

144

145 +* **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`.

146 * **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.

147 * **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.

148 * **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

+13 -3

@@ -1,5 +1,6 @@

1 ---

2 log:

3 +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.

4 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.

5 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.

6 ---

@@ -32,11 +33,20 @@ Defines the specific hardware model.

33 - `V5E1`: TPU v5e (1 core, optimized for inference/efficient training).

34 - `V6E1`: TPU v6e (1 core, high performance).

35 -### 3. CLI Mapping

36 +### 3. Machine shape (`shape`)

37 +Defines the RAM profile for runtimes that support a choice (CPU, T4, A100, etc.).

38 +- `STANDARD` (default): omit the `shape` query param on assign.

39 +- `HIGH_RAM`: send `shape=hm` on assign (requires Colab Pro/Pro+ entitlement).

40 +

41 +Accelerators with only one shape (L4, v5e1, v6e1) ignore `--high-mem`.

42 +

43 +### 4. CLI Mapping

44 The CLI maps user flags to these backend parameters:

45 - `colab new my-session` -> `variant=DEFAULT`, `accelerator=NONE`

38 -- `colab new my-session -gpu=L4` -> `variant=GPU`, `accelerator=L4`

39 -- `colab new my-session -tpu=v5e1` -> `variant=TPU`, `accelerator=V5E1`

46 +- `colab new my-session --gpu=L4` -> `variant=GPU`, `accelerator=L4`

47 +- `colab new my-session --tpu=v5e1` -> `variant=TPU`, `accelerator=V5E1`

48 +- `colab new my-session --high-mem` -> adds `shape=hm` (when supported)

49 +- `colab new my-session --gpu A100 --high-mem` -> `variant=GPU`, `accelerator=A100`, `shape=hm`

51 ## Approach

docs/05_run_command.md

@@ -1,5 +1,6 @@

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`.

@@ -29,6 +30,7 @@ colab run [OPTIONS] SCRIPT [SCRIPT_ARGS]...

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). |

docs/06_ssh_access.md

@@ -1,5 +1,6 @@

1 ---

2 log:

3 +2026-08-09: Added `--high-mem` passthrough when `colab ssh` auto-creates a runtime (forwards to `colab new --high-mem`).

4 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).

5 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`.

6 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]

32 | `--gpu` | str | None | GPU accelerator for a runtime this command creates (T4, L4, G4, H100, A100). |

33 | `--tpu` | str | None | TPU accelerator for a runtime this command creates (v5e1, v6e1). |

37 ### `~/.ssh/config` usage

src/colab_cli/client.py

+43 -4

@@ -91,6 +91,37 @@ class Shape(int, Enum):

91 HIGH_RAM = 1

94 +# Accelerators that only exist in a single (high-memory) shape; the assign

95 +# endpoint ignores shape=hm for these (colab-vscode forces STANDARD).

96 +HIGH_MEM_ONLY_ACCELERATORS = frozenset(

97 + {Accelerator.L4, Accelerator.V5E1, Accelerator.V6E1}

98 +)

99 +

100 +

101 +def resolve_assign_shape(

102 + accelerator: Optional[Accelerator],

103 + *,

104 + high_mem: bool = False,

105 +) -> Optional[Shape]:

106 + """Map CLI intent to the shape query param for /tun/m/assign.

107 +

108 + Returns ``Shape.HIGH_RAM`` when high memory was requested and the

109 + accelerator supports a choice; otherwise ``None`` (omit the URL param).

110 + """

111 + if not high_mem:

112 + return None

113 + if accelerator in HIGH_MEM_ONLY_ACCELERATORS:

114 + return None

115 + return Shape.HIGH_RAM

116 +

117 +

118 +def shape_display_label(shape: Union[Shape, str, int, None]) -> str:

119 + """Human-friendly label for sessions/status output."""

120 + if shape in (Shape.HIGH_RAM, "HIGH_RAM", 1):

121 + return "High-RAM"

122 + return "Standard"

123 +

124 +

125 class RuntimeProxyInfo(BaseModel):

126 token: str

127 token_expires_in_seconds: int = Field(..., alias="tokenExpiresInSeconds")

@@ -228,14 +259,17 @@ class Client:

259 notebook_hash: uuid.UUID,

260 variant: Optional[Variant] = None,

261 accelerator: Optional[Accelerator] = None,

262 + shape: Optional[Shape] = None,

263 ) -> Union[PostAssignmentResponse, Assignment]:

232 - assignment = self._get_assignment(notebook_hash, variant, accelerator)

264 + assignment = self._get_assignment(

265 + notebook_hash, variant, accelerator, shape

266 + )

267 if isinstance(assignment, Assignment):

268 return assignment

269

270 try:

271 res = self._post_assignment(

238 - notebook_hash, assignment.token, variant, accelerator

272 + notebook_hash, assignment.token, variant, accelerator, shape

273 )

274 except ColabRequestError as e:

275 if get_status_code(e) == 412:

@@ -249,6 +283,7 @@ class Client:

283 notebook_hash: uuid.UUID,

284 variant: Optional[Variant] = None,

285 accelerator: Optional[Accelerator] = None,

286 + shape: Optional[Shape] = None,

287 ) -> str:

288 url = urljoin(self.colab_domain, f"{TUN_ENDPOINT}/assign")

289 params = {"nbh": uuid_to_web_safe_base64(notebook_hash)}

@@ -256,6 +291,8 @@ class Client:

291 params["variant"] = variant.value

292 if accelerator:

293 params["accelerator"] = accelerator.value

294 + if shape == Shape.HIGH_RAM:

295 + params["shape"] = "hm"

296

297 req = requests.Request("GET", url, params=params)

298 prep = req.prepare()

@@ -266,8 +303,9 @@ class Client:

303 notebook_hash: uuid.UUID,

304 variant: Optional[Variant] = None,

305 accelerator: Optional[Accelerator] = None,

306 + shape: Optional[Shape] = None,

307 ) -> Union[GetAssignmentResponse, Assignment]:

270 - url = self._build_assign_url(notebook_hash, variant, accelerator)

308 + url = self._build_assign_url(notebook_hash, variant, accelerator, shape)

309 return self._issue_request(url, schema=Union[GetAssignmentResponse, Assignment])

310

311 def _post_assignment(

@@ -276,8 +314,9 @@ class Client:

314 xsrf_token: str,

315 variant: Optional[Variant] = None,

316 accelerator: Optional[Accelerator] = None,

317 + shape: Optional[Shape] = None,

318 ) -> PostAssignmentResponse:

280 - url = self._build_assign_url(notebook_hash, variant, accelerator)

319 + url = self._build_assign_url(notebook_hash, variant, accelerator, shape)

320 headers = {COLAB_XSRF_TOKEN_HEADER["key"]: xsrf_token}

321 return self._issue_request(

322 url, method="POST", headers=headers, schema=PostAssignmentResponse

src/colab_cli/commands/run.py

+28 -24

@@ -41,13 +41,15 @@ from typing_extensions import Annotated

41 from colab_cli.client import (

42 Accelerator,

43 ColabRequestError,

44 + HIGH_MEM_ONLY_ACCELERATORS,

45 PostAssignmentResponse,

45 - Variant,

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

@@ -55,27 +57,6 @@ from colab_cli.state import SessionState

57 from colab_cli.utils import get_status_code, is_terminal_error

58 -# TODO(sethtroisi): dedupe this logic with similar in session.py

59 -def _resolve_accelerator(gpu: Optional[str], tpu: Optional[str]):

60 - """Mirror the mapping logic in `commands.session.new`. Centralised so the

61 - two commands stay in lock-step on supported accelerator names.

62 - """

63 - if tpu:

64 - variant = Variant.TPU

65 - accelerator = Accelerator.V5E1 if tpu.lower() == "v5e1" else Accelerator.V6E1

66 - return variant, accelerator

67 - if gpu:

68 - mapping = {

69 - "a100": Accelerator.A100,

70 - "h100": Accelerator.H100,

71 - "l4": Accelerator.L4,

72 - "t4": Accelerator.T4,

73 - "g4": Accelerator.G4,

74 - }

75 - return Variant.GPU, mapping.get(gpu.lower(), Accelerator.A100)

76 - return Variant.DEFAULT, Accelerator.NONE

77 -

78 -

60 def _build_script_payload(

61 script_path: str, script_args: List[str], env_vars: Optional[dict[str, str]] = None

62 ) -> str:

@@ -259,6 +240,16 @@ def run_command(

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(

@@ -306,12 +297,21 @@ def run_command(

297 raise typer.Exit(2)

298

299 name = session or f"run-{uuid.uuid4().hex[:6]}"

309 - variant, accelerator = _resolve_accelerator(gpu, tpu)

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

314 + uuid.uuid4(), variant=variant, accelerator=accelerator, shape=shape

315 )

316 except ColabRequestError as e:

317 # Mirror `colab new`'s friendly accelerator-quota message.

@@ -346,6 +346,9 @@ def run_command(

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

@@ -384,6 +387,7 @@ def run_command(

387 "endpoint": endpoint,

388 "variant": variant.value,

389 "accelerator": accelerator.value,

390 + "machine_shape": s.machine_shape,

391 "via": "run",

392 },

393 )

src/colab_cli/commands/session.py

+60 -17

@@ -24,8 +24,12 @@ from typing_extensions import Annotated

24 from colab_cli.client import (

25 Accelerator,

26 ColabRequestError,

27 + HIGH_MEM_ONLY_ACCELERATORS,

28 PostAssignmentResponse,

29 + Shape,

30 Variant,

31 + resolve_assign_shape,

32 + shape_display_label,

33 )

34 from colab_cli.utils import get_status_code

35 from colab_cli.state import SessionState

@@ -95,16 +99,18 @@ def _format_session_line(

99 accelerator: str,

100 variant: str,

101 status: Optional[str] = None,

102 + machine_shape: Optional[str] = None,

103 ) -> str:

104 """Single source of truth for session display lines.

105

101 - Format: ``[name] endpoint | Hardware: X | Variant: Y[ | Status: Z]``.

106 + Format: ``[name] endpoint | Hardware: X | Shape: Y | Variant: Z[ | Status: W]``.

107 Use ``"?"`` as the name for orphaned server-side assignments with no local

108 state.

109 """

110 parts = [

111 f"[{name}] {endpoint}",

112 f"Hardware: {_hardware_label(accelerator)}",

113 + f"Shape: {shape_display_label(machine_shape)}",

114 f"Variant: {variant}",

115 ]

116 if status is not None:

@@ -112,6 +118,34 @@ def _format_session_line(

118 return " | ".join(parts)

119

120

121 +def resolve_runtime_options(

122 + gpu: Optional[str] = None,

123 + tpu: Optional[str] = None,

124 + *,

125 + high_mem: bool = False,

126 +) -> tuple[Variant, Accelerator, Optional[Shape]]:

127 + """Map CLI flags to backend variant, accelerator, and optional shape."""

128 + if tpu:

129 + variant = Variant.TPU

130 + accelerator = Accelerator.V5E1 if tpu.lower() == "v5e1" else Accelerator.V6E1

131 + elif gpu:

132 + variant = Variant.GPU

133 + mapping = {

134 + "a100": Accelerator.A100,

135 + "h100": Accelerator.H100,

136 + "l4": Accelerator.L4,

137 + "t4": Accelerator.T4,

138 + "g4": Accelerator.G4,

139 + }

140 + accelerator = mapping.get(gpu.lower(), Accelerator.A100)

141 + else:

142 + variant = Variant.DEFAULT

143 + accelerator = Accelerator.NONE

144 +

145 + shape = resolve_assign_shape(accelerator, high_mem=high_mem)

146 + return variant, accelerator, shape

147 +

148 +

149 def new(

150 session: Annotated[

151 Optional[str], typer.Option("-s", "--session", help="Session name")

@@ -132,32 +166,35 @@ def new(

166 ),

167 ),

168 ] = None,

169 + high_mem: Annotated[

170 + bool,

171 + typer.Option(

172 + "--high-mem",

173 + help=(

174 + "Request a high-RAM machine shape (CPU, T4, A100, etc.). "

175 + "Requires Colab Pro or Pro+ entitlement. Ignored for "

176 + "accelerators that only offer a single shape (L4, v5e1, v6e1)."

177 + ),

178 + ),

179 + ] = False,

180 ):

181 """Create a new session"""

182 from colab_cli.common import state

183

184 name = session or uuid.uuid4().hex[:6]

140 - variant = Variant.DEFAULT

141 - accelerator = Accelerator.NONE

185 + variant, accelerator, shape = resolve_runtime_options(gpu, tpu, high_mem=high_mem)

186

143 - if tpu:

144 - variant = Variant.TPU

145 - accelerator = Accelerator.V5E1 if tpu.lower() == "v5e1" else Accelerator.V6E1

146 - elif gpu:

147 - variant = Variant.GPU

148 - mapping = {

149 - "a100": Accelerator.A100,

150 - "h100": Accelerator.H100,

151 - "l4": Accelerator.L4,

152 - "t4": Accelerator.T4,

153 - "g4": Accelerator.G4,

154 - }

155 - accelerator = mapping.get(gpu.lower(), Accelerator.A100)

187 + if high_mem and accelerator in HIGH_MEM_ONLY_ACCELERATORS:

188 + typer.echo(

189 + "[colab] --high-mem ignored: this accelerator only offers one "

190 + "machine shape.",

191 + err=True,

192 + )

193

194 typer.echo(f"[colab] Creating session '{name}'...")

195 try:

196 res = state.client.assign(

160 - uuid.uuid4(), variant=variant, accelerator=accelerator

197 + uuid.uuid4(), variant=variant, accelerator=accelerator, shape=shape

198 )

199 except ColabRequestError as e:

200 # The Colab backend returns 400 when the caller is not entitled to the

@@ -198,6 +235,9 @@ def new(

235 endpoint=endpoint,

236 variant=variant.value,

237 accelerator=accelerator.value,

238 + machine_shape=(

239 + Shape.HIGH_RAM.name if shape == Shape.HIGH_RAM else Shape.STANDARD.name

240 + ),

241 )

242

243 # Pre-flight the keep-alive ping once. If it returns a 403 caused by

@@ -244,6 +284,7 @@ def new(

284 "endpoint": endpoint,

285 "variant": variant.value,

286 "accelerator": accelerator.value,

287 + "machine_shape": s.machine_shape,

288 },

289 )

290 typer.echo("[colab] Session READY.")

@@ -305,6 +346,7 @@ def sessions_command():

346 endpoint=a.endpoint,

347 accelerator=a.accelerator.value,

348 variant=a.variant.name,

349 + machine_shape=a.machine_shape.name,

350 )

351 )

352

@@ -319,6 +361,7 @@ def _print_status_for(s: SessionState) -> None:

361 accelerator=s.accelerator,

362 variant=s.variant,

363 status=status,

364 + machine_shape=s.machine_shape,

365 )

366 )

367 if s.last_execution:

src/colab_cli/commands/ssh.py

+43 -14

@@ -177,7 +177,11 @@ def _has_local_sessions() -> bool:

177

178

179 def _auto_create_session(

180 - gpu: Optional[str], tpu: Optional[str], name: Optional[str] = None

180 + gpu: Optional[str],

181 + tpu: Optional[str],

182 + name: Optional[str] = None,

183 + *,

184 + high_mem: bool = False,

185 ) -> SessionState:

186 """Creates a runtime via ``colab new`` and returns its session.

187

@@ -196,7 +200,7 @@ def _auto_create_session(

200

201 name = name or uuid.uuid4().hex[:6]

202 typer.echo(f"[colab] Creating runtime '{name}'...")

199 - session_cmd.new(session=name, gpu=gpu, tpu=tpu)

203 + session_cmd.new(session=name, gpu=gpu, tpu=tpu, high_mem=high_mem)

204 return _resolve_session(name)

205

206

@@ -456,7 +460,11 @@ def _run_interactive_ssh(session: SessionState, identity: Optional[str]) -> int:

460

461

462 def _select_proxy_session(

459 - session: Optional[str], gpu: Optional[str], tpu: Optional[str]

463 + session: Optional[str],

464 + gpu: Optional[str],

465 + tpu: Optional[str],

466 + *,

467 + high_mem: bool = False,

468 ) -> tuple[SessionState, bool]:

469 """Resolves (or creates) the session for ``--proxy-mode``.

470

@@ -475,12 +483,18 @@ def _select_proxy_session(

483 """

484 if session and not _session_exists(session):

485 with contextlib.redirect_stdout(sys.stderr):

478 - return _auto_create_session(gpu, tpu, name=session), True

486 + return _auto_create_session(

487 + gpu, tpu, name=session, high_mem=high_mem

488 + ), True

489 return _resolve_session(session), False

490

491

492 def _select_interactive_session(

483 - session: Optional[str], gpu: Optional[str], tpu: Optional[str]

493 + session: Optional[str],

494 + gpu: Optional[str],

495 + tpu: Optional[str],

496 + *,

497 + high_mem: bool = False,

498 ) -> tuple[SessionState, bool]:

499 """Resolves (or auto-creates) the session for an interactive shell.

500

@@ -496,17 +510,22 @@ def _select_interactive_session(

510 A ``(session_state, created)`` pair.

511 """

512 if not session and not _has_local_sessions():

499 - return _auto_create_session(gpu, tpu), True

513 + return _auto_create_session(gpu, tpu, high_mem=high_mem), True

514 return _resolve_session(session), False

515

516

517 def _warn_accelerator_ignored(

504 - gpu: Optional[str], tpu: Optional[str], created: bool

518 + gpu: Optional[str],

519 + tpu: Optional[str],

520 + created: bool,

521 + *,

522 + high_mem: bool = False,

523 ) -> None:

506 - """Warns that ``--gpu/--tpu`` are no-ops when no runtime was created."""

507 - if (gpu or tpu) and not created:

524 + """Warns that ``--gpu/--tpu/--high-mem`` are no-ops when no runtime was created."""

525 + if (gpu or tpu or high_mem) and not created:

526 typer.echo(

509 - "[colab] --gpu/--tpu ignored: only applies to a created runtime.",

527 + "[colab] --gpu/--tpu/--high-mem ignored: only applies to a "

528 + "created runtime.",

529 err=True,

530 )

531

@@ -647,6 +666,16 @@ def ssh(

666 ),

667 ),

668 ] = None,

669 + high_mem: Annotated[

670 + bool,

671 + typer.Option(

672 + "--high-mem",

673 + help=(

674 + "Request a high-RAM machine shape when this command "

675 + "auto-creates a runtime."

676 + ),

677 + ),

678 + ] = False,

679 rm: Annotated[

680 bool,

681 typer.Option(

@@ -669,12 +698,12 @@ def ssh(

698 ``--gpu/--tpu`` set its accelerator, ``--rm`` stops it on disconnect).

699 """

700 if proxy_mode:

672 - s, created = _select_proxy_session(session, gpu, tpu)

673 - _warn_accelerator_ignored(gpu, tpu, created)

701 + s, created = _select_proxy_session(session, gpu, tpu, high_mem=high_mem)

702 + _warn_accelerator_ignored(gpu, tpu, created, high_mem=high_mem)

703 raise typer.Exit(code=_run_proxy_bridge(s, identity, rm))

704

676 - s, created = _select_interactive_session(session, gpu, tpu)

677 - _warn_accelerator_ignored(gpu, tpu, created)

705 + s, created = _select_interactive_session(session, gpu, tpu, high_mem=high_mem)

706 + _warn_accelerator_ignored(gpu, tpu, created, high_mem=high_mem)

707 raise typer.Exit(code=_run_interactive_shell(s, identity, created, rm))

708

709

src/colab_cli/state.py

@@ -29,6 +29,7 @@ class SessionState(BaseModel):

29 endpoint: str

30 variant: str = "DEFAULT"

31 accelerator: str = "NONE"

32 + machine_shape: str = "STANDARD"

33 kernel_id: Optional[str] = None

34 session_id: Optional[str] = None

35 last_execution: Optional[Tuple[str, Optional[str], str]] = None

tests/test_cli.py

+40 -6

@@ -104,12 +104,13 @@ def test_cli_new_gpu_variants(mock_client, mock_store, gpu_flag, expected_acc):

104

105 def test_cli_sessions_unified_format(mock_client, mock_common_state):

106 """`sessions` should lead each line with the local name when known:

107 - `[name] endpoint | Hardware: X | Variant: Y`.

107 + `[name] endpoint | Hardware: X | Shape: Y | Variant: Z`.

108 """

109 mock_assignment = MagicMock()

110 mock_assignment.endpoint = "e1"

111 mock_assignment.variant.name = "GPU"

112 mock_assignment.accelerator.value = "T4"

113 + mock_assignment.machine_shape.name = "STANDARD"

114

115 mock_session_state = MagicMock()

116 mock_session_state.name = "s1"

@@ -123,7 +124,7 @@ def test_cli_sessions_unified_format(mock_client, mock_common_state):

124

125 result = runner.invoke(app, ["sessions"])

126 assert result.exit_code == 0

126 - assert "[s1] e1 | Hardware: T4 | Variant: GPU" in result.output

127 + assert "[s1] e1 | Hardware: T4 | Shape: Standard | Variant: GPU" in result.output

128

129

130 def test_cli_sessions_orphaned_assignment_marked(mock_client, mock_common_state):

@@ -132,13 +133,17 @@ def test_cli_sessions_orphaned_assignment_marked(mock_client, mock_common_state)

133 mock_assignment.endpoint = "orphan-ep"

134 mock_assignment.variant.name = "DEFAULT"

135 mock_assignment.accelerator.value = "NONE"

136 + mock_assignment.machine_shape.name = "HIGH_RAM"

137

138 mock_common_state.sync_sessions.return_value = ({}, [mock_assignment])

139

140 result = runner.invoke(app, ["sessions"])

141 assert result.exit_code == 0

142 # CPU is the alias for accelerator NONE

141 - assert "[?] orphan-ep | Hardware: CPU | Variant: DEFAULT" in result.output

143 + assert (

144 + "[?] orphan-ep | Hardware: CPU | Shape: High-RAM | Variant: DEFAULT"

145 + in result.output

146 + )

147

148

149 def test_cli_sessions_no_assignments(mock_client, mock_common_state):

@@ -154,6 +159,7 @@ def test_cli_status(mock_store, mock_common_state):

159 mock_session_state.endpoint = "e1"

160 mock_session_state.accelerator = "NONE"

161 mock_session_state.variant = "DEFAULT"

162 + mock_session_state.machine_shape = "STANDARD"

163 mock_session_state.running = None

164 mock_session_state.last_execution = (

165 "my_notebook.ipynb",

@@ -167,7 +173,10 @@ def test_cli_status(mock_store, mock_common_state):

173 # Test with explicit session: uses unified format including endpoint and Status

174 result = runner.invoke(app, ["status", "-s", "s1"])

175 assert result.exit_code == 0

170 - assert "[s1] e1 | Hardware: CPU | Variant: DEFAULT | Status: IDLE" in result.output

176 + assert (

177 + "[s1] e1 | Hardware: CPU | Shape: Standard | Variant: DEFAULT | Status: IDLE"

178 + in result.output

179 + )

180 assert (

181 "Last Execution: my_notebook.ipynb | Cell: cell_1 at 2023-10-27 12:00:00"

182 in result.output

@@ -184,7 +193,10 @@ def test_cli_status(mock_store, mock_common_state):

193 mock_store.get.return_value = mock_session_state

194 result = runner.invoke(app, ["status"])

195 assert result.exit_code == 0

187 - assert "[s1] e1 | Hardware: CPU | Variant: DEFAULT | Status: IDLE" in result.output

196 + assert (

197 + "[s1] e1 | Hardware: CPU | Shape: Standard | Variant: DEFAULT | Status: IDLE"

198 + in result.output

199 + )

200

201 # Test without execution metadata

202 mock_session_state.last_execution = None

@@ -200,6 +212,7 @@ def test_cli_status_running_shows_busy(mock_store, mock_common_state):

212 mock_session_state.endpoint = "e1"

213 mock_session_state.accelerator = "T4"

214 mock_session_state.variant = "GPU"

215 + mock_session_state.machine_shape = "STANDARD"

216 mock_session_state.running = "exec.py"

217 mock_session_state.last_execution = None

218 mock_store.get.return_value = mock_session_state

@@ -208,11 +221,32 @@ def test_cli_status_running_shows_busy(mock_store, mock_common_state):

221 result = runner.invoke(app, ["status", "-s", "s1"])

222 assert result.exit_code == 0

223 assert (

211 - "[s1] e1 | Hardware: T4 | Variant: GPU | Status: BUSY (exec.py)"

224 + "[s1] e1 | Hardware: T4 | Shape: Standard | Variant: GPU | Status: BUSY (exec.py)"

225 in result.output

226 )

227

228

229 +def test_cli_new_high_mem(mock_client, mock_store):

230 + mock_res = MagicMock()

231 + mock_res.__class__ = PostAssignmentResponse

232 + mock_res.runtime_proxy_info.token = "t1"

233 + mock_res.runtime_proxy_info.url = "u1"

234 + mock_res.endpoint = "e1"

235 + mock_client.assign.return_value = mock_res

236 +

237 + result = runner.invoke(app, ["new", "-s", "hm-sess", "--gpu", "A100", "--high-mem"])

238 + assert result.exit_code == 0

239 +

240 + mock_client.assign.assert_called_once()

241 + _, kwargs = mock_client.assign.call_args

242 + from colab_cli.client import Shape

243 +

244 + assert kwargs["shape"] == Shape.HIGH_RAM

245 +

246 + added_state = mock_store.add.call_args[0][0]

247 + assert added_state.machine_shape == "HIGH_RAM"

248 +

249 +

250 def test_cli_session_resolution(mock_store, mock_common_state):

251 mock_session_state = MagicMock()

252 mock_session_state.name = "unique-session"

tests/test_client.py

+88 -1

@@ -16,7 +16,16 @@ import uuid

16 import json

17 import pytest

18 from unittest.mock import MagicMock

19 -from colab_cli.client import Client, Prod, PostAssignmentResponse, Assignment

19 +from colab_cli.client import (

20 + Client,

21 + Prod,

22 + PostAssignmentResponse,

23 + Assignment,

24 + Accelerator,

25 + Shape,

26 + Variant,

27 + resolve_assign_shape,

28 +)

31 @pytest.fixture

@@ -234,3 +243,81 @@ def test_client_keep_alive_assignment_propagates_http_error(client, mock_session

243

244 with pytest.raises(ColabRequestError):

245 client.keep_alive_assignment("m-s-test-endpoint")

246 +

247 +

248 +def test_client_assign_url_includes_shape_hm(client, mock_session):

249 + get_resp = MagicMock()

250 + get_resp.ok = True

251 + get_resp.text = ")]}'\n" + json.dumps(

252 + {"acc": "NONE", "nbh": "some_nbh", "token": "xsrf_token", "variant": "DEFAULT"}

253 + )

254 + post_resp = MagicMock()

255 + post_resp.ok = True

256 + post_resp.text = ")]}'\n" + json.dumps(

257 + {

258 + "accelerator": "A100",

259 + "endpoint": "new_endpoint",

260 + "runtimeProxyInfo": {

261 + "token": "proxy_token",

262 + "tokenExpiresInSeconds": 3600,

263 + "url": "http://backend",

264 + },

265 + "variant": 1,

266 + }

267 + )

268 + mock_session.request.side_effect = [get_resp, post_resp]

269 +

270 + client.assign(

271 + uuid.uuid4(),

272 + variant=Variant.GPU,

273 + accelerator=Accelerator.A100,

274 + shape=Shape.HIGH_RAM,

275 + )

276 +

277 + get_url = mock_session.request.call_args_list[0].args[1]

278 + assert "shape=hm" in get_url

279 + post_url = mock_session.request.call_args_list[1].args[1]

280 + assert "shape=hm" in post_url

281 +

282 +

283 +def test_client_assign_url_omits_shape_for_standard(client, mock_session):

284 + get_resp = MagicMock()

285 + get_resp.ok = True

286 + get_resp.text = ")]}'\n" + json.dumps(

287 + {"acc": "NONE", "nbh": "some_nbh", "token": "xsrf_token", "variant": "DEFAULT"}

288 + )

289 + post_resp = MagicMock()

290 + post_resp.ok = True

291 + post_resp.text = ")]}'\n" + json.dumps(

292 + {

293 + "accelerator": "NONE",

294 + "endpoint": "new_endpoint",

295 + "runtimeProxyInfo": {

296 + "token": "proxy_token",

297 + "tokenExpiresInSeconds": 3600,

298 + "url": "http://backend",

299 + },

300 + "variant": 0,

301 + }

302 + )

303 + mock_session.request.side_effect = [get_resp, post_resp]

304 +

305 + client.assign(uuid.uuid4())

306 +

307 + get_url = mock_session.request.call_args_list[0].args[1]

308 + assert "shape=" not in get_url

309 +

310 +

311 +@pytest.mark.parametrize(

312 + "accelerator,high_mem,expected",

313 + [

314 + (Accelerator.T4, True, Shape.HIGH_RAM),

315 + (Accelerator.A100, True, Shape.HIGH_RAM),

316 + (Accelerator.NONE, True, Shape.HIGH_RAM),

317 + (Accelerator.L4, True, None),

318 + (Accelerator.V5E1, True, None),

319 + (Accelerator.T4, False, None),

320 + ],

321 +)

322 +def test_resolve_assign_shape(accelerator, high_mem, expected):

323 + assert resolve_assign_shape(accelerator, high_mem=high_mem) == expected

tests/test_run.py

+27

@@ -115,6 +115,33 @@ def test_run_basic_flow(

115 mock_client.unassign.assert_called_once_with("ep-123")

116

117

118 +def test_run_high_mem_passes_shape_to_assign(

119 + mock_client,

120 + mock_store,

121 + mock_runtime_class,

122 + mock_spawn_keep_alive,

123 + assign_response,

124 + script_path,

125 +):

126 + mock_client.assign.return_value = assign_response

127 + mock_runtime_class.return_value.execute_code.return_value = []

128 +

129 + persisted = {}

130 + mock_store.add.side_effect = lambda s: persisted.setdefault("s", s)

131 + mock_store.get.side_effect = lambda name: persisted.get("s")

132 +

133 + result = runner.invoke(

134 + app, ["run", "--gpu", "A100", "--high-mem", str(script_path)]

135 + )

136 + assert result.exit_code == 0, result.output

137 +

138 + from colab_cli.client import Shape

139 +

140 + _, kwargs = mock_client.assign.call_args

141 + assert kwargs["shape"] == Shape.HIGH_RAM

142 + assert persisted["s"].machine_shape == "HIGH_RAM"

143 +

144 +

145 # ---------------------------------------------------------------------------

146 # --keep flag

147 # ---------------------------------------------------------------------------

tests/test_ssh_autocreate.py

+19

@@ -114,10 +114,29 @@ def test_bare_ssh_session_resolution(

114 new.assert_called_once()

115 assert new.call_args.kwargs.get("gpu") == gpu

116 assert new.call_args.kwargs.get("tpu") is None

117 + assert new.call_args.kwargs.get("high_mem") is False

118 else:

119 new.assert_not_called()

120

121

122 +def test_bare_ssh_forwards_high_mem_on_autocreate(mock_common_state, mocker):

123 + mock_common_state.store.list.return_value = {}

124 + sess = _make_session()

125 + mock_common_state.store.get.return_value = sess

126 + new = mocker.patch("colab_cli.commands.session.new")

127 + _patch_interactive(mocker)

128 +

129 + result = runner.invoke(app, ["ssh", "--gpu", "A100", "--high-mem"])

130 + assert result.exit_code == 0

131 + new.assert_called_once()

132 + assert new.call_args.kwargs == {

133 + "session": new.call_args.kwargs["session"],

134 + "gpu": "A100",

135 + "tpu": None,

136 + "high_mem": True,

137 + }

138 +

139 +

140 # --- --proxy-mode -s NAME: create-if-missing / reuse ------------------------

141

142