Make num_workers_per_node method-only and centralize worker-sizing validation

- num_workers_per_node is now strictly a method (like num_workers); drop the
  field/method duality and hard-code via with_(num_workers_per_node=...).
- Enforce at-most-one worker-sizing option (num_workers / num_workers_per_node /
  actor-pool min/max/initial_workers) once in ProcessingStage._validate_worker_sizing,
  run at construction (StageMeta.__call__) and in with_(). Remove the duplicated
  checks from the ray_actor_pool, ray_data, and xenna backends.
- Simplify get_stage_num_workers_per_node to call num_workers_per_node() directly.
- Simplify RayDataStageAdapter compute-strategy construction.
- Migrate ImageDuplicatesRemovalStage and CallHomeReaderStage off the legacy
  xenna_stage_spec num_workers_per_node key; drop legacy spec-key support.

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

Author: praateekmahajan

nemo_curator/backends/ray_actor_pool/utils.py

@@ -61,12 +61,6 @@ def calculate_optimal_actors_for_stage(
 
     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:

nemo_curator/backends/ray_data/adapter.py

@@ -97,49 +97,41 @@ 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:
+    def _per_node_pool_size(self) -> int | None:
+        """Actor/task pool size derived from ``num_workers_per_node()``, or None if unset.
+
+        Mutual exclusion with ``num_workers()`` and actor-pool sizing keys is enforced at
+        stage construction (``ProcessingStage._validate_worker_sizing``).
+        """
         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:
+    def _compute_kwargs(self, ray_stage_spec: dict, stage_is_actor: bool, per_node_pool_size: int | None) -> dict:
+        """Select the Ray Data compute strategy (actor pool vs task pool) and its sizing."""
+        if stage_is_actor:
+            compute = (
+                ActorPoolStrategy(size=per_node_pool_size)
+                if per_node_pool_size is not None
+                else get_actor_compute_strategy_for_stage(self.stage)
+            )
+            return {"compute": compute}
+
+        # Task stages: actor-pool sizing keys don't apply, so surface them if present.
+        sizing_keys = get_configured_actor_pool_sizing_keys(ray_stage_spec)
+        if sizing_keys:
             logger.warning(
-                f"Ignoring ray_stage_spec worker sizing keys {actor_pool_sizing_keys} "
+                f"Ignoring ray_stage_spec worker sizing keys {sizing_keys} "
                 f"for Ray Data task stage {self.stage.name}; these keys only apply to actor stages."
             )
 
         map_batches_kwargs: dict[str, object] = {}
+        num_workers = self.stage.num_workers()
         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:
@@ -161,15 +153,11 @@ 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)
+        per_node_pool_size = self._per_node_pool_size()
 
-        if stage_is_actor:
-            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, ignore_head_node=self.ignore_head_node)
-            map_batches_kwargs = self._build_task_compute_kwargs(ray_stage_spec, per_node_pool_size, num_workers)
+        stage_factory = create_actor_from_stage if stage_is_actor else create_task_from_stage
+        map_batches_fn = stage_factory(self.stage, ignore_head_node=self.ignore_head_node)
+        map_batches_kwargs = self._compute_kwargs(ray_stage_spec, stage_is_actor, per_node_pool_size)
 
         map_batches_kwargs.update(self._build_resource_kwargs(ray_stage_spec))
 

nemo_curator/backends/ray_data/utils.py

@@ -14,7 +14,6 @@
 
 from collections.abc import Mapping
 
-from loguru import logger
 from ray.data import ActorPoolStrategy
 
 from nemo_curator.backends.utils import RayStageSpecKeys
@@ -42,12 +41,6 @@ def get_actor_compute_strategy_for_stage(stage: ProcessingStage) -> ActorPoolStr
     """
     num_workers = stage.num_workers()
     if num_workers is not None and num_workers > 0:
-        actor_pool_sizing_keys = get_configured_actor_pool_sizing_keys(stage.ray_stage_spec())
-        if actor_pool_sizing_keys:
-            logger.warning(
-                f"Stage {stage.name} uses num_workers={num_workers}; ignoring ray_stage_spec "
-                f"actor-pool sizing keys {actor_pool_sizing_keys}."
-            )
         return ActorPoolStrategy(size=num_workers)
 
     ray_stage_spec = stage.ray_stage_spec()

nemo_curator/backends/utils.py

@@ -16,7 +16,7 @@
 import time
 from copy import deepcopy
 from enum import Enum
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 import ray
 from loguru import logger
@@ -138,14 +138,8 @@ class RayStageSpecKeys(str, Enum):
 
 
 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
+    """Return a stage's validated per-node worker request from ``num_workers_per_node()``."""
+    value = stage.num_workers_per_node()
     if value is None:
         return None
     if isinstance(value, bool) or not isinstance(value, (int, float)):

nemo_curator/backends/xenna/executor.py

@@ -78,28 +78,15 @@ def execute(self, stages: list[ProcessingStage], initial_tasks: list[Task] | Non
         for stage in stages:
             # Get stage configuration
             stage_config = stage.xenna_stage_spec()
-            if "num_workers" in stage_config:
-                msg = f"Stage {stage.name} sets num_workers in xenna_stage_spec(). Use num_workers() instead."
-                raise ValueError(msg)
+            for reserved in ("num_workers", "num_workers_per_node"):
+                if reserved in stage_config:
+                    msg = f"Stage {stage.name} sets {reserved} in xenna_stage_spec(). Use {reserved}() instead."
+                    raise ValueError(msg)
 
+            # Mutual exclusion of num_workers() vs num_workers_per_node() is enforced at
+            # stage construction (ProcessingStage._validate_worker_sizing).
             num_workers = stage.num_workers()
-            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 "
-                    "num_workers_per_node(). Use only one worker sizing option."
-                )
-                raise ValueError(msg)
+            num_workers_per_node = get_stage_num_workers_per_node(stage)
 
             # Create Xenna stage adapter with the original stage's name
             xenna_stage = create_named_xenna_stage_adapter(

nemo_curator/stages/base.py

@@ -65,29 +65,6 @@ def get_num_workers_per_node() -> float | None:
     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
@@ -111,6 +88,12 @@ def __new__(mcls, name, bases, namespace, **kwargs):  # noqa: ANN001
 
         return cls
 
+    def __call__(cls, *args, **kwargs):
+        """Validate worker sizing once every stage instance is fully constructed."""
+        instance = super().__call__(*args, **kwargs)
+        instance._validate_worker_sizing()
+        return instance
+
 
 def get_stage_class(name: str) -> type[ProcessingStage]:
     """Retrieve a registered stage class by its *class name*.
@@ -179,16 +162,17 @@ def __init_subclass__(cls, **kwargs):
                 msg = f"{cls.__name__} must not override '{attr}'"
                 raise TypeError(msg)
 
-        num_workers = cls.__dict__.get("num_workers")
-        if (num_workers is not None and not callable(num_workers)) or (
-            num_workers is None and "num_workers" in cls.__dict__.get("__annotations__", {})
-        ):
-            msg = (
-                f"{cls.__name__} must not define 'num_workers' as a stage attribute. "
-                "Override num_workers() for backend worker sizing, or use a different field name "
-                "for stage-specific worker counts."
-            )
-            raise TypeError(msg)
+        for worker_attr in ("num_workers", "num_workers_per_node"):
+            value = cls.__dict__.get(worker_attr)
+            if (value is not None and not callable(value)) or (
+                value is None and worker_attr in cls.__dict__.get("__annotations__", {})
+            ):
+                msg = (
+                    f"{cls.__name__} must not define '{worker_attr}' as a stage attribute. "
+                    f"Override {worker_attr}() for backend worker sizing, or use a different field name "
+                    "for stage-specific worker counts."
+                )
+                raise TypeError(msg)
 
         for attr in ("name", "resources", "batch_size", "runtime_env"):
             if isinstance(cls.__dict__.get(attr), property):
@@ -207,6 +191,38 @@ 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_worker_sizing(self) -> None:
+        """Enforce that at most one worker-sizing mechanism is configured.
+
+        A stage may size its worker pool with exactly one of ``num_workers()``,
+        ``num_workers_per_node()``, or the Ray Data actor-pool sizing keys
+        (``min_workers``/``max_workers``/``initial_workers``) in ``ray_stage_spec()``.
+        Combining them is ambiguous, so we reject it once here — at construction and in
+        ``with_()`` — rather than defensively re-checking in every backend.
+        """
+        configured = [
+            name
+            for name, value in (
+                ("num_workers()", self.num_workers()),
+                ("num_workers_per_node()", self.num_workers_per_node()),
+            )
+            if value is not None
+        ]
+        ray_stage_spec = self.ray_stage_spec()
+        if ray_stage_spec:
+            # Lazy import: only stages that populate ray_stage_spec touch the Ray Data backend.
+            from nemo_curator.backends.ray_data.utils import get_configured_actor_pool_sizing_keys
+
+            if get_configured_actor_pool_sizing_keys(ray_stage_spec):
+                configured.append("actor-pool sizing keys (min/max/initial_workers)")
+
+        if len(configured) > 1:
+            msg = (
+                f"Stage {self.name} sets multiple worker-sizing options ({', '.join(configured)}); "
+                "use only one of num_workers, num_workers_per_node, or actor-pool sizing keys."
+            )
+            raise ValueError(msg)
+
     def validate_input(self, task: Task) -> bool:
         """Validate input task meets requirements.
         Args:
@@ -408,9 +424,11 @@ 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))
+            new_instance.num_workers_per_node = _num_workers_per_node_method(
+                cast("float | None", num_workers_per_node)
+            )
 
-        _check_worker_sizing_options(new_instance)
+        new_instance._validate_worker_sizing()
 
         return new_instance
 

nemo_curator/stages/image/deduplication/removal.py

@@ -14,7 +14,6 @@
 
 import os
 from dataclasses import dataclass, field
-from typing import Any
 
 import pyarrow as pa
 import pyarrow.dataset as ds
@@ -36,15 +35,14 @@ class ImageDuplicatesRemovalStage(ProcessingStage[ImageBatch, ImageBatch]):
         removal_parquets_dir: Directory containing Parquet files with image IDs to remove
         duplicate_id_field: Name of the column containing image IDs to remove
         verbose: Whether to log verbose output
-        num_workers_per_node: Number of workers per node for the stage. This is sometimes needed
-            to avoid OOM when concurrently running actors on one node loading the same removal
-            parquet files into memory.
+
+    To cap concurrent actors per node (e.g. to avoid OOM when each actor loads the same
+    removal parquet files into memory), configure ``.with_(num_workers_per_node=...)``.
     """
 
     removal_parquets_dir: str
     duplicate_id_field: str = "id"
     verbose: bool = False
-    num_workers_per_node: int | None = None
 
     name: str = "image_dedup_filter"
 
@@ -96,9 +94,3 @@ def process(self, task: ImageBatch) -> ImageBatch:
             _metadata=task._metadata,
             _stage_perf=task._stage_perf,
         )
-
-    def xenna_stage_spec(self) -> dict[str, Any]:
-        spec: dict[str, Any] = {}
-        if self.num_workers_per_node is not None:
-            spec["num_workers_per_node"] = self.num_workers_per_node
-        return spec

tests/backends/ray_actor_pool/test_executor.py

@@ -95,20 +95,9 @@ def test_calculate_optimal_actors_respects_num_workers_per_node(
             mock_warning.assert_called_once()
             assert expected_warning in mock_warning.call_args.args[0]
 
-    @pytest.mark.parametrize("num_workers", [0, 4])
-    def test_calculate_optimal_actors_rejects_num_workers_with_num_workers_per_node(self, num_workers: int) -> None:
-        stage = _stage_with_worker_sizing(num_workers=num_workers, num_workers_per_node=2, cpus=1.0, batch_size=10)
-
-        with (
-            mock.patch(
-                "nemo_curator.backends.ray_actor_pool.utils.get_available_cpu_gpu_resources",
-                return_value=(8.0, 0.0),
-            ),
-            pytest.raises(ValueError, match=r"num_workers\(\).*num_workers_per_node\(\)"),
-        ):
-            calculate_optimal_actors_for_stage(stage, num_tasks=1)
-
 
+# Mutual exclusion of num_workers() and num_workers_per_node() is enforced at stage
+# construction (ProcessingStage._validate_worker_sizing); see tests/stages/common/test_base.py.
 def _stage_with_worker_sizing(
     *, num_workers: int | None, num_workers_per_node: float | None, cpus: float, batch_size: int
 ) -> mock.Mock: