Add nccl rdma installer in gke a4x max (#5962)

Author: parulbajaj01

examples/gke-a4x-max-bm/gke-a4x-max-bm.yaml

@@ -54,6 +54,10 @@ vars:
   system_node_pool_disk_size_gb: 200
   a4x_max_node_pool_disk_size_gb: 100
   k8s_service_account_name: workload-identity-k8s-sa
+
+  # Installs NCCL library and Google NCCL plugin
+  # Runs an init container on all GB300 GPU nodes with the NCCL plugin image
+  gib_installer_path: $(ghpc_stage("./nccl-installer.yaml.tftpl"))
   nvidia_dra_driver_path: $(ghpc_stage("./nvidia-dra-driver.yaml"))
   asapd_lite_installer_path: $(ghpc_stage("./asapd-lite-installer.yaml"))
   kueue_configuration_path: $(ghpc_stage("./kueue-configuration.yaml.tftpl"))
@@ -253,6 +257,25 @@ deployment_groups:
         install: true
         version: "v25.8.0"
         accelerator_type: $(vars.accelerator_type)
+      gib:
+        install: true  # NCCL gIB plugin via DaemonSet initContainer
+        path: $(vars.gib_installer_path)
+        template_vars:
+          image: "us-docker.pkg.dev/gce-ai-infra/gpudirect-gib/nccl-plugin-gib-arm64"
+          version: v1.1.2
+          accelerator_count: 4
+          node_affinity:
+            requiredDuringSchedulingIgnoredDuringExecution:
+              nodeSelectorTerms:
+              - matchExpressions:
+                - key: cloud.google.com/gke-accelerator
+                  operator: In
+                  values:
+                  - $(vars.accelerator_type)
+                - key: kubernetes.io/arch
+                  operator: In
+                  values:
+                  - arm64
       apply_manifests:
       - source: $(vars.nvidia_dra_driver_path)
 

examples/gke-a4x-max-bm/nccl-installer.yaml.tftpl

@@ -0,0 +1,78 @@
+# Copyright 2026 Google Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+apiVersion: apps/v1
+kind: DaemonSet
+metadata:
+  name: nccl-rdma-installer
+  namespace: kube-system
+  labels:
+    k8s-app: nccl-rdma-installer
+spec:
+  selector:
+    matchLabels:
+      k8s-app: nccl-rdma-installer
+  updateStrategy:
+    type: RollingUpdate
+    rollingUpdate:
+      maxUnavailable: ${yamlencode(max_unavailable)}
+  template:
+    metadata:
+      labels:
+        name: nccl-rdma-installer
+        k8s-app: nccl-rdma-installer
+    spec:
+      priorityClassName: system-node-critical
+      affinity:
+        nodeAffinity:
+          ${indent(10, yamlencode(node_affinity))}
+      tolerations:
+        - operator: "Exists"
+      hostNetwork: true
+      hostPID: true
+      volumes:
+        - name: library-dir-host
+          hostPath:
+            path: /home/kubernetes/bin/nvidia/lib64
+            type: DirectoryOrCreate
+        - name: gib
+          hostPath:
+            path: /home/kubernetes/bin/gib
+        - name: nvidia-dir
+          hostPath:
+            path: /home/kubernetes/bin/nvidia
+            type: Directory
+      initContainers:
+        - image: ${image}:${version}
+          name: nccl-rdma-installer
+          resources:
+            requests:
+              cpu: 150m
+          securityContext:
+            privileged: true
+          volumeMounts:
+            - name: library-dir-host
+              mountPath: /usr/local/home/kubernetes/bin/nvidia/lib64
+            - name: gib
+              mountPath: /usr/local/home/kubernetes/bin/gib
+          command: ["/bin/sh", "-c"]
+          args:
+            - |
+              set -ex
+              /scripts/container_entry.sh install
+              cp -r /var/lib/gib/. /usr/local/home/kubernetes/bin/gib
+              echo "installation finishes"
+      containers:
+        - image: "gke.gcr.io/pause:3.8@sha256:880e63f94b145e46f1b1082bb71b85e21f16b99b180b9996407d61240ceb9830"
+          name: pause