thalamus-serve 0.2.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
• 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
/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

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.