Fix issue 3744: Regions are not respected for buckets created with sky launch

sky/data/storage.py

@@ -294,6 +294,12 @@ def initialize(self):
         """
         pass
 
+    def sync_bucket(self, region: Optional[str] = None):
+        """Fetches bucket if exists, or creates it if
+        it does not.
+        """
+        pass
+
     def _validate(self) -> None:
         """Runs validation checks on class args"""
         pass
@@ -808,14 +814,15 @@ def _add_store_from_metadata(
                                  'to be reconstructed while the corresponding '
                                  'storage account '
                                  f'{s_metadata.storage_account_name!r} does '
-                                 'not exist.')
+                                 'not exist')
                 else:
                     logger.debug(f'Storage object {self.name!r} was attempted '
                                  'to be reconstructed while the corresponding '
                                  'bucket was externally deleted.')
                 continue
 
             self._add_store(store, is_reconstructed=True)
+            # self._sync_store(store)
 
     @classmethod
     def from_metadata(cls, metadata: StorageMetadata,
@@ -917,14 +924,13 @@ def add_store(self,
                          f'store with name {self.name!r}.')
             raise
 
-        # Add store to storage
-        self._add_store(store)
+        return store
 
-        # Upload source to store
+    def initialize_and_sync_store(self, store: AbstractStore) -> None:
+        """Public method to initialize and sync the given store."""
+        self._add_store(store)
         self._sync_store(store)
 
-        return store
-
     def _add_store(self, store: AbstractStore, is_reconstructed: bool = False):
         # Adds a store object to the storage
         store_type = StoreType.from_store(store)
@@ -1020,7 +1026,7 @@ def from_yaml_config(cls, config: Dict[str, Any]) -> 'Storage':
 
         name = config.pop('name', None)
         source = config.pop('source', None)
-        store = config.pop('store', None)
+        store = config.pop('store', None)  # where do we store this?!!?!?!?
         mode_str = config.pop('mode', None)
         force_delete = config.pop('_force_delete', None)
         if force_delete is None:
@@ -1043,8 +1049,8 @@ def from_yaml_config(cls, config: Dict[str, Any]) -> 'Storage':
                           source=source,
                           persistent=persistent,
                           mode=mode)
-        if store is not None:
-            storage_obj.add_store(StoreType(store.upper()))
+        if store:
+            storage_obj.add_store(store_type=store.upper())
 
         # Add force deletion flag
         storage_obj.force_delete = force_delete
@@ -1138,11 +1144,11 @@ def _validate(self):
         if not _is_storage_cloud_enabled(str(clouds.AWS())):
             with ux_utils.print_exception_no_traceback():
                 raise exceptions.ResourcesUnavailableError(
-                    'Storage \'store: s3\' specified, but ' \
-                    'AWS access is disabled. To fix, enable '\
-                    'AWS by running `sky check`. More info: '\
-                    'https://skypilot.readthedocs.io/en/latest/getting-started/installation.html.' # pylint: disable=line-too-long
-                    )
+                    'Storage \'store: s3\' specified, but '
+                    'AWS access is disabled. To fix, enable '
+                    'AWS by running `sky check`. More info: '
+                    'https://skypilot.readthedocs.io/en/latest/getting-started/installation.html.'  # pylint: disable=line-too-long
+                )
 
     @classmethod
     def validate_name(cls, name: str) -> str:
@@ -1215,13 +1221,21 @@ def initialize(self):
           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
+        self.bucket = None
+
+    def sync_bucket(self, region: Optional[str] = None):
+        if region:
+            self.region = region
+        if self.bucket is None:
+            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
+
+        return self.bucket
 
     def upload(self):
         """Uploads source to store bucket.
@@ -1256,7 +1270,7 @@ def delete(self) -> None:
             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.'
+                f'externally. Removing from local state.'
         logger.info(f'{colorama.Fore.GREEN}{msg_str}'
                     f'{colorama.Style.RESET_ALL}')
 
@@ -1629,13 +1643,20 @@ def initialize(self):
           StorageInitError: If general initialization fails.
         """
         self.client = gcp.storage_client()
-        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
+        self.bucket = None
+
+    def sync_bucket(self, region: Optional[str] = None):
+        if region:
+            self.region = region
+        if self.bucket is None:
+            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
+        return self.bucket
 
     def upload(self):
         """Uploads source to store bucket.
@@ -1672,7 +1693,7 @@ def delete(self) -> None:
             msg_str = f'Deleted GCS bucket {self.name}.'
         else:
             msg_str = f'GCS bucket {self.name} may have been deleted ' \
-                      f'externally. Removing from local state.'
+                f'externally. Removing from local state.'
         logger.info(f'{colorama.Fore.GREEN}{msg_str}'
                     f'{colorama.Style.RESET_ALL}')
 
@@ -2724,11 +2745,11 @@ def _validate(self):
         if not _is_storage_cloud_enabled(cloudflare.NAME):
             with ux_utils.print_exception_no_traceback():
                 raise exceptions.ResourcesUnavailableError(
-                    'Storage \'store: r2\' specified, but ' \
-                    'Cloudflare R2 access is disabled. To fix, '\
-                    'enable Cloudflare R2 by running `sky check`. '\
+                    'Storage \'store: r2\' specified, but '
+                    'Cloudflare R2 access is disabled. To fix, '
+                    'enable Cloudflare R2 by running `sky check`. '
                     'More info: https://skypilot.readthedocs.io/en/latest/getting-started/installation.html.'  # pylint: disable=line-too-long
-                    )
+                )
 
     def initialize(self):
         """Initializes the R2 store object on the cloud.
@@ -2742,13 +2763,19 @@ def initialize(self):
           StorageInitError: If general initialization fails.
         """
         self.client = data_utils.create_r2_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
+        self.bucket = None
+
+    def sync_bucket(self, region: Optional[str] = None):
+        if region:
+            self.region = region
+        if self.bucket is None:
+            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
 
     def upload(self):
         """Uploads source to store bucket.
@@ -2783,7 +2810,7 @@ def delete(self) -> None:
             msg_str = f'Deleted R2 bucket {self.name}.'
         else:
             msg_str = f'R2 bucket {self.name} may have been deleted ' \
-                      f'externally. Removing from local state.'
+                f'externally. Removing from local state.'
         logger.info(f'{colorama.Fore.GREEN}{msg_str}'
                     f'{colorama.Style.RESET_ALL}')
 
@@ -3064,8 +3091,8 @@ def __init__(self,
         super().__init__(name, source, region, is_sky_managed,
                          sync_on_reconstruction)
         self.bucket_rclone_profile = \
-          Rclone.generate_rclone_bucket_profile_name(
-            self.name, Rclone.RcloneClouds.IBM)
+            Rclone.generate_rclone_bucket_profile_name(
+                self.name, Rclone.RcloneClouds.IBM)
 
     def _validate(self):
         if self.source is not None and isinstance(self.source, str):
@@ -3166,13 +3193,19 @@ def initialize(self):
         """
         self.client = ibm.get_cos_client(self.region)
         self.s3_resource = ibm.get_cos_resource(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
+        self.bucket = None
+
+    def sync_bucket(self, region: Optional[str] = None):
+        if region:
+            self.region = region
+        if self.bucket is None:
+            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
 
     def upload(self):
         """Uploads files from local machine to bucket.
@@ -3340,7 +3373,7 @@ def _get_bucket(self) -> Tuple[StorageHandle, bool]:
 
         # bucket's region doesn't match specified region in URI
         if bucket_region and uri_region and uri_region != bucket_region\
-              and self.sync_on_reconstruction:
+                and self.sync_on_reconstruction:
             with ux_utils.print_exception_no_traceback():
                 raise exceptions.StorageBucketGetError(
                     f'Bucket {self.name} exists in '
@@ -3432,7 +3465,8 @@ def _create_cos_bucket(self,
                         f'with storage class smart tier')
             self.bucket = self.s3_resource.Bucket(bucket_name)
 
-        except ibm.ibm_botocore.exceptions.ClientError as e:  # type: ignore[union-attr]  # pylint: disable=line-too-long
+        # type: ignore[union-attr]  # pylint: disable=line-too-long
+        except ibm.ibm_botocore.exceptions.ClientError as e:
             with ux_utils.print_exception_no_traceback():
                 raise exceptions.StorageBucketCreateError(
                     f'Failed to create bucket: '

sky/task.py

@@ -436,6 +436,7 @@ def from_yaml_config(
             assert mount_path, 'Storage mount path cannot be empty.'
             try:
                 storage_obj = storage_lib.Storage.from_yaml_config(storage[1])
+
             except exceptions.StorageSourceError as e:
                 # Patch the error message to include the mount path, if included
                 e.args = (e.args[0].replace('<destination_path>',
@@ -959,9 +960,21 @@ def sync_storage_mounts(self) -> None:
             if len(storage.stores) == 0:
                 store_type, store_region = self._get_preferred_store()
                 self.storage_plans[storage] = store_type
-                storage.add_store(store_type, store_region)
+                new_store = storage.add_store(store_type, store_region)
+                new_store.sync_bucket(store_region)
+                storage.initialize_and_sync_store(new_store)
             else:
                 # We will download the first store that is added to remote.
+                for store_type, store in storage.stores.items():
+                    if store_type == storage_lib.StoreType.AZURE:
+                        continue
+
+                    s_type, s_region = self._get_preferred_store()
+                    if store_type == s_type:
+                        store.sync_bucket(s_region)
+                    else:
+                        store.sync_bucket()
+                    storage.initialize_and_sync_store(new_store)
                 self.storage_plans[storage] = list(storage.stores.keys())[0]
 
         storage_mounts = self.storage_mounts