axlearn-job.yml

apiVersion: batch/v1
kind: Job
metadata:
    name: PLACEHOLDER
    labels:
        kueue.x-k8s.io/queue-name: p5-queue
spec:
    completions: 1
    parallelism: 1
    template:
        spec:
            restartPolicy: Never
            containers:
                - name: axlearn
                  image: PLACEHOLDER
                  command:
                    - bash
                    - -xo
                    - pipefail
                    - -c
                    - |
                      test-axlearn.sh \
                        --directory "." \
                        --output "/opt/output/" \
                        --test-files "/opt/axlearn/axlearn/common/*_test.py"

                      sync  
                      wait 
                      # 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
                  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