sheaf-serve 0.1.0__tar.gz

sheaf_serve-0.1.0/.github/workflows/ci.yml

@@ -0,0 +1,30 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v5
+      - run: uv sync --extra dev
+      - run: uv run ruff check src/ tests/
+      - run: uv run ruff format --check src/ tests/
+
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.10", "3.11", "3.12"]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+      - run: uv sync --extra dev
+      - run: uv run pytest tests/ -v

sheaf_serve-0.1.0/.github/workflows/publish.yml

@@ -0,0 +1,18 @@
+name: Publish to PyPI
+
+on:
+  push:
+    tags:
+      - "v*"
+
+jobs:
+  publish:
+    runs-on: ubuntu-latest
+    environment: pypi
+    permissions:
+      id-token: write  # required for trusted publishing
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v5
+      - run: uv build
+      - uses: pypa/gh-action-pypi-publish@release/v1

sheaf_serve-0.1.0/.gitignore

@@ -0,0 +1,13 @@
+__pycache__/
+*.py[cod]
+*.egg-info/
+.eggs/
+dist/
+build/
+.venv/
+venv/
+.mypy_cache/
+.ruff_cache/
+.pytest_cache/
+*.so
+.env

sheaf_serve-0.1.0/PKG-INFO

@@ -0,0 +1,52 @@
+Metadata-Version: 2.4
+Name: sheaf-serve
+Version: 0.1.0
+Summary: Unified serving layer for non-text foundation models
+Author-email: Alex Korbonits <alexkorbonits@gmail.com>
+License: Apache-2.0
+Keywords: foundation-models,inference,mlops,ray,serving
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Requires-Dist: numpy>=1.24.0
+Requires-Dist: pydantic>=2.0.0
+Requires-Dist: ray[serve]>=2.10.0
+Provides-Extra: all
+Requires-Dist: chronos-forecasting>=1.0.0; extra == 'all'
+Requires-Dist: fair-esm>=2.0.0; extra == 'all'
+Requires-Dist: feast>=0.40.0; extra == 'all'
+Requires-Dist: pymilvus>=2.4.0; extra == 'all'
+Requires-Dist: tabpfn>=2.0.0; extra == 'all'
+Requires-Dist: timesfm>=1.0.0; extra == 'all'
+Provides-Extra: dev
+Requires-Dist: mypy>=1.9.0; extra == 'dev'
+Requires-Dist: pre-commit>=3.7.0; extra == 'dev'
+Requires-Dist: pytest-asyncio>=0.23.0; extra == 'dev'
+Requires-Dist: pytest>=8.0.0; extra == 'dev'
+Requires-Dist: ruff>=0.4.0; extra == 'dev'
+Provides-Extra: feast
+Requires-Dist: feast>=0.40.0; extra == 'feast'
+Provides-Extra: milvus
+Requires-Dist: pymilvus>=2.4.0; extra == 'milvus'
+Provides-Extra: molecular
+Requires-Dist: fair-esm>=2.0.0; extra == 'molecular'
+Provides-Extra: tabular
+Requires-Dist: tabpfn>=2.0.0; extra == 'tabular'
+Provides-Extra: time-series
+Requires-Dist: chronos-forecasting>=1.0.0; extra == 'time-series'
+Requires-Dist: timesfm>=1.0.0; extra == 'time-series'
+Description-Content-Type: text/markdown
+
+# Sheaf
+
+Unified serving layer for non-text foundation models.
+
+```bash
+pip install sheaf-serve
+```

sheaf_serve-0.1.0/README.md

@@ -0,0 +1,7 @@
+# Sheaf
+
+Unified serving layer for non-text foundation models.
+
+```bash
+pip install sheaf-serve
+```

sheaf_serve-0.1.0/pyproject.toml

@@ -0,0 +1,75 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "sheaf-serve"
+version = "0.1.0"
+description = "Unified serving layer for non-text foundation models"
+readme = "README.md"
+requires-python = ">=3.10"
+license = { text = "Apache-2.0" }
+authors = [{ name = "Alex Korbonits", email = "alexkorbonits@gmail.com" }]
+keywords = ["serving", "inference", "foundation-models", "ray", "mlops"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: Apache Software License",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+    "ray[serve]>=2.10.0",
+    "pydantic>=2.0.0",
+    "numpy>=1.24.0",
+]
+
+[project.optional-dependencies]
+time-series = [
+    "chronos-forecasting>=1.0.0",
+    "timesfm>=1.0.0",
+]
+tabular = [
+    "tabpfn>=2.0.0",
+]
+molecular = [
+    "fair-esm>=2.0.0",
+]
+feast = [
+    "feast>=0.40.0",
+]
+milvus = [
+    "pymilvus>=2.4.0",
+]
+dev = [
+    "pytest>=8.0.0",
+    "pytest-asyncio>=0.23.0",
+    "ruff>=0.4.0",
+    "mypy>=1.9.0",
+    "pre-commit>=3.7.0",
+]
+all = [
+    "sheaf-serve[time-series,tabular,molecular,feast,milvus]",
+]
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/sheaf"]
+
+[tool.ruff]
+line-length = 88
+src = ["src"]
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "UP"]
+
+[tool.mypy]
+python_version = "3.10"
+strict = true
+files = ["src/sheaf"]
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"
+testpaths = ["tests"]

sheaf_serve-0.1.0/src/sheaf/__init__.py

@@ -0,0 +1,7 @@
+"""Sheaf — unified serving layer for non-text foundation models."""
+
+from sheaf.server import ModelServer
+from sheaf.spec import ModelSpec
+
+__version__ = "0.1.0"
+__all__ = ["ModelServer", "ModelSpec"]

sheaf_serve-0.1.0/src/sheaf/api/__init__.py

@@ -0,0 +1,10 @@
+from sheaf.api.base import BaseRequest, BaseResponse, ModelType
+from sheaf.api.time_series import TimeSeriesRequest, TimeSeriesResponse
+
+__all__ = [
+    "BaseRequest",
+    "BaseResponse",
+    "ModelType",
+    "TimeSeriesRequest",
+    "TimeSeriesResponse",
+]

sheaf_serve-0.1.0/src/sheaf/api/base.py

@@ -0,0 +1,35 @@
+"""Base request/response contracts and model type registry."""
+
+from enum import Enum
+from typing import Any
+from uuid import UUID, uuid4
+
+from pydantic import BaseModel, Field
+
+
+class ModelType(str, Enum):
+    TIME_SERIES = "time_series"
+    TABULAR = "tabular"
+    MOLECULAR = "molecular"
+    GEOSPATIAL = "geospatial"
+    DIFFUSION = "diffusion"
+    NEURAL_OPERATOR = "neural_operator"
+    AUDIO = "audio"
+    EMBEDDING = "embedding"
+
+
+class BaseRequest(BaseModel):
+    request_id: UUID = Field(default_factory=uuid4)
+    model_type: ModelType
+    model_name: str
+
+    model_config = {"arbitrary_types_allowed": True}
+
+
+class BaseResponse(BaseModel):
+    request_id: UUID
+    model_type: ModelType
+    model_name: str
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+    model_config = {"arbitrary_types_allowed": True}

sheaf_serve-0.1.0/src/sheaf/api/time_series.py

@@ -0,0 +1,80 @@
+"""API contract for time series foundation models (Chronos2, TimesFM, etc.)."""
+
+from enum import Enum
+from typing import Annotated, Any
+from uuid import UUID
+
+import numpy as np
+from pydantic import Field, model_validator
+
+from sheaf.api.base import BaseRequest, BaseResponse, ModelType
+
+
+class Frequency(str, Enum):
+    MINUTELY = "1min"
+    FIVE_MINUTELY = "5min"
+    FIFTEEN_MINUTELY = "15min"
+    HOURLY = "1h"
+    DAILY = "1d"
+    WEEKLY = "1W"
+    MONTHLY = "1M"
+
+
+class OutputMode(str, Enum):
+    MEAN = "mean"
+    QUANTILES = "quantiles"
+    SAMPLES = "samples"
+
+
+class TimeSeriesRequest(BaseRequest):
+    """Request contract for time series foundation models.
+
+    Either `history` (raw values) or `feature_ref` (Feast entity reference)
+    must be provided, not both.
+    """
+
+    model_type: ModelType = ModelType.TIME_SERIES
+
+    # Input: raw history or feature store reference
+    history: Annotated[list[float] | None, Field(default=None)]
+    feature_ref: Annotated[dict[str, str] | None, Field(default=None)]
+    # e.g. {"feature_view": "asset_prices", "entity_id": "AAPL"}
+
+    horizon: Annotated[int, Field(gt=0)]
+    frequency: Frequency
+    output_mode: OutputMode = OutputMode.MEAN
+    quantile_levels: list[float] = Field(
+        default=[0.1, 0.5, 0.9],
+        description="Quantile levels — only used when output_mode=quantiles",
+    )
+    num_samples: int = Field(
+        default=20,
+        description="Number of samples — only used when output_mode=samples",
+    )
+
+    @model_validator(mode="after")
+    def validate_input_source(self) -> "TimeSeriesRequest":
+        if self.history is None and self.feature_ref is None:
+            raise ValueError("One of `history` or `feature_ref` must be provided.")
+        if self.history is not None and self.feature_ref is not None:
+            raise ValueError("Provide either `history` or `feature_ref`, not both.")
+        return self
+
+
+class TimeSeriesResponse(BaseResponse):
+    """Response contract for time series foundation models."""
+
+    model_type: ModelType = ModelType.TIME_SERIES
+
+    # Mean forecast — always populated
+    mean: list[float]
+
+    # Populated when output_mode=quantiles
+    quantiles: dict[str, list[float]] | None = None
+    # e.g. {"0.1": [...], "0.5": [...], "0.9": [...]}
+
+    # Populated when output_mode=samples
+    samples: list[list[float]] | None = None
+
+    horizon: int
+    frequency: str

sheaf_serve-0.1.0/src/sheaf/backends/__init__.py

@@ -0,0 +1,4 @@
+from sheaf.backends.base import ModelBackend
+from sheaf.backends.chronos import Chronos2Backend
+
+__all__ = ["ModelBackend", "Chronos2Backend"]

sheaf_serve-0.1.0/src/sheaf/backends/base.py

@@ -0,0 +1,41 @@
+"""Abstract base class for model backends."""
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+from sheaf.api.base import BaseRequest, BaseResponse
+
+
+class ModelBackend(ABC):
+    """Pluggable model backend.
+
+    Implement this to add a new model to Sheaf. The backend owns:
+    - Model loading and initialization
+    - Preprocessing raw inputs
+    - Running inference
+    - Postprocessing outputs into a typed response
+    """
+
+    @abstractmethod
+    def load(self) -> None:
+        """Load model weights and initialize any required state."""
+        ...
+
+    @abstractmethod
+    def predict(self, request: BaseRequest) -> BaseResponse:
+        """Run inference for a single request."""
+        ...
+
+    def batch_predict(self, requests: list[BaseRequest]) -> list[BaseResponse]:
+        """Run inference over a batch.
+
+        Default implementation runs requests sequentially.
+        Override to implement model-type-aware batching.
+        """
+        return [self.predict(r) for r in requests]
+
+    @property
+    @abstractmethod
+    def model_type(self) -> str:
+        """The ModelType this backend serves."""
+        ...

sheaf_serve-0.1.0/src/sheaf/backends/chronos.py

@@ -0,0 +1,143 @@
+"""Chronos2 backend for time series forecasting.
+
+Requires: pip install sheaf-serve[time-series]
+Supports: amazon/chronos-t5-{tiny,mini,small,base,large}
+          amazon/chronos-bolt-{tiny,mini,small,base}
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+
+from sheaf.api.base import BaseRequest, BaseResponse, ModelType
+from sheaf.api.time_series import OutputMode, TimeSeriesRequest, TimeSeriesResponse
+from sheaf.backends.base import ModelBackend
+from sheaf.registry import register_backend
+
+if TYPE_CHECKING:
+    import torch
+
+
+@register_backend("chronos2")
+class Chronos2Backend(ModelBackend):
+    """ModelBackend implementation for Chronos / Chronos-Bolt models.
+
+    Args:
+        model_id: HuggingFace model ID, e.g. "amazon/chronos-bolt-small"
+        device_map: "cpu", "cuda", "mps", or "auto"
+        torch_dtype: "bfloat16", "float32", etc. Passed to from_pretrained.
+        num_samples: Default number of samples for probabilistic output.
+    """
+
+    def __init__(
+        self,
+        model_id: str = "amazon/chronos-bolt-small",
+        device_map: str = "cpu",
+        torch_dtype: str = "bfloat16",
+        num_samples: int = 20,
+    ) -> None:
+        self._model_id = model_id
+        self._device_map = device_map
+        self._torch_dtype = torch_dtype
+        self._default_num_samples = num_samples
+        self._pipeline: Any = None
+
+    @property
+    def model_type(self) -> str:
+        return ModelType.TIME_SERIES
+
+    def load(self) -> None:
+        try:
+            import torch
+            from chronos import BaseChronosPipeline
+        except ImportError as e:
+            raise ImportError(
+                "chronos-forecasting is required for the Chronos2 backend. "
+                "Install it with: pip install sheaf-serve[time-series]"
+            ) from e
+
+        dtype_map = {
+            "bfloat16": torch.bfloat16,
+            "float32": torch.float32,
+            "float16": torch.float16,
+        }
+        torch_dtype = dtype_map.get(self._torch_dtype, torch.bfloat16)
+
+        self._pipeline = BaseChronosPipeline.from_pretrained(
+            self._model_id,
+            device_map=self._device_map,
+            torch_dtype=torch_dtype,
+        )
+
+    def predict(self, request: BaseRequest) -> BaseResponse:
+        if not isinstance(request, TimeSeriesRequest):
+            raise TypeError(f"Expected TimeSeriesRequest, got {type(request)}")
+        return self._run([request])[0]
+
+    def batch_predict(self, requests: list[BaseRequest]) -> list[BaseResponse]:
+        ts_requests = [r for r in requests if isinstance(r, TimeSeriesRequest)]
+        if len(ts_requests) != len(requests):
+            raise TypeError("All requests must be TimeSeriesRequest for Chronos2Backend")
+        return self._run(ts_requests)
+
+    def _run(self, requests: list[TimeSeriesRequest]) -> list[TimeSeriesResponse]:
+        import torch
+
+        if self._pipeline is None:
+            raise RuntimeError("Backend not loaded. Call load() first.")
+
+        # Build context tensors — pad to same length within batch
+        histories = [r.history or [] for r in requests]
+        max_len = max(len(h) for h in histories)
+        padded = [
+            [float("nan")] * (max_len - len(h)) + h
+            for h in histories
+        ]
+        context = torch.tensor(padded, dtype=torch.float32)
+
+        # Bucket by horizon: all requests in a batch must share a horizon
+        # (caller is responsible for bucketing — see BatchPolicy(bucket_by="horizon"))
+        horizon = requests[0].horizon
+        num_samples = requests[0].num_samples
+
+        # predict returns [batch, num_samples, horizon]
+        forecast = self._pipeline.predict(
+            context=context,
+            prediction_length=horizon,
+            num_samples=num_samples,
+        )
+        forecast_np: np.ndarray = forecast.numpy()
+
+        responses = []
+        for i, req in enumerate(requests):
+            samples = forecast_np[i]  # [num_samples, horizon]
+            responses.append(self._build_response(req, samples))
+        return responses
+
+    def _build_response(
+        self, req: TimeSeriesRequest, samples: np.ndarray
+    ) -> TimeSeriesResponse:
+        mean = samples.mean(axis=0).tolist()
+
+        quantiles = None
+        if req.output_mode == OutputMode.QUANTILES:
+            quantiles = {
+                str(q): np.quantile(samples, q, axis=0).tolist()
+                for q in req.quantile_levels
+            }
+
+        raw_samples = None
+        if req.output_mode == OutputMode.SAMPLES:
+            raw_samples = samples.tolist()
+
+        return TimeSeriesResponse(
+            request_id=req.request_id,
+            model_name=req.model_name,
+            horizon=req.horizon,
+            frequency=req.frequency.value,
+            mean=mean,
+            quantiles=quantiles,
+            samples=raw_samples,
+        )

sheaf_serve-0.1.0/src/sheaf/registry.py

@@ -0,0 +1,24 @@
+"""Backend registry — separated to avoid circular imports."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from sheaf.backends.base import ModelBackend
+
+_BACKEND_REGISTRY: dict[str, type["ModelBackend"]] = {}
+
+
+def register_backend(name: str):
+    """Decorator to register a ModelBackend implementation by name.
+
+    Example:
+        @register_backend("chronos2")
+        class Chronos2Backend(ModelBackend):
+            ...
+    """
+    def decorator(cls: type["ModelBackend"]) -> type["ModelBackend"]:
+        _BACKEND_REGISTRY[name] = cls
+        return cls
+    return decorator

sheaf_serve-0.1.0/src/sheaf/scheduling/__init__.py

@@ -0,0 +1,3 @@
+from sheaf.scheduling.batch import BatchPolicy
+
+__all__ = ["BatchPolicy"]

sheaf_serve-0.1.0/src/sheaf/scheduling/batch.py

@@ -0,0 +1,17 @@
+"""Batching policies for model-type-aware request scheduling."""
+
+from pydantic import BaseModel, Field
+
+
+class BatchPolicy(BaseModel):
+    """Controls how requests are batched before hitting the model backend.
+
+    max_batch_size: hard cap on requests per batch
+    timeout_ms: max time to wait for a full batch before flushing
+    bucket_by: field name to bucket on before batching (e.g. "horizon" for
+               time series, so variable-length forecasts don't get mixed)
+    """
+
+    max_batch_size: int = Field(default=32, gt=0)
+    timeout_ms: int = Field(default=50, gt=0)
+    bucket_by: str | None = None

sheaf_serve-0.1.0/src/sheaf/server.py

@@ -0,0 +1,71 @@
+"""ModelServer — the Ray Serve entry point for Sheaf."""
+
+from typing import Any
+
+import ray
+from ray import serve
+
+from sheaf.api.base import BaseRequest, BaseResponse
+from sheaf.backends.base import ModelBackend
+from sheaf.registry import _BACKEND_REGISTRY, register_backend  # noqa: F401
+from sheaf.spec import ModelSpec
+
+
+@serve.deployment
+class _SheafDeployment:
+    def __init__(self, spec: ModelSpec) -> None:
+        backend_cls = _BACKEND_REGISTRY.get(spec.backend)
+        if backend_cls is None:
+            raise ValueError(
+                f"Unknown backend '{spec.backend}'. "
+                f"Registered backends: {list(_BACKEND_REGISTRY)}"
+            )
+        self._backend: ModelBackend = backend_cls(**spec.backend_kwargs)
+        self._backend.load()
+        self._spec = spec
+
+    async def __call__(self, request: BaseRequest) -> BaseResponse:
+        return self._backend.predict(request)
+
+
+class ModelServer:
+    """Top-level serving orchestrator.
+
+    Example:
+        server = ModelServer(
+            models=[chronos_spec, tabpfn_spec],
+        )
+        server.run()
+    """
+
+    def __init__(
+        self,
+        models: list[ModelSpec],
+        host: str = "0.0.0.0",
+        port: int = 8000,
+    ) -> None:
+        self._models = models
+        self._host = host
+        self._port = port
+        self._deployments: dict[str, Any] = {}
+
+    def run(self) -> None:
+        if not ray.is_initialized():
+            ray.init()
+
+        serve.start(http_options={"host": self._host, "port": self._port})
+
+        for spec in self._models:
+            deployment = _SheafDeployment.options(
+                name=spec.name,
+                num_replicas=spec.resources.replicas,
+                ray_actor_options={
+                    "num_cpus": spec.resources.num_cpus,
+                    "num_gpus": spec.resources.num_gpus,
+                },
+            ).bind(spec)
+            handle = serve.run(deployment, name=spec.name, route_prefix=f"/{spec.name}")
+            self._deployments[spec.name] = handle
+
+    def shutdown(self) -> None:
+        serve.shutdown()

sheaf_serve-0.1.0/src/sheaf/spec.py

@@ -0,0 +1,35 @@
+"""ModelSpec — declares what Sheaf should serve."""
+
+from pydantic import BaseModel, Field
+
+from sheaf.api.base import ModelType
+from sheaf.scheduling.batch import BatchPolicy
+
+
+class ResourceConfig(BaseModel):
+    num_cpus: float = 1.0
+    num_gpus: float = 0.0
+    memory_gb: float | None = None
+    replicas: int = 1
+
+
+class ModelSpec(BaseModel):
+    """Declares a model to be served by Sheaf.
+
+    Example:
+        spec = ModelSpec(
+            name="chronos2-small",
+            model_type=ModelType.TIME_SERIES,
+            backend="chronos2",
+            backend_kwargs={"model_size": "small"},
+            resources=ResourceConfig(num_gpus=1, replicas=2),
+            batch_policy=BatchPolicy(max_batch_size=64, bucket_by="horizon"),
+        )
+    """
+
+    name: str
+    model_type: ModelType
+    backend: str
+    backend_kwargs: dict = Field(default_factory=dict)
+    resources: ResourceConfig = Field(default_factory=ResourceConfig)
+    batch_policy: BatchPolicy = Field(default_factory=BatchPolicy)