positronic-model-server 0.2.2__py3-none-any.whl

positronic_model_server/grpc_wire.py

@@ -0,0 +1,166 @@
+"""The server side of the gRPC wire: one bidirectional stream per session, which carries the ``protocol`` frames,
+and the unary keepalive call beside it."""
+
+import json
+import logging
+from collections.abc import AsyncIterator, Mapping
+from urllib.parse import parse_qsl
+
+import grpc
+import grpc.aio
+from positronic_wire import wire
+from positronic_wire.grpc import (
+    KEEPALIVE_METHOD,
+    MESSAGE_SIZE_OPTIONS,
+    METHOD,
+    PING_EVERY_MS,
+    SERVICE,
+    SESSION_PATH_HEADER,
+    SESSION_QUERY_HEADER,
+    target,
+)
+
+from . import server_wire
+
+logger = logging.getLogger(__name__)
+
+
+class GrpcServerConnection(server_wire.ServerConnection):
+    """A server's end of one gRPC session."""
+
+    def __init__(
+        self,
+        requests: AsyncIterator[bytes],
+        context: grpc.aio.ServicerContext,
+        headers: Mapping[str, str],
+        served_address: server_wire.ServedHostPort,
+    ):
+        self._requests = requests
+        self._context = context
+        self._headers = headers
+        self._served_address = served_address
+
+    @property
+    def peer(self) -> str:
+        return self._context.peer()
+
+    @property
+    def served_address(self) -> server_wire.ServedHostPort:
+        return self._served_address
+
+    @property
+    def session_path(self) -> str:
+        return self._headers.get(SESSION_PATH_HEADER, wire.SESSION_PATH)
+
+    @property
+    def query_params(self) -> list[tuple[str, str]]:
+        return parse_qsl(self._headers.get(SESSION_QUERY_HEADER, ''), keep_blank_values=True)
+
+    async def send(self, message: bytes) -> None:
+        try:
+            await self._context.write(message)
+        except grpc.RpcError as e:
+            raise wire.PeerDisconnected(f'{self.peer} ended the session: {e}') from e
+
+    async def receive(self) -> bytes:
+        try:
+            return await anext(self._requests)
+        except StopAsyncIteration:
+            raise wire.PeerDisconnected(f'{self.peer} ended the session') from None
+
+    async def refuse(self, reason: str) -> None:
+        self._context.set_code(grpc.StatusCode.ABORTED)
+        self._context.set_details(reason)
+
+
+def _headers(context: grpc.aio.ServicerContext) -> dict[str, str]:
+    """The session metadata, as the header names both wires share. A ``-bin`` key carries no header."""
+    return {key: value for key, value in (context.invocation_metadata() or ()) if isinstance(value, str)}
+
+
+# The shortest ping interval the server answers without a strike: half of `PING_EVERY_MS`, so a client
+# ping that arrives early is never one.
+_PING_TOLERATED_EVERY_MS = PING_EVERY_MS // 2
+
+
+def _server_options() -> list[tuple[str, int]]:
+    return [
+        *MESSAGE_SIZE_OPTIONS,
+        # gRPC's own defaults, a five-minute floor and two strikes, answer a 20s ping with GOAWAY.
+        ('grpc.http2.min_ping_interval_without_data_ms', _PING_TOLERATED_EVERY_MS),
+        ('grpc.http2.max_ping_strikes', 0),
+    ]
+
+
+class GrpcWire(server_wire.Wire):
+    """The gRPC wire: sessions on a port of their own, one bidirectional stream each, and the keepalive call.
+
+    ``served_address`` is the host and the port it binds; a port of 0 binds any free one. The port is
+    plaintext, and a TLS edge in front of it serves an authenticated endpoint.
+    """
+
+    def __init__(self, served_address: server_wire.ServedHostPort):
+        self._binds = served_address
+        self._server: grpc.aio.Server | None = None
+        self._served_address: server_wire.ServedHostPort | None = None
+
+    @property
+    def served_address(self) -> server_wire.ServedHostPort:
+        assert self._served_address is not None, 'The gRPC wire has not started'
+        return self._served_address
+
+    async def start(
+        self,
+        session: server_wire.SessionHandler,
+        keepalive: server_wire.KeepaliveHandler,
+        authorized: server_wire.Authorized,
+    ) -> None:
+        """Bind the gRPC port, which carries sessions and the keepalive call."""
+
+        async def serve_one(requests: AsyncIterator[bytes], context: grpc.aio.ServicerContext) -> None:
+            headers = _headers(context)
+            if not authorized(headers):
+                await context.abort(grpc.StatusCode.PERMISSION_DENIED, 'Invalid or missing bearer token')
+            conn = GrpcServerConnection(requests, context, headers, self.served_address)
+            if conn.session_path != wire.SESSION_PATH:
+                refusal = f'No session on {conn.session_path!r}; sessions open on {wire.SESSION_PATH}'
+                await context.abort(grpc.StatusCode.NOT_FOUND, refusal)
+            try:
+                await session(conn)
+            except Exception as e:
+                # The session reports its own errors over the stream. One that reaches here reaches the client
+                # as the status alone.
+                logger.error(f'Failed gRPC session: {e}', exc_info=True)
+                await context.abort(grpc.StatusCode.INTERNAL, str(e))
+
+        async def answer_keepalive(_request: bytes, context: grpc.aio.ServicerContext) -> bytes:
+            if not authorized(_headers(context)):
+                await context.abort(grpc.StatusCode.PERMISSION_DENIED, 'Invalid or missing bearer token')
+            return json.dumps({wire.ALIVE_SECONDS: keepalive()}).encode()
+
+        methods = {
+            METHOD: grpc.stream_stream_rpc_method_handler(
+                serve_one, request_deserializer=None, response_serializer=None
+            ),
+            KEEPALIVE_METHOD: grpc.unary_unary_rpc_method_handler(
+                answer_keepalive, request_deserializer=None, response_serializer=None
+            ),
+        }
+        server = grpc.aio.server(options=_server_options())
+        server.add_generic_rpc_handlers((grpc.method_handlers_generic_handler(SERVICE, methods),))
+        bound = server.add_insecure_port(target(self._binds.host, self._binds.port))
+        if bound == 0:
+            # gRPC reports a refused bind as port 0, and a server started on it accepts nothing and says nothing.
+            raise OSError(f'gRPC could not bind {target(self._binds.host, self._binds.port)}')
+        self._server = server
+        self._served_address = server_wire.ServedHostPort(self._binds.host, bound)
+        await server.start()
+        logger.info(f'gRPC sessions on {self._binds.host}:{bound}')
+
+    async def serve(self) -> None:
+        assert self._server is not None, 'The gRPC wire has not started'
+        await self._server.wait_for_termination()
+
+    async def stop(self) -> None:
+        if self._server is not None:
+            await self._server.stop(grace=None)

positronic_model_server/keys.py

@@ -0,0 +1,10 @@
+"""Metadata fields written by the model-server wrapper."""
+
+HOST = 'host'
+PORT = 'port'
+# The socket path a server bound instead of a host and a port. One of the two pairs is present, never both.
+UDS = 'uds'
+LOCAL_STACK = 'local_stack'
+MODEL_SERVER_VERSION = 'model_server_version'
+SESSION_PARAMS = 'session_params'
+EFFECTIVE_PARAMS = 'effective_params'

positronic_model_server/protocol.py

@@ -0,0 +1,60 @@
+"""Shared inference message fields, versions, authentication and timing."""
+
+from enum import IntEnum, StrEnum
+
+AUTH_HEADER = 'Authorization'
+
+
+def bearer(token: str) -> str:
+    """The authorization header value for a bearer token."""
+    return f'Bearer {token}'
+
+
+# The top-level keys of every server-to-client message: ``STATUS`` until the server reports itself ready
+# and hands over its ``META``, then one ``RESULT`` or ``ERROR`` per inference.
+STATUS = 'status'
+META = 'meta'
+RESULT = 'result'
+ERROR = 'error'
+# The ready handshake issues an ID. Requests carry it beside the observation, and an end request
+# is acknowledged with the same ID after the model releases the session's state.
+SESSION_ID = 'session_id'
+OBSERVATION = 'observation'
+END_SESSION = 'end_session'
+PROTOCOL_VERSION = 'protocol_version'
+
+
+class ProtocolVersion(IntEnum):
+    V1 = 1
+    V2 = 2
+    V3 = 3
+
+
+# What the server spent on one inference, beside the ``RESULT`` it answers with: durations in
+# milliseconds on the server's own clock. A server that sends none leaves the round trip undivided.
+TIMING = 'timing'
+
+# The phases ``TIMING`` reports. `SERVED` brackets the others: it opens on the observation
+# arriving and closes before the answer is encoded.
+TIMING_SERVED = 'served_ms'
+TIMING_DECODE = 'decode_ms'
+TIMING_INFER = 'infer_ms'
+# Time the observation waited for the inference slot, inside `SERVED`.
+TIMING_QUEUED = 'queued_ms'
+
+
+def timing_key(name: str) -> str:
+    return f'{name}_ms'
+
+
+# The loaded model's call, excluding codec conversions.
+MODEL_CALL = 'model'
+# Time the model's own call took, inside `INFER`.
+TIMING_MODEL = timing_key(MODEL_CALL)
+
+
+class ServerStatus(StrEnum):
+    READY = 'ready'
+    WAITING = 'waiting'
+    LOADING = 'loading'
+    ERROR = 'error'

positronic_model_server/serialization.py

@@ -0,0 +1,144 @@
+"""Wire serialization helpers for numpy arrays, images and standard Python types.
+
+Supports:
+- built-in scalars: `str`, `int`, `float`, `bool`, `None`
+- containers: `dict` / `list` / `tuple` recursively composed of supported values
+- numeric numpy values: `numpy.ndarray` and `numpy` scalar types
+- images JPEG-encoded with ``encode_jpeg``
+
+Nothing here knows a domain type. A boundary that carries one writes its own msgpack hooks and
+delegates to ``pack`` / ``unpack`` for the rest.
+"""
+
+import collections.abc as cabc
+import functools
+import io
+from dataclasses import dataclass
+from typing import Any
+
+import msgpack
+import numpy as np
+from PIL import Image as PilImage
+
+# The envelope each encoded value travels in: a marker naming the type, and the fields carrying it.
+_NDARRAY = b'__ndarray__'
+_NPGENERIC = b'__npgeneric__'
+_JPEG = b'__jpeg__'
+_DATA = b'data'
+_DTYPE = b'dtype'
+_SHAPE = b'shape'
+_FRAMES = b'frames'
+_NDIM = b'ndim'
+
+DEFAULT_JPEG_QUALITY = 90
+
+
+@dataclass(frozen=True)
+class JpegEncoding:
+    """An explicit mapping/list path to an RGB image; an empty path selects the whole value."""
+
+    path: tuple[str | int, ...]
+    quality: int = DEFAULT_JPEG_QUALITY
+
+    def __post_init__(self) -> None:
+        if any(type(part) not in (str, int) for part in self.path):
+            raise ValueError('Image paths contain only mapping keys and list indices')
+        if type(self.quality) is not int or not 0 <= self.quality <= 100:
+            raise ValueError('JPEG quality must be an integer between 0 and 100')
+
+    def encode(self, value: Any) -> Any:
+        """Copy the selected container path and encode its leaf, preserving unrelated values."""
+
+        def at(node: Any, path: tuple[str | int, ...]) -> Any:
+            if not path:
+                if not isinstance(node, np.ndarray):
+                    raise ValueError(f'JPEG path {self.path!r} does not select an array')
+                return encode_jpeg(node, self.quality)
+            key, *rest = path
+            if isinstance(node, cabc.Mapping):
+                return {**node, key: at(node[key], tuple(rest))}
+            if isinstance(node, list | tuple) and type(key) is int:
+                copy = list(node)
+                copy[key] = at(node[key], tuple(rest))
+                return tuple(copy) if isinstance(node, tuple) else copy
+            raise ValueError(f'JPEG path {self.path!r} does not match the value')
+
+        return at(value, self.path)
+
+
+def encode_images(value: Any, encodings: cabc.Sequence[JpegEncoding]) -> Any:
+    """Encode only the explicitly selected images, in observations or model results."""
+    for encoding in encodings:
+        value = encoding.encode(value)
+    return value
+
+
+def encode_jpeg(image: np.ndarray, quality: int = DEFAULT_JPEG_QUALITY) -> dict[bytes, Any]:
+    """JPEG-encode RGB images shaped ``(..., H, W, 3)``, casting pixels to uint8.
+
+    Nonempty 3D/4D inputs use the v1/v2 marker. Extra leading dimensions and empty batches carry
+    their full shape. JPEG changes pixels; ordinary arrays passed to ``serialise`` stay lossless.
+    """
+    if image.ndim < 3 or image.shape[-1] != 3 or min(image.shape[-3:-1]) < 1:
+        raise ValueError('JPEG images must have shape (..., H, W, 3) with positive height and width')
+    if type(quality) is not int or not 0 <= quality <= 100:
+        raise ValueError('JPEG quality must be an integer between 0 and 100')
+    frames = image.reshape((-1, *image.shape[-3:]))
+    bufs = []
+    for frame in frames:
+        buf = io.BytesIO()
+        PilImage.fromarray(np.ascontiguousarray(frame, dtype=np.uint8)).save(buf, format='JPEG', quality=quality)
+        bufs.append(buf.getvalue())
+    if image.ndim > 4 or not bufs:
+        return {_JPEG: True, _FRAMES: bufs, _SHAPE: image.shape}
+    return {_JPEG: True, _FRAMES: bufs, _NDIM: int(image.ndim)}
+
+
+def _decode_jpeg(marker: dict) -> np.ndarray:
+    """Inverse of ``encode_jpeg``: decode per-frame JPEGs and restore the original shape."""
+    if _SHAPE in marker and not marker[_FRAMES]:
+        shape = tuple(marker[_SHAPE])
+        if len(shape) < 4 or shape[-1] != 3 or min(shape[-3:-1]) < 1 or 0 not in shape[:-3]:
+            raise ValueError('An empty JPEG batch must have an empty leading dimension')
+        return np.empty(shape, dtype=np.uint8)
+    frames = np.stack([np.asarray(PilImage.open(io.BytesIO(buf))) for buf in marker[_FRAMES]])
+    if _SHAPE in marker:
+        return frames.reshape(marker[_SHAPE])
+    return frames if marker[_NDIM] == 4 else frames[0]
+
+
+def pack(obj):
+    """msgpack's ``default`` hook: one value in its wire form, or unchanged when msgpack handles it."""
+    if isinstance(obj, cabc.Mapping):
+        return dict(obj)
+    if isinstance(obj, np.ndarray | np.generic) and obj.dtype.kind in ('V', 'O', 'c'):
+        raise ValueError(f'Unsupported dtype: {obj.dtype}')
+    if isinstance(obj, np.ndarray):
+        return {_NDARRAY: True, _DATA: obj.tobytes(), _DTYPE: obj.dtype.str, _SHAPE: obj.shape}
+    if isinstance(obj, np.generic):
+        return {_NPGENERIC: True, _DATA: obj.item(), _DTYPE: obj.dtype.str}
+    return obj
+
+
+def unpack(obj):
+    """msgpack's ``object_hook``: one decoded mapping restored to the value it encodes."""
+    if _NDARRAY in obj:
+        return np.ndarray(buffer=obj[_DATA], dtype=np.dtype(obj[_DTYPE]), shape=obj[_SHAPE])
+    if _NPGENERIC in obj:
+        return np.dtype(obj[_DTYPE]).type(obj[_DATA])
+    if _JPEG in obj:
+        return _decode_jpeg(obj)
+    return obj
+
+
+def serialise(obj: Any) -> bytes:
+    packed = msgpack.packb(obj, default=pack)
+    assert packed is not None
+    return packed
+
+
+deserialise = functools.partial(msgpack.unpackb, object_hook=unpack)
+
+# Aliases for consistency
+serialize = serialise
+deserialize = deserialise

positronic_model_server/server.py

@@ -0,0 +1,309 @@
+"""Serve a warmed native model and its sessions in one process."""
+
+import asyncio
+import hmac
+import json
+import logging
+import math
+import time
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Mapping, Sequence
+from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import asynccontextmanager
+from importlib.metadata import version
+from typing import Any, TypeVar
+from uuid import uuid4
+
+from positronic_wire import wire
+
+from . import keys, protocol, serialization, server_wire, spec
+
+logger = logging.getLogger(__name__)
+T = TypeVar('T')
+
+
+class Session(ABC):
+    """A prepared, warm session that owns its state while sharing loaded weights.
+
+    ``close`` releases that state. ``output_images`` selects result values for JPEG serialization.
+    """
+
+    def __init__(
+        self,
+        client_stack: dict[str, Any],
+        *,
+        metadata: dict[str, Any] | None = None,
+        output_images: Sequence[serialization.JpegEncoding] = (),
+    ):
+        self.client_stack = client_stack
+        self.metadata = {} if metadata is None else metadata
+        self.output_images = output_images
+
+    @abstractmethod
+    def infer(self, observation: Any) -> Any: ...
+
+    def close(self) -> None:
+        """Release this session's resources after its operations finish."""
+        return None
+
+
+class Model(ABC):
+    """Loaded, warm model resources and the operation that prepares each session.
+
+    ``parameters`` declares accepted session names and their defaults. ``prepare_session`` receives
+    resolved values, validates their model-specific meaning and returns only once that session is warm.
+    A preparation that raises must release any resources it acquired before returning a session.
+    """
+
+    def __init__(self, *, parameters: dict[str, Any] | None = None, metadata: dict[str, Any] | None = None):
+        self.parameters = {} if parameters is None else parameters
+        self.metadata = {} if metadata is None else metadata
+
+    @abstractmethod
+    def prepare_session(self, params: dict[str, Any]) -> Session: ...
+
+    def close(self) -> None:
+        """Release the model's resources after its sessions finish."""
+        return None
+
+
+class ModelServer:
+    """One model, serialized operations, and independently configured transport listeners.
+
+    Loading, session preparation, inference and cleanup all run on one worker thread. The network
+    loop remains available during these operations. The factory must finish shared warm-up before
+    returning its model; listeners bind afterwards. Each session must retain its own parameters and
+    history, and its warm-up must leave the session in its intended initial state.
+    """
+
+    WAITING_INTERVAL_SEC = 5.0
+
+    def __init__(
+        self, load_model: Callable[[], Model], *, idle_timeout_min: float | None = None, auth_token: str | None = None
+    ):
+        if auth_token is not None and not (auth_token and all('!' <= c <= '~' for c in auth_token)):
+            raise ValueError('auth_token must be non-empty printable ASCII without spaces; pass None to serve open')
+        if idle_timeout_min is not None and not 0 < idle_timeout_min < math.inf:
+            raise ValueError('idle_timeout_min must be positive and finite, or None')
+        self._load_model = load_model
+        self._auth_token = auth_token
+        self._idle_seconds = None if idle_timeout_min is None else idle_timeout_min * 60
+        self._model: Model | None = None
+        self._worker: ThreadPoolExecutor | None = None
+        self._stop: Future[None] = Future()
+        self._sessions: set[asyncio.Task] = set()
+        self._last_activity = time.monotonic()
+
+    def _authorized(self, headers: Mapping[str, str]) -> bool:
+        if self._auth_token is None:
+            return True
+        value = headers.get(protocol.AUTH_HEADER.lower(), '')
+        return hmac.compare_digest(value.encode(), protocol.bearer(self._auth_token).encode())
+
+    def _keepalive(self) -> int | None:
+        self._last_activity = time.monotonic()
+        return None if self._idle_seconds is None else math.floor(self._idle_seconds)
+
+    async def _call(self, function: Callable[[], T]) -> T:
+        """Keep an operation alive until it finishes, including while its caller is cancelled."""
+        assert self._worker is not None
+        pending = asyncio.get_running_loop().run_in_executor(self._worker, function)
+        try:
+            return await asyncio.shield(pending)
+        except asyncio.CancelledError:
+            await pending
+            raise
+
+    async def _answer(self, conn: server_wire.ServerConnection, session: Session, session_id: str) -> None:
+        while True:
+            raw = await conn.receive()
+            self._last_activity = time.monotonic()
+            started = time.perf_counter_ns()
+            request = serialization.deserialise(raw)
+            if request[protocol.SESSION_ID] != session_id:
+                raise ValueError('The session ID does not belong to this connection')
+            if request.get(protocol.END_SESSION) is True:
+                return
+            observation = request[protocol.OBSERVATION]
+            decoded = time.perf_counter_ns()
+            await conn.send(await self._infer_response(session, observation, started, decoded))
+
+    async def _infer_response(self, session: Session, observation: Any, started: int, decoded: int) -> bytes:
+        timing = {protocol.TIMING_DECODE: (decoded - started) / 1e6}
+
+        def infer() -> bytes:
+            called = time.perf_counter_ns()
+            timing[protocol.TIMING_QUEUED] = (called - decoded) / 1e6
+            try:
+                result = session.infer(observation)
+            finally:
+                timing[protocol.TIMING_MODEL] = timing[protocol.TIMING_INFER] = (time.perf_counter_ns() - called) / 1e6
+            timing[protocol.TIMING_SERVED] = (time.perf_counter_ns() - started) / 1e6
+            # Pack before releasing the worker: native results can reference reusable model buffers.
+            result = serialization.encode_images(result, session.output_images)
+            return serialization.serialise({protocol.RESULT: result, protocol.TIMING: timing})
+
+        try:
+            return await self._call(infer)
+        except Exception as error:
+            logger.exception('Inference failed')
+            return serialization.serialise({protocol.ERROR: str(error)})
+
+    async def _prepare(self, conn: server_wire.ServerConnection, prepare: Callable[[], None]) -> None:
+        pending = asyncio.create_task(self._call(prepare))
+        try:
+            while not pending.done():
+                await conn.send(serialization.serialise({protocol.STATUS: protocol.ServerStatus.WAITING}))
+                await asyncio.wait({pending}, timeout=self.WAITING_INTERVAL_SEC)
+            await pending
+        finally:
+            # A disconnected opener still owns the session being prepared until its cleanup completes.
+            await asyncio.shield(pending)
+
+    @asynccontextmanager
+    async def _prepared_session(self, conn: server_wire.ServerConnection, params: dict[str, Any]):
+        model = self._model
+        assert model is not None
+        session: Session | None = None
+
+        def prepare() -> None:
+            nonlocal session
+            session = model.prepare_session(json.loads(json.dumps(params, allow_nan=False)))
+            spec.validate(session.client_stack)
+
+        try:
+            await self._prepare(conn, prepare)
+            assert session is not None
+            yield session
+        finally:
+            if session is not None:
+                await self._call(session.close)
+
+    async def _exchange_session(self, conn: server_wire.ServerConnection) -> None:
+        model = self._model
+        assert model is not None
+        requested = spec.parse_params(conn.query_params)
+        params = spec.resolve_params(model.parameters, requested)
+        session_id = uuid4().hex
+        async with self._prepared_session(conn, params) as session:
+            metadata = {
+                **conn.served_address.meta,
+                **model.metadata,
+                **session.metadata,
+                keys.LOCAL_STACK: session.client_stack,
+                keys.SESSION_PARAMS: requested,
+                keys.EFFECTIVE_PARAMS: params,
+                keys.MODEL_SERVER_VERSION: version('positronic-model-server'),
+            }
+            await conn.send(
+                serialization.serialise({
+                    protocol.STATUS: protocol.ServerStatus.READY,
+                    protocol.PROTOCOL_VERSION: protocol.ProtocolVersion.V3,
+                    protocol.SESSION_ID: session_id,
+                    protocol.META: metadata,
+                })
+            )
+            await self._answer(conn, session, session_id)
+        await conn.send(serialization.serialise({protocol.SESSION_ID: session_id, protocol.END_SESSION: True}))
+
+    async def _serve_session(self, conn: server_wire.ServerConnection) -> None:
+        task = asyncio.current_task()
+        assert task is not None
+        self._sessions.add(task)
+        self._last_activity = time.monotonic()
+        try:
+            await self._exchange_session(conn)
+        except wire.PeerDisconnected:
+            logger.info('Session disconnected: %s', conn.peer)
+        except Exception as error:
+            logger.exception('Session failed: %s', conn.peer)
+            try:
+                await conn.send(
+                    serialization.serialise({protocol.STATUS: protocol.ServerStatus.ERROR, protocol.ERROR: str(error)})
+                )
+                await conn.refuse(str(error))
+            except wire.PeerDisconnected:
+                logger.info('Session disconnected before receiving its error: %s', conn.peer)
+        finally:
+            self._sessions.discard(task)
+            self._last_activity = time.monotonic()
+
+    async def _idle_watchdog(self) -> None:
+        assert self._idle_seconds is not None
+        while True:
+            await asyncio.sleep(min(self._idle_seconds, 30))
+            if not self._sessions and time.monotonic() - self._last_activity >= self._idle_seconds:
+                return
+
+    def _load(self) -> None:
+        self._model = self._load_model()
+        spec.resolve_params(self._model.parameters, {})
+
+    async def _run(self, wires: Sequence[server_wire.Wire], on_ready: Callable[[], None] | None) -> None:
+        started: list[server_wire.Wire] = []
+        serving: list[asyncio.Task] = []
+        ending: list[asyncio.Future] = []
+
+        try:
+            await self._call(self._load)
+            if self._stop.cancelled():
+                return
+            for transport in wires:
+                started.append(transport)
+                await transport.start(self._serve_session, self._keepalive, self._authorized)
+            self._last_activity = time.monotonic()
+            serving = [asyncio.create_task(transport.serve()) for transport in started]
+            ending = [asyncio.wrap_future(self._stop)]
+            if self._idle_seconds is not None:
+                ending.append(asyncio.create_task(self._idle_watchdog()))
+            if on_ready is not None:
+                on_ready()
+            await asyncio.wait(serving + ending, return_when=asyncio.FIRST_COMPLETED)
+        finally:
+            for task in ending:
+                task.cancel()
+            await asyncio.gather(*ending, return_exceptions=True)
+            await self._finish(started, serving)
+
+    async def _finish(self, started: Sequence[server_wire.Wire], serving: Sequence[asyncio.Task]) -> None:
+        stops = await asyncio.gather(*(transport.stop() for transport in started), return_exceptions=True)
+        for stopped, task in zip(stops, serving, strict=False):
+            if isinstance(stopped, BaseException):
+                task.cancel()
+        outcomes = await asyncio.gather(*serving, return_exceptions=True)
+        for task in self._sessions:
+            if not task.cancelling():
+                task.cancel()
+        sessions = await asyncio.gather(*self._sessions, return_exceptions=True)
+        failures = [value for value in [*stops, *sessions, *outcomes] if isinstance(value, Exception)]
+        try:
+            if self._model is not None:
+                await self._call(self._model.close)
+        except Exception as error:
+            failures.append(error)
+        finally:
+            self._model = None
+        for error in failures[1:]:
+            logger.error('Additional shutdown failure', exc_info=error)
+        if failures:
+            raise failures[0]
+
+    def serve(self, wires: Sequence[server_wire.Wire], on_ready: Callable[[], None] | None = None) -> None:
+        """Load and warm once, then serve until shutdown, idle expiry, or a transport ends."""
+        if not wires:
+            raise ValueError('At least one transport is required')
+        if self._worker is not None:
+            raise RuntimeError('This server is already running')
+        with ThreadPoolExecutor(max_workers=1, thread_name_prefix='model') as worker:
+            self._worker = worker
+            try:
+                asyncio.run(self._run(wires, on_ready))
+            except KeyboardInterrupt:
+                logger.info('Server interrupted')
+            finally:
+                self._worker = None
+                self._stop = Future()
+
+    def shutdown(self) -> None:
+        """Request shutdown from any thread; running model work finishes before cleanup."""
+        self._stop.cancel()