Add per-node worker sizing for download stages

Signed-off-by: Praateek <praateekm@gmail.com>

Author: praateekmahajan

nemo_curator/backends/ray_actor_pool/utils.py

@@ -18,7 +18,8 @@
 import ray
 from loguru import logger
 
-from nemo_curator.backends.utils import get_available_cpu_gpu_resources
+from nemo_curator.backends.utils import get_available_cpu_gpu_resources, get_stage_num_workers_per_node
+from nemo_curator.utils.ray_utils import get_alive_ray_node_count
 
 if TYPE_CHECKING:
     from ray.actor import ActorClass
@@ -59,6 +60,14 @@ def calculate_optimal_actors_for_stage(
         raise ValueError(msg)
 
     num_workers = stage.num_workers()
+    num_workers_per_node = get_stage_num_workers_per_node(stage)
+    if num_workers is not None and num_workers_per_node is not None:
+        msg = (
+            f"Stage {stage.name} defines both num_workers()={num_workers} and "
+            f"num_workers_per_node()={num_workers_per_node}. Use only one worker sizing option."
+        )
+        raise ValueError(msg)
+
     if num_workers is not None and num_workers > 0:
         if num_workers > max_actors_resources:
             msg = (
@@ -70,6 +79,19 @@ def calculate_optimal_actors_for_stage(
             return max_actors_resources
         return num_workers
 
+    if num_workers_per_node is not None:
+        num_nodes = get_alive_ray_node_count(ignore_head_node=ignore_head_node)
+        requested_actors = max(1, math.ceil(num_workers_per_node * num_nodes))
+        if requested_actors > max_actors_resources:
+            msg = (
+                f"Stage {stage.name} requires {requested_actors} actors from num_workers_per_node(), "
+                f"but only {max_actors_resources} fit with available resources. "
+                f"Capping actor count to {max_actors_resources}."
+            )
+            logger.warning(msg)
+            return max_actors_resources
+        return requested_actors
+
     number_of_batches = (
         math.ceil(num_tasks / stage.batch_size) if stage.batch_size is not None and stage.batch_size > 0 else num_tasks
     )

nemo_curator/backends/ray_data/adapter.py

@@ -13,15 +13,21 @@
 # limitations under the License.
 
 import copy
+import math
 from collections.abc import Callable
 from typing import Any
 
 from loguru import logger
-from ray.data import Dataset, TaskPoolStrategy
+from ray.data import ActorPoolStrategy, Dataset, TaskPoolStrategy
 
 from nemo_curator.backends.base import BaseStageAdapter
-from nemo_curator.backends.utils import RayStageSpecKeys, get_worker_metadata_and_node_id
+from nemo_curator.backends.utils import (
+    RayStageSpecKeys,
+    get_stage_num_workers_per_node,
+    get_worker_metadata_and_node_id,
+)
 from nemo_curator.stages.base import ProcessingStage
+from nemo_curator.utils.ray_utils import get_alive_ray_node_count
 
 from .utils import get_actor_compute_strategy_for_stage, get_configured_actor_pool_sizing_keys, is_actor_stage
 
@@ -39,8 +45,9 @@ class RayDataStageAdapter(BaseStageAdapter):
         c. Else we use tasks
     """
 
-    def __init__(self, stage: ProcessingStage):
+    def __init__(self, stage: ProcessingStage, ignore_head_node: bool = False):
         super().__init__(stage)
+        self.ignore_head_node = ignore_head_node
 
         self._batch_size = self.stage.batch_size
         if self._batch_size is None and self.stage.resources.gpus > 0:
@@ -90,6 +97,59 @@ def _build_resource_kwargs(self, ray_stage_spec: dict) -> dict[str, float]:
             kwargs["num_gpus"] = self.stage.resources.gpus  # type: ignore[reportArgumentType]
         return kwargs
 
+    def _get_per_node_pool_size(self, ray_stage_spec: dict, num_workers: int | None) -> int | None:
+        num_workers_per_node = get_stage_num_workers_per_node(self.stage)
+        if num_workers_per_node is None:
+            return None
+
+        if num_workers is not None:
+            msg = (
+                f"Stage {self.stage.name} defines both num_workers()={num_workers} and "
+                f"num_workers_per_node()={num_workers_per_node}. "
+                "Use only one worker sizing option."
+            )
+            raise ValueError(msg)
+
+        actor_pool_sizing_keys = get_configured_actor_pool_sizing_keys(ray_stage_spec)
+        if actor_pool_sizing_keys:
+            msg = (
+                f"Stage {self.stage.name} defines num_workers_per_node() "
+                f"and actor-pool sizing keys {actor_pool_sizing_keys}. Use only one worker sizing option."
+            )
+            raise ValueError(msg)
+
+        node_count = get_alive_ray_node_count(ignore_head_node=self.ignore_head_node)
+        if node_count <= 0:
+            msg = f"No alive Ray nodes available for num_workers_per_node sizing on stage {self.stage.name}."
+            raise ValueError(msg)
+        return max(1, math.ceil(num_workers_per_node * node_count))
+
+    def _build_actor_compute_kwargs(self, per_node_pool_size: int | None) -> dict[str, object]:
+        if per_node_pool_size is not None:
+            return {"compute": ActorPoolStrategy(size=per_node_pool_size)}
+        return {"compute": get_actor_compute_strategy_for_stage(self.stage)}
+
+    def _build_task_compute_kwargs(
+        self, ray_stage_spec: dict, per_node_pool_size: int | None, num_workers: int | None
+    ) -> dict[str, object]:
+        actor_pool_sizing_keys = get_configured_actor_pool_sizing_keys(ray_stage_spec)
+        if actor_pool_sizing_keys:
+            logger.warning(
+                f"Ignoring ray_stage_spec worker sizing keys {actor_pool_sizing_keys} "
+                f"for Ray Data task stage {self.stage.name}; these keys only apply to actor stages."
+            )
+
+        map_batches_kwargs: dict[str, object] = {}
+        if per_node_pool_size is not None:
+            map_batches_kwargs["compute"] = TaskPoolStrategy(size=per_node_pool_size)
+        elif num_workers is not None and num_workers > 0:
+            map_batches_kwargs["compute"] = TaskPoolStrategy(size=num_workers)
+
+        max_calls = ray_stage_spec.get(RayStageSpecKeys.MAX_CALLS_PER_WORKER)
+        if max_calls is not None:
+            map_batches_kwargs["max_calls"] = max_calls
+        return map_batches_kwargs
+
     def process_dataset(self, dataset: Dataset) -> Dataset:
         """Process a Ray Data dataset through this stage.
 
@@ -101,28 +161,15 @@ def process_dataset(self, dataset: Dataset) -> Dataset:
         """
         ray_stage_spec = self.stage.ray_stage_spec()
         stage_is_actor = ray_stage_spec.get(RayStageSpecKeys.IS_ACTOR_STAGE, is_actor_stage(self.stage))
+        num_workers = self.stage.num_workers()
+        per_node_pool_size = self._get_per_node_pool_size(ray_stage_spec, num_workers)
 
         if stage_is_actor:
-            map_batches_fn = create_actor_from_stage(self.stage)
-            map_batches_kwargs = {"compute": get_actor_compute_strategy_for_stage(self.stage)}
+            map_batches_fn = create_actor_from_stage(self.stage, ignore_head_node=self.ignore_head_node)
+            map_batches_kwargs = self._build_actor_compute_kwargs(per_node_pool_size)
         else:
-            map_batches_fn = create_task_from_stage(self.stage)
-            map_batches_kwargs = {}
-
-            actor_pool_sizing_keys = get_configured_actor_pool_sizing_keys(ray_stage_spec)
-            if actor_pool_sizing_keys:
-                logger.warning(
-                    f"Ignoring ray_stage_spec worker sizing keys {actor_pool_sizing_keys} "
-                    f"for Ray Data task stage {self.stage.name}; these keys only apply to actor stages."
-                )
-
-            num_workers = self.stage.num_workers()
-            if num_workers is not None and num_workers > 0:
-                map_batches_kwargs["compute"] = TaskPoolStrategy(size=num_workers)
-
-            max_calls = ray_stage_spec.get(RayStageSpecKeys.MAX_CALLS_PER_WORKER)
-            if max_calls is not None:
-                map_batches_kwargs["max_calls"] = max_calls
+            map_batches_fn = create_task_from_stage(self.stage, ignore_head_node=self.ignore_head_node)
+            map_batches_kwargs = self._build_task_compute_kwargs(ray_stage_spec, per_node_pool_size, num_workers)
 
         map_batches_kwargs.update(self._build_resource_kwargs(ray_stage_spec))
 
@@ -141,6 +188,9 @@ def process_dataset(self, dataset: Dataset) -> Dataset:
             )
             raise ValueError(msg)
 
+        if per_node_pool_size is not None:
+            map_batches_kwargs.setdefault("scheduling_strategy", "SPREAD")
+
         map_batches_kwargs.update(ray_remote_args)
 
         # Let Ray Data apply the selected compute strategy and resource requirements.
@@ -154,15 +204,15 @@ def process_dataset(self, dataset: Dataset) -> Dataset:
         return processed_dataset
 
 
-def create_actor_from_stage(stage: ProcessingStage) -> type[RayDataStageAdapter]:
+def create_actor_from_stage(stage: ProcessingStage, ignore_head_node: bool = False) -> type[RayDataStageAdapter]:
     """Create a StageProcessor class with the proper stage name for display."""
 
     class RayDataStageActorAdapter(RayDataStageAdapter):
         """Simplified stateful processor that wraps a ProcessingStage for Ray Data."""
 
         def __init__(self):
             """Initialize the stage processor."""
-            super().__init__(stage)
+            super().__init__(stage, ignore_head_node=ignore_head_node)
             self.setup_done = False
             node_info, worker_metadata = get_worker_metadata_and_node_id()
             self.setup_on_node(node_info, worker_metadata)
@@ -179,7 +229,9 @@ def __call__(self, batch: dict[str, Any]) -> dict[str, Any]:
     return RayDataStageActorAdapter
 
 
-def create_task_from_stage(stage: ProcessingStage) -> Callable[[dict[str, Any]], dict[str, Any]]:
+def create_task_from_stage(
+    stage: ProcessingStage, ignore_head_node: bool = False
+) -> Callable[[dict[str, Any]], dict[str, Any]]:
     """Create a named Ray Data stage adapter function.
 
     This creates a standalone function that wraps the stage processing logic
@@ -192,7 +244,7 @@ def create_task_from_stage(stage: ProcessingStage) -> Callable[[dict[str, Any]],
         Callable: A function that can be used directly with Ray Data's map_batches
     """
     # Create the adapter instance
-    adapter = RayDataStageAdapter(stage)
+    adapter = RayDataStageAdapter(stage, ignore_head_node=ignore_head_node)
 
     # Create a standalone function that wraps the adapter's processing logic
     def stage_map_fn(batch: dict[str, Any]) -> dict[str, Any]:

nemo_curator/backends/ray_data/executor.py

@@ -93,7 +93,7 @@ def execute(self, stages: list["ProcessingStage"], initial_tasks: list[Task] | N
                 logger.info(f"  CPU cores: {stage.resources.cpus}, GPU ratio: {stage.resources.gpus}")
 
                 # Create adapter for this stage
-                adapter = RayDataStageAdapter(stage)
+                adapter = RayDataStageAdapter(stage, ignore_head_node=self.ignore_head_node)
 
                 # Apply stage transformation
                 current_dataset = adapter.process_dataset(current_dataset)

nemo_curator/backends/utils.py

@@ -16,14 +16,14 @@
 import time
 from copy import deepcopy
 from enum import Enum
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import ray
 from loguru import logger
 
 from nemo_curator.backends.base import NodeInfo, WorkerMetadata
 from nemo_curator.stages.base import ProcessingStage
-from nemo_curator.utils.ray_utils import get_head_node_id, submit_on_each_node
+from nemo_curator.utils.ray_utils import get_alive_ray_nodes, get_head_node_id, submit_on_each_node
 
 if TYPE_CHECKING:
     import loguru
@@ -137,6 +137,26 @@ class RayStageSpecKeys(str, Enum):
     RAY_NUM_CPUS = "ray_num_cpus"
 
 
+def get_stage_num_workers_per_node(stage: ProcessingStage) -> int | float | None:
+    """Return a stage's generic per-node worker request.
+
+    Prefer ``ProcessingStage.num_workers_per_node()``. For older stages that still expose
+    a plain ``num_workers_per_node`` field, accept the field value so backend sizing stays
+    compatible while stages migrate to the method form.
+    """
+    value_or_method: Any = getattr(stage, "num_workers_per_node", None)
+    value = value_or_method() if callable(value_or_method) else value_or_method
+    if value is None:
+        return None
+    if isinstance(value, bool) or not isinstance(value, (int, float)):
+        msg = f"num_workers_per_node() for stage {stage.name} must be a positive number."
+        raise TypeError(msg)
+    if value <= 0:
+        msg = f"num_workers_per_node() for stage {stage.name} must be > 0."
+        raise ValueError(msg)
+    return value
+
+
 def get_worker_metadata_and_node_id() -> tuple[NodeInfo, WorkerMetadata]:
     """Get the worker metadata and node id from the runtime context."""
     ray_context = ray.get_runtime_context()
@@ -212,13 +232,8 @@ def execute_setup_on_node(stages: list[ProcessingStage], ignore_head_node: bool
     the sum of per-stage times — important when setup is heavy (model downloads, weight
     loads) and stages don't contend for the same resources.
     """
-    head_node_id = get_head_node_id() if ignore_head_node else None
-    for node in ray.nodes():
-        if not node.get("Alive"):
-            continue
+    for node in get_alive_ray_nodes(ignore_head_node=ignore_head_node):
         node_id = node["NodeID"]
-        if ignore_head_node and node_id == head_node_id:
-            continue
         logger.info(f"Executing setup on node {node_id} for {len(stages)} stages")
 
     refs: list = []

nemo_curator/backends/xenna/executor.py

@@ -20,7 +20,7 @@
 from loguru import logger
 
 from nemo_curator.backends.base import BaseExecutor
-from nemo_curator.backends.utils import register_loguru_serializer
+from nemo_curator.backends.utils import get_stage_num_workers_per_node, register_loguru_serializer
 from nemo_curator.backends.xenna.adapter import create_named_xenna_stage_adapter
 from nemo_curator.stages.base import ProcessingStage
 from nemo_curator.tasks import EmptyTask, Task
@@ -83,11 +83,21 @@ def execute(self, stages: list[ProcessingStage], initial_tasks: list[Task] | Non
                 raise ValueError(msg)
 
             num_workers = stage.num_workers()
-            num_workers_per_node = stage_config.get("num_workers_per_node")
+            stage_num_workers_per_node = get_stage_num_workers_per_node(stage)
+            spec_num_workers_per_node = stage_config.get("num_workers_per_node")
+            if stage_num_workers_per_node is not None and spec_num_workers_per_node is not None:
+                msg = (
+                    f"Stage {stage.name} sets both num_workers_per_node() and "
+                    "xenna_stage_spec()['num_workers_per_node']. Use only one worker sizing option."
+                )
+                raise ValueError(msg)
+            num_workers_per_node = (
+                stage_num_workers_per_node if stage_num_workers_per_node is not None else spec_num_workers_per_node
+            )
             if num_workers is not None and num_workers_per_node is not None:
                 msg = (
                     f"Stage {stage.name} sets both num_workers() and "
-                    "xenna_stage_spec()['num_workers_per_node']. Use only one worker sizing option."
+                    "num_workers_per_node(). Use only one worker sizing option."
                 )
                 raise ValueError(msg)
 

nemo_curator/stages/base.py

@@ -58,6 +58,36 @@ def get_num_workers() -> int | None:
     return get_num_workers
 
 
+def _num_workers_per_node_method(num_workers_per_node: float | None) -> Callable[[], float | None]:
+    def get_num_workers_per_node() -> float | None:
+        return num_workers_per_node
+
+    return get_num_workers_per_node
+
+
+def _get_num_workers_per_node_value(stage: ProcessingStage) -> float | None:
+    value_or_method = stage.num_workers_per_node
+    return value_or_method() if callable(value_or_method) else value_or_method
+
+
+def _set_num_workers_per_node_override(stage: ProcessingStage, num_workers_per_node: float | None) -> None:
+    if callable(stage.num_workers_per_node):
+        stage.num_workers_per_node = _num_workers_per_node_method(num_workers_per_node)
+    else:
+        stage.num_workers_per_node = num_workers_per_node
+
+
+def _check_worker_sizing_options(stage: ProcessingStage) -> None:
+    effective_num_workers = stage.num_workers()
+    effective_num_workers_per_node = _get_num_workers_per_node_value(stage)
+    if effective_num_workers is not None and effective_num_workers_per_node is not None:
+        msg = (
+            "Use only one worker sizing option: num_workers or num_workers_per_node. "
+            "Set one of them to None before configuring the other."
+        )
+        raise ValueError(msg)
+
+
 class StageMeta(ABCMeta):
     """Metaclass that automatically registers concrete Stage subclasses.
     A class is considered *concrete* if it directly inherits from
@@ -173,6 +203,10 @@ def num_workers(self) -> int | None:
         """Number of workers required. If None, then executor will determine the number of workers."""
         return None
 
+    def num_workers_per_node(self) -> float | None:
+        """Number of workers required per Ray node. If None, executor default sizing is used."""
+        return None
+
     def validate_input(self, task: Task) -> bool:
         """Validate input task meets requirements.
         Args:
@@ -321,6 +355,7 @@ def with_(  # noqa: PLR0913
         ray_stage_spec: dict[str, Any] | None = None,
         xenna_stage_spec: dict[str, Any] | None = None,
         num_workers: int | None | _UnsetType = _UNSET,
+        num_workers_per_node: float | None | _UnsetType = _UNSET,
     ) -> ProcessingStage:
         """Apply configuration changes to this stage with overridden properties.
 
@@ -335,6 +370,8 @@ def with_(  # noqa: PLR0913
             xenna_stage_spec: Merge overrides into the Xenna stage spec. User-provided keys win.
                 Use num_workers instead of setting num_workers in xenna_stage_spec.
             num_workers: Override the num_workers() result. Passing None explicitly resets to executor default behavior.
+            num_workers_per_node: Override the num_workers_per_node() result. Passing None explicitly resets to
+                executor default behavior.
         """
         new_instance = copy.deepcopy(self)
 
@@ -370,6 +407,10 @@ def with_(  # noqa: PLR0913
 
         if num_workers is not _UNSET:
             new_instance.num_workers = _num_workers_method(cast("int | None", num_workers))
+        if num_workers_per_node is not _UNSET:
+            _set_num_workers_per_node_override(new_instance, cast("float | None", num_workers_per_node))
+
+        _check_worker_sizing_options(new_instance)
 
         return new_instance