SPB Git forge

spb/llm-api

Public
0commits 0branches 0releases
0 Bsize
maindefault branch
—last push
5.8 KB · 150 lines python
Raw Blame History
1"""Runtime adapter interface. Every runtime runs a model in a dedicated local worker process that2speaks the OpenAI HTTP API on 127.0.0.1:<port>. The manager talks to workers uniformly."""34from __future__ import annotations56import asyncio7import os8import signal9import subprocess10import time11from dataclasses import dataclass, field12from pathlib import Path13from typing import Any1415import httpx16import psutil1718from ..hardware import process_footprint_bytes19from ..models.estimator import MemoryEstimate, estimate202122@dataclass23class WorkerHandle:24    model_id: str25    runtime: str26    port: int27    process: subprocess.Popen28    log_path: Path29    started_at: float = field(default_factory=time.time)30    context: int = 031    extra: dict[str, Any] = field(default_factory=dict)3233    @property34    def base_url(self) -> str:35        return f"http://127.0.0.1:{self.port}"3637    @property38    def pid(self) -> int:39        return self.process.pid4041    def alive(self) -> bool:42        return self.process.poll() is None4344    def memory_bytes(self) -> int:45        if not self.alive():46            return 047        return process_footprint_bytes(self.process.pid)484950class RuntimeAdapter:51    name = "base"5253    def __init__(self, settings, log_dir: Path):54        self.settings = settings55        self.log_dir = log_dir5657    # ---- to implement ------------------------------------------------------58    def build_command(self, model: dict, port: int, context: int) -> list[str]:59        raise NotImplementedError6061    async def is_ready(self, handle: WorkerHandle, client: httpx.AsyncClient) -> tuple[str, str | None]:62        """Return (status, error) where status in loading|ready|error."""63        raise NotImplementedError6465    def available(self) -> bool:66        raise NotImplementedError6768    # ---- shared -------------------------------------------------------------69    def estimate_memory(self, model: dict, context: int) -> MemoryEstimate:70        return estimate(model["weights_bytes"], self.name, model.get("kv_bytes_per_token") or 0, context,71                        bool(model.get("vision")))7273    def spawn(self, model: dict, port: int, context: int) -> WorkerHandle:74        cmd = self.build_command(model, port, context)75        self.log_dir.mkdir(parents=True, exist_ok=True)76        log_path = self.log_dir / f"worker-{model['id']}.log"77        env = dict(os.environ)78        env.setdefault("PYTHONUNBUFFERED", "1")79        env.setdefault("TOKENIZERS_PARALLELISM", "false")80        env.setdefault("HF_HUB_OFFLINE", "1")81        logf = open(log_path, "ab")82        logf.write(f"\n=== {time.strftime('%Y-%m-%d %H:%M:%S')} spawn: {' '.join(cmd)}\n".encode())83        proc = subprocess.Popen(cmd, stdout=logf, stderr=subprocess.STDOUT, stdin=subprocess.DEVNULL, env=env,84                                start_new_session=True, cwd=str(Path(model["path"]).parent))85        return WorkerHandle(model_id=model["id"], runtime=self.name, port=port, process=proc, log_path=log_path,86                            context=context)8788    async def warmup(self, handle: WorkerHandle, client: httpx.AsyncClient, model: dict) -> dict:89        """Tiny inference to prove the model works. Returns timings."""90        t0 = time.time()91        if model.get("embedding"):92            r = await client.post(f"{handle.base_url}/v1/embeddings", json={"model": model["id"], "input": "warm up"},93                                  timeout=600)94            r.raise_for_status()95            data = r.json()96            return {"ttft_ms": round((time.time() - t0) * 1000, 1), "dims": len(data["data"][0]["embedding"])}97        if model.get("reranker"):98            r = await client.post(f"{handle.base_url}/v1/rerank",99                                  json={"model": model["id"], "query": "warm", "documents": ["warm up"]}, timeout=600)100            r.raise_for_status()101            return {"ttft_ms": round((time.time() - t0) * 1000, 1)}102        r = await client.post(f"{handle.base_url}/v1/chat/completions",103                              json={"model": model["id"], "messages": [{"role": "user", "content": "Say OK."}],104                                    "max_tokens": 4, "temperature": 0.0, "stream": False}, timeout=600)105        r.raise_for_status()106        data = r.json()107        text = (data.get("choices") or [{}])[0].get("message", {}).get("content")108        timings = data.get("timings") or {}109        ttft = timings.get("ttft_ms") or timings.get("prompt_ms") or round((time.time() - t0) * 1000, 1)110        return {"ttft_ms": ttft, "text": text, "total_ms": round((time.time() - t0) * 1000, 1)}111112    async def stop(self, handle: WorkerHandle, client: httpx.AsyncClient | None = None, grace: float = 8.0) -> None:113        proc = handle.process114        if proc.poll() is not None:115            return116        # 1. polite shutdown117        if client is not None:118            try:119                await client.post(f"{handle.base_url}/shutdown", timeout=2)120            except Exception:121                pass122        # 2. SIGTERM the whole session123        try:124            os.killpg(proc.pid, signal.SIGTERM)125        except Exception:126            try:127                proc.terminate()128            except Exception:129                pass130        t0 = time.time()131        while proc.poll() is None and time.time() - t0 < grace:132            await asyncio.sleep(0.1)133        if proc.poll() is None:134            try:135                os.killpg(proc.pid, signal.SIGKILL)136            except Exception:137                try:138                    proc.kill()139                except Exception:140                    pass141            t0 = time.time()142            while proc.poll() is None and time.time() - t0 < 5:143                await asyncio.sleep(0.1)144        # kill stray children (llama-server forks none, but be safe)145        try:146            for c in psutil.Process(proc.pid).children(recursive=True):147                c.kill()148        except psutil.Error:149            pass150