API spec update and new CPU field for sandbox create/info

baseten/client/sandboxapi/__init__.py

@@ -18,6 +18,8 @@
     ArchiveManifest,
     ContentSearchMatch,
     ContentSearchResponse,
+    CopyRequest,
+    CopyResponse,
     DeleteFilesystemParams,
     DeleteFilesystemTreeParams,
     DeleteNetworkProcessMonitorResponse,
@@ -49,6 +51,7 @@
     GetFilesystemTreeParams,
     GetNetworkProcessPortsResponse,
     GetProcessResponse,
+    GetProcessWaitParams,
     GetWatchFilesystemParams,
     HandlerReloadResponse,
     HealthResponse,
@@ -87,6 +90,8 @@
     "AsyncApiClient",
     "ContentSearchMatch",
     "ContentSearchResponse",
+    "CopyRequest",
+    "CopyResponse",
     "DeleteFilesystemParams",
     "DeleteFilesystemTreeParams",
     "DeleteNetworkProcessMonitorResponse",
@@ -118,6 +123,7 @@
     "GetFilesystemTreeParams",
     "GetNetworkProcessPortsResponse",
     "GetProcessResponse",
+    "GetProcessWaitParams",
     "GetWatchFilesystemParams",
     "HandlerReloadResponse",
     "HealthResponse",

baseten/client/sandboxapi/_client.py

@@ -15,6 +15,8 @@
     ApplyEditRequest,
     ApplyEditResponse,
     ContentSearchResponse,
+    CopyRequest,
+    CopyResponse,
     DeleteFilesystemParams,
     DeleteFilesystemTreeParams,
     DeleteNetworkProcessMonitorResponse,
@@ -39,6 +41,7 @@
     GetFilesystemTreeParams,
     GetNetworkProcessPortsResponse,
     GetProcessResponse,
+    GetProcessWaitParams,
     GetWatchFilesystemParams,
     HandlerReloadResponse,
     HealthResponse,
@@ -628,6 +631,23 @@ def get_process_logs_stream(
             )
         )
 
+    def get_process_wait(
+        self, *, identifier: str, params: GetProcessWaitParams | None = None
+    ) -> ProcessResponse:
+        """Wait for a process to finish"""
+        return self._do_json(
+            ProcessResponse,
+            _ApiRequest(
+                method="GET",
+                path_fmt="/process/{}/wait",
+                path_args=[identifier],
+                body=None,
+                query=params,
+                success_codes=[200],
+                error_codes={400: "ErrorResponse", 404: "ErrorResponse"},
+            ),
+        )
+
     def get_watch_filesystem(
         self, *, path: str, params: GetWatchFilesystemParams | None = None
     ) -> httpx.Response:
@@ -731,6 +751,25 @@ def post_environment_reload(self) -> HandlerReloadResponse:
             ),
         )
 
+    def post_filesystem_copy(self, *, request: CopyRequest) -> CopyResponse:
+        """Copy a file or directory"""
+        return self._do_json(
+            CopyResponse,
+            _ApiRequest(
+                method="POST",
+                path_fmt="/filesystem-copy",
+                path_args=[],
+                body=request,
+                query=None,
+                success_codes=[200],
+                error_codes={
+                    400: "ErrorResponse",
+                    409: "ErrorResponse",
+                    422: "ErrorResponse",
+                },
+            ),
+        )
+
     def post_filesystem_multipart_complete(
         self, *, upload_id: str, request: MultipartCompleteRequest
     ) -> SuccessResponse:
@@ -868,7 +907,7 @@ def post_upgrade(self, *, request: UpgradeRequest) -> SuccessResponse:
                 body=request,
                 query=None,
                 success_codes=[200],
-                error_codes={500: "ErrorResponse"},
+                error_codes={400: "ErrorResponse", 500: "ErrorResponse"},
             ),
         )
 
@@ -1632,6 +1671,23 @@ async def get_process_logs_stream(
             )
         )
 
+    async def get_process_wait(
+        self, *, identifier: str, params: GetProcessWaitParams | None = None
+    ) -> ProcessResponse:
+        """Wait for a process to finish"""
+        return await self._do_json(
+            ProcessResponse,
+            _ApiRequest(
+                method="GET",
+                path_fmt="/process/{}/wait",
+                path_args=[identifier],
+                body=None,
+                query=params,
+                success_codes=[200],
+                error_codes={400: "ErrorResponse", 404: "ErrorResponse"},
+            ),
+        )
+
     async def get_watch_filesystem(
         self, *, path: str, params: GetWatchFilesystemParams | None = None
     ) -> httpx.Response:
@@ -1737,6 +1793,25 @@ async def post_environment_reload(self) -> HandlerReloadResponse:
             ),
         )
 
+    async def post_filesystem_copy(self, *, request: CopyRequest) -> CopyResponse:
+        """Copy a file or directory"""
+        return await self._do_json(
+            CopyResponse,
+            _ApiRequest(
+                method="POST",
+                path_fmt="/filesystem-copy",
+                path_args=[],
+                body=request,
+                query=None,
+                success_codes=[200],
+                error_codes={
+                    400: "ErrorResponse",
+                    409: "ErrorResponse",
+                    422: "ErrorResponse",
+                },
+            ),
+        )
+
     async def post_filesystem_multipart_complete(
         self, *, upload_id: str, request: MultipartCompleteRequest
     ) -> SuccessResponse:
@@ -1874,7 +1949,7 @@ async def post_upgrade(self, *, request: UpgradeRequest) -> SuccessResponse:
                 body=request,
                 query=None,
                 success_codes=[200],
-                error_codes={500: "ErrorResponse"},
+                error_codes={400: "ErrorResponse", 500: "ErrorResponse"},
             ),
         )
 

baseten/client/sandboxapi/_models.py

@@ -114,6 +114,18 @@ class ContentSearchResponse(BaseModel):
     total: Annotated[int, Field(examples=[5])]
 
 
+class CopyRequest(BaseModel):
+    destination: Annotated[str, Field(examples=["/app/destination"])]
+    noOverwrite: Annotated[bool, Field(examples=[False])] = False
+    source: Annotated[str, Field(examples=["/app/source"])]
+
+
+class CopyResponse(BaseModel):
+    destination: str
+    message: str
+    source: str
+
+
 class DriveMountInfo(BaseModel):
     driveName: str | None = None
     drivePath: str | None = None
@@ -168,6 +180,7 @@ class DriveUnmountResponse(BaseModel):
 
 
 class ErrorResponse(BaseModel):
+    code: Annotated[str | None, Field(examples=["FILE_ALREADY_EXISTS"])] = None
     error: Annotated[str, Field(examples=["Error message"])]
 
 
@@ -302,6 +315,20 @@ class ProcessLogs(BaseModel):
 
 
 class ProcessRequest(BaseModel):
+    callbackSecret: Annotated[
+        str | None,
+        Field(
+            description="Key the webhook is signed with (HMAC-SHA256, X-Blaxel-Signature header). At least 32 characters; required with callbackUrl. Never returned by the API.",
+            examples=["a-random-secret-of-at-least-32-chars"],
+        ),
+    ] = None
+    callbackUrl: Annotated[
+        str | None,
+        Field(
+            description="URL that receives a signed process.exited webhook when the process ends for good (not on auto-restart). Requires callbackSecret.",
+            examples=["https://example.com/hooks/process"],
+        ),
+    ] = None
     command: Annotated[str, Field(examples=["ls -la"])]
     env: Annotated[dict[str, str] | None, Field(examples=[{'{"PORT"': ' "3000"}'}])] = (
         None
@@ -431,15 +458,12 @@ class TunnelConfigRequest(BaseModel):
 class UpgradeRequest(BaseModel):
     baseUrl: Annotated[
         str | None,
-        Field(
-            description="Base URL for releases (useful for forks)",
-            examples=["https://github.com/blaxel-ai/sandbox/releases"],
-        ),
+        Field(description="Deprecated: custom binary sources are not accepted"),
     ] = None
     version: Annotated[
         str | None,
         Field(
-            description='Version to upgrade to: "latest" (default), "develop", "main", or specific tag like "v1.0.0"',
+            description='Version to upgrade to: "latest" (default), "develop", "main", or a release tag like "v1.0.0"',
             examples=["latest"],
         ),
     ] = None
@@ -609,6 +633,15 @@ class GetFilesystemTreeParams(BaseModel):
     ] = None
 
 
+class GetProcessWaitParams(BaseModel):
+    timeout: Annotated[
+        int | None,
+        Field(
+            description="Seconds to wait (default and max 119; larger values are capped)"
+        ),
+    ] = None
+
+
 class GetWatchFilesystemParams(BaseModel):
     ignore: Annotated[
         str | None,

baseten/sandbox/_client.py

@@ -296,6 +296,7 @@ def create(
         create_if_not_exists: bool | None = None,
         image: str | None = None,
         memory: int | None = None,
+        cpu: int | None = None,
         region: str | None = None,
         envs: Mapping[str, SandboxEnvValue] | None = None,
         labels: Mapping[str, str] | None = None,
@@ -317,8 +318,11 @@ def create(
                 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.
+            memory: Memory in megabytes. Defaults to 4096, or to 8192 per vCPU
+                when only *cpu* is set. When *cpu* is ``None``, it also sets
+                the CPU allocation.
+            cpu: Number of vCPUs, from 1 to 32. ``None`` derives the CPU
+                allocation from *memory*.
             region: Region to run in. ``None`` picks the closest one.
             envs: Environment variables, by name.
             labels: Labels, by key.
@@ -344,6 +348,7 @@ def create(
                         create_if_not_exists=create_if_not_exists,
                         image=image,
                         memory=memory,
+                        cpu=cpu,
                         region=region,
                         envs=None if envs is None else envs_to_api(envs),
                         labels=None if labels is None else dict(labels),
@@ -700,6 +705,7 @@ async def create(
         create_if_not_exists: bool | None = None,
         image: str | None = None,
         memory: int | None = None,
+        cpu: int | None = None,
         region: str | None = None,
         envs: Mapping[str, SandboxEnvValue] | None = None,
         labels: Mapping[str, str] | None = None,
@@ -722,6 +728,7 @@ async def create(
                         create_if_not_exists=create_if_not_exists,
                         image=image,
                         memory=memory,
+                        cpu=cpu,
                         region=region,
                         envs=None if envs is None else envs_to_api(envs),
                         labels=None if labels is None else dict(labels),

baseten/sandbox/_errors.py

@@ -218,7 +218,7 @@ def api_error_from(
     """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")
+        body = err.error_response.model_dump(mode="json", exclude_unset=True)
     else:
         try:
             body = json.loads(err.body)

baseten/sandbox/_info.py

@@ -247,7 +247,10 @@ class SandboxInfo:
     """Image reference, including its tag."""
 
     memory: int | None
-    """Memory in megabytes, which also sets the CPU allocation."""
+    """Memory in megabytes."""
+
+    cpu: int | None
+    """Number of vCPUs the sandbox runs with."""
 
     region: str | None
 
@@ -296,6 +299,7 @@ def sandbox_info_from_api(sandbox: baseten.client.managementapi.Sandbox) -> Sand
         status=sandbox.status,
         image=sandbox.image,
         memory=sandbox.memory,
+        cpu=sandbox.cpu,
         region=sandbox.region,
         envs=_envs_from_api(sandbox.envs),
         labels=dict(sandbox.labels or {}),