Add multi-chip concurrency

MaxKernel/auto_agent/config.py

@@ -52,3 +52,50 @@ def get_thinking_planner(level: str = "high") -> BuiltInPlanner:
       thinking_level=level,
     )
   )
+
+
+# MONKEY PATCH GENERATE_CONTENT to handle rate limits
+try:
+  import asyncio
+  import logging
+
+  import tenacity
+  from google import genai
+
+  def get_retry_decorator():
+    return tenacity.retry(
+      wait=tenacity.wait_exponential(multiplier=1, min=4, max=60),
+      stop=tenacity.stop_after_attempt(10),
+      retry=tenacity.retry_if_exception_type(Exception),
+      before_sleep=tenacity.before_sleep_log(
+        logging.getLogger(__name__), logging.WARNING
+      ),
+    )
+
+  if not hasattr(genai.models.Models, "_original_generate_content"):
+    orig_sync = genai.models.Models.generate_content
+    genai.models.Models._original_generate_content = orig_sync
+
+    @get_retry_decorator()
+    def wrapped_sync(self, *args, **kwargs):
+      return orig_sync(self, *args, **kwargs)
+
+    genai.models.Models.generate_content = wrapped_sync
+
+  if not hasattr(genai.models.AsyncModels, "_original_generate_content"):
+    orig_async = genai.models.AsyncModels.generate_content
+    genai.models.AsyncModels._original_generate_content = orig_async
+
+    @get_retry_decorator()
+    async def wrapped_async(self, *args, **kwargs):
+      # Gemini API occasionally hangs indefinitely on concurrent quotas.
+      # Force a 90 second hard timeout so it triggers a tenacity retry
+      # instead of infinitely blocking the orchestrator.
+      return await asyncio.wait_for(
+        orig_async(self, *args, **kwargs), timeout=90
+      )
+
+    genai.models.AsyncModels.generate_content = wrapped_async
+except ImportError:
+  pass
+# END MONKEY PATCH

MaxKernel/auto_agent/server_utils/server_config.py

@@ -35,35 +35,28 @@ def _resolve_config_path(cfg_path: str) -> Optional[str]:
   return None
 
 
-def get_local_tpu_port(cfg_path: str = "eval_config.yaml") -> Optional[int]:
-  """Checks eval_config.yaml and returns the port if a local TPU server is needed."""
+def get_local_tpu_ports(cfg_path: str = "eval_config.yaml") -> list[int]:
+  """Checks eval_config.yaml and returns the ports if local TPU servers are needed."""
   resolved_path = _resolve_config_path(cfg_path)
   if not resolved_path:
-    return None
+    return []
 
   try:
     with open(resolved_path, "r") as file:
       config = yaml.safe_load(file) or {}
   except Exception as e:
     logging.error(f"Config file {resolved_path} error: {e}")
-    return None
+    return []
 
   if not isinstance(config, dict):
-    raise ValueError(
-      f"Invalid configuration format in {resolved_path}: "
-      "Expected a YAML dictionary at the root level."
-    )
+    return []
 
   backends = config.get("backends", [])
   if not isinstance(backends, list):
-    raise ValueError(
-      f"Invalid configuration format in {resolved_path}: "
-      "'backends' must be a list."
-    )
+    return []
 
   local_ip = get_local_ip()
 
-  # Find all backends that are local TPUs
   local_tpu_backends = [
     b
     for b in backends
@@ -73,11 +66,12 @@ def get_local_tpu_port(cfg_path: str = "eval_config.yaml") -> Optional[int]:
     and "tpu_vm" not in b
   ]
 
-  if not local_tpu_backends:
-    return None
+  return [b.get("port", TPU_SERVER_PORT) for b in local_tpu_backends]
+
 
-  port = local_tpu_backends[0].get("port")
-  return port if port is not None else TPU_SERVER_PORT
+def get_local_tpu_port(cfg_path: str = "eval_config.yaml") -> Optional[int]:
+  ports = get_local_tpu_ports(cfg_path)
+  return ports[0] if ports else None
 
 
 def get_local_cpu_port(cfg_path: str = "eval_config.yaml") -> Optional[int]:
@@ -159,7 +153,9 @@ def get_bastion_config(
 
 
 if __name__ == "__main__":
-  tpu_p = get_local_tpu_port()
+  tpu_ports = get_local_tpu_ports()
+  tpu_p = tpu_ports[0] if tpu_ports else None
+  print(f"LOCAL_TPU_PORTS='{' '.join(map(str, tpu_ports))}'")
   cpu_p = get_local_cpu_port()
   b = get_bastion_config()
 

MaxKernel/auto_agent/server_utils/setup.sh

@@ -93,19 +93,57 @@ elif [ "$1" = "--start-gke" ]; then
         exit 1
     fi
 elif [ "$1" = "--start-local" ] || [ "$1" = "--start-gce" ]; then
+    CHIPS=1
+    while [[ "$#" -gt 0 ]]; do
+        case "$1" in
+            --chips) CHIPS="$2"; shift ;;
+        esac
+        shift
+    done
+
+    # Generate eval_config.yaml dynamically if CHIPS is specified
+    if [ "$CHIPS" -gt 1 ]; then
+        echo "Dynamically generating eval_config.yaml for $CHIPS TPU chips..."
+        HOSTNAME_IP="127.0.0.1"
+        target="$SCRIPT_DIR/eval_config.yaml"
+        echo "backends:" > "$target"
+        tpu_port=5463
+        for (( i=0; i<CHIPS; i++ )); do
+            echo "  - name: tpu-$i" >> "$target"
+            echo "    ip: $HOSTNAME_IP" >> "$target"
+            echo "    port: $tpu_port" >> "$target"
+            echo "    type: tpu" >> "$target"
+            ((tpu_port++))
+        done
+        cpu_port=5464
+        if [ $CHIPS -gt 1 ]; then
+            cpu_port=$tpu_port
+        fi
+        echo "  - name: cpu-0" >> "$target"
+        echo "    ip: $HOSTNAME_IP" >> "$target"
+        echo "    port: $cpu_port" >> "$target"
+        echo "    type: cpu" >> "$target"
+    fi
+
     load_config
     # Start all local execution/evaluation servers (needed for local or GCE cases)
     echo "Starting local background servers (CPU, TPU, Eval)..."
-    if [ -n "$LOCAL_TPU_PORT" ]; then
-        nohup python3 tpu_server.py > output_tpu_server.txt 2>&1 &
+    if [ -n "$LOCAL_TPU_PORTS" ]; then
+        tpu_index=0
+        for port in $LOCAL_TPU_PORTS; do
+            TPU_VISIBLE_DEVICES=$tpu_index TPU_CHIPS_PER_HOST_BOUNDS=1,1,1 TPU_HOST_BOUNDS=1,1,1 PORT=$port nohup python3 tpu_server.py > output_tpu_server_${port}.txt 2>&1 &
+            ((tpu_index++))
+        done
     fi
     if [ -n "$LOCAL_CPU_PORT" ]; then
         nohup python3 cpu_server.py > output_cpu_server.txt 2>&1 &
     fi
     nohup python3 eval_server.py > output_eval_server.txt 2>&1 &
 
-    if [ -n "$LOCAL_TPU_PORT" ]; then
-        wait_for_server_health "TPU server" "$LOCAL_TPU_PORT" "output_tpu_server.txt" || exit 1
+    if [ -n "$LOCAL_TPU_PORTS" ]; then
+        for port in $LOCAL_TPU_PORTS; do
+            wait_for_server_health "TPU server" "$port" "output_tpu_server_${port}.txt" || exit 1
+        done
     fi
     if [ -n "$LOCAL_CPU_PORT" ]; then
         wait_for_server_health "CPU server" "$LOCAL_CPU_PORT" "output_cpu_server.txt" || exit 1

MaxKernel/auto_agent/server_utils/tpu_server.py

@@ -515,7 +515,17 @@ def get_tpu_version() -> dict:
 
 
 if __name__ == "__main__":
-  tpu_port = get_local_tpu_port()
+  port_env = os.environ.get("PORT")
+  if port_env:
+    try:
+      tpu_port = int(port_env)
+    except ValueError:
+      logging.error(
+        f"Invalid PORT environment variable: {port_env}. Must be an integer."
+      )
+      sys.exit(1)
+  else:
+    tpu_port = get_local_tpu_port()
 
   if tpu_port is None:
     logging.info(

MaxKernel/auto_search/run_search.py

@@ -105,6 +105,7 @@ async def run_search(
   **kwargs: Any,
 ) -> Tuple[str, str]:
   """Executes the search algorithm asynchronously for a single reference file."""
+  global_start_time = time.time()
   problem_dir = os.path.dirname(os.path.abspath(reference_file_path))
   default_problem_id, ext = os.path.splitext(
     os.path.basename(reference_file_path)
@@ -172,7 +173,8 @@ async def run_search(
 
       try:
         logger.info("Generating timing summary...")
-        summary_text = analyze_path(dest_dir)
+        run_duration = time.time() - global_start_time
+        summary_text = analyze_path(dest_dir, real_wall_time=run_duration)
         out_file = os.path.join(dest_dir, "timing_summary.md")
         with open(out_file, "w") as f:
           f.write("```text\n" + summary_text + "\n```\n")

MaxKernel/auto_search/utils/analyze_timing.py

@@ -164,7 +164,7 @@ def process_file_metrics(file_path):
   return file_stats
 
 
-def analyze_path(target_path: str):
+def analyze_path(target_path: str, real_wall_time: float = None):
   path = Path(target_path)
   if path.is_file():
     files_to_process = [path]
@@ -254,9 +254,19 @@ def log(msg=""):
   log("        MACRO SUMMARY (ACROSS ALL DISCOVERED NODES)         ")
   log("============================================================")
   log(f"Total Nodes/Attempts Analyzed : {total_runs}")
-  log(
-    f"Aggregated Pipeline Time   : {global_pipeline:>7.2f}s computation-hours"
-  )
+  if real_wall_time:
+    concurrency_factor = (
+      (global_pipeline / real_wall_time) if real_wall_time > 0 else 0
+    )
+    log(
+      f"Aggregated Pipeline Time   : {global_pipeline / 60:>7.2f} computation-minutes"
+    )
+    log(f"Real-World Wall Time       : {real_wall_time / 60:>7.2f} minutes")
+    log(f"Concurrency Acceleration   : {concurrency_factor:>7.2f}x speedup")
+  else:
+    log(
+      f"Aggregated Pipeline Time   : {global_pipeline / 3600:>7.2f} computation-hours"
+    )
 
   if global_pipeline > 0:
     log(

MaxKernel/hitl_agent/config.py

@@ -28,3 +28,51 @@
     thinking_level="high",
   )
 )
+
+
+# MONKEY PATCH GENERATE_CONTENT to handle rate limits
+try:
+  import logging
+
+  import tenacity
+  from google import genai
+
+  def get_retry_decorator():
+    return tenacity.retry(
+      wait=tenacity.wait_exponential(multiplier=1, min=4, max=60),
+      stop=tenacity.stop_after_attempt(10),
+      retry=tenacity.retry_if_exception_type(Exception),
+      before_sleep=tenacity.before_sleep_log(
+        logging.getLogger(__name__), logging.WARNING
+      ),
+    )
+
+  if not hasattr(genai.models.Models, "_original_generate_content"):
+    orig_sync = genai.models.Models.generate_content
+    genai.models.Models._original_generate_content = orig_sync
+
+    @get_retry_decorator()
+    def wrapped_sync(self, *args, **kwargs):
+      return orig_sync(self, *args, **kwargs)
+
+    genai.models.Models.generate_content = wrapped_sync
+
+  if not hasattr(genai.models.AsyncModels, "_original_generate_content"):
+    orig_async = genai.models.AsyncModels.generate_content
+    genai.models.AsyncModels._original_generate_content = orig_async
+
+    @get_retry_decorator()
+    async def wrapped_async(self, *args, **kwargs):
+      import asyncio
+
+      # Gemini API occasionally hangs indefinitely on concurrent quotas.
+      # Force a 90 second hard timeout so it triggers a tenacity retry
+      # instead of infinitely blocking the orchestrator.
+      return await asyncio.wait_for(
+        orig_async(self, *args, **kwargs), timeout=90
+      )
+
+    genai.models.AsyncModels.generate_content = wrapped_async
+except ImportError:
+  pass
+# END MONKEY PATCH

MaxKernel/prepare_maxkernel.sh

@@ -12,6 +12,22 @@ YELLOW='\033[1;33m'
 BLUE='\033[0;34m'
 NC='\033[0m' # No Color
 
+
+CHIPS=1
+# Parse command line arguments
+while [[ "$#" -gt 0 ]]; do
+    case "$1" in
+        --chips) CHIPS="$2"; shift ;;
+        *) echo "Unknown parameter passed: $1"; exit 1 ;;
+    esac
+    shift
+done
+
+if ! [[ "$CHIPS" =~ ^[0-9]+$ ]] || [ "$CHIPS" -lt 1 ]; then
+    echo "Error: --chips must be a positive integer."
+    exit 1
+fi
+
 # Function to print colored output
 print_info() {
     echo -e "${BLUE}[INFO]${NC} $1"
@@ -213,6 +229,12 @@ install_dependencies() {
         pip install -e "$REPO_ROOT"
     fi
 
+    # Install ruff formatter/linter
+    if ! command -v ruff &> /dev/null; then
+        print_info "Installing ruff..."
+        pip install ruff
+    fi
+
     # Check if npx is installed
     if ! command -v npx &> /dev/null; then
         print_info "npx not found. Installing nodejs and npm via nvm..."
@@ -463,6 +485,7 @@ EOF
 }
 
 
+
 # Function to create eval_config.yaml for both auto_agent and hitl_agent
 create_eval_config() {
     print_info "Creating eval_config.yaml for evaluation servers..."
@@ -484,24 +507,31 @@ create_eval_config() {
         local target_dir
         target_dir="$(dirname "$target")"
         if [ -d "$target_dir" ]; then
-            cat > "$target" << EOF
-backends:
-  - name: tpu-0
-    ip: $HOSTNAME_IP
-    port: 5463
-    type: tpu
-  - name: cpu-0
-    ip: $HOSTNAME_IP
-    port: 5464
-    type: cpu
-EOF
+            echo "backends:" > "$target"
+            local tpu_port=5463
+            for (( i=0; i<CHIPS; i++ )); do
+                echo "  - name: tpu-$i" >> "$target"
+                echo "    ip: $HOSTNAME_IP" >> "$target"
+                echo "    port: $tpu_port" >> "$target"
+                echo "    type: tpu" >> "$target"
+                ((tpu_port++))
+            done
+            local cpu_port=5464
+            if [ $CHIPS -gt 1 ]; then
+                cpu_port=$tpu_port
+            fi
+            echo "  - name: cpu-0" >> "$target"
+            echo "    ip: $HOSTNAME_IP" >> "$target"
+            echo "    port: $cpu_port" >> "$target"
+            echo "    type: cpu" >> "$target"
         fi
     done
 
-    print_success "Created eval_config.yaml for both auto_agent and hitl_agent (TPU:5463, CPU:5464, IP: $HOSTNAME_IP)"
+    print_success "Created eval_config.yaml with $CHIPS TPU chips."
 }
 
 
+
 # Main execution
 main() {
     print_info "Starting MaxKernel Agent setup..."
@@ -589,6 +619,7 @@ STEP 2: Install Dependencies
 pip install -r dependency/main_requirements.txt
 pip install -r dependency/agent_requirements.txt
 pip install -e .
+pip install ruff
 
 STEP 3: Set Environment Variables
 ----------------------------------