Inspect node group reservation, use attached placement policy

Author: mr0re1

CHANGELOG.md

@@ -2,6 +2,11 @@
 
 All notable changes to this project will be documented in this file.
 
+## \[5.9.1\]
+
+- Use reservation placement policy if placement is enabled, and a reservation is
+  specified.
+
 ## \[5.9.0\]
 
 - Remove spurious log message on resume, referring to "Reservation name".

scripts/resume.py

@@ -125,11 +125,21 @@ def instance_properties(partition, model, placement_group, labels=None):
         }
 
     if node_group.reservation_name:
+        reservation_name = node_group.reservation_name
+        reservation = lkp.reservation(reservation_name)
         props.reservationAffinity = {
             "consumeReservationType": "SPECIFIC_RESERVATION",
             "key": "compute.googleapis.com/reservation-name",
-            "values": [node_group.reservation_name],
+            "values": [reservation_name],
         }
+        
+        placement = util.reservation_placement(reservation)
+        if placement:
+            props.resourcePolicies = [placement]
+            log.info(f"reservation {reservation_name} is being used with placement policy {placement}")
+        else:
+            props.resourcePolicies = []
+            log.info(f"reservation {reservation_name} is being used without a placement policy")
 
     return props
 

scripts/util.py

@@ -386,6 +386,16 @@ def is_exclusive_node(node):
     )
 
 
+def reservation_placement(reservation):
+    """
+    Inspects reservation object, returns placement policy name if one is found.
+    """
+    url = reservation.get("resourcePolicies", {}).get("placement")
+    if url:
+        return trim_self_link(url)
+    return None
+
+
 def compute_service(credentials=None, user_agent=USER_AGENT, version="v1"):
     """Make thread-safe compute service handle
     creates a new Http for each request
@@ -1629,6 +1639,29 @@ def clear_template_info_cache(self):
             cache.clear()
         self.template_info.cache_clear()
 
+    @lru_cache()
+    def reservation(self, name):
+        """
+        See https://cloud.google.com/compute/docs/reference/rest/v1/reservations
+        """
+        resp = ensure_execute(
+            self.compute.reservations().aggregatedList(
+                project=self.project, filter=f"name={name}"
+            )
+        )
+        reservation = None
+        for _, e in resp["items"].items():
+            for r in e.get("reservations", []):
+                assert (
+                    reservation is None
+                ), f"multiple reservations '{name}' found in '{self.project}'."
+                reservation = r
+
+        assert (
+            reservation is not None
+        ), f"reservation '{name}' not found in '{self.project}'."
+        return reservation
+
 
 # Define late globals
 cfg = load_config_file(CONFIG_FILE)