"Select TPU by topology (#525)" + Fix errors (#563)

* fix yaml formatting for workloads with TPU and NAP

* refactor tpu system characteristics

internal representation of TPU machines has changed, so now grep that used the old format fails

* fix device_type issue

---------

Co-authored-by: gcie <[email protected]>

Author: sharabiani

.github/workflows/build_tests.yaml

@@ -40,6 +40,7 @@ jobs:
       group-name: ${{ steps.set-group-name.outputs.group-name }}
       zone: ${{ steps.set-zone.outputs.zone }}
       tpu-type: ${{ steps.set-tpu-type.outputs.tpu-type }}
+      tpu-type-topology: ${{ steps.set-tpu-type-topology.outputs.tpu-type-topology }}
       location: ${{steps.set-location.outputs.location}}
       run-id: ${{steps.set-run-id.outputs.run-id}}
     steps:
@@ -76,6 +77,10 @@ jobs:
       id: set-tpu-type
       run: |
         echo tpu-type=v4-8 >> $GITHUB_OUTPUT
+    - name: set tpu-type-topology
+      id: set-tpu-type-topology
+      run: |
+        echo tpu-type-topology=v4-2x2x1 >> $GITHUB_OUTPUT
     - name: set location
       id: set-location
       run: |    
@@ -152,7 +157,7 @@ jobs:
     with:
       run-id: '${{needs.set-variables.outputs.run-id}}'
       cluster-name: '${{needs.set-variables.outputs.cluster-name}}'
-      tpu-type: '${{needs.set-variables.outputs.tpu-type || inputs.tpu-type}}'
+      tpu-type: '${{needs.set-variables.outputs.tpu-type-topology || inputs.tpu-type}}'
       zone: '${{needs.set-variables.outputs.zone}}'
       location: '${{needs.set-variables.outputs.location}}'
     secrets: inherit
@@ -165,7 +170,7 @@ jobs:
     with:
       cluster-name-dws: '${{needs.set-variables.outputs.cluster-name-dws}}'
       cluster-name: '${{needs.set-variables.outputs.cluster-name}}'
-      tpu-type: '${{needs.set-variables.outputs.tpu-type || inputs.tpu-type}}'
+      tpu-type: '${{needs.set-variables.outputs.tpu-type-topology || inputs.tpu-type}}'
       zone: '${{needs.set-variables.outputs.zone}}'
       location: '${{needs.set-variables.outputs.location}}'
       run-id: '${{needs.set-variables.outputs.run-id}}'
@@ -180,6 +185,7 @@ jobs:
       cluster-name: ${{needs.set-variables.outputs.cluster-name}}
       cluster-name-dws: '${{needs.set-variables.outputs.cluster-name-dws}}'
       tpu-type: ${{needs.set-variables.outputs.tpu-type}}
+      tpu-type-topology: ${{needs.set-variables.outputs.tpu-type-topology}}
       zone: ${{needs.set-variables.outputs.zone}}
       run-id: '${{needs.set-variables.outputs.run-id}}'
     secrets: inherit

.github/workflows/reusable_workload_tests.yaml

@@ -24,6 +24,9 @@ on:
       tpu-type:
         required: true
         type: string
+      tpu-type-topology:
+        required: true
+        type: string
       tpu-type-dws:
         required: false
         type: string
@@ -108,7 +111,7 @@ jobs:
         --docker-password='${{secrets.GCP_SA_KEY}}' \
         --docker-email='${{secrets.GCP_SA_EMAIL}}'
     - name: Run workload with private image
-      run: python xpk.py workload create --cluster ${{inputs.cluster-name}} --workload $PRIVATE_IMAGE_WORKLOAD_NAME  --command "echo foo"  --tpu-type=${{inputs.tpu-type}} --num-slices=1 --zone=${{inputs.zone}} --docker-image=${{secrets.DOCKER_REPO_SERVER}}ubuntu2004 --docker-image-pull-secret=gcr-key
+      run: python xpk.py workload create --cluster ${{inputs.cluster-name}} --workload $PRIVATE_IMAGE_WORKLOAD_NAME  --command "echo foo"  --tpu-type=${{inputs.tpu-type-topology}} --num-slices=1 --zone=${{inputs.zone}} --docker-image=${{secrets.DOCKER_REPO_SERVER}}ubuntu2004 --docker-image-pull-secret=gcr-key
     - name: Wait for private image workload completion and confirm it succeeded
       run: python3 xpk.py workload list --cluster ${{inputs.cluster-name}} --zone=${{inputs.zone}} --wait-for-job-completion $PRIVATE_IMAGE_WORKLOAD_NAME --timeout 300
     - name: Delete kubectl secret

src/xpk/core/capacity.py

@@ -232,9 +232,9 @@ def get_capacity_node_selectors_from_capacity_type(
     case CapacityType.ON_DEMAND.name:
       node_selector = ''
     case CapacityType.FLEX_START.name:
-      node_selector = 'cloud.google.com/gke-queued="true"'
+      node_selector = 'cloud.google.com/gke-queued: "true"'
     case CapacityType.SPOT.name:
-      node_selector = 'cloud.google.com/gke-spot="true"'
+      node_selector = 'cloud.google.com/gke-spot: "true"'
     case CapacityType.RESERVATION.name:
       node_selector = f'cloud.google.com/reservation-name: {args.reservation}'
     case _:

src/xpk/core/scheduling.py

@@ -49,7 +49,7 @@ def check_if_workload_can_schedule(args, system: SystemCharacteristics) -> bool:
   missing_gke_accelerator_type = False
   if not cluster_config_map.get(system.gke_accelerator):
     xpk_print(
-        f'Gke Accelerator Type Check: {args.workload} is requesting'
+        f'GKE Accelerator Type Check: {args.workload} is requesting'
         f' {system.gke_accelerator} but cluster only contains'
         f' {cluster_config_map.keys()}. '
     )