From 7315b62de9270236b087957944f7f5c7cff972ff Mon Sep 17 00:00:00 2001 From: shangkunwang Date: Thu, 2 Jul 2026 17:20:46 +0000 Subject: [PATCH 01/11] feat: refactor auto agent client to use InMemorySessionService and batch agent execution to use asyncio --- .../agent_client/auto_agent_client.py | 135 +++++----- .../agent_client/run_batch_agent_call.py | 236 ++++++++++-------- 2 files changed, 194 insertions(+), 177 deletions(-) diff --git a/MaxKernel/auto_agent/agent_client/auto_agent_client.py b/MaxKernel/auto_agent/agent_client/auto_agent_client.py index e1eb2fb..c80cfef 100644 --- a/MaxKernel/auto_agent/agent_client/auto_agent_client.py +++ b/MaxKernel/auto_agent/agent_client/auto_agent_client.py @@ -1,16 +1,25 @@ import argparse +import asyncio import json import logging +from typing import Any, Optional -import requests +# Load environment variables from .env file if available +try: + from dotenv import load_dotenv -REQUEST_TIMEOUT = 60 * 60 * 5 + load_dotenv() +except ImportError: + logging.warning( + "dotenv not installed, skipping loading environment variables" + ) + +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 -# Configure logging -logging.basicConfig( - level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" -) logger = logging.getLogger(__name__) @@ -18,7 +27,6 @@ class AutoAgentClient: user_id: str session_id: str query: str - base_url: str app_name: str = "auto_agent" def __init__( @@ -26,54 +34,35 @@ def __init__( user_id: str, session_id: str, query: str, - base_url: str = "http://localhost:8000", + agent: Optional[Any] = None, ): self.user_id = user_id self.session_id = session_id self.query = query - self.base_url = base_url - - def create_session(self): - session_url = f"{self.base_url}/apps/{self.app_name}/users/{self.user_id}/sessions/{self.session_id}" - response = requests.post( - session_url, headers={"Content-Type": "application/json"} - ) - return response - - def send_query(self): - run_url = f"{self.base_url}/run" - payload = { - "appName": self.app_name, - "userId": self.user_id, - "sessionId": self.session_id, - "newMessage": {"role": "user", "parts": [{"text": self.query}]}, - } - response = requests.post( - run_url, - headers={"Content-Type": "application/json"}, - data=json.dumps(payload), - timeout=REQUEST_TIMEOUT, + self.agent = agent or root_agent + self.session_service = InMemorySessionService() + self.session = None + + async def create_session(self) -> None: + self.session = await self.session_service.create_session( + app_name=self.app_name, + user_id=self.user_id, + session_id=self.session_id, ) - return response def _get_session_data(self) -> dict: - session_url = f"{self.base_url}/apps/{self.app_name}/users/{self.user_id}/sessions/{self.session_id}" - response = requests.get( - session_url, headers={"Content-Type": "application/json"} + if not self.session: + raise ValueError("Session has not been created yet.") + + return json.loads( + self.session.model_dump_json(by_alias=True, exclude_none=True) ) - if response.status_code != 200: - raise ValueError(f"Failed to get state: {response.status_code}") - try: - return response.json() - except json.JSONDecodeError: - raise ValueError("Failed to decode JSON response from server") - def get_state(self, key: str = None): - state_data = self._get_session_data() - state = state_data.get("state") - if not state: - raise ValueError("State not found in response") + def get_state(self, key: Optional[str] = None) -> Any: + if not self.session: + raise ValueError("Session has not been created yet.") + state = self.session.state if key is None: return state @@ -81,17 +70,36 @@ def get_state(self, key: str = None): raise ValueError(f"Key '{key}' not found in state") return state.get(key) + async def run_async(self) -> None: + if not self.session: + await self.create_session() + + runner = Runner( + app_name=self.app_name, + agent=self.agent, + session_service=self.session_service, + ) -def run_agent(client: AutoAgentClient): - # Create session - session_response = client.create_session() - if session_response.status_code != 200: - raise ValueError(f"Failed to create session: {session_response.text}") + new_message = Content(parts=[Part(text=self.query)]) - # Send query - query_response = client.send_query() - if query_response.status_code != 200: - raise ValueError(f"ADK server returned error: {query_response.text}") + logger.info(f"Starting in-process agent run for session {self.session_id}") + try: + async for event in runner.run_async( + user_id=self.user_id, + session_id=self.session_id, + new_message=new_message, + ): + pass + logger.info( + f"Finished in-process agent run for session {self.session_id}" + ) + finally: + # Retrieve the updated session from the session service, even if the run crashed + self.session = await self.session_service.get_session( + app_name=self.app_name, + user_id=self.user_id, + session_id=self.session_id, + ) def read_query_from_file(file_path: str) -> str: @@ -113,11 +121,6 @@ def main(): default="client_query.txt", help="File containing the query to send", ) - parser.add_argument( - "--already-generated", - action="store_true", - help="Whether to use an already generated script", - ) args = parser.parse_args() user_id = args.user_id @@ -128,16 +131,10 @@ def main(): # Create client instance client = AutoAgentClient(user_id, session_id, query) - if not args.already_generated: - # Generate and return script - logger.info( - f"Generating script for user {user_id} in session {session_id} with query: {query}" - ) - run_agent(client) - else: - logger.info( - f"Using already generated script for user {user_id} in session {session_id}" - ) + logger.info( + f"Generating script for user {user_id} in session {session_id} with query: {query}" + ) + asyncio.run(client.run_async()) if __name__ == "__main__": 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 f6927e3..bec07bb 100644 --- a/MaxKernel/auto_agent/agent_client/run_batch_agent_call.py +++ b/MaxKernel/auto_agent/agent_client/run_batch_agent_call.py @@ -1,20 +1,13 @@ import argparse +import asyncio import json import logging import os -import random import time -from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import Any -from auto_agent.agent_client.auto_agent_client import ( - AutoAgentClient, - run_agent, -) +from auto_agent.agent_client.auto_agent_client import AutoAgentClient -# Configure logging -logging.basicConfig( - level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" -) logger = logging.getLogger(__name__) @@ -28,9 +21,6 @@ def parse_args(): parser.add_argument( "--num_concurrent", type=int, default=1, help="Number of concurrent calls" ) - parser.add_argument( - "--port", type=int, default=8000, help="Port for AutoAgentClient" - ) parser.add_argument( "--max_retries", type=int, @@ -58,7 +48,7 @@ def get_optimized_kernel_path(session_file_path: str): return optimized_kernel_path -def save_session_to_file(client: AutoAgentClient, file_path: str): +def save_session_to_file(client: Any, file_path: str): """Fetches full session data and saves it to a local JSON file.""" session_data = client._get_session_data() with open(file_path, "w") as f: @@ -66,94 +56,89 @@ def save_session_to_file(client: AutoAgentClient, file_path: str): logger.info(f"Saved session data to {file_path}") -def process_problem( - problem_id: str, data_dir: str, port: int, max_retries: int +async def process_problem( + problem_id: str, + data_dir: str, + max_retries: int, + sem: asyncio.Semaphore, ): - problem_dir = os.path.join(data_dir, problem_id) - reference_file = os.path.join(problem_dir, "reference.py") + async with sem: + problem_dir = os.path.join(data_dir, problem_id) + reference_file = os.path.join(problem_dir, "reference.py") - if not os.path.exists(reference_file): - logger.error(f"reference.py not found in {problem_dir}") - return problem_id, "Failed: reference.py missing" + if not os.path.exists(reference_file): + logger.error(f"reference.py not found in {problem_dir}") + return problem_id, "Failed: reference.py missing" - with open(reference_file, "r") as f: - reference_code = f.read() + with open(reference_file, "r") as f: + reference_code = f.read() - query = ( - "Optimize the code below for peak performance with pallas kernel\n\n" - + reference_code - ) - - # Retry logic - for attempt in range(1, max_retries + 1): - logger.info(f"Processing {problem_id} - Attempt {attempt}/{max_retries}") - user_id = "user_0" - session_id = f"session_{problem_id}_attempt_{attempt}_{int(time.time())}" - - # Add random jitter to avoid SQLite database lock contention - jitter = random.uniform(0.1, 2.0) - logger.info(f"Sleeping for {jitter:.2f}s (jitter) to avoid DB lock.") - time.sleep(jitter) - - client = AutoAgentClient( - user_id=user_id, - session_id=session_id, - query=query, - base_url=f"http://localhost:{port}", + query = ( + "Optimize the code below for peak performance with pallas kernel\n\n" + + reference_code ) - session_file = os.path.join(problem_dir, f"session_attempt_{attempt}.json") - try: - run_agent(client) - - # Save session file - save_session_to_file(client, session_file) - - # Save optimized code - optimized_path = get_optimized_kernel_path(session_file) - if not optimized_path: - logger.error( - f"Optimized kernel path not found in state for {problem_id}" - ) - continue # Try next attempt - - # Read the code from the path - with open(optimized_path, "r") as f: - optimized_code = f.read() - - # Paste the code to the problem folder in the dataset - output_file = os.path.join(problem_dir, "optimized.py") - with open(output_file, "w") as f: - f.write(optimized_code) - - logger.info(f"Successfully saved optimized kernel to {output_file}") - return problem_id, "Success" - - except Exception as e: - logger.error(f"Error on attempt {attempt} for {problem_id}: {e}") + # Retry logic + for attempt in range(1, max_retries + 1): + logger.info(f"Processing {problem_id} - Attempt {attempt}/{max_retries}") + user_id = "user_0" + session_id = f"session_{problem_id}_attempt_{attempt}_{int(time.time())}" + + client = AutoAgentClient( + user_id=user_id, + session_id=session_id, + query=query, + ) + + session_file = os.path.join( + problem_dir, f"session_attempt_{attempt}.json" + ) try: + # Run in-process asynchronously + await client.run_async() + + # Save session file save_session_to_file(client, session_file) - logger.info( - f"Saved session file for {problem_id} after attempt failure" - ) - except Exception as se: - logger.error(f"Failed to save session file after attempt failure: {se}") - if attempt == max_retries: - return problem_id, f"Failed after {max_retries} attempts: {e}" - return problem_id, "Failed" + # Save optimized code + optimized_path = get_optimized_kernel_path(session_file) + if not optimized_path: + logger.error( + f"Optimized kernel path not found in state for {problem_id}" + ) + continue # Try next attempt + # Read the code from the path + with open(optimized_path, "r") as f: + optimized_code = f.read() -def main(): - args = parse_args() + # Paste the code to the problem folder in the dataset + output_file = os.path.join(problem_dir, "optimized.py") + with open(output_file, "w") as f: + f.write(optimized_code) - if args.log_file: - file_handler = logging.FileHandler(args.log_file) - file_handler.setFormatter( - logging.Formatter("%(asctime)s - %(levelname)s - %(message)s") - ) - logging.getLogger().addHandler(file_handler) + logger.info(f"Successfully saved optimized kernel to {output_file}") + return problem_id, "Success" + except Exception as e: + logger.error(f"Error on attempt {attempt} for {problem_id}: {e}") + try: + save_session_to_file(client, session_file) + logger.info( + f"Saved session file for {problem_id} after attempt failure" + ) + except Exception as se: + logger.error( + f"Failed to save session file after attempt failure: {se}" + ) + if attempt == max_retries: + return problem_id, f"Failed after {max_retries} attempts: {e}" + + return problem_id, "Failed" + + +async def main_async(args): + """Async main function coordinating the concurrent problem runs.""" if not os.path.exists(args.data_dir): logger.error(f"Data directory not found: {args.data_dir}") return @@ -170,25 +155,60 @@ def main(): f"Concurrent workers: {args.num_concurrent}, Max retries: {args.max_retries}" ) - with ThreadPoolExecutor(max_workers=args.num_concurrent) as executor: - future_to_problem = { - executor.submit( - process_problem, problem, args.data_dir, args.port, args.max_retries - ): problem - for problem in problems - } - - completed = 0 - for future in as_completed(future_to_problem): - problem = future_to_problem[future] - try: - prob_id, status = future.result() - completed += 1 - logger.info( - f"[{completed}/{len(problems)}] Problem {prob_id}: {status}" - ) - except Exception as e: - logger.error(f"Problem {problem} raised an exception: {e}") + # Semaphore to limit the number of concurrent active sessions + sem = asyncio.Semaphore(args.num_concurrent) + + tasks = [ + process_problem(problem, args.data_dir, args.max_retries, sem) + for problem in problems + ] + + completed = 0 + for future in asyncio.as_completed(tasks): + try: + prob_id, status = await future + completed += 1 + logger.info(f"[{completed}/{len(problems)}] Problem {prob_id}: {status}") + except Exception as e: + logger.error(f"A task raised an exception: {e}") + + +def main(): + args = parse_args() + + # Configure logging dynamically + log_format = "%(asctime)s - %(levelname)s - %(message)s" + if args.log_file: + # 1. Configure the batch runner logger (write only batch progress to the main log file) + main_logger = logging.getLogger("__main__") + main_logger.propagate = ( + False # Prevent batch logs from going to the root logger + ) + + batch_handler = logging.FileHandler(args.log_file) + batch_handler.setFormatter(logging.Formatter(log_format)) + main_logger.addHandler(batch_handler) + main_logger.setLevel(logging.INFO) + + # 2. Configure the root logger (write all agent internal logs to a separate file) + base, ext = os.path.splitext(args.log_file) + agent_log_file = f"{base}_agent{ext}" + logging.basicConfig( + filename=agent_log_file, + level=logging.INFO, + format=log_format, + force=True, + ) + else: + # Route all logs to the console/terminal + logging.basicConfig( + level=logging.INFO, + format=log_format, + force=True, + ) + + # Start the asyncio event loop + asyncio.run(main_async(args)) if __name__ == "__main__": From 22af1aa87ee16dfc73aa02a050c62139613171a3 Mon Sep 17 00:00:00 2001 From: shangkunwang Date: Thu, 2 Jul 2026 19:31:18 +0000 Subject: [PATCH 02/11] feat: add SearchGraph and Node classes for representing and persisting search tree states --- MaxKernel/auto_search/graph.py | 146 +++++++++++++++++++++++++++++++++ 1 file changed, 146 insertions(+) create mode 100644 MaxKernel/auto_search/graph.py diff --git a/MaxKernel/auto_search/graph.py b/MaxKernel/auto_search/graph.py new file mode 100644 index 0000000..8f7bfb1 --- /dev/null +++ b/MaxKernel/auto_search/graph.py @@ -0,0 +1,146 @@ +import json +import logging +import os +from typing import Any, Dict, Optional + +from pydantic import BaseModel, Field + +logger = logging.getLogger(__name__) + + +class EvaluationResult(BaseModel): + compiled: bool = False + correct: bool = False + latency_ms: Optional[float] = None + profiling_summary: Optional[str] = None + compilation_error: Optional[str] = None + test_error: Optional[str] = None + + +class Node(BaseModel): + node_id: str + parent_id: Optional[str] + session_id: str + code: str + plan: str + depth: int + strategy_applied: Optional[str] = None + execution_status: str = "FAIL" # SUCCESS, FAIL + execution_error: Optional[str] = None + evaluation: EvaluationResult = Field(default_factory=EvaluationResult) + + @property + def is_valid_candidate(self) -> bool: + return ( + self.execution_status == "SUCCESS" + and self.evaluation.compiled + and self.evaluation.correct + ) + + +class SearchGraph: + def __init__(self, problem_id: str, graph_db_path: Optional[str] = None): + self.problem_id = problem_id + self.graph_db_path = graph_db_path + self.nodes: Dict[str, Node] = {} + self.root_id: Optional[str] = None + self.best_node_id: Optional[str] = None + self.metadata: Dict[str, Any] = {} # Stores orchestrator state + logger.info(f"Initializing SearchGraph for problem: {problem_id}") + + if self.graph_db_path and os.path.exists(self.graph_db_path): + self.load() + + def add_node(self, node: Node) -> None: + self.nodes[node.node_id] = node + logger.info( + f"Added node {node.node_id} (parent: {node.parent_id}, " + f"depth: {node.depth}, status: {node.execution_status}) to graph." + ) + + if node.parent_id is None: + self.root_id = node.node_id + + self._update_best_node(node) + self.save() + + def get_node(self, node_id: str) -> Optional[Node]: + """Retrieves a node by its id. Returns None if not found.""" + if node_id not in self.nodes: + logger.warning(f"Node {node_id} not found in graph.") + return None + return self.nodes[node_id] + + def _update_best_node(self, node: Node) -> None: + if not node.is_valid_candidate: + logger.warning(f"Node {node.node_id} is not a valid candidate. Skip.") + return + + latency = node.evaluation.latency_ms + if latency is None: + logger.warning(f"Node {node.node_id} has no latency.") + return + + if self.best_node_id is None: + self.best_node_id = node.node_id + logger.info( + f"Set initial best node: {node.node_id} (latency: {latency} ms)" + ) + return + + best_node = self.nodes[self.best_node_id] + best_latency = best_node.evaluation.latency_ms + if best_latency is None or latency < best_latency: + old_best_id = self.best_node_id + old_latency = best_latency + self.best_node_id = node.node_id + logger.info( + f"New best node found: {node.node_id} ({latency} ms). " + f"Previous best: {old_best_id} ({old_latency} ms)." + ) + + def save(self) -> None: + if not self.graph_db_path: + return + logger.info(f"Saving search graph to {self.graph_db_path}...") + + dir_name = os.path.dirname(self.graph_db_path) + if dir_name: + os.makedirs(dir_name, exist_ok=True) + + data = { + "problem_id": self.problem_id, + "root_id": self.root_id, + "best_node_id": self.best_node_id, + "nodes": {nid: node.model_dump() for nid, node in self.nodes.items()}, + "metadata": self.metadata, + } + + # Atomic write for the graph DB + tmp_path = self.graph_db_path + ".tmp" + try: + with open(tmp_path, "w") as f: + json.dump(data, f, indent=2) + os.replace(tmp_path, self.graph_db_path) + except Exception as e: + logger.error(f"Failed to save search graph atomically: {e}") + if os.path.exists(tmp_path): + try: + os.remove(tmp_path) + except OSError: + pass + + def load(self) -> None: + logger.info(f"Loading search graph from {self.graph_db_path}...") + with open(self.graph_db_path, "r") as f: + data = json.load(f) + self.problem_id = data["problem_id"] + self.root_id = data["root_id"] + self.best_node_id = data["best_node_id"] + self.metadata = data.get("metadata", {}) + self.nodes = { + nid: Node.model_validate(ndat) for nid, ndat in data["nodes"].items() + } + logger.info( + f"Successfully loaded search graph with {len(self.nodes)} nodes." + ) From 6d790859f2840c0bddca0788a49b189f447bcfe4 Mon Sep 17 00:00:00 2001 From: shangkunwang Date: Thu, 2 Jul 2026 22:26:11 +0000 Subject: [PATCH 03/11] feat: implement ADKSessionWorker to execute isolated pipeline agent sessions for search tree node expansion and pre/post processing --- MaxKernel/auto_search/worker.py | 173 ++++++++++++++++++++++++++++++++ 1 file changed, 173 insertions(+) create mode 100644 MaxKernel/auto_search/worker.py diff --git a/MaxKernel/auto_search/worker.py b/MaxKernel/auto_search/worker.py new file mode 100644 index 0000000..9c6b389 --- /dev/null +++ b/MaxKernel/auto_search/worker.py @@ -0,0 +1,173 @@ +import logging +import os +from typing import Any, Dict, Optional + +from auto_agent.agent import root_agent +from auto_agent.agent_client.auto_agent_client import AutoAgentClient +from auto_agent.subagents.pipeline_agent import AutonomousPipelineAgent +from auto_search.graph import EvaluationResult, Node + +logger = logging.getLogger(__name__) + + +class ADKSessionWorker: + async def expand_node( + self, + session_id: str, + parent_node: Node, + session_dir: str, + strategy: Optional[str] = None, + agent_config: Optional[Dict[str, Any]] = None, + ) -> Node: + """Expand an ADK session to get a new optimized kernel.""" + os.makedirs(session_dir, exist_ok=True) + + # Prepare inputs files for the agent based on parent_node + self._prepare_inputs(session_dir, parent_node) + + # Run the agent and process results + try: + state = await self._run_agent( + session_id=session_id, + session_dir=session_dir, + strategy=strategy, + agent_config=agent_config, + ) + return self._process_results( + session_id=session_id, + parent_node=parent_node, + strategy=strategy, + state=state, + ) + except Exception as e: + logger.exception(f"Session {session_id} failed: {e}") + return Node( + node_id=session_id, + parent_id=parent_node.node_id, + session_id=session_id, + code="", + plan="", + depth=parent_node.depth + 1, + strategy_applied=strategy, + execution_status="FAIL", + execution_error=str(e), + evaluation=EvaluationResult(correct=False), + ) + + def _prepare_inputs(self, session_dir: str, parent_node: Node) -> None: + """Writes the parent node code and plan to the session directory.""" + base_kernel_path = os.path.join(session_dir, "base_kernel.py") + with open(base_kernel_path, "w") as f: + f.write(parent_node.code) + + optimized_kernel_path = os.path.join(session_dir, "optimized_kernel.py") + with open(optimized_kernel_path, "w") as f: + f.write(parent_node.code) + + kernel_plan_path = None + if parent_node.plan: + kernel_plan_path = os.path.join(session_dir, "base_kernel_plan.md") + with open(kernel_plan_path, "w") as f: + f.write(parent_node.plan) + + async def _run_agent( + self, + session_id: str, + session_dir: str, + strategy: Optional[str], + agent_config: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Sets up a custom AutonomousPipelineAgent and runs the client.""" + agent_config = agent_config or {} + + custom_agent = AutonomousPipelineAgent( + name="AutonomousPipelineAgent", + plan_agent=root_agent.plan_agent, + implement_agent=root_agent.implement_agent, + validate_agent=root_agent.validate_agent, + test_gen_agent=root_agent.test_gen_agent, + test_run_agent=root_agent.test_run_agent, + autotune_agent=root_agent.autotune_agent, + profile_agent=root_agent.profile_agent, + session_dir=session_dir, + **agent_config, + ) + + strategy_query = f"Focus on: {strategy}. " if strategy else "" + query = ( + "Optimize the code for peak performance with pallas kernel. " + f"{strategy_query}" + "Base code is at base_kernel.py." + ) + client = AutoAgentClient( + user_id="orchestrator", + session_id=session_id, + query=query, + agent=custom_agent, + ) + await client.create_session() + await client.run_async() + return client.get_state() + + def _process_results( + self, + session_id: str, + parent_node: Node, + strategy: Optional[str], + state: Dict[str, Any], + ) -> Node: + """Reads optimized code/plan and extracts metrics to construct a new Node.""" + history = state.get("history", []) + best_iter = state.get("best_iteration", -1) + + best_run = {} + if best_iter != -1: + best_run = next( + (run for run in history if run.get("iteration") == best_iter), {} + ) + elif history: + best_run = history[-1] + + comp_status = best_run.get("compilation_status", {}) + compiled = comp_status.get("success", False) + compilation_error = comp_status.get("message") if not compiled else None + + test_status = best_run.get("test_status", {}) + correct = test_status.get("success", False) + test_error = test_status.get("output") if not correct else None + + latency_ms = best_run.get("latency_ms") + profiling_summary = best_run.get("profiling_summary") + + # Read optimized code if it exists + opt_path = state.get("optimized_kernel_path") + optimized_code = "" + if opt_path and os.path.exists(opt_path): + with open(opt_path, "r") as f: + optimized_code = f.read() + + # Read optimized plan if it exists + plan_path = state.get("kernel_plan_path") + optimized_plan = "" + if plan_path and os.path.exists(plan_path): + with open(plan_path, "r") as f: + optimized_plan = f.read() + + return Node( + node_id=session_id, + parent_id=parent_node.node_id, + session_id=session_id, + code=optimized_code, + plan=optimized_plan, + depth=parent_node.depth + 1, + strategy_applied=strategy, + execution_status="SUCCESS", + evaluation=EvaluationResult( + compiled=compiled, + correct=correct, + latency_ms=latency_ms, + profiling_summary=profiling_summary, + compilation_error=compilation_error, + test_error=test_error, + ), + ) From 0a70cd20f6c7e422602b719ff9c1b6f3dd273045 Mon Sep 17 00:00:00 2001 From: shangkunwang Date: Mon, 6 Jul 2026 04:22:17 +0000 Subject: [PATCH 04/11] fix: make function get_session_data from auto agent client public --- MaxKernel/auto_agent/agent_client/auto_agent_client.py | 2 +- MaxKernel/auto_agent/agent_client/run_batch_agent_call.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/MaxKernel/auto_agent/agent_client/auto_agent_client.py b/MaxKernel/auto_agent/agent_client/auto_agent_client.py index c80cfef..dff0bd3 100644 --- a/MaxKernel/auto_agent/agent_client/auto_agent_client.py +++ b/MaxKernel/auto_agent/agent_client/auto_agent_client.py @@ -50,7 +50,7 @@ async def create_session(self) -> None: session_id=self.session_id, ) - def _get_session_data(self) -> dict: + def get_session_data(self) -> dict: if not self.session: raise ValueError("Session has not been created yet.") 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 bec07bb..17958f1 100644 --- a/MaxKernel/auto_agent/agent_client/run_batch_agent_call.py +++ b/MaxKernel/auto_agent/agent_client/run_batch_agent_call.py @@ -50,7 +50,7 @@ def get_optimized_kernel_path(session_file_path: str): def save_session_to_file(client: Any, file_path: str): """Fetches full session data and saves it to a local JSON file.""" - session_data = client._get_session_data() + session_data = client.get_session_data() with open(file_path, "w") as f: json.dump(session_data, f, indent=2) logger.info(f"Saved session data to {file_path}") From 6e656bca78fd91d4b86c69a8c7f38d2617fb324d Mon Sep 17 00:00:00 2001 From: shangkunwang Date: Mon, 6 Jul 2026 04:29:03 +0000 Subject: [PATCH 05/11] refactor: replace session_id with session_dir in Node and save session data in worker --- MaxKernel/auto_search/graph.py | 2 +- MaxKernel/auto_search/worker.py | 32 +++++++++++++++++++++----------- 2 files changed, 22 insertions(+), 12 deletions(-) diff --git a/MaxKernel/auto_search/graph.py b/MaxKernel/auto_search/graph.py index 8f7bfb1..fe1c39b 100644 --- a/MaxKernel/auto_search/graph.py +++ b/MaxKernel/auto_search/graph.py @@ -20,7 +20,7 @@ class EvaluationResult(BaseModel): class Node(BaseModel): node_id: str parent_id: Optional[str] - session_id: str + session_dir: Optional[str] = None code: str plan: str depth: int diff --git a/MaxKernel/auto_search/worker.py b/MaxKernel/auto_search/worker.py index 9c6b389..2ccf7bd 100644 --- a/MaxKernel/auto_search/worker.py +++ b/MaxKernel/auto_search/worker.py @@ -1,3 +1,4 @@ +import json import logging import os from typing import Any, Dict, Optional @@ -13,7 +14,7 @@ class ADKSessionWorker: async def expand_node( self, - session_id: str, + node_id: str, parent_node: Node, session_dir: str, strategy: Optional[str] = None, @@ -28,23 +29,24 @@ async def expand_node( # Run the agent and process results try: state = await self._run_agent( - session_id=session_id, + node_id=node_id, session_dir=session_dir, strategy=strategy, agent_config=agent_config, ) return self._process_results( - session_id=session_id, + node_id=node_id, parent_node=parent_node, strategy=strategy, state=state, + session_dir=session_dir, ) except Exception as e: - logger.exception(f"Session {session_id} failed: {e}") + logger.exception(f"Node {node_id} expansion failed: {e}") return Node( - node_id=session_id, + node_id=node_id, parent_id=parent_node.node_id, - session_id=session_id, + session_dir=session_dir, code="", plan="", depth=parent_node.depth + 1, @@ -72,7 +74,7 @@ def _prepare_inputs(self, session_dir: str, parent_node: Node) -> None: async def _run_agent( self, - session_id: str, + node_id: str, session_dir: str, strategy: Optional[str], agent_config: Optional[Dict[str, Any]] = None, @@ -101,20 +103,28 @@ async def _run_agent( ) client = AutoAgentClient( user_id="orchestrator", - session_id=session_id, + session_id=node_id, query=query, agent=custom_agent, ) await client.create_session() await client.run_async() + try: + session_json_path = os.path.join(session_dir, "session.json") + with open(session_json_path, "w") as f: + json.dump(client.get_session_data(), f, indent=2) + logger.info(f"Saved session data to {session_json_path}") + except Exception as se: + logger.warning(f"Failed to save session data for {node_id}: {se}") return client.get_state() def _process_results( self, - session_id: str, + node_id: str, parent_node: Node, strategy: Optional[str], state: Dict[str, Any], + session_dir: str, ) -> Node: """Reads optimized code/plan and extracts metrics to construct a new Node.""" history = state.get("history", []) @@ -154,9 +164,9 @@ def _process_results( optimized_plan = f.read() return Node( - node_id=session_id, + node_id=node_id, parent_id=parent_node.node_id, - session_id=session_id, + session_dir=session_dir, code=optimized_code, plan=optimized_plan, depth=parent_node.depth + 1, From 710a5032009591412b7c9ddd3797eb76f750d228 Mon Sep 17 00:00:00 2001 From: shangkunwang Date: Mon, 6 Jul 2026 17:39:10 +0000 Subject: [PATCH 06/11] refactor: pass reference_code to ADKSessionWorker for base kernel input --- MaxKernel/auto_search/worker.py | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/MaxKernel/auto_search/worker.py b/MaxKernel/auto_search/worker.py index 2ccf7bd..6ac8e14 100644 --- a/MaxKernel/auto_search/worker.py +++ b/MaxKernel/auto_search/worker.py @@ -17,14 +17,15 @@ async def expand_node( node_id: str, parent_node: Node, session_dir: str, + reference_code: str, strategy: Optional[str] = None, agent_config: Optional[Dict[str, Any]] = None, ) -> Node: """Expand an ADK session to get a new optimized kernel.""" os.makedirs(session_dir, exist_ok=True) - # Prepare inputs files for the agent based on parent_node - self._prepare_inputs(session_dir, parent_node) + # Prepare inputs files for the agent based on reference_code and parent_node + self._prepare_inputs(session_dir, parent_node, reference_code) # Run the agent and process results try: @@ -56,15 +57,18 @@ async def expand_node( evaluation=EvaluationResult(correct=False), ) - def _prepare_inputs(self, session_dir: str, parent_node: Node) -> None: - """Writes the parent node code and plan to the session directory.""" + def _prepare_inputs( + self, session_dir: str, parent_node: Node, reference_code: str + ) -> None: + """Writes the reference code and parent's optimized code (if any) to session directory.""" base_kernel_path = os.path.join(session_dir, "base_kernel.py") with open(base_kernel_path, "w") as f: - f.write(parent_node.code) + f.write(reference_code) - optimized_kernel_path = os.path.join(session_dir, "optimized_kernel.py") - with open(optimized_kernel_path, "w") as f: - f.write(parent_node.code) + if parent_node.depth > 0 and parent_node.code: + optimized_kernel_path = os.path.join(session_dir, "optimized_kernel.py") + with open(optimized_kernel_path, "w") as f: + f.write(parent_node.code) kernel_plan_path = None if parent_node.plan: From 102595c065650593b21404109cbb3b5bc353bbb4 Mon Sep 17 00:00:00 2001 From: shangkunwang Date: Tue, 7 Jul 2026 00:17:19 +0000 Subject: [PATCH 07/11] feat: implement the base orchestrator abstract class --- MaxKernel/auto_search/orchestrator.py | 213 ++++++++++++++++++++++++++ 1 file changed, 213 insertions(+) create mode 100644 MaxKernel/auto_search/orchestrator.py diff --git a/MaxKernel/auto_search/orchestrator.py b/MaxKernel/auto_search/orchestrator.py new file mode 100644 index 0000000..8aed02d --- /dev/null +++ b/MaxKernel/auto_search/orchestrator.py @@ -0,0 +1,213 @@ +import abc +import asyncio +import logging +import os +import time +from typing import Any, List, Optional, Tuple + +from auto_search.graph import EvaluationResult, Node, SearchGraph + +logger = logging.getLogger(__name__) + + +class SearchOrchestrator(abc.ABC): + def __init__( + self, + problem_id: str, + reference_code: str, + graph_db_path: Optional[str] = None, + max_concurrency: int = 2, + ): + # Resolve graph db path and run directory + if not graph_db_path: + workdir = os.environ.get("WORKDIR", os.getcwd()) + timestamp = time.strftime("%Y%m%d_%H%M%S") + graph_db_path = os.path.join( + workdir, f"graph_{problem_id}_{timestamp}.json" + ) + logger.info(f"No graph_db_path provided. Generated: {graph_db_path}") + + # The run directory is the root directory for this search run. + self.run_dir = os.path.dirname(os.path.abspath(graph_db_path)) + logger.info(f"Using run directory: {self.run_dir}") + os.makedirs(self.run_dir, exist_ok=True) + + # Initialization of internal states + self._semaphore = asyncio.Semaphore(max_concurrency) + self._node_counter = 0 + + self.reference_code = reference_code + + # Initialization of search graph + self.graph = SearchGraph(problem_id, graph_db_path) + if self.graph.root_id: + logger.info( + "Found existing graph. Call `.resume()` explicitly to restore" + " search state." + ) + else: + # Start from scratch, create root node + node_id, session_dir = self.get_next_session_node() + os.makedirs(session_dir, exist_ok=True) + # Write reference code to the root node directory + with open(os.path.join(session_dir, "base_kernel.py"), "w") as f: + f.write(reference_code) + + root_node = Node( + node_id=node_id, + parent_id=None, + session_dir=session_dir, + code=reference_code, + plan="", + depth=0, + strategy_applied="baseline", + execution_status="SUCCESS", + evaluation=EvaluationResult( + compiled=True, + correct=True, + latency_ms=float("inf"), + profiling_summary="", + ), + ) + + self.graph.add_node(root_node) + + def get_next_session_node( + self, suffix: Optional[str] = None + ) -> Tuple[str, str]: + """Returns (node_id, session_dir) for the next node in the graph.""" + node_id = f"node_{self._node_counter:03d}" + self._node_counter += 1 + dir_name = f"{node_id}_{suffix}" if suffix else node_id + session_dir = os.path.join(self.run_dir, "nodes", dir_name) + return node_id, session_dir + + def resume(self) -> None: + """Restores the orchestrator state from the loaded graph's metadata.""" + logger.info("Resuming search orchestration from persisted graph state...") + + # Calculate the node counter based on existing nodes in the graph + existing_indices = [ + int(node_id.split("_")[1]) + for node_id in self.graph.nodes.keys() + if node_id.startswith("node_") and node_id[5:].isdigit() + ] + self._node_counter = max(existing_indices) + 1 if existing_indices else 0 + + self._resume() + + def update_metadata(self, key: str, value: Any) -> None: + """Writes orchestrator metadata to the graph and triggers a save.""" + self.graph.metadata[key] = value + self.graph.save() + + # --- Template Methods --- + + @abc.abstractmethod + def _resume(self) -> None: + """Restores subclass-specific instance variables from metadata. + + Reads from self.graph.metadata to restore internal algorithm state. + """ + pass + + @abc.abstractmethod + def _select_nodes_to_expand(self) -> List[Node]: + """Selects candidate nodes from the search graph for expansion. + + Returns a list of nodes to be expanded in the current step. + """ + pass + + @abc.abstractmethod + def _generate_expansion_tasks(self, nodes: List[Node]) -> Any: + """Generates expansion tasks for the selected nodes. + + Creates tasks to be passed directly into _execute_expansions. + """ + pass + + @abc.abstractmethod + def _update_search_state(self, new_nodes: List[Node]) -> None: + """Updates internal algorithm state using evaluated nodes. + + Updates state such as beam frontier or learnings based on newly evaluated + nodes from worker expansions. + """ + pass + + @abc.abstractmethod + def _should_terminate(self) -> bool: + """Determines whether the search loop should terminate. + + Returns True if termination criteria (e.g., max depth reached or no + candidates) are met. + """ + pass + + def _post_step_hook(self) -> None: + """Optional hook executed at the end of each search step. + + Can be overridden for cleanup, state updates, or logging. + """ + pass + + # --- Concrete Flow --- + + @abc.abstractmethod + async def _execute_expansions(self, tasks: Any) -> List[Node]: + """Executes expansion tasks to evaluate new kernels. + + Runs tasks (typically in parallel via workers) to compile, test, and profile + candidate kernels, returning new Nodes. + """ + pass + + async def run(self) -> Node: + logger.info( + f"Starting search orchestration for problem: {self.graph.problem_id}" + ) + + while not self._should_terminate(): + # 1. Select nodes to expand + nodes_to_expand = self._select_nodes_to_expand() + if not nodes_to_expand: + break + + # 2. Generate strategies + tasks = self._generate_expansion_tasks(nodes_to_expand) + if not tasks: + break + + # 3. Parallel Execution (Workers compile, test, and profile) + new_nodes = await self._execute_expansions(tasks) + + # 4. Add any nodes not already saved during execution + for node in new_nodes: + if node.node_id not in self.graph.nodes: + self.graph.add_node(node) + + # 5. Update Search Algorithm State + self._update_search_state(new_nodes) + + # Log progress + best_id = self.graph.best_node_id + best_latency = "N/A" + if best_id: + best_node = self.graph.get_node(best_id) + if best_node and best_node.evaluation.latency_ms is not None: + best_latency = f"{best_node.evaluation.latency_ms:.3f} ms" + logger.info( + f"Step completed. Total nodes in graph: {len(self.graph.nodes)}. " + f"Current best latency: {best_latency}" + ) + + # 6. Step Cleanup + self._post_step_hook() + + best_id = self.graph.best_node_id + return ( + self.graph.get_node(best_id) + if best_id + else self.graph.get_node(self.graph.root_id) + ) From d2caea27b4ceb40b1295b72401a336a9a27bab83 Mon Sep 17 00:00:00 2001 From: shangkunwang Date: Tue, 7 Jul 2026 00:17:54 +0000 Subject: [PATCH 08/11] feat: implement the parallel search algorithm --- .../auto_search/algorithms/parallel_search.py | 145 ++++++++++++++++++ 1 file changed, 145 insertions(+) create mode 100644 MaxKernel/auto_search/algorithms/parallel_search.py diff --git a/MaxKernel/auto_search/algorithms/parallel_search.py b/MaxKernel/auto_search/algorithms/parallel_search.py new file mode 100644 index 0000000..37502a8 --- /dev/null +++ b/MaxKernel/auto_search/algorithms/parallel_search.py @@ -0,0 +1,145 @@ +import asyncio +import logging +from typing import List, Optional, Tuple + +from auto_search.graph import Node +from auto_search.orchestrator import SearchOrchestrator +from auto_search.worker import ADKSessionWorker + +logger = logging.getLogger(__name__) + + +class SimpleParallelSearchOrchestrator(SearchOrchestrator): + """Orchestrates single-iteration parallel kernel explorations from root.""" + + def __init__( + self, + num_parallel_runs: int = 2, + strategies: Optional[List[str]] = None, + max_worker_retries: int = 1, + agent_config: Optional[dict] = None, + **kwargs, + ): + if num_parallel_runs <= 0: + raise ValueError( + f"num_parallel_runs must be a positive integer, got {num_parallel_runs}." + ) + if max_worker_retries < 1: + raise ValueError( + f"max_worker_retries must be at least 1, got {max_worker_retries}." + ) + + self.num_parallel_runs = num_parallel_runs + if strategies is not None: + if len(strategies) > num_parallel_runs: + raise ValueError( + f"Number of user specified strategies ({len(strategies)}) cannot" + f" exceed num_parallel_runs ({num_parallel_runs})." + ) + self.strategies = [s or "" for s in strategies] + [""] * ( + num_parallel_runs - len(strategies) + ) + else: + self.strategies = [""] * num_parallel_runs + + self.remaining_strategies = list(self.strategies) + self.max_worker_retries = max_worker_retries + self.agent_config = agent_config + self.worker = ADKSessionWorker() + + super().__init__(**kwargs) + + def _resume(self) -> None: + if not self.graph.root_id: + logger.warning("Cannot resume: root_id is not set in graph.") + return + + existing_children = [ + node + for node in self.graph.nodes.values() + if node.parent_id == self.graph.root_id + ] + + for node in existing_children: + strategy = node.strategy_applied or "" + if strategy in self.remaining_strategies: + self.remaining_strategies.remove(strategy) + logger.info( + f"Resumed Simple Parallel Search. Found {len(existing_children)} existing" + f" evaluations. {len(self.remaining_strategies)} runs remaining." + ) + + def _select_nodes_to_expand(self) -> List[Node]: + if self._should_terminate(): + return [] + if not self.graph.root_id: + logger.error( + "Cannot select nodes to expand: root_id is not set in graph." + ) + return [] + root_node = self.graph.get_node(self.graph.root_id) + return [root_node] if root_node else [] + + def _generate_expansion_tasks( + self, nodes: List[Node] + ) -> List[Tuple[Node, str]]: + if not self.remaining_strategies: + logger.info( + f"All {self.num_parallel_runs} parallel runs have already completed." + ) + return [] + + root_node = self.graph.get_node(self.graph.root_id) + if not root_node: + return [] + + logger.info( + f"Launching {len(self.remaining_strategies)} parallel runs for root node." + ) + return [(root_node, strategy) for strategy in self.remaining_strategies] + + def _update_search_state(self, new_nodes: List[Node]) -> None: + for node in new_nodes: + strategy = node.strategy_applied or "" + if strategy in self.remaining_strategies: + self.remaining_strategies.remove(strategy) + + def _should_terminate(self) -> bool: + return len(self.remaining_strategies) == 0 + + async def _execute_expansions( + self, tasks: List[Tuple[Node, str]] + ) -> List[Node]: + async def run_task(task_idx: int, parent_node: Node, strategy: str) -> Node: + node_id, base_dir = self.get_next_session_node() + for attempt in range(1, self.max_worker_retries + 1): + async with self._semaphore: + session_dir = f"{base_dir}_attempt_{attempt}" + node = await self.worker.expand_node( + node_id, + parent_node, + strategy=strategy, + session_dir=session_dir, + reference_code=self.reference_code, + agent_config=self.agent_config, + ) + if ( + node.execution_status == "SUCCESS" + or attempt == self.max_worker_retries + ): + self.graph.add_node(node) + return node + + strat_info = f" (strategy: {strategy})" if strategy else "" + logger.warning( + f"Task {task_idx}{strat_info} failed attempt" + f" {attempt}/{self.max_worker_retries}: {node.execution_error}." + " Retrying..." + ) + await asyncio.sleep(2**attempt) + + futures = [ + run_task(i, parent, strategy) + for i, (parent, strategy) in enumerate(tasks) + ] + return await asyncio.gather(*futures) From 6cbad68a225bf790505865837a2a266a99287112 Mon Sep 17 00:00:00 2001 From: shangkunwang Date: Tue, 7 Jul 2026 00:19:34 +0000 Subject: [PATCH 09/11] feat: prepare the scripts for single and batch run of the auto search --- MaxKernel/auto_search/run_batch_search.py | 202 ++++++++++++++++ MaxKernel/auto_search/run_search.py | 273 ++++++++++++++++++++++ 2 files changed, 475 insertions(+) create mode 100644 MaxKernel/auto_search/run_batch_search.py create mode 100644 MaxKernel/auto_search/run_search.py diff --git a/MaxKernel/auto_search/run_batch_search.py b/MaxKernel/auto_search/run_batch_search.py new file mode 100644 index 0000000..abeb724 --- /dev/null +++ b/MaxKernel/auto_search/run_batch_search.py @@ -0,0 +1,202 @@ +import argparse +import asyncio +import json +import logging +import os +from typing import Any, Tuple + +from auto_search.run_search import run_search, setup_logging + +logger = logging.getLogger(__name__) + + +async def process_problem( + problem_dir: str, + algorithm: str, + sem: asyncio.Semaphore, + **kwargs: Any, +) -> Tuple[str, str]: + """Executes the search algorithm for a single benchmark problem in the batch.""" + async with sem: + try: + return await run_search( + problem_dir=problem_dir, + algorithm=algorithm, + **kwargs, + ) + except Exception as e: + problem_id = os.path.basename(os.path.normpath(problem_dir)) + logger.error( + f"Error executing search on {problem_id}: {e}", exc_info=True + ) + return problem_id, f"Failed with exception: {e}" + + +async def run_batch_search( + data_dir: str, + algorithm: str = "parallel", + num_problem_concurrency: int = 1, + **kwargs: Any, +): + """Coordinates concurrent problem execution across the dataset.""" + if not os.path.isdir(data_dir): + logger.error( + f"Dataset directory not found or not a directory: {data_dir}" + ) + return + + data_dir_valid = [ + os.path.join(data_dir, d) + for d in os.listdir(data_dir) + if os.path.isfile(os.path.join(data_dir, d, "reference.py")) + ] + data_dir_valid.sort() + + if not data_dir_valid: + logger.warning( + f"No valid benchmark problems (directories containing reference.py) found in {data_dir}" + ) + return + + logger.info(f"Found {len(data_dir_valid)} problems to process.") + max_concurrency = kwargs.get("max_concurrency", 2) + logger.info( + f"Algorithm: {algorithm}, Problem Concurrency:" + f" {num_problem_concurrency}, Worker Concurrency: {max_concurrency}" + ) + + sem = asyncio.Semaphore(num_problem_concurrency) + tasks = [ + process_problem( + problem_dir=problem_dir, + algorithm=algorithm, + sem=sem, + **kwargs, + ) + for problem_dir in data_dir_valid + ] + + completed = 0 + results = [] + for future in asyncio.as_completed(tasks): + try: + prob_id, status = await future + completed += 1 + results.append((prob_id, status)) + logger.info( + f"[{completed}/{len(data_dir_valid)}] Problem {prob_id}: {status}" + ) + except Exception as e: + logger.error(f"A task raised an exception: {e}") + + logger.info("\n--- Search Execution Summary ---") + for prob_id, status in sorted(results): + logger.info(f"{prob_id}: {status}") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run Auto-Search algorithms against batch dataset." + ) + + # General & Orchestration Arguments + orch_group = parser.add_argument_group( + "General & Orchestration Arguments", + "Arguments shared across all algorithms and orchestrators.", + ) + orch_group.add_argument( + "--data_dir", + type=str, + required=True, + help="Path to benchmark dataset directory", + ) + orch_group.add_argument( + "--algorithm", + type=str, + choices=["parallel", "beam", "agentic"], + default="parallel", + help="Search algorithm to execute", + ) + orch_group.add_argument( + "--max_concurrency", + type=int, + default=2, + help="Max concurrent worker expansions", + ) + orch_group.add_argument( + "--num_problem_concurrency", + type=int, + default=1, + help="Number of dataset problems to run concurrently", + ) + orch_group.add_argument( + "--log_file", + type=str, + default=None, + help="File to save logs to", + ) + # Parallel Search Arguments + parallel_group = parser.add_argument_group( + "Parallel Search Arguments", + "Parameters specific to the 'parallel' search algorithm.", + ) + parallel_group.add_argument( + "--num_parallel_runs", + type=int, + default=2, + help="Number of parallel runs", + ) + parallel_group.add_argument( + "--max_retries", + type=int, + default=1, + help="Max worker retries per expansion task", + ) + parallel_group.add_argument( + "--strategies", + nargs="+", + type=str, + default=None, + help="List of strategy strings to explore", + ) + parallel_group.add_argument( + "--agent_config", + type=str, + default=None, + help="JSON string of agent config parameters (e.g. '{\"max_iterations\": 5}')", + ) + return parser.parse_args() + + +def main(): + args = parse_args() + setup_logging(args.log_file) + + parsed_agent_config = None + if args.agent_config: + try: + parsed_agent_config = json.loads(args.agent_config) + except json.JSONDecodeError as e: + logger.error(f"Invalid JSON string for --agent_config: {e}") + return + + kwargs = { + "max_concurrency": args.max_concurrency, + "num_parallel_runs": args.num_parallel_runs, + "max_worker_retries": args.max_retries, + "strategies": args.strategies, + "agent_config": parsed_agent_config, + } + + asyncio.run( + run_batch_search( + data_dir=args.data_dir, + algorithm=args.algorithm, + num_problem_concurrency=args.num_problem_concurrency, + **kwargs, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/MaxKernel/auto_search/run_search.py b/MaxKernel/auto_search/run_search.py new file mode 100644 index 0000000..fd717e4 --- /dev/null +++ b/MaxKernel/auto_search/run_search.py @@ -0,0 +1,273 @@ +import argparse +import asyncio +import json +import logging +import os +import time +from typing import Any, Optional, Tuple + +try: + from dotenv import load_dotenv + + load_dotenv() +except ImportError: + logging.warning( + "dotenv not installed, skipping loading environment variables" + ) + +from auto_search.algorithms.parallel_search import ( + SimpleParallelSearchOrchestrator, +) +from auto_search.orchestrator import SearchOrchestrator + +logger = logging.getLogger(__name__) + + +def get_orchestrator( + algorithm: str, + problem_id: str, + reference_code: str, + graph_db_path: str, + **kwargs: Any, +) -> SearchOrchestrator: + """Factory function to instantiate the requested search orchestrator.""" + if algorithm == "parallel": + return SimpleParallelSearchOrchestrator( + problem_id=problem_id, + reference_code=reference_code, + graph_db_path=graph_db_path, + max_concurrency=kwargs.get("max_concurrency", 2), + num_parallel_runs=kwargs.get("num_parallel_runs", 2), + strategies=kwargs.get("strategies"), + max_worker_retries=kwargs.get("max_worker_retries", 1), + agent_config=kwargs.get("agent_config"), + ) + elif algorithm in ("beam", "agentic"): + raise NotImplementedError( + f"Algorithm '{algorithm}' is currently a placeholder." + ) + else: + raise ValueError(f"Unknown algorithm: {algorithm}") + + +def save_optimized_kernel( + orchestrator: SearchOrchestrator, + problem_dir: str, + algorithm: str, +) -> Optional[Tuple[str, float]]: + """Saves the best kernel code found during search to the problem directory.""" + best_id = orchestrator.graph.best_node_id + if not best_id or best_id not in orchestrator.graph.nodes: + return None + + best_node = orchestrator.graph.nodes[best_id] + if not best_node.code: + return None + + output_file = os.path.join(problem_dir, f"optimized_{algorithm}.py") + with open(output_file, "w") as f: + f.write(best_node.code) + + logger.info(f"Saved best kernel ({best_id}) to {output_file}") + latency = best_node.evaluation.latency_ms or -1.0 + return best_id, latency + + +async def run_search( + problem_dir: str, + algorithm: str = "parallel", + graph_db_path: Optional[str] = None, + **kwargs: Any, +) -> Tuple[str, str]: + """Executes the search algorithm asynchronously for a single problem directory.""" + problem_id = os.path.basename(os.path.normpath(problem_dir)) + reference_file = os.path.join(problem_dir, "reference.py") + + if not os.path.exists(reference_file): + logger.error(f"reference.py not found in {problem_dir}") + return problem_id, "Failed: reference.py missing" + + with open(reference_file, "r") as f: + reference_code = f.read() + + if not graph_db_path: + timestamp = time.strftime("%Y%m%d_%H%M%S") + run_subdir = os.path.join( + os.environ.get("WORKDIR"), + "search_runs", + f"{problem_id}_run_{algorithm}_{timestamp}", + ) + os.makedirs(run_subdir, exist_ok=True) + graph_db_path = os.path.join(run_subdir, "search_graph.json") + else: + os.makedirs(os.path.dirname(os.path.abspath(graph_db_path)), exist_ok=True) + + logger.info( + f"Starting '{algorithm}' search on {problem_id} (graph: {graph_db_path})" + ) + + try: + orchestrator = get_orchestrator( + algorithm=algorithm, + problem_id=problem_id, + reference_code=reference_code, + graph_db_path=graph_db_path, + **kwargs, + ) + + if os.path.exists(graph_db_path): + logger.info(f"Existing graph found for {problem_id}. Resuming...") + orchestrator.resume() + + await orchestrator.run() + + best_result = save_optimized_kernel(orchestrator, problem_dir, algorithm) + if best_result: + best_id, latency = best_result + return problem_id, f"Success (Best: {best_id}, Latency: {latency} ms)" + + return problem_id, "Completed (No valid candidate found)" + + except Exception as e: + logger.error(f"Error executing search on {problem_id}: {e}", exc_info=True) + return problem_id, f"Failed: {e}" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run Auto-Search algorithm on a single problem directory." + ) + # General & Orchestration Arguments + orch_group = parser.add_argument_group( + "General & Orchestration Arguments", + "Arguments shared across all algorithms and orchestrators.", + ) + orch_group.add_argument( + "--problem_dir", + type=str, + required=True, + help="Path to problem directory containing reference.py", + ) + orch_group.add_argument( + "--algorithm", + type=str, + choices=["parallel", "beam", "agentic"], + default="parallel", + help="Search algorithm to execute", + ) + orch_group.add_argument( + "--graph_db_path", + type=str, + default=None, + help="Explicit path to existing graph JSON to resume from (or write to)", + ) + orch_group.add_argument( + "--max_concurrency", + type=int, + default=2, + help="Max concurrent worker expansions", + ) + orch_group.add_argument( + "--log_file", + type=str, + default=None, + help="File to save logs to", + ) + # Parallel Search Arguments + parallel_group = parser.add_argument_group( + "Parallel Search Arguments", + "Parameters specific to the 'parallel' search algorithm.", + ) + parallel_group.add_argument( + "--num_parallel_runs", + type=int, + default=2, + help="Number of parallel runs", + ) + parallel_group.add_argument( + "--max_retries", + type=int, + default=1, + help="Max worker retries per expansion task", + ) + parallel_group.add_argument( + "--strategies", + nargs="+", + type=str, + default=None, + help="List of strategy strings to explore", + ) + parallel_group.add_argument( + "--agent_config", + type=str, + default=None, + help="JSON string of agent config parameters (e.g. '{\"max_iterations\": 5}')", + ) + return parser.parse_args() + + +def setup_logging(log_file: Optional[str]): + """Configures console and file logging.""" + log_format = "%(asctime)s - %(levelname)s - %(message)s" + if log_file: + # 1. Configure the main runner logger (write only main script progress to the main log file) + main_logger = logging.getLogger("__main__") + main_logger.propagate = False + + batch_handler = logging.FileHandler(log_file) + batch_handler.setFormatter(logging.Formatter(log_format)) + main_logger.addHandler(batch_handler) + main_logger.setLevel(logging.INFO) + + # 2. Configure the root logger (write all internal and dependency logs to a separate file) + base, ext = os.path.splitext(log_file) + agent_log_file = f"{base}_agent{ext}" + logging.basicConfig( + filename=agent_log_file, + level=logging.INFO, + format=log_format, + force=True, + ) + else: + # Route all logs to the console/terminal + logging.basicConfig( + level=logging.INFO, + format=log_format, + force=True, + ) + + +def main(): + args = parse_args() + setup_logging(args.log_file) + + agent_config = None + if args.agent_config: + try: + agent_config = json.loads(args.agent_config) + except json.JSONDecodeError as e: + logger.error(f"Invalid JSON string for --agent_config: {e}") + return + + kwargs = { + "max_concurrency": args.max_concurrency, + # Parallel Search Arguments + "num_parallel_runs": args.num_parallel_runs, + "max_worker_retries": args.max_retries, + "strategies": args.strategies, + "agent_config": agent_config, + } + + prob_id, status = asyncio.run( + run_search( + problem_dir=args.problem_dir, + algorithm=args.algorithm, + graph_db_path=args.graph_db_path, + **kwargs, + ) + ) + logger.info(f"Result for {prob_id}: {status}") + + +if __name__ == "__main__": + main() From 08b56fffc836167c401007fe81e48f346cde7685 Mon Sep 17 00:00:00 2001 From: shangkunwang Date: Fri, 10 Jul 2026 21:34:51 +0000 Subject: [PATCH 10/11] feat: add graph visualization --- .../auto_search/utils/visualize_graph.py | 72 +++++++++++++++++++ 1 file changed, 72 insertions(+) create mode 100644 MaxKernel/auto_search/utils/visualize_graph.py diff --git a/MaxKernel/auto_search/utils/visualize_graph.py b/MaxKernel/auto_search/utils/visualize_graph.py new file mode 100644 index 0000000..16e2f87 --- /dev/null +++ b/MaxKernel/auto_search/utils/visualize_graph.py @@ -0,0 +1,72 @@ +import json +import logging +import os +import sys +from typing import Optional + +import graphviz + +logger = logging.getLogger(__name__) + + +def visualize_graph( + graph_json_path: str, output_path: Optional[str] = None +) -> None: + """Reads a search graph JSON file and exports it to a PNG.""" + if not os.path.exists(graph_json_path): + logger.error(f"Graph JSON file not found: {graph_json_path}") + return + + if output_path is None: + output_base = os.path.splitext(graph_json_path)[0] + else: + output_base = os.path.splitext(output_path)[0] + + with open(graph_json_path, "r") as f: + data = json.load(f) + + nodes_data = data.get("nodes", {}) + best_node_id = data.get("best_node_id") + + dot = graphviz.Digraph(name="SearchGraph") + dot.attr("node", shape="box", style="filled", fontname="Helvetica") + + for node_id, node_data in nodes_data.items(): + eval_data = node_data.get("evaluation") or {} + latency_ms = eval_data.get("latency_ms") + latency = f"{latency_ms:.4f} ms" if latency_ms is not None else "N/A" + + execution_status = node_data.get("execution_status", "UNKNOWN") + depth = node_data.get("depth", 0) + + color = "lightblue" + if node_id == best_node_id: + color = "lightgreen" + elif execution_status != "SUCCESS": + color = "lightpink" + + label = f"ID: {node_id}\nDepth: {depth}\nLatency: {latency}" + dot.node(node_id, label, fillcolor=color) + + parent_id = node_data.get("parent_id") + if parent_id: + dot.edge(parent_id, node_id) + + try: + dot.render(filename=output_base, format="png", cleanup=True) + logger.info(f"Successfully rendered graph to PNG at {output_base}.png") + except Exception as e: + logger.error(f"Failed to generate PNG: {e}") + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + + if len(sys.argv) < 2: + print( + "Usage: python -m auto_search.utils.visualize_graph [output_graph_path]" + ) + sys.exit(1) + + output = sys.argv[2] if len(sys.argv) > 2 else None + visualize_graph(sys.argv[1], output) From 06af4f0583827ec506c5a4f274c978b70dc6a760 Mon Sep 17 00:00:00 2001 From: shangkunwang Date: Fri, 10 Jul 2026 21:47:23 +0000 Subject: [PATCH 11/11] refactor: move max_worker_retries as a common arg in the base orchestrator --- MaxKernel/auto_search/algorithms/parallel_search.py | 6 ------ MaxKernel/auto_search/orchestrator.py | 9 +++++++-- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/MaxKernel/auto_search/algorithms/parallel_search.py b/MaxKernel/auto_search/algorithms/parallel_search.py index 37502a8..368f0ab 100644 --- a/MaxKernel/auto_search/algorithms/parallel_search.py +++ b/MaxKernel/auto_search/algorithms/parallel_search.py @@ -16,7 +16,6 @@ def __init__( self, num_parallel_runs: int = 2, strategies: Optional[List[str]] = None, - max_worker_retries: int = 1, agent_config: Optional[dict] = None, **kwargs, ): @@ -24,10 +23,6 @@ def __init__( raise ValueError( f"num_parallel_runs must be a positive integer, got {num_parallel_runs}." ) - if max_worker_retries < 1: - raise ValueError( - f"max_worker_retries must be at least 1, got {max_worker_retries}." - ) self.num_parallel_runs = num_parallel_runs if strategies is not None: @@ -43,7 +38,6 @@ def __init__( self.strategies = [""] * num_parallel_runs self.remaining_strategies = list(self.strategies) - self.max_worker_retries = max_worker_retries self.agent_config = agent_config self.worker = ADKSessionWorker() diff --git a/MaxKernel/auto_search/orchestrator.py b/MaxKernel/auto_search/orchestrator.py index 8aed02d..a5c2a36 100644 --- a/MaxKernel/auto_search/orchestrator.py +++ b/MaxKernel/auto_search/orchestrator.py @@ -17,7 +17,13 @@ def __init__( reference_code: str, graph_db_path: Optional[str] = None, max_concurrency: int = 2, + max_worker_retries: int = 1, ): + if max_worker_retries < 1: + raise ValueError( + f"max_worker_retries must be at least 1, got {max_worker_retries}." + ) + self.max_worker_retries = max_worker_retries # Resolve graph db path and run directory if not graph_db_path: workdir = os.environ.get("WORKDIR", os.getcwd()) @@ -152,8 +158,6 @@ def _post_step_hook(self) -> None: """ pass - # --- Concrete Flow --- - @abc.abstractmethod async def _execute_expansions(self, tasks: Any) -> List[Node]: """Executes expansion tasks to evaluate new kernels. @@ -163,6 +167,7 @@ async def _execute_expansions(self, tasks: Any) -> List[Node]: """ pass + # --- Workflow --- async def run(self) -> Node: logger.info( f"Starting search orchestration for problem: {self.graph.problem_id}"