feat: expose max_parallel_workers parameter on workflow.node decorator and Node class

Expose `max_parallel_workers` on the `@workflow.node` decorator and `Node` class when `parallel_worker=True`. This allows throttling the concurrency limit when executing parallel workers.

Co-authored-by: Shangjie Chen <deanchen@google.com>
PiperOrigin-RevId: 940663519

Author: DeanChensj

docs/guides/workflow/parallel_worker/index.md

@@ -7,7 +7,7 @@
 When processing a list of items (e.g., a list of documents to analyze, queries to run, or topics to explain), running them sequentially can be slow, especially if the processing node performs I/O-bound operations like calling an LLM or an external API. `ParallelWorker` solves this by executing the wrapped node concurrently for all items in the input list, significantly reducing total execution time.
 
 Key features:
-- **Concurrency**: Runs multiple instances of the wrapped node in parallel.
+- **Concurrency**: Runs multiple instances of the wrapped node in parallel (optionally throttled via `max_parallel_workers`).
 - **Aggregation**: Gathers all outputs and returns them as a single list, maintaining the original order of the inputs.
 - **Error Propagation**: If any parallel task fails, all other pending tasks are cancelled, and the error is raised immediately.
 
@@ -64,7 +64,6 @@ If the wrapped node triggers an interrupt (e.g., `RequestInput`), the parallel w
 
 - **List Input**: The worker always expects a list and returns a list. If your upstream node doesn't produce a list, it will be treated as a list of one item.
 - **Fail-Fast**: The failure of a single item fails the entire worker and cancels all other items. There is currently no "continue on error" option to collect partial results.
-- **No Concurrency Limit**: Setting `max_concurrency` to limit the number of parallel tasks is not supported yet. All items in the list will be processed concurrently.
 
 ## Related samples
 

src/google/adk/cli/agent_test_runner.py

@@ -257,7 +257,7 @@ def _make_nodes_sequential(obj, visited=None):
       for node in obj.graph.nodes:
         _make_nodes_sequential(node, visited)
   elif isinstance(obj, _ParallelWorker):
-    obj.max_concurrency = 1
+    obj.max_parallel_workers = 1
     if hasattr(obj, "_node"):
       _make_nodes_sequential(obj._node, visited)
 

src/google/adk/workflow/_node.py

@@ -25,6 +25,7 @@
 from typing import TypeVar
 
 from pydantic import Field
+from pydantic import model_validator
 from pydantic import PrivateAttr
 from typing_extensions import override
 
@@ -50,6 +51,7 @@ def node(
     retry_config: RetryConfig | None = None,
     timeout: float | None = None,
     parallel_worker: bool = False,
+    max_parallel_workers: int | None = None,
     auth_config: AuthConfig | None = None,
     parameter_binding: Literal['state', 'node_input'] = 'state',
 ) -> Callable[
@@ -67,6 +69,7 @@ def node(
     retry_config: RetryConfig | None = None,
     timeout: float | None = None,
     parallel_worker: bool = False,
+    max_parallel_workers: int | None = None,
     auth_config: AuthConfig | None = None,
     parameter_binding: Literal['state', 'node_input'] = 'state',
 ) -> base_node.BaseNode:
@@ -81,6 +84,7 @@ def node(
     retry_config: RetryConfig | None = None,
     timeout: float | None = None,
     parallel_worker: bool = False,
+    max_parallel_workers: int | None = None,
     auth_config: AuthConfig | None = None,
     parameter_binding: Literal['state', 'node_input'] = 'state',
 ) -> Any:
@@ -123,6 +127,16 @@ async def my_func3(): ...
     a BaseNode instance.
   """
 
+  if max_parallel_workers is not None:
+    if not parallel_worker:
+      raise ValueError(
+          'max_parallel_workers can only be set when parallel_worker is True.'
+      )
+    if max_parallel_workers < 1:
+      raise ValueError(
+          'max_parallel_workers must be greater than or equal to 1.'
+      )
+
   def wrapper(
       func: T,
   ) -> function_node.FunctionNode | parallel_worker_lib._ParallelWorker:
@@ -138,7 +152,9 @@ def wrapper(
         parameter_binding=parameter_binding,
     )
     if parallel_worker:
-      return parallel_worker_lib._ParallelWorker(node=built_node)
+      return parallel_worker_lib._ParallelWorker(
+          node=built_node, max_parallel_workers=max_parallel_workers
+      )
     return built_node
 
   if node_like is None:
@@ -155,7 +171,9 @@ def wrapper(
         parameter_binding=parameter_binding,
     )
     if parallel_worker:
-      return parallel_worker_lib._ParallelWorker(node=built_node)
+      return parallel_worker_lib._ParallelWorker(
+          node=built_node, max_parallel_workers=max_parallel_workers
+      )
     return built_node
 
 
@@ -167,8 +185,22 @@ class Node(base_node.BaseNode):
   """
 
   parallel_worker: bool = Field(default=False, frozen=True)
+  max_parallel_workers: int | None = Field(default=None, frozen=True)
   _inner_node: base_node.BaseNode | None = PrivateAttr(default=None)
 
+  @model_validator(mode='after')
+  def _validate_parallel_worker_config(self) -> Node:
+    if self.max_parallel_workers is not None:
+      if not self.parallel_worker:
+        raise ValueError(
+            'max_parallel_workers can only be set when parallel_worker is True.'
+        )
+      if self.max_parallel_workers < 1:
+        raise ValueError(
+            'max_parallel_workers must be greater than or equal to 1.'
+        )
+    return self
+
   def model_post_init(self, __context: Any) -> None:
     super().model_post_init(__context)
     if self.parallel_worker:
@@ -179,7 +211,9 @@ def model_post_init(self, __context: Any) -> None:
       # original (essential for LlmAgent and Workflow subclasses).
       worker_node = self.model_copy(update={'parallel_worker': False})
 
-      inner = parallel_worker_lib._ParallelWorker(node=worker_node)
+      inner = parallel_worker_lib._ParallelWorker(
+          node=worker_node, max_parallel_workers=self.max_parallel_workers
+      )
       self._inner_node = inner
       # Synchronize rerun_on_resume with the inner node.
       self.rerun_on_resume = inner.rerun_on_resume
@@ -193,7 +227,9 @@ def model_copy(
 
     if copied.parallel_worker:
       worker_node = copied.model_copy(update={'parallel_worker': False})
-      copied._inner_node = parallel_worker_lib._ParallelWorker(node=worker_node)
+      copied._inner_node = parallel_worker_lib._ParallelWorker(
+          node=worker_node, max_parallel_workers=copied.max_parallel_workers
+      )
       copied.rerun_on_resume = copied._inner_node.rerun_on_resume
 
     return copied

src/google/adk/workflow/_parallel_worker.py

@@ -36,21 +36,21 @@ class _ParallelWorker(BaseNode):
   """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.
+    max_parallel_workers: The maximum number of parallel tasks to run. If None,
+      there is no limit on concurrency.
   """
 
   model_config = ConfigDict(arbitrary_types_allowed=True)
 
-  max_concurrency: int | None = Field(default=None)
+  max_parallel_workers: int | None = Field(default=None)
 
   _node: BaseNode = PrivateAttr()
 
   def __init__(
       self,
       *,
       node: NodeLike,
-      max_concurrency: int | None = None,
+      max_parallel_workers: int | None = None,
       retry_config: RetryConfig | None = None,
       timeout: float | None = None,
   ):
@@ -63,8 +63,12 @@ def __init__(
         retry_config=retry_config,
         timeout=timeout,
     )
+    if max_parallel_workers is not None and max_parallel_workers < 1:
+      raise ValueError(
+          'max_parallel_workers must be greater than or equal to 1.'
+      )
     self._node = built_node
-    self.max_concurrency = max_concurrency
+    self.max_parallel_workers = max_parallel_workers
 
   @override
   async def _run_impl(
@@ -89,8 +93,8 @@ async def _run_impl(
     while input_index < len(node_input) or pending_tasks:
       # Check for any inputs waiting to be processed.
       while input_index < len(node_input) and (
-          self.max_concurrency is None
-          or len(pending_tasks) < self.max_concurrency
+          self.max_parallel_workers is None
+          or len(pending_tasks) < self.max_parallel_workers
       ):
         item = node_input[input_index]
         task = asyncio.create_task(

tests/unittests/workflow/test_workflow_node.py

@@ -354,3 +354,97 @@ async def producer():
       "subclass",
       ["subclass: workflow -> input1", "subclass: workflow -> input2"],
   ) in by_node
+
+
+def test_node_decorator_parallel_worker_max_parallel_workers():
+  """Tests that node() correctly sets max_parallel_workers on ParallelWorker."""
+
+  @node(parallel_worker=True, max_parallel_workers=3)
+  def my_func(node_input):
+    return node_input
+
+  assert isinstance(my_func, ParallelWorker)
+  assert my_func.max_parallel_workers == 3
+
+
+def test_node_decorator_invalid_max_parallel_workers():
+  """Tests that node() raises ValueError if max_parallel_workers is set without parallel_worker."""
+  with pytest.raises(
+      ValueError,
+      match="max_parallel_workers can only be set when parallel_worker is True",
+  ):
+
+    @node(parallel_worker=False, max_parallel_workers=3)
+    def my_func(node_input):
+      return node_input
+
+
+def test_node_subclass_invalid_max_parallel_workers():
+  """Tests that Node subclass raises ValidationError if max_parallel_workers is set without parallel_worker."""
+  from pydantic import ValidationError
+
+  with pytest.raises(ValidationError) as exc_info:
+    _CustomNode(name="subclass", parallel_worker=False, max_parallel_workers=3)
+
+  assert (
+      "max_parallel_workers can only be set when parallel_worker is True"
+      in str(exc_info.value)
+  )
+
+
+def test_node_subclassing_model_copy_preserves_max_parallel_workers():
+  """Tests that Node.model_copy preserves max_parallel_workers."""
+  node_inst = _CustomNode(
+      name="subclass",
+      parallel_worker=True,
+      max_parallel_workers=5,
+      custom_val="barrier",
+  )
+  assert node_inst.parallel_worker is True
+  assert node_inst.max_parallel_workers == 5
+  assert node_inst._inner_node.max_parallel_workers == 5
+
+  cloned = node_inst.model_copy()
+  assert isinstance(cloned, _CustomNode)
+  assert cloned.parallel_worker is True
+  assert cloned.max_parallel_workers == 5
+  assert isinstance(cloned._inner_node, ParallelWorker)
+  assert cloned._inner_node.max_parallel_workers == 5
+  assert cloned._inner_node._node.parallel_worker is False
+
+
+def test_node_decorator_invalid_max_parallel_workers_less_than_one():
+  """Tests that node() raises ValueError if max_parallel_workers is less than 1."""
+  with pytest.raises(
+      ValueError,
+      match="max_parallel_workers must be greater than or equal to 1",
+  ):
+
+    @node(parallel_worker=True, max_parallel_workers=0)
+    def my_func(node_input):
+      return node_input
+
+
+def test_node_subclass_invalid_max_parallel_workers_less_than_one():
+  """Tests that Node subclass raises ValidationError if max_parallel_workers is less than 1."""
+  from pydantic import ValidationError
+
+  with pytest.raises(ValidationError) as exc_info:
+    _CustomNode(name="subclass", parallel_worker=True, max_parallel_workers=0)
+
+  assert "max_parallel_workers must be greater than or equal to 1" in str(
+      exc_info.value
+  )
+
+
+def test_parallel_worker_invalid_max_parallel_workers_less_than_one():
+  """Tests that ParallelWorker constructor raises ValueError if max_parallel_workers is less than 1."""
+
+  def dummy_func(x):
+    return x
+
+  with pytest.raises(
+      ValueError,
+      match="max_parallel_workers must be greater than or equal to 1",
+  ):
+    ParallelWorker(node=dummy_func, max_parallel_workers=0)