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

sky/clouds/aws.py

@@ -1,13 +1,14 @@
 """Amazon Web Services."""
 import enum
-import functools
+import fnmatch
+import hashlib
 import json
 import os
 import re
 import subprocess
 import time
 import typing
-from typing import Any, Dict, Iterator, List, Optional, Tuple
+from typing import Any, Dict, Iterator, List, Optional, Set, Tuple, Union
 
 from sky import clouds
 from sky import exceptions
@@ -16,8 +17,12 @@ from sky import sky_logging
 from sky import skypilot_config
 from sky.adaptors import aws
 from sky.clouds import service_catalog
+from sky.clouds.service_catalog import common as catalog_common
+from sky.clouds.utils import aws_utils
 from sky.skylet import constants
+from sky.utils import annotations
 from sky.utils import common_utils
+from sky.utils import registry
 from sky.utils import resources_utils
 from sky.utils import rich_utils
 from sky.utils import subprocess_utils
@@ -26,10 +31,18 @@ from sky.utils import ux_utils
 if typing.TYPE_CHECKING:
     # renaming to avoid shadowing variables
     from sky import resources as resources_lib
-    from sky import status_lib
+    from sky.utils import status_lib
 
 logger = sky_logging.init_logger(__name__)
 
+# Image ID tags
+_DEFAULT_CPU_IMAGE_ID = 'skypilot:custom-cpu-ubuntu'
+# For GPU-related package version,
+# see sky/clouds/service_catalog/images/provisioners/cuda.sh
+_DEFAULT_GPU_IMAGE_ID = 'skypilot:custom-gpu-ubuntu'
+_DEFAULT_GPU_K80_IMAGE_ID = 'skypilot:k80-ubuntu-2004'
+_DEFAULT_NEURON_IMAGE_ID = 'skypilot:neuron-ubuntu-2204'
+
 # This local file (under ~/.aws/) will be uploaded to remote nodes (any
 # cloud), if all of the following conditions hold:
 #   - the current user identity is not using AWS SSO
@@ -83,6 +96,10 @@ class AWSIdentityType(enum.Enum):
 
     CONTAINER_ROLE = 'container-role'
 
+    CUSTOM_PROCESS = 'custom-process'
+
+    ASSUME_ROLE = 'assume-role'
+
     #       Name                    Value             Type    Location
     #       ----                    -----             ----    --------
     #    profile                <not set>             None    None
@@ -91,8 +108,26 @@ class AWSIdentityType(enum.Enum):
     #     region                us-east-1      config-file    ~/.aws/config
     SHARED_CREDENTIALS_FILE = 'shared-credentials-file'
 
+    def can_credential_expire(self) -> bool:
+        """Check if the AWS identity type can expire.
+
+        SSO,IAM_ROLE and CONTAINER_ROLE are temporary credentials and refreshed
+        automatically. ENV and SHARED_CREDENTIALS_FILE are short-lived
+        credentials without refresh.
+        IAM ROLE:
+        https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html
+        SSO/Container-role refresh token:
+        https://docs.aws.amazon.com/solutions/latest/dea-api/auth-refreshtoken.html
+        """
+        # TODO(hong): Add a CLI based check for the expiration of the temporary
+        #  credentials
+        expirable_types = {
+            AWSIdentityType.ENV, AWSIdentityType.SHARED_CREDENTIALS_FILE
+        }
+        return self in expirable_types
+
 
-@clouds.CLOUD_REGISTRY.register
+@registry.CLOUD_REGISTRY.register
 class AWS(clouds.Cloud):
     """Amazon Web Services."""
 
@@ -172,6 +207,10 @@ class AWS(clouds.Cloud):
             regions = [r for r in regions if r.zones]
         return regions
 
+    @classmethod
+    def optimize_by_zone(cls) -> bool:
+        return aws_utils.use_reservations()
+
     @classmethod
     def zones_provision_loop(
         cls,
@@ -196,11 +235,13 @@ class AWS(clouds.Cloud):
                                             zone=None)
         for r in regions:
             assert r.zones is not None, r
-            if num_nodes > 1:
+            if num_nodes > 1 or aws_utils.use_reservations():
                 # When num_nodes > 1, we shouldn't pass a list of zones to the
                 # AWS NodeProvider to try, because it may then place the nodes of
                 # the same cluster in different zones. This is an artifact of the
                 # current AWS NodeProvider implementation.
+                # Also, when using reservations, they are zone-specific, so we
+                # should return one zone at a time.
                 for z in r.zones:
                     yield [z]
             else:
@@ -209,14 +250,20 @@ class AWS(clouds.Cloud):
     @classmethod
     def _get_default_ami(cls, region_name: str, instance_type: str) -> str:
         acc = cls.get_accelerators_from_instance_type(instance_type)
-        image_id = service_catalog.get_image_id_from_tag(
-            'skypilot:gpu-ubuntu-2004', region_name, clouds='aws')
+        image_id = service_catalog.get_image_id_from_tag(_DEFAULT_CPU_IMAGE_ID,
+                                                         region_name,
+                                                         clouds='aws')
         if acc is not None:
+            image_id = service_catalog.get_image_id_from_tag(
+                _DEFAULT_GPU_IMAGE_ID, region_name, clouds='aws')
             assert len(acc) == 1, acc
             acc_name = list(acc.keys())[0]
             if acc_name == 'K80':
                 image_id = service_catalog.get_image_id_from_tag(
-                    'skypilot:k80-ubuntu-2004', region_name, clouds='aws')
+                    _DEFAULT_GPU_K80_IMAGE_ID, region_name, clouds='aws')
+            if acc_name in ['Trainium', 'Inferentia']:
+                image_id = service_catalog.get_image_id_from_tag(
+                    _DEFAULT_NEURON_IMAGE_ID, region_name, clouds='aws')
         if image_id is not None:
             return image_id
         # Raise ResourcesUnavailableError to make sure the failover in
@@ -259,12 +306,12 @@ class AWS(clouds.Cloud):
         if image_id.startswith('skypilot:'):
             return DEFAULT_AMI_GB
         assert region is not None, (image_id, region)
-        client = aws.client('ec2', region_name=region)
         image_not_found_message = (
             f'Image {image_id!r} not found in AWS region {region}.\n'
             f'\nTo find AWS AMI IDs: https://docs.aws.amazon.com/cli/latest/reference/ec2/describe-images.html#examples\n'  # pylint: disable=line-too-long
             'Example: ami-0729d913a335efca7')
         try:
+            client = aws.client('ec2', region_name=region)
             image_info = client.describe_images(ImageIds=[image_id])
             image_info = image_info.get('Images', [])
             if not image_info:
@@ -273,7 +320,8 @@ class AWS(clouds.Cloud):
             image_info = image_info[0]
             image_size = image_info['BlockDeviceMappings'][0]['Ebs'][
                 'VolumeSize']
-        except aws.botocore_exceptions().NoCredentialsError:
+        except (aws.botocore_exceptions().NoCredentialsError,
+                aws.botocore_exceptions().ProfileNotFound):
             # Fallback to default image size if no credentials are available.
             # The credentials issue will be caught when actually provisioning
             # the instance and appropriate errors will be raised there.
@@ -288,7 +336,10 @@ class AWS(clouds.Cloud):
         # The command for getting the current zone is from:
         # https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instance-identity-documents.html  # pylint: disable=line-too-long
         command_str = (
-            'curl -s http://169.254.169.254/latest/dynamic/instance-identity/document'  # pylint: disable=line-too-long
+            'TOKEN=`curl -X PUT "http://169.254.169.254/latest/api/token" '
+            '-H "X-aws-ec2-metadata-token-ttl-seconds: 21600"` && '
+            'curl -H "X-aws-ec2-metadata-token: $TOKEN" -s '
+            'http://169.254.169.254/latest/dynamic/instance-identity/document'
             f' | {constants.SKY_PYTHON_CMD} -u -c "import sys, json; '
             'print(json.load(sys.stdin)[\'availabilityZone\'])"')
         return command_str
@@ -358,7 +409,7 @@ class AWS(clouds.Cloud):
     def get_accelerators_from_instance_type(
         cls,
         instance_type: str,
-    ) -> Optional[Dict[str, int]]:
+    ) -> Optional[Dict[str, Union[int, float]]]:
         return service_catalog.get_accelerators_from_instance_type(
             instance_type, clouds='aws')
 
@@ -370,12 +421,14 @@ class AWS(clouds.Cloud):
         return service_catalog.get_vcpus_mem_from_instance_type(instance_type,
                                                                 clouds='aws')
 
-    def make_deploy_resources_variables(self,
-                                        resources: 'resources_lib.Resources',
-                                        cluster_name_on_cloud: str,
-                                        region: 'clouds.Region',
-                                        zones: Optional[List['clouds.Zone']],
-                                        dryrun: bool = False) -> Dict[str, Any]:
+    def make_deploy_resources_variables(
+            self,
+            resources: 'resources_lib.Resources',
+            cluster_name: resources_utils.ClusterName,
+            region: 'clouds.Region',
+            zones: Optional[List['clouds.Zone']],
+            num_nodes: int,
+            dryrun: bool = False) -> Dict[str, Any]:
         del dryrun  # unused
         assert zones is not None, (region, zones)
 
@@ -385,10 +438,8 @@ class AWS(clouds.Cloud):
         r = resources
         # r.accelerators is cleared but .instance_type encodes the info.
         acc_dict = self.get_accelerators_from_instance_type(r.instance_type)
-        if acc_dict is not None:
-            custom_resources = json.dumps(acc_dict, separators=(',', ':'))
-        else:
-            custom_resources = None
+        custom_resources = resources_utils.make_ray_custom_resources_str(
+            acc_dict)
 
         if r.extract_docker_image() is not None:
             image_id_to_use = None
@@ -397,22 +448,39 @@ class AWS(clouds.Cloud):
         image_id = self._get_image_id(image_id_to_use, region_name,
                                       r.instance_type)
 
-        user_security_group = skypilot_config.get_nested(
+        disk_encrypted = skypilot_config.get_nested(('aws', 'disk_encrypted'),
+                                                    False)
+        user_security_group_config = skypilot_config.get_nested(
             ('aws', 'security_group_name'), None)
-        if resources.ports is not None:
-            # Already checked in Resources._try_validate_ports
-            assert user_security_group is None
-            security_group = USER_PORTS_SECURITY_GROUP_NAME.format(
-                cluster_name_on_cloud)
-        elif user_security_group is not None:
-            assert resources.ports is None
-            security_group = user_security_group
-        else:
+        user_security_group = None
+        if isinstance(user_security_group_config, str):
+            user_security_group = user_security_group_config
+        elif isinstance(user_security_group_config, list):
+            for profile in user_security_group_config:
+                if fnmatch.fnmatchcase(cluster_name.display_name,
+                                       list(profile.keys())[0]):
+                    user_security_group = list(profile.values())[0]
+                    break
+        security_group = user_security_group
+        if security_group is None:
             security_group = DEFAULT_SECURITY_GROUP_NAME
+            if resources.ports is not None:
+                # Already checked in Resources._try_validate_ports
+                security_group = USER_PORTS_SECURITY_GROUP_NAME.format(
+                    cluster_name.display_name)
+        elif resources.ports is not None:
+            with ux_utils.print_exception_no_traceback():
+                logger.warning(
+                    f'Skip opening ports {resources.ports} for cluster {cluster_name!r}, '
+                    'as `aws.security_group_name` in `~/.sky/config.yaml` is specified as '
+                    f' {security_group!r}. Please make sure the specified security group '
+                    'has requested ports setup; or, leave out `aws.security_group_name` '
+                    'in `~/.sky/config.yaml`.')
 
         return {
             'instance_type': r.instance_type,
             'custom_resources': custom_resources,
+            'disk_encrypted': disk_encrypted,
             'use_spot': r.use_spot,
             'region': region_name,
             'zones': ','.join(zone_names),
@@ -425,7 +493,7 @@ class AWS(clouds.Cloud):
 
     def _get_feasible_launchable_resources(
         self, resources: 'resources_lib.Resources'
-    ) -> Tuple[List['resources_lib.Resources'], List[str]]:
+    ) -> resources_utils.FeasibleResources:
         if resources.instance_type is not None:
             assert resources.is_launchable(), resources
             # Check the instance type is valid in the cloud
@@ -436,10 +504,12 @@ class AWS(clouds.Cloud):
                 region=resources.region,
                 zone=resources.zone)
             if not regions:
-                return ([], [])
+                # TODO: Add hints to all return values in this method to help
+                #  users understand why the resources are not launchable.
+                return resources_utils.FeasibleResources([], [], None)
             # Treat Resources(AWS, p3.2x, V100) as Resources(AWS, p3.2x).
             resources = resources.copy(accelerators=None)
-            return ([resources], [])
+            return resources_utils.FeasibleResources([resources], [], None)
 
         def _make(instance_list):
             resource_list = []
@@ -465,9 +535,10 @@ class AWS(clouds.Cloud):
                 memory=resources.memory,
                 disk_tier=resources.disk_tier)
             if default_instance_type is None:
-                return ([], [])
+                return resources_utils.FeasibleResources([], [], None)
             else:
-                return (_make([default_instance_type]), [])
+                return resources_utils.FeasibleResources(
+                    _make([default_instance_type]), [], None)
 
         assert len(accelerators) == 1, resources
         acc, acc_count = list(accelerators.items())[0]
@@ -482,11 +553,14 @@ class AWS(clouds.Cloud):
             zone=resources.zone,
             clouds='aws')
         if instance_list is None:
-            return ([], fuzzy_candidate_list)
-        return (_make(instance_list), fuzzy_candidate_list)
+            return resources_utils.FeasibleResources([], fuzzy_candidate_list,
+                                                     None)
+        return resources_utils.FeasibleResources(_make(instance_list),
+                                                 fuzzy_candidate_list, None)
 
     @classmethod
-    @functools.lru_cache(maxsize=1)  # Cache since getting identity is slow.
+    @annotations.lru_cache(scope='global',
+                           maxsize=1)  # Cache since getting identity is slow.
     def check_credentials(cls) -> Tuple[bool, Optional[str]]:
         """Checks if the user has access credentials to this cloud."""
 
@@ -516,7 +590,7 @@ class AWS(clouds.Cloud):
         # Checks if AWS credentials 1) exist and 2) are valid.
         # https://stackoverflow.com/questions/53548737/verify-aws-credentials-with-boto3
         try:
-            identity_str = cls.get_current_user_identity_str()
+            identity_str = cls.get_active_user_identity_str()
         except exceptions.CloudUserIdentityError as e:
             return False, str(e)
 
@@ -546,14 +620,31 @@ class AWS(clouds.Cloud):
             hints = f'AWS IAM role is set.{single_cloud_hint}'
         elif identity_type == AWSIdentityType.CONTAINER_ROLE:
             # Similar to the IAM ROLE, an ECS container may not store credentials
-            # in the~/.aws/credentials file. So we don't check for the existence of
+            # in the ~/.aws/credentials file. So we don't check for the existence of
             # the file. i.e. the container will be assigned the IAM role of the
             # task: skypilot-v1.
             hints = f'AWS container-role is set.{single_cloud_hint}'
+        elif identity_type == AWSIdentityType.CUSTOM_PROCESS:
+            # Similar to the IAM ROLE, a custom process may not store credentials
+            # in the ~/.aws/credentials file. So we don't check for the existence of
+            # the file. i.e. the custom process will be assigned the IAM role of the
+            # task: skypilot-v1.
+            hints = f'AWS custom-process is set.{single_cloud_hint}'
+        elif identity_type == AWSIdentityType.ASSUME_ROLE:
+            # When using ASSUME ROLE, the credentials are coming from a different
+            # source profile. So we don't check for the existence of ~/.aws/credentials.
+            # i.e. the assumed role will be assigned the IAM role of the
+            # task: skypilot-v1.
+            hints = f'AWS assume-role is set.{single_cloud_hint}'
+        elif identity_type == AWSIdentityType.ENV:
+            # When using ENV vars, the credentials are coming from the environment
+            # variables. So we don't check for the existence of ~/.aws/credentials.
+            # i.e. the identity is not determined by the file.
+            hints = f'AWS env is set.{single_cloud_hint}'
         else:
             # This file is required because it is required by the VMs launched on
             # other clouds to access private s3 buckets and resources like EC2.
-            # `get_current_user_identity` does not guarantee this file exists.
+            # `get_active_user_identity` does not guarantee this file exists.
             if not static_credential_exists:
                 return (False, '~/.aws/credentials does not exist. ' +
                         cls._STATIC_CREDENTIAL_HELP_STR)
@@ -570,21 +661,17 @@ class AWS(clouds.Cloud):
                 'Failed to fetch the availability zones for the account '
                 f'{identity_str}. It is likely due to permission issues, please'
                 ' check the minimal permission required for AWS: '
-                'https://skypilot.readthedocs.io/en/latest/cloud-setup/cloud-permissions/aws.html'  # pylint: disable=
+                'https://docs.skypilot.co/en/latest/cloud-setup/cloud-permissions/aws.html'  # pylint: disable=
                 f'\n{cls._INDENT_PREFIX}Details: '
                 f'{common_utils.format_exception(e, use_bracket=True)}')
         return True, hints
 
     @classmethod
     def _current_identity_type(cls) -> Optional[AWSIdentityType]:
-        proc = subprocess.run('aws configure list',
-                              shell=True,
-                              check=False,
-                              stdout=subprocess.PIPE,
-                              stderr=subprocess.PIPE)
-        if proc.returncode != 0:
+        stdout = cls._aws_configure_list()
+        if stdout is None:
             return None
-        stdout = proc.stdout.decode()
+        output = stdout.decode()
 
         # We determine the identity type by looking at the output of
         # `aws configure list`. The output looks like:
@@ -599,56 +686,35 @@ class AWS(clouds.Cloud):
 
         def _is_access_key_of_type(type_str: str) -> bool:
             # The dot (.) does not match line separators.
-            results = re.findall(fr'access_key.*{type_str}', stdout)
+            results = re.findall(fr'access_key.*{type_str}', output)
             if len(results) > 1:
                 raise RuntimeError(
-                    f'Unexpected `aws configure list` output:\n{stdout}')
+                    f'Unexpected `aws configure list` output:\n{output}')
             return len(results) == 1
 
-        if _is_access_key_of_type(AWSIdentityType.SSO.value):
-            return AWSIdentityType.SSO
-        elif _is_access_key_of_type(AWSIdentityType.IAM_ROLE.value):
-            return AWSIdentityType.IAM_ROLE
-        elif _is_access_key_of_type(AWSIdentityType.CONTAINER_ROLE.value):
-            return AWSIdentityType.CONTAINER_ROLE
-        elif _is_access_key_of_type(AWSIdentityType.ENV.value):
-            return AWSIdentityType.ENV
-        else:
-            return AWSIdentityType.SHARED_CREDENTIALS_FILE
+        for identity_type in AWSIdentityType:
+            if _is_access_key_of_type(identity_type.value):
+                return identity_type
+        return AWSIdentityType.SHARED_CREDENTIALS_FILE
 
     @classmethod
-    def get_current_user_identity(cls) -> Optional[List[str]]:
-        """Returns a [UserId, Account] list that uniquely identifies the user.
-
-        These fields come from `aws sts get-caller-identity`. We permit the same
-        actual user to:
-
-          - switch between different root accounts (after which both elements
-            of the list will be different) and have their clusters owned by
-            each account be protected; or
-
-          - within the same root account, switch between different IAM
-            users, and treat [user_id=1234, account=A] and
-            [user_id=4567, account=A] to be the *same*. Namely, switching
-            between these IAM roles within the same root account will cause
-            the first element of the returned list to differ, and will allow
-            the same actual user to continue to interact with their clusters.
-            Note: this is not 100% safe, since the IAM users can have very
-            specific permissions, that disallow them to access the clusters
-            but it is a reasonable compromise as that could be rare.
-
-        Returns:
-            A list of strings that uniquely identifies the user on this cloud.
-            For identity check, we will fallback through the list of strings
-            until we find a match, and print a warning if we fail for the
-            first string.
+    @annotations.lru_cache(scope='global', maxsize=1)
+    def _aws_configure_list(cls) -> Optional[bytes]:
+        proc = subprocess.run('aws configure list',
+                              shell=True,
+                              check=False,
+                              stdout=subprocess.PIPE,
+                              stderr=subprocess.PIPE)
+        if proc.returncode != 0:
+            return None
+        return proc.stdout
 
-        Raises:
-            exceptions.CloudUserIdentityError: if the user identity cannot be
-                retrieved.
-        """
+    @classmethod
+    @annotations.lru_cache(scope='global',
+                           maxsize=1)  # Cache since getting identity is slow.
+    def _sts_get_caller_identity(cls) -> Optional[List[List[str]]]:
         try:
-            sts = aws.client('sts')
+            sts = aws.client('sts', check_credentials=False)
             # The caller identity contains 3 fields: UserId, Account, Arn.
             # 1. 'UserId' is unique across all AWS entity, which looks like
             # "AROADBQP57FF2AEXAMPLE:role-session-name"
@@ -721,11 +787,80 @@ class AWS(clouds.Cloud):
                     f'Failed to get AWS user.\n'
                     f'  Reason: {common_utils.format_exception(e, use_bracket=True)}.'
                 ) from None
-        return user_ids
+        # TODO: Return a list of identities in the profile when we support
+        #   automatic switching for AWS. Currently we only support one identity.
+        return [user_ids]
 
     @classmethod
-    def get_current_user_identity_str(cls) -> Optional[str]:
-        user_identity = cls.get_current_user_identity()
+    @annotations.lru_cache(scope='global',
+                           maxsize=1)  # Cache since getting identity is slow.
+    def get_user_identities(cls) -> Optional[List[List[str]]]:
+        """Returns a [UserId, Account] list that uniquely identifies the user.
+
+        These fields come from `aws sts get-caller-identity` and are cached
+        locally by `aws configure list` output. The identities are assumed to
+        be stable for the duration of the `sky` process. Modifying the
+        credentials while the `sky` process is running will not affect the
+        identity returned by this function.
+
+        We permit the same actual user to:
+
+          - switch between different root accounts (after which both elements
+            of the list will be different) and have their clusters owned by
+            each account be protected; or
+
+          - within the same root account, switch between different IAM
+            users, and treat [user_id=1234, account=A] and
+            [user_id=4567, account=A] to be the *same*. Namely, switching
+            between these IAM roles within the same root account will cause
+            the first element of the returned list to differ, and will allow
+            the same actual user to continue to interact with their clusters.
+            Note: this is not 100% safe, since the IAM users can have very
+            specific permissions, that disallow them to access the clusters
+            but it is a reasonable compromise as that could be rare.
+
+        Returns:
+            A list of strings that uniquely identifies the user on this cloud.
+            For identity check, we will fallback through the list of strings
+            until we find a match, and print a warning if we fail for the
+            first string.
+
+        Raises:
+            exceptions.CloudUserIdentityError: if the user identity cannot be
+                retrieved.
+        """
+        stdout = cls._aws_configure_list()
+        if stdout is None:
+            # `aws configure list` is not available, possible reasons:
+            # - awscli is not installed but credentials are valid, e.g. run from
+            #   an EC2 instance with IAM role
+            # - aws credentials are not set, proceed anyway to get unified error
+            #   message for users
+            return cls._sts_get_caller_identity()
+        config_hash = hashlib.md5(stdout).hexdigest()[:8]
+        # Getting aws identity cost ~1s, so we cache the result with the output of
+        # `aws configure list` as cache key. Different `aws configure list` output
+        # can have same aws identity, our assumption is the output would be stable
+        # in real world, so the number of cache files would be limited.
+        # TODO(aylei): consider using a more stable cache key and evalute eviction.
+        cache_path = catalog_common.get_catalog_path(
+            f'aws/.cache/user-identity-{config_hash}.txt')
+        if os.path.exists(cache_path):
+            try:
+                with open(cache_path, 'r', encoding='utf-8') as f:
+                    return json.loads(f.read())
+            except json.JSONDecodeError:
+                # cache is invalid, ignore it and fetch identity again
+                pass
+
+        result = cls._sts_get_caller_identity()
+        with open(cache_path, 'w', encoding='utf-8') as f:
+            f.write(json.dumps(result))
+        return result
+
+    @classmethod
+    def get_active_user_identity_str(cls) -> Optional[str]:
+        user_identity = cls.get_active_user_identity()
         if user_identity is None:
             return None
         identity_str = f'{user_identity[0]} [account={user_identity[1]}]'
@@ -762,12 +897,22 @@ class AWS(clouds.Cloud):
             if os.path.exists(os.path.expanduser(f'~/.aws/{filename}'))
         }
 
+    @annotations.lru_cache(scope='global', maxsize=1)
+    def can_credential_expire(self) -> bool:
+        identity_type = self._current_identity_type()
+        return (identity_type is not None and
+                identity_type.can_credential_expire())
+
     def instance_type_exists(self, instance_type):
         return service_catalog.instance_type_exists(instance_type, clouds='aws')
 
     @classmethod
     def _get_disk_type(cls, disk_tier: resources_utils.DiskTier) -> str:
-        return 'standard' if disk_tier == resources_utils.DiskTier.LOW else 'gp3'
+        if disk_tier == resources_utils.DiskTier.LOW:
+            return 'standard'
+        if disk_tier == resources_utils.DiskTier.ULTRA:
+            return 'io2'
+        return 'gp3'
 
     @classmethod
     def _get_disk_specs(
@@ -775,15 +920,19 @@ class AWS(clouds.Cloud):
             disk_tier: Optional[resources_utils.DiskTier]) -> Dict[str, Any]:
         tier = cls._translate_disk_tier(disk_tier)
         tier2iops = {
+            resources_utils.DiskTier.ULTRA: 20000,
             resources_utils.DiskTier.HIGH: 7000,
             resources_utils.DiskTier.MEDIUM: 3500,
-            resources_utils.DiskTier.LOW: 0,  # only gp3 is required to set iops
+            resources_utils.DiskTier.LOW: 0,  # iops is not required on standard disk
         }
         return {
             'disk_tier': cls._get_disk_type(tier),
-            'disk_iops': tier2iops[tier],
-            'disk_throughput': tier2iops[tier] // 16,
-            'custom_disk_perf': tier != resources_utils.DiskTier.LOW,
+            'disk_iops': tier2iops[tier]
+                         if cls._get_disk_type(tier) != 'standard' else None,
+            # Custom disk throughput is only available for gp3
+            # see https://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-properties-ec2-launchtemplate-ebs.html
+            'disk_throughput': tier2iops[tier] // 16
+                               if cls._get_disk_type(tier) == 'gp3' else None,
         }
 
     @classmethod
@@ -800,7 +949,8 @@ class AWS(clouds.Cloud):
         Returns:
             False if the quota is found to be zero, and True otherwise.
         Raises:
-            ImportError: if the dependencies for AWS are not able to be installed.
+            ImportError: if the dependencies for AWS are not able to be
+                installed.
             botocore.exceptions.ClientError: error in Boto3 client request.
         """
 
@@ -814,7 +964,14 @@ class AWS(clouds.Cloud):
         quota_code = aws_catalog.get_quota_code(instance_type, use_spot)
 
         if quota_code is None:
-            # Quota code not found in the catalog for the chosen instance_type, try provisioning anyway
+            # Quota code not found in the catalog for the chosen instance_type,
+            # try provisioning anyway.
+            return True
+
+        if aws_utils.use_reservations():
+            # When reservations are used, it is possible that a user has
+            # reservations for an instance type, but does not have the quota
+            # for that instance type. Skipping the quota check in this case.
             return True
 
         client = aws.client('service-quotas', region_name=region)
@@ -822,7 +979,8 @@ class AWS(clouds.Cloud):
             response = client.get_service_quota(ServiceCode='ec2',
                                                 QuotaCode=quota_code)
         except aws.botocore_exceptions().ClientError:
-            # Botocore client connection not established, try provisioning anyways
+            # Botocore client connection not established, try provisioning
+            # anyways
             return True
 
         if response['Quota']['Value'] == 0:
@@ -832,6 +990,37 @@ class AWS(clouds.Cloud):
         # Quota found to be greater than zero, try provisioning
         return True
 
+    def get_reservations_available_resources(
+        self,
+        instance_type: str,
+        region: str,
+        zone: Optional[str],
+        specific_reservations: Set[str],
+    ) -> Dict[str, int]:
+        if zone is None:
+            # For backward compatibility, the cluster in INIT state launched
+            # before #2352 may not have zone information. In this case, we
+            # return 0 for all reservations.
+            return {reservation: 0 for reservation in specific_reservations}
+        reservations = aws_utils.list_reservations_for_instance_type(
+            instance_type, region)
+
+        filtered_reservations = []
+        for r in reservations:
+            if zone != r.zone:
+                continue
+            if r.targeted:
+                if r.name in specific_reservations:
+                    filtered_reservations.append(r)
+            else:
+                filtered_reservations.append(r)
+        reservation_available_resources = {
+            r.name: r.available_resources for r in filtered_reservations
+        }
+        logger.debug('Get AWS reservations available resources:'
+                     f'{region}-{zone}: {reservation_available_resources}')
+        return reservation_available_resources
+
     @classmethod
     def query_status(cls, name: str, tag_filters: Dict[str, str],
                      region: Optional[str], zone: Optional[str],
@@ -840,22 +1029,24 @@ class AWS(clouds.Cloud):
         assert False, 'This code path should not be used.'
 
     @classmethod
-    def create_image_from_cluster(cls, cluster_name: str,
-                                  cluster_name_on_cloud: str,
+    def create_image_from_cluster(cls,
+                                  cluster_name: resources_utils.ClusterName,
                                   region: Optional[str],
                                   zone: Optional[str]) -> str:
-        assert region is not None, (cluster_name, cluster_name_on_cloud, region)
+        assert region is not None, (cluster_name.display_name,
+                                    cluster_name.name_on_cloud, region)
         del zone  # unused
 
-        image_name = f'skypilot-{cluster_name}-{int(time.time())}'
+        image_name = f'skypilot-{cluster_name.display_name}-{int(time.time())}'
 
-        status = provision_lib.query_instances('AWS', cluster_name_on_cloud,
+        status = provision_lib.query_instances('AWS',
+                                               cluster_name.name_on_cloud,
                                                {'region': region})
         instance_ids = list(status.keys())
         if not instance_ids:
             with ux_utils.print_exception_no_traceback():
                 raise RuntimeError(
-                    f'Failed to find the source cluster {cluster_name!r} on '
+                    f'Failed to find the source cluster {cluster_name.display_name!r} on '
                     'AWS.')
 
         if len(instance_ids) != 1:
@@ -882,7 +1073,7 @@ class AWS(clouds.Cloud):
             stream_logs=True)
 
         rich_utils.force_update_status(
-            f'Waiting for the source image {cluster_name!r} from {region} to be available on AWS.'
+            f'Waiting for the source image {cluster_name.display_name!r} from {region} to be available on AWS.'
         )
         # Wait for the image to be available
         wait_image_cmd = (
@@ -973,7 +1164,7 @@ class AWS(clouds.Cloud):
     @classmethod
     def is_label_valid(cls, label_key: str,
                        label_value: str) -> Tuple[bool, Optional[str]]:
-        key_regex = re.compile(r'^[^aws:][\S]{0,127}$')
+        key_regex = re.compile(r'^(?!aws:)[\S]{1,127}$')
         value_regex = re.compile(r'^[\S]{0,255}$')
         key_valid = bool(key_regex.match(label_key))
         value_valid = bool(value_regex.match(label_value))