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

sky/clouds/aws.py

@@ -1,6 +1,7 @@
 """Amazon Web Services."""
 import enum
 import fnmatch
+import functools
 import hashlib
 import json
 import os
@@ -8,62 +9,55 @@ import re
 import subprocess
 import time
 import typing
-from typing import Any, Dict, Iterator, List, Optional, Set, Tuple, Union
+from typing import (Any, Callable, Dict, Iterator, List, Literal, Optional, Set,
+                    Tuple, TypeVar, Union)
 
+from typing_extensions import ParamSpec
+
+from sky import catalog
 from sky import clouds
 from sky import exceptions
 from sky import provision as provision_lib
 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.adaptors import common
+from sky.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 env_options
 from sky.utils import registry
 from sky.utils import resources_utils
 from sky.utils import rich_utils
 from sky.utils import subprocess_utils
 from sky.utils import ux_utils
+from sky.utils.db import kv_cache
 
 if typing.TYPE_CHECKING:
+    from mypy_boto3_ec2 import type_defs as ec2_type_defs
+
     # renaming to avoid shadowing variables
     from sky import resources as resources_lib
     from sky.utils import status_lib
+    from sky.utils import volume as volume_lib
 
 logger = sky_logging.init_logger(__name__)
 
 # Image ID tags
 _DEFAULT_CPU_IMAGE_ID = 'skypilot:custom-cpu-ubuntu'
+_DEFAULT_CPU_ARM64_IMAGE_ID = 'skypilot:custom-cpu-ubuntu-arm64'
 # For GPU-related package version,
-# see sky/clouds/service_catalog/images/provisioners/cuda.sh
+# see sky/catalog/images/provisioners/cuda.sh
 _DEFAULT_GPU_IMAGE_ID = 'skypilot:custom-gpu-ubuntu'
+_DEFAULT_GPU_ARM64_IMAGE_ID = 'skypilot:custom-gpu-ubuntu-arm64'
 _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
-#   - this file exists
-# It has the following purposes:
-#   - make all nodes (any cloud) able to access private S3 buckets
-#   - make some remote nodes able to launch new nodes on AWS (i.e., makes
-#     AWS head node able to launch AWS workers, or any-cloud jobs controller
-#     able to launch spot clusters on AWS).
-#
-# If we detect the current user identity is AWS SSO, we will not upload this
-# file to any remote nodes (any cloud). Instead, a SkyPilot IAM role is
-# assigned to both AWS head and workers.
-# TODO(skypilot): This also means we leave open a bug for AWS SSO users that
-# use multiple clouds. The non-AWS nodes will have neither the credential
-# file nor the ability to understand AWS IAM.
-_CREDENTIAL_FILES = [
-    'credentials',
-]
-
 DEFAULT_AMI_GB = 45
+DEFAULT_SSH_USER = 'ubuntu'
+DEFAULT_ROOT_DEVICE_NAME = '/dev/sda1'
 
 # Temporary measure, as deleting per-cluster SGs is too slow.
 # See https://github.com/skypilot-org/skypilot/pull/742.
@@ -74,6 +68,151 @@ DEFAULT_SECURITY_GROUP_NAME = f'sky-sg-{common_utils.user_and_hostname_hash()}'
 # Security group to use when user specified ports in their resources.
 USER_PORTS_SECURITY_GROUP_NAME = 'sky-sg-{}'
 
+# GPU instance types that support EFA
+# TODO(hailong): Some CPU instance types also support EFA, may need to support
+# all of them later.
+# TODO(hailong): Add the EFA info in catalog.
+_EFA_INSTANCE_TYPE_PREFIXES = [
+    'g4dn.',
+    'g5.',
+    'g6.',
+    'gr6.',
+    'g6e.',
+    'p4d.',
+    'p4de.',
+    'p5.',
+    'p5e.',
+    'p5en.',
+    'p6-b200.',
+]
+
+# Docker run options for EFA.
+# Refer to https://github.com/ofiwg/libfabric/issues/6437 for updating
+# memlock ulimit
+_EFA_DOCKER_RUN_OPTIONS = [
+    '--cap-add=IPC_LOCK',
+    '--device=/dev/infiniband',
+    '--ulimit memlock=-1:-1',
+]
+
+# AWS EFA image name.
+# Refer to https://docs.aws.amazon.com/dlami/latest/devguide/aws-deep-learning-base-gpu-ami-ubuntu-22-04.html for latest version. # pylint: disable=line-too-long
+# TODO(hailong): may need to update the version later.
+_EFA_IMAGE_NAME = 'Deep Learning Base OSS Nvidia Driver GPU AMI' \
+' (Ubuntu 22.04) 20250808'
+
+# For functions that needs caching per AWS profile.
+_AWS_PROFILE_SCOPED_FUNC_CACHE_SIZE = 5
+
+# Ref: https://docs.aws.amazon.com/cli/v1/userguide/cli-configure-envvars.html
+_DEFAULT_AWS_CONFIG_PATH = '~/.aws/credentials'
+_AWS_CONFIG_FILE_ENV_VAR = 'AWS_CONFIG_FILE'
+
+T = TypeVar('T')
+P = ParamSpec('P')
+
+
+def _get_credentials_path() -> str:
+    cred_path = os.getenv(_AWS_CONFIG_FILE_ENV_VAR, None)
+    if cred_path is not None:
+        if not os.path.isfile(os.path.expanduser(cred_path)):
+            raise FileNotFoundError(f'{_AWS_CONFIG_FILE_ENV_VAR}={cred_path},'
+                                    ' but the file does not exist.')
+        return cred_path
+    # Fallback to the default config path.
+    return _DEFAULT_AWS_CONFIG_PATH
+
+
+def aws_profile_aware_lru_cache(*lru_cache_args,
+                                scope: Literal['global', 'request'] = 'request',
+                                **lru_cache_kwargs) -> Callable:
+    """Similar to annotations.lru_cache, but automatically includes the
+    AWS profile (if set in the workspace config) in the cache key.
+    """
+
+    def decorator(func: Callable[P, T]) -> Callable[P, T]:
+
+        @annotations.lru_cache(scope, *lru_cache_args, **lru_cache_kwargs)
+        def cached_impl(aws_profile, *args, **kwargs):
+            del aws_profile  # Only used as part of the cache key.
+            return func(*args, **kwargs)
+
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            aws_profile = aws.get_workspace_profile()
+            return cached_impl(aws_profile, *args, **kwargs)
+
+        wrapper.cache_clear = cached_impl.cache_clear  # type: ignore[attr-defined]
+        return wrapper
+
+    return decorator
+
+
+def _is_efa_instance_type(instance_type: str) -> bool:
+    """Check if the instance type is in EFA supported instance family."""
+    return any(
+        instance_type.startswith(prefix)
+        for prefix in _EFA_INSTANCE_TYPE_PREFIXES)
+
+
+@annotations.lru_cache(scope='global', maxsize=128)
+def _get_efa_image_id(region_name: str) -> Optional[str]:
+    """Get the EFA image id for the given region."""
+    try:
+        client = aws.client('ec2', region_name=region_name)
+        response = client.describe_images(Filters=[{
+            'Name': 'name',
+            'Values': [_EFA_IMAGE_NAME]
+        }])
+        if 'Images' not in response:
+            return None
+        if len(response['Images']) == 0:
+            return None
+        available_images = [
+            img for img in response['Images'] if img['State'] == 'available'
+        ]
+        if len(available_images) == 0:
+            return None
+        sorted_images = sorted(available_images,
+                               key=lambda x: x['CreationDate'],
+                               reverse=True)
+        return sorted_images[0]['ImageId']
+    except (aws.botocore_exceptions().NoCredentialsError,
+            aws.botocore_exceptions().ProfileNotFound,
+            aws.botocore_exceptions().ClientError) as e:
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(f'Failed to get EFA image id: {e}') from None
+
+
+@annotations.lru_cache(scope='global', maxsize=128)
+def _get_max_efa_interfaces(instance_type: str, region_name: str) -> int:
+    """Get the maximum number of EFA interfaces for the given instance type."""
+    if not _is_efa_instance_type(instance_type):
+        return 0
+    try:
+        client = aws.client('ec2', region_name=region_name)
+        response = client.describe_instance_types(
+            # TODO(cooperc): fix the types for mypy 1.16
+            # Boto3 type stubs expect Literal instance types; using str list here.
+            InstanceTypes=[instance_type],  # type: ignore
+            Filters=[{
+                'Name': 'network-info.efa-supported',
+                'Values': ['true']
+            }])
+        if 'InstanceTypes' in response and len(response['InstanceTypes']) > 0:
+            network_info = response['InstanceTypes'][0]['NetworkInfo']
+            if ('EfaInfo' in network_info and
+                    'MaximumEfaInterfaces' in network_info['EfaInfo']):
+                return network_info['EfaInfo']['MaximumEfaInterfaces']
+        return 0
+    except (aws.botocore_exceptions().NoCredentialsError,
+            aws.botocore_exceptions().ProfileNotFound,
+            aws.botocore_exceptions().ClientError) as e:
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(
+                f'Failed to get max EFA interfaces for {instance_type}: {e}'
+            ) from None
+
 
 class AWSIdentityType(enum.Enum):
     """AWS identity type.
@@ -159,7 +298,9 @@ class AWS(clouds.Cloud):
 
     @classmethod
     def _unsupported_features_for_resources(
-        cls, resources: 'resources_lib.Resources'
+        cls,
+        resources: 'resources_lib.Resources',
+        region: Optional[str] = None,
     ) -> Dict[clouds.CloudImplementationFeatures, str]:
         unsupported_features = {}
         if resources.use_spot:
@@ -173,6 +314,11 @@ class AWS(clouds.Cloud):
                 f'High availability controllers are not supported on {cls._REPR}.'
             )
 
+        unsupported_features[
+            clouds.CloudImplementationFeatures.CUSTOM_MULTI_NETWORK] = (
+                f'Customized multiple network interfaces are not supported on {cls._REPR}.'
+            )
+
         return unsupported_features
 
     @classmethod
@@ -196,12 +342,17 @@ class AWS(clouds.Cloud):
     #### Regions/Zones ####
 
     @classmethod
-    def regions_with_offering(cls, instance_type: str,
-                              accelerators: Optional[Dict[str, int]],
-                              use_spot: bool, region: Optional[str],
-                              zone: Optional[str]) -> List[clouds.Region]:
+    def regions_with_offering(
+        cls,
+        instance_type: str,
+        accelerators: Optional[Dict[str, int]],
+        use_spot: bool,
+        region: Optional[str],
+        zone: Optional[str],
+        resources: Optional['resources_lib.Resources'] = None,
+    ) -> List[clouds.Region]:
         del accelerators  # unused
-        regions = service_catalog.get_region_zones_for_instance_type(
+        regions = catalog.get_region_zones_for_instance_type(
             instance_type, use_spot, 'aws')
 
         if region is not None:
@@ -256,19 +407,30 @@ 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(_DEFAULT_CPU_IMAGE_ID,
+        arch = cls.get_arch_from_instance_type(instance_type)
+        if arch == constants.ARM64_ARCH:
+            image_id = catalog.get_image_id_from_tag(
+                _DEFAULT_CPU_ARM64_IMAGE_ID, region_name, clouds='aws')
+        else:
+            image_id = catalog.get_image_id_from_tag(_DEFAULT_CPU_IMAGE_ID,
+                                                     region_name,
+                                                     clouds='aws')
+        if acc is not None:
+            if arch == constants.ARM64_ARCH:
+                image_id = catalog.get_image_id_from_tag(
+                    _DEFAULT_GPU_ARM64_IMAGE_ID, region_name, clouds='aws')
+            else:
+                image_id = catalog.get_image_id_from_tag(_DEFAULT_GPU_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(
+                image_id = catalog.get_image_id_from_tag(
                     _DEFAULT_GPU_K80_IMAGE_ID, region_name, clouds='aws')
-            if acc_name in ['Trainium', 'Inferentia']:
-                image_id = service_catalog.get_image_id_from_tag(
+            if acc_name.startswith('Trainium') or acc_name.startswith(
+                    'Inferentia'):
+                image_id = catalog.get_image_id_from_tag(
                     _DEFAULT_NEURON_IMAGE_ID, region_name, clouds='aws')
         if image_id is not None:
             return image_id
@@ -286,8 +448,13 @@ class AWS(clouds.Cloud):
         image_id: Optional[Dict[Optional[str], str]],
         region_name: str,
         instance_type: str,
+        enable_efa: bool,
     ) -> str:
         if image_id is None:
+            if enable_efa:
+                efa_image_id = _get_efa_image_id(region_name)
+                if efa_image_id:
+                    return efa_image_id
             return cls._get_default_ami(region_name, instance_type)
         if None in image_id:
             image_id_str = image_id[None]
@@ -295,9 +462,9 @@ class AWS(clouds.Cloud):
             assert region_name in image_id, image_id
             image_id_str = image_id[region_name]
         if image_id_str.startswith('skypilot:'):
-            image_id_str = service_catalog.get_image_id_from_tag(image_id_str,
-                                                                 region_name,
-                                                                 clouds='aws')
+            image_id_str = catalog.get_image_id_from_tag(image_id_str,
+                                                         region_name,
+                                                         clouds='aws')
             if image_id_str is None:
                 # Raise ResourcesUnavailableError to make sure the failover
                 # in CloudVMRayBackend will be correctly triggered.
@@ -307,35 +474,157 @@ class AWS(clouds.Cloud):
                     f'No image found for region {region_name}')
         return image_id_str
 
+    @classmethod
+    def _describe_image_with_retry(
+        cls,
+        image_id: str,
+        region: str,
+        log_context: str,
+    ) -> Optional['ec2_type_defs.ImageTypeDef']:
+        image_not_found_message = (
+            f'Image {image_id!r} not found in AWS region {region} - '
+            f'can\'t get {log_context}.\n\n'
+            f'To 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')
+        max_retries = 3
+        debug_message = 'no describe_images response'
+        for iteration in range(1, max_retries + 1):
+            try:
+                client = aws.client('ec2', region_name=region)
+                response = client.describe_images(ImageIds=[image_id])
+                # These values are not optional, but we will use .get() to avoid
+                # crashing on a malformed response from AWS.
+                metadata = response.get('ResponseMetadata', {})
+                image_info = response.get('Images')
+                debug_message = (
+                    'describe_images response:\n'
+                    f'  status code: {metadata.get("HTTPStatusCode")}\n'
+                    f'  retry attempts: {metadata.get("RetryAttempts")}\n'
+                    f'  len(images): {len(image_info) if image_info else -1}\n'
+                    f'  next token: {response.get("NextToken")}')
+                logger.debug(debug_message)
+                if not image_info:
+                    # image_info is [] (can't find image) or None (invalid
+                    # response from AWS)
+                    with ux_utils.print_exception_no_traceback():
+                        if env_options.Options.SHOW_DEBUG_INFO.get():
+                            image_not_found_message += f'\n{debug_message}'
+                        raise ValueError(image_not_found_message)
+                image = image_info[0]
+                return image
+            except (aws.botocore_exceptions().NoCredentialsError,
+                    aws.botocore_exceptions().ProfileNotFound) as e:
+                # The caller will fall back to its own default value when we
+                # return None. Mention that explicitly in the shared log line.
+                logger.debug(
+                    f'Failed to get {log_context} for {image_id} in region '
+                    f'{region}: {e}. Using default value.')
+                return None
+            except aws.botocore_exceptions().ClientError as e:
+                # This shared log message replaces two attribute-specific
+                # messages (image size/root device) for simplicity.
+                logger.debug(f'Failed to get {log_context} for image '
+                             f'{image_id!r} in region {region}: {e}')
+                if iteration == max_retries:
+                    with ux_utils.print_exception_no_traceback():
+                        if env_options.Options.SHOW_DEBUG_INFO.get():
+                            image_not_found_message += f'\n{debug_message}'
+                            # Note: the ClientError's exception message should
+                            # include most useful info:
+                            # https://github.com/boto/botocore/blob/260a8b91cedae895165984d2102bcbc487de3027/botocore/exceptions.py#L518-L532
+                            additional_info = f'  ClientError: {e}'
+                            logger.debug(additional_info)
+                            image_not_found_message += '\n' + additional_info
+                        raise ValueError(image_not_found_message) from None
+            # linear backoff starting from 0.5 seconds
+            time.sleep(iteration * 0.5)
+        # Should never reach here, but keep type checker happy.
+        raise RuntimeError('Unreachable')
+
     @classmethod
     def get_image_size(cls, image_id: str, region: Optional[str]) -> float:
         if image_id.startswith('skypilot:'):
             return DEFAULT_AMI_GB
         assert region is not None, (image_id, 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]).get(
-                'Images', [])
-            if not image_info:
-                with ux_utils.print_exception_no_traceback():
-                    raise ValueError(image_not_found_message)
-            image_size = image_info[0]['BlockDeviceMappings'][0]['Ebs'][
-                'VolumeSize']
-        except (aws.botocore_exceptions().NoCredentialsError,
-                aws.botocore_exceptions().ProfileNotFound):
+        # first try the cache
+        workspace_profile = aws.get_workspace_profile()
+        kv_cache_key = f'aws:ami:size:{workspace_profile}:{region}:{image_id}'
+        image_size = kv_cache.get_cache_entry(kv_cache_key)
+        if image_size is not None:
+            logger.debug(
+                f'Image size {image_size} found in cache {kv_cache_key}')
+            return float(image_size)
+        # if not found in cache, query the cloud
+        image = cls._describe_image_with_retry(
+            image_id,
+            region,
+            log_context='image size',
+        )
+        if image is None:
             # 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.
             return DEFAULT_AMI_GB
-        except aws.botocore_exceptions().ClientError:
-            with ux_utils.print_exception_no_traceback():
-                raise ValueError(image_not_found_message) from None
+        image_size = image['BlockDeviceMappings'][0]['Ebs']['VolumeSize']
+        # cache the result for a day.
+        # AMIs are immutable, so we can cache the result for a long time.
+        # While AMIs can be deleted, if the AMI is deleted before cache expiration,
+        # the actual VM launch still fails.
+        day_in_seconds = 60 * 60 * 24  # 1 day, 60s * 60m * 24h
+        try:
+            kv_cache.add_or_update_cache_entry(kv_cache_key, str(image_size),
+                                               time.time() + day_in_seconds)
+        except Exception as e:  # pylint: disable=broad-except
+            # Catch the error and continue.
+            # Failure to cache the result is not critical to the
+            # success of this function.
+            logger.debug(
+                f'Failed to cache image size for {image_id} in region {region}: {e}'
+            )
         return image_size
 
+    @classmethod
+    @aws_profile_aware_lru_cache(scope='request',
+                                 maxsize=_AWS_PROFILE_SCOPED_FUNC_CACHE_SIZE)
+    def get_image_root_device_name(cls, image_id: str,
+                                   region: Optional[str]) -> str:
+        if image_id.startswith('skypilot:'):
+            return DEFAULT_ROOT_DEVICE_NAME
+        assert region is not None, (image_id, region)
+        workspace_profile = aws.get_workspace_profile()
+        kv_cache_key = f'aws:ami:root_device_name:{workspace_profile}:{region}:{image_id}'
+        root_device_name = kv_cache.get_cache_entry(kv_cache_key)
+        if root_device_name is not None:
+            logger.debug(f'Image root device name {root_device_name} found in '
+                         f'cache {kv_cache_key}')
+            return root_device_name
+        # if not found in cache, query the cloud
+        image = cls._describe_image_with_retry(
+            image_id,
+            region,
+            log_context='image root device name',
+        )
+        if image is None:
+            return DEFAULT_ROOT_DEVICE_NAME
+        if 'RootDeviceName' not in image:
+            logger.debug(f'Image {image_id!r} does not have a root '
+                         f'device name. '
+                         f'Using {DEFAULT_ROOT_DEVICE_NAME}.')
+            return DEFAULT_ROOT_DEVICE_NAME
+        root_device_name = image['RootDeviceName']
+        day_in_seconds = 60 * 60 * 24  # 1 day, 60s * 60m * 24h
+        try:
+            kv_cache.add_or_update_cache_entry(kv_cache_key, root_device_name,
+                                               time.time() + day_in_seconds)
+        except Exception as e:  # pylint: disable=broad-except
+            # Catch the error and continue.
+            # Failure to cache the result is not critical to the
+            # success of this function.
+            logger.debug(
+                f'Failed to cache image root device name for {image_id} in region {region}: {e}'
+            )
+        return root_device_name
+
     @classmethod
     def get_zone_shell_cmd(cls) -> Optional[str]:
         # The command for getting the current zone is from:
@@ -356,11 +645,11 @@ class AWS(clouds.Cloud):
                                      use_spot: bool,
                                      region: Optional[str] = None,
                                      zone: Optional[str] = None) -> float:
-        return service_catalog.get_hourly_cost(instance_type,
-                                               use_spot=use_spot,
-                                               region=region,
-                                               zone=zone,
-                                               clouds='aws')
+        return catalog.get_hourly_cost(instance_type,
+                                       use_spot=use_spot,
+                                       region=region,
+                                       zone=zone,
+                                       clouds='aws')
 
     def accelerators_to_hourly_cost(self,
                                     accelerators: Dict[str, int],
@@ -397,16 +686,19 @@ class AWS(clouds.Cloud):
         return cost
 
     @classmethod
-    def get_default_instance_type(
-            cls,
-            cpus: Optional[str] = None,
-            memory: Optional[str] = None,
-            disk_tier: Optional[resources_utils.DiskTier] = None
-    ) -> Optional[str]:
-        return service_catalog.get_default_instance_type(cpus=cpus,
-                                                         memory=memory,
-                                                         disk_tier=disk_tier,
-                                                         clouds='aws')
+    def get_default_instance_type(cls,
+                                  cpus: Optional[str] = None,
+                                  memory: Optional[str] = None,
+                                  disk_tier: Optional[
+                                      resources_utils.DiskTier] = None,
+                                  region: Optional[str] = None,
+                                  zone: Optional[str] = None) -> Optional[str]:
+        return catalog.get_default_instance_type(cpus=cpus,
+                                                 memory=memory,
+                                                 disk_tier=disk_tier,
+                                                 region=region,
+                                                 zone=zone,
+                                                 clouds='aws')
 
     # TODO: factor the following three methods, as they are the same logic
     # between Azure and AWS.
@@ -415,48 +707,86 @@ class AWS(clouds.Cloud):
         cls,
         instance_type: str,
     ) -> Optional[Dict[str, Union[int, float]]]:
-        return service_catalog.get_accelerators_from_instance_type(
-            instance_type, clouds='aws')
+        return catalog.get_accelerators_from_instance_type(instance_type,
+                                                           clouds='aws')
+
+    @classmethod
+    def get_arch_from_instance_type(
+        cls,
+        instance_type: str,
+    ) -> Optional[str]:
+        return catalog.get_arch_from_instance_type(instance_type, clouds='aws')
 
     @classmethod
     def get_vcpus_mem_from_instance_type(
         cls,
         instance_type: str,
     ) -> Tuple[Optional[float], Optional[float]]:
-        return service_catalog.get_vcpus_mem_from_instance_type(instance_type,
-                                                                clouds='aws')
+        return catalog.get_vcpus_mem_from_instance_type(instance_type,
+                                                        clouds='aws')
 
     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]:
+        self,
+        resources: 'resources_lib.Resources',
+        cluster_name: resources_utils.ClusterName,
+        region: 'clouds.Region',
+        zones: Optional[List['clouds.Zone']],
+        num_nodes: int,
+        dryrun: bool = False,
+        volume_mounts: Optional[List['volume_lib.VolumeMount']] = None,
+    ) -> Dict[str, Any]:
         del dryrun  # unused
         assert zones is not None, (region, zones)
 
         region_name = region.name
         zone_names = [zone.name for zone in zones]
 
-        r = resources
-        # r.accelerators is cleared but .instance_type encodes the info.
-        acc_dict = self.get_accelerators_from_instance_type(r.instance_type)
+        resources = resources.assert_launchable()
+        # resources.accelerators is cleared but .instance_type encodes the info.
+        acc_dict = self.get_accelerators_from_instance_type(
+            resources.instance_type)
         custom_resources = resources_utils.make_ray_custom_resources_str(
             acc_dict)
 
-        if r.extract_docker_image() is not None:
+        network_tier = (resources.network_tier if resources.network_tier
+                        is not None else resources_utils.NetworkTier.STANDARD)
+        if network_tier == resources_utils.NetworkTier.BEST:
+            max_efa_interfaces = _get_max_efa_interfaces(
+                resources.instance_type, region_name)
+            enable_efa = max_efa_interfaces > 0
+        else:
+            max_efa_interfaces = 0
+            enable_efa = False
+
+        docker_run_options = []
+        if resources.extract_docker_image() is not None:
             image_id_to_use = None
+            if enable_efa:
+                docker_run_options = _EFA_DOCKER_RUN_OPTIONS
         else:
-            image_id_to_use = r.image_id
+            image_id_to_use = resources.image_id
         image_id = self._get_image_id(image_id_to_use, region_name,
-                                      r.instance_type)
-
-        disk_encrypted = skypilot_config.get_nested(('aws', 'disk_encrypted'),
-                                                    False)
-        user_security_group_config = skypilot_config.get_nested(
-            ('aws', 'security_group_name'), None)
+                                      resources.instance_type, enable_efa)
+
+        root_device_name = self.get_image_root_device_name(
+            image_id, region_name)
+
+        ssh_user = skypilot_config.get_effective_region_config(
+            cloud='aws',
+            region=region_name,
+            keys=('ssh_user',),
+            default_value=DEFAULT_SSH_USER)
+
+        disk_encrypted = skypilot_config.get_effective_region_config(
+            cloud='aws',
+            region=region_name,
+            keys=('disk_encrypted',),
+            default_value=False)
+        user_security_group_config = skypilot_config.get_effective_region_config(
+            cloud='aws',
+            region=region_name,
+            keys=('security_group_name',),
+            default_value=None)
         user_security_group = None
         if isinstance(user_security_group_config, str):
             user_security_group = user_security_group_config
@@ -483,17 +813,21 @@ class AWS(clouds.Cloud):
                     'in `~/.sky/config.yaml`.')
 
         return {
-            'instance_type': r.instance_type,
+            'instance_type': resources.instance_type,
             'custom_resources': custom_resources,
             'disk_encrypted': disk_encrypted,
-            'use_spot': r.use_spot,
+            'use_spot': resources.use_spot,
             'region': region_name,
             'zones': ','.join(zone_names),
             'image_id': image_id,
+            'root_device_name': root_device_name,
+            'ssh_user': ssh_user,
             'security_group': security_group,
             'security_group_managed_by_skypilot':
                 str(security_group != user_security_group).lower(),
-            **AWS._get_disk_specs(r.disk_tier)
+            'max_efa_interfaces': max_efa_interfaces,
+            'docker_run_options': docker_run_options,
+            **AWS._get_disk_specs(resources.disk_tier)
         }
 
     def _get_feasible_launchable_resources(
@@ -538,7 +872,9 @@ class AWS(clouds.Cloud):
             default_instance_type = AWS.get_default_instance_type(
                 cpus=resources.cpus,
                 memory=resources.memory,
-                disk_tier=resources.disk_tier)
+                disk_tier=resources.disk_tier,
+                region=resources.region,
+                zone=resources.zone)
             if default_instance_type is None:
                 return resources_utils.FeasibleResources([], [], None)
             else:
@@ -547,16 +883,16 @@ class AWS(clouds.Cloud):
 
         assert len(accelerators) == 1, resources
         acc, acc_count = list(accelerators.items())[0]
-        (instance_list, fuzzy_candidate_list
-        ) = service_catalog.get_instance_type_for_accelerator(
-            acc,
-            acc_count,
-            use_spot=resources.use_spot,
-            cpus=resources.cpus,
-            memory=resources.memory,
-            region=resources.region,
-            zone=resources.zone,
-            clouds='aws')
+        (instance_list,
+         fuzzy_candidate_list) = catalog.get_instance_type_for_accelerator(
+             acc,
+             acc_count,
+             use_spot=resources.use_spot,
+             cpus=resources.cpus,
+             memory=resources.memory,
+             region=resources.region,
+             zone=resources.zone,
+             clouds='aws')
         if instance_list is None:
             return resources_utils.FeasibleResources([], fuzzy_candidate_list,
                                                      None)
@@ -564,22 +900,70 @@ class AWS(clouds.Cloud):
                                                  fuzzy_candidate_list, None)
 
     @classmethod
-    def _check_compute_credentials(cls) -> Tuple[bool, Optional[str]]:
+    def _check_compute_credentials(
+            cls) -> Tuple[bool, Optional[Union[str, Dict[str, str]]]]:
         """Checks if the user has access credentials to this AWS's compute service."""
-        return cls._check_credentials()
+        credentials_exist, identity_str, hints = cls._check_credentials_exist()
+        if not credentials_exist:
+            return False, hints
+
+        # Fetch the AWS catalogs
+        # pylint: disable=import-outside-toplevel
+        from sky.catalog import aws_catalog
+
+        # Trigger the fetch of the availability zones mapping.
+        try:
+            aws_catalog.get_default_instance_type()
+        except RuntimeError as e:
+            return False, (
+                '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://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 _check_storage_credentials(cls) -> Tuple[bool, Optional[str]]:
+    def _check_storage_credentials(
+            cls) -> Tuple[bool, Optional[Union[str, Dict[str, str]]]]:
         """Checks if the user has access credentials to this AWS's storage service."""
-        # TODO(seungjin): Implement separate check for
-        # if the user has access to S3.
-        return cls._check_credentials()
+        credentials_exist, identity_str, hints = cls._check_credentials_exist()
+        if not credentials_exist:
+            return False, hints
+
+        try:
+            # Create an S3 client
+            s3_client = aws.client('s3')
+
+            # Try to list buckets
+            s3_client.list_buckets()
+        except aws.botocore_exceptions().ClientError as e:
+            return False, (
+                'Failed to list buckets for the account '
+                f'{identity_str}. It is likely due to permission issues, please'
+                ' check the storage permission required for AWS: '
+                '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
-    @annotations.lru_cache(scope='request',
-                           maxsize=1)  # Cache since getting identity is slow.
-    def _check_credentials(cls) -> Tuple[bool, Optional[str]]:
-        """Checks if the user has access credentials to AWS."""
+    # Cache since getting identity is slow.
+    @aws_profile_aware_lru_cache(scope='request',
+                                 maxsize=_AWS_PROFILE_SCOPED_FUNC_CACHE_SIZE)
+    def _check_credentials_exist(
+            cls) -> Tuple[bool, Optional[str], Optional[str]]:
+        """Checks if the user has access credentials to AWS.
+
+        Returns:
+            bool: True if credentials exist and are valid.
+            str: Identity string of the user. None if credentials do not exist.
+                 (i.e. the first boolean is False)
+            str: Hints for the user to set up credentials.
+        """
 
         dependency_installation_hints = (
             'AWS dependencies are not installed. '
@@ -595,24 +979,22 @@ class AWS(clouds.Cloud):
                               stdout=subprocess.PIPE,
                               stderr=subprocess.PIPE)
         if proc.returncode != 0:
-            return False, dependency_installation_hints
-        try:
-            # Checks if aws boto is installed properly
-            # pylint: disable=import-outside-toplevel, unused-import
-            import boto3
-            import botocore
-        except ImportError:
-            return False, dependency_installation_hints
+            return False, None, dependency_installation_hints
+
+        # Checks if aws boto is installed properly
+        if not common.can_import_modules(['boto3', 'botocore']):
+            return False, None, dependency_installation_hints
 
         # 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_active_user_identity_str()
         except exceptions.CloudUserIdentityError as e:
-            return False, str(e)
+            return False, None, str(e)
 
+        credentials_path = _get_credentials_path()
         static_credential_exists = os.path.isfile(
-            os.path.expanduser('~/.aws/credentials'))
+            os.path.expanduser(credentials_path))
         hints = None
         identity_type = cls._current_identity_type()
         single_cloud_hint = (
@@ -663,25 +1045,10 @@ class AWS(clouds.Cloud):
             # other clouds to access private s3 buckets and resources like EC2.
             # `get_active_user_identity` does not guarantee this file exists.
             if not static_credential_exists:
-                return (False, '~/.aws/credentials does not exist. ' +
+                return (False, None, f'{credentials_path} does not exist. ' +
                         cls._STATIC_CREDENTIAL_HELP_STR)
 
-        # Fetch the AWS catalogs
-        # pylint: disable=import-outside-toplevel
-        from sky.clouds.service_catalog import aws_catalog
-
-        # Trigger the fetch of the availability zones mapping.
-        try:
-            aws_catalog.get_default_instance_type()
-        except RuntimeError as e:
-            return False, (
-                '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://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
+        return True, identity_str, hints
 
     @classmethod
     def _current_identity_type(cls) -> Optional[AWSIdentityType]:
@@ -715,20 +1082,28 @@ class AWS(clouds.Cloud):
         return AWSIdentityType.SHARED_CREDENTIALS_FILE
 
     @classmethod
-    @annotations.lru_cache(scope='request', maxsize=1)
+    @aws_profile_aware_lru_cache(scope='request',
+                                 maxsize=_AWS_PROFILE_SCOPED_FUNC_CACHE_SIZE)
     def _aws_configure_list(cls) -> Optional[bytes]:
-        proc = subprocess.run('aws configure list',
+        cmd = 'aws configure list'
+        # Profile takes precedence over default configs.
+        profile = aws.get_workspace_profile()
+        if profile is not None:
+            # If profile does not exist, we will get returncode 255.
+            cmd += f' --profile {profile}'
+        proc = subprocess.run(cmd,
                               shell=True,
                               check=False,
                               stdout=subprocess.PIPE,
-                              stderr=subprocess.PIPE)
+                              stderr=subprocess.DEVNULL)
         if proc.returncode != 0:
             return None
         return proc.stdout
 
     @classmethod
-    @annotations.lru_cache(scope='request',
-                           maxsize=1)  # Cache since getting identity is slow.
+    # Cache since getting identity is slow.
+    @aws_profile_aware_lru_cache(scope='request',
+                                 maxsize=_AWS_PROFILE_SCOPED_FUNC_CACHE_SIZE)
     def _sts_get_caller_identity(cls) -> Optional[List[List[str]]]:
         try:
             sts = aws.client('sts', check_credentials=False)
@@ -790,7 +1165,8 @@ class AWS(clouds.Cloud):
                     f'Invalid AWS configuration.\n'
                     f'  Reason: {common_utils.format_exception(e, use_bracket=True)}.'
                 ) from None
-        except aws.botocore_exceptions().TokenRetrievalError:
+        except aws.botocore_exceptions().TokenRetrievalError as e:
+            logger.debug(f'Failed to get AWS caller identity: {e}.')
             # This is raised when the access token is expired, which mainly
             # happens when the user is using temporary credentials or SSO
             # login.
@@ -809,8 +1185,9 @@ class AWS(clouds.Cloud):
         return [user_ids]
 
     @classmethod
-    @annotations.lru_cache(scope='request',
-                           maxsize=1)  # Cache since getting identity is slow.
+    # Cache since getting identity is slow.
+    @aws_profile_aware_lru_cache(scope='request',
+                                 maxsize=_AWS_PROFILE_SCOPED_FUNC_CACHE_SIZE)
     def get_user_identities(cls) -> Optional[List[List[str]]]:
         """Returns a [UserId, Account] list that uniquely identifies the user.
 
@@ -859,7 +1236,7 @@ class AWS(clouds.Cloud):
         # `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.
+        # TODO(aylei): consider using a more stable cache key and evaluate eviction.
         cache_path = catalog_common.get_catalog_path(
             f'aws/.cache/user-identity-{config_hash}.txt')
         if os.path.exists(cache_path):
@@ -905,23 +1282,45 @@ class AWS(clouds.Cloud):
         # provider of the cluster to be launched in this function and make sure
         # the cluster will not be used for launching clusters in other clouds,
         # e.g. jobs controller.
+
         if self._current_identity_type(
         ) != AWSIdentityType.SHARED_CREDENTIALS_FILE:
             return {}
-        return {
-            f'~/.aws/{filename}': f'~/.aws/{filename}'
-            for filename in _CREDENTIAL_FILES
-            if os.path.exists(os.path.expanduser(f'~/.aws/{filename}'))
-        }
 
-    @annotations.lru_cache(scope='request', maxsize=1)
+        # This local credentials file (default to ~/.aws/credentials and can be
+        # overridden by AWS_CONFIG_FILE environment variable) will be uploaded
+        # to remote nodes (any cloud), if all of the following conditions hold:
+        #   - the current user identity is not using AWS SSO
+        #   - this file exists
+        # It has the following purposes:
+        #   - make all nodes (any cloud) able to access private S3 buckets
+        #   - make some remote nodes able to launch new nodes on AWS (i.e., makes
+        #     AWS head node able to launch AWS workers, or any-cloud jobs controller
+        #     able to launch spot clusters on AWS).
+        #
+        # If we detect the current user identity is AWS SSO, we will not upload this
+        # file to any remote nodes (any cloud). Instead, a SkyPilot IAM role is
+        # assigned to both AWS head and workers.
+        # TODO(skypilot): This also means we leave open a bug for AWS SSO users that
+        # use multiple clouds. The non-AWS nodes will have neither the credential
+        # file nor the ability to understand AWS IAM.
+        credentials_path = os.path.expanduser(_get_credentials_path())
+        if os.path.exists(credentials_path):
+            return {
+                # Upload to the default config location on remote cluster.
+                _DEFAULT_AWS_CONFIG_PATH: credentials_path
+            }
+        return {}
+
+    @aws_profile_aware_lru_cache(scope='request',
+                                 maxsize=_AWS_PROFILE_SCOPED_FUNC_CACHE_SIZE)
     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')
+        return catalog.instance_type_exists(instance_type, clouds='aws')
 
     @classmethod
     def _get_disk_type(cls, disk_tier: resources_utils.DiskTier) -> str:
@@ -971,12 +1370,13 @@ class AWS(clouds.Cloud):
             botocore.exceptions.ClientError: error in Boto3 client request.
         """
 
+        resources = resources.assert_launchable()
         instance_type = resources.instance_type
         region = resources.region
         use_spot = resources.use_spot
 
         # pylint: disable=import-outside-toplevel,unused-import
-        from sky.clouds.service_catalog import aws_catalog
+        from sky.catalog import aws_catalog
 
         quota_code = aws_catalog.get_quota_code(instance_type, use_spot)
 
@@ -1056,7 +1456,7 @@ class AWS(clouds.Cloud):
 
         image_name = f'skypilot-{cluster_name.display_name}-{int(time.time())}'
 
-        status = provision_lib.query_instances('AWS',
+        status = provision_lib.query_instances('AWS', cluster_name.display_name,
                                                cluster_name.name_on_cloud,
                                                {'region': region})
         instance_ids = list(status.keys())