Add supporting max_calls to the ray data executor

nemo_curator/backends/experimental/ray_data/adapter.py

@@ -96,6 +96,9 @@ def process_dataset(self, dataset: Dataset, ignore_head_node: bool = False) -> D
         else:
             map_batches_fn = create_task_from_stage(self.stage)
             concurrency_kwargs = {"concurrency": None}
+            max_calls = self.stage.ray_stage_spec().get(RayStageSpecKeys.MAX_CALLS_PER_WORKER, None)
+            if max_calls is not None:
+                concurrency_kwargs["max_calls"] = max_calls
 
         if self.stage.resources.cpus > 0:
             concurrency_kwargs["num_cpus"] = self.stage.resources.cpus  # type: ignore[reportArgumentType]

nemo_curator/backends/experimental/utils.py

@@ -60,6 +60,7 @@ class RayStageSpecKeys(str, Enum):
     IS_RAFT_ACTOR = "is_raft_actor"
     IS_LSH_STAGE = "is_lsh_stage"
     IS_SHUFFLE_STAGE = "is_shuffle_stage"
+    MAX_CALLS_PER_WORKER = "max_calls_per_worker"
 
 
 def get_worker_metadata_and_node_id() -> tuple[NodeInfo, WorkerMetadata]:

nemo_curator/stages/text/download/base/iterator.py

@@ -21,6 +21,7 @@
 import pandas as pd
 from loguru import logger
 
+from nemo_curator.backends.experimental.utils import RayStageSpecKeys
 from nemo_curator.stages.base import ProcessingStage
 from nemo_curator.tasks import DocumentBatch, FileGroupTask
 from nemo_curator.utils.column_utils import resolve_filename_column
@@ -59,6 +60,8 @@ class DocumentIterateExtractStage(ProcessingStage[FileGroupTask, DocumentBatch])
     extractor: DocumentExtractor | None = None
     record_limit: int | None = None
     add_filename_column: bool | str = True
+    # Restart worker Process every N tasks to mitigate memory fragmentation
+    max_calls_per_worker: int | None = None  # Only used if executor is Ray Data
 
     def __post_init__(self):
         """Initialize the stage."""
@@ -75,9 +78,22 @@ def inputs(self) -> tuple[list[str], list[str]]:
     def outputs(self) -> tuple[list[str], list[str]]:
         """Define output - produces DocumentBatch with processed records."""
         if self.extractor:
-            return (["data"], self.extractor.output_columns() + ([self.filename_col] if self.add_filename_column else []))
+            return (
+                ["data"],
+                self.extractor.output_columns() + ([self.filename_col] if self.add_filename_column else []),
+            )
         else:
-            return (["data"], self.iterator.output_columns() + ([self.filename_col] if self.add_filename_column else []))
+            return (
+                ["data"],
+                self.iterator.output_columns() + ([self.filename_col] if self.add_filename_column else []),
+            )
+
+    def ray_stage_spec(self) -> dict[str, Any]:
+        """Get Ray configuration for this stage."""
+        spec = {}
+        if self.max_calls_per_worker is not None:
+            spec[RayStageSpecKeys.MAX_CALLS_PER_WORKER] = self.max_calls_per_worker
+        return spec
 
     def process(self, task: FileGroupTask) -> DocumentBatch:
         """Iterate through files and extract structured content.

nemo_curator/stages/text/download/base/stage.py

@@ -41,6 +41,9 @@ class DocumentDownloadExtractStage(CompositeStage[_EmptyTask, DocumentBatch]):
     url_limit: int | None = None
     record_limit: int | None = None
     add_filename_column: bool | str = True
+    # Restart worker Process every N tasks to mitigate memory fragmentation
+    # Only used if executor is Ray Data
+    extractor_max_calls_per_worker: int | None = None
 
     def __post_init__(self):
         """Initialize the constituent stages."""
@@ -61,6 +64,7 @@ def __post_init__(self):
             extractor=self.extractor,
             record_limit=self.record_limit,
             add_filename_column=self.add_filename_column,
+            max_calls_per_worker=self.extractor_max_calls_per_worker,
         )
 
         stages = [url_stage, download_stage, iterate_extract_stage]

nemo_curator/stages/text/download/common_crawl/stage.py

@@ -14,9 +14,12 @@
 
 from typing import Literal
 
+from loguru import logger
+
 from nemo_curator.stages.base import ProcessingStage
 from nemo_curator.stages.text.download import DocumentDownloadExtractStage
 from nemo_curator.stages.text.download.html_extractors import HTMLExtractorAlgorithm
+from nemo_curator.stages.text.download.html_extractors.justext import JusTextExtractor
 
 from .download import CommonCrawlWARCDownloader
 from .extract import CommonCrawlHTMLExtractor
@@ -48,6 +51,7 @@ def __init__(  # noqa: PLR0913
         url_limit: int | None = None,
         record_limit: int | None = None,
         add_filename_column: bool | str = True,
+        extractor_max_calls_per_worker: int | None = None,
     ):
         self.crawl_type = crawl_type
         self.start_snapshot = start_snapshot
@@ -71,6 +75,13 @@ def __init__(  # noqa: PLR0913
             algorithm_kwargs=html_extraction_kwargs,
             stop_lists=stop_lists,
         )
+        if extractor_max_calls_per_worker is None and isinstance(self.extractor.algorithm, JusTextExtractor):
+            extractor_max_calls_per_worker = 2
+            logger.info(
+                "jusText extraction can cause memory fragmentation and lead to OOM errors. "
+                "Setting extractor_max_calls_per_worker=2 for the iterate-extract stage. "
+                "Pass extractor_max_calls_per_worker explicitly to override."
+            )
         super().__init__(
             url_generator=self.url_generator,
             downloader=self.downloader,
@@ -79,6 +90,7 @@ def __init__(  # noqa: PLR0913
             url_limit=url_limit,
             record_limit=record_limit,
             add_filename_column=add_filename_column,
+            extractor_max_calls_per_worker=extractor_max_calls_per_worker,
         )
         self.name = f"common_crawl_{self.crawl_type}_pipeline"
 

tests/backends/experimental/ray_data/test_max_calls_pid.py

@@ -0,0 +1,268 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+import os
+import re
+import shutil
+import subprocess
+import tempfile
+
+import pandas as pd
+import pytest
+import ray
+from loguru import logger
+
+from nemo_curator.backends.experimental.ray_data.executor import RayDataExecutor
+from nemo_curator.backends.experimental.utils import RayStageSpecKeys
+from nemo_curator.stages.base import ProcessingStage, Resources
+from nemo_curator.tasks import DocumentBatch, EmptyTask
+from tests.backends.utils import capture_logs
+from tests.conftest import build_ray_command
+
+
+@pytest.fixture(scope="module")
+def single_cpu_ray_cluster():
+    """Start an isolated 1-CPU Ray cluster for deterministic PID testing.
+
+    Uses a standalone Ray cluster instead of the session-scoped conftest cluster
+    to ensure a single-CPU cluster for testing.
+    Uses tempfile.mkdtemp for a short path to avoid hitting the Unix socket
+    path length limit (108 chars) that pytest's tmp_path_factory can exceed.
+    """
+    original_ray_address = os.environ.pop("RAY_ADDRESS", None)
+
+    temp_dir = tempfile.mkdtemp(prefix="ray1cpu_")
+    cmd, ray_port = build_ray_command(str(temp_dir), num_cpus=1, num_gpus=0, object_store_memory=2 * (1024**3))
+    ray_process = subprocess.Popen(cmd, shell=False)  # noqa: S603
+
+    ray_address = f"localhost:{ray_port}"
+    os.environ["RAY_ADDRESS"] = ray_address
+
+    try:
+        yield ray_address
+    finally:
+        ray_process.kill()
+        ray_process.wait()
+        shutil.rmtree(temp_dir, ignore_errors=True)
+        if original_ray_address is not None:
+            os.environ["RAY_ADDRESS"] = original_ray_address
+        elif "RAY_ADDRESS" in os.environ:
+            del os.environ["RAY_ADDRESS"]
+
+
+@pytest.fixture
+def single_cpu_ray_client(single_cpu_ray_cluster: str) -> None:
+    """Initialize Ray client for tests that need Ray API access."""
+    ray.init(
+        address=single_cpu_ray_cluster,
+        ignore_reinit_error=True,
+        log_to_driver=True,
+        local_mode=False,
+    )
+
+    try:
+        yield
+    finally:
+        logger.info("Shutting down Ray client")
+        ray.shutdown()
+
+
+class PidRecorderStage(ProcessingStage[DocumentBatch, DocumentBatch]):
+    """Test stage that records the worker PID in the output data."""
+
+    name = "pid_recorder"
+
+    def __init__(self, max_calls_per_worker: int | None = None):
+        self._max_calls_per_worker = max_calls_per_worker
+
+    def ray_stage_spec(self) -> dict:
+        spec = {}
+        if self._max_calls_per_worker is not None:
+            spec[RayStageSpecKeys.MAX_CALLS_PER_WORKER] = self._max_calls_per_worker
+        return spec
+
+    def process(self, task: DocumentBatch) -> DocumentBatch:
+        return DocumentBatch(
+            task_id=task.task_id,
+            dataset_name=task.dataset_name,
+            data=pd.DataFrame({"worker_pid": [os.getpid()]}),
+        )
+
+
+class PassthroughActorStage(ProcessingStage[DocumentBatch, DocumentBatch]):
+    """Actor stage (has setup()) that passes data through unchanged.
+
+    Overriding setup() causes is_actor_stage() to return True, so Ray Data
+    will execute this as an actor-based map_batches call.
+    """
+
+    name = "passthrough_actor"
+
+    def ray_stage_spec(self) -> dict:
+        return {
+            RayStageSpecKeys.IS_ACTOR_STAGE: True,
+        }
+
+    def process(self, task: DocumentBatch) -> DocumentBatch:
+        return task
+
+
+class PassthroughTaskStage(ProcessingStage[DocumentBatch, DocumentBatch]):
+    name = "passthrough_task"
+
+    def process(self, task: DocumentBatch) -> DocumentBatch:
+        return task
+
+
+@pytest.mark.parametrize("max_calls_per_worker", [2, None])
+@pytest.mark.usefixtures("single_cpu_ray_client")
+def test_pid_recycling(max_calls_per_worker: int | None):
+    tasks = [EmptyTask] * 8
+
+    stage = PidRecorderStage(max_calls_per_worker=max_calls_per_worker)
+    executor = RayDataExecutor()
+    results = executor.execute(stages=[stage], initial_tasks=tasks)
+
+    pids = [r.data["worker_pid"].iloc[0] for r in results]
+    expected_unique_pids = math.ceil(len(tasks) / max_calls_per_worker) if max_calls_per_worker else 1
+
+    assert len(set(pids)) == expected_unique_pids, (
+        f"Expected {expected_unique_pids} unique PIDs with max_calls={max_calls_per_worker}, "
+        f"got {len(set(pids))}: {pids}"
+    )
+
+
+@pytest.mark.parametrize("max_calls_per_worker", [2, None])
+@pytest.mark.usefixtures("single_cpu_ray_client")
+def test_max_calls_not_fused_with_actor_stage(max_calls_per_worker: int | None):
+    """Verify that a task stage with max_calls is not fused with a following actor stage.
+
+    Ray Data's operator fusion optimization can merge consecutive map_batches
+    operations. If PidRecorderStage (task-based, max_calls=1) were fused with
+    PassthroughActorStage (actor-based), the max_calls setting would be lost
+    and we'd see only 1 unique PID instead of one per task.
+
+    By chaining the two stages and asserting PID recycling still occurs, we
+    confirm that Ray Data keeps them as separate operators. We also verify
+    the execution plan directly to ensure no fusion occurred.
+    """
+    num_tasks = 4
+    tasks = [EmptyTask] * num_tasks
+
+    pid_stage = PidRecorderStage(max_calls_per_worker=max_calls_per_worker).with_(resources=Resources(cpus=0.5))
+    actor_stage = PassthroughActorStage().with_(resources=Resources(cpus=0.5))
+
+    executor = RayDataExecutor()
+    with capture_logs() as log_buffer:
+        results = executor.execute(stages=[pid_stage, actor_stage], initial_tasks=tasks)
+        all_logs = log_buffer.getvalue()
+
+    # Verify PID recycling
+    pids = [r.data["worker_pid"].iloc[0] for r in results]
+    expected_unique_pids = math.ceil(num_tasks / max_calls_per_worker) if max_calls_per_worker else 1
+
+    assert len(set(pids)) == expected_unique_pids, (
+        f"Expected {expected_unique_pids} unique PIDs (max_calls={max_calls_per_worker}, "
+        f"{num_tasks} tasks), got {len(set(pids))}: {pids}. "
+        f"Stages may have been fused by Ray Data, defeating max_calls."
+    )
+
+    # Verify execution plan fusion behavior
+    matches = re.findall(r"Execution plan of Dataset.*?:\s*(.+)", all_logs, re.MULTILINE)
+    assert matches, f"No execution plan found in Ray Data logs. Full logs:\n{all_logs}"
+    plan_stages = [s.strip() for s in matches[-1].split(" -> ")]
+    map_batches_stages = [s for s in plan_stages if "MapBatches" in s]
+
+    if max_calls_per_worker is not None:
+        # When max_calls is set, stages must NOT be fused — each should be a separate operator.
+        assert len(map_batches_stages) == 2, (
+            f"Expected 2 separate MapBatches operators, got {len(map_batches_stages)}: {map_batches_stages}. "
+            f"Full execution plan: {matches[-1]}"
+        )
+        for stage in map_batches_stages:
+            assert stage.count("MapBatches") == 1, (
+                f"Stages were fused into a single operator: {stage}. Full execution plan: {matches[-1]}"
+            )
+    else:
+        # When max_calls is None, Ray Data is free to fuse — confirm fusion happened.
+        assert len(map_batches_stages) == 1, (
+            f"Expected 1 fused MapBatches operator, got {len(map_batches_stages)}: {map_batches_stages}. "
+            f"Full execution plan: {matches[-1]}"
+        )
+        assert map_batches_stages[0].count("MapBatches") == 2, (
+            f"Expected fused operator with 2 MapBatches, got: {map_batches_stages[0]}. "
+            f"Full execution plan: {matches[-1]}"
+        )
+
+
+@pytest.mark.parametrize("max_calls_per_worker", [2, None])
+@pytest.mark.usefixtures("single_cpu_ray_client")
+def test_max_calls_not_fused_with_task_stage(max_calls_per_worker: int | None):
+    """Verify that a task stage with max_calls is not fused with a following task stage.
+
+    Ray Data's operator fusion optimization can merge consecutive map_batches
+    operations. If PidRecorderStage (task-based, max_calls=1) were fused with
+    PassthroughTaskStage (task-based), the max_calls setting would be lost
+    and we'd see only 1 unique PID instead of one per task.
+
+    We also verify the execution plan directly to ensure no fusion occurred.
+    """
+    num_tasks = 4
+    tasks = [EmptyTask] * num_tasks
+
+    pid_stage = PidRecorderStage(max_calls_per_worker=max_calls_per_worker).with_(resources=Resources(cpus=0.5))
+    task_stage = PassthroughTaskStage().with_(resources=Resources(cpus=0.5))
+
+    executor = RayDataExecutor()
+    with capture_logs() as log_buffer:
+        results = executor.execute(stages=[task_stage, pid_stage], initial_tasks=tasks)
+        all_logs = log_buffer.getvalue()
+
+    # Verify PID recycling
+    pids = [r.data["worker_pid"].iloc[0] for r in results]
+    expected_unique_pids = math.ceil(num_tasks / max_calls_per_worker) if max_calls_per_worker else 1
+
+    assert len(set(pids)) == expected_unique_pids, (
+        f"Expected {expected_unique_pids} unique PIDs (max_calls={max_calls_per_worker}, "
+        f"{num_tasks} tasks), got {len(set(pids))}: {pids}. "
+        f"Stages may have been fused by Ray Data, defeating max_calls."
+    )
+
+    # Verify execution plan fusion behavior
+    matches = re.findall(r"Execution plan of Dataset.*?:\s*(.+)", all_logs, re.MULTILINE)
+    assert matches, f"No execution plan found in Ray Data logs. Full logs:\n{all_logs}"
+    plan_stages = [s.strip() for s in matches[-1].split(" -> ")]
+    map_batches_stages = [s for s in plan_stages if "MapBatches" in s]
+
+    if max_calls_per_worker is not None:
+        # When max_calls is set, stages must NOT be fused — each should be a separate operator.
+        assert len(map_batches_stages) == 2, (
+            f"Expected 2 separate MapBatches operators, got {len(map_batches_stages)}: {map_batches_stages}. "
+            f"Full execution plan: {matches[-1]}"
+        )
+        for stage in map_batches_stages:
+            assert stage.count("MapBatches") == 1, (
+                f"Stages were fused into a single operator: {stage}. Full execution plan: {matches[-1]}"
+            )
+    else:
+        # When max_calls is None, Ray Data is free to fuse — confirm fusion happened.
+        assert len(map_batches_stages) == 1, (
+            f"Expected 1 fused MapBatches operator, got {len(map_batches_stages)}: {map_batches_stages}. "
+            f"Full execution plan: {matches[-1]}"
+        )
+        assert map_batches_stages[0].count("MapBatches") == 2, (
+            f"Expected fused operator with 2 MapBatches, got: {map_batches_stages[0]}. "
+            f"Full execution plan: {matches[-1]}"
+        )