diff --git a/nemo_curator/backends/experimental/ray_data/adapter.py b/nemo_curator/backends/experimental/ray_data/adapter.py index fef0dd256c..55c250e1df 100644 --- a/nemo_curator/backends/experimental/ray_data/adapter.py +++ b/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] diff --git a/nemo_curator/backends/experimental/utils.py b/nemo_curator/backends/experimental/utils.py index 9dd64d6066..d9e753bd90 100644 --- a/nemo_curator/backends/experimental/utils.py +++ b/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]: diff --git a/nemo_curator/stages/text/download/base/iterator.py b/nemo_curator/stages/text/download/base/iterator.py index 546d783a6f..6cdd8d513e 100644 --- a/nemo_curator/stages/text/download/base/iterator.py +++ b/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. diff --git a/nemo_curator/stages/text/download/base/stage.py b/nemo_curator/stages/text/download/base/stage.py index c55b110176..3d2e5a63a8 100644 --- a/nemo_curator/stages/text/download/base/stage.py +++ b/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] diff --git a/nemo_curator/stages/text/download/common_crawl/stage.py b/nemo_curator/stages/text/download/common_crawl/stage.py index b6df5a76a4..d8463ce67f 100644 --- a/nemo_curator/stages/text/download/common_crawl/stage.py +++ b/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" diff --git a/tests/backends/experimental/ray_data/test_max_calls_pid.py b/tests/backends/experimental/ray_data/test_max_calls_pid.py new file mode 100644 index 0000000000..e060eeea77 --- /dev/null +++ b/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]}" + ) diff --git a/tests/backends/experimental/ray_data/test_utils.py b/tests/backends/experimental/ray_data/test_utils.py index 9b3caf39c5..8958457b35 100644 --- a/tests/backends/experimental/ray_data/test_utils.py +++ b/tests/backends/experimental/ray_data/test_utils.py @@ -21,6 +21,7 @@ get_available_cpu_gpu_resources, ) from nemo_curator.stages.resources import Resources +from tests.backends.experimental.test_utils import reset_head_node_cache # noqa: F401 class TestGetAvailableCpuGpuResources: @@ -36,6 +37,7 @@ def test_get_available_cpu_gpu_resources_conftest(self, shared_ray_client: None) # Can be 0 (CPU-only) or 2 (GPU-enabled) depending on test selection assert gpus in [0.0, 2.0] + @pytest.mark.usefixtures("reset_head_node_cache") def test_get_resources_with_ignore_head_node( self, shared_ray_client: None, diff --git a/tests/conftest.py b/tests/conftest.py index 2491c0b2d8..d7450b78d1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# 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. @@ -167,7 +167,7 @@ def pytest_ignore_collect(collection_path: Path, config: pytest.Config) -> bool: return False -def _build_ray_command(temp_dir: str, num_cpus: int, num_gpus: int, object_store_memory: int) -> tuple[list[str], int]: +def build_ray_command(temp_dir: str, num_cpus: int, num_gpus: int, object_store_memory: int) -> tuple[list[str], int]: """Build the Ray start command with the given configuration.""" ray_port = find_free_port() dashboard_port = find_free_port() @@ -236,7 +236,7 @@ def shared_ray_cluster(tmp_path_factory: pytest.TempPathFactory, pytestconfig: p temp_dir = tmp_path_factory.mktemp("ray") # Build and execute Ray command - cmd_to_run, ray_port = _build_ray_command(str(temp_dir), num_cpus, num_gpus, object_store_memory) + cmd_to_run, ray_port = build_ray_command(str(temp_dir), num_cpus, num_gpus, object_store_memory) logger.info(f"Starting Ray cluster with {num_gpus} GPUs") logger.info(f"Running Ray command: {' '.join(cmd_to_run)}") diff --git a/tests/stages/text/download/base/test_iterator.py b/tests/stages/text/download/base/test_iterator.py index 4143115ea6..ab9023ab4b 100644 --- a/tests/stages/text/download/base/test_iterator.py +++ b/tests/stages/text/download/base/test_iterator.py @@ -19,6 +19,7 @@ import pytest +from nemo_curator.backends.experimental.utils import RayStageSpecKeys from nemo_curator.stages.resources import Resources from nemo_curator.stages.text.download.base.iterator import DocumentIterateExtractStage, DocumentIterator from nemo_curator.tasks import DocumentBatch, FileGroupTask @@ -520,3 +521,19 @@ def test_process_all_files_fail(self, tmp_path: Path, caplog: pytest.LogCaptureF # Check that error was logged assert "Error iterating" in caplog.text + + def test_ray_stage_spec_default_no_max_calls(self) -> None: + """Test that ray_stage_spec returns empty dict when max_calls_per_worker is not set.""" + iterator = MockDocumentIterator() + stage = DocumentIterateExtractStage(iterator=iterator) + + assert stage.ray_stage_spec() == {} + + @pytest.mark.parametrize("max_calls", [1, 5]) + def test_ray_stage_spec_with_max_calls(self, max_calls: int) -> None: + """Test that ray_stage_spec returns max_calls_per_worker when set.""" + iterator = MockDocumentIterator() + stage = DocumentIterateExtractStage(iterator=iterator, max_calls_per_worker=max_calls) + + spec = stage.ray_stage_spec() + assert spec[RayStageSpecKeys.MAX_CALLS_PER_WORKER] == max_calls diff --git a/tests/stages/text/download/base/test_stage.py b/tests/stages/text/download/base/test_stage.py index 943a5529ad..c71917c904 100644 --- a/tests/stages/text/download/base/test_stage.py +++ b/tests/stages/text/download/base/test_stage.py @@ -371,4 +371,5 @@ def test_stage_initialization_mocking( extractor=extractor, record_limit=10, add_filename_column="test_file", + max_calls_per_worker=None, )