vllm_support/src/kalm_t5gemma2_vllm_plugin/server.py

from __future__ import annotations

import argparse
import threading
import time
from pathlib import Path
from typing import Any, Optional

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field

from .constants import (
    DEFAULT_DOCUMENT_MAX_LENGTH,
    DEFAULT_ENCODER_CHUNK_SIZE,
    DEFAULT_MAX_MODEL_LEN,
    DEFAULT_QUERY_MAX_LENGTH,
    MODEL_ID,
    SUPPORTED_ENCODER_CHUNK_SIZES,
    parse_encoder_chunk_size,
)
from .reranker import KaLMVLLMReranker


class RerankRequest(BaseModel):
    query: str
    documents: list[str] = Field(min_length=1)
    instruction: Optional[str] = None
    top_k: Optional[int] = None
    return_margin: bool = False


class ScorePair(BaseModel):
    query: str
    document: str
    id: Optional[Any] = None


class ScoreRequest(BaseModel):
    pairs: list[ScorePair] = Field(min_length=1)
    instruction: Optional[str] = None
    return_margin: bool = False


class OnlineRerankerService:
    def __init__(self, args: argparse.Namespace) -> None:
        self.model = str(args.model)
        self.query_max_length = int(args.query_max_length)
        self.document_max_length = int(args.document_max_length)
        self.encoder_chunk_size = parse_encoder_chunk_size(args.encoder_chunk_size)
        self.max_model_len = int(args.max_model_len)
        self.batch_size = int(args.batch_size)
        self.dtype = str(args.dtype)
        self.gpu_memory_utilization = float(args.gpu_memory_utilization)
        self.started_at = time.time()
        self.lock = threading.Lock()
        self.reranker = KaLMVLLMReranker(
            self.model,
            query_max_length=self.query_max_length,
            document_max_length=self.document_max_length,
            encoder_chunk_size=self.encoder_chunk_size,
            max_model_len=self.max_model_len,
            batch_size=self.batch_size,
            dtype=self.dtype,
            gpu_memory_utilization=self.gpu_memory_utilization,
            tensor_parallel_size=args.tensor_parallel_size,
        )

    def config(self) -> dict[str, Any]:
        return {
            "model": self.model,
            "query_max_length": self.query_max_length,
            "document_max_length": self.document_max_length,
            "encoder_chunk_size": self.encoder_chunk_size,
            "max_model_len": self.max_model_len,
            "batch_size": self.batch_size,
            "dtype": self.dtype,
            "gpu_memory_utilization": self.gpu_memory_utilization,
            "tensor_parallel_size": 1,
            "supported_encoder_chunk_sizes": sorted(
                SUPPORTED_ENCODER_CHUNK_SIZES
            ),
        }

    def health(self) -> dict[str, Any]:
        return {
            "status": "ok",
            "uptime_seconds": round(time.time() - self.started_at, 3),
            **self.config(),
        }

    def close(self) -> None:
        self.reranker.close()

    def rerank(self, request: RerankRequest) -> dict[str, Any]:
        if request.top_k is not None and request.top_k < 0:
            raise ValueError("top_k must be non-negative or null.")
        with self.lock:
            rankings = self.reranker.rank(
                request.query,
                request.documents,
                instruction=request.instruction,
                top_k=request.top_k,
                return_margin=request.return_margin,
            )
        results: list[dict[str, Any]] = []
        for item in rankings:
            result = {
                "index": int(item["corpus_id"]),
                "score": float(item["score"]),
            }
            if request.return_margin:
                result["margin"] = float(item["margin"])
            results.append(result)
        return {"object": "rerank", "results": results}

    def score(self, request: ScoreRequest) -> dict[str, Any]:
        pairs = [(item.query, item.document) for item in request.pairs]
        with self.lock:
            predictions = self.reranker.predict(
                pairs,
                instruction=request.instruction,
                return_margin=True,
            )
        results: list[dict[str, Any]] = []
        for index, (pair, prediction) in enumerate(zip(request.pairs, predictions)):
            assert isinstance(prediction, dict)
            result = {
                "index": index,
                "score": float(prediction["score"]),
            }
            if pair.id is not None:
                result["id"] = pair.id
            if request.return_margin:
                result["margin"] = float(prediction["margin"])
            results.append(result)
        return {"object": "score", "results": results}


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="FastAPI server for the KaLM vLLM adapter.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=8000)
    parser.add_argument("--model", default=MODEL_ID)
    parser.add_argument("--query-max-length", type=int, default=DEFAULT_QUERY_MAX_LENGTH)
    parser.add_argument(
        "--document-max-length", type=int, default=DEFAULT_DOCUMENT_MAX_LENGTH
    )
    parser.add_argument(
        "--encoder-chunk-size", default=str(DEFAULT_ENCODER_CHUNK_SIZE)
    )
    parser.add_argument("--max-model-len", type=int, default=DEFAULT_MAX_MODEL_LEN)
    parser.add_argument("--batch-size", type=int, default=32)
    parser.add_argument("--dtype", default="bfloat16")
    parser.add_argument("--gpu-memory-utilization", type=float, default=0.85)
    parser.add_argument("--tensor-parallel-size", type=int, default=1)
    parser.add_argument("--log-level", default="info")
    return parser


def create_app(service: OnlineRerankerService) -> FastAPI:
    app = FastAPI(title="KaLM vLLM Online Reranker", version="0.1.0")
    app.state.service = service
    app.add_event_handler("shutdown", service.close)

    @app.get("/health")
    def health():
        return app.state.service.health()

    @app.post("/rerank")
    def rerank(request: RerankRequest):
        try:
            return app.state.service.rerank(request)
        except (TypeError, ValueError) as error:
            raise HTTPException(status_code=400, detail=str(error)) from error

    @app.post("/score")
    def score(request: ScoreRequest):
        try:
            return app.state.service.score(request)
        except (TypeError, ValueError) as error:
            raise HTTPException(status_code=400, detail=str(error)) from error

    return app


def main() -> int:
    args = build_parser().parse_args()
    args.encoder_chunk_size = parse_encoder_chunk_size(args.encoder_chunk_size)
    print("=== KaLM vLLM Online Reranker ===", flush=True)
    print(f"model: {args.model}", flush=True)
    print(f"query_max_length: {args.query_max_length}", flush=True)
    print(f"document_max_length: {args.document_max_length}", flush=True)
    print(f"encoder_chunk_size: {args.encoder_chunk_size}", flush=True)
    print(f"max_model_len: {args.max_model_len}", flush=True)
    print(f"listen: http://{args.host}:{args.port}", flush=True)

    service = OnlineRerankerService(args)
    app = create_app(service)
    import uvicorn

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


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