Skip to main content
Dynamic request batching for GPU inference. Implements a two-stage pipeline:
  1. Deserialization (parallel, per-request): Each incoming request is decoded, transformed, and validated in a thread pool. Failures are returned immediately as per-request errors (HTTP 400) without entering the batch queue.
  2. Batching (GPU, shared): Successfully prepared samples are accumulated into batches and dispatched to the GPU. A GPU-level failure (OOM, CUDA error) is shared by all requests in the batch and mapped to HTTP 503 via _ForwardError.
This separation guarantees that one user’s malformed payload can never cause failures for other users’ valid requests.

BatchScheduler

Two-stage pipeline: parallel deserialization then batched GPU inference. Stage 1 — Deserialization (per-request, parallel): Each predict call submits the request to a thread pool for image decoding, transform, shape validation, and (for M-Optimus) bulk RNA preprocessing. If any step fails, the request’s future is immediately resolved with a ValueError (→ HTTP 400) and the request never enters the batch queue. Stage 2 — Batching (shared GPU): Valid PreparedSample objects are collected into batches of up to max_batch_size and dispatched to the GPU. A forward failure raises _ForwardError (a RuntimeError subclass → HTTP 503) to all requests in that batch.
The server-side model to batch requests for.
Maximum tiles per GPU forward pass.
Maximum milliseconds to wait before dispatching a partial batch.
Thread pool size for parallel deserialization.
"prediction", "embedding", or "prediction_with_embedding".

shutdown

Releases resources owned by the scheduler. Shuts down the preprocessing thread pool (without waiting) and cancels the batch-loop task if it is still pending. Safe to call multiple times and from outside the scheduler’s event loop (e.g. when pruning a scheduler whose loop has already closed).

predict

Deserializes, validates, and enqueues a request for batched inference. Deserialization runs in a thread pool so the event loop is not blocked. If deserialization fails, the error is raised directly to the caller (HTTP 400). Otherwise the prepared sample enters the batch queue and awaits GPU execution.
The tile request to process.
Returns: The model response for this tile.