feat: Add timeout params (#36)

* feat: Add step-timeout param

* feat: Add overall timeout

* fix: use eval-timeout instead of step-timeout

* Update weco/api.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* fix: Add a constant.py

* fix: default to None (no limit)

* Provide clearer response when evaluation times out

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* fix: lint

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Dhruv Srikanth <51223342+DhruvSrikanth@users.noreply.github.com>

Author: 3 people

README.md

@@ -134,6 +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/`        |
+| `--eval-timeout`       | Timeout in seconds for each step in evaluation.                                                                                                                                                                             | No timeout (unlimited)                                                                                                                                                  | `--eval-timeout 3600`             |
 
 ---
 

weco/api.py

@@ -7,6 +7,7 @@
 from rich.console import Console
 
 from weco import __pkg_version__, __base_url__
+from .constants import DEFAULT_API_TIMEOUT
 
 
 # --- Session Configuration ---
@@ -48,7 +49,7 @@ def start_optimization_run(
     additional_instructions: str = None,
     api_keys: Dict[str, Any] = {},
     auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = 800,
+    timeout: Union[int, Tuple[int, int]] = DEFAULT_API_TIMEOUT,
 ) -> Dict[str, Any]:
     """Start the optimization run."""
     with console.status("[bold green]Starting Optimization..."):
@@ -87,7 +88,7 @@ def evaluate_feedback_then_suggest_next_solution(
     additional_instructions: str = None,
     api_keys: Dict[str, Any] = {},
     auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = 800,
+    timeout: Union[int, Tuple[int, int]] = DEFAULT_API_TIMEOUT,
 ) -> Dict[str, Any]:
     """Evaluate the feedback and suggest the next solution."""
     try:
@@ -114,7 +115,10 @@ def evaluate_feedback_then_suggest_next_solution(
 
 
 def get_optimization_run_status(
-    run_id: str, include_history: bool = False, auth_headers: dict = {}, timeout: Union[int, Tuple[int, int]] = 800
+    run_id: str,
+    include_history: bool = False,
+    auth_headers: dict = {},
+    timeout: Union[int, Tuple[int, int]] = DEFAULT_API_TIMEOUT,
 ) -> Dict[str, Any]:
     """Get the current status of the optimization run."""
     try:
@@ -132,7 +136,7 @@ def get_optimization_run_status(
         raise  # Re-raise
 
 
-def send_heartbeat(run_id: str, auth_headers: dict = {}, timeout: Union[int, Tuple[int, int]] = 10) -> bool:
+def send_heartbeat(run_id: str, auth_headers: dict = {}, timeout: Union[int, Tuple[int, int]] = (10, 10)) -> bool:
     """Send a heartbeat signal to the backend."""
     try:
         session = _get_weco_session()
@@ -156,7 +160,7 @@ def report_termination(
     reason: str,
     details: Optional[str] = None,
     auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = 30,
+    timeout: Union[int, Tuple[int, int]] = (10, 30),
 ) -> bool:
     """Report the termination reason to the backend."""
     try:
@@ -206,7 +210,7 @@ def get_optimization_suggestions_from_codebase(
     gitingest_content_str: str,
     console: Console,
     auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = 800,
+    timeout: Union[int, Tuple[int, int]] = DEFAULT_API_TIMEOUT,
 ) -> Optional[List[Dict[str, Any]]]:
     """Analyze codebase and get optimization suggestions using the model-agnostic backend API."""
     try:
@@ -245,7 +249,7 @@ def generate_evaluation_script_and_metrics(
     gitingest_content_str: str,
     console: Console,
     auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = 800,
+    timeout: Union[int, Tuple[int, int]] = DEFAULT_API_TIMEOUT,
 ) -> Tuple[Optional[str], Optional[str], Optional[str], Optional[str]]:
     """Generate evaluation script and determine metrics using the model-agnostic backend API."""
     try:
@@ -286,7 +290,7 @@ def analyze_evaluation_environment(
     gitingest_content_str: str,
     console: Console,
     auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = 800,
+    timeout: Union[int, Tuple[int, int]] = DEFAULT_API_TIMEOUT,
 ) -> Optional[Dict[str, Any]]:
     """Analyze existing evaluation scripts and environment using the model-agnostic backend API."""
     try:
@@ -326,7 +330,7 @@ def analyze_script_execution_requirements(
     target_file: str,
     console: Console,
     auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = 800,
+    timeout: Union[int, Tuple[int, int]] = DEFAULT_API_TIMEOUT,
 ) -> Optional[str]:
     """Analyze script to determine proper execution command using the model-agnostic backend API."""
     try:

weco/cli.py

@@ -61,6 +61,12 @@ def configure_run_parser(run_parser: argparse.ArgumentParser) -> None:
         type=str,
         help="Description of additional instruction or path to a file containing additional instructions. Defaults to None.",
     )
+    run_parser.add_argument(
+        "--eval-timeout",
+        type=int,
+        default=None,
+        help="Timeout in seconds for each evaluation. No timeout by default. Example: --eval-timeout 3600",
+    )
 
 
 def execute_run_command(args: argparse.Namespace) -> None:
@@ -77,6 +83,7 @@ def execute_run_command(args: argparse.Namespace) -> None:
         log_dir=args.log_dir,
         additional_instructions=args.additional_instructions,
         console=console,
+        eval_timeout=args.eval_timeout,
     )
     exit_code = 0 if success else 1
     sys.exit(exit_code)

weco/constants.py

@@ -0,0 +1,7 @@
+# weco/constants.py
+"""
+Constants for the Weco CLI package.
+"""
+
+# API timeout configuration (connect_timeout, read_timeout) in seconds
+DEFAULT_API_TIMEOUT = (10, 800)

weco/optimizer.py

@@ -37,6 +37,7 @@
     smooth_update,
     format_number,
 )
+from .constants import DEFAULT_API_TIMEOUT
 
 
 # --- Heartbeat Sender Class ---
@@ -78,6 +79,7 @@ def execute_optimization(
     log_dir: str = ".runs",
     additional_instructions: Optional[str] = None,
     console: Optional[Console] = None,
+    eval_timeout: Optional[int] = None,
 ) -> bool:
     """
     Execute the core optimization logic.
@@ -153,7 +155,7 @@ def signal_handler(signum, frame):
             "debug_prob": 0.5,
             "max_debug_depth": max(1, math.ceil(0.1 * steps)),
         }
-        timeout = 800
+        api_timeout = DEFAULT_API_TIMEOUT
         processed_additional_instructions = read_additional_instructions(additional_instructions=additional_instructions)
         source_fp = pathlib.Path(source)
         source_code = read_from_path(fp=source_fp, is_json=False)
@@ -181,7 +183,7 @@ def signal_handler(signum, frame):
             additional_instructions=processed_additional_instructions,
             api_keys=llm_api_keys,
             auth_headers=auth_headers,
-            timeout=timeout,
+            timeout=api_timeout,
         )
         run_id = run_response["run_id"]
         current_run_id_for_heartbeat = run_id
@@ -248,7 +250,7 @@ def signal_handler(signum, frame):
             )
 
             # Run evaluation on the initial solution
-            term_out = run_evaluation(eval_command=eval_command)
+            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(
@@ -265,7 +267,7 @@ def signal_handler(signum, frame):
                 if run_id:
                     try:
                         current_status_response = get_optimization_run_status(
-                            run_id=run_id, include_history=False, timeout=30, auth_headers=auth_headers
+                            run_id=run_id, include_history=False, timeout=(10, 30), auth_headers=auth_headers
                         )
                         current_run_status_val = current_status_response.get("status")
                         if current_run_status_val == "stopping":
@@ -284,14 +286,14 @@ def signal_handler(signum, frame):
                     additional_instructions=current_additional_instructions,
                     api_keys=llm_api_keys,
                     auth_headers=auth_headers,
-                    timeout=timeout,
+                    timeout=api_timeout,
                 )
                 # Save next solution (.runs/<run-id>/step_<step>.<extension>)
                 write_to_path(fp=runs_dir / f"step_{step}{source_fp.suffix}", content=eval_and_next_solution_response["code"])
                 # Write the next solution to the source file
                 write_to_path(fp=source_fp, content=eval_and_next_solution_response["code"])
                 status_response = get_optimization_run_status(
-                    run_id=run_id, include_history=True, timeout=timeout, auth_headers=auth_headers
+                    run_id=run_id, include_history=True, timeout=api_timeout, auth_headers=auth_headers
                 )
                 # Update the step of the progress bar, token counts, plan and metric tree
                 summary_panel.set_step(step=step)
@@ -347,7 +349,7 @@ def signal_handler(signum, frame):
                     ],
                     transition_delay=0.08,  # Slightly longer delay for more noticeable transitions
                 )
-                term_out = run_evaluation(eval_command=eval_command)
+                term_out = run_evaluation(eval_command=eval_command, timeout=eval_timeout)
                 eval_output_panel.update(output=term_out)
                 smooth_update(
                     live=live,
@@ -365,13 +367,13 @@ def signal_handler(signum, frame):
                     execution_output=term_out,
                     additional_instructions=current_additional_instructions,
                     api_keys=llm_api_keys,
-                    timeout=timeout,
+                    timeout=api_timeout,
                     auth_headers=auth_headers,
                 )
                 summary_panel.set_step(step=steps)
                 summary_panel.update_token_counts(usage=eval_and_next_solution_response["usage"])
                 status_response = get_optimization_run_status(
-                    run_id=run_id, include_history=True, timeout=timeout, auth_headers=auth_headers
+                    run_id=run_id, include_history=True, timeout=api_timeout, auth_headers=auth_headers
                 )
                 # No need to update the plan panel since we have finished the optimization
                 # Get the optimization run status for

weco/utils.py

@@ -157,20 +157,21 @@ def truncate_output(output: str, max_lines: int = DEFAULT_MAX_LINES, max_chars:
     return output
 
 
-def run_evaluation(eval_command: str) -> str:
+def run_evaluation(eval_command: str, timeout: int | None = None) -> str:
     """Run the evaluation command on the code and return the output."""
 
     # Run the eval command as is
-    result = subprocess.run(eval_command, shell=True, capture_output=True, text=True, check=False)
-
-    # Combine stdout and stderr for complete output
-    output = result.stderr if result.stderr else ""
-    if result.stdout:
-        if len(output) > 0:
-            output += "\n"
-        output += result.stdout
-
-    return truncate_output(output)
+    try:
+        result = subprocess.run(eval_command, shell=True, capture_output=True, text=True, check=False, timeout=timeout)
+        # Combine stdout and stderr for complete output
+        output = result.stderr if result.stderr else ""
+        if result.stdout:
+            if len(output) > 0:
+                output += "\n"
+            output += result.stdout
+        return truncate_output(output)
+    except subprocess.TimeoutExpired:
+        return f"Evaluation timed out after {'an unspecified duration' if timeout is None else f'{timeout} seconds'}."
 
 
 # Update Check Function