diff --git a/MaxKernel/.gitignore b/MaxKernel/.gitignore index ba04cde..d7c770f 100644 --- a/MaxKernel/.gitignore +++ b/MaxKernel/.gitignore @@ -4,4 +4,8 @@ session_info/ *egg-info *.txt .adk -eval_config.yaml \ No newline at end of file +eval_config.yaml + +# Local execution environments and outputs +.venv/ +beam_search/output/ \ No newline at end of file diff --git a/MaxKernel/auto_agent/dependency/adk_cli_patch.py b/MaxKernel/auto_agent/dependency/adk_cli_patch.py index 77fd873..2a5a52e 100644 --- a/MaxKernel/auto_agent/dependency/adk_cli_patch.py +++ b/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)") diff --git a/MaxKernel/auto_agent/subagents/pipeline_agent.py b/MaxKernel/auto_agent/subagents/pipeline_agent.py index cea9a88..b78f255 100644 --- a/MaxKernel/auto_agent/subagents/pipeline_agent.py +++ b/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 diff --git a/MaxKernel/auto_agent/subagents/profiling/agent.py b/MaxKernel/auto_agent/subagents/profiling/agent.py index 9aedb37..2b611af 100644 --- a/MaxKernel/auto_agent/subagents/profiling/agent.py +++ b/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 ], ) diff --git a/MaxKernel/auto_agent/subagents/profiling/offline_tools.py b/MaxKernel/auto_agent/subagents/profiling/offline_tools.py index 19b01f0..8570203 100644 --- a/MaxKernel/auto_agent/subagents/profiling/offline_tools.py +++ b/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: diff --git a/MaxKernel/auto_agent/tools/file_tools.py b/MaxKernel/auto_agent/tools/file_tools.py index 8fc47ae..af2d50b 100644 --- a/MaxKernel/auto_agent/tools/file_tools.py +++ b/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