diff --git a/priv/python/albedo_plugins/colab.py b/priv/python/albedo_plugins/colab.py index 6d46e278..b0e693eb 100644 --- a/priv/python/albedo_plugins/colab.py +++ b/priv/python/albedo_plugins/colab.py @@ -9,7 +9,11 @@ Colab exposes, run as ssh's ProxyCommand by `albedo_plugins.ws` (see `connect`). from __future__ import annotations import asyncio +import base64 +import http.client import json +import re +import secrets import shlex import sys import urllib.error @@ -21,7 +25,7 @@ from typing import Any import albedo_ssh from albedo_api import PythonApi, Record -from albedo_plugins.remote import remote +from albedo_plugins.remote import RemoteError, remote API_BASE = "https://colab.research.google.com" CONNECT_TIMEOUT = 60.0 @@ -43,6 +47,10 @@ class ColabError(RuntimeError): self.status = status +class ColabSshSessionError(ColabError): + """The runtime refused ssh because its single session slot is occupied.""" + + class ColabInstance(Record): """One Google Colab runtime instance.""" @@ -390,13 +398,178 @@ async def connect( "ALBEDO_WS_HEADERS": json.dumps({"X-Colab-Ssh-Pubkey": public}), }, ) - return await remote.connect( - host, - remote_cwd=remote_cwd or "/content", - python=python or "/usr/bin/python3", - modules=modules, - timeout=timeout, + try: + return await remote.connect( + host, + remote_cwd=remote_cwd or "/content", + python=python or "/usr/bin/python3", + modules=modules, + timeout=timeout, + ) + except RemoteError as error: + try: + occupied = await asyncio.to_thread( + _ssh_session_occupied, socket_url, public, min(timeout, 10.0) + ) + except (OSError, http.client.HTTPException): + occupied = False + if occupied: + raise ColabSshSessionError( + "Colab allows one ssh session per runtime; a previously held session " + "was not released. This is server-side and does not clear on retry. " + "Use the HTTP file routes to keep working: colab.contents(), " + "colab.read(), colab.pull() (and colab.runtime_status() for liveness). " + "stop()/unassign is the only forceful reset and destroys the VM.", + 429, + ) from error + raise + + +def _ssh_session_occupied(url: str, public: str, timeout: float) -> bool: + target = urllib.parse.urlsplit(url) + transport = ( + http.client.HTTPSConnection + if target.scheme == "wss" + else http.client.HTTPConnection + ) + connection = transport(target.netloc, timeout=timeout) + try: + connection.request( + "GET", + f"{target.path}?{target.query}", + headers={ + "Upgrade": "websocket", + "Connection": "Upgrade", + "Sec-WebSocket-Key": base64.b64encode(secrets.token_bytes(16)).decode(), + "Sec-WebSocket-Version": "13", + "X-Colab-Ssh-Pubkey": public, + }, + ) + response = connection.getresponse() + return response.status == 429 and b"already-active-session" in response.read( + 4096 + ) + finally: + connection.close() + + +async def _runtime_url(instance: str | ColabInstance, path: str, **query: str) -> str: + live = _find(await instances(), instance) + if not (live.url and live.token): + raise ColabError(f"Colab instance {live.id} has no runtime proxy yet") + target = urllib.parse.urlsplit(live.url) + return urllib.parse.urlunsplit( + ( + target.scheme, + target.netloc, + path, + urllib.parse.urlencode({"colab-runtime-proxy-token": live.token, **query}), + "", + ) + ) + + +def _runtime_open(url: str, headers: dict[str, str] | None = None) -> Any: + request = urllib.request.Request( + url, headers={"X-Colab-Client-Agent": "colab-cli", **(headers or {})} + ) + try: + return urllib.request.urlopen(request, timeout=API_TIMEOUT) + except urllib.error.HTTPError as error: + status = error.code + error.close() + raise ColabError( + f"Colab runtime HTTP request failed ({status})", status + ) from None + except urllib.error.URLError: + raise ColabError( + "network error connecting to the Colab runtime proxy" + ) from None + + +def _runtime_json(url: str) -> Any: + with _runtime_open(url) as response: + return json.load(response) + + +async def runtime_status(instance: str | ColabInstance) -> Record: + """Jupyter liveness over HTTPS, independent of the ssh session slot.""" + url = await _runtime_url(instance, "/api/status") + return Record(await asyncio.to_thread(_runtime_json, url)) + + +async def contents( + instance: str | ColabInstance, path: str = "", *, content: bool = True +) -> Record: + """Jupyter metadata, directory entries, or text-file content over HTTPS. + + Paths are relative to Jupyter's root, not the ssh working directory. + Use content=False for file metadata without loading the file. + """ + route = "/api/contents/" + urllib.parse.quote(path.lstrip("/"), safe="/") + url = await _runtime_url( + instance, route, content="1" if content else "0", format="text" ) + return Record(await asyncio.to_thread(_runtime_json, url)) + + +async def read(instance: str | ColabInstance, path: str) -> str: + """Read one text file over HTTPS; use pull() for large or binary files.""" + result = await contents(instance, path) + if result.get("type") != "file" or not isinstance(result.get("content"), str): + raise ColabError(f"not a text file: {path!r}") + return result["content"] + + +def _pull(url: str, local: Path, resume: bool) -> Path: + have = local.stat().st_size if resume and local.exists() else 0 + headers = {"Range": f"bytes={have}-"} if have else {} + with _runtime_open(url, headers) as response: + total = None + if response.status == 206: + match = re.fullmatch( + r"bytes (\d+)-(\d+)/(\d+)", response.headers.get("Content-Range", "") + ) + if not match or int(match[1]) != have or int(match[2]) + 1 != int(match[3]): + raise ColabError( + "invalid Content-Range from Colab; local file left unchanged" + ) + total = int(match[3]) + elif response.status == 200: + # An ignored Range must replace the partial file, not append to it. + have = 0 + length = response.headers.get("Content-Length") + total = int(length) if length is not None else None + else: + raise ColabError(f"unexpected Colab download response ({response.status})") + local.parent.mkdir(parents=True, exist_ok=True) + with local.open("ab" if have else "wb") as output: + while chunk := response.read(1 << 20): + output.write(chunk) + have += len(chunk) + if total is not None and have != total: + raise ColabError( + f"incomplete Colab download: got {have} of {total} bytes; retry to resume" + ) + return local + + +async def pull( + instance: str | ColabInstance, + remote: str, + local: str | Path, + *, + resume: bool = True, +) -> Path: + """Stream a Jupyter-root-relative file to disk over HTTPS, without ssh. + + Partial downloads remain on disk for retry. Resume assumes the remote file + has not changed; use resume=False for mutable files. A 416 response is an + error, not proof that the local file is complete. + """ + route = "/files/" + urllib.parse.quote(remote.lstrip("/"), safe="/") + url = await _runtime_url(instance, route) + return await asyncio.to_thread(_pull, url, Path(local).expanduser(), resume) class Colab: @@ -407,6 +580,10 @@ class Colab: launch = staticmethod(launch) stop = staticmethod(stop) connect = staticmethod(connect) + runtime_status = staticmethod(runtime_status) + contents = staticmethod(contents) + read = staticmethod(read) + pull = staticmethod(pull) colab = Colab() @@ -419,6 +596,7 @@ def setup(api: PythonApi) -> dict[str, object]: return { "colab": colab, "ColabError": ColabError, + "ColabSshSessionError": ColabSshSessionError, "ColabInstance": ColabInstance, "ColabStatus": ColabStatus, } diff --git a/robot-docs/colab.md b/robot-docs/colab.md index 39b4ec5b..269b4dee 100644 --- a/robot-docs/colab.md +++ b/robot-docs/colab.md @@ -24,6 +24,24 @@ runtimes from the kernel and reaches one through `remote`. written to `~/.ssh/config`. a missing ssh public key (`~/.ssh/id_ed25519.pub`, `id_ecdsa.pub`, or `ssh_key=`) is an error before anything connects. `stop()` drops the alias; it does not close a connection already open on it. +- **occupied ssh slot**: only after `remote.connect()` fails, a bounded websocket + handshake probe checks for HTTP 429 with `already-active-session`. this raises + `ColabSshSessionError` (a `ColabError`, `status=429`). unrelated refusals and + probe failures preserve the original `RemoteError`; successful connects cost + no extra request. a previously held server-side session may survive a bridge + restart. retrying does not release it. `stop()`/unassign is the only forceful + reset and **destroys the VM**; never do that to recover files from a live job. +- **HTTP recovery**: `await colab.runtime_status(vm)` reads Jupyter `/api/status`; + `await colab.contents(vm, "content")` lists a directory; + `await colab.contents(vm, path, content=False)` reads metadata (including size); + `await colab.read(vm, path)` returns a text file in memory; + `await colab.pull(vm, path, local, resume=True)` streams `/files/` to disk + in 1 MiB chunks and returns its local `Path`. these read-only routes refresh + the runtime proxy token and work independently of ssh. paths are relative to + **Jupyter's root**, not `/content`. partial downloads remain for retry; resume + assumes unchanged remote content, so use `resume=False` for mutable files. + an ignored Range restarts rather than appending; invalid ranges and HTTP 416 + fail instead of claiming completeness. runtime HTTP errors omit credential URLs. - **tests**: `test/harness/colab_plugin_test.py` (the plugin against a faked HTTP layer), `ws_test.py` (frames, and the bridge in a real subprocess) and `google_oauth_test.gleam` (the refresh against a local token diff --git a/src/albedo/harness/extensions/colab/extension.gleam b/src/albedo/harness/extensions/colab/extension.gleam index 01806b65..db56efa4 100644 --- a/src/albedo/harness/extensions/colab/extension.gleam +++ b/src/albedo/harness/extensions/colab/extension.gleam @@ -30,7 +30,13 @@ const instructions = <> " and await colab.status() checks credentials." <> " rem = await colab.connect(vm) (or await vm.connect()) returns what remote.connect() does," <> " for the VM in /content over Colab's ssh proxy; it needs an ssh public key in ~/.ssh or" - <> " ssh_key=." + <> " ssh_key=. ColabSshSessionError (a ColabError, status=429) means the runtime's single ssh" + <> " session is still held server-side; retrying does not release it. HTTP works independently:" + <> " await colab.runtime_status(vm) checks Jupyter liveness; await colab.contents(vm, path)" + <> " lists a directory (content=False gives metadata); await colab.read(vm, path) reads text;" + <> " await colab.pull(vm, remote, local, resume=True) streams a file to disk in 1 MiB chunks." + <> " Paths are relative to Jupyter's root, not /content; resume assumes unchanged content." + <> " stop()/unassign forcefully resets ssh but destroys the VM; use HTTP to recover files first." <> " Always call await vm.stop() (or await colab.stop(vm)) when done" <> " to release the runtime and preserve compute units. Failures raise ColabError." diff --git a/test/harness/colab_plugin_test.py b/test/harness/colab_plugin_test.py index f6d4038d..02e459a6 100644 --- a/test/harness/colab_plugin_test.py +++ b/test/harness/colab_plugin_test.py @@ -1,13 +1,21 @@ -"""Unit tests for the Google Colab tool plugin (albedo_plugins.colab).""" +"""Colab protocol tests: local refusals and interrupted downloads need no live GPU. + +The daemon cannot deterministically reproduce an occupied Colab ssh slot or a +proxy that ignores Range; these tests exercise those responses locally. +""" from __future__ import annotations import asyncio +from contextlib import contextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer import json import os from pathlib import Path import sys import tempfile +import threading +import urllib.parse import unittest from unittest.mock import AsyncMock, patch @@ -46,6 +54,35 @@ def instance(**fields) -> colab.ColabInstance: ) +@contextmanager +def runtime_proxy(status, body, headers=None): + requests = [] + + class Handler(BaseHTTPRequestHandler): + def do_GET(self): + requests.append((self.path, self.headers)) + self.send_response(status) + for key, value in (headers or {}).items(): + self.send_header(key, value) + self.end_headers() + self.wfile.write(body) + + def log_message(self, format, *args): + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread( + target=server.serve_forever, kwargs={"poll_interval": 0.01}, daemon=True + ) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_port}", requests + finally: + server.shutdown() + server.server_close() + thread.join() + + class ColabPluginTest(unittest.IsolatedAsyncioTestCase): def setUp(self): self.home = tempfile.mkdtemp(prefix="albedo-colab-home-") @@ -289,6 +326,168 @@ class ColabPluginTest(unittest.IsolatedAsyncioTestCase): self.assertIs(colab.HostErrorType, RuntimeError) loop.close() self.assertEqual( - sorted(exports), ["ColabError", "ColabInstance", "ColabStatus", "colab"] + sorted(exports), + [ + "ColabError", + "ColabInstance", + "ColabSshSessionError", + "ColabStatus", + "colab", + ], ) self.assertIs(exports["colab"], colab.colab) + + async def test_failed_connect_classifies_only_the_specific_refusal(self): + (Path(self.home) / ".ssh" / "id_ed25519.pub").write_text( + "ssh-ed25519 AAAA test" + ) + for status, body, occupied in ( + (429, b'{"error":"already-active-session"}', True), + (429, b'{"error":"rate-limit"}', False), + (502, b"already-active-session", False), + (101, b"", False), + ): + with ( + self.subTest(status=status, body=body), + runtime_proxy(status, body) as (url, requests), + ): + live = instance(url=url, token="secret/never-log") + original = colab.RemoteError("original failure") + with ( + patch.object(colab, "instances", AsyncMock(return_value=[live])), + patch.object( + colab.remote, "connect", AsyncMock(side_effect=original) + ), + ): + with self.assertRaises( + colab.ColabSshSessionError if occupied else colab.RemoteError + ) as caught: + await colab.connect(live) + if occupied: + assert isinstance(caught.exception, colab.ColabSshSessionError) + self.assertEqual(caught.exception.status, 429) + self.assertNotIn(live.token, str(caught.exception)) + self.assertIn("destroys the VM", str(caught.exception)) + else: + self.assertIs(caught.exception, original) + path, headers = requests[0] + self.assertIn("/colab/ssh?", path) + self.assertEqual(headers["Upgrade"], "websocket") + self.assertEqual(headers["X-Colab-Ssh-Pubkey"], "ssh-ed25519 AAAA test") + self.assertEqual(headers["Sec-WebSocket-Version"], "13") + + async def test_probe_network_failure_preserves_original_error(self): + (Path(self.home) / ".ssh" / "id_ed25519.pub").write_text( + "ssh-ed25519 AAAA test" + ) + live = instance(url="https://vm", token="secret") + original = colab.RemoteError("original") + with ( + patch.object(colab, "instances", AsyncMock(return_value=[live])), + patch.object(colab.remote, "connect", AsyncMock(side_effect=original)), + patch.object( + colab, "_ssh_session_occupied", side_effect=OSError("probe failure") + ), + ): + with self.assertRaises(colab.RemoteError) as caught: + await colab.connect(live) + self.assertIs(caught.exception, original) + + async def test_success_does_not_probe(self): + (Path(self.home) / ".ssh" / "id_ed25519.pub").write_text( + "ssh-ed25519 AAAA test" + ) + live = instance(url="https://vm", token="secret") + with ( + patch.object(colab, "instances", AsyncMock(return_value=[live])), + patch.object(colab.remote, "connect", AsyncMock(return_value="connection")), + patch.object(colab, "_ssh_session_occupied") as probe, + ): + self.assertEqual(await colab.connect(live), "connection") + probe.assert_not_called() + + async def test_runtime_helpers_refresh_and_quote_paths_without_ssh(self): + response = {"type": "file", "content": "hello"} + for name, path, route in ( + ("status", "", "/api/status"), + ("contents", "a ?#%/b", "/api/contents/a%20%3F%23%25/b"), + ("read", "text", "/api/contents/text"), + ): + with runtime_proxy(200, json.dumps(response).encode()) as (url, requests): + live = instance(url=url, token="secret") + with ( + patch.object( + colab, "instances", AsyncMock(return_value=[live]) + ) as listing, + patch.object(colab.remote, "connect", AsyncMock()) as ssh, + ): + if name == "status": + result = await colab.runtime_status(live) + elif name == "contents": + result = await colab.contents(live, path) + else: + result = await colab.read(live, path) + self.assertEqual(result, "hello" if name == "read" else response) + listing.assert_awaited_once() + ssh.assert_not_called() + self.assertEqual(urllib.parse.urlsplit(requests[0][0]).path, route) + + async def test_pull_resume_and_ignored_range(self): + for status, body, headers in ( + (206, b"def", {"Content-Range": "bytes 3-5/6"}), + (200, b"abcdef", {"Content-Length": "6"}), + ): + with ( + self.subTest(status=status), + runtime_proxy(status, body, headers) as (url, requests), + ): + live = instance(url=url, token="secret") + local = Path(self.home) / "download" + local.write_bytes(b"abc") + with patch.object(colab, "instances", AsyncMock(return_value=[live])): + result = await colab.pull(live, "remote ?#", local) + self.assertEqual(result, local) + self.assertEqual(local.read_bytes(), b"abcdef") + self.assertEqual(requests[0][1]["Range"], "bytes=3-") + self.assertIn("/files/remote%20%3F%23?", requests[0][0]) + + async def test_pull_rejects_bad_ranges_and_http_errors_without_overwriting(self): + for status, headers in ( + (206, {"Content-Range": "bytes 0-2/3"}), + (206, {"Content-Range": "garbage"}), + (416, {}), + (403, {}), + ): + with ( + self.subTest(status=status), + runtime_proxy(status, b"xyz", headers) as (url, _), + ): + local = Path(self.home) / "download" + local.write_bytes(b"abc") + with self.assertRaises(colab.ColabError) as caught: + await asyncio.to_thread( + colab._pull, url + "?token=secret", local, True + ) + self.assertEqual(local.read_bytes(), b"abc") + self.assertNotIn("secret", str(caught.exception)) + + async def test_pull_detects_truncation_and_keeps_partial_file(self): + with runtime_proxy(206, b"d", {"Content-Range": "bytes 3-5/6"}) as (url, _): + local = Path(self.home) / "download" + local.write_bytes(b"abc") + with self.assertRaisesRegex(colab.ColabError, "incomplete"): + await asyncio.to_thread(colab._pull, url, local, True) + self.assertEqual(local.read_bytes(), b"abcd") + + async def test_pull_fresh_multichunk_file_replaces_local_when_resume_disabled(self): + body = b"x" * ((2 << 20) + 17) + with runtime_proxy(200, body, {"Content-Length": str(len(body))}) as ( + url, + requests, + ): + local = Path(self.home) / "nested" / "download" + local.parent.mkdir() + local.write_bytes(b"old") + await asyncio.to_thread(colab._pull, url, local, False) + self.assertEqual(local.read_bytes(), body) + self.assertNotIn("Range", requests[0][1])