ENG-2736 - Remove Sandbox Unresponsive

packages/prime-sandboxes/src/prime_sandboxes/__init__.py

@@ -20,7 +20,6 @@
     SandboxNotRunningError,
     SandboxOOMError,
     SandboxTimeoutError,
-    SandboxUnresponsiveError,
     UploadTimeoutError,
 )
 from .models import (
@@ -45,7 +44,7 @@
 )
 from .sandbox import AsyncSandboxClient, AsyncTemplateClient, SandboxClient, TemplateClient
 
-__version__ = "0.2.11"
+__version__ = "0.2.12"
 
 # Deprecated alias for backward compatibility
 TimeoutError = APITimeoutError
@@ -90,7 +89,6 @@
     "SandboxTimeoutError",
     "SandboxImagePullError",
     "SandboxNotRunningError",
-    "SandboxUnresponsiveError",
     "CommandTimeoutError",
     "UploadTimeoutError",
     "DownloadTimeoutError",

packages/prime-sandboxes/src/prime_sandboxes/exceptions.py

@@ -68,19 +68,3 @@ class SandboxImagePullError(SandboxNotRunningError):
     """Raised when Docker image cannot be pulled."""
 
     pass
-
-
-class SandboxUnresponsiveError(CommandTimeoutError):
-    """Raised when sandbox appears running but commands time out unexpectedly."""
-
-    def __init__(
-        self,
-        sandbox_id: str,
-        command: str,
-        message: str,
-        sandbox_status: str | None = None,
-    ):
-        self.sandbox_id = sandbox_id
-        self.command = command
-        self.sandbox_status = sandbox_status
-        RuntimeError.__init__(self, message)

packages/prime-sandboxes/src/prime_sandboxes/sandbox.py

@@ -28,7 +28,6 @@
     SandboxNotRunningError,
     SandboxOOMError,
     SandboxTimeoutError,
-    SandboxUnresponsiveError,
     UploadTimeoutError,
 )
 from .models import (
@@ -109,14 +108,6 @@ def _build_terminated_message(command: str, ctx: dict) -> str:
     return " ".join(parts)
 
 
-def _build_unresponsive_message(command: str, timeout: int) -> str:
-    """Build helpful error message for unresponsive sandbox."""
-    cmd_preview = command[:50] + "..." if len(command) > 50 else command
-    msg = f"Command '{cmd_preview}' timed out after {timeout}s."
-
-    return msg
-
-
 def _raise_not_running_error(
     sandbox_id: str,
     ctx: dict,
@@ -447,6 +438,8 @@ def execute_command(
         }
 
         try:
+            # The + 2 accounts for connection creation and closing. Prevents any command
+            # running close to its `effective_timeout` from being killed prematurely
             client_timeout = effective_timeout + 2
             response = self._gateway_post(
                 url, headers=headers, timeout=client_timeout, json=payload
@@ -457,12 +450,7 @@ def execute_command(
             ctx = self._get_sandbox_error_context(sandbox_id)
             if ctx["status"] in ("TERMINATED", "ERROR", "TIMEOUT"):
                 _raise_not_running_error(sandbox_id, ctx, command=command, cause=e)
-            raise SandboxUnresponsiveError(
-                sandbox_id=sandbox_id,
-                command=command,
-                message=_build_unresponsive_message(command, effective_timeout),
-                sandbox_status=ctx["status"],
-            ) from e
+            raise CommandTimeoutError(sandbox_id, command, effective_timeout) from e
         except httpx.HTTPStatusError as e:
             resp = getattr(e, "response", None)
             status = getattr(resp, "status_code", "?")
@@ -1035,12 +1023,7 @@ async def execute_command(
             ctx = await self._get_sandbox_error_context(sandbox_id)
             if ctx["status"] in ("TERMINATED", "ERROR", "TIMEOUT"):
                 _raise_not_running_error(sandbox_id, ctx, command=command, cause=e)
-            raise SandboxUnresponsiveError(
-                sandbox_id=sandbox_id,
-                command=command,
-                message=_build_unresponsive_message(command, effective_timeout),
-                sandbox_status=ctx["status"],
-            ) from e
+            raise CommandTimeoutError(sandbox_id, command, effective_timeout) from e
         except httpx.HTTPStatusError as e:
             resp = getattr(e, "response", None)
             status = getattr(resp, "status_code", "?")

packages/prime/src/prime_cli/api/rl.py

@@ -311,9 +311,7 @@ def get_distributions(
             if step is not None:
                 params["step"] = step
 
-            response = self.client.get(
-                f"/rft/runs/{run_id}/distributions", params=params
-            )
+            response = self.client.get(f"/rft/runs/{run_id}/distributions", params=params)
             return {
                 "bins": response.get("bins", []),
                 "step": response.get("step"),

packages/prime/src/prime_cli/commands/rl.py

@@ -440,11 +440,9 @@ def warn(msg: str) -> None:
     wandb_configured = cfg.wandb.entity or cfg.wandb.project
     if wandb_configured and (not secrets or "WANDB_API_KEY" not in secrets):
         console.print("[red]Configuration Error:[/red]")
-        console.print(
-            "  WANDB_API_KEY is required when W&B monitoring is configured.\n"
-        )
+        console.print("  WANDB_API_KEY is required when W&B monitoring is configured.\n")
         console.print("Provide it via:")
-        console.print("  - env_files in your config: env_files = [\"secrets.env\"]")
+        console.print('  - env_files in your config: env_files = ["secrets.env"]')
         console.print("  - CLI flag: --env-file secrets.env")
         console.print("  - CLI flag: -e WANDB_API_KEY=your-key")
         console.print(

packages/prime/src/prime_cli/utils/env_metadata.py

@@ -1,21 +1,22 @@
 """Utilities for reading and managing environment metadata."""
+
 import json
 from pathlib import Path
 from typing import Any, Dict, Optional
 
 
 def get_environment_metadata(env_path: Path) -> Optional[Dict[str, Any]]:
     """Read environment metadata from .prime/.env-metadata.json with backwards compatibility.
-    
+
     Checks both the new location (.prime/.env-metadata.json) and old location
     (.env-metadata.json) for backwards compatibility.
-    
+
     This function only checks the provided path - it does not search multiple directories.
     Use find_environment_metadata() if you need to search multiple locations.
-    
+
     Args:
         env_path: Path to the environment directory
-        
+
     Returns:
         Dictionary containing environment metadata, or None if not found
     """
@@ -24,10 +25,10 @@ def get_environment_metadata(env_path: Path) -> Optional[Dict[str, Any]]:
     if not metadata_path.exists():
         # Fall back to old location for backwards compatibility
         metadata_path = env_path / ".env-metadata.json"
-    
+
     if not metadata_path.exists():
         return None
-    
+
     try:
         with open(metadata_path, "r") as f:
             return json.load(f)
@@ -41,29 +42,29 @@ def find_environment_metadata(
     module_name: Optional[str] = None,
 ) -> Optional[Dict[str, Any]]:
     """Search for environment metadata in multiple common locations.
-    
+
     Searches directories in the following order:
     1. env_path (if provided)
     2. ./environments/{module_name} (if module_name provided)
     3. ./environments/{env_name} (if env_name provided)
     4. ./{env_name} (if env_name provided)
     5. ./{module_name} (if module_name provided)
     6. ./ (current directory)
-    
+
     Args:
         env_name: Environment name (e.g., "simpleqa")
         env_path: Optional explicit path to check first
         module_name: Optional module name (e.g., "simple_qa" for env_name "simpleqa")
-        
+
     Returns:
         Dictionary containing environment metadata from the first location found, or None
     """
     possible_env_dirs = []
-    
+
     # 1. Check explicit path first if provided
     if env_path:
         possible_env_dirs.append(env_path)
-    
+
     # 2-5. Check common locations based on provided names
     if module_name:
         possible_env_dirs.append(Path("./environments") / module_name)
@@ -72,15 +73,14 @@ def find_environment_metadata(
         possible_env_dirs.append(Path(".") / env_name)
     if module_name:
         possible_env_dirs.append(Path(".") / module_name)
-    
+
     # 6. Check current directory last
     possible_env_dirs.append(Path("."))
-    
+
     # Search through directories and return first match
     for env_dir in possible_env_dirs:
         metadata = get_environment_metadata(env_dir)
         if metadata:
             return metadata
-
-    return None
 
+    return None

packages/prime/src/prime_cli/utils/env_vars.py

@@ -35,8 +35,7 @@ def parse_env_file(
             if "=" not in line:
                 if on_warning:
                     on_warning(
-                        f"Skipping invalid line {line_num} in {file_path}: "
-                        f"missing '=' separator"
+                        f"Skipping invalid line {line_num} in {file_path}: missing '=' separator"
                     )
                 continue
             key, _, value = line.partition("=")
@@ -112,8 +111,7 @@ def parse_env_arg(
     value = os.environ.get(key)
     if value is None:
         raise EnvParseError(
-            f"Environment variable '{key}' is not set. "
-            f"Either set it or use KEY=VALUE syntax."
+            f"Environment variable '{key}' is not set. Either set it or use KEY=VALUE syntax."
         )
     return {key: value}
 
@@ -154,4 +152,3 @@ def collect_env_vars(
             result.update(parse_env_arg(arg, on_warning=on_warning))
 
     return result
-