Merge branch 'main' into zxhe/multi-container

Author: zxhe-sean

.github/actions/install-kjob/action.yml

@@ -44,7 +44,6 @@ runs:
       run: gcloud auth configure-docker --quiet
       shell: bash
     - uses: ./.github/actions/install-kueue
-    - uses: ./.github/actions/install-kjob
     - name: Install XPK
       run: pip install dist/xpk-*.whl
       shell: bash

.github/actions/setup-test-env/action.yml

@@ -49,14 +49,13 @@ jobs:
         lookup-only: true
     - name: install dependencies
       if : steps.check-cache.outputs.cache-hit != 'true'
-      run: make install-dev && cp ./bin/kubectl-kueue /usr/local/bin/kubectl-kueue && cp ./bin/kubectl-kjob /usr/local/bin/kubectl-kjob
+      run: make install-dev && cp ./bin/kubectl-kueue /usr/local/bin/kubectl-kueue
     - name: Cache dependencies
       if : steps.check-cache.outputs.cache-hit != 'true'
       uses: actions/cache/save@v3
       with:
         path: |
           /usr/local/bin/kubectl-kueue
-          /usr/local/bin/kubectl-kjob
           ~/.cache/pip
           ${{env.pythonLocation}}
         key: xpk-deps-${{ matrix.python-version }}-${{github.run_id}}-${{github.run_attempt}}

.github/workflows/build_tests.yaml

@@ -152,62 +152,6 @@ jobs:
               run: xpk info --cluster $TPU_CLUSTER_NAME --zone=us-central2-b
             - name: Delete the workload on the cluster
               run: xpk workload delete --workload $WORKLOAD_NAME --cluster $TPU_CLUSTER_NAME --zone=us-central2-b
-            - name: Create test script to execute in batch
-              run: echo -e '#!/bin/bash \n#SBATCH --unknown-flag=value\n echo "Hello world from a test script!"' > batch.sh
-            - name: Run a batch job on the cluster
-              run: xpk batch --cluster $TPU_CLUSTER_NAME --zone=us-central2-b batch.sh --ignore-unknown-flags --array 1-5 --nodes 2 --ntasks 3
-            - name: List out the jobs on the cluster
-              run: xpk job ls --cluster $TPU_CLUSTER_NAME --zone=us-central2-b | grep 'xpk-def-app-profile-slurm-'
-            - name: Get created job name
-              run: |
-                  JOB_NAME=$(xpk job ls --cluster $TPU_CLUSTER_NAME --zone=us-central2-b | grep 'xpk-def-app-profile-slurm-' | grep 'multislice-queue' | head -1 | awk '{print $1}')
-                  echo "JOB_NAME=${JOB_NAME}" >> $GITHUB_ENV
-            - name: Check job spec
-              run: |
-                  job_spec=$(kubectl get job ${JOB_NAME} -o jsonpath='{.spec}')
-                  echo "$job_spec" | grep '"completions":2'
-                  echo "$job_spec" | grep '"parallelism":2'
-                  echo "$job_spec" | jq '.template.spec.containers | length' | grep 3
-            - name: Get job info for the last job created on the cluster
-              run: xpk job info ${JOB_NAME} --cluster $TPU_CLUSTER_NAME --zone=us-central2-b | grep -e "Entrypoint environment variables template:" -e "Job name:" -e "Labels:" -e "Mounts:" -e "Pods:" -e "Profile:" -e "Script name:" | wc -l | grep "7"
-            - name: Cancel the batch job on the cluster
-              run: xpk job cancel ${JOB_NAME} --cluster $TPU_CLUSTER_NAME --zone=us-central2-b | grep "job.batch/${JOB_NAME} deleted"
-            - name: Create shell and exit it immediately
-              run: |
-                  cat <<EOF > create-shell.exp
-                  #!/usr/bin/expect
-                  set timeout 180
-                  spawn sh -c "xpk shell --cluster $TPU_CLUSTER_NAME --zone=us-central2-b | tee shell.log"
-                  send "\n"
-                  expect {
-                    "/ # " {
-                        send "exit\n"
-                        # Wait for EOF after exit
-                        expect eof
-                        exit 0
-                    }
-                    timeout {
-                        puts "Timed out waiting for pod to be running"
-                        exit 1
-                    }
-                    eof {
-                        puts "Unexpected EOF before getting prompt"
-                        exit 1
-                    }
-                  }
-                  EOF
-                  chmod +x ./create-shell.exp
-                  expect ./create-shell.exp
-            - name: Check if shell exists and is running
-              run: |
-                  pod_name=$(grep 'waiting for pod' shell.log | awk -F'"' '{print $2}')
-                  kubectl wait --for='jsonpath={.status.conditions[?(@.type=="Ready")].status}=True' --timeout=1m pod/${pod_name}
-            - name: Stop the shell
-              run: xpk shell stop --cluster $TPU_CLUSTER_NAME --zone=us-central2-b
-            - name: Delete create-shell.exp file
-              run: rm create-shell.exp
-            - name: Delete shell.log file
-              run: rm shell.log
             - name: Delete the cluster created
               if: always()
               run: xpk cluster delete --cluster $TPU_CLUSTER_NAME --zone=us-central2-b --force

.github/workflows/integration_basic_cluster_create.yaml

@@ -0,0 +1,78 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+name: Basic GPU cluster create
+
+on:
+    workflow_call:
+
+permissions:
+  contents: read
+
+jobs:
+    gpu-cluster-create-and-delete:
+        runs-on: [ubuntu-22.04]
+        concurrency:
+            group: nightly-test-cluster-group-gpu
+            cancel-in-progress: false
+        env:
+            GPU_CLUSTER_NAME: nightly-xpk-b200
+            WORKLOAD_NAME: xpktest-gpu-nightly-${{ github.run_attempt }}
+        steps:
+            - uses: actions/download-artifact@v4
+              with:
+                  name: custom-scripts
+            - name: Setup environment
+              uses: ./.github/actions/setup-test-env
+              with:
+                  credentials_json: "${{ secrets.GCP_SA_KEY }}"
+            - name: Check xpk installation
+              run: xpk version
+            - name: 'Setup Service Account for XPK'
+              run: |
+                # 1. Clear any existing WIF configurations to avoid conflicts
+                rm -rf $HOME/.config/gcloud
+                mkdir -p $HOME/.config/gcloud
+
+                # 2. Write the Key File
+                echo '${{ secrets.GCP_SA_KEY }}' > $HOME/.config/gcloud/application_default_credentials.json
+
+                # 3. Activate the Service Account
+                # This updates the internal config files to point to the key file.
+                # When Docker mounts the directory, it will now see "Active Account: Service Account"
+                gcloud auth activate-service-account --key-file=$HOME/.config/gcloud/application_default_credentials.json --project=cloud-tpu-multipod-dev
+
+                # 4. Set Env Var for the host (GitHub Runner)
+                echo "GOOGLE_APPLICATION_CREDENTIALS=$HOME/.config/gcloud/application_default_credentials.json" >> $GITHUB_ENV
+            - name: Create an XPK Cluster with 1 x b200 GPU
+              run: xpk cluster create --cluster $GPU_CLUSTER_NAME --device-type=b200-8 --zone=asia-northeast1-b --default-pool-cpu-machine-type=n1-standard-16 --spot
+            - name: Authenticate Docker
+              run: gcloud auth configure-docker --quiet
+            - name: Run a base-docker-image workload
+              run: xpk workload create --cluster $GPU_CLUSTER_NAME --workload $WORKLOAD_NAME --docker-image='nvidia/cuda:12.1.0-base-ubuntu22.04' --command "nvidia-smi" --zone=asia-northeast1-b --device-type=b200-8
+            - name: List out the workloads on the cluster
+              run: xpk workload list --cluster $GPU_CLUSTER_NAME --zone=asia-northeast1-b
+            - name: Wait for workload completion and confirm it succeeded
+              run: xpk workload list --cluster $GPU_CLUSTER_NAME --zone=asia-northeast1-b --wait-for-job-completion $WORKLOAD_NAME --timeout 600
+            - name: Delete the workload on the cluster
+              run: xpk workload delete --workload $WORKLOAD_NAME --cluster $GPU_CLUSTER_NAME --zone=asia-northeast1-b
+            - name: Delete the cluster created
+              if: always()
+              run: xpk cluster delete --cluster $GPU_CLUSTER_NAME --zone=asia-northeast1-b --force
+            - name: Upload cluster nodepool creation log
+              if: always()
+              uses: actions/upload-artifact@v4
+              with:
+                  name: gpu-cluster-nodepool-log-${{github.run_id}}
+                  path: /tmp/NodepoolCreate-${{ env.GPU_CLUSTER_NAME }}-np-*

.github/workflows/integration_gpu_cluster_create.yaml

@@ -47,14 +47,13 @@ jobs:
                   lookup-only: true
             - name: install dependencies
               if: steps.check-cache.outputs.cache-hit != 'true'
-              run: make install-dev && cp ./bin/kubectl-kueue /usr/local/bin/kubectl-kueue && cp ./bin/kubectl-kjob /usr/local/bin/kubectl-kjob
+              run: make install-dev && cp ./bin/kubectl-kueue /usr/local/bin/kubectl-kueue
             - name: Cache dependencies
               if: steps.check-cache.outputs.cache-hit != 'true'
               uses: actions/cache/save@v3
               with:
                   path: |
                       /usr/local/bin/kubectl-kueue
-                      /usr/local/bin/kubectl-kjob
                       ~/.cache/pip
                       ${{env.pythonLocation}}
                   key: xpk-deps-${{ matrix.python-version }}-${{github.run_id}}-${{github.run_attempt}}

.github/workflows/integration_legacy_tests.yaml

@@ -36,8 +36,8 @@ jobs:
       with:
         mode: minimum
         count: 1
-        labels: "release-improvements, release-bugfix, release-features"
-        message: "This PR is being prevented from merging because it is not labeled.  Please add a label to this PR.  Accepted labels: release-improvements, release-bugfix, release-features"
+        labels: "release-improvements, release-bugfix, release-features, release-breaking"
+        message: "This PR is being prevented from merging because it is not labeled.  Please add a label to this PR.  Accepted labels: release-improvements, release-bugfix, release-features, release-breaking"
     - id: do-not-merge
       uses: mheap/github-action-required-labels@v5
       with:

.github/workflows/label-validation.yaml

@@ -16,38 +16,41 @@ name: Nightly Tests
 
 on:
   workflow_dispatch:
-  schedule: # Schedule the job run at 12AM PST daily.
-    - cron: "0 8 * * *"
+  schedule: # Schedule the job run at 6AM UTC daily.
+    - cron: "0 6 * * *"
 
 permissions:
   contents: read
 
 jobs:
-  build_kjob:
-    uses: ./.github/workflows/reusable_build_kjob.yaml
   build_wheel:
     uses: ./.github/workflows/reusable_build_wheel.yaml
   build_actions:
     uses: ./.github/workflows/reusable_build_scripts.yaml
   basic_cluster_create:
-    needs: [build_kjob, build_actions, build_wheel]
+    needs: [build_actions, build_wheel]
     uses: ./.github/workflows/integration_basic_cluster_create.yaml
     secrets: inherit
 
+  gpu_cluster_create:
+    needs: [build_actions, build_wheel]
+    uses: ./.github/workflows/integration_gpu_cluster_create.yaml
+    secrets: inherit
+
   pathways_cluster_create:
-    needs: [build_kjob, build_actions, build_wheel]
+    needs: [build_actions, build_wheel]
     uses: ./.github/workflows/integration_pathways_cluster_create.yaml
     secrets: inherit
 
   ray_cluster_create:
-    needs: [build_kjob, build_actions, build_wheel]
+    needs: [build_actions, build_wheel]
     uses: ./.github/workflows/integration_ray_cluster_create.yaml
     secrets: inherit
   legacy_integration:
-    needs: [build_kjob, build_actions, build_wheel]
+    needs: [build_actions, build_wheel]
     uses: ./.github/workflows/integration_legacy_tests.yaml
     secrets: inherit
   storage-tests:
-    needs: [build_kjob, build_actions, build_wheel]
+    needs: [build_actions, build_wheel]
     uses: ./.github/workflows/integration_storage_tests.yaml
     secrets: inherit

.github/workflows/nightly_tests.yaml

@@ -33,7 +33,6 @@ jobs:
       with:
         path: |
           /usr/local/bin/kubectl-kueue
-          /usr/local/bin/kubectl-kjob
           ~/.cache/pip
           ${{env.pythonLocation}}
         key: xpk-deps-3.10-${{github.run_id}}-${{github.run_attempt}}