skypilot-nightly 1.0.0.dev20250521__py3-none-any.whl → 1.0.0.dev20250523__py3-none-any.whl

sky/optimizer.py

@@ -21,6 +21,7 @@ from sky.usage import usage_lib
 from sky.utils import common
 from sky.utils import env_options
 from sky.utils import log_utils
+from sky.utils import registry
 from sky.utils import resources_utils
 from sky.utils import rich_utils
 from sky.utils import subprocess_utils
@@ -73,8 +74,8 @@ class Optimizer:
     def _egress_cost(src_cloud: clouds.Cloud, dst_cloud: clouds.Cloud,
                      gigabytes: float) -> float:
         """Returns estimated egress cost."""
-        if isinstance(src_cloud, DummyCloud) or isinstance(
-                dst_cloud, DummyCloud):
+        if isinstance(src_cloud, clouds.DummyCloud) or isinstance(
+                dst_cloud, clouds.DummyCloud):
             return 0.0
 
         if not src_cloud.is_same_cloud(dst_cloud):
@@ -88,8 +89,8 @@ class Optimizer:
                      gigabytes: float) -> float:
         """Returns estimated egress time in seconds."""
         # FIXME: estimate bandwidth between each cloud-region pair.
-        if isinstance(src_cloud, DummyCloud) or isinstance(
-                dst_cloud, DummyCloud):
+        if isinstance(src_cloud, clouds.DummyCloud) or isinstance(
+                dst_cloud, clouds.DummyCloud):
             return 0.0
         if not src_cloud.is_same_cloud(dst_cloud):
             # 10Gbps is close to the average of observed b/w from S3
@@ -167,7 +168,7 @@ class Optimizer:
 
         def make_dummy(name):
             dummy = task_lib.Task(name)
-            dummy.set_resources({DummyResources(cloud=DummyCloud())})
+            dummy.set_resources({DummyResources(cloud=clouds.DummyCloud())})
             dummy.set_time_estimator(lambda _: 0)
             return dummy
 
@@ -197,7 +198,7 @@ class Optimizer:
         node: task_lib.Task,
         resources: resources_lib.Resources,
     ) -> Tuple[Optional[clouds.Cloud], Optional[clouds.Cloud], Optional[float]]:
-        if isinstance(parent_resources.cloud, DummyCloud):
+        if isinstance(parent_resources.cloud, clouds.DummyCloud):
             # Special case.  The current 'node' is a real
             # source node, and its input may be on a different
             # cloud from 'resources'.
@@ -376,6 +377,10 @@ class Optimizer:
                     if any(orig_resources.cloud is None
                            for orig_resources in node.resources):
                         source_hint = 'catalog and kubernetes cluster'
+                    elif all(
+                            isinstance(orig_resources.cloud, clouds.SSH)
+                            for orig_resources in node.resources):
+                        source_hint = 'node pool'
                     elif all(
                             isinstance(orig_resources.cloud, clouds.Kubernetes)
                             for orig_resources in node.resources):
@@ -858,11 +863,19 @@ class Optimizer:
                 'accelerators': f'{resources.accelerators}',
                 'use_spot': resources.use_spot
             }
+
+            # Handle special case for Kubernetes and SSH clouds
             if isinstance(resources.cloud, clouds.Kubernetes):
-                # Region for Kubernetes is the context name, i.e. different
-                # Kubernetes clusters. We add region to the key to show all the
-                # Kubernetes clusters in the optimizer table for better UX.
+                # Region for Kubernetes-like clouds (SSH, Kubernetes) is the
+                # context name, i.e. different Kubernetes clusters. We add
+                # region to the key to show all the Kubernetes clusters in the
+                # optimizer table for better UX.
+
+                if resources.cloud.__class__.__name__ == 'SSH':
+                    resource_key_dict[
+                        'cloud'] = 'SSH'  # Force the cloud name to be SSH
                 resource_key_dict['region'] = resources.region
+
             return json.dumps(resource_key_dict, sort_keys=True)
 
         # Print the list of resouces that the optimizer considered.
@@ -1158,11 +1171,6 @@ class DummyResources(resources_lib.Resources):
         return 0
 
 
-class DummyCloud(clouds.Cloud):
-    """A dummy Cloud that has zero egress cost from/to."""
-    pass
-
-
 def _filter_out_blocked_launchable_resources(
         launchable_resources: Iterable[resources_lib.Resources],
         blocked_resources: Iterable[resources_lib.Resources]):
@@ -1221,7 +1229,13 @@ def _check_specified_clouds(dag: 'dag_lib.Dag') -> None:
         if disabled_clouds:
             is_or_are = 'is' if len(disabled_clouds) == 1 else 'are'
             task_name = f' {task.name!r}' if task.name is not None else ''
-            msg = (f'Task{task_name} requires {", ".join(disabled_clouds)} '
+            disabled_display_names = []
+            for c in disabled_clouds:
+                cloud_obj_one = registry.CLOUD_REGISTRY.from_str(c)
+                if cloud_obj_one is not None:
+                    disabled_display_names.append(cloud_obj_one.display_name())
+            cloud_names = ', '.join(disabled_display_names)
+            msg = (f'Task{task_name} requires {cloud_names} '
                    f'which {is_or_are} not enabled. To enable access, change '
                    f'the task cloud requirement or run: {colorama.Style.BRIGHT}'
                    f'sky check {" ".join(c.lower() for c in disabled_clouds)}'

sky/provision/__init__.py

@@ -23,6 +23,7 @@ from sky.provision import lambda_cloud
 from sky.provision import nebius
 from sky.provision import oci
 from sky.provision import runpod
+from sky.provision import ssh
 from sky.provision import vast
 from sky.provision import vsphere
 from sky.utils import command_runner

sky/provision/aws/instance.py

@@ -836,7 +836,23 @@ def open_ports(
 
     # For the case when every new ports is already opened.
     if ip_permissions:
-        sg.authorize_ingress(IpPermissions=ip_permissions)
+        # Filter out any permissions that already exist in the security group
+        existing_permissions = set()
+        for rule in sg.ip_permissions:
+            if rule['IpProtocol'] == 'tcp':
+                for ip_range in rule.get('IpRanges', []):
+                    if ip_range.get('CidrIp') == '0.0.0.0/0':
+                        existing_permissions.add(
+                            (rule['FromPort'], rule['ToPort']))
+
+        # Remove any permissions that already exist
+        filtered_permissions = []
+        for perm in ip_permissions:
+            if (perm['FromPort'], perm['ToPort']) not in existing_permissions:
+                filtered_permissions.append(perm)
+
+        if filtered_permissions:
+            sg.authorize_ingress(IpPermissions=filtered_permissions)
 
 
 def cleanup_ports(

sky/provision/gcp/constants.py

@@ -1,4 +1,5 @@
 """Constants used by the GCP provisioner."""
+import textwrap
 
 VERSION = 'v1'
 # Using v2 according to
@@ -53,9 +54,7 @@ CLUSTER_PREFIX_LENGTH = 10
 
 COMPACT_GROUP_PLACEMENT_POLICY = 'compact'
 COLLOCATED_COLLOCATION = 'COLLOCATED'
-GPU_DIRECT_TCPX_USER_DATA = """#!/bin/bash
-    set -e
-    set -x
+GPU_DIRECT_TCPX_USER_DATA = textwrap.dedent("""
     # Install GPU Direct TCPX
     cos-extensions install gpu -- --version=latest;
     sudo mount --bind /var/lib/nvidia /var/lib/nvidia;
@@ -83,7 +82,7 @@ GPU_DIRECT_TCPX_USER_DATA = """#!/bin/bash
     sudo mount --bind /var/lib/tcpx /var/lib/tcpx;
     sudo mount -o remount,exec /var/lib/tcpx;
     echo "GPU Direct TCPX installed"
-    """
+    """)
 
 GPU_DIRECT_TCPX_SPECIFIC_OPTIONS = [
     '--cap-add=IPC_LOCK',
@@ -106,6 +105,150 @@ GPU_DIRECT_TCPX_SPECIFIC_OPTIONS = [
     '--env LD_LIBRARY_PATH=/usr/local/nvidia/lib64:/usr/local/tcpx/lib64',
 ]
 
+PD_EXTREME_IOPS = 20000
+DEFAULT_DISK_SIZE = 100
+NETWORK_STORAGE_TYPE = 'PERSISTENT'
+INSTANCE_STORAGE_TYPE = 'SCRATCH'
+INSTANCE_STORAGE_DISK_TYPE = 'local-ssd'
+INSTANCE_STORAGE_INTERFACE_TYPE = 'NVME'
+INSTANCE_STORAGE_DEVICE_NAME_PREFIX = '/dev/disk/by-id/google-local-nvme-ssd-'
+DEVICE_NAME_PREFIX = '/dev/disk/by-id/google-'
+
+BASH_SCRIPT_START = textwrap.dedent("""#!/bin/bash
+set -e
+set -x
+""")
+DISK_MOUNT_USER_DATA_TEMPLATE = textwrap.dedent("""
+    # Define arrays for devices and mount points
+    declare -A device_mounts=(
+        {device_mounts}
+    )
+
+    # Function to format and mount a single device
+    format_and_mount() {{
+        local device_name="$1"
+        local mount_point="$2"
+
+        if [ ! -e "$device_name" ]; then
+            echo "Error: Device $device_name does not exist."
+            return 1
+        fi
+
+        # Check if filesystem is already formatted (ext4)
+        if ! sudo blkid "$device_name" | grep -q 'TYPE="ext4"'; then
+            if [[ "$device_name" == "/dev/disk/by-id/google-local-nvme-ssd"* ]]; then
+                echo "Formatting local SSD $device_name..."
+                if ! sudo mkfs.ext4 -F "$device_name"; then
+                    echo "Error: Failed to format $device_name"
+                    return 1
+                fi
+            else
+                echo "Formatting persistent disk $device_name..."
+                if ! sudo mkfs.ext4 -m 0 -E lazy_itable_init=0,lazy_journal_init=0,discard "$device_name"; then
+                    echo "Error: Failed to format $device_name"
+                    return 1
+                fi
+            fi
+        else
+            echo "$device_name is already formatted."
+        fi
+
+        # Check if already mounted
+        if ! grep -q "$mount_point" /proc/mounts; then
+            echo "Mounting $device_name to $mount_point..."
+            if ! sudo mkdir -p "$mount_point"; then
+                echo "Error: Failed to create mount point $mount_point"
+                return 1
+            fi
+
+            if ! sudo mount "$device_name" "$mount_point"; then
+                echo "Error: Failed to mount $device_name to $mount_point"
+                return 1
+            fi
+
+            # Add to fstab if not already present
+            if ! grep -q " $mount_point " /etc/fstab; then
+                echo "Adding mount entry to /etc/fstab..."
+                echo "UUID=`sudo blkid -s UUID -o value $device_name` $mount_point ext4 defaults,nofail 0 2" | sudo tee -a /etc/fstab
+            else
+                echo "Mount entry already exists in /etc/fstab"
+            fi
+        else
+            echo "$device_name is already mounted at $mount_point"
+        fi
+    }}
+
+    # Main execution
+    echo "Starting device mounting process..."
+
+    # Process each device-mount pair
+    for device in "${{!device_mounts[@]}}"; do
+        mount_point="${{device_mounts[$device]}}"
+        echo "Processing device: $device -> $mount_point"
+        if ! format_and_mount "$device" "$mount_point"; then
+            echo "Failed to process device $device"
+            # Continue with other devices even if one fails
+            continue
+        fi
+    done
+
+    echo "Device mounting process completed."
+""")
+
+# The local SSDs will be attached automatically to the following
+# machine types with the following number of disks.
+# Refer to https://cloud.google.com/compute/docs/disks/local-ssd#lssd_disks_fixed
+SSD_AUTO_ATTACH_MACHINE_TYPES = {
+    'c4a-standard-4-lssd': 1,
+    'c4a-highmem-4-lssd': 1,
+    'c4a-standard-8-lssd': 2,
+    'c4a-highmem-8-lssd': 2,
+    'c4a-standard-16-lssd': 4,
+    'c4a-highmem-16-lssd': 4,
+    'c4a-standard-32-lssd': 6,
+    'c4a-highmem-32-lssd': 6,
+    'c4a-standard-48-lssd': 10,
+    'c4a-highmem-48-lssd': 10,
+    'c4a-standard-64-lssd': 14,
+    'c4a-highmem-64-lssd': 14,
+    'c4a-standard-72-lssd': 16,
+    'c4a-highmem-72-lssd': 16,
+    'c3-standard-4-lssd': 1,
+    'c3-standard-8-lssd': 2,
+    'c3-standard-22-lssd': 4,
+    'c3-standard-44-lssd': 8,
+    'c3-standard-88-lssd': 16,
+    'c3-standard-176-lssd': 32,
+    'c3d-standard-8-lssd': 1,
+    'c3d-highmem-8-lssd': 1,
+    'c3d-standard-16-lssd': 1,
+    'c3d-highmem-16-lssd': 1,
+    'c3d-standard-30-lssd': 2,
+    'c3d-highmem-30-lssd': 2,
+    'c3d-standard-60-lssd': 4,
+    'c3d-highmem-60-lssd': 4,
+    'c3d-standard-90-lssd': 8,
+    'c3d-highmem-90-lssd': 8,
+    'c3d-standard-180-lssd': 16,
+    'c3d-highmem-180-lssd': 16,
+    'c3d-standard-360-lssd': 32,
+    'c3d-highmem-360-lssd': 32,
+    'a4-highgpu-8g': 32,
+    'a3-ultragpu-8g': 32,
+    'a3-megagpu-8g': 16,
+    'a3-highgpu-1g': 2,
+    'a3-highgpu-2g': 4,
+    'a3-highgpu-4g': 8,
+    'a3-highgpu-8g': 16,
+    'a3-edgegpu-8g': 16,
+    'a2-ultragpu-1g': 1,
+    'a2-ultragpu-2g': 2,
+    'a2-ultragpu-4g': 4,
+    'a2-ultragpu-8g': 8,
+    'z3-highmem-88': 12,
+    'z3-highmem-176': 12,
+}
+
 # Below parameters are from the default VPC on GCP.
 # https://cloud.google.com/vpc/docs/firewalls#more_rules_default_vpc
 VPC_TEMPLATE: dict = {

sky/provision/gcp/instance_utils.py

@@ -826,6 +826,16 @@ class GCPComputeInstance(GCPInstance):
         # https://cloud.google.com/compute/docs/reference/rest/v1/instances/bulkInsert # pylint: disable=line-too-long
         if config.get('sourceMachineImage') is not None:
             return False
+        # bulkInsert does not support attaching existing
+        # disks to the instances with READ_WRITE mode.
+        if config.get('disks') is not None:
+            for disk in config['disks']:
+                if disk.get('source') is not None and disk.get(
+                        'mode', 'READ_WRITE') == 'READ_WRITE':
+                    return False
+                if disk.get('initializeParams') is not None and disk.get(
+                        'initializeParams', {}).get('diskName') is not None:
+                    return False
         return True
 
     @classmethod

sky/provision/gcp/volume_utils.py

@@ -0,0 +1,247 @@
+"""Utilities for GCP volumes."""
+from typing import Any, Dict, List, Optional
+
+from sky import clouds
+from sky import exceptions
+from sky import sky_logging
+from sky.adaptors import gcp
+from sky.provision.gcp import constants
+from sky.utils import resources_utils
+from sky.utils import ux_utils
+
+logger = sky_logging.init_logger(__name__)
+
+
+def get_data_disk_tier_mapping(
+    instance_type: Optional[str],) -> Dict[resources_utils.DiskTier, str]:
+    # Define the default mapping from disk tiers to disk types.
+    # Refer to https://cloud.google.com/compute/docs/disks/hyperdisks
+    # and https://cloud.google.com/compute/docs/disks/persistent-disks
+    tier2name = {
+        resources_utils.DiskTier.ULTRA: 'pd-extreme',
+        resources_utils.DiskTier.HIGH: 'pd-ssd',
+        resources_utils.DiskTier.MEDIUM: 'pd-balanced',
+        resources_utils.DiskTier.LOW: 'pd-standard',
+    }
+
+    if instance_type is None:
+        return tier2name
+
+    # Remap series-specific disk types.
+    series = instance_type.split('-')[0]
+
+    if series in ['a4', 'x4']:
+        tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-extreme'
+        tier2name[resources_utils.DiskTier.HIGH] = 'hyperdisk-balanced'
+        tier2name[resources_utils.DiskTier.MEDIUM] = 'hyperdisk-balanced'
+        tier2name[resources_utils.DiskTier.LOW] = 'hyperdisk-balanced'
+    elif series in ['m4']:
+        tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-extreme'
+        tier2name[resources_utils.DiskTier.HIGH] = 'hyperdisk-balanced'
+        tier2name[resources_utils.DiskTier.MEDIUM] = 'hyperdisk-balanced'
+        tier2name[resources_utils.DiskTier.LOW] = 'hyperdisk-balanced'
+        num_cpus = int(instance_type.split('-')[2])  # type: ignore
+        if num_cpus < 112:
+            tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-balanced'
+    elif series in ['c4', 'c4a', 'c4d']:
+        tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-extreme'
+        tier2name[resources_utils.DiskTier.HIGH] = 'hyperdisk-balanced'
+        tier2name[resources_utils.DiskTier.MEDIUM] = 'hyperdisk-balanced'
+        tier2name[resources_utils.DiskTier.LOW] = 'hyperdisk-balanced'
+        num_cpus = int(instance_type.split('-')[2])  # type: ignore
+        if num_cpus < 64:
+            tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-balanced'
+    elif series in ['a3']:
+        if (instance_type.startswith('a3-ultragpu') or
+                instance_type.startswith('a3-megagpu') or
+                instance_type.startswith('a3-edgegpu')):
+            tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-extreme'
+            tier2name[resources_utils.DiskTier.HIGH] = 'hyperdisk-balanced'
+            tier2name[resources_utils.DiskTier.MEDIUM] = 'hyperdisk-balanced'
+            tier2name[resources_utils.DiskTier.LOW] = 'hyperdisk-balanced'
+        elif instance_type.startswith('a3-highgpu'):
+            tier2name[resources_utils.DiskTier.LOW] = 'pd-balanced'
+            if instance_type.startswith('a3-highgpu-8g'):
+                tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-extreme'
+                tier2name[resources_utils.DiskTier.HIGH] = 'hyperdisk-balanced'
+                tier2name[resources_utils.DiskTier.MEDIUM] = 'pd-ssd'
+            elif instance_type.startswith('a3-highgpu-4g'):
+                tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-extreme'
+            else:
+                tier2name[resources_utils.DiskTier.ULTRA] = 'pd-ssd'
+    elif series in ['c3d']:
+        tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-extreme'
+        tier2name[resources_utils.DiskTier.HIGH] = 'hyperdisk-balanced'
+        tier2name[resources_utils.DiskTier.MEDIUM] = 'pd-ssd'
+        tier2name[resources_utils.DiskTier.LOW] = 'pd-balanced'
+        num_cpus = int(instance_type.split('-')[2])  # type: ignore
+        if num_cpus < 60:
+            tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-balanced'
+    elif series in ['c3']:
+        tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-extreme'
+        tier2name[resources_utils.DiskTier.HIGH] = 'hyperdisk-balanced'
+        tier2name[resources_utils.DiskTier.MEDIUM] = 'pd-ssd'
+        tier2name[resources_utils.DiskTier.LOW] = 'pd-balanced'
+        num_cpus = int(instance_type.split('-')[2])  # type: ignore
+        if num_cpus < 88:
+            tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-balanced'
+    elif series in ['n4']:
+        tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-balanced'
+        tier2name[resources_utils.DiskTier.HIGH] = 'hyperdisk-balanced'
+        tier2name[resources_utils.DiskTier.MEDIUM] = 'hyperdisk-balanced'
+        tier2name[resources_utils.DiskTier.LOW] = 'hyperdisk-balanced'
+    elif series in ['n2d', 'n1', 't2d', 't2a', 'e2', 'c2', 'c2d', 'a2']:
+        tier2name[resources_utils.DiskTier.ULTRA] = 'pd-ssd'
+    elif series in ['z3']:
+        tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-extreme'
+        tier2name[resources_utils.DiskTier.LOW] = 'pd-balanced'
+    elif series in ['h3']:
+        tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-balanced'
+        tier2name[resources_utils.DiskTier.HIGH] = 'hyperdisk-balanced'
+        tier2name[resources_utils.DiskTier.LOW] = 'pd-balanced'
+    elif series in ['m3']:
+        tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-extreme'
+        tier2name[resources_utils.DiskTier.HIGH] = 'hyperdisk-balanced'
+        tier2name[resources_utils.DiskTier.MEDIUM] = 'pd-ssd'
+        tier2name[resources_utils.DiskTier.LOW] = 'pd-balanced'
+        num_cpus = int(instance_type.split('-')[2])  # type: ignore
+        if num_cpus < 64:
+            tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-balanced'
+    elif series in ['m2']:
+        tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-extreme'
+        tier2name[resources_utils.DiskTier.HIGH] = 'hyperdisk-balanced'
+    elif series in ['m1']:
+        tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-extreme'
+        tier2name[resources_utils.DiskTier.HIGH] = 'hyperdisk-balanced'
+        num_cpus = int(instance_type.split('-')[2])  # type: ignore
+        if num_cpus < 80:
+            tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-balanced'
+    elif series in ['g2']:
+        tier2name[resources_utils.DiskTier.ULTRA] = 'pd-ssd'
+        tier2name[resources_utils.DiskTier.LOW] = 'pd-balanced'
+    elif series in ['n2']:
+        num_cpus = int(instance_type.split('-')[2])  # type: ignore
+        if num_cpus < 64:
+            tier2name[resources_utils.DiskTier.ULTRA] = 'pd-ssd'
+        elif num_cpus >= 80:
+            tier2name[resources_utils.DiskTier.ULTRA] = 'hyperdisk-extreme'
+
+    return tier2name
+
+
+def validate_instance_volumes(
+    instance_type: Optional[str],
+    volumes: Optional[List[Dict[str, Any]]],
+) -> None:
+    if not volumes:
+        return
+    if instance_type is None:
+        logger.warning('Instance type is not specified,'
+                       ' skipping instance volume validation')
+        return
+    instance_volume_count = 0
+    for volume in volumes:
+        if volume['storage_type'] == resources_utils.StorageType.INSTANCE:
+            instance_volume_count += 1
+    if (instance_type in constants.SSD_AUTO_ATTACH_MACHINE_TYPES and
+            instance_volume_count >
+            constants.SSD_AUTO_ATTACH_MACHINE_TYPES[instance_type]):
+        raise exceptions.ResourcesUnavailableError(
+            f'The instance type {instance_type} supports'
+            f' {constants.SSD_AUTO_ATTACH_MACHINE_TYPES[instance_type]}'
+            f'  instance storage, but {instance_volume_count} are specified')
+    # TODO(hailong):
+    # check the instance storage count for the other instance types,
+    # refer to https://cloud.google.com/compute/docs/disks/local-ssd
+
+
+def translate_attach_mode(attach_mode: resources_utils.DiskAttachMode) -> str:
+    if attach_mode == resources_utils.DiskAttachMode.READ_ONLY:
+        return 'READ_ONLY'
+    return 'READ_WRITE'
+
+
+def check_volume_name_exist_in_region(
+        project_id: str, region: clouds.Region, use_mig: bool,
+        volume_name: str) -> Optional[Dict[str, Any]]:
+    """Check if the volume name exists and return the volume info."""
+    logger.debug(f'Checking volume {volume_name} in region {region}')
+    try:
+        compute = gcp.build('compute',
+                            'v1',
+                            credentials=None,
+                            cache_discovery=False)
+    except gcp.credential_error_exception():
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError('Not able to build compute client') from None
+
+    # Get all the zones in the region
+    all_zones = compute.zones().list(project=project_id).execute()
+    region_zones = []
+    if 'items' in all_zones:
+        for zone in all_zones['items']:
+            if zone['region'].split('/')[-1] == region.name:
+                region_zones.append(zone['name'])
+    volume_info = None
+    for zone in region_zones:
+        try:
+            volume_info = compute.disks().get(project=project_id,
+                                              zone=zone,
+                                              disk=volume_name).execute()
+            if volume_info is not None:
+                if use_mig:
+                    # With MIG, instance template will be used, in this case,
+                    # the `selfLink` for zonal disk needs to be the volume name
+                    # Refer to https://cloud.google.com/compute/docs/
+                    # reference/rest/v1/instances/insert
+                    volume_info['selfLink'] = volume_name
+                volume_info['available_zones'] = [zone]
+                return volume_info
+        except gcp.http_error_exception() as e:
+            if e.resp.status == 403:
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError('Not able to access the volume '
+                                     f'{volume_name!r}') from None
+            if e.resp.status == 404:
+                continue  # Try next zone
+            raise
+
+    # If not found in any zone, check region disk
+    try:
+        volume_info = compute.regionDisks().get(project=project_id,
+                                                region=region.name,
+                                                disk=volume_name).execute()
+        # 'replicaZones':
+        #  ['https://xxx/compute/v1/projects/sky-dev-465/zones/us-central1-a',
+        # 'https://xxx/compute/v1/projects/sky-dev-465/zones/us-central1-c']
+        if volume_info is not None and 'replicaZones' in volume_info:
+            replica_zones = [
+                zone.split('/')[-1] for zone in volume_info['replicaZones']
+            ]
+            volume_info['available_zones'] = replica_zones
+        return volume_info
+    except gcp.http_error_exception() as e:
+        if e.resp.status == 403:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError('Not able to access the volume '
+                                 f'{volume_name!r}') from None
+        if e.resp.status == 404:
+            logger.warning(
+                f'Volume {volume_name} is not found in region {region}.'
+                f' It will be created.')
+            return volume_info
+        raise
+
+
+def check_volume_zone_match(volume_name: str,
+                            zones: Optional[List[clouds.Zone]],
+                            available_zones: List[str]):
+    if zones is None:
+        return None
+    for zone in zones:
+        if zone.name in available_zones:
+            return None
+    with ux_utils.print_exception_no_traceback():
+        # Return a ResourcesUnavailableError to trigger failover
+        raise exceptions.ResourcesUnavailableError(
+            f'Volume {volume_name} not available in zones {zones}') from None

sky/provision/kubernetes/instance.py

@@ -1265,6 +1265,8 @@ def query_instances(
     assert provider_config is not None
     namespace = kubernetes_utils.get_namespace_from_config(provider_config)
     context = kubernetes_utils.get_context_from_config(provider_config)
+    is_ssh = context.startswith('ssh-') if context else False
+    identity = 'SSH Node Pool' if is_ssh else 'Kubernetes cluster'
 
     # Get all the pods with the label skypilot-cluster: <cluster_name>
     try:
@@ -1274,15 +1276,24 @@ def query_instances(
             _request_timeout=kubernetes.API_TIMEOUT).items
     except kubernetes.max_retry_error():
         with ux_utils.print_exception_no_traceback():
-            ctx = kubernetes_utils.get_current_kube_config_context_name()
+            if is_ssh:
+                node_pool = context.lstrip('ssh-') if context else ''
+                msg = (
+                    f'Cannot connect to SSH Node Pool {node_pool}. '
+                    'Please check if the SSH Node Pool is up and accessible. '
+                    'To debug, run `sky check ssh` to check the status of '
+                    'the SSH Node Pool.')
+            else:
+                ctx = kubernetes_utils.get_current_kube_config_context_name()
+                msg = (f'Network error - check if the {identity} in '
+                       f'context {ctx} is up and accessible.')
             raise exceptions.ClusterStatusFetchingError(
-                f'Failed to query cluster {cluster_name_on_cloud!r} status. '
-                'Network error - check if the Kubernetes cluster in '
-                f'context {ctx} is up and accessible.') from None
+                f'Failed to query cluster {cluster_name_on_cloud!r} status. ' +
+                msg) from None
     except Exception as e:  # pylint: disable=broad-except
         with ux_utils.print_exception_no_traceback():
             raise exceptions.ClusterStatusFetchingError(
-                f'Failed to query Kubernetes cluster {cluster_name_on_cloud!r} '
+                f'Failed to query {identity} {cluster_name_on_cloud!r} '
                 f'status: {common_utils.format_exception(e)}')
 
     # Check if the pods are running or pending