diff --git a/MaxKernel/auto_agent/agent_client/auto_agent_client.py b/MaxKernel/auto_agent/agent_client/auto_agent_client.py index 28eee18..56ae414 100644 --- a/MaxKernel/auto_agent/agent_client/auto_agent_client.py +++ b/MaxKernel/auto_agent/agent_client/auto_agent_client.py @@ -14,27 +14,26 @@ "dotenv not installed, skipping loading environment variables" ) +from google.adk.apps.app import App from google.adk.runners import Runner from google.adk.sessions.in_memory_session_service import InMemorySessionService from google.genai.types import Content, Part from auto_agent.agent import root_agent +from auto_agent.config import get_compaction_config logger = logging.getLogger(__name__) class AutoAgentClient: - user_id: str - session_id: str - query: str - app_name: str = "auto_agent" - def __init__( self, user_id: str, session_id: str, query: str, agent: Optional[Any] = None, + app_name: str = "auto_agent", + events_compaction: bool = False, ): self.user_id = user_id self.session_id = session_id @@ -42,6 +41,8 @@ def __init__( self.agent = agent or root_agent self.session_service = InMemorySessionService() self.session = None + self.app_name = app_name + self.events_compaction = events_compaction async def create_session( self, initial_state: Optional[dict[str, Any]] = None @@ -77,9 +78,19 @@ async def run_async(self) -> None: if not self.session: await self.create_session() + # Compaction configuration must be passed via an App object + if self.events_compaction: + compaction_config = get_compaction_config() + else: + compaction_config = None + app = App( + name=self.app_name, + root_agent=self.agent, + events_compaction_config=compaction_config, + ) + runner = Runner( - app_name=self.app_name, - agent=self.agent, + app=app, session_service=self.session_service, ) @@ -124,6 +135,11 @@ def main(): default="client_query.txt", help="File containing the query to send", ) + parser.add_argument( + "--events_compaction", + action="store_true", + help="Enable event compaction", + ) args = parser.parse_args() user_id = args.user_id @@ -132,7 +148,12 @@ def main(): query = read_query_from_file(args.query_file) # Create client instance - client = AutoAgentClient(user_id, session_id, query) + client = AutoAgentClient( + user_id=user_id, + session_id=session_id, + query=query, + events_compaction=args.events_compaction, + ) logger.info( f"Generating script for user {user_id} in session {session_id} with query: {query}" diff --git a/MaxKernel/auto_agent/agent_client/run_batch_agent_call.py b/MaxKernel/auto_agent/agent_client/run_batch_agent_call.py index 17958f1..b8b3426 100644 --- a/MaxKernel/auto_agent/agent_client/run_batch_agent_call.py +++ b/MaxKernel/auto_agent/agent_client/run_batch_agent_call.py @@ -27,6 +27,11 @@ def parse_args(): default=2, help="Maximum number of retries for failed tasks", ) + parser.add_argument( + "--events_compaction", + action="store_true", + help="Enable event compaction", + ) parser.add_argument( "--log_file", type=str, default=None, help="File to save logs to" ) @@ -61,6 +66,7 @@ async def process_problem( data_dir: str, max_retries: int, sem: asyncio.Semaphore, + events_compaction: bool = False, ): async with sem: problem_dir = os.path.join(data_dir, problem_id) @@ -88,6 +94,7 @@ async def process_problem( user_id=user_id, session_id=session_id, query=query, + events_compaction=events_compaction, ) session_file = os.path.join( @@ -159,7 +166,9 @@ async def main_async(args): sem = asyncio.Semaphore(args.num_concurrent) tasks = [ - process_problem(problem, args.data_dir, args.max_retries, sem) + process_problem( + problem, args.data_dir, args.max_retries, sem, args.events_compaction + ) for problem in problems ] diff --git a/MaxKernel/auto_agent/config.py b/MaxKernel/auto_agent/config.py index 002b757..ba85a9e 100644 --- a/MaxKernel/auto_agent/config.py +++ b/MaxKernel/auto_agent/config.py @@ -19,9 +19,9 @@ # Set events compaction policy to avoid memory overflow def get_compaction_config(): return EventsCompactionConfig( - token_threshold=200000, - event_retention_size=100, - compaction_interval=1, + token_threshold=300000, + event_retention_size=5, + compaction_interval=0, overlap_size=0, ) diff --git a/MaxKernel/auto_agent/custom_types.py b/MaxKernel/auto_agent/custom_types.py index 37398a5..f0f1fa6 100644 --- a/MaxKernel/auto_agent/custom_types.py +++ b/MaxKernel/auto_agent/custom_types.py @@ -8,10 +8,6 @@ from google.adk.models.google_llm import Gemini from google.genai import Client, types -from auto_agent.constants import ( - MODEL_NAME, -) - class TimeoutGemini(Gemini): @cached_property @@ -46,7 +42,7 @@ def __init__(self, *args, **kwargs): # If model is a string, use the pre-configured gemini_model with retry support if "model" in kwargs and isinstance(kwargs["model"], str): gemini_model = TimeoutGemini( - model=MODEL_NAME, + model=kwargs["model"], retry_options=types.HttpRetryOptions( initial_delay=1, attempts=10, diff --git a/MaxKernel/auto_agent/subagents/autotuning/agent.py b/MaxKernel/auto_agent/subagents/autotuning/agent.py index 084b7f4..18659ea 100644 --- a/MaxKernel/auto_agent/subagents/autotuning/agent.py +++ b/MaxKernel/auto_agent/subagents/autotuning/agent.py @@ -25,18 +25,25 @@ ) from auto_agent.tools.search_api_tool import search_api_tool + # 1. Planner Agent # This agent identifies parameters, creates the template, and defines the search space. # It saves them to session state instead of calling the tool directly. -autotune_planner_agent = CustomLlmAgent( - name="AutotunePlannerAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=get_thinking_planner("high"), - instruction=autotune_prompt.PROMPT, - description="Prepares code template and search space for auto-tuning Pallas kernels.", - tools=[filesystem_tool_r, write_autotune_specs_tool, search_api_tool], -) +def create_autotune_planner_agent( + model_name: str = MODEL_NAME, +) -> CustomLlmAgent: + return CustomLlmAgent( + name="AutotunePlannerAgent", + model=model_name, + generate_content_config=model_config, + planner=get_thinking_planner("high"), + instruction=autotune_prompt.PROMPT, + description="Prepares code template and search space for auto-tuning Pallas kernels.", + tools=[filesystem_tool_r, write_autotune_specs_tool, search_api_tool], + ) + + +autotune_planner_agent = create_autotune_planner_agent() # 2. Runner Agent @@ -154,52 +161,86 @@ async def _run_async_impl( ) -autotune_runner = AutotuneRunner( - name="AutotuneRunner", - output_key="autotune_results", -) +def create_autotune_runner() -> AutotuneRunner: + return AutotuneRunner( + name="AutotuneRunner", + output_key="autotune_results", + ) + + +autotune_runner = create_autotune_runner() # 3. Apply Best Config Agent -apply_best_config_agent = CustomLlmAgent( - name="ApplyBestConfigAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=get_thinking_planner("high"), - instruction=apply_best_config_prompt.PROMPT, - description="Applies autotuning results to the optimized kernel file.", - tools=[filesystem_tool_r, write_optimized_kernel_tool], -) +def create_apply_best_config_agent( + model_name: str = MODEL_NAME, +) -> CustomLlmAgent: + return CustomLlmAgent( + name="ApplyBestConfigAgent", + model=model_name, + generate_content_config=model_config, + planner=get_thinking_planner("high"), + instruction=apply_best_config_prompt.PROMPT, + description="Applies autotuning results to the optimized kernel file.", + tools=[filesystem_tool_r, write_optimized_kernel_tool], + ) + + +apply_best_config_agent = create_apply_best_config_agent() # 4. Summarizer Agent # This agent reads results from state and talks to the user. -autotune_summary_agent = CustomLlmAgent( - name="AutotuneSummaryAgent", - model=MODEL_NAME, - generate_content_config=model_config, - instruction=summary_prompt.PROMPT, - description="Summarizes autotuning results.", - tools=[filesystem_tool_r], - output_key="autotuning_summary", -) +def create_autotune_summary_agent( + model_name: str = MODEL_NAME, +) -> CustomLlmAgent: + return CustomLlmAgent( + name="AutotuneSummaryAgent", + model=model_name, + generate_content_config=model_config, + instruction=summary_prompt.PROMPT, + description="Summarizes autotuning results.", + tools=[filesystem_tool_r], + output_key="autotuning_summary", + ) + + +autotune_summary_agent = create_autotune_summary_agent() class CombinedAutotuneAgent(BaseAgent): """Chains autotuning steps and conditionally applies best config.""" - def __init__(self, name: str): - super().__init__(name=name) + planner_agent: Optional[BaseAgent] = None + runner_agent: Optional[BaseAgent] = None + apply_config_agent: Optional[BaseAgent] = None + summary_agent: Optional[BaseAgent] = None + + def __init__( + self, + name: str, + planner_agent: BaseAgent, + runner_agent: BaseAgent, + apply_config_agent: BaseAgent, + summary_agent: BaseAgent, + ): + super().__init__( + name=name, + planner_agent=planner_agent, + runner_agent=runner_agent, + apply_config_agent=apply_config_agent, + summary_agent=summary_agent, + ) async def _run_async_impl( self, ctx: InvocationContext ) -> AsyncGenerator[Event, None]: logging.info(f"[{self.name}] Running AutotunePlannerAgent...") - async for event in autotune_planner_agent.run_async(ctx): + async for event in self.planner_agent.run_async(ctx): yield event logging.info(f"[{self.name}] Running AutotuneRunner...") - async for event in autotune_runner.run_async(ctx): + async for event in self.runner_agent.run_async(ctx): yield event autotune_results = ctx.session.state.get("autotune_results", {}) @@ -209,7 +250,7 @@ async def _run_async_impl( and autotune_results.get("best_time_ms") is not None ): logging.info(f"[{self.name}] Running ApplyBestConfigAgent...") - async for event in apply_best_config_agent.run_async(ctx): + async for event in self.apply_config_agent.run_async(ctx): yield event else: logging.warning( @@ -218,10 +259,22 @@ async def _run_async_impl( ) logging.info(f"[{self.name}] Running AutotuneSummaryAgent...") - async for event in autotune_summary_agent.run_async(ctx): + async for event in self.summary_agent.run_async(ctx): yield event -autotune_agent = CombinedAutotuneAgent(name="AutotuneAgent") +def create_autotune_agent( + model_name: str = MODEL_NAME, +) -> CombinedAutotuneAgent: + return CombinedAutotuneAgent( + name="AutotuneAgent", + planner_agent=create_autotune_planner_agent(model_name), + runner_agent=create_autotune_runner(), + apply_config_agent=create_apply_best_config_agent(model_name), + summary_agent=create_autotune_summary_agent(model_name), + ) + + +autotune_agent = create_autotune_agent() -__all__ = ["autotune_agent"] +__all__ = ["autotune_agent", "create_autotune_agent"] diff --git a/MaxKernel/auto_agent/subagents/kernel_writing/agent.py b/MaxKernel/auto_agent/subagents/kernel_writing/agent.py index bdda2ac..0bcd51d 100644 --- a/MaxKernel/auto_agent/subagents/kernel_writing/agent.py +++ b/MaxKernel/auto_agent/subagents/kernel_writing/agent.py @@ -14,7 +14,6 @@ load_single_kernel_to_state, ) from auto_agent.config import ( - MAX_COMPILATION_RETRIES, get_thinking_planner, model_config, ) @@ -265,130 +264,201 @@ async def _run_async_impl( # Plan-based kernel writing agents (separate, not sequential) # These are called independently by the root orchestrator to allow user interaction between steps -plan_kernel_agent = CustomLlmAgent( - name="PlanKernelAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=get_thinking_planner("high"), - instruction=kernel_planning_prompt.PROMPT, - description="Creates or revises a detailed optimization plan for a Pallas kernel.", - tools=( - [search_api_tool, filesystem_tool_rw, vertex_ai_rag_tool] - if vertex_ai_rag_tool - else [search_api_tool, filesystem_tool_rw] - ), -) + +def create_plan_kernel_agent(model_name: str = MODEL_NAME) -> CustomLlmAgent: + return CustomLlmAgent( + name="PlanKernelAgent", + model=model_name, + generate_content_config=model_config, + planner=get_thinking_planner("high"), + instruction=kernel_planning_prompt.PROMPT, + description="Creates or revises a detailed optimization plan for a Pallas kernel.", + tools=( + [search_api_tool, filesystem_tool_rw, vertex_ai_rag_tool] + if vertex_ai_rag_tool + else [search_api_tool, filesystem_tool_rw] + ), + ) + + +plan_kernel_agent = create_plan_kernel_agent() + # Kernel compilation validation agents # Read file agent for validation - extracts kernel path from user message or state -read_file_for_validation_agent = CustomLlmAgent( - name="ReadFileForValidationAgent", - model=MODEL_NAME, - generate_content_config=model_config, - instruction=read_file_prompt.PROMPT, - description="Reads the kernel file mentioned by the user or from state for validation.", - tools=[filesystem_tool_r], +def create_read_file_for_validation_agent( + model_name: str = MODEL_NAME, +) -> CustomLlmAgent: + return CustomLlmAgent( + name="ReadFileForValidationAgent", + model=model_name, + generate_content_config=model_config, + instruction=read_file_prompt.PROMPT, + description="Reads the kernel file mentioned by the user or from state for validation.", + tools=[filesystem_tool_r], + ) + + +read_file_for_validation_agent = create_read_file_for_validation_agent() + + +def create_fix_kernel_compilation_agent( + model_name: str = MODEL_NAME, +) -> CustomLlmAgent: + return CustomLlmAgent( + name="FixKernelCompilationAgent", + model=model_name, + generate_content_config=model_config, + planner=get_thinking_planner("high"), + instruction=fix_kernel_compilation.PROMPT, + description="Fixes compilation errors in the generated kernel while preserving optimization strategy.", + tools=( + [ + search_api_tool, + filesystem_tool_r, + write_optimized_kernel_tool, + vertex_ai_rag_tool, + ] + if vertex_ai_rag_tool + else [search_api_tool, filesystem_tool_r, write_optimized_kernel_tool] + ), + before_agent_callback=load_kernel_and_plan_to_state, + after_model_callback=extract_fix_summary, + include_contents="none", + ) + + +fix_kernel_compilation_agent = create_fix_kernel_compilation_agent() + + +def create_add_debug_statements_agent( + model_name: str = MODEL_NAME, +) -> CustomLlmAgent: + return CustomLlmAgent( + name="AddDebugStatementsAgent", + model=model_name, + generate_content_config=model_config, + planner=get_thinking_planner("high"), + instruction=add_debug_statements.PROMPT, + description="Adds strategic debugging statements to diagnose persistent compilation issues.", + tools=[filesystem_tool_r, write_optimized_kernel_tool], + before_agent_callback=load_kernel_and_plan_to_state, + include_contents="none", + ) + + +add_debug_statements_agent = create_add_debug_statements_agent() + + +def create_cleanup_debug_statements_agent( + model_name: str = MODEL_NAME, +) -> CustomLlmAgent: + return CustomLlmAgent( + name="CleanupDebugStatementsAgent", + model=model_name, + generate_content_config=model_config, + planner=get_thinking_planner("high"), + instruction=cleanup_debug_statements.PROMPT, + description="Removes debugging statements from successfully compiled kernel.", + tools=[filesystem_tool_r, write_optimized_kernel_tool], + before_agent_callback=load_single_kernel_to_state, + include_contents="none", + ) + + +cleanup_debug_statements_agent = create_cleanup_debug_statements_agent() + + +def create_kernel_compilation_checker_for_validation() -> ( + KernelCompilationChecker +): + return KernelCompilationChecker( + name="KernelCompilationCheckerForValidation", + input_key="kernel_code", + output_key="compilation_results", + before_agent_callback=load_single_kernel_to_state, + ) + + +kernel_compilation_checker_for_validation = ( + create_kernel_compilation_checker_for_validation() ) -fix_kernel_compilation_agent = CustomLlmAgent( - name="FixKernelCompilationAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=get_thinking_planner("high"), - instruction=fix_kernel_compilation.PROMPT, - description="Fixes compilation errors in the generated kernel while preserving optimization strategy.", - tools=( - [ - search_api_tool, - filesystem_tool_r, - write_optimized_kernel_tool, - vertex_ai_rag_tool, - ] - if vertex_ai_rag_tool - else [search_api_tool, filesystem_tool_r, write_optimized_kernel_tool] - ), - before_agent_callback=load_kernel_and_plan_to_state, - after_model_callback=extract_fix_summary, - include_contents="none", -) -add_debug_statements_agent = CustomLlmAgent( - name="AddDebugStatementsAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=get_thinking_planner("high"), - instruction=add_debug_statements.PROMPT, - description="Adds strategic debugging statements to diagnose persistent compilation issues.", - tools=[filesystem_tool_r, write_optimized_kernel_tool], - before_agent_callback=load_kernel_and_plan_to_state, - include_contents="none", -) +def create_kernel_compilation_validation_loop( + model_name: str = MODEL_NAME, +) -> KernelCompilationValidationLoop: + return KernelCompilationValidationLoop( + name="KernelCompilationValidationLoop", + compilation_checker=create_kernel_compilation_checker_for_validation(), + fix_agent=create_fix_kernel_compilation_agent(model_name), + debug_agent=create_add_debug_statements_agent(model_name), + max_retries=6, + ) -cleanup_debug_statements_agent = CustomLlmAgent( - name="CleanupDebugStatementsAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=get_thinking_planner("high"), - instruction=cleanup_debug_statements.PROMPT, - description="Removes debugging statements from successfully compiled kernel.", - tools=[filesystem_tool_r, write_optimized_kernel_tool], - before_agent_callback=load_single_kernel_to_state, - include_contents="none", -) -kernel_compilation_checker_for_validation = KernelCompilationChecker( - name="KernelCompilationCheckerForValidation", - input_key="kernel_code", - output_key="compilation_results", - before_agent_callback=load_single_kernel_to_state, -) +kernel_compilation_validation_loop = create_kernel_compilation_validation_loop() -kernel_compilation_validation_loop = KernelCompilationValidationLoop( - name="KernelCompilationValidationLoop", - compilation_checker=kernel_compilation_checker_for_validation, - fix_agent=fix_kernel_compilation_agent, - debug_agent=add_debug_statements_agent, - max_retries=MAX_COMPILATION_RETRIES, -) -kernel_compilation_summary_agent = CustomLlmAgent( - name="KernelCompilationSummaryAgent", - model=MODEL_NAME, - generate_content_config=model_config, - instruction=kernel_compilation_summary.PROMPT, - description="Summarizes kernel compilation validation results with full trace on failure.", - include_contents="none", -) +def create_kernel_compilation_summary_agent( + model_name: str = MODEL_NAME, +) -> CustomLlmAgent: + return CustomLlmAgent( + name="KernelCompilationSummaryAgent", + model=model_name, + generate_content_config=model_config, + instruction=kernel_compilation_summary.PROMPT, + description="Summarizes kernel compilation validation results with full trace on failure.", + include_contents="none", + ) + + +kernel_compilation_summary_agent = create_kernel_compilation_summary_agent() + # Standalone validation orchestration agent (invoked by root when user requests validation) -validate_kernel_compilation_agent = ValidateKernelCompilationAgent( - name="ValidateKernelCompilationAgent", - read_file_agent=read_file_for_validation_agent, - validation_loop_agent=kernel_compilation_validation_loop, - cleanup_agent=cleanup_debug_statements_agent, - summary_agent=kernel_compilation_summary_agent, - description="Validates kernel compilation with automatic error fixing, debugging, and provides summary. Invoked when user requests validation.", -) +def create_validate_kernel_compilation_agent( + model_name: str = MODEL_NAME, +) -> ValidateKernelCompilationAgent: + return ValidateKernelCompilationAgent( + name="ValidateKernelCompilationAgent", + read_file_agent=create_read_file_for_validation_agent(model_name), + validation_loop_agent=create_kernel_compilation_validation_loop(model_name), + cleanup_agent=create_cleanup_debug_statements_agent(model_name), + summary_agent=create_kernel_compilation_summary_agent(model_name), + description="Validates kernel compilation with automatic error fixing, debugging, and provides summary. Invoked when user requests validation.", + ) + + +validate_kernel_compilation_agent = create_validate_kernel_compilation_agent() + # Implementation agent - implements kernel -implement_kernel_agent = CustomLlmAgent( - name="ImplementKernelAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=get_thinking_planner("high"), - instruction=kernel_implementation_prompt.PROMPT, - description="Implements the optimized Pallas kernel following the plan.", - tools=( - [ - search_api_tool, - filesystem_tool_r, - write_optimized_kernel_tool, - vertex_ai_rag_tool, - ] - if vertex_ai_rag_tool - else [search_api_tool, filesystem_tool_r, write_optimized_kernel_tool] - ), -) +def create_implement_kernel_agent( + model_name: str = MODEL_NAME, +) -> CustomLlmAgent: + return CustomLlmAgent( + name="ImplementKernelAgent", + model=model_name, + generate_content_config=model_config, + planner=get_thinking_planner("high"), + instruction=kernel_implementation_prompt.PROMPT, + description="Implements the optimized Pallas kernel following the plan.", + tools=( + [ + search_api_tool, + filesystem_tool_r, + write_optimized_kernel_tool, + vertex_ai_rag_tool, + ] + if vertex_ai_rag_tool + else [search_api_tool, filesystem_tool_r, write_optimized_kernel_tool] + ), + ) + + +implement_kernel_agent = create_implement_kernel_agent() + __all__ = [ "KernelCompilationValidationLoop", diff --git a/MaxKernel/auto_agent/subagents/kernel_writing/prompts/kernel_planning_prompt.py b/MaxKernel/auto_agent/subagents/kernel_writing/prompts/kernel_planning_prompt.py index a6ce62c..3451f8b 100644 --- a/MaxKernel/auto_agent/subagents/kernel_writing/prompts/kernel_planning_prompt.py +++ b/MaxKernel/auto_agent/subagents/kernel_writing/prompts/kernel_planning_prompt.py @@ -14,12 +14,13 @@ Identify whether you are creating a **NEW plan** or performing a **REVISION**. * **NEW Plan:** - * No existing plan file is mentioned. - * User provides a kernel filename or pastes code to optimize. + * The plan path (`{kernel_plan_path?}`) is either not provided, or the file at that path is empty/missing. + * Your primary input is the base kernel code that needs optimization. * **REVISION:** - * An existing plan path is provided: `{kernel_plan_path?}`. - * You receive results from other subagents (compilation status, test results, profiling summary) indicating issues. - + * The plan file at `{kernel_plan_path?}` already contains an existing plan. + * The `{optimized_kernel_path?}` already contains an optimized kernel implementation. + * You receive execution results (compilation status, test results, profiling summary) from a previous attempt that need to be addressed. + ### Step 2: Gather Context (Conditional) **For NEW Plans:** @@ -32,7 +33,8 @@ **For REVISIONS:** 1. **Read current plan:** Use `filesystem_tool` to read the existing plan at `{kernel_plan_path?}`. -2. **Review execution results:** Analyze the following to identify what needs improvement: +2. **Read optimized kernel:** Use `filesystem_tool` to read the `{optimized_kernel_path?}`. +3. **Review execution results:** Analyze the following to identify what needs improvement: * Compilation Status: `{kernel_compilation_status?}` * Test Results: `{test_results?}` *. Autotune Summary: `{autotuning_summary?}` @@ -48,7 +50,8 @@ ## 1. Current Kernel Analysis - Brief description of what the kernel does -- Current implementation approach +- Current implementation approach + - Analyze the `{base_kernel_path}` for NEW plans, or the `{optimized_kernel_path?}` for REVISIONS - Identified performance bottlenecks or issues ## 2. Optimization Strategy diff --git a/MaxKernel/auto_agent/subagents/pipeline_agent.py b/MaxKernel/auto_agent/subagents/pipeline_agent.py index 35c68f0..221b1bb 100644 --- a/MaxKernel/auto_agent/subagents/pipeline_agent.py +++ b/MaxKernel/auto_agent/subagents/pipeline_agent.py @@ -79,6 +79,7 @@ async def _run_async_impl( self._clear_iteration_metrics(ctx) if self._should_end_at_step(ctx, iteration, "plan"): + yield self._create_history_event(ctx) iteration += 1 continue @@ -87,6 +88,7 @@ async def _run_async_impl( async for event in self.implement_agent.run_async(ctx): yield event if self._should_end_at_step(ctx, iteration, "implement"): + yield self._create_history_event(ctx) iteration += 1 continue @@ -106,10 +108,12 @@ async def _run_async_impl( self._save_iteration_files_and_snapshot( ctx, iteration, step_name="validate" ) + yield self._create_history_event(ctx) iteration += 1 continue if self._should_end_at_step(ctx, iteration, "validate"): + yield self._create_history_event(ctx) iteration += 1 continue @@ -146,10 +150,12 @@ async def _run_async_impl( self._save_iteration_files_and_snapshot( ctx, iteration, step_name="test_run" ) + yield self._create_history_event(ctx) iteration += 1 continue if self._should_end_at_step(ctx, iteration, "test_run"): + yield self._create_history_event(ctx) iteration += 1 continue @@ -158,6 +164,7 @@ async def _run_async_impl( async for event in self.autotune_agent.run_async(ctx): yield event if self._should_end_at_step(ctx, iteration, "autotune"): + yield self._create_history_event(ctx) iteration += 1 continue @@ -166,17 +173,13 @@ async def _run_async_impl( async for event in self.profile_agent.run_async(ctx): yield event if self._should_end_at_step(ctx, iteration, "profile"): + yield self._create_history_event(ctx) iteration += 1 continue self._save_iteration_files_and_snapshot(ctx, iteration) - yield Event( - author=self.name, - actions=EventActions( - state_delta={"history": ctx.session.state.get("history", [])} - ), - ) + yield self._create_history_event(ctx) # Step 7: Check if improvement is needed # needs_improvement = ctx.session.state.get("needs_improvement", False) @@ -210,6 +213,14 @@ async def _run_async_impl( ), ) + def _create_history_event(self, ctx: InvocationContext) -> Event: + return Event( + author=self.name, + actions=EventActions( + state_delta={"history": ctx.session.state.get("history", [])} + ), + ) + def _should_end_at_step( self, ctx: InvocationContext, iteration: int, step_name: str ) -> bool: @@ -251,6 +262,7 @@ def _record_history_snapshot(self, ctx: InvocationContext, iteration: int): ), "test_status": ctx.session.state.get("test_results", {}), "latency_ms": latency, + "autotuning_summary": ctx.session.state.get("autotuning_summary", ""), "profiling_summary": ctx.session.state.get("profiling_summary", ""), } diff --git a/MaxKernel/auto_agent/subagents/profiling/agent.py b/MaxKernel/auto_agent/subagents/profiling/agent.py index cf31c42..d7de792 100644 --- a/MaxKernel/auto_agent/subagents/profiling/agent.py +++ b/MaxKernel/auto_agent/subagents/profiling/agent.py @@ -27,35 +27,54 @@ ) from auto_agent.tools.tools import vertex_ai_rag_tool + # Profiling script generation agent - writes profiling script to file -generate_profiling_script_agent = CustomLlmAgent( - name="GenerateProfilingScriptAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=get_thinking_planner("high"), - instruction=gen_profiling_script.PROMPT, - description="Generates a profiling script to identify performance bottlenecks in the kernel code and writes it to a file.", - tools=[filesystem_tool_r, write_profiling_script_tool], - before_agent_callback=load_single_kernel_to_state, -) +def create_generate_profiling_script_agent( + model_name: str = MODEL_NAME, +) -> CustomLlmAgent: + return CustomLlmAgent( + name="GenerateProfilingScriptAgent", + model=model_name, + generate_content_config=model_config, + planner=get_thinking_planner("medium"), + instruction=gen_profiling_script.PROMPT, + description="Generates a profiling script to identify performance bottlenecks in the kernel code and writes it to a file.", + tools=[filesystem_tool_r, write_profiling_script_tool], + before_agent_callback=load_single_kernel_to_state, + ) + + +generate_profiling_script_agent = create_generate_profiling_script_agent() + # Read profiling script agent - loads the generated profiling script file contents into state -read_profiling_script_agent = CustomLlmAgent( - name="ReadProfilingScriptAgent", - model=MODEL_NAME, - generate_content_config=model_config, - instruction=read_profiling_script_prompt.PROMPT, - description="Loads the generated profiling script file contents from disk into memory for execution.", - before_agent_callback=load_profiling_script_to_state, - include_contents="none", -) +def create_read_profiling_script_agent( + model_name: str = MODEL_NAME, +) -> CustomLlmAgent: + return CustomLlmAgent( + name="ReadProfilingScriptAgent", + model=model_name, + generate_content_config=model_config, + instruction=read_profiling_script_prompt.PROMPT, + description="Loads the generated profiling script file contents from disk into memory for execution.", + before_agent_callback=load_profiling_script_to_state, + include_contents="none", + ) + + +read_profiling_script_agent = create_read_profiling_script_agent() + # Profiling execution agent -eval_profile_agent = KernelProfiler( - name="ProfileEvalAgent", - input_key="profiling_script", - output_key="profiling_results", -) +def create_eval_profile_agent() -> KernelProfiler: + return KernelProfiler( + name="ProfileEvalAgent", + input_key="profiling_script", + output_key="profiling_results", + ) + + +eval_profile_agent = create_eval_profile_agent() class SummarizeProfileAgent(CustomLlmAgent): @@ -105,38 +124,50 @@ async def _run_async_impl( # Profiling summary agent -summarize_profile_agent = SummarizeProfileAgent( - name="SummarizeProfileAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=get_thinking_planner("high"), - instruction=analyze_profile_prompt.PROMPT, - description=( - "Summarizes the profiling results of the kernel and performs deep" - " analysis using offline XProf tools." - ), - 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] if vertex_ai_rag_tool else []), -) +def create_summarize_profile_agent( + model_name: str = MODEL_NAME, +) -> SummarizeProfileAgent: + return SummarizeProfileAgent( + name="SummarizeProfileAgent", + model=model_name, + generate_content_config=model_config, + planner=get_thinking_planner("high"), + instruction=analyze_profile_prompt.PROMPT, + description=( + "Summarizes the profiling results of the kernel and performs deep" + " analysis using offline XProf tools." + ), + 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] if vertex_ai_rag_tool else []), + ) + + +summarize_profile_agent = create_summarize_profile_agent() + # Main profiling orchestrator agent -profile_agent = SequentialAgent( - name="ProfileAgentOrchestrator", - sub_agents=[ - generate_profiling_script_agent, - read_profiling_script_agent, - eval_profile_agent, - summarize_profile_agent, - ], - description="Profiles the Pallas kernel to identify performance bottlenecks.", -) +def create_profile_agent(model_name: str = MODEL_NAME) -> SequentialAgent: + return SequentialAgent( + name="ProfileAgentOrchestrator", + sub_agents=[ + create_generate_profiling_script_agent(model_name), + create_read_profiling_script_agent(model_name), + create_eval_profile_agent(), + create_summarize_profile_agent(model_name), + ], + description="Profiles the Pallas kernel to identify performance bottlenecks.", + ) + + +profile_agent = create_profile_agent() + __all__ = [ "profile_agent", diff --git a/MaxKernel/auto_agent/subagents/testing/agent.py b/MaxKernel/auto_agent/subagents/testing/agent.py index d9088ea..dff2807 100644 --- a/MaxKernel/auto_agent/subagents/testing/agent.py +++ b/MaxKernel/auto_agent/subagents/testing/agent.py @@ -856,123 +856,195 @@ async def _run_async_impl( # Validation Summary Agent -validation_summary_agent = CustomLlmAgent( - name="ValidationSummaryAgent", - model=MODEL_NAME, - generate_content_config=model_config, - instruction=validation_summary.PROMPT, - description="Summarizes validation results and provides next steps to the user.", -) +def create_validation_summary_agent( + model_name: str = MODEL_NAME, +) -> CustomLlmAgent: + return CustomLlmAgent( + name="ValidationSummaryAgent", + model=model_name, + generate_content_config=model_config, + instruction=validation_summary.PROMPT, + description="Summarizes validation results and provides next steps to the user.", + ) + + +validation_summary_agent = create_validation_summary_agent() + # Test file generation agent -generate_test_file_agent = CustomLlmAgent( - name="GenerateTestFileAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=get_thinking_planner("high"), - instruction=gen_test_file.PROMPT, - description="Generates a comprehensive pytest test file.", - tools=( - [ - search_api_tool, - filesystem_tool_r, - write_test_file_tool, - vertex_ai_rag_tool, - ] - if vertex_ai_rag_tool - else [search_api_tool, filesystem_tool_r, write_test_file_tool] - ), -) +def create_generate_test_file_agent( + model_name: str = MODEL_NAME, +) -> CustomLlmAgent: + return CustomLlmAgent( + name="GenerateTestFileAgent", + model=model_name, + generate_content_config=model_config, + planner=get_thinking_planner("high"), + instruction=gen_test_file.PROMPT, + description="Generates a comprehensive pytest test file.", + tools=( + [ + search_api_tool, + filesystem_tool_r, + write_test_file_tool, + vertex_ai_rag_tool, + ] + if vertex_ai_rag_tool + else [search_api_tool, filesystem_tool_r, write_test_file_tool] + ), + ) + + +generate_test_file_agent = create_generate_test_file_agent() + # Validation agents -syntax_validation_agent = SyntaxValidationAgent( - name="SyntaxValidationAgent", - input_key="test_file_path", - output_key="syntax_validation", -) +def create_syntax_validation_agent() -> SyntaxValidationAgent: + return SyntaxValidationAgent( + name="SyntaxValidationAgent", + input_key="test_file_path", + output_key="syntax_validation", + ) -import_validation_agent = ImportValidationAgent( - name="ImportValidationAgent", - input_key="test_file_path", - output_key="import_validation", -) -structure_validation_agent = TestStructureValidationAgent( - name="TestStructureValidationAgent", - input_key="test_file_path", - output_key="structure_validation", -) +syntax_validation_agent = create_syntax_validation_agent() -mock_execution_validation_agent = MockTestExecutionAgent( - name="MockTestExecutionAgent", - input_key="test_file_path", - output_key="mock_execution_validation", -) -fix_test_script_agent = CustomLlmAgent( - name="FixTestScriptAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=get_thinking_planner("high"), - instruction=fix_test_script.PROMPT, - description="Fixes validation errors in the generated test file.", - tools=[filesystem_tool_r, write_test_file_tool, search_api_tool], - include_contents="none", -) +def create_import_validation_agent() -> ImportValidationAgent: + return ImportValidationAgent( + name="ImportValidationAgent", + input_key="test_file_path", + output_key="import_validation", + ) + + +import_validation_agent = create_import_validation_agent() + + +def create_structure_validation_agent() -> TestStructureValidationAgent: + return TestStructureValidationAgent( + name="TestStructureValidationAgent", + input_key="test_file_path", + output_key="structure_validation", + ) + + +structure_validation_agent = create_structure_validation_agent() + + +def create_mock_execution_validation_agent() -> MockTestExecutionAgent: + return MockTestExecutionAgent( + name="MockTestExecutionAgent", + input_key="test_file_path", + output_key="mock_execution_validation", + ) + + +mock_execution_validation_agent = create_mock_execution_validation_agent() + + +def create_fix_test_script_agent( + model_name: str = MODEL_NAME, +) -> CustomLlmAgent: + return CustomLlmAgent( + name="FixTestScriptAgent", + model=model_name, + generate_content_config=model_config, + planner=get_thinking_planner("high"), + instruction=fix_test_script.PROMPT, + description="Fixes validation errors in the generated test file.", + tools=[filesystem_tool_r, write_test_file_tool, search_api_tool], + include_contents="none", + ) + + +fix_test_script_agent = create_fix_test_script_agent() + # Validation loop agent -validation_loop_agent = TestValidationLoopAgent( - name="TestValidationLoopAgent", - syntax_agent=syntax_validation_agent, - import_agent=import_validation_agent, - structure_agent=structure_validation_agent, - mock_execution_agent=mock_execution_validation_agent, - fix_agent=fix_test_script_agent, - max_retries=6, -) +def create_validation_loop_agent( + model_name: str = MODEL_NAME, +) -> TestValidationLoopAgent: + return TestValidationLoopAgent( + name="TestValidationLoopAgent", + syntax_agent=create_syntax_validation_agent(), + import_agent=create_import_validation_agent(), + structure_agent=create_structure_validation_agent(), + mock_execution_agent=create_mock_execution_validation_agent(), + fix_agent=create_fix_test_script_agent(model_name), + max_retries=6, + ) + + +validation_loop_agent = create_validation_loop_agent() + # Sequential agent that generates and validates test files -validated_test_generation_agent = SequentialAgent( - name="ValidatedTestGenerationAgent", - sub_agents=[ - generate_test_file_agent, - validation_loop_agent, - validation_summary_agent, - ], - description="Generates a validated pytest test file with automatic iterative error detection and fixing.", -) +def create_validated_test_generation_agent( + model_name: str = MODEL_NAME, +) -> SequentialAgent: + return SequentialAgent( + name="ValidatedTestGenerationAgent", + sub_agents=[ + create_generate_test_file_agent(model_name), + create_validation_loop_agent(model_name), + create_validation_summary_agent(model_name), + ], + description="Generates a validated pytest test file with automatic iterative error detection and fixing.", + ) + + +validated_test_generation_agent = create_validated_test_generation_agent() + # Test execution agents -run_tests_agent = TestRunner( - name="RunTestsAgent", - input_key="test_file_path", - output_key="test_results", -) +def create_run_tests_agent() -> TestRunner: + return TestRunner( + name="RunTestsAgent", + input_key="test_file_path", + output_key="test_results", + ) -summarize_test_results_agent = CustomLlmAgent( - name="SummarizeTestResultsAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=get_thinking_planner("high"), - instruction=summarize_test_results_prompt.PROMPT, - description="Analyzes pytest test results and provides recommendations.", - tools=( - [search_api_tool, vertex_ai_rag_tool] - if vertex_ai_rag_tool - else [search_api_tool] - ), - output_key="test_summary", - include_contents="none", -) -unified_test_agent = SequentialAgent( - name="UnifiedTestAgent", - sub_agents=[ - run_tests_agent, - summarize_test_results_agent, - ], - description="Executes the generated pytest test file and provides a comprehensive summary.", -) +run_tests_agent = create_run_tests_agent() + + +def create_summarize_test_results_agent( + model_name: str = MODEL_NAME, +) -> CustomLlmAgent: + return CustomLlmAgent( + name="SummarizeTestResultsAgent", + model=model_name, + generate_content_config=model_config, + planner=get_thinking_planner("high"), + instruction=summarize_test_results_prompt.PROMPT, + description="Analyzes pytest test results and provides recommendations.", + tools=( + [search_api_tool, vertex_ai_rag_tool] + if vertex_ai_rag_tool + else [search_api_tool] + ), + output_key="test_summary", + include_contents="none", + ) + + +summarize_test_results_agent = create_summarize_test_results_agent() + + +def create_unified_test_agent(model_name: str = MODEL_NAME) -> SequentialAgent: + return SequentialAgent( + name="UnifiedTestAgent", + sub_agents=[ + create_run_tests_agent(), + create_summarize_test_results_agent(model_name), + ], + description="Executes the generated pytest test file and provides a comprehensive summary.", + ) + + +unified_test_agent = create_unified_test_agent() + __all__ = [ "TestRunner", diff --git a/MaxKernel/auto_agent/subagents/testing/prompts/gen_test_file.py b/MaxKernel/auto_agent/subagents/testing/prompts/gen_test_file.py index 6e2a141..3090045 100644 --- a/MaxKernel/auto_agent/subagents/testing/prompts/gen_test_file.py +++ b/MaxKernel/auto_agent/subagents/testing/prompts/gen_test_file.py @@ -82,11 +82,12 @@ - **Note**: During validation, the optimized kernel import will be temporarily disabled to verify the test structure works with baseline only 3. **TestPerformance class**: Tests that benchmark performance - - If the base kernel file (`{base_kernel_path?}`) contains a function to generate inputs (e.g., `get_inputs` or similar). If it is defined, you should directly copy and reuse it in the test file. - - Compare execution time between base and optimized kernels + - If the base kernel file (`{base_kernel_path?}`) contains a function to generate inputs (e.g., `get_inputs`), you must directly copy and reuse it. + - **CRITICAL**: If the input generation function returns MULTIPLE configurations (e.g., a list of inputs), you MUST select exactly ONE configuration (the last one) to run the performance benchmark on. Do NOT use `@pytest.mark.parametrize` for the performance test to ensure we get a single, consistent baseline measurement. + - Compare execution time between base and optimized kernels on this single configuration - Include warmup runs before timing - Use .block_until_ready() for accurate JAX timing - - Run 10 iterations for each benchmark to get reliable timing measurements + - Run 10 iterations for the benchmark to get reliable timing measurements ## Requirements @@ -136,9 +137,10 @@ if 'optimized_kernel' not in globals(): optimized_kernel = base_kernel -def report_perf_metrics(execution_time_ms): +def report_perf_metrics(execution_time_ms, speedup): import sys sys.__stdout__.write("PERF_METRICS: " + str(execution_time_ms) + "\n") + sys.__stdout__.write("SPEEDUP: " + str(speedup) + "\n") class TestCompilation: def test_optimized_kernel_compiles(self): @@ -161,7 +163,7 @@ def test_performance_comparison(self): pass if __name__ == "__main__": - sys.exit(pytest.main([__file__, "-v", "-s", "--tb=short"])) + sys.exit(pytest.main([__file__, "-v", "--tb=short"])) ``` ### Important Guidelines @@ -188,7 +190,7 @@ def test_performance_comparison(self): - Report speedup ratios - Include warmup iterations - Use multiple runs for statistical stability - - **Structured Output**: At the end of the performance test, you MUST call the `report_perf_metrics(execution_time_ms)` helper function provided in the template to report the final average execution time of the optimized kernel. Do not use standard `print()` for this. + - **Structured Output**: At the end of the performance test, you MUST call the `report_perf_metrics(execution_time_ms, speedup)` helper function provided in the template to report the final average execution time and the speedup of the optimized kernel over the base kernel. Do not use standard `print()` for this. ## Output Format When you have generated the test file: diff --git a/MaxKernel/auto_search/graph.py b/MaxKernel/auto_search/graph.py index fe1c39b..8c59d2e 100644 --- a/MaxKernel/auto_search/graph.py +++ b/MaxKernel/auto_search/graph.py @@ -13,6 +13,7 @@ class EvaluationResult(BaseModel): correct: bool = False latency_ms: Optional[float] = None profiling_summary: Optional[str] = None + autotuning_summary: Optional[str] = None compilation_error: Optional[str] = None test_error: Optional[str] = None diff --git a/MaxKernel/auto_search/worker.py b/MaxKernel/auto_search/worker.py index a5f9adc..ddbb6fb 100644 --- a/MaxKernel/auto_search/worker.py +++ b/MaxKernel/auto_search/worker.py @@ -179,6 +179,7 @@ def _process_results( test_error = test_status.get("output") if not correct else None latency_ms = best_run.get("latency_ms") + autotuning_summary = best_run.get("autotuning_summary") profiling_summary = best_run.get("profiling_summary") # Read optimized code if it exists @@ -208,6 +209,7 @@ def _process_results( compiled=compiled, correct=correct, latency_ms=latency_ms, + autotuning_summary=autotuning_summary, profiling_summary=profiling_summary, compilation_error=compilation_error, test_error=test_error,