[Feat] Per Stage Runtime Env

[suiyoubi]

Overview

Enables per-stage runtime environments in NeMo Curator pipelines. Each stage can declare a distinct set of Python packages via a `runtime_env` class variable, and Ray will create an isolated virtualenv for that stage's workers — allowing incompatible library versions to coexist in the same pipeline.

Changes

Stage API

• Added `runtime_env: ClassVar[dict | None] = None` to `ProcessingStage`
• Added `runtime_env` support to `with_()` for per-instance overrides

class MyStage(ProcessingStage[DocumentBatch, DocumentBatch]):
    runtime_env = {"pip": ["transformers==4.40.0"]}  # or {"uv": [...]}
    ...

# or override at instantiation time:
stage = MyStage().with_(runtime_env={"pip": ["transformers==4.45.0"]})

Backend support (all three backends)

Backend	Mechanism
RayData	runtime_env forwarded via ray_remote_args to map_batches
Xenna	CuratorRuntimeEnv duck-type bridges Xenna's env_info to Ray's full runtime_env dict
RayActorPool	runtime_env passed to actor .options() at pool creation

Infrastructure

• docker/Dockerfile: uv venv --seed so pip is present in the container venv and Ray's cloned worker venvs inherit it
• cicd-main.yml: uv venv --seed before uv sync so pip is available in CI's uv-managed .venv
• .rayignore: excludes uv.lock to prevent uv from enforcing locked transitive dep versions that would override the stage's declared versions

How it works

Ray's native runtime_env creates an isolated virtualenv per unique spec set under /tmp/ray/session_latest/runtime_resources/pip/<hash>/virtualenv, cached for the lifetime of the Ray session. No driver-side venv creation or PYTHONPATH manipulation is needed.

Tests

tests/pipelines/test_per_stage_runtime_env.py — parametrized over RayData (uv) and Xenna streaming (pip), runs a 3-stage pipeline with:

• base stage (no runtime_env)
• packaging==23.2 stage
• packaging==24.0 stage

Verifies each stage sees its declared version and that base-env packages (loguru) remain importable in isolated envs.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

Adds per-stage runtime_env support across all three backends (RayData, RayActorPool, Xenna), letting each stage declare an isolated Ray virtualenv via a runtime_env: ClassVar[dict | None] attribute. Infrastructure changes (uv venv --seed in Dockerfile and CI) ensure pip is available inside cloned worker venvs.

Confidence Score: 5/5

Safe to merge; all remaining findings are P2 style/consistency issues.

Previous rounds of review addressed the substantive concerns. The two remaining gaps (RAFT/shuffle paths not propagating runtime_env, and overwrite vs. merge in the Ray Data adapter) are unlikely to affect real users in this iteration since RAFT/shuffle stages don't need pip isolation and mixed use of both mechanisms is not a documented pattern. The test fragility is low-risk given modern packaging versions.

nemo_curator/backends/ray_data/adapter.py (overwrite vs. merge of runtime_env) and nemo_curator/backends/experimental/ray_actor_pool/executor.py (RAFT/shuffle paths).

Important Files Changed

Filename	Overview
nemo_curator/stages/base.py	Adds `runtime_env: ClassVar[dict
nemo_curator/backends/ray_data/adapter.py	Forwards runtime_env into ray_remote_args for map_batches; unconditionally overwrites rather than merging with any runtime_env already present in ray_stage_spec() RAY_REMOTE_ARGS.
nemo_curator/backends/ray_data/executor.py	No functional changes; comment added explaining the uv-venv/pip interaction with Ray's runtime_env.
nemo_curator/backends/xenna/adapter.py	New CuratorRuntimeEnv duck-type bridges Ray's full runtime_env dict to Xenna, and correctly merges Xenna-injected extra_env_vars back before handing off to Ray.
nemo_curator/backends/experimental/ray_actor_pool/executor.py	_create_actor_pool propagates runtime_env but _create_raft_actor_pool and _create_rapidsmpf_actors do not, leaving an inconsistency for those specialised actor paths.
nemo_curator/backends/experimental/utils.py	No functional changes to review.
tests/pipelines/test_per_stage_runtime_env.py	Good end-to-end coverage; test_all_three_versions_differ could fail if the base env already has one of the pinned packaging versions.
docker/Dockerfile	uv venv --seed ensures pip is present in the container venv so Ray's cloned worker venvs can install pip packages.
.github/workflows/cicd-main.yml	uv venv --seed added before uv sync so pip is available in CI's managed venv.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[ProcessingStage\nruntime_env: ClassVar] --> B{Which backend?}
    B -->|RayData| C[RayDataStageAdapter.process_dataset]
    C --> D[Build ray_remote_args\nfrom ray_stage_spec]
    D --> E{stage.runtime_env set?}
    E -->|Yes| F[ray_remote_args\n\x5b'runtime_env'\x5d = stage.runtime_env]
    E -->|No| G[Pass as-is]
    F --> H[dataset.map_batches\n**concurrency_kwargs]
    G --> H
    B -->|RayActorPool| I[_create_actor_pool]
    I --> J{stage.runtime_env set?}
    J -->|Yes| K[actor.options\nruntime_env=stage.runtime_env]
    J -->|No| L[actor.options\nno runtime_env]
    K --> M[ActorPool]
    L --> M
    B -->|Xenna| N[XennaStageAdapter.env_info]
    N --> O{stage.runtime_env set?}
    O -->|Yes| P[CuratorRuntimeEnv\nwraps full Ray dict]
    O -->|No| Q[return None]
    P --> R[to_ray_runtime_env\nmerges extra_env_vars]
    R --> S[Ray actor with\nisolated venv]

Loading

_{Reviews (19): Last reviewed commit: "Merge branch 'main' into aot/runtime_env" | Re-trigger Greptile}

[suiyoubi]

/claude review

[suiyoubi]

/ok to test beb11ac

[greptile-apps]

to_ray_runtime_env() silently drops Xenna's extra_env_vars mutations

__init__ copies runtime_env["env_vars"] into self.extra_env_vars as a separate dict. The comment confirms Xenna both reads and writes that attribute at actor startup (e.g., injecting per-actor CUDA device IDs). But to_ray_runtime_env() reconstructs RuntimeEnv from the original self._runtime_env snapshot and never consults self.extra_env_vars, so any env vars Xenna injects are silently discarded before Ray ever receives the runtime environment.

Suggested change

	def __init__(self, runtime_env: dict[str, Any]) -> None:
	self._runtime_env = runtime_env
	# Xenna's actor pool both reads and writes extra_env_vars on the runtime env object,
	# so this must be a plain settable attribute, not a read-only property.
	self.extra_env_vars: dict[str, str] = dict(runtime_env.get("env_vars", {}))
	def to_ray_runtime_env(self) -> ray.runtime_env.RuntimeEnv:
	return ray.runtime_env.RuntimeEnv(**self._runtime_env)
	def to_ray_runtime_env(self) -> ray.runtime_env.RuntimeEnv:
	merged = {**self._runtime_env, "env_vars": self.extra_env_vars}
	return ray.runtime_env.RuntimeEnv(**merged)

[suiyoubi]

/ok to test 952e77f

[praateekmahajan]

Hmmm why do we need this here? @thomasdhc shouldn't the ci tests run on the docker image itself?

[suiyoubi]

I don't think so... unit test are not running on the docker image (without this I got pip not found error)

[praateekmahajan]

Why do we need these two stages too, and not just do BaseEnvStage().with_(...)?

[suiyoubi]

So I want to have both stages need to write to different columns (stage1_packaging_version vs stage2_packaging_version) to preserve both versions in the final result for comparison.

Unless we drop the "both versions in one result" design and run two separate pipelines instead:

  stage = BaseEnvStage()  # one reusable class

  result1 = Pipeline(..., stages=[stage.with_(runtime_env={spec_type: ["packaging==23.2"]})]).run(...)
  result2 = Pipeline(..., stages=[stage.with_(runtime_env={spec_type: ["packaging==24.0"]})]).run(...)

  assert result1[0].to_pandas()["base_packaging_version"].iloc[0] == "23.2"
  assert result2[0].to_pandas()["base_packaging_version"].iloc[0] == "24.0"

Is this what you mean ?

[suiyoubi]

/ok to test a882ad4

[suiyoubi]

/ok to test 6792826

[praateekmahajan]

🙏 thank you for your work here

[suiyoubi]

/ok to test 00dd3b9