fix: Don't check flex in is_TAS_possible

src/xpk/commands/common.py

@@ -51,16 +51,8 @@ def set_cluster_command(args) -> int:
 def is_TAS_possible(
     system_characteristics: SystemCharacteristics | None,
     capacity_type: CapacityType | None,
-    flex: bool,
 ) -> bool:
-  """Check cluster's machine_type and capacity type to determine if Kueue TAS is possible
-
-  Args:
-    args: user provided arguments for running the command.
-
-  Returns:
-    True if possible and False otherwise.
-  """
+  """Check cluster's machine_type and capacity type to determine if Kueue TAS is possible"""
 
   if is_dry_run():
     return True
@@ -73,13 +65,7 @@ def is_TAS_possible(
     xpk_print('capacity_type data was not found in configmaps.')
     xpk_exit(1)
 
-  if not flex:
-    return False
-
-  if (
-      system_characteristics.device_type == H100_MEGA_DEVICE_TYPE
-      and capacity_type != CapacityType.RESERVATION
-  ):
-    return False
-
-  return True
+  return (
+      system_characteristics.device_type != H100_MEGA_DEVICE_TYPE
+      or capacity_type == CapacityType.RESERVATION
+  )

src/xpk/commands/kjob_common.py

@@ -54,7 +54,7 @@ def add_gpu_networking_annotations_to_command(args, cmd: str) -> str:
 def add_TAS_annotations_to_command(args, cmd: str) -> str:
   system_characteristics = get_cluster_system_characteristics(args)
   capacity_type = get_cluster_capacity_type(args)
-  if is_TAS_possible(system_characteristics, capacity_type, flex=False):
+  if is_TAS_possible(system_characteristics, capacity_type):
     cmd += f" --pod-template-annotation {Kueue_TAS_annotation}"
 
   return cmd

src/xpk/commands/workload.py

@@ -53,9 +53,6 @@
     try_to_delete_pathwaysjob_first,
 )
 from ..core.resources import get_cluster_capacity_type, get_cluster_system_characteristics
-from ..core.capacity import (
-    CapacityType,
-)
 from ..core.resources import CLUSTER_METADATA_CONFIGMAP, get_cluster_configmap
 from ..core.scheduling import (
     check_if_workload_can_schedule,
@@ -484,16 +481,12 @@ def workload_create(args) -> None:
     capacity_type = get_cluster_capacity_type(args)
 
     annotations = (
-        ''
-        if not is_TAS_possible(
-            system_characteristics,
-            capacity_type,
-            flex=True if capacity_type == CapacityType.FLEX_START else False,
-        )
-        else (
+        (
             'kueue.x-k8s.io/podset-preferred-topology:'
             ' "cloud.google.com/gce-topology-host"'
         )
+        if is_TAS_possible(system_characteristics, capacity_type)
+        else ''
     )
 
     if (