diff --git a/src/blaxel/core/client/models/sandbox_runtime.py b/src/blaxel/core/client/models/sandbox_runtime.py index 20dc667c..3b10250b 100644 --- a/src/blaxel/core/client/models/sandbox_runtime.py +++ b/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, ) diff --git a/src/blaxel/core/image/image.py b/src/blaxel/core/image/image.py index e9f61959..2a9c864a 100644 --- a/src/blaxel/core/image/image.py +++ b/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) diff --git a/src/blaxel/core/sandbox/default/sandbox.py b/src/blaxel/core/sandbox/default/sandbox.py index 419ad99e..73c24504 100644 --- a/src/blaxel/core/sandbox/default/sandbox.py +++ b/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, ), diff --git a/src/blaxel/core/sandbox/sync/sandbox.py b/src/blaxel/core/sandbox/sync/sandbox.py index b1bf5b7f..9cb180f1 100644 --- a/src/blaxel/core/sandbox/sync/sandbox.py +++ b/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, ), diff --git a/src/blaxel/core/sandbox/types.py b/src/blaxel/core/sandbox/types.py index 22628d69..85e73fc9 100644 --- a/src/blaxel/core/sandbox/types.py +++ b/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"), diff --git a/tests/core/test_image.py b/tests/core/test_image.py index 963288f0..fc40a407 100644 --- a/tests/core/test_image.py +++ b/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 diff --git a/tests/core/test_sandbox.py b/tests/core/test_sandbox.py index 6be260b8..3ad57255 100644 --- a/tests/core/test_sandbox.py +++ b/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")