fix nightly tests (#214)

Signed-off-by: Piotr Pawłowski <[email protected]>

Author: pawloch00

.github/workflows/nightly_tests.yaml

@@ -53,15 +53,15 @@ jobs:
         gcloud config get compute/zone
     - name: Install kueuectl
       run: curl -Lo ./kubectl-kueue https://github.com/kubernetes-sigs/kueue/releases/download/v0.8.1/kubectl-kueue-linux-amd64 && chmod +x ./kubectl-kueue && mv ./kubectl-kueue /usr/local/bin/kubectl-kueue
+    - name: Install xpk with pip and verify it executes corretly
+      run: |
+        pip install .
+        xpk --help
     - name: Create an XPK Cluster with zero node pools
       run: python xpk.py cluster create --cluster $EMPTY_CLUSTER_NAME --tpu-type=v4-8  --num-slices=0 --zone=us-central2-b --default-pool-cpu-machine-type=n1-standard-16 --reservation='${{ secrets.GCP_TPU_V4_RESERVATION }}' --custom-cluster-arguments='${{ secrets.CLUSTER_ARGUMENTS }}'
     - name: Delete the cluster created
       run: python xpk.py cluster delete --cluster $EMPTY_CLUSTER_NAME --zone=us-central2-b
       if: always()
-    - name: Install xpk with pip and verify it executes corretly
-      run: |
-        pip install .
-        xpk --help
     - name: Create an XPK Cluster with 2x v4-8 nodepools
       run: python xpk.py cluster create --cluster $TPU_CLUSTER_NAME --tpu-type=v4-8  --num-slices=2 --zone=us-central2-b --default-pool-cpu-machine-type=n1-standard-16 --reservation='${{ secrets.GCP_TPU_V4_RESERVATION }}' --custom-cluster-arguments='${{ secrets.CLUSTER_ARGUMENTS }}'
     - name: Authenticate Docker
@@ -85,7 +85,7 @@ jobs:
       run: python xpk.py cluster delete --cluster $TPU_CLUSTER_NAME --zone=us-central2-b
 
   pw-cluster-and-workload:
-    runs-on: [ubuntu-20.04]
+    runs-on: [ubuntu-22.04]
     concurrency: # We support one build test to run at a time currently.
       group: nightly-pw-test-cluster-group
       cancel-in-progress: false
@@ -105,6 +105,12 @@ jobs:
       run: |
         gcloud config set compute/zone us-east4-a
         gcloud config get compute/zone
+    - name: Install xpk with pip and verify it executes corretly
+      run: |
+        pip install .
+        xpk --help
+    - name: Install kueuectl
+      run: curl -Lo ./kubectl-kueue https://github.com/kubernetes-sigs/kueue/releases/download/v0.8.1/kubectl-kueue-linux-amd64 && chmod +x ./kubectl-kueue && mv ./kubectl-kueue /usr/local/bin/kubectl-kueue
     - name: Create an Pathways-enabled XPK Cluster with 2 x v4-8 nodepools
       run: python xpk.py cluster create-pathways --cluster $PATHWAYS_TPU_CLUSTER_NAME --tpu-type=v4-8  --num-slices=2 --zone=us-central2-b --default-pool-cpu-machine-type=n1-standard-16 --default-pool-cpu-num-nodes=16 --reservation='${{ secrets.GCP_TPU_V4_RESERVATION }}'
     - name: Create test script to execute in workloads