[Storage] Avoid opt-in regions for S3 (#4239)

* S3 fix + timeout * S3 fix + timeout * lint
skypilot-org · Nov 1, 2024 · bf17e87 · bf17e87
1 parent e8d0130
commit bf17e87
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 3 deletions.
diff --git a/sky/data/storage.py b/sky/data/storage.py
@@ -1082,16 +1082,31 @@ class S3Store(AbstractStore):
     for S3 buckets.
     """
 
+    _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'
+    ]
 
     def __init__(self,
                  name: str,
                  source: str,
-                 region: Optional[str] = 'us-east-2',
+                 region: Optional[str] = _DEFAULT_REGION,
                  is_sky_managed: Optional[bool] = None,
                  sync_on_reconstruction: bool = True):
         self.client: 'boto3.client.Client'
         self.bucket: 'StorageHandle'
+        # TODO(romilb): This is purely a stopgap fix for
+        #  https://github.com/skypilot-org/skypilot/issues/3405
+        # We should eventually make all opt-in regions also work for S3 by
+        # passing the right endpoint flags.
+        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)
 
@@ -1424,7 +1439,7 @@ def mount_command(self, mount_path: str) -> str:
 
     def _create_s3_bucket(self,
                           bucket_name: str,
-                          region='us-east-2') -> StorageHandle:
+                          region=_DEFAULT_REGION) -> StorageHandle:
         """Creates S3 bucket with specific name in specific region
 
         Args:

diff --git a/tests/test_smoke.py b/tests/test_smoke.py
@@ -3102,7 +3102,7 @@ def test_managed_jobs_cancellation_aws(aws_config_region):
             # Test cancellation during spot cluster being launched.
             f'sky jobs launch --cloud aws --region {region} -n {name} --use-spot "sleep 1000"  -y -d',
             'sleep 60',
-            f'{_GET_JOB_QUEUE} | grep {name} | head -n1 | grep "STARTING"',
+            f'{_GET_JOB_QUEUE} | grep {name} | head -n1 | grep "STARTING\|RUNNING"',
             f'sky jobs cancel -y -n {name}',
             'sleep 5',
             f'{_GET_JOB_QUEUE} | grep {name} | head -n1 | grep "CANCELLING\|CANCELLED"',