|
32 | 32 | create_or_update_cluster_configmap,
|
33 | 33 | )
|
34 | 34 | from .system_characteristics import AcceleratorType
|
| 35 | +from functools import reduce |
| 36 | +from operator import mul |
35 | 37 |
|
36 | 38 | CLOUD_PLATFORM_AUTH_SCOPE_URL = (
|
37 | 39 | '"https://www.googleapis.com/auth/cloud-platform"'
|
@@ -275,20 +277,24 @@ def run_gke_node_pool_create_command(
|
275 | 277 | f' --host-maintenance-interval={args.host_maintenance_interval}'
|
276 | 278 | f' {capacity_args}'
|
277 | 279 | ' --enable-gvnic'
|
278 |
| - f' {args.custom_nodepool_arguments}' |
279 | 280 | )
|
280 | 281 | if system.accelerator_type == AcceleratorType['TPU']:
|
281 | 282 | command += f' --node-version={gke_node_pool_version}'
|
| 283 | + topology_product = reduce( |
| 284 | + mul, (int(x) for x in system.topology.split('x')), 1 |
| 285 | + ) |
282 | 286 | if capacity_type == CapacityType.FLEX_START:
|
283 | 287 | command += ' --num-nodes=0'
|
284 |
| - else: |
| 288 | + elif topology_product > 1: |
285 | 289 | command += f' --num-nodes={system.vms_per_slice}'
|
286 |
| - command += ' --placement-type=COMPACT --max-pods-per-node 15' |
287 | 290 | command += (
|
288 | 291 | f' --scopes=storage-full,gke-default,{CLOUD_PLATFORM_AUTH_SCOPE_URL}'
|
289 | 292 | )
|
290 |
| - command += f' --tpu-topology={system.topology}' |
291 |
| - command += f' {args.custom_tpu_nodepool_arguments}' |
| 293 | + |
| 294 | + if topology_product > 1: |
| 295 | + command += ' --placement-type=COMPACT --max-pods-per-node 15' |
| 296 | + command += f' --tpu-topology={system.topology}' |
| 297 | + command += f' {args.custom_tpu_nodepool_arguments}' |
292 | 298 | elif system.accelerator_type == AcceleratorType['GPU']:
|
293 | 299 | subnet_prefix = f'{args.cluster}-{zone_to_region(args.zone)}'
|
294 | 300 | if capacity_type == CapacityType.FLEX_START:
|
@@ -319,6 +325,8 @@ def run_gke_node_pool_create_command(
|
319 | 325 | if args.enable_workload_identity or args.enable_gcsfuse_csi_driver:
|
320 | 326 | command += ' --workload-metadata=GKE_METADATA'
|
321 | 327 |
|
| 328 | + command += args.custom_nodepool_arguments |
| 329 | + |
322 | 330 | task = f'NodepoolCreate-{node_pool_name}'
|
323 | 331 | create_commands.append(command)
|
324 | 332 | create_task_names.append(task)
|
|
0 commit comments