NVIDIA / NVIDIA/NeMo-Retriever

[FEA]: Add Dynamic Lambda Stage Injection + YAML Pipeline Construction

Open
#767 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

feature request
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.