test deepseek perf tests on GKE

Author: zhengkezhou1

.github/workflows/gke-connectivity-smoke.yml

@@ -36,10 +36,32 @@ on:
         default: "default"
         type: string
       run_workload:
-        description: "Create and run a short multi-host TPU Kubernetes Job."
+        description: "Create and run a multi-host TPU Kubernetes Job."
         required: true
         default: false
         type: boolean
+      workload_type:
+        description: "Workload to run when run_workload is true."
+        required: true
+        default: "connectivity"
+        type: choice
+        options:
+          - connectivity
+          - deepseek_v2_lite_perf
+      bucket_name:
+        description: "GCS bucket mounted at /models for model cache."
+        required: true
+        default: "inference-model-storage-sgl"
+        type: string
+      service_account_name:
+        description: "Kubernetes service account for GCSFuse access."
+        required: true
+        default: "gcs-account"
+        type: string
+      repo_ref:
+        description: "Git ref to test. Defaults to current workflow SHA."
+        required: false
+        type: string
 
 concurrency:
   group: gke-connectivity-smoke-${{ github.ref }}
@@ -52,7 +74,7 @@ permissions:
 jobs:
   connectivity:
     runs-on: ubuntu-latest
-    timeout-minutes: 60
+    timeout-minutes: 180
     env:
       PROJECT_ID: ${{ inputs.project_id || vars.GCP_PROJECT_ID }}
       CLUSTER_NAME: ${{ inputs.cluster_name || vars.GKE_CLUSTER_NAME }}
@@ -62,6 +84,10 @@ jobs:
       NODE_POOL: ${{ inputs.node_pool || 'v6e-4x4-flex-perf-tests' }}
       NAMESPACE: ${{ inputs.namespace || 'default' }}
       RUN_WORKLOAD: ${{ inputs.run_workload || github.event_name == 'push' }}
+      WORKLOAD_TYPE: ${{ inputs.workload_type || 'connectivity' }}
+      BUCKET_NAME: ${{ inputs.bucket_name || 'inference-model-storage-sgl' }}
+      K8S_SERVICE_ACCOUNT: ${{ inputs.service_account_name || 'gcs-account' }}
+      REPO_REF: ${{ inputs.repo_ref || github.sha }}
       WORKLOAD_NAME: gke-tpu-smoke-${{ github.run_id }}
     steps:
       - name: Validate smoke test configuration
@@ -77,6 +103,15 @@ jobs:
             fi
           done
 
+          if [[ "${RUN_WORKLOAD}" == "true" && "${WORKLOAD_TYPE}" == "deepseek_v2_lite_perf" ]]; then
+            for name in BUCKET_NAME K8S_SERVICE_ACCOUNT; do
+              if [[ -z "${!name}" ]]; then
+                echo "::error::${name} is required for deepseek_v2_lite_perf."
+                missing=1
+              fi
+            done
+          fi
+
           if [[ "${missing}" -ne 0 ]]; then
             exit 1
           fi
@@ -86,6 +121,9 @@ jobs:
           echo "location=${CLUSTER_LOCATION}"
           echo "node_pool=${NODE_POOL}"
           echo "namespace=${NAMESPACE}"
+          echo "run_workload=${RUN_WORKLOAD}"
+          echo "workload_type=${WORKLOAD_TYPE}"
+          echo "repo_ref=${REPO_REF}"
 
       - name: Authenticate to Google Cloud
         if: env.SERVICE_ACCOUNT == ''
@@ -143,11 +181,15 @@ jobs:
         run: |
           set -euo pipefail
 
-          cat > /tmp/tpu-smoke.yaml <<EOF
+          HEADLESS_SERVICE_NAME="${WORKLOAD_NAME}-headless"
+          JOB_COMPLETE_TIMEOUT="5m"
+
+          if [[ "${WORKLOAD_TYPE}" == "connectivity" ]]; then
+            cat > /tmp/tpu-smoke.yaml <<EOF
           apiVersion: v1
           kind: Service
           metadata:
-            name: ${WORKLOAD_NAME}-headless
+            name: ${HEADLESS_SERVICE_NAME}
             namespace: ${NAMESPACE}
           spec:
             clusterIP: None
@@ -168,7 +210,7 @@ jobs:
             ttlSecondsAfterFinished: 600
             template:
               spec:
-                subdomain: ${WORKLOAD_NAME}-headless
+                subdomain: ${HEADLESS_SERVICE_NAME}
                 restartPolicy: Never
                 nodeSelector:
                   cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
@@ -197,20 +239,162 @@ jobs:
                     limits:
                       google.com/tpu: "4"
           EOF
+          elif [[ "${WORKLOAD_TYPE}" == "deepseek_v2_lite_perf" ]]; then
+            JOB_COMPLETE_TIMEOUT="30m"
+            TPU_PROCESS_ADDRESSES="${WORKLOAD_NAME}-0.${HEADLESS_SERVICE_NAME}:8471,${WORKLOAD_NAME}-1.${HEADLESS_SERVICE_NAME}:8471,${WORKLOAD_NAME}-2.${HEADLESS_SERVICE_NAME}:8471,${WORKLOAD_NAME}-3.${HEADLESS_SERVICE_NAME}:8471"
+            TPU_WORKER_HOSTNAMES="${WORKLOAD_NAME}-0.${HEADLESS_SERVICE_NAME},${WORKLOAD_NAME}-1.${HEADLESS_SERVICE_NAME},${WORKLOAD_NAME}-2.${HEADLESS_SERVICE_NAME},${WORKLOAD_NAME}-3.${HEADLESS_SERVICE_NAME}"
+
+            cat > /tmp/tpu-smoke.yaml <<EOF
+          apiVersion: v1
+          kind: Service
+          metadata:
+            name: ${HEADLESS_SERVICE_NAME}
+            namespace: ${NAMESPACE}
+          spec:
+            clusterIP: None
+            selector:
+              job-name: ${WORKLOAD_NAME}
+            ports:
+            - name: http
+              port: 30000
+            - name: dist-init
+              port: 10011
+            - name: tpu-process
+              port: 8471
+            - name: control
+              port: 18080
+          ---
+          apiVersion: batch/v1
+          kind: Job
+          metadata:
+            name: ${WORKLOAD_NAME}
+            namespace: ${NAMESPACE}
+          spec:
+            completionMode: Indexed
+            parallelism: 4
+            completions: 4
+            backoffLimit: 0
+            activeDeadlineSeconds: 7200
+            ttlSecondsAfterFinished: 600
+            template:
+              metadata:
+                annotations:
+                  gke-gcsfuse/volumes: "true"
+              spec:
+                subdomain: ${HEADLESS_SERVICE_NAME}
+                restartPolicy: Never
+                serviceAccountName: ${K8S_SERVICE_ACCOUNT}
+                nodeSelector:
+                  cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
+                  cloud.google.com/gke-tpu-topology: 4x4
+                  cloud.google.com/gke-nodepool: ${NODE_POOL}
+                tolerations:
+                - key: google.com/tpu
+                  operator: Equal
+                  value: present
+                  effect: NoSchedule
+                containers:
+                - name: deepseek-v2-lite-perf
+                  image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:jax0.8.1-rev1
+                  command:
+                  - /bin/bash
+                  - -lc
+                  - |
+                    set -euxo pipefail
+                    git clone https://github.com/sgl-project/sglang-jax.git /tmp/sglang-jax
+                    cd /tmp/sglang-jax
+                    git fetch origin "${REPO_REF}" || true
+                    git checkout "${REPO_REF}"
+                    python3 -m pip install --upgrade pip
+                    python3 -m pip install -e "python[all]"
+                    python3 test/srt/nightly-test/gke_deepseek_v2_lite_perf_entrypoint.py
+                  env:
+                  - name: JOB_COMPLETION_INDEX
+                    valueFrom:
+                      fieldRef:
+                        fieldPath: metadata.labels['batch.kubernetes.io/job-completion-index']
+                  - name: TPU_WORKER_ID
+                    valueFrom:
+                      fieldRef:
+                        fieldPath: metadata.labels['batch.kubernetes.io/job-completion-index']
+                  - name: TPU_PROCESS_PORT
+                    value: "8471"
+                  - name: TPU_PROCESS_ADDRESSES
+                    value: "${TPU_PROCESS_ADDRESSES}"
+                  - name: TPU_WORKER_HOSTNAMES
+                    value: "${TPU_WORKER_HOSTNAMES}"
+                  - name: WORKLOAD_NAME
+                    value: "${WORKLOAD_NAME}"
+                  - name: HEADLESS_SERVICE_NAME
+                    value: "${HEADLESS_SERVICE_NAME}"
+                  - name: NNODES
+                    value: "4"
+                  - name: CI_MOUNT_ROOT
+                    value: "/models"
+                  - name: SGLANG_JAX_MODEL_CACHE
+                    value: "/models/model_scope"
+                  - name: JAX_COMPILATION_CACHE_DIR
+                    value: "/tmp/jax_compilation_cache"
+                  - name: HF_HUB_DOWNLOAD_TIMEOUT
+                    value: "600"
+                  ports:
+                  - containerPort: 30000
+                    name: http
+                  - containerPort: 10011
+                    name: dist-init
+                  - containerPort: 8471
+                    name: tpu-process
+                  - containerPort: 18080
+                    name: control
+                  resources:
+                    requests:
+                      google.com/tpu: "4"
+                    limits:
+                      google.com/tpu: "4"
+                  volumeMounts:
+                  - mountPath: /models
+                    name: model-storage
+                    readOnly: true
+                  - mountPath: /dev/shm
+                    name: dev-shm
+                  - mountPath: /tmp/jax_compilation_cache
+                    name: jax-cache
+                volumes:
+                - name: dev-shm
+                  emptyDir:
+                    medium: Memory
+                    sizeLimit: 64Gi
+                - name: jax-cache
+                  emptyDir: {}
+                - name: gke-gcsfuse-cache
+                  emptyDir: {}
+                - name: model-storage
+                  csi:
+                    driver: gcsfuse.csi.storage.gke.io
+                    readOnly: true
+                    volumeAttributes:
+                      bucketName: ${BUCKET_NAME}
+                      mountOptions: implicit-dirs,metadata-cache:ttl-secs:-1,metadata-cache:stat-cache-max-size-mb:-1,metadata-cache:type-cache-max-size-mb:-1,file-cache:max-size-mb:409600,file-cache:cache-file-for-range-read:true,file-system:kernel-list-cache-ttl-secs:-1,file-cache:enable-parallel-downloads:true,file-cache:parallel-downloads-per-file:128,read_ahead_kb=1024
+          EOF
+          else
+            echo "::error::Unsupported WORKLOAD_TYPE=${WORKLOAD_TYPE}"
+            exit 1
+          fi
 
           kubectl apply -f /tmp/tpu-smoke.yaml
           kubectl get job "${WORKLOAD_NAME}" --namespace "${NAMESPACE}" -o wide
           kubectl wait --for=condition=Ready pod \
             --selector="job-name=${WORKLOAD_NAME}" \
             --namespace "${NAMESPACE}" \
-            --timeout=30m
+            --timeout=45m
           kubectl logs --selector="job-name=${WORKLOAD_NAME}" \
             --namespace "${NAMESPACE}" \
             --all-containers=true \
-            --prefix=true
+            --prefix=true \
+            --follow=true
           kubectl wait --for=condition=complete "job/${WORKLOAD_NAME}" \
             --namespace "${NAMESPACE}" \
-            --timeout=5m
+            --timeout="${JOB_COMPLETE_TIMEOUT}"
 
       - name: Diagnose TPU workload failure
         if: failure() && env.RUN_WORKLOAD == 'true'