diff --git a/runpod/api/sandboxes.py b/runpod/api/sandboxes.py index e5e6663d9..932e1e371 100644 --- a/runpod/api/sandboxes.py +++ b/runpod/api/sandboxes.py @@ -471,6 +471,20 @@ async def list_registry_auths(self) -> list[dict[str, Any]]: raise ValueError("Registry listing returned no response") return payload["registries"] + async def options(self) -> dict[str, Any]: + path = f"{_SANDBOX_PATH}/options" + payload = await self._request("GET", path) + if ( + not isinstance(payload, dict) + or not all( + isinstance(payload.get(key), list) + for key in ("dataCenterIds", "cpuFlavorIds") + ) + or not isinstance(payload.get("limits"), dict) + ): + raise error.QueryError("sandbox options response is malformed", f"GET {path}") + return payload + async def create(self, body: Mapping[str, Any]) -> dict[str, Any]: payload = await self._request( "POST", diff --git a/runpod/sandbox/__init__.py b/runpod/sandbox/__init__.py index f6d2fd128..ec7604e6b 100644 --- a/runpod/sandbox/__init__.py +++ b/runpod/sandbox/__init__.py @@ -10,6 +10,8 @@ SandboxCompute, SandboxExecutionError, SandboxInfo, + SandboxLimits, + SandboxOptions, SandboxStartupTimeout, SandboxState, SandboxStateError, @@ -28,6 +30,8 @@ "SandboxExecutionError", "SandboxCreationError", "SandboxInfo", + "SandboxLimits", + "SandboxOptions", "SandboxStartupTimeout", "SandboxState", "SandboxStateError", diff --git a/runpod/sandbox/asyncio.py b/runpod/sandbox/asyncio.py index e1c602791..e5e12824a 100644 --- a/runpod/sandbox/asyncio.py +++ b/runpod/sandbox/asyncio.py @@ -23,6 +23,7 @@ SandboxCompute, SandboxExecutionError, SandboxInfo, + SandboxOptions, SandboxStartupTimeout, SandboxState, SandboxStateError, @@ -244,6 +245,25 @@ async def get( raise return sandbox + @classmethod + async def options( + cls, + *, + api_key: Optional[str] = None, + base_url: Optional[str] = None, + request_timeout: float = 30, + ) -> SandboxOptions: + """Return the data centers and CPU flavors a create may choose from.""" + _timeout(request_timeout, "request_timeout") + api = AsyncSandboxAPI(api_key, base_url, request_timeout) + try: + result = SandboxOptions.from_dict(await api.options()) + except BaseException as error: + await _cleanup(api.close(), request_timeout, error) + raise + await _cleanup(api.close(), request_timeout) + return result + @classmethod async def list( cls, diff --git a/runpod/sandbox/models.py b/runpod/sandbox/models.py index f65da1fed..c10ee7ea5 100644 --- a/runpod/sandbox/models.py +++ b/runpod/sandbox/models.py @@ -110,6 +110,40 @@ def from_dict(cls, data: Mapping[str, Any]) -> "SandboxInfo": ) +@dataclass(frozen=True) +class SandboxLimits: + """limits for this account; capacity and other caps can still refuse a create.""" + + max_concurrent_sandboxes: int + max_idle_timeout_seconds: int + max_lifetime_seconds: int + + @classmethod + def from_dict(cls, payload: Mapping[str, Any]) -> "SandboxLimits": + return cls( + max_concurrent_sandboxes=payload["maxConcurrentSandboxes"], + max_idle_timeout_seconds=payload["maxIdleTimeoutSeconds"], + max_lifetime_seconds=payload["maxLifetimeSeconds"], + ) + + +@dataclass(frozen=True) +class SandboxOptions: + """what this account may choose when creating a sandbox.""" + + data_center_ids: tuple[str, ...] + cpu_flavor_ids: tuple[str, ...] + limits: SandboxLimits + + @classmethod + def from_dict(cls, payload: Mapping[str, Any]) -> "SandboxOptions": + return cls( + data_center_ids=tuple(payload["dataCenterIds"]), + cpu_flavor_ids=tuple(payload["cpuFlavorIds"]), + limits=SandboxLimits.from_dict(payload["limits"]), + ) + + @dataclass(frozen=True) class ExecResult: """command output with nullable host-reported execution details.""" diff --git a/runpod/sandbox/sync.py b/runpod/sandbox/sync.py index f01acf625..1fffe68b2 100644 --- a/runpod/sandbox/sync.py +++ b/runpod/sandbox/sync.py @@ -23,6 +23,7 @@ LogSource, SandboxCompute, SandboxInfo, + SandboxOptions, SandboxState, ) @@ -400,6 +401,27 @@ async def get() -> AsyncioSandbox: raise original from error raise + @classmethod + def options( + cls, + *, + api_key: Optional[str] = None, + base_url: Optional[str] = None, + request_timeout: float = 30, + ) -> SandboxOptions: + """Return the data centers and CPU flavors a create may choose from.""" + runner = _LoopRunner() + try: + return runner.run( + AsyncioSandbox.options( + api_key=api_key, + base_url=base_url, + request_timeout=request_timeout, + ) + ) + finally: + runner.stop() + @classmethod def list( cls, diff --git a/tests/test_sandbox.py b/tests/test_sandbox.py index af24e9523..88358e4e0 100644 --- a/tests/test_sandbox.py +++ b/tests/test_sandbox.py @@ -21,6 +21,7 @@ from runpod.sandbox import ( SandboxCreationError, SandboxExecutionError, + SandboxLimits, SandboxStartupTimeout, SandboxStateError, ) @@ -140,6 +141,18 @@ async def handle(request): ) if request.headers.get("Authorization") != "Bearer sandbox-test-key": return failure(401, "invalid key") + if request.path == "/v2/sandboxes/options": + return web.json_response( + { + "dataCenterIds": ["EU-RO-1", "US-CA-2"], + "cpuFlavorIds": ["cpu3c"], + "limits": { + "maxConcurrentSandboxes": 10, + "maxIdleTimeoutSeconds": 7200, + "maxLifetimeSeconds": 86400, + }, + } + ) if request.path == "/v2/global-volumes": if service.global_error: return failure(403, service.global_error) @@ -300,6 +313,7 @@ async def serve(): "/v2/sandboxes", "/v2/network-volumes", "/v2/registries", + "/v2/sandboxes/options", "/v2/sandboxes/{id}", "/v2/sandboxes/{id}/{operation}", ): @@ -1381,3 +1395,34 @@ async def test_same_resolved_volume_id_cannot_be_mounted_as_two_kinds(peer): **peer.options, ) assert peer.records == {} + + +def _connection(peer): + return {key: value for key, value in peer.options.items() if key != "startup_timeout"} + + +@pytest.mark.asyncio +async def test_options_returns_what_create_may_choose(peer): + options = await AsyncioSandbox.options(**_connection(peer)) + assert options.data_center_ids == ("EU-RO-1", "US-CA-2") + assert options.cpu_flavor_ids == ("cpu3c",) + assert options.limits == SandboxLimits( + max_concurrent_sandboxes=10, + max_idle_timeout_seconds=7200, + max_lifetime_seconds=86400, + ) + assert peer.requests[-1][:2] == ("GET", "/v2/sandboxes/options") + + +def test_sync_options_matches_async(peer): + options = Sandbox.options(**_connection(peer)) + assert options.data_center_ids == ("EU-RO-1", "US-CA-2") + assert options.cpu_flavor_ids == ("cpu3c",) + assert options.limits.max_concurrent_sandboxes == 10 + + +@pytest.mark.asyncio +async def test_options_authentication_failure_propagates(peer): + bad = {**_connection(peer), "api_key": "wrong"} + with pytest.raises(AuthenticationError): + await AsyncioSandbox.options(**bad)