diff --git a/packages/prime-sandboxes/README.md b/packages/prime-sandboxes/README.md index aa1c6851..2b41410f 100644 --- a/packages/prime-sandboxes/README.md +++ b/packages/prime-sandboxes/README.md @@ -256,7 +256,10 @@ sandbox_client.download_file(sandbox.id, "/app/model.pt", "./model.pt") ``` `get_background_jobs` is VM-only. Container sandboxes retain the existing -`get_background_job` polling behavior. +`get_background_job` polling behavior. Once an exit code is observed, completion +remains authoritative even if output retrieval exhausts its bounded retry +deadline: the unavailable stream is `None` and its `stdout_error` or +`stderr_error` field describes the retrieval failure. #### Async version diff --git a/packages/prime-sandboxes/src/prime_sandboxes/models.py b/packages/prime-sandboxes/src/prime_sandboxes/models.py index 141bc335..619efbfd 100644 --- a/packages/prime-sandboxes/src/prime_sandboxes/models.py +++ b/packages/prime-sandboxes/src/prime_sandboxes/models.py @@ -775,6 +775,8 @@ class BackgroundJobStatus(BaseModel): exit_code: Optional[int] = None stdout: Optional[str] = None stderr: Optional[str] = None + stdout_error: Optional[str] = None + stderr_error: Optional[str] = None stdout_truncated: bool = False stderr_truncated: bool = False diff --git a/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py b/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py index 6f3870a6..34e712d3 100644 --- a/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py +++ b/packages/prime-sandboxes/src/prime_sandboxes/sandbox.py @@ -3,6 +3,7 @@ import asyncio import functools import json +import math import os import random import re @@ -44,6 +45,7 @@ retry_if_exception_type, stop_after_attempt, wait_exponential, + wait_random_exponential, ) from ._connectrpc import GOOGLE_PROTOBUF_BINARY_CODEC @@ -193,6 +195,12 @@ def _network_update_payload( # Max bytes of stdout/stderr returned per background-job status check JOB_OUTPUT_TAIL_BYTES = 10 * 1024 * 1024 +# Keep a batch of simultaneously completed jobs from turning into a burst of +# gateway connections. Output reads share this client-wide limit and each job +# gets one deadline for its sequential stdout/stderr retrieval. +MAX_CONCURRENT_BACKGROUND_JOB_OUTPUT_READS = 20 +BACKGROUND_JOB_OUTPUT_FETCH_TIMEOUT_SECONDS = 45.0 + # Platform batch-status contracts cap one request at 100 identifiers. Concurrent # single-item waits are collected briefly so callers share a request without # adding a persistent worker to the client lifecycle. @@ -379,6 +387,31 @@ def _exception_chain(exc: BaseException) -> List[BaseException]: return chain +def _format_exception_diagnostic(exc: BaseException) -> str: + """Include nested OS errno details that transport wrappers otherwise hide.""" + diagnostic = f"{exc.__class__.__name__}: {exc}" + os_causes: List[str] = [] + seen: set[int] = set() + pending: List[BaseException] = [exc] + while pending: + error = pending.pop() + if id(error) in seen: + continue + seen.add(id(error)) + if isinstance(error, OSError) and error.errno is not None: + detail = f"{error.__class__.__name__}(errno={error.errno}): {error}" + if detail not in os_causes: + os_causes.append(detail) + cause = error.__cause__ or error.__context__ + if cause is not None: + pending.append(cause) + grouped = getattr(error, "exceptions", ()) + pending.extend(child for child in reversed(grouped) if isinstance(child, BaseException)) + if os_causes: + diagnostic += f"; OS cause: {'; '.join(os_causes)}" + return diagnostic + + def _is_retryable_reachability_error(exc: BaseException) -> bool: """Retry transport readiness failures, but surface local SDK defects.""" chain = _exception_chain(exc) @@ -473,7 +506,7 @@ def _is_retryable_read_file_error(exc: BaseException) -> bool: _read_file_retry = retry( retry=retry_if_exception(_is_retryable_read_file_error), stop=stop_after_attempt(4), - wait=wait_exponential(multiplier=1, min=1, max=30), + wait=wait_random_exponential(multiplier=1, min=1, max=30), reraise=True, ) @@ -888,6 +921,9 @@ def __init__(self, api_client: APIClient): self._background_job_status_batcher = _SyncRequestBatcher( self._fetch_background_job_statuses ) + self._background_job_output_semaphore = threading.BoundedSemaphore( + MAX_CONCURRENT_BACKGROUND_JOB_OUTPUT_READS + ) self._sandbox_status_batch_supported: Optional[bool] = None self._background_job_status_batch_supported: Optional[bool] = None @@ -1470,9 +1506,8 @@ def get_background_job( Args: sandbox_id: The sandbox ID job: The BackgroundJob handle from start_background_job() - timeout: Optional per-call timeout (in seconds) forwarded to the - underlying read_file calls. When None, the APIClient default - applies. + timeout: Optional output-retrieval deadline in seconds after the + exit file is observed. When None, a bounded SDK default applies. Returns: BackgroundJobStatus with completed flag, and exit_code/stdout if @@ -1486,20 +1521,6 @@ def read_or_empty(path: str) -> str: except SandboxFileNotFoundError: return "" - def read_output_tail(path: str) -> "tuple[str, bool]": - try: - response = self.read_file( - sandbox_id, - path, - timeout=timeout, - offset=-JOB_OUTPUT_TAIL_BYTES, - length=JOB_OUTPUT_TAIL_BYTES, - ) - # Servers without windowed-read support omit `truncated`. - return response.content, bool(response.truncated) - except SandboxFileNotFoundError: - return "", False - exit_content = read_or_empty(job.exit_file) if not exit_content.strip(): return BackgroundJobStatus(job_id=job.job_id, completed=False) @@ -1509,16 +1530,11 @@ def read_output_tail(path: str) -> "tuple[str, bool]": except ValueError: return BackgroundJobStatus(job_id=job.job_id, completed=False) - stdout, stdout_truncated = read_output_tail(job.stdout_log_file) - stderr, stderr_truncated = read_output_tail(job.stderr_log_file) - return BackgroundJobStatus( - job_id=job.job_id, - completed=True, - exit_code=exit_code, - stdout=stdout, - stderr=stderr, - stdout_truncated=stdout_truncated, - stderr_truncated=stderr_truncated, + return self._get_completed_background_job_output( + sandbox_id, + job, + exit_code, + timeout, ) def _request_background_job_status_batch( @@ -1612,29 +1628,55 @@ def _get_completed_background_job_output( exit_code: int, timeout: Optional[int], ) -> BackgroundJobStatus: - """Read bounded stdout and stderr tails for one completed job.""" + """Read sequential output tails without invalidating known completion.""" + + output_timeout = ( + float(timeout) if timeout is not None else BACKGROUND_JOB_OUTPUT_FETCH_TIMEOUT_SECONDS + ) + deadline = time.monotonic() + max(0.0, output_timeout) - def read_output_tail(path: str) -> tuple[str, bool]: + def deadline_error() -> str: + return f"Output retrieval deadline exceeded after {output_timeout:g}s" + + def read_output_tail(path: str) -> tuple[Optional[str], bool, Optional[str]]: + remaining = deadline - time.monotonic() + if remaining <= 0: + return None, False, deadline_error() + if not self._background_job_output_semaphore.acquire(timeout=remaining): + return None, False, deadline_error() try: + remaining = deadline - time.monotonic() + if remaining <= 0: + return None, False, deadline_error() + request_timeout = max( + 1, + min(timeout if timeout is not None else 30, math.ceil(remaining)), + ) response = self.read_file( sandbox_id, path, - timeout=timeout, + timeout=request_timeout, offset=-JOB_OUTPUT_TAIL_BYTES, length=JOB_OUTPUT_TAIL_BYTES, ) - return response.content, bool(response.truncated) + return response.content, bool(response.truncated), None except SandboxFileNotFoundError: - return "", False + return "", False, None + except APIError as exc: + return None, False, _format_exception_diagnostic(exc) + finally: + self._background_job_output_semaphore.release() - stdout, stdout_truncated = read_output_tail(job.stdout_log_file) - stderr, stderr_truncated = read_output_tail(job.stderr_log_file) + stdout, stdout_truncated, stdout_error = read_output_tail(job.stdout_log_file) + stderr, stderr_truncated, stderr_error = read_output_tail(job.stderr_log_file) return BackgroundJobStatus( job_id=job.job_id, completed=True, exit_code=exit_code, stdout=stdout, stderr=stderr, + stdout_error=stdout_error, + stderr_error=stderr_error, stdout_truncated=stdout_truncated, stderr_truncated=stderr_truncated, ) @@ -2182,7 +2224,7 @@ def read_file( method = getattr(req, "method", "?") u = getattr(req, "url", "?") raise APIError( - f"Read file failed: {e.__class__.__name__} at {method} {u}: {e}" + f"Read file failed at {method} {u}: {_format_exception_diagnostic(e)}" ) from e except Exception as e: raise APIError(f"Read file failed: {e.__class__.__name__}: {e}") from e @@ -2276,6 +2318,9 @@ def __init__( self._background_job_status_batcher = _AsyncRequestBatcher( self._fetch_background_job_statuses ) + self._background_job_output_semaphore = asyncio.Semaphore( + MAX_CONCURRENT_BACKGROUND_JOB_OUTPUT_READS + ) self._sandbox_status_batch_supported: Optional[bool] = None self._background_job_status_batch_supported: Optional[bool] = None @@ -2999,9 +3044,8 @@ async def get_background_job( Args: sandbox_id: The sandbox ID job: The BackgroundJob handle from start_background_job() - timeout: Optional per-call timeout (in seconds) forwarded to the - underlying read_file calls. When None, the APIClient default - applies. + timeout: Optional output-retrieval deadline in seconds after the + exit file is observed. When None, a bounded SDK default applies. Returns: BackgroundJobStatus with completed flag, and exit_code/stdout if @@ -3015,20 +3059,6 @@ async def read_or_empty(path: str) -> str: except SandboxFileNotFoundError: return "" - async def read_output_tail(path: str) -> "tuple[str, bool]": - try: - response = await self.read_file( - sandbox_id, - path, - timeout=timeout, - offset=-JOB_OUTPUT_TAIL_BYTES, - length=JOB_OUTPUT_TAIL_BYTES, - ) - # Servers without windowed-read support omit `truncated`. - return response.content, bool(response.truncated) - except SandboxFileNotFoundError: - return "", False - exit_content = await read_or_empty(job.exit_file) if not exit_content.strip(): return BackgroundJobStatus(job_id=job.job_id, completed=False) @@ -3038,16 +3068,11 @@ async def read_output_tail(path: str) -> "tuple[str, bool]": except ValueError: return BackgroundJobStatus(job_id=job.job_id, completed=False) - stdout, stdout_truncated = await read_output_tail(job.stdout_log_file) - stderr, stderr_truncated = await read_output_tail(job.stderr_log_file) - return BackgroundJobStatus( - job_id=job.job_id, - completed=True, - exit_code=exit_code, - stdout=stdout, - stderr=stderr, - stdout_truncated=stdout_truncated, - stderr_truncated=stderr_truncated, + return await self._get_completed_background_job_output( + sandbox_id, + job, + exit_code, + timeout, ) async def _request_background_job_status_batch( @@ -3137,33 +3162,66 @@ async def _get_completed_background_job_output( exit_code: int, timeout: Optional[int], ) -> BackgroundJobStatus: - """Read bounded stdout and stderr tails for one completed job.""" + """Read sequential output tails without invalidating known completion.""" + + output_timeout = ( + float(timeout) if timeout is not None else BACKGROUND_JOB_OUTPUT_FETCH_TIMEOUT_SECONDS + ) + loop = asyncio.get_running_loop() + deadline = loop.time() + max(0.0, output_timeout) + + def deadline_error() -> str: + return f"Output retrieval deadline exceeded after {output_timeout:g}s" + + async def read_output_tail( + path: str, + ) -> tuple[Optional[str], bool, Optional[str]]: + remaining = deadline - loop.time() + if remaining <= 0: + return None, False, deadline_error() + + async def fetch() -> ReadFileResponse: + async with self._background_job_output_semaphore: + remaining_after_acquire = deadline - loop.time() + if remaining_after_acquire <= 0: + raise asyncio.TimeoutError + request_timeout = max( + 1, + min( + timeout if timeout is not None else 30, + math.ceil(remaining_after_acquire), + ), + ) + return await self.read_file( + sandbox_id, + path, + timeout=request_timeout, + offset=-JOB_OUTPUT_TAIL_BYTES, + length=JOB_OUTPUT_TAIL_BYTES, + ) - async def read_output_tail(path: str) -> tuple[str, bool]: try: - response = await self.read_file( - sandbox_id, - path, - timeout=timeout, - offset=-JOB_OUTPUT_TAIL_BYTES, - length=JOB_OUTPUT_TAIL_BYTES, - ) - return response.content, bool(response.truncated) + response = await asyncio.wait_for(fetch(), timeout=remaining) + return response.content, bool(response.truncated), None except SandboxFileNotFoundError: - return "", False - - stdout, stderr = await asyncio.gather( - read_output_tail(job.stdout_log_file), - read_output_tail(job.stderr_log_file), - ) + return "", False, None + except asyncio.TimeoutError: + return None, False, deadline_error() + except APIError as exc: + return None, False, _format_exception_diagnostic(exc) + + stdout, stdout_truncated, stdout_error = await read_output_tail(job.stdout_log_file) + stderr, stderr_truncated, stderr_error = await read_output_tail(job.stderr_log_file) return BackgroundJobStatus( job_id=job.job_id, completed=True, exit_code=exit_code, - stdout=stdout[0], - stderr=stderr[0], - stdout_truncated=stdout[1], - stderr_truncated=stderr[1], + stdout=stdout, + stderr=stderr, + stdout_error=stdout_error, + stderr_error=stderr_error, + stdout_truncated=stdout_truncated, + stderr_truncated=stderr_truncated, ) async def _fetch_background_job_statuses( @@ -3737,7 +3795,7 @@ async def read_file( method = getattr(req, "method", "?") u = getattr(req, "url", "?") raise APIError( - f"Read file failed: {e.__class__.__name__} at {method} {u}: {e}" + f"Read file failed at {method} {u}: {_format_exception_diagnostic(e)}" ) from e except Exception as e: raise APIError(f"Read file failed: {e.__class__.__name__}: {e}") from e diff --git a/packages/prime-sandboxes/tests/test_background_job_timeout.py b/packages/prime-sandboxes/tests/test_background_job_timeout.py index cebd4d8d..65ac12ea 100644 --- a/packages/prime-sandboxes/tests/test_background_job_timeout.py +++ b/packages/prime-sandboxes/tests/test_background_job_timeout.py @@ -4,7 +4,7 @@ import pytest -from prime_sandboxes.core.client import APIClient +from prime_sandboxes.core.client import APIClient, APIError from prime_sandboxes.models import BackgroundJob, ReadFileResponse from prime_sandboxes.sandbox import AsyncSandboxClient, SandboxClient @@ -141,6 +141,35 @@ def fake_read_file( assert status.stderr_truncated is False +def test_sync_output_error_preserves_completed_exit_code(): + client = SandboxClient(APIClient(api_key="test-key")) + client_any = cast(Any, client) + + def fake_read_file( + sandbox_id: str, + file_path: str, + timeout: Optional[int] = None, + offset: Optional[int] = None, + length: Optional[int] = None, + ) -> ReadFileResponse: + if file_path.endswith(".exit"): + return _whole_file("7\n") + if file_path.endswith(".stdout"): + raise APIError("Read file failed: ConnectError") + return _whole_file("stderr") + + client_any.read_file = fake_read_file + + status = client.get_background_job("sbx-123", _make_job()) + + assert status.completed + assert status.exit_code == 7 + assert status.stdout is None + assert status.stdout_error == "APIError: Read file failed: ConnectError" + assert status.stderr == "stderr" + assert status.stderr_error is None + + @pytest.mark.asyncio async def test_async_get_background_job_handles_legacy_read_file_response(): client = AsyncSandboxClient(api_key="test-key") diff --git a/packages/prime-sandboxes/tests/test_batch_status.py b/packages/prime-sandboxes/tests/test_batch_status.py index 6c35a782..d5a7a21b 100644 --- a/packages/prime-sandboxes/tests/test_batch_status.py +++ b/packages/prime-sandboxes/tests/test_batch_status.py @@ -1,6 +1,7 @@ """Focused tests for platform lifecycle and VM background-job batch calls.""" import asyncio +import errno from concurrent.futures import ThreadPoolExecutor from types import SimpleNamespace from typing import Any, Optional, cast @@ -145,15 +146,24 @@ def request(self, method: str, path: str, **kwargs: Any) -> dict[str, Any]: class _AsyncBackgroundJobPlatformClient: - def __init__(self, error_job_id: Optional[str] = None) -> None: + def __init__( + self, + error_job_id: Optional[str] = None, + complete_all: bool = False, + ) -> None: self.calls: list[tuple[str, str, dict[str, Any]]] = [] self.error_job_id = error_job_id + self.complete_all = complete_all async def request(self, method: str, path: str, **kwargs: Any) -> dict[str, Any]: self.calls.append((method, path, kwargs)) return { "statuses": [ - {**job, "completed": False, "exit_code": None} + { + **job, + "completed": self.complete_all, + "exit_code": 0 if self.complete_all else None, + } for job in kwargs["json"]["jobs"] if job["job_id"] != self.error_job_id ], @@ -486,6 +496,118 @@ async def test_async_get_background_jobs_uses_one_platform_batch() -> None: assert not statuses[1].completed +@pytest.mark.asyncio +async def test_async_completed_output_reads_are_client_bounded_and_sequential() -> None: + client = AsyncSandboxClient(api_key="test-key") + await client.client.aclose() + platform = _AsyncBackgroundJobPlatformClient(complete_all=True) + cast(Any, client).client = platform + + active_reads = 0 + peak_reads = 0 + paths_by_sandbox: dict[str, list[str]] = {} + + async def read_file( + sandbox_id: str, + path: str, + timeout: Optional[int] = None, + offset: Optional[int] = None, + length: Optional[int] = None, + ) -> ReadFileResponse: + nonlocal active_reads, peak_reads + paths_by_sandbox.setdefault(sandbox_id, []).append(path) + active_reads += 1 + peak_reads = max(peak_reads, active_reads) + try: + await asyncio.sleep(0.01) + return ReadFileResponse(content=path, size=len(path), truncated=False) + finally: + active_reads -= 1 + + cast(Any, client).read_file = read_file + jobs = [_job(f"sandbox-{index}", f"{index:08x}") for index in range(100)] + try: + statuses = await client.get_background_jobs(jobs) + finally: + await client.aclose() + + assert all(status.completed and status.exit_code == 0 for status in statuses) + assert peak_reads == 20 + assert sum(len(paths) for paths in paths_by_sandbox.values()) == 200 + for job in jobs: + assert paths_by_sandbox[job.sandbox_id] == [ + job.stdout_log_file, + job.stderr_log_file, + ] + + +@pytest.mark.asyncio +async def test_async_output_error_preserves_completion_and_surfaces_errno() -> None: + client = AsyncSandboxClient(api_key="test-key") + await client.client.aclose() + cast(Any, client).client = _AsyncBackgroundJobPlatformClient(complete_all=True) + + async def read_file( + _sandbox_id: str, + path: str, + timeout: Optional[int] = None, + offset: Optional[int] = None, + length: Optional[int] = None, + ) -> ReadFileResponse: + if path.endswith("stdout.log"): + try: + try: + raise OSError(errno.EMFILE, "Too many open files") + except OSError as os_error: + raise RuntimeError("All connection attempts failed") from os_error + except RuntimeError as transport_error: + raise APIError("Read file failed: ConnectError") from transport_error + return ReadFileResponse(content="stderr", size=6, truncated=False) + + cast(Any, client).read_file = read_file + try: + status = (await client.get_background_jobs([_job("sandbox-a", "feedface")]))[0] + finally: + await client.aclose() + + assert status.completed + assert status.exit_code == 0 + assert status.stdout is None + assert status.stdout_error is not None + assert "errno=24" in status.stdout_error + assert "Too many open files" in status.stdout_error + assert status.stderr == "stderr" + assert status.stderr_error is None + + +@pytest.mark.asyncio +async def test_async_output_deadline_preserves_completion() -> None: + client = AsyncSandboxClient(api_key="test-key") + await client.client.aclose() + cast(Any, client).client = _AsyncBackgroundJobPlatformClient(complete_all=True) + + async def unexpected_read(*_args: Any, **_kwargs: Any) -> ReadFileResponse: + raise AssertionError("expired output deadline must not start a read") + + cast(Any, client).read_file = unexpected_read + try: + status = ( + await client.get_background_jobs( + [_job("sandbox-a", "feedface")], + timeout=0, + ) + )[0] + finally: + await client.aclose() + + assert status.completed + assert status.exit_code == 0 + assert status.stdout is None + assert status.stderr is None + assert status.stdout_error == "Output retrieval deadline exceeded after 0s" + assert status.stderr_error == "Output retrieval deadline exceeded after 0s" + + @pytest.mark.asyncio async def test_concurrent_async_background_waiters_share_one_platform_batch() -> None: client = AsyncSandboxClient(api_key="test-key") diff --git a/packages/prime-sandboxes/tests/test_client_retry.py b/packages/prime-sandboxes/tests/test_client_retry.py index d196363f..1ba362f1 100644 --- a/packages/prime-sandboxes/tests/test_client_retry.py +++ b/packages/prime-sandboxes/tests/test_client_retry.py @@ -886,3 +886,37 @@ def raise_timeout(url, headers, params, timeout): assert "Read file timed out after 12s" in message assert f"({expected_name})" in message assert "/tmp/job.exit" in message + + def test_read_file_connection_error_surfaces_nested_os_errno(self, monkeypatch): + """Transport wrappers retain the actionable nested socket errno.""" + import errno + + from prime_sandboxes.sandbox import SandboxClient + + request = httpx.Request("GET", "https://gateway.example/read-file") + + def raise_connect_error(url, headers, params, timeout): + try: + raise OSError(errno.EMFILE, "Too many open files") + except OSError as os_error: + raise httpx.ConnectError( + "All connection attempts failed", + request=request, + ) from os_error + + monkeypatch.setattr( + SandboxClient, + "_gateway_read_file_get", + staticmethod(raise_connect_error), + ) + + client = SandboxClient.__new__(SandboxClient) + client._auth_cache = DummySandboxAuthCache() + + with pytest.raises(APIError) as exc_info: + client.read_file("sandbox-123", "/tmp/job.stdout.log") + + message = str(exc_info.value) + assert "ConnectError" in message + assert "errno=24" in message + assert "Too many open files" in message