lablink-cli 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (71) hide show
  1. lablink_cli-0.1.0/PKG-INFO +76 -0
  2. lablink_cli-0.1.0/README.md +55 -0
  3. lablink_cli-0.1.0/pyproject.toml +59 -0
  4. lablink_cli-0.1.0/setup.cfg +4 -0
  5. lablink_cli-0.1.0/src/lablink_cli/__init__.py +8 -0
  6. lablink_cli-0.1.0/src/lablink_cli/api.py +428 -0
  7. lablink_cli-0.1.0/src/lablink_cli/app.py +938 -0
  8. lablink_cli-0.1.0/src/lablink_cli/byo_detect.py +112 -0
  9. lablink_cli-0.1.0/src/lablink_cli/commands/__init__.py +0 -0
  10. lablink_cli-0.1.0/src/lablink_cli/commands/cleanup.py +647 -0
  11. lablink_cli-0.1.0/src/lablink_cli/commands/deploy.py +863 -0
  12. lablink_cli-0.1.0/src/lablink_cli/commands/deploy_compose.py +1203 -0
  13. lablink_cli-0.1.0/src/lablink_cli/commands/doctor.py +549 -0
  14. lablink_cli-0.1.0/src/lablink_cli/commands/export_metrics.py +244 -0
  15. lablink_cli-0.1.0/src/lablink_cli/commands/launch.py +236 -0
  16. lablink_cli-0.1.0/src/lablink_cli/commands/logs.py +434 -0
  17. lablink_cli-0.1.0/src/lablink_cli/commands/register.py +839 -0
  18. lablink_cli-0.1.0/src/lablink_cli/commands/reset_overlay.py +109 -0
  19. lablink_cli-0.1.0/src/lablink_cli/commands/setup.py +347 -0
  20. lablink_cli-0.1.0/src/lablink_cli/commands/stats.py +133 -0
  21. lablink_cli-0.1.0/src/lablink_cli/commands/status.py +934 -0
  22. lablink_cli-0.1.0/src/lablink_cli/commands/unregister.py +188 -0
  23. lablink_cli-0.1.0/src/lablink_cli/commands/utils.py +552 -0
  24. lablink_cli-0.1.0/src/lablink_cli/config/__init__.py +0 -0
  25. lablink_cli-0.1.0/src/lablink_cli/config/schema.py +212 -0
  26. lablink_cli-0.1.0/src/lablink_cli/deployment_metrics.py +94 -0
  27. lablink_cli-0.1.0/src/lablink_cli/docker.py +419 -0
  28. lablink_cli-0.1.0/src/lablink_cli/log_shipper.py +441 -0
  29. lablink_cli-0.1.0/src/lablink_cli/templates/docker-compose.tailscale-override.yml +55 -0
  30. lablink_cli-0.1.0/src/lablink_cli/templates/docker-compose.yml +67 -0
  31. lablink_cli-0.1.0/src/lablink_cli/tofu_source.py +169 -0
  32. lablink_cli-0.1.0/src/lablink_cli/tui/__init__.py +0 -0
  33. lablink_cli-0.1.0/src/lablink_cli/tui/logs_viewer.py +413 -0
  34. lablink_cli-0.1.0/src/lablink_cli/tui/wizard.py +1814 -0
  35. lablink_cli-0.1.0/src/lablink_cli.egg-info/PKG-INFO +76 -0
  36. lablink_cli-0.1.0/src/lablink_cli.egg-info/SOURCES.txt +69 -0
  37. lablink_cli-0.1.0/src/lablink_cli.egg-info/dependency_links.txt +1 -0
  38. lablink_cli-0.1.0/src/lablink_cli.egg-info/entry_points.txt +2 -0
  39. lablink_cli-0.1.0/src/lablink_cli.egg-info/requires.txt +12 -0
  40. lablink_cli-0.1.0/src/lablink_cli.egg-info/top_level.txt +1 -0
  41. lablink_cli-0.1.0/tests/test_api.py +342 -0
  42. lablink_cli-0.1.0/tests/test_app.py +651 -0
  43. lablink_cli-0.1.0/tests/test_byo_detect.py +116 -0
  44. lablink_cli-0.1.0/tests/test_cleanup.py +590 -0
  45. lablink_cli-0.1.0/tests/test_cleanup_full.py +144 -0
  46. lablink_cli-0.1.0/tests/test_deploy.py +836 -0
  47. lablink_cli-0.1.0/tests/test_deploy_compose.py +2893 -0
  48. lablink_cli-0.1.0/tests/test_deployment_metrics.py +133 -0
  49. lablink_cli-0.1.0/tests/test_docker.py +297 -0
  50. lablink_cli-0.1.0/tests/test_docker_isolation.py +118 -0
  51. lablink_cli-0.1.0/tests/test_doctor.py +387 -0
  52. lablink_cli-0.1.0/tests/test_doctor_full.py +233 -0
  53. lablink_cli-0.1.0/tests/test_export_metrics.py +772 -0
  54. lablink_cli-0.1.0/tests/test_launch.py +400 -0
  55. lablink_cli-0.1.0/tests/test_log_shipper.py +673 -0
  56. lablink_cli-0.1.0/tests/test_logs.py +371 -0
  57. lablink_cli-0.1.0/tests/test_register.py +1776 -0
  58. lablink_cli-0.1.0/tests/test_registration_client.py +189 -0
  59. lablink_cli-0.1.0/tests/test_reset_overlay.py +134 -0
  60. lablink_cli-0.1.0/tests/test_schema.py +260 -0
  61. lablink_cli-0.1.0/tests/test_setup.py +240 -0
  62. lablink_cli-0.1.0/tests/test_setup_full.py +97 -0
  63. lablink_cli-0.1.0/tests/test_stats.py +231 -0
  64. lablink_cli-0.1.0/tests/test_status.py +724 -0
  65. lablink_cli-0.1.0/tests/test_status_full.py +373 -0
  66. lablink_cli-0.1.0/tests/test_status_manual.py +334 -0
  67. lablink_cli-0.1.0/tests/test_tofu_source.py +241 -0
  68. lablink_cli-0.1.0/tests/test_unregister.py +193 -0
  69. lablink_cli-0.1.0/tests/test_utils.py +528 -0
  70. lablink_cli-0.1.0/tests/test_utils_full.py +263 -0
  71. lablink_cli-0.1.0/tests/test_wizard.py +1049 -0
@@ -0,0 +1,76 @@
1
+ Metadata-Version: 2.4
2
+ Name: lablink-cli
3
+ Version: 0.1.0
4
+ Summary: CLI tool for deploying and managing LabLink infrastructure
5
+ Author-email: Elizabeth Berrigan <eberrigan@salk.edu>, Talmo Pereira <talmo@salk.edu>, Andrew Park <hep003@ucsd.edu>
6
+ Project-URL: Homepage, https://github.com/talmolab/lablink
7
+ Project-URL: Issues, https://github.com/talmolab/lablink/issues
8
+ Requires-Python: >=3.10
9
+ Description-Content-Type: text/markdown
10
+ Requires-Dist: lablink-allocator-service[config]>=0.2.0
11
+ Requires-Dist: typer>=0.15
12
+ Requires-Dist: textual>=3.0
13
+ Requires-Dist: rich>=13.0
14
+ Requires-Dist: pyyaml>=6.0
15
+ Requires-Dist: boto3>=1.35
16
+ Requires-Dist: psutil>=5.9
17
+ Provides-Extra: dev
18
+ Requires-Dist: pytest; extra == "dev"
19
+ Requires-Dist: pytest-cov; extra == "dev"
20
+ Requires-Dist: ruff; extra == "dev"
21
+
22
+ # LabLink CLI
23
+
24
+ Command-line tool for deploying and managing LabLink teaching lab infrastructure —
25
+ on AWS (EC2 via Terraform) or on a machine you already have (`provider: manual`,
26
+ deployed with docker-compose).
27
+
28
+ ## Installation
29
+
30
+ This package is not yet published to PyPI. Install it from source — the repo is a
31
+ `uv` workspace, so sync all three packages into the shared root venv:
32
+
33
+ ```bash
34
+ git clone https://github.com/talmolab/lablink.git
35
+ cd lablink
36
+ uv sync --all-packages
37
+ ```
38
+
39
+ ## Usage
40
+
41
+ ```bash
42
+ lablink --help
43
+ lablink --version # or -v
44
+ ```
45
+
46
+ ### Commands
47
+
48
+ | Command | Description |
49
+ | ---------------- | ---------------------------------------------------------------------------------------- |
50
+ | `configure` | Create or edit LabLink configuration (interactive TUI) |
51
+ | `setup` | Provision provider-specific bootstrap resources (AWS: S3 + DynamoDB for Terraform state) |
52
+ | `doctor` | Check prerequisites and configuration |
53
+ | `deploy` | Deploy LabLink infrastructure (AWS Terraform or docker-compose) |
54
+ | `destroy` | Tear down LabLink infrastructure |
55
+ | `status` | Show deployment health and inventory |
56
+ | `logs` | View allocator and client logs |
57
+ | `export-metrics` | Export deployment metrics to CSV or JSON |
58
+ | `stats` | Show a cohort session-metrics summary in the terminal |
59
+ | `cleanup` | Remove deployment resources and local state |
60
+ | `show-config` | View the current LabLink configuration |
61
+ | `cache-clear` | Clear LabLink caches (Terraform templates, deployment metrics) |
62
+
63
+ ### Client fleet commands
64
+
65
+ | Command | Description |
66
+ | ---------------------- | -------------------------------------------------------------------------------- |
67
+ | `client launch` | Launch client VMs via the allocator service (AWS provider) |
68
+ | `client register` | Register this bring-your-own box as a manual client and run the client container |
69
+ | `client unregister` | Tear down a registered BYO box |
70
+ | `client reset-overlay` | Discard this box's persisted mesh-overlay node identity |
71
+
72
+ Run `lablink <command> --help` for details on any command.
73
+
74
+ ## Documentation
75
+
76
+ Full CLI documentation: https://talmolab.github.io/lablink/cli/
@@ -0,0 +1,55 @@
1
+ # LabLink CLI
2
+
3
+ Command-line tool for deploying and managing LabLink teaching lab infrastructure —
4
+ on AWS (EC2 via Terraform) or on a machine you already have (`provider: manual`,
5
+ deployed with docker-compose).
6
+
7
+ ## Installation
8
+
9
+ This package is not yet published to PyPI. Install it from source — the repo is a
10
+ `uv` workspace, so sync all three packages into the shared root venv:
11
+
12
+ ```bash
13
+ git clone https://github.com/talmolab/lablink.git
14
+ cd lablink
15
+ uv sync --all-packages
16
+ ```
17
+
18
+ ## Usage
19
+
20
+ ```bash
21
+ lablink --help
22
+ lablink --version # or -v
23
+ ```
24
+
25
+ ### Commands
26
+
27
+ | Command | Description |
28
+ | ---------------- | ---------------------------------------------------------------------------------------- |
29
+ | `configure` | Create or edit LabLink configuration (interactive TUI) |
30
+ | `setup` | Provision provider-specific bootstrap resources (AWS: S3 + DynamoDB for Terraform state) |
31
+ | `doctor` | Check prerequisites and configuration |
32
+ | `deploy` | Deploy LabLink infrastructure (AWS Terraform or docker-compose) |
33
+ | `destroy` | Tear down LabLink infrastructure |
34
+ | `status` | Show deployment health and inventory |
35
+ | `logs` | View allocator and client logs |
36
+ | `export-metrics` | Export deployment metrics to CSV or JSON |
37
+ | `stats` | Show a cohort session-metrics summary in the terminal |
38
+ | `cleanup` | Remove deployment resources and local state |
39
+ | `show-config` | View the current LabLink configuration |
40
+ | `cache-clear` | Clear LabLink caches (Terraform templates, deployment metrics) |
41
+
42
+ ### Client fleet commands
43
+
44
+ | Command | Description |
45
+ | ---------------------- | -------------------------------------------------------------------------------- |
46
+ | `client launch` | Launch client VMs via the allocator service (AWS provider) |
47
+ | `client register` | Register this bring-your-own box as a manual client and run the client container |
48
+ | `client unregister` | Tear down a registered BYO box |
49
+ | `client reset-overlay` | Discard this box's persisted mesh-overlay node identity |
50
+
51
+ Run `lablink <command> --help` for details on any command.
52
+
53
+ ## Documentation
54
+
55
+ Full CLI documentation: https://talmolab.github.io/lablink/cli/
@@ -0,0 +1,59 @@
1
+ [project]
2
+ name = "lablink-cli"
3
+ authors = [
4
+ {name = "Elizabeth Berrigan", email = "eberrigan@salk.edu"},
5
+ {name = "Talmo Pereira", email = "talmo@salk.edu"},
6
+ {name = "Andrew Park", email = "hep003@ucsd.edu"},
7
+ ]
8
+ version = "0.1.0"
9
+ description = "CLI tool for deploying and managing LabLink infrastructure"
10
+ readme = "README.md"
11
+ requires-python = ">=3.10"
12
+ dependencies = [
13
+ # >=0.2.0: deploy_compose imports PUBLIC_HOSTNAME_HINT,
14
+ # is_valid_public_hostname and is_weak_admin_password, none of which
15
+ # exist in 0.1.2 — an unpinned resolve breaks `provider: manual`.
16
+ "lablink-allocator-service[config]>=0.2.0",
17
+ "typer>=0.15",
18
+ "textual>=3.0",
19
+ "rich>=13.0",
20
+ "pyyaml>=6.0",
21
+ "boto3>=1.35",
22
+ "psutil>=5.9",
23
+ ]
24
+
25
+ [project.optional-dependencies]
26
+ dev = ["pytest", "pytest-cov", "ruff"]
27
+
28
+ [project.scripts]
29
+ lablink = "lablink_cli.app:main"
30
+
31
+ [build-system]
32
+ requires = ["setuptools>=75"]
33
+ build-backend = "setuptools.build_meta"
34
+
35
+ [tool.setuptools.packages.find]
36
+ where = ["src"]
37
+
38
+ [tool.setuptools.package-data]
39
+ lablink_cli = ["templates/*.yml"]
40
+
41
+ [tool.pytest.ini_options]
42
+ testpaths = ["tests"]
43
+ markers = ["integration: real network tests (deselect with '-m not integration')"]
44
+ addopts = "-m 'not integration'"
45
+
46
+ [tool.coverage.run]
47
+ omit = [
48
+ "src/lablink_cli/tui/*",
49
+ ]
50
+
51
+ [tool.coverage.report]
52
+ exclude_lines = [
53
+ "pragma: no cover",
54
+ "if __name__",
55
+ ]
56
+
57
+ [project.urls]
58
+ Homepage = "https://github.com/talmolab/lablink"
59
+ Issues = "https://github.com/talmolab/lablink/issues"
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,8 @@
1
+ """LabLink CLI - Deploy and manage LabLink infrastructure."""
2
+
3
+ TEMPLATE_REPO = "talmolab/lablink-template"
4
+ TEMPLATE_VERSION = "v0.3.0"
5
+ # SHA-256 of the GitHub release tarball for TEMPLATE_VERSION.
6
+ # Update this when bumping TEMPLATE_VERSION.
7
+ # To compute: curl -sL <tarball_url> | sha256sum
8
+ TEMPLATE_SHA256 = "3d3f8803990f80e6bcda0b3098178d909bc454315fdf2304e70dbee5d43e3bb6"
@@ -0,0 +1,428 @@
1
+ """HTTP client for the LabLink allocator service."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import base64
6
+ import json
7
+ import ssl
8
+ import time
9
+ from typing import Callable, NoReturn
10
+ from urllib.error import HTTPError, URLError
11
+ from urllib.parse import urlencode
12
+ from urllib.request import Request, urlopen
13
+
14
+ try:
15
+ from importlib.metadata import version as _pkg_version
16
+
17
+ _CLI_VERSION = _pkg_version("lablink-cli")
18
+ except Exception: # pragma: no cover - package metadata unavailable
19
+ _CLI_VERSION = "0.0.0"
20
+
21
+ # Product User-Agent for all CLI HTTP requests. urllib's default
22
+ # "Python-urllib/x.y" is blocked with HTTP 403 by Cloudflare-proxied
23
+ # allocators, so every Request must identify itself with this instead.
24
+ USER_AGENT = f"lablink-cli/{_CLI_VERSION}"
25
+
26
+
27
+ def ssl_context(ssl_provider: str) -> ssl.SSLContext:
28
+ """Verifying context, except under `self_signed` where there is no CA
29
+ to verify against and the operator has opted into that trade."""
30
+ ctx = ssl.create_default_context()
31
+ if ssl_provider == "self_signed":
32
+ ctx.check_hostname = False
33
+ ctx.verify_mode = ssl.CERT_NONE
34
+ return ctx
35
+
36
+
37
+ def basic_auth_header(admin_user: str, admin_password: str) -> str:
38
+ credentials = base64.b64encode(
39
+ f"{admin_user}:{admin_password}".encode()
40
+ ).decode()
41
+ return f"Basic {credentials}"
42
+
43
+
44
+ def _read_body(
45
+ url: str,
46
+ *,
47
+ method: str,
48
+ auth_header: str,
49
+ ssl_ctx: ssl.SSLContext,
50
+ timeout: int,
51
+ data: bytes | None = None,
52
+ content_type: str | None = None,
53
+ ) -> str:
54
+ """Send one authenticated request and return the decoded body.
55
+
56
+ Owns the plumbing every caller in this module shares — the header
57
+ set, the context manager, the decode — and nothing else. HTTPError
58
+ and URLError propagate untranslated: the three callers each map them
59
+ differently (raw for `authenticated_json_request`, typed exceptions
60
+ with 404/502 handling for `AllocatorAPI`, typed with 409/400 for
61
+ `RegistrationClient`), and that mapping is the part that must stay
62
+ per-caller.
63
+ """
64
+ req = Request(url, data=data, method=method)
65
+ req.add_header("User-Agent", USER_AGENT)
66
+ req.add_header("Authorization", auth_header)
67
+ req.add_header("Accept", "application/json")
68
+ if content_type:
69
+ req.add_header("Content-Type", content_type)
70
+
71
+ # S310: URL scheme is operator-supplied by design (allocator base URL
72
+ # from the CLI config); skipping the scheme allowlist check.
73
+ with urlopen(req, timeout=timeout, context=ssl_ctx) as resp: # noqa: S310
74
+ return resp.read().decode()
75
+
76
+
77
+ def authenticated_json_request(
78
+ url: str,
79
+ admin_user: str,
80
+ admin_password: str,
81
+ *,
82
+ ssl_provider: str = "none",
83
+ timeout: int = 60,
84
+ ) -> dict:
85
+ """GET `url` with HTTP Basic auth and return the parsed JSON body.
86
+
87
+ Raises the underlying HTTPError/URLError/json.JSONDecodeError
88
+ unwrapped — stats/logs/export_metrics each need different
89
+ presentation for the same failures (a red console message here, a
90
+ TUI error dict there), so this only owns the request plumbing, not
91
+ the error translation.
92
+ """
93
+ raw = _read_body(
94
+ url,
95
+ method="GET",
96
+ auth_header=basic_auth_header(admin_user, admin_password),
97
+ ssl_ctx=ssl_context(ssl_provider),
98
+ timeout=timeout,
99
+ )
100
+ return json.loads(raw)
101
+
102
+ # /destroy and /api/launch are async (they return a job id immediately;
103
+ # the allocator runs OpenTofu on a background thread). Individual
104
+ # requests (submit or poll) are fast, so a short per-request timeout
105
+ # replaces the old single-request timeouts of up to 1800s;
106
+ # _POLL_TIMEOUT_SECONDS is the new overall ceiling for how long we'll
107
+ # keep polling before giving up (matches the old ceiling).
108
+ _REQUEST_TIMEOUT_SECONDS = 30
109
+ _POLL_INTERVAL_SECONDS = 2.0
110
+ _POLL_TIMEOUT_SECONDS = 1800
111
+
112
+ # The exact message providers/aws.py's destroy_hosts raises (wrapped by
113
+ # main.py's /destroy route into a failed operation with this same text)
114
+ # when no client VMs were ever launched. Matched here so destroy_vms()
115
+ # can keep raising AllocatorNotFoundError for this one case, preserving
116
+ # the contract deploy.py's callers already handle (non-fatal — continue
117
+ # tearing down the allocator).
118
+ _NO_VMS_LAUNCHED_ERROR = (
119
+ "tfvars does not exist — no client VMs were launched"
120
+ )
121
+
122
+
123
+ class AllocatorError(Exception):
124
+ """Base exception for allocator API errors."""
125
+
126
+
127
+ class AllocatorAuthError(AllocatorError):
128
+ """401 Unauthorized."""
129
+
130
+
131
+ class AllocatorNotFoundError(AllocatorError):
132
+ """404 Not Found (no client VMs launched)."""
133
+
134
+
135
+ class AllocatorUnavailableError(AllocatorError):
136
+ """502 Bad Gateway or connection failure."""
137
+
138
+
139
+ class AllocatorOperationTimeout(AllocatorError):
140
+ """An operation did not reach a terminal status within the poll deadline."""
141
+
142
+
143
+ class AllocatorAPI:
144
+ """HTTP client for the allocator service."""
145
+
146
+ def __init__(
147
+ self,
148
+ base_url: str,
149
+ admin_user: str,
150
+ admin_password: str,
151
+ ssl_provider: str = "none",
152
+ ) -> None:
153
+ self.base_url = base_url.rstrip("/")
154
+ self._auth_header = basic_auth_header(admin_user, admin_password)
155
+ self._ssl_ctx = ssl_context(ssl_provider)
156
+
157
+ def destroy_vms(
158
+ self,
159
+ on_progress: Callable[[int | None, int | None], None] | None = None,
160
+ ) -> dict | None:
161
+ """POST /destroy to tear down client VMs, then poll until the
162
+ job reaches a terminal status.
163
+
164
+ Returns {"status": "success", "output": <str>}. Raises
165
+ AllocatorNotFoundError if no client VMs were ever launched
166
+ (same as the old synchronous contract), or
167
+ AllocatorError/AllocatorOperationTimeout otherwise.
168
+ """
169
+ return self._submit_and_poll(
170
+ "POST", "/destroy", on_progress=on_progress,
171
+ )
172
+
173
+ def launch_vms(
174
+ self,
175
+ num_vms: int,
176
+ on_progress: Callable[[int | None, int | None], None] | None = None,
177
+ ) -> dict | None:
178
+ """POST /api/launch to provision num_vms new client VMs, then
179
+ poll until the job reaches a terminal status.
180
+
181
+ Returns {"status": "success", "output": <str>} on success.
182
+ """
183
+ data = urlencode({"num_vms": str(num_vms)}).encode()
184
+ return self._submit_and_poll(
185
+ "POST",
186
+ "/api/launch",
187
+ data,
188
+ content_type="application/x-www-form-urlencoded",
189
+ on_progress=on_progress,
190
+ )
191
+
192
+ def _submit_and_poll(
193
+ self,
194
+ method: str,
195
+ path: str,
196
+ data: bytes | None = b"",
197
+ *,
198
+ content_type: str | None = None,
199
+ on_progress: Callable[[int | None, int | None], None] | None = None,
200
+ ) -> dict | None:
201
+ """Submit an async apply/destroy job and poll GET
202
+ /api/operations/<id> until it reaches a terminal status.
203
+
204
+ If on_progress is given, it's called once per poll tick with
205
+ (resources_completed, resources_total) read off the operation
206
+ response. Both are None if the allocator predates progress
207
+ reporting (the keys are simply absent from its response) —
208
+ callers must handle that, not assume real numbers."""
209
+ submitted = self._request(
210
+ method, path, data, content_type=content_type
211
+ )
212
+ if not submitted or "job_id" not in submitted:
213
+ raise AllocatorError(
214
+ f"Unexpected response from {path}: {submitted}"
215
+ )
216
+ job_id = submitted["job_id"]
217
+
218
+ deadline = time.monotonic() + _POLL_TIMEOUT_SECONDS
219
+ while True:
220
+ try:
221
+ op = self._request("GET", f"/api/operations/{job_id}")
222
+ except AllocatorUnavailableError:
223
+ # Transient network blip while polling — the job keeps
224
+ # running server-side regardless of whether our poll
225
+ # request succeeds, so retry rather than abort an
226
+ # operation that may still finish successfully.
227
+ op = None
228
+
229
+ if op is not None:
230
+ if on_progress:
231
+ on_progress(
232
+ op.get("resources_completed"),
233
+ op.get("resources_total"),
234
+ )
235
+ status = op.get("status")
236
+ if status == "succeeded":
237
+ return {
238
+ "status": "success",
239
+ "output": op.get("output") or "",
240
+ }
241
+ if status in ("failed", "interrupted"):
242
+ error_text = (
243
+ op.get("error") or f"Operation #{job_id} {status}"
244
+ )
245
+ if _NO_VMS_LAUNCHED_ERROR in error_text:
246
+ raise AllocatorNotFoundError(error_text)
247
+ raise AllocatorError(error_text)
248
+
249
+ if time.monotonic() > deadline:
250
+ raise AllocatorOperationTimeout(
251
+ f"Operation #{job_id} did not finish within "
252
+ f"{_POLL_TIMEOUT_SECONDS}s"
253
+ )
254
+ time.sleep(_POLL_INTERVAL_SECONDS)
255
+
256
+ def _request(
257
+ self,
258
+ method: str,
259
+ path: str,
260
+ data: bytes | None = None,
261
+ *,
262
+ content_type: str | None = None,
263
+ ) -> dict | None:
264
+ """Send an HTTP request to the allocator."""
265
+ try:
266
+ raw = _read_body(
267
+ f"{self.base_url}{path}",
268
+ method=method,
269
+ auth_header=self._auth_header,
270
+ ssl_ctx=self._ssl_ctx,
271
+ timeout=_REQUEST_TIMEOUT_SECONDS,
272
+ data=data,
273
+ content_type=content_type,
274
+ )
275
+ except HTTPError as e:
276
+ self._handle_http_error(e)
277
+ except URLError as e:
278
+ raise AllocatorUnavailableError(str(e.reason)) from e
279
+
280
+ try:
281
+ body = json.loads(raw)
282
+ except (json.JSONDecodeError, ValueError):
283
+ return None
284
+
285
+ if isinstance(body, dict) and body.get("status") == "error":
286
+ raise AllocatorError(
287
+ body.get("error", "unknown error")
288
+ )
289
+
290
+ return body
291
+
292
+ def _handle_http_error(self, e: HTTPError) -> NoReturn:
293
+ """Translate HTTPError to typed exception."""
294
+ if e.code == 401:
295
+ raise AllocatorAuthError("Authentication failed") from e
296
+ if e.code == 404:
297
+ raise AllocatorNotFoundError(
298
+ "No client VMs found"
299
+ ) from e
300
+ if e.code == 502:
301
+ raise AllocatorUnavailableError(
302
+ "Allocator is unhealthy (502)"
303
+ ) from e
304
+
305
+ try:
306
+ body = json.loads(e.read().decode())
307
+ msg = body.get("error", str(e))
308
+ except (json.JSONDecodeError, UnicodeDecodeError):
309
+ msg = str(e)
310
+
311
+ raise AllocatorError(
312
+ f"HTTP {e.code}: {msg}"
313
+ ) from e
314
+
315
+
316
+ class AllocatorConflictError(AllocatorError):
317
+ """409 Conflict (already-registered machine_identity)."""
318
+
319
+
320
+ class RegistrationClient:
321
+ """HTTP client for the registration endpoint (Bearer auth, no admin creds).
322
+
323
+ Used by `lablink client register` on the BYO box. Unlike AllocatorAPI which
324
+ requires admin Basic auth, RegistrationClient authenticates with the
325
+ bootstrap register_token only — which is what a BYO box has at
326
+ onboarding time.
327
+ """
328
+
329
+ def __init__(
330
+ self,
331
+ base_url: str,
332
+ register_token: str,
333
+ *,
334
+ ssl_provider: str = "none",
335
+ ) -> None:
336
+ self.base_url = base_url.rstrip("/")
337
+ self._auth_header = f"Bearer {register_token}"
338
+ self._ssl_ctx = ssl_context(ssl_provider)
339
+
340
+ def register(
341
+ self,
342
+ *,
343
+ hostname: str,
344
+ machine_identity: str,
345
+ gpu_present: bool,
346
+ gpu_model: str | None,
347
+ lan_ip: str | None = None,
348
+ overlay_hostname: str | None = None,
349
+ reverse_tunnel: bool = False,
350
+ ) -> dict:
351
+ """POST /api/v1/clients/register; return parsed JSON.
352
+
353
+ Exactly one of ``lan_ip`` (real BYO box, LAN-direct connectivity),
354
+ ``overlay_hostname`` (mesh-overlay connectivity, e.g. a
355
+ Run:AI-hosted workload), or ``reverse_tunnel=True`` (reverse-tunnel
356
+ connectivity — the client dials out, so it has no address for the
357
+ allocator to reach) is expected — enforced by the caller
358
+ (``run_register``), not here.
359
+
360
+ Raises AllocatorAuthError on 401, AllocatorConflictError on 409,
361
+ AllocatorUnavailableError on connection failure, AllocatorError
362
+ on other HTTP error codes / malformed response.
363
+ """
364
+ if reverse_tunnel:
365
+ # The sentinel the allocator validates against its configured
366
+ # connectivity. No endpoint_url: a tunnel client has no address
367
+ # the allocator can dial -- the tunnel is the path.
368
+ provider_metadata = {"reverse_tunnel": True}
369
+ endpoint_url = None
370
+ elif overlay_hostname is not None:
371
+ provider_metadata = {"overlay_hostname": overlay_hostname}
372
+ endpoint_url = None
373
+ else:
374
+ provider_metadata = {"lan_ip": lan_ip}
375
+ endpoint_url = f"http://{lan_ip}:7070"
376
+ body = {
377
+ "hostname": hostname,
378
+ "machine_identity": machine_identity,
379
+ "provider": "manual",
380
+ "endpoint_url": endpoint_url,
381
+ "provider_metadata": provider_metadata,
382
+ "gpu_present": gpu_present,
383
+ "gpu_model": gpu_model,
384
+ }
385
+ return self._post("/api/v1/clients/register", body)
386
+
387
+ def _post(self, path: str, body: dict) -> dict:
388
+ try:
389
+ raw = _read_body(
390
+ f"{self.base_url}{path}",
391
+ method="POST",
392
+ auth_header=self._auth_header,
393
+ ssl_ctx=self._ssl_ctx,
394
+ timeout=60,
395
+ data=json.dumps(body).encode(),
396
+ content_type="application/json",
397
+ )
398
+ except HTTPError as e:
399
+ self._handle_http_error(e)
400
+ except URLError as e:
401
+ raise AllocatorUnavailableError(str(e.reason)) from e
402
+
403
+ try:
404
+ return json.loads(raw)
405
+ except (json.JSONDecodeError, ValueError) as e:
406
+ raise AllocatorError("Allocator returned non-JSON response") from e
407
+
408
+ def _handle_http_error(self, e: HTTPError) -> NoReturn:
409
+ try:
410
+ err_body = json.loads(e.read().decode())
411
+ msg = err_body.get("error", str(e))
412
+ except (json.JSONDecodeError, UnicodeDecodeError):
413
+ msg = str(e)
414
+
415
+ if e.code == 401:
416
+ raise AllocatorAuthError(
417
+ "Registration rejected — register_token may be stale "
418
+ "(allocator restarts mint a new one) or wrong."
419
+ ) from e
420
+ if e.code == 409:
421
+ raise AllocatorConflictError(
422
+ "Already registered with this machine_identity. "
423
+ "Re-register with --force or pass a different "
424
+ "--machine-identity."
425
+ ) from e
426
+ if e.code == 400:
427
+ raise AllocatorError(f"Bad request: {msg}") from e
428
+ raise AllocatorError(f"HTTP {e.code}: {msg}") from e