feat: Add step-timeout param

Author: dexhunter

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/`        |
+| `--step-timeout`       | Timeout in seconds for each optimization step.                                                                                                                                                                             | `3600`                                                                                                                                                  | `--step-timeout 60`             |
 
 ---
 

weco/cli.py

@@ -61,6 +61,9 @@ 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(
+        "--step-timeout", type=int, default=3600, help="Timeout in seconds for each optimization step. Defaults to 3600."
+    )
 
 
 def execute_run_command(args: argparse.Namespace) -> None:
@@ -77,6 +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,
     )
     exit_code = 0 if success else 1
     sys.exit(exit_code)

weco/optimizer.py

@@ -78,6 +78,7 @@ def execute_optimization(
     log_dir: str = ".runs",
     additional_instructions: Optional[str] = None,
     console: Optional[Console] = None,
+    step_timeout: int = 3600,
 ) -> bool:
     """
     Execute the core optimization logic.
@@ -248,7 +249,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=step_timeout)
             # Update the evaluation output panel
             eval_output_panel.update(output=term_out)
             smooth_update(
@@ -347,7 +348,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=step_timeout)
                 eval_output_panel.update(output=term_out)
                 smooth_update(
                     live=live,

weco/utils.py

@@ -124,19 +124,21 @@ def smooth_update(
 
 
 # Other helper functions
-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 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 output
+    except subprocess.TimeoutExpired:
+        return f"Evaluation timed out after {timeout} seconds."
 
 
 # Update Check Function