A node that runs a wrapped node in parallel for each item in the input list. Attributes: max_concurrency: The maximum number of parallel tasks to run. If None, there is no limit on concurrency.
| 33 | |
| 34 | |
| 35 | class _ParallelWorker(BaseNode): |
| 36 | """A node that runs a wrapped node in parallel for each item in the input list. |
| 37 | |
| 38 | Attributes: |
| 39 | max_concurrency: The maximum number of parallel tasks to run. If None, there |
| 40 | is no limit on concurrency. |
| 41 | """ |
| 42 | |
| 43 | model_config = ConfigDict(arbitrary_types_allowed=True) |
| 44 | |
| 45 | max_concurrency: int | None = Field(default=None) |
| 46 | |
| 47 | _node: BaseNode = PrivateAttr() |
| 48 | |
| 49 | def __init__( |
| 50 | self, |
| 51 | *, |
| 52 | node: NodeLike, |
| 53 | max_concurrency: int | None = None, |
| 54 | retry_config: RetryConfig | None = None, |
| 55 | timeout: float | None = None, |
| 56 | ): |
| 57 | if node == 'START': |
| 58 | raise ValueError('ParallelWorker cannot wrap a START node.') |
| 59 | built_node = build_node(node) |
| 60 | super().__init__( |
| 61 | name=built_node.name, |
| 62 | rerun_on_resume=True, |
| 63 | retry_config=retry_config, |
| 64 | timeout=timeout, |
| 65 | ) |
| 66 | self._node = built_node |
| 67 | self.max_concurrency = max_concurrency |
| 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 |