diff --git a/.github/workflows/cicd-main.yml b/.github/workflows/cicd-main.yml index e03e58d44e..75e3ec93c2 100644 --- a/.github/workflows/cicd-main.yml +++ b/.github/workflows/cicd-main.yml @@ -95,6 +95,7 @@ jobs: - name: Run tests ${{ matrix.folder }} (CPU) timeout-minutes: 40 run: | + uv venv --seed uv sync --link-mode copy --locked --extra audio_cpu --extra sdg_cpu --extra text_cpu --extra video_cpu --group test source .venv/bin/activate FOLDER="${{ matrix.folder }}" diff --git a/docker/Dockerfile b/docker/Dockerfile index a00bcb3268..1dddc78005 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -56,7 +56,7 @@ ENV UV_PROJECT_ENVIRONMENT=/opt/venv ENV UV_CACHE_DIR=/opt/uv_cache ENV PATH="$UV_PROJECT_ENVIRONMENT/bin:$PATH" ENV UV_LINK_MODE=copy -RUN uv venv ${UV_PROJECT_ENVIRONMENT} --system-site-packages +RUN uv venv ${UV_PROJECT_ENVIRONMENT} --system-site-packages --seed FROM build AS nemo_curator_dep diff --git a/nemo_curator/backends/experimental/ray_actor_pool/executor.py b/nemo_curator/backends/experimental/ray_actor_pool/executor.py index ec584f648e..f2e2273db7 100644 --- a/nemo_curator/backends/experimental/ray_actor_pool/executor.py +++ b/nemo_curator/backends/experimental/ray_actor_pool/executor.py @@ -168,14 +168,16 @@ def execute(self, stages: list["ProcessingStage"], initial_tasks: list[Task] | N def _create_actor_pool(self, stage: "ProcessingStage", num_actors: int) -> ActorPool: """Create an ActorPool for a specific stage.""" actors = [] + actor_options: dict = { + "num_cpus": stage.resources.cpus, + "num_gpus": stage.resources.gpus, + } + if stage.runtime_env: + actor_options["runtime_env"] = stage.runtime_env for i in range(num_actors): actor = ( create_named_ray_actor_pool_stage_adapter(stage, RayActorPoolStageAdapter) - .options( - num_cpus=stage.resources.cpus, - num_gpus=stage.resources.gpus, - name=f"{stage.name}-{i}", - ) + .options(**actor_options, name=f"{stage.name}-{i}") .remote(stage) ) actors.append(actor) diff --git a/nemo_curator/backends/experimental/utils.py b/nemo_curator/backends/experimental/utils.py index 2b1db368ae..6b81115ef5 100644 --- a/nemo_curator/backends/experimental/utils.py +++ b/nemo_curator/backends/experimental/utils.py @@ -62,6 +62,7 @@ class RayStageSpecKeys(str, Enum): IS_LSH_STAGE = "is_lsh_stage" IS_SHUFFLE_STAGE = "is_shuffle_stage" MAX_CALLS_PER_WORKER = "max_calls_per_worker" + RAY_REMOTE_ARGS = "ray_remote_args" def get_worker_metadata_and_node_id() -> tuple[NodeInfo, WorkerMetadata]: diff --git a/nemo_curator/backends/ray_data/adapter.py b/nemo_curator/backends/ray_data/adapter.py index d4f74a12b9..e3abe2847d 100644 --- a/nemo_curator/backends/ray_data/adapter.py +++ b/nemo_curator/backends/ray_data/adapter.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy from collections.abc import Callable from typing import Any @@ -101,6 +102,15 @@ def process_dataset(self, dataset: Dataset, ignore_head_node: bool = False) -> D if self.stage.resources.gpus > 0: concurrency_kwargs["num_gpus"] = self.stage.resources.gpus # type: ignore[reportArgumentType] + # Per-stage ray_remote_args (e.g. runtime_env with different pip versions per stage). + ray_remote_args = copy.deepcopy(self.stage.ray_stage_spec().get(RayStageSpecKeys.RAY_REMOTE_ARGS) or {}) + # If the stage declares runtime_env, forward it directly to Ray so Ray creates and + # caches an isolated virtualenv for this stage's workers. + if self.stage.runtime_env: + ray_remote_args["runtime_env"] = self.stage.runtime_env + + concurrency_kwargs.update(ray_remote_args) + # Calculate concurrency based on available resources logger.info(f"{self.stage.__class__.__name__} {is_actor_stage_=} with {concurrency_kwargs=}") diff --git a/nemo_curator/backends/ray_data/executor.py b/nemo_curator/backends/ray_data/executor.py index 63670d20cc..dbf9870f9e 100644 --- a/nemo_curator/backends/ray_data/executor.py +++ b/nemo_curator/backends/ray_data/executor.py @@ -61,6 +61,9 @@ def execute(self, stages: list["ProcessingStage"], initial_tasks: list[Task] | N # Initialize with initial tasks if provided, otherwise start with EmptyTask tasks: list[Task] = initial_tasks if initial_tasks else [EmptyTask] output_tasks: list[Task] = [] + # When runtime_env with pip is used, Ray's pip plugin sets up per-stage virtualenvs + # lazily on first task dispatch by cloning the current virtualenv. The NeMo Curator + # container's /opt/venv is created with `uv venv --seed` so pip is available in clones. try: # Initialize ray and explicitly set NOSET to empty # This ensures if Xenna was used before which was setting NOSET, we end up overriding it. diff --git a/nemo_curator/backends/xenna/adapter.py b/nemo_curator/backends/xenna/adapter.py index 84b36b4cdc..0d4e17d3d3 100644 --- a/nemo_curator/backends/xenna/adapter.py +++ b/nemo_curator/backends/xenna/adapter.py @@ -12,6 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +from typing import Any + +import ray.runtime_env from cosmos_xenna.pipelines import v1 as pipelines_v1 from cosmos_xenna.pipelines.private.resources import NodeInfo as XennaNodeInfo from cosmos_xenna.pipelines.private.resources import Resources as XennaResources @@ -23,6 +26,31 @@ from nemo_curator.tasks import Task +class CuratorRuntimeEnv: + """Duck-typed replacement for Xenna's RuntimeEnv that supports the full Ray runtime_env dict. + + Xenna's RuntimeEnv only supports conda + env_vars. This class accepts a raw + Ray-format runtime_env dict and implements the interface Xenna calls: + ``to_ray_runtime_env()``, ``format()``, and ``extra_env_vars``. + """ + + def __init__(self, runtime_env: dict[str, Any]) -> None: + self._runtime_env = runtime_env + # Xenna's actor pool both reads and writes extra_env_vars on the runtime env object, + # so this must be a plain settable attribute, not a read-only property. + self.extra_env_vars: dict[str, str] = dict(runtime_env.get("env_vars", {})) + + def to_ray_runtime_env(self) -> ray.runtime_env.RuntimeEnv: + # Merge self.extra_env_vars (which Xenna may have mutated with per-actor env vars + # such as CUDA_VISIBLE_DEVICES) back into the runtime_env dict before handing it + # to Ray, so those injected vars are not silently dropped. + merged = {**self._runtime_env, "env_vars": self.extra_env_vars} + return ray.runtime_env.RuntimeEnv(**merged) + + def format(self) -> str: + return f"runtime_env_keys: {', '.join(self._runtime_env.keys())}" + + class XennaStageAdapter(BaseStageAdapter, pipelines_v1.Stage): """Adapts ProcessingStage to Xenna. Args: @@ -54,9 +82,14 @@ def stage_batch_size(self) -> int: @property def env_info(self) -> pipelines_v1.RuntimeEnv | None: - """Runtime environment for this stage.""" - # Can be customized per stage if needed - return None + """Runtime environment for this stage. + + Converts the ProcessingStage.runtime_env dict (Ray-format) to a + CuratorRuntimeEnv that Xenna can forward to Ray actors. + """ + if not self.processing_stage.runtime_env: + return None + return CuratorRuntimeEnv(self.processing_stage.runtime_env) def process_data(self, tasks: list[Task]) -> list[Task] | None: """Process batch of tasks with automatic performance tracking. diff --git a/nemo_curator/stages/base.py b/nemo_curator/stages/base.py index 19e0b4f193..5761dfeb18 100644 --- a/nemo_curator/stages/base.py +++ b/nemo_curator/stages/base.py @@ -19,7 +19,7 @@ import time from abc import ABC, ABCMeta, abstractmethod from inspect import isabstract -from typing import TYPE_CHECKING, Any, Generic, TypeVar, final +from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, final from loguru import logger @@ -85,6 +85,7 @@ class ProcessingStage(ABC, Generic[X, Y], metaclass=StageMeta): name = "ProcessingStage" resources = Resources(cpus=1.0) batch_size = 1 + runtime_env: ClassVar[dict[str, Any] | None] = None @property @final @@ -114,7 +115,7 @@ def __init_subclass__(cls, **kwargs): msg = f"{cls.__name__} must not override '_batch_size'" raise TypeError(msg) - for attr in ("name", "resources", "batch_size"): + for attr in ("name", "resources", "batch_size", "runtime_env"): if isinstance(cls.__dict__.get(attr), property): msg = ( f"{cls.__name__} must not define '{attr}' as a @property. " @@ -259,7 +260,11 @@ def xenna_stage_spec(self) -> dict[str, Any]: return {} def with_( - self, name: str | None = None, resources: Resources | None = None, batch_size: int | None = None + self, + name: str | None = None, + resources: Resources | None = None, + batch_size: int | None = None, + runtime_env: dict[str, Any] | None = None, ) -> ProcessingStage: """Apply configuration changes to this stage with overridden properties. @@ -269,6 +274,7 @@ def with_( name: Override the name property resources: Override the resources property batch_size: Override the batch_size property + runtime_env: Override the runtime_env (Ray runtime environment dict) """ new_instance = copy.deepcopy(self) @@ -279,6 +285,8 @@ def with_( new_instance.resources = resources if batch_size is not None: new_instance.batch_size = batch_size + if runtime_env is not None: + new_instance.runtime_env = runtime_env return new_instance diff --git a/tests/pipelines/test_per_stage_runtime_env.py b/tests/pipelines/test_per_stage_runtime_env.py new file mode 100644 index 0000000000..8bdac05a18 --- /dev/null +++ b/tests/pipelines/test_per_stage_runtime_env.py @@ -0,0 +1,151 @@ +# Copyright (c) 2025, 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. + +"""Tests for per-stage runtime_env support. + +Verifies that stages can declare different runtime_env (pip/uv packages) and +that Ray's native runtime_env creates isolated venvs per actor on each node. +Also verifies that runtime_env is additive: base-env packages remain importable. +""" + +from typing import Any + +import pandas as pd +import pytest + +from nemo_curator.backends.base import BaseExecutor +from nemo_curator.backends.ray_data import RayDataExecutor +from nemo_curator.backends.xenna import XennaExecutor +from nemo_curator.pipeline.pipeline import Pipeline +from nemo_curator.stages.base import ProcessingStage +from nemo_curator.stages.resources import Resources +from nemo_curator.tasks import DocumentBatch + + +class RecordPackagingVersionStage(ProcessingStage[DocumentBatch, DocumentBatch]): + """Records the packaging library version visible to this worker. + + The column name is derived from self.name so multiple instances with + different runtime_env can coexist in the same pipeline. + Also checks whether loguru (a Curator dep, not a Ray dep) is importable. + """ + + name = "record_packaging_version" + resources = Resources(cpus=0.5) + batch_size = 1 + + def inputs(self) -> tuple[list[str], list[str]]: + return ["data"], [] + + def outputs(self) -> tuple[list[str], list[str]]: + return ["data"], [] + + def process(self, task: DocumentBatch) -> DocumentBatch: + import packaging + + try: + from loguru import logger + + loguru_available = logger is not None + except ImportError: + loguru_available = False + + batch = task.to_pandas().copy() + batch[f"{self.name}_version"] = packaging.__version__ + batch[f"{self.name}_loguru_available"] = loguru_available + return DocumentBatch( + task_id=task.task_id, + dataset_name=task.dataset_name, + data=batch, + _metadata=task._metadata, + _stage_perf=task._stage_perf, + ) + + +def _make_initial_task() -> DocumentBatch: + return DocumentBatch( + task_id="runtime_env_test", + dataset_name="test", + data=pd.DataFrame({"text": ["hello"]}), + ) + + +@pytest.mark.parametrize( + "backend_config", + [ + pytest.param((RayDataExecutor, {}), id="ray_data"), + pytest.param((XennaExecutor, {"execution_mode": "streaming"}), id="xenna_streaming"), + ], + indirect=True, +) +class TestPerStageRuntimeEnv: + """Stages with different runtime_env see different package versions.""" + + backend_cls: type[BaseExecutor] | None = None + config: dict[str, Any] | None = None + results: list[DocumentBatch] | None = None + + @pytest.fixture(scope="class", autouse=True) + def backend_config(self, request: pytest.FixtureRequest, shared_ray_cluster: str): + """Execute a 3-stage pipeline: base env, packaging==23.2 (pip), packaging==24.0 (uv).""" + backend_cls, config = request.param + request.cls.backend_cls = backend_cls + request.cls.config = config + + base_stage = RecordPackagingVersionStage().with_(name="base_env") + stage_v232 = RecordPackagingVersionStage().with_( + name="pinned_v232", + runtime_env={"pip": ["packaging==23.2"]}, + ) + stage_v240 = RecordPackagingVersionStage().with_( + name="pinned_v240", + runtime_env={"uv": ["packaging==24.0"]}, + ) + + pipeline = Pipeline(name="runtime_env_test", stages=[base_stage, stage_v232, stage_v240]) + request.cls.results = pipeline.run(backend_cls(config), initial_tasks=[_make_initial_task()]) + + def test_output_count(self): + assert self.results is not None + assert len(self.results) == 1 + + def test_base_env_uses_installed_version(self): + """Stage with no runtime_env should see the base environment's packaging version.""" + df = self.results[0].to_pandas() + assert "base_env_version" in df.columns + assert df["base_env_version"].iloc[0] # non-empty + + def test_pinned_v232(self): + df = self.results[0].to_pandas() + assert df["pinned_v232_version"].iloc[0] == "23.2" + + def test_pinned_v240(self): + df = self.results[0].to_pandas() + assert df["pinned_v240_version"].iloc[0] == "24.0" + + def test_all_three_versions_differ(self): + """Base env, 23.2, and 24.0 should all be distinct.""" + df = self.results[0].to_pandas() + versions = { + df["base_env_version"].iloc[0], + df["pinned_v232_version"].iloc[0], + df["pinned_v240_version"].iloc[0], + } + assert len(versions) == 3, f"Expected 3 distinct versions, got {versions}" + + def test_runtime_env_is_additive(self): + """Stages with runtime_env can still import base-env packages (loguru is a Curator dep, not Ray).""" + df = self.results[0].to_pandas() + assert bool(df["pinned_v232_loguru_available"].iloc[0]) + assert bool(df["pinned_v240_loguru_available"].iloc[0])