skypilot-nightly 1.0.0.dev20250502__py3-none-any.whl → 1.0.0.dev20251203__py3-none-any.whl

sky/provision/gcp/config.py

@@ -5,11 +5,14 @@ import time
 import typing
 from typing import Any, Dict, List, Set, Tuple
 
+from typing_extensions import TypedDict
+
 from sky.adaptors import gcp
 from sky.clouds.utils import gcp_utils
 from sky.provision import common
 from sky.provision.gcp import constants
 from sky.provision.gcp import instance_utils
+from sky.utils import resources_utils
 
 logger = logging.getLogger(__name__)
 
@@ -75,6 +78,30 @@ def wait_for_compute_global_operation(project_name, operation, compute):
     return result
 
 
+def wait_for_compute_region_operation(project_name, region, operation, compute):
+    """Poll for region compute operation until finished."""
+    logger.info('wait_for_compute_region_operation: '
+                'Waiting for operation {} to finish...'.format(
+                    operation['name']))
+
+    for _ in range(constants.MAX_POLLS):
+        result = (compute.regionOperations().get(
+            project=project_name,
+            region=region,
+            operation=operation['name'],
+        ).execute())
+        if 'error' in result:
+            raise Exception(result['error'])
+
+        if result['status'] == 'DONE':
+            logger.info('wait_for_compute_region_operation: Operation done.')
+            break
+
+        time.sleep(constants.POLL_INTERVAL)
+
+    return result
+
+
 def _create_crm(gcp_credentials=None):
     return gcp.build('cloudresourcemanager',
                      'v1',
@@ -168,6 +195,7 @@ def bootstrap_instances(
     iam_role = _configure_iam_role(config, crm, iam)
     config.node_config.update(iam_role)
     config = _configure_subnet(region, cluster_name, config, compute)
+    config = _configure_placement_policy(region, cluster_name, config, compute)
 
     return config
 
@@ -248,7 +276,7 @@ def _is_permission_satisfied(service_account, crm, iam, required_permissions,
     # For example, `roles/iam.serviceAccountUser` can be granted at the
     # skypilot-v1 service account level, which can be checked with
     # service_account_policy = iam.projects().serviceAccounts().getIamPolicy(
-    #    resource=f'projects/{project_id}/serviceAcccounts/{email}').execute()
+    #    resource=f'projects/{project_id}/serviceAccounts/{email}').execute()
     # We now skip the check for `iam.serviceAccounts.actAs` permission for
     # simplicity as it can be granted at the service account level.
     def check_permissions(policy, required_permissions):
@@ -389,6 +417,9 @@ def _configure_iam_role(config: common.ProvisionConfig, crm, iam) -> dict:
     return iam_role
 
 
+AllowedList = TypedDict('AllowedList', {'IPProtocol': str, 'ports': List[str]})
+
+
 def _check_firewall_rules(cluster_name: str, vpc_name: str, project_id: str,
                           compute):
     """Check if the firewall rules in the VPC are sufficient."""
@@ -440,7 +471,7 @@ def _check_firewall_rules(cluster_name: str, vpc_name: str, project_id: str,
             }
         """
         source2rules: Dict[Tuple[str, str], Dict[str, Set[int]]] = {}
-        source2allowed_list: Dict[Tuple[str, str], List[Dict[str, str]]] = {}
+        source2allowed_list: Dict[Tuple[str, str], List[AllowedList]] = {}
         for rule in rules:
             # Rules applied to specific VM (targetTags) may not work for the
             # current VM, so should be skipped.
@@ -506,7 +537,23 @@ def _check_firewall_rules(cluster_name: str, vpc_name: str, project_id: str,
     return True
 
 
-def _create_rules(project_id: str, compute, rules, vpc_name):
+def _delete_rules(project_id: str, compute, rules, vpc_name: str):
+    for rule_ori in rules:
+        # Query firewall rule by its name (unique in a project).
+        rule_name = rule_ori['name'].format(VPC_NAME=vpc_name)
+        rule_list = _list_firewall_rules(project_id,
+                                         compute,
+                                         filter=f'(name={rule_name})')
+        for rule in rule_list:
+            logger.info(f'Deleting firewall rule {rule["name"]}')
+            _delete_firewall_rule(project_id, compute, rule['name'])
+
+
+def _create_rules(project_id: str,
+                  compute,
+                  rules,
+                  vpc_name,
+                  recreate: bool = True):
     opertaions = []
     for rule in rules:
         # Query firewall rule by its name (unique in a project).
@@ -516,7 +563,11 @@ def _create_rules(project_id: str, compute, rules, vpc_name):
                                          compute,
                                          filter=f'(name={rule_name})')
         if rule_list:
-            _delete_firewall_rule(project_id, compute, rule_name)
+            if recreate:
+                _delete_firewall_rule(project_id, compute, rule_name)
+            else:
+                logger.info(f'Rule {rule_name} already exists')
+                continue
 
         body = rule.copy()
         body['name'] = body['name'].format(VPC_NAME=vpc_name)
@@ -660,6 +711,149 @@ def get_usable_vpc_and_subnet(
     return usable_vpc_name, usable_subnet
 
 
+def get_gpu_direct_usable_vpcs_and_subnets(
+    cluster_name: str,
+    region: str,
+    config: common.ProvisionConfig,
+    compute,
+) -> List[Tuple[str, 'google.cloud.compute_v1.types.compute.Subnetwork']]:
+    """Return a list of usable VPCs and subnets for GPU Direct."""
+    project_id = config.provider_config['project_id']
+    vpc_subnet_pairs = []
+
+    # TODO(hailong): Determine the num_vpcs per different GPU Direct types
+    num_vpcs = constants.SKYPILOT_GPU_DIRECT_VPC_NUM
+
+    cidr_prefix = constants.SKYPILOT_GPU_DIRECT_VPC_CIDR_PREFIX
+    for i in range(num_vpcs):
+        vpc_name = get_gpu_direct_vpc_name(cluster_name, i)
+        subnet_name = f'{vpc_name}-sub'
+        subnet_cidr_range = f'{cidr_prefix}.{i}.0/24'
+        # Check if VPC exists
+        vpc_list = _list_vpcnets(project_id, compute, filter=f'name={vpc_name}')
+        if not vpc_list:
+            body = constants.VPC_TEMPLATE.copy()
+            body['mtu'] = 8244
+            body['autoCreateSubnetworks'] = False
+            body['name'] = vpc_name
+            body['selfLink'] = body['selfLink'].format(PROJ_ID=project_id,
+                                                       VPC_NAME=vpc_name)
+            _create_vpcnet(project_id, compute, body)
+        # Check if subnet exists
+        subnets = _list_subnets(project_id, region, compute, network=vpc_name)
+        if not subnets:
+            _create_subnet(project_id, region, compute, vpc_name, subnet_name,
+                           subnet_cidr_range)
+            subnets = _list_subnets(project_id,
+                                    region,
+                                    compute,
+                                    network=vpc_name)
+        # Apply firewall rules
+        # No need to recreate the rules if exist,
+        # as they are totally managed by SkyPilot,
+        # in this case, we can skip the rules creation during failover
+        _create_rules(project_id,
+                      compute,
+                      constants.FIREWALL_RULES_TEMPLATE,
+                      vpc_name,
+                      recreate=False)
+        vpc_subnet_pairs.append((vpc_name, subnets[0]))
+    return vpc_subnet_pairs
+
+
+def get_gpu_direct_vpc_name(cluster_name: str, i: int) -> str:
+    """Get the name of the GPU Direct VPC."""
+    if i == 0:
+        return f'{cluster_name}-mgmt-net'
+    else:
+        return f'{cluster_name}-data-net-{i}'
+
+
+def delete_gpu_direct_vpcs_and_subnets(
+    cluster_name: str,
+    project_id: str,
+    region: str,
+    keep_global_resources: bool = False,
+):
+    """Delete GPU Direct subnets, firewalls, and VPCs.
+
+    Args:
+        cluster_name: The name of the cluster.
+        project_id: The ID of the project.
+        region: The region of the cluster.
+        keep_global_resources: Whether to keep the global resources. If True,
+            only delete the subnets. Otherwise, delete all the firewalls,
+            subnets, and VPCs.
+    """
+    compute = _create_compute()
+
+    # TODO(hailong): Determine the num_vpcs per different GPU Direct types
+    num_vpcs = constants.SKYPILOT_GPU_DIRECT_VPC_NUM
+
+    for i in range(num_vpcs):
+        vpc_name = get_gpu_direct_vpc_name(cluster_name, i)
+        # Check if VPC exists
+        vpc_list = _list_vpcnets(project_id, compute, filter=f'name={vpc_name}')
+        if not vpc_list:
+            continue
+        for vpc in vpc_list:
+            subnets = _list_subnets(project_id,
+                                    region,
+                                    compute,
+                                    network=vpc['name'])
+            for subnet in subnets:
+                logger.info(f'Deleting subnet {subnet["name"]}')
+                _delete_subnet(project_id, region, compute, subnet['name'])
+
+            if not keep_global_resources:
+                # For failover, keep_global_resources would be true,
+                # we don't delete the rules and VPCs,
+                # which are global resources and can be reused.
+                _delete_rules(project_id, compute,
+                              constants.FIREWALL_RULES_TEMPLATE, vpc['name'])
+                logger.info(f'Deleting VPC {vpc["name"]}')
+                _delete_vpcnet(project_id, compute, vpc['name'])
+
+
+def _configure_placement_policy(region: str, cluster_name: str,
+                                config: common.ProvisionConfig, compute):
+    """Configure placement group for GPU Direct."""
+    node_config = config.node_config
+    project_id = config.provider_config['project_id']
+    group_placement_policy = config.provider_config.get('placement_policy',
+                                                        None)
+    # If the placement policy is not compact,
+    # or the managed instance group is specified,
+    # skip the placement policy creation.
+    # If placement policy is specified together with managed instance group,
+    # it will cause the following error:
+    # Reason: [{'code': 'UNSUPPORTED_OPERATION',
+    # 'message': 'Creating queued resource with
+    # resource policies is not supported.'}]
+    mig_configuration = config.provider_config.get('use_managed_instance_group',
+                                                   False)
+    if (group_placement_policy is None or group_placement_policy.lower() !=
+            constants.COMPACT_GROUP_PLACEMENT_POLICY or mig_configuration):
+        return config
+
+    policy_name = f'{cluster_name}-placement-policy'
+    resource_policy = {
+        'name': policy_name,
+        'groupPlacementPolicy': {
+            'collocation': constants.COLLOCATED_COLLOCATION,
+        }
+    }
+    # Try to get the placement policy first, if not found, create it
+    placement_policy = _get_placement_policy(project_id, region, compute,
+                                             policy_name)
+    if not placement_policy:
+        logger.info(f'Creating placement policy {policy_name}'
+                    f' for cluster {cluster_name}')
+        _create_placement_policy(project_id, region, compute, resource_policy)
+    node_config['resourcePolicies'] = [policy_name]
+    return config
+
+
 def _configure_subnet(region: str, cluster_name: str,
                       config: common.ProvisionConfig, compute):
     """Pick a reasonable subnet if not specified by the config."""
@@ -671,25 +865,56 @@ def _configure_subnet(region: str, cluster_name: str,
     if 'networkInterfaces' in node_config or 'networkConfig' in node_config:
         return config
 
-    # SkyPilot: make sure there's a usable VPC
-    _, default_subnet = get_usable_vpc_and_subnet(cluster_name, region, config,
-                                                  compute)
-
-    default_interfaces = [{
-        'subnetwork': default_subnet['selfLink'],
-        'accessConfigs': [{
-            'name': 'External NAT',
-            'type': 'ONE_TO_ONE_NAT',
-        }]
-    }]
-    # Add gVNIC if specified in config
+    default_interfaces = []
+    enable_gpu_direct = config.provider_config.get('enable_gpu_direct', False)
     enable_gvnic = config.provider_config.get('enable_gvnic', False)
-    if enable_gvnic:
-        default_interfaces[0]['nicType'] = 'gVNIC'
+    network_tier = config.provider_config.get('network_tier', 'standard')
+    if (enable_gpu_direct or
+            network_tier == resources_utils.NetworkTier.BEST.value):
+        if not enable_gvnic:
+            logger.warning(
+                'Enable GPU Direct requires gvnic to be enabled, enabling gvnic'
+            )
+            config.provider_config['enable_gvnic'] = True
+            enable_gvnic = True
+        if 'machineType' not in node_config or node_config[
+                'machineType'] not in constants.GPU_DIRECT_TCPX_INSTANCE_TYPES:
+            raise ValueError(
+                'Enable GPU Direct requires machineType to be one of '
+                f'{constants.GPU_DIRECT_TCPX_INSTANCE_TYPES}')
+        logger.info(f'Enable GPU Direct for cluster {cluster_name} '
+                    f'with machineType {node_config["machineType"]}')
+        vpc_subnet_pairs = get_gpu_direct_usable_vpcs_and_subnets(
+            cluster_name, region, config, compute)
+        for _, subnet in vpc_subnet_pairs:
+            default_interfaces.append({
+                'subnetwork': subnet['selfLink'],
+                'accessConfigs': [{
+                    'name': 'External NAT',
+                    'type': 'ONE_TO_ONE_NAT',
+                }],
+                'nicType': 'gVNIC'
+            })
+    else:
+        # SkyPilot: make sure there's a usable VPC
+        _, default_subnet = get_usable_vpc_and_subnet(cluster_name, region,
+                                                      config, compute)
+
+        default_interfaces = [{
+            'subnetwork': default_subnet['selfLink'],
+            'accessConfigs': [{
+                'name': 'External NAT',
+                'type': 'ONE_TO_ONE_NAT',
+            }]
+        }]
+        # Add gVNIC if specified in config
+        if enable_gvnic:
+            default_interfaces[0]['nicType'] = 'gVNIC'
     enable_external_ips = _enable_external_ips(config)
     if not enable_external_ips:
         # Removing this key means the VM will not be assigned an external IP.
-        default_interfaces[0].pop('accessConfigs')
+        for interface in default_interfaces:
+            interface.pop('accessConfigs')
 
     # The not applicable key will be removed during node creation
 
@@ -747,6 +972,14 @@ def _list_vpcnets(project_id: str, compute, filter=None):  # pylint: disable=red
             if 'items' in response else [])
 
 
+def _delete_vpcnet(project_id: str, compute, vpcnet_name: str):
+    operation = compute.networks().delete(
+        project=project_id,
+        network=vpcnet_name,
+    ).execute()
+    return wait_for_compute_global_operation(project_id, operation, compute)
+
+
 def _list_subnets(
         project_id: str,
         region: str,
@@ -840,3 +1073,52 @@ def _add_iam_policy_binding(service_account, policy, crm, iam):
     ).execute())
 
     return result
+
+
+def _create_subnet(project_id: str, region: str, compute, vpc_name: str,
+                   subnet_name: str, ip_cidr_range: str):
+    body = {
+        'name': subnet_name,
+        'ipCidrRange': ip_cidr_range,
+        'network': f'projects/{project_id}/global/networks/{vpc_name}',
+        'region': region,
+    }
+    operation = compute.subnetworks().insert(project=project_id,
+                                             region=region,
+                                             body=body).execute()
+    response = wait_for_compute_region_operation(project_id, region, operation,
+                                                 compute)
+    return response
+
+
+def _delete_subnet(project_id: str, region: str, compute, subnet_name: str):
+    operation = compute.subnetworks().delete(
+        project=project_id,
+        region=region,
+        subnetwork=subnet_name,
+    ).execute()
+    return wait_for_compute_region_operation(project_id, region, operation,
+                                             compute)
+
+
+def _create_placement_policy(project_id: str, region: str, compute,
+                             placement_policy: dict):
+    operation = compute.resourcePolicies().insert(
+        project=project_id, region=region, body=placement_policy).execute()
+    response = wait_for_compute_region_operation(project_id, region, operation,
+                                                 compute)
+    return response
+
+
+def _get_placement_policy(project_id: str, region: str, compute, name: str):
+    try:
+        placement_policy = (compute.resourcePolicies().get(
+            project=project_id,
+            region=region,
+            resourcePolicy=name,
+        ).execute())
+    except gcp.http_error_exception() as e:
+        if e.resp.status == 404:
+            return None
+        raise
+    return placement_policy

sky/provision/gcp/constants.py

@@ -1,4 +1,5 @@
 """Constants used by the GCP provisioner."""
+import textwrap
 
 VERSION = 'v1'
 # Using v2 according to
@@ -41,6 +42,223 @@ HAS_TPU_PROVIDER_FIELD = '_has_tpus'
 # with ServiceAccounts.
 
 SKYPILOT_VPC_NAME = 'skypilot-vpc'
+SKYPILOT_GPU_DIRECT_VPC_NUM = 5
+SKYPILOT_GPU_DIRECT_VPC_CIDR_PREFIX = '10.129'
+GPU_DIRECT_TCPX_INSTANCE_TYPES = [
+    'a3-edgegpu-8g',
+    'a3-highgpu-8g',
+]
+
+COMPACT_GROUP_PLACEMENT_POLICY = 'compact'
+COLLOCATED_COLLOCATION = 'COLLOCATED'
+
+# From https://cloud.google.com/compute/docs/gpus/gpudirect
+# A specific image is used to ensure that the the GPU is configured with TCPX support.
+GCP_GPU_DIRECT_IMAGE_ID = 'docker:us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/nccl-plugin-gpudirecttcpx'
+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;
+    sudo mount -o remount,exec /var/lib/nvidia;
+    docker ps -a | grep -q receive-datapath-manager || \
+    docker run \
+    --detach \
+    --pull=always \
+    --name receive-datapath-manager \
+    --privileged \
+    --cap-add=NET_ADMIN --network=host \
+    --volume /var/lib/nvidia/lib64:/usr/local/nvidia/lib64 \
+    --device /dev/nvidia0:/dev/nvidia0 --device /dev/nvidia1:/dev/nvidia1 \
+    --device /dev/nvidia2:/dev/nvidia2 --device /dev/nvidia3:/dev/nvidia3 \
+    --device /dev/nvidia4:/dev/nvidia4 --device /dev/nvidia5:/dev/nvidia5 \
+    --device /dev/nvidia6:/dev/nvidia6 --device /dev/nvidia7:/dev/nvidia7 \
+    --device /dev/nvidia-uvm:/dev/nvidia-uvm --device /dev/nvidiactl:/dev/nvidiactl \
+    --env LD_LIBRARY_PATH=/usr/local/nvidia/lib64 \
+    --volume /run/tcpx:/run/tcpx \
+    --entrypoint /tcpgpudmarxd/build/app/tcpgpudmarxd \
+    us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/tcpgpudmarxd \
+    --gpu_nic_preset a3vm --gpu_shmem_type fd --uds_path "/run/tcpx" --setup_param "--verbose 128 2 0";
+    sudo iptables -I INPUT -p tcp -m tcp -j ACCEPT;
+    docker run --rm -v /var/lib:/var/lib us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/nccl-plugin-gpudirecttcpx install --install-nccl;
+    sudo mount --bind /var/lib/tcpx /var/lib/tcpx;
+    sudo mount -o remount,exec /var/lib/tcpx;
+    echo "GPU Direct TCPX installed"
+    """)
+
+# Some NCCL options are from the following link.
+# https://docs.nvidia.com/dgx-cloud/run-ai/latest/appendix-gcp.html
+GPU_DIRECT_TCPX_SPECIFIC_OPTIONS = [
+    '--cap-add=IPC_LOCK',
+    '--userns=host',
+    '--volume /run/tcpx:/run/tcpx',
+    '--volume /var/lib/nvidia/lib64:/usr/local/nvidia/lib64',
+    '--volume /var/lib/tcpx/lib64:/usr/local/tcpx/lib64',
+    '--volume /var/lib/nvidia/bin:/usr/local/nvidia/bin',
+    '--shm-size=1g --ulimit memlock=-1 --ulimit stack=67108864',
+    '--device /dev/nvidia0:/dev/nvidia0',
+    '--device /dev/nvidia1:/dev/nvidia1',
+    '--device /dev/nvidia2:/dev/nvidia2',
+    '--device /dev/nvidia3:/dev/nvidia3',
+    '--device /dev/nvidia4:/dev/nvidia4',
+    '--device /dev/nvidia5:/dev/nvidia5',
+    '--device /dev/nvidia6:/dev/nvidia6',
+    '--device /dev/nvidia7:/dev/nvidia7',
+    '--device /dev/nvidia-uvm:/dev/nvidia-uvm',
+    '--device /dev/nvidiactl:/dev/nvidiactl',
+    '--env LD_LIBRARY_PATH=/usr/local/nvidia/lib64:/usr/local/tcpx/lib64',
+    '--env NCCL_GPUDIRECTTCPX_SOCKET_IFNAME=eth1,eth2,eth3,eth4',
+    '--env NCCL_GPUDIRECTTCPX_CTRL_DEV=eth0',
+    '--env NCCL_GPUDIRECTTCPX_TX_BINDINGS="eth1:8-21,112-125;eth2:8-21,112-125;eth3:60-73,164-177;eth4:60-73,164-177"',
+    '--env NCCL_GPUDIRECTTCPX_RX_BINDINGS="eth1:22-35,126-139;eth2:22-35,126-139;eth3:74-87,178-191;eth4:74-87,178-191"',
+    '--env NCCL_GPUDIRECTTCPX_PROGRAM_FLOW_STEERING_WAIT_MICROS=50000',
+    '--env NCCL_GPUDIRECTTCPX_UNIX_CLIENT_PREFIX="/run/tcpx"',
+    '--env NCCL_GPUDIRECTTCPX_FORCE_ACK=0',
+    '--env NCCL_SOCKET_IFNAME=eth0',
+]
+
+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