From d48fa6218095a78bb8da20ea905c1515a9367563 Mon Sep 17 00:00:00 2001 From: Justin Date: Fri, 16 Oct 2026 19:56:27 -0400 Subject: [PATCH 1/6] fix(apps): pin requested gpu devices instead of their whole pool Endpoints select GPUs by pool, so a device list like ["NVIDIA GeForce RTX 4090", "NVIDIA L4", "NVIDIA RTX A5000"] deployed as ADA_24,AMPERE_24 and workers could land on any device in those pools, including an RTX PRO 6000 MIG 1g.24gb slice the platform added to AMPERE_24. Exclude every other device in each pool selected through a device name with "-" tokens, as runpod-flash did. Membership comes from the platform's live GPU catalog, because the static pool table lags it (it has no MIG slices); the static table is the fallback when the catalog is unreachable. Pools named directly still allow all devices. --- runpod/apps/gpu.py | 80 +++++++++++++++++++++++++++++++++---- tests/test_apps/conftest.py | 12 ++++++ tests/test_apps/test_app.py | 63 +++++++++++++++++++++++++++++ 3 files changed, 148 insertions(+), 7 deletions(-) create mode 100644 tests/test_apps/conftest.py diff --git a/runpod/apps/gpu.py b/runpod/apps/gpu.py index 79048b229..6a0d3edb8 100644 --- a/runpod/apps/gpu.py +++ b/runpod/apps/gpu.py @@ -5,8 +5,13 @@ configured. CpuInstanceType values are backend cpu flavor ids. """ +import functools +import logging +import os from enum import Enum -from typing import List, Optional, Union +from typing import Dict, List, Optional, Set, Union + +log = logging.getLogger(__name__) class GpuGroup(Enum): @@ -208,24 +213,85 @@ def resolve_gpu_string(value: str) -> List[str]: ) +CATALOG_URL = "https://api.runpod.io/v2/catalog/gpus" + + +@functools.lru_cache(maxsize=1) +def live_pool_members() -> Optional[Dict[str, Set[str]]]: + """pool id -> device ids, from the platform's gpu catalog. + + the static POOLS_TO_TYPES table lags the platform (pools gain mig + slices and new skus), so exact-device pinning reads the live + membership. returns None when the catalog is unreachable. + """ + import requests + + from .utils.network import api_key + + url = os.environ.get("RUNPOD_GPU_CATALOG_URL", CATALOG_URL) + try: + resp = requests.get( + url, headers={"Authorization": f"Bearer {api_key()}"}, timeout=10 + ) + resp.raise_for_status() + data = resp.json() + gpus = data.get("gpus", []) if isinstance(data, dict) else data + members: Dict[str, Set[str]] = {} + for gpu in gpus: + if gpu.get("pool") and gpu.get("id"): + members.setdefault(gpu["pool"], set()).add(gpu["id"]) + return members or None + except Exception as exc: # any failure falls back to the static table + log.debug("gpu catalog unavailable (%s); using the static pool table", exc) + return None + + +def _pool_members(pool: str) -> Set[str]: + live = live_pool_members() + if live is not None and pool in live: + return live[pool] + try: + return {t.value for t in POOLS_TO_TYPES.get(GpuGroup(pool), [])} + except ValueError: + return set() + + def gpu_ids_value(gpu: Optional[List[str]]) -> str: """the gpuIds string for an endpoint payload. "any gpu" (no selection, or the ANY sentinel) means every pool id: - the api has no wildcard and rejects "any". device names map back - to their pool (endpoints select by pool, not device). + the api has no wildcard and rejects "any". endpoints select by + pool, so device names map to their pool, and every other device in + that pool is excluded with a "-" token so workers only + land on the requested devices. pools named directly allow all of + their devices. """ if not gpu or any(str(g).lower() == "any" for g in gpu): return ",".join(g.value for g in GpuGroup.all()) pools: List[str] = [] + named_pools: Set[str] = set() + devices: Set[str] = set() + device_pools: List[str] = [] for entry in gpu: value = str(entry) try: pool = pool_for_gpu_type(GpuType(value)) - value = pool.value if pool is not None else value except ValueError: - # not a known GpuType: pass the raw pool/id string through - pass + # not a known GpuType: a pool id or raw id passes through + named_pools.add(value) + else: + if pool is not None: + devices.add(value) + if pool.value not in device_pools: + device_pools.append(pool.value) + value = pool.value if value not in pools: pools.append(value) - return ",".join(pools) + exclusions = sorted( + member + for pool in device_pools + if pool not in named_pools + for member in _pool_members(pool) + if member not in devices + ) + return ",".join(pools + [f"-{m}" for m in exclusions]) diff --git a/tests/test_apps/conftest.py b/tests/test_apps/conftest.py new file mode 100644 index 000000000..234b8edfd --- /dev/null +++ b/tests/test_apps/conftest.py @@ -0,0 +1,12 @@ +"""shared fixtures for apps tests.""" + +import pytest + +from runpod.apps import gpu as _gpu + + +@pytest.fixture(autouse=True) +def offline_gpu_catalog(monkeypatch): + """payload builders read the live gpu catalog; keep tests offline and + deterministic on the static pool table unless a test overrides it.""" + monkeypatch.setattr(_gpu, "live_pool_members", lambda: None) diff --git a/tests/test_apps/test_app.py b/tests/test_apps/test_app.py index e23614be0..c8a3ff180 100644 --- a/tests/test_apps/test_app.py +++ b/tests/test_apps/test_app.py @@ -1,11 +1,14 @@ """tests for the App registry and decorators.""" +from unittest.mock import patch + import pytest import runpod from runpod.apps import App, ApiHandle, FunctionHandle, ResourceKind from runpod.apps.app import _REGISTRY, _clear_registry from runpod.apps.errors import InvalidResourceError +from runpod.apps.gpu import live_pool_members as _real_live_pool_members from runpod.apps.markers import get, init, post @@ -382,6 +385,66 @@ def test_unknown_raises_at_decoration(self): with pytest.raises(InvalidResourceError, match="unknown gpu 'B300'"): normalize_gpu("B300") + def test_gpu_ids_excludes_other_live_pool_members(self, monkeypatch): + from runpod.apps import gpu as gpu_mod + + live = { + "ADA_24": {"NVIDIA GeForce RTX 4090"}, + "AMPERE_24": { + "NVIDIA GeForce RTX 3090", + "NVIDIA L4", + "NVIDIA RTX A5000", + "NVIDIA RTX PRO 6000 Blackwell Server Edition MIG 1g.24gb", + }, + } + monkeypatch.setattr(gpu_mod, "live_pool_members", lambda: live) + value = gpu_mod.gpu_ids_value( + ["NVIDIA GeForce RTX 4090", "NVIDIA L4", "NVIDIA RTX A5000"] + ) + assert value == ( + "ADA_24,AMPERE_24,-NVIDIA GeForce RTX 3090," + "-NVIDIA RTX PRO 6000 Blackwell Server Edition MIG 1g.24gb" + ) + + def test_gpu_ids_named_pool_allows_all_members(self, monkeypatch): + from runpod.apps import gpu as gpu_mod + + monkeypatch.setattr( + gpu_mod, + "live_pool_members", + lambda: {"AMPERE_24": {"NVIDIA L4", "NVIDIA GeForce RTX 3090"}}, + ) + assert gpu_mod.gpu_ids_value(["NVIDIA L4", "AMPERE_24"]) == "AMPERE_24" + + def test_gpu_ids_falls_back_to_static_pools(self, monkeypatch): + from runpod.apps import gpu as gpu_mod + + monkeypatch.setattr(gpu_mod, "live_pool_members", lambda: None) + assert gpu_mod.gpu_ids_value(["NVIDIA L4"]) == ( + "AMPERE_24,-NVIDIA GeForce RTX 3090,-NVIDIA RTX A5000" + ) + + def test_live_pool_members_parses_catalog_and_fails_soft(self, monkeypatch): + from unittest.mock import MagicMock + + monkeypatch.setenv("RUNPOD_API_KEY", "k") + resp = MagicMock() + resp.json.return_value = { + "gpus": [ + {"id": "NVIDIA L4", "pool": "AMPERE_24"}, + {"id": "NVIDIA A100-SXM4-40GB", "pool": None}, + ] + } + fetch = _real_live_pool_members + fetch.cache_clear() + with patch("requests.get", return_value=resp) as get: + assert fetch() == {"AMPERE_24": {"NVIDIA L4"}} + assert get.call_args.kwargs["headers"]["Authorization"] == "Bearer k" + fetch.cache_clear() + with patch("requests.get", side_effect=OSError("offline")): + assert fetch() is None + fetch.cache_clear() + def test_gpu_ids_maps_devices_to_pools(self): from runpod.apps.gpu import gpu_ids_value From 911ea98f0383e2edc8d9e5eb82ea231c1ecb2158 Mon Sep 17 00:00:00 2001 From: Justin Date: Mon, 19 Oct 2026 23:09:08 -0400 Subject: [PATCH 2/6] fix(apps): warn when gpu pins fall back to the static pool table a failed catalog lookup was logged at debug and cached for the process, so pins silently lost exclusions the static table lacks. warn instead, and retry the catalog after a minute. --- runpod/apps/gpu.py | 34 ++++++++++++++++++++++++++++------ tests/test_apps/test_app.py | 11 ++++++++++- 2 files changed, 38 insertions(+), 7 deletions(-) diff --git a/runpod/apps/gpu.py b/runpod/apps/gpu.py index 6a0d3edb8..7fa9363b1 100644 --- a/runpod/apps/gpu.py +++ b/runpod/apps/gpu.py @@ -5,11 +5,11 @@ configured. CpuInstanceType values are backend cpu flavor ids. """ -import functools import logging import os +import time from enum import Enum -from typing import Dict, List, Optional, Set, Union +from typing import Any, Dict, List, Optional, Set, Union log = logging.getLogger(__name__) @@ -216,18 +216,29 @@ def resolve_gpu_string(value: str) -> List[str]: CATALOG_URL = "https://api.runpod.io/v2/catalog/gpus" -@functools.lru_cache(maxsize=1) +# live catalog result, or when it last failed +_catalog: Dict[str, Any] = {} +_CATALOG_RETRY_SECONDS = 60.0 + + def live_pool_members() -> Optional[Dict[str, Set[str]]]: """pool id -> device ids, from the platform's gpu catalog. the static POOLS_TO_TYPES table lags the platform (pools gain mig slices and new skus), so exact-device pinning reads the live - membership. returns None when the catalog is unreachable. + membership. returns None when the catalog is unreachable. a result + is kept for the process; a failure is retried after a minute. """ import requests from .utils.network import api_key + if _catalog.get("members") is not None: + return _catalog["members"] + failed_at = _catalog.get("failed_at") + if failed_at is not None and time.monotonic() - failed_at < _CATALOG_RETRY_SECONDS: + return None + url = os.environ.get("RUNPOD_GPU_CATALOG_URL", CATALOG_URL) try: resp = requests.get( @@ -240,12 +251,23 @@ def live_pool_members() -> Optional[Dict[str, Set[str]]]: for gpu in gpus: if gpu.get("pool") and gpu.get("id"): members.setdefault(gpu["pool"], set()).add(gpu["id"]) - return members or None + if not members: + raise ValueError("catalog lists no pools") + _catalog["members"] = members + return members except Exception as exc: # any failure falls back to the static table - log.debug("gpu catalog unavailable (%s); using the static pool table", exc) + _catalog["failed_at"] = time.monotonic() + log.warning( + "gpu catalog unavailable (%s); device pins use the static pool " + "table, which can miss newer devices in a pool", + exc, + ) return None +live_pool_members.cache_clear = _catalog.clear # type: ignore[attr-defined] + + def _pool_members(pool: str) -> Set[str]: live = live_pool_members() if live is not None and pool in live: diff --git a/tests/test_apps/test_app.py b/tests/test_apps/test_app.py index c8a3ff180..52bd4c1b6 100644 --- a/tests/test_apps/test_app.py +++ b/tests/test_apps/test_app.py @@ -427,6 +427,8 @@ def test_gpu_ids_falls_back_to_static_pools(self, monkeypatch): def test_live_pool_members_parses_catalog_and_fails_soft(self, monkeypatch): from unittest.mock import MagicMock + from runpod.apps import gpu as gpu_mod + monkeypatch.setenv("RUNPOD_API_KEY", "k") resp = MagicMock() resp.json.return_value = { @@ -441,8 +443,15 @@ def test_live_pool_members_parses_catalog_and_fails_soft(self, monkeypatch): assert fetch() == {"AMPERE_24": {"NVIDIA L4"}} assert get.call_args.kwargs["headers"]["Authorization"] == "Bearer k" fetch.cache_clear() - with patch("requests.get", side_effect=OSError("offline")): + with patch("requests.get", side_effect=OSError("offline")) as get: assert fetch() is None + # a failure is not retried on every lookup + assert fetch() is None + assert get.call_count == 1 + # it is retried once the window passes + gpu_mod._catalog["failed_at"] -= gpu_mod._CATALOG_RETRY_SECONDS + with patch("requests.get", return_value=resp): + assert fetch() == {"AMPERE_24": {"NVIDIA L4"}} fetch.cache_clear() def test_gpu_ids_maps_devices_to_pools(self): From bffbd2897cf6265942c7a5d39ce6df17c3504a8b Mon Sep 17 00:00:00 2001 From: Justin Date: Tue, 20 Oct 2026 12:28:31 -0400 Subject: [PATCH 3/6] refactor(apps): typed gpu catalog cache and a flatter gpu_ids_value --- runpod/apps/gpu.py | 87 +++++++++++++++++++++++-------------- tests/test_apps/test_app.py | 8 ++-- 2 files changed, 58 insertions(+), 37 deletions(-) diff --git a/runpod/apps/gpu.py b/runpod/apps/gpu.py index 7fa9363b1..6e340a9fb 100644 --- a/runpod/apps/gpu.py +++ b/runpod/apps/gpu.py @@ -8,6 +8,7 @@ import logging import os import time +from dataclasses import dataclass from enum import Enum from typing import Any, Dict, List, Optional, Set, Union @@ -214,11 +215,35 @@ def resolve_gpu_string(value: str) -> List[str]: CATALOG_URL = "https://api.runpod.io/v2/catalog/gpus" +CATALOG_RETRY_SECONDS = 60.0 -# live catalog result, or when it last failed -_catalog: Dict[str, Any] = {} -_CATALOG_RETRY_SECONDS = 60.0 +@dataclass +class _CatalogCache: + """the live catalog's pool membership, or when fetching it last failed.""" + + members: Optional[Dict[str, Set[str]]] = None + failed_at: Optional[float] = None + + +_catalog_cache = _CatalogCache() + + +def _clear_catalog_cache() -> None: + global _catalog_cache + _catalog_cache = _CatalogCache() + + +def _parse_catalog(data: Any) -> Dict[str, Set[str]]: + """pool id -> device ids from a catalog response.""" + gpus = data.get("gpus", []) if isinstance(data, dict) else data + members: Dict[str, Set[str]] = {} + for gpu in gpus: + if gpu.get("pool") and gpu.get("id"): + members.setdefault(gpu["pool"], set()).add(gpu["id"]) + if not members: + raise ValueError("catalog lists no pools") + return members def live_pool_members() -> Optional[Dict[str, Set[str]]]: @@ -233,10 +258,13 @@ def live_pool_members() -> Optional[Dict[str, Set[str]]]: from .utils.network import api_key - if _catalog.get("members") is not None: - return _catalog["members"] - failed_at = _catalog.get("failed_at") - if failed_at is not None and time.monotonic() - failed_at < _CATALOG_RETRY_SECONDS: + cache = _catalog_cache + if cache.members is not None: + return cache.members + if ( + cache.failed_at is not None + and time.monotonic() - cache.failed_at < CATALOG_RETRY_SECONDS + ): return None url = os.environ.get("RUNPOD_GPU_CATALOG_URL", CATALOG_URL) @@ -245,27 +273,16 @@ def live_pool_members() -> Optional[Dict[str, Set[str]]]: url, headers={"Authorization": f"Bearer {api_key()}"}, timeout=10 ) resp.raise_for_status() - data = resp.json() - gpus = data.get("gpus", []) if isinstance(data, dict) else data - members: Dict[str, Set[str]] = {} - for gpu in gpus: - if gpu.get("pool") and gpu.get("id"): - members.setdefault(gpu["pool"], set()).add(gpu["id"]) - if not members: - raise ValueError("catalog lists no pools") - _catalog["members"] = members - return members - except Exception as exc: # any failure falls back to the static table - _catalog["failed_at"] = time.monotonic() + cache.members = _parse_catalog(resp.json()) + except Exception as exc: # noqa: BLE001 - offline, auth, or malformed + cache.failed_at = time.monotonic() log.warning( "gpu catalog unavailable (%s); device pins use the static pool " "table, which can miss newer devices in a pool", exc, ) return None - - -live_pool_members.cache_clear = _catalog.clear # type: ignore[attr-defined] + return cache.members def _pool_members(pool: str) -> Set[str]: @@ -278,6 +295,15 @@ def _pool_members(pool: str) -> Set[str]: return set() +def _device_pool(value: str) -> Optional[str]: + """the pool id of a known device name, or None.""" + try: + pool = pool_for_gpu_type(GpuType(value)) + except ValueError: + return None + return pool.value if pool is not None else None + + def gpu_ids_value(gpu: Optional[List[str]]) -> str: """the gpuIds string for an endpoint payload. @@ -293,26 +319,21 @@ def gpu_ids_value(gpu: Optional[List[str]]) -> str: pools: List[str] = [] named_pools: Set[str] = set() devices: Set[str] = set() - device_pools: List[str] = [] for entry in gpu: value = str(entry) - try: - pool = pool_for_gpu_type(GpuType(value)) - except ValueError: - # not a known GpuType: a pool id or raw id passes through + pool = _device_pool(value) + if pool is None: + # a pool id allows every device in it; other ids pass through named_pools.add(value) else: - if pool is not None: - devices.add(value) - if pool.value not in device_pools: - device_pools.append(pool.value) - value = pool.value + devices.add(value) + value = pool if value not in pools: pools.append(value) + device_pools = {_device_pool(d) for d in devices} - named_pools exclusions = sorted( member for pool in device_pools - if pool not in named_pools for member in _pool_members(pool) if member not in devices ) diff --git a/tests/test_apps/test_app.py b/tests/test_apps/test_app.py index 52bd4c1b6..d42c61fc0 100644 --- a/tests/test_apps/test_app.py +++ b/tests/test_apps/test_app.py @@ -438,21 +438,21 @@ def test_live_pool_members_parses_catalog_and_fails_soft(self, monkeypatch): ] } fetch = _real_live_pool_members - fetch.cache_clear() + gpu_mod._clear_catalog_cache() with patch("requests.get", return_value=resp) as get: assert fetch() == {"AMPERE_24": {"NVIDIA L4"}} assert get.call_args.kwargs["headers"]["Authorization"] == "Bearer k" - fetch.cache_clear() + gpu_mod._clear_catalog_cache() with patch("requests.get", side_effect=OSError("offline")) as get: assert fetch() is None # a failure is not retried on every lookup assert fetch() is None assert get.call_count == 1 # it is retried once the window passes - gpu_mod._catalog["failed_at"] -= gpu_mod._CATALOG_RETRY_SECONDS + gpu_mod._catalog_cache.failed_at -= gpu_mod.CATALOG_RETRY_SECONDS with patch("requests.get", return_value=resp): assert fetch() == {"AMPERE_24": {"NVIDIA L4"}} - fetch.cache_clear() + gpu_mod._clear_catalog_cache() def test_gpu_ids_maps_devices_to_pools(self): from runpod.apps.gpu import gpu_ids_value From 41df93a4c4cd1648b8da55d511d02ebc8d5382b1 Mon Sep 17 00:00:00 2001 From: Justin Date: Tue, 20 Oct 2026 13:00:22 -0400 Subject: [PATCH 4/6] refactor(apps): separate catalog fetch from its cache --- runpod/apps/gpu.py | 59 +++++++++++++++++++++++++--------------------- 1 file changed, 32 insertions(+), 27 deletions(-) diff --git a/runpod/apps/gpu.py b/runpod/apps/gpu.py index 6e340a9fb..c91c23a82 100644 --- a/runpod/apps/gpu.py +++ b/runpod/apps/gpu.py @@ -225,6 +225,15 @@ class _CatalogCache: members: Optional[Dict[str, Set[str]]] = None failed_at: Optional[float] = None + def may_fetch(self) -> bool: + """true until a fetch succeeds, except within a minute of a failure.""" + if self.members is not None: + return False + return ( + self.failed_at is None + or time.monotonic() - self.failed_at >= CATALOG_RETRY_SECONDS + ) + _catalog_cache = _CatalogCache() @@ -246,6 +255,19 @@ def _parse_catalog(data: Any) -> Dict[str, Set[str]]: return members +def _fetch_catalog() -> Dict[str, Set[str]]: + import requests + + from .utils.network import api_key + + url = os.environ.get("RUNPOD_GPU_CATALOG_URL", CATALOG_URL) + resp = requests.get( + url, headers={"Authorization": f"Bearer {api_key()}"}, timeout=10 + ) + resp.raise_for_status() + return _parse_catalog(resp.json()) + + def live_pool_members() -> Optional[Dict[str, Set[str]]]: """pool id -> device ids, from the platform's gpu catalog. @@ -254,34 +276,17 @@ def live_pool_members() -> Optional[Dict[str, Set[str]]]: membership. returns None when the catalog is unreachable. a result is kept for the process; a failure is retried after a minute. """ - import requests - - from .utils.network import api_key - cache = _catalog_cache - if cache.members is not None: - return cache.members - if ( - cache.failed_at is not None - and time.monotonic() - cache.failed_at < CATALOG_RETRY_SECONDS - ): - return None - - url = os.environ.get("RUNPOD_GPU_CATALOG_URL", CATALOG_URL) - try: - resp = requests.get( - url, headers={"Authorization": f"Bearer {api_key()}"}, timeout=10 - ) - resp.raise_for_status() - cache.members = _parse_catalog(resp.json()) - except Exception as exc: # noqa: BLE001 - offline, auth, or malformed - cache.failed_at = time.monotonic() - log.warning( - "gpu catalog unavailable (%s); device pins use the static pool " - "table, which can miss newer devices in a pool", - exc, - ) - return None + if cache.may_fetch(): + try: + cache.members = _fetch_catalog() + except Exception as exc: # noqa: BLE001 - offline, auth, or malformed + cache.failed_at = time.monotonic() + log.warning( + "gpu catalog unavailable (%s); device pins use the static pool " + "table, which can miss newer devices in a pool", + exc, + ) return cache.members From 53a50909a909286a2846d3c1c9853fe6022407af Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Tue, 20 Oct 2026 13:07:49 -0700 Subject: [PATCH 5/6] fix(apps): fetch gpu catalogs asynchronously during provisioning --- runpod/apps/api.py | 5 + runpod/apps/deploy.py | 10 +- runpod/apps/dev.py | 29 +++- runpod/apps/gpu.py | 100 +++++++++--- tests/test_apps/conftest.py | 3 + tests/test_apps/test_app.py | 3 +- tests/test_apps/test_gpu_catalog.py | 227 ++++++++++++++++++++++++++++ 7 files changed, 346 insertions(+), 31 deletions(-) create mode 100644 tests/test_apps/test_gpu_catalog.py diff --git a/runpod/apps/api.py b/runpod/apps/api.py index 9b557800a..c518e8b01 100644 --- a/runpod/apps/api.py +++ b/runpod/apps/api.py @@ -164,6 +164,11 @@ class AppsApiClient: def __init__(self, api_key: Optional[str] = None): self._api_key = api_key + async def gpu_pool_members(self) -> Optional[Dict[str, Set[str]]]: + from .gpu import live_pool_members_async + + return await live_pool_members_async(self._api_key) + async def _execute( self, query: str, diff --git a/runpod/apps/deploy.py b/runpod/apps/deploy.py index a0ef47ca7..bddd63615 100644 --- a/runpod/apps/deploy.py +++ b/runpod/apps/deploy.py @@ -7,7 +7,7 @@ import tempfile from dataclasses import dataclass, field from pathlib import Path -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Set, Tuple from pathspec import GitIgnoreSpec @@ -285,6 +285,7 @@ def _deployed_endpoint_input( python_version: str, *, env_name: Optional[str] = None, + pool_members: Optional[Dict[str, Set[str]]] = None, ) -> Dict[str, Any]: """saveEndpoint payload for one deployed queue/api resource. @@ -372,7 +373,7 @@ def _deployed_endpoint_input( else: from .gpu import gpu_ids_value - payload["gpuIds"] = gpu_ids_value(spec.gpu) + payload["gpuIds"] = gpu_ids_value(spec.gpu, pool_members=pool_members) payload["gpuCount"] = spec.gpu_count if spec.min_cuda_version: payload["minCudaVersion"] = spec.min_cuda_version @@ -406,12 +407,16 @@ async def reconcile_endpoints( if h.spec.kind in (ResourceKind.QUEUE, ResourceKind.API) } + from .gpu import needs_gpu_catalog from .registry import resolve_registry_auth from .volume import VolumeResolver, attach_endpoint_volumes resolver = VolumeResolver(client) endpoints: Dict[str, str] = {} for name, handle in sorted(provisionable.items()): + members = None + if needs_gpu_catalog(handle.spec.gpu): + members = await client.gpu_pool_members() payload = _deployed_endpoint_input( app, handle.spec, @@ -419,6 +424,7 @@ async def reconcile_endpoints( build_id, python_version, env_name=environment["name"], + pool_members=members or {}, ) await attach_endpoint_volumes(payload, handle.spec, resolver, app) auth_id = await resolve_registry_auth(handle.spec.registry_auth, api=client) diff --git a/runpod/apps/dev.py b/runpod/apps/dev.py index c4e95ad1c..be7b951a2 100644 --- a/runpod/apps/dev.py +++ b/runpod/apps/dev.py @@ -8,11 +8,12 @@ import logging import os as _os -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional, Set, Union from .api import AppsApiClient from .app import App from .errors import AppError +from .gpu import needs_gpu_catalog from .handles import ApiHandle, FunctionHandle from .shim import runtime_launcher from .spec import ResourceKind, ResourceSpec @@ -107,7 +108,13 @@ def _cpu_locations(instance_ids: List[str]) -> str: GENERATION_ENV = "RUNPOD_DEV_GENERATION" -def _endpoint_input(app: App, spec: ResourceSpec, generation: int = 1) -> Dict: +def _endpoint_input( + app: App, + spec: ResourceSpec, + generation: int = 1, + *, + pool_members: Optional[Dict[str, Set[str]]] = None, +) -> Dict: """build a live endpoint payload with an atomically bound container config.""" spec.validate() @@ -179,7 +186,7 @@ def _endpoint_input(app: App, spec: ResourceSpec, generation: int = 1) -> Dict: else: from .gpu import gpu_ids_value - payload["gpuIds"] = gpu_ids_value(spec.gpu) + payload["gpuIds"] = gpu_ids_value(spec.gpu, pool_members=pool_members) payload["gpuCount"] = spec.gpu_count if spec.min_cuda_version: payload["minCudaVersion"] = spec.min_cuda_version @@ -274,7 +281,13 @@ async def start(self) -> None: if found: self._endpoints[name] = found["id"] self._resource_names[name] = spec.name - payload = _endpoint_input(app, spec, self.generation) + spec.validate() + members = None + if needs_gpu_catalog(spec.gpu): + members = await self.api.gpu_pool_members() + payload = _endpoint_input( + app, spec, self.generation, pool_members=members or {} + ) await self._attach_volumes(payload, spec, app) hardware = ",".join(spec.cpu or spec.gpu or ["any"]) @@ -351,7 +364,13 @@ async def refresh(self, apps: List[App]) -> None: # update survivors (config + generation bump) and create additions for name, (app, handle) in desired.items(): - payload = _endpoint_input(app, handle.spec, self.generation) + handle.spec.validate() + members = None + if needs_gpu_catalog(handle.spec.gpu): + members = await self.api.gpu_pool_members() + payload = _endpoint_input( + app, handle.spec, self.generation, pool_members=members or {} + ) await self._attach_volumes(payload, handle.spec, app) existing_id = self._endpoints.get(name) diff --git a/runpod/apps/gpu.py b/runpod/apps/gpu.py index c91c23a82..b3dd35b44 100644 --- a/runpod/apps/gpu.py +++ b/runpod/apps/gpu.py @@ -10,7 +10,7 @@ import time from dataclasses import dataclass from enum import Enum -from typing import Any, Dict, List, Optional, Set, Union +from typing import Any, Dict, List, Optional, Set, Tuple, Union log = logging.getLogger(__name__) @@ -234,13 +234,40 @@ def may_fetch(self) -> bool: or time.monotonic() - self.failed_at >= CATALOG_RETRY_SECONDS ) + def fail(self, exc: Exception) -> None: + self.failed_at = time.monotonic() + log.warning( + "gpu catalog unavailable (%s); device pins use the static pool " + "table, which can miss newer devices in a pool", + exc, + ) + -_catalog_cache = _CatalogCache() +_catalog_caches: Dict[Tuple[str, str], _CatalogCache] = {} def _clear_catalog_cache() -> None: - global _catalog_cache - _catalog_cache = _CatalogCache() + _catalog_caches.clear() + + +def _catalog_context(api_key: Optional[str] = None) -> Tuple[str, str]: + import runpod + + base = os.environ.get("RUNPOD_API_BASE_URL") + url = os.environ.get( + "RUNPOD_GPU_CATALOG_URL", + f"{base.rstrip('/')}/v2/catalog/gpus" if base else CATALOG_URL, + ) + key = api_key or os.environ.get("RUNPOD_API_KEY") or runpod.api_key or "" + return url, key + + +def _catalog_headers(key: str) -> Dict[str, str]: + if not key: + raise RuntimeError( + "no api key configured. run `rp login` or set RUNPOD_API_KEY." + ) + return {"Authorization": f"Bearer {key}"} def _parse_catalog(data: Any) -> Dict[str, Set[str]]: @@ -255,15 +282,10 @@ def _parse_catalog(data: Any) -> Dict[str, Set[str]]: return members -def _fetch_catalog() -> Dict[str, Set[str]]: +def _fetch_catalog(url: str, key: str) -> Dict[str, Set[str]]: import requests - from .utils.network import api_key - - url = os.environ.get("RUNPOD_GPU_CATALOG_URL", CATALOG_URL) - resp = requests.get( - url, headers={"Authorization": f"Bearer {api_key()}"}, timeout=10 - ) + resp = requests.get(url, headers=_catalog_headers(key), timeout=10) resp.raise_for_status() return _parse_catalog(resp.json()) @@ -276,22 +298,39 @@ def live_pool_members() -> Optional[Dict[str, Set[str]]]: membership. returns None when the catalog is unreachable. a result is kept for the process; a failure is retried after a minute. """ - cache = _catalog_cache + context = _catalog_context() + cache = _catalog_caches.setdefault(context, _CatalogCache()) if cache.may_fetch(): try: - cache.members = _fetch_catalog() + cache.members = _fetch_catalog(*context) except Exception as exc: # noqa: BLE001 - offline, auth, or malformed - cache.failed_at = time.monotonic() - log.warning( - "gpu catalog unavailable (%s); device pins use the static pool " - "table, which can miss newer devices in a pool", - exc, - ) + cache.fail(exc) return cache.members -def _pool_members(pool: str) -> Set[str]: - live = live_pool_members() +async def live_pool_members_async( + api_key: Optional[str] = None, +) -> Optional[Dict[str, Set[str]]]: + """fetch pool membership without blocking, cached by catalog url and credentials.""" + import aiohttp + + context = _catalog_context(api_key) + cache = _catalog_caches.setdefault(context, _CatalogCache()) + if cache.may_fetch(): + url, key = context + try: + headers = _catalog_headers(key) + timeout = aiohttp.ClientTimeout(total=10) + async with aiohttp.ClientSession(timeout=timeout) as session: + async with session.get(url, headers=headers) as response: + response.raise_for_status() + cache.members = _parse_catalog(await response.json()) + except Exception as exc: + cache.fail(exc) + return cache.members + + +def _pool_members(pool: str, live: Optional[Dict[str, Set[str]]]) -> Set[str]: if live is not None and pool in live: return live[pool] try: @@ -309,7 +348,17 @@ def _device_pool(value: str) -> Optional[str]: return pool.value if pool is not None else None -def gpu_ids_value(gpu: Optional[List[str]]) -> str: +def needs_gpu_catalog(gpu: Optional[List[str]]) -> bool: + if not gpu or any(str(g).lower() == "any" for g in gpu): + return False + return any( + (pool := _device_pool(str(g))) is not None and pool not in gpu for g in gpu + ) + + +def gpu_ids_value( + gpu: Optional[List[str]], *, pool_members: Optional[Dict[str, Set[str]]] = None +) -> str: """the gpuIds string for an endpoint payload. "any gpu" (no selection, or the ANY sentinel) means every pool id: @@ -318,6 +367,9 @@ def gpu_ids_value(gpu: Optional[List[str]]) -> str: that pool is excluded with a "-" token so workers only land on the requested devices. pools named directly allow all of their devices. + + supplying pool_members keeps serialization offline; an empty mapping + uses the static pool table. omitted membership is fetched synchronously. """ if not gpu or any(str(g).lower() == "any" for g in gpu): return ",".join(g.value for g in GpuGroup.all()) @@ -336,10 +388,12 @@ def gpu_ids_value(gpu: Optional[List[str]]) -> str: if value not in pools: pools.append(value) device_pools = {_device_pool(d) for d in devices} - named_pools + if pool_members is None and device_pools: + pool_members = live_pool_members() exclusions = sorted( member for pool in device_pools - for member in _pool_members(pool) + for member in _pool_members(pool, pool_members) if member not in devices ) return ",".join(pools + [f"-{m}" for m in exclusions]) diff --git a/tests/test_apps/conftest.py b/tests/test_apps/conftest.py index 234b8edfd..d95469fb3 100644 --- a/tests/test_apps/conftest.py +++ b/tests/test_apps/conftest.py @@ -1,5 +1,7 @@ """shared fixtures for apps tests.""" +from unittest.mock import AsyncMock + import pytest from runpod.apps import gpu as _gpu @@ -10,3 +12,4 @@ def offline_gpu_catalog(monkeypatch): """payload builders read the live gpu catalog; keep tests offline and deterministic on the static pool table unless a test overrides it.""" monkeypatch.setattr(_gpu, "live_pool_members", lambda: None) + monkeypatch.setattr(_gpu, "live_pool_members_async", AsyncMock(return_value=None)) diff --git a/tests/test_apps/test_app.py b/tests/test_apps/test_app.py index d42c61fc0..ed72bec0b 100644 --- a/tests/test_apps/test_app.py +++ b/tests/test_apps/test_app.py @@ -449,7 +449,8 @@ def test_live_pool_members_parses_catalog_and_fails_soft(self, monkeypatch): assert fetch() is None assert get.call_count == 1 # it is retried once the window passes - gpu_mod._catalog_cache.failed_at -= gpu_mod.CATALOG_RETRY_SECONDS + cache = next(iter(gpu_mod._catalog_caches.values())) + cache.failed_at -= gpu_mod.CATALOG_RETRY_SECONDS with patch("requests.get", return_value=resp): assert fetch() == {"AMPERE_24": {"NVIDIA L4"}} gpu_mod._clear_catalog_cache() diff --git a/tests/test_apps/test_gpu_catalog.py b/tests/test_apps/test_gpu_catalog.py new file mode 100644 index 000000000..69f2a040f --- /dev/null +++ b/tests/test_apps/test_gpu_catalog.py @@ -0,0 +1,227 @@ +"""gpu catalog fetching across async provisioning and synchronous serialization.""" + +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest +from aiohttp import web + +import runpod +from runpod.apps import App, gpu +from runpod.apps.api import AppsApiClient +from runpod.apps.app import _clear_registry +from runpod.apps.deploy import reconcile_endpoints +from runpod.apps.dev import DevSession +from runpod.apps.gpu import live_pool_members_async as _real_live_pool_members_async + + +MIG = "NVIDIA RTX PRO 6000 Blackwell Server Edition MIG 1g.24gb" +CATALOG = { + "gpus": [ + {"id": name, "pool": "AMPERE_24"} + for name in ("NVIDIA L4", "NVIDIA GeForce RTX 3090", "NVIDIA RTX A5000", MIG) + ] +} + + +@pytest.fixture(autouse=True) +def clean_catalog(monkeypatch, offline_gpu_catalog): + _clear_registry() + gpu._clear_catalog_cache() + monkeypatch.delenv("RUNPOD_API_KEY", raising=False) + monkeypatch.delenv("RUNPOD_GPU_CATALOG_URL", raising=False) + monkeypatch.setattr(runpod, "api_key", None) + monkeypatch.setattr(gpu, "live_pool_members_async", _real_live_pool_members_async) + monkeypatch.setattr( + gpu, "live_pool_members", Mock(side_effect=AssertionError("synchronous fetch")) + ) + yield + gpu._clear_catalog_cache() + _clear_registry() + + +@pytest.fixture +async def catalog(monkeypatch): + state = SimpleNamespace( + requests=[], + started=asyncio.Event(), + release=asyncio.Event(), + status=200, + data=CATALOG, + ) + state.release.set() + + async def handle(request): + state.requests.append((request.path, request.headers.get("Authorization"))) + state.started.set() + await state.release.wait() + return web.json_response(state.data, status=state.status) + + app = web.Application() + app.router.add_get("/v2/catalog/gpus", handle) + app.router.add_get("/custom-catalog", handle) + runner = web.AppRunner(app, shutdown_timeout=0.1) + await runner.setup() + try: + site = web.TCPSite(runner, "127.0.0.1", 0) + await site.start() + state.url = f"http://127.0.0.1:{site._server.sockets[0].getsockname()[1]}" + other = web.TCPSite(runner, "127.0.0.1", 0) + await other.start() + other_port = other._server.sockets[0].getsockname()[1] + state.other_url = f"http://127.0.0.1:{other_port}" + monkeypatch.setenv("RUNPOD_API_BASE_URL", state.url) + yield state + finally: + state.release.set() + await runner.cleanup() + + +async def provision(operation, client, gpu_selection="L4", cpu=None): + app = App("catalog-test") + + @app.queue(name="work", gpu=gpu_selection, cpu=cpu) + def work(value): + return value + + client.save_endpoint = AsyncMock(return_value={"id": "test-endpoint"}) + client.list_my_endpoints = AsyncMock(return_value=[]) + if operation == "reconcile": + await reconcile_endpoints( + client, app, {"id": "env", "name": "test"}, "build", "3.12" + ) + else: + session = DevSession([app], api=client) + session._attach_volumes = AsyncMock() + if operation == "start": + await session.start() + else: + await session.refresh([app]) + return client.save_endpoint.call_args.args[0] + + +@pytest.mark.parametrize("operation", ["start", "refresh", "reconcile"]) +@pytest.mark.parametrize("global_key", [None, "unrelated-global-key"]) +async def test_provisioning_uses_explicit_credentials( + catalog, monkeypatch, operation, global_key +): + monkeypatch.setattr(runpod, "api_key", global_key) + if global_key: + monkeypatch.setenv("RUNPOD_API_KEY", "unrelated-env-key") + client = AppsApiClient(api_key="explicit-key") + payload = await provision(operation, client) + assert f"-{MIG}" in payload["gpuIds"] + assert catalog.requests == [("/v2/catalog/gpus", "Bearer explicit-key")] + + +@pytest.mark.parametrize("operation", ["start", "refresh", "reconcile"]) +async def test_catalog_wait_keeps_event_loop_responsive(catalog, operation): + catalog.release.clear() + client = AppsApiClient(api_key="explicit-key") + task = asyncio.create_task(provision(operation, client)) + try: + await asyncio.wait_for(catalog.started.wait(), timeout=1) + heartbeat = asyncio.Event() + asyncio.get_running_loop().call_later(0.01, heartbeat.set) + await asyncio.wait_for(heartbeat.wait(), timeout=1) + assert not task.done() + catalog.release.set() + payload = await asyncio.wait_for(task, timeout=1) + assert f"-{MIG}" in payload["gpuIds"] + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + +@pytest.mark.parametrize("operation", ["start", "refresh", "reconcile"]) +async def test_catalog_wait_respects_provisioning_deadline(catalog, operation): + catalog.release.clear() + client = AppsApiClient(api_key="explicit-key") + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(provision(operation, client), timeout=0.1) + assert catalog.started.is_set() + client.save_endpoint.assert_not_called() + cache = next(iter(gpu._catalog_caches.values())) + assert cache.failed_at is None + assert cache.members is None + catalog.release.set() + assert MIG in (await client.gpu_pool_members())["AMPERE_24"] + + +@pytest.mark.parametrize("operation", ["start", "refresh", "reconcile"]) +async def test_failed_catalog_uses_static_payload_without_sync_fetch( + catalog, operation +): + catalog.status = 500 + payload = await provision(operation, AppsApiClient(api_key="explicit-key")) + assert payload["gpuIds"] == ( + "AMPERE_24,-NVIDIA GeForce RTX 3090,-NVIDIA RTX A5000" + ) + assert len(catalog.requests) == 1 + + +@pytest.mark.parametrize("malformed", [False, True]) +async def test_failed_catalog_retries_after_cooldown(catalog, malformed, caplog): + if malformed: + catalog.data = {"gpus": []} + else: + catalog.status = 500 + client = AppsApiClient(api_key="explicit-key") + assert await client.gpu_pool_members() is None + assert await client.gpu_pool_members() is None + assert len(catalog.requests) == 1 + assert "device pins use the static pool table" in caplog.text + assert "explicit-key" not in caplog.text + cache = next(iter(gpu._catalog_caches.values())) + cache.failed_at -= gpu.CATALOG_RETRY_SECONDS + catalog.status = 200 + catalog.data = CATALOG + assert MIG in (await client.gpu_pool_members())["AMPERE_24"] + assert len(catalog.requests) == 2 + another_client = AppsApiClient(api_key="explicit-key") + assert MIG in (await another_client.gpu_pool_members())["AMPERE_24"] + assert len(catalog.requests) == 2 + + +async def test_cache_is_scoped_to_credentials_and_catalog_url(catalog, monkeypatch): + first = AppsApiClient(api_key="first-key") + assert MIG in (await first.gpu_pool_members())["AMPERE_24"] + catalog.data = {"gpus": [{"id": "NVIDIA L4", "pool": "AMPERE_24"}]} + second = AppsApiClient(api_key="second-key") + assert await second.gpu_pool_members() == {"AMPERE_24": {"NVIDIA L4"}} + assert MIG in (await first.gpu_pool_members())["AMPERE_24"] + assert len(catalog.requests) == 2 + monkeypatch.setenv("RUNPOD_API_BASE_URL", catalog.other_url) + assert await first.gpu_pool_members() == {"AMPERE_24": {"NVIDIA L4"}} + assert len(catalog.requests) == 3 + monkeypatch.setenv("RUNPOD_GPU_CATALOG_URL", f"{catalog.url}/custom-catalog") + assert await first.gpu_pool_members() == {"AMPERE_24": {"NVIDIA L4"}} + assert catalog.requests[-1] == ("/custom-catalog", "Bearer first-key") + assert len(catalog.requests) == 4 + + +@pytest.mark.parametrize("operation", ["start", "refresh", "reconcile"]) +@pytest.mark.parametrize("selection", [None, "any", "AMPERE_24", ["L4", "AMPERE_24"]]) +async def test_unpinned_gpus_do_not_fetch_catalog(operation, selection): + client = AppsApiClient(api_key="explicit-key") + client.gpu_pool_members = AsyncMock(side_effect=AssertionError("unneeded catalog")) + await provision(operation, client, gpu_selection=selection) + client.gpu_pool_members.assert_not_called() + + +@pytest.mark.parametrize("operation", ["start", "refresh", "reconcile"]) +async def test_cpu_resources_do_not_fetch_catalog(operation): + client = AppsApiClient(api_key="explicit-key") + client.gpu_pool_members = AsyncMock(side_effect=AssertionError("unneeded catalog")) + await provision(operation, client, gpu_selection=None, cpu="cpu3c-1-2") + client.gpu_pool_members.assert_not_called() + + +def test_supplied_membership_serializes_without_network(): + assert f"-{MIG}" in gpu.gpu_ids_value( + ["NVIDIA L4"], pool_members=gpu._parse_catalog(CATALOG) + ) + assert gpu.gpu_ids_value(["NVIDIA L4"], pool_members={}) == ( + "AMPERE_24,-NVIDIA GeForce RTX 3090,-NVIDIA RTX A5000" + ) From d96f49f551c24997f79b8904327bf92d73bf8c33 Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Tue, 20 Oct 2026 14:06:59 -0700 Subject: [PATCH 6/6] refactor(apps): reuse REST clients for GPU catalogs --- runpod/apps/gpu.py | 51 +++++++++-------------------- tests/test_apps/test_app.py | 9 ++--- tests/test_apps/test_gpu_catalog.py | 27 ++++++++++++--- 3 files changed, 44 insertions(+), 43 deletions(-) diff --git a/runpod/apps/gpu.py b/runpod/apps/gpu.py index b3dd35b44..4b7ba72fc 100644 --- a/runpod/apps/gpu.py +++ b/runpod/apps/gpu.py @@ -6,7 +6,6 @@ """ import logging -import os import time from dataclasses import dataclass from enum import Enum @@ -214,7 +213,7 @@ def resolve_gpu_string(value: str) -> List[str]: ) -CATALOG_URL = "https://api.runpod.io/v2/catalog/gpus" +CATALOG_PATH = "/v2/catalog/gpus" CATALOG_RETRY_SECONDS = 60.0 @@ -251,23 +250,14 @@ def _clear_catalog_cache() -> None: def _catalog_context(api_key: Optional[str] = None) -> Tuple[str, str]: - import runpod + from ..api.rest import _build_url, _resolve_api_key + from ..error import AuthenticationError - base = os.environ.get("RUNPOD_API_BASE_URL") - url = os.environ.get( - "RUNPOD_GPU_CATALOG_URL", - f"{base.rstrip('/')}/v2/catalog/gpus" if base else CATALOG_URL, - ) - key = api_key or os.environ.get("RUNPOD_API_KEY") or runpod.api_key or "" - return url, key - - -def _catalog_headers(key: str) -> Dict[str, str]: - if not key: - raise RuntimeError( - "no api key configured. run `rp login` or set RUNPOD_API_KEY." - ) - return {"Authorization": f"Bearer {key}"} + try: + key = _resolve_api_key(api_key) + except AuthenticationError: + key = "" + return _build_url(CATALOG_PATH), key def _parse_catalog(data: Any) -> Dict[str, Set[str]]: @@ -282,14 +272,6 @@ def _parse_catalog(data: Any) -> Dict[str, Set[str]]: return members -def _fetch_catalog(url: str, key: str) -> Dict[str, Set[str]]: - import requests - - resp = requests.get(url, headers=_catalog_headers(key), timeout=10) - resp.raise_for_status() - return _parse_catalog(resp.json()) - - def live_pool_members() -> Optional[Dict[str, Set[str]]]: """pool id -> device ids, from the platform's gpu catalog. @@ -298,11 +280,13 @@ def live_pool_members() -> Optional[Dict[str, Set[str]]]: membership. returns None when the catalog is unreachable. a result is kept for the process; a failure is retried after a minute. """ + from ..api.ctl_commands import get_gpus + context = _catalog_context() cache = _catalog_caches.setdefault(context, _CatalogCache()) if cache.may_fetch(): try: - cache.members = _fetch_catalog(*context) + cache.members = _parse_catalog(get_gpus(api_key=context[1])) except Exception as exc: # noqa: BLE001 - offline, auth, or malformed cache.fail(exc) return cache.members @@ -312,19 +296,16 @@ async def live_pool_members_async( api_key: Optional[str] = None, ) -> Optional[Dict[str, Set[str]]]: """fetch pool membership without blocking, cached by catalog url and credentials.""" - import aiohttp + from ..api.rest import run_rest_request_async context = _catalog_context(api_key) cache = _catalog_caches.setdefault(context, _CatalogCache()) if cache.may_fetch(): - url, key = context try: - headers = _catalog_headers(key) - timeout = aiohttp.ClientTimeout(total=10) - async with aiohttp.ClientSession(timeout=timeout) as session: - async with session.get(url, headers=headers) as response: - response.raise_for_status() - cache.members = _parse_catalog(await response.json()) + data = await run_rest_request_async( + "GET", CATALOG_PATH, api_key=context[1], timeout=10 + ) + cache.members = _parse_catalog(data) except Exception as exc: cache.fail(exc) return cache.members diff --git a/tests/test_apps/test_app.py b/tests/test_apps/test_app.py index ed72bec0b..882b9570f 100644 --- a/tests/test_apps/test_app.py +++ b/tests/test_apps/test_app.py @@ -429,8 +429,9 @@ def test_live_pool_members_parses_catalog_and_fails_soft(self, monkeypatch): from runpod.apps import gpu as gpu_mod - monkeypatch.setenv("RUNPOD_API_KEY", "k") + monkeypatch.setattr(runpod, "api_key", "k") resp = MagicMock() + resp.status_code = 200 resp.json.return_value = { "gpus": [ {"id": "NVIDIA L4", "pool": "AMPERE_24"}, @@ -439,11 +440,11 @@ def test_live_pool_members_parses_catalog_and_fails_soft(self, monkeypatch): } fetch = _real_live_pool_members gpu_mod._clear_catalog_cache() - with patch("requests.get", return_value=resp) as get: + with patch("requests.request", return_value=resp) as get: assert fetch() == {"AMPERE_24": {"NVIDIA L4"}} assert get.call_args.kwargs["headers"]["Authorization"] == "Bearer k" gpu_mod._clear_catalog_cache() - with patch("requests.get", side_effect=OSError("offline")) as get: + with patch("requests.request", side_effect=OSError("offline")) as get: assert fetch() is None # a failure is not retried on every lookup assert fetch() is None @@ -451,7 +452,7 @@ def test_live_pool_members_parses_catalog_and_fails_soft(self, monkeypatch): # it is retried once the window passes cache = next(iter(gpu_mod._catalog_caches.values())) cache.failed_at -= gpu_mod.CATALOG_RETRY_SECONDS - with patch("requests.get", return_value=resp): + with patch("requests.request", return_value=resp): assert fetch() == {"AMPERE_24": {"NVIDIA L4"}} gpu_mod._clear_catalog_cache() diff --git a/tests/test_apps/test_gpu_catalog.py b/tests/test_apps/test_gpu_catalog.py index 69f2a040f..a4baf4484 100644 --- a/tests/test_apps/test_gpu_catalog.py +++ b/tests/test_apps/test_gpu_catalog.py @@ -13,6 +13,7 @@ from runpod.apps.app import _clear_registry from runpod.apps.deploy import reconcile_endpoints from runpod.apps.dev import DevSession +from runpod.apps.gpu import live_pool_members as _real_live_pool_members from runpod.apps.gpu import live_pool_members_async as _real_live_pool_members_async @@ -30,7 +31,6 @@ def clean_catalog(monkeypatch, offline_gpu_catalog): _clear_registry() gpu._clear_catalog_cache() monkeypatch.delenv("RUNPOD_API_KEY", raising=False) - monkeypatch.delenv("RUNPOD_GPU_CATALOG_URL", raising=False) monkeypatch.setattr(runpod, "api_key", None) monkeypatch.setattr(gpu, "live_pool_members_async", _real_live_pool_members_async) monkeypatch.setattr( @@ -60,7 +60,7 @@ async def handle(request): app = web.Application() app.router.add_get("/v2/catalog/gpus", handle) - app.router.add_get("/custom-catalog", handle) + app.router.add_get("/custom/v2/catalog/gpus", handle) runner = web.AppRunner(app, shutdown_timeout=0.1) await runner.setup() try: @@ -184,6 +184,25 @@ async def test_failed_catalog_retries_after_cooldown(catalog, malformed, caplog) assert len(catalog.requests) == 2 +async def test_catalog_uses_sdk_authentication_errors(catalog, caplog): + catalog.status = 401 + assert await AppsApiClient(api_key="explicit-key").gpu_pool_members() is None + assert "Unauthorized request, please check your API key." in caplog.text + assert "explicit-key" not in caplog.text + + +def test_sync_catalog_reuses_gpu_catalog_command(monkeypatch): + from runpod.api import ctl_commands + + monkeypatch.setattr(runpod, "api_key", "sdk-key") + monkeypatch.setenv("RUNPOD_API_KEY", "different-env-key") + fetch = Mock(return_value=CATALOG["gpus"]) + monkeypatch.setattr(ctl_commands, "get_gpus", fetch) + assert MIG in _real_live_pool_members()["AMPERE_24"] + assert MIG in _real_live_pool_members()["AMPERE_24"] + fetch.assert_called_once_with(api_key="sdk-key") + + async def test_cache_is_scoped_to_credentials_and_catalog_url(catalog, monkeypatch): first = AppsApiClient(api_key="first-key") assert MIG in (await first.gpu_pool_members())["AMPERE_24"] @@ -195,9 +214,9 @@ async def test_cache_is_scoped_to_credentials_and_catalog_url(catalog, monkeypat monkeypatch.setenv("RUNPOD_API_BASE_URL", catalog.other_url) assert await first.gpu_pool_members() == {"AMPERE_24": {"NVIDIA L4"}} assert len(catalog.requests) == 3 - monkeypatch.setenv("RUNPOD_GPU_CATALOG_URL", f"{catalog.url}/custom-catalog") + monkeypatch.setenv("RUNPOD_API_BASE_URL", f"{catalog.url}/custom") assert await first.gpu_pool_members() == {"AMPERE_24": {"NVIDIA L4"}} - assert catalog.requests[-1] == ("/custom-catalog", "Bearer first-key") + assert catalog.requests[-1] == ("/custom/v2/catalog/gpus", "Bearer first-key") assert len(catalog.requests) == 4