diff --git a/README.md b/README.md index 99efb01..8389c8b 100644 --- a/README.md +++ b/README.md @@ -68,6 +68,10 @@ with Sandbox.create(plan="small") as sb: sb = Sandbox.from_id("sb_...") etl_boxes = Sandbox.list(tags={"job": "etl-42"}) +# List available plans (id, vcpu, ram_gb, disk_gb, price_per_hour) +for plan in Sandbox.catalog(): + print(plan.id, plan.vcpu, plan.ram_gb, plan.price_per_hour) + # Large scripts: upload, then run sb.fs.write("/work/script.py", open("script.py").read()) sb.exec("python3", "/work/script.py", timeout="30m") diff --git a/deepinfra/__init__.py b/deepinfra/__init__.py index 2b18749..bfc53c7 100644 --- a/deepinfra/__init__.py +++ b/deepinfra/__init__.py @@ -21,7 +21,7 @@ SandboxWaitError, TooManySandboxesError, ) -from ._sandbox_models import ExecResult, SandboxInfo +from ._sandbox_models import ExecResult, SandboxInfo, SandboxPlan from ._version import __version__ from .clients import DeepInfraClient, RequestSpec from .models import ( @@ -39,6 +39,7 @@ "Sandbox", "SandboxFS", "SandboxInfo", + "SandboxPlan", "ExecResult", # client "DeepInfraClient", diff --git a/deepinfra/_sandbox_models.py b/deepinfra/_sandbox_models.py index 242423c..9a64995 100644 --- a/deepinfra/_sandbox_models.py +++ b/deepinfra/_sandbox_models.py @@ -19,6 +19,18 @@ class SandboxInfo(pydantic.BaseModel): model_config = pydantic.ConfigDict(extra="ignore") +class SandboxPlan(pydantic.BaseModel): + """A plan offered by GET /v1/sandboxes/catalog.""" + + id: str + vcpu: int + ram_gb: int + disk_gb: int + price_per_hour: float + + model_config = pydantic.ConfigDict(extra="ignore") + + class ExecResult(pydantic.BaseModel): """Aggregated result of a sandbox command.""" diff --git a/deepinfra/sandbox_api.py b/deepinfra/sandbox_api.py index 5ce5a88..89a3bdc 100644 --- a/deepinfra/sandbox_api.py +++ b/deepinfra/sandbox_api.py @@ -22,7 +22,7 @@ from urllib.parse import quote from ._exceptions import NotFoundError, SandboxFailedError, SandboxTimeoutError -from ._sandbox_models import ExecResult, SandboxInfo +from ._sandbox_models import ExecResult, SandboxInfo, SandboxPlan from ._streaming import afold_exec_events, aiter_ndjson, fold_exec_events, iter_ndjson from ._utils import backoff_delays, parse_duration, tags_match from .clients.deepinfra import DeepInfraClient, RequestSpec, default_client @@ -173,6 +173,27 @@ async def alist( items = (await client.arequest(_list_spec())).json() return cls._from_list(items, tags, client) + @classmethod + def catalog( + cls, + *, + client: DeepInfraClient | None = None, + ) -> builtins.list[SandboxPlan]: + """List available sandbox plans with their specs and hourly pricing.""" + client = client or default_client() + items = client.request(_catalog_spec()).json() + return [SandboxPlan.model_validate(item) for item in items] + + @classmethod + async def acatalog( + cls, + *, + client: DeepInfraClient | None = None, + ) -> builtins.list[SandboxPlan]: + client = client or default_client() + items = (await client.arequest(_catalog_spec())).json() + return [SandboxPlan.model_validate(item) for item in items] + # -- lifecycle -- def refresh(self) -> Sandbox: @@ -438,6 +459,10 @@ def _list_spec() -> RequestSpec: return RequestSpec("GET", _SANDBOXES, retry_connect=True) +def _catalog_spec() -> RequestSpec: + return RequestSpec("GET", f"{_SANDBOXES}/catalog", retry_connect=True) + + def _op_spec(sandbox_id: str, op: str) -> RequestSpec: return RequestSpec("POST", _id_path(sandbox_id, op)) diff --git a/tests/test_sandbox.py b/tests/test_sandbox.py index 1398352..471cd3d 100644 --- a/tests/test_sandbox.py +++ b/tests/test_sandbox.py @@ -226,6 +226,28 @@ async def test_async_lifecycle(client): assert delete.call_count == 1 +@respx.mock +def test_catalog(client): + plans = [ + {"id": "nano", "vcpu": 1, "ram_gb": 1, "disk_gb": 10, "price_per_hour": 0.01}, + {"id": "medium", "vcpu": 4, "ram_gb": 8, "disk_gb": 40, "price_per_hour": 0.08}, + ] + get = respx.get(f"{BASE_URL}/v1/sandboxes/catalog").respond(json=plans) + result = Sandbox.catalog(client=client) + assert get.call_count == 1 + assert [p.id for p in result] == ["nano", "medium"] + assert result[0].vcpu == 1 + assert result[1].price_per_hour == 0.08 + + +@respx.mock +async def test_acatalog(client): + plans = [{"id": "nano", "vcpu": 1, "ram_gb": 1, "disk_gb": 10, "price_per_hour": 0.01}] + respx.get(f"{BASE_URL}/v1/sandboxes/catalog").respond(json=plans) + result = await Sandbox.acatalog(client=client) + assert result[0].id == "nano" + + @respx.mock async def test_async_context_manager(client): respx.get(f"{BASE_URL}/v1/sandboxes/{SB_ID}").respond(json=_info("running"))