Fix max-nodes when creating flex queued nodepool of tpus (#541)

* fix max-nodes when creating tpu dws flex queued nodepools

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

---------

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

Author: pawloch00

src/xpk/core/capacity.py

@@ -173,7 +173,7 @@ def verify_reservation_exists(args) -> int:
 
 
 def get_capacity_arguments_from_capacity_type(
-    args, capacity_type: CapacityType
+    args, capacity_type: CapacityType, max_nodes: int
 ) -> tuple[str, int]:
   """Determine the Nodepool creation capacity arguments needed.
 
@@ -197,7 +197,7 @@ def get_capacity_arguments_from_capacity_type(
       capacity_args = (
           ' --flex-start --enable-queued-provisioning --enable-autoscaling'
           ' --location-policy=ANY --reservation-affinity=none'
-          ' --no-enable-autorepair --max-nodes=1'
+          f' --no-enable-autorepair --max-nodes={max_nodes}'
       )
     case CapacityType.RESERVATION:
       capacity_args = (

src/xpk/core/kjob.py

@@ -347,7 +347,6 @@ def create_pod_template_instance(args: Namespace, service_account: str) -> int:
   if pod_image is None or len(pod_image) == 0:
     pod_image = PodTemplateDefaults.IMAGE.value
   working_directory = config.get(KJOB_SHELL_WORKING_DIRECTORY)
-  xpk_print("working directory is: ", working_directory)
   if working_directory is None or len(working_directory) == 0:
     working_directory = PodTemplateDefaults.WORKING_DIRECTORY.value
 

src/xpk/core/nodepool.py

@@ -77,8 +77,12 @@ def run_gke_node_pool_create_command(
     if return_code > 0:
       xpk_print('Listing all reservations failed!')
     return_code = 1
+  if system.accelerator_type == AcceleratorType['TPU']:
+    max_nodes = system.vms_per_slice
+  else:
+    max_nodes = 1000
   capacity_args, return_code = get_capacity_arguments_from_capacity_type(
-      args, capacity_type
+      args, capacity_type, max_nodes
   )
   if return_code > 0:
     xpk_print('Parsing capacity arguments failed!')