diff --git a/examples/04_streaming_output.py b/examples/04_streaming_output.py index c89ef938..3f4fe9b0 100644 --- a/examples/04_streaming_output.py +++ b/examples/04_streaming_output.py @@ -28,13 +28,13 @@ def main(): # Stream output in real-time result = sandbox.exec( - '''python3 -c " + '''python3 -u -c " import time for i in range(5): print(f'Line {i+1}') time.sleep(0.5) "''', - on_stdout=lambda data: print(data.strip(), end=" "), + on_stdout=lambda data: print(data.strip()), on_stderr=lambda data: print(f"ERR: {data.strip()}"), ) print(f"\nExit code: {result.exit_code}") diff --git a/examples/04_streaming_output_async.py b/examples/04_streaming_output_async.py index f3e24f29..8301a1c4 100644 --- a/examples/04_streaming_output_async.py +++ b/examples/04_streaming_output_async.py @@ -30,13 +30,13 @@ async def main(): # Stream output in real-time result = await sandbox.exec( - '''python3 -c " + '''python3 -u -c " import time for i in range(5): print(f'Line {i+1}') time.sleep(0.5) "''', - on_stdout=lambda data: print(data.strip(), end=" "), + on_stdout=lambda data: print(data.strip()), on_stderr=lambda data: print(f"ERR: {data.strip()}"), ) print(f"\nExit code: {result.exit_code}") diff --git a/koyeb/sandbox/__init__.py b/koyeb/sandbox/__init__.py index e9185f95..4688c7c6 100644 --- a/koyeb/sandbox/__init__.py +++ b/koyeb/sandbox/__init__.py @@ -19,7 +19,7 @@ ) from .filesystem import FileInfo, SandboxFilesystem from .sandbox import AsyncSandbox, ExposedPort, ProcessInfo, Sandbox -from .utils import SandboxDeploymentError, SandboxError, SandboxTimeoutError +from .utils import SandboxDeploymentError, SandboxError, SandboxServiceError, SandboxTimeoutError __all__ = [ "Sandbox", @@ -33,6 +33,7 @@ "SandboxStatus", "SandboxDeploymentError", "SandboxError", + "SandboxServiceError", "SandboxTimeoutError", "CommandResult", "CommandStatus", diff --git a/koyeb/sandbox/exec.py b/koyeb/sandbox/exec.py index 6328f5a9..c357fd7f 100644 --- a/koyeb/sandbox/exec.py +++ b/koyeb/sandbox/exec.py @@ -117,86 +117,66 @@ def __call__( stderr_buffer = [] exit_code = 0 - try: - client = self._get_client() - for event in client.run_streaming( - cmd=command, cwd=cwd, env=env, timeout=float(timeout) - ): - if "stream" in event: - stream_type = event["stream"] - data = event["data"] - - if stream_type == "stdout": - stdout_buffer.append(data) - if on_stdout: - on_stdout(data) - elif stream_type == "stderr": - stderr_buffer.append(data) - if on_stderr: - on_stderr(data) - elif "code" in event: - exit_code = event["code"] - elif "error" in event and isinstance(event["error"], str): - # Error starting command - return CommandResult( - stdout="", - stderr=event["error"], - exit_code=1, - status=CommandStatus.FAILED, - duration=time.time() - start_time, - command=command, - ) - - return CommandResult( - stdout="".join(stdout_buffer), - stderr="".join(stderr_buffer), - exit_code=exit_code, - status=( - CommandStatus.FINISHED - if exit_code == 0 - else CommandStatus.FAILED - ), - duration=time.time() - start_time, - command=command, - ) - except Exception as e: - return CommandResult( - stdout="", - stderr=f"Command execution failed: {str(e)}", - exit_code=1, - status=CommandStatus.FAILED, - duration=time.time() - start_time, - command=command, - ) - - # Use regular run for non-streaming execution - try: client = self._get_client() - response = client.run(cmd=command, cwd=cwd, env=env, timeout=float(timeout)) - - stdout = response.get("stdout", "") - stderr = response.get("stderr", "") - exit_code = response.get("code", 0) + for event in client.run_streaming( + cmd=command, cwd=cwd, env=env, timeout=float(timeout) + ): + if "stream" in event: + stream_type = event["stream"] + data = event["data"] + + if stream_type == "stdout": + stdout_buffer.append(data) + if on_stdout: + on_stdout(data) + elif stream_type == "stderr": + stderr_buffer.append(data) + if on_stderr: + on_stderr(data) + elif "code" in event: + exit_code = event["code"] + elif "error" in event and isinstance(event["error"], str): + # Error starting command + return CommandResult( + stdout="", + stderr=event["error"], + exit_code=1, + status=CommandStatus.FAILED, + duration=time.time() - start_time, + command=command, + ) return CommandResult( - stdout=stdout, - stderr=stderr, + stdout="".join(stdout_buffer), + stderr="".join(stderr_buffer), exit_code=exit_code, status=( - CommandStatus.FINISHED if exit_code == 0 else CommandStatus.FAILED + CommandStatus.FINISHED + if exit_code == 0 + else CommandStatus.FAILED ), duration=time.time() - start_time, command=command, ) - except Exception as e: - return CommandResult( - stdout="", - stderr=f"Command execution failed: {str(e)}", - exit_code=1, - status=CommandStatus.FAILED, - duration=time.time() - start_time, - command=command, - ) + + # Use regular run for non-streaming execution + client = self._get_client() + response = client.run(cmd=command, cwd=cwd, env=env, timeout=float(timeout)) + + stdout = response.get("stdout", "") + stderr = response.get("stderr", "") + exit_code = response.get("code", 0) + + return CommandResult( + stdout=stdout, + stderr=stderr, + exit_code=exit_code, + status=( + CommandStatus.FINISHED if exit_code == 0 else CommandStatus.FAILED + ), + duration=time.time() - start_time, + command=command, + ) class AsyncSandboxExecutor(SandboxExecutor): @@ -256,85 +236,65 @@ async def __call__( stderr_buffer: List[str] = [] exit_code = 0 - try: - client = self._get_async_client() - - async for event in client.run_streaming( - cmd=command, cwd=cwd, env=env, timeout=float(timeout) - ): - if "stream" in event: - stream_type = event["stream"] - data = event["data"] - - if stream_type == "stdout": - stdout_buffer.append(data) - if on_stdout: - on_stdout(data) - elif stream_type == "stderr": - stderr_buffer.append(data) - if on_stderr: - on_stderr(data) - elif "code" in event: - exit_code = event["code"] - elif "error" in event and isinstance(event["error"], str): - return CommandResult( - stdout="", - stderr=event["error"], - exit_code=1, - status=CommandStatus.FAILED, - duration=time.time() - start_time, - command=command, - ) - - return CommandResult( - stdout="".join(stdout_buffer), - stderr="".join(stderr_buffer), - exit_code=exit_code, - status=( - CommandStatus.FINISHED - if exit_code == 0 - else CommandStatus.FAILED - ), - duration=time.time() - start_time, - command=command, - ) - except Exception as e: - return CommandResult( - stdout="", - stderr=f"Command execution failed: {str(e)}", - exit_code=1, - status=CommandStatus.FAILED, - duration=time.time() - start_time, - command=command, - ) - - # Use native async for non-streaming execution - try: client = self._get_async_client() - response = await client.run( - cmd=command, cwd=cwd, env=env, timeout=float(timeout) - ) - stdout = response.get("stdout", "") - stderr = response.get("stderr", "") - exit_code = response.get("code", 0) + async for event in client.run_streaming( + cmd=command, cwd=cwd, env=env, timeout=float(timeout) + ): + if "stream" in event: + stream_type = event["stream"] + data = event["data"] + + if stream_type == "stdout": + stdout_buffer.append(data) + if on_stdout: + on_stdout(data) + elif stream_type == "stderr": + stderr_buffer.append(data) + if on_stderr: + on_stderr(data) + elif "code" in event: + exit_code = event["code"] + elif "error" in event and isinstance(event["error"], str): + return CommandResult( + stdout="", + stderr=event["error"], + exit_code=1, + status=CommandStatus.FAILED, + duration=time.time() - start_time, + command=command, + ) return CommandResult( - stdout=stdout, - stderr=stderr, + stdout="".join(stdout_buffer), + stderr="".join(stderr_buffer), exit_code=exit_code, status=( - CommandStatus.FINISHED if exit_code == 0 else CommandStatus.FAILED + CommandStatus.FINISHED + if exit_code == 0 + else CommandStatus.FAILED ), duration=time.time() - start_time, command=command, ) - except Exception as e: - return CommandResult( - stdout="", - stderr=f"Command execution failed: {str(e)}", - exit_code=1, - status=CommandStatus.FAILED, - duration=time.time() - start_time, - command=command, - ) + + # Use native async for non-streaming execution + client = self._get_async_client() + response = await client.run( + cmd=command, cwd=cwd, env=env, timeout=float(timeout) + ) + + stdout = response.get("stdout", "") + stderr = response.get("stderr", "") + exit_code = response.get("code", 0) + + return CommandResult( + stdout=stdout, + stderr=stderr, + exit_code=exit_code, + status=( + CommandStatus.FINISHED if exit_code == 0 else CommandStatus.FAILED + ), + duration=time.time() - start_time, + command=command, + ) diff --git a/koyeb/sandbox/executor_client.py b/koyeb/sandbox/executor_client.py index afa38c01..f398957e 100644 --- a/koyeb/sandbox/executor_client.py +++ b/koyeb/sandbox/executor_client.py @@ -13,7 +13,7 @@ import httpx -from .utils import DEFAULT_HTTP_TIMEOUT +from .utils import DEFAULT_HTTP_TIMEOUT, SandboxServiceError, SandboxTimeoutError logger = logging.getLogger(__name__) @@ -80,7 +80,7 @@ def __init__( def close(self) -> None: """Close the HTTP client and release resources.""" - if not self._closed and hasattr(self, "_client"): + if not getattr(self, "_closed", True) and hasattr(self, "_client"): self._client.close() self._closed = True @@ -94,7 +94,7 @@ def __exit__(self, exc_type, exc_val, exc_tb) -> None: def __del__(self): """Clean up client on deletion (fallback, not guaranteed to run).""" - if not self._closed: + if not getattr(self, "_closed", True): self.close() def _request_with_retry( @@ -156,17 +156,26 @@ def _request_with_retry( backoff *= 2 last_exception = e continue + if e.response.status_code >= 500: + raise SandboxServiceError( + status_code=e.response.status_code, + message=e.response.text, + ) from e raise except httpx.TimeoutException as e: - logger.warning(f"Request timeout after {self.timeout}s: {e}") - raise + raise SandboxTimeoutError( + f"Request timed out after {kwargs.get('timeout', self.timeout)}s" + ) from e except httpx.RequestError as e: logger.warning(f"Request failed: {e}") raise # If we exhausted all retries, raise the last exception if last_exception: - raise last_exception + raise SandboxServiceError( + status_code=last_exception.response.status_code, + message=last_exception.response.text, + ) from last_exception def health(self) -> Dict[str, str]: """ @@ -268,17 +277,29 @@ def run_streaming( if env is not None: payload["env"] = env - with self._client.stream( - "POST", - f"{self.base_url}/run_streaming", - json=payload, - timeout=timeout if timeout is not None else self.timeout, - ) as response: - response.raise_for_status() - for line in response.iter_lines(): - event = _parse_sse_line(line) - if event is not None: - yield event + request_timeout = timeout if timeout is not None else self.timeout + try: + with self._client.stream( + "POST", + f"{self.base_url}/run_streaming", + json=payload, + timeout=request_timeout, + ) as response: + if response.status_code >= 500: + response.read() + raise SandboxServiceError( + status_code=response.status_code, + message=response.text, + ) + response.raise_for_status() + for line in response.iter_lines(): + event = _parse_sse_line(line) + if event is not None: + yield event + except httpx.TimeoutException as e: + raise SandboxTimeoutError( + f"Request timed out after {request_timeout}s" + ) from e def write_file(self, path: str, content: str) -> Dict[str, Any]: """ @@ -540,7 +561,7 @@ def __init__( async def close(self) -> None: """Close the HTTP client and release resources.""" - if not self._closed and hasattr(self, "_client"): + if not getattr(self, "_closed", True) and hasattr(self, "_client"): await self._client.aclose() self._closed = True @@ -611,17 +632,26 @@ async def _request_with_retry( backoff *= 2 last_exception = e continue + if e.response.status_code >= 500: + raise SandboxServiceError( + status_code=e.response.status_code, + message=e.response.text, + ) from e raise except httpx.TimeoutException as e: - logger.warning(f"Request timeout after {self.timeout}s: {e}") - raise + raise SandboxTimeoutError( + f"Request timed out after {kwargs.get('timeout', self.timeout)}s" + ) from e except httpx.RequestError as e: logger.warning(f"Request failed: {e}") raise # If we exhausted all retries, raise the last exception if last_exception: - raise last_exception + raise SandboxServiceError( + status_code=last_exception.response.status_code, + message=last_exception.response.text, + ) from last_exception async def health(self) -> Dict[str, str]: """ @@ -723,17 +753,29 @@ async def run_streaming( if env is not None: payload["env"] = env - async with self._client.stream( - "POST", - f"{self.base_url}/run_streaming", - json=payload, - timeout=timeout if timeout is not None else self.timeout, - ) as response: - response.raise_for_status() - async for line in response.aiter_lines(): - event = _parse_sse_line(line) - if event is not None: - yield event + request_timeout = timeout if timeout is not None else self.timeout + try: + async with self._client.stream( + "POST", + f"{self.base_url}/run_streaming", + json=payload, + timeout=request_timeout, + ) as response: + if response.status_code >= 500: + await response.aread() + raise SandboxServiceError( + status_code=response.status_code, + message=response.text, + ) + response.raise_for_status() + async for line in response.aiter_lines(): + event = _parse_sse_line(line) + if event is not None: + yield event + except httpx.TimeoutException as e: + raise SandboxTimeoutError( + f"Request timed out after {request_timeout}s" + ) from e async def write_file(self, path: str, content: str) -> Dict[str, Any]: """ diff --git a/koyeb/sandbox/filesystem.py b/koyeb/sandbox/filesystem.py index 63eaf66f..74806873 100644 --- a/koyeb/sandbox/filesystem.py +++ b/koyeb/sandbox/filesystem.py @@ -14,6 +14,7 @@ from .executor_client import AsyncSandboxClient, SandboxClient from .utils import ( SandboxError, + SandboxServiceError, check_error_message, escape_shell_arg, ) @@ -95,9 +96,9 @@ def write_file( if response.get("error"): error_msg = response.get("error", "Unknown error") raise SandboxFilesystemError(f"Failed to write file: {error_msg}") + except (SandboxServiceError, SandboxFilesystemError): + raise except Exception as e: - if isinstance(e, SandboxFilesystemError): - raise raise SandboxFilesystemError(f"Failed to write file: {str(e)}") from e def read_file(self, path: str, encoding: str = "utf-8") -> FileInfo: @@ -129,7 +130,7 @@ def read_file(self, path: str, encoding: str = "utf-8") -> FileInfo: else: content = content_str return FileInfo(content=content, encoding=encoding) - except (SandboxFileNotFoundError, SandboxFilesystemError): + except (SandboxServiceError, SandboxFileNotFoundError, SandboxFilesystemError): raise except Exception as e: error_msg = str(e) @@ -155,7 +156,7 @@ def mkdir(self, path: str) -> None: if check_error_message(error_msg, "FILE_EXISTS"): raise SandboxFileExistsError(f"Directory already exists: {path}") raise SandboxFilesystemError(f"Failed to create directory: {error_msg}") - except (SandboxFileExistsError, SandboxFilesystemError): + except (SandboxServiceError, SandboxFileExistsError, SandboxFilesystemError): raise except Exception as e: error_msg = str(e) @@ -186,7 +187,7 @@ def list_dir(self, path: str = ".") -> List[str]: raise SandboxFilesystemError(f"Failed to list directory: {error_msg}") entries = response.get("entries", []) return entries - except (SandboxFileNotFoundError, SandboxFilesystemError): + except (SandboxServiceError, SandboxFileNotFoundError, SandboxFilesystemError): raise except Exception as e: error_msg = str(e) @@ -212,7 +213,7 @@ def delete_file(self, path: str) -> None: if check_error_message(error_msg, "NO_SUCH_FILE"): raise SandboxFileNotFoundError(f"File not found: {path}") raise SandboxFilesystemError(f"Failed to delete file: {error_msg}") - except (SandboxFileNotFoundError, SandboxFilesystemError): + except (SandboxServiceError, SandboxFileNotFoundError, SandboxFilesystemError): raise except Exception as e: error_msg = str(e) @@ -238,7 +239,7 @@ def delete_dir(self, path: str) -> None: if check_error_message(error_msg, "DIR_NOT_EMPTY"): raise SandboxFilesystemError(f"Directory not empty: {path}") raise SandboxFilesystemError(f"Failed to delete directory: {error_msg}") - except (SandboxFileNotFoundError, SandboxFilesystemError): + except (SandboxServiceError, SandboxFileNotFoundError, SandboxFilesystemError): raise except Exception as e: error_msg = str(e) @@ -467,9 +468,9 @@ async def write_file( if response.get("error"): error_msg = response.get("error", "Unknown error") raise SandboxFilesystemError(f"Failed to write file: {error_msg}") + except (SandboxServiceError, SandboxFilesystemError): + raise except Exception as e: - if isinstance(e, SandboxFilesystemError): - raise raise SandboxFilesystemError(f"Failed to write file: {str(e)}") from e async def read_file(self, path: str, encoding: str = "utf-8") -> FileInfo: @@ -501,7 +502,7 @@ async def read_file(self, path: str, encoding: str = "utf-8") -> FileInfo: else: content = content_str return FileInfo(content=content, encoding=encoding) - except (SandboxFileNotFoundError, SandboxFilesystemError): + except (SandboxServiceError, SandboxFileNotFoundError, SandboxFilesystemError): raise except Exception as e: error_msg = str(e) @@ -527,7 +528,7 @@ async def mkdir(self, path: str) -> None: if check_error_message(error_msg, "FILE_EXISTS"): raise SandboxFileExistsError(f"Directory already exists: {path}") raise SandboxFilesystemError(f"Failed to create directory: {error_msg}") - except (SandboxFileExistsError, SandboxFilesystemError): + except (SandboxServiceError, SandboxFileExistsError, SandboxFilesystemError): raise except Exception as e: error_msg = str(e) @@ -558,7 +559,7 @@ async def list_dir(self, path: str = ".") -> List[str]: raise SandboxFilesystemError(f"Failed to list directory: {error_msg}") entries = response.get("entries", []) return entries - except (SandboxFileNotFoundError, SandboxFilesystemError): + except (SandboxServiceError, SandboxFileNotFoundError, SandboxFilesystemError): raise except Exception as e: error_msg = str(e) @@ -584,7 +585,7 @@ async def delete_file(self, path: str) -> None: if check_error_message(error_msg, "NO_SUCH_FILE"): raise SandboxFileNotFoundError(f"File not found: {path}") raise SandboxFilesystemError(f"Failed to delete file: {error_msg}") - except (SandboxFileNotFoundError, SandboxFilesystemError): + except (SandboxServiceError, SandboxFileNotFoundError, SandboxFilesystemError): raise except Exception as e: error_msg = str(e) @@ -610,7 +611,7 @@ async def delete_dir(self, path: str) -> None: if check_error_message(error_msg, "DIR_NOT_EMPTY"): raise SandboxFilesystemError(f"Directory not empty: {path}") raise SandboxFilesystemError(f"Failed to delete directory: {error_msg}") - except (SandboxFileNotFoundError, SandboxFilesystemError): + except (SandboxServiceError, SandboxFileNotFoundError, SandboxFilesystemError): raise except Exception as e: error_msg = str(e) diff --git a/koyeb/sandbox/utils.py b/koyeb/sandbox/utils.py index 2956616c..4482386f 100644 --- a/koyeb/sandbox/utils.py +++ b/koyeb/sandbox/utils.py @@ -628,3 +628,11 @@ class SandboxTimeoutError(SandboxError): class SandboxDeploymentError(SandboxError): """Raised when a sandbox deployment reaches an error state""" + + +class SandboxServiceError(SandboxError): + """Raised when the sandbox executor returns an HTTP 5xx error""" + + def __init__(self, status_code: int, message: str): + self.status_code = status_code + super().__init__(f"Sandbox service error ({status_code}): {message}")