beam_search: parallel auto_agent workers in a beam_search orchestrator loop

MaxKernel/auto_agent/beam_worker.py

@@ -0,0 +1,28 @@
+"""Beam Search Worker Agent registration.
+
+This module instantiates the correctness-only pipeline agent
+used to generate and verify kernel candidates for Beam Search.
+"""
+
+from auto_agent.subagents.beam_worker_pipeline import BeamWorkerPipeline
+from auto_agent.subagents.kernel_writing import (
+  implement_kernel_agent,
+  plan_kernel_agent,
+  validate_kernel_compilation_agent,
+)
+from auto_agent.subagents.testing import (
+  unified_test_agent,
+  validated_test_generation_agent,
+)
+
+beam_worker_agent = BeamWorkerPipeline(
+  name="BeamWorkerPipeline",
+  plan_agent=plan_kernel_agent,
+  implement_agent=implement_kernel_agent,
+  validate_agent=validate_kernel_compilation_agent,
+  test_gen_agent=validated_test_generation_agent,
+  test_run_agent=unified_test_agent,
+  max_iterations=5,
+)
+
+__all__ = ["beam_worker_agent"]

MaxKernel/auto_agent/subagents/beam_worker_pipeline.py

@@ -0,0 +1,135 @@
+"""Specialized correctness-only pipeline for Beam Search workers."""
+
+import logging
+import os
+from typing import AsyncGenerator
+
+from auto_agent.subagents.pipeline_agent import AutonomousPipelineAgent
+from google.adk.events import Event, EventActions
+from google.adk.agents.invocation_context import InvocationContext
+
+class BeamWorkerPipeline(AutonomousPipelineAgent):
+  """Subclass of AutonomousPipelineAgent that exits early on correctness success."""
+
+  async def _run_async_impl(
+    self, ctx: InvocationContext
+  ) -> AsyncGenerator[Event, None]:
+    iteration = 0
+    yield self._initialize_state(ctx)
+
+    while iteration < self.max_iterations:
+      logging.info(
+        f"[{self.name}] Starting pipeline iteration {iteration + 1}/{self.max_iterations}"
+      )
+
+      # Step 1: Plan
+      logging.info(f"[{self.name}] Running PlanKernelAgent...")
+      async for event in self.plan_agent.run_async(ctx):
+        yield event
+
+      # Step 2: Implement
+      logging.info(f"[{self.name}] Running ImplementKernelAgent...")
+      async for event in self.implement_agent.run_async(ctx):
+        yield event
+
+      # Step 3: Validate (Compilation)
+      logging.info(f"[{self.name}] Running ValidateKernelCompilationAgent...")
+      async for event in self.validate_agent.run_async(ctx):
+        yield event
+
+      # Check if compilation succeeded
+      compilation_status = ctx.session.state.get("kernel_compilation_status", {})
+      if not compilation_status.get("success", False):
+        logging.error(
+          f"[{self.name}] Compilation failed. Looping back to planning."
+        )
+        self._save_iteration_files(
+          ctx, iteration, keys_to_save=["optimized_kernel_path"]
+        )
+        iteration += 1
+        continue
+
+      # Step 4: Test Gen
+      logging.info(f"[{self.name}] Running ValidatedTestGenerationAgent...")
+      async for event in self.test_gen_agent.run_async(ctx):
+        yield event
+
+      # Check if test generation succeeded
+      validation_status = ctx.session.state.get("validation_loop_status", {})
+      if not validation_status.get("success", False):
+        logging.error(
+          f"[{self.name}] Test generation/validation failed. Looping back to planning."
+        )
+        self._save_iteration_files(
+          ctx,
+          iteration,
+          keys_to_save=["optimized_kernel_path", "test_file_path"],
+        )
+        iteration += 1
+        continue
+
+      # Step 5: Test Run
+      logging.info(f"[{self.name}] Running UnifiedTestAgent...")
+      async for event in self.test_run_agent.run_async(ctx):
+        yield event
+
+      # Check if correctness tests passed (Early Exit check)
+      test_results = ctx.session.state.get("test_results", {})
+      if not test_results.get("success", False):
+        logging.error(f"[{self.name}] Tests failed. Looping back to planning.")
+        self._save_iteration_files(
+          ctx,
+          iteration,
+          keys_to_save=["optimized_kernel_path", "test_file_path"]
+        )
+        iteration += 1
+        continue
+
+      kernel_path = ctx.session.state.get("optimized_kernel_path")
+      # Snapshot the successful implementation
+      kernel_code = ""
+      if kernel_path and os.path.exists(kernel_path):
+        try:
+          with open(kernel_path, "r") as f:
+            kernel_code = f.read()
+        except Exception as e:
+          logging.error(
+            f"[{self.name}] Failed to read kernel file for snapshot: {e}"
+          )
+
+      # Extract latency
+      latency = self._extract_latency(ctx)
+
+      yield Event(
+        author=self.name,
+        actions=EventActions(
+          state_delta={
+            "worker_status": "Success",
+            "kernel_code": kernel_code,
+            "compilation_status": ctx.session.state.get(
+              "kernel_compilation_status", {}
+            ),
+            "test_status": ctx.session.state.get("test_results", {}),
+            "latency_ms": latency,
+            "pipeline_status": "Completed",
+            "pipeline_iteration": iteration,
+            "best_iteration": iteration,
+          }
+        ),
+      )
+      self._save_iteration_files(ctx, iteration)
+      return  # Terminate worker early since code is correct
+
+    # If we exit the loop, it means we failed to reach correctness
+    logging.warning(
+      f"[{self.name}] Failed to generate correct code within iteration limit."
+    )
+    yield Event(
+      author=self.name,
+      actions=EventActions(
+        state_delta={
+          "worker_status": "Failed",
+          "pipeline_status": "Failed",
+        }
+      ),
+    )

MaxKernel/auto_agent/tests/test_beam_worker_pipeline.py

@@ -0,0 +1,211 @@
+"""Unit tests for BeamWorkerPipeline execution logic."""
+
+import asyncio
+from typing import AsyncGenerator
+from unittest.mock import MagicMock, patch, mock_open
+import pytest
+
+from auto_agent.subagents.beam_worker_pipeline import BeamWorkerPipeline
+from google.adk.agents.base_agent import BaseAgent
+from google.adk.events import Event
+from google.adk.agents.invocation_context import InvocationContext
+
+
+# Define a dummy agent subclass that inherits from BaseAgent to satisfy Pydantic validations
+class DummySubAgent(BaseAgent):
+  # Custom mock callback runner
+  mock_run: MagicMock = None
+
+  async def _run_async_impl(self, ctx) -> AsyncGenerator[Event, None]:
+    if self.mock_run:
+      async for event in self.mock_run(ctx):
+        yield event
+    else:
+      yield Event(author=self.name)
+
+
+# Helper to construct basic yield generator
+async def default_mock_run(ctx):
+  yield Event(author="mock_agent")
+
+
+@pytest.fixture
+def mock_subagents():
+  """Construct valid subagent instances using the DummySubAgent subclass."""
+  subagents = {
+      "plan": DummySubAgent(name="PlanAgent"),
+      "implement": DummySubAgent(name="ImplementAgent"),
+      "validate": DummySubAgent(name="ValidateAgent"),
+      "test_gen": DummySubAgent(name="TestGenAgent"),
+      "test_run": DummySubAgent(name="TestRunAgent")
+  }
+
+  # Set default side effects
+  for agent in subagents.values():
+    agent.mock_run = MagicMock(side_effect=default_mock_run)
+
+  return subagents
+
+
+@pytest.fixture
+def mock_context():
+  """Create a mock ADK InvocationContext."""
+  from unittest.mock import AsyncMock
+
+  ctx = MagicMock(spec=InvocationContext)
+
+  # Ensure copies of context return the same configured mock context
+  ctx.model_copy.return_value = ctx
+  ctx.end_invocation = False
+
+  ctx.session = MagicMock()
+
+  ctx.session.id = "mock_session_id"
+  ctx.session.state = {
+      "history": [],
+      "kernel_compilation_status": {},
+      "validation_loop_status": {},
+      "test_results": {},
+      "autotune_results": {}
+  }
+
+  # Mock async plugin manager callbacks to be awaitable
+  ctx.plugin_manager = MagicMock()
+  ctx.plugin_manager.run_before_agent_callback = AsyncMock(return_value=None)
+  ctx.plugin_manager.run_after_agent_callback = AsyncMock(return_value=None)
+
+  return ctx
+
+
+
+@pytest.fixture
+def worker_pipeline(mock_subagents):
+  """Instantiate BeamWorkerPipeline with mocked subagents."""
+  pipeline = BeamWorkerPipeline(
+      name="TestBeamWorker",
+      plan_agent=mock_subagents["plan"],
+      implement_agent=mock_subagents["implement"],
+      validate_agent=mock_subagents["validate"],
+      test_gen_agent=mock_subagents["test_gen"],
+      test_run_agent=mock_subagents["test_run"],
+      max_iterations=3
+  )
+  return pipeline
+
+
+@pytest.mark.asyncio
+@patch("os.makedirs")
+@patch("os.path.exists", return_value=True)
+@patch("shutil.copy2")
+@patch("builtins.open", new_callable=mock_open, read_data="def optimized_kernel(): pass")
+async def test_early_exit_on_success(mock_file_open, mock_copy, mock_exists, mock_makedirs, worker_pipeline, mock_context):
+  """Test Case 1: Verifies worker pipeline exits early and yields Success event on correctness success."""
+  # Setup: Set compilation, test generation, and test execution successes to True
+  mock_context.session.state["kernel_compilation_status"] = {"success": True}
+  mock_context.session.state["validation_loop_status"] = {"success": True}
+  mock_context.session.state["test_results"] = {
+      "success": True,
+      "output": "PERF_METRICS: 8.52 ms"
+  }
+  mock_context.session.state["optimized_kernel_path"] = "/dummy/optimized_kernel.py"
+
+  events = []
+  async for event in worker_pipeline._run_async_impl(mock_context):
+    events.append(event)
+
+  # Check that each subagent's mock callback was invoked once
+  assert worker_pipeline.plan_agent.mock_run.call_count == 1
+  assert worker_pipeline.implement_agent.mock_run.call_count == 1
+  assert worker_pipeline.validate_agent.mock_run.call_count == 1
+  assert worker_pipeline.test_gen_agent.mock_run.call_count == 1
+  assert worker_pipeline.test_run_agent.mock_run.call_count == 1
+
+  # Check final emitted success state delta
+  delta_event = events[-1]
+  assert delta_event.actions.state_delta["worker_status"] == "Success"
+  assert delta_event.actions.state_delta["pipeline_status"] == "Completed"
+  assert delta_event.actions.state_delta["latency_ms"] == 8.52
+  assert delta_event.actions.state_delta["kernel_code"] == "def optimized_kernel(): pass"
+
+
+@pytest.mark.asyncio
+@patch("os.makedirs")
+@patch("os.path.exists", return_value=True)
+@patch("shutil.copy2")
+@patch("builtins.open", new_callable=mock_open, read_data="def code(): pass")
+async def test_iteration_loop_retries(mock_file_open, mock_copy, mock_exists, mock_makedirs, worker_pipeline, mock_context):
+  """Test Case 2: Verifies pipeline loops back to planning and retries on failures."""
+  states_history = [
+      # Iteration 0 (Compilation failure)
+      {
+          "kernel_compilation_status": {"success": False},
+          "validation_loop_status": {},
+          "test_results": {}
+      },
+      # Iteration 1 (Test generation failure)
+      {
+          "kernel_compilation_status": {"success": True},
+          "validation_loop_status": {"success": False},
+          "test_results": {}
+      },
+      # Iteration 2 (Correctness Success)
+      {
+          "kernel_compilation_status": {"success": True},
+          "validation_loop_status": {"success": True},
+          "test_results": {"success": True, "output": "PERF_METRICS: 9.1 ms"}
+      }
+  ]
+
+  # Mock run_async of plan_agent to dynamically adjust state context
+  call_count = 0
+  async def plan_side_effect(ctx):
+    nonlocal call_count
+    # Inject state variables for this iteration
+    ctx.session.state["kernel_compilation_status"] = states_history[call_count]["kernel_compilation_status"]
+    ctx.session.state["validation_loop_status"] = states_history[call_count]["validation_loop_status"]
+    ctx.session.state["test_results"] = states_history[call_count]["test_results"]
+    ctx.session.state["optimized_kernel_path"] = "/dummy/optimized_kernel.py"
+    call_count += 1
+    yield Event(author="PlanAgent")
+
+  worker_pipeline.plan_agent.mock_run = MagicMock(side_effect=plan_side_effect)
+
+  events = []
+  async for event in worker_pipeline._run_async_impl(mock_context):
+    events.append(event)
+
+  # Asserts it ran exactly 3 iterations (0, 1, 2)
+  assert call_count == 3
+  assert worker_pipeline.plan_agent.mock_run.call_count == 3
+
+  # Check final status
+  delta_event = events[-1]
+  assert delta_event.actions.state_delta["worker_status"] == "Success"
+  assert delta_event.actions.state_delta["pipeline_status"] == "Completed"
+  assert delta_event.actions.state_delta["latency_ms"] == 9.1
+
+
+@pytest.mark.asyncio
+@patch("os.makedirs")
+@patch("os.path.exists", return_value=True)
+@patch("shutil.copy2")
+@patch("builtins.open", new_callable=mock_open, read_data="def code(): pass")
+async def test_exhaust_iterations_failure(mock_file_open, mock_copy, mock_exists, mock_makedirs, worker_pipeline, mock_context):
+  """Test Case 3: Verifies failure event emitted when iteration limit is hit without correctness success."""
+  # Set max_iterations = 2
+  worker_pipeline.max_iterations = 2
+
+  # Set compilation and test generation success, but correctness tests always fail
+  mock_context.session.state["kernel_compilation_status"] = {"success": True}
+  mock_context.session.state["validation_loop_status"] = {"success": True}
+  mock_context.session.state["test_results"] = {"success": False}
+  mock_context.session.state["optimized_kernel_path"] = "/dummy/optimized_kernel.py"
+
+  events = []
+  async for event in worker_pipeline._run_async_impl(mock_context):
+    events.append(event)
+
+  # Check final emitted event is failure
+  delta_event = events[-1]
+  assert delta_event.actions.state_delta["worker_status"] == "Failed"
+  assert delta_event.actions.state_delta["pipeline_status"] == "Failed"