Merge pull request #256 from AI-Hypercomputer:yuyan-fix-test-hang

PiperOrigin-RevId: 752361458

Author: jetstream authors

.github/workflows/run_maxtext_jetstream_tests.yaml

@@ -35,7 +35,7 @@ jobs:
     - name: Cleanup old docker images
       run: docker system prune --all --force
     - name: Authenticate gcloud
-      run: gcloud auth configure-docker us-docker.pkg.dev --quiet
+      run: gcloud auth configure-docker gcr.io --quiet
 
   build_stable_stack:
     name: Build Stable Stack
@@ -45,6 +45,8 @@ jobs:
       LOCAL_IMAGE_TAG: jetstream-maxtext-stable-stack/tpu:github_${{ github.run_id }}
     steps:
     - uses: actions/checkout@v4
+    - name: Authenticate gcloud
+      run: gcloud auth configure-docker gcr.io --quiet
     - name: Build
       run: |
           pushd experimental/jetstream-maxtext-stable-stack
@@ -59,61 +61,87 @@ jobs:
           popd
     - name: Upload image
       run: |
-          UPLOAD_IMAGE_TAG=gcr.io/cloud-tpu-inference-test/${LOCAL_IMAGE_TAG}
+          UPLOAD_IMAGE_TAG=gcr.io/cloud-ml-auto-solutions/${LOCAL_IMAGE_TAG}
           docker tag ${LOCAL_IMAGE_TAG} ${UPLOAD_IMAGE_TAG}
           docker push ${UPLOAD_IMAGE_TAG}
-          NIGHTLY_TAG=${UPLOAD_IMAGE_TAG%:*}:nightly
-          NIGHTLY_TAG_DATE=${NIGHTLY_TAG}-$(date +"%Y%m%d")
-          docker tag ${LOCAL_IMAGE_TAG} ${NIGHTLY_TAG}
-          docker tag ${LOCAL_IMAGE_TAG} ${NIGHTLY_TAG_DATE}
-          docker push ${NIGHTLY_TAG}
-          docker push ${NIGHTLY_TAG_DATE}
 
   benchmark_report:
     name: Benchmark Report
     needs: build_stable_stack
     runs-on: ["self-hosted", "tpu", "v6e-8"]
-    container:
-      # sync with the image uploaded from build_stable_stack stage
-      image: gcr.io/cloud-tpu-inference-test/jetstream-maxtext-stable-stack/tpu:github_${{ github.run_id }}
-      options: "--net=host --privileged"
     env:
-      OUTPUT_DIR: /workspace/test_dir/
+      OUTPUT_DIR: ./test_dir
     steps:
-    - name: Create output directory # Ensure directory exists in container
-      run: mkdir -p ${OUTPUT_DIR}
     - name: Test MOEBenchmarks
-      # Report should generated in OUTPUT_DIR depend on ENV
-      run: bash JetStream/.github/workflows/test_moe_benchmarks.sh
+      run: |
+        rm -rf ${OUTPUT_DIR}
+        mkdir -p ${OUTPUT_DIR}
+        # sync with the image uploaded from build_stable_stack stage
+        # Report should generated in OUTPUT_DIR depend on ENV
+        DOCKER_OUTPUT_DIR=/output
+        docker run \
+          -v ${OUTPUT_DIR}:${DOCKER_OUTPUT_DIR} \
+          --env OUTPUT_DIR=${DOCKER_OUTPUT_DIR} \
+          --privileged --net=host --rm -i \
+          gcr.io/cloud-ml-auto-solutions/jetstream-maxtext-stable-stack/tpu:github_${{ github.run_id }} \
+          bash -c "
+            bash JetStream/.github/workflows/test_moe_benchmarks.sh
+          "
     - name: Upload build artifact
       uses: actions/upload-artifact@v4
       with:
         name: benchmark_report
         path: ${{ env.OUTPUT_DIR }}
 
-  clean_up:
-    if: ${{ always() }}  # always execute, regardless of previous jobs or steps.
+  clean_up_on_fail:
+    if: ${{ failure() }}
     needs: [build_stable_stack, benchmark_report]
     name: "Clean up"
     runs-on: ["self-hosted"]
     permissions:
       contents: read
       issues: write  # for failed-build-issue
     steps:
+    - name: Authenticate gcloud
+      run: gcloud auth configure-docker gcr.io --quiet
     - name: Delete TPU image
       # sync with the image uploaded from build_stable_stack stage
-      run: gcloud container images delete gcr.io/cloud-tpu-inference-test/jetstream-maxtext-stable-stack/tpu:github_${{ github.run_id }} --force-delete-tags --quiet
+      run: gcloud container images delete gcr.io/cloud-ml-auto-solutions/jetstream-maxtext-stable-stack/tpu:github_${{ github.run_id }} --force-delete-tags --quiet
+
+  tag_night_image:
+    needs: [build_stable_stack, benchmark_report]
+    name: "Tag night image"
+    runs-on: ["self-hosted"]
+    permissions:
+      contents: read
+      issues: write  # for failed-build-issue
+    steps:
+    - name: Authenticate gcloud
+      run: gcloud auth configure-docker gcr.io --quiet
+    - name: Upload night image
+      # sync with the image uploaded from build_stable_stack stage
+      run: |
+        UPLOAD_IMAGE_TAG=gcr.io/cloud-ml-auto-solutions/jetstream-maxtext-stable-stack/tpu:github_${{ github.run_id }}
+        NIGHTLY_TAG=${UPLOAD_IMAGE_TAG%:*}:nightly
+        NIGHTLY_TAG_DATE=${NIGHTLY_TAG}-$(date +"%Y%m%d")
+        gcloud container images add-tag ${UPLOAD_IMAGE_TAG} ${NIGHTLY_TAG} --quiet
+        gcloud container images add-tag ${UPLOAD_IMAGE_TAG} ${NIGHTLY_TAG_DATE} --quiet
+        gcloud container images untag ${UPLOAD_IMAGE_TAG} --quiet
 
   notify:
     name: Notify test build # creates an issue or modifies last open existing issue for failed build
     needs: [build_stable_stack, benchmark_report]
     runs-on: ["self-hosted", "tpu", "v6e-8"]
+    env:
+      BENCHMARK_REPORT_DIR: ./benchmark_report
     steps:
+    - name: Clean previous artifact
+      run: rm -rf ${{ env.OUTPUT_DIR }}
     - name: Download benchmark artifact
       uses: actions/download-artifact@v4
       with:
         name: benchmark_report
-        path: ./benchmark_report
+        path: ${{ env.OUTPUT_DIR }}
     - name: Check whether one of the jobs failed
       if: ${{ failure() }}
       uses: jayqi/failed-build-issue-action@1a893bbf43ef1c2a8705e2b115cd4f0fe3c5649b  # v1.2.0
@@ -133,5 +161,5 @@ jobs:
         to: [email protected], [email protected], [email protected]
         from: JetStream Runs
         secure: true
-        attachments: ./benchmark_report/moe_8x7b.txt,./benchmark_report/moe_8x22b.txt,./benchmark_report/moe_8x22b_long_context_8k_prefill.txt,./benchmark_report/moe_8x7b_jetstream.txt
+        attachments: ${{ env.OUTPUT_DIR }}/moe_8x7b.txt,${{ env.OUTPUT_DIR }}/moe_8x22b.txt,${{ env.OUTPUT_DIR }}/moe_8x22b_long_context_8k_prefill.txt,${{ env.OUTPUT_DIR }}/moe_8x7b_jetstream.txt
         body: workflow for ${{github.repository}} completed successfully!

experimental/jetstream-maxtext-stable-stack/test_script/benchmark_serving_example.sh

@@ -1,5 +1,7 @@
-# TODO: need a public path
-export PARAM_PATH=${PARAM_PATH}
+#!/bin/bash
+
+SERVER_PID=""
+CLIENT_PID=""
 
 python -c "import nltk; nltk.download('punkt')"
 
@@ -26,8 +28,10 @@ python -m MaxText.maxengine_server \
   model_call_mode=inference \
   sparse_matmul=False \
   use_chunked_prefill=true \
-  prefill_chunk_size=64 \
-  load_parameters_path=${PARAM_PATH} &
+  prefill_chunk_size=256 \
+  load_parameters_path=gs://jetstream-runner/8-7B-int8 &
+
+SERVER_PID=$!
 
 popd
 
@@ -41,4 +45,21 @@ python ./JetStream/benchmarks/benchmark_serving.py \
   --num-prompts 100 \
   --max-output-length 2048 \
   --dataset openorca \
-  --run-eval True
+  --run-eval True &
+
+CLIENT_PID=$!
+
+while true; do
+  # If server is not running, it is crash. Terminate the script.
+  if ! kill -0 "${SERVER_PID}" 2>/dev/null; then
+    exit 1
+  fi
+
+  # If client is done
+  if ! kill -0 "${CLIENT_PID}" 2>/dev/null; then
+    wait $CLIENT_PID
+    exit $?
+  fi
+
+  sleep 1
+done