Principle:Huggingface Datatrove LLM Inference Execution
| Knowledge Sources | |
|---|---|
| Domains | Inference, Synthetic_Data |
| Last Updated | 2026-02-14 00:00 GMT |
Overview
Orchestrating concurrent LLM inference across documents using user-defined rollout functions for scalable synthetic data generation.
Description
Inference execution is the core processing step of the datatrove synthetic data generation pipeline. It manages the concurrent processing of documents through user-defined rollout functions, which define the application-specific logic for interacting with an LLM. Each rollout function receives a document and a generate callback, sends one or more requests to the inference server, and returns structured results that are stored in the document's metadata.
The execution model is built around several key design decisions:
- Rollout function protocol: Rollout functions follow a defined async protocol: Template:Code. The Template:Code callback is a Template:Code that sends a request payload to the inference server and returns the parsed result. This design decouples the application logic (what to ask the LLM) from the infrastructure (how to send requests).
- Multi-rollout aggregation: Each document can be processed through multiple independent rollouts (controlled by Template:Code). All rollouts for a single document execute concurrently via Template:Code, and results are collected into a list stored under the configured metadata key. If all rollouts return Template:Code, the document is marked for removal.
- Concurrency control: Two levels of concurrency are managed: (1) document-level concurrency via a task pool bounded by Template:Code, and (2) request-level concurrency via an Template:Code bounded by Template:Code. This two-tier approach prevents overwhelming the inference server while maximizing throughput.
- Shared context: Rollout functions can receive additional keyword arguments from a shared context (e.g., process pools, coding environments). The shared context can be a plain dict, a callable, or a context manager, providing flexibility for resource initialization and cleanup.
- Error handling: The runner distinguishes between retryable errors (Template:Code, Template:Code, Template:Code) that trigger exponential backoff, non-retryable inference errors (Template:Code) that fail the request immediately, and server errors (Template:Code) that indicate the server has died.
- Metrics collection: Real-time metrics are tracked for token throughput (prompt and completion), request success/failure rates, and queue sizes (waiting vs. running requests). Metrics are reported at configurable intervals with both lifetime and windowed (5-minute) rates.
Usage
Inference execution is the core step of synthetic data generation pipelines:
- After server launch, as the main document processing loop
- When generating annotations, labels, summaries, or other synthetic data for each document
- When running multi-turn or multi-step reasoning workflows that require multiple LLM calls per document
- In combination with checkpointing for long-running jobs that may be interrupted
Theoretical Basis
The execution model follows an async concurrent processing pattern with configurable parallelism:
Rollout Function Protocol:
class RolloutFunction(Protocol):
def __call__(
self,
document: Document,
generate: GenerateFunction,
**kwargs: Any,
) -> Awaitable[RolloutResult]: ...
The rollout function is the user's entry point into the inference pipeline. It receives:
- document: The input document to process
- generate: An async callback that sends a request payload (dict) to the inference server and returns an Template:Code containing text, finish reason, and usage statistics
- **kwargs: Additional keyword arguments from the shared context
Processing Pipeline:
- Data ingestion: Synchronous document iterators are wrapped in an async generator via a thread pool executor, allowing the pipeline to start processing without blocking on I/O
- Document dispatch: Documents are dispatched to the task pool up to Template:Code; when the pool is full, the pipeline waits for at least one task to complete
- Rollout execution: For each document, Template:Code rollout tasks are created and executed concurrently
- Request caching: Each generate call checks the request cache before sending to the server, preventing redundant inference on job restart
- Result aggregation: Rollout results are collected, filtered (removing Template:Code values), and stored in Template:Code
- Checkpoint and write: Processed documents are written to both the output writer and checkpoint files atomically
Retry Strategy:
Retryable errors use exponential backoff with a formula of Template:Code seconds, up to 6 retries. This yields delays of 5, 10, 20, 40, 80, 160 seconds across attempts.