feat: add DWS flex tests

Author: scaliby

.github/workflows/nightly_tests.yaml

@@ -24,6 +24,7 @@ permissions:
 
 env:
   CLUSTER_NETWORK_ARGUMENTS: "--network=${{secrets.NETWORK_NAME}} --subnetwork=${{secrets.SUBNETWORK_NAME}}"
+  CLUSTER_NETWORK_ARGUMENTS_DWS: "--network=${{secrets.NETWORK_NAME}} --subnetwork=${{secrets.SUBNETWORK_NAME_DWS}}"
 
 jobs:
   build_wheel:
@@ -39,7 +40,8 @@ jobs:
     env:
       EMPTY_CLUSTER_NAME: nightly-xpk-zero-nodepools
       PRIVATE_CLUSTER_NAME: nightly-xpk-private-2-v4-8-nodepools
-      TPU_CLUSTER_NAME: nightly-xpk-2-v4-8-nodepools
+      DWS_FLEX_CLUSTER_NAME: xpk-dws-nightly-test-2-v4-8-nodepools
+      TPU_CLUSTER_NAME: nightly-xpk-2-v5p-8-nodepools
       WORKLOAD_NAME: xpktest-nightly-${{ github.run_attempt }}
     steps:
     - uses: actions/download-artifact@v4
@@ -140,6 +142,14 @@ jobs:
     - name: Delete the cluster created
       if: always()
       run: xpk cluster delete --cluster $TPU_CLUSTER_NAME --zone=us-central2-b --force
+    - name: Create a DWS flex queued xpk cluster
+      run: xpk cluster create --cluster ${DWS_FLEX_CLUSTER_NAME} --tpu-type=v5p-8  --num-slices=1 --zone=us-east5-a --default-pool-cpu-num-nodes=2 --flex --custom-cluster-arguments="${CLUSTER_NETWORK_ARGUMENTS_DWS}"
+    - name: Run dws flex queued TPU workload
+      run: xpk workload create --workload xpktest-build-${{ github.run_attempt }}-dws --cluster ${DWS_FLEX_CLUSTER_NAME} --zone=us-east5-a --tpu-type=v5p-8  --flex --command "echo foo" --num-slices=1
+    - name: Wait for workload completion and confirm it succeeded
+      run: xpk workload list --cluster ${DWS_FLEX_CLUSTER_NAME} --zone=us-east5-a --wait-for-job-completion xpktest-build-${{ github.run_attempt }}-dws --timeout 1000
+    - name: Delete the DWS flex queued cluster
+      run: xpk cluster delete --cluster ${DWS_FLEX_CLUSTER_NAME}
     - name: Upload cluster nodepool creation log
       if: always()
       uses: actions/upload-artifact@v4