[Feat] Per Stage Runtime Env

nemo_curator/backends/experimental/ray_data/adapter.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import copy
+import os
 from collections.abc import Callable
 from typing import Any
 
@@ -101,6 +103,18 @@ 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 pipeline resolved pip_specs to a venv, prepend to PYTHONPATH so workers use that env.
+        resolved_path = getattr(self.stage, "_resolved_site_packages_path", None)
+        if resolved_path is not None:
+            env_vars = ray_remote_args.setdefault("runtime_env", {}).setdefault("env_vars", {})
+            existing = env_vars.get("PYTHONPATH") or os.environ.get("PYTHONPATH", "")
+            env_vars["PYTHONPATH"] = f"{resolved_path}:{existing}" if existing else str(resolved_path)
+        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/experimental/utils.py

@@ -61,6 +61,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/xenna/adapter.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+
 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
@@ -54,8 +56,16 @@ 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
+        """Runtime environment for this stage.
+
+        When the pipeline has resolved pip_specs to a venv (_resolved_site_packages_path),
+        we prepend to PYTHONPATH so workers use that env while keeping any existing path.
+        """
+        resolved_path = getattr(self.processing_stage, "_resolved_site_packages_path", None)
+        if resolved_path is not None:
+            existing = os.environ.get("PYTHONPATH", "")
+            new_path = f"{resolved_path}:{existing}" if existing else str(resolved_path)
+            return pipelines_v1.RuntimeEnv(extra_env_vars={"PYTHONPATH": new_path})
         return None
 
     def process_data(self, tasks: list[Task]) -> list[Task] | None:

nemo_curator/pipeline/pipeline.py

@@ -19,6 +19,7 @@
 from nemo_curator.backends.base import BaseExecutor
 from nemo_curator.stages.base import CompositeStage, ProcessingStage
 from nemo_curator.tasks import Task
+from nemo_curator.utils.stage_pip_env import resolve_stage_pip_envs
 
 
 class Pipeline:
@@ -186,6 +187,16 @@ def run(self, executor: BaseExecutor | None = None, initial_tasks: list[Task] |
         """
         self.build()
 
+        # Resolve per-stage pip_specs into venvs so executors can use PYTHONPATH (e.g. Ray Data).
+        # Skip if already resolved (e.g. second run() on same pipeline) to avoid recreating venvs.
+        needs_resolve = any(
+            getattr(s, "pip_specs", None)
+            and not getattr(s, "_resolved_site_packages_path", None)
+            for s in self.stages
+        )
+        if needs_resolve:
+            resolve_stage_pip_envs(self.stages)
+
         if executor is None:
             from nemo_curator.backends.xenna import XennaExecutor
 

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
+    pip_specs: ClassVar[list[str] | None] = None
 
     @property
     @final

nemo_curator/utils/stage_pip_env.py

@@ -0,0 +1,140 @@
+# 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.
+
+"""Resolve per-stage pip_specs into virtualenvs and set _resolved_site_packages_path on stages.
+
+Uses the `uv` CLI to create venvs (no Python dependency on uv). When a stage defines
+pip_specs (e.g. ["vllm==0.6.0"]), the resolver creates a venv with those packages
+and sets the stage's _resolved_site_packages_path so executors can inject PYTHONPATH.
+"""
+
+import atexit
+import shutil
+import sysconfig
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from loguru import logger
+
+if TYPE_CHECKING:
+    from nemo_curator.stages.base import ProcessingStage
+
+
+def _site_packages_for_venv(venv_dir: Path) -> Path:
+    # Windows uses "Lib" (capital), Unix uses "lib"
+    for lib_name in ("lib", "Lib"):
+        lib = venv_dir / lib_name
+        if not lib.exists():
+            continue
+        # On Windows, site-packages may live directly under Lib\ (no python3.x subdir)
+        site_direct = lib / "site-packages"
+        if site_direct.exists():
+            return site_direct.resolve()
+        for p in lib.iterdir():
+            if p.is_dir() and p.name.startswith("python"):
+                site = p / "site-packages"
+                if site.exists():
+                    return site.resolve()
+    msg = f"no site-packages found under {venv_dir}"
+    raise RuntimeError(msg)
+
+
+def _venv_python_exe(venv_path: Path) -> Path:
+    """Path to the venv Python interpreter (portable: Windows Scripts/python.exe, Unix bin/python)."""
+    if sysconfig.get_platform().startswith("win"):
+        return venv_path / "Scripts" / "python.exe"
+    return venv_path / "bin" / "python"
+
+
+def _create_venv_for_specs(
+    specs: list[str], base_dir: Path, subdir_name: str, uv_exe: str
+) -> Path:
+    import subprocess
+
+    venv_root = base_dir / subdir_name
+    venv_root.mkdir(parents=True, exist_ok=True)
+    venv_path = venv_root / ".venv"
+    python_path = _venv_python_exe(venv_path)
+    subprocess.run(  # noqa: S603
+        [uv_exe, "venv", str(venv_path)],
+        check=True,
+        capture_output=True,
+        text=True,
+    )
+    # specs come from stage pip_specs (pipeline author configuration)
+    subprocess.run(  # noqa: S603
+        [uv_exe, "pip", "install", "--python", str(python_path), *specs],
+        check=True,
+        capture_output=True,
+        text=True,
+    )
+    return _site_packages_for_venv(venv_path)
+
+
+def resolve_stage_pip_envs(
+    stages: list["ProcessingStage"],
+    base_dir: Path | None = None,
+) -> None:
+    """Create venvs for stages that have pip_specs and set _resolved_site_packages_path.
+
+    Stages with the same pip_specs (order-independent) share one venv. Uses the `uv`
+    CLI; must be on PATH. Modifies each stage instance in place.
+
+    Args:
+        stages: Flat list of execution stages (e.g. after pipeline build).
+        base_dir: Directory for venv roots. If None, a temp directory is used.
+    """
+    import subprocess
+    import tempfile
+
+    uv_exe = shutil.which("uv")
+    if not uv_exe:
+        logger.warning(
+            "resolve_stage_pip_envs requires `uv` on PATH; skipping.",
+        )
+        return
+    try:
+        subprocess.run(  # noqa: S603
+            [uv_exe, "--version"],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+    except subprocess.CalledProcessError as e:
+        logger.warning(
+            "resolve_stage_pip_envs: uv check failed; skipping. Error: %s",
+            e,
+        )
+        return
+
+    if base_dir is None:
+        tmp_dir = tempfile.mkdtemp(prefix="curator_pip_envs_")
+        atexit.register(shutil.rmtree, tmp_dir, ignore_errors=True)
+        base = Path(tmp_dir)
+    else:
+        base = Path(base_dir)
+    base.mkdir(parents=True, exist_ok=True)
+
+    # Dedupe: same sorted specs -> one venv
+    unique_specs: dict[tuple[str, ...], Path] = {}
+    for stage in stages:
+        specs = getattr(stage, "pip_specs", None)
+        if not specs or not isinstance(specs, list):
+            continue
+        key = tuple(sorted(specs))
+        if key not in unique_specs:
+            subdir = f"env_{len(unique_specs)}"
+            unique_specs[key] = _create_venv_for_specs(list(specs), base, subdir, uv_exe)
+            logger.info("Created venv for pip_specs %s at %s", list(key), unique_specs[key])
+        stage._resolved_site_packages_path = unique_specs[key]