[Ardiegon]

Adding recipe for training Maxtext Qwen3.5-35B-A3B with v6e-256, both for kubernetes and xpk.

- UBench Run ID: maxtext_training_with_xpk-llama3_1_8b-2026-10-13_122827-0adb8e8a-dfb7-4c16-ac25-79fba638966b

This PR should be merged after successfully merging those two PRs:

• Add a TPU-friendly chunked delta rule for Gated DeltaNet (gdn_chunk_impl=fast) maxtext#5525
• Add Pallas MoE route/combine kernels, router top-k VJP, and Megablox v1 gmm/tgmm optimizations maxtext#5733

[gemini-code-assist]

Code Review

This pull request introduces recipes for pretraining the Qwen3.5-35B-A3B model using MaxText on Trillium GKE clusters (v6e-256) via both Kubernetes JobSet and XPK. The feedback highlights a few critical improvements: first, the ttlSecondsAfterFinished field in the JobSet manifest requires v0.14.0, which conflicts with the v0.13.0 version specified in the prerequisites; second, the signal trap in the manifest should use the captured $PID instead of $! to avoid volatility; and finally, the python command in the XPK recipe should be prepended with exec to ensure proper SIGTERM propagation for graceful shutdown.

[gemini-code-assist]

The ttlSecondsAfterFinished field was introduced in JobSet v0.14.0. Since the prerequisites in the README specify JobSet v0.13.0, applying this manifest will fail with a validation error on clusters running v0.13.0. Consider commenting this field out or upgrading the JobSet version in the README to v0.14.0 or later.

  # ttlSecondsAfterFinished: 43200

[gemini-code-assist]

Using $! inside the trap function is volatile because $! always refers to the most recently started background process in the shell, which can change if other background commands are executed. Since the background process PID is explicitly captured in $PID on line 55, it is much safer and more robust to use $PID in the signal handler. Additionally, using standard curly braces {} instead of subshell parentheses () is preferred for standard function definitions.

                  _sigterm() { kill -SIGTERM "$PID" 2>/dev/null; }

[gemini-code-assist]

The python training process is run in the foreground of the container's entrypoint shell without exec. When GKE/Kueue sends a SIGTERM (e.g., during preemption or eviction), the shell (PID 1) receives the signal but does not propagate it to the foreground python process, causing the training job to be abruptly killed without triggering MaxText's graceful shutdown and checkpointing handlers. Prepending exec ensures that the python process replaces the shell as PID 1 and receives termination signals directly.

Suggested change

	python3 -u -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml ${MAXTEXT_ARGS}"
	exec python3 -u -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml ${MAXTEXT_ARGS}"