[QUESTION] How to Cancel the Running of a Nested Workflow?

[jinhu2659]

Context:
I have a parent workflow (SupWorkflow) that initializes a sub-workflow (SubWorkflow) in its __init__ method. The sub-workflow is executed as part of a step in the parent workflow. When I call cancel_run() on the handle of the parent workflow, the sub-workflow continues running, which is unexpected. I would like guidance on how to properly cancel both the parent and nested workflows to ensure the sub-workflow stops execution as well.

Steps to Reproduce:

1. Define a parent workflow (SupWorkflow) that initializes a sub-workflow (SubWorkflow) in its __init__ method.
2. Run the sub-workflow as part of a step in the parent workflow.
3. Attempt to cancel the parent workflow using cancel_run() on its handle.
4. Observe that the sub-workflow continues running despite the cancellation.

Code Example:

from llama_index.core.workflow import Workflow, step, StartEvent, StopEvent
import asyncio

class SubWorkflow(Workflow):
    @step
    async def sub_step(self, ev: StartEvent) -> StopEvent:
        print("SubWorkflow: Starting long-running task")
        for i in range(10):
            print(f"SubWorkflow: Processing {i}")
            await asyncio.sleep(1)  # Simulate long-running task
        print("SubWorkflow: Completed")
        return StopEvent(result="SubWorkflow done")

class SupWorkflow(Workflow):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.sub_workflow = SubWorkflow()

    @step
    async def parent_step(self, ev: StartEvent) -> StopEvent:
        print("SupWorkflow: Starting parent step")
        result = await self.sub_workflow.run()  # Run sub-workflow
        print(f"SupWorkflow: SubWorkflow result: {result}")
        return StopEvent(result="Parent workflow done")

async def main():
    parent_workflow = SupWorkflow(timeout=20)
    handle = parent_workflow.run()
    await asyncio.sleep(3)  # Let it run for a bit
    print("Attempting to cancel parent workflow")
    await handle.cancel_run()  # Cancel the parent workflow
    print("Parent workflow cancellation requested")
    await asyncio.sleep(5)  # Wait to observe behavior
    print("Main: Done")

if __name__ == "__main__":
    asyncio.run(main())

Expected Behavior:
When cancel_run() is called on the parent workflow's handle, both the parent workflow (SupWorkflow) and the nested sub-workflow (SubWorkflow) should stop execution immediately.

Actual Behavior:
The parent workflow stops, but the sub-workflow continues running, as evidenced by the continued printing of SubWorkflow: Processing ... messages after cancellation.

Question:

• How can I ensure that calling cancel_run() on the parent workflow also cancels the nested sub-workflow?
• Is there a recommended approach to propagate cancellation to nested workflows in LlamaIndex?
• If this is a limitation, are there workarounds to achieve proper cancellation of nested workflows?

Environment:

• Llama-index-workflows version: [1.3.0]
• Python version: [3.11]
• Operating System: [Windows10]

Additional Notes:

• I suspect the issue may be related to the sub-workflow not being aware of the parent workflow's cancellation state.
• Any code examples or best practices for handling nested workflow cancellation would be greatly appreciated.

---

[dosubot[bot]]

Hi @jinhu2659! I'm Dosu and I’m helping the llama_index team.

LlamaIndex does not automatically propagate cancellation from a parent workflow to nested sub-workflows. When you call cancel_run() on the parent workflow's handle, only the parent workflow stops; any sub-workflows started within its steps continue running independently. This is a known limitation and matches common async Python patterns, where cancellation must be explicitly managed for each running task or coroutine [source].

Recommended approach:
To ensure both parent and sub-workflows are cancelled, you need to track the sub-workflow's handler and manually call cancel_run() on it when the parent is cancelled. Instead of directly awaiting self.sub_workflow.run(), get the handler and manage its lifecycle:

@step
async def parent_step(self, ev: StartEvent) -> StopEvent:
    print("SupWorkflow: Starting parent step")
    sub_handle = self.sub_workflow.run()
    # Optionally, stream events from the sub-workflow here
    result = await sub_handle
    print(f"SupWorkflow: SubWorkflow result: {result}")
    return StopEvent(result="Parent workflow done")

Then, in your cancellation logic, make sure to cancel both handles:

parent_handle = parent_workflow.run()
sub_handle = parent_workflow.sub_workflow.run()
await asyncio.sleep(3)
print("Attempting to cancel parent workflow and sub-workflow")
await parent_handle.cancel_run()
await sub_handle.cancel_run()

Alternatively, you can create an async task that monitors the parent workflow's cancellation and propagates it to the sub-workflow [source].

Best practices:

• Always track the handles of any sub-workflows you start.
• Manually propagate cancellation by calling cancel_run() on each handle.
• If you use event streaming (e.g., handler.stream_events()), you can monitor and control sub-workflow execution more granularly [source].

Note:
Recent versions of LlamaIndex have removed some deprecated sub-workflow features and event streaming methods, so always check the documentation and changelogs for the version you are using [source].

For more details and examples, see the discussion here and this issue.

To reply, just mention @dosu.

---

^{How did I do? Good | Irrelevant | Incorrect | Verbose | Hallucination | Report 🐛 | Other}    

[omni-front]

LlamaIndex doesn't propagate cancellation to nested workflows automatically, but you can implement this yourself. Here are two approaches:

Approach 1: Track and cancel sub-workflow handles

Store the sub-workflow handle and cancel it explicitly when the parent is cancelled:

from llama_index.core.workflow import Workflow, step, StartEvent, StopEvent
import asyncio

class SubWorkflow(Workflow):
    @step
    async def do_work(self, ev: StartEvent) -> StopEvent:
        for i in range(100):
            await asyncio.sleep(1)  # simulate work
            print(f"Sub-workflow step {i}")
        return StopEvent(result="done")

class SupWorkflow(Workflow):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.sub_workflow = SubWorkflow()
        self._sub_handle = None

    @step
    async def run_sub(self, ev: StartEvent) -> StopEvent:
        self._sub_handle = self.sub_workflow.run()
        result = await self._sub_handle
        return StopEvent(result=result)

    async def cancel(self):
        """Cancel both parent and sub-workflow."""
        if self._sub_handle:
            self._sub_handle.cancel()
        # cancel parent
        return await super().cancel()

Approach 2: Use a shared cancellation event

More robust for deeply nested workflows:

import asyncio
from contextvars import ContextVar

cancel_event: ContextVar[asyncio.Event] = ContextVar("cancel_event")

class SubWorkflow(Workflow):
    @step
    async def do_work(self, ev: StartEvent) -> StopEvent:
        cancel = cancel_event.get()
        for i in range(100):
            if cancel.is_set():
                print("Sub-workflow cancelled")
                return StopEvent(result="cancelled")
            await asyncio.sleep(1)
        return StopEvent(result="done")

class SupWorkflow(Workflow):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.sub_workflow = SubWorkflow()
        self._cancel_event = asyncio.Event()

    @step
    async def run_sub(self, ev: StartEvent) -> StopEvent:
        token = cancel_event.set(self._cancel_event)
        try:
            handler = self.sub_workflow.run()
            result = await handler
            return StopEvent(result=result)
        finally:
            cancel_event.reset(token)

    async def cancel(self):
        self._cancel_event.set()  # signals all sub-workflows

Approach 2 scales better when you have multiple levels of nesting. The shared asyncio.Event propagates cancellation to any depth without needing to track individual handles.

Hope this helps!