diff --git a/MaxCode/agents/base.py b/MaxCode/agents/base.py index 176aa72..dadfb67 100644 --- a/MaxCode/agents/base.py +++ b/MaxCode/agents/base.py @@ -1,6 +1,7 @@ """Base class for all agents.""" import abc +import logging from typing import Any, Dict, Optional from agents import utils @@ -35,7 +36,22 @@ def generate( prompt = prompt_template.format(**prompt_vars) else: prompt = prompt_template - return self._model.generate(prompt) + logging.info( + "--- %s PROMPT ---\n%s\n--- END PROMPT ---", + self.agent_type.name, + prompt, + ) + try: + response = self._model.generate(prompt) + logging.info( + "--- %s RESPONSE ---\n%s\n--- END RESPONSE ---", + self.agent_type.name, + response, + ) + return response + except Exception as e: + logging.exception("LLM generation failed: %s", e) + raise @abc.abstractmethod def run(self, *args, **kwargs): diff --git a/MaxCode/agents/migration/orchestrator.py b/MaxCode/agents/migration/orchestrator.py deleted file mode 100644 index 85977d4..0000000 --- a/MaxCode/agents/migration/orchestrator.py +++ /dev/null @@ -1 +0,0 @@ -"""Orchestrates the specific kernel sub-agents.""" diff --git a/MaxCode/agents/migration/primary_agent.py b/MaxCode/agents/migration/primary_agent.py index bf6e843..ca9d2cc 100644 --- a/MaxCode/agents/migration/primary_agent.py +++ b/MaxCode/agents/migration/primary_agent.py @@ -1,18 +1,29 @@ """Primary orchestration agent for repository migration.""" import logging import os -from typing import Any +import re +import subprocess +import tempfile +from typing import Any, Tuple import models from agents import base from agents import utils from agents.migration import model_conversion_agent from agents.migration import single_file_agent -from agents.migration import validation_agent +from agents.migration.prompts import prompts from rag import rag_agent +MAX_DEBUG_ITERATIONS = 10 logger = logging.getLogger(__name__) +def _strip_markdown_formatting(text: str) -> str: + """Strips markdown and returns only the first python code block.""" + code_block_match = re.search(r"```(?:python)?\n?(.*?)\n?```", text, re.DOTALL) + if code_block_match: + return code_block_match.group(1).strip() + return text + class PrimaryAgent(base.Agent): """Primary orchestration agent for repository migration.""" @@ -46,6 +57,35 @@ def _convert_file(self, pytorch_code: str, file_path: str) -> str: return self._model_conversion_agent.run(pytorch_code) return self._single_file_agent.run(pytorch_code) + def _execute_test( + self, pytorch_code: str, jax_code: str, test_code: str + ) -> Tuple[bool, str]: + """Executes the test script and returns success status and output.""" + with tempfile.TemporaryDirectory() as tempdir: + torch_module_path = os.path.join(tempdir, "torch_module.py") + jax_module_path = os.path.join(tempdir, "jax_module.py") + test_script_path = os.path.join(tempdir, "test_script.py") + + with open(torch_module_path, "w") as f: + f.write(pytorch_code) + with open(jax_module_path, "w") as f: + f.write(jax_code) + with open(test_script_path, "w") as f: + f.write(test_code) + + try: + result = subprocess.run( + ["python3", test_script_path], + capture_output=True, + text=True, + check=True, + cwd=tempdir, + timeout=600, + ) + return True, result.stdout + except subprocess.CalledProcessError as e: + return False, e.stderr + def _validate_and_repair(self, pytorch_code: str, converted_code: str, file_path: str) -> str: """Validates converted code and repairs deviations if found. @@ -100,11 +140,16 @@ def run(self, repo_path: str) -> dict[str, str]: Args: repo_path: The path to the repository file or directory. + context: Optional raw context to use instead of RAG retrieval. Returns: A dictionary mapping original file paths to converted JAX code. + + Raises: + RuntimeError: If the code conversion and validation fails after + `MAX_DEBUG_ITERATIONS` attempts. """ - try: + if os.path.isfile(repo_path): with open(repo_path, "r", encoding="utf-8", errors="replace") as f: pytorch_code = f.read() logger.info("Converting %s ...", repo_path) @@ -126,10 +171,78 @@ def run(self, repo_path: str) -> dict[str, str]: repo_path: f"# Error: path {repo_path} is not a file or directory." } - graph = utils.build_dependency_graph(repo_path) - ordered_files = utils.topological_sort(graph) - converted_files: dict[str, str] = {} + if context is None: + rag_context_list = self._rag_agent.retrieve_context( + pytorch_code, top_k=7 + ) + rag_context = "\n\n".join([ + f"File: {c['file']}\n```python\n{c['text']}\n```" + for c in rag_context_list + ]) + else: + rag_context = context + + jax_code = _strip_markdown_formatting( + self.generate( + prompts.MIGRATE_MODULE_TO_JAX_PROMPT, + {"pytorch_code": pytorch_code, "rag_context": rag_context}, + ) + ) + + for i in range(MAX_DEBUG_ITERATIONS): + logging.info("Starting testing iteration %d.", i) + test_code = _strip_markdown_formatting( + self.generate( + prompts.EVALUATE_CODE_PROMPT, + {"pytorch_code": pytorch_code, "jax_code": jax_code}, + ) + ) + + if "NOTESTCASE" in test_code: + print( + "Test generation returned NOTESTCASE, assuming conversion is ok." + ) + return {repo_path: jax_code} + success, output = self._execute_test(pytorch_code, jax_code, test_code) + + if success: + print(f"Validation successful after {i} debugging iterations.") + logging.info( + "Validation successful after %d debugging iterations.", i + ) + return {repo_path: jax_code} + else: + traceback = output + logging.error( + "Validation failed on iteration %d. Traceback:\n%s", i, traceback + ) + logging.info("Starting debug iteration %d.", i + 1) + bug_analysis = self.generate( + prompts.BUG_ANALYSIS_PROMPT, + { + "pytorch_code": pytorch_code, + "jax_code": jax_code, + "test_code": test_code, + "traceback": traceback, + }, + ) + print(f"Bug analysis:\n{bug_analysis}") + logging.info("Bug analysis:\n%s", bug_analysis) + jax_code = _strip_markdown_formatting( + self.generate( + prompts.SELF_DEBUGGING_PROMPT, + { + "pytorch_code": pytorch_code, + "jax_code": jax_code, + "test_code": test_code, + "traceback": traceback, + "bug_analysis": bug_analysis, + "rag_context": rag_context, + }, + ) + ) + print(f"Attempting fix with new JAX code for iteration {i+1}.") for i, file_rel_path in enumerate(ordered_files, 1): file_path = os.path.join(repo_path, file_rel_path) logger.info("Converting file %d/%d: %s ...", i, len(ordered_files), @@ -143,4 +256,23 @@ def run(self, repo_path: str) -> dict[str, str]: ) converted_files[file_path] = converted_code - return converted_files + raise RuntimeError( + "Failed to convert and validate code after" + f" {MAX_DEBUG_ITERATIONS} iterations." + ) + elif os.path.isdir(repo_path): + graph = utils.build_dependency_graph(repo_path) + ordered_files = utils.topological_sort(graph) + converted_files: dict[str, str] = {} + + for file_rel_path in ordered_files: + file_path = os.path.join(repo_path, file_rel_path) + with open(file_path, "r", encoding="utf-8", errors="replace") as f: + pytorch_code = f.read() + converted_code = self._convert_file(pytorch_code, file_path) + converted_files[file_path] = converted_code + return converted_files + else: + return { + repo_path: f"# Error: path {repo_path} is not a file or directory." + } diff --git a/MaxCode/agents/migration/prompts/prompts.py b/MaxCode/agents/migration/prompts/prompts.py index 64219cd..a8f6933 100644 --- a/MaxCode/agents/migration/prompts/prompts.py +++ b/MaxCode/agents/migration/prompts/prompts.py @@ -5,6 +5,7 @@ 1. **Use Flax Linen with @nn.compact**: Define all submodules inline inside `@nn.compact def __call__`. Do NOT use a separate `setup()` method or NNX. + All Flax modules must be defined using the `@nn.compact` decorator. 2. **KV Cache**: Use pre-allocated fixed-size caches updated via `jax.lax.dynamic_update_slice`. NEVER grow the cache with `jnp.concatenate` or Python list appends -- that breaks XLA compilation. @@ -44,7 +45,31 @@ ordering exactly. Reshape to [B, T, num_k_heads, per_head_size] and split within each group. NEVER flatten to a single dimension and do a flat split -- this produces wrong tensors when num_k_heads != num_v_heads. -13. **Weight Initialization**: Match PyTorch initialization exactly. +13. **Hallucination Prevention**: Never use `num_feature_axes` as an attribute or argument. It is not a valid Flax/JAX parameter. Instead, use `axis` (e.g., `axis=-1` or `axis=(-2, -1)`) for normalization layers, `features` for linear layers, or `in_axes`/`out_axes` for `nn.scan`. +14. **Flax Scoping and Naming**: + - All Flax layers (e.g., `nn.Dense`, `nn.RNN`, `nn.GRUCell`, `nn.LSTMCell`) must be explicitly named using the `name=` argument in their constructor (e.g., `nn.Dense(..., name='fc')`). + - To avoid `flax.errors.NameInUseError`, every submodule created within a loop or list comprehension MUST have a unique name that includes the loop index (e.g., `name=f'layer_{{i}}'`). + - When using `nn.scan` inside a loop, provide a unique `name` to the scanned module instantiation, NOT the scan transformation itself: `nn.scan(...)(..., name=f'scan_{{i}}')`. +15. **Recurrent Layers (RNN/GRU/LSTM)**: + - If you define a custom RNN cell (e.g., to match PyTorch GRU/LSTM math), prefer using `nn.scan` directly over `nn.RNN` for better control. + - If using `nn.scan(ModuleClass, ...)`, the `in_axes` and `out_axes` parameters apply to the arguments and return values of the module's `__call__` method. By default, the first argument is treated as the `carry` and is NOT included in the `in_axes` count. For example, if `__call__(self, carry, x)`, use `in_axes=1`. + - The `out_axes` parameter applies only to the `output` part of the returned `(carry, output)` tuple. If the cell returns `(new_h, new_h)`, then `out_axes=1` indicates the second `new_h` is scanned. + - If you are defining a custom cell to be used specifically with `nn.RNN`, you MUST define `num_feature_axes = 1` (or the appropriate rank) as a class attribute. If using `nn.scan`, this attribute is not needed. + - `flax.linen` does not contain an `nn.GRU` layer. Use `flax.linen.RNN(flax.linen.GRUCell(...)` instead. + - GRU Math Accuracy: In PyTorch, the GRU candidate state calculation applies the "reset gate" to the entire hidden transformation, including its bias: `candidate = tanh(W_in * x + bias_in + reset_gate * (W_hn * h + bias_hn))`. Ensure the JAX implementation explicitly separates the input_to_hidden and hidden_to_hidden biases to match this calculation exactly; you MUST NOT sum `bias_ih` and `bias_hh` for GRU layers, as this prevents correct gating of `bias_hh` by the reset gate. The gate order for GRU weights/biases is Reset, Update, New. + - LSTM Math Accuracy: If the PyTorch LSTM has `bias=True`, it uses `bias_ih_l` and `bias_hh_l`. When mapping to a Flax `LSTMCell` which has a single `bias` parameter per gate, the Flax bias for each gate must be the SUM of the corresponding slices from `bias_ih_l` and `bias_hh_l`. The gate order for LSTM weights/biases is Input, Forget, Cell/Gate, Output. + - Return Values and Unpacking: PyTorch's `nn.GRU`/`nn.LSTM` returns `(output_sequence, final_hidden_state)`. In Flax, `nn.RNN` only returns `output_sequence` by default. You must set `return_carry=True` and ensure the code correctly unpacks both the carry and the output to avoid "too many values to unpack" errors. + - Multi-layer Logic: PyTorch's `nn.GRU`/`nn.LSTM` with `num_layers > 1` applies dropout between layers but not on the output of the final layer. To replicate this in Flax, if `num_layers > 1`, you must define a list of RNNs or cells and manually iterate, applying dropout only between layers 0..N-2. For example: `self.layers = [nn.RNN(nn.GRUCell(features,...), name=f'rnn_{{i}}') for i in range(num_layers)]`. If dropout > 0, define `self.dropouts = [nn.Dropout(rate=..., name=f'dropout_{{i}}') for i in range(num_layers - 1)]` and apply `self.dropouts[i]` to the output of `self.layers[i]` before passing it to `self.layers[i+1]`. If an initial hidden state `h0` for a multi-layer RNN is provided, it will have shape `(num_layers, batch_size, hidden_size)`, and you must pass `h0[i]` when calling the i-th layer. Ensure `__call__` accepts a `training: bool` argument to control dropout via `deterministic=not training`. + - Parameter Naming for Equivalence: If using `nn.RNN` with `GRUCell` or `LSTMCell`, name the internal cell `RNNCell_0` for testing alignment. + - Custom RNN cells used with `nn.RNN` must implement `initialize_carry(self, rng, input_shape)` to handle initial state, otherwise `nn.RNN` will fail during initialization. +16. **Numerical Parity**: To avoid subtle mismatches ("silent killers"), you MUST ensure: + - All Flax layers are defined with `dtype=jnp.float32` and `param_dtype=jnp.float32`. + - For recurrent layers (GRU/LSTM), always use `precision=jax.lax.Precision.HIGHEST` in all internal dot products to match PyTorch's 64-bit accumulation behavior during 32-bit inference. + - All matrix multiplications (e.g., `jnp.einsum`, `jnp.dot`) and convolutions specify `precision=jax.lax.Precision.HIGHEST`. + - Every layer explicitly sets `use_bias=True` or `use_bias=False` to exactly match the PyTorch layer. +17. **BatchNorm Momentum**: JAX momentum is the decay factor for old statistics (`x_new = momentum * x_old + (1 - momentum) * x_batch`), but PyTorch uses `1 - decay`. To ensure parity, you MUST set JAX momentum to `1 - pytorch_momentum`. +18. **Data Layout**: Standardize on `NHWC` (Channels Last) for JAX performance, but include necessary `jnp.transpose` operations at input/output boundaries to match PyTorch's `NCHW` oracle outputs. +19. **Weight Initialization**: Match PyTorch initialization exactly. When the source explicitly calls `nn.init.zeros_` on a layer, use `nn.initializers.zeros_init()`. When the source uses bare `nn.Linear()` with no explicit init, use the Flax default (lecun_normal) or @@ -53,30 +78,30 @@ RMSNorm (1+w): `nn.initializers.zeros_init()`. RMSNorm (w): `nn.initializers.ones_init()`. Check each nn.Parameter in the source and match its init. -14. **Train/Eval Mode**: Flax modules do NOT have a `.train` attribute or +20. **Train/Eval Mode**: Flax modules do NOT have a `.train` attribute or `.eval()` / `.train()` methods. NEVER write `model.train = True` or `model.train = False` -- this does nothing in Flax and silently produces incorrect behavior. Instead, pass `deterministic=False` for training and `deterministic=True` for evaluation as an argument to `__call__` / `model.apply()`. All stochastic layers (Dropout, router noise) must check the `deterministic` flag. -15. **Preserve ALL Source Components**: Convert EVERY class, function, and +21. **Preserve ALL Source Components**: Convert EVERY class, function, and method from the source. Do NOT merge base classes into subclasses, do NOT drop utility classes or metric functions, and do NOT omit `get_config()` or serialization methods. If the source has `ExpertBase` and `FFNExpert`, convert both. If the source has a `MoEMetrics` class, convert it. -16. **Preserve Default Values Exactly**: All default parameter values in the +22. **Preserve Default Values Exactly**: All default parameter values in the JAX output must match the PyTorch source EXACTLY. Do NOT change any numeric default -- not capacity factors, not dropout rates, not epsilon values, not learning rates, not layer counts. Even if you believe a different value is "better" or "more stable", use the source value. Changed defaults silently alter model behavior and break reproducibility. -17. **Preserve Exact Reduction Operations**: When the source uses `.mean()`, +23. **Preserve Exact Reduction Operations**: When the source uses `.mean()`, use `jnp.mean()`. When the source uses `.sum()`, use `jnp.sum()`. NEVER substitute one reduction for another. `torch.mean(x, dim=N)` maps to `jnp.mean(x, axis=N)`. `torch.sum(x, dim=N)` maps to `jnp.sum(x, axis=N)`. The dim/axis integer stays the same. -18. **Preserve Method Placement**: If the source defines a method or attribute +24. **Preserve Method Placement**: If the source defines a method or attribute on a specific class, keep it on that class in the JAX output. Do NOT relocate methods between classes or replace instance methods with standalone functions unless the JAX idiom requires it. @@ -99,31 +124,200 @@ classes). If the source has it, the output must have it. """ -PYTORCH_TO_JAX_SINGLE_FILE_PROMPT = """You are an expert in JAX and PyTorch. -Your task is to convert the following PyTorch code to JAX. -If it is helpful, you can use the following JAX code snippets as context for -functionality that might be similar to your conversion task: +MIGRATE_MODULE_TO_JAX_PROMPT = """ +You are an expert AI code translator specializing in converting PyTorch code to JAX. +Your task is to convert code written in PyTorch, NumPy, or similar frameworks into +functionally equivalent JAX code using appropriate JAX libraries (jax.numpy, +Flax, Optax, etc.). + +Use the following repository locations as high-quality JAX code context to inform the conversion. When sufficient tokens are not available, prioritize them in the following order: +- Main: https://github.com/AI-Hypercomputer/maxtext/tree/main/src/MaxText +- Layers folder: https://github.com/AI-Hypercomputer/maxtext/tree/main/src/MaxText/layers +- Kernels folder: https://github.com/AI-Hypercomputer/maxtext/tree/main/src/MaxText/kernels +- Multimodal folder: https://github.com/AI-Hypercomputer/maxtext/tree/main/src/MaxText/multimodal +- Inference folder: https://github.com/AI-Hypercomputer/maxtext/tree/main/src/MaxText/inference + +The rest of the repository files can be ignored. --- {rag_context} --- + +Guidelines: +- Preserve the original code structure (functions, classes, variable names) unless modification is necessary for compatibility. +- Think step-by-step before generating code: first, identify all PyTorch layers, operations, and data transformations; second, determine their JAX/Flax counterparts; and finally, generate the equivalent JAX code based on this analysis. +- Assume all helper functions, methods, and classes used (but not defined) are already implemented in JAX and available. +- Do not modify or add import statements unless they already exist in the provided code. +- Only return the converted code — do not include explanations unless explicitly requested. +- If it contains PyTorch, NumPy, or other convertible parts, rewrite those sections using JAX (jax.numpy, Flax, Optax). +- Return no code change if the provided code is purely generic Python (i.e., no PyTorch/NumPy/JAX operations to convert). +- Ensure that the generated code: + - Is functionally equivalent to the original PyTorch code block. + - Uses idiomatic JAX practices (e.g., jax.numpy instead of numpy, vectorization where possible). + - Maintains the original architecture and logic, just rewritten in JAX. + - Preserves original function/class names unless absolutely necessary to change. +- Do not generate function calls or tool calls. Your response should only contain the JAX code block. + The PyTorch code to convert is as follows: ```python {pytorch_code} ``` -Please think step by step about the conversion process before generating the code. Then, provide the JAX equivalent of the PyTorch code above. -Ensure that the JAX code is idiomatic and follows best practices, such as using -pure functions and handling random number generation correctly with JAX's PRNG -keys. Only return the Python code block for the JAX implementation. +Ensure all imports are included at the top of the generated code. +Only return the Python code block for the JAX implementation. +""" + JAX_BEST_PRACTICES + +EVALUATE_CODE_PROMPT = """You are an expert machine learning engineer and automated testing specialist with deep +knowledge of Python, NumPy, PyTorch, JAX (Including libraries such as Flax, Flax.nnx and Optax). + +Your role is to generate a comprehensive test suite that compares a PyTorch code block and a JAX code block for functional equivalence. +The test suite should: +1. Validate the PyTorch module independently. +2. Validate the JAX module independently. +3. Compare their outputs across multiple randomized inputs using `numpy.allclose`, and verify the JAX model's parameter structure using a dummy initialization before generating the final equivalence test to ensure correct parameter mapping. If custom classes are present in JAX code (e.g. CustomLSTMCell), ensure tests use them correctly and that parameter mapping from PyTorch state_dict to JAX is accurate. + +Guidelines: +- Assume helper functions and classes not defined in the code are already implemented and available. +- Do not add or modify import statements unless they exist in the provided code. +- Only return test code (no explanations) unless explicitly asked. +- For trivial or untestable code, return `NOTESTCASE`. +- When comparing PyTorch and JAX: + - You will be given Pytorch code and JAX code snippets. Assume they can be imported or used directly. + - Accept an optional `#entry_point` that identifies the function or class to invoke. + - Automatically generate randomized test inputs for shapes like `(2,3)`, `(4,)`, etc. + - Write clear assertions for: + - Output validity (no errors or exceptions) + - Output comparison (`np.allclose`) + - For GRU layers, PyTorch's `nn.GRU` uses separate `bias_ih_l` and `bias_hh_l`. When mapping to Flax, these biases MUST remain separate and be assigned to the correct kernel transformations (e.g. input and hidden transformations) to ensure correct gating: n_t = tanh(W_in*x_t + b_in + r_t * (W_hn*h_{{t-1}} + b_hh)). Unlike LSTM, GRU input and hidden biases MUST NOT be summed. + - For LSTM layers, PyTorch's `nn.LSTM` concatenates gate weights (i, f, g, o) in `weight_ih_l` and `weight_hh_l`, while Flax's `LSTMCell` may store them as separate parameters (e.g., `ii/kernel`, `if/kernel`, `ig/kernel`, `io/kernel` for input weights and `hi/kernel`, `hf/kernel`, `hg/kernel`, `ho/kernel` for recurrent weights). When mapping PyTorch `state_dict` to JAX parameters for equivalence testing, you MUST split the PyTorch weights into 4 parts for each gate and assign them to the corresponding Flax parameters. For a hidden size `H`, slice PyTorch weights like `weight_ih_l[0:H, :]`, `weight_ih_l[H:2*H, :]`, etc. for gates i, f, g, o respectively. PyTorch's `bias_ih_l` and `bias_hh_l` must also be split into 4 slices each, and the corresponding slices must be SUMMED (`bias_ih_l_gate + bias_hh_l_gate`) to form the single bias parameter for each JAX gate. If `flax.linen.RNN` or `nn.scan` is used with `LSTMCell`, parameters may be nested inside a `scan` scope (e.g., `params['lstm']['scan(LSTMCell_0)']['...']`); ensure parameter mapping accounts for this nesting by inspecting the parameter tree via `jax.tree_util.tree_map(lambda x: x.shape, variables['params'])` and adjusting the mapping logic accordingly. If the assumed mapping structure doesn't match the initialized JAX model, raise an error. + - For Transformer layers (`nn.MultiheadAttention`), PyTorch combines weights into `in_proj_weight`. You MUST generate test code that correctly splits and reshapes this combined weight into the separate `query`, `key`, and `value` kernels and biases expected by Flax's `MultiHeadDotProductAttention` for weight mapping. +- Dynamic Parameter Inspection: + - The generated test script MUST first initialize the JAX model and print its parameter structure using `jax.tree_util.tree_map(lambda x: x.shape, variables['params'])`. + - Use this structure to dynamically verify that the paths used in the weight mapping actually exist. For multi-layer models, check for both `params['rnn_{{i}}']` and `params['layer_{{i}}']` patterns. + - If a `LayerWrapper` is used, the cell parameters will be under `params['layer_{{i}}']['cell']`. + - provide a helpful error message showing the expected vs. actual structure if they don't match. +- Do not generate function calls or tool calls. Your response should only contain the Python test script. + +Here is the PyTorch code: +```python +{pytorch_code} +``` + +Here is the JAX code: +```python +{jax_code} +``` + +Please generate a Python test script that saves the pytorch code to 'torch_module.py', +the jax code to 'jax_module.py', imports them, and runs comparison tests. +Only return the Python code block for the test script. +""" + +BUG_ANALYSIS_PROMPT = """You are an expert bug analyzer. +You are tasked with debugging a script failure, likely from a test comparing +PyTorch and JAX code conversion. +You will be given the PyTorch code, the converted JAX code, the test script that failed, +and the execution traceback from the test. +Your goal is to summarize the execution traceback and explain the root cause of the errors. +You do not need to propose solutions to fix the errors. + +PyTorch code: +```python +{pytorch_code} +``` + +JAX code: +```python +{jax_code} +``` + +Test script: +```python +{test_code} +``` + +Execution traceback: +``` +{traceback} +``` + +Please summarize the execution traceback and explain the root cause of the errors. +Do not generate function calls or tool calls. Your response should only contain the text analysis. +""" + +SELF_DEBUGGING_PROMPT = """ +You are an expert JAX programmer tasked with debugging JAX code based on PyTorch-to-JAX conversion errors. +You are continuing a debugging session. The previous JAX code you generated failed validation against the original PyTorch code. +Your job is to fix the JAX code based on the failing test traceback and bug analysis. + +Your task is to: +- Identify the cause of the error from the provided bug analysis and stack trace. +- Modify only the necessary parts of the previous JAX code to fix the error shown in Execution Traceback, following the JAX/Flax best practices below. +- The fix must be targeted. Do not change the core logic or intended functionality of the original code. +- You must import and use the functions or classes from the provided library files. Do not copy or redefine them in your main script. +- If you see `AttributeError: module 'flax.linen' has no attribute 'GRU'`, replace usage of `nn.GRU` with `nn.RNN(nn.GRUCell(...))`. +- If test failures are due to small numerical discrepancies, check rules for **Numerical Parity** and **BatchNorm Momentum** below. +- If the error is due to unavailable or incompatible external dependencies, replace them with equivalent or minimal alternatives. +- Do not rewrite working parts unless required for the fix. +- Do not use `try...except` blocks to catch, suppress, or ignore the original error. The fix must address the root cause of the problem. + +Ensure the generated code is correct and directly runnable. +Only return the fixed JAX code block (no explanations). + +Use the following repository locations as high-quality JAX code context to inform the conversion. When sufficient tokens are not available, prioritize them in the following order: +- Main: https://github.com/AI-Hypercomputer/maxtext/tree/main/src/MaxText +- Layers folder: https://github.com/AI-Hypercomputer/maxtext/tree/main/src/MaxText/layers +- Kernels folder: https://github.com/AI-Hypercomputer/maxtext/tree/main/src/MaxText/kernels +- Multimodal folder: https://github.com/AI-Hypercomputer/maxtext/tree/main/src/MaxText/multimodal +- Inference folder: https://github.com/AI-Hypercomputer/maxtext/tree/main/src/MaxText/inference + +The rest of the repository files can be ignored. +--- +{rag_context} +--- + +Original PyTorch code: +```python +{pytorch_code} +``` + +Previous JAX code (failed): +```python +{jax_code} +``` + +Test script: +```python +{test_code} +``` + +Execution traceback: +``` +{traceback} +``` + +Bug analysis: +``` +{bug_analysis} +``` + +Please provide the corrected JAX code. +Do not generate function calls or tool calls. Your response should only contain the fixed JAX code block. +Only return the Python code block for the JAX implementation. """ + JAX_BEST_PRACTICES PYTORCH_TO_JAX_REPO_PROMPT = """You are an expert in JAX and PyTorch. Your task is to convert a repository from PyTorch to JAX. You will be given a file path and the content of the file. You need to convert the given file from PyTorch to JAX, considering its context within the repository. -If it is helpful, you can use the following JAX code snippets as context for -functionality that might be similar to your conversion task: +Use the following repository locations as high-quality JAX code context to inform the conversion. When sufficient tokens are not available, prioritize them in the following order: +- Main: https://github.com/AI-Hypercomputer/maxtext/tree/main/src/MaxText +- Layers folder: https://github.com/AI-Hypercomputer/maxtext/tree/main/src/MaxText/layers +- Kernels folder: https://github.com/AI-Hypercomputer/maxtext/tree/main/src/MaxText/kernels +- Multimodal folder: https://github.com/AI-Hypercomputer/maxtext/tree/main/src/MaxText/multimodal +- Inference folder: https://github.com/AI-Hypercomputer/maxtext/tree/main/src/MaxText/inference + +The rest of the repository files can be ignored. --- {rag_context} --- @@ -145,8 +339,14 @@ HF_TO_JAX_SINGLE_FILE_PROMPT = """You are an expert in JAX and PyTorch, with special expertise in HuggingFace Transformers. Your task is to convert the following HuggingFace Transformers code (which uses PyTorch) to JAX. -If it is helpful, you can use the following JAX code snippets as context for -functionality that might be similar to your conversion task: +Use the following repository locations as high-quality JAX code context to inform the conversion. When sufficient tokens are not available, prioritize them in the following order: +- Main: https://github.com/AI-Hypercomputer/maxtext/tree/main/src/MaxText +- Layers folder: https://github.com/AI-Hypercomputer/maxtext/tree/main/src/MaxText/layers +- Kernels folder: https://github.com/AI-Hypercomputer/maxtext/tree/main/src/MaxText/kernels +- Multimodal folder: https://github.com/AI-Hypercomputer/maxtext/tree/main/src/MaxText/multimodal +- Inference folder: https://github.com/AI-Hypercomputer/maxtext/tree/main/src/MaxText/inference + +The rest of the repository files can be ignored. --- {rag_context} --- @@ -168,8 +368,14 @@ single JAX/Flax file. You MUST convert ALL classes, helper functions, constants, and configuration dataclasses -- not just one class. -If it is helpful, you can use the following JAX code snippets as context for -functionality that might be similar to your conversion task: +Use the following repository locations as high-quality JAX code context to inform the conversion. When sufficient tokens are not available, prioritize them in the following order: +- Main: https://github.com/AI-Hypercomputer/maxtext/tree/main/src/MaxText +- Layers folder: https://github.com/AI-Hypercomputer/maxtext/tree/main/src/MaxText/layers +- Kernels folder: https://github.com/AI-Hypercomputer/maxtext/tree/main/src/MaxText/kernels +- Multimodal folder: https://github.com/AI-Hypercomputer/maxtext/tree/main/src/MaxText/multimodal +- Inference folder: https://github.com/AI-Hypercomputer/maxtext/tree/main/src/MaxText/inference + +The rest of the repository files can be ignored. --- {rag_context} --- diff --git a/MaxCode/agents/migration/single_file_agent.py b/MaxCode/agents/migration/single_file_agent.py index 733c027..7bc991a 100644 --- a/MaxCode/agents/migration/single_file_agent.py +++ b/MaxCode/agents/migration/single_file_agent.py @@ -52,7 +52,7 @@ def run(self, pytorch_code: str) -> str: for c in rag_context_list ]) generated_code = self.generate( - prompts.PYTORCH_TO_JAX_SINGLE_FILE_PROMPT, + prompts.MIGRATE_MODULE_TO_JAX_PROMPT, {"pytorch_code": pytorch_code, "rag_context": rag_context}, ) return self._strip_markdown_formatting(generated_code) diff --git a/MaxCode/mcp_server/adk_agents.py b/MaxCode/mcp_server/adk_agents.py index 8730fe4..31eeb45 100644 --- a/MaxCode/mcp_server/adk_agents.py +++ b/MaxCode/mcp_server/adk_agents.py @@ -1,5 +1,6 @@ """ADK agent definitions.""" +import models from tools import evaluation_tool from tools import migration_tool from google.adk.agents.llm_agent import LlmAgent as Agent @@ -7,7 +8,7 @@ migration_agent = Agent( name="migration_agent", - model=Gemini(), + model=Gemini(model=models.GeminiModel.GEMINI_3_1_PRO_PREVIEW.value), description=( "Handles end-to-end code migration tasks, such as converting PyTorch" " to JAX, generating oracle data, and creating equivalence tests." @@ -15,6 +16,9 @@ instruction="""You are the migration specialist. Your task is to perform end-to-end migrations by orchestrating tools in a STRICT sequential order. You must extract the `api_key` and an optional `model_name` from the user prompt. Pass `api_key` to all tools that require it. If `model_name` is provided in the prompt, pass it as the `model_name` argument to `convert_code`, `generate_model_configs`, and `run_equivalence_tests`. +**Critical Naming Rule:** To avoid `flax.errors.NameInUseError`, ensure that any submodules created within a loop or list comprehension (like `CustomGRUCell` inside `nn.RNN`) have unique names that include the loop index (e.g., use an f-string like name=f'cell_index'). This applies to all layers and cells. + + Here is the sequence: 1. Call `convert_code` with `source_path`, `destination`, and `api_key` to translate PyTorch code to JAX. This tool returns a JSON string like: `{"dest_path": "/path/to/dest/timestamp", "mapping_path": "/path/to/dest/timestamp/mapping.json", "original_source_dir": "/path/to/dest/timestamp/original_source"}`. @@ -43,7 +47,7 @@ evaluation_agent = Agent( name="evaluation_agent", - model=Gemini(), + model=Gemini(model=models.GeminiModel.GEMINI_3_1_PRO_PREVIEW.value), description=( "Handles the generation of evaluation configurations and scripts." ), diff --git a/MaxCode/mcp_server/gemini-extension.json b/MaxCode/mcp_server/gemini-extension.json index 69fc216..b329a04 100644 --- a/MaxCode/mcp_server/gemini-extension.json +++ b/MaxCode/mcp_server/gemini-extension.json @@ -5,7 +5,7 @@ "dev-server": { "command": "python3.11", "args": ["-m", "mcp_server.primary_agent_server"], - "timeout": 120000, + "timeout": 1800000, "env": { "GOOGLE_API_KEY": "${env:GOOGLE_API_KEY}" } diff --git a/MaxCode/mcp_server/primary_agent_server.py b/MaxCode/mcp_server/primary_agent_server.py index 86fded2..77a7e8d 100644 --- a/MaxCode/mcp_server/primary_agent_server.py +++ b/MaxCode/mcp_server/primary_agent_server.py @@ -7,6 +7,7 @@ from absl import app from mcp_server import adk_agents from google.adk.agents.llm_agent import LlmAgent as Agent +from google.adk.models.google_llm import Gemini from google.adk.runners import Runner from google.adk.sessions.in_memory_session_service import InMemorySessionService from google.genai import types @@ -60,6 +61,21 @@ async def _execute_adk_agent( # 1. Handle API Key (Set in env so ADK model client finds it) if effective_api_key: os.environ["GOOGLE_API_KEY"] = effective_api_key + + model_match = re.search( + r"(?:model_name=|using model |using the )([\w.-]+)", + prompt, + re.IGNORECASE, + ) + if model_match: + model_name = model_match.group(1).strip().rstrip(".") + logging.info("Model name extracted from prompt: %s", model_name) + agent.model = Gemini(model=model_name) + + if isinstance(agent.model, Gemini): + logging.info("Using model: %s", agent.model.model) + else: + logging.info("Using model: %s", agent.model) session_service = InMemorySessionService() runner = Runner( agent=agent, diff --git a/MaxCode/models.py b/MaxCode/models.py index 8de7371..240c934 100644 --- a/MaxCode/models.py +++ b/MaxCode/models.py @@ -13,7 +13,7 @@ class GeminiModel(enum.Enum): GEMINI_2_5_PRO = "gemini-2.5-pro" GEMINI_2_5_FLASH = "gemini-2.5-flash" GEMINI_3_0_PRO = "gemini-3.0-pro" - GEMINI_3_0_FLASH = "gemini-3.0-flash" + GEMINI_3_0_FLASH = "gemini-3-flash-preview" GEMINI_3_1_PRO_PREVIEW = "gemini-3.1-pro-preview" @@ -30,7 +30,7 @@ class GeminiTool: def __init__( self, - model_name: GeminiModel | str = GeminiModel.GEMINI_2_5_PRO, + model_name: GeminiModel | str = GeminiModel.GEMINI_3_1_PRO_PREVIEW, system_instruction=None, api_key=None, ): diff --git a/MaxCode/setup_env.sh b/MaxCode/setup_env.sh index ec6da58..04a89c6 100755 --- a/MaxCode/setup_env.sh +++ b/MaxCode/setup_env.sh @@ -49,7 +49,7 @@ source "$VENV_DIR"/bin/activate # Install dependencies pip install --upgrade pip --index-url https://pypi.org/simple -pip install --upgrade google-genai numpy google-adk absl-py faiss-cpu torch flax jax[cpu] --index-url https://pypi.org/simple +pip install --upgrade google-genai numpy google-adk absl-py faiss-cpu torch flax jax[cpu] pytest --index-url https://pypi.org/simple # Check for GOOGLE_API_KEY if [ -z "$GOOGLE_API_KEY" ]; then diff --git a/MaxKernel/hitl_agent/agent.py b/MaxKernel/hitl_agent/agent.py index c50b6cc..b34fe68 100644 --- a/MaxKernel/hitl_agent/agent.py +++ b/MaxKernel/hitl_agent/agent.py @@ -4,58 +4,60 @@ for the human-in-the-loop kernel generation process. """ -from hitl_agent.custom_types import CustomLlmAgent - -from hitl_agent.constants import MODEL_NAME +from hitl_agent.callbacks import ( + add_pallas_docs, + add_workdir_callback, + get_tpu_version_callback, +) from hitl_agent.config import ( - model_config, - thinking_planner, + model_config, + thinking_planner, ) -from hitl_agent.tools.tools import filesystem_tool_r -from hitl_agent.callbacks import ( - get_tpu_version_callback, - add_workdir_callback, - add_pallas_docs, +from hitl_agent.constants import MODEL_NAME +from hitl_agent.custom_types import CustomLlmAgent +from hitl_agent.prompts import interactive_prompt +from hitl_agent.subagents.explanation import explanation_agent +from hitl_agent.subagents.gpu_to_jax_agent.agent import ( + gpu_to_jax_agent, ) from hitl_agent.subagents.kernel_writing import ( - plan_kernel_agent, - implement_kernel_agent, - validate_kernel_compilation_agent, + implement_kernel_agent, + plan_kernel_agent, + validate_kernel_compilation_agent, ) +from hitl_agent.subagents.profiling import profile_agent from hitl_agent.subagents.testing import ( - validated_test_generation_agent, - unified_test_agent, + unified_test_agent, + validated_test_generation_agent, ) -from hitl_agent.subagents.profiling import profile_agent -from hitl_agent.subagents.explanation import explanation_agent -from hitl_agent.subagents.gpu_to_jax_agent.agent import ( - gpu_to_jax_agent,) -from hitl_agent.prompts import interactive_prompt +from hitl_agent.tools.tools import filesystem_tool_r # Root orchestration agent root_agent = CustomLlmAgent( - name="KernelGenerationOrchestrationAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=thinking_planner, - before_agent_callback=[ - add_pallas_docs, get_tpu_version_callback, add_workdir_callback - ], - sub_agents=[ - explanation_agent, # Provides explanations - plan_kernel_agent, # Step 1: Create/revise plan - implement_kernel_agent, # Step 2: Implement kernel - validate_kernel_compilation_agent, # Step 3: Validate compilation - validated_test_generation_agent, # Step 4: Generate and validate tests - unified_test_agent, # Step 5: Run tests and provide summary - profile_agent, # Step 6: Profile for bottlenecks - gpu_to_jax_agent, # GPU-to-JAX conversion - ], - tools=[filesystem_tool_r - ], # Read-only access - orchestrator delegates writes to sub-agents - instruction=interactive_prompt.PROMPT, - description= - "Orchestrates the human-in-the-loop kernel generation process with GPU to JAX conversion capability.", + name="KernelGenerationOrchestrationAgent", + model=MODEL_NAME, + generate_content_config=model_config, + planner=thinking_planner, + before_agent_callback=[ + add_pallas_docs, + get_tpu_version_callback, + add_workdir_callback, + ], + sub_agents=[ + explanation_agent, # Provides explanations + plan_kernel_agent, # Step 1: Create/revise plan + implement_kernel_agent, # Step 2: Implement kernel + validate_kernel_compilation_agent, # Step 3: Validate compilation + validated_test_generation_agent, # Step 4: Generate and validate tests + unified_test_agent, # Step 5: Run tests and provide summary + profile_agent, # Step 6: Profile for bottlenecks + gpu_to_jax_agent, # GPU-to-JAX conversion + ], + tools=[ + filesystem_tool_r + ], # Read-only access - orchestrator delegates writes to sub-agents + instruction=interactive_prompt.PROMPT, + description="Orchestrates the human-in-the-loop kernel generation process with GPU to JAX conversion capability.", ) -__all__ = ['root_agent'] +__all__ = ["root_agent"] diff --git a/MaxKernel/hitl_agent/callbacks.py b/MaxKernel/hitl_agent/callbacks.py index 3b47cf8..0e73049 100644 --- a/MaxKernel/hitl_agent/callbacks.py +++ b/MaxKernel/hitl_agent/callbacks.py @@ -1,32 +1,34 @@ """Callback utilities for HITL kernel generation agents.""" -import os import json import logging +import os from typing import Any, Dict, Optional -from google.adk.tools import BaseTool, ToolContext + from google.adk.agents.callback_context import CallbackContext from google.adk.models import LlmResponse -from hitl_agent.config import WORKDIR, TPU_VERSION +from google.adk.tools import BaseTool, ToolContext + +from hitl_agent.config import TPU_VERSION, WORKDIR from hitl_agent.knowledge_base import pallas_docs, pallas_profiling_docs def create_path_saver(state_key: str): """ - Factory function that creates a callback to save file paths to a specific state key. + Factory function that creates a callback to save file paths to a specific state key. - Args: - state_key: The key in tool_context.state where the file path will be saved. + Args: + state_key: The key in tool_context.state where the file path will be saved. - Returns: - A callback function compatible with after_tool_callback signature. - """ + Returns: + A callback function compatible with after_tool_callback signature. + """ def save_path( - tool: BaseTool, - args: Dict[str, Any], - tool_context: ToolContext, - tool_response: Optional[Dict], + tool: BaseTool, + args: Dict[str, Any], + tool_context: ToolContext, + tool_response: Optional[Dict], ) -> Optional[Dict]: # MCP filesystem tools may have different naming patterns # Check for both snake_case and potential prefixed versions @@ -35,7 +37,7 @@ def save_path( if file_path: tool_context.state[state_key] = file_path logging.info( - f"Saved file path to {state_key}: {file_path} (from tool: {tool.name})" + f"Saved file path to {state_key}: {file_path} (from tool: {tool.name})" ) return None @@ -43,22 +45,24 @@ def save_path( def save_kernel_file_paths( - tool: BaseTool, - args: Dict[str, Any], - tool_context: ToolContext, - tool_response: Optional[Dict], + tool: BaseTool, + args: Dict[str, Any], + tool_context: ToolContext, + tool_response: Optional[Dict], ) -> Optional[Dict]: """ - Saves kernel file paths with semantic naming based on read order. - First file read = base_kernel_path, Second file read = optimized_kernel_path. - This callback is used by agents that need to compare two kernels. - """ + Saves kernel file paths with semantic naming based on read order. + First file read = base_kernel_path, Second file read = optimized_kernel_path. + This callback is used by agents that need to compare two kernels. + """ if tool.name == "read_file": file_path = args.get("path", None) # If base_kernel_path not set, this is the first file (base) - if ("base_kernel_path" not in tool_context.state or - not tool_context.state["base_kernel_path"]): + if ( + "base_kernel_path" not in tool_context.state + or not tool_context.state["base_kernel_path"] + ): tool_context.state["base_kernel_path"] = file_path logging.info(f"Set base kernel path: {file_path}") # Otherwise, this is the second file (optimized) @@ -70,10 +74,10 @@ def save_kernel_file_paths( def save_kernel_and_plan_paths( - tool: BaseTool, - args: Dict[str, Any], - tool_context: ToolContext, - tool_response: Optional[Dict], + tool: BaseTool, + args: Dict[str, Any], + tool_context: ToolContext, + tool_response: Optional[Dict], ) -> Optional[Dict]: """Saves both optimized_kernel_path and kernel_plan_path during implementation.""" if "read" in tool.name.lower() or "write" in tool.name.lower(): @@ -83,23 +87,23 @@ def save_kernel_and_plan_paths( if "plan" in file_path.lower() and file_path.endswith(".md"): tool_context.state["kernel_plan_path"] = file_path logging.info( - f"Saved plan path to kernel_plan_path: {file_path} (from tool: {tool.name})" + f"Saved plan path to kernel_plan_path: {file_path} (from tool: {tool.name})" ) # Otherwise assume it's the kernel file else: tool_context.state["optimized_kernel_path"] = file_path logging.info( - f"Saved kernel path to optimized_kernel_path: {file_path} (from tool: {tool.name})" + f"Saved kernel path to optimized_kernel_path: {file_path} (from tool: {tool.name})" ) return None def load_single_kernel_to_state(callback_context: CallbackContext): """ - Loads a single kernel file content into state. - Uses kernel_file_path to find the file. - Stores content in 'kernel_code' for use by compilation/profiling agents. - """ + Loads a single kernel file content into state. + Uses kernel_file_path to find the file. + Stores content in 'kernel_code' for use by compilation/profiling agents. + """ file_path = callback_context.state.get("kernel_file_path", None) if file_path: @@ -117,10 +121,10 @@ def load_single_kernel_to_state(callback_context: CallbackContext): def load_profiling_script_to_state(callback_context: CallbackContext): """ - Loads profiling script file content into state. - Uses profiling_script_path to find the file. - Stores content in 'profiling_script' for use by profiling execution agent. - """ + Loads profiling script file content into state. + Uses profiling_script_path to find the file. + Stores content in 'profiling_script' for use by profiling execution agent. + """ file_path = callback_context.state.get("profiling_script_path", None) if file_path: @@ -138,10 +142,10 @@ def load_profiling_script_to_state(callback_context: CallbackContext): def load_two_kernels_to_state(callback_context: CallbackContext): """ - Loads two kernel files (base and optimized) into state for comparison. - Reads from base_kernel_path and optimized_kernel_path. - Stores contents in base_kernel_code and optimized_kernel_code. - """ + Loads two kernel files (base and optimized) into state for comparison. + Reads from base_kernel_path and optimized_kernel_path. + Stores contents in base_kernel_code and optimized_kernel_code. + """ base_path = callback_context.state.get("base_kernel_path", None) optimized_path = callback_context.state.get("optimized_kernel_path", None) @@ -172,11 +176,11 @@ def load_two_kernels_to_state(callback_context: CallbackContext): def load_kernel_and_plan_to_state(callback_context: CallbackContext): """ - Loads kernel file and optimization plan into state for compilation fixing. - Uses optimized_kernel_path and kernel_plan_path to find files. - Stores content in 'kernel_code' and 'kernel_plan' for use by fix agent. - Also formats compilation_history for better readability. - """ + Loads kernel file and optimization plan into state for compilation fixing. + Uses optimized_kernel_path and kernel_plan_path to find files. + Stores content in 'kernel_code' and 'kernel_plan' for use by fix agent. + Also formats compilation_history for better readability. + """ # Load kernel code kernel_path = callback_context.state.get("optimized_kernel_path", None) if kernel_path and os.path.exists(kernel_path): @@ -205,7 +209,8 @@ def load_kernel_and_plan_to_state(callback_context: CallbackContext): callback_context.state["kernel_plan"] = None else: logging.info( - "No optimization plan path found (this is okay for some workflows)") + "No optimization plan path found (this is okay for some workflows)" + ) callback_context.state["kernel_plan"] = None # Format compilation history for readability @@ -230,10 +235,12 @@ def load_kernel_and_plan_to_state(callback_context: CallbackContext): # Store formatted version in a separate key for the prompt callback_context.state["compilation_history_formatted"] = "\n".join( - formatted_history) + formatted_history + ) else: callback_context.state["compilation_history_formatted"] = ( - "No previous attempts (this is the first attempt)") + "No previous attempts (this is the first attempt)" + ) def get_tpu_version_callback(callback_context: CallbackContext): @@ -251,7 +258,8 @@ def get_tpu_version_callback(callback_context: CallbackContext): callback_context.state["tpu_specs"] = tpu_specs[tpu_version] else: callback_context.state["tpu_specs"] = ( - "TPU specs not found for detected version.") + "TPU specs not found for detected version." + ) logging.info(f"Loaded TPU specs for {tpu_version}") except Exception as e: logging.error(f"Failed to load TPU specs: {e}") @@ -264,12 +272,13 @@ def add_workdir_callback(callback_context: CallbackContext): logging.info(f"Set working directory to: {WORKDIR}") -def extract_fix_summary(callback_context: CallbackContext, - llm_response: LlmResponse) -> LlmResponse: +def extract_fix_summary( + callback_context: CallbackContext, llm_response: LlmResponse +) -> LlmResponse: """Extract the agent's response and store it as the fix summary. - This is an after_model_callback that receives the LlmResponse directly. - """ + This is an after_model_callback that receives the LlmResponse directly. + """ if llm_response.content is None or not llm_response.content.parts: logging.warning("No content in LlmResponse to extract fix summary from") return llm_response @@ -298,15 +307,15 @@ def add_pallas_docs(callback_context: CallbackContext): __all__ = [ - "create_path_saver", - "save_kernel_file_paths", - "save_kernel_and_plan_paths", - "load_single_kernel_to_state", - "load_profiling_script_to_state", - "load_two_kernels_to_state", - "load_kernel_and_plan_to_state", - "get_tpu_version_callback", - "add_workdir_callback", - "extract_fix_summary", - "add_pallas_docs", + "create_path_saver", + "save_kernel_file_paths", + "save_kernel_and_plan_paths", + "load_single_kernel_to_state", + "load_profiling_script_to_state", + "load_two_kernels_to_state", + "load_kernel_and_plan_to_state", + "get_tpu_version_callback", + "add_workdir_callback", + "extract_fix_summary", + "add_pallas_docs", ] diff --git a/MaxKernel/hitl_agent/config.py b/MaxKernel/hitl_agent/config.py index 6970e35..82b5f44 100644 --- a/MaxKernel/hitl_agent/config.py +++ b/MaxKernel/hitl_agent/config.py @@ -1,8 +1,10 @@ """Shared configuration for HITL kernel generation agents.""" import os -from google.genai import types + from google.adk.planners import BuiltInPlanner +from google.genai import types + from hitl_agent.constants import TOP_K, TOP_P # Environment variables @@ -13,13 +15,15 @@ # Model configuration model_config = types.GenerateContentConfig( - temperature=0.5, - top_p=TOP_P, - top_k=TOP_K, + temperature=0.5, + top_p=TOP_P, + top_k=TOP_K, ) # Planner configuration with thinking/reasoning traces -thinking_planner = BuiltInPlanner(thinking_config=types.ThinkingConfig( +thinking_planner = BuiltInPlanner( + thinking_config=types.ThinkingConfig( include_thoughts=INCLUDE_THOUGHTS, thinking_level="high", -)) + ) +) diff --git a/MaxKernel/hitl_agent/custom_types.py b/MaxKernel/hitl_agent/custom_types.py index 69a0866..43d621b 100644 --- a/MaxKernel/hitl_agent/custom_types.py +++ b/MaxKernel/hitl_agent/custom_types.py @@ -1,55 +1,60 @@ import logging +from typing import AsyncGenerator + +from google.adk.agents import LlmAgent from google.adk.agents.invocation_context import InvocationContext from google.adk.events import Event, EventActions - -from typing import AsyncGenerator from google.adk.models.google_llm import Gemini from google.genai import types -from google.adk.agents import LlmAgent from hitl_agent.constants import ( - MODEL_NAME,) + MODEL_NAME, +) class CustomLlmAgent(LlmAgent): """Agent that allows early exit from the loop if a condition is met. - Automatically uses gemini_model (with retry support) when a string model name is provided. - """ + Automatically uses gemini_model (with retry support) when a string model name is provided. + """ def __init__(self, *args, **kwargs): """Initialize CustomLlmAgent with automatic Gemini model (with retry) wrapping.""" # If model is a string, use the pre-configured gemini_model with retry support if "model" in kwargs and isinstance(kwargs["model"], str): gemini_model = Gemini( - model=MODEL_NAME, - retry_options=types.HttpRetryOptions( - initial_delay=1, - attempts=5, - ), + model=MODEL_NAME, + retry_options=types.HttpRetryOptions( + initial_delay=1, + attempts=5, + ), ) kwargs["model"] = gemini_model super().__init__(*args, **kwargs) async def _run_async_impl( - self, ctx: InvocationContext) -> AsyncGenerator[Event, None]: + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: # Reset go_to_end flag when new user input is detected - if (hasattr(ctx.session, "contents") and ctx.session.contents and - len(ctx.session.contents) > 0): + if ( + hasattr(ctx.session, "contents") + and ctx.session.contents + and len(ctx.session.contents) > 0 + ): last_message = ctx.session.contents[-1] # Check if the last message is from the user (role='user') if hasattr(last_message, "role") and last_message.role == "user": if ctx.session.state.get("go_to_end", False): logging.info( - f"[{self.name}] New user input detected. Resetting go_to_end flag." + f"[{self.name}] New user input detected. Resetting go_to_end flag." ) ctx.session.state["go_to_end"] = False if ctx.session.state.get("go_to_end", False): logging.info(f"[{self.name}] Early exit condition met. Skipping loop.") yield Event( - author=self.name, - actions=EventActions(escalate=True), + author=self.name, + actions=EventActions(escalate=True), ) else: # Delegate to parent implementation (with native retry support at API level) diff --git a/MaxKernel/hitl_agent/dependency/adk_cli_patch.py b/MaxKernel/hitl_agent/dependency/adk_cli_patch.py index 6eb90fa..77fd873 100644 --- a/MaxKernel/hitl_agent/dependency/adk_cli_patch.py +++ b/MaxKernel/hitl_agent/dependency/adk_cli_patch.py @@ -18,8 +18,9 @@ def apply_patch(): # Find the ADK installation import google.adk + adk_path = Path(google.adk.__file__).parent - cli_file = adk_path / 'cli' / 'cli.py' + cli_file = adk_path / "cli" / "cli.py" if not cli_file.exists(): print(f"Error: Could not find {cli_file}") @@ -114,7 +115,7 @@ async def run_cli(""" content = content.replace(old_code_4, new_code_4) # Backup the original - backup_file = cli_file.with_suffix('.py.backup') + backup_file = cli_file.with_suffix(".py.backup") if not backup_file.exists(): backup_file.write_text(cli_file.read_text()) print(f"Backup created: {backup_file}") @@ -128,9 +129,10 @@ async def run_cli(""" def revert_patch(): """Revert the patch.""" import google.adk + adk_path = Path(google.adk.__file__).parent - cli_file = adk_path / 'cli' / 'cli.py' - backup_file = cli_file.with_suffix('.py.backup') + cli_file = adk_path / "cli" / "cli.py" + backup_file = cli_file.with_suffix(".py.backup") if not backup_file.exists(): print("No backup found!") @@ -141,8 +143,8 @@ def revert_patch(): return True -if __name__ == '__main__': - if len(sys.argv) > 1 and sys.argv[1] == 'revert': +if __name__ == "__main__": + if len(sys.argv) > 1 and sys.argv[1] == "revert": revert_patch() else: apply_patch() diff --git a/MaxKernel/hitl_agent/isolate_object.py b/MaxKernel/hitl_agent/isolate_object.py index f866b85..2efd0b0 100644 --- a/MaxKernel/hitl_agent/isolate_object.py +++ b/MaxKernel/hitl_agent/isolate_object.py @@ -18,9 +18,9 @@ - Standard Python libraries (os, sys, typing, etc.) - Third-party packages installed via pip (jax, flax, numpy, torch, etc.) - Any module NOT found in your workspace - + Example: "import jax.numpy as jnp" → Kept as-is in output - + How it's detected: - Module name is NOT in the local_package_names list - AND module files don't exist in workspace directories @@ -30,15 +30,15 @@ - Packages in your workspace (MaxText, tpu_commons, vllm, etc.) - Files in third_party/ directories - Custom local modules - + Example: "from MaxText.layers import Attention" → Extracts Attention class source code - + How it's detected: - Module name starts with a known local package (see local_package_names list) - OR module files exist in workspace paths: * workspace_root/module_name/ * workspace_root/third_party/module_name/ - + What gets extracted: - The imported object's source code - All functions/classes/constants it depends on @@ -96,17 +96,17 @@ python isolate_object.py third_party/google3/path/to/module.py MyFunction -o output.py """ +import argparse import ast import os import re import sys -import argparse import textwrap import warnings -from typing import Set, List, Dict, Tuple, Optional +from typing import List, Optional, Set, Tuple # Suppress all deprecation warnings from ast module -warnings.filterwarnings('ignore', category=DeprecationWarning) +warnings.filterwarnings("ignore", category=DeprecationWarning) class ImportCollector(ast.NodeVisitor): @@ -122,8 +122,9 @@ def __init__(self): def visit_Import(self, node): for alias in node.names: name = alias.asname if alias.asname else alias.name - self.imports.append(f"import {alias.name}" + - (f" as {alias.asname}" if alias.asname else "")) + self.imports.append( + f"import {alias.name}" + (f" as {alias.asname}" if alias.asname else "") + ) self.defined_names.add(name) # Track alias mapping: alias -> (None, original_module_name) if alias.asname: @@ -138,8 +139,9 @@ def visit_ImportFrom(self, node): names = [] for alias in node.names: name = alias.asname if alias.asname else alias.name - names.append(f"{alias.name}" + - (f" as {alias.asname}" if alias.asname else "")) + names.append( + f"{alias.name}" + (f" as {alias.asname}" if alias.asname else "") + ) self.defined_names.add(name) # Track alias mapping: alias -> (module, original_name) if alias.asname: @@ -174,11 +176,13 @@ def visit_Attribute(self, node): full_path = f"{module}.{original_name}.{node.attr}" self.module_attributes.add(full_path) elif isinstance(node.value, ast.Attribute) and isinstance( - node.value.value, ast.Name): + node.value.value, ast.Name + ): # Handle nested attributes like module.submodule.function self.used_names.add(node.value.value.id) self.module_attributes.add( - f"{node.value.value.id}.{node.value.attr}.{node.attr}") + f"{node.value.value.id}.{node.value.attr}.{node.attr}" + ) self.generic_visit(node) def visit_FunctionDef(self, node): @@ -218,8 +222,9 @@ def __init__(self): def visit_Import(self, node): for alias in node.names: name = alias.asname if alias.asname else alias.name - self.imports.append(f"import {alias.name}" + - (f" as {alias.asname}" if alias.asname else "")) + self.imports.append( + f"import {alias.name}" + (f" as {alias.asname}" if alias.asname else "") + ) self.defined_names.add(name) self.generic_visit(node) @@ -231,8 +236,9 @@ def visit_ImportFrom(self, node): names = [] for alias in node.names: name = alias.asname if alias.asname else alias.name - names.append(f"{alias.name}" + - (f" as {alias.asname}" if alias.asname else "")) + names.append( + f"{alias.name}" + (f" as {alias.asname}" if alias.asname else "") + ) self.defined_names.add(name) self.imports.append(f"from {module} import {', '.join(names)}") self.generic_visit(node) @@ -251,15 +257,17 @@ def __init__(self, filename: str, debug: bool = False): if self.debug: print(f"[DEBUG] Opening file: {abs_path}", file=sys.stderr) - with open(filename, 'r', encoding='utf-8') as f: + with open(filename, "r", encoding="utf-8") as f: self.source = f.read() self.tree = ast.parse(self.source) def find_object(self, object_name: str) -> Optional[ast.AST]: """Find the AST node for the given object name.""" for node in ast.walk(self.tree): - if isinstance( - node, (ast.FunctionDef, ast.ClassDef)) and node.name == object_name: + if ( + isinstance(node, (ast.FunctionDef, ast.ClassDef)) + and node.name == object_name + ): return node # Also search for assignments elif isinstance(node, ast.Assign): @@ -267,8 +275,11 @@ def find_object(self, object_name: str) -> Optional[ast.AST]: if isinstance(target, ast.Name) and target.id == object_name: return node # And annotated assignments - elif isinstance(node, ast.AnnAssign) and isinstance( - node.target, ast.Name) and node.target.id == object_name: + elif ( + isinstance(node, ast.AnnAssign) + and isinstance(node.target, ast.Name) + and node.target.id == object_name + ): return node return None @@ -283,23 +294,36 @@ def get_object_source(self, object_name: str) -> Optional[str]: # For functions and classes with decorators, we need to start from the first decorator start_line = target_node.lineno - 1 # Convert to 0-indexed - if isinstance(target_node, (ast.FunctionDef, ast.ClassDef)) and hasattr( - target_node, 'decorator_list') and target_node.decorator_list: + if ( + isinstance(target_node, (ast.FunctionDef, ast.ClassDef)) + and hasattr(target_node, "decorator_list") + and target_node.decorator_list + ): # Find the line number of the first decorator first_decorator = target_node.decorator_list[0] start_line = first_decorator.lineno - 1 # Convert to 0-indexed # Use end_lineno if available (Python 3.8+) for all node types - if hasattr(target_node, 'end_lineno') and target_node.end_lineno: + if hasattr(target_node, "end_lineno") and target_node.end_lineno: end_line = target_node.end_lineno elif isinstance(target_node, (ast.FunctionDef, ast.ClassDef)): # Fallback for functions and classes: find the end by looking for the next top-level definition end_line = len(lines) for node in ast.walk(self.tree): - if (isinstance(node, (ast.FunctionDef, ast.ClassDef, ast.Import, - ast.ImportFrom, ast.Assign)) and - node.lineno > target_node.lineno and - node.col_offset == 0): # Top-level only + if ( + isinstance( + node, + ( + ast.FunctionDef, + ast.ClassDef, + ast.Import, + ast.ImportFrom, + ast.Assign, + ), + ) + and node.lineno > target_node.lineno + and node.col_offset == 0 + ): # Top-level only end_line = node.lineno - 1 break else: @@ -314,7 +338,7 @@ def get_object_source(self, object_name: str) -> Optional[str]: object_lines.pop() # Join the lines - source_code = '\n'.join(object_lines) + source_code = "\n".join(object_lines) # Dedent the source code to remove extra indentation (e.g., from if TYPE_CHECKING blocks) # Only dedent if all lines have consistent indentation @@ -385,13 +409,24 @@ def get_additional_objects(self, object_name: str) -> List[Tuple[str, str]]: top_level_definitions[node.target.id] = node # Find names that are used but not defined in the object itself - external_names = object_collector.used_names - object_collector.defined_names + external_names = ( + object_collector.used_names - object_collector.defined_names + ) # Filter to only include names that are actually defined at module level # and exclude common built-in names and imported names builtin_names = { - 'self', 'True', 'False', 'None', 'int', 'float', 'str', 'list', 'dict', - 'tuple', 'set' + "self", + "True", + "False", + "None", + "int", + "float", + "str", + "list", + "dict", + "tuple", + "set", } # Get all imported names to exclude them from dependencies @@ -434,15 +469,18 @@ def get_additional_objects(self, object_name: str) -> List[Tuple[str, str]]: if len(node.targets) == 1 and isinstance(node.targets[0], ast.Name): target_name = node.targets[0].id if isinstance(node.value, ast.Attribute) and isinstance( - node.value.value, ast.Name): + node.value.value, ast.Name + ): # This looks like "name = module.something" module_ref = node.value.value.id attr_name = node.value.attr # Check if the module name is likely an imported module - external_module_patterns = ['_partitioning', '_utils', '_layers'] - if any(pattern in module_ref.lower() - for pattern in external_module_patterns): + external_module_patterns = ["_partitioning", "_utils", "_layers"] + if any( + pattern in module_ref.lower() + for pattern in external_module_patterns + ): # Skip this assignment since it should be handled by direct imports is_external_assignment = True @@ -471,15 +509,17 @@ def get_module_usage_patterns(self, object_name: str) -> Set[str]: module_patterns = set() for attr_pattern in collector.module_attributes: # Extract the module name (first part before the dot) - module_name = attr_pattern.split('.')[0] + module_name = attr_pattern.split(".")[0] module_patterns.add(module_name) return module_patterns - def get_local_import_files(self, - imports: List[str], - module_usage_patterns: Set[str] = None, - module_attributes: Set[str] = None) -> List[str]: + def get_local_import_files( + self, + imports: List[str], + module_usage_patterns: Set[str] = None, + module_attributes: Set[str] = None, + ) -> List[str]: """Get list of local Python files that are imported.""" local_files = [] base_dir = os.path.dirname(os.path.abspath(self.filename)) @@ -493,15 +533,15 @@ def get_local_import_files(self, # This handles cases like: from google3.a.b.c import d as e; f = e.f if module_attributes: for attr_pattern in sorted(module_attributes): - parts = attr_pattern.split('.') + parts = attr_pattern.split(".") # For patterns with multiple parts, try to resolve as module paths # We need at least 2 parts (module.attribute), but we'll try any pattern that might resolve if len(parts) >= 2: # e.g., fa_util.exp2, google3.a.b, etc. # Try progressively shorter module paths (e.g., a.b.c.d.f -> try a.b.c.d, a.b.c, etc.) for i in range( - len(parts) - 1, 0, - -1): # Changed from 1 to 0 to try even single-part modules - module_path = '.'.join(parts[:i]) + len(parts) - 1, 0, -1 + ): # Changed from 1 to 0 to try even single-part modules + module_path = ".".join(parts[:i]) # Check if this looks like a local package import is_local = self._is_local_package(module_path, workspace_root) @@ -509,32 +549,45 @@ def get_local_import_files(self, possible_paths = [] # Special handling for google3 - if module_path.startswith('google3.'): + if module_path.startswith("google3."): google3_root = self._get_google3_root() if google3_root: module_path_without_google3 = module_path[ - 8:] # Remove 'google3.' + 8: + ] # Remove 'google3.' possible_paths.append( - os.path.join( - google3_root, 'google3', - module_path_without_google3.replace('.', os.sep) + - '.py')) + os.path.join( + google3_root, + "google3", + module_path_without_google3.replace(".", os.sep) + ".py", + ) + ) possible_paths.append( - os.path.join( - google3_root, 'google3', - module_path_without_google3.replace('.', os.sep), - '__init__.py')) + os.path.join( + google3_root, + "google3", + module_path_without_google3.replace(".", os.sep), + "__init__.py", + ) + ) else: # Try standard paths - module_as_path = module_path.replace('.', os.sep) - possible_paths.extend([ - os.path.join(workspace_root, module_as_path + '.py'), - os.path.join(workspace_root, module_as_path, '__init__.py'), - os.path.join(workspace_root, 'third_party', - module_as_path + '.py'), - os.path.join(workspace_root, 'third_party', module_as_path, - '__init__.py'), - ]) + module_as_path = module_path.replace(".", os.sep) + possible_paths.extend( + [ + os.path.join(workspace_root, module_as_path + ".py"), + os.path.join(workspace_root, module_as_path, "__init__.py"), + os.path.join( + workspace_root, "third_party", module_as_path + ".py" + ), + os.path.join( + workspace_root, + "third_party", + module_as_path, + "__init__.py", + ), + ] + ) for path in possible_paths: if os.path.exists(path) and path not in local_files: @@ -551,13 +604,14 @@ def get_local_import_files(self, for module_name in sorted(module_usage_patterns): # Try direct paths instead of walking the entire tree possible_paths = [ - os.path.join(base_dir, f"{module_name}.py"), - os.path.join(base_dir, f"{module_name.lower()}.py"), - os.path.join(workspace_root, f"{module_name}.py"), - os.path.join(workspace_root, f"{module_name.lower()}.py"), - os.path.join(workspace_root, 'third_party', f"{module_name}.py"), - os.path.join(workspace_root, 'third_party', - f"{module_name.lower()}.py"), + os.path.join(base_dir, f"{module_name}.py"), + os.path.join(base_dir, f"{module_name.lower()}.py"), + os.path.join(workspace_root, f"{module_name}.py"), + os.path.join(workspace_root, f"{module_name.lower()}.py"), + os.path.join(workspace_root, "third_party", f"{module_name}.py"), + os.path.join( + workspace_root, "third_party", f"{module_name.lower()}.py" + ), ] # Only check google3 if we might actually need it (don't call g4 g4d unnecessarily) @@ -570,23 +624,35 @@ def get_local_import_files(self, # If current file is in a package, check that package for pkg in [ - 'google3', 'tpu_commons', 'MaxText', 'maxtext', 'JetStream', - 'jetstream', 'vllm', 'whisper_jax', 'torchprime', 'maxdiffusion', - 'RecML' + "google3", + "tpu_commons", + "MaxText", + "maxtext", + "JetStream", + "jetstream", + "vllm", + "whisper_jax", + "torchprime", + "maxdiffusion", + "RecML", ]: if pkg.lower() in file_path_lower: relevant_packages.append(pkg) # If no relevant packages found, don't search package-specific locations for pkg in relevant_packages: - possible_paths.extend([ + possible_paths.extend( + [ os.path.join(workspace_root, pkg, f"{module_name}.py"), os.path.join(workspace_root, pkg, f"{module_name.lower()}.py"), - os.path.join(workspace_root, 'third_party', pkg, - f"{module_name}.py"), - os.path.join(workspace_root, 'third_party', pkg, - f"{module_name.lower()}.py"), - ]) + os.path.join( + workspace_root, "third_party", pkg, f"{module_name}.py" + ), + os.path.join( + workspace_root, "third_party", pkg, f"{module_name.lower()}.py" + ), + ] + ) for path in possible_paths: if os.path.exists(path) and path not in local_files: @@ -602,92 +668,124 @@ def get_local_import_files(self, module_path = node.module # Handle relative imports (from .module import ...) - if module_path.startswith('.'): + if module_path.startswith("."): module_path = module_path[1:] # Remove leading dot search_dir = base_dir else: search_dir = workspace_root # Check if this is a local package (like tpu_commons, MaxText, etc.) - is_local_package = self._is_local_package(module_path, - workspace_root) + is_local_package = self._is_local_package( + module_path, workspace_root + ) if is_local_package: # Try multiple search strategies for local files possible_paths = [] # Handle google3 imports specially - ONLY call g4 g4d if we see google3 - if module_path.startswith('google3.'): - google3_root = self._get_google3_root( + if module_path.startswith("google3."): + google3_root = ( + self._get_google3_root() ) # Lazy: only called for google3 imports if google3_root: # For google3.foo.bar, look for /path/to/google3/foo/bar.py module_path_without_google3 = module_path[ - 8:] # Remove 'google3.' + 8: + ] # Remove 'google3.' possible_paths.append( - os.path.join( - google3_root, 'google3', - module_path_without_google3.replace('.', os.sep) + - '.py')) + os.path.join( + google3_root, + "google3", + module_path_without_google3.replace(".", os.sep) + + ".py", + ) + ) possible_paths.append( - os.path.join( - google3_root, 'google3', - module_path_without_google3.replace('.', os.sep), - '__init__.py')) + os.path.join( + google3_root, + "google3", + module_path_without_google3.replace(".", os.sep), + "__init__.py", + ) + ) # Direct module path conversion possible_paths.append( - os.path.join(search_dir, - module_path.replace('.', os.sep) + '.py')) + os.path.join( + search_dir, module_path.replace(".", os.sep) + ".py" + ) + ) possible_paths.append( - os.path.join(search_dir, module_path.replace('.', os.sep), - '__init__.py')) + os.path.join( + search_dir, module_path.replace(".", os.sep), "__init__.py" + ) + ) # Search specifically in third_party directories - for third_party_root in ['third_party']: + for third_party_root in ["third_party"]: third_party_path = os.path.join( - workspace_root, third_party_root, - module_path.replace('.', os.sep) + '.py') + workspace_root, + third_party_root, + module_path.replace(".", os.sep) + ".py", + ) possible_paths.append(third_party_path) third_party_pkg_path = os.path.join( - workspace_root, third_party_root, - module_path.replace('.', os.sep), '__init__.py') + workspace_root, + third_party_root, + module_path.replace(".", os.sep), + "__init__.py", + ) possible_paths.append(third_party_pkg_path) # Handle packages with duplicated directory structure (e.g., third_party/pkg/pkg/...) # Extract the first component of the module path - first_component = module_path.split('.')[0] + first_component = module_path.split(".")[0] duplicated_path = os.path.join( - workspace_root, third_party_root, first_component, - module_path.replace('.', os.sep) + '.py') + workspace_root, + third_party_root, + first_component, + module_path.replace(".", os.sep) + ".py", + ) possible_paths.append(duplicated_path) duplicated_pkg_path = os.path.join( - workspace_root, third_party_root, first_component, - module_path.replace('.', os.sep), '__init__.py') + workspace_root, + third_party_root, + first_component, + module_path.replace(".", os.sep), + "__init__.py", + ) possible_paths.append(duplicated_pkg_path) # For packages with different structures, try common patterns - package_parts = module_path.split('.') + package_parts = module_path.split(".") # Special handling for google3 - DON'T search in other packages! - if module_path.startswith('google3.'): + if module_path.startswith("google3."): # For google3, we already handled it above, no need to check other packages known_packages = [] else: # Determine which known packages are relevant for this module # Only check packages that match the module prefix or current file location all_known_packages = [ - 'tpu_commons', 'MaxText', 'maxtext', 'JetStream', - 'jetstream', 'vllm', 'whisper_jax', 'torchprime', - 'maxdiffusion', 'RecML' + "tpu_commons", + "MaxText", + "maxtext", + "JetStream", + "jetstream", + "vllm", + "whisper_jax", + "torchprime", + "maxdiffusion", + "RecML", ] relevant_packages = [] # Check if module starts with a known package name for pkg in all_known_packages: - if module_path.startswith(pkg + '.') or module_path == pkg: + if module_path.startswith(pkg + ".") or module_path == pkg: relevant_packages.append(pkg) # If no match, check if current file is within a known package @@ -698,76 +796,107 @@ def get_local_import_files(self, relevant_packages.append(pkg) # Use relevant packages, or fall back to first component as hint - known_packages = relevant_packages if relevant_packages else [ - package_parts[0] - ] + known_packages = ( + relevant_packages + if relevant_packages + else [package_parts[0]] + ) for i in range(len(package_parts)): - partial_path = os.path.join(*package_parts[:i + 1]) - remaining_path = os.path.join( - *package_parts[i + - 1:]) if i + 1 < len(package_parts) else "" + partial_path = os.path.join(*package_parts[: i + 1]) + remaining_path = ( + os.path.join(*package_parts[i + 1 :]) + if i + 1 < len(package_parts) + else "" + ) # Try in workspace root if remaining_path: - full_file_path = os.path.join(workspace_root, partial_path, - remaining_path + '.py') + full_file_path = os.path.join( + workspace_root, partial_path, remaining_path + ".py" + ) if os.path.exists(full_file_path): possible_paths.append(full_file_path) - full_pkg_path = os.path.join(workspace_root, partial_path, - remaining_path, '__init__.py') + full_pkg_path = os.path.join( + workspace_root, + partial_path, + remaining_path, + "__init__.py", + ) if os.path.exists(full_pkg_path): possible_paths.append(full_pkg_path) else: - target_file_path = os.path.join(workspace_root, - partial_path + '.py') + target_file_path = os.path.join( + workspace_root, partial_path + ".py" + ) if os.path.exists(target_file_path): possible_paths.append(target_file_path) # Try in third_party if remaining_path: - full_file_path = os.path.join(workspace_root, 'third_party', - partial_path, - remaining_path + '.py') + full_file_path = os.path.join( + workspace_root, + "third_party", + partial_path, + remaining_path + ".py", + ) if os.path.exists(full_file_path): possible_paths.append(full_file_path) - full_pkg_path = os.path.join(workspace_root, 'third_party', - partial_path, remaining_path, - '__init__.py') + full_pkg_path = os.path.join( + workspace_root, + "third_party", + partial_path, + remaining_path, + "__init__.py", + ) if os.path.exists(full_pkg_path): possible_paths.append(full_pkg_path) else: - target_file_path = os.path.join(workspace_root, - 'third_party', - partial_path + '.py') + target_file_path = os.path.join( + workspace_root, "third_party", partial_path + ".py" + ) if os.path.exists(target_file_path): possible_paths.append(target_file_path) # Try within known packages for pkg in known_packages: if remaining_path: - full_file_path = os.path.join(workspace_root, pkg, - partial_path, - remaining_path + '.py') + full_file_path = os.path.join( + workspace_root, + pkg, + partial_path, + remaining_path + ".py", + ) if os.path.exists(full_file_path): possible_paths.append(full_file_path) - full_pkg_path = os.path.join(workspace_root, pkg, - partial_path, remaining_path, - '__init__.py') + full_pkg_path = os.path.join( + workspace_root, + pkg, + partial_path, + remaining_path, + "__init__.py", + ) if os.path.exists(full_pkg_path): possible_paths.append(full_pkg_path) # Also try in third_party/pkg - full_file_path = os.path.join(workspace_root, - 'third_party', pkg, - partial_path, - remaining_path + '.py') + full_file_path = os.path.join( + workspace_root, + "third_party", + pkg, + partial_path, + remaining_path + ".py", + ) if os.path.exists(full_file_path): possible_paths.append(full_file_path) - full_pkg_path = os.path.join(workspace_root, - 'third_party', pkg, - partial_path, remaining_path, - '__init__.py') + full_pkg_path = os.path.join( + workspace_root, + "third_party", + pkg, + partial_path, + remaining_path, + "__init__.py", + ) if os.path.exists(full_pkg_path): possible_paths.append(full_pkg_path) @@ -778,7 +907,8 @@ def get_local_import_files(self, else: # Not recognized as local package, check for mock implementations mock_file = self._find_mock_implementation( - module_path, workspace_root) + module_path, workspace_root + ) if mock_file: local_files.append(mock_file) @@ -787,63 +917,93 @@ def get_local_import_files(self, module_name = alias.name # Check if this is a local package - is_local_package = self._is_local_package(module_name, - workspace_root) + is_local_package = self._is_local_package( + module_name, workspace_root + ) if is_local_package: possible_paths = [] # Special handling for google3 - resolve and skip other package checks - if module_name.startswith('google3.'): + if module_name.startswith("google3."): google3_root = self._get_google3_root() if google3_root: module_path_without_google3 = module_name[ - 8:] # Remove 'google3.' + 8: + ] # Remove 'google3.' possible_paths.append( - os.path.join( - google3_root, 'google3', - module_path_without_google3.replace('.', os.sep) + - '.py')) + os.path.join( + google3_root, + "google3", + module_path_without_google3.replace(".", os.sep) + + ".py", + ) + ) possible_paths.append( - os.path.join( - google3_root, 'google3', - module_path_without_google3.replace('.', os.sep), - '__init__.py')) + os.path.join( + google3_root, + "google3", + module_path_without_google3.replace(".", os.sep), + "__init__.py", + ) + ) # Don't search in other packages for google3 imports known_packages = [] else: # Try to find local file for this import possible_paths.append( - os.path.join(base_dir, - module_name.replace('.', os.sep) + '.py')) + os.path.join( + base_dir, module_name.replace(".", os.sep) + ".py" + ) + ) possible_paths.append( - os.path.join(workspace_root, - module_name.replace('.', os.sep) + '.py')) + os.path.join( + workspace_root, module_name.replace(".", os.sep) + ".py" + ) + ) possible_paths.append( - os.path.join(workspace_root, - module_name.replace('.', os.sep), - '__init__.py')) + os.path.join( + workspace_root, + module_name.replace(".", os.sep), + "__init__.py", + ) + ) # Search in third_party possible_paths.append( - os.path.join(workspace_root, 'third_party', - module_name.replace('.', os.sep) + '.py')) + os.path.join( + workspace_root, + "third_party", + module_name.replace(".", os.sep) + ".py", + ) + ) possible_paths.append( - os.path.join(workspace_root, 'third_party', - module_name.replace('.', os.sep), - '__init__.py')) + os.path.join( + workspace_root, + "third_party", + module_name.replace(".", os.sep), + "__init__.py", + ) + ) # Only search in relevant package directories (not all packages) # Determine relevant packages based on module name or current file location all_known_packages = [ - 'tpu_commons', 'MaxText', 'maxtext', 'JetStream', - 'jetstream', 'vllm', 'whisper_jax', 'torchprime', - 'maxdiffusion', 'RecML' + "tpu_commons", + "MaxText", + "maxtext", + "JetStream", + "jetstream", + "vllm", + "whisper_jax", + "torchprime", + "maxdiffusion", + "RecML", ] relevant_packages = [] for pkg in all_known_packages: - if module_name.startswith(pkg + '.') or module_name == pkg: + if module_name.startswith(pkg + ".") or module_name == pkg: relevant_packages.append(pkg) # If no match, check current file location @@ -855,21 +1015,34 @@ def get_local_import_files(self, break # Only add the one we're in # Use relevant packages or just check the module's first component - known_packages = relevant_packages if relevant_packages else [] + known_packages = ( + relevant_packages if relevant_packages else [] + ) - module_as_path = module_name.replace('.', os.sep) + module_as_path = module_name.replace(".", os.sep) for pkg in known_packages: possible_paths.append( - os.path.join(workspace_root, pkg, module_as_path + '.py')) + os.path.join(workspace_root, pkg, module_as_path + ".py") + ) possible_paths.append( - os.path.join(workspace_root, pkg, module_as_path, - '__init__.py')) + os.path.join( + workspace_root, pkg, module_as_path, "__init__.py" + ) + ) possible_paths.append( - os.path.join(workspace_root, 'third_party', pkg, - module_as_path + '.py')) + os.path.join( + workspace_root, "third_party", pkg, module_as_path + ".py" + ) + ) possible_paths.append( - os.path.join(workspace_root, 'third_party', pkg, - module_as_path, '__init__.py')) + os.path.join( + workspace_root, + "third_party", + pkg, + module_as_path, + "__init__.py", + ) + ) # Add any existing paths for path in possible_paths: @@ -894,7 +1067,11 @@ def _find_workspace_root(self, start_dir: str) -> str: # Look for common workspace markers markers = [ - '.git', '.vscode', 'pyproject.toml', 'setup.py', 'requirements.txt' + ".git", + ".vscode", + "pyproject.toml", + "setup.py", + "requirements.txt", ] while current_dir != os.path.dirname(current_dir): # Not at filesystem root @@ -917,17 +1094,17 @@ def _get_google3_root(self) -> Optional[str]: try: import subprocess - result = subprocess.run(['g4', 'g4d'], - capture_output=True, - text=True, - timeout=5) + + result = subprocess.run( + ["g4", "g4d"], capture_output=True, text=True, timeout=5 + ) if result.returncode == 0: g4d_output = result.stdout.strip() # Remove the 'google3' suffix to get the parent directory - if g4d_output.endswith('google3'): - self._google3_root_cache = g4d_output[:-len('google3')].rstrip('/') - elif g4d_output.endswith('/google3'): - self._google3_root_cache = g4d_output[:-len('/google3')] + if g4d_output.endswith("google3"): + self._google3_root_cache = g4d_output[: -len("google3")].rstrip("/") + elif g4d_output.endswith("/google3"): + self._google3_root_cache = g4d_output[: -len("/google3")] else: self._google3_root_cache = None return self._google3_root_cache @@ -940,14 +1117,22 @@ def _is_local_package(self, module_path: str, workspace_root: str) -> bool: """Check if a module path refers to a local package in the workspace.""" # Check if this is a google3 import - if module_path.startswith('google3.'): + if module_path.startswith("google3."): return True # List of known local packages in common ML workspaces local_package_names = [ - 'google3', 'tpu_commons', 'MaxText', 'maxtext', 'JetStream', - 'jetstream', 'vllm', 'whisper_jax', 'torchprime', 'maxdiffusion', - 'RecML' + "google3", + "tpu_commons", + "MaxText", + "maxtext", + "JetStream", + "jetstream", + "vllm", + "whisper_jax", + "torchprime", + "maxdiffusion", + "RecML", ] # Check if the module starts with any known local package @@ -956,46 +1141,49 @@ def _is_local_package(self, module_path: str, workspace_root: str) -> bool: return True # Check if there's actually a directory structure for this module in the workspace - module_parts = module_path.split('.') + module_parts = module_path.split(".") # Try various combinations to see if this could be a local module for i in range(len(module_parts)): - partial_path = os.path.join(*module_parts[:i + 1]) + partial_path = os.path.join(*module_parts[: i + 1]) # Check in common locations potential_locations = [ - os.path.join(workspace_root, partial_path), - os.path.join(workspace_root, 'third_party', partial_path), + os.path.join(workspace_root, partial_path), + os.path.join(workspace_root, "third_party", partial_path), ] # Only check google3 root if the module could plausibly be in google3 # This avoids calling g4 g4d for every single import - if module_path.startswith( - 'google3') or 'google3' in self.filename.lower(): + if ( + module_path.startswith("google3") or "google3" in self.filename.lower() + ): google3_root = self._get_google3_root() if google3_root: potential_locations.append( - os.path.join(google3_root, 'google3', partial_path)) + os.path.join(google3_root, "google3", partial_path) + ) for location in potential_locations: if os.path.exists(location): return True # Also check with .py extension - if os.path.exists(location + '.py'): + if os.path.exists(location + ".py"): return True # If we can't find any evidence this is local, assume it's external return False - def _find_mock_implementation(self, module_path: str, - workspace_root: str) -> Optional[str]: + def _find_mock_implementation( + self, module_path: str, workspace_root: str + ) -> Optional[str]: """Find mock implementations of external modules (e.g., vllm.logger -> tpu_commons.mock.vllm_logger).""" # Common mappings for mock implementations mock_mappings = { - 'vllm.logger': ['tpu_commons.mock.vllm_logger', 'tpu_commons.logger'], - 'vllm.config': ['tpu_commons.mock.vllm_config_utils'], - 'vllm.envs': ['tpu_commons.mock.vllm_envs'], - 'vllm.logging': ['tpu_commons.mock.vllm_logging_utils'], + "vllm.logger": ["tpu_commons.mock.vllm_logger", "tpu_commons.logger"], + "vllm.config": ["tpu_commons.mock.vllm_config_utils"], + "vllm.envs": ["tpu_commons.mock.vllm_envs"], + "vllm.logging": ["tpu_commons.mock.vllm_logging_utils"], } # Check if we have a known mock mapping @@ -1007,38 +1195,48 @@ def _find_mock_implementation(self, module_path: str, return mock_file # Try generic pattern: vllm.X -> tpu_commons.mock.vllm_X - if module_path.startswith('vllm.'): - module_suffix = module_path.split('vllm.', 1)[1] - mock_module = f'tpu_commons.mock.vllm_{module_suffix.replace(".", "_")}' + if module_path.startswith("vllm."): + module_suffix = module_path.split("vllm.", 1)[1] + mock_module = f"tpu_commons.mock.vllm_{module_suffix.replace('.', '_')}" mock_file = self._find_module_file(mock_module, workspace_root) if mock_file: return mock_file return None - def _find_module_file(self, module_path: str, - workspace_root: str) -> Optional[str]: + def _find_module_file( + self, module_path: str, workspace_root: str + ) -> Optional[str]: """Find the file path for a given module path.""" - module_as_path = module_path.replace('.', os.sep) + module_as_path = module_path.replace(".", os.sep) # Try various common locations possible_paths = [ - os.path.join(workspace_root, module_as_path + '.py'), - os.path.join(workspace_root, module_as_path, '__init__.py'), - os.path.join(workspace_root, 'third_party', module_as_path + '.py'), - os.path.join(workspace_root, 'third_party', module_as_path, - '__init__.py'), + os.path.join(workspace_root, module_as_path + ".py"), + os.path.join(workspace_root, module_as_path, "__init__.py"), + os.path.join(workspace_root, "third_party", module_as_path + ".py"), + os.path.join( + workspace_root, "third_party", module_as_path, "__init__.py" + ), ] # For nested packages, also try with duplicated directory structure # e.g., tpu_commons.mock.vllm_logger -> third_party/tpu_commons/tpu_commons/mock/vllm_logger.py - first_component = module_path.split('.')[0] - possible_paths.extend([ - os.path.join(workspace_root, 'third_party', first_component, - module_as_path + '.py'), - os.path.join(workspace_root, 'third_party', first_component, - module_as_path, '__init__.py'), - ]) + first_component = module_path.split(".")[0] + possible_paths.extend( + [ + os.path.join( + workspace_root, "third_party", first_component, module_as_path + ".py" + ), + os.path.join( + workspace_root, + "third_party", + first_component, + module_as_path, + "__init__.py", + ), + ] + ) for path in possible_paths: if os.path.exists(path): @@ -1051,9 +1249,18 @@ def _is_local_package_file(self, file_path: str) -> bool: # List of known local package patterns local_package_patterns = [ - 'google3', 'tpu_commons', 'MaxText', 'maxtext', 'JetStream', - 'jetstream', 'vllm', 'whisper_jax', 'torchprime', 'maxdiffusion', - 'RecML', 'third_party' + "google3", + "tpu_commons", + "MaxText", + "maxtext", + "JetStream", + "jetstream", + "vllm", + "whisper_jax", + "torchprime", + "maxdiffusion", + "RecML", + "third_party", ] # Check if the file path contains any local package patterns @@ -1064,11 +1271,11 @@ def _is_local_package_file(self, file_path: str) -> bool: return False def extract_from_local_file( - self, - file_path: str, - needed_names: Set[str], - module_attributes: Set[str] = None, - external_imports: Set[str] = None + self, + file_path: str, + needed_names: Set[str], + module_attributes: Set[str] = None, + external_imports: Set[str] = None, ) -> Tuple[List[str], List[Tuple[str, str]]]: """Extract needed objects from a local file.""" if module_attributes is None: @@ -1092,7 +1299,8 @@ def collect_definitions_recursive(nodes, depth=0): if isinstance(target, ast.Name): definitions[target.id] = node elif isinstance(node, ast.AnnAssign) and isinstance( - node.target, ast.Name): + node.target, ast.Name + ): definitions[node.target.id] = node # Recursively check inside if blocks (for typing.TYPE_CHECKING, etc.) elif isinstance(node, ast.If): @@ -1112,15 +1320,16 @@ def collect_definitions_recursive(nodes, depth=0): # Extract attributes that are referenced via module.attribute patterns for attr_pattern in module_attributes: - parts = attr_pattern.split('.') + parts = attr_pattern.split(".") if len(parts) >= 2: # Check both exact module name match and partial matches first_part = parts[0] attr_name = parts[1] # Direct module name match or file path contains the module reference - if (first_part == module_name or - first_part in file_path) and attr_name in definitions: + if ( + first_part == module_name or first_part in file_path + ) and attr_name in definitions: node = definitions[attr_name] if isinstance(node, (ast.FunctionDef, ast.ClassDef)): obj_source = extractor.get_object_source(attr_name) @@ -1140,39 +1349,50 @@ def collect_definitions_recursive(nodes, depth=0): for name, node in definitions.items(): if isinstance(node, ast.Assign): # Extract all assignments that look like constants or are needed - if (name in needed_names or isinstance(node.value, ast.Constant) or - (hasattr(ast, 'Str') and isinstance(node.value, ast.Str)) - or # Python < 3.8 compatibility - (hasattr(ast, 'Num') and isinstance(node.value, ast.Num)) - or # Python < 3.8 compatibility - (isinstance(node.value, ast.Attribute) and - isinstance(node.value.value, ast.Name)) or - (name.isupper() or '_' in name)): + if ( + name in needed_names + or isinstance(node.value, ast.Constant) + or ( + hasattr(ast, "Str") and isinstance(node.value, ast.Str) + ) # Python < 3.8 compatibility + or ( + hasattr(ast, "Num") and isinstance(node.value, ast.Num) + ) # Python < 3.8 compatibility + or ( + isinstance(node.value, ast.Attribute) + and isinstance(node.value.value, ast.Name) + ) + or (name.isupper() or "_" in name) + ): # Use get_object_source to properly handle multi-line assignments obj_source = extractor.get_object_source(name) - if obj_source and not any(obj_name == name - for obj_name, _ in extracted_objects): + if obj_source and not any( + obj_name == name for obj_name, _ in extracted_objects + ): extracted_objects.append((name, obj_source)) elif isinstance(node, ast.AnnAssign): # Extract annotated assignments (type aliases) - use get_object_source for multi-line if name in needed_names or name.isupper(): obj_source = extractor.get_object_source(name) - if obj_source and not any(obj_name == name - for obj_name, _ in extracted_objects): + if obj_source and not any( + obj_name == name for obj_name, _ in extracted_objects + ): extracted_objects.append((name, obj_source)) elif isinstance(node, ast.ClassDef): # Extract all class definitions, especially enums, if needed if name in needed_names: obj_source = extractor.get_object_source(name) - if obj_source and not any(obj_name == name - for obj_name, _ in extracted_objects): + if obj_source and not any( + obj_name == name for obj_name, _ in extracted_objects + ): extracted_objects.append((name, obj_source)) elif isinstance(node, ast.FunctionDef): # Extract function definitions if needed if name in needed_names: obj_source = extractor.get_object_source(name) - if obj_source and not any(obj_name == name - for obj_name, _ in extracted_objects): + if obj_source and not any( + obj_name == name for obj_name, _ in extracted_objects + ): extracted_objects.append((name, obj_source)) # For local package files, extract imports that are still needed @@ -1190,22 +1410,26 @@ def collect_definitions_recursive(nodes, depth=0): if isinstance(node, ast.ImportFrom) and node.module: # First check if there's a mock implementation mock_file = extractor._find_mock_implementation( - node.module, - extractor._find_workspace_root(os.path.dirname(file_path))) + node.module, + extractor._find_workspace_root(os.path.dirname(file_path)), + ) if mock_file: has_mock_impl = True - is_external_import = True # Treat as external so it gets checked for mocks + is_external_import = ( + True # Treat as external so it gets checked for mocks + ) elif extractor._is_local_package( - node.module, - extractor._find_workspace_root(os.path.dirname(file_path))): + node.module, + extractor._find_workspace_root(os.path.dirname(file_path)), + ): is_external_import = False break elif isinstance(node, ast.Import): for alias in node.names: if extractor._is_local_package( - alias.name, - extractor._find_workspace_root( - os.path.dirname(file_path))): + alias.name, + extractor._find_workspace_root(os.path.dirname(file_path)), + ): is_external_import = False break except: @@ -1226,8 +1450,9 @@ def collect_definitions_recursive(nodes, depth=0): node = definitions[name] if isinstance(node, (ast.FunctionDef, ast.ClassDef)): obj_source = extractor.get_object_source(name) - if obj_source and not any(obj_name == name - for obj_name, _ in extracted_objects): + if obj_source and not any( + obj_name == name for obj_name, _ in extracted_objects + ): extracted_objects.append((name, obj_source)) # Get imports needed by this object obj_imports = extractor.get_required_imports(name) @@ -1236,22 +1461,28 @@ def collect_definitions_recursive(nodes, depth=0): # Include all assignments and annotated assignments (constants, etc.) # Use get_object_source to properly handle multi-line assignments obj_source = extractor.get_object_source(name) - if obj_source and not any(obj_name == name - for obj_name, _ in extracted_objects): + if obj_source and not any( + obj_name == name for obj_name, _ in extracted_objects + ): extracted_objects.append((name, obj_source)) # Also check for module.attribute patterns, but skip externally imported names for attr_pattern in module_attributes: - if '.' in attr_pattern: - module_part, attr_part = attr_pattern.split('.', 1) - if module_part == module_name and attr_part in definitions and attr_part not in external_imports: + if "." in attr_pattern: + module_part, attr_part = attr_pattern.split(".", 1) + if ( + module_part == module_name + and attr_part in definitions + and attr_part not in external_imports + ): node = definitions[attr_part] if isinstance(node, (ast.FunctionDef, ast.ClassDef)): obj_source = extractor.get_object_source(attr_part) if obj_source: # Check if we haven't already added this object if not any( - name == attr_part for name, _ in extracted_objects): + name == attr_part for name, _ in extracted_objects + ): extracted_objects.append((attr_part, obj_source)) # Get imports needed by this object obj_imports = extractor.get_required_imports(attr_part) @@ -1259,18 +1490,21 @@ def collect_definitions_recursive(nodes, depth=0): elif isinstance(node, (ast.Assign, ast.AnnAssign)): # Include all assignments - use get_object_source to handle multi-line obj_source = extractor.get_object_source(attr_part) - if obj_source and not any(name == attr_part - for name, _ in extracted_objects): + if obj_source and not any( + name == attr_part for name, _ in extracted_objects + ): extracted_objects.append((attr_part, obj_source)) return extracted_imports, extracted_objects except Exception as e: - print(f"Warning: Could not extract from {file_path}: {e}", - file=sys.stderr) + print( + f"Warning: Could not extract from {file_path}: {e}", file=sys.stderr + ) return [], [] - def _import_provides_needed_names(self, import_stmt: str, - needed_names: Set[str]) -> bool: + def _import_provides_needed_names( + self, import_stmt: str, needed_names: Set[str] + ) -> bool: """Check if an import statement provides any of the needed names.""" try: import_tree = ast.parse(import_stmt) @@ -1283,37 +1517,36 @@ def _import_provides_needed_names(self, import_stmt: str, def _is_complete_assignment(self, line: str) -> bool: """Check if a line contains a complete assignment (no unclosed parentheses, brackets, etc.).""" # Count opening and closing parentheses, brackets, braces - paren_count = line.count('(') - line.count(')') - bracket_count = line.count('[') - line.count(']') - brace_count = line.count('{') - line.count('}') + paren_count = line.count("(") - line.count(")") + bracket_count = line.count("[") - line.count("]") + brace_count = line.count("{") - line.count("}") # If any are not balanced, it's an incomplete assignment if paren_count != 0 or bracket_count != 0 or brace_count != 0: return False # Also check for common incomplete patterns - if line.endswith('(') or line.endswith('[') or line.endswith('{'): + if line.endswith("(") or line.endswith("[") or line.endswith("{"): return False return True -def isolate_object(filename: str, - object_name: str, - max_depth: int = 5, - debug: bool = False) -> str: +def isolate_object( + filename: str, object_name: str, max_depth: int = 5, debug: bool = False +) -> str: + """ + Extract an object and all its dependencies from a Python file, including recursive imports. + + Args: + filename: Path to the Python file + object_name: Name of the class or function to extract + max_depth: Maximum recursion depth for following imports + debug: If True, print debug information including file paths being opened + + Returns: + Complete standalone Python code as a string """ - Extract an object and all its dependencies from a Python file, including recursive imports. - - Args: - filename: Path to the Python file - object_name: Name of the class or function to extract - max_depth: Maximum recursion depth for following imports - debug: If True, print debug information including file paths being opened - - Returns: - Complete standalone Python code as a string - """ extractor = ObjectExtractor(filename, debug=debug) # Get the main object source @@ -1323,9 +1556,9 @@ def isolate_object(filename: str, # Track processed files to avoid circular dependencies and duplicates processed_files = set() - processed_file_needed_names = { - } # Track what names we've extracted from each file - processed_objects = set( + processed_file_needed_names = {} # Track what names we've extracted from each file + processed_objects = ( + set() ) # Track (object_name, source_content) to avoid duplicates processed_files.add(os.path.abspath(filename)) @@ -1336,7 +1569,7 @@ def isolate_object(filename: str, def add_unique_object(name: str, source: str, file_origin: str = ""): """Add an object only if it hasn't been added before.""" # Normalize the source content for comparison - source_lines = [line.strip() for line in source.split('\n') if line.strip()] + source_lines = [line.strip() for line in source.split("\n") if line.strip()] source_key = tuple(source_lines) object_key = (name, source_key) @@ -1347,15 +1580,16 @@ def add_unique_object(name: str, source: str, file_origin: str = ""): return True return False - def process_file_recursive(file_extractor: ObjectExtractor, - target_obj: str, - depth: int = 0): + def process_file_recursive( + file_extractor: ObjectExtractor, target_obj: str, depth: int = 0 + ): if depth >= max_depth: return # Get workspace root for this file workspace_root = file_extractor._find_workspace_root( - os.path.dirname(file_extractor.filename)) + os.path.dirname(file_extractor.filename) + ) # Get required imports for this object imports = file_extractor.get_required_imports(target_obj) @@ -1386,7 +1620,8 @@ def process_file_recursive(file_extractor: ObjectExtractor, # Separate local and external imports local_files = file_extractor.get_local_import_files( - imports, module_usage_patterns, object_collector.module_attributes) + imports, module_usage_patterns, object_collector.module_attributes + ) external_imports = [] external_imported_names = set() @@ -1404,17 +1639,21 @@ def process_file_recursive(file_extractor: ObjectExtractor, if isinstance(node, ast.ImportFrom) and node.module: # More comprehensive matching for local modules if self._is_local_package( - node.module, - file_extractor._find_workspace_root( - os.path.dirname(file_extractor.filename))): + node.module, + file_extractor._find_workspace_root( + os.path.dirname(file_extractor.filename) + ), + ): is_local = True break elif isinstance(node, ast.Import): for alias in node.names: if self._is_local_package( - alias.name, - file_extractor._find_workspace_root( - os.path.dirname(file_extractor.filename))): + alias.name, + file_extractor._find_workspace_root( + os.path.dirname(file_extractor.filename) + ), + ): is_local = True break except: @@ -1429,84 +1668,114 @@ def process_file_recursive(file_extractor: ObjectExtractor, for node in ast.walk(import_tree): if isinstance(node, ast.ImportFrom) and node.module: mock_file = file_extractor._find_mock_implementation( - node.module, workspace_root) + node.module, workspace_root + ) if mock_file: has_mock = True # Extract from mock file for imported_name in node.names: - if imported_name.name != '*': + if imported_name.name != "*": try: mock_extractor = ObjectExtractor( - mock_file, debug=file_extractor.debug) + mock_file, debug=file_extractor.debug + ) original_name = imported_name.name - alias_name = imported_name.asname if imported_name.asname else original_name + alias_name = ( + imported_name.asname + if imported_name.asname + else original_name + ) obj_source = mock_extractor.get_object_source( - original_name) + original_name + ) if obj_source: # Rename if aliased if alias_name != original_name: obj_source = re.sub( - r'\bdef\s+' + re.escape(original_name) + r'\b', - f'def {alias_name}', obj_source) + r"\bdef\s+" + re.escape(original_name) + r"\b", + f"def {alias_name}", + obj_source, + ) obj_source = re.sub( - r'\bclass\s+' + re.escape(original_name) + r'\b', - f'class {alias_name}', obj_source) + r"\bclass\s+" + re.escape(original_name) + r"\b", + f"class {alias_name}", + obj_source, + ) add_unique_object( - alias_name, obj_source, - f"{os.path.basename(mock_file)} (mock)") + alias_name, + obj_source, + f"{os.path.basename(mock_file)} (mock)", + ) # Get imports needed by this mock object try: mock_imports = mock_extractor.get_required_imports( - original_name) + original_name + ) for mock_import in mock_imports: # Check if import is external or has its own mock is_mock_import_external = True try: import_tree = ast.parse(mock_import) for node in ast.walk(import_tree): - if isinstance(node, - ast.ImportFrom) and node.module: + if ( + isinstance(node, ast.ImportFrom) + and node.module + ): if mock_extractor._is_local_package( - node.module, workspace_root): + node.module, workspace_root + ): is_mock_import_external = False break except: pass - if is_mock_import_external and mock_import not in all_imports: + if ( + is_mock_import_external + and mock_import not in all_imports + ): all_imports.append(mock_import) except: pass # Get dependencies try: - mock_additional = mock_extractor.get_additional_objects( - original_name) + mock_additional = ( + mock_extractor.get_additional_objects(original_name) + ) for add_name, add_source in mock_additional: add_unique_object( - add_name, add_source, - f"{os.path.basename(mock_file)} (mock)") + add_name, + add_source, + f"{os.path.basename(mock_file)} (mock)", + ) # Also get imports for each additional dependency try: add_imports = mock_extractor.get_required_imports( - add_name) + add_name + ) for add_import in add_imports: is_add_import_external = True try: import_tree = ast.parse(add_import) for node in ast.walk(import_tree): - if isinstance( - node, ast.ImportFrom) and node.module: + if ( + isinstance(node, ast.ImportFrom) + and node.module + ): if mock_extractor._is_local_package( - node.module, workspace_root): + node.module, workspace_root + ): is_add_import_external = False break except: pass - if is_add_import_external and add_import not in all_imports: + if ( + is_add_import_external + and add_import not in all_imports + ): all_imports.append(add_import) except: pass @@ -1514,8 +1783,9 @@ def process_file_recursive(file_extractor: ObjectExtractor, pass except Exception as e: print( - f"Warning: Could not extract {imported_name.name} from mock {mock_file}: {e}", - file=sys.stderr) + f"Warning: Could not extract {imported_name.name} from mock {mock_file}: {e}", + file=sys.stderr, + ) break except: pass @@ -1565,12 +1835,14 @@ def process_file_recursive(file_extractor: ObjectExtractor, obj_collector.visit(obj_node) obj_module_patterns = file_extractor.get_module_usage_patterns( - obj_name) + obj_name + ) obj_module_attributes = obj_collector.module_attributes # Get local files for this object's imports obj_local_files = file_extractor.get_local_import_files( - obj_imports, obj_module_patterns, obj_module_attributes) + obj_imports, obj_module_patterns, obj_module_attributes + ) for obj_import in obj_imports: # Check if it's a local import @@ -1580,17 +1852,21 @@ def process_file_recursive(file_extractor: ObjectExtractor, for node in ast.walk(import_tree): if isinstance(node, ast.ImportFrom) and node.module: if file_extractor._is_local_package( - node.module, - file_extractor._find_workspace_root( - os.path.dirname(file_extractor.filename))): + node.module, + file_extractor._find_workspace_root( + os.path.dirname(file_extractor.filename) + ), + ): is_local_import = True break elif isinstance(node, ast.Import): for alias in node.names: if file_extractor._is_local_package( - alias.name, - file_extractor._find_workspace_root( - os.path.dirname(file_extractor.filename))): + alias.name, + file_extractor._find_workspace_root( + os.path.dirname(file_extractor.filename) + ), + ): is_local_import = True break except: @@ -1642,13 +1918,17 @@ def process_file_recursive(file_extractor: ObjectExtractor, for node in ast.walk(import_tree): if isinstance(node, ast.ImportFrom): if node.module and file_extractor._is_local_package( - node.module, workspace_root): - module_as_path = node.module.replace('.', os.sep) + node.module, workspace_root + ): + module_as_path = node.module.replace(".", os.sep) local_file_normalized = os.path.normpath(local_file) - if (local_file_normalized.endswith(module_as_path + '.py') or - local_file_normalized.endswith( - os.path.join(module_as_path, '__init__.py')) or - module_as_path in local_file_normalized): + if ( + local_file_normalized.endswith(module_as_path + ".py") + or local_file_normalized.endswith( + os.path.join(module_as_path, "__init__.py") + ) + or module_as_path in local_file_normalized + ): for alias in node.names: name = alias.asname if alias.asname else alias.name if name != "*": @@ -1658,7 +1938,8 @@ def process_file_recursive(file_extractor: ObjectExtractor, # Check if there are new names previously_needed = processed_file_needed_names.get( - abs_local_file, set()) + abs_local_file, set() + ) new_names = temp_needed_names - previously_needed if new_names: should_process = True @@ -1675,45 +1956,59 @@ def process_file_recursive(file_extractor: ObjectExtractor, if isinstance(node, ast.ImportFrom): # Check if this import statement refers to the current local file if node.module and file_extractor._is_local_package( - node.module, workspace_root): + node.module, workspace_root + ): # Convert module path to file path and check if it matches current local_file - module_as_path = node.module.replace('.', os.sep) + module_as_path = node.module.replace(".", os.sep) local_file_normalized = os.path.normpath(local_file) # Check if this import matches the local file is_match = False # Standard checks - if (local_file_normalized.endswith(module_as_path + '.py') or - local_file_normalized.endswith( - os.path.join(module_as_path, '__init__.py'))): + if local_file_normalized.endswith( + module_as_path + ".py" + ) or local_file_normalized.endswith( + os.path.join(module_as_path, "__init__.py") + ): is_match = True # For google3 imports, also check without the google3 prefix - if not is_match and node.module.startswith('google3.'): + if not is_match and node.module.startswith("google3."): # Remove google3. prefix and check again module_without_google3 = node.module[ - 8:] # Remove 'google3.' + 8: + ] # Remove 'google3.' module_as_path_no_g3 = module_without_google3.replace( - '.', os.sep) - if (local_file_normalized.endswith(module_as_path_no_g3 + - '.py') or - local_file_normalized.endswith( - os.path.join(module_as_path_no_g3, '__init__.py'))): + ".", os.sep + ) + if local_file_normalized.endswith( + module_as_path_no_g3 + ".py" + ) or local_file_normalized.endswith( + os.path.join(module_as_path_no_g3, "__init__.py") + ): is_match = True # Also check if the file contains google3/ in the path - if not is_match and 'google3' + os.sep in local_file_normalized: + if ( + not is_match + and "google3" + os.sep in local_file_normalized + ): # Extract the part after google3/ - google3_idx = local_file_normalized.find('google3' + - os.sep) + google3_idx = local_file_normalized.find( + "google3" + os.sep + ) if google3_idx != -1: path_after_google3 = local_file_normalized[ - google3_idx + 8:] # Skip 'google3/' - if (path_after_google3 == module_as_path_no_g3 + '.py' - or path_after_google3 == os.path.join( - module_as_path_no_g3, '__init__.py') or - path_after_google3.startswith(module_as_path_no_g3 + - os.sep)): + google3_idx + 8 : + ] # Skip 'google3/' + if ( + path_after_google3 == module_as_path_no_g3 + ".py" + or path_after_google3 + == os.path.join(module_as_path_no_g3, "__init__.py") + or path_after_google3.startswith( + module_as_path_no_g3 + os.sep + ) + ): is_match = True # Fallback: check if module_as_path is contained in the file path @@ -1729,44 +2024,57 @@ def process_file_recursive(file_extractor: ObjectExtractor, needed_names.add(name) elif isinstance(node, ast.Import): for alias in node.names: - if file_extractor._is_local_package(alias.name, - workspace_root): + if file_extractor._is_local_package( + alias.name, workspace_root + ): # For direct imports, check if the module name relates to this file - module_as_path = alias.name.replace('.', os.sep) + module_as_path = alias.name.replace(".", os.sep) local_file_normalized = os.path.normpath(local_file) is_match = False # Standard checks - if (local_file_normalized.endswith(module_as_path + '.py') - or local_file_normalized.endswith( - os.path.join(module_as_path, '__init__.py'))): + if local_file_normalized.endswith( + module_as_path + ".py" + ) or local_file_normalized.endswith( + os.path.join(module_as_path, "__init__.py") + ): is_match = True # For google3 imports, also check without the google3 prefix - if not is_match and alias.name.startswith('google3.'): + if not is_match and alias.name.startswith("google3."): module_without_google3 = alias.name[ - 8:] # Remove 'google3.' + 8: + ] # Remove 'google3.' module_as_path_no_g3 = module_without_google3.replace( - '.', os.sep) - if (local_file_normalized.endswith(module_as_path_no_g3 + - '.py') or - local_file_normalized.endswith( - os.path.join(module_as_path_no_g3, - '__init__.py'))): + ".", os.sep + ) + if local_file_normalized.endswith( + module_as_path_no_g3 + ".py" + ) or local_file_normalized.endswith( + os.path.join(module_as_path_no_g3, "__init__.py") + ): is_match = True # Also check if the file contains google3/ in the path - if not is_match and 'google3' + os.sep in local_file_normalized: - google3_idx = local_file_normalized.find('google3' + - os.sep) + if ( + not is_match + and "google3" + os.sep in local_file_normalized + ): + google3_idx = local_file_normalized.find( + "google3" + os.sep + ) if google3_idx != -1: path_after_google3 = local_file_normalized[ - google3_idx + 8:] # Skip 'google3/' - if (path_after_google3 == module_as_path_no_g3 + '.py' - or path_after_google3 == os.path.join( - module_as_path_no_g3, '__init__.py') or - path_after_google3.startswith( - module_as_path_no_g3 + os.sep)): + google3_idx + 8 : + ] # Skip 'google3/' + if ( + path_after_google3 == module_as_path_no_g3 + ".py" + or path_after_google3 + == os.path.join(module_as_path_no_g3, "__init__.py") + or path_after_google3.startswith( + module_as_path_no_g3 + os.sep + ) + ): is_match = True # Fallback: check if module_as_path is contained in the file path @@ -1787,33 +2095,44 @@ def process_file_recursive(file_extractor: ObjectExtractor, # Also add module attributes that might be needed # Sort to ensure deterministic order for attr_pattern in sorted(object_collector.module_attributes): - if '.' in attr_pattern: - module_part, attr_part = attr_pattern.split('.', 1) + if "." in attr_pattern: + module_part, attr_part = attr_pattern.split(".", 1) needed_names.add(attr_part) # Extract needed objects from the local file try: - local_extractor = ObjectExtractor(local_file, - debug=file_extractor.debug) - local_imports, local_objects = local_extractor.extract_from_local_file( - local_file, needed_names, object_collector.module_attributes, - external_imported_names) + local_extractor = ObjectExtractor( + local_file, debug=file_extractor.debug + ) + local_imports, local_objects = ( + local_extractor.extract_from_local_file( + local_file, + needed_names, + object_collector.module_attributes, + external_imported_names, + ) + ) # Track what names we extracted from this file extracted_names_from_file = set() for obj_name, _ in local_objects: # Extract just the base name (without file annotation) - base_name = obj_name.split( - ' (from ')[0] if ' (from ' in obj_name else obj_name + base_name = ( + obj_name.split(" (from ")[0] + if " (from " in obj_name + else obj_name + ) extracted_names_from_file.add(base_name) # Update the tracking dictionary if abs_local_file in processed_file_needed_names: processed_file_needed_names[abs_local_file].update( - extracted_names_from_file) + extracted_names_from_file + ) else: - processed_file_needed_names[ - abs_local_file] = extracted_names_from_file + processed_file_needed_names[abs_local_file] = ( + extracted_names_from_file + ) # Add imports from local file (filter external vs local) for local_import in local_imports: @@ -1827,24 +2146,32 @@ def process_file_recursive(file_extractor: ObjectExtractor, if isinstance(node, ast.ImportFrom) and node.module: # Check for mock implementations FIRST mock_file = local_extractor._find_mock_implementation( - node.module, workspace_root) + node.module, workspace_root + ) if mock_file: has_mock = True # Extract the needed names from this mock file for imported_name in node.names: - if imported_name.name != '*': + if imported_name.name != "*": original_name = imported_name.name - alias_name = imported_name.asname if imported_name.asname else original_name + alias_name = ( + imported_name.asname + if imported_name.asname + else original_name + ) mock_extractions.append( - (mock_file, original_name, alias_name)) + (mock_file, original_name, alias_name) + ) elif local_extractor._is_local_package( - node.module, workspace_root): + node.module, workspace_root + ): is_external = False break elif isinstance(node, ast.Import): for alias in node.names: if local_extractor._is_local_package( - alias.name, workspace_root): + alias.name, workspace_root + ): is_external = False break except Exception as e: @@ -1854,26 +2181,35 @@ def process_file_recursive(file_extractor: ObjectExtractor, if mock_extractions: for mock_file, original_name, alias_name in mock_extractions: try: - mock_extractor = ObjectExtractor(mock_file, - debug=file_extractor.debug) + mock_extractor = ObjectExtractor( + mock_file, debug=file_extractor.debug + ) obj_source = mock_extractor.get_object_source(original_name) if obj_source: # If there's an alias, we need to rename the function/class in the source if alias_name != original_name: # Simple replacement for function/class names obj_source = re.sub( - r'\bdef\s+' + re.escape(original_name) + r'\b', - f'def {alias_name}', obj_source) + r"\bdef\s+" + re.escape(original_name) + r"\b", + f"def {alias_name}", + obj_source, + ) obj_source = re.sub( - r'\bclass\s+' + re.escape(original_name) + r'\b', - f'class {alias_name}', obj_source) - add_unique_object(alias_name, obj_source, - f"{os.path.basename(mock_file)} (mock)") + r"\bclass\s+" + re.escape(original_name) + r"\b", + f"class {alias_name}", + obj_source, + ) + add_unique_object( + alias_name, + obj_source, + f"{os.path.basename(mock_file)} (mock)", + ) # Get imports needed by this mock object try: mock_imports = mock_extractor.get_required_imports( - original_name) + original_name + ) for mock_import in mock_imports: # Check if import is external or has its own mock is_mock_import_external = True @@ -1882,13 +2218,17 @@ def process_file_recursive(file_extractor: ObjectExtractor, for node in ast.walk(import_tree): if isinstance(node, ast.ImportFrom) and node.module: if mock_extractor._is_local_package( - node.module, workspace_root): + node.module, workspace_root + ): is_mock_import_external = False break except: pass - if is_mock_import_external and mock_import not in all_imports: + if ( + is_mock_import_external + and mock_import not in all_imports + ): all_imports.append(mock_import) except: pass @@ -1896,31 +2236,41 @@ def process_file_recursive(file_extractor: ObjectExtractor, # Also get dependencies of this mock object try: mock_additional = mock_extractor.get_additional_objects( - original_name) + original_name + ) for add_name, add_source in mock_additional: add_unique_object( - add_name, add_source, - f"{os.path.basename(mock_file)} (mock)") + add_name, + add_source, + f"{os.path.basename(mock_file)} (mock)", + ) # Also get imports for each additional dependency try: add_imports = mock_extractor.get_required_imports( - add_name) + add_name + ) for add_import in add_imports: is_add_import_external = True try: import_tree = ast.parse(add_import) for node in ast.walk(import_tree): - if isinstance(node, - ast.ImportFrom) and node.module: + if ( + isinstance(node, ast.ImportFrom) + and node.module + ): if mock_extractor._is_local_package( - node.module, workspace_root): + node.module, workspace_root + ): is_add_import_external = False break except: pass - if is_add_import_external and add_import not in all_imports: + if ( + is_add_import_external + and add_import not in all_imports + ): all_imports.append(add_import) except: pass @@ -1928,23 +2278,26 @@ def process_file_recursive(file_extractor: ObjectExtractor, pass except Exception as e: print( - f"Warning: Could not extract {original_name} from mock {mock_file}: {e}", - file=sys.stderr) + f"Warning: Could not extract {original_name} from mock {mock_file}: {e}", + file=sys.stderr, + ) if is_external and not has_mock: all_imports.append(local_import) # Add objects from local file for obj_name, obj_source in local_objects: - add_unique_object(obj_name, obj_source, - os.path.basename(local_file)) + add_unique_object( + obj_name, obj_source, os.path.basename(local_file) + ) # Recursively process dependencies of this extracted object within the same file try: obj_additional = local_extractor.get_additional_objects(obj_name) for add_name, add_source in obj_additional: - add_unique_object(add_name, add_source, - os.path.basename(local_file)) + add_unique_object( + add_name, add_source, os.path.basename(local_file) + ) except: pass # Continue if we can't analyze dependencies @@ -1955,8 +2308,9 @@ def process_file_recursive(file_extractor: ObjectExtractor, pass # Continue if we can't process recursively except Exception as e: - print(f"Warning: Could not process {local_file}: {e}", - file=sys.stderr) + print( + f"Warning: Could not process {local_file}: {e}", file=sys.stderr + ) # Start the recursive processing process_file_recursive(extractor, object_name) @@ -1972,8 +2326,7 @@ def process_file_recursive(file_extractor: ObjectExtractor, # Filter out imports that we're already providing as extracted objects extracted_object_names = set() - extracted_from_modules = { - } # Track which module each extracted object came from + extracted_from_modules = {} # Track which module each extracted object came from # Include the main object being isolated extracted_object_names.add(object_name) @@ -1982,12 +2335,12 @@ def process_file_recursive(file_extractor: ObjectExtractor, # Include all additional extracted objects for name, _ in all_objects: # Extract the base name (without the "(from file.py)" part) - base_name = name.split(' (from ')[0] + base_name = name.split(" (from ")[0] extracted_object_names.add(base_name) # Track which file this object came from - if ' (from ' in name: - file_part = name.split(' (from ')[1].rstrip(')') + if " (from " in name: + file_part = name.split(" (from ")[1].rstrip(")") extracted_from_modules[base_name] = file_part filtered_imports = [] @@ -2000,7 +2353,9 @@ def process_file_recursive(file_extractor: ObjectExtractor, import_collector.visit(import_tree) # Check both alias names and original names against extracted objects - imported_alias_names = import_collector.defined_names # The names available after import (aliases) + imported_alias_names = ( + import_collector.defined_names + ) # The names available after import (aliases) # Also collect the original imported names (before aliasing) original_imported_names = set() @@ -2031,14 +2386,19 @@ def process_file_recursive(file_extractor: ObjectExtractor, original_name = alias.name # Keep this import if neither the alias nor original name conflicts - if (alias_name not in extracted_object_names and - original_name not in extracted_object_names): - remaining_aliases.append(f"{alias.name}" + ( - f" as {alias.asname}" if alias.asname else "")) + if ( + alias_name not in extracted_object_names + and original_name not in extracted_object_names + ): + remaining_aliases.append( + f"{alias.name}" + + (f" as {alias.asname}" if alias.asname else "") + ) if remaining_aliases: filtered_imports.append( - f"from {node.module} import {', '.join(remaining_aliases)}") + f"from {node.module} import {', '.join(remaining_aliases)}" + ) # If no remaining names, the entire import is filtered out else: # Keep imports without a module (shouldn't happen, but be safe) @@ -2050,8 +2410,10 @@ def process_file_recursive(file_extractor: ObjectExtractor, for alias in node.names: name = alias.asname if alias.asname else alias.name if name not in extracted_object_names: - remaining_aliases.append(f"import {alias.name}" + ( - f" as {alias.asname}" if alias.asname else "")) + remaining_aliases.append( + f"import {alias.name}" + + (f" as {alias.asname}" if alias.asname else "") + ) if remaining_aliases: filtered_imports.extend(remaining_aliases) @@ -2063,20 +2425,22 @@ def process_file_recursive(file_extractor: ObjectExtractor, output_lines = [] # Fix module attribute references in all object sources - fixed_object_source = fix_module_attribute_references(object_source, - extracted_object_names) + fixed_object_source = fix_module_attribute_references( + object_source, extracted_object_names + ) # Also fix references in all extracted objects fixed_all_objects = [] for name, source in all_objects: - fixed_source = fix_module_attribute_references(source, - extracted_object_names) + fixed_source = fix_module_attribute_references( + source, extracted_object_names + ) fixed_all_objects.append((name, fixed_source)) # After fixing references, check if any module imports are still needed - final_imports = filter_unused_module_imports(filtered_imports, - fixed_object_source, - fixed_all_objects) + final_imports = filter_unused_module_imports( + filtered_imports, fixed_object_source, fixed_all_objects + ) # Sort objects by dependencies (topological sort) sorted_objects = topologically_sort_objects(fixed_all_objects) @@ -2086,7 +2450,8 @@ def process_file_recursive(file_extractor: ObjectExtractor, # Add header comment output_lines.append( - f"# Isolated {object_name} from {filename} (with recursive imports)") + f"# Isolated {object_name} from {filename} (with recursive imports)" + ) output_lines.append("") # Add final imports, filtering out any remaining google3 imports and consolidating duplicates @@ -2095,8 +2460,11 @@ def process_file_recursive(file_extractor: ObjectExtractor, non_google3_imports = [] for import_stmt in final_imports: # Skip any imports from google3 - if 'from google3.' in import_stmt or 'import google3.' in import_stmt or import_stmt.strip( - ).startswith('import google3'): + if ( + "from google3." in import_stmt + or "import google3." in import_stmt + or import_stmt.strip().startswith("import google3") + ): continue non_google3_imports.append(import_stmt) @@ -2109,7 +2477,7 @@ def process_file_recursive(file_extractor: ObjectExtractor, import_tree = ast.parse(import_stmt) for node in ast.walk(import_tree): if isinstance(node, ast.ImportFrom): - module = node.module or '' + module = node.module or "" if module not in from_imports: from_imports[module] = set() @@ -2160,18 +2528,19 @@ def process_file_recursive(file_extractor: ObjectExtractor, output_lines.append(f"# Main object: {object_name}") output_lines.append(fixed_object_source) - return '\n'.join(output_lines) + return "\n".join(output_lines) def topologically_sort_objects( - objects: List[Tuple[str, str]]) -> List[Tuple[str, str]]: + objects: List[Tuple[str, str]], +) -> List[Tuple[str, str]]: """Sort objects so that dependencies come before the objects that use them.""" # Build a dependency graph # Extract just the base names for analysis object_names = {} for name, source in objects: - base_name = name.split(' (from ')[0] + base_name = name.split(" (from ")[0] object_names[base_name] = (name, source) # Find dependencies for each object @@ -2182,7 +2551,6 @@ def topologically_sort_objects( tree = ast.parse(source) class DependencyFinder(ast.NodeVisitor): - def __init__(self): self.local_names = set() # Track locally defined names @@ -2203,8 +2571,11 @@ def visit_Name(self, node): if isinstance(node.ctx, ast.Load): # This name is being used, check if it's one of our extracted objects # and not a local variable/parameter - if (node.id in object_names and node.id != base_name and - node.id not in self.local_names): + if ( + node.id in object_names + and node.id != base_name + and node.id not in self.local_names + ): deps.add(node.id) elif isinstance(node.ctx, ast.Store): # This name is being defined locally @@ -2213,9 +2584,11 @@ def visit_Name(self, node): def visit_Attribute(self, node): # For type annotations like model_config: ModelConfig - if (isinstance(node.value, ast.Name) and - node.value.id in object_names and - node.value.id not in self.local_names): + if ( + isinstance(node.value, ast.Name) + and node.value.id in object_names + and node.value.id not in self.local_names + ): deps.add(node.value.id) self.generic_visit(node) @@ -2266,8 +2639,8 @@ def visit_Attribute(self, node): def filter_unused_module_imports( - imports: List[str], main_source: str, - all_objects: List[Tuple[str, str]]) -> List[str]: + imports: List[str], main_source: str, all_objects: List[Tuple[str, str]] +) -> List[str]: """Filter out module imports that are no longer used after fixing attribute references.""" used_names = set() @@ -2281,7 +2654,6 @@ def filter_unused_module_imports( tree = ast.parse(all_source_code) class NameUsageFinder(ast.NodeVisitor): - def visit_Attribute(self, node): # Track module.attribute usage (e.g., os.path, jnp.array) if isinstance(node.value, ast.Name): @@ -2339,11 +2711,12 @@ def visit_Name(self, node): if used_aliases: if node.module: final_imports.append( - f"from {node.module} import {', '.join(sorted(used_aliases))}" + f"from {node.module} import {', '.join(sorted(used_aliases))}" ) else: final_imports.append( - f"from . import {', '.join(sorted(used_aliases))}") + f"from . import {', '.join(sorted(used_aliases))}" + ) elif isinstance(node, ast.Import): # For regular imports, keep if used for alias in node.names: @@ -2361,18 +2734,18 @@ def visit_Name(self, node): return final_imports -def fix_module_attribute_references(source_code: str, - extracted_names: Set[str]) -> str: +def fix_module_attribute_references( + source_code: str, extracted_names: Set[str] +) -> str: + """ + Fix module.attribute references when the attribute has been extracted to the same scope. + For example, change test_helper.helper_function() to helper_function() if helper_function is extracted. + But avoid replacing module.attribute in assignment statements where we're assigning from module.attribute to the same attribute name. """ - Fix module.attribute references when the attribute has been extracted to the same scope. - For example, change test_helper.helper_function() to helper_function() if helper_function is extracted. - But avoid replacing module.attribute in assignment statements where we're assigning from module.attribute to the same attribute name. - """ try: tree = ast.parse(source_code) class AttributeReferenceFixer(ast.NodeTransformer): - def visit_Assign(self, node): # For assignment nodes, we need to be careful not to replace the right-hand side # if it's assigning from module.attribute to the same attribute name @@ -2383,10 +2756,12 @@ def visit_Assign(self, node): target_name = node.targets[0].id # Check if the value is a module.attribute where attribute matches the target - if (isinstance(node.value, ast.Attribute) and - isinstance(node.value.value, ast.Name) and - node.value.attr == target_name and - target_name in extracted_names): + if ( + isinstance(node.value, ast.Attribute) + and isinstance(node.value.value, ast.Name) + and node.value.attr == target_name + and target_name in extracted_names + ): # Don't transform this assignment - keep the original module.attribute reference return node @@ -2406,21 +2781,27 @@ def visit_Attribute(self, node): # Convert back to source code try: import astor + return astor.to_source(fixed_tree).strip() except ImportError: # astor not available, fall back to simple string replacement - print("Warning: astor not available, using simple string replacement", - file=sys.stderr) + print( + "Warning: astor not available, using simple string replacement", + file=sys.stderr, + ) return simple_fix_module_references(source_code, extracted_names) except Exception as e: - print(f"Warning: Could not parse source for fixing references: {e}", - file=sys.stderr) + print( + f"Warning: Could not parse source for fixing references: {e}", + file=sys.stderr, + ) return simple_fix_module_references(source_code, extracted_names) -def simple_fix_module_references(source_code: str, - extracted_names: Set[str]) -> str: +def simple_fix_module_references( + source_code: str, extracted_names: Set[str] +) -> str: """Simple string-based replacement for module.attribute references.""" import re @@ -2430,12 +2811,13 @@ def simple_fix_module_references(source_code: str, # More sophisticated pattern to avoid replacing assignment statements # where we're assigning from module.name to name # Pattern: Don't replace if it's "name = something.name" - assignment_pattern = r'^(\s*' + re.escape( - name) + r'\s*=\s*)\w+\.' + re.escape(name) + r'\b' + assignment_pattern = ( + r"^(\s*" + re.escape(name) + r"\s*=\s*)\w+\." + re.escape(name) + r"\b" + ) # Replace patterns like "module.name" with just "name", but not in assignment statements # Split by lines to handle assignments carefully - lines = fixed_code.split('\n') + lines = fixed_code.split("\n") for i, line in enumerate(lines): # Check if this line is an assignment from module.name to name if re.match(assignment_pattern, line.strip()): @@ -2443,10 +2825,10 @@ def simple_fix_module_references(source_code: str, continue else: # Apply normal replacement - pattern = r'\b\w+\.' + re.escape(name) + r'\b' + pattern = r"\b\w+\." + re.escape(name) + r"\b" lines[i] = re.sub(pattern, name, line) - fixed_code = '\n'.join(lines) + fixed_code = "\n".join(lines) return fixed_code @@ -2454,34 +2836,35 @@ def simple_fix_module_references(source_code: str, def main(): """Command line interface for the isolate_object function.""" parser = argparse.ArgumentParser( - description= - 'Isolate an object from a Python file with all dependencies (including recursive imports)' + description="Isolate an object from a Python file with all dependencies (including recursive imports)" + ) + parser.add_argument("filename", help="Path to the Python file") + parser.add_argument( + "object_name", help="Name of the class or function to extract" ) - parser.add_argument('filename', help='Path to the Python file') - parser.add_argument('object_name', - help='Name of the class or function to extract') - parser.add_argument('-o', '--output', help='Output file (default: stdout)') + parser.add_argument("-o", "--output", help="Output file (default: stdout)") parser.add_argument( - '-d', - '--max-depth', - type=int, - default=5, - help='Maximum recursion depth for following imports (default: 5)') + "-d", + "--max-depth", + type=int, + default=5, + help="Maximum recursion depth for following imports (default: 5)", + ) parser.add_argument( - '--debug', - action='store_true', - help='Show debug information including exact paths of files being opened') + "--debug", + action="store_true", + help="Show debug information including exact paths of files being opened", + ) args = parser.parse_args() try: - result = isolate_object(args.filename, - args.object_name, - args.max_depth, - debug=args.debug) + result = isolate_object( + args.filename, args.object_name, args.max_depth, debug=args.debug + ) if args.output: - with open(args.output, 'w', encoding='utf-8') as f: + with open(args.output, "w", encoding="utf-8") as f: f.write(result) print(f"Isolated {args.object_name} written to {args.output}") else: diff --git a/MaxKernel/hitl_agent/prompts/__init__.py b/MaxKernel/hitl_agent/prompts/__init__.py index 7c2ccf6..3d1c3a7 100644 --- a/MaxKernel/hitl_agent/prompts/__init__.py +++ b/MaxKernel/hitl_agent/prompts/__init__.py @@ -6,4 +6,4 @@ from . import interactive_prompt -__all__ = ['interactive_prompt'] +__all__ = ["interactive_prompt"] diff --git a/MaxKernel/hitl_agent/server_utils/cpu_server.py b/MaxKernel/hitl_agent/server_utils/cpu_server.py index f9fb199..ec44f3c 100644 --- a/MaxKernel/hitl_agent/server_utils/cpu_server.py +++ b/MaxKernel/hitl_agent/server_utils/cpu_server.py @@ -8,13 +8,14 @@ from fastapi import FastAPI, HTTPException from pydantic import BaseModel -from hitl_agent.tools.analyze_profile import analyze_trace + from hitl_agent.constants import CPU_SERVER_PORT +from hitl_agent.tools.analyze_profile import analyze_trace logging.basicConfig( - level=logging.INFO, - format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", - datefmt="%Y-%m-%d %H:%M:%S", + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", ) app = FastAPI(title="CPU Code Execution Server", version="1.0.0") @@ -43,8 +44,8 @@ class GetBackendVersionResponse(BaseModel): def get_cpu_env(): """ - Returns environment variables that force JAX to use CPU backend. - """ + Returns environment variables that force JAX to use CPU backend. + """ env = os.environ.copy() env["JAX_PLATFORMS"] = "cpu" env["JAX_PLATFORM_NAME"] = "cpu" @@ -61,9 +62,9 @@ async def health_check(): @app.post("/compilation_test", response_model=CodeResponse) async def compilation_test(request: CodeRequest): """ - Try to execute kernel safely in a subprocess with CPU backend and return the output. - """ - logging.info(f"Starting compilation test on CPU backend") + Try to execute kernel safely in a subprocess with CPU backend and return the output. + """ + logging.info("Starting compilation test on CPU backend") async with compilation_semaphore: try: # Extract code from markdown format if present @@ -87,31 +88,33 @@ async def compilation_test(request: CodeRequest): request.code = code_content # Create a temporary file to store the code - with tempfile.NamedTemporaryFile(mode="w", suffix=".py", - delete=False) as temp_file: + with tempfile.NamedTemporaryFile( + mode="w", suffix=".py", delete=False + ) as temp_file: temp_file.write(request.code) temp_file_path = temp_file.name # Execute the code in a subprocess with CPU-only environment process = await asyncio.create_subprocess_exec( - sys.executable, - temp_file_path, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - cwd=tempfile.gettempdir(), - env=get_cpu_env(), # Force CPU backend + sys.executable, + temp_file_path, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=tempfile.gettempdir(), + env=get_cpu_env(), # Force CPU backend ) try: - stdout, stderr = await asyncio.wait_for(process.communicate(), - timeout=request.timeout) + stdout, stderr = await asyncio.wait_for( + process.communicate(), timeout=request.timeout + ) output = stdout.decode("utf-8") if stdout else "" error = stderr.decode("utf-8") if stderr else None exit_code = process.returncode logging.info( - f"Compilation test completed successfully on CPU with exit_code: {exit_code}" + f"Compilation test completed successfully on CPU with exit_code: {exit_code}" ) return CodeResponse(output=output, error=error, exit_code=exit_code) @@ -136,15 +139,15 @@ async def compilation_test(request: CodeRequest): os.unlink(temp_file_path) except OSError: pass - logging.info(f"Compilation test finished") + logging.info("Compilation test finished") @app.post("/correctness_test", response_model=CodeResponse) async def correctness_test(request: CodeRequest): """ - Test the correctness of the kernel code by executing it on CPU and comparing the output. - """ - logging.info(f"Starting correctness test on CPU backend") + Test the correctness of the kernel code by executing it on CPU and comparing the output. + """ + logging.info("Starting correctness test on CPU backend") async with correctness_semaphore: try: # Extract code from markdown format if present @@ -168,31 +171,33 @@ async def correctness_test(request: CodeRequest): request.code = code_content # Create a temporary file to store the code - with tempfile.NamedTemporaryFile(mode="w", suffix=".py", - delete=False) as temp_file: + with tempfile.NamedTemporaryFile( + mode="w", suffix=".py", delete=False + ) as temp_file: temp_file.write(request.code) temp_file_path = temp_file.name # Execute the code in a subprocess with CPU-only environment process = await asyncio.create_subprocess_exec( - sys.executable, - temp_file_path, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - cwd=tempfile.gettempdir(), - env=get_cpu_env(), # Force CPU backend + sys.executable, + temp_file_path, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=tempfile.gettempdir(), + env=get_cpu_env(), # Force CPU backend ) try: - stdout, stderr = await asyncio.wait_for(process.communicate(), - timeout=request.timeout) + stdout, stderr = await asyncio.wait_for( + process.communicate(), timeout=request.timeout + ) output = stdout.decode("utf-8") if stdout else "" error = stderr.decode("utf-8") if stderr else None exit_code = process.returncode logging.info( - f"Correctness test completed successfully on CPU with exit_code: {exit_code}" + f"Correctness test completed successfully on CPU with exit_code: {exit_code}" ) return CodeResponse(output=output, error=error, exit_code=exit_code) @@ -216,15 +221,15 @@ async def correctness_test(request: CodeRequest): os.unlink(temp_file_path) except OSError: pass - logging.info(f"Correctness test finished") + logging.info("Correctness test finished") @app.post("/performance_test", response_model=CodeResponse) async def performance_test(request: CodeRequest): """ - Test the performance of the kernel code by executing it on CPU and measuring the execution time. - """ - logging.info(f"Starting performance test on CPU backend") + Test the performance of the kernel code by executing it on CPU and measuring the execution time. + """ + logging.info("Starting performance test on CPU backend") async with performance_semaphore: try: # Extract code from markdown format if present @@ -248,31 +253,33 @@ async def performance_test(request: CodeRequest): request.code = code_content # Create a temporary file to store the code - with tempfile.NamedTemporaryFile(mode="w", suffix=".py", - delete=False) as temp_file: + with tempfile.NamedTemporaryFile( + mode="w", suffix=".py", delete=False + ) as temp_file: temp_file.write(request.code) temp_file_path = temp_file.name # Execute the code in a subprocess with CPU-only environment process = await asyncio.create_subprocess_exec( - sys.executable, - temp_file_path, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - cwd=tempfile.gettempdir(), - env=get_cpu_env(), # Force CPU backend + sys.executable, + temp_file_path, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=tempfile.gettempdir(), + env=get_cpu_env(), # Force CPU backend ) try: - stdout, stderr = await asyncio.wait_for(process.communicate(), - timeout=request.timeout) + stdout, stderr = await asyncio.wait_for( + process.communicate(), timeout=request.timeout + ) output = stdout.decode("utf-8") if stdout else "" error = stderr.decode("utf-8") if stderr else None exit_code = process.returncode logging.info( - f"Performance test completed successfully on CPU with exit_code: {exit_code}" + f"Performance test completed successfully on CPU with exit_code: {exit_code}" ) return CodeResponse(output=output, error=error, exit_code=exit_code) @@ -296,12 +303,12 @@ async def performance_test(request: CodeRequest): os.unlink(temp_file_path) except OSError: pass - logging.info(f"Performance test finished") + logging.info("Performance test finished") @app.post("/profile", response_model=CodeResponse) async def profile(request: CodeRequest): - logging.info(f"Starting profile on CPU backend") + logging.info("Starting profile on CPU backend") async with profile_semaphore: try: # Extract code from markdown format if present @@ -336,23 +343,24 @@ async def profile(request: CodeRequest): # Execute the code in a subprocess with CPU-only environment process = await asyncio.create_subprocess_exec( - sys.executable, - temp_file_path, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - cwd=temp_dir, - env=get_cpu_env(), # Force CPU backend + sys.executable, + temp_file_path, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=temp_dir, + env=get_cpu_env(), # Force CPU backend ) try: - stdout, stderr = await asyncio.wait_for(process.communicate(), - timeout=request.timeout) + stdout, stderr = await asyncio.wait_for( + process.communicate(), timeout=request.timeout + ) output = stdout.decode("utf-8") if stdout else "" error = stderr.decode("utf-8") if stderr else None exit_code = process.returncode - logging.info(f"Profile code executed, now analyzing trace.") + logging.info("Profile code executed, now analyzing trace.") # Recursively search for .xplane.pb file under temp_file_path directory xplane_pb_file = None for root, _, files in os.walk(temp_dir): @@ -368,16 +376,13 @@ async def profile(request: CodeRequest): ratio = analyze_trace(xplane_pb_file) logging.info( - f"Profile analysis completed successfully on CPU with exit_code: {exit_code}" + f"Profile analysis completed successfully on CPU with exit_code: {exit_code}" ) return CodeResponse( - output=json.dumps({ - "ratio": ratio, - "xplane_path": xplane_pb_file - }), - error=error, - exit_code=exit_code, + output=json.dumps({"ratio": ratio, "xplane_path": xplane_pb_file}), + error=error, + exit_code=exit_code, ) except asyncio.TimeoutError: @@ -399,17 +404,17 @@ async def profile(request: CodeRequest): # shutil.rmtree(temp_dir) # except Exception: # pass - logging.info(f"Profile analysis finished") + logging.info("Profile analysis finished") @app.post("/get_backend_version", response_model=GetBackendVersionResponse) async def get_backend_version() -> str: """ - Returns the backend version for CPU execution. + Returns the backend version for CPU execution. - Returns: - A string indicating CPU backend. - """ + Returns: + A string indicating CPU backend. + """ return GetBackendVersionResponse(backend_version="CPU") diff --git a/MaxKernel/hitl_agent/server_utils/eval_server.py b/MaxKernel/hitl_agent/server_utils/eval_server.py index bc759ff..d614524 100644 --- a/MaxKernel/hitl_agent/server_utils/eval_server.py +++ b/MaxKernel/hitl_agent/server_utils/eval_server.py @@ -2,21 +2,22 @@ import logging from enum import Enum from typing import Optional -import asyncio + import aiohttp import yaml from fastapi import FastAPI, HTTPException from pydantic import BaseModel + from hitl_agent.constants import ( - EVAL_SERVER_PORT, - TPU_TIMEOUT, + EVAL_SERVER_PORT, + TPU_TIMEOUT, ) from hitl_agent.server_utils.tpu_server import CodeResponse, get_tpu_version logging.basicConfig( - level=logging.INFO, - format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", - datefmt="%Y-%m-%d %H:%M:%S", + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", ) app = FastAPI(title="Agent Evaluation Server", version="1.0.0") @@ -25,7 +26,6 @@ class Backend: - def __init__(self, name: str, ip: str, port: int, backend_type: str = "tpu"): self.name = name self.ip = ip @@ -84,7 +84,6 @@ class EvalRequest(BaseModel): class Evaluator: - def __init__(self, cfg_path="eval_config.yaml"): with open(cfg_path, "r") as file: self.config = yaml.safe_load(file) @@ -95,26 +94,26 @@ def __init__(self, cfg_path="eval_config.yaml"): logging.info("Using 'backends' configuration format") for backend_config in self.config["backends"]: backend_obj = Backend( - name=backend_config["name"], - ip=backend_config["ip"], - port=backend_config["port"], - backend_type=backend_config.get("type", "tpu"), + name=backend_config["name"], + ip=backend_config["ip"], + port=backend_config["port"], + backend_type=backend_config.get("type", "tpu"), ) self.backends.append(backend_obj) else: raise ValueError( - "No backends configured in eval_config.yaml. Please use the 'backends' format." + "No backends configured in eval_config.yaml. Please use the 'backends' format." ) logging.info(f"Evaluator initialized with backends: {self.backends}") async def get_available_backend(self, backend_type: Optional[str] = None): """ - Get an available backend, optionally filtered by type. + Get an available backend, optionally filtered by type. - Args: - backend_type: If specified, only return backends of this type ("tpu" or "cpu") - """ + Args: + backend_type: If specified, only return backends of this type ("tpu" or "cpu") + """ while True: for backend in self.backends: if backend.get_status() == "available": @@ -140,7 +139,8 @@ async def evaluate(request: EvalRequest): async with tpu_semaphore: # Get available backend, optionally filtered by requested backend type backend = await evaluator.get_available_backend( - backend_type=request.backend_type) + backend_type=request.backend_type + ) backend_ip = backend.ip backend_port = backend.port backend_name = backend.name @@ -149,41 +149,43 @@ async def evaluate(request: EvalRequest): try: # Start evaluation process - requested_type_msg = (f" (requested: {request.backend_type})" - if request.backend_type else "") + requested_type_msg = ( + f" (requested: {request.backend_type})" if request.backend_type else "" + ) logging.info( - f"Starting evaluation on {backend_type} backend '{backend_name}' ({backend_ip}:{backend_port}) for {request.eval_type.value}{requested_type_msg}" + f"Starting evaluation on {backend_type} backend '{backend_name}' ({backend_ip}:{backend_port}) for {request.eval_type.value}{requested_type_msg}" ) # Send request to backend server async with aiohttp.ClientSession() as session: async with session.post( - f"http://{backend_ip}:{backend_port}/{request.eval_type.value}", - json={ - "eval_type": request.eval_type.value, - "code": request.code, - "timeout": request.timeout, - }, + f"http://{backend_ip}:{backend_port}/{request.eval_type.value}", + json={ + "eval_type": request.eval_type.value, + "code": request.code, + "timeout": request.timeout, + }, ) as response: result = await response.json() logging.info( - f"Received response from {backend_name}: {response.status}") + f"Received response from {backend_name}: {response.status}" + ) if response.status != 200: if response.status == 408: raise HTTPException( - status_code=408, - detail= - f"Backend evaluation timed out. Timeout was set to {TPU_TIMEOUT} seconds.", + status_code=408, + detail=f"Backend evaluation timed out. Timeout was set to {TPU_TIMEOUT} seconds.", ) raise HTTPException( - status_code=response.status, - detail=result.get("detail", "Backend evaluation failed"), + status_code=response.status, + detail=result.get("detail", "Backend evaluation failed"), ) # Mark backend as available after evaluation logging.info( - f"Evaluation completed on {backend_name}, marking as available") + f"Evaluation completed on {backend_name}, marking as available" + ) backend.set_status("available") return result diff --git a/MaxKernel/hitl_agent/server_utils/server_manager_mixin.py b/MaxKernel/hitl_agent/server_utils/server_manager_mixin.py index 3f2a837..aa45b18 100644 --- a/MaxKernel/hitl_agent/server_utils/server_manager_mixin.py +++ b/MaxKernel/hitl_agent/server_utils/server_manager_mixin.py @@ -1,45 +1,46 @@ """Mixin class for managing TPU and eval server lifecycle in kernel evaluation agents.""" import asyncio -import aiohttp import logging import os import subprocess +import aiohttp + class ServerManagerMixin: """Mixin that provides server lifecycle management for evaluation agents. - This mixin adds the ability to automatically start and stop TPU and eval servers - before and after agent execution. Agents that inherit from this mixin should: + This mixin adds the ability to automatically start and stop TPU and eval servers + before and after agent execution. Agents that inherit from this mixin should: - 1. Set auto_manage_servers=True to enable automatic server management - 2. Initialize self._servers_started = [] in their __init__ - 3. Call await self._ensure_servers_running() before operations that need servers - 4. Call await self._cleanup_servers() in a finally block after operations complete + 1. Set auto_manage_servers=True to enable automatic server management + 2. Initialize self._servers_started = [] in their __init__ + 3. Call await self._ensure_servers_running() before operations that need servers + 4. Call await self._cleanup_servers() in a finally block after operations complete - The mixin tracks which servers it started and only tears down those servers, - so it won't interfere with already-running servers or servers started by other agents. + The mixin tracks which servers it started and only tears down those servers, + so it won't interfere with already-running servers or servers started by other agents. - Attributes: - auto_manage_servers: Should be set to True in child classes to enable server management - _servers_started: List tracking which servers this instance started (for cleanup) - """ + Attributes: + auto_manage_servers: Should be set to True in child classes to enable server management + _servers_started: List tracking which servers this instance started (for cleanup) + """ def _is_server_running(self, server_name: str) -> bool: """Check if a server process is running. - Args: - server_name: Name of the server process to check (e.g., "tpu_server.py") + Args: + server_name: Name of the server process to check (e.g., "tpu_server.py") - Returns: - True if the server is running, False otherwise - """ + Returns: + True if the server is running, False otherwise + """ try: result = subprocess.run( - ["pgrep", "-f", server_name], - capture_output=True, - text=True, + ["pgrep", "-f", server_name], + capture_output=True, + text=True, ) return result.returncode == 0 except Exception as e: @@ -47,50 +48,51 @@ def _is_server_running(self, server_name: str) -> bool: return False async def _wait_for_server_ready( - self, - port: int, - server_type: str, - max_retries: int = 30, - retry_delay: float = 1.0, + self, + port: int, + server_type: str, + max_retries: int = 30, + retry_delay: float = 1.0, ) -> bool: """Wait for a server's HTTP health endpoint to be ready. - Args: - port: Port number the server is listening on - server_type: Type of server for logging (e.g., "cpu", "tpu", "eval") - max_retries: Maximum number of health check attempts - retry_delay: Delay in seconds between retry attempts + Args: + port: Port number the server is listening on + server_type: Type of server for logging (e.g., "cpu", "tpu", "eval") + max_retries: Maximum number of health check attempts + retry_delay: Delay in seconds between retry attempts - Returns: - True if server health check succeeds, False otherwise - """ + Returns: + True if server health check succeeds, False otherwise + """ health_url = f"http://localhost:{port}/health" for attempt in range(max_retries): try: async with aiohttp.ClientSession() as session: async with session.get( - health_url, timeout=aiohttp.ClientTimeout(total=2)) as response: + health_url, timeout=aiohttp.ClientTimeout(total=2) + ) as response: if response.status == 200: data = await response.json() if data.get("status") == "healthy": logging.info( - f"[ServerManager] {server_type} server health check passed on attempt {attempt + 1}" + f"[ServerManager] {server_type} server health check passed on attempt {attempt + 1}" ) return True except (aiohttp.ClientError, asyncio.TimeoutError) as e: if attempt < max_retries - 1: logging.debug( - f"[ServerManager] {server_type} server not ready yet (attempt {attempt + 1}/{max_retries}), retrying..." + f"[ServerManager] {server_type} server not ready yet (attempt {attempt + 1}/{max_retries}), retrying..." ) await asyncio.sleep(retry_delay) else: logging.error( - f"[ServerManager] {server_type} server health check failed after {max_retries} attempts: {e}" + f"[ServerManager] {server_type} server health check failed after {max_retries} attempts: {e}" ) except Exception as e: logging.error( - f"[ServerManager] Unexpected error during {server_type} health check: {e}" + f"[ServerManager] Unexpected error during {server_type} health check: {e}" ) return False @@ -99,25 +101,25 @@ async def _wait_for_server_ready( async def _start_server(self, server_type: str, setup_script: str) -> bool: """Start a specific server (tpu, cpu, or eval). - Args: - server_type: Type of server to start ("tpu", "cpu", or "eval") - setup_script: Path to the setup.sh script + Args: + server_type: Type of server to start ("tpu", "cpu", or "eval") + setup_script: Path to the setup.sh script - Returns: - True if server started successfully, False otherwise - """ + Returns: + True if server started successfully, False otherwise + """ try: logging.info(f"Starting {server_type} server...") # Determine the directory containing the setup script setup_dir = os.path.dirname(setup_script) process = await asyncio.create_subprocess_exec( - "bash", - setup_script, - f"--start-{server_type}", - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - cwd=setup_dir, # Run from the directory containing setup.sh + "bash", + setup_script, + f"--start-{server_type}", + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=setup_dir, # Run from the directory containing setup.sh ) await process.wait() @@ -132,9 +134,9 @@ async def _start_server(self, server_type: str, setup_script: str) -> bool: # Map server type to port (from constants.py) port_map = { - "tpu": 5463, # TPU_SERVER_PORT - "cpu": 5464, # CPU_SERVER_PORT - "eval": 1245, # EVAL_SERVER_PORT + "tpu": 5463, # TPU_SERVER_PORT + "cpu": 5464, # CPU_SERVER_PORT + "eval": 1245, # EVAL_SERVER_PORT } port = port_map.get(server_type) @@ -144,14 +146,16 @@ async def _start_server(self, server_type: str, setup_script: str) -> bool: # Wait for HTTP health endpoint to be ready logging.info( - f"Waiting for {server_type} server HTTP endpoint to be ready...") + f"Waiting for {server_type} server HTTP endpoint to be ready..." + ) if await self._wait_for_server_ready(port, server_type): logging.info(f"{server_type} server started successfully and is ready") self._servers_started.append(server_type) return True else: logging.error( - f"Failed to start {server_type} server - health check failed") + f"Failed to start {server_type} server - health check failed" + ) return False except Exception as e: logging.error(f"Exception starting {server_type} server: {e}") @@ -160,45 +164,47 @@ async def _start_server(self, server_type: str, setup_script: str) -> bool: def _stop_server_sync(self, process_name: str): """Stop a specific server synchronously using pkill. - Args: - process_name: Name of the server process to stop (e.g., "tpu_server.py") - """ + Args: + process_name: Name of the server process to stop (e.g., "tpu_server.py") + """ try: logging.info(f"Stopping {process_name}...") result = subprocess.run( - ["pkill", "-f", process_name], - capture_output=True, - text=True, - timeout=5, + ["pkill", "-f", process_name], + capture_output=True, + text=True, + timeout=5, ) - if (result.returncode == 0 or - result.returncode == 1): # 1 means no process found + if ( + result.returncode == 0 or result.returncode == 1 + ): # 1 means no process found logging.info(f"{process_name} stopped") else: logging.warning( - f"pkill returned {result.returncode} for {process_name}") + f"pkill returned {result.returncode} for {process_name}" + ) except Exception as e: logging.error(f"Exception stopping {process_name}: {e}") async def _ensure_servers_running(self) -> tuple[bool, str]: """Ensure TPU and eval servers are running. - Checks if servers are running and starts them if needed. Only starts servers - if auto_manage_servers is True. Tracks which servers were started by this agent - in self._servers_started so they can be cleaned up later. + Checks if servers are running and starts them if needed. Only starts servers + if auto_manage_servers is True. Tracks which servers were started by this agent + in self._servers_started so they can be cleaned up later. - Returns: - Tuple of (success: bool, error_message: str) - - success: True if servers are running or were started successfully - - error_message: Empty string on success, error description on failure - """ + Returns: + Tuple of (success: bool, error_message: str) + - success: True if servers are running or were started successfully + - error_message: Empty string on success, error description on failure + """ logging.info( - f"[ServerManager] _ensure_servers_running called, auto_manage_servers={getattr(self, 'auto_manage_servers', 'NOT_SET')}" + f"[ServerManager] _ensure_servers_running called, auto_manage_servers={getattr(self, 'auto_manage_servers', 'NOT_SET')}" ) if not getattr(self, "auto_manage_servers", False): logging.info( - "[ServerManager] auto_manage_servers is False, skipping server management" + "[ServerManager] auto_manage_servers is False, skipping server management" ) return True, "" @@ -215,7 +221,8 @@ async def _ensure_servers_running(self) -> tuple[bool, str]: # Check and start CPU server (required by eval server) if self._is_server_running("cpu_server.py"): logging.info( - "CPU server already running, restarting to ensure fresh state") + "CPU server already running, restarting to ensure fresh state" + ) self._stop_server_sync("cpu_server.py") await asyncio.sleep(1) # Wait for graceful shutdown @@ -226,7 +233,8 @@ async def _ensure_servers_running(self) -> tuple[bool, str]: # Check and start TPU server if self._is_server_running("tpu_server.py"): logging.info( - "TPU server already running, restarting to ensure fresh state") + "TPU server already running, restarting to ensure fresh state" + ) self._stop_server_sync("tpu_server.py") await asyncio.sleep(1) # Wait for graceful shutdown @@ -237,7 +245,8 @@ async def _ensure_servers_running(self) -> tuple[bool, str]: # Check and start eval server if self._is_server_running("eval_server.py"): logging.info( - "Eval server already running, restarting to ensure fresh state") + "Eval server already running, restarting to ensure fresh state" + ) self._stop_server_sync("eval_server.py") await asyncio.sleep(1) # Wait for graceful shutdown @@ -250,10 +259,10 @@ async def _ensure_servers_running(self) -> tuple[bool, str]: async def _cleanup_servers(self): """Stop servers that were started by this agent. - Only stops servers that this agent instance started (tracked in self._servers_started). - This ensures we don't accidentally tear down servers that were already running - or started by other agents. - """ + Only stops servers that this agent instance started (tracked in self._servers_started). + This ensures we don't accidentally tear down servers that were already running + or started by other agents. + """ if not self._servers_started: return diff --git a/MaxKernel/hitl_agent/server_utils/tpu_server.py b/MaxKernel/hitl_agent/server_utils/tpu_server.py index 94b1e3f..b73d2b5 100644 --- a/MaxKernel/hitl_agent/server_utils/tpu_server.py +++ b/MaxKernel/hitl_agent/server_utils/tpu_server.py @@ -3,7 +3,6 @@ import logging import os import re -import shutil import subprocess import sys import tempfile @@ -11,13 +10,14 @@ from fastapi import FastAPI, HTTPException from pydantic import BaseModel -from hitl_agent.tools.analyze_profile import analyze_trace + from hitl_agent.constants import TPU_SERVER_PORT +from hitl_agent.tools.analyze_profile import analyze_trace logging.basicConfig( - level=logging.INFO, - format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", - datefmt="%Y-%m-%d %H:%M:%S", + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", ) app = FastAPI(title="TPU Code Execution Server", version="1.0.0") @@ -52,9 +52,9 @@ async def health_check(): @app.post("/compilation_test", response_model=CodeResponse) async def compilation_test(request: CodeRequest): """ - Try to execute kernel safely in a subprocess and return the output. - """ - logging.info(f"Starting compilation test") + Try to execute kernel safely in a subprocess and return the output. + """ + logging.info("Starting compilation test") async with compilation_semaphore: try: # Extract code from markdown format if present @@ -78,30 +78,32 @@ async def compilation_test(request: CodeRequest): request.code = code_content # Create a temporary file to store the code - with tempfile.NamedTemporaryFile(mode="w", suffix=".py", - delete=False) as temp_file: + with tempfile.NamedTemporaryFile( + mode="w", suffix=".py", delete=False + ) as temp_file: temp_file.write(request.code) temp_file_path = temp_file.name # Execute the code in a subprocess process = await asyncio.create_subprocess_exec( - sys.executable, - temp_file_path, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - cwd=tempfile.gettempdir(), + sys.executable, + temp_file_path, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=tempfile.gettempdir(), ) try: - stdout, stderr = await asyncio.wait_for(process.communicate(), - timeout=request.timeout) + stdout, stderr = await asyncio.wait_for( + process.communicate(), timeout=request.timeout + ) output = stdout.decode("utf-8") if stdout else "" error = stderr.decode("utf-8") if stderr else None exit_code = process.returncode logging.info( - f"Compilation test completed successfully with exit_code: {exit_code}" + f"Compilation test completed successfully with exit_code: {exit_code}" ) return CodeResponse(output=output, error=error, exit_code=exit_code) @@ -126,15 +128,15 @@ async def compilation_test(request: CodeRequest): os.unlink(temp_file_path) except OSError: pass - logging.info(f"Compilation test finished") + logging.info("Compilation test finished") @app.post("/correctness_test", response_model=CodeResponse) async def correctness_test(request: CodeRequest): """ - Test the correctness of the kernel code by executing it and comparing the output. - """ - logging.info(f"Starting correctness test") + Test the correctness of the kernel code by executing it and comparing the output. + """ + logging.info("Starting correctness test") async with correctness_semaphore: try: # Extract code from markdown format if present @@ -158,30 +160,32 @@ async def correctness_test(request: CodeRequest): request.code = code_content # Create a temporary file to store the code - with tempfile.NamedTemporaryFile(mode="w", suffix=".py", - delete=False) as temp_file: + with tempfile.NamedTemporaryFile( + mode="w", suffix=".py", delete=False + ) as temp_file: temp_file.write(request.code) temp_file_path = temp_file.name # Execute the code in a subprocess process = await asyncio.create_subprocess_exec( - sys.executable, - temp_file_path, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - cwd=tempfile.gettempdir(), + sys.executable, + temp_file_path, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=tempfile.gettempdir(), ) try: - stdout, stderr = await asyncio.wait_for(process.communicate(), - timeout=request.timeout) + stdout, stderr = await asyncio.wait_for( + process.communicate(), timeout=request.timeout + ) output = stdout.decode("utf-8") if stdout else "" error = stderr.decode("utf-8") if stderr else None exit_code = process.returncode logging.info( - f"Correctness test completed successfully with exit_code: {exit_code}" + f"Correctness test completed successfully with exit_code: {exit_code}" ) return CodeResponse(output=output, error=error, exit_code=exit_code) @@ -205,15 +209,15 @@ async def correctness_test(request: CodeRequest): os.unlink(temp_file_path) except OSError: pass - logging.info(f"Correctness test finished") + logging.info("Correctness test finished") @app.post("/performance_test", response_model=CodeResponse) async def performance_test(request: CodeRequest): """ - Test the performance of the kernel code by executing it and measuring the execution time. - """ - logging.info(f"Starting performance test") + Test the performance of the kernel code by executing it and measuring the execution time. + """ + logging.info("Starting performance test") async with performance_semaphore: try: # Extract code from markdown format if present @@ -237,30 +241,32 @@ async def performance_test(request: CodeRequest): request.code = code_content # Create a temporary file to store the code - with tempfile.NamedTemporaryFile(mode="w", suffix=".py", - delete=False) as temp_file: + with tempfile.NamedTemporaryFile( + mode="w", suffix=".py", delete=False + ) as temp_file: temp_file.write(request.code) temp_file_path = temp_file.name # Execute the code in a subprocess process = await asyncio.create_subprocess_exec( - sys.executable, - temp_file_path, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - cwd=tempfile.gettempdir(), + sys.executable, + temp_file_path, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=tempfile.gettempdir(), ) try: - stdout, stderr = await asyncio.wait_for(process.communicate(), - timeout=request.timeout) + stdout, stderr = await asyncio.wait_for( + process.communicate(), timeout=request.timeout + ) output = stdout.decode("utf-8") if stdout else "" error = stderr.decode("utf-8") if stderr else None exit_code = process.returncode logging.info( - f"Performance test completed successfully with exit_code: {exit_code}" + f"Performance test completed successfully with exit_code: {exit_code}" ) return CodeResponse(output=output, error=error, exit_code=exit_code) @@ -284,12 +290,12 @@ async def performance_test(request: CodeRequest): os.unlink(temp_file_path) except OSError: pass - logging.info(f"Performance test finished") + logging.info("Performance test finished") @app.post("/profile", response_model=CodeResponse) async def profile(request: CodeRequest): - logging.info(f"Starting profile") + logging.info("Starting profile") async with profile_semaphore: try: # Extract code from markdown format if present @@ -324,22 +330,23 @@ async def profile(request: CodeRequest): # Execute the code in a subprocess process = await asyncio.create_subprocess_exec( - sys.executable, - temp_file_path, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - cwd=temp_dir, + sys.executable, + temp_file_path, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=temp_dir, ) try: - stdout, stderr = await asyncio.wait_for(process.communicate(), - timeout=request.timeout) + stdout, stderr = await asyncio.wait_for( + process.communicate(), timeout=request.timeout + ) output = stdout.decode("utf-8") if stdout else "" error = stderr.decode("utf-8") if stderr else None exit_code = process.returncode - logging.info(f"Profile code executed, now analyzing trace.") + logging.info("Profile code executed, now analyzing trace.") # Recursively search for .xplane.pb file under temp_file_path directory xplane_pb_file = None for root, _, files in os.walk(temp_dir): @@ -357,9 +364,7 @@ async def profile(request: CodeRequest): for fname in files: all_files.append(os.path.join(root, fname)) - error_msg = ( - "No .xplane.pb trace file found after profiling. Files in" - f" temp_dir: {all_files[:10]}") + error_msg = f"No .xplane.pb trace file found after profiling. Files in temp_dir: {all_files[:10]}" logging.error(error_msg) # Return the execution output/error to help diagnose @@ -374,16 +379,14 @@ async def profile(request: CodeRequest): ratio = analyze_trace(xplane_pb_file) - logging.info("Profile analysis completed successfully with exit_code:" - f" {exit_code}") + logging.info( + f"Profile analysis completed successfully with exit_code: {exit_code}" + ) return CodeResponse( - output=json.dumps({ - "ratio": ratio, - "xplane_path": xplane_pb_file - }), - error=error, - exit_code=exit_code, + output=json.dumps({"ratio": ratio, "xplane_path": xplane_pb_file}), + error=error, + exit_code=exit_code, ) except asyncio.TimeoutError: @@ -405,7 +408,7 @@ async def profile(request: CodeRequest): # shutil.rmtree(temp_dir) # except Exception: # pass - logging.info(f"Profile analysis finished") + logging.info("Profile analysis finished") @app.post("/get_tpu_version", response_model=GetTpuVersionResponse) @@ -426,10 +429,9 @@ async def get_tpu_version() -> str: # --- Method 1: Try running the `tpu-info` command-line tool (No resource conflicts) --- try: # Run the tpu-info command - result = subprocess.run(["tpu-info"], - capture_output=True, - text=True, - check=True) + result = subprocess.run( + ["tpu-info"], capture_output=True, text=True, check=True + ) # Regex to find a pattern like "TPU v4", "TPU v5e", "TPU v3 chip", etc. # We search the entire output diff --git a/MaxKernel/hitl_agent/subagents/__init__.py b/MaxKernel/hitl_agent/subagents/__init__.py index 1378e32..91a6aaa 100644 --- a/MaxKernel/hitl_agent/subagents/__init__.py +++ b/MaxKernel/hitl_agent/subagents/__init__.py @@ -8,16 +8,12 @@ - gpu_to_jax: GPU-to-JAX code conversion """ -from . import kernel_writing -from . import testing -from . import profiling -from . import explanation -from . import gpu_to_jax_agent +from . import explanation, gpu_to_jax_agent, kernel_writing, profiling, testing __all__ = [ - 'kernel_writing', - 'testing', - 'profiling', - 'explanation', - 'gpu_to_jax_agent', + "kernel_writing", + "testing", + "profiling", + "explanation", + "gpu_to_jax_agent", ] diff --git a/MaxKernel/hitl_agent/subagents/explanation/__init__.py b/MaxKernel/hitl_agent/subagents/explanation/__init__.py index e8e7334..423ebc3 100644 --- a/MaxKernel/hitl_agent/subagents/explanation/__init__.py +++ b/MaxKernel/hitl_agent/subagents/explanation/__init__.py @@ -1,11 +1,11 @@ """Explanation subagent module.""" from .agent import ( - explanation_agent, - explanation_llm_agent, + explanation_agent, + explanation_llm_agent, ) __all__ = [ - 'explanation_agent', - 'explanation_llm_agent', + "explanation_agent", + "explanation_llm_agent", ] diff --git a/MaxKernel/hitl_agent/subagents/explanation/agent.py b/MaxKernel/hitl_agent/subagents/explanation/agent.py index 3de6b2b..d9adb41 100644 --- a/MaxKernel/hitl_agent/subagents/explanation/agent.py +++ b/MaxKernel/hitl_agent/subagents/explanation/agent.py @@ -2,36 +2,36 @@ from google.adk.agents import SequentialAgent -from hitl_agent.custom_types import CustomLlmAgent -from hitl_agent.constants import MODEL_NAME from hitl_agent.config import model_config, thinking_planner +from hitl_agent.constants import MODEL_NAME +from hitl_agent.custom_types import CustomLlmAgent +from hitl_agent.subagents.explanation.prompts import explanation_prompt from hitl_agent.tools.tools import ( - filesystem_tool_rw, - vertex_ai_rag_tool, + filesystem_tool_rw, + vertex_ai_rag_tool, ) -from hitl_agent.subagents.explanation.prompts import explanation_prompt # Explanation LLM agent explanation_llm_agent = CustomLlmAgent( - name="ExplanationLlmAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=thinking_planner, - instruction=explanation_prompt.PROMPT, - description="Provides explanations for the kernel generation process.", - tools=[filesystem_tool_rw, vertex_ai_rag_tool] - if vertex_ai_rag_tool else [filesystem_tool_rw], + name="ExplanationLlmAgent", + model=MODEL_NAME, + generate_content_config=model_config, + planner=thinking_planner, + instruction=explanation_prompt.PROMPT, + description="Provides explanations for the kernel generation process.", + tools=[filesystem_tool_rw, vertex_ai_rag_tool] + if vertex_ai_rag_tool + else [filesystem_tool_rw], ) # Explanation orchestrator agent explanation_agent = SequentialAgent( - name="ExplanationAgent", - sub_agents=[explanation_llm_agent], - description= - "Provides explanations for the kernel generation process and returns control to orchestration.", + name="ExplanationAgent", + sub_agents=[explanation_llm_agent], + description="Provides explanations for the kernel generation process and returns control to orchestration.", ) __all__ = [ - 'explanation_agent', - 'explanation_llm_agent', + "explanation_agent", + "explanation_llm_agent", ] diff --git a/MaxKernel/hitl_agent/subagents/explanation/prompts/__init__.py b/MaxKernel/hitl_agent/subagents/explanation/prompts/__init__.py index a5611ee..1806c65 100644 --- a/MaxKernel/hitl_agent/subagents/explanation/prompts/__init__.py +++ b/MaxKernel/hitl_agent/subagents/explanation/prompts/__init__.py @@ -2,4 +2,4 @@ from . import explanation_prompt -__all__ = ['explanation_prompt'] +__all__ = ["explanation_prompt"] diff --git a/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/__init__.py b/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/__init__.py index 01fb272..6d57fed 100644 --- a/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/__init__.py +++ b/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/__init__.py @@ -5,7 +5,8 @@ """ from hitl_agent.subagents.gpu_to_jax_agent.agent import ( - gpu_to_jax_agent,) + gpu_to_jax_agent, +) __version__ = "1.0.0" __all__ = ["gpu_to_jax_agent"] diff --git a/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/agent.py b/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/agent.py index 4a09ad8..fc527fa 100644 --- a/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/agent.py +++ b/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/agent.py @@ -8,51 +8,51 @@ warnings.filterwarnings("ignore", message=".*EXPERIMENTAL.*") warnings.filterwarnings("ignore", category=UserWarning) +import os + from google.adk.agents.callback_context import CallbackContext from google.adk.tools import AgentTool, BaseTool, FunctionTool, ToolContext -from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset from google.adk.tools.mcp_tool.mcp_session_manager import StdioConnectionParams +from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset from google.genai import types from mcp import StdioServerParameters -import os +from hitl_agent.custom_types import CustomLlmAgent from hitl_agent.subagents.gpu_to_jax_agent.constants import ( - MODEL_NAME, - TOP_P, - TOP_K, -) -from hitl_agent.subagents.gpu_to_jax_agent.prompts import ( - convert_simplified_to_jax_prompt, - fix_conversion_prompt, - fix_conversion_extended_prompt, - identify_framework_prompt, - summary_prompt, - generate_summary_extended_prompt, - orchestrator_prompt, - analyze_plan_prompt, - simplify_gpu_code_prompt, - write_readme_prompt, - validate_syntax_routing_prompt, - validate_compilation_routing_prompt, - validate_shapes_routing_prompt, - generate_test_prompt, - generate_test_extended_prompt, - run_test_routing_prompt, + MODEL_NAME, + TOP_K, + TOP_P, ) from hitl_agent.subagents.gpu_to_jax_agent.evaluators import ( - JaxSyntaxChecker, - ShapeValidator, - JaxCompilationChecker, - JaxCorrectnessChecker, + JaxCompilationChecker, + JaxCorrectnessChecker, + JaxSyntaxChecker, + ShapeValidator, +) +from hitl_agent.subagents.gpu_to_jax_agent.prompts import ( + analyze_plan_prompt, + convert_simplified_to_jax_prompt, + fix_conversion_extended_prompt, + fix_conversion_prompt, + generate_summary_extended_prompt, + generate_test_extended_prompt, + generate_test_prompt, + identify_framework_prompt, + orchestrator_prompt, + run_test_routing_prompt, + simplify_gpu_code_prompt, + summary_prompt, + validate_compilation_routing_prompt, + validate_shapes_routing_prompt, + validate_syntax_routing_prompt, + write_readme_prompt, ) -from hitl_agent.custom_types import CustomLlmAgent -from hitl_agent.tools.search_api_tool import search_api_tool # Model configuration model_config = types.GenerateContentConfig( - temperature=0.1, - top_p=TOP_P, - top_k=TOP_K, + temperature=0.1, + top_p=TOP_P, + top_k=TOP_K, ) WORKDIR = os.environ.get("WORKDIR", os.path.dirname(os.path.abspath(__file__))) @@ -62,24 +62,24 @@ # We wrap this to handle errors gracefully try: filesystem_tool_rw = MCPToolset( - connection_params=StdioConnectionParams( - server_params=StdioServerParameters( - command='npx', - args=[ - "-y", - "@modelcontextprotocol/server-filesystem@0.5.1", - os.path.abspath(WORKDIR), - ], - env={ - **os.environ, "MCP_LOG_LEVEL": "error" - }, # Suppress info messages - ),), - tool_filter=['list_directory', 'read_file', 'write_file']) + connection_params=StdioConnectionParams( + server_params=StdioServerParameters( + command="npx", + args=[ + "-y", + "@modelcontextprotocol/server-filesystem@0.5.1", + os.path.abspath(WORKDIR), + ], + env={**os.environ, "MCP_LOG_LEVEL": "error"}, # Suppress info messages + ), + ), + tool_filter=["list_directory", "read_file", "write_file"], + ) logging.info("MCP filesystem toolset initialized successfully") except Exception as e: logging.error(f"Failed to initialize MCP filesystem toolset: {e}") logging.warning( - "MCP tools may not be available. write_file_direct will be used as fallback." + "MCP tools may not be available. write_file_direct will be used as fallback." ) # Create a minimal placeholder - agents should prefer write_file_direct anyway filesystem_tool_rw = None @@ -89,17 +89,17 @@ def get_output_directory(tool_context: ToolContext, file_path: str) -> str: """Determine the output directory for generated files. - Priority: - 1. Directory of original GPU code file (if available and file_path is a simple filename) - 2. WORKDIR (fallback for paths with subdirectories) + Priority: + 1. Directory of original GPU code file (if available and file_path is a simple filename) + 2. WORKDIR (fallback for paths with subdirectories) - Args: - tool_context: Context containing state with original_gpu_code_path - file_path: The relative path being written (to check if it has subdirectories) + Args: + tool_context: Context containing state with original_gpu_code_path + file_path: The relative path being written (to check if it has subdirectories) - Returns: - Absolute path to output directory - """ + Returns: + Absolute path to output directory + """ # Check if we have an original GPU code path original_gpu_path = tool_context.state.get("original_gpu_code_path") @@ -110,11 +110,12 @@ def get_output_directory(tool_context: ToolContext, file_path: str) -> str: # This prevents doubling paths like /path/to/gpu/dir/subdir/file.py if original_gpu_path and os.path.exists(original_gpu_path): # Check if file_path has directory components - if os.path.dirname(file_path) == '': + if os.path.dirname(file_path) == "": # Simple filename - use GPU file's directory output_dir = os.path.dirname(os.path.abspath(original_gpu_path)) logging.info( - f"Using output directory from original GPU file: {output_dir}") + f"Using output directory from original GPU file: {output_dir}" + ) return output_dir else: # Path has subdirectories - use WORKDIR to avoid path duplication @@ -127,13 +128,14 @@ def get_output_directory(tool_context: ToolContext, file_path: str) -> str: # Direct file writing tool (bypasses MCP for reliability) -def write_file_direct(path: str, content: str, - tool_context: ToolContext) -> str: +def write_file_direct( + path: str, content: str, tool_context: ToolContext +) -> str: """Write content to a file. Writes to the directory of the original GPU file if available, otherwise WORKDIR.""" try: # Ensure path is relative if os.path.isabs(path): - error_msg = f"Error: Absolute paths not allowed. Use relative path like 'filename.md'" + error_msg = "Error: Absolute paths not allowed. Use relative path like 'filename.md'" logging.error(error_msg) return error_msg @@ -154,9 +156,10 @@ def write_file_direct(path: str, content: str, # Check if path is within any allowed directory is_allowed = any( - abs_full_path.startswith(allowed_dir) for allowed_dir in allowed_dirs) + abs_full_path.startswith(allowed_dir) for allowed_dir in allowed_dirs + ) if not is_allowed: - error_msg = f"Error: Path escapes allowed directories" + error_msg = "Error: Path escapes allowed directories" logging.error(error_msg) return error_msg @@ -164,13 +167,15 @@ def write_file_direct(path: str, content: str, try: os.makedirs(os.path.dirname(full_path), exist_ok=True) - with open(full_path, 'w', encoding='utf-8') as f: + with open(full_path, "w", encoding="utf-8") as f: f.write(content) # Save to state tool_context.state["most_recent_file_path"] = full_path - success_msg = f"Successfully wrote {len(content)} characters to {full_path}" + success_msg = ( + f"Successfully wrote {len(content)} characters to {full_path}" + ) logging.info(success_msg) return success_msg @@ -178,12 +183,12 @@ def write_file_direct(path: str, content: str, # If we failed to write to the GPU file directory, try WORKDIR as fallback if output_dir != WORKDIR: logging.warning( - f"Failed to write to {output_dir}, falling back to WORKDIR: {perm_error}" + f"Failed to write to {output_dir}, falling back to WORKDIR: {perm_error}" ) fallback_path = os.path.join(WORKDIR, path) os.makedirs(os.path.dirname(fallback_path), exist_ok=True) - with open(fallback_path, 'w', encoding='utf-8') as f: + with open(fallback_path, "w", encoding="utf-8") as f: f.write(content) tool_context.state["most_recent_file_path"] = fallback_path @@ -208,14 +213,14 @@ def write_file_direct(path: str, content: str, # Framework detection tool (saves framework to state) def save_framework_detection(framework: str, tool_context: ToolContext) -> str: """Save the detected GPU framework to state for use by downstream agents. - - Args: - framework: The detected framework name (e.g., 'CUDA', 'Triton', 'PyTorch CUDA') - tool_context: Context for accessing agent state - - Returns: - Success message confirming the framework was saved - """ + + Args: + framework: The detected framework name (e.g., 'CUDA', 'Triton', 'PyTorch CUDA') + tool_context: Context for accessing agent state + + Returns: + Success message confirming the framework was saved + """ try: tool_context.state["framework_detected"] = framework success_msg = f"Framework '{framework}' saved to state" @@ -235,9 +240,12 @@ def get_available_tools(*tools): return [t for t in tools if t is not None] -def save_path_from_tool_run(tool: BaseTool, args: Dict[str, Any], - tool_context: ToolContext, - tool_response: Optional[Dict]) -> Optional[Dict]: +def save_path_from_tool_run( + tool: BaseTool, + args: Dict[str, Any], + tool_context: ToolContext, + tool_response: Optional[Dict], +) -> Optional[Dict]: """Save file path to state after file operations and log tool execution results.""" # Log tool execution for debugging tool_name = tool.name @@ -251,7 +259,8 @@ def save_path_from_tool_run(tool: BaseTool, args: Dict[str, Any], # If response is a dict with error field, log it elif isinstance(tool_response, dict) and "error" in tool_response: logging.error( - f"Tool '{tool_name}' returned error: {tool_response.get('error')}") + f"Tool '{tool_name}' returned error: {tool_response.get('error')}" + ) else: logging.info(f"Tool '{tool_name}' response: {str(tool_response)[:200]}") @@ -260,30 +269,46 @@ def save_path_from_tool_run(tool: BaseTool, args: Dict[str, Any], tool_context.state["most_recent_file_path"] = file_path # Preserve the original GPU code file path (only set once, when reading GPU files) - if tool.name == "read_file" and file_path and not tool_context.state.get( - "original_gpu_code_path"): + if ( + tool.name == "read_file" + and file_path + and not tool_context.state.get("original_gpu_code_path") + ): # Check if this looks like a GPU source file (not a plan or output file) # Include common GPU source/header extensions: CUDA (.cu, .cuh), HIP (.hip), C/C++ (.c, .cpp, .h, .hpp), Python (.py) - gpu_extensions = ('.cu', '.cuh', '.py', '.cpp', '.c', '.h', '.hpp', - '.hip', '.cc', '.cxx') - if file_path.endswith(gpu_extensions) and 'PLAN' not in file_path.upper( - ) and 'SUMMARY' not in file_path.upper(): + gpu_extensions = ( + ".cu", + ".cuh", + ".py", + ".cpp", + ".c", + ".h", + ".hpp", + ".hip", + ".cc", + ".cxx", + ) + if ( + file_path.endswith(gpu_extensions) + and "PLAN" not in file_path.upper() + and "SUMMARY" not in file_path.upper() + ): # Check if the tool response indicates success (no error) is_success = True if tool_response: # Check for error indicators in the response if isinstance(tool_response, dict): # Check isError field if present - if tool_response.get('isError', False): + if tool_response.get("isError", False): is_success = False # Check if content contains error messages - elif 'content' in tool_response: - content = tool_response['content'] + elif "content" in tool_response: + content = tool_response["content"] if isinstance(content, list) and len(content) > 0: first_content = content[0] - if isinstance(first_content, dict) and 'text' in first_content: - text = first_content['text'] - if text.startswith('Error:'): + if isinstance(first_content, dict) and "text" in first_content: + text = first_content["text"] + if text.startswith("Error:"): is_success = False # Only save the path if the read was successful @@ -300,7 +325,8 @@ def save_path_from_tool_run(tool: BaseTool, args: Dict[str, Any], logging.info(f"Preserved original GPU code path: {abs_file_path}") else: logging.warning( - f"File does not exist, not saving path: {abs_file_path}") + f"File does not exist, not saving path: {abs_file_path}" + ) else: logging.warning(f"Read file failed, not saving path: {file_path}") @@ -310,30 +336,34 @@ def save_path_from_tool_run(tool: BaseTool, args: Dict[str, Any], def save_gpu_code_to_state(callback_context: CallbackContext): """Save GPU code from original GPU source file to state. Only loads once and caches.""" # If gpu_code already exists, we've already loaded it - don't reload - if "gpu_code" in callback_context.state and callback_context.state["gpu_code"]: + if ( + "gpu_code" in callback_context.state and callback_context.state["gpu_code"] + ): logging.info("GPU code already loaded in state, using cached version") return # Use original_gpu_code_path (preserved from initial file read), not most_recent_file_path # most_recent_file_path can change as files are written (e.g., SIMPLIFICATION_PLAN.md) file_path = callback_context.state.get( - "original_gpu_code_path") or callback_context.state.get( - "gpu_code_file_path") + "original_gpu_code_path" + ) or callback_context.state.get("gpu_code_file_path") if file_path: try: with open(file_path, "r") as f: gpu_code = f.read() callback_context.state["gpu_code"] = gpu_code - callback_context.state[ - "gpu_code_file_path"] = file_path # Store path for future reference + callback_context.state["gpu_code_file_path"] = ( + file_path # Store path for future reference + ) logging.info(f"Loaded GPU code from {file_path} and cached in state") except Exception as e: logging.error(f"Failed to read GPU code file: {e}") callback_context.state["gpu_code"] = None else: logging.warning( - "No original GPU code path found in state - cannot load GPU code") + "No original GPU code path found in state - cannot load GPU code" + ) def save_jax_code_to_state(callback_context: CallbackContext): @@ -375,11 +405,11 @@ def save_test_code_to_state(callback_context: CallbackContext): def ensure_summary_state_defaults(callback_context: CallbackContext): """Ensure all state variables needed for summary generation have defaults.""" defaults = { - "framework_detected": "Unknown", - "compilation_results": "Not available", - "syntax_validation_results": "Not available", - "shape_validation_results": "Not available", - "correctness_test_results": "Not available" + "framework_detected": "Unknown", + "compilation_results": "Not available", + "syntax_validation_results": "Not available", + "shape_validation_results": "Not available", + "correctness_test_results": "Not available", } for key, default_value in defaults.items(): if key not in callback_context.state or not callback_context.state[key]: @@ -405,231 +435,223 @@ def load_simplification_plan_from_file(callback_context: CallbackContext): plan_content = f.read() callback_context.state["simplification_plan"] = plan_content logging.info( - f"Loaded simplification plan from {plan_file_path} (includes any user edits)" + f"Loaded simplification plan from {plan_file_path} (includes any user edits)" ) except Exception as e: logging.error(f"Failed to read simplification plan file: {e}") else: logging.info( - f"No existing SIMPLIFICATION_PLAN.md found at {plan_file_path}") + f"No existing SIMPLIFICATION_PLAN.md found at {plan_file_path}" + ) # Step 1: Read GPU code and identify framework identify_framework_agent = CustomLlmAgent( - name="IdentifyFrameworkAgent", - model=MODEL_NAME, - generate_content_config=model_config, - instruction=identify_framework_prompt.PROMPT, - description= - "Reads GPU code file and identifies which GPU framework is being used", - output_key="framework_detected", - tools=get_available_tools(filesystem_tool_rw, save_framework_tool), - after_tool_callback=save_path_from_tool_run, + name="IdentifyFrameworkAgent", + model=MODEL_NAME, + generate_content_config=model_config, + instruction=identify_framework_prompt.PROMPT, + description="Reads GPU code file and identifies which GPU framework is being used", + output_key="framework_detected", + tools=get_available_tools(filesystem_tool_rw, save_framework_tool), + after_tool_callback=save_path_from_tool_run, ) # Step 2: Analyze, plan, and write plan to file (combined) analyze_plan_and_write_agent = CustomLlmAgent( - name="AnalyzePlanAndWriteAgent", - model=MODEL_NAME, - generate_content_config=model_config, - instruction=analyze_plan_prompt.PROMPT, - description= - "Analyzes GPU code, creates simplification plan, and writes it to file", - tools=get_available_tools( - filesystem_tool_rw, - write_file_tool), # Include both read and write tools - before_agent_callback= - load_simplification_plan_from_file, # Load existing plan from file before regenerating + name="AnalyzePlanAndWriteAgent", + model=MODEL_NAME, + generate_content_config=model_config, + instruction=analyze_plan_prompt.PROMPT, + description="Analyzes GPU code, creates simplification plan, and writes it to file", + tools=get_available_tools( + filesystem_tool_rw, write_file_tool + ), # Include both read and write tools + before_agent_callback=load_simplification_plan_from_file, # Load existing plan from file before regenerating ) # Step 3: Execute simplification based on approved plan and write to file organize_gpu_code_agent = CustomLlmAgent( - name="OrganizeGpuCodeAgent", - model=MODEL_NAME, - generate_content_config=model_config, - instruction=simplify_gpu_code_prompt.PROMPT, - description= - "Simplifies GPU code based on the approved plan and writes it to file with appropriate extension", - output_key="organized_code", - tools=get_available_tools(filesystem_tool_rw, write_file_tool), - after_tool_callback=save_path_from_tool_run, + name="OrganizeGpuCodeAgent", + model=MODEL_NAME, + generate_content_config=model_config, + instruction=simplify_gpu_code_prompt.PROMPT, + description="Simplifies GPU code based on the approved plan and writes it to file with appropriate extension", + output_key="organized_code", + tools=get_available_tools(filesystem_tool_rw, write_file_tool), + after_tool_callback=save_path_from_tool_run, ) # Step 4: Write simplification README documenting the process write_simplification_readme_agent = CustomLlmAgent( - name="WriteSimplificationReadmeAgent", - model=MODEL_NAME, - generate_content_config=model_config, - instruction=write_readme_prompt.PROMPT, - description= - "Writes README explaining original code and simplification steps", - tools=get_available_tools(filesystem_tool_rw, write_file_tool), + name="WriteSimplificationReadmeAgent", + model=MODEL_NAME, + generate_content_config=model_config, + instruction=write_readme_prompt.PROMPT, + description="Writes README explaining original code and simplification steps", + tools=get_available_tools(filesystem_tool_rw, write_file_tool), ) # Step 5: Convert to JAX and write to file (combined) convert_to_jax_agent = CustomLlmAgent( - name="ConvertToJaxAgent", - model=MODEL_NAME, - generate_content_config=model_config, - instruction=convert_simplified_to_jax_prompt.PROMPT, - description= - "Converts organized code to JAX and writes it to converted_jax.py", - tools=[write_file_tool], - after_tool_callback=save_path_from_tool_run, + name="ConvertToJaxAgent", + model=MODEL_NAME, + generate_content_config=model_config, + instruction=convert_simplified_to_jax_prompt.PROMPT, + description="Converts organized code to JAX and writes it to converted_jax.py", + tools=[write_file_tool], + after_tool_callback=save_path_from_tool_run, ) # Step 6: Validate syntax (evaluator wrapped as tool) _syntax_checker = JaxSyntaxChecker( - name="check_jax_syntax", - input_key="jax_code", - output_key="syntax_validation_results", + name="check_jax_syntax", + input_key="jax_code", + output_key="syntax_validation_results", ) syntax_checker_tool = AgentTool(agent=_syntax_checker) # Step 6b: Syntax validation with routing logic validate_syntax_agent = CustomLlmAgent( - name="ValidateSyntaxAgent", - model=MODEL_NAME, - generate_content_config=model_config, - tools=[syntax_checker_tool], - instruction=validate_syntax_routing_prompt.PROMPT, - description= - "Validates JAX syntax and routes to fix or compilation based on results", - before_agent_callback= - save_jax_code_to_state, # Load JAX code from file before validation + name="ValidateSyntaxAgent", + model=MODEL_NAME, + generate_content_config=model_config, + tools=[syntax_checker_tool], + instruction=validate_syntax_routing_prompt.PROMPT, + description="Validates JAX syntax and routes to fix or compilation based on results", + before_agent_callback=save_jax_code_to_state, # Load JAX code from file before validation ) # Step 7: Fix conversion errors fix_conversion_agent = CustomLlmAgent( - name="FixConversionAgent", - model=MODEL_NAME, - generate_content_config=model_config, - instruction=fix_conversion_prompt.PROMPT.replace( - "{jax_code}", "{jax_code}").replace("{error_messages}", - "{syntax_validation_results}") + - fix_conversion_extended_prompt.PROMPT, - description= - "Fixes syntax errors in JAX conversion and writes to converted_jax.py", - tools=[write_file_tool], + name="FixConversionAgent", + model=MODEL_NAME, + generate_content_config=model_config, + instruction=fix_conversion_prompt.PROMPT.replace( + "{jax_code}", "{jax_code}" + ).replace("{error_messages}", "{syntax_validation_results}") + + fix_conversion_extended_prompt.PROMPT, + description="Fixes syntax errors in JAX conversion and writes to converted_jax.py", + tools=[write_file_tool], ) # Step 8: Validate compilation (evaluator wrapped as tool) _compilation_checker = JaxCompilationChecker( - name="check_jax_compilation", - input_key="jax_code", - output_key="compilation_results", - auto_manage_servers=True, + name="check_jax_compilation", + input_key="jax_code", + output_key="compilation_results", + auto_manage_servers=True, ) compilation_checker_tool = AgentTool(agent=_compilation_checker) # Step 8b: Compilation validation with routing validate_compilation_agent = CustomLlmAgent( - name="ValidateCompilationAgent", - model=MODEL_NAME, - generate_content_config=model_config, - tools=[compilation_checker_tool], - instruction=validate_compilation_routing_prompt.PROMPT, - description="Validates JAX compilation and proceeds to shape validation", + name="ValidateCompilationAgent", + model=MODEL_NAME, + generate_content_config=model_config, + tools=[compilation_checker_tool], + instruction=validate_compilation_routing_prompt.PROMPT, + description="Validates JAX compilation and proceeds to shape validation", ) # Step 9: Validate shapes (evaluator wrapped as tool) _shape_validator = ShapeValidator( - name="validate_shapes", - input_key="jax_code", - output_key="shape_validation_results", + name="validate_shapes", + input_key="jax_code", + output_key="shape_validation_results", ) shape_validator_tool = AgentTool(agent=_shape_validator) # Step 9b: Shape validation with routing validate_shapes_agent = CustomLlmAgent( - name="ValidateShapesAgent", - model=MODEL_NAME, - generate_content_config=model_config, - tools=[shape_validator_tool], - instruction=validate_shapes_routing_prompt.PROMPT, - description="Validates tensor shapes and proceeds to test generation", + name="ValidateShapesAgent", + model=MODEL_NAME, + generate_content_config=model_config, + tools=[shape_validator_tool], + instruction=validate_shapes_routing_prompt.PROMPT, + description="Validates tensor shapes and proceeds to test generation", ) # Step 10: Generate correctness test and write to file (combined) generate_correctness_test_agent = CustomLlmAgent( - name="GenerateCorrectnessTestAgent", - model=MODEL_NAME, - generate_content_config=model_config, - instruction=generate_test_prompt.PROMPT + "\n" + - generate_test_extended_prompt.PROMPT, - description= - "Generates a validation test for JAX code and writes it to test_correctness.py", - # Note: No output_key - test is written to file and read by RunCorrectnessTestAgent from file - tools=[write_file_tool], - after_tool_callback=save_path_from_tool_run, + name="GenerateCorrectnessTestAgent", + model=MODEL_NAME, + generate_content_config=model_config, + instruction=generate_test_prompt.PROMPT + + "\n" + + generate_test_extended_prompt.PROMPT, + description="Generates a validation test for JAX code and writes it to test_correctness.py", + # Note: No output_key - test is written to file and read by RunCorrectnessTestAgent from file + tools=[write_file_tool], + after_tool_callback=save_path_from_tool_run, ) # Step 11: Run correctness test (evaluator wrapped as tool) _correctness_checker = JaxCorrectnessChecker( - name="run_correctness_test", - input_key="correctness_test_code", - output_key="correctness_test_results", - auto_manage_servers=True, + name="run_correctness_test", + input_key="correctness_test_code", + output_key="correctness_test_results", + auto_manage_servers=True, ) correctness_checker_tool = AgentTool(agent=_correctness_checker) # Step 11b: Run test with routing run_correctness_test_agent = CustomLlmAgent( - name="RunCorrectnessTestAgent", - model=MODEL_NAME, - generate_content_config=model_config, - tools=[correctness_checker_tool], - instruction=run_test_routing_prompt.PROMPT, - description= - "Loads test from file, runs correctness check, and proceeds to summary generation", - before_agent_callback= - save_test_code_to_state, # Load test file into state before running + name="RunCorrectnessTestAgent", + model=MODEL_NAME, + generate_content_config=model_config, + tools=[correctness_checker_tool], + instruction=run_test_routing_prompt.PROMPT, + description="Loads test from file, runs correctness check, and proceeds to summary generation", + before_agent_callback=save_test_code_to_state, # Load test file into state before running ) # Step 12: Generate summary and write to file (combined) generate_and_write_summary_agent = CustomLlmAgent( - name="GenerateAndWriteSummaryAgent", - model=MODEL_NAME, - generate_content_config=model_config, - instruction=summary_prompt.PROMPT.replace( - "{framework_detected}", "{framework_detected}").replace( - "{conversion_status}", "Success").replace( - "{test_results}", """Compilation: {compilation_results} + name="GenerateAndWriteSummaryAgent", + model=MODEL_NAME, + generate_content_config=model_config, + instruction=summary_prompt.PROMPT.replace( + "{framework_detected}", "{framework_detected}" + ) + .replace("{conversion_status}", "Success") + .replace( + "{test_results}", + """Compilation: {compilation_results} Syntax Validation: {syntax_validation_results} Shape Validation: {shape_validation_results} -Numerical Correctness: {correctness_test_results}""") + - generate_summary_extended_prompt.PROMPT, - description= - "Generates conversion summary and writes it to CONVERSION_SUMMARY.md", - output_key="conversion_summary", - tools=[write_file_tool], - before_agent_callback=ensure_summary_state_defaults, +Numerical Correctness: {correctness_test_results}""", + ) + + generate_summary_extended_prompt.PROMPT, + description="Generates conversion summary and writes it to CONVERSION_SUMMARY.md", + output_key="conversion_summary", + tools=[write_file_tool], + before_agent_callback=ensure_summary_state_defaults, ) # Main GPU-to-JAX orchestrator agent gpu_to_jax_agent = CustomLlmAgent( - name="GpuToJaxAgent", - model=MODEL_NAME, - generate_content_config=model_config, - instruction=orchestrator_prompt.PROMPT, - description= - "Routes user requests to appropriate GPU-to-JAX conversion workflow phase", - sub_agents=[ - # Framework identification and planning - identify_framework_agent, - analyze_plan_and_write_agent, - # Simplification execution - organize_gpu_code_agent, - write_simplification_readme_agent, - # JAX conversion and validation - convert_to_jax_agent, - validate_syntax_agent, - fix_conversion_agent, - validate_compilation_agent, - validate_shapes_agent, - # Testing and summary - generate_correctness_test_agent, - run_correctness_test_agent, - generate_and_write_summary_agent, - ], + name="GpuToJaxAgent", + model=MODEL_NAME, + generate_content_config=model_config, + instruction=orchestrator_prompt.PROMPT, + description="Routes user requests to appropriate GPU-to-JAX conversion workflow phase", + sub_agents=[ + # Framework identification and planning + identify_framework_agent, + analyze_plan_and_write_agent, + # Simplification execution + organize_gpu_code_agent, + write_simplification_readme_agent, + # JAX conversion and validation + convert_to_jax_agent, + validate_syntax_agent, + fix_conversion_agent, + validate_compilation_agent, + validate_shapes_agent, + # Testing and summary + generate_correctness_test_agent, + run_correctness_test_agent, + generate_and_write_summary_agent, + ], ) diff --git a/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/evaluators/__init__.py b/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/evaluators/__init__.py index 1916697..1bb5749 100644 --- a/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/evaluators/__init__.py +++ b/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/evaluators/__init__.py @@ -1,17 +1,21 @@ """Evaluators for GPU to JAX conversion agent.""" -from hitl_agent.subagents.gpu_to_jax_agent.evaluators.jax_syntax_checker import ( - JaxSyntaxChecker,) -from hitl_agent.subagents.gpu_to_jax_agent.evaluators.shape_validator import ( - ShapeValidator,) from hitl_agent.subagents.gpu_to_jax_agent.evaluators.compilation_checker import ( - JaxCompilationChecker,) + JaxCompilationChecker, +) from hitl_agent.subagents.gpu_to_jax_agent.evaluators.correctness_checker import ( - JaxCorrectnessChecker,) + JaxCorrectnessChecker, +) +from hitl_agent.subagents.gpu_to_jax_agent.evaluators.jax_syntax_checker import ( + JaxSyntaxChecker, +) +from hitl_agent.subagents.gpu_to_jax_agent.evaluators.shape_validator import ( + ShapeValidator, +) __all__ = [ - "JaxSyntaxChecker", - "ShapeValidator", - "JaxCompilationChecker", - "JaxCorrectnessChecker", + "JaxSyntaxChecker", + "ShapeValidator", + "JaxCompilationChecker", + "JaxCorrectnessChecker", ] diff --git a/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/evaluators/compilation_checker.py b/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/evaluators/compilation_checker.py index de88bfa..b29a2fa 100644 --- a/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/evaluators/compilation_checker.py +++ b/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/evaluators/compilation_checker.py @@ -7,12 +7,13 @@ from google.adk.agents import BaseAgent from google.adk.agents.invocation_context import InvocationContext from google.adk.events import Event, EventActions + +from hitl_agent.server_utils.server_manager_mixin import ServerManagerMixin from hitl_agent.subagents.gpu_to_jax_agent.constants import ( - EVAL_SERVER_PORT, - CONVERSION_TIMEOUT, - PREFERRED_BACKEND, + CONVERSION_TIMEOUT, + EVAL_SERVER_PORT, + PREFERRED_BACKEND, ) -from hitl_agent.server_utils.server_manager_mixin import ServerManagerMixin class JaxCompilationChecker(ServerManagerMixin, BaseAgent): @@ -23,11 +24,11 @@ class JaxCompilationChecker(ServerManagerMixin, BaseAgent): auto_manage_servers: bool = False def __init__( - self, - name: str, - input_key: str, - output_key: str, - auto_manage_servers: bool = False, + self, + name: str, + input_key: str, + output_key: str, + auto_manage_servers: bool = False, ): super().__init__(name=name) self.input_key = input_key @@ -36,7 +37,8 @@ def __init__( self._servers_started = [] # Track which servers this instance started async def _run_async_impl( - self, ctx: InvocationContext) -> AsyncGenerator[Event, None]: + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: # Ensure servers are running if auto_manage_servers is True await self._ensure_servers_running() @@ -44,24 +46,25 @@ async def _run_async_impl( if not code: logging.warning(f"[{self.name}] No {self.input_key} found in context") yield Event( - author=self.name, - actions=EventActions(state_delta={self.output_key: None}), + author=self.name, + actions=EventActions(state_delta={self.output_key: None}), ) return try: # Call the eval server to compile and run the code logging.info(f"[{self.name}] Compiling JAX code") - async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout( - total=CONVERSION_TIMEOUT)) as session: + async with aiohttp.ClientSession( + timeout=aiohttp.ClientTimeout(total=CONVERSION_TIMEOUT) + ) as session: async with session.post( - f"http://localhost:{EVAL_SERVER_PORT}/evaluate", - json={ - "eval_type": "compilation_test", - "code": code, - "timeout": CONVERSION_TIMEOUT, - "backend_type": PREFERRED_BACKEND, - }, + f"http://localhost:{EVAL_SERVER_PORT}/evaluate", + json={ + "eval_type": "compilation_test", + "code": code, + "timeout": CONVERSION_TIMEOUT, + "backend_type": PREFERRED_BACKEND, + }, ) as response: if response.status == 200: result = await response.json() @@ -69,19 +72,20 @@ async def _run_async_impl( if result["exit_code"] == 0: logging.info(f"[{self.name}] JAX code compilation successful.") yield Event( - author=self.name, - actions=EventActions( - state_delta={self.output_key: "Success"}), + author=self.name, + actions=EventActions(state_delta={self.output_key: "Success"}), ) - elif (result["error"] is None and result["output"] == "" and - result["exit_code"] == 1): + elif ( + result["error"] is None + and result["output"] == "" + and result["exit_code"] == 1 + ): logging.info( - f"[{self.name}] Code execution had exit code 1, but no error, indicating success." + f"[{self.name}] Code execution had exit code 1, but no error, indicating success." ) yield Event( - author=self.name, - actions=EventActions( - state_delta={self.output_key: "Success"}), + author=self.name, + actions=EventActions(state_delta={self.output_key: "Success"}), ) else: logging.error(f"[{self.name}] JAX code compilation failed.") @@ -92,37 +96,35 @@ async def _run_async_impl( full_error += f"\n\nOutput:\n{output_msg}" yield Event( - author=self.name, - actions=EventActions( - state_delta={self.output_key: full_error}), + author=self.name, + actions=EventActions(state_delta={self.output_key: full_error}), ) else: error_detail = await response.text() logging.error( - f"[{self.name}] HTTP error {response.status}: {error_detail}") + f"[{self.name}] HTTP error {response.status}: {error_detail}" + ) yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: - f"HTTP error {response.status}: {error_detail}" - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: f"HTTP error {response.status}: {error_detail}" + } + ), ) except aiohttp.ClientConnectorError: - error_msg = ( - f"Cannot connect to evaluation server at localhost:{EVAL_SERVER_PORT}. " - "Make sure the eval server is running.") + error_msg = f"Cannot connect to evaluation server at localhost:{EVAL_SERVER_PORT}. Make sure the eval server is running." logging.error(f"[{self.name}] {error_msg}") yield Event( - author=self.name, - actions=EventActions(state_delta={self.output_key: error_msg}), + author=self.name, + actions=EventActions(state_delta={self.output_key: error_msg}), ) except Exception as e: error_msg = f"Exception during compilation check: {str(e)}" logging.error(f"[{self.name}] {error_msg}") yield Event( - author=self.name, - actions=EventActions(state_delta={self.output_key: error_msg}), + author=self.name, + actions=EventActions(state_delta={self.output_key: error_msg}), ) finally: await self._cleanup_servers() diff --git a/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/evaluators/correctness_checker.py b/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/evaluators/correctness_checker.py index 9648d5a..be96f1e 100644 --- a/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/evaluators/correctness_checker.py +++ b/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/evaluators/correctness_checker.py @@ -7,13 +7,14 @@ from google.adk.agents import BaseAgent from google.adk.agents.invocation_context import InvocationContext from google.adk.events import Event, EventActions + +from hitl_agent.server_utils.server_manager_mixin import ServerManagerMixin from hitl_agent.subagents.gpu_to_jax_agent.constants import ( - EVAL_SERVER_PORT, - CONVERSION_TIMEOUT, - NUMERICAL_TOLERANCE, - PREFERRED_BACKEND, + CONVERSION_TIMEOUT, + EVAL_SERVER_PORT, + NUMERICAL_TOLERANCE, + PREFERRED_BACKEND, ) -from hitl_agent.server_utils.server_manager_mixin import ServerManagerMixin class JaxCorrectnessChecker(ServerManagerMixin, BaseAgent): @@ -23,11 +24,13 @@ class JaxCorrectnessChecker(ServerManagerMixin, BaseAgent): output_key: Optional[str] = None auto_manage_servers: bool = False - def __init__(self, - name: str, - input_key: str, - output_key: str, - auto_manage_servers: bool = False): + def __init__( + self, + name: str, + input_key: str, + output_key: str, + auto_manage_servers: bool = False, + ): super().__init__(name=name) self.input_key = input_key self.output_key = output_key @@ -35,7 +38,8 @@ def __init__(self, self._servers_started = [] # Track which servers this instance started async def _run_async_impl( - self, ctx: InvocationContext) -> AsyncGenerator[Event, None]: + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: # Ensure servers are running if auto_manage_servers is True await self._ensure_servers_running() @@ -43,26 +47,29 @@ async def _run_async_impl( if not test_code: logging.warning(f"[{self.name}] No {self.input_key} found in context") yield Event( - author=self.name, - actions=EventActions(state_delta={ - self.output_key: "No test code provided for correctness check" - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: "No test code provided for correctness check" + } + ), ) return try: # Call the eval server to run the correctness test logging.info(f"[{self.name}] Running correctness test") - async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout( - total=CONVERSION_TIMEOUT)) as session: + async with aiohttp.ClientSession( + timeout=aiohttp.ClientTimeout(total=CONVERSION_TIMEOUT) + ) as session: async with session.post( - f"http://localhost:{EVAL_SERVER_PORT}/evaluate", - json={ - "eval_type": "correctness_test", - "code": test_code, - "timeout": CONVERSION_TIMEOUT, - "backend_type": PREFERRED_BACKEND, - }, + f"http://localhost:{EVAL_SERVER_PORT}/evaluate", + json={ + "eval_type": "correctness_test", + "code": test_code, + "timeout": CONVERSION_TIMEOUT, + "backend_type": PREFERRED_BACKEND, + }, ) as response: if response.status == 200: result = await response.json() @@ -74,45 +81,49 @@ async def _run_async_impl( # Check for success indicators if "Identical" in output or "PASSED" in output.upper(): logging.info( - f"[{self.name}] Correctness test passed - outputs are identical" + f"[{self.name}] Correctness test passed - outputs are identical" ) yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: - f"Success: Numerical outputs match within tolerance ({NUMERICAL_TOLERANCE})" - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: f"Success: Numerical outputs match within tolerance ({NUMERICAL_TOLERANCE})" + } + ), ) elif "Different" in output or "FAILED" in output.upper(): logging.warning( - f"[{self.name}] Correctness test failed - outputs differ") + f"[{self.name}] Correctness test failed - outputs differ" + ) error_msg = ( - f"Correctness test failed: Outputs are not identical.\n" - f"Expected tolerance: {NUMERICAL_TOLERANCE}\n" - f"Details: {output}") + f"Correctness test failed: Outputs are not identical.\n" + f"Expected tolerance: {NUMERICAL_TOLERANCE}\n" + f"Details: {output}" + ) yield Event( - author=self.name, - actions=EventActions( - state_delta={self.output_key: error_msg}), + author=self.name, + actions=EventActions(state_delta={self.output_key: error_msg}), ) elif result["exit_code"] == 0 and not error: # Success with no specific marker logging.info( - f"[{self.name}] Correctness test passed (exit code 0)") - success_msg = f"Success: Test executed without errors" + f"[{self.name}] Correctness test passed (exit code 0)" + ) + success_msg = "Success: Test executed without errors" if output: success_msg += f"\nOutput: {output}" yield Event( - author=self.name, - actions=EventActions( - state_delta={self.output_key: success_msg}), + author=self.name, + actions=EventActions( + state_delta={self.output_key: success_msg} + ), ) else: # Error case logging.error( - f"[{self.name}] Correctness test encountered an error") - error_msg = f"Correctness test failed with errors:\n" + f"[{self.name}] Correctness test encountered an error" + ) + error_msg = "Correctness test failed with errors:\n" if error: error_msg += f"Error: {error}\n" if output: @@ -120,37 +131,35 @@ async def _run_async_impl( error_msg += f"Exit code: {result['exit_code']}" yield Event( - author=self.name, - actions=EventActions( - state_delta={self.output_key: error_msg}), + author=self.name, + actions=EventActions(state_delta={self.output_key: error_msg}), ) else: error_detail = await response.text() logging.error( - f"[{self.name}] HTTP error {response.status}: {error_detail}") + f"[{self.name}] HTTP error {response.status}: {error_detail}" + ) yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: - f"HTTP error {response.status}: {error_detail}" - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: f"HTTP error {response.status}: {error_detail}" + } + ), ) except aiohttp.ClientConnectorError: - error_msg = ( - f"Cannot connect to evaluation server at localhost:{EVAL_SERVER_PORT}. " - "Make sure the eval server is running.") + error_msg = f"Cannot connect to evaluation server at localhost:{EVAL_SERVER_PORT}. Make sure the eval server is running." logging.error(f"[{self.name}] {error_msg}") yield Event( - author=self.name, - actions=EventActions(state_delta={self.output_key: error_msg}), + author=self.name, + actions=EventActions(state_delta={self.output_key: error_msg}), ) except Exception as e: error_msg = f"Exception during correctness check: {str(e)}" logging.error(f"[{self.name}] {error_msg}") yield Event( - author=self.name, - actions=EventActions(state_delta={self.output_key: error_msg}), + author=self.name, + actions=EventActions(state_delta={self.output_key: error_msg}), ) finally: await self._cleanup_servers() diff --git a/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/evaluators/jax_syntax_checker.py b/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/evaluators/jax_syntax_checker.py index e071184..4c125f6 100644 --- a/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/evaluators/jax_syntax_checker.py +++ b/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/evaluators/jax_syntax_checker.py @@ -21,14 +21,16 @@ def __init__(self, name: str, input_key: str, output_key: str): self.output_key = output_key async def _run_async_impl( - self, ctx: InvocationContext) -> AsyncGenerator[Event, None]: + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: code = ctx.session.state.get(self.input_key, "") if not code: logging.warning(f"[{self.name}] No {self.input_key} found in context") yield Event( - author=self.name, - actions=EventActions( - state_delta={self.output_key: "No code to validate"}), + author=self.name, + actions=EventActions( + state_delta={self.output_key: "No code to validate"} + ), ) return @@ -44,8 +46,8 @@ async def _run_async_impl( logging.error(f"[{self.name}] {error_msg}") errors.append(error_msg) yield Event( - author=self.name, - actions=EventActions(state_delta={self.output_key: error_msg}), + author=self.name, + actions=EventActions(state_delta={self.output_key: error_msg}), ) return @@ -64,12 +66,14 @@ async def _run_async_impl( gpu_remnants.append(".to(device) call found (should be removed in JAX)") if "torch.Tensor" in code: gpu_remnants.append( - "torch.Tensor reference found (should be jnp.ndarray)") + "torch.Tensor reference found (should be jnp.ndarray)" + ) if "<<<" in code or ">>>" in code: gpu_remnants.append("CUDA kernel launch syntax found (<<< >>>)") if "__global__" in code or "__device__" in code: gpu_remnants.append( - "CUDA kernel decorators found (__global__, __device__)") + "CUDA kernel decorators found (__global__, __device__)" + ) if "@triton.jit" in code: gpu_remnants.append("Triton decorator found (@triton.jit)") @@ -80,17 +84,18 @@ async def _run_async_impl( if "random.normal" in code or "random.uniform" in code: if "random.PRNGKey" not in code and "PRNGKey" not in code: warnings.append( - "Using JAX random functions but no PRNGKey initialization found") + "Using JAX random functions but no PRNGKey initialization found" + ) # Check 5: Common API mismatches if " dim=" in code: warnings.append( - "Found 'dim=' parameter - JAX uses 'axis=' instead (PyTorch uses 'dim=')" + "Found 'dim=' parameter - JAX uses 'axis=' instead (PyTorch uses 'dim=')" ) if ".numpy()" in code or ".item()" in code: errors.append( - "Found .numpy() or .item() method call - JAX arrays don't have these methods" + "Found .numpy() or .item() method call - JAX arrays don't have these methods" ) # Check 6: Verify three-section structure @@ -111,8 +116,8 @@ async def _run_async_impl( logging.error(f"[{self.name}] {error_msg}") yield Event( - author=self.name, - actions=EventActions(state_delta={self.output_key: error_msg}), + author=self.name, + actions=EventActions(state_delta={self.output_key: error_msg}), ) else: success_msg = "JAX Syntax Validation Passed" @@ -123,6 +128,6 @@ async def _run_async_impl( logging.info(f"[{self.name}] {success_msg}") yield Event( - author=self.name, - actions=EventActions(state_delta={self.output_key: success_msg}), + author=self.name, + actions=EventActions(state_delta={self.output_key: success_msg}), ) diff --git a/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/evaluators/shape_validator.py b/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/evaluators/shape_validator.py index efc7b20..4c0a176 100644 --- a/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/evaluators/shape_validator.py +++ b/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/evaluators/shape_validator.py @@ -27,14 +27,14 @@ def extract_shapes_from_code(self, code: str) -> Dict[str, List[Tuple]]: # Pattern 1: Shape from random/zeros/ones calls: (N, M, K) shape_patterns = [ - r"random\.normal\([^,]+,\s*\(([^)]+)\)\)", - r"random\.uniform\([^,]+,\s*\(([^)]+)\)\)", - r"jnp\.zeros\(\(([^)]+)\)\)", - r"jnp\.ones\(\(([^)]+)\)\)", - r"np\.zeros\(\(([^)]+)\)\)", - r"np\.ones\(\(([^)]+)\)\)", - r"torch\.randn\(\(([^)]+)\)\)", - r"torch\.zeros\(\(([^)]+)\)\)", + r"random\.normal\([^,]+,\s*\(([^)]+)\)\)", + r"random\.uniform\([^,]+,\s*\(([^)]+)\)\)", + r"jnp\.zeros\(\(([^)]+)\)\)", + r"jnp\.ones\(\(([^)]+)\)\)", + r"np\.zeros\(\(([^)]+)\)\)", + r"np\.ones\(\(([^)]+)\)\)", + r"torch\.randn\(\(([^)]+)\)\)", + r"torch\.zeros\(\(([^)]+)\)\)", ] for pattern in shape_patterns: @@ -63,7 +63,8 @@ def extract_shapes_from_code(self, code: str) -> Dict[str, List[Tuple]]: return shapes def validate_shape_consistency( - self, code: str) -> Tuple[bool, List[str], List[str]]: + self, code: str + ) -> Tuple[bool, List[str], List[str]]: """Validate shape consistency in the code.""" errors = [] warnings = [] @@ -73,7 +74,7 @@ def validate_shape_consistency( # Check if we found any shape information if not shapes["inputs"]: warnings.append( - "Could not extract shape information from code. Manual verification recommended." + "Could not extract shape information from code. Manual verification recommended." ) return True, errors, warnings @@ -85,20 +86,20 @@ def validate_shape_consistency( # Look for common mistakes like (N, M) @ (N, K) instead of (N, M) @ (M, K) if "@" in line: warnings.append( - f"Line {i}: Found @ operator - verify matrix dimensions are compatible" + f"Line {i}: Found @ operator - verify matrix dimensions are compatible" ) # Check for reshape operations that might change dimensions if "reshape" in line or "view" in line: warnings.append( - f"Line {i}: Found reshape operation - verify shape transformation is correct" + f"Line {i}: Found reshape operation - verify shape transformation is correct" ) # Check for reduce operations that change dimensions if any(op in line for op in ["sum(", "mean(", "max(", "min(", "reduce("]): if "axis=" not in line and "dim=" not in line: warnings.append( - f"Line {i}: Found reduction operation without explicit axis - output shape may differ from expected" + f"Line {i}: Found reduction operation without explicit axis - output shape may differ from expected" ) # Verify computation function signature @@ -108,16 +109,18 @@ def validate_shape_consistency( if isinstance(node, ast.FunctionDef) and node.name == "computation": # Check if function has type hints has_type_hints = any( - arg.annotation is not None for arg in node.args.args) + arg.annotation is not None for arg in node.args.args + ) if not has_type_hints: warnings.append( - "computation() function missing type hints - add jnp.ndarray annotations for clarity" + "computation() function missing type hints - add jnp.ndarray annotations for clarity" ) # Check return type if node.returns is None: warnings.append( - "computation() function missing return type annotation") + "computation() function missing return type annotation" + ) except: pass @@ -125,14 +128,16 @@ def validate_shape_consistency( return is_valid, errors, warnings async def _run_async_impl( - self, ctx: InvocationContext) -> AsyncGenerator[Event, None]: + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: code = ctx.session.state.get(self.input_key, "") if not code: logging.warning(f"[{self.name}] No {self.input_key} found in context") yield Event( - author=self.name, - actions=EventActions( - state_delta={self.output_key: "No code to validate"}), + author=self.name, + actions=EventActions( + state_delta={self.output_key: "No code to validate"} + ), ) return @@ -150,8 +155,8 @@ async def _run_async_impl( logging.error(f"[{self.name}] {error_msg}") yield Event( - author=self.name, - actions=EventActions(state_delta={self.output_key: error_msg}), + author=self.name, + actions=EventActions(state_delta={self.output_key: error_msg}), ) else: success_msg = "Shape Validation Passed" @@ -165,14 +170,14 @@ async def _run_async_impl( logging.info(f"[{self.name}] {success_msg}") yield Event( - author=self.name, - actions=EventActions(state_delta={self.output_key: success_msg}), + author=self.name, + actions=EventActions(state_delta={self.output_key: success_msg}), ) except Exception as e: error_msg = f"Exception during shape validation: {str(e)}" logging.error(f"[{self.name}] {error_msg}") yield Event( - author=self.name, - actions=EventActions(state_delta={self.output_key: error_msg}), + author=self.name, + actions=EventActions(state_delta={self.output_key: error_msg}), ) diff --git a/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/prompts/__init__.py b/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/prompts/__init__.py index 00cc0d6..5fb0f88 100644 --- a/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/prompts/__init__.py +++ b/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/prompts/__init__.py @@ -1,44 +1,44 @@ """Prompts for GPU to JAX conversion agent.""" from . import ( - identify_framework_prompt, - organize_gpu_code_prompt, - simplify_gpu_code_prompt, - convert_to_jax_prompt, - convert_simplified_to_jax_prompt, - fix_conversion_prompt, - fix_conversion_extended_prompt, - summary_prompt, - generate_summary_extended_prompt, - orchestrator_prompt, - analyze_plan_prompt, - write_readme_prompt, - validate_syntax_routing_prompt, - validate_compilation_routing_prompt, - validate_shapes_routing_prompt, - generate_test_prompt, - generate_test_extended_prompt, - run_test_routing_prompt, + analyze_plan_prompt, + convert_simplified_to_jax_prompt, + convert_to_jax_prompt, + fix_conversion_extended_prompt, + fix_conversion_prompt, + generate_summary_extended_prompt, + generate_test_extended_prompt, + generate_test_prompt, + identify_framework_prompt, + orchestrator_prompt, + organize_gpu_code_prompt, + run_test_routing_prompt, + simplify_gpu_code_prompt, + summary_prompt, + validate_compilation_routing_prompt, + validate_shapes_routing_prompt, + validate_syntax_routing_prompt, + write_readme_prompt, ) __all__ = [ - "identify_framework_prompt", - "organize_gpu_code_prompt", - "simplify_gpu_code_prompt", - "convert_to_jax_prompt", - "convert_simplified_to_jax_prompt", - "fix_conversion_prompt", - "fix_conversion_extended_prompt", - "summary_prompt", - "generate_summary_extended_prompt", - "orchestrator_prompt", - "analyze_plan_prompt", - "identify_framework_prompt", - "write_readme_prompt", - "validate_syntax_routing_prompt", - "validate_compilation_routing_prompt", - "validate_shapes_routing_prompt", - "generate_test_prompt", - "generate_test_extended_prompt", - "run_test_routing_prompt", + "identify_framework_prompt", + "organize_gpu_code_prompt", + "simplify_gpu_code_prompt", + "convert_to_jax_prompt", + "convert_simplified_to_jax_prompt", + "fix_conversion_prompt", + "fix_conversion_extended_prompt", + "summary_prompt", + "generate_summary_extended_prompt", + "orchestrator_prompt", + "analyze_plan_prompt", + "identify_framework_prompt", + "write_readme_prompt", + "validate_syntax_routing_prompt", + "validate_compilation_routing_prompt", + "validate_shapes_routing_prompt", + "generate_test_prompt", + "generate_test_extended_prompt", + "run_test_routing_prompt", ] diff --git a/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/test_agent.py b/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/test_agent.py index 77fcc75..b58afbd 100644 --- a/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/test_agent.py +++ b/MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/test_agent.py @@ -1,11 +1,10 @@ """Unit tests for GPU to JAX conversion agent.""" import pytest -from unittest.mock import patch from hitl_agent.subagents.gpu_to_jax_agent.evaluators import ( - JaxSyntaxChecker, - ShapeValidator, + JaxSyntaxChecker, + ShapeValidator, ) @@ -31,9 +30,9 @@ def computation(A: jnp.ndarray) -> jnp.ndarray: result = jax.block_until_ready(computation(A)) """ - checker = JaxSyntaxChecker(name="TestChecker", - input_key="test_code", - output_key="test_results") + checker = JaxSyntaxChecker( + name="TestChecker", input_key="test_code", output_key="test_results" + ) # Note: This is a synchronous test - in practice, you'd need to run the async method # For now, we just verify the checker can be instantiated @@ -50,9 +49,9 @@ def test_invalid_python_syntax(self): def broken( return None """ - checker = JaxSyntaxChecker(name="TestChecker", - input_key="test_code", - output_key="test_results") + checker = JaxSyntaxChecker( + name="TestChecker", input_key="test_code", output_key="test_results" + ) assert checker is not None @@ -67,9 +66,9 @@ def test_extract_shapes_from_code(self): A = random.normal(key, (1024, 512)) B = random.normal(key, (512, 256)) """ - validator = ShapeValidator(name="TestValidator", - input_key="test_code", - output_key="test_results") + validator = ShapeValidator( + name="TestValidator", input_key="test_code", output_key="test_results" + ) shapes = validator.extract_shapes_from_code(code) assert "inputs" in shapes @@ -83,9 +82,9 @@ def test_validate_shape_consistency(self): def computation(A: jnp.ndarray, B: jnp.ndarray) -> jnp.ndarray: return jnp.matmul(A, B) """ - validator = ShapeValidator(name="TestValidator", - input_key="test_code", - output_key="test_results") + validator = ShapeValidator( + name="TestValidator", input_key="test_code", output_key="test_results" + ) is_valid, errors, warnings = validator.validate_shape_consistency(code) assert isinstance(is_valid, bool) @@ -100,7 +99,8 @@ class TestGpuToJaxAgent: async def test_agent_initialization(self): """Test that the agent can be initialized.""" from hitl_agent.subagents.gpu_to_jax_agent.agent import ( - gpu_to_jax_conversion_agent,) + gpu_to_jax_conversion_agent, + ) assert gpu_to_jax_conversion_agent is not None assert gpu_to_jax_conversion_agent.name == "GpuToJaxConversionAgent" diff --git a/MaxKernel/hitl_agent/subagents/kernel_writing/__init__.py b/MaxKernel/hitl_agent/subagents/kernel_writing/__init__.py index 1cf70be..226ccb0 100644 --- a/MaxKernel/hitl_agent/subagents/kernel_writing/__init__.py +++ b/MaxKernel/hitl_agent/subagents/kernel_writing/__init__.py @@ -1,15 +1,15 @@ """Kernel writing subagent module.""" from .agent import ( - KernelCompilationValidationLoop, - plan_kernel_agent, - implement_kernel_agent, - validate_kernel_compilation_agent, + KernelCompilationValidationLoop, + implement_kernel_agent, + plan_kernel_agent, + validate_kernel_compilation_agent, ) __all__ = [ - 'KernelCompilationValidationLoop', - 'plan_kernel_agent', - 'implement_kernel_agent', - 'validate_kernel_compilation_agent', + "KernelCompilationValidationLoop", + "plan_kernel_agent", + "implement_kernel_agent", + "validate_kernel_compilation_agent", ] diff --git a/MaxKernel/hitl_agent/subagents/kernel_writing/agent.py b/MaxKernel/hitl_agent/subagents/kernel_writing/agent.py index 78edf42..33cb8f9 100644 --- a/MaxKernel/hitl_agent/subagents/kernel_writing/agent.py +++ b/MaxKernel/hitl_agent/subagents/kernel_writing/agent.py @@ -1,37 +1,38 @@ """Kernel writing and compilation validation agents.""" -import os import logging -from typing import Optional, AsyncGenerator +import os +from typing import AsyncGenerator, Optional -from google.adk.agents import SequentialAgent, BaseAgent +from google.adk.agents import BaseAgent, SequentialAgent from google.adk.agents.invocation_context import InvocationContext from google.adk.events import Event, EventActions -from hitl_agent.custom_types import CustomLlmAgent +from hitl_agent.callbacks import ( + create_path_saver, + extract_fix_summary, + load_kernel_and_plan_to_state, + load_single_kernel_to_state, + save_kernel_and_plan_paths, +) +from hitl_agent.config import model_config, thinking_planner from hitl_agent.constants import MODEL_NAME +from hitl_agent.custom_types import CustomLlmAgent from hitl_agent.subagents.kernel_writing.kernel_compilation import ( - KernelCompilationChecker,) -from hitl_agent.config import model_config, thinking_planner -from hitl_agent.tools.search_api_tool import search_api_tool -from hitl_agent.tools.tools import filesystem_tool_rw, vertex_ai_rag_tool -from hitl_agent.callbacks import ( - create_path_saver, - save_kernel_and_plan_paths, - load_single_kernel_to_state, - load_kernel_and_plan_to_state, - extract_fix_summary, + KernelCompilationChecker, ) from hitl_agent.subagents.kernel_writing.prompts import ( - kernel_planning_prompt, - kernel_implementation_prompt, - ask_validation_prompt, - fix_kernel_compilation, - kernel_compilation_summary, - add_debug_statements, - cleanup_debug_statements, - read_file_prompt, + add_debug_statements, + ask_validation_prompt, + cleanup_debug_statements, + fix_kernel_compilation, + kernel_compilation_summary, + kernel_implementation_prompt, + kernel_planning_prompt, + read_file_prompt, ) +from hitl_agent.tools.search_api_tool import search_api_tool +from hitl_agent.tools.tools import filesystem_tool_rw, vertex_ai_rag_tool class KernelCompilationValidationLoop(BaseAgent): @@ -43,23 +44,24 @@ class KernelCompilationValidationLoop(BaseAgent): max_retries: int = 4 def __init__( - self, - name: str, - compilation_checker: BaseAgent, - fix_agent: BaseAgent, - debug_agent: Optional[BaseAgent] = None, - max_retries: int = 4, + self, + name: str, + compilation_checker: BaseAgent, + fix_agent: BaseAgent, + debug_agent: Optional[BaseAgent] = None, + max_retries: int = 4, ): super().__init__( - name=name, - compilation_checker=compilation_checker, - fix_agent=fix_agent, - debug_agent=debug_agent, - max_retries=max_retries, + name=name, + compilation_checker=compilation_checker, + fix_agent=fix_agent, + debug_agent=debug_agent, + max_retries=max_retries, ) async def _run_async_impl( - self, ctx: InvocationContext) -> AsyncGenerator[Event, None]: + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: """Validation loop: compile -> fix -> repeat until valid or max retries.""" # Check if a kernel file was actually generated @@ -71,23 +73,20 @@ async def _run_async_impl( if not kernel_path: logging.error( - f"[{self.name}] No kernel file path found in state. Kernel implementation may have failed." + f"[{self.name}] No kernel file path found in state. Kernel implementation may have failed." ) yield Event( - author=self.name, - actions=EventActions( - state_delta={ - "kernel_compilation_status": { - "success": - False, - "retries": - 0, - "message": - "No kernel file was generated. Cannot validate compilation.", - "valid": - False, - } - }), + author=self.name, + actions=EventActions( + state_delta={ + "kernel_compilation_status": { + "success": False, + "retries": 0, + "message": "No kernel file was generated. Cannot validate compilation.", + "valid": False, + } + } + ), ) return @@ -95,16 +94,17 @@ async def _run_async_impl( error_msg = f"Kernel file not found at {kernel_path}" logging.error(f"[{self.name}] {error_msg}") yield Event( - author=self.name, - actions=EventActions( - state_delta={ - "kernel_compilation_status": { - "success": False, - "retries": 0, - "message": error_msg, - "valid": False, - } - }), + author=self.name, + actions=EventActions( + state_delta={ + "kernel_compilation_status": { + "success": False, + "retries": 0, + "message": error_msg, + "valid": False, + } + } + ), ) return @@ -116,7 +116,7 @@ async def _run_async_impl( while retry_count < self.max_retries: logging.info( - f"[{self.name}] Compilation validation attempt {retry_count + 1}/{self.max_retries}" + f"[{self.name}] Compilation validation attempt {retry_count + 1}/{self.max_retries}" ) # Set kernel_file_path for the compilation checker to use @@ -133,14 +133,14 @@ async def _run_async_impl( # Record this attempt in history attempt_record = { - "attempt": retry_count + 1, - "result": compilation_result, - "success": compilation_valid, - "fix_summary": ctx.session.state.get("fix_summary", None), + "attempt": retry_count + 1, + "result": compilation_result, + "success": compilation_valid, + "fix_summary": ctx.session.state.get("fix_summary", None), } ctx.session.state["compilation_history"].append(attempt_record) logging.info( - f"[{self.name}] Recorded attempt {retry_count + 1} in compilation history" + f"[{self.name}] Recorded attempt {retry_count + 1} in compilation history" ) # Clear fix_summary for next iteration @@ -149,16 +149,17 @@ async def _run_async_impl( if compilation_valid: logging.info(f"[{self.name}] ✓ Kernel compilation succeeded!") yield Event( - author=self.name, - actions=EventActions( - state_delta={ - "kernel_compilation_status": { - "success": True, - "retries": retry_count, - "message": "Kernel compiled successfully", - "valid": True, - } - }), + author=self.name, + actions=EventActions( + state_delta={ + "kernel_compilation_status": { + "success": True, + "retries": retry_count, + "message": "Kernel compiled successfully", + "valid": True, + } + } + ), ) return @@ -171,7 +172,7 @@ async def _run_async_impl( # Add debugging statements after 2nd retry if debug agent is available if self.debug_agent and retry_count >= 1: logging.info( - f"[{self.name}] Adding debugging statements to diagnose persistent issues..." + f"[{self.name}] Adding debugging statements to diagnose persistent issues..." ) async for event in self.debug_agent.run_async(ctx): yield event @@ -180,27 +181,24 @@ async def _run_async_impl( else: # Max retries reached logging.error( - f"[{self.name}] ✗ Max retries reached. Kernel still has compilation errors." + f"[{self.name}] ✗ Max retries reached. Kernel still has compilation errors." + ) + compilation_error_msg = ctx.session.state.get( + "compilation_results", "Unknown error" ) - compilation_error_msg = ctx.session.state.get("compilation_results", - "Unknown error") yield Event( - author=self.name, - actions=EventActions( - state_delta={ - "kernel_compilation_status": { - "success": - False, - "retries": - retry_count, - "message": - f"Kernel compilation failed after {self.max_retries} attempts", - "valid": - False, - "final_errors": - compilation_error_msg, - } - }), + author=self.name, + actions=EventActions( + state_delta={ + "kernel_compilation_status": { + "success": False, + "retries": retry_count, + "message": f"Kernel compilation failed after {self.max_retries} attempts", + "valid": False, + "final_errors": compilation_error_msg, + } + } + ), ) return @@ -214,25 +212,26 @@ class ValidateKernelCompilationAgent(BaseAgent): summary_agent: BaseAgent def __init__( - self, - name: str, - read_file_agent: BaseAgent, - validation_loop_agent: BaseAgent, - cleanup_agent: BaseAgent, - summary_agent: BaseAgent, - description: str = "", + self, + name: str, + read_file_agent: BaseAgent, + validation_loop_agent: BaseAgent, + cleanup_agent: BaseAgent, + summary_agent: BaseAgent, + description: str = "", ): super().__init__( - name=name, - description=description, - read_file_agent=read_file_agent, - validation_loop_agent=validation_loop_agent, - cleanup_agent=cleanup_agent, - summary_agent=summary_agent, + name=name, + description=description, + read_file_agent=read_file_agent, + validation_loop_agent=validation_loop_agent, + cleanup_agent=cleanup_agent, + summary_agent=summary_agent, ) async def _run_async_impl( - self, ctx: InvocationContext) -> AsyncGenerator[Event, None]: + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: # Step 1: Ensure we have a kernel file path = ctx.session.state.get("optimized_kernel_path") @@ -265,150 +264,149 @@ async def _run_async_impl( # These are called independently by the root orchestrator to allow user interaction between steps plan_kernel_agent = CustomLlmAgent( - name="PlanKernelAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=thinking_planner, - instruction=kernel_planning_prompt.PROMPT, - description= - "Creates or revises a detailed optimization plan for a Pallas kernel.", - tools=([search_api_tool, filesystem_tool_rw, vertex_ai_rag_tool] - if vertex_ai_rag_tool else [search_api_tool, filesystem_tool_rw]), - after_tool_callback=create_path_saver("kernel_plan_path"), + name="PlanKernelAgent", + model=MODEL_NAME, + generate_content_config=model_config, + planner=thinking_planner, + instruction=kernel_planning_prompt.PROMPT, + description="Creates or revises a detailed optimization plan for a Pallas kernel.", + tools=( + [search_api_tool, filesystem_tool_rw, vertex_ai_rag_tool] + if vertex_ai_rag_tool + else [search_api_tool, filesystem_tool_rw] + ), + after_tool_callback=create_path_saver("kernel_plan_path"), ) # Inner LLM agent that does the actual implementation work implement_kernel_llm_agent = CustomLlmAgent( - name="ImplementKernelLlmAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=thinking_planner, - instruction=kernel_implementation_prompt.PROMPT, - description= - "Implements the optimized Pallas kernel following the approved plan.", - tools=([search_api_tool, filesystem_tool_rw, vertex_ai_rag_tool] - if vertex_ai_rag_tool else [search_api_tool, filesystem_tool_rw]), - after_tool_callback=save_kernel_and_plan_paths, + name="ImplementKernelLlmAgent", + model=MODEL_NAME, + generate_content_config=model_config, + planner=thinking_planner, + instruction=kernel_implementation_prompt.PROMPT, + description="Implements the optimized Pallas kernel following the approved plan.", + tools=( + [search_api_tool, filesystem_tool_rw, vertex_ai_rag_tool] + if vertex_ai_rag_tool + else [search_api_tool, filesystem_tool_rw] + ), + after_tool_callback=save_kernel_and_plan_paths, ) # Agent that asks user what to do after kernel implementation ask_validation_agent = CustomLlmAgent( - name="AskValidationAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=thinking_planner, - instruction=ask_validation_prompt.PROMPT, - description= - "Asks the user whether they want to validate kernel compilation or do something else.", - include_contents="none", + name="AskValidationAgent", + model=MODEL_NAME, + generate_content_config=model_config, + planner=thinking_planner, + instruction=ask_validation_prompt.PROMPT, + description="Asks the user whether they want to validate kernel compilation or do something else.", + include_contents="none", ) # Kernel compilation validation agents # Read file agent for validation - extracts kernel path from user message or state read_file_for_validation_agent = CustomLlmAgent( - name="ReadFileForValidationAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=thinking_planner, - instruction=read_file_prompt.PROMPT, - description= - "Reads the kernel file mentioned by the user or from state for validation.", - tools=[filesystem_tool_rw], - after_tool_callback=create_path_saver("optimized_kernel_path"), + name="ReadFileForValidationAgent", + model=MODEL_NAME, + generate_content_config=model_config, + planner=thinking_planner, + instruction=read_file_prompt.PROMPT, + description="Reads the kernel file mentioned by the user or from state for validation.", + tools=[filesystem_tool_rw], + after_tool_callback=create_path_saver("optimized_kernel_path"), ) fix_kernel_compilation_agent = CustomLlmAgent( - name="FixKernelCompilationAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=thinking_planner, - instruction=fix_kernel_compilation.PROMPT, - description= - "Fixes compilation errors in the generated kernel while preserving optimization strategy.", - tools=([search_api_tool, filesystem_tool_rw, vertex_ai_rag_tool] - if vertex_ai_rag_tool else [search_api_tool, filesystem_tool_rw]), - before_agent_callback=load_kernel_and_plan_to_state, - after_model_callback=extract_fix_summary, - include_contents="none", + name="FixKernelCompilationAgent", + model=MODEL_NAME, + generate_content_config=model_config, + planner=thinking_planner, + instruction=fix_kernel_compilation.PROMPT, + description="Fixes compilation errors in the generated kernel while preserving optimization strategy.", + tools=( + [search_api_tool, filesystem_tool_rw, vertex_ai_rag_tool] + if vertex_ai_rag_tool + else [search_api_tool, filesystem_tool_rw] + ), + before_agent_callback=load_kernel_and_plan_to_state, + after_model_callback=extract_fix_summary, + include_contents="none", ) add_debug_statements_agent = CustomLlmAgent( - name="AddDebugStatementsAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=thinking_planner, - instruction=add_debug_statements.PROMPT, - description= - "Adds strategic debugging statements to diagnose persistent compilation issues.", - tools=[filesystem_tool_rw], - before_agent_callback=load_kernel_and_plan_to_state, - include_contents="none", + name="AddDebugStatementsAgent", + model=MODEL_NAME, + generate_content_config=model_config, + planner=thinking_planner, + instruction=add_debug_statements.PROMPT, + description="Adds strategic debugging statements to diagnose persistent compilation issues.", + tools=[filesystem_tool_rw], + before_agent_callback=load_kernel_and_plan_to_state, + include_contents="none", ) cleanup_debug_statements_agent = CustomLlmAgent( - name="CleanupDebugStatementsAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=thinking_planner, - instruction=cleanup_debug_statements.PROMPT, - description= - "Removes debugging statements from successfully compiled kernel.", - tools=[filesystem_tool_rw], - before_agent_callback=load_single_kernel_to_state, - include_contents="none", + name="CleanupDebugStatementsAgent", + model=MODEL_NAME, + generate_content_config=model_config, + planner=thinking_planner, + instruction=cleanup_debug_statements.PROMPT, + description="Removes debugging statements from successfully compiled kernel.", + tools=[filesystem_tool_rw], + before_agent_callback=load_single_kernel_to_state, + include_contents="none", ) kernel_compilation_checker_for_validation = KernelCompilationChecker( - name="KernelCompilationCheckerForValidation", - input_key="kernel_code", - output_key="compilation_results", - before_agent_callback=load_single_kernel_to_state, - auto_manage_servers=True, + name="KernelCompilationCheckerForValidation", + input_key="kernel_code", + output_key="compilation_results", + before_agent_callback=load_single_kernel_to_state, + auto_manage_servers=True, ) kernel_compilation_validation_loop = KernelCompilationValidationLoop( - name="KernelCompilationValidationLoop", - compilation_checker=kernel_compilation_checker_for_validation, - fix_agent=fix_kernel_compilation_agent, - debug_agent=add_debug_statements_agent, - max_retries=4, + name="KernelCompilationValidationLoop", + compilation_checker=kernel_compilation_checker_for_validation, + fix_agent=fix_kernel_compilation_agent, + debug_agent=add_debug_statements_agent, + max_retries=4, ) kernel_compilation_summary_agent = CustomLlmAgent( - name="KernelCompilationSummaryAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=thinking_planner, - instruction=kernel_compilation_summary.PROMPT, - description= - "Summarizes kernel compilation validation results with full trace on failure.", - include_contents="none", + name="KernelCompilationSummaryAgent", + model=MODEL_NAME, + generate_content_config=model_config, + planner=thinking_planner, + instruction=kernel_compilation_summary.PROMPT, + description="Summarizes kernel compilation validation results with full trace on failure.", + include_contents="none", ) # Standalone validation orchestration agent (invoked by root when user requests validation) validate_kernel_compilation_agent = ValidateKernelCompilationAgent( - name="ValidateKernelCompilationAgent", - read_file_agent=read_file_for_validation_agent, - validation_loop_agent=kernel_compilation_validation_loop, - cleanup_agent=cleanup_debug_statements_agent, - summary_agent=kernel_compilation_summary_agent, - description= - "Validates kernel compilation with automatic error fixing, debugging, and provides summary. Invoked when user requests validation.", + name="ValidateKernelCompilationAgent", + read_file_agent=read_file_for_validation_agent, + validation_loop_agent=kernel_compilation_validation_loop, + cleanup_agent=cleanup_debug_statements_agent, + summary_agent=kernel_compilation_summary_agent, + description="Validates kernel compilation with automatic error fixing, debugging, and provides summary. Invoked when user requests validation.", ) # Implementation agent - implements kernel then asks user about validation implement_kernel_agent = SequentialAgent( - name="ImplementKernelAgent", - sub_agents=[implement_kernel_llm_agent, ask_validation_agent], - description= - "Implements the optimized Pallas kernel and asks user about next steps.", + name="ImplementKernelAgent", + sub_agents=[implement_kernel_llm_agent, ask_validation_agent], + description="Implements the optimized Pallas kernel and asks user about next steps.", ) __all__ = [ - "KernelCompilationValidationLoop", - "ValidateKernelCompilationAgent", - "plan_kernel_agent", - "implement_kernel_agent", - "validate_kernel_compilation_agent", - "read_file_for_validation_agent", + "KernelCompilationValidationLoop", + "ValidateKernelCompilationAgent", + "plan_kernel_agent", + "implement_kernel_agent", + "validate_kernel_compilation_agent", + "read_file_for_validation_agent", ] diff --git a/MaxKernel/hitl_agent/subagents/kernel_writing/kernel_compilation.py b/MaxKernel/hitl_agent/subagents/kernel_writing/kernel_compilation.py index 2f83533..82189fc 100644 --- a/MaxKernel/hitl_agent/subagents/kernel_writing/kernel_compilation.py +++ b/MaxKernel/hitl_agent/subagents/kernel_writing/kernel_compilation.py @@ -5,34 +5,39 @@ from google.adk.agents import BaseAgent from google.adk.agents.invocation_context import InvocationContext from google.adk.events import Event, EventActions + from hitl_agent.constants import ( - EVAL_SERVER_PORT, - REQUEST_TIMEOUT, - TPU_TIMEOUT, + EVAL_SERVER_PORT, + REQUEST_TIMEOUT, + TPU_TIMEOUT, ) from hitl_agent.server_utils.server_manager_mixin import ServerManagerMixin class KernelCompilationChecker(ServerManagerMixin, BaseAgent): """Checks whether kernel compiles and escalates to stop the loop if grade is 'pass'. - - Automatically manages eval server lifecycle: - - Starts TPU and eval servers if not running - - Runs compilation check - - Tears down servers after completion if auto_manage_servers is True - """ + + Automatically manages eval server lifecycle: + - Starts TPU and eval servers if not running + - Runs compilation check + - Tears down servers after completion if auto_manage_servers is True + """ input_key: Optional[str] = None output_key: Optional[str] = None before_agent_callback: Optional[Callable] = None - auto_manage_servers: bool = False # Default to False to preserve existing behavior + auto_manage_servers: bool = ( + False # Default to False to preserve existing behavior + ) - def __init__(self, - name: str, - input_key: str, - output_key: str, - before_agent_callback: Optional[Callable] = None, - auto_manage_servers: bool = False): + def __init__( + self, + name: str, + input_key: str, + output_key: str, + before_agent_callback: Optional[Callable] = None, + auto_manage_servers: bool = False, + ): super().__init__(name=name, before_agent_callback=before_agent_callback) self.input_key = input_key self.output_key = output_key @@ -40,13 +45,14 @@ def __init__(self, self._servers_started = [] # Track which servers we started async def _run_async_impl( - self, ctx: InvocationContext) -> AsyncGenerator[Event, None]: + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: code = ctx.session.state.get(self.input_key, "") if not code: logging.warning(f"[{self.name}] No {self.input_key} found in context") yield Event( - author=self.name, - actions=EventActions(state_delta={self.output_key: None}), + author=self.name, + actions=EventActions(state_delta={self.output_key: None}), ) return @@ -56,23 +62,24 @@ async def _run_async_impl( if not servers_ok: logging.error(f"[{self.name}] Server startup failed: {error_msg}") yield Event( - author=self.name, - actions=EventActions(state_delta={ - self.output_key: f"Server startup failed: {error_msg}" - }), + author=self.name, + actions=EventActions( + state_delta={self.output_key: f"Server startup failed: {error_msg}"} + ), ) return # Call the TPU server to execute the code logging.info(f"[{self.name}] Running code") - async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout( - total=REQUEST_TIMEOUT)) as session: + async with aiohttp.ClientSession( + timeout=aiohttp.ClientTimeout(total=REQUEST_TIMEOUT) + ) as session: async with session.post( - f"http://localhost:{EVAL_SERVER_PORT}/evaluate", - json={ - "eval_type": "compilation_test", - "code": code, - "timeout": TPU_TIMEOUT, - }, + f"http://localhost:{EVAL_SERVER_PORT}/evaluate", + json={ + "eval_type": "compilation_test", + "code": code, + "timeout": TPU_TIMEOUT, + }, ) as response: if response.status == 200: result = await response.json() @@ -80,51 +87,57 @@ async def _run_async_impl( if result["exit_code"] == 0: logging.info(f"[{self.name}] Code execution successful.") yield Event( - author=self.name, - actions=EventActions( - state_delta={self.output_key: "Success"}), + author=self.name, + actions=EventActions(state_delta={self.output_key: "Success"}), ) - elif (result["error"] is None and result["output"] == "" and - result["exit_code"] == 1): + elif ( + result["error"] is None + and result["output"] == "" + and result["exit_code"] == 1 + ): logging.info( - f"[{self.name}] Code execution had exit code 1, but no error, indicating success." + f"[{self.name}] Code execution had exit code 1, but no error, indicating success." ) yield Event( - author=self.name, - actions=EventActions( - state_delta={self.output_key: "Success"}), + author=self.name, + actions=EventActions(state_delta={self.output_key: "Success"}), ) else: logging.info( - f"[{self.name}] Code execution failed. Loop will continue.") + f"[{self.name}] Code execution failed. Loop will continue." + ) # Use 'or' to handle None case - result.get() returns None if key exists with None value - error_msg = result.get("error") or result.get( - "output" - ) or "Unknown error: No error message or output available" + error_msg = ( + result.get("error") + or result.get("output") + or "Unknown error: No error message or output available" + ) # Add diagnostic logging when error field is None or empty if result.get("error") is None: logging.warning( - f"[{self.name}] Error field is None. " - f"exit_code: {result.get('exit_code')}, " - f"output length: {len(result.get('output', ''))}, " - f"Using output as fallback: {result.get('output', '')[:200]}" + f"[{self.name}] Error field is None. " + f"exit_code: {result.get('exit_code')}, " + f"output length: {len(result.get('output', ''))}, " + f"Using output as fallback: {result.get('output', '')[:200]}" ) yield Event( - author=self.name, - actions=EventActions( - state_delta={self.output_key: error_msg}), + author=self.name, + actions=EventActions(state_delta={self.output_key: error_msg}), ) else: error_detail = await response.text() logging.error( - f"[{self.name}] HTTP error {response.status}: {error_detail}") + f"[{self.name}] HTTP error {response.status}: {error_detail}" + ) ctx.session.state[self.output_key] = ( - f"HTTP error {response.status}: {error_detail}") + f"HTTP error {response.status}: {error_detail}" + ) yield Event(author=self.name) except Exception as e: logging.error(f"[{self.name}] Exception during code execution: {str(e)}") ctx.session.state[self.output_key] = ( - f"Exception during code execution: {str(e)}") + f"Exception during code execution: {str(e)}" + ) yield Event(author=self.name) finally: # Cleanup servers if we started them diff --git a/MaxKernel/hitl_agent/subagents/kernel_writing/prompts/__init__.py b/MaxKernel/hitl_agent/subagents/kernel_writing/prompts/__init__.py index 6966c65..afea6c6 100644 --- a/MaxKernel/hitl_agent/subagents/kernel_writing/prompts/__init__.py +++ b/MaxKernel/hitl_agent/subagents/kernel_writing/prompts/__init__.py @@ -1,25 +1,25 @@ """Prompts for kernel writing subagent.""" from . import ( - kernel_planning_prompt, - kernel_implementation_prompt, - ask_validation_prompt, - fix_kernel_compilation, - kernel_compilation_summary, - add_debug_statements, - cleanup_debug_statements, - summary_prompt, - read_file_prompt, + add_debug_statements, + ask_validation_prompt, + cleanup_debug_statements, + fix_kernel_compilation, + kernel_compilation_summary, + kernel_implementation_prompt, + kernel_planning_prompt, + read_file_prompt, + summary_prompt, ) __all__ = [ - 'kernel_planning_prompt', - 'kernel_implementation_prompt', - 'ask_validation_prompt', - 'fix_kernel_compilation', - 'kernel_compilation_summary', - 'add_debug_statements', - 'cleanup_debug_statements', - 'summary_prompt', - 'read_file_prompt', + "kernel_planning_prompt", + "kernel_implementation_prompt", + "ask_validation_prompt", + "fix_kernel_compilation", + "kernel_compilation_summary", + "add_debug_statements", + "cleanup_debug_statements", + "summary_prompt", + "read_file_prompt", ] diff --git a/MaxKernel/hitl_agent/subagents/profiling/__init__.py b/MaxKernel/hitl_agent/subagents/profiling/__init__.py index 37199b6..b4f7787 100644 --- a/MaxKernel/hitl_agent/subagents/profiling/__init__.py +++ b/MaxKernel/hitl_agent/subagents/profiling/__init__.py @@ -1,17 +1,17 @@ """Profiling subagent module.""" from .agent import ( - profile_agent, - read_file_for_profiling_agent, - generate_profiling_script_agent, - eval_profile_agent, - summarize_profile_agent, + eval_profile_agent, + generate_profiling_script_agent, + profile_agent, + read_file_for_profiling_agent, + summarize_profile_agent, ) __all__ = [ - 'profile_agent', - 'read_file_for_profiling_agent', - 'generate_profiling_script_agent', - 'eval_profile_agent', - 'summarize_profile_agent', + "profile_agent", + "read_file_for_profiling_agent", + "generate_profiling_script_agent", + "eval_profile_agent", + "summarize_profile_agent", ] diff --git a/MaxKernel/hitl_agent/subagents/profiling/agent.py b/MaxKernel/hitl_agent/subagents/profiling/agent.py index 78ba940..74fbcd1 100644 --- a/MaxKernel/hitl_agent/subagents/profiling/agent.py +++ b/MaxKernel/hitl_agent/subagents/profiling/agent.py @@ -1,113 +1,109 @@ """Profiling subagent - performance profiling and analysis.""" from google.adk.agents import SequentialAgent + from hitl_agent.callbacks import ( - create_path_saver, - load_profiling_script_to_state, - load_single_kernel_to_state, + create_path_saver, + load_profiling_script_to_state, + load_single_kernel_to_state, ) from hitl_agent.config import model_config, thinking_planner - +from hitl_agent.constants import MODEL_NAME +from hitl_agent.custom_types import CustomLlmAgent from hitl_agent.subagents.profiling import offline_tools +from hitl_agent.subagents.profiling.kernel_profile import KernelProfiler from hitl_agent.subagents.profiling.prompts import ( - analyze_profile_prompt, - gen_profiling_script, - read_file_prompt, - read_profiling_script_prompt, + analyze_profile_prompt, + gen_profiling_script, + read_file_prompt, + read_profiling_script_prompt, ) from hitl_agent.tools.tools import filesystem_tool_rw, vertex_ai_rag_tool -from hitl_agent.custom_types import CustomLlmAgent -from hitl_agent.constants import MODEL_NAME -from hitl_agent.subagents.profiling.kernel_profile import KernelProfiler # Read file agent for profiling read_file_for_profiling_agent = CustomLlmAgent( - name="ReadFileForProfilingAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=thinking_planner, - instruction=read_file_prompt.PROMPT, - description= - "Reads the kernel file mentioned by the user for profiling analysis.", - tools=[filesystem_tool_rw], - after_tool_callback=create_path_saver("kernel_file_path"), + name="ReadFileForProfilingAgent", + model=MODEL_NAME, + generate_content_config=model_config, + planner=thinking_planner, + instruction=read_file_prompt.PROMPT, + description="Reads the kernel file mentioned by the user for profiling analysis.", + tools=[filesystem_tool_rw], + after_tool_callback=create_path_saver("kernel_file_path"), ) # Profiling script generation agent - writes profiling script to file generate_profiling_script_agent = CustomLlmAgent( - name="GenerateProfilingScriptAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=thinking_planner, - instruction=gen_profiling_script.PROMPT, - description= - "Generates a profiling script to identify performance bottlenecks in the kernel code and writes it to a file.", - tools=[filesystem_tool_rw], - before_agent_callback=load_single_kernel_to_state, - after_tool_callback=create_path_saver("profiling_script_path"), + name="GenerateProfilingScriptAgent", + model=MODEL_NAME, + generate_content_config=model_config, + planner=thinking_planner, + instruction=gen_profiling_script.PROMPT, + description="Generates a profiling script to identify performance bottlenecks in the kernel code and writes it to a file.", + tools=[filesystem_tool_rw], + before_agent_callback=load_single_kernel_to_state, + after_tool_callback=create_path_saver("profiling_script_path"), ) # Read profiling script agent - loads the generated profiling script file contents into state read_profiling_script_agent = CustomLlmAgent( - name="ReadProfilingScriptAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=thinking_planner, - instruction=read_profiling_script_prompt.PROMPT, - description= - "Loads the generated profiling script file contents from disk into memory for execution.", - before_agent_callback=load_profiling_script_to_state, - include_contents="none", + name="ReadProfilingScriptAgent", + model=MODEL_NAME, + generate_content_config=model_config, + planner=thinking_planner, + instruction=read_profiling_script_prompt.PROMPT, + description="Loads the generated profiling script file contents from disk into memory for execution.", + before_agent_callback=load_profiling_script_to_state, + include_contents="none", ) # Profiling execution agent eval_profile_agent = KernelProfiler( - name="ProfileEvalAgent", - input_key="profiling_script", - output_key="profiling_results", - auto_manage_servers=True, + name="ProfileEvalAgent", + input_key="profiling_script", + output_key="profiling_results", + auto_manage_servers=True, ) # Profiling summary agent summarize_profile_agent = CustomLlmAgent( - name="SummarizeProfileAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=thinking_planner, - instruction=analyze_profile_prompt.PROMPT, - description=( - "Summarizes the profiling results of the kernel and performs deep" - " analysis using offline XProf tools."), - output_key="profiling_summary", - include_contents="none", - tools=[ - offline_tools.load_xplane_and_query, - offline_tools.get_hlo_dump, - offline_tools.create_chart_from_xplane, - offline_tools.get_overview_page_metrics, - vertex_ai_rag_tool, - ], + name="SummarizeProfileAgent", + model=MODEL_NAME, + generate_content_config=model_config, + planner=thinking_planner, + instruction=analyze_profile_prompt.PROMPT, + description=( + "Summarizes the profiling results of the kernel and performs deep analysis using offline XProf tools." + ), + output_key="profiling_summary", + include_contents="none", + tools=[ + offline_tools.load_xplane_and_query, + offline_tools.get_hlo_dump, + offline_tools.create_chart_from_xplane, + offline_tools.get_overview_page_metrics, + vertex_ai_rag_tool, + ], ) # Main profiling orchestrator agent profile_agent = SequentialAgent( - name="ProfileAgentOrchestrator", - sub_agents=[ - read_file_for_profiling_agent, - generate_profiling_script_agent, - read_profiling_script_agent, - eval_profile_agent, - summarize_profile_agent, - ], - description= - "Profiles the Pallas kernel to identify performance bottlenecks.", + name="ProfileAgentOrchestrator", + sub_agents=[ + read_file_for_profiling_agent, + generate_profiling_script_agent, + read_profiling_script_agent, + eval_profile_agent, + summarize_profile_agent, + ], + description="Profiles the Pallas kernel to identify performance bottlenecks.", ) __all__ = [ - "profile_agent", - "read_file_for_profiling_agent", - "generate_profiling_script_agent", - "read_profiling_script_agent", - "eval_profile_agent", - "summarize_profile_agent", + "profile_agent", + "read_file_for_profiling_agent", + "generate_profiling_script_agent", + "read_profiling_script_agent", + "eval_profile_agent", + "summarize_profile_agent", ] diff --git a/MaxKernel/hitl_agent/subagents/profiling/kernel_profile.py b/MaxKernel/hitl_agent/subagents/profiling/kernel_profile.py index 3a25e0b..ecd7237 100644 --- a/MaxKernel/hitl_agent/subagents/profiling/kernel_profile.py +++ b/MaxKernel/hitl_agent/subagents/profiling/kernel_profile.py @@ -5,10 +5,11 @@ from google.adk.agents import BaseAgent from google.adk.agents.invocation_context import InvocationContext from google.adk.events import Event, EventActions + from hitl_agent.constants import ( - EVAL_SERVER_PORT, - REQUEST_TIMEOUT, - TPU_TIMEOUT, + EVAL_SERVER_PORT, + REQUEST_TIMEOUT, + TPU_TIMEOUT, ) from hitl_agent.server_utils.server_manager_mixin import ServerManagerMixin @@ -16,26 +17,28 @@ class KernelProfiler(ServerManagerMixin, BaseAgent): """Profiles the kernel to identify performance bottlenecks. - Automatically manages eval server lifecycle: - - Starts TPU and eval servers if not running - - Runs profiling - - Tears down servers after completion if auto_manage_servers is True - """ + Automatically manages eval server lifecycle: + - Starts TPU and eval servers if not running + - Runs profiling + - Tears down servers after completion if auto_manage_servers is True + """ input_key: Optional[str] = None output_key: Optional[str] = None before_agent_callback: Optional[Callable] = None raise_exception_upon_success: bool = True - auto_manage_servers: bool = False # Default to False to preserve existing behavior + auto_manage_servers: bool = ( + False # Default to False to preserve existing behavior + ) def __init__( - self, - name: str, - input_key: str, - output_key: str, - before_agent_callback: Optional[Callable] = None, - raise_exception_upon_success: bool = True, - auto_manage_servers: bool = False, + self, + name: str, + input_key: str, + output_key: str, + before_agent_callback: Optional[Callable] = None, + raise_exception_upon_success: bool = True, + auto_manage_servers: bool = False, ): super().__init__(name=name, before_agent_callback=before_agent_callback) self.input_key = input_key @@ -45,13 +48,14 @@ def __init__( self._servers_started = [] # Track which servers we started async def _run_async_impl( - self, ctx: InvocationContext) -> AsyncGenerator[Event, None]: + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: profile_code = ctx.session.state.get(self.input_key, "") if not profile_code: logging.warning(f"[{self.name}] No profile_code found in context") yield Event( - author=self.name, - actions=EventActions(state_delta={self.output_key: None}), + author=self.name, + actions=EventActions(state_delta={self.output_key: None}), ) return @@ -61,23 +65,24 @@ async def _run_async_impl( if not servers_ok: logging.error(f"[{self.name}] Server startup failed: {error_msg}") yield Event( - author=self.name, - actions=EventActions(state_delta={ - self.output_key: f"Server startup failed: {error_msg}" - }), + author=self.name, + actions=EventActions( + state_delta={self.output_key: f"Server startup failed: {error_msg}"} + ), ) return # Call the TPU server to execute the code logging.info(f"[{self.name}] Running code") - async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout( - total=REQUEST_TIMEOUT)) as session: + async with aiohttp.ClientSession( + timeout=aiohttp.ClientTimeout(total=REQUEST_TIMEOUT) + ) as session: async with session.post( - f"http://localhost:{EVAL_SERVER_PORT}/evaluate", - json={ - "eval_type": "profile", - "code": profile_code, - "timeout": TPU_TIMEOUT, - }, + f"http://localhost:{EVAL_SERVER_PORT}/evaluate", + json={ + "eval_type": "profile", + "code": profile_code, + "timeout": TPU_TIMEOUT, + }, ) as response: if response.status == 200: result = await response.json() @@ -91,15 +96,13 @@ async def _run_async_impl( # Profiling succeeds if exit_code is 0 and we have output # Stderr may contain warnings (like TensorFlow import warnings) which are not failures if exit_code != 0: - full_error = ( - f"Profiling script failed with exit code {exit_code}") + full_error = f"Profiling script failed with exit code {exit_code}" if error_msg: full_error += f": {error_msg}" logging.error(f"[{self.name}] {full_error}") yield Event( - author=self.name, - actions=EventActions( - state_delta={self.output_key: full_error}), + author=self.name, + actions=EventActions(state_delta={self.output_key: full_error}), ) elif not output or output.strip() == "": full_error = "Profiling script produced no output" @@ -107,9 +110,8 @@ async def _run_async_impl( full_error += f". Stderr: {error_msg}" logging.error(f"[{self.name}] {full_error}") yield Event( - author=self.name, - actions=EventActions( - state_delta={self.output_key: full_error}), + author=self.name, + actions=EventActions(state_delta={self.output_key: full_error}), ) else: # Successful profiling - parse the ratio and xplane path @@ -128,52 +130,58 @@ async def _run_async_impl( # Log warnings if present, but don't fail if error_msg: logging.warning( - f"[{self.name}] Profiling succeeded but had warnings in" - f" stderr: {error_msg[:200]}") + f"[{self.name}] Profiling succeeded but had warnings in stderr: {error_msg[:200]}" + ) logging.info( - f"[{self.name}] Profiling succeeded with ratio: {ratio}," - f" xplane_path: {xplane_path}") + f"[{self.name}] Profiling succeeded with ratio: {ratio}, xplane_path: {xplane_path}" + ) yield Event( - author=self.name, - actions=EventActions( - escalate=False, - state_delta={ - self.output_key: { - "DMAs_and_memory_transfers_ratio": ratio, - "compute_ratio": 1 - ratio, - "xplane_path": xplane_path, - } - }, - ), + author=self.name, + actions=EventActions( + escalate=False, + state_delta={ + self.output_key: { + "DMAs_and_memory_transfers_ratio": ratio, + "compute_ratio": 1 - ratio, + "xplane_path": xplane_path, + } + }, + ), ) except (ValueError, KeyError) as e: - error_msg_full = f"Failed to parse profiling output: '{output}'. Error: {e}" + error_msg_full = ( + f"Failed to parse profiling output: '{output}'. Error: {e}" + ) logging.error(f"[{self.name}] {error_msg_full}") yield Event( - author=self.name, - actions=EventActions( - state_delta={self.output_key: error_msg_full}), + author=self.name, + actions=EventActions( + state_delta={self.output_key: error_msg_full} + ), ) else: error_detail = await response.text() logging.error( - f"[{self.name}] HTTP error {response.status}: {error_detail}") + f"[{self.name}] HTTP error {response.status}: {error_detail}" + ) yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: - f"HTTP error {response.status}: {error_detail}" - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: f"HTTP error {response.status}: {error_detail}" + } + ), ) except Exception as e: logging.error(f"[{self.name}] Exception during code execution: {str(e)}") yield Event( - author=self.name, - actions=EventActions(state_delta={ - self.output_key: f"Exception during code execution: {str(e)}" - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: f"Exception during code execution: {str(e)}" + } + ), ) finally: # Cleanup servers if we started them diff --git a/MaxKernel/hitl_agent/subagents/profiling/offline_tools.py b/MaxKernel/hitl_agent/subagents/profiling/offline_tools.py index 93c862f..19b01f0 100644 --- a/MaxKernel/hitl_agent/subagents/profiling/offline_tools.py +++ b/MaxKernel/hitl_agent/subagents/profiling/offline_tools.py @@ -1,11 +1,10 @@ """Standalone XProf tools for analyzing xplane.pb files without external services.""" import gzip -import io import json -import os import sqlite3 -from typing import Any, Dict, List, Optional +from typing import Optional + import matplotlib.pyplot as plt import pandas as pd from tensorflow.tsl.profiler.protobuf import xplane_pb2 @@ -20,18 +19,18 @@ def _get_xplane_path(profiling_results: str) -> str: def load_xplane_and_query(xplane_path: str, sql_query: str) -> str: """Loads an xplane.pb file into an in-memory SQLite DB and runs a SQL query. - The database schema is: - - planes (id, name) - - lines (id, plane_id, display_id, name, timestamp_ns) - - events (plane_id, line_id, name, offset_ps, duration_ps, start_ps, end_ps) + The database schema is: + - planes (id, name) + - lines (id, plane_id, display_id, name, timestamp_ns) + - events (plane_id, line_id, name, offset_ps, duration_ps, start_ps, end_ps) - Args: - xplane_path: Path to the .xplane.pb file. - sql_query: The SQL query to execute against the loaded data. + Args: + xplane_path: Path to the .xplane.pb file. + sql_query: The SQL query to execute against the loaded data. - Returns: - A markdown-formatted table of the query results. - """ + Returns: + A markdown-formatted table of the query results. + """ try: # Open file (handle gz if needed) open_func = gzip.open if xplane_path.endswith(".gz") else open @@ -63,8 +62,8 @@ def get_meta_name(meta_map, mid): for line in plane.lines: c.execute( - "INSERT INTO lines VALUES (?, ?, ?, ?, ?)", - (line.id, plane.id, line.display_id, line.name, line.timestamp_ns), + "INSERT INTO lines VALUES (?, ?, ?, ?, ?)", + (line.id, plane.id, line.display_id, line.name, line.timestamp_ns), ) for event in line.events: @@ -72,16 +71,16 @@ def get_meta_name(meta_map, mid): start_ps = event.offset_ps end_ps = start_ps + event.duration_ps c.execute( - "INSERT INTO events VALUES (?, ?, ?, ?, ?, ?, ?)", - ( - plane.id, - line.id, - name, - event.offset_ps, - event.duration_ps, - start_ps, - end_ps, - ), + "INSERT INTO events VALUES (?, ?, ?, ?, ?, ?, ?)", + ( + plane.id, + line.id, + name, + event.offset_ps, + event.duration_ps, + start_ps, + end_ps, + ), ) conn.commit() @@ -96,17 +95,18 @@ def get_meta_name(meta_map, mid): return f"Error executing query: {e}" -def get_hlo_dump(xplane_path: str, - hlo_module_name: Optional[str] = None) -> str: +def get_hlo_dump( + xplane_path: str, hlo_module_name: Optional[str] = None +) -> str: """Extracts HLO proto from xplane.pb if available. - Args: - xplane_path: Path to .xplane.pb file. - hlo_module_name: Optional name filter. + Args: + xplane_path: Path to .xplane.pb file. + hlo_module_name: Optional name filter. - Returns: - Status string indicating where HLO was saved or if not found. - """ + Returns: + Status string indicating where HLO was saved or if not found. + """ try: open_func = gzip.open if xplane_path.endswith(".gz") else open with open_func(xplane_path, "rb") as f: @@ -125,32 +125,33 @@ def get_hlo_dump(xplane_path: str, # For now, returning a placeholder as true extraction requires inspecting specific metadata IDs return ( - "HLO extraction not fully implemented in this standalone version yet" - " (requires metadata ID mapping). Please use `load_xplane_and_query` to" - " explore 'hlo' related events.") + "HLO extraction not fully implemented in this standalone version yet" + " (requires metadata ID mapping). Please use `load_xplane_and_query` to" + " explore 'hlo' related events." + ) except Exception as e: return f"Error extracting HLO: {e}" def create_chart_from_xplane( - xplane_path: str, - sql_query: str, - chart_type: str = "bar", - x_col: str = "name", - y_col: str = "value", - title: str = "", + xplane_path: str, + sql_query: str, + chart_type: str = "bar", + x_col: str = "name", + y_col: str = "value", + title: str = "", ) -> str: """Generates a chart from xplane data using SQL query. - Args: - xplane_path: Path to .xplane.pb - sql_query: SQL query to get data. - chart_type: 'bar' or 'pie'. - x_col: Column for X axis (bar). - y_col: Column for Y axis (bar) or values (pie). - title: Chart title. - """ + Args: + xplane_path: Path to .xplane.pb + sql_query: SQL query to get data. + chart_type: 'bar' or 'pie'. + x_col: Column for X axis (bar). + y_col: Column for Y axis (bar) or values (pie). + title: Chart title. + """ try: # Re-use loading logic (inefficient but stateless) # TODO: we might want to cache the DB connection or pass it around. @@ -179,23 +180,23 @@ def get_meta_name(meta_map, mid): for line in plane.lines: c.execute( - "INSERT INTO lines VALUES (?, ?, ?, ?, ?)", - (line.id, plane.id, line.display_id, line.name, line.timestamp_ns), + "INSERT INTO lines VALUES (?, ?, ?, ?, ?)", + (line.id, plane.id, line.display_id, line.name, line.timestamp_ns), ) for event in line.events: name = get_meta_name(plane.event_metadata, event.metadata_id) start_ps = event.offset_ps c.execute( - "INSERT INTO events VALUES (?, ?, ?, ?, ?, ?, ?)", - ( - plane.id, - line.id, - name, - event.offset_ps, - event.duration_ps, - start_ps, - start_ps + event.duration_ps, - ), + "INSERT INTO events VALUES (?, ?, ?, ?, ?, ?, ?)", + ( + plane.id, + line.id, + name, + event.offset_ps, + event.duration_ps, + start_ps, + start_ps + event.duration_ps, + ), ) conn.commit() @@ -230,15 +231,15 @@ def get_meta_name(meta_map, mid): def get_overview_page_metrics(xplane_path: str) -> str: """Returns metrics and metadata from overview page for a given Xprof session. - Mimics the behavior of overview_page_tool.get_overview_page_metrics by - extracting high-level metrics from the xplane.pb file directly. + Mimics the behavior of overview_page_tool.get_overview_page_metrics by + extracting high-level metrics from the xplane.pb file directly. - Args: - xplane_path: Path to the .xplane.pb file. + Args: + xplane_path: Path to the .xplane.pb file. - Returns: - A JSON string containing metrics and metadata. - """ + Returns: + A JSON string containing metrics and metadata. + """ try: open_func = gzip.open if xplane_path.endswith(".gz") else open with open_func(xplane_path, "rb") as f: @@ -252,8 +253,11 @@ def get_overview_page_metrics(xplane_path: str) -> str: device_planes = [] for plane in xspace.planes: - if ("device" in plane.name.lower() or "tpu" in plane.name.lower() or - "gpu" in plane.name.lower()): + if ( + "device" in plane.name.lower() + or "tpu" in plane.name.lower() + or "gpu" in plane.name.lower() + ): device_planes.append(plane) else: host_planes.append(plane) @@ -305,8 +309,9 @@ def get_overview_page_metrics(xplane_path: str) -> str: # Assume full parallelism potential = device_count * total_duration potential_ps = len(device_planes) * total_duration_ps if potential_ps > 0: - metrics["device_duty_cycle_percent"] = (total_device_busy_ps / - potential_ps) * 100 + metrics["device_duty_cycle_percent"] = ( + total_device_busy_ps / potential_ps + ) * 100 else: metrics["device_duty_cycle_percent"] = 0 diff --git a/MaxKernel/hitl_agent/subagents/profiling/prompts/__init__.py b/MaxKernel/hitl_agent/subagents/profiling/prompts/__init__.py index 0829442..782a939 100644 --- a/MaxKernel/hitl_agent/subagents/profiling/prompts/__init__.py +++ b/MaxKernel/hitl_agent/subagents/profiling/prompts/__init__.py @@ -1,13 +1,13 @@ """Prompts for profiling subagent.""" from . import ( - gen_profiling_script, - read_file_prompt, - read_profiling_script_prompt, + gen_profiling_script, + read_file_prompt, + read_profiling_script_prompt, ) __all__ = [ - 'gen_profiling_script', - 'read_file_prompt', - 'read_profiling_script_prompt', + "gen_profiling_script", + "read_file_prompt", + "read_profiling_script_prompt", ] diff --git a/MaxKernel/hitl_agent/subagents/profiling/prompts/analyze_profile_prompt.py b/MaxKernel/hitl_agent/subagents/profiling/prompts/analyze_profile_prompt.py index 9aab494..3d0a1d6 100644 --- a/MaxKernel/hitl_agent/subagents/profiling/prompts/analyze_profile_prompt.py +++ b/MaxKernel/hitl_agent/subagents/profiling/prompts/analyze_profile_prompt.py @@ -1,4 +1,4 @@ -#hitl_agent/subagents/profiling/prompts/analyze_profile_prompt.py +# hitl_agent/subagents/profiling/prompts/analyze_profile_prompt.py """Prompt for analyzing profiling results using offline XProf tools.""" PROMPT = """ diff --git a/MaxKernel/hitl_agent/subagents/testing/__init__.py b/MaxKernel/hitl_agent/subagents/testing/__init__.py index eae94cc..d43e5e3 100644 --- a/MaxKernel/hitl_agent/subagents/testing/__init__.py +++ b/MaxKernel/hitl_agent/subagents/testing/__init__.py @@ -1,23 +1,23 @@ """Testing subagent module.""" from .agent import ( - TestRunner, - SyntaxValidationAgent, - ImportValidationAgent, - TestStructureValidationAgent, - MockTestExecutionAgent, - TestValidationLoopAgent, - validated_test_generation_agent, - unified_test_agent, + ImportValidationAgent, + MockTestExecutionAgent, + SyntaxValidationAgent, + TestRunner, + TestStructureValidationAgent, + TestValidationLoopAgent, + unified_test_agent, + validated_test_generation_agent, ) __all__ = [ - 'TestRunner', - 'SyntaxValidationAgent', - 'ImportValidationAgent', - 'TestStructureValidationAgent', - 'MockTestExecutionAgent', - 'TestValidationLoopAgent', - 'validated_test_generation_agent', - 'unified_test_agent', + "TestRunner", + "SyntaxValidationAgent", + "ImportValidationAgent", + "TestStructureValidationAgent", + "MockTestExecutionAgent", + "TestValidationLoopAgent", + "validated_test_generation_agent", + "unified_test_agent", ] diff --git a/MaxKernel/hitl_agent/subagents/testing/agent.py b/MaxKernel/hitl_agent/subagents/testing/agent.py index de0a52a..bdd4926 100644 --- a/MaxKernel/hitl_agent/subagents/testing/agent.py +++ b/MaxKernel/hitl_agent/subagents/testing/agent.py @@ -3,43 +3,42 @@ This module contains all agents related to generating, validating, and executing tests. """ -import os import ast -import logging import asyncio +import logging +import os +import re import subprocess import tempfile -import re -from typing import Optional, Callable, AsyncGenerator +from typing import AsyncGenerator, Callable, Optional -from google.adk.agents import SequentialAgent, BaseAgent +from google.adk.agents import BaseAgent, SequentialAgent from google.adk.agents.invocation_context import InvocationContext from google.adk.events import Event, EventActions -from hitl_agent.custom_types import CustomLlmAgent -from hitl_agent.constants import MODEL_NAME -from hitl_agent.config import model_config, thinking_planner -from hitl_agent.tools.search_api_tool import search_api_tool -from hitl_agent.tools.tools import filesystem_tool_rw, vertex_ai_rag_tool - from hitl_agent.callbacks import create_path_saver +from hitl_agent.config import model_config, thinking_planner +from hitl_agent.constants import MODEL_NAME +from hitl_agent.custom_types import CustomLlmAgent from hitl_agent.subagents.testing.prompts import ( - gen_test_file, - fix_test_script, - validation_summary, - summarize_test_results_prompt, - read_file_prompt, + fix_test_script, + gen_test_file, + read_file_prompt, + summarize_test_results_prompt, + validation_summary, ) +from hitl_agent.tools.search_api_tool import search_api_tool +from hitl_agent.tools.tools import filesystem_tool_rw, vertex_ai_rag_tool class TestRunner(BaseAgent): """Executes pytest on a generated test file and captures results with full tracebacks. - Automatically manages eval server lifecycle: - - Starts TPU and eval servers if not running - - Runs tests - - Always tears down servers after completion - """ + Automatically manages eval server lifecycle: + - Starts TPU and eval servers if not running + - Runs tests + - Always tears down servers after completion + """ input_key: Optional[str] = None output_key: Optional[str] = None @@ -47,12 +46,12 @@ class TestRunner(BaseAgent): auto_manage_servers: bool = True def __init__( - self, - name: str, - input_key: str, - output_key: str, - before_agent_callback: Optional[Callable] = None, - auto_manage_servers: bool = True, + self, + name: str, + input_key: str, + output_key: str, + before_agent_callback: Optional[Callable] = None, + auto_manage_servers: bool = True, ): super().__init__(name=name, before_agent_callback=before_agent_callback) self.input_key = input_key @@ -64,9 +63,9 @@ def _is_server_running(self, server_name: str) -> bool: """Check if a server process is running.""" try: result = subprocess.run( - ["pgrep", "-f", server_name], - capture_output=True, - text=True, + ["pgrep", "-f", server_name], + capture_output=True, + text=True, ) return result.returncode == 0 except Exception as e: @@ -78,11 +77,11 @@ async def _start_server(self, server_type: str, setup_script: str) -> bool: try: logging.info(f"Starting {server_type} server...") process = await asyncio.create_subprocess_exec( - "bash", - setup_script, - f"--start-{server_type}", - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, + "bash", + setup_script, + f"--start-{server_type}", + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, ) await process.wait() await asyncio.sleep(3) @@ -104,11 +103,11 @@ async def _stop_server(self, server_type: str, setup_script: str): try: logging.info(f"Stopping {server_type} server...") process = await asyncio.create_subprocess_exec( - "bash", - setup_script, - f"--stop-{server_type}", - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, + "bash", + setup_script, + f"--stop-{server_type}", + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, ) await process.wait() logging.info(f"{server_type} server stopped") @@ -120,10 +119,12 @@ async def _ensure_servers_running(self) -> tuple[bool, str]: if not self.auto_manage_servers: return True, "" - kernel_gen_path = os.path.join(os.path.dirname(__file__), "..", "..", - "kernel_gen_agent", "kernel_eval") - setup_script = os.path.join(os.path.dirname(__file__), "..", "..", - "server_utils", "setup.sh") + kernel_gen_path = os.path.join( + os.path.dirname(__file__), "..", "..", "kernel_gen_agent", "kernel_eval" + ) + setup_script = os.path.join( + os.path.dirname(__file__), "..", "..", "server_utils", "setup.sh" + ) if not os.path.exists(setup_script): error_msg = f"Setup script not found at {setup_script}" @@ -151,29 +152,32 @@ async def _cleanup_servers(self): if not self._servers_started: return - setup_script = os.path.join(os.path.dirname(__file__), "..", "..", - "server_utils", "setup.sh") + setup_script = os.path.join( + os.path.dirname(__file__), "..", "..", "server_utils", "setup.sh" + ) for server_type in self._servers_started: await self._stop_server(server_type, setup_script) async def _run_async_impl( - self, ctx: InvocationContext) -> AsyncGenerator[Event, None]: + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: test_file_path = ctx.session.state.get(self.input_key, "") if not test_file_path: error_msg = "No test file was generated. Please generate a test file first using the GenerateTestFileAgent." logging.warning(f"[{self.name}] {error_msg}") yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: { - "exit_code": -1, - "output": error_msg, - "success": False, - } - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: { + "exit_code": -1, + "output": error_msg, + "success": False, + } + } + ), ) return @@ -181,15 +185,16 @@ async def _run_async_impl( error_msg = f"Test file not found at {test_file_path}. Please ensure the file exists." logging.warning(f"[{self.name}] {error_msg}") yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: { - "exit_code": -1, - "output": error_msg, - "success": False, - } - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: { + "exit_code": -1, + "output": error_msg, + "success": False, + } + } + ), ) return @@ -197,71 +202,74 @@ async def _run_async_impl( servers_ok, error_msg = await self._ensure_servers_running() if not servers_ok: yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: { - "exit_code": -1, - "output": f"Server startup failed: {error_msg}", - "success": False, - } - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: { + "exit_code": -1, + "output": f"Server startup failed: {error_msg}", + "success": False, + } + } + ), ) return logging.info(f"[{self.name}] Running pytest on {test_file_path}") result = subprocess.run( - ["pytest", test_file_path, "-v", "--tb=long", "--maxfail=1"], - capture_output=True, - text=True, - cwd=os.path.dirname(test_file_path), - timeout=300, + ["pytest", test_file_path, "-v", "--tb=long", "--maxfail=1"], + capture_output=True, + text=True, + cwd=os.path.dirname(test_file_path), + timeout=300, ) full_output = f"STDOUT:\n{result.stdout}\n\nSTDERR:\n{result.stderr}" test_results = { - "exit_code": result.returncode, - "output": full_output, - "success": result.returncode == 0, + "exit_code": result.returncode, + "output": full_output, + "success": result.returncode == 0, } logging.info( - f"[{self.name}] Test execution completed with exit code {result.returncode}" + f"[{self.name}] Test execution completed with exit code {result.returncode}" ) yield Event( - author=self.name, - actions=EventActions(state_delta={self.output_key: test_results}), + author=self.name, + actions=EventActions(state_delta={self.output_key: test_results}), ) except subprocess.TimeoutExpired: error_msg = "Test execution timed out after 5 minutes" logging.error(f"[{self.name}] {error_msg}") yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: { - "exit_code": -1, - "output": error_msg, - "success": False, - } - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: { + "exit_code": -1, + "output": error_msg, + "success": False, + } + } + ), ) except Exception as e: error_msg = f"Exception during test execution: {str(e)}" logging.error(f"[{self.name}] {error_msg}") yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: { - "exit_code": -1, - "output": error_msg, - "success": False, - } - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: { + "exit_code": -1, + "output": error_msg, + "success": False, + } + } + ), ) finally: await self._cleanup_servers() @@ -279,22 +287,24 @@ def __init__(self, name: str, input_key: str, output_key: str): self.output_key = output_key async def _run_async_impl( - self, ctx: InvocationContext) -> AsyncGenerator[Event, None]: + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: test_file_path = ctx.session.state.get(self.input_key, "") if not test_file_path or not os.path.exists(test_file_path): error_msg = f"Test file not found at {test_file_path}" logging.error(f"[{self.name}] {error_msg}") yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: { - "valid": False, - "errors": [error_msg], - "validation_type": "syntax", - } - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: { + "valid": False, + "errors": [error_msg], + "validation_type": "syntax", + } + } + ), ) return @@ -305,47 +315,51 @@ async def _run_async_impl( ast.parse(code) logging.info( - f"[{self.name}] Syntax validation passed for {test_file_path}") + f"[{self.name}] Syntax validation passed for {test_file_path}" + ) yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: { - "valid": True, - "errors": [], - "validation_type": "syntax", - } - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: { + "valid": True, + "errors": [], + "validation_type": "syntax", + } + } + ), ) except SyntaxError as e: error_msg = f"Syntax error at line {e.lineno}: {e.msg}" logging.error(f"[{self.name}] {error_msg}") yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: { - "valid": False, - "errors": [error_msg], - "validation_type": "syntax", - "details": str(e), - } - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: { + "valid": False, + "errors": [error_msg], + "validation_type": "syntax", + "details": str(e), + } + } + ), ) except Exception as e: error_msg = f"Unexpected error during syntax validation: {str(e)}" logging.error(f"[{self.name}] {error_msg}") yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: { - "valid": False, - "errors": [error_msg], - "validation_type": "syntax", - } - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: { + "valid": False, + "errors": [error_msg], + "validation_type": "syntax", + } + } + ), ) @@ -361,91 +375,98 @@ def __init__(self, name: str, input_key: str, output_key: str): self.output_key = output_key async def _run_async_impl( - self, ctx: InvocationContext) -> AsyncGenerator[Event, None]: + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: test_file_path = ctx.session.state.get(self.input_key, "") if not test_file_path or not os.path.exists(test_file_path): error_msg = f"Test file not found at {test_file_path}" logging.error(f"[{self.name}] {error_msg}") yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: { - "valid": False, - "errors": [error_msg], - "validation_type": "import", - } - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: { + "valid": False, + "errors": [error_msg], + "validation_type": "import", + } + } + ), ) return try: result = subprocess.run( - ["python", "-m", "py_compile", test_file_path], - capture_output=True, - text=True, - cwd=os.path.dirname(test_file_path), - timeout=30, + ["python", "-m", "py_compile", test_file_path], + capture_output=True, + text=True, + cwd=os.path.dirname(test_file_path), + timeout=30, ) if result.returncode == 0: logging.info( - f"[{self.name}] Import validation passed for {test_file_path}") + f"[{self.name}] Import validation passed for {test_file_path}" + ) yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: { - "valid": True, - "errors": [], - "validation_type": "import", - } - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: { + "valid": True, + "errors": [], + "validation_type": "import", + } + } + ), ) else: error_msg = f"Import validation failed: {result.stderr}" logging.error(f"[{self.name}] {error_msg}") yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: { - "valid": False, - "errors": [error_msg], - "validation_type": "import", - "details": result.stderr, - } - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: { + "valid": False, + "errors": [error_msg], + "validation_type": "import", + "details": result.stderr, + } + } + ), ) except subprocess.TimeoutExpired: error_msg = "Import validation timed out" logging.error(f"[{self.name}] {error_msg}") yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: { - "valid": False, - "errors": [error_msg], - "validation_type": "import", - } - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: { + "valid": False, + "errors": [error_msg], + "validation_type": "import", + } + } + ), ) except Exception as e: error_msg = f"Unexpected error during import validation: {str(e)}" logging.error(f"[{self.name}] {error_msg}") yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: { - "valid": False, - "errors": [error_msg], - "validation_type": "import", - } - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: { + "valid": False, + "errors": [error_msg], + "validation_type": "import", + } + } + ), ) @@ -461,126 +482,128 @@ def __init__(self, name: str, input_key: str, output_key: str): self.output_key = output_key async def _run_async_impl( - self, ctx: InvocationContext) -> AsyncGenerator[Event, None]: + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: test_file_path = ctx.session.state.get(self.input_key, "") if not test_file_path or not os.path.exists(test_file_path): error_msg = f"Test file not found at {test_file_path}" logging.error(f"[{self.name}] {error_msg}") yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: { - "valid": False, - "errors": [error_msg], - "validation_type": "structure", - } - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: { + "valid": False, + "errors": [error_msg], + "validation_type": "structure", + } + } + ), ) return try: collect_result = subprocess.run( - ["pytest", test_file_path, "--collect-only", "-q"], - capture_output=True, - text=True, - cwd=os.path.dirname(test_file_path), - timeout=30, + ["pytest", test_file_path, "--collect-only", "-q"], + capture_output=True, + text=True, + cwd=os.path.dirname(test_file_path), + timeout=30, ) - if ("no tests ran" in collect_result.stdout.lower() or - collect_result.returncode != 0): + if ( + "no tests ran" in collect_result.stdout.lower() + or collect_result.returncode != 0 + ): error_msg = f"No valid pytest tests found or collection failed: {collect_result.stdout}" logging.warning(f"[{self.name}] {error_msg}") yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: { - "valid": - False, - "errors": [error_msg], - "validation_type": - "structure", - "details": - collect_result.stdout + "\n" + - collect_result.stderr, - } - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: { + "valid": False, + "errors": [error_msg], + "validation_type": "structure", + "details": collect_result.stdout + "\n" + collect_result.stderr, + } + } + ), ) return setup_result = subprocess.run( - ["pytest", test_file_path, "--setup-only", "-q"], - capture_output=True, - text=True, - cwd=os.path.dirname(test_file_path), - timeout=30, + ["pytest", test_file_path, "--setup-only", "-q"], + capture_output=True, + text=True, + cwd=os.path.dirname(test_file_path), + timeout=30, ) if setup_result.returncode != 0: error_msg = f"Test setup failed (imports or fixtures broken): {setup_result.stdout}" logging.error(f"[{self.name}] {error_msg}") yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: { - "valid": - False, - "errors": [error_msg], - "validation_type": - "structure", - "details": - setup_result.stdout + "\n" + setup_result.stderr, - } - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: { + "valid": False, + "errors": [error_msg], + "validation_type": "structure", + "details": setup_result.stdout + "\n" + setup_result.stderr, + } + } + ), ) else: logging.info( - f"[{self.name}] Test structure validation passed for {test_file_path}" + f"[{self.name}] Test structure validation passed for {test_file_path}" ) yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: { - "valid": True, - "errors": [], - "validation_type": "structure", - "tests_collected": collect_result.stdout, - } - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: { + "valid": True, + "errors": [], + "validation_type": "structure", + "tests_collected": collect_result.stdout, + } + } + ), ) except subprocess.TimeoutExpired: error_msg = "Test structure validation timed out" logging.error(f"[{self.name}] {error_msg}") yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: { - "valid": False, - "errors": [error_msg], - "validation_type": "structure", - } - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: { + "valid": False, + "errors": [error_msg], + "validation_type": "structure", + } + } + ), ) except Exception as e: error_msg = f"Unexpected error during structure validation: {str(e)}" logging.error(f"[{self.name}] {error_msg}") yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: { - "valid": False, - "errors": [error_msg], - "validation_type": "structure", - } - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: { + "valid": False, + "errors": [error_msg], + "validation_type": "structure", + } + } + ), ) @@ -596,22 +619,24 @@ def __init__(self, name: str, input_key: str, output_key: str): self.output_key = output_key async def _run_async_impl( - self, ctx: InvocationContext) -> AsyncGenerator[Event, None]: + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: test_file_path = ctx.session.state.get(self.input_key, "") if not test_file_path or not os.path.exists(test_file_path): error_msg = f"Test file not found at {test_file_path}" logging.error(f"[{self.name}] {error_msg}") yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: { - "valid": False, - "errors": [error_msg], - "validation_type": "mock_execution", - } - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: { + "valid": False, + "errors": [error_msg], + "validation_type": "mock_execution", + } + } + ), ) return @@ -619,29 +644,33 @@ async def _run_async_impl( with open(test_file_path, "r") as f: test_content = f.read() - has_baseline_ref = any(keyword in test_content.lower() for keyword in [ + has_baseline_ref = any( + keyword in test_content.lower() + for keyword in [ "baseline", "jax_baseline", "reference_impl", "converted_jax", - ]) + ] + ) if not has_baseline_ref: logging.warning( - f"[{self.name}] No baseline reference found in test file. Skipping mock execution validation." + f"[{self.name}] No baseline reference found in test file. Skipping mock execution validation." ) yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: { - "valid": True, - "errors": [], - "validation_type": "mock_execution", - "skipped": True, - "reason": "No baseline reference found in test", - } - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: { + "valid": True, + "errors": [], + "validation_type": "mock_execution", + "skipped": True, + "reason": "No baseline reference found in test", + } + } + ), ) return @@ -657,59 +686,62 @@ async def _run_async_impl( mock_content = mock_prefix + mock_content with tempfile.NamedTemporaryFile( - mode="w", - suffix="_mock_test.py", - delete=False, - dir=os.path.dirname(test_file_path), + mode="w", + suffix="_mock_test.py", + delete=False, + dir=os.path.dirname(test_file_path), ) as tmp_file: tmp_file.write(mock_content) tmp_test_path = tmp_file.name try: result = subprocess.run( - ["pytest", tmp_test_path, "-v", "--tb=short", "--maxfail=3"], - capture_output=True, - text=True, - cwd=os.path.dirname(test_file_path), - timeout=60, + ["pytest", tmp_test_path, "-v", "--tb=short", "--maxfail=3"], + capture_output=True, + text=True, + cwd=os.path.dirname(test_file_path), + timeout=60, ) if result.returncode == 0: logging.info(f"[{self.name}] Mock execution validation passed") yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: { - "valid": - True, - "errors": [], - "validation_type": - "mock_execution", - "tests_passed": - True, - "output_summary": (result.stdout[-500:] if len( - result.stdout) > 500 else result.stdout), - } - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: { + "valid": True, + "errors": [], + "validation_type": "mock_execution", + "tests_passed": True, + "output_summary": ( + result.stdout[-500:] + if len(result.stdout) > 500 + else result.stdout + ), + } + } + ), ) else: error_msg = f"Mock execution failed. Tests may have structural issues.\n{result.stdout}\n{result.stderr}" logging.warning(f"[{self.name}] {error_msg}") yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: { - "valid": - False, - "errors": [error_msg[:1000]], - "validation_type": - "mock_execution", - "details": (result.stdout[-1000:] if len( - result.stdout) > 1000 else result.stdout), - } - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: { + "valid": False, + "errors": [error_msg[:1000]], + "validation_type": "mock_execution", + "details": ( + result.stdout[-1000:] + if len(result.stdout) > 1000 + else result.stdout + ), + } + } + ), ) finally: try: @@ -721,29 +753,31 @@ async def _run_async_impl( error_msg = "Mock test execution timed out after 60 seconds" logging.error(f"[{self.name}] {error_msg}") yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: { - "valid": False, - "errors": [error_msg], - "validation_type": "mock_execution", - } - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: { + "valid": False, + "errors": [error_msg], + "validation_type": "mock_execution", + } + } + ), ) except Exception as e: error_msg = f"Unexpected error during mock execution validation: {str(e)}" logging.error(f"[{self.name}] {error_msg}") yield Event( - author=self.name, - actions=EventActions( - state_delta={ - self.output_key: { - "valid": False, - "errors": [error_msg], - "validation_type": "mock_execution", - } - }), + author=self.name, + actions=EventActions( + state_delta={ + self.output_key: { + "valid": False, + "errors": [error_msg], + "validation_type": "mock_execution", + } + } + ), ) @@ -758,27 +792,28 @@ class TestValidationLoopAgent(BaseAgent): max_retries: int = 3 def __init__( - self, - name: str, - syntax_agent: BaseAgent, - import_agent: BaseAgent, - structure_agent: BaseAgent, - mock_execution_agent: BaseAgent, - fix_agent: BaseAgent, - max_retries: int = 3, + self, + name: str, + syntax_agent: BaseAgent, + import_agent: BaseAgent, + structure_agent: BaseAgent, + mock_execution_agent: BaseAgent, + fix_agent: BaseAgent, + max_retries: int = 3, ): super().__init__( - name=name, - syntax_agent=syntax_agent, - import_agent=import_agent, - structure_agent=structure_agent, - mock_execution_agent=mock_execution_agent, - fix_agent=fix_agent, - max_retries=max_retries, + name=name, + syntax_agent=syntax_agent, + import_agent=import_agent, + structure_agent=structure_agent, + mock_execution_agent=mock_execution_agent, + fix_agent=fix_agent, + max_retries=max_retries, ) async def _run_async_impl( - self, ctx: InvocationContext) -> AsyncGenerator[Event, None]: + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: """Validation loop: validate -> fix -> repeat until valid or max retries.""" test_file_path = ctx.session.state.get("test_file_path", "") @@ -789,19 +824,20 @@ async def _run_async_impl( if not test_file_path: logging.error(f"[{self.name}] No test file path found in state.") yield Event( - author=self.name, - actions=EventActions( - state_delta={ - "validation_loop_status": { - "success": False, - "retries": 0, - "message": "No test file was generated. Cannot validate.", - "syntax_valid": False, - "import_valid": False, - "structure_valid": False, - "mock_execution_valid": False, - } - }), + author=self.name, + actions=EventActions( + state_delta={ + "validation_loop_status": { + "success": False, + "retries": 0, + "message": "No test file was generated. Cannot validate.", + "syntax_valid": False, + "import_valid": False, + "structure_valid": False, + "mock_execution_valid": False, + } + } + ), ) return @@ -809,7 +845,7 @@ async def _run_async_impl( while retry_count < self.max_retries: logging.info( - f"[{self.name}] Validation attempt {retry_count + 1}/{self.max_retries}" + f"[{self.name}] Validation attempt {retry_count + 1}/{self.max_retries}" ) async for event in self.syntax_agent.run_async(ctx): @@ -824,17 +860,25 @@ async def _run_async_impl( async for event in self.mock_execution_agent.run_async(ctx): yield event - syntax_valid = ctx.session.state.get("syntax_validation", - {}).get("valid", False) - import_valid = ctx.session.state.get("import_validation", - {}).get("valid", False) - structure_valid = ctx.session.state.get("structure_validation", - {}).get("valid", False) - mock_execution_valid = ctx.session.state.get("mock_execution_validation", - {}).get("valid", False) - - if (syntax_valid and import_valid and structure_valid and - mock_execution_valid): + syntax_valid = ctx.session.state.get("syntax_validation", {}).get( + "valid", False + ) + import_valid = ctx.session.state.get("import_validation", {}).get( + "valid", False + ) + structure_valid = ctx.session.state.get("structure_validation", {}).get( + "valid", False + ) + mock_execution_valid = ctx.session.state.get( + "mock_execution_validation", {} + ).get("valid", False) + + if ( + syntax_valid + and import_valid + and structure_valid + and mock_execution_valid + ): logging.info(f"[{self.name}] ✓ All validations passed!") if test_file_path and os.path.exists(test_file_path): @@ -842,33 +886,37 @@ async def _run_async_impl( with open(test_file_path, "r") as f: content = f.read() - content = re.sub(r"# (from .+ import .+ as optimized_kernel)", - r"\1", content) + content = re.sub( + r"# (from .+ import .+ as optimized_kernel)", r"\1", content + ) - content = re.sub(r"\noptimized_kernel = base_kernel.*(?=\n)", "", - content) + content = re.sub( + r"\noptimized_kernel = base_kernel.*(?=\n)", "", content + ) with open(test_file_path, "w") as f: f.write(content) logging.info( - f"[{self.name}] Successfully uncommented optimized kernel import" + f"[{self.name}] Successfully uncommented optimized kernel import" ) except Exception as e: logging.warning( - f"[{self.name}] Failed to uncomment kernel import: {e}") + f"[{self.name}] Failed to uncomment kernel import: {e}" + ) yield Event( - author=self.name, - actions=EventActions( - state_delta={ - "validation_loop_status": { - "success": True, - "retries": retry_count, - "message": "Test file validated successfully", - "all_checks_passed": True, - } - }), + author=self.name, + actions=EventActions( + state_delta={ + "validation_loop_status": { + "success": True, + "retries": retry_count, + "message": "Test file validated successfully", + "all_checks_passed": True, + } + } + ), ) return @@ -880,164 +928,160 @@ async def _run_async_impl( else: logging.error(f"[{self.name}] ✗ Max retries reached.") yield Event( - author=self.name, - actions=EventActions( - state_delta={ - "validation_loop_status": { - "success": - False, - "retries": - retry_count, - "message": - f"Test file validation failed after {self.max_retries} attempts", - "syntax_valid": - syntax_valid, - "import_valid": - import_valid, - "structure_valid": - structure_valid, - "mock_execution_valid": - mock_execution_valid, - "all_checks_passed": - False, - } - }), + author=self.name, + actions=EventActions( + state_delta={ + "validation_loop_status": { + "success": False, + "retries": retry_count, + "message": f"Test file validation failed after {self.max_retries} attempts", + "syntax_valid": syntax_valid, + "import_valid": import_valid, + "structure_valid": structure_valid, + "mock_execution_valid": mock_execution_valid, + "all_checks_passed": False, + } + } + ), ) return # Validation Summary Agent validation_summary_agent = CustomLlmAgent( - name="ValidationSummaryAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=thinking_planner, - instruction=validation_summary.PROMPT, - description= - "Summarizes validation results and provides next steps to the user.", + name="ValidationSummaryAgent", + model=MODEL_NAME, + generate_content_config=model_config, + planner=thinking_planner, + instruction=validation_summary.PROMPT, + description="Summarizes validation results and provides next steps to the user.", ) # Test file generation agent generate_test_file_agent = CustomLlmAgent( - name="GenerateTestFileAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=thinking_planner, - instruction=gen_test_file.PROMPT, - description="Generates a comprehensive pytest test file.", - tools=([search_api_tool, filesystem_tool_rw, vertex_ai_rag_tool] - if vertex_ai_rag_tool else [search_api_tool, filesystem_tool_rw]), - after_tool_callback=create_path_saver("test_file_path"), + name="GenerateTestFileAgent", + model=MODEL_NAME, + generate_content_config=model_config, + planner=thinking_planner, + instruction=gen_test_file.PROMPT, + description="Generates a comprehensive pytest test file.", + tools=( + [search_api_tool, filesystem_tool_rw, vertex_ai_rag_tool] + if vertex_ai_rag_tool + else [search_api_tool, filesystem_tool_rw] + ), + after_tool_callback=create_path_saver("test_file_path"), ) # Validation agents syntax_validation_agent = SyntaxValidationAgent( - name="SyntaxValidationAgent", - input_key="test_file_path", - output_key="syntax_validation", + name="SyntaxValidationAgent", + input_key="test_file_path", + output_key="syntax_validation", ) import_validation_agent = ImportValidationAgent( - name="ImportValidationAgent", - input_key="test_file_path", - output_key="import_validation", + name="ImportValidationAgent", + input_key="test_file_path", + output_key="import_validation", ) structure_validation_agent = TestStructureValidationAgent( - name="TestStructureValidationAgent", - input_key="test_file_path", - output_key="structure_validation", + name="TestStructureValidationAgent", + input_key="test_file_path", + output_key="structure_validation", ) mock_execution_validation_agent = MockTestExecutionAgent( - name="MockTestExecutionAgent", - input_key="test_file_path", - output_key="mock_execution_validation", + name="MockTestExecutionAgent", + input_key="test_file_path", + output_key="mock_execution_validation", ) fix_test_script_agent = CustomLlmAgent( - name="FixTestScriptAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=thinking_planner, - instruction=fix_test_script.PROMPT, - description="Fixes validation errors in the generated test file.", - tools=[filesystem_tool_rw], - include_contents="none", + name="FixTestScriptAgent", + model=MODEL_NAME, + generate_content_config=model_config, + planner=thinking_planner, + instruction=fix_test_script.PROMPT, + description="Fixes validation errors in the generated test file.", + tools=[filesystem_tool_rw], + include_contents="none", ) # Validation loop agent validation_loop_agent = TestValidationLoopAgent( - name="TestValidationLoopAgent", - syntax_agent=syntax_validation_agent, - import_agent=import_validation_agent, - structure_agent=structure_validation_agent, - mock_execution_agent=mock_execution_validation_agent, - fix_agent=fix_test_script_agent, - max_retries=3, + name="TestValidationLoopAgent", + syntax_agent=syntax_validation_agent, + import_agent=import_validation_agent, + structure_agent=structure_validation_agent, + mock_execution_agent=mock_execution_validation_agent, + fix_agent=fix_test_script_agent, + max_retries=3, ) # Sequential agent that generates and validates test files validated_test_generation_agent = SequentialAgent( - name="ValidatedTestGenerationAgent", - sub_agents=[ - generate_test_file_agent, - validation_loop_agent, - validation_summary_agent, - ], - description= - "Generates a validated pytest test file with automatic iterative error detection and fixing.", + name="ValidatedTestGenerationAgent", + sub_agents=[ + generate_test_file_agent, + validation_loop_agent, + validation_summary_agent, + ], + description="Generates a validated pytest test file with automatic iterative error detection and fixing.", ) # Test execution agents read_file_for_testing_agent = CustomLlmAgent( - name="ReadFileForTestingAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=thinking_planner, - instruction=read_file_prompt.PROMPT, - description="Reads the test file mentioned by the user.", - tools=[filesystem_tool_rw], - after_tool_callback=create_path_saver("test_file_path"), + name="ReadFileForTestingAgent", + model=MODEL_NAME, + generate_content_config=model_config, + planner=thinking_planner, + instruction=read_file_prompt.PROMPT, + description="Reads the test file mentioned by the user.", + tools=[filesystem_tool_rw], + after_tool_callback=create_path_saver("test_file_path"), ) run_tests_agent = TestRunner( - name="RunTestsAgent", - input_key="test_file_path", - output_key="test_results", + name="RunTestsAgent", + input_key="test_file_path", + output_key="test_results", ) summarize_test_results_agent = CustomLlmAgent( - name="SummarizeTestResultsAgent", - model=MODEL_NAME, - generate_content_config=model_config, - planner=thinking_planner, - instruction=summarize_test_results_prompt.PROMPT, - description="Analyzes pytest test results and provides recommendations.", - tools=([search_api_tool, vertex_ai_rag_tool] - if vertex_ai_rag_tool else [search_api_tool]), - output_key="test_summary", - include_contents="none", + name="SummarizeTestResultsAgent", + model=MODEL_NAME, + generate_content_config=model_config, + planner=thinking_planner, + instruction=summarize_test_results_prompt.PROMPT, + description="Analyzes pytest test results and provides recommendations.", + tools=( + [search_api_tool, vertex_ai_rag_tool] + if vertex_ai_rag_tool + else [search_api_tool] + ), + output_key="test_summary", + include_contents="none", ) unified_test_agent = SequentialAgent( - name="UnifiedTestAgent", - sub_agents=[ - read_file_for_testing_agent, - run_tests_agent, - summarize_test_results_agent, - ], - description= - "Executes the generated pytest test file and provides a comprehensive summary.", + name="UnifiedTestAgent", + sub_agents=[ + read_file_for_testing_agent, + run_tests_agent, + summarize_test_results_agent, + ], + description="Executes the generated pytest test file and provides a comprehensive summary.", ) __all__ = [ - "TestRunner", - "SyntaxValidationAgent", - "ImportValidationAgent", - "TestStructureValidationAgent", - "MockTestExecutionAgent", - "TestValidationLoopAgent", - "validated_test_generation_agent", - "unified_test_agent", + "TestRunner", + "SyntaxValidationAgent", + "ImportValidationAgent", + "TestStructureValidationAgent", + "MockTestExecutionAgent", + "TestValidationLoopAgent", + "validated_test_generation_agent", + "unified_test_agent", ] diff --git a/MaxKernel/hitl_agent/subagents/testing/prompts/__init__.py b/MaxKernel/hitl_agent/subagents/testing/prompts/__init__.py index d62be9a..b399b28 100644 --- a/MaxKernel/hitl_agent/subagents/testing/prompts/__init__.py +++ b/MaxKernel/hitl_agent/subagents/testing/prompts/__init__.py @@ -1,17 +1,17 @@ """Prompts for testing subagent.""" from . import ( - gen_test_file, - fix_test_script, - validation_summary, - summarize_test_results_prompt, - read_file_prompt, + fix_test_script, + gen_test_file, + read_file_prompt, + summarize_test_results_prompt, + validation_summary, ) __all__ = [ - 'gen_test_file', - 'fix_test_script', - 'validation_summary', - 'summarize_test_results_prompt', - 'read_file_prompt', + "gen_test_file", + "fix_test_script", + "validation_summary", + "summarize_test_results_prompt", + "read_file_prompt", ] diff --git a/MaxKernel/hitl_agent/tests/conftest.py b/MaxKernel/hitl_agent/tests/conftest.py index 228073b..e0bcca3 100644 --- a/MaxKernel/hitl_agent/tests/conftest.py +++ b/MaxKernel/hitl_agent/tests/conftest.py @@ -1,16 +1,16 @@ """Shared fixtures and test utilities for HITL agent tests.""" -import pytest import tempfile -from typing import Optional, Callable, AsyncGenerator +from typing import AsyncGenerator, Callable, Optional from unittest.mock import Mock -from pydantic import PrivateAttr -from google.adk.sessions import Session +import pytest from google.adk.agents import BaseAgent from google.adk.agents.invocation_context import InvocationContext from google.adk.events import Event, EventActions +from google.adk.sessions import Session from google.genai.types import Content, Part +from pydantic import PrivateAttr @pytest.fixture @@ -23,9 +23,9 @@ def temp_workdir(): @pytest.fixture def mock_session(temp_workdir): """Create a proper ADK session with initialized state.""" - session = Session(id="test-session-123", - user_id="test-user", - appName="test-hitl-agent") + session = Session( + id="test-session-123", user_id="test-user", appName="test-hitl-agent" + ) # Initialize state with test values session.state["workdir"] = temp_workdir session.state["tpu_version"] = "v5e" @@ -123,6 +123,7 @@ def kernel(x_ref, y_ref, o_ref): class MockAgent(BaseAgent): """Generic mock agent for tests.""" + _run_func: Optional[Callable] = PrivateAttr(default=None) _call_count: int = PrivateAttr(default=0) @@ -147,8 +148,10 @@ async def _run_async_impl(self, ctx): async for event in self._run_func(ctx): yield event else: - yield Event(author=self.name, - content=Content(parts=[Part(text=f"Mock {self.name} run")])) + yield Event( + author=self.name, + content=Content(parts=[Part(text=f"Mock {self.name} run")]), + ) class MockFixAgent(MockAgent): @@ -160,6 +163,7 @@ def __init__(self, name: str, fix_func=None): class MockCompilationChecker(BaseAgent): """Mock compilation checker that returns predefined results.""" + _output_key: str = PrivateAttr() _results: list = PrivateAttr() _call_count: int = PrivateAttr(default=0) @@ -170,14 +174,16 @@ def __init__(self, name, output_key, results): self._results = results if isinstance(results, list) else [results] self._call_count = 0 - async def run_async(self, - ctx: InvocationContext) -> AsyncGenerator[Event, None]: + async def run_async( + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: """Override run_async to directly call _run_async_impl for tests.""" async for event in self._run_async_impl(ctx): yield event async def _run_async_impl( - self, ctx: InvocationContext) -> AsyncGenerator[Event, None]: + self, ctx: InvocationContext + ) -> AsyncGenerator[Event, None]: if self._call_count < len(self._results): result = self._results[self._call_count] else: @@ -188,8 +194,10 @@ async def _run_async_impl( # Update state directly so the loop sees it immediately ctx.session.state[self._output_key] = result - yield Event(author=self.name, - actions=EventActions(state_delta={self._output_key: result})) + yield Event( + author=self.name, + actions=EventActions(state_delta={self._output_key: result}), + ) class CompilationCheckerWrapper(BaseAgent): diff --git a/MaxKernel/hitl_agent/tests/test_compilation_validation_loop.py b/MaxKernel/hitl_agent/tests/test_compilation_validation_loop.py index 623fda4..74f437d 100644 --- a/MaxKernel/hitl_agent/tests/test_compilation_validation_loop.py +++ b/MaxKernel/hitl_agent/tests/test_compilation_validation_loop.py @@ -12,15 +12,19 @@ pytest -m "not integration" # Skip these tests """ -import pytest import os + +import pytest +from conftest import CompilationCheckerWrapper, MockFixAgent from google.adk.events import Event from google.genai.types import Content, Part -from hitl_agent.subagents.kernel_writing.kernel_compilation import KernelCompilationChecker from hitl_agent.subagents.kernel_writing.agent import ( - KernelCompilationValidationLoop as _KernelCompilationValidationLoop,) -from conftest import MockFixAgent, CompilationCheckerWrapper + KernelCompilationValidationLoop as _KernelCompilationValidationLoop, +) +from hitl_agent.subagents.kernel_writing.kernel_compilation import ( + KernelCompilationChecker, +) # Test-only wrapper that bypasses BaseAgent.run_async scaffolding for simplified test setup @@ -37,8 +41,9 @@ class TestCompilationValidationLoop: @pytest.mark.asyncio @pytest.mark.integration - async def test_compilation_success(self, mock_invocation_context, - temp_workdir, kernel_code_valid): + async def test_compilation_success( + self, mock_invocation_context, temp_workdir, kernel_code_valid + ): """Test compilation validation - successful case.""" # Setup: Create a valid kernel file kernel_path = os.path.join(temp_workdir, "test_kernel.py") @@ -51,10 +56,10 @@ async def test_compilation_success(self, mock_invocation_context, # Use real KernelCompilationChecker with auto_manage_servers compilation_checker = KernelCompilationChecker( - name="RealCompilationChecker", - input_key="kernel_code", - output_key="compilation_results", - auto_manage_servers=True # Let the checker manage servers + name="RealCompilationChecker", + input_key="kernel_code", + output_key="compilation_results", + auto_manage_servers=True, # Let the checker manage servers ) # Wrap it to call _run_async_impl directly wrapped_checker = CompilationCheckerWrapper(compilation_checker) @@ -70,46 +75,56 @@ async def mock_fix(ctx): # Create validation loop validation_loop = KernelCompilationValidationLoop( - name="TestValidationLoop", - compilation_checker=wrapped_checker, - fix_agent=mock_fix_agent, - max_retries=4) + name="TestValidationLoop", + compilation_checker=wrapped_checker, + fix_agent=mock_fix_agent, + max_retries=4, + ) # Run validation events = [] async for event in validation_loop._run_async_impl(mock_invocation_context): events.append(event) # Apply state_delta if present - if hasattr(event, 'actions') and event.actions and hasattr( - event.actions, 'state_delta'): + if ( + hasattr(event, "actions") + and event.actions + and hasattr(event.actions, "state_delta") + ): if event.actions.state_delta: mock_invocation_context.session.state.update( - event.actions.state_delta) + event.actions.state_delta + ) # Assertions assert len(events) > 0, "Should have received events" # Check compilation succeeded compilation_result = mock_invocation_context.session.state.get( - "compilation_results") - assert compilation_result == "Success", f"Expected Success but got: {compilation_result}" + "compilation_results" + ) + assert compilation_result == "Success", ( + f"Expected Success but got: {compilation_result}" + ) # Fix agent should not have been called - assert fix_called[ - 0] is False, "Fix agent should not be called for valid code" + assert fix_called[0] is False, ( + "Fix agent should not be called for valid code" + ) # Check final status status = mock_invocation_context.session.state.get( - "kernel_compilation_status") + "kernel_compilation_status" + ) assert status is not None assert status["success"] is True assert status["retries"] == 0 @pytest.mark.asyncio @pytest.mark.integration - async def test_compilation_syntax_error(self, mock_invocation_context, - temp_workdir, - kernel_code_syntax_error): + async def test_compilation_syntax_error( + self, mock_invocation_context, temp_workdir, kernel_code_syntax_error + ): """Test compilation validation - syntax error case.""" # Setup: Create a kernel file with syntax error kernel_path = os.path.join(temp_workdir, "test_kernel.py") @@ -118,15 +133,17 @@ async def test_compilation_syntax_error(self, mock_invocation_context, # Set up state mock_invocation_context.session.state["optimized_kernel_path"] = kernel_path - mock_invocation_context.session.state[ - "kernel_code"] = kernel_code_syntax_error + mock_invocation_context.session.state["kernel_code"] = ( + kernel_code_syntax_error + ) # Use real KernelCompilationChecker compilation_checker = KernelCompilationChecker( - name="RealCompilationChecker", - input_key="kernel_code", - output_key="compilation_results", - auto_manage_servers=True) + name="RealCompilationChecker", + input_key="kernel_code", + output_key="compilation_results", + auto_manage_servers=True, + ) # Wrap it to call _run_async_impl directly wrapped_checker = CompilationCheckerWrapper(compilation_checker) @@ -136,17 +153,18 @@ async def test_compilation_syntax_error(self, mock_invocation_context, async def mock_fix(ctx): fix_call_count[0] += 1 # Don't actually fix the code, just yield an event - yield Event(author="Fix", - content=Content(parts=[Part(text="Attempted fix")])) + yield Event( + author="Fix", content=Content(parts=[Part(text="Attempted fix")]) + ) mock_fix_agent = MockFixAgent(name="MockFixAgent", fix_func=mock_fix) # Create validation loop with limited retries validation_loop = KernelCompilationValidationLoop( - name="TestValidationLoop", - compilation_checker=wrapped_checker, - fix_agent=mock_fix_agent, - max_retries=2 # Limit retries for faster test + name="TestValidationLoop", + compilation_checker=wrapped_checker, + fix_agent=mock_fix_agent, + max_retries=2, # Limit retries for faster test ) # Run validation @@ -154,38 +172,49 @@ async def mock_fix(ctx): async for event in validation_loop._run_async_impl(mock_invocation_context): events.append(event) # Apply state_delta if present - if hasattr(event, 'actions') and event.actions and hasattr( - event.actions, 'state_delta'): + if ( + hasattr(event, "actions") + and event.actions + and hasattr(event.actions, "state_delta") + ): if event.actions.state_delta: mock_invocation_context.session.state.update( - event.actions.state_delta) + event.actions.state_delta + ) # Assertions assert len(events) > 0, "Should have received events" # Check compilation failed compilation_result = mock_invocation_context.session.state.get( - "compilation_results") + "compilation_results" + ) assert compilation_result != "Success", "Should have failed compilation" - assert "syntax" in compilation_result.lower( - ) or "error" in compilation_result.lower() + assert ( + "syntax" in compilation_result.lower() + or "error" in compilation_result.lower() + ) # Fix agent should have been called multiple times assert fix_call_count[0] > 0, "Fix agent should have been called" # Check final status shows failure status = mock_invocation_context.session.state.get( - "kernel_compilation_status") + "kernel_compilation_status" + ) assert status is not None assert status["success"] is False assert status["retries"] == 1 # Should have retried once before giving up @pytest.mark.asyncio @pytest.mark.integration - async def test_compilation_retry_success(self, mock_invocation_context, - temp_workdir, - kernel_code_syntax_error, - kernel_code_valid): + async def test_compilation_retry_success( + self, + mock_invocation_context, + temp_workdir, + kernel_code_syntax_error, + kernel_code_valid, + ): """Test compilation validation - retry success case (initially wrong, then fixed).""" # Setup: Create a kernel file with syntax error initially kernel_path = os.path.join(temp_workdir, "test_kernel.py") @@ -194,15 +223,17 @@ async def test_compilation_retry_success(self, mock_invocation_context, # Set up state mock_invocation_context.session.state["optimized_kernel_path"] = kernel_path - mock_invocation_context.session.state[ - "kernel_code"] = kernel_code_syntax_error + mock_invocation_context.session.state["kernel_code"] = ( + kernel_code_syntax_error + ) # Use real KernelCompilationChecker compilation_checker = KernelCompilationChecker( - name="RealCompilationChecker", - input_key="kernel_code", - output_key="compilation_results", - auto_manage_servers=True) + name="RealCompilationChecker", + input_key="kernel_code", + output_key="compilation_results", + auto_manage_servers=True, + ) # Wrap it to call _run_async_impl directly wrapped_checker = CompilationCheckerWrapper(compilation_checker) @@ -216,35 +247,42 @@ async def mock_fix(ctx): with open(kernel_path, "w") as f: f.write(kernel_code_valid) - yield Event(author="Fix", - content=Content(parts=[Part(text="Fixed the code")])) + yield Event( + author="Fix", content=Content(parts=[Part(text="Fixed the code")]) + ) mock_fix_agent = MockFixAgent(name="MockFixAgent", fix_func=mock_fix) # Create validation loop validation_loop = KernelCompilationValidationLoop( - name="TestValidationLoop", - compilation_checker=wrapped_checker, - fix_agent=mock_fix_agent, - max_retries=2) + name="TestValidationLoop", + compilation_checker=wrapped_checker, + fix_agent=mock_fix_agent, + max_retries=2, + ) # Run validation events = [] async for event in validation_loop._run_async_impl(mock_invocation_context): events.append(event) # Apply state_delta if present - if hasattr(event, 'actions') and event.actions and hasattr( - event.actions, 'state_delta'): + if ( + hasattr(event, "actions") + and event.actions + and hasattr(event.actions, "state_delta") + ): if event.actions.state_delta: mock_invocation_context.session.state.update( - event.actions.state_delta) + event.actions.state_delta + ) # Assertions assert len(events) > 0 # Check compilation succeeded compilation_result = mock_invocation_context.session.state.get( - "compilation_results") + "compilation_results" + ) assert compilation_result == "Success" # Fix agent should have been called once @@ -252,7 +290,8 @@ async def mock_fix(ctx): # Check final status status = mock_invocation_context.session.state.get( - "kernel_compilation_status") + "kernel_compilation_status" + ) assert status is not None assert status["success"] is True assert status["retries"] == 1 diff --git a/MaxKernel/hitl_agent/tests/test_validate_kernel_compilation_agent.py b/MaxKernel/hitl_agent/tests/test_validate_kernel_compilation_agent.py index 4c2a8b8..d881876 100644 --- a/MaxKernel/hitl_agent/tests/test_validate_kernel_compilation_agent.py +++ b/MaxKernel/hitl_agent/tests/test_validate_kernel_compilation_agent.py @@ -10,16 +10,19 @@ For integration tests with real TPU compilation, see test_compilation_validation_loop.py. """ -import pytest import os + +import pytest +from conftest import MockAgent, MockCompilationChecker, MockFixAgent from google.adk.events import Event from google.genai.types import Content, Part from hitl_agent.subagents.kernel_writing.agent import ( - KernelCompilationValidationLoop as _KernelCompilationValidationLoop, - ValidateKernelCompilationAgent as _ValidateKernelCompilationAgent, + KernelCompilationValidationLoop as _KernelCompilationValidationLoop, +) +from hitl_agent.subagents.kernel_writing.agent import ( + ValidateKernelCompilationAgent as _ValidateKernelCompilationAgent, ) -from conftest import MockAgent, MockFixAgent, MockCompilationChecker # Test-only wrappers that bypass BaseAgent.run_async scaffolding for simplified test setup @@ -44,9 +47,9 @@ class TestValidateKernelCompilationAgent: @pytest.mark.asyncio @pytest.mark.unit - async def test_full_agent_compilation_success(self, mock_invocation_context, - temp_workdir, - kernel_code_valid): + async def test_full_agent_compilation_success( + self, mock_invocation_context, temp_workdir, kernel_code_valid + ): """Test compilation validation - successful case.""" # Setup: Create a valid kernel file kernel_path = os.path.join(temp_workdir, "test_kernel.py") @@ -59,9 +62,10 @@ async def test_full_agent_compilation_success(self, mock_invocation_context, # Use Mock KernelCompilationChecker compilation_checker = MockCompilationChecker( - name="MockCompilationChecker", - output_key="compilation_results", - results="Success") + name="MockCompilationChecker", + output_key="compilation_results", + results="Success", + ) # Mock fix agent (should not be called for valid code) fix_called = [False] @@ -74,10 +78,11 @@ async def mock_fix(ctx): # Create validation loop validation_loop = KernelCompilationValidationLoop( - name="TestValidationLoop", - compilation_checker=compilation_checker, - fix_agent=mock_fix_agent, - max_retries=4) + name="TestValidationLoop", + compilation_checker=compilation_checker, + fix_agent=mock_fix_agent, + max_retries=4, + ) # Create other mock agents mock_read_file = MockAgent("MockReadFile") @@ -86,38 +91,48 @@ async def mock_fix(ctx): # Create ValidateKernelCompilationAgent validate_agent = ValidateKernelCompilationAgent( - name="TestValidateAgent", - read_file_agent=mock_read_file, - validation_loop_agent=validation_loop, - cleanup_agent=mock_cleanup, - summary_agent=mock_summary) + name="TestValidateAgent", + read_file_agent=mock_read_file, + validation_loop_agent=validation_loop, + cleanup_agent=mock_cleanup, + summary_agent=mock_summary, + ) # Run validation events = [] async for event in validate_agent._run_async_impl(mock_invocation_context): events.append(event) # Apply state_delta if present - if hasattr(event, 'actions') and event.actions and hasattr( - event.actions, 'state_delta'): + if ( + hasattr(event, "actions") + and event.actions + and hasattr(event.actions, "state_delta") + ): if event.actions.state_delta: mock_invocation_context.session.state.update( - event.actions.state_delta) + event.actions.state_delta + ) # Assertions assert len(events) > 0, "Should have received events" # Check compilation succeeded compilation_result = mock_invocation_context.session.state.get( - "compilation_results") - assert compilation_result == "Success", f"Expected Success but got: {compilation_result}" + "compilation_results" + ) + assert compilation_result == "Success", ( + f"Expected Success but got: {compilation_result}" + ) # Fix agent should not have been called - assert fix_called[ - 0] is False, "Fix agent should not be called for valid code" + assert fix_called[0] is False, ( + "Fix agent should not be called for valid code" + ) # Check final status status = mock_invocation_context.session.state.get( - "kernel_compilation_status") + "kernel_compilation_status" + ) assert status is not None assert status["success"] is True assert status["retries"] == 0 @@ -130,9 +145,9 @@ async def mock_fix(ctx): @pytest.mark.asyncio @pytest.mark.unit - async def test_full_agent_compilation_failure(self, mock_invocation_context, - temp_workdir, - kernel_code_syntax_error): + async def test_full_agent_compilation_failure( + self, mock_invocation_context, temp_workdir, kernel_code_syntax_error + ): """Test full agent with compilation failure and retry exhaustion.""" # Setup: Create a kernel file with syntax error kernel_path = os.path.join(temp_workdir, "test_kernel.py") @@ -141,14 +156,16 @@ async def test_full_agent_compilation_failure(self, mock_invocation_context, # Set up state mock_invocation_context.session.state["optimized_kernel_path"] = kernel_path - mock_invocation_context.session.state[ - "kernel_code"] = kernel_code_syntax_error + mock_invocation_context.session.state["kernel_code"] = ( + kernel_code_syntax_error + ) # Use Mock KernelCompilationChecker compilation_checker = MockCompilationChecker( - name="MockCompilationChecker", - output_key="compilation_results", - results="Syntax Error: missing parenthesis") + name="MockCompilationChecker", + output_key="compilation_results", + results="Syntax Error: missing parenthesis", + ) # Mock fix agent that doesn't actually fix (to test retry exhaustion) fix_call_count = [0] @@ -156,17 +173,18 @@ async def test_full_agent_compilation_failure(self, mock_invocation_context, async def mock_fix(ctx): fix_call_count[0] += 1 # Don't actually fix the code, just yield an event - yield Event(author="Fix", - content=Content(parts=[Part(text="Attempted fix")])) + yield Event( + author="Fix", content=Content(parts=[Part(text="Attempted fix")]) + ) mock_fix_agent = MockFixAgent(name="MockFixAgent", fix_func=mock_fix) # Create validation loop with limited retries validation_loop = KernelCompilationValidationLoop( - name="TestValidationLoop", - compilation_checker=compilation_checker, - fix_agent=mock_fix_agent, - max_retries=2 # Limit retries for faster test + name="TestValidationLoop", + compilation_checker=compilation_checker, + fix_agent=mock_fix_agent, + max_retries=2, # Limit retries for faster test ) # Create other mock agents @@ -176,39 +194,48 @@ async def mock_fix(ctx): # Create ValidateKernelCompilationAgent validate_agent = ValidateKernelCompilationAgent( - name="TestValidateAgent", - read_file_agent=mock_read_file, - validation_loop_agent=validation_loop, - cleanup_agent=mock_cleanup, - summary_agent=mock_summary) + name="TestValidateAgent", + read_file_agent=mock_read_file, + validation_loop_agent=validation_loop, + cleanup_agent=mock_cleanup, + summary_agent=mock_summary, + ) # Run validation events = [] async for event in validate_agent._run_async_impl(mock_invocation_context): events.append(event) # Apply state_delta if present - if hasattr(event, 'actions') and event.actions and hasattr( - event.actions, 'state_delta'): + if ( + hasattr(event, "actions") + and event.actions + and hasattr(event.actions, "state_delta") + ): if event.actions.state_delta: mock_invocation_context.session.state.update( - event.actions.state_delta) + event.actions.state_delta + ) # Assertions assert len(events) > 0, "Should have received events" # Check compilation failed compilation_result = mock_invocation_context.session.state.get( - "compilation_results") + "compilation_results" + ) assert compilation_result != "Success", "Should have failed compilation" - assert "syntax" in compilation_result.lower( - ) or "error" in compilation_result.lower() + assert ( + "syntax" in compilation_result.lower() + or "error" in compilation_result.lower() + ) # Fix agent should have been called multiple times assert fix_call_count[0] > 0, "Fix agent should have been called" # Check final status shows failure status = mock_invocation_context.session.state.get( - "kernel_compilation_status") + "kernel_compilation_status" + ) assert status is not None assert status["success"] is False assert status["retries"] == 1 # Should have retried once before giving up @@ -219,10 +246,13 @@ async def mock_fix(ctx): @pytest.mark.asyncio @pytest.mark.unit - async def test_full_agent_retry_success(self, mock_invocation_context, - temp_workdir, - kernel_code_syntax_error, - kernel_code_valid): + async def test_full_agent_retry_success( + self, + mock_invocation_context, + temp_workdir, + kernel_code_syntax_error, + kernel_code_valid, + ): """Test full agent with retry success (initially fails, then succeeds after fix).""" # Setup: Create a kernel file with syntax error initially kernel_path = os.path.join(temp_workdir, "test_kernel.py") @@ -231,14 +261,16 @@ async def test_full_agent_retry_success(self, mock_invocation_context, # Set up state mock_invocation_context.session.state["optimized_kernel_path"] = kernel_path - mock_invocation_context.session.state[ - "kernel_code"] = kernel_code_syntax_error + mock_invocation_context.session.state["kernel_code"] = ( + kernel_code_syntax_error + ) # Use Mock KernelCompilationChecker with sequential results compilation_checker = MockCompilationChecker( - name="MockCompilationChecker", - output_key="compilation_results", - results=["Syntax Error", "Success"]) + name="MockCompilationChecker", + output_key="compilation_results", + results=["Syntax Error", "Success"], + ) # Mock fix agent that fixes the code fix_call_count = [0] @@ -250,17 +282,19 @@ async def mock_fix(ctx): with open(kernel_path, "w") as f: f.write(kernel_code_valid) - yield Event(author="Fix", - content=Content(parts=[Part(text="Fixed the code")])) + yield Event( + author="Fix", content=Content(parts=[Part(text="Fixed the code")]) + ) mock_fix_agent = MockFixAgent(name="MockFixAgent", fix_func=mock_fix) # Create validation loop validation_loop = KernelCompilationValidationLoop( - name="TestValidationLoop", - compilation_checker=compilation_checker, - fix_agent=mock_fix_agent, - max_retries=2) + name="TestValidationLoop", + compilation_checker=compilation_checker, + fix_agent=mock_fix_agent, + max_retries=2, + ) # Create other mock agents mock_read_file = MockAgent("MockReadFile") @@ -269,29 +303,35 @@ async def mock_fix(ctx): # Create ValidateKernelCompilationAgent validate_agent = ValidateKernelCompilationAgent( - name="TestValidateAgent", - read_file_agent=mock_read_file, - validation_loop_agent=validation_loop, - cleanup_agent=mock_cleanup, - summary_agent=mock_summary) + name="TestValidateAgent", + read_file_agent=mock_read_file, + validation_loop_agent=validation_loop, + cleanup_agent=mock_cleanup, + summary_agent=mock_summary, + ) # Run validation events = [] async for event in validate_agent._run_async_impl(mock_invocation_context): events.append(event) # Apply state_delta if present - if hasattr(event, 'actions') and event.actions and hasattr( - event.actions, 'state_delta'): + if ( + hasattr(event, "actions") + and event.actions + and hasattr(event.actions, "state_delta") + ): if event.actions.state_delta: mock_invocation_context.session.state.update( - event.actions.state_delta) + event.actions.state_delta + ) # Assertions assert len(events) > 0 # Check compilation succeeded compilation_result = mock_invocation_context.session.state.get( - "compilation_results") + "compilation_results" + ) assert compilation_result == "Success" # Fix agent should have been called once @@ -299,7 +339,8 @@ async def mock_fix(ctx): # Check final status status = mock_invocation_context.session.state.get( - "kernel_compilation_status") + "kernel_compilation_status" + ) assert status is not None assert status["success"] is True assert status["retries"] == 1 diff --git a/MaxKernel/hitl_agent/tools/analyze_profile.py b/MaxKernel/hitl_agent/tools/analyze_profile.py index ca22742..51eba0d 100644 --- a/MaxKernel/hitl_agent/tools/analyze_profile.py +++ b/MaxKernel/hitl_agent/tools/analyze_profile.py @@ -1,12 +1,14 @@ import json + from xprof.convert import raw_to_tool_data def analyze_trace(path): - tool_data_result, _ = raw_to_tool_data.xspace_to_tool_data([path], - 'trace_viewer', {}) + tool_data_result, _ = raw_to_tool_data.xspace_to_tool_data( + [path], "trace_viewer", {} + ) trace_data = json.loads(tool_data_result) - events = trace_data.get('traceEvents', []) + events = trace_data.get("traceEvents", []) pid = None start_last = None end_last = None @@ -15,32 +17,35 @@ def analyze_trace(path): events_for_tpu_0 = [] jit_computation_events = [] for event in events: - if 'args' in event and event['args'].get('name', None) == '/device:TPU:0': - pid = event.get('pid', -1) - if event.get('pid', -1) == pid: + if "args" in event and event["args"].get("name", None) == "/device:TPU:0": + pid = event.get("pid", -1) + if event.get("pid", -1) == pid: events_for_tpu_0.append(event) - if 'jit_computation' in event.get('name', None): + if "jit_computation" in event.get("name", None): jit_computation_events.append(event) - start_last = jit_computation_events[-2]['ts'] + jit_computation_events[-2][ - 'dur'] - end_last = jit_computation_events[-1]['ts'] + jit_computation_events[-1]['dur'] + start_last = ( + jit_computation_events[-2]["ts"] + jit_computation_events[-2]["dur"] + ) + end_last = ( + jit_computation_events[-1]["ts"] + jit_computation_events[-1]["dur"] + ) for event in events_for_tpu_0: - if 'dur' in event: - if event['ts'] >= start_last and (event['ts'] + event['dur']) <= end_last: - if 'SyncWait' in event.get('name', None): - sync_wait_total += event['dur'] + if "dur" in event: + if event["ts"] >= start_last and (event["ts"] + event["dur"]) <= end_last: + if "SyncWait" in event.get("name", None): + sync_wait_total += event["dur"] total_computation_time = end_last - start_last if total_computation_time > 0: ratio = sync_wait_total / total_computation_time print( - f"We see that kernel spends {ratio*100:.4f}% waiting for synchronization and {(1 - ratio)*100:.4f}% computing." + f"We see that kernel spends {ratio * 100:.4f}% waiting for synchronization and {(1 - ratio) * 100:.4f}% computing." ) return ratio if __name__ == "__main__": - path = '/tmp/profile-data/plugins/profile/2025_09_17_23_21_59/t1v-n-ab373229-w-0.xplane.pb' + path = "/tmp/profile-data/plugins/profile/2025_09_17_23_21_59/t1v-n-ab373229-w-0.xplane.pb" analyze_trace(path) diff --git a/MaxKernel/hitl_agent/tools/api_rag/get_apis.py b/MaxKernel/hitl_agent/tools/api_rag/get_apis.py index 337e11a..bf3ce65 100644 --- a/MaxKernel/hitl_agent/tools/api_rag/get_apis.py +++ b/MaxKernel/hitl_agent/tools/api_rag/get_apis.py @@ -2,8 +2,9 @@ import argparse import importlib import inspect -import textwrap import pkgutil # <-- Switched to the more robust pkgutil for discovery +import textwrap + from docstring_parser import parse as parse_docstring # --- API Discovery (Now using pkgutil for reliability) --- @@ -11,24 +12,24 @@ def discover_apis(api_prefix, recursive): """ - Discovers APIs. If recursive, performs a deep walk of all submodules - using the robust pkgutil method. - """ + Discovers APIs. If recursive, performs a deep walk of all submodules + using the robust pkgutil method. + """ if not recursive: try: resolve_api(api_prefix) return [api_prefix] except ImportError as e: print( - f"❌ Error: Could not resolve the single API '{api_prefix}'. Details: {e}" + f"❌ Error: Could not resolve the single API '{api_prefix}'. Details: {e}" ) return [] print( - f"🔎 Deep-recursively discovering APIs under '{api_prefix}' using pkgutil..." + f"🔎 Deep-recursively discovering APIs under '{api_prefix}' using pkgutil..." ) found_apis = set() - root_package = api_prefix.split('.')[0] + root_package = api_prefix.split(".")[0] try: # Import the top-level package to get its filesystem path for pkgutil @@ -37,7 +38,7 @@ def discover_apis(api_prefix, recursive): search_path = prefix_module.__path__ except (ImportError, AttributeError): print( - f"❌ Error: '{api_prefix}' is not a valid package or could not be found." + f"❌ Error: '{api_prefix}' is not a valid package or could not be found." ) return [] @@ -45,8 +46,9 @@ def discover_apis(api_prefix, recursive): modules_to_inspect = {api_prefix} # pkgutil.walk_packages finds all submodules recursively - for module_info in pkgutil.walk_packages(path=search_path, - prefix=api_prefix + '.'): + for module_info in pkgutil.walk_packages( + path=search_path, prefix=api_prefix + "." + ): modules_to_inspect.add(module_info.name) # Now, iterate through the complete list of discovered modules @@ -54,14 +56,16 @@ def discover_apis(api_prefix, recursive): try: module = importlib.import_module(module_name) for member_name, member in inspect.getmembers(module): - if member_name.startswith('_'): + if member_name.startswith("_"): continue # Add functions or classes that belong to the root package - if (inspect.isfunction(member) or inspect.isclass(member)): - if hasattr(member, '__module__' - ) and member.__module__ and member.__module__.startswith( - root_package): + if inspect.isfunction(member) or inspect.isclass(member): + if ( + hasattr(member, "__module__") + and member.__module__ + and member.__module__.startswith(root_package) + ): full_api_path = f"{module_name}.{member_name}" found_apis.add(full_api_path) except Exception: @@ -70,7 +74,7 @@ def discover_apis(api_prefix, recursive): if not found_apis: print( - f"⚠️ Warning: No public functions or classes found under '{api_prefix}'." + f"⚠️ Warning: No public functions or classes found under '{api_prefix}'." ) else: print(f"✅ Found {len(found_apis)} APIs to document.") @@ -113,9 +117,10 @@ def get_attributes(obj): """Gets the public attributes of a class.""" if inspect.isclass(obj): return [ - name for name, _ in inspect.getmembers(obj) - if not name.startswith("_") and - not inspect.isroutine(getattr(obj, name, None)) + name + for name, _ in inspect.getmembers(obj) + if not name.startswith("_") + and not inspect.isroutine(getattr(obj, name, None)) ] return [] @@ -131,8 +136,11 @@ def format_docstring_sections(doc): returns = parsed.returns examples = parsed.examples - param_strs = [f" - **{p.arg_name}**: {p.description}" for p in parameters - ] if parameters else [] + param_strs = ( + [f" - **{p.arg_name}**: {p.description}" for p in parameters] + if parameters + else [] + ) return_str = f"{returns.description}" if returns else "" example_str = "\n".join(ex.description for ex in examples) if examples else "" return description, param_strs, return_str, example_str @@ -140,9 +148,9 @@ def format_docstring_sections(doc): def generate_definition(api_str): """ - Resolves a JAX API, parses its documentation, and returns the - formatted definition as a string. - """ + Resolves a JAX API, parses its documentation, and returns the + formatted definition as a string. + """ obj = resolve_api(api_str) doc = inspect.getdoc(obj) signature = get_signature(obj) @@ -190,22 +198,22 @@ def add_line(text=""): if __name__ == "__main__": parser = argparse.ArgumentParser( - description="Parse JAX API definitions and write to a file.") + description="Parse JAX API definitions and write to a file." + ) parser.add_argument( - "--api", - type=str, - required=True, - help= - "JAX API string (e.g., jax.numpy.dot) or prefix for recursion (e.g., jax.numpy)." + "--api", + type=str, + required=True, + help="JAX API string (e.g., jax.numpy.dot) or prefix for recursion (e.g., jax.numpy).", ) - parser.add_argument("--output", - type=str, - required=True, - help="Path to the output file.") parser.add_argument( - "--recursive", - action="store_true", - help="If enabled, finds all public APIs under the given --api prefix.") + "--output", type=str, required=True, help="Path to the output file." + ) + parser.add_argument( + "--recursive", + action="store_true", + help="If enabled, finds all public APIs under the given --api prefix.", + ) args = parser.parse_args() try: diff --git a/MaxKernel/hitl_agent/tools/search_api_tool.py b/MaxKernel/hitl_agent/tools/search_api_tool.py index d120a27..88ee6d7 100644 --- a/MaxKernel/hitl_agent/tools/search_api_tool.py +++ b/MaxKernel/hitl_agent/tools/search_api_tool.py @@ -9,39 +9,39 @@ """ +from google.adk.tools import FunctionTool +from google.genai import types + from hitl_agent.constants import ( - TEMPERATURE, - TOP_K, - TOP_P, + TEMPERATURE, + TOP_K, + TOP_P, ) - -from google.genai import types -from google.adk.tools import FunctionTool from hitl_agent.tools.api_rag.get_apis import generate_definition model_config = types.GenerateContentConfig( - temperature=TEMPERATURE, - top_p=TOP_P, - top_k=TOP_K, + temperature=TEMPERATURE, + top_p=TOP_P, + top_k=TOP_K, ) def search_api(api_name: str) -> dict: """ - Search for API documentation and generate its definition. + Search for API documentation and generate its definition. - This function attempts to retrieve and generate a definition for a given API name. - It serves as a tool for looking up API specifications and returning structured - information about the requested API. + This function attempts to retrieve and generate a definition for a given API name. + It serves as a tool for looking up API specifications and returning structured + information about the requested API. - Args: - api_name: The name of the API to search for and generate documentation. + Args: + api_name: The name of the API to search for and generate documentation. - Returns: - A dictionary containing the operation result with the following structure: - - If successful: {"status": "success", "message": } - - If failed: {"status": "error", "message": "API provided is not a valid API"} - """ + Returns: + A dictionary containing the operation result with the following structure: + - If successful: {"status": "success", "message": } + - If failed: {"status": "error", "message": "API provided is not a valid API"} + """ try: definition = generate_definition(api_name) return {"status": "success", "message": definition} diff --git a/MaxKernel/hitl_agent/tools/tools.py b/MaxKernel/hitl_agent/tools/tools.py index f1e268e..8aa108a 100644 --- a/MaxKernel/hitl_agent/tools/tools.py +++ b/MaxKernel/hitl_agent/tools/tools.py @@ -1,16 +1,20 @@ """Tool setup for HITL kernel generation agents.""" -import os import logging -from google.adk.tools import ToolContext -from google.adk.tools.retrieval.vertex_ai_rag_retrieval import VertexAiRagRetrieval +import os + from google.adk.models import LlmRequest -from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset +from google.adk.tools import ToolContext from google.adk.tools.mcp_tool.mcp_session_manager import StdioConnectionParams +from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset +from google.adk.tools.retrieval.vertex_ai_rag_retrieval import ( + VertexAiRagRetrieval, +) from mcp import StdioServerParameters from vertexai.preview import rag + +from hitl_agent.config import RAG_CORPUS, WORKDIR from hitl_agent.tools.search_api_tool import search_api_tool -from hitl_agent.config import WORKDIR, RAG_CORPUS # Custom VertexAiRagRetrieval that forces function_declarations mode to avoid @@ -18,72 +22,78 @@ class CompatibleVertexAiRagRetrieval(VertexAiRagRetrieval): """VertexAiRagRetrieval that uses function_declarations instead of retrieval mode. - This avoids the 400 INVALID_ARGUMENT error when mixing with MCPToolset. - """ + This avoids the 400 INVALID_ARGUMENT error when mixing with MCPToolset. + """ async def process_llm_request( - self, - *, - tool_context: ToolContext, - llm_request: LlmRequest, + self, + *, + tool_context: ToolContext, + llm_request: LlmRequest, ) -> None: # Always use function_declarations mode, even for Gemini 2+ # to maintain compatibility with MCPToolset from google.adk.tools.retrieval.base_retrieval_tool import BaseRetrievalTool - await BaseRetrievalTool.process_llm_request(self, - tool_context=tool_context, - llm_request=llm_request) + + await BaseRetrievalTool.process_llm_request( + self, tool_context=tool_context, llm_request=llm_request + ) # Read-only filesystem tool for orchestration agent (no write access) filesystem_tool_r = MCPToolset( - connection_params=StdioConnectionParams( - server_params=StdioServerParameters( - command='npx', - args=[ - "-y", # Argument for npx to auto-confirm install - "@modelcontextprotocol/server-filesystem@0.5.1", - os.path.abspath(WORKDIR), - ], - ),), - # Optional: Filter which tools from the MCP server are exposed - tool_filter=['list_directory', 'read_file']) + connection_params=StdioConnectionParams( + server_params=StdioServerParameters( + command="npx", + args=[ + "-y", # Argument for npx to auto-confirm install + "@modelcontextprotocol/server-filesystem@0.5.1", + os.path.abspath(WORKDIR), + ], + ), + ), + # Optional: Filter which tools from the MCP server are exposed + tool_filter=["list_directory", "read_file"], +) # Read-write filesystem tool for sub-agents filesystem_tool_rw = MCPToolset( - connection_params=StdioConnectionParams( - server_params=StdioServerParameters( - command='npx', - args=[ - "-y", # Argument for npx to auto-confirm install - "@modelcontextprotocol/server-filesystem@0.5.1", - os.path.abspath(WORKDIR), - ], - ),), - # Optional: Filter which tools from the MCP server are exposed - tool_filter=['list_directory', 'read_file', 'write_file']) + connection_params=StdioConnectionParams( + server_params=StdioServerParameters( + command="npx", + args=[ + "-y", # Argument for npx to auto-confirm install + "@modelcontextprotocol/server-filesystem@0.5.1", + os.path.abspath(WORKDIR), + ], + ), + ), + # Optional: Filter which tools from the MCP server are exposed + tool_filter=["list_directory", "read_file", "write_file"], +) # Vertex AI RAG Engine tool vertex_ai_rag_tool = None if RAG_CORPUS: vertex_ai_rag_tool = CompatibleVertexAiRagRetrieval( - name='retrieval_tool', - description= - 'Use this tool to retrieve Pallas/JAX/TPU documentation and examples from the RAG corpus. This is helpful for answering questions about Pallas concepts, JAX APIs, TPU architecture, best practices, and implementation patterns.', - rag_resources=[rag.RagResource(rag_corpus=RAG_CORPUS)], - similarity_top_k=10, - vector_distance_threshold=0.6, + name="retrieval_tool", + description="Use this tool to retrieve Pallas/JAX/TPU documentation and examples from the RAG corpus. This is helpful for answering questions about Pallas concepts, JAX APIs, TPU architecture, best practices, and implementation patterns.", + rag_resources=[rag.RagResource(rag_corpus=RAG_CORPUS)], + similarity_top_k=10, + vector_distance_threshold=0.6, ) logging.info( - f"Initialized CompatibleVertexAiRagRetrieval with corpus: {RAG_CORPUS}") + f"Initialized CompatibleVertexAiRagRetrieval with corpus: {RAG_CORPUS}" + ) else: logging.warning( - "RAG_CORPUS not set. VertexAiRagRetrieval tool will not be available.") + "RAG_CORPUS not set. VertexAiRagRetrieval tool will not be available." + ) __all__ = [ - 'search_api_tool', - 'filesystem_tool_r', - 'filesystem_tool_rw', - 'vertex_ai_rag_tool', - 'CompatibleVertexAiRagRetrieval', + "search_api_tool", + "filesystem_tool_r", + "filesystem_tool_rw", + "vertex_ai_rag_tool", + "CompatibleVertexAiRagRetrieval", ] diff --git a/MaxKernel/pyproject.toml b/MaxKernel/pyproject.toml index dc9dce8..6cec9c1 100644 --- a/MaxKernel/pyproject.toml +++ b/MaxKernel/pyproject.toml @@ -1,5 +1,5 @@ [tool.ruff] -line-length = 120 +line-length = 80 target-version = "py310" indent-width = 2 diff --git a/MaxKernel/setup.py b/MaxKernel/setup.py index 362380d..38b0bc6 100644 --- a/MaxKernel/setup.py +++ b/MaxKernel/setup.py @@ -1,13 +1,12 @@ -from setuptools import setup, find_packages +from setuptools import find_packages, setup setup( - name="hitl-agent", - version="0.1.0", - description="human-in-the-loop agent for TPU kernel generation, orchestrating the entire process and integrating GPU to JAX conversion capabilities.", - author="Your Name", - author_email="your.email@example.com", - packages=find_packages(), - python_requires=">=3.7", - install_requires=[ - ], + name="hitl-agent", + version="0.1.0", + description="human-in-the-loop agent for TPU kernel generation, orchestrating the entire process and integrating GPU to JAX conversion capabilities.", + author="Your Name", + author_email="your.email@example.com", + packages=find_packages(), + python_requires=">=3.7", + install_requires=[], )