Universal Runtime best practices for PyTorch inference, Transformers models, and FastAPI serving. Covers device management, model loading, memory optimization, and performance tuning.
Best practices and code review checklists for the Universal Runtime - LlamaFarm's local ML inference server.
The Universal Runtime provides OpenAI-compatible endpoints for HuggingFace models:
Directory: runtimes/universal/
Python: 3.11+
Key Dependencies: PyTorch, Transformers, FastAPI, llama-cpp-python
This skill extends the shared Python practices. Always apply these first:
| Topic | File | Priority |
|---|---|---|
| Patterns | python-skills/patterns.md | Medium |
| Async | python-skills/async.md | High |
| Typing | python-skills/typing.md | Medium |
| Testing | python-skills/testing.md | Medium |
| Errors | python-skills/error-handling.md | High |
| Security | python-skills/security.md | Critical |
| Topic | File | Key Points |
|---|---|---|
| PyTorch | pytorch.md | Device management, dtype, memory cleanup |
| Transformers | transformers.md | Model loading, tokenization, inference |
| FastAPI | fastapi.md | API design, streaming, lifespan |
| Performance | performance.md | Batching, caching, optimizations |
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/
_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)
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)
_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()
# 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)
When reviewing Universal Runtime code:
Critical - Security
High - Memory & Device
Medium - Performance
Low - Code Style