Approach: use a single background batching thread that collects incoming requests into a batch protected by a Condition. submit() returns a concurrent.futures.Future per request. The worker flushes when batch size reaches max_batch_size or when max_latency_ms elapses since first item in batch. Backpressure via max_queue_size (submit blocks or raises). Never call model.predict while holding the main lock to avoid deadlocks. Support cancellation by checking future.cancelled() before including it in batch and setting exceptions for cancelled items.python
import threading
import time
from concurrent.futures import Future
from typing import Any, List
class BatchRequest:
def __init__(self, item):
self.item = item
self.future = Future()
class BatchingPredictor:
def __init__(self, model, max_batch_size=32, max_latency_ms=50, max_queue_size=1000):
self.model = model
self.max_batch_size = max_batch_size
self.max_latency = max_latency_ms / 1000.0
self.max_queue_size = max_queue_size
self.lock = threading.Lock()
self.cv = threading.Condition(self.lock)
self.queue: List[BatchRequest] = []
self.stopped = False
self.worker = threading.Thread(target=self._worker, daemon=True)
self.worker.start()
def submit(self, item, block=True, timeout=None) -> Future:
req = BatchRequest(item)
with self.cv:
start = time.time()
while len(self.queue) >= self.max_queue_size:
if not block:
raise RuntimeError("Queue full")
remaining = None if timeout is None else timeout - (time.time() - start)
if remaining is not None and remaining <= 0:
raise TimeoutError("submit timeout due to backpressure")
self.cv.wait(remaining)
self.queue.append(req)
# If first item, record arrival time on container via attribute
if len(self.queue) == 1:
self._first_arrival = time.time()
# notify worker to re-evaluate flush conditions
self.cv.notify()
return req.future
def shutdown(self, wait=True):
with self.cv:
self.stopped = True
self.cv.notify_all()
if wait:
self.worker.join()
def _worker(self):
while True:
with self.cv:
if self.stopped and not self.queue:
return
now = time.time()
if not self.queue:
# wait until new request or stop
self.cv.wait()
continue
# compute remaining time until latency deadline
elapsed = now - getattr(self, "_first_arrival", now)
time_left = max(0.0, self.max_latency - elapsed)
# flush if batch full
if len(self.queue) >= self.max_batch_size:
batch = self._pop_batch(self.max_batch_size)
else:
# wait until either more arrives or latency expires or stop signalled
self.cv.wait(time_left)
# re-evaluate conditions
now2 = time.time()
elapsed2 = now2 - getattr(self, "_first_arrival", now2)
if len(self.queue) >= self.max_batch_size or elapsed2 >= self.max_latency or self.stopped:
batch = self._pop_batch(self.max_batch_size)
else:
continue
# At this point we have a batch list and are outside lock
if not batch:
continue
inputs = []
futures = []
for req in batch:
if req.future.cancelled():
# skip cancelled requests
req.future.set_exception(RuntimeError("Cancelled before dispatch"))
continue
inputs.append(req.item)
futures.append(req.future)
try:
if inputs:
outputs = self.model.predict(inputs)
# expect outputs align with inputs
for fut, out in zip(futures, outputs):
if not fut.cancelled():
fut.set_result(out)
# for any skipped inputs (cancelled) nothing to do
except Exception as e:
for fut in futures:
if not fut.cancelled():
fut.set_exception(e)
def _pop_batch(self, max_n):
# must be called with lock held
n = min(len(self.queue), max_n)
batch = self.queue[:n]
self.queue = self.queue[n:]
# reset first arrival time if queue still has items
if self.queue:
self._first_arrival = time.time()
else:
self._first_arrival = None
# wake up any submitters blocked by backpressure
self.cv.notify_all()
return batch
