Make max compilation retries configurable via environment variable and chat commands

Author: LukeA12

MaxKernel/auto_agent/config.py

@@ -12,6 +12,7 @@
 TPU_VERSION = os.environ.get("TPU_VERSION", "")
 RAG_CORPUS = os.environ.get("RAG_CORPUS", "")
 INCLUDE_THOUGHTS = os.environ.get("INCLUDE_THOUGHTS", "true").lower() == "true"
+MAX_COMPILATION_RETRIES = int(os.environ.get("MAX_COMPILATION_RETRIES", "6"))
 
 # Model configuration
 model_config = types.GenerateContentConfig(

MaxKernel/auto_agent/subagents/kernel_writing/agent.py

@@ -16,7 +16,7 @@
   load_kernel_and_plan_to_state,
   load_single_kernel_to_state,
 )
-from auto_agent.config import model_config, thinking_planner
+from auto_agent.config import model_config, thinking_planner, MAX_COMPILATION_RETRIES
 from auto_agent.constants import MODEL_NAME
 from auto_agent.custom_types import CustomLlmAgent
 from auto_agent.subagents.kernel_writing.kernel_compilation import (
@@ -46,15 +46,15 @@ class KernelCompilationValidationLoop(BaseAgent):
   compilation_checker: Optional[BaseAgent] = None
   fix_agent: Optional[BaseAgent] = None
   debug_agent: Optional[BaseAgent] = None
-  max_retries: int = 4
+  max_retries: int = 6
 
   def __init__(
     self,
     name: str,
     compilation_checker: BaseAgent,
     fix_agent: BaseAgent,
     debug_agent: Optional[BaseAgent] = None,
-    max_retries: int = 4,
+    max_retries: int = 6,
   ):
     super().__init__(
       name=name,
@@ -353,7 +353,7 @@ async def _run_async_impl(
   compilation_checker=kernel_compilation_checker_for_validation,
   fix_agent=fix_kernel_compilation_agent,
   debug_agent=add_debug_statements_agent,
-  max_retries=6,
+  max_retries=MAX_COMPILATION_RETRIES,
 )
 
 kernel_compilation_summary_agent = CustomLlmAgent(

MaxKernel/auto_agent/subagents/kernel_writing/prompts/ask_validation_prompt.py

@@ -8,7 +8,7 @@
 YOUR TASK:
 Present the user with clear options for what to do next:
 
-1. **Validate & Auto-Fix Compilation**: Run automatic compilation validation with error fixing (up to 4 attempts)
+1. **Validate & Auto-Fix Compilation**: Run automatic compilation validation with error fixing (includes up to 6 automatic error-fixing attempts; you can ask me to change this limit at any time)
 2. **Generate Tests**: Skip validation and proceed directly to test generation
 3. **Something Else**: Ask a question or perform another action
 

MaxKernel/auto_agent/subagents/profiling/agent.py

@@ -118,13 +118,22 @@ async def _run_async_impl(
   ),
   output_key="profiling_summary",
   include_contents="none",
-  tools=[
-    offline_tools.load_xplane_and_query,
-    offline_tools.get_hlo_dump,
-    offline_tools.create_chart_from_xplane,
-    offline_tools.get_overview_page_metrics,
-    vertex_ai_rag_tool,
-  ],
+  tools=(
+    [
+      offline_tools.load_xplane_and_query,
+      offline_tools.get_hlo_dump,
+      offline_tools.create_chart_from_xplane,
+      offline_tools.get_overview_page_metrics,
+      vertex_ai_rag_tool,
+    ]
+    if vertex_ai_rag_tool
+    else [
+      offline_tools.load_xplane_and_query,
+      offline_tools.get_hlo_dump,
+      offline_tools.create_chart_from_xplane,
+      offline_tools.get_overview_page_metrics,
+    ]
+  ),
 )
 
 # Main profiling orchestrator agent

MaxKernel/hitl_agent/agent.py

@@ -31,7 +31,11 @@
   unified_test_agent,
   validated_test_generation_agent,
 )
-from hitl_agent.tools.tools import filesystem_tool_r, set_working_directory
+from hitl_agent.tools.tools import (
+  filesystem_tool_r,
+  set_working_directory,
+  set_max_compilation_retries,
+)
 
 # Root orchestration agent
 root_agent = CustomLlmAgent(
@@ -58,6 +62,7 @@
   tools=[
     filesystem_tool_r,
     set_working_directory,
+    set_max_compilation_retries,
   ],  # 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.",

MaxKernel/hitl_agent/config.py

@@ -12,6 +12,7 @@
 TPU_VERSION = os.environ.get("TPU_VERSION", "")
 RAG_CORPUS = os.environ.get("RAG_CORPUS", "")
 INCLUDE_THOUGHTS = os.environ.get("INCLUDE_THOUGHTS", "true").lower() == "true"
+MAX_COMPILATION_RETRIES = int(os.environ.get("MAX_COMPILATION_RETRIES", "6"))
 
 # Model configuration
 model_config = types.GenerateContentConfig(

MaxKernel/hitl_agent/prompts/interactive_prompt.py

@@ -11,10 +11,11 @@
 
 1.  **Filesystem Tool (Read-Only)**: You have a `filesystem_tool` that can **only read and list files** - you cannot write files. All file writing must be delegated to the appropriate sub-agents. Assume that all files you need to read are located in {workdir}. For example, if the user asks to read file `script.py`, you should call the tool with the path `{workdir}/script.py`.
 2.  **Set Working Directory Tool**: You have a `set_working_directory` tool that can update the active workspace/working directory path for the session. Use this tool when the user requests to switch, set, or change the active working directory or workspace folder.
-3.  **Specialist Sub-Agents**: You have a team of sub-agents you can delegate tasks to:
+3.  **Set Max Compilation Retries Tool**: You have a `set_max_compilation_retries` tool that can update the maximum number of compilation validation and auto-fixing attempts. Use this tool when the user requests to set or change the maximum compilation retries or fixing attempts count.
+4.  **Specialist Sub-Agents**: You have a team of sub-agents you can delegate tasks to:
     * **PlanKernelAgent**: Creates or revises optimization plans for Pallas kernels. Automatically presents plans to the user with approval options. Use for both new plans ("optimize X") and revisions ("change Y in the plan").
     * **ImplementKernelAgent**: Implements a Pallas kernel following an approved plan. Use ONLY after user approves the plan. Automatically asks the user about validation options after completion.
-    * **ValidateKernelCompilationAgent**: Validates kernel compilation with automatic error fixing and debugging (up to 4 attempts). Use when user requests validation or wants to check compilation.
+    * **ValidateKernelCompilationAgent**: Validates kernel compilation with automatic error fixing and debugging (up to 6 attempts by default; the user can ask you to change this limit). Use when user requests validation or wants to check compilation.
     * **ExplanationAgent**: Explains TPU or Pallas concepts.
     * **GenerateTestFileAgent**: Generates a comprehensive pytest test file with compilation, correctness, and performance tests for kernel files.
     * **UnifiedTestAgent**: Executes the generated pytest test file on TPU and provides comprehensive results including full tracebacks. Automatically manages server lifecycle (starts/stops TPU and eval servers as needed).
@@ -35,12 +36,16 @@
     * Is it profiling? (e.g., "Profile the kernel", "What are the bottlenecks?")
     * Is it GPU conversion? (e.g., "Convert CUDA to JAX", "Write JAX from PyTorch")
     * Is it setting/changing the workspace or working directory? (e.g., "Change working directory to /path/to/folder", "Set workspace folder to Y")
+    * Is it changing the maximum compilation retries or attempts? (e.g., "Set max retries to 8", "Change max fixing attempts to 10")
 
 2.  **Execute the Plan**:
 
     * **If the request is to CHANGE/SET the workspace or working directory**:
         * **Action**: Call `set_working_directory` tool with the absolute path.
 
+    * **If the request is to CHANGE/SET the maximum compilation retries**:
+        * **Action**: Call `set_max_compilation_retries` tool with the integer value.
+
     * **If the request is simple explanation** (like "What is a TPU?"):
         * **Action**: Delegate to `ExplanationAgent`.
 

MaxKernel/hitl_agent/subagents/kernel_writing/agent.py

@@ -15,7 +15,7 @@
   load_single_kernel_to_state,
   save_kernel_and_plan_paths,
 )
-from hitl_agent.config import model_config, thinking_planner
+from hitl_agent.config import model_config, thinking_planner, MAX_COMPILATION_RETRIES
 from hitl_agent.constants import MODEL_NAME
 from hitl_agent.custom_types import CustomLlmAgent
 from hitl_agent.subagents.kernel_writing.kernel_compilation import (
@@ -41,15 +41,15 @@ class KernelCompilationValidationLoop(BaseAgent):
   compilation_checker: Optional[BaseAgent] = None
   fix_agent: Optional[BaseAgent] = None
   debug_agent: Optional[BaseAgent] = None
-  max_retries: int = 4
+  max_retries: int = 6
 
   def __init__(
     self,
     name: str,
     compilation_checker: BaseAgent,
     fix_agent: BaseAgent,
     debug_agent: Optional[BaseAgent] = None,
-    max_retries: int = 4,
+    max_retries: int = 6,
   ):
     super().__init__(
       name=name,
@@ -372,7 +372,7 @@ async def _run_async_impl(
   compilation_checker=kernel_compilation_checker_for_validation,
   fix_agent=fix_kernel_compilation_agent,
   debug_agent=add_debug_statements_agent,
-  max_retries=4,
+  max_retries=MAX_COMPILATION_RETRIES,
 )
 
 kernel_compilation_summary_agent = CustomLlmAgent(

MaxKernel/hitl_agent/subagents/kernel_writing/prompts/ask_validation_prompt.py

@@ -8,7 +8,7 @@
 YOUR TASK:
 Present the user with clear options for what to do next:
 
-1. **Validate & Auto-Fix Compilation**: Run automatic compilation validation with error fixing (up to 4 attempts)
+1. **Validate & Auto-Fix Compilation**: Run automatic compilation validation with error fixing (includes up to 6 automatic error-fixing attempts; you can ask me to change this limit at any time)
 2. **Generate Tests**: Skip validation and proceed directly to test generation
 3. **Something Else**: Ask a question or perform another action
 

MaxKernel/hitl_agent/tools/retries_tool.py

@@ -0,0 +1,64 @@
+"""Tool for dynamically updating the max compilation retries setting."""
+
+import logging
+import os
+from google.adk.tools import FunctionTool, ToolContext
+
+
+async def set_max_compilation_retries_fn(retries: int, tool_context: ToolContext) -> str:
+  """Set the maximum number of compilation validation and auto-fixing attempts.
+
+  Args:
+      retries: The maximum number of attempts (must be positive).
+
+  Returns:
+      A string containing the operation result description:
+      - If successful: "Successfully updated maximum compilation retries to: <retries>"
+      - If failed: "Error: The number of retries must be a positive integer."
+  """
+  if retries <= 0:
+    return "Error: The number of retries must be a positive integer."
+
+  # Update in-memory config variables
+  import hitl_agent.config as hitl_cfg
+  import auto_agent.config as auto_cfg
+  hitl_cfg.MAX_COMPILATION_RETRIES = retries
+  auto_cfg.MAX_COMPILATION_RETRIES = retries
+  os.environ["MAX_COMPILATION_RETRIES"] = str(retries)
+
+  # Update .env file
+  try:
+    env_path = ".env"
+    if os.path.exists(env_path):
+      with open(env_path, "r") as f:
+        lines = f.readlines()
+      with open(env_path, "w") as f:
+        updated = False
+        for line in lines:
+          if line.strip().startswith("MAX_COMPILATION_RETRIES="):
+            f.write(f"MAX_COMPILATION_RETRIES={retries}\n")
+            updated = True
+          else:
+            f.write(line)
+        if not updated:
+          f.write(f"MAX_COMPILATION_RETRIES={retries}\n")
+  except Exception as e:
+    logging.error(f"Failed to update .env: {e}")
+
+  # Update the active validation loops in-memory
+  try:
+    from hitl_agent.subagents.kernel_writing.agent import kernel_compilation_validation_loop as hitl_loop
+    hitl_loop.max_retries = retries
+  except Exception as e:
+    logging.warning(f"Could not update hitl_loop max_retries: {e}")
+
+  try:
+    from auto_agent.subagents.kernel_writing.agent import kernel_compilation_validation_loop as auto_loop
+    auto_loop.max_retries = retries
+  except Exception as e:
+    logging.warning(f"Could not update auto_loop max_retries: {e}")
+
+  return f"Successfully updated maximum compilation retries to: {retries}"
+
+
+set_max_compilation_retries = FunctionTool(set_max_compilation_retries_fn)