diff --git a/community/modules/compute/schedmd-slurm-gcp-v6-nodeset/README.md b/community/modules/compute/schedmd-slurm-gcp-v6-nodeset/README.md index cea9e8e862..c024ea0235 100644 --- a/community/modules/compute/schedmd-slurm-gcp-v6-nodeset/README.md +++ b/community/modules/compute/schedmd-slurm-gcp-v6-nodeset/README.md @@ -168,6 +168,7 @@ No modules. | [disk\_size\_gb](#input\_disk\_size\_gb) | Size of boot disk to create for the partition compute nodes. | `number` | `50` | no | | [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 | | [enable\_confidential\_vm](#input\_enable\_confidential\_vm) | Enable the Confidential VM configuration. Note: the instance image must support option. | `bool` | `false` | no | +| [enable\_maintenance\_reservation](#input\_enable\_maintenance\_reservation) | Enables slurm reservation for scheduled maintenance. | `bool` | `true` | no | | [enable\_oslogin](#input\_enable\_oslogin) | Enables Google Cloud os-login for user login and authentication for VMs.
See https://cloud.google.com/compute/docs/oslogin | `bool` | `true` | no | | [enable\_placement](#input\_enable\_placement) | Enable placement groups. | `bool` | `true` | no | | [enable\_public\_ips](#input\_enable\_public\_ips) | If set to true. The node group VMs will have a random public IP assigned to it. Ignored if access\_config is set. | `bool` | `false` | no | diff --git a/community/modules/compute/schedmd-slurm-gcp-v6-nodeset/main.tf b/community/modules/compute/schedmd-slurm-gcp-v6-nodeset/main.tf index 491ea64419..61e08f9fb6 100644 --- a/community/modules/compute/schedmd-slurm-gcp-v6-nodeset/main.tf +++ b/community/modules/compute/schedmd-slurm-gcp-v6-nodeset/main.tf @@ -105,6 +105,8 @@ locals { startup_script = local.ghpc_startup_script network_storage = var.network_storage + + enable_maintenance_reservation = var.enable_maintenance_reservation } } diff --git a/community/modules/compute/schedmd-slurm-gcp-v6-nodeset/variables.tf b/community/modules/compute/schedmd-slurm-gcp-v6-nodeset/variables.tf index 0a53ef95e2..1e6f5a076d 100644 --- a/community/modules/compute/schedmd-slurm-gcp-v6-nodeset/variables.tf +++ b/community/modules/compute/schedmd-slurm-gcp-v6-nodeset/variables.tf @@ -505,3 +505,10 @@ variable "instance_properties" { type = any default = null } + + +variable "enable_maintenance_reservation" { + type = bool + description = "Enables slurm reservation for scheduled maintenance." + default = true +} diff --git a/community/modules/scheduler/schedmd-slurm-gcp-v6-controller/README.md b/community/modules/scheduler/schedmd-slurm-gcp-v6-controller/README.md index 30ee38d084..2504525fa9 100644 --- a/community/modules/scheduler/schedmd-slurm-gcp-v6-controller/README.md +++ b/community/modules/scheduler/schedmd-slurm-gcp-v6-controller/README.md @@ -271,7 +271,7 @@ limitations under the License. | [metadata](#input\_metadata) | Metadata, provided as a map. | `map(string)` | `{}` | no | | [min\_cpu\_platform](#input\_min\_cpu\_platform) | Specifies a minimum CPU platform. Applicable values are the friendly names of
CPU platforms, such as Intel Haswell or Intel Skylake. See the complete list:
https://cloud.google.com/compute/docs/instances/specify-min-cpu-platform | `string` | `null` | no | | [network\_storage](#input\_network\_storage) | An array of network attached storage mounts to be configured on all instances. | 
list(object({
    server_ip             = string,
    remote_mount          = string,
    local_mount           = string,
    fs_type               = string,
    mount_options         = string,
    client_install_runner = optional(map(string))
    mount_runner          = optional(map(string))
  }))
| `[]` | no | -| [nodeset](#input\_nodeset) | Define nodesets, as a list. | 
list(object({
    node_count_static      = optional(number, 0)
    node_count_dynamic_max = optional(number, 1)
    node_conf              = optional(map(string), {})
    nodeset_name           = string
    additional_disks = optional(list(object({
      disk_name    = optional(string)
      device_name  = optional(string)
      disk_size_gb = optional(number)
      disk_type    = optional(string)
      disk_labels  = optional(map(string), {})
      auto_delete  = optional(bool, true)
      boot         = optional(bool, false)
    })), [])
    bandwidth_tier         = optional(string, "platform_default")
    can_ip_forward         = optional(bool, false)
    disable_smt            = optional(bool, false)
    disk_auto_delete       = optional(bool, true)
    disk_labels            = optional(map(string), {})
    disk_size_gb           = optional(number)
    disk_type              = optional(string)
    enable_confidential_vm = optional(bool, false)
    enable_placement       = optional(bool, false)
    enable_oslogin         = optional(bool, true)
    enable_shielded_vm     = optional(bool, false)
    gpu = optional(object({
      count = number
      type  = string
    }))
    labels                   = optional(map(string), {})
    machine_type             = optional(string)
    maintenance_interval     = optional(string)
    instance_properties_json = string
    metadata                 = optional(map(string), {})
    min_cpu_platform         = optional(string)
    network_tier             = optional(string, "STANDARD")
    network_storage = optional(list(object({
      server_ip             = string
      remote_mount          = string
      local_mount           = string
      fs_type               = string
      mount_options         = string
      client_install_runner = optional(map(string))
      mount_runner          = optional(map(string))
    })), [])
    on_host_maintenance = optional(string)
    preemptible         = optional(bool, false)
    region              = optional(string)
    service_account = optional(object({
      email  = optional(string)
      scopes = optional(list(string), ["https://www.googleapis.com/auth/cloud-platform"])
    }))
    shielded_instance_config = optional(object({
      enable_integrity_monitoring = optional(bool, true)
      enable_secure_boot          = optional(bool, true)
      enable_vtpm                 = optional(bool, true)
    }))
    source_image_family  = optional(string)
    source_image_project = optional(string)
    source_image         = optional(string)
    subnetwork_self_link = string
    additional_networks = optional(list(object({
      network            = string
      subnetwork         = string
      subnetwork_project = string
      network_ip         = string
      nic_type           = string
      stack_type         = string
      queue_count        = number
      access_config = list(object({
        nat_ip       = string
        network_tier = string
      }))
      ipv6_access_config = list(object({
        network_tier = string
      }))
      alias_ip_range = list(object({
        ip_cidr_range         = string
        subnetwork_range_name = string
      }))
    })))
    access_config = optional(list(object({
      nat_ip       = string
      network_tier = string
    })))
    spot               = optional(bool, false)
    tags               = optional(list(string), [])
    termination_action = optional(string)
    reservation_name   = optional(string)
    startup_script = optional(list(object({
      filename = string
    content = string })), [])

    zone_target_shape = string
    zone_policy_allow = set(string)
    zone_policy_deny  = set(string)
  }))
| `[]` | no | +| [nodeset](#input\_nodeset) | Define nodesets, as a list. | 
list(object({
    node_count_static      = optional(number, 0)
    node_count_dynamic_max = optional(number, 1)
    node_conf              = optional(map(string), {})
    nodeset_name           = string
    additional_disks = optional(list(object({
      disk_name    = optional(string)
      device_name  = optional(string)
      disk_size_gb = optional(number)
      disk_type    = optional(string)
      disk_labels  = optional(map(string), {})
      auto_delete  = optional(bool, true)
      boot         = optional(bool, false)
    })), [])
    bandwidth_tier                 = optional(string, "platform_default")
    can_ip_forward                 = optional(bool, false)
    disable_smt                    = optional(bool, false)
    disk_auto_delete               = optional(bool, true)
    disk_labels                    = optional(map(string), {})
    disk_size_gb                   = optional(number)
    disk_type                      = optional(string)
    enable_confidential_vm         = optional(bool, false)
    enable_placement               = optional(bool, false)
    enable_oslogin                 = optional(bool, true)
    enable_shielded_vm             = optional(bool, false)
    enable_maintenance_reservation = optional(bool, true)
    gpu = optional(object({
      count = number
      type  = string
    }))
    labels                   = optional(map(string), {})
    machine_type             = optional(string)
    maintenance_interval     = optional(string)
    instance_properties_json = string
    metadata                 = optional(map(string), {})
    min_cpu_platform         = optional(string)
    network_tier             = optional(string, "STANDARD")
    network_storage = optional(list(object({
      server_ip             = string
      remote_mount          = string
      local_mount           = string
      fs_type               = string
      mount_options         = string
      client_install_runner = optional(map(string))
      mount_runner          = optional(map(string))
    })), [])
    on_host_maintenance = optional(string)
    preemptible         = optional(bool, false)
    region              = optional(string)
    service_account = optional(object({
      email  = optional(string)
      scopes = optional(list(string), ["https://www.googleapis.com/auth/cloud-platform"])
    }))
    shielded_instance_config = optional(object({
      enable_integrity_monitoring = optional(bool, true)
      enable_secure_boot          = optional(bool, true)
      enable_vtpm                 = optional(bool, true)
    }))
    source_image_family  = optional(string)
    source_image_project = optional(string)
    source_image         = optional(string)
    subnetwork_self_link = string
    additional_networks = optional(list(object({
      network            = string
      subnetwork         = string
      subnetwork_project = string
      network_ip         = string
      nic_type           = string
      stack_type         = string
      queue_count        = number
      access_config = list(object({
        nat_ip       = string
        network_tier = string
      }))
      ipv6_access_config = list(object({
        network_tier = string
      }))
      alias_ip_range = list(object({
        ip_cidr_range         = string
        subnetwork_range_name = string
      }))
    })))
    access_config = optional(list(object({
      nat_ip       = string
      network_tier = string
    })))
    spot               = optional(bool, false)
    tags               = optional(list(string), [])
    termination_action = optional(string)
    reservation_name   = optional(string)
    startup_script = optional(list(object({
      filename = string
    content = string })), [])

    zone_target_shape = string
    zone_policy_allow = set(string)
    zone_policy_deny  = set(string)
  }))
| `[]` | no | | [nodeset\_dyn](#input\_nodeset\_dyn) | Defines dynamic nodesets, as a list. | 
list(object({
    nodeset_name    = string
    nodeset_feature = string
  }))
| `[]` | no | | [nodeset\_tpu](#input\_nodeset\_tpu) | Define TPU nodesets, as a list. | 
list(object({
    node_count_static      = optional(number, 0)
    node_count_dynamic_max = optional(number, 5)
    nodeset_name           = string
    enable_public_ip       = optional(bool, false)
    node_type              = string
    accelerator_config = optional(object({
      topology = string
      version  = string
      }), {
      topology = ""
      version  = ""
    })
    tf_version   = string
    preemptible  = optional(bool, false)
    preserve_tpu = optional(bool, false)
    zone         = string
    data_disks   = optional(list(string), [])
    docker_image = optional(string, "")
    network_storage = optional(list(object({
      server_ip             = string
      remote_mount          = string
      local_mount           = string
      fs_type               = string
      mount_options         = string
      client_install_runner = optional(map(string))
      mount_runner          = optional(map(string))
    })), [])
    subnetwork = string
    service_account = optional(object({
      email  = optional(string)
      scopes = optional(list(string), ["https://www.googleapis.com/auth/cloud-platform"])
    }))
    project_id = string
    reserved   = optional(string, false)
  }))
| `[]` | no | | [on\_host\_maintenance](#input\_on\_host\_maintenance) | Instance availability Policy. | `string` | `"MIGRATE"` | no | diff --git a/community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/slurmsync.py b/community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/slurmsync.py index 002057a3b1..eebc747532 100755 --- a/community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/slurmsync.py +++ b/community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/slurmsync.py @@ -27,7 +27,7 @@ import yaml import datetime as dt from datetime import datetime -from typing import Dict, Tuple +from typing import Dict, Tuple, Set import util from util import ( @@ -409,14 +409,14 @@ def reconfigure_slurm(): if lookup().cfg.hybrid: # terraform handles generating the config.yaml, don't do it here return - + upd, cfg_new = util.fetch_config() if not upd: log.debug("No changes in config detected.") return log.debug("Changes in config detected. Reconfiguring Slurm now.") util.update_config(cfg_new) - + if lookup().is_controller: conf.gen_controller_configs(lookup()) log.info("Restarting slurmctld to make changes take effect.") @@ -454,7 +454,7 @@ def create_reservation(lkp: util.Lookup, reservation_name: str, node: str, start util.run(f"{lkp.scontrol} create reservation user=slurm starttime={formatted_start_time} duration=180 nodes={node} reservationname={reservation_name}") -def get_slurm_reservation_maintenance(lkp: util.Lookup) -> Dict[str, datetime]: +def get_slurm_reservation_maintenance(lkp: util.Lookup, nodeset_enable_reservation: Set[str]) -> Dict[str, datetime]: res = util.run(f"{lkp.scontrol} show reservation --json") all_reservations = json.loads(res.stdout) reservation_map = {} @@ -473,27 +473,47 @@ def get_slurm_reservation_maintenance(lkp: util.Lookup) -> Dict[str, datetime]: if name != f"{nodes}_maintenance": continue + if nodes not in nodeset_enable_reservation: + continue + reservation_map[name] = datetime.fromtimestamp(time_epoch) return reservation_map -def get_upcoming_maintenance(lkp: util.Lookup) -> Dict[str, Tuple[str, datetime]]: +def get_upcoming_maintenance(lkp: util.Lookup, nodeset_enable_reservation: Set[str]) -> Dict[str, Tuple[str, datetime]]: upc_maint_map = {} for node, properties in lkp.instances().items(): - if 'upcomingMaintenance' in properties: + if 'upcomingMaintenance' in properties and node in nodeset_enable_reservation: start_time = datetime.strptime(properties['upcomingMaintenance']['startTimeWindow']['earliest'], '%Y-%m-%dT%H:%M:%S%z') upc_maint_map[node + "_maintenance"] = (node, start_time) return upc_maint_map +def get_nodeset_enable_reservation(lkp: util.Lookup) -> Set[str]: + nodeset_enable_reservation = set() + for nodeset in lkp.cfg.nodeset.values(): + if nodeset.enable_maintenance_reservation and nodeset.node_count_static: + static, _ = lkp.nodenames(nodeset) + nodeset_enable_reservation.update(static) + + return nodeset_enable_reservation + + def sync_maintenance_reservation(lkp: util.Lookup) -> None: - upc_maint_map = get_upcoming_maintenance(lkp) # map reservation_name -> (node_name, time) + nodeset_enable_reservation = get_nodeset_enable_reservation(lkp) + log.info(f"nodeset enabled for reservation for scheduled maintenance: {nodeset_enable_reservation}") + + if not nodeset_enable_reservation: + log.debug("no nodesets are enabled for reservation for scheduled maintenance.") + return + + upc_maint_map = get_upcoming_maintenance(lkp, nodeset_enable_reservation) # map reservation_name -> (node_name, time) log.debug(f"upcoming-maintenance-vms: {upc_maint_map}") - curr_reservation_map = get_slurm_reservation_maintenance(lkp) # map reservation_name -> time + curr_reservation_map = get_slurm_reservation_maintenance(lkp, nodeset_enable_reservation) # map reservation_name -> time log.debug(f"curr-reservation-map: {curr_reservation_map}") del_reservation = set(curr_reservation_map.keys() - upc_maint_map.keys()) @@ -541,14 +561,13 @@ def main(): except Exception: log.exception("failed to update topology") - ## TODO: Enable reservation for scheduled maintenance. - # try: - # sync_maintenance_reservation(lookup()) - # except Exception: - # log.exception("failed to sync slurm reservation for scheduled maintenance") + try: + sync_maintenance_reservation(lookup()) + except Exception: + log.exception("failed to sync slurm reservation for scheduled maintenance") try: - # TODO: it performs 1 to 4 GCS list requests, + # TODO: it performs 1 to 4 GCS list requests, # use cached version, combine with `_list_config_blobs` install_custom_scripts(check_hash=True) except Exception: diff --git a/community/modules/scheduler/schedmd-slurm-gcp-v6-controller/partition.tf b/community/modules/scheduler/schedmd-slurm-gcp-v6-controller/partition.tf index c3448df181..5799515a01 100644 --- a/community/modules/scheduler/schedmd-slurm-gcp-v6-controller/partition.tf +++ b/community/modules/scheduler/schedmd-slurm-gcp-v6-controller/partition.tf @@ -81,20 +81,21 @@ module "nodeset_cleanup" { locals { nodesets = [for name, ns in local.nodeset_map : { - nodeset_name = ns.nodeset_name - node_conf = ns.node_conf - instance_template = module.slurm_nodeset_template[ns.nodeset_name].self_link - node_count_dynamic_max = ns.node_count_dynamic_max - node_count_static = ns.node_count_static - subnetwork = ns.subnetwork_self_link - reservation_name = ns.reservation_name - maintenance_interval = ns.maintenance_interval - instance_properties_json = ns.instance_properties_json - enable_placement = ns.enable_placement - network_storage = ns.network_storage - zone_target_shape = ns.zone_target_shape - zone_policy_allow = ns.zone_policy_allow - zone_policy_deny = ns.zone_policy_deny + nodeset_name = ns.nodeset_name + node_conf = ns.node_conf + instance_template = module.slurm_nodeset_template[ns.nodeset_name].self_link + node_count_dynamic_max = ns.node_count_dynamic_max + node_count_static = ns.node_count_static + subnetwork = ns.subnetwork_self_link + reservation_name = ns.reservation_name + maintenance_interval = ns.maintenance_interval + instance_properties_json = ns.instance_properties_json + enable_placement = ns.enable_placement + network_storage = ns.network_storage + zone_target_shape = ns.zone_target_shape + zone_policy_allow = ns.zone_policy_allow + zone_policy_deny = ns.zone_policy_deny + enable_maintenance_reservation = ns.enable_maintenance_reservation }] } diff --git a/community/modules/scheduler/schedmd-slurm-gcp-v6-controller/variables.tf b/community/modules/scheduler/schedmd-slurm-gcp-v6-controller/variables.tf index 26dfd21d49..2fc7bebb4b 100644 --- a/community/modules/scheduler/schedmd-slurm-gcp-v6-controller/variables.tf +++ b/community/modules/scheduler/schedmd-slurm-gcp-v6-controller/variables.tf @@ -196,17 +196,18 @@ variable "nodeset" { auto_delete = optional(bool, true) boot = optional(bool, false) })), []) - bandwidth_tier = optional(string, "platform_default") - can_ip_forward = optional(bool, false) - disable_smt = optional(bool, false) - disk_auto_delete = optional(bool, true) - disk_labels = optional(map(string), {}) - disk_size_gb = optional(number) - disk_type = optional(string) - enable_confidential_vm = optional(bool, false) - enable_placement = optional(bool, false) - enable_oslogin = optional(bool, true) - enable_shielded_vm = optional(bool, false) + bandwidth_tier = optional(string, "platform_default") + can_ip_forward = optional(bool, false) + disable_smt = optional(bool, false) + disk_auto_delete = optional(bool, true) + disk_labels = optional(map(string), {}) + disk_size_gb = optional(number) + disk_type = optional(string) + enable_confidential_vm = optional(bool, false) + enable_placement = optional(bool, false) + enable_oslogin = optional(bool, true) + enable_shielded_vm = optional(bool, false) + enable_maintenance_reservation = optional(bool, true) gpu = optional(object({ count = number type = string