Use pip-compile to help with consistent Python dependency resolution

.github/container/Dockerfile.base

@@ -1,4 +1,5 @@
 ARG BASE_IMAGE=nvidia/cuda:12.2.0-devel-ubuntu22.04
+
 FROM ${BASE_IMAGE}
 
 ###############################################################################
@@ -17,13 +18,28 @@ RUN apt-get update && \
         git \
         lld \
         vim \
+        bat \
+        curl \
+        git \
+        gnupg \
+        rsync \
         python-is-python3 \
         python3-pip \
+        liblzma-dev \
         wget \
     && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
-RUN pip install --upgrade --no-cache-dir pip
+RUN <<"EOF" bash -ex
+git config --global user.name "JAX Toolbox"
+git config --global user.email "jax@nvidia.com"
+EOF
+RUN pip install --upgrade --no-cache-dir pip pip-tools && rm -rf ~/.cache/*
+RUN mkdir -p /opt/pip-tools.d
+ADD --chmod=777 \
+    get-source.sh \
+    pip-finalize.sh \
+    /usr/local/bin/
 
 ###############################################################################
 ## Install cuDNN

.github/container/Dockerfile.jax

@@ -1,18 +1,25 @@
 ARG BASE_IMAGE=ghcr.io/nvidia/jax-toolbox:base
 ARG REPO_JAX="https://github.com/google/jax.git"
 ARG REPO_XLA="https://github.com/openxla/xla.git"
+ARG REPO_FLAX="https://github.com/google/flax.git"
+ARG REPO_TE="https://github.com/NVIDIA/TransformerEngine.git"
 ARG REF_JAX=main
 ARG REF_XLA=main
-ARG SRC_PATH_JAX=/opt/jax-source
+ARG REF_FLAX=main
+ARG REF_TE=main
+ARG SRC_PATH_JAX=/opt/jax
 ARG SRC_PATH_XLA=/opt/xla-source
+ARG SRC_PATH_FLAX=/opt/flax
+ARG SRC_PATH_TE=/opt/transformer-engine-source
+
 ARG BAZEL_CACHE=/tmp
 ARG BUILD_DATE
 
 ###############################################################################
 ## Build JAX
 ###############################################################################
 
-FROM ${BASE_IMAGE} as jax-builder
+FROM ${BASE_IMAGE} as builder
 ARG REPO_JAX
 ARG REPO_XLA
 ARG REF_JAX
@@ -47,15 +54,12 @@ RUN build-jax.sh \
     --xla-arm64-patch /opt/xla-arm64-neon.patch \ 
     --clean
 
-RUN cp -r ${SRC_PATH_JAX} ${SRC_PATH_JAX}-no-git && rm -rf ${SRC_PATH_JAX}-no-git/.git
-RUN cp -r ${SRC_PATH_XLA} ${SRC_PATH_XLA}-no-git && rm -rf ${SRC_PATH_XLA}-no-git/.git
-
 ###############################################################################
-## Build 'runtime' flavor without the git metadata
+## Pack jaxlib wheel and various source dirs into a pre-installation image
 ###############################################################################
 
 ARG BASE_IMAGE
-FROM ${BASE_IMAGE} as runtime-image
+FROM ${BASE_IMAGE} as mealkit
 ARG SRC_PATH_JAX
 ARG SRC_PATH_XLA
 ARG BUILD_DATE
@@ -66,29 +70,43 @@ ENV CUDA_DEVICE_MAX_CONNECTIONS=1
 ENV NCCL_IB_SL=1
 ENV CUDA_MODULE_LOADING=EAGER
 
-COPY --from=jax-builder ${SRC_PATH_JAX}-no-git ${SRC_PATH_JAX}
-COPY --from=jax-builder ${SRC_PATH_XLA}-no-git ${SRC_PATH_XLA}
+COPY --from=builder ${SRC_PATH_JAX} ${SRC_PATH_JAX}
+COPY --from=builder ${SRC_PATH_XLA} ${SRC_PATH_XLA}
+ADD build-jax.sh local_cuda_arch test-jax.sh /usr/local/bin/
 
-RUN pip --disable-pip-version-check install ${SRC_PATH_JAX}/dist/*.whl && \
-    pip --disable-pip-version-check install -e ${SRC_PATH_JAX} && \
-    rm -rf ~/.cache/pip/
+RUN mkdir -p /opt/pip-tools.d
+RUN <<"EOF" bash -ex
+echo "-e file://${SRC_PATH_JAX}"                        >> /opt/pip-tools.d/manifest.jax
+echo "jaxlib @ file://$(ls ${SRC_PATH_JAX}/dist/*.whl)" >> /opt/pip-tools.d/manifest.jax
+EOF
 
-# Install software stack in JAX ecosystem
-# Made this optional since tensorstore cannot build on Ubuntu 20.04 + ARM
-RUN { pip install flax || true; } && rm -rf ~/.cache/*
+## Flax
+ARG REPO_FLAX
+ARG REF_FLAX
+ARG SRC_PATH_FLAX
+RUN get-source.sh -f ${REPO_FLAX} -r ${REF_FLAX} -d ${SRC_PATH_FLAX} -m /opt/pip-tools.d/manifest.flax
+
+## Transformer engine: check out source and build wheel
+ARG REPO_TE
+ARG REF_TE
+ARG SRC_PATH_TE
+ENV NVTE_FRAMEWORK=jax
+ENV SRC_PATH_TE=${SRC_PATH_TE}
+RUN <<"EOF" bash -ex
+set -o pipefail
+pip install ninja && rm -rf ~/.cache/pip
+get-source.sh -f ${REPO_TE} -r ${REF_TE} -d ${SRC_PATH_TE}
+pushd ${SRC_PATH_TE}
+python setup.py bdist_wheel && rm -rf build
+echo "transformer-engine @ file://$(ls ${SRC_PATH_TE}/dist/*.whl)" >> /opt/pip-tools.d/manifest.te
+EOF
 
 # TODO: properly configure entrypoint
-# COPY entrypoint.d/ /opt/nvidia/entrypoint.d/
 
 ###############################################################################
-## Build 'devel' image with build scripts and git metadata
+## Install primary packages and transitive dependencies
 ###############################################################################
 
-FROM runtime-image as devel-image
-ARG SRC_PATH_JAX
-ARG SRC_PATH_XLA
-
-ADD build-jax.sh local_cuda_arch test-jax.sh /usr/local/bin/
+FROM mealkit as final
 
-COPY --from=jax-builder ${SRC_PATH_JAX}/.git ${SRC_PATH_JAX}/.git
-COPY --from=jax-builder ${SRC_PATH_XLA}/.git ${SRC_PATH_XLA}/.git
+RUN pip-finalize.sh

.github/container/Dockerfile.pax.amd64

@@ -1,33 +1,54 @@
 # syntax=docker/dockerfile:1-labs
-###############################################################################
-## Pax
-###############################################################################
 
 ARG BASE_IMAGE=ghcr.io/nvidia/jax:latest
-FROM ${BASE_IMAGE}
-
-ADD install-pax.sh /usr/local/bin
-ADD install-flax.sh /usr/local/bin
-ADD install-te.sh /usr/local/bin
-
-ENV NVTE_FRAMEWORK=jax
 ARG REPO_PAXML=https://github.com/google/paxml.git
 ARG REPO_PRAXIS=https://github.com/google/praxis.git
 ARG REF_PAXML=main
 ARG REF_PRAXIS=main
+ARG SRC_PATH_PAXML=/opt/paxml
+ARG SRC_PATH_PRAXIS=/opt/praxis
+
+###############################################################################
+## Download source and add auxiliary scripts
+###############################################################################
+
+FROM ${BASE_IMAGE} as mealkit
+ARG REPO_PAXML
+ARG REPO_PRAXIS
+ARG REF_PAXML
+ARG REF_PRAXIS
+ARG SRC_PATH_PAXML
+ARG SRC_PATH_PRAXIS
+
+# update TE manifest file to install the [test] extras
+RUN sed -i "s/transformer-engine @/transformer-engine[test] @/g" /opt/pip-tools.d/manifest.te
+
 RUN <<"EOF" bash -ex
-install-pax.sh --defer --from_paxml ${REPO_PAXML} --from_praxis ${REPO_PRAXIS} --ref_paxml ${REF_PAXML} --ref_praxis ${REF_PRAXIS}
-install-flax.sh --defer
-install-te.sh --defer
-
-if [[ -f /opt/requirements-defer.txt ]]; then
-  # SKIP_HEAD_INSTALLS avoids having to install jax from Github source so that
-  # we do not overwrite the jax that was already installed.
-  SKIP_HEAD_INSTALLS=true pip install -r /opt/requirements-defer.txt
-fi
-if [[ -f /opt/cleanup.sh ]]; then
-  bash -ex /opt/cleanup.sh
-fi
+get-source.sh -f ${REPO_PAXML}  -r ${REF_PAXML}  -d ${SRC_PATH_PAXML}
+get-source.sh -f ${REPO_PRAXIS} -r ${REF_PRAXIS} -d ${SRC_PATH_PRAXIS}
+echo "-e file://${SRC_PATH_PAXML}[gpu]" >> /opt/pip-tools.d/manifest.pax
+echo "-e file://${SRC_PATH_PRAXIS}"     >> /opt/pip-tools.d/manifest.pax
+
+for src in ${SRC_PATH_PAXML} ${SRC_PATH_PRAXIS}; do
+  pushd ${src}
+  sed -i "s| @ git+https://github.com/google/flax||g" requirements.in
+  sed -i "s| @ git+https://github.com/google/jax||g" requirements.in
+  if git diff --quiet; then
+      echo "URL specs no longer present in select dependencies for ${src}"
+      exit 1
+  else
+      git commit -a -m "remove URL specs from select dependencies for ${src}"
+  fi
+  popd
+done
 EOF
 
 ADD test-pax.sh /usr/local/bin
+
+###############################################################################
+## Install accumulated packages from the base image and the previous stage
+###############################################################################
+
+FROM mealkit as final
+
+RUN pip-finalize.sh