refactor: use REST for global volume operations

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!) {

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) {

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}

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

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