feat(llm-serving): manage clef-serve in git, generated from app.py
The Clef runtime was running from hand-applied objects that existed only in the cluster.
A rebuild would have lost it.
Three things move into git:
- clef-serve.yaml: the Deployment and the Service the gateway routes to
- configMapGenerator for clef-serve/app.py, so an edit to the code rolls the pod
instead of leaving a stale copy mounted
- clef-flash.yaml is deleted: it described an SGLang predictor that cannot run on a
V100, and its slot is now clef-serve.yaml
Recorded in the manifest, because both cost hours: SGLang needs CUDA 13 for Clef and
CUDA 13 has no sm_70 kernels, and device_map over two cards breaks the forward pass with
'indices should be either on cpu or on the same device as the indexed tensor'.
Still open and written down in the file: the image cannot be pushed, because Cloudflare
answers 413 for layers above 100 MB. So each start pip-installs and needs PyPI.
This commit is contained in:
@@ -0,0 +1,22 @@
|
||||
# Base image: torch 2.8.0+cu128 still lists sm_70 in torch.cuda.get_arch_list().
|
||||
# The CUDA 13 builds of torch drop sm_70. worker-1 holds Tesla V100 cards with
|
||||
# compute capability 7.0, so a newer base kills the kernel launch at load time.
|
||||
# Keep this pin until the model moves to newer hardware.
|
||||
FROM pytorch/pytorch:2.8.0-cuda12.8-cudnn9-runtime
|
||||
|
||||
# Unbuffered output lets the pod log show the startup lines in order.
|
||||
ENV PYTHONUNBUFFERED=1 \
|
||||
PIP_NO_CACHE_DIR=1
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# torch comes from the base image. Never reinstall it here, or pip pulls a
|
||||
# wheel built for another CUDA version.
|
||||
COPY requirements.txt .
|
||||
RUN pip install --no-cache-dir -r requirements.txt
|
||||
|
||||
COPY app.py .
|
||||
|
||||
EXPOSE 8080
|
||||
|
||||
CMD ["uvicorn", "app:create_app", "--factory", "--host", "0.0.0.0", "--port", "8080"]
|
||||
@@ -0,0 +1,201 @@
|
||||
"""HTTP server for the Clef-flash decision model.
|
||||
|
||||
Clef is a decision model and not a chat model. It reads one state plus a schema
|
||||
of typed questions and returns one probability per allowed option. Consumers
|
||||
call POST /v1/systemone and never /v1/chat/completions.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import traceback
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
from typing import Any, AsyncIterator
|
||||
|
||||
import torch
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.responses import JSONResponse
|
||||
from safetensors.torch import load_file
|
||||
from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration
|
||||
|
||||
# uvicorn configures its own loggers only. Without this call the root logger
|
||||
# stays at WARNING and the startup lines below never reach the pod log.
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||||
LOG = logging.getLogger("clef-serve")
|
||||
|
||||
MODEL_PATH = os.environ.get("MODEL_PATH", "/mnt/models/clef-flash")
|
||||
DTYPE = os.environ.get("DTYPE", "float16")
|
||||
MAX_LEN = int(os.environ.get("MAX_LEN", "65536"))
|
||||
PORT = int(os.environ.get("PORT", "8080"))
|
||||
|
||||
# HF_HOME stays untouched. The k8s manifest sets it, and the loader must reuse
|
||||
# that cache path instead of downloading a second copy of an 18.8 GB checkpoint.
|
||||
|
||||
_TORCH_DTYPES = {
|
||||
"float16": torch.float16,
|
||||
"bfloat16": torch.bfloat16,
|
||||
"float32": torch.float32,
|
||||
}
|
||||
|
||||
|
||||
def _resolve_dtype() -> torch.dtype:
|
||||
"""Map the DTYPE env string to a torch dtype. Fail fast on a typo."""
|
||||
try:
|
||||
return _TORCH_DTYPES[DTYPE]
|
||||
except KeyError:
|
||||
raise ValueError(
|
||||
f"The server does not support DTYPE={DTYPE!r}. Use one of {sorted(_TORCH_DTYPES)}."
|
||||
) from None
|
||||
|
||||
|
||||
def _load_runtime() -> tuple[Any, Any, list[str], Any]:
|
||||
"""Load backbone, joint head and processor. Runs in a worker thread."""
|
||||
# The released code imports its sibling modules by bare name. The model
|
||||
# directory is therefore the only place on sys.path that resolves them.
|
||||
sys.path.insert(0, MODEL_PATH)
|
||||
|
||||
# This import must follow the sys.path.insert above. joint_schema_model.py
|
||||
# ships inside MODEL_PATH and not in site-packages. An import before that
|
||||
# insert raises ModuleNotFoundError. A local import binds the name in this
|
||||
# function scope only, so the caller must receive the callable as a value.
|
||||
from joint_schema_model import (
|
||||
ClefModel,
|
||||
JointSchemaHead,
|
||||
collate_records,
|
||||
encode_record,
|
||||
systemone,
|
||||
)
|
||||
|
||||
dtype = _resolve_dtype()
|
||||
LOG.info("Loading backbone. path=%s dtype=%s", MODEL_PATH, DTYPE)
|
||||
|
||||
# load_release_model passes device_map={"": str(device)} and so pins the
|
||||
# whole backbone on one card. This loader asks accelerate to shard instead.
|
||||
backbone = Qwen3_5ForConditionalGeneration.from_pretrained(
|
||||
MODEL_PATH,
|
||||
dtype=dtype,
|
||||
device_map="auto",
|
||||
low_cpu_mem_usage=True,
|
||||
)
|
||||
|
||||
# The head must sit on the device that produces the final hidden state,
|
||||
# because ClefModel applies it to the output of the last layer.
|
||||
head_cfg = json.loads((Path(MODEL_PATH) / "joint_head_config.json").read_text())
|
||||
head = JointSchemaHead(**head_cfg)
|
||||
head.load_state_dict(load_file(str(Path(MODEL_PATH) / "joint_head.safetensors")))
|
||||
# The output embeddings carry the last layer's device. Without them, take
|
||||
# the last parameter, which lives there too.
|
||||
output_embeddings = backbone.get_output_embeddings()
|
||||
if output_embeddings is not None:
|
||||
head_device = output_embeddings.weight.device
|
||||
else:
|
||||
head_device = next(reversed(list(backbone.parameters()))).device
|
||||
head.to(device=head_device, dtype=dtype)
|
||||
|
||||
model = ClefModel(backbone, head).eval()
|
||||
processor = AutoProcessor.from_pretrained(MODEL_PATH)
|
||||
|
||||
# hf_device_map maps layer names to ints for CUDA shards. Report it so a
|
||||
# human can see whether one card or two carried the weights.
|
||||
raw_map = getattr(backbone, "hf_device_map", None) or {}
|
||||
if not raw_map:
|
||||
raw_map = {"": str(next(backbone.parameters()).device)}
|
||||
devices = list(
|
||||
dict.fromkeys(
|
||||
f"cuda:{name}" if isinstance(name, int) else str(name)
|
||||
for name in raw_map.values()
|
||||
)
|
||||
)
|
||||
LOG.info(
|
||||
"Model ready. dtype=%s devices=%s device_map=%s head_device=%s",
|
||||
DTYPE,
|
||||
devices,
|
||||
raw_map,
|
||||
head_device,
|
||||
)
|
||||
return model, processor, devices, systemone
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _lifespan(app: FastAPI) -> AsyncIterator[None]:
|
||||
"""Load the model on startup without blocking the event loop."""
|
||||
# A 9.4B load takes minutes. asyncio.to_thread keeps /health answerable
|
||||
# during that window, so the startup probe sees 503 and not a dead socket.
|
||||
model, processor, devices, systemone_fn = await asyncio.to_thread(_load_runtime)
|
||||
app.state.model = model
|
||||
app.state.processor = processor
|
||||
app.state.devices = devices
|
||||
# The route handler runs at module scope and cannot see the local import
|
||||
# inside _load_runtime. Publish the callable here, or the handler raises
|
||||
# NameError on the first request.
|
||||
app.state.systemone = systemone_fn
|
||||
app.state.lock = asyncio.Lock()
|
||||
LOG.info("Runtime published. systemone_callable=%s", systemone_fn is not None)
|
||||
yield
|
||||
app.state.model = None
|
||||
app.state.processor = None
|
||||
app.state.systemone = None
|
||||
|
||||
|
||||
def create_app() -> FastAPI:
|
||||
app = FastAPI(title="clef-flash", lifespan=_lifespan)
|
||||
|
||||
@app.get("/health")
|
||||
async def health() -> JSONResponse:
|
||||
# A ready flag alone is not enough. The pod must not take traffic until
|
||||
# the weights and the joint head are both resident.
|
||||
if getattr(app.state, "model", None) is None:
|
||||
return JSONResponse({"status": "loading"}, status_code=503)
|
||||
return JSONResponse(
|
||||
{
|
||||
"status": "ok",
|
||||
"dtype": DTYPE,
|
||||
"devices": app.state.devices,
|
||||
},
|
||||
status_code=200,
|
||||
)
|
||||
|
||||
@app.post("/v1/systemone")
|
||||
async def systemone_route(request: Request) -> Any:
|
||||
model = getattr(app.state, "model", None)
|
||||
if model is None:
|
||||
return JSONResponse({"error": "model is not loaded"}, status_code=503)
|
||||
|
||||
body = await request.json()
|
||||
processor = app.state.processor
|
||||
|
||||
def _forward() -> dict:
|
||||
# inference_mode is thread-local. Enter it inside this call, or the
|
||||
# worker thread records no autograd state saving and still builds
|
||||
# the graph.
|
||||
with torch.inference_mode():
|
||||
return app.state.systemone(model, processor, body, max_length=MAX_LEN)
|
||||
|
||||
# Clef runs one prefill with shared buffers and a shared joint head.
|
||||
# Two calls at once corrupt each other, so serialize them here.
|
||||
async with app.state.lock:
|
||||
try:
|
||||
result = await asyncio.to_thread(_forward)
|
||||
except ValueError as exc:
|
||||
# A malformed request body is the caller's fault.
|
||||
return JSONResponse({"error": str(exc)}, status_code=400)
|
||||
except Exception as exc:
|
||||
LOG.error("systemone failed\n%s", traceback.format_exc())
|
||||
return JSONResponse(
|
||||
{"error": type(exc).__name__, "detail": str(exc)},
|
||||
status_code=500,
|
||||
)
|
||||
return result
|
||||
|
||||
return app
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
|
||||
uvicorn.run("app:create_app", factory=True, host="0.0.0.0", port=PORT)
|
||||
@@ -0,0 +1,12 @@
|
||||
# Runtime dependencies for the clef-flash server.
|
||||
# torch ships with the base image and stays out of this list.
|
||||
transformers==5.10.2
|
||||
# device_map="auto" sharding is an accelerate feature, not a transformers one
|
||||
accelerate>=1.1.0
|
||||
fastapi
|
||||
uvicorn[standard]
|
||||
pillow
|
||||
safetensors
|
||||
huggingface_hub
|
||||
sentencepiece
|
||||
protobuf
|
||||
Reference in New Issue
Block a user