From 5271c7c6cae9fe5622515d8b04c56d6f7666aaf9 Mon Sep 17 00:00:00 2001 From: Ao Tang Date: Tue, 17 Mar 2026 11:21:32 -0700 Subject: [PATCH 01/22] Add support for per-stage pip specifications and virtual environments Signed-off-by: Ao Tang --- .../backends/experimental/ray_data/adapter.py | 13 +++ nemo_curator/backends/experimental/utils.py | 1 + nemo_curator/backends/xenna/adapter.py | 10 +- nemo_curator/pipeline/pipeline.py | 5 + nemo_curator/stages/base.py | 1 + nemo_curator/utils/stage_pip_env.py | 104 ++++++++++++++++++ 6 files changed, 132 insertions(+), 2 deletions(-) create mode 100644 nemo_curator/utils/stage_pip_env.py diff --git a/nemo_curator/backends/experimental/ray_data/adapter.py b/nemo_curator/backends/experimental/ray_data/adapter.py index d4f74a12b9..77dbd9734c 100644 --- a/nemo_curator/backends/experimental/ray_data/adapter.py +++ b/nemo_curator/backends/experimental/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,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, inject PYTHONPATH so workers use that env. + resolved_path = getattr(self.stage, "_resolved_site_packages_path", None) + if resolved_path is not None: + ray_remote_args.setdefault("runtime_env", {}).setdefault("env_vars", {})[ + "PYTHONPATH" + ] = 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=}") diff --git a/nemo_curator/backends/experimental/utils.py b/nemo_curator/backends/experimental/utils.py index d9e753bd90..9d62e98d19 100644 --- a/nemo_curator/backends/experimental/utils.py +++ b/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]: diff --git a/nemo_curator/backends/xenna/adapter.py b/nemo_curator/backends/xenna/adapter.py index 84b36b4cdc..b3a662d557 100644 --- a/nemo_curator/backends/xenna/adapter.py +++ b/nemo_curator/backends/xenna/adapter.py @@ -54,8 +54,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 + """Runtime environment for this stage. + + When the pipeline has resolved pip_specs to a venv (_resolved_site_packages_path), + we return a RuntimeEnv with PYTHONPATH so this stage's workers use that env. + """ + resolved_path = getattr(self.processing_stage, "_resolved_site_packages_path", None) + if resolved_path is not None: + return pipelines_v1.RuntimeEnv(extra_env_vars={"PYTHONPATH": str(resolved_path)}) return None def process_data(self, tasks: list[Task]) -> list[Task] | None: diff --git a/nemo_curator/pipeline/pipeline.py b/nemo_curator/pipeline/pipeline.py index 2f42b505e9..fce1136a56 100644 --- a/nemo_curator/pipeline/pipeline.py +++ b/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,10 @@ 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). + if any(getattr(s, "pip_specs", None) for s in self.stages): + resolve_stage_pip_envs(self.stages) + if executor is None: from nemo_curator.backends.xenna import XennaExecutor diff --git a/nemo_curator/stages/base.py b/nemo_curator/stages/base.py index cbf7652ac0..6066b2e7a4 100644 --- a/nemo_curator/stages/base.py +++ b/nemo_curator/stages/base.py @@ -85,6 +85,7 @@ class ProcessingStage(ABC, Generic[X, Y], metaclass=StageMeta): name = "ProcessingStage" resources = Resources(cpus=1.0) batch_size = 1 + pip_specs: list[str] | None = None @property @final diff --git a/nemo_curator/utils/stage_pip_env.py b/nemo_curator/utils/stage_pip_env.py new file mode 100644 index 0000000000..0ac0ce9685 --- /dev/null +++ b/nemo_curator/utils/stage_pip_env.py @@ -0,0 +1,104 @@ +# 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. +""" + +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: + lib = venv_dir / "lib" + if not lib.exists(): + msg = f"venv has no lib/: {venv_dir}" + raise RuntimeError(msg) + 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 _create_venv_for_specs(specs: list[str], base_dir: Path, subdir_name: 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_path / "bin" / "python" + subprocess.run( + ["uv", "venv", str(venv_path)], + check=True, + capture_output=True, + ) + subprocess.run( + ["uv", "pip", "install", "--python", str(python_path), *specs], + check=True, + capture_output=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 + + try: + subprocess.run(["uv", "--version"], check=True, capture_output=True) + except (subprocess.CalledProcessError, FileNotFoundError) as e: + logger.warning( + "resolve_stage_pip_envs requires `uv` on PATH; skipping. Error: %s", + e, + ) + return + + base = Path(base_dir) if base_dir else Path(tempfile.mkdtemp(prefix="curator_pip_envs_")) + base.mkdir(parents=True, exist_ok=True) + + # Dedupe: same sorted specs -> one venv + unique_specs: dict[tuple[str, ...], Path] = {} + for i, stage in enumerate(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) + logger.info("Created venv for pip_specs %s at %s", list(key), unique_specs[key]) + setattr(stage, "_resolved_site_packages_path", unique_specs[key]) From 9d1fd4e11e5c8527c6e67ae971b9d70690539b88 Mon Sep 17 00:00:00 2001 From: Ao Tang Date: Wed, 18 Mar 2026 09:06:58 -0700 Subject: [PATCH 02/22] add test Signed-off-by: Ao Tang --- tests/pipelines/test_per_stage_runtime_env.py | 162 ++++++++++++++++++ 1 file changed, 162 insertions(+) create mode 100644 tests/pipelines/test_per_stage_runtime_env.py 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..0c241f9ccf --- /dev/null +++ b/tests/pipelines/test_per_stage_runtime_env.py @@ -0,0 +1,162 @@ +# 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. + +Uses pip_specs + resolve_stage_pip_envs (uv CLI) so each stage runs with its own +packaging version. Requires `uv` on PATH; test is skipped if uv is not available. +See tutorials/per_stage_runtime_env_example.py and docs/design/per-stage-runtime-environment.md. +""" + +import subprocess + +import pandas as pd +import pytest + +from nemo_curator.backends.experimental.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 _uv_available() -> bool: + try: + subprocess.run(["uv", "--version"], check=True, capture_output=True) + return True + except (subprocess.CalledProcessError, FileNotFoundError): + return False + + +@pytest.fixture +def require_uv(): + if not _uv_available(): + pytest.skip("uv not on PATH; per-stage pip_specs tests require uv") + + +class VersionStage1(ProcessingStage[DocumentBatch, DocumentBatch]): + """Stage 1: packaging==23.2 via pip_specs (resolver creates venv, PYTHONPATH).""" + + name = "version_stage_1" + resources = Resources(cpus=0.5) + batch_size = 1 + pip_specs = ["packaging==23.2"] + + def inputs(self) -> tuple[list[str], list[str]]: + return ["data"], [] + + def outputs(self) -> tuple[list[str], list[str]]: + return ["data"], ["stage1_packaging_version"] + + def process(self, task: DocumentBatch) -> DocumentBatch: + import packaging # noqa: PLC0415 + + df = task.to_pandas().copy() + df["stage1_packaging_version"] = packaging.__version__ + return DocumentBatch( + task_id=task.task_id, + dataset_name=task.dataset_name, + data=df, + _metadata=task._metadata, + _stage_perf=task._stage_perf, + ) + + +class VersionStage2(ProcessingStage[DocumentBatch, DocumentBatch]): + """Stage 2: packaging==24.0 via pip_specs (resolver creates venv, PYTHONPATH).""" + + name = "version_stage_2" + resources = Resources(cpus=0.5) + batch_size = 1 + pip_specs = ["packaging==24.0"] + + def inputs(self) -> tuple[list[str], list[str]]: + return ["data"], ["stage1_packaging_version"] + + def outputs(self) -> tuple[list[str], list[str]]: + return ["data"], ["stage2_packaging_version"] + + def process(self, task: DocumentBatch) -> DocumentBatch: + import packaging # noqa: PLC0415 + + df = task.to_pandas().copy() + df["stage2_packaging_version"] = packaging.__version__ + return DocumentBatch( + task_id=task.task_id, + dataset_name=task.dataset_name, + data=df, + _metadata=task._metadata, + _stage_perf=task._stage_perf, + ) + + +@pytest.mark.usefixtures("shared_ray_client", "require_uv") +def test_per_stage_different_package_versions_ray_data() -> None: + """Run two stages with different packaging versions via pip_specs; assert each sees its own version. + + Pipeline.run() calls resolve_stage_pip_envs() to create venvs with uv; Ray Data adapter + injects PYTHONPATH so workers load the correct site-packages per stage. + """ + initial = DocumentBatch( + task_id="per_stage_version_test", + dataset_name="test", + data=pd.DataFrame({"text": ["hello"]}), + ) + pipeline = Pipeline( + name="per_stage_version_test", + stages=[VersionStage1(), VersionStage2()], + ) + results = pipeline.run( + executor=RayDataExecutor(), + initial_tasks=[initial], + ) + assert results is not None + assert len(results) == 1 + out = results[0] + df = out.to_pandas() + assert "stage1_packaging_version" in df.columns + assert "stage2_packaging_version" in df.columns + assert df["stage1_packaging_version"].iloc[0] == "23.2", "Stage 1 should see packaging 23.2" + assert df["stage2_packaging_version"].iloc[0] == "24.0", "Stage 2 should see packaging 24.0" + + +@pytest.mark.usefixtures("shared_ray_client", "require_uv") +def test_per_stage_different_package_versions_xenna() -> None: + """Run two stages with different packaging versions via pip_specs using XennaExecutor. + + Pipeline.run() calls resolve_stage_pip_envs() to create venvs with uv; Xenna adapter + uses env_info() to set PYTHONPATH so workers load the correct site-packages per stage. + """ + initial = DocumentBatch( + task_id="per_stage_version_test_xenna", + dataset_name="test", + data=pd.DataFrame({"text": ["hello"]}), + ) + pipeline = Pipeline( + name="per_stage_version_test_xenna", + stages=[VersionStage1(), VersionStage2()], + ) + results = pipeline.run( + executor=XennaExecutor(config={"execution_mode": "streaming"}), + initial_tasks=[initial], + ) + assert results is not None + assert len(results) == 1 + out = results[0] + df = out.to_pandas() + assert "stage1_packaging_version" in df.columns + assert "stage2_packaging_version" in df.columns + assert df["stage1_packaging_version"].iloc[0] == "23.2", "Stage 1 should see packaging 23.2" + assert df["stage2_packaging_version"].iloc[0] == "24.0", "Stage 2 should see packaging 24.0" From b3624fce3dc3e6e1653655b3f01675b1dc00c046 Mon Sep 17 00:00:00 2001 From: Ao Tang Date: Wed, 18 Mar 2026 09:25:22 -0700 Subject: [PATCH 03/22] ruff Signed-off-by: Ao Tang --- tests/pipelines/test_per_stage_runtime_env.py | 20 ++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/tests/pipelines/test_per_stage_runtime_env.py b/tests/pipelines/test_per_stage_runtime_env.py index 0c241f9ccf..7c369b4a4e 100644 --- a/tests/pipelines/test_per_stage_runtime_env.py +++ b/tests/pipelines/test_per_stage_runtime_env.py @@ -19,7 +19,9 @@ See tutorials/per_stage_runtime_env_example.py and docs/design/per-stage-runtime-environment.md. """ +import shutil import subprocess +from typing import ClassVar import pandas as pd import pytest @@ -33,11 +35,15 @@ def _uv_available() -> bool: + uv_exe = shutil.which("uv") + if not uv_exe: + return False try: - subprocess.run(["uv", "--version"], check=True, capture_output=True) - return True - except (subprocess.CalledProcessError, FileNotFoundError): + subprocess.run([uv_exe, "--version"], check=True, capture_output=True) # noqa: S603 + except subprocess.CalledProcessError: return False + else: + return True @pytest.fixture @@ -52,7 +58,7 @@ class VersionStage1(ProcessingStage[DocumentBatch, DocumentBatch]): name = "version_stage_1" resources = Resources(cpus=0.5) batch_size = 1 - pip_specs = ["packaging==23.2"] + pip_specs: ClassVar[list[str]] = ["packaging==23.2"] def inputs(self) -> tuple[list[str], list[str]]: return ["data"], [] @@ -61,7 +67,7 @@ def outputs(self) -> tuple[list[str], list[str]]: return ["data"], ["stage1_packaging_version"] def process(self, task: DocumentBatch) -> DocumentBatch: - import packaging # noqa: PLC0415 + import packaging df = task.to_pandas().copy() df["stage1_packaging_version"] = packaging.__version__ @@ -80,7 +86,7 @@ class VersionStage2(ProcessingStage[DocumentBatch, DocumentBatch]): name = "version_stage_2" resources = Resources(cpus=0.5) batch_size = 1 - pip_specs = ["packaging==24.0"] + pip_specs: ClassVar[list[str]] = ["packaging==24.0"] def inputs(self) -> tuple[list[str], list[str]]: return ["data"], ["stage1_packaging_version"] @@ -89,7 +95,7 @@ def outputs(self) -> tuple[list[str], list[str]]: return ["data"], ["stage2_packaging_version"] def process(self, task: DocumentBatch) -> DocumentBatch: - import packaging # noqa: PLC0415 + import packaging df = task.to_pandas().copy() df["stage2_packaging_version"] = packaging.__version__ From e2ae9a22d75b2d9cff804965cbee1950e624a0a4 Mon Sep 17 00:00:00 2001 From: Ao Tang Date: Wed, 18 Mar 2026 10:01:07 -0700 Subject: [PATCH 04/22] comments resolved Signed-off-by: Ao Tang --- .../backends/experimental/ray_data/adapter.py | 9 +-- nemo_curator/backends/xenna/adapter.py | 8 ++- nemo_curator/stages/base.py | 4 +- nemo_curator/utils/stage_pip_env.py | 57 ++++++++++++++----- 4 files changed, 57 insertions(+), 21 deletions(-) diff --git a/nemo_curator/backends/experimental/ray_data/adapter.py b/nemo_curator/backends/experimental/ray_data/adapter.py index 77dbd9734c..87a9cb0c17 100644 --- a/nemo_curator/backends/experimental/ray_data/adapter.py +++ b/nemo_curator/backends/experimental/ray_data/adapter.py @@ -13,6 +13,7 @@ # limitations under the License. import copy +import os from collections.abc import Callable from typing import Any @@ -106,12 +107,12 @@ def process_dataset(self, dataset: Dataset, ignore_head_node: bool = False) -> D ray_remote_args = copy.deepcopy( self.stage.ray_stage_spec().get(RayStageSpecKeys.RAY_REMOTE_ARGS) or {} ) - # If pipeline resolved pip_specs to a venv, inject PYTHONPATH so workers use that env. + # 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: - ray_remote_args.setdefault("runtime_env", {}).setdefault("env_vars", {})[ - "PYTHONPATH" - ] = str(resolved_path) + 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 diff --git a/nemo_curator/backends/xenna/adapter.py b/nemo_curator/backends/xenna/adapter.py index b3a662d557..fffaec55d0 100644 --- a/nemo_curator/backends/xenna/adapter.py +++ b/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 @@ -57,11 +59,13 @@ def env_info(self) -> pipelines_v1.RuntimeEnv | None: """Runtime environment for this stage. When the pipeline has resolved pip_specs to a venv (_resolved_site_packages_path), - we return a RuntimeEnv with PYTHONPATH so this stage's workers use that env. + 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: - return pipelines_v1.RuntimeEnv(extra_env_vars={"PYTHONPATH": str(resolved_path)}) + 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: diff --git a/nemo_curator/stages/base.py b/nemo_curator/stages/base.py index 6066b2e7a4..4ac93d76e8 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,7 +85,7 @@ class ProcessingStage(ABC, Generic[X, Y], metaclass=StageMeta): name = "ProcessingStage" resources = Resources(cpus=1.0) batch_size = 1 - pip_specs: list[str] | None = None + pip_specs: ClassVar[list[str] | None] = None @property @final diff --git a/nemo_curator/utils/stage_pip_env.py b/nemo_curator/utils/stage_pip_env.py index 0ac0ce9685..be05e7d2e1 100644 --- a/nemo_curator/utils/stage_pip_env.py +++ b/nemo_curator/utils/stage_pip_env.py @@ -19,6 +19,9 @@ 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 @@ -42,22 +45,34 @@ def _site_packages_for_venv(venv_dir: Path) -> Path: raise RuntimeError(msg) -def _create_venv_for_specs(specs: list[str], base_dir: Path, subdir_name: str) -> Path: +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_path / "bin" / "python" - subprocess.run( - ["uv", "venv", str(venv_path)], + python_path = _venv_python_exe(venv_path) + subprocess.run( # noqa: S603 + [uv_exe, "venv", str(venv_path)], check=True, capture_output=True, + text=True, ) - subprocess.run( - ["uv", "pip", "install", "--python", str(python_path), *specs], + # 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) @@ -78,27 +93,43 @@ def resolve_stage_pip_envs( 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(["uv", "--version"], check=True, capture_output=True) - except (subprocess.CalledProcessError, FileNotFoundError) as e: + 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 requires `uv` on PATH; skipping. Error: %s", + "resolve_stage_pip_envs: uv check failed; skipping. Error: %s", e, ) return - base = Path(base_dir) if base_dir else Path(tempfile.mkdtemp(prefix="curator_pip_envs_")) + 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 i, stage in enumerate(stages): + 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) + 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]) - setattr(stage, "_resolved_site_packages_path", unique_specs[key]) + stage._resolved_site_packages_path = unique_specs[key] From b25df9f0755bb48dc120f5b26f9ff7e404880700 Mon Sep 17 00:00:00 2001 From: Ao Tang Date: Wed, 18 Mar 2026 10:13:31 -0700 Subject: [PATCH 05/22] comments resolve Signed-off-by: Ao Tang --- nemo_curator/pipeline/pipeline.py | 8 +++++++- nemo_curator/utils/stage_pip_env.py | 23 ++++++++++++++--------- 2 files changed, 21 insertions(+), 10 deletions(-) diff --git a/nemo_curator/pipeline/pipeline.py b/nemo_curator/pipeline/pipeline.py index fce1136a56..19e5e2673d 100644 --- a/nemo_curator/pipeline/pipeline.py +++ b/nemo_curator/pipeline/pipeline.py @@ -188,7 +188,13 @@ 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). - if any(getattr(s, "pip_specs", None) for s in self.stages): + # 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: diff --git a/nemo_curator/utils/stage_pip_env.py b/nemo_curator/utils/stage_pip_env.py index be05e7d2e1..6c1fbb1807 100644 --- a/nemo_curator/utils/stage_pip_env.py +++ b/nemo_curator/utils/stage_pip_env.py @@ -32,15 +32,20 @@ def _site_packages_for_venv(venv_dir: Path) -> Path: - lib = venv_dir / "lib" - if not lib.exists(): - msg = f"venv has no lib/: {venv_dir}" - raise RuntimeError(msg) - for p in lib.iterdir(): - if p.is_dir() and p.name.startswith("python"): - site = p / "site-packages" - if site.exists(): - return site.resolve() + # 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) From e24375552fcf1c61da3f7237938c147321593a4b Mon Sep 17 00:00:00 2001 From: Ao Tang Date: Wed, 18 Mar 2026 10:40:57 -0700 Subject: [PATCH 06/22] comments resolve Signed-off-by: Ao Tang --- nemo_curator/pipeline/pipeline.py | 15 ++++++----- nemo_curator/utils/stage_pip_env.py | 42 ++++++++++++++++++++--------- 2 files changed, 37 insertions(+), 20 deletions(-) diff --git a/nemo_curator/pipeline/pipeline.py b/nemo_curator/pipeline/pipeline.py index 19e5e2673d..d6f5de55a7 100644 --- a/nemo_curator/pipeline/pipeline.py +++ b/nemo_curator/pipeline/pipeline.py @@ -188,14 +188,15 @@ 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) + # Only pass unresolved stages so already-resolved stages keep their venvs. + unresolved = [ + s for s in self.stages - ) - if needs_resolve: - resolve_stage_pip_envs(self.stages) + if getattr(s, "pip_specs", None) + and not getattr(s, "_resolved_site_packages_path", None) + ] + if unresolved: + resolve_stage_pip_envs(unresolved) if executor is None: from nemo_curator.backends.xenna import XennaExecutor diff --git a/nemo_curator/utils/stage_pip_env.py b/nemo_curator/utils/stage_pip_env.py index 6c1fbb1807..d2804323c4 100644 --- a/nemo_curator/utils/stage_pip_env.py +++ b/nemo_curator/utils/stage_pip_env.py @@ -73,12 +73,17 @@ def _create_venv_for_specs( 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, - ) + try: + subprocess.run( # noqa: S603 + [uv_exe, "pip", "install", "--python", str(python_path), *specs], + check=True, + capture_output=True, + text=True, + ) + except subprocess.CalledProcessError as e: + logger.error("uv pip install failed for specs %s:\n%s", specs, e.stderr) + raise + return _site_packages_for_venv(venv_path) @@ -98,11 +103,17 @@ def resolve_stage_pip_envs( import subprocess import tempfile + stages_with_pip = [s for s in stages if getattr(s, "pip_specs", None)] uv_exe = shutil.which("uv") if not uv_exe: - logger.warning( - "resolve_stage_pip_envs requires `uv` on PATH; skipping.", - ) + if stages_with_pip: + names = ", ".join(getattr(s, "name", s.__class__.__name__) for s in stages_with_pip) + msg = ( + "Stages with pip_specs require `uv` on PATH to resolve dependencies, " + f"but `uv` was not found. Affected stages: {names}. " + "Install uv (e.g. pip install uv) or add it to PATH." + ) + raise RuntimeError(msg) return try: subprocess.run( # noqa: S603 @@ -112,10 +123,13 @@ def resolve_stage_pip_envs( text=True, ) except subprocess.CalledProcessError as e: - logger.warning( - "resolve_stage_pip_envs: uv check failed; skipping. Error: %s", - e, - ) + if stages_with_pip: + names = ", ".join(getattr(s, "name", s.__class__.__name__) for s in stages_with_pip) + msg = ( + f"Stages with pip_specs require a working `uv` CLI; uv check failed. " + f"Affected stages: {names}. Error: {e}" + ) + raise RuntimeError(msg) from e return if base_dir is None: @@ -129,6 +143,8 @@ def resolve_stage_pip_envs( # Dedupe: same sorted specs -> one venv unique_specs: dict[tuple[str, ...], Path] = {} for stage in stages: + if getattr(stage, "_resolved_site_packages_path", None): + continue # already resolved; skip specs = getattr(stage, "pip_specs", None) if not specs or not isinstance(specs, list): continue From 98470cad0a0fbb64d846b664c3446a4f893328f2 Mon Sep 17 00:00:00 2001 From: Ao Tang Date: Thu, 19 Mar 2026 07:23:51 -0700 Subject: [PATCH 07/22] normalize pip specifications before creating virtual environments. Signed-off-by: Ao Tang --- nemo_curator/utils/stage_pip_env.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/nemo_curator/utils/stage_pip_env.py b/nemo_curator/utils/stage_pip_env.py index d2804323c4..a11ec5fbe6 100644 --- a/nemo_curator/utils/stage_pip_env.py +++ b/nemo_curator/utils/stage_pip_env.py @@ -140,7 +140,7 @@ def resolve_stage_pip_envs( base = Path(base_dir) base.mkdir(parents=True, exist_ok=True) - # Dedupe: same sorted specs -> one venv + # Dedupe by normalized specs unique_specs: dict[tuple[str, ...], Path] = {} for stage in stages: if getattr(stage, "_resolved_site_packages_path", None): @@ -148,9 +148,9 @@ def resolve_stage_pip_envs( specs = getattr(stage, "pip_specs", None) if not specs or not isinstance(specs, list): continue - key = tuple(sorted(specs)) + key = tuple(sorted(s.strip().lower() for s in 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]) + logger.info("Created venv for pip_specs %s at %s", list(specs), unique_specs[key]) stage._resolved_site_packages_path = unique_specs[key] From fc8f51007f4b2683b53663fe0bb29b3a49e255bb Mon Sep 17 00:00:00 2001 From: Ao Tang Date: Thu, 19 Mar 2026 07:54:26 -0700 Subject: [PATCH 08/22] Add Path import and update ProcessingStage class Signed-off-by: Ao Tang --- nemo_curator/stages/base.py | 2 ++ tests/pipelines/test_per_stage_runtime_env.py | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/nemo_curator/stages/base.py b/nemo_curator/stages/base.py index b3791a7f20..34ef03face 100644 --- a/nemo_curator/stages/base.py +++ b/nemo_curator/stages/base.py @@ -19,6 +19,7 @@ import time from abc import ABC, ABCMeta, abstractmethod from inspect import isabstract +from pathlib import Path from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, final from loguru import logger @@ -86,6 +87,7 @@ class ProcessingStage(ABC, Generic[X, Y], metaclass=StageMeta): resources = Resources(cpus=1.0) batch_size = 1 pip_specs: ClassVar[list[str] | None] = None + _resolved_site_packages_path: Path | None = None # set by resolve_stage_pip_envs @property @final diff --git a/tests/pipelines/test_per_stage_runtime_env.py b/tests/pipelines/test_per_stage_runtime_env.py index 7c369b4a4e..ad2cb1ec4e 100644 --- a/tests/pipelines/test_per_stage_runtime_env.py +++ b/tests/pipelines/test_per_stage_runtime_env.py @@ -39,7 +39,7 @@ def _uv_available() -> bool: if not uv_exe: return False try: - subprocess.run([uv_exe, "--version"], check=True, capture_output=True) # noqa: S603 + subprocess.run([uv_exe, "--version"], check=True, capture_output=True, text=True) # noqa: S603 except subprocess.CalledProcessError: return False else: From 2666b2411699faad424d8b5b11bbe78653319be6 Mon Sep 17 00:00:00 2001 From: Ao Tang Date: Thu, 19 Mar 2026 07:57:18 -0700 Subject: [PATCH 09/22] ruff check Signed-off-by: Ao Tang --- nemo_curator/stages/base.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/nemo_curator/stages/base.py b/nemo_curator/stages/base.py index 34ef03face..ecc778d147 100644 --- a/nemo_curator/stages/base.py +++ b/nemo_curator/stages/base.py @@ -19,7 +19,6 @@ import time from abc import ABC, ABCMeta, abstractmethod from inspect import isabstract -from pathlib import Path from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, final from loguru import logger @@ -28,6 +27,8 @@ from nemo_curator.tasks import Task if TYPE_CHECKING: + from pathlib import Path + from nemo_curator.backends.base import NodeInfo, WorkerMetadata X = TypeVar("X", bound=Task) # Input task type From f4ef2097104f6a216ff74b1abf9c3a4d1d01e793 Mon Sep 17 00:00:00 2001 From: Ao Tang Date: Thu, 19 Mar 2026 08:08:59 -0700 Subject: [PATCH 10/22] ruff Signed-off-by: Ao Tang --- tests/pipelines/test_per_stage_runtime_env.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/pipelines/test_per_stage_runtime_env.py b/tests/pipelines/test_per_stage_runtime_env.py index ad2cb1ec4e..eb1553210e 100644 --- a/tests/pipelines/test_per_stage_runtime_env.py +++ b/tests/pipelines/test_per_stage_runtime_env.py @@ -21,7 +21,6 @@ import shutil import subprocess -from typing import ClassVar import pandas as pd import pytest @@ -58,7 +57,7 @@ class VersionStage1(ProcessingStage[DocumentBatch, DocumentBatch]): name = "version_stage_1" resources = Resources(cpus=0.5) batch_size = 1 - pip_specs: ClassVar[list[str]] = ["packaging==23.2"] + pip_specs = ["packaging==23.2"] # noqa: RUF012 def inputs(self) -> tuple[list[str], list[str]]: return ["data"], [] @@ -86,7 +85,7 @@ class VersionStage2(ProcessingStage[DocumentBatch, DocumentBatch]): name = "version_stage_2" resources = Resources(cpus=0.5) batch_size = 1 - pip_specs: ClassVar[list[str]] = ["packaging==24.0"] + pip_specs = ["packaging==24.0"] # noqa: RUF012 def inputs(self) -> tuple[list[str], list[str]]: return ["data"], ["stage1_packaging_version"] From e971625aa4f905fa0a19ccc3e93774d458fa03ed Mon Sep 17 00:00:00 2001 From: Ao Tang Date: Thu, 19 Mar 2026 08:11:13 -0700 Subject: [PATCH 11/22] ruff check Signed-off-by: Ao Tang --- tests/pipelines/test_per_stage_runtime_env.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/pipelines/test_per_stage_runtime_env.py b/tests/pipelines/test_per_stage_runtime_env.py index eb1553210e..799818e1f0 100644 --- a/tests/pipelines/test_per_stage_runtime_env.py +++ b/tests/pipelines/test_per_stage_runtime_env.py @@ -19,7 +19,7 @@ See tutorials/per_stage_runtime_env_example.py and docs/design/per-stage-runtime-environment.md. """ -import shutil +import shutil # noqa: I001 import subprocess import pandas as pd From fab1b90e79c668113db503eb81b9fcdc1a508ecb Mon Sep 17 00:00:00 2001 From: Ao Tang Date: Thu, 19 Mar 2026 08:52:30 -0700 Subject: [PATCH 12/22] Add runtime environment conflict check in RayDataStageAdapter and clean up imports in base.py Signed-off-by: Ao Tang --- nemo_curator/backends/ray_data/adapter.py | 6 ++++++ nemo_curator/stages/base.py | 3 +-- nemo_curator/utils/stage_pip_env.py | 16 ---------------- 3 files changed, 7 insertions(+), 18 deletions(-) diff --git a/nemo_curator/backends/ray_data/adapter.py b/nemo_curator/backends/ray_data/adapter.py index 87a9cb0c17..04bce984ea 100644 --- a/nemo_curator/backends/ray_data/adapter.py +++ b/nemo_curator/backends/ray_data/adapter.py @@ -110,6 +110,12 @@ def process_dataset(self, dataset: Dataset, ignore_head_node: bool = False) -> D # 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: + if ray_remote_args.get("runtime_env", {}).get("pip"): + msg = ( + f"Stage {self.stage.__class__.__name__} defines both pip_specs and ray_remote_args runtime_env.pip; " + "package versions may conflict. Please use only one of them." + ) + raise RuntimeError(msg) 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) diff --git a/nemo_curator/stages/base.py b/nemo_curator/stages/base.py index ecc778d147..9a3ad1e7aa 100644 --- a/nemo_curator/stages/base.py +++ b/nemo_curator/stages/base.py @@ -19,6 +19,7 @@ import time from abc import ABC, ABCMeta, abstractmethod from inspect import isabstract +from pathlib import Path # noqa: TC003 from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, final from loguru import logger @@ -27,8 +28,6 @@ from nemo_curator.tasks import Task if TYPE_CHECKING: - from pathlib import Path - from nemo_curator.backends.base import NodeInfo, WorkerMetadata X = TypeVar("X", bound=Task) # Input task type diff --git a/nemo_curator/utils/stage_pip_env.py b/nemo_curator/utils/stage_pip_env.py index a11ec5fbe6..449f143863 100644 --- a/nemo_curator/utils/stage_pip_env.py +++ b/nemo_curator/utils/stage_pip_env.py @@ -115,22 +115,6 @@ def resolve_stage_pip_envs( ) raise RuntimeError(msg) return - try: - subprocess.run( # noqa: S603 - [uv_exe, "--version"], - check=True, - capture_output=True, - text=True, - ) - except subprocess.CalledProcessError as e: - if stages_with_pip: - names = ", ".join(getattr(s, "name", s.__class__.__name__) for s in stages_with_pip) - msg = ( - f"Stages with pip_specs require a working `uv` CLI; uv check failed. " - f"Affected stages: {names}. Error: {e}" - ) - raise RuntimeError(msg) from e - return if base_dir is None: tmp_dir = tempfile.mkdtemp(prefix="curator_pip_envs_") From e89badbf6b7d7f94cfb98aaf1a615b085b9f1dc2 Mon Sep 17 00:00:00 2001 From: Ao Tang Date: Thu, 19 Mar 2026 08:57:24 -0700 Subject: [PATCH 13/22] fix Signed-off-by: Ao Tang --- nemo_curator/utils/stage_pip_env.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/nemo_curator/utils/stage_pip_env.py b/nemo_curator/utils/stage_pip_env.py index 449f143863..aab5f1b0c9 100644 --- a/nemo_curator/utils/stage_pip_env.py +++ b/nemo_curator/utils/stage_pip_env.py @@ -100,7 +100,6 @@ def resolve_stage_pip_envs( 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 stages_with_pip = [s for s in stages if getattr(s, "pip_specs", None)] @@ -135,6 +134,6 @@ def resolve_stage_pip_envs( key = tuple(sorted(s.strip().lower() for s in 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(specs), unique_specs[key]) + unique_specs[key] = _create_venv_for_specs(list(key), 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] From 806242740e277e2b90a96a4b38414d08474994ec Mon Sep 17 00:00:00 2001 From: Ao Tang Date: Fri, 20 Mar 2026 10:32:49 -0700 Subject: [PATCH 14/22] Refactor pip_specs to use ClassVar for type hinting in VersionStage1 and VersionStage2 classes Signed-off-by: Ao Tang --- tests/pipelines/test_per_stage_runtime_env.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/pipelines/test_per_stage_runtime_env.py b/tests/pipelines/test_per_stage_runtime_env.py index 799818e1f0..e184dd54ec 100644 --- a/tests/pipelines/test_per_stage_runtime_env.py +++ b/tests/pipelines/test_per_stage_runtime_env.py @@ -25,6 +25,8 @@ import pandas as pd import pytest +from typing import ClassVar + from nemo_curator.backends.experimental.ray_data import RayDataExecutor from nemo_curator.backends.xenna import XennaExecutor from nemo_curator.pipeline.pipeline import Pipeline @@ -57,7 +59,7 @@ class VersionStage1(ProcessingStage[DocumentBatch, DocumentBatch]): name = "version_stage_1" resources = Resources(cpus=0.5) batch_size = 1 - pip_specs = ["packaging==23.2"] # noqa: RUF012 + pip_specs: ClassVar[list[str]] = ["packaging==23.2"] def inputs(self) -> tuple[list[str], list[str]]: return ["data"], [] @@ -85,7 +87,7 @@ class VersionStage2(ProcessingStage[DocumentBatch, DocumentBatch]): name = "version_stage_2" resources = Resources(cpus=0.5) batch_size = 1 - pip_specs = ["packaging==24.0"] # noqa: RUF012 + pip_specs: ClassVar[list[str]] = ["packaging==24.0"] def inputs(self) -> tuple[list[str], list[str]]: return ["data"], ["stage1_packaging_version"] From 297b8e4826f708097ccff8ebbcafbb8f49855d32 Mon Sep 17 00:00:00 2001 From: Ao Tang Date: Mon, 6 Apr 2026 11:14:19 -0700 Subject: [PATCH 15/22] refactor to use runtime_env Signed-off-by: Ao Tang --- .../experimental/ray_actor_pool/executor.py | 12 +- nemo_curator/backends/ray_data/adapter.py | 22 +-- nemo_curator/backends/ray_data/executor.py | 10 ++ nemo_curator/backends/xenna/adapter.py | 40 +++-- nemo_curator/backends/xenna/executor.py | 7 + nemo_curator/pipeline/pipeline.py | 12 -- nemo_curator/stages/base.py | 15 +- nemo_curator/utils/stage_pip_env.py | 139 ------------------ tests/pipelines/test_per_stage_runtime_env.py | 119 ++++----------- 9 files changed, 100 insertions(+), 276 deletions(-) delete mode 100644 nemo_curator/utils/stage_pip_env.py 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/ray_data/adapter.py b/nemo_curator/backends/ray_data/adapter.py index 04bce984ea..e3abe2847d 100644 --- a/nemo_curator/backends/ray_data/adapter.py +++ b/nemo_curator/backends/ray_data/adapter.py @@ -13,7 +13,6 @@ # limitations under the License. import copy -import os from collections.abc import Callable from typing import Any @@ -104,21 +103,12 @@ def process_dataset(self, dataset: Dataset, ignore_head_node: bool = False) -> D 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: - if ray_remote_args.get("runtime_env", {}).get("pip"): - msg = ( - f"Stage {self.stage.__class__.__name__} defines both pip_specs and ray_remote_args runtime_env.pip; " - "package versions may conflict. Please use only one of them." - ) - raise RuntimeError(msg) - 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) + 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 diff --git a/nemo_curator/backends/ray_data/executor.py b/nemo_curator/backends/ray_data/executor.py index 63670d20cc..ca211f3a4b 100644 --- a/nemo_curator/backends/ray_data/executor.py +++ b/nemo_curator/backends/ray_data/executor.py @@ -61,7 +61,17 @@ 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. ensurepip.bootstrap() below ensures pip is available + # in the cloned virtualenv (see comment there for details). try: + # Ray clones the current virtualenv when creating per-stage pip virtualenvs. + # The NeMo Curator container's /opt/venv is created by uv, which does not include + # pip as a seed package, so the clone also lacks pip and `python -m pip install` + # fails inside the worker virtualenv. Bootstrap pip first so the clone inherits it. + import ensurepip + + ensurepip.bootstrap(upgrade=True) # Initialize ray and explicitly set NOSET to empty # This ensures if Xenna was used before which was setting NOSET, we end up overriding it. ray.init( diff --git a/nemo_curator/backends/xenna/adapter.py b/nemo_curator/backends/xenna/adapter.py index fffaec55d0..6ae02b0b7e 100644 --- a/nemo_curator/backends/xenna/adapter.py +++ b/nemo_curator/backends/xenna/adapter.py @@ -12,12 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os +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 -from cosmos_xenna.pipelines.private.resources import WorkerMetadata as XennaWorkerMetadata +from cosmos_xenna.pipelines.v1 import NodeInfo as XennaNodeInfo +from cosmos_xenna.pipelines.v1 import Resources as XennaResources +from cosmos_xenna.pipelines.v1 import WorkerMetadata as XennaWorkerMetadata from loguru import logger from nemo_curator.backends.base import BaseStageAdapter, NodeInfo, WorkerMetadata @@ -25,6 +26,24 @@ 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 two methods Xenna calls: + ``to_ray_runtime_env()`` and ``format()``. + """ + + def __init__(self, runtime_env: dict[str, Any]) -> None: + self._runtime_env = runtime_env + + 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: @@ -58,15 +77,12 @@ def stage_batch_size(self) -> int: def env_info(self) -> pipelines_v1.RuntimeEnv | None: """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. + Converts the ProcessingStage.runtime_env dict (Ray-format) to a + CuratorRuntimeEnv that Xenna can forward to Ray actors. """ - 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 + 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/backends/xenna/executor.py b/nemo_curator/backends/xenna/executor.py index aaf51c8383..ddfc7bfe16 100644 --- a/nemo_curator/backends/xenna/executor.py +++ b/nemo_curator/backends/xenna/executor.py @@ -135,6 +135,13 @@ def execute(self, stages: list[ProcessingStage], initial_tasks: list[Task] | Non logger.info(f"Execution mode: {exec_mode.name}") try: + # Ray clones the current virtualenv when creating per-stage pip virtualenvs. + # The NeMo Curator container's /opt/venv is created by uv, which does not include + # pip as a seed package, so the clone also lacks pip and `python -m pip install` + # fails inside the worker virtualenv. Bootstrap pip first so the clone inherits it. + import ensurepip + + ensurepip.bootstrap(upgrade=True) register_loguru_serializer() # Prevent Ray from overriding accelerator env vars when num_gpus=0, letting Xenna manage them instead. ray.init( diff --git a/nemo_curator/pipeline/pipeline.py b/nemo_curator/pipeline/pipeline.py index d6f5de55a7..2f42b505e9 100644 --- a/nemo_curator/pipeline/pipeline.py +++ b/nemo_curator/pipeline/pipeline.py @@ -19,7 +19,6 @@ 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: @@ -187,17 +186,6 @@ 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). - # Only pass unresolved stages so already-resolved stages keep their venvs. - unresolved = [ - s - for s in self.stages - if getattr(s, "pip_specs", None) - and not getattr(s, "_resolved_site_packages_path", None) - ] - if unresolved: - resolve_stage_pip_envs(unresolved) - if executor is None: from nemo_curator.backends.xenna import XennaExecutor diff --git a/nemo_curator/stages/base.py b/nemo_curator/stages/base.py index 9a3ad1e7aa..5761dfeb18 100644 --- a/nemo_curator/stages/base.py +++ b/nemo_curator/stages/base.py @@ -19,7 +19,6 @@ import time from abc import ABC, ABCMeta, abstractmethod from inspect import isabstract -from pathlib import Path # noqa: TC003 from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, final from loguru import logger @@ -86,8 +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 - _resolved_site_packages_path: Path | None = None # set by resolve_stage_pip_envs + runtime_env: ClassVar[dict[str, Any] | None] = None @property @final @@ -117,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. " @@ -262,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. @@ -272,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) @@ -282,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/nemo_curator/utils/stage_pip_env.py b/nemo_curator/utils/stage_pip_env.py deleted file mode 100644 index aab5f1b0c9..0000000000 --- a/nemo_curator/utils/stage_pip_env.py +++ /dev/null @@ -1,139 +0,0 @@ -# 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) - try: - subprocess.run( # noqa: S603 - [uv_exe, "pip", "install", "--python", str(python_path), *specs], - check=True, - capture_output=True, - text=True, - ) - except subprocess.CalledProcessError as e: - logger.error("uv pip install failed for specs %s:\n%s", specs, e.stderr) - raise - - 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 tempfile - - stages_with_pip = [s for s in stages if getattr(s, "pip_specs", None)] - uv_exe = shutil.which("uv") - if not uv_exe: - if stages_with_pip: - names = ", ".join(getattr(s, "name", s.__class__.__name__) for s in stages_with_pip) - msg = ( - "Stages with pip_specs require `uv` on PATH to resolve dependencies, " - f"but `uv` was not found. Affected stages: {names}. " - "Install uv (e.g. pip install uv) or add it to PATH." - ) - raise RuntimeError(msg) - 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 by normalized specs - unique_specs: dict[tuple[str, ...], Path] = {} - for stage in stages: - if getattr(stage, "_resolved_site_packages_path", None): - continue # already resolved; skip - specs = getattr(stage, "pip_specs", None) - if not specs or not isinstance(specs, list): - continue - key = tuple(sorted(s.strip().lower() for s in specs)) - if key not in unique_specs: - subdir = f"env_{len(unique_specs)}" - unique_specs[key] = _create_venv_for_specs(list(key), 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] diff --git a/tests/pipelines/test_per_stage_runtime_env.py b/tests/pipelines/test_per_stage_runtime_env.py index e184dd54ec..0da301e1f8 100644 --- a/tests/pipelines/test_per_stage_runtime_env.py +++ b/tests/pipelines/test_per_stage_runtime_env.py @@ -14,52 +14,29 @@ """Tests for per-stage runtime environment: different Python package versions per stage. -Uses pip_specs + resolve_stage_pip_envs (uv CLI) so each stage runs with its own -packaging version. Requires `uv` on PATH; test is skipped if uv is not available. -See tutorials/per_stage_runtime_env_example.py and docs/design/per-stage-runtime-environment.md. +Each stage declares runtime_env and Ray creates an isolated virtualenv per unique spec set. +Both the RayData and Xenna backends use Ray's native runtime_env mechanism. """ -import shutil # noqa: I001 -import subprocess +from typing import ClassVar import pandas as pd import pytest -from typing import ClassVar - -from nemo_curator.backends.experimental.ray_data import RayDataExecutor -from nemo_curator.backends.xenna import XennaExecutor +from nemo_curator.backends.ray_data import RayDataExecutor 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 _uv_available() -> bool: - uv_exe = shutil.which("uv") - if not uv_exe: - return False - try: - subprocess.run([uv_exe, "--version"], check=True, capture_output=True, text=True) # noqa: S603 - except subprocess.CalledProcessError: - return False - else: - return True - - -@pytest.fixture -def require_uv(): - if not _uv_available(): - pytest.skip("uv not on PATH; per-stage pip_specs tests require uv") - - class VersionStage1(ProcessingStage[DocumentBatch, DocumentBatch]): - """Stage 1: packaging==23.2 via pip_specs (resolver creates venv, PYTHONPATH).""" + """Stage 1: packaging==23.2.""" name = "version_stage_1" resources = Resources(cpus=0.5) batch_size = 1 - pip_specs: ClassVar[list[str]] = ["packaging==23.2"] + runtime_env: ClassVar[dict] = {"pip": ["packaging==23.2"]} def inputs(self) -> tuple[list[str], list[str]]: return ["data"], [] @@ -70,24 +47,24 @@ def outputs(self) -> tuple[list[str], list[str]]: def process(self, task: DocumentBatch) -> DocumentBatch: import packaging - df = task.to_pandas().copy() - df["stage1_packaging_version"] = packaging.__version__ + batch = task.to_pandas().copy() + batch["stage1_packaging_version"] = packaging.__version__ return DocumentBatch( task_id=task.task_id, dataset_name=task.dataset_name, - data=df, + data=batch, _metadata=task._metadata, _stage_perf=task._stage_perf, ) class VersionStage2(ProcessingStage[DocumentBatch, DocumentBatch]): - """Stage 2: packaging==24.0 via pip_specs (resolver creates venv, PYTHONPATH).""" + """Stage 2: packaging==24.0.""" name = "version_stage_2" resources = Resources(cpus=0.5) batch_size = 1 - pip_specs: ClassVar[list[str]] = ["packaging==24.0"] + runtime_env: ClassVar[dict] = {"pip": ["packaging==24.0"]} def inputs(self) -> tuple[list[str], list[str]]: return ["data"], ["stage1_packaging_version"] @@ -98,72 +75,40 @@ def outputs(self) -> tuple[list[str], list[str]]: def process(self, task: DocumentBatch) -> DocumentBatch: import packaging - df = task.to_pandas().copy() - df["stage2_packaging_version"] = packaging.__version__ + batch = task.to_pandas().copy() + batch["stage2_packaging_version"] = packaging.__version__ return DocumentBatch( task_id=task.task_id, dataset_name=task.dataset_name, - data=df, + data=batch, _metadata=task._metadata, _stage_perf=task._stage_perf, ) -@pytest.mark.usefixtures("shared_ray_client", "require_uv") +@pytest.mark.usefixtures("shared_ray_client") def test_per_stage_different_package_versions_ray_data() -> None: - """Run two stages with different packaging versions via pip_specs; assert each sees its own version. - - Pipeline.run() calls resolve_stage_pip_envs() to create venvs with uv; Ray Data adapter - injects PYTHONPATH so workers load the correct site-packages per stage. - """ - initial = DocumentBatch( - task_id="per_stage_version_test", - dataset_name="test", - data=pd.DataFrame({"text": ["hello"]}), - ) - pipeline = Pipeline( - name="per_stage_version_test", - stages=[VersionStage1(), VersionStage2()], - ) - results = pipeline.run( - executor=RayDataExecutor(), - initial_tasks=[initial], + """Run two stages with different packaging versions; assert each sees its own version.""" + initial = DocumentBatch(task_id="test", dataset_name="test", data=pd.DataFrame({"text": ["hello"]})) + results = Pipeline(name="test", stages=[VersionStage1(), VersionStage2()]).run( + executor=RayDataExecutor(), initial_tasks=[initial] ) assert results is not None - assert len(results) == 1 - out = results[0] - df = out.to_pandas() - assert "stage1_packaging_version" in df.columns - assert "stage2_packaging_version" in df.columns - assert df["stage1_packaging_version"].iloc[0] == "23.2", "Stage 1 should see packaging 23.2" - assert df["stage2_packaging_version"].iloc[0] == "24.0", "Stage 2 should see packaging 24.0" + result = results[0].to_pandas() + assert result["stage1_packaging_version"].iloc[0] == "23.2" + assert result["stage2_packaging_version"].iloc[0] == "24.0" -@pytest.mark.usefixtures("shared_ray_client", "require_uv") +@pytest.mark.usefixtures("shared_ray_client") def test_per_stage_different_package_versions_xenna() -> None: - """Run two stages with different packaging versions via pip_specs using XennaExecutor. - - Pipeline.run() calls resolve_stage_pip_envs() to create venvs with uv; Xenna adapter - uses env_info() to set PYTHONPATH so workers load the correct site-packages per stage. - """ - initial = DocumentBatch( - task_id="per_stage_version_test_xenna", - dataset_name="test", - data=pd.DataFrame({"text": ["hello"]}), - ) - pipeline = Pipeline( - name="per_stage_version_test_xenna", - stages=[VersionStage1(), VersionStage2()], - ) - results = pipeline.run( - executor=XennaExecutor(config={"execution_mode": "streaming"}), - initial_tasks=[initial], + """Run two stages with different packaging versions using XennaExecutor.""" + from nemo_curator.backends.xenna import XennaExecutor + + initial = DocumentBatch(task_id="test", dataset_name="test", data=pd.DataFrame({"text": ["hello"]})) + results = Pipeline(name="test", stages=[VersionStage1(), VersionStage2()]).run( + executor=XennaExecutor(config={"execution_mode": "streaming"}), initial_tasks=[initial] ) assert results is not None - assert len(results) == 1 - out = results[0] - df = out.to_pandas() - assert "stage1_packaging_version" in df.columns - assert "stage2_packaging_version" in df.columns - assert df["stage1_packaging_version"].iloc[0] == "23.2", "Stage 1 should see packaging 23.2" - assert df["stage2_packaging_version"].iloc[0] == "24.0", "Stage 2 should see packaging 24.0" + result = results[0].to_pandas() + assert result["stage1_packaging_version"].iloc[0] == "23.2" + assert result["stage2_packaging_version"].iloc[0] == "24.0" From d88d678ae0ea0b53988c30d064ce2809ad8af0a5 Mon Sep 17 00:00:00 2001 From: Ao Tang Date: Mon, 6 Apr 2026 11:27:15 -0700 Subject: [PATCH 16/22] add more tests Signed-off-by: Ao Tang --- nemo_curator/backends/xenna/adapter.py | 6 +- tests/pipelines/test_per_stage_runtime_env.py | 170 +++++++++++++----- 2 files changed, 126 insertions(+), 50 deletions(-) diff --git a/nemo_curator/backends/xenna/adapter.py b/nemo_curator/backends/xenna/adapter.py index 6ae02b0b7e..aaadcdae41 100644 --- a/nemo_curator/backends/xenna/adapter.py +++ b/nemo_curator/backends/xenna/adapter.py @@ -16,9 +16,9 @@ import ray.runtime_env from cosmos_xenna.pipelines import v1 as pipelines_v1 -from cosmos_xenna.pipelines.v1 import NodeInfo as XennaNodeInfo -from cosmos_xenna.pipelines.v1 import Resources as XennaResources -from cosmos_xenna.pipelines.v1 import WorkerMetadata as XennaWorkerMetadata +from cosmos_xenna.pipelines.private.resources import NodeInfo as XennaNodeInfo +from cosmos_xenna.pipelines.private.resources import Resources as XennaResources +from cosmos_xenna.pipelines.private.resources import WorkerMetadata as XennaWorkerMetadata from loguru import logger from nemo_curator.backends.base import BaseStageAdapter, NodeInfo, WorkerMetadata diff --git a/tests/pipelines/test_per_stage_runtime_env.py b/tests/pipelines/test_per_stage_runtime_env.py index 0da301e1f8..e76af4a305 100644 --- a/tests/pipelines/test_per_stage_runtime_env.py +++ b/tests/pipelines/test_per_stage_runtime_env.py @@ -15,40 +15,65 @@ """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. -Both the RayData and Xenna backends use Ray's native runtime_env mechanism. +RayData, Xenna, and RayActorPool backends are all tested via parametrization. """ -from typing import ClassVar +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 -class VersionStage1(ProcessingStage[DocumentBatch, DocumentBatch]): - """Stage 1: packaging==23.2.""" +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 - name = "version_stage_1" + 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 - runtime_env: ClassVar[dict] = {"pip": ["packaging==23.2"]} def inputs(self) -> tuple[list[str], list[str]]: return ["data"], [] def outputs(self) -> tuple[list[str], list[str]]: - return ["data"], ["stage1_packaging_version"] + return ["data"], ["base_packaging_version"] def process(self, task: DocumentBatch) -> DocumentBatch: import packaging batch = task.to_pandas().copy() - batch["stage1_packaging_version"] = packaging.__version__ + batch["base_packaging_version"] = packaging.__version__ return DocumentBatch( task_id=task.task_id, dataset_name=task.dataset_name, @@ -58,57 +83,108 @@ def process(self, task: DocumentBatch) -> DocumentBatch: ) -class VersionStage2(ProcessingStage[DocumentBatch, DocumentBatch]): - """Stage 2: packaging==24.0.""" +class VersionStage1(ProcessingStage[DocumentBatch, DocumentBatch]): + """Stage 1: packaging==23.2.""" - name = "version_stage_2" + name = "version_stage_1" resources = Resources(cpus=0.5) batch_size = 1 - runtime_env: ClassVar[dict] = {"pip": ["packaging==24.0"]} + runtime_env: ClassVar[dict] = {"pip": ["packaging==23.2"]} def inputs(self) -> tuple[list[str], list[str]]: - return ["data"], ["stage1_packaging_version"] + return ["data"], ["base_packaging_version"] def outputs(self) -> tuple[list[str], list[str]]: - return ["data"], ["stage2_packaging_version"] + return ["data"], ["stage1_packaging_version", "stage1_loguru_available"] def process(self, task: DocumentBatch) -> DocumentBatch: - import packaging + return _record_packaging_and_loguru(task, "stage1_packaging_version", "stage1_loguru_available") - batch = task.to_pandas().copy() - batch["stage2_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 VersionStage2(ProcessingStage[DocumentBatch, DocumentBatch]): + """Stage 2: packaging==24.0.""" -@pytest.mark.usefixtures("shared_ray_client") -def test_per_stage_different_package_versions_ray_data() -> None: - """Run two stages with different packaging versions; assert each sees its own version.""" - initial = DocumentBatch(task_id="test", dataset_name="test", data=pd.DataFrame({"text": ["hello"]})) - results = Pipeline(name="test", stages=[VersionStage1(), VersionStage2()]).run( - executor=RayDataExecutor(), initial_tasks=[initial] - ) - assert results is not None - result = results[0].to_pandas() - assert result["stage1_packaging_version"].iloc[0] == "23.2" - assert result["stage2_packaging_version"].iloc[0] == "24.0" + 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"] -@pytest.mark.usefixtures("shared_ray_client") -def test_per_stage_different_package_versions_xenna() -> None: - """Run two stages with different packaging versions using XennaExecutor.""" - from nemo_curator.backends.xenna import XennaExecutor + def outputs(self) -> tuple[list[str], list[str]]: + return ["data"], ["stage2_packaging_version", "stage2_loguru_available"] - initial = DocumentBatch(task_id="test", dataset_name="test", data=pd.DataFrame({"text": ["hello"]})) - results = Pipeline(name="test", stages=[VersionStage1(), VersionStage2()]).run( - executor=XennaExecutor(config={"execution_mode": "streaming"}), initial_tasks=[initial] - ) - assert results is not None - result = results[0].to_pandas() - assert result["stage1_packaging_version"].iloc[0] == "23.2" - assert result["stage2_packaging_version"].iloc[0] == "24.0" + 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, {}), id="ray_data"), + pytest.param((XennaExecutor, {"execution_mode": "streaming"}), id="xenna_streaming"), + pytest.param((XennaExecutor, {"execution_mode": "batch"}), id="xenna_batch"), + pytest.param((RayActorPoolExecutor, {}), id="ray_actor_pool"), + ], + indirect=True, +) +class TestPerStageRuntimeEnv: + """Stages with different runtime_env see different package versions across all backends.""" + + 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 and store results.""" + backend_cls, config = 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"]}), + ) + pipeline = Pipeline( + name="per_stage_runtime_env_test", + stages=[BaseEnvStage(), VersionStage1(), VersionStage2()], + ) + 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" From dcf71430d00eed41c25cfe1a09d5c9f7a83f2508 Mon Sep 17 00:00:00 2001 From: Ao Tang Date: Mon, 6 Apr 2026 11:46:55 -0700 Subject: [PATCH 17/22] remove ensurepip and fix it with --seed in dockerfile Signed-off-by: Ao Tang --- docker/Dockerfile | 2 +- nemo_curator/backends/ray_data/executor.py | 11 ++--------- nemo_curator/backends/xenna/executor.py | 7 ------- 3 files changed, 3 insertions(+), 17 deletions(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index 1a75b2eca3..db0f087a79 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/ray_data/executor.py b/nemo_curator/backends/ray_data/executor.py index ca211f3a4b..dbf9870f9e 100644 --- a/nemo_curator/backends/ray_data/executor.py +++ b/nemo_curator/backends/ray_data/executor.py @@ -62,16 +62,9 @@ def execute(self, stages: list["ProcessingStage"], initial_tasks: list[Task] | N 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. ensurepip.bootstrap() below ensures pip is available - # in the cloned virtualenv (see comment there for details). + # 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: - # Ray clones the current virtualenv when creating per-stage pip virtualenvs. - # The NeMo Curator container's /opt/venv is created by uv, which does not include - # pip as a seed package, so the clone also lacks pip and `python -m pip install` - # fails inside the worker virtualenv. Bootstrap pip first so the clone inherits it. - import ensurepip - - ensurepip.bootstrap(upgrade=True) # Initialize ray and explicitly set NOSET to empty # This ensures if Xenna was used before which was setting NOSET, we end up overriding it. ray.init( diff --git a/nemo_curator/backends/xenna/executor.py b/nemo_curator/backends/xenna/executor.py index ddfc7bfe16..aaf51c8383 100644 --- a/nemo_curator/backends/xenna/executor.py +++ b/nemo_curator/backends/xenna/executor.py @@ -135,13 +135,6 @@ def execute(self, stages: list[ProcessingStage], initial_tasks: list[Task] | Non logger.info(f"Execution mode: {exec_mode.name}") try: - # Ray clones the current virtualenv when creating per-stage pip virtualenvs. - # The NeMo Curator container's /opt/venv is created by uv, which does not include - # pip as a seed package, so the clone also lacks pip and `python -m pip install` - # fails inside the worker virtualenv. Bootstrap pip first so the clone inherits it. - import ensurepip - - ensurepip.bootstrap(upgrade=True) register_loguru_serializer() # Prevent Ray from overriding accelerator env vars when num_gpus=0, letting Xenna manage them instead. ray.init( From 7164dd0dc60ebb020233287daa54eae68d41f9d3 Mon Sep 17 00:00:00 2001 From: Ao Tang Date: Mon, 6 Apr 2026 11:55:10 -0700 Subject: [PATCH 18/22] add extra_env_vars Signed-off-by: Ao Tang --- nemo_curator/backends/xenna/adapter.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/nemo_curator/backends/xenna/adapter.py b/nemo_curator/backends/xenna/adapter.py index aaadcdae41..5b14582b00 100644 --- a/nemo_curator/backends/xenna/adapter.py +++ b/nemo_curator/backends/xenna/adapter.py @@ -30,13 +30,18 @@ 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 two methods Xenna calls: - ``to_ray_runtime_env()`` and ``format()``. + 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 + @property + def extra_env_vars(self) -> dict[str, str]: + """Environment variables from the runtime_env dict, as Xenna's actor pool expects.""" + return self._runtime_env.get("env_vars", {}) + def to_ray_runtime_env(self) -> ray.runtime_env.RuntimeEnv: return ray.runtime_env.RuntimeEnv(**self._runtime_env) From beb11acdee7da489b0cfc4a417b8034e23f6d9c3 Mon Sep 17 00:00:00 2001 From: Ao Tang Date: Mon, 6 Apr 2026 12:24:13 -0700 Subject: [PATCH 19/22] Refactor extra_env_vars to be a settable attribute in CuratorRuntimeEnv Signed-off-by: Ao Tang --- nemo_curator/backends/xenna/adapter.py | 8 ++--- tests/pipelines/test_per_stage_runtime_env.py | 30 ++++++++++++------- 2 files changed, 22 insertions(+), 16 deletions(-) diff --git a/nemo_curator/backends/xenna/adapter.py b/nemo_curator/backends/xenna/adapter.py index 5b14582b00..3f21211537 100644 --- a/nemo_curator/backends/xenna/adapter.py +++ b/nemo_curator/backends/xenna/adapter.py @@ -36,11 +36,9 @@ class CuratorRuntimeEnv: def __init__(self, runtime_env: dict[str, Any]) -> None: self._runtime_env = runtime_env - - @property - def extra_env_vars(self) -> dict[str, str]: - """Environment variables from the runtime_env dict, as Xenna's actor pool expects.""" - return self._runtime_env.get("env_vars", {}) + # 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) diff --git a/tests/pipelines/test_per_stage_runtime_env.py b/tests/pipelines/test_per_stage_runtime_env.py index e76af4a305..e5a2910769 100644 --- a/tests/pipelines/test_per_stage_runtime_env.py +++ b/tests/pipelines/test_per_stage_runtime_env.py @@ -15,7 +15,8 @@ """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. +RayData, Xenna, and RayActorPool backends are all tested via parametrization, for both +pip and uv spec types. """ from typing import Any, ClassVar @@ -84,7 +85,7 @@ def process(self, task: DocumentBatch) -> DocumentBatch: class VersionStage1(ProcessingStage[DocumentBatch, DocumentBatch]): - """Stage 1: packaging==23.2.""" + """Stage 1: packaging==23.2 (default pip; overridable via with_()).""" name = "version_stage_1" resources = Resources(cpus=0.5) @@ -102,7 +103,7 @@ def process(self, task: DocumentBatch) -> DocumentBatch: class VersionStage2(ProcessingStage[DocumentBatch, DocumentBatch]): - """Stage 2: packaging==24.0.""" + """Stage 2: packaging==24.0 (default pip; overridable via with_()).""" name = "version_stage_2" resources = Resources(cpus=0.5) @@ -122,15 +123,19 @@ def process(self, task: DocumentBatch) -> DocumentBatch: @pytest.mark.parametrize( "backend_config", [ - pytest.param((RayDataExecutor, {}), id="ray_data"), - pytest.param((XennaExecutor, {"execution_mode": "streaming"}), id="xenna_streaming"), - pytest.param((XennaExecutor, {"execution_mode": "batch"}), id="xenna_batch"), - pytest.param((RayActorPoolExecutor, {}), id="ray_actor_pool"), + 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.""" + """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 @@ -138,8 +143,8 @@ class TestPerStageRuntimeEnv: @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 and store results.""" - backend_cls, config = request.param + """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 @@ -148,9 +153,12 @@ def backend_config(self, request: pytest.FixtureRequest, shared_ray_cluster: str 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(), VersionStage1(), VersionStage2()], + stages=[BaseEnvStage(), stage1, stage2], ) request.cls.results = pipeline.run(executor=backend_cls(config), initial_tasks=[initial]) From 952e77fe0022b55637593f13d2f9179576de12be Mon Sep 17 00:00:00 2001 From: Ao Tang Date: Mon, 6 Apr 2026 12:55:13 -0700 Subject: [PATCH 20/22] Update CI workflow to include 'uv venv --seed' command for improved test setup. Signed-off-by: Ao Tang --- .github/workflows/cicd-main.yml | 1 + nemo_curator/backends/xenna/adapter.py | 6 +++++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/.github/workflows/cicd-main.yml b/.github/workflows/cicd-main.yml index 2a01996f39..50b56fa2d7 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/nemo_curator/backends/xenna/adapter.py b/nemo_curator/backends/xenna/adapter.py index 3f21211537..0d4e17d3d3 100644 --- a/nemo_curator/backends/xenna/adapter.py +++ b/nemo_curator/backends/xenna/adapter.py @@ -41,7 +41,11 @@ def __init__(self, runtime_env: dict[str, Any]) -> None: 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) + # 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())}" From a882ad43a8270633efcafa1f66353b84e2bc9000 Mon Sep 17 00:00:00 2001 From: Ao Tang Date: Mon, 6 Apr 2026 13:57:34 -0700 Subject: [PATCH 21/22] limit to two test case --- tests/pipelines/test_per_stage_runtime_env.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/tests/pipelines/test_per_stage_runtime_env.py b/tests/pipelines/test_per_stage_runtime_env.py index e5a2910769..acea021fa9 100644 --- a/tests/pipelines/test_per_stage_runtime_env.py +++ b/tests/pipelines/test_per_stage_runtime_env.py @@ -25,7 +25,6 @@ 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 @@ -123,14 +122,8 @@ def process(self, task: DocumentBatch) -> DocumentBatch: @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, ) From 679282657a165d54b3f0b6ea50c5b959bb35ada6 Mon Sep 17 00:00:00 2001 From: Ao Tang Date: Mon, 6 Apr 2026 16:24:04 -0700 Subject: [PATCH 22/22] use _with for unit test Signed-off-by: Ao Tang --- tests/pipelines/test_per_stage_runtime_env.py | 170 +++++++----------- 1 file changed, 65 insertions(+), 105 deletions(-) diff --git a/tests/pipelines/test_per_stage_runtime_env.py b/tests/pipelines/test_per_stage_runtime_env.py index acea021fa9..8bdac05a18 100644 --- a/tests/pipelines/test_per_stage_runtime_env.py +++ b/tests/pipelines/test_per_stage_runtime_env.py @@ -12,14 +12,14 @@ # 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. +"""Tests for per-stage runtime_env support. -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. +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, ClassVar +from typing import Any import pandas as pd import pytest @@ -33,33 +33,15 @@ 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 +class RecordPackagingVersionStage(ProcessingStage[DocumentBatch, DocumentBatch]): + """Records the packaging library version visible to this worker. - try: - from loguru import logger + 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. + """ - 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" + name = "record_packaging_version" resources = Resources(cpus=0.5) batch_size = 1 @@ -67,13 +49,21 @@ def inputs(self) -> tuple[list[str], list[str]]: return ["data"], [] def outputs(self) -> tuple[list[str], list[str]]: - return ["data"], ["base_packaging_version"] + 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["base_packaging_version"] = packaging.__version__ + 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, @@ -83,52 +73,24 @@ def process(self, task: DocumentBatch) -> DocumentBatch: ) -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") +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, {}, "uv"), id="ray_data-uv"), - pytest.param((XennaExecutor, {"execution_mode": "streaming"}, "pip"), id="xenna_streaming-pip"), + 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 across all backends and spec types.""" + """Stages with different runtime_env see different package versions.""" backend_cls: type[BaseExecutor] | None = None config: dict[str, Any] | None = None @@ -136,56 +98,54 @@ class TestPerStageRuntimeEnv: @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 + """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 - initial = DocumentBatch( - task_id="test", - dataset_name="test", - data=pd.DataFrame({"text": ["hello"]}), + base_stage = RecordPackagingVersionStage().with_(name="base_env") + stage_v232 = RecordPackagingVersionStage().with_( + name="pinned_v232", + runtime_env={"pip": ["packaging==23.2"]}, ) - # 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], + stage_v240 = RecordPackagingVersionStage().with_( + name="pinned_v240", + runtime_env={"uv": ["packaging==24.0"]}, ) - 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" + 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_stage2_packaging_version(self): + def test_output_count(self): assert self.results is not None - result = self.results[0].to_pandas() - assert result["stage2_packaging_version"].iloc[0] == "24.0" + assert len(self.results) == 1 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 + """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 must all be distinct.""" - assert self.results is not None - result = self.results[0].to_pandas() + """Base env, 23.2, and 24.0 should all be distinct.""" + df = self.results[0].to_pandas() versions = { - result["base_packaging_version"].iloc[0], - result["stage1_packaging_version"].iloc[0], - result["stage2_packaging_version"].iloc[0], + 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 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" + """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])