From 1bdb8f68ad9c8ae3b70e0b848cdf273d7e54d815 Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Mon, 9 Nov 2026 15:41:56 -0800 Subject: [PATCH 1/3] fix(apps): tag HTTP traffic with apps user agent --- runpod/api/ctl_commands.py | 9 +- runpod/api/rest.py | 12 +- runpod/apps/api.py | 44 ++- runpod/apps/auth.py | 5 +- runpod/apps/gpu.py | 12 +- runpod/apps/logs.py | 3 +- runpod/apps/targets.py | 4 +- runpod/apps/tasks.py | 7 +- runpod/apps/utils/network.py | 6 +- runpod/network_volume_transfer.py | 7 +- runpod/user_agent.py | 17 +- tests/test_api/test_ctl_commands.py | 10 +- tests/test_apps/test_agent_auth.py | 4 +- tests/test_apps/test_api_client.py | 23 +- tests/test_apps/test_deployed_tasks.py | 8 +- tests/test_apps/test_gpu_catalog.py | 8 +- .../test_apps/test_runtime_call_validation.py | 3 +- tests/test_apps/test_user_agent.py | 338 ++++++++++++++++++ tests/test_user_agent.py | 38 +- 19 files changed, 513 insertions(+), 45 deletions(-) create mode 100644 tests/test_apps/test_user_agent.py diff --git a/runpod/api/ctl_commands.py b/runpod/api/ctl_commands.py index 1291df490..27685cfb1 100644 --- a/runpod/api/ctl_commands.py +++ b/runpod/api/ctl_commands.py @@ -7,6 +7,7 @@ from urllib.parse import quote from runpod import error +from runpod.user_agent import USER_AGENT from .graphql import run_graphql_query from .mutations import container_register_auth as container_register_auth_mutations @@ -63,9 +64,13 @@ def update_user_settings(pubkey: str, api_key: Optional[str] = None) -> dict: return get_user(api_key=api_key) -def get_gpus(api_key: Optional[str] = None) -> list[dict]: +def get_gpus( + api_key: Optional[str] = None, *, user_agent: str = USER_AGENT +) -> list[dict]: """Get all GPU types.""" - response = run_rest_request("GET", "/v2/catalog/gpus", api_key=api_key) + response = run_rest_request( + "GET", "/v2/catalog/gpus", api_key=api_key, user_agent=user_agent + ) return response["gpus"] diff --git a/runpod/api/rest.py b/runpod/api/rest.py index 63851d712..4cfb4a935 100644 --- a/runpod/api/rest.py +++ b/runpod/api/rest.py @@ -98,6 +98,7 @@ def run_rest_request( timeout: float = 30, max_retries: int = 3, max_rate_limit_wait: float = 60, + user_agent: str = USER_AGENT, ) -> Optional[dict[str, Any]]: """Send a REST request with bounded 429 retries and shared quota waits. @@ -109,6 +110,7 @@ def run_rest_request( key = _resolve_api_key(api_key) url = _build_url(path) headers = _build_headers(key) + headers["User-Agent"] = user_agent budget = _RetryBudget(max_retries, max_rate_limit_wait, f"{method.upper()} {path}") with rate_limits.bucket(url, key) as bucket: for attempt in range(max_retries + 1): @@ -167,6 +169,7 @@ async def run_rest_request_async( max_retries: int = 3, max_rate_limit_wait: float = 60, session: Optional[aiohttp.ClientSession] = None, + user_agent: str = USER_AGENT, ) -> Optional[dict[str, Any]]: """Send a REST request with cancellable quota waits and bounded 429 retries. @@ -178,6 +181,7 @@ async def run_rest_request_async( key = _resolve_api_key(api_key) url = _build_url(path) headers = _build_headers(key) + headers["User-Agent"] = user_agent budget = _RetryBudget(max_retries, max_rate_limit_wait, f"{method.upper()} {path}") client_timeout = aiohttp.ClientTimeout(total=timeout) context = nullcontext(session) if session is not None else aiohttp.ClientSession() @@ -233,6 +237,7 @@ async def run_rest_paginated_request_async( collection: str, *, api_key: Optional[str] = None, + user_agent: str = USER_AGENT, ) -> list[dict[str, Any]]: """read all cursor pages with one session and four transport attempts per page.""" records = [] @@ -244,7 +249,12 @@ async def run_rest_paginated_request_async( for attempt in range(4): try: data = await run_rest_request_async( - "GET", path, api_key=api_key, params=params, session=session + "GET", + path, + api_key=api_key, + params=params, + session=session, + user_agent=user_agent, ) break except (aiohttp.ClientError, OSError, asyncio.TimeoutError): diff --git a/runpod/apps/api.py b/runpod/apps/api.py index c6195c398..49dd1d6ca 100644 --- a/runpod/apps/api.py +++ b/runpod/apps/api.py @@ -16,6 +16,7 @@ from ..api.queries import user as user_queries from ..api.rest import run_rest_paginated_request_async, run_rest_request_async from ..error import AuthenticationError, QueryError +from ..user_agent import APPS_USER_AGENT _TRANSPORT_RETRIES = 4 @@ -31,7 +32,7 @@ async def _execute(self, query: str, task_input: Dict[str, Any]) -> Dict[str, An query, variables={"input": task_input}, api_key=(headers["Authorization"][7:] if headers else api_key()), - headers=headers, + headers={**(headers or {}), "User-Agent": APPS_USER_AGENT}, ) return response["data"] @@ -228,6 +229,7 @@ async def _execute( api_key=self._api_key, variables=variables, anonymous=anonymous, + headers={"User-Agent": APPS_USER_AGENT}, ) return response["data"] except (aiohttp.ClientError, OSError, asyncio.TimeoutError): @@ -237,6 +239,7 @@ async def _execute( api_key=self._api_key, variables=variables, anonymous=anonymous, + headers={"User-Agent": APPS_USER_AGENT}, ) return response["data"] @@ -248,6 +251,7 @@ async def save_endpoint(self, endpoint_input: Dict[str, Any]) -> Dict[str, Any]: current = await run_rest_request_async( "GET", f"/v2/serverless/{_path_segment(endpoint_id)}", + user_agent=APPS_USER_AGENT, api_key=self._api_key, ) if not current: @@ -281,6 +285,7 @@ async def save_endpoint(self, endpoint_input: Dict[str, Any]) -> Dict[str, Any]: return await run_rest_request_async( "PATCH" if endpoint_id else "POST", path, + user_agent=APPS_USER_AGENT, api_key=self._api_key, json=body, ) @@ -299,13 +304,14 @@ async def delete_endpoint(self, endpoint_id: str) -> bool: await run_rest_request_async( "DELETE", f"/v2/serverless/{_path_segment(endpoint_id)}", + user_agent=APPS_USER_AGENT, api_key=self._api_key, ) return True async def list_my_endpoints(self) -> List[Dict[str, Any]]: data = await run_rest_request_async( - "GET", "/v2/serverless", api_key=self._api_key + "GET", "/v2/serverless", user_agent=APPS_USER_AGENT, api_key=self._api_key ) return data["endpoints"] @@ -314,6 +320,7 @@ async def endpoint_workers(self, endpoint_id: str) -> Dict[str, Any]: return await run_rest_request_async( "GET", f"/v2/serverless/{_path_segment(endpoint_id)}/workers", + user_agent=APPS_USER_AGENT, api_key=self._api_key, timeout=10, ) @@ -352,7 +359,10 @@ async def deploy_task_pod( async def terminate_pod(self, pod_id: str) -> None: await run_rest_request_async( - "DELETE", f"/v2/pods/{_path_segment(pod_id)}", api_key=self._api_key + "DELETE", + f"/v2/pods/{_path_segment(pod_id)}", + user_agent=APPS_USER_AGENT, + api_key=self._api_key, ) async def get_app_by_name(self, name: str) -> Optional[Dict[str, Any]]: @@ -388,6 +398,7 @@ async def gpu_stock_status( data = await run_rest_request_async( "GET", f"/v2/catalog/gpus/{_path_segment(gpu_id)}", + user_agent=APPS_USER_AGENT, api_key=self._api_key, params={ "include": "AVAILABILITY", @@ -424,6 +435,7 @@ async def cpu_stock_status( data = await run_rest_request_async( "GET", f"/v2/catalog/cpus/{_path_segment(cpu['id'])}", + user_agent=APPS_USER_AGENT, api_key=self._api_key, params={ "include": "AVAILABILITY", @@ -436,7 +448,10 @@ async def cpu_stock_status( async def network_volume_datacenters(self) -> Set[str]: """datacenters supporting any network volume tier chosen by the backend.""" data = await run_rest_request_async( - "GET", "/v2/catalog/datacenters", api_key=self._api_key + "GET", + "/v2/catalog/datacenters", + user_agent=APPS_USER_AGENT, + api_key=self._api_key, ) datacenters = data.get("dataCenters") if isinstance(data, dict) else None if not isinstance(datacenters, list): @@ -461,7 +476,10 @@ async def network_volume_datacenters(self) -> Set[str]: async def list_global_volumes(self) -> List[Dict[str, Any]]: """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 + "/v2/global-volumes", + "globalVolumes", + user_agent=APPS_USER_AGENT, + api_key=self._api_key, ) async def create_global_volume(self, name: str) -> Dict[str, Any]: @@ -469,6 +487,7 @@ async def create_global_volume(self, name: str) -> Dict[str, Any]: data = await run_rest_request_async( "POST", "/v2/global-volumes", + user_agent=APPS_USER_AGENT, api_key=self._api_key, json={"name": name}, ) @@ -487,7 +506,10 @@ async def create_global_volume(self, name: str) -> Dict[str, Any]: async def list_network_volumes(self) -> List[Dict[str, Any]]: """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 + "/v2/network-volumes", + "networkVolumes", + user_agent=APPS_USER_AGENT, + api_key=self._api_key, ) async def create_network_volume( @@ -496,13 +518,14 @@ async def create_network_volume( return await run_rest_request_async( "POST", "/v2/network-volumes", + user_agent=APPS_USER_AGENT, api_key=self._api_key, json={"name": name, "size": size, "dataCenter": data_center_id}, ) async def list_registry_auths(self) -> List[Dict[str, Any]]: data = await run_rest_request_async( - "GET", "/v2/registries", api_key=self._api_key + "GET", "/v2/registries", user_agent=APPS_USER_AGENT, api_key=self._api_key ) return data["registries"] @@ -512,6 +535,7 @@ async def create_registry_auth( return await run_rest_request_async( "POST", "/v2/registries", + user_agent=APPS_USER_AGENT, api_key=self._api_key, json={"name": name, "username": username, "password": password}, ) @@ -520,6 +544,7 @@ async def delete_registry_auth(self, auth_id: str) -> bool: await run_rest_request_async( "DELETE", f"/v2/registries/{_path_segment(auth_id)}", + user_agent=APPS_USER_AGENT, api_key=self._api_key, ) return True @@ -603,7 +628,9 @@ async def verify_api_key(self, api_key: str) -> bool: """true when api_key authenticates; transport failures still raise.""" try: response = await run_graphql_query_async( - user_queries.QUERY_USER, api_key=api_key + user_queries.QUERY_USER, + api_key=api_key, + headers={"User-Agent": APPS_USER_AGENT}, ) except (AuthenticationError, QueryError): return False @@ -697,6 +724,7 @@ async def _put_tarball(self, upload_url, reader, total, timeout) -> None: upload_url, data=reader, headers={ + "User-Agent": APPS_USER_AGENT, "Content-Type": "application/gzip", "Content-Length": str(total), }, diff --git a/runpod/apps/auth.py b/runpod/apps/auth.py index 015b16db2..0082cf42a 100644 --- a/runpod/apps/auth.py +++ b/runpod/apps/auth.py @@ -13,7 +13,7 @@ from runpod._hosts import console_url from runpod.api.rest import _build_url -from runpod.user_agent import USER_AGENT +from runpod.user_agent import APPS_USER_AGENT import aiohttp @@ -21,7 +21,6 @@ from .api import AppsApiClient from .errors import AppError - MAX_KEY_ATTEMPTS = 3 _CLOSED_REASONS = { @@ -101,7 +100,7 @@ async def agent_login(email: str, *, read_code: Callable[[], str]) -> str: async with aiohttp.ClientSession( timeout=aiohttp.ClientTimeout(total=30), - headers={"Accept": "application/json", "User-Agent": USER_AGENT}, + headers={"Accept": "application/json", "User-Agent": APPS_USER_AGENT}, cookie_jar=aiohttp.DummyCookieJar(), ) as session: attempt = await _post_authentication_attempt( diff --git a/runpod/apps/gpu.py b/runpod/apps/gpu.py index 4b7ba72fc..1365dc83b 100644 --- a/runpod/apps/gpu.py +++ b/runpod/apps/gpu.py @@ -11,6 +11,8 @@ from enum import Enum from typing import Any, Dict, List, Optional, Set, Tuple, Union +from ..user_agent import APPS_USER_AGENT + log = logging.getLogger(__name__) @@ -286,7 +288,9 @@ def live_pool_members() -> Optional[Dict[str, Set[str]]]: cache = _catalog_caches.setdefault(context, _CatalogCache()) if cache.may_fetch(): try: - cache.members = _parse_catalog(get_gpus(api_key=context[1])) + cache.members = _parse_catalog( + get_gpus(api_key=context[1], user_agent=APPS_USER_AGENT) + ) except Exception as exc: # noqa: BLE001 - offline, auth, or malformed cache.fail(exc) return cache.members @@ -303,7 +307,11 @@ async def live_pool_members_async( if cache.may_fetch(): try: data = await run_rest_request_async( - "GET", CATALOG_PATH, api_key=context[1], timeout=10 + "GET", + CATALOG_PATH, + api_key=context[1], + timeout=10, + user_agent=APPS_USER_AGENT, ) cache.members = _parse_catalog(data) except Exception as exc: diff --git a/runpod/apps/logs.py b/runpod/apps/logs.py index 6b0aed8a7..11489fb75 100644 --- a/runpod/apps/logs.py +++ b/runpod/apps/logs.py @@ -13,6 +13,7 @@ from ..api.ctl_commands import _path_segment from ..api.rest import _build_url +from ..user_agent import APPS_USER_AGENT HAPI_BASE = os.environ.get("RUNPOD_HAPI_URL", "https://hapi.runpod.net") @@ -22,7 +23,7 @@ def _headers() -> Dict[str, str]: from .utils.network import api_key - return {"Authorization": f"Bearer {api_key()}"} + return {"Authorization": f"Bearer {api_key()}", "User-Agent": APPS_USER_AGENT} async def pod_logs( diff --git a/runpod/apps/targets.py b/runpod/apps/targets.py index 711ab3109..93c761a5d 100644 --- a/runpod/apps/targets.py +++ b/runpod/apps/targets.py @@ -14,7 +14,7 @@ import aiohttp from .._hosts import endpoint_domain -from ..user_agent import USER_AGENT +from ..user_agent import APPS_USER_AGENT from .errors import EndpointNotFound, RemoteExecutionError from .protocol import FORMAT_JSON, FunctionRequest from .serialization import require_json_value @@ -43,7 +43,7 @@ def _headers(extra: Optional[Dict[str, str]] = None) -> Dict[str, str]: headers = { "Content-Type": "application/json", - "User-Agent": USER_AGENT, + "User-Agent": APPS_USER_AGENT, } runtime = runtime_headers() headers.update( diff --git a/runpod/apps/tasks.py b/runpod/apps/tasks.py index 7e7792911..08d4bd0cb 100644 --- a/runpod/apps/tasks.py +++ b/runpod/apps/tasks.py @@ -20,6 +20,7 @@ import aiohttp from .._hosts import proxy_url +from ..user_agent import APPS_USER_AGENT from .api import AppsApiClient, is_capacity_error from .context import CLEANUP_TIMEOUT from .spec import ResourceSpec @@ -206,7 +207,7 @@ def __init__(self, api): @property def _headers(self) -> Dict[str, str]: - return {"Authorization": f"Bearer {self.token}"} + return {"Authorization": f"Bearer {self.token}", "User-Agent": APPS_USER_AGENT} async def wait_ready(self, timeout: float = READY_TIMEOUT) -> None: """poll the runner's /ping through the pod proxy until it answers.""" @@ -216,7 +217,9 @@ async def wait_ready(self, timeout: float = READY_TIMEOUT) -> None: while time.monotonic() < deadline: try: async with session.get( - url, timeout=aiohttp.ClientTimeout(total=5) + url, + timeout=aiohttp.ClientTimeout(total=5), + headers={"User-Agent": APPS_USER_AGENT}, ) as resp: if resp.status == 200: return diff --git a/runpod/apps/utils/network.py b/runpod/apps/utils/network.py index 0b1818b59..af4316fbc 100644 --- a/runpod/apps/utils/network.py +++ b/runpod/apps/utils/network.py @@ -5,6 +5,8 @@ import aiohttp +from ...user_agent import APPS_USER_AGENT + def runtime_headers() -> Optional[Dict[str, str]]: from ..context import deployed_identity @@ -60,6 +62,8 @@ def endpoint_url_base() -> str: async def get_json(url: str, headers: Dict[str, str], timeout: float) -> Dict[str, Any]: client_timeout = aiohttp.ClientTimeout(total=timeout) async with aiohttp.ClientSession(timeout=client_timeout) as session: - async with session.get(url, headers=headers) as response: + async with session.get( + url, headers={**headers, "User-Agent": APPS_USER_AGENT} + ) as response: response.raise_for_status() return await response.json() diff --git a/runpod/network_volume_transfer.py b/runpod/network_volume_transfer.py index 1fcdbdfc1..4c74b1e41 100644 --- a/runpod/network_volume_transfer.py +++ b/runpod/network_volume_transfer.py @@ -30,6 +30,8 @@ from pathlib import Path, PurePosixPath from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar +from .user_agent import APPS_USER_AGENT + from ._volume_transfer_utils import ( # noqa: F401 - re-exported names MIB, PARTIAL_SUFFIX, @@ -852,7 +854,9 @@ def _list_volumes() -> List[Dict[str, Any]]: try: return block( - run_rest_paginated_request_async("/v2/network-volumes", "networkVolumes") + run_rest_paginated_request_async( + "/v2/network-volumes", "networkVolumes", user_agent=APPS_USER_AGENT + ) ) except error.AuthenticationError as exc: from runpod import api_key # pylint: disable=import-outside-toplevel @@ -927,6 +931,7 @@ def _client_config(concurrency: int) -> Any: from botocore.config import Config # pylint: disable=import-outside-toplevel return Config( + user_agent_extra=APPS_USER_AGENT, s3={"addressing_style": "path"}, request_checksum_calculation="when_required", response_checksum_validation="when_required", diff --git a/runpod/user_agent.py b/runpod/user_agent.py index 021745de8..f2a9c9c6d 100644 --- a/runpod/user_agent.py +++ b/runpod/user_agent.py @@ -1,18 +1,26 @@ -""" User-Agent for RunPod-Python-SDK """ +"""User-Agent for RunPod-Python-SDK""" import os import platform +import re +from typing import Optional from runpod import agent from runpod.version import __version__ as runpod_version -def construct_user_agent(): +def construct_user_agent(integration: Optional[str] = None): """Constructs the User-Agent string for the RunPod-Python-SDK Example: RunPod-Python-SDK/0.1.0 (Linux 5.4.0-54-generic; x86_64) Language/Python 3.8.5 + + explicit integrations precede the user's RUNPOD_UA_INTEGRATION tag. """ + if integration is not None and not re.fullmatch(r"[A-Za-z0-9._-]+", integration): + raise ValueError( + "integration must be a non-empty token using A-Z, a-z, 0-9, ., _, or -" + ) os_info = f"{platform.system()} {platform.release()}; {platform.machine()}" python_version = platform.python_version() integration_method = os.getenv("RUNPOD_UA_INTEGRATION") @@ -23,7 +31,9 @@ def construct_user_agent(): f"Language/Python {python_version}", ] - if integration_method: + if integration: + ua_components.append(f"Integration/{integration}") + if integration_method and integration_method != integration: ua_components.append(f"Integration/{integration_method}") agent_suffix = agent.suffix() @@ -35,3 +45,4 @@ def construct_user_agent(): USER_AGENT = construct_user_agent() +APPS_USER_AGENT = construct_user_agent(integration="apps") diff --git a/tests/test_api/test_ctl_commands.py b/tests/test_api/test_ctl_commands.py index d3134abc8..c388731d9 100644 --- a/tests/test_api/test_ctl_commands.py +++ b/tests/test_api/test_ctl_commands.py @@ -49,7 +49,11 @@ def test_get_gpus_unwraps_response(): ) as request: assert ctl_commands.get_gpus(api_key="key") == gpus - request.assert_called_once_with("GET", "/v2/catalog/gpus", api_key="key") + from runpod.user_agent import USER_AGENT + + request.assert_called_once_with( + "GET", "/v2/catalog/gpus", api_key="key", user_agent=USER_AGENT + ) def test_get_gpu_requests_pod_availability(): @@ -296,7 +300,9 @@ def test_create_pod_validates_image_and_cloud(): ({"min_upload": 100}, "min_upload"), ], ) -def test_create_gpu_pod_rejects_unsupported_constraints(template_backend, kwargs, field): +def test_create_gpu_pod_rejects_unsupported_constraints( + template_backend, kwargs, field +): with pytest.raises(ValueError, match=field): ctl_commands.create_pod("pod", "image", gpu_type_id="NVIDIA A100", **kwargs) diff --git a/tests/test_apps/test_agent_auth.py b/tests/test_apps/test_agent_auth.py index 0df1a4ce9..7f3c67937 100644 --- a/tests/test_apps/test_agent_auth.py +++ b/tests/test_apps/test_agent_auth.py @@ -15,7 +15,7 @@ import runpod from runpod.apps.auth import LoginError, _post_authentication_attempt, agent_login -from runpod.user_agent import USER_AGENT +from runpod.user_agent import APPS_USER_AGENT @pytest.fixture @@ -94,7 +94,7 @@ def read_code(): for path, headers, _ in (start, verify): assert "Authorization" not in headers assert "Cookie" not in headers - assert headers["User-Agent"] == USER_AGENT + assert headers["User-Agent"] == APPS_USER_AGENT assert verifier not in path diff --git a/tests/test_apps/test_api_client.py b/tests/test_apps/test_api_client.py index 4c56429d3..3bbbfbce0 100644 --- a/tests/test_apps/test_api_client.py +++ b/tests/test_apps/test_api_client.py @@ -7,6 +7,7 @@ from runpod.apps.api import AppsApiClient from runpod.error import AuthenticationError, QueryError +from runpod.user_agent import APPS_USER_AGENT def _respond(data): @@ -95,7 +96,12 @@ async def test_list(self, volumes): 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 + "GET", + "/v2/global-volumes", + api_key="test-key", + params={}, + session=ANY, + user_agent=APPS_USER_AGENT, ) graphql.assert_not_awaited() @@ -123,6 +129,7 @@ async def test_list_follows_all_cursors(self): call( "GET", "/v2/global-volumes", + user_agent=APPS_USER_AGENT, api_key="test-key", params=params, session=ANY, @@ -180,6 +187,7 @@ async def test_create(self): rest.assert_awaited_once_with( "POST", "/v2/global-volumes", + user_agent=APPS_USER_AGENT, api_key="test-key", json={"name": "training-data"}, ) @@ -243,6 +251,7 @@ async def test_list_follows_all_cursors(self): call( "GET", "/v2/network-volumes", + user_agent=APPS_USER_AGENT, api_key="test-key", params=params, session=ANY, @@ -272,7 +281,8 @@ async def test_dev_endpoint_preserves_runtime_and_placement(self, monkeypatch): payload["template"]["containerRegistryAuthId"] = "registry1" payload["networkVolumeIds"] = [{"networkVolumeId": "volume1"}] - async def create(method, path, *, api_key, json): + async def create(method, path, *, api_key, json, user_agent): + assert user_agent == APPS_USER_AGENT assert (method, path) == ("POST", "/v2/serverless") assert json["gpu"]["count"] == 2 assert json["dataCenterIds"] == ["EU-RO-1"] @@ -608,7 +618,8 @@ async def test_get_environment_not_found(self): class TestStock: async def test_gpu_availability_respects_count_product_and_datacenter(self): - async def catalog(method, path, *, api_key, params): + async def catalog(method, path, *, api_key, params, user_agent): + assert user_agent == APPS_USER_AGENT available = ( params["product"] == "SERVERLESS" and params["count"] == 1 @@ -629,7 +640,8 @@ async def catalog(method, path, *, api_key, params): assert await client.gpu_stock_status("4090", "US-KS-2") == "NONE" async def test_cpu_availability_respects_flavor_size_and_product(self): - async def catalog(method, path, *, api_key, params): + async def catalog(method, path, *, api_key, params, user_agent): + assert user_agent == APPS_USER_AGENT available = ( path == "/v2/catalog/cpus/cpu3c" and params["vcpuCount"] == 2 @@ -657,7 +669,8 @@ async def catalog(method, path, *, api_key, params): @pytest.mark.parametrize("pods", [False, True]) @pytest.mark.parametrize("instance_id", ["cpu3c-1-2", "cpu3c-3-6"]) async def test_legacy_cpu_stock_keeps_exact_configuration(self, pods, instance_id): - async def graphql(query, *, api_key, variables, anonymous): + async def graphql(query, *, api_key, variables, anonymous, headers): + assert headers == {"User-Agent": APPS_USER_AGENT} flavor = variables["cpuFlavorInput"] specifics = variables["specificsInput"] available = flavor == {"id": "cpu3c", "isSls": not pods} and specifics == { diff --git a/tests/test_apps/test_deployed_tasks.py b/tests/test_apps/test_deployed_tasks.py index 4fa5ce1ae..45478ab3b 100644 --- a/tests/test_apps/test_deployed_tasks.py +++ b/tests/test_apps/test_deployed_tasks.py @@ -6,6 +6,7 @@ import pytest import runpod +from runpod.user_agent import APPS_USER_AGENT from runpod.apps import App, Task from runpod.apps.api import FlashTaskApiClient from runpod.apps.app import _clear_registry @@ -177,7 +178,7 @@ async def test_explicit_local_task_uses_account_named_api(monkeypatch): with patch("runpod.apps.api.run_graphql_query_async", execute): await FlashTaskApiClient().create({"resourceName": "work"}) assert execute.await_args.kwargs["api_key"] == "local-key" - assert execute.await_args.kwargs["headers"] is None + assert execute.await_args.kwargs["headers"] == {"User-Agent": APPS_USER_AGENT} async def test_task_target_selection_is_production_only(production, monkeypatch): @@ -324,7 +325,10 @@ async def test_execution_creation_contains_only_named_identity_and_timeout(): "executionTimeoutSeconds": 12.5, } ) - assert execution._headers == {"Authorization": f"Bearer {execution.token}"} + assert execution._headers == { + "Authorization": f"Bearer {execution.token}", + "User-Agent": APPS_USER_AGENT, + } await execution.terminate() diff --git a/tests/test_apps/test_gpu_catalog.py b/tests/test_apps/test_gpu_catalog.py index c6f7b54c2..ecfc72698 100644 --- a/tests/test_apps/test_gpu_catalog.py +++ b/tests/test_apps/test_gpu_catalog.py @@ -15,7 +15,7 @@ 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 - +from runpod.user_agent import APPS_USER_AGENT MIG = "NVIDIA RTX PRO 6000 Blackwell Server Edition MIG 1g.24gb" CATALOG = { @@ -165,9 +165,7 @@ async def test_failed_catalog_uses_static_payload_without_sync_fetch( ): 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 payload["gpuIds"] == ("AMPERE_24,-NVIDIA GeForce RTX 3090,-NVIDIA RTX A5000") assert len(catalog.requests) == 1 @@ -210,7 +208,7 @@ def test_sync_catalog_reuses_gpu_catalog_command(monkeypatch): 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") + fetch.assert_called_once_with(api_key="sdk-key", user_agent=APPS_USER_AGENT) async def test_cache_is_scoped_to_credentials_and_catalog_url(catalog, monkeypatch): diff --git a/tests/test_apps/test_runtime_call_validation.py b/tests/test_apps/test_runtime_call_validation.py index 782618c3d..af60d02f9 100644 --- a/tests/test_apps/test_runtime_call_validation.py +++ b/tests/test_apps/test_runtime_call_validation.py @@ -7,6 +7,7 @@ import pytest import runpod +from runpod.user_agent import APPS_USER_AGENT from runpod.apps import App from runpod.apps.api import FlashTaskApiClient from runpod.apps.app import _clear_registry @@ -110,7 +111,7 @@ async def test_account_clients_can_target_other_apps_and_environments(monkeypatc } ) assert execute.await_args.kwargs["api_key"] == "account-key" - assert execute.await_args.kwargs["headers"] is None + assert execute.await_args.kwargs["headers"] == {"User-Agent": APPS_USER_AGENT} @dataclass diff --git a/tests/test_apps/test_user_agent.py b/tests/test_apps/test_user_agent.py new file mode 100644 index 000000000..6c467f5d1 --- /dev/null +++ b/tests/test_apps/test_user_agent.py @@ -0,0 +1,338 @@ +"""apps attribution on control-plane and data-plane HTTP requests.""" + +import json +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch + +import aiohttp +import pytest +import requests + +import runpod +from runpod.api.graphql import run_graphql_query, run_graphql_query_async +from runpod.api.rest import run_rest_request, run_rest_request_async +from runpod.apps import gpu, logs, tasks +from runpod.apps.api import AppsApiClient, FlashTaskApiClient +from runpod.apps.auth import agent_login +from runpod.apps.gpu import live_pool_members, live_pool_members_async +from runpod.apps.targets import LiveTarget, SentinelTarget +from runpod.apps.utils.network import get_json +from runpod.user_agent import APPS_USER_AGENT, USER_AGENT + + +@pytest.fixture +def server(monkeypatch): + assert "Integration/apps" in APPS_USER_AGENT.split() + state = SimpleNamespace(requests=[], responses=[{}], user_agent=APPS_USER_AGENT) + + class Handler(BaseHTTPRequestHandler): + def respond(self): + self.rfile.read(int(self.headers.get("Content-Length", 0))) + state.requests.append((self.command, self.path, dict(self.headers))) + data = state.responses.pop(0) + status, data = data if isinstance(data, tuple) else (200, data) + body = data if isinstance(data, bytes) else json.dumps(data).encode() + self.send_response(status) + self.send_header("Content-Length", str(len(body))) + self.send_header("Content-Type", "application/json") + self.end_headers() + self.wfile.write(body) + + do_GET = do_POST = do_PATCH = do_PUT = do_DELETE = respond + + def log_message(self, *args): + pass + + httpd = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread( + target=httpd.serve_forever, kwargs={"poll_interval": 0.01}, daemon=True + ) + thread.start() + state.url = f"http://127.0.0.1:{httpd.server_port}" + monkeypatch.setenv("RUNPOD_API_BASE_URL", state.url) + monkeypatch.setenv("RUNPOD_API_KEY", "test-key") + monkeypatch.setattr(runpod, "api_key", "test-key") + monkeypatch.setattr(runpod, "endpoint_url_base", state.url) + monkeypatch.setattr(logs, "HAPI_BASE", state.url) + monkeypatch.setattr(tasks, "_proxy_url", lambda pod_id: state.url) + monkeypatch.setattr("runpod.apps.context.deployed_identity", lambda: None) + gpu._clear_catalog_cache() + try: + yield state + assert state.requests + assert not state.responses + for _, _, headers in state.requests: + assert headers["User-Agent"] == state.user_agent + finally: + gpu._clear_catalog_cache() + httpd.shutdown() + httpd.server_close() + thread.join() + + +@pytest.mark.parametrize("anonymous", [False, True]) +@pytest.mark.parametrize("retry", [False, True]) +async def test_graphql_lifecycle(server, anonymous, retry): + server.responses = [{"data": {"ok": True}}] + client = AppsApiClient(api_key="test-key") + assert await client._execute("query {}", anonymous=anonymous, retry=retry) + headers = server.requests[0][2] + assert ("Authorization" not in headers) == anonymous + + +async def test_graphql_retry_keeps_attribution(server): + server.responses = [{"data": {"ok": True}}] + original = run_graphql_query_async + attempts = [] + + async def query(*args, **kwargs): + attempts.append(kwargs["headers"]) + if len(attempts) < 4: + raise aiohttp.ClientConnectionError("reset") + return await original(*args, **kwargs) + + with ( + patch("runpod.apps.api.run_graphql_query_async", query), + patch("runpod.apps.api._TRANSPORT_RETRIES", 4), + patch("asyncio.sleep", AsyncMock()), + ): + await AppsApiClient()._execute("query {}", retry=True) + assert attempts == [{"User-Agent": APPS_USER_AGENT}] * 4 + + +async def test_verify_api_key(server): + server.responses = [{"data": {"myself": {"id": "user"}}}] + assert await AppsApiClient().verify_api_key("test-key") + + +@pytest.mark.parametrize("runtime", [None, "endpoint", "pod"]) +async def test_task_lifecycle_preserves_runtime_auth(server, runtime, monkeypatch): + headers = ( + { + "Authorization": "Bearer runtime-key", + f"X-RunPod-{runtime.title()}-Id": "source", + } + if runtime + else None + ) + monkeypatch.setattr("runpod.apps.utils.network.runtime_headers", lambda: headers) + server.responses = [{"data": {"ok": True}}] + await FlashTaskApiClient()._execute("query {}", {}) + sent = server.requests[0][2] + assert sent["Authorization"] == ( + "Bearer runtime-key" if runtime else "Bearer test-key" + ) + for key, value in (headers or {}).items(): + assert sent[key] == value + + +@pytest.mark.parametrize( + "method,args,response", + [ + ( + "save_endpoint", + ( + { + "name": "ep", + "template": {"imageName": "image"}, + "scalerType": "QUEUE_DELAY", + "scalerValue": 4, + }, + ), + {"id": "ep"}, + ), + ("delete_endpoint", ("ep",), {}), + ("list_my_endpoints", (), {"endpoints": []}), + ("endpoint_workers", ("ep",), {"workers": []}), + ("terminate_pod", ("pod",), {}), + ("gpu_stock_status", ("gpu", "dc"), {"dataCenters": []}), + ("cpu_stock_status", ("cpu3g-2-8", "dc"), {"dataCenters": []}), + ("network_volume_datacenters", (), {"dataCenters": []}), + ("create_global_volume", ("data",), {"id": "volume"}), + ("create_network_volume", ("data", 10, "dc"), {"id": "volume"}), + ("list_registry_auths", (), {"registries": []}), + ("create_registry_auth", ("registry", "user", "password"), {"id": "registry"}), + ("delete_registry_auth", ("registry",), {}), + ], +) +async def test_rest_resource_operations(server, method, args, response): + server.responses = [response] + await getattr(AppsApiClient(), method)(*args) + + +async def test_rest_endpoint_update(server): + server.responses = [{"id": "ep", "type": "QUEUE"}, {"id": "ep"}] + await AppsApiClient().save_endpoint({"id": "ep", "workersMin": 1}) + assert [request[0] for request in server.requests] == ["GET", "PATCH"] + + +@pytest.mark.parametrize( + "kind,collection", [("global", "globalVolumes"), ("network", "networkVolumes")] +) +async def test_rest_pagination(server, kind, collection): + server.responses = [ + { + collection: [{"id": "one"}], + "pagination": {"hasNextPage": True, "nextCursor": "next"}, + }, + {collection: [{"id": "two"}], "pagination": {"hasNextPage": False}}, + ] + assert len(await getattr(AppsApiClient(), f"list_{kind}_volumes")()) == 2 + assert server.requests[1][1].endswith("?cursor=next") + + +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_gpu_catalog(server, asynchronous): + server.responses = [{"gpus": [{"id": "gpu", "pool": "pool"}]}] + result = await live_pool_members_async() if asynchronous else live_pool_members() + assert result == {"pool": {"gpu"}} + + +async def test_artifact_upload(server, tmp_path): + tarball = tmp_path / "app.tar.gz" + tarball.write_bytes(b"archive") + await AppsApiClient().upload_tarball(server.url, str(tarball)) + assert server.requests[0][2]["Content-Type"] == "application/gzip" + + +async def test_email_authentication(server): + server.responses = [ + (201, {"attemptId": "attempt"}), + {"credential": {"apiKey": "key"}}, + ] + assert await agent_login("owner@example.com", read_code=lambda: "123456") == "key" + assert all("Authorization" not in r[2] for r in server.requests) + + +@pytest.mark.parametrize( + "target", [LiveTarget("dev"), SentinelTarget("app", "env", "work")] +) +@pytest.mark.parametrize( + "operation", ["run", "invoke", "status", "cancel", "retry", "stream", "http"] +) +async def test_queue_and_load_balancer(server, target, operation, monkeypatch): + server.responses = [{"id": "job", "status": "COMPLETED", "output": {}}] + client = target._client + if operation == "http": + original = aiohttp.ClientSession._request + + async def request(session, method, url, **kwargs): + assert url.startswith("https://") + return await original(session, method, server.url + "/route", **kwargs) + + monkeypatch.setattr(aiohttp.ClientSession, "_request", request) + await client.http("POST", "/route", {}) + elif operation == "stream": + assert [chunk async for chunk in client.stream("job", timeout=10)] == [] + elif operation == "invoke": + await client.invoke({"input": {}}, timeout=10) + else: + await getattr(client, operation)({"input": {}} if operation == "run" else "job") + + +@pytest.mark.parametrize( + "operation", ["wait_ready", "submit", "set_timeout", "poll_result"] +) +async def test_task_proxy(server, operation): + transport = tasks._TaskTransport(None) + transport.pod_id = "pod" + if operation == "submit": + await transport.submit({"args": []}) + elif operation == "set_timeout": + server.responses = [{"timeout": 5}] + await transport.set_timeout(5) + else: + await getattr(transport, operation)() + if operation == "wait_ready": + assert "Authorization" not in server.requests[0][2] + else: + assert server.requests[0][2]["Authorization"] == f"Bearer {transport.token}" + + +@pytest.mark.parametrize( + "log_type,endpoint", [("all", None), ("container", None), ("system", "ep")] +) +async def test_log_streams(server, log_type, endpoint): + server.responses = [b'data: {"line":"ready"}\n\n'] + assert [ + event + async for event in logs.stream_pod_logs( + "pod", log_type=log_type, endpoint_id=endpoint + ) + ] == [{"line": "ready"}] + + +async def test_log_snapshot(server): + server.responses = [{"container": ["ready"]}] + assert await logs.pod_logs("pod") == {"container": ["ready"]} + + +async def test_lightweight_json_client(server): + await get_json(server.url, {"Authorization": "Bearer test-key"}, 10) + assert server.requests[0][2]["Authorization"] == "Bearer test-key" + + +def test_volume_lookup(server): + from runpod.network_volume_transfer import _locate + + server.responses = [ + { + "networkVolumes": [{"id": "vol", "name": "data", "dataCenter": "EU-RO-1"}], + "pagination": {"hasNextPage": False}, + } + ] + assert _locate("data")[0] == "vol" + + +def test_s3_transfer_user_agent(): + from runpod.network_volume_transfer import _make_transfer + + transfer = _make_transfer( + "vol", + "test-access", + "test-secret", + None, + 1, + None, + location=("vol", "EU-RO-1", "https://s3.example.test"), + ) + requests = [] + + def capture(request, **kwargs): + requests.append(request) + raise RuntimeError("captured request") + + transfer.client.meta.events.register("before-send.s3", capture) + try: + with pytest.raises(RuntimeError, match="captured request"): + transfer.client.head_object(Bucket="vol", Key="file") + assert APPS_USER_AGENT in requests[0].headers["User-Agent"].decode() + finally: + transfer.client.close() + + +async def test_plain_sdk_transports_keep_default_user_agent(server): + server.user_agent = USER_AGENT + server.responses = [{}, {}, {"data": {}}, {"data": {}}] + run_rest_request("GET", "/v2/pods") + await run_rest_request_async("GET", "/v2/pods") + run_graphql_query("query {}") + await run_graphql_query_async("query {}") + for _, _, headers in server.requests: + assert headers["User-Agent"] == USER_AGENT + + +def test_plain_serverless_calls_keep_requests_user_agent(server): + server.user_agent = requests.utils.default_user_agent() + server.responses = [ + {"id": "job"}, + {"id": "job", "status": "COMPLETED", "output": {}}, + ] + endpoint = runpod.Endpoint("ep") + try: + endpoint.run({"value": 1}) + assert endpoint.run_sync({"value": 1}) == {} + finally: + endpoint.rp_client.rp_session.close() diff --git a/tests/test_user_agent.py b/tests/test_user_agent.py index 354d2a1c9..6589a7eb6 100644 --- a/tests/test_user_agent.py +++ b/tests/test_user_agent.py @@ -1,14 +1,49 @@ -""" Tests for the user_agent module. """ +"""Tests for the user_agent module.""" import os import unittest from unittest.mock import patch +import pytest + from runpod import __version__ as runpod_version from runpod import agent from runpod.user_agent import construct_user_agent +@pytest.mark.parametrize("integration", ["apps", "apps-dev", "SDK_2.1"]) +def test_explicit_integration_is_a_separate_token(monkeypatch, integration): + monkeypatch.delenv("RUNPOD_UA_INTEGRATION", raising=False) + with patch("runpod.user_agent.agent.suffix", return_value=" (via claude-code)"): + default = construct_user_agent() + tagged = construct_user_agent(integration=integration) + assert tagged == default.replace(" (via", f" Integration/{integration} (via") + assert construct_user_agent() == default + + +@pytest.mark.parametrize( + "integration", ["", "apps dev", "apps/dev", "apps\r\nX-Header:value", "äpp"] +) +def test_explicit_integration_rejects_non_tokens(integration): + with pytest.raises(ValueError, match="integration must"): + construct_user_agent(integration=integration) + + +@pytest.mark.parametrize("custom", ["SkyPilot", "apps", ""]) +def test_apps_preserves_user_integration(monkeypatch, custom): + monkeypatch.setenv("RUNPOD_UA_INTEGRATION", custom) + default = construct_user_agent() + tagged = construct_user_agent(integration="apps") + tokens = [token for token in tagged.split() if token.startswith("Integration/")] + assert tokens == ( + ["Integration/apps", "Integration/SkyPilot"] + if custom == "SkyPilot" + else ["Integration/apps"] + ) + assert construct_user_agent() == default + assert os.environ["RUNPOD_UA_INTEGRATION"] == custom + + def _agent_env_keys(): """Every env var that could inject an agent tag, plus the integration var.""" return agent.known_env_vars() + ["RUNPOD_UA_INTEGRATION"] @@ -59,7 +94,6 @@ def test_user_agent_with_integration( assert mock_release.called assert mock_system.called - @patch("runpod.user_agent.platform.system", return_value="Linux") @patch("runpod.user_agent.platform.release", return_value="5.4") @patch("runpod.user_agent.platform.machine", return_value="x86_64") From 679be197ab4cdb7dea26d5a55b2a9b316c7478c8 Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Wed, 11 Nov 2026 11:23:31 -0800 Subject: [PATCH 2/3] docs: focus README on quickstarts --- README.md | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/README.md b/README.md index 819b47422..9273152f3 100644 --- a/README.md +++ b/README.md @@ -84,18 +84,6 @@ make setup make test ``` -## Remote values - -Queue and task calls use finite JSON-compatible arguments, return values, -and stream chunks. JSON conversion turns tuples into lists and supported -dictionary keys into strings, including in same-worker remote calls. -Bytes, custom objects, circular references, and non-finite numbers produce -explicit errors. `.local()` remains a direct Python call. - -Successful function responses contain `json_result`, including for a null -return value. SDK and worker runtimes share the JSON protocol validators; -use compatible package releases when building worker images. - ## ⚡ | Serverless Worker (SDK) This python package can also be used to create a serverless worker that can be deployed to Runpod as a custom endpoint API. From 70ab87d08d7e56ea044d1fdf79a3acc853b51c7f Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Wed, 11 Nov 2026 12:56:13 -0800 Subject: [PATCH 3/3] fix(volume): distinguish CLI and Apps transfer traffic --- runpod/apps/volume.py | 8 ++- runpod/network_volume_transfer.py | 31 +++++---- runpod/rp_cli/volume.py | 2 + tests/test_apps/test_user_agent.py | 92 +++++++++++++++++++++++++-- tests/test_network_volume_transfer.py | 18 ++++-- 5 files changed, 128 insertions(+), 23 deletions(-) diff --git a/runpod/apps/volume.py b/runpod/apps/volume.py index 3ae904a38..ab93661a5 100644 --- a/runpod/apps/volume.py +++ b/runpod/apps/volume.py @@ -9,6 +9,7 @@ from types import MappingProxyType from typing import Any, Dict, List, Optional, Set, Tuple +from ..user_agent import APPS_USER_AGENT from .errors import AppError from .utils.client import default_client from .utils.events import emit @@ -77,6 +78,7 @@ def __init__( self.datacenter = datacenter self.create = create self._s3_locations: Dict[Optional[str], Tuple[str, str, str]] = {} + self._user_agent = APPS_USER_AGENT @property def name(self) -> str: @@ -93,7 +95,9 @@ def _s3_location(self, endpoint_url: Optional[str] = None) -> Tuple[str, str, st if location is None: from runpod import network_volume_transfer as nvt - location = nvt._locate(self.reference, endpoint_url) + location = nvt._locate( + self.reference, endpoint_url, user_agent=self._user_agent + ) self._s3_locations[endpoint_url] = location return location @@ -137,6 +141,7 @@ def upload( progress, force, location=self._s3_location(endpoint_url), + user_agent=self._user_agent, ) return transfer.upload(Path(local_path), remote_path) @@ -174,6 +179,7 @@ def download( progress, force, location=self._s3_location(endpoint_url), + user_agent=self._user_agent, ) # passed through as given: Path() would drop a trailing "/" return transfer.download(remote_path, local_path) diff --git a/runpod/network_volume_transfer.py b/runpod/network_volume_transfer.py index 4c74b1e41..c55a37822 100644 --- a/runpod/network_volume_transfer.py +++ b/runpod/network_volume_transfer.py @@ -30,7 +30,7 @@ from pathlib import Path, PurePosixPath from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar -from .user_agent import APPS_USER_AGENT +from .user_agent import USER_AGENT from ._volume_transfer_utils import ( # noqa: F401 - re-exported names MIB, @@ -844,7 +844,7 @@ def _read_range(self, start: int, length: int) -> bytes: # ---------------------------------------------------------------- client -def _list_volumes() -> List[Dict[str, Any]]: +def _list_volumes(*, user_agent: str = USER_AGENT) -> List[Dict[str, Any]]: """every network volume of the account.""" from runpod import error # pylint: disable=import-outside-toplevel from runpod.api.rest import ( # pylint: disable=import-outside-toplevel @@ -855,7 +855,7 @@ def _list_volumes() -> List[Dict[str, Any]]: try: return block( run_rest_paginated_request_async( - "/v2/network-volumes", "networkVolumes", user_agent=APPS_USER_AGENT + "/v2/network-volumes", "networkVolumes", user_agent=user_agent ) ) except error.AuthenticationError as exc: @@ -866,14 +866,14 @@ def _list_volumes() -> List[Dict[str, Any]]: raise NetworkVolumeTransferError(API_KEY_REJECTED) from exc -def _resolve_volume(reference: str) -> Tuple[str, str]: +def _resolve_volume(reference: str, *, user_agent: str = USER_AGENT) -> Tuple[str, str]: """the id and datacenter of the network volume with this id or name.""" from runpod.apps.utils.lookup import ( # pylint: disable=import-outside-toplevel find_by_id_or_name, ) record = find_by_id_or_name( - _list_volumes(), + _list_volumes(user_agent=user_agent), reference, noun="network volumes", error=NetworkVolumeTransferError, @@ -927,11 +927,11 @@ def _s3_keys(access_key: Optional[str], secret_key: Optional[str]) -> Tuple[str, return keys -def _client_config(concurrency: int) -> Any: +def _client_config(concurrency: int, *, user_agent: str = USER_AGENT) -> Any: from botocore.config import Config # pylint: disable=import-outside-toplevel return Config( - user_agent_extra=APPS_USER_AGENT, + user_agent_extra=user_agent, s3={"addressing_style": "path"}, request_checksum_calculation="when_required", response_checksum_validation="when_required", @@ -945,10 +945,15 @@ def _client_config(concurrency: int) -> Any: ) -def _locate(reference: str, endpoint_url: Optional[str] = None) -> Tuple[str, str, str]: +def _locate( + reference: str, + endpoint_url: Optional[str] = None, + *, + user_agent: str = USER_AGENT, +) -> Tuple[str, str, str]: """the id, datacenter, and s3 endpoint of the network volume with this id or name.""" - volume_id, datacenter = _resolve_volume(reference) + volume_id, datacenter = _resolve_volume(reference, user_agent=user_agent) endpoint = endpoint_url or f"https://s3api-{datacenter.lower()}.runpod.io/" return volume_id, datacenter, endpoint @@ -962,13 +967,17 @@ def _make_transfer( progress: Optional[Progress], force: bool = False, location: Optional[Tuple[str, str, str]] = None, + *, + user_agent: str = USER_AGENT, ) -> _Transfer: """a transfer to the network volume with this id or name. backs NetworkVolume.upload and NetworkVolume.download. `location` is a result of `_locate` the caller already holds.""" import boto3 # pylint: disable=import-outside-toplevel - volume_id, datacenter, endpoint = location or _locate(volume, endpoint_url) + volume_id, datacenter, endpoint = location or _locate( + volume, endpoint_url, user_agent=user_agent + ) access_key, secret_key = _s3_keys(access_key, secret_key) concurrency = max(1, int(concurrency)) client = boto3.session.Session().client( @@ -977,7 +986,7 @@ def _make_transfer( endpoint_url=endpoint, aws_access_key_id=access_key, aws_secret_access_key=secret_key, - config=_client_config(concurrency), + config=_client_config(concurrency, user_agent=user_agent), ) transfer = _Transfer( client, volume_id, endpoint, concurrency, progress, force=force diff --git a/runpod/rp_cli/volume.py b/runpod/rp_cli/volume.py index ace49af28..bdefc85ea 100644 --- a/runpod/rp_cli/volume.py +++ b/runpod/rp_cli/volume.py @@ -556,12 +556,14 @@ def _transfer(ctx, action_label, method_name, reference, source, target, opts) - """the checks every transfer needs, in order: login, volume, s3 keys.""" from runpod.apps import NetworkVolume from runpod.apps.volume import VolumeError + from runpod.user_agent import USER_AGENT import runpod.network_volume_transfer as nvt _ensure_logged_in(ctx) reference = reference or _pick_volume() try: target_volume = NetworkVolume(reference, create=False) + target_volume._user_agent = USER_AGENT # the volume caches this lookup, so the transfer does not repeat it target_volume._s3_location(opts["endpoint_url"]) except (nvt.NetworkVolumeTransferError, VolumeError) as exc: diff --git a/tests/test_apps/test_user_agent.py b/tests/test_apps/test_user_agent.py index 6c467f5d1..4b1103c04 100644 --- a/tests/test_apps/test_user_agent.py +++ b/tests/test_apps/test_user_agent.py @@ -274,19 +274,27 @@ async def test_lightweight_json_client(server): assert server.requests[0][2]["Authorization"] == "Bearer test-key" -def test_volume_lookup(server): +@pytest.mark.parametrize("user_agent", [None, APPS_USER_AGENT]) +def test_volume_lookup(server, user_agent): from runpod.network_volume_transfer import _locate + server.user_agent = user_agent or USER_AGENT server.responses = [ { "networkVolumes": [{"id": "vol", "name": "data", "dataCenter": "EU-RO-1"}], "pagination": {"hasNextPage": False}, } ] - assert _locate("data")[0] == "vol" + assert ( + _locate( + "data", **({"user_agent": user_agent} if user_agent is not None else {}) + )[0] + == "vol" + ) -def test_s3_transfer_user_agent(): +@pytest.mark.parametrize("user_agent", [None, APPS_USER_AGENT]) +def test_s3_transfer_user_agent(user_agent): from runpod.network_volume_transfer import _make_transfer transfer = _make_transfer( @@ -297,6 +305,7 @@ def test_s3_transfer_user_agent(): 1, None, location=("vol", "EU-RO-1", "https://s3.example.test"), + **({"user_agent": user_agent} if user_agent is not None else {}), ) requests = [] @@ -308,11 +317,86 @@ def capture(request, **kwargs): try: with pytest.raises(RuntimeError, match="captured request"): transfer.client.head_object(Bucket="vol", Key="file") - assert APPS_USER_AGENT in requests[0].headers["User-Agent"].decode() + header = requests[0].headers["User-Agent"].decode() + assert (user_agent or USER_AGENT) in header + assert ("Integration/apps" in header.split()) == (user_agent is not None) finally: transfer.client.close() +@pytest.mark.parametrize("caller", ["apps", "cli"]) +@pytest.mark.parametrize("operation", ["upload", "download"]) +def test_volume_transfer_caller_attribution( + server, monkeypatch, tmp_path, caller, operation +): + import boto3.session + from click.testing import CliRunner + + from runpod.apps import NetworkVolume + from runpod.rp_cli.main import cli + + server.user_agent = APPS_USER_AGENT if caller == "apps" else USER_AGENT + server.responses = [ + { + "networkVolumes": [{"id": "vol", "name": "data", "dataCenter": "EU-RO-1"}], + "pagination": {"hasNextPage": False}, + } + ] + monkeypatch.setattr("runpod.rp_cli.update.start_background_check", lambda: None) + local = tmp_path / "model.bin" + local.write_bytes(b"model") + source, target = ( + (str(local), "model.bin") + if operation == "upload" + else ("model.bin", str(local)) + ) + headers = [] + clients = [] + real_client = boto3.session.Session.client + + def capture(request, **kwargs): + headers.append(request.headers["User-Agent"].decode()) + raise RuntimeError("captured s3 request") + + def client(session, *args, **kwargs): + result = real_client(session, *args, **kwargs) + result.meta.events.register("before-send.s3", capture) + clients.append(result) + return result + + monkeypatch.setattr(boto3.session.Session, "client", client) + try: + if caller == "apps": + with pytest.raises(RuntimeError, match="captured s3 request"): + getattr(NetworkVolume("data"), operation)( + source, target, access_key="test-access", secret_key="test-secret" + ) + else: + result = CliRunner().invoke( + cli, + [ + "volume", + operation, + "data", + source, + target, + "--s3-access-key", + "test-access", + "--s3-secret-key", + "test-secret", + ], + ) + assert result.exit_code == 1 + assert "captured s3 request" in result.output + assert len(server.requests) == 1 + assert len(headers) == 1 + assert server.user_agent in headers[0] + assert ("Integration/apps" in headers[0].split()) == (caller == "apps") + finally: + for client in clients: + client.close() + + async def test_plain_sdk_transports_keep_default_user_agent(server): server.user_agent = USER_AGENT server.responses = [{}, {}, {"data": {}}, {"data": {}}] diff --git a/tests/test_network_volume_transfer.py b/tests/test_network_volume_transfer.py index d18f1c473..58cc2cde1 100644 --- a/tests/test_network_volume_transfer.py +++ b/tests/test_network_volume_transfer.py @@ -691,7 +691,9 @@ def test_make_transfer_requires_s3_keys(monkeypatch): monkeypatch.delenv(nvt.SECRET_KEY_ENV, raising=False) monkeypatch.setenv("AWS_ACCESS_KEY_ID", "aws") monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "aws") - monkeypatch.setattr(nvt, "_resolve_volume", lambda volume: ("vol", "EU-RO-1")) + monkeypatch.setattr( + nvt, "_resolve_volume", lambda volume, **kwargs: ("vol", "EU-RO-1") + ) with pytest.raises(nvt.NetworkVolumeTransferError, match="RUNPOD_S3_ACCESS_KEY"): nvt._make_transfer("vol", None, None, None, 4, None) @@ -781,7 +783,7 @@ def make(*args, **kwargs): volume = runpod.NetworkVolume("models", create=False) located = [] - def locate(reference, endpoint_url=None): + def locate(reference, endpoint_url=None, **kwargs): located.append(reference) return ("vol123", "EU-RO-1", "https://s3api-eu-ro-1.runpod.io/") @@ -798,7 +800,9 @@ def locate(reference, endpoint_url=None): def test_network_volume_download_keeps_trailing_slash(tmp_path, monkeypatch): fake = FakeS3() fake.put("w.bin", b"w") - monkeypatch.setattr(nvt, "_locate", lambda *a: ("vol", "EU-RO-1", "https://x/")) + monkeypatch.setattr( + nvt, "_locate", lambda *a, **k: ("vol", "EU-RO-1", "https://x/") + ) monkeypatch.setattr( nvt, "_make_transfer", lambda *a, **k: _transfer(fake, tmp_path) ) @@ -1348,7 +1352,7 @@ def interactive(monkeypatch): def cli_preflight(monkeypatch): """Keep CLI presentation tests independent of the account and DNS.""" monkeypatch.setattr( - nvt, "_resolve_volume", lambda reference: (reference, "EU-RO-1") + nvt, "_resolve_volume", lambda reference, **kwargs: (reference, "EU-RO-1") ) @@ -1473,7 +1477,7 @@ def test_cli_prompts_for_keys_and_offers_to_save( def test_cli_checks_volume_before_asking_for_keys(monkeypatch, tmp_path, interactive): - def missing_volume(reference): + def missing_volume(reference, **kwargs): raise nvt.NetworkVolumeTransferError("no network volume named 'missing'") monkeypatch.setattr(nvt, "_resolve_volume", missing_volume) @@ -1588,7 +1592,7 @@ def test_cli_rejects_a_blank_volume_cleanly(tmp_path): def test_network_volume_caches_one_location_per_endpoint(monkeypatch): located = [] - def locate(reference, endpoint_url=None): + def locate(reference, endpoint_url=None, **kwargs): located.append(endpoint_url) return ("vol123", "EU-RO-1", endpoint_url or "https://default/") @@ -1627,7 +1631,7 @@ def test_cli_looks_up_the_volume_once(monkeypatch, tmp_path): resolved = [] built = [] - def resolve(reference): + def resolve(reference, **kwargs): resolved.append(reference) return "vol123", "EU-RO-1"