Add support for axlearn

[Steboss]

No description provided.

[Steboss]

@olupton
Some updates on the last push:

• in store-delete-k8s-ghcr we're creating a token through ${RANDOM} (echo "token-name=${RANDOM}-${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT}" >> $GITHUB_OUTPUT) to create a unique token per process

TODO:

• Implement the s3 mount https://github.com/awslabs/mountpoint-s3-csi-driver in k8s jobs
• at the moment we're testing fuji-1B model, we need --model-name to be added to test-fuji-model.sh to allow tests on other models
• add pre-commit file that implements: 1) yq formatting, ruff and newline checks

[olupton]

Looking good. A few more comments, but we can refine later.

I think some of the error handling from

JAX-Toolbox/.github/workflows/_test_maxtext_k8s.yaml

Lines 74 to 90 in b0ec72a

	- name: Retrieve Kubernetes job status
	shell: bash -exo pipefail {0}
	run: |
	while readarray -d : -t status < <(kubectl get job/${JOB_NAME} -o 'jsonpath={.status.failed}:{.status.succeeded}'); do
	failure=${status[0]:-0}
	success=${status[1]:-0}
	total=$((failure+success))
	if [[ ${total} < 2 ]]; then
	sleep 1
	elif [[ ${total} == 2 ]]; then
	break
	else
	# FIXME
	exit 255
	fi
	done
	exit ${failure}

might be worth including in a follow-up.

Can you also add AXLearn to the README and plumb through its badge etc. ?

[olupton]

🚀