[AWS] Failover Optimization from Quota

sky/backends/cloud_vm_ray_backend.py

@@ -1331,6 +1331,16 @@ def _retry_zones(
             to_provision, 'region should have been set by the optimizer.')
         region = clouds.Region(to_provision.region)
 
+        if (not to_provision.cloud.check_quota_not_zero(
+                to_provision.region, to_provision.instance_type,
+                to_provision.use_spot)):
+
+            # if quota is found to be zero, raise exception and skip to
+            # the next region
+            raise exceptions.ResourcesUnavailableError(
+                'Zero quota in the region with '
+                'attempted provisioning')
+
         for zones in self._yield_zones(to_provision, num_nodes, cluster_name,
                                        prev_cluster_status):
             # Filter out zones that are blocked, if any.

sky/clouds/aws.py

@@ -687,3 +687,57 @@ def _get_disk_specs(cls, disk_tier: Optional[str]) -> Dict[str, Any]:
             'disk_throughput': tier2iops[tier] // 16,
             'custom_disk_perf': tier != 'low',
         }
+
+    @classmethod
+    def check_quota_not_zero(cls,
+                             region: str,
+                             instance_type: str,
+                             use_spot: bool = False) -> bool:
+        """
+        AWS-specific implmentation of check_quota_not_zero. The function works by
+        matching the instance_type to the corresponding AWS quota code, and then using
+        the boto3 Python API to query the region for the specific quota code.
+
+        Returns:
+            False if the quota is found to be zero, and True otherwise.
+        Raises:
+            ImportError: if the dependencies for AWS are not able to be installed.
+            botocore.exceptions.ClientError: error in Boto3 client request.
+        """
+
+        try:
+            # pylint: disable=import-outside-toplevel
+            import boto3
+            import botocore
+        except ImportError:
+            raise ImportError('Fail to import dependencies for AWS.'
+                              'Try pip install "skypilot[aws]"') from None
+
+        from sky.clouds.service_catalog import aws_catalog  # pylint: disable=import-outside-toplevel,unused-import
+
+        spot_header = ''
+        if use_spot:
+            spot_header = 'SpotInstanceCode'
+        else:
+            spot_header = 'OnDemandInstanceCode'
+
+        quota_code = aws_catalog.get_quota_code(instance_type, spot_header)
+
+        if quota_code is None:
+            # Quota code not found in the catalog for the chosen instance_type, try provisioning anyway
+            return True
+
+        client = boto3.client('service-quotas', region_name=region)
+        try:
+            response = client.get_service_quota(ServiceCode='ec2',
+                                                QuotaCode=quota_code)
+        except botocore.exceptions.ClientError:
+            # Botocore client connection not established, try provisioning anyways
+            return True
+
+        if response['Quota']['Value'] == 0:
+            # Quota found to be zero, do not try provisioning
+            return False
+
+        # Quota found to be greater than zero, try provisioning
+        return True

sky/clouds/cloud.py

@@ -467,5 +467,58 @@ def check_disk_tier_enabled(cls, instance_type: str,
         """
         raise NotImplementedError
 
+    @classmethod
+    def check_quota_not_zero(cls,
+                             region: str,
+                             instance_type: str,
+                             use_spot: bool = False) -> bool:
+        """
+        Checks to ensure that a particular accelerator has a nonzero quota
+        in a region.
+
+        (As of right now, check_quota_not_zero is only implemented for AWS.)
+
+        The _retry_zones funtion in cloud_vm_ray_backend goes through different
+        candidate regions and attempts to provision the requested instance_type
+        accelerators in the region, until a successful provisioning happens
+        or all regions with the requested accelerator have been looked at.
+        Previously, SkyPilot would attempt to provision resources in all of
+        these regions. However, many regions would have a zero quota or
+        inadequate quota, meaning these attempted provisionings were destined
+        to fail from the get-go.
+
+        Checking the quota is substantially faster than attempting a failed
+        provision (~1 second vs 30+ seconds) so this function attempts to
+        check the resource quota and return False if it is found to be zero,
+        or True otherwise. If False is returned, _retry_zones will not attempt
+        a provision in the region, saving time.
+
+        We are only checking for a nonzero quota, instead of also factoring in
+        quota utilization because many cloud providers' APIs don't have a
+        built-in command for checking the real-time utilization. Checking
+        real-time utilization is a more difficult endeavor that involves
+        monitoring etc., so we are holding off on that for now.
+
+        If for at any point the function fails, whether it's because we can't
+        import the necessary dependencies or a query using a cloud provider's
+        API fails, we will return True, because we cannot conclusively say the
+        relevant quota is zero in these cases, and we don't want to
+        preemptively exclude regions from an attempted provision if they may
+        have an adequate quota.
+
+        Design choice: We chose a just-in-time approach where
+        check_quota_not_zero is called immediately before a potential
+        attempted provision, rather than checking all region quotas
+        beforehand, storing them, and using those values on-demand. This is
+        because, for example, _retry_zones may only need to go through one or
+        a few regions before a successful region, and running a query to check
+        *every* region's quota beforehand would cause an unnecessary delay.
+
+        Returns:
+            False if the quota is found to be zero, and true otherwise.
+        """
+
+        return True
+
     def __repr__(self):
         return self._REPR

sky/clouds/service_catalog/aws_catalog.py

@@ -64,6 +64,9 @@
 _image_df = common.read_catalog('aws/images.csv',
                                 pull_frequency_hours=_PULL_FREQUENCY_HOURS)
 
+_quotas_df = common.read_catalog('aws/instance_quota_mapping.csv',
+                                 pull_frequency_hours=_PULL_FREQUENCY_HOURS)
+
 
 def _fetch_and_apply_az_mapping(df: pd.DataFrame) -> pd.DataFrame:
     """Maps zone IDs (use1-az1) to zone names (us-east-1x).
@@ -148,6 +151,19 @@ def _get_df() -> pd.DataFrame:
         return _user_df
 
 
+def get_quota_code(instance_type: str, spot_header: str) -> Optional[str]:
+    # Get the quota code from the accelerator instance type
+    # This will be used in the botocore command to check for
+    # a non-zero quota
+    try:
+        quota_code = _quotas_df.loc[_quotas_df['InstanceType'] == instance_type,
+                                    spot_header].values[0]
+        return quota_code
+
+    except IndexError:
+        return None
+
+
 def instance_type_exists(instance_type: str) -> bool:
     return common.instance_type_exists_impl(_get_df(), instance_type)
 

sky/clouds/service_catalog/constants.py

@@ -1,6 +1,7 @@
 """Constants used for service catalog."""
 import os
 
-HOSTED_CATALOG_DIR_URL = 'https://raw.githubusercontent.com/skypilot-org/skypilot-catalog/master/catalogs'  # pylint: disable=line-too-long
+# TODO (Hriday Sheth) change URL once catalog branch is merged in
+HOSTED_CATALOG_DIR_URL = 'https://raw.githubusercontent.com/shethhriday29/skypilot-catalog/master/catalogs'  # pylint: disable=line-too-long
 CATALOG_SCHEMA_VERSION = 'v5'
 LOCAL_CATALOG_DIR = os.path.expanduser('~/.sky/catalogs/')