diff --git a/runpod/api/mutations/apps.py b/runpod/api/mutations/apps.py index eb8f4db80..2e3c73a44 100644 --- a/runpod/api/mutations/apps.py +++ b/runpod/api/mutations/apps.py @@ -1,11 +1,5 @@ """GraphQL mutations for capabilities absent from the resource REST API.""" -MUTATION_CREATE_GLOBAL_VOLUME = """ -mutation createGlobalVolume($input: GlobalStoreBucketCreateInput!) { - globalStoreBucketCreate(input: $input) { id name } -} -""" - # endpoint saves retain flash bindings, cached models, and schedules atomically. MUTATION_SAVE_ENDPOINT = """ mutation saveEndpoint($input: EndpointInput!) { diff --git a/runpod/api/queries/apps.py b/runpod/api/queries/apps.py index 9362bb2fd..9da1b2cde 100644 --- a/runpod/api/queries/apps.py +++ b/runpod/api/queries/apps.py @@ -1,13 +1,5 @@ """GraphQL queries for capabilities absent from the resource REST API.""" -QUERY_GLOBAL_VOLUMES = """ -query myGlobalVolumes { - myself { - globalStoreBuckets { id name } - } -} -""" - # rest stock filters require a power-of-two count of at least two. QUERY_CPU_STOCK = """ query CpuStock($cpuFlavorInput: CpuFlavorInput, $specificsInput: SpecificsInput) { diff --git a/runpod/api/rest.py b/runpod/api/rest.py index 22e519c72..5f77b6190 100644 --- a/runpod/api/rest.py +++ b/runpod/api/rest.py @@ -1,7 +1,9 @@ """Runpod REST API transport.""" +import asyncio import json as json_module import os +from contextlib import nullcontext from typing import Any, Mapping, Optional import aiohttp @@ -17,7 +19,9 @@ def _resolve_api_key(api_key: Optional[str]) -> str: - from runpod import api_key as global_api_key # pylint: disable=import-outside-toplevel,cyclic-import + from runpod import ( + api_key as global_api_key, + ) # pylint: disable=import-outside-toplevel,cyclic-import effective_api_key = api_key or global_api_key if not effective_api_key: @@ -39,9 +43,7 @@ def _build_headers(api_key: str) -> dict[str, str]: } -def _raise_for_error( - status_code: int, method: str, path: str, text: str = "" -) -> None: +def _raise_for_error(status_code: int, method: str, path: str, text: str = "") -> None: if status_code == HTTP_STATUS_UNAUTHORIZED: raise error.AuthenticationError( "Unauthorized request, please check your API key." @@ -102,11 +104,13 @@ async def run_rest_request_async( params: Optional[Mapping[str, Any]] = None, json: Optional[Mapping[str, Any]] = None, timeout: float = 30, + session: Optional[aiohttp.ClientSession] = None, ) -> Optional[dict[str, Any]]: - """Send an authenticated REST request without blocking the event loop.""" + """send an authenticated request, leaving a supplied session open.""" headers = _build_headers(_resolve_api_key(api_key)) client_timeout = aiohttp.ClientTimeout(total=timeout) - async with aiohttp.ClientSession(timeout=client_timeout) as session: + context = nullcontext(session) if session is not None else aiohttp.ClientSession() + async with context as session: # aiohttp otherwise replays put/delete requests after an ambiguous disconnect. session._retry_connection = False # pylint: disable=protected-access async with session.request( @@ -115,6 +119,7 @@ async def run_rest_request_async( headers=headers, params=params, json=json, + timeout=client_timeout, ) as response: if response.status >= HTTP_STATUS_BAD_REQUEST: text = ( @@ -127,3 +132,54 @@ async def run_rest_request_async( if response.status == HTTP_STATUS_NO_CONTENT or not await response.read(): return None return await response.json(content_type=None) + + +async def run_rest_paginated_request_async( + path: str, + collection: str, + *, + api_key: Optional[str] = None, +) -> list[dict[str, Any]]: + """read all cursor pages with one session and four transport attempts per page.""" + records = [] + params = {} + seen_cursors = set() + query = f"GET {path}" + async with aiohttp.ClientSession() as session: + while True: + for attempt in range(4): + try: + data = await run_rest_request_async( + "GET", path, api_key=api_key, params=params, session=session + ) + break + except (aiohttp.ClientError, OSError, asyncio.TimeoutError): + if attempt == 3: + raise + await asyncio.sleep(2**attempt) + page = data.get(collection) if isinstance(data, dict) else None + if not isinstance(page, list) or any( + not isinstance(record, dict) for record in page + ): + raise error.QueryError(f"invalid {collection} list response", query) + pagination = data.get("pagination") + if not isinstance(pagination, dict) or not isinstance( + pagination.get("hasNextPage"), bool + ): + raise error.QueryError( + f"invalid {collection} pagination response", query + ) + records.extend(page) + if not pagination["hasNextPage"]: + return records + cursor = pagination.get("nextCursor") + if not isinstance(cursor, str) or not cursor: + raise error.QueryError( + f"{collection} pagination is missing a valid nextCursor", query + ) + if cursor in seen_cursors: + raise error.QueryError( + f"{collection} pagination repeated a cursor", query + ) + seen_cursors.add(cursor) + params = {"cursor": cursor} diff --git a/runpod/apps/api.py b/runpod/apps/api.py index 30b6c301d..9b557800a 100644 --- a/runpod/apps/api.py +++ b/runpod/apps/api.py @@ -1,7 +1,7 @@ """control-plane calls for app provisioning. rest handles equivalent resource operations; graphql handles flash lifecycle, -auth, secrets, global storage, task provisioning, and endpoint capabilities +auth, secrets, task provisioning, and endpoint capabilities absent from rest. management verbs for the wider sdk stay in runpod.api.ctl_commands. """ @@ -13,7 +13,7 @@ from ..api.graphql import run_graphql_query_async from ..api.mutations import apps as app_mutations from ..api.queries import apps as app_queries -from ..api.rest import run_rest_request_async +from ..api.rest import run_rest_paginated_request_async, run_rest_request_async from ..error import QueryError _TRANSPORT_RETRIES = 4 @@ -413,20 +413,36 @@ async def network_volume_datacenters(self) -> Set[str]: return supported async def list_global_volumes(self) -> List[Dict[str, Any]]: - data = await self._execute(app_queries.QUERY_GLOBAL_VOLUMES, retry=True) - return data["myself"]["globalStoreBuckets"] + """list every rest page, retrying transient read failures with backoff.""" + return await run_rest_paginated_request_async( + "/v2/global-volumes", "globalVolumes", api_key=self._api_key + ) async def create_global_volume(self, name: str) -> Dict[str, Any]: - data = await self._execute( - app_mutations.MUTATION_CREATE_GLOBAL_VOLUME, {"input": {"name": name}} + """create once via rest and require a usable volume id in the response.""" + data = await run_rest_request_async( + "POST", + "/v2/global-volumes", + api_key=self._api_key, + json={"name": name}, ) - return data["globalStoreBucketCreate"] + if ( + not isinstance(data, dict) + or not isinstance(data.get("id"), str) + or not data["id"] + ): + raise QueryError( + "global volume creation returned an invalid response; " + "the volume may have been created, check existing volumes before retrying", + "POST /v2/global-volumes", + ) + return data async def list_network_volumes(self) -> List[Dict[str, Any]]: - data = await run_rest_request_async( - "GET", "/v2/network-volumes", api_key=self._api_key + """list every rest page, retrying transient read failures with backoff.""" + return await run_rest_paginated_request_async( + "/v2/network-volumes", "networkVolumes", api_key=self._api_key ) - return data["networkVolumes"] async def create_network_volume( self, name: str, size: int, data_center_id: str diff --git a/tests/test_api/test_rest.py b/tests/test_api/test_rest.py index 38a7c5836..5e133aedf 100644 --- a/tests/test_api/test_rest.py +++ b/tests/test_api/test_rest.py @@ -5,13 +5,18 @@ import socket import threading from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from unittest.mock import AsyncMock, call, patch import aiohttp import pytest import requests import runpod -from runpod.api.rest import run_rest_request, run_rest_request_async +from runpod.api.rest import ( + run_rest_paginated_request_async, + run_rest_request, + run_rest_request_async, +) from runpod.error import AuthenticationError, QueryError @@ -45,7 +50,9 @@ def _respond(self): status, response_body = 401, b"invalid key" elif self.path != state.get("path", "/prefix/v2/pods"): status, response_body = 404, b"wrong path" - elif "expected_body" in state and json.loads(body) != state["expected_body"]: + elif ( + "expected_body" in state and json.loads(body) != state["expected_body"] + ): status, response_body = 400, b"wrong body" self.send_response(status) self.send_header("Content-Type", state["content_type"]) @@ -180,7 +187,9 @@ async def poll(): @pytest.mark.asyncio -async def test_cancellation_closes_pending_request_without_repeating_mutation(rest_server): +async def test_cancellation_closes_pending_request_without_repeating_mutation( + rest_server, +): rest_server["stall"] = True pending = asyncio.create_task( run_rest_request_async( @@ -208,3 +217,146 @@ async def test_timeout_closes_pending_request_without_repeating_mutation(rest_se ) await _wait_for_event(rest_server["disconnected"]) assert rest_server["hits"] == 1 + + +async def test_supplied_session_remains_open(rest_server): + async with aiohttp.ClientSession() as session: + for _ in range(2): + assert await run_rest_request_async( + "GET", "/v2/pods", api_key="key", session=session + ) == {"id": "pod-1"} + assert not session.closed + assert rest_server["hits"] == 2 + + +@pytest.mark.parametrize("method", ["POST", "PUT", "DELETE"]) +async def test_supplied_session_does_not_replay_mutations(rest_server, method): + async with aiohttp.ClientSession() as session: + await run_rest_request_async("GET", "/v2/pods", api_key="key", session=session) + rest_server["disconnect"] = True + with pytest.raises(aiohttp.ClientConnectionError): + await run_rest_request_async( + method, "/v2/pods", api_key="key", session=session + ) + assert rest_server["hits"] == 2 + + +class TestCursorPagination: + async def test_retries_current_page_with_backoff_and_one_session(self): + first = {"id": "one"} + second = {"id": "two"} + rest = AsyncMock( + side_effect=[ + { + "items": [first], + "pagination": {"hasNextPage": True, "nextCursor": "next"}, + }, + aiohttp.ClientError("reset"), + OSError("dns"), + asyncio.TimeoutError(), + {"items": [second], "pagination": {"hasNextPage": False}}, + ] + ) + with ( + patch("runpod.api.rest.run_rest_request_async", rest), + patch("runpod.api.rest.asyncio.sleep", new_callable=AsyncMock) as sleep, + ): + assert await run_rest_paginated_request_async( + "/v2/items", "items", api_key="key" + ) == [first, second] + assert sleep.await_args_list == [call(1), call(2), call(4)] + assert [item.kwargs["params"] for item in rest.await_args_list] == [ + {}, + *[{"cursor": "next"}] * 4, + ] + session = rest.await_args_list[0].kwargs["session"] + assert all(item.kwargs["session"] is session for item in rest.await_args_list) + assert session.closed + + async def test_exhausted_retries_close_session(self): + rest = AsyncMock(side_effect=aiohttp.ClientError("down")) + with ( + patch("runpod.api.rest.run_rest_request_async", rest), + patch("runpod.api.rest.asyncio.sleep", new_callable=AsyncMock) as sleep, + ): + with pytest.raises(aiohttp.ClientError, match="down"): + await run_rest_paginated_request_async("/v2/items", "items") + assert rest.await_count == 4 + assert sleep.await_args_list == [call(1), call(2), call(4)] + assert rest.await_args.kwargs["session"].closed + + @pytest.mark.parametrize( + "failure", + [ + QueryError("denied"), + AuthenticationError("bad key"), + asyncio.CancelledError(), + ], + ) + async def test_non_transport_errors_are_not_retried(self, failure): + rest = AsyncMock(side_effect=failure) + with patch("runpod.api.rest.run_rest_request_async", rest): + with pytest.raises(type(failure)): + await run_rest_paginated_request_async("/v2/items", "items") + rest.assert_awaited_once() + assert rest.await_args.kwargs["session"].closed + + @pytest.mark.parametrize( + "data", + [ + None, + [], + {}, + {"items": None}, + {"items": {}}, + {"items": [None]}, + {"items": []}, + {"items": [], "pagination": None}, + {"items": [], "pagination": []}, + {"items": [], "pagination": {}}, + {"items": [], "pagination": {"hasNextPage": "false"}}, + {"items": [], "pagination": {"hasNextPage": 0}}, + ], + ) + async def test_invalid_envelopes_raise_query_error(self, data): + rest = AsyncMock(return_value=data) + with patch("runpod.api.rest.run_rest_request_async", rest): + with pytest.raises(QueryError, match="invalid items") as raised: + await run_rest_paginated_request_async("/v2/items", "items") + assert raised.value.query == "GET /v2/items" + rest.assert_awaited_once() + assert rest.await_args.kwargs["session"].closed + + @pytest.mark.parametrize( + "pagination", + [ + {"hasNextPage": True}, + *[ + {"hasNextPage": True, "nextCursor": cursor} + for cursor in (None, "", 1, [], {}) + ], + ], + ) + async def test_missing_or_invalid_cursor_raises(self, pagination): + rest = AsyncMock(return_value={"items": [], "pagination": pagination}) + with patch("runpod.api.rest.run_rest_request_async", rest): + with pytest.raises(QueryError, match="nextCursor"): + await run_rest_paginated_request_async("/v2/items", "items") + rest.assert_awaited_once() + + @pytest.mark.parametrize("cursors", [["a", "a"], ["a", "b", "a"]]) + async def test_cursor_cycles_raise(self, cursors): + rest = AsyncMock( + side_effect=[ + { + "items": [{"id": cursor}], + "pagination": {"hasNextPage": True, "nextCursor": cursor}, + } + for cursor in cursors + ] + ) + with patch("runpod.api.rest.run_rest_request_async", rest): + with pytest.raises(QueryError, match="repeated a cursor"): + await run_rest_paginated_request_async("/v2/items", "items") + assert rest.await_count == len(cursors) + assert rest.await_args.kwargs["session"].closed diff --git a/tests/test_apps/test_api_client.py b/tests/test_apps/test_api_client.py index 0303e84f8..d3ea81e9a 100644 --- a/tests/test_apps/test_api_client.py +++ b/tests/test_apps/test_api_client.py @@ -1,6 +1,6 @@ """unit tests for the apps control-plane client.""" -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import ANY, AsyncMock, MagicMock, call, patch import aiohttp import pytest @@ -67,12 +67,190 @@ async def test_graphql_errors_propagate_immediately(self): async def test_mutation_transport_failure_is_not_retried(self, failure, operation): client = AppsApiClient(api_key="test-key") transport = AsyncMock(side_effect=failure) - with patch("runpod.apps.api.run_graphql_query_async", transport): + transport_name = ( + "run_rest_request_async" + if operation == "create_global_volume" + else "run_graphql_query_async" + ) + with patch(f"runpod.apps.api.{transport_name}", transport): with pytest.raises(type(failure)): await getattr(client, operation)("demo") assert transport.await_count == 1 +class TestGlobalVolumes: + @pytest.mark.parametrize("volumes", [[], [{"id": "gsb_abc", "name": "data"}]]) + async def test_list(self, volumes): + rest = AsyncMock( + return_value={ + "globalVolumes": volumes, + "pagination": {"hasNextPage": False, "nextCursor": None}, + } + ) + with ( + patch("runpod.api.rest.run_rest_request_async", rest), + patch("runpod.apps.api.run_graphql_query_async", AsyncMock()) as graphql, + ): + assert ( + await AppsApiClient(api_key="test-key").list_global_volumes() == volumes + ) + rest.assert_awaited_once_with( + "GET", "/v2/global-volumes", api_key="test-key", params={}, session=ANY + ) + graphql.assert_not_awaited() + + async def test_list_follows_all_cursors(self): + volumes = [ + {"id": f"gsb_{index}", "name": f"data-{index}"} for index in range(3) + ] + rest = AsyncMock( + side_effect=[ + { + "globalVolumes": [volume], + "pagination": { + "hasNextPage": index < 2, + "nextCursor": f"opaque-{index}" if index < 2 else None, + }, + } + for index, volume in enumerate(volumes) + ] + ) + with patch("runpod.api.rest.run_rest_request_async", rest): + assert ( + await AppsApiClient(api_key="test-key").list_global_volumes() == volumes + ) + assert rest.await_args_list == [ + call( + "GET", + "/v2/global-volumes", + api_key="test-key", + params=params, + session=ANY, + ) + for params in ({}, {"cursor": "opaque-0"}, {"cursor": "opaque-1"}) + ] + + @pytest.mark.parametrize("duplicate", [False, True]) + async def test_resolver_checks_names_across_pages(self, duplicate): + from runpod.apps.volume import GlobalVolume, VolumeError, VolumeResolver + + rest = AsyncMock( + side_effect=[ + { + "globalVolumes": [ + {"id": "gsb_first", "name": "data" if duplicate else "other"} + ], + "pagination": {"hasNextPage": True, "nextCursor": "next"}, + }, + { + "globalVolumes": [{"id": "gsb_second", "name": "data"}], + "pagination": {"hasNextPage": False, "nextCursor": None}, + }, + ] + ) + resolver = VolumeResolver(api=AppsApiClient(api_key="test-key")) + with patch("runpod.api.rest.run_rest_request_async", rest): + if duplicate: + with pytest.raises(VolumeError, match="multiple"): + await resolver.resolve(GlobalVolume("data", create=True)) + else: + assert await resolver.resolve(GlobalVolume("data", create=True)) == { + "id": "gsb_second" + } + assert [request.args[0] for request in rest.await_args_list] == ["GET", "GET"] + + async def test_create(self): + volume = { + "id": "gsb_abc", + "name": "training-data", + "createdAt": "2026-09-24T00:00:00Z", + "createdByUserId": "user_123", + } + rest = AsyncMock(return_value=volume) + with ( + patch("runpod.apps.api.run_rest_request_async", rest), + patch("runpod.apps.api.run_graphql_query_async", AsyncMock()) as graphql, + ): + assert ( + await AppsApiClient(api_key="test-key").create_global_volume( + "training-data" + ) + == volume + ) + rest.assert_awaited_once_with( + "POST", + "/v2/global-volumes", + api_key="test-key", + json={"name": "training-data"}, + ) + graphql.assert_not_awaited() + + @pytest.mark.parametrize( + "operation", ["list_global_volumes", "create_global_volume"] + ) + async def test_rest_error_propagates_without_graphql_fallback(self, operation): + failure = QueryError("denied", status_code=403) + rest = AsyncMock(side_effect=failure) + transport = ( + "runpod.api.rest.run_rest_request_async" + if operation == "list_global_volumes" + else "runpod.apps.api.run_rest_request_async" + ) + with ( + patch(transport, rest), + patch("runpod.apps.api.run_graphql_query_async", AsyncMock()) as graphql, + ): + args = ("data",) if operation == "create_global_volume" else () + with pytest.raises(QueryError) as exc: + await getattr(AppsApiClient(api_key="test-key"), operation)(*args) + assert exc.value is failure + assert rest.await_count == 1 + graphql.assert_not_awaited() + + @pytest.mark.parametrize( + "data", [None, [], {}, {"id": None}, {"id": ""}, {"id": 1}] + ) + async def test_create_rejects_invalid_response_without_retry(self, data): + rest = AsyncMock(return_value=data) + with patch("runpod.apps.api.run_rest_request_async", rest): + with pytest.raises(QueryError, match="may have been created") as raised: + await AppsApiClient(api_key="test-key").create_global_volume("data") + assert raised.value.query == "POST /v2/global-volumes" + rest.assert_awaited_once() + + +class TestNetworkVolumePagination: + async def test_list_follows_all_cursors(self): + volumes = [{"id": "first", "name": "other"}, {"id": "second", "name": "data"}] + rest = AsyncMock( + side_effect=[ + { + "networkVolumes": [volumes[0]], + "pagination": {"hasNextPage": True, "nextCursor": "next"}, + }, + { + "networkVolumes": [volumes[1]], + "pagination": {"hasNextPage": False, "nextCursor": None}, + }, + ] + ) + with patch("runpod.api.rest.run_rest_request_async", rest): + assert ( + await AppsApiClient(api_key="test-key").list_network_volumes() + == volumes + ) + assert rest.await_args_list == [ + call( + "GET", + "/v2/network-volumes", + api_key="test-key", + params=params, + session=ANY, + ) + for params in ({}, {"cursor": "next"}) + ] + + class TestEndpoints: async def test_dev_endpoint_preserves_runtime_and_placement(self, monkeypatch): from runpod.apps import App