auto_agent: Pre-requisites of auto_agent beam search prototype using mock backends

MaxKernel/.gitignore

@@ -4,4 +4,8 @@ session_info/
 *egg-info
 *.txt
 .adk
-eval_config.yaml
+eval_config.yaml
+
+# Local execution environments and outputs
+.venv/
+beam_search/output/

MaxKernel/auto_agent/dependency/adk_cli_patch.py

@@ -37,16 +37,14 @@ def apply_patch():
     return True
 
   # Patch 1: Modify run_interactively to return runner instead of closing it
-  old_code_1 = """  await runner.close()
+  old_code_1 = """      resume_invocation_id = invocation_id
 
+  await runner.close()"""
 
-async def run_cli("""
+  new_code_1 = """      resume_invocation_id = invocation_id
 
-  new_code_1 = """  # PATCHED: MCP cleanup fix - return runner instead of closing
-  return runner
-
-
-async def run_cli("""
+  # PATCHED: MCP cleanup fix - return runner instead of closing
+  return runner"""
 
   if old_code_1 not in content:
     print("Error: Could not find code to patch (section 1)")
@@ -55,26 +53,8 @@ async def run_cli("""
   content = content.replace(old_code_1, new_code_1)
 
   # Patch 2: Update run_cli to close runner after session saving
-  old_code_2 = """    await run_interactively(
-        agent_or_app,
-        artifact_service,
-        session,
-        session_service,
-        credential_service,
-    )
-
-  if save_session:"""
-
-  new_code_2 = """    # PATCHED: MCP cleanup fix - get runner to close it later
-    runner = await run_interactively(
-        agent_or_app,
-        artifact_service,
-        session,
-        session_service,
-        credential_service,
-    )
-
-  if save_session:"""
+  old_code_2 = "    await run_interactively("
+  new_code_2 = "    runner = await run_interactively("
 
   if old_code_2 not in content:
     print("Error: Could not find code to patch (section 2)")

MaxKernel/auto_agent/subagents/pipeline_agent.py

@@ -23,9 +23,10 @@ class AutonomousPipelineAgent(BaseAgent):
   validate_agent: BaseAgent
   test_gen_agent: BaseAgent
   test_run_agent: BaseAgent
-  autotune_agent: BaseAgent
-  profile_agent: BaseAgent
+  autotune_agent: BaseAgent | None = None
+  profile_agent: BaseAgent | None = None
   max_iterations: int = 2
+  stop_on_first_valid: bool = False
 
   def __init__(
     self,
@@ -35,9 +36,10 @@ def __init__(
     validate_agent: BaseAgent,
     test_gen_agent: BaseAgent,
     test_run_agent: BaseAgent,
-    autotune_agent: BaseAgent,
-    profile_agent: BaseAgent,
+    autotune_agent: BaseAgent | None = None,
+    profile_agent: BaseAgent | None = None,
     max_iterations: int = 2,
+    stop_on_first_valid: bool = False,
   ):
     super().__init__(
       name=name,
@@ -49,6 +51,7 @@ def __init__(
       autotune_agent=autotune_agent,
       profile_agent=profile_agent,
       max_iterations=max_iterations,
+      stop_on_first_valid=stop_on_first_valid,
     )
 
   async def _run_async_impl(
@@ -129,14 +132,16 @@ async def _run_async_impl(
         continue
 
       # Step 6: Autotune
-      logging.info(f"[{self.name}] Running AutotuneAgent...")
-      async for event in self.autotune_agent.run_async(ctx):
-        yield event
+      if self.autotune_agent:
+        logging.info(f"[{self.name}] Running AutotuneAgent...")
+        async for event in self.autotune_agent.run_async(ctx):
+          yield event
 
       # Step 7: Profile
-      logging.info(f"[{self.name}] Running ProfileAgentOrchestrator...")
-      async for event in self.profile_agent.run_async(ctx):
-        yield event
+      if self.profile_agent:
+        logging.info(f"[{self.name}] Running ProfileAgentOrchestrator...")
+        async for event in self.profile_agent.run_async(ctx):
+          yield event
 
       # Snapshot intermediate result
       kernel_path = ctx.session.state.get("optimized_kernel_path")
@@ -177,6 +182,19 @@ async def _run_async_impl(
 
       self._save_iteration_files(ctx, iteration)
 
+      # Check if we should stop on first successful validation
+      if self.stop_on_first_valid:
+        has_valid_candidate = any(
+            s.get("compilation_status", {}).get("success")
+            and s.get("test_status", {}).get("success")
+            for s in ctx.session.state.get("history", [])
+        )
+        if has_valid_candidate:
+          logging.info(
+              f"[{self.name}] stop_on_first_valid is True and a valid, profiled candidate was generated. Stopping pipeline."
+          )
+          break
+
       # Step 7: Check if improvement is needed
       # needs_improvement = ctx.session.state.get("needs_improvement", False)
       needs_improvement = True

MaxKernel/auto_agent/subagents/profiling/agent.py

@@ -119,11 +119,15 @@ async def _run_async_impl(
   output_key="profiling_summary",
   include_contents="none",
   tools=[
-    offline_tools.load_xplane_and_query,
-    offline_tools.get_hlo_dump,
-    offline_tools.create_chart_from_xplane,
-    offline_tools.get_overview_page_metrics,
-    vertex_ai_rag_tool,
+      t
+      for t in [
+          offline_tools.load_xplane_and_query,
+          offline_tools.get_hlo_dump,
+          offline_tools.create_chart_from_xplane,
+          offline_tools.get_overview_page_metrics,
+          vertex_ai_rag_tool,
+      ]
+      if t is not None
   ],
 )
 

MaxKernel/auto_agent/subagents/profiling/offline_tools.py

@@ -2,6 +2,7 @@
 
 import gzip
 import json
+import os
 import sqlite3
 from typing import Optional
 
@@ -10,6 +11,22 @@
 from tensorflow.tsl.profiler.protobuf import xplane_pb2
 
 
+def is_mock_path(path: str) -> bool:
+  return "mock_profile" in path or os.environ.get("MOCK_COMPILER", "false").lower() == "true"
+
+
+def to_markdown_simple(df: pd.DataFrame) -> str:
+  cols = list(df.columns)
+  headers = "| " + " | ".join(str(c) for c in cols) + " |"
+  separator = "| " + " | ".join("---" for _ in cols) + " |"
+  lines = [headers, separator]
+  for _, row in df.iterrows():
+    lines.append("| " + " | ".join(str(row[c]) for c in cols) + " |")
+  return "\n".join(lines)
+
+
+
+
 def _get_xplane_path(profiling_results: str) -> str:
   """Extracts xplane path from profiling results string or context."""
 
@@ -31,6 +48,16 @@ def load_xplane_and_query(xplane_path: str, sql_query: str) -> str:
   Returns:
       A markdown-formatted table of the query results.
   """
+  if is_mock_path(xplane_path):
+    data = {
+        "plane_name": ["Device 0", "Device 0"],
+        "line_name": ["Compute", "Memory"],
+        "event_name": ["MatMul", "HBM_to_VMEM"],
+        "duration_us": [1500.0, 420.0],
+    }
+    df = pd.DataFrame(data)
+    return to_markdown_simple(df)
+
   try:
     # Open file (handle gz if needed)
     open_func = gzip.open if xplane_path.endswith(".gz") else open
@@ -89,7 +116,7 @@ def get_meta_name(meta_map, mid):
     df = pd.read_sql_query(sql_query, conn)
     conn.close()
 
-    return df.to_markdown(index=False)
+    return to_markdown_simple(df)
 
   except Exception as e:
     return f"Error executing query: {e}"
@@ -152,6 +179,12 @@ def create_chart_from_xplane(
       y_col: Column for Y axis (bar) or values (pie).
       title: Chart title.
   """
+  if is_mock_path(xplane_path):
+    output_filename = f"{xplane_path}.png"
+    with open(output_filename, "wb") as f:
+      f.write(b"MOCK_PNG_DATA")
+    return f"Chart saved to {output_filename}"
+
   try:
     # Re-use loading logic (inefficient but stateless)
     # TODO: we might want to cache the DB connection or pass it around.
@@ -240,6 +273,19 @@ def get_overview_page_metrics(xplane_path: str) -> str:
   Returns:
       A JSON string containing metrics and metadata.
   """
+  if is_mock_path(xplane_path):
+    metrics = {
+        "device_count": 1,
+        "host_count": 1,
+        "total_duration_ms": 1.92,
+        "device_duty_cycle_percent": 78.0,
+        "average_step_time_ms": 1.92,
+        "step_count": 1,
+        "build_target": "N/A (Offline Mock)",
+        "xid": "N/A (Offline Mock)",
+    }
+    return json.dumps(metrics, indent=2)
+
   try:
     open_func = gzip.open if xplane_path.endswith(".gz") else open
     with open_func(xplane_path, "rb") as f:

MaxKernel/auto_agent/tools/file_tools.py

@@ -1,6 +1,7 @@
 """File-related tools for subagents."""
 
 import json
+import logging
 import os
 from pathlib import Path
 from typing import Any, Dict, List
@@ -55,21 +56,35 @@ def restricted_write_file(state_key: str, description: str) -> FunctionTool:
 
   def _write_file(content: str, tool_context: ToolContext) -> str:
     target_path = tool_context.state.get(state_key)
+    logging.info(f"[restricted_write_file] Called for {state_key}. Target path from state: {target_path}")
     if not target_path:
-      return f"Error: Path variable '{state_key}' not found in session state."
+      err = f"Error: Path variable '{state_key}' not found in session state."
+      logging.error(f"[restricted_write_file] {err}")
+      return err
 
     base = Path(WORKDIR).resolve()
     target = Path(target_path).resolve()
 
     try:
       if not target.is_relative_to(base):
-        return f"Error: Access denied. Path is outside {WORKDIR}"
-    except ValueError:
-      return "Error: Invalid path or access denied."
+        err = f"Error: Access denied. Path {target} is outside base {base}"
+        logging.error(f"[restricted_write_file] {err}")
+        return err
+    except ValueError as e:
+      err = f"Error: Invalid path or access denied: {e}"
+      logging.error(f"[restricted_write_file] {err}")
+      return err
 
-    target.parent.mkdir(parents=True, exist_ok=True)
-    target.write_text(content)
-    return f"Successfully wrote to {target}"
+    try:
+      target.parent.mkdir(parents=True, exist_ok=True)
+      target.write_text(content)
+      res = f"Successfully wrote to {target}"
+      logging.info(f"[restricted_write_file] {res}")
+      return res
+    except Exception as e:
+      err = f"Error writing file {target}: {e}"
+      logging.error(f"[restricted_write_file] {err}")
+      return err
 
   _write_file.__name__ = "restricted_write_file"
   _write_file.__doc__ = description