[Feat] Per Stage Runtime Env

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,27 @@
 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:
+        return ray.runtime_env.RuntimeEnv(**self._runtime_env)
+
+    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 +78,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,198 @@
+# 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 environment: different Python package versions per stage.
+
+Each stage declares runtime_env and Ray creates an isolated virtualenv per unique spec set.
+RayData, Xenna, and RayActorPool backends are all tested via parametrization, for both
+pip and uv spec types.
+"""
+
+from typing import Any, ClassVar
+
+import pandas as pd
+import pytest
+
+from nemo_curator.backends.base import BaseExecutor
+from nemo_curator.backends.experimental.ray_actor_pool import RayActorPoolExecutor
+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
+
+
+def _record_packaging_and_loguru(task: DocumentBatch, version_col: str, loguru_col: str) -> DocumentBatch:
+    """Helper: record packaging.__version__ and whether loguru is importable."""
+    import packaging
+
+    try:
+        from loguru import logger
+
+        loguru_available = logger is not None
+    except ImportError:
+        loguru_available = False
+
+    batch = task.to_pandas().copy()
+    batch[version_col] = packaging.__version__
+    batch[loguru_col] = loguru_available
+    return DocumentBatch(
+        task_id=task.task_id,
+        dataset_name=task.dataset_name,
+        data=batch,
+        _metadata=task._metadata,
+        _stage_perf=task._stage_perf,
+    )
+
+
+class BaseEnvStage(ProcessingStage[DocumentBatch, DocumentBatch]):
+    """Stage with no runtime_env — runs in the base environment."""
+
+    name = "base_env"
+    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"], ["base_packaging_version"]
+
+    def process(self, task: DocumentBatch) -> DocumentBatch:
+        import packaging
+
+        batch = task.to_pandas().copy()
+        batch["base_packaging_version"] = packaging.__version__
+        return DocumentBatch(
+            task_id=task.task_id,
+            dataset_name=task.dataset_name,
+            data=batch,
+            _metadata=task._metadata,
+            _stage_perf=task._stage_perf,
+        )
+
+
+class VersionStage1(ProcessingStage[DocumentBatch, DocumentBatch]):
+    """Stage 1: packaging==23.2 (default pip; overridable via with_())."""
+
+    name = "version_stage_1"
+    resources = Resources(cpus=0.5)
+    batch_size = 1
+    runtime_env: ClassVar[dict] = {"pip": ["packaging==23.2"]}
+
+    def inputs(self) -> tuple[list[str], list[str]]:
+        return ["data"], ["base_packaging_version"]
+
+    def outputs(self) -> tuple[list[str], list[str]]:
+        return ["data"], ["stage1_packaging_version", "stage1_loguru_available"]
+
+    def process(self, task: DocumentBatch) -> DocumentBatch:
+        return _record_packaging_and_loguru(task, "stage1_packaging_version", "stage1_loguru_available")
+
+
+class VersionStage2(ProcessingStage[DocumentBatch, DocumentBatch]):
+    """Stage 2: packaging==24.0 (default pip; overridable via with_())."""
+
+    name = "version_stage_2"
+    resources = Resources(cpus=0.5)
+    batch_size = 1
+    runtime_env: ClassVar[dict] = {"pip": ["packaging==24.0"]}
+
+    def inputs(self) -> tuple[list[str], list[str]]:
+        return ["data"], ["stage1_packaging_version", "stage1_loguru_available"]
+
+    def outputs(self) -> tuple[list[str], list[str]]:
+        return ["data"], ["stage2_packaging_version", "stage2_loguru_available"]
+
+    def process(self, task: DocumentBatch) -> DocumentBatch:
+        return _record_packaging_and_loguru(task, "stage2_packaging_version", "stage2_loguru_available")
+
+
+@pytest.mark.parametrize(
+    "backend_config",
+    [
+        pytest.param((RayDataExecutor, {}, "pip"), id="ray_data-pip"),
+        pytest.param((RayDataExecutor, {}, "uv"), id="ray_data-uv"),
+        pytest.param((XennaExecutor, {"execution_mode": "streaming"}, "pip"), id="xenna_streaming-pip"),
+        pytest.param((XennaExecutor, {"execution_mode": "streaming"}, "uv"), id="xenna_streaming-uv"),
+        pytest.param((XennaExecutor, {"execution_mode": "batch"}, "pip"), id="xenna_batch-pip"),
+        pytest.param((XennaExecutor, {"execution_mode": "batch"}, "uv"), id="xenna_batch-uv"),
+        pytest.param((RayActorPoolExecutor, {}, "pip"), id="ray_actor_pool-pip"),
+        pytest.param((RayActorPoolExecutor, {}, "uv"), id="ray_actor_pool-uv"),
+    ],
+    indirect=True,
+)
+class TestPerStageRuntimeEnv:
+    """Stages with different runtime_env see different package versions across all backends and spec types."""
+
+    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):
+        """Run the three-stage pipeline once per backend/spec_type and store results."""
+        backend_cls, config, spec_type = request.param
+        request.cls.backend_cls = backend_cls
+        request.cls.config = config
+
+        initial = DocumentBatch(
+            task_id="test",
+            dataset_name="test",
+            data=pd.DataFrame({"text": ["hello"]}),
+        )
+        # Use with_() to override runtime_env spec type (pip or uv) per parametrize run.
+        stage1 = VersionStage1().with_(runtime_env={spec_type: ["packaging==23.2"]})
+        stage2 = VersionStage2().with_(runtime_env={spec_type: ["packaging==24.0"]})
+        pipeline = Pipeline(
+            name="per_stage_runtime_env_test",
+            stages=[BaseEnvStage(), stage1, stage2],
+        )
+        request.cls.results = pipeline.run(executor=backend_cls(config), initial_tasks=[initial])
+
+    def test_stage1_packaging_version(self):
+        assert self.results is not None
+        result = self.results[0].to_pandas()
+        assert result["stage1_packaging_version"].iloc[0] == "23.2"
+
+    def test_stage2_packaging_version(self):
+        assert self.results is not None
+        result = self.results[0].to_pandas()
+        assert result["stage2_packaging_version"].iloc[0] == "24.0"
+
+    def test_base_env_uses_installed_version(self):
+        """Stage with no runtime_env sees the base environment's packaging version."""
+        assert self.results is not None
+        result = self.results[0].to_pandas()
+        assert "base_packaging_version" in result.columns
+        assert result["base_packaging_version"].iloc[0]  # non-empty
+
+    def test_all_three_versions_differ(self):
+        """Base env, 23.2, and 24.0 must all be distinct."""
+        assert self.results is not None
+        result = self.results[0].to_pandas()
+        versions = {
+            result["base_packaging_version"].iloc[0],
+            result["stage1_packaging_version"].iloc[0],
+            result["stage2_packaging_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 a Ray dep)."""
+        assert self.results is not None
+        result = self.results[0].to_pandas()
+        assert bool(result["stage1_loguru_available"].iloc[0]), "loguru not importable in stage1 runtime_env"
+        assert bool(result["stage2_loguru_available"].iloc[0]), "loguru not importable in stage2 runtime_env"