Spaces:
Paused
Paused
| """ | |
| Správa vLLM inference enginu jako subprocess. | |
| Proč subprocess a ne in-process vllm.LLM: | |
| - `vllm serve` = oficiální OpenAI-kompatibilní server s nativním tool-callingem | |
| (--enable-auto-tool-choice) a reasoning parsery => odpadá ruční parsování | |
| <tool_call> tagů. | |
| - Model lze VYMĚNIT ZA BĚHU: stop subprocess -> start s novým modelem. | |
| Aplikace (UI, API, /health) běží nepřetržitě, Space se nerestartuje. | |
| - Pád enginu neshodí aplikaci; watchdog ho automaticky nahodí. | |
| Testovatelnost: binárku lze podvrhnout přes env VLLM_BINARY (testy používají | |
| mock server), port přes ENGINE_PORT. | |
| """ | |
| from __future__ import annotations | |
| import logging | |
| import os | |
| import shlex | |
| import shutil | |
| import signal | |
| import subprocess | |
| import threading | |
| import time | |
| import uuid | |
| from pathlib import Path | |
| import httpx | |
| from presets import MODEL_VOLUME_ROOT, infer_parsers | |
| from settings import Settings | |
| logger = logging.getLogger("codeagent.engine") | |
| ENGINE_HOST = "127.0.0.1" | |
| ENGINE_PORT = int(os.environ.get("ENGINE_PORT", "8001")) | |
| ENGINE_LOG = Path(os.environ.get("ENGINE_LOG", "/tmp/vllm-engine.log")) | |
| # Interní klíč mezi aplikací a lokálním vLLM serverem. Bez explicitního env | |
| # se generuje náhodně při startu aplikace (vLLM ho loguje v non-default args, | |
| # fixní default by byl zbytečně předvídatelný). | |
| ENGINE_API_KEY = os.environ.get("ENGINE_API_KEY") or f"ca-{uuid.uuid4().hex}" | |
| SERVED_MODEL_NAME = "code-agent-llm" | |
| # Velké modely se na 4 GPU načítají dlouho (a mountované váhy jdou po síti). | |
| STARTUP_TIMEOUT = int(os.environ.get("ENGINE_STARTUP_TIMEOUT", "2700")) | |
| STOP_TIMEOUT = int(os.environ.get("ENGINE_STOP_TIMEOUT", "60")) | |
| MAX_AUTO_RESTARTS = int(os.environ.get("ENGINE_MAX_AUTO_RESTARTS", "3")) | |
| LOCAL_MODEL_DIR = Path(os.environ.get("LOCAL_MODEL_DIR", "/app/models")) | |
| # RW storage bucket bývá mountovaný na /data — pokud existuje, stahujeme | |
| # váhy tam (přežijí restart Space a neomezuje je ~50GB ephemeral disk). | |
| BUCKET_DATA_DIR = "/data" | |
| DEFAULT_DOWNLOAD_DIR = "/app/cache/models" | |
| def resolve_download_dir(s: Settings) -> Path: | |
| """Adresář pro stahované váhy (vLLM --download-dir). | |
| Priorita: settings.download_dir (runtime, může mířit na bucket mount) | |
| > /data/models pokud je /data zapisovatelné (RW bucket) > ephemeral cache. | |
| """ | |
| if s.download_dir: | |
| return Path(s.download_dir) | |
| data = Path(BUCKET_DATA_DIR) | |
| try: | |
| if data.is_dir() and os.access(data, os.W_OK): | |
| return data / "models" | |
| except OSError: | |
| pass | |
| return Path(DEFAULT_DOWNLOAD_DIR) | |
| def resolve_model(spec: str) -> tuple[str, str]: | |
| """Přeloží spec na (cesta_nebo_repo_id, zdroj). | |
| Pořadí: absolutní cesta > volume mount /repos/<repo_id> > /app/models > | |
| HF repo id (stáhne se při startu enginu do DOWNLOAD_DIR). | |
| """ | |
| spec = spec.strip() | |
| if spec.startswith("/"): | |
| return spec, "path" | |
| volume = Path(MODEL_VOLUME_ROOT) / spec | |
| if volume.is_dir(): | |
| return str(volume), "volume" | |
| local = LOCAL_MODEL_DIR / spec.replace("/", "--") | |
| if local.is_dir(): | |
| return str(local), "local" | |
| return spec, "hub" | |
| def available_gpu_count() -> int: | |
| """Počet GPU, které engine subprocess reálně uvidí. | |
| Pořadí: CUDA_VISIBLE_DEVICES (dědí se do subprocessu) > torch > | |
| nvidia-smi. 0 = nezjištěno/žádné (např. lokální vývoj bez GPU). | |
| """ | |
| cvd = os.environ.get("CUDA_VISIBLE_DEVICES") | |
| if cvd is not None: | |
| return len([x for x in cvd.split(",") if x.strip() != ""]) | |
| try: | |
| import torch | |
| if torch.cuda.is_available(): | |
| return torch.cuda.device_count() | |
| except Exception: # torch chybí nebo bez CUDA — zkus nvidia-smi | |
| pass | |
| try: | |
| out = subprocess.run(["nvidia-smi", "-L"], capture_output=True, | |
| text=True, timeout=10) | |
| if out.returncode == 0: | |
| return len([line for line in out.stdout.splitlines() | |
| if line.strip().startswith("GPU ")]) | |
| except (OSError, subprocess.TimeoutExpired): | |
| pass | |
| return 0 | |
| def build_child_env(base: dict | None = None) -> dict: | |
| """Prostředí pro vLLM subprocess. | |
| - HF_XET_HIGH_PERFORMANCE nahrazuje deprecated HF_HUB_ENABLE_HF_TRANSFER | |
| (hub 1.x už hf_transfer nepoužívá — FutureWarning v každém workeru). | |
| - Při namountovaném RW bucketu (/data) se tam persistují cache: | |
| VLLM_CACHE_ROOT (torch.compile artefakty — bez toho každý start Space | |
| platí kompilaci znovu) a HF_HOME (tokenizer/config/remote-code moduly). | |
| - VLLM_BUILD_* / VLLM_IMAGE_TAG jsou metadata oficiálního image, vLLM je | |
| hlásí jako neznámé proměnné — odstraňují se (jen šum v logu). | |
| """ | |
| env = dict(base if base is not None else os.environ) | |
| for key in list(env): | |
| if key.startswith("VLLM_BUILD_") or key == "VLLM_IMAGE_TAG": | |
| env.pop(key) | |
| env.pop("HF_HUB_ENABLE_HF_TRANSFER", None) | |
| env.setdefault("HF_XET_HIGH_PERFORMANCE", "1") | |
| env.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") | |
| data = Path(BUCKET_DATA_DIR) | |
| try: | |
| if data.is_dir() and os.access(data, os.W_OK): | |
| env.setdefault("VLLM_CACHE_ROOT", str(data / "cache" / "vllm")) | |
| env.setdefault("HF_HOME", str(data / "cache" / "hf")) | |
| except OSError: | |
| pass | |
| return env | |
| class EngineState: | |
| STOPPED = "stopped" | |
| STARTING = "starting" | |
| READY = "ready" | |
| ERROR = "error" | |
| class VLLMEngine: | |
| """Životní cyklus jednoho `vllm serve` procesu + health/watchdog.""" | |
| def __init__(self, host: str = ENGINE_HOST, port: int = ENGINE_PORT): | |
| self.host = host | |
| self.port = port | |
| self.base_url = f"http://{host}:{port}" | |
| self._lock = threading.RLock() | |
| self._proc: subprocess.Popen | None = None | |
| self._log_file = None | |
| self.state = EngineState.STOPPED | |
| self.last_error: str | None = None | |
| self.current_model: str | None = None | |
| self.model_source: str | None = None | |
| self.active_download_dir: str | None = None | |
| self.detected_gpus: int | None = None | |
| self.effective_tp: int | None = None | |
| self.tp_clamped_from: int | None = None | |
| self.started_at: float | None = None | |
| self.ready_at: float | None = None | |
| self._generation = 0 # roste s každým start/stop (ruší staré waitery) | |
| self._auto_restarts = 0 | |
| self._watchdog_started = False | |
| self._pending_settings: Settings | None = None | |
| # ------------------------------------------------------------- příkaz | |
| def build_command(self, s: Settings) -> list[str]: | |
| model_path, source = resolve_model(s.model) | |
| self.model_source = source | |
| # TP nesmí překročit počet reálně dostupných GPU — jinak vLLM spadne | |
| # hned při startu (ValidationError). Typicky po změně HW tieru Space. | |
| tp = s.tensor_parallel_size | |
| gpus = available_gpu_count() | |
| self.detected_gpus = gpus or None | |
| self.tp_clamped_from = None | |
| if gpus and tp > gpus: | |
| logger.warning("tensor_parallel_size=%s > dostupných GPU=%s — " | |
| "snižuji TP na %s (zkontroluj hardware Space!)", | |
| tp, gpus, gpus) | |
| self.tp_clamped_from = tp | |
| tp = gpus | |
| self.effective_tp = tp | |
| binary = os.environ.get("VLLM_BINARY", "vllm") | |
| cmd = [binary, "serve", model_path, | |
| "--host", self.host, | |
| "--port", str(self.port), | |
| "--api-key", ENGINE_API_KEY, | |
| "--served-model-name", SERVED_MODEL_NAME, | |
| "--tensor-parallel-size", str(tp), | |
| "--gpu-memory-utilization", str(s.gpu_memory_utilization), | |
| "--max-model-len", str(s.max_model_len), | |
| "--dtype", s.dtype, | |
| "--trust-remote-code", | |
| ] | |
| if source == "hub" and s.model_revision.strip(): | |
| # Pin na konkrétní revizi: reprodukovatelné buildy a žádné tiché | |
| # aktualizace remote-code souborů (trust_remote_code=True). | |
| cmd += ["--revision", s.model_revision.strip()] | |
| if source == "hub": | |
| download_dir = resolve_download_dir(s) | |
| try: | |
| download_dir.mkdir(parents=True, exist_ok=True) | |
| cmd += ["--download-dir", str(download_dir)] | |
| self.active_download_dir = str(download_dir) | |
| except OSError as e: | |
| logger.warning("Download dir %s nedostupný (%s) — použije se " | |
| "výchozí HF cache.", download_dir, e) | |
| self.active_download_dir = None | |
| else: | |
| self.active_download_dir = None | |
| if s.quantization not in ("auto", "", "none"): | |
| cmd += ["--quantization", s.quantization] | |
| if s.kv_cache_dtype != "auto": | |
| cmd += ["--kv-cache-dtype", s.kv_cache_dtype] | |
| if s.enforce_eager: | |
| cmd += ["--enforce-eager"] | |
| if not s.enable_prefix_caching: | |
| cmd += ["--no-enable-prefix-caching"] | |
| if s.max_num_seqs > 0: | |
| cmd += ["--max-num-seqs", str(s.max_num_seqs)] | |
| tool_parser, reasoning_parser = s.tool_call_parser, s.reasoning_parser | |
| if tool_parser == "auto" or reasoning_parser == "auto": | |
| inferred_tool, inferred_reasoning = infer_parsers(s.model) | |
| if tool_parser == "auto": | |
| tool_parser = inferred_tool | |
| if reasoning_parser == "auto": | |
| reasoning_parser = inferred_reasoning | |
| if tool_parser: | |
| cmd += ["--enable-auto-tool-choice", "--tool-call-parser", tool_parser] | |
| if reasoning_parser: | |
| cmd += ["--reasoning-parser", reasoning_parser] | |
| if s.engine_extra_args.strip(): | |
| cmd += shlex.split(s.engine_extra_args) | |
| return cmd | |
| # ------------------------------------------------------------- start/stop | |
| def start(self, s: Settings, block: bool = False): | |
| """Spustí engine na pozadí. Při block=True čeká na ready/error.""" | |
| with self._lock: | |
| if self.state == EngineState.STARTING: | |
| logger.info("Engine už startuje — požadavek ignorován") | |
| return | |
| self._terminate_locked() | |
| self._generation += 1 | |
| generation = self._generation | |
| self.state = EngineState.STARTING | |
| self.last_error = None | |
| self.current_model = s.model | |
| self.started_at = time.time() | |
| self.ready_at = None | |
| cmd = self.build_command(s) | |
| env = build_child_env() | |
| logger.info("Engine start: GPU=%s, TP=%s%s, CUDA_VISIBLE_DEVICES=%r", | |
| self.detected_gpus, self.effective_tp, | |
| f" (sníženo z {self.tp_clamped_from})" if self.tp_clamped_from else "", | |
| os.environ.get("CUDA_VISIBLE_DEVICES")) | |
| logger.info("Startuji engine (gen %s): %s", generation, " ".join(cmd)) | |
| try: | |
| ENGINE_LOG.parent.mkdir(parents=True, exist_ok=True) | |
| self._log_file = open(ENGINE_LOG, "ab", buffering=0) | |
| self._log_file.write( | |
| f"\n===== engine start gen={generation} model={s.model} " | |
| f"{time.strftime('%Y-%m-%d %H:%M:%S')} =====\n".encode()) | |
| self._proc = subprocess.Popen( | |
| cmd, stdout=self._log_file, stderr=subprocess.STDOUT, | |
| env=env, start_new_session=True) | |
| except OSError as e: | |
| self.state = EngineState.ERROR | |
| self.last_error = f"Spuštění selhalo: {e}" | |
| logger.error(self.last_error) | |
| return | |
| self._start_watchdog() | |
| waiter = threading.Thread(target=self._wait_until_ready, | |
| args=(generation,), daemon=True, | |
| name=f"engine-waiter-{generation}") | |
| waiter.start() | |
| if block: | |
| waiter.join() | |
| def _wait_until_ready(self, generation: int): | |
| deadline = time.time() + STARTUP_TIMEOUT | |
| while time.time() < deadline: | |
| with self._lock: | |
| if generation != self._generation: | |
| return # mezitím proběhl nový start/stop | |
| proc = self._proc | |
| if proc is None or proc.poll() is not None: | |
| with self._lock: | |
| if generation == self._generation: | |
| code = proc.poll() if proc else "?" | |
| self.state = EngineState.ERROR | |
| self.last_error = (f"Engine proces skončil (exit {code}). " | |
| f"Detail: {self.log_tail(15)}") | |
| logger.error("Engine zemřel při startu (gen %s)", generation) | |
| return | |
| if self._health_ok(): | |
| with self._lock: | |
| if generation == self._generation: | |
| self.state = EngineState.READY | |
| self.ready_at = time.time() | |
| self._auto_restarts = 0 | |
| logger.info("Engine READY (gen %s, model %s, %.0fs)", | |
| generation, self.current_model, | |
| time.time() - (self.started_at or time.time())) | |
| return | |
| time.sleep(3) | |
| with self._lock: | |
| if generation == self._generation: | |
| self.state = EngineState.ERROR | |
| self.last_error = f"Timeout startu ({STARTUP_TIMEOUT}s)" | |
| logger.error("Engine start timeout (gen %s)", generation) | |
| def _health_ok(self) -> bool: | |
| try: | |
| r = httpx.get(f"{self.base_url}/health", timeout=5) | |
| return r.status_code == 200 | |
| except httpx.HTTPError: | |
| return False | |
| def stop(self): | |
| with self._lock: | |
| self._generation += 1 | |
| self._terminate_locked() | |
| self.state = EngineState.STOPPED | |
| self.current_model = None | |
| def _terminate_locked(self): | |
| proc, self._proc = self._proc, None | |
| log_file, self._log_file = self._log_file, None | |
| if proc and proc.poll() is None: | |
| logger.info("Ukončuji engine pid=%s", proc.pid) | |
| try: | |
| os.killpg(os.getpgid(proc.pid), signal.SIGTERM) | |
| except ProcessLookupError: | |
| pass | |
| try: | |
| proc.wait(timeout=STOP_TIMEOUT) | |
| except subprocess.TimeoutExpired: | |
| logger.warning("Engine nereaguje na SIGTERM, posílám SIGKILL") | |
| try: | |
| os.killpg(os.getpgid(proc.pid), signal.SIGKILL) | |
| except ProcessLookupError: | |
| pass | |
| proc.wait(timeout=10) | |
| if log_file: | |
| try: | |
| log_file.close() | |
| except OSError: | |
| pass | |
| def reload(self, s: Settings, block: bool = False): | |
| """Vymění model/konfiguraci za běhu — Space se nerestartuje.""" | |
| logger.info("Reload enginu: model=%s TP=%s quant=%s", | |
| s.model, s.tensor_parallel_size, s.quantization) | |
| self.start(s, block=block) | |
| # ------------------------------------------------------------- watchdog | |
| def _start_watchdog(self): | |
| with self._lock: | |
| if self._watchdog_started: | |
| return | |
| self._watchdog_started = True | |
| threading.Thread(target=self._watchdog_loop, daemon=True, | |
| name="engine-watchdog").start() | |
| def _watchdog_loop(self): | |
| while True: | |
| time.sleep(15) | |
| with self._lock: | |
| proc = self._proc | |
| state = self.state | |
| settings = self._pending_settings | |
| if state != EngineState.READY or proc is None: | |
| continue | |
| if proc.poll() is not None: | |
| logger.error("Engine zemřel za běhu (exit %s)", proc.poll()) | |
| with self._lock: | |
| self.state = EngineState.ERROR | |
| self.last_error = (f"Engine spadl za běhu (exit {proc.poll()}). " | |
| f"Detail: {self.log_tail(15)}") | |
| restarts = self._auto_restarts | |
| if settings is not None and restarts < MAX_AUTO_RESTARTS: | |
| with self._lock: | |
| self._auto_restarts += 1 | |
| n = self._auto_restarts | |
| logger.info("Auto-restart enginu (%s/%s)", n, MAX_AUTO_RESTARTS) | |
| self.start(settings) | |
| def remember_settings(self, s: Settings): | |
| """Uloží poslední použitá nastavení pro watchdog auto-restart.""" | |
| with self._lock: | |
| self._pending_settings = s | |
| # ------------------------------------------------------------- introspekce | |
| def is_ready(self) -> bool: | |
| return self.state == EngineState.READY | |
| def status_dict(self) -> dict: | |
| with self._lock: | |
| proc = self._proc | |
| uptime = time.time() - self.ready_at if self.ready_at else None | |
| loading = (time.time() - self.started_at | |
| if self.started_at and self.state == EngineState.STARTING | |
| else None) | |
| return { | |
| "state": self.state, | |
| "model": self.current_model, | |
| "model_source": self.model_source, | |
| "detected_gpus": self.detected_gpus, | |
| "effective_tensor_parallel": self.effective_tp, | |
| "tensor_parallel_clamped_from": self.tp_clamped_from, | |
| "download_dir": self.active_download_dir, | |
| "download_dir_free_gb": (disk_free_gb(self.active_download_dir) | |
| if self.active_download_dir else None), | |
| "pid": proc.pid if proc and proc.poll() is None else None, | |
| "uptime_seconds": round(uptime) if uptime else None, | |
| "loading_seconds": round(loading) if loading else None, | |
| "last_error": self.last_error, | |
| "base_url": self.base_url, | |
| "served_model_name": SERVED_MODEL_NAME, | |
| } | |
| def log_tail(self, lines: int = 60) -> str: | |
| try: | |
| if not ENGINE_LOG.exists(): | |
| return "(log neexistuje)" | |
| with open(ENGINE_LOG, "rb") as f: | |
| f.seek(0, os.SEEK_END) | |
| size = f.tell() | |
| f.seek(max(0, size - 65536)) | |
| data = f.read().decode("utf-8", errors="replace") | |
| return "\n".join(data.splitlines()[-lines:]) | |
| except OSError as e: | |
| return f"(chyba čtení logu: {e})" | |
| # ------------------------------------------------------------- klient | |
| def openai_client(self): | |
| from openai import OpenAI | |
| return OpenAI(base_url=f"{self.base_url}/v1", api_key=ENGINE_API_KEY, | |
| timeout=600, max_retries=1) | |
| def disk_free_gb(path: str = "/") -> float: | |
| try: | |
| return round(shutil.disk_usage(path).free / 1e9, 1) | |
| except OSError: | |
| return 0.0 | |