Streaming workflow with nested workflow(s)

[anirbanbasu]

Hi,

I am trying to stream events from a workflow, as documented in https://docs.llamaindex.ai/en/stable/understanding/workflows/stream/. And, it works!

However, when I have one or more nested workflow(s) invoked from one of the steps in the main workflow (as shown in https://docs.llamaindex.ai/en/stable/understanding/workflows/nested/ but using asyncio.create_task(nested_workflow.run(parameters)) and eventually await task), I get the error:

llama_index.core.workflow.errors.WorkflowRuntimeError: This workflow has multiple concurrent runs in progress and cannot stream events. To be able to stream events, make sure you call `run()` on this workflow only once.

I am trying to stream events from the nested workflow upwards through the main workflow as follows, before calling response = await task:

async for nested_ev in react_workflow.stream_events():
    ctx.write_event_to_stream(Event(msg=nested_ev.msg))

I think I know what is going wrong because the step in the main workflow that calls the nested workflow as an async task can be triggered multiple times with different data while the nested workflow tasks have not finished. That is why the error is about calling run() on the nested workflow once.

How can I pass different data to different calls to the nested workflow if I can call run() on it only once?

Or, is there a far better way of doing this?

[dosubot[bot]]

To handle different data for multiple calls to the nested workflow while ensuring that each invocation of the nested workflow is awaited before starting a new one, you can use the add_workflows method to add nested workflows and then call run() on each nested workflow separately with the required data. This approach ensures that each workflow instance is only run once, thus avoiding the WorkflowRuntimeError.

Here is an example of how you can achieve this:

from llama_index.core.workflow.workflow import Workflow, StartEvent, StopEvent, step, Context, Event

# Define your main workflow
class MainWorkflow(Workflow):
    @step
    async def start(self, ctx: Context, ev: StartEvent, reflection_workflow: Workflow) -> Event:
        print("Need to run reflection")
        res = await reflection_workflow.run(query=ev.query)
        return Event(query=res)

# Define your nested workflow
class NestedWorkflow(Workflow):
    @step
    async def sub_start(self, ctx: Context, ev: StartEvent) -> StopEvent:
        print("Doing custom reflection")
        return StopEvent(result="Improved query")

# Create instances of the workflows
main_workflow = MainWorkflow()
nested_workflow1 = NestedWorkflow()
nested_workflow2 = NestedWorkflow()

# Add the nested workflows to the main workflow
main_workflow.add_workflows(reflection_workflow=nested_workflow1)

# Run the nested workflows with different data
await nested_workflow1.run(query="data_for_first_call")
await nested_workflow2.run(query="data_for_second_call")

# Stream events from the main workflow
async for event in main_workflow.stream_events():
    print(event)

In this example, nested_workflow1 and nested_workflow2 are run with different data, and the stream_events method is used to stream events from the main workflow. This approach ensures that each nested workflow is only run once, thus avoiding the WorkflowRuntimeError. Additionally, the add_workflows method is used to supply the nested workflows to the main workflow, allowing for flexible and modular workflow design.

Additionally, you can refer to the test_streaming.py file in the LlamaIndex core tests for more examples on how to handle streaming events and avoid WorkflowRuntimeError ^[1][2][3].

To continue talking to Dosu, mention @dosu.

[anirbanbasu]

What am I doing wrong here? The events written to the context in this nested workflow are not available in the main workflow.

Using the nested workflow as you have suggested

class NestedWorkflow(Workflow):
    @step
    async def sub_start(self, ctx: Context, ev: StartEvent) -> StopEvent:
        ctx.write_event_to_stream(Event(msg="nested_status"))
        # Additional logic for the sub_start step
        return StopEvent(result="Sub workflow step complete.")

and the call to the main workflow run

task = asyncio.create_task(
      main_workflow.run(parameters_1)
)

async for ev in main_workflow.stream_events():
    print(ev.msg)

result = await task

I only get the print(ev.msg) to output messages that are written to the context ctx.write_event_to_stream(Event(msg="main_status")) by a step in the main workflow, but not the one from the nested workflow.

Of course, I have added the nested workflow to the main workflow as follows.

main_workflow.add_workflows(
        nested_workflow=NestedWorkflow(parameters_2)
)

In some step in the main workflow, I am running the nested workflow similar to what you suggested:

await nested_workflow.run(some_parameters)

[dosubot[bot]]

The issue arises because each workflow, including nested workflows, has its own Context instance. Events written to the context of a nested workflow are not automatically available in the main workflow's context. To propagate events from the nested workflow to the main workflow, you need to explicitly stream the events from the nested workflow to the main workflow.

Here's how you can modify your code to ensure that events from the nested workflow are streamed to the main workflow:

1. Stream Events from Nested Workflow to Main Workflow: Use async for to iterate over events from the nested workflow and write each event to the main workflow's context.

2. Ensure Sequential Execution: Make sure that the nested workflow is awaited before starting a new one to avoid concurrent runs.

Here is an updated example:

Nested Workflow

class NestedWorkflow(Workflow):
    @step
    async def sub_start(self, ctx: Context, ev: StartEvent) -> StopEvent:
        ctx.write_event_to_stream(Event(msg="nested_status"))
        # Additional logic for the sub_start step
        return StopEvent(result="Sub workflow step complete.")

Main Workflow

class MainWorkflow(Workflow):
    @step
    async def start(self, ctx: Context, ev: StartEvent, nested_workflow: Workflow) -> Event:
        print("Need to run nested workflow")
        task = asyncio.create_task(nested_workflow.run(some_parameters))

        # Stream events from the nested workflow to the main workflow
        async for nested_ev in nested_workflow.stream_events():
            ctx.write_event_to_stream(Event(msg=nested_ev.msg))

        response = await task
        return Event(query=response)

Running the Workflows

main_workflow = MainWorkflow()
nested_workflow = NestedWorkflow()

main_workflow.add_workflows(nested_workflow=nested_workflow)

task = asyncio.create_task(main_workflow.run(parameters_1))

async for ev in main_workflow.stream_events():
    print(ev.msg)

result = await task

In this setup:

• The start step in the MainWorkflow runs the NestedWorkflow and streams its events to the main workflow's context.
• The async for nested_ev in nested_workflow.stream_events() loop ensures that events from the nested workflow are captured and written to the main workflow's event stream.

This approach ensures that events from the nested workflow are available in the main workflow's event stream, allowing you to capture and process them as needed ^[1][2][3].

To continue talking to Dosu, mention @dosu.

[anirbanbasu]

@dosu, this is back to where I started -- this is exactly what I did. It causes the llama_index.core.workflow.errors.WorkflowRuntimeError: This workflow has multiple concurrent runs in progress and cannot stream events. To be able to stream events, make sure you call run() on this workflow only once. because the step in the main workflow that calls the nested workflow can be triggered in parallel, i.e., because multiple events that it consumes are emitted by another step in the main workflow.

The result is that nested_workflow.run() is already running as an async task when another call to it is made. I cannot predict how many of such parallel triggers will be there, so I cannot create a predetermined number of nested workflow instances.

What is the way around it?

[dosubot[bot]]

To handle multiple parallel triggers in the main workflow that call the nested workflow without causing the WorkflowRuntimeError due to concurrent runs, you can dynamically create and manage multiple instances of the nested workflow. This approach ensures that each nested workflow instance is only run once, thus avoiding the WorkflowRuntimeError.

Here is an example of how you can achieve this:

Main Workflow

import asyncio
from llama_index.core.workflow import (
    StartEvent,
    StopEvent,
    Workflow,
    step,
    Event,
    Context,
)

class ProcessEvent(Event):
    data: str

class ResultEvent(Event):
    result: str

class MainWorkflow(Workflow):
    @step
    async def start(self, ctx: Context, ev: StartEvent) -> ProcessEvent:
        data_list = ["A", "B", "C"]  # This can be dynamically generated
        await ctx.set("num_to_collect", len(data_list))
        for item in data_list:
            ctx.send_event(ProcessEvent(data=item))
        return None

    @step(num_workers=3)  # Adjust num_workers based on your resource limits
    async def process_data(self, ctx: Context, ev: ProcessEvent) -> ResultEvent:
        nested_workflow = NestedWorkflow()
        task = asyncio.create_task(nested_workflow.run(data=ev.data))

        # Stream events from the nested workflow to the main workflow
        async for nested_ev in nested_workflow.stream_events():
            ctx.write_event_to_stream(Event(msg=nested_ev.msg))

        response = await task
        return ResultEvent(result=response)

    @step
    async def combine_results(self, ctx: Context, ev: ResultEvent) -> StopEvent | None:
        num_to_collect = await ctx.get("num_to_collect")
        results = ctx.collect_events(ev, [ResultEvent] * num_to_collect)
        if results is None:
            return None

        combined_result = ", ".join([event.result for event in results])
        return StopEvent(result=combined_result)

Nested Workflow

class NestedWorkflow(Workflow):
    @step
    async def sub_start(self, ctx: Context, ev: StartEvent) -> StopEvent:
        ctx.write_event_to_stream(Event(msg="nested_status"))
        # Additional logic for the sub_start step
        return StopEvent(result="Sub workflow step complete.")

Running the Workflows

import time

main_workflow = MainWorkflow()

print("Start the main workflow with parallel processing")
start_time = time.time()
result = await main_workflow.run()
end_time = time.time()
print(f"Workflow result: {result}")
print(f"Time taken: {end_time - start_time} seconds")

In this setup:

• The process_data step in the MainWorkflow dynamically creates a new instance of NestedWorkflow for each ProcessEvent.
• The async for nested_ev in nested_workflow.stream_events() loop ensures that events from each nested workflow instance are captured and written to the main workflow's event stream.
• The num_workers parameter in the @step decorator controls the number of steps executed simultaneously, enabling efficient parallel processing.

This approach ensures that each nested workflow instance is only run once, thus avoiding the WorkflowRuntimeError due to concurrent runs ^[1][2].

To continue talking to Dosu, mention @dosu.

[anirbanbasu]

@dosu, thanks a lot. That works! So, there's absolutely no need of main_workflow.add_workflow(nested_workflow) as the nested_workflow is an instance created dynamically in the step of the main workflow that calls the nested workflow.

[dosubot[bot]]

Exactly, there's no need to use main_workflow.add_workflow(nested_workflow) when you are dynamically creating instances of the NestedWorkflow within the step of the main workflow. This approach ensures that each nested workflow instance is managed independently, avoiding the WorkflowRuntimeError and allowing for parallel processing. Glad to hear it works for you! If you have any more questions or need further assistance, feel free to ask.

To continue talking to Dosu, mention @dosu.