SPB Git forge

spb/llm-api

Public
0commits 0branches 0releases
0 Bsize
maindefault branch
—last push
2.2 KB · 59 lines python
Raw Blame History
1"""MLX runtime: spawns llm_api.worker.mlx_worker in the server's own Python environment."""23from __future__ import annotations45import sys67import httpx89from .base import RuntimeAdapter, WorkerHandle101112class MLXAdapter(RuntimeAdapter):13    name = "mlx"1415    def available(self) -> bool:16        try:17            import mlx.core  # noqa: F40118            import mlx_lm  # noqa: F40119            return True20        except Exception:21            return False2223    def build_command(self, model: dict, port: int, context: int) -> list[str]:24        py = self.settings.worker_python or sys.executable25        task = "embedding" if model.get("embedding") else "reranking" if model.get("reranker") else "text"26        cmd = [py, "-m", "llm_api.worker.mlx_worker", "--model-path", model["path"], "--model-id", model["id"],27               "--port", str(port), "--max-context", str(context),28               "--default-max-tokens", str(self.settings.default_max_tokens), "--task", task,29               "--generation-timeout", str(self.settings.generation_timeout_seconds)]30        if model.get("vision"):31            cmd.append("--vision")32        overrides = model.get("overrides") or {}33        if overrides.get("kv_bits"):34            cmd += ["--kv-bits", str(overrides["kv_bits"])]35        return cmd3637    async def warmup(self, handle: WorkerHandle, client: httpx.AsyncClient, model: dict) -> dict:38        r = await client.post(f"{handle.base_url}/warmup", timeout=600)39        r.raise_for_status()40        return r.json()4142    async def memory_gb(self, handle: WorkerHandle, client: httpx.AsyncClient) -> float | None:43        try:44            r = await client.get(f"{handle.base_url}/health", timeout=3)45            m = r.json().get("memory") or {}46            return (m.get("active_gb") or 0) + (m.get("cache_gb") or 0) or None47        except Exception:48            return None4950    async def is_ready(self, handle: WorkerHandle, client: httpx.AsyncClient) -> tuple[str, str | None]:51        try:52            r = await client.get(f"{handle.base_url}/health", timeout=3)53        except Exception:54            return "loading", None55        if r.status_code != 200:56            return "loading", None57        d = r.json()58        return d.get("status", "loading"), d.get("error")59