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"