Improve error handling, s3 mounting, distributed tests for axlearn

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

@@ -16,7 +16,7 @@ runs:
       uses: ./.github/actions/with-post-step 
       with: 
         main: |
-          echo "Submit K8s job" 
+          echo "Submit K8s job ${{ inputs.job-config-file }}" 
           kubectl apply -f "${{ inputs.job-config-file }}"
 
           # Wait for job to be craeted
@@ -32,6 +32,47 @@ runs:
 
           # 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 
+          if [ -z "${parallelism}" ]; then
+            echo "No parallelism specified, defaulting to 1"
+            parallelism=1
+          fi
+
+          # Check whether the job succeeded or failed
+          while IFS=: read -r failures successes; do
+            failures="${status[0]:-0}"
+            successes="${status[1]:-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 -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/symlnk-cudnn.sh

@@ -10,7 +10,7 @@ CUDNN_MAJOR_VERSION=9
 prefix=/opt/nvidia/cudnn
 if [[ -d "${prefix}" ]]; then
   echo "Skipping link farm creation"
-  exit 1
+  exit 0
 fi
 
 arch=$(uname -m)-linux-gnu

.github/container/symlnk-nccl.sh

@@ -8,7 +8,7 @@ set -ex -o pipefail
 prefix=/opt/nvidia/nccl
 if [[ -d "${prefix}" ]]; then
   echo "Skipping link farm creation"
-  exit 1
+  exit 0
 fi
 arch=$(uname -m)-linux-gnu
 nccl_packages=$(dpkg -l 'libnccl*' | awk '/^ii/ {print $2}')

.github/eks-workflow-files/axlearn/axlearn-job.yml

@@ -26,35 +26,27 @@ spec:
 
                       sync  
                       wait 
-                      # after execution flag the results have been produced 
-                      touch /opt/output/done
-                  resources:
-                    limits:
-                        nvidia.com/gpu: 8
-                  volumeMounts:
-                    - name: output
-                      mountPath: /opt/output
-                - name: upload
-                  image: amazon/aws-cli
+                      # copy results to the mounted s3 bucket 
+                      mkdir -p /jax-toolbox-eks-output/axlearn/${RUN_ID}
+                      cp /opt/output/summary.txt /jax-toolbox-eks-output/axlearn/${RUN_ID}/summary.txt
+                      # copy all the log files 
+                      cp /opt/output/*.log /jax-toolbox-eks-output/axlearn/${RUN_ID}/.
                   env:
                     - name: RUN_ID
                       value: PLACEHOLDER
-                  command:
-                    - sh
-                    - -c
-                    - |
-                      while [ ! -f /opt/output/done ]; do
-                        sleep 5
-                      done
-                      # Upload to S3 bucket
-                      aws s3 cp /opt/output/summary.txt s3://jax-toolbox-eks-output/axlearn/${RUN_ID}/summary.txt
-                      # Upload logs to S3 bucket
-                      aws s3 cp /opt/output/ s3://jax-toolbox-eks-output/axlearn/${RUN_ID}/ --recursive --exclude "*" --include "*.log"
+                  resources:
+                    limits:
+                        nvidia.com/gpu: 8
                   volumeMounts:
                     - name: output
                       mountPath: /opt/output
+                    - name: s3-storage 
+                      mountPath: /jax-toolbox-eks-output
             imagePullSecrets:
                 - name: PLACEHOLDER
             volumes:
                 - name: output
                   emptyDir: {}
+                - name: s3-storage
+                  persistentVolumeClaim:
+                    claimName: s3-pvc