General cleaning and fixing of ci

.github/actions/build-container/action.yml

@@ -150,7 +150,6 @@ runs:
       with:
         builder: ${{ env.BUILDX_BUILDER_NAME }}
         context: ${{ inputs.DOCKER_CONTEXT }}
-        load: ${{inputs.registry-build-context-name == ''}}
         push: true
         file: ${{ inputs.DOCKERFILE }}
         platforms: linux/${{ inputs.ARCHITECTURE }}

.github/actions/eks-jobset/action.yml

@@ -163,7 +163,7 @@ runs:
               '{name: $name, value: $value}'
           done <<< "${{ inputs.ENVS }}" | jq -s -c .
         )
-        
+
         yq -iP '.spec.replicatedJobs[].template.spec.template.spec.containers[0].env += env(envs_json)' ${JOBSET_YAML}
         cat ${JOBSET_YAML}
 
@@ -178,10 +178,10 @@ runs:
       POLL_TIMEOUT: 10800
     run: |
       set -euo pipefail
-  
+
       START=$(date +%s)
       JOBSET_ACTIVE=""
-  
+
       while [[ "${JOBSET_ACTIVE:-}" != "true" ]]; do
         JOBSET_ACTIVE="$(
           kubectl get jobset -o json | jq -r '
@@ -190,18 +190,18 @@ runs:
             | .status.replicatedJobsStatus[0].active == 1
           '
         )"
-  
+
         NOW=$(date +%s)
         ELAPSED=$(( NOW - START ))
         if (( ELAPSED > POLL_TIMEOUT )); then
           echo "Timeout after waiting for JobSet ${JOBSET_NAME} to become active in the cluster"
           exit 1
         fi
-  
+
         echo "Waiting for JobSet ${JOBSET_NAME} to become active in the cluster"
         sleep 10
       done
-  
+
       echo "JobSet ${JOBSET_NAME} has just become active in the cluster"
 
   - name: Set JobSet Pods
@@ -235,25 +235,26 @@ runs:
     run: |
       parse_pod_exit_code() {
         local pod=$1
-        MAYBE_JOBSET_EXIT_CODE="$(cat ${JOBSET_NAME}/${pod}.log | grep -oE 'EXIT\_CODE=[0-9]+$')"
-
-        if [ $? -ne 0 ]; then
-          echo "The JobSet ${JOBSET_NAME} did not complete as expected " >&2
-          echo "JOBSET_EXIT_CODE=1" >> ${GITHUB_ENV}
-          exit 1
+        local match
+        # tail -n1: a pod can print EXIT_CODE= more than once; the last one wins
+        match="$(grep -oE 'EXIT_CODE=[0-9]+$' ${JOBSET_NAME}/${pod}.log | tail -n1)"
+        if [ -z "${match}" ]; then
+          return 1
         fi
-
-        echo "Pod ${pod} exited with ${MAYBE_JOBSET_EXIT_CODE}" >&2
-
-        eval "export ${MAYBE_JOBSET_EXIT_CODE}"
-        echo ${EXIT_CODE}
+        echo "${match#EXIT_CODE=}"
       }
 
       ALL_EXIT_CODES=0
       for jobset_pod in ${JOBSET_PODS//[()]/}; do
-        POD_EXIT_CODE=$(parse_pod_exit_code ${jobset_pod})
-        echo "Pod ${jobset_pod} had exit code ${POD_EXIT_CODE}"
-        ALL_EXIT_CODES=$(( ALL_EXIT_CODES + POD_EXIT_CODE ))
+        # check the substitution's status in the parent shell: an exit inside
+        # $( ) only kills the subshell and cannot fail this step by itself
+        if POD_EXIT_CODE=$(parse_pod_exit_code ${jobset_pod}); then
+          echo "Pod ${jobset_pod} had exit code ${POD_EXIT_CODE}"
+          ALL_EXIT_CODES=$(( ALL_EXIT_CODES + POD_EXIT_CODE ))
+        else
+          echo "Pod ${jobset_pod} has no EXIT_CODE= marker in its log: the JobSet ${JOBSET_NAME} did not complete as expected" >&2
+          ALL_EXIT_CODES=$(( ALL_EXIT_CODES + 1 ))
+        fi
       done
 
       echo "JOBSET_EXIT_CODE=${ALL_EXIT_CODES}" >> ${GITHUB_ENV}
@@ -272,16 +273,16 @@ runs:
     shell: bash -x -u {0}
     if: ${{ !cancelled() }}
     run: |
-      aws s3 cp s3://${{ inputs.S3_BUCKET }}/${{ inputs.NAME }}/${JOBSET_NAME}/ ${JOBSET_NAME}/ --recursive 
-  
+      aws s3 cp s3://${{ inputs.S3_BUCKET }}/${{ inputs.NAME }}/${JOBSET_NAME}/ ${JOBSET_NAME}/ --recursive
+
   - name: Emit unit test statistics
     if: ${{ !cancelled() }}
     shell: bash -xe -u {0}
     run: |
       STATISTICS_SCRIPT="${{ inputs.STATISTICS_SCRIPT }}"
       STATISTICS_SCRIPT=$(echo ${STATISTICS_SCRIPT} | sed 's/\n/\ /g')
       eval "${STATISTICS_SCRIPT}"
-      
+
       if [[ ${NUM_ERRORS} > 0 ]] || [[ ${NUM_TESTS} == 0 ]]; then
         BADGE_COLOR=red
         BADGE_MESSAGE=error
@@ -296,15 +297,15 @@ runs:
         BADGE_MESSAGE="${NUM_PASSED}/${NUM_TESTS} passed"
         SUMMARY="${{ inputs.NAME }} on ${{ inputs.GPU_ARCH }}: ${NUM_TESTS} total tests, ${NUM_ERRORS} errors, ${NUM_PASSED} passed, ${NUM_FAILED} failed."
       fi
-      
+
       source .github/workflows/scripts/to_json.sh
       BADGE_LABEL="${{ inputs.NAME }} ${{ inputs.GPU_ARCH }} Unit"
       to_json \
         SUMMARY \
         NUM_ERRORS NUM_TESTS NUM_PASSED NUM_FAILED \
         BADGE_LABEL BADGE_COLOR BADGE_MESSAGE \
       > sitrep.json
-      
+
       schemaVersion=1 \
       label="${BADGE_LABEL}" \
       message="${BADGE_MESSAGE}" \

.github/actions/gke-xpk/action.yml

@@ -113,7 +113,7 @@ runs:
     id: check
     shell: bash -x -u {0}
     run: |
-      gcloud container clusters list | grep ${{ inputs.GKE_CLUSTER }} | grep RUNNING > /dev/null 
+      gcloud container clusters list | grep ${{ inputs.GKE_CLUSTER }} | grep RUNNING > /dev/null
 
       if [[ "$?" == "1" ]]; then
         echo "Cluster ${{ inputs.GKE_CLUSTER }} does not exist"
@@ -168,9 +168,9 @@ runs:
       mkdir -p ${JOBSET_NAME}
       uv venv --verbose --python=${{ inputs.XPK_PYTHON }} --directory=${JOBSET_NAME}
       source ${JOBSET_NAME}/.venv/bin/activate
-     
+
       git clone --depth=1 --branch=${{ inputs.XPK_VERSION }} https://github.com/AI-Hypercomputer/xpk.git ${JOBSET_NAME}/xpk
-     
+
       # apply XPK workload patch
       PATCH_PATH=.github/gke-workflow/xpk/${{ inputs.XPK_VERSION}}
       ls ${PATCH_PATH}/*.patch | xargs -I {} git apply --unsafe-paths {} --directory ${JOBSET_NAME}/xpk
@@ -291,10 +291,10 @@ runs:
       POLL_TIMEOUT: 10800
     run: |
       set -euo pipefail
-  
+
       START=$(date +%s)
       JOBSET_ACTIVE=""
-  
+
       while [[ "${JOBSET_ACTIVE:-}" != "true" ]]; do
         JOBSET_ACTIVE="$(
           kubectl get jobset -o json | jq -r '
@@ -303,18 +303,18 @@ runs:
             | .status.replicatedJobsStatus[0].active == 1
           '
         )"
-  
+
         NOW=$(date +%s)
         ELAPSED=$(( NOW - START ))
         if (( ELAPSED > POLL_TIMEOUT )); then
           echo "Timeout after waiting for JobSet ${JOBSET_NAME} to become active in the cluster"
           exit 1
         fi
-  
+
         echo "Waiting for JobSet ${JOBSET_NAME} to become active in the cluster"
         sleep 10
       done
-  
+
       echo "JobSet ${JOBSET_NAME} has just become active in the cluster"
 
   - name: Set JobSet Pods
@@ -351,25 +351,26 @@ runs:
     run: |
       parse_pod_exit_code() {
         local pod=$1
-        MAYBE_JOBSET_EXIT_CODE="$(cat ${JOBSET_NAME}/${pod}.log | grep -oE 'EXIT\_CODE=[0-9]+$')"
-        echo ${MAYBE_JOBSET_EXIT_CODE} | grep -E 'EXIT\_CODE=[0-9]+$' > /dev/null
-
-        if [ $? -ne 0 ]; then
-          echo "The JobSet ${JOBSET_NAME} did not complete as expected " >&2
-          echo "JOBSET_EXIT_CODE=1" >> ${GITHUB_ENV}
-          exit 1
+        local match
+        # tail -n1: a pod can print EXIT_CODE= more than once; the last one wins
+        match="$(grep -oE 'EXIT_CODE=[0-9]+$' ${JOBSET_NAME}/${pod}.log | tail -n1)"
+        if [ -z "${match}" ]; then
+          return 1
         fi
-
-        echo "Pod ${pod} exited with ${MAYBE_JOBSET_EXIT_CODE}" >&2
-
-        eval "export ${MAYBE_JOBSET_EXIT_CODE}"
-        echo ${EXIT_CODE}
+        echo "${match#EXIT_CODE=}"
       }
 
       ALL_EXIT_CODES=0
       for jobset_pod in ${JOBSET_PODS//[()]/}; do
-        POD_EXIT_CODE=$(parse_pod_exit_code ${jobset_pod})
-        ALL_EXIT_CODES=$(( ALL_EXIT_CODES + POD_EXIT_CODE ))
+        # check the substitution's status in the parent shell: an exit inside
+        # $( ) only kills the subshell and cannot fail this step by itself
+        if POD_EXIT_CODE=$(parse_pod_exit_code ${jobset_pod}); then
+          echo "Pod ${jobset_pod} had exit code ${POD_EXIT_CODE}"
+          ALL_EXIT_CODES=$(( ALL_EXIT_CODES + POD_EXIT_CODE ))
+        else
+          echo "Pod ${jobset_pod} has no EXIT_CODE= marker in its log: the JobSet ${JOBSET_NAME} did not complete as expected" >&2
+          ALL_EXIT_CODES=$(( ALL_EXIT_CODES + 1 ))
+        fi
       done
 
       echo "JOBSET_EXIT_CODE=${ALL_EXIT_CODES}" >> ${GITHUB_ENV}
@@ -429,7 +430,7 @@ runs:
         NUM_ERRORS NUM_TESTS NUM_PASSED NUM_FAILED \
         BADGE_LABEL BADGE_COLOR BADGE_MESSAGE \
       > sitrep.json
-      
+
       schemaVersion=1 \
       label="${BADGE_LABEL}" \
       message="${BADGE_MESSAGE}" \

.github/container/Dockerfile.alphafold

@@ -39,6 +39,9 @@ COPY alphafold_files/model_parameters.md /opt/alphafold_input_data/model_paramet
 ###############################################################################
 
 FROM mealkit AS final
+# Global ARGs are not visible inside a stage unless re-declared: without this
+# the WORKDIR below expands to "" and is a silent no-op
+ARG SRC_PATH_ALPHAFOLD
 
 RUN pip-finalize.sh
 

.github/container/Dockerfile.axlearn

@@ -3,7 +3,7 @@ ARG BASE_IMAGE=ghcr.io/nvidia/jax-mealkit:jax
 ARG URLREF_AXLEARN=https://github.com/apple/axlearn.git#main
 ARG SRC_PATH_AXLEARN=/opt/axlearn
 ARG DEST_MANIFEST_DIR=/opt/manifest.d
-ARG GIT_USER_NAME=JAX Toolbox
+ARG GIT_USER_NAME="JAX Toolbox"
 ARG GIT_USER_EMAIL=jax@nvidia.com
 
 ###############################################################################
@@ -15,7 +15,6 @@ ARG SRC_PATH_AXLEARN
 ARG GIT_USER_NAME
 ARG GIT_USER_EMAIL
 ARG URLREF_AXLEARN
-ARG SRC_PATH_AXLEARN
 
 # Run the patch with cloning
 RUN <<"EOF" bash -exu
@@ -59,10 +58,6 @@ pytest-reportlog
 REQUIREMENTS
 EOF
 
-RUN echo "$(cat /opt/pip-tools.d/requirements-axlearn.in)"
-
-
-
 ###############################################################################
 ## Add test script to the path
 ###############################################################################
@@ -73,6 +68,10 @@ ADD test-axlearn.sh fuji-train-perf.py /usr/local/bin/
 ## Install accumulated packages from the base image and the previous stage
 ###############################################################################
 FROM mealkit AS final
+# Global ARGs are not visible inside a stage unless re-declared: without this
+# the WORKDIR below expands to "" and is a silent no-op (it only worked before
+# because the mealkit stage already set WORKDIR /opt/axlearn)
+ARG SRC_PATH_AXLEARN
 
 RUN pip-finalize.sh
 

.github/container/Dockerfile.base

@@ -45,12 +45,14 @@ 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 \
+ADD --chmod=755 \
   git-clone.sh \
   pip-finalize.sh \
   /usr/local/bin/
+# Pinned (rather than releases/latest) for reproducible builds
+ARG YQ_VERSION=v4.53.3
 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
+  chmod 755 /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"
@@ -71,7 +73,7 @@ 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
+# 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
@@ -112,7 +114,14 @@ ADD parallel-launch /usr/local/bin/
 ##############################################################################
 
 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
+# Run under bash explicitly: with the default /bin/sh (dash) the [[ test fails
+# with status 127 and the else-less `if` still returns 0, silently skipping the
+# install even when the override is set.
+RUN <<"EOF" bash -ex
+if [[ -n "${NSIGHT_SYSTEMS_VERSION_OVERRIDE}" ]]; then
+  install-nsight-systems.sh "${NSIGHT_SYSTEMS_VERSION_OVERRIDE}"
+fi
+EOF
 
 ##############################################################################
 ## Remove CUDA NVSHMEM symlinks

.github/container/Dockerfile.jax

@@ -48,7 +48,7 @@ RUN --mount=type=ssh \
 EOF
 
 ADD build-jax.sh build-te.sh local_cuda_arch /usr/local/bin/
-# Install bazelisk
+# Install bazelisk, pinned (rather than releases/latest) for reproducible builds
 RUN ARCH="$(dpkg --print-architecture)" && \
     wget -O /usr/local/bin/bazel https://github.com/bazelbuild/bazelisk/releases/latest/download/bazelisk-linux-${ARCH} && \
     chmod +x /usr/local/bin/bazel

.github/container/Dockerfile.maxtext

@@ -51,6 +51,9 @@ ADD test-maxtext.sh /usr/local/bin
 ###############################################################################
 
 FROM mealkit AS final
+# Global ARGs are not visible inside a stage unless re-declared: without this
+# the WORKDIR below expands to "" and is a silent no-op
+ARG SRC_PATH_MAXTEXT
 
 RUN pip-finalize.sh
 

.github/container/gcp-autoconfig.sh

@@ -1,6 +1,6 @@
 #!/bin/bash
 
-if [ "${DISABLE_GCP_TCPX_SETUP}" == 1 ]; then 
+if [ "${DISABLE_GCP_TCPX_SETUP}" == 1 ]; then
   return
 fi
 
@@ -12,20 +12,16 @@ NOCOLOR='\033[0m'
 if host metadata.google.internal &> /dev/null; then
     echo -e "${GREEN}
 ============================== JAX-ToolBox on GCP ==============================
-
 It looks like you're running on GCP. In order to maximize your multi-node
 performance, you'll need to use Google's TCPx NCCL plugin. This container ships
 the plugin at ${TCPX_LIBRARY_PATH}, which is already added to LD_LIBRARY_PATH by
 this script $2.
-
 For A3 instances, this script also automatically determines the best NICs to use
 by NCCL for maximum performance.
-
 However, there are additional steps you will need to take. Mainly, you'll need
 to run a separate receive-datapath-manager daemon on each of your nodes, and
 correctly configure your networks and NICs. For more information, please see the
 guide at: https://cloud.google.com/compute/docs/gpus/gpudirect#provide-access
-
 If you believe this setup is causing undesired effects, you can disable it by
 setting DISABLE_GCP_TCPX_SETUP=1
 ${NOCOLOR}"