skypilot-nightly 1.0.0.dev2024053101__py3-none-any.whl → 1.0.0.dev2025022801__py3-none-any.whl

sky/provision/runpod/api/commands.py

@@ -0,0 +1,119 @@
+"""This module provides functions to generate GraphQL mutations for deploying
+spot instance Pods on RunPod.
+
+Reference:
+    https://github.com/runpod/runpod-python/blob/main/runpod/api/ctl_commands.py
+
+Functions:
+    generate_spot_pod_deployment_mutation: Generates a GraphQL mutation string
+        for deploying a spot instance Pod on RunPod.
+
+Example:
+    >>> mutation = generate_spot_pod_deployment_mutation(
+            name='test',
+            image_name='runpod/stack',
+            gpu_type_id='NVIDIA GeForce RTX 3070',
+            bid_per_gpu=0.3
+        )
+"""
+from typing import List, Optional
+
+from sky.adaptors import runpod
+from sky.provision.runpod.api.pods import generate_spot_pod_deployment_mutation
+
+_INTERRUPTABLE_POD_FIELD: str = 'podRentInterruptable'
+_RESPONSE_DATA_FIELD: str = 'data'
+
+
+def create_spot_pod(
+    name: str,
+    image_name: str,
+    gpu_type_id: str,
+    bid_per_gpu: float,
+    cloud_type: str = 'ALL',
+    volume_mount_path: str = '/runpod-volume',
+    gpu_count: Optional[int] = 1,
+    min_memory_in_gb: Optional[int] = 1,
+    min_vcpu_count: Optional[int] = 1,
+    container_disk_in_gb: Optional[int] = None,
+    volume_in_gb: Optional[int] = 0,
+    ports: Optional[str] = None,
+    start_ssh: Optional[bool] = True,
+    start_jupyter: Optional[bool] = False,
+    env: Optional[dict] = None,
+    docker_args: Optional[str] = '',
+    support_public_ip: Optional[bool] = True,
+    terminate_after: Optional[str] = None,
+    stop_after: Optional[str] = None,
+    data_center_id: Optional[str] = None,
+    country_code: Optional[str] = None,
+    network_volume_id: Optional[str] = None,
+    allowed_cuda_versions: Optional[List[str]] = None,
+    min_download: Optional[int] = None,
+    min_upload: Optional[int] = None,
+    cuda_version: Optional[str] = None,
+    template_id: Optional[str] = None,
+    volume_key: Optional[str] = None,
+) -> dict:
+    """This module provides functions to generate GraphQL mutations for
+    deploying spot instance Pods on RunPod.
+
+    Functions:
+        generate_spot_pod_deployment_mutation: Generates a GraphQL mutation
+            string for deploying a spot instance Pod on RunPod.
+
+    Example:
+        >>> mutation = generate_spot_pod_deployment_mutation(
+                name='test',
+                image_name='runpod/stack',
+                gpu_type_id='NVIDIA GeForce RTX 3070',
+                bid_per_gpu=0.3
+            )
+    """
+    runpod.runpod.get_gpu(gpu_type_id)
+    # refer to https://graphql-spec.runpod.io/#definition-CloudTypeEnum
+    if cloud_type not in ['ALL', 'COMMUNITY', 'SECURE']:
+        raise ValueError('cloud_type must be one of ALL, COMMUNITY or SECURE')
+
+    if network_volume_id and data_center_id is None:
+        user_info = runpod.runpod.get_user()
+        for network_volume in user_info['networkVolumes']:
+            if network_volume['id'] == network_volume_id:
+                data_center_id = network_volume['dataCenterId']
+                break
+
+    if container_disk_in_gb is None and template_id is None:
+        container_disk_in_gb = 10
+
+    mutation = generate_spot_pod_deployment_mutation(
+        name=name,
+        image_name=image_name,
+        gpu_type_id=gpu_type_id,
+        bid_per_gpu=bid_per_gpu,
+        cloud_type=cloud_type,
+        gpu_count=gpu_count,
+        min_memory_in_gb=min_memory_in_gb,
+        min_vcpu_count=min_vcpu_count,
+        container_disk_in_gb=container_disk_in_gb,
+        volume_in_gb=volume_in_gb,
+        volume_mount_path=volume_mount_path,
+        ports=ports,
+        start_ssh=start_ssh,
+        start_jupyter=start_jupyter,
+        env=env,
+        docker_args=docker_args,
+        support_public_ip=support_public_ip,
+        terminate_after=terminate_after,
+        stop_after=stop_after,
+        data_center_id=data_center_id,
+        country_code=country_code,
+        network_volume_id=network_volume_id,
+        allowed_cuda_versions=allowed_cuda_versions,
+        min_download=min_download,
+        min_upload=min_upload,
+        cuda_version=cuda_version,
+        template_id=template_id,
+        volume_key=volume_key,
+    )
+    response = runpod.runpod.api.graphql.run_graphql_query(mutation)
+    return response[_RESPONSE_DATA_FIELD][_INTERRUPTABLE_POD_FIELD]

sky/provision/runpod/api/pods.py

@@ -0,0 +1,142 @@
+"""This module provides functions to generate GraphQL mutations for deploying
+spot instance Pods on RunPod.
+
+Reference:
+    https://github.com/runpod/runpod-python/blob/main/runpod/api/mutations/pods.py
+
+Functions:
+    generate_spot_pod_deployment_mutation: Generates a GraphQL mutation string
+        for deploying a spot instance Pod on RunPod.
+Example:
+    >>> mutation = generate_spot_pod_deployment_mutation(
+            name='test',
+            image_name='runpod/stack',
+            gpu_type_id='NVIDIA GeForce RTX 3070',
+            bid_per_gpu=0.3
+        )
+"""
+
+from typing import List, Optional
+
+
+# refer to https://graphql-spec.runpod.io/#definition-PodRentInterruptableInput
+def generate_spot_pod_deployment_mutation(
+    name: str,
+    image_name: str,
+    gpu_type_id: str,
+    bid_per_gpu: float,
+    volume_mount_path: str,
+    cloud_type: str = 'ALL',
+    gpu_count: Optional[int] = None,
+    min_memory_in_gb: Optional[int] = None,
+    min_vcpu_count: Optional[int] = None,
+    container_disk_in_gb: Optional[int] = None,
+    volume_in_gb: Optional[int] = None,
+    ports: Optional[str] = None,
+    start_ssh: Optional[bool] = True,
+    start_jupyter: Optional[bool] = False,
+    env: Optional[dict] = None,
+    docker_args: Optional[str] = None,
+    support_public_ip: Optional[bool] = True,
+    terminate_after: Optional[str] = None,
+    stop_after: Optional[str] = None,
+    data_center_id: Optional[str] = None,
+    country_code: Optional[str] = None,
+    network_volume_id: Optional[str] = None,
+    allowed_cuda_versions: Optional[List[str]] = None,
+    min_download: Optional[int] = None,
+    min_upload: Optional[int] = None,
+    cuda_version: Optional[str] = None,
+    template_id: Optional[str] = None,
+    volume_key: Optional[str] = None,
+) -> str:
+    input_fields = []
+
+    # Required Fields
+    input_fields.append(f'name: "{name}"')
+    input_fields.append(f'imageName: "{image_name}"')
+    input_fields.append(f'gpuTypeId: "{gpu_type_id}"')
+    input_fields.append(f'bidPerGpu: {bid_per_gpu}')
+    input_fields.append(f'volumeMountPath: "{volume_mount_path}"')
+
+    # Default Fields
+    input_fields.append(f'cloudType: {cloud_type}')
+
+    if start_ssh:
+        input_fields.append('startSsh: true')
+    if start_jupyter:
+        input_fields.append('startJupyter: true')
+    if support_public_ip:
+        input_fields.append('supportPublicIp: true')
+    else:
+        input_fields.append('supportPublicIp: false')
+
+    # Optional Fields
+    if gpu_count is not None:
+        input_fields.append(f'gpuCount: {gpu_count}')
+    if min_memory_in_gb is not None:
+        input_fields.append(f'minMemoryInGb: {min_memory_in_gb}')
+    if min_vcpu_count is not None:
+        input_fields.append(f'minVcpuCount: {min_vcpu_count}')
+    if container_disk_in_gb is not None:
+        input_fields.append(f'containerDiskInGb: {container_disk_in_gb}')
+    if volume_in_gb is not None:
+        input_fields.append(f'volumeInGb: {volume_in_gb}')
+    if ports is not None:
+        ports = ports.replace(' ', '')
+        input_fields.append(f'ports: "{ports}"')
+    if docker_args is not None:
+        input_fields.append(f'dockerArgs: "{docker_args}"')
+    if terminate_after is not None:
+        input_fields.append(f'terminateAfter: "{terminate_after}"')
+    if stop_after is not None:
+        input_fields.append(f'stopAfter: "{stop_after}"')
+    if data_center_id is not None:
+        input_fields.append(f'dataCenterId: "{data_center_id}"')
+    if country_code is not None:
+        input_fields.append(f'countryCode: "{country_code}"')
+    if network_volume_id is not None:
+        input_fields.append(f'networkVolumeId: "{network_volume_id}"')
+    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 min_download is not None:
+        input_fields.append(f'minDownload: {min_download}')
+    if min_upload is not None:
+        input_fields.append(f'minUpload: {min_upload}')
+    if cuda_version is not None:
+        input_fields.append(f'cudaVersion: "{cuda_version}"')
+    if template_id is not None:
+        input_fields.append(f'templateId: "{template_id}"')
+    if volume_key is not None:
+        input_fields.append(f'volumeKey: "{volume_key}"')
+
+    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}]')
+
+    # Format input fields
+    input_string = ', '.join(input_fields)
+    return f"""
+    mutation {{
+      podRentInterruptable(
+        input: {{
+          {input_string}
+        }}
+      ) {{
+        id
+        desiredStatus
+        imageName
+        env
+        machineId
+        machine {{
+          podHostId
+        }}
+      }}
+    }}
+    """

sky/provision/runpod/instance.py

@@ -3,24 +3,27 @@ import time
 from typing import Any, Dict, List, Optional
 
 from sky import sky_logging
-from sky import status_lib
 from sky.provision import common
 from sky.provision.runpod import utils
 from sky.utils import common_utils
+from sky.utils import resources_utils
+from sky.utils import status_lib
 from sky.utils import ux_utils
 
 POLL_INTERVAL = 5
+QUERY_PORTS_TIMEOUT_SECONDS = 30
 
 logger = sky_logging.init_logger(__name__)
 
 
 def _filter_instances(cluster_name_on_cloud: str,
-                      status_filters: Optional[List[str]]) -> Dict[str, Any]:
+                      status_filters: Optional[List[str]],
+                      head_only: bool = False) -> Dict[str, Any]:
 
     instances = utils.list_instances()
-    possible_names = [
-        f'{cluster_name_on_cloud}-head', f'{cluster_name_on_cloud}-worker'
-    ]
+    possible_names = [f'{cluster_name_on_cloud}-head']
+    if not head_only:
+        possible_names.append(f'{cluster_name_on_cloud}-worker')
 
     filtered_instances = {}
     for instance_id, instance in instances.items():
@@ -80,10 +83,19 @@ def run_instances(region: str, cluster_name_on_cloud: str,
         node_type = 'head' if head_instance_id is None else 'worker'
         try:
             instance_id = utils.launch(
-                name=f'{cluster_name_on_cloud}-{node_type}',
+                cluster_name=cluster_name_on_cloud,
+                node_type=node_type,
                 instance_type=config.node_config['InstanceType'],
                 region=region,
-                disk_size=config.node_config['DiskSize'])
+                disk_size=config.node_config['DiskSize'],
+                image_name=config.node_config['ImageId'],
+                ports=config.ports_to_open_on_launch,
+                public_key=config.node_config['PublicKey'],
+                preemptible=config.node_config['Preemptible'],
+                bid_per_gpu=config.node_config['BidPerGPU'],
+                docker_login_config=config.provider_config.get(
+                    'docker_login_config'),
+            )
         except Exception as e:  # pylint: disable=broad-except
             logger.warning(f'run_instances error: {e}')
             raise
@@ -136,6 +148,8 @@ def terminate_instances(
     """See sky/provision/__init__.py"""
     del provider_config  # unused
     instances = _filter_instances(cluster_name_on_cloud, None)
+    template_name, registry_auth_id = utils.get_registry_auth_resources(
+        cluster_name_on_cloud)
     for inst_id, inst in instances.items():
         logger.debug(f'Terminating instance {inst_id}: {inst}')
         if worker_only and inst['name'].endswith('-head'):
@@ -148,6 +162,10 @@ def terminate_instances(
                     f'Failed to terminate instance {inst_id}: '
                     f'{common_utils.format_exception(e, use_bracket=False)}'
                 ) from e
+    if template_name is not None:
+        utils.delete_pod_template(template_name)
+    if registry_auth_id is not None:
+        utils.delete_register_auth(registry_auth_id)
 
 
 def get_cluster_info(
@@ -205,6 +223,44 @@ def query_instances(
 
 def cleanup_ports(
     cluster_name_on_cloud: str,
+    ports: List[str],
     provider_config: Optional[Dict[str, Any]] = None,
 ) -> None:
-    del cluster_name_on_cloud, provider_config
+    del cluster_name_on_cloud, ports, provider_config  # Unused.
+
+
+def query_ports(
+    cluster_name_on_cloud: str,
+    ports: List[str],
+    head_ip: Optional[str] = None,
+    provider_config: Optional[Dict[str, Any]] = None,
+) -> Dict[int, List[common.Endpoint]]:
+    """See sky/provision/__init__.py"""
+    del head_ip, provider_config  # Unused.
+    # RunPod ports sometimes take a while to be ready.
+    start_time = time.time()
+    ports_to_query = resources_utils.port_ranges_to_set(ports)
+    while True:
+        instances = _filter_instances(cluster_name_on_cloud,
+                                      None,
+                                      head_only=True)
+        assert len(instances) <= 1
+        # It is possible that the instance is terminated on console by
+        # the user. In this case, the instance will not be found and we
+        # should return an empty dict.
+        if not instances:
+            return {}
+        head_inst = list(instances.values())[0]
+        ready_ports: Dict[int, List[common.Endpoint]] = {
+            port: [common.SocketEndpoint(**endpoint)]
+            for port, endpoint in head_inst['port2endpoint'].items()
+            if port in ports_to_query
+        }
+        not_ready_ports = ports_to_query - set(ready_ports.keys())
+        if not not_ready_ports:
+            return ready_ports
+        if time.time() - start_time > QUERY_PORTS_TIMEOUT_SECONDS:
+            logger.warning(f'Querying ports {ports} timed out. Ports '
+                           f'{not_ready_ports} are not ready.')
+            return ready_ports
+        time.sleep(1)