#!/bin/sh

NV_DRIVER_VERS=$(sed -n 's/^NVRM.*Kernel Module\( for [a-z0-9_]*\| \) *\([^() ]*\).*$/\2/p' /proc/driver/nvidia/version 2>/dev/null | sed 's/^$/unknown/')

export _CUDA_COMPAT_PATH=${_CUDA_COMPAT_PATH:-/usr/local/cuda/compat}

# Set paths to use
_CUDA_COMPAT_CHECKFILE="${_CUDA_COMPAT_PATH}/.${NV_DRIVER_VERS}.$(hostname).checked"
_CUDA_COMPAT_SYMLINK="${_CUDA_COMPAT_PATH}/lib"
_CUDA_COMPAT_REALLIB="${_CUDA_COMPAT_PATH}/lib.real"
if [ -e /dev/nvgpu ]; then
  _CUDA_COMPAT_REALLIB="${_CUDA_COMPAT_REALLIB}-nvgpu"
fi

_RUNNING_CUDA_VERSION="$(nvidia-smi -q -d COMPUTE 2>/dev/null | grep "^CUDA Version" | sed 's/^.*: //')"
_CONTAINER_CUDA_VERSION="$(echo "${CUDA_VERSION}" | cut -d . -f 1-2)"

# If the CUDA driver was detected and the compat check hasn't been flagged as done yet, proceed
if [ -n "${NV_DRIVER_VERS}" -a ! -e "${_CUDA_COMPAT_CHECKFILE}" ]; then

  # Make note that compat check has already run in this environment
  touch "${_CUDA_COMPAT_CHECKFILE}" 2>/dev/null

  # Remove symlink that might have been left behind by a previous 'docker commit';
  # it will be recreated below if the compat check passes
  rm -f "${_CUDA_COMPAT_SYMLINK}"

  # Run cudaCheck with the compat library on LD_LIBRARY_PATH to see if it will initialize
  export _CUDA_COMPAT_STATUS="$(LD_LIBRARY_PATH="${_CUDA_COMPAT_REALLIB}" \
                                LD_PRELOAD="" \
                                timeout -s KILL $((${_CUDA_COMPAT_TIMEOUT:-30} + 5)) \
                                    /usr/local/bin/cudaCheck 2>/dev/null)"
  export _CUDA_COMPAT_STATUS="${_CUDA_COMPAT_STATUS:-"CUDA Driver UNAVAILABLE (compatibility check timed out)"}"

  # If compat check passed, bring the compat lib into the path (several ways for safety)
  if [ "${_CUDA_COMPAT_STATUS}" = "CUDA Driver OK" ]; then
    # symlink the compat lib into a location that was preset to be on LD_LIBRARY_PATH via ENV
    ln -sf "${_CUDA_COMPAT_REALLIB}" "${_CUDA_COMPAT_SYMLINK}" 2>/dev/null
    # Additionally prepend _CUDA_COMPAT_REALLIB onto LD_LIBRARY_PATH in case _CUDA_COMPAT_PATH was not writable
    export LD_LIBRARY_PATH="${_CUDA_COMPAT_REALLIB}${LD_LIBRARY_PATH:+":${LD_LIBRARY_PATH}"}"
  fi

fi

_DETECTED_MOFED=$(cat /sys/module/mlx5_core/version 2>/dev/null || true)
if [ -n "${_DETECTED_MOFED}" ] && \
   [ "${_DETECTED_MOFED}" != "$(uname -r)" ] && \
   dpkg --compare-versions "${_DETECTED_MOFED:-0}" lt 4.9
then
    # If the Mellanox driver is detected, appears to be from MOFED rather than inbox, is from MOFED < 4.9,
    # and OMPI_MCA_pml is not set, then we should just set it to disable UCX, because our UCX built on inbox
    # drivers won't work with that old MOFED driver
    export OMPI_MCA_pml=${OMPI_MCA_pml:-^ucx}
fi

if [ -d "/sys/bus/pci/drivers/efa" ] && \
   { [ -f "/opt/amazon/ofi-nccl/lib/libnccl-net-ofi.so" ] || \
     [ -f "/opt/amazon/aws-ofi-nccl/lib/libnccl-net-ofi.so" ]; }
then
    # It appears we're running on AWS EFA, so tell NCCL to try to use EFA plugins
    export NCCL_NET_PLUGIN=ofi
    export NCCL_TUNER_PLUGIN=ofi
fi

# GCP detection: try multiple independent signals because newer machine types
# (A3-mega, A4) don't expose SMBIOS to the guest, so the DMI path alone misses
# them. Order: cheapest+most-specific first.
_is_gcp_host() {
    # gVNIC driver bound to at least one PCI device — works on A3/A3-mega/A4.
    for _d in /sys/bus/pci/drivers/gve/*:*; do
        [ -e "$_d" ] && return 0
    done
    # Any PCI device with Google's vendor ID (0x1ae0) — covers the case where
    # the gve driver is shipped but not yet bound.
    for _v in /sys/bus/pci/devices/*/vendor; do
        [ -r "$_v" ] && [ "$(cat "$_v")" = "0x1ae0" ] && return 0
    done
    # SMBIOS — original signal, still works on N1/N2/A2/A3 families.
    [ "$(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
}

# The gIB plugin is RoCE over the MRDMA NICs (ConnectX-7/CX-8, mlx5 class)
# that only gIB-capable instance families expose (today: A3 Ultra/H200,
# A4/B200, A4X/GB200, A4X Max/GB300). gVNIC-only GPU instances (A3 Mega/High,
# A2, G2, G4) must not get the gIB settings: the forced plugin can't
# initialize there, drops NCCL to sockets, and overrides Google's TCPX/TCPXO
# plugins (DLR-5803).
_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
    # It appears we're running on GCP with the GCP NCCL plugins installed.
    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

# Clean up
unset _DETECTED_MOFED
unset -f _is_gcp_host
unset -f _has_nvidia_rdma_nic
unset _d
unset _v
unset _CUDA_COMPAT_CHECKFILE
unset _CUDA_COMPAT_REALLIB
unset _CUDA_COMPAT_SYMLINK
unset _RUNNING_CUDA_VERSION
unset _CONTAINER_CUDA_VERSION
