[Feat] Per Stage Runtime Env

.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 }}"

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
 

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)

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]:

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=}")
 

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.

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.

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
 

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])