fix: use eval-timeout instead of step-timeout

Author: dexhunter

README.md

@@ -134,8 +134,7 @@ For more advanced examples, including [Triton](/examples/triton/README.md), [CUD
 | `-M, --model`                  | Model identifier for the LLM to use (e.g., `o4-mini`, `claude-sonnet-4-0`).                                                                                                        | `o4-mini` when `OPENAI_API_KEY` is set; `claude-sonnet-4-0` when `ANTHROPIC_API_KEY` is set; `gemini-2.5-pro` when `GEMINI_API_KEY` is set. | `-M o4-mini`         |
 | `-i, --additional-instructions`| Natural language description of specific instructions **or** path to a file containing detailed instructions to guide the LLM.                                                                                             | `None`                                                                                                                                                  | `-i instructions.md` or `-i "Optimize the model for faster inference"`|
 | `-l, --log-dir`                | Path to the directory to log intermediate steps and final optimization result.                                                                                                                                             | `.runs/`                                                                                                                                               | `-l ./logs/`        |
-| `--step-timeout`       | Timeout in seconds for each optimization step.                                                                                                                                                                             | `3600`                                                                                                                                                  | `--step-timeout 60`             |
-| `--overall-timeout`       | Overall timeout in seconds for the entire optimization run.                                                                                                                                                                             | `None`                                                                                                                                                  | `--overall-timeout 3600`             |
+| `--eval-timeout`       | Timeout in seconds for each step in evaluation.                                                                                                                                                                             | `3600`                                                                                                                                                  | `--eval-timeout 60`             |
 
 ---
 

weco/cli.py

@@ -62,13 +62,7 @@ def configure_run_parser(run_parser: argparse.ArgumentParser) -> None:
         help="Description of additional instruction or path to a file containing additional instructions. Defaults to None.",
     )
     run_parser.add_argument(
-        "--step-timeout", type=int, default=3600, help="Timeout in seconds for each optimization step. Defaults to 3600."
-    )
-    run_parser.add_argument(
-        "--overall-timeout",
-        type=int,
-        default=None,
-        help="Overall timeout in seconds for the entire optimization run. Defaults to None (no limit).",
+        "--eval-timeout", type=int, default=3600, help="Timeout in seconds for each evaluation. Defaults to 3600."
     )
 
 
@@ -86,8 +80,7 @@ def execute_run_command(args: argparse.Namespace) -> None:
         log_dir=args.log_dir,
         additional_instructions=args.additional_instructions,
         console=console,
-        step_timeout=args.step_timeout,
-        overall_timeout=args.overall_timeout,
+        eval_timeout=args.eval_timeout,
     )
     exit_code = 0 if success else 1
     sys.exit(exit_code)

weco/optimizer.py

@@ -78,8 +78,7 @@ def execute_optimization(
     log_dir: str = ".runs",
     additional_instructions: Optional[str] = None,
     console: Optional[Console] = None,
-    step_timeout: int = 3600,
-    overall_timeout: Optional[int] = None,
+    eval_timeout: int = 3600,
 ) -> bool:
     """
     Execute the core optimization logic.
@@ -120,39 +119,10 @@ def signal_handler(signum, frame):
         # Exit gracefully
         sys.exit(0)
 
-    def timeout_handler():
-        console.print(f"\n[bold yellow]Overall timeout of {overall_timeout} seconds reached. Shutting down...[/]")
-
-        # Stop heartbeat thread
-        stop_heartbeat_event.set()
-        if heartbeat_thread and heartbeat_thread.is_alive():
-            heartbeat_thread.join(timeout=2)
-
-        # Report termination
-        if current_run_id_for_heartbeat:
-            report_termination(
-                run_id=current_run_id_for_heartbeat,
-                status_update="terminated",
-                reason="timeout_reached",
-                details=f"Overall optimization timeout of {overall_timeout} seconds reached.",
-                auth_headers=current_auth_headers_for_heartbeat,
-                timeout=3,
-            )
-
-        # Exit gracefully
-        sys.exit(0)
-
     # Set up signal handlers for this run
     original_sigint_handler = signal.signal(signal.SIGINT, signal_handler)
     original_sigterm_handler = signal.signal(signal.SIGTERM, signal_handler)
 
-    # Set up the overall timeout if specified
-    if overall_timeout is not None:
-        overall_timeout_timer = threading.Timer(overall_timeout, timeout_handler)
-        overall_timeout_timer.start()
-    else:
-        overall_timeout_timer = None
-
     run_id = None
     optimization_completed_normally = False
     user_stop_requested_flag = False
@@ -279,7 +249,7 @@ def timeout_handler():
             )
 
             # Run evaluation on the initial solution
-            term_out = run_evaluation(eval_command=eval_command, timeout=step_timeout)
+            term_out = run_evaluation(eval_command=eval_command, timeout=eval_timeout)
             # Update the evaluation output panel
             eval_output_panel.update(output=term_out)
             smooth_update(
@@ -378,7 +348,7 @@ def timeout_handler():
                     ],
                     transition_delay=0.08,  # Slightly longer delay for more noticeable transitions
                 )
-                term_out = run_evaluation(eval_command=eval_command, timeout=step_timeout)
+                term_out = run_evaluation(eval_command=eval_command, timeout=eval_timeout)
                 eval_output_panel.update(output=term_out)
                 smooth_update(
                     live=live,
@@ -478,10 +448,6 @@ def timeout_handler():
         # Ensure optimization_completed_normally is False
         optimization_completed_normally = False
     finally:
-        # Cancel the overall timeout timer if it was set
-        if overall_timeout_timer:
-            overall_timeout_timer.cancel()
-
         # Restore original signal handlers
         signal.signal(signal.SIGINT, original_sigint_handler)
         signal.signal(signal.SIGTERM, original_sigterm_handler)