Model Registry

from engine.models.loader import load_hf_model

class MyProcessor:
    def __init__(self, config, **kwargs):
        self._model = None
        self._tokenizer = None
        self.config = config

    def _ensure_model_loaded(self):
        if self._model is not None:
            return

        # Loads from HF (first time) or Ray cache (subsequent calls)
        cached = load_hf_model(
            hf_model_id="intfloat/multilingual-e5-large-instruct",
            model_class="AutoModel",
            tokenizer_class="AutoTokenizer",
            torch_dtype="float16",
        )

        # Instantiate from cached state dict
        from transformers import AutoModel, AutoConfig, AutoTokenizer

        config = AutoConfig.from_dict(cached["config"])
        self._model = AutoModel.from_config(config)
        self._model.load_state_dict(cached["state_dict"])
        self._model.to("cuda")
        self._model.eval()

        # Load tokenizer
        self._tokenizer = AutoTokenizer.from_pretrained(
            cached["tokenizer_config"]["tokenizer_dir"]
        )

    def __call__(self, batch):
        self._ensure_model_loaded()
        # Use self._model and self._tokenizer...

{
    "state_dict": dict,          # Model weights (load with model.load_state_dict())
    "config": dict,              # Model config (use with AutoConfig.from_dict())
    "tokenizer_config": {        # Tokenizer info (if tokenizer_class provided)
        "tokenizer_dir": str,    # Path to tokenizer files
        "class": str,            # Tokenizer class name
    },
    "model_class": str,          # The model class used
    "hf_model_id": str,          # The HuggingFace model ID
}

from engine.models.loader import load_hf_model_async
import ray

# Start loading both models in parallel
refs = [
    load_hf_model_async("intfloat/multilingual-e5-large-instruct"),
    load_hf_model_async("openai/clip-vit-large-patch14"),
]

# Wait for both to complete
e5_data, clip_data = ray.get(refs)

from engine.models.loader import is_hf_model_cached

if is_hf_model_cached("intfloat/multilingual-e5-large-instruct"):
    print("Model already cached in cluster")

from engine.models.lazy import LazyModelMixin
from engine.inference.services import BaseBatchInferenceService

class MyEmbeddingProcessor(LazyModelMixin, BaseBatchInferenceService):
    # Configure model via class attributes
    model_id = "intfloat/multilingual-e5-large-instruct"
    model_class = "AutoModel"
    tokenizer_class = "AutoTokenizer"
    torch_dtype = "float16"
    model_source = "huggingface"  # or "namespace" for S3 models

    def __init__(self, config, **kwargs):
        super().__init__(**kwargs)
        self.config = config
        # Don't load model here - LazyModelMixin handles it

    def _process_batch(self, batch):
        # Model is automatically loaded on first call
        model, tokenizer = self.get_model()

        # Process batch
        inputs = tokenizer(
            batch["text"].tolist(),
            padding=True,
            truncation=True,
            return_tensors="pt",
        )

        with torch.no_grad():
            outputs = model(**inputs)

        batch["embedding"] = outputs.last_hidden_state.mean(dim=1).tolist()
        return batch

class MyCustomProcessor(LazyModelMixin, BaseBatchInferenceService):
    model_id = "my-custom-model"
    model_source = "namespace"  # Load from S3

    def _instantiate_model(self, cached_data):
        """Custom instantiation for non-standard models."""
        import torch

        # Create your model architecture
        model = torch.nn.Sequential(
            torch.nn.Linear(768, 512),
            torch.nn.ReLU(),
            torch.nn.Linear(512, 256),
        )

        # Load weights from cached data
        model.load_state_dict(cached_data)
        model.to(self._detect_device())
        model.eval()

        return model, None  # No tokenizer

from engine.models.lazy import lazy_model

class MyProcessor:
    def __init__(self, config, **kwargs):
        self.config = config
        self._model = None
        self._tokenizer = None

    @lazy_model(
        model_id="intfloat/multilingual-e5-large-instruct",
        model_class="AutoModel",
        tokenizer_class="AutoTokenizer",
        torch_dtype="float16",
    )
    def __call__(self, batch):
        # self._model and self._tokenizer are automatically available
        inputs = self._tokenizer(
            batch["text"].tolist(),
            return_tensors="pt",
            padding=True,
        )

        with torch.no_grad():
            outputs = self._model(**inputs)

        batch["embedding"] = outputs.last_hidden_state.mean(dim=1).tolist()
        return batch

┌─────────────┐     ┌─────────────┐     ┌─────────────┐
│   Upload    │────▶│   Deploy    │────▶│    Use      │
│  (S3 store) │     │ (Ray cache) │     │ (Inference) │
└─────────────┘     └─────────────┘     └─────────────┘

# Create archive
tar -czvf my_model.tar.gz ./model_weights/

# Upload
curl -X POST "$MIXPEEK_API_URL/v1/namespaces/$NAMESPACE_ID/models" \
  -H "Authorization: Bearer $MIXPEEK_API_KEY" \
  -F "file=@my_model.tar.gz" \
  -F "name=my-embedding-model" \
  -F "version=1.0.0" \
  -F "model_format=pytorch" \
  -F "framework=pytorch" \
  -F "task_type=embedding" \
  -F "num_gpus=0" \
  -F "memory_gb=4.0"

curl -X POST "$MIXPEEK_API_URL/v1/namespaces/$NAMESPACE_ID/models/my-embedding-model_1_0_0/deploy" \
  -H "Authorization: Bearer $MIXPEEK_API_KEY"

from engine.models.lazy import LazyModelMixin

class MyCustomProcessor(LazyModelMixin, BaseBatchInferenceService):
    model_id = "my-embedding-model_1_0_0"
    model_source = "namespace"  # Load from S3

    def _instantiate_model(self, weights):
        """Custom instantiation for S3 models."""
        import torch

        model = torch.nn.Linear(768, 256)
        model.load_state_dict(weights)
        model.to(self._detect_device())
        model.eval()
        return model, None

    def _process_batch(self, batch):
        model, _ = self.get_model()
        # Use model...

from engine.models.loader import load_namespace_model

class MyCustomProcessor:
    def __init__(self, config):
        self._model = None

    def _ensure_model_loaded(self):
        if self._model is not None:
            return

        # Load from S3 (cached in Ray object store)
        weights = load_namespace_model("my-embedding-model_1_0_0")

        self._model = torch.nn.Linear(768, 256)
        self._model.load_state_dict(weights)
        self._model.eval()

    def __call__(self, batch):
        self._ensure_model_loaded()
        # Use self._model...

Need a model?
├── Is it a common task (embeddings, transcription, etc.)?
│   └── Use built-in models
├── Is it a standard HuggingFace model?
│   ├── Simple usage → LazyModelMixin or @lazy_model
│   └── Custom logic → load_hf_model() + manual instantiation
└── Is it custom/fine-tuned weights?
    └── Upload to namespace → LazyModelMixin with model_source="namespace"

my-embedding-model_1_0_0  (production)
my-embedding-model_1_1_0  (staging)
my-embedding-model_2_0_0  (development)

# GPU model with high memory
-F "num_gpus=1" \
-F "memory_gb=16.0"

# CPU-only lightweight model
-F "num_cpus=0.5" \
-F "num_gpus=0" \
-F "memory_gb=2.0"

from mixpeek import Mixpeek

client = Mixpeek(api_key="sk_...")

# Upload model
with open("my_model.tar.gz", "rb") as f:
    result = client.models.upload(
        namespace_id="ns_abc123",
        file=f,
        name="my-reranker",
        version="1.0.0",
        model_format="pytorch",
        task_type="reranking",
    )

# Deploy model
client.models.deploy(
    namespace_id="ns_abc123",
    model_id=result.model_id,
)

# List models
models = client.models.list(namespace_id="ns_abc123")

# Wrong - loads at __init__
def __init__(self):
    self._model = load_hf_model(...)  # Don't do this

# Right - loads on first use
def _process_batch(self, batch):
    self._ensure_model_loaded()  # Loads here