Add Phase 3 Internal Load Balancer and active-passive health check agent for Slurm HA (#5959)

Author: arpit974

community/modules/scheduler/schedmd-slurm-gcp-v6-controller/README.md

@@ -148,6 +148,30 @@ Example configuration overrides:
       create_firewall_rule: true
 ```
 
+### Internal Load Balancer (ILB) and Virtual IP (VIP)
+
+To provide stable network failover and avoid hardcoding individual controller IPs, an Internal Load Balancer (ILB) can be provisioned in front of the primary and backup controllers.
+
+To enable the Internal Load Balancer:
+1. Ensure `enable_backup_controller` is set to `true`.
+2. Set `enable_controller_load_balancer` to `true`.
+3. (Optional) Provide a static IP address using `controller_load_balancer_ip` to assign to the ILB's forwarding rule (VIP). If omitted, a dynamic IP is allocated from the subnetwork.
+
+#### Active-Passive HTTP Health Check Agent
+
+An Active-Passive HTTP health check agent automatically runs on port `6821` on each controller instance. The health check agent executes `scontrol ping` locally and reports status to the ILB:
+* **Primary Controller (`-0`)**: Returns HTTP `200 OK` if the primary controller is up and active, and `503 Service Unavailable` if offline.
+* **Backup Controller (`-1`)**: Returns HTTP `200 OK` only when the primary controller is down/offline and the backup controller is up and active. If the primary controller is up, the backup returns `503 Service Unavailable` (Standby) to prevent traffic routing to the standby instance.
+
+This setup ensures that all Slurm traffic is directed exclusively to the active controller node, enabling seamless and stable failover.
+
+#### Health Check and Network Integration Details
+
+* **Bypassing Loopback Routing Loops**: When `enable_controller_load_balancer` is active, the primary and backup controllers are configured with their real internal IPs mapped in `/etc/hosts`. This ensures that they can communicate directly with each other and bypasses potential loopback routing loops through the load balancer VIP. Non-controller nodes (such as login and compute nodes) map the controller hostname directly to the Load Balancer VIP.
+* **Multi-Threaded Probing**: The health check service uses a multi-threaded HTTPServer (`ThreadingHTTPServer`) to support simultaneous, concurrent probes from multiple Google Cloud load-balancer IP ranges without blocking or timing out.
+* **Optimized Local Timeout**: The health check agent executes `scontrol ping` with a reduced timeout of `3` seconds (down from `5` seconds). This ensures that if the local Slurm daemon hangs, the health check agent times out *before* the load balancer's typical 5-second probe timeout, returning a clean `503 Service Unavailable` response and preventing abrupt TCP teardowns.
+* **Security & Least Privilege**: The `slurm_health_check.service` systemd unit runs under the dedicated `slurm` user and group rather than `root`, adhering to security best practices.
+
 ### Backward Compatibility
 
 The HA feature is **opt-in**. If `enable_backup_controller` is set to `false`
@@ -394,9 +418,12 @@ limitations under the License.
 | [google_compute_address.controller_ips](https://registry.terraform.io/providers/hashicorp/google/latest/docs/resources/compute_address) | resource |
 | [google_compute_disk.controller_disk](https://registry.terraform.io/providers/hashicorp/google/latest/docs/resources/compute_disk) | resource |
 | [google_compute_firewall.health_check_firewall_rule](https://registry.terraform.io/providers/hashicorp/google/latest/docs/resources/compute_firewall) | resource |
+| [google_compute_forwarding_rule.slurm_controller_vip](https://registry.terraform.io/providers/hashicorp/google/latest/docs/resources/compute_forwarding_rule) | resource |
 | [google_compute_health_check.controller_health_check](https://registry.terraform.io/providers/hashicorp/google/latest/docs/resources/compute_health_check) | resource |
 | [google_compute_instance_group_manager.controller_zonal_mig](https://registry.terraform.io/providers/hashicorp/google/latest/docs/resources/compute_instance_group_manager) | resource |
 | [google_compute_per_instance_config.controller_zonal_stateful_ips](https://registry.terraform.io/providers/hashicorp/google/latest/docs/resources/compute_per_instance_config) | resource |
+| [google_compute_region_backend_service.slurm_controller_backend](https://registry.terraform.io/providers/hashicorp/google/latest/docs/resources/compute_region_backend_service) | resource |
+| [google_compute_region_health_check.slurm_health_check](https://registry.terraform.io/providers/hashicorp/google/latest/docs/resources/compute_region_health_check) | resource |
 | [google_compute_region_instance_group_manager.controller_regional_mig](https://registry.terraform.io/providers/hashicorp/google/latest/docs/resources/compute_region_instance_group_manager) | resource |
 | [google_compute_region_per_instance_config.controller_regional_stateful_ips](https://registry.terraform.io/providers/hashicorp/google/latest/docs/resources/compute_region_per_instance_config) | resource |
 | [google_secret_manager_secret.cloudsql](https://registry.terraform.io/providers/hashicorp/google/latest/docs/resources/secret_manager_secret) | resource |
@@ -427,6 +454,7 @@ limitations under the License.
 | <a name="input_compute_startup_script"></a> [compute\_startup\_script](#input\_compute\_startup\_script) | DEPRECATED: `compute_startup_script` has been deprecated.<br/>Use `startup_script` of nodeset module instead. | `any` | `null` | no |
 | <a name="input_compute_startup_scripts_timeout"></a> [compute\_startup\_scripts\_timeout](#input\_compute\_startup\_scripts\_timeout) | The timeout (seconds) applied to each startup script in compute nodes. If<br/>any script exceeds this timeout, then the instance setup process is considered<br/>failed and handled accordingly.<br/><br/>NOTE: When set to 0, the timeout is considered infinite and thus disabled. | `number` | `300` | no |
 | <a name="input_controller_ha_type"></a> [controller\_ha\_type](#input\_controller\_ha\_type) | Type of Managed Instance Group for controllers: 'zonal' or 'regional'. | `string` | `"zonal"` | no |
+| <a name="input_controller_load_balancer_ip"></a> [controller\_load\_balancer\_ip](#input\_controller\_load\_balancer\_ip) | Optional static IP address to assign to the controller Internal Load Balancer (VIP). If null, one will be dynamically assigned from the subnetwork. | `string` | `null` | no |
 | <a name="input_controller_network_attachment"></a> [controller\_network\_attachment](#input\_controller\_network\_attachment) | SelfLink for NetworkAttachment to be attached to the controller, if any. | `string` | `null` | no |
 | <a name="input_controller_project_id"></a> [controller\_project\_id](#input\_controller\_project\_id) | Optionally. Provision controller and config bucket in the different project | `string` | `null` | no |
 | <a name="input_controller_startup_script"></a> [controller\_startup\_script](#input\_controller\_startup\_script) | Startup script used by the controller VM. | `string` | `"# no-op"` | no |
@@ -450,6 +478,7 @@ limitations under the License.
 | <a name="input_enable_chs_gpu_health_check_prolog"></a> [enable\_chs\_gpu\_health\_check\_prolog](#input\_enable\_chs\_gpu\_health\_check\_prolog) | Enable a Cluster Health Sacnner(CHS) GPU health check that slurmd executes as a prolog script whenever it is asked to run a job step from a new job allocation. Compute nodes that fail GPU health check during prolog will be marked as drained. Find more details at:<br/>https://github.com/GoogleCloudPlatform/cluster-toolkit/tree/main/docs/CHS-Slurm.md | `bool` | `false` | no |
 | <a name="input_enable_cleanup_compute"></a> [enable\_cleanup\_compute](#input\_enable\_cleanup\_compute) | Enables automatic cleanup of compute nodes and resource policies (e.g.<br/>placement groups) managed by this module, when cluster is destroyed.<br/><br/>*WARNING*: Toggling this off will impact the running workload.<br/>Deployed compute nodes will be destroyed. | `bool` | `true` | 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_controller_load_balancer"></a> [enable\_controller\_load\_balancer](#input\_enable\_controller\_load\_balancer) | Enables an Internal Load Balancer (ILB) in front of the controllers for stable Virtual IP and network-level failover. | `bool` | `false` | no |
 | <a name="input_enable_controller_public_ips"></a> [enable\_controller\_public\_ips](#input\_enable\_controller\_public\_ips) | If set to true. The controller will have a random public IP assigned to it. Ignored if access\_config is set. | `bool` | `false` | no |
 | <a name="input_enable_debug_logging"></a> [enable\_debug\_logging](#input\_enable\_debug\_logging) | Enables debug logging mode. | `bool` | `false` | no |
 | <a name="input_enable_default_mounts"></a> [enable\_default\_mounts](#input\_enable\_default\_mounts) | Enable default global network storage from the controller<br/>- /home<br/>- /opt/apps | `bool` | `true` | no |

community/modules/scheduler/schedmd-slurm-gcp-v6-controller/controller.tf

@@ -240,7 +240,7 @@ resource "google_compute_firewall" "health_check_firewall_rule" {
 
   allow {
     protocol = "tcp"
-    ports    = [tostring(var.health_check.port)]
+    ports    = concat([tostring(var.health_check.port)], var.enable_controller_load_balancer ? ["6821"] : [])
   }
 }
 
@@ -417,6 +417,70 @@ resource "google_compute_region_per_instance_config" "controller_regional_statef
   }
 }
 
+# INTERNAL LOAD BALANCER: REGIONAL HTTP HEALTH CHECK (PORT 6821)
+resource "google_compute_region_health_check" "slurm_health_check" {
+  count       = (var.enable_backup_controller && var.enable_controller_load_balancer) ? 1 : 0
+  name        = "${local.slurm_cluster_name}-slurm-hc"
+  description = "Active-Passive HTTP health check for Slurm controllers (scontrol ping on port 6821)"
+  project     = local.controller_project_id
+  region      = var.region
+
+  check_interval_sec  = 10
+  timeout_sec         = 5
+  healthy_threshold   = 2
+  unhealthy_threshold = 3
+
+  http_health_check {
+    port         = 6821
+    request_path = "/"
+  }
+}
+
+# INTERNAL LOAD BALANCER: REGIONAL BACKEND SERVICE (L4 TCP)
+resource "google_compute_region_backend_service" "slurm_controller_backend" {
+  count                           = (var.enable_backup_controller && var.enable_controller_load_balancer) ? 1 : 0
+  name                            = "${local.slurm_cluster_name}-controller-backend"
+  description                     = "Internal L4 backend service for Slurm controllers VIP"
+  project                         = local.controller_project_id
+  region                          = var.region
+  load_balancing_scheme           = "INTERNAL"
+  protocol                        = "TCP"
+  connection_draining_timeout_sec = 10
+  health_checks                   = google_compute_region_health_check.slurm_health_check[*].id
+
+  dynamic "backend" {
+    for_each = var.controller_ha_type == "regional" ? [1] : []
+    content {
+      group          = one(google_compute_region_instance_group_manager.controller_regional_mig[*].instance_group)
+      balancing_mode = "CONNECTION"
+    }
+  }
+
+  dynamic "backend" {
+    for_each = var.controller_ha_type == "zonal" ? [1] : []
+    content {
+      group          = one(google_compute_instance_group_manager.controller_zonal_mig[*].instance_group)
+      balancing_mode = "CONNECTION"
+    }
+  }
+}
+
+# INTERNAL LOAD BALANCER: FORWARDING RULE (VIP)
+resource "google_compute_forwarding_rule" "slurm_controller_vip" {
+  count                 = (var.enable_backup_controller && var.enable_controller_load_balancer) ? 1 : 0
+  name                  = "${local.slurm_cluster_name}-controller-vip"
+  description           = "Virtual IP forwarding rule for Slurm controllers (ports 6817/6818)"
+  project               = local.controller_project_id
+  region                = var.region
+  load_balancing_scheme = "INTERNAL"
+  backend_service       = one(google_compute_region_backend_service.slurm_controller_backend[*].id)
+  ip_protocol           = "TCP"
+  all_ports             = true
+  subnetwork            = var.subnetwork_self_link
+  ip_address            = var.controller_load_balancer_ip
+  allow_global_access   = true
+}
+
 moved {
   from = module.slurm_controller_instance.google_compute_instance_from_template.slurm_instance[0]
   to   = google_compute_instance_from_template.controller

community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/README.md

@@ -81,6 +81,7 @@ No modules.
 | <a name="input_enable_bigquery_load"></a> [enable\_bigquery\_load](#input\_enable\_bigquery\_load) | Enables loading of cluster job usage into big query.<br/><br/>NOTE: Requires Google Bigquery API. | `bool` | `false` | no |
 | <a name="input_enable_chs_gpu_health_check_epilog"></a> [enable\_chs\_gpu\_health\_check\_epilog](#input\_enable\_chs\_gpu\_health\_check\_epilog) | Enable a Cluster Health Sacnner(CHS) GPU health check that slurmd executes as an epilog script after completing a job step from a new job allocation.<br/>Compute nodes that fail GPU health check during epilog will be marked as drained. Find more details at:<br/>https://github.com/GoogleCloudPlatform/cluster-toolkit/tree/main/docs/CHS-Slurm.md | `bool` | `false` | no |
 | <a name="input_enable_chs_gpu_health_check_prolog"></a> [enable\_chs\_gpu\_health\_check\_prolog](#input\_enable\_chs\_gpu\_health\_check\_prolog) | Enable a Cluster Health Sacnner(CHS) GPU health check that slurmd executes as a prolog script whenever it is asked to run a job step from a new job allocation. Compute nodes that fail GPU health check during prolog will be marked as drained. Find more details at:<br/>https://github.com/GoogleCloudPlatform/cluster-toolkit/tree/main/docs/CHS-Slurm.md | `bool` | `false` | no |
+| <a name="input_enable_controller_load_balancer"></a> [enable\_controller\_load\_balancer](#input\_enable\_controller\_load\_balancer) | Whether an Internal Load Balancer is configured in front of the controllers. | `bool` | `false` | no |
 | <a name="input_enable_debug_logging"></a> [enable\_debug\_logging](#input\_enable\_debug\_logging) | Enables debug logging mode. Not for production use. | `bool` | `false` | no |
 | <a name="input_enable_expedited_requeue"></a> [enable\_expedited\_requeue](#input\_enable\_expedited\_requeue) | Enables Expedited Requeue, which automatically requeues eligible jobs and grants them the highest priority upon node failure. (Usage: sbatch --requeue=expedite) | `bool` | `true` | no |
 | <a name="input_enable_external_prolog_epilog"></a> [enable\_external\_prolog\_epilog](#input\_enable\_external\_prolog\_epilog) | Automatically enable a script that will execute prolog and epilog scripts<br/>shared by NFS from the controller to compute nodes. Find more details at:<br/>https://github.com/GoogleCloudPlatform/slurm-gcp/blob/v5/tools/prologs-epilogs/README.md | `bool` | `false` | no |

community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/main.tf

@@ -41,19 +41,20 @@ resource "random_uuid" "cluster_id" {
 
 locals {
   config = {
-    enable_bigquery_load           = var.enable_bigquery_load
-    cloudsql_secret                = var.cloudsql_secret
-    cluster_id                     = random_uuid.cluster_id.result
-    project                        = var.project_id
-    slurm_cluster_name             = var.slurm_cluster_name
-    slurm_backup_controller_name   = var.slurm_backup_controller_name
-    slurm_backup_controller_ip     = var.slurm_backup_controller_ip
-    accounting_storage_backup_host = var.accounting_storage_backup_host
-    enable_slurm_auth              = var.enable_slurm_auth
-    bucket_path                    = local.bucket_path
-    enable_debug_logging           = var.enable_debug_logging
-    extra_logging_flags            = var.extra_logging_flags
-    controller_state_disk          = var.controller_state_disk
+    enable_bigquery_load            = var.enable_bigquery_load
+    cloudsql_secret                 = var.cloudsql_secret
+    cluster_id                      = random_uuid.cluster_id.result
+    project                         = var.project_id
+    slurm_cluster_name              = var.slurm_cluster_name
+    slurm_backup_controller_name    = var.slurm_backup_controller_name
+    slurm_backup_controller_ip      = var.slurm_backup_controller_ip
+    enable_controller_load_balancer = var.enable_controller_load_balancer
+    accounting_storage_backup_host  = var.accounting_storage_backup_host
+    enable_slurm_auth               = var.enable_slurm_auth
+    bucket_path                     = local.bucket_path
+    enable_debug_logging            = var.enable_debug_logging
+    extra_logging_flags             = var.extra_logging_flags
+    controller_state_disk           = var.controller_state_disk
 
     # storage
     disable_default_mounts = var.disable_default_mounts
@@ -203,6 +204,7 @@ locals {
     "repair.py",
     "setup_network_storage.py",
     "setup.py",
+    "slurm_health_check.py",
     "slurmsync.py",
     "sort_nodes.py",
     "suspend_wrapper.sh",

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

@@ -466,11 +466,14 @@ def install_slurm_conf(lkp: util.Lookup) -> None:
     # Resolve IPs for SlurmctldHost to ensure correct HA configuration (requires name(IP) format)
     # and use multiple lines for Slurm 23.02+ syntax compliance if HA is enabled.
     control_host = lkp.control_host
-    try:
-        control_ip = socket.gethostbyname(control_host)
-    except Exception as e:
-        log.warning(f"Failed to resolve control_host {control_host}: {e}")
-        control_ip = lkp.control_addr if lkp.control_addr else control_host
+    if lkp.cfg.get("enable_controller_load_balancer"):
+        control_ip = control_host
+    else:
+        try:
+            control_ip = socket.gethostbyname(control_host)
+        except Exception as e:
+            log.warning(f"Failed to resolve control_host {control_host}: {e}")
+            control_ip = lkp.control_addr if lkp.control_addr else control_host
 
     slurmctld_hosts_str = f"SlurmctldHost={control_host}({control_ip})"