From 769e0a38087db03868b3f6e3bd6ceb88441f14d2 Mon Sep 17 00:00:00 2001 From: mesakhcienet Date: Mon, 6 Jul 2026 02:08:14 +0000 Subject: [PATCH 1/3] fix: Revert notebook CI to use v6e-4 runners --- .github/workflows/ci_pipeline.yml | 4 ++-- .github/workflows/run_jupyter_notebooks.yml | 18 ++++++++++++++++-- 2 files changed, 18 insertions(+), 4 deletions(-) diff --git a/.github/workflows/ci_pipeline.yml b/.github/workflows/ci_pipeline.yml index 9e1d8fd22f..c066ecd25b 100644 --- a/.github/workflows/ci_pipeline.yml +++ b/.github/workflows/ci_pipeline.yml @@ -147,9 +147,9 @@ jobs: fail-fast: false with: device_type: tpu - device_name: v6e-8 + device_name: v6e-4 base_image: maxtext-unit-test-tpu:py312 - cloud_runner: linux-x86-ct6e-180-8tpu + cloud_runner: linux-x86-ct6e-180-4tpu maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }} secrets: HF_TOKEN: ${{ secrets.HF_TOKEN }} diff --git a/.github/workflows/run_jupyter_notebooks.yml b/.github/workflows/run_jupyter_notebooks.yml index 0b383e3b81..a32162a949 100644 --- a/.github/workflows/run_jupyter_notebooks.yml +++ b/.github/workflows/run_jupyter_notebooks.yml @@ -107,11 +107,15 @@ jobs: $PYTHON_EXE -m ipykernel install --user --name maxtext_venv # Run Hugging Face authentication - hf auth login --token "$HF_TOKEN" + if [ -n "$HF_TOKEN" ]; then + hf auth login --token "$HF_TOKEN" + else + echo "Warning: HF_TOKEN is empty (likely because this is a PR from a fork). Skipping Hugging Face authentication." + fi for notebook in "$MAXTEXT_NOTEBOOKS_ROOT"/*.ipynb; do filename=$(basename "$notebook") - if [[ "$filename" == "sft_llama3_demo_gpu.ipynb" || "$filename" == "maxtext_with_gepa.ipynb" || "$filename" == "demo_decoding.ipynb" ]]; then + if [[ "$filename" == "sft_llama3_demo_gpu.ipynb" || "$filename" == "maxtext_with_gepa.ipynb" || "$filename" == "rl_llama3_demo.ipynb" || "$filename" == "demo_decoding.ipynb" ]]; then echo "Skipping $filename" continue fi @@ -127,6 +131,16 @@ jobs: echo "Post-notebook disk cleanup for $filename ..." rm -rf "$MAXTEXT_PKG_DIR"/*_output rm -rf "$HOME/.cache/huggingface/hub" + + # papermill's kernel teardown can time out on heavy notebooks (e.g. + # sft_multimodal_gemma3_demo), leaving the ipykernel process alive and + # still holding the TPU + /tmp/libtpu_lockfile. That breaks TPU init for + # the next notebook ("ABORTED: ... libtpu multi-process lockfile"). Reap + # any surviving kernel and clear the lock before the next iteration. + echo "Post-notebook TPU cleanup for $filename ..." + pkill -9 -f ipykernel_launcher 2>/dev/null || true + sleep 2 + rm -f /tmp/libtpu_lockfile done - name: Upload Outputs if: always() From 4a6aff22d117d4607346c81c67beeea7d0ccb0ec Mon Sep 17 00:00:00 2001 From: mesakhcienet Date: Mon, 6 Jul 2026 13:24:16 +0800 Subject: [PATCH 2/3] Update run_jupyter_notebooks.yml --- .github/workflows/run_jupyter_notebooks.yml | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/.github/workflows/run_jupyter_notebooks.yml b/.github/workflows/run_jupyter_notebooks.yml index a32162a949..6059836f17 100644 --- a/.github/workflows/run_jupyter_notebooks.yml +++ b/.github/workflows/run_jupyter_notebooks.yml @@ -132,15 +132,21 @@ jobs: rm -rf "$MAXTEXT_PKG_DIR"/*_output rm -rf "$HOME/.cache/huggingface/hub" - # papermill's kernel teardown can time out on heavy notebooks (e.g. - # sft_multimodal_gemma3_demo), leaving the ipykernel process alive and - # still holding the TPU + /tmp/libtpu_lockfile. That breaks TPU init for - # the next notebook ("ABORTED: ... libtpu multi-process lockfile"). Reap - # any surviving kernel and clear the lock before the next iteration. + # Kill EVERY process still holding a vfio/accel fd (by device, not by name), + # clear the lock, then wait for the device to actually be released before the next notebook. echo "Post-notebook TPU cleanup for $filename ..." + for pid in $(ls /proc 2>/dev/null | grep -E '^[0-9]+$'); do + if ls -l /proc/"$pid"/fd 2>/dev/null | grep -qE '/dev/(vfio|accel)'; then + echo " killing pid $pid ($(cat /proc/"$pid"/comm 2>/dev/null)) holding TPU" + kill -9 "$pid" 2>/dev/null || true + fi + done pkill -9 -f ipykernel_launcher 2>/dev/null || true - sleep 2 rm -f /tmp/libtpu_lockfile + for _ in $(seq 1 30); do + ls -l /proc/*/fd 2>/dev/null | grep -qE '/dev/(vfio|accel)' || break + sleep 2 + done done - name: Upload Outputs if: always() From 66f4925b553cadc87b9995a9f700779d93fb8c88 Mon Sep 17 00:00:00 2001 From: mesakhcienet Date: Mon, 6 Jul 2026 15:54:36 +0800 Subject: [PATCH 3/3] Update run_jupyter_notebooks.yml --- .github/workflows/run_jupyter_notebooks.yml | 15 --------------- 1 file changed, 15 deletions(-) diff --git a/.github/workflows/run_jupyter_notebooks.yml b/.github/workflows/run_jupyter_notebooks.yml index 6059836f17..d33dda602a 100644 --- a/.github/workflows/run_jupyter_notebooks.yml +++ b/.github/workflows/run_jupyter_notebooks.yml @@ -132,21 +132,6 @@ jobs: rm -rf "$MAXTEXT_PKG_DIR"/*_output rm -rf "$HOME/.cache/huggingface/hub" - # Kill EVERY process still holding a vfio/accel fd (by device, not by name), - # clear the lock, then wait for the device to actually be released before the next notebook. - echo "Post-notebook TPU cleanup for $filename ..." - for pid in $(ls /proc 2>/dev/null | grep -E '^[0-9]+$'); do - if ls -l /proc/"$pid"/fd 2>/dev/null | grep -qE '/dev/(vfio|accel)'; then - echo " killing pid $pid ($(cat /proc/"$pid"/comm 2>/dev/null)) holding TPU" - kill -9 "$pid" 2>/dev/null || true - fi - done - pkill -9 -f ipykernel_launcher 2>/dev/null || true - rm -f /tmp/libtpu_lockfile - for _ in $(seq 1 30); do - ls -l /proc/*/fd 2>/dev/null | grep -qE '/dev/(vfio|accel)' || break - sleep 2 - done done - name: Upload Outputs if: always()