y

skillopt/envs/spreadsheetbench/codegen_agent.py

@@ -188,7 +188,7 @@ def _build_user(
 
 # ── LLM call with retry ────────────────────────────────────────────────────
 
-def _llm_call_with_retry(call_fn, *, retries: int = 5, timeout: int = 120):
+def _llm_call_with_retry(call_fn, *, retries: int = 5, timeout: int | None = 120):
     """Wrap an LLM API call with retry and per-call timeout."""
     last_err = None
     for attempt in range(retries):
@@ -335,7 +335,7 @@ def _chat_call(
     deployment: str,
     messages: list[dict],
     max_output_tokens: int,
-    llm_timeout: int = 120,
+    llm_timeout: int | None = 120,
 ) -> str:
     """Single LLM call, no tools. Returns raw text."""
     reasoning_effort = get_reasoning_effort()
@@ -402,8 +402,8 @@ def run_single(
     answer_position: str = "",
     skill_content: str = "",
     max_output_tokens: int = 16384,
-    llm_timeout: int = 120,
-    task_timeout: int = 300,
+    llm_timeout: int | None = 120,
+    task_timeout: int | None = 300,
     diagnostic_mode: bool = False,
     diagnostic_instruction: str = "",
     diagnostic_trace_context: str = "",
@@ -416,8 +416,9 @@ def run_single(
 
     Returns ``{"code": str, "raw": str, "n_turns": 1}``.
     """
+    no_task_timeout = task_timeout is None or task_timeout <= 0
     if is_target_exec_backend():
-        deadline = time.time() + task_timeout
+        deadline = None if no_task_timeout else time.time() + task_timeout
         deployment = _get_deployment()
         work_dir, skill_md, task_md, prompt = _prepare_codex_workspace(
             instruction=instruction,
@@ -430,8 +431,11 @@ def run_single(
             diagnostic_instruction=diagnostic_instruction,
             diagnostic_trace_context=diagnostic_trace_context,
         )
-        remaining = max(10, int(deadline - time.time()))
-        effective_timeout = min(task_timeout, remaining)
+        if deadline is None:
+            effective_timeout = 10**9
+        else:
+            remaining = max(10, int(deadline - time.time()))
+            effective_timeout = min(task_timeout, remaining)
         final_message, raw = _run_exec_backend(
             work_dir=work_dir,
             prompt=prompt,
@@ -453,7 +457,7 @@ def run_single(
             "target_user_prompt": f"{prompt}\n\n## Task File\n\n{task_md}",
         }
 
-    deadline = time.time() + task_timeout
+    deadline = None if no_task_timeout else time.time() + task_timeout
     client = get_target_client()
     deployment = _get_deployment()
     system = _build_system(skill_content)
@@ -472,8 +476,11 @@ def run_single(
         {"role": "user", "content": user},
     ]
 
-    remaining = max(10, int(deadline - time.time()))
-    effective_timeout = min(llm_timeout, remaining)
+    if deadline is None:
+        effective_timeout = None
+    else:
+        remaining = max(10, int(deadline - time.time()))
+        effective_timeout = min(llm_timeout or remaining, remaining)
     raw = _chat_call(client, deployment, messages, max_output_tokens, llm_timeout=effective_timeout)
     time.sleep(3)  # Rate-limit cooldown after successful LLM call
     code = extract_code(raw)
@@ -497,8 +504,8 @@ def run_multi(
     skill_content: str = "",
     max_turns: int = 5,
     max_output_tokens: int = 16384,
-    llm_timeout: int = 120,
-    task_timeout: int = 600,
+    llm_timeout: int | None = 120,
+    task_timeout: int | None = 600,
     gold_path: str = "",
     diagnostic_mode: bool = False,
     diagnostic_instruction: str = "",
@@ -520,8 +527,9 @@ def run_multi(
 
     Returns ``{"code": str, "raw": str, "n_turns": int, "conversation": [...]}``.
     """
+    no_task_timeout = task_timeout is None or task_timeout <= 0
     if is_target_exec_backend():
-        deadline = time.time() + task_timeout
+        deadline = None if no_task_timeout else time.time() + task_timeout
         deployment = _get_deployment()
         work_dir, skill_md, task_md, initial_prompt = _prepare_codex_workspace(
             instruction=instruction,
@@ -549,11 +557,13 @@ def run_multi(
         solution_path = os.path.join(work_dir, "solution.py")
 
         for turn in range(max_turns):
-            remaining = deadline - time.time()
-            if remaining <= 10:
-                break
-
-            effective_timeout = max(10, int(remaining))
+            if deadline is None:
+                effective_timeout = 10**9
+            else:
+                remaining = deadline - time.time()
+                if remaining <= 10:
+                    break
+                effective_timeout = max(10, int(remaining))
             final_message, raw = _run_exec_backend(
                 work_dir=work_dir,
                 prompt=prompt,
@@ -577,7 +587,12 @@ def run_multi(
                     "Write a complete `solution.py` that reads `INPUT_PATH` and saves `OUTPUT_PATH`."
                 )
             else:
-                ok, err = run_generated_code(code, input_xlsx, output_path)
+                ok, err = run_generated_code(
+                    code,
+                    input_xlsx,
+                    output_path,
+                    timeout=None if no_task_timeout else 120,
+                )
                 if ok:
                     if gold_path and answer_position:
                         from skillopt.envs.spreadsheetbench.rollout import _auto_verify_output
@@ -617,7 +632,7 @@ def run_multi(
             "target_user_prompt": f"{initial_prompt}\n\n## Task File\n\n{task_md}",
         }
 
-    deadline = time.time() + task_timeout
+    deadline = None if no_task_timeout else time.time() + task_timeout
     client = get_target_client()
     deployment = _get_deployment()
     system = _build_system(skill_content)
@@ -640,12 +655,14 @@ def run_multi(
     raw = ""
 
     for turn in range(max_turns):
-        remaining = deadline - time.time()
-        if remaining <= 10:
-            # Not enough time for another round
-            break
-
-        effective_timeout = min(llm_timeout, int(remaining))
+        if deadline is None:
+            effective_timeout = None
+        else:
+            remaining = deadline - time.time()
+            if remaining <= 10:
+                # Not enough time for another round
+                break
+            effective_timeout = min(llm_timeout or int(remaining), int(remaining))
         raw = _chat_call(client, deployment, messages, max_output_tokens, llm_timeout=effective_timeout)
         time.sleep(3)  # Rate-limit cooldown after successful LLM call
         code = extract_code(raw)
@@ -663,7 +680,12 @@ def run_multi(
             continue
 
         # Execute the code
-        ok, err = run_generated_code(code, input_xlsx, output_path)
+        ok, err = run_generated_code(
+            code,
+            input_xlsx,
+            output_path,
+            timeout=None if no_task_timeout else 120,
+        )
         if ok:
             # Execution succeeded — check correctness if gold_path available
             if gold_path and answer_position:

skillopt/envs/spreadsheetbench/executor.py

@@ -34,7 +34,7 @@ def _strip_path_assignments(code: str) -> str:
     return _PATH_ASSIGN_RE.sub("", code)
 
 
-def run_generated_code(code: str, input_path: str, output_path: str, timeout: int = 120) -> tuple[bool, str]:
+def run_generated_code(code: str, input_path: str, output_path: str, timeout: int | None = 120) -> tuple[bool, str]:
     os.makedirs(os.path.dirname(output_path), exist_ok=True)
     cleaned = _strip_path_assignments(code)
     indented = textwrap.indent(cleaned, "    ")
@@ -51,7 +51,7 @@ def run_generated_code(code: str, input_path: str, output_path: str, timeout: in
             [sys.executable, tmp],
             capture_output=True,
             text=True,
-            timeout=timeout,
+            timeout=timeout if timeout and timeout > 0 else None,
         )
         if proc.returncode != 0:
             return False, (proc.stdout + "\n" + proc.stderr).strip()

skillopt/envs/spreadsheetbench/rollout.py

@@ -547,6 +547,7 @@ def process_one_codegen(
     mode: str = "single",
     max_turns: int = 5,
     max_completion_tokens: int = 16384,
+    task_timeout: int = 600,
     use_eval_feedback: bool = False,
     diagnostic_mode: bool = False,
     diagnostic_instruction: str = "",
@@ -659,6 +660,7 @@ def process_one_codegen(
                     skill_content=skill_content,
                     max_turns=max_turns,
                     max_output_tokens=max_completion_tokens,
+                    task_timeout=task_timeout,
                     gold_path=first_gold if use_eval_feedback else "",
                     diagnostic_mode=diagnostic_mode,
                     diagnostic_instruction=diagnostic_instruction,
@@ -673,6 +675,7 @@ def process_one_codegen(
                     answer_position=answer_position_eval,
                     skill_content=skill_content,
                     max_output_tokens=max_completion_tokens,
+                    task_timeout=task_timeout,
                     diagnostic_mode=diagnostic_mode,
                     diagnostic_instruction=diagnostic_instruction,
                     diagnostic_trace_context=diagnostic_trace_context,
@@ -810,10 +813,10 @@ def run_spreadsheet_batch_codegen(
     Args:
         mode: "single" or "multi".
         task_timeout: Hard per-task timeout in seconds at the future level.
-            0 = auto (single: 300s, multi: 600s).
+            0 or negative disables the per-task timeout.
     """
-    if task_timeout <= 0:
-        task_timeout = 300 if mode == "single" else 600
+    no_task_timeout = task_timeout <= 0
+    task_timeout_label = "none" if no_task_timeout else f"{task_timeout}s"
 
     os.makedirs(out_root, exist_ok=True)
 
@@ -833,7 +836,7 @@ def run_spreadsheet_batch_codegen(
     pending = [it for it in items if str(it["id"]) not in done_ids]
     print(
         f"  [spreadsheet codegen-{mode}] total={len(items)} done={len(done_ids)} "
-        f"pending={len(pending)} workers={max_api_workers} task_timeout={task_timeout}s"
+        f"pending={len(pending)} workers={max_api_workers} task_timeout={task_timeout_label}"
     )
 
     if not pending:
@@ -854,6 +857,7 @@ def run_spreadsheet_batch_codegen(
             mode,
             max_turns,
             max_completion_tokens,
+            task_timeout,
             use_eval_feedback,
             diagnostic_mode,
             diagnostic_instruction,
@@ -903,7 +907,7 @@ def run_spreadsheet_batch_codegen(
         while pending_futs:
             done, _ = wait(pending_futs, timeout=5, return_when=FIRST_COMPLETED)
             now = time.time()
-            timed_out = [
+            timed_out = [] if no_task_timeout else [
                 fut for fut in pending_futs - done
                 if str(futs[fut]["id"]) in started_at
                 and now - started_at[str(futs[fut]["id"])] >= task_timeout