diff --git a/stacklets/ai/hooks/on_install.py b/stacklets/ai/hooks/on_install.py index 6263a1d5..18dad130 100644 --- a/stacklets/ai/hooks/on_install.py +++ b/stacklets/ai/hooks/on_install.py @@ -108,6 +108,22 @@ def _install_omlx_formula(ctx) -> None: ctx.shell_live(f"{BREW_INSTALL} omlx --with-grammar") +def wait_for_omlx(ctx, base_url: str, key: str, timeout: float = 30) -> bool: + """Wait until oMLX lists its models; False when it does not in time. + + The request carries the key the install wrote into oMLX's settings: + oMLX answers 401 to a request without it. + """ + deadline = time.time() + timeout + while time.time() < deadline: + try: + ctx.shell(f'curl -sf --max-time 2 -H "Authorization: Bearer {key}" {base_url}/models') + return True + except RuntimeError: + time.sleep(2) + return False + + def install_omlx(ctx, state_dir: Path): section("oMLX", "MLX inference (Metal GPU)") @@ -153,14 +169,7 @@ def install_omlx(ctx, state_dir: Path): (state_dir / "omlx-managed").touch() ctx.step("Waiting for oMLX to start...") - deadline = time.time() + 30 - while time.time() < deadline: - try: - ctx.shell(f'curl -sf --max-time 2 http://localhost:{OMLX_PORT}/v1/models') - break - except RuntimeError: - time.sleep(2) - else: + if not wait_for_omlx(ctx, f"http://localhost:{OMLX_PORT}/v1", settings["auth"]["api_key"]): warn(f"oMLX not responding on port {OMLX_PORT} — it may still be starting") done(f"oMLX running on port {OMLX_PORT}") diff --git a/tests/unit/stacklets/test_ai_omlx_install.py b/tests/unit/stacklets/test_ai_omlx_install.py index 3be260d3..8d0b50a3 100644 --- a/tests/unit/stacklets/test_ai_omlx_install.py +++ b/tests/unit/stacklets/test_ai_omlx_install.py @@ -115,3 +115,36 @@ def test_a_homebrew_without_a_trust_gate_still_installs(self): ctx = FakeCtx(no_trust_subcommand=True) _install_omlx_formula(ctx) # must not raise assert any("brew install omlx" in c for c in ctx.commands) + + +class ShellCtx: + """Runs the hook's shell commands for real, as `ctx.shell` does.""" + + def shell(self, cmd): + import subprocess + result = subprocess.run(cmd, shell=True, capture_output=True, text=True) + if result.returncode != 0: + raise RuntimeError(cmd) + return result.stdout + + +class TestWaitingForOmlx: + """oMLX 0.7 answers `/v1/models` with 401 unless the request carries the key. + + The install writes that key into oMLX's settings, so its readiness + check has to send it, or it waits out its whole timeout on a server + that is up and warns that it is not. + """ + + def test_a_server_that_wants_the_key_counts_as_up(self, httpserver): + httpserver.expect_request( + "/v1/models", headers={"Authorization": "Bearer local"}, + ).respond_with_json({"data": []}) + httpserver.expect_request("/v1/models").respond_with_data("API key required", status=401) + + assert _ai_on_install.wait_for_omlx(ShellCtx(), httpserver.url_for("/v1"), "local", timeout=5) + + def test_a_server_that_refuses_the_key_is_not_up(self, httpserver): + httpserver.expect_request("/v1/models").respond_with_data("API key required", status=401) + + assert not _ai_on_install.wait_for_omlx(ShellCtx(), httpserver.url_for("/v1"), "local", timeout=1)