Repository: skypilot-org/skypilot
Pull request: [Storage] Error Handle Attempts to Mount Non-sky Managed Bucket without Source
<!-- Describe the changes in this PR -->
This closes #2779, #2950

`Storage` is designed in a way to not allow externally created cloud storages to be mounted without specifying the `source` field. The only way to mount non-sky managed bucket is to specify the bucket's URI as the `source` field. This fix raises `StorageSpecError` when a non-sky managed bucket is attempted to be mounted just with `name` field and without specifying the `source` field.

<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [x] Code formatting: `bash format.sh`
- [x] All smoke tests: `pytest tests/test_smoke.py::TestStorageWithCredentials` 
- [x] Manually test: Create storage externally and `MOUNT` with `name` specified and `source` unspecified. Should raise error
- [x] Manually test: Create storage externally and `MOUNT` with `source` specified and `name` unspecified. Should go through
- [x] `pytest tests/test_smoke.py::test_gcp_storage_mounts_with_stop --gcp`
- [x] `pytest tests/test_smoke.py::test_aws_storage_mounts_with_stop --aws`



Problem statement:
[Storage] sky start failing to re-mount on externally created Storage
User reported to @Michaelvll :

> Got this error while restarting an instance:
> `sky.exceptions.StorageExternalDeletionError: The bucket, 'bucket-name', could not be mounted on cluster 'cluster-name'. Please verify that the bucket exists. The cluster started successfully without mounting the bucket.`

User also confirmed that the storage exists in the console. Did some investigation and found that this is an edge case for cloud storages that are not sky managed (created externally). 

repro:
1. `sky launch` with a task yaml specifying `MOUNT` mode on externally created cloud storage(create storage manually from cloud service provider console)
2. `sky stop` the cluster
3. `sky start` the cluster

This edge case should be fixed so the externally created storages can be re-mounted when `sky start` is ran.

Update: This edge case is rather rooting from an undesired behavior of current `Storage` class implementation. To raise the above error by attempting to re-mount an externally created bucket with `sky start`, user specified a `name` of the externally created bucket under `file_mount`'s section when first launched the instance. However, we only want to allow mounting an externally created bucket by user specifying the bucket's url to the `source` field. The fix to this issue should be to raise an error when a `name` of externally created bucket is specified with `MOUNT` mode.

Hints from issue discussion:

Commit message: add new smoke test

Patch to review:
diff --git a/sky/data/storage.py b/sky/data/storage.py
index d9f8156fa89..89dc67ade3f 100644
--- a/sky/data/storage.py
+++ b/sky/data/storage.py
@@ -819,6 +819,10 @@ def add_store(self, store_type: Union[str, StoreType]) -> AbstractStore:
             logger.error(f'Could not initialize {store_type} store with '
                          f'name {self.name}. General initialization error.')
             raise
+        except exceptions.StorageSpecError:
+            logger.error(f'Could not mount externally created {store_type}'
+                         f'store with name {self.name!r}.')
+            raise
 
         # Add store to storage
         self._add_store(store)
@@ -1233,6 +1237,8 @@ def _get_bucket(self) -> Tuple[StorageHandle, bool]:
           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
@@ -1248,6 +1254,26 @@ def _get_bucket(self) -> Tuple[StorageHandle, bool]:
             # bucket or if it is a user's bucket that is publicly
             # accessible.
             self.client.head_bucket(Bucket=self.name)
+            # Check if 'source' is None, indicating Storage is in MOUNT mode.
+            # Note: In COPY mode, a 'source' being None is already handled
+            # as an error in _validate_storage_spec.
+            if self.source is None:
+                # Retrieve a handle associated with the storage name.
+                # This handle links to sky managed storage if it exists.
+                handle = global_user_state.get_handle_from_storage_name(
+                    self.name)
+                # If handle is None, it implies the bucket is created
+                # externally and not managed by Skypilot. For mounting such
+                # externally created buckets, users must provide the
+                # bucket's URL as 'source'.
+                if handle is None:
+                    with ux_utils.print_exception_no_traceback():
+                        raise exceptions.StorageSpecError(
+                            'Attempted to mount a non-sky managed bucket '
+                            f'{self.name!r} from S3 without specifying a '
+                            'storage source. To mount an externally created '
+                            'bucket, please specify the bucket URL as '
+                            'the storage source.')
             return bucket, False
         except aws.botocore_exceptions().ClientError as e:
             error_code = e.response['Error']['Code']
@@ -1678,6 +1704,8 @@ def _get_bucket(self) -> Tuple[StorageHandle, bool]:
           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
@@ -1686,6 +1714,26 @@ def _get_bucket(self) -> Tuple[StorageHandle, bool]:
         """
         try:
             bucket = self.client.get_bucket(self.name)
+            # Check if 'source' is None, indicating Storage is in MOUNT mode.
+            # Note: In COPY mode, a 'source' being None is already handled
+            # as an error in _validate_storage_spec.
+            if self.source is None:
+                # Retrieve a handle associated with the storage name.
+                # This handle links to sky managed storage if it exists.
+                handle = global_user_state.get_handle_from_storage_name(
+                    self.name)
+                # If handle is None, it implies the bucket is created
+                # externally and not managed by Skypilot. For mounting such
+                # externally created buckets, users must provide the
+                # bucket's URL as 'source'.
+                if handle is None:
+                    with ux_utils.print_exception_no_traceback():
+                        raise exceptions.StorageSpecError(
+                            'Attempted to mount a non-sky managed bucket '
+                            f'{self.name!r} from GCS without specifying a '
+                            'storage source. To mount an externally created '
+                            'bucket, please specify the bucket URL as '
+                            'the storage source.')
             return bucket, False
         except gcp.not_found_exception() as e:
             if isinstance(self.source, str) and self.source.startswith('gs://'):
@@ -2031,6 +2079,8 @@ def _get_bucket(self) -> Tuple[StorageHandle, bool]:
           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
@@ -2046,6 +2096,26 @@ def _get_bucket(self) -> Tuple[StorageHandle, bool]:
             # bucket or if it is a user's bucket that is publicly
             # accessible.
             self.client.head_bucket(Bucket=self.name)
+            # Check if 'source' is None, indicating Storage is in MOUNT mode.
+            # Note: In COPY mode, a 'source' being None is already handled
+            # as an error in _validate_storage_spec.
+            if self.source is None:
+                # Retrieve a handle associated with the storage name.
+                # This handle links to sky managed storage if it exists.
+                handle = global_user_state.get_handle_from_storage_name(
+                    self.name)
+                # If handle is None, it implies the bucket is created
+                # externally and not managed by Skypilot. For mounting such
+                # externally created buckets, users must provide the
+                # bucket's URL as 'source'.
+                if handle is None:
+                    with ux_utils.print_exception_no_traceback():
+                        raise exceptions.StorageSpecError(
+                            'Attempted to mount a non-sky managed bucket '
+                            f'{self.name!r} from R2 without specifying a '
+                            'storage source. To mount an externally created '
+                            'bucket, please specify the bucket URL as '
+                            'the storage source.')
             return bucket, False
         except aws.botocore_exceptions().ClientError as e:
             error_code = e.response['Error']['Code']
@@ -2446,6 +2516,8 @@ def _get_bucket(self) -> Tuple[StorageHandle, bool]:
           bool: indicates whether a new bucket was created.
 
         Raises:
+            StorageSpecError: If externally created bucket is attempted to be
+                mounted without specifying storage source.
             StorageBucketCreateError: If bucket creation fails.
             StorageBucketGetError: If fetching a bucket fails
             StorageExternalDeletionError: If externally deleted storage is
@@ -2505,7 +2577,28 @@ def _get_bucket(self) -> Tuple[StorageHandle, bool]:
                 f'{self.name}')
         else:
             # bucket exists
-            return self.s3_resource.Bucket(self.name), False
+            bucket = self.s3_resource.Bucket(self.name)
+            # Check if 'source' is None, indicating Storage is in MOUNT mode.
+            # Note: In COPY mode, a 'source' being None is already handled
+            # as an error in _validate_storage_spec.
+            if self.source is None:
+                # Retrieve a handle associated with the storage name.
+                # This handle links to sky managed storage if it exists.
+                handle = global_user_state.get_handle_from_storage_name(
+                    self.name)
+                # If handle is None, it implies the bucket is created
+                # externally and not managed by Skypilot. For mounting such
+                # externally created buckets, users must provide the
+                # bucket's URL as 'source'.
+                if handle is None:
+                    with ux_utils.print_exception_no_traceback():
+                        raise exceptions.StorageSpecError(
+                            'Attempted to mount a non-sky managed bucket '
+                            f'{self.name!r} from IBM COS without specifying a '
+                            'storage source. To mount an externally created '
+                            'bucket, please specify the bucket URL as '
+                            'the storage source.')
+            return bucket, False
 
     def _download_file(self, remote_path: str, local_path: str) -> None:
         """Downloads file from remote to local on s3 bucket
diff --git a/tests/test_smoke.py b/tests/test_smoke.py
index 62490e44fd7..aa133c5ce0e 100644
--- a/tests/test_smoke.py
+++ b/tests/test_smoke.py
@@ -3390,28 +3390,30 @@ def tmp_gitignore_storage_obj(self, tmp_bucket_name, gitignore_structure):
     @pytest.fixture
     def tmp_awscli_bucket(self, tmp_bucket_name):
         # Creates a temporary bucket using awscli
-        subprocess.check_call(['aws', 's3', 'mb', f's3://{tmp_bucket_name}'])
-        yield tmp_bucket_name
-        subprocess.check_call(
-            ['aws', 's3', 'rb', f's3://{tmp_bucket_name}', '--force'])
+        bucket_uri = f's3://{tmp_bucket_name}'
+        subprocess.check_call(['aws', 's3', 'mb', bucket_uri])
+        yield tmp_bucket_name, bucket_uri
+        subprocess.check_call(['aws', 's3', 'rb', bucket_uri, '--force'])
 
     @pytest.fixture
     def tmp_gsutil_bucket(self, tmp_bucket_name):
         # Creates a temporary bucket using gsutil
-        subprocess.check_call(['gsutil', 'mb', f'gs://{tmp_bucket_name}'])
-        yield tmp_bucket_name
-        subprocess.check_call(['gsutil', 'rm', '-r', f'gs://{tmp_bucket_name}'])
+        bucket_uri = f'gs://{tmp_bucket_name}'
+        subprocess.check_call(['gsutil', 'mb', bucket_uri])
+        yield tmp_bucket_name, bucket_uri
+        subprocess.check_call(['gsutil', 'rm', '-r', bucket_uri])
 
     @pytest.fixture
     def tmp_awscli_bucket_r2(self, tmp_bucket_name):
         # Creates a temporary bucket using awscli
         endpoint_url = cloudflare.create_endpoint()
+        bucket_uri = f's3://{tmp_bucket_name}'
         subprocess.check_call(
-            f'AWS_SHARED_CREDENTIALS_FILE={cloudflare.R2_CREDENTIALS_PATH} aws s3 mb s3://{tmp_bucket_name} --endpoint {endpoint_url} --profile=r2',
+            f'AWS_SHARED_CREDENTIALS_FILE={cloudflare.R2_CREDENTIALS_PATH} aws s3 mb {bucket_uri} --endpoint {endpoint_url} --profile=r2',
             shell=True)
-        yield tmp_bucket_name
+        yield tmp_bucket_name, bucket_uri
         subprocess.check_call(
-            f'AWS_SHARED_CREDENTIALS_FILE={cloudflare.R2_CREDENTIALS_PATH} aws s3 rb s3://{tmp_bucket_name} --force --endpoint {endpoint_url} --profile=r2',
+            f'AWS_SHARED_CREDENTIALS_FILE={cloudflare.R2_CREDENTIALS_PATH} aws s3 rb {bucket_uri} --force --endpoint {endpoint_url} --profile=r2',
             shell=True)
 
     @pytest.fixture
@@ -3644,7 +3646,7 @@ def test_upload_to_existing_bucket(self, ext_bucket_fixture, request,
                                        tmp_source, store_type):
         # Tries uploading existing files to newly created bucket (outside of
         # sky) and verifies that files are written.
-        bucket_name = request.getfixturevalue(ext_bucket_fixture)
+        bucket_name, _ = request.getfixturevalue(ext_bucket_fixture)
         storage_obj = storage_lib.Storage(name=bucket_name, source=tmp_source)
         storage_obj.add_store(store_type)
 
@@ -3760,6 +3762,43 @@ def test_excluded_file_cloud_storage_upload_copy(self, gitignore_structure,
         assert '4' in cnt_output.decode('utf-8'), \
                'Some items listed in .gitignore and .git/info/exclude are not excluded.'
 
+    @pytest.mark.parametrize('ext_bucket_fixture, store_type',
+                             [('tmp_awscli_bucket', storage_lib.StoreType.S3),
+                              ('tmp_gsutil_bucket', storage_lib.StoreType.GCS),
+                              pytest.param('tmp_awscli_bucket_r2',
+                                           storage_lib.StoreType.R2,
+                                           marks=pytest.mark.cloudflare)])
+    def test_externally_created_bucket_mount_without_source(
+            self, ext_bucket_fixture, request, store_type):
+        # Non-sky managed buckets(buckets created outside of Skypilot CLI)
+        # are allowed to be MOUNTed by specifying the URI of the bucket to
+        # source field only. When it is attempted by specifying the name of
+        # the bucket only, it should error out.
+        # TODO(doyoung): Add test for IBM COS. Currently, this is blocked
+        # as rclone used to interact with IBM COS does not support feature to
+        # create a bucket, and the ibmcloud CLI is not supported in Skypilot.
+        # Either of the feature is necessary to simulate an external bucket
+        # createion for IBM COS.
+        # https://github.com/skypilot-org/skypilot/pull/1966/files#r1253439837
+
+        ext_bucket_name, ext_bucket_uri = request.getfixturevalue(
+            ext_bucket_fixture)
+        # invalid spec
+        with pytest.raises(sky.exceptions.StorageSpecError) as e:
+            storage_obj = storage_lib.Storage(
+                name=ext_bucket_name, mode=storage_lib.StorageMode.MOUNT)
+            storage_obj.add_store(store_type)
+
+        assert 'Attempted to mount a non-sky managed bucket' in str(e)
+
+        # valid spec
+        storage_obj = storage_lib.Storage(source=ext_bucket_uri,
+                                          mode=storage_lib.StorageMode.MOUNT)
+        handle = global_user_state.get_handle_from_storage_name(
+            storage_obj.name)
+        if handle:
+            storage_obj.delete()
+
 
 # ---------- Testing YAML Specs ----------
 # Our sky storage requires credentials to check the bucket existance when


Code review comments:
diff_hunk: @@ -1248,6 +1254,26 @@ def _get_bucket(self) -> Tuple[StorageHandle, bool]:
             # bucket or if it is a user's bucket that is publicly
             # accessible.
             self.client.head_bucket(Bucket=self.name)
+            # Check if 'source' is None, indicating Storage is in MOUNT mode.
+            # Note: In COPY mode, a 'source' being None is already handled
+            # as an error in _validate_storage_spec.
+            if self.source is None:
+                # Retrieve a handle associated with the storage name.
+                # This handle links to sky managed storage if it exists.
+                handle = global_user_state.get_handle_from_storage_name(
+                    self.name)
+                # If handle is None, it implies the bucket is created
+                # externally and not managed by Skypilot. For mounting such
+                # externally created buckets, users must provide the
+                # bucket's URL as 'source'.
+                if handle is None:
+                    with ux_utils.print_exception_no_traceback():
+                        raise exceptions.StorageSpecError(
+                            'Attempted to mount a non-sky managed bucket '
+                            f'{self.name!r} from S3 without specifying a '
+                            'storage source. To mount an externally created '
+                            'bucket, please specify the bucket URL as '
+                            'the storage source.')
path: sky/data/storage.py
text: @user1:
Why do we duplicate the same codes in three places? Should we do some refactoring?

@author:
Refactored at [a602af3](https://github.com/skypilot-org/skypilot/pull/2804/commits/a602af39e6c45f49d90371c11fc24e8e787afaeb)