serve/server.py

"""GPU-oriented FastAPI server for SAGE."""

from __future__ import annotations

import logging
import os
from pathlib import Path
from typing import Any, Optional

import torch
from fastapi import FastAPI
from pydantic import BaseModel

from model.config import ModelConfig
from model.model import SageTransformer
from serve.control_plane import build_control_router, get_runtime_access_info
from train.checkpoint import load_latest_checkpoint
from train.hardware import HardwareConfig


app = FastAPI(title="SAGE Server")
_MODEL: SageTransformer | None = None
_TOKENIZER = None
_MODEL_DEVICE: torch.device | None = None
_MODEL_STATE: dict[str, Any] = {
    "model_config": None,
    "checkpoint_dir": None,
    "checkpoint_loaded": False,
    "checkpoint_step": 0,
    "tokenizer_path": None,
}
_LOGGER = logging.getLogger("uvicorn.error")


def _print_startup_banner() -> None:
    """Print the login details for the browser control UI."""
    access = get_runtime_access_info()
    local_url = (access["local_url"] or "http://127.0.0.1:8000").rstrip("/")
    public_url = access["public_url"]
    _LOGGER.info("SAGE local URL: %s/", local_url)
    if public_url:
        _LOGGER.info("SAGE public URL: %s/", public_url.rstrip("/"))
    _LOGGER.info("SAGE login password: %s", access["password"])


class GenerationRequest(BaseModel):
    """Request schema for text generation."""

    input_ids: list[int]
    max_new_tokens: int = 32


class ChatRequest(BaseModel):
    """Request schema for text generation through the tokenizer."""

    prompt: str
    max_new_tokens: int = 64


def get_generation_device() -> torch.device:
    """Pick the active inference device."""
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


def _resolve_model_config_path() -> Path:
    configured = Path(os.environ.get("SAGE_MODEL_CONFIG", "configs/model/1b.yaml"))
    return configured if configured.exists() else Path("configs/model/1b.yaml")


def _resolve_checkpoint_dir() -> Path:
    return Path(os.environ.get("SAGE_CHECKPOINT_DIR", "runs/default"))


def _resolve_tokenizer_path() -> Path:
    return Path(os.environ.get("SAGE_TOKENIZER_MODEL", "tokenizer/tokenizer.model"))


def configure_runtime_paths(model_config: str | None, checkpoint_dir: str | None, tokenizer_model: str | None) -> None:
    """Configure runtime paths via environment variables."""
    if model_config:
        os.environ["SAGE_MODEL_CONFIG"] = model_config
    if checkpoint_dir:
        os.environ["SAGE_CHECKPOINT_DIR"] = checkpoint_dir
    if tokenizer_model:
        os.environ["SAGE_TOKENIZER_MODEL"] = tokenizer_model


def get_model() -> SageTransformer:
    """Lazily create the model for server startup."""
    global _MODEL, _MODEL_DEVICE
    if _MODEL is None:
        config_path = _resolve_model_config_path()
        config = ModelConfig.from_yaml(config_path) if config_path.exists() else ModelConfig()
        _MODEL = SageTransformer(config)
        checkpoint_dir = _resolve_checkpoint_dir()
        checkpoint_step = 0
        if checkpoint_dir.exists():
            checkpoint_step = load_latest_checkpoint(_MODEL, None, None, None, str(checkpoint_dir), device="cpu")
        _MODEL_STATE.update(
            {
                "model_config": str(config_path),
                "checkpoint_dir": str(checkpoint_dir),
                "checkpoint_loaded": checkpoint_step > 0,
                "checkpoint_step": checkpoint_step,
            }
        )
        _MODEL.eval()
    device = get_generation_device()
    if _MODEL_DEVICE != device:
        _MODEL = _MODEL.to(device)
        _MODEL_DEVICE = device
    return _MODEL


def get_tokenizer():
    """Lazily load the SentencePiece tokenizer if present."""
    global _TOKENIZER
    if _TOKENIZER is None:
        tokenizer_path = _resolve_tokenizer_path()
        _MODEL_STATE["tokenizer_path"] = str(tokenizer_path)
        if not tokenizer_path.exists():
            return None
        from tokenizer.validate_tokenizer import load_processor

        _TOKENIZER = load_processor(str(tokenizer_path))
    return _TOKENIZER


def _generate_token_ids(input_ids: list[int], max_new_tokens: int) -> list[int]:
    """Run greedy decoding from input token ids."""
    model = get_model()
    device = get_generation_device()
    context_length = model.config.context_length
    generated = list(input_ids[-context_length:])
    with torch.inference_mode():
        for _ in range(max(0, int(max_new_tokens))):
            window = generated[-context_length:]
            tensor_ids = torch.tensor([window], dtype=torch.long, device=device)
            logits, _ = model(tensor_ids)
            next_token = int(torch.argmax(logits[:, -1, :], dim=-1).item())
            generated.append(next_token)
    return generated


def chat_status() -> dict[str, object]:
    """Return whether text chat is configured for the current server."""
    tokenizer = get_tokenizer()
    checkpoint_dir = _MODEL_STATE["checkpoint_dir"] or str(_resolve_checkpoint_dir())
    checkpoint_loaded = bool(_MODEL_STATE["checkpoint_loaded"])
    if not checkpoint_loaded:
        checkpoint_loaded = any(Path(checkpoint_dir).glob("ckpt_step_*.pt"))
    checkpoint_step = int(_MODEL_STATE["checkpoint_step"] or 0)
    if checkpoint_step == 0 and checkpoint_loaded:
        latest = sorted(Path(checkpoint_dir).glob("ckpt_step_*.pt"))
        if latest:
            checkpoint_step = int(latest[-1].stem.split("_")[-1])
    available = tokenizer is not None
    warning = None
    if tokenizer is None:
        warning = "Tokenizer model not found. Train or place tokenizer/tokenizer.model before using browser chat."
    elif not checkpoint_loaded:
        warning = "No checkpoint loaded. Chat will run with randomly initialized model weights until you train or configure SAGE_CHECKPOINT_DIR."
    return {
        "available": available,
        "tokenizer_path": _MODEL_STATE["tokenizer_path"],
        "checkpoint_dir": checkpoint_dir,
        "checkpoint_loaded": checkpoint_loaded,
        "checkpoint_step": checkpoint_step,
        "warning": warning,
    }


@app.get("/health")
def health() -> dict[str, object]:
    """Return basic health and hardware information."""
    hw = HardwareConfig(model_size_b=1.0, context_length=4096)
    return {"status": "ok", "hardware": hw.summary(), "chat": chat_status()}


@app.post("/generate")
def generate(request: GenerationRequest) -> dict[str, object]:
    """Generate continuation token ids from an input token list."""
    return {"tokens": _generate_token_ids(request.input_ids, request.max_new_tokens)}


@app.get("/chat/status")
def get_chat_status() -> dict[str, object]:
    """Expose browser-chat readiness."""
    return chat_status()


@app.post("/chat")
def chat(request: ChatRequest) -> dict[str, object]:
    """Generate text from a prompt using the local tokenizer."""
    tokenizer = get_tokenizer()
    if tokenizer is None:
        return {
            "success": False,
            "detail": "Tokenizer model not found. Train the tokenizer first or set SAGE_TOKENIZER_MODEL.",
            **chat_status(),
        }
    prompt = request.prompt.strip()
    if not prompt:
        return {"success": False, "detail": "Prompt cannot be empty.", **chat_status()}
    prompt_ids = list(tokenizer.encode(prompt, out_type=int))
    generated = _generate_token_ids(prompt_ids, request.max_new_tokens)
    completion_ids = generated[len(prompt_ids) :]
    return {
        "success": True,
        "prompt": prompt,
        "response": tokenizer.decode(completion_ids),
        "input_ids": prompt_ids,
        "output_ids": generated,
        "new_token_ids": completion_ids,
        **chat_status(),
    }


def _health_action(_: dict[str, object]) -> dict[str, object]:
    return health()


def _generate_action(args: dict[str, object]) -> dict[str, object]:
    input_ids_raw = args.get("input_ids", [])
    if not isinstance(input_ids_raw, list):
        input_ids_raw = [input_ids_raw] if input_ids_raw is not None else []
    input_ids = [int(item) for item in input_ids_raw]  # type: ignore
    request = GenerationRequest(
        input_ids=input_ids,
        max_new_tokens=int(args.get("max_new_tokens", 32)),  # type: ignore
    )
    return generate(request)


app.include_router(build_control_router({"health_check": _health_action, "generate": _generate_action}))


@app.on_event("startup")
def _startup_banner() -> None:
    _print_startup_banner()