TokenBatcher

Token-aware batcher for LLM inference workloads.

Constructor

Signature

def TokenBatcher(
        self,
        inference_fn: ProcessFn[RecordT, ResultT] | None = None,
        *,
        process_fn: ProcessFn[RecordT, ResultT] | None = None,
        token_estimator: CostEstimatorFn[RecordT] | None = None,
        cost_estimator: CostEstimatorFn[RecordT] | None = None,
        target_batch_tokens: int | None = None,
        target_batch_cost: int = 32_000,
        default_token_estimate: int | None = None,
        default_cost: int = 1,
        max_batch_size: int = 256,
        min_batch_size: int = 1,
        batch_timeout_s: float = 0.05,
        max_queue_size: int = 5_000,
        prefetch_batches: int = 2,
    ): ...

Parameters

Name	Type	Description
inference_fn	`ProcessFn[RecordT, ResultT]	None` = None
process_fn	`ProcessFn[RecordT, ResultT]	None` = None
token_estimator	`CostEstimatorFn[RecordT]	None` = None
cost_estimator	`CostEstimatorFn[RecordT]	None` = None
target_batch_tokens	`int	None` = None
target_batch_cost	int = 32000	The target total cost per batch.
default_token_estimate	`int	None` = None
default_cost	int = 1	Default cost if estimation fails.
max_batch_size	int = 256	Maximum number of records in a single batch.
min_batch_size	int = 1	Minimum number of records required to trigger processing before timeout.
batch_timeout_s	float = 0.05	Time to wait in seconds before processing a partial batch.
max_queue_size	int = 5000	Maximum number of records allowed in the queue.
prefetch_batches	int = 2	Number of batches to prefetch for processing.

---

Methods

---

submit()

async def submit(
        self,
        record: RecordT,
        *,
        estimated_tokens: int | None = None,
        estimated_cost: int | None = None,
    ) -> asyncio.Future[ResultT]: ...

Submit a single record for batched inference. Accepts either estimated_tokens or estimated_cost.

Parameters

Name	Type	Description
record	RecordT	The input record.
estimated_tokens	`int	None`
estimated_cost	`int	None`

Returns

Type	Description
asyncio.Future[ResultT]	A future whose result is the corresponding entry from the list returned by the inference function.

---