diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f76292c..76c6b6d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -8,6 +8,8 @@ on: jobs: check: strategy: + # One OS or Python version failing should not hide the others' results. + fail-fast: false matrix: os: [ubuntu-latest, macos-latest, windows-latest] python-version: ["3.14"] diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index e4d8ad7..1f8799b 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -18,7 +18,7 @@ All tasks are run via `uv run poe `: ## End-to-End Tests -E2e tests in `tests/test_e2e.py` run against a live Baseten environment. They are skipped automatically when `BASETEN_E2E_TEST_API_KEY` is not set. +E2e tests in `tests/client/test_e2e.py` and `tests/sandbox/test_e2e.py` run against a live Baseten environment. They are skipped automatically when `BASETEN_E2E_TEST_API_KEY` is not set. The sandbox tests need no model, and print every request and response with `BASETEN_E2E_TEST_DEBUG=1`. ### Bootstrap diff --git a/README.md b/README.md index b4e386e..d2699eb 100644 --- a/README.md +++ b/README.md @@ -34,6 +34,31 @@ async with AsyncManagementClient(api_key="my-api-key") as client: print(model.name) ``` +## Sandboxes + +`baseten.sandbox` creates and works with sandboxes. HTTP/2 is used when the `h2` package is installed, which is strongly recommended: + +```bash +pip install "baseten" "httpx[http2]" +``` + +```python +from baseten.sandbox import SandboxClient + +with SandboxClient(api_key="my-api-key") as client: + sandbox = client.create() + try: + sandbox.fs.write(path="/tmp/hello.txt", content="hello") + result = sandbox.process.exec( + command="cat /tmp/hello.txt", wait_for_completion=True + ) + print(result.stdout) + finally: + client.delete(name=sandbox.name) +``` + +`AsyncSandboxClient` is the same for async code. Operations the SDK does not wrap are available on the generated clients, through `client.raw_api` for the control plane and `sandbox.raw_api` for a sandbox. See the [API documentation](https://basetenlabs.github.io/baseten-python/) for the rest. + ## Upgrading from 0.8.2 and earlier Version 0.9.0 is a rewrite and shares no API with the earlier `baseten` releases. diff --git a/baseten/client/inferenceapi/_client.py b/baseten/client/inferenceapi/_client.py index b60ddae..947ac1a 100644 --- a/baseten/client/inferenceapi/_client.py +++ b/baseten/client/inferenceapi/_client.py @@ -4,6 +4,7 @@ import contextlib import urllib.parse +from collections.abc import Mapping from dataclasses import dataclass from typing import Any, Literal, TypeVar @@ -72,9 +73,33 @@ class ApiClient: They may change without notice between versions. """ - def __init__(self, http_client: httpx.Client) -> None: - """Create a new client. The caller is responsible for closing *http_client*.""" + def __init__( + self, + http_client: httpx.Client, + *, + base_url: str | None = None, + auth: httpx.Auth | None = None, + headers: Mapping[str, str] | None = None, + timeout: httpx.Timeout | None = None, + ) -> None: + """Create a new client. The caller is responsible for closing *http_client*. + + The keyword arguments apply to this client's requests only, so one + HTTP client can serve several APIs. Each one left ``None`` falls back + to the HTTP client's own setting. + + Args: + http_client: HTTP client to send requests with. + base_url: Base URL that request paths are appended to. + auth: Authentication for each request. + headers: Headers for each request, added to the HTTP client's own. + timeout: Timeouts for each request. + """ self._http_client = http_client + self._base_url = base_url + self._auth = auth + self._headers = headers + self._timeout = timeout def async_predict( self, *, env_name: str, request: AsyncPredictRequest @@ -848,7 +873,7 @@ def _build_request(self, request: _ApiRequest) -> httpx.Request: json_body = None content_body = None files_body = None - headers: dict[str, str] = {} + headers: dict[str, str] = {**(self._headers or {})} if request.accept is not None: headers["Accept"] = request.accept if request.body is not None: @@ -882,16 +907,22 @@ def _build_request(self, request: _ApiRequest) -> httpx.Request: params = request.query return self._http_client.build_request( request.method, - path, + path if self._base_url is None else self._base_url.rstrip("/") + path, json=json_body, content=content_body, files=files_body, params=params, headers=headers, + timeout=httpx.USE_CLIENT_DEFAULT + if self._timeout is None + else self._timeout, ) def _do(self, request: _ApiRequest) -> httpx.Response: - response = self._http_client.send(self._build_request(request)) + response = self._http_client.send( + self._build_request(request), + auth=httpx.USE_CLIENT_DEFAULT if self._auth is None else self._auth, + ) if response.status_code not in request.success_codes: if request.error_codes and response.status_code in request.error_codes: error_name = request.error_codes[response.status_code] @@ -911,7 +942,11 @@ def _do(self, request: _ApiRequest) -> httpx.Response: return response def _do_raw(self, request: _ApiRequest) -> httpx.Response: - response = self._http_client.send(self._build_request(request), stream=True) + response = self._http_client.send( + self._build_request(request), + auth=httpx.USE_CLIENT_DEFAULT if self._auth is None else self._auth, + stream=True, + ) if response.status_code not in request.success_codes: response.read() if request.error_codes and response.status_code in request.error_codes: @@ -952,9 +987,33 @@ class AsyncApiClient: They may change without notice between versions. """ - def __init__(self, http_client: httpx.AsyncClient) -> None: - """Create a new client. The caller is responsible for closing *http_client*.""" + def __init__( + self, + http_client: httpx.AsyncClient, + *, + base_url: str | None = None, + auth: httpx.Auth | None = None, + headers: Mapping[str, str] | None = None, + timeout: httpx.Timeout | None = None, + ) -> None: + """Create a new client. The caller is responsible for closing *http_client*. + + The keyword arguments apply to this client's requests only, so one + HTTP client can serve several APIs. Each one left ``None`` falls back + to the HTTP client's own setting. + + Args: + http_client: HTTP client to send requests with. + base_url: Base URL that request paths are appended to. + auth: Authentication for each request. + headers: Headers for each request, added to the HTTP client's own. + timeout: Timeouts for each request. + """ self._http_client = http_client + self._base_url = base_url + self._auth = auth + self._headers = headers + self._timeout = timeout async def async_predict( self, *, env_name: str, request: AsyncPredictRequest @@ -1738,7 +1797,7 @@ def _build_request(self, request: _ApiRequest) -> httpx.Request: json_body = None content_body = None files_body = None - headers: dict[str, str] = {} + headers: dict[str, str] = {**(self._headers or {})} if request.accept is not None: headers["Accept"] = request.accept if request.body is not None: @@ -1772,16 +1831,22 @@ def _build_request(self, request: _ApiRequest) -> httpx.Request: params = request.query return self._http_client.build_request( request.method, - path, + path if self._base_url is None else self._base_url.rstrip("/") + path, json=json_body, content=content_body, files=files_body, params=params, headers=headers, + timeout=httpx.USE_CLIENT_DEFAULT + if self._timeout is None + else self._timeout, ) async def _do(self, request: _ApiRequest) -> httpx.Response: - response = await self._http_client.send(self._build_request(request)) + response = await self._http_client.send( + self._build_request(request), + auth=httpx.USE_CLIENT_DEFAULT if self._auth is None else self._auth, + ) if response.status_code not in request.success_codes: if request.error_codes and response.status_code in request.error_codes: error_name = request.error_codes[response.status_code] @@ -1802,7 +1867,9 @@ async def _do(self, request: _ApiRequest) -> httpx.Response: async def _do_raw(self, request: _ApiRequest) -> httpx.Response: response = await self._http_client.send( - self._build_request(request), stream=True + self._build_request(request), + auth=httpx.USE_CLIENT_DEFAULT if self._auth is None else self._auth, + stream=True, ) if response.status_code not in request.success_codes: await response.aread() diff --git a/baseten/client/managementapi/__init__.py b/baseten/client/managementapi/__init__.py index 8e465c2..7273e32 100644 --- a/baseten/client/managementapi/__init__.py +++ b/baseten/client/managementapi/__init__.py @@ -89,8 +89,7 @@ AwsIamDockerAuth, AwsOidcDockerAuth, Background, - Background1, - BackgroundHarnessModels, + BackgroundHarnessDefaults, BasetenLatestCheckpointConfig, BasetenNamedCheckpointConfig, BenchmarkSnapshot, @@ -111,8 +110,9 @@ Chains, ChainTombstone, CheckpointFile, + ClassifierEnvironmentName, CleanupImagesParams, - CleanupImagesResponse, + CleanupSandboxImagesResponse, Component, ComputeCost, ComputeCost1, @@ -140,8 +140,10 @@ CreateLoopsSamplerRequest, CreateLoopsSamplerResponse, CreateLoopsSessionResponse, + CreateLoopsTrainerRequest, CreateModelDeploymentRequest, CreateModelRequest, + CreateRouteConnectionRequest, CreateRouteRequest, CreateSandboxParams, CreateSandboxRequest, @@ -171,7 +173,6 @@ DedicatedUsage, DeleteImageParams, DeleteImageTagParams, - DeleteRoutesHarnessConfigsParams, DeleteSandboxParams, DeleteVolumeRequest, DeleteVolumeResponse, @@ -218,6 +219,7 @@ EnvironmentGroupUser, Environments, EnvironmentTombstone, + ErrorRate, ExploreCost, ExploreCostValues, ExploreMetadata, @@ -274,9 +276,12 @@ GetModelsEnvironmentsLogsParams, GetModelsEnvironmentsMetricsParams, GetModelsParams, - GetRoutesHarnessConfigsParams, + GetRoutesConnectionsParams, GetRoutesParams, GetRoutesUsageParams, + GetSandboxConfigurationResponse, + GetSandboxLogsParams, + GetSandboxMetricsParams, GetSandboxParams, GetTeamsLoopsRunsParams, GetTeamsLoopsSamplersParams, @@ -304,10 +309,6 @@ GroupHierarchy, GroupMetadata, GroupsResponse, - Image, - ImageBuildLog, - ImageBuildLogsResponse, - ImageTag, InferenceVolumeByStatusDatapoint, InProgressPromotion, InstanceType, @@ -330,16 +331,15 @@ Limit6, ListAuditLogsResponse, ListImagesParams, - ListImagesResponse, ListImageTagsParams, - ListImageTagsResponse, ListLoopsCheckpointsResponse, ListLoopsDeploymentsResponse, ListLoopsRunsResponse, ListLoopsSamplersResponse, ListSandboxesParams, ListSandboxesResponse, - ListSandboxLibraryImagesParams, + ListSandboxImagesResponse, + ListSandboxImageTagsResponse, ListSandboxLibraryImagesResponse, ListTrainingJobsResponse, ListTrainingProjectsResponse, @@ -391,7 +391,8 @@ Name1, Name2, Name3, - Name4, + Name5, + NumReplicas, OneTimeAutoscalingSchedule, OneTimeAutoscalingScheduleUpsert, OrderBy, @@ -407,15 +408,14 @@ PrepareModelUploadRequest, PrepareModelUploadResponse, Primary, - Primary2, - PrimaryHarnessModels, + PrimaryHarnessDefaults, PromoteRequest, PromoteToChainEnvironmentRequest, PromoteToEnvironmentRequest, PromotionSettings, PushImageParams, - PushImageRequest, - PushImageResponse, + PushSandboxImageRequest, + PushSandboxImageResponse, QueueEvent, RateLimit, RecreateTrainingJobResponse, @@ -433,46 +433,73 @@ RetryDeploymentResponse, RollingDeployConfig, Route, - RouteHarnessConfig, - RouteHarnessConfigsResponse, - RouteHarnessConfigTombstone, + RouteConnection, + RouteConnectionAnthropic, + RouteConnectionConfigAnthropic, + RouteConnectionConfigOpenAI, + RouteConnectionConfigXAI, + RouteConnectionOpenAI, + RouteConnectionsResponse, + RouteConnectionTombstone, + RouteConnectionXAI, + RouteEffectiveSpendLimit, + RouteHarnessDefaults, RouteHarnessModel, - RouteSpendLimit, + RouteRef, + RouteSpendLimitSetting, + RouteSpendLimitTeamDefault, RoutesResponse, RoutesUsageBucket, RoutesUsageResponse, RoutesUsageResult, RouteTargetAnthropic, RouteTargetBasetenModelAPI, + RouteTargetClassifierModelBased, + RouteTargetConfigAnthropic, + RouteTargetConfigBasetenModelAPI, + RouteTargetConfigClassifierModelBased, + RouteTargetConfigOpenAI, + RouteTargetConfigXAI, RouteTargetOpenAI, RouteTargetXAI, + RouteTeamSettings, + RouteTeamSpendLimitSetting, RouteTombstone, + RouteUserSettings, + SamplerNumReplicas, Sandbox, SandboxApiPagination, SandboxConfiguration, SandboxDateExpirationPolicy, SandboxEnv, + SandboxImage, + SandboxImageBuildLog, + SandboxImageBuildLogsResponse, + SandboxImageSummary, + SandboxImageTag, SandboxLibraryImage, SandboxLibraryImageCreationOptions, SandboxLibraryImageVolume, SandboxLifecycle, + SandboxLogEntry, + SandboxLogs, + SandboxMetrics, + SandboxMetricsPoint, SandboxNetwork, SandboxPort, SandboxProxyConfig, SandboxProxyTarget, + SandboxRegion, SandboxTTLIdleExpirationPolicy, SandboxTTLMaxAgeExpirationPolicy, SearchPattern, SearchTrainingJobsRequest, SearchTrainingJobsResponse, Secret, + SecretId, SecretReference, Secrets, SecretTombstone, - SetClaudeCodeHarnessConfig, - SetCodexHarnessConfig, - SetOpenCodeHarnessConfig, - SetRouteHarnessConfigRequest, SignalPromotionResponse, SignSSHCertificateRequest, SignSSHCertificateResponse, @@ -496,6 +523,7 @@ Team, TeamId, TeamId1, + TeamId2, Teams, TeamTrainingGpuCapacityItem, TerminateReplicaResponse, @@ -527,8 +555,6 @@ UpdateChainletEnvironmentAutoscalingSettingsRequest, UpdateChainletEnvironmentInstanceTypeRequest, UpdateChainletEnvironmentInstanceTypeResponse, - UpdateClaudeCodeHarnessConfig, - UpdateCodexHarnessConfig, UpdateDeploymentRequest, UpdateEndpointRequest, UpdateEnvironmentGroupManageAccess, @@ -540,14 +566,20 @@ UpdateLibraryListingRequest, UpdateLibraryListingVersionRequest, UpdateModelRequest, - UpdateOpenCodeHarnessConfig, UpdatePrimaryHarnessModels, UpdatePromotionSettings, UpdateRequestBackpressureSettings, UpdateRollingDeployConfig, - UpdateRouteHarnessConfigRequest, + UpdateRouteConnectionConfigAnthropic, + UpdateRouteConnectionConfigOpenAI, + UpdateRouteConnectionConfigXAI, + UpdateRouteConnectionRequest, + UpdateRouteHarnessDefaults, UpdateRouteRequest, - UpdateRouteSpendLimitRequest, + UpdateRouteSpendLimitSetting, + UpdateRouteTeamSettingsRequest, + UpdateRouteTeamSpendLimitSetting, + UpdateRouteUserSettingsRequest, UpdateSandboxParams, UpdateSandboxRequest, UpdateTrainingJobRequest, @@ -665,8 +697,7 @@ "AwsIamDockerAuth", "AwsOidcDockerAuth", "Background", - "Background1", - "BackgroundHarnessModels", + "BackgroundHarnessDefaults", "BasetenLatestCheckpointConfig", "BasetenNamedCheckpointConfig", "BenchmarkSnapshot", @@ -687,8 +718,9 @@ "ChainletEnvironmentSettingsRequest", "Chains", "CheckpointFile", + "ClassifierEnvironmentName", "CleanupImagesParams", - "CleanupImagesResponse", + "CleanupSandboxImagesResponse", "Component", "ComputeCost", "ComputeCost1", @@ -715,8 +747,10 @@ "CreateLoopsSamplerRequest", "CreateLoopsSamplerResponse", "CreateLoopsSessionResponse", + "CreateLoopsTrainerRequest", "CreateModelDeploymentRequest", "CreateModelRequest", + "CreateRouteConnectionRequest", "CreateRouteRequest", "CreateSandboxParams", "CreateSandboxRequest", @@ -747,7 +781,6 @@ "DedicatedUsage", "DeleteImageParams", "DeleteImageTagParams", - "DeleteRoutesHarnessConfigsParams", "DeleteSandboxParams", "DeleteVolumeRequest", "DeleteVolumeResponse", @@ -794,6 +827,7 @@ "EnvironmentGroups", "EnvironmentTombstone", "Environments", + "ErrorRate", "ExploreCost", "ExploreCostValues", "ExploreMetadata", @@ -850,9 +884,12 @@ "GetModelsEnvironmentsLogsParams", "GetModelsEnvironmentsMetricsParams", "GetModelsParams", - "GetRoutesHarnessConfigsParams", + "GetRoutesConnectionsParams", "GetRoutesParams", "GetRoutesUsageParams", + "GetSandboxConfigurationResponse", + "GetSandboxLogsParams", + "GetSandboxMetricsParams", "GetSandboxParams", "GetTeamsLoopsRunsParams", "GetTeamsLoopsSamplersParams", @@ -880,10 +917,6 @@ "GroupHierarchy", "GroupMetadata", "GroupsResponse", - "Image", - "ImageBuildLog", - "ImageBuildLogsResponse", - "ImageTag", "InProgressPromotion", "InferenceVolumeByStatusDatapoint", "InstanceType", @@ -908,14 +941,13 @@ "Limit6", "ListAuditLogsResponse", "ListImageTagsParams", - "ListImageTagsResponse", "ListImagesParams", - "ListImagesResponse", "ListLoopsCheckpointsResponse", "ListLoopsDeploymentsResponse", "ListLoopsRunsResponse", "ListLoopsSamplersResponse", - "ListSandboxLibraryImagesParams", + "ListSandboxImageTagsResponse", + "ListSandboxImagesResponse", "ListSandboxLibraryImagesResponse", "ListSandboxesParams", "ListSandboxesResponse", @@ -967,7 +999,8 @@ "Name1", "Name2", "Name3", - "Name4", + "Name5", + "NumReplicas", "OneTimeAutoscalingSchedule", "OneTimeAutoscalingScheduleUpsert", "OrderBy", @@ -983,15 +1016,14 @@ "PrepareModelUploadRequest", "PrepareModelUploadResponse", "Primary", - "Primary2", - "PrimaryHarnessModels", + "PrimaryHarnessDefaults", "PromoteRequest", "PromoteToChainEnvironmentRequest", "PromoteToEnvironmentRequest", "PromotionSettings", "PushImageParams", - "PushImageRequest", - "PushImageResponse", + "PushSandboxImageRequest", + "PushSandboxImageResponse", "QueueEvent", "RateLimit", "RecreateTrainingJobResponse", @@ -1010,46 +1042,73 @@ "RetryDeploymentResponse", "RollingDeployConfig", "Route", - "RouteHarnessConfig", - "RouteHarnessConfigTombstone", - "RouteHarnessConfigsResponse", + "RouteConnection", + "RouteConnectionAnthropic", + "RouteConnectionConfigAnthropic", + "RouteConnectionConfigOpenAI", + "RouteConnectionConfigXAI", + "RouteConnectionOpenAI", + "RouteConnectionTombstone", + "RouteConnectionXAI", + "RouteConnectionsResponse", + "RouteEffectiveSpendLimit", + "RouteHarnessDefaults", "RouteHarnessModel", - "RouteSpendLimit", + "RouteRef", + "RouteSpendLimitSetting", + "RouteSpendLimitTeamDefault", "RouteTargetAnthropic", "RouteTargetBasetenModelAPI", + "RouteTargetClassifierModelBased", + "RouteTargetConfigAnthropic", + "RouteTargetConfigBasetenModelAPI", + "RouteTargetConfigClassifierModelBased", + "RouteTargetConfigOpenAI", + "RouteTargetConfigXAI", "RouteTargetOpenAI", "RouteTargetXAI", + "RouteTeamSettings", + "RouteTeamSpendLimitSetting", "RouteTombstone", + "RouteUserSettings", "RoutesResponse", "RoutesUsageBucket", "RoutesUsageResponse", "RoutesUsageResult", + "SamplerNumReplicas", "Sandbox", "SandboxApiPagination", "SandboxConfiguration", "SandboxDateExpirationPolicy", "SandboxEnv", + "SandboxImage", + "SandboxImageBuildLog", + "SandboxImageBuildLogsResponse", + "SandboxImageSummary", + "SandboxImageTag", "SandboxLibraryImage", "SandboxLibraryImageCreationOptions", "SandboxLibraryImageVolume", "SandboxLifecycle", + "SandboxLogEntry", + "SandboxLogs", + "SandboxMetrics", + "SandboxMetricsPoint", "SandboxNetwork", "SandboxPort", "SandboxProxyConfig", "SandboxProxyTarget", + "SandboxRegion", "SandboxTTLIdleExpirationPolicy", "SandboxTTLMaxAgeExpirationPolicy", "SearchPattern", "SearchTrainingJobsRequest", "SearchTrainingJobsResponse", "Secret", + "SecretId", "SecretReference", "SecretTombstone", "Secrets", - "SetClaudeCodeHarnessConfig", - "SetCodexHarnessConfig", - "SetOpenCodeHarnessConfig", - "SetRouteHarnessConfigRequest", "SignSSHCertificateRequest", "SignSSHCertificateResponse", "SignalPromotionResponse", @@ -1074,6 +1133,7 @@ "Team", "TeamId", "TeamId1", + "TeamId2", "TeamTrainingGpuCapacityItem", "Teams", "TerminateReplicaResponse", @@ -1104,8 +1164,6 @@ "UpdateChainletEnvironmentAutoscalingSettingsRequest", "UpdateChainletEnvironmentInstanceTypeRequest", "UpdateChainletEnvironmentInstanceTypeResponse", - "UpdateClaudeCodeHarnessConfig", - "UpdateCodexHarnessConfig", "UpdateDeploymentRequest", "UpdateEndpointRequest", "UpdateEnvironmentGroupManageAccess", @@ -1117,14 +1175,20 @@ "UpdateLibraryListingRequest", "UpdateLibraryListingVersionRequest", "UpdateModelRequest", - "UpdateOpenCodeHarnessConfig", "UpdatePrimaryHarnessModels", "UpdatePromotionSettings", "UpdateRequestBackpressureSettings", "UpdateRollingDeployConfig", - "UpdateRouteHarnessConfigRequest", + "UpdateRouteConnectionConfigAnthropic", + "UpdateRouteConnectionConfigOpenAI", + "UpdateRouteConnectionConfigXAI", + "UpdateRouteConnectionRequest", + "UpdateRouteHarnessDefaults", "UpdateRouteRequest", - "UpdateRouteSpendLimitRequest", + "UpdateRouteSpendLimitSetting", + "UpdateRouteTeamSettingsRequest", + "UpdateRouteTeamSpendLimitSetting", + "UpdateRouteUserSettingsRequest", "UpdateSandboxParams", "UpdateSandboxRequest", "UpdateTrainingJobRequest", diff --git a/baseten/client/managementapi/_client.py b/baseten/client/managementapi/_client.py index 1e083e6..e295dd8 100644 --- a/baseten/client/managementapi/_client.py +++ b/baseten/client/managementapi/_client.py @@ -3,6 +3,7 @@ from __future__ import annotations import urllib.parse +from collections.abc import Mapping from dataclasses import dataclass from typing import Any, TypeVar @@ -23,7 +24,7 @@ Chains, ChainTombstone, CleanupImagesParams, - CleanupImagesResponse, + CleanupSandboxImagesResponse, CreateApiKeyForGroupRequest, CreateApiKeyForGroupResponse, CreateAPIKeyRequest, @@ -43,8 +44,10 @@ CreateLoopsSamplerRequest, CreateLoopsSamplerResponse, CreateLoopsSessionResponse, + CreateLoopsTrainerRequest, CreateModelDeploymentRequest, CreateModelRequest, + CreateRouteConnectionRequest, CreateRouteRequest, CreateSandboxParams, CreateSandboxRequest, @@ -60,7 +63,6 @@ DeactivateResponse, DeleteImageParams, DeleteImageTagParams, - DeleteRoutesHarnessConfigsParams, DeleteSandboxParams, DeleteVolumeRequest, DeleteVolumeResponse, @@ -129,9 +131,12 @@ GetModelsEnvironmentsLogsParams, GetModelsEnvironmentsMetricsParams, GetModelsParams, - GetRoutesHarnessConfigsParams, + GetRoutesConnectionsParams, GetRoutesParams, GetRoutesUsageParams, + GetSandboxConfigurationResponse, + GetSandboxLogsParams, + GetSandboxMetricsParams, GetSandboxParams, GetTeamsLoopsRunsParams, GetTeamsLoopsSamplersParams, @@ -156,8 +161,6 @@ GetVolumesVersionsParams, Group, GroupsResponse, - Image, - ImageBuildLogsResponse, InstanceTypePrices, InstanceTypes, KeysForGroupResponse, @@ -169,16 +172,15 @@ LibraryListingVersionTombstone, ListAuditLogsResponse, ListImagesParams, - ListImagesResponse, ListImageTagsParams, - ListImageTagsResponse, ListLoopsCheckpointsResponse, ListLoopsDeploymentsResponse, ListLoopsRunsResponse, ListLoopsSamplersResponse, ListSandboxesParams, ListSandboxesResponse, - ListSandboxLibraryImagesParams, + ListSandboxImagesResponse, + ListSandboxImageTagsResponse, ListSandboxLibraryImagesResponse, ListTrainingJobsResponse, ListTrainingProjectsResponse, @@ -209,8 +211,8 @@ PromoteToChainEnvironmentRequest, PromoteToEnvironmentRequest, PushImageParams, - PushImageRequest, - PushImageResponse, + PushSandboxImageRequest, + PushSandboxImageResponse, RecreateTrainingJobResponse, Regions, RegisterAPIKeyRequest, @@ -220,20 +222,24 @@ RestoreVolumeVersionResponse, RetryDeploymentResponse, Route, - RouteHarnessConfig, - RouteHarnessConfigsResponse, - RouteHarnessConfigTombstone, - RouteSpendLimit, + RouteConnection, + RouteConnectionsResponse, + RouteConnectionTombstone, RoutesResponse, RoutesUsageResponse, + RouteTeamSettings, RouteTombstone, + RouteUserSettings, Sandbox, + SandboxImage, + SandboxImageBuildLogsResponse, + SandboxLogs, + SandboxMetrics, SearchTrainingJobsRequest, SearchTrainingJobsResponse, Secret, Secrets, SecretTombstone, - SetRouteHarnessConfigRequest, SignalPromotionResponse, SignSSHCertificateRequest, SignSSHCertificateResponse, @@ -265,9 +271,10 @@ UpdateLibraryListingVersionRequest, UpdateModelRequest, UpdateRequestBackpressureSettings, - UpdateRouteHarnessConfigRequest, + UpdateRouteConnectionRequest, UpdateRouteRequest, - UpdateRouteSpendLimitRequest, + UpdateRouteTeamSettingsRequest, + UpdateRouteUserSettingsRequest, UpdateSandboxParams, UpdateSandboxRequest, UpdateTrainingJobRequest, @@ -319,16 +326,40 @@ class ApiClient: They may change without notice between versions. """ - def __init__(self, http_client: httpx.Client) -> None: - """Create a new client. The caller is responsible for closing *http_client*.""" + def __init__( + self, + http_client: httpx.Client, + *, + base_url: str | None = None, + auth: httpx.Auth | None = None, + headers: Mapping[str, str] | None = None, + timeout: httpx.Timeout | None = None, + ) -> None: + """Create a new client. The caller is responsible for closing *http_client*. + + The keyword arguments apply to this client's requests only, so one + HTTP client can serve several APIs. Each one left ``None`` falls back + to the HTTP client's own setting. + + Args: + http_client: HTTP client to send requests with. + base_url: Base URL that request paths are appended to. + auth: Authentication for each request. + headers: Headers for each request, added to the HTTP client's own. + timeout: Timeouts for each request. + """ self._http_client = http_client + self._base_url = base_url + self._auth = auth + self._headers = headers + self._timeout = timeout def cleanup_images( self, *, params: CleanupImagesParams | None = None - ) -> CleanupImagesResponse: + ) -> CleanupSandboxImagesResponse: """Clean up unused images""" return self._do_json( - CleanupImagesResponse, + CleanupSandboxImagesResponse, _ApiRequest( method="POST", path_fmt="/v1/sandboxes/cleanup_images", @@ -424,10 +455,10 @@ def delete_gateway_endpoints(self, *, endpoint_id: str) -> EndpointTombstone: def delete_image( self, *, image_name: str, params: DeleteImageParams | None = None - ) -> Image: + ) -> SandboxImage: """Delete a sandbox image""" return self._do_json( - Image, + SandboxImage, _ApiRequest( method="DELETE", path_fmt="/v1/sandboxes/images/{}", @@ -445,10 +476,10 @@ def delete_image_tag( image_name: str, tag_name: str, params: DeleteImageTagParams | None = None, - ) -> Image: + ) -> SandboxImage: """Delete an image tag""" return self._do_json( - Image, + SandboxImage, _ApiRequest( method="DELETE", path_fmt="/v1/sandboxes/images/{}/tags/{}", @@ -575,18 +606,18 @@ def delete_routes(self, *, route_id: str) -> RouteTombstone: ), ) - def delete_routes_harness_configs( - self, *, harness: str, params: DeleteRoutesHarnessConfigsParams - ) -> RouteHarnessConfigTombstone: - """Clears default models for a coding harness""" + def delete_routes_connections( + self, *, connection_id: str + ) -> RouteConnectionTombstone: + """Deletes a connection""" return self._do_json( - RouteHarnessConfigTombstone, + RouteConnectionTombstone, _ApiRequest( method="DELETE", - path_fmt="/v1/routes/harness-configs/{}", - path_args=[harness], + path_fmt="/v1/routes/connections/{}", + path_args=[connection_id], body=None, - query=params, + query=None, success_codes=[200], error_codes=None, ), @@ -1118,10 +1149,10 @@ def get_gateway_groups_group_id(self, *, group_id: str) -> Group: def get_image( self, *, image_name: str, params: GetImageParams | None = None - ) -> Image: + ) -> SandboxImage: """Get a sandbox image""" return self._do_json( - Image, + SandboxImage, _ApiRequest( method="GET", path_fmt="/v1/sandboxes/images/{}", @@ -1135,10 +1166,10 @@ def get_image( def get_image_build_logs( self, *, image_name: str, params: GetImageBuildLogsParams | None = None - ) -> ImageBuildLogsResponse: + ) -> SandboxImageBuildLogsResponse: """Get image build logs""" return self._do_json( - ImageBuildLogsResponse, + SandboxImageBuildLogsResponse, _ApiRequest( method="GET", path_fmt="/v1/sandboxes/images/{}/logs", @@ -1860,15 +1891,15 @@ def get_routes(self, *, params: GetRoutesParams | None = None) -> RoutesResponse ), ) - def get_routes_harness_configs( - self, *, params: GetRoutesHarnessConfigsParams - ) -> RouteHarnessConfigsResponse: - """Lists default models for coding harnesses""" + def get_routes_connections( + self, *, params: GetRoutesConnectionsParams | None = None + ) -> RouteConnectionsResponse: + """Lists connections""" return self._do_json( - RouteHarnessConfigsResponse, + RouteConnectionsResponse, _ApiRequest( method="GET", - path_fmt="/v1/routes/harness-configs", + path_fmt="/v1/routes/connections", path_args=[], body=None, query=params, @@ -1892,13 +1923,28 @@ def get_routes_route_id(self, *, route_id: str) -> Route: ), ) - def get_routes_spend_limits(self, *, user_id: str) -> RouteSpendLimit: - """Gets a user's spend limits""" + def get_routes_settings_teams(self, *, team_id: str) -> RouteTeamSettings: + """Gets a team's route settings""" + return self._do_json( + RouteTeamSettings, + _ApiRequest( + method="GET", + path_fmt="/v1/routes/settings/teams/{}", + path_args=[team_id], + body=None, + query=None, + success_codes=[200], + error_codes=None, + ), + ) + + def get_routes_settings_users(self, *, user_id: str) -> RouteUserSettings: + """Gets a user's route settings""" return self._do_json( - RouteSpendLimit, + RouteUserSettings, _ApiRequest( method="GET", - path_fmt="/v1/routes/spend_limits/{}", + path_fmt="/v1/routes/settings/users/{}", path_args=[user_id], body=None, query=None, @@ -1941,6 +1987,55 @@ def get_sandbox( ), ) + def get_sandbox_configuration(self) -> GetSandboxConfigurationResponse: + """Get sandbox configuration""" + return self._do_json( + GetSandboxConfigurationResponse, + _ApiRequest( + method="GET", + path_fmt="/v1/sandboxes/configuration", + path_args=[], + body=None, + query=None, + success_codes=[200], + error_codes=None, + ), + ) + + def get_sandbox_logs( + self, *, sandbox_name: str, params: GetSandboxLogsParams + ) -> SandboxLogs: + """Get sandbox logs""" + return self._do_json( + SandboxLogs, + _ApiRequest( + method="GET", + path_fmt="/v1/sandboxes/instances/{}/logs", + path_args=[sandbox_name], + body=None, + query=params, + success_codes=[200], + error_codes=None, + ), + ) + + def get_sandbox_metrics( + self, *, sandbox_name: str, params: GetSandboxMetricsParams + ) -> SandboxMetrics: + """Get sandbox metrics""" + return self._do_json( + SandboxMetrics, + _ApiRequest( + method="GET", + path_fmt="/v1/sandboxes/instances/{}/metrics", + path_args=[sandbox_name], + body=None, + query=params, + success_codes=[200], + error_codes=None, + ), + ) + def get_secrets(self) -> Secrets: """Gets all secrets (metadata only, no plain text keys)""" return self._do_json( @@ -2494,10 +2589,10 @@ def get_volumes_volume_name( def list_image_tags( self, *, image_name: str, params: ListImageTagsParams | None = None - ) -> ListImageTagsResponse: + ) -> ListSandboxImageTagsResponse: """List image tags""" return self._do_json( - ListImageTagsResponse, + ListSandboxImageTagsResponse, _ApiRequest( method="GET", path_fmt="/v1/sandboxes/images/{}/tags", @@ -2511,10 +2606,10 @@ def list_image_tags( def list_images( self, *, params: ListImagesParams | None = None - ) -> ListImagesResponse: + ) -> ListSandboxImagesResponse: """List sandbox images""" return self._do_json( - ListImagesResponse, + ListSandboxImagesResponse, _ApiRequest( method="GET", path_fmt="/v1/sandboxes/images", @@ -2526,9 +2621,7 @@ def list_images( ), ) - def list_sandbox_library_images( - self, *, params: ListSandboxLibraryImagesParams | None = None - ) -> ListSandboxLibraryImagesResponse: + def list_sandbox_library_images(self) -> ListSandboxLibraryImagesResponse: """List built-in sandbox images""" return self._do_json( ListSandboxLibraryImagesResponse, @@ -2537,7 +2630,7 @@ def list_sandbox_library_images( path_fmt="/v1/sandboxes/library_images", path_args=[], body=None, - query=params, + query=None, success_codes=[200], error_codes=None, ), @@ -2826,7 +2919,7 @@ def patch_models_environments( ) def patch_routes(self, *, route_id: str, request: UpdateRouteRequest) -> Route: - """Updates a route's display name or description""" + """Updates a route""" return self._do_json( Route, _ApiRequest( @@ -2840,16 +2933,33 @@ def patch_routes(self, *, route_id: str, request: UpdateRouteRequest) -> Route: ), ) - def patch_routes_harness_configs( - self, *, request: UpdateRouteHarnessConfigRequest - ) -> RouteHarnessConfig: - """Updates default models for a coding harness""" + def patch_routes_connections( + self, *, connection_id: str, request: UpdateRouteConnectionRequest + ) -> RouteConnection: + """Updates a connection""" return self._do_json( - RouteHarnessConfig, + RouteConnection, _ApiRequest( method="PATCH", - path_fmt="/v1/routes/harness-configs", - path_args=[], + path_fmt="/v1/routes/connections/{}", + path_args=[connection_id], + body=request, + query=None, + success_codes=[200], + error_codes=None, + ), + ) + + def patch_routes_settings_teams( + self, *, team_id: str, request: UpdateRouteTeamSettingsRequest + ) -> RouteTeamSettings: + """Updates a team's route settings""" + return self._do_json( + RouteTeamSettings, + _ApiRequest( + method="PATCH", + path_fmt="/v1/routes/settings/teams/{}", + path_args=[team_id], body=request, query=None, success_codes=[200], @@ -2857,15 +2967,15 @@ def patch_routes_harness_configs( ), ) - def patch_routes_spend_limits( - self, *, user_id: str, request: UpdateRouteSpendLimitRequest - ) -> RouteSpendLimit: - """Updates a user's spend limits""" + def patch_routes_settings_users( + self, *, user_id: str, request: UpdateRouteUserSettingsRequest + ) -> RouteUserSettings: + """Updates a user's route settings""" return self._do_json( - RouteSpendLimit, + RouteUserSettings, _ApiRequest( method="PATCH", - path_fmt="/v1/routes/spend_limits/{}", + path_fmt="/v1/routes/settings/users/{}", path_args=[user_id], body=request, query=None, @@ -3321,7 +3431,7 @@ def post_loops_sessions(self) -> CreateLoopsSessionResponse: ) def post_loops_trainers( - self, *, request: CreateLoopsRunRequest + self, *, request: CreateLoopsTrainerRequest ) -> CreateLoopsRunResponse: """Creates a Loops trainer""" return self._do_json( @@ -3817,6 +3927,23 @@ def post_routes(self, *, request: CreateRouteRequest) -> Route: ), ) + def post_routes_connections( + self, *, request: CreateRouteConnectionRequest + ) -> RouteConnection: + """Creates a connection""" + return self._do_json( + RouteConnection, + _ApiRequest( + method="POST", + path_fmt="/v1/routes/connections", + path_args=[], + body=request, + query=None, + success_codes=[200], + error_codes=None, + ), + ) + def post_secrets(self, *, request: UpsertSecretRequest) -> Secret: """Upserts a secret""" return self._do_json( @@ -3916,7 +4043,7 @@ def post_teams_loops_sessions(self, *, team_id: str) -> CreateLoopsSessionRespon ) def post_teams_loops_trainers( - self, *, team_id: str, request: CreateLoopsRunRequest + self, *, team_id: str, request: CreateLoopsTrainerRequest ) -> CreateLoopsRunResponse: """Creates a Loops trainer""" return self._do_json( @@ -4220,11 +4347,11 @@ def post_volumes_versions_restore( ) def push_image( - self, *, params: PushImageParams | None = None, request: PushImageRequest - ) -> PushImageResponse: + self, *, params: PushImageParams | None = None, request: PushSandboxImageRequest + ) -> PushSandboxImageResponse: """Push a sandbox image""" return self._do_json( - PushImageResponse, + PushSandboxImageResponse, _ApiRequest( method="POST", path_fmt="/v1/sandboxes/images", @@ -4236,23 +4363,6 @@ def push_image( ), ) - def put_routes_harness_configs( - self, *, request: SetRouteHarnessConfigRequest - ) -> RouteHarnessConfig: - """Sets default models for a coding harness""" - return self._do_json( - RouteHarnessConfig, - _ApiRequest( - method="PUT", - path_fmt="/v1/routes/harness-configs", - path_args=[], - body=request, - query=None, - success_codes=[200], - error_codes=None, - ), - ) - def update_sandbox( self, *, @@ -4281,7 +4391,7 @@ def _build_request(self, request: _ApiRequest) -> httpx.Request: json_body = None content_body = None files_body = None - headers: dict[str, str] = {} + headers: dict[str, str] = {**(self._headers or {})} if request.accept is not None: headers["Accept"] = request.accept if request.body is not None: @@ -4315,16 +4425,22 @@ def _build_request(self, request: _ApiRequest) -> httpx.Request: params = request.query return self._http_client.build_request( request.method, - path, + path if self._base_url is None else self._base_url.rstrip("/") + path, json=json_body, content=content_body, files=files_body, params=params, headers=headers, + timeout=httpx.USE_CLIENT_DEFAULT + if self._timeout is None + else self._timeout, ) def _do(self, request: _ApiRequest) -> httpx.Response: - response = self._http_client.send(self._build_request(request)) + response = self._http_client.send( + self._build_request(request), + auth=httpx.USE_CLIENT_DEFAULT if self._auth is None else self._auth, + ) if response.status_code not in request.success_codes: raise ResponseError(status_code=response.status_code, body=response.text) return response @@ -4347,16 +4463,40 @@ class AsyncApiClient: They may change without notice between versions. """ - def __init__(self, http_client: httpx.AsyncClient) -> None: - """Create a new client. The caller is responsible for closing *http_client*.""" + def __init__( + self, + http_client: httpx.AsyncClient, + *, + base_url: str | None = None, + auth: httpx.Auth | None = None, + headers: Mapping[str, str] | None = None, + timeout: httpx.Timeout | None = None, + ) -> None: + """Create a new client. The caller is responsible for closing *http_client*. + + The keyword arguments apply to this client's requests only, so one + HTTP client can serve several APIs. Each one left ``None`` falls back + to the HTTP client's own setting. + + Args: + http_client: HTTP client to send requests with. + base_url: Base URL that request paths are appended to. + auth: Authentication for each request. + headers: Headers for each request, added to the HTTP client's own. + timeout: Timeouts for each request. + """ self._http_client = http_client + self._base_url = base_url + self._auth = auth + self._headers = headers + self._timeout = timeout async def cleanup_images( self, *, params: CleanupImagesParams | None = None - ) -> CleanupImagesResponse: + ) -> CleanupSandboxImagesResponse: """Clean up unused images""" return await self._do_json( - CleanupImagesResponse, + CleanupSandboxImagesResponse, _ApiRequest( method="POST", path_fmt="/v1/sandboxes/cleanup_images", @@ -4452,10 +4592,10 @@ async def delete_gateway_endpoints(self, *, endpoint_id: str) -> EndpointTombsto async def delete_image( self, *, image_name: str, params: DeleteImageParams | None = None - ) -> Image: + ) -> SandboxImage: """Delete a sandbox image""" return await self._do_json( - Image, + SandboxImage, _ApiRequest( method="DELETE", path_fmt="/v1/sandboxes/images/{}", @@ -4473,10 +4613,10 @@ async def delete_image_tag( image_name: str, tag_name: str, params: DeleteImageTagParams | None = None, - ) -> Image: + ) -> SandboxImage: """Delete an image tag""" return await self._do_json( - Image, + SandboxImage, _ApiRequest( method="DELETE", path_fmt="/v1/sandboxes/images/{}/tags/{}", @@ -4603,18 +4743,18 @@ async def delete_routes(self, *, route_id: str) -> RouteTombstone: ), ) - async def delete_routes_harness_configs( - self, *, harness: str, params: DeleteRoutesHarnessConfigsParams - ) -> RouteHarnessConfigTombstone: - """Clears default models for a coding harness""" + async def delete_routes_connections( + self, *, connection_id: str + ) -> RouteConnectionTombstone: + """Deletes a connection""" return await self._do_json( - RouteHarnessConfigTombstone, + RouteConnectionTombstone, _ApiRequest( method="DELETE", - path_fmt="/v1/routes/harness-configs/{}", - path_args=[harness], + path_fmt="/v1/routes/connections/{}", + path_args=[connection_id], body=None, - query=params, + query=None, success_codes=[200], error_codes=None, ), @@ -5150,10 +5290,10 @@ async def get_gateway_groups_group_id(self, *, group_id: str) -> Group: async def get_image( self, *, image_name: str, params: GetImageParams | None = None - ) -> Image: + ) -> SandboxImage: """Get a sandbox image""" return await self._do_json( - Image, + SandboxImage, _ApiRequest( method="GET", path_fmt="/v1/sandboxes/images/{}", @@ -5167,10 +5307,10 @@ async def get_image( async def get_image_build_logs( self, *, image_name: str, params: GetImageBuildLogsParams | None = None - ) -> ImageBuildLogsResponse: + ) -> SandboxImageBuildLogsResponse: """Get image build logs""" return await self._do_json( - ImageBuildLogsResponse, + SandboxImageBuildLogsResponse, _ApiRequest( method="GET", path_fmt="/v1/sandboxes/images/{}/logs", @@ -5894,15 +6034,15 @@ async def get_routes( ), ) - async def get_routes_harness_configs( - self, *, params: GetRoutesHarnessConfigsParams - ) -> RouteHarnessConfigsResponse: - """Lists default models for coding harnesses""" + async def get_routes_connections( + self, *, params: GetRoutesConnectionsParams | None = None + ) -> RouteConnectionsResponse: + """Lists connections""" return await self._do_json( - RouteHarnessConfigsResponse, + RouteConnectionsResponse, _ApiRequest( method="GET", - path_fmt="/v1/routes/harness-configs", + path_fmt="/v1/routes/connections", path_args=[], body=None, query=params, @@ -5926,13 +6066,28 @@ async def get_routes_route_id(self, *, route_id: str) -> Route: ), ) - async def get_routes_spend_limits(self, *, user_id: str) -> RouteSpendLimit: - """Gets a user's spend limits""" + async def get_routes_settings_teams(self, *, team_id: str) -> RouteTeamSettings: + """Gets a team's route settings""" + return await self._do_json( + RouteTeamSettings, + _ApiRequest( + method="GET", + path_fmt="/v1/routes/settings/teams/{}", + path_args=[team_id], + body=None, + query=None, + success_codes=[200], + error_codes=None, + ), + ) + + async def get_routes_settings_users(self, *, user_id: str) -> RouteUserSettings: + """Gets a user's route settings""" return await self._do_json( - RouteSpendLimit, + RouteUserSettings, _ApiRequest( method="GET", - path_fmt="/v1/routes/spend_limits/{}", + path_fmt="/v1/routes/settings/users/{}", path_args=[user_id], body=None, query=None, @@ -5975,6 +6130,55 @@ async def get_sandbox( ), ) + async def get_sandbox_configuration(self) -> GetSandboxConfigurationResponse: + """Get sandbox configuration""" + return await self._do_json( + GetSandboxConfigurationResponse, + _ApiRequest( + method="GET", + path_fmt="/v1/sandboxes/configuration", + path_args=[], + body=None, + query=None, + success_codes=[200], + error_codes=None, + ), + ) + + async def get_sandbox_logs( + self, *, sandbox_name: str, params: GetSandboxLogsParams + ) -> SandboxLogs: + """Get sandbox logs""" + return await self._do_json( + SandboxLogs, + _ApiRequest( + method="GET", + path_fmt="/v1/sandboxes/instances/{}/logs", + path_args=[sandbox_name], + body=None, + query=params, + success_codes=[200], + error_codes=None, + ), + ) + + async def get_sandbox_metrics( + self, *, sandbox_name: str, params: GetSandboxMetricsParams + ) -> SandboxMetrics: + """Get sandbox metrics""" + return await self._do_json( + SandboxMetrics, + _ApiRequest( + method="GET", + path_fmt="/v1/sandboxes/instances/{}/metrics", + path_args=[sandbox_name], + body=None, + query=params, + success_codes=[200], + error_codes=None, + ), + ) + async def get_secrets(self) -> Secrets: """Gets all secrets (metadata only, no plain text keys)""" return await self._do_json( @@ -6530,10 +6734,10 @@ async def get_volumes_volume_name( async def list_image_tags( self, *, image_name: str, params: ListImageTagsParams | None = None - ) -> ListImageTagsResponse: + ) -> ListSandboxImageTagsResponse: """List image tags""" return await self._do_json( - ListImageTagsResponse, + ListSandboxImageTagsResponse, _ApiRequest( method="GET", path_fmt="/v1/sandboxes/images/{}/tags", @@ -6547,10 +6751,10 @@ async def list_image_tags( async def list_images( self, *, params: ListImagesParams | None = None - ) -> ListImagesResponse: + ) -> ListSandboxImagesResponse: """List sandbox images""" return await self._do_json( - ListImagesResponse, + ListSandboxImagesResponse, _ApiRequest( method="GET", path_fmt="/v1/sandboxes/images", @@ -6562,9 +6766,7 @@ async def list_images( ), ) - async def list_sandbox_library_images( - self, *, params: ListSandboxLibraryImagesParams | None = None - ) -> ListSandboxLibraryImagesResponse: + async def list_sandbox_library_images(self) -> ListSandboxLibraryImagesResponse: """List built-in sandbox images""" return await self._do_json( ListSandboxLibraryImagesResponse, @@ -6573,7 +6775,7 @@ async def list_sandbox_library_images( path_fmt="/v1/sandboxes/library_images", path_args=[], body=None, - query=params, + query=None, success_codes=[200], error_codes=None, ), @@ -6866,7 +7068,7 @@ async def patch_models_environments( async def patch_routes( self, *, route_id: str, request: UpdateRouteRequest ) -> Route: - """Updates a route's display name or description""" + """Updates a route""" return await self._do_json( Route, _ApiRequest( @@ -6880,16 +7082,33 @@ async def patch_routes( ), ) - async def patch_routes_harness_configs( - self, *, request: UpdateRouteHarnessConfigRequest - ) -> RouteHarnessConfig: - """Updates default models for a coding harness""" + async def patch_routes_connections( + self, *, connection_id: str, request: UpdateRouteConnectionRequest + ) -> RouteConnection: + """Updates a connection""" return await self._do_json( - RouteHarnessConfig, + RouteConnection, _ApiRequest( method="PATCH", - path_fmt="/v1/routes/harness-configs", - path_args=[], + path_fmt="/v1/routes/connections/{}", + path_args=[connection_id], + body=request, + query=None, + success_codes=[200], + error_codes=None, + ), + ) + + async def patch_routes_settings_teams( + self, *, team_id: str, request: UpdateRouteTeamSettingsRequest + ) -> RouteTeamSettings: + """Updates a team's route settings""" + return await self._do_json( + RouteTeamSettings, + _ApiRequest( + method="PATCH", + path_fmt="/v1/routes/settings/teams/{}", + path_args=[team_id], body=request, query=None, success_codes=[200], @@ -6897,15 +7116,15 @@ async def patch_routes_harness_configs( ), ) - async def patch_routes_spend_limits( - self, *, user_id: str, request: UpdateRouteSpendLimitRequest - ) -> RouteSpendLimit: - """Updates a user's spend limits""" + async def patch_routes_settings_users( + self, *, user_id: str, request: UpdateRouteUserSettingsRequest + ) -> RouteUserSettings: + """Updates a user's route settings""" return await self._do_json( - RouteSpendLimit, + RouteUserSettings, _ApiRequest( method="PATCH", - path_fmt="/v1/routes/spend_limits/{}", + path_fmt="/v1/routes/settings/users/{}", path_args=[user_id], body=request, query=None, @@ -7367,7 +7586,7 @@ async def post_loops_sessions(self) -> CreateLoopsSessionResponse: ) async def post_loops_trainers( - self, *, request: CreateLoopsRunRequest + self, *, request: CreateLoopsTrainerRequest ) -> CreateLoopsRunResponse: """Creates a Loops trainer""" return await self._do_json( @@ -7865,6 +8084,23 @@ async def post_routes(self, *, request: CreateRouteRequest) -> Route: ), ) + async def post_routes_connections( + self, *, request: CreateRouteConnectionRequest + ) -> RouteConnection: + """Creates a connection""" + return await self._do_json( + RouteConnection, + _ApiRequest( + method="POST", + path_fmt="/v1/routes/connections", + path_args=[], + body=request, + query=None, + success_codes=[200], + error_codes=None, + ), + ) + async def post_secrets(self, *, request: UpsertSecretRequest) -> Secret: """Upserts a secret""" return await self._do_json( @@ -7966,7 +8202,7 @@ async def post_teams_loops_sessions( ) async def post_teams_loops_trainers( - self, *, team_id: str, request: CreateLoopsRunRequest + self, *, team_id: str, request: CreateLoopsTrainerRequest ) -> CreateLoopsRunResponse: """Creates a Loops trainer""" return await self._do_json( @@ -8272,11 +8508,11 @@ async def post_volumes_versions_restore( ) async def push_image( - self, *, params: PushImageParams | None = None, request: PushImageRequest - ) -> PushImageResponse: + self, *, params: PushImageParams | None = None, request: PushSandboxImageRequest + ) -> PushSandboxImageResponse: """Push a sandbox image""" return await self._do_json( - PushImageResponse, + PushSandboxImageResponse, _ApiRequest( method="POST", path_fmt="/v1/sandboxes/images", @@ -8288,23 +8524,6 @@ async def push_image( ), ) - async def put_routes_harness_configs( - self, *, request: SetRouteHarnessConfigRequest - ) -> RouteHarnessConfig: - """Sets default models for a coding harness""" - return await self._do_json( - RouteHarnessConfig, - _ApiRequest( - method="PUT", - path_fmt="/v1/routes/harness-configs", - path_args=[], - body=request, - query=None, - success_codes=[200], - error_codes=None, - ), - ) - async def update_sandbox( self, *, @@ -8333,7 +8552,7 @@ def _build_request(self, request: _ApiRequest) -> httpx.Request: json_body = None content_body = None files_body = None - headers: dict[str, str] = {} + headers: dict[str, str] = {**(self._headers or {})} if request.accept is not None: headers["Accept"] = request.accept if request.body is not None: @@ -8367,16 +8586,22 @@ def _build_request(self, request: _ApiRequest) -> httpx.Request: params = request.query return self._http_client.build_request( request.method, - path, + path if self._base_url is None else self._base_url.rstrip("/") + path, json=json_body, content=content_body, files=files_body, params=params, headers=headers, + timeout=httpx.USE_CLIENT_DEFAULT + if self._timeout is None + else self._timeout, ) async def _do(self, request: _ApiRequest) -> httpx.Response: - response = await self._http_client.send(self._build_request(request)) + response = await self._http_client.send( + self._build_request(request), + auth=httpx.USE_CLIENT_DEFAULT if self._auth is None else self._auth, + ) if response.status_code not in request.success_codes: raise ResponseError(status_code=response.status_code, body=response.text) return response diff --git a/baseten/client/managementapi/_models.py b/baseten/client/managementapi/_models.py index df85624..e662f5e 100644 --- a/baseten/client/managementapi/_models.py +++ b/baseten/client/managementapi/_models.py @@ -789,8 +789,7 @@ class GetRoutesParams(BaseModel): team_id: Annotated[ str | None, Field( - description="Filter by owning team ID. Preserved by the cursor; if repeated, must match the original filter.", - title="Team Id", + description="Identifier of the team whose routes to list.", title="Team Id" ), ] = None name: Annotated[ @@ -802,26 +801,64 @@ class GetRoutesParams(BaseModel): ] = None -class GetRoutesHarnessConfigsParams(BaseModel): +class GetRoutesConnectionsParams(BaseModel): + cursor: Annotated[ + str | None, + Field( + description="Opaque cursor returned by a previous page. Omit to fetch the first page.", + title="Cursor", + ), + ] = None + limit: Annotated[ + int, + Field( + description="Maximum number of items to return.", + ge=1, + le=1000, + title="Limit", + ), + ] = 100 team_id: Annotated[ - str, + str | None, Field( - description="Identifier of the team whose default models to list.", - min_length=1, + description="Identifier of the team whose connections to list.", title="Team Id", ), - ] + ] = None -class DeleteRoutesHarnessConfigsParams(BaseModel): +class GetSandboxMetricsParams(BaseModel): team_id: Annotated[ - str, + str | None, Field( - description="Identifier of the team whose default models to clear.", + description="Optional team ID. Must match X-Team-Id when both are supplied. If neither selector is supplied, defaults to the caller's only accessible team. Callers with multiple accessible teams must select a team. Requests without access to any team are forbidden.", min_length=1, - title="Team Id", ), - ] + ] = None + start_time: AwareDatetime + end_time: AwareDatetime + interval_seconds: Literal[10, 30, 60, 300, 900, 3600] | str + + +class GetSandboxLogsParams(BaseModel): + team_id: Annotated[ + str | None, + Field( + description="Optional team ID. Must match X-Team-Id when both are supplied. If neither selector is supplied, defaults to the caller's only accessible team. Callers with multiple accessible teams must select a team. Requests without access to any team are forbidden.", + min_length=1, + ), + ] = None + start_time: AwareDatetime + end_time: AwareDatetime + limit: Annotated[int, Field(ge=1, le=1000)] = 100 + cursor: Annotated[ + str | None, + Field( + description="Opaque next_cursor from the previous page. Omit for the first page.", + max_length=4096, + min_length=1, + ), + ] = None class StatusItem(RootModel[str]): @@ -924,16 +961,6 @@ class DeleteSandboxParams(BaseModel): ] = None -class ListSandboxLibraryImagesParams(BaseModel): - team_id: Annotated[ - str | None, - Field( - description="Optional team ID. Must match X-Team-Id when both are supplied. If neither selector is supplied, defaults to the caller's only accessible team. Callers with multiple accessible teams must select a team. Requests without access to any team are forbidden.", - min_length=1, - ), - ] = None - - class ListImagesParams(BaseModel): team_id: Annotated[ str | None, @@ -1260,9 +1287,6 @@ class ListVolumeNamespacesResponse(BaseModel): class VolumeSyncAuthenticationAWSAssumeRole(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) role_arn: Annotated[ str, Field( @@ -1280,9 +1304,6 @@ class VolumeSyncAuthenticationAWSAssumeRole(BaseModel): class VolumeSyncAuthenticationAWSOIDC(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) role_arn: Annotated[ str, Field( @@ -1302,9 +1323,6 @@ class VolumeSyncAuthenticationAWSOIDC(BaseModel): class VolumeSyncAuthenticationGCPOIDC(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) service_account: Annotated[ str, Field( @@ -1324,9 +1342,6 @@ class VolumeSyncAuthenticationGCPOIDC(BaseModel): class VolumeSyncDestination(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) ref: Annotated[ str, Field( @@ -1360,9 +1375,6 @@ class AuthSecretName(RootModel[str]): class VolumeSyncSourceAzure(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) uri: Annotated[ str, Field( @@ -1392,9 +1404,6 @@ class VolumeSyncSourceAzure(BaseModel): class VolumeSyncSourceBasetenTraining(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) uri: Annotated[ str, Field( @@ -1417,9 +1426,6 @@ class VolumeSyncSourceBasetenTraining(BaseModel): class VolumeSyncSourceCoreWeave(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) uri: Annotated[ str, Field( @@ -1449,9 +1455,6 @@ class VolumeSyncSourceCoreWeave(BaseModel): class VolumeSyncSourceGCS(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) uri: Annotated[ str, Field( @@ -1487,9 +1490,6 @@ class VolumeSyncSourceGCS(BaseModel): class VolumeSyncSourceHuggingFace(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) uri: Annotated[ str, Field( @@ -1519,9 +1519,6 @@ class VolumeSyncSourceHuggingFace(BaseModel): class VolumeSyncSourceR2(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) uri: Annotated[ str, Field( @@ -1550,9 +1547,6 @@ class VolumeSyncSourceR2(BaseModel): class VolumeSyncSourceS3(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) uri: Annotated[ str, Field( @@ -2513,9 +2507,6 @@ class AuditLogActor(BaseModel): class AuditLogEventApiKeyCreated(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[Literal["API_KEY_CREATED"], Field(title="Event Type")] api_key_id: Annotated[str, Field(title="Api Key Id")] api_key_type: Annotated[ @@ -2542,9 +2533,6 @@ class AuditLogEventApiKeyCreated(BaseModel): class AuditLogEventApiKeyDeleted(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[Literal["API_KEY_DELETED"], Field(title="Event Type")] api_key_id: Annotated[str, Field(title="Api Key Id")] api_key_type: Annotated[ @@ -2571,9 +2559,6 @@ class AuditLogEventApiKeyDeleted(BaseModel): class AuditLogEventAutoscalingScheduleSettings(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) min_replica: Annotated[int, Field(title="Min Replica")] max_replica: Annotated[int, Field(title="Max Replica")] concurrency_target: Annotated[int | None, Field(title="Concurrency Target")] @@ -2600,9 +2585,6 @@ class AuditLogEventAutoscalingScheduleSettings(BaseModel): class AuditLogEventAutoscalingSettings(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) min_replica: Annotated[int, Field(title="Min Replica")] max_replica: Annotated[int, Field(title="Max Replica")] concurrency_target: Annotated[int, Field(title="Concurrency Target")] @@ -2618,9 +2600,6 @@ class AuditLogEventAutoscalingSettings(BaseModel): class AuditLogEventChainDeleted(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[Literal["CHAIN_DELETED"], Field(title="Event Type")] chain_id: Annotated[str, Field(title="Chain Id")] chain_name: Annotated[str, Field(title="Chain Name")] @@ -2628,9 +2607,6 @@ class AuditLogEventChainDeleted(BaseModel): class AuditLogEventChainDeployed(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[Literal["CHAIN_DEPLOYED"], Field(title="Event Type")] chain_id: Annotated[str, Field(title="Chain Id")] chain_name: Annotated[str, Field(title="Chain Name")] @@ -2641,9 +2617,6 @@ class AuditLogEventChainDeployed(BaseModel): class AuditLogEventChainDeploymentActivated(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[ Literal["CHAIN_DEPLOYMENT_ACTIVATED"], Field(title="Event Type") ] @@ -2654,9 +2627,6 @@ class AuditLogEventChainDeploymentActivated(BaseModel): class AuditLogEventChainDeploymentDeactivated(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[ Literal["CHAIN_DEPLOYMENT_DEACTIVATED"], Field(title="Event Type") ] @@ -2667,9 +2637,6 @@ class AuditLogEventChainDeploymentDeactivated(BaseModel): class AuditLogEventChainDeploymentDeleted(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[ Literal["CHAIN_DEPLOYMENT_DELETED"], Field(title="Event Type") ] @@ -2680,9 +2647,6 @@ class AuditLogEventChainDeploymentDeleted(BaseModel): class AuditLogEventChainDeploymentPromoted(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[ Literal["CHAIN_DEPLOYMENT_PROMOTED"], Field(title="Event Type") ] @@ -2695,9 +2659,6 @@ class AuditLogEventChainDeploymentPromoted(BaseModel): class AuditLogEventChainEnvironmentCreated(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[ Literal["CHAIN_ENVIRONMENT_CREATED"], Field(title="Event Type") ] @@ -2714,9 +2675,6 @@ class AuditLogEventChainEnvironmentCreated(BaseModel): class AuditLogEventChainEnvironmentUpdated(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[ Literal["CHAIN_ENVIRONMENT_UPDATED"], Field(title="Event Type") ] @@ -2733,9 +2691,6 @@ class AuditLogEventChainEnvironmentUpdated(BaseModel): class AuditLogEventChainletAutoscalingSettingsChanged(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) min_replica: Annotated[int, Field(title="Min Replica")] max_replica: Annotated[int, Field(title="Max Replica")] concurrency_target: Annotated[int, Field(title="Concurrency Target")] @@ -2761,9 +2716,6 @@ class AuditLogEventChainletAutoscalingSettingsChanged(BaseModel): class AuditLogEventChainletInstanceTypeChanged(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[ Literal["CHAINLET_INSTANCE_TYPE_CHANGED"], Field(title="Event Type") ] @@ -2777,9 +2729,6 @@ class AuditLogEventChainletInstanceTypeChanged(BaseModel): class AuditLogEventDirectoryGroupRoleUpdated(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[ Literal["DIRECTORY_GROUP_ROLE_UPDATED"], Field(title="Event Type") ] @@ -2791,9 +2740,6 @@ class AuditLogEventDirectoryGroupRoleUpdated(BaseModel): class AuditLogEventEnvironmentCreated(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) min_replica: Annotated[int, Field(title="Min Replica")] max_replica: Annotated[int, Field(title="Max Replica")] concurrency_target: Annotated[int, Field(title="Concurrency Target")] @@ -2842,9 +2788,6 @@ class AuditLogEventEnvironmentCreated(BaseModel): class AuditLogEventEnvironmentDeleted(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[Literal["ENVIRONMENT_DELETED"], Field(title="Event Type")] model_id: Annotated[str, Field(title="Model Id")] model_name: Annotated[str, Field(title="Model Name")] @@ -2852,9 +2795,6 @@ class AuditLogEventEnvironmentDeleted(BaseModel): class AuditLogEventEnvironmentSettings(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) min_replica: Annotated[int, Field(title="Min Replica")] max_replica: Annotated[int, Field(title="Max Replica")] concurrency_target: Annotated[int, Field(title="Concurrency Target")] @@ -2898,9 +2838,6 @@ class AuditLogEventEnvironmentSettings(BaseModel): class AuditLogEventGatewayEndpointCreated(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[ Literal["GATEWAY_ENDPOINT_CREATED"], Field(title="Event Type") ] @@ -2909,9 +2846,6 @@ class AuditLogEventGatewayEndpointCreated(BaseModel): class AuditLogEventGatewayEndpointDeleted(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[ Literal["GATEWAY_ENDPOINT_DELETED"], Field(title="Event Type") ] @@ -2920,9 +2854,6 @@ class AuditLogEventGatewayEndpointDeleted(BaseModel): class AuditLogEventGatewayEndpointUpdated(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[ Literal["GATEWAY_ENDPOINT_UPDATED"], Field(title="Event Type") ] @@ -2932,18 +2863,12 @@ class AuditLogEventGatewayEndpointUpdated(BaseModel): class AuditLogEventModelDeleted(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[Literal["MODEL_DELETED"], Field(title="Event Type")] model_id: Annotated[str, Field(title="Model Id")] model_name: Annotated[str, Field(title="Model Name")] class AuditLogEventModelDeployed(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[Literal["MODEL_DEPLOYED"], Field(title="Event Type")] model_id: Annotated[str, Field(title="Model Id")] model_name: Annotated[str, Field(title="Model Name")] @@ -2956,9 +2881,6 @@ class AuditLogEventModelDeployed(BaseModel): class AuditLogEventModelDeploymentActivated(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[ Literal["MODEL_DEPLOYMENT_ACTIVATED"], Field(title="Event Type") ] @@ -2969,9 +2891,6 @@ class AuditLogEventModelDeploymentActivated(BaseModel): class AuditLogEventModelDeploymentDeactivated(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[ Literal["MODEL_DEPLOYMENT_DEACTIVATED"], Field(title="Event Type") ] @@ -2982,9 +2901,6 @@ class AuditLogEventModelDeploymentDeactivated(BaseModel): class AuditLogEventModelDeploymentDeleted(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[ Literal["MODEL_DEPLOYMENT_DELETED"], Field(title="Event Type") ] @@ -2995,9 +2911,6 @@ class AuditLogEventModelDeploymentDeleted(BaseModel): class AuditLogEventModelDeploymentInstanceTypeChanged(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[ Literal["MODEL_DEPLOYMENT_INSTANCE_TYPE_CHANGED"], Field(title="Event Type") ] @@ -3009,9 +2922,6 @@ class AuditLogEventModelDeploymentInstanceTypeChanged(BaseModel): class AuditLogEventModelDeploymentPromoted(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[ Literal["MODEL_DEPLOYMENT_PROMOTED"], Field(title="Event Type") ] @@ -3024,9 +2934,6 @@ class AuditLogEventModelDeploymentPromoted(BaseModel): class AuditLogEventModelDeploymentRequestBackpressureSettingsChanged(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[ Literal["MODEL_DEPLOYMENT_REQUEST_BACKPRESSURE_SETTINGS_CHANGED"], Field(title="Event Type"), @@ -3040,9 +2947,6 @@ class AuditLogEventModelDeploymentRequestBackpressureSettingsChanged(BaseModel): class AuditLogEventModelDeploymentRetried(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[ Literal["MODEL_DEPLOYMENT_RETRIED"], Field(title="Event Type") ] @@ -3054,9 +2958,6 @@ class AuditLogEventModelDeploymentRetried(BaseModel): class AuditLogEventModelRenamed(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[Literal["MODEL_RENAMED"], Field(title="Event Type")] model_id: Annotated[str, Field(title="Model Id")] model_name: Annotated[str, Field(title="Model Name")] @@ -3064,9 +2965,6 @@ class AuditLogEventModelRenamed(BaseModel): class AuditLogEventProviderConnectionCreated(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[ Literal["PROVIDER_CONNECTION_CREATED"], Field(title="Event Type") ] @@ -3076,9 +2974,6 @@ class AuditLogEventProviderConnectionCreated(BaseModel): class AuditLogEventProviderConnectionDeleted(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[ Literal["PROVIDER_CONNECTION_DELETED"], Field(title="Event Type") ] @@ -3088,9 +2983,6 @@ class AuditLogEventProviderConnectionDeleted(BaseModel): class AuditLogEventProviderConnectionUpdated(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[ Literal["PROVIDER_CONNECTION_UPDATED"], Field(title="Event Type") ] @@ -3100,9 +2992,6 @@ class AuditLogEventProviderConnectionUpdated(BaseModel): class AuditLogEventReplicaTerminated(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[Literal["REPLICA_TERMINATED"], Field(title="Event Type")] model_id: Annotated[str, Field(title="Model Id")] model_name: Annotated[str, Field(title="Model Name")] @@ -3112,9 +3001,6 @@ class AuditLogEventReplicaTerminated(BaseModel): class AuditLogEventRequireGroupBasedAdminsEnabled(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[ Literal["REQUIRE_GROUP_BASED_ADMINS_ENABLED"], Field(title="Event Type") ] @@ -3122,27 +3008,18 @@ class AuditLogEventRequireGroupBasedAdminsEnabled(BaseModel): class AuditLogEventSecretDeleted(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[Literal["SECRET_DELETED"], Field(title="Event Type")] secret_id: Annotated[str, Field(title="Secret Id")] secret_name: Annotated[str, Field(title="Secret Name")] class AuditLogEventSecretUpdated(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[Literal["SECRET_UPDATED"], Field(title="Event Type")] secret_id: Annotated[str, Field(title="Secret Id")] secret_name: Annotated[str, Field(title="Secret Name")] class AuditLogEventSshCertificateSigned(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[Literal["SSH_CERTIFICATE_SIGNED"], Field(title="Event Type")] workload_type: Annotated[str, Field(title="Workload Type")] workload_id: Annotated[str, Field(title="Workload Id")] @@ -3153,18 +3030,12 @@ class AuditLogEventSshCertificateSigned(BaseModel): class AuditLogEventUserInvited(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[Literal["USER_INVITED"], Field(title="Event Type")] invited_user_email: Annotated[str, Field(title="Invited User Email")] role_name: Annotated[str, Field(title="Role Name")] class AuditLogEventUserJoinedOrganization(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[ Literal["USER_JOINED_ORGANIZATION"], Field(title="Event Type") ] @@ -3173,17 +3044,11 @@ class AuditLogEventUserJoinedOrganization(BaseModel): class AuditLogEventUserRemoved(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[Literal["USER_REMOVED"], Field(title="Event Type")] removed_user_email: Annotated[str, Field(title="Removed User Email")] class AuditLogEventUserRoleUpdated(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[Literal["USER_ROLE_UPDATED"], Field(title="Event Type")] user_id: Annotated[str, Field(title="User Id")] user_email: Annotated[str, Field(title="User Email")] @@ -3191,9 +3056,6 @@ class AuditLogEventUserRoleUpdated(BaseModel): class AuditLogEventUserTeamRoleUpdated(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[Literal["USER_TEAM_ROLE_UPDATED"], Field(title="Event Type")] user_id: Annotated[str, Field(title="User Id")] user_email: Annotated[str, Field(title="User Email")] @@ -3203,9 +3065,6 @@ class AuditLogEventUserTeamRoleUpdated(BaseModel): class AuditLogEventVolumeDeleted(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[Literal["VOLUME_DELETED"], Field(title="Event Type")] volume_ref: Annotated[str, Field(title="Volume Ref")] namespace: Annotated[str, Field(title="Namespace")] @@ -3214,9 +3073,6 @@ class AuditLogEventVolumeDeleted(BaseModel): class AuditLogEventVolumeVersionDeleted(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[Literal["VOLUME_VERSION_DELETED"], Field(title="Event Type")] volume_ref: Annotated[str, Field(title="Volume Ref")] namespace: Annotated[str, Field(title="Namespace")] @@ -3226,9 +3082,6 @@ class AuditLogEventVolumeVersionDeleted(BaseModel): class AuditLogEventVolumeVersionRestored(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[Literal["VOLUME_VERSION_RESTORED"], Field(title="Event Type")] volume_ref: Annotated[str, Field(title="Volume Ref")] namespace: Annotated[str, Field(title="Namespace")] @@ -3238,9 +3091,6 @@ class AuditLogEventVolumeVersionRestored(BaseModel): class AuditLogEventWebhookSigningSecretCreated(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[ Literal["WEBHOOK_SIGNING_SECRET_CREATED"], Field(title="Event Type") ] @@ -3248,9 +3098,6 @@ class AuditLogEventWebhookSigningSecretCreated(BaseModel): class AuditLogEventWebhookSigningSecretDeleted(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[ Literal["WEBHOOK_SIGNING_SECRET_DELETED"], Field(title="Event Type") ] @@ -3258,9 +3105,6 @@ class AuditLogEventWebhookSigningSecretDeleted(BaseModel): class AuditLogEventWebhookSigningSecretRotated(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[ Literal["WEBHOOK_SIGNING_SECRET_ROTATED"], Field(title="Event Type") ] @@ -5764,6 +5608,17 @@ class Name3(RootModel[str]): ] = None +class SamplerNumReplicas(RootModel[int]): + root: Annotated[ + int | None, + Field( + description="Number of replicas the run's sampler runs, applied as both its minimum and maximum. Must be at least 1. If omitted, a new sampler uses the platform defaults and a sampler reused from an earlier run keeps its settings. A run that already has a sampler keeps it unchanged. When the run ends, its sampler is scaled down.", + ge=1, + title="Sampler Num Replicas", + ), + ] = None + + class CreateLoopsRunRequest(BaseModel): session_id: Annotated[ str, @@ -5843,6 +5698,13 @@ class CreateLoopsRunRequest(BaseModel): title="Reuse From Session Id", ), ] = None + sampler_num_replicas: Annotated[ + SamplerNumReplicas | None, + Field( + description="Number of replicas the run's sampler runs, applied as both its minimum and maximum. Must be at least 1. If omitted, a new sampler uses the platform defaults and a sampler reused from an earlier run keeps its settings. A run that already has a sampler keeps it unchanged. When the run ends, its sampler is scaled down.", + title="Sampler Num Replicas", + ), + ] = None class DeactivateLoopsRunResponse(BaseModel): @@ -5857,12 +5719,104 @@ class DeactivateLoopsRunResponse(BaseModel): user: Annotated[User, Field(description="The user who owns the Loops run.")] -class ListLoopsSamplersResponse(BaseModel): - samplers: Annotated[ - list[LoopsSampler], Field(description="List of samplers.", title="Samplers") +class CreateLoopsTrainerRequest(BaseModel): + session_id: Annotated[ + str, + Field( + description="ID of the Loops session this run belongs to.", + title="Session Id", + ), ] + base_model: Annotated[ + str, + Field(description="Base model ID (e.g. 'Qwen/Qwen3-8B').", title="Base Model"), + ] + name: Annotated[ + Name3 | None, + Field( + description="Optional display name for the run. Defaults to the base model name when omitted.", + title="Name", + ), + ] = None + max_seq_len: Annotated[ + int | None, + Field( + description="Maximum prompt length (in tokens) the run must handle. Set this to the longest training example you plan to send. Defaults to the maximum supported by the model configuration.", + title="Max Seq Len", + ), + ] = None + lora_rank: Annotated[ + int, Field(description="LoRA rank.", ge=1, title="Lora Rank") + ] = 64 + seed: Annotated[ + int | None, Field(description="Random seed for reproducibility.", title="Seed") + ] = None + scale_down_delay_seconds: Annotated[ + int, + Field( + description="Seconds of inactivity before the run scales to zero. Must be between 1 and 3600 (1 hour). Defaults to 900 (15 minutes).", + gt=0, + le=3600, + title="Scale Down Delay Seconds", + ), + ] = 900 + availability_model: Annotated[ + Literal["dedicated", "spot"] | str, + Field( + description="Capacity the trainer runs on. 'dedicated' is not preempted. 'spot' runs below inference and reaches idle reserved capacity, but the run is stopped if its GPUs are reclaimed and cannot be resumed.", + examples=["spot"], + title="V1AvailabilityModel", + ), + ] = "dedicated" + replicas: Annotated[ + int, + Field( + description="Number of data-parallel trainer replicas. Each replica is one full copy of the model's preset node group, so the trainer deployment runs (preset node_count * replicas) nodes (e.g. replicas=4 on a 4-node preset → 16 nodes, 4 DP workers). Must be a positive integer. Defaults to 1.", + ge=1, + title="Replicas", + ), + ] = 1 + path: Annotated[ + str | None, + Field( + description="Optional bt:// URI of an existing checkpoint to load weights from on startup. Form: bt://loops:/weights/.", + examples=["bt://loops:k4q95w5/weights/step-100"], + title="Path", + ), + ] = None + reuse_from_run_id: Annotated[ + str | None, + Field( + description="Optional ID of a prior Loops run whose trainer and/or sampler should be reused for this run instead of provisioning fresh. The prior run must use the same base model and belong to the same team.", + title="Reuse From Run Id", + ), + ] = None + reuse_from_session_id: Annotated[ + str | None, + Field( + description="Optional ID of a prior Loops session whose trainer and/or sampler should be reused for this run. Deprecated in favor of reuse_from_run_id.", + title="Reuse From Session Id", + ), + ] = None + + +class ListLoopsSamplersResponse(BaseModel): + samplers: Annotated[ + list[LoopsSampler], Field(description="List of samplers.", title="Samplers") + ] + + +class NumReplicas(RootModel[int]): + root: Annotated[ + int | None, + Field( + description="Number of replicas the sampler runs, applied as both its minimum and maximum. Must be at least 1. If omitted, a new sampler uses the platform defaults and a paired sampler reused from an earlier run keeps its settings. A run that already has a sampler keeps it unchanged, and when a run ends, its sampler is scaled down.", + ge=1, + title="Num Replicas", + ), + ] = None + - class CreateLoopsSamplerRequest(BaseModel): availability_model: Annotated[ Literal["dedicated", "spot"] | str | None, @@ -5913,6 +5867,13 @@ class CreateLoopsSamplerRequest(BaseModel): title="Reuse From Session Id", ), ] = None + num_replicas: Annotated[ + NumReplicas | None, + Field( + description="Number of replicas the sampler runs, applied as both its minimum and maximum. Must be at least 1. If omitted, a new sampler uses the platform defaults and a paired sampler reused from an earlier run keeps its settings. A run that already has a sampler keeps it unchanged, and when a run ends, its sampler is scaled down.", + title="Num Replicas", + ), + ] = None class CreateLoopsSamplerResponse(BaseModel): @@ -7040,9 +7001,6 @@ class CreateLLMModelVersionRequest(BaseModel): class LibraryListingMetadata(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) parameter_count: Annotated[int | None, Field(title="Parameter Count")] = None context_length: Annotated[int | None, Field(title="Context Length")] = None input_modalities: Annotated[ @@ -7175,7 +7133,7 @@ class UpdateLibraryListingRequest(BaseModel): metadata: Annotated[ LibraryListingMetadata | None, Field( - description="Model-level metadata for the listing. When provided, replaces the stored metadata. Unknown fields are rejected." + description="Model-level metadata for the listing. When provided, replaces the stored metadata." ), ] = None @@ -7825,7 +7783,125 @@ class ExploreMetadataAPIFormats(BaseModel): ] = False +class RouteRef(BaseModel): + id: Annotated[str, Field(description="Stable route identifier.", title="Id")] + slug: Annotated[ + str, + Field( + description="Name of the route to send in the inference request's model field.", + title="Slug", + ), + ] + + class RouteTargetAnthropic(BaseModel): + type: Annotated[ + Literal["ANTHROPIC"], + Field(description="Target kind for Anthropic.", title="Type"), + ] + model: Annotated[ + str, + Field( + description="Model name sent to the provider.", min_length=1, title="Model" + ), + ] + secret_name: Annotated[ + str, + Field( + description="Name of a credential secret owned by the route's team.", + min_length=1, + title="Secret Name", + ), + ] + + +class RouteTargetBasetenModelAPI(BaseModel): + type: Annotated[ + Literal["BASETEN_MODEL_API"], + Field(description="Target kind for a Baseten Model API.", title="Type"), + ] + model: Annotated[ + str, + Field(description="Name of the target Model API.", min_length=1, title="Model"), + ] + + +class RouteTargetClassifierModelBased(BaseModel): + type: Annotated[ + Literal["CLASSIFIER_MODEL_BASED"], + Field( + description="Target kind for a classifier-model-based route.", title="Type" + ), + ] + classifier_model_id: Annotated[ + str, + Field( + description="ID of the Baseten model that picks a route for each request.", + examples=["abcd123"], + title="Classifier Model Id", + ), + ] + classifier_environment_name: Annotated[ + str | None, + Field( + description="Environment of the classifier model. Null for production.", + examples=["production"], + title="Classifier Environment Name", + ), + ] + default_route: Annotated[ + RouteRef, Field(description="Route used when the classifier picks none.") + ] + allowed_routes: Annotated[ + list[RouteRef], + Field( + description="Routes the classifier may pick, in creation order.", + title="Allowed Routes", + ), + ] + + +class RouteTargetOpenAI(BaseModel): + type: Annotated[ + Literal["OPENAI"], Field(description="Target kind for OpenAI.", title="Type") + ] + model: Annotated[ + str, + Field( + description="Model name sent to the provider.", min_length=1, title="Model" + ), + ] + secret_name: Annotated[ + str, + Field( + description="Name of a credential secret owned by the route's team.", + min_length=1, + title="Secret Name", + ), + ] + + +class RouteTargetXAI(BaseModel): + type: Annotated[ + Literal["XAI"], Field(description="Target kind for xAI.", title="Type") + ] + model: Annotated[ + str, + Field( + description="Model name sent to the provider.", min_length=1, title="Model" + ), + ] + secret_name: Annotated[ + str, + Field( + description="Name of a credential secret owned by the route's team.", + min_length=1, + title="Secret Name", + ), + ] + + +class RouteTargetConfigAnthropic(BaseModel): model_config = ConfigDict( extra="forbid", ) @@ -7849,7 +7925,7 @@ class RouteTargetAnthropic(BaseModel): ] -class RouteTargetBasetenModelAPI(BaseModel): +class RouteTargetConfigBasetenModelAPI(BaseModel): model_config = ConfigDict( extra="forbid", ) @@ -7863,7 +7939,65 @@ class RouteTargetBasetenModelAPI(BaseModel): ] -class RouteTargetOpenAI(BaseModel): +class ClassifierEnvironmentName(RootModel[str]): + root: Annotated[ + str | None, + Field( + description="Environment of the classifier model. Omit for production, which is returned as null.", + examples=["production"], + min_length=1, + title="Classifier Environment Name", + ), + ] = None + + +class RouteTargetConfigClassifierModelBased(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + type: Annotated[ + Literal["CLASSIFIER_MODEL_BASED"], + Field( + description="Target kind for a classifier-model-based route. Not intended for general use: classifiers are deployed by Baseten's post-training team.", + title="Type", + ), + ] + classifier_model_id: Annotated[ + str, + Field( + description="ID of the Baseten model that picks a route for each request. Only classifiers deployed by Baseten's post-training team are supported.", + examples=["abcd123"], + min_length=1, + title="Classifier Model Id", + ), + ] + classifier_environment_name: Annotated[ + ClassifierEnvironmentName | None, + Field( + description="Environment of the classifier model. Omit for production, which is returned as null.", + examples=["production"], + title="Classifier Environment Name", + ), + ] = None + allowed_route_ids: Annotated[ + list[str], + Field( + description="IDs of the routes the classifier may pick. All must belong to the route's team and must not be routers themselves.", + examples=[["abc1234", "def5678"]], + title="Allowed Route Ids", + ), + ] + default_route_id: Annotated[ + str, + Field( + description="ID of the allowed route used when the classifier picks none.", + examples=["abc1234"], + title="Default Route Id", + ), + ] + + +class RouteTargetConfigOpenAI(BaseModel): model_config = ConfigDict( extra="forbid", ) @@ -7886,7 +8020,7 @@ class RouteTargetOpenAI(BaseModel): ] -class RouteTargetXAI(BaseModel): +class RouteTargetConfigXAI(BaseModel): model_config = ConfigDict( extra="forbid", ) @@ -7913,7 +8047,7 @@ class TeamId1(RootModel[str]): root: Annotated[ str | None, Field( - description="Identifier of the team that owns the route. When omitted, uses your organization's default team.", + description="Identifier of the team that owns the route.", examples=["abc1234"], min_length=1, title="Team Id", @@ -7953,7 +8087,7 @@ class CreateRouteRequest(BaseModel): team_id: Annotated[ TeamId1 | None, Field( - description="Identifier of the team that owns the route. When omitted, uses your organization's default team.", + description="Identifier of the team that owns the route.", examples=["abc1234"], title="Team Id", ), @@ -7968,10 +8102,11 @@ class CreateRouteRequest(BaseModel): ] = None target: Annotated[ Annotated[ - RouteTargetBasetenModelAPI - | RouteTargetAnthropic - | RouteTargetOpenAI - | RouteTargetXAI, + RouteTargetConfigBasetenModelAPI + | RouteTargetConfigAnthropic + | RouteTargetConfigOpenAI + | RouteTargetConfigXAI + | RouteTargetConfigClassifierModelBased, Field( description="Upstream target for the route.", discriminator="type", @@ -8056,47 +8191,67 @@ class RoutesUsageResult(BaseModel): ] -class RouteSpendLimit(BaseModel): - user_id: Annotated[ - str, Field(description="ID of the user.", examples=["abc1234"], title="User Id") +class RouteSpendLimitTeamDefault(BaseModel): + team_id: Annotated[ + str, + Field( + description="ID of the team the user's active Code key belongs to.", + examples=["abc1234"], + title="Team Id", + ), ] - email: Annotated[ - str | None, + per_member_monthly_limit_usd: Annotated[ + str, Field( - description="Email address of the user.", - examples=["dev@example.com"], - title="Email", + description="The team's per-member spend limit in USD for each UTC calendar month.", + examples=["100"], + title="Per Member Monthly Limit Usd", ), ] - monthly_limit_usd: Annotated[ + + +class UpdateRouteSpendLimitSetting(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + user_monthly_limit_usd: Annotated[ str | None, Field( - description="Standing spend limit in USD for each UTC calendar month, returned as an exact decimal string. Null when the user has no standing limit.", + description="Standing spend limit in USD for each UTC calendar month. Send null to remove the limit; omit to leave it unchanged.", examples=["200"], - title="Monthly Limit Usd", + title="User Monthly Limit Usd", ), - ] + ] = None -class UpdateRouteSpendLimitRequest(BaseModel): +class UpdateRouteUserSettingsRequest(BaseModel): model_config = ConfigDict( extra="forbid", ) - monthly_limit_usd: Annotated[ + spend_limit: Annotated[ + UpdateRouteSpendLimitSetting | None, + Field( + description="Spend limit fields to change. Pass null to remove the user's limit; omit to leave it unchanged." + ), + ] = None + + +class RouteTeamSpendLimitSetting(BaseModel): + per_member_monthly_limit_usd: Annotated[ str | None, Field( - description="Standing spend limit in USD for each UTC calendar month, as a non-negative decimal string with at most 9 decimal places. Send null to remove the limit; omit to leave it unchanged.", + description="Spend limit in USD for each UTC calendar month that applies to each member whose active Code key belongs to this team, unless the member has a limit of their own. Returned as an exact decimal string. Null when the team has no per-member limit.", examples=["200"], - title="Monthly Limit Usd", + title="Per Member Monthly Limit Usd", ), - ] = None + ] class Primary(RootModel[str]): root: Annotated[ str | None, Field( - description="Route ID for the primary model, which new sessions use. Omit to keep the current route, or pass null to use Baseten's default.", + description="Route ID for the primary model, which new sessions use. Omit to keep the current route, or pass null to use the team's default.", examples=["abc1234"], min_length=1, title="Primary", @@ -8108,7 +8263,7 @@ class Background(RootModel[str]): root: Annotated[ str | None, Field( - description="Route ID for background tasks, such as session titles. Omit to keep the current route, or pass null to use Baseten's default.", + description="Route ID for background tasks, such as session titles. Omit to keep the current route, or pass null to use the team's default.", examples=["def5678"], min_length=1, title="Background", @@ -8123,7 +8278,7 @@ class UpdateBackgroundHarnessModels(BaseModel): primary: Annotated[ Primary | None, Field( - description="Route ID for the primary model, which new sessions use. Omit to keep the current route, or pass null to use Baseten's default.", + description="Route ID for the primary model, which new sessions use. Omit to keep the current route, or pass null to use the team's default.", examples=["abc1234"], title="Primary", ), @@ -8131,256 +8286,334 @@ class UpdateBackgroundHarnessModels(BaseModel): background: Annotated[ Background | None, Field( - description="Route ID for background tasks, such as session titles. Omit to keep the current route, or pass null to use Baseten's default.", + description="Route ID for background tasks, such as session titles. Omit to keep the current route, or pass null to use the team's default.", examples=["def5678"], title="Background", ), ] = None -class UpdateClaudeCodeHarnessConfig(BaseModel): +class UpdatePrimaryHarnessModels(BaseModel): model_config = ConfigDict( extra="forbid", ) - team_id: Annotated[ - str, - Field( - description="Identifier of the team whose default models to set. Every route must belong to this team.", - min_length=1, - title="Team Id", - ), - ] - harness: Annotated[ - Literal["claude-code"], - Field( - description="Claude Code, which supports the `primary` and `background` roles.", - title="Harness", - ), - ] - models: Annotated[ - UpdateBackgroundHarnessModels, + primary: Annotated[ + Primary | None, Field( - description="Route IDs for the model roles to change. Roles left out are unchanged." + description="Route ID for the primary model, which new sessions use. Omit to keep the current route, or pass null to use the team's default.", + examples=["abc1234"], + title="Primary", ), - ] + ] = None -class UpdateOpenCodeHarnessConfig(BaseModel): +class UpdateRouteHarnessDefaults(BaseModel): model_config = ConfigDict( extra="forbid", ) - team_id: Annotated[ - str, + claude_code: Annotated[ + UpdateBackgroundHarnessModels | None, Field( - description="Identifier of the team whose default models to set. Every route must belong to this team.", - min_length=1, - title="Team Id", + description="Route IDs for the model roles to change; roles left out are unchanged. Every route must belong to the team. Pass null to clear every role, or omit to leave the harness unchanged." ), - ] - harness: Annotated[ - Literal["opencode"], + ] = None + opencode: Annotated[ + UpdateBackgroundHarnessModels | None, Field( - description="OpenCode, which supports the `primary` and `background` roles.", - title="Harness", + description="Route IDs for the model roles to change; roles left out are unchanged. Every route must belong to the team. Pass null to clear every role, or omit to leave the harness unchanged." ), - ] - models: Annotated[ - UpdateBackgroundHarnessModels, + ] = None + codex: Annotated[ + UpdatePrimaryHarnessModels | None, Field( - description="Route IDs for the model roles to change. Roles left out are unchanged." + description="Route IDs for the model roles to change; roles left out are unchanged. Every route must belong to the team. Pass null to clear every role, or omit to leave the harness unchanged." ), - ] + ] = None + pi: Annotated[ + UpdatePrimaryHarnessModels | None, + Field( + description="Route IDs for the model roles to change; roles left out are unchanged. Every route must belong to the team. Pass null to clear every role, or omit to leave the harness unchanged." + ), + ] = None -class UpdatePrimaryHarnessModels(BaseModel): +class UpdateRouteTeamSpendLimitSetting(BaseModel): model_config = ConfigDict( extra="forbid", ) - primary: Annotated[ - Primary | None, + per_member_monthly_limit_usd: Annotated[ + str | None, Field( - description="Route ID for the primary model, which new sessions use. Omit to keep the current route, or pass null to use Baseten's default.", - examples=["abc1234"], - title="Primary", + description="Spend limit in USD for each UTC calendar month that applies to each member whose active Code key belongs to this team, unless the member has a limit of their own. Send null to remove it; omit to leave it unchanged.", + examples=["200"], + title="Per Member Monthly Limit Usd", ), ] = None -class Primary2(RootModel[str]): - root: Annotated[ - str | None, +class UpdateRouteTeamSettingsRequest(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + harness_defaults: Annotated[ + UpdateRouteHarnessDefaults | None, Field( - description="Route ID for the primary model, which new sessions use. Omit to use Baseten's default.", - examples=["abc1234"], - min_length=1, - title="Primary", + description="Harnesses to change; harnesses left out are unchanged. Pass null to clear every harness, so each role uses the team's default." + ), + ] = None + spend_limit: Annotated[ + UpdateRouteTeamSpendLimitSetting | None, + Field( + description="Spend limit fields to change. Pass null to remove the team's per-member limit; omit to leave it unchanged." ), ] = None -class Background1(RootModel[str]): - root: Annotated[ - str | None, +class RouteConnectionAnthropic(BaseModel): + provider: Annotated[ + Literal["ANTHROPIC"], + Field(description="Provider kind for Anthropic.", title="Provider"), + ] + secret_id: Annotated[ + str, Field( - description="Route ID for background tasks, such as session titles. Omit to use Baseten's default.", - examples=["def5678"], - min_length=1, - title="Background", + description="Identifier of the team secret holding the provider API key.", + title="Secret Id", ), - ] = None + ] + secret_name: Annotated[ + str, + Field( + description="Name of the team secret holding the provider API key.", + title="Secret Name", + ), + ] -class BackgroundHarnessModels(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) - primary: Annotated[ - Primary2 | None, +class RouteConnectionOpenAI(BaseModel): + provider: Annotated[ + Literal["OPENAI"], + Field(description="Provider kind for OpenAI.", title="Provider"), + ] + secret_id: Annotated[ + str, Field( - description="Route ID for the primary model, which new sessions use. Omit to use Baseten's default.", - examples=["abc1234"], - title="Primary", + description="Identifier of the team secret holding the provider API key.", + title="Secret Id", ), - ] = None - background: Annotated[ - Background1 | None, + ] + secret_name: Annotated[ + str, Field( - description="Route ID for background tasks, such as session titles. Omit to use Baseten's default.", - examples=["def5678"], - title="Background", + description="Name of the team secret holding the provider API key.", + title="Secret Name", ), - ] = None + ] -class PrimaryHarnessModels(BaseModel): +class RouteConnectionXAI(BaseModel): + provider: Annotated[ + Literal["XAI"], Field(description="Provider kind for xAI.", title="Provider") + ] + secret_id: Annotated[ + str, + Field( + description="Identifier of the team secret holding the provider API key.", + title="Secret Id", + ), + ] + secret_name: Annotated[ + str, + Field( + description="Name of the team secret holding the provider API key.", + title="Secret Name", + ), + ] + + +class RouteConnectionConfigAnthropic(BaseModel): model_config = ConfigDict( extra="forbid", ) - primary: Annotated[ - Primary2 | None, + provider: Annotated[ + Literal["ANTHROPIC"], + Field(description="Provider kind for Anthropic.", title="Provider"), + ] + secret_id: Annotated[ + str, Field( - description="Route ID for the primary model, which new sessions use. Omit to use Baseten's default.", - examples=["abc1234"], - title="Primary", + description="Identifier of an existing secret, owned by the same team, that holds the provider API key.", + min_length=1, + title="Secret Id", ), - ] = None + ] -class SetClaudeCodeHarnessConfig(BaseModel): +class RouteConnectionConfigOpenAI(BaseModel): model_config = ConfigDict( extra="forbid", ) - team_id: Annotated[ + provider: Annotated[ + Literal["OPENAI"], + Field(description="Provider kind for OpenAI.", title="Provider"), + ] + secret_id: Annotated[ str, Field( - description="Identifier of the team whose default models to set. Every route must belong to this team.", + description="Identifier of an existing secret, owned by the same team, that holds the provider API key.", min_length=1, - title="Team Id", + title="Secret Id", ), ] - harness: Annotated[ - Literal["claude-code"], + + +class RouteConnectionConfigXAI(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + provider: Annotated[ + Literal["XAI"], Field(description="Provider kind for xAI.", title="Provider") + ] + secret_id: Annotated[ + str, Field( - description="Claude Code, which supports the `primary` and `background` roles.", - title="Harness", + description="Identifier of an existing secret, owned by the same team, that holds the provider API key.", + min_length=1, + title="Secret Id", ), ] - models: Annotated[ - BackgroundHarnessModels, + + +class TeamId2(RootModel[str]): + root: Annotated[ + str | None, Field( - description="Route ID for each model role. Roles left out use Baseten's defaults." + description="Identifier of the team that owns the connection.", + examples=["abc1234"], + min_length=1, + title="Team Id", ), - ] + ] = None -class SetCodexHarnessConfig(BaseModel): +class CreateRouteConnectionRequest(BaseModel): model_config = ConfigDict( extra="forbid", ) team_id: Annotated[ - str, + TeamId2 | None, Field( - description="Identifier of the team whose default models to set. Every route must belong to this team.", - min_length=1, + description="Identifier of the team that owns the connection.", + examples=["abc1234"], title="Team Id", ), + ] = None + config: Annotated[ + Annotated[ + RouteConnectionConfigAnthropic + | RouteConnectionConfigOpenAI + | RouteConnectionConfigXAI, + Field( + description="Provider the connection authenticates with, and the team secret holding its API key.", + discriminator="provider", + examples=[{"provider": "ANTHROPIC", "secret_id": "abc1234"}], + title="Config", + ), + ] + | dict[str, Any], + Field(union_mode="left_to_right"), + ] + + +class RouteConnectionTombstone(BaseModel): + id: Annotated[ + str, + Field(description="Stable identifier of the deleted connection.", title="Id"), ] - harness: Annotated[ - Literal["codex"], + + +class SecretId(RootModel[str]): + root: Annotated[ + str | None, Field( - description="Codex, which supports only the `primary` role.", - title="Harness", + description="Identifier of the new secret, owned by the connection's team, that holds the provider API key. Omit to keep the current secret.", + min_length=1, + title="Secret Id", ), + ] = None + + +class UpdateRouteConnectionConfigAnthropic(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + provider: Annotated[ + Literal["ANTHROPIC"], + Field(description="Provider kind for Anthropic.", title="Provider"), ] - models: Annotated[ - PrimaryHarnessModels, + secret_id: Annotated[ + SecretId | None, Field( - description="Route ID for each model role. Roles left out use Baseten's defaults." + description="Identifier of the new secret, owned by the connection's team, that holds the provider API key. Omit to keep the current secret.", + title="Secret Id", ), - ] + ] = None -class SetOpenCodeHarnessConfig(BaseModel): +class UpdateRouteConnectionConfigOpenAI(BaseModel): model_config = ConfigDict( extra="forbid", ) - team_id: Annotated[ - str, - Field( - description="Identifier of the team whose default models to set. Every route must belong to this team.", - min_length=1, - title="Team Id", - ), + provider: Annotated[ + Literal["OPENAI"], + Field(description="Provider kind for OpenAI.", title="Provider"), ] - harness: Annotated[ - Literal["opencode"], + secret_id: Annotated[ + SecretId | None, Field( - description="OpenCode, which supports the `primary` and `background` roles.", - title="Harness", + description="Identifier of the new secret, owned by the connection's team, that holds the provider API key. Omit to keep the current secret.", + title="Secret Id", ), + ] = None + + +class UpdateRouteConnectionConfigXAI(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + provider: Annotated[ + Literal["XAI"], Field(description="Provider kind for xAI.", title="Provider") ] - models: Annotated[ - BackgroundHarnessModels, + secret_id: Annotated[ + SecretId | None, Field( - description="Route ID for each model role. Roles left out use Baseten's defaults." + description="Identifier of the new secret, owned by the connection's team, that holds the provider API key. Omit to keep the current secret.", + title="Secret Id", ), - ] + ] = None -class SetRouteHarnessConfigRequest( - RootModel[ - SetClaudeCodeHarnessConfig | SetOpenCodeHarnessConfig | SetCodexHarnessConfig - ] -): - root: Annotated[ +class UpdateRouteConnectionRequest(BaseModel): + model_config = ConfigDict( + extra="forbid", + ) + config: Annotated[ Annotated[ - SetClaudeCodeHarnessConfig - | SetOpenCodeHarnessConfig - | SetCodexHarnessConfig, - Field(discriminator="harness", title="SetRouteHarnessConfigRequestV1"), + UpdateRouteConnectionConfigAnthropic + | UpdateRouteConnectionConfigOpenAI + | UpdateRouteConnectionConfigXAI, + Field( + description="Connection fields to change. The provider must match the connection and is immutable.", + discriminator="provider", + examples=[{"provider": "ANTHROPIC", "secret_id": "abc1234"}], + title="Config", + ), ] | dict[str, Any], Field(union_mode="left_to_right"), ] -class RouteHarnessConfigTombstone(BaseModel): - harness: Annotated[ - Literal["claude-code", "opencode", "codex"] | str, - Field( - description="Harness whose default models were cleared.", - title="RouteHarness", - ), - ] - team_id: Annotated[ - str, - Field( - description="Identifier of the team whose default models were cleared.", - title="Team Id", - ), - ] - - class RouteTombstone(BaseModel): id: Annotated[ str, Field(description="Stable identifier of the deleted route.", title="Id") @@ -8433,6 +8666,18 @@ class UpdateRouteRequest(BaseModel): title="Display Name", ), ] = None + target: Annotated[ + RouteTargetConfigBasetenModelAPI + | RouteTargetConfigAnthropic + | RouteTargetConfigOpenAI + | RouteTargetConfigXAI + | RouteTargetConfigClassifierModelBased + | None, + Field( + description="Complete new target configuration. Omit to keep the current configuration; null is not accepted. Only the CLASSIFIER_MODEL_BASED configuration of a router route is mutable.", + title="Target", + ), + ] = None class VertexTargetConfig(BaseModel): @@ -8693,7 +8938,7 @@ class EffectiveRateLimit(BaseModel): ] -class Name4(RootModel[str]): +class Name5(RootModel[str]): root: Annotated[ str | None, Field( @@ -8707,7 +8952,7 @@ class Name4(RootModel[str]): class GroupMetadata(BaseModel): name: Annotated[ - Name4 | None, + Name5 | None, Field( description="Optional display name for the group.", examples=["Acme prod"], @@ -8776,7 +9021,7 @@ class CreateGroupHierarchy(BaseModel): class UpdateGroupMetadata(BaseModel): name: Annotated[ - Name4 | None, + Name5 | None, Field( description="Optional display name for the group.", examples=["Acme prod"], @@ -8854,14 +9099,90 @@ class RegisterAPIKeyResponse(BaseModel): ] -class ImageBuildLog(BaseModel): +class SandboxRegion(BaseModel): + name: Annotated[ + str, + Field( + description="Public region identifier to use when creating a sandbox.", + examples=["us-was-1"], + ), + ] + country: Annotated[str, Field(description="Country code.", examples=["us"])] + continent: Annotated[str, Field(description="Continent code.", examples=["na"])] + location: Annotated[ + str, Field(description="Region location.", examples=["Washington"]) + ] + info_generation: Annotated[ + Literal["CARBON"] | str, + Field( + description="Runtime generation supported by this region, using the public name CARBON. Actual runtime selection depends on the team and sandbox configuration." + ), + ] + + +class SandboxLogEntry(BaseModel): + timestamp: AwareDatetime + message: str + severity: Annotated[int, Field(description="Numeric log severity.")] + trace_id: Annotated[ + str, + Field( + description="Associated trace identifier, or an empty string when unavailable." + ), + ] + action: Annotated[ + str | None, + Field(description="Action or command when present on the log entry."), + ] = None + + +class ErrorRate(RootModel[float]): + root: Annotated[ + float, + Field( + description="Fraction of requests returning 4xx or 5xx. Null when there are no requests.", + ge=0.0, + le=1.0, + ), + ] + + +class SandboxMetricsPoint(BaseModel): + timestamp: Annotated[AwareDatetime, Field(description="Start of this interval.")] + requests: Annotated[ + float | None, + Field( + description="Request count. Zero without traffic; null for intervals entirely before creation." + ), + ] + cpu_percent: Annotated[ + float | None, + Field( + description="Peak CPU usage in this interval, where 100 is one full CPU core. Null without a sample." + ), + ] + memory_bytes: Annotated[ + float | None, + Field( + description="Memory usage in bytes, averaged per series then maximum across series. Null without a sample." + ), + ] + error_rate: Annotated[ + ErrorRate | None, + Field( + description="Fraction of requests returning 4xx or 5xx. Null when there are no requests." + ), + ] + + +class SandboxImageBuildLog(BaseModel): timestamp: AwareDatetime message: str severity: Annotated[int, Field(description="Numeric OpenTelemetry severity level.")] -class ImageBuildLogsResponse(BaseModel): - logs: list[ImageBuildLog] +class SandboxImageBuildLogsResponse(BaseModel): + logs: list[SandboxImageBuildLog] total_count: Annotated[ int, Field( @@ -8876,7 +9197,7 @@ class SandboxTTLIdleExpirationPolicy(BaseModel): value: Annotated[ str, Field( - description="Duration using seconds, minutes, hours, or composite durations such as 1h30m. Whole days and weeks are also supported, for example 7d or 2w.", + description="Duration using seconds, minutes, hours, or composite durations such as 1h30m. Whole days and weeks are also supported, for example 7d or 2w, where d is 24h and w is 7 × 24h. Days and weeks cannot be combined with other units, so 1d12h is rejected; use 36h instead. Values are returned exactly as sent, without normalization.", examples=["24h"], pattern="^[+-]?(0|[0-9]+[dw]|([0-9]+(\\.[0-9]*)?|\\.[0-9]+)(ns|us|µs|μs|ms|s|m|h)(([0-9]+(\\.[0-9]*)?|\\.[0-9]+)(ns|us|µs|μs|ms|s|m|h))*)$", ), @@ -8889,7 +9210,7 @@ class SandboxTTLMaxAgeExpirationPolicy(BaseModel): value: Annotated[ str, Field( - description="Duration using seconds, minutes, hours, or composite durations such as 1h30m. Whole days and weeks are also supported, for example 7d or 2w.", + description="Duration using seconds, minutes, hours, or composite durations such as 1h30m. Whole days and weeks are also supported, for example 7d or 2w, where d is 24h and w is 7 × 24h. Days and weeks cannot be combined with other units, so 1d12h is rejected; use 36h instead. Values are returned exactly as sent, without normalization.", examples=["24h"], pattern="^[+-]?(0|[0-9]+[dw]|([0-9]+(\\.[0-9]*)?|\\.[0-9]+)(ns|us|µs|μs|ms|s|m|h)(([0-9]+(\\.[0-9]*)?|\\.[0-9]+)(ns|us|µs|μs|ms|s|m|h))*)$", ), @@ -8983,7 +9304,7 @@ class SandboxApiPagination(BaseModel): ] = None -class ImageTag(BaseModel): +class SandboxImageTag(BaseModel): name: Annotated[str, Field(description="Image tag name.", examples=["latest"])] created_at: Annotated[ AwareDatetime | None, @@ -9004,7 +9325,7 @@ class ImageTag(BaseModel): ] = None -class Image(BaseModel): +class SandboxImageSummary(BaseModel): name: Annotated[ str, Field( @@ -9053,31 +9374,19 @@ class Image(BaseModel): ge=0, ), ] = None + + +class SandboxImage(SandboxImageSummary): tags: Annotated[ - list[ImageTag], + list[SandboxImageTag] | None, Field( - description="Empty for list and get summary responses. Use GET /sandboxes/images/{image_name}/tags to retrieve paginated image versions.", - examples=[ - [ - { - "name": "latest", - "created_at": "2026-09-16T21:20:00Z", - "updated_at": "2026-09-16T21:25:00Z", - "size": 134217728, - }, - { - "name": "20260915212000", - "created_at": "2026-09-15T21:20:00Z", - "updated_at": "2026-09-15T21:25:00Z", - "size": 125829120, - }, - ] - ], + description="Empty for get summary responses. Use GET /sandboxes/images/{image_name}/tags to retrieve paginated image versions.", + examples=[[]], ), - ] + ] = None -class PushImageRequest(BaseModel): +class PushSandboxImageRequest(BaseModel): name: Annotated[ str, Field( @@ -9104,7 +9413,7 @@ class PushImageRequest(BaseModel): ] = None -class PushImageResponse(BaseModel): +class PushSandboxImageResponse(BaseModel): name: Annotated[ str, Field(description="Target image repository name.", examples=["base-image"]) ] @@ -9176,7 +9485,7 @@ class SandboxLibraryImageCreationOptions(BaseModel): ] = None -class CleanupImagesResponse(BaseModel): +class CleanupSandboxImagesResponse(BaseModel): deleted: Annotated[ int, Field(description="Number of image versions removed.", examples=[3], ge=0) ] @@ -10308,9 +10617,6 @@ class EnvironmentGroups(BaseModel): class AuditLogEventAutoscalingScheduleChange(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) action: Annotated[ Literal["CREATED", "UPDATED", "DELETED", "UNCHANGED"] | str, Field( @@ -10324,9 +10630,6 @@ class AuditLogEventAutoscalingScheduleChange(BaseModel): class AuditLogEventEnvironmentUpdated(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) schedules: Annotated[ list[AuditLogEventAutoscalingScheduleChange] | None, Field(title="Schedules") ] @@ -10379,9 +10682,6 @@ class AuditLogEventEnvironmentUpdated(BaseModel): class AuditLogEventModelDeploymentAutoscalingSettingsChanged(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) schedules: Annotated[ list[AuditLogEventAutoscalingScheduleChange] | None, Field(title="Schedules") ] @@ -10410,9 +10710,6 @@ class AuditLogEventModelDeploymentAutoscalingSettingsChanged(BaseModel): class AuditLogEventModelPromotionControlAction(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) event_type: Annotated[ Literal["MODEL_PROMOTION_CONTROL_ACTION"], Field(title="Event Type") ] @@ -12276,7 +12573,8 @@ class Route(BaseModel): RouteTargetBasetenModelAPI | RouteTargetAnthropic | RouteTargetOpenAI - | RouteTargetXAI, + | RouteTargetXAI + | RouteTargetClassifierModelBased, Field( description="Configured upstream target.", discriminator="type", @@ -12289,7 +12587,7 @@ class Route(BaseModel): metadata: Annotated[ ExploreMetadata | None, Field( - description="Resolved model metadata; null when the route has no linked metadata row." + description="Resolved model metadata; for a router, the envelope of its allowed routes' metadata. Null when nothing is linked." ), ] invoke_url: Annotated[ @@ -12339,61 +12637,112 @@ class RoutesUsageResponse(BaseModel): ] -class RouteHarnessModel(BaseModel): +class RouteEffectiveSpendLimit(BaseModel): + monthly_limit_usd: Annotated[ + str | None, + Field( + description="Spend limit in USD enforced for the current UTC calendar month: the user's own limit, else the team default. Null when no limit applies.", + examples=["200"], + title="Monthly Limit Usd", + ), + ] source: Annotated[ - Literal["team", "baseten"] | str, + Literal["user", "team"] | str | None, Field( - description="Who chose this role's route: `team` if a team admin set it, or `baseten` if it is Baseten's default, chosen from the team's Model API routes.", - title="RouteHarnessModelSource", + description="Where the limit comes from: `user` when it is set on the user, `team` when it is the team default. Null when no limit applies." ), ] - route: Annotated[Route, Field(description="Route to use for this role.")] -class UpdateCodexHarnessConfig(BaseModel): - model_config = ConfigDict( - extra="forbid", - ) - team_id: Annotated[ - str, +class RouteSpendLimitSetting(BaseModel): + user_monthly_limit_usd: Annotated[ + str | None, Field( - description="Identifier of the team whose default models to set. Every route must belong to this team.", - min_length=1, - title="Team Id", + description="Standing spend limit in USD for each UTC calendar month set on the user. Null when the user has no limit of their own.", + examples=["200"], + title="User Monthly Limit Usd", ), ] - harness: Annotated[ - Literal["codex"], + team_default: Annotated[ + RouteSpendLimitTeamDefault | None, Field( - description="Codex, which supports only the `primary` role.", - title="Harness", + description="Per-member limit of the team the user's active Code key belongs to. This limit applies when the user has no limit of their own. Null when there is none." ), ] - models: Annotated[ - UpdatePrimaryHarnessModels, + effective: Annotated[ + RouteEffectiveSpendLimit, + Field(description="The effective limit enforced for the current month."), + ] + + +class RouteUserSettings(BaseModel): + user_id: Annotated[ + str, Field(description="ID of the user.", examples=["abc1234"], title="User Id") + ] + email: Annotated[ + str | None, + Field( + description="Email address of the user.", + examples=["dev@example.com"], + title="Email", + ), + ] + spend_limit: Annotated[ + RouteSpendLimitSetting, Field( - description="Route IDs for the model roles to change. Roles left out are unchanged." + description="Spend limit for Baseten Code. Applies to requests with Routes keys the user created; personal API keys are not limited." ), ] -class UpdateRouteHarnessConfigRequest( - RootModel[ - UpdateClaudeCodeHarnessConfig - | UpdateOpenCodeHarnessConfig - | UpdateCodexHarnessConfig +class RouteHarnessModel(BaseModel): + source: Annotated[ + Literal["user", "team"] | str, + Field( + description="Where this role's route comes from. `team` covers both a route a team admin set and the default chosen from the team's Model API routes.", + title="RouteSettingSourceV1", + ), ] -): - root: Annotated[ + route: Annotated[Route, Field(description="Route to use for this role.")] + + +class RouteConnection(BaseModel): + id: Annotated[str, Field(description="Stable connection identifier.", title="Id")] + team_id: Annotated[ + str, + Field( + description="Identifier of the team that owns the connection.", + title="Team Id", + ), + ] + config: Annotated[ Annotated[ - UpdateClaudeCodeHarnessConfig - | UpdateOpenCodeHarnessConfig - | UpdateCodexHarnessConfig, - Field(discriminator="harness", title="UpdateRouteHarnessConfigRequestV1"), + RouteConnectionAnthropic | RouteConnectionOpenAI | RouteConnectionXAI, + Field( + description="Provider the connection authenticates with, and the team secret holding its API key.", + discriminator="provider", + title="Config", + ), ] | dict[str, Any], Field(union_mode="left_to_right"), ] + created_at: Annotated[ + AwareDatetime, Field(description="Creation time, ISO 8601.", title="Created At") + ] + updated_at: Annotated[ + AwareDatetime, + Field(description="Last update time, ISO 8601.", title="Updated At"), + ] + + +class RouteConnectionsResponse(BaseModel): + items: Annotated[ + list[RouteConnection], Field(description="Items in this page.", title="Items") + ] + pagination: Annotated[ + PaginationResponse, Field(description="Pagination metadata for the page.") + ] class EndpointTarget(BaseModel): @@ -12576,6 +12925,36 @@ class UpdateGroupRequest(BaseModel): ] = None +class GetSandboxConfigurationResponse(BaseModel): + regions: Annotated[ + list[SandboxRegion], + Field( + description="Available Carbon-compatible regions, sorted by name. Empty when none are available." + ), + ] + + +class SandboxLogs(BaseModel): + sandbox_name: str + start_time: AwareDatetime + end_time: AwareDatetime + logs: Annotated[list[SandboxLogEntry], Field(max_length=1000)] + next_cursor: Annotated[ + str | None, + Field( + description="Opaque cursor for the next page; null when there are no more entries." + ), + ] + + +class SandboxMetrics(BaseModel): + sandbox_name: str + start_time: AwareDatetime + end_time: AwareDatetime + interval_seconds: int + data: Annotated[list[SandboxMetricsPoint], Field(max_length=10000)] + + class SandboxProxyConfig(BaseModel): allowed_domains: Annotated[ list[str] | None, @@ -12623,13 +13002,15 @@ class SandboxProxyConfig(BaseModel): ] = None -class ListImagesResponse(BaseModel): - items: Annotated[list[Image], Field(description="Image repositories on this page.")] +class ListSandboxImagesResponse(BaseModel): + items: Annotated[ + list[SandboxImageSummary], Field(description="Image repositories on this page.") + ] pagination: SandboxApiPagination -class ListImageTagsResponse(BaseModel): - items: list[ImageTag] +class ListSandboxImageTagsResponse(BaseModel): + items: list[SandboxImageTag] pagination: SandboxApiPagination @@ -13091,19 +13472,57 @@ class GetLoopsDeploymentMetricsResponse(BaseModel): ] -class RouteHarnessConfig(BaseModel): - models: Annotated[ - dict[Literal["primary", "background"] | str, RouteHarnessModel], - Field(description="Route for each model role, keyed by role.", title="Models"), +class BackgroundHarnessDefaults(BaseModel): + primary: Annotated[ + RouteHarnessModel | None, + Field( + description="Route for the primary model, which new sessions use. Null when the team has no route to use." + ), + ] + background: Annotated[ + RouteHarnessModel | None, + Field( + description="Route for background tasks, such as session titles. Null when the team has no route to use." + ), + ] + + +class PrimaryHarnessDefaults(BaseModel): + primary: Annotated[ + RouteHarnessModel | None, + Field( + description="Route for the primary model, which new sessions use. Null when the team has no route to use." + ), + ] + + +class RouteHarnessDefaults(BaseModel): + claude_code: Annotated[ + BackgroundHarnessDefaults, Field(description="Default models for Claude Code.") + ] + opencode: Annotated[ + BackgroundHarnessDefaults, Field(description="Default models for OpenCode.") ] + codex: Annotated[ + PrimaryHarnessDefaults, Field(description="Default models for Codex.") + ] + pi: Annotated[PrimaryHarnessDefaults, Field(description="Default models for Pi.")] -class RouteHarnessConfigsResponse(BaseModel): - harness_configs: Annotated[ - dict[Literal["claude-code", "opencode", "codex"] | str, RouteHarnessConfig], +class RouteTeamSettings(BaseModel): + team_id: Annotated[ + str, Field(description="ID of the team.", examples=["abc1234"], title="Team Id") + ] + harness_defaults: Annotated[ + RouteHarnessDefaults, + Field( + description="Route for each model role of each coding harness. A role with no route set by a team admin uses the default chosen from the team's Model API routes: the recommended model for `primary` and the lowest-priced model for `background`. These defaults never use external-provider routes." + ), + ] + spend_limit: Annotated[ + RouteTeamSpendLimitSetting, Field( - description="Default models for each harness, keyed by harness. A harness is omitted when none of its roles has a route.", - title="Harness Configs", + description="Spend limits for Baseten Code for members whose active Code key belongs to this team." ), ] @@ -13185,6 +13604,7 @@ class SandboxLifecycle(BaseModel): Field( description="Duration to keep the sandbox record after termination for log access (e.g., '1h', '24h', '7d'). Defaults to 5m. Subject to maximum quota limits.", examples=["24h"], + pattern="^[+-]?(0|[0-9]+[dw]|([0-9]+(\\.[0-9]*)?|\\.[0-9]+)(ns|us|µs|μs|ms|s|m|h)(([0-9]+(\\.[0-9]*)?|\\.[0-9]+)(ns|us|µs|μs|ms|s|m|h))*)$", ), ] = None @@ -13230,13 +13650,6 @@ class SandboxNetwork(BaseModel): class SandboxConfiguration(BaseModel): - enabled: Annotated[ - bool, - Field( - description="When false, the sandbox is disabled and will not accept connections", - examples=[True], - ), - ] = True lifecycle: Annotated[ SandboxLifecycle | None, Field( @@ -13481,13 +13894,6 @@ class CreateSandboxRequest(BaseModel): class UpdateSandboxRequest(BaseModel): - enabled: Annotated[ - bool | None, - Field( - description="When false, the sandbox is disabled and will not accept connections", - examples=[True], - ), - ] = None lifecycle: Annotated[ SandboxLifecycle | None, Field( @@ -13508,13 +13914,6 @@ class UpdateSandboxRequest(BaseModel): ], ), ] = None - region: Annotated[ - str | None, - Field( - description="Region where the sandbox runs (for example us-pdx-1 or eu-lon-1). When omitted at creation, the closest region is selected.", - examples=["us-pdx-1"], - ), - ] = None envs: Annotated[ list[SandboxEnv] | None, Field( @@ -13527,20 +13926,6 @@ class UpdateSandboxRequest(BaseModel): ], ), ] = None - image: Annotated[ - str | None, - Field( - description="Image reference including its tag. Built-in image references are returned in the canonical baseten/ namespace. Use baseten/base-image:latest to get started with the built-in sandbox execution API. This image is available directly without building, pushing, or listing images through GET /v1/sandboxes/images.", - examples=["baseten/base-image:latest"], - ), - ] = None - ports: Annotated[ - list[SandboxPort] | None, - Field( - description="Set of ports for a resource", - examples=[[{"name": "http", "protocol": "HTTP", "target": 3000}]], - ), - ] = None external_id: Annotated[ str | None, Field( @@ -13652,7 +14037,6 @@ class ListSandboxesResponse(BaseModel): examples=[ [ { - "enabled": True, "lifecycle": { "expiration_policies": [ { diff --git a/baseten/client/modelconfig/_models.py b/baseten/client/modelconfig/_models.py index 002dc0f..7a92f40 100644 --- a/baseten/client/modelconfig/_models.py +++ b/baseten/client/modelconfig/_models.py @@ -496,7 +496,9 @@ class ModelTRTLLMRuntimeConfiguration(BaseModel): served_model_name: Annotated[str | None, Field(title="Served Model Name")] = None total_token_limit: Annotated[int, Field(title="Total Token Limit")] = 500000 webserver_default_route: Annotated[ - Literal["/v1/embeddings", "/rerank", "/predict", "/predict_tokens"] + Literal[ + "/v1/embeddings", "/v1/systemone", "/rerank", "/predict", "/predict_tokens" + ] | str | None, Field(title="Webserver Default Route"), @@ -957,15 +959,15 @@ class BaseImage(BaseModel): str, Field( description="The path to the Docker image.", - examples=["vllm/vllm-openai:v0.7.3", "nvcr.io/nvidia/nemo:23.03"], + examples=["vllm/vllm-openai:v0.29.0", "nvcr.io/nvidia/nemo:23.03"], title="Image", ), ] = "" python_executable_path: Annotated[ str, Field( - description="A path to the Python executable on the image.", - examples=["/usr/bin/python"], + description="A path to the Python executable on the image. Truss uses `python3` from the image's PATH when this is empty.", + examples=["/usr/bin/python3"], title="Python Executable Path", ), ] = "" diff --git a/baseten/client/sandboxapi/__init__.py b/baseten/client/sandboxapi/__init__.py index a09b9a3..410287c 100644 --- a/baseten/client/sandboxapi/__init__.py +++ b/baseten/client/sandboxapi/__init__.py @@ -46,7 +46,6 @@ GetFilesystemParams, GetFilesystemResponse, GetFilesystemSearchParams, - GetFilesystemTreeResponse, GetNetworkProcessPortsResponse, GetProcessResponse, GetWatchFilesystemParams, @@ -66,7 +65,6 @@ ProcessRequest, ProcessResponse, PutFilesystemMultipartPartParams, - PutFilesystemTreeResponse, QuiesceStatus, RankedFile, RerankingResponse, @@ -116,7 +114,6 @@ "GetFilesystemParams", "GetFilesystemResponse", "GetFilesystemSearchParams", - "GetFilesystemTreeResponse", "GetNetworkProcessPortsResponse", "GetProcessResponse", "GetWatchFilesystemParams", @@ -136,7 +133,6 @@ "ProcessRequest", "ProcessResponse", "PutFilesystemMultipartPartParams", - "PutFilesystemTreeResponse", "QuiesceStatus", "RankedFile", "RerankingResponse", diff --git a/baseten/client/sandboxapi/_client.py b/baseten/client/sandboxapi/_client.py index 261162e..4f0af2b 100644 --- a/baseten/client/sandboxapi/_client.py +++ b/baseten/client/sandboxapi/_client.py @@ -18,6 +18,7 @@ DeleteFilesystemParams, DeleteFilesystemTreeParams, DeleteNetworkProcessMonitorResponse, + Directory, DriveListResponse, DriveMountRequest, DriveMountResponse, @@ -35,7 +36,6 @@ GetFilesystemParams, GetFilesystemResponse, GetFilesystemSearchParams, - GetFilesystemTreeResponse, GetNetworkProcessPortsResponse, GetProcessResponse, GetWatchFilesystemParams, @@ -53,7 +53,6 @@ ProcessRequest, ProcessResponse, PutFilesystemMultipartPartParams, - PutFilesystemTreeResponse, QuiesceStatus, RerankingResponse, SuccessResponse, @@ -109,9 +108,33 @@ class ApiClient: They may change without notice between versions. """ - def __init__(self, http_client: httpx.Client) -> None: - """Create a new client. The caller is responsible for closing *http_client*.""" + def __init__( + self, + http_client: httpx.Client, + *, + base_url: str | None = None, + auth: httpx.Auth | None = None, + headers: Mapping[str, str] | None = None, + timeout: httpx.Timeout | None = None, + ) -> None: + """Create a new client. The caller is responsible for closing *http_client*. + + The keyword arguments apply to this client's requests only, so one + HTTP client can serve several APIs. Each one left ``None`` falls back + to the HTTP client's own setting. + + Args: + http_client: HTTP client to send requests with. + base_url: Base URL that request paths are appended to. + auth: Authentication for each request. + headers: Headers for each request, added to the HTTP client's own. + timeout: Timeouts for each request. + """ self._http_client = http_client + self._base_url = base_url + self._auth = auth + self._headers = headers + self._timeout = timeout def delete_drives_mount(self, *, mount_path: str) -> DriveUnmountResponse: """Detach a drive from a local path""" @@ -474,10 +497,10 @@ def get_filesystem_search( ), ) - def get_filesystem_tree(self, *, path: str) -> GetFilesystemTreeResponse: + def get_filesystem_tree(self, *, path: str) -> Directory: """Get directory tree""" return self._do_json( - GetFilesystemTreeResponse, + Directory, _ApiRequest( method="GET", path_fmt="/filesystem/tree/{}", @@ -576,8 +599,10 @@ def get_process_logs(self, *, identifier: str) -> ProcessLogs: ), ) - def get_process_logs_stream(self, *, identifier: str) -> httpx.Response: - """Stream process logs in real time. Returns the response unread. The caller must close the response.""" + def get_process_logs_stream( + self, *, identifier: str, accept: Literal["application/x-ndjson", "text/plain"] + ) -> httpx.Response: + """Stream process logs in real time. Returns the response unread, in the requested content type. The caller must close the response.""" return self._do_raw( _ApiRequest( method="GET", @@ -586,8 +611,13 @@ def get_process_logs_stream(self, *, identifier: str) -> httpx.Response: body=None, query=None, success_codes=[200], - error_codes=None, - accept="text/plain", + error_codes={ + 400: "ErrorResponse", + 404: "ErrorResponse", + 409: "ErrorResponse", + 500: "ErrorResponse", + }, + accept=accept, ) ) @@ -608,6 +638,20 @@ def get_watch_filesystem( ) ) + def head_filesystem(self, *, path: str) -> httpx.Response: + """Stat a file or directory. Returns the response unread, since its headers are the result. The caller must close the response.""" + return self._do_raw( + _ApiRequest( + method="HEAD", + path_fmt="/filesystem/{}", + path_args=[path], + body=None, + query=None, + success_codes=[200], + error_codes=None, + ) + ) + def post_archive_export( self, *, request: ExportOptions ) -> ExportResult | ExportProgress: @@ -761,7 +805,7 @@ def post_process_raw( self, *, request: ProcessRequest, - accept: Literal["application/json", "text/event-stream"], + accept: Literal["application/json", "application/x-ndjson"], ) -> httpx.Response: """Execute a command. Returns the response unread, in the requested content type. The caller must close the response.""" return self._do_raw( @@ -882,12 +926,10 @@ def put_filesystem_multipart_part( ), ) - def put_filesystem_tree( - self, *, path: str, request: TreeRequest - ) -> PutFilesystemTreeResponse: + def put_filesystem_tree(self, *, path: str, request: TreeRequest) -> Directory: """Create or update directory tree""" return self._do_json( - PutFilesystemTreeResponse, + Directory, _ApiRequest( method="PUT", path_fmt="/filesystem/tree/{}", @@ -931,7 +973,7 @@ def _build_request(self, request: _ApiRequest) -> httpx.Request: json_body = None content_body = None files_body = None - headers: dict[str, str] = {} + headers: dict[str, str] = {**(self._headers or {})} if request.accept is not None: headers["Accept"] = request.accept if request.body is not None: @@ -965,16 +1007,22 @@ def _build_request(self, request: _ApiRequest) -> httpx.Request: params = request.query return self._http_client.build_request( request.method, - path, + path if self._base_url is None else self._base_url.rstrip("/") + path, json=json_body, content=content_body, files=files_body, params=params, headers=headers, + timeout=httpx.USE_CLIENT_DEFAULT + if self._timeout is None + else self._timeout, ) def _do(self, request: _ApiRequest) -> httpx.Response: - response = self._http_client.send(self._build_request(request)) + response = self._http_client.send( + self._build_request(request), + auth=httpx.USE_CLIENT_DEFAULT if self._auth is None else self._auth, + ) if response.status_code not in request.success_codes: if request.error_codes and response.status_code in request.error_codes: error_name = request.error_codes[response.status_code] @@ -994,7 +1042,11 @@ def _do(self, request: _ApiRequest) -> httpx.Response: return response def _do_raw(self, request: _ApiRequest) -> httpx.Response: - response = self._http_client.send(self._build_request(request), stream=True) + response = self._http_client.send( + self._build_request(request), + auth=httpx.USE_CLIENT_DEFAULT if self._auth is None else self._auth, + stream=True, + ) if response.status_code not in request.success_codes: response.read() if request.error_codes and response.status_code in request.error_codes: @@ -1049,9 +1101,33 @@ class AsyncApiClient: They may change without notice between versions. """ - def __init__(self, http_client: httpx.AsyncClient) -> None: - """Create a new client. The caller is responsible for closing *http_client*.""" + def __init__( + self, + http_client: httpx.AsyncClient, + *, + base_url: str | None = None, + auth: httpx.Auth | None = None, + headers: Mapping[str, str] | None = None, + timeout: httpx.Timeout | None = None, + ) -> None: + """Create a new client. The caller is responsible for closing *http_client*. + + The keyword arguments apply to this client's requests only, so one + HTTP client can serve several APIs. Each one left ``None`` falls back + to the HTTP client's own setting. + + Args: + http_client: HTTP client to send requests with. + base_url: Base URL that request paths are appended to. + auth: Authentication for each request. + headers: Headers for each request, added to the HTTP client's own. + timeout: Timeouts for each request. + """ self._http_client = http_client + self._base_url = base_url + self._auth = auth + self._headers = headers + self._timeout = timeout async def delete_drives_mount(self, *, mount_path: str) -> DriveUnmountResponse: """Detach a drive from a local path""" @@ -1416,10 +1492,10 @@ async def get_filesystem_search( ), ) - async def get_filesystem_tree(self, *, path: str) -> GetFilesystemTreeResponse: + async def get_filesystem_tree(self, *, path: str) -> Directory: """Get directory tree""" return await self._do_json( - GetFilesystemTreeResponse, + Directory, _ApiRequest( method="GET", path_fmt="/filesystem/tree/{}", @@ -1520,8 +1596,10 @@ async def get_process_logs(self, *, identifier: str) -> ProcessLogs: ), ) - async def get_process_logs_stream(self, *, identifier: str) -> httpx.Response: - """Stream process logs in real time. Returns the response unread. The caller must close the response.""" + async def get_process_logs_stream( + self, *, identifier: str, accept: Literal["application/x-ndjson", "text/plain"] + ) -> httpx.Response: + """Stream process logs in real time. Returns the response unread, in the requested content type. The caller must close the response.""" return await self._do_raw( _ApiRequest( method="GET", @@ -1530,8 +1608,13 @@ async def get_process_logs_stream(self, *, identifier: str) -> httpx.Response: body=None, query=None, success_codes=[200], - error_codes=None, - accept="text/plain", + error_codes={ + 400: "ErrorResponse", + 404: "ErrorResponse", + 409: "ErrorResponse", + 500: "ErrorResponse", + }, + accept=accept, ) ) @@ -1552,6 +1635,20 @@ async def get_watch_filesystem( ) ) + async def head_filesystem(self, *, path: str) -> httpx.Response: + """Stat a file or directory. Returns the response unread, since its headers are the result. The caller must close the response.""" + return await self._do_raw( + _ApiRequest( + method="HEAD", + path_fmt="/filesystem/{}", + path_args=[path], + body=None, + query=None, + success_codes=[200], + error_codes=None, + ) + ) + async def post_archive_export( self, *, request: ExportOptions ) -> ExportResult | ExportProgress: @@ -1707,7 +1804,7 @@ async def post_process_raw( self, *, request: ProcessRequest, - accept: Literal["application/json", "text/event-stream"], + accept: Literal["application/json", "application/x-ndjson"], ) -> httpx.Response: """Execute a command. Returns the response unread, in the requested content type. The caller must close the response.""" return await self._do_raw( @@ -1832,10 +1929,10 @@ async def put_filesystem_multipart_part( async def put_filesystem_tree( self, *, path: str, request: TreeRequest - ) -> PutFilesystemTreeResponse: + ) -> Directory: """Create or update directory tree""" return await self._do_json( - PutFilesystemTreeResponse, + Directory, _ApiRequest( method="PUT", path_fmt="/filesystem/tree/{}", @@ -1879,7 +1976,7 @@ def _build_request(self, request: _ApiRequest) -> httpx.Request: json_body = None content_body = None files_body = None - headers: dict[str, str] = {} + headers: dict[str, str] = {**(self._headers or {})} if request.accept is not None: headers["Accept"] = request.accept if request.body is not None: @@ -1913,16 +2010,22 @@ def _build_request(self, request: _ApiRequest) -> httpx.Request: params = request.query return self._http_client.build_request( request.method, - path, + path if self._base_url is None else self._base_url.rstrip("/") + path, json=json_body, content=content_body, files=files_body, params=params, headers=headers, + timeout=httpx.USE_CLIENT_DEFAULT + if self._timeout is None + else self._timeout, ) async def _do(self, request: _ApiRequest) -> httpx.Response: - response = await self._http_client.send(self._build_request(request)) + response = await self._http_client.send( + self._build_request(request), + auth=httpx.USE_CLIENT_DEFAULT if self._auth is None else self._auth, + ) if response.status_code not in request.success_codes: if request.error_codes and response.status_code in request.error_codes: error_name = request.error_codes[response.status_code] @@ -1943,7 +2046,9 @@ async def _do(self, request: _ApiRequest) -> httpx.Response: async def _do_raw(self, request: _ApiRequest) -> httpx.Response: response = await self._http_client.send( - self._build_request(request), stream=True + self._build_request(request), + auth=httpx.USE_CLIENT_DEFAULT if self._auth is None else self._auth, + stream=True, ) if response.status_code not in request.success_codes: await response.aread() diff --git a/baseten/client/sandboxapi/_models.py b/baseten/client/sandboxapi/_models.py index a692a2c..1c116ee 100644 --- a/baseten/client/sandboxapi/_models.py +++ b/baseten/client/sandboxapi/_models.py @@ -97,7 +97,11 @@ class ArchiveManifest(BaseModel): class ContentSearchMatch(BaseModel): column: Annotated[int, Field(examples=[10])] context: Annotated[ - str | None, Field(examples=["previous line\ncurrent line\nnext line"]) + str | None, + Field( + description="The matching line with up to contextLines lines before and after it, newline-separated; omitted when contextLines is 0", + examples=["previous line\ncurrent line\nnext line"], + ), ] = None line: Annotated[int, Field(examples=[42])] path: Annotated[str, Field(examples=["src/main.go"])] @@ -280,7 +284,13 @@ class PortMonitorRequest(BaseModel): class ProcessLogs(BaseModel): - logs: Annotated[str, Field(examples=["logs output"])] + logs: Annotated[ + str, + Field( + description="Concatenation of the returned stdout followed by the returned stderr. Output from the two streams is not interleaved.", + examples=["logs output"], + ), + ] stderr: Annotated[str, Field(examples=["stderr output"])] stdout: Annotated[str, Field(examples=["stdout output"])] @@ -327,7 +337,13 @@ class ProcessRequest(BaseModel): class ProcessResponse(BaseModel): command: Annotated[str, Field(examples=["ls -la"])] - completedAt: Annotated[str, Field(examples=["Wed, 01 Jan 2023 12:01:00 GMT"])] + completedAt: Annotated[ + str, + Field( + description="Completion time, same format as startedAt. Empty string while the process runs", + examples=["Wed, 01 Jan 2023 12:01:00 GMT"], + ), + ] exitCode: Annotated[int, Field(examples=[0])] keepAlive: Annotated[ bool | None, @@ -342,7 +358,13 @@ class ProcessResponse(BaseModel): pid: Annotated[str, Field(examples=["1234"])] restartCount: Annotated[int | None, Field(examples=[2])] = None restartOnFailure: Annotated[bool | None, Field(examples=[True])] = None - startedAt: Annotated[str, Field(examples=["Wed, 01 Jan 2023 12:00:00 GMT"])] + startedAt: Annotated[ + str, + Field( + description="Start time as an HTTP date (RFC 1123, e.g. Wed, 01 Jan 2023 12:00:00 GMT)", + examples=["Wed, 01 Jan 2023 12:00:00 GMT"], + ), + ] status: Annotated[ Literal["failed", "killed", "stopped", "running", "completed"] | str, Field(examples=["running"]), @@ -451,6 +473,12 @@ class GetFilesystemContentSearchParams(BaseModel): description="Comma-separated directory names to skip (default: node_modules,vendor,.git,dist,build,target,__pycache__,.venv,.next,coverage)" ), ] = None + contextLines: Annotated[ + int | None, + Field( + description="Lines to include before and after each match in its context field (default: 0, max: 20; invalid values count as 0)" + ), + ] = None class GetFilesystemFindParams(BaseModel): @@ -484,13 +512,21 @@ class PutFilesystemMultipartPartParams(BaseModel): class GetFilesystemSearchParams(BaseModel): + query: Annotated[ + str | None, + Field( + description="Fuzzy pattern matched against each relative path (e.g., mngo for src/main.go). When omitted, the search path itself is used as the pattern." + ), + ] = None maxResults: Annotated[ int | None, Field(description="Maximum number of results to return (default: 20)"), ] = None patterns: Annotated[ str | None, - Field(description="Comma-separated file patterns to include (e.g., *.go,*.js)"), + Field( + description="Accepted for compatibility but currently ignored; use filesystem-find for glob filtering" + ), ] = None excludeDirs: Annotated[ str | None, @@ -524,7 +560,10 @@ class DeleteFilesystemTreeParams(BaseModel): class GetWatchFilesystemParams(BaseModel): ignore: Annotated[ - str | None, Field(description="Ignore patterns (comma-separated)") + str | None, + Field( + description="Comma-separated substrings; events whose full path contains one are skipped" + ), ] = None @@ -775,14 +814,6 @@ class GetFilesystemResponse(RootModel[Directory | FileWithContent | bytes]): root: Directory | FileWithContent | bytes -class GetFilesystemTreeResponse(RootModel[Directory | FileWithContent | bytes]): - root: Directory | FileWithContent | bytes - - -class PutFilesystemTreeResponse(RootModel[Directory | FileWithContent | bytes]): - root: Directory | FileWithContent | bytes - - class ArchiveChange(BaseModel): kind: Annotated[ Literal["added", "modified", "deleted"] | str, Field(examples=["added"]) @@ -843,12 +874,24 @@ class ExportResult(BaseModel): class HealthResponse(BaseModel): arch: Annotated[str, Field(examples=["amd64"])] - buildTime: Annotated[str, Field(examples=["2026-01-29T17:36:52+00:00"])] + buildTime: Annotated[ + str, + Field( + description='Build time in RFC 3339 (UTC), or "unknown" for builds without it', + examples=["2026-01-29T17:36:52+00:00"], + ), + ] gitCommit: Annotated[str, Field(examples=["abc123"])] goVersion: Annotated[str, Field(examples=["go1.25.0"])] lastUpgrade: UpgradeStatus os: Annotated[str, Field(examples=["linux"])] - startedAt: Annotated[str, Field(examples=["2026-01-29T18:45:49+00:00"])] + startedAt: Annotated[ + str, + Field( + description="API start time in RFC 3339", + examples=["2026-01-29T18:45:49+00:00"], + ), + ] status: Annotated[str, Field(examples=["ok"])] upgradeCount: Annotated[int, Field(examples=[0])] uptime: Annotated[str, Field(examples=["1h30m"])] diff --git a/baseten/sandbox/__init__.py b/baseten/sandbox/__init__.py new file mode 100644 index 0000000..158d7bf --- /dev/null +++ b/baseten/sandbox/__init__.py @@ -0,0 +1,186 @@ +"""Create and work with Baseten sandboxes. + +Create one :class:`SandboxClient` up front and reuse it, then create a sandbox +and run a command in it:: + + from baseten.sandbox import SandboxClient + + with SandboxClient(api_key="my-api-key") as client: + sandbox = client.create() + process = sandbox.process.exec( + command="echo hello", wait_for_completion=True + ) + print(process.stdout) + +:class:`AsyncSandboxClient` is the same for async code. HTTP/2 is used when the +``h2`` package is installed, which is recommended: ``pip install +"httpx[http2]"``. +""" + +from baseten.sandbox._auth import ( + AsyncSandboxTokenProvider, + SandboxTokenProvider, + SandboxTokenProviderContext, +) +from baseten.sandbox._client import ( + AsyncSandboxClient, + AsyncSandboxCreateResult, + SandboxClient, + SandboxClientOptions, + SandboxCreateResult, +) +from baseten.sandbox._common import AsyncSandboxStream, SandboxStream +from baseten.sandbox._errors import ( + ImageBuildError, + ImageUploadError, + SandboxAPIError, + SandboxError, + SandboxGatewayError, + SandboxProcessWaitTimeoutError, +) +from baseten.sandbox._fs import ( + AsyncSandboxFileSystem, + SandboxFileSystem, + SandboxFileSystemCopyError, + SandboxFileSystemDirectory, + SandboxFileSystemEntryType, + SandboxFileSystemFileInfo, + SandboxFileSystemFindMatch, + SandboxFileSystemFindResult, + SandboxFileSystemGrepMatch, + SandboxFileSystemGrepResult, + SandboxFileSystemSubdirectoryInfo, + SandboxFileSystemWatchEvent, + SandboxFileSystemWatchOp, +) +from baseten.sandbox._images import ( + AsyncImageClient, + ImageCleanupResult, + ImageClient, + ImageInfo, + ImageLibraryInfo, + ImageLibraryVolume, + ImageLogLine, + ImageSort, + ImageStatus, + ImageTagInfo, + ImageTagSort, +) +from baseten.sandbox._images_builder import ImageBuilder +from baseten.sandbox._images_ignore import ( + AsyncImageIgnoreFileFunc, + AsyncImageIgnoreFileProcessor, + ImageIgnoreFileFunc, + ImageIgnoreFileOptions, + ImageIgnoreFileProcessor, + ImageIgnoreFileProcessorOptions, + default_image_ignore_file, +) +from baseten.sandbox._info import ( + SandboxEnvValue, + SandboxExpirationAction, + SandboxExpirationPolicy, + SandboxExpirationPolicyDate, + SandboxExpirationPolicyIdle, + SandboxExpirationPolicyMaxAge, + SandboxExpirationPolicyUnknown, + SandboxInfo, + SandboxLifecycle, + SandboxNetwork, + SandboxNetworkProxy, + SandboxNetworkProxyRoute, + SandboxPort, + SandboxPortProtocol, + SandboxStatus, +) +from baseten.sandbox._process import ( + AsyncSandboxProcess, + SandboxProcess, + SandboxProcessExecEvent, + SandboxProcessExecEventExit, + SandboxProcessExecEventOutput, + SandboxProcessInfo, + SandboxProcessLogLine, + SandboxProcessLogs, + SandboxProcessStatus, +) +from baseten.sandbox._retry import SandboxRetryOptions +from baseten.sandbox._sandbox import AsyncSandbox, Sandbox + +__all__ = [ + "AsyncImageClient", + "AsyncImageIgnoreFileFunc", + "AsyncImageIgnoreFileProcessor", + "AsyncSandbox", + "AsyncSandboxClient", + "AsyncSandboxCreateResult", + "AsyncSandboxFileSystem", + "AsyncSandboxProcess", + "AsyncSandboxStream", + "AsyncSandboxTokenProvider", + "ImageBuildError", + "ImageBuilder", + "ImageCleanupResult", + "ImageClient", + "ImageIgnoreFileFunc", + "ImageIgnoreFileOptions", + "ImageIgnoreFileProcessor", + "ImageIgnoreFileProcessorOptions", + "ImageInfo", + "ImageLibraryInfo", + "ImageLibraryVolume", + "ImageLogLine", + "ImageSort", + "ImageStatus", + "ImageTagInfo", + "ImageTagSort", + "ImageUploadError", + "Sandbox", + "SandboxAPIError", + "SandboxClient", + "SandboxClientOptions", + "SandboxCreateResult", + "SandboxEnvValue", + "SandboxError", + "SandboxExpirationAction", + "SandboxExpirationPolicy", + "SandboxExpirationPolicyDate", + "SandboxExpirationPolicyIdle", + "SandboxExpirationPolicyMaxAge", + "SandboxExpirationPolicyUnknown", + "SandboxFileSystem", + "SandboxFileSystemCopyError", + "SandboxFileSystemDirectory", + "SandboxFileSystemEntryType", + "SandboxFileSystemFileInfo", + "SandboxFileSystemFindMatch", + "SandboxFileSystemFindResult", + "SandboxFileSystemGrepMatch", + "SandboxFileSystemGrepResult", + "SandboxFileSystemSubdirectoryInfo", + "SandboxFileSystemWatchEvent", + "SandboxFileSystemWatchOp", + "SandboxGatewayError", + "SandboxInfo", + "SandboxLifecycle", + "SandboxNetwork", + "SandboxNetworkProxy", + "SandboxNetworkProxyRoute", + "SandboxPort", + "SandboxPortProtocol", + "SandboxProcess", + "SandboxProcessExecEvent", + "SandboxProcessExecEventExit", + "SandboxProcessExecEventOutput", + "SandboxProcessInfo", + "SandboxProcessLogLine", + "SandboxProcessLogs", + "SandboxProcessStatus", + "SandboxProcessWaitTimeoutError", + "SandboxRetryOptions", + "SandboxStatus", + "SandboxStream", + "SandboxTokenProvider", + "SandboxTokenProviderContext", + "default_image_ignore_file", +] diff --git a/baseten/sandbox/_auth.py b/baseten/sandbox/_auth.py new file mode 100644 index 0000000..740bcdf --- /dev/null +++ b/baseten/sandbox/_auth.py @@ -0,0 +1,283 @@ +from __future__ import annotations + +import asyncio +import threading +import time +from collections.abc import AsyncGenerator, Awaitable, Callable, Generator +from dataclasses import dataclass +from datetime import UTC, datetime, timedelta +from typing import TypeAlias + +import httpx + +import baseten.client.managementapi + + +@dataclass(kw_only=True) +class SandboxTokenProviderContext: + """What a token provider is told about the token it is asked for.""" + + revoked_token: str | None = None + """A token the server just rejected as revoked. + + Set when the request is being sent again with a new token. Revocation can + happen well before a token's expiry, for example when any API key of the + token's user changes. + """ + + +SandboxTokenProvider: TypeAlias = Callable[[SandboxTokenProviderContext], str] +"""Returns the bearer token to send on a request. Called for every request. + +A provider that caches tokens must drop its cached token when it matches +:attr:`SandboxTokenProviderContext.revoked_token`, since the server rejects +that token for good. +""" + +AsyncSandboxTokenProvider: TypeAlias = Callable[ + [SandboxTokenProviderContext], Awaitable[str] +] +"""Async form of :data:`SandboxTokenProvider`.""" + +# Refresh this long before the token's stated expiry, to allow for clock skew +# between this machine and the server, and for time spent in flight. +_TOKEN_EXPIRY_LEEWAY = timedelta(seconds=60) + +# A few resends cover a new token being invalidated in the same event as the +# old one. Past that, the rejection is returned, since tokens that keep getting +# rejected point at something a new token cannot fix. +_TOKEN_INVALIDATION_MAX_RETRIES = 2 + +# Revocation rejects every token issued before a cutoff, the time of the +# revoking event 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 = 1.0 + +TOKEN_MINT_TIMEOUT = 30.0 +"""Seconds that one token exchange may take. + +No caller's cancellation reaches a shared exchange, so this keeps a stalled one +from holding up every later caller. +""" + + +@dataclass(frozen=True) +class _MintedToken: + """A token minted from an API key, with its expiry.""" + + token: str + expires_at: datetime + + def fresh(self) -> bool: + """Report whether the token is far enough from expiry to use.""" + return datetime.now(UTC) < self.expires_at - _TOKEN_EXPIRY_LEEWAY + + +class TokenSource: + """Where each request's bearer token comes from. + + Either a caller's token provider, or a token minted from an API key and + 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_api: baseten.client.managementapi.ApiClient, + ) -> None: + """Create a token source. ``mint_api`` authenticates with the API key.""" + 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_api = mint_api if api_key != "" else None + self._lock = threading.Lock() + self._cached: _MintedToken | None = None + + def token(self, revoked_token: str | None = None) -> str | None: + """Return the token for a request. + + ``None`` means no Authorization is sent, which is the advanced opt-out + of an empty API key with no token provider. + """ + if self._provider is not None: + return self._provider( + SandboxTokenProviderContext(revoked_token=revoked_token) + ) + if self._mint_api is None: + return None + # Held across the mint, so a burst on a cold cache sends one exchange. + with self._lock: + cached = self._cached + if cached is not None and cached.fresh(): + return cached.token + minted = self._mint_api.post_token( + request=baseten.client.managementapi.CreateTokenRequest( + scopes=["sandboxes"] + ) + ) + self._cached = _MintedToken( + token=minted.token, expires_at=minted.expires_at + ) + return minted.token + + def invalidate(self, token: str) -> None: + """Drop a token the server rejected, so the next request gets a new one.""" + with self._lock: + # Compared by value, so a rejection of an older token never + # discards a newer one another request already minted. + if self._cached is not None and self._cached.token == token: + self._cached = None + + +class AsyncTokenSource: + """Async form of :class:`TokenSource`.""" + + def __init__( + self, + *, + api_key: str, + token_provider: AsyncSandboxTokenProvider | None, + mint_api: baseten.client.managementapi.AsyncApiClient, + ) -> None: + """Create a token source. ``mint_api`` authenticates with the API key.""" + 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_api = mint_api if api_key != "" else None + self._cached: _MintedToken | None = None + self._minting: asyncio.Task[_MintedToken] | None = None + + async def token(self, revoked_token: str | None = None) -> str | None: + """Return the token for a request, as :meth:`TokenSource.token` does.""" + if self._provider is not None: + return await self._provider( + SandboxTokenProviderContext(revoked_token=revoked_token) + ) + if self._mint_api is None: + return None + cached = self._cached + if cached is not None and cached.fresh(): + return cached.token + # Every caller shares one mint, so a burst on a cold cache sends one + # exchange. Shielded, since one caller being cancelled must not cancel + # the mint the others are waiting on. + if self._minting is None: + self._minting = asyncio.create_task(self._mint(self._mint_api)) + # Retrieved here so a failure no caller is left to see is not + # reported as never retrieved. + self._minting.add_done_callback( + lambda task: task.cancelled() or task.exception() + ) + minted = await asyncio.shield(self._minting) + return minted.token + + async def invalidate(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. + if self._cached is not None and self._cached.token == token: + self._cached = None + + async def _mint( + self, mint_api: baseten.client.managementapi.AsyncApiClient + ) -> _MintedToken: + """Exchange the API key for a token and cache it.""" + try: + minted = await mint_api.post_token( + request=baseten.client.managementapi.CreateTokenRequest( + scopes=["sandboxes"] + ) + ) + self._cached = _MintedToken( + token=minted.token, expires_at=minted.expires_at + ) + return self._cached + finally: + self._minting = None + + +class TokenAuth(httpx.Auth): + """Authenticates every request with a token from a :class:`TokenSource`.""" + + def __init__(self, tokens: TokenSource) -> None: + """Create the auth.""" + self.tokens = tokens + + def sync_auth_flow( + self, request: httpx.Request + ) -> Generator[httpx.Request, httpx.Response, None]: + """Send the request, again with a new token if its token was revoked.""" + replayable = _is_replayable(request) + revoked_token: str | None = None + for retry in range(_TOKEN_INVALIDATION_MAX_RETRIES + 1): + token = self.tokens.token(revoked_token) + _set_authorization(request, token) + response = yield request + if token is None or not replayable or not _is_token_revoked(response): + return + if retry == _TOKEN_INVALIDATION_MAX_RETRIES: + return + self.tokens.invalidate(token) + revoked_token = token + if retry > 0: + time.sleep(_token_revoked_retry_delay) + + +class AsyncTokenAuth(httpx.Auth): + """Authenticates every request with a token from an :class:`AsyncTokenSource`.""" + + def __init__(self, tokens: AsyncTokenSource) -> None: + """Create the auth.""" + self.tokens = tokens + + async def async_auth_flow( + self, request: httpx.Request + ) -> AsyncGenerator[httpx.Request, httpx.Response]: + """Send the request, again with a new token if its token was revoked.""" + replayable = _is_replayable(request) + revoked_token: str | None = None + for retry in range(_TOKEN_INVALIDATION_MAX_RETRIES + 1): + token = await self.tokens.token(revoked_token) + _set_authorization(request, token) + response = yield request + if token is None or not replayable or not _is_token_revoked(response): + return + if retry == _TOKEN_INVALIDATION_MAX_RETRIES: + return + await self.tokens.invalidate(token) + revoked_token = token + if retry > 0: + await asyncio.sleep(_token_revoked_retry_delay) + + +def _is_replayable(request: httpx.Request) -> bool: + """Report whether a request's body can be sent again.""" + # A stream body can be read only once, so a request carrying one is never + # sent again. The SDK itself never sends one; only raw API calls can. + return isinstance(request.stream, httpx.ByteStream) + + +def _set_authorization(request: httpx.Request, token: str | None) -> None: + """Set the request's bearer token, or leave Authorization unset for ``None``.""" + if token is not None: + request.headers["Authorization"] = f"Bearer {token}" + + +def _is_token_revoked(response: httpx.Response) -> bool: + """Report whether the server rejected a request's token as revoked. + + The request was rejected before any work was done, so sending it again + with a new token is safe even when it has side effects. + """ + # Both the control plane and sandboxes send this header, while their + # bodies differ. + return ( + response.status_code == 401 + and response.headers.get("x-blaxel-error-code") == "TOKEN_REVOKED" + ) diff --git a/baseten/sandbox/_client.py b/baseten/sandbox/_client.py new file mode 100644 index 0000000..937833c --- /dev/null +++ b/baseten/sandbox/_client.py @@ -0,0 +1,853 @@ +from __future__ import annotations + +import importlib.util +from collections.abc import ( + AsyncIterator, + Iterator, + Mapping, + Sequence, +) +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Self, overload + +import httpx + +import baseten.client +import baseten.client.managementapi +import baseten.client.sandboxapi +from baseten.client._user_agent import with_user_agent +from baseten.sandbox._auth import ( + TOKEN_MINT_TIMEOUT, + AsyncSandboxTokenProvider, + AsyncTokenAuth, + AsyncTokenSource, + SandboxTokenProvider, + TokenAuth, + TokenSource, +) +from baseten.sandbox._common import ( + AsyncSandboxContext, + SandboxContext, + paginate, + paginate_async, +) +from baseten.sandbox._errors import converted_errors +from baseten.sandbox._images import ( + AsyncImageClient, + AsyncImageClientContext, + ImageClient, + ImageClientContext, +) +from baseten.sandbox._info import ( + SandboxEnvValue, + SandboxInfo, + SandboxLifecycle, + SandboxNetwork, + SandboxPort, + SandboxStatus, + envs_to_api, + lifecycle_to_api, + network_to_api, + ports_to_api, + sandbox_info_from_api, + set_fields, +) +from baseten.sandbox._retry import SandboxRetryOptions, resolve_retry_options +from baseten.sandbox._sandbox import AsyncSandbox, Sandbox + +# Timeouts of the HTTP client a sandbox client creates. No read timeout: a +# process run that waits for completion sends response headers only when the +# process exits. No write timeout either: httpcore's sync HTTP/2 sets each +# timeout on the shared socket before every read and write, so a write +# timeout set by one thread turns another thread's blocking read non-blocking, +# failing it and the connection. Connecting gets longer than httpx's 5s +# default, while pool waits keep it. +_DEFAULT_TIMEOUT = httpx.Timeout(5.0, connect=10.0, read=None, write=None) + + +@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.""" + + team_id: str | None = None + """Team to act in, or ``None`` for the caller's only team.""" + + token_provider: SandboxTokenProvider | AsyncSandboxTokenProvider | None = None + """Provider of each request's bearer token, in place of the API key.""" + + base_url_override: str | None = None + """Explicit base URL override, or ``None`` to use the default.""" + + headers: Mapping[str, str] | None = None + """Additional headers to send on every request.""" + + retry: SandboxRetryOptions | None = None + """Retry budgets for transient failures, or ``None`` for the defaults.""" + + http2: bool | None = None + """Whether the client's own HTTP client uses HTTP/2.""" + + @property + def base_url(self) -> str: + """The resolved base URL for the management API.""" + if self.base_url_override is not None: + return self.base_url_override + return baseten.client.ManagementClient.default_base_url() + + +class SandboxCreateResult(Sandbox): + """A newly created sandbox, with its record as of creation.""" + + info: SandboxInfo + """The sandbox as reported when creation returned.""" + + def __init__(self, context: SandboxContext, info: SandboxInfo) -> None: + """Internal. Returned by :meth:`SandboxClient.create`. + + :meta private: + """ + super().__init__(context) + self.info = info + + +class AsyncSandboxCreateResult(AsyncSandbox): + """Async form of :class:`SandboxCreateResult`.""" + + info: SandboxInfo + """The sandbox as reported when creation returned.""" + + def __init__(self, context: AsyncSandboxContext, info: SandboxInfo) -> None: + """Internal. Returned by :meth:`AsyncSandboxClient.create`. + + :meta private: + """ + super().__init__(context) + self.info = info + + +class SandboxClient: + """Client for Baseten sandboxes. + + Creates, finds, and deletes sandboxes, and gets a :class:`Sandbox` to work + in one. Can be used as a context manager to ensure the underlying HTTP + client is closed on exit. + """ + + def __init__( + self, + *, + api_key: str, + team_id: str | None = None, + token_provider: SandboxTokenProvider | None = None, + base_url_override: str | None = None, + headers: Mapping[str, str] | None = None, + retry: SandboxRetryOptions | None = None, + http2: bool | None = None, + http_client_override: httpx.Client | None = None, + close_http_client_on_close: bool | None = None, + executor: ThreadPoolExecutor | None = None, + ) -> None: + """Create a new sandbox client. + + Args: + api_key: API key for authentication. Empty, together with + *token_provider*, is the advanced opt-out of API key + authentication. + team_id: Team to act in. ``None`` uses the caller's only team, + and is an error for callers in more than one team. + token_provider: Returns the bearer token for each request, instead + of one derived from the API key, which must then be empty. + base_url_override: Override the management API's base URL, which + defaults to ``https://api.baseten.co``. + headers: Additional headers to send on every request. + retry: Retry budgets for transient failures. ``None`` uses the + defaults. + http2: Whether the HTTP client this client creates uses HTTP/2, + which matters for performance. ``None`` uses it when the + ``h2`` package is installed, as with ``httpx[http2]``. + ``True`` requires it, and ``False`` never uses it. Has no + effect with *http_client_override*. + http_client_override: Pre-configured httpx client to send every + request through, such as another sandbox client's + :attr:`http_client` to share its connections. Requests carry + absolute URLs and their own authentication and headers, so it + needs none of its own. Its timeouts apply: a read timeout + cuts off long process runs and streams, and with HTTP/2, read + and write timeouts that differ break requests made from + several threads at once, such as a large file's upload parts. + close_http_client_on_close: Whether :meth:`close` should close + the underlying HTTP client. Defaults to ``True`` when the + client is created internally, ``False`` when + *http_client_override* is provided. + executor: Thread pool for uploading a large file's parts in + parallel, on top of the calling thread. Never shut down by + this client. ``None`` gives each :class:`Sandbox` its own pool + of 2 threads, created on its first large upload, whose idle + threads end once the sandbox is no longer referenced. + """ + self._executor = executor + self._options = SandboxClientOptions( + api_key=api_key, + team_id=team_id, + token_provider=token_provider, + base_url_override=base_url_override, + headers=headers, + retry=retry, + http2=http2, + ) + if http_client_override is None: + self._http_client = httpx.Client( + http2=_use_http2(http2), timeout=_DEFAULT_TIMEOUT + ) + self.close_http_client_on_close = ( + True + if close_http_client_on_close is None + else close_http_client_on_close + ) + else: + self._http_client = http_client_override + self.close_http_client_on_close = ( + False + if close_http_client_on_close is None + else close_http_client_on_close + ) + self._request_headers = with_user_agent(headers or {}) + mint_api = baseten.client.managementapi.ApiClient( + self._http_client, + base_url=self._options.base_url, + headers={**self._request_headers, "Authorization": f"Bearer {api_key}"}, + timeout=httpx.Timeout(TOKEN_MINT_TIMEOUT), + ) + self._auth = TokenAuth( + TokenSource( + api_key=api_key, token_provider=token_provider, mint_api=mint_api + ) + ) + self._raw_api = baseten.client.managementapi.ApiClient( + self._http_client, + base_url=self._options.base_url, + auth=self._auth, + headers=self._request_headers, + ) + self._retry = resolve_retry_options(retry) + self._images = ImageClient( + ImageClientContext( + raw_api=self._raw_api, + team_id=team_id, + http_client=self._http_client, + ) + ) + + @property + def options(self) -> SandboxClientOptions: + """Client options.""" + return self._options + + @property + def http_client(self) -> httpx.Client: + """The underlying HTTP client.""" + return self._http_client + + @property + def images(self) -> ImageClient: + """Images that sandboxes are created from.""" + return self._images + + @property + def raw_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._raw_api + + def create( + self, + *, + name: str | None = None, + create_if_not_exists: bool | None = None, + image: str | None = None, + memory: int | None = None, + region: str | None = None, + envs: Mapping[str, SandboxEnvValue] | None = None, + labels: Mapping[str, str] | None = None, + external_id: str | None = None, + lifecycle: SandboxLifecycle | None = None, + ports: Sequence[SandboxPort] | None = None, + network: SandboxNetwork | None = None, + ) -> SandboxCreateResult: + """Create a sandbox. + + Args: + name: Unique name of the sandbox. Assigned by the server when + ``None``. + create_if_not_exists: When true, return the existing live sandbox + with this name, or recreate it if it is failed, terminated, or + being deleted. Requires *name*. Existing configuration is + preserved. + image: Image reference including its tag. Defaults to + ``baseten/base-image:latest``. + memory: Memory in megabytes, which also sets the CPU allocation. + Defaults to 4096. + region: Region to run in. ``None`` picks the closest one. + envs: Environment variables, by name. + labels: Labels, by key. + external_id: Caller-owned identifier for external lookups. + lifecycle: When the sandbox is deleted automatically. + ports: Ports the sandbox exposes. + network: Network configuration, which cannot be changed after + creation. + """ + with self._control_api() as api: + sandbox = api.create_sandbox( + params=baseten.client.managementapi.CreateSandboxParams( + **set_fields(team_id=self._options.team_id) + ), + request=baseten.client.managementapi.CreateSandboxRequest( + **set_fields( + name=name, + create_if_not_exists=create_if_not_exists, + image=image, + memory=memory, + region=region, + envs=None if envs is None else envs_to_api(envs), + labels=None if labels is None else dict(labels), + external_id=external_id, + lifecycle=None + if lifecycle is None + else lifecycle_to_api(lifecycle), + ports=None if ports is None else ports_to_api(ports), + network=None if network is None else network_to_api(network), + ) + ), + ) + info = sandbox_info_from_api(sandbox) + return SandboxCreateResult(self._sandbox_context(info.name, info.url), info) + + def get_info(self, *, name: str, show_secrets: bool | None = None) -> SandboxInfo: + """Get a sandbox's current record. + + Args: + name: Name of the sandbox. + show_secrets: Reveal environment variable values for workspace + administrators. Defaults to false. Callers without the admin + role receive masked values even when true. + """ + with self._control_api() as api: + sandbox = api.get_sandbox( + sandbox_name=name, + params=baseten.client.managementapi.GetSandboxParams( + **set_fields( + team_id=self._options.team_id, show_secrets=show_secrets + ) + ), + ) + return sandbox_info_from_api(sandbox) + + @overload + def get(self, *, name: str) -> Sandbox: ... + + @overload + def get(self, *, info: SandboxInfo) -> Sandbox: ... + + def get( + self, *, name: str | None = None, info: SandboxInfo | None = None + ) -> Sandbox: + """Get a :class:`Sandbox` to work in. + + Fetches the sandbox's record when given a name, and makes no call when + given the record. + + Args: + name: Name of the sandbox. + info: The sandbox's record, already in hand. + """ + if info is None: + if name is None: + raise TypeError("exactly one of name or info must be given") + info = self.get_info(name=name) + elif name is not None: + raise TypeError("exactly one of name or info must be given") + return Sandbox(self._sandbox_context(info.name, info.url)) + + def sandbox_from_url(self, *, url: str) -> Sandbox: + """Get a :class:`Sandbox` for a sandbox's execution API URL, making no call. + + Its :attr:`Sandbox.name` is empty. + + Args: + url: Base URL of the sandbox's execution API, as reported in + :attr:`SandboxInfo.url`. + """ + return Sandbox(self._sandbox_context("", url)) + + def list( + self, + *, + query: str | None = None, + statuses: Sequence[SandboxStatus] | None = None, + external_id: str | None = None, + page_size: int | None = None, + ) -> Iterator[SandboxInfo]: + """List sandboxes, fetching further pages as iteration reaches them. + + Args: + query: Searches sandbox names and labels. + statuses: Only sandboxes with one of these statuses. Cannot be + combined with *external_id*. + external_id: Only the sandbox with this external identifier. + Cannot be combined with *statuses*. + page_size: How many sandboxes to fetch per underlying request. + """ + + def fetch_page( + cursor: str | None, + ) -> baseten.client.managementapi.ListSandboxesResponse: + with self._control_api() as api: + return api.list_sandboxes( + params=baseten.client.managementapi.ListSandboxesParams( + **set_fields( + team_id=self._options.team_id, + cursor=cursor, + limit=page_size, + q=query, + status=None if statuses is None else [*statuses], + external_id=external_id, + ) + ) + ) + + for sandbox in paginate("sandbox list", fetch_page): + yield sandbox_info_from_api(sandbox) + + def update( + self, + *, + name: str, + lifecycle: SandboxLifecycle | None = None, + envs: Mapping[str, SandboxEnvValue] | None = None, + external_id: str | None = None, + labels: Mapping[str, str] | None = None, + ) -> SandboxInfo: + """Update a sandbox's configuration, returning its record after the change. + + Arguments left ``None`` stay unchanged, and set ones replace their + previous values, including every entry of *envs*, *labels*, and the + lifecycle's expiration policies. Other lifecycle fields change only + when set. At least one argument besides *name* must be set. Other + settings cannot be changed after creation. + + Args: + name: Name of the sandbox. + lifecycle: When the sandbox is deleted automatically. + envs: Environment variables, by name. + external_id: Caller-owned identifier for external lookups. + labels: Labels, by key. + """ + with self._control_api() as api: + sandbox = api.update_sandbox( + sandbox_name=name, + params=baseten.client.managementapi.UpdateSandboxParams( + **set_fields(team_id=self._options.team_id) + ), + request=baseten.client.managementapi.UpdateSandboxRequest( + **set_fields( + lifecycle=None + if lifecycle is None + else lifecycle_to_api(lifecycle), + envs=None if envs is None else envs_to_api(envs), + external_id=external_id, + labels=None if labels is None else dict(labels), + ) + ), + ) + return sandbox_info_from_api(sandbox) + + def delete(self, *, name: str) -> SandboxInfo: + """Delete a sandbox. + + Deletion continues after this returns, so the returned record usually + shows it as still deleting. + + Args: + name: Name of the sandbox. + """ + with self._control_api() as api: + sandbox = api.delete_sandbox( + sandbox_name=name, + params=baseten.client.managementapi.DeleteSandboxParams( + **set_fields(team_id=self._options.team_id) + ), + ) + return sandbox_info_from_api(sandbox) + + def close(self) -> None: + """Close the client, optionally closing the underlying HTTP client.""" + if self.close_http_client_on_close: + self._http_client.close() + + def __enter__(self) -> Self: + return self + + def __exit__(self, *args: object) -> None: + self.close() + + @contextmanager + def _control_api(self) -> Iterator[baseten.client.managementapi.ApiClient]: + """Give the generated client, raising its error responses as ours.""" + with converted_errors("control"): + yield self._raw_api + + def _sandbox_context(self, name: str, url: str) -> SandboxContext: + """Build what a :class:`Sandbox` for the given sandbox sends with.""" + return SandboxContext( + name=name, + url=url, + http_client=self._http_client, + auth=self._auth, + headers=self._request_headers, + raw_api=baseten.client.sandboxapi.ApiClient( + self._http_client, + base_url=url, + auth=self._auth, + headers=self._request_headers, + ), + retry=self._retry, + executor=self._executor, + ) + + +class AsyncSandboxClient: + """Async form of :class:`SandboxClient`. + + Can be used as an async context manager to ensure the underlying HTTP + client is closed on exit. + """ + + def __init__( + self, + *, + api_key: str, + team_id: str | None = None, + token_provider: AsyncSandboxTokenProvider | None = None, + base_url_override: str | None = None, + headers: Mapping[str, str] | None = None, + retry: SandboxRetryOptions | None = None, + http2: bool | None = None, + http_client_override: httpx.AsyncClient | None = None, + close_http_client_on_close: bool | None = None, + ) -> None: + """Create a new async sandbox client. + + Args: + api_key: API key for authentication. Empty, together with + *token_provider*, is the advanced opt-out of API key + authentication. + team_id: Team to act in. ``None`` uses the caller's only team, + and is an error for callers in more than one team. + token_provider: Returns the bearer token for each request, instead + of one derived from the API key, which must then be empty. + base_url_override: Override the management API's base URL, which + defaults to ``https://api.baseten.co``. + headers: Additional headers to send on every request. + retry: Retry budgets for transient failures. ``None`` uses the + defaults. + http2: Whether the HTTP client this client creates uses HTTP/2, + which matters for performance. ``None`` uses it when the + ``h2`` package is installed, as with ``httpx[http2]``. + ``True`` requires it, and ``False`` never uses it. Has no + effect with *http_client_override*. + http_client_override: Pre-configured httpx async client to send + every request through, such as another sandbox client's + :attr:`http_client` to share its connections. Requests carry + absolute URLs and their own authentication and headers, so it + needs none of its own. Its timeouts apply, and a read timeout + cuts off long process runs and streams. + close_http_client_on_close: Whether :meth:`close` should close + the underlying HTTP client. Defaults to ``True`` when the + client is created internally, ``False`` when + *http_client_override* is provided. + """ + self._options = SandboxClientOptions( + api_key=api_key, + team_id=team_id, + token_provider=token_provider, + base_url_override=base_url_override, + headers=headers, + retry=retry, + http2=http2, + ) + if http_client_override is None: + self._http_client = httpx.AsyncClient( + http2=_use_http2(http2), timeout=_DEFAULT_TIMEOUT + ) + self.close_http_client_on_close = ( + True + if close_http_client_on_close is None + else close_http_client_on_close + ) + else: + self._http_client = http_client_override + self.close_http_client_on_close = ( + False + if close_http_client_on_close is None + else close_http_client_on_close + ) + self._request_headers = with_user_agent(headers or {}) + mint_api = baseten.client.managementapi.AsyncApiClient( + self._http_client, + base_url=self._options.base_url, + headers={**self._request_headers, "Authorization": f"Bearer {api_key}"}, + timeout=httpx.Timeout(TOKEN_MINT_TIMEOUT), + ) + self._auth = AsyncTokenAuth( + AsyncTokenSource( + api_key=api_key, token_provider=token_provider, mint_api=mint_api + ) + ) + self._raw_api = baseten.client.managementapi.AsyncApiClient( + self._http_client, + base_url=self._options.base_url, + auth=self._auth, + headers=self._request_headers, + ) + self._retry = resolve_retry_options(retry) + self._images = AsyncImageClient( + AsyncImageClientContext( + raw_api=self._raw_api, + team_id=team_id, + http_client=self._http_client, + ) + ) + + @property + def options(self) -> SandboxClientOptions: + """Client options.""" + return self._options + + @property + def http_client(self) -> httpx.AsyncClient: + """The underlying HTTP client.""" + return self._http_client + + @property + def images(self) -> AsyncImageClient: + """Images that sandboxes are created from.""" + return self._images + + @property + def raw_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._raw_api + + async def create( + self, + *, + name: str | None = None, + create_if_not_exists: bool | None = None, + image: str | None = None, + memory: int | None = None, + region: str | None = None, + envs: Mapping[str, SandboxEnvValue] | None = None, + labels: Mapping[str, str] | None = None, + external_id: str | None = None, + lifecycle: SandboxLifecycle | None = None, + ports: Sequence[SandboxPort] | None = None, + network: SandboxNetwork | None = None, + ) -> AsyncSandboxCreateResult: + """Create a sandbox. As :meth:`SandboxClient.create`.""" + with self._control_api() as api: + sandbox = await api.create_sandbox( + params=baseten.client.managementapi.CreateSandboxParams( + **set_fields(team_id=self._options.team_id) + ), + request=baseten.client.managementapi.CreateSandboxRequest( + **set_fields( + name=name, + create_if_not_exists=create_if_not_exists, + image=image, + memory=memory, + region=region, + envs=None if envs is None else envs_to_api(envs), + labels=None if labels is None else dict(labels), + external_id=external_id, + lifecycle=None + if lifecycle is None + else lifecycle_to_api(lifecycle), + ports=None if ports is None else ports_to_api(ports), + network=None if network is None else network_to_api(network), + ) + ), + ) + info = sandbox_info_from_api(sandbox) + return AsyncSandboxCreateResult( + self._sandbox_context(info.name, info.url), info + ) + + async def get_info( + self, *, name: str, show_secrets: bool | None = None + ) -> SandboxInfo: + """Get a sandbox's current record. As :meth:`SandboxClient.get_info`.""" + with self._control_api() as api: + sandbox = await api.get_sandbox( + sandbox_name=name, + params=baseten.client.managementapi.GetSandboxParams( + **set_fields( + team_id=self._options.team_id, show_secrets=show_secrets + ) + ), + ) + return sandbox_info_from_api(sandbox) + + @overload + async def get(self, *, name: str) -> AsyncSandbox: ... + + @overload + async def get(self, *, info: SandboxInfo) -> AsyncSandbox: ... + + async def get( + self, *, name: str | None = None, info: SandboxInfo | None = None + ) -> AsyncSandbox: + """Get an :class:`AsyncSandbox` to work in. As :meth:`SandboxClient.get`.""" + if info is None: + if name is None: + raise TypeError("exactly one of name or info must be given") + info = await self.get_info(name=name) + elif name is not None: + raise TypeError("exactly one of name or info must be given") + return AsyncSandbox(self._sandbox_context(info.name, info.url)) + + def sandbox_from_url(self, *, url: str) -> AsyncSandbox: + """Get an :class:`AsyncSandbox` for a URL. As :meth:`SandboxClient.sandbox_from_url`.""" + return AsyncSandbox(self._sandbox_context("", url)) + + async def list( + self, + *, + query: str | None = None, + statuses: Sequence[SandboxStatus] | None = None, + external_id: str | None = None, + page_size: int | None = None, + ) -> AsyncIterator[SandboxInfo]: + """List sandboxes. As :meth:`SandboxClient.list`.""" + + async def fetch_page( + cursor: str | None, + ) -> baseten.client.managementapi.ListSandboxesResponse: + with self._control_api() as api: + return await api.list_sandboxes( + params=baseten.client.managementapi.ListSandboxesParams( + **set_fields( + team_id=self._options.team_id, + cursor=cursor, + limit=page_size, + q=query, + status=None if statuses is None else [*statuses], + external_id=external_id, + ) + ) + ) + + async for sandbox in paginate_async("sandbox list", fetch_page): + yield sandbox_info_from_api(sandbox) + + async def update( + self, + *, + name: str, + lifecycle: SandboxLifecycle | None = None, + envs: Mapping[str, SandboxEnvValue] | None = None, + external_id: str | None = None, + labels: Mapping[str, str] | None = None, + ) -> SandboxInfo: + """Update a sandbox's configuration. As :meth:`SandboxClient.update`.""" + with self._control_api() as api: + sandbox = await api.update_sandbox( + sandbox_name=name, + params=baseten.client.managementapi.UpdateSandboxParams( + **set_fields(team_id=self._options.team_id) + ), + request=baseten.client.managementapi.UpdateSandboxRequest( + **set_fields( + lifecycle=None + if lifecycle is None + else lifecycle_to_api(lifecycle), + envs=None if envs is None else envs_to_api(envs), + external_id=external_id, + labels=None if labels is None else dict(labels), + ) + ), + ) + return sandbox_info_from_api(sandbox) + + async def delete(self, *, name: str) -> SandboxInfo: + """Delete a sandbox. As :meth:`SandboxClient.delete`.""" + with self._control_api() as api: + sandbox = await api.delete_sandbox( + sandbox_name=name, + params=baseten.client.managementapi.DeleteSandboxParams( + **set_fields(team_id=self._options.team_id) + ), + ) + return sandbox_info_from_api(sandbox) + + async def close(self) -> None: + """Close the client, optionally closing the underlying HTTP client.""" + if self.close_http_client_on_close: + await self._http_client.aclose() + + async def __aenter__(self) -> Self: + return self + + async def __aexit__(self, *args: object) -> None: + await self.close() + + @contextmanager + def _control_api(self) -> Iterator[baseten.client.managementapi.AsyncApiClient]: + """Give the generated client, raising its error responses as ours.""" + # A plain context manager suffices, since it only wraps awaited calls. + with converted_errors("control"): + yield self._raw_api + + def _sandbox_context(self, name: str, url: str) -> AsyncSandboxContext: + """Build what an :class:`AsyncSandbox` for the given sandbox sends with.""" + return AsyncSandboxContext( + name=name, + url=url, + http_client=self._http_client, + auth=self._auth, + headers=self._request_headers, + raw_api=baseten.client.sandboxapi.AsyncApiClient( + self._http_client, + base_url=url, + auth=self._auth, + headers=self._request_headers, + ), + retry=self._retry, + ) + + +def _use_http2(http2: bool | None) -> bool: + """Resolve the ``http2`` option against whether ``h2`` is installed.""" + if http2 is False: + return False + available = importlib.util.find_spec("h2") is not None + if http2 and not available: + raise ImportError( + 'http2=True requires the h2 package, installed with "httpx[http2]"' + ) + return available diff --git a/baseten/sandbox/_common.py b/baseten/sandbox/_common.py new file mode 100644 index 0000000..27649c0 --- /dev/null +++ b/baseten/sandbox/_common.py @@ -0,0 +1,440 @@ +from __future__ import annotations + +import codecs +import email.utils +import re +from collections.abc import ( + AsyncGenerator, + AsyncIterator, + Awaitable, + Callable, + Generator, + Iterator, + Mapping, + Sequence, +) +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from dataclasses import dataclass +from datetime import datetime, timedelta +from typing import Any, Generic, Protocol, Self, TypeVar + +import httpx + +import baseten.client.sandboxapi +from baseten.sandbox._errors import SandboxError, converted_errors +from baseten.sandbox._retry import ( + ResolvedRetryOptions, + retry_idempotent, + retry_idempotent_async, +) + +_T = TypeVar("_T") +_ItemT = TypeVar("_ItemT") +_ItemT_co = TypeVar("_ItemT_co", covariant=True) + +_LINE_BREAK = re.compile(r"\r?\n") + +_HTTP_DATE = re.compile( + r"(Mon|Tue|Wed|Thu|Fri|Sat|Sun), [0-9]{2} " + r"(Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec) [0-9]{4} " + r"[0-9]{2}:[0-9]{2}:[0-9]{2} GMT" +) + +_DURATION_UNIT_MICROSECONDS = { + "ns": 0.001, + "us": 1.0, + "µs": 1.0, + "μs": 1.0, + "ms": 1_000.0, + "s": 1_000_000.0, + "m": 60_000_000.0, + "h": 3_600_000_000.0, +} + +_WHOLE_DAYS_OR_WEEKS = re.compile(r"([0-9]+)([dw])") + +_DURATION_PART = re.compile(r"([0-9]+(?:\.[0-9]*)?|\.[0-9]+)(ns|us|µs|μs|ms|s|m|h)") + + +def format_duration(value: timedelta) -> str: + """Format a control plane duration, rounded to whole milliseconds.""" + return f"{round(value / timedelta(milliseconds=1))}ms" + + +def parse_duration(value: str, what: str) -> timedelta: + """Parse a control plane duration, naming it as ``what`` in errors.""" + # Go duration syntax, or a whole number of days or weeks on its own, at + # 24 hours per day. + sign = -1 if value.startswith("-") else 1 + body = value[1:] if value.startswith(("-", "+")) else value + if body == "0": + return timedelta(0) + whole = _WHOLE_DAYS_OR_WEEKS.fullmatch(body) + if whole is not None: + days = int(whole[1]) * (7 if whole[2] == "w" else 1) + return sign * timedelta(days=days) + total_microseconds = 0.0 + index = 0 + while index < len(body): + part = _DURATION_PART.match(body, index) + if part is None: + break + total_microseconds += float(part[1]) * _DURATION_UNIT_MICROSECONDS[part[2]] + index = part.end() + if body == "" or index < len(body): + raise SandboxError(f"{what} is not a valid duration: {value!r}") + return sign * timedelta(microseconds=total_microseconds) + + +def parse_timestamp(value: str, what: str) -> datetime: + """Parse an execution API timestamp, naming it as ``what`` in errors.""" + # The API declares no format, and its examples differ: HTTP dates, RFC + # 3339, and RFC 3339 with a space for the T. A time without a zone is + # rejected rather than guessed at. + if _HTTP_DATE.fullmatch(value): + return email.utils.parsedate_to_datetime(value) + try: + parsed = datetime.fromisoformat(value) + except ValueError: + parsed = None + if parsed is None or parsed.tzinfo is None: + raise SandboxError(f"{what} is not a valid timestamp: {value!r}") + return parsed + + +@dataclass(kw_only=True) +class SandboxContext: + """What a :class:`Sandbox` and its subsystems send requests with.""" + + name: str + url: str + http_client: httpx.Client + auth: httpx.Auth + headers: Mapping[str, str] + raw_api: baseten.client.sandboxapi.ApiClient + """The sandbox's execution API, authenticated.""" + retry: ResolvedRetryOptions + executor: ThreadPoolExecutor | None + """The caller's thread pool for upload parts, if given.""" + + @contextmanager + def api(self) -> Iterator[baseten.client.sandboxapi.ApiClient]: + """Give the generated client, raising its error responses as ours.""" + with converted_errors("exec"): + yield self.raw_api + + @contextmanager + def api_with_timeout( + self, timeout: timedelta + ) -> Iterator[baseten.client.sandboxapi.ApiClient]: + """As :meth:`api`, with requests timing out after ``timeout``.""" + with converted_errors("exec"): + yield baseten.client.sandboxapi.ApiClient( + self.http_client, + base_url=self.url, + auth=self.auth, + headers=self.headers, + timeout=httpx.Timeout(timeout.total_seconds()), + ) + + def call_read( + self, call: Callable[[baseten.client.sandboxapi.ApiClient], _T] + ) -> _T: + """Run an idempotent call, retrying it on the read budgets.""" + + def attempt() -> _T: + with self.api() as api: + return call(api) + + return retry_idempotent( + attempt, + max_retries=self.retry.read_max_retries, + gateway_max_retries=self.retry.gateway_max_retries, + ) + + def call_upload( + self, call: Callable[[baseten.client.sandboxapi.ApiClient], _T] + ) -> _T: + """Run an idempotent upload, retrying it on the upload budgets.""" + + def attempt() -> _T: + with self.api() as api: + return call(api) + + return retry_idempotent( + attempt, + max_retries=self.retry.upload_max_retries, + gateway_max_retries=self.retry.gateway_max_retries, + ) + + def send(self, request: httpx.Request) -> None: + """Send a request the generated client cannot build, raising an error response as ours.""" + with converted_errors("exec"): + response = self.http_client.send(request, auth=self.auth) + if not response.is_success: + raise baseten.client.sandboxapi.ResponseError( + status_code=response.status_code, body=response.text + ) + + def build_request(self, method: str, path: str, **kwargs: Any) -> httpx.Request: + """Build a request to a path of the sandbox, with the client's headers.""" + return self.http_client.build_request( + method, self.url.rstrip("/") + path, headers=self.headers, **kwargs + ) + + +@dataclass(kw_only=True) +class AsyncSandboxContext: + """Async form of :class:`SandboxContext`.""" + + name: str + url: str + http_client: httpx.AsyncClient + auth: httpx.Auth + headers: Mapping[str, str] + raw_api: baseten.client.sandboxapi.AsyncApiClient + """The sandbox's execution API, authenticated.""" + retry: ResolvedRetryOptions + + @contextmanager + def api(self) -> Iterator[baseten.client.sandboxapi.AsyncApiClient]: + """Give the generated client, raising its error responses as ours.""" + # A plain context manager suffices, since it only wraps awaited calls. + with converted_errors("exec"): + yield self.raw_api + + @contextmanager + def api_with_timeout( + self, timeout: timedelta + ) -> Iterator[baseten.client.sandboxapi.AsyncApiClient]: + """As :meth:`api`, with requests timing out after ``timeout``.""" + with converted_errors("exec"): + yield baseten.client.sandboxapi.AsyncApiClient( + self.http_client, + base_url=self.url, + auth=self.auth, + headers=self.headers, + timeout=httpx.Timeout(timeout.total_seconds()), + ) + + async def call_read( + self, call: Callable[[baseten.client.sandboxapi.AsyncApiClient], Awaitable[_T]] + ) -> _T: + """Run an idempotent call, retrying it on the read budgets.""" + + async def attempt() -> _T: + with self.api() as api: + return await call(api) + + return await retry_idempotent_async( + attempt, + max_retries=self.retry.read_max_retries, + gateway_max_retries=self.retry.gateway_max_retries, + ) + + async def call_upload( + self, call: Callable[[baseten.client.sandboxapi.AsyncApiClient], Awaitable[_T]] + ) -> _T: + """Run an idempotent upload, retrying it on the upload budgets.""" + + async def attempt() -> _T: + with self.api() as api: + return await call(api) + + return await retry_idempotent_async( + attempt, + max_retries=self.retry.upload_max_retries, + gateway_max_retries=self.retry.gateway_max_retries, + ) + + async def send(self, request: httpx.Request) -> None: + """Send a request the generated client cannot build, raising an error response as ours.""" + with converted_errors("exec"): + response = await self.http_client.send(request, auth=self.auth) + if not response.is_success: + raise baseten.client.sandboxapi.ResponseError( + status_code=response.status_code, body=response.text + ) + + def build_request(self, method: str, path: str, **kwargs: Any) -> httpx.Request: + """Build a request to a path of the sandbox, with the client's headers.""" + return self.http_client.build_request( + method, self.url.rstrip("/") + path, headers=self.headers, **kwargs + ) + + +class SandboxStream(Generic[_ItemT]): + """Items a sandbox sends as they happen, read by iterating. + + Reading to the end, or an error while reading, closes the stream. To stop + early, use it as a context manager or call :meth:`close`. Closing stops + only the stream, never what it reports on, such as a running process. + """ + + def __init__( + self, + response: httpx.Response, + parse: Callable[[Iterator[str]], Generator[_ItemT, None, None]], + ) -> None: + """Internal. Returned by the methods that stream. + + :meta private: + """ + self._response = response + self._items = parse(response_lines(response)) + + def __iter__(self) -> Self: + return self + + def __next__(self) -> _ItemT: + try: + return next(self._items) + except BaseException: + # Includes the StopIteration at the end. + self.close() + raise + + def close(self) -> None: + """Stop reading and release the connection. Does nothing if already closed.""" + self._items.close() + # A generator closed before it started never runs its own cleanup, + # so the response is closed here rather than by the generator. + self._response.close() + + def __enter__(self) -> Self: + return self + + def __exit__(self, *args: object) -> None: + self.close() + + +class AsyncSandboxStream(Generic[_ItemT]): + """Async form of :class:`SandboxStream`.""" + + def __init__( + self, + response: httpx.Response, + parse: Callable[[AsyncIterator[str]], AsyncGenerator[_ItemT, None]], + ) -> None: + """Internal. Returned by the methods that stream. + + :meta private: + """ + self._response = response + self._items = parse(response_lines_async(response)) + + def __aiter__(self) -> Self: + return self + + async def __anext__(self) -> _ItemT: + try: + return await anext(self._items) + except BaseException: + # Includes the StopAsyncIteration at the end. + await self.close() + raise + + async def close(self) -> None: + """Stop reading and release the connection. Does nothing if already closed.""" + await self._items.aclose() + await self._response.aclose() + + async def __aenter__(self) -> Self: + return self + + async def __aexit__(self, *args: object) -> None: + await self.close() + + +class _Pagination(Protocol): + """Pagination fields of a control plane listing page.""" + + @property + def has_more(self) -> bool: ... + + @property + def cursor(self) -> str | None: ... + + +class _Page(Protocol[_ItemT_co]): + """A control plane listing page.""" + + @property + def items(self) -> Sequence[_ItemT_co]: ... + + @property + def pagination(self) -> _Pagination: ... + + +def paginate( + what: str, fetch_page: Callable[[str | None], _Page[_ItemT_co]] +) -> Iterator[_ItemT_co]: + """Yield every item of a cursor-paginated control plane listing. + + Each page is fetched as iteration reaches it. ``what`` names the listing + in errors. + """ + seen_cursors: set[str] = set() + cursor: str | None = None + while True: + page = fetch_page(cursor) + yield from page.items + cursor = _next_cursor(what, page, seen_cursors) + if cursor is None: + return + + +async def paginate_async( + what: str, fetch_page: Callable[[str | None], Awaitable[_Page[_ItemT_co]]] +) -> AsyncIterator[_ItemT_co]: + """Async form of :func:`paginate`.""" + seen_cursors: set[str] = set() + cursor: str | None = None + while True: + page = await fetch_page(cursor) + for item in page.items: + yield item + cursor = _next_cursor(what, page, seen_cursors) + if cursor is None: + return + + +def _next_cursor(what: str, page: _Page[Any], seen_cursors: set[str]) -> str | None: + """Return the cursor of the page after ``page``, or ``None`` after the last.""" + cursor = page.pagination.cursor if page.pagination.has_more else None + # A server repeating a cursor would otherwise page forever. + if cursor is not None: + if cursor in seen_cursors: + raise SandboxError(f"{what} returned a repeated cursor") + seen_cursors.add(cursor) + return cursor + + +def response_lines(response: httpx.Response) -> Generator[str, None, None]: + """Yield a response body's lines without line endings, including a last one with none.""" + # Only \n and \r\n end a line, unlike httpx's iter_lines, which also + # splits at a lone \r, such as a progress bar's. + decoder = codecs.getincrementaldecoder("utf-8")(errors="replace") + buffer = "" + for chunk in response.iter_bytes(): + buffer += decoder.decode(chunk) + *lines, buffer = _LINE_BREAK.split(buffer) + yield from lines + buffer += decoder.decode(b"", final=True) + if buffer: + yield buffer + + +async def response_lines_async(response: httpx.Response) -> AsyncGenerator[str, None]: + """Async form of :func:`response_lines`.""" + decoder = codecs.getincrementaldecoder("utf-8")(errors="replace") + buffer = "" + async for chunk in response.aiter_bytes(): + buffer += decoder.decode(chunk) + *lines, buffer = _LINE_BREAK.split(buffer) + for line in lines: + yield line + buffer += decoder.decode(b"", final=True) + if buffer: + yield buffer diff --git a/baseten/sandbox/_errors.py b/baseten/sandbox/_errors.py new file mode 100644 index 0000000..b641dd5 --- /dev/null +++ b/baseten/sandbox/_errors.py @@ -0,0 +1,231 @@ +from __future__ import annotations + +import json +from collections.abc import Iterator +from contextlib import contextmanager +from typing import Any, Literal, Self + +import baseten.client.managementapi +import baseten.client.sandboxapi + +GENERATED_RESPONSE_ERRORS = ( + baseten.client.managementapi.ResponseError, + baseten.client.sandboxapi.ResponseError, + baseten.client.sandboxapi.ResponseErrorResponse, +) +"""Errors the generated clients raise for an error response.""" + +_GATEWAY_STATUSES = frozenset({502, 503, 504}) + + +class SandboxError(Exception): + """Base class for every error the sandbox SDK raises itself.""" + + +class SandboxAPIError(SandboxError): + """An error response from the sandbox control plane or from a sandbox.""" + + status: int + """HTTP status of the response.""" + + code: str | None + """Machine-readable error code, when the response has one.""" + + message: str | None + """Description of the error, when the response has one.""" + + details: Any + """Additional error details, when the response has them.""" + + body: Any + """Response body, parsed as JSON when possible, otherwise the raw text.""" + + def __init__( + self, + *, + status: int, + code: str | None, + message: str | None, + details: Any, + body: Any, + ) -> None: + """Create an error. + + Args: + status: HTTP status of the response. + code: Machine-readable error code, when the response has one. + message: Description of the error, when the response has one. + details: Additional error details, when the response has them. + body: Response body, parsed as JSON when possible, otherwise the + raw text. + """ + text = f"sandbox API error (HTTP {status})" + if code is not None: + text += f" {code}" + if message is not None: + text += f": {message}" + super().__init__(text) + self.status = status + self.code = code + self.message = message + self.details = details + self.body = body + + @classmethod + def _from_response(cls, *, status: int, body: Any) -> Self: + """Create an error from a response, reading its fields from the body.""" + fields: dict[str, Any] = body if isinstance(body, dict) else {} + body_message = fields.get("message") + if not isinstance(body_message, str): + body_message = None + body_error = fields.get("error") + if not isinstance(body_error, str): + body_error = None + # The control plane puts the code in error and the description in + # message. A sandbox puts the description in error. + code = fields.get("code") + if not isinstance(code, str): + code = body_error if body_message is not None else None + return cls( + status=status, + code=code, + message=body_message if body_message is not None else body_error, + details=fields.get("details"), + body=body, + ) + + +class SandboxGatewayError(SandboxAPIError): + """A 502, 503, or 504 from the edge in front of a sandbox. + + Comes from the edge rather than from the sandbox itself, and is usually + transient, for example while a sandbox wakes from standby or when a + request outlasts the edge's timeout. + """ + + +class SandboxProcessWaitTimeoutError(SandboxError): + """A process still running when a wait for it ran out of time. + + The process keeps running regardless. The last error the wait rode out, + if any, is chained as the cause. + """ + + identifier: str + """PID or name of the process waited on.""" + + def __init__(self, *, identifier: str) -> None: + """Create an error. + + Args: + identifier: PID or name of the process waited on. + """ + super().__init__( + f"process {identifier} had not finished when the wait timed out; " + "it may still be running" + ) + self.identifier = identifier + + +class ImageUploadError(SandboxError): + """A failed upload of an image's build context. + + The upload goes to storage rather than to the API, so its error body has + no fixed shape. The image stays as it was left, so push it again or + delete it. + """ + + image_name: str + """Name of the image whose build context was being uploaded.""" + + status: int + """HTTP status of the storage response.""" + + body: str + """Raw text of the storage response body.""" + + def __init__(self, *, image_name: str, status: int, body: str) -> None: + """Create an error. + + Args: + image_name: Name of the image whose build context was being + uploaded. + status: HTTP status of the storage response. + body: Raw text of the storage response body. + """ + super().__init__( + f"uploading the source of image {image_name} failed (HTTP {status}); " + "the image stays as it was left, so push it again or delete it" + ) + self.image_name = image_name + self.status = status + self.body = body + + +class ImageBuildError(SandboxError): + """An image that did not become ready to use. + + Either its build failed, or it was still processing when the wait ran out + of time. :meth:`ImageClient.logs` shows the build's output. + """ + + image_name: str + """Name of the image.""" + + status: str + """Status last seen, ``FAILED`` when the build failed.""" + + timed_out: bool + """Whether the wait ran out of time. Processing continues regardless.""" + + def __init__(self, *, image_name: str, status: str, timed_out: bool) -> None: + """Create an error. + + Args: + image_name: Name of the image. + status: Status last seen, ``FAILED`` when the build failed. + timed_out: Whether the wait ran out of time. + """ + super().__init__( + f"image {image_name} was still {status} when the wait timed out; " + "it may still finish" + if timed_out + else f"image {image_name} failed to build (status {status})" + ) + self.image_name = image_name + self.status = status + self.timed_out = timed_out + + +@contextmanager +def converted_errors(plane: Literal["control", "exec"]) -> Iterator[None]: + """Raise generated clients' error responses as :class:`SandboxAPIError`. + + Anything else, such as a network failure, passes through unchanged. + """ + try: + yield + except GENERATED_RESPONSE_ERRORS as err: + raise api_error_from(err, plane) from err + + +def api_error_from( + err: baseten.client.managementapi.ResponseError + | baseten.client.sandboxapi.ResponseError + | baseten.client.sandboxapi.ResponseErrorResponse, + plane: Literal["control", "exec"], +) -> SandboxAPIError: + """Convert an error a generated client raised for an error response.""" + body: Any + if isinstance(err, baseten.client.sandboxapi.ResponseErrorResponse): + body = err.error_response.model_dump(mode="json") + else: + try: + body = json.loads(err.body) + except ValueError: + body = err.body + # Only a sandbox sits behind the edge, so the same statuses from the + # control plane are ordinary errors. + if plane == "exec" and err.status_code in _GATEWAY_STATUSES: + return SandboxGatewayError._from_response(status=err.status_code, body=body) + return SandboxAPIError._from_response(status=err.status_code, body=body) diff --git a/baseten/sandbox/_fs.py b/baseten/sandbox/_fs.py new file mode 100644 index 0000000..bd95f90 --- /dev/null +++ b/baseten/sandbox/_fs.py @@ -0,0 +1,1115 @@ +from __future__ import annotations + +import asyncio +import builtins +import concurrent.futures +import contextlib +import json +import posixpath +import threading +import urllib.parse +from collections.abc import AsyncGenerator, AsyncIterator, Generator, Iterator, Mapping +from dataclasses import dataclass +from datetime import datetime, timedelta +from typing import Literal, TypeAlias + +import httpx + +import baseten.client.sandboxapi +from baseten.sandbox._common import ( + AsyncSandboxContext, + AsyncSandboxStream, + SandboxContext, + SandboxStream, + parse_timestamp, +) +from baseten.sandbox._errors import SandboxError +from baseten.sandbox._info import set_fields +from baseten.sandbox._process import ( + AsyncSandboxProcess, + SandboxProcess, + SandboxProcessInfo, +) +from baseten.sandbox._retry import retry_idempotent, retry_idempotent_async + +_DEFAULT_COPY_TIMEOUT = timedelta(seconds=180) +_COPY_POLL_INTERVAL = timedelta(milliseconds=100) + +_MULTIPART_THRESHOLD = 5 * 1024 * 1024 +_MULTIPART_PART_SIZE = 5 * 1024 * 1024 + +# Bounds the abort of a failed multipart upload, which runs even when the +# caller gave up, so a stalled abort cannot hang the write it cleans up after. +_MULTIPART_ABORT_TIMEOUT = timedelta(seconds=10) + +# Most upload parts in flight at once per sandbox, across all of its uploads. +# Many at once on one HTTP/2 connection can trip the server's rapid reset +# limit, so this holds on every transport rather than guessing which is used. +_UPLOAD_PARTS_IN_FLIGHT = 2 + +SandboxFileSystemEntryType: TypeAlias = Literal["file", "directory"] | str +"""Kind of a filesystem entry. + +Other values may be added, so do not treat this list as exhaustive. +""" + +SandboxFileSystemWatchOp: TypeAlias = ( + Literal["CREATE", "WRITE", "REMOVE", "RENAME", "CHMOD"] | str +) +"""A kind of change in a :class:`SandboxFileSystemWatchEvent`. + +Other values may be added, so do not treat this list as exhaustive. +""" + + +class SandboxFileSystemCopyError(SandboxError): + """A copy by :meth:`SandboxFileSystem.copy` that did not succeed.""" + + source: str + """Path copied from.""" + + destination: str + """Path copied to.""" + + process: SandboxProcessInfo + """The ``cp`` process as it ended, whose ``logs`` hold its error output.""" + + def __init__( + self, *, source: str, destination: str, process: SandboxProcessInfo + ) -> None: + """Create an error. + + Args: + source: Path copied from. + destination: Path copied to. + process: The ``cp`` process as it ended. + """ + detail = process.stderr.strip() or f"exit code {process.exit_code}" + super().__init__( + f"copying {source} to {destination} ended {process.status}: {detail}" + ) + self.source = source + self.destination = destination + self.process = process + + +@dataclass(kw_only=True) +class SandboxFileSystemFileInfo: + """A file in a :class:`SandboxFileSystemDirectory`.""" + + name: str + + path: str + + size_bytes: int + + permissions: str + """File mode, as the sandbox reports it.""" + + owner: str + + group: str + + last_modified: datetime + + +@dataclass(kw_only=True) +class SandboxFileSystemSubdirectoryInfo: + """A subdirectory in a :class:`SandboxFileSystemDirectory`.""" + + name: str + + path: str + + +@dataclass(kw_only=True) +class SandboxFileSystemDirectory: + """A directory's contents, one level deep, from :meth:`SandboxFileSystem.list`.""" + + name: str + + path: str + + files: list[SandboxFileSystemFileInfo] + + subdirectories: list[SandboxFileSystemSubdirectoryInfo] + + +@dataclass(kw_only=True) +class SandboxFileSystemFindMatch: + """An entry found by :meth:`SandboxFileSystem.find`.""" + + path: str + """Path relative to the searched path.""" + + type: SandboxFileSystemEntryType + + +@dataclass(kw_only=True) +class SandboxFileSystemFindResult: + """Result of :meth:`SandboxFileSystem.find`.""" + + matches: list[SandboxFileSystemFindMatch] + + total: int + """Number of entries found, which may exceed the matches returned.""" + + +@dataclass(kw_only=True) +class SandboxFileSystemGrepMatch: + """A matching line found by :meth:`SandboxFileSystem.grep`.""" + + path: str + """Path relative to the searched path.""" + + line: int + """Line number, starting at 1.""" + + column: int + + text: str + """The matching line.""" + + context: str | None + """The matching line with the lines around it, newline-separated. + + Holds up to ``context_lines`` lines before and after the match, and is + ``None`` when ``context_lines`` was not set. + """ + + +@dataclass(kw_only=True) +class SandboxFileSystemGrepResult: + """Result of :meth:`SandboxFileSystem.grep`.""" + + matches: list[SandboxFileSystemGrepMatch] + + total: int + """Number of matching lines, which may exceed the matches returned.""" + + +@dataclass(kw_only=True) +class SandboxFileSystemWatchEvent: + """A change seen by :meth:`SandboxFileSystem.watch`.""" + + ops: list[SandboxFileSystemWatchOp] + """Kinds of change, more than one when the sandbox reports several together.""" + + path: str + """Full path of the changed entry.""" + + +class SandboxFileSystem: + """Files and directories in a sandbox. + + Relative paths are resolved against the sandbox's working directory. + """ + + def __init__(self, context: SandboxContext, process: SandboxProcess) -> None: + """Internal. Obtained from :attr:`Sandbox.fs`. + + :meta private: + """ + self._context = context + self._process = process + # One per sandbox, since a sandbox builds its file system once. + self._part_slots = threading.Semaphore(_UPLOAD_PARTS_IN_FLIGHT) + self._executor_lock = threading.Lock() + self._own_executor: concurrent.futures.ThreadPoolExecutor | None = None + + def read(self, *, path: str) -> str: + """Read a text file. + + Args: + path: Path of the file. + """ + result = self._context.call_read(lambda api: api.get_filesystem(path=path)) + return _file_content(path, result) + + def read_bytes(self, *, path: str) -> bytes: + """Read a file as bytes. + + Args: + path: Path of the file. + """ + + # The body is read inside the retried call, so a connection dropped + # partway through the body is retried too. + def read(api: baseten.client.sandboxapi.ApiClient) -> bytes: + response = api.get_filesystem_raw( + path=path, accept="application/octet-stream" + ) + try: + _check_not_directory(path, response) + return response.read() + finally: + response.close() + + return self._context.call_read(read) + + def write(self, *, path: str, content: str) -> None: + """Write a text file, creating it or replacing its content. + + Content over 5MB is uploaded in 5MB parts through the multipart upload + endpoints, at most 2 parts at a time per sandbox. + + Args: + path: Path of the file, created or replaced. + content: Text content to write. + """ + encoded = _encoded_if_large(content) + if encoded is not None: + self._write_multipart(path, encoded, None) + return + # Retried, since the spec documents the write as idempotent. + self._context.call_upload( + lambda api: api.put_filesystem( + path=path, + request=baseten.client.sandboxapi.FileRequest(content=content), + ) + ) + + def write_bytes( + self, *, path: str, content: bytes, permissions: str | None = None + ) -> None: + """Write a file from bytes, creating it or replacing its content. + + Content over 5MB is uploaded in 5MB parts through the multipart upload + endpoints, at most 2 parts at a time per sandbox. + + Args: + path: Path of the file, created or replaced. + content: Content to write. + permissions: Octal file mode, such as ``"0755"``, applied when the + file is created. An existing file keeps its mode. ``None`` + uses the sandbox's default. + """ + if len(content) > _MULTIPART_THRESHOLD: + self._write_multipart(path, content, permissions) + return + # The generated client has no multipart form for this endpoint, since + # it also takes a JSON body. Retried, since the spec documents the + # write as idempotent, with the request rebuilt for each attempt. + retry_idempotent( + lambda: self._context.send( + _write_bytes_request(self._context, path, content, permissions) + ), + max_retries=self._context.retry.upload_max_retries, + gateway_max_retries=self._context.retry.gateway_max_retries, + ) + + def mkdir(self, *, path: str, permissions: str | None = None) -> None: + """Create a directory, along with any missing parents. + + Args: + path: Path of the directory. + permissions: Octal directory mode, such as ``"0755"``. ``None`` + uses the sandbox's default. + """ + # Retried, since the spec documents the write as idempotent. + self._context.call_upload( + lambda api: api.put_filesystem( + path=path, request=_mkdir_request(permissions) + ) + ) + + def list(self, *, path: str) -> SandboxFileSystemDirectory: + """List a directory's files and subdirectories, one level deep. + + Args: + path: Path of the directory. + """ + result = self._context.call_read(lambda api: api.get_filesystem(path=path)) + return _directory(path, result) + + def remove(self, *, path: str, recursive: bool | None = None) -> None: + """Remove a file or directory. + + Args: + path: Path of the file or directory. + recursive: Whether to remove a directory and everything in it. + A directory that is not empty fails without it. + """ + # Not retried, since a repeat of a removal whose response was lost + # would fail for a path that is already gone. + with self._context.api() as api: + api.delete_filesystem( + path=path, + params=baseten.client.sandboxapi.DeleteFilesystemParams( + **set_fields(recursive=recursive) + ), + ) + + def find( + self, + *, + path: str, + type: SandboxFileSystemEntryType | None = None, + patterns: builtins.list[str] | None = None, + max_results: int | None = None, + exclude_dirs: builtins.list[str] | None = None, + exclude_hidden: bool | None = None, + ) -> SandboxFileSystemFindResult: + """Find files and directories by name under a directory. + + Args: + path: Path of the directory to search in. + type: Whether to find only files or only directories. ``None`` + finds both. + patterns: File patterns to include, such as ``"*.py"``. + max_results: Maximum number of results. 0 means all. The sandbox + defaults to 20. + exclude_dirs: Directory names to skip. The sandbox defaults to + common dependency and build directories, such as + ``node_modules`` and ``.git``. + exclude_hidden: Whether to skip hidden files and directories. The + sandbox defaults to true. + """ + params = _find_params(type, patterns, max_results, exclude_dirs, exclude_hidden) + result = self._context.call_read( + lambda api: api.get_filesystem_find(path=path, params=params) + ) + return _find_result(result) + + def grep( + self, + *, + path: str, + query: str, + case_sensitive: bool | None = None, + max_results: int | None = None, + file_pattern: str | None = None, + exclude_dirs: builtins.list[str] | None = None, + context_lines: int | None = None, + ) -> SandboxFileSystemGrepResult: + """Search the contents of files under a directory for text. + + Args: + path: Path of the directory to search in. + query: Text to search for. + case_sensitive: Whether the search is case sensitive. The sandbox + defaults to false. + max_results: Maximum number of results. The sandbox defaults to + 100. + file_pattern: File pattern to include, such as ``"*.py"``. + exclude_dirs: Directory names to skip. The sandbox defaults to + common dependency and build directories, such as + ``node_modules`` and ``.git``. + context_lines: How many lines before and after each match to + include in its :attr:`SandboxFileSystemGrepMatch.context`, at + most 20. ``None`` includes none. + """ + params = _grep_params( + query, + case_sensitive, + max_results, + file_pattern, + exclude_dirs, + context_lines, + ) + result = self._context.call_read( + lambda api: api.get_filesystem_content_search(path=path, params=params) + ) + return _grep_result(result) + + def copy( + self, + *, + source: str, + destination: str, + timeout: timedelta | None = _DEFAULT_COPY_TIMEOUT, + ) -> None: + """Copy a file or directory. + + Runs ``cp -r`` in the sandbox through :meth:`SandboxProcess.exec`, then + waits for it with :meth:`SandboxProcess.wait`, so ``cp``'s rules apply, + such as copying into a destination that is an existing directory. + + Args: + source: Path of the file or directory to copy. + destination: Path to copy to. + timeout: How long to wait for the copy, which keeps running + regardless. ``None`` waits with no limit. + + Raises: + SandboxFileSystemCopyError: The copy did not succeed. + SandboxProcessWaitTimeoutError: The copy was still running when + the wait timed out. + """ + started = self._process.exec(command=_copy_command(source, destination)) + finished = self._process.wait( + identifier=started.pid, timeout=timeout, poll_interval=_COPY_POLL_INTERVAL + ) + if finished.status != "completed" or finished.exit_code != 0: + raise SandboxFileSystemCopyError( + source=source, destination=destination, process=finished + ) + + def write_tree(self, *, path: str, files: Mapping[str, str]) -> None: + """Write several text files under a directory in one call, creating or replacing each. + + Args: + path: Path of the directory to write under. + files: Text content of each file, by path relative to *path*. + """ + # Retried, since the spec documents the write as idempotent. + self._context.call_upload( + lambda api: api.put_filesystem_tree( + path=path, + request=baseten.client.sandboxapi.TreeRequest(files=dict(files)), + ) + ) + + def watch( + self, + *, + path: str, + recursive: bool | None = None, + ignore: builtins.list[str] | None = None, + ) -> SandboxStream[SandboxFileSystemWatchEvent]: + """Stream changes in a directory as they happen. + + Only the directory's direct entries are watched unless *recursive* is + set. Closing the stream stops watching. + + Args: + path: Path of the directory to watch. + recursive: Whether to also watch every subdirectory, including + ones created later. + ignore: Skips events whose full path contains any of these + substrings. + """ + # Not retried, since a resent watch would miss changes in between. + with self._context.api() as api: + response = api.get_watch_filesystem( + path=_watch_path(path, recursive), params=_watch_params(ignore) + ) + return SandboxStream(response, _watch_events) + + def _write_multipart( + self, path: str, content: bytes, permissions: str | None + ) -> None: + """Upload content in parts through the multipart upload endpoints.""" + # Not retried, since a repeat would start a second upload. + with self._context.api() as api: + initiated = api.post_filesystem_multipart_initiate( + path=path, request=_multipart_initiate_request(permissions) + ) + upload_id = _upload_id(path, initiated) + try: + parts = self._send_parts(upload_id, content) + # Not retried, since a repeat of a completion whose response was + # lost would fail for an upload that is already complete. + with self._context.api() as api: + api.post_filesystem_multipart_complete( + upload_id=upload_id, + request=baseten.client.sandboxapi.MultipartCompleteRequest( + parts=parts + ), + ) + except BaseException: + # Sent even on an interrupt, which may be what ended the upload. + # Its own failure is dropped in favor of the error that ended it. + with ( + contextlib.suppress(Exception), + self._context.api_with_timeout(_MULTIPART_ABORT_TIMEOUT) as api, + ): + api.delete_filesystem_multipart_abort(upload_id=upload_id) + raise + + def _send_parts( + self, upload_id: str, content: bytes + ) -> builtins.list[baseten.client.sandboxapi.MultipartPartInfo]: + """Upload every part of content. The first part to fail stops the rest.""" + part_count = -(-len(content) // _MULTIPART_PART_SIZE) + parts: builtins.list[baseten.client.sandboxapi.MultipartPartInfo | None] = [ + None + ] * part_count + lock = threading.Lock() + next_index = 0 + stopped = threading.Event() + + def work() -> None: + nonlocal next_index + while True: + with lock: + if stopped.is_set() or next_index >= part_count: + return + index = next_index + next_index += 1 + part_number = index + 1 + try: + etag = self._send_part( + upload_id, + part_number, + content[ + index * _MULTIPART_PART_SIZE : part_number + * _MULTIPART_PART_SIZE + ], + ) + except BaseException: + stopped.set() + raise + parts[index] = baseten.client.sandboxapi.MultipartPartInfo( + partNumber=part_number, etag=etag + ) + + # The calling thread is one of the workers, so a caller's pool that is + # busy, even with this very call, slows the upload rather than + # deadlocking it. + other = self._executor().submit(work) if part_count > 1 else None + try: + work() + except BaseException: + # Lets none outlive this call. One that never started is + # cancelled instead, since it may be queued behind this very call. + if other is not None and not other.cancel(): + concurrent.futures.wait([other]) + raise + # The calling thread sent every part when the other worker never + # started. + if other is not None and not other.cancel(): + try: + other.result() + except BaseException: + # Also when interrupted while waiting, so the other worker + # stops after its current part. + stopped.set() + raise + return [part for part in parts if part is not None] + + def _send_part(self, upload_id: str, part_number: int, part: bytes) -> str | None: + """Upload one part and return its ETag.""" + + # The slot is taken per attempt, so a part waiting to retry does not + # hold one. + def send( + api: baseten.client.sandboxapi.ApiClient, + ) -> baseten.client.sandboxapi.MultipartUploadPartResponse: + with self._part_slots: + return api.put_filesystem_multipart_part( + upload_id=upload_id, + params=baseten.client.sandboxapi.PutFilesystemMultipartPartParams( + partNumber=part_number + ), + # The filename a browser's FormData gives a part sent + # without one. + files={"file": ("blob", part)}, + ) + + # Retried, since a repeat sends the same part. + return self._context.call_upload(send).etag + + def _executor(self) -> concurrent.futures.ThreadPoolExecutor: + """Give the caller's thread pool, or this sandbox's own, created on first use.""" + if self._context.executor is not None: + return self._context.executor + with self._executor_lock: + if self._own_executor is None: + # Its idle threads end once this sandbox is no longer + # referenced, since the pool is then collected. + self._own_executor = concurrent.futures.ThreadPoolExecutor( + max_workers=_UPLOAD_PARTS_IN_FLIGHT, + thread_name_prefix="baseten-sandbox-upload", + ) + return self._own_executor + + +class AsyncSandboxFileSystem: + """Async form of :class:`SandboxFileSystem`.""" + + def __init__( + self, context: AsyncSandboxContext, process: AsyncSandboxProcess + ) -> None: + """Internal. Obtained from :attr:`AsyncSandbox.fs`. + + :meta private: + """ + self._context = context + self._process = process + # One per sandbox, since a sandbox builds its file system once. + self._part_slots = asyncio.Semaphore(_UPLOAD_PARTS_IN_FLIGHT) + + async def read(self, *, path: str) -> str: + """Read a text file. As :meth:`SandboxFileSystem.read`.""" + result = await self._context.call_read( + lambda api: api.get_filesystem(path=path) + ) + return _file_content(path, result) + + async def read_bytes(self, *, path: str) -> bytes: + """Read a file as bytes. As :meth:`SandboxFileSystem.read_bytes`.""" + + async def read(api: baseten.client.sandboxapi.AsyncApiClient) -> bytes: + response = await api.get_filesystem_raw( + path=path, accept="application/octet-stream" + ) + try: + _check_not_directory(path, response) + return await response.aread() + finally: + await response.aclose() + + return await self._context.call_read(read) + + async def write(self, *, path: str, content: str) -> None: + """Write a text file. As :meth:`SandboxFileSystem.write`.""" + encoded = _encoded_if_large(content) + if encoded is not None: + await self._write_multipart(path, encoded, None) + return + await self._context.call_upload( + lambda api: api.put_filesystem( + path=path, + request=baseten.client.sandboxapi.FileRequest(content=content), + ) + ) + + async def write_bytes( + self, *, path: str, content: bytes, permissions: str | None = None + ) -> None: + """Write a file from bytes. As :meth:`SandboxFileSystem.write_bytes`.""" + if len(content) > _MULTIPART_THRESHOLD: + await self._write_multipart(path, content, permissions) + return + await retry_idempotent_async( + lambda: self._context.send( + _write_bytes_request(self._context, path, content, permissions) + ), + max_retries=self._context.retry.upload_max_retries, + gateway_max_retries=self._context.retry.gateway_max_retries, + ) + + async def mkdir(self, *, path: str, permissions: str | None = None) -> None: + """Create a directory. As :meth:`SandboxFileSystem.mkdir`.""" + await self._context.call_upload( + lambda api: api.put_filesystem( + path=path, request=_mkdir_request(permissions) + ) + ) + + async def list(self, *, path: str) -> SandboxFileSystemDirectory: + """List a directory. As :meth:`SandboxFileSystem.list`.""" + result = await self._context.call_read( + lambda api: api.get_filesystem(path=path) + ) + return _directory(path, result) + + async def remove(self, *, path: str, recursive: bool | None = None) -> None: + """Remove a file or directory. As :meth:`SandboxFileSystem.remove`.""" + with self._context.api() as api: + await api.delete_filesystem( + path=path, + params=baseten.client.sandboxapi.DeleteFilesystemParams( + **set_fields(recursive=recursive) + ), + ) + + async def find( + self, + *, + path: str, + type: SandboxFileSystemEntryType | None = None, + patterns: builtins.list[str] | None = None, + max_results: int | None = None, + exclude_dirs: builtins.list[str] | None = None, + exclude_hidden: bool | None = None, + ) -> SandboxFileSystemFindResult: + """Find files and directories by name. As :meth:`SandboxFileSystem.find`.""" + params = _find_params(type, patterns, max_results, exclude_dirs, exclude_hidden) + result = await self._context.call_read( + lambda api: api.get_filesystem_find(path=path, params=params) + ) + return _find_result(result) + + async def grep( + self, + *, + path: str, + query: str, + case_sensitive: bool | None = None, + max_results: int | None = None, + file_pattern: str | None = None, + exclude_dirs: builtins.list[str] | None = None, + context_lines: int | None = None, + ) -> SandboxFileSystemGrepResult: + """Search the contents of files. As :meth:`SandboxFileSystem.grep`.""" + params = _grep_params( + query, + case_sensitive, + max_results, + file_pattern, + exclude_dirs, + context_lines, + ) + result = await self._context.call_read( + lambda api: api.get_filesystem_content_search(path=path, params=params) + ) + return _grep_result(result) + + async def copy( + self, + *, + source: str, + destination: str, + timeout: timedelta | None = _DEFAULT_COPY_TIMEOUT, + ) -> None: + """Copy a file or directory. As :meth:`SandboxFileSystem.copy`.""" + started = await self._process.exec(command=_copy_command(source, destination)) + finished = await self._process.wait( + identifier=started.pid, timeout=timeout, poll_interval=_COPY_POLL_INTERVAL + ) + if finished.status != "completed" or finished.exit_code != 0: + raise SandboxFileSystemCopyError( + source=source, destination=destination, process=finished + ) + + async def write_tree(self, *, path: str, files: Mapping[str, str]) -> None: + """Write several text files. As :meth:`SandboxFileSystem.write_tree`.""" + await self._context.call_upload( + lambda api: api.put_filesystem_tree( + path=path, + request=baseten.client.sandboxapi.TreeRequest(files=dict(files)), + ) + ) + + async def watch( + self, + *, + path: str, + recursive: bool | None = None, + ignore: builtins.list[str] | None = None, + ) -> AsyncSandboxStream[SandboxFileSystemWatchEvent]: + """Stream changes in a directory. As :meth:`SandboxFileSystem.watch`.""" + with self._context.api() as api: + response = await api.get_watch_filesystem( + path=_watch_path(path, recursive), params=_watch_params(ignore) + ) + return AsyncSandboxStream(response, _watch_events_async) + + async def _write_multipart( + self, path: str, content: bytes, permissions: str | None + ) -> None: + """Upload content in parts through the multipart upload endpoints.""" + with self._context.api() as api: + initiated = await api.post_filesystem_multipart_initiate( + path=path, request=_multipart_initiate_request(permissions) + ) + upload_id = _upload_id(path, initiated) + try: + parts = await self._send_parts(upload_id, content) + with self._context.api() as api: + await api.post_filesystem_multipart_complete( + upload_id=upload_id, + request=baseten.client.sandboxapi.MultipartCompleteRequest( + parts=parts + ), + ) + except BaseException: + # Shielded, so the caller's cancellation, which may be what ended + # the upload, does not stop it. Its own failure is dropped in + # favor of the error that ended the upload. + with contextlib.suppress(Exception): + await asyncio.shield(self._abort_multipart(upload_id)) + raise + + async def _abort_multipart(self, upload_id: str) -> None: + """Abort a multipart upload, on its own timeout.""" + with self._context.api_with_timeout(_MULTIPART_ABORT_TIMEOUT) as api: + await api.delete_filesystem_multipart_abort(upload_id=upload_id) + + async def _send_parts( + self, upload_id: str, content: bytes + ) -> builtins.list[baseten.client.sandboxapi.MultipartPartInfo]: + """Upload every part of content. The first part to fail stops the rest.""" + part_count = -(-len(content) // _MULTIPART_PART_SIZE) + parts: builtins.list[baseten.client.sandboxapi.MultipartPartInfo | None] = [ + None + ] * part_count + next_index = 0 + + async def work() -> None: + nonlocal next_index + while next_index < part_count: + index = next_index + next_index += 1 + part_number = index + 1 + etag = await self._send_part( + upload_id, + part_number, + content[ + index * _MULTIPART_PART_SIZE : part_number + * _MULTIPART_PART_SIZE + ], + ) + parts[index] = baseten.client.sandboxapi.MultipartPartInfo( + partNumber=part_number, etag=etag + ) + + # More workers than slots would only wait for one. + workers = [ + asyncio.create_task(work()) + for _ in range(min(_UPLOAD_PARTS_IN_FLIGHT, part_count)) + ] + try: + await asyncio.wait(workers, return_when=asyncio.FIRST_EXCEPTION) + finally: + # Stops the rest after a failure or the caller's cancellation, and + # lets none outlive this call. + for worker in workers: + worker.cancel() + await asyncio.wait(workers) + for worker in workers: + if not worker.cancelled() and (error := worker.exception()) is not None: + raise error + return [part for part in parts if part is not None] + + async def _send_part( + self, upload_id: str, part_number: int, part: bytes + ) -> str | None: + """Upload one part and return its ETag.""" + + async def send( + api: baseten.client.sandboxapi.AsyncApiClient, + ) -> baseten.client.sandboxapi.MultipartUploadPartResponse: + async with self._part_slots: + return await api.put_filesystem_multipart_part( + upload_id=upload_id, + params=baseten.client.sandboxapi.PutFilesystemMultipartPartParams( + partNumber=part_number + ), + files={"file": ("blob", part)}, + ) + + return (await self._context.call_upload(send)).etag + + +def _file_content( + path: str, result: baseten.client.sandboxapi.GetFilesystemResponse +) -> str: + """Give a file's text from a filesystem read, failing for a directory.""" + entry = result.root + if isinstance(entry, baseten.client.sandboxapi.FileWithContent): + return entry.content + if isinstance(entry, bytes): + return entry.decode(errors="replace") + raise SandboxError(f"{path} is a directory, not a file") + + +def _check_not_directory(path: str, response: httpx.Response) -> None: + """Fail when a byte read returned a directory listing.""" + # A directory comes back as its JSON listing whatever was asked for. + if "application/json" in response.headers.get("content-type", ""): + raise SandboxError(f"{path} is a directory, not a file") + + +def _encoded_if_large(content: str) -> bytes | None: + """Give text as UTF-8 if that is over the multipart threshold, else ``None``.""" + # A code point is at most 4 bytes of UTF-8, so only a long string can be + # over the threshold, and shorter ones skip encoding just to measure. + if len(content) * 4 <= _MULTIPART_THRESHOLD: + return None + encoded = content.encode() + return encoded if len(encoded) > _MULTIPART_THRESHOLD else None + + +def _write_bytes_request( + context: SandboxContext | AsyncSandboxContext, + path: str, + content: bytes, + permissions: str | None, +) -> httpx.Request: + """Build the multipart form request that writes a file from bytes.""" + data = {"permissions": permissions} if permissions is not None else {} + data["path"] = path + return context.build_request( + "PUT", + "/filesystem/" + urllib.parse.quote(path, safe=""), + files={"file": (posixpath.basename(path) or "file", content)}, + data=data, + ) + + +def _mkdir_request(permissions: str | None) -> baseten.client.sandboxapi.FileRequest: + """Build the request that creates a directory.""" + return baseten.client.sandboxapi.FileRequest( + **set_fields(isDirectory=True, permissions=permissions) + ) + + +def _directory( + path: str, result: baseten.client.sandboxapi.GetFilesystemResponse +) -> SandboxFileSystemDirectory: + """Give a directory's listing from a filesystem read, failing for a file.""" + entry = result.root + if not isinstance(entry, baseten.client.sandboxapi.Directory): + raise SandboxError(f"{path} is a file, not a directory") + return SandboxFileSystemDirectory( + name=entry.name, + path=entry.path, + files=[ + SandboxFileSystemFileInfo( + name=file.name, + path=file.path, + size_bytes=file.size, + permissions=file.permissions, + owner=file.owner, + group=file.group, + last_modified=parse_timestamp( + file.lastModified, "file modification time" + ), + ) + for file in entry.files + ], + subdirectories=[ + SandboxFileSystemSubdirectoryInfo( + name=subdirectory.name, path=subdirectory.path + ) + for subdirectory in entry.subdirectories + ], + ) + + +def _find_params( + type: str | None, + patterns: builtins.list[str] | None, + max_results: int | None, + exclude_dirs: builtins.list[str] | None, + exclude_hidden: bool | None, +) -> baseten.client.sandboxapi.GetFilesystemFindParams: + """Build the query of a find.""" + return baseten.client.sandboxapi.GetFilesystemFindParams( + **set_fields( + type=type, + patterns=None if patterns is None else ",".join(patterns), + maxResults=max_results, + excludeDirs=None if exclude_dirs is None else ",".join(exclude_dirs), + excludeHidden=exclude_hidden, + ) + ) + + +def _find_result( + result: baseten.client.sandboxapi.FindResponse, +) -> SandboxFileSystemFindResult: + """Convert a find's response.""" + return SandboxFileSystemFindResult( + matches=[ + SandboxFileSystemFindMatch(path=match.path, type=match.type) + for match in result.matches + ], + total=result.total, + ) + + +def _grep_params( + query: str, + case_sensitive: bool | None, + max_results: int | None, + file_pattern: str | None, + exclude_dirs: builtins.list[str] | None, + context_lines: int | None, +) -> baseten.client.sandboxapi.GetFilesystemContentSearchParams: + """Build the query of a content search.""" + return baseten.client.sandboxapi.GetFilesystemContentSearchParams( + query=query, + **set_fields( + caseSensitive=case_sensitive, + maxResults=max_results, + filePattern=file_pattern, + excludeDirs=None if exclude_dirs is None else ",".join(exclude_dirs), + contextLines=context_lines, + ), + ) + + +def _grep_result( + result: baseten.client.sandboxapi.ContentSearchResponse, +) -> SandboxFileSystemGrepResult: + """Convert a content search's response.""" + return SandboxFileSystemGrepResult( + matches=[ + SandboxFileSystemGrepMatch( + path=match.path, + line=match.line, + column=match.column, + text=match.text, + context=match.context, + ) + for match in result.matches + ], + total=result.total, + ) + + +def _copy_command(source: str, destination: str) -> str: + """Build the shell command that copies source to destination.""" + # Each path quoted as one literal POSIX shell argument. + quoted = [ + "'" + value.replace("'", "'\\''") + "'" for value in (source, destination) + ] + return f"cp -r {quoted[0]} {quoted[1]}" + + +def _watch_path(path: str, recursive: bool | None) -> str: + """Give the watch endpoint's path, with the suffix that makes it recursive.""" + return path.rstrip("/") + "/**" if recursive else path + + +def _watch_params( + ignore: builtins.list[str] | None, +) -> baseten.client.sandboxapi.GetWatchFilesystemParams: + """Build the query of a watch.""" + return baseten.client.sandboxapi.GetWatchFilesystemParams( + **set_fields(ignore=None if ignore is None else ",".join(ignore)) + ) + + +def _watch_event(line: str) -> SandboxFileSystemWatchEvent | None: + """Parse one line of a watch, or ``None`` for a line to skip.""" + if line.strip() == "" or line.startswith("[keepalive]"): + return None + event = json.loads(line) + directory = event["path"] + return SandboxFileSystemWatchEvent( + ops=event["op"].split("|"), + path=directory.rstrip("/") + "/" + event["name"], + ) + + +def _watch_events( + lines: Iterator[str], +) -> Generator[SandboxFileSystemWatchEvent, None, None]: + """Parse a watch into events.""" + for line in lines: + event = _watch_event(line) + if event is not None: + yield event + + +async def _watch_events_async( + lines: AsyncIterator[str], +) -> AsyncGenerator[SandboxFileSystemWatchEvent, None]: + """Async form of :func:`_watch_events`.""" + async for line in lines: + event = _watch_event(line) + if event is not None: + yield event + + +def _multipart_initiate_request( + permissions: str | None, +) -> baseten.client.sandboxapi.MultipartInitiateRequest: + """Build the request that starts a multipart upload.""" + return baseten.client.sandboxapi.MultipartInitiateRequest( + **set_fields(permissions=permissions) + ) + + +def _upload_id( + path: str, initiated: baseten.client.sandboxapi.MultipartInitiateResponse +) -> str: + """Give a started multipart upload's ID, failing without one.""" + if not initiated.uploadId: + raise SandboxError(f"multipart upload of {path} returned no upload ID") + return initiated.uploadId diff --git a/baseten/sandbox/_images.py b/baseten/sandbox/_images.py new file mode 100644 index 0000000..1bf5cc2 --- /dev/null +++ b/baseten/sandbox/_images.py @@ -0,0 +1,1351 @@ +from __future__ import annotations + +import asyncio +import builtins +import contextlib +import inspect +import io +import tempfile +import time +from collections.abc import AsyncIterator, Iterator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass +from datetime import UTC, datetime, timedelta +from typing import IO, Literal, TypeAlias, overload + +import httpx + +import baseten.client.managementapi +from baseten.sandbox._common import paginate, paginate_async +from baseten.sandbox._errors import ( + ImageBuildError, + ImageUploadError, + SandboxAPIError, + SandboxError, + converted_errors, +) +from baseten.sandbox._images_builder import ImageBuilder +from baseten.sandbox._images_ignore import ( + AsyncImageIgnoreFileFunc, + AsyncImageIgnoreFileProcessor, + ImageIgnoreFileFunc, + ImageIgnoreFileProcessor, + default_image_ignore_file, + missing_processor_error, + read_ignore_file, +) +from baseten.sandbox._info import SandboxPort, ports_from_api, set_fields +from baseten.sandbox._retry import is_transient_reset_error +from baseten.sandbox._zip import ( + ZipEntry, + check_directory_dockerfile, + check_zip_dockerfile, + files_zip_entries, + run_walk, + run_walk_async, + walk_directory, + write_zip, +) + +_DEFAULT_WAIT_TIMEOUT = timedelta(minutes=15) +_DEFAULT_POLL_INTERVAL = timedelta(seconds=3) + +# The largest page and offset the build log endpoint accepts, which together +# cap a range at the 11,000 most recent lines. +_LOGS_PAGE_SIZE = 1000 +_LOGS_MAX_OFFSET = 10_000 + +# Statuses a wait rides out, since the image may not be visible yet or the +# gateway may briefly fail. +_WAIT_RETRY_STATUSES = frozenset({404, 502, 503, 504}) + +# Size of each read of a zip in a temporary file while uploading it. +_UPLOAD_CHUNK_SIZE = 1024 * 1024 + +ImageStatus: TypeAlias = Literal["UPLOADING", "BUILDING", "BUILT", "FAILED"] | str +"""Processing status of an image. Only ``BUILT`` images are ready to use. + +Other values may be added, so do not treat this list as exhaustive. +""" + +ImageSort: TypeAlias = ( + Literal["name:asc", "name:desc", "createdAt:asc", "createdAt:desc"] | str +) +"""Sort order of an image listing. + +Other values may be added, so do not treat this list as exhaustive. +""" + +ImageTagSort: TypeAlias = Literal["name:asc", "name:desc"] | str +"""Sort order of a tag listing. + +Other values may be added, so do not treat this list as exhaustive. +""" + + +@dataclass(kw_only=True) +class ImageInfo: + """An image repository as last reported by the control plane.""" + + name: str + """Repository name given when pushing. + + To create a sandbox from the image, pass ``:latest`` as its + ``image``, or ``:`` for a specific version. + """ + + status: ImageStatus + + created_at: datetime | None + + updated_at: datetime | None + + last_deployed_at: datetime | None + """Most recent time any version of the image was used by a sandbox.""" + + size_bytes: int | None + """Total size of all versions.""" + + tag_count: int | None + """Number of versions, each with its own tag.""" + + +@dataclass(kw_only=True) +class ImageTagInfo: + """One version of an image.""" + + name: str + """Tag name, assigned by the service.""" + + created_at: datetime | None + + updated_at: datetime | None + + size_bytes: int | None + + +@dataclass(kw_only=True) +class ImageCleanupResult: + """Result of :meth:`ImageClient.cleanup`.""" + + deleted: int + """Number of image versions removed.""" + + message: str + """Human-readable description of the result.""" + + +@dataclass(kw_only=True) +class ImageLogLine: + """One line of an image's build log, from :meth:`ImageClient.logs`.""" + + timestamp: datetime + + severity: int + """Numeric OpenTelemetry severity level.""" + + text: str + + +@dataclass(kw_only=True) +class ImageLibraryVolume: + """A volume attachment a built-in image suggests, in :class:`ImageLibraryInfo`.""" + + name: str + """Volume name, or an internal identifier for an ephemeral volume.""" + + mount_path: str + """Absolute path the volume is mounted at.""" + + type: str | None + """Volume type, persistent when empty or ``None``.""" + + size_mb: int | None + """Storage capacity in megabytes of an ephemeral volume.""" + + read_only: bool + """Whether the volume is mounted read-only.""" + + +@dataclass(kw_only=True) +class ImageLibraryInfo: + """A built-in image available to every team. + + Usable as a sandbox's image without building or pushing it. From + :meth:`ImageClient.list_library`. + """ + + name: str + """Stable identifier of the image.""" + + display_name: str | None + """Human-readable name.""" + + description: str | None + """Short description.""" + + long_description: str | None + """Detailed description.""" + + image: str + """Image reference including its tag, to pass as a sandbox's ``image``.""" + + memory: int | None + """Recommended memory allocation in megabytes.""" + + ports: list[SandboxPort] + """Ports the image exposes.""" + + categories: list[str] + + tags: list[str] + + url: str | None + """Documentation URL.""" + + icon: str | None + """Icon URL.""" + + icon_light: str | None + """Light-mode icon URL.""" + + icon_dark: str | None + """Dark-mode icon URL.""" + + enterprise: bool + """Whether the image requires an enterprise plan.""" + + creation_extra_args: dict[str, str] | None + """Kernel selection arguments suggested when creating a sandbox from this image. + + Creating a sandbox cannot pass these yet. + """ + + creation_volumes: list[ImageLibraryVolume] + """Volume attachments suggested when creating a sandbox from this image. + + Creating a sandbox cannot pass these yet. + """ + + +@dataclass(kw_only=True) +class ImageClientContext: + """What an :class:`ImageClient` uses from the client that made it.""" + + raw_api: baseten.client.managementapi.ApiClient + team_id: str | None + http_client: httpx.Client + """Sends storage uploads, which take none of the API's authentication.""" + + @contextmanager + def api(self) -> Iterator[baseten.client.managementapi.ApiClient]: + """Give the generated client, raising its error responses as ours.""" + with converted_errors("control"): + yield self.raw_api + + +@dataclass(kw_only=True) +class AsyncImageClientContext: + """Async form of :class:`ImageClientContext`.""" + + raw_api: baseten.client.managementapi.AsyncApiClient + team_id: str | None + http_client: httpx.AsyncClient + """Sends storage uploads, which take none of the API's authentication.""" + + @contextmanager + def api(self) -> Iterator[baseten.client.managementapi.AsyncApiClient]: + """Give the generated client, raising its error responses as ours.""" + with converted_errors("control"): + yield self.raw_api + + +@dataclass(kw_only=True) +class _Archive: + """A zipped build context, in memory or in an anonymous temporary file.""" + + file: IO[bytes] + """Positioned at the start. Closing it removes a temporary file.""" + + size: int + + +class ImageClient: + """Images that sandboxes are created from. + + Pushes them from a build context or a registry, and finds and deletes + them. Obtained from :attr:`SandboxClient.images`. + """ + + def __init__(self, context: ImageClientContext) -> None: + """Internal. Obtained from :attr:`SandboxClient.images`. + + :meta private: + """ + self._context = context + + @overload + def push( + self, + *, + name: str, + builder: ImageBuilder, + temp_file: bool | None = ..., + wait: bool = ..., + timeout: timedelta | None = ..., + poll_interval: timedelta = ..., + ) -> ImageInfo: ... + + @overload + def push( + self, + *, + name: str, + registry_image: str, + docker_config: str | None = ..., + wait: bool = ..., + timeout: timedelta | None = ..., + poll_interval: timedelta = ..., + ) -> ImageInfo: ... + + @overload + def push( + self, + *, + name: str, + zip: bytes, + wait: bool = ..., + timeout: timedelta | None = ..., + poll_interval: timedelta = ..., + ) -> ImageInfo: ... + + @overload + def push( + self, + *, + name: str, + directory: str, + ignore_file_processor: ImageIgnoreFileProcessor | None = ..., + default_ignore_file: ImageIgnoreFileFunc | None = ..., + temp_file: bool | None = ..., + wait: bool = ..., + timeout: timedelta | None = ..., + poll_interval: timedelta = ..., + ) -> ImageInfo: ... + + @overload + def push( + self, + *, + name: str, + files: Mapping[str, str | bytes], + wait: bool = ..., + timeout: timedelta | None = ..., + poll_interval: timedelta = ..., + ) -> ImageInfo: ... + + def push( + self, + *, + name: str, + builder: ImageBuilder | None = None, + registry_image: str | None = None, + docker_config: str | None = None, + zip: bytes | None = None, + directory: str | None = None, + ignore_file_processor: ImageIgnoreFileProcessor | None = None, + default_ignore_file: ImageIgnoreFileFunc | None = None, + files: Mapping[str, str | bytes] | None = None, + temp_file: bool | None = None, + wait: bool = True, + timeout: timedelta | None = _DEFAULT_WAIT_TIMEOUT, + poll_interval: timedelta = _DEFAULT_POLL_INTERVAL, + ) -> ImageInfo: + """Push an image, from a registry or from a build context built into one. + + Takes exactly one source. Waits for the image to be ready to use + unless *wait* is false. To create a sandbox from it, pass + ``:latest`` as the sandbox's ``image``. + + A build context read from local files, from *directory* or a + builder's local files and directories, is zipped into an anonymous + temporary file unless *temp_file* is false; any other is zipped in + memory. It is uploaded in one attempt. If the upload fails, the image + is left as the service has it, so push again or delete it. + + Args: + name: Repository name. Pushing to an existing name adds a new + version to it. + builder: Source: a builder whose Dockerfile and files make up the + build context. + registry_image: Source: a registry image reference to import, + including its registry host. + docker_config: Serialized registry credentials for a private + *registry_image*. + zip: Source: a zip archive of the build context, with a + ``Dockerfile`` at its root. + directory: Source: a local directory holding the build context, + with a ``Dockerfile`` at its root. Files keep their + permissions, a link to a file within the directory stores a + copy of the file, a link to a file outside it fails the push, a + link to a directory stores an empty directory, and a broken + link is left out. Paths are left out as a ``.dockerignore`` at + the directory's root says, read through + *ignore_file_processor*, or else by *default_ignore_file*. The + root ``Dockerfile`` and ``.dockerignore`` are always kept, as + Docker keeps them. + ignore_file_processor: Parses the ``.dockerignore`` at the root of + *directory*. Required if there is one; otherwise the push + fails before anything is sent. This package has no + ``.dockerignore`` parser of its own. + default_ignore_file: Filters *directory* when it has no + ``.dockerignore``. Defaults to + :func:`default_image_ignore_file`. Pass ``lambda _: False`` to + push everything. + files: Source: build context files by relative path, with forward + slashes. Must include ``Dockerfile``. A string is written as + UTF-8, and files get mode 0644. + temp_file: Whether to zip a *directory* or a builder with local + files into a temporary file, rather than in memory. Defaults to + true. + wait: Whether to return only once the image is ready to use. + timeout: How long to wait for the image. ``None`` waits with no + limit. Processing continues regardless. + poll_interval: How long between checks while waiting. + + Raises: + ImageUploadError: Uploading the build context failed. + ImageBuildError: The build failed, or the wait timed out. + """ + # Zipped before anything is sent, so a bad source leaves nothing behind. + archive = self._source_archive( + builder=builder, + registry_image=registry_image, + docker_config=docker_config, + zip=zip, + directory=directory, + ignore_file_processor=ignore_file_processor, + default_ignore_file=default_ignore_file, + files=files, + temp_file=temp_file, + ) + try: + with self._context.api() as api: + response = api.push_image( + params=baseten.client.managementapi.PushImageParams( + **set_fields(team_id=self._context.team_id) + ), + request=_push_request(name, registry_image, docker_config), + ) + if archive is not None: + self._upload(name, _upload_url(name, response), archive) + finally: + if archive is not None: + archive.file.close() + if not wait: + return _pushed_info(response) + return self._wait_built(name, timeout, poll_interval) + + def wait_built( + self, + *, + name: str, + timeout: timedelta | None = _DEFAULT_WAIT_TIMEOUT, + poll_interval: timedelta = _DEFAULT_POLL_INTERVAL, + ) -> ImageInfo: + """Wait for an image to be ready to use. + + Checks the image's current status, so right after pushing a new + version of an image that was already built, this may return before the + new version is processed; push with *wait* left on to avoid that. + + Args: + name: Name of the image. + timeout: How long to wait. ``None`` waits with no limit. + Processing continues regardless. + poll_interval: How long between checks. + + Raises: + ImageBuildError: The build failed, or the wait timed out. + """ + return self._wait_built(name, timeout, poll_interval) + + def logs( + self, + *, + name: str, + start_time: datetime | None = None, + end_time: datetime | None = None, + ) -> builtins.list[ImageLogLine]: + """Get the build log of an image's latest build, including a failed one. + + Oldest line first. Logs may take a short time to appear. Returns at + most the 11,000 most recent lines in the range; narrow the range to see + earlier ones. + + Args: + name: Name of the image. + start_time: Start of the range, inclusive. Defaults to 24 hours + before *end_time*. + end_time: End of the range. Defaults to now. The range must not + exceed 7 days. + """ + # Pinned on the first page, so the default range does not move while + # paging through it. + end = datetime.now(UTC) if end_time is None else end_time + lines: builtins.list[ImageLogLine] = [] + for offset in range(0, _LOGS_MAX_OFFSET + 1, _LOGS_PAGE_SIZE): + with self._context.api() as api: + page = api.get_image_build_logs( + image_name=name, + params=_logs_params(self._context.team_id, start_time, end, offset), + ) + if _add_log_page(lines, page, offset): + break + # Pages come newest first. + lines.reverse() + return lines + + def get_info(self, *, name: str) -> ImageInfo: + """Get an image's current record. + + Args: + name: Name of the image. + """ + with self._context.api() as api: + image = api.get_image( + image_name=name, + params=baseten.client.managementapi.GetImageParams( + **set_fields(team_id=self._context.team_id) + ), + ) + return _image_info(image) + + def list( + self, + *, + name_prefix: str | None = None, + sort: ImageSort | None = None, + page_size: int | None = None, + ) -> Iterator[ImageInfo]: + """List images, fetching further pages as iteration reaches them. + + Args: + name_prefix: Only images whose names start with this, + case-sensitively. + sort: Defaults to ``createdAt:desc``. + page_size: How many images to fetch per underlying request. + """ + + def fetch_page( + cursor: str | None, + ) -> baseten.client.managementapi.ListSandboxImagesResponse: + with self._context.api() as api: + return api.list_images( + params=baseten.client.managementapi.ListImagesParams( + **set_fields( + team_id=self._context.team_id, + cursor=cursor, + limit=page_size, + sort=sort, + q=name_prefix, + ) + ) + ) + + for image in paginate("image list", fetch_page): + yield _image_info(image) + + def delete(self, *, name: str) -> ImageInfo: + """Delete an image and all its versions. Fails while a sandbox uses any version. + + Args: + name: Name of the image. + """ + with self._context.api() as api: + image = api.delete_image( + image_name=name, + params=baseten.client.managementapi.DeleteImageParams( + **set_fields(team_id=self._context.team_id) + ), + ) + return _image_info(image) + + def list_tags( + self, + *, + name: str, + name_prefix: str | None = None, + sort: ImageTagSort | None = None, + page_size: int | None = None, + ) -> Iterator[ImageTagInfo]: + """List an image's versions, fetching further pages as iteration reaches them. + + Args: + name: Name of the image. + name_prefix: Only tags whose names start with this, + case-sensitively. Sorts by name ascending. + sort: Sort order of the tags. + page_size: How many tags to fetch per underlying request. + """ + + def fetch_page( + cursor: str | None, + ) -> baseten.client.managementapi.ListSandboxImageTagsResponse: + with self._context.api() as api: + return api.list_image_tags( + image_name=name, + params=baseten.client.managementapi.ListImageTagsParams( + **set_fields( + team_id=self._context.team_id, + cursor=cursor, + limit=page_size, + sort=sort, + q=name_prefix, + ) + ), + ) + + for tag in paginate("image tag list", fetch_page): + yield _tag_info(tag) + + def delete_tag(self, *, name: str, tag: str) -> ImageInfo: + """Delete one version of an image, returning the image as it is afterwards. + + Fails while a sandbox uses the version. + + Args: + name: Name of the image. + tag: Tag of the version to delete. + """ + with self._context.api() as api: + image = api.delete_image_tag( + image_name=name, + tag_name=tag, + params=baseten.client.managementapi.DeleteImageTagParams( + **set_fields(team_id=self._context.team_id) + ), + ) + return _image_info(image) + + def cleanup(self) -> ImageCleanupResult: + """Remove image versions that no sandbox uses.""" + with self._context.api() as api: + result = api.cleanup_images( + params=baseten.client.managementapi.CleanupImagesParams( + **set_fields(team_id=self._context.team_id) + ) + ) + return ImageCleanupResult(deleted=result.deleted, message=result.message) + + def list_library(self) -> builtins.list[ImageLibraryInfo]: + """List the built-in images available to every team.""" + with self._context.api() as api: + result = api.list_sandbox_library_images() + return [_library_info(image) for image in result.items] + + def _source_archive( + self, + *, + builder: ImageBuilder | None, + registry_image: str | None, + docker_config: str | None, + zip: bytes | None, + directory: str | None, + ignore_file_processor: ImageIgnoreFileProcessor | None, + default_ignore_file: ImageIgnoreFileFunc | None, + files: Mapping[str, str | bytes] | None, + temp_file: bool | None, + ) -> _Archive | None: + """Check that a push has exactly one source, and zip it unless it is a registry image.""" + _check_push_sources( + builder=builder, + registry_image=registry_image, + docker_config=docker_config, + zip=zip, + directory=directory, + ignore_file_processor=ignore_file_processor, + default_ignore_file=default_ignore_file, + files=files, + temp_file=temp_file, + ) + if builder is not None: + entries = builder._zip_entries() + return _zip_archive(entries, _builder_to_file(entries, temp_file)) + if zip is not None: + check_zip_dockerfile(zip) + return _Archive(file=io.BytesIO(zip), size=len(zip)) + if directory is not None: + check_directory_dockerfile(directory) + ignore_file = read_ignore_file(directory) + if ignore_file is None: + ignore = default_ignore_file or default_image_ignore_file + elif ignore_file_processor is None: + raise missing_processor_error(ignore_file) + else: + ignore = ignore_file_processor(ignore_file) + entries = run_walk(walk_directory(directory, ""), ignore) + return _zip_archive(entries, temp_file is not False) + if files is not None: + return _zip_archive(files_zip_entries(files), False) + return None + + def _upload(self, name: str, url: str, archive: _Archive) -> None: + """Put a build context to the signed storage URL a push returned.""" + # The URL is signed for storage, so none of the API's or the HTTP + # client's headers, cookies, or credentials go with it: the request is + # built apart from the client. Storage rejects an upload without a + # length. + response = self._context.http_client.send( + httpx.Request( + "PUT", url, content=archive.file, headers=_upload_headers(archive) + ), + auth=None, + ) + if not response.is_success: + raise ImageUploadError( + image_name=name, status=response.status_code, body=response.text + ) + + # A push resets the image's status, so after one, BUILT or FAILED is that + # push's outcome. + def _wait_built( + self, name: str, timeout: timedelta | None, poll_interval: timedelta + ) -> ImageInfo: + """Poll an image until it is built or failed.""" + # A check already in flight cannot be interrupted, so the wait can + # run over by up to one check. + deadline = ( + None if timeout is None else time.monotonic() + timeout.total_seconds() + ) + last_status = "unknown" + last_error: Exception | None = None + while True: + try: + with self._context.api() as api: + image = api.get_image( + image_name=name, + params=baseten.client.managementapi.GetImageParams( + **set_fields(team_id=self._context.team_id) + ), + ) + except Exception as err: + if not _is_retryable_wait_error(err): + raise + last_error = err + else: + last_error = None + last_status = image.status + info = _built_image(name, image) + if info is not None: + return info + delay = poll_interval.total_seconds() + if deadline is not None: + remaining = deadline - time.monotonic() + if remaining <= delay: + time.sleep(max(remaining, 0)) + raise ImageBuildError( + image_name=name, status=last_status, timed_out=True + ) from last_error + time.sleep(delay) + + +class AsyncImageClient: + """Async form of :class:`ImageClient`.""" + + def __init__(self, context: AsyncImageClientContext) -> None: + """Internal. Obtained from :attr:`AsyncSandboxClient.images`. + + :meta private: + """ + self._context = context + + @overload + async def push( + self, + *, + name: str, + builder: ImageBuilder, + temp_file: bool | None = ..., + wait: bool = ..., + timeout: timedelta | None = ..., + poll_interval: timedelta = ..., + ) -> ImageInfo: ... + + @overload + async def push( + self, + *, + name: str, + registry_image: str, + docker_config: str | None = ..., + wait: bool = ..., + timeout: timedelta | None = ..., + poll_interval: timedelta = ..., + ) -> ImageInfo: ... + + @overload + async def push( + self, + *, + name: str, + zip: bytes, + wait: bool = ..., + timeout: timedelta | None = ..., + poll_interval: timedelta = ..., + ) -> ImageInfo: ... + + @overload + async def push( + self, + *, + name: str, + directory: str, + ignore_file_processor: AsyncImageIgnoreFileProcessor | None = ..., + default_ignore_file: AsyncImageIgnoreFileFunc | None = ..., + temp_file: bool | None = ..., + wait: bool = ..., + timeout: timedelta | None = ..., + poll_interval: timedelta = ..., + ) -> ImageInfo: ... + + @overload + async def push( + self, + *, + name: str, + files: Mapping[str, str | bytes], + wait: bool = ..., + timeout: timedelta | None = ..., + poll_interval: timedelta = ..., + ) -> ImageInfo: ... + + async def push( + self, + *, + name: str, + builder: ImageBuilder | None = None, + registry_image: str | None = None, + docker_config: str | None = None, + zip: bytes | None = None, + directory: str | None = None, + ignore_file_processor: AsyncImageIgnoreFileProcessor | None = None, + default_ignore_file: AsyncImageIgnoreFileFunc | None = None, + files: Mapping[str, str | bytes] | None = None, + temp_file: bool | None = None, + wait: bool = True, + timeout: timedelta | None = _DEFAULT_WAIT_TIMEOUT, + poll_interval: timedelta = _DEFAULT_POLL_INTERVAL, + ) -> ImageInfo: + """Push an image. As :meth:`ImageClient.push`. + + Local files are read and zipped on a worker thread. The ignore + function and processor may also be coroutine functions. + """ + archive = await self._source_archive( + builder=builder, + registry_image=registry_image, + docker_config=docker_config, + zip=zip, + directory=directory, + ignore_file_processor=ignore_file_processor, + default_ignore_file=default_ignore_file, + files=files, + temp_file=temp_file, + ) + try: + with self._context.api() as api: + response = await api.push_image( + params=baseten.client.managementapi.PushImageParams( + **set_fields(team_id=self._context.team_id) + ), + request=_push_request(name, registry_image, docker_config), + ) + if archive is not None: + await self._upload(name, _upload_url(name, response), archive) + finally: + if archive is not None: + archive.file.close() + if not wait: + return _pushed_info(response) + return await self._wait_built(name, timeout, poll_interval) + + async def wait_built( + self, + *, + name: str, + timeout: timedelta | None = _DEFAULT_WAIT_TIMEOUT, + poll_interval: timedelta = _DEFAULT_POLL_INTERVAL, + ) -> ImageInfo: + """Wait for an image to be ready to use. As :meth:`ImageClient.wait_built`.""" + return await self._wait_built(name, timeout, poll_interval) + + async def logs( + self, + *, + name: str, + start_time: datetime | None = None, + end_time: datetime | None = None, + ) -> builtins.list[ImageLogLine]: + """Get the build log of an image's latest build. As :meth:`ImageClient.logs`.""" + end = datetime.now(UTC) if end_time is None else end_time + lines: builtins.list[ImageLogLine] = [] + for offset in range(0, _LOGS_MAX_OFFSET + 1, _LOGS_PAGE_SIZE): + with self._context.api() as api: + page = await api.get_image_build_logs( + image_name=name, + params=_logs_params(self._context.team_id, start_time, end, offset), + ) + if _add_log_page(lines, page, offset): + break + lines.reverse() + return lines + + async def get_info(self, *, name: str) -> ImageInfo: + """Get an image's current record. As :meth:`ImageClient.get_info`.""" + with self._context.api() as api: + image = await api.get_image( + image_name=name, + params=baseten.client.managementapi.GetImageParams( + **set_fields(team_id=self._context.team_id) + ), + ) + return _image_info(image) + + async def list( + self, + *, + name_prefix: str | None = None, + sort: ImageSort | None = None, + page_size: int | None = None, + ) -> AsyncIterator[ImageInfo]: + """List images. As :meth:`ImageClient.list`.""" + + async def fetch_page( + cursor: str | None, + ) -> baseten.client.managementapi.ListSandboxImagesResponse: + with self._context.api() as api: + return await api.list_images( + params=baseten.client.managementapi.ListImagesParams( + **set_fields( + team_id=self._context.team_id, + cursor=cursor, + limit=page_size, + sort=sort, + q=name_prefix, + ) + ) + ) + + async for image in paginate_async("image list", fetch_page): + yield _image_info(image) + + async def delete(self, *, name: str) -> ImageInfo: + """Delete an image. As :meth:`ImageClient.delete`.""" + with self._context.api() as api: + image = await api.delete_image( + image_name=name, + params=baseten.client.managementapi.DeleteImageParams( + **set_fields(team_id=self._context.team_id) + ), + ) + return _image_info(image) + + async def list_tags( + self, + *, + name: str, + name_prefix: str | None = None, + sort: ImageTagSort | None = None, + page_size: int | None = None, + ) -> AsyncIterator[ImageTagInfo]: + """List an image's versions. As :meth:`ImageClient.list_tags`.""" + + async def fetch_page( + cursor: str | None, + ) -> baseten.client.managementapi.ListSandboxImageTagsResponse: + with self._context.api() as api: + return await api.list_image_tags( + image_name=name, + params=baseten.client.managementapi.ListImageTagsParams( + **set_fields( + team_id=self._context.team_id, + cursor=cursor, + limit=page_size, + sort=sort, + q=name_prefix, + ) + ), + ) + + async for tag in paginate_async("image tag list", fetch_page): + yield _tag_info(tag) + + async def delete_tag(self, *, name: str, tag: str) -> ImageInfo: + """Delete one version of an image. As :meth:`ImageClient.delete_tag`.""" + with self._context.api() as api: + image = await api.delete_image_tag( + image_name=name, + tag_name=tag, + params=baseten.client.managementapi.DeleteImageTagParams( + **set_fields(team_id=self._context.team_id) + ), + ) + return _image_info(image) + + async def cleanup(self) -> ImageCleanupResult: + """Remove image versions that no sandbox uses. As :meth:`ImageClient.cleanup`.""" + with self._context.api() as api: + result = await api.cleanup_images( + params=baseten.client.managementapi.CleanupImagesParams( + **set_fields(team_id=self._context.team_id) + ) + ) + return ImageCleanupResult(deleted=result.deleted, message=result.message) + + async def list_library(self) -> builtins.list[ImageLibraryInfo]: + """List the built-in images. As :meth:`ImageClient.list_library`.""" + with self._context.api() as api: + result = await api.list_sandbox_library_images() + return [_library_info(image) for image in result.items] + + async def _source_archive( + self, + *, + builder: ImageBuilder | None, + registry_image: str | None, + docker_config: str | None, + zip: bytes | None, + directory: str | None, + ignore_file_processor: AsyncImageIgnoreFileProcessor | None, + default_ignore_file: AsyncImageIgnoreFileFunc | None, + files: Mapping[str, str | bytes] | None, + temp_file: bool | None, + ) -> _Archive | None: + """Check that a push has exactly one source, and zip it unless it is a registry image.""" + _check_push_sources( + builder=builder, + registry_image=registry_image, + docker_config=docker_config, + zip=zip, + directory=directory, + ignore_file_processor=ignore_file_processor, + default_ignore_file=default_ignore_file, + files=files, + temp_file=temp_file, + ) + if builder is not None: + entries = await asyncio.to_thread(builder._zip_entries) + return await asyncio.to_thread( + _zip_archive, entries, _builder_to_file(entries, temp_file) + ) + if zip is not None: + # Reading a large zip's entries would block the loop. + await asyncio.to_thread(check_zip_dockerfile, zip) + return _Archive(file=io.BytesIO(zip), size=len(zip)) + if directory is not None: + await asyncio.to_thread(check_directory_dockerfile, directory) + ignore_file = await asyncio.to_thread(read_ignore_file, directory) + ignore: AsyncImageIgnoreFileFunc + if ignore_file is None: + ignore = default_ignore_file or default_image_ignore_file + elif ignore_file_processor is None: + raise missing_processor_error(ignore_file) + else: + processed = ignore_file_processor(ignore_file) + ignore = ( + await processed if inspect.isawaitable(processed) else processed + ) + entries = await run_walk_async(walk_directory(directory, ""), ignore) + return await asyncio.to_thread( + _zip_archive, entries, temp_file is not False + ) + if files is not None: + # Encoding and compressing a large build context would block the + # loop. + return await asyncio.to_thread( + lambda source: _zip_archive(files_zip_entries(source), False), files + ) + return None + + async def _upload(self, name: str, url: str, archive: _Archive) -> None: + """Put a build context to the signed storage URL a push returned.""" + + # Read on a worker thread, since the archive may be a file on disk. + async def chunks() -> AsyncIterator[bytes]: + while chunk := await asyncio.to_thread( + archive.file.read, _UPLOAD_CHUNK_SIZE + ): + yield chunk + + # Built apart from the HTTP client, as in ImageClient._upload. + response = await self._context.http_client.send( + httpx.Request( + "PUT", url, content=chunks(), headers=_upload_headers(archive) + ), + auth=None, + ) + if not response.is_success: + raise ImageUploadError( + image_name=name, status=response.status_code, body=response.text + ) + + async def _wait_built( + self, name: str, timeout: timedelta | None, poll_interval: timedelta + ) -> ImageInfo: + """Poll an image until it is built or failed.""" + last_status = "unknown" + last_error: Exception | None = None + scope = asyncio.timeout(None if timeout is None else timeout.total_seconds()) + try: + async with scope: + while True: + try: + with self._context.api() as api: + image = await api.get_image( + image_name=name, + params=baseten.client.managementapi.GetImageParams( + **set_fields(team_id=self._context.team_id) + ), + ) + except Exception as err: + if not _is_retryable_wait_error(err): + raise + last_error = err + else: + last_error = None + last_status = image.status + info = _built_image(name, image) + if info is not None: + return info + await asyncio.sleep(poll_interval.total_seconds()) + except TimeoutError: + # Only this wait's own timeout, not a caller's or one from a + # check, becomes the build error. + if not scope.expired(): + raise + raise ImageBuildError( + image_name=name, status=last_status, timed_out=True + ) from last_error + + +def _check_push_sources( + *, + builder: object, + registry_image: object, + docker_config: object, + zip: object, + directory: object, + ignore_file_processor: object, + default_ignore_file: object, + files: object, + temp_file: object, +) -> None: + """Raise unless a push has exactly one source and only options that apply to it.""" + sources = [ + source + for source, value in ( + ("builder", builder), + ("registry_image", registry_image), + ("zip", zip), + ("directory", directory), + ("files", files), + ) + if value is not None + ] + if len(sources) != 1: + raise TypeError( + "an image push needs exactly one of builder, registry_image, zip, " + f"directory, or files, got {', '.join(sources) or 'none'}" + ) + if docker_config is not None and registry_image is None: + raise TypeError("docker_config applies only to registry_image") + if ( + ignore_file_processor is not None or default_ignore_file is not None + ) and directory is None: + raise TypeError( + "ignore_file_processor and default_ignore_file apply only to directory" + ) + if temp_file is not None and directory is None and builder is None: + raise TypeError("temp_file applies only to directory and builder") + + +def _builder_to_file(entries: list[ZipEntry], temp_file: bool | None) -> bool: + """Report whether a builder's build context is zipped into a temporary file.""" + # Only local files are read while zipping, so only they can make a build + # context too large for memory. + return temp_file is not False and any( + entry.source_path is not None for entry in entries + ) + + +def _zip_archive(entries: list[ZipEntry], to_file: bool) -> _Archive: + """Zip entries into an anonymous temporary file when ``to_file`` is set, or else in memory.""" + with contextlib.ExitStack() as on_failure: + file: IO[bytes] + if to_file: + # A temporary file with no name, which the OS removes once it is + # closed, even if the process dies first. + file = on_failure.enter_context(tempfile.TemporaryFile()) + else: + file = io.BytesIO() + write_zip(entries, file) + size = file.tell() + file.seek(0) + # Kept open for the upload, which closes it. + on_failure.pop_all() + return _Archive(file=file, size=size) + + +def _push_request( + name: str, registry_image: str | None, docker_config: str | None +) -> baseten.client.managementapi.PushSandboxImageRequest: + """Build the request that pushes an image.""" + return baseten.client.managementapi.PushSandboxImageRequest( + name=name, + **set_fields(image=registry_image, docker_config=docker_config), + ) + + +def _upload_url( + name: str, response: baseten.client.managementapi.PushSandboxImageResponse +) -> str: + """Give the storage URL a push returned for its build context.""" + if response.upload_url is None: + raise SandboxError(f"pushing image {name} returned no upload URL") + return str(response.upload_url) + + +def _upload_headers(archive: _Archive) -> dict[str, str]: + """Give the headers of a build context upload.""" + return {"Content-Type": "application/zip", "Content-Length": str(archive.size)} + + +def _pushed_info( + response: baseten.client.managementapi.PushSandboxImageResponse, +) -> ImageInfo: + """Give the record a push returns without waiting.""" + return ImageInfo( + name=response.name, + status=response.status, + created_at=None, + updated_at=None, + last_deployed_at=None, + size_bytes=None, + tag_count=None, + ) + + +def _built_image( + name: str, image: baseten.client.managementapi.SandboxImage +) -> ImageInfo | None: + """Give a built image's record, raise for a failed one, or ``None`` while it processes.""" + if image.status == "FAILED": + raise ImageBuildError(image_name=name, status=image.status, timed_out=False) + if image.status == "BUILT": + return _image_info(image) + # UPLOADING, BUILDING, or a status added later, all still processing. + return None + + +def _is_retryable_wait_error(err: Exception) -> bool: + """Report whether a wait rides out an error. + + A dropped connection, a gateway error, or a not-found while a push is + still becoming visible. + """ + if isinstance(err, SandboxAPIError): + return err.status in _WAIT_RETRY_STATUSES + return is_transient_reset_error(err) + + +def _logs_params( + team_id: str | None, start_time: datetime | None, end_time: datetime, offset: int +) -> baseten.client.managementapi.GetImageBuildLogsParams: + """Build the query of one page of build logs.""" + return baseten.client.managementapi.GetImageBuildLogsParams( + end_time=end_time, + limit=_LOGS_PAGE_SIZE, + offset=offset, + **set_fields(team_id=team_id, start_time=start_time), + ) + + +def _add_log_page( + lines: list[ImageLogLine], + page: baseten.client.managementapi.SandboxImageBuildLogsResponse, + offset: int, +) -> bool: + """Add a page of build logs to ``lines``, reporting whether it was the last.""" + lines.extend( + ImageLogLine(timestamp=log.timestamp, severity=log.severity, text=log.message) + for log in page.logs + ) + return ( + len(page.logs) < _LOGS_PAGE_SIZE or offset + _LOGS_PAGE_SIZE >= page.total_count + ) + + +def _image_info( + image: baseten.client.managementapi.SandboxImage + | baseten.client.managementapi.SandboxImageSummary, +) -> ImageInfo: + """Convert an image record, from a listing or not.""" + return ImageInfo( + name=image.name, + status=image.status, + created_at=image.created_at, + updated_at=image.updated_at, + last_deployed_at=image.last_deployed_at, + size_bytes=image.size, + tag_count=image.tag_count, + ) + + +def _tag_info(tag: baseten.client.managementapi.SandboxImageTag) -> ImageTagInfo: + """Convert an image version's record.""" + return ImageTagInfo( + name=tag.name, + created_at=tag.created_at, + updated_at=tag.updated_at, + size_bytes=tag.size, + ) + + +# hidden and coming_soon are left out, since listings always have them false. +def _library_info( + image: baseten.client.managementapi.SandboxLibraryImage, +) -> ImageLibraryInfo: + """Convert a built-in image's record.""" + options = image.creation_options + return ImageLibraryInfo( + name=image.name, + display_name=image.display_name, + description=image.description, + long_description=image.long_description, + image=image.image, + memory=image.memory, + ports=ports_from_api(image.ports), + categories=list(image.categories or []), + tags=list(image.tags or []), + url=image.url, + icon=image.icon, + icon_light=image.icon_light, + icon_dark=image.icon_dark, + enterprise=bool(image.enterprise), + creation_extra_args=None + if options is None or options.extra_args is None + else dict(options.extra_args), + creation_volumes=[] + if options is None + else [ + ImageLibraryVolume( + 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 options.volumes or [] + ], + ) diff --git a/baseten/sandbox/_images_builder.py b/baseten/sandbox/_images_builder.py new file mode 100644 index 0000000..677d9c4 --- /dev/null +++ b/baseten/sandbox/_images_builder.py @@ -0,0 +1,491 @@ +from __future__ import annotations + +import json +import os +import re +import stat +from collections.abc import Mapping +from dataclasses import dataclass, replace +from typing import Literal, Self + +from baseten.sandbox._zip import ZipEntry, run_walk, walk_directory + +_DEFAULT_SANDBOX_API_IMAGE = "ghcr.io/blaxel-ai/sandbox:latest" +_SANDBOX_API_PATH = "/usr/local/bin/sandbox-api" + +# Names the build itself reads at the root of its context, so a context entry +# must never take them. +_RESERVED_CONTEXT_NAMES = frozenset({"Dockerfile", ".dockerignore"}) + +# Characters a package manager argument may have and still be passed to the +# shell unquoted. +_SHELL_SAFE = re.compile(r"[A-Za-z0-9._\-+=:/@~^*]+") + +_ENTRYPOINT = re.compile(r"\s*ENTRYPOINT\b", re.IGNORECASE) + + +@dataclass(frozen=True, kw_only=True) +class _ContextEntry: + """A file or directory the builder adds to the build context.""" + + kind: Literal["content", "file", "directory"] + name: str + content: bytes = b"" + mode: int = 0o644 + source_path: str = "" + + +@dataclass(frozen=True, kw_only=True) +class _BuilderState: + base_image: str + sandbox_api_image: str + instructions: tuple[str, ...] = () + context: tuple[_ContextEntry, ...] = () + + +class ImageBuilder: + """An image described as a Dockerfile and the files it copies in. + + Push it with :meth:`ImageClient.push`. Start one with + :meth:`ImageBuilder.from_registry`. + + A builder never changes: each method returns a new builder with the + instruction appended, so one builder can safely be the base of several. + Invalid input raises :class:`ValueError` at the call that takes it. + + The Dockerfile it pushes ends by copying in the sandbox API binary and + making it the entrypoint, unless the builder already did either, since a + sandbox needs it running to be reached. :meth:`dockerfile` shows the + Dockerfile exactly as pushed. + """ + + def __init__(self, state: _BuilderState) -> None: + """Internal. Use :meth:`from_registry` instead. + + :meta private: + """ + self._state = state + + @classmethod + def from_registry( + cls, base_image: str, *, sandbox_api_image: str | None = None + ) -> ImageBuilder: + """Start a builder from a registry image, as ``FROM ``. + + Args: + base_image: Registry image to start from. + sandbox_api_image: Image to copy the sandbox API binary from, when + the Dockerfile does not bring it in itself. Defaults to + ``ghcr.io/blaxel-ai/sandbox:latest``. + """ + return cls( + _BuilderState( + base_image=_single_line("base_image", base_image), + sandbox_api_image=_single_line( + "sandbox_api_image", + _DEFAULT_SANDBOX_API_IMAGE + if sandbox_api_image is None + else sandbox_api_image, + ), + ) + ) + + @property + def base_image(self) -> str: + """The image this builder starts from.""" + return self._state.base_image + + def dockerfile(self) -> str: + """Give the Dockerfile, exactly as it is pushed.""" + state = self._state + lines = [f"FROM {state.base_image}", *state.instructions] + text = "\n".join(lines) + if "sandbox-api" not in text and "blaxel-ai/sandbox" not in text: + lines.append( + f"COPY --from={state.sandbox_api_image} /sandbox-api {_SANDBOX_API_PATH}" + ) + # Read from the lines rather than tracked by entrypoint(), so one + # written with dockerfile_lines() counts too. + has_entrypoint = any( + _ENTRYPOINT.match(line) + for instruction in state.instructions + for line in instruction.split("\n") + ) + if not has_entrypoint: + lines.append(f"ENTRYPOINT {_json([_SANDBOX_API_PATH])}") + return "\n".join(lines) + "\n" + + def workdir(self, path: str) -> ImageBuilder: + """Append ``WORKDIR ``.""" + return self._append(f"WORKDIR {_single_line('path', path)}") + + def run_commands(self, *commands: str) -> ImageBuilder: + """Append ``RUN `` for each command, in shell form.""" + return self._append( + *(f"RUN {_single_line('command', command)}" for command in commands) + ) + + def env(self, variables: Mapping[str, str]) -> ImageBuilder: + """Append ``ENV =""`` for each variable. + + ``"`` and ``\\`` in the value are escaped. ``$`` is left as is, so the + build expands variables in it, as in ``{"PATH": "/opt/bin:$PATH"}``. + """ + return self._append( + *( + f"ENV {_key_name('env name', name)}={_quoted('env value', value)}" + for name, value in variables.items() + ) + ) + + def copy(self, source: str, destination: str) -> ImageBuilder: + """Append ``COPY ["", ""]``, copying from the build context.""" + _single_line("source", source) + _single_line("destination", destination) + return self._append(f"COPY {_json([source, destination])}") + + def expose(self, *ports: int) -> ImageBuilder: + """Append ``EXPOSE `` for each port.""" + for port in ports: + if not 1 <= port <= 65535: + raise ValueError(f"port {port} is not between 1 and 65535") + return self._append(*(f"EXPOSE {port}" for port in ports)) + + def entrypoint(self, *args: str) -> ImageBuilder: + """Append ``ENTRYPOINT ["", ...]``, in exec form. + + The image then runs this instead of the sandbox API binary, so it must + start that itself for the sandbox to be reachable. + """ + if not args: + return self + for arg in args: + _single_line("entrypoint argument", arg) + return self._append(f"ENTRYPOINT {_json(list(args))}") + + def user(self, user: str) -> ImageBuilder: + """Append ``USER ``.""" + return self._append(f"USER {_single_line('user', user)}") + + def label(self, labels: Mapping[str, str]) -> ImageBuilder: + """Append ``LABEL =""`` for each label, escaped as in :meth:`env`.""" + return self._append( + *( + f"LABEL {_key_name('label key', key)}={_quoted('label value', value)}" + for key, value in labels.items() + ) + ) + + def arg(self, name: str, default: str | None = None) -> ImageBuilder: + """Append ``ARG ``, or ``ARG =""`` escaped as in :meth:`env`.""" + declared = _key_name("arg name", name) + if default is None: + return self._append(f"ARG {declared}") + return self._append(f"ARG {declared}={_quoted('arg default', default)}") + + def dockerfile_lines(self, *lines: str) -> ImageBuilder: + """Append each line verbatim. + + For anything the other methods do not cover, such as comments, ``CMD``, + or a heredoc. A line may contain newlines, and nothing is checked or + escaped. + """ + return self._append(*lines) + + def pip_install(self, *args: str) -> ImageBuilder: + """Append ``RUN pip install ``.""" + return self._run("pip install", args) + + def apt_install(self, *args: str) -> ImageBuilder: + """Append an ``apt-get install`` of the packages as one step. + + The step is ``RUN apt-get update && apt-get install -y + --no-install-recommends && rm -rf /var/lib/apt/lists/*``, so the + package lists add nothing to the image. + """ + if not args: + return self + return self._append( + "RUN apt-get update && apt-get install -y --no-install-recommends " + f"{_shell_args(args)} && rm -rf /var/lib/apt/lists/*" + ) + + def apk_add(self, *args: str) -> ImageBuilder: + """Append ``RUN apk add --no-cache ``.""" + return self._run("apk add --no-cache", args) + + def npm_install(self, *args: str) -> ImageBuilder: + """Append ``RUN npm install ``. + + With no arguments, appends ``RUN npm install``, to install from the + working directory's ``package.json``. + """ + if not args: + return self._append("RUN npm install") + return self._run("npm install", args) + + def gem_install(self, *args: str) -> ImageBuilder: + """Append ``RUN gem install --no-document ``.""" + return self._run("gem install --no-document", args) + + def cargo_install(self, *args: str) -> ImageBuilder: + """Append ``RUN cargo install ``.""" + return self._run("cargo install", args) + + def go_install(self, *packages: str) -> ImageBuilder: + """Append ``RUN go install && go install ...``. + + One ``go install`` per package, since one call can only install + packages from a single module. + """ + return self._run_each("go install", packages) + + def composer_install(self, *args: str) -> ImageBuilder: + """Append ``RUN composer require ``.""" + return self._run("composer require", args) + + def uv_install(self, *args: str) -> ImageBuilder: + """Append ``RUN uv pip install --system ``.""" + return self._run("uv pip install --system", args) + + def pipx_install(self, *packages: str) -> ImageBuilder: + """Append ``RUN pipx install && pipx install ...``, one per package.""" + return self._run_each("pipx install", packages) + + def add_local_file( + self, source_path: str, destination: str, *, context_name: str | None = None + ) -> ImageBuilder: + """Add a local file to the build context and append ``COPY ["", ""]``. + + The name is the file's name in the context. The file is read when + pushed, keeping its permissions, and a symbolic link is read as the + file it points to. + + Args: + source_path: Local file to add. + destination: Path to copy it to in the image. + context_name: Name in the build context. Defaults to the source's + basename, with a suffix such as `` (1)`` if another entry + already has that name. + """ + return self._add_local("file", source_path, destination, context_name) + + def add_local_dir( + self, source_path: str, destination: str, *, context_name: str | None = None + ) -> ImageBuilder: + """Add a local directory to the build context and append ``COPY ["", ""]``. + + The name is the directory's name in the context. As with any ``COPY`` + of a directory, its contents are copied into the destination, not the + directory itself. It is read when pushed, with files keeping their + permissions, a link to a file within the directory stored as a copy of + the file, a link to a file outside it failing the push, and a link to + a directory stored as an empty directory. Everything in it is added: no + ``.dockerignore`` or default ignore rules apply. + + Args: + source_path: Local directory to add. + destination: Path to copy its contents to in the image. + context_name: Name in the build context. Defaults to the source's + basename, with a suffix such as `` (1)`` if another entry + already has that name. + """ + return self._add_local("directory", source_path, destination, context_name) + + def add_file( + self, + destination: str, + content: str | bytes, + *, + context_name: str | None = None, + mode: int = 0o644, + ) -> ImageBuilder: + """Add a file with the given content to the build context and append ``COPY ["", ""]``. + + The name is the file's name in the context. + + Args: + destination: Path of the file in the image. + content: Content of the file. A string is written as UTF-8. + context_name: Name in the build context. Defaults to the + destination's basename, with a suffix such as `` (1)`` if + another entry already has that name. + mode: Unix permission bits of the file in the image. + """ + _single_line("destination", destination) + default_name = destination.split("/")[-1] + if context_name is None and default_name == "": + raise ValueError( + f"destination {destination} has no file name; set context_name" + ) + name = self._context_name(context_name, default_name) + return self._append( + f"COPY {_json([name, destination])}", + context_entry=_ContextEntry( + kind="content", + name=name, + content=content.encode() if isinstance(content, str) else content, + mode=mode, + ), + ) + + def _zip_entries(self) -> list[ZipEntry]: + """Give the build context to zip, with local entries checked now and read while zipping.""" + entries = [ + ZipEntry(path="Dockerfile", data=self.dockerfile().encode(), mode=0o644) + ] + for entry in self._state.context: + if entry.kind == "content": + entries.append( + ZipEntry(path=entry.name, data=entry.content, mode=entry.mode) + ) + continue + try: + info = os.stat(entry.source_path) + except FileNotFoundError: + info = None + if entry.kind == "file": + if info is None or not stat.S_ISREG(info.st_mode): + raise ValueError(f"local file {entry.source_path} is not a file") + entries.append( + ZipEntry( + path=entry.name, + mode=stat.S_IMODE(info.st_mode) & 0o777, + source_path=entry.source_path, + ) + ) + continue + if info is None or not stat.S_ISDIR(info.st_mode): + raise ValueError( + f"local directory {entry.source_path} is not a directory" + ) + entries.append( + ZipEntry( + path=entry.name, + mode=stat.S_IMODE(info.st_mode) & 0o777, + is_dir=True, + ) + ) + # No ignore rules: adding a directory is an explicit choice of its + # files, as Docker's COPY of a directory takes them all. + entries.extend( + run_walk(walk_directory(entry.source_path, entry.name + "/"), None) + ) + return entries + + def _append( + self, *instructions: str, context_entry: _ContextEntry | None = None + ) -> Self: + """Give a new builder with the instructions and context entry appended.""" + if not instructions: + return self + state = self._state + return type(self)( + replace( + state, + instructions=state.instructions + instructions, + context=state.context + if context_entry is None + else (*state.context, context_entry), + ) + ) + + def _run(self, command: str, args: tuple[str, ...]) -> ImageBuilder: + """Append ``RUN ``, or nothing without arguments.""" + if not args: + return self + return self._append(f"RUN {command} {_shell_args(args)}") + + def _run_each(self, command: str, packages: tuple[str, ...]) -> ImageBuilder: + """Append one ``RUN`` running the command once per package.""" + if not packages: + return self + runs = [f"{command} {_shell_args((package,))}" for package in packages] + return self._append("RUN " + " && ".join(runs)) + + def _add_local( + self, + kind: Literal["file", "directory"], + source_path: str, + destination: str, + context_name: str | None, + ) -> ImageBuilder: + """Add a local file or directory, read when pushed.""" + _single_line("destination", destination) + # Resolved now, so a later change of working directory does not change + # what is pushed. + resolved = os.path.abspath(source_path) + name = self._context_name(context_name, os.path.basename(resolved)) + return self._append( + f"COPY {_json([name, destination])}", + context_entry=_ContextEntry(kind=kind, name=name, source_path=resolved), + ) + + def _context_name(self, requested: str | None, default_name: str) -> str: + """Give a context entry's name. + + An explicit name is used exactly or raises; a default one is suffixed + the way browsers name a repeated download, before any extension. + """ + taken = {entry.name for entry in self._state.context} | _RESERVED_CONTEXT_NAMES + if requested is not None: + _check_context_name(requested) + if requested in taken: + raise ValueError( + f"context name {requested} is reserved or already used" + ) + return requested + _check_context_name(default_name) + if default_name not in taken: + return default_name + dot = default_name.rfind(".") + stem, extension = ( + (default_name[:dot], default_name[dot:]) if dot > 0 else (default_name, "") + ) + number = 1 + while f"{stem} ({number}){extension}" in taken: + number += 1 + return f"{stem} ({number}){extension}" + + +def _single_line(what: str, value: str) -> str: + """Raise if a value would span lines in the Dockerfile, otherwise return it.""" + if "\n" in value or "\r" in value: + raise ValueError(f"{what} must not contain a newline") + return value + + +def _key_name(what: str, value: str) -> str: + """Check an env, label, or arg name, which is written unquoted.""" + if value == "" or re.search(r'[\s="]', value): + raise ValueError(f'{what} {value!r} must be non-empty, without spaces, = or "') + return value + + +def _quoted(what: str, value: str) -> str: + """Double-quote a value, escaping ``"`` and ``\\``.""" + return '"' + re.sub(r'(["\\])', r"\\\1", _single_line(what, value)) + '"' + + +def _check_context_name(name: str) -> None: + """Raise unless a context name is one path segment.""" + if name in ("", ".", "..") or re.search(r"[/\\\r\n]", name): + raise ValueError( + f"context name {name!r} must be one path segment, without slashes or newlines" + ) + + +def _shell_args(args: tuple[str, ...]) -> str: + """Join arguments for the shell, single-quoting any that need it.""" + quoted = [] + for arg in args: + _single_line("argument", arg) + quoted.append( + arg if _SHELL_SAFE.fullmatch(arg) else "'" + arg.replace("'", "'\\''") + "'" + ) + return " ".join(quoted) + + +def _json(values: list[str]) -> str: + """Write a JSON array as a Dockerfile's exec form takes it, compact like JavaScript's.""" + return json.dumps(values, ensure_ascii=False, separators=(",", ":")) diff --git a/baseten/sandbox/_images_ignore.py b/baseten/sandbox/_images_ignore.py new file mode 100644 index 0000000..fcc1da6 --- /dev/null +++ b/baseten/sandbox/_images_ignore.py @@ -0,0 +1,129 @@ +from __future__ import annotations + +import os +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from typing import TypeAlias + +DOCKERIGNORE_FILE_NAME = ".dockerignore" + + +@dataclass(kw_only=True) +class ImageIgnoreFileOptions: + """A candidate path in a directory being pushed, for an ignore function.""" + + rel_path: str + """Path relative to the directory's root, with forward slashes on all platforms.""" + + is_directory: bool + """Whether the path is a directory, not a link to one.""" + + +ImageIgnoreFileFunc: TypeAlias = Callable[[ImageIgnoreFileOptions], bool] +"""Reports whether a path should be left out of a pushed build context. + +When it returns true for a directory, the directory's entire subtree is left +out. Raising fails the push before anything is sent. +""" + +AsyncImageIgnoreFileFunc: TypeAlias = Callable[ + [ImageIgnoreFileOptions], bool | Awaitable[bool] +] +"""Async form of :data:`ImageIgnoreFileFunc`, which may also be a coroutine function.""" + + +@dataclass(kw_only=True) +class ImageIgnoreFileProcessorOptions: + """A ``.dockerignore`` found at the root of a directory being pushed.""" + + path: str + """Absolute path of the ``.dockerignore``.""" + + contents: str + """Contents of the ``.dockerignore``, decoded as UTF-8.""" + + +ImageIgnoreFileProcessor: TypeAlias = Callable[ + [ImageIgnoreFileProcessorOptions], ImageIgnoreFileFunc +] +"""Parses a ``.dockerignore`` into an :data:`ImageIgnoreFileFunc`. + +The function should match Docker's ``.dockerignore`` semantics. +""" + +AsyncImageIgnoreFileProcessor: TypeAlias = Callable[ + [ImageIgnoreFileProcessorOptions], + AsyncImageIgnoreFileFunc | Awaitable[AsyncImageIgnoreFileFunc], +] +"""Async form of :data:`ImageIgnoreFileProcessor`, which may also be a coroutine function.""" + +# Left out at any depth by default_image_ignore_file. +_DEFAULT_IGNORED_NAMES = frozenset( + { + ".blaxel", + ".env.build", + ".docker", + ".git", + "dist", + ".venv", + "venv", + "node_modules", + ".env", + ".next", + "__pycache__", + } +) + + +def default_image_ignore_file(options: ImageIgnoreFileOptions) -> bool: + """Report whether a path is left out of a pushed build context by default. + + Applied when the directory has no ``.dockerignore``. Its result is the same + as this ``.dockerignore``:: + + **/.blaxel + **/.env.build + **/.docker + **/.git + **/dist + **/.venv + **/venv + **/node_modules + **/.env + .env* + **/.next + **/__pycache__ + + So each of those names is left out at any depth, as a file or as a + directory with everything in it, and a name starting with ``.env`` is left + out at the root. To extend the defaults, copy them into a + ``.dockerignore``. + + Args: + options: The candidate path. + """ + components = options.rel_path.split("/") + # Any component, so a path under an ignored directory is ignored even when + # its directory was not pruned first. + if any(component in _DEFAULT_IGNORED_NAMES for component in components): + return True + # .env* has no **/, so it matches at the root only, as Docker reads it. + return components[0].startswith(".env") + + +def read_ignore_file(directory: str) -> ImageIgnoreFileProcessorOptions | None: + """Read the ``.dockerignore`` at a directory's root, or ``None`` without one.""" + path = os.path.abspath(os.path.join(directory, DOCKERIGNORE_FILE_NAME)) + try: + with open(path, encoding="utf-8", errors="replace") as file: + contents = file.read() + except FileNotFoundError: + return None + return ImageIgnoreFileProcessorOptions(path=path, contents=contents) + + +def missing_processor_error(ignore_file: ImageIgnoreFileProcessorOptions) -> ValueError: + """Give the error for a ``.dockerignore`` with no processor to read it.""" + # Failing, rather than pushing everything, keeps files the caller meant to + # leave out, such as secrets, from being uploaded. + return ValueError(f"{ignore_file.path} exists but ignore_file_processor is not set") diff --git a/baseten/sandbox/_info.py b/baseten/sandbox/_info.py new file mode 100644 index 0000000..ac28719 --- /dev/null +++ b/baseten/sandbox/_info.py @@ -0,0 +1,490 @@ +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from datetime import datetime, timedelta +from typing import Any, Literal, TypeAlias + +import baseten.client.managementapi +from baseten.sandbox._common import format_duration, parse_duration + +SandboxStatus: TypeAlias = ( + Literal[ + "DEPLOYING", + "DEPLOYED", + "FAILED", + "DEACTIVATING", + "DEACTIVATED", + "DELETING", + "TERMINATED", + "ARCHIVING", + "ARCHIVED", + "UNARCHIVING", + "BUILDING", + "UPLOADING", + ] + | str +) +"""Deployment status of a sandbox. Only ``DEPLOYED`` sandboxes run commands. + +Other values may be added, so do not treat this list as exhaustive. +""" + +SandboxExpirationAction: TypeAlias = Literal["DELETE"] | str +"""What an expiration policy does when met. + +Other values may be added, so do not treat this list as exhaustive. +""" + +SandboxPortProtocol: TypeAlias = Literal["HTTP", "TCP", "UDP", "TLS"] | str +"""Protocol of a sandbox port. + +Other values may be added, so do not treat this list as exhaustive. +""" + + +@dataclass(kw_only=True) +class SandboxEnvValue: + """Value of an environment variable in a sandbox.""" + + value: str + + secret: bool | None = None + """Whether the value is a secret. Defaults to true. + + Secret values come back masked as ``"****"``, unless read with + ``show_secrets``. + """ + + +@dataclass(kw_only=True) +class SandboxExpirationPolicyIdle: + """Expires a sandbox after a period without activity.""" + + after: timedelta + """How long the sandbox may go without activity.""" + + action: SandboxExpirationAction | None = None + """What happens when the policy is met. ``DELETE`` when unset.""" + + +@dataclass(kw_only=True) +class SandboxExpirationPolicyMaxAge: + """Expires a sandbox a period after its creation, regardless of activity.""" + + after: timedelta + """How long after creation the sandbox expires.""" + + action: SandboxExpirationAction | None = None + """What happens when the policy is met. ``DELETE`` when unset.""" + + +@dataclass(kw_only=True) +class SandboxExpirationPolicyDate: + """Expires a sandbox at a fixed time.""" + + at: datetime + """When the sandbox expires. Must be timezone-aware.""" + + action: SandboxExpirationAction | None = None + """What happens when the policy is met. ``DELETE`` when unset.""" + + +@dataclass(kw_only=True) +class SandboxExpirationPolicyUnknown: + """An expiration policy of a kind this SDK does not know. + + Kept as the server sent it, and sent back unchanged when given in a + lifecycle, so updating a sandbox does not lose it. + """ + + type: str | None + """The policy's type, when the server sent one as a string.""" + + raw: dict[str, Any] + """The policy exactly as the server sent it.""" + + +SandboxExpirationPolicy: TypeAlias = ( + SandboxExpirationPolicyIdle + | SandboxExpirationPolicyMaxAge + | SandboxExpirationPolicyDate + | SandboxExpirationPolicyUnknown +) +"""A condition for acting on a sandbox automatically. + +Durations the server reports in days or weeks convert at 24 hours per day. +""" + + +@dataclass(kw_only=True) +class SandboxLifecycle: + """When a sandbox is deleted automatically, and how long its record stays after.""" + + expiration_policies: Sequence[SandboxExpirationPolicy] | None = None + """Conditions for deleting the sandbox. Whichever is met first applies. + + Replaces all previous policies when updated. + """ + + terminated_retention: timedelta | None = None + """How long the record stays after the sandbox terminates, for log access. + + The server defaults to 5 minutes. + """ + + +@dataclass(kw_only=True) +class SandboxPort: + """A port the sandbox exposes.""" + + target: int + """Port number in the sandbox, 1 to 65535.""" + + name: str | None = None + + protocol: SandboxPortProtocol | None = None + + +@dataclass(kw_only=True) +class SandboxNetworkProxyRoute: + """A proxy rule injecting headers and body fields into requests to its destinations.""" + + destinations: Sequence[str] | None = None + """Destination domains the rule applies to. ``["*"]`` matches all.""" + + headers: Mapping[str, str] | None = None + """Headers to inject. Values may reference ``{{SECRET:name}}`` from ``secrets``.""" + + body: Mapping[str, str] | None = None + """Body fields to inject. Values may reference ``{{SECRET:name}}`` from ``secrets``.""" + + secrets: Mapping[str, str] | None = None + """Named secret values for this rule. + + Write-only, so always unset on :class:`SandboxInfo`. + """ + + +@dataclass(kw_only=True) +class SandboxNetworkProxy: + """Proxy configuration of a sandbox's network.""" + + allowed_domains: Sequence[str] | None = None + """When set, only these external domains are reachable. + + Supports wildcards such as ``*.example.com``. Takes precedence over + ``forbidden_domains``. + """ + + forbidden_domains: Sequence[str] | None = None + """When set, every external domain except these is reachable. Supports wildcards.""" + + bypass: Sequence[str] | None = None + """Domains reached directly rather than through the proxy. + + Supports wildcards. Local and private addresses always bypass it. + """ + + routing: Sequence[SandboxNetworkProxyRoute] | None = None + """Rules injecting headers and body fields into matching requests.""" + + +@dataclass(kw_only=True) +class SandboxNetwork: + """Network configuration of a sandbox, fixed at creation.""" + + subnet: str | None = None + """Subnet name. The server defaults to ``default``.""" + + proxy: SandboxNetworkProxy | None = None + """Routes the sandbox's HTTP traffic through the platform proxy.""" + + +@dataclass(kw_only=True) +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 at creation.""" + + url: str + """Base URL of the sandbox's execution API.""" + + status: SandboxStatus + + image: str | None + """Image reference, including its tag.""" + + memory: int | None + """Memory in megabytes, which also sets the CPU allocation.""" + + region: str | None + + envs: dict[str, SandboxEnvValue] + """Environment variables, by name. + + Secret values are masked, so passing these back to an update overwrites + their real values with the masks. + """ + + labels: dict[str, str] + + external_id: str | None + """Caller-owned identifier for external lookups.""" + + lifecycle: SandboxLifecycle | None + + ports: list[SandboxPort] + + network: SandboxNetwork | None + + created_at: datetime + + updated_at: datetime | None + + created_by: str | None + + updated_by: str | None + + last_used_at: datetime | None + + expires_in: timedelta | None + """Time left before automatic deletion, when expiration is configured.""" + + +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), + status=sandbox.status, + image=sandbox.image, + memory=sandbox.memory, + region=sandbox.region, + envs=_envs_from_api(sandbox.envs), + labels=dict(sandbox.labels or {}), + external_id=sandbox.external_id, + lifecycle=None + if sandbox.lifecycle is None + else _lifecycle_from_api(sandbox.lifecycle), + ports=ports_from_api(sandbox.ports), + network=None if sandbox.network is None else _network_from_api(sandbox.network), + 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=None + if sandbox.expires_in is None + else timedelta(seconds=sandbox.expires_in), + ) + + +def ports_from_api( + ports: Sequence[baseten.client.managementapi.SandboxPort] | None, +) -> list[SandboxPort]: + """Convert the control plane's ports, ``None`` meaning none.""" + return [ + SandboxPort(target=port.target, name=port.name, protocol=port.protocol) + for port in ports or [] + ] + + +def envs_to_api( + envs: Mapping[str, SandboxEnvValue], +) -> list[baseten.client.managementapi.SandboxEnv]: + """Convert environment variables to the control plane's form.""" + return [ + baseten.client.managementapi.SandboxEnv( + **set_fields(name=name, value=env.value, secret=env.secret) + ) + for name, env in envs.items() + ] + + +def lifecycle_to_api( + lifecycle: SandboxLifecycle, +) -> baseten.client.managementapi.SandboxLifecycle: + """Convert a lifecycle to the control plane's form.""" + return baseten.client.managementapi.SandboxLifecycle( + **set_fields( + expiration_policies=None + if lifecycle.expiration_policies is None + else [_policy_to_api(policy) for policy in lifecycle.expiration_policies], + terminated_retention=None + if lifecycle.terminated_retention is None + else format_duration(lifecycle.terminated_retention), + ) + ) + + +def ports_to_api( + ports: Sequence[SandboxPort], +) -> list[baseten.client.managementapi.SandboxPort]: + """Convert ports to the control plane's form.""" + return [ + baseten.client.managementapi.SandboxPort( + **set_fields(target=port.target, name=port.name, protocol=port.protocol) + ) + for port in ports + ] + + +def network_to_api( + network: SandboxNetwork, +) -> baseten.client.managementapi.SandboxNetwork: + """Convert a network configuration to the control plane's form.""" + proxy = network.proxy + return baseten.client.managementapi.SandboxNetwork( + **set_fields( + subnet=network.subnet, + proxy=None + if proxy is None + else baseten.client.managementapi.SandboxProxyConfig( + **set_fields( + allowed_domains=_optional_list(proxy.allowed_domains), + forbidden_domains=_optional_list(proxy.forbidden_domains), + bypass=_optional_list(proxy.bypass), + routing=None + if proxy.routing is None + else [ + baseten.client.managementapi.SandboxProxyTarget( + **set_fields( + destinations=_optional_list(route.destinations), + headers=_optional_dict(route.headers), + body=_optional_dict(route.body), + secrets=_optional_dict(route.secrets), + ) + ) + for route in proxy.routing + ], + ) + ), + ) + ) + + +def set_fields(**fields: Any) -> dict[str, Any]: + """Return the fields that are not ``None``. + + Passing only these to a generated model leaves the rest unset, so they are + left out of the request rather than sent as null or as the model's default. + """ + return {name: value for name, value in fields.items() if value is not None} + + +def _policy_to_api( + policy: SandboxExpirationPolicy, +) -> ( + baseten.client.managementapi.SandboxTTLIdleExpirationPolicy + | baseten.client.managementapi.SandboxTTLMaxAgeExpirationPolicy + | baseten.client.managementapi.SandboxDateExpirationPolicy + | dict[str, Any] +): + """Convert an expiration policy to the control plane's form.""" + if isinstance(policy, SandboxExpirationPolicyUnknown): + return policy.raw + action = "DELETE" if policy.action is None else policy.action + match policy: + case SandboxExpirationPolicyIdle(): + return baseten.client.managementapi.SandboxTTLIdleExpirationPolicy( + type="TTL_IDLE", action=action, value=format_duration(policy.after) + ) + case SandboxExpirationPolicyMaxAge(): + return baseten.client.managementapi.SandboxTTLMaxAgeExpirationPolicy( + type="TTL_MAX_AGE", action=action, value=format_duration(policy.after) + ) + case SandboxExpirationPolicyDate(): + return baseten.client.managementapi.SandboxDateExpirationPolicy( + type="DATE", action=action, value=policy.at + ) + + +def _lifecycle_from_api( + lifecycle: baseten.client.managementapi.SandboxLifecycle, +) -> SandboxLifecycle: + retention = lifecycle.terminated_retention + return SandboxLifecycle( + expiration_policies=None + if lifecycle.expiration_policies is None + else [_policy_from_api(policy) for policy in lifecycle.expiration_policies], + terminated_retention=None + if retention is None + else parse_duration(retention, "lifecycle terminated retention"), + ) + + +def _policy_from_api( + policy: baseten.client.managementapi.SandboxTTLIdleExpirationPolicy + | baseten.client.managementapi.SandboxTTLMaxAgeExpirationPolicy + | baseten.client.managementapi.SandboxDateExpirationPolicy + | dict[str, Any], +) -> SandboxExpirationPolicy: + match policy: + case baseten.client.managementapi.SandboxTTLIdleExpirationPolicy(): + return SandboxExpirationPolicyIdle( + after=parse_duration(policy.value, "TTL_IDLE expiration policy"), + action=policy.action, + ) + case baseten.client.managementapi.SandboxTTLMaxAgeExpirationPolicy(): + return SandboxExpirationPolicyMaxAge( + after=parse_duration(policy.value, "TTL_MAX_AGE expiration policy"), + action=policy.action, + ) + case baseten.client.managementapi.SandboxDateExpirationPolicy(): + return SandboxExpirationPolicyDate(at=policy.value, action=policy.action) + case _: + policy_type = policy.get("type") + return SandboxExpirationPolicyUnknown( + type=policy_type if isinstance(policy_type, str) else None, + raw=dict(policy), + ) + + +def _network_from_api( + network: baseten.client.managementapi.SandboxNetwork, +) -> SandboxNetwork: + proxy = network.proxy + return SandboxNetwork( + subnet=network.subnet, + proxy=None + if proxy is None + else SandboxNetworkProxy( + allowed_domains=proxy.allowed_domains, + forbidden_domains=proxy.forbidden_domains, + bypass=proxy.bypass, + routing=None + if proxy.routing is None + else [ + SandboxNetworkProxyRoute( + destinations=route.destinations, + headers=route.headers, + body=route.body, + secrets=route.secrets, + ) + for route in proxy.routing + ], + ), + ) + + +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="" if env.value is None else env.value, secret=env.secret + ) + return result + + +def _optional_list(values: Sequence[str] | None) -> list[str] | None: + return None if values is None else list(values) + + +def _optional_dict(values: Mapping[str, str] | None) -> dict[str, str] | None: + return None if values is None else dict(values) diff --git a/baseten/sandbox/_process.py b/baseten/sandbox/_process.py new file mode 100644 index 0000000..897fa2a --- /dev/null +++ b/baseten/sandbox/_process.py @@ -0,0 +1,751 @@ +from __future__ import annotations + +import asyncio +import builtins +import json +import time +from collections.abc import ( + AsyncGenerator, + AsyncIterator, + Generator, + Iterator, + Mapping, + Sequence, +) +from dataclasses import dataclass +from datetime import datetime, timedelta +from typing import Literal, TypeAlias + +import baseten.client.sandboxapi +from baseten.sandbox._common import ( + AsyncSandboxContext, + AsyncSandboxStream, + SandboxContext, + SandboxStream, + parse_timestamp, +) +from baseten.sandbox._errors import ( + SandboxAPIError, + SandboxError, + SandboxProcessWaitTimeoutError, +) +from baseten.sandbox._info import set_fields +from baseten.sandbox._retry import is_transient_reset_error + +_DEFAULT_WAIT_TIMEOUT = timedelta(seconds=60) +_DEFAULT_WAIT_POLL_INTERVAL = timedelta(seconds=1) + +_TERMINAL_STATUSES = frozenset({"completed", "failed", "killed", "stopped"}) + +# Statuses a wait rides out, since the sandbox or its edge may briefly fail. +_WAIT_RETRY_STATUSES = frozenset({408, 429, 500, 502, 503, 504}) + +SandboxProcessStatus: TypeAlias = ( + Literal["running", "completed", "failed", "killed", "stopped"] | str +) +"""Status of a process in a sandbox. + +Other values may be added, so do not treat this list as exhaustive. +""" + + +@dataclass(kw_only=True) +class SandboxProcessInfo: + """A process in a sandbox.""" + + pid: str + """Process identifier, usable wherever a process name is.""" + + 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 + """When the process exited. ``None`` while it runs.""" + + keep_alive: bool | None + """Whether scale-to-zero is disabled for this process.""" + + max_restarts: int | None + """Maximum number of restarts on failure. Negative means unlimited.""" + + restart_count: int | None + """Number of times the process has been restarted.""" + + restart_on_failure: bool | None + """Whether the process is restarted when it fails.""" + + stdin: bool | None + """Whether the process was started with a writable stdin pipe.""" + + +@dataclass(kw_only=True) +class SandboxProcessLogs: + """Output of a process, as returned by :meth:`SandboxProcess.logs`.""" + + stdout: str + + stderr: str + + logs: str + """Standard output and standard error, interleaved.""" + + +@dataclass(kw_only=True) +class SandboxProcessLogLine: + """One line of a process's output.""" + + stream: Literal["stdout", "stderr"] | None + """Which output the line came from, when the sandbox reports it.""" + + text: str + """The line, without its line ending.""" + + +@dataclass(kw_only=True) +class SandboxProcessExecEventOutput: + """Output from :meth:`SandboxProcess.exec_stream`. + + As the process wrote it, which may hold several lines, part of one, or + text with no newline at all, such as a prompt. If the process finished + before streaming started, each line instead arrives as its own event, + without its newline. + """ + + stream: Literal["stdout", "stderr"] + """Which output the text came from.""" + + text: str + + +@dataclass(kw_only=True) +class SandboxProcessExecEventExit: + """The exit of the process run by :meth:`SandboxProcess.exec_stream`. + + Always the last event. + """ + + process: SandboxProcessInfo + """The process as it exited.""" + + +SandboxProcessExecEvent: TypeAlias = ( + SandboxProcessExecEventOutput | SandboxProcessExecEventExit +) +"""An event from :meth:`SandboxProcess.exec_stream`.""" + + +class SandboxProcess: + """Processes in a sandbox.""" + + def __init__(self, context: SandboxContext) -> None: + """Internal. Obtained from :attr:`Sandbox.process`. + + :meta private: + """ + self._context = context + + def exec( + self, + *, + command: str, + name: str | None = None, + working_dir: str | None = None, + env: Mapping[str, str] | None = None, + timeout: timedelta | None = None, + wait_for_completion: bool | None = None, + keep_alive: bool | None = None, + max_restarts: int | None = None, + restart_on_failure: bool | None = None, + stdin: bool | None = None, + wait_for_ports: Sequence[int] | None = None, + ) -> SandboxProcessInfo: + """Start a command. + + Args: + command: Shell command to run. + name: Name to refer to the process by, instead of its PID. + working_dir: Directory to run the command in. + env: Environment variables for the process, on top of the + sandbox's own. + timeout: How long until the process is killed, rounded up to + whole seconds. Zero means never. When *keep_alive* is true, + defaults to 600 seconds. + wait_for_completion: Whether to return only once the process has + exited. Defaults to false. + keep_alive: Disable scale-to-zero while the process runs. + max_restarts: Maximum number of restarts on failure. A negative + value, such as -1, means unlimited. + restart_on_failure: Whether to restart the process when it fails. + stdin: Open a writable stdin pipe, fed by :meth:`write_stdin` and + closed by :meth:`close_stdin`. The pipe does not survive a + restart of the sandbox's execution API: the process then + reads end of file. + wait_for_ports: Ports the process must be listening on before + the call returns. + """ + # Not retried, since a command may have side effects even when its + # response is lost. + with self._context.api() as api: + process = api.post_process( + request=_process_request( + command=command, + name=name, + working_dir=working_dir, + env=env, + timeout=timeout, + wait_for_completion=wait_for_completion, + keep_alive=keep_alive, + max_restarts=max_restarts, + restart_on_failure=restart_on_failure, + stdin=stdin, + wait_for_ports=wait_for_ports, + ) + ) + return process_info_from_api(process) + + def exec_stream( + self, + *, + command: str, + name: str | None = None, + working_dir: str | None = None, + env: Mapping[str, str] | None = None, + timeout: timedelta | None = None, + keep_alive: bool | None = None, + max_restarts: int | None = None, + restart_on_failure: bool | None = None, + stdin: bool | None = None, + wait_for_ports: Sequence[int] | None = None, + ) -> SandboxStream[SandboxProcessExecEvent]: + """Run a command, streaming its output as it arrives and then its exit. + + Closing the stream early stops only the stream, not the process. + + Args: + command: Shell command to run. + name: Name to refer to the process by, instead of its PID. + working_dir: Directory to run the command in. + env: Environment variables for the process, on top of the + sandbox's own. + timeout: How long until the process is killed, rounded up to + whole seconds. Zero means never. When *keep_alive* is true, + defaults to 600 seconds. + keep_alive: Disable scale-to-zero while the process runs. + max_restarts: Maximum number of restarts on failure. A negative + value, such as -1, means unlimited. + restart_on_failure: Whether to restart the process when it fails. + stdin: Open a writable stdin pipe, fed by :meth:`write_stdin` and + closed by :meth:`close_stdin`. The pipe does not survive a + restart of the sandbox's execution API: the process then + reads end of file. + wait_for_ports: Ports the process must be listening on before + output starts. + """ + with self._context.api() as api: + response = api.post_process_raw( + accept="application/x-ndjson", + request=_process_request( + command=command, + name=name, + working_dir=working_dir, + env=env, + timeout=timeout, + wait_for_completion=True, + keep_alive=keep_alive, + max_restarts=max_restarts, + restart_on_failure=restart_on_failure, + stdin=stdin, + wait_for_ports=wait_for_ports, + ), + ) + return SandboxStream(response, _exec_events) + + def get(self, *, identifier: str) -> SandboxProcessInfo: + """Get a process. + + Args: + identifier: PID or name of the process. + """ + process = self._context.call_read( + lambda api: api.get_process_identifier(identifier=identifier) + ) + return process_info_from_api(process) + + # builtins.list, since the method shadows it in the class body. + def list(self) -> builtins.list[SandboxProcessInfo]: + """List the sandbox's processes, running and finished.""" + processes = self._context.call_read(lambda api: api.get_process()) + return [process_info_from_api(process) for process in processes.root] + + def stop(self, *, identifier: str) -> None: + """Request that a process exit gracefully, without waiting. + + Use :meth:`wait` for it to end, with status ``stopped``. + + Args: + identifier: PID or name of the process. + """ + with self._context.api() as api: + api.delete_process(identifier=identifier) + + def kill(self, *, identifier: str) -> None: + """Request that a process exit forcefully, without waiting. + + Use :meth:`wait` for it to end, with status ``killed``. + + Args: + identifier: PID or name of the process. + """ + with self._context.api() as api: + api.delete_process_kill(identifier=identifier) + + def logs(self, *, identifier: str) -> SandboxProcessLogs: + """Get a process's output so far. + + Args: + identifier: PID or name of the process. + """ + logs = self._context.call_read( + lambda api: api.get_process_logs(identifier=identifier) + ) + return SandboxProcessLogs( + stdout=logs.stdout, stderr=logs.stderr, logs=logs.logs + ) + + def stream_logs(self, *, identifier: str) -> SandboxStream[SandboxProcessLogLine]: + """Stream a process's output line by line as it arrives. + + Starts from the process's first output and ends when the process + exits. Works for any process, including one started elsewhere. A + partial line, such as a prompt, is yielded once its newline arrives + or the stream ends. Use :meth:`exec_stream` for output as it is + written. Closing the stream early stops only the stream, not the + process. + + Args: + identifier: PID or name of the process. + """ + with self._context.api() as api: + response = api.get_process_logs_stream( + identifier=identifier, accept="text/plain" + ) + return SandboxStream(response, _log_lines) + + def wait( + self, + *, + identifier: str, + timeout: timedelta | None = _DEFAULT_WAIT_TIMEOUT, + poll_interval: timedelta = _DEFAULT_WAIT_POLL_INTERVAL, + ) -> SandboxProcessInfo: + """Wait for a process to finish, returning it whether it succeeded or not. + + Checks the process at each poll interval, riding out errors the + sandbox may answer differently on the next check. + + Args: + identifier: PID or name of the process. + timeout: How long to wait. ``None`` waits with no limit. + poll_interval: How long between checks. + + Raises: + SandboxProcessWaitTimeoutError: The process was still running + when the wait timed out. The process keeps running. + """ + # A check already in flight cannot be interrupted, so the wait can + # run over by up to one check. + deadline = ( + None if timeout is None else time.monotonic() + timeout.total_seconds() + ) + last_error: Exception | None = None + while True: + try: + with self._context.api() as api: + process = api.get_process_identifier(identifier=identifier) + except Exception as err: + if not _is_retryable_wait_error(err): + raise + last_error = err + else: + last_error = None + info = _finished_process(identifier, process) + if info is not None: + return info + delay = poll_interval.total_seconds() + if deadline is not None: + remaining = deadline - time.monotonic() + if remaining <= delay: + time.sleep(max(remaining, 0)) + raise SandboxProcessWaitTimeoutError( + identifier=identifier + ) from last_error + time.sleep(delay) + + def write_stdin(self, *, identifier: str, data: str | bytes) -> None: + """Write to the standard input of a process started with ``stdin=True``. + + Args: + identifier: PID or name of the process. + data: Bytes to write, sent as is. A string is written as UTF-8. + Include any trailing newline the process expects. + """ + # Not retried, since a repeat could write the bytes twice. + with self._context.api() as api: + api.post_process_stdin(identifier=identifier, content=data) + + def close_stdin(self, *, identifier: str) -> None: + """Close a process's standard input, so it reads end of file. + + Args: + identifier: PID or name of the process. + """ + self._context.call_read( + lambda api: api.delete_process_stdin(identifier=identifier) + ) + + +class AsyncSandboxProcess: + """Async form of :class:`SandboxProcess`.""" + + def __init__(self, context: AsyncSandboxContext) -> None: + """Internal. Obtained from :attr:`AsyncSandbox.process`. + + :meta private: + """ + self._context = context + + async def exec( + self, + *, + command: str, + name: str | None = None, + working_dir: str | None = None, + env: Mapping[str, str] | None = None, + timeout: timedelta | None = None, + wait_for_completion: bool | None = None, + keep_alive: bool | None = None, + max_restarts: int | None = None, + restart_on_failure: bool | None = None, + stdin: bool | None = None, + wait_for_ports: Sequence[int] | None = None, + ) -> SandboxProcessInfo: + """Start a command. As :meth:`SandboxProcess.exec`.""" + with self._context.api() as api: + process = await api.post_process( + request=_process_request( + command=command, + name=name, + working_dir=working_dir, + env=env, + timeout=timeout, + wait_for_completion=wait_for_completion, + keep_alive=keep_alive, + max_restarts=max_restarts, + restart_on_failure=restart_on_failure, + stdin=stdin, + wait_for_ports=wait_for_ports, + ) + ) + return process_info_from_api(process) + + async def exec_stream( + self, + *, + command: str, + name: str | None = None, + working_dir: str | None = None, + env: Mapping[str, str] | None = None, + timeout: timedelta | None = None, + keep_alive: bool | None = None, + max_restarts: int | None = None, + restart_on_failure: bool | None = None, + stdin: bool | None = None, + wait_for_ports: Sequence[int] | None = None, + ) -> AsyncSandboxStream[SandboxProcessExecEvent]: + """Run a command, streaming its output. As :meth:`SandboxProcess.exec_stream`.""" + with self._context.api() as api: + response = await api.post_process_raw( + accept="application/x-ndjson", + request=_process_request( + command=command, + name=name, + working_dir=working_dir, + env=env, + timeout=timeout, + wait_for_completion=True, + keep_alive=keep_alive, + max_restarts=max_restarts, + restart_on_failure=restart_on_failure, + stdin=stdin, + wait_for_ports=wait_for_ports, + ), + ) + return AsyncSandboxStream(response, _exec_events_async) + + async def get(self, *, identifier: str) -> SandboxProcessInfo: + """Get a process. As :meth:`SandboxProcess.get`.""" + process = await self._context.call_read( + lambda api: api.get_process_identifier(identifier=identifier) + ) + return process_info_from_api(process) + + # builtins.list, since the method shadows it in the class body. + async def list(self) -> builtins.list[SandboxProcessInfo]: + """List the sandbox's processes. As :meth:`SandboxProcess.list`.""" + processes = await self._context.call_read(lambda api: api.get_process()) + return [process_info_from_api(process) for process in processes.root] + + async def stop(self, *, identifier: str) -> None: + """Request that a process exit gracefully. As :meth:`SandboxProcess.stop`.""" + with self._context.api() as api: + await api.delete_process(identifier=identifier) + + async def kill(self, *, identifier: str) -> None: + """Request that a process exit forcefully. As :meth:`SandboxProcess.kill`.""" + with self._context.api() as api: + await api.delete_process_kill(identifier=identifier) + + async def logs(self, *, identifier: str) -> SandboxProcessLogs: + """Get a process's output so far. As :meth:`SandboxProcess.logs`.""" + logs = await self._context.call_read( + lambda api: api.get_process_logs(identifier=identifier) + ) + return SandboxProcessLogs( + stdout=logs.stdout, stderr=logs.stderr, logs=logs.logs + ) + + async def stream_logs( + self, *, identifier: str + ) -> AsyncSandboxStream[SandboxProcessLogLine]: + """Stream a process's output. As :meth:`SandboxProcess.stream_logs`.""" + with self._context.api() as api: + response = await api.get_process_logs_stream( + identifier=identifier, accept="text/plain" + ) + return AsyncSandboxStream(response, _log_lines_async) + + async def wait( + self, + *, + identifier: str, + timeout: timedelta | None = _DEFAULT_WAIT_TIMEOUT, + poll_interval: timedelta = _DEFAULT_WAIT_POLL_INTERVAL, + ) -> SandboxProcessInfo: + """Wait for a process to finish. As :meth:`SandboxProcess.wait`.""" + last_error: Exception | None = None + scope = asyncio.timeout(None if timeout is None else timeout.total_seconds()) + try: + async with scope: + while True: + try: + with self._context.api() as api: + process = await api.get_process_identifier( + identifier=identifier + ) + except Exception as err: + if not _is_retryable_wait_error(err): + raise + last_error = err + else: + last_error = None + info = _finished_process(identifier, process) + if info is not None: + return info + await asyncio.sleep(poll_interval.total_seconds()) + except TimeoutError: + # Only this wait's own timeout, not a caller's or one from a + # check, becomes the wait timeout error. + if not scope.expired(): + raise + raise SandboxProcessWaitTimeoutError(identifier=identifier) from last_error + + async def write_stdin(self, *, identifier: str, data: str | bytes) -> None: + """Write to a process's standard input. As :meth:`SandboxProcess.write_stdin`.""" + with self._context.api() as api: + await api.post_process_stdin(identifier=identifier, content=data) + + async def close_stdin(self, *, identifier: str) -> None: + """Close a process's standard input. As :meth:`SandboxProcess.close_stdin`.""" + await self._context.call_read( + lambda api: api.delete_process_stdin(identifier=identifier) + ) + + +def process_info_from_api( + process: baseten.client.sandboxapi.ProcessResponse, +) -> SandboxProcessInfo: + """Convert a process from the execution API.""" + return SandboxProcessInfo( + pid=process.pid, + name=process.name, + command=process.command, + status=process.status, + exit_code=process.exitCode, + stdout=process.stdout, + stderr=process.stderr, + logs=process.logs, + working_dir=process.workingDir, + started_at=parse_timestamp(process.startedAt, "process startedAt"), + # Empty while the process runs. + completed_at=parse_timestamp(process.completedAt, "process completedAt") + if process.completedAt + else None, + keep_alive=process.keepAlive, + max_restarts=process.maxRestarts, + restart_count=process.restartCount, + restart_on_failure=process.restartOnFailure, + stdin=process.stdin, + ) + + +def _process_request( + *, + command: str, + name: str | None, + working_dir: str | None, + env: Mapping[str, str] | None, + timeout: timedelta | None, + wait_for_completion: bool | None, + keep_alive: bool | None, + max_restarts: int | None, + restart_on_failure: bool | None, + stdin: bool | None, + wait_for_ports: Sequence[int] | None, +) -> baseten.client.sandboxapi.ProcessRequest: + """Build the execution API request to run a command.""" + return baseten.client.sandboxapi.ProcessRequest( + command=command, + **set_fields( + name=name, + workingDir=working_dir, + env=None if env is None else dict(env), + # Whole seconds, rounded up. + timeout=None if timeout is None else -(-timeout // timedelta(seconds=1)), + waitForCompletion=wait_for_completion, + keepAlive=keep_alive, + maxRestarts=max_restarts, + restartOnFailure=restart_on_failure, + stdin=stdin, + waitForPorts=None if wait_for_ports is None else [*wait_for_ports], + ), + ) + + +def _finished_process( + identifier: str, process: baseten.client.sandboxapi.ProcessResponse +) -> SandboxProcessInfo | None: + """Return a process a wait is done with, or ``None`` while it runs.""" + if process.status in _TERMINAL_STATUSES: + return process_info_from_api(process) + # Unlike elsewhere, an unknown status raises rather than being waited on, + # keeping existing behavior. + if process.status != "running": + raise SandboxError(f"process {identifier} has unknown status {process.status}") + return None + + +def _is_retryable_wait_error(err: Exception) -> bool: + """Report whether a wait should ride out an error from a check.""" + if isinstance(err, SandboxAPIError): + return err.status in _WAIT_RETRY_STATUSES + return is_transient_reset_error(err) + + +def _exec_event(line: str) -> SandboxProcessExecEvent | None: + """Parse one line of a run's stream, or ``None`` for a line to skip.""" + if line == "": + return None + record = json.loads(line) + if not isinstance(record, dict): + return None + kind = record.get("type") + data = record.get("data") + if kind == "stdout" or kind == "stderr": + return SandboxProcessExecEventOutput( + stream=kind, text="" if data is None else str(data) + ) + if kind == "result": + process = baseten.client.sandboxapi.ProcessResponse.model_validate_json( + str(data) + ) + return SandboxProcessExecEventExit(process=process_info_from_api(process)) + if kind == "error": + # The response status is already sent, so a failure after streaming + # starts arrives as a record that ends the stream. + raise SandboxError(f"process stream failed: {'' if data is None else data}") + # keepalive and any other record type are skipped, so new ones don't break + # streaming. + return None + + +def _exec_events( + lines: Iterator[str], +) -> Generator[SandboxProcessExecEvent, None, None]: + """Parse a run's stream into events, ending with its exit.""" + for line in lines: + event = _exec_event(line) + if event is None: + continue + yield event + if isinstance(event, SandboxProcessExecEventExit): + return + raise SandboxError("the process stream ended before reporting the process's exit") + + +async def _exec_events_async( + lines: AsyncIterator[str], +) -> AsyncGenerator[SandboxProcessExecEvent, None]: + """Async form of :func:`_exec_events`.""" + async for line in lines: + event = _exec_event(line) + if event is None: + continue + yield event + if isinstance(event, SandboxProcessExecEventExit): + return + raise SandboxError("the process stream ended before reporting the process's exit") + + +def _log_line(line: str) -> SandboxProcessLogLine | None: + """Parse one line of a log stream, or ``None`` for a line to skip.""" + if line.startswith("[keepalive]"): + return None + if line.startswith("stdout:"): + return SandboxProcessLogLine(stream="stdout", text=line[len("stdout:") :]) + if line.startswith("stderr:"): + return SandboxProcessLogLine(stream="stderr", text=line[len("stderr:") :]) + return SandboxProcessLogLine(stream=None, text=line) + + +def _log_lines(lines: Iterator[str]) -> Generator[SandboxProcessLogLine, None, None]: + """Parse a log stream into lines.""" + for line in lines: + log_line = _log_line(line) + if log_line is not None: + yield log_line + + +async def _log_lines_async( + lines: AsyncIterator[str], +) -> AsyncGenerator[SandboxProcessLogLine, None]: + """Async form of :func:`_log_lines`.""" + async for line in lines: + log_line = _log_line(line) + if log_line is not None: + yield log_line diff --git a/baseten/sandbox/_retry.py b/baseten/sandbox/_retry.py new file mode 100644 index 0000000..b388c65 --- /dev/null +++ b/baseten/sandbox/_retry.py @@ -0,0 +1,178 @@ +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 SandboxAPIError, SandboxGatewayError + +_T = TypeVar("_T") + +_DEFAULT_READ_MAX_RETRIES = 5 +_DEFAULT_GATEWAY_MAX_RETRIES = 2 +_DEFAULT_UPLOAD_MAX_RETRIES = 3 + +# Retry backoff bounds in seconds. Module variables so tests can shorten them. +_backoff_base = 0.2 +_backoff_max = 2.0 + +# Texts of the protocol errors httpx raises when a connection dropped before +# a response arrived, or HTTP/2 reset a stream or closed the connection. Other +# protocol errors, such as a malformed response, are not transient. +_TRANSIENT_PROTOCOL_ERROR_MARKERS = ( + # HTTP/1.1 and HTTP/2: the connection closed before a response. + "Server disconnected", + # HTTP/1.1: the connection closed partway through the body. + "peer closed connection without sending complete message body", + # HTTP/2: the server reset one stream, including when rate limiting + # stream resets with ENHANCE_YOUR_CALM. + "StreamReset", + # HTTP/2: the server sent GOAWAY. + "ConnectionTerminated", +) + +# How far down an error's cause chain to look. +_MAX_CAUSE_DEPTH = 10 + + +@dataclass(kw_only=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 | None = None + """Retries of a read after a dropped or reset connection. Defaults to 5.""" + + gateway_max_retries: int | None = None + """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. Defaults to 2. + """ + + upload_max_retries: int | None = None + """Retries of a file upload after a dropped or reset connection. Defaults to 3.""" + + +@dataclass(frozen=True) +class ResolvedRetryOptions: + """Retry budgets with the defaults filled in.""" + + read_max_retries: int + gateway_max_retries: int + upload_max_retries: int + + +def resolve_retry_options(options: SandboxRetryOptions | None) -> ResolvedRetryOptions: + """Fill in the default retry budgets.""" + options = options or SandboxRetryOptions() + return ResolvedRetryOptions( + read_max_retries=_DEFAULT_READ_MAX_RETRIES + if options.read_max_retries is None + else options.read_max_retries, + gateway_max_retries=_DEFAULT_GATEWAY_MAX_RETRIES + if options.gateway_max_retries is None + else options.gateway_max_retries, + upload_max_retries=_DEFAULT_UPLOAD_MAX_RETRIES + if options.upload_max_retries is None + else options.upload_max_retries, + ) + + +def is_transient_reset_error(err: BaseException) -> bool: + """Report whether an error is a dropped or reset connection. + + Only then, as opposed to a response from the server, is repeating an + idempotent call safe. + """ + # A response came back, so the server handled the request, even when its + # body happens to contain one of the markers. + if isinstance(err, SandboxAPIError): + return False + current: BaseException | None = err + for _ in range(_MAX_CAUSE_DEPTH): + if current is None: + break + # Reset, refused, aborted, and broken pipe, on every platform. + if isinstance(current, ConnectionError): + return True + # Only the operating system's connection timeout carries an errno; a + # socket timeout the caller configured has none, and is theirs to + # handle. + if isinstance(current, TimeoutError) and current.errno is not None: + return True + if isinstance(current, httpx.RemoteProtocolError) and any( + marker in str(current) for marker in _TRANSIENT_PROTOCOL_ERROR_MARKERS + ): + return True + current = current.__cause__ or current.__context__ + return False + + +def backoff_delay(attempt: int) -> float: + """Return the seconds to wait before a retry. + + Exponential backoff capped at a maximum, plus up to one base delay of + jitter so that many clients failing together do not retry together. + """ + capped = min(_backoff_base * 2 ** (attempt - 1), _backoff_max) + return capped + random.uniform(0, _backoff_base) + + +def retry_idempotent( + call: Callable[[], _T], *, max_retries: int, gateway_max_retries: int +) -> _T: + """Run an idempotent call, retrying transient failures. + + Dropped connections and edge gateway errors are retried on separate + budgets. 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(gateway_attempt) + except Exception as err: + attempt += 1 + if attempt > max_retries or not is_transient_reset_error(err): + raise + delay = backoff_delay(attempt) + time.sleep(delay) + + +async def retry_idempotent_async( + call: Callable[[], Awaitable[_T]], *, max_retries: int, gateway_max_retries: int +) -> _T: + """Async form 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(gateway_attempt) + except Exception as err: + attempt += 1 + if attempt > max_retries or not is_transient_reset_error(err): + raise + delay = backoff_delay(attempt) + await asyncio.sleep(delay) diff --git a/baseten/sandbox/_sandbox.py b/baseten/sandbox/_sandbox.py new file mode 100644 index 0000000..7128cb6 --- /dev/null +++ b/baseten/sandbox/_sandbox.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +import baseten.client.sandboxapi +from baseten.sandbox._common import AsyncSandboxContext, SandboxContext +from baseten.sandbox._fs import AsyncSandboxFileSystem, SandboxFileSystem +from baseten.sandbox._process import AsyncSandboxProcess, SandboxProcess + + +class Sandbox: + """One sandbox, to run processes and work with files in. + + Get one from :meth:`SandboxClient.create`, :meth:`SandboxClient.get`, or + :meth:`SandboxClient.sandbox_from_url`. + """ + + def __init__(self, context: SandboxContext) -> None: + """Internal. Use the :class:`SandboxClient` methods that return one instead. + + :meta private: + """ + self._context = context + self._process = SandboxProcess(context) + self._fs = SandboxFileSystem(context, self._process) + + @property + def name(self) -> str: + """Name of the sandbox. Empty when made from a URL alone.""" + return self._context.name + + @property + def url(self) -> str: + """Base URL of the sandbox's execution API.""" + return self._context.url + + @property + def fs(self) -> SandboxFileSystem: + """Files and directories in the sandbox.""" + return self._fs + + @property + def process(self) -> SandboxProcess: + """Processes in the sandbox.""" + return self._process + + @property + def raw_api(self) -> baseten.client.sandboxapi.ApiClient: + """The generated client for this sandbox's execution API, authenticated. + + For operations this SDK does not wrap. The generated API surface is not + covered by stability guarantees and may change between versions. + """ + return self._context.raw_api + + +class AsyncSandbox: + """Async form of :class:`Sandbox`. + + Get one from :meth:`AsyncSandboxClient.create`, + :meth:`AsyncSandboxClient.get`, or + :meth:`AsyncSandboxClient.sandbox_from_url`. + """ + + def __init__(self, context: AsyncSandboxContext) -> None: + """Internal. Use the :class:`AsyncSandboxClient` methods that return one instead. + + :meta private: + """ + self._context = context + self._process = AsyncSandboxProcess(context) + self._fs = AsyncSandboxFileSystem(context, self._process) + + @property + def name(self) -> str: + """Name of the sandbox. Empty when made from a URL alone.""" + return self._context.name + + @property + def url(self) -> str: + """Base URL of the sandbox's execution API.""" + return self._context.url + + @property + def fs(self) -> AsyncSandboxFileSystem: + """Files and directories in the sandbox.""" + return self._fs + + @property + def process(self) -> AsyncSandboxProcess: + """Processes in the sandbox.""" + return self._process + + @property + def raw_api(self) -> baseten.client.sandboxapi.AsyncApiClient: + """The generated client for this sandbox's execution API, authenticated. + + For operations this SDK does not wrap. The generated API surface is not + covered by stability guarantees and may change between versions. + """ + return self._context.raw_api diff --git a/baseten/sandbox/_zip.py b/baseten/sandbox/_zip.py new file mode 100644 index 0000000..5e676da --- /dev/null +++ b/baseten/sandbox/_zip.py @@ -0,0 +1,238 @@ +from __future__ import annotations + +import asyncio +import inspect +import io +import os +import shutil +import stat +import zipfile +from collections.abc import Generator, Mapping +from dataclasses import dataclass +from typing import IO + +from baseten.sandbox._images_ignore import ( + DOCKERIGNORE_FILE_NAME, + AsyncImageIgnoreFileFunc, + ImageIgnoreFileFunc, + ImageIgnoreFileOptions, +) + +# Asks for an ignore decision on each path, and gets the entries at the end. +DirectoryWalk = Generator[ImageIgnoreFileOptions, bool, "list[ZipEntry]"] + +# The MS-DOS directory attribute, set alongside the Unix mode on a directory. +_MSDOS_DIRECTORY = 0x10 + +# The creator system that tells readers the high bits of an entry's external +# attributes hold a Unix mode. +_CREATOR_UNIX = 3 + + +@dataclass(kw_only=True) +class ZipEntry: + """One file or directory of a build context.""" + + path: str + """Path in the zip, with forward slashes and no trailing slash.""" + + mode: int + """Unix permission bits.""" + + is_dir: bool = False + + data: bytes = b"" + """A file's content, unless ``source_path`` is set.""" + + source_path: str | None = None + """A local file read while zipping, so a large build context is never held in memory whole.""" + + +def write_zip(entries: list[ZipEntry], file: IO[bytes]) -> None: + """Zip entries, in order, into a file.""" + with zipfile.ZipFile(file, "w") as archive: + for entry in entries: + if entry.is_dir: + info = zipfile.ZipInfo(entry.path + "/") + info.external_attr = ( + (stat.S_IFDIR | entry.mode) << 16 + ) | _MSDOS_DIRECTORY + else: + info = zipfile.ZipInfo(entry.path) + info.external_attr = (stat.S_IFREG | entry.mode) << 16 + info.compress_type = zipfile.ZIP_DEFLATED + info.create_system = _CREATOR_UNIX + if entry.source_path is None: + archive.writestr(info, entry.data) + continue + with open(entry.source_path, "rb") as source: + # ZIP64 must be chosen before writing an entry of unknown + # size, so it goes by the file's size now. + large = os.fstat(source.fileno()).st_size >= zipfile.ZIP64_LIMIT + with archive.open(info, "w", force_zip64=large) as target: + shutil.copyfileobj(source, target) + + +def check_zip_dockerfile(data: bytes) -> None: + """Check that a caller's zip has a ``Dockerfile`` file at its root.""" + try: + archive = zipfile.ZipFile(io.BytesIO(data)) + except zipfile.BadZipFile as err: + raise ValueError(f"the zip source is not a readable zip: {err}") from err + with archive: + for info in archive.infolist(): + file_type = stat.S_IFMT(info.external_attr >> 16) + # A link named Dockerfile would only fail later, in the build. A + # type of 0 is a zip made without Unix modes. + if ( + info.filename == "Dockerfile" + and not info.is_dir() + and file_type in (0, stat.S_IFREG) + ): + return + raise ValueError("the zip source has no Dockerfile file at its root") + + +def files_zip_entries(files: Mapping[str, str | bytes]) -> list[ZipEntry]: + """Build entries from in-memory files, each with mode 0644.""" + if "Dockerfile" not in files: + raise ValueError("the files source has no Dockerfile") + entries: list[ZipEntry] = [] + # Sorted, so the same files always zip the same way. + for path in sorted(files): + segments = path.split("/") + if "\\" in path or any(segment in ("", ".", "..") for segment in segments): + raise ValueError( + f"files path {path!r} must be relative, with forward slashes and " + "no empty, . or .. segments" + ) + content = files[path] + data = content.encode() if isinstance(content, str) else content + entries.append(ZipEntry(path=path, data=data, mode=0o644)) + return entries + + +def check_directory_dockerfile(directory: str) -> None: + """Check that a directory has a ``Dockerfile`` file at its root.""" + # Checked before reading anything, so a mistaken directory, such as a home + # directory, fails fast instead of being read in full. + if not os.path.isfile(os.path.join(directory, "Dockerfile")): + raise ValueError(f"directory {directory} has no Dockerfile at its root") + + +def walk_directory(directory: str, prefix: str) -> DirectoryWalk: + """Walk a local directory into entries, each path prefixed with ``prefix``. + + Yields every path but the root ``Dockerfile`` and ``.dockerignore``, which + Docker always keeps in a build context, and leaves out each one it is sent + true for. Files keep their modes, a link to a file within the directory + stores a copy of the file, a link to a file outside it is an error, a link + to a directory stores an empty directory, and a broken link is left out. + """ + # Resolved, so a link's resolved target can be checked against it. + root = os.path.realpath(directory) + entries: list[ZipEntry] = [] + yield from _walk_tree(root, directory, "", prefix, entries) + return entries + + +def _walk_tree( + root: str, directory: str, rel_dir: str, prefix: str, entries: list[ZipEntry] +) -> Generator[ImageIgnoreFileOptions, bool, None]: + """Walk the children of a directory whose path relative to the walked one is ``rel_dir``.""" + with os.scandir(directory) as scan: + # Sorted by name, so the same directory always zips the same way. + children = sorted(scan, key=lambda child: child.name) + for child in children: + rel_path = rel_dir + child.name + always_kept = rel_dir == "" and child.name in ( + "Dockerfile", + DOCKERIGNORE_FILE_NAME, + ) + if not always_kept: + ignored = yield ImageIgnoreFileOptions( + rel_path=rel_path, is_directory=child.is_dir(follow_symlinks=False) + ) + if ignored: + continue + archive_path = prefix + rel_path + try: + info = os.stat(child.path) + except FileNotFoundError: + if child.is_symlink(): + # A broken link has nothing to archive. + continue + raise + mode = stat.S_IMODE(info.st_mode) & 0o777 + if stat.S_ISDIR(info.st_mode): + entries.append(ZipEntry(path=archive_path, mode=mode, is_dir=True)) + # A link to a directory is stored as an empty directory, not + # followed, so a link cannot pull in the rest of the machine. + if not child.is_symlink(): + yield from _walk_tree(root, child.path, rel_path + "/", prefix, entries) + elif stat.S_ISREG(info.st_mode): + # A link to a file is stored as a copy of the file, but only one + # within the directory, so a link cannot pull in a file from + # elsewhere on the machine. + if child.is_symlink(): + target = os.path.realpath(child.path) + if not _is_within(root, target): + raise ValueError( + f"{child.path} is a link to {target}, outside {root}; only " + "links to files within it can be pushed" + ) + entries.append( + ZipEntry(path=archive_path, mode=mode, source_path=child.path) + ) + # Anything else, such as a socket, has no content to archive. + + +def _is_within(root: str, path: str) -> bool: + """Report whether a resolved path is inside a resolved root.""" + try: + relative = os.path.relpath(path, root) + except ValueError: + # On another Windows drive. + return False + return ( + relative != os.pardir + and not relative.startswith(os.pardir + os.sep) + and not os.path.isabs(relative) + ) + + +def run_walk(walk: DirectoryWalk, ignore: ImageIgnoreFileFunc | None) -> list[ZipEntry]: + """Run a directory walk, deciding each path with ``ignore``, which keeps all when ``None``.""" + try: + options = next(walk) + while True: + options = walk.send(ignore is not None and ignore(options)) + except StopIteration as stop: + return stop.value + + +async def run_walk_async( + walk: DirectoryWalk, ignore: AsyncImageIgnoreFileFunc | None +) -> list[ZipEntry]: + """Async form of :func:`run_walk`, reading the filesystem on a worker thread.""" + done, step = await asyncio.to_thread(_walk_step, walk, None) + while not done: + assert isinstance(step, ImageIgnoreFileOptions) + ignored = False + if ignore is not None: + result = ignore(step) + ignored = await result if inspect.isawaitable(result) else result + done, step = await asyncio.to_thread(_walk_step, walk, ignored) + assert isinstance(step, list) + return step + + +def _walk_step( + walk: DirectoryWalk, ignored: bool | None +) -> tuple[bool, ImageIgnoreFileOptions | list[ZipEntry]]: + """Advance a walk, giving whether it is done and its next path or its entries.""" + # Caught here, since a StopIteration cannot cross into a future. + try: + return False, next(walk) if ignored is None else walk.send(ignored) + except StopIteration as stop: + return True, stop.value diff --git a/pyproject.toml b/pyproject.toml index 3342aa0..c269cc4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ dev = [ # bump would silently rewrite the models. "datamodel-code-generator>=0.81.0,<0.82", "furo>=2025.12.19", + "h2>=4.4.1", "poethepoet>=0.48", "pytest>=9.1.1", "pytest-asyncio>=1.4.0", diff --git a/scripts/apigen/clientgen.py b/scripts/apigen/clientgen.py index f00246b..84525d7 100644 --- a/scripts/apigen/clientgen.py +++ b/scripts/apigen/clientgen.py @@ -30,6 +30,9 @@ class _Operation: body_content_type: str json_responses: list[tuple[int, str]] raw_accepts: list[str] + # Whether a success response declares headers, which an operation without + # a body can only return as the raw response. + response_headers: bool success_codes: list[int] error_codes: dict[int, str] | None summary: str @@ -122,6 +125,7 @@ def _extract_operations(spec: dict) -> list[_Operation]: body_content_type=_body_content_type(spec, op_data), json_responses=json_responses, raw_accepts=_raw_response_accepts(spec, op_data), + response_headers=_has_response_headers(spec, op_data), success_codes=_extract_success_codes(op_data, http_method, path), error_codes=_error_code_map(spec, op_data), summary=op_data.get("summary", ""), @@ -240,6 +244,14 @@ def _raw_response_accepts(spec: dict, op: dict) -> list[str]: return sorted(accepts) +def _has_response_headers(spec: dict, op: dict) -> bool: + for resp_node in _success_responses(op).values(): + resolved = _resolve_ref(spec, resp_node) + if resolved is not None and resolved.get("headers"): + return True + return False + + def _has_json_content(node: dict | None) -> bool: if node is None: return False @@ -270,8 +282,13 @@ def _render_client(ops: list[_Operation]) -> str: # Status dispatch covers an operation with more than one success code and # at least one JSON schema among them, since the others may be bodyless. has_status_resp = any(op.json_responses and len(op.success_codes) > 1 for op in ops) - has_raw = any(op.raw_accepts for op in ops) - has_no_resp = any(not op.json_responses and not op.raw_accepts for op in ops) + has_raw = any( + op.raw_accepts or (not op.json_responses and op.response_headers) for op in ops + ) + has_no_resp = any( + not op.json_responses and not op.raw_accepts and not op.response_headers + for op in ops + ) error_refs = sorted({ref for op in ops for ref in (op.error_codes or {}).values()}) @@ -389,9 +406,33 @@ class {cls}: They may change without notice between versions. \""" - def __init__(self, http_client: {http_cls}) -> None: - \"""Create a new client. The caller is responsible for closing *http_client*.\""" + def __init__( + self, + http_client: {http_cls}, + *, + base_url: str | None = None, + auth: httpx.Auth | None = None, + headers: Mapping[str, str] | None = None, + timeout: httpx.Timeout | None = None, + ) -> None: + \"""Create a new client. The caller is responsible for closing *http_client*. + + The keyword arguments apply to this client's requests only, so one + HTTP client can serve several APIs. Each one left ``None`` falls back + to the HTTP client's own setting. + + Args: + http_client: HTTP client to send requests with. + base_url: Base URL that request paths are appended to. + auth: Authentication for each request. + headers: Headers for each request, added to the HTTP client's own. + timeout: Timeouts for each request. + \""" self._http_client = http_client + self._base_url = base_url + self._auth = auth + self._headers = headers + self._timeout = timeout """ for op in ops: @@ -399,7 +440,9 @@ def __init__(self, http_client: {http_cls}) -> None: src += "\n" + _render_method(op, is_async=is_async, is_raw=False) if op.raw_accepts: src += "\n" + _render_method(op, is_async=is_async, is_raw=True) - elif op.raw_accepts: + elif op.raw_accepts or op.response_headers: + # With no body, declared headers are the result, so only the raw + # response carries them. src += "\n" + _render_method(op, is_async=is_async, is_raw=True) else: src += "\n" + _render_method(op, is_async=is_async, is_raw=False) @@ -431,7 +474,7 @@ def _build_request(self, request: _ApiRequest) -> httpx.Request: json_body = None content_body = None files_body = None - headers: dict[str, str] = {{}} + headers: dict[str, str] = {{**(self._headers or {{}})}} if request.accept is not None: headers["Accept"] = request.accept if request.body is not None: @@ -465,16 +508,20 @@ def _build_request(self, request: _ApiRequest) -> httpx.Request: params = request.query return self._http_client.build_request( request.method, - path, + path if self._base_url is None else self._base_url.rstrip("/") + path, json=json_body, content=content_body, files=files_body, params=params, headers=headers, + timeout=httpx.USE_CLIENT_DEFAULT if self._timeout is None else self._timeout, ) {adef} _do(self, request: _ApiRequest) -> httpx.Response: - response = {aw}self._http_client.send(self._build_request(request)) + response = {aw}self._http_client.send( + self._build_request(request), + auth=httpx.USE_CLIENT_DEFAULT if self._auth is None else self._auth, + ) if response.status_code not in request.success_codes: {error_dispatch}\ raise ResponseError(status_code=response.status_code, body=response.text) @@ -485,7 +532,11 @@ def _build_request(self, request: _ApiRequest) -> httpx.Request: resp_read = "await response.aread()" if is_async else "response.read()" src += f""" {adef} _do_raw(self, request: _ApiRequest) -> httpx.Response: - response = {aw}self._http_client.send(self._build_request(request), stream=True) + response = {aw}self._http_client.send( + self._build_request(request), + auth=httpx.USE_CLIENT_DEFAULT if self._auth is None else self._auth, + stream=True, + ) if response.status_code not in request.success_codes: {resp_read} {error_dispatch}\ @@ -563,7 +614,7 @@ def _render_method(op: _Operation, *, is_async: bool, is_raw: bool) -> str: accept_type = "Literal[" + ", ".join(repr(c) for c in content_types) + "]" kwargs.append(f"accept: {accept_type}") accept_expr = "accept" - elif is_raw: + elif is_raw and content_types: accept_expr = repr(content_types[0]) else: accept_expr = "None" @@ -638,6 +689,8 @@ def _render_method(op: _Operation, *, is_async: bool, is_raw: bool) -> str: summary = summary.rstrip(".") + ". Returns the response unread" if accept_expr == "accept": summary += ", in the requested content type" + elif accept_expr == "None": + summary += ", since its headers are the result" summary += ". The caller must close the response." if body_arg == "files": summary = summary.rstrip(".") + ". *files* is sent as httpx ``files``." diff --git a/scripts/apigen/preprocess.py b/scripts/apigen/preprocess.py index 8f0ade6..1c1db7c 100644 --- a/scripts/apigen/preprocess.py +++ b/scripts/apigen/preprocess.py @@ -90,6 +90,18 @@ def preprocess_spec(data: bytes) -> bytes: # generate dict[str, Any] instead. _fix_bare_object_schemas(doc) + # TODO: Remove once the spec drops tags from SandboxImageV1, as the server + # already has. Image get and delete responses omit tags, which the spec + # still requires, so the generated model rejects them. + for part in ( + doc.get("components", {}) + .get("schemas", {}) + .get("SandboxImageV1", {}) + .get("allOf", []) + ): + if "tags" in part.get("required", []): + part["required"].remove("tags") + # Schema renames: a trailing V1 is stripped (ModelV1 -> Model) and names # that are not valid Python identifiers are folded to PascalCase # (archive.Change -> ArchiveChange). diff --git a/scripts/apigen/specs/config.schema.json b/scripts/apigen/specs/config.schema.json index 560fe8c..b87664c 100644 --- a/scripts/apigen/specs/config.schema.json +++ b/scripts/apigen/specs/config.schema.json @@ -199,7 +199,7 @@ "default": "", "description": "The path to the Docker image.", "examples": [ - "vllm/vllm-openai:v0.7.3", + "vllm/vllm-openai:v0.29.0", "nvcr.io/nvidia/nemo:23.03" ], "title": "Image", @@ -207,9 +207,9 @@ }, "python_executable_path": { "default": "", - "description": "A path to the Python executable on the image.", + "description": "A path to the Python executable on the image. Truss uses `python3` from the image's PATH when this is empty.", "examples": [ - "/usr/bin/python" + "/usr/bin/python3" ], "title": "Python Executable Path", "type": "string" @@ -1878,6 +1878,7 @@ { "enum": [ "/v1/embeddings", + "/v1/systemone", "/rerank", "/predict", "/predict_tokens" diff --git a/scripts/apigen/specs/management.json b/scripts/apigen/specs/management.json index 7e1635c..a7ea40e 100644 --- a/scripts/apigen/specs/management.json +++ b/scripts/apigen/specs/management.json @@ -6549,11 +6549,11 @@ "x-codeSamples": [ { "lang": "bash", - "source": "curl --request POST \\\n--url https://api.baseten.co/v1/loops/runs \\\n--header \"Authorization: Bearer $BASETEN_API_KEY\" \\\n--data '{\n \"session_id\": null,\n \"base_model\": null,\n \"name\": null,\n \"max_seq_len\": null,\n \"seed\": null,\n \"availability_model\": \"spot\",\n \"path\": \"bt://loops:k4q95w5/weights/step-100\",\n \"reuse_from_run_id\": null,\n \"reuse_from_session_id\": null\n}'" + "source": "curl --request POST \\\n--url https://api.baseten.co/v1/loops/runs \\\n--header \"Authorization: Bearer $BASETEN_API_KEY\" \\\n--data '{\n \"session_id\": null,\n \"base_model\": null,\n \"name\": null,\n \"max_seq_len\": null,\n \"seed\": null,\n \"availability_model\": \"spot\",\n \"path\": \"bt://loops:k4q95w5/weights/step-100\",\n \"reuse_from_run_id\": null,\n \"reuse_from_session_id\": null,\n \"sampler_num_replicas\": null\n}'" }, { "lang": "python", - "source": "import requests\nimport os\nAPI_KEY = os.environ.get(\"BASETEN_API_KEY\", \"\")\nurl = \"https://api.baseten.co/v1/loops/runs\"\n\nheaders = {\"Authorization\": f\"Bearer {API_KEY}\"}\n\nresponse = requests.request(\n \"POST\",\n url,\n headers=headers,\n json={'session_id': None, 'base_model': None, 'name': None, 'max_seq_len': None, 'seed': None, 'availability_model': 'spot', 'path': 'bt://loops:k4q95w5/weights/step-100', 'reuse_from_run_id': None, 'reuse_from_session_id': None}\n)\n\nprint(response.text)" + "source": "import requests\nimport os\nAPI_KEY = os.environ.get(\"BASETEN_API_KEY\", \"\")\nurl = \"https://api.baseten.co/v1/loops/runs\"\n\nheaders = {\"Authorization\": f\"Bearer {API_KEY}\"}\n\nresponse = requests.request(\n \"POST\",\n url,\n headers=headers,\n json={'session_id': None, 'base_model': None, 'name': None, 'max_seq_len': None, 'seed': None, 'availability_model': 'spot', 'path': 'bt://loops:k4q95w5/weights/step-100', 'reuse_from_run_id': None, 'reuse_from_session_id': None, 'sampler_num_replicas': None}\n)\n\nprint(response.text)" } ], "description": "Creates a Loops run with an associated sampler in the given session.", @@ -6665,7 +6665,7 @@ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/CreateLoopsRunRequestV1" + "$ref": "#/components/schemas/CreateLoopsTrainerRequestV1" } } }, @@ -6740,11 +6740,11 @@ "x-codeSamples": [ { "lang": "bash", - "source": "curl --request POST \\\n--url https://api.baseten.co/v1/loops/samplers \\\n--header \"Authorization: Bearer $BASETEN_API_KEY\" \\\n--data '{\n \"availability_model\": null,\n \"session_id\": null,\n \"base_model\": null,\n \"run_id\": null,\n \"max_seq_length\": null,\n \"model_path\": \"bt://loops:k4q95w5/sampler_weights/step-100\",\n \"reuse_from_session_id\": null\n}'" + "source": "curl --request POST \\\n--url https://api.baseten.co/v1/loops/samplers \\\n--header \"Authorization: Bearer $BASETEN_API_KEY\" \\\n--data '{\n \"availability_model\": null,\n \"session_id\": null,\n \"base_model\": null,\n \"run_id\": null,\n \"max_seq_length\": null,\n \"model_path\": \"bt://loops:k4q95w5/sampler_weights/step-100\",\n \"reuse_from_session_id\": null,\n \"num_replicas\": null\n}'" }, { "lang": "python", - "source": "import requests\nimport os\nAPI_KEY = os.environ.get(\"BASETEN_API_KEY\", \"\")\nurl = \"https://api.baseten.co/v1/loops/samplers\"\n\nheaders = {\"Authorization\": f\"Bearer {API_KEY}\"}\n\nresponse = requests.request(\n \"POST\",\n url,\n headers=headers,\n json={'availability_model': None, 'session_id': None, 'base_model': None, 'run_id': None, 'max_seq_length': None, 'model_path': 'bt://loops:k4q95w5/sampler_weights/step-100', 'reuse_from_session_id': None}\n)\n\nprint(response.text)" + "source": "import requests\nimport os\nAPI_KEY = os.environ.get(\"BASETEN_API_KEY\", \"\")\nurl = \"https://api.baseten.co/v1/loops/samplers\"\n\nheaders = {\"Authorization\": f\"Bearer {API_KEY}\"}\n\nresponse = requests.request(\n \"POST\",\n url,\n headers=headers,\n json={'availability_model': None, 'session_id': None, 'base_model': None, 'run_id': None, 'max_seq_length': None, 'model_path': 'bt://loops:k4q95w5/sampler_weights/step-100', 'reuse_from_session_id': None, 'num_replicas': None}\n)\n\nprint(response.text)" } ], "description": "Creates a standalone Loops sampler not linked to a run.", @@ -7585,11 +7585,11 @@ "x-codeSamples": [ { "lang": "bash", - "source": "curl --request POST \\\n--url https://api.baseten.co/v1/teams/{team_id}/loops/runs \\\n--header \"Authorization: Bearer $BASETEN_API_KEY\" \\\n--data '{\n \"session_id\": null,\n \"base_model\": null,\n \"name\": null,\n \"max_seq_len\": null,\n \"seed\": null,\n \"availability_model\": \"spot\",\n \"path\": \"bt://loops:k4q95w5/weights/step-100\",\n \"reuse_from_run_id\": null,\n \"reuse_from_session_id\": null\n}'" + "source": "curl --request POST \\\n--url https://api.baseten.co/v1/teams/{team_id}/loops/runs \\\n--header \"Authorization: Bearer $BASETEN_API_KEY\" \\\n--data '{\n \"session_id\": null,\n \"base_model\": null,\n \"name\": null,\n \"max_seq_len\": null,\n \"seed\": null,\n \"availability_model\": \"spot\",\n \"path\": \"bt://loops:k4q95w5/weights/step-100\",\n \"reuse_from_run_id\": null,\n \"reuse_from_session_id\": null,\n \"sampler_num_replicas\": null\n}'" }, { "lang": "python", - "source": "import requests\nimport os\nAPI_KEY = os.environ.get(\"BASETEN_API_KEY\", \"\")\nurl = \"https://api.baseten.co/v1/teams/{team_id}/loops/runs\"\n\nheaders = {\"Authorization\": f\"Bearer {API_KEY}\"}\n\nresponse = requests.request(\n \"POST\",\n url,\n headers=headers,\n json={'session_id': None, 'base_model': None, 'name': None, 'max_seq_len': None, 'seed': None, 'availability_model': 'spot', 'path': 'bt://loops:k4q95w5/weights/step-100', 'reuse_from_run_id': None, 'reuse_from_session_id': None}\n)\n\nprint(response.text)" + "source": "import requests\nimport os\nAPI_KEY = os.environ.get(\"BASETEN_API_KEY\", \"\")\nurl = \"https://api.baseten.co/v1/teams/{team_id}/loops/runs\"\n\nheaders = {\"Authorization\": f\"Bearer {API_KEY}\"}\n\nresponse = requests.request(\n \"POST\",\n url,\n headers=headers,\n json={'session_id': None, 'base_model': None, 'name': None, 'max_seq_len': None, 'seed': None, 'availability_model': 'spot', 'path': 'bt://loops:k4q95w5/weights/step-100', 'reuse_from_run_id': None, 'reuse_from_session_id': None, 'sampler_num_replicas': None}\n)\n\nprint(response.text)" } ], "description": "Creates a Loops run with an associated sampler in the given session; the run's infrastructure belongs to the given team.", @@ -7635,7 +7635,7 @@ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/CreateLoopsRunRequestV1" + "$ref": "#/components/schemas/CreateLoopsTrainerRequestV1" } } }, @@ -7720,11 +7720,11 @@ "x-codeSamples": [ { "lang": "bash", - "source": "curl --request POST \\\n--url https://api.baseten.co/v1/teams/{team_id}/loops/samplers \\\n--header \"Authorization: Bearer $BASETEN_API_KEY\" \\\n--data '{\n \"availability_model\": null,\n \"session_id\": null,\n \"base_model\": null,\n \"run_id\": null,\n \"max_seq_length\": null,\n \"model_path\": \"bt://loops:k4q95w5/sampler_weights/step-100\",\n \"reuse_from_session_id\": null\n}'" + "source": "curl --request POST \\\n--url https://api.baseten.co/v1/teams/{team_id}/loops/samplers \\\n--header \"Authorization: Bearer $BASETEN_API_KEY\" \\\n--data '{\n \"availability_model\": null,\n \"session_id\": null,\n \"base_model\": null,\n \"run_id\": null,\n \"max_seq_length\": null,\n \"model_path\": \"bt://loops:k4q95w5/sampler_weights/step-100\",\n \"reuse_from_session_id\": null,\n \"num_replicas\": null\n}'" }, { "lang": "python", - "source": "import requests\nimport os\nAPI_KEY = os.environ.get(\"BASETEN_API_KEY\", \"\")\nurl = \"https://api.baseten.co/v1/teams/{team_id}/loops/samplers\"\n\nheaders = {\"Authorization\": f\"Bearer {API_KEY}\"}\n\nresponse = requests.request(\n \"POST\",\n url,\n headers=headers,\n json={'availability_model': None, 'session_id': None, 'base_model': None, 'run_id': None, 'max_seq_length': None, 'model_path': 'bt://loops:k4q95w5/sampler_weights/step-100', 'reuse_from_session_id': None}\n)\n\nprint(response.text)" + "source": "import requests\nimport os\nAPI_KEY = os.environ.get(\"BASETEN_API_KEY\", \"\")\nurl = \"https://api.baseten.co/v1/teams/{team_id}/loops/samplers\"\n\nheaders = {\"Authorization\": f\"Bearer {API_KEY}\"}\n\nresponse = requests.request(\n \"POST\",\n url,\n headers=headers,\n json={'availability_model': None, 'session_id': None, 'base_model': None, 'run_id': None, 'max_seq_length': None, 'model_path': 'bt://loops:k4q95w5/sampler_weights/step-100', 'reuse_from_session_id': None, 'num_replicas': None}\n)\n\nprint(response.text)" } ], "description": "Creates a standalone Loops sampler not linked to a run; the sampler belongs to the given team.", @@ -9482,7 +9482,7 @@ "source": "import requests\nimport os\nAPI_KEY = os.environ.get(\"BASETEN_API_KEY\", \"\")\nurl = \"https://api.baseten.co/v1/routes\"\n\nheaders = {\"Authorization\": f\"Bearer {API_KEY}\"}\n\nresponse = requests.request(\n \"GET\",\n url,\n headers=headers,\n params={'cursor': None, 'team_id': None, 'name': None}\n)\n\nprint(response.text)" } ], - "description": "Lists routes you can invoke, newest first, optionally filtered by team and exact name.", + "description": "Lists routes you can invoke in the team, newest first, optionally filtered by exact name.", "parameters": [ { "name": "cursor", @@ -9519,7 +9519,7 @@ "name": "team_id", "in": "query", "required": false, - "description": "Filter by owning team ID. Preserved by the cursor; if repeated, must match the original filter.", + "description": "Identifier of the team whose routes to list.", "schema": { "anyOf": [ { @@ -9756,27 +9756,27 @@ "x-stability": "unstable" } }, - "/v1/routes/spend_limits/{user_id}": { + "/v1/routes/settings/users/{user_id}": { "get": { - "summary": "Gets a user's spend limits", + "summary": "Gets a user's route settings", "x-codeSamples": [ { "lang": "bash", - "source": "curl --request GET \\\n--url https://api.baseten.co/v1/routes/spend_limits/{user_id} \\\n--header \"Authorization: Bearer $BASETEN_API_KEY\"\n" + "source": "curl --request GET \\\n--url https://api.baseten.co/v1/routes/settings/users/{user_id} \\\n--header \"Authorization: Bearer $BASETEN_API_KEY\"\n" }, { "lang": "python", - "source": "import requests\nimport os\nAPI_KEY = os.environ.get(\"BASETEN_API_KEY\", \"\")\nurl = \"https://api.baseten.co/v1/routes/spend_limits/{user_id}\"\n\nheaders = {\"Authorization\": f\"Bearer {API_KEY}\"}\n\nresponse = requests.request(\n \"GET\",\n url,\n headers=headers,\n json={}\n)\n\nprint(response.text)" + "source": "import requests\nimport os\nAPI_KEY = os.environ.get(\"BASETEN_API_KEY\", \"\")\nurl = \"https://api.baseten.co/v1/routes/settings/users/{user_id}\"\n\nheaders = {\"Authorization\": f\"Bearer {API_KEY}\"}\n\nresponse = requests.request(\n \"GET\",\n url,\n headers=headers,\n json={}\n)\n\nprint(response.text)" } ], - "description": "Spend limits cap what a Baseten Code user can spend each month. Personal API keys are not limited. Users can read their own spend limits; only organization admins can read other users' spend limits or change them.", + "description": "Returns the settings that apply to the user, including the spend limit that caps what they can spend on Baseten Code each month: the limit set on the user, their team's per-member limit, and the effective limit for this user. The user limit always overrides the team's per-member limit.", "responses": { "200": { "description": "", "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/RouteSpendLimitV1" + "$ref": "#/components/schemas/RouteUserSettingsV1" } } } @@ -9790,23 +9790,23 @@ } ], "patch": { - "summary": "Updates a user's spend limits", + "summary": "Updates a user's route settings", "x-codeSamples": [ { "lang": "bash", - "source": "curl --request PATCH \\\n--url https://api.baseten.co/v1/routes/spend_limits/{user_id} \\\n--header \"Authorization: Bearer $BASETEN_API_KEY\" \\\n--data '{\n \"monthly_limit_usd\": \"200\"\n}'" + "source": "curl --request PATCH \\\n--url https://api.baseten.co/v1/routes/settings/users/{user_id} \\\n--header \"Authorization: Bearer $BASETEN_API_KEY\" \\\n--data '{\n \"spend_limit\": null\n}'" }, { "lang": "python", - "source": "import requests\nimport os\nAPI_KEY = os.environ.get(\"BASETEN_API_KEY\", \"\")\nurl = \"https://api.baseten.co/v1/routes/spend_limits/{user_id}\"\n\nheaders = {\"Authorization\": f\"Bearer {API_KEY}\"}\n\nresponse = requests.request(\n \"PATCH\",\n url,\n headers=headers,\n json={'monthly_limit_usd': '200'}\n)\n\nprint(response.text)" + "source": "import requests\nimport os\nAPI_KEY = os.environ.get(\"BASETEN_API_KEY\", \"\")\nurl = \"https://api.baseten.co/v1/routes/settings/users/{user_id}\"\n\nheaders = {\"Authorization\": f\"Bearer {API_KEY}\"}\n\nresponse = requests.request(\n \"PATCH\",\n url,\n headers=headers,\n json={'spend_limit': None}\n)\n\nprint(response.text)" } ], - "description": "Once the user's metered spend in a month reaches the limit, requests with Routes keys they created are rejected until the limit is raised or the next month starts. A limit of 0 blocks the user after their first metered request. Spend is metered every 15 minutes and can lag, so a user can go over the limit. If the user has spend this month, changes take effect within seconds; otherwise at the next metering run.", + "description": "Changes only the fields in the request. Once the user's metered spend in a month reaches the spend limit, requests with Routes keys they created are rejected until the limit is raised or the next month starts. Spend is metered every 15 minutes and can lag, so a user can go over the limit. The user's own limit supersedes their team's per-member limit.", "requestBody": { "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/UpdateRouteSpendLimitRequestV1" + "$ref": "#/components/schemas/UpdateRouteUserSettingsRequestV1" } } }, @@ -9818,7 +9818,7 @@ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/RouteSpendLimitV1" + "$ref": "#/components/schemas/RouteUserSettingsV1" } } } @@ -9827,40 +9827,27 @@ "x-stability": "unstable" } }, - "/v1/routes/harness-configs": { + "/v1/routes/settings/teams/{team_id}": { "get": { - "summary": "Lists default models for coding harnesses", + "summary": "Gets a team's route settings", "x-codeSamples": [ { "lang": "bash", - "source": "curl --request GET \\\n--url https://api.baseten.co/v1/routes/harness-configs \\\n--header \"Authorization: Bearer $BASETEN_API_KEY\"\n" + "source": "curl --request GET \\\n--url https://api.baseten.co/v1/routes/settings/teams/{team_id} \\\n--header \"Authorization: Bearer $BASETEN_API_KEY\"\n" }, { "lang": "python", - "source": "import requests\nimport os\nAPI_KEY = os.environ.get(\"BASETEN_API_KEY\", \"\")\nurl = \"https://api.baseten.co/v1/routes/harness-configs\"\n\nheaders = {\"Authorization\": f\"Bearer {API_KEY}\"}\n\nresponse = requests.request(\n \"GET\",\n url,\n headers=headers,\n params={'team_id': None}\n)\n\nprint(response.text)" - } - ], - "description": "Each role uses the route a team admin set for it (`source` is `team`). A role with no route set uses Baseten's default (`source` is `baseten`), chosen from the team's Model API routes: the recommended model for `primary` and the lowest-priced model for `background`. Baseten's defaults never use external-provider routes.", - "parameters": [ - { - "name": "team_id", - "in": "query", - "required": true, - "description": "Identifier of the team whose default models to list.", - "schema": { - "minLength": 1, - "title": "Team Id", - "type": "string" - } + "source": "import requests\nimport os\nAPI_KEY = os.environ.get(\"BASETEN_API_KEY\", \"\")\nurl = \"https://api.baseten.co/v1/routes/settings/teams/{team_id}\"\n\nheaders = {\"Authorization\": f\"Bearer {API_KEY}\"}\n\nresponse = requests.request(\n \"GET\",\n url,\n headers=headers,\n json={}\n)\n\nprint(response.text)" } ], + "description": "Returns the settings that apply to the team, including the default models for each coding harness and the team's spend limits. Any member who can invoke the team's routes can read them.", "responses": { "200": { "description": "", "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/RouteHarnessConfigsResponseV1" + "$ref": "#/components/schemas/RouteTeamSettingsV1" } } } @@ -9868,24 +9855,29 @@ }, "x-stability": "unstable" }, + "parameters": [ + { + "$ref": "#/components/parameters/team_id" + } + ], "patch": { - "summary": "Updates default models for a coding harness", + "summary": "Updates a team's route settings", "x-codeSamples": [ { "lang": "bash", - "source": "curl --request PATCH \\\n--url https://api.baseten.co/v1/routes/harness-configs \\\n--header \"Authorization: Bearer $BASETEN_API_KEY\" \\\n--data '{\n \"team_id\": \"abc1234\",\n \"harness\": \"claude-code\",\n \"models\": {\n \"background\": \"ghi9012\"\n }\n}'" + "source": "curl --request PATCH \\\n--url https://api.baseten.co/v1/routes/settings/teams/{team_id} \\\n--header \"Authorization: Bearer $BASETEN_API_KEY\" \\\n--data '{\n \"harness_defaults\": null,\n \"spend_limit\": null\n}'" }, { "lang": "python", - "source": "import requests\nimport os\nAPI_KEY = os.environ.get(\"BASETEN_API_KEY\", \"\")\nurl = \"https://api.baseten.co/v1/routes/harness-configs\"\n\nheaders = {\"Authorization\": f\"Bearer {API_KEY}\"}\n\nresponse = requests.request(\n \"PATCH\",\n url,\n headers=headers,\n json={'team_id': 'abc1234', 'harness': 'claude-code', 'models': {'background': 'ghi9012'}}\n)\n\nprint(response.text)" + "source": "import requests\nimport os\nAPI_KEY = os.environ.get(\"BASETEN_API_KEY\", \"\")\nurl = \"https://api.baseten.co/v1/routes/settings/teams/{team_id}\"\n\nheaders = {\"Authorization\": f\"Bearer {API_KEY}\"}\n\nresponse = requests.request(\n \"PATCH\",\n url,\n headers=headers,\n json={'harness_defaults': None, 'spend_limit': None}\n)\n\nprint(response.text)" } ], - "description": "Changes only the model roles in the request and leaves the others unchanged. Set a role to null to use Baseten's default. Requires team admin, and every route must belong to the team.", + "description": "Changes only the fields in the request, all or nothing. Requires team admin.", "requestBody": { "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/UpdateRouteHarnessConfigRequestV1" + "$ref": "#/components/schemas/UpdateRouteTeamSettingsRequestV1" } } }, @@ -9897,7 +9889,87 @@ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/RouteHarnessConfigV1" + "$ref": "#/components/schemas/RouteTeamSettingsV1" + } + } + } + } + }, + "x-stability": "unstable" + } + }, + "/v1/routes/connections": { + "get": { + "summary": "Lists connections", + "x-codeSamples": [ + { + "lang": "bash", + "source": "curl --request GET \\\n--url https://api.baseten.co/v1/routes/connections \\\n--header \"Authorization: Bearer $BASETEN_API_KEY\"\n" + }, + { + "lang": "python", + "source": "import requests\nimport os\nAPI_KEY = os.environ.get(\"BASETEN_API_KEY\", \"\")\nurl = \"https://api.baseten.co/v1/routes/connections\"\n\nheaders = {\"Authorization\": f\"Bearer {API_KEY}\"}\n\nresponse = requests.request(\n \"GET\",\n url,\n headers=headers,\n params={'cursor': None, 'team_id': None}\n)\n\nprint(response.text)" + } + ], + "description": "Lists the team's connections, newest first.", + "parameters": [ + { + "name": "cursor", + "in": "query", + "required": false, + "description": "Opaque cursor returned by a previous page. Omit to fetch the first page.", + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Cursor" + } + }, + { + "name": "limit", + "in": "query", + "required": false, + "description": "Maximum number of items to return.", + "schema": { + "default": 100, + "maximum": 1000, + "minimum": 1, + "title": "Limit", + "type": "integer" + } + }, + { + "name": "team_id", + "in": "query", + "required": false, + "description": "Identifier of the team whose connections to list.", + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Team Id" + } + } + ], + "responses": { + "200": { + "description": "", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/RouteConnectionsResponseV1" } } } @@ -9905,24 +9977,24 @@ }, "x-stability": "unstable" }, - "put": { - "summary": "Sets default models for a coding harness", + "post": { + "summary": "Creates a connection", "x-codeSamples": [ { "lang": "bash", - "source": "curl --request PUT \\\n--url https://api.baseten.co/v1/routes/harness-configs \\\n--header \"Authorization: Bearer $BASETEN_API_KEY\" \\\n--data '{\n \"team_id\": \"abc1234\",\n \"harness\": \"claude-code\",\n \"models\": {\n \"primary\": \"def5678\",\n \"background\": \"ghi9012\"\n }\n}'" + "source": "curl --request POST \\\n--url https://api.baseten.co/v1/routes/connections \\\n--header \"Authorization: Bearer $BASETEN_API_KEY\" \\\n--data '{\n \"team_id\": \"abc1234\",\n \"config\": {\n \"provider\": \"ANTHROPIC\",\n \"secret_id\": \"abc1234\"\n }\n}'" }, { "lang": "python", - "source": "import requests\nimport os\nAPI_KEY = os.environ.get(\"BASETEN_API_KEY\", \"\")\nurl = \"https://api.baseten.co/v1/routes/harness-configs\"\n\nheaders = {\"Authorization\": f\"Bearer {API_KEY}\"}\n\nresponse = requests.request(\n \"PUT\",\n url,\n headers=headers,\n json={'team_id': 'abc1234', 'harness': 'claude-code', 'models': {'primary': 'def5678', 'background': 'ghi9012'}}\n)\n\nprint(response.text)" + "source": "import requests\nimport os\nAPI_KEY = os.environ.get(\"BASETEN_API_KEY\", \"\")\nurl = \"https://api.baseten.co/v1/routes/connections\"\n\nheaders = {\"Authorization\": f\"Bearer {API_KEY}\"}\n\nresponse = requests.request(\n \"POST\",\n url,\n headers=headers,\n json={'team_id': 'abc1234', 'config': {'provider': 'ANTHROPIC', 'secret_id': 'abc1234'}}\n)\n\nprint(response.text)" } ], - "description": "Replaces every model role of one harness. Roles left out use Baseten's defaults. Requires team admin, and every route must belong to the team. The `harness` in the request body determines which roles are accepted.", + "description": "Creates a connection from the team to an upstream provider.", "requestBody": { "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/SetRouteHarnessConfigRequestV1" + "$ref": "#/components/schemas/CreateRouteConnectionRequestV1" } } }, @@ -9934,7 +10006,7 @@ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/RouteHarnessConfigV1" + "$ref": "#/components/schemas/RouteConnectionV1" } } } @@ -9943,40 +10015,27 @@ "x-stability": "unstable" } }, - "/v1/routes/harness-configs/{harness}": { + "/v1/routes/connections/{connection_id}": { "delete": { - "summary": "Clears default models for a coding harness", + "summary": "Deletes a connection", "x-codeSamples": [ { "lang": "bash", - "source": "curl --request DELETE \\\n--url https://api.baseten.co/v1/routes/harness-configs/{harness} \\\n--header \"Authorization: Bearer $BASETEN_API_KEY\"\n" + "source": "curl --request DELETE \\\n--url https://api.baseten.co/v1/routes/connections/{connection_id} \\\n--header \"Authorization: Bearer $BASETEN_API_KEY\"\n" }, { "lang": "python", - "source": "import requests\nimport os\nAPI_KEY = os.environ.get(\"BASETEN_API_KEY\", \"\")\nurl = \"https://api.baseten.co/v1/routes/harness-configs/{harness}\"\n\nheaders = {\"Authorization\": f\"Bearer {API_KEY}\"}\n\nresponse = requests.request(\n \"DELETE\",\n url,\n headers=headers,\n params={'team_id': None}\n)\n\nprint(response.text)" - } - ], - "description": "Removes the team's routes for every role of the harness, so each role uses Baseten's default. Requires team admin.", - "parameters": [ - { - "name": "team_id", - "in": "query", - "required": true, - "description": "Identifier of the team whose default models to clear.", - "schema": { - "minLength": 1, - "title": "Team Id", - "type": "string" - } + "source": "import requests\nimport os\nAPI_KEY = os.environ.get(\"BASETEN_API_KEY\", \"\")\nurl = \"https://api.baseten.co/v1/routes/connections/{connection_id}\"\n\nheaders = {\"Authorization\": f\"Bearer {API_KEY}\"}\n\nresponse = requests.request(\n \"DELETE\",\n url,\n headers=headers,\n json={}\n)\n\nprint(response.text)" } ], + "description": "Deletes the connection. The team secret it points at is left in place.", "responses": { "200": { "description": "", "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/RouteHarnessConfigTombstoneV1" + "$ref": "#/components/schemas/RouteConnectionTombstoneV1" } } } @@ -9986,9 +10045,46 @@ }, "parameters": [ { - "$ref": "#/components/parameters/harness" + "$ref": "#/components/parameters/connection_id" } - ] + ], + "patch": { + "summary": "Updates a connection", + "x-codeSamples": [ + { + "lang": "bash", + "source": "curl --request PATCH \\\n--url https://api.baseten.co/v1/routes/connections/{connection_id} \\\n--header \"Authorization: Bearer $BASETEN_API_KEY\" \\\n--data '{\n \"config\": {\n \"provider\": \"ANTHROPIC\",\n \"secret_id\": \"abc1234\"\n }\n}'" + }, + { + "lang": "python", + "source": "import requests\nimport os\nAPI_KEY = os.environ.get(\"BASETEN_API_KEY\", \"\")\nurl = \"https://api.baseten.co/v1/routes/connections/{connection_id}\"\n\nheaders = {\"Authorization\": f\"Bearer {API_KEY}\"}\n\nresponse = requests.request(\n \"PATCH\",\n url,\n headers=headers,\n json={'config': {'provider': 'ANTHROPIC', 'secret_id': 'abc1234'}}\n)\n\nprint(response.text)" + } + ], + "description": "Updates the connection's fields. The provider and owning team are immutable; point the connection at a new team secret to change its API key, or rotate the key itself by updating the secret.", + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/UpdateRouteConnectionRequestV1" + } + } + }, + "required": true + }, + "responses": { + "200": { + "description": "", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/RouteConnectionV1" + } + } + } + } + }, + "x-stability": "unstable" + } }, "/v1/routes/{route_id}": { "delete": { @@ -10049,18 +10145,18 @@ "x-stability": "unstable" }, "patch": { - "summary": "Updates a route's display name or description", + "summary": "Updates a route", "x-codeSamples": [ { "lang": "bash", - "source": "curl --request PATCH \\\n--url https://api.baseten.co/v1/routes/{route_id} \\\n--header \"Authorization: Bearer $BASETEN_API_KEY\" \\\n--data '{\n \"description\": \"Assistant for code review and debugging.\",\n \"display_name\": \"Assistant\"\n}'" + "source": "curl --request PATCH \\\n--url https://api.baseten.co/v1/routes/{route_id} \\\n--header \"Authorization: Bearer $BASETEN_API_KEY\" \\\n--data '{\n \"description\": \"Assistant for code review and debugging.\",\n \"display_name\": \"Assistant\",\n \"target\": null\n}'" }, { "lang": "python", - "source": "import requests\nimport os\nAPI_KEY = os.environ.get(\"BASETEN_API_KEY\", \"\")\nurl = \"https://api.baseten.co/v1/routes/{route_id}\"\n\nheaders = {\"Authorization\": f\"Bearer {API_KEY}\"}\n\nresponse = requests.request(\n \"PATCH\",\n url,\n headers=headers,\n json={'description': 'Assistant for code review and debugging.', 'display_name': 'Assistant'}\n)\n\nprint(response.text)" + "source": "import requests\nimport os\nAPI_KEY = os.environ.get(\"BASETEN_API_KEY\", \"\")\nurl = \"https://api.baseten.co/v1/routes/{route_id}\"\n\nheaders = {\"Authorization\": f\"Bearer {API_KEY}\"}\n\nresponse = requests.request(\n \"PATCH\",\n url,\n headers=headers,\n json={'description': 'Assistant for code review and debugging.', 'display_name': 'Assistant', 'target': None}\n)\n\nprint(response.text)" } ], - "description": "The route name, target, and owning team are immutable.", + "description": "The route name, target type, and owning team are immutable.", "requestBody": { "content": { "application/json": { @@ -10519,11 +10615,60 @@ } ] }, - "/v1/sandboxes/instances": { + "/v1/sandboxes/configuration": { "get": { - "operationId": "ListSandboxes", - "summary": "List sandboxes", - "description": "Return a cursor-paginated page of sandboxes. Terminated sandboxes are hidden by default.", + "operationId": "GetSandboxConfiguration", + "summary": "Get sandbox configuration", + "description": "Return available Carbon-compatible sandbox regions. Configuration is the same for every team, so no team is selected. Only enabled regions supporting Carbon are returned. This describes regional capability, not the runtime selected for a particular team or sandbox. Region configuration is cached for up to 30 seconds.", + "tags": [ + "Sandboxes" + ], + "responses": { + "200": { + "description": "Successful operation.", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/GetSandboxConfigurationResponseV1" + }, + "example": { + "regions": [ + { + "name": "us-was-1", + "country": "us", + "continent": "na", + "location": "Washington", + "info_generation": "CARBON" + } + ] + } + } + } + }, + "400": { + "$ref": "#/components/responses/SandboxError400" + }, + "401": { + "$ref": "#/components/responses/SandboxError401" + }, + "403": { + "$ref": "#/components/responses/SandboxError403" + }, + "429": { + "$ref": "#/components/responses/SandboxError429" + }, + "500": { + "$ref": "#/components/responses/SandboxError500" + } + }, + "x-stability": "unstable" + } + }, + "/v1/sandboxes/instances/{sandbox_name}/metrics": { + "get": { + "operationId": "GetSandboxMetrics", + "summary": "Get sandbox metrics", + "description": "Returns every interval in the requested half-open range [start_time, end_time).\nBuckets align to UTC interval boundaries; partial boundary buckets include only samples\nwithin the requested range. The first bucket may start before start_time. Ranges are limited to 31 days\nand 10000 buckets. Request counts are zero when there is no traffic, and null\nfor buckets entirely before sandbox creation. CPU and memory are null without\nsamples, including standby. Error rate is the fraction of requests returning\n4xx or 5xx, null when there are no requests. Telemetry failures return an error.\n", "tags": [ "Sandboxes" ], @@ -10535,108 +10680,286 @@ "$ref": "#/components/parameters/TeamIdHeader" }, { - "$ref": "#/components/parameters/Cursor" - }, - { - "$ref": "#/components/parameters/Limit" + "$ref": "#/components/parameters/SandboxName" }, { - "name": "q", + "name": "start_time", "in": "query", - "description": "Search indexed sandbox names and labels. Search is applied before pagination.", - "required": false, + "required": true, "schema": { "type": "string", - "example": "api-review" + "format": "date-time" } }, { - "name": "status", + "name": "end_time", "in": "query", - "description": "Deployment statuses. Repeat the query parameter for each status, for example status=DEPLOYED&status=FAILED. Unknown values are rejected. Cannot be combined with external_id.", - "required": false, - "style": "form", - "explode": true, + "required": true, "schema": { - "type": "array", - "minItems": 1, - "items": { - "type": "string", - "minLength": 1 - }, - "example": [ - "DEPLOYED", - "FAILED" - ] + "type": "string", + "format": "date-time" } }, { - "name": "external_id", + "name": "interval_seconds", "in": "query", - "description": "Filter by a caller-owned external identifier. Cannot be combined with status.", - "required": false, + "required": true, "schema": { - "type": "string", - "example": "api-review-20260916-001" + "type": "integer", + "enum": [ + 10, + 30, + 60, + 300, + 900, + 3600 + ] } } ], "responses": { "200": { - "description": "Successful operation.", + "description": "Complete sandbox metrics time series.", "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/ListSandboxesResponseV1" - }, - "examples": { - "page": { - "$ref": "#/components/examples/SandboxesPage" - }, - "last_page": { - "$ref": "#/components/examples/SandboxesLastPage" - } + "$ref": "#/components/schemas/SandboxMetricsV1" } } } }, "400": { - "$ref": "#/components/responses/SandboxError400" + "description": "Invalid range, interval, or query parameter." }, "401": { - "$ref": "#/components/responses/SandboxError401" + "description": "Authentication required." }, "403": { - "$ref": "#/components/responses/SandboxError403" + "description": "Team access denied." }, - "429": { - "$ref": "#/components/responses/SandboxError429" + "404": { + "description": "Sandbox not found." }, - "500": { - "$ref": "#/components/responses/SandboxError500" + "502": { + "description": "Metrics query failed." + }, + "503": { + "description": "Metrics are unavailable." } }, "x-stability": "unstable" - }, - "post": { + } + }, + "/v1/sandboxes/instances/{sandbox_name}/logs": { + "get": { + "operationId": "GetSandboxLogs", + "summary": "Get sandbox logs", + "description": "Returns sandbox runtime logs, newest first, within the half-open range\n[start_time, end_time). Ranges are limited to 7 days. Follow next_cursor\nuntil it is null to retrieve all pages, keeping the same range and limit.\nLogs before sandbox creation are excluded. Unavailable telemetry returns an error.\n", + "tags": [ + "Sandboxes" + ], "parameters": [ { "$ref": "#/components/parameters/TeamId" }, { "$ref": "#/components/parameters/TeamIdHeader" - } - ], - "operationId": "CreateSandbox", - "summary": "Create a sandbox", - "description": "Create a sandbox with an optional client-provided name, generated by the server when omitted. The request waits briefly for deployment, then returns the created resource. Deployment may continue after the response; inspect status to determine readiness. Names belonging to failed or terminated sandboxes can be reused; an active sandbox reserves its name. Every field is optional. The server applies image and memory defaults; SDKs should omit fields the caller did not set. A request with no body creates a sandbox from baseten/base-image:latest, the built-in image with the sandbox execution API ready to use, with 4096 MB of memory. It is available directly; no image build, push, or GET /v1/sandboxes/images is needed. Use the returned url to execute commands with your authentication token as a Bearer token, not your Baseten API key.", - "tags": [ - "Sandboxes" - ], - "requestBody": { - "required": false, - "content": { - "application/json": { + }, + { + "$ref": "#/components/parameters/SandboxName" + }, + { + "name": "start_time", + "in": "query", + "required": true, + "schema": { + "type": "string", + "format": "date-time" + } + }, + { + "name": "end_time", + "in": "query", + "required": true, + "schema": { + "type": "string", + "format": "date-time" + } + }, + { + "name": "limit", + "in": "query", + "required": false, + "schema": { + "type": "integer", + "minimum": 1, + "maximum": 1000, + "default": 100 + } + }, + { + "name": "cursor", + "in": "query", + "required": false, + "description": "Opaque next_cursor from the previous page. Omit for the first page.", + "schema": { + "type": "string", + "minLength": 1, + "maxLength": 4096 + } + } + ], + "responses": { + "200": { + "description": "A page of sandbox logs.", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/SandboxLogsV1" + } + } + } + }, + "400": { + "description": "Invalid range, limit, cursor, or query parameter." + }, + "401": { + "description": "Authentication required." + }, + "403": { + "description": "Team access denied." + }, + "404": { + "description": "Sandbox not found." + }, + "502": { + "description": "Logs query failed." + }, + "503": { + "description": "Logs are unavailable." + } + }, + "x-stability": "unstable" + } + }, + "/v1/sandboxes/instances": { + "get": { + "operationId": "ListSandboxes", + "summary": "List sandboxes", + "description": "Return a cursor-paginated page of sandboxes. Terminated sandboxes are hidden by default.", + "tags": [ + "Sandboxes" + ], + "parameters": [ + { + "$ref": "#/components/parameters/TeamId" + }, + { + "$ref": "#/components/parameters/TeamIdHeader" + }, + { + "$ref": "#/components/parameters/Cursor" + }, + { + "$ref": "#/components/parameters/Limit" + }, + { + "name": "q", + "in": "query", + "description": "Search indexed sandbox names and labels. Search is applied before pagination.", + "required": false, + "schema": { + "type": "string", + "example": "api-review" + } + }, + { + "name": "status", + "in": "query", + "description": "Deployment statuses. Repeat the query parameter for each status, for example status=DEPLOYED&status=FAILED. Unknown values are rejected. Cannot be combined with external_id.", + "required": false, + "style": "form", + "explode": true, + "schema": { + "type": "array", + "minItems": 1, + "items": { + "type": "string", + "minLength": 1 + }, + "example": [ + "DEPLOYED", + "FAILED" + ] + } + }, + { + "name": "external_id", + "in": "query", + "description": "Filter by a caller-owned external identifier. Cannot be combined with status.", + "required": false, + "schema": { + "type": "string", + "example": "api-review-20260916-001" + } + } + ], + "responses": { + "200": { + "description": "Successful operation.", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ListSandboxesResponseV1" + }, + "examples": { + "page": { + "$ref": "#/components/examples/SandboxesPage" + }, + "last_page": { + "$ref": "#/components/examples/SandboxesLastPage" + } + } + } + } + }, + "400": { + "$ref": "#/components/responses/SandboxError400" + }, + "401": { + "$ref": "#/components/responses/SandboxError401" + }, + "403": { + "$ref": "#/components/responses/SandboxError403" + }, + "429": { + "$ref": "#/components/responses/SandboxError429" + }, + "500": { + "$ref": "#/components/responses/SandboxError500" + } + }, + "x-stability": "unstable" + }, + "post": { + "parameters": [ + { + "$ref": "#/components/parameters/TeamId" + }, + { + "$ref": "#/components/parameters/TeamIdHeader" + } + ], + "operationId": "CreateSandbox", + "summary": "Create a sandbox", + "description": "Create a sandbox with an optional client-provided name, generated by the server when omitted. The request waits briefly for deployment, then returns the created resource. Deployment may continue after the response; inspect status to determine readiness. Names belonging to failed or terminated sandboxes can be reused; an active sandbox reserves its name. Every field is optional. The server applies image and memory defaults; SDKs should omit fields the caller did not set. A request with no body creates a sandbox from baseten/base-image:latest, the built-in image with the sandbox execution API ready to use, with 4096 MB of memory. It is available directly; no image build, push, or GET /v1/sandboxes/images is needed. Use the returned url to execute commands with your authentication token as a Bearer token, not your Baseten API key.", + "tags": [ + "Sandboxes" + ], + "requestBody": { + "required": false, + "content": { + "application/json": { "schema": { "$ref": "#/components/schemas/CreateSandboxRequestV1" }, @@ -10775,7 +11098,7 @@ ], "operationId": "UpdateSandbox", "summary": "Update a sandbox", - "description": "Partially update configuration. Omitted fields remain unchanged; supplied arrays and maps replace their previous values. Structured objects update only supplied fields. Changes to image may reset running state. The name, memory, and network configuration are immutable after creation. Supplying memory or network returns 400, including unchanged, empty, or null values.", + "description": "Partially update configuration. Omitted fields remain unchanged; supplied arrays and maps replace their previous values. Structured objects update only supplied fields. The name, memory, network, region, image, and ports are immutable after creation. Supplying any of these fields returns 400, including unchanged, empty, or null values.", "tags": [ "Sandboxes" ], @@ -10899,18 +11222,10 @@ "get": { "operationId": "ListSandboxLibraryImages", "summary": "List built-in sandbox images", - "description": "Return the built-in sandbox images available to every team. Use the image field directly as the image of a sandbox, without building or pushing it. Hidden and upcoming images are excluded.", + "description": "Return the built-in sandbox images, the same for every team, so no team is selected. Use the image field directly as the image of a sandbox, without building or pushing it. Hidden and upcoming images are excluded.", "tags": [ "Images" ], - "parameters": [ - { - "$ref": "#/components/parameters/TeamId" - }, - { - "$ref": "#/components/parameters/TeamIdHeader" - } - ], "responses": { "200": { "description": "Successful operation.", @@ -10988,7 +11303,7 @@ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/ListImagesResponseV1" + "$ref": "#/components/schemas/ListSandboxImagesResponseV1" }, "examples": { "complete": { @@ -11036,7 +11351,7 @@ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/PushImageRequestV1" + "$ref": "#/components/schemas/PushSandboxImageRequestV1" }, "examples": { "archive": { @@ -11055,7 +11370,7 @@ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/PushImageResponseV1" + "$ref": "#/components/schemas/PushSandboxImageResponseV1" }, "examples": { "archive": { @@ -11109,7 +11424,7 @@ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/CleanupImagesResponseV1" + "$ref": "#/components/schemas/CleanupSandboxImagesResponseV1" }, "examples": { "complete": { @@ -11141,7 +11456,7 @@ "/v1/sandboxes/images/{image_name}": { "parameters": [ { - "$ref": "#/components/parameters/ImageName" + "$ref": "#/components/parameters/SandboxImageName" } ], "get": { @@ -11165,7 +11480,7 @@ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/ImageV1" + "$ref": "#/components/schemas/SandboxImageV1" }, "examples": { "complete": { @@ -11214,7 +11529,7 @@ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/ImageV1" + "$ref": "#/components/schemas/SandboxImageV1" }, "examples": { "complete": { @@ -11249,7 +11564,7 @@ "/v1/sandboxes/images/{image_name}/logs": { "parameters": [ { - "$ref": "#/components/parameters/ImageName" + "$ref": "#/components/parameters/SandboxImageName" } ], "get": { @@ -11313,7 +11628,7 @@ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/ImageBuildLogsResponseV1" + "$ref": "#/components/schemas/SandboxImageBuildLogsResponseV1" }, "example": { "logs": [ @@ -11356,7 +11671,7 @@ "/v1/sandboxes/images/{image_name}/tags": { "parameters": [ { - "$ref": "#/components/parameters/ImageName" + "$ref": "#/components/parameters/SandboxImageName" } ], "get": { @@ -11411,7 +11726,7 @@ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/ListImageTagsResponseV1" + "$ref": "#/components/schemas/ListSandboxImageTagsResponseV1" } } } @@ -11441,7 +11756,7 @@ "/v1/sandboxes/images/{image_name}/tags/{tag_name}": { "parameters": [ { - "$ref": "#/components/parameters/ImageName" + "$ref": "#/components/parameters/SandboxImageName" }, { "$ref": "#/components/parameters/TagName" @@ -11468,7 +11783,7 @@ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/ImageV1" + "$ref": "#/components/schemas/SandboxImageV1" }, "examples": { "complete": { @@ -11945,7 +12260,6 @@ "type": "object" }, "VolumeSyncAuthenticationAWSAssumeRoleV1": { - "additionalProperties": false, "description": "Authentication using an AWS IAM role assumed by Baseten.", "properties": { "role_arn": { @@ -11969,7 +12283,6 @@ "type": "object" }, "VolumeSyncAuthenticationAWSOIDCV1": { - "additionalProperties": false, "description": "Authentication using an AWS IAM role and Baseten workload identity.", "properties": { "role_arn": { @@ -11993,7 +12306,6 @@ "type": "object" }, "VolumeSyncAuthenticationGCPOIDCV1": { - "additionalProperties": false, "description": "Authentication using GCP Workload Identity Federation.", "properties": { "service_account": { @@ -12017,7 +12329,6 @@ "type": "object" }, "VolumeSyncDestinationV1": { - "additionalProperties": false, "description": "BDN destination for a volume sync.", "properties": { "ref": { @@ -12054,7 +12365,6 @@ "type": "object" }, "VolumeSyncSourceAzureV1": { - "additionalProperties": false, "description": "Azure Blob Storage source.", "properties": { "uri": { @@ -12107,7 +12417,6 @@ "type": "object" }, "VolumeSyncSourceBasetenTrainingV1": { - "additionalProperties": false, "description": "Baseten training artifact source.", "properties": { "uri": { @@ -12146,7 +12455,6 @@ "type": "object" }, "VolumeSyncSourceCoreWeaveV1": { - "additionalProperties": false, "description": "CoreWeave object storage source.", "properties": { "uri": { @@ -12199,7 +12507,6 @@ "type": "object" }, "VolumeSyncSourceGCSV1": { - "additionalProperties": false, "description": "Google Cloud Storage source.", "properties": { "uri": { @@ -12264,7 +12571,6 @@ "type": "object" }, "VolumeSyncSourceHuggingFaceV1": { - "additionalProperties": false, "description": "Hugging Face source.", "properties": { "uri": { @@ -12317,7 +12623,6 @@ "type": "object" }, "VolumeSyncSourceR2V1": { - "additionalProperties": false, "description": "Cloudflare R2 source.", "properties": { "uri": { @@ -12370,7 +12675,6 @@ "type": "object" }, "VolumeSyncSourceS3V1": { - "additionalProperties": false, "description": "Amazon S3 source.", "properties": { "uri": { @@ -14552,7 +14856,6 @@ "type": "object" }, "AuditLogEventApiKeyCreatedV1": { - "additionalProperties": false, "description": "An API key was created.", "properties": { "event_type": { @@ -14582,7 +14885,6 @@ "type": "object" }, "AuditLogEventApiKeyDeletedV1": { - "additionalProperties": false, "description": "An API key was revoked.", "properties": { "event_type": { @@ -14623,7 +14925,6 @@ "type": "string" }, "AuditLogEventAutoscalingScheduleChangeV1": { - "additionalProperties": false, "description": "What an autoscaling change did to one schedule, and the schedule on either side of it.\n`previous` is null on a create, `current` on a delete, and an unchanged schedule carries only\n`current`. Not itself a payload in the discriminated union.", "properties": { "action": { @@ -14664,7 +14965,6 @@ "type": "object" }, "AuditLogEventAutoscalingScheduleSettingsV1": { - "additionalProperties": false, "description": "One autoscaling schedule's timing and autoscaling settings. Not itself a union payload.", "properties": { "min_replica": { @@ -14863,7 +15163,6 @@ "type": "object" }, "AuditLogEventAutoscalingSettingsV1": { - "additionalProperties": false, "description": "Autoscaling settings for a deployment or environment.", "properties": { "min_replica": { @@ -14948,7 +15247,6 @@ "type": "object" }, "AuditLogEventChainDeletedV1": { - "additionalProperties": false, "description": "A chain was deleted.", "properties": { "event_type": { @@ -14986,7 +15284,6 @@ "type": "object" }, "AuditLogEventChainDeployedV1": { - "additionalProperties": false, "description": "A chain deployment was created.", "properties": { "event_type": { @@ -15039,7 +15336,6 @@ "type": "object" }, "AuditLogEventChainDeploymentActivatedV1": { - "additionalProperties": false, "description": "A chain deployment was activated.", "properties": { "event_type": { @@ -15082,7 +15378,6 @@ "type": "object" }, "AuditLogEventChainDeploymentDeactivatedV1": { - "additionalProperties": false, "description": "A chain deployment was deactivated.", "properties": { "event_type": { @@ -15125,7 +15420,6 @@ "type": "object" }, "AuditLogEventChainDeploymentDeletedV1": { - "additionalProperties": false, "description": "A chain deployment was deleted.", "properties": { "event_type": { @@ -15168,7 +15462,6 @@ "type": "object" }, "AuditLogEventChainDeploymentPromotedV1": { - "additionalProperties": false, "description": "A chain deployment was promoted to an environment.", "properties": { "event_type": { @@ -15235,7 +15528,6 @@ "type": "object" }, "AuditLogEventChainEnvironmentCreatedV1": { - "additionalProperties": false, "description": "A chain environment was created.", "properties": { "event_type": { @@ -15302,7 +15594,6 @@ "type": "object" }, "AuditLogEventChainEnvironmentUpdatedV1": { - "additionalProperties": false, "description": "A chain environment was updated.", "properties": { "event_type": { @@ -15369,7 +15660,6 @@ "type": "object" }, "AuditLogEventChainletAutoscalingSettingsChangedV1": { - "additionalProperties": false, "description": "A chainlet's autoscaling settings were changed.", "properties": { "min_replica": { @@ -15508,7 +15798,6 @@ "type": "object" }, "AuditLogEventChainletInstanceTypeChangedV1": { - "additionalProperties": false, "description": "A chainlet's instance type was changed.", "properties": { "event_type": { @@ -15566,7 +15855,6 @@ "type": "object" }, "AuditLogEventDirectoryGroupRoleUpdatedV1": { - "additionalProperties": false, "description": "A directory group's role was updated.", "properties": { "event_type": { @@ -15621,7 +15909,6 @@ "type": "object" }, "AuditLogEventEnvironmentCreatedV1": { - "additionalProperties": false, "description": "A model environment was created.", "properties": { "min_replica": { @@ -15883,7 +16170,6 @@ "type": "object" }, "AuditLogEventEnvironmentDeletedV1": { - "additionalProperties": false, "description": "A model environment was deleted.", "properties": { "event_type": { @@ -15914,7 +16200,6 @@ "type": "object" }, "AuditLogEventEnvironmentSettingsV1": { - "additionalProperties": false, "description": "Full environment settings (autoscaling + rolling promotion + deprecated canary);\nshared base for the environment events and the type of their previous_settings\nsnapshots. Not itself a payload in the discriminated union.", "properties": { "min_replica": { @@ -16143,7 +16428,6 @@ "type": "object" }, "AuditLogEventEnvironmentUpdatedV1": { - "additionalProperties": false, "description": "A model environment's settings were updated.", "properties": { "schedules": { @@ -16431,7 +16715,6 @@ "type": "object" }, "AuditLogEventGatewayEndpointCreatedV1": { - "additionalProperties": false, "description": "A Frontier Gateway endpoint was created.", "properties": { "event_type": { @@ -16457,7 +16740,6 @@ "type": "object" }, "AuditLogEventGatewayEndpointDeletedV1": { - "additionalProperties": false, "description": "A Frontier Gateway endpoint was deleted.", "properties": { "event_type": { @@ -16483,7 +16765,6 @@ "type": "object" }, "AuditLogEventGatewayEndpointUpdatedV1": { - "additionalProperties": false, "description": "A Frontier Gateway endpoint was updated.", "properties": { "event_type": { @@ -16521,7 +16802,6 @@ "type": "object" }, "AuditLogEventModelDeletedV1": { - "additionalProperties": false, "description": "A model was deleted.", "properties": { "event_type": { @@ -16547,7 +16827,6 @@ "type": "object" }, "AuditLogEventModelDeployedV1": { - "additionalProperties": false, "description": "A model deployment was created.", "properties": { "event_type": { @@ -16610,7 +16889,6 @@ "type": "object" }, "AuditLogEventModelDeploymentActivatedV1": { - "additionalProperties": false, "description": "A model deployment was activated.", "properties": { "event_type": { @@ -16646,7 +16924,6 @@ "type": "object" }, "AuditLogEventModelDeploymentAutoscalingSettingsChangedV1": { - "additionalProperties": false, "description": "A model deployment's autoscaling settings were changed.", "properties": { "schedules": { @@ -16795,7 +17072,6 @@ "type": "object" }, "AuditLogEventModelDeploymentDeactivatedV1": { - "additionalProperties": false, "description": "A model deployment was deactivated.", "properties": { "event_type": { @@ -16831,7 +17107,6 @@ "type": "object" }, "AuditLogEventModelDeploymentDeletedV1": { - "additionalProperties": false, "description": "A model deployment was deleted.", "properties": { "event_type": { @@ -16867,7 +17142,6 @@ "type": "object" }, "AuditLogEventModelDeploymentInstanceTypeChangedV1": { - "additionalProperties": false, "description": "A model deployment's instance type was changed.", "properties": { "event_type": { @@ -16908,7 +17182,6 @@ "type": "object" }, "AuditLogEventModelDeploymentPromotedV1": { - "additionalProperties": false, "description": "A model deployment was promoted to an environment.", "properties": { "event_type": { @@ -16968,7 +17241,6 @@ "type": "object" }, "AuditLogEventModelDeploymentRequestBackpressureSettingsChangedV1": { - "additionalProperties": false, "description": "A model deployment's request backpressure settings were changed.", "properties": { "event_type": { @@ -17028,7 +17300,6 @@ "type": "object" }, "AuditLogEventModelDeploymentRetriedV1": { - "additionalProperties": false, "description": "A model deployment build was retried.", "properties": { "event_type": { @@ -17069,7 +17340,6 @@ "type": "object" }, "AuditLogEventModelPromotionControlActionV1": { - "additionalProperties": false, "description": "A user-initiated promotion control signal was sent to a rolling promotion.", "properties": { "event_type": { @@ -17126,7 +17396,6 @@ "type": "object" }, "AuditLogEventModelRenamedV1": { - "additionalProperties": false, "description": "A model was renamed. `model_name` is the new name.", "properties": { "event_type": { @@ -17164,7 +17433,6 @@ "type": "object" }, "AuditLogEventProviderConnectionCreatedV1": { - "additionalProperties": false, "properties": { "event_type": { "const": "PROVIDER_CONNECTION_CREATED", @@ -17194,7 +17462,6 @@ "type": "object" }, "AuditLogEventProviderConnectionDeletedV1": { - "additionalProperties": false, "properties": { "event_type": { "const": "PROVIDER_CONNECTION_DELETED", @@ -17224,7 +17491,6 @@ "type": "object" }, "AuditLogEventProviderConnectionUpdatedV1": { - "additionalProperties": false, "properties": { "event_type": { "const": "PROVIDER_CONNECTION_UPDATED", @@ -17254,7 +17520,6 @@ "type": "object" }, "AuditLogEventReplicaTerminatedV1": { - "additionalProperties": false, "description": "A replica of a model deployment was terminated.", "properties": { "event_type": { @@ -17295,7 +17560,6 @@ "type": "object" }, "AuditLogEventRequireGroupBasedAdminsEnabledV1": { - "additionalProperties": false, "description": "Group-based admin enforcement was enabled for the organization.", "properties": { "event_type": { @@ -17316,7 +17580,6 @@ "type": "object" }, "AuditLogEventSecretDeletedV1": { - "additionalProperties": false, "description": "A secret was deleted.", "properties": { "event_type": { @@ -17342,7 +17605,6 @@ "type": "object" }, "AuditLogEventSecretUpdatedV1": { - "additionalProperties": false, "description": "A secret was created or updated.", "properties": { "event_type": { @@ -17368,7 +17630,6 @@ "type": "object" }, "AuditLogEventSshCertificateSignedV1": { - "additionalProperties": false, "description": "An SSH certificate was signed for a workload.", "properties": { "event_type": { @@ -17471,7 +17732,6 @@ "type": "string" }, "AuditLogEventUserInvitedV1": { - "additionalProperties": false, "description": "A user was invited to the organization.", "properties": { "event_type": { @@ -17497,7 +17757,6 @@ "type": "object" }, "AuditLogEventUserJoinedOrganizationV1": { - "additionalProperties": false, "description": "A user joined the organization.", "properties": { "event_type": { @@ -17523,7 +17782,6 @@ "type": "object" }, "AuditLogEventUserRemovedV1": { - "additionalProperties": false, "description": "A user was removed from the organization.", "properties": { "event_type": { @@ -17544,7 +17802,6 @@ "type": "object" }, "AuditLogEventUserRoleUpdatedV1": { - "additionalProperties": false, "description": "A user's organization role was updated.", "properties": { "event_type": { @@ -17575,7 +17832,6 @@ "type": "object" }, "AuditLogEventUserTeamRoleUpdatedV1": { - "additionalProperties": false, "description": "A user's team role was updated.", "properties": { "event_type": { @@ -17616,7 +17872,6 @@ "type": "object" }, "AuditLogEventVolumeDeletedV1": { - "additionalProperties": false, "description": "A volume was deleted, tombstoning every version it still held.", "properties": { "event_type": { @@ -17652,7 +17907,6 @@ "type": "object" }, "AuditLogEventVolumeVersionDeletedV1": { - "additionalProperties": false, "description": "One version of a volume was deleted.", "properties": { "event_type": { @@ -17693,7 +17947,6 @@ "type": "object" }, "AuditLogEventVolumeVersionRestoredV1": { - "additionalProperties": false, "description": "A deleted version of a volume was restored.", "properties": { "event_type": { @@ -17734,7 +17987,6 @@ "type": "object" }, "AuditLogEventWebhookSigningSecretCreatedV1": { - "additionalProperties": false, "description": "A webhook signing secret was created.", "properties": { "event_type": { @@ -17755,7 +18007,6 @@ "type": "object" }, "AuditLogEventWebhookSigningSecretDeletedV1": { - "additionalProperties": false, "description": "A webhook signing secret was deleted.", "properties": { "event_type": { @@ -17776,7 +18027,6 @@ "type": "object" }, "AuditLogEventWebhookSigningSecretRotatedV1": { - "additionalProperties": false, "description": "A webhook signing secret was rotated.", "properties": { "event_type": { @@ -25787,6 +26037,7 @@ "type": "object" }, "CreateLoopsRunRequestV1": { + "description": "Request to create a Loops run together with its paired sampler.", "properties": { "session_id": { "description": "ID of the Loops session this run belongs to.", @@ -25910,6 +26161,20 @@ "default": null, "description": "Optional ID of a prior Loops session whose trainer and/or sampler should be reused for this run. Deprecated in favor of reuse_from_run_id.", "title": "Reuse From Session Id" + }, + "sampler_num_replicas": { + "anyOf": [ + { + "minimum": 1, + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Number of replicas the run's sampler runs, applied as both its minimum and maximum. Must be at least 1. If omitted, a new sampler uses the platform defaults and a sampler reused from an earlier run keeps its settings. A run that already has a sampler keeps it unchanged. When the run ends, its sampler is scaled down.", + "title": "Sampler Num Replicas" } }, "required": [ @@ -25971,6 +26236,139 @@ "title": "GetLoopsRunResponseV1", "type": "object" }, + "CreateLoopsTrainerRequestV1": { + "properties": { + "session_id": { + "description": "ID of the Loops session this run belongs to.", + "title": "Session Id", + "type": "string" + }, + "base_model": { + "description": "Base model ID (e.g. 'Qwen/Qwen3-8B').", + "title": "Base Model", + "type": "string" + }, + "name": { + "anyOf": [ + { + "maxLength": 255, + "minLength": 1, + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Optional display name for the run. Defaults to the base model name when omitted.", + "title": "Name" + }, + "max_seq_len": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Maximum prompt length (in tokens) the run must handle. Set this to the longest training example you plan to send. Defaults to the maximum supported by the model configuration.", + "title": "Max Seq Len" + }, + "lora_rank": { + "default": 64, + "description": "LoRA rank.", + "minimum": 1, + "title": "Lora Rank", + "type": "integer" + }, + "seed": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Random seed for reproducibility.", + "title": "Seed" + }, + "scale_down_delay_seconds": { + "default": 900, + "description": "Seconds of inactivity before the run scales to zero. Must be between 1 and 3600 (1 hour). Defaults to 900 (15 minutes).", + "exclusiveMinimum": 0, + "maximum": 3600, + "title": "Scale Down Delay Seconds", + "type": "integer" + }, + "availability_model": { + "$ref": "#/components/schemas/V1AvailabilityModel", + "default": "dedicated", + "description": "Capacity the trainer runs on. 'dedicated' is not preempted. 'spot' runs below inference and reaches idle reserved capacity, but the run is stopped if its GPUs are reclaimed and cannot be resumed.", + "examples": [ + "spot" + ] + }, + "replicas": { + "default": 1, + "description": "Number of data-parallel trainer replicas. Each replica is one full copy of the model's preset node group, so the trainer deployment runs (preset node_count * replicas) nodes (e.g. replicas=4 on a 4-node preset \u2192 16 nodes, 4 DP workers). Must be a positive integer. Defaults to 1.", + "minimum": 1, + "title": "Replicas", + "type": "integer" + }, + "path": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Optional bt:// URI of an existing checkpoint to load weights from on startup. Form: bt://loops:/weights/.", + "examples": [ + "bt://loops:k4q95w5/weights/step-100" + ], + "title": "Path" + }, + "reuse_from_run_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Optional ID of a prior Loops run whose trainer and/or sampler should be reused for this run instead of provisioning fresh. The prior run must use the same base model and belong to the same team.", + "title": "Reuse From Run Id" + }, + "reuse_from_session_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Optional ID of a prior Loops session whose trainer and/or sampler should be reused for this run. Deprecated in favor of reuse_from_run_id.", + "title": "Reuse From Session Id" + } + }, + "required": [ + "session_id", + "base_model" + ], + "title": "CreateLoopsTrainerRequestV1", + "type": "object" + }, "ListLoopsSamplersResponseV1": { "description": "Response for ``GET /v1/loops/samplers``.\n\nReturns the caller's samplers, including those paired to runs and\nstandalone samplers. Ordered newest-first.", "properties": { @@ -26098,6 +26496,20 @@ "default": null, "description": "Optional ID of a prior Loops session to reuse a trainer and/or sampler from. Deprecated.", "title": "Reuse From Session Id" + }, + "num_replicas": { + "anyOf": [ + { + "minimum": 1, + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Number of replicas the sampler runs, applied as both its minimum and maximum. Must be at least 1. If omitted, a new sampler uses the platform defaults and a paired sampler reused from an earlier run keeps its settings. A run that already has a sampler keeps it unchanged, and when a run ends, its sampler is scaled down.", + "title": "Num Replicas" } }, "required": [ @@ -28964,7 +29376,6 @@ "type": "object" }, "LibraryListingMetadataV1": { - "additionalProperties": false, "properties": { "parameter_count": { "anyOf": [ @@ -29294,7 +29705,7 @@ } ], "default": null, - "description": "Model-level metadata for the listing. When provided, replaces the stored metadata. Unknown fields are rejected." + "description": "Model-level metadata for the listing. When provided, replaces the stored metadata." } }, "title": "UpdateLibraryListingRequestV1", @@ -31582,8 +31993,27 @@ "title": "ExploreMetadataRequestV1", "type": "object" }, + "RouteRefV1": { + "properties": { + "id": { + "description": "Stable route identifier.", + "title": "Id", + "type": "string" + }, + "slug": { + "description": "Name of the route to send in the inference request's model field.", + "title": "Slug", + "type": "string" + } + }, + "required": [ + "id", + "slug" + ], + "title": "RouteRefV1", + "type": "object" + }, "RouteTargetAnthropicV1": { - "additionalProperties": false, "properties": { "type": { "const": "ANTHROPIC", @@ -31613,7 +32043,6 @@ "type": "object" }, "RouteTargetBasetenModelAPIV1": { - "additionalProperties": false, "properties": { "type": { "const": "BASETEN_MODEL_API", @@ -31635,8 +32064,61 @@ "title": "RouteTargetBasetenModelAPIV1", "type": "object" }, + "RouteTargetClassifierModelBasedV1": { + "properties": { + "type": { + "const": "CLASSIFIER_MODEL_BASED", + "description": "Target kind for a classifier-model-based route.", + "title": "Type", + "type": "string" + }, + "classifier_model_id": { + "description": "ID of the Baseten model that picks a route for each request.", + "examples": [ + "abcd123" + ], + "title": "Classifier Model Id", + "type": "string" + }, + "classifier_environment_name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Environment of the classifier model. Null for production.", + "examples": [ + "production" + ], + "title": "Classifier Environment Name" + }, + "default_route": { + "$ref": "#/components/schemas/RouteRefV1", + "description": "Route used when the classifier picks none." + }, + "allowed_routes": { + "description": "Routes the classifier may pick, in creation order.", + "items": { + "$ref": "#/components/schemas/RouteRefV1" + }, + "title": "Allowed Routes", + "type": "array" + } + }, + "required": [ + "type", + "classifier_model_id", + "classifier_environment_name", + "default_route", + "allowed_routes" + ], + "title": "RouteTargetClassifierModelBasedV1", + "type": "object" + }, "RouteTargetOpenAIV1": { - "additionalProperties": false, "properties": { "type": { "const": "OPENAI", @@ -31666,7 +32148,6 @@ "type": "object" }, "RouteTargetXAIV1": { - "additionalProperties": false, "properties": { "type": { "const": "XAI", @@ -31733,6 +32214,7 @@ "mapping": { "ANTHROPIC": "#/components/schemas/RouteTargetAnthropicV1", "BASETEN_MODEL_API": "#/components/schemas/RouteTargetBasetenModelAPIV1", + "CLASSIFIER_MODEL_BASED": "#/components/schemas/RouteTargetClassifierModelBasedV1", "OPENAI": "#/components/schemas/RouteTargetOpenAIV1", "XAI": "#/components/schemas/RouteTargetXAIV1" }, @@ -31750,6 +32232,9 @@ }, { "$ref": "#/components/schemas/RouteTargetXAIV1" + }, + { + "$ref": "#/components/schemas/RouteTargetClassifierModelBasedV1" } ], "title": "Target" @@ -31763,7 +32248,7 @@ "type": "null" } ], - "description": "Resolved model metadata; null when the route has no linked metadata row." + "description": "Resolved model metadata; for a router, the envelope of its allowed routes' metadata. Null when nothing is linked." }, "invoke_url": { "description": "Base URL for inference requests using this route.", @@ -31847,7 +32332,7 @@ } ], "default": null, - "description": "Filter by owning team ID. Preserved by the cursor; if repeated, must match the original filter.", + "description": "Identifier of the team whose routes to list.", "title": "Team Id" }, "name": { @@ -31867,52 +32352,233 @@ "title": "RoutesRequestV1", "type": "object" }, - "CreateRouteRequestV1": { + "RouteTargetConfigAnthropicV1": { "additionalProperties": false, "properties": { - "team_id": { - "anyOf": [ - { - "minLength": 1, - "type": "string" - }, - { - "type": "null" - } - ], - "default": null, - "description": "Identifier of the team that owns the route. When omitted, uses your organization's default team.", - "examples": [ - "abc1234" - ], - "title": "Team Id" + "type": { + "const": "ANTHROPIC", + "description": "Target kind for Anthropic.", + "title": "Type", + "type": "string" }, - "display_name": { - "anyOf": [ - { - "maxLength": 255, - "minLength": 1, - "type": "string" - }, - { - "type": "null" - } - ], - "default": null, - "description": "Display label. Omit to use the route name; null is not accepted.", - "examples": [ - "Assistant" - ], - "title": "Display Name" + "model": { + "description": "Model name sent to the provider.", + "minLength": 1, + "title": "Model", + "type": "string" }, - "target": { - "description": "Upstream target for the route.", - "discriminator": { - "mapping": { - "ANTHROPIC": "#/components/schemas/RouteTargetAnthropicV1", - "BASETEN_MODEL_API": "#/components/schemas/RouteTargetBasetenModelAPIV1", - "OPENAI": "#/components/schemas/RouteTargetOpenAIV1", - "XAI": "#/components/schemas/RouteTargetXAIV1" + "secret_name": { + "description": "Name of a credential secret owned by the route's team.", + "minLength": 1, + "title": "Secret Name", + "type": "string" + } + }, + "required": [ + "type", + "model", + "secret_name" + ], + "title": "RouteTargetConfigAnthropicV1", + "type": "object" + }, + "RouteTargetConfigBasetenModelAPIV1": { + "additionalProperties": false, + "properties": { + "type": { + "const": "BASETEN_MODEL_API", + "description": "Target kind for a Baseten Model API.", + "title": "Type", + "type": "string" + }, + "model": { + "description": "Name of the target Model API.", + "minLength": 1, + "title": "Model", + "type": "string" + } + }, + "required": [ + "type", + "model" + ], + "title": "RouteTargetConfigBasetenModelAPIV1", + "type": "object" + }, + "RouteTargetConfigClassifierModelBasedV1": { + "additionalProperties": false, + "properties": { + "type": { + "const": "CLASSIFIER_MODEL_BASED", + "description": "Target kind for a classifier-model-based route. Not intended for general use: classifiers are deployed by Baseten's post-training team.", + "title": "Type", + "type": "string" + }, + "classifier_model_id": { + "description": "ID of the Baseten model that picks a route for each request. Only classifiers deployed by Baseten's post-training team are supported.", + "examples": [ + "abcd123" + ], + "minLength": 1, + "title": "Classifier Model Id", + "type": "string" + }, + "classifier_environment_name": { + "anyOf": [ + { + "minLength": 1, + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Environment of the classifier model. Omit for production, which is returned as null.", + "examples": [ + "production" + ], + "title": "Classifier Environment Name" + }, + "allowed_route_ids": { + "description": "IDs of the routes the classifier may pick. All must belong to the route's team and must not be routers themselves.", + "examples": [ + [ + "abc1234", + "def5678" + ] + ], + "items": { + "type": "string" + }, + "title": "Allowed Route Ids", + "type": "array" + }, + "default_route_id": { + "description": "ID of the allowed route used when the classifier picks none.", + "examples": [ + "abc1234" + ], + "title": "Default Route Id", + "type": "string" + } + }, + "required": [ + "type", + "classifier_model_id", + "allowed_route_ids", + "default_route_id" + ], + "title": "RouteTargetConfigClassifierModelBasedV1", + "type": "object" + }, + "RouteTargetConfigOpenAIV1": { + "additionalProperties": false, + "properties": { + "type": { + "const": "OPENAI", + "description": "Target kind for OpenAI.", + "title": "Type", + "type": "string" + }, + "model": { + "description": "Model name sent to the provider.", + "minLength": 1, + "title": "Model", + "type": "string" + }, + "secret_name": { + "description": "Name of a credential secret owned by the route's team.", + "minLength": 1, + "title": "Secret Name", + "type": "string" + } + }, + "required": [ + "type", + "model", + "secret_name" + ], + "title": "RouteTargetConfigOpenAIV1", + "type": "object" + }, + "RouteTargetConfigXAIV1": { + "additionalProperties": false, + "properties": { + "type": { + "const": "XAI", + "description": "Target kind for xAI.", + "title": "Type", + "type": "string" + }, + "model": { + "description": "Model name sent to the provider.", + "minLength": 1, + "title": "Model", + "type": "string" + }, + "secret_name": { + "description": "Name of a credential secret owned by the route's team.", + "minLength": 1, + "title": "Secret Name", + "type": "string" + } + }, + "required": [ + "type", + "model", + "secret_name" + ], + "title": "RouteTargetConfigXAIV1", + "type": "object" + }, + "CreateRouteRequestV1": { + "additionalProperties": false, + "properties": { + "team_id": { + "anyOf": [ + { + "minLength": 1, + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Identifier of the team that owns the route.", + "examples": [ + "abc1234" + ], + "title": "Team Id" + }, + "display_name": { + "anyOf": [ + { + "maxLength": 255, + "minLength": 1, + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Display label. Omit to use the route name; null is not accepted.", + "examples": [ + "Assistant" + ], + "title": "Display Name" + }, + "target": { + "description": "Upstream target for the route.", + "discriminator": { + "mapping": { + "ANTHROPIC": "#/components/schemas/RouteTargetConfigAnthropicV1", + "BASETEN_MODEL_API": "#/components/schemas/RouteTargetConfigBasetenModelAPIV1", + "CLASSIFIER_MODEL_BASED": "#/components/schemas/RouteTargetConfigClassifierModelBasedV1", + "OPENAI": "#/components/schemas/RouteTargetConfigOpenAIV1", + "XAI": "#/components/schemas/RouteTargetConfigXAIV1" }, "propertyName": "type" }, @@ -31924,16 +32590,19 @@ ], "oneOf": [ { - "$ref": "#/components/schemas/RouteTargetBasetenModelAPIV1" + "$ref": "#/components/schemas/RouteTargetConfigBasetenModelAPIV1" }, { - "$ref": "#/components/schemas/RouteTargetAnthropicV1" + "$ref": "#/components/schemas/RouteTargetConfigAnthropicV1" }, { - "$ref": "#/components/schemas/RouteTargetOpenAIV1" + "$ref": "#/components/schemas/RouteTargetConfigOpenAIV1" }, { - "$ref": "#/components/schemas/RouteTargetXAIV1" + "$ref": "#/components/schemas/RouteTargetConfigXAIV1" + }, + { + "$ref": "#/components/schemas/RouteTargetConfigClassifierModelBasedV1" } ], "title": "Target" @@ -32196,17 +32865,9 @@ "title": "RoutesUsageRequestV1", "type": "object" }, - "RouteSpendLimitV1": { + "RouteEffectiveSpendLimitV1": { "properties": { - "user_id": { - "description": "ID of the user.", - "examples": [ - "abc1234" - ], - "title": "User Id", - "type": "string" - }, - "email": { + "monthly_limit_usd": { "anyOf": [ { "type": "string" @@ -32215,41 +32876,42 @@ "type": "null" } ], - "description": "Email address of the user.", + "description": "Spend limit in USD enforced for the current UTC calendar month: the user's own limit, else the team default. Null when no limit applies.", "examples": [ - "dev@example.com" + "200" ], - "title": "Email" + "title": "Monthly Limit Usd" }, - "monthly_limit_usd": { + "source": { "anyOf": [ { - "type": "string" + "$ref": "#/components/schemas/RouteSettingSourceV1" }, { "type": "null" } ], - "description": "Standing spend limit in USD for each UTC calendar month, returned as an exact decimal string. Null when the user has no standing limit.", - "examples": [ - "200" - ], - "title": "Monthly Limit Usd" + "description": "Where the limit comes from: `user` when it is set on the user, `team` when it is the team default. Null when no limit applies." } }, "required": [ - "user_id", - "email", - "monthly_limit_usd" + "monthly_limit_usd", + "source" ], - "title": "RouteSpendLimitV1", + "title": "RouteEffectiveSpendLimitV1", "type": "object" }, - "UpdateRouteSpendLimitRequestV1": { - "additionalProperties": false, - "minProperties": 1, + "RouteSettingSourceV1": { + "enum": [ + "user", + "team" + ], + "title": "RouteSettingSourceV1", + "type": "string" + }, + "RouteSpendLimitSettingV1": { "properties": { - "monthly_limit_usd": { + "user_monthly_limit_usd": { "anyOf": [ { "type": "string" @@ -32258,291 +32920,298 @@ "type": "null" } ], - "default": null, - "description": "Standing spend limit in USD for each UTC calendar month, as a non-negative decimal string with at most 9 decimal places. Send null to remove the limit; omit to leave it unchanged.", + "description": "Standing spend limit in USD for each UTC calendar month set on the user. Null when the user has no limit of their own.", "examples": [ "200" ], - "title": "Monthly Limit Usd", - "x-null-distinct": true + "title": "User Monthly Limit Usd" + }, + "team_default": { + "anyOf": [ + { + "$ref": "#/components/schemas/RouteSpendLimitTeamDefaultV1" + }, + { + "type": "null" + } + ], + "description": "Per-member limit of the team the user's active Code key belongs to. This limit applies when the user has no limit of their own. Null when there is none." + }, + "effective": { + "$ref": "#/components/schemas/RouteEffectiveSpendLimitV1", + "description": "The effective limit enforced for the current month." } }, - "title": "UpdateRouteSpendLimitRequestV1", - "type": "object" - }, - "RouteHarness": { - "enum": [ - "claude-code", - "opencode", - "codex" + "required": [ + "user_monthly_limit_usd", + "team_default", + "effective" ], - "title": "RouteHarness", - "type": "string" + "title": "RouteSpendLimitSettingV1", + "type": "object" }, - "RouteHarnessConfigV1": { + "RouteSpendLimitTeamDefaultV1": { "properties": { - "models": { - "additionalProperties": { - "$ref": "#/components/schemas/RouteHarnessModelV1" - }, - "description": "Route for each model role, keyed by role.", - "propertyNames": { - "$ref": "#/components/schemas/RouteHarnessRole" - }, - "title": "Models", - "type": "object" + "team_id": { + "description": "ID of the team the user's active Code key belongs to.", + "examples": [ + "abc1234" + ], + "title": "Team Id", + "type": "string" + }, + "per_member_monthly_limit_usd": { + "description": "The team's per-member spend limit in USD for each UTC calendar month.", + "examples": [ + "100" + ], + "title": "Per Member Monthly Limit Usd", + "type": "string" } }, "required": [ - "models" + "team_id", + "per_member_monthly_limit_usd" ], - "title": "RouteHarnessConfigV1", + "title": "RouteSpendLimitTeamDefaultV1", "type": "object" }, - "RouteHarnessModelSource": { - "enum": [ - "team", - "baseten" - ], - "title": "RouteHarnessModelSource", - "type": "string" - }, - "RouteHarnessModelV1": { + "RouteUserSettingsV1": { "properties": { - "source": { - "$ref": "#/components/schemas/RouteHarnessModelSource", - "description": "Who chose this role's route: `team` if a team admin set it, or `baseten` if it is Baseten's default, chosen from the team's Model API routes." + "user_id": { + "description": "ID of the user.", + "examples": [ + "abc1234" + ], + "title": "User Id", + "type": "string" }, - "route": { - "$ref": "#/components/schemas/RouteV1", - "description": "Route to use for this role." + "email": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Email address of the user.", + "examples": [ + "dev@example.com" + ], + "title": "Email" + }, + "spend_limit": { + "$ref": "#/components/schemas/RouteSpendLimitSettingV1", + "description": "Spend limit for Baseten Code. Applies to requests with Routes keys the user created; personal API keys are not limited." } }, "required": [ - "source", - "route" + "user_id", + "email", + "spend_limit" ], - "title": "RouteHarnessModelV1", + "title": "RouteUserSettingsV1", "type": "object" }, - "RouteHarnessRole": { - "enum": [ - "primary", - "background" - ], - "title": "RouteHarnessRole", - "type": "string" - }, - "RouteHarnessConfigsResponseV1": { + "UpdateRouteSpendLimitSettingV1": { + "additionalProperties": false, + "minProperties": 1, "properties": { - "harness_configs": { - "additionalProperties": { - "$ref": "#/components/schemas/RouteHarnessConfigV1" - }, - "description": "Default models for each harness, keyed by harness. A harness is omitted when none of its roles has a route.", - "propertyNames": { - "$ref": "#/components/schemas/RouteHarness" - }, - "title": "Harness Configs", - "type": "object" + "user_monthly_limit_usd": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Standing spend limit in USD for each UTC calendar month. Send null to remove the limit; omit to leave it unchanged.", + "examples": [ + "200" + ], + "title": "User Monthly Limit Usd", + "x-null-distinct": true } }, - "required": [ - "harness_configs" - ], - "title": "RouteHarnessConfigsResponseV1", + "title": "UpdateRouteSpendLimitSettingV1", "type": "object" }, - "RouteHarnessConfigsRequestV1": { + "UpdateRouteUserSettingsRequestV1": { + "additionalProperties": false, + "minProperties": 1, "properties": { - "team_id": { - "description": "Identifier of the team whose default models to list.", - "minLength": 1, - "title": "Team Id", - "type": "string" + "spend_limit": { + "anyOf": [ + { + "$ref": "#/components/schemas/UpdateRouteSpendLimitSettingV1" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Spend limit fields to change. Pass null to remove the user's limit; omit to leave it unchanged.", + "x-null-distinct": true } }, - "required": [ - "team_id" - ], - "title": "RouteHarnessConfigsRequestV1", + "title": "UpdateRouteUserSettingsRequestV1", "type": "object" }, - "UpdateBackgroundHarnessModelsV1": { - "additionalProperties": false, - "minProperties": 1, + "BackgroundHarnessDefaultsV1": { "properties": { "primary": { "anyOf": [ { - "minLength": 1, - "type": "string" + "$ref": "#/components/schemas/RouteHarnessModelV1" }, { "type": "null" } ], - "default": null, - "description": "Route ID for the primary model, which new sessions use. Omit to keep the current route, or pass null to use Baseten's default.", - "examples": [ - "abc1234" - ], - "title": "Primary" + "description": "Route for the primary model, which new sessions use. Null when the team has no route to use." }, "background": { "anyOf": [ { - "minLength": 1, - "type": "string" + "$ref": "#/components/schemas/RouteHarnessModelV1" }, { "type": "null" } ], - "default": null, - "description": "Route ID for background tasks, such as session titles. Omit to keep the current route, or pass null to use Baseten's default.", - "examples": [ - "def5678" - ], - "title": "Background" + "description": "Route for background tasks, such as session titles. Null when the team has no route to use." } }, - "title": "UpdateBackgroundHarnessModelsV1", + "required": [ + "primary", + "background" + ], + "title": "BackgroundHarnessDefaultsV1", "type": "object" }, - "UpdateClaudeCodeHarnessConfigV1": { - "additionalProperties": false, + "PrimaryHarnessDefaultsV1": { "properties": { - "team_id": { - "description": "Identifier of the team whose default models to set. Every route must belong to this team.", - "minLength": 1, - "title": "Team Id", - "type": "string" - }, - "harness": { - "const": "claude-code", - "description": "Claude Code, which supports the `primary` and `background` roles.", - "title": "Harness", - "type": "string" - }, - "models": { - "$ref": "#/components/schemas/UpdateBackgroundHarnessModelsV1", - "description": "Route IDs for the model roles to change. Roles left out are unchanged." + "primary": { + "anyOf": [ + { + "$ref": "#/components/schemas/RouteHarnessModelV1" + }, + { + "type": "null" + } + ], + "description": "Route for the primary model, which new sessions use. Null when the team has no route to use." } }, "required": [ - "team_id", - "harness", - "models" + "primary" ], - "title": "UpdateClaudeCodeHarnessConfigV1", + "title": "PrimaryHarnessDefaultsV1", "type": "object" }, - "UpdateCodexHarnessConfigV1": { - "additionalProperties": false, + "RouteHarnessDefaultsV1": { "properties": { - "team_id": { - "description": "Identifier of the team whose default models to set. Every route must belong to this team.", - "minLength": 1, - "title": "Team Id", - "type": "string" + "claude_code": { + "$ref": "#/components/schemas/BackgroundHarnessDefaultsV1", + "description": "Default models for Claude Code." }, - "harness": { - "const": "codex", - "description": "Codex, which supports only the `primary` role.", - "title": "Harness", - "type": "string" + "opencode": { + "$ref": "#/components/schemas/BackgroundHarnessDefaultsV1", + "description": "Default models for OpenCode." }, - "models": { - "$ref": "#/components/schemas/UpdatePrimaryHarnessModelsV1", - "description": "Route IDs for the model roles to change. Roles left out are unchanged." + "codex": { + "$ref": "#/components/schemas/PrimaryHarnessDefaultsV1", + "description": "Default models for Codex." + }, + "pi": { + "$ref": "#/components/schemas/PrimaryHarnessDefaultsV1", + "description": "Default models for Pi." } }, "required": [ - "team_id", - "harness", - "models" + "claude_code", + "opencode", + "codex", + "pi" ], - "title": "UpdateCodexHarnessConfigV1", + "title": "RouteHarnessDefaultsV1", "type": "object" }, - "UpdateOpenCodeHarnessConfigV1": { - "additionalProperties": false, + "RouteHarnessModelV1": { "properties": { - "team_id": { - "description": "Identifier of the team whose default models to set. Every route must belong to this team.", - "minLength": 1, - "title": "Team Id", - "type": "string" - }, - "harness": { - "const": "opencode", - "description": "OpenCode, which supports the `primary` and `background` roles.", - "title": "Harness", - "type": "string" + "source": { + "$ref": "#/components/schemas/RouteSettingSourceV1", + "description": "Where this role's route comes from. `team` covers both a route a team admin set and the default chosen from the team's Model API routes." }, - "models": { - "$ref": "#/components/schemas/UpdateBackgroundHarnessModelsV1", - "description": "Route IDs for the model roles to change. Roles left out are unchanged." + "route": { + "$ref": "#/components/schemas/RouteV1", + "description": "Route to use for this role." } }, "required": [ - "team_id", - "harness", - "models" + "source", + "route" ], - "title": "UpdateOpenCodeHarnessConfigV1", + "title": "RouteHarnessModelV1", "type": "object" }, - "UpdatePrimaryHarnessModelsV1": { - "additionalProperties": false, - "minProperties": 1, + "RouteTeamSpendLimitSettingV1": { "properties": { - "primary": { + "per_member_monthly_limit_usd": { "anyOf": [ { - "minLength": 1, "type": "string" }, { "type": "null" } ], - "default": null, - "description": "Route ID for the primary model, which new sessions use. Omit to keep the current route, or pass null to use Baseten's default.", + "description": "Spend limit in USD for each UTC calendar month that applies to each member whose active Code key belongs to this team, unless the member has a limit of their own. Returned as an exact decimal string. Null when the team has no per-member limit.", "examples": [ - "abc1234" + "200" ], - "title": "Primary" + "title": "Per Member Monthly Limit Usd" } }, - "title": "UpdatePrimaryHarnessModelsV1", + "required": [ + "per_member_monthly_limit_usd" + ], + "title": "RouteTeamSpendLimitSettingV1", "type": "object" }, - "UpdateRouteHarnessConfigRequestV1": { - "discriminator": { - "mapping": { - "claude-code": "#/components/schemas/UpdateClaudeCodeHarnessConfigV1", - "codex": "#/components/schemas/UpdateCodexHarnessConfigV1", - "opencode": "#/components/schemas/UpdateOpenCodeHarnessConfigV1" - }, - "propertyName": "harness" - }, - "oneOf": [ - { - "$ref": "#/components/schemas/UpdateClaudeCodeHarnessConfigV1" + "RouteTeamSettingsV1": { + "properties": { + "team_id": { + "description": "ID of the team.", + "examples": [ + "abc1234" + ], + "title": "Team Id", + "type": "string" }, - { - "$ref": "#/components/schemas/UpdateOpenCodeHarnessConfigV1" + "harness_defaults": { + "$ref": "#/components/schemas/RouteHarnessDefaultsV1", + "description": "Route for each model role of each coding harness. A role with no route set by a team admin uses the default chosen from the team's Model API routes: the recommended model for `primary` and the lowest-priced model for `background`. These defaults never use external-provider routes." }, - { - "$ref": "#/components/schemas/UpdateCodexHarnessConfigV1" + "spend_limit": { + "$ref": "#/components/schemas/RouteTeamSpendLimitSettingV1", + "description": "Spend limits for Baseten Code for members whose active Code key belongs to this team." } + }, + "required": [ + "team_id", + "harness_defaults", + "spend_limit" ], - "title": "UpdateRouteHarnessConfigRequestV1" + "title": "RouteTeamSettingsV1", + "type": "object" }, - "BackgroundHarnessModelsV1": { + "UpdateBackgroundHarnessModelsV1": { "additionalProperties": false, + "minProperties": 1, "properties": { "primary": { "anyOf": [ @@ -32555,11 +33224,12 @@ } ], "default": null, - "description": "Route ID for the primary model, which new sessions use. Omit to use Baseten's default.", + "description": "Route ID for the primary model, which new sessions use. Omit to keep the current route, or pass null to use the team's default.", "examples": [ "abc1234" ], - "title": "Primary" + "title": "Primary", + "x-null-distinct": true }, "background": { "anyOf": [ @@ -32572,18 +33242,20 @@ } ], "default": null, - "description": "Route ID for background tasks, such as session titles. Omit to use Baseten's default.", + "description": "Route ID for background tasks, such as session titles. Omit to keep the current route, or pass null to use the team's default.", "examples": [ "def5678" ], - "title": "Background" + "title": "Background", + "x-null-distinct": true } }, - "title": "BackgroundHarnessModelsV1", + "title": "UpdateBackgroundHarnessModelsV1", "type": "object" }, - "PrimaryHarnessModelsV1": { + "UpdatePrimaryHarnessModelsV1": { "additionalProperties": false, + "minProperties": 1, "properties": { "primary": { "anyOf": [ @@ -32596,184 +33268,302 @@ } ], "default": null, - "description": "Route ID for the primary model, which new sessions use. Omit to use Baseten's default.", + "description": "Route ID for the primary model, which new sessions use. Omit to keep the current route, or pass null to use the team's default.", "examples": [ "abc1234" ], - "title": "Primary" + "title": "Primary", + "x-null-distinct": true } }, - "title": "PrimaryHarnessModelsV1", + "title": "UpdatePrimaryHarnessModelsV1", "type": "object" }, - "SetClaudeCodeHarnessConfigV1": { + "UpdateRouteHarnessDefaultsV1": { "additionalProperties": false, + "minProperties": 1, "properties": { - "team_id": { - "description": "Identifier of the team whose default models to set. Every route must belong to this team.", - "minLength": 1, - "title": "Team Id", - "type": "string" + "claude_code": { + "anyOf": [ + { + "$ref": "#/components/schemas/UpdateBackgroundHarnessModelsV1" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Route IDs for the model roles to change; roles left out are unchanged. Every route must belong to the team. Pass null to clear every role, or omit to leave the harness unchanged.", + "x-null-distinct": true }, - "harness": { - "const": "claude-code", - "description": "Claude Code, which supports the `primary` and `background` roles.", - "title": "Harness", - "type": "string" + "opencode": { + "anyOf": [ + { + "$ref": "#/components/schemas/UpdateBackgroundHarnessModelsV1" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Route IDs for the model roles to change; roles left out are unchanged. Every route must belong to the team. Pass null to clear every role, or omit to leave the harness unchanged.", + "x-null-distinct": true }, - "models": { - "$ref": "#/components/schemas/BackgroundHarnessModelsV1", - "description": "Route ID for each model role. Roles left out use Baseten's defaults." + "codex": { + "anyOf": [ + { + "$ref": "#/components/schemas/UpdatePrimaryHarnessModelsV1" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Route IDs for the model roles to change; roles left out are unchanged. Every route must belong to the team. Pass null to clear every role, or omit to leave the harness unchanged.", + "x-null-distinct": true + }, + "pi": { + "anyOf": [ + { + "$ref": "#/components/schemas/UpdatePrimaryHarnessModelsV1" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Route IDs for the model roles to change; roles left out are unchanged. Every route must belong to the team. Pass null to clear every role, or omit to leave the harness unchanged.", + "x-null-distinct": true } }, - "required": [ - "team_id", - "harness", - "models" - ], - "title": "SetClaudeCodeHarnessConfigV1", + "title": "UpdateRouteHarnessDefaultsV1", "type": "object" }, - "SetCodexHarnessConfigV1": { + "UpdateRouteTeamSpendLimitSettingV1": { "additionalProperties": false, + "minProperties": 1, "properties": { - "team_id": { - "description": "Identifier of the team whose default models to set. Every route must belong to this team.", - "minLength": 1, - "title": "Team Id", - "type": "string" - }, - "harness": { - "const": "codex", - "description": "Codex, which supports only the `primary` role.", - "title": "Harness", - "type": "string" - }, - "models": { - "$ref": "#/components/schemas/PrimaryHarnessModelsV1", - "description": "Route ID for each model role. Roles left out use Baseten's defaults." + "per_member_monthly_limit_usd": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Spend limit in USD for each UTC calendar month that applies to each member whose active Code key belongs to this team, unless the member has a limit of their own. Send null to remove it; omit to leave it unchanged.", + "examples": [ + "200" + ], + "title": "Per Member Monthly Limit Usd", + "x-null-distinct": true } }, - "required": [ - "team_id", - "harness", - "models" - ], - "title": "SetCodexHarnessConfigV1", + "title": "UpdateRouteTeamSpendLimitSettingV1", "type": "object" }, - "SetOpenCodeHarnessConfigV1": { + "UpdateRouteTeamSettingsRequestV1": { "additionalProperties": false, + "minProperties": 1, "properties": { - "team_id": { - "description": "Identifier of the team whose default models to set. Every route must belong to this team.", - "minLength": 1, - "title": "Team Id", + "harness_defaults": { + "anyOf": [ + { + "$ref": "#/components/schemas/UpdateRouteHarnessDefaultsV1" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Harnesses to change; harnesses left out are unchanged. Pass null to clear every harness, so each role uses the team's default.", + "x-null-distinct": true + }, + "spend_limit": { + "anyOf": [ + { + "$ref": "#/components/schemas/UpdateRouteTeamSpendLimitSettingV1" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Spend limit fields to change. Pass null to remove the team's per-member limit; omit to leave it unchanged.", + "x-null-distinct": true + } + }, + "title": "UpdateRouteTeamSettingsRequestV1", + "type": "object" + }, + "RouteConnectionAnthropicV1": { + "properties": { + "provider": { + "const": "ANTHROPIC", + "description": "Provider kind for Anthropic.", + "title": "Provider", "type": "string" }, - "harness": { - "const": "opencode", - "description": "OpenCode, which supports the `primary` and `background` roles.", - "title": "Harness", + "secret_id": { + "description": "Identifier of the team secret holding the provider API key.", + "title": "Secret Id", "type": "string" }, - "models": { - "$ref": "#/components/schemas/BackgroundHarnessModelsV1", - "description": "Route ID for each model role. Roles left out use Baseten's defaults." + "secret_name": { + "description": "Name of the team secret holding the provider API key.", + "title": "Secret Name", + "type": "string" } }, "required": [ - "team_id", - "harness", - "models" + "provider", + "secret_id", + "secret_name" ], - "title": "SetOpenCodeHarnessConfigV1", + "title": "RouteConnectionAnthropicV1", "type": "object" }, - "SetRouteHarnessConfigRequestV1": { - "discriminator": { - "mapping": { - "claude-code": "#/components/schemas/SetClaudeCodeHarnessConfigV1", - "codex": "#/components/schemas/SetCodexHarnessConfigV1", - "opencode": "#/components/schemas/SetOpenCodeHarnessConfigV1" - }, - "propertyName": "harness" - }, - "oneOf": [ - { - "$ref": "#/components/schemas/SetClaudeCodeHarnessConfigV1" + "RouteConnectionOpenAIV1": { + "properties": { + "provider": { + "const": "OPENAI", + "description": "Provider kind for OpenAI.", + "title": "Provider", + "type": "string" }, - { - "$ref": "#/components/schemas/SetOpenCodeHarnessConfigV1" + "secret_id": { + "description": "Identifier of the team secret holding the provider API key.", + "title": "Secret Id", + "type": "string" }, - { - "$ref": "#/components/schemas/SetCodexHarnessConfigV1" + "secret_name": { + "description": "Name of the team secret holding the provider API key.", + "title": "Secret Name", + "type": "string" } + }, + "required": [ + "provider", + "secret_id", + "secret_name" ], - "title": "SetRouteHarnessConfigRequestV1" + "title": "RouteConnectionOpenAIV1", + "type": "object" }, - "RouteHarnessConfigTombstoneV1": { + "RouteConnectionV1": { "properties": { - "harness": { - "$ref": "#/components/schemas/RouteHarness", - "description": "Harness whose default models were cleared." + "id": { + "description": "Stable connection identifier.", + "title": "Id", + "type": "string" }, "team_id": { - "description": "Identifier of the team whose default models were cleared.", + "description": "Identifier of the team that owns the connection.", "title": "Team Id", "type": "string" + }, + "config": { + "description": "Provider the connection authenticates with, and the team secret holding its API key.", + "discriminator": { + "mapping": { + "ANTHROPIC": "#/components/schemas/RouteConnectionAnthropicV1", + "OPENAI": "#/components/schemas/RouteConnectionOpenAIV1", + "XAI": "#/components/schemas/RouteConnectionXAIV1" + }, + "propertyName": "provider" + }, + "oneOf": [ + { + "$ref": "#/components/schemas/RouteConnectionAnthropicV1" + }, + { + "$ref": "#/components/schemas/RouteConnectionOpenAIV1" + }, + { + "$ref": "#/components/schemas/RouteConnectionXAIV1" + } + ], + "title": "Config" + }, + "created_at": { + "description": "Creation time, ISO 8601.", + "format": "date-time", + "title": "Created At", + "type": "string" + }, + "updated_at": { + "description": "Last update time, ISO 8601.", + "format": "date-time", + "title": "Updated At", + "type": "string" } }, "required": [ - "harness", - "team_id" + "id", + "team_id", + "config", + "created_at", + "updated_at" ], - "title": "RouteHarnessConfigTombstoneV1", + "title": "RouteConnectionV1", "type": "object" }, - "ClearRouteHarnessConfigRequestV1": { + "RouteConnectionXAIV1": { "properties": { - "team_id": { - "description": "Identifier of the team whose default models to clear.", - "minLength": 1, - "title": "Team Id", + "provider": { + "const": "XAI", + "description": "Provider kind for xAI.", + "title": "Provider", + "type": "string" + }, + "secret_id": { + "description": "Identifier of the team secret holding the provider API key.", + "title": "Secret Id", + "type": "string" + }, + "secret_name": { + "description": "Name of the team secret holding the provider API key.", + "title": "Secret Name", "type": "string" } }, "required": [ - "team_id" + "provider", + "secret_id", + "secret_name" ], - "title": "ClearRouteHarnessConfigRequestV1", + "title": "RouteConnectionXAIV1", "type": "object" }, - "RouteTombstoneV1": { + "RouteConnectionsResponseV1": { "properties": { - "id": { - "description": "Stable identifier of the deleted route.", - "title": "Id", - "type": "string" + "items": { + "description": "Items in this page.", + "items": { + "$ref": "#/components/schemas/RouteConnectionV1" + }, + "title": "Items", + "type": "array" }, - "name": { - "description": "Name of the deleted route.", - "title": "Name", - "type": "string" + "pagination": { + "$ref": "#/components/schemas/PaginationResponseV1", + "description": "Pagination metadata for the page." } }, "required": [ - "id", - "name" + "items", + "pagination" ], - "title": "RouteTombstoneV1", + "title": "RouteConnectionsResponseV1", "type": "object" }, - "UpdateRouteRequestV1": { - "additionalProperties": false, - "minProperties": 1, + "RouteConnectionsRequestV1": { "properties": { - "description": { + "cursor": { "anyOf": [ { - "maxLength": 1000, "type": "string" }, { @@ -32781,17 +33571,20 @@ } ], "default": null, - "description": "New description. Omit to keep the current description; use an empty string to clear it. Null is not accepted.", - "examples": [ - "Assistant for code review and debugging." - ], - "title": "Description" + "description": "Opaque cursor returned by a previous page. Omit to fetch the first page.", + "title": "Cursor" }, - "display_name": { + "limit": { + "default": 100, + "description": "Maximum number of items to return.", + "maximum": 1000, + "minimum": 1, + "title": "Limit", + "type": "integer" + }, + "team_id": { "anyOf": [ { - "maxLength": 255, - "minLength": 1, "type": "string" }, { @@ -32799,26 +33592,89 @@ } ], "default": null, - "description": "New display label. Omit to keep the current label; null is not accepted.", - "examples": [ - "Assistant" - ], - "title": "Display Name" + "description": "Identifier of the team whose connections to list.", + "title": "Team Id" } }, - "title": "UpdateRouteRequestV1", + "title": "RouteConnectionsRequestV1", "type": "object" }, - "EndpointTargetV1": { - "description": "One configured upstream target of an endpoint.", + "RouteConnectionConfigAnthropicV1": { + "additionalProperties": false, "properties": { "provider": { - "$ref": "#/components/schemas/GatewayProvider", - "description": "Upstream provider." + "const": "ANTHROPIC", + "description": "Provider kind for Anthropic.", + "title": "Provider", + "type": "string" + }, + "secret_id": { + "description": "Identifier of an existing secret, owned by the same team, that holds the provider API key.", + "minLength": 1, + "title": "Secret Id", + "type": "string" + } + }, + "required": [ + "provider", + "secret_id" + ], + "title": "RouteConnectionConfigAnthropicV1", + "type": "object" + }, + "RouteConnectionConfigOpenAIV1": { + "additionalProperties": false, + "properties": { + "provider": { + "const": "OPENAI", + "description": "Provider kind for OpenAI.", + "title": "Provider", + "type": "string" + }, + "secret_id": { + "description": "Identifier of an existing secret, owned by the same team, that holds the provider API key.", + "minLength": 1, + "title": "Secret Id", + "type": "string" + } + }, + "required": [ + "provider", + "secret_id" + ], + "title": "RouteConnectionConfigOpenAIV1", + "type": "object" + }, + "RouteConnectionConfigXAIV1": { + "additionalProperties": false, + "properties": { + "provider": { + "const": "XAI", + "description": "Provider kind for xAI.", + "title": "Provider", + "type": "string" }, "secret_id": { + "description": "Identifier of an existing secret, owned by the same team, that holds the provider API key.", + "minLength": 1, + "title": "Secret Id", + "type": "string" + } + }, + "required": [ + "provider", + "secret_id" + ], + "title": "RouteConnectionConfigXAIV1", + "type": "object" + }, + "CreateRouteConnectionRequestV1": { + "additionalProperties": false, + "properties": { + "team_id": { "anyOf": [ { + "minLength": 1, "type": "string" }, { @@ -32826,23 +33682,330 @@ } ], "default": null, - "description": "Referenced secret, if any.", - "title": "Secret Id" + "description": "Identifier of the team that owns the connection.", + "examples": [ + "abc1234" + ], + "title": "Team Id" }, - "target_model": { - "anyOf": [ - { - "type": "string" + "config": { + "description": "Provider the connection authenticates with, and the team secret holding its API key.", + "discriminator": { + "mapping": { + "ANTHROPIC": "#/components/schemas/RouteConnectionConfigAnthropicV1", + "OPENAI": "#/components/schemas/RouteConnectionConfigOpenAIV1", + "XAI": "#/components/schemas/RouteConnectionConfigXAIV1" }, + "propertyName": "provider" + }, + "examples": [ { - "type": "null" + "provider": "ANTHROPIC", + "secret_id": "abc1234" } ], - "default": null, - "description": "Upstream model name, if any.", - "title": "Target Model" - }, - "base_url": { + "oneOf": [ + { + "$ref": "#/components/schemas/RouteConnectionConfigAnthropicV1" + }, + { + "$ref": "#/components/schemas/RouteConnectionConfigOpenAIV1" + }, + { + "$ref": "#/components/schemas/RouteConnectionConfigXAIV1" + } + ], + "title": "Config" + } + }, + "required": [ + "config" + ], + "title": "CreateRouteConnectionRequestV1", + "type": "object" + }, + "RouteConnectionTombstoneV1": { + "properties": { + "id": { + "description": "Stable identifier of the deleted connection.", + "title": "Id", + "type": "string" + } + }, + "required": [ + "id" + ], + "title": "RouteConnectionTombstoneV1", + "type": "object" + }, + "UpdateRouteConnectionConfigAnthropicV1": { + "additionalProperties": false, + "properties": { + "provider": { + "const": "ANTHROPIC", + "description": "Provider kind for Anthropic.", + "title": "Provider", + "type": "string" + }, + "secret_id": { + "anyOf": [ + { + "minLength": 1, + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Identifier of the new secret, owned by the connection's team, that holds the provider API key. Omit to keep the current secret.", + "title": "Secret Id", + "x-null-distinct": true + } + }, + "required": [ + "provider" + ], + "title": "UpdateRouteConnectionConfigAnthropicV1", + "type": "object" + }, + "UpdateRouteConnectionConfigOpenAIV1": { + "additionalProperties": false, + "properties": { + "provider": { + "const": "OPENAI", + "description": "Provider kind for OpenAI.", + "title": "Provider", + "type": "string" + }, + "secret_id": { + "anyOf": [ + { + "minLength": 1, + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Identifier of the new secret, owned by the connection's team, that holds the provider API key. Omit to keep the current secret.", + "title": "Secret Id", + "x-null-distinct": true + } + }, + "required": [ + "provider" + ], + "title": "UpdateRouteConnectionConfigOpenAIV1", + "type": "object" + }, + "UpdateRouteConnectionConfigXAIV1": { + "additionalProperties": false, + "properties": { + "provider": { + "const": "XAI", + "description": "Provider kind for xAI.", + "title": "Provider", + "type": "string" + }, + "secret_id": { + "anyOf": [ + { + "minLength": 1, + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Identifier of the new secret, owned by the connection's team, that holds the provider API key. Omit to keep the current secret.", + "title": "Secret Id", + "x-null-distinct": true + } + }, + "required": [ + "provider" + ], + "title": "UpdateRouteConnectionConfigXAIV1", + "type": "object" + }, + "UpdateRouteConnectionRequestV1": { + "additionalProperties": false, + "properties": { + "config": { + "description": "Connection fields to change. The provider must match the connection and is immutable.", + "discriminator": { + "mapping": { + "ANTHROPIC": "#/components/schemas/UpdateRouteConnectionConfigAnthropicV1", + "OPENAI": "#/components/schemas/UpdateRouteConnectionConfigOpenAIV1", + "XAI": "#/components/schemas/UpdateRouteConnectionConfigXAIV1" + }, + "propertyName": "provider" + }, + "examples": [ + { + "provider": "ANTHROPIC", + "secret_id": "abc1234" + } + ], + "oneOf": [ + { + "$ref": "#/components/schemas/UpdateRouteConnectionConfigAnthropicV1" + }, + { + "$ref": "#/components/schemas/UpdateRouteConnectionConfigOpenAIV1" + }, + { + "$ref": "#/components/schemas/UpdateRouteConnectionConfigXAIV1" + } + ], + "title": "Config" + } + }, + "required": [ + "config" + ], + "title": "UpdateRouteConnectionRequestV1", + "type": "object" + }, + "RouteTombstoneV1": { + "properties": { + "id": { + "description": "Stable identifier of the deleted route.", + "title": "Id", + "type": "string" + }, + "name": { + "description": "Name of the deleted route.", + "title": "Name", + "type": "string" + } + }, + "required": [ + "id", + "name" + ], + "title": "RouteTombstoneV1", + "type": "object" + }, + "UpdateRouteRequestV1": { + "additionalProperties": false, + "minProperties": 1, + "properties": { + "description": { + "anyOf": [ + { + "maxLength": 1000, + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "New description. Omit to keep the current description; use an empty string to clear it. Null is not accepted.", + "examples": [ + "Assistant for code review and debugging." + ], + "title": "Description" + }, + "display_name": { + "anyOf": [ + { + "maxLength": 255, + "minLength": 1, + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "New display label. Omit to keep the current label; null is not accepted.", + "examples": [ + "Assistant" + ], + "title": "Display Name" + }, + "target": { + "anyOf": [ + { + "discriminator": { + "mapping": { + "ANTHROPIC": "#/components/schemas/RouteTargetConfigAnthropicV1", + "BASETEN_MODEL_API": "#/components/schemas/RouteTargetConfigBasetenModelAPIV1", + "CLASSIFIER_MODEL_BASED": "#/components/schemas/RouteTargetConfigClassifierModelBasedV1", + "OPENAI": "#/components/schemas/RouteTargetConfigOpenAIV1", + "XAI": "#/components/schemas/RouteTargetConfigXAIV1" + }, + "propertyName": "type" + }, + "oneOf": [ + { + "$ref": "#/components/schemas/RouteTargetConfigBasetenModelAPIV1" + }, + { + "$ref": "#/components/schemas/RouteTargetConfigAnthropicV1" + }, + { + "$ref": "#/components/schemas/RouteTargetConfigOpenAIV1" + }, + { + "$ref": "#/components/schemas/RouteTargetConfigXAIV1" + }, + { + "$ref": "#/components/schemas/RouteTargetConfigClassifierModelBasedV1" + } + ] + }, + { + "type": "null" + } + ], + "description": "Complete new target configuration. Omit to keep the current configuration; null is not accepted. Only the CLASSIFIER_MODEL_BASED configuration of a router route is mutable.", + "title": "Target", + "x-null-distinct": true + } + }, + "title": "UpdateRouteRequestV1", + "type": "object" + }, + "EndpointTargetV1": { + "description": "One configured upstream target of an endpoint.", + "properties": { + "provider": { + "$ref": "#/components/schemas/GatewayProvider", + "description": "Upstream provider." + }, + "secret_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Referenced secret, if any.", + "title": "Secret Id" + }, + "target_model": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Upstream model name, if any.", + "title": "Target Model" + }, + "base_url": { "anyOf": [ { "type": "string" @@ -33721,163 +34884,384 @@ "description": "Per-model rate and usage limit configuration.", "title": "Models" } - }, - "title": "UpdateGroupRequestV1", - "type": "object" + }, + "title": "UpdateGroupRequestV1", + "type": "object" + }, + "GatewayKeyInfoV1": { + "properties": { + "prefix": { + "description": "The prefix of the Model API key.", + "title": "Prefix", + "type": "string" + }, + "name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Optional display name.", + "title": "Name" + } + }, + "required": [ + "prefix" + ], + "title": "GatewayKeyInfoV1", + "type": "object" + }, + "KeysForGroupResponseV1": { + "properties": { + "items": { + "description": "Items in this page.", + "items": { + "$ref": "#/components/schemas/GatewayKeyInfoV1" + }, + "title": "Items", + "type": "array" + }, + "pagination": { + "$ref": "#/components/schemas/PaginationResponseV1", + "description": "Pagination metadata for the page." + } + }, + "required": [ + "items", + "pagination" + ], + "title": "KeysForGroupResponseV1", + "type": "object" + }, + "CreateApiKeyForGroupRequestV1": { + "properties": { + "name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Optional display name for the new key.", + "examples": [ + "prod-key-1" + ], + "title": "Name" + } + }, + "title": "CreateApiKeyForGroupRequestV1", + "type": "object" + }, + "CreateApiKeyForGroupResponseV1": { + "properties": { + "api_key": { + "description": "Plaintext key string, returned exactly once.", + "title": "Api Key", + "type": "string" + }, + "prefix": { + "description": "Key prefix (the part before the dot).", + "title": "Prefix", + "type": "string" + }, + "name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Display name of the key.", + "title": "Name" + } + }, + "required": [ + "api_key", + "prefix" + ], + "title": "CreateApiKeyForGroupResponseV1", + "type": "object" + }, + "RegisterAPIKeyRequestV1": { + "description": "Request to register a caller-supplied API key against an existing FederatedGroup.", + "properties": { + "name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Optional name for the Model API key", + "examples": [ + "my-model-api-key" + ], + "title": "Name" + }, + "key": { + "description": "Value of the API key to register", + "examples": [ + "my-secure-api-key-value" + ], + "title": "Key", + "type": "string" + } + }, + "required": [ + "key" + ], + "title": "RegisterAPIKeyRequestV1", + "type": "object" + }, + "RegisterAPIKeyResponseV1": { + "properties": { + "ok": { + "description": "Whether the registration was successful", + "title": "Ok", + "type": "boolean" + } + }, + "required": [ + "ok" + ], + "title": "RegisterAPIKeyResponseV1", + "type": "object" + }, + "GetSandboxConfigurationResponseV1": { + "type": "object", + "required": [ + "regions" + ], + "properties": { + "regions": { + "type": "array", + "description": "Available Carbon-compatible regions, sorted by name. Empty when none are available.", + "items": { + "$ref": "#/components/schemas/SandboxRegionV1" + } + } + } + }, + "SandboxRegionV1": { + "type": "object", + "required": [ + "name", + "country", + "continent", + "location", + "info_generation" + ], + "properties": { + "name": { + "type": "string", + "description": "Public region identifier to use when creating a sandbox.", + "example": "us-was-1" + }, + "country": { + "type": "string", + "description": "Country code.", + "example": "us" + }, + "continent": { + "type": "string", + "description": "Continent code.", + "example": "na" + }, + "location": { + "type": "string", + "description": "Region location.", + "example": "Washington" + }, + "info_generation": { + "type": "string", + "description": "Runtime generation supported by this region, using the public name CARBON. Actual runtime selection depends on the team and sandbox configuration.", + "enum": [ + "CARBON" + ] + } + } + }, + "SandboxLogsV1": { + "type": "object", + "required": [ + "sandbox_name", + "start_time", + "end_time", + "logs", + "next_cursor" + ], + "properties": { + "sandbox_name": { + "type": "string" + }, + "start_time": { + "type": "string", + "format": "date-time" + }, + "end_time": { + "type": "string", + "format": "date-time" + }, + "logs": { + "type": "array", + "maxItems": 1000, + "items": { + "$ref": "#/components/schemas/SandboxLogEntryV1" + } + }, + "next_cursor": { + "description": "Opaque cursor for the next page; null when there are no more entries.", + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ] + } + } + }, + "SandboxLogEntryV1": { + "type": "object", + "required": [ + "timestamp", + "message", + "severity", + "trace_id" + ], + "properties": { + "timestamp": { + "type": "string", + "format": "date-time" + }, + "message": { + "type": "string" + }, + "severity": { + "type": "integer", + "description": "Numeric log severity." + }, + "trace_id": { + "type": "string", + "description": "Associated trace identifier, or an empty string when unavailable." + }, + "action": { + "type": "string", + "description": "Action or command when present on the log entry." + } + } + }, + "SandboxMetricsV1": { + "type": "object", + "required": [ + "sandbox_name", + "start_time", + "end_time", + "interval_seconds", + "data" + ], + "properties": { + "sandbox_name": { + "type": "string" + }, + "start_time": { + "type": "string", + "format": "date-time" + }, + "end_time": { + "type": "string", + "format": "date-time" + }, + "interval_seconds": { + "type": "integer" + }, + "data": { + "type": "array", + "maxItems": 10000, + "items": { + "$ref": "#/components/schemas/SandboxMetricsPointV1" + } + } + } }, - "GatewayKeyInfoV1": { + "SandboxMetricsPointV1": { + "type": "object", + "required": [ + "timestamp", + "requests", + "cpu_percent", + "memory_bytes", + "error_rate" + ], "properties": { - "prefix": { - "description": "The prefix of the Model API key.", - "title": "Prefix", - "type": "string" + "timestamp": { + "type": "string", + "format": "date-time", + "description": "Start of this interval." }, - "name": { + "requests": { + "description": "Request count. Zero without traffic; null for intervals entirely before creation.", "anyOf": [ { - "type": "string" + "type": "number" }, { "type": "null" } - ], - "default": null, - "description": "Optional display name.", - "title": "Name" - } - }, - "required": [ - "prefix" - ], - "title": "GatewayKeyInfoV1", - "type": "object" - }, - "KeysForGroupResponseV1": { - "properties": { - "items": { - "description": "Items in this page.", - "items": { - "$ref": "#/components/schemas/GatewayKeyInfoV1" - }, - "title": "Items", - "type": "array" + ] }, - "pagination": { - "$ref": "#/components/schemas/PaginationResponseV1", - "description": "Pagination metadata for the page." - } - }, - "required": [ - "items", - "pagination" - ], - "title": "KeysForGroupResponseV1", - "type": "object" - }, - "CreateApiKeyForGroupRequestV1": { - "properties": { - "name": { + "cpu_percent": { + "description": "Peak CPU usage in this interval, where 100 is one full CPU core. Null without a sample.", "anyOf": [ { - "type": "string" + "type": "number" }, { "type": "null" } - ], - "default": null, - "description": "Optional display name for the new key.", - "examples": [ - "prod-key-1" - ], - "title": "Name" - } - }, - "title": "CreateApiKeyForGroupRequestV1", - "type": "object" - }, - "CreateApiKeyForGroupResponseV1": { - "properties": { - "api_key": { - "description": "Plaintext key string, returned exactly once.", - "title": "Api Key", - "type": "string" - }, - "prefix": { - "description": "Key prefix (the part before the dot).", - "title": "Prefix", - "type": "string" + ] }, - "name": { + "memory_bytes": { + "description": "Memory usage in bytes, averaged per series then maximum across series. Null without a sample.", "anyOf": [ { - "type": "string" + "type": "number" }, { "type": "null" } - ], - "default": null, - "description": "Display name of the key.", - "title": "Name" - } - }, - "required": [ - "api_key", - "prefix" - ], - "title": "CreateApiKeyForGroupResponseV1", - "type": "object" - }, - "RegisterAPIKeyRequestV1": { - "description": "Request to register a caller-supplied API key against an existing FederatedGroup.", - "properties": { - "name": { + ] + }, + "error_rate": { + "description": "Fraction of requests returning 4xx or 5xx. Null when there are no requests.", "anyOf": [ { - "type": "string" + "type": "number", + "minimum": 0, + "maximum": 1 }, { "type": "null" } - ], - "default": null, - "description": "Optional name for the Model API key", - "examples": [ - "my-model-api-key" - ], - "title": "Name" - }, - "key": { - "description": "Value of the API key to register", - "examples": [ - "my-secure-api-key-value" - ], - "title": "Key", - "type": "string" - } - }, - "required": [ - "key" - ], - "title": "RegisterAPIKeyRequestV1", - "type": "object" - }, - "RegisterAPIKeyResponseV1": { - "properties": { - "ok": { - "description": "Whether the registration was successful", - "title": "Ok", - "type": "boolean" + ] } - }, - "required": [ - "ok" - ], - "title": "RegisterAPIKeyResponseV1", - "type": "object" + } }, - "ImageBuildLogV1": { + "SandboxImageBuildLogV1": { "type": "object", "required": [ "timestamp", @@ -33898,7 +35282,7 @@ } } }, - "ImageBuildLogsResponseV1": { + "SandboxImageBuildLogsResponseV1": { "type": "object", "required": [ "logs", @@ -33908,7 +35292,7 @@ "logs": { "type": "array", "items": { - "$ref": "#/components/schemas/ImageBuildLogV1" + "$ref": "#/components/schemas/SandboxImageBuildLogV1" } }, "total_count": { @@ -33947,7 +35331,11 @@ ] }, "terminated_retention": { - "type": "string", + "allOf": [ + { + "$ref": "#/components/schemas/SandboxDurationV1" + } + ], "description": "Duration to keep the sandbox record after termination for log access (e.g., '1h', '24h', '7d'). Defaults to 5m. Subject to maximum quota limits.", "example": "24h" } @@ -33997,7 +35385,7 @@ }, "SandboxDurationV1": { "type": "string", - "description": "Duration using seconds, minutes, hours, or composite durations such as 1h30m. Whole days and weeks are also supported, for example 7d or 2w.", + "description": "Duration using seconds, minutes, hours, or composite durations such as 1h30m. Whole days and weeks are also supported, for example 7d or 2w, where d is 24h and w is 7 \u00d7 24h. Days and weeks cannot be combined with other units, so 1d12h is rejected; use 36h instead. Values are returned exactly as sent, without normalization.", "pattern": "^[+-]?(0|[0-9]+[dw]|([0-9]+(\\.[0-9]*)?|\\.[0-9]+)(ns|us|\u00b5s|\u03bcs|ms|s|m|h)(([0-9]+(\\.[0-9]*)?|\\.[0-9]+)(ns|us|\u00b5s|\u03bcs|ms|s|m|h))*)$", "example": "24h" }, @@ -34418,12 +35806,6 @@ "type": "object", "description": "Writable sandbox configuration. Fields are serialized at the root of the request or resource.", "properties": { - "enabled": { - "type": "boolean", - "description": "When false, the sandbox is disabled and will not accept connections", - "example": true, - "default": true - }, "lifecycle": { "description": "Lifecycle configuration controlling automatic sandbox deletion based on idle time, max age, or specific dates", "allOf": [ @@ -34566,7 +35948,6 @@ } }, "example": { - "enabled": true, "lifecycle": { "expiration_policies": [ { @@ -34896,13 +36277,8 @@ }, "UpdateSandboxRequestV1": { "type": "object", - "description": "Partial sandbox update. Omitted fields remain unchanged. Supplied arrays and maps (including labels) replace their previous values; supplied structured objects update only their supplied fields. Null is not accepted. The name, memory, and network configuration are immutable after creation. Supplying memory or network returns 400, including unchanged, empty, or null values.", + "description": "Partial sandbox update. Omitted fields remain unchanged. Supplied arrays and maps (including labels) replace their previous values; supplied structured objects update only their supplied fields. Null is not accepted. The name, memory, network, region, image, and ports are immutable after creation. Supplying any of these fields returns 400, including unchanged, empty, or null values.", "properties": { - "enabled": { - "type": "boolean", - "description": "When false, the sandbox is disabled and will not accept connections", - "example": true - }, "lifecycle": { "description": "Lifecycle configuration controlling automatic sandbox deletion based on idle time, max age, or specific dates", "allOf": [ @@ -34931,11 +36307,6 @@ "terminated_retention": "24h" } }, - "region": { - "type": "string", - "description": "Region where the sandbox runs (for example us-pdx-1 or eu-lon-1). When omitted at creation, the closest region is selected.", - "example": "us-pdx-1" - }, "envs": { "type": "array", "description": "Environment variables injected into the sandbox.", @@ -34955,26 +36326,6 @@ } ] }, - "image": { - "type": "string", - "description": "Image reference including its tag. Built-in image references are returned in the canonical baseten/ namespace. Use baseten/base-image:latest to get started with the built-in sandbox execution API. This image is available directly without building, pushing, or listing images through GET /v1/sandboxes/images.", - "example": "baseten/base-image:latest" - }, - "ports": { - "description": "Set of ports for a resource", - "allOf": [ - { - "$ref": "#/components/schemas/SandboxPortsV1" - } - ], - "example": [ - { - "name": "http", - "protocol": "HTTP", - "target": 3000 - } - ] - }, "external_id": { "type": "string", "description": "Caller-owned identifier for external lookups. Max 64 chars, alphanumeric + dash.", @@ -34999,7 +36350,6 @@ }, "minProperties": 1, "example": { - "enabled": true, "lifecycle": { "expiration_policies": [ { @@ -35020,7 +36370,6 @@ ], "terminated_retention": "24h" }, - "region": "us-pdx-1", "envs": [ { "name": "NODE_ENV", @@ -35033,14 +36382,6 @@ "value": "3000" } ], - "image": "baseten/base-image:latest", - "ports": [ - { - "name": "http", - "protocol": "HTTP", - "target": 3000 - } - ], "external_id": "api-review-20260916-001", "labels": { "env": "development", @@ -35143,7 +36484,6 @@ } ], "example": { - "enabled": true, "lifecycle": { "expiration_policies": [ { @@ -35290,7 +36630,6 @@ }, "example": [ { - "enabled": true, "lifecycle": { "expiration_policies": [ { @@ -35402,7 +36741,6 @@ "example": { "items": [ { - "enabled": true, "lifecycle": { "expiration_policies": [ { @@ -35499,7 +36837,7 @@ } } }, - "ListImagesResponseV1": { + "ListSandboxImagesResponseV1": { "type": "object", "description": "One page of image repository summaries. Fetch tags through the separate tag listing endpoint.", "properties": { @@ -35507,7 +36845,7 @@ "type": "array", "description": "Image repositories on this page.", "items": { - "$ref": "#/components/schemas/ImageV1" + "$ref": "#/components/schemas/SandboxImageSummaryV1" } }, "pagination": { @@ -35519,14 +36857,14 @@ "pagination" ] }, - "ListImageTagsResponseV1": { + "ListSandboxImageTagsResponseV1": { "type": "object", "description": "One page of image tags.", "properties": { "items": { "type": "array", "items": { - "$ref": "#/components/schemas/ImageTagV1" + "$ref": "#/components/schemas/SandboxImageTagV1" } }, "pagination": { @@ -35538,7 +36876,7 @@ "pagination" ] }, - "ImageStatusV1": { + "SandboxImageStatusV1": { "type": "string", "description": "Image processing status. Only BUILT images are ready to use.", "enum": [ @@ -35550,7 +36888,7 @@ "readOnly": true, "example": "BUILT" }, - "ImageTagV1": { + "SandboxImageTagV1": { "type": "object", "description": "A tag identifying a version of a sandbox image.", "properties": { @@ -35592,9 +36930,9 @@ "size": 134217728 } }, - "ImageV1": { + "SandboxImageSummaryV1": { "type": "object", - "description": "Sandbox image repository. List and get operations return a summary without embedded tags.", + "description": "Sandbox image repository summary. Fetch tags through GET /sandboxes/images/{image_name}/tags.", "properties": { "name": { "type": "string", @@ -35605,7 +36943,7 @@ "description": "Image processing status. Only BUILT images are ready to use.", "allOf": [ { - "$ref": "#/components/schemas/ImageStatusV1" + "$ref": "#/components/schemas/SandboxImageStatusV1" } ], "example": "BUILT" @@ -35646,33 +36984,11 @@ "readOnly": true, "description": "Number of image versions in the repository.", "example": 2 - }, - "tags": { - "type": "array", - "description": "Empty for list and get summary responses. Use GET /sandboxes/images/{image_name}/tags to retrieve paginated image versions.", - "items": { - "$ref": "#/components/schemas/ImageTagV1" - }, - "example": [ - { - "name": "latest", - "created_at": "2026-09-16T21:20:00Z", - "updated_at": "2026-09-16T21:25:00Z", - "size": 134217728 - }, - { - "name": "20260915212000", - "created_at": "2026-09-15T21:20:00Z", - "updated_at": "2026-09-15T21:25:00Z", - "size": 125829120 - } - ] } }, "required": [ "name", - "status", - "tags" + "status" ], "example": { "name": "base-image", @@ -35681,23 +36997,44 @@ "updated_at": "2026-09-16T21:25:00Z", "last_deployed_at": "2026-09-16T21:26:58.545765901Z", "size": 260046848, - "tags": [ - { - "name": "latest", - "created_at": "2026-09-16T21:20:00Z", - "updated_at": "2026-09-16T21:25:00Z", - "size": 134217728 + "tag_count": 2 + } + }, + "SandboxImageV1": { + "description": "Sandbox image repository. Get operations return a summary with an empty tags array.", + "allOf": [ + { + "$ref": "#/components/schemas/SandboxImageSummaryV1" + }, + { + "type": "object", + "properties": { + "tags": { + "type": "array", + "description": "Empty for get summary responses. Use GET /sandboxes/images/{image_name}/tags to retrieve paginated image versions.", + "items": { + "$ref": "#/components/schemas/SandboxImageTagV1" + }, + "example": [] + } }, - { - "name": "20260915212000", - "created_at": "2026-09-15T21:20:00Z", - "updated_at": "2026-09-15T21:25:00Z", - "size": 125829120 - } - ] + "required": [ + "tags" + ] + } + ], + "example": { + "name": "base-image", + "status": "BUILT", + "created_at": "2026-09-15T21:20:00Z", + "updated_at": "2026-09-16T21:25:00Z", + "last_deployed_at": "2026-09-16T21:26:58.545765901Z", + "size": 260046848, + "tag_count": 2, + "tags": [] } }, - "PushImageRequestV1": { + "PushSandboxImageRequestV1": { "type": "object", "description": "Push a sandbox image from a source archive or an existing registry image.", "properties": { @@ -35728,7 +37065,7 @@ "docker_config": "{\"auths\":{\"https://index.docker.io/v1/\":{\"auth\":\"YjEwLXJldmlldzpkZW1vLW5vdC1hLXZhbGlkLXJlZ2lzdHJ5LXRva2Vu\"}}}" } }, - "PushImageResponseV1": { + "PushSandboxImageResponseV1": { "type": "object", "description": "Accepted image push. Acceptance does not imply readiness; poll GET /sandboxes/images/{image_name} until status is BUILT or FAILED.", "properties": { @@ -35741,7 +37078,7 @@ "description": "Image processing status. Only BUILT images are ready to use.", "allOf": [ { - "$ref": "#/components/schemas/ImageStatusV1" + "$ref": "#/components/schemas/SandboxImageStatusV1" } ], "example": "BUILDING" @@ -35929,7 +37266,7 @@ "items" ] }, - "CleanupImagesResponseV1": { + "CleanupSandboxImagesResponseV1": { "type": "object", "description": "Result of cleaning up unused sandbox images.", "properties": { @@ -36148,11 +37485,11 @@ "in": "path", "required": true }, - "harness": { + "connection_id": { "schema": { "type": "string" }, - "name": "harness", + "name": "connection_id", "in": "path", "required": true }, @@ -36211,7 +37548,7 @@ "example": "baseten-api-review-0916" } }, - "ImageName": { + "SandboxImageName": { "name": "image_name", "in": "path", "required": true, @@ -36466,7 +37803,6 @@ "UpdateSandbox": { "summary": "Update every mutable field", "value": { - "enabled": true, "lifecycle": { "expiration_policies": [ { @@ -36487,7 +37823,6 @@ ], "terminated_retention": "24h" }, - "region": "us-pdx-1", "envs": [ { "name": "NODE_ENV", @@ -36500,14 +37835,6 @@ "value": "3000" } ], - "image": "baseten/base-image:latest", - "ports": [ - { - "name": "http", - "protocol": "HTTP", - "target": 3000 - } - ], "external_id": "api-review-20260916-001", "labels": { "env": "development", @@ -36521,7 +37848,6 @@ "SandboxReady": { "summary": "Running sandbox with complete configuration", "value": { - "enabled": true, "lifecycle": { "expiration_policies": [ { @@ -36616,7 +37942,6 @@ "SandboxUpdated": { "summary": "Updated sandbox with complete configuration", "value": { - "enabled": true, "lifecycle": { "expiration_policies": [ { @@ -36711,7 +38036,6 @@ "SandboxDeleting": { "summary": "Sandbox deletion accepted", "value": { - "enabled": true, "lifecycle": { "expiration_policies": [ { @@ -36807,7 +38131,6 @@ "value": { "items": [ { - "enabled": true, "lifecycle": { "expiration_policies": [ { @@ -36910,7 +38233,6 @@ "value": { "items": [ { - "enabled": true, "lifecycle": { "expiration_policies": [ { @@ -37073,8 +38395,7 @@ "updated_at": "2026-09-16T21:25:00Z", "last_deployed_at": "2026-09-16T21:26:58.545765901Z", "size": 260046848, - "tag_count": 2, - "tags": [] + "tag_count": 2 } ], "pagination": { diff --git a/scripts/apigen/specs/sandbox.yml b/scripts/apigen/specs/sandbox.yml index f669352..f633299 100644 --- a/scripts/apigen/specs/sandbox.yml +++ b/scripts/apigen/specs/sandbox.yml @@ -12,6 +12,8 @@ paths: The sandbox is quiesced first: the process list is saved (unless saveProcesses is false), every process is stopped, and the API then refuses the calls that would write to the filesystem. The freeze is not lifted afterwards, since an exported sandbox is meant to be restored elsewhere; call POST /archive/resume to lift it. + A sandbox still frozen by an earlier export is exported as it is, with the process list that export saved. + Use dryRun to get the archive's content and exact size without stopping anything and without uploading. Set async to start the export and answer immediately, which is what archiving a large filesystem needs: the export then reports itself through GET /archive/status. @@ -42,7 +44,7 @@ paths: schema: $ref: "#/components/schemas/ErrorResponse" "409": - description: An export is already in progress + description: An export is already in progress, or the sandbox is frozen for a restore content: application/json: schema: @@ -377,7 +379,7 @@ paths: - system "/filesystem-content-search/{path}": get: - description: Searches for text content inside files using ripgrep. Returns matching lines with context. + description: Searches for text content inside files. Returns each matching line, with the lines around it when contextLines is set. parameters: - description: Directory path to search in in: path @@ -411,6 +413,11 @@ paths: name: excludeDirs schema: type: string + - description: "Lines to include before and after each match in its context field (default: 0, max: 20; invalid values count as 0)" + in: query + name: contextLines + schema: + type: integer responses: "200": description: Content search results @@ -606,7 +613,7 @@ paths: - filesystem "/filesystem-multipart/{uploadId}/part": put: - description: Upload a single part of a multipart upload + description: Upload a single part of a multipart upload. Re-uploading a part number replaces that part. Wait for the previous request for that part to finish before retrying. parameters: - description: Upload ID in: path @@ -739,7 +746,10 @@ paths: - filesystem "/filesystem-search/{path}": get: - description: Performs fuzzy search on filesystem paths using fuzzy matching algorithm. Optimized alternative to find and grep commands. + description: >- + Ranks the files and directories under a path by how well their relative path fuzzy-matches `query` (fzf algorithm: the query's characters must appear in order, not necessarily next to each other), best match first. + + Fuzzy search is for "jump to file" lookups from a partial name. The `patterns` parameter is currently ignored by this endpoint; use find for exact glob filtering. parameters: - description: Path to search in (e.g., /home/user/projects) in: path @@ -747,12 +757,17 @@ paths: required: true schema: type: string + - description: Fuzzy pattern matched against each relative path (e.g., mngo for src/main.go). When omitted, the search path itself is used as the pattern. + in: query + name: query + schema: + type: string - description: "Maximum number of results to return (default: 20)" in: query name: maxResults schema: type: integer - - description: Comma-separated file patterns to include (e.g., *.go,*.js) + - description: Accepted for compatibility but currently ignored; use filesystem-find for glob filtering in: query name: patterns schema: @@ -841,7 +856,7 @@ paths: get: description: Get content of a file or listing of a directory. Use Accept header to control response format for files. parameters: - - description: File or directory path + - description: "File or directory path. Relative paths are resolved against the filesystem working directory. For example, with a working directory of /app, GET /filesystem/workspace resolves to /app/workspace. To access the absolute path /workspace over HTTP, encode its leading slash as %2F: GET /filesystem/%2Fworkspace. The double-slash form GET /filesystem//workspace also addresses /workspace." in: path name: path required: true @@ -859,17 +874,17 @@ paths: application/json: schema: oneOf: - - $ref: "#/components/schemas/Directory" - - $ref: "#/components/schemas/FileWithContent" - - type: string - format: binary + - $ref: '#/components/schemas/Directory' + - $ref: '#/components/schemas/FileWithContent' + - type: string + format: binary application/octet-stream: schema: oneOf: - - $ref: "#/components/schemas/Directory" - - $ref: "#/components/schemas/FileWithContent" - - type: string - format: binary + - $ref: '#/components/schemas/Directory' + - $ref: '#/components/schemas/FileWithContent' + - type: string + format: binary "404": description: File or directory not found content: @@ -900,8 +915,47 @@ paths: summary: Get file or directory information tags: - filesystem + head: + description: Returns the metadata of a file or directory as headers, with no body. This checks stat availability, not permission to read file contents or list a directory. When the path does not exist or its metadata cannot be accessed, the response is an empty 200 without the X-File-Type header. + parameters: + - description: File or directory path + in: path + name: path + required: true + schema: + type: string + responses: + "200": + description: Path metadata (X-File-Type is set only when stat succeeds) + headers: + Content-Length: + description: File size in bytes (files only) + schema: + type: integer + Last-Modified: + description: Modification time (HTTP date) + schema: + type: string + X-File-Mode: + description: Permission bits including sticky, setgid and setuid in octal (e.g., 644 or 1777) + schema: + type: string + X-File-Type: + description: file or directory + schema: + type: string + summary: Stat a file or directory + tags: + - filesystem put: - description: Create or update a file or directory + description: >- + Create or update a file or directory. + + + Idempotent: an existing file is overwritten (truncated, not appended to) and an existing directory is kept, so retrying the same request is safe. + + + Send either a JSON body (FileRequest) or `multipart/form-data` for binary files. Multipart fields, in any order: `file` (required, the file content), `permissions` (optional octal mode such as `0755`, applied when the file is created, default `0644`; an existing file keeps its mode), `path` (optional, ignored: the target is always the URL path). parameters: - description: File or directory path in: path @@ -914,6 +968,23 @@ paths: application/json: schema: $ref: "#/components/schemas/FileRequest" + multipart/form-data: + schema: + type: object + required: + - file + properties: + file: + type: string + format: binary + description: File content + permissions: + type: string + example: "0755" + description: Octal mode applied when the file is created (default 0644); an existing file keeps its mode + path: + type: string + description: 'Ignored: the target is always the URL path' description: File or directory details required: true responses: @@ -1002,11 +1073,7 @@ paths: content: application/json: schema: - oneOf: - - $ref: "#/components/schemas/Directory" - - $ref: "#/components/schemas/FileWithContent" - - type: string - format: binary + $ref: "#/components/schemas/Directory" "400": description: Bad request content: @@ -1029,7 +1096,7 @@ paths: tags: - filesystem put: - description: Create or update multiple files within a directory tree structure + description: "Create or update multiple files within a directory tree structure. Idempotent: existing files are overwritten, so retrying the same request is safe." parameters: - description: Root directory path in: path @@ -1050,11 +1117,7 @@ paths: content: application/json: schema: - oneOf: - - $ref: "#/components/schemas/Directory" - - $ref: "#/components/schemas/FileWithContent" - - type: string - format: binary + $ref: "#/components/schemas/Directory" "400": description: Bad request content: @@ -1294,7 +1357,13 @@ paths: tags: - process post: - description: Execute a command and return process information. If Accept header is text/event-stream, streams logs in SSE format and returns the process response as a final event. + description: >- + Execute a command and return process information. + + + Streaming: with `Accept: application/x-ndjson` (or `Accept: text/event-stream`, kept for compatibility) the response is NDJSON (`Content-Type: application/x-ndjson`), not SSE: one JSON object per line, `{"type": "...", "data": "..."}`. + + `type` is `stdout` or `stderr` (`data` is a raw output chunk, sent as soon as the process writes it, newlines included; if the process finished before any chunk was streamed, its output is sent instead as one event per line, without the newline), `keepalive` (every 5 seconds, no data), `error` (`data` is the message, ends the stream) or `result` (last event, `data` is the ProcessResponse as a JSON string). requestBody: content: application/json: @@ -1309,16 +1378,17 @@ paths: application/json: schema: $ref: "#/components/schemas/ProcessResponse" - text/event-stream: + application/x-ndjson: schema: - $ref: "#/components/schemas/ProcessResponse" + type: string + description: Newline-delimited JSON event objects. Each record has type (stdout, stderr, result, error, or keepalive) and optional string data. Result data is a JSON-encoded ProcessResponse; output data preserves raw chunks. "400": description: Invalid request content: application/json: schema: $ref: "#/components/schemas/ErrorResponse" - text/event-stream: + application/x-ndjson: schema: $ref: "#/components/schemas/ErrorResponse" "422": @@ -1327,7 +1397,7 @@ paths: application/json: schema: $ref: "#/components/schemas/ErrorResponse" - text/event-stream: + application/x-ndjson: schema: $ref: "#/components/schemas/ErrorResponse" "500": @@ -1336,7 +1406,7 @@ paths: application/json: schema: $ref: "#/components/schemas/ErrorResponse" - text/event-stream: + application/x-ndjson: schema: $ref: "#/components/schemas/ErrorResponse" summary: Execute a command @@ -1344,7 +1414,7 @@ paths: - process "/process/{identifier}": delete: - description: Gracefully stop a running process + description: Request graceful termination. Poll GET /process/{identifier} until terminal status confirms the managed process has exited. parameters: - description: Process identifier (PID or name) in: path @@ -1354,7 +1424,7 @@ paths: type: string responses: "200": - description: Process stopped + description: Process stop requested content: application/json: schema: @@ -1407,7 +1477,7 @@ paths: - process "/process/{identifier}/kill": delete: - description: Forcefully kill a running process + description: Request forceful termination. Poll GET /process/{identifier} until terminal status confirms the managed process has exited. parameters: - description: Process identifier (PID or name) in: path @@ -1417,7 +1487,7 @@ paths: type: string responses: "200": - description: Process killed + description: Process kill requested content: application/json: schema: @@ -1483,7 +1553,16 @@ paths: - process "/process/{identifier}/logs/stream": get: - description: Streams the stdout and stderr output of a process in real time, one line per log, prefixed with 'stdout:' or 'stderr:'. Closes when the process exits or the client disconnects. + description: >- + Streams the stdout and stderr output of a process in real time: the output so far, then live output as the process writes it. Closes when the process exits or the client disconnects. + + By default, output is plain text with `stdout:` or `stderr:` prefixes at the start of each source stream's lines. Partial lines (e.g. prompts) are sent without waiting for a newline; their continuations have no new prefix. The other stream is not held back while a line is incomplete, and chronological order across stdout and stderr is not guaranteed. + + This format does not provide unambiguous framing: output from the other stream, or a `[keepalive]` marker sent every 30 seconds, can appear inside an unfinished line. A later continuation can therefore lack a prefix identifying its source. Do not rely on this text stream to reconstruct stdout and stderr separately; request NDJSON for source-preserving streaming, or use GET /process/{identifier}/logs for separate output snapshots. + + With `Accept: application/x-ndjson`, each line is a JSON object with `type` (`stdout`, `stderr`, `keepalive`, `restart`, `truncated`, or `error`) and optional `data`. Output records retain source identity and original bytes, including partial lines and newlines, for both retained history and live output. For `encoding: "base64"`, decode `data` before concatenating bytes per source; this occurs for binary data or a UTF-8 character split across chunks. Chunk boundaries are arbitrary and ordering is collection order, not a strict chronology across streams. + + `keepalive` has no output data. A quiet active stream sends an initial keepalive after attachment so clients can receive the response before sending stdin; subsequent keepalives are sent every 30 seconds. `restart` carries a supervisor restart notice, not process output. `truncated` reports a retention or slow-reader gap and must not be appended to stdout/stderr. `error` reports a streaming failure. The connection ends after the final process exit, including automatic restarts; there is no `result` record. Older processes without structured history return HTTP 409 for NDJSON; their text stream remains available. parameters: - description: Process identifier (PID or name) in: path @@ -1491,31 +1570,46 @@ paths: required: true schema: type: string + - description: Explicit application/x-ndjson opts into structured records; absent or wildcard Accept retains text/plain. Supported explicit media types honor q weights, preferring NDJSON on a tie. + in: header + name: Accept + schema: + type: string responses: "200": - description: Stream of process logs, one line per log (prefixed with stdout:/stderr:) + description: Process output as prefixed text or NDJSON records content: text/plain: schema: type: string + application/x-ndjson: + schema: + type: string + description: Newline-delimited JSON records with type (stdout, stderr, keepalive, restart, truncated, error), optional data, and optional encoding=base64 for non-UTF-8 chunks. Decode each output record before concatenating bytes by source. No result record. + "400": + description: Invalid process identifier + content: + application/json: + schema: + $ref: '#/components/schemas/ErrorResponse' "404": description: Process not found content: - text/plain: + application/json: schema: - $ref: "#/components/schemas/ErrorResponse" - "422": - description: Unprocessable entity + $ref: '#/components/schemas/ErrorResponse' + "409": + description: Structured history unavailable for this older process content: - text/plain: + application/json: schema: - $ref: "#/components/schemas/ErrorResponse" + $ref: '#/components/schemas/ErrorResponse' "500": description: Internal server error content: - text/plain: + application/json: schema: - $ref: "#/components/schemas/ErrorResponse" + $ref: '#/components/schemas/ErrorResponse' summary: Stream process logs in real time tags: - process @@ -1643,14 +1737,20 @@ paths: - system "/watch/filesystem/{path}": get: - description: Streams the path of modified files (one per line) in the given directory. Closes when the client disconnects. + description: >- + Streams change events for a directory until the client disconnects. + + + The body is JSON lines (sent with `Content-Type: text/plain`): one event object per line, e.g. `{"op":"WRITE","name":"main.go","path":"/app/src","error":null}`. `op` is the fsnotify operation (CREATE, WRITE, REMOVE, RENAME or CHMOD, several can be joined with `|`), `name` the base name of the changed entry, `path` the directory containing it, and `error` is always null. A `[keepalive]` line (not JSON) is sent every 30 seconds. + + Only the directory's direct entries are watched. To also watch every subdirectory, including ones created later, end the path with `/**`. The stream reports subsequent changes without an initial snapshot. When a new directory is detected, synthetic CREATE events are emitted for entries already present inside it to cover changes before its watch was registered. parameters: - - description: Ignore patterns (comma-separated) + - description: Comma-separated substrings; events whose full path contains one are skipped in: query name: ignore schema: type: string - - description: Directory path to watch + - description: Directory path to watch (append /** to watch subdirectories) in: path name: path required: true @@ -1658,7 +1758,7 @@ paths: type: string responses: "200": - description: Stream of modified file paths, one per line + description: JSON lines stream of change events content: text/plain: schema: @@ -1797,6 +1897,7 @@ components: example: 10 type: integer context: + description: The matching line with up to contextLines lines before and after it, newline-separated; omitted when contextLines is 0 example: |- previous line current line @@ -2200,6 +2301,7 @@ components: example: amd64 type: string buildTime: + description: Build time in RFC 3339 (UTC), or "unknown" for builds without it example: 2026-01-29T17:36:52Z type: string gitCommit: @@ -2214,6 +2316,7 @@ components: example: linux type: string startedAt: + description: API start time in RFC 3339 example: 2026-01-29T18:45:49Z type: string status: @@ -2341,6 +2444,7 @@ components: ProcessLogs: properties: logs: + description: Concatenation of the returned stdout followed by the returned stderr. Output from the two streams is not interleaved. example: logs output type: string stderr: @@ -2409,6 +2513,7 @@ components: example: ls -la type: string completedAt: + description: Completion time, same format as startedAt. Empty string while the process runs example: Wed, 01 Jan 2023 12:01:00 GMT type: string exitCode: @@ -2437,6 +2542,7 @@ components: example: true type: boolean startedAt: + description: Start time as an HTTP date (RFC 1123, e.g. Wed, 01 Jan 2023 12:00:00 GMT) example: Wed, 01 Jan 2023 12:00:00 GMT type: string status: diff --git a/scripts/docgen/conf.py b/scripts/docgen/conf.py index 3888bd0..fd83b4a 100644 --- a/scripts/docgen/conf.py +++ b/scripts/docgen/conf.py @@ -10,6 +10,8 @@ import inspect import pathlib +from sphinx.util.docstrings import separate_metadata + project = "Baseten Python SDK" html_title = "Baseten SDK" html_theme = "furo" @@ -83,5 +85,24 @@ def _generate_class_pages(app): ) +def _hide_private_constructor( + app, what, name, obj, options, signature, return_annotation +): + """Drop a class's signature when its ``__init__`` is marked ``:meta private:``. + + Autodoc honors the marker on members but not on constructors, whose + signature it shows on the class itself. Users get instances of such classes + from SDK methods rather than constructing them. + """ + if what != "class": + return None + init_doc = inspect.getdoc(obj.__dict__.get("__init__")) or "" + _, metadata = separate_metadata(init_doc) + if "private" in metadata: + return "", None + return None + + def setup(app): app.connect("builder-inited", _generate_class_pages) + app.connect("autodoc-process-signature", _hide_private_constructor) 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/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/client/__init__.py b/tests/client/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_e2e.py b/tests/client/test_e2e.py similarity index 100% rename from tests/test_e2e.py rename to tests/client/test_e2e.py diff --git a/tests/client/test_sandbox.py b/tests/client/test_sandbox.py index 0e7afac..47c9618 100644 --- a/tests/client/test_sandbox.py +++ b/tests/client/test_sandbox.py @@ -100,15 +100,35 @@ def test_raw_returns_response_unread_with_accept() -> None: def test_raw_only_stream_keeps_plain_name() -> None: - fake = FakeTransport(200, raw_content=b"log line\n", content_type="text/plain") + fake = FakeTransport( + 200, + raw_content=b'{"type":"stdout","data":"log line\\n"}\n', + content_type="application/x-ndjson", + ) client = make_sync_client(fake) - response = client.api.get_process_logs_stream(identifier="proc-1") + response = client.api.get_process_logs_stream( + identifier="proc-1", accept="application/x-ndjson" + ) response.read() - assert response.content == b"log line\n" + assert response.content == b'{"type":"stdout","data":"log line\\n"}\n' + response.close() + assert fake.capture.headers["accept"] == "application/x-ndjson" + client.close() + + +def test_raw_only_headers_sends_no_accept() -> None: + fake = FakeTransport(200, raw_content=b"") + client = make_sync_client(fake) + + response = client.api.head_filesystem(path="/app") + response.close() - assert fake.capture.headers["accept"] == "text/plain" + assert fake.capture.method == "HEAD" + assert fake.capture.path == "/filesystem/%2Fapp" + # httpx's default, since the generated client sets none. + assert fake.capture.headers["accept"] == "*/*" client.close() @@ -210,17 +230,17 @@ async def test_async_client_round_trip() -> None: @pytest.mark.asyncio async def test_async_raw_sibling() -> None: fake = FakeTransport( - 200, raw_content=b"event: done", content_type="text/event-stream" + 200, raw_content=b'{"type":"stdout"}\n', content_type="application/x-ndjson" ) client = make_async_client(fake) response = await client.api.post_process_raw( request=baseten.client.sandboxapi.ProcessRequest(command="ls"), - accept="text/event-stream", + accept="application/x-ndjson", ) await response.aread() - assert response.content == b"event: done" + assert response.content == b'{"type":"stdout"}\n' await response.aclose() - assert fake.capture.headers["accept"] == "text/event-stream" + assert fake.capture.headers["accept"] == "application/x-ndjson" await client.close() diff --git a/tests/sandbox/__init__.py b/tests/sandbox/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/sandbox/fixtures/localhost-cert.pem b/tests/sandbox/fixtures/localhost-cert.pem new file mode 100644 index 0000000..7f2bf2f --- /dev/null +++ b/tests/sandbox/fixtures/localhost-cert.pem @@ -0,0 +1,11 @@ +-----BEGIN CERTIFICATE----- +MIIBmzCCAUGgAwIBAgIUby1OBqEhBhIMMJzBoz57jrLwRyswCgYIKoZIzj0EAwIw +FDESMBAGA1UEAwwJbG9jYWxob3N0MCAXDTI2MDkyOTE0MTIwMVoYDzIxMjYwOTA1 +MTQxMjAxWjAUMRIwEAYDVQQDDAlsb2NhbGhvc3QwWTATBgcqhkjOPQIBBggqhkjO +PQMBBwNCAAQCtQ8p/aMvfQo0r4hweQIsHZA2QkIUbKUCsUjPjMyzCaqZIQoYO2d8 +jk0qHlxr2FZ9C8LmvRIBtnunRkIIFwleo28wbTAdBgNVHQ4EFgQUIijxIvQDPQRY +NZ729TnM2QClKYQwHwYDVR0jBBgwFoAUIijxIvQDPQRYNZ729TnM2QClKYQwDwYD +VR0TAQH/BAUwAwEB/zAaBgNVHREEEzARgglsb2NhbGhvc3SHBH8AAAEwCgYIKoZI +zj0EAwIDSAAwRQIgP7gFl+uPbiaaCylc95mKxluCgNieFUuss2pLp0oIADECIQDY +QhK7Qt0zNAU0UFfcI2TCuyWCCfSS7zUBj4bdPWQhfQ== +-----END CERTIFICATE----- diff --git a/tests/sandbox/fixtures/localhost-key.pem b/tests/sandbox/fixtures/localhost-key.pem new file mode 100644 index 0000000..02af64f --- /dev/null +++ b/tests/sandbox/fixtures/localhost-key.pem @@ -0,0 +1,5 @@ +-----BEGIN PRIVATE KEY----- +MIGHAgEAMBMGByqGSM49AgEGCCqGSM49AwEHBG0wawIBAQQgeRm87i+OeHlRArVZ +q90rhm9L+27e+oC9m8hZ5lYRmXKhRANCAAQCtQ8p/aMvfQo0r4hweQIsHZA2QkIU +bKUCsUjPjMyzCaqZIQoYO2d8jk0qHlxr2FZ9C8LmvRIBtnunRkIIFwle +-----END PRIVATE KEY----- diff --git a/tests/sandbox/test_auth.py b/tests/sandbox/test_auth.py new file mode 100644 index 0000000..3969b56 --- /dev/null +++ b/tests/sandbox/test_auth.py @@ -0,0 +1,345 @@ +from __future__ import annotations + +import asyncio +import json +import threading +from collections.abc import AsyncIterator, Callable +from datetime import UTC, datetime, timedelta + +import httpx +import pytest + +import baseten.client.managementapi +import baseten.sandbox._auth +from baseten.sandbox import SandboxTokenProviderContext +from baseten.sandbox._auth import ( + AsyncTokenAuth, + AsyncTokenSource, + TokenAuth, + TokenSource, +) + + +def _revoked(request: httpx.Request | None = None) -> httpx.Response: + return httpx.Response(401, headers={"x-blaxel-error-code": "TOKEN_REVOKED"}) + + +@pytest.fixture(autouse=True) +def no_revoked_retry_delay(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(baseten.sandbox._auth, "_token_revoked_retry_delay", 0.0) + + +class _Server: + """Mints numbered tokens at ``/v1/token`` and answers other requests.""" + + def __init__( + self, + *, + expires_in: timedelta = timedelta(hours=2), + respond: Callable[[httpx.Request], httpx.Response] | None = None, + ) -> None: + self.expires_in = expires_in + self.respond = respond or (lambda request: httpx.Response(200)) + self.mints: list[httpx.Request] = [] + self.requests: list[httpx.Request] = [] + self.authorizations: list[str | None] = [] + self.mint_failures = 0 + self.mint_gate: threading.Event | None = None + """When set, each mint waits for it.""" + + def handle(self, request: httpx.Request) -> httpx.Response: + if request.url.path == "/v1/token": + if self.mint_gate is not None: + self.mint_gate.wait() + self.mints.append(request) + if self.mint_failures > 0: + self.mint_failures -= 1 + return httpx.Response(500, text="mint failed") + return httpx.Response( + 200, + json={ + "token": f"token-{len(self.mints)}", + "expires_at": (datetime.now(UTC) + self.expires_in).isoformat(), + }, + ) + self.requests.append(request) + # A resend reuses the request object, so its header is read now. + self.authorizations.append(request.headers.get("Authorization")) + return self.respond(request) + + def tokens(self, *, api_key: str = "key", provider=None) -> TokenSource: + return TokenSource( + api_key=api_key, + token_provider=provider, + mint_api=baseten.client.managementapi.ApiClient( + httpx.Client(transport=httpx.MockTransport(self.handle)), + base_url="https://api.test", + headers={"Authorization": f"Bearer {api_key}"}, + ), + ) + + +class TestTokenSource: + def test_mints_with_api_key_and_reuses_token(self) -> None: + server = _Server() + tokens = server.tokens() + assert tokens.token() == "token-1" + assert tokens.token() == "token-1" + assert len(server.mints) == 1 + assert server.mints[0].headers["Authorization"] == "Bearer key" + assert json.loads(server.mints[0].content) == {"scopes": ["sandboxes"]} + + def test_refreshes_within_a_minute_of_expiry(self) -> None: + server = _Server(expires_in=timedelta(seconds=30)) + tokens = server.tokens() + assert tokens.token() == "token-1" + assert tokens.token() == "token-2" + + def test_shares_one_mint_across_threads(self) -> None: + release = threading.Event() + server = _Server() + server.mint_gate = release + tokens = server.tokens() + results: list[str | None] = [] + started = [threading.Event() for _ in range(5)] + + def get(started_event: threading.Event) -> None: + started_event.set() + results.append(tokens.token()) + + threads = [threading.Thread(target=get, args=(event,)) for event in started] + for thread in threads: + thread.start() + for event in started: + event.wait() + release.set() + for thread in threads: + thread.join() + assert results == ["token-1"] * 5 + assert len(server.mints) == 1 + + def test_mints_again_after_failed_mint(self) -> None: + server = _Server() + server.mint_failures = 1 + tokens = server.tokens() + with pytest.raises(baseten.client.managementapi.ResponseError): + tokens.token() + assert tokens.token() == "token-2" + + def test_invalidate_drops_only_matching_token(self) -> None: + server = _Server() + tokens = server.tokens() + assert tokens.token() == "token-1" + tokens.invalidate("older-token") + assert tokens.token() == "token-1" + tokens.invalidate("token-1") + assert tokens.token() == "token-2" + + def test_calls_provider_on_every_request(self) -> None: + contexts: list[SandboxTokenProviderContext] = [] + + def provider(context: SandboxTokenProviderContext) -> str: + contexts.append(context) + return f"provided-{len(contexts)}" + + tokens = _Server().tokens(api_key="", provider=provider) + assert tokens.token() == "provided-1" + assert tokens.token(revoked_token="provided-1") == "provided-2" + assert [context.revoked_token for context in contexts] == [None, "provided-1"] + + def test_rejects_api_key_alongside_provider(self) -> None: + with pytest.raises(ValueError, match="api_key must be empty"): + _Server().tokens(provider=lambda context: "token") + + def test_no_token_for_empty_api_key_without_provider(self) -> None: + server = _Server() + assert server.tokens(api_key="").token() is None + assert server.mints == [] + + +class TestAsyncTokenSource: + @pytest.mark.asyncio + async def test_shares_one_mint_across_tasks(self) -> None: + release = asyncio.Event() + server = _Server() + tokens = _async_tokens(server, release=release) + tasks = [asyncio.create_task(tokens.token()) for _ in range(5)] + # Lets every task reach the shared mint before it completes. + await asyncio.sleep(0) + release.set() + assert await asyncio.gather(*tasks) == ["token-1"] * 5 + assert len(server.mints) == 1 + + @pytest.mark.asyncio + async def test_cancelled_caller_does_not_fail_shared_mint(self) -> None: + release = asyncio.Event() + server = _Server() + tokens = _async_tokens(server, release=release) + cancelled = asyncio.create_task(tokens.token()) + waiting = asyncio.create_task(tokens.token()) + await asyncio.sleep(0) + cancelled.cancel() + release.set() + assert await waiting == "token-1" + with pytest.raises(asyncio.CancelledError): + await cancelled + assert len(server.mints) == 1 + + @pytest.mark.asyncio + async def test_mints_again_after_failed_mint(self) -> None: + server = _Server() + server.mint_failures = 1 + tokens = _async_tokens(server) + with pytest.raises(baseten.client.managementapi.ResponseError): + await tokens.token() + assert await tokens.token() == "token-2" + + @pytest.mark.asyncio + async def test_calls_provider_with_revoked_token(self) -> None: + contexts: list[SandboxTokenProviderContext] = [] + + async def provider(context: SandboxTokenProviderContext) -> str: + contexts.append(context) + return "provided" + + tokens = _async_tokens(_Server(), api_key="", provider=provider) + assert await tokens.token(revoked_token="old") == "provided" + assert contexts[0].revoked_token == "old" + + +class TestTokenAuth: + def test_sends_bearer_and_keeps_other_headers(self) -> None: + server = _Server() + _get(server, server.tokens(), headers={"X-Other": "kept"}) + assert server.requests[0].headers["Authorization"] == "Bearer token-1" + assert server.requests[0].headers["X-Other"] == "kept" + + def test_resends_with_new_token_after_revocation(self) -> None: + server = _Server() + server.respond = lambda request: ( + _revoked() if len(server.requests) == 1 else httpx.Response(200) + ) + response = _get(server, server.tokens()) + assert response.status_code == 200 + assert server.authorizations == [ + "Bearer token-1", + "Bearer token-2", + ] + + def test_tells_caching_provider_which_token_was_revoked(self) -> None: + server = _Server() + server.respond = lambda request: ( + _revoked() if len(server.requests) == 1 else httpx.Response(200) + ) + contexts: list[SandboxTokenProviderContext] = [] + + def provider(context: SandboxTokenProviderContext) -> str: + contexts.append(context) + return f"provided-{len(contexts)}" + + _get(server, server.tokens(api_key="", provider=provider)) + assert [context.revoked_token for context in contexts] == [None, "provided-1"] + + def test_returns_revocation_after_two_resends( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + delays: list[float] = [] + monkeypatch.setattr( + baseten.sandbox._auth.time, "sleep", lambda seconds: delays.append(seconds) + ) + server = _Server(respond=_revoked) + response = _get(server, server.tokens()) + assert response.status_code == 401 + assert len(server.requests) == 3 + # Only the second resend waits, to be minted past the revocation. + assert delays == [baseten.sandbox._auth._token_revoked_retry_delay] + + def test_does_not_resend_other_401(self) -> None: + server = _Server(respond=lambda request: httpx.Response(401)) + assert _get(server, server.tokens()).status_code == 401 + assert len(server.requests) == 1 + + def test_does_not_resend_stream_body(self) -> None: + server = _Server(respond=_revoked) + with httpx.Client(transport=httpx.MockTransport(server.handle)) as client: + response = client.post( + "https://sbx.test/upload", + content=iter([b"chunk"]), + auth=TokenAuth(server.tokens()), + ) + assert response.status_code == 401 + assert len(server.requests) == 1 + + def test_sends_no_authorization_without_token(self) -> None: + server = _Server() + _get(server, server.tokens(api_key="")) + assert "Authorization" not in server.requests[0].headers + + +class TestAsyncTokenAuth: + @pytest.mark.asyncio + async def test_resends_with_new_token_after_revocation(self) -> None: + server = _Server() + server.respond = lambda request: ( + _revoked() if len(server.requests) == 1 else httpx.Response(200) + ) + tokens = _async_tokens(server) + async with httpx.AsyncClient( + transport=httpx.MockTransport(server.handle) + ) as client: + response = await client.get( + "https://sbx.test/x", auth=AsyncTokenAuth(tokens) + ) + assert response.status_code == 200 + assert server.authorizations == [ + "Bearer token-1", + "Bearer token-2", + ] + + @pytest.mark.asyncio + async def test_does_not_resend_stream_body(self) -> None: + server = _Server(respond=_revoked) + + async def body() -> AsyncIterator[bytes]: + yield b"chunk" + + async with httpx.AsyncClient( + transport=httpx.MockTransport(server.handle) + ) as client: + response = await client.post( + "https://sbx.test/upload", + content=body(), + auth=AsyncTokenAuth(_async_tokens(server)), + ) + assert response.status_code == 401 + assert len(server.requests) == 1 + + +def _get( + server: _Server, tokens: TokenSource, headers: dict[str, str] | None = None +) -> httpx.Response: + with httpx.Client(transport=httpx.MockTransport(server.handle)) as client: + return client.get("https://sbx.test/x", headers=headers, auth=TokenAuth(tokens)) + + +def _async_tokens( + server: _Server, + *, + api_key: str = "key", + provider=None, + release: asyncio.Event | None = None, +) -> AsyncTokenSource: + async def handle(request: httpx.Request) -> httpx.Response: + if release is not None and request.url.path == "/v1/token": + await release.wait() + return server.handle(request) + + return AsyncTokenSource( + api_key=api_key, + token_provider=provider, + mint_api=baseten.client.managementapi.AsyncApiClient( + httpx.AsyncClient(transport=httpx.MockTransport(handle)), + base_url="https://api.test", + headers={"Authorization": f"Bearer {api_key}"}, + ), + ) diff --git a/tests/sandbox/test_client.py b/tests/sandbox/test_client.py new file mode 100644 index 0000000..19bf9e0 --- /dev/null +++ b/tests/sandbox/test_client.py @@ -0,0 +1,467 @@ +from __future__ import annotations + +import importlib.util +import json +from collections.abc import Callable +from datetime import UTC, datetime, timedelta +from typing import Any + +import httpx +import pytest + +import baseten.sandbox._client +from baseten.sandbox import ( + AsyncSandboxClient, + SandboxAPIError, + SandboxClient, + SandboxEnvValue, + SandboxError, + SandboxExpirationPolicyIdle, + SandboxInfo, + SandboxLifecycle, + SandboxNetwork, + SandboxPort, +) + +_SANDBOX_URL = "https://sbx.example.dev/" + + +class _Api: + """Fake control plane, token exchange, and sandbox, recording requests.""" + + def __init__(self) -> None: + self.requests: list[httpx.Request] = [] + self.authorizations: list[str | None] = [] + self.mints = 0 + self.routes: dict[ + tuple[str, str], Callable[[httpx.Request], httpx.Response] + ] = {} + + def route( + self, + method: str, + path: str, + respond: Callable[[httpx.Request], httpx.Response] + | httpx.Response + | dict[str, Any], + *, + status: int = 200, + ) -> None: + if isinstance(respond, dict): + body = respond + self.routes[(method, path)] = lambda request: httpx.Response( + status, json=body + ) + elif isinstance(respond, httpx.Response): + # A response is single use, so each request gets a copy. + template = respond + self.routes[(method, path)] = lambda request: httpx.Response( + template.status_code, + headers=template.headers, + content=template.content, + ) + else: + self.routes[(method, path)] = respond + + def handle(self, request: httpx.Request) -> httpx.Response: + if request.url.path == "/v1/token": + self.mints += 1 + return httpx.Response( + 200, + json={ + "token": f"token-{self.mints}", + "expires_at": (datetime.now(UTC) + timedelta(hours=2)).isoformat(), + }, + ) + self.requests.append(request) + self.authorizations.append(request.headers.get("Authorization")) + return self.routes[(request.method, request.url.path)](request) + + def client(self, **options: Any) -> SandboxClient: + return SandboxClient( + api_key="key", + base_url_override="https://api.test", + http_client_override=httpx.Client( + transport=httpx.MockTransport(self.handle) + ), + **options, + ) + + def async_client(self, **options: Any) -> AsyncSandboxClient: + async def handle(request: httpx.Request) -> httpx.Response: + return self.handle(request) + + return AsyncSandboxClient( + api_key="key", + base_url_override="https://api.test", + http_client_override=httpx.AsyncClient( + transport=httpx.MockTransport(handle) + ), + **options, + ) + + +def _record(name: str = "sbx", **fields: Any) -> dict[str, Any]: + return { + "name": name, + "url": _SANDBOX_URL, + "status": "DEPLOYED", + "created_at": "2026-10-06T12:00:00Z", + **fields, + } + + +def _body(request: httpx.Request) -> Any: + return json.loads(request.content) + + +class TestSandboxClient: + def test_creates_with_nothing_set_leaving_defaults_to_server(self) -> None: + api = _Api() + api.route("POST", "/v1/sandboxes/instances", _record(), status=201) + result = api.client().create() + request = api.requests[0] + assert str(request.url) == "https://api.test/v1/sandboxes/instances" + assert _body(request) == {} + assert result.info.name == "sbx" + assert (result.name, result.url) == ("sbx", _SANDBOX_URL) + + def test_creates_with_every_field_mapped(self) -> None: + api = _Api() + api.route("POST", "/v1/sandboxes/instances", _record(), status=201) + api.client().create( + name="sbx", + create_if_not_exists=True, + image="img:1", + memory=2048, + region="us-was-1", + envs={"A": SandboxEnvValue(value="1", secret=False)}, + labels={"team": "infra"}, + external_id="ext-1", + lifecycle=SandboxLifecycle( + expiration_policies=[ + SandboxExpirationPolicyIdle(after=timedelta(minutes=5)) + ] + ), + ports=[SandboxPort(target=80)], + network=SandboxNetwork(subnet="default"), + ) + assert _body(api.requests[0]) == { + "name": "sbx", + "create_if_not_exists": True, + "image": "img:1", + "memory": 2048, + "region": "us-was-1", + "envs": [{"name": "A", "value": "1", "secret": False}], + "labels": {"team": "infra"}, + "external_id": "ext-1", + "lifecycle": { + "expiration_policies": [ + {"type": "TTL_IDLE", "action": "DELETE", "value": "300000ms"} + ] + }, + "ports": [{"target": 80}], + "network": {"subnet": "default"}, + } + + def test_sends_team_id_on_every_call_when_set(self) -> None: + api = _Api() + api.route("POST", "/v1/sandboxes/instances", _record(), status=201) + api.route("GET", "/v1/sandboxes/instances/sbx", _record()) + api.route( + "GET", + "/v1/sandboxes/instances", + {"items": [], "pagination": {"has_more": False}}, + ) + api.route("PATCH", "/v1/sandboxes/instances/sbx", _record()) + api.route("DELETE", "/v1/sandboxes/instances/sbx", _record(), status=202) + client = api.client(team_id="team-1") + client.create() + client.get_info(name="sbx") + list(client.list()) + client.update(name="sbx", labels={"a": "b"}) + client.delete(name="sbx") + assert [r.url.params.get("team_id") for r in api.requests] == ["team-1"] * 5 + + def test_omits_team_id_when_unset(self) -> None: + api = _Api() + api.route("GET", "/v1/sandboxes/instances/sbx", _record()) + api.client().get_info(name="sbx") + assert "team_id" not in api.requests[0].url.params + + def test_gets_info_with_show_secrets_only_when_set(self) -> None: + api = _Api() + api.route("GET", "/v1/sandboxes/instances/sbx", _record()) + client = api.client() + client.get_info(name="sbx") + client.get_info(name="sbx", show_secrets=True) + assert [r.url.params.get("show_secrets") for r in api.requests] == [ + None, + "true", + ] + + def test_gets_sandbox_by_name_with_call_and_from_record_without(self) -> None: + api = _Api() + api.route("GET", "/v1/sandboxes/instances/sbx", _record()) + client = api.client() + by_name = client.get(name="sbx") + assert len(api.requests) == 1 + from_record = client.get(info=client.get_info(name="sbx")) + assert len(api.requests) == 2 + assert (by_name.name, by_name.url) == ("sbx", _SANDBOX_URL) + assert (from_record.name, from_record.url) == ("sbx", _SANDBOX_URL) + + def test_get_needs_exactly_one_of_name_or_info(self) -> None: + client = _Api().client() + with pytest.raises(TypeError, match="exactly one"): + client.get() # ty: ignore[no-matching-overload] + info = _info() + with pytest.raises(TypeError, match="exactly one"): + client.get(name="sbx", info=info) # ty: ignore[no-matching-overload] + + def test_builds_sandbox_from_url_without_call(self) -> None: + api = _Api() + sandbox = api.client().sandbox_from_url(url=_SANDBOX_URL) + assert (sandbox.name, sandbox.url) == ("", _SANDBOX_URL) + assert api.requests == [] + + def test_gives_sandboxes_client_token(self) -> None: + api = _Api() + api.route("GET", "/v1/sandboxes/instances/sbx", _record()) + api.route("GET", "/process", httpx.Response(200, json=[])) + client = api.client(headers={"X-Extra": "1"}) + client.get(name="sbx").raw_api.get_process() + exec_request = api.requests[1] + assert str(exec_request.url) == "https://sbx.example.dev/process" + assert api.authorizations == ["Bearer token-1", "Bearer token-1"] + assert api.mints == 1 + for request in api.requests: + assert request.headers["X-Extra"] == "1" + assert request.headers["User-Agent"].startswith("baseten-python/") + + def test_drops_client_token_when_sandbox_finds_it_revoked(self) -> None: + api = _Api() + revoked = {"done": False} + + def revoke_once(request: httpx.Request) -> httpx.Response: + if revoked["done"]: + return httpx.Response(200, json=[]) + revoked["done"] = True + return httpx.Response(401, headers={"x-blaxel-error-code": "TOKEN_REVOKED"}) + + api.route("GET", "/process", revoke_once) + api.route("GET", "/v1/sandboxes/instances/sbx", _record()) + client = api.client() + client.sandbox_from_url(url=_SANDBOX_URL).raw_api.get_process() + client.get_info(name="sbx") + assert api.authorizations == [ + "Bearer token-1", + "Bearer token-2", + "Bearer token-2", + ] + + def test_lists_lazily_across_pages_with_filters_on_each(self) -> None: + api = _Api() + pages = iter( + [ + { + "items": [_record("a"), _record("b")], + "pagination": {"has_more": True, "cursor": "c1"}, + }, + {"items": [_record("c")], "pagination": {"has_more": False}}, + ] + ) + api.route( + "GET", + "/v1/sandboxes/instances", + lambda request: httpx.Response(200, json=next(pages)), + ) + listed = api.client().list( + query="q", statuses=["DEPLOYED", "FAILED"], page_size=2 + ) + assert next(listed).name == "a" + assert len(api.requests) == 1 + assert [info.name for info in listed] == ["b", "c"] + first, second = (request.url.params for request in api.requests) + assert "cursor" not in first + assert second["cursor"] == "c1" + for params in (first, second): + assert params.get_list("status") == ["DEPLOYED", "FAILED"] + assert (params["q"], params["limit"]) == ("q", "2") + + def test_stops_listing_on_repeated_cursor(self) -> None: + api = _Api() + api.route( + "GET", + "/v1/sandboxes/instances", + {"items": [_record()], "pagination": {"has_more": True, "cursor": "same"}}, + ) + with pytest.raises( + SandboxError, match="sandbox list returned a repeated cursor" + ): + list(api.client().list()) + assert len(api.requests) == 2 + + def test_updates_with_every_field_mapped(self) -> None: + api = _Api() + api.route("PATCH", "/v1/sandboxes/instances/sbx", _record(labels={"rev": "2"})) + info = api.client().update( + name="sbx", + lifecycle=SandboxLifecycle(terminated_retention=timedelta(hours=1)), + envs={"A": SandboxEnvValue(value="1")}, + external_id="ext-2", + labels={"rev": "2"}, + ) + assert _body(api.requests[0]) == { + "lifecycle": {"terminated_retention": "3600000ms"}, + "envs": [{"name": "A", "value": "1"}], + "external_id": "ext-2", + "labels": {"rev": "2"}, + } + assert info.labels == {"rev": "2"} + + def test_updates_only_fields_that_are_set(self) -> None: + api = _Api() + api.route("PATCH", "/v1/sandboxes/instances/sbx", _record()) + api.client().update(name="sbx", labels={}) + assert _body(api.requests[0]) == {"labels": {}} + + def test_fails_update_server_rejects(self) -> None: + api = _Api() + api.route( + "PATCH", + "/v1/sandboxes/instances/sbx", + httpx.Response( + 400, content=b'{"error": "BAD_REQUEST", "message": "nothing to update"}' + ), + ) + with pytest.raises(SandboxAPIError) as info: + api.client().update(name="sbx") + assert (info.value.status, info.value.code, info.value.message) == ( + 400, + "BAD_REQUEST", + "nothing to update", + ) + + def test_never_retries_control_plane(self) -> None: + api = _Api() + api.route( + "GET", "/v1/sandboxes/instances/sbx", httpx.Response(503, content=b"") + ) + with pytest.raises(SandboxAPIError): + api.client().get_info(name="sbx") + assert len(api.requests) == 1 + + def test_deletes_and_returns_record(self) -> None: + api = _Api() + api.route( + "DELETE", + "/v1/sandboxes/instances/sbx", + _record(status="DELETING"), + status=202, + ) + assert api.client().delete(name="sbx").status == "DELETING" + + def test_closes_only_its_own_http_client(self) -> None: + own = SandboxClient(api_key="key", http2=False) + own.close() + assert own.http_client.is_closed + override = httpx.Client() + SandboxClient(api_key="key", http_client_override=override).close() + assert not override.is_closed + override.close() + + +class TestAsyncSandboxClient: + @pytest.mark.asyncio + async def test_creates_and_gets_sandbox(self) -> None: + api = _Api() + api.route("POST", "/v1/sandboxes/instances", _record(), status=201) + api.route("GET", "/v1/sandboxes/instances/sbx", _record()) + api.route("GET", "/process", httpx.Response(200, json=[])) + client = api.async_client(team_id="team-1") + created = await client.create(name="sbx") + assert created.info.name == "sbx" + sandbox = await client.get(name="sbx") + await sandbox.raw_api.get_process() + assert _body(api.requests[0]) == {"name": "sbx"} + assert [r.url.params.get("team_id") for r in api.requests[:2]] == ["team-1"] * 2 + assert api.authorizations == ["Bearer token-1"] * 3 + + @pytest.mark.asyncio + async def test_lists_across_pages(self) -> None: + api = _Api() + pages = iter( + [ + { + "items": [_record("a")], + "pagination": {"has_more": True, "cursor": "c1"}, + }, + {"items": [_record("b")], "pagination": {"has_more": False}}, + ] + ) + api.route( + "GET", + "/v1/sandboxes/instances", + lambda request: httpx.Response(200, json=next(pages)), + ) + names = [info.name async for info in api.async_client().list()] + assert names == ["a", "b"] + + @pytest.mark.asyncio + async def test_converts_error_responses(self) -> None: + api = _Api() + api.route( + "DELETE", + "/v1/sandboxes/instances/sbx", + httpx.Response(404, content=b'{"error": "NOT_FOUND", "message": "gone"}'), + ) + with pytest.raises(SandboxAPIError) as info: + await api.async_client().delete(name="sbx") + assert info.value.code == "NOT_FOUND" + + @pytest.mark.asyncio + async def test_closes_only_its_own_http_client(self) -> None: + own = AsyncSandboxClient(api_key="key", http2=False) + await own.close() + assert own.http_client.is_closed + override = httpx.AsyncClient() + await AsyncSandboxClient(api_key="key", http_client_override=override).close() + assert not override.is_closed + await override.aclose() + + +class TestUseHttp2: + def test_uses_http2_when_h2_installed(self) -> None: + assert baseten.sandbox._client._use_http2(None) + assert baseten.sandbox._client._use_http2(True) + assert not baseten.sandbox._client._use_http2(False) + + def test_without_h2(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(importlib.util, "find_spec", lambda name: None) + assert not baseten.sandbox._client._use_http2(None) + with pytest.raises(ImportError, match="httpx\\[http2\\]"): + baseten.sandbox._client._use_http2(True) + + +def _info() -> SandboxInfo: + return SandboxInfo( + name="sbx", + url=_SANDBOX_URL, + status="DEPLOYED", + image=None, + memory=None, + region=None, + envs={}, + labels={}, + external_id=None, + lifecycle=None, + ports=[], + network=None, + created_at=datetime(2026, 10, 6, tzinfo=UTC), + updated_at=None, + created_by=None, + updated_by=None, + last_used_at=None, + expires_in=None, + ) diff --git a/tests/sandbox/test_common.py b/tests/sandbox/test_common.py new file mode 100644 index 0000000..b8e4949 --- /dev/null +++ b/tests/sandbox/test_common.py @@ -0,0 +1,237 @@ +from __future__ import annotations + +from collections.abc import AsyncGenerator, AsyncIterator, Generator, Iterator +from datetime import UTC, datetime, timedelta, timezone + +import httpx +import pytest + +from baseten.sandbox import AsyncSandboxStream, SandboxError, SandboxStream +from baseten.sandbox._common import ( + format_duration, + parse_duration, + parse_timestamp, + response_lines, + response_lines_async, +) + + +def test_format_duration_rounds_to_milliseconds() -> None: + assert format_duration(timedelta(hours=1)) == "3600000ms" + assert format_duration(timedelta(microseconds=1500)) == "2ms" + assert format_duration(timedelta(0)) == "0ms" + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + ("0", timedelta(0)), + ("90s", timedelta(seconds=90)), + ("1h30m", timedelta(hours=1, minutes=30)), + ("1.5h", timedelta(hours=1, minutes=30)), + (".5s", timedelta(milliseconds=500)), + ("3600000ms", timedelta(hours=1)), + ("250us", timedelta(microseconds=250)), + ("250µs", timedelta(microseconds=250)), + ("1500ns", timedelta(microseconds=1.5)), + ("+5m", timedelta(minutes=5)), + ("-5m", timedelta(minutes=-5)), + # Days and weeks only on their own, at 24 hours per day. + ("7d", timedelta(days=7)), + ("2w", timedelta(days=14)), + ], +) +def test_parse_duration(value: str, expected: timedelta) -> None: + assert parse_duration(value, "test duration") == expected + + +@pytest.mark.parametrize("value", ["", "5", "1d12h", "abc", "5m x", "1.5d"]) +def test_parse_duration_rejects_invalid(value: str) -> None: + with pytest.raises(SandboxError, match="test duration is not a valid duration"): + parse_duration(value, "test duration") + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + # HTTP date, as process times come. + ( + "Wed, 01 Jan 2023 12:00:00 GMT", + datetime(2023, 1, 1, 12, 0, 0, tzinfo=UTC), + ), + # RFC 3339 with a space for the T, as health times come. + ( + "2026-01-29 17:36:52+00:00", + datetime(2026, 1, 29, 17, 36, 52, tzinfo=UTC), + ), + ("2026-10-06T12:00:00Z", datetime(2026, 10, 6, 12, 0, 0, tzinfo=UTC)), + ( + "2026-10-06T12:00:00+02:00", + datetime(2026, 10, 6, 12, 0, 0, tzinfo=timezone(timedelta(hours=2))), + ), + ( + "2026-10-06T12:00:00.123456Z", + datetime(2026, 10, 6, 12, 0, 0, 123456, tzinfo=UTC), + ), + # Nanoseconds, as Go writes them, truncated to microseconds. + ( + "2026-10-06T12:00:00.123456789Z", + datetime(2026, 10, 6, 12, 0, 0, 123456, tzinfo=UTC), + ), + ], +) +def test_parse_timestamp(value: str, expected: datetime) -> None: + parsed = parse_timestamp(value, "test time") + assert parsed == expected + assert parsed.utcoffset() == expected.utcoffset() + + +@pytest.mark.parametrize( + "value", + [ + "", + "garbage", + # No zone, so the time would be a guess. + "2026-10-06T12:00:00", + "2026-10-06 12:00:00", + "2026-10-06", + # HTTP dates in any zone but GMT, or missing the weekday. + "Wed, 01 Jan 2023 12:00:00 EST", + "Wed, 01 Jan 2023 12:00:00 +0000", + "01 Jan 2023 12:00:00 GMT", + # Unix seconds. + "1696593600", + ], +) +def test_parse_timestamp_rejects_invalid(value: str) -> None: + with pytest.raises(SandboxError, match="test time is not a valid timestamp"): + parse_timestamp(value, "test time") + + +class _ChunkStream(httpx.SyncByteStream, httpx.AsyncByteStream): + """Response body sent in the given chunks, recording whether it was closed.""" + + def __init__(self, *chunks: bytes) -> None: + self.chunks = chunks + self.closed = False + + def __iter__(self) -> Iterator[bytes]: + yield from self.chunks + + async def __aiter__(self) -> AsyncIterator[bytes]: + for chunk in self.chunks: + yield chunk + + def close(self) -> None: + self.closed = True + + async def aclose(self) -> None: + self.closed = True + + +@pytest.mark.parametrize( + ("chunks", "expected"), + [ + ((b"",), []), + ((b"a\nb\n",), ["a", "b"]), + # A last line without a line ending is still a line. + ((b"a\nb",), ["a", "b"]), + ((b"a\r\nb\r\n",), ["a", "b"]), + # Lines split across chunks, including between \r and \n. + ((b"ab", b"c\nd", b"e\n"), ["abc", "de"]), + ((b"a\r", b"\nb"), ["a", "b"]), + # A lone \r is part of the line, as a progress bar writes it. + ((b"10%\r20%\n",), ["10%\r20%"]), + ((b"a\n\nb\n",), ["a", "", "b"]), + # A character split across chunks, and invalid UTF-8 replaced. + (("é\n".encode()[:1], "é\n".encode()[1:]), ["é"]), + ((b"a\xffb\n",), ["a�b"]), + ], +) +class TestResponseLines: + def test_sync(self, chunks: tuple[bytes, ...], expected: list[str]) -> None: + response = httpx.Response(200, stream=_ChunkStream(*chunks)) + assert list(response_lines(response)) == expected + + @pytest.mark.asyncio + async def test_async(self, chunks: tuple[bytes, ...], expected: list[str]) -> None: + response = httpx.Response(200, stream=_ChunkStream(*chunks)) + assert [line async for line in response_lines_async(response)] == expected + + +def _upper(lines: Iterator[str]) -> Generator[str, None, None]: + for line in lines: + if line == "fail": + raise ValueError("bad line") + yield line.upper() + + +async def _upper_async(lines: AsyncIterator[str]) -> AsyncGenerator[str, None]: + async for line in lines: + if line == "fail": + raise ValueError("bad line") + yield line.upper() + + +class TestSandboxStream: + def test_yields_items_and_closes_at_end(self) -> None: + body = _ChunkStream(b"a\nb\n") + stream = SandboxStream(httpx.Response(200, stream=body), _upper) + assert list(stream) == ["A", "B"] + assert body.closed + + def test_closes_on_error(self) -> None: + body = _ChunkStream(b"a\nfail\nb\n") + stream = SandboxStream(httpx.Response(200, stream=body), _upper) + assert next(stream) == "A" + with pytest.raises(ValueError, match="bad line"): + next(stream) + assert body.closed + + def test_closes_when_left_early(self) -> None: + body = _ChunkStream(b"a\nb\n") + with SandboxStream(httpx.Response(200, stream=body), _upper) as stream: + assert next(stream) == "A" + assert body.closed + with pytest.raises(StopIteration): + next(stream) + + def test_closes_before_first_item(self) -> None: + body = _ChunkStream(b"a\n") + SandboxStream(httpx.Response(200, stream=body), _upper).close() + assert body.closed + + +class TestAsyncSandboxStream: + @pytest.mark.asyncio + async def test_yields_items_and_closes_at_end(self) -> None: + body = _ChunkStream(b"a\nb\n") + stream = AsyncSandboxStream(httpx.Response(200, stream=body), _upper_async) + assert [item async for item in stream] == ["A", "B"] + assert body.closed + + @pytest.mark.asyncio + async def test_closes_on_error(self) -> None: + body = _ChunkStream(b"a\nfail\nb\n") + stream = AsyncSandboxStream(httpx.Response(200, stream=body), _upper_async) + assert await anext(stream) == "A" + with pytest.raises(ValueError, match="bad line"): + await anext(stream) + assert body.closed + + @pytest.mark.asyncio + async def test_closes_when_left_early(self) -> None: + body = _ChunkStream(b"a\nb\n") + async with AsyncSandboxStream( + httpx.Response(200, stream=body), _upper_async + ) as stream: + assert await anext(stream) == "A" + assert body.closed + with pytest.raises(StopAsyncIteration): + await anext(stream) + + @pytest.mark.asyncio + async def test_closes_before_first_item(self) -> None: + body = _ChunkStream(b"a\n") + await AsyncSandboxStream(httpx.Response(200, stream=body), _upper_async).close() + assert body.closed diff --git a/tests/sandbox/test_e2e.py b/tests/sandbox/test_e2e.py new file mode 100644 index 0000000..92e680d --- /dev/null +++ b/tests/sandbox/test_e2e.py @@ -0,0 +1,1459 @@ +"""End-to-end tests of baseten.sandbox against a live Baseten environment. + +Requires the following environment variables: + BASETEN_E2E_TEST_API_KEY — API key for authentication + BASETEN_E2E_TEST_DOMAIN — Domain (e.g. baseten.co) + +Optional: + BASETEN_E2E_TEST_DEBUG=1 — Print every request and response, credentials + masked + BASETEN_E2E_TEST_KEEP_SANDBOX — Leave the shared sandbox in place after the run + +Skipped when BASETEN_E2E_TEST_API_KEY is empty or unset. If the API key is set +but DOMAIN is missing, the tests error. + +Every test shares one sandbox, created once per run and always deleted. Each +test runs on both the sync and the async client. One image is pushed per run, +except on Python 3.11, and a sandbox created from it; both are always deleted. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import os +import re +import secrets +import sys +import threading +import time +from collections.abc import AsyncIterator, Callable, Iterator +from dataclasses import dataclass +from datetime import UTC, datetime, timedelta +from typing import Any, cast + +import httpx +import pydantic +import pytest +import pytest_asyncio + +from baseten.sandbox import ( + AsyncSandbox, + AsyncSandboxClient, + ImageBuilder, + Sandbox, + SandboxAPIError, + SandboxClient, + SandboxEnvValue, + SandboxError, + SandboxExpirationPolicyIdle, + SandboxFileSystemCopyError, + SandboxFileSystemWatchEvent, + SandboxInfo, + SandboxLifecycle, + SandboxProcessExecEvent, + SandboxProcessExecEventExit, + SandboxProcessExecEventOutput, + SandboxProcessWaitTimeoutError, +) + +API_KEY = os.environ.get("BASETEN_E2E_TEST_API_KEY", "") +DOMAIN = os.environ.get("BASETEN_E2E_TEST_DOMAIN", "") +DEBUG = os.environ.get("BASETEN_E2E_TEST_DEBUG") == "1" +KEEP_SANDBOX = bool(os.environ.get("BASETEN_E2E_TEST_KEEP_SANDBOX")) + +# Labels on every sandbox the suite creates. +_LABELS = {"created_by": "e2e"} + +# Env on the shared sandbox, by name. +_SHARED_ENVS = { + "E2E_PLAIN": SandboxEnvValue(value="plain-value", secret=False), + "E2E_SECRET": SandboxEnvValue(value="secret-value", secret=True), + # Left to the server's default, which is secret. + "E2E_DEFAULT": SandboxEnvValue(value="default-value"), +} + +# Allowance between this machine's clock and the server's. +_CLOCK_SKEW = timedelta(minutes=5) + +# The SDK's own timeouts, for the debug clients that replace its HTTP client. +_DEBUG_TIMEOUT = httpx.Timeout(5.0, connect=10.0, read=None) + + +@dataclass +class _Shared: + """The sandbox every test shares.""" + + name: str + info: SandboxInfo + """Its record as created.""" + started: datetime + """When the run started creating it.""" + + +@pytest.fixture(scope="session") +def shared() -> Iterator[_Shared]: + _require_env() + client = _client() + name = _unique() + started = datetime.now(UTC) + # The name is known before creating, so even a partial create is removed. + try: + created = client.create( + name=name, + # TODO: Temporary, while other regions' exec planes reject our + # tokens. + region="us-was-1", + labels=_LABELS, + envs=_SHARED_ENVS, + ) + # Create returns once the sandbox is up, so there is nothing to wait + # for. + assert created.info.status == "DEPLOYED" + except BaseException: + try: + client.delete(name=name) + except SandboxAPIError as err: + # A create that never got through leaves nothing to delete. The + # original error is raised either way. + if err.status != 404: + print(f"failed to delete sandbox {name} after setup failed: {err}") + client.close() + raise + try: + yield _Shared(name=name, info=created.info, started=started) + finally: + try: + if KEEP_SANDBOX: + print(f"BASETEN_E2E_TEST_KEEP_SANDBOX set; leaving sandbox {name}") + else: + deleted = client.delete(name=name) + assert (deleted.name, deleted.status) == (name, "DELETING") + finally: + client.close() + + +# Each test runs on both clients: the sync one's calls go to a thread, so test +# bodies stay async for both. +@pytest_asyncio.fixture(params=["sync", "async"]) +async def client( + request: pytest.FixtureRequest, +) -> AsyncIterator[SandboxClient | AsyncSandboxClient]: + if request.param == "async": + async with _async_client() as async_client: + yield async_client + else: + with _client() as sync_client: + yield sync_client + + +@pytest.mark.asyncio +async def test_gets_record( + client: SandboxClient | AsyncSandboxClient, shared: _Shared +) -> None: + if isinstance(client, SandboxClient): + info = await asyncio.to_thread(client.get_info, name=shared.name) + else: + info = await client.get_info(name=shared.name) + assert info.name == shared.name + assert info.status == "DEPLOYED" + assert info.url.startswith("https://") + assert info.labels.items() >= _LABELS.items() + assert shared.started - _CLOCK_SKEW <= info.created_at + assert info.created_at <= datetime.now(UTC) + _CLOCK_SKEW + # Only an env set as not secret comes back unmasked, and one left unset is + # secret. + assert info.envs == { + "E2E_PLAIN": SandboxEnvValue(value="plain-value", secret=False), + "E2E_SECRET": SandboxEnvValue(value="****", secret=True), + "E2E_DEFAULT": SandboxEnvValue(value="****", secret=True), + } + + +@pytest.mark.asyncio +async def test_gets_record_with_secrets_revealed( + client: SandboxClient | AsyncSandboxClient, shared: _Shared +) -> None: + # The e2e key is a workspace admin, so secrets are revealed. + if isinstance(client, SandboxClient): + info = await asyncio.to_thread( + client.get_info, name=shared.name, show_secrets=True + ) + else: + info = await client.get_info(name=shared.name, show_secrets=True) + assert info.envs == { + "E2E_PLAIN": SandboxEnvValue(value="plain-value", secret=False), + "E2E_SECRET": SandboxEnvValue(value="secret-value", secret=True), + "E2E_DEFAULT": SandboxEnvValue(value="default-value", secret=True), + } + + +@pytest.mark.asyncio +async def test_gets_sandbox_by_name_record_and_url( + client: SandboxClient | AsyncSandboxClient, shared: _Shared +) -> None: + # Reaching the sandbox from its URL checks the minted token works there too. + if isinstance(client, SandboxClient): + info = await asyncio.to_thread(client.get_info, name=shared.name) + by_name = await asyncio.to_thread(client.get, name=shared.name) + from_info = await asyncio.to_thread(client.get, info=info) + from_url = client.sandbox_from_url(url=info.url) + processes = await asyncio.to_thread(from_url.process.list) + else: + info = await client.get_info(name=shared.name) + by_name = await client.get(name=shared.name) + from_info = await client.get(info=info) + from_url = client.sandbox_from_url(url=info.url) + processes = await from_url.process.list() + for sandbox in (by_name, from_info): + assert (sandbox.name, sandbox.url) == (shared.name, info.url) + assert from_url.name == "" + assert isinstance(processes, list) + + +@pytest.mark.asyncio +async def test_lists( + client: SandboxClient | AsyncSandboxClient, shared: _Shared +) -> None: + if isinstance(client, SandboxClient): + infos = await asyncio.to_thread(lambda: list(client.list(query=shared.name))) + else: + infos = [info async for info in client.list(query=shared.name)] + assert shared.name in [info.name for info in infos] + + +@pytest.mark.asyncio +async def test_create_if_not_exists_returns_existing( + client: SandboxClient | AsyncSandboxClient, shared: _Shared +) -> None: + if isinstance(client, SandboxClient): + before = await asyncio.to_thread(client.get_info, name=shared.name) + created = await asyncio.to_thread( + client.create, + name=shared.name, + create_if_not_exists=True, + labels={"other": "1"}, + ) + else: + before = await client.get_info(name=shared.name) + created = await client.create( + name=shared.name, create_if_not_exists=True, labels={"other": "1"} + ) + assert (created.name, created.url) == (shared.name, before.url) + assert created.info.created_at == before.created_at + # The existing configuration is kept, not replaced by the request's. + assert created.info.labels == before.labels + + +# Changes only fields no other test on the shared sandbox depends on. +@pytest.mark.asyncio +async def test_update_replaces_labels( + client: SandboxClient | AsyncSandboxClient, shared: _Shared +) -> None: + labels = {**_LABELS, "e2e_extra": "1"} + lifecycle = SandboxLifecycle( + expiration_policies=[SandboxExpirationPolicyIdle(after=timedelta(hours=1))], + terminated_retention=timedelta(minutes=10), + ) + if isinstance(client, SandboxClient): + updated = await asyncio.to_thread( + client.update, name=shared.name, lifecycle=lifecycle, labels=labels + ) + fetched = await asyncio.to_thread(client.get_info, name=shared.name) + replaced = await asyncio.to_thread( + client.update, name=shared.name, labels=_LABELS + ) + else: + updated = await client.update( + name=shared.name, lifecycle=lifecycle, labels=labels + ) + fetched = await client.get_info(name=shared.name) + replaced = await client.update(name=shared.name, labels=_LABELS) + for info in (updated, fetched): + assert info.lifecycle == SandboxLifecycle( + expiration_policies=[ + SandboxExpirationPolicyIdle(after=timedelta(hours=1), action="DELETE") + ], + terminated_retention=timedelta(minutes=10), + ) + assert info.labels == labels + assert replaced.labels == _LABELS + + +@pytest.mark.asyncio +async def test_update_with_nothing_fails( + client: SandboxClient | AsyncSandboxClient, shared: _Shared +) -> None: + with pytest.raises(SandboxAPIError) as info: + if isinstance(client, SandboxClient): + await asyncio.to_thread(client.update, name=shared.name) + else: + await client.update(name=shared.name) + assert info.value.status == 400 + + +@pytest.mark.asyncio +async def test_get_missing_fails( + client: SandboxClient | AsyncSandboxClient, shared: _Shared +) -> None: + missing = f"{shared.name}-missing" + with pytest.raises(SandboxAPIError) as info: + if isinstance(client, SandboxClient): + await asyncio.to_thread(client.get_info, name=missing) + else: + await client.get_info(name=missing) + assert info.value.status == 404 + assert info.value.code is not None + + +@pytest_asyncio.fixture +async def sandbox( + client: SandboxClient | AsyncSandboxClient, shared: _Shared +) -> Sandbox | AsyncSandbox: + if isinstance(client, SandboxClient): + return await asyncio.to_thread(client.get, name=shared.name) + return await client.get(name=shared.name) + + +@pytest.mark.asyncio +async def test_lists_library_images( + client: SandboxClient | AsyncSandboxClient, +) -> None: + _require_env() + if isinstance(client, SandboxClient): + images = await asyncio.to_thread(client.images.list_library) + else: + images = await client.images.list_library() + assert all(image.image for image in images) + base = [image for image in images if image.name == "base-image"] + assert [image.image for image in base] == ["baseten/base-image:latest"] + + +@pytest.mark.asyncio +async def test_process_runs_to_completion(sandbox: Sandbox | AsyncSandbox) -> None: + started = datetime.now(UTC) + command = "sh -c 'echo out; echo err >&2; exit 3'" + if isinstance(sandbox, Sandbox): + process = await asyncio.to_thread( + sandbox.process.exec, command=command, wait_for_completion=True + ) + else: + process = await sandbox.process.exec(command=command, wait_for_completion=True) + assert process.exit_code == 3 + assert (process.stdout.strip(), process.stderr.strip()) == ("out", "err") + _assert_within_run(started, process.started_at) + assert process.completed_at is not None + _assert_within_run(started, process.completed_at) + + +@pytest.mark.asyncio +async def test_process_working_dir_and_env(sandbox: Sandbox | AsyncSandbox) -> None: + command = "sh -c 'pwd; echo $E2E_PLAIN; echo $E2E_SECRET; echo $E2E_EXEC'" + options: dict[str, Any] = { + "command": command, + "working_dir": "/tmp", + "env": {"E2E_EXEC": "exec-value"}, + "wait_for_completion": True, + } + if isinstance(sandbox, Sandbox): + process = await asyncio.to_thread(lambda: sandbox.process.exec(**options)) + else: + process = await sandbox.process.exec(**options) + # The sandbox's envs, secret ones too, reach its processes. + assert process.stdout.split() == [ + "/tmp", + "plain-value", + "secret-value", + "exec-value", + ] + + +@pytest.mark.asyncio +async def test_process_gets_lists_and_waits_for_background_process( + sandbox: Sandbox | AsyncSandbox, +) -> None: + name = _unique() + command = "sh -c 'sleep 1; echo done'" + try: + if isinstance(sandbox, Sandbox): + await asyncio.to_thread(sandbox.process.exec, command=command, name=name) + got = await asyncio.to_thread(sandbox.process.get, identifier=name) + processes = await asyncio.to_thread(sandbox.process.list) + finished = await asyncio.to_thread(sandbox.process.wait, identifier=name) + logs = await asyncio.to_thread(sandbox.process.logs, identifier=got.pid) + else: + await sandbox.process.exec(command=command, name=name) + got = await sandbox.process.get(identifier=name) + processes = await sandbox.process.list() + finished = await sandbox.process.wait(identifier=name) + logs = await sandbox.process.logs(identifier=got.pid) + finally: + await _kill_quietly(sandbox, name) + assert (got.name, got.status, got.completed_at) == (name, "running", None) + assert name in [process.name for process in processes] + assert (finished.status, finished.exit_code) == ("completed", 0) + assert logs.stdout.strip() == "done" + + +@pytest.mark.asyncio +async def test_process_stops_and_kills(sandbox: Sandbox | AsyncSandbox) -> None: + stopped, killed = _unique(), _unique() + try: + if isinstance(sandbox, Sandbox): + for name in (stopped, killed): + await asyncio.to_thread( + sandbox.process.exec, command="sleep 300", name=name + ) + await asyncio.to_thread(sandbox.process.stop, identifier=stopped) + await asyncio.to_thread(sandbox.process.kill, identifier=killed) + statuses = [ + (await asyncio.to_thread(sandbox.process.wait, identifier=name)).status + for name in (stopped, killed) + ] + else: + for name in (stopped, killed): + await sandbox.process.exec(command="sleep 300", name=name) + await sandbox.process.stop(identifier=stopped) + await sandbox.process.kill(identifier=killed) + statuses = [ + (await sandbox.process.wait(identifier=name)).status + for name in (stopped, killed) + ] + finally: + for name in (stopped, killed): + await _kill_quietly(sandbox, name) + assert statuses == ["stopped", "killed"] + + +@pytest.mark.asyncio +async def test_process_writes_and_closes_stdin(sandbox: Sandbox | AsyncSandbox) -> None: + name = _unique() + try: + if isinstance(sandbox, Sandbox): + await asyncio.to_thread( + sandbox.process.exec, command="cat", name=name, stdin=True + ) + await asyncio.to_thread( + sandbox.process.write_stdin, identifier=name, data="hello\n" + ) + await asyncio.to_thread( + sandbox.process.write_stdin, identifier=name, data=b"more\n" + ) + await asyncio.to_thread(sandbox.process.close_stdin, identifier=name) + finished = await asyncio.to_thread(sandbox.process.wait, identifier=name) + else: + await sandbox.process.exec(command="cat", name=name, stdin=True) + await sandbox.process.write_stdin(identifier=name, data="hello\n") + await sandbox.process.write_stdin(identifier=name, data=b"more\n") + await sandbox.process.close_stdin(identifier=name) + finished = await sandbox.process.wait(identifier=name) + finally: + await _kill_quietly(sandbox, name) + assert finished.stdin + assert (finished.status, finished.stdout) == ("completed", "hello\nmore\n") + + +@pytest.mark.asyncio +async def test_process_write_stdin_without_stdin_fails( + sandbox: Sandbox | AsyncSandbox, +) -> None: + name = _unique() + try: + with pytest.raises(SandboxAPIError) as raised: + if isinstance(sandbox, Sandbox): + await asyncio.to_thread( + sandbox.process.exec, command="sleep 300", name=name + ) + await asyncio.to_thread( + sandbox.process.write_stdin, identifier=name, data="hello\n" + ) + else: + await sandbox.process.exec(command="sleep 300", name=name) + await sandbox.process.write_stdin(identifier=name, data="hello\n") + finally: + await _kill_quietly(sandbox, name) + assert raised.value.status == 409 + + +@pytest.mark.asyncio +async def test_process_streams_logs_of_running_and_finished_processes( + sandbox: Sandbox | AsyncSandbox, +) -> None: + running, finished = _unique(), _unique() + try: + if isinstance(sandbox, Sandbox): + await asyncio.to_thread( + sandbox.process.exec, + command="sh -c 'sleep 1; echo a; echo b >&2; echo c'", + name=running, + ) + running_lines = await asyncio.to_thread( + lambda: list(sandbox.process.stream_logs(identifier=running)) + ) + await asyncio.to_thread( + sandbox.process.exec, command="sh -c 'echo a; echo b'", name=finished + ) + await asyncio.to_thread(sandbox.process.wait, identifier=finished) + finished_lines = await asyncio.to_thread( + lambda: list(sandbox.process.stream_logs(identifier=finished)) + ) + else: + await sandbox.process.exec( + command="sh -c 'sleep 1; echo a; echo b >&2; echo c'", name=running + ) + stream = await sandbox.process.stream_logs(identifier=running) + running_lines = [line async for line in stream] + await sandbox.process.exec(command="sh -c 'echo a; echo b'", name=finished) + await sandbox.process.wait(identifier=finished) + stream = await sandbox.process.stream_logs(identifier=finished) + finished_lines = [line async for line in stream] + finally: + for name in (running, finished): + await _kill_quietly(sandbox, name) + # Order holds within each stream, not between them. + by_stream = { + stream: [line.text for line in running_lines if line.stream == stream] + for stream in ("stdout", "stderr") + } + assert by_stream == {"stdout": ["a", "c"], "stderr": ["b"]} + assert len(running_lines) == 3 + # A finished process's logs replay from its start. + assert [(line.stream, line.text) for line in finished_lines] == [ + ("stdout", "a"), + ("stdout", "b"), + ] + + +@pytest.mark.asyncio +async def test_process_wait_timeout_leaves_process_running( + sandbox: Sandbox | AsyncSandbox, +) -> None: + name = _unique() + try: + with pytest.raises(SandboxProcessWaitTimeoutError) as raised: + if isinstance(sandbox, Sandbox): + await asyncio.to_thread( + sandbox.process.exec, command="sleep 300", name=name + ) + await asyncio.to_thread( + sandbox.process.wait, identifier=name, timeout=timedelta(seconds=1) + ) + else: + await sandbox.process.exec(command="sleep 300", name=name) + await sandbox.process.wait( + identifier=name, timeout=timedelta(seconds=1) + ) + if isinstance(sandbox, Sandbox): + process = await asyncio.to_thread(sandbox.process.get, identifier=name) + else: + process = await sandbox.process.get(identifier=name) + finally: + await _kill_quietly(sandbox, name) + assert raised.value.identifier == name + assert process.status == "running" + + +@pytest.mark.asyncio +async def test_process_exec_stream_output_then_exit( + sandbox: Sandbox | AsyncSandbox, +) -> None: + command = "sh -c 'printf \"a\\n\\nb\\n\"; echo e >&2; exit 2'" + if isinstance(sandbox, Sandbox): + + def run() -> list[SandboxProcessExecEvent]: + with sandbox.process.exec_stream(command=command) as stream: + return list(stream) + + events = await asyncio.to_thread(run) + else: + async with await sandbox.process.exec_stream(command=command) as stream: + events = [event async for event in stream] + *outputs, exit_event = events + assert isinstance(exit_event, SandboxProcessExecEventExit) + assert exit_event.process.exit_code == 2 + # Chunks follow the process's writes, so only what they join to is fixed. + joined = {"stdout": "", "stderr": ""} + for output in outputs: + assert isinstance(output, SandboxProcessExecEventOutput) + joined[output.stream] += output.text + assert joined == {"stdout": "a\n\nb\n", "stderr": "e\n"} + + +@pytest.mark.asyncio +async def test_process_exec_stream_prompt_before_newline( + sandbox: Sandbox | AsyncSandbox, +) -> None: + name = _unique() + command = 'sh -c \'printf "Continue? "; read answer; echo "got $answer"\'' + stdout = "" + exit_code: int | None = None + try: + if isinstance(sandbox, Sandbox): + + def run() -> tuple[str, int | None]: + text = "" + with sandbox.process.exec_stream( + command=command, name=name, stdin=True + ) as stream: + for event in stream: + if isinstance(event, SandboxProcessExecEventExit): + return text, event.process.exit_code + if event.stream != "stdout": + continue + text += event.text + # The process blocks on the read, so this only happens + # if the prompt arrived without its line finished. + if text == "Continue? ": + sandbox.process.write_stdin(identifier=name, data="yes\n") + return text, None + + stdout, exit_code = await asyncio.to_thread(run) + else: + async with await sandbox.process.exec_stream( + command=command, name=name, stdin=True + ) as stream: + async for event in stream: + if isinstance(event, SandboxProcessExecEventExit): + exit_code = event.process.exit_code + break + if event.stream != "stdout": + continue + stdout += event.text + if stdout == "Continue? ": + await sandbox.process.write_stdin(identifier=name, data="yes\n") + finally: + await _kill_quietly(sandbox, name) + assert (stdout, exit_code) == ("Continue? got yes\n", 0) + + +# The sandbox sends output as JSON strings, replacing invalid UTF-8 with U+FFFD +# before it reaches the SDK. +@pytest.mark.asyncio +async def test_process_non_utf8_output_replaced( + sandbox: Sandbox | AsyncSandbox, +) -> None: + command = "printf 'a\\377b'" + if isinstance(sandbox, Sandbox): + process = await asyncio.to_thread( + sandbox.process.exec, command=command, wait_for_completion=True + ) + logs = await asyncio.to_thread(sandbox.process.logs, identifier=process.pid) + + def stream_text() -> str: + with sandbox.process.exec_stream(command=command) as stream: + return "".join( + event.text + for event in stream + if isinstance(event, SandboxProcessExecEventOutput) + ) + + streamed = await asyncio.to_thread(stream_text) + else: + process = await sandbox.process.exec(command=command, wait_for_completion=True) + logs = await sandbox.process.logs(identifier=process.pid) + async with await sandbox.process.exec_stream(command=command) as stream: + streamed = "".join( + [ + event.text + async for event in stream + if isinstance(event, SandboxProcessExecEventOutput) + ] + ) + assert process.stdout == logs.stdout == streamed == "a�b" + + +@pytest.mark.asyncio +async def test_process_unknown_fails(sandbox: Sandbox | AsyncSandbox) -> None: + for call in ("logs", "get"): + with pytest.raises(SandboxAPIError) as raised: + if isinstance(sandbox, Sandbox): + await asyncio.to_thread( + getattr(sandbox.process, call), identifier=_unique() + ) + else: + await getattr(sandbox.process, call)(identifier=_unique()) + assert raised.value.status == 404 + + +@pytest_asyncio.fixture +async def fs_dir(sandbox: Sandbox | AsyncSandbox) -> str: + # A directory of the test's own, so tests on the shared sandbox never see + # each other's files. + path = f"/tmp/e2e-fs/{_unique()}" + if isinstance(sandbox, Sandbox): + await asyncio.to_thread(sandbox.fs.mkdir, path=path) + else: + await sandbox.fs.mkdir(path=path) + return path + + +@pytest.mark.asyncio +async def test_fs_writes_and_reads_including_missing_directories_and_relative( + sandbox: Sandbox | AsyncSandbox, fs_dir: str +) -> None: + relative = f"{_unique()}.txt" + paths = { + f"{fs_dir}/hello.txt": "hello", + f"{fs_dir}/a/b/c.txt": "nested", + # Resolved against the sandbox's working directory. + relative: "relative", + } + try: + if isinstance(sandbox, Sandbox): + for path, content in paths.items(): + await asyncio.to_thread(sandbox.fs.write, path=path, content=content) + read = { + path: await asyncio.to_thread(sandbox.fs.read, path=path) + for path in paths + } + else: + for path, content in paths.items(): + await sandbox.fs.write(path=path, content=content) + read = {path: await sandbox.fs.read(path=path) for path in paths} + finally: + if isinstance(sandbox, Sandbox): + await asyncio.to_thread(sandbox.fs.remove, path=relative) + else: + await sandbox.fs.remove(path=relative) + assert read == paths + + +@pytest.mark.asyncio +async def test_fs_read_missing_fails( + sandbox: Sandbox | AsyncSandbox, fs_dir: str +) -> None: + with pytest.raises(SandboxAPIError) as raised: + if isinstance(sandbox, Sandbox): + await asyncio.to_thread(sandbox.fs.read, path=f"{fs_dir}/missing.txt") + else: + await sandbox.fs.read(path=f"{fs_dir}/missing.txt") + assert raised.value.status == 404 + + +@pytest.mark.asyncio +async def test_fs_names_needing_escaping( + sandbox: Sandbox | AsyncSandbox, fs_dir: str +) -> None: + path = f"{fs_dir}/a b%#?.txt" + if isinstance(sandbox, Sandbox): + await asyncio.to_thread(sandbox.fs.write, path=path, content="escaped") + text = await asyncio.to_thread(sandbox.fs.read, path=path) + await asyncio.to_thread(sandbox.fs.write_bytes, path=path, content=b"\x01\x02") + data = await asyncio.to_thread(sandbox.fs.read_bytes, path=path) + listing = await asyncio.to_thread(sandbox.fs.list, path=fs_dir) + else: + await sandbox.fs.write(path=path, content="escaped") + text = await sandbox.fs.read(path=path) + await sandbox.fs.write_bytes(path=path, content=b"\x01\x02") + data = await sandbox.fs.read_bytes(path=path) + listing = await sandbox.fs.list(path=fs_dir) + assert (text, data) == ("escaped", b"\x01\x02") + assert [file.name for file in listing.files] == ["a b%#?.txt"] + + +@pytest.mark.asyncio +async def test_fs_round_trips_every_byte_value( + sandbox: Sandbox | AsyncSandbox, fs_dir: str +) -> None: + content = bytes(range(256)) + path = f"{fs_dir}/all.bin" + if isinstance(sandbox, Sandbox): + await asyncio.to_thread(sandbox.fs.write_bytes, path=path, content=content) + read = await asyncio.to_thread(sandbox.fs.read_bytes, path=path) + else: + await sandbox.fs.write_bytes(path=path, content=content) + read = await sandbox.fs.read_bytes(path=path) + assert read == content + + +@pytest.mark.asyncio +async def test_fs_write_bytes_with_permissions_kept_on_overwrite( + sandbox: Sandbox | AsyncSandbox, fs_dir: str +) -> None: + path = f"{fs_dir}/run.sh" + if isinstance(sandbox, Sandbox): + await asyncio.to_thread( + sandbox.fs.write_bytes, + path=path, + content=b"#!/bin/sh\necho ran\n", + permissions="0755", + ) + created_mode = await _mode(sandbox, path) + ran = await asyncio.to_thread( + sandbox.process.exec, command=path, wait_for_completion=True + ) + await asyncio.to_thread( + sandbox.fs.write_bytes, + path=path, + content=b"#!/bin/sh\necho again\n", + permissions="0700", + ) + else: + await sandbox.fs.write_bytes( + path=path, content=b"#!/bin/sh\necho ran\n", permissions="0755" + ) + created_mode = await _mode(sandbox, path) + ran = await sandbox.process.exec(command=path, wait_for_completion=True) + await sandbox.fs.write_bytes( + path=path, content=b"#!/bin/sh\necho again\n", permissions="0700" + ) + assert created_mode == "755" + assert ran.stdout.strip() == "ran" + # An existing file keeps its mode. + assert await _mode(sandbox, path) == "755" + + +@pytest.mark.asyncio +async def test_fs_writes_over_5mb_in_parts( + sandbox: Sandbox | AsyncSandbox, fs_dir: str +) -> None: + # Three parts, the last one partial, in a pattern that does not line up + # with the part size, so a part off by any number of bytes shows. + content = bytes(index % 251 for index in range(12 * 1024 * 1024 + 17)) + # Two bytes each in UTF-8, with a line number so parts out of order show. + text = "".join(f"{index} é\n" for index in range(1_000_000)) + binary_path, text_path = f"{fs_dir}/big.bin", f"{fs_dir}/big.txt" + if isinstance(sandbox, Sandbox): + await asyncio.to_thread( + sandbox.fs.write_bytes, + path=binary_path, + content=content, + permissions="0755", + ) + await asyncio.to_thread(sandbox.fs.write, path=text_path, content=text) + read = await asyncio.to_thread(sandbox.fs.read_bytes, path=binary_path) + read_text = await asyncio.to_thread(sandbox.fs.read, path=text_path) + else: + await sandbox.fs.write_bytes( + path=binary_path, content=content, permissions="0755" + ) + await sandbox.fs.write(path=text_path, content=text) + read = await sandbox.fs.read_bytes(path=binary_path) + read_text = await sandbox.fs.read(path=text_path) + assert read == content + assert read_text == text + # Unlike a single upload's, a multipart upload's permissions apply. + assert await _mode(sandbox, binary_path) == "755" + + +@pytest.mark.asyncio +async def test_fs_mkdir_with_permissions_making_parents_and_repeating( + sandbox: Sandbox | AsyncSandbox, fs_dir: str +) -> None: + if isinstance(sandbox, Sandbox): + await asyncio.to_thread( + sandbox.fs.mkdir, path=f"{fs_dir}/private", permissions="0700" + ) + for _ in range(2): + await asyncio.to_thread(sandbox.fs.mkdir, path=f"{fs_dir}/x/y/z") + listing = await asyncio.to_thread(sandbox.fs.list, path=f"{fs_dir}/x/y") + else: + await sandbox.fs.mkdir(path=f"{fs_dir}/private", permissions="0700") + for _ in range(2): + await sandbox.fs.mkdir(path=f"{fs_dir}/x/y/z") + listing = await sandbox.fs.list(path=f"{fs_dir}/x/y") + assert await _mode(sandbox, f"{fs_dir}/private") == "700" + assert [subdirectory.name for subdirectory in listing.subdirectories] == ["z"] + + +@pytest.mark.asyncio +async def test_fs_read_directory_and_list_file_fail( + sandbox: Sandbox | AsyncSandbox, fs_dir: str +) -> None: + file = f"{fs_dir}/f.txt" + if isinstance(sandbox, Sandbox): + await asyncio.to_thread(sandbox.fs.write, path=file, content="f") + with pytest.raises(SandboxError, match="is a directory"): + await asyncio.to_thread(sandbox.fs.read, path=fs_dir) + with pytest.raises(SandboxError, match="is a directory"): + await asyncio.to_thread(sandbox.fs.read_bytes, path=fs_dir) + with pytest.raises(SandboxError, match="is a file"): + await asyncio.to_thread(sandbox.fs.list, path=file) + else: + await sandbox.fs.write(path=file, content="f") + with pytest.raises(SandboxError, match="is a directory"): + await sandbox.fs.read(path=fs_dir) + with pytest.raises(SandboxError, match="is a directory"): + await sandbox.fs.read_bytes(path=fs_dir) + with pytest.raises(SandboxError, match="is a file"): + await sandbox.fs.list(path=file) + + +@pytest.mark.asyncio +async def test_fs_lists(sandbox: Sandbox | AsyncSandbox, fs_dir: str) -> None: + before = datetime.now(UTC) + if isinstance(sandbox, Sandbox): + await asyncio.to_thread( + sandbox.fs.write, path=f"{fs_dir}/a.txt", content="12345" + ) + await asyncio.to_thread(sandbox.fs.mkdir, path=f"{fs_dir}/sub") + listing = await asyncio.to_thread(sandbox.fs.list, path=fs_dir) + else: + await sandbox.fs.write(path=f"{fs_dir}/a.txt", content="12345") + await sandbox.fs.mkdir(path=f"{fs_dir}/sub") + listing = await sandbox.fs.list(path=fs_dir) + assert listing.path == fs_dir + assert [subdirectory.name for subdirectory in listing.subdirectories] == ["sub"] + [file] = listing.files + assert (file.name, file.path, file.size_bytes) == ("a.txt", f"{fs_dir}/a.txt", 5) + _assert_within_run(before, file.last_modified) + + +@pytest.mark.asyncio +async def test_fs_removes_non_empty_directory_only_recursively( + sandbox: Sandbox | AsyncSandbox, fs_dir: str +) -> None: + file, directory = f"{fs_dir}/f.txt", f"{fs_dir}/d" + if isinstance(sandbox, Sandbox): + await asyncio.to_thread(sandbox.fs.write, path=file, content="f") + await asyncio.to_thread(sandbox.fs.remove, path=file) + with pytest.raises(SandboxAPIError) as removed_file: + await asyncio.to_thread(sandbox.fs.read, path=file) + await asyncio.to_thread( + sandbox.fs.write, path=f"{directory}/g.txt", content="g" + ) + with pytest.raises(SandboxAPIError): + await asyncio.to_thread(sandbox.fs.remove, path=directory) + await asyncio.to_thread(sandbox.fs.remove, path=directory, recursive=True) + with pytest.raises(SandboxAPIError) as removed_directory: + await asyncio.to_thread(sandbox.fs.list, path=directory) + else: + await sandbox.fs.write(path=file, content="f") + await sandbox.fs.remove(path=file) + with pytest.raises(SandboxAPIError) as removed_file: + await sandbox.fs.read(path=file) + await sandbox.fs.write(path=f"{directory}/g.txt", content="g") + with pytest.raises(SandboxAPIError): + await sandbox.fs.remove(path=directory) + await sandbox.fs.remove(path=directory, recursive=True) + with pytest.raises(SandboxAPIError) as removed_directory: + await sandbox.fs.list(path=directory) + assert removed_file.value.status == removed_directory.value.status == 404 + + +@pytest.mark.asyncio +async def test_fs_finds_and_greps(sandbox: Sandbox | AsyncSandbox, fs_dir: str) -> None: + files = { + "a.ts": "a", + "b.js": "b", + "sub/c.ts": "c", + "n.txt": "one\nfind Needle here\nthree\n", + } + if isinstance(sandbox, Sandbox): + await asyncio.to_thread(sandbox.fs.write_tree, path=fs_dir, files=files) + found = await asyncio.to_thread(sandbox.fs.find, path=fs_dir, patterns=["*.ts"]) + directories = await asyncio.to_thread( + sandbox.fs.find, path=fs_dir, type="directory" + ) + grepped = await asyncio.to_thread(sandbox.fs.grep, path=fs_dir, query="needle") + with_context = await asyncio.to_thread( + sandbox.fs.grep, path=fs_dir, query="needle", context_lines=1 + ) + exact = await asyncio.to_thread( + sandbox.fs.grep, path=fs_dir, query="needle", case_sensitive=True + ) + else: + await sandbox.fs.write_tree(path=fs_dir, files=files) + found = await sandbox.fs.find(path=fs_dir, patterns=["*.ts"]) + directories = await sandbox.fs.find(path=fs_dir, type="directory") + grepped = await sandbox.fs.grep(path=fs_dir, query="needle") + with_context = await sandbox.fs.grep( + path=fs_dir, query="needle", context_lines=1 + ) + exact = await sandbox.fs.grep(path=fs_dir, query="needle", case_sensitive=True) + # Relative to the searched path, as observed from the server. + assert sorted(match.path for match in found.matches) == ["a.ts", "sub/c.ts"] + assert found.total == 2 + assert all(match.type == "directory" for match in directories.matches) + assert "sub" in [match.path for match in directories.matches] + [match] = grepped.matches + assert (match.path, match.line, match.text) == ("n.txt", 2, "find Needle here") + assert not match.context + assert [match.context for match in with_context.matches] == [ + "one\nfind Needle here\nthree" + ] + assert exact.total == 0 + + +@pytest.mark.asyncio +async def test_fs_copies_files_and_directories( + sandbox: Sandbox | AsyncSandbox, fs_dir: str +) -> None: + if isinstance(sandbox, Sandbox): + await asyncio.to_thread( + sandbox.fs.write, path=f"{fs_dir}/it's.txt", content="quoted" + ) + await asyncio.to_thread( + sandbox.fs.copy, source=f"{fs_dir}/it's.txt", destination=f"{fs_dir}/c.txt" + ) + await asyncio.to_thread( + sandbox.fs.write_tree, + path=f"{fs_dir}/src", + files={"a.txt": "a", "d/b.txt": "b"}, + ) + await asyncio.to_thread( + sandbox.fs.copy, source=f"{fs_dir}/src", destination=f"{fs_dir}/dst" + ) + copied = await asyncio.to_thread(sandbox.fs.read, path=f"{fs_dir}/c.txt") + nested = await asyncio.to_thread(sandbox.fs.read, path=f"{fs_dir}/dst/d/b.txt") + with pytest.raises(SandboxFileSystemCopyError) as raised: + await asyncio.to_thread( + sandbox.fs.copy, source=f"{fs_dir}/missing", destination=f"{fs_dir}/x" + ) + else: + await sandbox.fs.write(path=f"{fs_dir}/it's.txt", content="quoted") + await sandbox.fs.copy( + source=f"{fs_dir}/it's.txt", destination=f"{fs_dir}/c.txt" + ) + await sandbox.fs.write_tree( + path=f"{fs_dir}/src", files={"a.txt": "a", "d/b.txt": "b"} + ) + await sandbox.fs.copy(source=f"{fs_dir}/src", destination=f"{fs_dir}/dst") + copied = await sandbox.fs.read(path=f"{fs_dir}/c.txt") + nested = await sandbox.fs.read(path=f"{fs_dir}/dst/d/b.txt") + with pytest.raises(SandboxFileSystemCopyError) as raised: + await sandbox.fs.copy(source=f"{fs_dir}/missing", destination=f"{fs_dir}/x") + assert (copied, nested) == ("quoted", "b") + assert "missing" in raised.value.process.stderr + + +@pytest.mark.asyncio +async def test_fs_write_tree_leaves_other_files( + sandbox: Sandbox | AsyncSandbox, fs_dir: str +) -> None: + files = {"a.txt": "a", "d/e/b.txt": "b"} + paths = ["a.txt", "d/e/b.txt", "keep.txt"] + if isinstance(sandbox, Sandbox): + await asyncio.to_thread( + sandbox.fs.write, path=f"{fs_dir}/keep.txt", content="keep" + ) + await asyncio.to_thread(sandbox.fs.write_tree, path=fs_dir, files=files) + read = [ + await asyncio.to_thread(sandbox.fs.read, path=f"{fs_dir}/{path}") + for path in paths + ] + else: + await sandbox.fs.write(path=f"{fs_dir}/keep.txt", content="keep") + await sandbox.fs.write_tree(path=fs_dir, files=files) + read = [await sandbox.fs.read(path=f"{fs_dir}/{path}") for path in paths] + assert read == ["a", "b", "keep"] + + +# The watch covers only the directory itself unless recursive. +@pytest.mark.asyncio +async def test_fs_watches_directory( + sandbox: Sandbox | AsyncSandbox, fs_dir: str +) -> None: + event = await _watch_until_written(sandbox, f"{fs_dir}/top.txt", path=fs_dir) + assert {"CREATE", "WRITE"} & set(event.ops) + + +@pytest.mark.asyncio +async def test_fs_watches_recursively( + sandbox: Sandbox | AsyncSandbox, fs_dir: str +) -> None: + # The subdirectory exists before the watch, so the event shows the watch + # covers it. + if isinstance(sandbox, Sandbox): + await asyncio.to_thread( + sandbox.fs.write, path=f"{fs_dir}/sub/existing.txt", content="e" + ) + else: + await sandbox.fs.write(path=f"{fs_dir}/sub/existing.txt", content="e") + event = await _watch_until_written( + sandbox, f"{fs_dir}/sub/nested.txt", path=fs_dir, recursive=True + ) + assert {"CREATE", "WRITE"} & set(event.ops) + + +# Every push starts a build on the server, and CI runs this suite on every +# push, so a run pushes one image, built once and checked from both clients. +@dataclass +class _PushedImage: + name: str + builder: ImageBuilder + sandbox_name: str + """A sandbox created from the image.""" + + +@pytest.fixture(scope="session") +def pushed_image() -> Iterator[_PushedImage]: + _require_env() + # Three CI jobs push per run already; the 3.11 one leaves it to them. + if sys.version_info[:2] == (3, 11): + pytest.skip("the image push runs only on the newest Python's CI jobs") + client = _client() + image_name = _unique() + sandbox_name = _unique() + sandbox_created = False + # A pullable base, since the default sandbox image name only resolves on + # sandbox create, not in a build. The label makes each run's build context + # unique: identical pushes share one build on the server, which links + # their images so none can be deleted before the others. The sandbox API + # binary and entrypoint come from the builder's injection. + builder = ( + ImageBuilder.from_registry("debian:bookworm-slim") + .label({"e2e-image": image_name}) + .add_file("/hello.txt", "hello from the image") + .run_commands("echo built by RUN > /run.txt") + ) + try: + pushed = client.images.push(name=image_name, builder=builder) + assert (pushed.name, pushed.status) == (image_name, "BUILT") + sandbox_created = True + created = client.create( + name=sandbox_name, + image=f"{image_name}:latest", + # TODO: Same temporary region pin as the shared sandbox. + region="us-was-1", + labels=_LABELS, + ) + assert created.info.status == "DEPLOYED" + yield _PushedImage(name=image_name, builder=builder, sandbox_name=sandbox_name) + finally: + # The image cannot be deleted while a sandbox uses it. Cleanup errors + # are printed rather than raised, so they never hide the test's own + # failure. + try: + if sandbox_created: + _delete_and_wait(client, sandbox_name) + client.images.delete(name=image_name) + except ( + SandboxError, + httpx.HTTPError, + TimeoutError, + pydantic.ValidationError, + ) as err: + if not (isinstance(err, SandboxAPIError) and err.status == 404): + print(f"cleanup of image {image_name} failed: {err!r}") + finally: + client.close() + + +@pytest.mark.asyncio +async def test_image_pushed_is_found( + client: SandboxClient | AsyncSandboxClient, pushed_image: _PushedImage +) -> None: + name = pushed_image.name + if isinstance(client, SandboxClient): + info = await asyncio.to_thread(client.images.get_info, name=name) + listed = await asyncio.to_thread( + lambda: list(client.images.list(name_prefix=name)) + ) + tags = await asyncio.to_thread(lambda: list(client.images.list_tags(name=name))) + logs = await asyncio.to_thread(client.images.logs, name=name) + else: + info = await client.images.get_info(name=name) + listed = [image async for image in client.images.list(name_prefix=name)] + tags = [tag async for tag in client.images.list_tags(name=name)] + logs = await client.images.logs(name=name) + assert info.status == "BUILT" + assert [image.name for image in listed] == [name] + assert tags + # Not checked for content: some builds have no logs on the server, so only + # the order is asserted. + assert [line.timestamp for line in logs] == sorted(line.timestamp for line in logs) + assert pushed_image.builder.dockerfile().endswith( + "COPY --from=ghcr.io/blaxel-ai/sandbox:latest /sandbox-api " + "/usr/local/bin/sandbox-api\n" + 'ENTRYPOINT ["/usr/local/bin/sandbox-api"]\n' + ) + + +@pytest.mark.asyncio +async def test_image_sandbox_has_its_files_and_sandbox_api( + client: SandboxClient | AsyncSandboxClient, pushed_image: _PushedImage +) -> None: + paths = ["/hello.txt", "/run.txt"] + commands = ["test -x /usr/local/bin/sandbox-api", "tr '\\0' ' ' < /proc/1/cmdline"] + if isinstance(client, SandboxClient): + sandbox = await asyncio.to_thread(client.get, name=pushed_image.sandbox_name) + contents = [ + await asyncio.to_thread(sandbox.fs.read, path=path) for path in paths + ] + processes = [ + await asyncio.to_thread( + sandbox.process.exec, command=command, wait_for_completion=True + ) + for command in commands + ] + else: + async_sandbox = await client.get(name=pushed_image.sandbox_name) + contents = [await async_sandbox.fs.read(path=path) for path in paths] + processes = [ + await async_sandbox.process.exec(command=command, wait_for_completion=True) + for command in commands + ] + assert contents == ["hello from the image", "built by RUN\n"] + # The copied binary is where the copy put it, and runs as the entrypoint, + # which is what answers these calls at all. + assert processes[0].exit_code == 0 + assert re.match(r"/usr/local/bin/sandbox-api\b", processes[1].stdout) + + +def test_uses_http2_by_default(shared: _Shared) -> None: + versions: list[str] = [] + with _client() as client: + client.http_client.event_hooks["response"].append( + lambda response: versions.append(response.http_version) + ) + client.get_info(name=shared.name) + # The token exchange and the call itself. + assert versions == ["HTTP/2", "HTTP/2"] + + +def _unique() -> str: + return f"py-e2e-{secrets.token_hex(4)}" + + +def _assert_within_run(since: datetime, timestamp: datetime) -> None: + # Catches timestamp format and zone mistakes. + assert since - _CLOCK_SKEW <= timestamp <= datetime.now(UTC) + _CLOCK_SKEW + + +async def _kill_quietly(sandbox: Sandbox | AsyncSandbox, identifier: str) -> None: + # The test may have ended the process already, and cleanup never masks the + # test's own failure. + with contextlib.suppress(Exception): + if isinstance(sandbox, Sandbox): + await asyncio.to_thread(sandbox.process.kill, identifier=identifier) + else: + await sandbox.process.kill(identifier=identifier) + + +async def _mode(sandbox: Sandbox | AsyncSandbox, path: str) -> str: + command = f"stat -c %a '{path}'" + if isinstance(sandbox, Sandbox): + process = await asyncio.to_thread( + sandbox.process.exec, command=command, wait_for_completion=True + ) + else: + process = await sandbox.process.exec(command=command, wait_for_completion=True) + return process.stdout.strip() + + +# How many times a watch test writes its file before giving up, and how long it +# waits for the event in all. +_WATCH_ATTEMPTS = 20 +_WATCH_TIMEOUT = 60 + + +async def _watch_until_written( + sandbox: Sandbox | AsyncSandbox, + file: str, + *, + path: str, + recursive: bool | None = None, +) -> SandboxFileSystemWatchEvent: + # The server's watch starts at some unknown point after the request, so a + # change made too early is missed. Rather than sleeping, a writer repeats + # the change, each attempt a full round trip, until the watch reports it. + if isinstance(sandbox, Sandbox): + return await asyncio.to_thread( + _watch_until_written_sync, sandbox, file, path, recursive + ) + found = asyncio.Event() + + async def write() -> None: + for attempt in range(_WATCH_ATTEMPTS): + if found.is_set(): + return + await sandbox.fs.write(path=file, content=str(attempt)) + + stream = await sandbox.fs.watch(path=path, recursive=recursive) + writer = asyncio.create_task(write()) + try: + async with stream, asyncio.timeout(_WATCH_TIMEOUT): + async for event in stream: + if event.path == file: + return event + pytest.fail("the watch ended before an event for the file") + finally: + found.set() + await writer + + +def _watch_until_written_sync( + sandbox: Sandbox, file: str, path: str, recursive: bool | None +) -> SandboxFileSystemWatchEvent: + found = threading.Event() + with sandbox.fs.watch(path=path, recursive=recursive) as stream: + + def write() -> None: + for attempt in range(_WATCH_ATTEMPTS): + if found.is_set(): + return + sandbox.fs.write(path=file, content=str(attempt)) + # Every attempt went unseen, so the watch is ended from here, since + # it would otherwise wait for an event forever. Only the response + # can be closed from another thread, not the stream reading it. + if not found.is_set(): + stream._response.close() + + writer = threading.Thread(target=write) + writer.start() + try: + for event in stream: + if event.path == file: + return event + pytest.fail("the watch ended before an event for the file") + finally: + found.set() + writer.join() + + +def _delete_and_wait(client: SandboxClient, name: str) -> None: + # Cleanup only: a sandbox must be gone before its image can be deleted. + client.delete(name=name) + deadline = time.monotonic() + 120 + while True: + try: + info = client.get_info(name=name) + except SandboxAPIError as err: + if err.status == 404: + return + raise + # A deleted sandbox stays listed as TERMINATED rather than + # disappearing. + if info.status == "TERMINATED": + return + if time.monotonic() > deadline: + raise TimeoutError(f"sandbox {name} is still {info.status}") + time.sleep(2) + + +def _require_env() -> None: + if not API_KEY: + pytest.skip("BASETEN_E2E_TEST_API_KEY not set") + if not DOMAIN: + raise OSError( + "BASETEN_E2E_TEST_API_KEY is set but BASETEN_E2E_TEST_DOMAIN is missing" + ) + + +def _client() -> SandboxClient: + if not DEBUG: + return SandboxClient(api_key=API_KEY, base_url_override=f"https://api.{DOMAIN}") + return SandboxClient( + api_key=API_KEY, + base_url_override=f"https://api.{DOMAIN}", + http_client_override=httpx.Client( + transport=_DebugTransport(http2=True), timeout=_DEBUG_TIMEOUT + ), + close_http_client_on_close=True, + ) + + +def _async_client() -> AsyncSandboxClient: + if not DEBUG: + return AsyncSandboxClient( + api_key=API_KEY, base_url_override=f"https://api.{DOMAIN}" + ) + return AsyncSandboxClient( + api_key=API_KEY, + base_url_override=f"https://api.{DOMAIN}", + http_client_override=httpx.AsyncClient( + transport=_AsyncDebugTransport(http2=True), timeout=_DEBUG_TIMEOUT + ), + close_http_client_on_close=True, + ) + + +class _LoggedStream(httpx.SyncByteStream): + """Passes a response body through, logging it once fully read or closed.""" + + def __init__( + self, stream: httpx.SyncByteStream, log: Callable[[bytes], None] + ) -> None: + self._stream = stream + self._log = log + self._chunks: list[bytes] = [] + + def __iter__(self) -> Iterator[bytes]: + for chunk in self._stream: + self._chunks.append(chunk) + yield chunk + + def close(self) -> None: + self._stream.close() + self._log(b"".join(self._chunks)) + + +class _AsyncLoggedStream(httpx.AsyncByteStream): + """Async form of :class:`_LoggedStream`.""" + + def __init__( + self, stream: httpx.AsyncByteStream, log: Callable[[bytes], None] + ) -> None: + self._stream = stream + self._log = log + self._chunks: list[bytes] = [] + + async def __aiter__(self) -> AsyncIterator[bytes]: + async for chunk in self._stream: + self._chunks.append(chunk) + yield chunk + + async def aclose(self) -> None: + await self._stream.aclose() + self._log(b"".join(self._chunks)) + + +class _DebugTransport(httpx.HTTPTransport): + """Prints each request and its response, credentials masked.""" + + def handle_request(self, request: httpx.Request) -> httpx.Response: + response = super().handle_request(request) + response.stream = _LoggedStream( + cast(httpx.SyncByteStream, response.stream), + _response_logger(request, response), + ) + return response + + +class _AsyncDebugTransport(httpx.AsyncHTTPTransport): + """Async form of :class:`_DebugTransport`.""" + + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + response = await super().handle_async_request(request) + response.stream = _AsyncLoggedStream( + cast(httpx.AsyncByteStream, response.stream), + _response_logger(request, response), + ) + return response + + +def _response_logger( + request: httpx.Request, response: httpx.Response +) -> Callable[[bytes], None]: + """Return a function printing the request with the given response body.""" + headers = dict(request.headers) + if "authorization" in headers: + headers["authorization"] = "" + try: + request_body: Any = request.content.decode(errors="replace") + except httpx.RequestNotRead: + request_body = "" + + def log(body: bytes) -> None: + # The token exchange's response body is the minted token. + shown = ( + "" + if request.url.path == "/v1/token" + else body.decode(errors="replace") + ) + print( + request.method, + request.url, + { + "request": {"headers": headers, "body": request_body}, + "response": { + "status": response.status_code, + "http_version": response.http_version, + "headers": dict(response.headers), + "body": shown, + }, + }, + ) + + return log diff --git a/tests/sandbox/test_errors.py b/tests/sandbox/test_errors.py new file mode 100644 index 0000000..fe76974 --- /dev/null +++ b/tests/sandbox/test_errors.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +import json + +import httpx +import pytest + +import baseten.client.managementapi +import baseten.client.sandboxapi +from baseten.sandbox import SandboxAPIError, SandboxError, SandboxGatewayError +from baseten.sandbox._errors import api_error_from, converted_errors + + +class TestControlPlaneErrors: + def test_carries_status_code_message_and_details(self) -> None: + body = {"code": "NOT_FOUND", "message": "no such sandbox", "details": {"a": 1}} + err = _control_error(404, json.dumps(body)) + assert type(err) is SandboxAPIError + assert err.status == 404 + assert err.code == "NOT_FOUND" + assert err.message == "no such sandbox" + assert err.details == {"a": 1} + assert err.body == body + assert str(err) == "sandbox API error (HTTP 404) NOT_FOUND: no such sandbox" + assert isinstance(err, SandboxError) + + def test_takes_code_from_error_field_alongside_message(self) -> None: + err = _control_error( + 409, json.dumps({"error": "CONFLICT", "message": "already exists"}) + ) + assert err.code == "CONFLICT" + assert err.message == "already exists" + + def test_keeps_body_that_is_not_json_as_text(self) -> None: + err = _control_error(500, "upstream failed") + assert err.body == "upstream failed" + assert err.code is None + assert err.message is None + assert str(err) == "sandbox API error (HTTP 500)" + + def test_gateway_status_is_not_gateway_error(self) -> None: + assert type(_control_error(503, "")) is SandboxAPIError + + +class TestSandboxErrors: + def test_takes_message_from_error_field(self) -> None: + err = api_error_from( + baseten.client.sandboxapi.ResponseErrorResponse( + status_code=404, + error_response=baseten.client.sandboxapi.ErrorResponse( + error="process not found" + ), + ), + "exec", + ) + assert err.code is None + assert err.message == "process not found" + assert err.body == {"error": "process not found"} + assert str(err) == "sandbox API error (HTTP 404): process not found" + + @pytest.mark.parametrize("status", [502, 503, 504]) + def test_converts_edge_gateway_status_to_gateway_error(self, status: int) -> None: + err = api_error_from( + baseten.client.sandboxapi.ResponseError(status_code=status, body=""), + "exec", + ) + assert type(err) is SandboxGatewayError + assert err.status == status + assert err.body == "" + + +def test_converted_errors_chains_generated_error() -> None: + generated = baseten.client.managementapi.ResponseError(status_code=400, body="{}") + with pytest.raises(SandboxAPIError) as info, converted_errors("control"): + raise generated + assert info.value.__cause__ is generated + + +def test_converted_errors_passes_network_failure_through() -> None: + failure = httpx.ConnectError("refused") + with pytest.raises(httpx.ConnectError) as info, converted_errors("exec"): + raise failure + assert info.value is failure + + +def _control_error(status: int, body: str) -> SandboxAPIError: + return api_error_from( + baseten.client.managementapi.ResponseError(status_code=status, body=body), + "control", + ) diff --git a/tests/sandbox/test_fs.py b/tests/sandbox/test_fs.py new file mode 100644 index 0000000..a7e6d7c --- /dev/null +++ b/tests/sandbox/test_fs.py @@ -0,0 +1,853 @@ +from __future__ import annotations + +import asyncio +import email.parser +import email.policy +import inspect +import itertools +import json +import queue +import threading +from collections.abc import Awaitable, Callable +from concurrent.futures import ThreadPoolExecutor +from datetime import UTC, datetime, timedelta +from typing import Any + +import httpx +import pytest + +import baseten.sandbox._retry +from baseten.sandbox import ( + AsyncSandbox, + AsyncSandboxClient, + Sandbox, + SandboxAPIError, + SandboxClient, + SandboxError, + SandboxFileSystemCopyError, + SandboxFileSystemDirectory, + SandboxFileSystemFileInfo, + SandboxFileSystemFindMatch, + SandboxFileSystemFindResult, + SandboxFileSystemGrepMatch, + SandboxFileSystemGrepResult, + SandboxFileSystemSubdirectoryInfo, + SandboxFileSystemWatchEvent, +) + +_SANDBOX_URL = "https://sbx.test" + +_PART_SIZE = 5 * 1024 * 1024 + +_Handler = Callable[[httpx.Request], httpx.Response | Awaitable[httpx.Response]] + + +@pytest.fixture(autouse=True) +def short_backoff(monkeypatch: pytest.MonkeyPatch) -> None: + # Retries take milliseconds instead of seconds. + monkeypatch.setattr(baseten.sandbox._retry, "_backoff_base", 0.001) + monkeypatch.setattr(baseten.sandbox._retry, "_backoff_max", 0.002) + + +class _Fake: + """Fake sandbox answering through a handler, recording requests. + + Assertions belong in the test, not in the handler. + """ + + def __init__(self, handler: _Handler) -> None: + self.handler = handler + self.requests: list[httpx.Request] = [] + self._lock = threading.Lock() + + def _record(self, request: httpx.Request) -> httpx.Response | None: + if request.url.path == "/v1/token": + return httpx.Response( + 200, + json={ + "token": "token", + "expires_at": (datetime.now(UTC) + timedelta(hours=2)).isoformat(), + }, + ) + with self._lock: + self.requests.append(request) + return None + + def sandbox(self, executor: ThreadPoolExecutor | None = None) -> Sandbox: + def handle(request: httpx.Request) -> httpx.Response: + request.read() + token = self._record(request) + if token is not None: + return token + response = self.handler(request) + assert isinstance(response, httpx.Response) + return response + + client = SandboxClient( + api_key="key", + base_url_override="https://api.test", + http_client_override=httpx.Client(transport=httpx.MockTransport(handle)), + executor=executor, + ) + return client.sandbox_from_url(url=_SANDBOX_URL) + + def async_sandbox(self) -> AsyncSandbox: + async def handle(request: httpx.Request) -> httpx.Response: + await request.aread() + token = self._record(request) + if token is not None: + return token + response = self.handler(request) + if inspect.isawaitable(response): + response = await response + return response + + client = AsyncSandboxClient( + api_key="key", + base_url_override="https://api.test", + http_client_override=httpx.AsyncClient( + transport=httpx.MockTransport(handle) + ), + ) + return client.sandbox_from_url(url=_SANDBOX_URL) + + +def _ok(body: Any = None) -> httpx.Response: + return httpx.Response(200, json={"message": "ok"} if body is None else body) + + +def _replies(*responses: httpx.Response) -> _Handler: + """Answer with each response in turn, repeating the last.""" + remaining = list(responses) + + def handle(_: httpx.Request) -> httpx.Response: + response = remaining.pop(0) if len(remaining) > 1 else remaining[0] + # A response is single use, so each request gets a copy. + return httpx.Response( + response.status_code, headers=response.headers, content=response.content + ) + + return handle + + +def _multipart( + part: Callable[[int], httpx.Response | None] | None = None, +) -> _Handler: + """Fake the multipart upload endpoints. + + Uploads are numbered from 1 and each part is answered with + ``etag-``. *part*, when set, runs first for each part and answers it + instead when it returns a response. + """ + uploads = itertools.count(1) + + def handle(request: httpx.Request) -> httpx.Response: + path = request.url.path + if path.startswith("/filesystem-multipart/initiate/"): + return _ok({"uploadId": f"upload-{next(uploads)}"}) + if path.endswith("/part"): + part_number = int(request.url.params["partNumber"]) + response = None if part is None else part(part_number) + if response is not None: + return response + return _ok({"etag": f"etag-{part_number}", "partNumber": part_number}) + return _ok() + + return handle + + +def _path(request: httpx.Request) -> str: + """Give a request's path as sent, percent-encoding included.""" + return request.url.raw_path.decode().partition("?")[0] + + +def _body(request: httpx.Request) -> Any: + return json.loads(request.content) + + +def _form(request: httpx.Request) -> tuple[str | None, bytes, dict[str, str]]: + """Give a multipart request's file name, file content, and other fields.""" + message = email.parser.BytesParser(policy=email.policy.HTTP).parsebytes( + f"Content-Type: {request.headers['content-type']}\r\n\r\n".encode() + + request.content + ) + filename: str | None = None + content = b"" + fields: dict[str, str] = {} + for part in message.iter_parts(): + name = part.get_param("name", header="content-disposition") + payload = part.get_payload(decode=True) + assert isinstance(payload, bytes) + if name == "file": + filename, content = part.get_filename(), payload + else: + fields[str(name)] = payload.decode() + return filename, content, fields + + +def _uploaded(requests: list[httpx.Request]) -> bytes: + """Join the content of every part request in part number order.""" + parts: dict[int, bytes] = {} + for request in requests: + if _path(request).endswith("/part"): + _, content, _ = _form(request) + assert len(content) <= _PART_SIZE + parts[int(request.url.params["partNumber"])] = content + assert sorted(parts) == list(range(1, len(parts) + 1)) + return b"".join(parts[number] for number in sorted(parts)) + + +def _content(size: int) -> bytes: + return bytes(index % 251 for index in range(size)) + + +_DIRECTORY = { + "name": "tmp", + "path": "/tmp", + "files": [ + { + "name": "a.txt", + "path": "/tmp/a.txt", + "size": 5, + "permissions": "-rw-r--r--", + "owner": "root", + "group": "root", + "lastModified": "Wed, 01 Jan 2023 12:00:00 GMT", + } + ], + "subdirectories": [{"name": "d", "path": "/tmp/d"}], +} + + +def _process(**fields: Any) -> dict[str, Any]: + return { + "pid": "7", + "name": "", + "command": "cp", + "status": "completed", + "exitCode": 0, + "stdout": "", + "stderr": "", + "logs": "", + "workingDir": "/", + "startedAt": "Wed, 01 Jan 2023 12:00:00 GMT", + "completedAt": "Wed, 01 Jan 2023 12:00:01 GMT", + **fields, + } + + +def _watch_lines(*lines: str) -> httpx.Response: + return httpx.Response( + 200, + headers={"Content-Type": "text/plain"}, + content="".join(line + "\n" for line in lines), + ) + + +class TestSandboxFileSystem: + def test_reads_text_escaping_the_path(self) -> None: + fake = _Fake( + _replies( + _ok( + { + "name": "a", + "path": "/tmp/a b%.txt", + "content": "hello", + "group": "g", + "lastModified": "", + "owner": "o", + "permissions": "p", + "size": 5, + } + ) + ) + ) + assert fake.sandbox().fs.read(path="/tmp/a b%.txt") == "hello" + assert _path(fake.requests[0]) == "/filesystem/%2Ftmp%2Fa%20b%25.txt" + + def test_fails_to_read_a_directory(self) -> None: + fake = _Fake(_replies(_ok(_DIRECTORY))) + with pytest.raises(SandboxError, match="is a directory"): + fake.sandbox().fs.read(path="/tmp") + + def test_reads_bytes_as_octets(self) -> None: + fake = _Fake( + _replies( + httpx.Response( + 200, + headers={"Content-Type": "application/octet-stream"}, + content=b"\x00\x01\xff", + ) + ) + ) + assert fake.sandbox().fs.read_bytes(path="/tmp/bin") == b"\x00\x01\xff" + assert fake.requests[0].headers["Accept"] == "application/octet-stream" + + def test_fails_to_read_a_directory_as_bytes(self) -> None: + fake = _Fake(_replies(_ok(_DIRECTORY))) + with pytest.raises(SandboxError, match="is a directory"): + fake.sandbox().fs.read_bytes(path="/tmp") + + def test_writes_text_retrying_a_gateway_error(self) -> None: + fake = _Fake(_replies(httpx.Response(502, json={}), _ok())) + fake.sandbox().fs.write(path="/tmp/a", content="") + assert len(fake.requests) == 2 + assert fake.requests[1].method == "PUT" + # Empty content is still sent, so the file ends up empty. + assert _body(fake.requests[1]) == {"content": ""} + + def test_writes_text_up_to_5mb_of_utf8_in_one_request(self) -> None: + fake = _Fake(_replies(_ok())) + fake.sandbox().fs.write(path="/tmp/a", content="é" * (_PART_SIZE // 2)) + assert len(fake.requests) == 1 + + def test_writes_text_over_5mb_of_utf8_in_parts_without_permissions(self) -> None: + fake = _Fake(_multipart()) + content = "é" * (_PART_SIZE // 2 + 1) + fake.sandbox().fs.write(path="/tmp/a", content=content) + assert _path(fake.requests[0]) == "/filesystem-multipart/initiate/%2Ftmp%2Fa" + assert _body(fake.requests[0]) == {} + assert _uploaded(fake.requests) == content.encode() + + def test_writes_bytes_as_a_form_rebuilt_for_each_attempt(self) -> None: + fake = _Fake(_replies(httpx.Response(503, json={}), _ok())) + fake.sandbox().fs.write_bytes( + path="/tmp/run.sh", content=b"\x01\x02\x03", permissions="0755" + ) + assert len(fake.requests) == 2 + for request in fake.requests: + assert request.method == "PUT" + assert _path(request) == "/filesystem/%2Ftmp%2Frun.sh" + assert request.headers["Authorization"] == "Bearer token" + assert _form(request) == ( + "run.sh", + b"\x01\x02\x03", + {"permissions": "0755", "path": "/tmp/run.sh"}, + ) + + def test_writes_bytes_without_permissions_when_unset(self) -> None: + fake = _Fake(_replies(_ok())) + fake.sandbox().fs.write_bytes(path="/tmp/a", content=b"x") + assert _form(fake.requests[0])[2] == {"path": "/tmp/a"} + + def test_fails_a_byte_write_with_the_sandboxs_error(self) -> None: + fake = _Fake(_replies(httpx.Response(422, json={"error": "bad path"}))) + with pytest.raises(SandboxAPIError, match="bad path") as raised: + fake.sandbox().fs.write_bytes(path="/tmp/a", content=b"x") + assert raised.value.status == 422 + + def test_writes_exactly_5mb_of_bytes_in_one_request(self) -> None: + fake = _Fake(_replies(_ok())) + fake.sandbox().fs.write_bytes(path="/tmp/a", content=bytes(_PART_SIZE)) + assert len(fake.requests) == 1 + + def test_writes_bytes_over_5mb_in_parts_every_byte_once(self) -> None: + fake = _Fake(_multipart()) + content = _content(2 * _PART_SIZE + 17) + fake.sandbox().fs.write_bytes( + path="/tmp/big", content=content, permissions="0755" + ) + assert _body(fake.requests[0]) == {"permissions": "0755"} + assert _uploaded(fake.requests) == content + complete = fake.requests[-1] + assert _path(complete) == "/filesystem-multipart/upload-1/complete" + assert _body(complete) == { + "parts": [ + {"partNumber": 1, "etag": "etag-1"}, + {"partNumber": 2, "etag": "etag-2"}, + {"partNumber": 3, "etag": "etag-3"}, + ] + } + + def test_retries_a_part_after_a_gateway_error_rebuilding_its_form(self) -> None: + failed = threading.Event() + + def part(part_number: int) -> httpx.Response | None: + if part_number == 2 and not failed.is_set(): + failed.set() + return httpx.Response(502, json={}) + return None + + fake = _Fake(_multipart(part)) + content = bytes(_PART_SIZE) + b"\x07" + fake.sandbox().fs.write_bytes(path="/tmp/a", content=content) + part_two = [ + _form(request)[1] + for request in fake.requests + if _path(request).endswith("/part") + and request.url.params["partNumber"] == "2" + ] + assert part_two == [b"\x07", b"\x07"] + + def test_aborts_the_upload_when_a_part_fails_with_the_parts_error(self) -> None: + fake = _Fake( + _multipart( + lambda part_number: ( + httpx.Response(500, json={"error": "disk full"}) + if part_number == 1 + else None + ) + ) + ) + with pytest.raises(SandboxAPIError, match="disk full"): + fake.sandbox().fs.write_bytes(path="/tmp/a", content=bytes(_PART_SIZE + 1)) + paths = [_path(request) for request in fake.requests] + assert paths[-1] == "/filesystem-multipart/upload-1/abort" + assert not any(path.endswith("/complete") for path in paths) + + def test_aborts_when_completing_fails_keeping_that_error(self) -> None: + multipart = _multipart() + + def handle( + request: httpx.Request, + ) -> httpx.Response | Awaitable[httpx.Response]: + if request.url.path.endswith("/complete"): + return httpx.Response(500, json={"error": "complete failed"}) + if request.url.path.endswith("/abort"): + return httpx.Response(500, json={"error": "abort failed"}) + return multipart(request) + + fake = _Fake(handle) + with pytest.raises(SandboxAPIError, match="complete failed"): + fake.sandbox().fs.write_bytes(path="/tmp/a", content=bytes(_PART_SIZE + 1)) + assert _path(fake.requests[-1]) == "/filesystem-multipart/upload-1/abort" + + def test_keeps_at_most_2_parts_in_flight_per_sandbox_across_its_uploads( + self, + ) -> None: + uploads = 2 + parts_per_upload = 3 + lock = threading.Lock() + in_flight = 0 + max_in_flight = 0 + arrived: queue.Queue[None] = queue.Queue() + release = threading.Semaphore(0) + + def part(_: int) -> None: + nonlocal in_flight, max_in_flight + with lock: + in_flight += 1 + max_in_flight = max(max_in_flight, in_flight) + arrived.put(None) + release.acquire() + with lock: + in_flight -= 1 + + sandbox = _Fake(_multipart(part)).sandbox() + with ThreadPoolExecutor(max_workers=uploads) as callers: + writes = [ + callers.submit( + sandbox.fs.write_bytes, + path=f"/tmp/{index}", + content=bytes(2 * _PART_SIZE + 1), + ) + for index in range(uploads) + ] + # Two parts are held from the start, and each release lets one + # waiting part in, so parts beyond the limit would arrive while + # two are held. + arrived.get() + arrived.get() + for _ in range(uploads * parts_per_upload - 2): + release.release() + arrived.get() + release.release(2) + for write in writes: + write.result() + assert max_in_flight == 2 + + def test_gives_each_sandbox_its_own_limit_of_parts_in_flight(self) -> None: + arrived: queue.Queue[None] = queue.Queue() + release = threading.Event() + + def part(_: int) -> None: + arrived.put(None) + release.wait() + + fake = _Fake(_multipart(part)) + sandboxes = [fake.sandbox(), fake.sandbox()] + with ThreadPoolExecutor(max_workers=2) as callers: + writes = [ + callers.submit( + sandbox.fs.write_bytes, path="/tmp/a", content=bytes(_PART_SIZE + 1) + ) + for sandbox in sandboxes + ] + # Four parts in flight at once, two per sandbox. + for _ in range(4): + arrived.get() + release.set() + for write in writes: + write.result() + + def test_sends_parts_through_the_callers_thread_pool(self) -> None: + threads: set[str] = set() + lock = threading.Lock() + + def part(_: int) -> None: + with lock: + threads.add(threading.current_thread().name) + + with ThreadPoolExecutor(thread_name_prefix="caller-pool") as pool: + sandbox = _Fake(_multipart(part)).sandbox(executor=pool) + sandbox.fs.write_bytes(path="/tmp/a", content=bytes(3 * _PART_SIZE)) + assert any(name.startswith("caller-pool") for name in threads) + + def test_uploads_from_a_busy_callers_pool_without_deadlocking(self) -> None: + # The pool's one thread runs the write itself, so only the calling + # thread is left to send parts. + with ThreadPoolExecutor(max_workers=1) as pool: + fake = _Fake(_multipart()) + sandbox = fake.sandbox(executor=pool) + content = _content(2 * _PART_SIZE + 1) + pool.submit(sandbox.fs.write_bytes, path="/tmp/a", content=content).result( + timeout=30 + ) + assert _uploaded(fake.requests) == content + + def test_makes_a_directory_with_permissions_only_when_set(self) -> None: + fake = _Fake(_replies(_ok())) + fake.sandbox().fs.mkdir(path="/tmp/d") + fake.sandbox().fs.mkdir(path="/tmp/d", permissions="0700") + assert _body(fake.requests[0]) == {"isDirectory": True} + assert _body(fake.requests[1]) == {"isDirectory": True, "permissions": "0700"} + + def test_lists_a_directory(self) -> None: + fake = _Fake(_replies(_ok(_DIRECTORY))) + assert fake.sandbox().fs.list(path="/tmp") == SandboxFileSystemDirectory( + name="tmp", + path="/tmp", + files=[ + SandboxFileSystemFileInfo( + name="a.txt", + path="/tmp/a.txt", + size_bytes=5, + permissions="-rw-r--r--", + owner="root", + group="root", + last_modified=datetime(2023, 1, 1, 12, tzinfo=UTC), + ) + ], + subdirectories=[SandboxFileSystemSubdirectoryInfo(name="d", path="/tmp/d")], + ) + + def test_fails_to_list_a_file(self) -> None: + fake = _Fake( + _replies( + _ok( + { + "name": "a", + "path": "/tmp/a", + "content": "x", + "group": "g", + "lastModified": "", + "owner": "o", + "permissions": "p", + "size": 1, + } + ) + ) + ) + with pytest.raises(SandboxError, match="is a file"): + fake.sandbox().fs.list(path="/tmp/a") + + def test_removes_without_retrying_sending_recursive_only_when_set(self) -> None: + fake = _Fake(_replies(httpx.Response(502, json={}), _ok())) + with pytest.raises(SandboxAPIError): + fake.sandbox().fs.remove(path="/tmp/a") + fake.sandbox().fs.remove(path="/tmp/d", recursive=True) + assert len(fake.requests) == 2 + assert fake.requests[0].url.query == b"" + assert fake.requests[1].url.params["recursive"] == "true" + + def test_finds_entries_with_every_option_mapped(self) -> None: + fake = _Fake( + _replies(_ok({"matches": [{"path": "a.py", "type": "file"}], "total": 3})) + ) + result = fake.sandbox().fs.find( + path="/app", + type="file", + patterns=["*.py", "*.txt"], + max_results=0, + exclude_dirs=["node_modules", ".git"], + exclude_hidden=False, + ) + assert result == SandboxFileSystemFindResult( + matches=[SandboxFileSystemFindMatch(path="a.py", type="file")], total=3 + ) + assert dict(fake.requests[0].url.params) == { + "type": "file", + "patterns": "*.py,*.txt", + "maxResults": "0", + "excludeDirs": "node_modules,.git", + "excludeHidden": "false", + } + + def test_finds_with_no_options_sent_when_none_are_set(self) -> None: + fake = _Fake(_replies(_ok({"matches": [], "total": 0}))) + fake.sandbox().fs.find(path="/app") + assert fake.requests[0].url.query == b"" + + def test_greps_with_every_option_mapped(self) -> None: + fake = _Fake( + _replies( + _ok( + { + "matches": [ + { + "path": "a.py", + "line": 2, + "column": 5, + "text": "find Needle here", + "context": "one\nfind Needle here\nthree", + } + ], + "total": 1, + "query": "Needle", + } + ) + ) + ) + result = fake.sandbox().fs.grep( + path="/app", + query="Needle", + case_sensitive=True, + max_results=5, + file_pattern="*.py", + exclude_dirs=["dist"], + context_lines=1, + ) + assert result == SandboxFileSystemGrepResult( + matches=[ + SandboxFileSystemGrepMatch( + path="a.py", + line=2, + column=5, + text="find Needle here", + context="one\nfind Needle here\nthree", + ) + ], + total=1, + ) + assert dict(fake.requests[0].url.params) == { + "query": "Needle", + "caseSensitive": "true", + "maxResults": "5", + "filePattern": "*.py", + "excludeDirs": "dist", + "contextLines": "1", + } + + def test_copies_by_running_cp_with_both_paths_quoted_then_waiting(self) -> None: + def handle(request: httpx.Request) -> httpx.Response: + if request.method == "POST": + return _ok(_process(status="running", completedAt="")) + return _ok(_process()) + + fake = _Fake(handle) + fake.sandbox().fs.copy(source="/tmp/it's", destination="/tmp/b c") + assert _body(fake.requests[0])["command"] == "cp -r '/tmp/it'\\''s' '/tmp/b c'" + assert _path(fake.requests[1]) == "/process/7" + + @pytest.mark.parametrize( + ("finished", "message"), + [ + ( + {"status": "failed", "exitCode": 1, "stderr": "cp: no such file\n"}, + "ended failed: cp: no such file", + ), + ({"status": "killed", "exitCode": 0}, "ended killed: exit code 0"), + ], + ) + def test_fails_a_copy_whose_cp_does_not_complete( + self, finished: dict[str, Any], message: str + ) -> None: + fake = _Fake(_replies(_ok(_process(**finished)))) + with pytest.raises(SandboxFileSystemCopyError, match=message) as raised: + fake.sandbox().fs.copy(source="/a", destination="/b") + assert (raised.value.source, raised.value.destination) == ("/a", "/b") + assert raised.value.process.status == finished["status"] + + def test_writes_a_tree_retrying_a_gateway_error(self) -> None: + fake = _Fake(_replies(httpx.Response(504, json={}), _ok(_DIRECTORY))) + fake.sandbox().fs.write_tree(path="/app", files={"a/b.txt": "hi"}) + assert len(fake.requests) == 2 + assert _path(fake.requests[1]) == "/filesystem/tree/%2Fapp" + assert _body(fake.requests[1]) == {"files": {"a/b.txt": "hi"}} + + def test_watches_joining_each_events_directory_and_name(self) -> None: + fake = _Fake( + _replies( + _watch_lines( + '{"op":"CREATE","path":"/app/","name":"a"}', + "[keepalive]", + "", + '{"op":"WRITE|CHMOD","path":"/app","name":"b"}', + ) + ) + ) + with fake.sandbox().fs.watch( + path="/app", ignore=["node_modules", ".git"] + ) as stream: + events = list(stream) + assert events == [ + SandboxFileSystemWatchEvent(ops=["CREATE"], path="/app/a"), + SandboxFileSystemWatchEvent(ops=["WRITE", "CHMOD"], path="/app/b"), + ] + assert _path(fake.requests[0]) == "/watch/filesystem/%2Fapp" + assert fake.requests[0].url.params["ignore"] == "node_modules,.git" + + def test_watches_recursively_by_appending_a_suffix_to_the_path(self) -> None: + fake = _Fake(_replies(_watch_lines())) + with fake.sandbox().fs.watch(path="/app/", recursive=True) as stream: + list(stream) + assert _path(fake.requests[0]) == "/watch/filesystem/%2Fapp%2F%2A%2A" + + def test_fails_a_watch_on_an_unparseable_line(self) -> None: + fake = _Fake(_replies(_watch_lines("not json"))) + with ( + fake.sandbox().fs.watch(path="/app") as stream, + pytest.raises(json.JSONDecodeError), + ): + list(stream) + + +class TestAsyncSandboxFileSystem: + @pytest.mark.asyncio + async def test_reads_writes_and_lists(self) -> None: + def handle(request: httpx.Request) -> httpx.Response: + if request.url.path.startswith("/filesystem/tree/"): + return _ok(_DIRECTORY) + if request.method == "PUT": + return _ok() + if request.headers["Accept"] == "application/octet-stream": + return httpx.Response( + 200, + headers={"Content-Type": "application/octet-stream"}, + content=b"\x00\xff", + ) + return _ok(_DIRECTORY) + + fs = _Fake(handle).async_sandbox().fs + await fs.write(path="/tmp/a", content="hi") + await fs.write_bytes(path="/tmp/b", content=b"\x00\xff") + await fs.mkdir(path="/tmp/d") + await fs.write_tree(path="/tmp", files={"c": "x"}) + assert await fs.read_bytes(path="/tmp/b") == b"\x00\xff" + assert (await fs.list(path="/tmp")).subdirectories == [ + SandboxFileSystemSubdirectoryInfo(name="d", path="/tmp/d") + ] + + @pytest.mark.asyncio + async def test_writes_bytes_over_5mb_in_parts_every_byte_once(self) -> None: + fake = _Fake(_multipart()) + content = _content(2 * _PART_SIZE + 17) + await fake.async_sandbox().fs.write_bytes(path="/tmp/big", content=content) + assert _uploaded(fake.requests) == content + assert _path(fake.requests[-1]) == "/filesystem-multipart/upload-1/complete" + + @pytest.mark.asyncio + async def test_aborts_the_upload_when_a_part_fails(self) -> None: + fake = _Fake( + _multipart( + lambda part_number: ( + httpx.Response(500, json={"error": "disk full"}) + if part_number == 2 + else None + ) + ) + ) + with pytest.raises(SandboxAPIError, match="disk full"): + await fake.async_sandbox().fs.write_bytes( + path="/tmp/a", content=bytes(_PART_SIZE + 1) + ) + paths = [_path(request) for request in fake.requests] + assert paths[-1] == "/filesystem-multipart/upload-1/abort" + assert not any(path.endswith("/complete") for path in paths) + + @pytest.mark.asyncio + async def test_aborts_the_upload_when_cancelled(self) -> None: + arrived = asyncio.Event() + multipart = _multipart() + + async def held_part() -> httpx.Response: + arrived.set() + await asyncio.Event().wait() + raise AssertionError("never released") + + def handle( + request: httpx.Request, + ) -> httpx.Response | Awaitable[httpx.Response]: + if request.url.path.endswith("/part"): + return held_part() + return multipart(request) + + fake = _Fake(handle) + write = asyncio.create_task( + fake.async_sandbox().fs.write_bytes( + path="/tmp/a", content=bytes(_PART_SIZE + 1) + ) + ) + await arrived.wait() + write.cancel() + with pytest.raises(asyncio.CancelledError): + await write + assert _path(fake.requests[-1]) == "/filesystem-multipart/upload-1/abort" + + @pytest.mark.asyncio + async def test_keeps_at_most_2_parts_in_flight_per_sandbox(self) -> None: + in_flight = 0 + max_in_flight = 0 + arrived: asyncio.Queue[None] = asyncio.Queue() + release = asyncio.Semaphore(0) + multipart = _multipart() + + async def held_part(request: httpx.Request) -> httpx.Response: + nonlocal in_flight, max_in_flight + in_flight += 1 + max_in_flight = max(max_in_flight, in_flight) + arrived.put_nowait(None) + await release.acquire() + in_flight -= 1 + response = multipart(request) + assert isinstance(response, httpx.Response) + return response + + def handle( + request: httpx.Request, + ) -> httpx.Response | Awaitable[httpx.Response]: + if request.url.path.endswith("/part"): + return held_part(request) + return multipart(request) + + fs = _Fake(handle).async_sandbox().fs + writes = asyncio.gather( + *( + fs.write_bytes(path=f"/tmp/{index}", content=bytes(2 * _PART_SIZE + 1)) + for index in range(2) + ) + ) + # As the sync test: parts beyond the limit would arrive while two are + # held. + await arrived.get() + await arrived.get() + for _ in range(4): + release.release() + await arrived.get() + release.release() + release.release() + await writes + assert max_in_flight == 2 + + @pytest.mark.asyncio + async def test_copies_and_watches(self) -> None: + def handle(request: httpx.Request) -> httpx.Response: + if request.url.path.startswith("/watch/"): + return _watch_lines('{"op":"REMOVE","path":"/app","name":"a"}') + return _ok(_process()) + + fs = _Fake(handle).async_sandbox().fs + await fs.copy(source="/a", destination="/b") + async with await fs.watch(path="/app") as stream: + events = [event async for event in stream] + assert events == [SandboxFileSystemWatchEvent(ops=["REMOVE"], path="/app/a")] diff --git a/tests/sandbox/test_images.py b/tests/sandbox/test_images.py new file mode 100644 index 0000000..287209f --- /dev/null +++ b/tests/sandbox/test_images.py @@ -0,0 +1,762 @@ +from __future__ import annotations + +import asyncio +import io +import json +import tempfile +import zipfile +from collections.abc import Callable +from datetime import UTC, datetime, timedelta +from pathlib import Path +from typing import IO, Any + +import httpx +import pytest + +import baseten.sandbox._images +import baseten.sandbox._retry +from baseten.sandbox import ( + AsyncSandboxClient, + ImageBuilder, + ImageBuildError, + ImageCleanupResult, + ImageIgnoreFileOptions, + ImageIgnoreFileProcessorOptions, + ImageInfo, + ImageLibraryInfo, + ImageLibraryVolume, + ImageLogLine, + ImageTagInfo, + ImageUploadError, + SandboxAPIError, + SandboxClient, + SandboxError, + SandboxPort, +) + +_UPLOAD_URL = "https://storage.test/upload?signature=abc" + +_NO_WAIT = timedelta(0) + +_Reply = httpx.Response | Exception | dict[str, Any] + + +@pytest.fixture(autouse=True) +def short_backoff(monkeypatch: pytest.MonkeyPatch) -> None: + # Retries take milliseconds instead of seconds. + monkeypatch.setattr(baseten.sandbox._retry, "_backoff_base", 0.001) + monkeypatch.setattr(baseten.sandbox._retry, "_backoff_max", 0.002) + + +class _Api: + """Fake control plane and storage replying from per-route queues, recording requests. + + Each route replies in order and repeats its last reply. + """ + + def __init__(self) -> None: + self.requests: list[httpx.Request] = [] + self.routes: dict[tuple[str, str], list[_Reply]] = {} + + def route(self, method: str, path: str, *replies: _Reply) -> None: + self.routes[(method, path)] = list(replies) + + def handle(self, request: httpx.Request) -> httpx.Response: + if request.url.path == "/v1/token": + return httpx.Response( + 200, + json={ + "token": "token", + "expires_at": (datetime.now(UTC) + timedelta(hours=2)).isoformat(), + }, + ) + request.read() + self.requests.append(request) + replies = self.routes[(request.method, request.url.path)] + reply = replies.pop(0) if len(replies) > 1 else replies[0] + if isinstance(reply, Exception): + raise reply + if isinstance(reply, httpx.Response): + return httpx.Response( + reply.status_code, headers=reply.headers, content=reply.content + ) + return httpx.Response(200, json=reply) + + def client(self, **options: Any) -> SandboxClient: + return SandboxClient( + api_key="key", + base_url_override="https://api.test", + http_client_override=httpx.Client( + transport=httpx.MockTransport(self.handle) + ), + **options, + ) + + def async_client(self) -> AsyncSandboxClient: + async def handle(request: httpx.Request) -> httpx.Response: + await request.aread() + return self.handle(request) + + return AsyncSandboxClient( + api_key="key", + base_url_override="https://api.test", + http_client_override=httpx.AsyncClient( + transport=httpx.MockTransport(handle) + ), + ) + + def paths(self) -> list[str]: + return [f"{request.method} {request.url.path}" for request in self.requests] + + def upload(self) -> httpx.Request: + return next(r for r in self.requests if r.url.host == "storage.test") + + +def _image(status: str = "BUILT", **fields: Any) -> dict[str, Any]: + return { + "name": "img", + "status": status, + "created_at": "2026-10-07T12:00:00Z", + "size": 1024, + "tag_count": 2, + "tags": [], + **fields, + } + + +def _pushing(api: _Api, *statuses: str) -> None: + api.route( + "POST", + "/v1/sandboxes/images", + httpx.Response( + 202, json={"name": "img", "status": "UPLOADING", "upload_url": _UPLOAD_URL} + ), + ) + api.route("PUT", "/upload", httpx.Response(200)) + api.route( + "GET", "/v1/sandboxes/images/img", *(_image(status) for status in statuses) + ) + + +def _zip_names(request: httpx.Request) -> dict[str, bytes]: + with zipfile.ZipFile(io.BytesIO(request.content)) as archive: + return {name: archive.read(name) for name in archive.namelist()} + + +def _error(status: int) -> httpx.Response: + return httpx.Response(status, json={"error": "NOPE", "message": "nope"}) + + +class TestImageClientPush: + def test_pushes_files_uploading_a_zip_without_auth_then_waits(self) -> None: + api = _Api() + _pushing(api, "BUILDING", "BUILT") + info = api.client(team_id="team-1").images.push( + name="img", + files={"Dockerfile": "FROM scratch", "a.txt": b"a"}, + poll_interval=_NO_WAIT, + ) + assert info.status == "BUILT" + assert api.paths() == [ + "POST /v1/sandboxes/images", + "PUT /upload", + "GET /v1/sandboxes/images/img", + "GET /v1/sandboxes/images/img", + ] + push = api.requests[0] + assert push.url.params["team_id"] == "team-1" + assert json.loads(push.content) == {"name": "img"} + upload = api.upload() + assert "Authorization" not in upload.headers + assert upload.headers["Content-Type"] == "application/zip" + assert upload.headers["Content-Length"] == str(len(upload.content)) + assert str(upload.url) == _UPLOAD_URL + assert _zip_names(upload) == {"Dockerfile": b"FROM scratch", "a.txt": b"a"} + + def test_uploads_without_the_http_clients_own_headers_or_cookies(self) -> None: + api = _Api() + _pushing(api, "BUILT") + client = SandboxClient( + api_key="key", + base_url_override="https://api.test", + http_client_override=httpx.Client( + transport=httpx.MockTransport(api.handle), + headers={"Authorization": "Bearer other", "X-Other": "1"}, + cookies={"session": "secret"}, + ), + ) + client.images.push(name="img", files={"Dockerfile": "x"}) + upload = api.upload() + assert "Authorization" not in upload.headers + assert "X-Other" not in upload.headers + assert "Cookie" not in upload.headers + + def test_takes_built_at_the_first_poll(self) -> None: + api = _Api() + _pushing(api, "BUILT") + info = api.client().images.push(name="img", files={"Dockerfile": "x"}) + assert info == ImageInfo( + name="img", + status="BUILT", + created_at=datetime(2026, 10, 7, 12, tzinfo=UTC), + updated_at=None, + last_deployed_at=None, + size_bytes=1024, + tag_count=2, + ) + + def test_raises_image_build_error_when_the_build_fails(self) -> None: + api = _Api() + _pushing(api, "FAILED") + with pytest.raises(ImageBuildError, match="failed to build") as raised: + api.client().images.push(name="img", files={"Dockerfile": "x"}) + assert (raised.value.image_name, raised.value.status) == ("img", "FAILED") + assert not raised.value.timed_out + + def test_raises_a_timed_out_image_build_error_when_the_wait_runs_out(self) -> None: + api = _Api() + _pushing(api, "BUILDING") + with pytest.raises(ImageBuildError, match="still BUILDING") as raised: + api.client().images.push( + name="img", files={"Dockerfile": "x"}, timeout=_NO_WAIT + ) + assert raised.value.timed_out + + def test_rides_out_404_and_gateway_errors_while_waiting(self) -> None: + api = _Api() + _pushing(api) + api.route( + "GET", + "/v1/sandboxes/images/img", + _error(404), + _error(503), + httpx.RemoteProtocolError("Server disconnected"), + _image("BUILT"), + ) + info = api.client().images.wait_built(name="img", poll_interval=_NO_WAIT) + assert info.status == "BUILT" + assert len(api.requests) == 4 + + def test_stops_waiting_on_other_errors(self) -> None: + api = _Api() + api.route("GET", "/v1/sandboxes/images/img", _error(403)) + with pytest.raises(SandboxAPIError) as raised: + api.client().images.wait_built(name="img") + assert raised.value.status == 403 + + def test_returns_the_push_response_without_waiting_when_wait_is_false( + self, + ) -> None: + api = _Api() + _pushing(api) + info = api.client().images.push( + name="img", files={"Dockerfile": "x"}, wait=False + ) + assert (info.name, info.status) == ("img", "UPLOADING") + assert api.paths() == ["POST /v1/sandboxes/images", "PUT /upload"] + + def test_imports_a_registry_image_without_uploading(self) -> None: + api = _Api() + api.route( + "POST", + "/v1/sandboxes/images", + httpx.Response( + 202, + json={"name": "img", "status": "BUILDING", "image": "registry/img:1"}, + ), + ) + api.route("GET", "/v1/sandboxes/images/img", _image("BUILT")) + api.client().images.push( + name="img", registry_image="docker.io/library/python:3", docker_config="{}" + ) + assert json.loads(api.requests[0].content) == { + "name": "img", + "image": "docker.io/library/python:3", + "docker_config": "{}", + } + assert api.paths()[1:] == ["GET /v1/sandboxes/images/img"] + + def test_raises_image_upload_error_when_storage_rejects_without_waiting( + self, + ) -> None: + api = _Api() + _pushing(api) + api.route("PUT", "/upload", httpx.Response(403, text="Denied")) + with pytest.raises(ImageUploadError, match="push it again") as raised: + api.client().images.push(name="img", files={"Dockerfile": "x"}) + assert (raised.value.status, raised.value.body) == ( + 403, + "Denied", + ) + assert api.paths() == ["POST /v1/sandboxes/images", "PUT /upload"] + + def test_uploads_a_given_zip_as_is(self) -> None: + api = _Api() + _pushing(api, "BUILT") + archive = io.BytesIO() + with zipfile.ZipFile(archive, "w") as writer: + writer.writestr("Dockerfile", "FROM x") + api.client().images.push(name="img", zip=archive.getvalue()) + assert api.upload().content == archive.getvalue() + + @pytest.mark.parametrize( + ("sources", "message"), + [ + ({}, "got none"), + ( + {"files": {"Dockerfile": "x"}, "registry_image": "r"}, + "got registry_image, files", + ), + ( + {"files": {"Dockerfile": "x"}, "docker_config": "{}"}, + "docker_config applies only", + ), + ( + {"files": {"Dockerfile": "x"}, "default_ignore_file": lambda _: False}, + "apply only to directory", + ), + ({"files": {"Dockerfile": "x"}, "temp_file": False}, "temp_file applies"), + ({"files": {"a": "x"}}, "has no Dockerfile"), + ({"zip": b"not a zip"}, "not a readable zip"), + ], + ) + def test_rejects_a_bad_source_before_sending_anything( + self, sources: dict[str, Any], message: str + ) -> None: + api = _Api() + with pytest.raises((TypeError, ValueError), match=message): + api.client().images.push(name="img", **sources) + assert api.requests == [] + + def test_pushes_a_builder_as_a_zip_of_its_dockerfile_and_files(self) -> None: + api = _Api() + _pushing(api, "BUILT") + builder = ImageBuilder.from_registry("debian:bookworm-slim").add_file( + "/hello.txt", "hello" + ) + api.client().images.push(name="img", builder=builder) + assert _zip_names(api.upload()) == { + "Dockerfile": builder.dockerfile().encode(), + "hello.txt": b"hello", + } + + +class TestImageClientPushDirectory: + def _directory(self, root: Path) -> None: + (root / "Dockerfile").write_text("FROM x") + (root / "app.py").write_text("print()") + (root / ".git").mkdir() + (root / ".git" / "HEAD").write_text("ref") + (root / "build").mkdir() + (root / "build" / "out.bin").write_text("o") + + def test_leaves_out_the_defaults_without_a_dockerignore( + self, tmp_path: Path + ) -> None: + self._directory(tmp_path) + api = _Api() + _pushing(api, "BUILT") + api.client().images.push(name="img", directory=str(tmp_path)) + assert sorted(_zip_names(api.upload())) == [ + "Dockerfile", + "app.py", + "build/", + "build/out.bin", + ] + + def test_uses_default_ignore_file_instead_of_the_defaults( + self, tmp_path: Path + ) -> None: + self._directory(tmp_path) + api = _Api() + _pushing(api, "BUILT") + api.client().images.push( + name="img", + directory=str(tmp_path), + default_ignore_file=lambda options: options.rel_path == "app.py", + ) + assert "app.py" not in _zip_names(api.upload()) + assert ".git/HEAD" in _zip_names(api.upload()) + + def test_reads_a_dockerignore_through_the_processor_pruning_directories( + self, tmp_path: Path + ) -> None: + self._directory(tmp_path) + (tmp_path / ".dockerignore").write_text("build\n") + seen: list[ImageIgnoreFileProcessorOptions] = [] + + def processor( + options: ImageIgnoreFileProcessorOptions, + ) -> Callable[[ImageIgnoreFileOptions], bool]: + seen.append(options) + return lambda candidate: candidate.rel_path in options.contents.split() + + api = _Api() + _pushing(api, "BUILT") + api.client().images.push( + name="img", directory=str(tmp_path), ignore_file_processor=processor + ) + assert seen == [ + ImageIgnoreFileProcessorOptions( + path=str(tmp_path / ".dockerignore"), contents="build\n" + ) + ] + assert sorted(_zip_names(api.upload())) == [ + ".dockerignore", + ".git/", + ".git/HEAD", + "Dockerfile", + "app.py", + ] + + def test_fails_before_sending_when_a_dockerignore_has_no_processor( + self, tmp_path: Path + ) -> None: + self._directory(tmp_path) + (tmp_path / ".dockerignore").write_text("build\n") + api = _Api() + with pytest.raises(ValueError, match="ignore_file_processor is not set"): + api.client().images.push(name="img", directory=str(tmp_path)) + assert api.requests == [] + + def test_rejects_a_directory_without_a_dockerfile_before_sending( + self, tmp_path: Path + ) -> None: + api = _Api() + with pytest.raises(ValueError, match="has no Dockerfile at its root"): + api.client().images.push(name="img", directory=str(tmp_path)) + assert api.requests == [] + + @pytest.mark.parametrize(("temp_file", "spooled"), [(None, True), (False, False)]) + def test_zips_a_directory_into_a_temporary_file_closed_after_upload( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + temp_file: bool | None, + spooled: bool, + ) -> None: + self._directory(tmp_path) + files: list[IO[bytes]] = [] + original = tempfile.TemporaryFile + + def recording() -> IO[bytes]: + file = original() + files.append(file) + return file + + monkeypatch.setattr( + baseten.sandbox._images.tempfile, "TemporaryFile", recording + ) + api = _Api() + _pushing(api, "BUILT") + api.client().images.push( + name="img", directory=str(tmp_path), temp_file=temp_file + ) + assert len(files) == (1 if spooled else 0) + assert all(file.closed for file in files) + upload = api.upload() + assert upload.headers["Content-Length"] == str(len(upload.content)) + + def test_closes_the_temporary_file_when_the_upload_fails( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + self._directory(tmp_path) + files: list[IO[bytes]] = [] + original = tempfile.TemporaryFile + + def recording() -> IO[bytes]: + file = original() + files.append(file) + return file + + monkeypatch.setattr( + baseten.sandbox._images.tempfile, "TemporaryFile", recording + ) + api = _Api() + _pushing(api) + api.route("PUT", "/upload", httpx.Response(500)) + with pytest.raises(ImageUploadError): + api.client().images.push(name="img", directory=str(tmp_path)) + assert len(files) == 1 + assert files[0].closed + + def test_spools_a_builder_only_when_it_has_local_files( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + (tmp_path / "a.txt").write_text("a") + files: list[IO[bytes]] = [] + original = tempfile.TemporaryFile + + def recording() -> IO[bytes]: + file = original() + files.append(file) + return file + + monkeypatch.setattr( + baseten.sandbox._images.tempfile, "TemporaryFile", recording + ) + api = _Api() + _pushing(api, "BUILT") + base = ImageBuilder.from_registry("debian:bookworm-slim") + api.client().images.push(name="img", builder=base.add_file("/a", "a")) + assert files == [] + api.client().images.push( + name="img", builder=base.add_local_file(str(tmp_path / "a.txt"), "/a") + ) + assert len(files) == 1 + + +class TestImageClientListing: + def test_lists_images_across_pages_with_filters_and_team(self) -> None: + api = _Api() + api.route( + "GET", + "/v1/sandboxes/images", + { + "items": [{"name": "a", "status": "BUILT"}], + "pagination": {"has_more": True, "cursor": "c1"}, + }, + { + "items": [{"name": "b", "status": "FAILED"}], + "pagination": {"has_more": False}, + }, + ) + images = list( + api.client(team_id="t").images.list( + name_prefix="a", sort="name:asc", page_size=1 + ) + ) + assert [(image.name, image.status) for image in images] == [ + ("a", "BUILT"), + ("b", "FAILED"), + ] + assert dict(api.requests[0].url.params) == { + "team_id": "t", + "limit": "1", + "sort": "name:asc", + "q": "a", + } + assert api.requests[1].url.params["cursor"] == "c1" + + def test_stops_listing_on_a_repeated_cursor(self) -> None: + api = _Api() + api.route( + "GET", + "/v1/sandboxes/images", + {"items": [], "pagination": {"has_more": True, "cursor": "same"}}, + ) + with pytest.raises(SandboxError, match="image list returned a repeated cursor"): + list(api.client().images.list()) + + def test_lists_tags_deletes_a_tag_gets_deletes_and_cleans_up(self) -> None: + api = _Api() + api.route( + "GET", + "/v1/sandboxes/images/img/tags", + { + "items": [{"name": "abc123", "size": 10}], + "pagination": {"has_more": False}, + }, + ) + api.route("DELETE", "/v1/sandboxes/images/img/tags/abc123", _image()) + api.route("GET", "/v1/sandboxes/images/img", _image()) + api.route("DELETE", "/v1/sandboxes/images/img", _image("BUILT", tag_count=0)) + api.route( + "POST", "/v1/sandboxes/cleanup_images", {"deleted": 3, "message": "done"} + ) + images = api.client().images + assert list(images.list_tags(name="img", name_prefix="abc")) == [ + ImageTagInfo(name="abc123", created_at=None, updated_at=None, size_bytes=10) + ] + assert api.requests[0].url.params["q"] == "abc" + assert images.delete_tag(name="img", tag="abc123").tag_count == 2 + assert images.get_info(name="img").size_bytes == 1024 + assert images.delete(name="img").tag_count == 0 + assert images.cleanup() == ImageCleanupResult(deleted=3, message="done") + + def test_lists_library_images(self) -> None: + api = _Api() + api.route( + "GET", + "/v1/sandboxes/library_images", + { + "items": [ + { + "name": "base-image", + "image": "baseten/base-image:latest", + "ports": [{"target": 8080, "protocol": "HTTP"}], + "hidden": False, + "creation_options": { + "extra_args": {"kernel": "x"}, + "volumes": [{"name": "v", "mount_path": "/data"}], + }, + } + ] + }, + ) + assert api.client().images.list_library() == [ + ImageLibraryInfo( + name="base-image", + display_name=None, + description=None, + long_description=None, + image="baseten/base-image:latest", + memory=None, + ports=[SandboxPort(target=8080, protocol="HTTP")], + categories=[], + tags=[], + url=None, + icon=None, + icon_light=None, + icon_dark=None, + enterprise=False, + creation_extra_args={"kernel": "x"}, + creation_volumes=[ + ImageLibraryVolume( + name="v", + mount_path="/data", + type=None, + size_mb=None, + read_only=False, + ) + ], + ) + ] + assert "team_id" not in api.requests[0].url.params + + +class TestImageClientLogs: + def _log(self, second: int) -> dict[str, Any]: + return { + "timestamp": f"2026-10-07T12:00:{second:02d}Z", + "message": f"line {second}", + "severity": 9, + } + + def test_gets_every_page_oldest_first_with_the_end_time_pinned( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(baseten.sandbox._images, "_LOGS_PAGE_SIZE", 2) + api = _Api() + api.route( + "GET", + "/v1/sandboxes/images/img/logs", + {"logs": [self._log(3), self._log(2)], "total_count": 3}, + {"logs": [self._log(1)], "total_count": 3}, + ) + lines = api.client().images.logs(name="img") + assert lines == [ + ImageLogLine( + timestamp=datetime(2026, 10, 7, 12, 0, second, tzinfo=UTC), + severity=9, + text=f"line {second}", + ) + for second in (1, 2, 3) + ] + params = [request.url.params for request in api.requests] + assert [p["offset"] for p in params] == ["0", "2"] + assert params[0]["end_time"] == params[1]["end_time"] + assert "start_time" not in params[0] + + def test_sends_the_given_range_and_stops_at_a_short_page(self) -> None: + api = _Api() + api.route( + "GET", + "/v1/sandboxes/images/img/logs", + {"logs": [self._log(1)], "total_count": 5000}, + ) + start = datetime(2026, 10, 7, 11, tzinfo=UTC) + end = datetime(2026, 10, 7, 12, tzinfo=UTC) + api.client().images.logs(name="img", start_time=start, end_time=end) + assert len(api.requests) == 1 + params = api.requests[0].url.params + assert datetime.fromisoformat(params["start_time"]) == start + assert datetime.fromisoformat(params["end_time"]) == end + assert params["limit"] == "1000" + + def test_stops_at_the_11000_most_recent_lines(self) -> None: + api = _Api() + api.route( + "GET", + "/v1/sandboxes/images/img/logs", + {"logs": [self._log(1)] * 1000, "total_count": 50_000}, + ) + lines = api.client().images.logs(name="img") + assert len(lines) == 11_000 + assert api.requests[-1].url.params["offset"] == "10000" + + +class TestAsyncImageClient: + @pytest.mark.asyncio + async def test_pushes_a_directory_with_an_async_ignore_then_waits( + self, tmp_path: Path + ) -> None: + (tmp_path / "Dockerfile").write_text("FROM x") + (tmp_path / "keep.txt").write_text("k") + (tmp_path / "skip.txt").write_text("s") + (tmp_path / ".dockerignore").write_text("skip.txt") + + async def processor( + options: ImageIgnoreFileProcessorOptions, + ) -> Callable[[ImageIgnoreFileOptions], Any]: + async def ignore(candidate: ImageIgnoreFileOptions) -> bool: + return candidate.rel_path == options.contents + + return ignore + + api = _Api() + _pushing(api, "BUILDING", "BUILT") + info = await api.async_client().images.push( + name="img", + directory=str(tmp_path), + ignore_file_processor=processor, + poll_interval=_NO_WAIT, + ) + assert info.status == "BUILT" + upload = api.upload() + assert sorted(_zip_names(upload)) == [".dockerignore", "Dockerfile", "keep.txt"] + assert upload.headers["Content-Length"] == str(len(upload.content)) + assert "Authorization" not in upload.headers + + @pytest.mark.asyncio + async def test_times_out_waiting_as_an_image_build_error(self) -> None: + api = _Api() + api.route("GET", "/v1/sandboxes/images/img", _image("BUILDING")) + with pytest.raises(ImageBuildError) as raised: + await api.async_client().images.wait_built(name="img", timeout=_NO_WAIT) + assert raised.value.timed_out + + @pytest.mark.asyncio + async def test_leaves_the_callers_own_timeout_alone(self) -> None: + api = _Api() + api.route("GET", "/v1/sandboxes/images/img", _image("BUILDING")) + with pytest.raises(TimeoutError): + async with asyncio.timeout(0.05): + await api.async_client().images.wait_built( + name="img", timeout=None, poll_interval=timedelta(seconds=10) + ) + + @pytest.mark.asyncio + async def test_lists_and_gets_logs(self) -> None: + api = _Api() + api.route( + "GET", + "/v1/sandboxes/images", + { + "items": [{"name": "a", "status": "BUILT"}], + "pagination": {"has_more": False}, + }, + ) + api.route( + "GET", + "/v1/sandboxes/images/a/logs", + { + "logs": [ + {"timestamp": "2026-10-07T12:00:00Z", "message": "m", "severity": 9} + ], + "total_count": 1, + }, + ) + images = api.async_client().images + assert [image.name async for image in images.list()] == ["a"] + assert [line.text for line in await images.logs(name="a")] == ["m"] diff --git a/tests/sandbox/test_images_builder.py b/tests/sandbox/test_images_builder.py new file mode 100644 index 0000000..ea25ce6 --- /dev/null +++ b/tests/sandbox/test_images_builder.py @@ -0,0 +1,351 @@ +from __future__ import annotations + +import os +import sys +from collections.abc import Callable +from pathlib import Path + +import pytest + +from baseten.sandbox import ImageBuilder +from baseten.sandbox._zip import ZipEntry + +_INJECTED = [ + "COPY --from=ghcr.io/blaxel-ai/sandbox:latest /sandbox-api /usr/local/bin/sandbox-api", + 'ENTRYPOINT ["/usr/local/bin/sandbox-api"]', +] + +# Windows reads every file as 0666 and has no exec bits to keep. +_KEEPS_MODES = sys.platform != "win32" + + +def _expected(lines: list[str], base: str = "debian:bookworm-slim") -> str: + # The Dockerfile a builder from base pushes with lines, injection included. + return "\n".join([f"FROM {base}", *lines, *_INJECTED]) + "\n" + + +def _base() -> ImageBuilder: + return ImageBuilder.from_registry("debian:bookworm-slim") + + +def test_writes_each_core_instruction_then_injects_the_sandbox_api() -> None: + builder = ( + _base() + .workdir("/app") + .run_commands("echo one", "echo two") + .env({"A": "1", "B": "two words"}) + .copy("src dir", "/app/src") + .expose(8080, 9090) + .user("app") + .label({"team": "infra"}) + .arg("VERSION") + .arg("CHANNEL", "stable") + ) + assert builder.base_image == "debian:bookworm-slim" + assert builder.dockerfile() == _expected( + [ + "WORKDIR /app", + "RUN echo one", + "RUN echo two", + 'ENV A="1"', + 'ENV B="two words"', + 'COPY ["src dir","/app/src"]', + "EXPOSE 8080", + "EXPOSE 9090", + "USER app", + 'LABEL team="infra"', + "ARG VERSION", + 'ARG CHANNEL="stable"', + ] + ) + + +def test_escapes_quotes_and_backslashes_in_values_leaving_dollar() -> None: + builder = ( + _base() + .env({"PATH": "/opt/bin:$PATH", "QUOTE": 'say "hi" \\ bye'}) + .label({"note": 'a "b"'}) + .arg("X", 'c"d') + ) + assert builder.dockerfile() == _expected( + [ + 'ENV PATH="/opt/bin:$PATH"', + 'ENV QUOTE="say \\"hi\\" \\\\ bye"', + 'LABEL note="a \\"b\\""', + 'ARG X="c\\"d"', + ] + ) + + +@pytest.mark.parametrize( + ("build", "message"), + [ + (lambda: ImageBuilder.from_registry("a\nb"), "base_image must not"), + (lambda: _base().workdir("/a\n/b"), "path must not"), + (lambda: _base().run_commands("a\nb"), "command must not"), + (lambda: _base().env({"A": "1\r\n2"}), "env value must not"), + (lambda: _base().copy("a\n", "/b"), "source must not"), + (lambda: _base().entrypoint("a\nb"), "entrypoint argument must not"), + (lambda: _base().pip_install("a\nb"), "argument must not"), + ], +) +def test_raises_on_a_newline_anywhere_but_raw_lines( + build: Callable[[], ImageBuilder], message: str +) -> None: + with pytest.raises(ValueError, match=message): + build() + + +def test_raises_on_a_name_that_cannot_be_written_unquoted() -> None: + with pytest.raises(ValueError, match="env name 'A B'"): + _base().env({"A B": "1"}) + with pytest.raises(ValueError, match="label key 'a=b'"): + _base().label({"a=b": "1"}) + with pytest.raises(ValueError, match="arg name ''"): + _base().arg("") + + +def test_rejects_ports_that_are_not_between_1_and_65535() -> None: + with pytest.raises(ValueError, match="port 0 is not between 1 and 65535"): + _base().expose(0) + with pytest.raises(ValueError, match="port 65536"): + _base().expose(65536) + + +def test_appends_raw_lines_verbatim_newlines_included() -> None: + heredoc = "RUN < None: + builder = _base().entrypoint("/bin/sh", "-c", 'echo "hi"') + assert builder.dockerfile() == ( + "\n".join( + [ + "FROM debian:bookworm-slim", + 'ENTRYPOINT ["/bin/sh","-c","echo \\"hi\\""]', + _INJECTED[0], + ] + ) + + "\n" + ) + assert _base().entrypoint().dockerfile() == _expected([]) + + +def test_counts_an_entrypoint_in_raw_lines_in_any_case_but_not_in_a_comment() -> None: + assert ( + _INJECTED[1] not in _base().dockerfile_lines('entrypoint ["/run"]').dockerfile() + ) + assert ( + _INJECTED[1] in _base().dockerfile_lines('# ENTRYPOINT ["/run"]').dockerfile() + ) + + +def test_skips_copying_the_sandbox_api_when_already_brought_in() -> None: + assert ( + _INJECTED[0] + not in _base() + .run_commands("curl -o /usr/local/bin/sandbox-api https://x") + .dockerfile() + ) + assert ( + _INJECTED[0] + not in ImageBuilder.from_registry("ghcr.io/blaxel-ai/sandbox:v1").dockerfile() + ) + + +def test_copies_the_sandbox_api_from_the_given_image() -> None: + builder = ImageBuilder.from_registry( + "debian:bookworm-slim", sandbox_api_image="registry.example/sandbox:v2" + ) + assert ( + "COPY --from=registry.example/sandbox:v2 /sandbox-api /usr/local/bin/sandbox-api" + in builder.dockerfile() + ) + + +def test_never_changes_each_call_returns_a_new_builder() -> None: + shared = _base().pip_install("numpy") + with_torch = shared.pip_install("torch") + with_jax = shared.pip_install("jax") + assert shared.dockerfile() == _expected(["RUN pip install numpy"]) + assert with_torch.dockerfile() == _expected( + ["RUN pip install numpy", "RUN pip install torch"] + ) + assert with_jax.dockerfile() == _expected( + ["RUN pip install numpy", "RUN pip install jax"] + ) + + +def test_writes_each_package_managers_command_quoting_arguments() -> None: + builder = ( + _base() + .pip_install("--pre", "numpy>=2", "it's") + .apt_install("git", "curl") + .apk_add("git") + .npm_install("-g", "typescript@5") + .npm_install() + .gem_install("rails") + .cargo_install("--locked", "ripgrep") + .go_install("golang.org/x/tools/gopls@latest", "example.com/cmd@v1") + .composer_install("laravel/framework") + .uv_install("ruff") + .pipx_install("black", "httpie") + ) + assert builder.dockerfile() == _expected( + [ + "RUN pip install --pre 'numpy>=2' 'it'\\''s'", + "RUN apt-get update && apt-get install -y --no-install-recommends git curl && rm -rf /var/lib/apt/lists/*", + "RUN apk add --no-cache git", + "RUN npm install -g typescript@5", + "RUN npm install", + "RUN gem install --no-document rails", + "RUN cargo install --locked ripgrep", + "RUN go install golang.org/x/tools/gopls@latest && go install example.com/cmd@v1", + "RUN composer require laravel/framework", + "RUN uv pip install --system ruff", + "RUN pipx install black && pipx install httpie", + ] + ) + + +def test_appends_nothing_for_a_package_manager_without_arguments_except_npm() -> None: + builder = ( + _base() + .pip_install() + .apt_install() + .apk_add() + .gem_install() + .cargo_install() + .go_install() + .composer_install() + .uv_install() + .pipx_install() + .run_commands() + .env({}) + .expose() + ) + assert builder.dockerfile() == _expected([]) + + +def test_names_context_entries_by_basename_suffixing_repeats_and_reserved() -> None: + builder = ( + _base() + .add_file("/etc/a/config.json", "1") + .add_file("/etc/b/config.json", "2") + .add_file("/etc/c/config.json", "3") + .add_file("/srv/Makefile", "4") + .add_file("/opt/Makefile", "5") + .add_file("/x/Dockerfile", "6") + .add_file("/x/.dockerignore", "7") + ) + assert builder.dockerfile() == _expected( + [ + 'COPY ["config.json","/etc/a/config.json"]', + 'COPY ["config (1).json","/etc/b/config.json"]', + 'COPY ["config (2).json","/etc/c/config.json"]', + 'COPY ["Makefile","/srv/Makefile"]', + 'COPY ["Makefile (1)","/opt/Makefile"]', + 'COPY ["Dockerfile (1)","/x/Dockerfile"]', + 'COPY [".dockerignore (1)","/x/.dockerignore"]', + ] + ) + + +def test_uses_an_explicit_context_name_exactly_raising_if_taken_or_reserved() -> None: + builder = _base().add_file("/app/run.sh", "x", context_name="start.sh") + assert 'COPY ["start.sh","/app/run.sh"]' in builder.dockerfile() + with pytest.raises( + ValueError, match="context name start.sh is reserved or already used" + ): + builder.add_file("/b", "y", context_name="start.sh") + with pytest.raises(ValueError, match="context name Dockerfile is reserved"): + _base().add_file("/b", "y", context_name="Dockerfile") + with pytest.raises(ValueError, match="context name 'a/b' must be one path segment"): + _base().add_file("/b", "y", context_name="a/b") + + +def test_needs_a_context_name_for_a_destination_ending_in_a_slash() -> None: + with pytest.raises(ValueError, match="destination /app/ has no file name"): + _base().add_file("/app/", "x") + assert ( + 'COPY ["x.txt","/app/"]' + in _base().add_file("/app/", "x", context_name="x.txt").dockerfile() + ) + + +def test_zips_the_dockerfile_and_every_entry_reading_local_files_at_push( + tmp_path: Path, +) -> None: + tool = tmp_path / "tool.sh" + tool.write_text("before") + tool.chmod(0o755) + (tmp_path / "conf" / "nested").mkdir(parents=True) + (tmp_path / "conf" / "nested" / "a.txt").write_text("a") + builder = ( + _base() + .add_file("/etc/text.txt", "text") + .add_file("/etc/bytes.bin", b"\x00\xff") + .add_file("/usr/local/bin/run", b"#!/bin/sh", mode=0o755) + .add_local_file(str(tool), "/usr/local/bin/tool") + .add_local_dir(str(tmp_path / "conf"), "/etc/conf") + ) + # Read when pushed, not when added. + tool.write_text("after") + entries = {entry.path: entry for entry in builder._zip_entries()} + assert list(entries) == [ + "Dockerfile", + "text.txt", + "bytes.bin", + "run", + "tool.sh", + "conf", + "conf/nested", + "conf/nested/a.txt", + ] + assert entries["Dockerfile"].data == builder.dockerfile().encode() + assert entries["text.txt"] == ZipEntry(path="text.txt", data=b"text", mode=0o644) + assert entries["bytes.bin"].data == b"\x00\xff" + assert entries["run"].mode == 0o755 + assert entries["tool.sh"].source_path == str(tool) + assert Path(str(entries["tool.sh"].source_path)).read_text() == "after" + assert entries["conf"].is_dir + assert entries["conf/nested/a.txt"].source_path == str( + tmp_path / "conf" / "nested" / "a.txt" + ) + if _KEEPS_MODES: + assert entries["tool.sh"].mode == 0o755 + + +def test_fails_at_push_when_a_local_file_or_directory_is_gone(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="is not a file"): + _base().add_local_file(str(tmp_path / "gone"), "/a")._zip_entries() + with pytest.raises(ValueError, match="is not a directory"): + _base().add_local_dir(str(tmp_path / "gone"), "/a")._zip_entries() + + +def test_adds_everything_in_a_local_directory_with_no_ignore_rules( + tmp_path: Path, +) -> None: + (tmp_path / "src" / ".git").mkdir(parents=True) + (tmp_path / "src" / ".git" / "HEAD").write_text("ref") + (tmp_path / "src" / ".env").write_text("SECRET=1") + entries = _base().add_local_dir(str(tmp_path / "src"), "/app")._zip_entries() + assert [entry.path for entry in entries] == [ + "Dockerfile", + "src", + "src/.env", + "src/.git", + "src/.git/HEAD", + ] + + +def test_resolves_a_local_path_when_added( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + (tmp_path / "a.txt").write_text("a") + monkeypatch.chdir(tmp_path) + builder = _base().add_local_file("a.txt", "/a.txt") + monkeypatch.chdir(os.path.dirname(tmp_path)) + assert builder._zip_entries()[1].source_path == str(tmp_path / "a.txt") diff --git a/tests/sandbox/test_images_ignore.py b/tests/sandbox/test_images_ignore.py new file mode 100644 index 0000000..63e7104 --- /dev/null +++ b/tests/sandbox/test_images_ignore.py @@ -0,0 +1,48 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +from baseten.sandbox import ImageIgnoreFileOptions, default_image_ignore_file +from baseten.sandbox._images_ignore import read_ignore_file + + +@pytest.mark.parametrize( + ("rel_path", "ignored"), + [ + (".git", True), + (".git/config", True), + ("app/node_modules", True), + ("app/node_modules/x/y.js", True), + ("a/b/__pycache__", True), + ("dist", True), + ("src/dist", True), + (".env", True), + ("app/.env", True), + (".env.local", True), + (".envrc", True), + (".blaxel", True), + (".env.build", True), + ("app/.env.build", True), + # .env* matches at the root only. + ("app/.env.local", False), + ("distribution", False), + ("main.py", False), + ("src/app.go", False), + ("Dockerfile", False), + ("git", False), + ], +) +def test_default_image_ignore_file(rel_path: str, ignored: bool) -> None: + options = ImageIgnoreFileOptions(rel_path=rel_path, is_directory=False) + assert default_image_ignore_file(options) is ignored + + +def test_reads_a_dockerignore_as_utf8_replacing_invalid_bytes(tmp_path: Path) -> None: + assert read_ignore_file(str(tmp_path)) is None + (tmp_path / ".dockerignore").write_bytes(b"node_modules\n\xff\n") + ignore_file = read_ignore_file(str(tmp_path)) + assert ignore_file is not None + assert ignore_file.path == str(tmp_path / ".dockerignore") + assert ignore_file.contents == "node_modules\n�\n" diff --git a/tests/sandbox/test_info.py b/tests/sandbox/test_info.py new file mode 100644 index 0000000..678fd1d --- /dev/null +++ b/tests/sandbox/test_info.py @@ -0,0 +1,189 @@ +from __future__ import annotations + +from datetime import UTC, datetime, timedelta +from typing import Any + +import baseten.client.managementapi +from baseten.sandbox import ( + SandboxEnvValue, + SandboxExpirationPolicyDate, + SandboxExpirationPolicyIdle, + SandboxExpirationPolicyMaxAge, + SandboxExpirationPolicyUnknown, + SandboxLifecycle, + SandboxNetwork, + SandboxNetworkProxy, + SandboxNetworkProxyRoute, + SandboxPort, +) +from baseten.sandbox._info import ( + envs_to_api, + lifecycle_to_api, + network_to_api, + ports_to_api, + sandbox_info_from_api, +) + +_FULL_RECORD: dict[str, Any] = { + "name": "sbx", + "url": "https://sbx.example.dev/", + "status": "DEPLOYED", + "image": "baseten/base-image:latest", + "memory": 4096, + "region": "us-was-1", + "envs": [ + {"name": "PLAIN", "value": "shown", "secret": False}, + {"name": "SECRET", "value": "****", "secret": True}, + {"value": "nameless"}, + ], + "labels": {"team": "infra"}, + "external_id": "ext-1", + "lifecycle": { + "expiration_policies": [ + {"type": "TTL_IDLE", "action": "DELETE", "value": "1h30m"}, + {"type": "TTL_MAX_AGE", "action": "DELETE", "value": "7d"}, + {"type": "DATE", "action": "DELETE", "value": "2026-10-07T12:00:00Z"}, + {"type": "ON_WEEKEND", "action": "ARCHIVE", "value": "sat"}, + ], + "terminated_retention": "24h", + }, + "ports": [{"target": 8080, "name": "http", "protocol": "HTTP"}], + "network": { + "subnet": "default", + "proxy": { + "allowed_domains": ["pypi.org"], + "routing": [{"destinations": ["*"], "headers": {"X-A": "b"}}], + }, + }, + "created_at": "2026-10-06T12:00:00Z", + "updated_at": "2026-10-06T13:00:00Z", + "created_by": "creator", + "updated_by": "updater", + "last_used_at": "2026-10-06T14:00:00Z", + "expires_in": 3600, +} + + +def test_maps_sandbox_record() -> None: + info = sandbox_info_from_api( + baseten.client.managementapi.Sandbox.model_validate(_FULL_RECORD) + ) + assert info.name == "sbx" + assert info.url == "https://sbx.example.dev/" + assert info.status == "DEPLOYED" + assert (info.image, info.memory, info.region) == ( + "baseten/base-image:latest", + 4096, + "us-was-1", + ) + # An env without a name is left out. + assert info.envs == { + "PLAIN": SandboxEnvValue(value="shown", secret=False), + "SECRET": SandboxEnvValue(value="****", secret=True), + } + assert info.labels == {"team": "infra"} + assert info.external_id == "ext-1" + assert info.lifecycle == SandboxLifecycle( + expiration_policies=[ + SandboxExpirationPolicyIdle( + after=timedelta(hours=1, minutes=30), action="DELETE" + ), + SandboxExpirationPolicyMaxAge(after=timedelta(days=7), action="DELETE"), + SandboxExpirationPolicyDate( + at=datetime(2026, 10, 7, 12, tzinfo=UTC), action="DELETE" + ), + SandboxExpirationPolicyUnknown( + type="ON_WEEKEND", + raw={"type": "ON_WEEKEND", "action": "ARCHIVE", "value": "sat"}, + ), + ], + terminated_retention=timedelta(hours=24), + ) + assert info.ports == [SandboxPort(target=8080, name="http", protocol="HTTP")] + assert info.network == SandboxNetwork( + subnet="default", + proxy=SandboxNetworkProxy( + allowed_domains=["pypi.org"], + routing=[ + SandboxNetworkProxyRoute(destinations=["*"], headers={"X-A": "b"}) + ], + ), + ) + assert info.created_at == datetime(2026, 10, 6, 12, tzinfo=UTC) + assert info.updated_at == datetime(2026, 10, 6, 13, tzinfo=UTC) + assert (info.created_by, info.updated_by) == ("creator", "updater") + assert info.last_used_at == datetime(2026, 10, 6, 14, tzinfo=UTC) + assert info.expires_in == timedelta(hours=1) + + +def test_maps_sandbox_record_without_optional_fields() -> None: + info = sandbox_info_from_api( + baseten.client.managementapi.Sandbox.model_validate( + { + "name": "sbx", + "url": "https://sbx.example.dev/", + "status": "SOMETHING_NEW", + "created_at": "2026-10-06T12:00:00Z", + } + ) + ) + assert info.status == "SOMETHING_NEW" + assert info.envs == {} + assert info.labels == {} + assert info.lifecycle is None + assert info.ports == [] + assert info.network is None + assert info.expires_in is None + + +def test_lifecycle_to_api() -> None: + lifecycle = lifecycle_to_api( + SandboxLifecycle( + expiration_policies=[ + SandboxExpirationPolicyIdle(after=timedelta(minutes=10)), + SandboxExpirationPolicyMaxAge( + after=timedelta(days=1), action="ARCHIVE" + ), + SandboxExpirationPolicyDate(at=datetime(2026, 10, 7, tzinfo=UTC)), + SandboxExpirationPolicyUnknown( + type="ON_WEEKEND", raw={"type": "ON_WEEKEND", "value": "sat"} + ), + ], + terminated_retention=timedelta(hours=1), + ) + ) + assert _dump(lifecycle) == { + "expiration_policies": [ + {"type": "TTL_IDLE", "action": "DELETE", "value": "600000ms"}, + {"type": "TTL_MAX_AGE", "action": "ARCHIVE", "value": "86400000ms"}, + {"type": "DATE", "action": "DELETE", "value": "2026-10-07T00:00:00Z"}, + # Sent back exactly as the server sent it. + {"type": "ON_WEEKEND", "value": "sat"}, + ], + "terminated_retention": "3600000ms", + } + + +def test_inputs_send_only_set_fields() -> None: + assert _dump(lifecycle_to_api(SandboxLifecycle())) == {} + assert [_dump(env) for env in envs_to_api({"A": SandboxEnvValue(value="1")})] == [ + {"name": "A", "value": "1"} + ] + assert [_dump(port) for port in ports_to_api([SandboxPort(target=80)])] == [ + {"target": 80} + ] + assert _dump( + network_to_api( + SandboxNetwork( + proxy=SandboxNetworkProxy( + bypass=("example.com",), + routing=[SandboxNetworkProxyRoute(secrets={"k": "v"})], + ) + ) + ) + ) == {"proxy": {"bypass": ["example.com"], "routing": [{"secrets": {"k": "v"}}]}} + + +def _dump(model: Any) -> Any: + # As the generated client serializes a request body. + return model.model_dump(mode="json", exclude_unset=True, by_alias=True) diff --git a/tests/sandbox/test_parity.py b/tests/sandbox/test_parity.py new file mode 100644 index 0000000..1357c05 --- /dev/null +++ b/tests/sandbox/test_parity.py @@ -0,0 +1,337 @@ +from __future__ import annotations + +import dataclasses +import inspect +from collections.abc import Callable +from typing import Any + +import pytest +from pydantic import BaseModel + +import baseten.client.managementapi +import baseten.client.sandboxapi +from baseten.sandbox import ( + AsyncImageClient, + AsyncSandbox, + AsyncSandboxClient, + AsyncSandboxCreateResult, + AsyncSandboxFileSystem, + AsyncSandboxProcess, + AsyncSandboxStream, + ImageCleanupResult, + ImageClient, + ImageInfo, + ImageLibraryInfo, + ImageLogLine, + ImageTagInfo, + Sandbox, + SandboxClient, + SandboxCreateResult, + SandboxEnvValue, + SandboxExpirationPolicyDate, + SandboxExpirationPolicyIdle, + SandboxExpirationPolicyMaxAge, + SandboxFileSystem, + SandboxFileSystemDirectory, + SandboxFileSystemFileInfo, + SandboxFileSystemFindMatch, + SandboxFileSystemFindResult, + SandboxFileSystemGrepMatch, + SandboxFileSystemGrepResult, + SandboxFileSystemSubdirectoryInfo, + SandboxInfo, + SandboxLifecycle, + SandboxNetwork, + SandboxNetworkProxy, + SandboxNetworkProxyRoute, + SandboxPort, + SandboxProcess, + SandboxProcessInfo, + SandboxProcessLogs, + SandboxStream, +) + +# Methods of an async class that are plain functions, since they make no call. +_SYNC_IN_ASYNC = {(AsyncSandboxClient, "sandbox_from_url")} + + +@pytest.mark.parametrize( + ("sync_cls", "async_cls"), + [ + (SandboxClient, AsyncSandboxClient), + (Sandbox, AsyncSandbox), + (SandboxCreateResult, AsyncSandboxCreateResult), + (SandboxProcess, AsyncSandboxProcess), + (SandboxFileSystem, AsyncSandboxFileSystem), + (ImageClient, AsyncImageClient), + (SandboxStream, AsyncSandboxStream), + ], +) +def test_sync_and_async_match(sync_cls: type, async_cls: type) -> None: + sync_members = _public_members(sync_cls) + async_members = _public_members(async_cls) + assert sync_members.keys() == async_members.keys() + for name, sync_member in sync_members.items(): + async_member = async_members[name] + if isinstance(sync_member, property): + assert isinstance(async_member, property), name + continue + assert _parameters(sync_member) == _parameters(async_member), name + if inspect.isgeneratorfunction(sync_member): + assert inspect.isasyncgenfunction(async_member), name + elif (async_cls, name) in _SYNC_IN_ASYNC: + assert not inspect.iscoroutinefunction(async_member), name + else: + assert inspect.iscoroutinefunction(async_member), name + + +@dataclasses.dataclass(kw_only=True) +class _Parity: + """How one of our types or calls corresponds to a generated model.""" + + ours: set[str] + generated: type[BaseModel] + renamed: dict[str, str] = dataclasses.field(default_factory=dict) + """Generated field name to ours.""" + ours_only: dict[str, str] = dataclasses.field(default_factory=dict) + """Our field to why the generated model lacks it.""" + generated_only: dict[str, str] = dataclasses.field(default_factory=dict) + """Generated field to why we leave it out.""" + + +def _dataclass_fields(cls: Any) -> set[str]: + return {field.name for field in dataclasses.fields(cls)} + + +def _kwargs(method: Callable[..., Any]) -> set[str]: + return { + name + for name, parameter in inspect.signature(method).parameters.items() + if parameter.kind is inspect.Parameter.KEYWORD_ONLY + } + + +_PROCESS_RENAMES = { + "workingDir": "working_dir", + "keepAlive": "keep_alive", + "maxRestarts": "max_restarts", + "restartOnFailure": "restart_on_failure", +} +_PROCESS_REQUEST_RENAMES = { + **_PROCESS_RENAMES, + "waitForCompletion": "wait_for_completion", + "waitForPorts": "wait_for_ports", +} + +_POLICY_RENAMES = {"value": "after"} +_POLICY_TYPE = {"type": "set by which policy class is used"} + +# A field the server adds fails here until it is mapped or listed. +_PARITY = { + "SandboxInfo": _Parity( + ours=_dataclass_fields(SandboxInfo), + generated=baseten.client.managementapi.Sandbox, + ), + "SandboxEnvValue": _Parity( + ours=_dataclass_fields(SandboxEnvValue), + generated=baseten.client.managementapi.SandboxEnv, + generated_only={"name": "the key envs are mapped by"}, + ), + "SandboxLifecycle": _Parity( + ours=_dataclass_fields(SandboxLifecycle), + generated=baseten.client.managementapi.SandboxLifecycle, + ), + "SandboxExpirationPolicyIdle": _Parity( + ours=_dataclass_fields(SandboxExpirationPolicyIdle), + generated=baseten.client.managementapi.SandboxTTLIdleExpirationPolicy, + renamed=_POLICY_RENAMES, + generated_only=_POLICY_TYPE, + ), + "SandboxExpirationPolicyMaxAge": _Parity( + ours=_dataclass_fields(SandboxExpirationPolicyMaxAge), + generated=baseten.client.managementapi.SandboxTTLMaxAgeExpirationPolicy, + renamed=_POLICY_RENAMES, + generated_only=_POLICY_TYPE, + ), + "SandboxExpirationPolicyDate": _Parity( + ours=_dataclass_fields(SandboxExpirationPolicyDate), + generated=baseten.client.managementapi.SandboxDateExpirationPolicy, + renamed={"value": "at"}, + generated_only=_POLICY_TYPE, + ), + "SandboxPort": _Parity( + ours=_dataclass_fields(SandboxPort), + generated=baseten.client.managementapi.SandboxPort, + ), + "SandboxNetwork": _Parity( + ours=_dataclass_fields(SandboxNetwork), + generated=baseten.client.managementapi.SandboxNetwork, + ), + "SandboxNetworkProxy": _Parity( + ours=_dataclass_fields(SandboxNetworkProxy), + generated=baseten.client.managementapi.SandboxProxyConfig, + ), + "SandboxNetworkProxyRoute": _Parity( + ours=_dataclass_fields(SandboxNetworkProxyRoute), + generated=baseten.client.managementapi.SandboxProxyTarget, + ), + "SandboxClient.create": _Parity( + ours=_kwargs(SandboxClient.create), + generated=baseten.client.managementapi.CreateSandboxRequest, + ), + "SandboxClient.update": _Parity( + ours=_kwargs(SandboxClient.update) - {"name"}, + generated=baseten.client.managementapi.UpdateSandboxRequest, + ), + "SandboxProcessInfo": _Parity( + ours=_dataclass_fields(SandboxProcessInfo), + generated=baseten.client.sandboxapi.ProcessResponse, + renamed={ + **_PROCESS_RENAMES, + "exitCode": "exit_code", + "startedAt": "started_at", + "completedAt": "completed_at", + "restartCount": "restart_count", + }, + ), + "SandboxProcessLogs": _Parity( + ours=_dataclass_fields(SandboxProcessLogs), + generated=baseten.client.sandboxapi.ProcessLogs, + ), + "SandboxProcess.exec": _Parity( + ours=_kwargs(SandboxProcess.exec), + generated=baseten.client.sandboxapi.ProcessRequest, + renamed=_PROCESS_REQUEST_RENAMES, + ), + "SandboxProcess.exec_stream": _Parity( + ours=_kwargs(SandboxProcess.exec_stream), + generated=baseten.client.sandboxapi.ProcessRequest, + renamed=_PROCESS_REQUEST_RENAMES, + generated_only={"waitForCompletion": "always sent true, to stream"}, + ), + "SandboxFileSystemDirectory": _Parity( + ours=_dataclass_fields(SandboxFileSystemDirectory), + generated=baseten.client.sandboxapi.Directory, + ), + "SandboxFileSystemFileInfo": _Parity( + ours=_dataclass_fields(SandboxFileSystemFileInfo), + generated=baseten.client.sandboxapi.File, + renamed={"size": "size_bytes", "lastModified": "last_modified"}, + ), + "SandboxFileSystemSubdirectoryInfo": _Parity( + ours=_dataclass_fields(SandboxFileSystemSubdirectoryInfo), + generated=baseten.client.sandboxapi.Subdirectory, + ), + "SandboxFileSystemFindResult": _Parity( + ours=_dataclass_fields(SandboxFileSystemFindResult), + generated=baseten.client.sandboxapi.FindResponse, + ), + "SandboxFileSystemFindMatch": _Parity( + ours=_dataclass_fields(SandboxFileSystemFindMatch), + generated=baseten.client.sandboxapi.FindMatch, + ), + "SandboxFileSystemGrepResult": _Parity( + ours=_dataclass_fields(SandboxFileSystemGrepResult), + generated=baseten.client.sandboxapi.ContentSearchResponse, + generated_only={"query": "the caller's own query, echoed"}, + ), + "SandboxFileSystemGrepMatch": _Parity( + ours=_dataclass_fields(SandboxFileSystemGrepMatch), + generated=baseten.client.sandboxapi.ContentSearchMatch, + ), + "ImageInfo": _Parity( + ours=_dataclass_fields(ImageInfo), + generated=baseten.client.managementapi.SandboxImage, + renamed={"size": "size_bytes"}, + generated_only={"tags": "always empty; list_tags lists them"}, + ), + "ImageInfo/Summary": _Parity( + ours=_dataclass_fields(ImageInfo), + generated=baseten.client.managementapi.SandboxImageSummary, + renamed={"size": "size_bytes"}, + ), + "ImageTagInfo": _Parity( + ours=_dataclass_fields(ImageTagInfo), + generated=baseten.client.managementapi.SandboxImageTag, + renamed={"size": "size_bytes"}, + ), + "ImageLogLine": _Parity( + ours=_dataclass_fields(ImageLogLine), + generated=baseten.client.managementapi.SandboxImageBuildLog, + renamed={"message": "text"}, + ), + "ImageCleanupResult": _Parity( + ours=_dataclass_fields(ImageCleanupResult), + generated=baseten.client.managementapi.CleanupSandboxImagesResponse, + ), + "ImageLibraryInfo": _Parity( + ours=_dataclass_fields(ImageLibraryInfo), + generated=baseten.client.managementapi.SandboxLibraryImage, + ours_only={ + "creation_extra_args": "flattened from creation_options", + "creation_volumes": "flattened from creation_options", + }, + generated_only={ + "hidden": "always false in listings", + "coming_soon": "always false in listings", + "creation_options": "flattened into creation_* fields", + }, + ), + "ImageClient.list": _Parity( + ours=_kwargs(ImageClient.list), + generated=baseten.client.managementapi.ListImagesParams, + renamed={"limit": "page_size", "q": "name_prefix"}, + generated_only={ + "team_id": "the client's team", + "cursor": "paged internally", + }, + ), + "SandboxFileSystem.find": _Parity( + ours=_kwargs(SandboxFileSystem.find) - {"path"}, + generated=baseten.client.sandboxapi.GetFilesystemFindParams, + renamed={ + "maxResults": "max_results", + "excludeDirs": "exclude_dirs", + "excludeHidden": "exclude_hidden", + }, + ), + "SandboxFileSystem.grep": _Parity( + ours=_kwargs(SandboxFileSystem.grep) - {"path"}, + generated=baseten.client.sandboxapi.GetFilesystemContentSearchParams, + renamed={ + "caseSensitive": "case_sensitive", + "maxResults": "max_results", + "filePattern": "file_pattern", + "excludeDirs": "exclude_dirs", + "contextLines": "context_lines", + }, + ), +} + + +@pytest.mark.parametrize("name", _PARITY) +def test_fields_match_generated(name: str) -> None: + parity = _PARITY[name] + generated = { + parity.renamed.get(field, field) + for field in parity.generated.model_fields + if field not in parity.generated_only + } + assert generated == parity.ours - parity.ours_only.keys() + + +def _public_members(cls: type) -> dict[str, Any]: + return { + name: member + for name, member in inspect.getmembers(cls) + if not name.startswith("_") + } + + +def _parameters(method: Callable[..., Any]) -> list[tuple[str, Any, Any]]: + # Annotations differ between the two forms, so only names, kinds, and + # defaults are compared. + return [ + (parameter.name, parameter.kind, parameter.default) + for parameter in inspect.signature(method).parameters.values() + ] diff --git a/tests/sandbox/test_process.py b/tests/sandbox/test_process.py new file mode 100644 index 0000000..d94b23c --- /dev/null +++ b/tests/sandbox/test_process.py @@ -0,0 +1,485 @@ +from __future__ import annotations + +import asyncio +import json +from datetime import UTC, datetime, timedelta +from typing import Any + +import httpx +import pytest + +import baseten.sandbox._retry +from baseten.sandbox import ( + AsyncSandbox, + AsyncSandboxClient, + Sandbox, + SandboxAPIError, + SandboxClient, + SandboxError, + SandboxGatewayError, + SandboxProcessExecEventExit, + SandboxProcessExecEventOutput, + SandboxProcessInfo, + SandboxProcessLogLine, + SandboxProcessLogs, + SandboxProcessWaitTimeoutError, +) + +_SANDBOX_URL = "https://sbx.test" + +_RESET = httpx.RemoteProtocolError("Server disconnected without sending a response.") + +_Reply = httpx.Response | Exception | dict[str, Any] | list[Any] + + +@pytest.fixture(autouse=True) +def short_backoff(monkeypatch: pytest.MonkeyPatch) -> None: + # Retries take milliseconds instead of seconds. + monkeypatch.setattr(baseten.sandbox._retry, "_backoff_base", 0.001) + monkeypatch.setattr(baseten.sandbox._retry, "_backoff_max", 0.002) + + +class _Api: + """Fake sandbox replying from per-route queues, recording requests. + + Each route replies in order and repeats its last reply. + """ + + def __init__(self) -> None: + self.requests: list[httpx.Request] = [] + self.routes: dict[tuple[str, str], list[_Reply]] = {} + + def route(self, method: str, path: str, *replies: _Reply) -> None: + self.routes[(method, path)] = list(replies) + + def handle(self, request: httpx.Request) -> httpx.Response: + if request.url.path == "/v1/token": + return httpx.Response( + 200, + json={ + "token": "token", + "expires_at": (datetime.now(UTC) + timedelta(hours=2)).isoformat(), + }, + ) + self.requests.append(request) + replies = self.routes[(request.method, request.url.path)] + reply = replies.pop(0) if len(replies) > 1 else replies[0] + if isinstance(reply, Exception): + raise reply + if isinstance(reply, httpx.Response): + # A response is single use, so each request gets a copy. + return httpx.Response( + reply.status_code, headers=reply.headers, content=reply.content + ) + return httpx.Response(200, json=reply) + + def sandbox(self) -> Sandbox: + client = SandboxClient( + api_key="key", + base_url_override="https://api.test", + http_client_override=httpx.Client( + transport=httpx.MockTransport(self.handle) + ), + ) + return client.sandbox_from_url(url=_SANDBOX_URL) + + def async_sandbox(self) -> AsyncSandbox: + async def handle(request: httpx.Request) -> httpx.Response: + return self.handle(request) + + client = AsyncSandboxClient( + api_key="key", + base_url_override="https://api.test", + http_client_override=httpx.AsyncClient( + transport=httpx.MockTransport(handle) + ), + ) + return client.sandbox_from_url(url=_SANDBOX_URL) + + +def _process(**fields: Any) -> dict[str, Any]: + return { + "pid": "42", + "name": "job", + "command": "echo hi", + "status": "completed", + "exitCode": 0, + "stdout": "hi\n", + "stderr": "", + "logs": "hi\n", + "workingDir": "/app", + "startedAt": "Wed, 01 Jan 2023 12:00:00 GMT", + "completedAt": "Wed, 01 Jan 2023 12:00:01 GMT", + **fields, + } + + +def _error(status: int, message: str = "nope") -> httpx.Response: + return httpx.Response(status, json={"error": message}) + + +def _ndjson(*records: dict[str, Any]) -> httpx.Response: + content = "".join(json.dumps(record) + "\n" for record in records) + return httpx.Response( + 200, headers={"Content-Type": "application/x-ndjson"}, content=content + ) + + +def _text(content: str) -> httpx.Response: + return httpx.Response(200, headers={"Content-Type": "text/plain"}, content=content) + + +def _body(request: httpx.Request) -> Any: + return json.loads(request.content) + + +class TestSandboxProcess: + def test_execs_with_every_field_mapped(self) -> None: + api = _Api() + api.route("POST", "/process", _process()) + info = api.sandbox().process.exec( + command="echo hi", + name="job", + working_dir="/app", + env={"A": "1"}, + timeout=timedelta(milliseconds=1500), + wait_for_completion=True, + keep_alive=True, + max_restarts=-1, + restart_on_failure=True, + stdin=True, + wait_for_ports=[3000], + ) + assert _body(api.requests[0]) == { + "command": "echo hi", + "name": "job", + "workingDir": "/app", + "env": {"A": "1"}, + # Rounded up to whole seconds. + "timeout": 2, + "waitForCompletion": True, + "keepAlive": True, + "maxRestarts": -1, + "restartOnFailure": True, + "stdin": True, + "waitForPorts": [3000], + } + assert info == SandboxProcessInfo( + pid="42", + name="job", + command="echo hi", + status="completed", + exit_code=0, + stdout="hi\n", + stderr="", + logs="hi\n", + working_dir="/app", + started_at=datetime(2023, 1, 1, 12, 0, 0, tzinfo=UTC), + completed_at=datetime(2023, 1, 1, 12, 0, 1, tzinfo=UTC), + keep_alive=None, + max_restarts=None, + restart_count=None, + restart_on_failure=None, + stdin=None, + ) + + def test_execs_with_only_set_fields(self) -> None: + api = _Api() + api.route("POST", "/process", _process(status="running", completedAt="")) + info = api.sandbox().process.exec(command="sleep 1") + assert _body(api.requests[0]) == {"command": "sleep 1"} + assert info.completed_at is None + + def test_never_retries_exec(self) -> None: + api = _Api() + api.route("POST", "/process", _error(502), _process()) + with pytest.raises(SandboxGatewayError): + api.sandbox().process.exec(command="echo hi") + assert len(api.requests) == 1 + + def test_converts_exec_error_response(self) -> None: + api = _Api() + api.route("POST", "/process", _error(422, "bad command")) + with pytest.raises(SandboxAPIError, match="bad command") as raised: + api.sandbox().process.exec(command="") + assert raised.value.status == 422 + + def test_exec_streams_output_then_exit(self) -> None: + api = _Api() + api.route( + "POST", + "/process", + _ndjson( + {"type": "stdout", "data": "Name? "}, + {"type": "stderr", "data": "warn\n"}, + {"type": "keepalive"}, + {"type": "unknown", "data": "skipped"}, + {"type": "result", "data": json.dumps(_process(exitCode=3))}, + {"type": "stdout", "data": "never read"}, + ), + ) + with api.sandbox().process.exec_stream( + command="ask", timeout=timedelta(0) + ) as stream: + events = list(stream) + request = api.requests[0] + assert request.headers["Accept"] == "application/x-ndjson" + assert _body(request) == { + "command": "ask", + "timeout": 0, + "waitForCompletion": True, + } + assert events[:2] == [ + SandboxProcessExecEventOutput(stream="stdout", text="Name? "), + SandboxProcessExecEventOutput(stream="stderr", text="warn\n"), + ] + assert len(events) == 3 + assert isinstance(events[2], SandboxProcessExecEventExit) + assert events[2].process.exit_code == 3 + + def test_exec_stream_fails_without_exit(self) -> None: + api = _Api() + api.route("POST", "/process", _ndjson({"type": "stdout", "data": "hi"})) + stream = api.sandbox().process.exec_stream(command="echo hi") + assert next(stream) == SandboxProcessExecEventOutput(stream="stdout", text="hi") + with pytest.raises(SandboxError, match="ended before reporting"): + next(stream) + + def test_exec_stream_fails_on_error_record(self) -> None: + api = _Api() + api.route( + "POST", + "/process", + _ndjson( + {"type": "stdout", "data": "hi"}, + {"type": "error", "data": "process lost"}, + {"type": "stdout", "data": "never read"}, + ), + ) + stream = api.sandbox().process.exec_stream(command="echo hi") + assert next(stream) == SandboxProcessExecEventOutput(stream="stdout", text="hi") + with pytest.raises(SandboxError, match="process stream failed: process lost"): + next(stream) + + def test_exec_stream_raises_error_response_at_call(self) -> None: + api = _Api() + api.route("POST", "/process", _error(400, "bad request")) + with pytest.raises(SandboxAPIError, match="bad request"): + api.sandbox().process.exec_stream(command="echo hi") + + def test_gets_process_retrying_dropped_connection(self) -> None: + api = _Api() + api.route("GET", "/process/job", _RESET, _process()) + info = api.sandbox().process.get(identifier="job") + assert info.pid == "42" + assert len(api.requests) == 2 + + def test_lists_processes(self) -> None: + api = _Api() + api.route("GET", "/process", [_process(pid="1"), _process(pid="2")]) + processes = api.sandbox().process.list() + assert [process.pid for process in processes] == ["1", "2"] + + def test_stops_and_kills(self) -> None: + api = _Api() + api.route("DELETE", "/process/job", {"message": "ok"}) + api.route("DELETE", "/process/job/kill", {"message": "ok"}) + sandbox = api.sandbox() + sandbox.process.stop(identifier="job") + sandbox.process.kill(identifier="job") + assert [request.url.path for request in api.requests] == [ + "/process/job", + "/process/job/kill", + ] + + def test_never_retries_stop(self) -> None: + api = _Api() + api.route("DELETE", "/process/job", _RESET, {"message": "ok"}) + with pytest.raises(httpx.RemoteProtocolError): + api.sandbox().process.stop(identifier="job") + assert len(api.requests) == 1 + + def test_gets_logs(self) -> None: + api = _Api() + api.route( + "GET", + "/process/job/logs", + {"stdout": "a\n", "stderr": "b\n", "logs": "a\nb\n"}, + ) + logs = api.sandbox().process.logs(identifier="job") + assert logs == SandboxProcessLogs(stdout="a\n", stderr="b\n", logs="a\nb\n") + + def test_streams_logs_by_line(self) -> None: + api = _Api() + api.route( + "GET", + "/process/job/logs/stream", + _text("stdout:one\n[keepalive]\nstderr:two\nplain\nstdout:"), + ) + with api.sandbox().process.stream_logs(identifier="job") as stream: + lines = list(stream) + assert lines == [ + SandboxProcessLogLine(stream="stdout", text="one"), + SandboxProcessLogLine(stream="stderr", text="two"), + SandboxProcessLogLine(stream=None, text="plain"), + SandboxProcessLogLine(stream="stdout", text=""), + ] + + def test_waits_until_finished_riding_out_transient_errors(self) -> None: + api = _Api() + api.route( + "GET", + "/process/job", + _process(status="running", completedAt=""), + _error(503), + _RESET, + _error(500), + _process(status="failed", exitCode=1), + ) + info = api.sandbox().process.wait( + identifier="job", poll_interval=timedelta(milliseconds=1) + ) + # A failed process is returned, not raised. + assert (info.status, info.exit_code) == ("failed", 1) + assert len(api.requests) == 5 + + def test_wait_raises_error_it_cannot_ride_out(self) -> None: + api = _Api() + api.route("GET", "/process/job", _error(404, "process not found")) + with pytest.raises(SandboxAPIError, match="process not found"): + api.sandbox().process.wait(identifier="job") + + def test_wait_raises_on_unknown_status(self) -> None: + api = _Api() + api.route("GET", "/process/job", _process(status="paused")) + with pytest.raises(SandboxError, match="process job has unknown status paused"): + api.sandbox().process.wait(identifier="job") + + def test_wait_times_out_with_last_error_as_cause(self) -> None: + api = _Api() + api.route("GET", "/process/job", _process(status="running"), _error(503)) + with pytest.raises(SandboxProcessWaitTimeoutError) as raised: + api.sandbox().process.wait( + identifier="job", + timeout=timedelta(milliseconds=20), + poll_interval=timedelta(milliseconds=1), + ) + assert raised.value.identifier == "job" + assert isinstance(raised.value.__cause__, SandboxGatewayError) + + def test_writes_stdin_as_bytes(self) -> None: + api = _Api() + api.route("POST", "/process/job/stdin", {"message": "ok"}) + sandbox = api.sandbox() + sandbox.process.write_stdin(identifier="job", data="é\n") + sandbox.process.write_stdin(identifier="job", data=b"\x00\xff") + assert [request.content for request in api.requests] == [ + "é\n".encode(), + b"\x00\xff", + ] + assert api.requests[0].headers["Content-Type"] == "application/octet-stream" + + def test_never_retries_write_stdin(self) -> None: + api = _Api() + api.route("POST", "/process/job/stdin", _error(502), {"message": "ok"}) + with pytest.raises(SandboxGatewayError): + api.sandbox().process.write_stdin(identifier="job", data="x") + assert len(api.requests) == 1 + + def test_closes_stdin_retrying_gateway_error(self) -> None: + api = _Api() + api.route("DELETE", "/process/job/stdin", _error(502), {"message": "ok"}) + api.sandbox().process.close_stdin(identifier="job") + assert len(api.requests) == 2 + + +class TestAsyncSandboxProcess: + @pytest.mark.asyncio + async def test_execs_and_gets(self) -> None: + api = _Api() + api.route("POST", "/process", _process()) + api.route("GET", "/process/job", _RESET, _process()) + sandbox = api.async_sandbox() + info = await sandbox.process.exec(command="echo hi", name="job") + assert _body(api.requests[0]) == {"command": "echo hi", "name": "job"} + assert info.stdout == "hi\n" + assert (await sandbox.process.get(identifier="job")).pid == "42" + + @pytest.mark.asyncio + async def test_exec_streams_output_then_exit(self) -> None: + api = _Api() + api.route( + "POST", + "/process", + _ndjson( + {"type": "stdout", "data": "hi\n"}, + {"type": "result", "data": json.dumps(_process())}, + ), + ) + async with await api.async_sandbox().process.exec_stream( + command="echo hi" + ) as stream: + events = [event async for event in stream] + assert events[0] == SandboxProcessExecEventOutput(stream="stdout", text="hi\n") + assert isinstance(events[1], SandboxProcessExecEventExit) + + @pytest.mark.asyncio + async def test_streams_logs_by_line(self) -> None: + api = _Api() + api.route("GET", "/process/job/logs/stream", _text("stdout:one\n[keepalive]\n")) + stream = await api.async_sandbox().process.stream_logs(identifier="job") + assert [line async for line in stream] == [ + SandboxProcessLogLine(stream="stdout", text="one") + ] + + @pytest.mark.asyncio + async def test_waits_until_finished(self) -> None: + api = _Api() + api.route( + "GET", + "/process/job", + _process(status="running"), + _error(429), + _process(status="stopped"), + ) + info = await api.async_sandbox().process.wait( + identifier="job", poll_interval=timedelta(milliseconds=1) + ) + assert info.status == "stopped" + + @pytest.mark.asyncio + async def test_wait_times_out_with_last_error_as_cause(self) -> None: + api = _Api() + api.route("GET", "/process/job", _process(status="running"), _RESET) + with pytest.raises(SandboxProcessWaitTimeoutError) as raised: + await api.async_sandbox().process.wait( + identifier="job", + timeout=timedelta(milliseconds=20), + poll_interval=timedelta(milliseconds=1), + ) + assert raised.value.identifier == "job" + assert isinstance(raised.value.__cause__, httpx.RemoteProtocolError) + + @pytest.mark.asyncio + async def test_wait_leaves_callers_own_timeout_alone(self) -> None: + api = _Api() + api.route("GET", "/process/job", _process(status="running")) + sandbox = api.async_sandbox() + with pytest.raises(TimeoutError) as raised: + async with asyncio.timeout(0.02): + await sandbox.process.wait( + identifier="job", + timeout=None, + poll_interval=timedelta(milliseconds=1), + ) + assert not isinstance(raised.value, SandboxProcessWaitTimeoutError) + + @pytest.mark.asyncio + async def test_writes_and_closes_stdin(self) -> None: + api = _Api() + api.route("POST", "/process/job/stdin", {"message": "ok"}) + api.route("DELETE", "/process/job/stdin", {"message": "ok"}) + sandbox = api.async_sandbox() + await sandbox.process.write_stdin(identifier="job", data="x") + await sandbox.process.close_stdin(identifier="job") + assert [request.method for request in api.requests] == ["POST", "DELETE"] diff --git a/tests/sandbox/test_retry.py b/tests/sandbox/test_retry.py new file mode 100644 index 0000000..c9cfc81 --- /dev/null +++ b/tests/sandbox/test_retry.py @@ -0,0 +1,371 @@ +from __future__ import annotations + +import asyncio +import contextlib +import socket +import ssl +import struct +import sys +import threading +from collections.abc import Callable, Iterator +from pathlib import Path +from typing import Any, Literal + +import h2.config +import h2.connection +import h2.errors +import h2.events +import httpx +import pytest + +import baseten.sandbox._retry +from baseten.sandbox import SandboxAPIError, SandboxGatewayError, SandboxRetryOptions +from baseten.sandbox._retry import ( + is_transient_reset_error, + resolve_retry_options, + retry_idempotent, + retry_idempotent_async, +) + +_FIXTURES = Path(__file__).parent / "fixtures" + + +@pytest.fixture(autouse=True) +def short_backoff(monkeypatch: pytest.MonkeyPatch) -> None: + # Retries take milliseconds instead of seconds. + monkeypatch.setattr(baseten.sandbox._retry, "_backoff_base", 0.001) + monkeypatch.setattr(baseten.sandbox._retry, "_backoff_max", 0.002) + + +def test_resolve_retry_options() -> None: + resolved = resolve_retry_options( + SandboxRetryOptions(gateway_max_retries=0, upload_max_retries=7) + ) + assert resolved.read_max_retries == 5 + assert resolved.gateway_max_retries == 0 + assert resolved.upload_max_retries == 7 + + +# Each case produces its failure on a real connection, since the errors that +# mark a dropped connection differ by platform, protocol, and sync or async +# I/O. +@pytest.mark.parametrize("use_async", [False, True], ids=["sync", "async"]) +class TestIsTransientResetError: + def test_closed_before_response(self, use_async: bool) -> None: + with _http1_server(lambda conn: None) as url: + err = _request_error(url, use_async=use_async) + assert isinstance(err, httpx.RemoteProtocolError) + assert is_transient_reset_error(err), repr(err) + + def test_closed_mid_body(self, use_async: bool) -> None: + def respond_partially(conn: socket.socket) -> None: + conn.sendall(b"HTTP/1.1 200 OK\r\nContent-Length: 100\r\n\r\npartial") + + with _http1_server(respond_partially) as url: + err = _request_error(url, use_async=use_async) + assert isinstance(err, httpx.RemoteProtocolError) + assert is_transient_reset_error(err), repr(err) + + def test_connection_reset(self, use_async: bool) -> None: + def reset(conn: socket.socket) -> None: + # Closing with no linger sends a reset instead of a clean close. + linger = struct.pack("HH" if sys.platform == "win32" else "ii", 1, 0) + conn.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, linger) + + with _http1_server(reset) as url: + err = _request_error(url, use_async=use_async) + assert is_transient_reset_error(err), repr(err) + + def test_connection_refused(self, use_async: bool) -> None: + listener = socket.create_server(("127.0.0.1", 0)) + port = listener.getsockname()[1] + listener.close() + err = _request_error(f"http://127.0.0.1:{port}", use_async=use_async) + assert isinstance(err, httpx.ConnectError) + assert is_transient_reset_error(err), repr(err) + + def test_http2_stream_reset(self, use_async: bool) -> None: + with _http2_server("reset") as url: + err = _request_error(url, use_async=use_async, http2=True) + assert "StreamReset" in str(err) + assert is_transient_reset_error(err), repr(err) + + def test_http2_goaway(self, use_async: bool) -> None: + with _http2_server("goaway") as url: + err = _request_error(url, use_async=use_async, http2=True) + assert "ConnectionTerminated" in str(err) + assert is_transient_reset_error(err), repr(err) + + def test_http2_connection_dropped(self, use_async: bool) -> None: + with _http2_server("drop") as url: + err = _request_error(url, use_async=use_async, http2=True) + assert is_transient_reset_error(err), repr(err) + + def test_read_timeout_never_transient(self, use_async: bool) -> None: + # A timeout the caller configured is theirs to handle, not a dropped + # connection. + responded = threading.Event() + with _http1_server(lambda conn: responded.wait()) as url: + try: + err = _request_error( + url, use_async=use_async, timeout=httpx.Timeout(5.0, read=0.05) + ) + finally: + responded.set() + assert isinstance(err, httpx.ReadTimeout) + assert not is_transient_reset_error(err), repr(err) + + +def test_api_error_never_transient() -> None: + # A response came back, even when its body reads like a reset. + err = SandboxAPIError( + status=500, + code="StreamReset", + message="Server disconnected", + details=None, + body="ConnectionTerminated", + ) + err.__cause__ = ConnectionResetError() + assert not is_transient_reset_error(err) + + +def test_non_reset_failures_never_transient() -> None: + assert not is_transient_reset_error(ValueError("bad value")) + # A malformed response is a protocol error, but not a dropped connection. + assert not is_transient_reset_error(httpx.RemoteProtocolError("illegal header")) + # A name that does not resolve fails to connect, but will not on a retry. + dns_failure = httpx.ConnectError("no such host") + dns_failure.__cause__ = socket.gaierror(socket.EAI_NONAME, "no such host") + assert not is_transient_reset_error(dns_failure) + + +def test_backoff_delay_doubles_up_to_max_plus_jitter( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(baseten.sandbox._retry, "_backoff_base", 0.2) + monkeypatch.setattr(baseten.sandbox._retry, "_backoff_max", 2.0) + for attempt, capped in [(1, 0.2), (2, 0.4), (4, 1.6), (5, 2.0), (9, 2.0)]: + delay = baseten.sandbox._retry.backoff_delay(attempt) + assert capped <= delay <= capped + 0.2 + + +class TestRetryIdempotent: + def test_retries_transient_up_to_limit(self) -> None: + calls = 0 + + def call() -> int: + nonlocal calls + calls += 1 + raise httpx.ReadError("reset") from ConnectionResetError() + + with pytest.raises(httpx.ReadError): + retry_idempotent(call, max_retries=2, gateway_max_retries=0) + assert calls == 3 + + def test_succeeds_after_transient(self) -> None: + calls = 0 + + def call() -> int: + nonlocal calls + calls += 1 + if calls == 1: + raise httpx.ReadError("reset") from ConnectionResetError() + return 42 + + assert retry_idempotent(call, max_retries=2, gateway_max_retries=0) == 42 + + def test_gateway_has_own_budget(self) -> None: + calls = 0 + + def call() -> int: + nonlocal calls + calls += 1 + raise _gateway_error() + + with pytest.raises(SandboxGatewayError): + retry_idempotent(call, max_retries=0, gateway_max_retries=1) + assert calls == 2 + + def test_non_transient_not_retried(self) -> None: + calls = 0 + + def call() -> int: + nonlocal calls + calls += 1 + raise SandboxAPIError( + status=404, code=None, message=None, details=None, body="" + ) + + with pytest.raises(SandboxAPIError): + retry_idempotent(call, max_retries=5, gateway_max_retries=5) + assert calls == 1 + + +class TestRetryIdempotentAsync: + @pytest.mark.asyncio + async def test_retries_transient_up_to_limit(self) -> None: + calls = 0 + + async def call() -> int: + nonlocal calls + calls += 1 + raise httpx.ReadError("reset") from ConnectionResetError() + + with pytest.raises(httpx.ReadError): + await retry_idempotent_async(call, max_retries=2, gateway_max_retries=0) + assert calls == 3 + + @pytest.mark.asyncio + async def test_gateway_has_own_budget(self) -> None: + calls = 0 + + async def call() -> int: + nonlocal calls + calls += 1 + raise _gateway_error() + + with pytest.raises(SandboxGatewayError): + await retry_idempotent_async(call, max_retries=0, gateway_max_retries=1) + assert calls == 2 + + @pytest.mark.asyncio + async def test_cancellation_stops_backoff( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + # A backoff far longer than the test, so cancellation ends the wait. + monkeypatch.setattr(baseten.sandbox._retry, "_backoff_base", 3600.0) + monkeypatch.setattr(baseten.sandbox._retry, "_backoff_max", 3600.0) + failed = asyncio.Event() + + async def call() -> int: + failed.set() + raise httpx.ReadError("reset") from ConnectionResetError() + + task = asyncio.create_task( + retry_idempotent_async(call, max_retries=5, gateway_max_retries=0) + ) + await failed.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + @pytest.mark.asyncio + async def test_cancellation_not_retried(self) -> None: + calls = 0 + + async def call() -> int: + nonlocal calls + calls += 1 + raise asyncio.CancelledError() + + with pytest.raises(asyncio.CancelledError): + await retry_idempotent_async(call, max_retries=5, gateway_max_retries=5) + assert calls == 1 + + +def _gateway_error() -> SandboxGatewayError: + return SandboxGatewayError( + status=503, code=None, message=None, details=None, body="" + ) + + +def _request_error(url: str, *, use_async: bool, **client_options: Any) -> Exception: + """Send a GET that must fail and return its error.""" + if "http2" in client_options: + client_options["verify"] = ssl.create_default_context( + cafile=str(_FIXTURES / "localhost-cert.pem") + ) + if use_async: + + async def send() -> None: + async with httpx.AsyncClient(**client_options) as client: + await client.get(url) + + with pytest.raises(Exception) as info: + asyncio.run(send()) + else: + with pytest.raises(Exception) as info, httpx.Client(**client_options) as client: + client.get(url) + return info.value + + +@contextlib.contextmanager +def _http1_server(on_request: Callable[[socket.socket], object]) -> Iterator[str]: + """Serve one HTTP/1.1 connection, handing it to ``on_request`` once the + request has arrived, then closing it.""" + listener = socket.create_server(("127.0.0.1", 0)) + + def serve() -> None: + conn, _ = listener.accept() + with conn: + request = b"" + while b"\r\n\r\n" not in request: + chunk = conn.recv(65536) + if not chunk: + return + request += chunk + on_request(conn) + + thread = threading.Thread(target=serve, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{listener.getsockname()[1]}" + finally: + thread.join() + listener.close() + + +@contextlib.contextmanager +def _http2_server(action: Literal["reset", "goaway", "drop"]) -> Iterator[str]: + """Serve one HTTP/2 connection over TLS, failing its first request. + + ``reset`` resets the request's stream, ``goaway`` closes the connection + with GOAWAY, and ``drop`` closes the socket without either. + """ + context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH) + context.load_cert_chain( + _FIXTURES / "localhost-cert.pem", _FIXTURES / "localhost-key.pem" + ) + context.set_alpn_protocols(["h2"]) + listener = socket.create_server(("127.0.0.1", 0)) + + def serve() -> None: + raw, _ = listener.accept() + with context.wrap_socket(raw, server_side=True) as sock: + conn = h2.connection.H2Connection( + config=h2.config.H2Configuration(client_side=False) + ) + conn.initiate_connection() + sock.sendall(conn.data_to_send()) + # Serves until the client goes away, however that surfaces here. + with contextlib.suppress(OSError): + while data := sock.recv(65536): + for event in conn.receive_data(data): + if not isinstance(event, h2.events.RequestReceived): + continue + if action == "drop": + return + if action == "reset": + conn.reset_stream( + event.stream_id, h2.errors.ErrorCodes.INTERNAL_ERROR + ) + else: + conn.close_connection() + sock.sendall(conn.data_to_send()) + # Drained until the client closes, without + # handing it to the connection, which accepts no + # more frames. Closing with the client's frames + # unread would reset the connection, which on + # Windows can reach the client before the GOAWAY. + while sock.recv(65536): + pass + return + sock.sendall(conn.data_to_send()) + + thread = threading.Thread(target=serve, daemon=True) + thread.start() + try: + yield f"https://127.0.0.1:{listener.getsockname()[1]}" + finally: + thread.join() + listener.close() diff --git a/tests/sandbox/test_zip.py b/tests/sandbox/test_zip.py new file mode 100644 index 0000000..71aa453 --- /dev/null +++ b/tests/sandbox/test_zip.py @@ -0,0 +1,209 @@ +from __future__ import annotations + +import io +import stat +import sys +import zipfile +from pathlib import Path + +import pytest + +from baseten.sandbox import ImageIgnoreFileOptions +from baseten.sandbox._zip import ( + ZipEntry, + check_directory_dockerfile, + check_zip_dockerfile, + files_zip_entries, + run_walk, + run_walk_async, + walk_directory, + write_zip, +) + +# Windows reads every file as 0666, and making links there needs privileges. +_POSIX = sys.platform != "win32" + + +def _zip(entries: list[ZipEntry]) -> bytes: + file = io.BytesIO() + write_zip(entries, file) + return file.getvalue() + + +def _tree(root: Path) -> None: + (root / "Dockerfile").write_text("FROM scratch") + (root / "run.sh").write_text("#!/bin/sh") + (root / "sub").mkdir() + (root / "sub" / "a.txt").write_text("a") + + +def test_builds_files_entries_sorted_with_mode_0644() -> None: + entries = files_zip_entries({"b/c.txt": b"\x00", "Dockerfile": "FROM x"}) + assert entries == [ + ZipEntry(path="Dockerfile", data=b"FROM x", mode=0o644), + ZipEntry(path="b/c.txt", data=b"\x00", mode=0o644), + ] + + +def test_files_entries_require_a_dockerfile_and_safe_paths() -> None: + with pytest.raises(ValueError, match="has no Dockerfile"): + files_zip_entries({"a": "x"}) + for path in ["/abs", "a//b", "./a", "a/../b", "a\\b"]: + with pytest.raises(ValueError, match="must be relative"): + files_zip_entries({"Dockerfile": "x", path: "y"}) + + +def test_zips_entries_read_back_identically_with_modes(tmp_path: Path) -> None: + local = tmp_path / "local.bin" + local.write_bytes(b"\x01" * 100_000) + data = _zip( + [ + ZipEntry(path="Dockerfile", data=b"FROM x", mode=0o644), + ZipEntry(path="bin", mode=0o755, is_dir=True), + ZipEntry(path="bin/tool", mode=0o755, source_path=str(local)), + ] + ) + with zipfile.ZipFile(io.BytesIO(data)) as archive: + infos = archive.infolist() + assert [info.filename for info in infos] == ["Dockerfile", "bin/", "bin/tool"] + assert archive.read("bin/tool") == b"\x01" * 100_000 + modes = {info.filename: info.external_attr >> 16 for info in infos} + assert modes == { + "Dockerfile": stat.S_IFREG | 0o644, + "bin/": stat.S_IFDIR | 0o755, + "bin/tool": stat.S_IFREG | 0o755, + } + + +def test_makes_the_same_bytes_for_the_same_entries() -> None: + entries = files_zip_entries({"Dockerfile": "FROM x", "a.txt": "a" * 1000}) + assert _zip(entries) == _zip(entries) + + +def test_checks_a_zip_for_a_root_dockerfile_file() -> None: + check_zip_dockerfile(_zip(files_zip_entries({"Dockerfile": "FROM x"}))) + with pytest.raises(ValueError, match="no Dockerfile file at its root"): + check_zip_dockerfile( + _zip(files_zip_entries({"Dockerfile": "x", "sub/Dockerfile": "x"})[1:]) + ) + with pytest.raises(ValueError, match="not a readable zip"): + check_zip_dockerfile(b"not a zip") + + +def test_rejects_a_zip_whose_dockerfile_is_a_link() -> None: + file = io.BytesIO() + with zipfile.ZipFile(file, "w") as archive: + info = zipfile.ZipInfo("Dockerfile") + info.create_system = 3 + info.external_attr = (stat.S_IFLNK | 0o777) << 16 + archive.writestr(info, "elsewhere") + with pytest.raises(ValueError, match="no Dockerfile file"): + check_zip_dockerfile(file.getvalue()) + + +def test_accepts_a_zip_made_without_unix_modes() -> None: + file = io.BytesIO() + with zipfile.ZipFile(file, "w") as archive: + archive.writestr(zipfile.ZipInfo("Dockerfile"), "FROM x") + check_zip_dockerfile(file.getvalue()) + + +def test_checks_a_directory_for_a_root_dockerfile(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="has no Dockerfile at its root"): + check_directory_dockerfile(str(tmp_path)) + (tmp_path / "Dockerfile").mkdir() + with pytest.raises(ValueError, match="has no Dockerfile at its root"): + check_directory_dockerfile(str(tmp_path)) + + +def test_walks_a_directory_sorted_with_modes(tmp_path: Path) -> None: + _tree(tmp_path) + (tmp_path / "run.sh").chmod(0o755) + entries = run_walk(walk_directory(str(tmp_path), "ctx/"), None) + assert [(entry.path, entry.is_dir) for entry in entries] == [ + ("ctx/Dockerfile", False), + ("ctx/run.sh", False), + ("ctx/sub", True), + ("ctx/sub/a.txt", False), + ] + assert entries[3].source_path == str(tmp_path / "sub" / "a.txt") + if _POSIX: + assert entries[1].mode == 0o755 + + +@pytest.mark.skipif(not _POSIX, reason="making links needs privileges on Windows") +def test_walks_links_copying_files_within_and_storing_empty_directories( + tmp_path: Path, +) -> None: + _tree(tmp_path) + (tmp_path / "file-link").symlink_to(tmp_path / "sub" / "a.txt") + (tmp_path / "dir-link").symlink_to(tmp_path / "sub") + (tmp_path / "broken-link").symlink_to(tmp_path / "gone") + entries = { + entry.path: entry for entry in run_walk(walk_directory(str(tmp_path), ""), None) + } + assert "broken-link" not in entries + assert entries["file-link"].source_path == str(tmp_path / "file-link") + assert entries["dir-link"].is_dir + assert not any(path.startswith("dir-link/") for path in entries) + + +@pytest.mark.skipif(not _POSIX, reason="making links needs privileges on Windows") +def test_fails_a_walk_with_a_link_to_a_file_outside(tmp_path: Path) -> None: + root = tmp_path / "root" + root.mkdir() + _tree(root) + (tmp_path / "secret").write_text("s") + (root / "leak").symlink_to(tmp_path / "secret") + with pytest.raises(ValueError, match="outside"): + run_walk(walk_directory(str(root), ""), None) + + +def test_walk_asks_about_every_path_but_the_root_dockerfile_and_prunes( + tmp_path: Path, +) -> None: + _tree(tmp_path) + (tmp_path / ".dockerignore").write_text("sub") + asked: list[ImageIgnoreFileOptions] = [] + + def ignore(options: ImageIgnoreFileOptions) -> bool: + asked.append(options) + return options.rel_path == "sub" + + entries = run_walk(walk_directory(str(tmp_path), ""), ignore) + assert [entry.path for entry in entries] == [ + ".dockerignore", + "Dockerfile", + "run.sh", + ] + assert asked == [ + ImageIgnoreFileOptions(rel_path="run.sh", is_directory=False), + ImageIgnoreFileOptions(rel_path="sub", is_directory=True), + ] + + +def test_walk_fails_when_the_ignore_function_raises(tmp_path: Path) -> None: + _tree(tmp_path) + + def ignore(_: ImageIgnoreFileOptions) -> bool: + raise RuntimeError("bad pattern") + + with pytest.raises(RuntimeError, match="bad pattern"): + run_walk(walk_directory(str(tmp_path), ""), ignore) + + +@pytest.mark.asyncio +async def test_walks_asynchronously_with_sync_or_async_ignore(tmp_path: Path) -> None: + _tree(tmp_path) + + async def ignore_async(options: ImageIgnoreFileOptions) -> bool: + return options.rel_path == "sub" + + entries = await run_walk_async(walk_directory(str(tmp_path), ""), ignore_async) + assert [entry.path for entry in entries] == ["Dockerfile", "run.sh"] + entries = await run_walk_async( + walk_directory(str(tmp_path), ""), lambda options: options.rel_path == "run.sh" + ) + assert [entry.path for entry in entries] == ["Dockerfile", "sub", "sub/a.txt"] + entries = await run_walk_async(walk_directory(str(tmp_path), ""), None) + assert len(entries) == 4 diff --git a/uv.lock b/uv.lock index 9dbab2f..f5f0f24 100644 --- a/uv.lock +++ b/uv.lock @@ -96,6 +96,7 @@ dev = [ { name = "autodoc-pydantic" }, { name = "datamodel-code-generator" }, { name = "furo" }, + { name = "h2" }, { name = "poethepoet" }, { name = "pytest" }, { name = "pytest-asyncio" }, @@ -117,6 +118,7 @@ dev = [ { name = "autodoc-pydantic", specifier = ">=2.2.0" }, { name = "datamodel-code-generator", specifier = ">=0.81.0,<0.82" }, { name = "furo", specifier = ">=2025.12.19" }, + { name = "h2", specifier = ">=4.4.1" }, { name = "poethepoet", specifier = ">=0.48" }, { name = "pytest", specifier = ">=9.1.1" }, { name = "pytest-asyncio", specifier = ">=1.4.0" }, @@ -414,6 +416,28 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/04/4b/29cac41a4d98d144bf5f6d33995617b185d14b22401f75ca86f384e87ff1/h11-0.16.0-py3-none-any.whl", hash = "sha256:63cf8bbe7522de3bf65932fda1d9c2772064ffb3dae62d55932da54b31cb6c86", size = 37515, upload-time = "2025-04-24T03:35:24.344Z" }, ] +[[package]] +name = "h2" +version = "4.4.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "hpack" }, + { name = "hyperframe" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e7/85/7c366e69d84c17bb778fe41419e1fbcce3033d5b7ce29bbffff0a98b859f/h2-4.4.1.tar.gz", hash = "sha256:4e866ffb1a869ae14dd9b5e6beb5c24a13da0495ad72b65925ded182521c1516", size = 2157281, upload-time = "2026-08-03T11:45:09.509Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7e/22/e85faf23bd72a92d1921e37d674ca56eb298a3c8be31fdecef0ff2b3aaac/h2-4.4.1-py3-none-any.whl", hash = "sha256:0e25f1462b23c9cb82d9eb02e28bc706dac2a68cb457c6a0d74d63c8a2a5d0e6", size = 62636, upload-time = "2026-08-03T11:44:59.164Z" }, +] + +[[package]] +name = "hpack" +version = "4.2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/26/5b/fcabf6028144a8723726318b07a32c2f3314acdff6265743cf08a344b18e/hpack-4.2.0.tar.gz", hash = "sha256:0895cfa3b5531fc65fe439c05eb65144f123bf7a394fcaa56aa423548d8e45c0", size = 51300, upload-time = "2026-06-23T18:34:46.667Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/71/b4/4a9fcfb2aef6ba44d9073ecd301443aa00b3dac95de5619f2a7de7ec8a91/hpack-4.2.0-py3-none-any.whl", hash = "sha256:858ac0b02280fa582b5080d68db0899c62a80375e0e5413a74970c5e518b6986", size = 34246, upload-time = "2026-06-23T18:34:45.472Z" }, +] + [[package]] name = "httpcore" version = "1.0.9" @@ -442,6 +466,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl", hash = "sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad", size = 73517, upload-time = "2024-12-06T15:37:21.509Z" }, ] +[[package]] +name = "hyperframe" +version = "6.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/02/e7/94f8232d4a74cc99514c13a9f995811485a6903d48e5d952771ef6322e30/hyperframe-6.1.0.tar.gz", hash = "sha256:f630908a00854a7adeabd6382b43923a4c4cd4b821fcb527e6ab9e15382a3b08", size = 26566, upload-time = "2025-01-22T21:41:49.302Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/48/30/47d0bf6072f7252e6521f3447ccfa40b421b6824517f82854703d0f5a98b/hyperframe-6.1.0-py3-none-any.whl", hash = "sha256:b03380493a519fce58ea5af42e4a42317bf9bd425596f7a0835ffce80f1a42e5", size = 13007, upload-time = "2025-01-22T21:41:47.295Z" }, +] + [[package]] name = "idna" version = "3.19"