engine.py

import os
import shutil
import subprocess
import tarfile
import threading
import time
import urllib.request
from pathlib import Path
from huggingface_hub import hf_hub_download
import config

_runtime_lock = threading.Lock()
_server_process = None
_server_ready = False
_server_error = None
_warmup_thread = None

def build_env():
    env = os.environ.copy()
    server_dir = str(resolve_server_path().parent)
    env["LD_LIBRARY_PATH"] = f"{server_dir}:{env.get('LD_LIBRARY_PATH', '')}".rstrip(":")
    env.setdefault("GGML_CPU_BACKEND", "x64")
    return env


def resolve_server_path() -> Path:
    explicit_path = os.environ.get("LLAMA_SERVER_PATH", "").strip()
    if explicit_path:
        candidate = Path(explicit_path)
        if candidate.exists():
            return candidate

    bundled_path = Path("/usr/local/bin/llama-server")
    if bundled_path.exists():
        return bundled_path

    return config.SERVER_PATH

def ensure_runtime_files():
    config.RUNTIME_ROOT.mkdir(parents=True, exist_ok=True)
    config.MODEL_DIR.mkdir(parents=True, exist_ok=True)

    server_path = resolve_server_path()

    if not server_path.exists():
        print(f"Downloading llama.cpp binary: {config.LLAMA_CPP_URL}")
        urllib.request.urlretrieve(config.LLAMA_CPP_URL, config.ARCHIVE_PATH)
        if config.BIN_DIR.exists():
            shutil.rmtree(config.BIN_DIR)
        config.BIN_DIR.mkdir(parents=True, exist_ok=True)
        with tarfile.open(config.ARCHIVE_PATH, "r:gz") as tar:
            tar.extractall(config.BIN_DIR)
        server_path = resolve_server_path()
        os.chmod(server_path, 0o755)

    print(f"Downloading model: {config.REPO_ID}/{config.FILENAME}")
    model_path = hf_hub_download(
        repo_id=config.REPO_ID,
        filename=config.FILENAME,
        local_dir=config.MODEL_DIR,
    )
    return str(server_path), model_path

def wait_for_server(timeout_seconds=180):
    deadline = time.time() + timeout_seconds
    health_url = f"http://127.0.0.1:{config.SERVER_PORT}/health"
    while time.time() < deadline:
        try:
            with urllib.request.urlopen(health_url, timeout=5) as response:
                if response.status == 200:
                    return
        except Exception:
            pass
        if _server_process is not None and _server_process.poll() is not None:
            raise RuntimeError(f"llama-server 已提前結束")
        time.sleep(2)
    raise RuntimeError("等待 llama-server 啟動逾時")

def ensure_llama_server():
    global _server_process, _server_ready, _server_error
    
    # 先做一次快速檢查，不拿鎖
    if _server_ready and _server_process is not None and _server_process.poll() is None:
        return

    with _runtime_lock:
        # 拿鎖後再檢查一次，防止重複啟動
        if _server_ready and _server_process is not None and _server_process.poll() is None:
            return
        
        try:
            _server_ready = False
            _server_error = None
            
            server_path, model_path = ensure_runtime_files()
            command = [
                server_path, "-m", model_path, "--alias", config.MODEL_ALIAS,
                "--host", "127.0.0.1", "--port", str(config.SERVER_PORT),
                "-c", "4096", "-t", "2", "--embedding"
            ]
            log_file = open(config.RUNTIME_ROOT / "llama.log", "w")
            _server_process = subprocess.Popen(
                command, cwd=str(Path(server_path).parent), env=build_env(),
                stdout=log_file, stderr=log_file
            )
            # 這裡會阻塞，但如果是從 trigger_warmup (thread) 呼叫就沒關係
            wait_for_server()
            _server_ready = True
            print("llama-server started successfully.")
        except Exception as exc:
            _server_ready = False
            _server_error = str(exc)
            print(f"Failed to start llama-server: {exc}")
            # 不要在這裡 raise，讓狀態維持在 error


def get_server_status():
    if _server_error: return "error"
    if _server_ready and _server_process and _server_process.poll() is None: return "ok"
    return "warming"


def get_server_error():
    return _server_error


def get_server_log_tail(max_chars: int = 1200):
    log_path = config.RUNTIME_ROOT / "llama.log"
    if not log_path.exists():
        return ""
    try:
        content = log_path.read_text(encoding="utf-8", errors="replace")
        return content[-max_chars:]
    except Exception as exc:
        return f"<read_log_failed: {exc}>"

def trigger_warmup():
    global _warmup_thread
    with _runtime_lock:
        if get_server_status() == "ok": return
        if _warmup_thread and _warmup_thread.is_alive(): return
        _warmup_thread = threading.Thread(target=ensure_llama_server, daemon=True)
        _warmup_thread.start()