feat: 🎸 Add --regions and --exclude flag to fetch_azure (#2553)

* feat: 🎸 Add --regions and --exclude flag to fetch_azure * rm .idea * rm .idea & update .gitignore * format * update the code based on the review feedback * fix * fix * fix * fix * Minor fixes --------- Co-authored-by: Isaac Ong <[email protected]>
skypilot-org · Sep 14, 2023 · 4cd0d45 · 4cd0d45
1 parent e145996
commit 4cd0d45
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 10 deletions.
diff --git a/.gitignore b/.gitignore
@@ -10,3 +10,5 @@ build/
 sky_logs/
 sky/clouds/service_catalog/data_fetchers/*.csv
 .vscode/
+.idea/
+
diff --git a/docs/source/reference/faq.rst b/docs/source/reference/faq.rst
@@ -103,12 +103,22 @@ By default, SkyPilot supports most global regions on AWS and only supports the U
   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
+  # GCP
   pip install lxml
+  # Fetch U.S. regions for GCP
+  python -m sky.clouds.service_catalog.data_fetchers.fetch_gcp
+  # Fetch all regions for GCP
   python -m sky.clouds.service_catalog.data_fetchers.fetch_gcp --all-regions
-  
+
+  # Azure
+  # Fetch U.S. regions for Azure
+  python -m sky.clouds.service_catalog.data_fetchers.fetch_azure
   # Fetch all regions for Azure
   python -m sky.clouds.service_catalog.data_fetchers.fetch_azure --all-regions
+  # Fetch the specified regions for Azure
+  python -m sky.clouds.service_catalog.data_fetchers.fetch_azure --regions japaneast australiaeast uksouth
+  # Fetch U.S. regions for Azure, excluding the specified regions
+  python -m sky.clouds.service_catalog.data_fetchers.fetch_azure --exclude centralus eastus
 
 To make your managed spot jobs potentially use all global regions, please log into the spot controller with ``ssh sky-spot-controller-<hash>`` 
 (the full name can be found in ``sky status``), and run the commands above.

diff --git a/sky/clouds/service_catalog/data_fetchers/fetch_azure.py b/sky/clouds/service_catalog/data_fetchers/fetch_azure.py
@@ -14,7 +14,7 @@
 import pandas as pd
 import requests
 
-US_REGIONS = [
+US_REGIONS = {
     'centralus',
     'eastus',
     'eastus2',
@@ -24,7 +24,7 @@
     'westus',
     'westus2',
     'westus3',
-]
+}
 
 # Exclude the following regions as they do not have ProductName in the
 # pricing table. Reference: #1768 #2548
@@ -249,14 +249,29 @@ def get_additional_columns(row):
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument(
-        '--all-regions',
-        action='store_true',
-        help='Fetch all global regions, not just the U.S. ones.')
+    group = parser.add_mutually_exclusive_group()
+    group.add_argument('--all-regions',
+                       action='store_true',
+                       help='Fetch all global regions, not just the U.S. ones.')
+    group.add_argument('--regions',
+                       nargs='+',
+                       help='Fetch the list of specified regions.')
+    parser.add_argument('--exclude',
+                        nargs='+',
+                        help='Exclude the list of specified regions.')
     args = parser.parse_args()
 
-    region_filter = get_regions() if args.all_regions else US_REGIONS
-    region_filter = set(region_filter) - EXCLUDED_REGIONS
+    if args.regions:
+        region_filter = set(args.regions) - EXCLUDED_REGIONS
+    elif args.all_regions:
+        region_filter = set(get_regions()) - EXCLUDED_REGIONS
+    else:
+        region_filter = US_REGIONS
+    region_filter = region_filter - set(
+        args.exclude) if args.exclude else region_filter
+
+    if not region_filter:
+        raise ValueError('No regions to fetch. Please check your arguments.')
 
     instance_df = get_all_regions_instance_types_df(region_filter)
     os.makedirs('azure', exist_ok=True)