From 9b21a9b6f9aa435b54050ccd8dfbcf300dd20b6f Mon Sep 17 00:00:00 2001 From: Dmitry Meyer Date: Wed, 21 Oct 2026 11:08:16 +0000 Subject: [PATCH] Migrate Runpod backend to REST API v2 Pods, clusters, network volumes, and container registry credentials are managed via REST API v2 instead of the deprecated GraphQL API, with the following exceptions and changes: * Community Cloud pods are still created via GraphQL, since v2 does not support placement by country code. gpuhunt also still uses GraphQL. * Clusters with `registry_auth` are no longer supported: v2 cluster creation does not accept registry credentials, and cluster member pods cannot be updated via `PATCH /v2/pods/{id}`. * Spot offers are rejected, since v2 has no interruptible pods. gpuhunt does not produce Runpod spot offers anyway. * The cluster internal IP is set in `update_provisioning_data()`, since v2 may assign it after the cluster is created. * A cluster is deleted if its pods cannot be listed after creation or their count does not match, as dstack does not track it yet. * A failed container registry credentials cleanup no longer fails provisioning of an already created pod. --- .../core/backends/runpod/api_client.py | 1003 ++++++++--------- .../_internal/core/backends/runpod/compute.py | 327 +++--- .../core/backends/runpod/test_api_client.py | 256 +++-- .../core/backends/runpod/test_compute.py | 258 +++++ 4 files changed, 1122 insertions(+), 722 deletions(-) create mode 100644 src/tests/_internal/core/backends/runpod/test_compute.py diff --git a/src/dstack/_internal/core/backends/runpod/api_client.py b/src/dstack/_internal/core/backends/runpod/api_client.py index 2feae83ca7..c550cb72da 100644 --- a/src/dstack/_internal/core/backends/runpod/api_client.py +++ b/src/dstack/_internal/core/backends/runpod/api_client.py @@ -1,18 +1,222 @@ -import time -from datetime import timedelta -from typing import Any, Dict, List, Optional +from collections.abc import Mapping +from http import HTTPStatus +from typing import Any, Dict, List, Literal, Optional, TypedDict, cast import requests from gpuhunt.providers.runpod import RunpodProvider -from requests import Response +from typing_extensions import NotRequired from dstack._internal.core.errors import BackendError, BackendInvalidCredentialsError -from dstack._internal.utils.common import get_current_datetime +from dstack._internal.utils.logging import get_logger -API_URL = "https://api.runpod.io/graphql" +logger = get_logger(__name__) +API_BASE_URL = "https://api.runpod.io" +API_GRAPHQL_URL = f"{API_BASE_URL}/graphql" +API_V2_URL = f"{API_BASE_URL}/v2" -class RunpodApiClientError(BackendError): +REQUEST_TIMEOUT = 120 + +AUTH_ERROR_STATUS_CODES = (HTTPStatus.UNAUTHORIZED, HTTPStatus.FORBIDDEN) + +# Applies to non-JSON error responses, which are not expected to contain anything useful past +# the first few lines +MAX_ERROR_DETAIL_LEN = 200 + + +class V2ErrorResponse(TypedDict): + title: str + status: int + detail: str + errors: NotRequired[list[str]] + + +V2Cloud = Literal["SECURE", "COMMUNITY"] +"""Cloud tier. +`SECURE` — Runpod-owned datacenter hardware, `COMMUNITY` — community-hosted hardware""" + + +class V2Pod(TypedDict): + id: str + status: Literal["PROVISIONING", "STARTING", "RUNNING", "EXITED", "ERROR", "TERMINATED"] + runtime: Optional["V2PodRuntime"] + mounts: "V2Mounts" + cluster: NotRequired["V2PodCluster"] + + +class V2PodList(TypedDict): + pods: list[V2Pod] + + +class V2PodCluster(TypedDict): + id: str + rank: Optional[int] + ip: NotRequired[str] + + +class V2PodRuntime(TypedDict): + ports: NotRequired[list["V2PodRuntimePort"]] + + +class V2PodRuntimePort(TypedDict): + private: NotRequired[int] + public: NotRequired[Optional[int]] + ip: NotRequired[Optional[str]] + + +class V2Mounts(TypedDict): + network: NotRequired[list["V2NetworkMount"]] + """The `network` field is an array for forward compatibility with eventual multi-network-volume + support, but `maxItems` is 1 today""" + + +class V2NetworkMount(TypedDict): + volumeId: str + path: str + + +class V2CreatePodRequest(TypedDict): + """ + Request body for creating a pod. Exactly one of `gpu` or `cpu` must be set. + """ + + name: str + cloud: NotRequired[V2Cloud] + """Cloud tier. Defaults to `SECURE` when omitted""" + dataCenterIds: NotRequired[list[str]] + """Preferred data centers for placement. Omit or pass an empty array to let + the scheduler choose.""" + gpu: NotRequired["V2CreateGpuConfig"] + cpu: NotRequired["V2CreateCpuConfig"] + templateId: NotRequired[str] + """ID of a pod template to base this pod on. The template is resolved at create time into + the same container settings you could otherwise spread into this body.""" + disk: NotRequired[int] + """Container disk in GB (ephemeral, wiped on restart)""" + mounts: NotRequired["V2Mounts"] + image: NotRequired[str] + """`image` is required unless `templateId` is set""" + registry: NotRequired[Optional[str]] + """Container registry credential ID (for private images)""" + entrypoint: NotRequired[list[str]] + """Container ENTRYPOINT in exec form""" + cmd: NotRequired[list[str]] + """Container CMD in exec form""" + ports: NotRequired[list[str]] + """Exposed ports, formatted as port/protocol""" + env: NotRequired[dict[str, str]] + + +class V2UpdatePodRequest(TypedDict): + registry: NotRequired[Optional[str]] + + +class V2CreateGpuConfig(TypedDict): + id: str + """GPU type identifier, as returned by `GET /v2/catalog/gpus`""" + count: NotRequired[int] + """Number of GPUs. Default: `1`""" + minCudaVersion: NotRequired[str] + """Lowest acceptable CUDA version for the host machine, as `major.minor`""" + minRamPerGpu: NotRequired[int] + """Minimum host system RAM in GB per requested GPU. This is a placement filter, not a resource + request. The actual total allocated system RAM may be higher and is returned in `gpu.memory`""" + minVcpuCountPerGpu: NotRequired[int] + """Minimum host vCPU count per requested GPU. This is a placement filter, not a resource + request. The actual total allocated vCPU count may be higher and is returned in `gpu.vcpuCount` + """ + + +class V2CreateCpuConfig(TypedDict): + id: str + """CPU flavor identifier, as returned by `GET /v2/catalog/cpus`""" + vcpuCount: int + """Number of vCPUs. Must be valid for the selected CPU flavor and must be a power of two.""" + + +class V2Cluster(TypedDict): + id: str + name: str + type: "V2ClusterType" + compute: "V2ClusterCompute" + + +V2ClusterType = Literal["APPLICATION", "TRAINING", "SLURM", "RAY"] +""""Cluster type. `TRAINING` is the generic distributed-training cluster; `SLURM` provisions +a managed Slurm controller/compute topology; `RAY` provisions a managed Ray head/worker topology; +`APPLICATION` is a general multi-node application cluster.""" + + +class V2ClusterCompute(TypedDict): + gpuTypeId: str + gpuCountPerPod: int + podCount: int + + +class V2CreateClusterRequest(TypedDict): + name: str + type: V2ClusterType + compute: V2ClusterCompute + dataCenterIds: NotRequired[list[str]] + """Preferred data centers for placement. Omit or pass an empty array to let the scheduler + choose. A cluster is always placed within a single data center.""" + templateId: NotRequired[str] + """ID of a pod template to provision every member pod from. The template supplies the container + settings (image, args, disk, env, ports) and the container registry credential for private + images — the only private-image path for clusters.""" + disk: NotRequired[int] + """Container disk in GB (ephemeral, wiped on restart)""" + mounts: NotRequired["V2Mounts"] + image: NotRequired[str] + """`image` is required unless `templateId` is set""" + entrypoint: NotRequired[list[str]] + """Container ENTRYPOINT in exec form""" + cmd: NotRequired[list[str]] + """Container CMD in exec form""" + ports: NotRequired[list[str]] + """Exposed ports, formatted as port/protocol""" + env: NotRequired[dict[str, str]] + + +class V2Registry(TypedDict): + id: str + name: str + + +class V2CreateRegistryRequest(TypedDict): + name: str + password: str + username: str + + +class V2ListRegistriesResponse(TypedDict): + registries: list[V2Registry] + + +V2VolumeType = Literal["STANDARD", "HIGH_PERFORMANCE"] + + +class V2NetworkVolume(TypedDict): + id: str + name: str + """Human-readable name (not required to be unique)""" + size: int + """Allocated storage in GB, [ 10 .. 4096 ]""" + dataCenter: str + type: V2VolumeType + + +class V2CreateNetworkVolumeRequest(TypedDict): + name: str + size: int + """Storage to allocate in GB, [ 10 .. 4096 ]""" + dataCenter: str + type: NotRequired[V2VolumeType] + """ When omitted, the volume is provisioned using the requested data center's default (primary) + storage tier.""" + + +class RunpodGraphqlClientError(BackendError): errors: List[Dict] def __init__(self, errors: List[Dict]): @@ -20,6 +224,27 @@ def __init__(self, errors: List[Dict]): super().__init__(errors) +class RunpodV2ClientError(BackendError): + def __init__( + self, *, status: int, title: str, detail: str, errors: Optional[list[str]] = None + ) -> None: + self.status = status + self.title = title + self.detail = detail + self.errors = errors or [] + + def __str__(self) -> str: + parts: list[str] = [str(self.status), self.title, self.detail.rstrip(".")] + # NOTE: Exceptions are exposed to the end user. We assume `errors` doesn't contain + # sensitive data + if self.errors: + parts.append(repr(self.errors)) + return ": ".join(parts) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}({self})" + + class RunpodApiClient: def __init__(self, api_key: str): self._session = requests.Session() @@ -27,28 +252,85 @@ def __init__(self, api_key: str): def validate_api_key(self) -> bool: try: - self.get_user_details() + self._make_v2_request("GET", "account", "ssh-keys") except BackendInvalidCredentialsError: return False return True - def get_user_details(self) -> Dict: - resp = self._make_request( - { - "query": """ - query myself { - myself { - id - authId - email - } - } - """ + def create_pod( + self, + *, + name: str, + cloud: V2Cloud, + data_center_ids: Optional[list[str]] = None, + gpu_type_id: Optional[str] = None, + gpu_count: Optional[int] = None, + min_cpu_count_per_gpu: Optional[int] = None, + min_memory_in_gb_per_gpu: Optional[int] = None, + cpu_type_id: Optional[str] = None, + cpu_count: Optional[int] = None, + container_disk_in_gb: int, + network_volume_id: Optional[str] = None, + volume_mount_path: Optional[str] = None, + image_name: str, + container_registry_credentials_id: Optional[str] = None, + entrypoint: list[str], + cmd: list[str], + ports: Optional[list[str]] = None, + env: Optional[dict[str, str]] = None, + ) -> V2Pod: + body: V2CreatePodRequest = { + "name": name, + "cloud": cloud, + "image": image_name, + "entrypoint": entrypoint, + "cmd": cmd, + "disk": container_disk_in_gb, + } + + if (gpu_type_id is None) == (cpu_type_id is None): + raise ValueError("Exactly one of CPU or GPU type must be set") + if gpu_type_id is not None: + if cpu_count is not None: + raise ValueError("cpu_count is not applicable for GPU pods") + gpu_config: V2CreateGpuConfig = { + "id": gpu_type_id, + "minCudaVersion": RunpodProvider.MIN_CUDA_VERSION, } - ) - return resp.json() + if gpu_count is not None: + gpu_config["count"] = gpu_count + if min_cpu_count_per_gpu is not None: + gpu_config["minVcpuCountPerGpu"] = min_cpu_count_per_gpu + if min_memory_in_gb_per_gpu is not None: + gpu_config["minRamPerGpu"] = min_memory_in_gb_per_gpu + body["gpu"] = gpu_config + else: + if cpu_count is None: + raise ValueError("cpu_count is required for CPU pods") + assert cpu_type_id is not None + cpu_config: V2CreateCpuConfig = { + "id": cpu_type_id, + "vcpuCount": cpu_count, + } + body["cpu"] = cpu_config - def create_pod( + mounts = _build_v2_mounts(network_volume_id, volume_mount_path) + if mounts is not None: + body["mounts"] = mounts + + if data_center_ids is not None: + body["dataCenterIds"] = data_center_ids + if container_registry_credentials_id is not None: + body["registry"] = container_registry_credentials_id + if ports is not None: + body["ports"] = ports + if env is not None: + body["env"] = env + + response = self._make_v2_request("POST", "pods", json=body) + return response.json() + + def create_gpu_pod_graphql( self, name: str, image_name: str, @@ -73,7 +355,7 @@ def create_pod( bid_per_gpu: Optional[float] = None, container_registry_auth_id: Optional[str] = None, ) -> Dict: - resp = self._make_request( + resp = self._make_graphql_request( { "query": _generate_pod_deployment_mutation( name=name, @@ -104,327 +386,209 @@ def create_pod( data = resp.json()["data"] return data["podRentInterruptable"] if bid_per_gpu else data["podFindAndDeployOnDemand"] - def create_cpu_pod( - self, - name: str, - image_name: str, - instance_id: str, - cloud_type: str, - deploy_cost: float, - start_ssh: bool = True, - data_center_id: Optional[str] = None, - container_disk_in_gb: Optional[int] = None, - docker_args: Optional[str] = None, - ports: Optional[str] = None, - volume_mount_path: Optional[str] = None, - env: Optional[Dict[str, Any]] = None, - template_id: Optional[str] = None, - network_volume_id: Optional[str] = None, - container_registry_auth_id: Optional[str] = None, - ) -> Dict: - resp = self._make_request( - { - "query": _generate_cpu_pod_deployment_mutation( - name=name, - image_name=image_name, - instance_id=instance_id, - cloud_type=cloud_type, - deploy_cost=deploy_cost, - start_ssh=start_ssh, - data_center_id=data_center_id, - container_disk_in_gb=container_disk_in_gb, - docker_args=docker_args, - ports=ports, - volume_mount_path=volume_mount_path, - env=env, - template_id=template_id, - network_volume_id=network_volume_id, - container_registry_auth_id=container_registry_auth_id, - ) - } - ) - return resp.json()["data"]["deployCpuPod"] - - def update_pod_container_registry_auth( - self, - pod_id: str, - container_registry_auth_id: str, - ) -> str: - resp = self._make_request( - { - "query": f""" - mutation {{ - podEditJob(input: {{ - podId: "{pod_id}" - containerRegistryAuthId: "{container_registry_auth_id}" - }}) {{ - id - }} - }} - """ - } - ) - return resp.json()["data"]["podEditJob"]["id"] - - def get_pod(self, pod_id: str) -> Dict: - resp = self._make_request({"query": _generate_pod_query(pod_id)}) - data = resp.json() - return data["data"]["pod"] - - def terminate_pod(self, pod_id: str) -> Dict: - resp = self._make_request({"query": _generate_pod_terminate_mutation(pod_id)}) - data = resp.json() - return data["data"] - - def get_container_registry_auths(self) -> List[Dict]: - resp = self._make_request( - { - "query": """ - query myself { - myself { - containerRegistryCreds { - id, - name - } - } - } - """ - } - ) - return resp.json()["data"]["myself"]["containerRegistryCreds"] + def get_pod(self, pod_id: str) -> Optional[V2Pod]: + try: + response = self._make_v2_request("GET", "pods", pod_id) + except RunpodV2ClientError as e: + if e.status == HTTPStatus.NOT_FOUND: + return None + raise + return response.json() - def add_container_registry_auth(self, name: str, username: str, password: str) -> str: - resp = self._make_request( - { - "query": f""" - mutation {{ - saveRegistryAuth( - input: {{ - name: "{name}", - username: "{username}", - password: "{password}" - }} - ) {{ - id - }} - }} - """ - } - ) - return resp.json()["data"]["saveRegistryAuth"]["id"] + def terminate_pod(self, pod_id: str) -> bool: + try: + self._make_v2_request("DELETE", "pods", pod_id) + except RunpodV2ClientError as e: + if e.status == HTTPStatus.NOT_FOUND: + return False + raise + return True - def delete_container_registry_auth(self, auth_id: str) -> None: - self._make_request( - { - "query": f""" - mutation {{ - deleteRegistryAuth( - registryAuthId: "{auth_id}" - ) - }} - """ - } - ) + def list_container_registry_credentials(self) -> List[V2Registry]: + response = self._make_v2_request("GET", "registries") + response_json: V2ListRegistriesResponse = response.json() + return response_json["registries"] - def get_network_volume(self, volume_id: str) -> Optional[Dict]: - response = self._make_request( - { - "query": """ - query getMyVolumes { - myself { - networkVolumes { - id, - name, - size, - dataCenter { - id - name - } - } - } - } - """ - } - ) - network_volumes = response.json()["data"]["myself"]["networkVolumes"] - for vol in network_volumes: - if vol["id"] == volume_id: - return vol - return None - - def create_network_volume(self, name: str, region: str, size: int) -> str: - response = self._make_request( - { - "query": f""" - mutation {{ - createNetworkVolume( - input: {{ - name: "{name}", - size: {size}, - dataCenterId: "{region}" - }} - ) {{ - id - }} - }} - """ - } - ) - return response.json()["data"]["createNetworkVolume"]["id"] - - def delete_network_volume(self, volume_id: str) -> None: - self._make_request( - { - "query": f""" - mutation {{ - deleteNetworkVolume( - input: {{ - id: "{volume_id}" - }} - ) - }} - """ - } - ) + def create_container_registry_credentials( + self, + *, + name: str, + username: str, + password: str, + ) -> V2Registry: + body: V2CreateRegistryRequest = { + "name": name, + "username": username, + "password": password, + } + response = self._make_v2_request("POST", "registries", json=body) + return response.json() + + def delete_container_registry_credentials(self, credentials_id: str) -> None: + # As of 2026-09-18, v2 API returns 400 with "detail": "Registry auth not found" + # instead of 404, hence we don't handle "not found" case + self._make_v2_request("DELETE", "registries", credentials_id) + + def get_network_volume(self, volume_id: str) -> Optional[V2NetworkVolume]: + try: + response = self._make_v2_request("GET", "network-volumes", volume_id) + except RunpodV2ClientError as e: + if e.status == HTTPStatus.NOT_FOUND: + return None + raise + return response.json() + + def create_network_volume( + self, *, name: str, data_center_id: str, size: int + ) -> V2NetworkVolume: + body: V2CreateNetworkVolumeRequest = { + "name": name, + "dataCenter": data_center_id, + "size": size, + } + response = self._make_v2_request("POST", "network-volumes", json=body) + return response.json() + + def delete_network_volume(self, volume_id: str) -> bool: + try: + self._make_v2_request("DELETE", "network-volumes", volume_id) + except RunpodV2ClientError as e: + if e.status == HTTPStatus.NOT_FOUND: + return False + raise + return True def create_cluster( self, - cluster_name: str, + *, + name: str, + cluster_type: V2ClusterType, gpu_type_id: str, - pod_count: int, gpu_count_per_pod: int, - image_name: str, - deploy_cost: str, - template_id: Optional[str] = None, - cluster_type: str = "TRAINING", + pod_count: int, + data_center_ids: Optional[list[str]] = None, + container_disk_in_gb: int, network_volume_id: Optional[str] = None, - volume_in_gb: Optional[int] = None, - throughput: Optional[int] = None, - allowed_cuda_versions: Optional[List[str]] = None, - volume_key: Optional[str] = None, - data_center_id: Optional[str] = None, - start_jupyter: bool = False, - start_ssh: bool = False, - container_disk_in_gb: Optional[int] = None, - docker_args: Optional[str] = None, - env: Optional[Dict[str, Any]] = None, volume_mount_path: Optional[str] = None, - ports: Optional[str] = None, - ) -> Dict: - resp = self._make_request( - { - "query": _generate_create_cluster_mutation( - cluster_name=cluster_name, - gpu_type_id=gpu_type_id, - pod_count=pod_count, - gpu_count_per_pod=gpu_count_per_pod, - image_name=image_name, - cluster_type=cluster_type, - deploy_cost=deploy_cost, - template_id=template_id, - network_volume_id=network_volume_id, - volume_in_gb=volume_in_gb, - throughput=throughput, - allowed_cuda_versions=allowed_cuda_versions, - volume_key=volume_key, - data_center_id=data_center_id, - start_jupyter=start_jupyter, - start_ssh=start_ssh, - container_disk_in_gb=container_disk_in_gb, - docker_args=docker_args, - env=env, - volume_mount_path=volume_mount_path, - ports=ports, - ) - } - ) - data = resp.json()["data"] - return data["createCluster"] + image_name: str, + entrypoint: list[str], + cmd: list[str], + ports: Optional[list[str]] = None, + env: Optional[dict[str, str]] = None, + ) -> V2Cluster: + body: V2CreateClusterRequest = { + "name": name, + "type": cluster_type, + "compute": { + "gpuTypeId": gpu_type_id, + "gpuCountPerPod": gpu_count_per_pod, + "podCount": pod_count, + }, + "image": image_name, + "entrypoint": entrypoint, + "cmd": cmd, + "disk": container_disk_in_gb, + } + mounts = _build_v2_mounts(network_volume_id, volume_mount_path) + if mounts is not None: + body["mounts"] = mounts + if data_center_ids is not None: + body["dataCenterIds"] = data_center_ids + if ports is not None: + body["ports"] = ports + if env is not None: + body["env"] = env + response = self._make_v2_request("POST", "clusters", json=body) + return response.json() def delete_cluster(self, cluster_id: str) -> bool: - resp = self._make_request({"query": _generate_delete_cluster_mutation(cluster_id)}) - data = resp.json()["data"] - return data["deleteCluster"] + try: + self._make_v2_request("DELETE", "clusters", cluster_id) + except RunpodV2ClientError as e: + if e.status == HTTPStatus.NOT_FOUND: + return False + raise + return True + + def list_cluster_pods(self, cluster_id: str) -> list[V2Pod]: + response = self._make_v2_request("GET", "clusters", cluster_id, "pods") + pod_list: V2PodList = response.json() + return pod_list["pods"] - def _make_request(self, data: Optional[Dict[str, Any]] = None) -> Response: + def _make_graphql_request(self, data: Optional[Dict[str, Any]] = None) -> requests.Response: try: response = self._session.request( method="POST", - url=API_URL, + url=API_GRAPHQL_URL, json=data, - timeout=120, + timeout=REQUEST_TIMEOUT, ) response.raise_for_status() response_json = response.json() # Runpod returns 200 on client errors if "errors" in response_json: - raise RunpodApiClientError(errors=response_json["errors"]) + raise RunpodGraphqlClientError(errors=response_json["errors"]) return response except requests.HTTPError as e: - if e.response is not None and e.response.status_code in ( - requests.codes.forbidden, - requests.codes.unauthorized, - ): + if e.response is not None and e.response.status_code in AUTH_ERROR_STATUS_CODES: raise BackendInvalidCredentialsError(e.response.text) raise - def wait_for_instance(self, instance_id: str) -> Optional[Dict]: - start = get_current_datetime() - wait_for_instance_interval = 5 - # To change the status to "running," the image must be pulled and then started. - # We have to wait for 20 minutes while a large image is pulled. - while get_current_datetime() < (start + timedelta(minutes=20)): - pod = self.get_pod(instance_id) - if pod["runtime"] is not None: - return pod - time.sleep(wait_for_instance_interval) - return - - -def _generate_pod_query(pod_id: str) -> str: - """ - Generate a query for a specific GPU type - """ - - return f""" - query pod {{ - pod(input: {{podId: "{pod_id}"}}) {{ - id - clusterIp - 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 _make_v2_request( + self, method: str, /, *path_parts: str, json: Optional[Mapping[str, Any]] = None + ) -> requests.Response: + url = "/".join((API_V2_URL, *path_parts)) + response = self._session.request( + method=method, + url=url, + timeout=REQUEST_TIMEOUT, + json=json, + ) + if HTTPStatus.BAD_REQUEST <= response.status_code < HTTPStatus.INTERNAL_SERVER_ERROR: + error = _parse_v2_error_response(response) + logger.debug("Runpod API v2 error: %s", error) + if response.status_code in AUTH_ERROR_STATUS_CODES: + raise BackendInvalidCredentialsError(f"{error['title']}: {error['detail']}") + raise RunpodV2ClientError( + status=response.status_code, + title=error["title"], + detail=error["detail"], + errors=error.get("errors"), + ) + response.raise_for_status() + return response + + +def _parse_v2_error_response(response: requests.Response) -> V2ErrorResponse: + try: + error = response.json() + except requests.JSONDecodeError: + error = None + if isinstance(error, dict) and "title" in error and "detail" in error: + return cast(V2ErrorResponse, error) + # Not a v2 error response, e.g., an HTML page served by a proxy in front of the API + return { + "title": response.reason or "Unknown error", + "status": response.status_code, + "detail": " ".join(response.text.split())[:MAX_ERROR_DETAIL_LEN], + } + + +def _build_v2_mounts( + network_volume_id: Optional[str], volume_mount_path: Optional[str] +) -> Optional[V2Mounts]: + if network_volume_id is not None: + if volume_mount_path is None: + raise ValueError("volume_mount_path must be set if network_volume_id is set") + network_mount: V2NetworkMount = { + "volumeId": network_volume_id, + "path": volume_mount_path, + } + mounts: V2Mounts = {"network": [network_mount]} + return mounts + elif volume_mount_path is not None: + raise ValueError("volume_mount_path must not be set if network_volume_id is not set") + return None def _generate_pod_deployment_mutation( @@ -526,206 +690,3 @@ def _generate_pod_deployment_mutation( }} }} """ - - -def _generate_cpu_pod_deployment_mutation( - name: str, - image_name: str, - instance_id: str, - cloud_type: str, - deploy_cost: float, - start_ssh: bool = True, - data_center_id: Optional[str] = None, - container_disk_in_gb: Optional[int] = None, - docker_args: Optional[str] = None, - ports: Optional[str] = None, - volume_mount_path: Optional[str] = None, - env: Optional[Dict[str, Any]] = None, - template_id: Optional[str] = None, - network_volume_id: Optional[str] = None, - container_registry_auth_id: Optional[str] = None, -) -> str: - """ - Generates a mutation to deploy CPU pod. - """ - input_fields = [] - input_fields.append(f'name: "{name}"') - input_fields.append(f'imageName: "{image_name}"') - input_fields.append(f'instanceId: "{instance_id}"') - input_fields.append(f"cloudType: {cloud_type}") - input_fields.append(f"deployCost: {deploy_cost}") - - if start_ssh: - input_fields.append("startSsh: true") - if data_center_id is not None: - input_fields.append(f'dataCenterId: "{data_center_id}"') - if container_disk_in_gb is not None: - input_fields.append(f"containerDiskInGb: {container_disk_in_gb}") - if docker_args is not None: - input_fields.append(f'dockerArgs: "{docker_args}"') - if ports is not None: - ports = ports.replace(" ", "") - input_fields.append(f'ports: "{ports}"') - if volume_mount_path is not None: - input_fields.append(f'volumeMountPath: "{volume_mount_path}"') - 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}]") - 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 container_registry_auth_id is not None: - input_fields.append(f'containerRegistryAuthId: "{container_registry_auth_id}"') - - input_string = ", ".join(input_fields) - return f""" - mutation {{ - deployCpuPod( - input: {{ - {input_string} - }} - ) {{ - id - lastStatusChange - imageName - 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}" }}) - }} - """ - - -def _generate_delete_cluster_mutation(cluster_id: str) -> str: - """ - Generates a mutation to delete a cluster. - """ - return f""" - mutation {{ - deleteCluster( - input: {{ - id: "{cluster_id}" - }} - ) - }} - """ - - -def _generate_create_cluster_mutation( - cluster_name: str, - gpu_type_id: str, - pod_count: int, - gpu_count_per_pod: int, - image_name: str, - cluster_type: str, - deploy_cost: str, - template_id: Optional[str] = None, - network_volume_id: Optional[str] = None, - volume_in_gb: Optional[int] = None, - throughput: Optional[int] = None, - allowed_cuda_versions: Optional[List[str]] = None, - volume_key: Optional[str] = None, - data_center_id: Optional[str] = None, - start_jupyter: bool = False, - start_ssh: bool = False, - container_disk_in_gb: Optional[int] = None, - docker_args: Optional[str] = None, - env: Optional[Dict[str, Any]] = None, - volume_mount_path: Optional[str] = None, - ports: Optional[str] = None, -) -> str: - """ - Generates a mutation to create a cluster. - """ - input_fields = [] - - # ------------------------------ Required Fields ----------------------------- # - input_fields.append(f'clusterName: "{cluster_name}"') - input_fields.append(f'gpuTypeId: "{gpu_type_id}"') - input_fields.append(f"podCount: {pod_count}") - input_fields.append(f'imageName: "{image_name}"') - input_fields.append(f"type: {cluster_type}") - input_fields.append(f"gpuCountPerPod: {gpu_count_per_pod}") - # If deploy_cost is not specified, Runpod returns Insufficient resources error. - input_fields.append(f"deployCost: {deploy_cost}") - - # ------------------------------ Optional Fields ----------------------------- # - 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 volume_in_gb is not None: - input_fields.append(f"volumeInGb: {volume_in_gb}") - if throughput is not None: - input_fields.append(f"throughput: {throughput}") - if allowed_cuda_versions is not None: - allowed_cuda_versions_string = ", ".join( - [f'"{version}"' for version in allowed_cuda_versions] - ) - input_fields.append(f"allowedCudaVersions: [{allowed_cuda_versions_string}]") - if volume_key is not None: - input_fields.append(f'volumeKey: "{volume_key}"') - if data_center_id is not None: - input_fields.append(f'dataCenterId: "{data_center_id}"') - if start_jupyter: - input_fields.append("startJupyter: true") - if start_ssh: - input_fields.append("startSsh: true") - if container_disk_in_gb is not None: - input_fields.append(f"containerDiskInGb: {container_disk_in_gb}") - if docker_args is not None: - input_fields.append(f'dockerArgs: "{docker_args}"') - 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}]") - 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}"') - - # Provisioning fails if minCudaVersion is specified for createCluster. - # See https://github.com/dstackai/dstack/issues/3910. - # TODO: Uncomment when RunPod fixes minCudaVersion for createCluster. - # - # input_fields.append(f'minCudaVersion: "{RunpodProvider.MIN_CUDA_VERSION}"') - - # Format input fields - input_string = ", ".join(input_fields) - return f""" - mutation {{ - createCluster( - input: {{ - {input_string} - }} - ) {{ - id - name - pods {{ - id - clusterIp - lastStatusChange - imageName - machine {{ - podHostId - }} - }} - }} - }} - """ diff --git a/src/dstack/_internal/core/backends/runpod/compute.py b/src/dstack/_internal/core/backends/runpod/compute.py index 5f248d3753..625c58a134 100644 --- a/src/dstack/_internal/core/backends/runpod/compute.py +++ b/src/dstack/_internal/core/backends/runpod/compute.py @@ -1,4 +1,5 @@ import json +import re import uuid from datetime import timedelta from typing import Callable, List, Optional @@ -17,10 +18,10 @@ ) from dstack._internal.core.backends.base.models import JobConfiguration from dstack._internal.core.backends.base.offers import get_catalog_offers -from dstack._internal.core.backends.runpod.api_client import RunpodApiClient, RunpodApiClientError +from dstack._internal.core.backends.runpod.api_client import RunpodApiClient, V2Pod from dstack._internal.core.backends.runpod.models import RunpodConfig from dstack._internal.core.consts import DSTACK_RUNNER_SSH_PORT -from dstack._internal.core.errors import ComputeError +from dstack._internal.core.errors import BackendError, ComputeError from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.common import CoreModel, RegistryAuth, validate_extra_ignore from dstack._internal.core.models.compute_groups import ComputeGroup, ComputeGroupProvisioningData @@ -50,6 +51,13 @@ # Runpod does not seem to have any limits on the disk size. CONFIGURABLE_DISK_SIZE = Range[Memory](min=Memory.parse("1GB"), max=None) +# gpuhunt, as of 0.1.30, builds CPU instance type names as +# {CPU flavor (aka type) identifier}-{vCPU count}-{memory in GB}, e.g., cpu3c-2-4 +# The deprecated GraphQL API (used by old dstack servers) accepts the whole string as +# `instanceId` field; REST API v2, which we currently use, accepts separate `id` and `vcpuCount` +# fields (memory is inferred) +CPU_INSTANCE_ID_REGEX = re.compile(r"(?P.+)-(?P\d+)-(?P\d+)") + class RunpodOfferBackendData(CoreModel): pod_counts: Optional[list[int]] = None @@ -128,8 +136,6 @@ def run_job( pod_name = generate_unique_instance_name(instance_config, max_length=MAX_RESOURCE_NAME_LEN) authorized_keys = build_authorized_keys(project_ssh_public_key, extra_authorized_keys) - memory_size = round(instance_offer.instance.resources.memory_mib / 1024) - disk_size = round(instance_offer.instance.resources.disk.size_mib / 1024) network_volume_id = None volume_mount_path = None @@ -139,71 +145,104 @@ def run_job( network_volume_id = volumes[0].volume_id volume_mount_path = run.run_spec.configuration.volumes[0].path - container_registry_auth_id = self._generate_container_registry_auth_id( + if instance_offer.instance.resources.spot: + raise ComputeError("Runpod spot offers are not supported") + + offer_gpu_count = len(instance_offer.instance.resources.gpus) + offer_cpu_count = instance_offer.instance.resources.cpus + offer_memory_gib = round(instance_offer.instance.resources.memory_mib / 1024) + offer_disk_gib = round(instance_offer.instance.resources.disk.size_mib / 1024) + + is_gpu_offer = offer_gpu_count > 0 + is_secure_cloud = _is_secure_cloud(instance_offer.region) + if not is_gpu_offer and not is_secure_cloud: + raise ComputeError("Runpod CPU offers are only supported in secure cloud regions") + + gpu_type_id: Optional[str] = None + gpu_count: Optional[int] = None + min_cpu_count_per_gpu: Optional[int] = None + min_memory_in_gb_per_gpu: Optional[int] = None + cpu_type_id: Optional[str] = None + cpu_count: Optional[int] = None + if is_gpu_offer: + gpu_type_id = instance_offer.instance.name + gpu_count = offer_gpu_count + # Offer resources are host totals. Round down, since these are lower bounds, and + # rounding up may exclude the host the offer is based on + min_cpu_count_per_gpu = offer_cpu_count // gpu_count + min_memory_in_gb_per_gpu = offer_memory_gib // gpu_count + else: + cpu_type_id = _get_cpu_type_id(instance_offer) + cpu_count = offer_cpu_count + + container_registry_credentials_id = self._generate_container_registry_credentials_id( job.job_spec.registry_auth ) - gpu_count = len(instance_offer.instance.resources.gpus) - if gpu_count == 0: - if not _is_secure_cloud(instance_offer.region): - raise ComputeError("Runpod CPU offers are only supported in secure cloud regions") - resp = self.api_client.create_cpu_pod( + + if is_secure_cloud: + # API v2 is used by default; its unified method supports both CPU and GPU pods + pod = self.api_client.create_pod( name=pod_name, - image_name=job.job_spec.image_name, - container_registry_auth_id=container_registry_auth_id, - instance_id=instance_offer.instance.name, - cloud_type="SECURE", - deploy_cost=instance_offer.price, - data_center_id=instance_offer.region, - container_disk_in_gb=disk_size, - start_ssh=True, - docker_args=_get_docker_args(authorized_keys), - ports=f"{DSTACK_RUNNER_SSH_PORT}/tcp", + cloud="SECURE", + data_center_ids=[instance_offer.region], + gpu_type_id=gpu_type_id, + gpu_count=gpu_count, + min_cpu_count_per_gpu=min_cpu_count_per_gpu, + min_memory_in_gb_per_gpu=min_memory_in_gb_per_gpu, + cpu_type_id=cpu_type_id, + cpu_count=cpu_count, + container_disk_in_gb=offer_disk_gib, network_volume_id=network_volume_id, volume_mount_path=volume_mount_path, + image_name=job.job_spec.image_name, + container_registry_credentials_id=container_registry_credentials_id, + entrypoint=["/bin/sh"], + cmd=["-c", " && ".join(get_docker_commands(authorized_keys))], + ports=[f"{DSTACK_RUNNER_SSH_PORT}/tcp"], env={"RUNPOD_POD_USER": "0"}, ) + instance_id = pod["id"] else: - bid_per_gpu = None - if instance_offer.instance.resources.spot: - bid_per_gpu = instance_offer.price / gpu_count - if _is_secure_cloud(instance_offer.region): - cloud_type = "SECURE" - data_center_id = instance_offer.region - country_code = None - else: - cloud_type = "COMMUNITY" - data_center_id = None - country_code = instance_offer.region - - resp = self.api_client.create_pod( + # Deprecated GraphQL API is used for Community Cloud, as API v2 does not support + # country codes + assert gpu_type_id is not None + assert gpu_count is not None + resp = self.api_client.create_gpu_pod_graphql( name=pod_name, - image_name=job.job_spec.image_name, - container_registry_auth_id=container_registry_auth_id, - gpu_type_id=instance_offer.instance.name, - cloud_type=cloud_type, - data_center_id=data_center_id, - country_code=country_code, - gpu_count=gpu_count, - container_disk_in_gb=disk_size, - min_vcpu_count=instance_offer.instance.resources.cpus, - min_memory_in_gb=memory_size, + cloud_type="COMMUNITY", + data_center_id=None, + country_code=instance_offer.region, + bid_per_gpu=None, support_public_ip=True, - docker_args=_get_docker_args(authorized_keys), - ports=f"{DSTACK_RUNNER_SSH_PORT}/tcp", - bid_per_gpu=bid_per_gpu, + gpu_type_id=gpu_type_id, + gpu_count=gpu_count, + min_vcpu_count=offer_cpu_count, + min_memory_in_gb=offer_memory_gib, + container_disk_in_gb=offer_disk_gib, network_volume_id=network_volume_id, volume_mount_path=volume_mount_path, + image_name=job.job_spec.image_name, + container_registry_auth_id=container_registry_credentials_id, + docker_args=_get_docker_args(authorized_keys), + ports=f"{DSTACK_RUNNER_SSH_PORT}/tcp", env={"RUNPOD_POD_USER": "0"}, ) - - instance_id = resp["id"] + instance_id = resp["id"] if ( self._last_cleanup_time is None or self._last_cleanup_time < get_current_datetime() - timedelta(seconds=CONTAINER_REGISTRY_AUTH_CLEANUP_INTERVAL) ): - self._clean_stale_container_registry_auths() + # The pod is already created, a failed cleanup must not fail provisioning + try: + self._clean_stale_container_registry_credentials() + except Exception as e: + logger.warning( + "Failed to clean up stale container registry credentials: %r", + e, + exc_info=not isinstance(e, BackendError), + ) self._last_cleanup_time = get_current_datetime() return JobProvisioningData( @@ -234,6 +273,8 @@ def run_jobs( ) -> ComputeGroupProvisioningData: master_job_configuration = job_configurations[0] master_job = master_job_configuration.job + if master_job.job_spec.registry_auth is not None: + raise ComputeError("Runpod clusters with private registries are not supported") master_job_volumes = master_job_configuration.volumes all_volumes_names = set(v.name for jc in job_configurations for v in jc.volumes) instance_config = InstanceConfiguration( @@ -243,7 +284,9 @@ def run_jobs( user=run.user, ) - pod_name = generate_unique_instance_name(instance_config, max_length=MAX_RESOURCE_NAME_LEN) + cluster_name = generate_unique_instance_name( + instance_config, max_length=MAX_RESOURCE_NAME_LEN + ) authorized_keys = build_authorized_keys(project_ssh_public_key, extra_authorized_keys) disk_size = round(instance_offer.instance.resources.disk.size_mib / 1024) @@ -269,37 +312,32 @@ def run_jobs( f"Failed to provision {pod_count} pods. Available pod counts: {offer_pod_counts}" ) - container_registry_auth_id = self._generate_container_registry_auth_id( - master_job.job_spec.registry_auth - ) - resp = self.api_client.create_cluster( - cluster_name=pod_name, + cluster = self.api_client.create_cluster( + name=cluster_name, + cluster_type="TRAINING", gpu_type_id=instance_offer.instance.name, - pod_count=pod_count, gpu_count_per_pod=gpu_count, - deploy_cost=f"{instance_offer.price * pod_count:.2f}", - image_name=master_job.job_spec.image_name, - cluster_type="TRAINING", - data_center_id=data_center_id, + pod_count=pod_count, + data_center_ids=[data_center_id], container_disk_in_gb=disk_size, - docker_args=_get_docker_args(authorized_keys), - ports=f"{DSTACK_RUNNER_SSH_PORT}/tcp", network_volume_id=network_volume_id, volume_mount_path=volume_mount_path, + image_name=master_job.job_spec.image_name, + entrypoint=["/bin/sh"], + cmd=["-c", " && ".join(get_docker_commands(authorized_keys))], + ports=[f"{DSTACK_RUNNER_SSH_PORT}/tcp"], env={"RUNPOD_POD_USER": "0"}, ) - - # Unlike create mutations for individual pods, createCluster mutation doesn't accept - # containerRegistryAuthId. - # The workaround is to inject containerRegistryAuthId into already created pods. - # Expect a long time (~5m) for the pods to pick up the creds. - # TODO: remove once createCluster supports containerRegistryAuthId - if container_registry_auth_id is not None: - for pod in resp["pods"]: - self.api_client.update_pod_container_registry_auth( - pod_id=pod["id"], - container_registry_auth_id=container_registry_auth_id, + try: + pods = self.api_client.list_cluster_pods(cluster["id"]) + if len(pods) != pod_count: + raise ComputeError( + f"Expected {pod_count} pods in cluster {cluster['id']}, got {len(pods)}" ) + except Exception: + # dstack does not track the cluster until run_jobs() returns + self._delete_cluster_after_failed_provisioning(cluster["id"]) + raise jpds = [ JobProvisioningData( @@ -307,17 +345,18 @@ def run_jobs( instance_type=instance_offer.instance, instance_id=pod["id"], hostname=None, - internal_ip=pod["clusterIp"], + # May not be assigned yet, see update_provisioning_data() + internal_ip=_get_pod_cluster_ip(pod), region=instance_offer.region, price=instance_offer.price, username="root", dockerized=False, ) - for pod in resp["pods"] + for pod in pods ] return ComputeGroupProvisioningData( - compute_group_id=resp["id"], - compute_group_name=resp["name"], + compute_group_id=cluster["id"], + compute_group_name=cluster["name"], backend=BackendType.RUNPOD, region=instance_offer.region, job_provisioning_datas=jpds, @@ -326,26 +365,15 @@ def run_jobs( def terminate_instance( self, instance_id: str, region: str, backend_data: Optional[str] = None ): - try: - self.api_client.terminate_pod(instance_id) - except RunpodApiClientError as e: - if len(e.errors) > 0 and e.errors[0]["message"] == "pod not found to terminate": - logger.debug("The instance %s not found. Skipping deletion.", instance_id) - return - raise + terminated = self.api_client.terminate_pod(instance_id) + if not terminated: + logger.debug("The instance %s not found. Skipping deletion.", instance_id) def terminate_compute_group(self, compute_group: ComputeGroup): - provisioning_data = compute_group.provisioning_data - try: - self.api_client.delete_cluster(provisioning_data.compute_group_id) - except RunpodApiClientError as e: - if len(e.errors) > 0 and e.errors[0]["extensions"]["code"] == "Cluster not found": - logger.debug( - "The cluster %s not found. Skipping deletion.", - provisioning_data.compute_group_id, - ) - return - raise + cluster_id = compute_group.provisioning_data.compute_group_id + deleted = self.api_client.delete_cluster(cluster_id) + if not deleted: + logger.debug("The cluster %s not found. Skipping deletion.", cluster_id) def update_provisioning_data( self, @@ -360,22 +388,28 @@ def update_provisioning_data( ports = pod["runtime"].get("ports") if ports is None: return - for port in pod["runtime"]["ports"]: - if port["privatePort"] == DSTACK_RUNNER_SSH_PORT: - provisioning_data.hostname = port["ip"] - provisioning_data.ssh_port = port["publicPort"] + if "cluster" in pod: + # The job copies the instance's provisioning data only until the hostname is set, + # so the internal IP must be set no later than the hostname + internal_ip = _get_pod_cluster_ip(pod) + if internal_ip is None: + return + provisioning_data.internal_ip = internal_ip + for port in ports: + if port.get("private") == DSTACK_RUNNER_SSH_PORT: + provisioning_data.hostname = port.get("ip") + provisioning_data.ssh_port = port.get("public") def register_volume(self, volume: Volume) -> VolumeProvisioningData: assert isinstance(volume.configuration, RunpodVolumeConfiguration) - volume_data = self.api_client.get_network_volume( - volume_id=get_or_error(volume.configuration.volume_id) - ) - if volume_data is None: - raise ComputeError(f"Volume {volume.configuration.volume_id} not found") - size_gb = volume_data["size"] + volume_id = get_or_error(volume.configuration.volume_id) + network_volume = self.api_client.get_network_volume(volume_id) + if network_volume is None: + raise ComputeError(f"Volume {volume_id} not found") + size_gb = network_volume["size"] return VolumeProvisioningData( backend=BackendType.RUNPOD, - volume_id=volume_data["id"], + volume_id=network_volume["id"], size_gb=size_gb, price=_get_volume_price(size_gb), attachable=False, @@ -388,15 +422,15 @@ def create_volume(self, volume: Volume) -> VolumeProvisioningData: size_gb = volume.configuration.size_gb # Runpod regions must be uppercase. # Lowercase regions are accepted in the API but they break Runpod in several ways. - region = volume.configuration.region.upper() - volume_id = self.api_client.create_network_volume( + data_center_id = volume.configuration.region.upper() + network_volume = self.api_client.create_network_volume( name=volume_name, - region=region, + data_center_id=data_center_id, size=size_gb, ) return VolumeProvisioningData( backend=BackendType.RUNPOD, - volume_id=volume_id, + volume_id=network_volume["id"], size_gb=size_gb, price=_get_volume_price(size_gb), attachable=False, @@ -404,40 +438,49 @@ def create_volume(self, volume: Volume) -> VolumeProvisioningData: ) def delete_volume(self, volume: Volume): - if volume.volume_id is not None: - try: - self.api_client.delete_network_volume(volume_id=volume.volume_id) - except RunpodApiClientError as e: - if ( - len(e.errors) > 0 - and "Tried to delete nonexistent network volume" in e.errors[0]["message"] - ): - logger.debug( - "The volume %s not found. Skipping deletion.", - volume.volume_id, - ) - return - raise - - def _generate_container_registry_auth_id( + if volume.volume_id is None: + return + deleted = self.api_client.delete_network_volume(volume.volume_id) + if not deleted: + logger.debug("The volume %s not found. Skipping deletion.", volume.volume_id) + + def _generate_container_registry_credentials_id( self, registry_auth: Optional[RegistryAuth] ) -> Optional[str]: if registry_auth is None: return None - return self.api_client.add_container_registry_auth( - uuid.uuid4().hex, registry_auth.username, registry_auth.password + creds = self.api_client.create_container_registry_credentials( + name=uuid.uuid4().hex, + username=registry_auth.username, + password=registry_auth.password, ) - - def _clean_stale_container_registry_auths(self) -> None: - container_registry_auths = self.api_client.get_container_registry_auths() - # Container_registry_auths sorted by creation time so try to delete the oldest first - # when we reach container_registry_auths that is still in use, we stop - for container_registry_auth in container_registry_auths: + return creds["id"] + + def _clean_stale_container_registry_credentials(self) -> None: + # XXX: Credentials are sorted by creation time so try to delete the oldest first. + # When we reach credentials that is still in use, we stop. + # > Permanently deletes a container registry credential by ID. Rejected if any pod + # > currently uses this credential to pull its image. Templates that reference it are not + # > part of that check — they silently lose the reference (registry becomes null) instead + # > of blocking the delete. + for creds in self.api_client.list_container_registry_credentials(): try: - self.api_client.delete_container_registry_auth(container_registry_auth["id"]) + self.api_client.delete_container_registry_credentials(creds["id"]) except Exception: break + def _delete_cluster_after_failed_provisioning(self, cluster_id: str) -> None: + try: + self.api_client.delete_cluster(cluster_id) + except Exception as e: + logger.error( + "Failed to delete cluster %s after failed provisioning." + " Please delete it manually to avoid unexpected charges. Error: %r", + cluster_id, + e, + exc_info=not isinstance(e, BackendError), + ) + def _get_docker_args(authorized_keys: List[str]) -> str: commands = get_docker_commands(authorized_keys) @@ -465,3 +508,17 @@ def _get_offer_pod_counts(offer: InstanceOfferWithAvailability) -> list[int]: backend_data = validate_extra_ignore(RunpodOfferBackendData, offer.backend_data) pod_counts = backend_data.pod_counts or [] return pod_counts + + +def _get_cpu_type_id(offer: InstanceOfferWithAvailability) -> str: + match = CPU_INSTANCE_ID_REGEX.fullmatch(offer.instance.name) + if match is None: + raise ComputeError(f"Failed to parse CPU instance type: {offer.instance.name!r}") + return match.group("flavor_id") + + +def _get_pod_cluster_ip(pod: V2Pod) -> Optional[str]: + cluster = pod.get("cluster") + if cluster is None: + return None + return cluster.get("ip") diff --git a/src/tests/_internal/core/backends/runpod/test_api_client.py b/src/tests/_internal/core/backends/runpod/test_api_client.py index 9de7ebed47..712aa3c124 100644 --- a/src/tests/_internal/core/backends/runpod/test_api_client.py +++ b/src/tests/_internal/core/backends/runpod/test_api_client.py @@ -1,71 +1,195 @@ +import pytest + from dstack._internal.core.backends.runpod.api_client import ( + API_V2_URL, RunpodApiClient, - _generate_cpu_pod_deployment_mutation, + RunpodV2ClientError, ) +from dstack._internal.core.errors import BackendInvalidCredentialsError + + +def _client() -> RunpodApiClient: + return RunpodApiClient(api_key="test") + + +def _error(status: int, detail: str, errors: list[str] | None = None) -> dict: + error = {"title": "Error", "status": status, "detail": detail} + if errors is not None: + error["errors"] = errors + return error + + +class TestRunpodApiClientCreatePod: + def test_gpu_pod(self, requests_mock): + requests_mock.post(f"{API_V2_URL}/pods", status_code=201, json={"id": "pod-1"}) + + pod = _client().create_pod( + name="test-pod", + cloud="SECURE", + data_center_ids=["EU-RO-1"], + gpu_type_id="NVIDIA H100 80GB HBM3", + gpu_count=2, + min_cpu_count_per_gpu=8, + min_memory_in_gb_per_gpu=125, + container_disk_in_gb=100, + network_volume_id="vol-1", + volume_mount_path="/data", + image_name="dstackai/base:latest", + container_registry_credentials_id="creds-1", + entrypoint=["/bin/sh"], + cmd=["-c", "echo hi"], + ports=["10022/tcp"], + env={"RUNPOD_POD_USER": "0"}, + ) + + assert pod["id"] == "pod-1" + assert requests_mock.last_request.headers["Authorization"] == "Bearer test" + assert requests_mock.last_request.json() == { + "name": "test-pod", + "cloud": "SECURE", + "dataCenterIds": ["EU-RO-1"], + "gpu": { + "id": "NVIDIA H100 80GB HBM3", + "count": 2, + "minCudaVersion": "12.8", + "minVcpuCountPerGpu": 8, + "minRamPerGpu": 125, + }, + "disk": 100, + "mounts": {"network": [{"volumeId": "vol-1", "path": "/data"}]}, + "image": "dstackai/base:latest", + "registry": "creds-1", + "entrypoint": ["/bin/sh"], + "cmd": ["-c", "echo hi"], + "ports": ["10022/tcp"], + "env": {"RUNPOD_POD_USER": "0"}, + } + + def test_cpu_pod(self, requests_mock): + requests_mock.post(f"{API_V2_URL}/pods", status_code=201, json={"id": "pod-1"}) + + _client().create_pod( + name="test-pod", + cloud="SECURE", + cpu_type_id="cpu3c", + cpu_count=4, + container_disk_in_gb=20, + image_name="dstackai/base:latest", + entrypoint=["/bin/sh"], + cmd=["-c", "echo hi"], + ) + + assert requests_mock.last_request.json() == { + "name": "test-pod", + "cloud": "SECURE", + "cpu": {"id": "cpu3c", "vcpuCount": 4}, + "disk": 20, + "image": "dstackai/base:latest", + "entrypoint": ["/bin/sh"], + "cmd": ["-c", "echo hi"], + } + + def test_rejects_both_gpu_and_cpu(self): + with pytest.raises(ValueError, match="Exactly one"): + _client().create_pod( + name="test-pod", + cloud="SECURE", + gpu_type_id="NVIDIA H100 80GB HBM3", + cpu_type_id="cpu3c", + container_disk_in_gb=20, + image_name="dstackai/base:latest", + entrypoint=["/bin/sh"], + cmd=["-c", "echo hi"], + ) + + +class TestRunpodApiClientErrors: + def test_v2_error_response(self, requests_mock): + requests_mock.post( + f"{API_V2_URL}/registries", + status_code=422, + json=_error(422, "Request validation failed.", ["$: missing property 'name'"]), + ) + + with pytest.raises(RunpodV2ClientError) as exc_info: + _client().create_container_registry_credentials(name="", username="u", password="p") + + assert exc_info.value.status == 422 + assert exc_info.value.detail == "Request validation failed." + assert exc_info.value.errors == ["$: missing property 'name'"] + assert str(exc_info.value) == ( + "422: Error: Request validation failed: [\"$: missing property 'name'\"]" + ) + + @pytest.mark.parametrize("status_code", [401, 403]) + def test_auth_error_response(self, requests_mock, status_code: int): + requests_mock.get( + f"{API_V2_URL}/account/ssh-keys", + status_code=status_code, + json=_error(status_code, "Invalid token"), + ) + + assert _client().validate_api_key() is False + + def test_non_json_auth_error_response(self, requests_mock): + requests_mock.get( + f"{API_V2_URL}/pods/pod-1", + status_code=403, + reason="Forbidden", + text="Access denied", + ) + + with pytest.raises(BackendInvalidCredentialsError, match="Forbidden: "): + _client().get_pod("pod-1") + + def test_non_json_error_response(self, requests_mock): + requests_mock.get( + f"{API_V2_URL}/pods/pod-1", + status_code=429, + reason="Too Many Requests", + text="\n" + "x" * 1000 + "\n", + ) + + with pytest.raises(RunpodV2ClientError) as exc_info: + _client().get_pod("pod-1") + + assert exc_info.value.status == 429 + assert exc_info.value.title == "Too Many Requests" + assert exc_info.value.detail.startswith(" xxx") + assert len(exc_info.value.detail) == 200 + + +class TestRunpodApiClientNotFound: + def test_get_pod(self, requests_mock): + requests_mock.get(f"{API_V2_URL}/pods/pod-1", status_code=404, json=_error(404, "-")) + + assert _client().get_pod("pod-1") is None + + def test_terminate_pod(self, requests_mock): + requests_mock.delete(f"{API_V2_URL}/pods/pod-1", status_code=404, json=_error(404, "-")) + + assert _client().terminate_pod("pod-1") is False + + def test_delete_cluster(self, requests_mock): + requests_mock.delete( + f"{API_V2_URL}/clusters/cluster-1", status_code=404, json=_error(404, "-") + ) + + assert _client().delete_cluster("cluster-1") is False + + def test_delete_network_volume(self, requests_mock): + requests_mock.delete( + f"{API_V2_URL}/network-volumes/vol-1", status_code=404, json=_error(404, "-") + ) + + assert _client().delete_network_volume("vol-1") is False + def test_other_errors_are_raised(self, requests_mock): + requests_mock.delete( + f"{API_V2_URL}/pods/pod-1", + status_code=409, + json=_error(409, "Pod belongs to a cluster"), + ) -class _Response: - def __init__(self, payload): - self._payload = payload - - def json(self): - return self._payload - - -def test_generate_cpu_pod_deployment_mutation(): - mutation = _generate_cpu_pod_deployment_mutation( - name="cpu-test", - image_name="python:3.11-slim", - instance_id="cpu3g-2-8", - cloud_type="SECURE", - deploy_cost=0.08, - start_ssh=True, - data_center_id="AP-JP-1", - container_disk_in_gb=5, - docker_args='{"cmd":["echo hi"]}', - ports="22/tcp, 8080/http", - volume_mount_path="/workspace", - env={"HELLO": "WORLD"}, - template_id="runpod-ubuntu", - network_volume_id="vol-1", - container_registry_auth_id="cred-1", - ) - - assert "deployCpuPod" in mutation - assert 'name: "cpu-test"' in mutation - assert 'imageName: "python:3.11-slim"' in mutation - assert 'instanceId: "cpu3g-2-8"' in mutation - assert "cloudType: SECURE" in mutation - assert "deployCost: 0.08" in mutation - assert "startSsh: true" in mutation - assert 'dataCenterId: "AP-JP-1"' in mutation - assert "containerDiskInGb: 5" in mutation - assert 'ports: "22/tcp,8080/http"' in mutation - assert 'volumeMountPath: "/workspace"' in mutation - assert 'env: [{ key: "HELLO", value: "WORLD" }]' in mutation - assert 'templateId: "runpod-ubuntu"' in mutation - assert 'networkVolumeId: "vol-1"' in mutation - assert 'containerRegistryAuthId: "cred-1"' in mutation - - -def test_create_cpu_pod_uses_deploy_cpu_pod(monkeypatch): - client = RunpodApiClient(api_key="test") - query = {} - - def fake_make_request(data): - query["value"] = data["query"] - return _Response({"data": {"deployCpuPod": {"id": "cpu-pod-1"}}}) - - monkeypatch.setattr(client, "_make_request", fake_make_request) - - response = client.create_cpu_pod( - name="cpu-test", - image_name="python:3.11-slim", - instance_id="cpu3g-2-8", - cloud_type="SECURE", - deploy_cost=0.08, - ) - - assert response["id"] == "cpu-pod-1" - assert "deployCpuPod" in query["value"] - assert "podFindAndDeployOnDemand" not in query["value"] + with pytest.raises(RunpodV2ClientError): + _client().terminate_pod("pod-1") diff --git a/src/tests/_internal/core/backends/runpod/test_compute.py b/src/tests/_internal/core/backends/runpod/test_compute.py new file mode 100644 index 0000000000..dd7e982ddb --- /dev/null +++ b/src/tests/_internal/core/backends/runpod/test_compute.py @@ -0,0 +1,258 @@ +from typing import Optional +from unittest.mock import MagicMock, patch + +import pytest + +from dstack._internal.core.backends.runpod.api_client import RunpodV2ClientError +from dstack._internal.core.backends.runpod.compute import RunpodCompute +from dstack._internal.core.backends.runpod.models import RunpodAPIKeyCreds, RunpodConfig +from dstack._internal.core.errors import ComputeError +from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.instances import InstanceOfferWithAvailability +from dstack._internal.core.models.resources import ResourcesSpec +from dstack._internal.core.models.runs import JobProvisioningData, Requirements +from dstack._internal.server.testing.common import ( + get_instance_offer_with_availability, + get_job_provisioning_data, +) + +# build_authorized_keys() rejects the project key unless it actually parses +PROJECT_SSH_PUBLIC_KEY = ( + "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAINOmx0T+hBRaJ6jCi21ZYe2NW3EZS8e0Mdwl+yZJt+kD project" +) + + +def _compute() -> RunpodCompute: + compute = RunpodCompute(RunpodConfig(creds=RunpodAPIKeyCreds(api_key="test"))) + compute.api_client = MagicMock() + return compute + + +def _gpu_offer( + *, gpu_count: int = 2, cpu_count: int = 15, memory_gib: float = 251 +) -> InstanceOfferWithAvailability: + return get_instance_offer_with_availability( + backend=BackendType.RUNPOD, + region="EU-RO-1", + gpu_count=gpu_count, + gpu_name="H100", + cpu_count=cpu_count, + memory_gib=memory_gib, + instance_type="NVIDIA H100 80GB HBM3", + ) + + +def _cluster_offer(pod_counts: list[int]) -> InstanceOfferWithAvailability: + offer = _gpu_offer(gpu_count=8, cpu_count=128, memory_gib=1024) + offer.backend_data = {"pod_counts": pod_counts} + return offer + + +def _job() -> MagicMock: + job = MagicMock() + job.job_spec.job_name = "test-job-0" + job.job_spec.image_name = "dstackai/base:latest" + job.job_spec.registry_auth = None + return job + + +def _run() -> MagicMock: + run = MagicMock() + run.project_name = "main" + run.user = "test-user" + return run + + +def _run_job(compute: RunpodCompute, offer: InstanceOfferWithAvailability) -> JobProvisioningData: + with patch( + "dstack._internal.core.backends.runpod.compute.get_docker_commands", + return_value=["echo hi"], + ): + return compute.run_job( + run=_run(), + job=_job(), + instance_offer=offer, + project_ssh_public_key=PROJECT_SSH_PUBLIC_KEY, + project_ssh_private_key="private-key", + volumes=[], + placement_group=None, + requirements=Requirements(resources=ResourcesSpec()), + extra_authorized_keys=[], + ) + + +def _run_jobs(compute: RunpodCompute, offer: InstanceOfferWithAvailability, job_count: int): + job_configurations = [] + for _ in range(job_count): + job_configuration = MagicMock() + job_configuration.job = _job() + job_configuration.volumes = [] + job_configurations.append(job_configuration) + with patch( + "dstack._internal.core.backends.runpod.compute.get_docker_commands", + return_value=["echo hi"], + ): + return compute.run_jobs( + run=_run(), + job_configurations=job_configurations, + instance_offer=offer, + project_ssh_public_key=PROJECT_SSH_PUBLIC_KEY, + project_ssh_private_key="private-key", + placement_group=None, + requirements=Requirements(resources=ResourcesSpec(), multinode=True), + extra_authorized_keys=[], + ) + + +def _pod( + pod_id: str = "pod-1", + *, + ports: Optional[list[dict]] = None, + cluster_ip: Optional[str] = None, + is_cluster_pod: bool = False, +) -> dict: + pod: dict = { + "id": pod_id, + "status": "RUNNING", + "runtime": None if ports is None else {"ports": ports}, + "mounts": {}, + } + if is_cluster_pod or cluster_ip is not None: + pod["cluster"] = {"id": "cluster-1", "rank": 0} + if cluster_ip is not None: + pod["cluster"]["ip"] = cluster_ip + return pod + + +SSH_PORTS = [ + {"ip": "100.65.0.1", "private": 8080, "public": None, "type": "http"}, + {"ip": "203.0.113.10", "private": 10022, "public": 40022, "type": "tcp"}, +] + + +class TestRunpodComputeRunJob: + def test_requests_per_gpu_minimums_rounded_down(self): + compute = _compute() + compute.api_client.create_pod.return_value = {"id": "pod-1"} + + jpd = _run_job(compute, _gpu_offer(gpu_count=2, cpu_count=15, memory_gib=251)) + + assert jpd.instance_id == "pod-1" + kwargs = compute.api_client.create_pod.call_args.kwargs + assert kwargs["cloud"] == "SECURE" + assert kwargs["data_center_ids"] == ["EU-RO-1"] + assert kwargs["gpu_type_id"] == "NVIDIA H100 80GB HBM3" + assert kwargs["gpu_count"] == 2 + assert kwargs["min_cpu_count_per_gpu"] == 7 + assert kwargs["min_memory_in_gb_per_gpu"] == 125 + + def test_failed_registry_credentials_cleanup_does_not_fail_provisioning(self): + compute = _compute() + compute.api_client.create_pod.return_value = {"id": "pod-1"} + compute.api_client.list_container_registry_credentials.side_effect = RunpodV2ClientError( + status=429, title="Too Many Requests", detail="Rate limit exceeded" + ) + + jpd = _run_job(compute, _gpu_offer()) + + assert jpd.instance_id == "pod-1" + compute.api_client.terminate_pod.assert_not_called() + + +class TestRunpodComputeRunJobs: + def test_sets_internal_ip_from_cluster_ip(self): + compute = _compute() + compute.api_client.create_cluster.return_value = {"id": "cluster-1", "name": "test"} + compute.api_client.list_cluster_pods.return_value = [ + _pod("pod-1", cluster_ip="10.65.0.2"), + _pod("pod-2", is_cluster_pod=True), + ] + + cgpd = _run_jobs(compute, _cluster_offer([2, 4]), job_count=2) + + assert cgpd.compute_group_id == "cluster-1" + assert [jpd.instance_id for jpd in cgpd.job_provisioning_datas] == ["pod-1", "pod-2"] + assert [jpd.internal_ip for jpd in cgpd.job_provisioning_datas] == ["10.65.0.2", None] + + def test_deletes_cluster_if_pod_count_mismatch(self): + compute = _compute() + compute.api_client.create_cluster.return_value = {"id": "cluster-1", "name": "test"} + compute.api_client.list_cluster_pods.return_value = [_pod("pod-1", is_cluster_pod=True)] + + with pytest.raises(ComputeError, match="Expected 2 pods"): + _run_jobs(compute, _cluster_offer([2, 4]), job_count=2) + + compute.api_client.delete_cluster.assert_called_once_with("cluster-1") + + def test_deletes_cluster_if_listing_pods_fails(self): + compute = _compute() + compute.api_client.create_cluster.return_value = {"id": "cluster-1", "name": "test"} + error = RunpodV2ClientError(status=429, title="Too Many Requests", detail="-") + compute.api_client.list_cluster_pods.side_effect = error + + with pytest.raises(RunpodV2ClientError): + _run_jobs(compute, _cluster_offer([2, 4]), job_count=2) + + compute.api_client.delete_cluster.assert_called_once_with("cluster-1") + + def test_raises_original_error_if_cluster_deletion_fails(self): + compute = _compute() + compute.api_client.create_cluster.return_value = {"id": "cluster-1", "name": "test"} + compute.api_client.list_cluster_pods.return_value = [] + compute.api_client.delete_cluster.side_effect = RuntimeError("delete failed") + + with pytest.raises(ComputeError, match="Expected 2 pods"): + _run_jobs(compute, _cluster_offer([2, 4]), job_count=2) + + +class TestRunpodComputeUpdateProvisioningData: + def _update(self, compute: RunpodCompute, jpd: JobProvisioningData) -> None: + compute.update_provisioning_data( + provisioning_data=jpd, + project_ssh_public_key=PROJECT_SSH_PUBLIC_KEY, + project_ssh_private_key="private-key", + ) + + def test_sets_hostname_and_ssh_port(self): + compute = _compute() + compute.api_client.get_pod.return_value = _pod(ports=SSH_PORTS) + jpd = get_job_provisioning_data(backend=BackendType.RUNPOD, hostname=None) + + self._update(compute, jpd) + + assert jpd.hostname == "203.0.113.10" + assert jpd.ssh_port == 40022 + + def test_waits_for_runtime(self): + compute = _compute() + compute.api_client.get_pod.return_value = _pod(ports=None) + jpd = get_job_provisioning_data(backend=BackendType.RUNPOD, hostname=None) + + self._update(compute, jpd) + + assert jpd.hostname is None + + def test_cluster_pod_waits_for_cluster_ip(self): + compute = _compute() + compute.api_client.get_pod.return_value = _pod(ports=SSH_PORTS, is_cluster_pod=True) + jpd = get_job_provisioning_data( + backend=BackendType.RUNPOD, hostname=None, internal_ip=None + ) + + self._update(compute, jpd) + + assert jpd.hostname is None + assert jpd.internal_ip is None + + def test_cluster_pod_sets_internal_ip(self): + compute = _compute() + compute.api_client.get_pod.return_value = _pod(ports=SSH_PORTS, cluster_ip="10.65.0.3") + jpd = get_job_provisioning_data( + backend=BackendType.RUNPOD, hostname=None, internal_ip=None + ) + + self._update(compute, jpd) + + assert jpd.hostname == "203.0.113.10" + assert jpd.ssh_port == 40022 + assert jpd.internal_ip == "10.65.0.3"