[Custom Image] Support tag for the images and global regions (skypilot-org#1366)

* Support image tag for AWS

* add gcp image support

* address comments

* fix

* remove pandas warning

* Add example for using ubuntu1804

* add ubuntu 1804 in the test

* Enforce trying us regions first

* format

* address comments

* address comments

* Add docs and rename methods

* Add fetch global regions for GCP

* Add all regions for Azure

* rename and add doc

* remvoe accidently added folder

* fix service_catalog

* remove extra line

* Address comments

* mkdir for catalog path

* increase waiting time in test

* fix test recovery

* format

Author: Michaelvll

docs/source/reference/faq.rst

@@ -76,6 +76,24 @@ If you have edited the ``file_mounts`` section (e.g., by adding some files) and
 To avoid rerunning the ``setup`` commands, pass the ``--no-setup`` flag to ``sky launch``.
 
 
+(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:
+
+.. code-block:: bash
+
+  cd ~/.sky/catalogs/v4
+  # Fetch all regions for AWS
+  python -m sky.clouds.service_catalog.data_fetchers.fetch_aws --all-regions
+  # Fetch all regions for GCP
+  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
+
+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.
+
 
 (Advanced) How to edit or update the regions or pricing information used by SkyPilot?
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

docs/source/reference/yaml-spec.rst

@@ -91,6 +91,11 @@ Available fields:
       #
       # AWS
       # To find AWS AMI ids: https://leaherb.com/how-to-find-an-aws-marketplace-ami-image-id
+      # You can also change the default OS version by choosing from the following image tags provided by SkyPilot:
+      #   image_id: skypilot:gpu-ubuntu-2004
+      #   image_id: skypilot:k80-ubuntu-2004
+      #   image_id: skypilot:gpu-ubuntu-1804
+      #   image_id: skypilot:k80-ubuntu-1804
       image_id: ami-0868a20f5a3bf9702
       # GCP
       # To find GCP images: https://cloud.google.com/compute/docs/images

examples/image_with_tag.yaml

@@ -0,0 +1,10 @@
+resources:
+  cloud: aws
+  image_id: skypilot:gpu-ubuntu-1804
+
+
+setup: |
+  echo "running setup"
+
+run: |
+  conda env list

sky/clouds/aws.py

@@ -9,6 +9,7 @@
 from typing import Dict, Iterator, List, Optional, Tuple
 
 from sky import clouds
+from sky import exceptions
 from sky.clouds import service_catalog
 
 if typing.TYPE_CHECKING:
@@ -94,45 +95,41 @@ def region_zones_provision_loop(
     @classmethod
     def get_default_ami(cls, region_name: str, instance_type: str) -> str:
         acc = cls.get_accelerators_from_instance_type(instance_type)
+        image_id = service_catalog.get_image_id_from_tag(
+            'skypilot:gpu-ubuntu-2004', region_name, clouds='aws')
         if acc is not None:
             assert len(acc) == 1, acc
             acc_name = list(acc.keys())[0]
             if acc_name == 'K80':
-                # Deep Learning AMI GPU PyTorch 1.10.0 (Ubuntu 20.04) 20211208
-                # Downgrade the AMI for K80 due as it is only compatible with
-                # NVIDIA driver lower than 470.
-                amis = {
-                    'us-east-1': 'ami-0868a20f5a3bf9702',
-                    'us-east-2': 'ami-09b8825010d4dc701',
-                    # This AMI is 20210623 as aws does not provide a newer one.
-                    'us-west-1': 'ami-0b3c34d643904a734',
-                    'us-west-2': 'ami-06b3479ab15aaeaf1',
-                }
-                assert region_name in amis, region_name
-                return amis[region_name]
-        # https://console.aws.amazon.com/ec2/v2/home?region=us-east-1#Images:visibility=public-images;v=3;search=:64,:Ubuntu%2020,:Deep%20Learning%20AMI%20GPU%20PyTorch # pylint: disable=line-too-long
-
-        # Commented below are newer AMIs, but as other clouds do not support
-        # torch==1.13.0+cu117 we do not use these AMIs to avoid frequent updates:
-        # Deep Learning AMI GPU PyTorch 1.12.1 (Ubuntu 20.04) 20221025
-        #   Nvidia driver: 510.47.03, CUDA Version: 11.6 (supports torch==1.13.0+cu117)
-        #     'us-east-1': 'ami-0eb1f91977a3fcc1b'
-        #     'us-east-2': 'ami-0274a6db2e19b7cc6'
-        #     'us-west-1': 'ami-0fb299af41d32cfd3'
-        #     'us-west-2': 'ami-04ba15f9bd464eb20'
-        #
-        # Current AMIs:
-        # Deep Learning AMI GPU PyTorch 1.10.0 (Ubuntu 20.04) 20220308
-        #   Nvidia driver: 510.47.03, CUDA Version: 11.6 (does not support torch==1.13.0+cu117)
-        amis = {
-            'us-east-1': 'ami-0729d913a335efca7',
-            'us-east-2': 'ami-070f4af81c19b41bf',
-            # This AMI is 20210623 as aws does not provide a newer one.
-            'us-west-1': 'ami-0b3c34d643904a734',
-            'us-west-2': 'ami-050814f384259894c',
-        }
-        assert region_name in amis, region_name
-        return amis[region_name]
+                image_id = service_catalog.get_image_id_from_tag(
+                    'skypilot:k80-ubuntu-2004', region_name, clouds='aws')
+        if image_id is not None:
+            return image_id
+        # Raise ResourcesUnavailableError to make sure the failover in
+        # CloudVMRayBackend will be correctly triggered.
+        # TODO(zhwu): This is a information leakage to the cloud implementor,
+        # we need to find a better way to handle this.
+        raise exceptions.ResourcesUnavailableError(
+            'No image found in catalog for region '
+            f'{region_name}. Try setting a valid image_id.')
+
+    @classmethod
+    def _get_image_id(cls, region_name: str, instance_type: str,
+                      image_id: Optional[str]) -> str:
+        if image_id is not None:
+            if image_id.startswith('skypilot:'):
+                image_id = service_catalog.get_image_id_from_tag(image_id,
+                                                                 region_name,
+                                                                 clouds='aws')
+                if image_id is None:
+                    # Raise ResourcesUnavailableError to make sure the failover
+                    # in CloudVMRayBackend will be correctly triggered.
+                    # TODO(zhwu): This is a information leakage to the cloud
+                    # implementor, we need to find a better way to handle this.
+                    raise exceptions.ResourcesUnavailableError(
+                        f'No image found for region {region_name}')
+            return image_id
+        return cls.get_default_ami(region_name, instance_type)
 
     @classmethod
     def get_zone_shell_cmd(cls) -> Optional[str]:
@@ -232,10 +229,7 @@ def make_deploy_resources_variables(
         else:
             custom_resources = None
 
-        if r.image_id is not None:
-            image_id = r.image_id
-        else:
-            image_id = self.get_default_ami(region_name, r.instance_type)
+        image_id = self._get_image_id(region_name, r.instance_type, r.image_id)
 
         return {
             'instance_type': r.instance_type,
@@ -342,9 +336,6 @@ def get_credential_file_mounts(self) -> Dict[str, str]:
     def instance_type_exists(self, instance_type):
         return service_catalog.instance_type_exists(instance_type, clouds='aws')
 
-    def validate_region_zone(self, region: Optional[str], zone: Optional[str]):
-        return service_catalog.validate_region_zone(region, zone, clouds='aws')
-
     def accelerator_in_region_or_zone(self,
                                       accelerator: str,
                                       acc_count: int,

sky/clouds/azure.py

@@ -297,11 +297,6 @@ def instance_type_exists(self, instance_type):
         return service_catalog.instance_type_exists(instance_type,
                                                     clouds='azure')
 
-    def validate_region_zone(self, region: Optional[str], zone: Optional[str]):
-        return service_catalog.validate_region_zone(region,
-                                                    zone,
-                                                    clouds='azure')
-
     def accelerator_in_region_or_zone(self,
                                       accelerator: str,
                                       acc_count: int,

sky/clouds/cloud.py

@@ -3,6 +3,7 @@
 import typing
 from typing import Dict, Iterator, List, Optional, Tuple
 
+from sky.clouds import service_catalog
 from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
@@ -162,6 +163,13 @@ def get_default_instance_type(cls) -> str:
     def _get_default_region(cls) -> Region:
         raise NotImplementedError
 
+    @classmethod
+    def is_image_tag_valid(cls, image_tag: str, region: Optional[str]) -> bool:
+        """Validates that the image tag is valid for this cloud."""
+        return service_catalog.is_image_tag_valid(image_tag,
+                                                  region,
+                                                  clouds=cls._REPR.lower())
+
     def get_feasible_launchable_resources(self, resources):
         """Returns a list of feasible and launchable resources.
 
@@ -194,7 +202,9 @@ def instance_type_exists(self, instance_type):
 
     def validate_region_zone(self, region: Optional[str], zone: Optional[str]):
         """Validates the region and zone."""
-        raise NotImplementedError
+        return service_catalog.validate_region_zone(region,
+                                                    zone,
+                                                    clouds=self._REPR.lower())
 
     def accelerator_in_region_or_zone(self,
                                       accelerator: str,

sky/clouds/gcp.py

@@ -34,8 +34,6 @@
     'active_config',
 ]
 
-_IMAGE_ID_PREFIX = ('projects/deeplearning-platform-release/global/images/')
-
 _GCLOUD_INSTALLATION_LOG = '~/.sky/logs/gcloud_installation.log'
 # Need to be run with /bin/bash
 # We factor out the installation logic to keep it align in both spot
@@ -217,7 +215,10 @@ def make_deploy_resources_variables(
         # gcloud compute images list \
         # --project deeplearning-platform-release \
         # --no-standard-images
-        image_id = _IMAGE_ID_PREFIX + 'common-cpu-v20220806'
+        # We use the debian image, as the ubuntu image has some connectivity
+        # issue when first booted.
+        image_id = service_catalog.get_image_id_from_tag(
+            'skypilot:cpu-debian-10', clouds='gcp')
 
         r = resources
         # Find GPU spec, if any.
@@ -261,17 +262,20 @@ def make_deploy_resources_variables(
                     # Though the image is called cu113, it actually has later
                     # versions of CUDA as noted below.
                     # CUDA driver version 470.57.02, CUDA Library 11.4
-                    image_id = _IMAGE_ID_PREFIX + 'common-cu113-v20220701'
+                    image_id = service_catalog.get_image_id_from_tag(
+                        'skypilot:k80-debian-10', clouds='gcp')
                 else:
                     # Though the image is called cu113, it actually has later
                     # versions of CUDA as noted below.
                     # CUDA driver version 510.47.03, CUDA Library 11.6
                     # Does not support torch==1.13.0 with cu117
-                    image_id = _IMAGE_ID_PREFIX + 'common-cu113-v20220806'
+                    image_id = service_catalog.get_image_id_from_tag(
+                        'skypilot:gpu-debian-10', clouds='gcp')
 
         if resources.image_id is not None:
             image_id = resources.image_id
 
+        assert image_id is not None, (image_id, r)
         resources_vars['image_id'] = image_id
 
         return resources_vars
@@ -452,9 +456,6 @@ def get_credential_file_mounts(self) -> Dict[str, str]:
     def instance_type_exists(self, instance_type):
         return service_catalog.instance_type_exists(instance_type, 'gcp')
 
-    def validate_region_zone(self, region: Optional[str], zone: Optional[str]):
-        return service_catalog.validate_region_zone(region, zone, clouds='gcp')
-
     def accelerator_in_region_or_zone(self,
                                       accelerator: str,
                                       acc_count: int,

sky/clouds/service_catalog/__init__.py

@@ -235,6 +235,20 @@ def get_tpus() -> List[str]:
     ]
 
 
+def get_image_id_from_tag(tag: str,
+                          region: Optional[str] = None,
+                          clouds: CloudFilter = None) -> str:
+    """Returns the image ID from the tag."""
+    return _map_clouds_catalog(clouds, 'get_image_id_from_tag', tag, region)
+
+
+def is_image_tag_valid(tag: str,
+                       region: Optional[str],
+                       clouds: CloudFilter = None) -> None:
+    """Validates the image tag."""
+    return _map_clouds_catalog(clouds, 'is_image_tag_valid', tag, region)
+
+
 __all__ = [
     'list_accelerators',
     'list_accelerator_counts',
@@ -246,6 +260,9 @@ def get_tpus() -> List[str]:
     'get_region_zones_for_accelerators',
     'get_common_gpus',
     'get_tpus',
+    # Images
+    'get_image_id_from_tag',
+    'is_image_tag_valid',
     # Constants
     'HOSTED_CATALOG_DIR_URL',
     'CATALOG_SCHEMA_VERSION',

sky/clouds/service_catalog/aws_catalog.py

@@ -3,12 +3,16 @@
 This module loads the service catalog file and can be used to query
 instance types and pricing information for AWS.
 """
+import typing
 from typing import Dict, List, Optional, Tuple
 
-from sky.clouds import cloud
 from sky.clouds.service_catalog import common
 
-_df = common.read_catalog('aws.csv')
+if typing.TYPE_CHECKING:
+    from sky.clouds import cloud
+
+_df = common.read_catalog('aws/vms.csv')
+_image_df = common.read_catalog('aws/images.csv')
 
 
 def instance_type_exists(instance_type: str) -> bool:
@@ -57,9 +61,19 @@ def get_instance_type_for_accelerator(
 
 
 def get_region_zones_for_instance_type(instance_type: str,
-                                       use_spot: bool) -> List[cloud.Region]:
+                                       use_spot: bool) -> List['cloud.Region']:
     df = _df[_df['InstanceType'] == instance_type]
-    return common.get_region_zones(df, use_spot)
+    region_list = common.get_region_zones(df, use_spot)
+    # Hack: Enforce US regions are always tried first:
+    #   [US regions sorted by price] + [non-US regions sorted by price]
+    us_region_list = []
+    other_region_list = []
+    for region in region_list:
+        if region.name.startswith('us-'):
+            us_region_list.append(region)
+        else:
+            other_region_list.append(region)
+    return us_region_list + other_region_list
 
 
 def list_accelerators(gpus_only: bool,
@@ -69,3 +83,13 @@ def list_accelerators(gpus_only: bool,
     """Returns all instance types in AWS offering accelerators."""
     return common.list_accelerators_impl('AWS', _df, gpus_only, name_filter,
                                          case_sensitive)
+
+
+def get_image_id_from_tag(tag: str, region: Optional[str]) -> Optional[str]:
+    """Returns the image id from the tag."""
+    return common.get_image_id_from_tag_impl(_image_df, tag, region)
+
+
+def is_image_tag_valid(tag: str, region: Optional[str]) -> bool:
+    """Returns whether the image tag is valid."""
+    return common.is_image_tag_valid_impl(_image_df, tag, region)

sky/clouds/service_catalog/azure_catalog.py

@@ -5,11 +5,11 @@
 """
 from typing import Dict, List, Optional, Tuple
 
-from sky.clouds import cloud
+from sky import clouds as cloud_lib
 from sky.clouds.service_catalog import common
 from sky.utils import ux_utils
 
-_df = common.read_catalog('azure.csv')
+_df = common.read_catalog('azure/vms.csv')
 
 
 def instance_type_exists(instance_type: str) -> bool:
@@ -60,8 +60,8 @@ def get_instance_type_for_accelerator(
                                                          acc_count=acc_count)
 
 
-def get_region_zones_for_instance_type(instance_type: str,
-                                       use_spot: bool) -> List[cloud.Region]:
+def get_region_zones_for_instance_type(
+        instance_type: str, use_spot: bool) -> List[cloud_lib.Region]:
     df = _df[_df['InstanceType'] == instance_type]
     return common.get_region_zones(df, use_spot)