Repository: skypilot-org/skypilot
Pull request: Fix show gpus when user does not have AWS credentials
<!-- Describe the changes in this PR -->

Fixes #1710 

<!-- 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):

Without credential
- [x] `sky show-gpus` (1.4 seconds; origin: 1.4 seconds)
- [x] `sky show-gpus --all`
- [x] `sky show-gpus --cloud aws --region us-east-1` (1.4 seconds; origin: 1.4 seconds)

With credential
- [x] `sky show-gpus` (1.4 seconds; origin: 1.4 seconds)
- [x] `sky show-gpus --all`
- [x] `sky show-gpus --cloud aws --region us-east-1` (1.2 seconds; origin: 1.2 seconds)


Tests:
- [x] Launch an AWS cluster; remove the AWS credentials; `sky status` and `sky status -r`
- [x] `pytest tests/test_smoke.py`
- [x] `pytest tests/test_smoke.py` with no AWS credentials
- [x] `pytest tests/test_smoke.py --aws`

Problem statement:
`sky show-gpus` should not assume AWS credentials to fetch AZ mappings
On a fresh docker install
```
pip install skypilot[lambda]
sky show-gpus
```
shows
```
...
I 02-21 20:24:39 aws_catalog.py:53] Fetching availability zones mapping for AWS...
--- Logging error ---
....
  File "/opt/conda/lib/python3.9/site-packages/sky/clouds/service_catalog/data_fetchers/fetch_aws.py", line 69, in get_enabled_regions
    user_cloud_regions = aws_client.describe_regions()['Regions']
  File "/opt/conda/lib/python3.9/site-packages/botocore/client.py", line 530, in _api_call
    return self._make_api_call(operation_name, kwargs)
  File "/opt/conda/lib/python3.9/site-packages/botocore/client.py", line 943, in _make_api_call
    http, parsed_response = self._make_request(
  File "/opt/conda/lib/python3.9/site-packages/botocore/client.py", line 966, in _make_request
    return self._endpoint.make_request(operation_model, request_dict)
  File "/opt/conda/lib/python3.9/site-packages/botocore/endpoint.py", line 119, in make_request
    return self._send_request(request_dict, operation_model)
  File "/opt/conda/lib/python3.9/site-packages/botocore/endpoint.py", line 198, in _send_request
    request = self.create_request(request_dict, operation_model)
  File "/opt/conda/lib/python3.9/site-packages/botocore/endpoint.py", line 134, in create_request
    self._event_emitter.emit(
  File "/opt/conda/lib/python3.9/site-packages/botocore/hooks.py", line 412, in emit
    return self._emitter.emit(aliased_event_name, **kwargs)
  File "/opt/conda/lib/python3.9/site-packages/botocore/hooks.py", line 256, in emit
    return self._emit(event_name, kwargs)
  File "/opt/conda/lib/python3.9/site-packages/botocore/hooks.py", line 239, in _emit
    response = handler(**kwargs)
  File "/opt/conda/lib/python3.9/site-packages/botocore/signers.py", line 105, in handler
    return self.sign(operation_name, request)
  File "/opt/conda/lib/python3.9/site-packages/botocore/signers.py", line 189, in sign
    auth.add_auth(request)
  File "/opt/conda/lib/python3.9/site-packages/botocore/auth.py", line 418, in add_auth
    raise NoCredentialsError()
botocore.exceptions.NoCredentialsError: Unable to locate credentials
```
We should perhaps check other call paths that will by default fetch AWS AZ mappings.

cc @WoosukKwon @Michaelvll

Hints from issue discussion:

Commit message: fix test

Patch to review:
diff --git a/sky/clouds/service_catalog/aws_catalog.py b/sky/clouds/service_catalog/aws_catalog.py
index dfcba50848a..a6458a0a1e8 100644
--- a/sky/clouds/service_catalog/aws_catalog.py
+++ b/sky/clouds/service_catalog/aws_catalog.py
@@ -5,6 +5,7 @@
 """
 import colorama
 import os
+import threading
 import typing
 from typing import Dict, List, Optional, Tuple
 
@@ -29,13 +30,17 @@
 # skypilot-catalog/.github/workflows/update-aws-catalog.yml
 _PULL_FREQUENCY_HOURS = 7
 
+# Flag to indicate whether the availability zone mapping has been applied.
+_az_mapping_applied = False
+_apply_az_mapping_lock = threading.Lock()
+
 _df = common.read_catalog('aws/vms.csv',
                           pull_frequency_hours=_PULL_FREQUENCY_HOURS)
 _image_df = common.read_catalog('aws/images.csv',
                                 pull_frequency_hours=_PULL_FREQUENCY_HOURS)
 
 
-def _apply_az_mapping(df: 'pd.DataFrame') -> 'pd.DataFrame':
+def _apply_az_mapping():
     """Maps zone IDs (use1-az1) to zone names (us-east-1x).
 
     Such mappings are account-specific and determined by AWS.
@@ -44,6 +49,10 @@ def _apply_az_mapping(df: 'pd.DataFrame') -> 'pd.DataFrame':
         A dataframe with column 'AvailabilityZone' that's correctly replaced
         with the zone name (e.g. us-east-1a).
     """
+    global _az_mapping_applied
+    if _az_mapping_applied:
+        return
+    global _df
     az_mapping_path = common.get_catalog_path('aws/az_mappings.csv')
     if not os.path.exists(az_mapping_path):
         # Fetch az mapping from AWS.
@@ -61,25 +70,52 @@ def _apply_az_mapping(df: 'pd.DataFrame') -> 'pd.DataFrame':
     # Use inner join to drop rows with unknown AZ IDs, which are likely
     # because the user does not have access to that Region. Otherwise,
     # there will be rows with NaN in the AvailabilityZone column.
-    df = df.merge(az_mappings, on=['AvailabilityZone'], how='inner')
-    df = df.drop(columns=['AvailabilityZone']).rename(
+    _df = _df.merge(az_mappings, on=['AvailabilityZone'], how='inner')
+    _df = _df.drop(columns=['AvailabilityZone']).rename(
         columns={'AvailabilityZoneName': 'AvailabilityZone'})
-    return df
+    _az_mapping_applied = True
+
+
+def _apply_az_mapping_decorator(func):
+    """Decorator to apply availability zone mapping to the dataframe."""
 
+    def wrapper(*args, **kwargs):
+        with _apply_az_mapping_lock:
+            _apply_az_mapping()
+        return func(*args, **kwargs)
 
-_df = _apply_az_mapping(_df)
+    return wrapper
 
 
+@_apply_az_mapping_decorator
 def instance_type_exists(instance_type: str) -> bool:
     return common.instance_type_exists_impl(_df, instance_type)
 
 
+# Helper functions for `validate_region_zone`.
+def _validate_region(
+        region: Optional[str]) -> Tuple[Optional[str], Optional[str]]:
+    # When zone is not specified, we don't need to apply the AZ mapping.
+    # This is useful for `sky show-gpus` where user may not have access to
+    # AWS.
+    return common.validate_region_zone_impl('aws', _df, region, None)
+
+
+@_apply_az_mapping_decorator
+def _validate_region_zone(region: Optional[str],
+                          zone: str) -> Tuple[Optional[str], Optional[str]]:
+    return common.validate_region_zone_impl('aws', _df, region, zone)
+
+
 def validate_region_zone(
         region: Optional[str],
         zone: Optional[str]) -> Tuple[Optional[str], Optional[str]]:
-    return common.validate_region_zone_impl('aws', _df, region, zone)
+    if zone is None:
+        return _validate_region(region)
+    return _validate_region_zone(region, zone)
 
 
+@_apply_az_mapping_decorator
 def accelerator_in_region_or_zone(acc_name: str,
                                   acc_count: int,
                                   region: Optional[str] = None,
@@ -88,6 +124,7 @@ def accelerator_in_region_or_zone(acc_name: str,
                                                      region, zone)
 
 
+@_apply_az_mapping_decorator
 def get_hourly_cost(instance_type: str,
                     use_spot: bool = False,
                     region: Optional[str] = None,
@@ -113,6 +150,7 @@ def get_accelerators_from_instance_type(
     return common.get_accelerators_from_instance_type_impl(_df, instance_type)
 
 
+@_apply_az_mapping_decorator
 def get_instance_type_for_accelerator(
     acc_name: str,
     acc_count: int,
@@ -134,6 +172,7 @@ def get_instance_type_for_accelerator(
                                                          zone=zone)
 
 
+@_apply_az_mapping_decorator
 def get_region_zones_for_instance_type(instance_type: str,
                                        use_spot: bool) -> List['cloud.Region']:
     df = _df[_df['InstanceType'] == instance_type]
diff --git a/tests/conftest.py b/tests/conftest.py
index b37b7541281..1d31f0c794c 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -161,6 +161,7 @@ def pytest_sessionstart(session):
         try:
             # Try to fetch the AZ mapping from AWS (if we have AWS access)
             from sky.clouds.service_catalog import aws_catalog
+            aws_catalog._apply_az_mapping()
         except:
             # If we don't have AWS access, create a dummy file
             with tempfile.NamedTemporaryFile(mode='w', delete=False) as f:


Code review comments:
diff_hunk: @@ -44,6 +49,10 @@ def _apply_az_mapping(df: 'pd.DataFrame') -> 'pd.DataFrame':
         A dataframe with column 'AvailabilityZone' that's correctly replaced
         with the zone name (e.g. us-east-1a).
     """
+    global _az_mapping_applied
+    if _az_mapping_applied:
path: sky/clouds/service_catalog/aws_catalog.py
text: @user1:
An outside module (conftest) calls into this without a lock. Should we take the lock inside this func, rather than in the decorator?

diff_hunk: @@ -61,25 +70,52 @@ def _apply_az_mapping(df: 'pd.DataFrame') -> 'pd.DataFrame':
     # Use inner join to drop rows with unknown AZ IDs, which are likely
     # because the user does not have access to that Region. Otherwise,
     # there will be rows with NaN in the AvailabilityZone column.
-    df = df.merge(az_mappings, on=['AvailabilityZone'], how='inner')
-    df = df.drop(columns=['AvailabilityZone']).rename(
+    _df = _df.merge(az_mappings, on=['AvailabilityZone'], how='inner')
+    _df = _df.drop(columns=['AvailabilityZone']).rename(
         columns={'AvailabilityZoneName': 'AvailabilityZone'})
-    return df
+    _az_mapping_applied = True
+
+
+def _apply_az_mapping_decorator(func):
+    """Decorator to apply availability zone mapping to the dataframe."""
 
+    def wrapper(*args, **kwargs):
+        with _apply_az_mapping_lock:
+            _apply_az_mapping()
+        return func(*args, **kwargs)
 
-_df = _apply_az_mapping(_df)
+    return wrapper
 
 
+@_apply_az_mapping_decorator
 def instance_type_exists(instance_type: str) -> bool:
     return common.instance_type_exists_impl(_df, instance_type)
 
 
+# Helper functions for `validate_region_zone`.
path: sky/clouds/service_catalog/aws_catalog.py
text: @user1:
Minor preference: I feel like it's clearer to keep these internal to validate_region_zone() (shallow methods). Up to you though.

diff_hunk: @@ -61,25 +70,52 @@ def _apply_az_mapping(df: 'pd.DataFrame') -> 'pd.DataFrame':
     # Use inner join to drop rows with unknown AZ IDs, which are likely
     # because the user does not have access to that Region. Otherwise,
     # there will be rows with NaN in the AvailabilityZone column.
-    df = df.merge(az_mappings, on=['AvailabilityZone'], how='inner')
-    df = df.drop(columns=['AvailabilityZone']).rename(
+    _df = _df.merge(az_mappings, on=['AvailabilityZone'], how='inner')
+    _df = _df.drop(columns=['AvailabilityZone']).rename(
         columns={'AvailabilityZoneName': 'AvailabilityZone'})
-    return df
+    _az_mapping_applied = True
+
+
+def _apply_az_mapping_decorator(func):
+    """Decorator to apply availability zone mapping to the dataframe."""
 
+    def wrapper(*args, **kwargs):
+        with _apply_az_mapping_lock:
+            _apply_az_mapping()
+        return func(*args, **kwargs)
 
-_df = _apply_az_mapping(_df)
+    return wrapper
 
 
+@_apply_az_mapping_decorator
path: sky/clouds/service_catalog/aws_catalog.py
text: @user1:
Does this func need the mapping?

diff_hunk: @@ -61,25 +70,52 @@ def _apply_az_mapping(df: 'pd.DataFrame') -> 'pd.DataFrame':
     # Use inner join to drop rows with unknown AZ IDs, which are likely
     # because the user does not have access to that Region. Otherwise,
     # there will be rows with NaN in the AvailabilityZone column.
-    df = df.merge(az_mappings, on=['AvailabilityZone'], how='inner')
-    df = df.drop(columns=['AvailabilityZone']).rename(
+    _df = _df.merge(az_mappings, on=['AvailabilityZone'], how='inner')
+    _df = _df.drop(columns=['AvailabilityZone']).rename(
         columns={'AvailabilityZoneName': 'AvailabilityZone'})
-    return df
+    _az_mapping_applied = True
+
+
+def _apply_az_mapping_decorator(func):
+    """Decorator to apply availability zone mapping to the dataframe."""
 
+    def wrapper(*args, **kwargs):
+        with _apply_az_mapping_lock:
+            _apply_az_mapping()
+        return func(*args, **kwargs)
 
-_df = _apply_az_mapping(_df)
+    return wrapper
 
 
+@_apply_az_mapping_decorator
 def instance_type_exists(instance_type: str) -> bool:
     return common.instance_type_exists_impl(_df, instance_type)
 
 
+# Helper functions for `validate_region_zone`.
+def _validate_region(
+        region: Optional[str]) -> Tuple[Optional[str], Optional[str]]:
+    # When zone is not specified, we don't need to apply the AZ mapping.
+    # This is useful for `sky show-gpus` where user may not have access to
+    # AWS.
+    return common.validate_region_zone_impl('aws', _df, region, None)
path: sky/clouds/service_catalog/aws_catalog.py
text: @user1:
```suggestion
    return common.validate_region_zone_impl('aws', _df, region, zone=None)
```

diff_hunk: @@ -61,25 +70,52 @@ def _apply_az_mapping(df: 'pd.DataFrame') -> 'pd.DataFrame':
     # Use inner join to drop rows with unknown AZ IDs, which are likely
     # because the user does not have access to that Region. Otherwise,
     # there will be rows with NaN in the AvailabilityZone column.
-    df = df.merge(az_mappings, on=['AvailabilityZone'], how='inner')
-    df = df.drop(columns=['AvailabilityZone']).rename(
+    _df = _df.merge(az_mappings, on=['AvailabilityZone'], how='inner')
+    _df = _df.drop(columns=['AvailabilityZone']).rename(
         columns={'AvailabilityZoneName': 'AvailabilityZone'})
-    return df
+    _az_mapping_applied = True
+
+
+def _apply_az_mapping_decorator(func):
+    """Decorator to apply availability zone mapping to the dataframe."""
 
+    def wrapper(*args, **kwargs):
+        with _apply_az_mapping_lock:
+            _apply_az_mapping()
+        return func(*args, **kwargs)
 
-_df = _apply_az_mapping(_df)
+    return wrapper
 
 
+@_apply_az_mapping_decorator
 def instance_type_exists(instance_type: str) -> bool:
     return common.instance_type_exists_impl(_df, instance_type)
 
 
+# Helper functions for `validate_region_zone`.
+def _validate_region(
+        region: Optional[str]) -> Tuple[Optional[str], Optional[str]]:
+    # When zone is not specified, we don't need to apply the AZ mapping.
path: sky/clouds/service_catalog/aws_catalog.py
text: @user1:
```suggestion
    # When zone is not specified, we don't need to apply the AZ mapping (note that no decorator is applied to this func).
```

diff_hunk: @@ -29,13 +30,17 @@
 # skypilot-catalog/.github/workflows/update-aws-catalog.yml
 _PULL_FREQUENCY_HOURS = 7
 
+# Flag to indicate whether the availability zone mapping has been applied.
+_az_mapping_applied = False
+_apply_az_mapping_lock = threading.Lock()
+
 _df = common.read_catalog('aws/vms.csv',
                           pull_frequency_hours=_PULL_FREQUENCY_HOURS)
 _image_df = common.read_catalog('aws/images.csv',
                                 pull_frequency_hours=_PULL_FREQUENCY_HOURS)
 
 
-def _apply_az_mapping(df: 'pd.DataFrame') -> 'pd.DataFrame':
+def _apply_az_mapping():
path: sky/clouds/service_catalog/aws_catalog.py
text: @user1:
Since we now have the decorator, how about `_apply_az_mapping_once()` or `_ensure_az_mapping_is_applied()`?