From 24cd392fd32d41ba6230f3321f923570c5d09c71 Mon Sep 17 00:00:00 2001
From: Yifeng Lu <yifenglu@google.com>
Date: Mon, 29 Jun 2026 13:42:26 -0700
Subject: [PATCH] Fan out independent eval leaf nodes concurrently in the eval
 runner

PiperOrigin-RevId: 940006704
---
 langfun/core/eval/base.py     | 14 +++++++++-----
 langfun/core/llms/__init__.py |  1 +
 langfun/core/llms/gemini.py   | 26 ++++++++++++++++++++++++++
 langfun/core/llms/vertexai.py |  7 +++++++
 4 files changed, 43 insertions(+), 5 deletions(-)

diff --git a/langfun/core/eval/base.py b/langfun/core/eval/base.py
index b9dc226d..c380f4ce 100644
--- a/langfun/core/eval/base.py
+++ b/langfun/core/eval/base.py
@@ -325,16 +325,19 @@ def _run_group(arg: tuple[int, list[_LeafNode]]) -> None:
               'LastCompleted': leaf.node.id
           })
 
-      # NOTE(daiyip): Run leaf nodes grouped by model resource id. This allows
-      # evaluations using the same resource to run sequentially, which favors
-      # completing evaluations over running evaluations sparsely.
+      # NOTE: Fan out all independent leaf evaluations concurrently. Each leaf
+      # is placed in its own group (keyed by its unique leaf id) and the run is
+      # dispatched with max_workers == number of leaves, so every leaf can be
+      # in flight at once. Per-model backpressure is still enforced by the
+      # LM-level rate-limit semaphore (Layer C, keyed by resource_id), which
+      # remains the only cap on concurrent requests to a given model.
       filter = filter or (lambda x: True)
       leaf_nodes: list[_LeafNode] = []
       leaf_groups: dict[str, list[_LeafNode]] = collections.defaultdict(list)
 
       for i, leaf in enumerate(self.leaf_nodes):
         node = _LeafNode(index=i + 1, node=leaf, enabled=filter(leaf))
-        leaf_groups[leaf.lm.resource_id].append(node)
+        leaf_groups[leaf.id].append(node)
         leaf_nodes.append(node)
 
       if leaf_groups:
@@ -354,7 +357,8 @@ def _run_group(arg: tuple[int, list[_LeafNode]]) -> None:
               _run_group,
               [(overview_bar, group) for group in leaf_groups.values()],
               silence_on_errors=None,
-              max_workers=len(leaf_groups)):
+              max_workers=len(leaf_nodes),
+          ):
             pass
 
           # Save results for non-leaf nodes.
diff --git a/langfun/core/llms/__init__.py b/langfun/core/llms/__init__.py
index c22bc027..0781657e 100644
--- a/langfun/core/llms/__init__.py
+++ b/langfun/core/llms/__init__.py
@@ -98,6 +98,7 @@
 from langfun.core.llms.vertexai import VertexAIGemini3ProPreview
 from langfun.core.llms.vertexai import VertexAIGemini3ProImagePreview
 from langfun.core.llms.vertexai import VertexAIGemini3FlashPreview
+from langfun.core.llms.vertexai import VertexAIGemini35Flash
 
 # Veo video generation models.
 from langfun.core.llms.veo import Veo
diff --git a/langfun/core/llms/gemini.py b/langfun/core/llms/gemini.py
index b2aacdcb..2488daea 100644
--- a/langfun/core/llms/gemini.py
+++ b/langfun/core/llms/gemini.py
@@ -309,6 +309,32 @@ def estimate_cost(self, usage: lf.LMSamplingUsage) -> float | None:
             max_tokens_per_minute=4_000_000,
         ),
     ),
+    # Gemini 3.5 Flash
+    GeminiModelInfo(
+        model_id='gemini-3.5-flash',
+        in_service=True,
+        provider=pg.oneof(['Google GenAI', 'VertexAI']),
+        model_type='instruction-tuned',
+        description=(
+            'Gemini 3.5 Flash: High-efficiency, low-latency multimodal'
+            ' model optimized for agentic workflows.'
+        ),
+        release_date=datetime.datetime(2026, 5, 19),
+        input_modalities=GeminiModelInfo.ALL_SUPPORTED_INPUT_TYPES,
+        context_length=lf.ModelInfo.ContextLength(
+            max_input_tokens=1_048_576,
+            max_output_tokens=65_536,
+        ),
+        pricing=GeminiModelInfo.Pricing(
+            cost_per_1m_cached_input_tokens=0.15,
+            cost_per_1m_input_tokens=1.50,
+            cost_per_1m_output_tokens=9.00,
+        ),
+        rate_limits=lf.ModelInfo.RateLimits(
+            max_requests_per_minute=2_000,
+            max_tokens_per_minute=4_000_000,
+        ),
+    ),
     # Gemini 2.5 Flash
     GeminiModelInfo(
         model_id='gemini-2.5-flash',
diff --git a/langfun/core/llms/vertexai.py b/langfun/core/llms/vertexai.py
index 03673323..cf743f15 100644
--- a/langfun/core/llms/vertexai.py
+++ b/langfun/core/llms/vertexai.py
@@ -294,6 +294,13 @@ class VertexAIGemini3FlashPreview(VertexAIGemini):  # pylint: disable=invalid-na
   location = 'global'
 
 
+class VertexAIGemini35Flash(VertexAIGemini):  # pylint: disable=invalid-name
+  """Gemini 3.5 Flash GA model launched on 05/19/2026."""
+
+  model = 'gemini-3.5-flash'
+  location = 'global'
+
+
 class VertexAIGemini31FlashLitePreview(VertexAIGemini):  # pylint: disable=invalid-name
   """Gemini 3.1 Flash Lite Preview model."""