Dockerfile.base

# syntax=docker/dockerfile:1-labs
ARG BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.02-cuda13.1-devel-ubuntu24.04
ARG GIT_USER_NAME="JAX Toolbox"
ARG GIT_USER_EMAIL=jax@nvidia.com
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 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=(
  file
  libxml2
  python-is-python3
  python3-dev
  xxd
  zip
)
apt-get install --no-install-recommends --no-install-suggests -y ${apt_packages[@]}
apt-get clean
rm -rf /var/lib/apt/lists/*
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
# Install the pip at /opt/pip using get-pip.py
RUN <<EOF bash -exo pipefail
curl -O https://bootstrap.pypa.io/get-pip.py
python get-pip.py -e /opt/pip
rm -rf get-pip.py ~/.cache/pip
EOF

# Install pip-tools, which is used in pip-finalize.sh
RUN pip install --no-cache-dir pip-tools

# 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

###############################################################################
## 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/