From 1ffb5a0433521439064337a2215ba05be1de41c0 Mon Sep 17 00:00:00 2001 From: Justin Date: Fri, 16 Oct 2026 19:56:27 -0400 Subject: [PATCH] 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