thalamus-serve 0.3.0

A Python ML model serving framework with built-in observability, caching, and GPU management

thalamus-serve

A Python ML model serving framework built on FastAPI with built-in observability, caching, and GPU management.

Features

• Simple Decorator API - Register models with @app.model() decorator
• Explicit Schema Types - Input/output types specified via input_type and output_type parameters
• Built-in Observability - Prometheus metrics, structured logging
• Weight Management - Automatic fetching from S3, HuggingFace Hub, or HTTP
• GPU Management - Automatic device detection and allocation (CUDA/MPS/CPU)
• Caching - LRU cache for model weights with configurable size
• Batch Processing - Native batch inference support
• Health Checks - /health, /ready, /status endpoints
• Capacity Reporting - /capacity advertises free slots and ideal batch size
• API Key Authentication - Protect endpoints with API key middleware

Installation

pip install thalamus-serve

# With GPU/PyTorch support
pip install thalamus-serve[gpu]

# For development
pip install thalamus-serve[dev]

Quick Start

from pathlib import Path
from pydantic import BaseModel
from thalamus_serve import Thalamus


class TextInput(BaseModel):
    text: str


class SentimentOutput(BaseModel):
    label: str
    confidence: float


app = Thalamus()


@app.model(
    model_id="sentiment",
    version="1.0.0",
    description="Sentiment analysis model",
    default=True,
    input_type=TextInput,
    output_type=SentimentOutput,
)
class SentimentModel:
    def load(self, weights: dict[str, Path], device: str) -> None:
        # Load your model weights here
        pass

    def predict(self, inputs: list[TextInput]) -> list[SentimentOutput]:
        return [
            SentimentOutput(label="positive", confidence=0.95)
            for _ in inputs
        ]


if __name__ == "__main__":
    app.serve(host="0.0.0.0", port=8000)

Run the server:

export THALAMUS_API_KEY=your-secret-key
python app.py

Test the endpoint:

curl -X POST http://localhost:8000/predict \
  -H "Content-Type: application/json" \
  -H "X-API-Key: your-secret-key" \
  -d '{"inputs": [{"text": "I love this!"}]}'

Examples

See the examples/ directory for complete implementations:

• vanilla/ - Basic example with minimal setup
• torch/ - PyTorch image classification model
• hf_model/ - HuggingFace Transformers integration
• sklearn/ - scikit-learn model serving
• deploy/ - Docker deployment with Docker Compose

API Endpoints

Endpoint	Method	Auth	Description
/health	GET	No	Service health status
/ready	GET	No	Readiness check (critical models loaded)
/status	GET	No	Detailed status with cache/GPU info
/metrics	GET	No	Prometheus metrics
/capacity	GET	Yes	Remaining request slots and ideal batch size
/schema	GET	Yes	List all model schemas
/schema/{model_id}	GET	Yes	Get specific model schema
/predict	POST	Yes	Run inference
/cache/clear	POST	Yes	Clear weight cache
/models/{model_id}/unload	POST	Yes	Unload model from memory

Model Interface

class MyModel:
    def load(self, weights: dict[str, Path], device: str) -> None:
        """Called during startup to load weights."""
        pass

    def predict(self, inputs: list[InputType]) -> list[OutputType]:
        """Required. Runs inference on a batch."""
        pass

    # Optional hooks
    def preprocess(self, inputs: list[InputType]) -> list[Any]: ...
    def postprocess(self, outputs: list[Any]) -> list[OutputType]: ...

    @property
    def is_ready(self) -> bool:
        """Optional. Used by /ready endpoint."""
        return True

    def capacity(self) -> ModelCapacity:
        """Optional. Used by /capacity endpoint."""
        ...

Capacity

/capacity tells a caller how much work this deployment will accept right now, so a batching client can size its next dispatch instead of guessing. Declare the numbers on the decorator:

@app.model(
    model_id="doc-type",
    input_type=Input,
    output_type=Output,
    max_batch_size=32,
    ideal_batch_size=16,
    max_concurrent_requests=2,
)
class DocTypeClassifier:
    ...

Argument	Default	Meaning
max_batch_size	1	Hard cap on inputs accepted in one /predict call
ideal_batch_size	max_batch_size	Throughput sweet spot
max_concurrent_requests	1	Parallel /predict calls the pod tolerates

The defaults are deliberately conservative — a model that has not declared anything advertises a batch size of 1. ideal_batch_size must be between 1 and max_batch_size, and both must be at least 1; violating that raises ValueError at import time.

A model that can measure its own headroom overrides the static numbers with a capacity() method:

from thalamus_serve import ModelCapacity, get_gpu_memory

class DocTypeClassifier:
    def capacity(self) -> ModelCapacity:
        used_mb, total_mb = get_gpu_memory(self.device) or (0.0, 1.0)
        headroom = 1.0 - (used_mb / total_mb)
        return ModelCapacity(
            accepting=headroom > 0.1,
            remaining_requests=2 if headroom > 0.4 else 1,
            ideal_batch_size=16 if headroom > 0.4 else 4,
            max_batch_size=32,
            reason=None if headroom > 0.1 else "vram_headroom_low",
        )

The hook is polled before every dispatch, so it must be O(1) — read a cached gauge, never run inference. A hook that raises is reported as accepting=false with reason="capacity_hook_error"; it never fails the request.

The response aggregates across models. Top-level accepting is the AND over critical models, so an unloaded non-critical model does not mark the pod unavailable. Top-level remaining_requests is the minimum across accepting models — the bottleneck — and is 0 whenever accepting is false, so the two can never disagree. Per-model slot counts stay visible in models regardless. Querying /capacity never triggers a lazy model load; an unloaded model reports reason="model_not_loaded".

{
  "accepting": true,
  "remaining_requests": 2,
  "models": {
    "doc-type@2.1.0": {
      "accepting": true,
      "remaining_requests": 2,
      "ideal_batch_size": 16,
      "max_batch_size": 32,
      "reason": null
    }
  },
  "inflight_requests": 0,
  "uptime_seconds": 128.4
}

Configuration

Weight Sources

Models can load weights from multiple sources by specifying them directly in the decorator:

from pydantic import BaseModel
from thalamus_serve import Thalamus, HFWeight, S3Weight, HTTPWeight

app = Thalamus()

class MyInput(BaseModel):
    text: str

class MyOutput(BaseModel):
    embedding: list[float]

@app.model(
    model_id="my-model",
    version="1.0.0",
    device="cuda:0",
    input_type=MyInput,
    output_type=MyOutput,
    weights={
        "model": HFWeight(repo="bert-base-uncased"),
        "tokenizer": S3Weight(bucket="my-bucket", key="models/tokenizer.json"),
    },
)
class MyModel:
    def load(self, weights: dict[str, Path], device: str) -> None:
        # weights["model"] is a directory path (full repo)
        # weights["tokenizer"] is a file path
        pass

Supported weight sources:

Source	Single File	Directory/Sharded
HuggingFace	HFWeight(repo="...", filename="model.bin")	HFWeight(repo="...")
S3	S3Weight(bucket="...", key="path/model.pt")	S3Weight(bucket="...", prefix="path/shards/")
HTTP	HTTPWeight(urls=["https://.../model.pt"])	HTTPWeight(urls=["https://.../shard1.pt", "https://.../shard2.pt"])

Examples:

# HuggingFace - single file
HFWeight(repo="bert-base-uncased", filename="pytorch_model.bin", revision="main")

# HuggingFace - full repo snapshot (for sharded models)
HFWeight(repo="meta-llama/Llama-2-7b-hf")

# S3 - single file
S3Weight(bucket="my-models", key="bert/model.pt", region="us-east-1")

# S3 - directory prefix (downloads all files under prefix)
S3Weight(bucket="my-models", prefix="llama/shards/")

# HTTP - single file
HTTPWeight(urls=["https://example.com/model.pt"])

# HTTP - multiple files (sharded)
HTTPWeight(urls=[
    "https://example.com/model-00001.pt",
    "https://example.com/model-00002.pt",
])

Environment Variables

Variable	Default	Description
THALAMUS_API_KEY	-	API key for protected endpoints
THALAMUS_LOG_LEVEL	INFO	Logging level
THALAMUS_CACHE_DIR	/tmp/thalamus	Weight cache directory
THALAMUS_CACHE_MAX_GB	50	Maximum cache size in GB
HF_TOKEN	-	HuggingFace authentication token
AWS_ACCESS_KEY_ID	-	AWS credentials for S3
AWS_SECRET_ACCESS_KEY	-	AWS credentials for S3

Built-in Schemas

Common ML types available from the package:

from thalamus_serve import (
    Base64Data,  # Base64-encoded binary data
    BBox,        # Bounding box (x1, y1, x2, y2)
    Label,       # Classification label with confidence
    Vector,      # Embedding vector
    Span,        # Text span with optional label
    Prob,        # Probability value (0-1)
    S3Ref,       # S3 reference (bucket, key)
    Url,         # URL type
)

Contributing

We welcome contributions! Please see CONTRIBUTING.md for guidelines.

License

Apache License 2.0. See LICENSE for details.