Ch 2 — Model Deployment¶

"""
serve.py — Production LLM serving endpoint with FastAPI.

Requirements:
    pip install fastapi uvicorn[standard] pydantic python-jose[cryptography] vllm
"""
from __future__ import annotations

import json
import logging
import os
import time
from typing import AsyncIterator

from fastapi import Depends, FastAPI, HTTPException, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from jose import JWTError, jwt
from pydantic import BaseModel, Field
from vllm import AsyncLLMEngine, SamplingParams
from vllm.engine.arg_utils import AsyncEngineArgs

logger = logging.getLogger(__name__)
security = HTTPBearer()

# ── Auth ──────────────────────────────────────────────────────────────────────
SECRET_KEY: str = os.environ["JWT_SECRET_KEY"]
ALGORITHM: str = "HS256"


def verify_token(
    credentials: HTTPAuthorizationCredentials = Depends(security),
) -> dict:
    """Validate JWT bearer token; raise 401 on failure."""
    try:
        return jwt.decode(credentials.credentials, SECRET_KEY, algorithms=[ALGORITHM])
    except JWTError as exc:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid or expired token",
            headers={"WWW-Authenticate": "Bearer"},
        ) from exc


# ── Request / response models ─────────────────────────────────────────────────
class GenerateRequest(BaseModel):
    prompt: str = Field(..., min_length=1, max_length=32_000)
    max_tokens: int = Field(default=512, ge=1, le=4096)
    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
    top_p: float = Field(default=0.95, ge=0.0, le=1.0)
    stream: bool = Field(default=False)


class GenerateResponse(BaseModel):
    text: str
    tokens_generated: int
    latency_ms: float


# ── App + engine ──────────────────────────────────────────────────────────────
app = FastAPI(title="LLM Inference API", version="1.0.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["https://app.example.com"],
    allow_methods=["POST", "GET"],
    allow_headers=["Authorization", "Content-Type"],
)

engine: AsyncLLMEngine | None = None


@app.on_event("startup")
async def load_model() -> None:
    global engine
    engine = AsyncLLMEngine.from_engine_args(
        AsyncEngineArgs(
            model=os.environ.get("MODEL_ID", "meta-llama/Llama-3-8B-Instruct"),
            dtype="bfloat16",
            max_model_len=8192,
            gpu_memory_utilization=0.90,
        )
    )
    logger.info("vLLM engine ready")


# ── SSE streaming helper ──────────────────────────────────────────────────────
async def token_stream(
    request_id: str,
    prompt: str,
    params: SamplingParams,
) -> AsyncIterator[str]:
    async for output in engine.generate(prompt, params, request_id):  # type: ignore[union-attr]
        for completion in output.outputs:
            yield f"data: {json.dumps({'token': completion.text})}\n\n"
    yield "data: [DONE]\n\n"


# ── Endpoints ─────────────────────────────────────────────────────────────────
@app.post("/generate", response_model=GenerateResponse)
async def generate(
    body: GenerateRequest,
    claims: dict = Depends(verify_token),
) -> GenerateResponse | StreamingResponse:
    if engine is None:
        raise HTTPException(status_code=503, detail="Model not loaded")

    params = SamplingParams(
        temperature=body.temperature,
        top_p=body.top_p,
        max_tokens=body.max_tokens,
    )
    request_id = f"{claims.get('sub', 'anon')}-{time.monotonic_ns()}"

    if body.stream:
        return StreamingResponse(
            token_stream(request_id, body.prompt, params),
            media_type="text/event-stream",
            headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
        )

    t0 = time.monotonic()
    outputs = []
    async for out in engine.generate(body.prompt, params, request_id):
        outputs = out.outputs
    text = outputs[0].text if outputs else ""
    return GenerateResponse(
        text=text,
        tokens_generated=len(outputs[0].token_ids) if outputs else 0,
        latency_ms=round((time.monotonic() - t0) * 1000, 2),
    )


@app.get("/health")
async def health() -> dict:
    return {"status": "ok", "model_loaded": engine is not None}

"""
export_onnx.py — Export a PyTorch classifier to ONNX with dynamic batch size.

Requirements:
    pip install torch onnx onnxruntime
"""
from __future__ import annotations

import numpy as np
import onnxruntime as ort
import torch
import torch.nn as nn


def export_to_onnx(
    model: nn.Module,
    sample_input: torch.Tensor,
    output_path: str = "model.onnx",
    opset_version: int = 17,
) -> None:
    """Export model to ONNX with dynamic batch size and sequence length."""
    model.eval()
    with torch.no_grad():
        torch.onnx.export(
            model,
            sample_input,
            output_path,
            opset_version=opset_version,
            input_names=["input_ids"],
            output_names=["logits"],
            dynamic_axes={
                "input_ids": {0: "batch_size", 1: "sequence_length"},
                "logits": {0: "batch_size"},
            },
            do_constant_folding=True,
        )
    print(f"Exported ONNX model → {output_path}")


def run_onnx_inference(
    model_path: str,
    input_ids: np.ndarray,
) -> np.ndarray:
    """Run inference with ONNX Runtime (CUDA preferred, CPU fallback)."""
    session = ort.InferenceSession(
        model_path,
        providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
    )
    return session.run(output_names=["logits"], input_feed={"input_ids": input_ids})[0]

import torch
import torch.nn as nn


def export_torchscript(
    model: nn.Module,
    sample_input: torch.Tensor,
    output_path: str = "model_scripted.pt",
) -> None:
    """Trace model to TorchScript for C++ or mobile deployment."""
    model.eval()
    with torch.no_grad():
        scripted = torch.jit.trace(model, sample_input)
    scripted.save(output_path)


def load_torchscript(model_path: str, x: torch.Tensor) -> torch.Tensor:
    loaded = torch.jit.load(model_path)
    loaded.eval()
    with torch.no_grad():
        return loaded(x)

from safetensors.torch import load_file, save_file
import torch

# Save
save_file(model.state_dict(), "model.safetensors")

# Load — zero-copy memory map, fast and safe
state_dict = load_file("model.safetensors", device="cuda")
model.load_state_dict(state_dict)

import torch
from torch.quantization import quantize_dynamic

# Dynamically quantise Linear layers to INT8
quantised = quantize_dynamic(
    model,
    qconfig_spec={torch.nn.Linear},
    dtype=torch.qint8,
)

# Memory comparison
fp32_mb = sum(p.numel() * 4 for p in model.parameters()) / 1e6
int8_mb = sum(p.numel() * 1 for p in model.parameters()) / 1e6
print(f"FP32: {fp32_mb:.0f} MB → INT8: {int8_mb:.0f} MB  (4× reduction)")

from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
from transformers import AutoTokenizer

quant_config = BaseQuantizeConfig(bits=4, group_size=128, desc_act=False)
tokeniser = AutoTokenizer.from_pretrained("meta-llama/Llama-3-8B-Instruct")
model = AutoGPTQForCausalLM.from_pretrained(
    "meta-llama/Llama-3-8B-Instruct",
    quantize_config=quant_config,
)

# Calibration data (representative production queries)
calibration = [
    tokeniser("Summarise the following document.", return_tensors="pt").input_ids,
    tokeniser("What are the key risks?", return_tensors="pt").input_ids,
]
model.quantize(calibration)
model.save_quantized("llama3-8b-int4-gptq")

from transformers import AutoModelForCausalLM
import torch

# Flash Attention 2 — enabled via attn_implementation flag
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3-8B-Instruct",
    attn_implementation="flash_attention_2",
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

from vllm import LLM, SamplingParams

llm = LLM(
    model="meta-llama/Llama-3-8B-Instruct",
    dtype="bfloat16",
    max_model_len=8192,
    gpu_memory_utilization=0.90,
)

prompts = [
    "Explain the attention mechanism in transformers.",
    "What is MLOps?",
    "Summarise the EU AI Act in three bullet points.",
]
params = SamplingParams(temperature=0.7, max_tokens=512)
outputs = llm.generate(prompts, params)

for output in outputs:
    print(output.outputs[0].text[:100])