Merge pull request #615 from AI-Hypercomputer/tas-support

Add TAS support for workloads on DWS clusters

Author: FIoannides

src/xpk/commands/common.py

@@ -71,7 +71,7 @@ def is_TAS_possible(
     xpk_print('capacity_type data was not found in configmaps.')
     xpk_exit(1)
 
-  if flex:
+  if not flex:
     return False
 
   if (

src/xpk/commands/workload.py

@@ -226,7 +226,8 @@
             metadata:
               labels:
                 xpk.google.com/workload: {args.workload}
-              annotations: {annotations}
+              annotations:
+                {annotations}
             spec:
               priorityClassName: {args.priority}
               restartPolicy: Never

src/xpk/core/kueue.py

@@ -89,6 +89,10 @@
   name: dws-config
 spec:
   provisioningClassName: queued-provisioning.gke.io
+  podSetUpdates:
+    nodeSelector:
+    - key: autoscaling.gke.io/provisioning-request
+      valueFromProvisioningClassDetail: ResizeRequestName
   managedResources:
   - {managed_resource}
 ---