Principle:Huggingface Datatrove Inference Checkpointing
| Knowledge Sources | |
|---|---|
| Domains | Fault_Tolerance, Distributed_Computing |
| Last Updated | 2026-02-14 00:00 GMT |
Overview
Providing fault-tolerant checkpointing and request-level caching for long-running inference jobs, enabling efficient recovery without redundant computation.
Description
Inference checkpointing addresses a critical challenge in large-scale synthetic data generation: jobs that process millions of documents over many hours must be able to resume efficiently after failures (hardware faults, preemption, OOM, etc.) without reprocessing already-completed work.
The checkpointing system has two complementary components:
- CheckpointManager: Saves processed documents in fixed-size chunks to local JSONL files. Each chunk corresponds to a contiguous range of input documents. When a chunk reaches Template:Code documents, it is marked as complete, its local checkpoint file is deleted (the data has already been written to the main output writer), and the Template:Code counter advances. On restart, the manager reads back any incomplete chunks, re-uploads their documents to the output writer, and skips already-completed chunks entirely.
- RequestCache: Stores individual API request/response pairs in a SQLite database, keyed by a combination of document ID, rollout index, and payload hash (computed via xxHash-128). When a request is made during a resumed job, the cache is checked first. If a matching entry exists, the cached result is returned immediately without contacting the inference server. This prevents redundant inference calls even for documents in partially-completed chunks.
Together, these two components provide a layered recovery strategy:
- Completed chunks are skipped entirely (no documents re-read, no requests re-sent)
- Partially-completed chunks have their documents re-read but their individual requests served from the SQLite cache
- Unprocessed documents proceed through normal inference
This approach minimizes both wasted compute (no redundant GPU inference) and wasted I/O (no redundant data reads for completed chunks).
Usage
Inference checkpointing is automatically engaged when:
- Template:Code is set in the InferenceRunner constructor
- The output writer's filename template includes Template:Code} (required for proper chunk-based file management)
Checkpointing is recommended for:
- Jobs expected to run for more than 1 hour
- Jobs running on preemptible or spot instances
- Jobs processing large datasets where reprocessing would be prohibitively expensive
- Multi-node jobs where individual node failures are more likely
Theoretical Basis
Chunk-based checkpointing:
Documents are grouped into sequential chunks of fixed size (default: 6000 documents). The chunk index is generated by a simple counter that yields the same index for Template:Code consecutive documents, then increments. This ensures that:
- Each chunk maps to exactly one output file (via the Template:Code} template variable)
- Chunks can be independently completed and cleaned up
- The Template:Code provides a single scalar that encodes how many documents have been fully processed
The Template:Code is persisted to a text file (Template:Code) within the checkpoint directory, providing an O(1) resume point.
Request-level caching:
The SQLite cache uses a composite primary key of Template:Code to uniquely identify each request. The payload hash is computed using xxHash-128 over the JSON-serialized, key-sorted request payload. This design handles:
- Multi-rollout documents: Different rollouts for the same document are cached independently
- Payload variations: If the rollout function sends different payloads for the same document on retry (e.g., different prompts), only exact matches are served from cache
- Error caching: BadRequestError responses are also cached, preventing the same invalid request from being retried
The cache uses an async writer queue to avoid blocking the main inference loop on SQLite writes. Reads are performed synchronously (within the async event loop) since SQLite reads from the same connection are fast.
Recovery sequence:
- Read Template:Code from persisted file -> skip that many chunks worth of input documents
- Scan local checkpoint directory for any remaining JSONL files (partially-completed chunks)
- Re-upload documents from partial chunks to the output writer
- Collect document IDs from partial chunks into a skip set
- Resume processing from the first unprocessed document, using the request cache for any in-flight requests