Update EFA/AWS-OFI-NCCL installation recipe (#1116)

This is not run automatically at present. EFA-enabled containers can be
prepared by adding a new `RUN install-efa.sh` layer on top of the
published JAX-Toolbox containers.

Author: olupton

.github/container/install-efa.sh

@@ -1,35 +1,37 @@
 #!/bin/bash
-
 set -ex
 
-# Update distro
-apt-get update
-
-# Install required packages
-apt-get install -y curl
-
-# clean up all previously installed library to avoid conflicts
-# while installing Amazon EFA version
-dpkg --purge efa-config efa-profile libfabric openmpi \
-             ibacm ibverbs-providers ibverbs-utils infiniband-diags \
-             libibmad-dev libibmad5 libibnetdisc-dev libibnetdisc5 \
-             libibumad-dev libibumad3 libibverbs-dev libibverbs1 librdmacm-dev \
-             librdmacm1 rdma-core rdmacm-utils
-
-# Download Amazon EFA package and install
-EFA_INSTALLER_VERSION=latest
-WORKDIR=$(mktemp -d)
-
-pushd ${WORKDIR}
-
-AMAZON_EFA_LINK="https://efa-installer.amazonaws.com/aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz"
-curl -O "$AMAZON_EFA_LINK" 
-tar -xf aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz && cd aws-efa-installer
-./efa_installer.sh -y -g -d --skip-kmod --skip-limit-conf --no-verify
-
+EFA_INSTALLER_VERSION=1.34.0 # or: latest
+AWS_OFI_NCCL_PREFIX=/opt/aws-ofi-nccl
+AWS_OFI_NCCL_VERSION=1.11.0
+
+apt update
+
+EFA_TMP=$(mktemp -d)
+pushd $EFA_TMP
+curl -O https://efa-installer.amazonaws.com/aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz
+tar -xf aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz
+cd aws-efa-installer
+rm -v DEBS/UBUNTU2204/x86_64/{libpmix,openmpi,prrte}* # block installation of MPI components
+apt-get purge -y ibverbs-providers libibverbs-dev libibverbs1 libibumad-dev libibumad3 librdmacm1 librdmacm-dev ibverbs-utils
+./efa_installer.sh -g -y --skip-kmod --skip-limit-conf --no-verify |& tee install.log
+mv -v install.log /opt/amazon/efa/install.log
 popd
+rm -rf $EFA_TMP
+
+AWS_OFI_NCCL_TMP=$(mktemp -d)
+pushd $AWS_OFI_NCCL_TMP
+apt-get install -y libhwloc-dev
+curl -OL https://github.com/aws/aws-ofi-nccl/releases/download/v${AWS_OFI_NCCL_VERSION}-aws/aws-ofi-nccl-${AWS_OFI_NCCL_VERSION}-aws.tar.gz
+tar -xf aws-ofi-nccl-${AWS_OFI_NCCL_VERSION}-aws.tar.gz
+cd aws-ofi-nccl-${AWS_OFI_NCCL_VERSION}-aws
+./configure --prefix=${AWS_OFI_NCCL_PREFIX} --with-libfabric=/opt/amazon/efa --with-cuda=/usr/local/cuda --with-mpi=/usr/local/mpi
+make -j$(nproc) install
+popd
+rm -rf $AWS_OFI_NCCL_TMP
 
-# Clean up
-apt-get clean
 rm -rf /var/lib/apt/lists/*
-rm -rf ${WORKDIR}
+
+# Ranks higher than HPC-X => newly-installed libnccl-net.so becomes the default
+echo "${AWS_OFI_NCCL_PREFIX}/lib" > /etc/ld.so.conf.d/000_aws_ofi_nccl.conf
+ldconfig