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. 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..f45349a5 --- /dev/null +++ b/runpod/sandbox/asyncio.py @@ -0,0 +1,534 @@ +"""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; callers propagate an existing error on success.""" + 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: + # task.result() below propagates the failure with its cleanup context. + break + try: + task.result() + except BaseException as error: + if interrupted is not None: + raise interrupted from error + raise + if interrupted is not None and original is 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") + self._initialize(api_key, base_url, request_timeout, startup_timeout) + self._create_body = { + "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, + } + + 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, **options: Any) -> "AsyncioSandbox": + """Create and return an owned handle, without claiming container readiness.""" + sandbox = cls(**options) + 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) + raise + 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: + # finish releasing the resource before propagating recovery errors. + 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) + raise + 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) + raise + 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: + # close every stream before propagating cancellation or another failure. + 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) + raise + + 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) + raise + 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) + 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() + + 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..73dbb856 --- /dev/null +++ b/runpod/sandbox/sync.py @@ -0,0 +1,402 @@ +"""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 runpod.sandbox.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: + # interrupts belong to the calling thread, not 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: + # the waiting caller receives the task's cancellation or failure. + 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, **options: Any) -> "Sandbox": + """Create immediately; the returned snapshot may still be CREATING.""" + sandbox = cls(**options) + try: + sandbox._runner.run(sandbox._sandbox._create()) + except BaseException: + sandbox.close() + raise + return sandbox + + @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_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.""" diff --git a/tests/test_sandbox.py b/tests/test_sandbox.py new file mode 100644 index 00000000..8b347512 --- /dev/null +++ b/tests/test_sandbox.py @@ -0,0 +1,430 @@ +"""Sandbox lifecycle and streaming regressions over an actual HTTP connection.""" + +import asyncio +import subprocess +import sys +import threading +from collections import deque +from contextlib import contextmanager + +import pytest +from aiohttp import web + +from runpod import AsyncioSandbox, Sandbox +from runpod.error import AuthenticationError, 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() + 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} + ) + 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: + 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: + # clients may stop reading before the fixture finishes streaming. + pass + 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() + 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() + service.loop.call_soon_threadsafe(stopped.set) + 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") + caught = None + try: + with Sandbox(image_name="python:3.12-slim", **peer.options): + raise failure + except failure_type as error: + caught = error + 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" + + +@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") + caught = None + try: + async with AsyncioSandbox(image_name="python:3.12-slim", **peer.options): + raise failure + except ValueError as error: + caught = error + 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 + + +@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) + await asyncio.gather(opening, return_exceptions=True) + assert opening.cancelled() + 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()) + + +@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)