diff --git a/baseten/sandbox/__init__.py b/baseten/sandbox/__init__.py new file mode 100644 index 0000000..0b9a618 --- /dev/null +++ b/baseten/sandbox/__init__.py @@ -0,0 +1,56 @@ +"""Baseten sandbox SDK. + +Use :class:`SandboxClient` or :class:`AsyncSandboxClient` to create, find, +and delete sandboxes, and the :class:`Sandbox` or :class:`AsyncSandbox` they +hand out to work in one. +""" + +from baseten.sandbox._auth import AsyncSandboxTokenProvider, SandboxTokenProvider +from baseten.sandbox._client import ( + AsyncSandboxClient, + SandboxClient, + SandboxClientOptions, +) +from baseten.sandbox._errors import SandboxApiError, SandboxGatewayError +from baseten.sandbox._info import ( + LibraryImage, + LibraryImagePort, + LibraryImageVolume, + SandboxEnvValue, + SandboxInfo, + SandboxProcessInfo, + SandboxProcessLogs, +) +from baseten.sandbox._retry import SandboxRetryOptions +from baseten.sandbox._sandbox import ( + AsyncSandbox, + AsyncSandboxFileSystem, + AsyncSandboxProcess, + Sandbox, + SandboxFileSystem, + SandboxProcess, +) + +__all__ = [ + "AsyncSandbox", + "AsyncSandboxClient", + "AsyncSandboxFileSystem", + "AsyncSandboxProcess", + "AsyncSandboxTokenProvider", + "LibraryImage", + "LibraryImagePort", + "LibraryImageVolume", + "Sandbox", + "SandboxApiError", + "SandboxClient", + "SandboxClientOptions", + "SandboxEnvValue", + "SandboxFileSystem", + "SandboxGatewayError", + "SandboxInfo", + "SandboxProcess", + "SandboxProcessInfo", + "SandboxProcessLogs", + "SandboxRetryOptions", + "SandboxTokenProvider", +] diff --git a/baseten/sandbox/_auth.py b/baseten/sandbox/_auth.py new file mode 100644 index 0000000..b24e0d5 --- /dev/null +++ b/baseten/sandbox/_auth.py @@ -0,0 +1,289 @@ +from __future__ import annotations + +import asyncio +import threading +import time +from collections.abc import Awaitable, Callable +from datetime import UTC, datetime, timedelta +from typing import Protocol + +import httpx + +# Refresh before the stated expiry: clock skew, plus time in flight. +_TOKEN_EXPIRY_LEEWAY = timedelta(seconds=60) + +# A few retries cover a new token invalidated in the same event as the old +# one; beyond that, rejection points at something a new token cannot fix. +_TOKEN_INVALIDATION_MAX_RETRIES = 2 + +# A revocation rejects every token issued before its cutoff, the event's time +# rounded up to the next whole second. A token minted right after the event +# can fall before the cutoff and be rejected too, so after the first resend, +# each one waits this long to be minted past it. A module variable so tests +# can shorten it. +_TOKEN_REVOKED_RETRY_DELAY_SECONDS = 1.0 + +# Both planes send this header on revocation, which can precede the stated +# expiry. The request was untouched, so re-sending with a new token is safe +# even when it has side effects. +_TOKEN_REVOKED_CODE = "TOKEN_REVOKED" + + +class SandboxTokenProvider(Protocol): + """Returns the bearer token to send on a request. Called for every request.""" + + def __call__(self) -> str: ... + + +class AsyncSandboxTokenProvider(Protocol): + """Returns the bearer token to send on a request, or an awaitable of one.""" + + def __call__(self) -> str | Awaitable[str]: ... + + +def resolve_http2_flag(http2: bool | None) -> bool: + """Resolve the HTTP/2 option to a concrete flag for httpx. + + ``None`` enables HTTP/2 when the optional ``h2`` package is importable; + ``True`` requires it and fails otherwise; ``False`` never uses it. + """ + if http2 is False: + return False + try: + import h2 # ty: ignore[unresolved-import] # noqa: F401 + except ImportError: + if http2 is True: + raise ValueError( + "http2=True requires the h2 package; install baseten with the" + " http2 extra (pip install baseten[http2])" + ) + return False + return True + + +class _TokenCache: + """A cached minted token, guarded so concurrent callers share one mint. + + The lock spans the miss check and the mint, so a burst on a cold cache + sends a single exchange. + """ + + def __init__(self) -> None: + self._lock = threading.Lock() + self._token: str | None = None + self._expires_at: datetime | None = None + + def get_or_mint( + self, mint: Callable[[], tuple[str, datetime]] + ) -> tuple[str, datetime]: + with self._lock: + if ( + self._token is not None + and self._expires_at is not None + and datetime.now(UTC) < self._expires_at - _TOKEN_EXPIRY_LEEWAY + ): + return self._token, self._expires_at + self._token, self._expires_at = mint() + return self._token, self._expires_at + + def drop(self, token: str) -> None: + """Drop a token the server rejected, so the next request gets a new one. + + Compared by value, so a rejection of an older token never discards a + newer one another request already minted. + """ + with self._lock: + if self._token == token: + self._token = None + self._expires_at = None + + +class SyncTokenSource: + """Where each request's bearer token comes from: a caller's token provider, + or a token minted from an API key, cached until shortly before it expires. + + The API key is exchanged for a token because the sandbox APIs do not + accept API keys directly. Should that change, only this class changes. + """ + + def __init__( + self, + *, + api_key: str, + token_provider: SandboxTokenProvider | None, + mint: Callable[[], tuple[str, datetime]] | None, + ) -> None: + if token_provider is not None and api_key != "": + raise ValueError("api_key must be empty when token_provider is set") + self._provider = token_provider + self._mint = mint + self._cache = _TokenCache() + + def token(self) -> str | None: + """The token for a request, or None to send no Authorization. + + None is the advanced opt-out of an empty API key with no provider. + """ + if self._provider is not None: + return self._provider() + if self._mint is None: + return None + token, _ = self._cache.get_or_mint(self._mint) + return token + + def invalidate(self, token: str) -> None: + self._cache.drop(token) + + +class AsyncTokenSource: + """Async variant of :class:`SyncTokenSource`. + + The token provider may return a string or an awaitable of one. The + asyncio lock makes concurrent callers share one mint. + """ + + def __init__( + self, + *, + api_key: str, + token_provider: SandboxTokenProvider | AsyncSandboxTokenProvider | None, + mint: Callable[[], Awaitable[tuple[str, datetime]]] | None, + ) -> None: + if token_provider is not None and api_key != "": + raise ValueError("api_key must be empty when token_provider is set") + self._provider = token_provider + self._mint = mint + self._lock = asyncio.Lock() + self._token: str | None = None + self._expires_at: datetime | None = None + + async def token(self) -> str | None: + if self._provider is not None: + provided = self._provider() + if isinstance(provided, str): + return provided + return await provided + if self._mint is None: + return None + async with self._lock: + if ( + self._token is not None + and self._expires_at is not None + and datetime.now(UTC) < self._expires_at - _TOKEN_EXPIRY_LEEWAY + ): + return self._token + token, expires_at = await self._mint() + self._token = token + self._expires_at = expires_at + return token + + async def invalidate(self, token: str) -> None: + async with self._lock: + if self._token == token: + self._token = None + self._expires_at = None + + +class _SyncAuthTransport(httpx.BaseTransport): + """Adds a bearer token from the source to every request. + + A request whose token was revoked is re-sent with a fresh one, up to + :data:`_TOKEN_INVALIDATION_MAX_RETRIES` times. Only a body already in + memory is re-sent: one still backed by a stream is sent once, since + replaying it would mean materializing it here. + """ + + def __init__( + self, inner: httpx.BaseTransport, token_source: SyncTokenSource + ) -> None: + self._inner = inner + self._token_source = token_source + + def handle_request(self, request: httpx.Request) -> httpx.Response: + # A body still backed by a stream, such as a raw multipart upload, + # cannot be replayed, and materializing it here is the memory + # exhaustion to avoid: it is sent once, without revocation resend. + try: + _ = request.content + replayable = True + except httpx.RequestNotRead: + replayable = False + for attempt in range(_TOKEN_INVALIDATION_MAX_RETRIES + 1): + token = self._token_source.token() + if token is not None: + request.headers["Authorization"] = f"Bearer {token}" + response = self._inner.handle_request(request) + revoked = ( + token is not None + and response.status_code == 401 + and response.headers.get("x-blaxel-error-code") == _TOKEN_REVOKED_CODE + ) + if revoked: + # A revoked token never stays cached, even when the request + # cannot be re-sent or the attempts are exhausted. + self._token_source.invalidate(token) + # The last attempt's response is returned even when still + # revoked, so the generated client raises a meaningful error. + if ( + not revoked + or not replayable + or attempt == _TOKEN_INVALIDATION_MAX_RETRIES + ): + return response + response.close() + if attempt > 0: + time.sleep(_TOKEN_REVOKED_RETRY_DELAY_SECONDS) + raise AssertionError # pragma: no cover - loop returns first + + def close(self) -> None: + self._inner.close() + + +class _AsyncAuthTransport(httpx.AsyncBaseTransport): + """Async variant of :class:`_SyncAuthTransport`.""" + + def __init__( + self, inner: httpx.AsyncBaseTransport, token_source: AsyncTokenSource + ) -> None: + self._inner = inner + self._token_source = token_source + + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + # A body still backed by an async stream, such as a raw multipart + # upload, cannot be replayed, and materializing it here is the + # memory exhaustion to avoid: it is sent once, without revocation + # resend. + try: + _ = request.content + replayable = True + except httpx.RequestNotRead: + replayable = False + for attempt in range(_TOKEN_INVALIDATION_MAX_RETRIES + 1): + token = await self._token_source.token() + if token is not None: + request.headers["Authorization"] = f"Bearer {token}" + response = await self._inner.handle_async_request(request) + revoked = ( + token is not None + and response.status_code == 401 + and response.headers.get("x-blaxel-error-code") == _TOKEN_REVOKED_CODE + ) + if revoked: + # A revoked token never stays cached, even when the request + # cannot be re-sent or the attempts are exhausted. + await self._token_source.invalidate(token) + # The last attempt's response is returned even when still + # revoked, so the generated client raises a meaningful error. + if ( + not revoked + or not replayable + or attempt == _TOKEN_INVALIDATION_MAX_RETRIES + ): + return response + await response.aclose() + if attempt > 0: + await asyncio.sleep(_TOKEN_REVOKED_RETRY_DELAY_SECONDS) + raise AssertionError # pragma: no cover - loop returns first + + async def aclose(self) -> None: + await self._inner.aclose() diff --git a/baseten/sandbox/_client.py b/baseten/sandbox/_client.py new file mode 100644 index 0000000..8ee349b --- /dev/null +++ b/baseten/sandbox/_client.py @@ -0,0 +1,751 @@ +from __future__ import annotations + +from collections.abc import AsyncIterator, Iterator, Mapping, Sequence +from dataclasses import dataclass +from datetime import datetime +from typing import Any, Self + +import httpx + +import baseten.client.managementapi +from baseten.client._user_agent import with_user_agent +from baseten.sandbox._auth import ( + AsyncSandboxTokenProvider, + AsyncTokenSource, + SandboxTokenProvider, + SyncTokenSource, + _AsyncAuthTransport, + _SyncAuthTransport, + resolve_http2_flag, +) +from baseten.sandbox._errors import to_sandbox_api_error +from baseten.sandbox._info import ( + LibraryImage, + SandboxEnvValue, + SandboxInfo, + library_image_from_api, + sandbox_envs_to_api, + sandbox_info_from_api, +) +from baseten.sandbox._retry import SandboxRetryOptions +from baseten.sandbox._sandbox import AsyncSandbox, Sandbox + +DEFAULT_MANAGEMENT_BASE_URL = "https://api.baseten.co" + +# Only connect time is bounded: a sandbox call runs as long as the operation +# needs. Callers pass a stricter `timeout` when wanted. +DEFAULT_TIMEOUT = httpx.Timeout(None, connect=10.0) + +# One token exchange is bounded on its own, so a stalled mint is not what +# every caller waits on. A module variable so tests can shorten it. +_TOKEN_MINT_TIMEOUT_SECONDS = 30.0 + + +@dataclass(frozen=True) +class SandboxClientOptions: + """Options for :class:`SandboxClient` and :class:`AsyncSandboxClient`. + + Obtain via :attr:`SandboxClient.options` to inspect the values a client + was constructed with. + """ + + api_key: str + """API key for authentication. Empty, together with *token_provider*, is + the advanced opt-out of API key authentication.""" + + token_provider: SandboxTokenProvider | AsyncSandboxTokenProvider | None = None + """Returns the bearer token for each request instead of one minted from + the API key, which must then be empty.""" + + team_id: str | None = None + """Team to act in. Unset uses the caller's only team, and is an error for + callers in more than one team.""" + + base_url_override: str | None = None + """Explicit management API base URL, or None to use the default.""" + + sandboxes_base_url_override: str | None = None + """Where /v1/sandboxes is served, when that differs from the management + API. The token exchange still goes to the management API.""" + + headers: Mapping[str, str] | None = None + """Additional headers to send on every request.""" + + http2: bool | None = None + """Use HTTP/2 when the optional h2 package allows. None enables it when + h2 is importable, True requires it, False never uses it.""" + + timeout: httpx.Timeout | None = None + """HTTP timeouts. Defaults to bounding only the connect time.""" + + retries: SandboxRetryOptions | None = None + """Retry budgets for transient failures.""" + + transport_override: httpx.BaseTransport | httpx.AsyncBaseTransport | None = None + """HTTP transport to send every request through, both planes. For tests + with a mock transport.""" + + +def _management_base_url(options: SandboxClientOptions) -> str: + return options.base_url_override or DEFAULT_MANAGEMENT_BASE_URL + + +def _request_headers( + options: SandboxClientOptions, authorization: str | None +) -> dict[str, str]: + headers: dict[str, str] = {**(options.headers or {})} + if authorization is not None: + headers["Authorization"] = authorization + return with_user_agent(headers) + + +def _sandbox_create_request( + *, + name: str | None, + image: str | None, + memory: int | None, + region: str | None, + envs: list[baseten.client.managementapi.SandboxEnv] | None, + labels: dict[str, str] | None, + external_id: str | None, + create_if_not_exists: bool, +) -> baseten.client.managementapi.CreateSandboxRequest: + """Build the create body from only the fields the caller set. + + The generated client keeps explicitly passed None values, since null can + mean "clear" on updates. Creation must send nothing instead, so the + server-side defaults apply. + """ + create_fields: dict[str, Any] = {} + if name is not None: + create_fields["name"] = name + if image is not None: + create_fields["image"] = image + if memory is not None: + create_fields["memory"] = memory + if region is not None: + create_fields["region"] = region + if envs is not None: + create_fields["envs"] = envs + if labels is not None: + create_fields["labels"] = labels + if external_id is not None: + create_fields["external_id"] = external_id + if create_if_not_exists: + create_fields["create_if_not_exists"] = True + return baseten.client.managementapi.CreateSandboxRequest(**create_fields) + + +class SandboxClient: + """Synchronous client for Baseten sandboxes. + + Creates, finds, and deletes sandboxes, and hands out :class:`Sandbox` + objects to work in one. Closing the client closes the sandboxes' shared + connections. Can be used as a context manager. + """ + + def __init__( + self, + *, + api_key: str, + token_provider: SandboxTokenProvider | None = None, + team_id: str | None = None, + base_url_override: str | None = None, + sandboxes_base_url_override: str | None = None, + headers: Mapping[str, str] | None = None, + http2: bool | None = None, + timeout: httpx.Timeout | None = None, + retries: SandboxRetryOptions | None = None, + transport_override: httpx.BaseTransport | None = None, + ) -> None: + """Create a new synchronous sandbox client. + + Every argument matches the :class:`SandboxClientOptions` field of the + same name. + """ + self._options = SandboxClientOptions( + api_key=api_key, + token_provider=token_provider, + team_id=team_id, + base_url_override=base_url_override, + sandboxes_base_url_override=sandboxes_base_url_override, + headers=headers, + http2=http2, + timeout=timeout, + retries=retries, + transport_override=transport_override, + ) + if transport_override is None: + self._pool: httpx.BaseTransport = httpx.HTTPTransport( + http2=resolve_http2_flag(self._options.http2) + ) + elif isinstance(transport_override, httpx.BaseTransport): + # MockTransport implements both interfaces; only an async-only + # transport fails this isinstance. + self._pool = transport_override + else: + raise TypeError( + "SandboxClient needs a sync transport in transport_override" + ) + self._timeout = self._options.timeout or DEFAULT_TIMEOUT + # API-key auth, not the auth transport: minting must not need a token. + self._mint_http_client = httpx.Client( + transport=self._pool, + base_url=_management_base_url(self._options), + headers=_request_headers( + self._options, + f"Bearer {api_key}" if api_key != "" else None, + ), + timeout=httpx.Timeout(_TOKEN_MINT_TIMEOUT_SECONDS), + ) + self._token_source = SyncTokenSource( + api_key=api_key, + token_provider=token_provider, + mint=self._mint if api_key != "" else None, + ) + self._auth_transport = _SyncAuthTransport(self._pool, self._token_source) + self._http_client = httpx.Client( + transport=self._auth_transport, + base_url=_management_base_url(self._options), + headers=_request_headers(self._options, None), + timeout=self._timeout, + ) + self._api = baseten.client.managementapi.ApiClient(self._http_client) + self._sandbox_api = self._api + if self._options.sandboxes_base_url_override is not None: + self._sandboxes_http_client = httpx.Client( + transport=self._auth_transport, + base_url=self._options.sandboxes_base_url_override, + headers=_request_headers(self._options, None), + timeout=self._timeout, + ) + self._sandbox_api = baseten.client.managementapi.ApiClient( + self._sandboxes_http_client + ) + self._sandbox_http_clients: list[httpx.Client] = [] + + def _mint(self) -> tuple[str, datetime]: + # The mint client carries the API key, so minting cannot recurse + # through the auth transport. + minted = baseten.client.managementapi.ApiClient( + self._mint_http_client + ).post_token( + request=baseten.client.managementapi.CreateTokenRequest( + scopes=[baseten.client.managementapi.TokenScope.sandboxes] + ) + ) + return minted.token, minted.expires_at + + @property + def options(self) -> SandboxClientOptions: + """The options this client was constructed with.""" + return self._options + + @property + def api(self) -> baseten.client.managementapi.ApiClient: + """The generated management API client, authenticated for sandboxes. + + For operations this SDK does not wrap. The generated API surface is + not covered by stability guarantees and may change between versions. + """ + return self._api + + @property + def sandbox_api(self) -> baseten.client.managementapi.ApiClient: + """The generated client bound to where /v1/sandboxes is served. + + With no override set this is :attr:`api`; with one set, raw sandbox + operations go to the override host while :attr:`api` stays on the + management domain. The generated API surface is not covered by + stability guarantees and may change between versions. + """ + return self._sandbox_api + + def create( + self, + *, + name: str | None = None, + image: str | None = None, + memory: int | None = None, + region: str | None = None, + envs: dict[str, SandboxEnvValue] | None = None, + labels: dict[str, str] | None = None, + external_id: str | None = None, + create_if_not_exists: bool = False, + ) -> Sandbox: + """Create a sandbox. + + Defaults for image and memory live server-side, so nothing is sent + for what the caller leaves unset. + """ + try: + sandbox = self._sandbox_api.create_sandbox( + params=baseten.client.managementapi.CreateSandboxParams( + team_id=self._options.team_id + ), + request=_sandbox_create_request( + name=name, + image=image, + memory=memory, + region=region, + envs=sandbox_envs_to_api(envs) if envs is not None else None, + labels=labels, + external_id=external_id, + create_if_not_exists=create_if_not_exists, + ), + ) + except Exception as error: + raise to_sandbox_api_error(error, "control") from error + info = sandbox_info_from_api(sandbox) + if info.url is None: + raise ValueError(f"sandbox {info.name} was created but has no URL yet") + return self._sandbox_for(info.name, info.url, info) + + def get_info(self, name: str, *, show_secrets: bool = False) -> SandboxInfo: + """Get a sandbox's current record. + + ``show_secrets`` returns environment variable values unmasked; + it needs the workspace administrator role, and callers without it + still see masked values. + """ + try: + sandbox = self._sandbox_api.get_sandbox( + sandbox_name=name, + params=baseten.client.managementapi.GetSandboxParams( + team_id=self._options.team_id, show_secrets=show_secrets + ), + ) + except Exception as error: + raise to_sandbox_api_error(error, "control") from error + return sandbox_info_from_api(sandbox) + + def get(self, sandbox: str | SandboxInfo) -> Sandbox: + """Get a :class:`Sandbox` to work in. + + Fetches the record when given a name; none is fetched when the + record is already in hand. + """ + info = self.get_info(sandbox) if isinstance(sandbox, str) else sandbox + if info.url is None: + raise ValueError(f"sandbox {info.name} has no URL yet") + return self._sandbox_for(info.name, info.url) + + def sandbox_from_url(self, url: str) -> Sandbox: + """Get a :class:`Sandbox` for a sandbox whose URL is already known. + + Makes no call. The sandbox shares this client's authentication, + connections, and retry budgets, and closes with it. + """ + return self._sandbox_for("", url) + + def list( + self, + *, + query: str | None = None, + statuses: Sequence[str] | None = None, + external_id: str | None = None, + page_size: int | None = None, + ) -> Iterator[SandboxInfo]: + """List sandboxes, fetching further pages as iteration reaches them.""" + seen_cursors: set[str] = set() + cursor: str | None = None + while True: + try: + response = self._sandbox_api.list_sandboxes( + params=baseten.client.managementapi.ListSandboxesParams( + team_id=self._options.team_id, + cursor=cursor, + q=query, + status=[ + baseten.client.managementapi.StatusItem(root=status) + for status in statuses + ] + if statuses is not None + else None, + external_id=external_id, + **({"limit": page_size} if page_size is not None else {}), + ) + ) + except Exception as error: + raise to_sandbox_api_error(error, "control") from error + yield from (sandbox_info_from_api(sb) for sb in response.items) + cursor = ( + response.pagination.cursor if response.pagination.has_more else None + ) + if cursor is None: + return + if cursor in seen_cursors: + raise ValueError("sandbox list returned a repeated cursor") + seen_cursors.add(cursor) + + def delete(self, name: str) -> SandboxInfo: + """Delete a sandbox. + + Deletion continues after this returns, so the returned record usually + shows the sandbox as still deleting. + """ + try: + sandbox = self._sandbox_api.delete_sandbox( + sandbox_name=name, + params=baseten.client.managementapi.DeleteSandboxParams( + team_id=self._options.team_id + ), + ) + except Exception as error: + raise to_sandbox_api_error(error, "control") from error + return sandbox_info_from_api(sandbox) + + def library_images(self) -> Sequence[LibraryImage]: + """List the platform's starter images from the starter-image library. + + The server drops hidden and coming-soon entries. + """ + try: + catalog = self._sandbox_api.list_sandbox_library_images( + params=baseten.client.managementapi.ListSandboxLibraryImagesParams( + team_id=self._options.team_id + ) + ) + except Exception as error: + raise to_sandbox_api_error(error, "control") from error + return [library_image_from_api(image) for image in catalog.items] + + def close(self) -> None: + """Close the client and the connections every sandbox shares.""" + for http_client in self._sandbox_http_clients: + http_client.close() + if self._options.sandboxes_base_url_override is not None: + self._sandboxes_http_client.close() + self._http_client.close() + self._mint_http_client.close() + + def __enter__(self) -> Self: + return self + + def __exit__(self, *args: object) -> None: + self.close() + + def _sandbox_for( + self, name: str, url: str, creation_info: SandboxInfo | None = None + ) -> Sandbox: + http_client = httpx.Client( + transport=self._auth_transport, + base_url=url, + headers=_request_headers(self._options, None), + timeout=self._timeout, + ) + self._sandbox_http_clients.append(http_client) + return Sandbox._shared( + name=name, + url=url, + http_client=http_client, + retries=self._options.retries, + creation_info=creation_info, + ) + + +class AsyncSandboxClient: + """Asynchronous client for Baseten sandboxes. + + Creates, finds, and deletes sandboxes, and hands out + :class:`AsyncSandbox` objects to work in one. Closing the client closes + the sandboxes' shared connections. Can be used as an async context + manager. + """ + + def __init__( + self, + *, + api_key: str, + token_provider: SandboxTokenProvider | AsyncSandboxTokenProvider | None = None, + team_id: str | None = None, + base_url_override: str | None = None, + sandboxes_base_url_override: str | None = None, + headers: Mapping[str, str] | None = None, + http2: bool | None = None, + timeout: httpx.Timeout | None = None, + retries: SandboxRetryOptions | None = None, + transport_override: httpx.AsyncBaseTransport | None = None, + ) -> None: + """Create a new asynchronous sandbox client. + + The token provider may return a string or an awaitable of one. Every + other argument matches the :class:`SandboxClientOptions` field of + the same name. + """ + self._options = SandboxClientOptions( + api_key=api_key, + token_provider=token_provider, + team_id=team_id, + base_url_override=base_url_override, + sandboxes_base_url_override=sandboxes_base_url_override, + headers=headers, + http2=http2, + timeout=timeout, + retries=retries, + transport_override=transport_override, + ) + if transport_override is None: + self._pool: httpx.AsyncBaseTransport = httpx.AsyncHTTPTransport( + http2=resolve_http2_flag(self._options.http2) + ) + elif isinstance(transport_override, httpx.AsyncBaseTransport): + self._pool = transport_override + else: + raise TypeError( + "AsyncSandboxClient needs an async transport in transport_override" + ) + self._timeout = self._options.timeout or DEFAULT_TIMEOUT + self._mint_http_client = httpx.AsyncClient( + transport=self._pool, + timeout=httpx.Timeout(_TOKEN_MINT_TIMEOUT_SECONDS), + base_url=_management_base_url(self._options), + headers=_request_headers( + self._options, + f"Bearer {api_key}" if api_key != "" else None, + ), + ) + + async def mint() -> tuple[str, datetime]: + minted = await baseten.client.managementapi.AsyncApiClient( + self._mint_http_client + ).post_token( + request=baseten.client.managementapi.CreateTokenRequest( + scopes=[baseten.client.managementapi.TokenScope.sandboxes] + ) + ) + return minted.token, minted.expires_at + + self._token_source = AsyncTokenSource( + api_key=api_key, + token_provider=self._options.token_provider, # type: ignore[arg-type] + mint=mint if api_key != "" else None, + ) + self._auth_transport = _AsyncAuthTransport(self._pool, self._token_source) + self._http_client = httpx.AsyncClient( + transport=self._auth_transport, + base_url=_management_base_url(self._options), + headers=_request_headers(self._options, None), + timeout=self._timeout, + ) + self._api = baseten.client.managementapi.AsyncApiClient(self._http_client) + self._sandbox_api = self._api + if self._options.sandboxes_base_url_override is not None: + self._sandboxes_http_client = httpx.AsyncClient( + transport=self._auth_transport, + base_url=self._options.sandboxes_base_url_override, + headers=_request_headers(self._options, None), + timeout=self._timeout, + ) + self._sandbox_api = baseten.client.managementapi.AsyncApiClient( + self._sandboxes_http_client + ) + self._sandbox_http_clients: list[httpx.AsyncClient] = [] + + @property + def options(self) -> SandboxClientOptions: + """The options this client was constructed with.""" + return self._options + + @property + def api(self) -> baseten.client.managementapi.AsyncApiClient: + """The generated management API client, authenticated for sandboxes. + + For operations this SDK does not wrap. The generated API surface is + not covered by stability guarantees and may change between versions. + """ + return self._api + + @property + def sandbox_api(self) -> baseten.client.managementapi.AsyncApiClient: + """The generated client bound to where /v1/sandboxes is served. + + With no override set this is :attr:`api`; with one set, raw sandbox + operations go to the override host while :attr:`api` stays on the + management domain. The generated API surface is not covered by + stability guarantees and may change between versions. + """ + return self._sandbox_api + + async def create( + self, + *, + name: str | None = None, + image: str | None = None, + memory: int | None = None, + region: str | None = None, + envs: dict[str, SandboxEnvValue] | None = None, + labels: dict[str, str] | None = None, + external_id: str | None = None, + create_if_not_exists: bool = False, + ) -> AsyncSandbox: + """Create a sandbox. + + Defaults for image and memory live server-side, so nothing is sent + for what the caller leaves unset. + """ + try: + sandbox = await self._sandbox_api.create_sandbox( + params=baseten.client.managementapi.CreateSandboxParams( + team_id=self._options.team_id + ), + request=_sandbox_create_request( + name=name, + image=image, + memory=memory, + region=region, + envs=sandbox_envs_to_api(envs) if envs is not None else None, + labels=labels, + external_id=external_id, + create_if_not_exists=create_if_not_exists, + ), + ) + except Exception as error: + raise to_sandbox_api_error(error, "control") from error + info = sandbox_info_from_api(sandbox) + if info.url is None: + raise ValueError(f"sandbox {info.name} was created but has no URL yet") + return self._sandbox_for(info.name, info.url, info) + + async def get_info(self, name: str, *, show_secrets: bool = False) -> SandboxInfo: + """Get a sandbox's current record. + + ``show_secrets`` returns environment variable values unmasked; + it needs the workspace administrator role, and callers without it + still see masked values. + """ + try: + sandbox = await self._sandbox_api.get_sandbox( + sandbox_name=name, + params=baseten.client.managementapi.GetSandboxParams( + team_id=self._options.team_id, show_secrets=show_secrets + ), + ) + except Exception as error: + raise to_sandbox_api_error(error, "control") from error + return sandbox_info_from_api(sandbox) + + async def get(self, sandbox: str | SandboxInfo) -> AsyncSandbox: + """Get an :class:`AsyncSandbox` to work in. + + Fetches the record when given a name; none is fetched when the + record is already in hand. + """ + info = await self.get_info(sandbox) if isinstance(sandbox, str) else sandbox + if info.url is None: + raise ValueError(f"sandbox {info.name} has no URL yet") + return self._sandbox_for(info.name, info.url) + + def sandbox_from_url(self, url: str) -> AsyncSandbox: + """Get an :class:`AsyncSandbox` for a sandbox whose URL is already known. + + Makes no call. The sandbox shares this client's authentication, + connections, and retry budgets, and closes with it. + """ + return self._sandbox_for("", url) + + async def list( + self, + *, + query: str | None = None, + statuses: Sequence[str] | None = None, + external_id: str | None = None, + page_size: int | None = None, + ) -> AsyncIterator[SandboxInfo]: + """List sandboxes, fetching further pages as iteration reaches them.""" + seen_cursors: set[str] = set() + cursor: str | None = None + while True: + try: + response = await self._sandbox_api.list_sandboxes( + params=baseten.client.managementapi.ListSandboxesParams( + team_id=self._options.team_id, + cursor=cursor, + q=query, + status=[ + baseten.client.managementapi.StatusItem(root=status) + for status in statuses + ] + if statuses is not None + else None, + external_id=external_id, + **({"limit": page_size} if page_size is not None else {}), + ) + ) + except Exception as error: + raise to_sandbox_api_error(error, "control") from error + for sb in response.items: + yield sandbox_info_from_api(sb) + cursor = ( + response.pagination.cursor if response.pagination.has_more else None + ) + if cursor is None: + return + if cursor in seen_cursors: + raise ValueError("sandbox list returned a repeated cursor") + seen_cursors.add(cursor) + + async def delete(self, name: str) -> SandboxInfo: + """Delete a sandbox. + + Deletion continues after this returns, so the returned record usually + shows the sandbox as still deleting. + """ + try: + sandbox = await self._sandbox_api.delete_sandbox( + sandbox_name=name, + params=baseten.client.managementapi.DeleteSandboxParams( + team_id=self._options.team_id + ), + ) + except Exception as error: + raise to_sandbox_api_error(error, "control") from error + return sandbox_info_from_api(sandbox) + + async def library_images(self) -> Sequence[LibraryImage]: + """List the platform's starter images from the starter-image library. + + The server drops hidden and coming-soon entries. + """ + try: + catalog = await self._sandbox_api.list_sandbox_library_images( + params=baseten.client.managementapi.ListSandboxLibraryImagesParams( + team_id=self._options.team_id + ) + ) + except Exception as error: + raise to_sandbox_api_error(error, "control") from error + return [library_image_from_api(image) for image in catalog.items] + + async def close(self) -> None: + """Close the client and the connections every sandbox shares.""" + for http_client in self._sandbox_http_clients: + await http_client.aclose() + if self._options.sandboxes_base_url_override is not None: + await self._sandboxes_http_client.aclose() + await self._http_client.aclose() + await self._mint_http_client.aclose() + + async def __aenter__(self) -> Self: + return self + + async def __aexit__(self, *args: object) -> None: + await self.close() + + def _sandbox_for( + self, name: str, url: str, creation_info: SandboxInfo | None = None + ) -> AsyncSandbox: + http_client = httpx.AsyncClient( + transport=self._auth_transport, + base_url=url, + headers=_request_headers(self._options, None), + timeout=self._timeout, + ) + self._sandbox_http_clients.append(http_client) + return AsyncSandbox._shared( + name=name, + url=url, + http_client=http_client, + retries=self._options.retries, + creation_info=creation_info, + ) diff --git a/baseten/sandbox/_errors.py b/baseten/sandbox/_errors.py new file mode 100644 index 0000000..9e918e2 --- /dev/null +++ b/baseten/sandbox/_errors.py @@ -0,0 +1,93 @@ +from __future__ import annotations + +import json + +from pydantic import BaseModel + +import baseten.client.managementapi +import baseten.client.sandboxapi + + +class SandboxApiError(Exception): + """An error response from the sandbox control plane or from a sandbox. + + The response body is parsed as JSON when possible, otherwise kept as + text. *code* holds the machine-readable error code when the response has + one; the control plane puts it in ``code`` or ``error``, a sandbox in + ``error``. + """ + + def __init__(self, status: int, body: object) -> None: + fields = body if isinstance(body, dict) else {} + body_message = fields.get("message") + body_error = fields.get("error") + body_code = fields.get("code") + code = ( + body_code + if isinstance(body_code, str) + else (body_error if isinstance(body_error, str) else None) + ) + description = None + if isinstance(body_message, str): + description = body_message + elif isinstance(body_error, str): + description = body_error + message = f"sandbox API error (HTTP {status})" + if isinstance(code, str): + message += f" {code}" + if description is not None: + message += f": {description}" + super().__init__(message) + self.status = status + self.code = code + self.details = fields.get("details") + self.body = body + + +class SandboxGatewayError(SandboxApiError): + """A 502, 503, or 504 from the edge in front of a sandbox. + + Not from the sandbox itself: usually transient, for example while a + sandbox wakes from standby or when a request outlasts the edge's timeout. + """ + + +_GATEWAY_STATUSES = frozenset({502, 503, 504}) + + +def to_sandbox_api_error(error: object, plane: str) -> Exception: + """Convert an error raised by a generated client, or return it unchanged. + + ``plane`` is ``"control"`` or ``"exec"``. Only a sandbox sits behind the + edge, so gateway statuses from the control plane stay ordinary errors. + """ + status = getattr(error, "status_code", None) + if not isinstance(status, int): + return error if isinstance(error, Exception) else Exception(str(error)) + if isinstance( + error, + ( + baseten.client.managementapi.ResponseError, + baseten.client.sandboxapi.ResponseError, + ), + ): + try: + body: object = json.loads(error.body) + except (json.JSONDecodeError, TypeError): + body = error.body + else: + # Generated typed errors are dataclasses of status_code plus the + # error's parsed model, keyed by the schema name. + model_fields = { + name: value for name, value in vars(error).items() if name != "status_code" + } + if len(model_fields) == 1: + body = next(iter(model_fields.values())) + else: # pragma: no cover - generated errors carry exactly one model + body = model_fields + if isinstance(body, BaseModel): + # The curated error reads code and details off a plain dict. + body = body.model_dump() + if plane == "exec" and status in _GATEWAY_STATUSES: + return SandboxGatewayError(status, body) + return SandboxApiError(status, body) diff --git a/baseten/sandbox/_info.py b/baseten/sandbox/_info.py new file mode 100644 index 0000000..2273eb8 --- /dev/null +++ b/baseten/sandbox/_info.py @@ -0,0 +1,305 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from datetime import datetime +from email.utils import parsedate_to_datetime + +import baseten.client.managementapi +import baseten.client.sandboxapi + +# Deployment status of a sandbox. Only DEPLOYED sandboxes run commands. +# Kept open: the server may add values, so do not treat this as exhaustive. +SANDBOX_STATUSES = ( + "DEPLOYING", + "DEPLOYED", + "FAILED", + "DEACTIVATING", + "DEACTIVATED", + "DELETING", + "TERMINATED", + "ARCHIVING", + "ARCHIVED", + "UNARCHIVING", + "BUILDING", + "UPLOADING", +) + +# Status of a process in a sandbox. Kept open. +SANDBOX_PROCESS_STATUSES = ("running", "completed", "failed", "killed", "stopped") + +SandboxStatus = str +SandboxProcessStatus = str + + +@dataclass +class SandboxEnvValue: + """Value of an environment variable in a sandbox.""" + + value: str + # The server treats every env value as secret by default, so the safe + # default here is secret too; False opts out and unmasks the value in + # responses. + secret: bool = True + + +@dataclass +class SandboxInfo: + """A sandbox as last reported by the control plane.""" + + name: str + """Unique name of the sandbox, assigned by the server when not given.""" + + status: SandboxStatus + enabled: bool + created_at: datetime + envs: dict[str, SandboxEnvValue] = field(default_factory=dict) + labels: dict[str, str] = field(default_factory=dict) + + url: str | None = None + """Base URL of the sandbox's execution API, once it has one.""" + + image: str | None = None + """Image reference, including its tag.""" + + memory: int | None = None + """Memory in megabytes, which also sets the CPU allocation.""" + + region: str | None = None + external_id: str | None = None + """Caller-owned identifier for external lookups.""" + + updated_at: datetime | None = None + created_by: str | None = None + updated_by: str | None = None + last_used_at: datetime | None = None + expires_in_seconds: int | None = None + """Seconds left before automatic deletion, when expiration is configured.""" + + +@dataclass +class SandboxProcessInfo: + """A process in a sandbox.""" + + pid: str + name: str + command: str + status: SandboxProcessStatus + exit_code: int + stdout: str + stderr: str + logs: str + """Standard output and standard error, interleaved.""" + + working_dir: str + started_at: datetime + completed_at: datetime | None = None + + +@dataclass +class SandboxProcessLogs: + """A process's captured output.""" + + stdout: str + stderr: str + logs: str + """Standard output and standard error, interleaved.""" + + +@dataclass +class LibraryImagePort: + """One port a starter image's services listen on.""" + + name: str | None = None + target: int | None = None + """Port number inside the sandbox.""" + + protocol: str | None = None + + +@dataclass +class LibraryImage: + """One starter image from the platform's starter-image library. + + Available to any sandbox without building or pushing; pass + :attr:`image` as the image when creating a sandbox. + """ + + name: str + """Stable identifier of the starter image.""" + + image: str + """Image reference, including its tag.""" + + display_name: str | None = None + description: str | None = None + long_description: str | None = None + memory: int | None = None + """Default memory allocation in megabytes.""" + + categories: list[str] = field(default_factory=list) + tags: list[str] = field(default_factory=list) + ports: list[LibraryImagePort] = field(default_factory=list) + icon_url: str | None = None + project_url: str | None = None + """Project page for the image's stack.""" + + enterprise: bool = False + """Whether the image is gated to enterprise workspaces.""" + + creation_extra_args: dict[str, str] = field(default_factory=dict) + """Kernel selection arguments suggested when creating a sandbox.""" + + creation_volumes: list[LibraryImageVolume] = field(default_factory=list) + """Volume attachments suggested when creating a sandbox.""" + + +@dataclass +class LibraryImageVolume: + """A volume attachment a starter image suggests.""" + + name: str | None = None + """Volume name, or an internal identifier for an ephemeral volume.""" + + mount_path: str | None = None + """Absolute filesystem path where the volume is mounted.""" + + type: str | None = None + """Volume type, empty for a persistent volume.""" + + size_mb: int | None = None + """Storage capacity in megabytes for an ephemeral volume.""" + + read_only: bool = False + + +def library_image_from_api( + image: baseten.client.managementapi.SandboxLibraryImage, +) -> LibraryImage: + """Convert the control plane's starter-image catalog entry.""" + return LibraryImage( + name=image.name, + image=image.image, + display_name=image.display_name, + description=image.description, + long_description=image.long_description, + memory=image.memory, + categories=image.categories or [], + tags=image.tags or [], + ports=[ + LibraryImagePort( + name=port.name, + target=port.target, + protocol=str(port.protocol) if port.protocol else None, + ) + for port in (image.ports.root if image.ports else []) + ], + icon_url=image.icon, + project_url=image.url, + enterprise=bool(image.enterprise), + creation_extra_args=dict(image.creation_options.extra_args or {}) + if image.creation_options + else {}, + creation_volumes=[ + LibraryImageVolume( + name=volume.name, + mount_path=volume.mount_path, + type=volume.type, + size_mb=volume.size_mb, + read_only=bool(volume.read_only), + ) + for volume in (image.creation_options.volumes or []) + ] + if image.creation_options + else [], + ) + + +def sandbox_info_from_api(sandbox: baseten.client.managementapi.Sandbox) -> SandboxInfo: + """Convert the control plane's sandbox record.""" + return SandboxInfo( + name=sandbox.name, + url=str(sandbox.url) if sandbox.url is not None else None, + status=str(sandbox.status), + image=sandbox.image, + memory=sandbox.memory, + region=sandbox.region, + enabled=sandbox.enabled, + envs=_envs_from_api(sandbox.envs), + labels=dict(sandbox.labels.root) if sandbox.labels is not None else {}, + external_id=sandbox.external_id, + created_at=sandbox.created_at, + updated_at=sandbox.updated_at, + created_by=sandbox.created_by, + updated_by=sandbox.updated_by, + last_used_at=sandbox.last_used_at, + expires_in_seconds=sandbox.expires_in, + ) + + +def sandbox_envs_to_api( + envs: dict[str, SandboxEnvValue], +) -> list[baseten.client.managementapi.SandboxEnv]: + """Convert environment variables to the control plane's form.""" + return [ + baseten.client.managementapi.SandboxEnv( + name=name, value=value.value, secret=value.secret + ) + for name, value in envs.items() + ] + + +def process_info_from_api( + process: baseten.client.sandboxapi.ProcessResponse, +) -> SandboxProcessInfo: + """Convert the execution plane's process record. + + The plane reports timestamps as strings in either ISO 8601 or RFC 1123 + form, and an empty completion time while the process runs. + """ + return SandboxProcessInfo( + pid=process.pid, + name=process.name, + command=process.command, + status=str(process.status), + exit_code=process.exitCode, + stdout=process.stdout, + stderr=process.stderr, + logs=process.logs, + working_dir=process.workingDir, + started_at=parse_exec_timestamp(process.startedAt), + completed_at=( + parse_exec_timestamp(process.completedAt) if process.completedAt else None + ), + ) + + +def process_logs_from_api( + logs: baseten.client.sandboxapi.ProcessLogs, +) -> SandboxProcessLogs: + """Convert the execution plane's captured process output.""" + return SandboxProcessLogs(stdout=logs.stdout, stderr=logs.stderr, logs=logs.logs) + + +def parse_exec_timestamp(value: str) -> datetime: + """Parse an execution-plane timestamp, ISO 8601 or RFC 1123.""" + try: + return datetime.fromisoformat(value) + except ValueError: + pass + parsed = parsedate_to_datetime(value) + if parsed is None: + raise ValueError(f"unparseable execution-plane timestamp {value!r}") from None + return parsed + + +def _envs_from_api( + envs: list[baseten.client.managementapi.SandboxEnv] | None, +) -> dict[str, SandboxEnvValue]: + result: dict[str, SandboxEnvValue] = {} + for env in envs or []: + if env.name is None: + continue + result[env.name] = SandboxEnvValue( + value=env.value or "", secret=bool(env.secret) + ) + return result diff --git a/baseten/sandbox/_retry.py b/baseten/sandbox/_retry.py new file mode 100644 index 0000000..86db63d --- /dev/null +++ b/baseten/sandbox/_retry.py @@ -0,0 +1,128 @@ +from __future__ import annotations + +import asyncio +import random +import time +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from typing import TypeVar + +import httpx + +from baseten.sandbox._errors import SandboxGatewayError + +_T = TypeVar("_T") + +DEFAULT_READ_MAX_RETRIES = 5 +DEFAULT_GATEWAY_MAX_RETRIES = 2 +DEFAULT_UPLOAD_MAX_RETRIES = 3 + +_BACKOFF_BASE_SECONDS = 0.2 +_BACKOFF_MAX_SECONDS = 2.0 + +# httpx exceptions meaning the connection dropped before a complete response +# arrived. Only then is repeating an idempotent call safe. +_TRANSIENT_ERRORS = ( + httpx.ConnectError, + httpx.ConnectTimeout, + httpx.ReadError, + httpx.WriteError, + httpx.RemoteProtocolError, +) + + +@dataclass(frozen=True) +class SandboxRetryOptions: + """Retry budgets for transient failures. + + Retries only ever apply to operations that are safe to repeat, such as + reads. Operations with side effects, such as running a process or + creating a sandbox, are never retried. Each budget counts retries after + the first attempt, so 0 disables that kind of retry. + """ + + read_max_retries: int = DEFAULT_READ_MAX_RETRIES + """Retries of a read after a dropped or reset connection.""" + + gateway_max_retries: int = DEFAULT_GATEWAY_MAX_RETRIES + """Retries after the edge in front of a sandbox answers 502, 503, or 504. + + Each attempt can take as long as the edge's own timeout, so this is kept + small. + """ + + upload_max_retries: int = DEFAULT_UPLOAD_MAX_RETRIES + """Retries of a file upload after a dropped or reset connection.""" + + +def is_transient_reset_error(error: object) -> bool: + """Whether the error is a dropped or reset connection, not a response. + + A response from the server, even an error, means the request was handled; + only a transport failure makes repeating an idempotent call safe. + """ + return isinstance(error, _TRANSIENT_ERRORS) + + +def backoff_delay_seconds(attempt: int) -> float: + """Exponential backoff capped at a maximum, plus jitter. + + The jitter is up to one base delay, so many clients failing together do + not retry together. + """ + capped = min(_BACKOFF_BASE_SECONDS * 2 ** (attempt - 1), _BACKOFF_MAX_SECONDS) + return capped + random.uniform(0, _BACKOFF_BASE_SECONDS) + + +def retry_idempotent( + call: Callable[[], _T], + *, + max_retries: int, + gateway_max_retries: int, +) -> _T: + """Run an idempotent call, retrying dropped connections and edge errors. + + Must never wrap a call with side effects, since a dropped connection does + not tell whether the server acted on it. + """ + attempt = 0 + gateway_attempt = 0 + while True: + try: + return call() + except SandboxGatewayError: + gateway_attempt += 1 + if gateway_attempt > gateway_max_retries: + raise + delay = backoff_delay_seconds(gateway_attempt) + except Exception as error: + attempt += 1 + if attempt > max_retries or not is_transient_reset_error(error): + raise + delay = backoff_delay_seconds(attempt) + time.sleep(delay) + + +async def retry_idempotent_async( + call: Callable[[], Awaitable[_T]], + *, + max_retries: int, + gateway_max_retries: int, +) -> _T: + """Async variant of :func:`retry_idempotent`.""" + attempt = 0 + gateway_attempt = 0 + while True: + try: + return await call() + except SandboxGatewayError: + gateway_attempt += 1 + if gateway_attempt > gateway_max_retries: + raise + delay = backoff_delay_seconds(gateway_attempt) + except Exception as error: + attempt += 1 + if attempt > max_retries or not is_transient_reset_error(error): + raise + delay = backoff_delay_seconds(attempt) + await asyncio.sleep(delay) diff --git a/baseten/sandbox/_sandbox.py b/baseten/sandbox/_sandbox.py new file mode 100644 index 0000000..e0e3c7c --- /dev/null +++ b/baseten/sandbox/_sandbox.py @@ -0,0 +1,519 @@ +from __future__ import annotations + +from collections.abc import Sequence + +import httpx + +import baseten.client.sandboxapi +from baseten.sandbox._errors import to_sandbox_api_error +from baseten.sandbox._info import ( + SandboxInfo, + SandboxProcessInfo, + SandboxProcessLogs, + process_info_from_api, + process_logs_from_api, +) +from baseten.sandbox._retry import ( + SandboxRetryOptions, + retry_idempotent, + retry_idempotent_async, +) + +_NOT_IMPLEMENTED = ( + "Not implemented yet. Reach the operation through {} in the meantime." +) + + +class Sandbox: + """One sandbox, reached directly at its own URL. + + Get one from :meth:`baseten.sandbox.SandboxClient.create`, + :meth:`baseten.sandbox.SandboxClient.get`, or + :meth:`baseten.sandbox.SandboxClient.sandbox_from_url`. A sandbox + always shares its client's authentication, connections, and retry + budgets, and closes with it. + """ + + _name: str + _url: str + _retries: SandboxRetryOptions + _http_client: httpx.Client + _api: baseten.client.sandboxapi.ApiClient + _fs: SandboxFileSystem | None + _process: SandboxProcess | None + _info: SandboxInfo | None + + def __init__(self) -> None: + # _shared builds instances; direct construction cannot produce a + # usable one, so refuse it rather than fail later on access. + raise TypeError( + "build a Sandbox through SandboxClient.create, .get, or .sandbox_from_url" + ) + + @classmethod + def _shared( + cls, + *, + name: str, + url: str, + http_client: httpx.Client, + retries: SandboxRetryOptions | None, + creation_info: SandboxInfo | None, + ) -> Sandbox: + sandbox = cls.__new__(cls) + sandbox._name = name + sandbox._url = url.rstrip("/") + sandbox._retries = retries or SandboxRetryOptions() + sandbox._http_client = http_client + sandbox._api = baseten.client.sandboxapi.ApiClient(http_client) + sandbox._fs = None + sandbox._process = None + sandbox._info = creation_info + return sandbox + + @property + def name(self) -> str: + """Name of the sandbox.""" + return self._name + + @property + def url(self) -> str: + """Base URL of the sandbox's execution API.""" + return self._url + + @property + def info(self) -> SandboxInfo | None: + """The record as of creation, when this sandbox came from + :meth:`baseten.sandbox.SandboxClient.create`. None otherwise.""" + return self._info + + @property + def api(self) -> baseten.client.sandboxapi.ApiClient: + """The generated client for this sandbox's execution API. + + For operations this SDK does not wrap. The generated API surface is + not covered by stability guarantees and may change between versions. + """ + return self._api + + @property + def fs(self) -> SandboxFileSystem: + """Files and directories in the sandbox.""" + if self._fs is None: + self._fs = SandboxFileSystem(self._api, self._retries) + return self._fs + + @property + def process(self) -> SandboxProcess: + """Processes in the sandbox.""" + if self._process is None: + self._process = SandboxProcess(self._api, self._retries) + return self._process + + @property + def drives(self) -> object: + """Drive mounts in the sandbox. Not implemented yet.""" + raise NotImplementedError(_NOT_IMPLEMENTED.format("sandbox.api")) + + @property + def codegen(self) -> object: + """Code editing helpers. Not implemented yet.""" + raise NotImplementedError(_NOT_IMPLEMENTED.format("sandbox.api")) + + @property + def system(self) -> object: + """Sandbox system operations. Not implemented yet.""" + raise NotImplementedError(_NOT_IMPLEMENTED.format("sandbox.api")) + + @property + def network(self) -> object: + """Network access to ports in the sandbox. Not implemented yet.""" + raise NotImplementedError(_NOT_IMPLEMENTED.format("sandbox.api")) + + +class AsyncSandbox: + """One sandbox, reached directly at its own URL. Async variant. + + Get one from :meth:`baseten.sandbox.AsyncSandboxClient.create`, + :meth:`baseten.sandbox.AsyncSandboxClient.get`, or + :meth:`baseten.sandbox.AsyncSandboxClient.sandbox_from_url`. The sandbox + shares its client's authentication, connections, and retry budgets, and + closes with it. + """ + + _name: str + _url: str + _retries: SandboxRetryOptions + _http_client: httpx.AsyncClient + _api: baseten.client.sandboxapi.AsyncApiClient + _fs: AsyncSandboxFileSystem | None + _process: AsyncSandboxProcess | None + _info: SandboxInfo | None + + def __init__(self) -> None: + # _shared builds instances; direct construction cannot produce a + # usable one, so refuse it rather than fail later on access. + raise TypeError( + "build an AsyncSandbox through AsyncSandboxClient.create, .get, " + "or .sandbox_from_url" + ) + + @classmethod + def _shared( + cls, + *, + name: str, + url: str, + http_client: httpx.AsyncClient, + retries: SandboxRetryOptions | None, + creation_info: SandboxInfo | None, + ) -> AsyncSandbox: + sandbox = cls.__new__(cls) + sandbox._name = name + sandbox._url = url.rstrip("/") + sandbox._retries = retries or SandboxRetryOptions() + sandbox._http_client = http_client + sandbox._api = baseten.client.sandboxapi.AsyncApiClient(http_client) + sandbox._fs = None + sandbox._process = None + sandbox._info = creation_info + return sandbox + + @property + def name(self) -> str: + """Name of the sandbox.""" + return self._name + + @property + def url(self) -> str: + """Base URL of the sandbox's execution API.""" + return self._url + + @property + def info(self) -> SandboxInfo | None: + """The record as of creation, when this sandbox came from + :meth:`baseten.sandbox.AsyncSandboxClient.create`. None otherwise.""" + return self._info + + @property + def api(self) -> baseten.client.sandboxapi.AsyncApiClient: + """The generated client for this sandbox's execution API. + + For operations this SDK does not wrap. The generated API surface is + not covered by stability guarantees and may change between versions. + """ + return self._api + + @property + def fs(self) -> AsyncSandboxFileSystem: + """Files and directories in the sandbox.""" + if self._fs is None: + self._fs = AsyncSandboxFileSystem(self._api, self._retries) + return self._fs + + @property + def process(self) -> AsyncSandboxProcess: + """Processes in the sandbox.""" + if self._process is None: + self._process = AsyncSandboxProcess(self._api, self._retries) + return self._process + + @property + def drives(self) -> object: + """Drive mounts in the sandbox. Not implemented yet.""" + raise NotImplementedError(_NOT_IMPLEMENTED.format("sandbox.api")) + + @property + def codegen(self) -> object: + """Code editing helpers. Not implemented yet.""" + raise NotImplementedError(_NOT_IMPLEMENTED.format("sandbox.api")) + + @property + def system(self) -> object: + """Sandbox system operations. Not implemented yet.""" + raise NotImplementedError(_NOT_IMPLEMENTED.format("sandbox.api")) + + @property + def network(self) -> object: + """Network access to ports in the sandbox. Not implemented yet.""" + raise NotImplementedError(_NOT_IMPLEMENTED.format("sandbox.api")) + + +class SandboxFileSystem: + """Files and directories in a sandbox.""" + + def __init__( + self, api: baseten.client.sandboxapi.ApiClient, retries: SandboxRetryOptions + ) -> None: + self._api = api + self._retries = retries + + def read(self, path: str) -> str: + """Read a text file. + + Retries dropped connections and edge gateway errors, since a read is + safe to repeat. + """ + + def attempt() -> baseten.client.sandboxapi.GetFilesystemResponse: + # Convert per attempt: the retry helper must see + # SandboxGatewayError, not the generated ResponseError. + try: + return self._api.get_filesystem(path=path) + except Exception as error: + raise to_sandbox_api_error(error, "exec") from error + + result = retry_idempotent( + attempt, + max_retries=self._retries.read_max_retries, + gateway_max_retries=self._retries.gateway_max_retries, + ) + root = result.root + if isinstance(root, baseten.client.sandboxapi.FileWithContent): + return root.content + if isinstance(root, bytes): + return root.decode() + raise ValueError(f"{path} is a directory, not a file") + + def write(self, path: str, content: str) -> None: + """Write a text file, creating it or replacing its content. + + A write of the same content is idempotent, so it is retried on the + upload budgets after a dropped connection. + """ + + def attempt() -> None: + try: + self._api.put_filesystem( + path=path, + request=baseten.client.sandboxapi.FileRequest(content=content), + ) + except Exception as error: + raise to_sandbox_api_error(error, "exec") from error + + retry_idempotent( + attempt, + max_retries=self._retries.upload_max_retries, + gateway_max_retries=self._retries.gateway_max_retries, + ) + + +class AsyncSandboxFileSystem: + """Files and directories in a sandbox. Async variant.""" + + def __init__( + self, + api: baseten.client.sandboxapi.AsyncApiClient, + retries: SandboxRetryOptions, + ) -> None: + self._api = api + self._retries = retries + + async def read(self, path: str) -> str: + """Read a text file. + + Retries dropped connections and edge gateway errors, since a read is + safe to repeat. + """ + + async def attempt() -> baseten.client.sandboxapi.GetFilesystemResponse: + # Convert per attempt: the retry helper must see + # SandboxGatewayError, not the generated ResponseError. + try: + return await self._api.get_filesystem(path=path) + except Exception as error: + raise to_sandbox_api_error(error, "exec") from error + + result = await retry_idempotent_async( + attempt, + max_retries=self._retries.read_max_retries, + gateway_max_retries=self._retries.gateway_max_retries, + ) + root = result.root + if isinstance(root, baseten.client.sandboxapi.FileWithContent): + return root.content + if isinstance(root, bytes): + return root.decode() + raise ValueError(f"{path} is a directory, not a file") + + async def write(self, path: str, content: str) -> None: + """Write a text file, creating it or replacing its content. + + A write of the same content is idempotent, so it is retried on the + upload budgets after a dropped connection. + """ + + async def attempt() -> None: + try: + await self._api.put_filesystem( + path=path, + request=baseten.client.sandboxapi.FileRequest(content=content), + ) + except Exception as error: + raise to_sandbox_api_error(error, "exec") from error + + await retry_idempotent_async( + attempt, + max_retries=self._retries.upload_max_retries, + gateway_max_retries=self._retries.gateway_max_retries, + ) + + +class SandboxProcess: + """Processes in a sandbox.""" + + def __init__( + self, + api: baseten.client.sandboxapi.ApiClient, + retries: SandboxRetryOptions, + ) -> None: + self._api = api + self._retries = retries + + def exec( + self, + command: str, + *, + working_dir: str | None = None, + env: dict[str, str] | None = None, + name: str | None = None, + timeout_seconds: int | None = None, + wait_for_completion: bool | None = None, + ) -> SandboxProcessInfo: + """Start a command. + + Never retried, since a command may have side effects even when its + response is lost. + """ + try: + response = self._api.post_process( + request=baseten.client.sandboxapi.ProcessRequest( + command=command, + workingDir=working_dir, + env=env, + name=name, + timeout=timeout_seconds, + waitForCompletion=wait_for_completion, + ) + ) + except Exception as error: + raise to_sandbox_api_error(error, "exec") from error + return process_info_from_api(response) + + def list(self) -> Sequence[SandboxProcessInfo]: + """List the sandbox's processes. + + Retried as an idempotent read, on the configured budgets. + """ + + def attempt() -> baseten.client.sandboxapi.GetProcessResponse: + try: + return self._api.get_process() + except Exception as error: + raise to_sandbox_api_error(error, "exec") from error + + response = retry_idempotent( + attempt, + max_retries=self._retries.read_max_retries, + gateway_max_retries=self._retries.gateway_max_retries, + ) + return [process_info_from_api(process) for process in response.root] + + def logs(self, identifier: str) -> SandboxProcessLogs: + """Get one process's captured output, by pid or name. + + Retried as an idempotent read, on the configured budgets. + """ + + def attempt() -> baseten.client.sandboxapi.ProcessLogs: + try: + return self._api.get_process_logs(identifier=identifier) + except Exception as error: + raise to_sandbox_api_error(error, "exec") from error + + response = retry_idempotent( + attempt, + max_retries=self._retries.read_max_retries, + gateway_max_retries=self._retries.gateway_max_retries, + ) + return process_logs_from_api(response) + + +class AsyncSandboxProcess: + """Processes in a sandbox. Async variant.""" + + def __init__( + self, + api: baseten.client.sandboxapi.AsyncApiClient, + retries: SandboxRetryOptions, + ) -> None: + self._api = api + self._retries = retries + + async def exec( + self, + command: str, + *, + working_dir: str | None = None, + env: dict[str, str] | None = None, + name: str | None = None, + timeout_seconds: int | None = None, + wait_for_completion: bool | None = None, + ) -> SandboxProcessInfo: + """Start a command. + + Never retried, since a command may have side effects even when its + response is lost. + """ + try: + response = await self._api.post_process( + request=baseten.client.sandboxapi.ProcessRequest( + command=command, + workingDir=working_dir, + env=env, + name=name, + timeout=timeout_seconds, + waitForCompletion=wait_for_completion, + ) + ) + except Exception as error: + raise to_sandbox_api_error(error, "exec") from error + return process_info_from_api(response) + + async def list(self) -> Sequence[SandboxProcessInfo]: + """List the sandbox's processes. + + Retried as an idempotent read, on the configured budgets. + """ + + async def attempt() -> baseten.client.sandboxapi.GetProcessResponse: + try: + return await self._api.get_process() + except Exception as error: + raise to_sandbox_api_error(error, "exec") from error + + response = await retry_idempotent_async( + attempt, + max_retries=self._retries.read_max_retries, + gateway_max_retries=self._retries.gateway_max_retries, + ) + return [process_info_from_api(process) for process in response.root] + + async def logs(self, identifier: str) -> SandboxProcessLogs: + """Get one process's captured output, by pid or name. + + Retried as an idempotent read, on the configured budgets. + """ + + async def attempt() -> baseten.client.sandboxapi.ProcessLogs: + try: + return await self._api.get_process_logs(identifier=identifier) + except Exception as error: + raise to_sandbox_api_error(error, "exec") from error + + response = await retry_idempotent_async( + attempt, + max_retries=self._retries.read_max_retries, + gateway_max_retries=self._retries.gateway_max_retries, + ) + return process_logs_from_api(response) diff --git a/scripts/docgen/conf.py b/scripts/docgen/conf.py index 3888bd0..f6153c8 100644 --- a/scripts/docgen/conf.py +++ b/scripts/docgen/conf.py @@ -55,6 +55,12 @@ "baseten.client.inferenceapi": ["ApiClient", "AsyncApiClient"], "baseten.client.sandboxapi": ["ApiClient", "AsyncApiClient"], "baseten.client.modelconfig": ["ModelConfig"], + "baseten.sandbox": [ + "SandboxClient", + "AsyncSandboxClient", + "Sandbox", + "AsyncSandbox", + ], } diff --git a/scripts/docgen/index.rst b/scripts/docgen/index.rst index b4e7482..6e16b07 100644 --- a/scripts/docgen/index.rst +++ b/scripts/docgen/index.rst @@ -11,3 +11,4 @@ the source. For installation, usage, and examples, see the project on GitHub: :caption: API Reference baseten.client + baseten.sandbox diff --git a/tests/sandbox/test_parity.py b/tests/sandbox/test_parity.py new file mode 100644 index 0000000..61a5ceb --- /dev/null +++ b/tests/sandbox/test_parity.py @@ -0,0 +1,151 @@ +"""Parity between the curated surface and the generated models it wraps. + +The surface never exposes generated types: methods take keyword arguments of +primitives and curated dataclasses, and the generated models are constructed +internally. These tests pin that translation, so a spec field added upstream +fails here until the surface either forwards it or lists it under +INTENTIONALLY_NOT_TRANSLATED_*. +""" + +from __future__ import annotations + +import baseten.client.managementapi +import baseten.client.sandboxapi + +# Control-plane fields left to the raw client. +INTENTIONALLY_NOT_TRANSLATED_SANDBOX = {"lifecycle", "network", "ports"} + +INTENTIONALLY_NOT_TRANSLATED_PROCESS = { + "keepAlive", + "maxRestarts", + "restartOnFailure", + "stdin", + "waitForPorts", +} + +# Deliberately matching the JS surface's process record. +INTENTIONALLY_NOT_TRANSLATED_PROCESS_RECORD = { + "keepAlive", + "maxRestarts", + "restartCount", + "restartOnFailure", + "stdin", +} + +# File-write fields left to the raw client. +INTENTIONALLY_NOT_TRANSLATED_FILE = {"isDirectory", "permissions"} + +# Starter-image fields not on the curated record: icon variants, plus the flags driving the library_images filter. +INTENTIONALLY_NOT_TRANSLATED_LIBRARY_IMAGE = { + "icon_light", + "icon_dark", + "hidden", + "coming_soon", +} + + +def test_create_forwards_every_translated_field() -> None: + forwarded = { + "name", + "image", + "memory", + "region", + "envs", + "labels", + "external_id", + "create_if_not_exists", + } + fields = set(baseten.client.managementapi.CreateSandboxRequest.model_fields) + assert forwarded <= fields + assert fields - forwarded == INTENTIONALLY_NOT_TRANSLATED_SANDBOX + + +def test_record_translation_covers_every_translated_field() -> None: + translated = { + "name", + "url", + "status", + "image", + "memory", + "region", + "enabled", + "envs", + "labels", + "external_id", + "created_at", + "updated_at", + "created_by", + "updated_by", + "last_used_at", + "expires_in", + } + fields = set(baseten.client.managementapi.Sandbox.model_fields) + assert translated <= fields + assert fields - translated == INTENTIONALLY_NOT_TRANSLATED_SANDBOX + + +def test_exec_forwards_every_translated_field() -> None: + forwarded = { + "command", + "workingDir", + "env", + "name", + "timeout", + "waitForCompletion", + } + fields = set(baseten.client.sandboxapi.ProcessRequest.model_fields) + assert forwarded <= fields + assert fields - forwarded == INTENTIONALLY_NOT_TRANSLATED_PROCESS + + +def test_process_record_translation_covers_every_translated_field() -> None: + translated = { + "command", + "completedAt", + "exitCode", + "logs", + "name", + "pid", + "startedAt", + "status", + "stderr", + "stdout", + "workingDir", + } + fields = set(baseten.client.sandboxapi.ProcessResponse.model_fields) + assert translated <= fields + assert fields - translated == INTENTIONALLY_NOT_TRANSLATED_PROCESS_RECORD + + +def test_process_logs_translation_covers_every_field() -> None: + translated = {"logs", "stderr", "stdout"} + fields = set(baseten.client.sandboxapi.ProcessLogs.model_fields) + assert fields == translated + + +def test_fs_write_forwards_every_translated_field() -> None: + forwarded = {"content"} + fields = set(baseten.client.sandboxapi.FileRequest.model_fields) + assert forwarded <= fields + assert fields - forwarded == INTENTIONALLY_NOT_TRANSLATED_FILE + + +def test_library_image_translation_covers_every_translated_field() -> None: + translated = { + "name", + "display_name", + "image", + "description", + "long_description", + "memory", + "categories", + "tags", + "ports", + "icon", + "url", + "enterprise", + "creation_options", + } + fields = set(baseten.client.managementapi.SandboxLibraryImage.model_fields) + assert translated <= fields + assert fields - translated == INTENTIONALLY_NOT_TRANSLATED_LIBRARY_IMAGE diff --git a/tests/sandbox/test_sandbox_e2e.py b/tests/sandbox/test_sandbox_e2e.py new file mode 100644 index 0000000..aa85006 --- /dev/null +++ b/tests/sandbox/test_sandbox_e2e.py @@ -0,0 +1,229 @@ +"""End-to-end tests against a live environment. + +Skipped unless all of BASETEN_E2E_TEST_API_KEY, BASETEN_E2E_TEST_DOMAIN, +and BASETEN_E2E_TEST_SANDBOXES_DOMAIN are set. The key must belong to an +organization with sandboxes enabled, and the sandbox management API must +be served on its own host: today no environment serves /v1/sandboxes on +the token-exchange domain, so the split is required rather than optional. +Mirrors the JS suite's environment contract. +""" + +from __future__ import annotations + +import os +import time + +import pytest + +import baseten.client.sandboxapi +from baseten.sandbox import AsyncSandboxClient, SandboxClient + +API_KEY = os.environ.get("BASETEN_E2E_TEST_API_KEY", "") +MANAGEMENT_DOMAIN = os.environ.get("BASETEN_E2E_TEST_DOMAIN", "") +SANDBOXES_DOMAIN = os.environ.get("BASETEN_E2E_TEST_SANDBOXES_DOMAIN", "") + +e2e_required = pytest.mark.skipif( + not (API_KEY and MANAGEMENT_DOMAIN and SANDBOXES_DOMAIN), + reason="sandbox e2e environment not set", +) + +# The one region where the execution plane accepts staging-minted tokens +# (eu-dub-1 rejects them with AUTHENTICATION_FAILED). +E2E_REGION = "us-was-1" + + +def unique_sandbox_name() -> str: + return f"baseten-python-e2e-{time.time():.0f}" + + +def make_client() -> SandboxClient: + return SandboxClient( + api_key=API_KEY, + base_url_override=f"https://{MANAGEMENT_DOMAIN}", + sandboxes_base_url_override=f"https://{SANDBOXES_DOMAIN}", + ) + + +@e2e_required +def test_create_exec_fs_delete() -> None: + name = unique_sandbox_name() + client = make_client() + try: + sandbox = client.create(name=name, region=E2E_REGION) + assert sandbox.info is not None + assert sandbox.info.status == "DEPLOYED" + assert sandbox.info.url is not None + + process = sandbox.process.exec("echo hello", wait_for_completion=True) + assert process.status == "completed" + assert process.exit_code == 0 + assert process.stdout == "hello\n" + + sandbox.fs.write("/tmp/e2e.txt", "python sdk") + assert sandbox.fs.read("/tmp/e2e.txt") == "python sdk" + finally: + deleted = client.delete(name) + assert deleted.status in ("DELETING", "TERMINATED") + client.close() + + +@e2e_required +@pytest.mark.asyncio +async def test_async_create_and_list() -> None: + name = unique_sandbox_name() + client = AsyncSandboxClient( + api_key=API_KEY, + base_url_override=f"https://{MANAGEMENT_DOMAIN}", + sandboxes_base_url_override=f"https://{SANDBOXES_DOMAIN}", + ) + try: + sandbox = await client.create(name=name, region=E2E_REGION) + assert sandbox.name == name + + listed = [info.name async for info in client.list(query=name)] + assert name in listed + finally: + await client.delete(name) + await client.close() + + +@e2e_required +def test_get_info_and_get() -> None: + name = unique_sandbox_name() + client = make_client() + try: + created = client.create(name=name, region=E2E_REGION) + info = client.get_info(name) + assert info.name == name + assert info.status == "DEPLOYED" + # The record keeps the server's trailing slash; Sandbox.url normalizes it. + assert info.url is not None + assert info.url.rstrip("/") == created.url + + sandbox = client.get(name) + assert sandbox.name == name + assert sandbox.url == created.url + finally: + client.delete(name) + client.close() + + +@e2e_required +def test_list_paginates_across_pages() -> None: + prefix = f"baseten-python-e2e-{time.time():.0f}" + names = [f"{prefix}-a", f"{prefix}-b"] + client = make_client() + try: + for sandbox_name in names: + client.create(name=sandbox_name, region=E2E_REGION) + + listed = [info.name for info in client.list(query=prefix, page_size=1)] + assert sorted(listed) == sorted(names) + finally: + for sandbox_name in names: + client.delete(sandbox_name) + client.close() + + +@e2e_required +@pytest.mark.xfail( + reason="server bug: POST /process with waitForCompletion=false sends" + " logs: null, but the spec declares logs required; pydantic rejects the" + " response (JS does not validate, so only Python fails). Tracked with Chad.", + strict=True, +) +def test_logs_stream_and_stdin_via_raw_api() -> None: + name = unique_sandbox_name() + client = make_client() + try: + sandbox = client.create(name=name, region=E2E_REGION) + # `cat` with a stdin pipe, driven through the raw execution client: + # the surface's exec does not expose stdin yet. + sandbox.api.post_process( + request=baseten.client.sandboxapi.ProcessRequest( + command="cat", stdin=True, name="piper", waitForCompletion=False + ) + ) + sandbox.api.post_process_stdin( + identifier="piper", content=b"stdin round trip\n" + ) + # Closing stdin makes cat exit; closing again is the no-op probe. + sandbox.api.delete_process_stdin(identifier="piper") + for _ in range(20): + time.sleep(0.5) + process = sandbox.api.get_process_identifier(identifier="piper") + if process.status == "completed": + assert "stdin round trip" in process.stdout + break + else: + raise AssertionError("cat did not complete after stdin closed") + + logs_response = sandbox.api.get_process_logs_stream(identifier="piper") + logs_response.read() + logs_response.close() + assert "stdin round trip" in logs_response.text + finally: + client.delete(name) + client.close() + + +@e2e_required +def test_multipart_upload_via_raw_api() -> None: + name = unique_sandbox_name() + client = make_client() + try: + sandbox = client.create(name=name, region=E2E_REGION) + remote_path = "/tmp/multipart-e2e.txt" + initiated = sandbox.api.post_filesystem_multipart_initiate( + path=remote_path, + request=baseten.client.sandboxapi.MultipartInitiateRequest(), + ) + upload_id = initiated.uploadId + assert upload_id is not None + + part = sandbox.api.put_filesystem_multipart_part( + upload_id=upload_id, + params=baseten.client.sandboxapi.PutFilesystemMultipartPartParams( + partNumber=1 + ), + files={"file": ("part.bin", b"multipart e2e content")}, + ) + assert part.partNumber == 1 + + completed = sandbox.api.post_filesystem_multipart_complete( + upload_id=upload_id, + request=baseten.client.sandboxapi.MultipartCompleteRequest( + parts=[ + baseten.client.sandboxapi.MultipartPartInfo( + partNumber=part.partNumber, etag=part.etag + ) + ] + ), + ) + assert completed.message + + assert sandbox.fs.read(remote_path) == "multipart e2e content" + finally: + client.delete(name) + client.close() + + +@e2e_required +def test_archive_export_status_dispatch() -> None: + name = unique_sandbox_name() + client = make_client() + try: + sandbox = client.create(name=name, region=E2E_REGION) + sandbox.fs.write("/tmp/archive-e2e.txt", "archive me") + + result = sandbox.api.post_archive_export( + request=baseten.client.sandboxapi.ExportOptions(dryRun=True) + ) + assert isinstance(result, baseten.client.sandboxapi.ExportResult) + assert result.manifest is not None + + # The async (202 ExportProgress) path needs a presigned upload URL, + # which the test cannot mint; it stays covered by unit tests. + finally: + client.delete(name) + client.close() diff --git a/tests/sandbox/test_surface.py b/tests/sandbox/test_surface.py new file mode 100644 index 0000000..d3b532c --- /dev/null +++ b/tests/sandbox/test_surface.py @@ -0,0 +1,815 @@ +from __future__ import annotations + +import json +import types +from typing import Any + +import httpx +import pytest + +import baseten.sandbox._auth +from baseten.sandbox import ( + AsyncSandboxClient, + SandboxApiError, + SandboxClient, + SandboxGatewayError, + SandboxInfo, +) + +MANAGEMENT_URL = "https://api.baseten.co" +SANDBOX_URL = "https://sb-1.b10.co" +FAR_FUTURE = "2100-01-01T00:00:00Z" + + +def sandbox_record(name: str = "sb-1", status: str = "DEPLOYED") -> dict[str, Any]: + return { + "name": name, + "url": SANDBOX_URL, + "status": status, + "state": "RUNNING", + "created_at": "2026-09-29T00:00:00Z", + "enabled": True, + } + + +def process_record() -> dict[str, Any]: + return { + "command": "echo hi", + "completedAt": "Tue, 30 Sep 2026 00:00:01 GMT", + "exitCode": 0, + "logs": "hi\n", + "name": "echo", + "pid": "123", + "startedAt": "2026-09-30T00:00:00Z", + "status": "completed", + "stderr": "", + "stdout": "hi\n", + "workingDir": "/tmp", + } + + +def file_record() -> dict[str, Any]: + return { + "content": "file text", + "group": "root", + "lastModified": "2026-09-29T00:00:00Z", + "name": "file.txt", + "owner": "root", + "path": "/file.txt", + "permissions": "0644", + "size": 9, + } + + +class RoutingTransport: + """Mock transport routing by path, with scripted per-path responses. + + Each route maps path to either a response factory or a list consumed one + entry per matching request. + """ + + def __init__( + self, + routes: dict[str, Any], + mint_tokens: list[str] | None = None, + ) -> None: + self.routes = routes + self.mint_tokens = mint_tokens or ["tok-1"] + self.requests: list[tuple[str, str, str | None, bytes]] = [] + self.mint_count = 0 + + def _respond(self, request: httpx.Request) -> httpx.Response: + method = request.method + path = request.url.path + authorization = request.headers.get("authorization") + self.requests.append((method, path, authorization, request.content)) + + def json_response(payload: Any, status: int = 200) -> httpx.Response: + return httpx.Response( + status, json=payload, headers={"content-type": "application/json"} + ) + + if path == "/v1/token": + token = self.mint_tokens[min(self.mint_count, len(self.mint_tokens) - 1)] + self.mint_count += 1 + return json_response({"token": token, "expires_at": FAR_FUTURE}) + + route = self.routes.get(f"{method} {path}") + if route is None: + route = self.routes.get(path) + if route is None: + return httpx.Response(404, text=f"no route for {method} {path}") + if isinstance(route, list): + if not route: + return httpx.Response(500, text="route exhausted") + route = route.pop(0) + if isinstance(route, Exception): + raise route + if callable(route): + return route(request) + if isinstance(route, httpx.Response): + return route + return json_response(route) + + @property + def sync_transport(self) -> httpx.MockTransport: + return httpx.MockTransport(self._respond) + + @property + def async_transport(self) -> httpx.MockTransport: + async def handler(request: httpx.Request) -> httpx.Response: + return self._respond(request) + + return httpx.MockTransport(handler) + + def authorizations(self, path: str) -> list[str | None]: + return [auth for _, p, auth, _ in self.requests if p == path and auth] + + def count(self, method: str, path: str) -> int: + return sum(1 for m, p, _, _ in self.requests if m == method and p == path) + + +def make_sync_client(transport: RoutingTransport) -> SandboxClient: + return SandboxClient( + api_key="test-key", + base_url_override=MANAGEMENT_URL, + transport_override=transport.sync_transport, + ) + + +def created_response(record: dict[str, Any] | None = None) -> httpx.Response: + return httpx.Response( + 201, + json=record or sandbox_record(), + headers={"content-type": "application/json"}, + ) + + +def revoked_token_response(token: str) -> httpx.Response: + return httpx.Response( + 401, + json={"code": "TOKEN_REVOKED", "message": "revoked"}, + headers={ + "content-type": "application/json", + "x-blaxel-error-code": "TOKEN_REVOKED", + }, + ) + + +def test_create_mints_token_and_returns_sandbox() -> None: + transport = RoutingTransport( + routes={ + "POST /v1/sandboxes/instances": created_response(), + } + ) + client = make_sync_client(transport) + + sandbox = client.create(name="sb-1") + + assert sandbox.name == "sb-1" + assert sandbox.url == SANDBOX_URL + assert isinstance(sandbox.info, SandboxInfo) + assert sandbox.info.status == "DEPLOYED" + body = json.loads(transport.requests[1][3]) + assert body == {"name": "sb-1"} + assert transport.authorizations("/v1/sandboxes/instances") == ["Bearer tok-1"] + client.close() + + +def test_token_is_minted_once_and_cached() -> None: + transport = RoutingTransport( + routes={ + "POST /v1/sandboxes/instances": [created_response(), created_response()], + "GET /v1/sandboxes/instances/sb-1": sandbox_record(), + } + ) + client = make_sync_client(transport) + + client.create() + client.get_info("sb-1") + + assert transport.mint_count == 1 + client.close() + + +def test_empty_api_key_sends_no_authorization() -> None: + transport = RoutingTransport( + routes={"GET /v1/sandboxes/instances/sb-1": sandbox_record()} + ) + client = SandboxClient( + api_key="", + base_url_override=MANAGEMENT_URL, + transport_override=transport.sync_transport, + ) + + client.get_info("sb-1") + + assert transport.authorizations("/v1/sandboxes/instances/sb-1") == [] + client.close() + + +def test_revoked_token_is_replaced_and_request_resent() -> None: + transport = RoutingTransport( + routes={ + "GET /v1/sandboxes/instances/sb-1": lambda request: ( + revoked_token_response("tok-1") + if request.headers.get("authorization") == "Bearer tok-1" + else httpx.Response( + 200, + json=sandbox_record(), + headers={"content-type": "application/json"}, + ) + ), + }, + mint_tokens=["tok-1", "tok-2"], + ) + client = make_sync_client(transport) + + info = client.get_info("sb-1") + + assert info.name == "sb-1" + assert transport.mint_count == 2 + assert transport.authorizations("/v1/sandboxes/instances/sb-1") == [ + "Bearer tok-1", + "Bearer tok-2", + ] + client.close() + + +def test_list_follows_cursor_pages() -> None: + transport = RoutingTransport( + routes={ + "GET /v1/sandboxes/instances": [ + { + "items": [sandbox_record("sb-1")], + "pagination": {"has_more": True, "cursor": "c1"}, + }, + { + "items": [sandbox_record("sb-2")], + "pagination": {"has_more": False}, + }, + ], + } + ) + client = make_sync_client(transport) + + names = [info.name for info in client.list()] + + assert names == ["sb-1", "sb-2"] + assert transport.count("GET", "/v1/sandboxes/instances") == 2 + client.close() + + +def test_list_repeated_cursor_raises() -> None: + transport = RoutingTransport( + routes={ + "GET /v1/sandboxes/instances": { + "items": [sandbox_record()], + "pagination": {"has_more": True, "cursor": "c1"}, + }, + } + ) + client = make_sync_client(transport) + + with pytest.raises(ValueError, match="repeated cursor"): + list(client.list()) + + client.close() + + +def test_delete_returns_record() -> None: + transport = RoutingTransport( + routes={ + "DELETE /v1/sandboxes/instances/sb-1": httpx.Response( + 202, + json=sandbox_record(status="DELETING"), + headers={"content-type": "application/json"}, + ) + } + ) + client = make_sync_client(transport) + + info = client.delete("sb-1") + + assert info.status == "DELETING" + client.close() + + +def test_control_plane_error_becomes_sandbox_api_error() -> None: + transport = RoutingTransport( + routes={ + "GET /v1/sandboxes/instances/sb-1": httpx.Response( + 404, json={"code": "NOT_FOUND", "message": "no such sandbox"} + ) + } + ) + client = make_sync_client(transport) + + with pytest.raises(SandboxApiError) as exc_info: + client.get_info("sb-1") + + assert exc_info.value.status == 404 + assert exc_info.value.code == "NOT_FOUND" + client.close() + + +def make_full_client(transport: RoutingTransport) -> SandboxClient: + transport.routes.setdefault("POST /v1/sandboxes/instances", created_response()) + transport.routes.setdefault("GET /filesystem/file.txt", file_record()) + transport.routes.setdefault("PUT /filesystem/file.txt", {"message": "written"}) + transport.routes.setdefault("POST /process", process_record()) + return make_sync_client(transport) + + +def test_fs_read_and_write_on_created_sandbox() -> None: + transport = RoutingTransport(routes={}) + client = make_full_client(transport) + + sandbox = client.create(name="sb-1") + assert sandbox.fs.read("file.txt") == "file text" + sandbox.fs.write("file.txt", "new text") + + assert transport.count("GET", "/filesystem/file.txt") == 1 + written = json.loads(transport.requests[-1][3]) + assert written == {"content": "new text"} + client.close() + + +def test_fs_read_retries_dropped_connection() -> None: + transport = RoutingTransport( + routes={ + "GET /filesystem/file.txt": [httpx.ConnectError("dropped"), file_record()] + } + ) + client = make_full_client(transport) + + sandbox = client.create(name="sb-1") + content = sandbox.fs.read("file.txt") + + assert content == "file text" + assert transport.count("GET", "/filesystem/file.txt") == 2 + client.close() + + +@pytest.mark.asyncio +async def test_async_fs_read_and_write_retry_transient_failures() -> None: + transport = RoutingTransport( + routes={ + "POST /v1/sandboxes/instances": created_response(), + "GET /filesystem/file.txt": [ + httpx.ConnectError("dropped"), + httpx.Response(502, text="bad gateway"), + file_record(), + ], + "PUT /filesystem/file.txt": [ + httpx.ConnectError("dropped"), + httpx.Response(503, text="unavailable"), + {"message": "written"}, + ], + } + ) + client = AsyncSandboxClient( + api_key="test-key", + base_url_override=MANAGEMENT_URL, + transport_override=transport.async_transport, + ) + sandbox = await client.create(name="sb-1") + + content = await sandbox.fs.read("file.txt") + await sandbox.fs.write("file.txt", content) + + assert content == "file text" + assert transport.count("GET", "/filesystem/file.txt") == 3 + assert transport.count("PUT", "/filesystem/file.txt") == 3 + await client.close() + + +def test_streaming_body_is_sent_once_on_revocation() -> None: + # A stream-backed body cannot be replayed, so a revoked token surfaces + # instead of a re-send that would have to materialize the stream. + sends: list[bytes] = [] + + def revoked_route(request: httpx.Request) -> httpx.Response: + sends.append(request.read()) + return revoked_token_response("revoked") + + transport = RoutingTransport( + routes={"POST /filesystem/big.bin": revoked_route}, + mint_tokens=["tok-1"], + ) + client = make_sync_client(transport) + + stream = httpx.Client( + transport=client._auth_transport, + base_url=MANAGEMENT_URL, + headers=_request_headers_for_test(), + ) + response = stream.post("/filesystem/big.bin", content=iter([b"part-1", b"part-2"])) + + assert response.status_code == 401 + assert sends == [b"part-1part-2"] + assert transport.mint_count == 1 + + # The revoked token left the cache: the next request mints a fresh one + # instead of reusing the known-bad token. + transport.routes["GET /v1/sandboxes/instances/sb-1"] = sandbox_record() + info = client.get_info("sb-1") + assert info.name == "sb-1" + assert transport.mint_count == 2 + client.close() + + +def _request_headers_for_test() -> dict[str, str]: + return {"user-agent": "test"} + + +def test_error_only_body_carries_its_error_as_code_and_description() -> None: + transport = RoutingTransport( + routes={ + "GET /v1/sandboxes/instances/sb-1": httpx.Response( + 409, + json={"error": "SANDBOX_NAME_TAKEN"}, + ) + } + ) + client = make_sync_client(transport) + + with pytest.raises(SandboxApiError) as exc_info: + client.get_info("sb-1") + + assert exc_info.value.code == "SANDBOX_NAME_TAKEN" + assert "SANDBOX_NAME_TAKEN" in str(exc_info.value) + client.close() + + +def test_fs_read_directory_raises() -> None: + transport = RoutingTransport( + routes={ + "GET /filesystem/dir": { + "files": [], + "name": "dir", + "path": "/dir", + "subdirectories": [], + } + } + ) + client = make_full_client(transport) + + sandbox = client.create(name="sb-1") + with pytest.raises(ValueError, match="is a directory"): + sandbox.fs.read("dir") + + client.close() + + +def test_process_list_and_logs() -> None: + def process_list_response(request: httpx.Request) -> httpx.Response: + # Callable route: a bare list payload would be consumed one entry per request. + return httpx.Response( + 200, json=[process_record()], headers={"content-type": "application/json"} + ) + + transport = RoutingTransport( + routes={ + "GET /process": process_list_response, + "GET /process/123/logs": {"logs": "hi\n", "stdout": "hi\n", "stderr": ""}, + } + ) + client = make_full_client(transport) + + sandbox = client.create(name="sb-1") + processes = sandbox.process.list() + logs = sandbox.process.logs("123") + + assert [process.pid for process in processes] == ["123"] + assert processes[0].status == "completed" + assert logs.logs == "hi\n" + assert logs.stderr == "" + client.close() + + +async def _async_sandbox(transport: RoutingTransport): + client = AsyncSandboxClient( + api_key="test-key", + base_url_override=MANAGEMENT_URL, + transport_override=transport.async_transport, + ) + return client, await client.create(name="sb-1") + + +@pytest.mark.asyncio +async def test_async_process_list_and_logs() -> None: + def process_list_response(request: httpx.Request) -> httpx.Response: + # Callable route: a bare list payload would be consumed one entry per request. + return httpx.Response( + 200, json=[process_record()], headers={"content-type": "application/json"} + ) + + transport = RoutingTransport( + routes={ + "POST /v1/sandboxes/instances": created_response(), + "GET /process": process_list_response, + "GET /process/123/logs": {"logs": "hi\n", "stdout": "hi\n", "stderr": ""}, + } + ) + client, sandbox = await _async_sandbox(transport) + + processes = await sandbox.process.list() + logs = await sandbox.process.logs("123") + + assert [process.pid for process in processes] == ["123"] + assert logs.logs == "hi\n" + await client.close() + + +def test_process_exec_returns_curated_info() -> None: + transport = RoutingTransport(routes={}) + client = make_full_client(transport) + + sandbox = client.create(name="sb-1") + info = sandbox.process.exec("echo hi", working_dir="/tmp", wait_for_completion=True) + + assert info.pid == "123" + assert info.status == "completed" + assert info.exit_code == 0 + assert info.working_dir == "/tmp" + assert info.completed_at is not None + assert info.completed_at.year == 2026 + body = json.loads(transport.requests[-1][3]) + assert body["command"] == "echo hi" + assert body["workingDir"] == "/tmp" + assert body["waitForCompletion"] is True + client.close() + + +def test_exec_gateway_error_is_sandbox_gateway_error() -> None: + transport = RoutingTransport( + routes={"POST /process": httpx.Response(502, text="bad gateway")} + ) + client = make_full_client(transport) + + sandbox = client.create(name="sb-1") + with pytest.raises(SandboxGatewayError) as exc_info: + sandbox.process.exec("echo hi") + + assert exc_info.value.status == 502 + client.close() + + +def test_get_by_name_and_by_info() -> None: + transport = RoutingTransport( + routes={"GET /v1/sandboxes/instances/sb-1": sandbox_record()} + ) + client = make_sync_client(transport) + + info = client.get_info("sb-1") + by_name = client.get("sb-1") + assert by_name.name == "sb-1" + assert by_name.info is None + + by_record = client.get(info) + assert by_record.name == "sb-1" + + # The by-record get made no call of its own. + assert transport.count("GET", "/v1/sandboxes/instances/sb-1") == 2 + client.close() + + +@pytest.mark.asyncio +async def test_async_client_create_and_exec() -> None: + transport = RoutingTransport(routes={}) + transport.routes.setdefault("POST /v1/sandboxes/instances", created_response()) + transport.routes.setdefault("POST /process", process_record()) + client = AsyncSandboxClient( + api_key="test-key", + base_url_override=MANAGEMENT_URL, + transport_override=transport.async_transport, + ) + + sandbox = await client.create(name="sb-1") + info = await sandbox.process.exec("echo hi") + + assert info.pid == "123" + assert transport.authorizations("/process") == ["Bearer tok-1"] + await client.close() + + +@pytest.mark.asyncio +async def test_async_list_paginates() -> None: + transport = RoutingTransport( + routes={ + "GET /v1/sandboxes/instances": [ + { + "items": [sandbox_record("sb-1")], + "pagination": {"has_more": True, "cursor": "c1"}, + }, + { + "items": [sandbox_record("sb-2")], + "pagination": {"has_more": False}, + }, + ], + } + ) + client = AsyncSandboxClient( + api_key="test-key", + base_url_override=MANAGEMENT_URL, + transport_override=transport.async_transport, + ) + + names = [info.name async for info in client.list()] + + assert names == ["sb-1", "sb-2"] + await client.close() + + +def test_create_without_name_omits_the_field() -> None: + transport = RoutingTransport( + routes={"POST /v1/sandboxes/instances": created_response()} + ) + client = make_sync_client(transport) + + client.create() + + body = json.loads(transport.requests[1][3]) + assert body == {} + client.close() + + +def test_fs_read_retries_gateway_errors() -> None: + transport = RoutingTransport( + routes={ + "GET /filesystem/file.txt": [ + httpx.Response(502, text="bad gateway"), + httpx.Response(503, text="unavailable"), + file_record(), + ] + } + ) + client = make_full_client(transport) + + sandbox = client.create(name="sb-1") + content = sandbox.fs.read("file.txt") + + assert content == "file text" + assert transport.count("GET", "/filesystem/file.txt") == 3 + client.close() + + +def test_persistent_revocation_returns_last_response(monkeypatch) -> None: + # The first resend is immediate; later ones wait out the revocation's + # rounded-up cutoff second, here shortened to keep the test fast. + delays: list[float] = [] + monkeypatch.setattr( + baseten.sandbox._auth, "_TOKEN_REVOKED_RETRY_DELAY_SECONDS", 0.01 + ) + monkeypatch.setattr( + baseten.sandbox._auth, + "time", + types.SimpleNamespace(sleep=lambda s: delays.append(s)), + ) + transport = RoutingTransport( + routes={ + "GET /v1/sandboxes/instances/sb-1": lambda request: revoked_token_response( + "revoked" + ) + }, + mint_tokens=["tok-1", "tok-2", "tok-3"], + ) + client = make_sync_client(transport) + + with pytest.raises(SandboxApiError) as exc_info: + client.get_info("sb-1") + + # Exhaustion surfaces the final 401 as a converted error, and every + # minted token was tried, the later resends past the cutoff. + assert exc_info.value.status == 401 + assert transport.mint_count == 3 + assert delays == [0.01] + client.close() + + +def test_get_info_show_secrets_reaches_the_query() -> None: + seen_queries: list[str] = [] + + def show_secrets_route(request: httpx.Request) -> httpx.Response: + seen_queries.append(str(request.url.params)) + return httpx.Response( + 200, json=sandbox_record(), headers={"content-type": "application/json"} + ) + + transport = RoutingTransport( + routes={"GET /v1/sandboxes/instances/sb-1": show_secrets_route} + ) + client = make_sync_client(transport) + + client.get_info("sb-1") + client.get_info("sb-1", show_secrets=True) + + assert "show_secrets=false" in seen_queries[0] + assert "show_secrets=true" in seen_queries[1] + client.close() + + +def test_sandbox_from_url_makes_no_call_and_shares_auth() -> None: + transport = RoutingTransport( + routes={ + "POST /process": process_record(), + } + ) + client = make_sync_client(transport) + + sandbox = client.sandbox_from_url("https://sb-1.invalid") + assert sandbox.name == "" + assert sandbox.url == "https://sb-1.invalid" + # No control-plane call was made to resolve it. + assert transport.count("GET", "/v1/sandboxes/instances/sb-1") == 0 + + # The sandbox shares the client's token: its exec call carries it. + sandbox.process.exec("echo hi", wait_for_completion=True) + assert transport.authorizations("/process") == ["Bearer tok-1"] + client.close() + + +def test_typed_error_body_becomes_a_dict() -> None: + transport = RoutingTransport( + routes={ + "GET /v1/sandboxes/instances/sb-1": httpx.Response( + 404, + json={"code": "NOT_FOUND", "message": "no such sandbox"}, + ) + } + ) + client = make_sync_client(transport) + + with pytest.raises(SandboxApiError) as exc_info: + client.get_info("sb-1") + + assert exc_info.value.code == "NOT_FOUND" + assert isinstance(exc_info.value.body, dict) + client.close() + + +def library_catalog_response(request: httpx.Request) -> httpx.Response: + # Callable route: a dict payload would be consumed as a scripted response. + return httpx.Response( + 200, + json={"items": library_catalog()}, + headers={"content-type": "application/json"}, + ) + + +def library_catalog() -> list[dict[str, Any]]: + # The server drops hidden and coming-soon entries before listing. + return [ + { + "name": "expo", + "image": "baseten/expo:latest", + "display_name": "Expo", + "description": "Expo app development", + "long_description": "Build and run Expo apps.", + "memory": 4096, + "categories": ["web"], + "tags": ["react-native"], + "ports": [{"name": "http", "target": 3000, "protocol": "HTTP"}], + "icon": "https://images.b10.co/expo.svg", + "url": "https://expo.dev", + "enterprise": False, + } + ] + + +def test_library_images_translates_the_served_catalog() -> None: + transport = RoutingTransport( + routes={"GET /v1/sandboxes/library_images": library_catalog_response} + ) + client = make_sync_client(transport) + + images = client.library_images() + + assert [image.name for image in images] == ["expo"] + image = images[0] + assert image.image == "baseten/expo:latest" + assert image.display_name == "Expo" + assert image.memory == 4096 + assert image.categories == ["web"] + assert image.ports[0].target == 3000 + assert image.icon_url == "https://images.b10.co/expo.svg" + assert image.project_url == "https://expo.dev" + client.close() + + +@pytest.mark.asyncio +async def test_async_library_images_translates() -> None: + transport = RoutingTransport( + routes={"GET /v1/sandboxes/library_images": library_catalog_response} + ) + client = AsyncSandboxClient( + api_key="test-key", + base_url_override=MANAGEMENT_URL, + transport_override=transport.async_transport, + ) + + images = await client.library_images() + + assert [image.name for image in images] == ["expo"] + await client.close()