"""Jev/SystemOne-compatible HTTP API for Clef models on vLLM. POST /v1/systemone same request/response bodies as joint_schema_model.systemone() GET /v1/models, /health Requests are encoded with the release's encode_record and scored by vLLM, which batches concurrent requests continuously. --dp N runs one engine per GPU behind one port. clef-systemone --model ./clef-flash-NVFP4 --dp 2 --port 8000 """ import argparse import asyncio import importlib.util import sys import uuid from pathlib import Path import uvicorn from fastapi import FastAPI, HTTPException from fastapi.responses import JSONResponse from transformers import AutoTokenizer from vllm import PoolingParams from vllm.engine.arg_utils import AsyncEngineArgs from vllm.v1.engine.async_llm import AsyncLLM from clef_vllm.schema import ARCHITECTURE, TASK, schema_layout def parse_args(argv=None): ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--model", required=True, help="local path of a Clef checkpoint (snapshot directory)") ap.add_argument("--served-model-name", default="clef-flash") ap.add_argument("--host", default="0.0.0.0") ap.add_argument("--port", type=int, default=8000) ap.add_argument("--dp", type=int, default=1, help="data-parallel engines, one per GPU") ap.add_argument("--tp", type=int, default=1, help="tensor-parallel GPUs per engine") ap.add_argument("--max-model-len", type=int, default=16384) ap.add_argument("--max-num-seqs", type=int, default=64) ap.add_argument("--max-num-batched-tokens", type=int, default=16384) ap.add_argument("--gpu-memory-utilization", type=float, default=0.90) ap.add_argument("--flashinfer-autotune", choices=["auto", "on", "off"], default="auto", help="FlashInfer kernel autotuning; auto = off when --tp > 1 (vLLM 0.28 can deadlock " "autotuning NVFP4 GEMMs across tensor-parallel ranks)") ap.add_argument("--no-warmup", action="store_true", help="skip the startup warm-up (first requests then pay kernel JIT/autotune time)") return ap.parse_args(argv) def _release_module(model_dir: Path): spec = importlib.util.spec_from_file_location("joint_schema_model", model_dir / "joint_schema_model.py") module = importlib.util.module_from_spec(spec) sys.modules[spec.name] = module spec.loader.exec_module(module) return module def create_app(args) -> FastAPI: model_dir = Path(args.model) release = _release_module(model_dir) tokenizer = AutoTokenizer.from_pretrained(model_dir) app = FastAPI(title="Clef SystemOne API") state = {"engine": None, "ready": False} async def score(encoded): params = PoolingParams(task=TASK, extra_kwargs={"clef_questions": schema_layout(encoded)}) final = None async for output in state["engine"].encode( {"prompt_token_ids": list(encoded.input_ids)}, params, f"so-{uuid.uuid4().hex}" ): final = output if final is None: raise HTTPException(500, "engine returned no output") return final.outputs.data.float() async def warmup() -> None: """Run batches of several sizes so kernel JIT/autotuning happens before serving.""" question = {"type": "choice", "instructions": "Pick one.", "criteria": {"a": "First.", "b": "Second."}} for words, concurrent in ((50, 1), (400, 8), (1500, 32), (6000, 4), (200, args.max_num_seqs)): record = {"state": "warm up " * words, "questions": {"q": question, "n": {"type": "noul"}}} encoded = release.encode_record(tokenizer, record, args.max_model_len) await asyncio.gather(*(score(encoded) for _ in range(concurrent))) def validate(request: dict) -> None: """The checks of joint_schema_model.systemone(), plus text-only input.""" questions = request.get("questions") if not isinstance(request.get("model"), str) or "state" not in request: raise ValueError("model and state are required") if not isinstance(questions, dict) or not questions: raise ValueError("at least one question is required") for question_id, question in questions.items(): if not isinstance(question, dict) or question.get("type") not in release.QUESTION_TYPES: raise ValueError(f"{question_id}: type must be noul, choice, or score") if question["type"] != "noul" and not question.get("criteria"): raise ValueError(f"{question_id}: criteria must not be empty") if request.get("images") or request.get("videos"): raise ValueError("images and videos are not supported by this server (text and JSON state only)") autotune = args.flashinfer_autotune == "on" or (args.flashinfer_autotune == "auto" and args.tp == 1) kernel_kwargs = {} if autotune else {"kernel_config": {"enable_flashinfer_autotune": False}} @app.on_event("startup") async def start_engine() -> None: state["engine"] = AsyncLLM.from_engine_args(AsyncEngineArgs( model=str(model_dir), hf_overrides={"architectures": [ARCHITECTURE]}, runner="pooling", tensor_parallel_size=args.tp, data_parallel_size=args.dp, enable_prefix_caching=False, language_model_only=True, max_model_len=args.max_model_len, max_num_seqs=args.max_num_seqs, max_num_batched_tokens=args.max_num_batched_tokens, gpu_memory_utilization=args.gpu_memory_utilization, disable_log_stats=True, **kernel_kwargs, )) if not args.no_warmup: await warmup() state["ready"] = True @app.on_event("shutdown") async def stop_engine() -> None: if state["engine"] is not None: state["engine"].shutdown() @app.post("/v1/systemone") async def systemone(request: dict): try: validate(request) encoded = await asyncio.to_thread(release.encode_record, tokenizer, request, args.max_model_len) except (ValueError, TypeError, KeyError) as exc: return JSONResponse(status_code=400, content={"error": {"message": str(exc), "type": "invalid_request_error"}}) flat = await score(encoded) answers, offset = {}, 0 for question in encoded.questions: count = len(question.option_ids) probabilities = flat[offset:offset + count].softmax(-1).tolist() offset += count answers[question.question_id] = release.systemone_answer( request["questions"][question.question_id], dict(zip(question.option_ids, probabilities))) return { "model": request["model"], "answers": answers, "usage": {"input_tokens": len(encoded.input_ids), "output_tokens": 0}, } @app.get("/v1/models") async def models(): return {"object": "list", "data": [{"id": args.served_model_name, "object": "model", "owned_by": "clef"}]} @app.get("/health") async def health(): engine = state["engine"] if engine is None or engine.errored or not state["ready"]: raise HTTPException(503, "engine not ready") return {"status": "ok"} return app def main(argv=None) -> None: args = parse_args(argv) uvicorn.run(create_app(args), host=args.host, port=args.port, log_level="info") if __name__ == "__main__": main()