Merge branch 'develop' into chzheng/docker_image_flag

ycchenzheng · web-flow · commit 88f97a97e88b · 2025-09-03T17:06:20.000-04:00
diff --git a/src/xpk/commands/common.py b/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 (
diff --git a/src/xpk/commands/workload.py b/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
diff --git a/src/xpk/core/blueprint/blueprint_generator.py b/src/xpk/core/blueprint/blueprint_generator.py
@@ -52,7 +52,7 @@
 }
 
 cluster_toolkit_url = "github.com/GoogleCloudPlatform/cluster-toolkit"
-cluster_toolkit_version = "v1.57.1"
+cluster_toolkit_version = "v1.62.2"
 
 
 class BlueprintGeneratorOutput:
diff --git a/src/xpk/core/config.py b/src/xpk/core/config.py
@@ -22,7 +22,7 @@
 from ..utils.console import xpk_print
 
 # This is the version for XPK PyPI package
-__version__ = 'v0.10.1'
+__version__ = 'v0.11.0'
 XPK_CURRENT_VERSION = __version__
 XPK_CONFIG_FILE = os.path.expanduser('~/.config/xpk/config.yaml')
 
diff --git a/src/xpk/core/docker_manager.py b/src/xpk/core/docker_manager.py
@@ -30,7 +30,7 @@
 DockerRunCommandExitCode = 135
 dockerBuildErrorCode = 134
 ctk_dockerfile_path = "Dockerfile"
-ctk_build_ref = "v1.57.1"
+ctk_build_ref = "v1.62.2"
 ctk_docker_image = "xpk-ctk"
 ctk_container_name = "xpk-ctk-container"
 gcloud_cfg_mount_path = "/root/.config/gcloud"
diff --git a/src/xpk/core/kueue.py b/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}
 ---
diff --git a/src/xpk/core/nodepool.py b/src/xpk/core/nodepool.py
@@ -14,6 +14,7 @@
 limitations under the License.
 """
 
+from typing import List
 from ..utils.console import get_user_input, xpk_print
 from .capacity import (
     AUTOPROVISIONING_CONFIG_VALUE,
@@ -90,20 +91,26 @@ def run_gke_node_pool_create_command(
     xpk_print('Parsing capacity arguments failed!')
     return return_code
 
-  if system.accelerator_type == AcceleratorType['GPU']:
-    xpk_print(
-        f'Creating 1 node pool with {args.num_nodes} nodes of'
-        f' {system.device_type}\nUnderlyingly, we assume that means: {system}'
-    )
-    desired_node_pool_names = [f'{args.cluster}-np-0']
-  else:
-    xpk_print(
-        f'Creating {args.num_slices} node pool or pools of'
-        f' {system.device_type}\nUnderlyingly, we assume that means: {system}'
-    )
-    desired_node_pool_names = [
-        f'{args.cluster}-np-{slice_num}' for slice_num in range(args.num_slices)
-    ]
+  desired_node_pool_count = (
+      1
+      if system.accelerator_type == AcceleratorType['GPU']
+      else args.num_slices
+  )
+  message = (
+      (
+          f'Creating 1 node pool with {args.num_nodes} nodes of'
+          f' {system.device_type}\nUnderlyingly, we assume that means: {system}'
+      )
+      if system.accelerator_type == AcceleratorType['GPU']
+      else (
+          f'Creating {args.num_slices} node pool or pools of'
+          f' {system.device_type}\nUnderlyingly, we assume that means: {system}'
+      )
+  )
+  xpk_print(message)
+  desired_node_pool_names = get_desired_node_pool_names(
+      existing_node_pool_names, args.cluster, desired_node_pool_count
+  )
 
   node_pools_to_remain = []
   delete_commands = []
@@ -602,3 +609,21 @@ def get_nodepool_workload_metadata_mode(
     return 1, None
 
   return 0, nodepool_WI_mode.strip()
+
+
+def get_desired_node_pool_names(
+    existing_node_pool_names: List[str],
+    cluster_name: str,
+    desired_node_pool_count: int,
+) -> List[str]:
+  cluster_node_pools = [
+      np
+      for np in existing_node_pool_names
+      if np.startswith(f'{cluster_name}-np-')
+  ]
+  result = set(cluster_node_pools[:desired_node_pool_count])
+  i = 0
+  while len(result) < desired_node_pool_count:
+    result.add(f'{cluster_name}-np-{i}')
+    i += 1
+  return list(result)
diff --git a/src/xpk/core/tests/data/a3_mega.yaml b/src/xpk/core/tests/data/a3_mega.yaml
@@ -16,7 +16,7 @@
 !Blueprint
 blueprint_name: xpk-gke-a3-megagpu
 toolkit_modules_url: github.com/GoogleCloudPlatform/cluster-toolkit
-toolkit_modules_version: v1.57.1
+toolkit_modules_version: v1.62.2
 
 vars:
   project_id: "foo"
diff --git a/src/xpk/core/tests/data/a3_mega_spot.yaml b/src/xpk/core/tests/data/a3_mega_spot.yaml
@@ -16,7 +16,7 @@
 !Blueprint
 blueprint_name: xpk-gke-a3-megagpu
 toolkit_modules_url: github.com/GoogleCloudPlatform/cluster-toolkit
-toolkit_modules_version: v1.57.1
+toolkit_modules_version: v1.62.2
 
 vars:
   project_id: "foo"
diff --git a/src/xpk/core/tests/data/a3_ultra.yaml b/src/xpk/core/tests/data/a3_ultra.yaml
@@ -14,7 +14,7 @@
 !Blueprint
 blueprint_name: xpk-gke-a3-ultra
 toolkit_modules_url: github.com/GoogleCloudPlatform/cluster-toolkit
-toolkit_modules_version: v1.57.1
+toolkit_modules_version: v1.62.2
 
 vars:
 
diff --git a/src/xpk/core/tests/data/a4.yaml b/src/xpk/core/tests/data/a4.yaml
@@ -14,7 +14,7 @@
 !Blueprint
 blueprint_name: xpk-gke-a4
 toolkit_modules_url: github.com/GoogleCloudPlatform/cluster-toolkit
-toolkit_modules_version: v1.57.1
+toolkit_modules_version: v1.62.2
 
 vars:
 
diff --git a/src/xpk/core/tests/unit/test_nodepool.py b/src/xpk/core/tests/unit/test_nodepool.py
@@ -0,0 +1,82 @@
+"""
+Copyright 2025 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from xpk.core.nodepool import get_desired_node_pool_names
+
+CLUSTER_NAME = "running-cucumber"
+
+
+def node_pool_name(number: int) -> str:
+  return f"{CLUSTER_NAME}-np-{number}"
+
+
+def test_compute_desired_node_pool_names_with_desired_larger_than_existing():
+  result = get_desired_node_pool_names(
+      existing_node_pool_names=[node_pool_name(0)],
+      cluster_name=CLUSTER_NAME,
+      desired_node_pool_count=2,
+  )
+
+  expected_result = [node_pool_name(0), node_pool_name(1)]
+  assert set(result) == set(expected_result)
+
+
+def test_compute_desired_node_pool_names_with_desired_smaller_than_existing():
+  result = get_desired_node_pool_names(
+      existing_node_pool_names=[node_pool_name(0), node_pool_name(1)],
+      cluster_name=CLUSTER_NAME,
+      desired_node_pool_count=1,
+  )
+
+  expected_result = [node_pool_name(0)]
+  assert set(result) == set(expected_result)
+
+
+def test_compute_desired_node_pool_names_with_consecutive_numbers_missing():
+  result = get_desired_node_pool_names(
+      existing_node_pool_names=[node_pool_name(0), node_pool_name(3)],
+      cluster_name=CLUSTER_NAME,
+      desired_node_pool_count=3,
+  )
+
+  expected_result = [node_pool_name(0), node_pool_name(1), node_pool_name(3)]
+  assert set(result) == set(expected_result)
+
+
+def test_compute_desired_node_pool_names_with_consecutive_numbers_missing_and_desired_equal_to_existing():
+  result = get_desired_node_pool_names(
+      existing_node_pool_names=[node_pool_name(0), node_pool_name(3)],
+      cluster_name=CLUSTER_NAME,
+      desired_node_pool_count=2,
+  )
+
+  expected_result = [node_pool_name(0), node_pool_name(3)]
+  assert set(result) == set(expected_result)
+
+
+def test_compute_desired_node_pool_names_with_unknown_node_pools():
+  result = get_desired_node_pool_names(
+      existing_node_pool_names=[
+          "unknown-node-pool",
+          node_pool_name(0),
+          node_pool_name(3),
+      ],
+      cluster_name=CLUSTER_NAME,
+      desired_node_pool_count=2,
+  )
+
+  expected_result = [node_pool_name(0), node_pool_name(3)]
+  assert set(result) == set(expected_result)

Original file line number	Diff line number	Diff line change
`@@ -52,7 +52,7 @@`
`52`	`52`	`}`
`53`	`53`
`54`	`54`	`cluster_toolkit_url = "github.com/GoogleCloudPlatform/cluster-toolkit"`
`55`		`-cluster_toolkit_version = "v1.57.1"`
	`55`	`+cluster_toolkit_version = "v1.62.2"`
`56`	`56`
`57`	`57`
`58`	`58`	`class BlueprintGeneratorOutput:`