custom rollout exception and termination_reason=interrupted

Author: mayinghan

eval_protocol/mcp/execution/manager.py

@@ -12,17 +12,16 @@
 import threading
 import time
 from dataclasses import asdict
-from typing import TYPE_CHECKING, Any, AsyncIterator, Callable, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
 
 import anyio
-import httpx
 from openai.types import CompletionUsage
 
 from vendor.tau2.data_model.message import AssistantMessage, UserMessage
 from vendor.tau2.user.user_simulator import UserSimulator
 
 from ...models import EvaluationRow, InputMetadata, Message
-from ...types import MCPSession, MCPToolCall, TerminationReason, Trajectory
+from ...types import TerminationReason, Trajectory, RolloutTerminationException
 
 if TYPE_CHECKING:
     from ..session.manager import GeneralMCPVectorEnv
@@ -107,7 +106,7 @@ async def _execute_with_semaphore(idx):
                 )
 
                 # Convert trajectory to EvaluationRow immediately
-                evaluation_row = evaluation_rows[idx]
+                evaluation_row: EvaluationRow = evaluation_rows[idx]
 
                 # Handle multimodal content by extracting text from complex content structures
                 messages = []
@@ -137,14 +136,16 @@ async def _execute_with_semaphore(idx):
                 }
 
                 if trajectory.terminated:
+                    evaluation_row.rollout_status.termination_reason = trajectory.termination_reason
+                    # preserve the true error mesage if there are any
+                    if trajectory.control_plane_summary.get("error_message"):
+                        evaluation_row.rollout_status.extra_info = {
+                            "error_message": trajectory.control_plane_summary.get("error_message")
+                        }
                     if trajectory.termination_reason == TerminationReason.ERROR:
                         evaluation_row.rollout_status.status = "error"
-                        evaluation_row.rollout_status.termination_reason = trajectory.control_plane_summary.get(
-                            "error_message", None
-                        )
                     else:
                         evaluation_row.rollout_status.status = "finished"
-                        evaluation_row.rollout_status.termination_reason = trajectory.termination_reason
                 else:
                     evaluation_row.rollout_status.status = "running"
 
@@ -437,17 +438,16 @@ async def _execute_rollout(
             logger.info(
                 f"✅ Rollout {rollout_idx} completed: {trajectory.steps} steps, reward: {trajectory.total_reward:.2f}, termination: {trajectory.termination_reason}, in thread {threading.current_thread().name}"
             )
-
-        except asyncio.CancelledError:
+        except RolloutTerminationException as e:
+            trajectory.terminated = True
+            trajectory.termination_reason = TerminationReason.INTERRUPTED
+            trajectory.control_plane_summary.update({"error_message": str(e)})
+            logger.error(f"🚨 Rollout {rollout_idx} terminated due to user exception: {str(e)}", exc_info=True)
+        except asyncio.CancelledError:  # need explicit catch since it doesn't inherit from Exception
             failure_reason = "asyncio context cancelled"
             logger.error(
                 f"🚨 Error in rollout {session.dataset_row.id} {rollout_idx}: {failure_reason}", exc_info=True
             )
-        except (anyio.ClosedResourceError, anyio.BrokenResourceError):
-            failure_reason = "anyioconnection/resource error"
-            logger.error(
-                f"🚨 Error in rollout {session.dataset_row.id} {rollout_idx}: {failure_reason}", exc_info=True
-            )
         except Exception as e:
             error_msg = str(e) if str(e) else f"{type(e).__name__}: Unexpected error"
             logger.error(f"🚨 Error in rollout {session.dataset_row.id} {rollout_idx}: {error_msg}", exc_info=True)
@@ -461,7 +461,6 @@ async def _execute_rollout(
                 await envs.connection_manager.reset_session(session)
             except Exception as e:
                 logger.warning(f"Failed to reset session {session.session_id}: {type(e).__name__}: {e}", exc_info=True)
-
             try:
                 await envs.connection_manager.close_session(session)
             except Exception as e:

eval_protocol/models.py

@@ -11,6 +11,7 @@
 
 from eval_protocol.get_pep440_version import get_pep440_version
 from eval_protocol.human_id import generate_id
+from eval_protocol.types import TerminationReason
 
 
 class ChatCompletionContentPartTextParam(BaseModel):
@@ -285,14 +286,14 @@ class RolloutStatus(BaseModel):
 
     """
     running: Unfinished rollout which is still in progress.
-    finished: Rollout finished successfully.
-    error: Rollout failed.
-    stopped: Rollout terminated unexpectedly (e.g. max step, control plane signal, user stop).
+    finished: Rollout finished.
+    error: Rollout failed due to unexpected error. The rollout record should be discard.
     """
     status: Literal["running", "finished", "error"] = Field("running", description="Status of the rollout.")
-    termination_reason: Optional[str] = Field(
-        "", description="reason of the rollout status, mapped to values in TerminationReason"
+    termination_reason: Optional[TerminationReason] = Field(
+        None, description="reason of the rollout status, mapped to values in TerminationReason"
     )
+    extra_info: Optional[Dict[str, Any]] = Field(None, description="Extra information about the rollout status.")
 
 
 class EvaluationRow(BaseModel):

eval_protocol/types/__init__.py

@@ -1,3 +1,4 @@
 from .types import DatasetRow, MCPSession, MCPToolCall, TerminationReason, Trajectory
+from .errors import RolloutTerminationException
 
-__all__ = ["MCPSession", "MCPToolCall", "TerminationReason", "Trajectory", "DatasetRow"]
+__all__ = ["MCPSession", "MCPToolCall", "TerminationReason", "Trajectory", "DatasetRow", "RolloutTerminationException"]

eval_protocol/types/errors.py

@@ -0,0 +1,13 @@
+class RolloutTerminationException(Exception):
+    """
+    Exception raised during rollout. This error means the rollout needs to be terminated and its not retriable,
+    and the trajectory need to be perserved for future evalution or analysis.
+
+    For example, if the policy (llm) returns 400 User error, we need to end the rollout but keep the trajectory.
+    It differs from other exceptions such as network error, which are retriable and the trajectory should be discarded if
+    it fails eventually after retries.
+
+    This will cause trajectory.termination_reason to be set to TerminationReason.INTERRUPTED.
+    """
+
+    pass

eval_protocol/types/types.py

@@ -13,6 +13,7 @@ class TerminationReason(str, Enum):
     CONTROL_PLANE_SIGNAL: Trajectory ends because the control plane signals termination (e.g. env goal reached or failure condition)
     USER_STOP: Trajectory ends because the simulated user signals to stop
     ERROR: Trajectory ends because of an error
+    INTERRUPTED: Trajectory is interrupted by some non-system related reason (e.g. policy returns unexpected response and we need to terminate the rollout).
     STOP: Trajectory ends by the policy (mapped to llm response stop reason "stop")
     LENGTH: Trajectory ends by the policy (mapped to llm response stop reason "length")
     TOOL_CALLS: Trajectory ends by the policy with a hanging tool call response (mapped to llm response stop reason "tool_calls")
@@ -22,6 +23,7 @@ class TerminationReason(str, Enum):
     CONTROL_PLANE_SIGNAL = "control_plane_signal"
     USER_STOP = "user_stop"
     ERROR = "error"
+    INTERRUPTED = "interrupted"
     STOP = "stop"
     LENGTH = "length"
     TOOL_CALLS = "tool_calls"