konduktor-nightly 0.1.0.dev20250409105017__py3-none-any.whl → 0.1.0.dev20250410104738__py3-none-any.whl

konduktor/data/aws/s3.py

@@ -0,0 +1,1114 @@
+# Proprietary Changes made for Trainy under the Trainy Software License
+# Original source: skypilot: https://github.com/skypilot-org/skypilot
+# which is Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Amazon Web Services (AWS) S3 Storage."""
+
+import enum
+import hashlib
+import os
+import re
+import shlex
+import subprocess
+import tempfile
+import time
+import typing
+import uuid
+from typing import Any, Dict, List, Optional, Tuple
+
+import boto3
+import colorama
+
+from konduktor import config, logging
+from konduktor.adaptors import aws
+from konduktor.data import constants, data_utils, storage_utils
+from konduktor.utils import (
+    annotations,
+    base64_utils,
+    common_utils,
+    exceptions,
+    kubernetes_utils,
+    rich_utils,
+    ux_utils,
+)
+
+logger = logging.get_logger(__name__)
+
+# Maximum number of concurrent rsync upload processes
+_MAX_CONCURRENT_UPLOADS = 32
+
+_CREDENTIAL_FILES = ['credentials', 'config']
+
+AWS_SECRET_NAME = 'awscredentials'
+AWS_CREDENTIALS_KEY = 'awscredentials'
+
+DEFAULT_AWS_CREDENTIALS_DIR = '~/.aws/'
+DEFAULT_AWS_CREDENTIAL_PATH = '~/.aws/credentials'
+DEFAULT_AWS_CONFIG_PATH = '~/.aws/config'
+
+
+class AWSIdentityType(enum.Enum):
+    """AWS identity type.
+
+    The account type is determined by the current user identity, based on `aws
+    configure list`. We will check the existence of the value in the output of
+    `aws configure list` to determine the account type.
+    """
+
+    #       Name                    Value             Type    Location
+    #       ----                    -----             ----    --------
+    #    profile                     1234              env    ...
+    # access_key     ****************abcd              sso
+    # secret_key     ****************abcd              sso
+    #     region                <not set>             None    None
+    SSO = 'sso'
+    ENV = 'env'
+    IAM_ROLE = 'iam-role'
+    CONTAINER_ROLE = 'container-role'
+    CUSTOM_PROCESS = 'custom-process'
+    ASSUME_ROLE = 'assume-role'
+
+    #       Name                    Value             Type    Location
+    #       ----                    -----             ----    --------
+    #    profile                <not set>             None    None
+    # access_key     ****************abcd shared-credentials-file
+    # secret_key     ****************abcd shared-credentials-file
+    #     region                us-east-1      config-file    ~/.aws/config
+    SHARED_CREDENTIALS_FILE = 'shared-credentials-file'
+
+    # IN GCS.PY
+    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
+
+
+class S3Store(storage_utils.AbstractStore):
+    """S3Store inherits from Storage Object and represents the backend
+    for S3 buckets.
+    """
+
+    # k8s secret name for aws credentials
+    _AWS_SECRET_NAME = f'{AWS_SECRET_NAME}-{common_utils.user_and_hostname_hash()}'
+    _AWS_CREDENTIALS_KEY = AWS_CREDENTIALS_KEY
+
+    _DEFAULT_REGION = 'us-east-1'
+    _ACCESS_DENIED_MESSAGE = 'Access Denied'
+    _CUSTOM_ENDPOINT_REGIONS = [
+        'ap-east-1',
+        'me-south-1',
+        'af-south-1',
+        'eu-south-1',
+        'eu-south-2',
+        'ap-south-2',
+        'ap-southeast-3',
+        'ap-southeast-4',
+        'me-central-1',
+        'il-central-1',
+    ]
+
+    _INDENT_PREFIX = '    '
+
+    _STATIC_CREDENTIAL_HELP_STR = (
+        'Run the following commands:'
+        f'\n{_INDENT_PREFIX}  $ aws configure'
+        f'\n{_INDENT_PREFIX}  $ aws configure list '
+        '# Ensure that this shows identity is set.'
+        f'\n{_INDENT_PREFIX}For more info: '
+        'https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-quickstart.html'  # pylint: disable=line-too-long
+    )
+
+    _REPR = 'S3Store'
+
+    def __init__(
+        self,
+        name: str,
+        source: str,
+        region: Optional[str] = _DEFAULT_REGION,
+        is_sky_managed: Optional[bool] = False,
+        sync_on_reconstruction: Optional[bool] = True,
+        _bucket_sub_path: Optional[str] = None,
+    ):
+        self.client: 'boto3.client.Client'
+        self.bucket: 'constants.StorageHandle'
+        if region in self._CUSTOM_ENDPOINT_REGIONS:
+            logger.warning(
+                'AWS opt-in regions are not supported for S3. '
+                f'Falling back to default region '
+                f'{self._DEFAULT_REGION} for bucket {name!r}.'
+            )
+            region = self._DEFAULT_REGION
+        super().__init__(
+            name,
+            source,
+            region,
+            is_sky_managed,
+            sync_on_reconstruction,
+            _bucket_sub_path,
+        )
+
+    def __repr__(self):
+        return self._REPR
+
+    # IN GCS.PY
+    def _validate(self):
+        if self.source is not None and isinstance(self.source, str):
+            if self.source.startswith('s3://'):
+                assert self.name == data_utils.split_s3_path(self.source)[0], (
+                    'S3 Bucket is specified as path, the name should be the'
+                    ' same as S3 bucket.'
+                )
+                assert data_utils.verify_s3_bucket(self.name), (
+                    f'Source specified as {self.source}, an S3 bucket. ',
+                    'S3 Bucket should exist.',
+                )
+            # if self.source.startswith('gs://'):
+            #     assert self.name == data_utils.split_gcs_path(self.source)[0], (
+            #         'GCS Bucket is specified as path, the name should be '
+            #         'the same as GCS bucket.'
+            #     )
+            # elif data_utils.is_az_container_endpoint(self.source):
+            #     storage_account_name, container_name, _ = (
+            #         data_utils.split_az_path(self.source))
+            #     assert self.name == container_name, (
+            #         'Azure bucket is specified as path, the name should be '
+            #         'the same as Azure bucket.')
+            #     assert data_utils.verify_az_bucket(
+            #         storage_account_name, self.name), (
+            #             f'Source specified as {self.source}, an Azure bucket. '
+            #             'Azure bucket should exist.')
+            # elif self.source.startswith('r2://'):
+            #     assert self.name == data_utils.split_r2_path(self.source)[0], (
+            #         'R2 Bucket is specified as path, the name should be '
+            #         'the same as R2 bucket.')
+            #     assert data_utils.verify_r2_bucket(self.name), (
+            #         f'Source specified as {self.source}, a R2 bucket. ',
+            #         'R2 Bucket should exist.')
+            # elif self.source.startswith('cos://'):
+            #     assert self.name == data_utils.split_cos_path(self.source)[0], (
+            #         'COS Bucket is specified as path, the name should be '
+            #         'the same as COS bucket.')
+            #     assert data_utils.verify_ibm_cos_bucket(self.name), (
+            #         f'Source specified as {self.source}, a COS bucket. ',
+            #         'COS Bucket should exist.')
+        # Validate name
+        self.name = self.validate_name(self.name)
+
+    # IN GCS.PY
+    @classmethod
+    def validate_name(cls, name: str) -> str:
+        """Validates the name of the S3 store.
+
+        Source for rules:
+        https://docs.aws.amazon.com/AmazonS3/latest/userguide/bucketnamingrules.html
+        """
+
+        def _raise_no_traceback_name_error(err_str):
+            with ux_utils.print_exception_no_traceback():
+                raise exceptions.StorageNameError(err_str)
+
+        if name is not None and isinstance(name, str):
+            # Check for overall length
+            if not 3 <= len(name) <= 63:
+                _raise_no_traceback_name_error(
+                    f'Invalid store name: name {name} must be between 3 (min) '
+                    'and 63 (max) characters long.'
+                )
+
+            # Check for valid characters and start/end with a number or letter
+            pattern = r'^[a-z0-9][-a-z0-9._]*[a-z0-9]$'
+            if not re.match(pattern, name):
+                _raise_no_traceback_name_error(
+                    f'Invalid store name: name {name} can consist only of '
+                    'lowercase letters, numbers, dots (.), and hyphens (-). '
+                    'It must begin and end with a letter or number.'
+                )
+
+            # Check for two adjacent periods
+            if '..' in name:
+                _raise_no_traceback_name_error(
+                    f'Invalid store name: name {name} must not contain '
+                    'two adjacent periods.'
+                )
+
+            # Check for IP address format
+            ip_pattern = r'^(?:\d{1,3}\.){3}\d{1,3}$'
+            if re.match(ip_pattern, name):
+                _raise_no_traceback_name_error(
+                    f'Invalid store name: name {name} must not be formatted as '
+                    'an IP address (for example, 192.168.5.4).'
+                )
+
+            # Check for 'xn--' prefix
+            if name.startswith('xn--'):
+                _raise_no_traceback_name_error(
+                    f'Invalid store name: name {name} must not start with the '
+                    'prefix "xn--".'
+                )
+
+            # Check for '-s3alias' suffix
+            if name.endswith('-s3alias'):
+                _raise_no_traceback_name_error(
+                    f'Invalid store name: name {name} must not end with the '
+                    'suffix "-s3alias".'
+                )
+
+            # Check for '--ol-s3' suffix
+            if name.endswith('--ol-s3'):
+                _raise_no_traceback_name_error(
+                    f'Invalid store name: name {name} must not end with the '
+                    'suffix "--ol-s3".'
+                )
+        else:
+            _raise_no_traceback_name_error('Store name must be specified.')
+        return name
+
+    # IN GCS.PY
+    def initialize(self):
+        """Initializes the S3 store object on the cloud.
+
+        Initialization involves fetching bucket if exists, or creating it if
+        it does not.
+
+        Raises:
+          StorageBucketCreateError: If bucket creation fails
+          StorageBucketGetError: If fetching existing bucket fails
+          StorageInitError: If general initialization fails.
+        """
+        self.client = data_utils.create_s3_client(self.region)
+        self.bucket, is_new_bucket = self._get_bucket()
+        if self.is_sky_managed is None:
+            # If is_sky_managed is not specified, then this is a new storage
+            # object (i.e., did not exist in global_user_state) and we should
+            # set the is_sky_managed property.
+            # If is_sky_managed is specified, then we take no action.
+            self.is_sky_managed = is_new_bucket
+
+    # IN GCS.PY
+    def upload(self):
+        """Uploads source to store bucket.
+
+        Upload must be called by the Storage handler - it is not called on
+        Store initialization.
+
+        Raises:
+            StorageUploadError: if upload fails.
+        """
+        try:
+            if isinstance(self.source, list):
+                self.batch_aws_rsync(self.source, create_dirs=True)
+            elif self.source is not None:
+                if self.source.startswith('s3://'):
+                    pass
+                # elif self.source.startswith('gs://'):
+                #     self._transfer_to_s3()
+                # elif self.source.startswith('r2://'):
+                #     self._transfer_to_s3()
+                else:
+                    self.batch_aws_rsync([self.source])
+        except exceptions.StorageUploadError:
+            raise
+        except Exception as e:
+            raise exceptions.StorageUploadError(
+                f'Upload failed for store {self.name}'
+            ) from e
+
+    # IN GCS.PY
+    def delete(self) -> None:
+        deleted_by_skypilot = self._delete_s3_bucket(self.name)
+        if deleted_by_skypilot:
+            msg_str = f'Deleted S3 bucket {self.name}.'
+        else:
+            msg_str = (
+                f'S3 bucket {self.name} may have been deleted '
+                f'externally. Removing from local state.'
+            )
+        logger.info(f'{colorama.Fore.GREEN}{msg_str}' f'{colorama.Style.RESET_ALL}')
+
+    # IN GCS.PY
+    def get_handle(self) -> 'constants.StorageHandle':
+        return aws.resource('s3').Bucket(self.name)
+
+    # FROM data/storage.py but matches GCS.PY batch_gsutil_rsync() (s3 specific)
+    def batch_aws_rsync(
+        self, source_path_list: List['constants.Path'], create_dirs: bool = False
+    ) -> None:
+        """Invokes aws s3 sync to batch upload a list of local paths to S3
+
+        AWS Sync by default uses 10 threads to upload files to the bucket.  To
+        increase parallelism, modify max_concurrent_requests in your aws config
+        file (Default path: ~/.aws/config).
+
+        Since aws s3 sync does not support batch operations, we construct
+        multiple commands to be run in parallel.
+
+        Args:
+            source_path_list: List of paths to local files or directories
+            create_dirs: If the local_path is a directory and this is set to
+                False, the contents of the directory are directly uploaded to
+                root of the bucket. If the local_path is a directory and this is
+                set to True, the directory is created in the bucket root and
+                contents are uploaded to it.
+        """
+        sub_path = f'/{self._bucket_sub_path}' if self._bucket_sub_path else ''
+
+        def get_file_sync_command(base_dir_path, file_names):
+            includes = ' '.join(
+                [f'--include {shlex.quote(file_name)}' for file_name in file_names]
+            )
+            base_dir_path = shlex.quote(base_dir_path)
+            sync_command = (
+                'aws s3 sync --no-follow-symlinks --exclude="*" '
+                f'{includes} {base_dir_path} '
+                f's3://{self.name}{sub_path}'
+            )
+            return sync_command
+
+        def get_dir_sync_command(src_dir_path, dest_dir_name):
+            # we exclude .git directory from the sync
+            excluded_list = storage_utils.get_excluded_files(src_dir_path)
+            excluded_list.append('.git/*')
+            excludes = ' '.join(
+                [f'--exclude {shlex.quote(file_name)}' for file_name in excluded_list]
+            )
+            src_dir_path = shlex.quote(src_dir_path)
+            sync_command = (
+                f'aws s3 sync --no-follow-symlinks {excludes} '
+                f'{src_dir_path} '
+                f's3://{self.name}{sub_path}/{dest_dir_name}'
+            )
+            return sync_command
+
+        # Generate message for upload
+        if len(source_path_list) > 1:
+            source_message = f'{len(source_path_list)} paths'
+        else:
+            source_message = source_path_list[0]
+
+        log_path = logging.generate_tmp_logging_file_path(
+            constants._STORAGE_LOG_FILE_NAME
+        )
+        sync_path = f'{source_message} -> s3://{self.name}{sub_path}/'
+        with rich_utils.safe_status(
+            ux_utils.spinner_message(f'Syncing {sync_path}', log_path=log_path)
+        ):
+            data_utils.parallel_upload(
+                source_path_list,
+                get_file_sync_command,
+                get_dir_sync_command,
+                log_path,
+                self.name,
+                self._ACCESS_DENIED_MESSAGE,
+                create_dirs=create_dirs,
+                max_concurrent_uploads=_MAX_CONCURRENT_UPLOADS,
+            )
+        logger.info(
+            ux_utils.finishing_message(f'Storage synced: {sync_path}', log_path)
+        )
+
+    # IN GCS.PY
+    def _get_bucket(self) -> Tuple['constants.StorageHandle', bool]:
+        """Obtains the S3 bucket.
+
+        If the bucket exists, this method will return the bucket.
+        If the bucket does not exist, there are three cases:
+          1) Raise an error if the bucket source starts with s3://
+          2) Return None if bucket has been externally deleted and
+             sync_on_reconstruction is False
+          3) Create and return a new bucket otherwise
+
+        Raises:
+            StorageSpecError: If externally created bucket is attempted to be
+                mounted without specifying storage source.
+            StorageBucketCreateError: If creating the bucket fails
+            StorageBucketGetError: If fetching a bucket fails
+            StorageExternalDeletionError: If externally deleted storage is
+                attempted to be fetched while reconstructing the storage for
+                'sky storage delete' or 'sky start'
+        """
+        s3 = aws.resource('s3')
+        bucket = s3.Bucket(self.name)
+
+        try:
+            # Try Public bucket case.
+            # This line does not error out if the bucket is an external public
+            # bucket or if it is a user's bucket that is publicly
+            # accessible.
+            self.client.head_bucket(Bucket=self.name)
+            self._validate_existing_bucket()
+            return bucket, False
+        except aws.botocore_exceptions().ClientError as e:
+            error_code = e.response['Error']['Code']
+            # AccessDenied error for buckets that are private and not owned by
+            # user.
+            if error_code == '403':
+                command = f'aws s3 ls {self.name}'
+                with ux_utils.print_exception_no_traceback():
+                    raise exceptions.StorageBucketGetError(
+                        f'Bucket {self.name} does not exist.'
+                        + f' To debug, consider running `{command}`.'
+                    ) from e
+
+        if isinstance(self.source, str) and self.source.startswith('s3://'):
+            with ux_utils.print_exception_no_traceback():
+                raise exceptions.StorageBucketGetError(
+                    'Attempted to use a non-existent bucket as a source: '
+                    f'{self.source}. Consider using `aws s3 ls '
+                    f'{self.source}` to debug.'
+                )
+
+        # If bucket cannot be found in both private and public settings,
+        # the bucket is to be created by Sky. However, creation is skipped if
+        # Store object is being reconstructed for deletion or re-mount with
+        # sky start, and error is raised instead.
+        if self.sync_on_reconstruction:
+            bucket = self._create_s3_bucket(self.name, self.region)
+            return bucket, True
+        else:
+            # Raised when Storage object is reconstructed for sky storage
+            # delete or to re-mount Storages with sky start but the storage
+            # is already removed externally.
+            raise exceptions.StorageExternalDeletionError(
+                'Attempted to fetch a non-existent bucket: ' f'{self.name}'
+            )
+
+    # IN GCS.PY
+    def _download_file(self, remote_path: str, local_path: str) -> None:
+        """Downloads file from remote to local on s3 bucket
+        using the boto3 API
+
+        Args:
+          remote_path: str; Remote path on S3 bucket
+          local_path: str; Local path on user's device
+        """
+        self.bucket.download_file(remote_path, local_path)
+
+    # IN GCS.PY
+    def _create_s3_bucket(
+        self, bucket_name: str, region=_DEFAULT_REGION
+    ) -> 'constants.StorageHandle':
+        """Creates S3 bucket with specific name in specific region
+
+        Args:
+          bucket_name: str; Name of bucket
+          region: str; Region name, e.g. us-west-1, us-east-2
+        Raises:
+          StorageBucketCreateError: If bucket creation fails.
+        """
+        s3_client = self.client
+        try:
+            create_bucket_config: Dict[str, Any] = {'Bucket': bucket_name}
+            # If default us-east-1 region of create_bucket API is used,
+            # the LocationConstraint must not be specified.
+            # Reference: https://stackoverflow.com/a/51912090
+            if region is not None and region != 'us-east-1':
+                create_bucket_config['CreateBucketConfiguration'] = {
+                    'LocationConstraint': region
+                }
+            s3_client.create_bucket(**create_bucket_config)
+            logger.info(
+                f'  {colorama.Style.DIM}Created S3 bucket {bucket_name!r} in '
+                f'{region or "us-east-1"}{colorama.Style.RESET_ALL}'
+            )
+
+            # Add AWS tags configured in config.yaml to the bucket.
+            # This is useful for cost tracking and external cleanup.
+            bucket_tags = config.get_nested(('aws', 'labels'), {})
+            if bucket_tags:
+                s3_client.put_bucket_tagging(
+                    Bucket=bucket_name,
+                    Tagging={
+                        'TagSet': [
+                            {'Key': k, 'Value': v} for k, v in bucket_tags.items()
+                        ]
+                    },
+                )
+
+        except aws.botocore_exceptions().ClientError as e:
+            with ux_utils.print_exception_no_traceback():
+                raise exceptions.StorageBucketCreateError(
+                    f'Attempted to create a bucket {self.name} but failed.'
+                ) from e
+        return aws.resource('s3').Bucket(bucket_name)
+
+    # NOT IN GCS.PY but FROM data/storage.py (s3 specific)
+    def _execute_s3_remove_command(
+        self, command: str, bucket_name: str, hint_operating: str, hint_failed: str
+    ) -> bool:
+        try:
+            with rich_utils.safe_status(ux_utils.spinner_message(hint_operating)):
+                subprocess.check_output(command.split(' '), stderr=subprocess.STDOUT)
+        except subprocess.CalledProcessError as e:
+            if 'NoSuchBucket' in e.output.decode('utf-8'):
+                logger.debug(f'Bucket {bucket_name} does not exist.')
+                return False
+            else:
+                with ux_utils.print_exception_no_traceback():
+                    raise exceptions.StorageBucketDeleteError(
+                        f'{hint_failed}' f'Detailed error: {e.output}'
+                    )
+        return True
+
+    # IN GCS.PY
+    def _delete_s3_bucket(self, bucket_name: str) -> bool:
+        """Deletes S3 bucket, including all objects in bucket
+
+        Args:
+            bucket_name: str; Name of bucket
+
+        Returns:
+            bool; True if bucket was deleted, False if it was deleted externally.
+
+        Raises:
+            StorageBucketDeleteError: If deleting the bucket fails.
+        """
+        # Deleting objects is very slow programatically
+        # (i.e. bucket.objects.all().delete() is slow).
+        # In addition, standard delete operations (i.e. via `aws s3 rm`)
+        # are slow, since AWS puts deletion markers.
+        # https://stackoverflow.com/questions/49239351/why-is-it-so-much-slower-to-delete-objects-in-aws-s3-than-it-is-to-create-them
+        # The fastest way to delete is to run `aws s3 rb --force`,
+        # which removes the bucket by force.
+        remove_command = f'aws s3 rb s3://{bucket_name} --force'
+        success = self._execute_s3_remove_command(
+            remove_command,
+            bucket_name,
+            f'Deleting S3 bucket [green]{bucket_name}[/]',
+            f'Failed to delete S3 bucket {bucket_name}.',
+        )
+        if not success:
+            return False
+
+        # Wait until bucket deletion propagates on AWS servers
+        while data_utils.verify_s3_bucket(bucket_name):
+            time.sleep(0.1)
+        return True
+
+    # NOT IN GCS.PY but FROM data/storage.py (s3 specific)
+    def _delete_s3_bucket_sub_path(self, bucket_name: str, sub_path: str) -> bool:
+        """Deletes the sub path from the bucket."""
+        remove_command = f'aws s3 rm s3://{bucket_name}/{sub_path}/ --recursive'
+        return self._execute_s3_remove_command(
+            remove_command,
+            bucket_name,
+            f'Removing objects from S3 bucket ' f'[green]{bucket_name}/{sub_path}[/]',
+            f'Failed to remove objects from S3 bucket {bucket_name}/{sub_path}.',
+        )
+
+    @classmethod
+    @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
+
+    @classmethod
+    def _sso_credentials_help_str(cls, expired: bool = False) -> str:
+        help_str = 'Run the following commands (must use AWS CLI v2):'
+        if not expired:
+            help_str += f'\n{cls._INDENT_PREFIX}  $ aws configure sso'
+        help_str += (
+            f'\n{cls._INDENT_PREFIX}  $ aws sso login --profile <profile_name>'
+            f'\n{cls._INDENT_PREFIX}For more info: '
+            'https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-sso.html'
+        )
+        return help_str
+
+    @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', 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"
+            # 2. 'Account' can be shared by multiple users under the same
+            # organization
+            # 3. 'Arn' is the full path to the user, which can be reused when
+            # the user is deleted and recreated.
+            # Refer to: <https://docs.aws.amazon.com/cli/latest/reference/sts/get-caller-identity.html>
+            # and <https://docs.aws.amazon.com/IAM/latest/UserGuide/reference_policies_variables.html#principaltable>
+            user_info = sts.get_caller_identity()
+            # Allow fallback to AccountId if UserId does not match, because:
+            # 1. In the case where multiple IAM users belong a single root account,
+            # those users normally share the visibility of the VMs, so we do not
+            # need to identity them with each other. (There can be some cases,
+            # when an IAM user is given a limited permission by the admin, we may
+            # ignore that case for now, or print out a warning if the underlying
+            # userid changed for a cluster).
+            # 2. In the case where the multiple users belong to an organization,
+            # those users will have different account id, so fallback works.
+            user_ids = [user_info['UserId'], user_info['Account']]
+        except aws.botocore_exceptions().NoCredentialsError as e:
+            with ux_utils.print_exception_no_traceback():
+                raise exceptions.CloudUserIdentityError(
+                    'AWS credentials are not set. '
+                    f'{cls._STATIC_CREDENTIAL_HELP_STR}\n'
+                    f'{cls._INDENT_PREFIX}Details: `aws sts '
+                    'get-caller-identity` failed with error:'
+                    f' {common_utils.format_exception(e, use_bracket=True)}.'
+                ) from None
+        except aws.botocore_exceptions().ClientError as e:
+            with ux_utils.print_exception_no_traceback():
+                raise exceptions.CloudUserIdentityError(
+                    'Failed to access AWS services with credentials. '
+                    'Make sure that the access and secret keys are correct.'
+                    f' {cls._STATIC_CREDENTIAL_HELP_STR}\n'
+                    f'{cls._INDENT_PREFIX}Details: `aws sts '
+                    'get-caller-identity` failed with error:'
+                    f' {common_utils.format_exception(e, use_bracket=True)}.'
+                ) from None
+        except aws.botocore_exceptions().InvalidConfigError as e:
+            # pylint: disable=import-outside-toplevel
+            import awscli
+            from packaging import version
+
+            awscli_version = version.parse(awscli.__version__)
+            if awscli_version < version.parse(
+                '1.27.10'
+            ) and 'configured to use SSO' in str(e):
+                with ux_utils.print_exception_no_traceback():
+                    raise exceptions.CloudUserIdentityError(
+                        'awscli is too old to use SSO.'
+                        'Run the following command to upgrade:'
+                        f'\n{cls._INDENT_PREFIX}  $ pip install awscli>=1.27.10'
+                        f'\n{cls._INDENT_PREFIX}You may need to log into SSO again '
+                        f'after upgrading. {cls._sso_credentials_help_str()}'
+                    ) from None
+            with ux_utils.print_exception_no_traceback():
+                raise exceptions.CloudUserIdentityError(
+                    f'Invalid AWS configuration.\n'
+                    f'  Reason: {common_utils.format_exception(e, use_bracket=True)}.'
+                ) from None
+        except aws.botocore_exceptions().TokenRetrievalError:
+            # This is raised when the access token is expired, which mainly
+            # happens when the user is using temporary credentials or SSO
+            # login.
+            with ux_utils.print_exception_no_traceback():
+                raise exceptions.CloudUserIdentityError(
+                    'AWS access token is expired.'
+                    f' {cls._sso_credentials_help_str(expired=True)}'
+                ) from None
+        except Exception as e:  # pylint: disable=broad-except
+            with ux_utils.print_exception_no_traceback():
+                raise exceptions.CloudUserIdentityError(
+                    f'Failed to get AWS user.\n'
+                    f'  Reason: {common_utils.format_exception(e, use_bracket=True)}.'
+                ) from None
+        # 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]
+
+    # IN GCS.PY
+    @classmethod
+    @annotations.lru_cache(
+        scope='global', maxsize=1
+    )  # Cache since getting identity is slow.
+    def get_user_identities(cls) -> 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]  # noqa: F841
+        # 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.
+        # TODO:(ryan) ??? Ignoring caching for now (returning early)
+        return cls._sts_get_caller_identity()
+        # 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
+
+    # IN GCS.PY
+    @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]}]'
+        return identity_str
+
+    # IN GCS.PY
+    @classmethod
+    @annotations.lru_cache(scope='global', maxsize=1)
+    def check_credentials(cls) -> Tuple[bool, Optional[str]]:
+        """Checks if the user has access credentials to AWS."""
+
+        dependency_installation_hints = (
+            'AWS dependencies are not installed. '
+            'Run the following commands:'
+            f'\n{cls._INDENT_PREFIX}  $ pip install boto3 botocore awscli'
+            f'\n{cls._INDENT_PREFIX}Credentials may also need to be set. '
+            f'{cls._STATIC_CREDENTIAL_HELP_STR}'
+        )
+
+        # Checks if the AWS CLI is installed properly
+        proc = subprocess.run(
+            'aws --version',
+            shell=True,
+            check=False,
+            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  # noqa: F401
+            import botocore  # noqa: F401
+        except ImportError:
+            return False, 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()  # noqa: F841
+        except exceptions.CloudUserIdentityError as e:
+            return False, str(e)
+
+        static_credential_exists = os.path.isfile(
+            os.path.expanduser('~/.aws/credentials')
+        )
+        hints = None
+        identity_type = cls._current_identity_type()
+        single_cloud_hint = (
+            ' It will work if you use AWS only, but will cause problems '
+            'if you want to use multiple clouds. To set up static credentials, '
+            'try: aws configure'
+        )
+        if identity_type == AWSIdentityType.SSO:
+            hints = 'AWS SSO is set.'
+            if static_credential_exists:
+                hints += (
+                    ' To ensure multiple clouds work correctly, please use SkyPilot '
+                    'with static credentials (e.g., ~/.aws/credentials) by unsetting '
+                    'the AWS_PROFILE environment variable.'
+                )
+            else:
+                hints += single_cloud_hint
+        elif identity_type == AWSIdentityType.IAM_ROLE:
+            # When using an IAM role, the credentials may not exist in the
+            # ~/.aws/credentials file. So we don't check for the existence of the
+            # file. This will happen when the user is on a VM (or
+            # jobs-controller) created by an SSO account, i.e. the VM will be
+            # assigned the IAM role: skypilot-v1.
+            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
+            # 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_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,
+                )
+
+        try:
+            s3 = boto3.client('s3')
+
+            suffix = uuid.uuid4().hex[:6]
+            test_bucket = f'konduktor-check-s3-{int(time.time())}-{suffix}'
+
+            try:
+                s3.create_bucket(Bucket=test_bucket)
+
+                time.sleep(1)
+
+                s3.get_bucket_location(Bucket=test_bucket)
+                s3.list_objects_v2(Bucket=test_bucket, MaxKeys=1)
+
+                # Object-related checks
+                s3.put_object(Bucket=test_bucket, Key='test.txt', Body='hello')
+                s3.get_object(Bucket=test_bucket, Key='test.txt')
+                s3.delete_object(Bucket=test_bucket, Key='test.txt')
+
+            finally:
+                # Always attempt to clean up, even if earlier checks failed
+                try:
+                    s3.delete_bucket(Bucket=test_bucket)
+                except Exception:
+                    raise RuntimeError(
+                        'AWS S3 delete bucket permission is missing. '
+                        'Please check your policy.\n'
+                    )
+
+        except Exception:
+            return False, (
+                'One or more AWS S3 permissions are missing. '
+                'Please check your policy.\n'
+            )
+
+        logger.info(
+            f'AWS credentials are valid '
+            f'for the current identity {logging.CHECK_MARK_EMOJI}'
+        )
+        logger.info('Creating k8s secret with AWS credentials...')
+        set_ok, result = cls.set_secret_credentials()
+        if not set_ok:
+            logger.error(f'Failed to create k8s secret with AWS credentials: {result}')
+            return False, result
+        return True, hints
+
+    @classmethod
+    def _current_identity_type(cls) -> Optional[AWSIdentityType]:
+        stdout = cls._aws_configure_list()
+        if stdout is None:
+            return None
+        output = stdout.decode()
+
+        # We determine the identity type by looking at the output of
+        # `aws configure list`. The output looks like:
+        #   Name                   Value         Type    Location
+        #   ----                   -----         ----    --------
+        #   profile                <not set>     None    None
+        #   access_key     *       <not set>     sso     None
+        #   secret_key     *       <not set>     sso     None
+        #   region                 <not set>     None    None
+        # We try to determine the identity type by looking for the
+        # string "sso"/"iam-role" in the output, i.e. the "Type" column.
+
+        def _is_access_key_of_type(type_str: str) -> bool:
+            # The dot (.) does not match line separators.
+            results = re.findall(rf'access_key.*{type_str}', output)
+            if len(results) > 1:
+                raise RuntimeError(f'Unexpected `aws configure list` output:\n{output}')
+            return len(results) == 1
+
+        for identity_type in AWSIdentityType:
+            if _is_access_key_of_type(identity_type.value):
+                return identity_type
+        return AWSIdentityType.SHARED_CREDENTIALS_FILE
+
+    # IN GCS.PY
+    @classmethod
+    def set_secret_credentials(cls) -> Tuple[bool, Optional[str]]:
+        """
+        Set the k8s secret storing the AWS credentials
+        """
+        context = kubernetes_utils.get_current_kube_config_context_name()
+        namespace = kubernetes_utils.get_kube_config_context_namespace()
+
+        # Check if credentials are provided via environment
+        access_key = os.environ.get('AWS_ACCESS_KEY_ID')
+        secret_key = os.environ.get('AWS_SECRET_ACCESS_KEY')
+        region = os.environ.get('AWS_DEFAULT_REGION', 'us-east-1')
+
+        if access_key and secret_key:
+            logger.info('Using AWS credentials from env')
+            credentials_dir = tempfile.mkdtemp()
+            credentials_path = os.path.join(credentials_dir, 'credentials')
+            config_path = os.path.join(credentials_dir, 'config')
+
+            with open(credentials_path, 'w') as f:
+                f.write(f"""[default]
+            aws_access_key_id = {access_key}
+            aws_secret_access_key = {secret_key}
+            """)
+
+            with open(config_path, 'w') as f:
+                f.write(f"""[default]
+            region = {region}
+            """)
+        else:
+            logger.info('Using AWS credentils from ~/.aws/')
+            credentials_dir = DEFAULT_AWS_CREDENTIALS_DIR
+
+        credentials_files = [
+            os.path.expanduser(os.path.join(credentials_dir, f))
+            for f in _CREDENTIAL_FILES
+        ]
+        ok, result = kubernetes_utils.set_secret(
+            secret_name=cls._AWS_SECRET_NAME,
+            namespace=namespace,
+            context=context,
+            secret_key=cls._AWS_CREDENTIALS_KEY,
+            secret_value=base64_utils.zip_base64encode(credentials_files),
+        )
+        if not ok:
+            logger.error(f'Failed to set AWS credentials in k8s secret: \n{result}')
+            return False, result
+        else:
+            logger.info(
+                f'AWS credentials set in k8s secret: {cls._AWS_SECRET_NAME} '
+                f'in namespace {namespace} in context {context} '
+                f'{logging.CHECK_MARK_EMOJI}'
+            )
+
+        try:
+            identity = boto3.client('sts').get_caller_identity()
+            logger.info(
+                f"Using AWS credentials for ARN: {identity['Arn']} "
+                f"(UserId: {identity['UserId']}, Account: {identity['Account']})"
+            )
+        except Exception as e:
+            logger.warning(f'Could not fetch caller identity: {e}')
+
+        return True, None
+
+    # IN GCS.PY
+    @classmethod
+    def get_k8s_credential_name(cls) -> str:
+        return cls._AWS_SECRET_NAME
+
+
+class S3CloudStorage(storage_utils.CloudStorage):
+    """S3 Storage."""
+
+    # List of commands to install AWS CLI
+    _GET_AWSCLI = [
+        'command -v aws >/dev/null 2>&1 || ('
+        'apt-get update && apt-get install -y curl unzip && '
+        'curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip" && '  # noqa: E501
+        'unzip awscliv2.zip && '
+        './aws/install -i ~/aws-cli -b ~/bin && '
+        'export PATH=$HOME/bin:$PATH && '
+        'rm -rf aws awscliv2.zip'
+        ') && export PATH=$HOME/bin:$PATH'
+    ]
+
+    _STORE: typing.Type[storage_utils.AbstractStore] = S3Store
+
+    # IN GCS.PY
+    def is_directory(self, url: str) -> bool:
+        """Returns whether S3 'url' is a directory.
+
+        In cloud object stores, a "directory" refers to a regular object whose
+        name is a prefix of other objects.
+        """
+        s3 = aws.resource('s3')
+        bucket_name, path = data_utils.split_s3_path(url)
+        bucket = s3.Bucket(bucket_name)
+
+        num_objects = 0
+        for obj in bucket.objects.filter(Prefix=path):
+            num_objects += 1
+            if obj.key == path:
+                return False
+            # If there are more than 1 object in filter, then it is a directory
+            if num_objects == 3:
+                return True
+
+        # A directory with few or no items
+        return True
+
+    # IN GCS.PY
+    def make_sync_dir_command(self, source: str, destination: str) -> str:
+        """Downloads using AWS CLI."""
+        # AWS Sync by default uses 10 threads to upload files to the bucket.
+        # To increase parallelism, modify max_concurrent_requests in your
+        # aws config file (Default path: ~/.aws/config).
+        all_commands = list(self._GET_AWSCLI)
+
+        all_commands.append(f'aws s3 sync --no-follow-symlinks {source} {destination}')
+        return ' && '.join(all_commands)
+
+    # IN GCS.PY
+    def make_sync_file_command(self, source: str, destination: str) -> str:
+        all_commands = list(self._GET_AWSCLI)
+
+        all_commands.append(f'aws s3 cp {source} {destination}')
+        return ' && '.join(all_commands)