Add sandbox storage_mb support

src/blaxel/core/client/models/sandbox_runtime.py

@@ -31,6 +31,8 @@ class SandboxRuntime:
         memory (Union[Unset, int]): Memory allocation in megabytes. Also determines CPU allocation (CPU cores = memory
             in MB / 2048, e.g., 4096MB = 2 CPUs). Example: 4096.
         ports (Union[Unset, list['Port']]): Set of ports for a resource
+        storage_mb (Union[Unset, int]): Disk-backed root storage capacity in megabytes. When omitted, the sandbox uses
+            the default tmpfs overlay sizing based on its memory allocation. Example: 102400.
         termination_grace_period_seconds (Union[Unset, int]): Duration in seconds the pod needs to terminate gracefully.
             Defaults to 0 for immediate termination. Example: 30.
         ttl (Union[Unset, str]): Max-age from creation: the sandbox is deleted this long after it is created, regardless
@@ -44,12 +46,12 @@ class SandboxRuntime:
     image: Union[Unset, str] = UNSET
     memory: Union[Unset, int] = UNSET
     ports: Union[Unset, list["Port"]] = UNSET
+    storage_mb: Union[Unset, int] = UNSET
     termination_grace_period_seconds: Union[Unset, int] = UNSET
     ttl: Union[Unset, str] = UNSET
     additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
 
     def to_dict(self) -> dict[str, Any]:
-
         envs: Union[Unset, list[dict[str, Any]]] = UNSET
         if not isinstance(self.envs, Unset):
             envs = []
@@ -86,6 +88,8 @@ def to_dict(self) -> dict[str, Any]:
                     componentsschemas_ports_item = componentsschemas_ports_item_data.to_dict()
                 ports.append(componentsschemas_ports_item)
 
+        storage_mb = self.storage_mb
+
         termination_grace_period_seconds = self.termination_grace_period_seconds
 
         ttl = self.ttl
@@ -105,6 +109,8 @@ def to_dict(self) -> dict[str, Any]:
             field_dict["memory"] = memory
         if ports is not UNSET:
             field_dict["ports"] = ports
+        if storage_mb is not UNSET:
+            field_dict["storageMb"] = storage_mb
         if termination_grace_period_seconds is not UNSET:
             field_dict["terminationGracePeriodSeconds"] = termination_grace_period_seconds
         if ttl is not UNSET:
@@ -148,6 +154,8 @@ def from_dict(cls: type[T], src_dict: dict[str, Any]) -> T | None:
 
             ports.append(componentsschemas_ports_item)
 
+        storage_mb = d.pop("storageMb", d.pop("storage_mb", UNSET))
+
         termination_grace_period_seconds = d.pop(
             "terminationGracePeriodSeconds", d.pop("termination_grace_period_seconds", UNSET)
         )
@@ -161,6 +169,7 @@ def from_dict(cls: type[T], src_dict: dict[str, Any]) -> T | None:
             image=image,
             memory=memory,
             ports=ports,
+            storage_mb=storage_mb,
             termination_grace_period_seconds=termination_grace_period_seconds,
             ttl=ttl,
         )

src/blaxel/core/image/image.py

@@ -804,13 +804,19 @@ def _create_zip(self, build_dir: Path) -> bytes:
         zip_buffer.seek(0)
         return zip_buffer.getvalue()
 
-    def _create_sandbox_payload(self, name: str, memory: int = 4096) -> Sandbox:
+    def _create_sandbox_payload(
+        self,
+        name: str,
+        memory: int = 4096,
+        storage_mb: int | None = None,
+    ) -> Sandbox:
         """
         Create the sandbox payload for deployment.
 
         Args:
             name: Name for the sandbox
             memory: Memory in MB (default 4096)
+            storage_mb: Disk-backed root storage capacity in MB
 
         Returns:
             Sandbox object
@@ -824,6 +830,8 @@ def _create_sandbox_payload(self, name: str, memory: int = 4096) -> Sandbox:
         )
 
         runtime = SandboxRuntime(memory=memory)
+        if storage_mb is not None:
+            runtime.storage_mb = storage_mb
         spec = SandboxSpec(runtime=runtime)
 
         return Sandbox(metadata=metadata, spec=spec)
@@ -1137,6 +1145,7 @@ def build_sync(
         self,
         name: str,
         memory: int = 4096,
+        storage_mb: int | None = None,
         timeout: float = 900.0,
         on_status_change: Callable[[str], None] | None = None,
         sandbox_version: str = "latest",
@@ -1155,6 +1164,7 @@ def build_sync(
         Args:
             name: Name for the sandbox
             memory: Memory in MB (default 4096)
+            storage_mb: Disk-backed root storage capacity in MB
             timeout: Maximum time to wait for deployment in seconds (default 15 minutes)
             on_status_change: Optional callback called when status changes
             sandbox_version: Version of sandbox-api to use (default "latest")
@@ -1177,7 +1187,7 @@ def build_sync(
             zip_content = self._create_zip(build_dir)
 
             # Create sandbox payload
-            sandbox_payload = self._create_sandbox_payload(name, memory)
+            sandbox_payload = self._create_sandbox_payload(name, memory, storage_mb)
 
             # Create/update sandbox and get upload URL
             response, upload_url = self._create_sandbox_with_upload_sync(sandbox_payload)
@@ -1220,6 +1230,7 @@ async def build(
         self,
         name: str,
         memory: int = 4096,
+        storage_mb: int | None = None,
         timeout: float = 900.0,
         on_status_change: Callable[[str], None] | None = None,
         sandbox_version: str = "latest",
@@ -1238,6 +1249,7 @@ async def build(
         Args:
             name: Name for the sandbox
             memory: Memory in MB (default 4096)
+            storage_mb: Disk-backed root storage capacity in MB
             timeout: Maximum time to wait for deployment in seconds (default 15 minutes)
             on_status_change: Optional callback called when status changes
             sandbox_version: Version of sandbox-api to use (default "latest")
@@ -1260,7 +1272,7 @@ async def build(
             zip_content = self._create_zip(build_dir)
 
             # Create sandbox payload
-            sandbox_payload = self._create_sandbox_payload(name, memory)
+            sandbox_payload = self._create_sandbox_payload(name, memory, storage_mb)
 
             # Create/update sandbox and get upload URL
             response, upload_url = await self._create_sandbox_with_upload(sandbox_payload)

src/blaxel/core/sandbox/default/sandbox.py

@@ -221,6 +221,8 @@ async def create(
                     or "name" in (sandbox if isinstance(sandbox, dict) else sandbox.__dict__)
                     or "image" in (sandbox if isinstance(sandbox, dict) else sandbox.__dict__)
                     or "memory" in (sandbox if isinstance(sandbox, dict) else sandbox.__dict__)
+                    or "storage_mb" in (sandbox if isinstance(sandbox, dict) else sandbox.__dict__)
+                    or "storageMb" in (sandbox if isinstance(sandbox, dict) else sandbox.__dict__)
                     or "ports" in (sandbox if isinstance(sandbox, dict) else sandbox.__dict__)
                     or "envs" in (sandbox if isinstance(sandbox, dict) else sandbox.__dict__)
                     or "volumes" in (sandbox if isinstance(sandbox, dict) else sandbox.__dict__)
@@ -249,6 +251,7 @@ async def create(
             name = config.name or default_name
             image = config.image or default_image
             memory = config.memory or default_memory
+            storage_mb = config.storage_mb if config.storage_mb is not None else UNSET
             ports = config._normalize_ports() or UNSET
             envs = config._normalize_envs() or UNSET
             volumes = config._normalize_volumes() or UNSET
@@ -278,6 +281,7 @@ async def create(
                     runtime=SandboxRuntime(
                         image=image,
                         memory=memory,
+                        storage_mb=storage_mb,
                         ports=ports,
                         envs=envs,
                     ),

src/blaxel/core/sandbox/sync/sandbox.py

@@ -165,6 +165,8 @@ def create(
                     or "name" in (sandbox if isinstance(sandbox, dict) else sandbox.__dict__)
                     or "image" in (sandbox if isinstance(sandbox, dict) else sandbox.__dict__)
                     or "memory" in (sandbox if isinstance(sandbox, dict) else sandbox.__dict__)
+                    or "storage_mb" in (sandbox if isinstance(sandbox, dict) else sandbox.__dict__)
+                    or "storageMb" in (sandbox if isinstance(sandbox, dict) else sandbox.__dict__)
                     or "ports" in (sandbox if isinstance(sandbox, dict) else sandbox.__dict__)
                     or "envs" in (sandbox if isinstance(sandbox, dict) else sandbox.__dict__)
                     or "volumes" in (sandbox if isinstance(sandbox, dict) else sandbox.__dict__)
@@ -191,6 +193,7 @@ def create(
             name = config.name or default_name
             image = config.image or default_image
             memory = config.memory or default_memory
+            storage_mb = config.storage_mb if config.storage_mb is not None else UNSET
             ports = config._normalize_ports() or UNSET
             envs = config._normalize_envs() or UNSET
             volumes = config._normalize_volumes() or UNSET
@@ -214,6 +217,7 @@ def create(
                     runtime=SandboxRuntime(
                         image=image,
                         memory=memory,
+                        storage_mb=storage_mb,
                         ports=ports,
                         envs=envs,
                     ),

src/blaxel/core/sandbox/types.py

@@ -164,6 +164,7 @@ def __init__(
         name: str | None = None,
         image: str | None = None,
         memory: int | None = None,
+        storage_mb: int | None = None,
         ports: Union[List[Port], List[Dict[str, Any]]] | None = None,
         envs: List[Dict[str, str]] | None = None,
         volumes: Union[List[VolumeBinding], List[VolumeAttachment], List[Dict[str, Any]]]
@@ -180,6 +181,7 @@ def __init__(
         self.name = name
         self.image = image
         self.memory = memory
+        self.storage_mb = storage_mb
         self.ports = ports
         self.envs = envs
         self.volumes = volumes
@@ -210,6 +212,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "SandboxCreateConfiguration":
             name=data.get("name"),
             image=data.get("image"),
             memory=data.get("memory"),
+            storage_mb=data.get("storage_mb", data.get("storageMb")),
             ports=data.get("ports"),
             envs=data.get("envs"),
             volumes=data.get("volumes"),

tests/core/test_image.py

@@ -1256,3 +1256,16 @@ def test_sandbox_api_added_at_end_of_dockerfile(self):
         assert entrypoint_idx > run_idx
         # Entrypoint should be last
         assert entrypoint_idx > copy_idx
+
+    def test_create_sandbox_payload_includes_storage_mb(self):
+        """Test that image builds can request disk-backed root storage."""
+        image = ImageInstance.from_registry("python:3.11-slim")
+
+        payload = image._create_sandbox_payload(
+            "image-build",
+            memory=4096,
+            storage_mb=102400,
+        )
+
+        assert payload.spec.runtime.storage_mb == 102400
+        assert payload.spec.runtime.to_dict()["storageMb"] == 102400

tests/core/test_sandbox.py

@@ -214,6 +214,40 @@ async def test_create_forwards_create_if_not_exist_to_generated_client():
         assert mock_create_sandbox.await_args.kwargs["create_if_not_exist"] is True
 
 
+@pytest.mark.asyncio
+async def test_create_forwards_storage_mb_to_runtime():
+    created = sandbox_instance("storage").sandbox
+
+    with patch(
+        "blaxel.core.sandbox.default.sandbox.create_sandbox",
+        new_callable=AsyncMock,
+    ) as mock_create_sandbox:
+        mock_create_sandbox.return_value = created
+
+        await SandboxInstance.create(
+            {"name": "storage", "region": "us-pdx-1", "storage_mb": 102400},
+        )
+
+        body = mock_create_sandbox.await_args.kwargs["body"]
+        assert body.spec.runtime.storage_mb == 102400
+        assert body.spec.runtime.to_dict()["storageMb"] == 102400
+
+
+def test_sync_create_forwards_storage_mb_to_runtime():
+    created = sandbox_instance("sync-storage", cls=SyncSandboxInstance).sandbox
+
+    with patch("blaxel.core.sandbox.sync.sandbox.create_sandbox") as mock_create_sandbox:
+        mock_create_sandbox.return_value = created
+
+        SyncSandboxInstance.create(
+            {"name": "sync-storage", "region": "us-pdx-1", "storage_mb": 102400},
+        )
+
+        body = mock_create_sandbox.call_args.kwargs["body"]
+        assert body.spec.runtime.storage_mb == 102400
+        assert body.spec.runtime.to_dict()["storageMb"] == 102400
+
+
 @pytest.mark.asyncio
 async def test_create_if_not_exists_returns_existing_after_conflict():
     existing = sandbox_instance("existing")