| 68 | |
| 69 | @override |
| 70 | async def _run_impl( |
| 71 | self, |
| 72 | *, |
| 73 | ctx: Context, |
| 74 | node_input: Any, |
| 75 | ) -> AsyncGenerator[Any, None]: |
| 76 | if not isinstance(node_input, list): |
| 77 | # Wrap the single input in a list to allow processing. |
| 78 | # This handles cases where the input is a single item. |
| 79 | node_input = [node_input] |
| 80 | |
| 81 | if not node_input: |
| 82 | yield [] |
| 83 | return |
| 84 | |
| 85 | results = [None] * len(node_input) |
| 86 | pending_tasks: set[asyncio.Task[Any]] = set() |
| 87 | input_index = 0 |
| 88 | |
| 89 | while input_index < len(node_input) or pending_tasks: |
| 90 | # Check for any inputs waiting to be processed. |
| 91 | while input_index < len(node_input) and ( |
| 92 | self.max_concurrency is None |
| 93 | or len(pending_tasks) < self.max_concurrency |
| 94 | ): |
| 95 | item = node_input[input_index] |
| 96 | task = asyncio.create_task( |
| 97 | ctx.run_node( |
| 98 | self._node, |
| 99 | node_input=item, |
| 100 | use_sub_branch=True, |
| 101 | ) |
| 102 | ) |
| 103 | # Store index on task so we can place result correctly when done |
| 104 | setattr(task, '_worker_index', input_index) |
| 105 | pending_tasks.add(task) |
| 106 | input_index += 1 |
| 107 | |
| 108 | # If there are pending tasks, wait for first one to complete. |
| 109 | # We only wait for the first one, because after it completes, we want |
| 110 | # to check if any new items are waiting to be processed. |
| 111 | if pending_tasks: |
| 112 | done, pending = await asyncio.wait( |
| 113 | pending_tasks, return_when=asyncio.FIRST_COMPLETED |
| 114 | ) |
| 115 | for task in done: |
| 116 | exc = task.exception() |
| 117 | if exc is not None: |
| 118 | # If a task failed, cancel all other pending tasks. |
| 119 | for p_task in pending: |
| 120 | p_task.cancel() |
| 121 | # Wait for all pending tasks to be cancelled |
| 122 | if pending: |
| 123 | await asyncio.wait(pending) |
| 124 | # Raise the exception from the failed task |
| 125 | raise exc |
| 126 | |
| 127 | index = getattr(task, '_worker_index') |