Skip to content

Actorpool progress bar #1457

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ayushdg merged 5 commits into from
Feb 4, 2026
Merged

Actorpool progress bar #1457

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
6610d22
Add progress bars for actor pool executor
ayushdg Feb 3, 2026
4b5d249
Merge branch 'main' of github.com:NVIDIA/NeMo-Curator into actorpool-…
ayushdg Feb 3, 2026
2521586
Merge branch 'main' into actorpool-progressbar
ayushdg Feb 3, 2026
0e42c3d
Merge branch 'main' into actorpool-progressbar
ayushdg Feb 3, 2026
b563bcb
Merge branch 'main' into actorpool-progressbar
ayushdg Feb 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 32 additions & 5 deletions nemo_curator/backends/experimental/ray_actor_pool/executor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
import ray
from loguru import logger
from ray.util.actor_pool import ActorPool
from tqdm import tqdm

@greptile-apps greptile-apps Bot Feb 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P1] tqdm import may not be declared dependency for this backend

This adds from tqdm import tqdm, but if tqdm isn’t in this repo’s declared runtime deps for the experimental ray actor pool backend, users will hit ModuleNotFoundError at import time. Consider either adding tqdm to the appropriate dependency set, or making the progress bar optional via a lazy import / fallback when show_progress=True.


from nemo_curator.backends.base import BaseExecutor
from nemo_curator.backends.experimental.utils import RayStageSpecKeys, execute_setup_on_node
Expand Down Expand Up @@ -59,8 +60,24 @@ class RayActorPoolExecutor(BaseExecutor):
4. Provides better backpressure management through ActorPool
"""

def __init__(self, config: dict | None = None, ignore_head_node: bool = False):
def __init__(
self,
config: dict | None = None,
ignore_head_node: bool = False,
show_progress: bool = True,
progress_interval: float = 10.0,
):
Comment on lines +63 to +69

@greptile-apps greptile-apps Bot Feb 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] progress_interval accepts invalid values

tqdm(..., mininterval=self.progress_interval) will raise or behave unexpectedly if progress_interval is <= 0 or non-finite. Since this is a public constructor arg, it’s worth either validating it (e.g., > 0) or documenting constraints.

"""Initialize the Ray Actor Pool executor.

Args:
config: Configuration dictionary for the executor.
ignore_head_node: If True, don't schedule tasks on the head node.
show_progress: If True, display tqdm progress bars during execution.
progress_interval: Minimum interval in seconds between progress bar updates.
"""
super().__init__(config, ignore_head_node)
self.show_progress = show_progress
self.progress_interval = progress_interval

def execute(self, stages: list["ProcessingStage"], initial_tasks: list[Task] | None = None) -> list[Task]: # noqa: PLR0912
"""Execute the pipeline stages using ActorPool.
Expand Down Expand Up @@ -293,8 +310,12 @@ def _process_stage_with_pool(

# Process each task and flatten the results since each task can produce multiple output tasks
all_results = []
for result_batch in actor_pool.map_unordered(
lambda actor, batch: actor.process_batch.remote(batch), task_batches
for result_batch in tqdm(
actor_pool.map_unordered(lambda actor, batch: actor.process_batch.remote(batch), task_batches),
total=len(task_batches),
desc=f"Processing {_stage.name}",
mininterval=self.progress_interval,
disable=not self.show_progress,
):
# result_batch is a list of tasks from processing a single input task
all_results.extend(result_batch)
Expand Down Expand Up @@ -322,8 +343,14 @@ def _process_shuffle_stage_with_rapidsmpf_actors(

# Step 1: Insert tasks into shuffler
_ = list(
actor_pool.map_unordered(
lambda actor, batch: actor.read_and_insert.remote(tasks=batch, **insert_kwargs), task_batches
tqdm(
actor_pool.map_unordered(
lambda actor, batch: actor.read_and_insert.remote(tasks=batch, **insert_kwargs), task_batches
),
total=len(task_batches),
desc="Inserting into shuffler",
mininterval=self.progress_interval,
disable=not self.show_progress,
)
)

Expand Down
Loading