NVIDIA / NVIDIA/NeMo-Retriever
[FEA]: Add support for creating lambda stages
Open
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
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
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