.github/container/Dockerfile.base

# syntax=docker/dockerfile:1-labs
ARG BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:25.09-cuda13.0-devel-ubuntu24.04
ARG GIT_USER_NAME="JAX Toolbox"
ARG GIT_USER_EMAIL=jax@nvidia.com
ARG CLANG_VERSION=18
ARG JAX_TOOLBOX_REF=main
ARG NSIGHT_SYSTEMS_VERSION_OVERRIDE

###############################################################################
## Build base image
###############################################################################

FROM ${BASE_IMAGE}
ARG BASE_IMAGE
ARG GIT_USER_EMAIL
ARG GIT_USER_NAME
ARG CLANG_VERSION
ARG JAX_TOOLBOX_REF
ARG NSIGHT_SYSTEMS_VERSION_OVERRIDE
ENV CUDA_BASE_IMAGE=${BASE_IMAGE}
ENV JAX_TOOLBOX_REF=${JAX_TOOLBOX_REF}

###############################################################################
## Install Python and essential tools
###############################################################################

RUN <<"EOF" bash -ex
export DEBIAN_FRONTEND=noninteractive
export TZ=America/Los_Angeles
apt-get update

apt_packages=(
  # generic
  bat
  build-essential
  checkinstall
  cmake
  curl
  git
  gnupg
  liblzma-dev
  ninja-build
  python-is-python3
  python3-pip
  rsync
  vim
  wget
  jq
  zip
  # llvm.sh
  lsb-release 
  software-properties-common
)
apt-get install -y ${apt_packages[@]}

# Install LLVM/Clang
bash -c "$(wget -O - https://apt.llvm.org/llvm.sh)" -- ${CLANG_VERSION}

# Make sure that clang and clang++ point to the new version. This list is based
# on the symlinks installed by the `clang` (as opposed to `clang-14`) and `lld`
# (as opposed to `lld-14`) packages available in Ubuntu 22.04. 100 is an
# arbitrary priority.
update-alternatives --verbose \
  --install /usr/bin/clang          clang          /usr/lib/llvm-${CLANG_VERSION}/bin/clang 100 \
  --slave   /usr/bin/asan_symbolize asan_symbolize /usr/bin/asan_symbolize-${CLANG_VERSION} \
  --slave   /usr/bin/clang++        clang++        /usr/lib/llvm-${CLANG_VERSION}/bin/clang++ \
  --slave   /usr/bin/ld.lld         ld.lld         /usr/lib/llvm-${CLANG_VERSION}/bin/lld \
  --slave   /usr/bin/lld            lld            /usr/lib/llvm-${CLANG_VERSION}/bin/lld \
  --slave   /usr/bin/lld-link       lld-link       /usr/lib/llvm-${CLANG_VERSION}/bin/lld \
  --slave   /usr/bin/wasm-ld        wasm-ld        /usr/lib/llvm-${CLANG_VERSION}/bin/lld

# Make sure that any later attempt to install `clang` or `lld` will fail.
cat >/etc/apt/preferences.d/no-unversioned-clang-lld <<EOL
# LLVM is installed from apt.llvm.org using versioned packages, whereas the
# unversioned clang and lld packages come from Ubuntu and refer to older
# versions of LLVM. Please use versioned packages in this container.
Package: clang
Pin: release *
Pin-Priority: -1

Package: lld
Pin: release *
Pin-Priority: -1
EOL

apt-get clean
rm -rf /var/lib/apt/lists/*

# There are several python packages (in the list below) that are installed with OS 
# package manager (the run of `apt-get install` above) and can not be uninstall
# using pip (in pip-finalize.sh script) during JAX installation. Remove then in
# advance to avoid JAX installation issue.
remove_packages=(
  python3-gi
  software-properties-common
  lsb-release
  python3-yaml
  python3-pygments
)

apt-get remove -y ${remove_packages[@]}
apt-get autoremove -y  # removes python3-blinker which conflicts with pip-compile in JAX
EOF

RUN <<"EOF" bash -ex
git config --global user.name "${GIT_USER_NAME}"
git config --global user.email "${GIT_USER_EMAIL}"
EOF
RUN mkdir -p /opt/pip-tools.d /opt/pip-tools-post-install.d
ADD --chmod=777 \
  git-clone.sh \
  pip-finalize.sh \
  /usr/local/bin/
RUN wget https://github.com/mikefarah/yq/releases/latest/download/yq_linux_$(dpkg --print-architecture) -O /usr/local/bin/yq && \
  chmod 777 /usr/local/bin/yq
RUN git clone -b 23.3.1 https://github.com/pypa/pip.git /opt/pip
# Patch is specific to 23.3.1
# Generated via: "git diff > pip-vcs-equivalency.patch"
ADD pip-vcs-equivalency.patch /opt/pip/
RUN <<EOF bash -e -x
cd /opt/pip
git apply </opt/pip/pip-vcs-equivalency.patch
git add -u
git commit -m 'Adds JAX_TOOLBOX_VCS_EQUIVALENCY as a trigger to treat all github VCS installs for a package as equivalent. The spec of the last encountered version will be used'
EOF

# install all python packages system-wide.
ENV PIP_BREAK_SYSTEM_PACKAGES=1
# An extra flag `--ignore-installed` is added below, because of the following reason: 
# after upgrading to ver 23.3.1 (from /opt/pip) `pip` tries to uninstall itself (default pip-24.0) 
# and fails due to pip-24.0 has been installed with system tool `apt` but not `python`. So we keep 
# both pip-24.0 and pip-23.3.1 in the system, but use 23.3.1 with equivalency patch (see above).
RUN pip install --upgrade --ignore-installed --no-cache-dir -e /opt/pip "pip-tools==7.4.1" && rm -rf ~/.cache/*

# The symlinks for CUDA/cuDNN/NCCL exist to make the container's installations
# of those components conform to XLA's expectations for local installations.

###############################################################################
## Symlink for CUDA
###############################################################################

RUN ln -s /usr/local/cuda/lib64 /usr/local/cuda/lib

###############################################################################
## Symlink for cuDNN
###############################################################################

ADD symlnk-cudnn.sh /usr/local/bin
RUN symlnk-cudnn.sh

###############################################################################
## Symlink for NCCL
###############################################################################

ADD symlnk-nccl.sh /usr/local/bin
RUN symlnk-nccl.sh

##############################################################################
## NCCL sanity check utility
##############################################################################

ADD install-nccl-sanity-check.sh /usr/local/bin
ADD nccl-sanity-check.cu /opt
RUN install-nccl-sanity-check.sh
ADD jax-nccl-test parallel-launch /usr/local/bin/


##############################################################################
## Install a newer Nsight Systems version than is included in cuda-dl-base
##############################################################################

ADD install-nsight-systems.sh /usr/local/bin
RUN if [[ -n "${NSIGHT_SYSTEMS_VERSION_OVERRIDE}" ]]; then install-nsight-systems.sh "${NSIGHT_SYSTEMS_VERSION_OVERRIDE}"; fi

##############################################################################
## Install NVSHMEM
##############################################################################
ADD install-nvshmem.sh /usr/local/bin
RUN install-nvshmem.sh

##############################################################################
## Create symlinks to help XLA find NVSHMEM
##############################################################################
ADD symlnk-nvshmem.sh /usr/local/bin
RUN symlnk-nvshmem.sh

###############################################################################
## Install the nsys-jax JAX/XLA-aware profiling scripts, patch Nsight Systems
###############################################################################

ADD install-nsys-jax.sh /usr/local/bin
RUN install-nsys-jax.sh ${JAX_TOOLBOX_REF}

###############################################################################
## Copy manifest file to the container
###############################################################################

ENV MANIFEST_FILE="/opt/manifest.d/manifest.yaml"
ADD manifest.yaml create-distribution.sh bump.sh /opt/manifest.d/

COPY patches/ /opt/manifest.d/patches/