From 003240c4a76e69f2d9bbdd4f8c3940faa31ad522 Mon Sep 17 00:00:00 2001
From: "Yu-Hang \"Maxin\" Tang" <Tang.Maxin@gmail.com>
Date: Sat, 7 Oct 2023 12:03:32 -0700
Subject: [PATCH] Install libnccl-dev in base image since JAX no longer vendors
 NCCL (#289)

---
 .github/container/Dockerfile.base | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/.github/container/Dockerfile.base b/.github/container/Dockerfile.base
index df73b8cd1..1aace5b53 100644
--- a/.github/container/Dockerfile.base
+++ b/.github/container/Dockerfile.base
@@ -30,6 +30,15 @@ RUN pip install --upgrade --no-cache-dir pip
 ADD install-cudnn.sh /usr/local/bin
 RUN install-cudnn.sh
 
+###############################################################################
+## Install NCCL
+###############################################################################
+
+RUN apt-get update && \
+    apt-get install -y libnccl-dev && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
 ###############################################################################
 ## RoCE and InfiniteBand support
 ###############################################################################