Improve error handling, s3 mounting, distributed tests for axlearn (#1332)

Co-authored-by: Olli Lupton <[email protected]>

Author: Steboss

.github/actions/submit-delete-k8s-job/action.yml

@@ -13,8 +13,8 @@ runs:
   using: "composite"
   steps:
     - name: Submit and Delete Kubernetes job
-      uses: ./.github/actions/with-post-step 
-      with: 
+      uses: ./.github/actions/with-post-step
+      with:
         main: |
           set -x 
           TIMEOUT_JOB_CREATION=60s
@@ -35,9 +35,51 @@ runs:
           kubectl wait --for=condition=Ready \
             --selector=batch.kubernetes.io/job-name=${{ inputs.job-name }} \
             --timeout=$TIMEOUT_JOB_START pod
-          
+
           # Stream logs
           kubectl logs --all-containers=true --all-pods=true --follow job/${{ inputs.job-name }}
-          
-        post: | 
+
+          # Detect job parallelism
+          parallelism=$(kubectl get job/"${{ inputs.job-name }}" -o jsonpath='{.spec.parallelism}')
+          # if parallelism is not set, use default value of 1
+          echo "Parallelism ${parallelism}"
+          if [ -z "${parallelism}" ]; then
+            echo "No parallelism specified, defaulting to 1"
+            parallelism=1
+          fi
+
+          while IFS=: read -r failures successes; do
+            failures="${failures:-0}"
+            successes="${successes:-0}"
+            total=$((failures + successes))
+
+            if [ $total -lt $parallelism ]; then
+              # neither "failed" nor "succeeded", so wait
+              sleep 1
+            elif [ $total -eq $parallelism ]; then
+              # we have total=parallelism => either X successes or X failures
+              # In any case, the job is done
+              break
+            else
+              # Log here
+              echo "Unexpected number of completed pods ${total} with parallelism ${parallelism}"
+              exit 255
+            fi
+          done <<EOF
+          $(kubectl get job/"${{ inputs.job-name }}" -o 'jsonpath={.status.failed}:{.status.succeeded}')
+          EOF
+
+           # If job indicates a failure try to print out the info
+          if [ "${failures:-0}" -gt 0 ]; then
+            echo "Job ${{ inputs.job-name }} has $failures failures"
+            # this is for batch jobs only
+            pods=$(kubectl get pods --selector=batch.kubernetes.io/job-name=${{ inputs.job-name }} -o name)
+            if [ -n "${pods}" ]; then
+              kubectl describe ${pods}
+            fi
+            exit 1
+          fi
+
+        post: |
+          echo "Deleting K8s job: ${{ inputs.job-name }}"
           kubectl delete -f "${{ inputs.job-config-file }}"

.github/container/Dockerfile.axlearn

@@ -1,6 +1,6 @@
 # syntax=docker/dockerfile:1-labs
 ARG BASE_IMAGE=ghcr.io/nvidia/jax-mealkit:jax
-ARG URLREF_AXLEARN=https://github.com/apple/axlearn.git#main
+ARG URLREF_AXLEARN=https://github.com/Steboss/axlearn.git#main
 ARG SRC_PATH_AXLEARN=/opt/axlearn
 
 ###############################################################################
@@ -22,8 +22,9 @@ einops==0.8.0
 nltk==3.7
 portpicker==1.6.0
 seqio==0.0.18
-protobuf==3.20.3  
+protobuf==3.20.3
 pytest>=7.4.3
+tensorflow==2.18.1
 REQUIREMENTS
 EOF
 
@@ -32,7 +33,7 @@ EOF
 ## Add test script to the path
 ###############################################################################
 
-ADD test-axlearn.sh /usr/local/bin/
+ADD test-axlearn.sh fuji-train-perf.py /usr/local/bin/
 
 ###############################################################################
 ## Install accumulated packages from the base image and the previous stage