Implementation:Huggingface Datatrove InferenceProgressMonitor

Knowledge Sources	InferenceProgressMonitor source Progress helper functions
Domains	Monitoring, Observability, HuggingFace Hub
Last Updated	2026-02-14 00:00 GMT

Overview

Pipeline step that monitors inference job progress by counting uploaded documents, calculating ETAs, and updating the HuggingFace dataset card with a visual progress bar.

Description

Template:Code is a Template:Code dataclass that runs a polling loop to track and report progress of a parallel inference job. It operates independently of the inference runner, typically in a separate process or Slurm job, and communicates progress by reading from and writing to the HuggingFace Hub.

The monitor's main loop:

1. Check if Template:Code exists (indicates job completion)
2. Optionally check if the Slurm inference job is still running via Template:Code
3. Count documents in the output repository by reading Parquet file metadata headers via Template:Code
4. Calculate progress percentage and ETA using linear throughput projection
5. Render a progress bar and update the dataset card via Template:Code
6. Sleep for Template:Code seconds and repeat

Key supporting functions:

• count_documents_in_repo(): Reads Parquet metadata headers (not data) from the Hub to count rows efficiently. Uses Template:Code with explicit cache invalidation to ensure fresh file listings.
• get_total_expected_documents(): Determines the total input dataset size, trying builder metadata first (fastest), then falling back to loading the dataset, then using Template:Code as a last resort.
• render_progress_bar(): Creates a 20-character visual progress bar with filled/empty dots, percentage, document counts, time remaining, and projected completion datetime.
• calculate_eta(): Linear ETA projection based on current throughput (docs/sec).
• format_time_remaining(): Converts seconds to human-readable format (e.g., "1h 30m").

Usage

Use InferenceProgressMonitor when:

• Running long inference jobs where progress visibility is needed on the Hub
• Monitoring Slurm-submitted inference jobs from a separate monitoring job
• Providing ETA estimates for resource planning and coordination

Code Reference

Source Location

• Repository: huggingface/datatrove
• InferenceProgressMonitor: src/datatrove/pipeline/inference/progress_monitor.py:L255-362
• Helper functions: src/datatrove/pipeline/inference/progress_monitor.py:L30-253

Signature

@dataclass
class InferenceProgressMonitor(PipelineStep):
    """Monitor dataset generation progress and update dataset card periodically."""

    # Dataset card parameters
    params: InferenceDatasetCardParams

    # Monitoring parameters
    inference_job_id: str | None = None
    max_examples: int = -1
    update_interval: int = 3600  # 1 hour

    name: str = "InferenceProgressMonitor"
    type: str = "Monitor"

    def _is_job_running(self, job_id: str) -> bool:
        """Check if a Slurm job is still running or pending via squeue."""
        ...

    def run(self, data=None, rank: int = 0, world_size: int = 1):
        """Monitor progress and update dataset card until completion. Only runs on rank 0."""
        ...

Import

from datatrove.pipeline.inference.progress_monitor import InferenceProgressMonitor
from datatrove.pipeline.inference.dataset_card_generator import InferenceDatasetCardParams

I/O Contract

Inputs

Name	Type	Required	Description
params	InferenceDatasetCardParams	Yes	Dataset card metadata including Template:Code (monitored for progress) and Template:Code (checked for completion)
inference_job_id	str / None	No	Slurm job ID to monitor for completion/failure detection via Template:Code
max_examples	int	No	Maximum expected documents to process (-1 for all documents in the input dataset)
update_interval	int	No	Seconds between progress checks and card updates (default: 3600)
data	Iterable / None	No	Optional passthrough data from previous pipeline step
rank	int	Yes	Process rank; monitoring only executes on rank 0

Outputs

Name	Type	Description
Updated dataset card	README.md on Hub	Periodically updated dataset card at Template:Code with progress bar, document counts, ETA, and projected completion time
Passthrough data	Iterable	Input data yielded unchanged (transparent pipeline step)
Log output	Log messages	Progress messages including document counts, percentages, and update status

Usage Examples

Example: Monitoring an inference job from a separate Slurm job

from datatrove.pipeline.inference.progress_monitor import InferenceProgressMonitor
from datatrove.pipeline.inference.dataset_card_generator import InferenceDatasetCardParams

card_params = InferenceDatasetCardParams(
    output_repo_id="my-org/synthetic-summaries",
    input_dataset_name="my-org/source-documents",
    input_dataset_split="train",
    input_dataset_config=None,
    prompt_column="text",
    prompt_template=None,
    system_prompt="You are a helpful assistant.",
    model_name="meta-llama/Llama-3.1-70B-Instruct",
    model_revision="main",
    generation_kwargs={
        "temperature": 0.7,
        "max_tokens": 1024,
        "model_max_context": 8192,
    },
    spec_config=None,
    stats_path="/shared/output/stats.json",
)

monitor = InferenceProgressMonitor(
    params=card_params,
    inference_job_id="12345678",  # Slurm job ID of the inference job
    max_examples=-1,  # Process all documents
    update_interval=1800,  # Update every 30 minutes
)

# Run as a pipeline step (blocks until inference completes or fails)
monitor.run(rank=0, world_size=1)

Example: Progress bar output

During monitoring, the dataset card is updated with a progress section like:

## Generation Progress

[XXXXXXXXXXXXXXXXXXOO] 60% * 3,000/5,000 documents processed * ETA 2h 15m * Nov 27, 18:30 UTC

*Last updated: 2026-02-14 12:30:00 UTC*

Example: Using helper functions directly

from datatrove.pipeline.inference.progress_monitor import (
    count_documents_in_repo,
    get_total_expected_documents,
    render_progress_bar,
)
import time

# Count documents currently in the output repo
completed = count_documents_in_repo("my-org/synthetic-summaries")

# Get expected total from source dataset
total = get_total_expected_documents(
    dataset_name="my-org/source-documents",
    split="train",
    config=None,
    max_examples=-1,
)

# Render progress bar
start_time = time.time() - 3600  # Started 1 hour ago
bar = render_progress_bar(completed, total, start_time, time.time())
print(bar)

Related Pages

Implements Principle

• Principle:Huggingface_Datatrove_Inference_Progress_Monitoring
• Environment:Huggingface_Datatrove_Inference_GPU_Environment