diff --git a/MaxKernel/auto_agent/config.py b/MaxKernel/auto_agent/config.py index 21c356a..13f9241 100644 --- a/MaxKernel/auto_agent/config.py +++ b/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 diff --git a/MaxKernel/auto_agent/server_utils/server_config.py b/MaxKernel/auto_agent/server_utils/server_config.py index b4c47b1..7386b69 100644 --- a/MaxKernel/auto_agent/server_utils/server_config.py +++ b/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() diff --git a/MaxKernel/auto_agent/server_utils/setup.sh b/MaxKernel/auto_agent/server_utils/setup.sh old mode 100644 new mode 100755 index 46e499f..6b99dbe --- a/MaxKernel/auto_agent/server_utils/setup.sh +++ b/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> "$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 diff --git a/MaxKernel/auto_agent/server_utils/tpu_server.py b/MaxKernel/auto_agent/server_utils/tpu_server.py index 97dff1e..23764b0 100644 --- a/MaxKernel/auto_agent/server_utils/tpu_server.py +++ b/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( diff --git a/MaxKernel/auto_search/run_search.py b/MaxKernel/auto_search/run_search.py index d0609a7..2ea4a49 100644 --- a/MaxKernel/auto_search/run_search.py +++ b/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") diff --git a/MaxKernel/auto_search/utils/analyze_timing.py b/MaxKernel/auto_search/utils/analyze_timing.py index 9d02435..6cafec5 100644 --- a/MaxKernel/auto_search/utils/analyze_timing.py +++ b/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( diff --git a/MaxKernel/hitl_agent/config.py b/MaxKernel/hitl_agent/config.py index 9951901..3d696b9 100644 --- a/MaxKernel/hitl_agent/config.py +++ b/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 diff --git a/MaxKernel/prepare_maxkernel.sh b/MaxKernel/prepare_maxkernel.sh old mode 100644 new mode 100755 index aea69c4..f170330 --- a/MaxKernel/prepare_maxkernel.sh +++ b/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> "$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 ----------------------------------