From f3da6e84fef41d4d25e9d6c507ecf1d075c7f055 Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Mon, 14 Sep 2026 10:57:16 -0700 Subject: [PATCH 1/6] feat: add sync and async sandbox SDK Implements CON-1523, CON-1524, CON-1525, and CON-1526. Based on REST migration PR #584. Keep this draft local until that PR is merged. Includes managed ownership, bounded startup handling, closeable SSE streams, and cancellation-safe synchronous and asynchronous resource cleanup. --- runpod/__init__.py | 16 +- runpod/api/rest.py | 38 ++- runpod/api/sandboxes.py | 386 +++++++++++++++++++++++++ runpod/sandbox/__init__.py | 31 ++ runpod/sandbox/asyncio.py | 565 +++++++++++++++++++++++++++++++++++++ runpod/sandbox/models.py | 136 +++++++++ runpod/sandbox/sync.py | 438 ++++++++++++++++++++++++++++ tests/test_sandbox.py | 458 ++++++++++++++++++++++++++++++ 8 files changed, 2052 insertions(+), 16 deletions(-) create mode 100644 runpod/api/sandboxes.py create mode 100644 runpod/sandbox/__init__.py create mode 100644 runpod/sandbox/asyncio.py create mode 100644 runpod/sandbox/models.py create mode 100644 runpod/sandbox/sync.py create mode 100644 tests/test_sandbox.py diff --git a/runpod/__init__.py b/runpod/__init__.py index 6d24180a..812fc866 100644 --- a/runpod/__init__.py +++ b/runpod/__init__.py @@ -1,4 +1,4 @@ -""" Allows runpod to be imported as a module. """ +"""Allows runpod to be imported as a module.""" import logging import os @@ -29,19 +29,20 @@ set_credentials, ) from .endpoint import AsyncioEndpoint, AsyncioJob, Endpoint +from .sandbox import AsyncioSandbox, Sandbox from .serverless.modules.rp_logger import RunPodLogger from .version import __version__ __all__ = [ # API functions "create_container_registry_auth", - "create_endpoint", + "create_endpoint", "create_pod", "create_template", "delete_container_registry_auth", "get_endpoints", "get_gpu", - "get_gpus", + "get_gpus", "get_pod", "get_pods", "get_user", @@ -53,12 +54,15 @@ "update_user_settings", # Config functions "check_credentials", - "get_credentials", + "get_credentials", "set_credentials", # Endpoint classes "AsyncioEndpoint", "AsyncioJob", "Endpoint", + # Sandbox classes + "AsyncioSandbox", + "Sandbox", # Serverless module "serverless", # Logger class @@ -68,8 +72,8 @@ # Module variables "SSH_KEY_PATH", "profile", - "api_key", - "endpoint_url_base" + "api_key", + "endpoint_url_base", ] # ------------------------------- Config Paths ------------------------------- # diff --git a/runpod/api/rest.py b/runpod/api/rest.py index 987eb143..fced2429 100644 --- a/runpod/api/rest.py +++ b/runpod/api/rest.py @@ -15,7 +15,9 @@ def _resolve_api_key(api_key: Optional[str]) -> str: - from runpod import api_key as global_api_key # pylint: disable=import-outside-toplevel,cyclic-import + from runpod import ( + api_key as global_api_key, + ) # pylint: disable=import-outside-toplevel,cyclic-import effective_api_key = api_key or global_api_key if not effective_api_key: @@ -23,8 +25,10 @@ def _resolve_api_key(api_key: Optional[str]) -> str: return effective_api_key -def _build_url(path: str) -> str: - api_url_base = os.environ.get("RUNPOD_API_BASE_URL", "https://api.runpod.io") +def _build_url(path: str, base_url: Optional[str] = None) -> str: + api_url_base = base_url or os.environ.get( + "RUNPOD_API_BASE_URL", "https://api.runpod.io" + ) return f"{api_url_base.rstrip('/')}/{path.lstrip('/')}" @@ -45,30 +49,44 @@ def _response_json(response: requests.Response) -> dict[str, Any]: return payload if isinstance(payload, dict) else {} -def _raise_for_error( - response: requests.Response, method: str, path: str +def _raise_for_status( + status_code: int, + payload: Mapping[str, Any], + text: str, + method: str, + path: str, ) -> None: - if response.status_code == HTTP_STATUS_UNAUTHORIZED: + """Map HTTP error details consistently across REST transports.""" + if status_code == HTTP_STATUS_UNAUTHORIZED: raise error.AuthenticationError( "Unauthorized request, please check your API key." ) - if response.status_code < HTTP_STATUS_BAD_REQUEST: + if status_code < HTTP_STATUS_BAD_REQUEST: return - payload = _response_json(response) message = payload.get("detail") or payload.get("title") if not message: - message = response.text or f"Request failed with status {response.status_code}" + message = text or f"Request failed with status {status_code}" raise error.QueryError( str(message), f"{method.upper()} {path}", - status_code=response.status_code, + status_code=status_code, errors=payload.get("errors"), ) +def _raise_for_error(response: requests.Response, method: str, path: str) -> None: + if response.status_code < HTTP_STATUS_BAD_REQUEST: + return + if response.status_code == HTTP_STATUS_UNAUTHORIZED: + _raise_for_status(response.status_code, {}, "", method, path) + _raise_for_status( + response.status_code, _response_json(response), response.text, method, path + ) + + def run_rest_request( method: str, path: str, diff --git a/runpod/api/sandboxes.py b/runpod/api/sandboxes.py new file mode 100644 index 00000000..74dac1fb --- /dev/null +++ b/runpod/api/sandboxes.py @@ -0,0 +1,386 @@ +"""Asynchronous REST and streaming transport for CPU sandboxes.""" + +from __future__ import annotations + +import asyncio +import codecs +import json +import math +import re +from collections.abc import AsyncIterator, Mapping, Sequence +from typing import Any +from urllib.parse import quote, urlencode + +import aiohttp +from yarl import URL + +from runpod import error +from runpod.api.rest import ( + _build_headers, + _build_url, + _raise_for_status, + _resolve_api_key, +) + +_SANDBOX_PATH = "/v2/sandboxes" +_LINE_END = re.compile(r"\r\n|\r|\n") + + +def _sandbox_path(sandbox_id: str) -> str: + return f"{_SANDBOX_PATH}/{quote(sandbox_id, safe='')}" + + +async def _raise_response_error( + response: aiohttp.ClientResponse, method: str, path: str +) -> None: + if response.status == 401: + _raise_for_status(response.status, {}, "", method, path) + if 200 <= response.status < 300: + return + raw = await response.read() + try: + payload = json.loads(raw) + except (ValueError, UnicodeError): + payload = {} + if not isinstance(payload, dict): + payload = {} + try: + text = raw.decode(response.charset or "utf-8", errors="replace") + except LookupError: + text = raw.decode("utf-8", errors="replace") + _raise_for_status(response.status, payload, text, method, path) + # Redirects are deliberately not followed, including on read operations. + raise error.QueryError( + text or f"Unexpected HTTP status {response.status}", + f"{method} {path}", + status_code=response.status, + errors=payload.get("errors"), + ) + + +async def _sse_lines(content: aiohttp.StreamReader) -> AsyncIterator[str]: + """Decode only the current line, allowing arbitrary UTF-8/chunk boundaries.""" + decoder = codecs.getincrementaldecoder("utf-8-sig")() + pending = "" + skip_lf = False + async for chunk in content.iter_chunked(8192): + text = decoder.decode(chunk) + if not text: + continue + if skip_lf: + if text.startswith("\n"): + text = text[1:] + skip_lf = False + text = pending + text + start = 0 + for match in _LINE_END.finditer(text): + yield text[start : match.start()] + start = match.end() + skip_lf = text.endswith("\r") + pending = text[start:] + pending += decoder.decode(b"", final=True) + if pending: + yield pending + + +class AsyncSandboxLogStream: + """One SSE connection; use an async context or ``aclose`` on early exit. + + ``open`` performs the HTTP handshake without waiting for a log event. + Exhaustion, cancellation, parsing errors and server timeout frames close + the response. No reconnect is attempted, even when a cursor is supplied. + """ + + def __init__( + self, + api: AsyncSandboxAPI, + path: str, + params: list[tuple[str, str]], + last_event_id: str | None, + ) -> None: + self._api = api + self._path = path + self._params = params + self._last_event_id = last_event_id + self._response: aiohttp.ClientResponse | None = None + self._open_task: asyncio.Task[None] | None = None + self._iterator: AsyncIterator[dict[str, Any]] | None = None + self._closed = False + + async def _perform_open(self) -> None: + try: + if self._closed: + raise RuntimeError("Sandbox log stream is closed") + headers = self._api._headers() + headers["Accept"] = "text/event-stream" + if self._last_event_id is not None: + headers["Last-Event-ID"] = self._last_event_id + self._response = await self._api._get_session().request( + "GET", + self._api._url(self._path, self._params), + headers=headers, + timeout=aiohttp.ClientTimeout( + total=None, + connect=self._api.request_timeout, + sock_connect=self._api.request_timeout, + sock_read=None, + ), + allow_redirects=False, + ) + await _raise_response_error(self._response, "GET", self._path) + if self._response.content_type != "text/event-stream": + raise error.QueryError( + "Expected a text/event-stream sandbox log response", + f"GET {self._path}", + status_code=self._response.status, + ) + self._iterator = self._events(self._response) + except BaseException: + self._dispose() + raise + + async def open(self) -> AsyncSandboxLogStream: + """Open once, with a bounded handshake but no total stream timeout.""" + if self._closed: + raise RuntimeError("Sandbox log stream is closed") + if self._open_task is None: + self._open_task = asyncio.create_task( + asyncio.wait_for( + self._perform_open(), timeout=self._api.request_timeout + ) + ) + try: + await self._open_task + except BaseException: + await self.aclose() + raise + return self + + def _dispose(self) -> None: + self._closed = True + if self._response is not None: + self._response.close() + self._response = None + self._iterator = None + self._api._streams.discard(self) + + async def aclose(self) -> None: + """Release this stream, including a handshake still in progress.""" + self._dispose() + if self._open_task is not None and not self._open_task.done(): + self._open_task.cancel() + await asyncio.gather(self._open_task, return_exceptions=True) + + async def __aenter__(self) -> AsyncSandboxLogStream: + return await self.open() + + async def __aexit__(self, exc_type, exc_value, traceback) -> None: + await self.aclose() + + def __aiter__(self) -> AsyncSandboxLogStream: + return self + + async def __anext__(self) -> dict[str, Any]: + if self._closed: + raise StopAsyncIteration + try: + await self.open() + assert self._iterator is not None + return await anext(self._iterator) + except BaseException: + await self.aclose() + raise + + async def _events( + self, response: aiohttp.ClientResponse + ) -> AsyncIterator[dict[str, Any]]: + data: list[str] = [] + event_type = "" + event_id: str | None = None + async for line in _sse_lines(response.content): + if not line: + if event_type == "timeout": + return + if event_type not in ("", "message"): + raise error.QueryError( + f"Sandbox log event {event_type!r}: " + "\n".join(data), + f"GET {self._path}", + ) + if data: + payload = json.loads("\n".join(data)) + if not isinstance(payload, dict): + raise ValueError("Sandbox log event must be a JSON object") + yield { + "source": payload["source"], + "line": payload["line"], + "ts": payload["ts"], + "id": event_id, + } + data = [] + event_type = "" + continue + if line.startswith(":"): + continue + field, separator, value = line.partition(":") + if separator and value.startswith(" "): + value = value[1:] + if field == "data": + data.append(value) + elif field == "event": + event_type = value + elif field == "id" and "\x00" not in value: + event_id = value + # SSE retry directives are intentionally ignored: no reconnects. + # Per SSE, EOF does not dispatch an event lacking its blank-line boundary. + + +class AsyncSandboxAPI: + """Lazy, owned aiohttp transport. Each operation makes one HTTP request. + + ``close`` releases streams and sockets, but a later request can open a new + session. Use the instance from one event loop while its session is active. + Readiness and retry policies belong to the higher-level sandbox handles. + """ + + def __init__( + self, + api_key: str | None = None, + base_url: str | None = None, + request_timeout: float = 30, + ) -> None: + if not math.isfinite(request_timeout) or request_timeout <= 0: + raise ValueError("request_timeout must be a finite positive number") + self.api_key = api_key + self.base_url = base_url + self.request_timeout = request_timeout + self._session: aiohttp.ClientSession | None = None + self._streams: set[AsyncSandboxLogStream] = set() + + def _headers(self) -> dict[str, str]: + return _build_headers(_resolve_api_key(self.api_key)) + + def _url(self, path: str, params: list[tuple[str, str]] | None = None) -> URL: + url = _build_url(path, self.base_url) + if params: + url += "?" + urlencode(params, safe="=") + # Preserve quoted ID segments (including encoded slashes) verbatim. + return URL(url, encoded=True) + + def _get_session(self) -> aiohttp.ClientSession: + if self._session is None or self._session.closed: + self._session = aiohttp.ClientSession( + timeout=aiohttp.ClientTimeout(total=self.request_timeout), + cookie_jar=aiohttp.DummyCookieJar(), + ) + # aiohttp otherwise transparently replays disconnected GET/DELETE + # requests. It currently exposes no public switch for this policy. + self._session._retry_connection = False + return self._session + + async def _request( + self, + method: str, + path: str, + *, + params: list[tuple[str, str]] | None = None, + body: Mapping[str, Any] | None = None, + ) -> dict[str, Any] | None: + headers = self._headers() + async with self._get_session().request( + method, + self._url(path, params), + headers=headers, + json=body, + allow_redirects=False, + ) as response: + await _raise_response_error(response, method, path) + if response.status == 204: + return None + payload = await response.json(content_type=None) + if not isinstance(payload, dict): + raise ValueError("Sandbox API response must be a JSON object") + return payload + + async def create(self, body: Mapping[str, Any]) -> dict[str, Any]: + payload = await self._request( + "POST", + _SANDBOX_PATH, + body={key: value for key, value in body.items() if value is not None}, + ) + if payload is None: + raise ValueError("Sandbox creation returned no snapshot") + return payload + + async def list( + self, + *, + state: str | None = None, + labels: Mapping[str, str] | None = None, + ) -> list[dict[str, Any]]: + params = [] + if state is not None: + params.append(("state", state)) + if labels is not None: + params.extend(("labels", f"{key}={value}") for key, value in labels.items()) + payload = await self._request("GET", _SANDBOX_PATH, params=params) + if payload is None: + raise ValueError("Sandbox listing returned no response") + sandboxes = payload["sandboxes"] + if not isinstance(sandboxes, list) or any( + not isinstance(item, dict) for item in sandboxes + ): + raise ValueError("Sandbox listing must contain an array of objects") + return sandboxes + + async def get(self, sandbox_id: str) -> dict[str, Any]: + payload = await self._request("GET", _sandbox_path(sandbox_id)) + if payload is None: + raise ValueError("Sandbox lookup returned no snapshot") + return payload + + async def terminate(self, sandbox_id: str) -> None: + await self._request("DELETE", _sandbox_path(sandbox_id)) + + async def exec(self, sandbox_id: str, command: Sequence[str]) -> dict[str, Any]: + payload = await self._request( + "POST", f"{_sandbox_path(sandbox_id)}/exec", body={"command": list(command)} + ) + if payload is None: + raise ValueError("Sandbox execution returned no result") + return payload + + def logs( + self, + sandbox_id: str, + *, + source: str | None = None, + tail: int | None = None, + since: str | None = None, + last_event_id: str | None = None, + ) -> AsyncSandboxLogStream: + params = [ + (key, str(value)) + for key, value in (("source", source), ("tail", tail), ("since", since)) + if value is not None + ] + stream = AsyncSandboxLogStream( + self, f"{_sandbox_path(sandbox_id)}/logs", params, last_event_id + ) + self._streams.add(stream) + return stream + + async def close(self) -> None: + """Close all active streams and the owned session; allow later reuse.""" + streams = tuple(self._streams) + session, self._session = self._session, None + # Dispose every response before the first await, including on cancellation. + for stream in streams: + stream._dispose() + if stream._open_task is not None and not stream._open_task.done(): + stream._open_task.cancel() + try: + for stream in streams: + await stream.aclose() + finally: + if session is not None: + await session.close() diff --git a/runpod/sandbox/__init__.py b/runpod/sandbox/__init__.py new file mode 100644 index 00000000..9bcb930d --- /dev/null +++ b/runpod/sandbox/__init__.py @@ -0,0 +1,31 @@ +"""Sync and async sandbox lifecycles, execution results, and log streams.""" + +from .asyncio import AsyncioSandbox, AsyncSandboxLogStream +from .models import ( + ExecResult, + LogEvent, + LogSource, + SandboxCompute, + SandboxExecutionError, + SandboxInfo, + SandboxStartupTimeout, + SandboxState, + SandboxStateError, +) +from .sync import Sandbox, SandboxLogs + +__all__ = [ + "AsyncioSandbox", + "AsyncSandboxLogStream", + "Sandbox", + "SandboxLogs", + "ExecResult", + "LogEvent", + "LogSource", + "SandboxCompute", + "SandboxExecutionError", + "SandboxInfo", + "SandboxStartupTimeout", + "SandboxState", + "SandboxStateError", +] diff --git a/runpod/sandbox/asyncio.py b/runpod/sandbox/asyncio.py new file mode 100644 index 00000000..849c9494 --- /dev/null +++ b/runpod/sandbox/asyncio.py @@ -0,0 +1,565 @@ +"""Asynchronous sandbox handles, lifecycle management, and typed log streams.""" + +import asyncio +import math +from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import replace +from datetime import datetime +from types import TracebackType +from typing import Any, Optional, TypeVar + +from runpod.api.sandboxes import AsyncSandboxAPI +from runpod.error import QueryError +from runpod.sandbox.models import ( + ExecResult, + LogEvent, + LogSource, + SandboxCompute, + SandboxExecutionError, + SandboxInfo, + SandboxStartupTimeout, + SandboxState, + SandboxStateError, +) + +_T = TypeVar("_T") + + +def _timeout(value: float, name: str, *, allow_zero: bool = False) -> float: + if not math.isfinite(value) or value < 0 or (value == 0 and not allow_zero): + raise ValueError( + f"{name} must be finite and {'nonnegative' if allow_zero else 'positive'}" + ) + return value + + +async def _cleanup( + operation: Awaitable[None], + timeout: float, + original: Optional[BaseException] = None, +) -> None: + """Finish bounded cleanup despite cancellation, keeping the original primary.""" + task = asyncio.create_task(asyncio.wait_for(operation, timeout)) + interrupted = original + while not task.done(): + try: + await asyncio.shield(task) + except asyncio.CancelledError as error: + if interrupted is None: + interrupted = error + except BaseException: + break + try: + task.result() + except BaseException as error: + if interrupted is not None: + raise interrupted from error + raise + if interrupted is not None: + raise interrupted + + +class AsyncioSandbox: + """A cached sandbox handle; construction never performs network I/O. + + Supply exactly one of ``image_name`` or ``template_id``. Unspecified + resource and lifetime options are left to the server's account policy. + Created handles own their remote sandbox in a context; ``get`` and ``list`` + return borrowed handles whose context only closes local resources. + """ + + def __init__( + self, + *, + image_name: Optional[str] = None, + template_id: Optional[str] = None, + name: Optional[str] = None, + cpu_flavor_id: Optional[str] = None, + vcpu_count: Optional[int] = None, + memory_in_gb: Optional[int] = None, + data_center_id: Optional[str] = None, + env: Optional[Mapping[str, str]] = None, + idle_timeout_seconds: Optional[int] = None, + max_lifetime_seconds: Optional[int] = None, + labels: Optional[Mapping[str, str]] = None, + api_key: Optional[str] = None, + base_url: Optional[str] = None, + request_timeout: float = 30, + startup_timeout: float = 60, + ) -> None: + if (image_name is None) == (template_id is None): + raise ValueError("Supply exactly one of image_name or template_id") + options = { + "imageName": image_name, + "templateId": template_id, + "name": name, + "cpuFlavorId": cpu_flavor_id, + "vcpuCount": vcpu_count, + "memoryInGb": memory_in_gb, + "dataCenterId": data_center_id, + "env": dict(env) if env is not None else None, + "idleTimeoutSeconds": idle_timeout_seconds, + "maxLifetimeSeconds": max_lifetime_seconds, + "labels": dict(labels) if labels is not None else None, + } + self._initialize(api_key, base_url, request_timeout, startup_timeout) + self._create_body = { + key: value for key, value in options.items() if value is not None + } + + def _initialize( + self, + api_key: Optional[str], + base_url: Optional[str], + request_timeout: float, + startup_timeout: float, + ) -> None: + self._request_timeout = _timeout(request_timeout, "request_timeout") + self._startup_timeout = _timeout( + startup_timeout, "startup_timeout", allow_zero=True + ) + self._api = AsyncSandboxAPI(api_key, base_url, request_timeout) + self._info: Optional[SandboxInfo] = None + self._sandbox_id: Optional[str] = None + self._create_body: Optional[dict[str, Any]] = None + self._owned = False + self._entering = False + self._creating = False + self._streams: set[AsyncSandboxLogStream] = set() + + @classmethod + async def create( + cls, + *, + image_name: Optional[str] = None, + template_id: Optional[str] = None, + name: Optional[str] = None, + cpu_flavor_id: Optional[str] = None, + vcpu_count: Optional[int] = None, + memory_in_gb: Optional[int] = None, + data_center_id: Optional[str] = None, + env: Optional[Mapping[str, str]] = None, + idle_timeout_seconds: Optional[int] = None, + max_lifetime_seconds: Optional[int] = None, + labels: Optional[Mapping[str, str]] = None, + api_key: Optional[str] = None, + base_url: Optional[str] = None, + request_timeout: float = 30, + startup_timeout: float = 60, + ) -> "AsyncioSandbox": + """Create and return an owned handle, without claiming container readiness.""" + sandbox = cls( + image_name=image_name, + template_id=template_id, + name=name, + cpu_flavor_id=cpu_flavor_id, + vcpu_count=vcpu_count, + memory_in_gb=memory_in_gb, + data_center_id=data_center_id, + env=env, + idle_timeout_seconds=idle_timeout_seconds, + max_lifetime_seconds=max_lifetime_seconds, + labels=labels, + api_key=api_key, + base_url=base_url, + request_timeout=request_timeout, + startup_timeout=startup_timeout, + ) + await sandbox._create() + return sandbox + + @classmethod + async def get( + cls, + sandbox_id: str, + *, + api_key: Optional[str] = None, + base_url: Optional[str] = None, + request_timeout: float = 30, + startup_timeout: float = 60, + ) -> "AsyncioSandbox": + """Fetch a borrowed handle; leaving its context does not terminate it.""" + sandbox = cls.__new__(cls) + sandbox._initialize(api_key, base_url, request_timeout, startup_timeout) + sandbox._sandbox_id = sandbox_id + try: + await sandbox.refresh() + except BaseException as error: + await _cleanup(sandbox.close(), request_timeout, error) + return sandbox + + @classmethod + async def list( + cls, + *, + state: Optional[SandboxState] = None, + labels: Optional[Mapping[str, str]] = None, + api_key: Optional[str] = None, + base_url: Optional[str] = None, + request_timeout: float = 30, + startup_timeout: float = 60, + ) -> list["AsyncioSandbox"]: + """Return borrowed snapshots with independent, not-yet-open sessions.""" + _timeout(request_timeout, "request_timeout") + _timeout(startup_timeout, "startup_timeout", allow_zero=True) + api = AsyncSandboxAPI(api_key, base_url, request_timeout) + try: + rows = await api.list(state=state, labels=labels) + snapshots = [SandboxInfo.from_dict(row) for row in rows] + except BaseException as error: + await _cleanup(api.close(), request_timeout, error) + raise + else: + await _cleanup(api.close(), request_timeout) + handles = [] + for info in snapshots: + sandbox = cls.__new__(cls) + sandbox._initialize(api_key, base_url, request_timeout, startup_timeout) + sandbox._info = info + sandbox._sandbox_id = info.id + handles.append(sandbox) + return handles + + @property + def info(self) -> SandboxInfo: + """The latest server snapshot; raises before creation, never fetches.""" + if self._info is None: + raise RuntimeError("Sandbox has not been created or fetched") + return self._info + + @property + def id(self) -> str: + """The cached sandbox identifier.""" + return self.info.id + + @property + def state(self) -> SandboxState: + """The cached lifecycle state, not a container-readiness guarantee.""" + return self.info.state + + @property + def compute(self) -> Optional[SandboxCompute]: + """The cached allocation, if one currently exists.""" + return self.info.compute + + @property + def expires_at(self) -> datetime: + """The cached maximum-lifetime deadline.""" + return self.info.expires_at + + async def _create(self) -> None: + if self._creating: + raise RuntimeError("Sandbox creation is already in progress") + if self._sandbox_id is not None: + return + if self._create_body is None: + raise RuntimeError("A borrowed sandbox cannot create a new resource") + self._creating = True + + async def create_remote() -> None: + data = await self._api.create(self._create_body) + sandbox_id = data["id"] + if not isinstance(sandbox_id, str) or not sandbox_id: + raise ValueError("Create response has no valid sandbox id") + self._sandbox_id = sandbox_id + self._owned = True + self._info = SandboxInfo.from_dict(data) + + task = asyncio.create_task(create_remote()) + try: + await asyncio.shield(task) + except BaseException as original: + + async def recover_and_release() -> None: + recovery_error = None + try: + if not task.done(): + await asyncio.wait_for(task, self._request_timeout) + else: + task.result() + except BaseException as error: + if error is not original: + recovery_error = error + try: + if self._sandbox_id is not None: + await self.terminate() + else: + await self.close() + except BaseException as error: + if recovery_error is not None: + raise error from recovery_error + raise + if recovery_error is not None: + raise recovery_error + + await _cleanup(recover_and_release(), 2 * self._request_timeout, original) + finally: + self._creating = False + + async def __aenter__(self) -> "AsyncioSandbox": + """Create if needed and enter a single, non-reentrant ownership scope.""" + if self._entering: + raise RuntimeError( + "Sandbox contexts cannot be nested or entered concurrently" + ) + self._entering = True + try: + await self._create() + return self + except BaseException: + self._entering = False + raise + + async def __aexit__( + self, + exc_type: Optional[type[BaseException]], + exc: Optional[BaseException], + traceback: Optional[TracebackType], + ) -> None: + """Release the owned resource, shielding bounded cleanup from cancellation.""" + try: + operation = self.terminate() if self._owned else self.close() + await _cleanup(operation, self._request_timeout, exc) + finally: + self._entering = False + + async def refresh(self) -> SandboxInfo: + """Fetch a fresh snapshot, including after local close or termination.""" + if self._sandbox_id is None: + raise RuntimeError("Sandbox has not been created or fetched") + self._info = SandboxInfo.from_dict(await self._api.get(self._sandbox_id)) + return self._info + + async def _ready_operation( + self, + operation: Callable[[], Awaitable[_T]], + startup_timeout: Optional[float], + *, + handshake: bool = False, + ) -> _T: + timeout = ( + self._startup_timeout + if startup_timeout is None + else _timeout(startup_timeout, "startup_timeout", allow_zero=True) + ) + sandbox_id = self.id + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + while True: + try: + if handshake and timeout > 0: + try: + return await asyncio.wait_for( + operation(), deadline - loop.time() + ) + except asyncio.TimeoutError as error: + if loop.time() >= deadline: + raise SandboxStartupTimeout(sandbox_id, timeout) from error + raise + return await operation() + except QueryError as error: + if error.status_code != 409 or timeout == 0: + raise + remaining = deadline - loop.time() + if remaining <= 0: + raise SandboxStartupTimeout(sandbox_id, timeout) from error + try: + info = await asyncio.wait_for(self.refresh(), remaining) + except asyncio.TimeoutError as refresh_error: + if loop.time() >= deadline: + raise SandboxStartupTimeout( + sandbox_id, timeout + ) from refresh_error + raise + if info.state in ("FAILED", "TERMINATED"): + raise SandboxStateError(info) from error + remaining = deadline - loop.time() + if remaining <= 0: + raise SandboxStartupTimeout(sandbox_id, timeout) from error + await asyncio.sleep(min(0.25, remaining)) + if loop.time() >= deadline: + raise SandboxStartupTimeout(sandbox_id, timeout) from error + + async def exec( + self, + command: Sequence[str], + *, + check: bool = False, + startup_timeout: Optional[float] = None, + ) -> ExecResult: + """Execute argv, retrying only explicit 409 startup rejections. + + The startup deadline limits retries, not execution of an accepted + command. Each request retains its request timeout. Transport errors, + timeouts, and other HTTP errors are never retried because the command + might already have executed. ``check=True`` preserves output on failure. + """ + if isinstance(command, (str, bytes)) or not isinstance(command, Sequence): + raise TypeError( + "command must be a nonempty sequence of strings, not a string" + ) + if not command or any(not isinstance(argument, str) for argument in command): + raise ValueError("command must be a nonempty sequence of strings") + argv = tuple(command) + data = await self._ready_operation( + lambda: self._api.exec(self.id, argv), startup_timeout + ) + result = ExecResult(output=data["output"], error=data.get("error")) + if check and result.error is not None: + raise SandboxExecutionError(self.id, result) + return result + + def logs( + self, + *, + source: Optional[LogSource] = None, + tail: Optional[int] = None, + since: Optional[str] = None, + last_event_id: Optional[str] = None, + startup_timeout: Optional[float] = None, + ) -> "AsyncSandboxLogStream": + """Return a lazy, closeable stream of main-process or lifecycle logs. + + This does not include exec output. Only the initial HTTP handshake is + retried on 409; disconnects and errors after opening are never replayed. + Use an async context or ``aclose`` when stopping iteration early. + """ + sandbox_id = self.id + stream = AsyncSandboxLogStream( + self, sandbox_id, source, tail, since, last_event_id, startup_timeout + ) + self._streams.add(stream) + return stream + + async def terminate(self) -> None: + """Delete remotely and always close locally, even if deletion fails. + + The readable remote record remains. Only successful deletion changes + the cached state; timestamps are left untouched until ``refresh``. + """ + try: + if self._sandbox_id is None: + raise RuntimeError("Sandbox has not been created or fetched") + await self._api.terminate(self._sandbox_id) + if self._info is not None: + self._info = replace(self._info, state="TERMINATED", compute=None) + except BaseException as error: + await _cleanup(self.close(), self._request_timeout, error) + else: + await _cleanup(self.close(), self._request_timeout) + + async def close(self) -> None: + """Close streams and the local session without deleting; later reuse is allowed.""" + error = None + for stream in tuple(self._streams): + try: + await stream.aclose() + except BaseException as stream_error: + if error is None: + error = stream_error + try: + await self._api.close() + except BaseException as close_error: + if error is not None: + raise error from close_error + raise + if error is not None: + raise error + + +class AsyncSandboxLogStream: + """A handle-owned typed log stream with a separately awaitable handshake.""" + + def __init__( + self, + sandbox: AsyncioSandbox, + sandbox_id: str, + source: Optional[LogSource], + tail: Optional[int], + since: Optional[str], + last_event_id: Optional[str], + startup_timeout: Optional[float], + ) -> None: + self._sandbox = sandbox + self._sandbox_id = sandbox_id + self._options = dict( + source=source, tail=tail, since=since, last_event_id=last_event_id + ) + self._startup_timeout = startup_timeout + self._stream: Any = None + self._opened = False + self._closed = False + self._opening: Optional[asyncio.Task[None]] = None + + async def _open_once(self) -> None: + stream = self._sandbox._api.logs(self._sandbox_id, **self._options) + self._stream = stream + try: + await stream.open() + except BaseException as error: + self._stream = None + await _cleanup(stream.aclose(), self._sandbox._request_timeout, error) + + async def open(self) -> "AsyncSandboxLogStream": + """Complete the HTTP handshake without waiting for the first log event.""" + if self._closed: + raise RuntimeError("Log stream is closed") + if not self._opened: + if self._opening is None: + self._opening = asyncio.create_task( + self._sandbox._ready_operation( + self._open_once, self._startup_timeout, handshake=True + ) + ) + try: + await self._opening + if self._closed: + raise RuntimeError("Log stream was closed while opening") + self._opened = True + except BaseException as error: + await _cleanup(self.aclose(), self._sandbox._request_timeout, error) + return self + + def __aiter__(self) -> "AsyncSandboxLogStream": + return self + + async def __anext__(self) -> LogEvent: + if self._closed: + raise StopAsyncIteration + try: + await self.open() + return LogEvent.from_dict(await self._stream.__anext__()) + except StopAsyncIteration: + await self.aclose() + raise + except BaseException as error: + await _cleanup(self.aclose(), self._sandbox._request_timeout, error) + raise + + async def aclose(self) -> None: + """Release the response and detach from the sandbox; idempotent.""" + if self._closed: + return + self._closed = True + self._sandbox._streams.discard(self) + if self._opening is not None and not self._opening.done(): + self._opening.cancel() + try: + await self._opening + except BaseException: + # The open caller receives its failure; close only owns disposal. + pass + stream, self._stream = self._stream, None + if stream is not None: + await stream.aclose() + + async def __aenter__(self) -> "AsyncSandboxLogStream": + return await self.open() + + async def __aexit__( + self, + exc_type: Optional[type[BaseException]], + exc: Optional[BaseException], + traceback: Optional[TracebackType], + ) -> None: + await _cleanup(self.aclose(), self._sandbox._request_timeout, exc) diff --git a/runpod/sandbox/models.py b/runpod/sandbox/models.py new file mode 100644 index 00000000..712989e1 --- /dev/null +++ b/runpod/sandbox/models.py @@ -0,0 +1,136 @@ +"""Typed snapshots and results returned by the sandbox domain.""" + +from dataclasses import dataclass +from datetime import datetime +from typing import Any, Literal, Mapping, Optional + +from runpod.error import RunPodError + +SandboxState = Literal["CREATING", "RUNNING", "TERMINATED", "FAILED"] +LogSource = Literal["container", "system"] + + +def _timestamp(value: str) -> datetime: + return datetime.fromisoformat(value.replace("Z", "+00:00")) + + +@dataclass(frozen=True) +class SandboxCompute: + vcpu_count: float + memory_in_gb: float + container_disk_in_gb: int + cost_per_hr: float + + +@dataclass(frozen=True) +class SandboxInfo: + """A server snapshot. Reading fields does not fetch or refresh anything.""" + + id: str + name: str + state: SandboxState + idle_timeout_seconds: int + max_lifetime_seconds: int + last_activity_at: datetime + idle_expires_at: datetime + expires_at: datetime + created_at: datetime + updated_at: datetime + template_id: Optional[str] + image_name: Optional[str] + cpu_flavor_id: Optional[str] + termination_reason: Optional[str] + terminated_at: Optional[datetime] + labels: Mapping[str, str] + compute: Optional[SandboxCompute] + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> "SandboxInfo": + compute = data.get("compute") + terminated_at = data.get("terminatedAt") + return cls( + id=data["id"], + name=data["name"], + state=data["state"], + idle_timeout_seconds=data["idleTimeoutSeconds"], + max_lifetime_seconds=data["maxLifetimeSeconds"], + last_activity_at=_timestamp(data["lastActivityAt"]), + idle_expires_at=_timestamp(data["idleExpiresAt"]), + expires_at=_timestamp(data["expiresAt"]), + created_at=_timestamp(data["createdAt"]), + updated_at=_timestamp(data["updatedAt"]), + template_id=data.get("templateId"), + image_name=data.get("imageName"), + cpu_flavor_id=data.get("cpuFlavorId"), + termination_reason=data.get("terminationReason"), + terminated_at=_timestamp(terminated_at) if terminated_at else None, + labels=dict(data.get("labels") or {}), + compute=( + SandboxCompute( + vcpu_count=compute["vcpuCount"], + memory_in_gb=compute["memoryInGb"], + container_disk_in_gb=compute["containerDiskInGb"], + cost_per_hr=compute["costPerHr"], + ) + if compute is not None + else None + ), + ) + + +@dataclass(frozen=True) +class ExecResult: + """Combined command output and the host's optional failure description.""" + + output: str + error: Optional[str] = None + + +@dataclass(frozen=True) +class LogEvent: + """One log record; preserve the opaque SSE id for subsequent resume requests.""" + + source: LogSource + line: str + timestamp: datetime + id: Optional[str] = None + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> "LogEvent": + return cls( + source=data["source"], + line=data["line"], + timestamp=_timestamp(data["ts"]), + id=data.get("id"), + ) + + +class SandboxExecutionError(RunPodError): + """A command failed with check=True; its partial output remains available.""" + + def __init__(self, sandbox_id: str, result: ExecResult): + super().__init__(result.error) + self.sandbox_id = sandbox_id + self.result = result + + +class SandboxStateError(RunPodError): + """An operation cannot proceed because the sandbox reached a terminal state.""" + + def __init__(self, info: SandboxInfo): + super().__init__( + f"Sandbox {info.id} is {info.state}" + + (f": {info.termination_reason}" if info.termination_reason else "") + ) + self.info = info + + +class SandboxStartupTimeout(TimeoutError): + """The sandbox did not accept the operation within the startup deadline.""" + + def __init__(self, sandbox_id: str, timeout: float): + super().__init__( + f"Sandbox {sandbox_id} did not become ready within {timeout:g}s" + ) + self.sandbox_id = sandbox_id + self.timeout = timeout diff --git a/runpod/sandbox/sync.py b/runpod/sandbox/sync.py new file mode 100644 index 00000000..b5f9db64 --- /dev/null +++ b/runpod/sandbox/sync.py @@ -0,0 +1,438 @@ +"""Blocking sandbox facade backed by one stable asyncio loop per active handle.""" + +import asyncio +import threading +from concurrent.futures import Future +from datetime import datetime +from typing import Any, Coroutine, Iterator, Mapping, Optional, Sequence, TypeVar + +from .asyncio import AsyncioSandbox +from .models import ( + ExecResult, + LogEvent, + LogSource, + SandboxCompute, + SandboxInfo, + SandboxState, +) + +_T = TypeVar("_T") + + +class _LoopRunner: + """Keep aiohttp sessions on their original loop, including across log reads.""" + + def __init__(self) -> None: + self._loop: Optional[asyncio.AbstractEventLoop] = None + self._thread: Optional[threading.Thread] = None + self._lock = threading.Lock() + + @property + def started(self) -> bool: + return self._thread is not None + + def _start(self) -> asyncio.AbstractEventLoop: + with self._lock: + if self._loop is None: + ready = threading.Event() + + def serve() -> None: + loop = asyncio.new_event_loop() + self._loop = loop + ready.set() + try: + loop.run_forever() + finally: + loop.run_until_complete(loop.shutdown_asyncgens()) + loop.run_until_complete(loop.shutdown_default_executor()) + loop.close() + + self._thread = threading.Thread( + target=serve, name="runpod-sandbox", daemon=True + ) + self._thread.start() + ready.wait() + assert self._loop is not None + return self._loop + + def run(self, coroutine: Coroutine[Any, Any, _T]) -> _T: + if threading.current_thread() is self._thread: + coroutine.close() + raise RuntimeError("Cannot call the sync sandbox API from its own loop") + loop = self._start() + future: Future[_T] = Future() + finished = threading.Event() + task: Optional[asyncio.Task[tuple[Optional[_T], Optional[BaseException]]]] = ( + None + ) + invoked = False + + async def invoke() -> tuple[Optional[_T], Optional[BaseException]]: + nonlocal invoked + invoked = True + try: + return await coroutine, None + except BaseException as error: + # KeyboardInterrupt/SystemExit must reach the calling thread, + # not escape Task execution and stop the background loop. + return None, error + + def complete(done: asyncio.Task) -> None: + try: + result, error = done.result() + if error is not None: + future.set_exception(error) + else: + future.set_result(result) + except BaseException as error: + future.set_exception(error) + finally: + if not invoked: + coroutine.close() + finished.set() + + def submit() -> None: + nonlocal task + task = loop.create_task(invoke()) + task.add_done_callback(complete) + + def cancel() -> None: + if task is not None: + task.cancel() + + loop.call_soon_threadsafe(submit) + try: + return future.result() + except BaseException: + if not future.done(): + # Cancelling a concurrent Future marks it done before the async + # task's finally blocks finish. Wait for the actual task before + # stopping its loop, or Ctrl-C can strand a newly created sandbox. + loop.call_soon_threadsafe(cancel) + while not finished.is_set(): + try: + finished.wait() + except KeyboardInterrupt: + # Cleanup is bounded by the async domain's timeouts. + continue + raise + + def stop(self) -> None: + with self._lock: + loop, thread = self._loop, self._thread + if loop is None or thread is None: + return + + async def drain() -> None: + current = asyncio.current_task() + tasks = [task for task in asyncio.all_tasks() if task is not current] + for task in tasks: + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + + try: + asyncio.run_coroutine_threadsafe(drain(), loop).result() + finally: + loop.call_soon_threadsafe(loop.stop) + thread.join() + self._loop = None + self._thread = None + + +class SandboxLogs(Iterator[LogEvent]): + """Closeable log iterator. Use a with block when stopping consumption early.""" + + def __init__(self, stream: Any, runner: _LoopRunner) -> None: + self._stream = stream + self._runner = runner + self._closed = False + + def __enter__(self) -> "SandboxLogs": + if self._closed: + raise RuntimeError("Log stream is closed") + self._runner.run(self._stream.__aenter__()) + return self + + def __exit__(self, exc_type: Any, exc: Any, traceback: Any) -> bool: + self._closed = True + if self._runner.started: + self._runner.run(self._stream.__aexit__(exc_type, exc, traceback)) + return False + + def __iter__(self) -> "SandboxLogs": + return self + + def __next__(self) -> LogEvent: + if self._closed: + raise StopIteration + try: + return self._runner.run(self._stream.__anext__()) + except StopAsyncIteration: + self._closed = True + raise StopIteration from None + except BaseException: + self.close() + raise + + def close(self) -> None: + if not self._closed: + self._closed = True + if self._runner.started: + self._runner.run(self._stream.aclose()) + + +class Sandbox: + """An isolated CPU sandbox with blocking methods. + + ``with Sandbox(image_name=...)`` creates on entry and terminates on exit. + ``create`` leaves the lifetime under your control; call ``terminate`` when + finished, or ``close`` to release only local connections. Handles retrieved + with ``get`` or ``list`` never terminate compute on context exit. + + A stable background event loop runs the async implementation. No thread or + connection is started by the constructor. Prefer AsyncioSandbox in async + applications so blocking calls do not pause the application's event loop. + """ + + def __init__( + self, + *, + image_name: Optional[str] = None, + template_id: Optional[str] = None, + name: Optional[str] = None, + cpu_flavor_id: Optional[str] = None, + vcpu_count: Optional[int] = None, + memory_in_gb: Optional[int] = None, + data_center_id: Optional[str] = None, + env: Optional[Mapping[str, str]] = None, + idle_timeout_seconds: Optional[int] = None, + max_lifetime_seconds: Optional[int] = None, + labels: Optional[Mapping[str, str]] = None, + api_key: Optional[str] = None, + base_url: Optional[str] = None, + request_timeout: float = 30, + startup_timeout: float = 60, + ) -> None: + self._runner = _LoopRunner() + self._entered = False + self._sandbox = AsyncioSandbox( + image_name=image_name, + template_id=template_id, + name=name, + cpu_flavor_id=cpu_flavor_id, + vcpu_count=vcpu_count, + memory_in_gb=memory_in_gb, + data_center_id=data_center_id, + env=env, + idle_timeout_seconds=idle_timeout_seconds, + max_lifetime_seconds=max_lifetime_seconds, + labels=labels, + api_key=api_key, + base_url=base_url, + request_timeout=request_timeout, + startup_timeout=startup_timeout, + ) + + @classmethod + def _from_async( + cls, sandbox: AsyncioSandbox, runner: Optional[_LoopRunner] = None + ) -> "Sandbox": + instance = cls.__new__(cls) + instance._sandbox = sandbox + instance._runner = runner if runner is not None else _LoopRunner() + instance._entered = False + return instance + + @classmethod + def create( + cls, + *, + image_name: Optional[str] = None, + template_id: Optional[str] = None, + name: Optional[str] = None, + cpu_flavor_id: Optional[str] = None, + vcpu_count: Optional[int] = None, + memory_in_gb: Optional[int] = None, + data_center_id: Optional[str] = None, + env: Optional[Mapping[str, str]] = None, + idle_timeout_seconds: Optional[int] = None, + max_lifetime_seconds: Optional[int] = None, + labels: Optional[Mapping[str, str]] = None, + api_key: Optional[str] = None, + base_url: Optional[str] = None, + request_timeout: float = 30, + startup_timeout: float = 60, + ) -> "Sandbox": + """Create immediately; the returned snapshot may still be CREATING.""" + runner = _LoopRunner() + try: + sandbox = runner.run( + AsyncioSandbox.create( + image_name=image_name, + template_id=template_id, + name=name, + cpu_flavor_id=cpu_flavor_id, + vcpu_count=vcpu_count, + memory_in_gb=memory_in_gb, + data_center_id=data_center_id, + env=env, + idle_timeout_seconds=idle_timeout_seconds, + max_lifetime_seconds=max_lifetime_seconds, + labels=labels, + api_key=api_key, + base_url=base_url, + request_timeout=request_timeout, + startup_timeout=startup_timeout, + ) + ) + except BaseException: + runner.stop() + raise + return cls._from_async(sandbox, runner) + + @classmethod + def get( + cls, + sandbox_id: str, + *, + api_key: Optional[str] = None, + base_url: Optional[str] = None, + request_timeout: float = 30, + startup_timeout: float = 60, + ) -> "Sandbox": + """Fetch a borrowed handle; closing its context does not terminate it.""" + runner = _LoopRunner() + try: + sandbox = runner.run( + AsyncioSandbox.get( + sandbox_id, + api_key=api_key, + base_url=base_url, + request_timeout=request_timeout, + startup_timeout=startup_timeout, + ) + ) + except BaseException: + runner.stop() + raise + return cls._from_async(sandbox, runner) + + @classmethod + def list( + cls, + *, + state: Optional[SandboxState] = None, + labels: Optional[Mapping[str, str]] = None, + api_key: Optional[str] = None, + base_url: Optional[str] = None, + request_timeout: float = 30, + startup_timeout: float = 60, + ) -> list["Sandbox"]: + """List borrowed handles. Filters combine with AND on the server.""" + runner = _LoopRunner() + try: + sandboxes = runner.run( + AsyncioSandbox.list( + state=state, + labels=labels, + api_key=api_key, + base_url=base_url, + request_timeout=request_timeout, + startup_timeout=startup_timeout, + ) + ) + finally: + runner.stop() + return [cls._from_async(sandbox) for sandbox in sandboxes] + + @property + def info(self) -> SandboxInfo: + return self._sandbox.info + + @property + def id(self) -> str: + return self._sandbox.id + + @property + def state(self) -> SandboxState: + return self._sandbox.state + + @property + def compute(self) -> Optional[SandboxCompute]: + return self._sandbox.compute + + @property + def expires_at(self) -> datetime: + return self._sandbox.expires_at + + def refresh(self) -> SandboxInfo: + """Fetch current server metadata and replace the local snapshot.""" + return self._runner.run(self._sandbox.refresh()) + + def exec( + self, + command: Sequence[str], + *, + check: bool = False, + startup_timeout: Optional[float] = None, + ) -> ExecResult: + """Execute argv, waiting only for explicit startup rejections.""" + return self._runner.run( + self._sandbox.exec(command, check=check, startup_timeout=startup_timeout) + ) + + def logs( + self, + *, + source: Optional[LogSource] = None, + tail: Optional[int] = None, + since: Optional[str] = None, + last_event_id: Optional[str] = None, + startup_timeout: Optional[float] = None, + ) -> SandboxLogs: + """Stream container/system logs, not exec output; close on early exit.""" + return SandboxLogs( + self._sandbox.logs( + source=source, + tail=tail, + since=since, + last_event_id=last_event_id, + startup_timeout=startup_timeout, + ), + self._runner, + ) + + def terminate(self) -> None: + """Release remote compute and local connections; safe to repeat.""" + try: + self._runner.run(self._sandbox.terminate()) + finally: + self._runner.stop() + + def close(self) -> None: + """Release local connections and the loop, without terminating compute.""" + if self._runner.started: + try: + self._runner.run(self._sandbox.close()) + finally: + self._runner.stop() + + def __enter__(self) -> "Sandbox": + if self._entered: + raise RuntimeError("Sandbox context is already entered") + try: + self._runner.run(self._sandbox.__aenter__()) + except BaseException: + self._runner.stop() + raise + self._entered = True + return self + + def __exit__(self, exc_type: Any, exc: Any, traceback: Any) -> bool: + try: + self._runner.run(self._sandbox.__aexit__(exc_type, exc, traceback)) + return False + finally: + self._entered = False + self._runner.stop() diff --git a/tests/test_sandbox.py b/tests/test_sandbox.py new file mode 100644 index 00000000..a326e0f3 --- /dev/null +++ b/tests/test_sandbox.py @@ -0,0 +1,458 @@ +"""Sandbox lifecycle and streaming regressions over an actual HTTP connection.""" + +import asyncio +import json +import subprocess +import sys +import threading +from collections import deque +from contextlib import contextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from urllib.parse import parse_qs, unquote, urlsplit + +import pytest + +from runpod import AsyncioSandbox, Sandbox +from runpod.error import QueryError +from runpod.sandbox import ( + SandboxExecutionError, + SandboxStartupTimeout, + SandboxStateError, +) + + +class SandboxService: + """A controllable REST peer for lifecycle races, not a mocked SDK transport.""" + + def __init__(self): + self.records = {} + self.requests = [] + self.executed = [] + self.exec_statuses = deque() + self.log_statuses = deque() + self.delete_status = 204 + self.create_started = threading.Event() + self.allow_create = threading.Event() + self.allow_create.set() + self.log_started = threading.Event() + self.allow_logs = threading.Event() + self.allow_logs.set() + self.log_disconnected = threading.Event() + self.hold_logs = False + self.stopping = threading.Event() + # Mix line endings, split a UTF-8 codepoint, preserve an opaque SSE id, + # and include both a multiline data event and a non-JSON timeout frame. + self.log_bytes = ( + ": heartbeat\r\n\r\nid: opaque/42\r\nevent: message\r\n" + 'data: {"source":"container",\r\n' + 'data: "line":"caf\u00e9","ts":"2026-09-14T12:00:00Z"}\r\n\r\n' + 'id: opaque/43\rdata: {"source":"system","line":"started",' + '"ts":"2026-09-14T12:00:01Z"}\r\r' + "event: timeout\ndata: max stream duration reached\n\n" + ).encode() + + def create_record(self, body): + sandbox_id = f"sandbox-{len(self.records) + 1}" + now = "2026-09-14T12:00:00Z" + record = { + "id": sandbox_id, + "name": body.get("name", sandbox_id), + "state": "RUNNING", + "imageName": body.get("imageName"), + "templateId": body.get("templateId"), + "idleTimeoutSeconds": body.get("idleTimeoutSeconds", 300), + "maxLifetimeSeconds": body.get("maxLifetimeSeconds", 900), + "lastActivityAt": now, + "idleExpiresAt": "2026-09-14T12:05:00Z", + "expiresAt": "2026-09-14T12:15:00Z", + "createdAt": now, + "updatedAt": now, + "labels": body.get("labels", {}), + "compute": { + "vcpuCount": body.get("vcpuCount", 2), + "memoryInGb": body.get("memoryInGb", 4), + "containerDiskInGb": 10, + "costPerHr": 0.026, + }, + } + self.records[sandbox_id] = record + return record + + +@contextmanager +def sandbox_peer(): + service = SandboxService() + + class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def log_message(self, *_args): + pass + + def reply(self, status, data=None): + payload = json.dumps(data).encode() if data is not None else b"" + self.send_response(status) + self.send_header( + "Content-Type", + "application/problem+json" if status >= 400 else "application/json", + ) + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + try: + self.wfile.write(payload) + except (BrokenPipeError, ConnectionResetError): + pass + + def handle_request(self): + parsed = urlsplit(self.path) + segments = [unquote(part) for part in parsed.path.split("/") if part] + content_length = int(self.headers.get("Content-Length", "0")) + body = ( + json.loads(self.rfile.read(content_length)) if content_length else None + ) + service.requests.append( + (self.command, parsed.path, body, dict(self.headers)) + ) + if self.headers.get("Authorization") != "Bearer sandbox-test-key": + self.reply(401, {"detail": "invalid key"}) + return + if segments == ["v2", "sandboxes"]: + if self.command == "POST": + record = service.create_record(body) + service.create_started.set() + service.allow_create.wait(5) + self.reply(201, record) + else: + query = parse_qs(parsed.query) + state = query.get("state", [None])[0] + records = [ + record + for record in service.records.values() + if ( + record["state"] == state + if state + else record["state"] != "TERMINATED" + ) + ] + for label in query.get("labels", []): + key, value = label.split("=", 1) + records = [ + record + for record in records + if record["labels"].get(key) == value + ] + self.reply(200, {"sandboxes": records}) + return + if ( + len(segments) < 3 + or segments[:2] != ["v2", "sandboxes"] + or segments[2] not in service.records + ): + self.reply(404, {"detail": "missing sandbox"}) + return + record = service.records[segments[2]] + if len(segments) == 3: + if self.command == "DELETE": + if service.delete_status == 204: + record.update( + state="TERMINATED", + terminatedAt="2026-09-14T12:01:00Z", + terminationReason="USER", + compute=None, + ) + self.reply(204) + else: + self.reply( + service.delete_status, {"detail": "cleanup unavailable"} + ) + else: + self.reply(200, record) + elif segments[3] == "exec": + status = ( + service.exec_statuses.popleft() if service.exec_statuses else 200 + ) + if record["state"] in ("FAILED", "TERMINATED"): + status = 409 + if status == 409: + self.reply(409, {"detail": "container not started"}) + return + service.executed.append(body["command"]) + if status >= 400: + self.reply( + status, {"detail": "response lost after command execution"} + ) + return + if body["command"] == ["fail"]: + self.reply( + 200, {"output": "partial output", "error": "command failed"} + ) + else: + self.reply(200, {"output": "completed", "error": None}) + elif segments[3] == "logs": + service.log_started.set() + service.allow_logs.wait(5) + status = service.log_statuses.popleft() if service.log_statuses else 200 + if status != 200: + self.reply(status, {"detail": "logs not ready"}) + return + self.send_response(200) + self.send_header("Content-Type", "text/event-stream") + self.send_header("Connection", "close") + self.end_headers() + self.close_connection = True + try: + # Byte-at-a-time writes ensure parser state is exercised, + # regardless of how the client batches socket reads. + for byte in service.log_bytes: + self.wfile.write(bytes([byte])) + self.wfile.flush() + while service.hold_logs and not service.stopping.wait(0.01): + self.wfile.write(b": heartbeat\n\n") + self.wfile.flush() + except (BrokenPipeError, ConnectionResetError): + pass + finally: + service.log_disconnected.set() + + do_GET = handle_request + do_POST = handle_request + do_DELETE = handle_request + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + service.options = { + "api_key": "sandbox-test-key", + "base_url": f"http://127.0.0.1:{server.server_port}", + "request_timeout": 2, + "startup_timeout": 2, + } + try: + yield service + finally: + service.allow_create.set() + service.stopping.set() + server.shutdown() + server.server_close() + thread.join() + + +@pytest.fixture +def peer(): + with sandbox_peer() as service: + yield service + + +def test_sync_startup_conflict_and_borrowed_context_ownership(peer): + peer.exec_statuses.extend([409, 200]) + with Sandbox(image_name="python:3.12-slim", **peer.options) as owner: + with Sandbox.get(owner.id, **peer.options) as borrowed: + assert borrowed.exec(["work"]).output == "completed" + assert peer.records[owner.id]["state"] == "RUNNING" + # A rejected nested entry must not destroy the outer context's client. + with pytest.raises(RuntimeError): + with owner: + pytest.fail("nested context entry succeeded") + assert owner.refresh().state == "RUNNING" + assert peer.records[owner.id]["state"] == "TERMINATED" + assert peer.executed == [["work"]] + + +@pytest.mark.asyncio +async def test_async_startup_conflict_and_borrowed_context_ownership(peer): + peer.exec_statuses.extend([409, 200]) + async with AsyncioSandbox(image_name="python:3.12-slim", **peer.options) as owner: + async with await AsyncioSandbox.get(owner.id, **peer.options) as borrowed: + assert (await borrowed.exec(["work"])).output == "completed" + assert peer.records[owner.id]["state"] == "RUNNING" + with pytest.raises(RuntimeError): + async with owner: + pytest.fail("nested context entry succeeded") + assert (await owner.refresh()).state == "RUNNING" + assert peer.records[owner.id]["state"] == "TERMINATED" + assert peer.executed == [["work"]] + + +@pytest.mark.asyncio +async def test_cancellation_during_creation_recovers_id_and_terminates(peer): + peer.allow_create.clear() + + async def create(): + async with AsyncioSandbox(image_name="python:3.12-slim", **peer.options): + pytest.fail("cancelled context was entered") + + task = asyncio.create_task(create()) + assert await asyncio.to_thread(peer.create_started.wait, 2) + task.cancel() + peer.allow_create.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, 4) + assert peer.records["sandbox-1"]["state"] == "TERMINATED" + + +def test_sync_interrupt_during_creation_waits_for_resource_cleanup(): + script = """ +import os +import signal +import threading +from runpod import Sandbox +from tests.test_sandbox import sandbox_peer + +with sandbox_peer() as peer: + peer.allow_create.clear() + def interrupt(): + assert peer.create_started.wait(2) + os.kill(os.getpid(), signal.SIGINT) + threading.Timer(0.1, peer.allow_create.set).start() + worker = threading.Thread(target=interrupt) + worker.start() + try: + Sandbox.create(image_name="python:3.12-slim", **peer.options) + raise AssertionError("SIGINT was swallowed") + except KeyboardInterrupt: + assert peer.records["sandbox-1"]["state"] == "TERMINATED" + finally: + worker.join() +""" + process = subprocess.run( + [sys.executable, "-c", script], capture_output=True, text=True, timeout=10 + ) + assert process.returncode == 0, process.stdout + process.stderr + + +@pytest.mark.parametrize("delete_status", [204, 503]) +@pytest.mark.parametrize("failure_type", [ValueError, KeyboardInterrupt, SystemExit]) +def test_sync_body_exception_remains_primary_during_cleanup( + peer, delete_status, failure_type +): + peer.delete_status = delete_status + failure = failure_type("application failed") + with pytest.raises(failure_type) as raised: + with Sandbox(image_name="python:3.12-slim", **peer.options): + raise failure + assert raised.value is failure + if delete_status == 204: + assert peer.records["sandbox-1"]["state"] == "TERMINATED" + else: + assert isinstance(raised.value.__cause__, QueryError) + assert raised.value.__cause__.status_code == 503 + assert peer.records["sandbox-1"]["state"] == "RUNNING" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("delete_status", [204, 503]) +async def test_async_body_exception_remains_primary_during_cleanup(peer, delete_status): + peer.delete_status = delete_status + failure = ValueError("application failed") + with pytest.raises(ValueError) as raised: + async with AsyncioSandbox(image_name="python:3.12-slim", **peer.options): + raise failure + assert raised.value is failure + if delete_status == 204: + assert peer.records["sandbox-1"]["state"] == "TERMINATED" + else: + assert isinstance(raised.value.__cause__, QueryError) + assert raised.value.__cause__.status_code == 503 + + +@pytest.mark.asyncio +async def test_exec_never_replays_ambiguous_failure_and_preserves_partial_output(peer): + async with AsyncioSandbox(image_name="python:3.12-slim", **peer.options) as sandbox: + peer.exec_statuses.append(500) + with pytest.raises(QueryError) as raised: + await sandbox.exec(["side-effect"]) + assert raised.value.status_code == 500 + assert peer.executed == [["side-effect"]] + result = await sandbox.exec(["fail"]) + assert result.output == "partial output" + assert result.error == "command failed" + with pytest.raises(SandboxExecutionError) as failure: + await sandbox.exec(["fail"], check=True) + assert failure.value.result == result + assert failure.value.sandbox_id == sandbox.id + + +@pytest.mark.asyncio +async def test_startup_deadline_and_terminal_state_do_not_execute_commands(peer): + async with AsyncioSandbox(image_name="python:3.12-slim", **peer.options) as sandbox: + peer.exec_statuses.extend([409] * 20) + with pytest.raises(QueryError) as rejected: + await sandbox.exec(["work"], startup_timeout=0) + assert rejected.value.status_code == 409 + with pytest.raises(SandboxStartupTimeout) as expired: + await sandbox.exec(["work"], startup_timeout=0.03) + assert expired.value.sandbox_id == sandbox.id + peer.records[sandbox.id]["state"] = "FAILED" + with pytest.raises(SandboxStateError) as failed: + await sandbox.exec(["work"]) + assert failed.value.info.state == "FAILED" + assert peer.executed == [] + + +def test_sync_sse_decodes_fragments_and_retains_resume_cursor(peer): + with Sandbox(image_name="python:3.12-slim", **peer.options) as sandbox: + with sandbox.logs(source="container", tail=0) as logs: + events = list(logs) + assert [(event.source, event.line, event.id) for event in events] == [ + ("container", "caf\u00e9", "opaque/42"), + ("system", "started", "opaque/43"), + ] + assert events[0].timestamp.isoformat() == "2026-09-14T12:00:00+00:00" + with sandbox.logs(last_event_id=events[-1].id) as logs: + assert next(logs).line == "caf\u00e9" + log_requests = [ + request for request in peer.requests if request[1].endswith("/logs") + ] + assert log_requests[-1][3]["Last-Event-ID"] == "opaque/43" + + +@pytest.mark.asyncio +async def test_async_log_early_close_releases_live_connection(peer): + peer.hold_logs = True + # Exclude the timeout frame so this stays open until the client closes it. + peer.log_bytes = peer.log_bytes.split(b"event: timeout")[0] + peer.log_statuses.extend([409, 200]) + async with AsyncioSandbox(image_name="python:3.12-slim", **peer.options) as sandbox: + async with sandbox.logs() as logs: + event = await logs.__anext__() + assert event.line == "caf\u00e9" + assert await asyncio.to_thread(peer.log_disconnected.wait, 2) + + +@pytest.mark.asyncio +async def test_closing_sandbox_cancels_pending_log_handshake(peer): + peer.allow_logs.clear() + sandbox = await AsyncioSandbox.create(image_name="python:3.12-slim", **peer.options) + logs = sandbox.logs() + opening = asyncio.create_task(logs.open()) + try: + assert await asyncio.to_thread(peer.log_started.wait, 2) + await asyncio.wait_for(sandbox.close(), 3) + with pytest.raises(asyncio.CancelledError): + await opening + assert peer.records[sandbox.id]["state"] == "RUNNING" + finally: + peer.allow_logs.set() + await sandbox.terminate() + + +@pytest.mark.asyncio +async def test_list_filters_and_independent_handles_do_not_delete_resources(peer): + first = peer.create_record( + {"imageName": "image", "labels": {"team": "a", "job": "42"}} + ) + second = peer.create_record( + {"imageName": "image", "labels": {"team": "a", "job": "43"}} + ) + peer.create_record({"imageName": "image", "labels": {"team": "b", "job": "42"}}) + matches = await AsyncioSandbox.list( + labels={"team": "a", "job": "42"}, **peer.options + ) + assert [sandbox.id for sandbox in matches] == [first["id"]] + await matches[0].close() + handles = await AsyncioSandbox.list(labels={"team": "a"}, **peer.options) + await handles[0].close() + async with handles[1] as borrowed: + assert (await borrowed.exec(["work"])).output == "completed" + assert borrowed.id == second["id"] + assert all(record["state"] == "RUNNING" for record in peer.records.values()) From 4fc05cbc21d7c1a6620541e48f05a587a8062985 Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Mon, 14 Sep 2026 11:07:20 -0700 Subject: [PATCH 2/6] refactor: simplify sandbox factories and test peer --- runpod/sandbox/asyncio.py | 47 +------ runpod/sandbox/sync.py | 46 +------ tests/test_sandbox.py | 250 ++++++++++++++++---------------------- 3 files changed, 113 insertions(+), 230 deletions(-) diff --git a/runpod/sandbox/asyncio.py b/runpod/sandbox/asyncio.py index 849c9494..17a6b860 100644 --- a/runpod/sandbox/asyncio.py +++ b/runpod/sandbox/asyncio.py @@ -89,7 +89,8 @@ def __init__( ) -> None: if (image_name is None) == (template_id is None): raise ValueError("Supply exactly one of image_name or template_id") - options = { + self._initialize(api_key, base_url, request_timeout, startup_timeout) + self._create_body = { "imageName": image_name, "templateId": template_id, "name": name, @@ -102,10 +103,6 @@ def __init__( "maxLifetimeSeconds": max_lifetime_seconds, "labels": dict(labels) if labels is not None else None, } - self._initialize(api_key, base_url, request_timeout, startup_timeout) - self._create_body = { - key: value for key, value in options.items() if value is not None - } def _initialize( self, @@ -128,43 +125,9 @@ def _initialize( self._streams: set[AsyncSandboxLogStream] = set() @classmethod - async def create( - cls, - *, - image_name: Optional[str] = None, - template_id: Optional[str] = None, - name: Optional[str] = None, - cpu_flavor_id: Optional[str] = None, - vcpu_count: Optional[int] = None, - memory_in_gb: Optional[int] = None, - data_center_id: Optional[str] = None, - env: Optional[Mapping[str, str]] = None, - idle_timeout_seconds: Optional[int] = None, - max_lifetime_seconds: Optional[int] = None, - labels: Optional[Mapping[str, str]] = None, - api_key: Optional[str] = None, - base_url: Optional[str] = None, - request_timeout: float = 30, - startup_timeout: float = 60, - ) -> "AsyncioSandbox": + async def create(cls, **options: Any) -> "AsyncioSandbox": """Create and return an owned handle, without claiming container readiness.""" - sandbox = cls( - image_name=image_name, - template_id=template_id, - name=name, - cpu_flavor_id=cpu_flavor_id, - vcpu_count=vcpu_count, - memory_in_gb=memory_in_gb, - data_center_id=data_center_id, - env=env, - idle_timeout_seconds=idle_timeout_seconds, - max_lifetime_seconds=max_lifetime_seconds, - labels=labels, - api_key=api_key, - base_url=base_url, - request_timeout=request_timeout, - startup_timeout=startup_timeout, - ) + sandbox = cls(**options) await sandbox._create() return sandbox @@ -208,7 +171,6 @@ async def list( snapshots = [SandboxInfo.from_dict(row) for row in rows] except BaseException as error: await _cleanup(api.close(), request_timeout, error) - raise else: await _cleanup(api.close(), request_timeout) handles = [] @@ -534,7 +496,6 @@ async def __anext__(self) -> LogEvent: raise except BaseException as error: await _cleanup(self.aclose(), self._sandbox._request_timeout, error) - raise async def aclose(self) -> None: """Release the response and detach from the sandbox; idempotent.""" diff --git a/runpod/sandbox/sync.py b/runpod/sandbox/sync.py index b5f9db64..9658eeb1 100644 --- a/runpod/sandbox/sync.py +++ b/runpod/sandbox/sync.py @@ -245,51 +245,15 @@ def _from_async( return instance @classmethod - def create( - cls, - *, - image_name: Optional[str] = None, - template_id: Optional[str] = None, - name: Optional[str] = None, - cpu_flavor_id: Optional[str] = None, - vcpu_count: Optional[int] = None, - memory_in_gb: Optional[int] = None, - data_center_id: Optional[str] = None, - env: Optional[Mapping[str, str]] = None, - idle_timeout_seconds: Optional[int] = None, - max_lifetime_seconds: Optional[int] = None, - labels: Optional[Mapping[str, str]] = None, - api_key: Optional[str] = None, - base_url: Optional[str] = None, - request_timeout: float = 30, - startup_timeout: float = 60, - ) -> "Sandbox": + def create(cls, **options: Any) -> "Sandbox": """Create immediately; the returned snapshot may still be CREATING.""" - runner = _LoopRunner() + sandbox = cls(**options) try: - sandbox = runner.run( - AsyncioSandbox.create( - image_name=image_name, - template_id=template_id, - name=name, - cpu_flavor_id=cpu_flavor_id, - vcpu_count=vcpu_count, - memory_in_gb=memory_in_gb, - data_center_id=data_center_id, - env=env, - idle_timeout_seconds=idle_timeout_seconds, - max_lifetime_seconds=max_lifetime_seconds, - labels=labels, - api_key=api_key, - base_url=base_url, - request_timeout=request_timeout, - startup_timeout=startup_timeout, - ) - ) + sandbox._runner.run(sandbox._sandbox._create()) except BaseException: - runner.stop() + sandbox.close() raise - return cls._from_async(sandbox, runner) + return sandbox @classmethod def get( diff --git a/tests/test_sandbox.py b/tests/test_sandbox.py index a326e0f3..88ba9d10 100644 --- a/tests/test_sandbox.py +++ b/tests/test_sandbox.py @@ -1,16 +1,14 @@ """Sandbox lifecycle and streaming regressions over an actual HTTP connection.""" import asyncio -import json import subprocess import sys import threading from collections import deque from contextlib import contextmanager -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from urllib.parse import parse_qs, unquote, urlsplit import pytest +from aiohttp import web from runpod import AsyncioSandbox, Sandbox from runpod.error import QueryError @@ -82,158 +80,118 @@ def create_record(self, body): @contextmanager def sandbox_peer(): service = SandboxService() - - class Handler(BaseHTTPRequestHandler): - protocol_version = "HTTP/1.1" - - def log_message(self, *_args): - pass - - def reply(self, status, data=None): - payload = json.dumps(data).encode() if data is not None else b"" - self.send_response(status) - self.send_header( - "Content-Type", - "application/problem+json" if status >= 400 else "application/json", + ready, stopped = threading.Event(), asyncio.Event() + + def failure(status, detail): + return web.json_response( + {"detail": detail}, status=status, content_type="application/problem+json" + ) + + async def handle(request): + body = await request.json() if request.can_read_body else None + service.requests.append( + (request.method, request.path, body, dict(request.headers)) + ) + if request.headers.get("Authorization") != "Bearer sandbox-test-key": + return failure(401, "invalid key") + sandbox_id = request.match_info.get("id") + if sandbox_id is None: + if request.method == "POST": + record = service.create_record(body) + service.create_started.set() + await asyncio.to_thread(service.allow_create.wait, 5) + return web.json_response(record, status=201) + state = request.query.get("state") + labels = [term.split("=", 1) for term in request.query.getall("labels", [])] + records = [ + record + for record in service.records.values() + if ( + record["state"] == state + if state + else record["state"] != "TERMINATED" + ) + and all(record["labels"].get(key) == value for key, value in labels) + ] + return web.json_response({"sandboxes": records}) + if sandbox_id not in service.records: + return failure(404, "missing sandbox") + record = service.records[sandbox_id] + operation = request.match_info.get("operation") + if request.method == "DELETE": + if service.delete_status != 204: + return failure(service.delete_status, "cleanup unavailable") + record.update(state="TERMINATED", compute=None) + return web.Response(status=204) + if operation == "exec": + status = service.exec_statuses.popleft() if service.exec_statuses else 200 + if status == 409 or record["state"] in ("FAILED", "TERMINATED"): + return failure(409, "container not started") + service.executed.append(body["command"]) + if status >= 400: + return failure(status, "response lost after command execution") + result = ( + {"output": "partial output", "error": "command failed"} + if body["command"] == ["fail"] + else {"output": "completed", "error": None} ) - self.send_header("Content-Length", str(len(payload))) - self.end_headers() + return web.json_response(result) + if operation == "logs": + service.log_started.set() + await asyncio.to_thread(service.allow_logs.wait, 5) + status = service.log_statuses.popleft() if service.log_statuses else 200 + if status != 200: + return failure(status, "logs not ready") + response = web.StreamResponse(headers={"Content-Type": "text/event-stream"}) try: - self.wfile.write(payload) - except (BrokenPipeError, ConnectionResetError): + await response.prepare(request) + for byte in service.log_bytes: + await response.write(bytes([byte])) + await asyncio.sleep(0) + while service.hold_logs and not service.stopping.is_set(): + await response.write(b": heartbeat\n\n") + await asyncio.sleep(0.01) + except ConnectionResetError: pass - - def handle_request(self): - parsed = urlsplit(self.path) - segments = [unquote(part) for part in parsed.path.split("/") if part] - content_length = int(self.headers.get("Content-Length", "0")) - body = ( - json.loads(self.rfile.read(content_length)) if content_length else None - ) - service.requests.append( - (self.command, parsed.path, body, dict(self.headers)) - ) - if self.headers.get("Authorization") != "Bearer sandbox-test-key": - self.reply(401, {"detail": "invalid key"}) - return - if segments == ["v2", "sandboxes"]: - if self.command == "POST": - record = service.create_record(body) - service.create_started.set() - service.allow_create.wait(5) - self.reply(201, record) - else: - query = parse_qs(parsed.query) - state = query.get("state", [None])[0] - records = [ - record - for record in service.records.values() - if ( - record["state"] == state - if state - else record["state"] != "TERMINATED" - ) - ] - for label in query.get("labels", []): - key, value = label.split("=", 1) - records = [ - record - for record in records - if record["labels"].get(key) == value - ] - self.reply(200, {"sandboxes": records}) - return - if ( - len(segments) < 3 - or segments[:2] != ["v2", "sandboxes"] - or segments[2] not in service.records - ): - self.reply(404, {"detail": "missing sandbox"}) - return - record = service.records[segments[2]] - if len(segments) == 3: - if self.command == "DELETE": - if service.delete_status == 204: - record.update( - state="TERMINATED", - terminatedAt="2026-09-14T12:01:00Z", - terminationReason="USER", - compute=None, - ) - self.reply(204) - else: - self.reply( - service.delete_status, {"detail": "cleanup unavailable"} - ) - else: - self.reply(200, record) - elif segments[3] == "exec": - status = ( - service.exec_statuses.popleft() if service.exec_statuses else 200 - ) - if record["state"] in ("FAILED", "TERMINATED"): - status = 409 - if status == 409: - self.reply(409, {"detail": "container not started"}) - return - service.executed.append(body["command"]) - if status >= 400: - self.reply( - status, {"detail": "response lost after command execution"} - ) - return - if body["command"] == ["fail"]: - self.reply( - 200, {"output": "partial output", "error": "command failed"} - ) - else: - self.reply(200, {"output": "completed", "error": None}) - elif segments[3] == "logs": - service.log_started.set() - service.allow_logs.wait(5) - status = service.log_statuses.popleft() if service.log_statuses else 200 - if status != 200: - self.reply(status, {"detail": "logs not ready"}) - return - self.send_response(200) - self.send_header("Content-Type", "text/event-stream") - self.send_header("Connection", "close") - self.end_headers() - self.close_connection = True - try: - # Byte-at-a-time writes ensure parser state is exercised, - # regardless of how the client batches socket reads. - for byte in service.log_bytes: - self.wfile.write(bytes([byte])) - self.wfile.flush() - while service.hold_logs and not service.stopping.wait(0.01): - self.wfile.write(b": heartbeat\n\n") - self.wfile.flush() - except (BrokenPipeError, ConnectionResetError): - pass - finally: - service.log_disconnected.set() - - do_GET = handle_request - do_POST = handle_request - do_DELETE = handle_request - - server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) - thread = threading.Thread(target=server.serve_forever, daemon=True) + finally: + service.log_disconnected.set() + return response + return web.json_response(record) + + async def serve(): + service.loop = asyncio.get_running_loop() + app = web.Application() + for path in ( + "/v2/sandboxes", + "/v2/sandboxes/{id}", + "/v2/sandboxes/{id}/{operation}", + ): + app.router.add_route("*", path, handle) + runner = web.AppRunner(app, access_log=None, shutdown_timeout=1) + await runner.setup() + try: + await web.TCPSite(runner, "127.0.0.1", 0).start() + service.options = { + "api_key": "sandbox-test-key", + "base_url": f"http://127.0.0.1:{runner.addresses[0][1]}", + "request_timeout": 2, + "startup_timeout": 2, + } + ready.set() + await stopped.wait() + finally: + await runner.cleanup() + + thread = threading.Thread(target=lambda: asyncio.run(serve()), daemon=True) thread.start() - service.options = { - "api_key": "sandbox-test-key", - "base_url": f"http://127.0.0.1:{server.server_port}", - "request_timeout": 2, - "startup_timeout": 2, - } + assert ready.wait(5), "local sandbox peer did not start" try: yield service finally: service.allow_create.set() + service.allow_logs.set() service.stopping.set() - server.shutdown() - server.server_close() + service.loop.call_soon_threadsafe(stopped.set) thread.join() From a495fbe94e2f2cb4c008bf9f35be4d922f111ace Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Mon, 14 Sep 2026 14:02:37 -0700 Subject: [PATCH 3/6] test: remove closed public export inventory --- tests/test_init.py | 24 ------------------------ 1 file changed, 24 deletions(-) diff --git a/tests/test_init.py b/tests/test_init.py index 91242328..694ec289 100644 --- a/tests/test_init.py +++ b/tests/test_init.py @@ -86,30 +86,6 @@ def test_private_imports_not_exported(self): for private_symbol in private_symbols: assert private_symbol not in all_symbols, f"Private symbol '{private_symbol}' should not be in __all__" - def test_all_covers_expected_public_api(self): - """Test that __all__ contains the expected public API symbols.""" - expected_symbols = { - # API functions - 'create_container_registry_auth', 'create_endpoint', 'create_pod', 'create_template', - 'delete_container_registry_auth', 'get_endpoints', 'get_gpu', 'get_gpus', - 'get_pod', 'get_pods', 'get_user', 'resume_pod', 'stop_pod', 'terminate_pod', - 'update_container_registry_auth', 'update_endpoint_template', 'update_user_settings', - # Config functions - 'check_credentials', 'get_credentials', 'set_credentials', - # Endpoint classes - 'AsyncioEndpoint', 'AsyncioJob', 'Endpoint', - # Serverless module - 'serverless', - # Logger class - 'RunPodLogger', - # Version - '__version__', - # Module variables - 'SSH_KEY_PATH', 'profile', 'api_key', 'endpoint_url_base' - } - - actual_symbols = set(runpod.__all__) - assert expected_symbols == actual_symbols, f"Expected {expected_symbols}, got {actual_symbols}" def test_no_duplicate_symbols_in_all(self): """Test that __all__ contains no duplicate symbols.""" From 73ffd7b3f42d84683a59e8570b2ee75063796e0c Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Tue, 15 Sep 2026 11:16:00 -0700 Subject: [PATCH 4/6] docs: describe sandbox lifecycle and streaming --- README.md | 36 ++++++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/README.md b/README.md index 11b6d9b3..1fd18ab7 100644 --- a/README.md +++ b/README.md @@ -23,6 +23,7 @@ Welcome to the official Python library for Runpod API & SDK. - [Quick Start](#quick-start) - [Local Test Worker](#local-test-worker) - [📚 | REST API v2 Wrapper](#--rest-api-v2-wrapper) + - [Sandboxes](#sandboxes) - [Endpoints](#endpoints) - [GPU Cloud (Pods)](#gpu-cloud-pods) - [📁 | Directory](#--directory) @@ -172,6 +173,41 @@ import runpod runpod.api_key = "your_runpod_api_key_found_under_settings" ``` +### Sandboxes + +`Sandbox` and `AsyncioSandbox` manage isolated CPU sandboxes through REST API v2. Set `RUNPOD_API_KEY`, assign `runpod.api_key`, or pass `api_key` to a handle. Supply exactly one of `image_name` or `template_id`. + +```python +from runpod import Sandbox + +with Sandbox(image_name="python:3.12-slim") as sandbox: + result = sandbox.exec(["python", "-c", "print('hello from a sandbox')"], check=True) + print(result.output) +``` + +Use `AsyncioSandbox` in asynchronous applications: + +```python +import asyncio +from runpod import AsyncioSandbox + +async def main(): + async with AsyncioSandbox(image_name="python:3.12-slim") as sandbox: + result = await sandbox.exec( + ["python", "-c", "print('hello from a sandbox')"], + check=True, + ) + print(result.output) + +asyncio.run(main()) +``` + +These contexts create a sandbox on entry and terminate it on exit, including when the body raises. `Sandbox.create(...)` and `await AsyncioSandbox.create(...)` return owned handles for explicit lifetime management. Call `terminate()` to release remote compute; `close()` releases local connections only. Handles returned by `get(sandbox_id)` or `list(state=..., labels=...)` are borrowed, so their contexts only close local resources. Async factories and lifecycle methods are awaited. + +Handle properties such as `state`, `compute`, and `expires_at` read cached metadata; `refresh()` fetches a current snapshot. Command execution accepts an argument sequence. `check=True` raises `SandboxExecutionError` with the partial output available in `error.result`. Explicit startup rejections are retried within `startup_timeout`; ambiguous transport failures are not replayed. + +`sandbox.logs(source="container", tail=10)` streams typed log events from the container's main process; `source="system"` selects lifecycle logs. Command output is returned by `exec`, not this stream. Use `with` and regular iteration for synchronous log streams, or `async with` and `async for` for asynchronous streams, to close the connection when stopping early. Preserve an event's `id` and pass it as `last_event_id` to resume a stream, or filter by `since`. + ### Endpoints You can interact with Runpod endpoints via a `run` or `run_sync` method. From 8d5f8a501e9990db74457fdef30ce6010aa2c1d0 Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Tue, 15 Sep 2026 11:26:35 -0700 Subject: [PATCH 5/6] fix: preserve sandbox errors through cleanup --- runpod/sandbox/asyncio.py | 32 +++++++++++++++++---------- runpod/sandbox/sync.py | 6 ++--- tests/test_sandbox.py | 46 +++++++++++++++++++++++++-------------- 3 files changed, 53 insertions(+), 31 deletions(-) diff --git a/runpod/sandbox/asyncio.py b/runpod/sandbox/asyncio.py index 17a6b860..f45349a5 100644 --- a/runpod/sandbox/asyncio.py +++ b/runpod/sandbox/asyncio.py @@ -38,7 +38,7 @@ async def _cleanup( timeout: float, original: Optional[BaseException] = None, ) -> None: - """Finish bounded cleanup despite cancellation, keeping the original primary.""" + """finish bounded cleanup; callers propagate an existing error on success.""" task = asyncio.create_task(asyncio.wait_for(operation, timeout)) interrupted = original while not task.done(): @@ -48,6 +48,7 @@ async def _cleanup( if interrupted is None: interrupted = error except BaseException: + # task.result() below propagates the failure with its cleanup context. break try: task.result() @@ -55,7 +56,7 @@ async def _cleanup( if interrupted is not None: raise interrupted from error raise - if interrupted is not None: + if interrupted is not None and original is None: raise interrupted @@ -149,6 +150,7 @@ async def get( await sandbox.refresh() except BaseException as error: await _cleanup(sandbox.close(), request_timeout, error) + raise return sandbox @classmethod @@ -171,6 +173,7 @@ async def list( snapshots = [SandboxInfo.from_dict(row) for row in rows] except BaseException as error: await _cleanup(api.close(), request_timeout, error) + raise else: await _cleanup(api.close(), request_timeout) handles = [] @@ -240,6 +243,7 @@ async def recover_and_release() -> None: else: task.result() except BaseException as error: + # finish releasing the resource before propagating recovery errors. if error is not original: recovery_error = error try: @@ -255,6 +259,7 @@ async def recover_and_release() -> None: raise recovery_error await _cleanup(recover_and_release(), 2 * self._request_timeout, original) + raise finally: self._creating = False @@ -407,6 +412,7 @@ async def terminate(self) -> None: self._info = replace(self._info, state="TERMINATED", compute=None) except BaseException as error: await _cleanup(self.close(), self._request_timeout, error) + raise else: await _cleanup(self.close(), self._request_timeout) @@ -417,6 +423,7 @@ async def close(self) -> None: try: await stream.aclose() except BaseException as stream_error: + # close every stream before propagating cancellation or another failure. if error is None: error = stream_error try: @@ -461,6 +468,7 @@ async def _open_once(self) -> None: except BaseException as error: self._stream = None await _cleanup(stream.aclose(), self._sandbox._request_timeout, error) + raise async def open(self) -> "AsyncSandboxLogStream": """Complete the HTTP handshake without waiting for the first log event.""" @@ -480,6 +488,7 @@ async def open(self) -> "AsyncSandboxLogStream": self._opened = True except BaseException as error: await _cleanup(self.aclose(), self._sandbox._request_timeout, error) + raise return self def __aiter__(self) -> "AsyncSandboxLogStream": @@ -496,6 +505,7 @@ async def __anext__(self) -> LogEvent: raise except BaseException as error: await _cleanup(self.aclose(), self._sandbox._request_timeout, error) + raise async def aclose(self) -> None: """Release the response and detach from the sandbox; idempotent.""" @@ -503,16 +513,14 @@ async def aclose(self) -> None: return self._closed = True self._sandbox._streams.discard(self) - if self._opening is not None and not self._opening.done(): - self._opening.cancel() - try: - await self._opening - except BaseException: - # The open caller receives its failure; close only owns disposal. - pass - stream, self._stream = self._stream, None - if stream is not None: - await stream.aclose() + try: + if self._opening is not None and not self._opening.done(): + self._opening.cancel() + await asyncio.gather(self._opening, return_exceptions=True) + finally: + stream, self._stream = self._stream, None + if stream is not None: + await stream.aclose() async def __aenter__(self) -> "AsyncSandboxLogStream": return await self.open() diff --git a/runpod/sandbox/sync.py b/runpod/sandbox/sync.py index 9658eeb1..73dbb856 100644 --- a/runpod/sandbox/sync.py +++ b/runpod/sandbox/sync.py @@ -6,7 +6,7 @@ from datetime import datetime from typing import Any, Coroutine, Iterator, Mapping, Optional, Sequence, TypeVar -from .asyncio import AsyncioSandbox +from runpod.sandbox.asyncio import AsyncioSandbox from .models import ( ExecResult, LogEvent, @@ -73,8 +73,7 @@ async def invoke() -> tuple[Optional[_T], Optional[BaseException]]: try: return await coroutine, None except BaseException as error: - # KeyboardInterrupt/SystemExit must reach the calling thread, - # not escape Task execution and stop the background loop. + # interrupts belong to the calling thread, not the background loop. return None, error def complete(done: asyncio.Task) -> None: @@ -85,6 +84,7 @@ def complete(done: asyncio.Task) -> None: else: future.set_result(result) except BaseException as error: + # the waiting caller receives the task's cancellation or failure. future.set_exception(error) finally: if not invoked: diff --git a/tests/test_sandbox.py b/tests/test_sandbox.py index 88ba9d10..286e99f4 100644 --- a/tests/test_sandbox.py +++ b/tests/test_sandbox.py @@ -11,7 +11,7 @@ from aiohttp import web from runpod import AsyncioSandbox, Sandbox -from runpod.error import QueryError +from runpod.error import AuthenticationError, QueryError from runpod.sandbox import ( SandboxExecutionError, SandboxStartupTimeout, @@ -152,6 +152,7 @@ async def handle(request): await response.write(b": heartbeat\n\n") await asyncio.sleep(0.01) except ConnectionResetError: + # clients may stop reading before the fixture finishes streaming. pass finally: service.log_disconnected.set() @@ -285,16 +286,19 @@ def test_sync_body_exception_remains_primary_during_cleanup( ): peer.delete_status = delete_status failure = failure_type("application failed") - with pytest.raises(failure_type) as raised: + try: with Sandbox(image_name="python:3.12-slim", **peer.options): raise failure - assert raised.value is failure - if delete_status == 204: - assert peer.records["sandbox-1"]["state"] == "TERMINATED" + except failure_type as caught: + assert caught is failure + if delete_status == 204: + assert peer.records["sandbox-1"]["state"] == "TERMINATED" + else: + assert isinstance(caught.__cause__, QueryError) + assert caught.__cause__.status_code == 503 + assert peer.records["sandbox-1"]["state"] == "RUNNING" else: - assert isinstance(raised.value.__cause__, QueryError) - assert raised.value.__cause__.status_code == 503 - assert peer.records["sandbox-1"]["state"] == "RUNNING" + pytest.fail("sandbox context suppressed the application failure") @pytest.mark.asyncio @@ -302,15 +306,18 @@ def test_sync_body_exception_remains_primary_during_cleanup( async def test_async_body_exception_remains_primary_during_cleanup(peer, delete_status): peer.delete_status = delete_status failure = ValueError("application failed") - with pytest.raises(ValueError) as raised: + try: async with AsyncioSandbox(image_name="python:3.12-slim", **peer.options): raise failure - assert raised.value is failure - if delete_status == 204: - assert peer.records["sandbox-1"]["state"] == "TERMINATED" + except ValueError as caught: + assert caught is failure + if delete_status == 204: + assert peer.records["sandbox-1"]["state"] == "TERMINATED" + else: + assert isinstance(caught.__cause__, QueryError) + assert caught.__cause__.status_code == 503 else: - assert isinstance(raised.value.__cause__, QueryError) - assert raised.value.__cause__.status_code == 503 + pytest.fail("sandbox context suppressed the application failure") @pytest.mark.asyncio @@ -386,8 +393,8 @@ async def test_closing_sandbox_cancels_pending_log_handshake(peer): try: assert await asyncio.to_thread(peer.log_started.wait, 2) await asyncio.wait_for(sandbox.close(), 3) - with pytest.raises(asyncio.CancelledError): - await opening + await asyncio.gather(opening, return_exceptions=True) + assert opening.cancelled() assert peer.records[sandbox.id]["state"] == "RUNNING" finally: peer.allow_logs.set() @@ -414,3 +421,10 @@ async def test_list_filters_and_independent_handles_do_not_delete_resources(peer assert (await borrowed.exec(["work"])).output == "completed" assert borrowed.id == second["id"] assert all(record["state"] == "RUNNING" for record in peer.records.values()) + + +@pytest.mark.asyncio +async def test_list_authentication_failure_propagates_after_cleanup(peer): + options = {**peer.options, "api_key": "invalid-key"} + with pytest.raises(AuthenticationError): + await AsyncioSandbox.list(**options) From 7211c6a324dbdf7dad9bfd3aa91a6087801da294 Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Tue, 15 Sep 2026 11:29:14 -0700 Subject: [PATCH 6/6] test: assert sandbox failures after context exit --- tests/test_sandbox.py | 34 +++++++++++++++++----------------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/tests/test_sandbox.py b/tests/test_sandbox.py index 286e99f4..8b347512 100644 --- a/tests/test_sandbox.py +++ b/tests/test_sandbox.py @@ -286,19 +286,19 @@ def test_sync_body_exception_remains_primary_during_cleanup( ): peer.delete_status = delete_status failure = failure_type("application failed") + caught = None try: with Sandbox(image_name="python:3.12-slim", **peer.options): raise failure - except failure_type as caught: - assert caught is failure - if delete_status == 204: - assert peer.records["sandbox-1"]["state"] == "TERMINATED" - else: - assert isinstance(caught.__cause__, QueryError) - assert caught.__cause__.status_code == 503 - assert peer.records["sandbox-1"]["state"] == "RUNNING" + except failure_type as error: + caught = error + assert caught is failure + if delete_status == 204: + assert peer.records["sandbox-1"]["state"] == "TERMINATED" else: - pytest.fail("sandbox context suppressed the application failure") + assert isinstance(caught.__cause__, QueryError) + assert caught.__cause__.status_code == 503 + assert peer.records["sandbox-1"]["state"] == "RUNNING" @pytest.mark.asyncio @@ -306,18 +306,18 @@ def test_sync_body_exception_remains_primary_during_cleanup( async def test_async_body_exception_remains_primary_during_cleanup(peer, delete_status): peer.delete_status = delete_status failure = ValueError("application failed") + caught = None try: async with AsyncioSandbox(image_name="python:3.12-slim", **peer.options): raise failure - except ValueError as caught: - assert caught is failure - if delete_status == 204: - assert peer.records["sandbox-1"]["state"] == "TERMINATED" - else: - assert isinstance(caught.__cause__, QueryError) - assert caught.__cause__.status_code == 503 + except ValueError as error: + caught = error + assert caught is failure + if delete_status == 204: + assert peer.records["sandbox-1"]["state"] == "TERMINATED" else: - pytest.fail("sandbox context suppressed the application failure") + assert isinstance(caught.__cause__, QueryError) + assert caught.__cause__.status_code == 503 @pytest.mark.asyncio