| 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 | import abc |
| 16 | from dataclasses import dataclass |
| 17 | from enum import Enum |
| 18 | import json |
| 19 | import logging |
| 20 | from typing import Dict, List, Optional, Union |
| 21 | from urllib.parse import urljoin, urlparse |
| 22 | import uuid |
| 23 | |
| 24 | from colab_cli.utils import get_status_code |
| 25 | from pydantic import BaseModel, Field, TypeAdapter |
| 26 | import requests |
| 27 | |
| 28 | # Standard Colab Headers |
| 29 | ACCEPT_JSON_HEADER = {"key": "Accept", "value": "application/json"} |
| 30 | COLAB_CLIENT_AGENT_HEADER = { |
| 31 | "key": "X-Colab-Client-Agent", |
| 32 | "value": "colab-cli", |
| 33 | } |
| 34 | COLAB_XSRF_TOKEN_HEADER = {"key": "X-Goog-Colab-Token", "value": ""} |
| 35 | # Marks a request as one that should be resolved through the Colab tunnel |
| 36 | # (Tunnel Frontend). Required by TFE-intercepted paths such as the keep-alive |
| 37 | # ping; without it the front-door rejects the request with HTTP 400. |
| 38 | COLAB_TUNNEL_HEADER = {"key": "X-Colab-Tunnel", "value": "Google"} |
| 39 | |
| 40 | # Per-request timeout (seconds) for the keep-alive tunnel ping. TFE records the |
| 41 | # activity as soon as the request arrives, so we do not need to wait long for |
| 42 | # the (often non-responding) VM. A short timeout keeps the keep-alive daemon |
| 43 | # responsive on its 60s cadence. |
| 44 | KEEP_ALIVE_TIMEOUT = 10 |
| 45 | |
| 46 | |
| 47 | @dataclass |
| 48 | class ColabEnvironment(abc.ABC): |
| 49 | domain: str |
| 50 | api: str |
| 51 | |
| 52 | |
| 53 | @dataclass |
| 54 | class Prod(ColabEnvironment): |
| 55 | domain: str = "https://colab.research.google.com" |
| 56 | api: str = "https://colab.pa.googleapis.com" |
| 57 | |
| 58 | |
| 59 | def uuid_to_web_safe_base64(uuid_val: uuid.UUID) -> str: |
| 60 | uuid_str = str(uuid_val) |
| 61 | transformed = uuid_str.replace("-", "_") |
| 62 | padding = "." * (44 - len(uuid_str)) |
| 63 | return transformed + padding |
| 64 | |
| 65 | |
| 66 | class Accelerator(str, Enum): |
| 67 | NONE = "NONE" |
| 68 | G4 = "G4" |
| 69 | T4 = "T4" |
| 70 | L4 = "L4" |
| 71 | A100 = "A100" |
| 72 | H100 = "H100" |
| 73 | V5E1 = "V5E1" |
| 74 | V6E1 = "V6E1" |
| 75 | |
| 76 | |
| 77 | class Variant(str, Enum): |
| 78 | DEFAULT = "DEFAULT" |
| 79 | GPU = "GPU" |
| 80 | TPU = "TPU" |
| 81 | |
| 82 | |
| 83 | class AssignmentVariant(int, Enum): |
| 84 | DEFAULT = 0 |
| 85 | GPU = 1 |
| 86 | TPU = 2 |
| 87 | |
| 88 | |
| 89 | class Shape(int, Enum): |
| 90 | STANDARD = 0 |
| 91 | HIGH_RAM = 1 |
| 92 | |
| 93 | |
| 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") |
| 128 | url: str |
| 129 | |
| 130 | |
| 131 | class ListedAssignment(BaseModel): |
| 132 | accelerator: Accelerator |
| 133 | endpoint: str |
| 134 | variant: AssignmentVariant |
| 135 | machine_shape: Shape = Field(..., alias="machineShape") |
| 136 | runtime_proxy_info: RuntimeProxyInfo = Field(..., alias="runtimeProxyInfo") |
| 137 | |
| 138 | |
| 139 | class ListedAssignments(BaseModel): |
| 140 | assignments: List[ListedAssignment] |
| 141 | |
| 142 | |
| 143 | class PostAssignmentResponse(BaseModel): |
| 144 | accelerator: Accelerator |
| 145 | endpoint: str |
| 146 | runtime_proxy_info: RuntimeProxyInfo = Field(..., alias="runtimeProxyInfo") |
| 147 | variant: AssignmentVariant |
| 148 | |
| 149 | |
| 150 | class GetAssignmentResponse(BaseModel): |
| 151 | acc: str = Field(..., alias="acc") |
| 152 | nbh: str = Field(..., alias="nbh") |
| 153 | token: str = Field(..., alias="token") |
| 154 | variant: Variant = Field(..., alias="variant") |
| 155 | |
| 156 | |
| 157 | class GetUnassignRequest(BaseModel): |
| 158 | token: str |
| 159 | |
| 160 | |
| 161 | class Assignment(BaseModel): |
| 162 | endpoint: str |
| 163 | runtime_proxy_info: RuntimeProxyInfo = Field(..., alias="runtimeProxyInfo") |
| 164 | |
| 165 | |
| 166 | XSSI_PREFIX = ")]}'\n" |
| 167 | TUN_ENDPOINT = "/tun/m" |
| 168 | |
| 169 | |
| 170 | class ColabRequestError(Exception): |
| 171 | def __init__(self, message, request, response, response_body=None): |
| 172 | super().__init__(message) |
| 173 | self.request = request |
| 174 | self.response = response |
| 175 | self.response_body = response_body |
| 176 | |
| 177 | |
| 178 | class TooManyAssignmentsError(Exception): |
| 179 | pass |
| 180 | |
| 181 | |
| 182 | class Client: |
| 183 | def __init__(self, env: ColabEnvironment, session, logger=None): |
| 184 | self.colab_domain = env.domain |
| 185 | self.colab_api_domain = env.api |
| 186 | self.session = session |
| 187 | self.logger = logger or logging.getLogger(__name__) |
| 188 | |
| 189 | def _strip_xssi_prefix(self, v: str) -> str: |
| 190 | if not v.startswith(XSSI_PREFIX): |
| 191 | return v |
| 192 | return v[len(XSSI_PREFIX) :] |
| 193 | |
| 194 | def _issue_request( |
| 195 | self, |
| 196 | endpoint: str, |
| 197 | method: str = "GET", |
| 198 | headers: Dict[str, str] = None, |
| 199 | params: Dict[str, str] = None, |
| 200 | schema: Optional[BaseModel] = None, |
| 201 | **kwargs, |
| 202 | ): |
| 203 | parsed_endpoint = urlparse(endpoint) |
| 204 | if parsed_endpoint.hostname in urlparse(self.colab_domain).hostname: |
| 205 | if params is None: |
| 206 | params = {} |
| 207 | params["authuser"] = "0" |
| 208 | |
| 209 | request_headers = headers.copy() if headers else {} |
| 210 | request_headers[ACCEPT_JSON_HEADER["key"]] = ACCEPT_JSON_HEADER["value"] |
| 211 | request_headers[COLAB_CLIENT_AGENT_HEADER["key"]] = COLAB_CLIENT_AGENT_HEADER[ |
| 212 | "value" |
| 213 | ] |
| 214 | |
| 215 | self.logger.debug(f"Request: {method} {endpoint}") |
| 216 | self.logger.debug(f"Params: {params}") |
| 217 | |
| 218 | response = self.session.request( |
| 219 | method, endpoint, headers=request_headers, params=params, **kwargs |
| 220 | ) |
| 221 | |
| 222 | self.logger.debug(f"Request Headers: {response.request.headers}") |
| 223 | self.logger.debug(f"Response: {response.status_code} {response.reason}") |
| 224 | self.logger.debug(f"Response Headers: {response.headers}") |
| 225 | self.logger.debug(f"Response Body: {response.text}") |
| 226 | if not response.ok: |
| 227 | raise ColabRequestError( |
| 228 | f"Failed to issue request {method} {endpoint}: {response.reason}", |
| 229 | request=response.request, |
| 230 | response=response, |
| 231 | response_body=response.text, |
| 232 | ) |
| 233 | |
| 234 | body = self._strip_xssi_prefix(response.text) |
| 235 | if not body: |
| 236 | return |
| 237 | # Some endpoints (e.g. KeepAliveAssignment) return a non-empty body |
| 238 | # but the caller doesn't care about the response content — skip |
| 239 | # pydantic validation entirely when no schema was supplied. |
| 240 | if schema is None: |
| 241 | return |
| 242 | return TypeAdapter(schema).validate_python(json.loads(body)) |
| 243 | |
| 244 | def list_assignments(self) -> List[ListedAssignment]: |
| 245 | url = urljoin(self.colab_domain, f"{TUN_ENDPOINT}/assignments") |
| 246 | assignments = self._issue_request(url, schema=ListedAssignments) |
| 247 | return assignments.assignments |
| 248 | |
| 249 | def unassign(self, endpoint: str): |
| 250 | url = urljoin(self.colab_domain, f"{TUN_ENDPOINT}/unassign/{endpoint}") |
| 251 | resp = self._issue_request(url, schema=GetUnassignRequest) |
| 252 | headers = {COLAB_XSRF_TOKEN_HEADER["key"]: resp.token} |
| 253 | return self._issue_request( |
| 254 | url, method="POST", headers=headers, schema=BaseModel |
| 255 | ) |
| 256 | |
| 257 | def assign( |
| 258 | self, |
| 259 | notebook_hash: uuid.UUID, |
| 260 | variant: Optional[Variant] = None, |
| 261 | accelerator: Optional[Accelerator] = None, |
| 262 | shape: Optional[Shape] = None, |
| 263 | ) -> Union[PostAssignmentResponse, Assignment]: |
| 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( |
| 272 | notebook_hash, assignment.token, variant, accelerator, shape |
| 273 | ) |
| 274 | except ColabRequestError as e: |
| 275 | if get_status_code(e) == 412: |
| 276 | raise TooManyAssignmentsError(str(e)) |
| 277 | raise e |
| 278 | |
| 279 | return res |
| 280 | |
| 281 | def _build_assign_url( |
| 282 | self, |
| 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)} |
| 290 | if variant: |
| 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() |
| 299 | return prep.url |
| 300 | |
| 301 | def _get_assignment( |
| 302 | self, |
| 303 | notebook_hash: uuid.UUID, |
| 304 | variant: Optional[Variant] = None, |
| 305 | accelerator: Optional[Accelerator] = None, |
| 306 | shape: Optional[Shape] = None, |
| 307 | ) -> Union[GetAssignmentResponse, Assignment]: |
| 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( |
| 312 | self, |
| 313 | notebook_hash: uuid.UUID, |
| 314 | xsrf_token: str, |
| 315 | variant: Optional[Variant] = None, |
| 316 | accelerator: Optional[Accelerator] = None, |
| 317 | shape: Optional[Shape] = None, |
| 318 | ) -> PostAssignmentResponse: |
| 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 |
| 323 | ) |
| 324 | |
| 325 | def keep_alive_assignment(self, endpoint: str): |
| 326 | """Refreshes the idle timer for the given assignment endpoint. |
| 327 | |
| 328 | TFE notes the activity as soon as the request arrives, then forwards it |
| 329 | to the VM, which does not always respond on this path — so the request |
| 330 | commonly read-times-out even though the keep-alive succeeded. A read |
| 331 | timeout is therefore treated as success; only an actual HTTP error |
| 332 | response (4xx/5xx, e.g. 404 for a deleted assignment) is surfaced. |
| 333 | """ |
| 334 | url = urljoin(self.colab_domain, f"{TUN_ENDPOINT}/{endpoint}/keep-alive/") |
| 335 | headers = {COLAB_TUNNEL_HEADER["key"]: COLAB_TUNNEL_HEADER["value"]} |
| 336 | try: |
| 337 | return self._issue_request( |
| 338 | url, method="GET", headers=headers, timeout=KEEP_ALIVE_TIMEOUT |
| 339 | ) |
| 340 | except requests.exceptions.ReadTimeout: |
| 341 | # The activity was recorded by TFE before the request was forwarded; |
| 342 | # the VM simply didn't answer in time. This is the normal, |
| 343 | # successful case for this path. |
| 344 | return None |