okami.mom/albedo

git clone https://git.nkp.pet/okami.mom/albedo

nekomimi.petcolab: refresh tokens in the daemon, share the websocket and google oauth code, drop the ssh config editsc9cb3a0

main
branchesmainall refs →
11.2 KiB294 linesraw
1"""Unit tests for the Google Colab tool plugin (albedo_plugins.colab)."""
2
3from __future__ import annotations
4
5import asyncio
6import json
7import os
8from pathlib import Path
9import sys
10import tempfile
11import unittest
12from unittest.mock import AsyncMock, patch
13
14ROOT = Path(__file__).resolve().parents[2]
15sys.path.insert(0, str(ROOT / "priv" / "python"))
16sys.path.insert(0, str(ROOT / "test"))
17
18import albedo_ssh  # noqa: E402
19from albedo_api import PythonApi  # noqa: E402
20from albedo_capture import Capture  # noqa: E402
21from albedo_plugins import colab  # noqa: E402
22import scratch  # noqa: E402
23
24scratch.claim("harness")
25
26REAL_ACCOUNT = colab._account
27ACCOUNT = {"access": "test-token", "email": "user@example.com"}
28
29
30def instance(**fields) -> colab.ColabInstance:
31    return colab.ColabInstance(
32        {
33            "id": "gpu-t4-1",
34            "endpoint": "gpu-t4-1",
35            "name": "dev-t4",
36            "variant": "GPU",
37            "accelerator": "T4",
38            "shape": "standard",
39            "status": "active",
40            "url": None,
41            "token": None,
42            "token_expires_in": None,
43            "raw": {},
44            **fields,
45        }
46    )
47
48
49class ColabPluginTest(unittest.IsolatedAsyncioTestCase):
50    def setUp(self):
51        self.home = tempfile.mkdtemp(prefix="albedo-colab-home-")
52        (Path(self.home) / ".ssh").mkdir()
53        env = patch.dict(os.environ, {"HOME": self.home})
54        env.start()
55        self.addCleanup(env.stop)
56        account = patch.object(colab, "_account", AsyncMock(return_value=ACCOUNT))
57        account.start()
58        self.addCleanup(account.stop)
59        self.addCleanup(albedo_ssh.ALIASES.clear)
60
61    async def test_account_comes_from_the_daemon_and_its_refusal_is_a_colab_error(self):
62        class Refused(Exception):
63            pass
64
65        calls = []
66
67        async def host(method, args):
68            calls.append(method)
69            raise Refused("Google Colab is not authenticated; run /login")
70
71        with (
72            patch.object(colab, "_account", REAL_ACCOUNT),
73            patch.multiple(colab, host_call=host, HostErrorType=Refused),
74        ):
75            with self.assertRaises(colab.ColabError) as cm:
76                await colab._account()
77        self.assertEqual(calls, ["colab.account"])
78        self.assertIn("not authenticated", str(cm.exception))
79
80    async def test_instances_query(self):
81        body = (
82            b")]}'\n"
83            + json.dumps(
84                {
85                    "assignments": [
86                        {
87                            "endpoint": "gpu-t4-s-12345",
88                            "name": "colab-t4",
89                            "variant": "GPU",
90                            "accelerator": "T4",
91                            "shape": "standard",
92                            "status": "active",
93                            "runtimeProxyInfo": {
94                                "url": "https://vm1.colab.googleusercontent.com",
95                                "token": "proxy-tok-123",
96                                "tokenExpiresInSeconds": 3600,
97                            },
98                        }
99                    ]
100                }
101            ).encode()
102        )
103        with patch.object(colab, "_send", return_value=(200, body)) as send:
104            found = await colab.instances()
105        self.assertEqual(len(found), 1)
106        inst = found[0]
107        self.assertEqual(inst.id, "gpu-t4-s-12345")
108        self.assertEqual(inst.gpu, "T4")
109        self.assertEqual(inst.url, "https://vm1.colab.googleusercontent.com")
110        self.assertEqual(inst.token, "proxy-tok-123")
111        method, url, headers, _, _ = send.call_args[0]
112        self.assertEqual(method, "GET")
113        self.assertIn("/tun/m/assignments?authuser=0", url)
114        self.assertEqual(headers["Authorization"], "Bearer test-token")
115        self.assertEqual(headers["X-Colab-Tunnel"], "Google")
116
117    async def test_status_summary_and_single_lookup(self):
118        with patch.object(colab, "instances", AsyncMock(return_value=[instance()])):
119            overall = await colab.status()
120            self.assertTrue(overall.authenticated)
121            self.assertEqual(overall.email, "user@example.com")
122            self.assertEqual(overall.active, 1)
123            self.assertIn("gpu-t4-1", str(overall))
124            self.assertEqual((await colab.status("gpu-t4-1")).id, "gpu-t4-1")
125            self.assertEqual((await colab.status("dev-t4")).id, "gpu-t4-1")
126            with self.assertRaises(colab.ColabError):
127                await colab.status("missing-instance")
128
129    def test_notebook_hash(self):
130        nbh = colab._notebook_hash()
131        self.assertEqual(len(nbh), 44)
132        self.assertNotIn("-", nbh)
133        self.assertTrue(nbh.endswith("."))
134
135    async def test_launch_provisions_with_the_xsrf_token(self):
136        xsrf = b')]}\'\n{"token": "xsrf-token-123"}'
137        assigned = json.dumps(
138            {
139                "assignment": {
140                    "endpoint": "gpu-t4-assigned-999",
141                    "variant": "GPU",
142                    "accelerator": "T4",
143                    "shape": "highmem",
144                    "runtimeProxyInfo": {"url": "https://vm", "token": "tok-t4"},
145                }
146            }
147        ).encode()
148        with patch.object(
149            colab, "_send", side_effect=[(200, xsrf), (200, assigned)]
150        ) as send:
151            inst = await colab.launch(gpu="T4", mem="high")
152        self.assertEqual(inst.id, "gpu-t4-assigned-999")
153        self.assertEqual(inst.mem, "high")
154        (m1, u1, _, b1, _), (m2, u2, h2, b2, _) = (c[0] for c in send.call_args_list)
155        self.assertEqual(m1, "GET")
156        for part in (
157            "/tun/m/assign",
158            "authuser=0",
159            "variant=GPU",
160            "accelerator=T4",
161            "shape=hm",
162        ):
163            self.assertIn(part, u1)
164        self.assertIsNone(b1)
165        self.assertEqual(m2, "POST")
166        self.assertIn("/tun/m/assign", u2)
167        self.assertEqual(h2["X-Goog-Colab-Token"], "xsrf-token-123")
168        self.assertIsNone(b2)
169
170    async def test_launch_returns_an_existing_assignment(self):
171        existing = json.dumps(
172            {
173                "endpoint": "gpu-t4-existing-111",
174                "variant": "GPU",
175                "accelerator": "T4",
176                "runtimeProxyInfo": {"url": "https://vm", "token": "tok"},
177            }
178        ).encode()
179        with patch.object(colab, "_send", return_value=(200, existing)) as send:
180            inst = await colab.launch(gpu="T4")
181        self.assertEqual(inst.id, "gpu-t4-existing-111")
182        self.assertEqual(send.call_count, 1)
183
184    async def test_launch_capacity_exhausted(self):
185        with patch.object(colab, "_send", return_value=(503, b"Resource exhausted")):
186            with self.assertRaises(colab.ColabError) as cm:
187                await colab.launch(gpu="A100")
188        self.assertIn("capacity exhausted", str(cm.exception))
189        self.assertEqual(cm.exception.status, 503)
190
191    async def test_stop_instance_and_already_stopped(self):
192        xsrf = b')]}\'\n{"token": "unassign-xsrf-456"}'
193        with patch.object(
194            colab, "_send", side_effect=[(200, xsrf), (204, b"")]
195        ) as send:
196            result = await colab.stop("gpu-t4-s-12345")
197        self.assertTrue(result.success)
198        (m1, u1, _, _, _), (m2, u2, h2, _, _) = (c[0] for c in send.call_args_list)
199        self.assertEqual((m1, m2), ("GET", "POST"))
200        for url in (u1, u2):
201            self.assertIn("/tun/m/unassign/gpu-t4-s-12345", url)
202        self.assertEqual(h2["X-Goog-Colab-Token"], "unassign-xsrf-456")
203
204        with patch.object(colab, "_send", return_value=(404, b"Not Found")):
205            result = await colab.stop("already-stopped")
206        self.assertFalse(result.success)
207        self.assertIn("not found or already stopped", str(result.message))
208
209    async def test_connect_names_the_host_in_memory_and_leaves_ssh_config_alone(self):
210        ssh = Path(self.home) / ".ssh"
211        (ssh / "id_ed25519.pub").write_text("ssh-ed25519 AAAA test@box\n")
212        live = instance(
213            url="https://vm.colab.googleusercontent.com", token="proxy tok/1"
214        )
215        reached = AsyncMock(return_value="connection")
216        with (
217            patch.object(colab, "instances", AsyncMock(return_value=[live])),
218            patch.object(colab.remote, "connect", reached),
219        ):
220            result = await live.connect(remote_cwd="/content/project")
221
222        self.assertEqual(result, "connection")
223        reached.assert_awaited_once_with(
224            "colab-gpu-t4-1",
225            remote_cwd="/content/project",
226            python="/usr/bin/python3",
227            modules=None,
228            timeout=60.0,
229        )
230        self.assertEqual(list(ssh.iterdir()), [ssh / "id_ed25519.pub"])
231
232        argv = albedo_ssh.base("colab-gpu-t4-1")
233        options = [argv[i + 1] for i, arg in enumerate(argv) if arg == "-o"]
234        self.assertTrue(options[0].startswith("ProxyCommand="))
235        self.assertIn("-m albedo_plugins.ws", options[0])
236        self.assertLess(
237            options.index("StrictHostKeyChecking=no"),
238            options.index("StrictHostKeyChecking=accept-new"),
239        )
240        self.assertIn("User=root", options)
241
242        environment = albedo_ssh.env("colab-gpu-t4-1")
243        self.assertEqual(
244            environment["ALBEDO_WS_URL"],
245            "wss://vm.colab.googleusercontent.com/colab/ssh?colab-runtime-proxy-token=proxy+tok%2F1",
246        )
247        self.assertEqual(
248            json.loads(environment["ALBEDO_WS_HEADERS"]),
249            {"X-Colab-Ssh-Pubkey": "ssh-ed25519 AAAA test@box"},
250        )
251        self.assertNotIn("wss://", " ".join(argv))
252        self.assertNotIn("proxy+tok", " ".join(argv))
253        self.assertNotIn("ALBEDO_WS_URL", albedo_ssh.env("other-host"))
254
255        with patch.object(colab, "_send", return_value=(404, b"")):
256            await colab.stop(live)
257        self.assertNotIn("colab-gpu-t4-1", albedo_ssh.ALIASES)
258
259    async def test_connect_needs_a_runtime_proxy_and_a_public_key(self):
260        waiting = instance()
261        with patch.object(colab, "instances", AsyncMock(return_value=[waiting])):
262            with self.assertRaises(colab.ColabError) as cm:
263                await colab.connect("gpu-t4-1")
264        self.assertIn("no runtime proxy", str(cm.exception))
265
266        ready = instance(url="https://vm", token="t")
267        with patch.object(colab, "instances", AsyncMock(return_value=[ready])):
268            with self.assertRaises(colab.ColabError) as cm:
269                await colab.connect("gpu-t4-1")
270        self.assertIn("no ssh public key", str(cm.exception))
271        self.assertEqual(albedo_ssh.ALIASES, {})
272
273    def test_setup_exports(self):
274        loop = asyncio.new_event_loop()
275        api = PythonApi(
276            loop,
277            lambda m, a: asyncio.sleep(0),
278            RuntimeError,
279            Capture,
280            64 * 1024,
281            lambda e: None,
282            lambda c: None,
283            lambda cls: None,
284        )
285        with patch.multiple(
286            colab, host_call=colab.host_call, HostErrorType=colab.HostErrorType
287        ):
288            exports = colab.setup(api)
289            self.assertIs(colab.HostErrorType, RuntimeError)
290        loop.close()
291        self.assertEqual(
292            sorted(exports), ["ColabError", "ColabInstance", "ColabStatus", "colab"]
293        )
294        self.assertIs(exports["colab"], colab.colab)