Blame for server.py

"""Run: uv run python server.py --port 8043."""
from __future__ import annotations

import argparse
import threading
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Annotated, Literal

import uvicorn
from fastapi import FastAPI, HTTPException, Query
from fastapi.responses import FileResponse, Response
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, ConfigDict, Field

from runtime import Runtime, MAX_TOKENS, MODEL_ID, GEMMA_ID, MODELS
from grpo import GRPOJobs
from critic import CriticRuntime
from ppo import PPOJobs

ROOT = Path(__file__).parent
ModelChoice = Literal[MODEL_ID, GEMMA_ID]
runtimes = {model_id: Runtime(model_id) for model_id in MODELS}
grpo_jobs = GRPOJobs()
critic = CriticRuntime()
ppo_jobs = PPOJobs(critic)


class Message(BaseModel):
    role: Literal["system", "user", "assistant"]
    content: str = Field(max_length=50000)


class Input(BaseModel):
    model_config = ConfigDict(extra="forbid")
    model: ModelChoice = MODEL_ID
    mode: Literal["chat", "raw"] = "chat"
    messages: list[Message] = Field(default_factory=list, max_length=40)
    text: str = Field(default="", max_length=100000)
    thinking: bool = False
    max_new_tokens: int = Field(default=0, ge=0, le=128)


class AppendInput(BaseModel):
    model_config = ConfigDict(extra="forbid")
    text: str | None = Field(default=None, max_length=20000)


class RolloutInput(BaseModel):
    model_config = ConfigDict(extra="forbid", allow_inf_nan=False)
    model: ModelChoice = MODEL_ID
    mode: Literal["chat", "raw"] = "chat"
    prefix: str = Field(min_length=1, max_length=20000)
    assistant_prefix: str = Field(default="", max_length=5000)
    answer: str = Field(min_length=1, max_length=5000)
    thinking: bool = False
    group_size: int = Field(default=5, ge=2, le=12)
    max_new_tokens: int = Field(default=24, ge=1, le=64)
    temperature: float = Field(default=1, ge=0.1, le=2)
    seed: int = Field(default=42, ge=0, le=4294967295)


class PreviewInput(BaseModel):
    model_config = ConfigDict(extra="forbid", allow_inf_nan=False)
    answer: str = Field(min_length=1, max_length=5000)
    reward_rule: Literal["exact", "contains"] = "exact"
    rewards: list[Annotated[float, Field(ge=-1000000, le=1000000)] | None] = Field(min_length=2, max_length=12)
    normalize: bool = True
    clipping: bool = True
    epsilon: float = Field(default=0.2, ge=0.01, le=0.99)
    kl: bool = True
    beta: float = Field(default=0.04, ge=0, le=2)
    step_size: float = Field(default=10, ge=0, le=100)
    drift: float = Field(default=0, ge=-3, le=3)


class PPOPreviewInput(PreviewInput):
    normalize: bool = False
    step_size: float = Field(default=0, ge=0, le=100)
    gamma: float = Field(default=1, ge=0, le=1)
    gae_lambda: float = Field(default=0.95, ge=0, le=1)
    critic: bool = True
    values: list[Annotated[list[Annotated[float, Field(ge=-1000000, le=1000000)] | None], Field(max_length=64)] | None] | None = Field(default=None, max_length=12)


@asynccontextmanager
async def lifespan(app):
    def load_models():
        for runtime in runtimes.values():
            runtime.load()
        critic.load()
    threading.Thread(target=load_models, daemon=True).start()
    yield


app = FastAPI(lifespan=lifespan)


def start_rollouts(payload, jobs):
    runtime = runtimes[payload.model]
    try:
        context, ids = runtime.prepare(Input(model=payload.model, mode=payload.mode, text=payload.prefix,
            messages=[Message(role="user", content=payload.prefix)], thinking=payload.thinking))
        if payload.mode == "chat" and payload.assistant_prefix:
            context, ids = runtime.prepare(Input(model=payload.model, mode="raw", text=context+payload.assistant_prefix))
        return jobs.start(runtime, payload, context, ids)
    except BlockingIOError as exc:
        raise HTTPException(409, str(exc)) from exc
    except ValueError as exc:
        raise HTTPException(422, str(exc)) from exc


@app.post("/api/grpo/runs")
def start_grpo(payload: RolloutInput):
    return start_rollouts(payload, grpo_jobs)


@app.post("/api/ppo/runs")
def start_ppo(payload: RolloutInput):
    return start_rollouts(payload, ppo_jobs)


@app.get("/api/grpo/runs/{sid}")
def get_grpo(sid: str):
    try:
        return grpo_jobs.snapshot(sid)
    except KeyError as exc:
        raise HTTPException(404, "This rollout group has expired. Sample a new group.") from exc


@app.post("/api/grpo/runs/{sid}/cancel")
def cancel_grpo(sid: str):
    try:
        return grpo_jobs.cancel(sid)
    except KeyError as exc:
        raise HTTPException(404, "This rollout group has expired.") from exc


@app.post("/api/grpo/runs/{sid}/preview")
def preview_grpo(sid: str, payload: PreviewInput):
    try:
        return grpo_jobs.preview(sid, payload)
    except KeyError as exc:
        raise HTTPException(404, "This rollout group has expired. Sample a new group.") from exc
    except ValueError as exc:
        raise HTTPException(422, str(exc)) from exc


@app.get("/api/ppo/runs/{sid}")
def get_ppo(sid: str):
    try:
        return ppo_jobs.snapshot(sid)
    except KeyError as exc:
        raise HTTPException(404, "This PPO group has expired. Sample a new group.") from exc


@app.post("/api/ppo/runs/{sid}/cancel")
def cancel_ppo(sid: str):
    try:
        return ppo_jobs.cancel(sid)
    except KeyError as exc:
        raise HTTPException(404, "This PPO group has expired.") from exc


@app.post("/api/ppo/runs/{sid}/preview")
def preview_ppo(sid: str, payload: PPOPreviewInput):
    try:
        return ppo_jobs.preview(sid, payload)
    except KeyError as exc:
        raise HTTPException(404, "This PPO group has expired. Sample a new group.") from exc
    except ValueError as exc:
        raise HTTPException(422, str(exc)) from exc


@app.get("/api/status")
def status(model: ModelChoice = MODEL_ID):
    return dict(**runtimes[model].status(), models=[dict(id=mid, **spec) for mid, spec in MODELS.items()], critic=critic.status())


@app.post("/api/prepare")
def prepare(payload: Input):
    runtime = runtimes[payload.model]
    try:
        text, ids = runtime.prepare(payload)
        return dict(text=text, count=len(ids), tokens=[runtime.token(i) for i in ids],
                    max_tokens=MAX_TOKENS, fits=len(ids) + payload.max_new_tokens <= MAX_TOKENS)
    except ValueError as exc:
        raise HTTPException(422, str(exc)) from exc


@app.post("/api/analyze")
def analyze(payload: Input):
    try:
        return runtimes[payload.model].analyze(payload)
    except BlockingIOError as exc:
        raise HTTPException(409, str(exc)) from exc
    except ValueError as exc:
        raise HTTPException(422, str(exc)) from exc
    except Exception as exc:
        import traceback
        traceback.print_exc()
        raise HTTPException(500, f"Computation failed: {exc}") from exc


def get_session(sid):
    for runtime in runtimes.values():
        try:
            return runtime, runtime.session(sid)
        except KeyError:
            continue
    raise HTTPException(404, "This result has been evicted from memory. Run the analysis again.")


@app.get("/api/runs/{sid}/inspect")
def inspect(sid: str, layer: int = Query(ge=0), query: int = Query(ge=0)):
    runtime, s = get_session(sid)
    if layer >= s["response"]["layers"]:
        raise HTTPException(422, "Layer is outside this model.")
    if query >= len(s["response"]["tokens"]):
        raise HTTPException(422, "Token position is outside the context.")
    return runtime.inspect(sid, layer, query)


@app.post("/api/runs/{sid}/append-preview")
def append_preview(sid: str, payload: AppendInput):
    runtime, session = get_session(sid)
    try:
        ids = runtime.append_ids(sid, payload.text)
    except ValueError as exc:
        raise HTTPException(422, str(exc)) from exc
    total = len(session["response"]["tokens"]) + len(ids)
    return dict(tokens=[runtime.token(tid) for tid in ids], count=len(ids), total=total,
                max_tokens=MAX_TOKENS, fits=total <= MAX_TOKENS)


@app.post("/api/runs/{sid}/continue")
def continue_run(sid: str, payload: AppendInput | None = None):
    runtime, s = get_session(sid)
    previous = s["response"]
    ids = [token["id"] for token in previous["tokens"]]
    try:
        # Preserve the exact token IDs, including a possible partial UTF-8 token.
        text = payload.text if payload is not None else None
        suffix = runtime.append_ids(sid, text)
        sources = previous["token_sources"] + ["predicted" if text is None else "manual"] * len(suffix)
        return runtime.analyze(Input(model=runtime.model_id, mode="raw"),
                               prefix=(previous["context"], ids + suffix),
                               input_count=previous["input_count"], token_sources=sources)
    except BlockingIOError as exc:
        raise HTTPException(409, str(exc)) from exc
    except ValueError as exc:
        raise HTTPException(422, str(exc)) from exc


@app.get("/api/runs/{sid}/attention")
def attention(sid: str, layer: int = Query(ge=0), head: int | Literal["mean"] = Query(...)):
    _, s = get_session(sid)
    if layer >= s["response"]["layers"] or (head != "mean" and not 0 <= head < s["response"]["heads"]):
        raise HTTPException(422, "Layer or head is outside this model.")
    heads = s["attention"][layer]
    # Average stored post-softmax probabilities, accumulating in float32.
    matrix = heads.mean(axis=0, dtype="<f4") if head == "mean" else heads[head].astype("<f4")
    return Response(matrix.tobytes(), media_type="application/octet-stream",
                    headers={"X-Token-Count": str(matrix.shape[0]), "Cache-Control": "private, max-age=3600"})


@app.get("/api/runs/{sid}/attention-layers")
def attention_layers(sid: str, query: int = Query(ge=0), head: int | Literal["mean"] = "mean"):
    runtime, session = get_session(sid)
    if query >= len(session["response"]["tokens"]):
        raise HTTPException(422, "Token position is outside the context.")
    if head != "mean" and not 0 <= head < session["response"]["heads"]:
        raise HTTPException(422, "Head is outside this model.")
    return runtime.attention_layers(sid, query, head)


@app.get("/")
def index():
    return FileResponse(ROOT / "static" / "index.html")


@app.get("/grpo")
def grpo_index():
    return FileResponse(ROOT / "static" / "grpo.html")


@app.get("/ppo")
def ppo_index():
    return FileResponse(ROOT / "static" / "ppo.html")


app.mount("/", StaticFiles(directory=ROOT / "static"), name="static")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--port", type=int, default=8043)
    parser.add_argument("--host", default="127.0.0.1")
    args = parser.parse_args()
    uvicorn.run(app, host=args.host, port=args.port)