Merge branch 'develop' into chzheng/docker_image_flag

Author: ycchenzheng

src/xpk/commands/cluster.py

@@ -710,6 +710,9 @@ def cluster_create_ray_cluster(args) -> None:
 
 def install_jq(args):
   """Installs 'jq' utility."""
+  if shutil.which('jq'):
+    xpk_print("Task: 'Install jq' skipped, jq already installed.")
+    return
   command_jq_install = 'sudo apt install jq -y'
   xpk_print("Task: 'Install jq' in progress.")
   return_code = run_command_with_updates(command_jq_install, 'Install jq', args)

src/xpk/core/nap.py

@@ -42,6 +42,8 @@
 management:
   autoRepair: true
   autoUpgrade: true
+scopes:
+  - "https://www.googleapis.com/auth/devstorage.read_write"
 autoprovisioningLocations:
   {zones}
 {resource_limits}
@@ -183,11 +185,11 @@ def create_autoprovisioning_config(
   # is not controlled by NAP.
   cpu_limits = """
   minimum: 1
-  maximum: 10000
+  maximum: 1000000
   """
   memory_limits = """
   minimum: 1
-  maximum: 10000
+  maximum: 10000000
   """
 
   # By default, the maximum chips is set to be the current number of resources used

src/xpk/core/nodepool.py

@@ -32,6 +32,8 @@
     create_or_update_cluster_configmap,
 )
 from .system_characteristics import AcceleratorType
+from functools import reduce
+from operator import mul
 
 CLOUD_PLATFORM_AUTH_SCOPE_URL = (
     '"https://www.googleapis.com/auth/cloud-platform"'
@@ -275,20 +277,24 @@ def run_gke_node_pool_create_command(
         f' --host-maintenance-interval={args.host_maintenance_interval}'
         f' {capacity_args}'
         ' --enable-gvnic'
-        f' {args.custom_nodepool_arguments}'
     )
     if system.accelerator_type == AcceleratorType['TPU']:
       command += f' --node-version={gke_node_pool_version}'
+      topology_product = reduce(
+          mul, (int(x) for x in system.topology.split('x')), 1
+      )
       if capacity_type == CapacityType.FLEX_START:
         command += ' --num-nodes=0'
-      else:
+      elif topology_product > 1:
         command += f' --num-nodes={system.vms_per_slice}'
-      command += ' --placement-type=COMPACT  --max-pods-per-node 15'
       command += (
           f' --scopes=storage-full,gke-default,{CLOUD_PLATFORM_AUTH_SCOPE_URL}'
       )
-      command += f' --tpu-topology={system.topology}'
-      command += f' {args.custom_tpu_nodepool_arguments}'
+
+      if topology_product > 1:
+        command += ' --placement-type=COMPACT  --max-pods-per-node 15'
+        command += f' --tpu-topology={system.topology}'
+        command += f' {args.custom_tpu_nodepool_arguments}'
     elif system.accelerator_type == AcceleratorType['GPU']:
       subnet_prefix = f'{args.cluster}-{zone_to_region(args.zone)}'
       if capacity_type == CapacityType.FLEX_START:
@@ -319,6 +325,8 @@ def run_gke_node_pool_create_command(
     if args.enable_workload_identity or args.enable_gcsfuse_csi_driver:
       command += ' --workload-metadata=GKE_METADATA'
 
+    command += f' {args.custom_nodepool_arguments}'
+
     task = f'NodepoolCreate-{node_pool_name}'
     create_commands.append(command)
     create_task_names.append(task)

src/xpk/core/system_characteristics.py

@@ -242,6 +242,115 @@ def get_tpu_system_characteristics_map(
         'h100-mega-80gb-8',
     ),
     # TPU system characteristics
+    **get_tpu_system_characteristics_map(
+        'tpu7x', 2, 'tpu7x', 'tpu7x-standard-1t', ['1x1x1']
+    ),
+    **get_tpu_system_characteristics_map(
+        'tpu7x',
+        2,
+        'tpu7x',
+        'tpu7x-standard-4t',
+        [
+            '12x12x12',
+            '12x12x16',
+            '12x12x20',
+            '12x12x24',
+            '12x12x28',
+            '12x12x36',
+            '12x16x16',
+            '12x16x20',
+            '12x16x24',
+            '12x16x28',
+            '12x20x20',
+            '12x20x24',
+            '12x24x24',
+            '16x16x16',
+            '16x16x20',
+            '16x16x24',
+            '16x16x32',
+            '16x20x28',
+            '16x24x24',
+            '2x2x1',
+            '2x2x2',
+            '2x2x4',
+            '2x4x4',
+            '4x12x116',
+            '4x12x12',
+            '4x12x124',
+            '4x12x20',
+            '4x12x28',
+            '4x12x44',
+            '4x12x52',
+            '4x12x68',
+            '4x12x76',
+            '4x12x92',
+            '4x20x20',
+            '4x20x28',
+            '4x20x44',
+            '4x20x52',
+            '4x20x68',
+            '4x20x76',
+            '4x28x28',
+            '4x28x44',
+            '4x28x52',
+            '4x4x116',
+            '4x4x12',
+            '4x4x124',
+            '4x4x148',
+            '4x4x164',
+            '4x4x172',
+            '4x4x188',
+            '4x4x20',
+            '4x4x212',
+            '4x4x236',
+            '4x4x244',
+            '4x4x28',
+            '4x4x4',
+            '4x4x44',
+            '4x4x52',
+            '4x4x68',
+            '4x4x76',
+            '4x4x8',
+            '4x4x92',
+            '4x8x116',
+            '4x8x12',
+            '4x8x124',
+            '4x8x148',
+            '4x8x164',
+            '4x8x172',
+            '4x8x188',
+            '4x8x20',
+            '4x8x28',
+            '4x8x44',
+            '4x8x52',
+            '4x8x68',
+            '4x8x76',
+            '4x8x8',
+            '4x8x92',
+            '8x12x12',
+            '8x12x16',
+            '8x12x20',
+            '8x12x28',
+            '8x12x44',
+            '8x12x52',
+            '8x16x16',
+            '8x16x20',
+            '8x16x28',
+            '8x16x44',
+            '8x20x20',
+            '8x20x28',
+            '8x8x12',
+            '8x8x16',
+            '8x8x20',
+            '8x8x28',
+            '8x8x44',
+            '8x8x52',
+            '8x8x68',
+            '8x8x76',
+            '8x8x8',
+            '8x8x92',
+        ],
+    ),
     **get_tpu_system_characteristics_map(
         'v6e', 1, 'tpu-v6e-slice', 'ct6e-standard-1t', ['1x1']
     ),