[AWS] Update catalog from host periodically (skypilot-org#1451)

* wip

* Add auto update and preferred area

* Add filters for AWS

* format

* fix

* Auto-update by default

* turn off non default aws regions for global

* address comments

* Map AZ ID to AZ name

* fix comment

* fix name

* fix

* fix

* fix path

* format

* Fix

* wip: fix

* address comments

* format

* Add ap-south-1

* update faq for download global regions

* modify docs

* update docs

* fix pytest

* format

* fix pytest

* Fetch the az mapping during checking

* format

* remove output for downloading mapping / catalogs

* fix

* fetching

* increase spot time

* increase timeout

* fix no follow for spot

* format

Author: Michaelvll

docs/source/reference/faq.rst

@@ -79,17 +79,17 @@ To avoid rerunning the ``setup`` commands, pass the ``--no-setup`` flag to ``sky
 (Advanced) How to make SkyPilot use all global regions?
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-By default, SkyPilot only supports the US regions on different clouds for convenience. If you want to utilize all global regions, please run the following command:
+By default, SkyPilot supports most global regions on AWS and only supports the US regions on GCP and Azure. If you want to utilize all global regions, please run the following command:
 
 .. code-block:: bash
 
-  mkdir -p ~/.sky/catalogs/v4
-  cd ~/.sky/catalogs/v4
-  # Fetch all regions for AWS
-  python -m sky.clouds.service_catalog.data_fetchers.fetch_aws --all-regions
+  version=$(python -c 'import sky; print(sky.clouds.service_catalog.constants.CATALOG_SCHEMA_VERSION)')
+  mkdir -p ~/.sky/catalogs/${version}
+  cd ~/.sky/catalogs/${version}
   # Fetch all regions for GCP
   pip install lxml
   python -m sky.clouds.service_catalog.data_fetchers.fetch_gcp --all-regions
+  
   # Fetch all regions for Azure
   python -m sky.clouds.service_catalog.data_fetchers.fetch_azure --all-regions
 

sky/backends/backend_utils.py

@@ -1,5 +1,4 @@
 """Util constants/functions for the backends."""
-import contextlib
 import copy
 from datetime import datetime
 import difflib
@@ -1429,15 +1428,6 @@ def _process_cli_query(
     ]
 
 
-@contextlib.contextmanager
-def suppress_output():
-    """Suppress stdout and stderr."""
-    with open(os.devnull, 'w') as devnull:
-        with contextlib.redirect_stdout(devnull), contextlib.redirect_stderr(
-                devnull):
-            yield
-
-
 def _ray_launch_hash(cluster_name: str, ray_config: Dict[str, Any]) -> Set[str]:
     """Returns a set of Ray launch config hashes, one per node type."""
     # Use the cached Ray launch hashes if they exist.
@@ -1447,7 +1437,7 @@ def _ray_launch_hash(cluster_name: str, ray_config: Dict[str, Any]) -> Set[str]:
     if ray_launch_hashes is not None:
         logger.debug('Using cached launch_caches')
         return set(ray_launch_hashes)
-    with suppress_output():
+    with ux_utils.suppress_output():
         ray_config = ray_commands._bootstrap_config(ray_config)  # pylint: disable=protected-access
     # Adopted from https://github.com/ray-project/ray/blob/ray-2.0.1/python/ray/autoscaler/_private/node_launcher.py#L87-L97
     # TODO(zhwu): this logic is duplicated from the ray code above (keep in sync).
@@ -1461,7 +1451,7 @@ def _ray_launch_hash(cluster_name: str, ray_config: Dict[str, Any]) -> Set[str]:
         launch_config = copy.deepcopy(launch_config)
 
         launch_config.update(node_config['node_config'])
-        with suppress_output():
+        with ux_utils.suppress_output():
             current_hash = ray_util.hash_launch_conf(launch_config,
                                                      ray_config['auth'])
         launch_hashes.add(current_hash)

sky/clouds/aws.py

@@ -330,6 +330,9 @@ def check_credentials(self) -> Tuple[bool, Optional[str]]:
                 ' Make sure that the access and secret keys are correct.'
                 ' To reconfigure the credentials, ' + help_str[1].lower() +
                 help_str[2:])
+
+        # Fetch the AWS availability zones mapping from ID to name.
+        from sky.clouds.service_catalog import aws_catalog  # pylint: disable=import-outside-toplevel,unused-import
         return True, None
 
     def get_credential_file_mounts(self) -> Dict[str, str]:

sky/clouds/service_catalog/aws_catalog.py

@@ -3,16 +3,62 @@
 This module loads the service catalog file and can be used to query
 instance types and pricing information for AWS.
 """
+import colorama
+import os
 import typing
 from typing import Dict, List, Optional, Tuple
 
+import pandas as pd
+
+from sky import sky_logging
 from sky.clouds.service_catalog import common
+from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
     from sky.clouds import cloud
 
-_df = common.read_catalog('aws/vms.csv')
-_image_df = common.read_catalog('aws/images.csv')
+logger = sky_logging.init_logger(__name__)
+
+# Keep it synced with the frequency in
+# skypilot-catalog/.github/workflows/update-aws-catalog.yml
+_PULL_FREQUENCY_HOURS = 7
+
+_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':
+    """Maps zone IDs (use1-az1) to zone names (us-east-1x).
+
+    Such mappings are account-specific and determined by AWS.
+
+    Returns:
+        A dataframe with column 'AvailabilityZone' that's correctly replaced
+        with the zone name (e.g. us-east-1a).
+    """
+    az_mapping_path = common.get_catalog_path('aws/az_mappings.csv')
+    if not os.path.exists(az_mapping_path):
+        # Fetch az mapping from AWS.
+        # pylint: disable=import-outside-toplevel
+        import ray
+        from sky.clouds.service_catalog.data_fetchers import fetch_aws
+        logger.info(f'{colorama.Style.DIM}Fetching availability zones mapping '
+                    f'for AWS...{colorama.Style.RESET_ALL}')
+        with ux_utils.suppress_output():
+            ray.init()
+            az_mappings = fetch_aws.fetch_availability_zone_mappings()
+        az_mappings.to_csv(az_mapping_path, index=False)
+    else:
+        az_mappings = pd.read_csv(az_mapping_path)
+    df = df.merge(az_mappings, on=['AvailabilityZone'], how='left')
+    df = df.drop(columns=['AvailabilityZone']).rename(
+        columns={'AvailabilityZoneName': 'AvailabilityZone'})
+    return df
+
+
+_df = _apply_az_mapping(_df)
 
 
 def instance_type_exists(instance_type: str) -> bool:

sky/clouds/service_catalog/common.py

@@ -1,8 +1,11 @@
 """Common utilities for service catalog."""
+import hashlib
 import os
+import time
 from typing import Dict, List, NamedTuple, Optional, Tuple
 
 import difflib
+import filelock
 import requests
 import pandas as pd
 
@@ -46,32 +49,70 @@ def get_catalog_path(filename: str) -> str:
     return os.path.join(_CATALOG_DIR, filename)
 
 
-def read_catalog(filename: str) -> pd.DataFrame:
+def read_catalog(filename: str,
+                 pull_frequency_hours: Optional[int] = None) -> pd.DataFrame:
     """Reads the catalog from a local CSV file.
 
     If the file does not exist, download the up-to-date catalog that matches
     the schema version.
+    If `pull_frequency_hours` is not None: pull the latest catalog with
+    possibly updated prices, if the local catalog file is older than
+    `pull_frequency_hours` and no changes to the local catalog file are
+    made after the last pull.
     """
     assert filename.endswith('.csv'), 'The catalog file must be a CSV file.'
+    assert (pull_frequency_hours is None or
+            pull_frequency_hours > 0), pull_frequency_hours
     catalog_path = get_catalog_path(filename)
     cloud = cloud_lib.CLOUD_REGISTRY.from_str(os.path.dirname(filename))
-    if not os.path.exists(catalog_path):
-        url = f'{constants.HOSTED_CATALOG_DIR_URL}/{constants.CATALOG_SCHEMA_VERSION}/{filename}'  # pylint: disable=line-too-long
-        with backend_utils.safe_console_status(
-                f'Downloading {cloud} catalog...'):
-            try:
-                r = requests.get(url)
-                r.raise_for_status()
-            except requests.exceptions.RequestException as e:
-                logger.error(f'Failed to download {cloud} catalog:')
-                with ux_utils.print_exception_no_traceback():
-                    raise e
-        # Save the catalog to a local file.
-        os.makedirs(os.path.dirname(catalog_path), exist_ok=True)
-        with open(catalog_path, 'w') as f:
-            f.write(r.text)
-        logger.info(f'A new {cloud} catalog has been downloaded to '
-                    f'{catalog_path}')
+
+    meta_path = os.path.join(_CATALOG_DIR, '.meta', filename)
+    os.makedirs(os.path.dirname(meta_path), exist_ok=True)
+
+    # Atomic check, to avoid conflicts with other processes.
+    # TODO(mraheja): remove pylint disabling when filelock version updated
+    # pylint: disable=abstract-class-instantiated
+    with filelock.FileLock(meta_path + '.lock'):
+
+        def _need_update() -> bool:
+            if not os.path.exists(catalog_path):
+                return True
+            if pull_frequency_hours is None:
+                return False
+            # Check the md5 of the file to see if it has changed.
+            with open(catalog_path, 'rb') as f:
+                file_md5 = hashlib.md5(f.read()).hexdigest()
+            md5_filepath = meta_path + '.md5'
+            if os.path.exists(md5_filepath):
+                with open(md5_filepath, 'r') as f:
+                    last_md5 = f.read()
+                if file_md5 != last_md5:
+                    # Do not update the file if the user modified it.
+                    return False
+
+            last_update = os.path.getmtime(catalog_path)
+            return last_update + pull_frequency_hours * 3600 < time.time()
+
+        if _need_update():
+            url = f'{constants.HOSTED_CATALOG_DIR_URL}/{constants.CATALOG_SCHEMA_VERSION}/{filename}'  # pylint: disable=line-too-long
+            update_frequency_str = ''
+            if pull_frequency_hours is not None:
+                update_frequency_str = f' (every {pull_frequency_hours} hours)'
+            with backend_utils.safe_console_status(
+                    f'Updating {cloud} catalog{update_frequency_str}'):
+                try:
+                    r = requests.get(url)
+                    r.raise_for_status()
+                except requests.exceptions.RequestException as e:
+                    logger.error(f'Failed to download {cloud} catalog:')
+                    with ux_utils.print_exception_no_traceback():
+                        raise e
+            # Save the catalog to a local file.
+            os.makedirs(os.path.dirname(catalog_path), exist_ok=True)
+            with open(catalog_path, 'w') as f:
+                f.write(r.text)
+            with open(meta_path + '.md5', 'w') as f:
+                f.write(hashlib.md5(r.text.encode()).hexdigest())
 
     try:
         df = pd.read_csv(catalog_path)

sky/clouds/service_catalog/constants.py

@@ -2,5 +2,5 @@
 import os
 
 HOSTED_CATALOG_DIR_URL = 'https://raw.githubusercontent.com/skypilot-org/skypilot-catalog/master/catalogs'  # pylint: disable=line-too-long
-CATALOG_SCHEMA_VERSION = 'v4'
+CATALOG_SCHEMA_VERSION = 'v5'
 LOCAL_CATALOG_DIR = os.path.expanduser('~/.sky/catalogs/')