69 lines
2.1 KiB
Python
69 lines
2.1 KiB
Python
"""Dynamic batching utilities for face detection inference."""
|
|
|
|
import logging
|
|
import time
|
|
from collections import deque
|
|
from typing import Callable, Deque, Generic, List, Optional, TypeVar
|
|
|
|
from gpu_manager import GPUMemoryManager
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
T = TypeVar("T")
|
|
R = TypeVar("R")
|
|
|
|
|
|
class DynamicBatcher(Generic[T, R]):
|
|
"""Accumulate items into batches and flush based on size or timeout."""
|
|
|
|
def __init__(
|
|
self,
|
|
process_fn: Callable[[List[T]], List[R]],
|
|
gpu_manager: Optional[GPUMemoryManager] = None,
|
|
max_batch_size: int = 16,
|
|
batch_timeout_ms: float = 100.0,
|
|
min_batch_size: int = 1,
|
|
):
|
|
self.process_fn = process_fn
|
|
self.gpu_manager = gpu_manager
|
|
self.max_batch_size = max_batch_size
|
|
self.batch_timeout_ms = batch_timeout_ms
|
|
self.min_batch_size = min_batch_size
|
|
self._queue: Deque[T] = deque()
|
|
self._last_flush = time.monotonic()
|
|
|
|
def add(self, item: T) -> List[R]:
|
|
"""Add an item and return any results if a batch was flushed."""
|
|
self._queue.append(item)
|
|
if len(self._queue) >= self._current_batch_size():
|
|
return self.flush()
|
|
|
|
elapsed_ms = (time.monotonic() - self._last_flush) * 1000
|
|
if elapsed_ms >= self.batch_timeout_ms and len(self._queue) >= self.min_batch_size:
|
|
return self.flush()
|
|
return []
|
|
|
|
def flush(self) -> List[R]:
|
|
"""Flush all queued items through the processor."""
|
|
if not self._queue:
|
|
return []
|
|
|
|
batch = [self._queue.popleft() for _ in range(min(len(self._queue), self._current_batch_size()))]
|
|
results = self.process_fn(batch)
|
|
self._last_flush = time.monotonic()
|
|
|
|
if self.gpu_manager is not None:
|
|
self.gpu_manager.adjust_batch_size()
|
|
self.gpu_manager.empty_cache()
|
|
|
|
return results
|
|
|
|
def _current_batch_size(self) -> int:
|
|
if self.gpu_manager is not None:
|
|
return min(self.gpu_manager.current_batch_size, self.max_batch_size)
|
|
return self.max_batch_size
|
|
|
|
@property
|
|
def queued_count(self) -> int:
|
|
return len(self._queue)
|