Implementation:Huggingface Datatrove CheckpointManager

Knowledge Sources	CheckpointManager source RequestCache source
Domains	Fault_Tolerance, Distributed_Computing, Data Persistence
Last Updated	2026-02-14 00:00 GMT

Overview

Manages chunk-based document checkpointing and SQLite-backed request caching for fault-tolerant inference job recovery.

Description

The checkpointing module provides two classes that work together for fault tolerance:

RequestCache is a lightweight SQLite-backed cache that stores individual inference request/response pairs. It uses an async writer queue (Template:Code) to batch SQLite writes without blocking the inference loop. Payloads are hashed using xxHash-128 for fast, collision-resistant key generation. The cache supports storing both successful results (as BLOB-serialized JSON via orjson) and error messages (as TEXT), enabling error replay on restart.

CheckpointManager manages chunk-based document persistence. It maintains per-chunk counters, file locks (one Template:Code per chunk), and a set of newly-completed chunks. When a chunk reaches its capacity, the manager closes the corresponding output file, records the chunk as complete, and cleans up local checkpoint state. On restart, it scans for existing checkpoint files, re-uploads their contents to the output writer, and returns skip information so the runner knows where to resume.

Key implementation details:

Async file I/O: Checkpoint JSONL files are written using Template:Code for non-blocking I/O
Per-chunk locking: Each chunk has its own Template:Code to allow concurrent writes to different chunks
Monotonic completion: Chunks must complete in order; if chunk N+1 finishes before chunk N, the manager waits until N completes before advancing Template:Code
Automatic cleanup: Completed chunk files and their request cache entries are deleted immediately to prevent unbounded disk usage

Usage

Use CheckpointManager when:

Building fault-tolerant inference pipelines that may be interrupted
Processing datasets too large to restart from the beginning on failure
Running on preemptible compute where job interruption is expected

Code Reference

Source Location

Repository: huggingface/datatrove
RequestCache: src/datatrove/pipeline/inference/checkpointing.py:L38-204
CheckpointManager: src/datatrove/pipeline/inference/checkpointing.py:L206-353

Signature

class RequestCache:
    def __init__(self, checkpoints_local_dir: str | None):
        ...

    async def initialize(self, rank: int) -> None:
        """Create SQLite database and start async writer loop."""
        ...

    def prepare_payload(self, payload: dict[str, Any]) -> str:
        """Compute xxHash-128 hex digest of JSON-serialized payload."""
        ...

    async def get_cached_response(
        self, doc_id: str, rollout_idx: int, *, payload_hash: str,
    ) -> tuple[dict[str, Any] | None, str | None]:
        """Return (result_dict, error_message) or (None, None) if not cached."""
        ...

    async def store_result(
        self, chunk_index: int, doc_id: str, rollout_idx: int,
        result: dict[str, Any], *, payload_hash: str,
    ) -> None:
        ...

    async def store_error(
        self, chunk_index: int, doc_id: str, rollout_idx: int,
        error_message: str, *, payload_hash: str,
    ) -> None:
        ...

    async def drop_chunk(self, chunk_index: int) -> None:
        """Delete all cache entries for a completed chunk."""
        ...

    async def mark_document_complete(self, doc_id: str) -> None:
        """Remove doc_id from in-memory cache index."""
        ...

    async def close(self, delete_file: bool = False) -> None:
        ...


class CheckpointManager:
    def __init__(
        self,
        checkpoints_local_dir: str | None = None,
        records_per_chunk: int = 6000,
        request_cache: RequestCache | None = None,
    ):
        ...

    async def write_document(
        self, document: Document, rank: int, chunk_index: int,
        output_writer_context: DiskWriter,
    ):
        """Write document to checkpoint and output writer. Close chunk if full."""
        ...

    async def parse_existing_checkpoints(
        self, rank: int, output_writer_context: DiskWriter,
    ) -> tuple[int, set[str]]:
        """Load existing checkpoints and return (docs_to_skip, processed_ids)."""
        ...

    async def cleanup_last_chunk(self, rank: int, chunk_index: int):
        """Finalize the last chunk after all documents are processed."""
        ...

    def chunk_index_gen(self):
        """Generator yielding chunk indices for sequential document assignment."""
        ...

Import

from datatrove.pipeline.inference.checkpointing import CheckpointManager, RequestCache

I/O Contract

Inputs

Name	Type	Required	Description
checkpoints_local_dir	str / None	No	Local directory for checkpoint files. Must be a local path (not remote). When None, checkpointing is disabled.
records_per_chunk	int	Yes	Number of documents per checkpoint chunk (must be positive, default: 6000)
request_cache	RequestCache / None	No	SQLite-backed request cache instance for request-level deduplication
document	Document	Yes (for write)	Processed document object to be checkpointed
rank	int	Yes	Process rank identifier, used for directory partitioning
chunk_index	int	Yes (for write)	Target chunk index for the document

Outputs

Name	Type	Description
Checkpoint JSONL files	Local files	Files at Template:Code containing serialized documents
SQLite request cache	Local database	File at Template:Code with cached request/response pairs
last_chunk_index file	Text file	File at Template:Code containing the index of the last fully completed chunk
docs_to_skip	int	Number of documents from completed chunks to skip on restart
processed_ids	set[str]	Set of document IDs from partially-completed chunks (already processed, skip on restart)

Usage Examples

Example: CheckpointManager integrated with InferenceRunner

from datatrove.pipeline.inference.checkpointing import CheckpointManager, RequestCache

# Initialize cache and manager
request_cache = RequestCache(checkpoints_local_dir="/tmp/inference_checkpoints")
checkpoint_mgr = CheckpointManager(
    checkpoints_local_dir="/tmp/inference_checkpoints",
    records_per_chunk=6000,
    request_cache=request_cache,
)

# On startup: parse existing checkpoints for resumption
docs_to_skip, processed_ids = await checkpoint_mgr.parse_existing_checkpoints(
    rank=0, output_writer_context=writer
)

# During processing: write each document
await checkpoint_mgr.write_document(doc, rank=0, chunk_index=0, output_writer_context=writer)

# After all documents: finalize
await checkpoint_mgr.cleanup_last_chunk(rank=0, chunk_index=last_chunk)

Example: RequestCache for deduplicating requests

from datatrove.pipeline.inference.checkpointing import RequestCache

cache = RequestCache(checkpoints_local_dir="/tmp/inference_checkpoints")
await cache.initialize(rank=0)

# Check cache before sending request
payload = {"messages": [{"role": "user", "content": "Hello"}], "max_tokens": 64}
payload_hash = cache.prepare_payload(payload)
cached_result, cached_error = await cache.get_cached_response(
    doc_id="doc_001", rollout_idx=0, payload_hash=payload_hash,
)

if cached_result is not None:
    # Use cached result, no inference needed
    pass
else:
    # Send request to server, then store result
    result = await send_to_server(payload)
    await cache.store_result(
        chunk_index=0, doc_id="doc_001", rollout_idx=0,
        result={"text": result.text, "finish_reason": result.finish_reason, "usage": result.usage},
        payload_hash=payload_hash,
    )

# Cleanup
await cache.close(delete_file=True)

Related Pages

Implements Principle

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment