NVIDIA / NVIDIA/NeMo-Retriever

[FEA]: Add support for creating lambda stages

Open
#766 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

Introduce a utility function that wraps a user-defined Callable[[ControlMessage], ControlMessage] as a RayActorStage. This enables ad hoc logic injection into the DAG without boilerplate subclassing.

Describe the feature, and optionally a solution or implementation and any alternatives
def wrap_callable_as_stage(fn: Callable[[ControlMessage], ControlMessage]) -> RayActorStage:
    class CallableWrapperStage(RayActorStage):
        def on_data(self, control_message: ControlMessage) -> ControlMessage:
            try:
                return fn(control_message)
            except Exception as e:
                self._logger.exception(f"Callable raised an error: {e}")
                return control_message  # or optionally None, based on design
    return CallableWrapperStage()
def my_transform(cm: ControlMessage) -> ControlMessage:
    payload = cm.payload()
    updated_payload = do_work(payload)
   cm.payload(updated_payload)
    return cm

transform_stage = wrap_callable_as_stage(my_transform)
Additional context
flowchart LR
    A[Input ControlMessage] --> B[wrap_callable_as_stage]
    B --> C[User Callable<br/>cm -> cm]
    C --> D[Transformed ControlMessage]

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

Start by locating RayActorStage, ControlMessage, and the existing on_data entry point. Trace how stages are instantiated and how callable results and exceptions are handled, then define the wrapper behavior and verify that a user Callable can produce a transformed ControlMessage without subclassing.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
distributed-systems
Issue type
Feature
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.