Support DWS canlendar with regular reservation

Author: jemish-google

community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/resume.py

@@ -105,7 +105,6 @@ def instance_properties(nodeset: NSDict, model:str, placement_group:Optional[str
         props.resourcePolicies = [placement_group]
 
     if reservation := lookup().nodeset_reservation(nodeset):
-
         update_reservation_props(reservation, props, placement_group, reservation.calendar)
 
     if (fr := lookup().future_reservation(nodeset)) and fr.specific:
@@ -295,12 +294,12 @@ def chunk_nodes(nodes: List[str]):
 def resume_nodes(nodes: List[str], resume_data: Optional[ResumeData]):
     """resume nodes in nodelist"""
     lkp = lookup()
-    # Prevent dormant nodes associated with a future reservation from being resumed
-    nodes, dormant_fr_nodes = util.separate(lkp.is_dormant_fr_node, nodes)
+    # Prevent dormant nodes associated with a reservation from being resumed
+    nodes, dormant_res_nodes = util.separate(lkp.is_dormant_res_node, nodes)
 
-    if dormant_fr_nodes:
-        log.warning(f"Resume was unable to resume future reservation nodes={dormant_fr_nodes}")
-        down_nodes_notify_jobs(dormant_fr_nodes, "Reservation is not active, nodes cannot be resumed", resume_data)
+    if dormant_res_nodes:
+        log.warning(f"Resume was unable to resume reservation nodes={dormant_res_nodes}")
+        down_nodes_notify_jobs(dormant_res_nodes, "Reservation is not active, nodes cannot be resumed", resume_data)
 
     nodes, flex_managed = util.separate(lkp.is_provisioning_flex_node, nodes)
     if flex_managed:

community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/tests/test_util.py

@@ -198,7 +198,7 @@ def test_nodeset_reservation_err(nodeset, err):
                     policies=[],
                     deployment_type=None,
                     reservation_mode=None,
-                    assured_count=None,
+                    assured_count=0,
                     delete_at_time=None,
                     bulk_insert_name="projects/bobin/reservations/robin")),
             (TstNodeset(
@@ -212,7 +212,7 @@ def test_nodeset_reservation_err(nodeset, err):
                     policies=["wanders", "apples", "yum"],
                     deployment_type=None,
                     reservation_mode=None,
-                    assured_count=None,
+                    assured_count=0,
                     delete_at_time=None,
                     bulk_insert_name="projects/bobin/reservations/robin")),
             (TstNodeset(
@@ -226,7 +226,7 @@ def test_nodeset_reservation_err(nodeset, err):
                     policies=[],
                     deployment_type=None,
                     reservation_mode=None,
-                    assured_count=1,
+                    assured_count=0,
                     delete_at_time=None,
                     bulk_insert_name="projects/bobin/reservations/robin/snek/cheese-brie-6")),
 
@@ -555,6 +555,8 @@ def test_future_reservation_active(_):
                 name='melon',
                 policies=[],
                 reservation_mode=None,
+                assured_count=0,
+                delete_at_time=None,
                 bulk_insert_name="projects/manhattan/reservations/melon",
                 deployment_type=None))
 
@@ -581,7 +583,7 @@ def test_future_reservation_inactive(_):
         TstNodeset(future_reservation="projects/manhattan/zones/danger/futureReservations/zebra")) == FutureReservation(
             project='manhattan', 
             zone='danger', 
-            name='zebra',
+            name='zebra', 
             specific=True, 
             start_time=datetime(2025, 1, 27, 23, 30, tzinfo=timezone.utc), 
             end_time=datetime(2025, 2, 21, 23, 30, tzinfo=timezone.utc), 

community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/util.py

@@ -1454,6 +1454,8 @@ class ReservationDetails:
     bulk_insert_name: str # name in format suitable for bulk insert (currently identical to user supplied name in long format)
     deployment_type: Optional[str]
     reservation_mode: Optional[str]
+    assured_count: int 
+    delete_at_time: Optional[datetime]
 
     @property
     def dense(self) -> bool:
@@ -1621,11 +1623,29 @@ def nodeset_is_tpu(self, nodeset_name=None) -> bool:
     def node_is_fr(self, node_name:str) -> bool:
         return bool(self.node_nodeset(node_name).future_reservation)
 
-    def is_dormant_fr_node(self, node_name:str) -> bool:
+    def is_dormant_res_node(self, node_name:str) -> bool:
         fr = self.future_reservation(self.node_nodeset(node_name))
-        if not fr:
+        res = self.nodeset_reservation(self.node_nodeset(node_name))
+        
+        if fr is None and res is None:
             return False
-        return fr.active_reservation is None
+        
+        if fr:
+            return fr.active_reservation is None
+        
+        if res:
+            if res.calendar:
+                # If reservation is calendar based, check if it is past the delete_at_time
+                if res.delete_at_time is not None and now() >= res.delete_at_time:
+                    log.info(f"DWS calendar reservation {res.bulk_insert_name} is past deletion time {res.delete_at_time}, skipping resume.")
+                    return True
+
+                # If assured_count is 0 do not resume nodes as they are not active yet
+                if  res.delete_at_time is not None and res.assured_count <= 0:
+                    log.info(f"DWS calendar reservation {res.bulk_insert_name} is not active yet, skipping resume.")
+                    return True
+                
+        return False    
 
     def node_is_dyn(self, node_name=None) -> bool:
         nodeset = self.node_nodeset_name(node_name)
@@ -1824,6 +1844,8 @@ def get_reservation_details(self, project:str, zone:str, name:str, bulk_insert_n
             policies=policies,
             deployment_type=reservation.get("deploymentType"),
             reservation_mode=reservation.get("reservationMode"),
+            assured_count=reservation.get("assuredCount") if reservation.get("assuredCount") else 0,
+            delete_at_time=parse_gcp_timestamp(reservation.get("deleteAtTime")) if reservation.get("deleteAtTime") else None,
             bulk_insert_name=bulk_insert_name)
 
     def nodeset_reservation(self, nodeset: NSDict) -> Optional[ReservationDetails]: