engine.py

"""
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

    @property
    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