diff --git a/packages/prime-sandboxes/src/prime_sandboxes/__init__.py b/packages/prime-sandboxes/src/prime_sandboxes/__init__.py index 0bd1b8cd..f619ec31 100644 --- a/packages/prime-sandboxes/src/prime_sandboxes/__init__.py +++ b/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", diff --git a/packages/prime-sandboxes/src/prime_sandboxes/exceptions.py b/packages/prime-sandboxes/src/prime_sandboxes/exceptions.py index c98f537c..357a0fb8 100644 --- a/packages/prime-sandboxes/src/prime_sandboxes/exceptions.py +++ b/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) diff --git a/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py b/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py index 645dbbb8..886abdbf 100644 --- a/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py +++ b/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", "?") diff --git a/packages/prime/src/prime_cli/api/rl.py b/packages/prime/src/prime_cli/api/rl.py index a8beba46..bf831afd 100644 --- a/packages/prime/src/prime_cli/api/rl.py +++ b/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"), diff --git a/packages/prime/src/prime_cli/commands/rl.py b/packages/prime/src/prime_cli/commands/rl.py index 4439147e..7da6e60b 100644 --- a/packages/prime/src/prime_cli/commands/rl.py +++ b/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( diff --git a/packages/prime/src/prime_cli/utils/env_metadata.py b/packages/prime/src/prime_cli/utils/env_metadata.py index fe07be4f..cde73fd7 100644 --- a/packages/prime/src/prime_cli/utils/env_metadata.py +++ b/packages/prime/src/prime_cli/utils/env_metadata.py @@ -1,4 +1,5 @@ """Utilities for reading and managing environment metadata.""" + import json from pathlib import Path from typing import Any, Dict, Optional @@ -6,16 +7,16 @@ 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,7 +42,7 @@ 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) @@ -49,21 +50,21 @@ def find_environment_metadata( 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 diff --git a/packages/prime/src/prime_cli/utils/env_vars.py b/packages/prime/src/prime_cli/utils/env_vars.py index a5fa4d31..b234e258 100644 --- a/packages/prime/src/prime_cli/utils/env_vars.py +++ b/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 -