Use pip-compile to help with consistent Python dependency resolution

.github/container/Dockerfile.base

@@ -1,6 +1,10 @@
 ARG BASE_IMAGE=nvidia/cuda:12.2.0-devel-ubuntu22.04
+ARG GIT_USER_NAME="JAX Toolbox"
+ARG GIT_USER_EMAIL=jax@nvidia.com
 
 FROM ${BASE_IMAGE}
+ARG GIT_USER_EMAIL
+ARG GIT_USER_NAME
 
 ###############################################################################
 ## Install Python and essential tools
@@ -31,8 +35,8 @@ RUN apt-get update && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 RUN <<"EOF" bash -ex
-git config --global user.name "JAX Toolbox"
-git config --global user.email "jax@nvidia.com"
+git config --global user.name "${GIT_USER_NAME}"
+git config --global user.email "${GIT_USER_EMAIL}"
 EOF
 RUN pip install --upgrade --no-cache-dir pip pip-tools && rm -rf ~/.cache/*
 RUN mkdir -p /opt/pip-tools.d

.github/container/Dockerfile.jax

@@ -11,6 +11,8 @@ 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 GIT_USER_NAME="JAX Toolbox"
+ARG GIT_USER_EMAIL=jax@nvidia.com
 
 ARG BAZEL_CACHE=/tmp
 ARG BUILD_DATE
@@ -27,6 +29,8 @@ ARG REF_XLA
 ARG SRC_PATH_JAX
 ARG SRC_PATH_XLA
 ARG BAZEL_CACHE
+ARG GIT_USER_NAME
+ARG GIT_USER_EMAIL
 
 RUN git clone "${REPO_JAX}" "${SRC_PATH_JAX}" && cd "${SRC_PATH_JAX}" && git checkout ${REF_JAX}
 RUN --mount=type=ssh \
@@ -37,8 +41,8 @@ RUN --mount=type=ssh \
 RUN <<EOF bash -ex
 cd ${SRC_PATH_XLA}
 
-git config user.name "JAX Toolbox"
-git config user.email "jax@nvidia.com"
+git config --global user.name "${GIT_USER_NAME}"
+git config --global user.email "${GIT_USER_EMAIL}"
 git remote add -f ashors1 https://github.com/ashors1/xla
 git cherry-pick --allow-empty $(git merge-base ashors/main ashors1/revert-84222)..ashors1/revert-84222
 git remote remove ashors1

.github/workflows/_build_base.yaml

@@ -27,6 +27,17 @@ on:
         description: 'Name of the endpoint JSON file for shields.io badge'
         required: false
         default: 'badge-base-build'
+      GIT_USER_NAME:
+        type: string
+        description: 'Username in GIT to perform git pull/push'
+        required: false
+        default: 'JAX Toolbox'
+      GIT_USER_EMAIL:
+        type: string
+        description: 'User email in GIT to perform git pull/push'
+        required: false
+        default: 'jax@nvidia.com'
+
     outputs:
       DOCKER_TAG:
         description: "Tag of the image built"
@@ -92,6 +103,8 @@ jobs:
           tags: ${{ steps.meta.outputs.tags }}
           labels: ${{ steps.meta.outputs.labels }}
           build-args: |
+            GIT_USER_NAME=${{ inputs.GIT_USER_NAME }}
+            GIT_USER_EMAIL=${{ inputs.GIT_USER_EMAIL }}
             BUILD_DATE=${{ inputs.BUILD_DATE }}
             ${{ inputs.BASE_IMAGE != 'latest' && format('BASE_IMAGE={0}', inputs.BASE_IMAGE) }}
 

.github/workflows/_build_jax.yaml

@@ -57,6 +57,16 @@ on:
         description: 'Name of the endpoint JSON file for shields.io badge'
         required: false
         default: 'badge-jax-build'
+      GIT_USER_NAME:
+        type: string
+        description: 'Username in GIT to perform git pull/push'
+        required: false
+        default: 'JAX Toolbox'
+      GIT_USER_EMAIL:
+        type: string
+        description: 'User email in GIT to perform git pull/push'
+        required: false
+        default: 'jax@nvidia.com'
     outputs:
       DOCKER_TAG_MEALKIT:
         description: "Tags of the 'mealkit' image built"
@@ -154,6 +164,8 @@ jobs:
             REF_JAX=${{ inputs.REF_JAX }}
             REF_XLA=${{ inputs.REF_XLA }}
             REF_TE=${{ inputs.REF_TE }}
+            GIT_USER_NAME=${{ inputs.GIT_USER_NAME }}
+            GIT_USER_EMAIL=${{ inputs.GIT_USER_EMAIL }}
 
       - name: Set docker metadata - final
         id: final-metadata

.github/workflows/_test_distribution.yaml

@@ -2,7 +2,30 @@ name: ~test core distribution logic
 
 on:
   workflow_call:
+    inputs:
+      GIT_USER_NAME:
+        type: string
+        description: 'Username in GIT to perform git pull/push'
+        required: false
+        default: 'JAX Toolbox'
+      GIT_USER_EMAIL:
+        type: string
+        description: 'User email in GIT to perform git pull/push'
+        required: false
+        default: 'jax@nvidia.com'
   workflow_dispatch:
+    inputs:
+      GIT_USER_NAME:
+        type: string
+        description: 'Username in GIT to perform git pull/push'
+        required: false
+        default: 'JAX Toolbox'
+      GIT_USER_EMAIL:
+        type: string
+        description: 'User email in GIT to perform git pull/push'
+        required: false
+        default: 'jax@nvidia.com'
+
 
 jobs:
   test-create-distribution:
@@ -17,8 +40,8 @@ jobs:
 
       - name: Set git login for tests
         run: |
-          git config --global user.email "jax@nvidia.com"
-          git config --global user.name "JAX-Toolbox CI"
+          git config --global user.name "${GIT_USER_NAME}"
+          git config --global user.email "${GIT_USER_EMAIL}"
 
       - name: Check out the repository under ${GITHUB_WORKSPACE}
         uses: actions/checkout@v3

.github/workflows/nightly-rosetta-pax-build.yaml

@@ -11,7 +11,7 @@ on:
       BASE_IMAGE:
         type: string
         description: 'PAX image built by NVIDIA/JAX-Toolbox'
-        default: 'ghcr.io/nvidia/upstream-pax:mealkit'
+        default: ''
         required: true
       PUBLISH:
         type: boolean

.github/workflows/nightly-rosetta-t5x-build-test.yaml

@@ -11,7 +11,7 @@ on:
       BASE_IMAGE:
         type: string
         description: 'T5x image built by NVIDIA/JAX-Toolbox'
-        default: 'ghcr.io/nvidia/upstream-t5x:latest'
+        default: ''
         required: true
       PUBLISH:
         type: boolean

rosetta/Dockerfile.pax

@@ -1,9 +1,12 @@
 # syntax=docker/dockerfile:1-labs
-ARG BASE_IMAGE=ghcr.io/nvidia/upstream-pax:latest
+ARG BASE_IMAGE=ghcr.io/nvidia/upstream-pax:mealkit
 ARG GIT_USER_EMAIL=jax@nvidia.com
 ARG GIT_USER_NAME=NVIDIA
 ARG SRC_PATH_PAXML=/opt/paxml
 ARG SRC_PATH_PRAXIS=/opt/praxis
+# These patchlist paths should be relative to this script
+ARG PAXML_PATCHLIST=patchlist-paxml.txt
+ARG PRAXIS_PATCHLIST=patchlist-praxis.txt
 
 FROM scratch as rosetta-source
 ARG SRC_PATH_PAXML
@@ -17,30 +20,43 @@ ADD --keep-git-dir=true https://github.com/google/paxml.git#main /
 FROM scratch as praxis-mirror-source
 ADD --keep-git-dir=true https://github.com/google/praxis.git#main /
 
-FROM ${BASE_IMAGE} AS rosetta
+###############################################################################
+### Download source and add auxiliary scripts
+################################################################################
+
+FROM ${BASE_IMAGE} AS mealkit
 ENV ENABLE_TE=1
 
 ARG GIT_USER_EMAIL
 ARG GIT_USER_NAME
-RUN <<EOF bash -e
-git config --global user.email "${GIT_USER_EMAIL}"
-git config --global user.name "${GIT_USER_NAME}"
-EOF
+ARG PAXML_PATCHLIST
+ARG PRAXIS_PATCHLIST
 
 COPY --from=rosetta-source / /opt/rosetta
 WORKDIR /opt/rosetta
 RUN --mount=target=/opt/pax-mirror,from=pax-mirror-source,readwrite \
-    --mount=target=/opt/praxis-mirror,from=praxis-mirror-source,readwrite <<EOF
+    --mount=target=/opt/praxis-mirror,from=praxis-mirror-source,readwrite <<EOF 
+bash -e
+git config --global user.email "${GIT_USER_EMAIL}"
+git config --global user.name "${GIT_USER_NAME}"
+
 bash create-distribution.sh \
-  -p patchlist-paxml.txt \
+  -p ${PAXML_PATCHLIST} \
   -m https://github.com/nvjax-svc-0/paxml.git \
-  -d $(dirname $(python -c "import paxml; print(*paxml.__path__)")) \
+  -d /opt/paxml \
   -e /opt/pax-mirror
 bash create-distribution.sh \
-  -p patchlist-praxis.txt \
+  -p ${PRAXIS_PATCHLIST} \
   -m https://github.com/nvjax-svc-0/praxis.git \
-  -d $(dirname $(python -c "import praxis; print(*praxis.__path__)")) \
+  -d /opt/praxis \
   -e /opt/praxis-mirror
-rm -rf $(find /opt -name "__pycache__")
+rm -rf $(find /opt -name "__pycache__") ~/.gitconfig
 EOF
 
+###############################################################################
+### Install accumulated packages from the base image and the previous stage
+################################################################################
+
+FROM mealkit as final
+
+RUN pip-finalize.sh

rosetta/Dockerfile.t5x

@@ -2,6 +2,9 @@
 ARG BASE_IMAGE=ghcr.io/nvidia/upstream-t5x:mealkit
 ARG GIT_USER_EMAIL=jax@nvidia.com
 ARG GIT_USER_NAME=NVIDIA
+# These patchlist paths should be relative to this script
+ARG T5X_PATCHLIST=patchlist-t5x.txt
+ARG FLAX_PATCHLIST=patchlist-flax.txt
 
 FROM scratch as rosetta-source
 COPY . /
@@ -12,45 +15,47 @@ ADD --keep-git-dir=true https://github.com/google-research/t5x.git#main /
 FROM scratch as flax-mirror-source
 ADD --keep-git-dir=true https://github.com/google/flax.git#main /
 
-FROM ${BASE_IMAGE} AS rosetta
+###############################################################################
+### Download source and add auxiliary scripts
+################################################################################
+
+FROM ${BASE_IMAGE} AS mealkit
 ENV ENABLE_TE=1
 
 ARG GIT_USER_EMAIL
 ARG GIT_USER_NAME
-RUN <<EOF bash -e
-git config --global user.email "${GIT_USER_EMAIL}"
-git config --global user.name "${GIT_USER_NAME}"
-EOF
+ARG T5X_PATCHLIST
+ARG FLAX_PATCHLIST
 
 COPY --from=rosetta-source / /opt/rosetta
 WORKDIR /opt/rosetta
 RUN --mount=target=/opt/t5x-mirror,from=t5x-mirror-source,readwrite \
-    --mount=target=/opt/flax-mirror,from=flax-mirror-source,readwrite <<EOF
+    --mount=target=/opt/flax-mirror,from=flax-mirror-source,readwrite <<EOF 
+bash -e
+git config --global user.email "${GIT_USER_EMAIL}"
+git config --global user.name "${GIT_USER_NAME}"
 bash create-distribution.sh \
-  -p patchlist-t5x.txt \
+  -p ${T5X_PATCHLIST} \
   -m https://github.com/nvjax-svc-0/t5x.git \
-  -d $(dirname $(python -c "import t5x; print(*t5x.__path__)")) \
+  -d /opt/t5x \
   -e /opt/t5x-mirror
 bash create-distribution.sh \
-  -p patchlist-flax.txt \
+  -p ${FLAX_PATCHLIST} \
   -m https://github.com/nvjax-svc-0/flax.git \
-  -d $(dirname $(python -c "import flax; print(*flax.__path__)")) \
+  -d /opt/t5x \
   -e /opt/flax-mirror
-rm -rf $(find /opt -name "__pycache__")
+rm -rf $(find /opt -name "__pycache__") ~/.gitconfig
+
+echo "--extra-index-url https://developer.download.nvidia.com/compute/redist" >> /opt/pip-tools.d/manifest.t5x
+echo "-e file:///opt/rosetta" >> /opt/pip-tools.d/manifest.t5x
 EOF
 
 WORKDIR /opt/rosetta
-RUN <<EOF bash -e
-pip install -e /opt/rosetta --extra-index-url https://developer.download.nvidia.com/compute/redist
-rm -rf ~/.cache/pip/
-EOF
 
-ADD test-vit.sh /usr/local/bin
+###############################################################################
+### Install accumulated packages from the base image and the previous stage
+################################################################################
 
-FROM rosetta AS rosetta-devel
+FROM mealkit as final
 
-WORKDIR /opt/rosetta
-RUN <<EOF bash -e
-pip install -e '/opt/rosetta[test,lint]' --extra-index-url https://developer.download.nvidia.com/compute/redist
-rm -rf ~/.cache/pip/
-EOF
+RUN pip-finalize.sh