diff --git a/MaxKernel/auto_agent/agent_client/auto_agent_client.py b/MaxKernel/auto_agent/agent_client/auto_agent_client.py index e1eb2fb..28eee18 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,38 @@ 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, initial_state: Optional[dict[str, Any]] = None + ) -> None: + self.session = await self.session_service.create_session( + app_name=self.app_name, + user_id=self.user_id, + session_id=self.session_id, + state=initial_state, ) - 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"} + def get_session_data(self) -> dict: + 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 +73,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 +124,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 +134,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..17958f1 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,102 +48,97 @@ 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() + 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}") -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__": diff --git a/MaxKernel/auto_search/README.md b/MaxKernel/auto_search/README.md new file mode 100644 index 0000000..debe2cc --- /dev/null +++ b/MaxKernel/auto_search/README.md @@ -0,0 +1,68 @@ +# MaxKernel Auto-Search + +This directory contains the automated search orchestration framework for `MaxKernel`. It leverages autonomous LLM-based worker agents to systematically search for optimized kernel implementations (like Pallas or Triton) for reference JAX benchmark problems. + +## Search Algorithms +The framework currently supports two primary search topologies: +1. **Parallel Search (`--algorithm parallel`)**: Dispatches multiple independent worker agents in parallel to optimize the reference kernel. The best performing candidate across all parallel runs is returned. +2. **Beam Search (`--algorithm beam`)**: A more structured tree-search approach that explores multiple optimization strategies (branches) concurrently, evaluates their latency, and keeps only the top `beam_size` candidates at each depth before expanding further. + +--- + +## Usage + +### 1. Run Search on a Single Problem +Use `run_search.py` to optimize a single benchmark directory. The directory must contain a valid `reference.py` file. + +**Example (Beam Search):** +```bash +python -m auto_search.run_search \ + --problem_dir /path/to/problem_dir \ + --algorithm beam \ + --beam_size 2 \ + --branches_per_node 2 \ + --max_depth 3 \ + --max_concurrency 4 +``` + +### 2. Run Batch Search on a Dataset +Use `run_batch_search.py` to optimize an entire dataset containing multiple problem directories concurrently. + +**Example (Parallel Search):** +```bash +python -m auto_search.run_batch_search \ + --data_dir /path/to/dataset_dir \ + --algorithm parallel \ + --num_parallel_runs 4 \ + --num_problem_concurrency 2 \ + --max_concurrency 8 +``` + +--- + +## Key Arguments + +### General Orchestration Arguments +* `--algorithm` : The algorithm to use (`parallel`, `beam`). Default is `parallel`. +* `--max_concurrency`: The maximum number of concurrent LLM agent workers to spawn at any given time. +* `--max_worker_retries`: Maximum times a worker is allowed to retry if the generated kernel crashes or fails syntax checks. +* `--strategies`: (Optional) Space-separated list of explicit strategies to explore (e.g., `--strategies "Fuse operations" "Vectorize loops"`). +* `--agent_config`: (Optional) JSON string of internal agent parameters. Available parameters: + * `"max_iterations"` (int): Maximum number of improvement loop iterations (default: 2). + * `"end_agent"` (string): The sub-agent pipeline stage at which to terminate early (e.g., `"validate_agent"`, `"test_run_agent"`). + * `"session_dir"` (string): Explicit path for session artifacts (usually auto-generated by the worker). + * Example: `'{"max_iterations": 5, "end_agent": "validate_agent"}'`. +* `--log_file`: File path to save the orchestration logs (creates separate `agent.log` and main logs). +* `--graph_db_path`: *(Single-problem only)* Explicit path to a `search_graph.json` file to resume a previously interrupted search. + +### Parallel Search Specific Arguments +* `--num_parallel_runs`: Number of independent agents to spawn to tackle the problem simultaneously. + +### Beam Search Specific Arguments +* `--beam_size`: Number of top-performing candidates to retain at each depth of the search tree. +* `--branches_per_node`: Number of new optimization strategies to explore from each retained node. +* `--max_depth`: Maximum depth of the beam search tree before terminating. +* `--keep_factor`: Factor of parent latency required to keep a candidate. For example, `1.0` means the child must be strictly faster (lower latency) than its parent to survive pruning. + +### Batch Search Specific Arguments +* `--num_problem_concurrency`: Number of distinct dataset problems to optimize concurrently. diff --git a/MaxKernel/auto_search/algorithms/beam_search.py b/MaxKernel/auto_search/algorithms/beam_search.py new file mode 100644 index 0000000..f3ed48b --- /dev/null +++ b/MaxKernel/auto_search/algorithms/beam_search.py @@ -0,0 +1,192 @@ +import asyncio +import logging +import random +from typing import List, Tuple + +from auto_search.graph import Node +from auto_search.orchestrator import SearchOrchestrator +from auto_search.strategies import TPU_PALLAS_OPTIMIZATION_STRATEGIES +from auto_search.worker import ADKSessionWorker + +logger = logging.getLogger(__name__) + + +class BeamSearchOrchestrator(SearchOrchestrator): + def __init__( + self, + beam_size: int = 2, + branches_per_node: int = 2, + max_depth: int = 2, + keep_factor: float = 1, + strategies: List[str] = TPU_PALLAS_OPTIMIZATION_STRATEGIES, + agent_config: dict = None, + **kwargs, + ): + self._validate_args(beam_size, branches_per_node, max_depth, keep_factor) + self.beam_size = beam_size + self.branches_per_node = branches_per_node + self.strategies = strategies + self.max_depth = max_depth + self.keep_factor = keep_factor + self.agent_config = agent_config or {"max_iterations": 1} + self.worker = ADKSessionWorker() + + self.current_depth = 0 + self.beam: List[Node] = [] + + super().__init__(**kwargs) + + if not self.graph.metadata: + self.beam = [self.graph.get_node(self.graph.root_id)] + self.update_metadata("current_depth", self.current_depth) + self.update_metadata( + "beam_node_ids", [node.node_id for node in self.beam] + ) + + def _validate_args( + self, + beam_size: int, + branches_per_node: int, + max_depth: int, + keep_factor: float, + ) -> None: + if beam_size < 1: + raise ValueError(f"beam_size must be at least 1, got {beam_size}.") + if branches_per_node < 1: + raise ValueError( + f"branches_per_node must be at least 1, got {branches_per_node}." + ) + if max_depth < 1: + raise ValueError(f"max_depth must be at least 1, got {max_depth}.") + if keep_factor <= 0: + raise ValueError(f"keep_factor must be positive, got {keep_factor}.") + + def _resume(self) -> None: + self.current_depth = self.graph.metadata.get("current_depth", 0) + beam_ids = self.graph.metadata.get("beam_node_ids", []) + self.beam = [ + node + for node in (self.graph.get_node(node_id) for node_id in beam_ids) + if node is not None + ] + logger.info( + f"Resumed Beam Search at depth {self.current_depth} with beam: {beam_ids}" + ) + + def _select_nodes_to_expand(self) -> List[Node]: + return self.beam + + def _generate_expansion_tasks( + self, nodes: List[Node] + ) -> List[Tuple[Node, str]]: + tasks = [] + for node in nodes: + selected_strategies = random.sample( + self.strategies, + min(self.branches_per_node, len(self.strategies)), + ) + for strategy in selected_strategies: + tasks.append((node, strategy)) + return tasks + + def _update_search_state(self, new_nodes: List[Node]) -> None: + candidates = [] + regressed_candidates = [] + for node in new_nodes: + if not node.is_valid_candidate: + logger.warning( + f"Node {node.node_id} failed Validity Check. " + "Adding to regressed candidates" + ) + regressed_candidates.append(node) + continue + + parent = self.graph.get_node(node.parent_id) + parent_latency = parent.evaluation.latency_ms + parent_latency = ( + parent_latency if parent_latency is not None else float("inf") + ) + + current_latency = node.evaluation.latency_ms + current_latency = ( + current_latency if current_latency is not None else float("inf") + ) + + if current_latency < parent_latency * self.keep_factor: + candidates.append(node) + else: + logger.info( + f"Node {node.node_id} failed Parent Regression Gate. " + "Adding to regressed candidates" + ) + regressed_candidates.append(node) + + candidates.sort(key=lambda n: n.evaluation.latency_ms) + + if len(candidates) < self.beam_size and regressed_candidates: + shortage = self.beam_size - len(candidates) + logger.warning( + f"Only {len(candidates)} candidates passed the Parent Regression Gate. " + f"Padding with the best {shortage} regressed candidates to keep search alive." + ) + regressed_candidates.sort( + key=lambda n: ( + n.evaluation.latency_ms + if n.evaluation.latency_ms is not None + else float("inf") + ) + ) + candidates.extend(regressed_candidates) + + self.beam = candidates[: self.beam_size] + self.update_metadata("beam_node_ids", [n.node_id for n in self.beam]) + + def _post_step_hook(self) -> None: + self.current_depth += 1 + self.update_metadata("current_depth", self.current_depth) + + def _should_terminate(self) -> bool: + return self.current_depth >= self.max_depth or not self.beam + + 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): + logger.info( + f"Task {task_idx}: Expanding {parent_node.node_id} using \n" + f" strategy '{strategy}' -> {node_id}. \n" + f"Attempt {attempt}/{self.max_worker_retries}." + ) + async with self._semaphore: + session_dir = f"{base_dir}_attempt_{attempt}" + node = await self.worker.expand_node( + node_id, + parent_node, + session_dir=session_dir, + reference_code=self.reference_code, + strategy=strategy, + agent_config=self.agent_config, + ) + if ( + node.execution_status == "SUCCESS" + or attempt == self.max_worker_retries + ): + self.graph.add_node(node) + logger.info( + f"Task {task_idx}: Finished {node_id} with status {node.execution_status} " + f"(Latency: {node.evaluation.latency_ms} ms)" + ) + return node + + logger.warning( + f"Task {task_idx} (strategy: {strategy}) failed attempt" + f" {attempt}/{self.max_worker_retries}: {node.execution_error}. Retrying..." + ) + await asyncio.sleep(2**attempt) + + futures = [ + run_task(i, parent, strat) for i, (parent, strat) in enumerate(tasks) + ] + return await asyncio.gather(*futures) diff --git a/MaxKernel/auto_search/algorithms/parallel_search.py b/MaxKernel/auto_search/algorithms/parallel_search.py new file mode 100644 index 0000000..368f0ab --- /dev/null +++ b/MaxKernel/auto_search/algorithms/parallel_search.py @@ -0,0 +1,139 @@ +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, + 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}." + ) + + 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.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) diff --git a/MaxKernel/auto_search/graph.py b/MaxKernel/auto_search/graph.py new file mode 100644 index 0000000..fe1c39b --- /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_dir: Optional[str] = None + 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." + ) diff --git a/MaxKernel/auto_search/orchestrator.py b/MaxKernel/auto_search/orchestrator.py new file mode 100644 index 0000000..a5c2a36 --- /dev/null +++ b/MaxKernel/auto_search/orchestrator.py @@ -0,0 +1,218 @@ +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, + 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()) + 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 + + @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 + + # --- Workflow --- + 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) + ) diff --git a/MaxKernel/auto_search/run_batch_search.py b/MaxKernel/auto_search/run_batch_search.py new file mode 100644 index 0000000..8ae8d43 --- /dev/null +++ b/MaxKernel/auto_search/run_batch_search.py @@ -0,0 +1,233 @@ +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", + ) + orch_group.add_argument( + "--max_worker_retries", + type=int, + default=1, + help="Max worker retries per expansion task", + ) + orch_group.add_argument( + "--strategies", + nargs="+", + type=str, + default=None, + help="List of strategy strings to explore", + ) + orch_group.add_argument( + "--agent_config", + type=str, + default=None, + help="JSON string of agent config parameters (e.g. '{\"max_iterations\": 5}')", + ) + # 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", + ) + # Beam Search Arguments + beam_group = parser.add_argument_group( + "Beam Search Arguments", + "Parameters specific to the 'beam' search algorithm.", + ) + beam_group.add_argument( + "--beam_size", + type=int, + default=2, + help="Size of the beam (number of candidates to keep per depth)", + ) + beam_group.add_argument( + "--branches_per_node", + type=int, + default=2, + help="Number of branches/strategies to explore per node in the beam", + ) + beam_group.add_argument( + "--max_depth", + type=int, + default=2, + help="Maximum depth of the beam search", + ) + beam_group.add_argument( + "--keep_factor", + type=float, + default=1.0, + help="Factor of parent latency to keep candidates (e.g. 1.0 means must not be worse than parent)", + ) + 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, + "max_worker_retries": args.max_worker_retries, + "strategies": args.strategies, + "agent_config": parsed_agent_config, + "num_parallel_runs": args.num_parallel_runs, + "beam_size": args.beam_size, + "branches_per_node": args.branches_per_node, + "max_depth": args.max_depth, + "keep_factor": args.keep_factor, + } + + 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..b3bd293 --- /dev/null +++ b/MaxKernel/auto_search/run_search.py @@ -0,0 +1,349 @@ +import argparse +import asyncio +import json +import logging +import os +import sys +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.beam_search import BeamSearchOrchestrator +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 == "beam": + strategies_kwargs = {} + if kwargs.get("strategies"): + strategies_kwargs["strategies"] = kwargs.get("strategies") + + return BeamSearchOrchestrator( + problem_id=problem_id, + reference_code=reference_code, + graph_db_path=graph_db_path, + max_concurrency=kwargs.get("max_concurrency", 2), + max_worker_retries=kwargs.get("max_worker_retries", 2), + beam_size=kwargs.get("beam_size", 2), + branches_per_node=kwargs.get("branches_per_node", 2), + max_depth=kwargs.get("max_depth", 2), + keep_factor=kwargs.get("keep_factor", 1.0), + agent_config=kwargs.get("agent_config"), + **strategies_kwargs, + ) + elif algorithm == "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() + + run_subdir = None + 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: + graph_exist = os.path.exists(graph_db_path) + + orchestrator = get_orchestrator( + algorithm=algorithm, + problem_id=problem_id, + reference_code=reference_code, + graph_db_path=graph_db_path, + **kwargs, + ) + + if graph_exist: + 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 run_subdir: + import shutil + + dest_dir = os.path.join(problem_dir, os.path.basename(run_subdir)) + shutil.copytree(run_subdir, dest_dir, dirs_exist_ok=True) + logger.info(f"Copied artifacts to {dest_dir}") + + 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", + ) + orch_group.add_argument( + "--max_worker_retries", + type=int, + default=1, + help="Max worker retries per expansion task", + ) + orch_group.add_argument( + "--strategies", + nargs="+", + type=str, + default=None, + help="List of strategy strings to explore", + ) + orch_group.add_argument( + "--agent_config", + type=str, + default=None, + help="JSON string of agent config parameters (e.g. '{\"max_iterations\": 5}')", + ) + # 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", + ) + # Beam Search Arguments + beam_group = parser.add_argument_group( + "Beam Search Arguments", + "Parameters specific to the 'beam' search algorithm.", + ) + beam_group.add_argument( + "--beam_size", + type=int, + default=2, + help="Size of the beam (number of candidates to keep per depth)", + ) + beam_group.add_argument( + "--branches_per_node", + type=int, + default=2, + help="Number of branches/strategies to explore per node in the beam", + ) + beam_group.add_argument( + "--max_depth", + type=int, + default=2, + help="Maximum depth of the beam search", + ) + beam_group.add_argument( + "--keep_factor", + type=float, + default=1.0, + help="Factor of parent latency to keep candidates (e.g. 1.0 means must not be worse than parent)", + ) + return parser.parse_args() + + +def setup_logging(log_file: Optional[str]): + """Configures console and file logging.""" + log_format = "%(asctime)s - %(levelname)s - %(message)s" + formatter = logging.Formatter(log_format) + + if log_file: + base, ext = os.path.splitext(log_file) + agent_log_path = f"{base}_agent{ext}" + main_log_path = log_file + else: + agent_log_path = "agent.log" + main_log_path = None + + # 1. Catch the rest of the logs in the Root Logger and route to agent_log_path + logging.basicConfig( + filename=agent_log_path, + level=logging.INFO, + format=log_format, + force=True, + ) + + # 2. Explicitly route only auto_search and __main__ to the main_log_path + auto_search_loggers = [ + logging.getLogger("auto_search"), + logging.getLogger("__main__"), + ] + + if main_log_path: + auto_search_handler = logging.FileHandler(main_log_path) + else: + auto_search_handler = logging.StreamHandler(sys.stdout) + + auto_search_handler.setFormatter(formatter) + + for logger in auto_search_loggers: + logger.propagate = False + logger.setLevel(logging.INFO) + if logger.hasHandlers(): + logger.handlers.clear() + logger.addHandler(auto_search_handler) + + +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, + "max_worker_retries": args.max_worker_retries, + "strategies": args.strategies, + "agent_config": agent_config, + # Parallel Search Arguments + "num_parallel_runs": args.num_parallel_runs, + # Beam Search Arguments + "beam_size": args.beam_size, + "branches_per_node": args.branches_per_node, + "max_depth": args.max_depth, + "keep_factor": args.keep_factor, + } + + 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() diff --git a/MaxKernel/auto_search/strategies.py b/MaxKernel/auto_search/strategies.py new file mode 100644 index 0000000..96d68a6 --- /dev/null +++ b/MaxKernel/auto_search/strategies.py @@ -0,0 +1,37 @@ +TPU_PALLAS_OPTIMIZATION_STRATEGIES = [ + "Reduce data movement", + "Overlap data movement and compute", + "Cache reused data in local memory instead of reloading from main memory", + "Loop tiling", + "Loop reordering and restructuring", + "Loop unrolling", + "Fuse operations", + "Use lower precision", + "Double buffering", + "Software pipelining", + "Hoist redundant operations out of loops", + "Eliminate redundant computation", + "Simplify or remove unnecessary code", + "Try new parameter values", + "Rewrite the algorithm to reduce total work", + "Place reduction axis last in grid to enable in-place SRAM accumulation without HBM round-trips", + "Align block dimensions to 8x128 tile boundaries to avoid wasted padding and register spills", + "Use scratch_shapes=[pltpu.VMEM(...)] for persistent high-precision accumulators during reduction loops", + "Maximize block sizes up to ~16 MB VMEM capacity to increase arithmetic intensity per pipeline step", + "Use scalar prefetch via PrefetchScalarGridSpec to load indices/metadata into SMEM without stalling vector core", + "Upcast bf16/int8 to float32 before elementwise ops, downcast only on final output write", + "Fuse transpose into lax.dot_general contraction dimensions instead of materializing transposed operands", + "Arrange grid iteration order so consecutive invocations reuse already-resident input slices", + "Increase pipeline buffer count beyond double buffering to hide memory latency for bandwidth-bound kernels", + "Generate random numbers inside kernel via hardware PRNG with key in SMEM instead of passing precomputed arrays", + "Avoid singleton dimensions in last two array axes to prevent full-tile waste per element", + "Reduce along second-to-last dimension rather than last dimension when possible", + "Prefer add/multiply over exp/tanh/division; restructure math to minimize expensive elementwise ops", + "Tune block sizes jointly — systematically vary BM, BN, BK together under the VMEM budget constraint (16 MiB including double-buffering)", + "Compute arithmetic intensity accounting for tiling amplification to predict compute-bound vs memory-bound regime", + "Minimize control flow inside kernels; consolidate into single basic blocks to avoid unrolling overhead", + "Pass all data as explicit kernel inputs with BlockSpec instead of closing over constants", + "Use pltpu.VMEM scoped scratch buffers for temporary storage within kernel lifetime", + "Balance block size against pipeline depth to amortize startup/drain bubble cost over enough iterations", + "Explicitly initialize accumulator buffers to zero on first reduction iteration since SRAM starts undefined", +] diff --git a/MaxKernel/auto_search/utils/visualize_graph.py b/MaxKernel/auto_search/utils/visualize_graph.py new file mode 100644 index 0000000..0b1eef5 --- /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) diff --git a/MaxKernel/auto_search/worker.py b/MaxKernel/auto_search/worker.py new file mode 100644 index 0000000..7329de1 --- /dev/null +++ b/MaxKernel/auto_search/worker.py @@ -0,0 +1,213 @@ +import json +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, + 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 reference_code and parent_node + self._prepare_inputs(session_dir, parent_node, reference_code) + # Prepare initial state + initial_state = self._prepare_initial_state(parent_node) + + # Run the agent and process results + try: + state = await self._run_agent( + node_id=node_id, + session_dir=session_dir, + strategy=strategy, + agent_config=agent_config, + initial_state=initial_state, + ) + return self._process_results( + node_id=node_id, + parent_node=parent_node, + strategy=strategy, + state=state, + session_dir=session_dir, + ) + except Exception as e: + logger.exception(f"Node {node_id} expansion failed: {e}") + return Node( + node_id=node_id, + parent_id=parent_node.node_id, + session_dir=session_dir, + 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, 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(reference_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: + 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) + + def _prepare_initial_state(self, parent_node: Node) -> Dict[str, Any]: + initial_state = {} + if parent_node and parent_node.evaluation: + initial_state = { + "kernel_compilation_status": { + "success": parent_node.evaluation.compiled, + "message": parent_node.evaluation.compilation_error, + }, + "test_results": { + "success": parent_node.evaluation.correct, + "output": parent_node.evaluation.test_error, + }, + "profiling_summary": parent_node.evaluation.profiling_summary, + } + return initial_state + + async def _run_agent( + self, + node_id: str, + session_dir: str, + strategy: Optional[str], + agent_config: Optional[Dict[str, Any]] = None, + initial_state: 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=node_id, + query=query, + agent=custom_agent, + ) + + await client.create_session(initial_state) + 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, + 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", []) + 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") or {} + compiled = comp_status.get("success", False) + compilation_error = None + if not compiled: + compilation_error = comp_status.get("message") or "" + final_errors = comp_status.get("final_errors") + if final_errors: + compilation_error += "\n\nFinal Errors:\n" + final_errors + + test_status = best_run.get("test_status") or {} + 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=node_id, + parent_id=parent_node.node_id, + session_dir=session_dir, + 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, + ), + )