Enhance Dockerfile to install CSP NCCL network plugins for AWS and GCP. Added installation steps for Google's gIB plugin and updated the NCCL versioning. Ensured compatibility with different architectures and improved initialization script handling.

Signed-off-by: Ajay Balasa <abalasa@nvidia.com>

Author: balasaajay

docker/Dockerfile.fw_base

@@ -134,21 +134,44 @@ RUN --mount=type=bind,from=vllm_wheel,source=/src/vllm/,target=/tmp/vllm/ \
 
 ##############################################################################
 ##
-## Reinstall the aws-ofi-nccl EFA plugin
+## Install CSP NCCL network plugins
 ##
 ## NVIDIA PyTorch >= 26.06 dropped aws-ofi-nccl and defaults
 ## NCCL_NET_PLUGIN=spcx (HPCX Spectrum-X), which cannot drive EFA -> NCCL falls
 ## back to TCP sockets on AWS. Rebuild the plugin against the reinstalled NCCL and
-## the image's libfabric so EFA jobs keep their fabric. Selection stays per-CSP in
-## the launcher (EKSEnvPlugin pins NCCL_NET_PLUGIN for AWS). Kept last so it does
-## not invalidate the CUDA-toolkit or vLLM layer caches.
+## the image's libfabric so EFA jobs keep their fabric. Also install Google's gIB
+## plugin for GCP A3 Ultra+ instances. Kept last so these plugins are built against
+## the final NCCL and do not invalidate the CUDA-toolkit or vLLM layer caches.
 ##############################################################################
 
 ARG AWS_OFI_NCCL_VERSION=v1.17.3
+ARG GCP_NCCL_PLUGINS_VERSION=1.1.2
+ARG TARGETARCH
 
 RUN --mount=type=bind,source=docker/common/install_aws_ofi_nccl.sh,target=/opt/install_aws_ofi_nccl.sh \
     bash /opt/install_aws_ofi_nccl.sh --AWS_OFI_NCCL_VER=$AWS_OFI_NCCL_VERSION
 
+RUN set -euxo pipefail && \
+    case "${TARGETARCH}" in amd64|arm64) ;; *) echo "Unsupported TARGETARCH: ${TARGETARCH}" >&2; exit 1;; esac && \
+    mkdir -p /etc/apt/keyrings && \
+    curl -fsSL https://packages.cloud.google.com/apt/doc/apt-key.gpg \
+        -o /etc/apt/keyrings/cloud-google.asc && \
+    echo "deb [signed-by=/etc/apt/keyrings/cloud-google.asc] https://packages.cloud.google.com/apt gpudirect-gib-apt main" \
+        > /etc/apt/sources.list.d/nccl-gib.list && \
+    apt-get update && \
+    apt-get install -y --no-install-recommends \
+        "nccl-gcp-plugins-ngc=${GCP_NCCL_PLUGINS_VERSION}-*" && \
+    test -f /opt/gcp/nccl-plugins/lib64/libnccl-net-gcp.so && \
+    echo "/opt/gcp/nccl-plugins/lib64" > /etc/ld.so.conf.d/gcp-nccl-plugins.conf && \
+    rm -f /etc/apt/sources.list.d/nccl-gib.list /etc/apt/keyrings/cloud-google.asc && \
+    rm -rf /var/lib/apt/lists/* && \
+    ldconfig
+
+RUN --mount=type=bind,source=docker/common/gcp_nccl_init.sh,target=/opt/gcp_nccl_init.sh \
+    install -m 0644 /opt/gcp_nccl_init.sh /etc/gcp_nccl_init.sh && \
+    { grep -qxF '. /etc/gcp_nccl_init.sh' /etc/shinit_v2 || \
+      echo '. /etc/gcp_nccl_init.sh' >> /etc/shinit_v2; }
+
 ##############################################################################
 ##
 ## Final layer of FW_BASE container

docker/common/gcp_nccl_init.sh

@@ -0,0 +1,69 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. 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.
+
+#!/bin/sh
+
+# GCP detection: try multiple independent signals because newer machine types
+# do not expose SMBIOS to the guest, so the DMI path alone misses them.
+_is_gcp_host() {
+    # gVNIC driver bound to at least one PCI device.
+    for _d in /sys/bus/pci/drivers/gve/*:*; do
+        [ -e "$_d" ] && return 0
+    done
+    # Any PCI device with Google's vendor ID.
+    for _v in /sys/bus/pci/devices/*/vendor; do
+        [ -r "$_v" ] && [ "$(cat "$_v")" = "0x1ae0" ] && return 0
+    done
+    # SMBIOS fallback for machine families that expose it.
+    [ "$(cat /sys/class/dmi/id/product_name 2>/dev/null)" = "Google Compute Engine" ] && return 0
+    [ "$(cat /sys/class/dmi/id/sys_vendor 2>/dev/null)" = "Google" ] && return 0
+    return 1
+}
+
+# Only gIB-capable GCP instance families expose mlx5 RDMA NICs.
+_has_nvidia_rdma_nic() {
+    for _d in /sys/bus/pci/drivers/mlx5_core/*:*; do
+        [ -e "$_d" ] && return 0
+    done
+    return 1
+}
+
+if _is_gcp_host && _has_nvidia_rdma_nic && \
+    [ -f /opt/gcp/nccl-plugins/lib64/libnccl-net-gcp.so ]; then
+    export NCCL_NET_PLUGIN=gcp
+    if [ -f /opt/gcp/nccl-plugins/lib64/libnccl-tuner-gcp.so ]; then
+        export NCCL_TUNER_PLUGIN=gcp
+    fi
+    if [ -f /opt/gcp/nccl-plugins/lib64/libnccl-profiler-gcp.so ]; then
+        export NCCL_PROFILER_PLUGIN=gcp
+    fi
+    if [ -f /opt/gcp/nccl-plugins/lib64/libnccl-env-gcp.so ]; then
+        export NCCL_ENV_PLUGIN=gcp
+    fi
+    if [ -d /opt/gcp/nccl-plugins/configs ]; then
+        export NCCL_TUNER_CONFIG_PATH=/opt/gcp/nccl-plugins/configs
+    fi
+fi
+
+if [ "${NCCL_PROFILER_PLUGIN:-}" = gcp ]; then
+    echo
+    echo "NOTE: GCP NCCL telemetry enabled."
+    echo "      See https://docs.cloud.google.com/ai-hypercomputer/docs/nccl/comma"
+    echo "      for further information."
+fi
+
+unset -f _is_gcp_host
+unset -f _has_nvidia_rdma_nic
+unset _d
+unset _v