serve.py

#!/usr/bin/env python3
"""OpenAI-compatible, microbatched RivetCoder FP8 serving entry point."""

from __future__ import annotations

import argparse
import asyncio
import concurrent.futures
import json
import queue
import threading
import time
import uuid
from collections import deque
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


@dataclass(slots=True)
class PendingCompletion:
    messages: list[dict[str, Any]]
    max_tokens: int
    temperature: float
    top_p: float
    future: concurrent.futures.Future[str]

    @property
    def batch_key(self) -> tuple[int, float, float]:
        return self.max_tokens, self.temperature, self.top_p


class MicrobatchEngine:
    """Collect compatible requests briefly, then generate them as one batch."""

    def __init__(
        self,
        model: Any,
        tokenizer: Any,
        *,
        max_batch_size: int,
        batch_wait_ms: float,
    ) -> None:
        self.model = model
        self.tokenizer = tokenizer
        self.max_batch_size = int(max_batch_size)
        self.batch_wait_seconds = float(batch_wait_ms) / 1000.0
        self.incoming: queue.Queue[PendingCompletion | None] = queue.Queue()
        self.deferred: deque[PendingCompletion] = deque()
        self.thread = threading.Thread(target=self._worker, name="rivetcoder-gpu", daemon=True)
        self.thread.start()

    def submit(
        self,
        messages: list[dict[str, Any]],
        *,
        max_tokens: int,
        temperature: float,
        top_p: float,
    ) -> concurrent.futures.Future[str]:
        future: concurrent.futures.Future[str] = concurrent.futures.Future()
        self.incoming.put(
            PendingCompletion(
                messages=messages,
                max_tokens=max_tokens,
                temperature=temperature,
                top_p=top_p,
                future=future,
            )
        )
        return future

    def close(self) -> None:
        self.incoming.put(None)

    def _next_request(self) -> PendingCompletion | None:
        if self.deferred:
            return self.deferred.popleft()
        return self.incoming.get()

    def _collect_batch(self, first: PendingCompletion) -> list[PendingCompletion]:
        batch = [first]
        key = first.batch_key
        deadline = time.perf_counter() + self.batch_wait_seconds
        while len(batch) < self.max_batch_size:
            remaining = deadline - time.perf_counter()
            if remaining <= 0:
                break
            try:
                item = self.incoming.get(timeout=remaining)
            except queue.Empty:
                break
            if item is None:
                self.incoming.put(None)
                break
            if item.batch_key == key:
                batch.append(item)
            else:
                self.deferred.append(item)
        return batch

    def _worker(self) -> None:
        while True:
            first = self._next_request()
            if first is None:
                return
            batch = self._collect_batch(first)
            try:
                results = self._generate(batch)
            except BaseException as error:
                for item in batch:
                    item.future.set_exception(error)
                continue
            for item, text in zip(batch, results, strict=True):
                item.future.set_result(text)

    def _generate(self, batch: list[PendingCompletion]) -> list[str]:
        rendered = [
            self.tokenizer.apply_chat_template(
                item.messages,
                add_generation_prompt=True,
                tokenize=False,
            )
            for item in batch
        ]
        encoded = self.tokenizer(
            rendered,
            add_special_tokens=False,
            padding=True,
            return_tensors="pt",
        ).to("cuda")
        prompt_width = encoded["input_ids"].shape[-1]
        temperature = batch[0].temperature
        generation_kwargs = {
            "max_new_tokens": batch[0].max_tokens,
            "do_sample": temperature > 0,
            "use_cache": True,
            "logits_to_keep": 1,
            "pad_token_id": self.tokenizer.pad_token_id,
            "eos_token_id": self.tokenizer.eos_token_id,
        }
        if temperature > 0:
            generation_kwargs.update(temperature=temperature, top_p=batch[0].top_p)
        with torch.no_grad():
            generated = self.model.generate(**encoded, **generation_kwargs)
        return self.tokenizer.batch_decode(
            generated[:, prompt_width:],
            skip_special_tokens=True,
        )


def load_runtime(args: argparse.Namespace) -> tuple[Any, Any, dict[str, Any]]:
    tokenizer = AutoTokenizer.from_pretrained(
        args.model,
        trust_remote_code=True,
        local_files_only=args.local_files_only,
    )
    tokenizer.padding_side = "left"
    model = AutoModelForCausalLM.from_pretrained(
        args.model,
        trust_remote_code=True,
        local_files_only=args.local_files_only,
        dtype=torch.bfloat16,
        device_map=0,
        attn_implementation="sdpa",
    ).eval()
    model.config.use_cache = True
    model.set_coding_enabled(True)
    if not hasattr(model, "enable_fast_fp8_serving"):
        raise RuntimeError(
            "The model package does not contain the grouped-FP8 runtime. "
            "Use the updated RivetCoder FP8 package."
        )
    report = model.enable_fast_fp8_serving()
    return tokenizer, model, report


def warmup(model: Any, tokenizer: Any, batch_sizes: list[int]) -> None:
    text = tokenizer.apply_chat_template(
        [{"role": "user", "content": "Return the integer 1."}],
        add_generation_prompt=True,
        tokenize=False,
    )
    for batch_size in batch_sizes:
        encoded = tokenizer(
            [text] * batch_size,
            add_special_tokens=False,
            padding=True,
            return_tensors="pt",
        ).to("cuda")
        with torch.no_grad():
            model.generate(
                **encoded,
                max_new_tokens=2,
                do_sample=False,
                use_cache=True,
                logits_to_keep=1,
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id,
            )


def build_app(engine: MicrobatchEngine, runtime_report: dict[str, Any], model_name: str) -> Any:
    try:
        from fastapi import FastAPI, HTTPException
    except ImportError as error:
        raise RuntimeError("Serving requires fastapi and uvicorn") from error

    app = FastAPI(title="RivetCoder FP8 Server")
    created = int(time.time())

    @app.get("/health")
    async def health() -> dict[str, Any]:
        return {
            "status": "ok",
            "model": model_name,
            "runtime": runtime_report,
            "max_batch_size": engine.max_batch_size,
            "batch_wait_ms": engine.batch_wait_seconds * 1000.0,
            "cuda_allocated_gib": torch.cuda.memory_allocated() / 1024**3,
        }

    @app.get("/v1/models")
    async def models() -> dict[str, Any]:
        return {
            "object": "list",
            "data": [{"id": model_name, "object": "model", "created": created, "owned_by": "HCHs"}],
        }

    @app.post("/v1/chat/completions")
    async def chat_completions(payload: dict[str, Any]) -> dict[str, Any]:
        if payload.get("stream", False):
            raise HTTPException(status_code=400, detail="Streaming is not implemented in the microbatch server")
        messages = payload.get("messages")
        if not isinstance(messages, list) or not messages:
            raise HTTPException(status_code=400, detail="messages must be a non-empty list")
        max_tokens = int(payload.get("max_tokens", 512))
        if max_tokens < 1 or max_tokens > 4096:
            raise HTTPException(status_code=400, detail="max_tokens must be between 1 and 4096")
        temperature = float(payload.get("temperature", 0.2))
        top_p = float(payload.get("top_p", 0.95))
        future = engine.submit(
            messages,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
        )
        try:
            text = await asyncio.wrap_future(future)
        except Exception as error:
            raise HTTPException(status_code=500, detail=str(error)) from error
        completion_id = f"chatcmpl-{uuid.uuid4().hex}"
        return {
            "id": completion_id,
            "object": "chat.completion",
            "created": int(time.time()),
            "model": model_name,
            "choices": [
                {
                    "index": 0,
                    "message": {"role": "assistant", "content": text},
                    "finish_reason": "stop",
                }
            ],
        }

    return app


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--model",
        default=str(Path("RivetCoder-9B-A4B-FP8")),
        help="Local FP8 model directory or Hugging Face model id",
    )
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8000)
    parser.add_argument("--max-batch-size", type=int, default=16)
    parser.add_argument("--batch-wait-ms", type=float, default=3.0)
    parser.add_argument("--warmup-batches", default="1,8,16")
    parser.add_argument("--local-files-only", action=argparse.BooleanOptionalAction, default=True)
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    if not torch.cuda.is_available():
        raise SystemExit("CUDA is required")
    torch.set_float32_matmul_precision("high")
    tokenizer, model, runtime_report = load_runtime(args)
    warmup_batches = [int(value) for value in args.warmup_batches.split(",") if value]
    warmup(model, tokenizer, warmup_batches)
    engine = MicrobatchEngine(
        model,
        tokenizer,
        max_batch_size=args.max_batch_size,
        batch_wait_ms=args.batch_wait_ms,
    )
    app = build_app(engine, runtime_report, str(args.model))
    print(json.dumps({"runtime": runtime_report, "listen": f"http://{args.host}:{args.port}"}, indent=2))
    import uvicorn

    uvicorn.run(app, host=args.host, port=args.port, workers=1)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())