Repository: skypilot-org/skypilot
Pull request #2836 (2023-12-03 08:54:28)

[GCP] Refactor the reserved instances cache
<!-- Describe the changes in this PR -->

The previous implementation of GCP's reserved instance cache has different cache for multiple GCP() cloud objects. This can cause long optimization time if multiple resources are specified with different GCP() objects.

This PR refactors the cache out to a module to avoid querying the reservation list multiple times.

It reduces the time for optimization from 13 seconds to 4 seconds for the following yaml:
```yaml
resources:
    # accelerators: A100:8
    any_of:
      # AWS:
      - region: us-east-1
      - region: us-east-2
      - region: us-west-1
      - region: us-west-2
      # GCP
      - region: us-central1
      - region: us-east1
      - region: us-east4
      - region: us-west1
      - region: us-west2
      - region: us-west3
      - region: us-west4
```
`sky launch -c test test.yaml`

<!-- 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`
- [x] Any manual or new tests for this PR (please specify below)
  - [x] `sky launch --cloud gcp -t n2-standard-2 echo hi` with reservation for `n2-standard-2`
  - [x] `sky launch -c test-gcp --cloud gcp echo hi` on master
  - [x] `sky exec test-gcp echo hi` on current PR
- [ ] 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`


concretevitamin: We may want to add a comment to `class Cloud` docstr, something like:

Lightweight stateless objects. SkyPilot may create multiple such objects; therefore, subclasses should take care to make methods inexpensive to call, and should not store heavy state. If state needs to be queried from the cloud provider and cached, create a module in sky/clouds/utils/ so they can be reused across cloud object creation.

Patch:
diff --git a/sky/clouds/cloud.py b/sky/clouds/cloud.py
index 625694f87..98f0048b8 100644
--- a/sky/clouds/cloud.py
+++ b/sky/clouds/cloud.py
@@ -1,4 +1,12 @@
-"""Interfaces: clouds, regions, and zones."""
+"""Interfaces: clouds, regions, and zones.
+
+clouds.Cloud is lightweight stateless objects. SkyPilot may create multiple such
+objects; therefore, subclasses should take care to make methods inexpensive to
+call, and should not store heavy state. If state needs to be queried from the
+cloud provider and cached, create a module in sky/clouds/utils/ and probably add
+caches for the return value (e.g., sky/clouds/utils/gcp_utils), so they can be
+reused across cloud object creation.
+"""
 import collections
 import enum
 import re
diff --git a/sky/clouds/gcp.py b/sky/clouds/gcp.py
index edb67a94f..a86946870 100644
--- a/sky/clouds/gcp.py
+++ b/sky/clouds/gcp.py
@@ -1,6 +1,4 @@
 """Google Cloud Platform."""
-import dataclasses
-import datetime
 import functools
 import json
 import os
@@ -8,9 +6,8 @@ import re
 import subprocess
 import time
 import typing
-from typing import Any, Dict, Iterator, List, Optional, Set, Tuple
+from typing import Dict, Iterator, List, Optional, Set, Tuple
 
-import cachetools
 import colorama
 
 from sky import clouds
@@ -19,6 +16,7 @@ from sky import sky_logging
 from sky import status_lib
 from sky.adaptors import gcp
 from sky.clouds import service_catalog
+from sky.clouds.utils import gcp_utils
 from sky.skylet import log_lib
 from sky.utils import common_utils
 from sky.utils import subprocess_utils
@@ -120,68 +118,6 @@ 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 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.
-
-        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 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."""
@@ -225,11 +161,6 @@ 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 = None
-
     @classmethod
     def _cloud_unsupported_features(
             cls) -> Dict[clouds.CloudImplementationFeatures, str]:
@@ -633,82 +564,6 @@ 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
-        if zone is None:
-            # For backward compatibility, the cluster in INIT state launched
-            # before #2352 may not have zone information. In this case, we
-            # return 0 for all reservations.
-            return {reservation: 0 for reservation in specific_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}')]
-
-    def _get_or_create_ttl_cache(self):
-        if getattr(self, '_list_reservations_cache', None) is None:
-            # Backward compatibility: Clusters created before #2352 has GCP
-            # objects serialized without the attribute. So we access it this
-            # way.
-            self._list_reservations_cache = cachetools.TTLCache(
-                maxsize=1,
-                ttl=datetime.timedelta(300),
-                timer=datetime.datetime.now)
-        return self._list_reservations_cache
-
-    @cachetools.cachedmethod(cache=lambda self: self._get_or_create_ttl_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,
-            # 1: means connection aborted (although it shows 22 in the error,
-            # but the actual error code is 1)
-            # Example: ERROR: gcloud crashed (ConnectionError): ('Connection aborted.', OSError(22, 'Invalid argument')) # pylint: disable=line-too-long
-            retry_returncode=[255, 1],
-        )
-        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.
@@ -1058,6 +913,28 @@ class GCP(clouds.Cloud):
         # Quota found to be greater than zero, try provisioning
         return True
 
+    def get_reservations_available_resources(
+        self,
+        instance_type: str,
+        region: str,
+        zone: Optional[str],
+        specific_reservations: Set[str],
+    ) -> Dict[str, int]:
+        del region  # Unused
+        if zone is None:
+            # For backward compatibility, the cluster in INIT state launched
+            # before #2352 may not have zone information. In this case, we
+            # return 0 for all reservations.
+            return {reservation: 0 for reservation in specific_reservations}
+        reservations = gcp_utils.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)
+        }
+
     @classmethod
     def query_status(cls, name: str, tag_filters: Dict[str, str],
                      region: Optional[str], zone: Optional[str],
@@ -1227,10 +1104,3 @@ class GCP(clouds.Cloud):
             error_msg=f'Failed to delete image {image_name!r}',
             stderr=stderr,
             stream_logs=True)
-
-    def __getstate__(self) -> Dict[str, Any]:
-        state = self.__dict__.copy()
-        # We should avoid saving third-party object to the state, as it may
-        # cause unpickling error when the third-party API is updated.
-        state.pop('_list_reservations_cache', None)
-        return state
diff --git a/sky/clouds/utils/gcp_utils.py b/sky/clouds/utils/gcp_utils.py
new file mode 100644
index 000000000..0b8256e42
--- /dev/null
+++ b/sky/clouds/utils/gcp_utils.py
@@ -0,0 +1,126 @@
+"""Utility functions for GCP.
+
+The functions that are used to access GCP APIs. We have the reservation-related
+functions here, so that the cache of the reservations can be shared across
+multiple clouds.GCP() objects.
+"""
+
+import dataclasses
+import json
+import time
+from typing import List, Set
+
+import cachetools
+
+from sky import sky_logging
+from sky.utils import subprocess_utils
+
+logger = sky_logging.init_logger(__name__)
+
+
+@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 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.
+
+        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 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:])
+
+
+def list_reservations_for_instance_type_in_zone(
+    instance_type: str,
+    zone: str,
+) -> List[GCPReservation]:
+    reservations = _list_reservations_for_instance_type(instance_type)
+    return [r for r in reservations if r.zone.endswith(f'/{zone}')]
+
+
+@cachetools.cached(cache=cachetools.TTLCache(maxsize=1,
+                                             ttl=300,
+                                             timer=time.time))
+def _list_reservations_for_instance_type(
+    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.
+    """
+    logger.debug(f'Querying GCP reservations for instance {instance_type!r}')
+    list_reservations_cmd = (
+        'gcloud compute reservations list '
+        '--filter="specificReservation.instanceProperties.machineType='
+        f'{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,
+        # 1: means connection aborted (although it shows 22 in the error,
+        # but the actual error code is 1)
+        # Example: ERROR: gcloud crashed (ConnectionError): ('Connection aborted.', OSError(22, 'Invalid argument')) # pylint: disable=line-too-long
+        retry_returncode=[255, 1],
+    )
+    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)]


Test patch:
diff --git a/tests/common.py b/tests/common.py
index 4844def6b..a8a33fa3f 100644
--- a/tests/common.py
+++ b/tests/common.py
@@ -48,7 +48,7 @@ def enable_all_clouds_in_monkeypatch(
                         lambda _: None)
 
     monkeypatch.setattr(
-        'sky.clouds.gcp.GCP._list_reservations_for_instance_type',
+        'sky.clouds.utils.gcp_utils.list_reservations_for_instance_type_in_zone',
         lambda *_args, **_kwargs: [])
 
     # Monkey patch Kubernetes resource detection since it queries
diff --git a/tests/unit_tests/test_gcp.py b/tests/unit_tests/test_gcp.py
index 95db1fef8..8fc224b33 100644
--- a/tests/unit_tests/test_gcp.py
+++ b/tests/unit_tests/test_gcp.py
@@ -3,44 +3,46 @@ from unittest.mock import patch
 import pytest
 
 from sky.clouds.gcp import GCP
-from sky.clouds.gcp import GCPReservation
-from sky.clouds.gcp import SpecificReservation
+from sky.clouds.utils import gcp_utils
 
 
 @pytest.mark.parametrize((
     'mock_return', 'expected'
 ), [([
-    GCPReservation(
+    gcp_utils.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=gcp_utils.SpecificReservation(count=1,
+                                                           in_use_count=0),
         specific_reservation_required=True,
         zone='zone')
 ], {
     'projects/<project>/reservations/<reservation>': 1
 }),
     ([
-        GCPReservation(
+        gcp_utils.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=gcp_utils.SpecificReservation(count=2,
+                                                               in_use_count=1),
             specific_reservation_required=False,
             zone='zone')
     ], {
         'projects/<project>/reservations/<reservation>': 1
     }),
     ([
-        GCPReservation(
+        gcp_utils.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=gcp_utils.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',
+    with patch.object(gcp_utils,
+                      'list_reservations_for_instance_type_in_zone',
                       return_value=mock_return):
         reservations = gcp.get_reservations_available_resources(
             'instance_type', 'region', 'zone',
@@ -49,7 +51,7 @@ def test_gcp_get_reservations_available_resources(mock_return, expected):
 
 
 def test_gcp_reservation_from_dict():
-    r = GCPReservation.from_dict({
+    r = gcp_utils.GCPReservation.from_dict({
         'selfLink': 'test',
         'specificReservation': {
             'count': '1',
@@ -69,20 +71,22 @@ def test_gcp_reservation_from_dict():
 @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')
+    r = gcp_utils.GCPReservation(
+        self_link='test',
+        specific_reservation=gcp_utils.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(
+    r = gcp_utils.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=gcp_utils.SpecificReservation(count=1,
+                                                           in_use_count=1),
         specific_reservation_required=True,
         zone='zone')
     assert r.name == 'projects/<project>/reservations/<reservation-name>'
@@ -97,10 +101,11 @@ def test_gcp_reservation_name():
     ])
 def test_gcp_reservation_is_consumable(specific_reservations,
                                        specific_reservation_required, expected):
-    r = GCPReservation(
+    r = gcp_utils.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=gcp_utils.SpecificReservation(count=1,
+                                                           in_use_count=1),
         specific_reservation_required=specific_reservation_required,
         zone='zone')
     assert r.is_consumable(