NVIDIA / NVIDIA/NeMo-Retriever
[FEA]: Add Dynamic Lambda Stage Injection + YAML Pipeline Construction
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 3k
- Forks
- 349
- Avg merge
- 1d 23h
- Merged PRs (30d)
- 116
Description
Is this a new feature, an improvement, or a change to existing functionality?
New Feature
How would you describe the priority of this feature request
Significant improvement
Please provide a clear description of problem this feature solves
Extend the RayPipeline class to support dynamic construction, injection, and modification of pipeline stages dynamically, either directly or using using a yaml-based configuration. This includes support for injecting lambda-based stages, referencing module paths, and modifying the DAG topology via programmatic or configuration-driven control.
Implementation Scope
add_stage(stage: Stage | Callable | str, stage_id: str, config: dict = None, ...)
Accepts a stage object, a callable, or a module path string.
If a config block is supplied:
If the stage is a class, config is passed as keyword arguments during instantiation.
If the stage is a callable, the stage is wrapped via wrap_callable_as_stage(callable, config=config) and the config applied during init or through a configure() method.
splice(before_stage: str, after_stage: str, new_stage: Stage | Callable | str, new_stage_id: str, config: dict = None)
Inserts a stage between before_stage and after_stage.
Instantiates the stage using the config if provided.
Rewires edges appropriately.
remove_stage(stage_id: str)
Removes the stage from the internal DAG and disconnects its edges.
Optionally reconnects predecessors to successors (configurable).
from_yaml(path: str)
Parses a YAML pipeline configuration.
Instantiates and connects all stages, applying configuration blocks as necessary.
Describe the feature, and optionally a solution or implementation and any alternatives
def resolve_stage(
module_path: str,
config: dict
) -> RayActorStage:
module_str, attr = module_path.split(":")
mod = importlib.import_module(module_str)
obj = getattr(mod, attr)
if isinstance(obj, type) and issubclass(obj, RayActorStage):
return obj(**config)
elif callable(obj):
return wrap_callable_as_stage(obj, config=config)
else:
raise TypeError(f"Unsupported stage definition at: {module_path}")
def add_stage(self, stage_input: Union[RayActorStage, Callable, str], stage_id: str, config: dict = None):
config = config or {}
if isinstance(stage_input, str):
stage = resolve_stage(stage_input, config)
elif callable(stage_input):
stage = wrap_callable_as_stage(stage_input, config=config)
else:
stage = stage_input
self.stages[stage_id] = stage
def splice(self, before: str, after: str, new_stage_id: str, stage_input: Union[RayActorStage, Callable, str], config: dict = None):
self.remove_edge(before, after)
self.add_stage(stage_input, new_stage_id, config=config)
self.add_edge(before, new_stage_id)
self.add_edge(new_stage_id, after)
def remove_stage(self, stage_id: str):
if stage_id not in self.stages:
raise ValueError(f"Stage {stage_id} does not exist.")
self.remove_all_edges(stage_id)
del self.stages[stage_id]
def from_yaml(self, yaml_path: str):
with open(yaml_path, "r") as f:
raw_config = yaml.safe_load(f)
stages = [StageConfig(**entry) for entry in raw_config["stages"]]
for stage_cfg in stages:
self.add_stage(stage_cfg.module, stage_cfg.stage_id, config=stage_cfg.config)
for inp in stage_cfg.inputs:
self.add_edge(inp, stage_cfg.stage_id)
Additional context
Pesudo config file
stages:
- stage_id: load_data
module: mylib.stages.load_data:LoadDataStage
min_replicas: 1
max_replicas: 1
description: Load raw data from source
config:
source_path: /data/input.csv
delimiter: ","
- stage_id: normalize
module: mylib.transforms.normalize_data:normalize_fn
min_replicas: 2
max_replicas: 4
description: Normalize numeric values
inputs: [load_data]
config:
scale_method: zscore
drop_nulls: true
- stage_id: write_output
module: mylib.sinks.output:SinkStage
min_replicas: 1
max_replicas: 1
description: Write output to disk
inputs: [normalize]
config:
output_path: /data/output.json
format: json
edges:
- [load_data, normalize]
- [normalize, write_output]
flowchart LR
A[load_data<br/>LoadDataStage<br/>source_path=/data/input.csv] --> B[normalize<br/>normalize_fn<br/>scale_method=zscore]
B --> C[write_output<br/>SinkStage<br/>output_path=/data/output.json]
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
No files or tests are named. Start by locating RayPipeline, RayActorStage, wrap_callable_as_stage, and the existing DAG edge methods; then determine how stage registration and configuration currently work. Done means supporting add_stage, splice, remove_stage, and from_yaml with callable, module-path, configuration, and topology behavior covered by tests.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, yaml
- Domain
- backend, data-engineering
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 25/100