Repository: skypilot-org/skypilot
Pull request #2352 (2023-08-04 16:21:42)

GCP - reservations
<!-- Describe the changes in this PR -->

Adding GCP reservation support

<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [ ] Code formatting: `bash format.sh`
- [ ] Any manual or new tests for this PR (please specify below)
- [ ] All smoke tests: `pytest tests/test_smoke.py` 
- [ ] Relevant individual smoke tests: `pytest tests/test_smoke.py::test_fill_in_the_name` 
- [ ] Backward compatibility tests: `bash tests/backward_comaptibility_tests.sh`


Michaelvll: Thanks for the awesome work @gbmarc1! This helps the user utilize the prioritize their reserved instances. Very excited to see this included in SkyPilot. The smoke tests just passed. I am merging this PR now.

Tested:
- [x] `pytest tests/test_smoke.py` (except for the TPU VM pod test, which is unrelated to this PR)

Patch:
diff --git a/.gitignore b/.gitignore
index 3172e74a3..285ab7676 100644
--- a/.gitignore
+++ b/.gitignore
@@ -9,3 +9,4 @@ docs/_build/
 build/
 sky_logs/
 sky/clouds/service_catalog/data_fetchers/*.csv
+.vscode/
diff --git a/requirements-dev.txt b/requirements-dev.txt
index 8f29a5b46..9e8427248 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -12,6 +12,7 @@ mypy==0.991
 types-PyYAML
 types-requests
 types-setuptools
+types-cachetools
 
 # testing
 pytest
diff --git a/sky/backends/backend_utils.py b/sky/backends/backend_utils.py
index 4b71ca4ac..a05a62fec 100644
--- a/sky/backends/backend_utils.py
+++ b/sky/backends/backend_utils.py
@@ -859,6 +859,22 @@ def write_cluster_config(
     if isinstance(cloud, clouds.GCP):
         gcp_project_id = cloud.get_project_id(dryrun=dryrun)
 
+    specific_reservations = set(
+        skypilot_config.get_nested(('gcp', 'specific_reservations'), set()))
+
+    reservations = to_provision.get_reservations_available_resources(
+        specific_reservations)
+
+    filtered_specific_reservations = [
+        r for r, available_resources in reservations.items()
+        if r in specific_reservations and available_resources > 0
+    ]
+    available_specific_reservations = sum(
+        available_resources for r, available_resources in reservations.items()
+        if r in specific_reservations)
+    num_specific_reserved_workers = max(
+        min(available_specific_reservations - 1, num_nodes - 1), 0)
+
     assert cluster_name is not None
     credentials = sky_check.get_cloud_credential_file_mounts()
 
@@ -917,7 +933,6 @@ def write_cluster_config(
     default_aws_sg_name = f'sky-sg-{common_utils.user_and_hostname_hash()}'
     if ports is not None:
         default_aws_sg_name += f'-{common_utils.truncate_and_hash_cluster_name(cluster_name)}'
-
     # Use a tmp file path to avoid incomplete YAML file being re-used in the
     # future.
     tmp_yaml_path = yaml_path + '.tmp'
@@ -960,6 +975,8 @@ def write_cluster_config(
 
                 # GCP only:
                 'gcp_project_id': gcp_project_id,
+                'specific_reservations': filtered_specific_reservations,
+                'num_specific_reserved_workers': num_specific_reserved_workers,
 
                 # Conda setup
                 'conda_installation_commands':
diff --git a/sky/clouds/__init__.py b/sky/clouds/__init__.py
index 1094a28bd..74fc0089d 100644
--- a/sky/clouds/__init__.py
+++ b/sky/clouds/__init__.py
@@ -1,6 +1,6 @@
 """Clouds in Sky."""
 from sky.clouds.cloud import Cloud
-from sky.clouds.cloud import CLOUD_REGISTRY
+from sky.clouds.cloud_registry import CLOUD_REGISTRY
 from sky.clouds.cloud import CloudImplementationFeatures
 from sky.clouds.cloud import Region
 from sky.clouds.cloud import Zone
diff --git a/sky/clouds/cloud.py b/sky/clouds/cloud.py
index 06757a468..7185d4d16 100644
--- a/sky/clouds/cloud.py
+++ b/sky/clouds/cloud.py
@@ -3,10 +3,11 @@ import collections
 import enum
 import re
 import typing
-from typing import Dict, Iterator, List, Optional, Set, Tuple, Type
+from typing import Dict, Iterator, List, Optional, Set, Tuple
 
 from sky import exceptions
 from sky.clouds import service_catalog
+from sky import skypilot_config
 from sky.utils import log_utils
 from sky.utils import ux_utils
 
@@ -51,28 +52,6 @@ class Zone(collections.namedtuple('Zone', ['name'])):
     region: Region
 
 
-class _CloudRegistry(dict):
-    """Registry of clouds."""
-
-    def from_str(self, name: Optional[str]) -> Optional['Cloud']:
-        if name is None:
-            return None
-        if name.lower() not in self:
-            with ux_utils.print_exception_no_traceback():
-                raise ValueError(f'Cloud {name!r} is not a valid cloud among '
-                                 f'{list(self.keys())}')
-        return self.get(name.lower())
-
-    def register(self, cloud_cls: Type['Cloud']) -> Type['Cloud']:
-        name = cloud_cls.__name__.lower()
-        assert name not in self, f'{name} already registered'
-        self[name] = cloud_cls()
-        return cloud_cls
-
-
-CLOUD_REGISTRY = _CloudRegistry()
-
-
 class Cloud:
     """A cloud provider."""
 
@@ -318,6 +297,21 @@ class Cloud:
     def _get_feasible_launchable_resources(self, resources):
         raise NotImplementedError
 
+    def get_reservations_available_resources(
+        self,
+        instance_type: str,
+        region: str,
+        zone: Optional[str],
+        specific_reservations: Set[str],
+    ) -> Dict[str, int]:
+        """"
+        Returns the number of available resources per reservation for the given
+        instance type in the given region/zone.
+        Default implementation returns 0 for non-implemented clouds.
+        """
+        del instance_type, region, zone
+        return {reservation: 0 for reservation in specific_reservations}
+
     @classmethod
     def check_credentials(cls) -> Tuple[bool, Optional[str]]:
         """Checks if the user has access credentials to this cloud.
@@ -445,6 +439,19 @@ class Cloud:
             requested features.
         """
         unsupported_features2reason = cls._cloud_unsupported_features()
+
+        # Multi-node is not supported when specific_reservations is specified
+        # because if the node count is greater than the number of available
+        # resources in the specific reservations, the Google API will throw
+        # an error when trying to create the instance.
+        multi_node = CloudImplementationFeatures.MULTI_NODE
+        specific_reservations = skypilot_config.get_nested(
+            (cls._REPR.lower(), 'specific_reservations'), set())
+        if (specific_reservations and multi_node in requested_features):
+            unsupported_features2reason[multi_node] = (
+                'Multi-node is not supported when the `specific_reservations` '
+                f'is specified: {specific_reservations}')
+
         unsupported_features = set(unsupported_features2reason.keys())
         unsupported_features = requested_features.intersection(
             unsupported_features)
diff --git a/sky/clouds/cloud_registry.py b/sky/clouds/cloud_registry.py
new file mode 100644
index 000000000..db9a84e6e
--- /dev/null
+++ b/sky/clouds/cloud_registry.py
@@ -0,0 +1,34 @@
+"""
+Cloud need to be registered in the CLOUD_REGISTRY so that they can be discovered
+"""
+import typing
+from typing import Optional, Type
+
+from sky.utils import ux_utils
+if typing.TYPE_CHECKING:
+    import sky
+
+
+class _CloudRegistry(dict):
+    """Registry of clouds."""
+
+    def from_str(self,
+                 name: Optional[str]) -> Optional['sky.clouds.cloud.Cloud']:
+        if name is None:
+            return None
+        if name.lower() not in self:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(f'Cloud {name!r} is not a valid cloud among '
+                                 f'{list(self.keys())}')
+        return self.get(name.lower())
+
+    def register(
+        self, cloud_cls: Type['sky.clouds.cloud.Cloud']
+    ) -> Type['sky.clouds.cloud.Cloud']:
+        name = cloud_cls.__name__.lower()
+        assert name not in self, f'{name} already registered'
+        self[name] = cloud_cls()
+        return cloud_cls
+
+
+CLOUD_REGISTRY: _CloudRegistry = _CloudRegistry()
diff --git a/sky/clouds/gcp.py b/sky/clouds/gcp.py
index 1b2d395ee..25558e7fb 100644
--- a/sky/clouds/gcp.py
+++ b/sky/clouds/gcp.py
@@ -1,4 +1,5 @@
 """Google Cloud Platform."""
+import dataclasses
 import functools
 import json
 import os
@@ -6,7 +7,11 @@ import re
 import subprocess
 import time
 import typing
-from typing import Dict, Iterator, List, Optional, Tuple
+import datetime
+import operator
+from typing import Dict, Iterator, List, Optional, Tuple, Set
+
+import cachetools
 
 from sky import clouds
 from sky import exceptions
@@ -99,6 +104,70 @@ def is_api_disabled(endpoint: str, project_id: str) -> bool:
     return proc.returncode != 0
 
 
+@dataclasses.dataclass
+class SpecificReservation:
+    count: int
+    in_use_count: int
+
+    @classmethod
+    def from_dict(cls, d: dict) -> 'SpecificReservation':
+        return cls(count=int(d['count']), in_use_count=int(d['inUseCount']))
+
+
+class GCPReservation:
+    """
+    GCP Reservation object that contains the reservation information.
+    """
+
+    def __init__(self, self_link: str, zone: str,
+                 specific_reservation: SpecificReservation,
+                 specific_reservation_required: bool) -> None:
+        self.self_link = self_link
+        self.zone = zone
+        self.specific_reservation = specific_reservation
+        self.specific_reservation_required = specific_reservation_required
+
+    @classmethod
+    def from_dict(cls, d: dict) -> 'GCPReservation':
+        return cls(
+            self_link=d['selfLink'],
+            zone=d['zone'],
+            specific_reservation=SpecificReservation.from_dict(
+                d['specificReservation']),
+            specific_reservation_required=d['specificReservationRequired'],
+        )
+
+    @property
+    def available_resources(self) -> int:
+        """
+        Count the resources available that can be used in this reservation.
+        """
+        return (self.specific_reservation.count -
+                self.specific_reservation.in_use_count)
+
+    def is_consumable(
+        self,
+        specific_reservations: Set[str],
+    ) -> bool:
+        """
+        Check if the reservation is consumable with the provided specific
+        reservation names. This is defined by the Consumption type.
+        For more details:
+        https://cloud.google.com/compute/docs/instances/reservations-overview#how-reservations-work
+        """
+        return (not self.specific_reservation_required or
+                self.name in specific_reservations)
+
+    @property
+    def name(self) -> str:
+        """Name is derived from reservation self link.
+        The naming convention can be found here:
+        https://cloud.google.com/compute/docs/instances/reservations-consume#consuming_a_specific_shared_reservation
+        """
+        parts = self.self_link.split('/')
+        return '/'.join(parts[-6:-4] + parts[-2:])
+
+
 @clouds.CLOUD_REGISTRY.register
 class GCP(clouds.Cloud):
     """Google Cloud Platform."""
@@ -142,6 +211,12 @@ class GCP(clouds.Cloud):
         'https://skypilot.readthedocs.io/en/latest/getting-started/installation.html#google-cloud-platform-gcp'  # pylint: disable=line-too-long
     )
 
+    def __init__(self):
+        super().__init__()
+
+        self._list_reservations_cache = cachetools.TTLCache(
+            maxsize=1, ttl=datetime.timedelta(300), timer=datetime.datetime.now)
+
     @classmethod
     def _cloud_unsupported_features(
             cls) -> Dict[clouds.CloudImplementationFeatures, str]:
@@ -524,6 +599,65 @@ class GCP(clouds.Cloud):
         return service_catalog.get_vcpus_mem_from_instance_type(instance_type,
                                                                 clouds='gcp')
 
+    def get_reservations_available_resources(
+        self,
+        instance_type: str,
+        region: str,
+        zone: Optional[str],
+        specific_reservations: Set[str],
+    ) -> Dict[str, int]:
+        del region  # Unused
+        assert zone is not None, 'GCP requires zone to get available reservations.'
+        reservations = self._list_reservations_for_instance_type_in_zone(
+            instance_type, zone)
+
+        return {
+            r.name: r.available_resources
+            for r in reservations
+            if r.is_consumable(specific_reservations)
+        }
+
+    def _list_reservations_for_instance_type_in_zone(
+        self,
+        instance_type: str,
+        zone: str,
+    ) -> List[GCPReservation]:
+        reservations = self._list_reservations_for_instance_type(instance_type)
+        return [r for r in reservations if r.zone.endswith(f'/{zone}')]
+
+    @cachetools.cachedmethod(
+        cache=operator.attrgetter('_list_reservations_cache'))
+    def _list_reservations_for_instance_type(
+        self,
+        instance_type: str,
+    ) -> List[GCPReservation]:
+        """
+        List all reservations for the given instance type.
+        TODO: We need to incorporate accelerators because the reserved instance
+        can be consumed only when the instance_type + GPU type matches, and in
+        GCP GPUs except for A100 and L4 do not have their own instance type.
+        For example, if we have a specific reservation with n1-highmem-8
+        in us-central1-c. `sky launch --gpus V100` will fail.
+        """
+        list_reservations_cmd = (
+            'gcloud compute reservations list '
+            f'--filter="specificReservation.instanceProperties.machineType={instance_type} AND status=READY" '
+            '--format="json(specificReservation.count, specificReservation.inUseCount, specificReservationRequired, selfLink, zone)"'
+        )
+        returncode, stdout, stderr = subprocess_utils.run_with_retries(
+            list_reservations_cmd,
+            retry_returncode=[255],
+        )
+        subprocess_utils.handle_returncode(
+            returncode,
+            list_reservations_cmd,
+            error_msg=
+            f'Failed to get list reservations for {instance_type!r}:\n{stderr}',
+            stderr=stderr,
+            stream_logs=True,
+        )
+        return [GCPReservation.from_dict(r) for r in json.loads(stdout)]
+
     @classmethod
     def _find_application_key_path(cls) -> str:
         # Check the application default credentials in the environment variable.
diff --git a/sky/clouds/service_catalog/common.py b/sky/clouds/service_catalog/common.py
index 4491e3391..456d0f623 100644
--- a/sky/clouds/service_catalog/common.py
+++ b/sky/clouds/service_catalog/common.py
@@ -12,6 +12,7 @@ import pandas as pd
 
 from sky import sky_logging
 from sky.clouds import cloud as cloud_lib
+from sky.clouds import cloud_registry
 from sky.clouds.service_catalog import constants
 from sky.utils import log_utils
 from sky.utils import ux_utils
@@ -69,7 +70,7 @@ def read_catalog(filename: str,
     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))
+    cloud = cloud_registry.CLOUD_REGISTRY.from_str(os.path.dirname(filename))
 
     meta_path = os.path.join(_CATALOG_DIR, '.meta', filename)
     os.makedirs(os.path.dirname(meta_path), exist_ok=True)
diff --git a/sky/optimizer.py b/sky/optimizer.py
index 6d7a4b0af..8d8e5682d 100644
--- a/sky/optimizer.py
+++ b/sky/optimizer.py
@@ -14,6 +14,7 @@ from sky import exceptions
 from sky import global_user_state
 from sky import resources as resources_lib
 from sky import sky_logging
+from sky import skypilot_config
 from sky import task as task_lib
 from sky.backends import backend_utils
 from sky.utils import env_options
@@ -231,6 +232,8 @@ class Optimizer:
 
         # node -> cloud -> list of resources that satisfy user's requirements.
         node_to_candidate_map: _TaskToPerCloudCandidates = {}
+        specific_reservations = set(
+            skypilot_config.get_nested(('gcp', 'specific_reservations'), set()))
 
         # Compute the estimated cost/time for each node.
         for node_i, node in enumerate(topo_order):
@@ -258,6 +261,7 @@ class Optimizer:
                 launchable_resources = {list(node_resources)[0]: node_resources}
 
             num_resources = len(node.get_resources())
+
             for orig_resources, launchable_list in launchable_resources.items():
                 if not launchable_list:
                     location_hint = ''
@@ -303,7 +307,17 @@ class Optimizer:
 
                     if minimize_cost:
                         cost_per_node = resources.get_cost(estimated_runtime)
-                        estimated_cost_or_time = cost_per_node * node.num_nodes
+                        num_available_reserved_nodes = sum(
+                            resources.get_reservations_available_resources(
+                                specific_reservations).values())
+
+                        # We consider the cost of the unused reservation
+                        # resources to be 0 since we are already paying for
+                        # them.
+                        # TODO: different policies can be applied here for
+                        # whether to choose reserved instances.
+                        estimated_cost_or_time = cost_per_node * max(
+                            node.num_nodes - num_available_reserved_nodes, 0)
                     else:
                         # Minimize run time.
                         estimated_cost_or_time = estimated_runtime
@@ -886,6 +900,9 @@ def _make_launchables_for_valid_region_zones(
     # on-demand instances in the same region regardless of the zones. On the
     # other hand, for spot instances, we do not batch the requests because the
     # "AWS" spot prices may vary across zones.
+    # For GCP, we do not batch the requests because GCP reservation system is
+    # zone based. Therefore, price estimation is potentially different across
+    # zones.
 
     # NOTE(woosuk): GCP does not support region-level provisioning APIs. Thus,
     # while we return per-region resources here, the provisioner will still
@@ -899,7 +916,8 @@ def _make_launchables_for_valid_region_zones(
     launchables = []
     regions = launchable_resources.get_valid_regions_for_launchable()
     for region in regions:
-        if launchable_resources.use_spot and region.zones is not None:
+        if (launchable_resources.use_spot and region.zones is not None or
+                isinstance(launchable_resources.cloud, clouds.GCP)):
             # Spot instances.
             # Do not batch the per-zone requests.
             for zone in region.zones:
@@ -1004,5 +1022,4 @@ def _fill_in_launchable_resources(
 
         launchable[resources] = _filter_out_blocked_launchable_resources(
             launchable[resources], blocked_resources)
-
     return launchable, cloud_candidates
diff --git a/sky/resources.py b/sky/resources.py
index 87f56fcc4..5f64fbc2b 100644
--- a/sky/resources.py
+++ b/sky/resources.py
@@ -801,6 +801,13 @@ class Resources:
                 self.accelerators, self.use_spot, self._region, self._zone)
         return hourly_cost * hours
 
+    def get_reservations_available_resources(
+            self, specific_reservations: Set[str]) -> Dict[str, int]:
+        """Returns the number of available reservation resources."""
+        return self.cloud.get_reservations_available_resources(
+            self._instance_type, self._region, self._zone,
+            specific_reservations)
+
     def less_demanding_than(self,
                             other: Union[List['Resources'], 'Resources'],
                             requested_num_nodes: int = 1) -> bool:
diff --git a/sky/skypilot_config.py b/sky/skypilot_config.py
index 7aeb3c9b8..6e1f04e41 100644
--- a/sky/skypilot_config.py
+++ b/sky/skypilot_config.py
@@ -48,7 +48,7 @@ from typing import Any, Dict, Iterable
 import yaml
 
 from sky import sky_logging
-from sky import clouds
+from sky.clouds import cloud_registry
 from sky.utils import common_utils
 
 # The config path is discovered in this order:
@@ -160,7 +160,7 @@ def _try_load_config() -> None:
         except yaml.YAMLError as e:
             logger.error(f'Error in loading config file ({config_path}):', e)
 
-        for cloud in clouds.CLOUD_REGISTRY:
+        for cloud in cloud_registry.CLOUD_REGISTRY:
             _syntax_check_for_ssh_proxy_command(cloud)
         logger.debug('Config syntax check passed.')
 
diff --git a/sky/templates/gcp-ray.yml.j2 b/sky/templates/gcp-ray.yml.j2
index 27252efd6..c3b0367a7 100644
--- a/sky/templates/gcp-ray.yml.j2
+++ b/sky/templates/gcp-ray.yml.j2
@@ -43,6 +43,12 @@ available_node_types:
     node_config:
       labels:
         skypilot-user: {{ user }}
+      {%- if specific_reservations %}
+      reservationAffinity:
+        consumeReservationType: SPECIFIC_RESERVATION
+        key: compute.googleapis.com/reservation-name
+        values: {{ specific_reservations }}
+      {%- endif %}
 {%- if tpu_vm %}
       acceleratorType: {{tpu_type}}
       runtimeVersion: {{runtime_version}}
@@ -96,10 +102,10 @@ available_node_types:
         - onHostMaintenance: TERMINATE  # Required for GPU-attached VMs.
   {%- endif %}
 {%- endif %}
-{% if num_nodes > 1 %}
+{% if num_nodes - 1 - num_specific_reserved_workers > 0 %}
   ray_worker_default:
-    min_workers: {{num_nodes - 1}}
-    max_workers: {{num_nodes - 1}}
+    min_workers: {{num_nodes - 1 - num_specific_reserved_workers}}
+    max_workers: {{num_nodes - 1 - num_specific_reserved_workers}}
     resources: {}
     node_config:
       labels:
@@ -159,6 +165,73 @@ available_node_types:
   {%- endif %}
 {%- endif %}
 
+{% if num_specific_reserved_workers > 0 %}
+  ray_worker_reserved:
+    min_workers: {{num_specific_reserved_workers}}
+    max_workers: {{num_specific_reserved_workers}}
+    resources: {}
+    node_config:
+      labels:
+        skypilot-user: {{ user }}
+      reservationAffinity:
+        consumeReservationType: SPECIFIC_RESERVATION
+        key: compute.googleapis.com/reservation-name
+        values: {{ specific_reservations }}
+  {%- if tpu_vm %}
+      acceleratorType: {{tpu_type}}
+      runtimeVersion: {{runtime_version}}
+      metadata:
+        # TPU VM's metadata has different format than normal VMs.
+        # After replacing the variables, this will become username:ssh_public_key_content.
+        # This is a specific syntax required by GCP https://cloud.google.com/compute/docs/connect/add-ssh-keys
+        ssh-keys: |
+              skypilot:ssh_user:skypilot:ssh_public_key_content
+    {%- if use_spot %}
+      schedulingConfig:
+        preemptible: true
+    {%- endif %}
+  {%- else %}
+      machineType: {{instance_type}}
+    {%- if machine_image is not none %}
+      sourceMachineImage: {{machine_image}}
+    {%- endif %}
+      disks:
+        - boot: true
+          autoDelete: true
+          type: PERSISTENT
+          initializeParams:
+            diskSizeGb: {{disk_size}}
+            # See https://cloud.google.com/deep-learning-vm/docs/images
+          {%- if machine_image is none %}
+            sourceImage: {{image_id}}
+          {%- endif %}
+            diskType: zones/{{zones}}/diskTypes/{{disk_tier}}
+    {%- if gpu is not none %}
+      guestAccelerators:
+        - acceleratorType: projects/{{gcp_project_id}}/zones/{{zones}}/acceleratorTypes/{{gpu}}
+          acceleratorCount: {{gpu_count}}
+    {%- endif %}
+      metadata:
+        items:
+          - key: ssh-keys
+            # After replacing the variables, this will become username:ssh_public_key_content.
+            # This is a specific syntax required by GCP https://cloud.google.com/compute/docs/connect/add-ssh-keys
+            value: |
+              skypilot:ssh_user:skypilot:ssh_public_key_content
+    {%- if gpu is not none %}
+          - key: install-nvidia-driver
+            value: "True"
+    {%- endif %}
+      scheduling:
+    {%- if use_spot %}
+        - preemptible: true
+    {%- endif %}
+    {%- if gpu is not none %}
+        - onHostMaintenance: TERMINATE  # Required for GPU-attached VMs.
+    {%- endif %}
+  {%- endif %}
+{%- endif %}
+
 head_node_type: ray_head_default
 
 # Format: `REMOTE_PATH : LOCAL_PATH`
diff --git a/sky/utils/schemas.py b/sky/utils/schemas.py
index 21fc7b421..597c51a63 100644
--- a/sky/utils/schemas.py
+++ b/sky/utils/schemas.py
@@ -4,7 +4,7 @@ Schemas conform to the JSON Schema specification as defined at
 https://json-schema.org/
 """
 
-from sky.clouds import cloud
+from sky.clouds import cloud_registry
 from sky.data import storage
 
 
@@ -17,7 +17,8 @@ def get_resources_schema():
         'properties': {
             'cloud': {
                 'type': 'string',
-                'case_insensitive_enum': list(cloud.CLOUD_REGISTRY.keys())
+                'case_insensitive_enum': list(
+                    cloud_registry.CLOUD_REGISTRY.keys())
             },
             'region': {
                 'type': 'string',


Test patch:
diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml
index 81a43cb49..77020384f 100644
--- a/.github/workflows/pytest.yml
+++ b/.github/workflows/pytest.yml
@@ -16,6 +16,7 @@ jobs:
       matrix:
         python-version: [3.8]
         test-path:
+          - tests/unit_tests
           - tests/test_cli.py
           - tests/test_config.py
           - tests/test_global_user_state.py
diff --git a/tests/conftest.py b/tests/conftest.py
index 299ec455e..d6f194702 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,6 +1,7 @@
 import pytest
 import tempfile
 from typing import List
+from unittest.mock import patch
 
 import pandas as pd
 
@@ -179,7 +180,7 @@ def generic_cloud(request) -> str:
 
 
 @pytest.fixture
-def enable_all_clouds(monkeypatch):
+def enable_all_clouds(monkeypatch, patch_gcloud_list_reservations):
     from sky import clouds
     # Monkey-patching is required because in the test environment, no cloud is
     # enabled. The optimizer checks the environment to find enabled clouds, and
@@ -226,3 +227,10 @@ def aws_config_region(monkeypatch) -> str:
         if isinstance(ssh_proxy_command, dict) and ssh_proxy_command:
             region = list(ssh_proxy_command.keys())[0]
     return region
+
+
+@pytest.fixture
+def patch_gcloud_list_reservations():
+    with patch('sky.clouds.gcp.GCP._list_reservations_for_instance_type',
+               return_value=[]) as mock:
+        yield mock
diff --git a/tests/test_jobs.py b/tests/test_jobs.py
index d61a851c4..970314548 100644
--- a/tests/test_jobs.py
+++ b/tests/test_jobs.py
@@ -21,7 +21,7 @@ class TestExecutionOnExistingClusters:
     @pytest.fixture
     def _mock_cluster_state(self, _mock_db_conn, enable_all_clouds):
         """Add clusters to the global state.
-        
+
         This fixture adss three clusters to the global state:
         - test-cluster1: AWS, 2x p4d.24xlarge (8x A100)
         - test-cluster2: GCP, 1x n1-highmem-64, 4x V100
@@ -74,7 +74,7 @@ class TestExecutionOnExistingClusters:
 
     def test_launch_exec(self, _mock_cluster_state, monkeypatch):
         """Test launch and exec on existing clusters.
-        
+
         This test runs launch and exec with less demanding resources
         than the existing clusters can pass the check.
         """
diff --git a/tests/test_optimizer_dryruns.py b/tests/test_optimizer_dryruns.py
index 8bfa4944a..4a82cbad6 100644
--- a/tests/test_optimizer_dryruns.py
+++ b/tests/test_optimizer_dryruns.py
@@ -118,23 +118,23 @@ def test_resources_azure(monkeypatch):
     _test_resources_launch(monkeypatch, sky.Azure(), 'Standard_NC24s_v3')
 
 
-def test_resources_gcp(monkeypatch):
+def test_resources_gcp(monkeypatch, patch_gcloud_list_reservations):
     _test_resources_launch(monkeypatch, sky.GCP(), 'n1-standard-16')
 
 
-def test_partial_cpus(monkeypatch):
+def test_partial_cpus(monkeypatch, patch_gcloud_list_reservations):
     _test_resources_launch(monkeypatch, cpus=4)
     _test_resources_launch(monkeypatch, cpus='4')
     _test_resources_launch(monkeypatch, cpus='7+')
 
 
-def test_partial_memory(monkeypatch):
+def test_partial_memory(monkeypatch, patch_gcloud_list_reservations):
     _test_resources_launch(monkeypatch, memory=32)
     _test_resources_launch(monkeypatch, memory='32')
     _test_resources_launch(monkeypatch, memory='32+')
 
 
-def test_partial_k80(monkeypatch):
+def test_partial_k80(monkeypatch, patch_gcloud_list_reservations):
     _test_resources_launch(monkeypatch, accelerators='K80')
 
 
@@ -142,16 +142,16 @@ def test_partial_m60(monkeypatch):
     _test_resources_launch(monkeypatch, accelerators='M60')
 
 
-def test_partial_p100(monkeypatch):
+def test_partial_p100(monkeypatch, patch_gcloud_list_reservations):
     _test_resources_launch(monkeypatch, accelerators='P100')
 
 
-def test_partial_t4(monkeypatch):
+def test_partial_t4(monkeypatch, patch_gcloud_list_reservations):
     _test_resources_launch(monkeypatch, accelerators='T4')
     _test_resources_launch(monkeypatch, accelerators={'T4': 8}, use_spot=True)
 
 
-def test_partial_tpu(monkeypatch):
+def test_partial_tpu(monkeypatch, patch_gcloud_list_reservations):
     _test_resources_launch(monkeypatch, accelerators='tpu-v3-8')
 
 
@@ -224,7 +224,8 @@ def test_instance_type_mismatches_memory(monkeypatch):
         assert 'does not have the requested memory' in str(e.value)
 
 
-def test_instance_type_matches_cpus(monkeypatch):
+def test_instance_type_matches_cpus(monkeypatch,
+                                    patch_gcloud_list_reservations):
     _test_resources_launch(monkeypatch,
                            sky.AWS(),
                            instance_type='c6i.8xlarge',
@@ -243,7 +244,8 @@ def test_instance_type_matches_cpus(monkeypatch):
                            cpus=8.0)
 
 
-def test_instance_type_matches_memory(monkeypatch):
+def test_instance_type_matches_memory(monkeypatch,
+                                      patch_gcloud_list_reservations):
     _test_resources_launch(monkeypatch,
                            sky.AWS(),
                            instance_type='c6i.8xlarge',
@@ -262,7 +264,8 @@ def test_instance_type_matches_memory(monkeypatch):
                            memory=32)
 
 
-def test_instance_type_from_cpu_memory(monkeypatch, capfd):
+def test_instance_type_from_cpu_memory(monkeypatch, capfd,
+                                       patch_gcloud_list_reservations):
     _test_resources_launch(monkeypatch, cpus=8)
     stdout, _ = capfd.readouterr()
     # Choose General Purpose instance types
@@ -358,7 +361,8 @@ def test_instance_type_mistmatches_accelerators(monkeypatch):
         assert 'Infeasible resource demands found' in str(e.value)
 
 
-def test_instance_type_matches_accelerators(monkeypatch):
+def test_instance_type_matches_accelerators(monkeypatch,
+                                            patch_gcloud_list_reservations):
     _test_resources_launch(monkeypatch,
                            sky.AWS(),
                            instance_type='p3.2xlarge',
diff --git a/tests/unit_tests/sky/clouds/test_cloud.py b/tests/unit_tests/sky/clouds/test_cloud.py
new file mode 100644
index 000000000..e611b5459
--- /dev/null
+++ b/tests/unit_tests/sky/clouds/test_cloud.py
@@ -0,0 +1,14 @@
+import pytest
+
+from sky.clouds.cloud import Cloud
+
+
+@pytest.mark.parametrize(("specific_reservations", "expected"), [({"a"}, {
+    "a": 0
+}), ((set(), {}))])
+def test_cloud_get_reservations_available_resources(specific_reservations,
+                                                    expected):
+
+    available_resources = Cloud().get_reservations_available_resources(
+        "instance_type", "region", "zone", specific_reservations)
+    assert available_resources == expected
diff --git a/tests/unit_tests/sky/clouds/test_gcp.py b/tests/unit_tests/sky/clouds/test_gcp.py
new file mode 100644
index 000000000..5d8136296
--- /dev/null
+++ b/tests/unit_tests/sky/clouds/test_gcp.py
@@ -0,0 +1,103 @@
+from unittest.mock import patch
+from sky.clouds.gcp import GCP, GCPReservation, SpecificReservation
+import pytest
+
+
+@pytest.mark.parametrize((
+    'mock_return', 'expected'
+), [([
+    GCPReservation(
+        self_link=
+        'https://www.googleapis.com/compute/v1/projects/<project>/zones/<zone>/reservations/<reservation>',
+        specific_reservation=SpecificReservation(count=1, in_use_count=0),
+        specific_reservation_required=True,
+        zone='zone')
+], {
+    'projects/<project>/reservations/<reservation>': 1
+}),
+    ([
+        GCPReservation(
+            self_link=
+            'https://www.googleapis.com/compute/v1/projects/<project>/zones/<zone>/reservations/<reservation>',
+            specific_reservation=SpecificReservation(count=2, in_use_count=1),
+            specific_reservation_required=False,
+            zone='zone')
+    ], {
+        'projects/<project>/reservations/<reservation>': 1
+    }),
+    ([
+        GCPReservation(
+            self_link=
+            'https://www.googleapis.com/compute/v1/projects/<project2>/zones/<zone>/reservations/<reservation>',
+            specific_reservation=SpecificReservation(count=1, in_use_count=0),
+            specific_reservation_required=True,
+            zone='zone')
+    ], {})])
+def test_gcp_get_reservations_available_resources(mock_return, expected):
+    gcp = GCP()
+    with patch.object(gcp,
+                      '_list_reservations_for_instance_type_in_zone',
+                      return_value=mock_return):
+        reservations = gcp.get_reservations_available_resources(
+            'instance_type', 'region', 'zone',
+            {'projects/<project>/reservations/<reservation>'})
+        assert reservations == expected
+
+
+def test_gcp_reservation_from_dict():
+    r = GCPReservation.from_dict({
+        'selfLink': 'test',
+        'specificReservation': {
+            'count': '1',
+            'inUseCount': '0'
+        },
+        'specificReservationRequired': True,
+        'zone': 'zone'
+    })
+
+    assert r.self_link == 'test'
+    assert r.specific_reservation.count == 1
+    assert r.specific_reservation.in_use_count == 0
+    assert r.specific_reservation_required == True
+    assert r.zone == 'zone'
+
+
+@pytest.mark.parametrize(('count', 'in_use_count', 'expected'), [(1, 0, 1),
+                                                                 (1, 1, 0)])
+def test_gcp_reservation_available_resources(count, in_use_count, expected):
+    r = GCPReservation(self_link='test',
+                       specific_reservation=SpecificReservation(
+                           count=count, in_use_count=in_use_count),
+                       specific_reservation_required=True,
+                       zone='zone')
+
+    assert r.available_resources == expected
+
+
+def test_gcp_reservation_name():
+    r = GCPReservation(
+        self_link=
+        'https://www.googleapis.com/compute/v1/projects/<project>/zones/<zone>/reservations/<reservation-name>',
+        specific_reservation=SpecificReservation(count=1, in_use_count=1),
+        specific_reservation_required=True,
+        zone='zone')
+    assert r.name == 'projects/<project>/reservations/<reservation-name>'
+
+
+@pytest.mark.parametrize(
+    ('specific_reservations', 'specific_reservation_required', 'expected'), [
+        ([], False, True),
+        ([], True, False),
+        (['projects/<project>/reservations/<reservation>'], True, True),
+        (['projects/<project>/reservations/<invalid>'], True, False),
+    ])
+def test_gcp_reservation_is_consumable(specific_reservations,
+                                       specific_reservation_required, expected):
+    r = GCPReservation(
+        self_link=
+        'https://www.googleapis.com/compute/v1/projects/<project>/zones/<zone>/reservations/<reservation>',
+        specific_reservation=SpecificReservation(count=1, in_use_count=1),
+        specific_reservation_required=specific_reservation_required,
+        zone='zone')
+    assert r.is_consumable(
+        specific_reservations=specific_reservations) is expected
diff --git a/tests/unit_tests/sky/test_resources.py b/tests/unit_tests/sky/test_resources.py
new file mode 100644
index 000000000..ae458aa2b
--- /dev/null
+++ b/tests/unit_tests/sky/test_resources.py
@@ -0,0 +1,12 @@
+from sky.resources import Resources
+from unittest.mock import Mock
+
+
+def test_get_reservations_available_resources():
+    mock = Mock()
+    r = Resources(cloud=mock, instance_type="instance_type")
+    r._region = "region"
+    r._zone = "zone"
+    r.get_reservations_available_resources(set())
+    mock.get_reservations_available_resources.assert_called_once_with(
+        "instance_type", "region", "zone", set())