[Image] Check image size and existence (skypilot-org#1508)

* Check image size and existence

* Check image existence and size

* Fix no credential

* Fix error message

* fix gcp image

* lint

* address comments

* yapf

Author: Michaelvll

sky/adaptors/aws.py

@@ -56,7 +56,7 @@ def session():
 
 
 @import_package
-def client_exception():
-    """Client exception."""
+def botocore_exceptions():
+    """AWS botocore exception."""
     from botocore import exceptions
-    return exceptions.ClientError
+    return exceptions

sky/adaptors/gcp.py

@@ -68,3 +68,17 @@ def forbidden_exception():
     """Forbidden exception."""
     from google.api_core import exceptions as gcs_exceptions
     return gcs_exceptions.Forbidden
+
+
+@import_package
+def http_error_exception():
+    """HttpError exception."""
+    from googleapiclient import errors
+    return errors.HttpError
+
+
+@import_package
+def credential_error_exception():
+    """CredentialError exception."""
+    from google.auth import exceptions
+    return exceptions.DefaultCredentialsError

sky/backends/cloud_vm_ray_backend.py

@@ -633,12 +633,18 @@ def _update_blocklist_on_gcp_error(self, region, zones, stdout, stderr):
                 else:
                     assert False, error
         elif len(httperror_str) >= 1:
-            # Parse HttpError for unauthorized regions. Example:
-            # googleapiclient.errors.HttpError: <HttpError 403 when requesting ... returned "Location us-east1-d is not found or access is unauthorized.". # pylint: disable=line-too-long
-            # Details: "Location us-east1-d is not found or access is
-            # unauthorized.">
             logger.info(f'Got {httperror_str[0]}')
-            self._blocked_zones.add(zone.name)
+            if ('Requested disk size cannot be smaller than the image size'
+                    in httperror_str[0]):
+                logger.info('Skipping all regions due to disk size issue.')
+                for r, _ in clouds.GCP.region_zones_provision_loop():
+                    self._blocked_regions.add(r.name)
+            else:
+                # Parse HttpError for unauthorized regions. Example:
+                # googleapiclient.errors.HttpError: <HttpError 403 when requesting ... returned "Location us-east1-d is not found or access is unauthorized.". # pylint: disable=line-too-long
+                # Details: "Location us-east1-d is not found or access is
+                # unauthorized.">
+                self._blocked_zones.add(zone.name)
         else:
             # No such structured error response found.
             assert not exception_str, stderr

sky/clouds/aws.py

@@ -10,7 +10,9 @@
 
 from sky import clouds
 from sky import exceptions
+from sky.adaptors import aws
 from sky.clouds import service_catalog
+from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
     # renaming to avoid shadowing variables
@@ -31,6 +33,10 @@ def _run_output(cmd):
     return proc.stdout.decode('ascii')
 
 
+# TODO(zhwu): Move the default AMI size to the catalog instead.
+DEFAULT_AMI_GB = 45
+
+
 @clouds.CLOUD_REGISTRY.register
 class AWS(clouds.Cloud):
     """Amazon Web Services."""
@@ -114,27 +120,51 @@ def get_default_ami(cls, region_name: str, instance_type: str) -> str:
             f'{region_name}. Try setting a valid image_id.')
 
     @classmethod
-    def _get_image_id(cls, region_name: str, instance_type: str,
-                      image_id: Optional[Dict[str, str]]) -> str:
-        if image_id is not None:
-            if None in image_id:
-                image_id = image_id[None]
-            else:
-                assert region_name in image_id, image_id
-                image_id = image_id[region_name]
-            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)
+    def _get_image_id(
+        cls,
+        image_id: Optional[Dict[str, str]],
+        region_name: str,
+    ) -> str:
+        if image_id is None:
+            return None
+        if None in image_id:
+            image_id = image_id[None]
+        else:
+            assert region_name in image_id, image_id
+            image_id = image_id[region_name]
+        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
+
+    def get_image_size(self, image_id: str, region: Optional[str]) -> float:
+        if image_id.startswith('skypilot:'):
+            return DEFAULT_AMI_GB
+        assert region is not None, (image_id, region)
+        client = aws.client('ec2', region_name=region)
+        try:
+            image_info = client.describe_images(ImageIds=[image_id])
+            image_info = image_info['Images'][0]
+            image_size = image_info['BlockDeviceMappings'][0]['Ebs'][
+                'VolumeSize']
+        except aws.botocore_exceptions().NoCredentialsError:
+            # Fallback to default image size if no credentials are available.
+            # The credentials issue will be caught when actually provisioning
+            # the instance and appropriate errors will be raised there.
+            return DEFAULT_AMI_GB
+        except aws.botocore_exceptions().ClientError:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(f'Image {image_id!r} not found in '
+                                 f'AWS region {region}') from None
+        return image_size
 
     @classmethod
     def get_zone_shell_cmd(cls) -> Optional[str]:
@@ -234,7 +264,9 @@ def make_deploy_resources_variables(
         else:
             custom_resources = None
 
-        image_id = self._get_image_id(region_name, r.instance_type, r.image_id)
+        image_id = self._get_image_id(r.image_id, region_name)
+        if image_id is None:
+            image_id = self.get_default_ami(region_name, r.instance_type)
 
         return {
             'instance_type': r.instance_type,

sky/clouds/cloud.py

@@ -196,6 +196,14 @@ def get_credential_file_mounts(self) -> Dict[str, str]:
         """
         raise NotImplementedError
 
+    def get_image_size(self, image_id: str, region: Optional[str]) -> float:
+        """Check the image size from the cloud.
+
+        Returns: the image size in GB.
+        Raises: ValueError if the image cannot be found.
+        """
+        raise NotImplementedError
+
     def instance_type_exists(self, instance_type):
         """Returns whether the instance type exists for this cloud."""
         raise NotImplementedError

sky/clouds/gcp.py

@@ -9,12 +9,16 @@
 from google import auth
 
 from sky import clouds
+from sky import sky_logging
+from sky.adaptors import gcp
 from sky.clouds import service_catalog
 from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
     from sky import resources
 
+logger = sky_logging.init_logger(__name__)
+
 DEFAULT_GCP_APPLICATION_CREDENTIAL_PATH = os.path.expanduser(
     '~/.config/gcloud/'
     'application_default_credentials.json')
@@ -51,6 +55,9 @@
     {{ cp {GCP_CONFIG_SKY_BACKUP_PATH} {GCP_CONFIG_PATH} > /dev/null 2>&1 || true; }} && \
     popd &>/dev/null'
 
+# TODO(zhwu): Move the default AMI size to the catalog instead.
+DEFAULT_GCP_IMAGE_GB = 50
+
 
 def _run_output(cmd):
     proc = subprocess.run(cmd,
@@ -187,6 +194,37 @@ def get_egress_cost(self, num_gigabytes):
     def is_same_cloud(self, other):
         return isinstance(other, GCP)
 
+    def get_image_size(self, image_id: str, region: Optional[str]) -> float:
+        del region  # unused
+        if image_id.startswith('skypilot:'):
+            return DEFAULT_GCP_IMAGE_GB
+        try:
+            compute = gcp.build('compute',
+                                'v1',
+                                credentials=None,
+                                cache_discovery=False)
+        except gcp.credential_error_exception() as e:
+            return DEFAULT_GCP_IMAGE_GB
+        try:
+            image_attrs = image_id.split('/')
+            if len(image_attrs) == 1:
+                raise ValueError(f'Image {image_id!r} not found in GCP.')
+            project = image_attrs[1]
+            image_name = image_attrs[-1]
+            image_infos = compute.images().get(project=project,
+                                               image=image_name).execute()
+            return float(image_infos['diskSizeGb'])
+        except gcp.http_error_exception() as e:
+            if e.resp.status == 403:
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError('Not able to access the image '
+                                     f'{image_id!r}') from None
+            if e.resp.status == 404:
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError(f'Image {image_id!r} not found in '
+                                     'GCP.') from None
+            raise
+
     @classmethod
     def get_default_instance_type(cls) -> str:
         # 8 vCpus, 52 GB RAM.  First-gen general purpose.

sky/data/storage.py

@@ -972,7 +972,7 @@ def _get_bucket(self) -> Tuple[StorageHandle, bool]:
             # accessible.
             self.client.head_bucket(Bucket=self.name)
             return bucket, False
-        except aws.client_exception() as e:
+        except aws.botocore_exceptions().ClientError as e:
             error_code = e.response['Error']['Code']
             # AccessDenied error for buckets that are private and not owned by
             # user.
@@ -1044,7 +1044,7 @@ def _create_s3_bucket(self,
                 s3_client.create_bucket(Bucket=bucket_name,
                                         CreateBucketConfiguration=location)
                 logger.info(f'Created S3 bucket {bucket_name} in {region}')
-        except aws.client_exception() as e:
+        except aws.botocore_exceptions().ClientError as e:
             with ux_utils.print_exception_no_traceback():
                 raise exceptions.StorageBucketCreateError(
                     f'Attempted to create a bucket '

sky/resources.py

@@ -74,11 +74,6 @@ def __init__(
             self._spot_recovery = spot_recovery.upper()
 
         if disk_size is not None:
-            if disk_size < 50:
-                with ux_utils.print_exception_no_traceback():
-                    raise ValueError(
-                        'OS disk size must be larger than 50GB. Got: '
-                        f'{disk_size}.')
             if round(disk_size) != disk_size:
                 with ux_utils.print_exception_no_traceback():
                     raise ValueError(
@@ -440,6 +435,19 @@ def _try_validate_image_id(self) -> None:
                         'image_id is only supported for AWS in a specific '
                         'region, please explicitly specify the region.')
 
+        # Validate the image exists and the size is smaller than the disk size.
+        for region, image_id in self._image_id.items():
+            # Check the image exists and get the image size.
+            # It will raise ValueError if the image does not exist.
+            image_size = self.cloud.get_image_size(image_id, region)
+            if image_size > self.disk_size:
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError(
+                        f'Image {image_id!r} is {image_size}GB, which is '
+                        f'larger than the specified disk_size: {self.disk_size}'
+                        ' GB. Please specify a larger disk_size to use this '
+                        'image.')
+
     def get_cost(self, seconds: float) -> float:
         """Returns cost in USD for the runtime in seconds."""
         hours = seconds / 3600