gigscan/server.py

import subprocess
import threading
import time
from pathlib import Path
import os

import requests
from huggingface_hub import hf_hub_download

from gigscan import state
from gigscan.config import (
    HEALTH_URL,
    HOST,
    LLAMA_SERVER,
    MMPROJ_FILE,
    MODEL_FILE,
    MODEL_REPO,
    PORT,
)
from gigscan.logging_utils import log, timed


def stream_server_logs():
    if state.server_process is None or state.server_process.stdout is None:
        return

    for line in state.server_process.stdout:
        line = line.rstrip()
        state.SERVER_LOG_LINES.append(line)
        if len(state.SERVER_LOG_LINES) > 1000:
            del state.SERVER_LOG_LINES[:200]
        print(f"[llama-server] {line}", flush=True)


def wait_for_server(timeout_seconds=240):
    start = time.time()

    while time.time() - start < timeout_seconds:
        if state.server_process is not None and state.server_process.poll() is not None:
            log("llama-server exited while waiting for readiness")
            return False

        try:
            r = requests.get(HEALTH_URL, timeout=3)
            log(f"/health -> {r.status_code} {r.text[:200]}")
            if r.status_code == 200:
                return True
        except Exception as e:
            log(f"/health probe error: {repr(e)}")

        time.sleep(3)

    return False


def download_model_files():
    log(f"Downloading model from {MODEL_REPO}")

    model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
    mmproj_path = hf_hub_download(repo_id=MODEL_REPO, filename=MMPROJ_FILE)

    for p in [model_path, mmproj_path]:
        size_gb = Path(p).stat().st_size / (1024 ** 3)
        log(f"Downloaded {p} ({size_gb:.3f} GB)")

    return model_path, mmproj_path


def start_llama_server():
    if state.server_process is not None and state.server_process.poll() is None:
        log("llama-server already running")
        return

    log(f"Using llama-server binary: {LLAMA_SERVER}")

    paths, download_s = timed("Model download", download_model_files)
    model_path, mmproj_path = paths

    use_gpu = os.environ.get("GIGSCAN_USE_GPU", "0") == "1"

    cmd = [
        LLAMA_SERVER,
        "-m",
        model_path,
        "--mmproj",
        mmproj_path,
        "-c",
        "4096",
        "--reasoning-budget",
        "0",
        "--host",
        HOST,
        "--port",
        PORT,
    ]

    if use_gpu:
        cmd.extend(["-ngl", "999"])
        log("GPU offload enabled")
    else:
        log("GPU offload disabled; running on CPU")

    log("Starting llama-server")
    log(" ".join(cmd))

    state.server_process = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        bufsize=1,
    )

    threading.Thread(target=stream_server_logs, daemon=True).start()

    ready, ready_s = timed("Server readiness wait", wait_for_server)

    if not ready:
        raise RuntimeError("llama-server did not become ready")

    log(f"Startup complete (downloads {download_s:.2f}s, readiness {ready_s:.2f}s)")