Merge pull request GoogleCloudPlatform#4696 from arpit974/dws_flex_start_non_exclusive_partition

Allowing setting use_job_duration with non-exclusive partitions.

Author: arpit974

community/modules/compute/schedmd-slurm-gcp-v6-nodeset/README.md

@@ -174,7 +174,7 @@ modules. For support with the underlying modules, see the instructions in the
 | <a name="input_disk_resource_manager_tags"></a> [disk\_resource\_manager\_tags](#input\_disk\_resource\_manager\_tags) | (Optional) A set of key/value resource manager tag pairs to bind to the instance disks. Keys must be in the format tagKeys/{tag\_key\_id}, and values are in the format tagValues/456. | `map(string)` | `{}` | no |
 | <a name="input_disk_size_gb"></a> [disk\_size\_gb](#input\_disk\_size\_gb) | Size of boot disk to create for the partition compute nodes. | `number` | `50` | no |
 | <a name="input_disk_type"></a> [disk\_type](#input\_disk\_type) | Boot disk type, can be either hyperdisk-balanced, pd-ssd, pd-standard, pd-balanced, or pd-extreme. | `string` | `"pd-standard"` | no |
-| <a name="input_dws_flex"></a> [dws\_flex](#input\_dws\_flex) | If set and `enabled = true`, will utilize the DWS Flex Start to provision nodes.<br/> See: https://cloud.google.com/blog/products/compute/introducing-dynamic-workload-scheduler<br/> Options:<br/> - enable: Enable DWS Flex Start<br/> - max\_run\_duration: Maximum duration in seconds for the job to run, should not exceed 604,800 (one week).<br/> - use\_job\_duration: Use the job duration to determine the max\_run\_duration, if job duration is not set, max\_run\_duration will be used.<br/> - use\_bulk\_insert: Uses the legacy implementation of DWS Flex Start with Bulk Insert for non-accelerator instances<br/><br/>Limitations:<br/> - CAN NOT be used with reservations;<br/> - CAN NOT be used with placement groups;<br/> - If `use_job_duration` is enabled nodeset can be used in "exclusive" partitions only | <pre>object({<br/>    enabled          = optional(bool, true)<br/>    max_run_duration = optional(number, 604800) # one week<br/>    use_job_duration = optional(bool, false)<br/>    use_bulk_insert  = optional(bool, false)<br/>  })</pre> | <pre>{<br/>  "enabled": false<br/>}</pre> | no |
+| <a name="input_dws_flex"></a> [dws\_flex](#input\_dws\_flex) | If set and `enabled = true`, will utilize the DWS Flex Start to provision nodes.<br/> See: https://cloud.google.com/blog/products/compute/introducing-dynamic-workload-scheduler<br/> Options:<br/> - enable: Enable DWS Flex Start<br/> - max\_run\_duration: Maximum duration in seconds for the job to run, should not exceed 604,800 (one week).<br/> - use\_job\_duration: Use the job duration to determine the max\_run\_duration, if job duration is not set, max\_run\_duration will be used.<br/> - use\_bulk\_insert: Uses the legacy implementation of DWS Flex Start with Bulk Insert for non-accelerator instances<br/><br/>Limitations:<br/> - CAN NOT be used with reservations;<br/> - CAN NOT be used with placement groups; | <pre>object({<br/>    enabled          = optional(bool, true)<br/>    max_run_duration = optional(number, 604800) # one week<br/>    use_job_duration = optional(bool, false)<br/>    use_bulk_insert  = optional(bool, false)<br/>  })</pre> | <pre>{<br/>  "enabled": false<br/>}</pre> | no |
 | <a name="input_enable_confidential_vm"></a> [enable\_confidential\_vm](#input\_enable\_confidential\_vm) | Enable the Confidential VM configuration. Note: the instance image must support option. | `bool` | `false` | no |
 | <a name="input_enable_maintenance_reservation"></a> [enable\_maintenance\_reservation](#input\_enable\_maintenance\_reservation) | Enables slurm reservation for scheduled maintenance. | `bool` | `false` | no |
 | <a name="input_enable_opportunistic_maintenance"></a> [enable\_opportunistic\_maintenance](#input\_enable\_opportunistic\_maintenance) | On receiving maintenance notification, maintenance will be performed as soon as nodes becomes idle. | `bool` | `false` | no |

community/modules/compute/schedmd-slurm-gcp-v6-nodeset/variables.tf

@@ -610,7 +610,6 @@ variable "dws_flex" {
  Limitations:
   - CAN NOT be used with reservations;
   - CAN NOT be used with placement groups;
-  - If `use_job_duration` is enabled nodeset can be used in "exclusive" partitions only
 
  EOD
 

community/modules/compute/schedmd-slurm-gcp-v6-partition/main.tf

@@ -13,8 +13,7 @@
 # limitations under the License.
 
 locals {
-  use_static        = [for ns in concat(var.nodeset, var.nodeset_tpu) : ns.nodeset_name if ns.node_count_static > 0]
-  uses_job_duration = length([for ns in var.nodeset : ns.dws_flex.use_job_duration if ns.dws_flex.use_job_duration]) > 0
+  use_static = [for ns in concat(var.nodeset, var.nodeset_tpu) : ns.nodeset_name if ns.node_count_static > 0]
 
   has_node = length(var.nodeset) > 0
   has_dyn  = length(var.nodeset_dyn) > 0

community/modules/compute/schedmd-slurm-gcp-v6-partition/outputs.tf

@@ -32,11 +32,6 @@ output "partitions" {
     condition     = sum([for b in [local.has_node, local.has_dyn, local.has_tpu] : b ? 1 : 0]) == 1
     error_message = "Partition must contain exactly one type of nodeset."
   }
-
-  precondition {
-    condition     = !local.uses_job_duration || var.exclusive
-    error_message = "`use_job_duration` can only be used in exclusive partitions"
-  }
 }
 
 output "nodeset" {

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

@@ -17,13 +17,14 @@
 import util
 import uuid
 from addict import Dict as NSDict # type: ignore
-from datetime import timedelta
+from datetime import datetime, timedelta
 from collections import defaultdict
 import logging
 from time import sleep
 
 log = logging.getLogger()
 
+DWS_EOL_RESERVATION_DURATION = 10 # minutes
 
 def _duration(flex_options: NSDict, job_id: Optional[int], lkp: util.Lookup) -> int:
     dur = flex_options.max_run_duration
@@ -40,6 +41,29 @@ def _duration(flex_options: NSDict, job_id: Optional[int], lkp: util.Lookup) ->
     log.info("Job TimeLimit cannot be less than 10 minutes or exceed one week")
     return dur
 
+def _create_slurm_reservation(node_name: str, boot_time: datetime, run_duration: int, lkp: util.Lookup):
+    """
+    Create a Slurm reservation starting at EOL - buffer time.
+    """
+    eol = boot_time + timedelta(seconds=run_duration)
+    start_str = eol.strftime("%Y-%m-%dT%H:%M:%S")
+    reservation_name = f"dws-eol-{node_name}"
+    log.debug(f"creating slurm reservation for {node_name}")
+    try:
+        util.run(f"{lkp.scontrol} create reservation user=slurm starttime={start_str} duration={DWS_EOL_RESERVATION_DURATION} nodes={node_name} reservationname={reservation_name} flags=maint,ignore_jobs")
+    except Exception as e:
+        log.error(f"Failed to create reservation for {node_name}: {e}")
+
+def _delete_slurm_reservation(node_name: str, lkp: util.Lookup):
+    """
+    Delete the Slurm reservation for the given node.
+    """
+    reservation_name = f"dws-eol-{node_name}"
+    try:
+        util.run(f"{lkp.scontrol} delete reservation {reservation_name}")
+        log.debug(f"Deleted Slurm reservation {reservation_name} for {node_name}")
+    except Exception as e:
+        log.error(f"Failed to delete reservation for {node_name}: {e}")
 
 def resume_flex_chunk(nodes: List[str], job_id: Optional[int], lkp: util.Lookup) -> None:
   assert nodes
@@ -77,8 +101,9 @@ def resume_flex_chunk(nodes: List[str], job_id: Optional[int], lkp: util.Lookup)
   op = req.execute()
   res = util.wait_for_operation(op)
   assert "error" not in res, f"{res}"
-  
+
   # Create resize request
+  duration_seconds = _duration(nodeset.dws_flex, job_id, lkp)
   req = lkp.compute.regionInstanceGroupManagerResizeRequests().insert(
     project=lkp.project,
     region=region,
@@ -87,13 +112,29 @@ def resume_flex_chunk(nodes: List[str], job_id: Optional[int], lkp: util.Lookup)
       name="initial-resize",
       instances=[dict(name=n) for n in nodes],
       requested_run_duration=dict(
-        seconds=_duration(nodeset.dws_flex, job_id, lkp)
+        seconds=duration_seconds
       )
     )
   )
   util.log_api_request(req)
   op = req.execute()
   res = util.wait_for_operation(op)
+
+  # Create Slurm reservations if use_job_duration is set
+  if nodeset.dws_flex.use_job_duration:
+      # Get run duration (seconds)
+      run_duration = duration_seconds
+      for node_name in nodes:
+          # Fetch instance creation time from GCP instance (via util.py)
+          instance = lkp.instance(node_name)
+          if(instance and instance.creation_timestamp):
+            log.debug("creating with creation_timestamp")
+            boot_time = instance.creation_timestamp  # Already a datetime object
+          else:
+            boot_time = datetime.utcnow()
+            log.debug("creating with utcnow time: {boot_time}")
+          _create_slurm_reservation(node_name, boot_time, run_duration, lkp)
+
   assert "error" not in res, f"{res}"
 
 def _suspend_flex_mig(mig_self_link: str, nodes: List[str], lkp: util.Lookup) -> None:
@@ -144,6 +185,12 @@ def _suspend_flex_mig(mig_self_link: str, nodes: List[str], lkp: util.Lookup) ->
   op = req.execute()
 
   res = util.wait_for_operation(op)
+
+  # Delete Slurm reservations for nodes being deprovisioned
+  for node_name in nodes:
+      log.info("delete dws reservation")
+      _delete_slurm_reservation(node_name, lkp)
+
   assert "error" not in res, f"{res}"
 
 def _suspend_provisioning_inst(nodes:List[str], node_template:str, lkp: util.Lookup) -> None: