Implementation:Huggingface Datatrove InferenceProgressMonitor
| Knowledge Sources | |
|---|---|
| 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:
- Check if Template:Code exists (indicates job completion)
- Optionally check if the Slurm inference job is still running via Template:Code
- Count documents in the output repository by reading Parquet file metadata headers via Template:Code
- Calculate progress percentage and ETA using linear throughput projection
- Render a progress bar and update the dataset card via Template:Code
- 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)