Finished Error Handling

xzrderek · xzrderek · commit 7356361c1beb · 2025-08-13T21:21:57.000-07:00
diff --git a/eval_protocol/mcp/execution/manager.py b/eval_protocol/mcp/execution/manager.py
@@ -138,7 +138,7 @@ async def _execute_with_semaphore(idx):
                 if trajectory.terminated:
                     if trajectory.termination_reason == TerminationReason.ERROR:
                         evaluation_row.rollout_status.status = "error"
-                        evaluation_row.rollout_status.error_message = trajectory.control_plane_summary.get(
+                        evaluation_row.rollout_status.termination_reason = trajectory.control_plane_summary.get(
                             "error_message", None
                         )
                     else:
diff --git a/eval_protocol/pytest/default_agent_rollout_processor.py b/eval_protocol/pytest/default_agent_rollout_processor.py
@@ -126,7 +126,7 @@ async def default_agent_rollout_processor(
     async def process_row(row: EvaluationRow) -> EvaluationRow:
         """Process a single row with agent rollout."""
         agent = Agent(
-            model=config.completion_params.model, row=row, config_path=config.mcp_config_path, logger=config.logger
+            model=config.completion_params["model"], row=row, config_path=config.mcp_config_path, logger=config.logger
         )
         try:
             await agent.setup()
@@ -141,7 +141,8 @@ async def _sem_wrapper(r: EvaluationRow) -> EvaluationRow:
             try:
                 return await process_row(r)
             except Exception as e:
-                logger.exception(f"Error processing row {r.input_metadata.row_id}: {e}")
+                r.rollout_status.status = "error"
+                r.rollout_status.termination_reason = str(e)
                 return r
 
     # Create all tasks
diff --git a/eval_protocol/pytest/default_mcp_gym_rollout_processor.py b/eval_protocol/pytest/default_mcp_gym_rollout_processor.py
@@ -6,12 +6,14 @@
 import subprocess
 import time
 from pathlib import Path
-from typing import AsyncIterator, List, Optional
+from typing import Any, AsyncIterator, Dict, List, Optional
 
 import eval_protocol as ep
 from eval_protocol.models import EvaluationRow, Message
 from eval_protocol.pytest.types import RolloutProcessorConfig
 
+CURRENT_RUN_STATE: Dict[str, Any] = {}
+
 
 class MCPServerManager:
     """Manages MCP server lifecycle for testing."""
@@ -204,41 +206,78 @@ async def default_mcp_gym_rollout_processor(
     Args:
         rows: List of EvaluationRow objects containing messages and dataset info in input_metadata
         config: RolloutProcessorConfig with model and other parameters
+            - config.kwargs can include:
+                - start_server (bool): If True, create fresh server and environments. If False, reuse existing ones. Default: True.
 
     Returns:
         AsyncIterator of EvaluationRow objects with completed conversations
     """
-    if config.server_script_path is None:
-        raise ValueError("server_script_path is required for default_mcp_gym_rollout_processor")
-    server = MCPServerManager(config.server_script_path, port=9700, **(config.kwargs or {}))
-
-    try:
-        server.start()
-
-        policy = ep.LiteLLMPolicy(
-            model_id=config.completion_params.model,
-            temperature=config.completion_params.get("temperature", 0.0),
-            max_tokens=config.completion_params.get("max_tokens", 4096),
-            reasoning_effort=config.completion_params.get("reasoning_effort", None),
-        )
+    start_server = config.kwargs.get("start_server", True) if config.kwargs else True
+    if start_server:
+        # Create fresh MCP server and environments for this run
+        if config.server_script_path is None:
+            raise ValueError("server_script_path is required for default_mcp_gym_rollout_processor")
 
-        # Create MCP environments directly from evaluation_rows
-        envs = ep.make(
-            "http://localhost:9700/mcp/",
-            evaluation_rows=rows,
-            model_id=policy.model_id,
-        )
+        server = MCPServerManager(config.server_script_path, port=9700, **(config.kwargs or {}))
 
-        # Run rollout with environments and policy
-        async for evaluation_row in ep.rollout(
-            envs,
-            policy=policy,
-            evaluation_rows=rows,
-            steps=config.steps,
-            max_concurrent_rollouts=config.max_concurrent_rollouts,
-        ):
-            yield evaluation_row
-
-    finally:
-        # Always clean up the server
-        server.stop()
+        try:
+            server.start()
+
+            policy = ep.LiteLLMPolicy(
+                model_id=config.completion_params.get("model", None),
+                temperature=config.completion_params.get("temperature", 0.0),
+                max_tokens=config.completion_params.get("max_tokens", 4096),
+                reasoning_effort=config.completion_params.get("reasoning_effort", None),
+            )
+
+            # Create MCP environments directly from evaluation_rows
+            envs = ep.make(
+                "http://localhost:9700/mcp/",
+                evaluation_rows=rows,
+                model_id=policy.model_id,
+            )
+
+            # Store in current run state for reuse within this run
+            CURRENT_RUN_STATE.update(
+                {
+                    "server": server,
+                    "envs": envs,
+                    "policy": policy,
+                }
+            )
+
+        except Exception as e:
+            server.stop()
+            CURRENT_RUN_STATE.clear()
+            raise e
+
+    else:
+        # Reuse existing MCP environments for retry
+        if not CURRENT_RUN_STATE:
+            raise RuntimeError("Cannot retry without existing server/environments. Call with start_server=True first.")
+
+        server = CURRENT_RUN_STATE["server"]
+        envs = CURRENT_RUN_STATE["envs"]
+        policy = CURRENT_RUN_STATE["policy"]
+
+    # Run rollout with environments and policy (automatically resets environments)
+    async for evaluation_row in ep.rollout(
+        envs,
+        policy=policy,
+        evaluation_rows=rows,
+        steps=config.steps,
+        max_concurrent_rollouts=config.max_concurrent_rollouts,
+    ):
+        yield evaluation_row
+
+
+# Add cleanup method directly to the function object
+def _cleanup_mcp_gym_rollout_processor():
+    """Cleanup function for MCP gym rollout processor"""
+    if CURRENT_RUN_STATE and "server" in CURRENT_RUN_STATE:
+        CURRENT_RUN_STATE["server"].stop()
+        CURRENT_RUN_STATE.clear()  # Clear for next run
+
+
+# Attach cleanup method to the processor function
+default_mcp_gym_rollout_processor.cleanup = _cleanup_mcp_gym_rollout_processor
diff --git a/eval_protocol/pytest/default_single_turn_rollout_process.py b/eval_protocol/pytest/default_single_turn_rollout_process.py
@@ -112,6 +112,8 @@ async def _sem_wrapper(r: EvaluationRow) -> EvaluationRow:
             try:
                 return await process_row(r)
             except Exception as e:
+                r.rollout_status.status = "error"
+                r.rollout_status.termination_reason = str(e)
                 return r
 
     # Create all tasks
diff --git a/eval_protocol/pytest/evaluation_test.py b/eval_protocol/pytest/evaluation_test.py
@@ -8,6 +8,7 @@
 import re
 import statistics
 import time
+from dataclasses import replace
 from typing import Any, Callable, Dict, List, Literal, Optional, Union
 
 import pytest
@@ -269,6 +270,82 @@ def generate_combinations():
 
             return combinations
 
+        async def rollout_processor_with_retry(
+            rollout_processor: RolloutProcessor,
+            fresh_dataset: List[EvaluationRow],
+            config: RolloutProcessorConfig,
+            max_retry: int,
+        ):
+            """
+            Wrapper around rollout_processor that handles retry logic internally.
+            Uses async queue pattern to yield results immediately as they become available.
+            Yields both successful and failed results, leaving it up to the user to handle them in test_func.
+            """
+
+            try:
+                queue = asyncio.Queue()
+                retry_counts = {r.execution_metadata.rollout_id: 0 for r in fresh_dataset}
+                failed_permanently = []
+
+                async def retry_handler(failed_row: EvaluationRow):
+                    rollout_id = failed_row.execution_metadata.rollout_id
+                    current_attempts = retry_counts.get(rollout_id, 0)
+
+                    if current_attempts >= max_retry:
+                        assert (
+                            failed_row.rollout_status and failed_row.rollout_status.status == "error"
+                        ), f"Rollout {failed_row.execution_metadata.rollout_id} did not fail with error status"
+                        failed_permanently.append(failed_row)
+                        await queue.put(failed_row)  # put failed row on queue
+                        return
+
+                    retry_counts[rollout_id] = current_attempts + 1
+
+                    # add kwargs start_server=False to config so we don't start new MCP server
+                    retry_config = replace(config, kwargs={**(config.kwargs or {}), "start_server": False})
+
+                    retry_call = rollout_processor([failed_row], retry_config)
+
+                    retry_result = await anext(retry_call)
+                    if retry_result.rollout_status and retry_result.rollout_status.status == "finished":
+                        await queue.put(retry_result)
+                    else:
+                        asyncio.create_task(retry_handler(retry_result))  # retry failed, spawn another retry
+
+                async def initial_processor():
+                    """Process initial batch and spawn retries for failures"""
+                    async for initial_row in rollout_processor(fresh_dataset, config):
+                        if initial_row.rollout_status and initial_row.rollout_status.status == "finished":
+                            await queue.put(initial_row)  # rollout succeeded, put on queue
+                        else:
+                            asyncio.create_task(retry_handler(initial_row))  # rollout errored, spawn retry task
+
+                processor_task = asyncio.create_task(initial_processor())
+
+                # yield results as they become available
+                completed_count = 0
+                total_expected = len(fresh_dataset)
+
+                while completed_count < total_expected:
+                    finished_row = await queue.get()
+
+                    # only permanent failure rows are put on the queue, so we can check for them here
+                    if finished_row.rollout_status and finished_row.rollout_status.status == "error":
+                        if os.getenv("EP_FAIL_ON_PERMANENT_FAILURE", "true") != "false":
+                            raise RuntimeError(
+                                f"Rollout {finished_row.execution_metadata.rollout_id} failed after {max_retry} retries. Errors: {finished_row.rollout_status.termination_reason}"
+                            )
+
+                    completed_count += 1
+                    yield finished_row
+
+                await processor_task  # explicitly wait for task completion and catch any exceptions
+
+            finally:
+                # processor clean up after themselves if they have a cleanup method
+                if hasattr(rollout_processor, "cleanup"):
+                    rollout_processor.cleanup()
+
         combinations = generate_combinations()
         if len(combinations) == 0:
             raise ValueError(
@@ -410,6 +487,8 @@ def _log_eval_error(
                         kwargs=rollout_processor_kwargs or {},
                     )
 
+                    max_retry = int(os.getenv("EP_MAX_RETRY", "0"))
+
                     for i in range(num_runs):
                         # Regenerate outputs each run by deep-copying the pristine dataset
                         # so model responses are not reused across runs.
@@ -428,15 +507,15 @@ def _log_eval_error(
                         for row in fresh_dataset:
                             active_logger.log(row)
 
-                        rollout_result = rollout_processor(fresh_dataset, config)
-
                         if mode == "pointwise":
                             # Pointwise mode, rollouts will return as they complete so we can pipeline evaluation_test execution
                             semaphore = asyncio.Semaphore(max_concurrent_rollouts)
                             tasks = []
 
                             async def _execute_with_semaphore(row):
                                 async with semaphore:
+                                    # NOTE: we will still evaluate errored rows (give users control over this)
+                                    # i.e., they can choose to give EvaluateResult.score = 0 for errored rows in their test_func
                                     result = await execute_with_params(
                                         test_func,
                                         processed_row=row,
@@ -448,17 +527,23 @@ async def _execute_with_semaphore(row):
                                         )
                                     return result
 
-                            async for row in rollout_processor(fresh_dataset, config):
+                            # Use wrapper that handles retry logic internally
+                            async for row in rollout_processor_with_retry(
+                                rollout_processor, fresh_dataset, config, max_retry
+                            ):
                                 tasks.append(asyncio.create_task(_execute_with_semaphore(row)))
 
                             all_results[i] = await asyncio.gather(*tasks)
 
                         else:
                             # Batch mode: collect all results first, then evaluate (no pipelining)
                             input_dataset = []
-                            async for row in rollout_result:
+                            async for row in rollout_processor_with_retry(
+                                rollout_processor, fresh_dataset, config, max_retry
+                            ):
                                 input_dataset.append(row)
-
+                            # NOTE: we will still evaluate errored rows (give users control over this)
+                            # i.e., they can choose to give EvaluateResult.score = 0 for errored rows in their test_func
                             results = await execute_with_params(
                                 test_func,
                                 processed_dataset=input_dataset,
@@ -530,7 +615,7 @@ async def _execute_with_semaphore(row):
                         should_print = os.getenv("EP_PRINT_SUMMARY") == "1"
                         summary_path = os.getenv("EP_SUMMARY_JSON")
                         suite_name = test_func.__name__
-                        model_used = config.completion_params.model
+                        model_used = config.completion_params["model"]
                         total_rows = len([item for sublist in all_results for item in sublist])
                         summary_obj = {
                             "suite": suite_name,
@@ -990,7 +1075,7 @@ def _deep_update_dict(base: dict, override: dict) -> dict:
             total_rows = len(all_results)
             summary_obj = {
                 "suite": suite_name,
-                "model": config.completion_params.model,
+                "model": config.completion_params["model"],
                 "agg_score": float(agg_score) if agg_score is not None else None,
                 "num_runs": num_runs,
                 "rows": total_rows,
@@ -1001,11 +1086,11 @@ def _deep_update_dict(base: dict, override: dict) -> dict:
             if should_print:
                 if ci_low is not None and ci_high is not None:
                     print(
-                        f"EP Summary | suite={suite_name} model={config.completion_params.model} agg={summary_obj['agg_score']:.3f} ci95=[{ci_low:.3f},{ci_high:.3f}] runs={num_runs} rows={total_rows}"
+                        f"EP Summary | suite={suite_name} model={config.completion_params['model']} agg={summary_obj['agg_score']:.3f} ci95=[{ci_low:.3f},{ci_high:.3f}] runs={num_runs} rows={total_rows}"
                     )
                 else:
                     print(
-                        f"EP Summary | suite={suite_name} model={config.completion_params.model} agg={summary_obj['agg_score']:.3f} runs={num_runs} rows={total_rows}"
+                        f"EP Summary | suite={suite_name} model={config.completion_params['model']} agg={summary_obj['agg_score']:.3f} runs={num_runs} rows={total_rows}"
                     )
             if summary_path:
                 import json as _json
@@ -1037,7 +1122,7 @@ def _extract_effort_tag(params: dict) -> str | None:
                         return None
                     return None
 
-                model_slug = _sanitize_filename(config.completion_params.model)
+                model_slug = _sanitize_filename(config.completion_params["model"])
                 effort_tag = _extract_effort_tag(completion_params) or ""
                 effort_suffix = f"__effort-{_sanitize_filename(effort_tag)}" if effort_tag else ""
                 base_name = f"{suite_name}__{model_slug}{effort_suffix}__{mode}__runs{num_runs}.json"
diff --git a/eval_protocol/pytest/plugin.py b/eval_protocol/pytest/plugin.py
@@ -59,6 +59,23 @@ def pytest_addoption(parser) -> None:
             "Values: low|medium|high"
         ),
     )
+    group.addoption(
+        "--ep-max-retry",
+        action="store",
+        type=int,
+        default=None,
+        help=("Failed rollouts (with rollout_status.status == 'error') will be retried up to this many times."),
+    )
+    group.addoption(
+        "--ep-fail-on-permanent-failure",
+        action="store",
+        default=None,
+        choices=["true", "false"],
+        help=(
+            "Whether to fail the entire rollout when permanent failures occur after max retries. "
+            "Default: true (fail on permanent failures). Set to 'false' to continue with remaining rollouts."
+        ),
+    )
 
 
 def _normalize_max_rows(val: Optional[str]) -> Optional[str]:
@@ -100,6 +117,14 @@ def pytest_configure(config) -> None:
     if summary_json_path:
         os.environ["EP_SUMMARY_JSON"] = summary_json_path
 
+    max_retry = config.getoption("--ep-max-retry")
+    if max_retry is not None:
+        os.environ["EP_MAX_RETRY"] = str(max_retry)
+
+    fail_on_permanent_failure = config.getoption("--ep-fail-on-permanent-failure")
+    if fail_on_permanent_failure is not None:
+        os.environ["EP_FAIL_ON_PERMANENT_FAILURE"] = fail_on_permanent_failure
+
     # Allow ad-hoc overrides of input params via CLI flags
     try:
         import json as _json
diff --git a/tests/test_retry_mechanism.py b/tests/test_retry_mechanism.py
diff --git a/tests/test_rollout_error_handling.py b/tests/test_rollout_error_handling.py

Original file line number	Diff line number	Diff line change
`@@ -138,7 +138,7 @@ async def _execute_with_semaphore(idx):`
`138`	`138`	`if trajectory.terminated:`
`139`	`139`	`if trajectory.termination_reason == TerminationReason.ERROR:`
`140`	`140`	`evaluation_row.rollout_status.status = "error"`
`141`		`- evaluation_row.rollout_status.error_message = trajectory.control_plane_summary.get(`
	`141`	`+ evaluation_row.rollout_status.termination_reason = trajectory.control_plane_summary.get(`
`142`	`142`	`"error_message", None`
`143`	`143`	`)`
`144`	`144`	`else:`