From 9ac948074f174707733f65e797e21ce30ee97b1b Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Wed, 2 Sep 2026 14:06:05 -0700 Subject: [PATCH 1/2] feat: migrate api wrapper to rest v2 --- README.md | 32 +- docs/api/handling_errors.md | 26 +- docs/api/queries.md | 38 +- .../{graphql_wrapper.py => rest_wrapper.py} | 21 +- runpod/api/__init__.py | 2 +- runpod/api/ctl_commands.py | 623 ++++++------- .../api/mutations/container_register_auth.py | 81 +- runpod/api/mutations/endpoints.py | 102 --- runpod/api/mutations/pods.py | 194 ----- runpod/api/mutations/templates.py | 88 -- runpod/api/mutations/user.py | 29 - runpod/api/queries/endpoints.py | 36 - runpod/api/queries/gpus.py | 45 - runpod/api/queries/pods.py | 89 -- runpod/api/rest.py | 92 ++ runpod/cli/groups/pod/commands.py | 2 +- runpod/cli/groups/project/functions.py | 5 +- runpod/cli/utils/rp_info.py | 27 +- runpod/error.py | 12 +- tests/test_api/test_ctl_commands.py | 815 +++++++++--------- .../test_mutation_container_registry_auth.py | 83 +- tests/test_api/test_mutation_endpoints.py | 45 - tests/test_api/test_mutations_pods.py | 89 -- tests/test_api/test_mutations_templates.py | 45 - tests/test_api/test_rest.py | 117 +++ .../test_cli_groups/test_pod_commands.py | 12 +- .../test_cli_groups/test_project_functions.py | 4 +- tests/test_cli/test_cli_utils/test_info.py | 12 +- tests/test_error.py | 21 +- 29 files changed, 1044 insertions(+), 1743 deletions(-) rename examples/{graphql_wrapper.py => rest_wrapper.py} (53%) delete mode 100644 runpod/api/mutations/endpoints.py delete mode 100644 runpod/api/mutations/pods.py delete mode 100644 runpod/api/mutations/templates.py delete mode 100644 runpod/api/mutations/user.py delete mode 100644 runpod/api/queries/endpoints.py delete mode 100644 runpod/api/queries/gpus.py delete mode 100644 runpod/api/queries/pods.py create mode 100644 runpod/api/rest.py delete mode 100644 tests/test_api/test_mutation_endpoints.py delete mode 100644 tests/test_api/test_mutations_pods.py delete mode 100644 tests/test_api/test_mutations_templates.py create mode 100644 tests/test_api/test_rest.py diff --git a/README.md b/README.md index 6ad9c6669..0effe146a 100644 --- a/README.md +++ b/README.md @@ -22,7 +22,7 @@ Welcome to the official Python library for Runpod API & SDK. - [⚡ | Serverless Worker (SDK)](#--serverless-worker-sdk) - [Quick Start](#quick-start) - [Local Test Worker](#local-test-worker) -- [📚 | API Language Library (GraphQL Wrapper)](#--api-language-library-graphql-wrapper) +- [📚 | REST API v2 Wrapper](#--rest-api-v2-wrapper) - [Endpoints](#endpoints) - [GPU Cloud (Pods)](#gpu-cloud-pods) - [📁 | Directory](#--directory) @@ -162,9 +162,9 @@ with VolumeCache(dirs=["/root/.cache/huggingface"]): See [Network-Volume Warm Cache](https://github.com/runpod/runpod-python/blob/main/docs/serverless/volume_cache.md) documentation for configuration and details. -## 📚 | API Language Library (GraphQL Wrapper) +## 📚 | REST API v2 Wrapper -When interacting with the Runpod API you can use this library to make requests to the API. +Use the API wrapper to manage Runpod resources through REST API v2. ```python import runpod @@ -281,26 +281,26 @@ import runpod runpod.api_key = "your_runpod_api_key_found_under_settings" -# Get all my pods +# get all my pods pods = runpod.get_pods() -# Get a specific pod -pod = runpod.get_pod(pod.id) +# get a specific pod +pod = runpod.get_pod(pods[0]["id"]) -# Create a pod with GPU -pod = runpod.create_pod("test", "runpod/stack", "NVIDIA GeForce RTX 3070") +# create a pod with a gpu +pod = runpod.create_pod("test", "runpod/stack", "NVIDIA GeForce RTX 4090") -# Create a pod with CPU +# create a pod with a cpu pod = runpod.create_pod("test", "runpod/stack", instance_id="cpu3c-2-4") -# Stop the pod -runpod.stop_pod(pod.id) +# stop the pod +runpod.stop_pod(pod["id"]) -# Resume the pod -runpod.resume_pod(pod.id) +# resume the pod +runpod.resume_pod(pod["id"], 1) -# Terminate the pod -runpod.terminate_pod(pod.id) +# terminate the pod +runpod.terminate_pod(pod["id"]) ``` ## 📁 | Directory @@ -310,7 +310,7 @@ runpod.terminate_pod(pod.id) ├── docs # Documentation ├── examples # Examples ├── runpod # Package source code -│ ├── api_wrapper # Language library - API (GraphQL) +│ ├── api # rest api v2 wrapper │ ├── cli # Command Line Interface Functions │ ├── endpoint # Language library - Endpoints │ └── serverless # SDK - Serverless Worker diff --git a/docs/api/handling_errors.md b/docs/api/handling_errors.md index 055e11eda..e8b1a28b1 100644 --- a/docs/api/handling_errors.md +++ b/docs/api/handling_errors.md @@ -1,10 +1,28 @@ -# Handling Errors +# Handling API errors -```Python +Authentication failures raise `AuthenticationError`: + +```python import runpod try: - # Use runpod to make a request + runpod.get_pods() except runpod.error.AuthenticationError as err: - # Authentication with the API failed + print(err) +``` + +REST API problem responses raise `QueryError`. The exception includes the HTTP +status code, request method and path, and request-validation errors when present: + +```python +try: + runpod.create_pod( + "training", + "runpod/pytorch:1.0.2-cu1281-torch280-ubuntu2404", + "NVIDIA GeForce RTX 4090", + ) +except runpod.error.QueryError as err: + print(err.status_code) + print(err.query) + print(err.errors) ``` diff --git a/docs/api/queries.md b/docs/api/queries.md index dc26b350e..b576c1e4f 100644 --- a/docs/api/queries.md +++ b/docs/api/queries.md @@ -11,25 +11,43 @@ for gpu in gpus: print(gpu) ``` -### get_gpus Output +### get_gpus output -```json -{'id': 'NVIDIA A100 80GB PCIe', 'displayName': 'A100 80GB', 'memoryInGb': 80} -{'id': 'NVIDIA A100-SXM4-80GB', 'displayName': 'A100 SXM 80GB', 'memoryInGb': 80} -{'id': 'NVIDIA A30', 'displayName': 'A30', 'memoryInGb': 24} +```python +{ + "id": "NVIDIA GeForce RTX 4090", + "name": "RTX 4090", + "pool": "ADA_24", + "manufacturer": "NVIDIA", + "memory": 24, + "secure": True, + "community": True, + "price": {"secure": 0.44, "community": 0.31, "serverless": 1.1}, + "maxCount": {"secure": 8, "community": 4}, +} ``` ## get_gpu ```python -gpu_id = "NVIDIA A100 80GB PCIe" -gpu = runpod.get_gpu(gpu_id) +gpu_id = "NVIDIA GeForce RTX 4090" +gpu = runpod.get_gpu(gpu_id, gpu_quantity=2) print(gpu) ``` -### get_gpu Output +`get_gpu` requests pod availability for the requested GPU count. + +### get_gpu output -```json -{'id': 'NVIDIA A100 80GB PCIe', 'displayName': 'A100 80GB', 'memoryInGb': 80, 'secureCloud': True, 'communityCloud': True, 'lowestPrice': {'minimumBidPrice': 1.158, 'uninterruptablePrice': 1.89}} +```python +{ + "id": "NVIDIA GeForce RTX 4090", + "name": "RTX 4090", + "memory": 24, + "availability": "HIGH", + "dataCenters": [ + {"id": "US-KS-2", "name": "US Kansas 2", "availability": "HIGH"} + ], +} ``` diff --git a/examples/graphql_wrapper.py b/examples/rest_wrapper.py similarity index 53% rename from examples/graphql_wrapper.py rename to examples/rest_wrapper.py index 218c3c1ab..2f7010539 100644 --- a/examples/graphql_wrapper.py +++ b/examples/rest_wrapper.py @@ -1,6 +1,4 @@ -"""' -GraphQL wrapper for the Runpod API -""" +"""REST v2 wrapper for the Runpod API.""" import time @@ -8,37 +6,32 @@ runpod.api_key = "YOUR_RUNPOD_API_KEY" -# Get all GPUs gpus = runpod.get_gpus() print(gpus) -# Get a specific GPU -gpu = runpod.get_gpu("NVIDIA GeForce RTX 3070") +gpu = runpod.get_gpu("NVIDIA GeForce RTX 4090") print(gpu) -# Create a pod -pod = runpod.create_pod("test", "runpod/stack", "NVIDIA GeForce RTX 3070") +pod = runpod.create_pod( + "test", + "runpod/pytorch:1.0.2-cu1281-torch280-ubuntu2404", + "NVIDIA GeForce RTX 4090", +) print(pod) -# Pause while the pod is being created print("Waiting for pod to be created...") time.sleep(10) -# Stop a pod pod = runpod.stop_pod(pod["id"]) print(pod) -# Pause while the pod is being stopped print("Waiting for pod to be stopped...") time.sleep(10) -# Resume a pod pod = runpod.resume_pod(pod["id"], 1) print(pod) -# Pause while the pod is being resumed print("Waiting for pod to be resumed...") time.sleep(10) -# Terminate a pod runpod.terminate_pod(pod["id"]) diff --git a/runpod/api/__init__.py b/runpod/api/__init__.py index 90bf6dc72..ffea4b4b3 100644 --- a/runpod/api/__init__.py +++ b/runpod/api/__init__.py @@ -1 +1 @@ -""" Allows api_wrapper to be imported as a module.""" +"""Runpod API wrapper.""" diff --git a/runpod/api/ctl_commands.py b/runpod/api/ctl_commands.py index df5500d05..5af4293a8 100644 --- a/runpod/api/ctl_commands.py +++ b/runpod/api/ctl_commands.py @@ -1,113 +1,117 @@ -""" -Runpod | API Wrapper | CTL Commands -""" +"""Runpod API wrapper commands.""" # pylint: disable=too-many-arguments,too-many-locals -from typing import Optional +from typing import Any, Iterable, Optional +from urllib.parse import quote + +from runpod import error from .graphql import run_graphql_query from .mutations import container_register_auth as container_register_auth_mutations -from .mutations import endpoints as endpoint_mutations -from .mutations import pods as pod_mutations -from .mutations import templates as template_mutations -from .mutations import user as user_mutations -from .queries import endpoints as endpoint_queries -from .queries import gpus -from .queries import pods as pod_queries from .queries import user as user_queries +from .rest import run_rest_request -def get_user(api_key: Optional[str] = None) -> dict: - """ - Get the current user with optional API key override. - - Args: - api_key: Optional API key to use for this query. - """ - raw_response = run_graphql_query(user_queries.QUERY_USER, api_key=api_key) - cleaned_return = raw_response["data"]["myself"] - return cleaned_return +def _path_segment(value: str) -> str: + return quote(value, safe="") -def update_user_settings(pubkey: str, api_key: Optional[str] = None) -> dict: - """ - Update the current user +def _split_values(value: Optional[Iterable[Any] | str]) -> list[str]: + if value is None: + return [] + if isinstance(value, str): + values = value.split(",") + else: + values = value + return [str(item).strip() for item in values if str(item).strip()] - Args: - pubkey: the public key of the user - api_key: Optional API key to use for this query. - """ - raw_response = run_graphql_query( - user_mutations.generate_user_mutation(pubkey), - api_key=api_key - ) - cleaned_return = raw_response["data"]["updateUserSettings"] - return cleaned_return - - -def get_gpus(api_key: Optional[str] = None) -> dict: - """ - Get all GPU types - - Args: - api_key: Optional API key to use for this query. - """ - raw_response = run_graphql_query(gpus.QUERY_GPU_TYPES, api_key=api_key) - cleaned_return = raw_response["data"]["gpuTypes"] - return cleaned_return - - -def get_gpu(gpu_id: str, gpu_quantity: int = 1, api_key: Optional[str] = None): - """ - Get a specific GPU type - - Args: - gpu_id: the id of the gpu - gpu_quantity: how many of the gpu should be returned - api_key: Optional API key to use for this query. - """ - raw_response = run_graphql_query( - gpus.generate_gpu_query(gpu_id, gpu_quantity), - api_key=api_key - ) - cleaned_return = raw_response["data"]["gpuTypes"] +def _environment(env: Optional[dict]) -> dict[str, str]: + return {str(key): str(value) for key, value in (env or {}).items()} - if len(cleaned_return) < 1: - raise ValueError( - "No GPU found with the specified ID, " - "run runpod.get_gpus() to get a list of all GPUs" - ) - return cleaned_return[0] +def _cpu_config(instance_id: Optional[str], min_vcpu_count: int) -> dict[str, Any]: + if not instance_id: + raise ValueError("instance_id must be provided for CPU pods") + parts = instance_id.split("-") + vcpu_count = max(2, min_vcpu_count) + if len(parts) > 1: + try: + vcpu_count = int(parts[1]) + except ValueError as exc: + raise ValueError( + "instance_id must use the format --" + ) from exc -def get_pods(api_key: Optional[str] = None) -> dict: - """ - Get all pods - - Args: - api_key: Optional API key to use for this query. - """ - raw_return = run_graphql_query(pod_queries.QUERY_POD, api_key=api_key) - cleaned_return = raw_return["data"]["myself"]["pods"] - return cleaned_return + return {"id": parts[0], "vcpuCount": vcpu_count} -def get_pod(pod_id: str, api_key: Optional[str] = None): - """ - Get a specific pod +def get_user(api_key: Optional[str] = None) -> dict: + """Get the current user.""" + raw_response = run_graphql_query(user_queries.QUERY_USER, api_key=api_key) + return raw_response["data"]["myself"] - Args: - pod_id: the id of the pod - api_key: Optional API key to use for this query. - """ - raw_response = run_graphql_query( - pod_queries.generate_pod_query(pod_id), - api_key=api_key + +def update_user_settings(pubkey: str, api_key: Optional[str] = None) -> dict: + """Replace the current user's SSH public keys.""" + keys = [key.strip() for key in pubkey.splitlines() if key.strip()] + run_rest_request( + "PUT", + "/v2/account/ssh-keys", + api_key=api_key, + json={"keys": keys}, ) - return raw_response["data"]["pod"] + return get_user(api_key=api_key) + + +def get_gpus(api_key: Optional[str] = None) -> list[dict]: + """Get all GPU types.""" + response = run_rest_request("GET", "/v2/catalog/gpus", api_key=api_key) + return response["gpus"] + + +def get_gpu( + gpu_id: str, gpu_quantity: int = 1, api_key: Optional[str] = None +) -> dict: + """Get a GPU type and its pod availability.""" + try: + return run_rest_request( + "GET", + f"/v2/catalog/gpus/{_path_segment(gpu_id)}", + api_key=api_key, + params={ + "include": "AVAILABILITY", + "product": "POD", + "count": gpu_quantity, + }, + ) + except error.QueryError as exc: + if exc.status_code == 404: + raise ValueError( + "No GPU found with the specified ID, " + "run runpod.get_gpus() to get a list of all GPUs" + ) from exc + raise + + +def get_pods(api_key: Optional[str] = None) -> list[dict]: + """Get all standalone pods.""" + response = run_rest_request("GET", "/v2/pods", api_key=api_key) + return response["pods"] + + +def get_pod(pod_id: str, api_key: Optional[str] = None) -> Optional[dict]: + """Get a pod by ID.""" + try: + return run_rest_request( + "GET", f"/v2/pods/{_path_segment(pod_id)}", api_key=api_key + ) + except error.QueryError as exc: + if exc.status_code == 404: + return None + raise def create_pod( @@ -131,144 +135,99 @@ def create_pod( template_id: Optional[str] = None, network_volume_id: Optional[str] = None, allowed_cuda_versions: Optional[list] = None, - min_download = None, - min_upload = None, + min_download=None, + min_upload=None, instance_id: Optional[str] = None, ) -> dict: - """ - Create a pod - - :param name: the name of the pod - :param image_name: the name of the docker image to be used by the pod - :param gpu_type_id: the gpu type wanted by the pod (retrievable by get_gpus). If None, creates a CPU-only pod - :param cloud_type: if secure cloud, community cloud or all is wanted - :param data_center_id: the id of the data center - :param country_code: the code for country to start the pod in - :param gpu_count: how many gpus should be attached to the pod (ignored for CPU-only pods) - :param volume_in_gb: how big should the pod volume be - :param ports: the ports to open in the pod, example format - "8888/http,666/tcp" - :param volume_mount_path: where to mount the volume? - :param env: the environment variables to inject into the pod, - for example {EXAMPLE_VAR:"example_value", EXAMPLE_VAR2:"example_value 2"}, will - inject EXAMPLE_VAR and EXAMPLE_VAR2 into the pod with the mentioned values - :param template_id: the id of the template to use for the pod - :param min_download: minimum download speed in Mbps - :param min_upload: minimum upload speed in Mbps - :param instance_id: the id of a specific instance to deploy to (for CPU pods) - :example: - - >>> # Create GPU pod - >>> pod_id = runpod.create_pod("test", "runpod/stack", "NVIDIA GeForce RTX 3070") - >>> # Create CPU pod - >>> pod_id = runpod.create_pod("test", "runpod/stack") - >>> # Create CPU pod on specific instance - >>> pod_id = runpod.create_pod("test", "runpod/stack", instance_id="cpu3c-2-4") - """ - # Input Validation - + """Create a GPU or CPU pod.""" if not image_name and not template_id: raise ValueError("Either image_name or template_id must be provided") - - if gpu_type_id is not None: - get_gpu(gpu_type_id) # Check if GPU exists, will raise ValueError if not. - if cloud_type not in ["ALL", "COMMUNITY", "SECURE"]: + if cloud_type not in {"ALL", "COMMUNITY", "SECURE"}: raise ValueError("cloud_type must be one of ALL, COMMUNITY or SECURE") - if network_volume_id and data_center_id is None: - user_info = get_user() - for network_volume in user_info["networkVolumes"]: - if network_volume["id"] == network_volume_id: - data_center_id = network_volume["dataCenterId"] - break - - if container_disk_in_gb is None and template_id is None: - container_disk_in_gb = 10 - - raw_response = run_graphql_query( - pod_mutations.generate_pod_deployment_mutation( - name, - image_name, - gpu_type_id, - cloud_type, - support_public_ip, - start_ssh, - data_center_id, - country_code, - gpu_count if gpu_type_id is not None else None, - volume_in_gb, - container_disk_in_gb, - min_vcpu_count, - min_memory_in_gb, - docker_args, - ports, - volume_mount_path, - env, - template_id, - network_volume_id, - allowed_cuda_versions, - min_download, - min_upload, - instance_id, - ) - ) - - if gpu_type_id is not None: - cleaned_response = raw_response["data"]["podFindAndDeployOnDemand"] + unsupported = [] + if support_public_ip is not True: + unsupported.append("support_public_ip") + if country_code is not None: + unsupported.append("country_code") + if gpu_type_id and min_vcpu_count != 1: + unsupported.append("min_vcpu_count") + if min_memory_in_gb != 1: + unsupported.append("min_memory_in_gb") + if min_download is not None: + unsupported.append("min_download") + if min_upload is not None: + unsupported.append("min_upload") + if unsupported: + fields = ", ".join(unsupported) + raise ValueError(f"REST API v2 does not support: {fields}") + + body: dict[str, Any] = { + "name": name, + "args": docker_args, + "startSsh": start_ssh, + } + if image_name: + body["image"] = image_name + if template_id: + body["templateId"] = template_id + if cloud_type != "ALL": + body["cloud"] = cloud_type + if data_center_id: + body["dataCenterIds"] = [data_center_id] + if container_disk_in_gb is not None: + body["disk"] = container_disk_in_gb + elif not template_id: + body["disk"] = 10 + if ports is not None: + body["ports"] = _split_values(ports) + if env is not None: + body["env"] = _environment(env) + + if network_volume_id: + body["mounts"] = { + "network": [ + {"volumeId": network_volume_id, "path": volume_mount_path} + ] + } + elif volume_in_gb: + body["mounts"] = { + "persistent": {"size": volume_in_gb, "path": volume_mount_path} + } + + if gpu_type_id: + gpu: dict[str, Any] = {"id": gpu_type_id, "count": gpu_count} + if allowed_cuda_versions is not None: + gpu["allowedCudaVersions"] = _split_values(allowed_cuda_versions) + body["gpu"] = gpu else: - cleaned_response = raw_response["data"]["deployCpuPod"] - return cleaned_response - - -def stop_pod(pod_id: str): - """ - Stop a pod - - :param pod_id: the id of the pod - - :example: - - >>> pod_id = runpod.create_pod("test", "runpod/stack", "NVIDIA GeForce RTX 3070") - >>> runpod.stop_pod(pod_id) - """ - raw_response = run_graphql_query(pod_mutations.generate_pod_stop_mutation(pod_id)) - - cleaned_response = raw_response["data"]["podStop"] - return cleaned_response - + body["cpu"] = _cpu_config(instance_id, min_vcpu_count) -def resume_pod(pod_id: str, gpu_count: int): - """ - Resume a pod + return run_rest_request("POST", "/v2/pods", json=body) - :param pod_id: the id of the pod - :param gpu_count: the number of GPUs to attach to the pod - :example: - - >>> pod_id = runpod.create_pod("test", "runpod/stack", "NVIDIA GeForce RTX 3070") - >>> runpod.stop_pod(pod_id) - >>> runpod.resume_pod(pod_id) - """ - raw_response = run_graphql_query( - pod_mutations.generate_pod_resume_mutation(pod_id, gpu_count) +def stop_pod(pod_id: str) -> dict: + """Stop a pod.""" + return run_rest_request( + "POST", + f"/v2/pods/{_path_segment(pod_id)}/action", + json={"action": "stop"}, ) - cleaned_response = raw_response["data"]["podResume"] - return cleaned_response - -def terminate_pod(pod_id: str): - """ - Terminate a pod - - :param pod_id: the id of the pod +def resume_pod(pod_id: str, gpu_count: int) -> dict: + """Start a stopped pod.""" + _ = gpu_count + return run_rest_request( + "POST", + f"/v2/pods/{_path_segment(pod_id)}/action", + json={"action": "start"}, + ) - :example: - >>> pod_id = runpod.create_pod("test", "runpod/stack", "NVIDIA GeForce RTX 3070") - >>> runpod.terminate_pod(pod_id) - """ - run_graphql_query(pod_mutations.generate_pod_terminate_mutation(pod_id)) +def terminate_pod(pod_id: str) -> None: + """Terminate a pod.""" + run_rest_request("DELETE", f"/v2/pods/{_path_segment(pod_id)}") def create_template( @@ -282,52 +241,37 @@ def create_template( env: dict = None, is_serverless: bool = False, registry_auth_id: str = None, -): - """ - Create a template - - :param name: the name of the template - :param image_name: the name of the docker image to be used by the template - :param docker_start_cmd: the command to start the docker container with - :param container_disk_in_gb: how big should the container disk be - :param volume_in_gb: how big should the volume be - :param ports: the ports to open in the pod, example format - "8888/http,666/tcp" - :param volume_mount_path: where to mount the volume? - :param env: the environment variables to inject into the pod, - for example {EXAMPLE_VAR:"example_value", EXAMPLE_VAR2:"example_value 2"}, will - inject EXAMPLE_VAR and EXAMPLE_VAR2 into the pod with the mentioned values - :param is_serverless: is the template serverless? - :param registry_auth_id: the id of the registry auth - - :example: - - >>> template_id = runpod.create_template("test", "runpod/stack", "python3 main.py") - """ - raw_response = run_graphql_query( - template_mutations.generate_pod_template( - name=name, - image_name=image_name, - docker_start_cmd=docker_start_cmd, - container_disk_in_gb=container_disk_in_gb, - volume_in_gb=volume_in_gb, - volume_mount_path=volume_mount_path, - ports=ports, - env=env, - is_serverless=is_serverless, - registry_auth_id=registry_auth_id, - ) - ) - - return raw_response["data"]["saveTemplate"] - - -def get_endpoints() -> dict: - """ - Get all endpoints - """ - raw_return = run_graphql_query(endpoint_queries.QUERY_ENDPOINT) - cleaned_return = raw_return["data"]["myself"]["endpoints"] - return cleaned_return +) -> dict: + """Create a pod or serverless template.""" + body: dict[str, Any] = { + "name": name, + "image": image_name, + "disk": container_disk_in_gb, + "serverless": is_serverless, + } + if docker_start_cmd is not None: + body["args"] = docker_start_cmd + if volume_in_gb is not None: + body["mounts"] = { + "persistent": { + "size": volume_in_gb, + "path": volume_mount_path or "/workspace", + } + } + if ports is not None: + body["ports"] = _split_values(ports) + if env is not None: + body["env"] = _environment(env) + if registry_auth_id is not None: + body["registry"] = registry_auth_id + + return run_rest_request("POST", "/v2/templates", json=body) + + +def get_endpoints() -> list[dict]: + """Get all serverless endpoints.""" + response = run_rest_request("GET", "/v2/serverless") + return response["endpoints"] def create_endpoint( @@ -344,98 +288,72 @@ def create_endpoint( flashboot=False, allowed_cuda_versions: str = None, gpu_count: int = 1, -): - """ - Create an endpoint - - :param name: the name of the endpoint - :param template_id: the id of the template to use for the endpoint - :param gpu_ids: the ids of the GPUs to use for the endpoint - :param network_volume_id: the id of the network volume to use for the endpoint - :param locations: the locations to use for the endpoint - :param idle_timeout: the idle timeout for the endpoint - :param scaler_type: the scaler type for the endpoint - :param scaler_value: the scaler value for the endpoint - :param workers_min: the minimum number of workers for the endpoint - :param workers_max: the maximum number of workers for the endpoint - :param allowed_cuda_versions: Comma-separated list of allowed CUDA versions (e.g., ["12.4", "12.5"]). - :param gpu_count: the number of GPUs to use for the endpoint - - :example: - - >>> endpoint_id = runpod.create_endpoint("test", "template_id") - """ - raw_response = run_graphql_query( - endpoint_mutations.generate_endpoint_mutation( - name, - template_id, - gpu_ids, - network_volume_id, - locations, - idle_timeout, - scaler_type, - scaler_value, - workers_min, - workers_max, - flashboot, - allowed_cuda_versions, - gpu_count - ) - ) - - return raw_response["data"]["saveEndpoint"] - - -def update_endpoint_template(endpoint_id: str, template_id: str): - """ - Update an endpoint template - - :param endpoint_id: the id of the endpoint - :param template_id: the id of the template to use for the endpoint - - :example: - - >>> endpoint_id = runpod.update_endpoint_template("test", "template_id") - """ - raw_response = run_graphql_query( - endpoint_mutations.update_endpoint_template_mutation(endpoint_id, template_id) +) -> dict: + """Create a queue-based serverless endpoint.""" + scaler_type = { + "QUEUE_DELAY": "QUEUE_DELAY", + "REQUEST_COUNT": "REQUEST_COUNT", + "WORKER_COUNT": "REQUEST_COUNT", + }.get(scaler_type, scaler_type) + if scaler_type == "QUEUE_DELAY": + scaling = {"type": scaler_type, "queueDelay": scaler_value} + elif scaler_type == "REQUEST_COUNT": + scaling = {"type": scaler_type, "requestCount": scaler_value} + else: + raise ValueError("scaler_type must be QUEUE_DELAY or REQUEST_COUNT") + + gpu: dict[str, Any] = { + "pools": _split_values(gpu_ids), + "count": gpu_count, + } + if allowed_cuda_versions is not None: + gpu["allowedCudaVersions"] = _split_values(allowed_cuda_versions) + + workers = {"min": workers_min, "max": workers_max} + if scaler_type == "QUEUE_DELAY": + workers["idleTimeout"] = idle_timeout + + body: dict[str, Any] = { + "name": name, + "templateId": template_id, + "type": "QUEUE", + "gpu": gpu, + "workers": workers, + "scaling": scaling, + "flashboot": "FLASHBOOT" if flashboot else "OFF", + } + if network_volume_id: + body["networkVolumes"] = [network_volume_id] + if locations: + body["dataCenterIds"] = _split_values(locations) + + return run_rest_request("POST", "/v2/serverless", json=body) + + +def update_endpoint_template(endpoint_id: str, template_id: str) -> dict: + """Apply a serverless template to an endpoint.""" + return run_rest_request( + "PATCH", + f"/v2/serverless/{_path_segment(endpoint_id)}", + json={"templateId": template_id}, ) - return raw_response["data"]["updateEndpointTemplate"] - - -def create_container_registry_auth(name: str, username: str, password: str): - """ - Create a container registry authentication. - - Args: - name (str): The name of the container registry. - username (str): The username for authentication. - password (str): The password for authentication. - Returns: - dict: The response data containing the saved container registry authentication. - """ - raw_response = run_graphql_query( - container_register_auth_mutations.generate_container_registry_auth( - name, username, password - ) +def create_container_registry_auth( + name: str, username: str, password: str +) -> dict: + """Create a container registry credential.""" + return run_rest_request( + "POST", + "/v2/registries", + json={"name": name, "username": username, "password": password}, ) - return raw_response["data"]["saveRegistryAuth"] - - -def update_container_registry_auth(registry_auth_id: str, username: str, password: str): - """ - Update a container registry authentication. - Args: - registry_auth_id (str): The id of the container registry authentication - username (str): The username for authentication. - password (str): The password for authentication. - Returns: - dict: The response data containing the updated container registry authentication. - """ +def update_container_registry_auth( + registry_auth_id: str, username: str, password: str +) -> dict: + """Update a container registry credential.""" raw_response = run_graphql_query( container_register_auth_mutations.update_container_registry_auth( registry_auth_id, username, password @@ -444,16 +362,9 @@ def update_container_registry_auth(registry_auth_id: str, username: str, passwor return raw_response["data"]["updateRegistryAuth"] -def delete_container_registry_auth(registry_auth_id: str): - """ - Delete a container registry authentication. - - Args: - registry_auth_id (str): The id of the container registry authentication - """ - raw_response = run_graphql_query( - container_register_auth_mutations.delete_container_registry_auth( - registry_auth_id - ) +def delete_container_registry_auth(registry_auth_id: str) -> bool: + """Delete a container registry credential.""" + run_rest_request( + "DELETE", f"/v2/registries/{_path_segment(registry_auth_id)}" ) - return raw_response["data"]["deleteRegistryAuth"] + return True diff --git a/runpod/api/mutations/container_register_auth.py b/runpod/api/mutations/container_register_auth.py index c2eac9897..d4ac5f7ec 100644 --- a/runpod/api/mutations/container_register_auth.py +++ b/runpod/api/mutations/container_register_auth.py @@ -1,51 +1,18 @@ -""" Runpod | API Wrapper | Mutations | Container Registry Auth """ - - -def generate_container_registry_auth(name: str, username: str, password: str): - """ - Generate a GraphQL mutation string to save container registry authentication details. - - Args: - name (str): The name of the container registry. - username (str): The username for authentication. - password (str): The password for authentication. - - Returns: - str: The GraphQL mutation string. - """ - # Prepare the input dictionary - input_dict = {"name": name, "username": username, "password": password} - - # Convert the input dictionary to a string, properly formatted for GraphQL - input_str = ", ".join(f'{key}: "{value}"' for key, value in input_dict.items()) - - return f""" - mutation SaveRegistryAuth {{ - saveRegistryAuth(input: {{{input_str}}}) {{ - id - name - }} - }} - """ - - -def update_container_registry_auth(registry_auth_id: str, username: str, password: str): - """ - Generate a GraphQL mutation string to update registry authentication details. - - Args: - registry_auth_id (str): The id of the container registry authentication - username (str): The username for authentication. - password (str): The password for authentication. - - Returns: - str: The GraphQL mutation string. - """ - # Prepare the input dictionary - input_dict = {"id": registry_auth_id, "username": username, "password": password} - - # Convert the input dictionary to a string, properly formatted for GraphQL - input_str = ", ".join(f'{key}: "{value}"' for key, value in input_dict.items()) +"""GraphQL operations for container registry credentials.""" + + +def update_container_registry_auth( + registry_auth_id: str, username: str, password: str +) -> str: + """Build the registry credential update mutation.""" + input_dict = { + "id": registry_auth_id, + "username": username, + "password": password, + } + input_str = ", ".join( + f'{key}: "{value}"' for key, value in input_dict.items() + ) return f""" mutation UpdateRegistryAuth {{ @@ -55,21 +22,3 @@ def update_container_registry_auth(registry_auth_id: str, username: str, passwor }} }} """ - - -def delete_container_registry_auth(registry_auth_id: str): - """ - Generate a GraphQL mutation string to delete registry authentication details. - - Args: - registry_auth_id (str): The id of the container registry authentication - - Returns: - str: The GraphQL mutation string. - """ - - return f""" - mutation DeleteRegistryAuth {{ - deleteRegistryAuth(registryAuthId: "{registry_auth_id}") - }} - """ diff --git a/runpod/api/mutations/endpoints.py b/runpod/api/mutations/endpoints.py deleted file mode 100644 index ab7c09df7..000000000 --- a/runpod/api/mutations/endpoints.py +++ /dev/null @@ -1,102 +0,0 @@ -f"""Runpod | API Wrapper | Mutations | Endpoints""" - -# pylint: disable=too-many-arguments - - -def generate_endpoint_mutation( - name: str, - template_id: str, - gpu_ids: str = "AMPERE_16", - network_volume_id: str = None, - locations: str = None, - idle_timeout: int = 5, - scaler_type: str = "QUEUE_DELAY", - scaler_value: int = 4, - workers_min: int = 0, - workers_max: int = 3, - flashboot=False, - allowed_cuda_versions: str = None, - gpu_count: int = None, -): - """Generate a string for a GraphQL mutation to create a new endpoint.""" - input_fields = [] - - # ------------------------------ Required Fields ----------------------------- # - if flashboot: - input_fields.append('flashBootType: FLASHBOOT') - - input_fields.append(f'name: "{name}"') - input_fields.append(f'templateId: "{template_id}"') - input_fields.append(f'gpuIds: "{gpu_ids}"') - - # ------------------------------ Optional Fields ----------------------------- # - if network_volume_id is not None: - input_fields.append(f'networkVolumeId: "{network_volume_id}"') - else: - input_fields.append('networkVolumeId: ""') - - if locations is not None: - input_fields.append(f'locations: "{locations}"') - else: - input_fields.append('locations: ""') - - input_fields.append(f"idleTimeout: {idle_timeout}") - input_fields.append(f'scalerType: "{scaler_type}"') - input_fields.append(f"scalerValue: {scaler_value}") - input_fields.append(f"workersMin: {workers_min}") - input_fields.append(f"workersMax: {workers_max}") - - if allowed_cuda_versions: - input_fields.append(f'allowedCudaVersions: "{allowed_cuda_versions}"') - - if gpu_count is not None: - input_fields.append(f"gpuCount: {gpu_count}") - - # Format the input fields into a string - input_fields_string = ", ".join(input_fields) - - return f""" - mutation {{ - saveEndpoint( - input: {{ - {input_fields_string} - }} - ) {{ - id - name - templateId - gpuIds - networkVolumeId - locations - idleTimeout - scalerType - scalerValue - workersMin - workersMax - allowedCudaVersions - gpuCount - flashBootType - }} - }} - """ - - -def update_endpoint_template_mutation(endpoint_id: str, template_id: str): - """Generate a string for a GraphQL mutation to update an existing endpoint's template.""" - input_fields = [] - - # ------------------------------ Required Fields ----------------------------- # - input_fields.append(f'templateId: "{template_id}"') - input_fields.append(f'endpointId: "{endpoint_id}"') - - # Format the input fields into a string - input_fields_string = ", ".join(input_fields) - result = f""" - mutation {{ - updateEndpointTemplate(input: {{{input_fields_string}}}) {{ - id - templateId - }} - }} - """ - return result diff --git a/runpod/api/mutations/pods.py b/runpod/api/mutations/pods.py deleted file mode 100644 index 48409a606..000000000 --- a/runpod/api/mutations/pods.py +++ /dev/null @@ -1,194 +0,0 @@ -""" -Runpod | API Wrapper | Mutations | Pods -""" - -# pylint: disable=too-many-arguments, too-many-locals, too-many-branches - -from typing import List, Optional - - -def generate_pod_deployment_mutation( - name: str, - image_name: str, - gpu_type_id: Optional[str] = None, - cloud_type: str = "ALL", - support_public_ip: bool = True, - start_ssh: bool = True, - data_center_id: Optional[str] = None, - country_code: Optional[str] = None, - gpu_count: Optional[int] = None, - volume_in_gb: Optional[int] = None, - container_disk_in_gb: Optional[int] = None, - min_vcpu_count: Optional[int] = None, - min_memory_in_gb: Optional[int] = None, - docker_args: Optional[str] = None, - ports: Optional[str] = None, - volume_mount_path: Optional[str] = None, - env: Optional[dict] = None, - template_id: Optional[str] = None, - network_volume_id: Optional[str] = None, - allowed_cuda_versions: Optional[List[str]] = None, - min_download: Optional[int] = None, - min_upload: Optional[int] = None, - instance_id: Optional[str] = None, -) -> str: - """ - Generates a mutation to deploy a pod on demand. - - Args: - name: Name of the pod - image_name: Docker image name - gpu_type_id: GPU type ID for GPU pods, None for CPU pods - cloud_type: Cloud type (ALL, COMMUNITY, or SECURE) - support_public_ip: Whether to support public IP - start_ssh: Whether to start SSH service - data_center_id: Data center ID - country_code: Country code for pod location - gpu_count: Number of GPUs (for GPU pods) - volume_in_gb: Volume size in GB - container_disk_in_gb: Container disk size in GB - min_vcpu_count: Minimum vCPU count - min_memory_in_gb: Minimum memory in GB - docker_args: Docker arguments - ports: Port mappings (e.g. "8080/tcp,22/tcp") - volume_mount_path: Volume mount path - env: Environment variables dict - template_id: Template ID - network_volume_id: Network volume ID - allowed_cuda_versions: List of allowed CUDA versions - min_download: Minimum download speed in Mbps - min_upload: Minimum upload speed in Mbps - instance_id: Instance ID for CPU pods - - Returns: - str: GraphQL mutation string - """ - input_fields = [] - - # Required Fields - input_fields.extend([ - f'name: "{name}"', - f'imageName: "{image_name}"', - f"cloudType: {cloud_type}" - ]) - - if start_ssh: - input_fields.append("startSsh: true") - - # GPU Pod Fields - if gpu_type_id is not None: - input_fields.append(f'gpuTypeId: "{gpu_type_id}"') - input_fields.append(f"supportPublicIp: {str(support_public_ip).lower()}") - - if gpu_count is not None: - input_fields.append(f"gpuCount: {gpu_count}") - if volume_in_gb is not None: - input_fields.append(f"volumeInGb: {volume_in_gb}") - if min_vcpu_count is not None: - input_fields.append(f"minVcpuCount: {min_vcpu_count}") - if min_memory_in_gb is not None: - input_fields.append(f"minMemoryInGb: {min_memory_in_gb}") - - if allowed_cuda_versions is not None: - cuda_versions = ", ".join(f'"{v}"' for v in allowed_cuda_versions) - input_fields.append(f"allowedCudaVersions: [{cuda_versions}]") - - # CPU Pod Fields - else: - if instance_id is not None: - input_fields.append(f'instanceId: "{instance_id}"') - template_id = template_id or "runpod-ubuntu" - - # Optional Fields - if data_center_id is not None: - input_fields.append(f'dataCenterId: "{data_center_id}"') - else: - input_fields.append("dataCenterId: null") - - if docker_args is not None: - input_fields.append(f'dockerArgs: "{docker_args}"') - if country_code is not None: - input_fields.append(f'countryCode: "{country_code}"') - if container_disk_in_gb is not None: - input_fields.append(f"containerDiskInGb: {container_disk_in_gb}") - if ports is not None: - input_fields.append(f'ports: "{ports.replace(" ", "")}"') - if volume_mount_path is not None: - input_fields.append(f'volumeMountPath: "{volume_mount_path}"') - if env is not None: - env_items = [f'{{ key: "{k}", value: "{v}" }}' for k, v in env.items()] - input_fields.append(f"env: [{', '.join(env_items)}]") - if template_id is not None: - input_fields.append(f'templateId: "{template_id}"') - if network_volume_id is not None: - input_fields.append(f'networkVolumeId: "{network_volume_id}"') - if min_download is not None: - input_fields.append(f'minDownload: {min_download}') - if min_upload is not None: - input_fields.append(f'minUpload: {min_upload}') - - mutation_type = "podFindAndDeployOnDemand" if gpu_type_id else "deployCpuPod" - input_string = ", ".join(input_fields) - - return f""" - mutation {{ - {mutation_type}( - input: {{ - {input_string} - }} - ) {{ - id - imageName - env - machineId - machine {{ - podHostId - }} - }} - }} - """ - - -def generate_pod_stop_mutation(pod_id: str) -> str: - """ - Generates a mutation to stop a pod. - """ - return f""" - mutation {{ - podStop(input: {{ podId: "{pod_id}" }}) {{ - id - desiredStatus - }} - }} - """ - - -def generate_pod_resume_mutation(pod_id: str, gpu_count: int) -> str: - """ - Generates a mutation to resume a pod. - """ - return f""" - mutation {{ - podResume(input: {{ podId: "{pod_id}", gpuCount: {gpu_count} }}) {{ - id - desiredStatus - imageName - env - machineId - machine {{ - podHostId - }} - }} - }} - """ - - -def generate_pod_terminate_mutation(pod_id: str) -> str: - """ - Generates a mutation to terminate a pod. - """ - return f""" - mutation {{ - podTerminate(input: {{ podId: "{pod_id}" }}) - }} - """ diff --git a/runpod/api/mutations/templates.py b/runpod/api/mutations/templates.py deleted file mode 100644 index 094b49bb3..000000000 --- a/runpod/api/mutations/templates.py +++ /dev/null @@ -1,88 +0,0 @@ -""" Runpod | API Wrapper | Mutations | Templates """ - -# pylint: disable=too-many-arguments, too-many-branches - - -def generate_pod_template( - name: str, - image_name: str, - docker_start_cmd: str = None, - container_disk_in_gb: int = 10, - volume_in_gb: int = None, - volume_mount_path: str = None, - ports: str = None, - env: dict = None, - is_serverless: bool = False, - registry_auth_id: str = None, -): - """Generate a string for a GraphQL mutation to create a new pod template.""" - input_fields = [f'name: "{name}"', f'imageName: "{image_name}"'] - - # ------------------------------ Optional Fields ----------------------------- # - if docker_start_cmd is not None: - docker_start_cmd = docker_start_cmd.replace('"', '\\"') - input_fields.append(f'dockerArgs: "{docker_start_cmd}"') - else: - input_fields.append('dockerArgs: ""') - - input_fields.append(f"containerDiskInGb: {container_disk_in_gb}") - - if volume_in_gb is not None: - input_fields.append(f"volumeInGb: {volume_in_gb}") - else: - input_fields.append("volumeInGb: 0") - - if volume_mount_path is not None: - input_fields.append(f'volumeMountPath: "{volume_mount_path}"') - - if ports is not None: - ports = ports.replace(" ", "") - input_fields.append(f'ports: "{ports}"') - else: - input_fields.append('ports: ""') - - if env is not None: - env_string = ", ".join( - [f'{{ key: "{key}", value: "{value}" }}' for key, value in env.items()] - ) - input_fields.append(f"env: [{env_string}]") - else: - input_fields.append("env: []") - - if is_serverless: - input_fields.append("isServerless: true") - else: - input_fields.append("isServerless: false") - - if registry_auth_id is not None: - input_fields.append(f'containerRegistryAuthId : "{registry_auth_id}"') - else: - input_fields.append('containerRegistryAuthId : ""') - - input_fields.extend(("startSsh: true", "isPublic: false", 'readme: ""')) - # Format the input fields into a string - input_fields_string = ", ".join(input_fields) - - return f""" - mutation {{ - saveTemplate( - input: {{ - {input_fields_string} - }} - ) {{ - id - name - imageName - dockerArgs - containerDiskInGb - volumeInGb - volumeMountPath - ports - env {{ - key - value - }} - isServerless - }} - }} - """ diff --git a/runpod/api/mutations/user.py b/runpod/api/mutations/user.py deleted file mode 100644 index 2d700a5ca..000000000 --- a/runpod/api/mutations/user.py +++ /dev/null @@ -1,29 +0,0 @@ -""" -Runpod | API | Mutations | User -""" - - -def generate_user_mutation(pubkey): - """' - Generates a mutation to edit a user. - """ - input_fields = [] - - escaped_pubkey = pubkey.replace("\n", "\\n") - input_fields.append(f'pubKey: "{escaped_pubkey}"') - - # Format input fields - input_string = ", ".join(input_fields) - - return f""" - mutation {{ - updateUserSettings( - input: {{ - {input_string} - }} - ) {{ - id - pubKey - }} - }} - """ diff --git a/runpod/api/queries/endpoints.py b/runpod/api/queries/endpoints.py deleted file mode 100644 index e262d760b..000000000 --- a/runpod/api/queries/endpoints.py +++ /dev/null @@ -1,36 +0,0 @@ -""" GraphQL queries for endpoints. """ - -QUERY_ENDPOINT = """ -query Query { - myself { - endpoints { - aiKey - gpuIds - id - idleTimeout - name - networkVolumeId - locations - scalerType - scalerValue - templateId - type - userId - version - workersMax - workersMin - workersStandby - gpuCount - env { - key - value - } - createdAt - networkVolume { - id - dataCenterId - } - } - } -} -""" diff --git a/runpod/api/queries/gpus.py b/runpod/api/queries/gpus.py deleted file mode 100644 index 726e8d7e3..000000000 --- a/runpod/api/queries/gpus.py +++ /dev/null @@ -1,45 +0,0 @@ -""" -Runpod | API | Queries | GPUs -""" - -QUERY_GPU_TYPES = """ -query GpuTypes { - gpuTypes { - id - displayName - memoryInGb - } -} -""" - - -def generate_gpu_query(gpu_id, gpu_count=1): - """ - Generate a query for a specific GPU type - """ - - return f""" - query GpuTypes {{ - gpuTypes(input: {{id: "{gpu_id}"}}) {{ - maxGpuCount - id - displayName - manufacturer - memoryInGb - cudaCores - secureCloud - communityCloud - securePrice - communityPrice - oneMonthPrice - threeMonthPrice - oneWeekPrice - communitySpotPrice - secureSpotPrice - lowestPrice(input: {{gpuCount: {gpu_count}}}) {{ - minimumBidPrice - uninterruptablePrice - }} - }} - }} - """ diff --git a/runpod/api/queries/pods.py b/runpod/api/queries/pods.py deleted file mode 100644 index 49690c82d..000000000 --- a/runpod/api/queries/pods.py +++ /dev/null @@ -1,89 +0,0 @@ -""" -Runpod | API Wrapper | Queries | GPUs -""" - -QUERY_POD = """ -query myPods { - myself { - pods { - id - containerDiskInGb - costPerHr - desiredStatus - dockerArgs - dockerId - env - gpuCount - imageName - lastStatusChange - machineId - memoryInGb - name - podType - port - ports - uptimeSeconds - vcpuCount - volumeInGb - volumeMountPath - runtime { - ports{ - ip - isIpPublic - privatePort - publicPort - type - } - } - machine { - gpuDisplayName - } - } - } -} -""" - - -def generate_pod_query(pod_id): - """ - Generate a query for a specific GPU type - """ - - return f""" - query pod {{ - pod(input: {{podId: "{pod_id}"}}) {{ - id - containerDiskInGb - costPerHr - desiredStatus - dockerArgs - dockerId - env - gpuCount - imageName - lastStatusChange - machineId - memoryInGb - name - podType - port - ports - uptimeSeconds - vcpuCount - volumeInGb - volumeMountPath - runtime {{ - ports {{ - ip - isIpPublic - privatePort - publicPort - type - }} - }} - machine {{ - gpuDisplayName - }} - }} - }} - """ diff --git a/runpod/api/rest.py b/runpod/api/rest.py new file mode 100644 index 000000000..abcf039cc --- /dev/null +++ b/runpod/api/rest.py @@ -0,0 +1,92 @@ +"""Runpod REST API transport.""" + +import os +from typing import Any, Mapping, Optional + +import requests + +from runpod import error +from runpod.user_agent import USER_AGENT + +HTTP_STATUS_NO_CONTENT = 204 +HTTP_STATUS_UNAUTHORIZED = 401 + + +def _resolve_api_key(api_key: Optional[str]) -> str: + from runpod import api_key as global_api_key # pylint: disable=import-outside-toplevel,cyclic-import + + effective_api_key = api_key or global_api_key + if not effective_api_key: + raise error.AuthenticationError("No API key provided") + return effective_api_key + + +def _build_url(path: str) -> str: + api_url_base = os.environ.get("RUNPOD_API_BASE_URL", "https://api.runpod.io") + return f"{api_url_base.rstrip('/')}/{path.lstrip('/')}" + + +def _build_headers(api_key: str) -> dict[str, str]: + return { + "Accept": "application/json", + "Content-Type": "application/json", + "User-Agent": USER_AGENT, + "Authorization": f"Bearer {api_key}", + } + + +def _response_json(response: requests.Response) -> dict[str, Any]: + try: + payload = response.json() + except ValueError: + return {} + return payload if isinstance(payload, dict) else {} + + +def _raise_for_error( + response: requests.Response, method: str, path: str +) -> None: + if response.status_code == HTTP_STATUS_UNAUTHORIZED: + raise error.AuthenticationError( + "Unauthorized request, please check your API key." + ) + + if response.status_code < 400: + return + + payload = _response_json(response) + message = payload.get("detail") or payload.get("title") + if not message: + message = response.text or f"Request failed with status {response.status_code}" + + raise error.QueryError( + str(message), + f"{method.upper()} {path}", + status_code=response.status_code, + errors=payload.get("errors"), + ) + + +def run_rest_request( + method: str, + path: str, + *, + api_key: Optional[str] = None, + params: Optional[Mapping[str, Any]] = None, + json: Optional[Mapping[str, Any]] = None, + timeout: float = 30, +) -> Optional[dict[str, Any]]: + """Send an authenticated request to the Runpod REST API.""" + response = requests.request( + method, + _build_url(path), + headers=_build_headers(_resolve_api_key(api_key)), + params=params, + json=json, + timeout=timeout, + ) + _raise_for_error(response, method, path) + + if response.status_code == HTTP_STATUS_NO_CONTENT or not response.content: + return None + return response.json() diff --git a/runpod/cli/groups/pod/commands.py b/runpod/cli/groups/pod/commands.py index e6f37c81c..5cd23ea89 100644 --- a/runpod/cli/groups/pod/commands.py +++ b/runpod/cli/groups/pod/commands.py @@ -25,7 +25,7 @@ def list_pods(): """ table = PrettyTable(["ID", "Name", "Status", "Image"]) for pod in get_pods(): - table.add_row((pod["id"], pod["name"], pod["desiredStatus"], pod["imageName"])) + table.add_row((pod["id"], pod["name"], pod["status"], pod["image"])) click.echo(table) diff --git a/runpod/cli/groups/project/functions.py b/runpod/cli/groups/project/functions.py index 4fe011573..0972685f1 100644 --- a/runpod/cli/groups/project/functions.py +++ b/runpod/cli/groups/project/functions.py @@ -60,10 +60,7 @@ def _launch_dev_pod(): sys.stdout.flush() # Wait for the pod to come online - while ( - new_pod.get("desiredStatus", None) != "RUNNING" - or new_pod.get("runtime") is None - ): + while new_pod.get("status") != "RUNNING" or new_pod.get("runtime") is None: new_pod = get_pod(new_pod["id"]) project_pod_id = new_pod["id"] diff --git a/runpod/cli/utils/rp_info.py b/runpod/cli/utils/rp_info.py index 545f79166..ecdd70117 100644 --- a/runpod/cli/utils/rp_info.py +++ b/runpod/cli/utils/rp_info.py @@ -15,27 +15,34 @@ def get_pod_ssh_ip_port(pod_id, timeout=300): start_time = time.time() pod_ip = None pod_port = None + status = None while time.time() - start_time < timeout and (pod_ip is None or pod_port is None): - pod = get_pod(pod_id) - desired_status = pod.get("desiredStatus", None) - runtime = pod.get("runtime", None) - - if desired_status == "RUNNING" and runtime and "ports" in pod["runtime"]: - for port in pod["runtime"]["ports"]: - if port["privatePort"] == 22: + pod = get_pod(pod_id) or {} + status = pod.get("status") + direct_ssh = (pod.get("ssh") or {}).get("direct") + + if status == "RUNNING" and direct_ssh: + pod_ip = direct_ssh["host"] + pod_port = int(direct_ssh["port"]) + break + + runtime = pod.get("runtime") or {} + if status == "RUNNING": + for port in runtime.get("ports", []): + if port["private"] == 22: pod_ip = port["ip"] - pod_port = int(port["publicPort"]) + pod_port = int(port["public"]) break time.sleep(1) - if desired_status != "RUNNING": + if status != "RUNNING": raise TimeoutError( f"Pod {pod_id} did not reach 'RUNNING' state within {timeout} seconds." ) - if runtime is None: + if pod_ip is None or pod_port is None: raise TimeoutError( f"Pod {pod_id} did not report runtime data within {timeout} seconds." ) diff --git a/runpod/error.py b/runpod/error.py index b715cb115..35a1a8f31 100644 --- a/runpod/error.py +++ b/runpod/error.py @@ -30,9 +30,17 @@ class AuthenticationError(RunPodError): class QueryError(RunPodError): """ - Raised when a GraphQL query fails + Raised when an API request fails """ - def __init__(self, message: Optional[str] = None, query: Optional[str] = None): + def __init__( + self, + message: Optional[str] = None, + query: Optional[str] = None, + status_code: Optional[int] = None, + errors: Optional[list[str]] = None, + ): super().__init__(message) self.query = query + self.status_code = status_code + self.errors = errors or [] diff --git a/tests/test_api/test_ctl_commands.py b/tests/test_api/test_ctl_commands.py index 472fb51ab..903a9228e 100644 --- a/tests/test_api/test_ctl_commands.py +++ b/tests/test_api/test_ctl_commands.py @@ -1,415 +1,430 @@ -""" Tests for ctl_commands.py """ +"""Tests for the API wrapper commands.""" -import unittest from unittest.mock import patch +import pytest + from runpod.api import ctl_commands +from runpod.error import QueryError + + +def test_get_user_uses_graphql(): + with patch( + "runpod.api.ctl_commands.run_graphql_query", + return_value={"data": {"myself": {"id": "user"}}}, + ) as request: + assert ctl_commands.get_user(api_key="key") == {"id": "user"} + + request.assert_called_once_with(ctl_commands.user_queries.QUERY_USER, api_key="key") + + +def test_update_user_settings_replaces_ssh_keys(): + with ( + patch("runpod.api.ctl_commands.run_rest_request") as request, + patch( + "runpod.api.ctl_commands.get_user", + return_value={"id": "user", "pubKey": "ssh-ed25519 key user"}, + ) as get_user, + ): + result = ctl_commands.update_user_settings( + "\nssh-ed25519 key user\n\n", api_key="key" + ) + + assert result == {"id": "user", "pubKey": "ssh-ed25519 key user"} + request.assert_called_once_with( + "PUT", + "/v2/account/ssh-keys", + api_key="key", + json={"keys": ["ssh-ed25519 key user"]}, + ) + get_user.assert_called_once_with(api_key="key") + + +def test_get_gpus_unwraps_response(): + gpus = [{"id": "NVIDIA A100", "name": "A100", "memory": 80}] + with patch( + "runpod.api.ctl_commands.run_rest_request", return_value={"gpus": gpus} + ) as request: + assert ctl_commands.get_gpus(api_key="key") == gpus + + request.assert_called_once_with("GET", "/v2/catalog/gpus", api_key="key") + + +def test_get_gpu_requests_pod_availability(): + gpu = {"id": "NVIDIA A100"} + with patch( + "runpod.api.ctl_commands.run_rest_request", return_value=gpu + ) as request: + assert ctl_commands.get_gpu("NVIDIA A100", 2, api_key="key") == gpu + + request.assert_called_once_with( + "GET", + "/v2/catalog/gpus/NVIDIA%20A100", + api_key="key", + params={"include": "AVAILABILITY", "product": "POD", "count": 2}, + ) + + +def test_get_gpu_converts_not_found_to_value_error(): + with ( + patch( + "runpod.api.ctl_commands.run_rest_request", + side_effect=QueryError("not found", status_code=404), + ), + pytest.raises(ValueError, match="No GPU found"), + ): + ctl_commands.get_gpu("missing") + + +def test_get_gpu_propagates_other_api_errors(): + with ( + patch( + "runpod.api.ctl_commands.run_rest_request", + side_effect=QueryError("forbidden", status_code=403), + ), + pytest.raises(QueryError, match="forbidden"), + ): + ctl_commands.get_gpu("NVIDIA A100") + + +def test_get_pods_unwraps_response(): + pods = [{"id": "pod"}] + with patch( + "runpod.api.ctl_commands.run_rest_request", return_value={"pods": pods} + ) as request: + assert ctl_commands.get_pods(api_key="key") == pods + + request.assert_called_once_with("GET", "/v2/pods", api_key="key") + + +def test_get_pod_escapes_id(): + with patch( + "runpod.api.ctl_commands.run_rest_request", return_value={"id": "pod/id"} + ) as request: + assert ctl_commands.get_pod("pod/id", api_key="key") == {"id": "pod/id"} + + request.assert_called_once_with("GET", "/v2/pods/pod%2Fid", api_key="key") + + +def test_get_pod_returns_none_when_not_found(): + with patch( + "runpod.api.ctl_commands.run_rest_request", + side_effect=QueryError("not found", status_code=404), + ): + assert ctl_commands.get_pod("missing") is None + + +def test_get_pod_propagates_other_api_errors(): + with ( + patch( + "runpod.api.ctl_commands.run_rest_request", + side_effect=QueryError("forbidden", status_code=403), + ), + pytest.raises(QueryError, match="forbidden"), + ): + ctl_commands.get_pod("pod") + + +def test_create_gpu_pod_translates_request(): + with patch( + "runpod.api.ctl_commands.run_rest_request", return_value={"id": "pod"} + ) as request: + result = ctl_commands.create_pod( + name="training", + image_name="runpod/pytorch:latest", + gpu_type_id="NVIDIA A100", + cloud_type="COMMUNITY", + start_ssh=True, + data_center_id="US-KS-2", + gpu_count=2, + volume_in_gb=20, + container_disk_in_gb=50, + docker_args="python main.py", + ports="8888/http, 22/tcp", + volume_mount_path="/workspace", + env={"COUNT": 2}, + allowed_cuda_versions=["12.8", "12.6"], + ) + + assert result == {"id": "pod"} + request.assert_called_once_with( + "POST", + "/v2/pods", + json={ + "name": "training", + "image": "runpod/pytorch:latest", + "args": "python main.py", + "startSsh": True, + "cloud": "COMMUNITY", + "dataCenterIds": ["US-KS-2"], + "disk": 50, + "ports": ["8888/http", "22/tcp"], + "env": {"COUNT": "2"}, + "mounts": { + "persistent": {"size": 20, "path": "/workspace"} + }, + "gpu": { + "id": "NVIDIA A100", + "count": 2, + "allowedCudaVersions": ["12.8", "12.6"], + }, + }, + ) + + +def test_create_gpu_pod_with_network_volume_and_template(): + with patch( + "runpod.api.ctl_commands.run_rest_request", return_value={"id": "pod"} + ) as request: + ctl_commands.create_pod( + name="training", + template_id="template", + gpu_type_id="NVIDIA A100", + network_volume_id="volume", + ) + + body = request.call_args.kwargs["json"] + assert body["templateId"] == "template" + assert body["mounts"] == { + "network": [{"volumeId": "volume", "path": "/runpod-volume"}] + } + assert "image" not in body + assert "disk" not in body + assert "cloud" not in body + + +def test_create_cpu_pod_translates_instance_id(): + with patch( + "runpod.api.ctl_commands.run_rest_request", return_value={"id": "pod"} + ) as request: + ctl_commands.create_pod( + "cpu-pod", "python:3.11", instance_id="cpu3c-4-8" + ) + + assert request.call_args.kwargs["json"]["cpu"] == { + "id": "cpu3c", + "vcpuCount": 4, + } + + +def test_create_cpu_pod_requires_instance_id(): + with pytest.raises(ValueError, match="instance_id"): + ctl_commands.create_pod("cpu-pod", "python:3.11") + + +def test_create_cpu_pod_validates_instance_id(): + with pytest.raises(ValueError, match="format"): + ctl_commands.create_pod( + "cpu-pod", "python:3.11", instance_id="cpu3c-invalid" + ) + +def test_create_pod_validates_image_and_cloud(): + with pytest.raises(ValueError, match="Either image_name or template_id"): + ctl_commands.create_pod("pod", gpu_type_id="NVIDIA A100") -class TestCTL(unittest.TestCase): - """Tests for CTL Commands""" - - def setUp(self): - """Set up test fixtures""" - import runpod - runpod.api_key = "MOCK_API_KEY" - - def test_get_user(self): - """ - Tests get_user - """ - with patch("runpod.api.graphql.requests.post") as patch_request: - patch_request.return_value.json.return_value = { - "data": { - "myself": { - "id": "USER_ID", - } - } - } - - user = ctl_commands.get_user() - self.assertEqual(user["id"], "USER_ID") - - def test_update_user_settings(self): - """ - Tests update_user_settings - """ - with patch("runpod.api.graphql.requests.post") as patch_request: - patch_request.return_value.json.return_value = { - "data": { - "updateUserSettings": {"id": "USER_ID", "publicKey": "PUBLIC_KEY"} - } - } - - user = ctl_commands.update_user_settings("PUBLIC_KEY") - self.assertEqual(user["id"], "USER_ID") - self.assertEqual(user["publicKey"], "PUBLIC_KEY") - - def test_get_gpus(self): - """ - Tests get_gpus - """ - with patch("runpod.api.graphql.requests.post") as patch_request: - patch_request.return_value.json.return_value = { - "data": { - "gpuTypes": [ - { - "id": "NVIDIA A100 80GB PCIe", - "displayName": "A100 80GB", - "memoryInGb": 80, - } - ] - } - } - - gpus = ctl_commands.get_gpus() - - self.assertEqual(len(gpus), 1) - self.assertEqual(gpus[0]["id"], "NVIDIA A100 80GB PCIe") - - def test_get_gpu(self): - """ - Tests get_gpu_by_id - """ - with patch("runpod.api.graphql.requests.post") as patch_request: - patch_request.return_value.json.return_value = { - "data": { - "gpuTypes": [ - { - "id": "NVIDIA A100 80GB PCIe", - "displayName": "A100 80GB", - "memoryInGb": 80, - } - ] - } - } - - gpu = ctl_commands.get_gpu("NVIDIA A100 80GB PCIe") - self.assertEqual(gpu["id"], "NVIDIA A100 80GB PCIe") - - patch_request.return_value.json.return_value = {"data": {"gpuTypes": []}} - - with self.assertRaises(ValueError) as context: - gpu = ctl_commands.get_gpu("Not a GPU") - - self.assertEqual( - str(context.exception), - "No GPU found with the specified ID, " - "run runpod.get_gpus() to get a list of all GPUs", - ) - - def test_create_pod(self): - """ - Tests create_pod - """ - with patch("runpod.api.graphql.requests.post") as patch_request, patch( - "runpod.api.ctl_commands.get_gpu" - ) as patch_get_gpu, patch("runpod.api.ctl_commands.get_user") as patch_get_user: - patch_request.return_value.json.return_value = { - "data": {"podFindAndDeployOnDemand": {"id": "POD_ID"}} - } - - patch_get_gpu.return_value = None - - patch_get_user.return_value = { - "networkVolumes": [ - {"id": "NETWORK_VOLUME_ID", "dataCenterId": "us-east-1"} - ] - } - - pod = ctl_commands.create_pod( - name="POD_NAME", - image_name="IMAGE_NAME", - support_public_ip=False, - gpu_type_id="NVIDIA A100 80GB PCIe", - network_volume_id="NETWORK_VOLUME_ID", - ) - - self.assertEqual(pod["id"], "POD_ID") - - with self.assertRaises(ValueError) as context: - pod = ctl_commands.create_pod( - name="POD_NAME", - cloud_type="NOT_A_CLOUD_TYPE", - image_name="IMAGE_NAME", - gpu_type_id="NVIDIA A100 80GB PCIe", - network_volume_id="NETWORK_VOLUME_ID", - ) - - self.assertEqual( - str(context.exception), - "cloud_type must be one of ALL, COMMUNITY or SECURE", - ) - - with self.assertRaises(ValueError) as context: - pod = ctl_commands.create_pod( - name="POD_NAME", - gpu_type_id="NVIDIA A100 80GB PCIe", - network_volume_id="NETWORK_VOLUME_ID", - ) - - self.assertEqual( - str(context.exception), - "Either image_name or template_id must be provided", - ) - - def test_stop_pod(self): - """ - Test stop_pod - """ - with patch("runpod.api.graphql.requests.post") as patch_request: - patch_request.return_value.json.return_value = { - "data": {"podStop": {"id": "POD_ID"}} - } - - pod = ctl_commands.stop_pod(pod_id="POD_ID") - - self.assertEqual(pod["id"], "POD_ID") - - def test_resume_pod(self): - """ - Test resume_pod - """ - with patch("runpod.api.graphql.requests.post") as patch_request: - patch_request.return_value.json.return_value = { - "data": {"podResume": {"id": "POD_ID"}} - } - - pod = ctl_commands.resume_pod(pod_id="POD_ID", gpu_count=1) - - self.assertEqual(pod["id"], "POD_ID") - - def test_terminate_pod(self): - """ - Test terminate_pod - """ - with patch("runpod.api.graphql.requests.post") as patch_request: - patch_request.return_value.json.return_value = { - "data": {"podTerminate": {"id": "POD_ID"}} - } - - self.assertIsNone(ctl_commands.terminate_pod(pod_id="POD_ID")) - - def test_raised_error(self): - """ - Test raised_error - """ - with patch("runpod.api.graphql.requests.post") as patch_request: - patch_request.return_value.json.return_value = { - "errors": [{"message": "Error Message"}] - } - - with self.assertRaises(Exception) as context: - ctl_commands.get_gpus() - - self.assertEqual(str(context.exception), "Error Message") - - # Test Unauthorized with status code 401 - with patch("runpod.api.graphql.requests.post") as patch_request: - patch_request.return_value.status_code = 401 - - with self.assertRaises(Exception) as context: - ctl_commands.get_gpus() - - self.assertEqual( - str(context.exception), - "Unauthorized request, please check your API key.", - ) - - def test_get_pods(self): - """ - Tests get_pods - """ - with patch("runpod.api.graphql.requests.post") as patch_request: - patch_request.return_value.json.return_value = { - "data": { - "myself": { - "pods": [ - { - "id": "POD_ID", - "containerDiskInGb": 5, - "costPerHr": 0.34, - "desiredStatus": "RUNNING", - "dockerArgs": None, - "dockerId": None, - "env": [], - "gpuCount": 1, - "imageName": "runpod/pytorch:2.0.1-py3.10-cuda11.8.0-devel", - "lastStatusChange": "Rented by User: Tue Aug 15 2023", - "machineId": "MACHINE_ID", - "memoryInGb": 83, - "name": "POD_NAME", - "podType": "RESERVED", - "port": None, - "ports": "80/http", - "uptimeSeconds": 0, - "vcpuCount": 21, - "volumeInGb": 200, - "volumeMountPath": "/workspace", - "machine": {"gpuDisplayName": "RTX 3090"}, - } - ] - } - } - } - - pods = ctl_commands.get_pods() - - self.assertEqual(len(pods), 1) - self.assertEqual(pods[0]["id"], "POD_ID") - - def test_get_pod(self): - """ - Tests get_pods - """ - with patch("runpod.api.graphql.requests.post") as patch_request: - patch_request.return_value.json.return_value = { - "data": { - "pod": { - "id": "POD_ID", - "containerDiskInGb": 5, - "costPerHr": 0.34, - "desiredStatus": "RUNNING", - "dockerArgs": None, - "dockerId": None, - "env": [], - "gpuCount": 1, - "imageName": "runpod/pytorch:2.0.1-py3.10-cuda11.8.0-devel", - "lastStatusChange": "Rented by User: Tue Aug 15 2023", - "machineId": "MACHINE_ID", - "memoryInGb": 83, - "name": "POD_NAME", - "podType": "RESERVED", - "port": None, - "ports": "80/http", - "uptimeSeconds": 0, - "vcpuCount": 21, - "volumeInGb": 200, - "volumeMountPath": "/workspace", - "machine": {"gpuDisplayName": "RTX 3090"}, - } - } - } - - pods = ctl_commands.get_pod("POD_ID") - - self.assertEqual(pods["id"], "POD_ID") - - def test_create_template(self): - """ - Tests create_template - """ - with patch("runpod.api.graphql.requests.post") as patch_request, patch( - "runpod.api.ctl_commands.get_gpu" - ) as patch_get_gpu: - patch_request.return_value.json.return_value = { - "data": {"saveTemplate": {"id": "TEMPLATE_ID"}} - } - - patch_get_gpu.return_value = None - - template = ctl_commands.create_template( - name="TEMPLATE_NAME", image_name="IMAGE_NAME" - ) - - self.assertEqual(template["id"], "TEMPLATE_ID") - - def test_get_endpoints(self): - """ - Tests get_endpoints - """ - with patch("runpod.api.graphql.requests.post") as patch_request: - patch_request.return_value.json.return_value = { - "data": { - "myself": { - "endpoints": [ - { - "id": "ENDPOINT_ID", - "name": "ENDPOINT_NAME", - "template": { - "id": "TEMPLATE_ID", - "imageName": "IMAGE_NAME", - }, - } - ] - } - } - } - - endpoints = ctl_commands.get_endpoints() - - self.assertEqual(len(endpoints), 1) - self.assertEqual(endpoints[0]["id"], "ENDPOINT_ID") - - def test_create_endpoint(self): - """ - Tests create_endpoint - """ - with patch("runpod.api.graphql.requests.post") as patch_request, patch( - "runpod.api.ctl_commands.get_gpu" - ) as patch_get_gpu: - patch_request.return_value.json.return_value = { - "data": {"saveEndpoint": {"id": "ENDPOINT_ID"}} - } - - patch_get_gpu.return_value = None - - endpoint = ctl_commands.create_endpoint( - name="ENDPOINT_NAME", template_id="TEMPLATE_ID" - ) - - self.assertEqual(endpoint["id"], "ENDPOINT_ID") - - def test_update_endpoint_template(self): - """ - Tests update_endpoint_template - """ - with patch("runpod.api.graphql.requests.post") as patch_request, patch( - "runpod.api.ctl_commands.get_gpu" - ) as patch_get_gpu: - patch_request.return_value.json.return_value = { - "data": {"updateEndpointTemplate": {"id": "ENDPOINT_ID"}} - } - - patch_get_gpu.return_value = None - - endpoint = ctl_commands.update_endpoint_template( - endpoint_id="ENDPOINT_ID", template_id="TEMPLATE_ID" - ) - - self.assertEqual(endpoint["id"], "ENDPOINT_ID") - - @patch("runpod.api.ctl_commands.run_graphql_query") - def test_create_container_registry_auth(self, mock_run_graphql_query): - """ - Tests create_container_registry_auth by mocking the run_graphql_query function - """ - # Set up the mock to return a predefined response - mock_run_graphql_query.return_value = { - "data": { - "saveRegistryAuth": {"id": "REGISTRY_AUTH_ID", "name": "REGISTRY_NAME"} - } + with pytest.raises(ValueError, match="cloud_type"): + ctl_commands.create_pod( + "pod", "image", gpu_type_id="NVIDIA A100", cloud_type="INVALID" + ) + + +@pytest.mark.parametrize( + ("kwargs", "field"), + [ + ({"support_public_ip": False}, "support_public_ip"), + ({"country_code": "US"}, "country_code"), + ({"min_vcpu_count": 8}, "min_vcpu_count"), + ({"min_memory_in_gb": 32}, "min_memory_in_gb"), + ({"min_download": 100}, "min_download"), + ({"min_upload": 100}, "min_upload"), + ], +) +def test_create_gpu_pod_rejects_unsupported_constraints(kwargs, field): + with pytest.raises(ValueError, match=field): + ctl_commands.create_pod( + "pod", "image", gpu_type_id="NVIDIA A100", **kwargs + ) + + +def test_stop_and_resume_pod_use_actions(): + with patch( + "runpod.api.ctl_commands.run_rest_request", return_value={"id": "pod"} + ) as request: + assert ctl_commands.stop_pod("pod") == {"id": "pod"} + assert ctl_commands.resume_pod("pod", 8) == {"id": "pod"} + + assert request.call_args_list[0].args == ( + "POST", + "/v2/pods/pod/action", + ) + assert request.call_args_list[0].kwargs == {"json": {"action": "stop"}} + assert request.call_args_list[1].kwargs == {"json": {"action": "start"}} + + +def test_terminate_pod_deletes_resource(): + with patch("runpod.api.ctl_commands.run_rest_request", return_value=None) as request: + assert ctl_commands.terminate_pod("pod/id") is None + + request.assert_called_once_with("DELETE", "/v2/pods/pod%2Fid") + + +def test_create_template_translates_request(): + with patch( + "runpod.api.ctl_commands.run_rest_request", return_value={"id": "template"} + ) as request: + result = ctl_commands.create_template( + name="template", + image_name="image", + docker_start_cmd="python main.py", + container_disk_in_gb=20, + volume_in_gb=50, + ports="8888/http,22/tcp", + env={"PORT": 8888}, + is_serverless=True, + registry_auth_id="registry", + ) + + assert result == {"id": "template"} + request.assert_called_once_with( + "POST", + "/v2/templates", + json={ + "name": "template", + "image": "image", + "args": "python main.py", + "disk": 20, + "mounts": { + "persistent": {"size": 50, "path": "/workspace"} + }, + "ports": ["8888/http", "22/tcp"], + "env": {"PORT": "8888"}, + "serverless": True, + "registry": "registry", + }, + ) + + +def test_get_endpoints_unwraps_response(): + endpoints = [{"id": "endpoint"}] + with patch( + "runpod.api.ctl_commands.run_rest_request", + return_value={"endpoints": endpoints}, + ) as request: + assert ctl_commands.get_endpoints() == endpoints + + request.assert_called_once_with("GET", "/v2/serverless") + + +def test_create_endpoint_translates_queue_scaling(): + with patch( + "runpod.api.ctl_commands.run_rest_request", return_value={"id": "endpoint"} + ) as request: + result = ctl_commands.create_endpoint( + name="endpoint", + template_id="template", + gpu_ids="AMPERE_16,ADA_24", + network_volume_id="volume", + locations="US-KS-2, EU-RO-1", + idle_timeout=10, + scaler_value=8, + workers_min=1, + workers_max=5, + flashboot=True, + allowed_cuda_versions="12.8,12.6", + gpu_count=2, + ) + + assert result == {"id": "endpoint"} + request.assert_called_once_with( + "POST", + "/v2/serverless", + json={ + "name": "endpoint", + "templateId": "template", + "type": "QUEUE", + "gpu": { + "pools": ["AMPERE_16", "ADA_24"], + "count": 2, + "allowedCudaVersions": ["12.8", "12.6"], + }, + "networkVolumes": ["volume"], + "dataCenterIds": ["US-KS-2", "EU-RO-1"], + "workers": {"min": 1, "max": 5, "idleTimeout": 10}, + "scaling": {"type": "QUEUE_DELAY", "queueDelay": 8}, + "flashboot": "FLASHBOOT", + }, + ) + + +def test_create_endpoint_translates_worker_count_scaling(): + with patch( + "runpod.api.ctl_commands.run_rest_request", return_value={"id": "endpoint"} + ) as request: + ctl_commands.create_endpoint( + "endpoint", "template", scaler_type="WORKER_COUNT", scaler_value=2 + ) + + body = request.call_args.kwargs["json"] + assert body["scaling"] == {"type": "REQUEST_COUNT", "requestCount": 2} + assert body["workers"] == {"min": 0, "max": 3} + + +def test_create_endpoint_rejects_invalid_scaler(): + with pytest.raises(ValueError, match="scaler_type"): + ctl_commands.create_endpoint( + "endpoint", "template", scaler_type="INVALID" + ) + + +def test_update_endpoint_template_uses_patch(): + with patch( + "runpod.api.ctl_commands.run_rest_request", return_value={"id": "endpoint"} + ) as request: + assert ctl_commands.update_endpoint_template("endpoint/id", "template") == { + "id": "endpoint" } - # Call the function under test with dummy arguments + request.assert_called_once_with( + "PATCH", + "/v2/serverless/endpoint%2Fid", + json={"templateId": "template"}, + ) + + +def test_create_container_registry_auth_uses_rest(): + with patch( + "runpod.api.ctl_commands.run_rest_request", return_value={"id": "registry"} + ) as request: result = ctl_commands.create_container_registry_auth( - name="REGISTRY_NAME", username="username", password="password" + "registry", "user", "password" ) - # Assertions to verify the function behavior - self.assertEqual(result["id"], "REGISTRY_AUTH_ID") - self.assertEqual(result["name"], "REGISTRY_NAME") + assert result == {"id": "registry"} + request.assert_called_once_with( + "POST", + "/v2/registries", + json={"name": "registry", "username": "user", "password": "password"}, + ) + + +def test_update_container_registry_auth_uses_graphql(): + with patch( + "runpod.api.ctl_commands.run_graphql_query", + return_value={"data": {"updateRegistryAuth": {"id": "registry"}}}, + ) as request: + result = ctl_commands.update_container_registry_auth( + "registry", "user", "password" + ) - # Verify that run_graphql_query was called with the correct parameters - mock_run_graphql_query.assert_called_once() # Ensure it was called exactly once + assert result == {"id": "registry"} + mutation = request.call_args.args[0] + assert "mutation UpdateRegistryAuth" in mutation + assert 'id: "registry"' in mutation - # Access the first (and only) call's arguments directly - called_args, _ = mock_run_graphql_query.call_args - # The GraphQL query is expected to be the first positional argument in the call - graphql_query = called_args[0] +def test_delete_container_registry_auth_uses_rest(): + with patch("runpod.api.ctl_commands.run_rest_request", return_value=None) as request: + assert ctl_commands.delete_container_registry_auth("registry/id") is True - self.assertIn("mutation SaveRegistryAuth", graphql_query) - self.assertIn("REGISTRY_NAME", graphql_query) - self.assertIn("username", graphql_query) - self.assertIn("password", graphql_query) + request.assert_called_once_with("DELETE", "/v2/registries/registry%2Fid") diff --git a/tests/test_api/test_mutation_container_registry_auth.py b/tests/test_api/test_mutation_container_registry_auth.py index fa508a848..5df0ffe3f 100644 --- a/tests/test_api/test_mutation_container_registry_auth.py +++ b/tests/test_api/test_mutation_container_registry_auth.py @@ -1,81 +1,14 @@ -""" Test suite for the generate_container_registry_auth function. """ - -import unittest +"""Tests for GraphQL registry credential operations.""" from runpod.api.mutations.container_register_auth import ( - delete_container_registry_auth, - generate_container_registry_auth, - update_container_registry_auth + update_container_registry_auth, ) -class TestGenerateContainerRegistryAuth(unittest.TestCase): - """Test suite for the generate_container_registry_auth function.""" - - def test_generate_container_registry_auth(self): - """ - Test that the generate_container_registry_auth function produces the correct - GraphQL mutation string with the provided name, username, and password. - """ - # Define test inputs - name = "testRegistry" - username = "testUser" - password = "testPass" - - # Generate the actual mutation string - actual_mutation = generate_container_registry_auth( - name, username, password - ).strip() - - # Verify key components of the mutation string - self.assertIn("mutation SaveRegistryAuth", actual_mutation) - self.assertIn( - 'saveRegistryAuth(input: {name: "testRegistry", username: "testUser", password: "testPass"})', # pylint: disable=line-too-long - actual_mutation, - ) - self.assertIn("id", actual_mutation) - self.assertIn("name", actual_mutation) - - def test_update_container_registry_auth(self): - """ - Test that the update_container_registry_auth function produces the correct - GraphQL mutation string with the provided registry_auth_id, username and password. - """ - # Define test inputs - registry_auth_id = "testAuthId" - username = "testUser" - password = "testPass" - - # Generate the actual mutation string - actual_mutation = update_container_registry_auth( - registry_auth_id, username, password - ).strip() - - # Verify key components of the mutation string - self.assertIn("mutation UpdateRegistryAuth", actual_mutation) - self.assertIn( - 'updateRegistryAuth(input: {id: "testAuthId", username: "testUser", password: "testPass"})', # pylint: disable=line-too-long - actual_mutation, - ) - self.assertIn("id", actual_mutation) - self.assertIn("name", actual_mutation) - - def test_delete_container_registry_auth(self): - """ - Test that the delete_container_registry_auth function produces the correct - GraphQL mutation string with the provided registry_auth_id - """ - # Define test inputs - registry_auth_id = "testAuthId" - - # Generate the actual mutation string - actual_mutation = delete_container_registry_auth( - registry_auth_id - ).strip() +def test_update_container_registry_auth(): + mutation = update_container_registry_auth("registry", "user", "password") - # Verify key components of the mutation string - self.assertIn("mutation DeleteRegistryAuth", actual_mutation) - self.assertIn( - 'deleteRegistryAuth(registryAuthId: "testAuthId")', # pylint: disable=line-too-long - actual_mutation, - ) + assert "mutation UpdateRegistryAuth" in mutation + assert 'id: "registry"' in mutation + assert 'username: "user"' in mutation + assert 'password: "password"' in mutation diff --git a/tests/test_api/test_mutation_endpoints.py b/tests/test_api/test_mutation_endpoints.py deleted file mode 100644 index 0def4f308..000000000 --- a/tests/test_api/test_mutation_endpoints.py +++ /dev/null @@ -1,45 +0,0 @@ -"""Tests for the endpoint mutation generation.""" - -import unittest - -from runpod.api.mutations.endpoints import generate_endpoint_mutation - - -class TestGenerateEndpointMutation(unittest.TestCase): - """Tests for the endpoint mutation generation.""" - - def test_required_fields(self): - """Test the required fields.""" - result = generate_endpoint_mutation("test_name", "test_template_id") - self.assertIn('name: "test_name"', result) - self.assertIn('templateId: "test_template_id"', result) - self.assertIn('gpuIds: "AMPERE_16"', result) # Default value - self.assertIn('networkVolumeId: ""', result) # Default value - self.assertIn('locations: ""', result) # Default value - - def test_all_fields(self): - """Test all the fields.""" - result = generate_endpoint_mutation( - "test_name", - "test_template_id", - "AMPERE_20", - "test_volume_id", - "US_WEST", - 10, - "WORKER_COUNT", - 5, - 2, - 4, - True, - ) - self.assertIn('name: "test_name"', result) - self.assertIn('templateId: "test_template_id"', result) - self.assertIn('gpuIds: "AMPERE_20"', result) - self.assertIn('networkVolumeId: "test_volume_id"', result) - self.assertIn('locations: "US_WEST"', result) - self.assertIn("idleTimeout: 10", result) - self.assertIn('scalerType: "WORKER_COUNT"', result) - self.assertIn("scalerValue: 5", result) - self.assertIn("workersMin: 2", result) - self.assertIn("workersMax: 4", result) - self.assertIn("flashBootType: FLASHBOOT", result) diff --git a/tests/test_api/test_mutations_pods.py b/tests/test_api/test_mutations_pods.py deleted file mode 100644 index 77f289864..000000000 --- a/tests/test_api/test_mutations_pods.py +++ /dev/null @@ -1,89 +0,0 @@ -""" Test API Wrapper Pod Mutations """ - -import unittest - -from runpod.api.mutations import pods - - -class TestPodMutations(unittest.TestCase): - """Test API Wrapper Pod Mutations""" - - def test_generate_pod_deployment_mutation(self): - """ - Test generate_pod_deployment_mutation for both GPU and CPU pods - """ - # Test GPU pod deployment - gpu_result = pods.generate_pod_deployment_mutation( - name="test", - image_name="test_image", - gpu_type_id="1", - cloud_type="cloud", - data_center_id="1", - country_code="US", - gpu_count=1, - volume_in_gb=100, - container_disk_in_gb=10, - min_vcpu_count=1, - min_memory_in_gb=1, - docker_args="args", - ports="8080", - volume_mount_path="/path", - env={"ENV": "test"}, - support_public_ip=True, - template_id="abcde", - allowed_cuda_versions=["11.8", "12.0"], - ) - - # Test CPU pod deployment - cpu_result = pods.generate_pod_deployment_mutation( - name="test-cpu", - image_name="test_image", - cloud_type="cloud", - data_center_id="1", - country_code="US", - volume_in_gb=100, - container_disk_in_gb=10, - min_vcpu_count=2, - min_memory_in_gb=4, - docker_args="args", - ports="8080", - volume_mount_path="/path", - env={"ENV": "test"}, - instance_id="cpu3c-2-4" - ) - - # Check GPU pod mutation structure - self.assertIn("mutation", gpu_result) - self.assertIn("podFindAndDeployOnDemand", gpu_result) - - # Check CPU pod mutation structure - self.assertIn("mutation", cpu_result) - self.assertIn("deployCpuPod", cpu_result) - - def test_generate_pod_stop_mutation(self): - """ - Test generate_pod_stop_mutation - """ - result = pods.generate_pod_stop_mutation("pod_id") - # Here you should check the correct structure of the result - self.assertIn("mutation", result) - - def test_generate_pod_resume_mutation(self): - """ - Test generate_pod_resume_mutation - """ - result = pods.generate_pod_resume_mutation("pod_id", 1) - # Here you should check the correct structure of the result - self.assertIn("mutation", result) - - def test_generate_pod_terminate_mutation(self): - """ - Test generate_pod_terminate_mutation - """ - result = pods.generate_pod_terminate_mutation("pod_id") - # Here you should check the correct structure of the result - self.assertIn("mutation", result) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/test_api/test_mutations_templates.py b/tests/test_api/test_mutations_templates.py deleted file mode 100644 index 9e5e06d52..000000000 --- a/tests/test_api/test_mutations_templates.py +++ /dev/null @@ -1,45 +0,0 @@ -""" Unit tests for the function generate_pod_template in the file api_wrapper.py """ - -import unittest - -from runpod.api.mutations.templates import generate_pod_template - - -class TestGeneratePodTemplate(unittest.TestCase): - """Unit tests for the function generate_pod_template in the file api_wrapper.py""" - - def test_basic_required_fields(self): - """Test the basic required fields are present in the generated template""" - result = generate_pod_template("test_name", "test_image_name") - self.assertIn('name: "test_name"', result) - self.assertIn('imageName: "test_image_name"', result) - self.assertIn('dockerArgs: ""', result) # Defaults - self.assertIn("containerDiskInGb: 10", result) # Defaults - self.assertIn("volumeInGb: 0", result) # Defaults - self.assertIn('ports: ""', result) # Defaults - self.assertIn("env: []", result) # Defaults - self.assertIn("isServerless: false", result) # Defaults - - def test_optional_fields(self): - """Test the optional fields are present in the generated template""" - result = generate_pod_template( - "test_name", - "test_image_name", - docker_start_cmd="test_cmd", - volume_in_gb=5, - volume_mount_path="/path/to/volume", - ports="8000, 8001", - env={"VAR1": "val1", "VAR2": "val2"}, - is_serverless=True, - registry_auth_id="test_auth", - ) - self.assertIn('dockerArgs: "test_cmd"', result) - self.assertIn("volumeInGb: 5", result) - self.assertIn('volumeMountPath: "/path/to/volume"', result) - self.assertIn('ports: "8000,8001"', result) - self.assertIn( - 'env: [{ key: "VAR1", value: "val1" }, { key: "VAR2", value: "val2" }]', - result, - ) - self.assertIn("isServerless: true", result) - self.assertIn('containerRegistryAuthId : "test_auth"', result) diff --git a/tests/test_api/test_rest.py b/tests/test_api/test_rest.py new file mode 100644 index 000000000..d00274687 --- /dev/null +++ b/tests/test_api/test_rest.py @@ -0,0 +1,117 @@ +"""Tests for the REST API transport.""" + +from unittest.mock import Mock, patch + +import pytest + +import runpod +from runpod.api.rest import run_rest_request +from runpod.error import AuthenticationError, QueryError +from runpod.user_agent import USER_AGENT + + +def _response(status_code=200, payload=None, content=b"{}", text=""): + response = Mock() + response.status_code = status_code + response.content = content + response.text = text + response.json.return_value = payload if payload is not None else {} + return response + + +def test_request_uses_explicit_api_key(): + response = _response(payload={"pods": []}) + with patch("runpod.api.rest.requests.request", return_value=response) as request: + result = run_rest_request( + "POST", + "/v2/pods", + api_key="key", + params={"include": "all"}, + json={"name": "pod"}, + ) + + assert result == {"pods": []} + request.assert_called_once_with( + "POST", + "https://api.runpod.io/v2/pods", + headers={ + "Accept": "application/json", + "Content-Type": "application/json", + "User-Agent": USER_AGENT, + "Authorization": "Bearer key", + }, + params={"include": "all"}, + json={"name": "pod"}, + timeout=30, + ) + + +def test_request_uses_global_api_key_and_custom_base_url(): + response = _response(payload={"gpus": []}) + with ( + patch.object(runpod, "api_key", "global-key"), + patch.dict("os.environ", {"RUNPOD_API_BASE_URL": "https://example.test/"}), + patch("runpod.api.rest.requests.request", return_value=response) as request, + ): + run_rest_request("GET", "/v2/catalog/gpus") + + assert request.call_args.args[1] == "https://example.test/v2/catalog/gpus" + assert request.call_args.kwargs["headers"]["Authorization"] == "Bearer global-key" + + +def test_request_requires_api_key(): + with ( + patch.object(runpod, "api_key", None), + patch("runpod.api.rest.requests.request") as request, + pytest.raises(AuthenticationError, match="No API key provided"), + ): + run_rest_request("GET", "/v2/pods") + + request.assert_not_called() + + +def test_request_returns_none_for_no_content(): + response = _response(status_code=204, content=b"") + with patch("runpod.api.rest.requests.request", return_value=response): + assert run_rest_request("DELETE", "/v2/pods/pod", api_key="key") is None + + +def test_request_raises_authentication_error_for_unauthorized_response(): + response = _response(status_code=401, payload={"detail": "invalid key"}) + with ( + patch("runpod.api.rest.requests.request", return_value=response), + pytest.raises(AuthenticationError, match="Unauthorized request"), + ): + run_rest_request("GET", "/v2/pods", api_key="key") + + +def test_request_raises_query_error_from_problem_response(): + response = _response( + status_code=422, + payload={ + "title": "Unprocessable Entity", + "detail": "request validation failed", + "errors": ["$.name is required"], + }, + ) + with ( + patch("runpod.api.rest.requests.request", return_value=response), + pytest.raises(QueryError, match="request validation failed") as raised, + ): + run_rest_request("POST", "/v2/pods", api_key="key", json={}) + + assert raised.value.query == "POST /v2/pods" + assert raised.value.status_code == 422 + assert raised.value.errors == ["$.name is required"] + + +def test_request_uses_text_for_non_json_error(): + response = _response(status_code=500, content=b"failure", text="upstream failure") + response.json.side_effect = ValueError + with ( + patch("runpod.api.rest.requests.request", return_value=response), + pytest.raises(QueryError, match="upstream failure") as raised, + ): + run_rest_request("GET", "/v2/pods", api_key="key") + + assert raised.value.status_code == 500 diff --git a/tests/test_cli/test_cli_groups/test_pod_commands.py b/tests/test_cli/test_cli_groups/test_pod_commands.py index 1eefbe791..47250a8a2 100644 --- a/tests/test_cli/test_cli_groups/test_pod_commands.py +++ b/tests/test_cli/test_cli_groups/test_pod_commands.py @@ -23,14 +23,14 @@ def test_list_pods(self, mock_echo, mock_get_pods): { "id": "1", "name": "Pod1", - "desiredStatus": "Running", - "imageName": "Image1", + "status": "RUNNING", + "image": "Image1", }, { "id": "2", "name": "Pod2", - "desiredStatus": "Stopped", - "imageName": "Image2", + "status": "STOPPED", + "image": "Image2", }, ] @@ -40,8 +40,8 @@ def test_list_pods(self, mock_echo, mock_get_pods): # Create expected table assert result.exit_code == 0, result.exception expected_table = PrettyTable(["ID", "Name", "Status", "Image"]) - expected_table.add_row(("1", "Pod1", "Running", "Image1")) - expected_table.add_row(("2", "Pod2", "Stopped", "Image2")) + expected_table.add_row(("1", "Pod1", "RUNNING", "Image1")) + expected_table.add_row(("2", "Pod2", "STOPPED", "Image2")) # Assert that click.echo was called with the correct table mock_echo.assert_called() diff --git a/tests/test_cli/test_cli_groups/test_project_functions.py b/tests/test_cli/test_cli_groups/test_project_functions.py index 667d71e60..0a56b38a6 100644 --- a/tests/test_cli/test_cli_groups/test_project_functions.py +++ b/tests/test_cli/test_cli_groups/test_project_functions.py @@ -170,13 +170,13 @@ def test_start_nonexistent_successfully( mock_attempt_pod_launch.return_value = { "id": "new_pod_id", - "desiredStatus": "PENDING", + "status": "PROVISIONING", "runtime": None, } mock_get_pod.return_value = { "id": "new_pod_id", - "desiredStatus": "RUNNING", + "status": "RUNNING", "runtime": "ONLINE", } diff --git a/tests/test_cli/test_cli_utils/test_info.py b/tests/test_cli/test_cli_utils/test_info.py index 8f2d4ecc0..5d96396b5 100644 --- a/tests/test_cli/test_cli_utils/test_info.py +++ b/tests/test_cli/test_cli_utils/test_info.py @@ -14,12 +14,8 @@ def test_get_pod_ssh_ip_port_normal(self): """Test get_pod_ssh_ip_port normal""" with patch("runpod.cli.utils.rp_info.get_pod") as mock_get_pod: mock_get_pod.return_value = { - "desiredStatus": "RUNNING", - "runtime": { - "ports": [ - {"privatePort": 22, "ip": "127.0.0.1", "publicPort": 2222} - ] - }, + "status": "RUNNING", + "ssh": {"direct": {"host": "127.0.0.1", "port": 2222}}, } ip, port = get_pod_ssh_ip_port("pod_id") @@ -29,12 +25,12 @@ def test_get_pod_ssh_ip_port_normal(self): def test_get_pod_ssh_ip_port_timeout(self): """Test get_pod_ssh_ip_port timeout""" with patch("runpod.cli.utils.rp_info.get_pod") as mock_get_pod: - mock_get_pod.return_value = {"desiredStatus": "RUNNING", "runtime": None} + mock_get_pod.return_value = {"status": "RUNNING", "runtime": None} with pytest.raises(TimeoutError): get_pod_ssh_ip_port("pod_id", timeout=0.1) - mock_get_pod.return_value = {"desiredStatus": "NOT_RUNNING"} + mock_get_pod.return_value = {"status": "PROVISIONING"} with pytest.raises(TimeoutError): get_pod_ssh_ip_port("pod_id", timeout=0.1) diff --git a/tests/test_error.py b/tests/test_error.py index 1820a4977..3e5a3a7ab 100644 --- a/tests/test_error.py +++ b/tests/test_error.py @@ -26,16 +26,25 @@ def test_authentication_error(self): err = AuthenticationError(error_msg) self.assertEqual(str(err), error_msg) - def test_query_error_with_message_and_query(self): - """Test the QueryError class with a message and query.""" + def test_query_error_with_request_details(self): + """Test the QueryError class with request details.""" error_msg = "Query failed" - query_str = "SELECT * FROM some_table WHERE condition" - err = QueryError(error_msg, query_str) + query_str = "POST /v2/pods" + err = QueryError( + error_msg, + query_str, + status_code=422, + errors=["$.name is required"], + ) self.assertEqual(str(err), error_msg) self.assertEqual(err.query, query_str) + self.assertEqual(err.status_code, 422) + self.assertEqual(err.errors, ["$.name is required"]) - def test_query_error_without_message_and_query(self): - """Test the QueryError class without a message or query.""" + def test_query_error_without_request_details(self): + """Test the QueryError class without request details.""" err = QueryError() self.assertEqual(str(err), "None") self.assertIsNone(err.query) + self.assertIsNone(err.status_code) + self.assertEqual(err.errors, []) From 00035b8938e6c3b44502da9460f46e6ac23427b8 Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Tue, 8 Sep 2026 13:25:24 -0700 Subject: [PATCH 2/2] fix(api): preserve REST v2 creation semantics --- README.md | 16 ++ runpod/api/ctl_commands.py | 65 ++++-- tests/test_api/test_ctl_commands.py | 335 +++++++++++++++------------- 3 files changed, 234 insertions(+), 182 deletions(-) diff --git a/README.md b/README.md index 0effe146a..7acaca1b6 100644 --- a/README.md +++ b/README.md @@ -303,6 +303,22 @@ runpod.resume_pod(pod["id"], 1) runpod.terminate_pod(pod["id"]) ``` +### Template and placement options + +- With `create_pod(template_id=...)`, omitting `docker_args` inherits the template's + command. Pass `docker_args=""` to clear that inherited command. +- Omitting `volume_mount_path` preserves an inherited GPU template volume's path. + Setting it explicitly changes the path while retaining the inherited volume size. + New persistent volumes and network volume mounts default to `/runpod-volume`. +- For GPU pods, `min_memory_in_gb` and `min_vcpu_count` specify minimum host RAM + and vCPUs **per GPU**, not GPU VRAM or totals for the pod. +- `create_template(volume_in_gb=0)` creates a template without a persistent volume. +- `create_endpoint(locations="US-KS-2,EU-RO-1")` accepts comma-separated + datacenter IDs. Country codes such as `US` and `RO` are not supported. + Endpoint creation sends a single REST request without catalog lookups. +- Endpoint `gpu_ids` accepts pool IDs and excluded GPU types, for example + `gpu_ids="ADA_48_PRO,-NVIDIA L40"` selects that pool without NVIDIA L40 GPUs. + ## 📁 | Directory ```BASH diff --git a/runpod/api/ctl_commands.py b/runpod/api/ctl_commands.py index 5af4293a8..475bbc818 100644 --- a/runpod/api/ctl_commands.py +++ b/runpod/api/ctl_commands.py @@ -72,9 +72,7 @@ def get_gpus(api_key: Optional[str] = None) -> list[dict]: return response["gpus"] -def get_gpu( - gpu_id: str, gpu_quantity: int = 1, api_key: Optional[str] = None -) -> dict: +def get_gpu(gpu_id: str, gpu_quantity: int = 1, api_key: Optional[str] = None) -> dict: """Get a GPU type and its pod availability.""" try: return run_rest_request( @@ -128,9 +126,9 @@ def create_pod( container_disk_in_gb: Optional[int] = None, min_vcpu_count: int = 1, min_memory_in_gb: int = 1, - docker_args: str = "", + docker_args: Optional[str] = None, ports: Optional[str] = None, - volume_mount_path: str = "/runpod-volume", + volume_mount_path: Optional[str] = None, env: Optional[dict] = None, template_id: Optional[str] = None, network_volume_id: Optional[str] = None, @@ -150,9 +148,7 @@ def create_pod( unsupported.append("support_public_ip") if country_code is not None: unsupported.append("country_code") - if gpu_type_id and min_vcpu_count != 1: - unsupported.append("min_vcpu_count") - if min_memory_in_gb != 1: + if not gpu_type_id and min_memory_in_gb != 1: unsupported.append("min_memory_in_gb") if min_download is not None: unsupported.append("min_download") @@ -164,9 +160,10 @@ def create_pod( body: dict[str, Any] = { "name": name, - "args": docker_args, "startSsh": start_ssh, } + if docker_args is not None: + body["args"] = docker_args if image_name: body["image"] = image_name if template_id: @@ -184,19 +181,31 @@ def create_pod( if env is not None: body["env"] = _environment(env) + mount_path = ( + volume_mount_path if volume_mount_path is not None else "/runpod-volume" + ) if network_volume_id: body["mounts"] = { - "network": [ - {"volumeId": network_volume_id, "path": volume_mount_path} - ] + "network": [{"volumeId": network_volume_id, "path": mount_path}] } elif volume_in_gb: - body["mounts"] = { - "persistent": {"size": volume_in_gb, "path": volume_mount_path} - } + body["mounts"] = {"persistent": {"size": volume_in_gb, "path": mount_path}} + elif template_id and gpu_type_id and volume_mount_path is not None: + template = run_rest_request( + "GET", f"/v2/templates/{_path_segment(template_id)}" + ) + persistent = template["mounts"].get("persistent") + if persistent is not None: + body["mounts"] = { + "persistent": {"size": persistent["size"], "path": volume_mount_path} + } if gpu_type_id: gpu: dict[str, Any] = {"id": gpu_type_id, "count": gpu_count} + if min_memory_in_gb != 1: + gpu["minRamPerGpu"] = min_memory_in_gb + if min_vcpu_count != 1: + gpu["minVcpuCountPerGpu"] = min_vcpu_count if allowed_cuda_versions is not None: gpu["allowedCudaVersions"] = _split_values(allowed_cuda_versions) body["gpu"] = gpu @@ -251,7 +260,7 @@ def create_template( } if docker_start_cmd is not None: body["args"] = docker_start_cmd - if volume_in_gb is not None: + if volume_in_gb: body["mounts"] = { "persistent": { "size": volume_in_gb, @@ -289,7 +298,7 @@ def create_endpoint( allowed_cuda_versions: str = None, gpu_count: int = 1, ) -> dict: - """Create a queue-based serverless endpoint.""" + """Create a queue-based serverless endpoint; locations are datacenter IDs.""" scaler_type = { "QUEUE_DELAY": "QUEUE_DELAY", "REQUEST_COUNT": "REQUEST_COUNT", @@ -302,10 +311,22 @@ def create_endpoint( else: raise ValueError("scaler_type must be QUEUE_DELAY or REQUEST_COUNT") + pools = [] + excluded_types = [] + for gpu_id in _split_values(gpu_ids): + if gpu_id.startswith("-"): + excluded_type = gpu_id[1:].strip() + if excluded_type and excluded_type not in excluded_types: + excluded_types.append(excluded_type) + else: + pools.append(gpu_id) + gpu: dict[str, Any] = { - "pools": _split_values(gpu_ids), + "pools": pools, "count": gpu_count, } + if excluded_types: + gpu["excludedTypes"] = excluded_types if allowed_cuda_versions is not None: gpu["allowedCudaVersions"] = _split_values(allowed_cuda_versions) @@ -339,9 +360,7 @@ def update_endpoint_template(endpoint_id: str, template_id: str) -> dict: ) -def create_container_registry_auth( - name: str, username: str, password: str -) -> dict: +def create_container_registry_auth(name: str, username: str, password: str) -> dict: """Create a container registry credential.""" return run_rest_request( "POST", @@ -364,7 +383,5 @@ def update_container_registry_auth( def delete_container_registry_auth(registry_auth_id: str) -> bool: """Delete a container registry credential.""" - run_rest_request( - "DELETE", f"/v2/registries/{_path_segment(registry_auth_id)}" - ) + run_rest_request("DELETE", f"/v2/registries/{_path_segment(registry_auth_id)}") return True diff --git a/tests/test_api/test_ctl_commands.py b/tests/test_api/test_ctl_commands.py index 903a9228e..c2e8929a1 100644 --- a/tests/test_api/test_ctl_commands.py +++ b/tests/test_api/test_ctl_commands.py @@ -1,5 +1,7 @@ """Tests for the API wrapper commands.""" +from copy import deepcopy +from urllib.parse import unquote from unittest.mock import patch import pytest @@ -52,9 +54,7 @@ def test_get_gpus_unwraps_response(): def test_get_gpu_requests_pod_availability(): gpu = {"id": "NVIDIA A100"} - with patch( - "runpod.api.ctl_commands.run_rest_request", return_value=gpu - ) as request: + with patch("runpod.api.ctl_commands.run_rest_request", return_value=gpu) as request: assert ctl_commands.get_gpu("NVIDIA A100", 2, api_key="key") == gpu request.assert_called_once_with( @@ -125,81 +125,128 @@ def test_get_pod_propagates_other_api_errors(): ctl_commands.get_pod("pod") -def test_create_gpu_pod_translates_request(): - with patch( - "runpod.api.ctl_commands.run_rest_request", return_value={"id": "pod"} - ) as request: - result = ctl_commands.create_pod( - name="training", - image_name="runpod/pytorch:latest", - gpu_type_id="NVIDIA A100", - cloud_type="COMMUNITY", - start_ssh=True, - data_center_id="US-KS-2", - gpu_count=2, - volume_in_gb=20, - container_disk_in_gb=50, - docker_args="python main.py", - ports="8888/http, 22/tcp", - volume_mount_path="/workspace", - env={"COUNT": 2}, - allowed_cuda_versions=["12.8", "12.6"], - ) +@pytest.fixture +def template_backend(): + """Model REST template expansion and its persistent-volume size floor.""" + templates = { + "template/id": { + "image": "training-image", + "args": "python train.py", + "mounts": {"persistent": {"size": 30, "path": "/training"}}, + } + } - assert result == {"id": "pod"} - request.assert_called_once_with( - "POST", - "/v2/pods", - json={ - "name": "training", - "image": "runpod/pytorch:latest", - "args": "python main.py", - "startSsh": True, - "cloud": "COMMUNITY", - "dataCenterIds": ["US-KS-2"], - "disk": 50, - "ports": ["8888/http", "22/tcp"], - "env": {"COUNT": "2"}, - "mounts": { - "persistent": {"size": 20, "path": "/workspace"} - }, - "gpu": { - "id": "NVIDIA A100", - "count": 2, - "allowedCudaVersions": ["12.8", "12.6"], - }, - }, + def request(method, path, *, json=None): + if method == "GET" and path.startswith("/v2/templates/"): + return deepcopy(templates[unquote(path.rsplit("/", 1)[1])]) + if method != "POST" or path not in {"/v2/templates", "/v2/pods"}: + raise AssertionError(f"Unexpected request: {method} {path}") + persistent = json.get("mounts", {}).get("persistent") + if persistent is not None and persistent["size"] < 10: + raise QueryError( + "Persistent volume must be at least 10 GB", status_code=400 + ) + if path == "/v2/templates": + templates["created-template"] = deepcopy(json) + return {"id": "created-template"} + + # REST replaces an explicitly supplied mount object; it does not merge + # a path-only override with the template's persistent-volume size. + effective = deepcopy(templates.get(json.get("templateId"), {})) + effective.update(deepcopy(json)) + effective.setdefault("mounts", {}) + return effective + + with patch("runpod.api.ctl_commands.run_rest_request", side_effect=request): + yield + + +@pytest.mark.parametrize( + ("kwargs", "expected"), + [({}, "python train.py"), ({"docker_args": ""}, "")], + ids=["inherit-command", "clear-command"], +) +def test_create_pod_template_command_precedence(template_backend, kwargs, expected): + pod = ctl_commands.create_pod( + "training", template_id="template/id", gpu_type_id="NVIDIA A100", **kwargs ) + assert pod["args"] == expected -def test_create_gpu_pod_with_network_volume_and_template(): - with patch( - "runpod.api.ctl_commands.run_rest_request", return_value={"id": "pod"} - ) as request: - ctl_commands.create_pod( - name="training", - template_id="template", + +@pytest.mark.parametrize( + ("kwargs", "expected_path"), + [({}, "/training"), ({"volume_mount_path": "/custom"}, "/custom")], + ids=["inherit-mount", "override-path-retain-size"], +) +def test_create_pod_template_mount_precedence(template_backend, kwargs, expected_path): + pod = ctl_commands.create_pod( + "training", template_id="template/id", gpu_type_id="NVIDIA A100", **kwargs + ) + + assert pod["mounts"] == {"persistent": {"size": 30, "path": expected_path}} + + +@pytest.mark.parametrize( + ("kwargs", "expected_mounts"), + [ + ( + {"image_name": "image", "volume_in_gb": 20}, + {"persistent": {"size": 20, "path": "/runpod-volume"}}, + ), + ( + {"template_id": "template/id", "network_volume_id": "volume"}, + {"network": [{"volumeId": "volume", "path": "/runpod-volume"}]}, + ), + ], + ids=["fresh-persistent-volume", "network-replaces-template-storage"], +) +def test_create_pod_new_storage_uses_default_mount( + template_backend, kwargs, expected_mounts +): + pod = ctl_commands.create_pod("training", gpu_type_id="NVIDIA A100", **kwargs) + + assert pod["mounts"] == expected_mounts + + +def test_create_gpu_pod_enforces_per_gpu_host_minima(): + # Both deficient offers can satisfy total-pod minima for two GPUs, but not + # the requested per-GPU minima. The exact boundary must remain eligible. + offers = [ + ("too-little-ram", 16, 8), + ("too-few-cpus", 32, 4), + ("at-boundary", 32, 8), + ("oversized", 64, 64), + ] + + def allocate(method, path, *, json): + assert (method, path) == ("POST", "/v2/pods") + gpu = json["gpu"] + for machine, ram, vcpus in offers: + if ram >= gpu.get("minRamPerGpu", 1) and vcpus >= gpu.get( + "minVcpuCountPerGpu", 1 + ): + return {"machineId": machine} + raise QueryError("No matching machine", status_code=400) + + with patch("runpod.api.ctl_commands.run_rest_request", side_effect=allocate): + pod = ctl_commands.create_pod( + "training", + "image", gpu_type_id="NVIDIA A100", - network_volume_id="volume", + gpu_count=2, + min_memory_in_gb=32, + min_vcpu_count=8, ) - body = request.call_args.kwargs["json"] - assert body["templateId"] == "template" - assert body["mounts"] == { - "network": [{"volumeId": "volume", "path": "/runpod-volume"}] - } - assert "image" not in body - assert "disk" not in body - assert "cloud" not in body + assert pod["machineId"] == "at-boundary" def test_create_cpu_pod_translates_instance_id(): with patch( "runpod.api.ctl_commands.run_rest_request", return_value={"id": "pod"} ) as request: - ctl_commands.create_pod( - "cpu-pod", "python:3.11", instance_id="cpu3c-4-8" - ) + ctl_commands.create_pod("cpu-pod", "python:3.11", instance_id="cpu3c-4-8") assert request.call_args.kwargs["json"]["cpu"] == { "id": "cpu3c", @@ -214,9 +261,7 @@ def test_create_cpu_pod_requires_instance_id(): def test_create_cpu_pod_validates_instance_id(): with pytest.raises(ValueError, match="format"): - ctl_commands.create_pod( - "cpu-pod", "python:3.11", instance_id="cpu3c-invalid" - ) + ctl_commands.create_pod("cpu-pod", "python:3.11", instance_id="cpu3c-invalid") def test_create_pod_validates_image_and_cloud(): @@ -234,17 +279,13 @@ def test_create_pod_validates_image_and_cloud(): [ ({"support_public_ip": False}, "support_public_ip"), ({"country_code": "US"}, "country_code"), - ({"min_vcpu_count": 8}, "min_vcpu_count"), - ({"min_memory_in_gb": 32}, "min_memory_in_gb"), ({"min_download": 100}, "min_download"), ({"min_upload": 100}, "min_upload"), ], ) def test_create_gpu_pod_rejects_unsupported_constraints(kwargs, field): with pytest.raises(ValueError, match=field): - ctl_commands.create_pod( - "pod", "image", gpu_type_id="NVIDIA A100", **kwargs - ) + ctl_commands.create_pod("pod", "image", gpu_type_id="NVIDIA A100", **kwargs) def test_stop_and_resume_pod_use_actions(): @@ -263,47 +304,34 @@ def test_stop_and_resume_pod_use_actions(): def test_terminate_pod_deletes_resource(): - with patch("runpod.api.ctl_commands.run_rest_request", return_value=None) as request: + with patch( + "runpod.api.ctl_commands.run_rest_request", return_value=None + ) as request: assert ctl_commands.terminate_pod("pod/id") is None request.assert_called_once_with("DELETE", "/v2/pods/pod%2Fid") -def test_create_template_translates_request(): - with patch( - "runpod.api.ctl_commands.run_rest_request", return_value={"id": "template"} - ) as request: - result = ctl_commands.create_template( - name="template", - image_name="image", - docker_start_cmd="python main.py", - container_disk_in_gb=20, - volume_in_gb=50, - ports="8888/http,22/tcp", - env={"PORT": 8888}, - is_serverless=True, - registry_auth_id="registry", - ) - - assert result == {"id": "template"} - request.assert_called_once_with( - "POST", - "/v2/templates", - json={ - "name": "template", - "image": "image", - "args": "python main.py", - "disk": 20, - "mounts": { - "persistent": {"size": 50, "path": "/workspace"} - }, - "ports": ["8888/http", "22/tcp"], - "env": {"PORT": "8888"}, - "serverless": True, - "registry": "registry", - }, +@pytest.mark.parametrize( + ("volume_in_gb", "expected_mounts"), + [ + (0, {}), + (10, {"persistent": {"size": 10, "path": "/data"}}), + ], + ids=["no-persistent-storage", "minimum-persistent-storage"], +) +def test_create_template_persistent_storage_boundary( + template_backend, volume_in_gb, expected_mounts +): + template = ctl_commands.create_template( + "template", "image", volume_in_gb=volume_in_gb, volume_mount_path="/data" + ) + pod = ctl_commands.create_pod( + "training", template_id=template["id"], gpu_type_id="NVIDIA A100" ) + assert pod["mounts"] == expected_mounts + def test_get_endpoints_unwraps_response(): endpoints = [{"id": "endpoint"}] @@ -316,65 +344,54 @@ def test_get_endpoints_unwraps_response(): request.assert_called_once_with("GET", "/v2/serverless") -def test_create_endpoint_translates_queue_scaling(): - with patch( - "runpod.api.ctl_commands.run_rest_request", return_value={"id": "endpoint"} - ) as request: - result = ctl_commands.create_endpoint( - name="endpoint", - template_id="template", - gpu_ids="AMPERE_16,ADA_24", - network_volume_id="volume", - locations="US-KS-2, EU-RO-1", - idle_timeout=10, - scaler_value=8, - workers_min=1, - workers_max=5, - flashboot=True, - allowed_cuda_versions="12.8,12.6", - gpu_count=2, - ) +@pytest.fixture +def endpoint_backend(): + data_centers = ["US-KS-2", "US-TX-3", "EU-RO-1", "CA-MTL-1"] + pools = { + "AMPERE_16": {"NVIDIA RTX A4000"}, + "ADA_48_PRO": {"NVIDIA L40", "NVIDIA L40S"}, + } - assert result == {"id": "endpoint"} - request.assert_called_once_with( - "POST", - "/v2/serverless", - json={ - "name": "endpoint", - "templateId": "template", - "type": "QUEUE", - "gpu": { - "pools": ["AMPERE_16", "ADA_24"], - "count": 2, - "allowedCudaVersions": ["12.8", "12.6"], - }, - "networkVolumes": ["volume"], - "dataCenterIds": ["US-KS-2", "EU-RO-1"], - "workers": {"min": 1, "max": 5, "idleTimeout": 10}, - "scaling": {"type": "QUEUE_DELAY", "queueDelay": 8}, - "flashboot": "FLASHBOOT", - }, + def request(method, path, *, json): + if (method, path) != ("POST", "/v2/serverless"): + raise AssertionError(f"Unexpected request: {method} {path}") + eligible = set() + for pool in json["gpu"]["pools"]: + if pool not in pools: + raise QueryError("GPU pool is not available", status_code=400) + eligible.update(pools[pool]) + eligible.difference_update(json["gpu"].get("excludedTypes", [])) + selection = json.get("dataCenterIds") or data_centers + placements = [center for center in selection if center in data_centers] + if not placements: + raise QueryError("No matching data centers", status_code=400) + return {"eligibleGpuTypes": sorted(eligible), "dataCenterIds": placements} + + with patch("runpod.api.ctl_commands.run_rest_request", side_effect=request): + yield + + +def test_create_endpoint_preserves_datacenter_selection(endpoint_backend): + endpoint = ctl_commands.create_endpoint( + "endpoint", "template", locations="US-KS-2, EU-RO-1" ) + assert set(endpoint["dataCenterIds"]) == {"US-KS-2", "EU-RO-1"} -def test_create_endpoint_translates_worker_count_scaling(): - with patch( - "runpod.api.ctl_commands.run_rest_request", return_value={"id": "endpoint"} - ) as request: - ctl_commands.create_endpoint( - "endpoint", "template", scaler_type="WORKER_COUNT", scaler_value=2 - ) - body = request.call_args.kwargs["json"] - assert body["scaling"] == {"type": "REQUEST_COUNT", "requestCount": 2} - assert body["workers"] == {"min": 0, "max": 3} +def test_create_endpoint_preserves_excluded_gpu_types(endpoint_backend): + endpoint = ctl_commands.create_endpoint( + "endpoint", + "template", + gpu_ids="ADA_48_PRO, AMPERE_16, -NVIDIA L40, - NVIDIA L40", + ) + + assert endpoint["eligibleGpuTypes"] == ["NVIDIA L40S", "NVIDIA RTX A4000"] def test_create_endpoint_rejects_invalid_scaler(): with pytest.raises(ValueError, match="scaler_type"): - ctl_commands.create_endpoint( - "endpoint", "template", scaler_type="INVALID" - ) + ctl_commands.create_endpoint("endpoint", "template", scaler_type="INVALID") def test_update_endpoint_template_uses_patch(): @@ -424,7 +441,9 @@ def test_update_container_registry_auth_uses_graphql(): def test_delete_container_registry_auth_uses_rest(): - with patch("runpod.api.ctl_commands.run_rest_request", return_value=None) as request: + with patch( + "runpod.api.ctl_commands.run_rest_request", return_value=None + ) as request: assert ctl_commands.delete_container_registry_auth("registry/id") is True request.assert_called_once_with("DELETE", "/v2/registries/registry%2Fid")