Runtime Skills

Universal Runtime Skills

Best practices and code review checklists for the Universal Runtime - LlamaFarm's local ML inference server.

Overview

The Universal Runtime provides OpenAI-compatible endpoints for HuggingFace models:

• Text generation (Causal LMs: GPT, Llama, Mistral, Qwen)
• Text embeddings (BERT, sentence-transformers, ModernBERT)
• Classification, NER, and reranking
• OCR and document understanding
• Anomaly detection

Directory: runtimes/universal/ Python: 3.11+ Key Dependencies: PyTorch, Transformers, FastAPI, llama-cpp-python

Links to Shared Skills

This skill extends the shared Python practices. Always apply these first:

Runtime-Specific Checklists

Architecture

runtimes/universal/
├── server.py              # FastAPI app, model caching, endpoints
├── core/
│   └── logging.py         # UniversalRuntimeLogger (structlog)
├── models/
│   ├── base.py            # BaseModel ABC with device management
│   ├── language_model.py  # Transformers text generation
│   ├── gguf_language_model.py  # llama-cpp-python for GGUF
│   ├── encoder_model.py   # Embeddings, classification, NER, reranking
│   └── ...                # OCR, anomaly, document models
├── routers/
│   └── chat_completions/  # Chat completions with streaming
├── utils/
│   ├── device.py          # Device detection (CUDA/MPS/CPU)
│   ├── model_cache.py     # TTL-based model caching
│   ├── model_format.py    # GGUF vs transformers detection
│   └── context_calculator.py  # GGUF context size computation
└── tests/

Key Patterns

1. Model Loading with Double-Checked Locking

_model_load_lock = asyncio.Lock()

async def load_encoder(model_id: str, task: str = "embedding"):
    cache_key = f"encoder:{task}:{model_id}"
    if cache_key not in _models:
        async with _model_load_lock:
            # Double-check after acquiring lock
            if cache_key not in _models:
                model = EncoderModel(model_id, device, task=task)
                await model.load()
                _models[cache_key] = model
    return _models.get(cache_key)

2. Device-Aware Tensor Operations

class BaseModel(ABC):
    def get_dtype(self, force_float32: bool = False):
        if force_float32:
            return torch.float32
        if self.device in ("cuda", "mps"):
            return torch.float16
        return torch.float32

    def to_device(self, tensor: torch.Tensor, dtype=None):
        # Don't change dtype for integer tensors
        if tensor.dtype in (torch.int32, torch.int64, torch.long):
            return tensor.to(device=self.device)
        dtype = dtype or self.get_dtype()
        return tensor.to(device=self.device, dtype=dtype)

3. TTL-Based Model Caching

_models: ModelCache[BaseModel] = ModelCache(ttl=300)  # 5 min TTL

async def _cleanup_idle_models():
    while True:
        await asyncio.sleep(CLEANUP_CHECK_INTERVAL)
        for cache_key, model in _models.pop_expired():
            await model.unload()

4. Async Generation with Thread Pools

# GGUF models use blocking llama-cpp, run in executor
self._executor = ThreadPoolExecutor(max_workers=1)

async def generate(self, messages, max_tokens=512, ...):
    loop = asyncio.get_running_loop()
    return await loop.run_in_executor(self._executor, self._generate_sync)

Review Priority

When reviewing Universal Runtime code:

1. Critical - Security

• Path traversal prevention in file endpoints
• Input sanitization for model IDs

2. High - Memory & Device

• Proper CUDA/MPS cache clearing on unload
• torch.no_grad() for inference
• Correct dtype for device

3. Medium - Performance

• Model caching patterns
• Batch processing where applicable
• Streaming implementation

4. Low - Code Style

• Consistent with patterns.md
• Proper type hints

FastAPI for Inference Serving Checklist

FastAPI patterns for ML inference endpoints, streaming, and lifecycle management.

---

Category: Application Lifecycle

Use Lifespan Context Manager

What to check: Manage startup/shutdown with lifespan

Search pattern:

rg "@asynccontextmanager" --type py runtimes/universal/server.py -A 15

Pass criteria:

• Lifespan function handles startup and shutdown
• Background tasks started at startup
• Resources cleaned up at shutdown
• Models unloaded gracefully

Severity: High

Good pattern (from server.py):

@asynccontextmanager
async def lifespan(app: FastAPI):
    global _cleanup_task

    # Startup
    logger.info("Starting Universal Runtime")
    _cleanup_task = asyncio.create_task(_cleanup_idle_models())

    yield  # Server is running

    # Shutdown
    logger.info("Shutting down Universal Runtime")
    if _cleanup_task is not None:
        _cleanup_task.cancel()
        with suppress(asyncio.CancelledError):
            await _cleanup_task

    # Unload all models
    for cache_key, model in list(_models.items()):
        await model.unload()
    _models.clear()

app = FastAPI(lifespan=lifespan)

---

Background Task Cleanup

What to check: Properly cancel and await background tasks

Search pattern:

rg "\.cancel\(\)|CancelledError" --type py runtimes/universal/

Pass criteria:

• Call task.cancel() on shutdown
• Use suppress(CancelledError) to await cleanly
• Background tasks have try/except CancelledError

Severity: High

Good pattern:

# In shutdown:
_cleanup_task.cancel()
with suppress(asyncio.CancelledError):
    await _cleanup_task

# In background task:
async def _cleanup_idle_models():
    while True:
        try:
            await asyncio.sleep(CLEANUP_CHECK_INTERVAL)
            # ... cleanup logic
        except asyncio.CancelledError:
            logger.info("Cleanup task cancelled")
            break

---

Category: Endpoint Design

Use Async Route Handlers

What to check: Define routes as async for I/O-bound operations

Search pattern:

rg "@app\.(get|post|put|delete)" --type py runtimes/universal/server.py -A 2

Pass criteria:

• All routes use async def
• Await async model operations
• No blocking calls in route handlers

Severity: High

Good pattern:

@app.post("/v1/embeddings")
async def create_embeddings(request: EmbeddingRequest):
    model = await load_encoder(request.model, task="embedding")
    embeddings = await model.embed(texts, normalize=True)
    return {"data": embeddings}

---

Pydantic Request/Response Models

What to check: Use Pydantic models for request validation

Search pattern:

rg "class.*Request\(.*BaseModel\)" --type py runtimes/universal/

Pass criteria:

• Request models inherit from pydantic.BaseModel
• Use Literal types for enum-like fields
• Document fields with descriptions

Severity: Medium

Good pattern:

class EmbeddingRequest(PydanticBaseModel):
    model: str
    input: str | list[str]
    encoding_format: Literal["float", "base64"] | None = "float"
    user: str | None = None
    extra_body: dict | None = None

---

HTTPException with Proper Status Codes

What to check: Use appropriate HTTP status codes

Search pattern:

rg "HTTPException\(status_code=" --type py runtimes/universal/

Pass criteria:

• 400 for bad request / validation errors
• 404 for resource not found
• 413 for payload too large
• 500 for internal errors
• Chain exceptions with from e

Severity: Medium

Good pattern:

if len(content) > MAX_UPLOAD_SIZE:
    raise HTTPException(
        status_code=413,
        detail=f"File too large. Maximum size is {MAX_UPLOAD_SIZE // (1024 * 1024)} MB",
    )

try:
    result = await model.process()
except Exception as e:
    logger.error(f"Error: {e}", exc_info=True)
    raise HTTPException(status_code=500, detail=str(e)) from e

---

Category: Streaming Responses

Server-Sent Events (SSE) for Streaming

What to check: Use StreamingResponse with proper media type

Search pattern:

rg "StreamingResponse|text/event-stream" --type py runtimes/universal/

Pass criteria:

• media_type="text/event-stream"
• Headers to disable buffering
• Proper SSE format: "data: {...}\n\n"

Severity: High

Good pattern:

return StreamingResponse(
    generate_sse(),
    media_type="text/event-stream",
    headers={
        "Cache-Control": "no-cache",
        "Connection": "keep-alive",
        "X-Accel-Buffering": "no",  # Disable nginx buffering
    },
)

---

SSE Format with JSON Data

What to check: Format SSE messages correctly

Search pattern:

rg "data:.*\\\\n\\\\n" --type py runtimes/universal/

Pass criteria:

• Format: data: {json}\n\n
• End with data: [DONE]\n\n
• Use .model_dump_json() for Pydantic models

Severity: High

Good pattern:

async def generate_sse():
    # Initial chunk
    yield f"data: {initial_chunk.model_dump_json(exclude_none=True)}\n\n".encode()

    # Content chunks
    async for token in model.generate_stream(...):
        chunk = ChatCompletionChunk(...)
        yield f"data: {chunk.model_dump_json(exclude_none=True)}\n\n".encode()

    # Final chunk
    yield f"data: {final_chunk.model_dump_json(exclude_none=True)}\n\n".encode()
    yield b"data: [DONE]\n\n"

---

Force Event Loop Yield for Real-Time Streaming

What to check: Use asyncio.sleep(0) to flush stream buffers

Search pattern:

rg "asyncio\.sleep\(0\)" --type py runtimes/universal/

Pass criteria:

• Call await asyncio.sleep(0) after each token yield
• Ensures immediate delivery of tokens
• Without this, tokens buffer and arrive in batches

Severity: High

Good pattern:

async for token in token_stream:
    chunk = ChatCompletionChunk(...)
    yield f"data: {chunk.model_dump_json()}\n\n".encode()

    # CRITICAL: Force event loop to yield for immediate delivery
    await asyncio.sleep(0)

---

Category: File Uploads

Validate File Size Before Processing

What to check: Check file size to prevent memory exhaustion

Search pattern:

rg "MAX_UPLOAD_SIZE|\.read\(\)" --type py runtimes/universal/server.py

Pass criteria:

• Read file content first
• Check size against limit
• Return 413 if too large

Severity: High

Good pattern:

MAX_UPLOAD_SIZE = int(os.environ.get("MAX_UPLOAD_SIZE", 100 * 1024 * 1024))

@app.post("/v1/files")
async def upload_file(file: UploadFile):
    content = await file.read()
    if len(content) > MAX_UPLOAD_SIZE:
        raise HTTPException(
            status_code=413,
            detail=f"File too large. Maximum is {MAX_UPLOAD_SIZE // (1024 * 1024)} MB",
        )

---

Use Form Parameters for File Metadata

What to check: Accept file options via Form parameters

Search pattern:

rg "Form\(default=" --type py runtimes/universal/

Pass criteria:

• Use Form() for multipart form data
• Provide sensible defaults
• Document parameters in docstring

Severity: Low

Good pattern:

@app.post("/v1/files")
async def upload_file(
    file: UploadFile,
    convert_pdf: bool = Form(default=True),
    pdf_dpi: int = Form(default=150),
):
    """Upload a file for use with OCR or document extraction."""

---

Category: OpenAI API Compatibility

Follow OpenAI Response Schema

What to check: Match OpenAI API response structure

Search pattern:

rg "\"object\":|\"id\":|\"created\":" --type py runtimes/universal/

Pass criteria:

• Include "object" field (e.g., "list", "embedding")
• Include "id" with unique identifier
• Include "created" timestamp
• Include "model" field

Severity: High

Good pattern:

return {
    "id": f"chatcmpl-{os.urandom(16).hex()}",
    "object": "chat.completion",
    "created": int(datetime.now().timestamp()),
    "model": request.model,
    "choices": [...],
    "usage": {...},
}

---

Support OpenAI Types from SDK

What to check: Use openai SDK types for responses

Search pattern:

rg "from openai\.types" --type py runtimes/universal/

Pass criteria:

• Import types from openai SDK
• Use for streaming chunk types
• Ensures compatibility with OpenAI client

Severity: Medium

Good pattern:

from openai.types.chat.chat_completion_chunk import (
    ChatCompletionChunk,
    ChoiceDelta,
    Choice as ChoiceChunk,
)

chunk = ChatCompletionChunk(
    id=completion_id,
    object="chat.completion.chunk",
    created=created_time,
    model=model_name,
    choices=[ChoiceChunk(index=0, delta=ChoiceDelta(content=token), finish_reason=None)],
)

---

Category: Router Organization

Use APIRouter for Endpoint Groups

What to check: Organize related endpoints in routers

Search pattern:

rg "APIRouter|include_router" --type py runtimes/universal/

Pass criteria:

• Related endpoints in separate router module
• Router included in main app
• Prefix for route grouping if needed

Severity: Low

Good pattern:

# routers/chat_completions/router.py
from fastapi import APIRouter
router = APIRouter()

@router.post("/v1/chat/completions")
async def chat_completions(...):
    ...

# server.py
from routers.chat_completions import router as chat_completions_router
app.include_router(chat_completions_router)

---

Category: Health and Monitoring

Health Check Endpoint

What to check: Provide health check with device info

Search pattern:

rg "@app\.get\(\"/health" --type py runtimes/universal/

Pass criteria:

• Return device information
• List loaded models
• Include timestamp and PID

Severity: Medium

Good pattern:

@app.get("/health")
async def health_check():
    device_info = get_device_info()
    return {
        "status": "healthy",
        "device": device_info,
        "loaded_models": list(_models.keys()),
        "timestamp": datetime.utcnow().isoformat(),
        "pid": os.getpid(),
    }

---

List Loaded Models Endpoint

What to check: Provide endpoint to list loaded models

Search pattern:

rg "@app\.get\(\"/v1/models" --type py runtimes/universal/

Pass criteria:

• Return model IDs and types
• Follow OpenAI list format
• Include metadata (created, owner)

Severity: Low

Good pattern:

@app.get("/v1/models")
async def list_models():
    models_list = [
        {
            "id": model_id,
            "object": "model",
            "created": int(datetime.now().timestamp()),
            "owned_by": "transformers-runtime",
            "type": model.model_type,
        }
        for model_id, model in _models.items()
    ]
    return {"object": "list", "data": models_list}

Inference Performance Checklist

Optimization patterns for model caching, batching, memory management, and inference speed.

---

Category: Model Caching

TTL-Based Model Cache

What to check: Use TTL cache to auto-unload idle models

Search pattern:

rg "ModelCache|TTLCache" --type py runtimes/universal/

Pass criteria:

• Models cached with configurable TTL (default: 5 minutes)
• TTL refreshed on each access
• Background task cleans up expired models

Severity: High

Good pattern (from utils/model_cache.py):

class ModelCache(Generic[T]):
    def __init__(self, ttl: float, maxsize: int = 1000):
        self._ttl = ttl
        self._cache: TTLCache[str, T] = TTLCache(maxsize=maxsize, ttl=ttl * 10)
        self._access: dict[str, float] = {}

    def get(self, key: str, default: T | None = None) -> T | None:
        if key not in self._cache:
            return default
        self._access[key] = self._timer()  # Refresh TTL on read
        return self._cache[key]

---

Configurable Unload Timeout

What to check: Allow environment variable control of model timeout

Search pattern:

rg "MODEL_UNLOAD_TIMEOUT|CLEANUP_CHECK_INTERVAL" --type py runtimes/universal/

Pass criteria:

• MODEL_UNLOAD_TIMEOUT configurable (default: 300 seconds)
• CLEANUP_CHECK_INTERVAL configurable (default: 30 seconds)
• Values read from environment at startup

Severity: Medium

Good pattern:

MODEL_UNLOAD_TIMEOUT = int(os.getenv("MODEL_UNLOAD_TIMEOUT", "300"))
CLEANUP_CHECK_INTERVAL = int(os.getenv("CLEANUP_CHECK_INTERVAL", "30"))

_models: ModelCache[BaseModel] = ModelCache(ttl=MODEL_UNLOAD_TIMEOUT)

---

Cache Key Design for Model Variants

What to check: Include all configuration in cache keys

Search pattern:

rg "_make_.*_cache_key" --type py runtimes/universal/

Pass criteria:

• Cache key includes model_id
• Cache key includes task/mode (embedding, classification)
• Cache key includes quantization preference
• Cache key includes context size (for GGUF)

Severity: High

Good pattern:

def _make_encoder_cache_key(
    model_id: str,
    task: str,
    model_format: str,
    preferred_quantization: str | None = None,
    max_length: int | None = None,
) -> str:
    quant_key = preferred_quantization or "default"
    len_key = max_length if max_length is not None else "auto"
    return f"encoder:{task}:{model_format}:{model_id}:quant{quant_key}:len{len_key}"

---

Category: Lazy Loading

Double-Checked Locking for Model Loading

What to check: Prevent duplicate model loading with locking

Search pattern:

rg "_model_load_lock|async with.*lock" --type py runtimes/universal/

Pass criteria:

• Check cache before acquiring lock
• Double-check cache after acquiring lock
• Only one instance of each model loaded

Severity: Critical

Good pattern:

_model_load_lock = asyncio.Lock()

async def load_encoder(model_id: str, task: str = "embedding"):
    cache_key = _make_encoder_cache_key(model_id, task, model_format)

    if cache_key not in _models:
        async with _model_load_lock:
            # Double-check after acquiring lock
            if cache_key not in _models:
                model = EncoderModel(model_id, device, task=task)
                await model.load()
                _models[cache_key] = model

    return _models.get(cache_key)  # get() refreshes TTL

---

Load on First Request

What to check: Models loaded lazily, not at startup

Search pattern:

rg "await.*load\(\)" --type py runtimes/universal/server.py

Pass criteria:

• No models loaded in lifespan startup
• Models loaded when first endpoint is called
• Fast server startup time

Severity: Medium

---

Category: Memory Optimization

Periodic Cleanup Task

What to check: Background task unloads idle models

Search pattern:

rg "_cleanup_idle_models|pop_expired" --type py runtimes/universal/

Pass criteria:

• Cleanup task runs periodically
• Uses pop_expired() to get idle models
• Calls await model.unload() for each
• Continues running despite individual errors

Severity: High

Good pattern:

async def _cleanup_idle_models() -> None:
    while True:
        try:
            await asyncio.sleep(CLEANUP_CHECK_INTERVAL)

            for cache, cache_name in [(_models, "models"), (_classifiers, "classifiers")]:
                expired_items = cache.pop_expired()
                for cache_key, model in expired_items:
                    try:
                        await model.unload()
                        logger.info(f"Unloaded: {cache_key}")
                    except Exception as e:
                        logger.error(f"Error unloading {cache_key}: {e}")

        except asyncio.CancelledError:
            break
        except Exception as e:
            logger.error(f"Error in cleanup: {e}")
            # Continue running

---

Graceful Shutdown Model Cleanup

What to check: Unload all models on server shutdown

Search pattern:

rg "Shutting down|Unloading.*remaining" --type py runtimes/universal/

Pass criteria:

• All cached models unloaded
• GPU memory freed
• Log each model unload

Severity: High

Good pattern:

# In lifespan shutdown:
if _models:
    logger.info(f"Unloading {len(_models)} remaining model(s)")
    for cache_key, model in list(_models.items()):
        try:
            await model.unload()
            logger.info(f"Unloaded: {cache_key}")
        except Exception as e:
            logger.error(f"Error unloading {cache_key}: {e}")
    _models.clear()

---

Category: GGUF Optimizations

Compute Optimal Context Size

What to check: Auto-calculate context size based on available memory

Search pattern:

rg "get_default_context_size|context_calculator" --type py runtimes/universal/

Pass criteria:

• Consider available GPU/system memory
• Read model's default from GGUF metadata
• Allow user override via parameter
• Log warnings if reduced

Severity: High

Good pattern (conceptual):

def get_default_context_size(
    model_id: str,
    gguf_path: str,
    device: str,
    config_n_ctx: int | None = None,
) -> tuple[int, list[str]]:
    warnings = []

    # User-specified takes priority
    if config_n_ctx is not None:
        return config_n_ctx, warnings

    # Read model's default from GGUF
    model_default = read_gguf_context_length(gguf_path)

    # Calculate based on available memory
    available_memory = get_available_memory(device)
    computed_ctx = compute_safe_context(model_default, available_memory)

    if computed_ctx < model_default:
        warnings.append(f"Reduced context from {model_default} to {computed_ctx}")

    return computed_ctx, warnings

---

Single-File Quantization Download

What to check: Only download the requested quantization

Search pattern:

rg "preferred_quantization|get_gguf_file_path" --type py runtimes/universal/

Pass criteria:

• Accept quantization preference (Q4_K_M, Q8_0, etc.)
• Download only the specified file
• Default to Q4_K_M (good balance of size/quality)

Severity: Medium

Good pattern:

async def load_language(
    model_id: str,
    n_ctx: int | None = None,
    preferred_quantization: str | None = None,
):
    model = GGUFLanguageModel(
        model_id,
        device,
        n_ctx=n_ctx,
        preferred_quantization=preferred_quantization,  # Only downloads this one
    )

---

Category: Batch Processing

Batch Tokenization

What to check: Tokenize multiple inputs together for efficiency

Search pattern:

rg "self\.tokenizer\(" --type py runtimes/universal/models/ -A 5

Pass criteria:

• Pass list of texts to tokenizer (not one at a time)
• Use padding=True for batch processing
• Reduces tokenizer overhead

Severity: Medium

Good pattern:

async def embed(self, texts: list[str], normalize: bool = True):
    # Batch tokenization - single call for all texts
    encoded = self.tokenizer(
        texts,  # List of texts
        padding=True,
        truncation=True,
        max_length=self.max_length,
        return_tensors="pt",
    )

---

Batch Inference Where Possible

What to check: Process multiple inputs in single forward pass

Search pattern:

rg "model\(\*\*encoded\)" --type py runtimes/universal/models/

Pass criteria:

• Embedding: batch multiple texts
• Classification: batch multiple texts
• Generation: typically single request (streaming)

Severity: Medium

Good pattern:

# Good: Batch processing for embeddings
async def embed(self, texts: list[str], normalize: bool = True):
    encoded = self.tokenizer(texts, padding=True, ...)
    with torch.no_grad():
        model_output = self.model(**encoded)  # All texts at once
    embeddings = self._mean_pooling(model_output, encoded["attention_mask"])
    return embeddings.cpu().tolist()

---

Category: Streaming Performance

Avoid Buffering in Streams

What to check: Yield tokens immediately without buffering

Search pattern:

rg "async for.*in.*stream" --type py runtimes/universal/

Pass criteria:

• Yield each token as it's generated
• Use asyncio.sleep(0) to force flush
• Set no-cache headers on response

Severity: High

Good pattern:

async for token in model.generate_stream(messages=messages, ...):
    chunk = ChatCompletionChunk(...)
    yield f"data: {chunk.model_dump_json()}\n\n".encode()
    await asyncio.sleep(0)  # Force immediate delivery

---

Queue-Based Streaming for Blocking Backends

What to check: Use asyncio.Queue for thread-based generation

Search pattern:

rg "asyncio\.Queue" --type py runtimes/universal/

Pass criteria:

• Create queue for producer/consumer pattern
• Producer runs in thread pool
• Consumer yields from queue in async context

Severity: High

Good pattern (from gguf_language_model.py):

async def generate_stream(self, messages, ...):
    queue: asyncio.Queue[str | Exception | None] = asyncio.Queue()
    loop = asyncio.get_running_loop()

    def _generate_stream():
        for chunk in self.llama.create_chat_completion(stream=True, ...):
            content = chunk["choices"][0].get("delta", {}).get("content", "")
            if content:
                future = asyncio.run_coroutine_threadsafe(queue.put(content), loop)
                future.result()
        asyncio.run_coroutine_threadsafe(queue.put(None), loop).result()

    loop.run_in_executor(self._executor, _generate_stream)

    while True:
        item = await queue.get()
        if item is None:
            break
        yield item

---

Category: Model Format Detection

Automatic GGUF vs Transformers Detection

What to check: Detect model format automatically

Search pattern:

rg "detect_model_format|model_format" --type py runtimes/universal/

Pass criteria:

• Check HuggingFace repo for .gguf files
• Use GGUF loader for quantized models
• Use Transformers loader for standard models

Severity: High

Good pattern:

# Detect model format
model_format = detect_model_format(model_id)

if model_format == "gguf":
    model = GGUFLanguageModel(model_id, device, n_ctx=n_ctx)
else:
    model = LanguageModel(model_id, device)

---

Parse Model:Quantization Syntax

What to check: Support model:Q4_K_M syntax for quantization selection

Search pattern:

rg "parse_model_with_quantization" --type py runtimes/universal/

Pass criteria:

• Parse "model:Q4_K_M" into (model, Q4_K_M)
• Handle models without quantization suffix
• Support HuggingFace model IDs with slashes

Severity: Medium

Good pattern:

# Parse model name to extract quantization if present
model_id, gguf_quantization = parse_model_with_quantization(request.model)

model = await self.load_language(
    model_id,
    n_ctx=n_ctx,
    preferred_quantization=gguf_quantization,
)

PyTorch Patterns Checklist

Device management, dtype handling, and memory optimization for PyTorch inference.

---

Category: Device Management

Use get_optimal_device() for Device Selection

What to check: Centralized device detection with proper fallbacks

Search pattern:

rg "get_optimal_device|get_device_info" --type py runtimes/universal/

Pass criteria:

• Device detected once and cached
• Supports CUDA, MPS (Apple Silicon), and CPU
• Environment variable overrides (TRANSFORMERS_FORCE_CPU, TRANSFORMERS_SKIP_MPS)

Severity: High

Good pattern (from utils/device.py):

def get_optimal_device() -> str:
    if os.environ.get("TRANSFORMERS_FORCE_CPU", "").lower() in ("1", "true"):
        return "cpu"
    if torch.cuda.is_available():
        return "cuda"
    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        return "mps"
    return "cpu"

---

Cache Device Selection at Module Level

What to check: Device should be detected once, not on every request

Search pattern:

rg "get_device\(\)|get_optimal_device\(\)" --type py runtimes/universal/server.py

Pass criteria:

• Global _current_device cached after first detection
• get_device() function returns cached value

Severity: Medium

Good pattern:

_current_device = None

def get_device():
    global _current_device
    if _current_device is None:
        _current_device = get_optimal_device()
        logger.info(f"Using device: {_current_device}")
    return _current_device

---

Category: Dtype Handling

Device-Appropriate Dtype Selection

What to check: Use float16 on GPU, float32 on CPU

Search pattern:

rg "get_dtype|torch\.float16|torch\.float32" --type py runtimes/universal/models/

Pass criteria:

• float16 for CUDA and MPS (memory efficient)
• float32 for CPU (better compatibility)
• force_float32 option for MPS compatibility issues

Severity: High

Good pattern (from models/base.py):

def get_dtype(self, force_float32: bool = False):
    if force_float32:
        return torch.float32
    if self.device == "cuda" or self.device == "mps":
        return torch.float16
    else:
        return torch.float32

---

Preserve Integer Tensor Dtypes

What to check: Don't convert input_ids and attention_mask to float

Search pattern:

rg "\.to\(device" --type py runtimes/universal/models/

Pass criteria:

• Integer tensors (input_ids, attention_mask) only moved to device
• Float tensors get appropriate dtype conversion

Severity: High

Good pattern:

def to_device(self, tensor: torch.Tensor, dtype: torch.dtype | None = None):
    # Don't change dtype for integer tensors
    if tensor.dtype in (torch.int32, torch.int64, torch.long, torch.bool):
        return tensor.to(device=self.device)
    dtype = dtype or self.get_dtype()
    return tensor.to(device=self.device, dtype=dtype)

---

Category: Memory Management

Clear GPU Cache on Model Unload

What to check: Free GPU memory when unloading models

Search pattern:

rg "empty_cache|cuda\.empty_cache|mps\.empty_cache" --type py runtimes/universal/

Pass criteria:

• torch.cuda.empty_cache() called after unload
• torch.mps.empty_cache() called on Apple Silicon
• Model moved to CPU before clearing references

Severity: Critical

Good pattern (from models/base.py):

async def unload(self) -> None:
    # Move model to CPU to free GPU memory
    if self.model is not None and hasattr(self.model, "to"):
        self.model = self.model.to("cpu")

    # Clear references
    self.model = None
    self.tokenizer = None

    # Clear CUDA cache
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    # Clear MPS cache
    if hasattr(torch, "mps") and hasattr(torch.mps, "empty_cache"):
        torch.mps.empty_cache()

---

Use torch.no_grad() for Inference

What to check: Disable gradient computation during inference

Search pattern:

rg "with torch\.no_grad\(\):" --type py runtimes/universal/models/

Pass criteria:

• All inference code wrapped in torch.no_grad()
• No gradient computation for embeddings, generation, classification

Severity: Critical

Good pattern:

async def embed(self, texts: list[str], normalize: bool = True):
    with torch.no_grad():
        model_output = self.model(**encoded)
    embeddings = self._mean_pooling(model_output, attention_mask)
    return embeddings.cpu().tolist()

Recommendation: torch.no_grad() reduces memory usage and speeds up inference

---

Return Tensors to CPU Before Python Conversion

What to check: Move tensors to CPU before .tolist() or numpy conversion

Search pattern:

rg "\.tolist\(\)|\.numpy\(\)" --type py runtimes/universal/models/

Pass criteria:

• Always call .cpu() before .tolist() or .numpy()
• Prevents GPU memory from being held by Python objects

Severity: High

Good pattern:

# Correct - move to CPU first
embeddings = F.normalize(embeddings, p=2, dim=1)
return embeddings.cpu().tolist()

# Incorrect - GPU tensor leaked to Python
return embeddings.tolist()  # May fail on MPS/CUDA

---

Category: Model Loading

Specify device_map for CUDA

What to check: Use device_map="auto" for multi-GPU and efficient loading

Search pattern:

rg "device_map=" --type py runtimes/universal/models/

Pass criteria:

• device_map="auto" for CUDA devices
• None for CPU/MPS (handle manually)

Severity: Medium

Good pattern:

self.model = AutoModelForCausalLM.from_pretrained(
    self.model_id,
    dtype=dtype,
    device_map="auto" if self.device == "cuda" else None,
    trust_remote_code=True,
)

if self.device != "cuda":
    self.model = self.model.to(self.device)

---

Use torch_dtype Parameter

What to check: Pass dtype to from_pretrained() for efficient loading

Search pattern:

rg "torch_dtype=|dtype=" --type py runtimes/universal/models/

Pass criteria:

• torch_dtype set in from_pretrained() for GPU devices
• Omit for CPU to use full precision

Severity: Medium

Good pattern:

model_kwargs = {
    "trust_remote_code": True,
    "token": self.token,
}

# Only set torch_dtype for GPU devices
if self.device != "cpu":
    model_kwargs["torch_dtype"] = dtype

self.model = AutoModel.from_pretrained(self.model_id, **model_kwargs)

---

Category: Platform Optimizations

Apply Platform-Specific Optimizations

What to check: Enable optimizations for each platform

Search pattern:

rg "enable_attention_slicing|enable_xformers|enable_model_cpu_offload" --type py runtimes/universal/

Pass criteria:

• MPS: attention_slicing (reduces memory pressure)
• CUDA: xformers memory efficient attention
• CUDA: model CPU offload for large models

Severity: Medium

Good pattern:

def apply_optimizations(self):
    if self.device == "mps":
        self.pipe.enable_attention_slicing()
    elif self.device == "cuda":
        try:
            self.pipe.enable_xformers_memory_efficient_attention()
        except Exception:
            pass  # xformers not available

---

Handle MPS Limitations

What to check: Work around MPS 4GB buffer limit and other issues

Search pattern:

rg "TRANSFORMERS_SKIP_MPS|mps.*limit|4GB" --type py runtimes/universal/

Pass criteria:

• Environment variable to skip MPS if needed
• Warning about MPS limitations logged
• force_float32 option for incompatible models

Severity: Medium

Good pattern:

if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
    skip_mps = os.environ.get("TRANSFORMERS_SKIP_MPS", "").lower() in ("1", "true")
    if skip_mps:
        return "cpu"
    logger.warning(
        "MPS has a 4GB temporary buffer limit. "
        "Set TRANSFORMERS_SKIP_MPS=1 to use CPU if you encounter errors."
    )
    return "mps"

---

Category: Thread Safety

ThreadPoolExecutor for Blocking Operations

What to check: Run blocking llama-cpp operations in thread pool

Search pattern:

rg "ThreadPoolExecutor|run_in_executor" --type py runtimes/universal/

Pass criteria:

• Thread pool created with max_workers=1 (serialized inference)
• Blocking Llama() calls run via executor
• Executor shutdown on model unload

Severity: High

Good pattern (from models/gguf_language_model.py):

def __init__(self, model_id: str, device: str, ...):
    self._executor = ThreadPoolExecutor(max_workers=1)

async def load(self) -> None:
    loop = asyncio.get_running_loop()
    self.llama = await loop.run_in_executor(self._executor, _load_model)

async def unload(self) -> None:
    self.llama = None
    if hasattr(self, "_executor"):
        self._executor.shutdown(wait=True, cancel_futures=True)
        self._executor = ThreadPoolExecutor(max_workers=1)

---

Proper Executor Cleanup

What to check: Shutdown thread pool executors properly

Search pattern:

rg "shutdown\(wait=" --type py runtimes/universal/

Pass criteria:

• shutdown(wait=True) for graceful cleanup
• cancel_futures=True to stop pending work
• Create new executor if model might be reloaded

Severity: High

Good pattern:

async def unload(self) -> None:
    if hasattr(self, "_executor"):
        self._executor.shutdown(wait=True, cancel_futures=True)
        # Recreate for potential future use
        self._executor = ThreadPoolExecutor(max_workers=1)

Transformers Model Patterns Checklist

HuggingFace Transformers model loading, tokenization, and inference patterns.

---

Category: Model Loading

Use trust_remote_code=True Consistently

What to check: Enable trust_remote_code for custom model architectures

Search pattern:

rg "trust_remote_code=" --type py runtimes/universal/models/

Pass criteria:

• All AutoModel/AutoTokenizer calls include trust_remote_code=True
• Enables loading of custom model code (required for many popular models)

Severity: High

Good pattern:

self.tokenizer = AutoTokenizer.from_pretrained(
    self.model_id,
    trust_remote_code=True,
    token=self.token,
)

self.model = AutoModelForCausalLM.from_pretrained(
    self.model_id,
    trust_remote_code=True,
    token=self.token,
)

---

Load Config Before Model for Capability Detection

What to check: Use AutoConfig to detect model capabilities before loading

Search pattern:

rg "AutoConfig\.from_pretrained" --type py runtimes/universal/models/

Pass criteria:

• Config loaded first to detect max_length, Flash Attention support
• Avoids loading model with incompatible settings

Severity: Medium

Good pattern (from models/encoder_model.py):

async def load(self) -> None:
    # Load config first to detect capabilities
    config = AutoConfig.from_pretrained(
        self.model_id, trust_remote_code=True, token=self.token
    )

    # Detect max sequence length
    self._detected_max_length = self._detect_max_length(config)

    # Check Flash Attention support
    if self._supports_flash_attention(config):
        model_kwargs["attn_implementation"] = "flash_attention_2"

---

Use Correct AutoModel Class for Task

What to check: Select the right AutoModel variant for the task

Search pattern:

rg "AutoModel|AutoModelFor" --type py runtimes/universal/models/

Pass criteria:

• AutoModelForCausalLM for text generation
• AutoModel for embeddings (no classification head)
• AutoModelForSequenceClassification for classification/reranking
• AutoModelForTokenClassification for NER

Severity: High

Good pattern:

if self.task == "classification":
    self.model = AutoModelForSequenceClassification.from_pretrained(...)
elif self.task == "reranking":
    self.model = AutoModelForSequenceClassification.from_pretrained(...)
elif self.task == "ner":
    self.model = AutoModelForTokenClassification.from_pretrained(...)
else:  # embedding
    self.model = AutoModel.from_pretrained(...)

---

Set Model to Eval Mode for Inference

What to check: Call model.eval() after loading for inference

Search pattern:

rg "\.eval\(\)" --type py runtimes/universal/models/

Pass criteria:

• model.eval() called after loading
• Disables dropout and batch normalization training behavior

Severity: High

Good pattern:

if self.model is not None:
    self.model = self.model.to(self.device)
    self.model.eval()  # Critical for inference

---

Category: Tokenization

Use Tokenizer's Chat Template

What to check: Apply model-specific chat template for message formatting

Search pattern:

rg "apply_chat_template" --type py runtimes/universal/models/

Pass criteria:

• Use tokenizer.apply_chat_template() when available
• Fallback to simple format if not available
• Set add_generation_prompt=True for assistant response

Severity: High

Good pattern:

def format_messages(self, messages: list[dict]) -> str:
    if self.tokenizer and hasattr(self.tokenizer, "apply_chat_template"):
        try:
            return self.tokenizer.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True
            )
        except Exception:
            pass

    # Fallback to simple concatenation
    return "\n".join(f"{msg['role']}: {msg['content']}" for msg in messages)

---

Configure Tokenizer Padding and Truncation

What to check: Set padding, truncation, and max_length consistently

Search pattern:

rg "padding=|truncation=|max_length=" --type py runtimes/universal/models/

Pass criteria:

• padding=True for batch processing
• truncation=True to handle long inputs
• max_length set from model's detected capability
• return_tensors="pt" for PyTorch tensors

Severity: High

Good pattern:

encoded = self.tokenizer(
    texts,
    padding=True,
    truncation=True,
    max_length=self.max_length,
    return_tensors="pt",
)
encoded = {k: v.to(self.device) for k, v in encoded.items()}

---

Handle Tokenizer Output Properly

What to check: Move tokenizer output to device correctly

Search pattern:

rg "\.to\(self\.device\)" --type py runtimes/universal/models/

Pass criteria:

• All tensor values moved to device
• Use dict comprehension for clean handling
• Remove non-tensor fields before model forward pass

Severity: High

Good pattern:

# For embeddings - keep all fields
encoded = {k: v.to(self.device) for k, v in encoded.items()}

# For NER - remove offset_mapping before model call
offset_mapping = encoded.pop("offset_mapping")[0].tolist()
encoded = {k: v.to(self.device) for k, v in encoded.items()}

---

Detect Max Sequence Length from Config

What to check: Extract max_length from model config attributes

Search pattern:

rg "max_position_embeddings|max_seq_length|n_positions" --type py runtimes/universal/

Pass criteria:

• Check multiple config attributes for compatibility
• Handle extended context models (ModernBERT: 8192)
• Default to 512 for classic BERT

Severity: Medium

Good pattern:

def _detect_max_length(self, config: AutoConfig) -> int:
    # Check for known extended context models
    for prefix, length in self.EXTENDED_CONTEXT_MODELS.items():
        if self.model_id.startswith(prefix):
            return length

    # Try config attributes
    if hasattr(config, "max_position_embeddings"):
        return config.max_position_embeddings
    if hasattr(config, "max_seq_length"):
        return config.max_seq_length
    if hasattr(config, "n_positions"):
        return config.n_positions

    return 512  # Default for classic BERT

---

Category: Text Generation

Use pad_token_id from Tokenizer

What to check: Set pad_token_id in generate() to avoid warnings

Search pattern:

rg "pad_token_id=" --type py runtimes/universal/models/

Pass criteria:

• pad_token_id set to eos_token_id (common pattern)
• Prevents "Setting pad_token_id to eos_token_id" warning

Severity: Low

Good pattern:

outputs = self.model.generate(
    **inputs,
    max_new_tokens=max_tokens,
    temperature=temperature,
    pad_token_id=self.tokenizer.eos_token_id,
)

---

Decode Only New Tokens

What to check: Skip input tokens when decoding generation output

Search pattern:

rg "skip_special_tokens" --type py runtimes/universal/models/

Pass criteria:

• Slice output to skip input_ids: outputs[0][inputs.input_ids.shape[1]:]
• Use skip_special_tokens=True to clean output

Severity: Medium

Good pattern:

# Generate
outputs = self.model.generate(**inputs, max_new_tokens=max_tokens)

# Decode only the NEW tokens (skip input)
generated_text = self.tokenizer.decode(
    outputs[0][inputs.input_ids.shape[1]:],
    skip_special_tokens=True
)
return generated_text.strip()

---

Use TextIteratorStreamer for Streaming

What to check: Stream tokens using TextIteratorStreamer

Search pattern:

rg "TextIteratorStreamer" --type py runtimes/universal/models/

Pass criteria:

• TextIteratorStreamer instantiated with tokenizer
• skip_prompt=True to not re-emit input
• Generation runs in separate thread

Severity: Medium

Good pattern:

from transformers import TextIteratorStreamer
from threading import Thread

streamer = TextIteratorStreamer(
    cast(AutoTokenizer, self.tokenizer),
    skip_prompt=True,
    skip_special_tokens=True,
)

thread = Thread(target=self.model.generate, kwargs={**inputs, "streamer": streamer})
thread.start()

for text in streamer:
    yield text

thread.join()

---

Category: Embeddings

Mean Pooling for Embeddings

What to check: Apply mean pooling with attention mask weighting

Search pattern:

rg "_mean_pooling|mean.*pooling" --type py runtimes/universal/models/

Pass criteria:

• Expand attention mask to match embedding dimensions
• Sum weighted by attention mask
• Clamp denominator to avoid division by zero

Severity: High

Good pattern:

def _mean_pooling(self, model_output, attention_mask):
    token_embeddings = model_output[0]
    input_mask_expanded = (
        attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    )
    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(
        input_mask_expanded.sum(1), min=1e-9
    )

---

L2 Normalize Embeddings

What to check: Normalize embeddings for cosine similarity

Search pattern:

rg "F\.normalize|normalize.*p=2" --type py runtimes/universal/models/

Pass criteria:

• Use F.normalize(embeddings, p=2, dim=1)
• Default to normalize=True for embedding endpoints

Severity: Medium

Good pattern:

if normalize:
    embeddings = F.normalize(embeddings, p=2, dim=1)
return embeddings.cpu().tolist()

---

Category: Flash Attention

Check Flash Attention Compatibility

What to check: Enable Flash Attention 2 only when supported

Search pattern:

rg "flash_attention|attn_implementation" --type py runtimes/universal/

Pass criteria:

• Only enable on CUDA devices
• Check torch version >= 2.0
• Check if flash_attn package is installed
• Set attn_implementation="flash_attention_2" in model kwargs

Severity: Medium

Good pattern:

def _supports_flash_attention(self, config: AutoConfig) -> bool:
    if self.device != "cuda":
        return False

    torch_version = tuple(map(int, torch.__version__.split(".")[:2]))
    if torch_version < (2, 0):
        return False

    try:
        import flash_attn
        return True
    except ImportError:
        return False

# In load():
if self._use_flash_attention and self._supports_flash_attention(config):
    model_kwargs["attn_implementation"] = "flash_attention_2"

---

Category: GGUF Models

Use llama-cpp-python's Chat Template

What to check: Use create_chat_completion() for GGUF models

Search pattern:

rg "create_chat_completion" --type py runtimes/universal/

Pass criteria:

• Use create_chat_completion() not manual prompt formatting
• GGUF metadata contains proper chat template
• Essential for models with special tokens (Qwen, Llama)

Severity: High

Good pattern:

# GGUF models use embedded chat template
result = self.llama.create_chat_completion(
    messages=messages,
    max_tokens=max_tokens,
    temperature=temperature,
    stream=False,
)
return result["choices"][0]["message"]["content"]

---

Configure GPU Layers for llama-cpp

What to check: Set n_gpu_layers based on device

Search pattern:

rg "n_gpu_layers" --type py runtimes/universal/

Pass criteria:

• n_gpu_layers=-1 for GPU (all layers)
• n_gpu_layers=0 for CPU only

Severity: High

Good pattern:

if self.device != "cpu":
    n_gpu_layers = -1  # All layers on GPU
else:
    n_gpu_layers = 0  # CPU only

self.llama = Llama(
    model_path=gguf_path,
    n_ctx=context_size,
    n_gpu_layers=n_gpu_layers,
)