inferential 0.1.0__py3-none-any.whl

inferential/__init__.py

@@ -0,0 +1,14 @@
+from inferential._version import __version__
+from inferential.client import Connection, Model
+from inferential.scheduler.base import register_policy
+from inferential.scheduler.request import InferenceRequest
+from inferential.server import Server
+
+__all__ = [
+    "__version__",
+    "Connection",
+    "InferenceRequest",
+    "Model",
+    "Server",
+    "register_policy",
+]

inferential/_version.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

inferential/client.py

@@ -0,0 +1,156 @@
+from __future__ import annotations
+
+import time
+from typing import Any
+
+import numpy as np
+import zmq
+import zmq.asyncio
+
+from inferential.proto import (
+    RAW,
+    Client,
+    ModelResponse,
+    Observation,
+    Tensor,
+    dtype_from_numpy,
+    dtype_to_numpy,
+)
+
+
+class Connection:
+    def __init__(
+        self,
+        server: str = "localhost:5555",
+        client_id: str | int = "",
+        client_type: str = "",
+        reconnect_ivl_ms: int = 100,
+        reconnect_max_ms: int = 5000,
+    ) -> None:
+        self._server = server
+        self._client_id = str(client_id)
+        self._client_type = client_type
+        self._ctx = zmq.Context()
+        self._socket = self._ctx.socket(zmq.DEALER)
+        self._socket.identity = str(client_id).encode()
+        self._socket.setsockopt(zmq.RECONNECT_IVL, reconnect_ivl_ms)
+        self._socket.setsockopt(zmq.RECONNECT_IVL_MAX, reconnect_max_ms)
+        self._socket.setsockopt(zmq.RCVHWM, 100)
+        self._socket.setsockopt(zmq.SNDHWM, 100)
+        self._socket.setsockopt(zmq.LINGER, 0)
+
+        if not server.startswith("tcp://"):
+            server = f"tcp://{server}"
+        self._socket.connect(server)
+
+    def model(
+        self,
+        model_id: str,
+        latency_budget_ms: float = 50.0,
+        priority: int = 1,
+    ) -> Model:
+        return Model(
+            connection=self,
+            model_id=model_id,
+            latency_budget_ms=latency_budget_ms,
+            priority=priority,
+        )
+
+    def _send(self, envelope: bytes, payload: bytes) -> None:
+        self._socket.send_multipart([b"", envelope, payload])
+
+    def _recv(self, timeout_ms: int = 0) -> tuple[bytes, bytes] | None:
+        if timeout_ms > 0:
+            if self._socket.poll(timeout_ms, zmq.POLLIN) == 0:
+                return None
+        elif not self._socket.poll(0, zmq.POLLIN):
+            return None
+
+        frames = self._socket.recv_multipart(zmq.NOBLOCK)
+        # frames = [b'', envelope, payload]
+        envelope = frames[1] if len(frames) > 1 else b""
+        payload = frames[2] if len(frames) > 2 else b""
+        return envelope, payload
+
+    def close(self) -> None:
+        self._socket.close()
+        self._ctx.term()
+
+
+class Model:
+    def __init__(
+        self,
+        connection: Connection,
+        model_id: str,
+        latency_budget_ms: float = 50.0,
+        priority: int = 1,
+    ) -> None:
+        self._conn = connection
+        self._model_id = model_id
+        self._latency_budget_ms = latency_budget_ms
+        self._priority = priority
+
+    def observe(
+        self,
+        urgency: float = 0.0,
+        steps_remaining: int | None = None,
+        **kwargs: Any,
+    ) -> None:
+        obs = Observation()
+        obs.client.CopyFrom(Client(id=self._conn._client_id, type=self._conn._client_type))
+        obs.model_id = self._model_id
+        obs.timestamp_ns = int(time.time() * 1_000_000_000)
+        obs.urgency = urgency
+        if steps_remaining is not None:
+            obs.steps_remaining = steps_remaining
+
+        payload_parts: list[bytes] = []
+        offset = 0
+
+        for key, value in kwargs.items():
+            if isinstance(value, np.ndarray):
+                data = value.tobytes()
+                tensor = Tensor()
+                tensor.key = key
+                tensor.dtype = dtype_from_numpy(value.dtype)
+                tensor.shape.extend(value.shape)
+                tensor.byte_offset = offset
+                tensor.byte_length = len(data)
+                tensor.timestamp_ns = obs.timestamp_ns
+                tensor.encoding = RAW
+                obs.tensors.append(tensor)
+                payload_parts.append(data)
+                offset += len(data)
+            elif isinstance(value, str):
+                obs.metadata[key] = value
+
+        payload = b"".join(payload_parts)
+        envelope = obs.SerializeToString()
+        self._conn._send(envelope, payload)
+
+    def get_result(self, timeout_ms: int = 100) -> dict | None:
+        result = self._conn._recv(timeout_ms)
+        if result is None:
+            return None
+
+        envelope, payload = result
+        resp = ModelResponse()
+        resp.ParseFromString(envelope)
+
+        output: dict[str, Any] = {
+            "response_id": resp.response_id,
+            "model_id": resp.model_id,
+            "inference_latency_ms": resp.inference_latency_ms,
+        }
+
+        for tensor in resp.tensors:
+            if tensor.byte_length > 0 and tensor.dtype:
+                np_dtype = dtype_to_numpy(tensor.dtype)
+                shape = tuple(tensor.shape) if tensor.shape else ()
+                data = payload[tensor.byte_offset : tensor.byte_offset + tensor.byte_length]
+                output[tensor.key] = np.frombuffer(data, dtype=np_dtype).reshape(shape)
+
+        for k, v in resp.metadata.items():
+            output[k] = v
+
+        return output

inferential/config/__init__.py

@@ -0,0 +1,31 @@
+from inferential.config.schema import (
+    BatchConfig,
+    ClientDefaults,
+    ClientEntry,
+    ClientsConfig,
+    DeadlineAwareConfig,
+    InferentialConfig,
+    MetricsConfig,
+    ObservationConfig,
+    ResponseTrackingConfig,
+    SchedulingConfig,
+    TieredConfig,
+    TransportConfig,
+)
+from inferential.config.watcher import ConfigWatcher
+
+__all__ = [
+    "BatchConfig",
+    "ClientDefaults",
+    "ClientEntry",
+    "ClientsConfig",
+    "ConfigWatcher",
+    "DeadlineAwareConfig",
+    "InferentialConfig",
+    "MetricsConfig",
+    "ObservationConfig",
+    "ResponseTrackingConfig",
+    "SchedulingConfig",
+    "TieredConfig",
+    "TransportConfig",
+]

inferential/config/schema.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+from typing import Literal
+
+from pydantic import BaseModel, Field
+
+
+class TransportConfig(BaseModel):
+    type: Literal["zmq"] = "zmq"
+    bind: str = "tcp://*:5555"
+    recv_hwm: int = Field(default=1000, ge=1)
+    send_hwm: int = Field(default=1000, ge=1)
+
+
+class DeadlineAwareConfig(BaseModel):
+    cadence_weight: float = 45.0
+    urgency_weight: float = 25.0
+    steps_weight: float = 15.0
+    priority_weight: float = 10.0
+    age_weight: float = 15.0
+
+
+class BatchConfig(BaseModel):
+    max_batch_size: int = Field(default=8, ge=1)
+    max_wait_ms: float = Field(default=10.0, gt=0)
+
+
+class TieredConfig(BaseModel):
+    num_tiers: int = Field(default=3, ge=1)
+
+
+class SchedulingConfig(BaseModel):
+    strategy: str = "deadline_aware"
+    max_queue_size: int = Field(default=1000, ge=1)
+    request_ttl_ms: float = Field(default=5000.0, gt=0)
+    overflow_policy: Literal["drop_oldest", "reject_newest"] = "drop_oldest"
+    max_retries: int = Field(default=0, ge=0)
+    deadline_aware: DeadlineAwareConfig = DeadlineAwareConfig()
+    batch_optimized: BatchConfig = BatchConfig()
+    priority_tiered: TieredConfig = TieredConfig()
+
+
+class ClientDefaults(BaseModel):
+    latency_budget_ms: float = Field(default=50.0, gt=0)
+    priority: int = Field(default=1, ge=0)
+
+
+class ClientEntry(BaseModel):
+    id: str
+    model: str
+    latency_budget_ms: float | None = None
+    priority: int | None = None
+
+
+class ClientsConfig(BaseModel):
+    defaults: ClientDefaults = ClientDefaults()
+    known: list[ClientEntry] = []
+    accept_unknown: bool = True
+
+
+class ResponseTrackingConfig(BaseModel):
+    overdue_multiplier: float = Field(default=1.5, gt=1.0)
+    disconnect_timeout_s: float = Field(default=10.0, gt=0)
+    cadence_alpha: float = Field(default=0.3, gt=0, lt=1.0)
+
+
+class ObservationConfig(BaseModel):
+    max_staleness_ms: float = Field(default=1000.0, gt=0)
+    temporal_alignment_threshold_ms: float = Field(default=20.0, gt=0)
+
+
+class MetricsConfig(BaseModel):
+    ring_buffer_size: int = Field(default=10000, ge=100)
+    enabled: bool = True
+
+
+class InferentialConfig(BaseModel):
+    ray: dict = Field(default_factory=dict)
+    transport: TransportConfig = TransportConfig()
+    scheduling: SchedulingConfig = SchedulingConfig()
+    clients: ClientsConfig = ClientsConfig()
+    response_tracking: ResponseTrackingConfig = ResponseTrackingConfig()
+    observations: ObservationConfig = ObservationConfig()
+    metrics: MetricsConfig = MetricsConfig()
+
+    def get_client_budget(self, client_id: str) -> float:
+        for entry in self.clients.known:
+            if entry.id == client_id and entry.latency_budget_ms is not None:
+                return entry.latency_budget_ms
+        return self.clients.defaults.latency_budget_ms
+
+    def get_client_priority(self, client_id: str) -> int:
+        for entry in self.clients.known:
+            if entry.id == client_id and entry.priority is not None:
+                return entry.priority
+        return self.clients.defaults.priority

inferential/config/watcher.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import os
+import signal
+from pathlib import Path
+from typing import Callable
+
+from inferential.config.schema import InferentialConfig
+
+logger = logging.getLogger("inferential.config")
+
+ReloadCallback = Callable[[InferentialConfig], None]
+
+
+class ConfigWatcher:
+    def __init__(
+        self,
+        config_path: str | Path | None = None,
+        poll_interval_s: float = 5.0,
+    ) -> None:
+        self._config_path = Path(config_path) if config_path else None
+        self._poll_interval_s = poll_interval_s
+        self._callbacks: list[ReloadCallback] = []
+        self._last_mtime: float = 0.0
+        self._sighup_received = False
+
+    def on_reload(self, callback: ReloadCallback) -> ReloadCallback:
+        self._callbacks.append(callback)
+        return callback
+
+    def _install_sighup_handler(self) -> None:
+        try:
+            loop = asyncio.get_running_loop()
+            loop.add_signal_handler(signal.SIGHUP, self._handle_sighup)
+        except (NotImplementedError, OSError):
+            pass  # Windows or restricted environment
+
+    def _handle_sighup(self) -> None:
+        self._sighup_received = True
+        logger.info("SIGHUP received, will reload config")
+
+    def _load_config(self) -> InferentialConfig | None:
+        if self._config_path is None or not self._config_path.exists():
+            return None
+        try:
+            data = json.loads(self._config_path.read_text())
+            return InferentialConfig(**data)
+        except Exception as e:
+            logger.error("Failed to reload config from %s: %s", self._config_path, e)
+            return None
+
+    def _fire(self, config: InferentialConfig) -> None:
+        for cb in self._callbacks:
+            try:
+                cb(config)
+            except Exception:
+                logger.exception("Error in config reload callback")
+
+    async def watch(self) -> None:
+        self._install_sighup_handler()
+
+        while True:
+            await asyncio.sleep(self._poll_interval_s)
+
+            should_reload = self._sighup_received
+
+            if self._config_path and self._config_path.exists():
+                try:
+                    mtime = os.path.getmtime(self._config_path)
+                    if mtime > self._last_mtime:
+                        should_reload = True
+                        self._last_mtime = mtime
+                except OSError:
+                    pass
+
+            if should_reload:
+                self._sighup_received = False
+                config = self._load_config()
+                if config is not None:
+                    logger.info("Config reloaded from %s", self._config_path)
+                    self._fire(config)

inferential/dispatch/__init__.py

@@ -0,0 +1,8 @@
+from inferential.dispatch.dispatcher import DispatchResult, RayDispatcher
+from inferential.dispatch.health import EndpointHealth
+
+__all__ = [
+    "DispatchResult",
+    "EndpointHealth",
+    "RayDispatcher",
+]

inferential/dispatch/dispatcher.py

@@ -0,0 +1,166 @@
+from __future__ import annotations
+
+import logging
+import time
+import uuid
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+
+from inferential.dispatch.health import EndpointHealth
+from inferential.proto import (
+    RAW,
+    Client,
+    ModelResponse,
+    Observation,
+    Tensor,
+    dtype_from_numpy,
+    dtype_to_numpy,
+)
+from inferential.scheduler.request import InferenceRequest
+
+logger = logging.getLogger("inferential.dispatch")
+
+
+@dataclass
+class DispatchResult:
+    client_id: str
+    response_id: str
+    identity: bytes
+    envelope: bytes
+    payload: bytes
+    latency_ms: float
+    success: bool
+    error: str | None = None
+
+
+class RayDispatcher:
+    def __init__(self, ray_config: dict | None = None) -> None:
+        self._ray_config = ray_config or {}
+        self._handles: dict[str, Any] = {}
+        self._health: dict[str, EndpointHealth] = {}
+
+    def _get_handle(self, model_id: str) -> Any:
+        if model_id not in self._handles:
+            from ray import serve
+
+            self._handles[model_id] = serve.get_deployment_handle(model_id)
+        return self._handles[model_id]
+
+    def _reconstruct_numpy(self, req: InferenceRequest) -> dict:
+        obs = Observation()
+        obs.ParseFromString(req.envelope)
+        result: dict[str, Any] = {}
+        for tensor in obs.tensors:
+            np_dtype = dtype_to_numpy(tensor.dtype)
+            shape = tuple(tensor.shape) if tensor.shape else ()
+            data = req.payload[tensor.byte_offset : tensor.byte_offset + tensor.byte_length]
+            array = np.frombuffer(data, dtype=np_dtype).reshape(shape)
+            result[tensor.key] = array
+        for k, v in obs.metadata.items():
+            result[k] = v
+        return result
+
+    def _ndarray_to_tensor(
+        self, key: str, arr: np.ndarray, offset: int
+    ) -> tuple[Tensor, bytes]:
+        data = arr.tobytes()
+        tensor = Tensor()
+        tensor.key = key
+        tensor.dtype = dtype_from_numpy(arr.dtype)
+        tensor.shape.extend(arr.shape)
+        tensor.byte_offset = offset
+        tensor.byte_length = len(data)
+        tensor.encoding = RAW
+        return tensor, data
+
+    def _build_result(
+        self,
+        req: InferenceRequest,
+        ray_result: Any,
+        latency_ms: float,
+    ) -> DispatchResult:
+        resp = ModelResponse()
+        resp.client.CopyFrom(Client(id=req.client_id))
+        resp.response_id = uuid.uuid4().hex
+        resp.timestamp_ns = int(time.time() * 1_000_000_000)
+        resp.inference_latency_ms = latency_ms
+        resp.model_id = req.model_id
+
+        payload_parts: list[bytes] = []
+        offset = 0
+
+        if isinstance(ray_result, np.ndarray):
+            tensor, data = self._ndarray_to_tensor("actions", ray_result, offset)
+            resp.tensors.append(tensor)
+            payload_parts.append(data)
+        elif isinstance(ray_result, dict):
+            for key, value in ray_result.items():
+                if isinstance(value, np.ndarray):
+                    tensor, data = self._ndarray_to_tensor(key, value, offset)
+                    resp.tensors.append(tensor)
+                    payload_parts.append(data)
+                    offset += len(data)
+                elif isinstance(value, str):
+                    resp.metadata[key] = value
+
+        payload = b"".join(payload_parts)
+
+        return DispatchResult(
+            client_id=req.client_id,
+            response_id=resp.response_id,
+            identity=req.identity,
+            envelope=resp.SerializeToString(),
+            payload=payload,
+            latency_ms=latency_ms,
+            success=True,
+        )
+
+    async def dispatch(
+        self, requests: list[InferenceRequest]
+    ) -> list[DispatchResult]:
+        results: list[DispatchResult] = []
+        by_model: dict[str, list[InferenceRequest]] = defaultdict(list)
+        for req in requests:
+            by_model[req.model_id].append(req)
+
+        for model_id, model_requests in by_model.items():
+            handle = self._get_handle(model_id)
+            for req in model_requests:
+                start = time.monotonic()
+                try:
+                    obs_dict = self._reconstruct_numpy(req)
+                    result = await handle.infer.remote(obs_dict)
+                    latency = (time.monotonic() - start) * 1000
+                    self._update_health(model_id, latency, True)
+                    results.append(self._build_result(req, result, latency))
+                except Exception as e:
+                    latency = (time.monotonic() - start) * 1000
+                    self._update_health(model_id, latency, False)
+                    logger.error("Dispatch failed for client %s: %s", req.client_id, e)
+                    results.append(
+                        DispatchResult(
+                            client_id=req.client_id,
+                            response_id=uuid.uuid4().hex,
+                            identity=req.identity,
+                            envelope=b"",
+                            payload=b"",
+                            latency_ms=latency,
+                            success=False,
+                            error=str(e),
+                        )
+                    )
+        return results
+
+    def _update_health(self, model_id: str, latency_ms: float, success: bool) -> None:
+        if model_id not in self._health:
+            self._health[model_id] = EndpointHealth(model_id=model_id)
+        self._health[model_id].update(latency_ms, success)
+
+    def endpoint_health(self, model_id: str) -> EndpointHealth | None:
+        return self._health.get(model_id)
+
+    def get_handle(self, model_id: str) -> Any | None:
+        return self._handles.get(model_id)

inferential/dispatch/health.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+
+@dataclass
+class EndpointHealth:
+    model_id: str
+    avg_latency_ms: float = 0.0
+    error_rate: float = 0.0
+    total_requests: int = 0
+    last_error: str | None = None
+    _alpha: float = field(default=0.2, repr=False)
+
+    def update(self, latency_ms: float, success: bool) -> None:
+        self.total_requests += 1
+        self.avg_latency_ms = (
+            self._alpha * latency_ms + (1 - self._alpha) * self.avg_latency_ms
+        )
+        error_val = 0.0 if success else 1.0
+        self.error_rate = self._alpha * error_val + (1 - self._alpha) * self.error_rate
+        if not success:
+            self.last_error = f"request #{self.total_requests}"

inferential/metrics/__init__.py

@@ -0,0 +1,13 @@
+from inferential.metrics.callbacks import CallbackRegistry, MetricCallback
+from inferential.metrics.collector import MetricsCollector, MetricSnapshot
+from inferential.metrics.store import MetricPoint, MetricSeries, MetricStats
+
+__all__ = [
+    "CallbackRegistry",
+    "MetricCallback",
+    "MetricPoint",
+    "MetricSeries",
+    "MetricSnapshot",
+    "MetricStats",
+    "MetricsCollector",
+]

inferential/metrics/callbacks.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+from typing import Callable
+
+MetricCallback = Callable[[str, float, dict[str, str]], None]
+
+
+class CallbackRegistry:
+    def __init__(self) -> None:
+        self._callbacks: list[MetricCallback] = []
+
+    def register(self, callback: MetricCallback) -> None:
+        self._callbacks.append(callback)
+
+    def fire(self, name: str, value: float, labels: dict[str, str]) -> None:
+        for cb in self._callbacks:
+            try:
+                cb(name, value, labels)
+            except Exception:
+                pass  # Callbacks must not crash the server

inferential/metrics/collector.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from inferential.metrics.callbacks import CallbackRegistry, MetricCallback
+from inferential.metrics.store import MetricSeries, MetricStats
+
+if TYPE_CHECKING:
+    from inferential.config.schema import MetricsConfig
+
+
+@dataclass
+class MetricSnapshot:
+    metrics: dict[str, list[dict]]
+
+
+class MetricsCollector:
+    def __init__(self, config: MetricsConfig | None = None) -> None:
+        self._ring_buffer_size = config.ring_buffer_size if config else 10000
+        self._enabled = config.enabled if config else True
+        self._series: dict[str, MetricSeries] = {}
+        self._callbacks = CallbackRegistry()
+
+    def record(
+        self,
+        name: str,
+        value: float,
+        labels: dict[str, str] | None = None,
+    ) -> None:
+        if not self._enabled:
+            return
+        labels = labels or {}
+        if name not in self._series:
+            self._series[name] = MetricSeries(max_size=self._ring_buffer_size)
+        self._series[name].record(value, labels)
+        self._callbacks.fire(name, value, labels)
+
+    def get_latest(
+        self,
+        name: str,
+        labels: dict[str, str] | None = None,
+    ) -> float | None:
+        series = self._series.get(name)
+        if series is None:
+            return None
+        return series.get_latest(labels)
+
+    def get_stats(
+        self,
+        name: str,
+        labels: dict[str, str] | None = None,
+        window_seconds: int = 300,
+    ) -> MetricStats | None:
+        series = self._series.get(name)
+        if series is None:
+            return None
+        return series.get_stats(labels, window_seconds)
+
+    def snapshot(self) -> MetricSnapshot:
+        result: dict[str, list[dict]] = {}
+        for name, series in self._series.items():
+            result[name] = [
+                {"timestamp": p.timestamp, "value": p.value, "labels": p.labels}
+                for p in series._buffer
+            ]
+        return MetricSnapshot(metrics=result)
+
+    def on_metric(self, callback: MetricCallback) -> MetricCallback:
+        self._callbacks.register(callback)
+        return callback