xpk 1.1.1__py3-none-any.whl → 1.2.0__py3-none-any.whl

xpk/commands/cluster.py

@@ -369,7 +369,7 @@ def cluster_create(args) -> None:
 
   get_cluster_credentials(args)
 
-  update_coredns_command_code = update_coredns_if_necessary()
+  update_coredns_command_code = update_coredns_if_necessary(args)
   if update_coredns_command_code != 0:
     xpk_exit(update_coredns_command_code)
 
@@ -927,7 +927,7 @@ def scale_down_deployment(deployment_name: str, namespace: str = 'kube-system'):
   xpk_print(f'{deployment_name} has been scaled down.')
 
 
-def scale_up_coredns(replicas: int = 15, namespace: str = 'kube-system'):
+def scale_up_coredns(replicas: int, namespace: str = 'kube-system'):
   """Scales up the CoreDNS deployment to a specified number of replicas."""
   command_coredns_scale = (
       f'kubectl scale deployment coredns --replicas={replicas} -n {namespace}'
@@ -1008,7 +1008,14 @@ def cleanup_coredns_repo(coredns_repo_full_path: str):
     xpk_print(f'Error deleting directory {coredns_repo_full_path}: {e}')
 
 
-def update_coredns() -> int:
+def _get_coredns_replica_count(args) -> int:
+  # XPK large scale guide recommends 15 coreDNS replicas for clusters with 5000 VMs.
+  # Otherwise, limit the replica count to the desired number of default pool nodes.
+  default_pool_node_count: int = args.default_pool_cpu_num_nodes
+  return min(15, default_pool_node_count)
+
+
+def update_coredns(args) -> int:
   """Updates and deploys CoreDNS within a cluster.
 
   Returns:
@@ -1018,6 +1025,8 @@ def update_coredns() -> int:
   coredns_repo_dir_name = 'deployment'
   coredns_repo_full_path = os.path.join(coredns_repo_dir, coredns_repo_dir_name)
   coredns_k8s_path = os.path.join(coredns_repo_full_path, 'kubernetes')
+  coredns_replica_count = _get_coredns_replica_count(args)
+
   # 1. Install jq
   install_jq()
 
@@ -1034,7 +1043,7 @@ def update_coredns() -> int:
   scale_down_deployment('kube-dns')
 
   # 6. Scale up coredns and verify readiness
-  scale_up_coredns(replicas=15)
+  scale_up_coredns(coredns_replica_count)
   verify_coredns_readiness()
 
   xpk_print('The CoreDNS setup process has been completed.')
@@ -1074,7 +1083,7 @@ def coredns_deployment_exists(namespace: str = 'kube-system') -> bool:
     return False
 
 
-def update_coredns_if_necessary() -> int:
+def update_coredns_if_necessary(args) -> int:
   """Updates and deploys CoreDNS within the cluster if it's not already present.
 
   This function checks for the existence of the CoreDNS deployment.
@@ -1089,7 +1098,7 @@ def update_coredns_if_necessary() -> int:
     return 0
   else:
     xpk_print('CoreDNS deployment not found. Proceeding with CoreDNS setup.')
-    return update_coredns()
+    return update_coredns(args)
 
 
 def create_cluster_if_necessary(

xpk/commands/cluster_test.py

@@ -22,7 +22,7 @@ from unittest.mock import MagicMock, patch
 import pytest
 
 from xpk.core.telemetry import MetricsCollector
-from xpk.commands.cluster import _install_kueue, _validate_cluster_create_args, run_gke_cluster_create_command, cluster_create, _log_cluster_create_telemetry
+from xpk.commands.cluster import _install_kueue, _validate_cluster_create_args, _get_coredns_replica_count, run_gke_cluster_create_command, cluster_create, _log_cluster_create_telemetry
 from xpk.core.capacity import CapacityType
 from xpk.core.system_characteristics import SystemCharacteristics, UserFacingNameToSystemCharacteristics
 from xpk.core.testing.commands_tester import CommandsTester
@@ -787,3 +787,18 @@ def test_validate_cluster_create_args_sets_correct_num_slices(
   _validate_cluster_create_args(args, SUPER_SLICING_SYSTEM)
 
   assert args.num_slices == expected
+
+
+def test_get_coredns_replica_count_lower_limit_is_number_of_nodes():
+  args = construct_args(
+      default_pool_cpu_num_nodes=7,
+  )
+
+  assert _get_coredns_replica_count(args) == 7
+
+
+def test_get_coredns_replica_count_upper_limit_is_15():
+  args = construct_args(
+      default_pool_cpu_num_nodes=20,
+  )
+  assert _get_coredns_replica_count(args) == 15

xpk/core/cluster.py

@@ -21,6 +21,8 @@ from kubernetes import client as k8s_client
 from kubernetes import config
 from kubernetes.client.exceptions import ApiException
 
+from .kubectl_common import PatchResources, patch_controller_manager_resources
+from ..utils.feature_flags import FeatureFlags
 from ..utils.console import xpk_exit, xpk_print
 from .capacity import H200_DEVICE_TYPE
 from .commands import (
@@ -33,6 +35,7 @@ from .gcloud_context import (
     get_cluster_location,
     zone_to_region,
 )
+from .nodepool import recreate_nodes_in_existing_node_pools
 from .resources import get_cluster_system_characteristics
 from .system_characteristics import INSTALLER_NCCL_TCPXO, SystemCharacteristics
 
@@ -72,7 +75,21 @@ def set_jobset_on_cluster(args) -> int:
         ' https://github.com/google/xpk/blob/main/README.md#troubleshooting for'
         ' instructions on how to fix these permissions.'
     )
-  return return_code
+    return return_code
+
+  if FeatureFlags.SUPER_SLICING_ENABLED and args.super_slicing:
+    return patch_controller_manager_resources(
+        name='jobset-controller-manager',
+        namespace='jobset-system',
+        patch_resources=PatchResources(
+            cpu_request=4,
+            cpu_limit=4,
+            memory_request='16Gi',
+            memory_limit='16Gi',
+        ),
+    )
+
+  return 0
 
 
 def set_pathways_job_on_cluster(args) -> int:
@@ -605,6 +622,19 @@ def update_gke_cluster_with_lustre_driver_enabled(args) -> int:
   if return_code != 0:
     xpk_print(f'GKE Cluster Update request returned ERROR {return_code}')
     return 1
+
+  xpk_print(
+      'Recreating existing nodes (if any) to complete the Lustre CSI driver'
+      ' installation.'
+  )
+  return_code = recreate_nodes_in_existing_node_pools(args)
+  if return_code != 0:
+    xpk_print(
+        f'Node recreation failed with ERROR {return_code}. You must recreate'
+        ' the nodes manually in order to access Lustre storage from your'
+        ' workloads.'
+    )
+    return 1
   return 0
 
 

xpk/core/cluster_test.py

@@ -14,10 +14,12 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+from unittest.mock import MagicMock
 import pytest
 from .testing.commands_tester import CommandsTester
-from .cluster import get_cluster_credentials, update_gke_cluster_with_lustre_driver_enabled, update_cluster_with_lustre_driver_if_necessary
+from .cluster import get_cluster_credentials, update_gke_cluster_with_lustre_driver_enabled, update_cluster_with_lustre_driver_if_necessary, set_jobset_on_cluster
 from pytest_mock import MockerFixture
+from ..utils.feature_flags import FeatureFlags
 
 
 @pytest.fixture(autouse=True)
@@ -26,6 +28,9 @@ def commands_tester(mocker: MockerFixture) -> CommandsTester:
       mocker=mocker,
       run_command_for_value_path="xpk.core.cluster.run_command_for_value",
       run_command_with_updates_path="xpk.core.cluster.run_command_with_updates",
+      run_command_with_updates_retry_path=(
+          "xpk.core.cluster.run_command_with_updates_retry"
+      ),
   )
 
 
@@ -38,7 +43,17 @@ def mock_location(mocker: MockerFixture):
 
 @pytest.fixture(autouse=True)
 def command_args(mocker: MockerFixture):
-  return mocker.Mock(cluster="cluster", project="project", zone="zone")
+  return mocker.Mock(
+      cluster="cluster", project="project", zone="zone", super_slicing=False
+  )
+
+
+@pytest.fixture(autouse=True)
+def mock_patch_controller_manager_resources(mocker: MockerFixture) -> MagicMock:
+  return mocker.patch(
+      "xpk.core.cluster.patch_controller_manager_resources",
+      return_value=0,
+  )
 
 
 def test_get_cluster_credentials_returns_1_when_retrieval_commands_fail(
@@ -166,11 +181,14 @@ def test_update_cluster_with_lustre_driver_if_necessary_with_legacy_port_runs_co
 
 
 def test_update_gke_cluster_with_lustre_driver_enabled_default_port(
-    commands_tester: CommandsTester, command_args
+    commands_tester: CommandsTester, command_args, mocker: MockerFixture
 ):
   commands_tester.set_result_for_command(
       (0, ""), "gcloud container clusters update"
   )
+  mocker.patch(
+      "xpk.core.cluster.recreate_nodes_in_existing_node_pools", return_value=0
+  )
   command_args.enable_legacy_lustre_port = None
   update_gke_cluster_with_lustre_driver_enabled(command_args)
 
@@ -181,12 +199,30 @@ def test_update_gke_cluster_with_lustre_driver_enabled_default_port(
   ]
 
 
+def test_update_gke_cluster_with_lustre_driver_enabled_fails_if_node_recreation_failed(
+    commands_tester: CommandsTester, command_args, mocker: MockerFixture
+):
+  commands_tester.set_result_for_command(
+      (0, ""), "gcloud container clusters update"
+  )
+  mocker.patch(
+      "xpk.core.cluster.recreate_nodes_in_existing_node_pools", return_value=123
+  )
+  command_args.enable_legacy_lustre_port = None
+  return_code = update_gke_cluster_with_lustre_driver_enabled(command_args)
+
+  assert return_code != 0
+
+
 def test_update_gke_cluster_with_lustre_driver_enabled_legacy_port(
-    commands_tester: CommandsTester, command_args
+    commands_tester: CommandsTester, command_args, mocker: MockerFixture
 ):
   commands_tester.set_result_for_command(
       (0, ""), "gcloud container clusters update"
   )
+  mocker.patch(
+      "xpk.core.cluster.recreate_nodes_in_existing_node_pools", return_value=0
+  )
   command_args.enable_legacy_lustre_port = True
   update_gke_cluster_with_lustre_driver_enabled(command_args)
 
@@ -195,3 +231,24 @@ def test_update_gke_cluster_with_lustre_driver_enabled_legacy_port(
       "gcloud container clusters update cluster --project=project"
       " --location=us-central1 --quiet --enable-legacy-lustre-port"
   ]
+
+
+def test_set_jobset_on_cluster_not_setting_resources_by_default(
+    mock_patch_controller_manager_resources: MagicMock, command_args
+):
+  result = set_jobset_on_cluster(command_args)
+
+  assert result == 0
+  mock_patch_controller_manager_resources.assert_not_called()
+
+
+def test_set_jobset_on_cluster_super_slicing_resources(
+    mock_patch_controller_manager_resources: MagicMock, command_args
+):
+  FeatureFlags.SUPER_SLICING_ENABLED = True
+  command_args.super_slicing = True
+
+  result = set_jobset_on_cluster(command_args)
+
+  assert result == 0
+  mock_patch_controller_manager_resources.assert_called()

xpk/core/docker_container.py

@@ -181,7 +181,9 @@ def get_main_container(
             tpu_stacktrace_terminate_command=tpu_stacktrace_terminate_command,
             gpu_workload_terminate_command=gpu_workload_terminate_command,
             xpk_internal_commands=xpk_internal_commands,
-            resources=get_main_container_resources(args, system, resource_type),
+            resources=get_main_container_resources(
+                args, system, resource_type, parallel_containers
+            ),
             volume_mounts=volume_mounts,
         )
     )

xpk/core/docker_resources.py

@@ -23,7 +23,10 @@ from ..utils.execution_context import is_dry_run
 
 
 def get_main_container_resources(
-    args, system: SystemCharacteristics, resource_type
+    args,
+    system: SystemCharacteristics,
+    resource_type: str,
+    parallel_containers: int,
 ) -> str:
   """Resources for the main container.
   Args:
@@ -53,10 +56,7 @@ def get_main_container_resources(
     offset_vCPUs = int(system.chips_per_vm) * 0.95
     return f'{resource_type}: {offset_vCPUs}'
 
-  return (
-      f'{resource_type}:'
-      f' {int(system.chips_per_vm / system.parallel_containers)}'
-  )
+  return f'{resource_type}: {int(system.chips_per_vm / parallel_containers)}'
 
 
 def get_env_container(args, system: SystemCharacteristics) -> str:

xpk/core/kubectl_common.py

@@ -0,0 +1,77 @@
+"""
+Copyright 2026 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 collections import defaultdict
+from dataclasses import dataclass
+import json
+from typing import Any
+
+from .commands import run_command_with_updates_retry
+
+
+@dataclass(frozen=True)
+class PatchResources:
+  cpu_request: int | None = None
+  cpu_limit: int | None = None
+  memory_request: str | None = None
+  memory_limit: str | None = None
+
+
+_EMPTY_PATCH_RESOURCES = PatchResources()
+
+
+def patch_controller_manager_resources(
+    name: str,
+    namespace: str,
+    replicas: int | None = None,
+    patch_resources: PatchResources | None = None,
+) -> int:
+  if replicas is None and patch_resources is None:
+    return 0
+
+  patch: dict[str, Any] = {"spec": {}}
+
+  if replicas is not None:
+    patch["spec"]["replicas"] = replicas
+
+  if patch_resources and patch_resources != _EMPTY_PATCH_RESOURCES:
+    resources: dict[str, dict[str, str]] = defaultdict(dict)
+    if patch_resources.cpu_request is not None:
+      resources["requests"]["cpu"] = str(patch_resources.cpu_request)
+    if patch_resources.cpu_limit is not None:
+      resources["limits"]["cpu"] = str(patch_resources.cpu_limit)
+    if patch_resources.memory_request is not None:
+      resources["requests"]["memory"] = patch_resources.memory_request
+    if patch_resources.memory_limit is not None:
+      resources["limits"]["memory"] = patch_resources.memory_limit
+    patch["spec"]["template"] = {
+        "spec": {
+            "containers": [{
+                "name": "manager",
+                "resources": resources,
+            }]
+        }
+    }
+
+  patch_str = json.dumps(patch)
+  patch_command = (
+      f"kubectl patch deployment {name} -n {namespace}"
+      f" --type='strategic' --patch='{patch_str}'"
+  )
+  return run_command_with_updates_retry(
+      patch_command,
+      "Updating Controller Manager resources",
+  )

xpk/core/kubectl_common_test.py

@@ -0,0 +1,174 @@
+"""
+Copyright 2026 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.
+"""
+
+import json
+import pytest
+from xpk.core.kubectl_common import PatchResources, patch_controller_manager_resources
+from xpk.core.testing.commands_tester import CommandsTester
+
+
+@pytest.fixture
+def commands_tester(mocker):
+  return CommandsTester(
+      mocker,
+      run_command_with_updates_retry_path=(
+          "xpk.core.kubectl_common.run_command_with_updates_retry"
+      ),
+  )
+
+
+def test_patch_controller_manager_resources_full(
+    commands_tester: CommandsTester,
+):
+  result = patch_controller_manager_resources(
+      name="name",
+      namespace="namespace",
+      replicas=7,
+      patch_resources=PatchResources(
+          cpu_request=1,
+          cpu_limit=2,
+          memory_request="10Gi",
+          memory_limit="20Gi",
+      ),
+  )
+
+  assert result == 0
+  commands_tester.assert_command_run(
+      "kubectl patch deployment", "name", "-n namespace"
+  )
+  expected_patch_dict = {
+      "spec": {
+          "replicas": 7,
+          "template": {
+              "spec": {
+                  "containers": [{
+                      "name": "manager",
+                      "resources": {
+                          "requests": {"cpu": "1", "memory": "10Gi"},
+                          "limits": {"cpu": "2", "memory": "20Gi"},
+                      },
+                  }]
+              }
+          },
+      }
+  }
+  commands_tester.assert_command_run(
+      "kubectl patch", json.dumps(expected_patch_dict)
+  )
+
+
+def test_patch_controller_manager_resources_only_replicas(
+    commands_tester: CommandsTester,
+):
+  result = patch_controller_manager_resources(
+      name="name",
+      namespace="namespace",
+      replicas=7,
+  )
+
+  assert result == 0
+  expected_patch_dict = {
+      "spec": {
+          "replicas": 7,
+      }
+  }
+  commands_tester.assert_command_run(
+      "kubectl patch", json.dumps(expected_patch_dict)
+  )
+
+
+def test_patch_controller_manager_resources_only_requests(
+    commands_tester: CommandsTester,
+):
+  result = patch_controller_manager_resources(
+      name="name",
+      namespace="namespace",
+      patch_resources=PatchResources(
+          cpu_request=1,
+          memory_request="10Gi",
+      ),
+  )
+
+  assert result == 0
+  commands_tester.assert_command_run(
+      "kubectl patch deployment", "name", "-n namespace"
+  )
+  expected_patch_dict = {
+      "spec": {
+          "template": {
+              "spec": {
+                  "containers": [{
+                      "name": "manager",
+                      "resources": {
+                          "requests": {"cpu": "1", "memory": "10Gi"},
+                      },
+                  }]
+              }
+          },
+      }
+  }
+  commands_tester.assert_command_run(
+      "kubectl patch", json.dumps(expected_patch_dict)
+  )
+
+
+def test_patch_controller_manager_resources_only_limits(
+    commands_tester: CommandsTester,
+):
+  result = patch_controller_manager_resources(
+      name="name",
+      namespace="namespace",
+      patch_resources=PatchResources(
+          cpu_limit=2,
+          memory_limit="20Gi",
+      ),
+  )
+
+  assert result == 0
+  commands_tester.assert_command_run(
+      "kubectl patch deployment", "name", "-n namespace"
+  )
+  expected_patch_dict = {
+      "spec": {
+          "template": {
+              "spec": {
+                  "containers": [{
+                      "name": "manager",
+                      "resources": {
+                          "limits": {"cpu": "2", "memory": "20Gi"},
+                      },
+                  }]
+              }
+          },
+      }
+  }
+  commands_tester.assert_command_run(
+      "kubectl patch", json.dumps(expected_patch_dict)
+  )
+
+
+def test_patch_controller_manager_resources_propagates_error(
+    commands_tester: CommandsTester,
+):
+  commands_tester.set_result_for_command((123, "kubectl patch"))
+
+  result = patch_controller_manager_resources(
+      name="name",
+      namespace="namespace",
+      replicas=7,
+  )
+
+  assert result == 123

xpk/core/kueue_manager.py

@@ -20,6 +20,7 @@ from typing import Optional, List, Dict, Any
 import json
 from jinja2 import Environment, FileSystemLoader
 
+from .kubectl_common import PatchResources, patch_controller_manager_resources
 from ..utils.topology import get_slice_topology_level, get_topology_product, is_topology_contained
 from ..utils.kueue import is_queued_cluster
 from kubernetes.utils import parse_quantity
@@ -304,7 +305,9 @@ class KueueManager:
     if return_code != 0:
       return return_code
 
-    return self.__update_kueue_resources_if_necessary()
+    return self.__update_kueue_resources_if_necessary(
+        configure_super_slicing=kueue_config.configure_super_slicing
+    )
 
   def __build_template_context(
       self,
@@ -452,8 +455,23 @@ class KueueManager:
     command = f"kubectl apply -f {tmp_file}"
     return run_command_with_updates(command, task)
 
-  def __update_kueue_resources_if_necessary(self) -> int:
+  def __update_kueue_resources_if_necessary(
+      self, configure_super_slicing: bool
+  ) -> int:
     """Patch memory size limit if necessary."""
+    if configure_super_slicing:
+      return patch_controller_manager_resources(
+          name="kueue-controller-manager",
+          namespace="kueue-system",
+          replicas=3,
+          patch_resources=PatchResources(
+              cpu_request=16,
+              cpu_limit=16,
+              memory_request="64Gi",
+              memory_limit="64Gi",
+          ),
+      )
+
     # Get total number of nodes
     cmd_total_node_num = "kubectl get node --no-headers | wc -l"
     return_code, out = run_command_for_value(
@@ -465,31 +483,13 @@ class KueueManager:
     new_memory_limit = (
         f"{max(math.ceil(int(out) * MEMORY_SIZE_PER_VM), MIN_MEMORY_LIMIT_SIZE)}Mi"
     )
-    patch = {
-        "spec": {
-            "template": {
-                "spec": {
-                    "containers": [{
-                        "name": "manager",
-                        "resources": {"limits": {"memory": new_memory_limit}},
-                    }]
-                }
-            }
-        }
-    }
-    patch_str = json.dumps(patch)
-    patch_command = (
-        "kubectl patch deployment kueue-controller-manager -n kueue-system"
-        f" --type='strategic' --patch='{patch_str}'"
-    )
-    task = "Updating Kueue Controller Manager resources"
-    return_code = run_command_with_updates_retry(
-        patch_command,
-        task,
+    return patch_controller_manager_resources(
+        name="kueue-controller-manager",
+        namespace="kueue-system",
+        patch_resources=PatchResources(
+            memory_limit=new_memory_limit,
+        ),
     )
-    if return_code != 0:
-      xpk_print(f"{task} returned ERROR {return_code}")
-    return return_code
 
   def __autocorrect_resource_limits(
       self, kueue_config: KueueConfig

xpk/core/kueue_manager_test.py

@@ -21,6 +21,7 @@ from pytest_mock import MockerFixture
 import yaml
 from unittest.mock import MagicMock, patch
 
+from xpk.core.kubectl_common import PatchResources
 from xpk.core.kueue_manager import KueueConfig, KueueManager, has_sub_slicing_enabled, has_super_slicing_enabled
 from xpk.core.system_characteristics import GpuConfig, DockerPlatform, AcceleratorType, SystemCharacteristics, UserFacingNameToSystemCharacteristics
 from xpk.core.testing.commands_tester import CommandsTester
@@ -86,6 +87,14 @@ def mock_commands(mocker: MockerFixture) -> CommandsTester:
   )
 
 
+@pytest.fixture(autouse=True)
+def mock_patch_controller_manager_resources(mocker: MockerFixture) -> MagicMock:
+  return mocker.patch(
+      "xpk.core.kueue_manager.patch_controller_manager_resources",
+      return_value=0,
+  )
+
+
 @pytest.fixture(autouse=True)
 @patch("jinja2.Environment", return_value=MagicMock())
 def kueue_manager(mock_env: MagicMock) -> KueueManager:
@@ -239,7 +248,9 @@ def test_installation_without_tolerations(
 
 
 def test_resource_update_for_small_cluster(
-    mock_commands: CommandsTester, kueue_manager: KueueManager
+    mock_commands: CommandsTester,
+    kueue_manager: KueueManager,
+    mock_patch_controller_manager_resources: MagicMock,
 ):
   """Test resource update logic for a small cluster."""
   set_installed_kueue_version(mock_commands, None)
@@ -248,17 +259,21 @@ def test_resource_update_for_small_cluster(
   result = kueue_manager.install_or_upgrade(KUEUE_CONFIG)
 
   assert result == 0
+
   # 100 * 1.2 = 120, which is less than 4096. So it should be 4096.
-  mock_commands.assert_command_run(
-      "kubectl patch deployment kueue-controller-manager -n kueue-system"
-      ' --type=\'strategic\' --patch=\'{"spec": {"template": {"spec":'
-      ' {"containers": [{"name": "manager", "resources": {"limits":'
-      ' {"memory": "4096Mi"}}}]}}}}\'',
+  mock_patch_controller_manager_resources.assert_called_with(
+      name="kueue-controller-manager",
+      namespace="kueue-system",
+      patch_resources=PatchResources(
+          memory_limit="4096Mi",
+      ),
   )
 
 
 def test_resource_update_for_large_cluster(
-    mock_commands: CommandsTester, kueue_manager: KueueManager
+    mock_commands: CommandsTester,
+    kueue_manager: KueueManager,
+    mock_patch_controller_manager_resources: MagicMock,
 ):
   """Test resource update logic for a large cluster."""
   set_installed_kueue_version(mock_commands, None)
@@ -268,11 +283,36 @@ def test_resource_update_for_large_cluster(
 
   assert result == 0
   # 5000 * 1.2 = 6000, which is > 4096.
-  mock_commands.assert_command_run(
-      "kubectl patch deployment kueue-controller-manager -n kueue-system"
-      ' --type=\'strategic\' --patch=\'{"spec": {"template": {"spec":'
-      ' {"containers": [{"name": "manager", "resources": {"limits":'
-      ' {"memory": "6000Mi"}}}]}}}}\'',
+  mock_patch_controller_manager_resources.assert_called_with(
+      name="kueue-controller-manager",
+      namespace="kueue-system",
+      patch_resources=PatchResources(
+          memory_limit="6000Mi",
+      ),
+  )
+
+
+def test_resource_update_for_super_slicing_cluster(
+    mock_commands: CommandsTester,
+    kueue_manager: KueueManager,
+    mock_patch_controller_manager_resources: MagicMock,
+):
+  set_installed_kueue_version(mock_commands, None)
+  kueue_config = dataclasses.replace(KUEUE_CONFIG, configure_super_slicing=True)
+
+  result = kueue_manager.install_or_upgrade(kueue_config)
+
+  assert result == 0
+  mock_patch_controller_manager_resources.assert_called_with(
+      name="kueue-controller-manager",
+      namespace="kueue-system",
+      replicas=3,
+      patch_resources=PatchResources(
+          cpu_request=16,
+          cpu_limit=16,
+          memory_request="64Gi",
+          memory_limit="64Gi",
+      ),
   )
 
 

xpk/core/nodepool.py

@@ -691,3 +691,37 @@ def _validate_reservation_count(
     )
     return 1
   return 0
+
+
+def recreate_nodes_in_existing_node_pools(args) -> int:
+  """Triggers a manual upgrade of nodepools to the same version to force recreation
+  of nodes.
+  """
+
+  existing_node_pool_names, return_code = get_all_nodepools_programmatic(args)
+  if return_code > 0:
+    xpk_print('Listing all node pools failed!')
+    return return_code
+
+  commands = [
+      (
+          f'gcloud container clusters upgrade {args.cluster}'
+          f' --project={args.project}'
+          f' --node-pool={node_pool_name}'
+          f' --location={get_cluster_location(args.project, args.cluster, args.zone)}'
+          ' --quiet'
+      )
+      for node_pool_name in existing_node_pool_names
+  ]
+  task_names = [
+      f'NodesRecreate-{node_pool_name}'
+      for node_pool_name in existing_node_pool_names
+  ]
+  for i, command in enumerate(commands):
+    xpk_print(f'To complete {task_names[i]} we are executing {command}')
+  maybe_failure = run_commands(
+      commands,
+      'Recreate nodes in nodepools',
+      task_names,
+  )
+  return maybe_failure.return_code if maybe_failure is not None else 0

xpk/core/nodepool_test.py

@@ -20,6 +20,7 @@ from xpk.core.nodepool import (
     ensure_resource_policy_exists,
     get_desired_node_pool_names,
     run_gke_node_pool_create_command,
+    recreate_nodes_in_existing_node_pools,
     _validate_reservation_count,
 )
 from xpk.core.system_characteristics import AcceleratorType, SystemCharacteristics, DockerPlatform, GpuConfig
@@ -568,3 +569,106 @@ def test_run_gke_node_pool_create_command_partial_reservations(
   commands_tester.assert_command_run(
       "gcloud", "node-pools create", "test-cluster-np-2", "--reservation=res2"
   )
+
+
+def test_recreate_nodes_in_existing_node_pools_upgrades_existing_nodepools(
+    mocker,
+    commands_tester: CommandsTester,
+):
+  mocker.patch(
+      "xpk.core.nodepool.get_all_nodepools_programmatic",
+      return_value=(["nodepool1", "nodepool2"], 0),
+  )
+  mocker.patch(
+      "xpk.core.nodepool.get_cluster_location", return_value="us-central1"
+  )
+  args = mocker.Mock(
+      cluster="test-cluster",
+      project="test-project",
+      zone="us-central1-a",
+  )
+
+  commands_tester.set_result_for_command(
+      (0, ""), "gcloud container clusters upgrade"
+  )
+
+  result = recreate_nodes_in_existing_node_pools(args)
+
+  assert result == 0
+  commands_tester.assert_command_run(
+      "gcloud",
+      "container clusters upgrade test-cluster",
+      "--project=test-project",
+      "--node-pool=nodepool1",
+      "--location=us-central1",
+      "--quiet",
+  )
+  commands_tester.assert_command_run(
+      "gcloud",
+      "container clusters upgrade test-cluster",
+      "--project=test-project",
+      "--node-pool=nodepool2",
+      "--location=us-central1",
+      "--quiet",
+  )
+
+
+def test_recreate_nodes_in_existing_node_pools_returns_error_code_if_listing_fails(
+    mocker,
+    commands_tester: CommandsTester,
+):
+  mocker.patch(
+      "xpk.core.nodepool.get_all_nodepools_programmatic",
+      return_value=([], 123),
+  )
+  args = mocker.Mock(
+      cluster="test-cluster",
+      project="test-project",
+      zone="us-central1-a",
+  )
+
+  result = recreate_nodes_in_existing_node_pools(args)
+
+  assert result == 123
+
+
+def test_recreate_nodes_in_existing_node_pools_returns_error_code_if_upgrade_fails(
+    mocker,
+    commands_tester: CommandsTester,
+):
+  mocker.patch(
+      "xpk.core.nodepool.get_all_nodepools_programmatic",
+      return_value=(["nodepool1", "nodepool2"], 0),
+  )
+  mocker.patch(
+      "xpk.core.nodepool.get_cluster_location", return_value="us-central1"
+  )
+  args = mocker.Mock(
+      cluster="test-cluster",
+      project="test-project",
+      zone="us-central1-a",
+  )
+
+  commands_tester.set_result_for_command(
+      (123, ""), "gcloud container clusters upgrade"
+  )
+
+  result = recreate_nodes_in_existing_node_pools(args)
+
+  assert result == 123
+  commands_tester.assert_command_run(
+      "gcloud",
+      "container clusters upgrade test-cluster",
+      "--project=test-project",
+      "--node-pool=nodepool1",
+      "--location=us-central1",
+      "--quiet",
+  )
+  commands_tester.assert_command_not_run(
+      "gcloud",
+      "container clusters upgrade test-cluster",
+      "--project=test-project",
+      "--node-pool=nodepool2",
+      "--location=us-central1",
+      "--quiet",
+  )

{xpk-1.1.1.dist-info → xpk-1.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xpk
-Version: 1.1.1
+Version: 1.2.0
 Summary: xpk helps Cloud developers to orchestrate training jobs on accelerators on GKE.
 Author-email: XPK team <xpk-code-reviewers@google.com>
 License: Apache-2.0

{xpk-1.1.1.dist-info → xpk-1.2.0.dist-info}/RECORD

@@ -13,10 +13,10 @@ xpk/blueprints/a4/config-map.yaml.tftpl,sha256=o6LeGIYUfFGyj3vj-8ztV5ildQ46QZVl7
 xpk/blueprints/a4/nccl-rdma-installer-a4.yaml,sha256=if3WOmNLVGTJIJHU76EWC1FyiIXDTRIXcwo4OsBxarQ,2113
 xpk/blueprints/a4/storage_crd.yaml,sha256=r4WFXnSJJ25EUF-t4Ljfbl-cJoSaiFiZkP8451eTub4,1260
 xpk/commands/__init__.py,sha256=YPwWBbgLAu7L-YlTVGB2r8ZV4TzypURMRBcehSHHlLY,561
-xpk/commands/cluster.py,sha256=9LGVlrXg5C2w7wRKYTWe6lh3wSl7ZFpaNXSlpBfd0tc,45311
+xpk/commands/cluster.py,sha256=9Eqj0d77_4FZ4ztsYFiIvX514jp-XvOTkZIVMdvitOI,45715
 xpk/commands/cluster_gcluster.py,sha256=H6_pphIrIzDCHZg-ZH3o-xA2AgVQKSjE2HKbhIW6-Xo,13727
 xpk/commands/cluster_gcluster_test.py,sha256=zdxz5gAMu3HRVNsj7F-VYRf4TYSPMjuOG7DolQN2Pb4,6263
-xpk/commands/cluster_test.py,sha256=va1ODWx0Y2fSZIGzIsqEIVN25tIsfYr95ubcyxTMDfA,23494
+xpk/commands/cluster_test.py,sha256=osBhaZTdB2kgQw1aY3W_ybFOqyb0UZncvVLyZNuKRpU,23878
 xpk/commands/common.py,sha256=fsM4Sud3y1RU6a8JHi99l13O4raYvW2oPahCBzvMwh4,3884
 xpk/commands/common_test.py,sha256=BDYFtN-cVfpEpj6akZy4R2KrnP53AIV1Lh1FEImhXx0,6106
 xpk/commands/config.py,sha256=L_zRpQTxMcSh6rxOT8gG263V6YGqzVoz4UxdWywTFdA,850
@@ -32,30 +32,32 @@ xpk/commands/workload_test.py,sha256=m79x6YDYn-36BX0CttTtAMdt_O-WJY40FLTGa6KwKg8
 xpk/core/__init__.py,sha256=YPwWBbgLAu7L-YlTVGB2r8ZV4TzypURMRBcehSHHlLY,561
 xpk/core/capacity.py,sha256=MGiNOwBCwg8Ci-hsssbZYIJ2xXTm6Y5yKTO4J5ozqEk,11053
 xpk/core/capacity_test.py,sha256=04ecANSECL3APmFCjdtkw2Wz6JxWkRZwE_QHa2m1_68,6026
-xpk/core/cluster.py,sha256=RC-91Dk4nx9F-jLmZPP78ALlUgSBq670geaIqnRMZxY,24184
+xpk/core/cluster.py,sha256=JLHFlrcc3Ch3ggLv2EJrOdQOOTn9MB-0IF5fQLIRPD0,25159
 xpk/core/cluster_private.py,sha256=43BEO49MA-p1KfKuchxVZk6wwIlRCUU17j-6MMMeq4I,6868
-xpk/core/cluster_test.py,sha256=VeC1C7kN0OJe6yeoL8GCaFk4uPhijP6CjvQAcE7q9xw,6653
+xpk/core/cluster_test.py,sha256=0ymagvKxyX909NjA5J23lutq7xyHXlIWwvGRFoVE0n8,8492
 xpk/core/commands.py,sha256=at73VJHdZ4rVA8uvW997tNrvnCjP9v6zaw96bU0kd74,10841
 xpk/core/config.py,sha256=U2JDXx-XBuqQpZJf2iUDoww5--E8ejZfgmIxKeGu-gU,4668
 xpk/core/config_test.py,sha256=POSuofK0LFbNNygDAo2fjtKY4NMrRjUFeGcpBh9JOS4,3569
-xpk/core/docker_container.py,sha256=SZ3msIMx-2gBIok3j5N73KO1wVcMfzw1pQpgTdyZOQA,8243
+xpk/core/docker_container.py,sha256=9kJpTEholW_d_GamjcqunCWT4XwrDyZs3fcvcPNCb8Y,8294
 xpk/core/docker_image.py,sha256=9vwqbb6Mc3C5ZEOph03WS-EWI5hxMYGGigqzIMkDTjE,6909
 xpk/core/docker_manager.py,sha256=vGPCWPDB507sxEsXvSD4IM-h5HqQzYLk7WSdCUmSDb4,10568
-xpk/core/docker_resources.py,sha256=7EXV1CvwCVogE5-m6utSE1GXxwf6EpB4QDYeuGXWHmI,12547
+xpk/core/docker_resources.py,sha256=bwHGNh_gOtprVOeoFC8NObgKGD9aDjNc2XBMS6syD2Q,12562
 xpk/core/filestore.py,sha256=mcuUzsAPARbnrBG4fIGsEoN8NmzjaQ6k0tvIwMtjO9k,8068
 xpk/core/gcloud_context.py,sha256=d1wQ76zp7QMdG5BxB3sJz4b4OF5Mc8OzmPd_m0xd-Ys,6810
 xpk/core/gcloud_context_test.py,sha256=M8rp6S1zaEcAI7u4Bt8ukWKzv82HH5h9oYVojBcKgHk,5987
 xpk/core/gcluster_manager.py,sha256=lyv_MvdnkByy9_PEBj_ugAEBwnCbFNiWTSrEFjrMlPc,6236
 xpk/core/gcsfuse.py,sha256=kg5pgxdTjgiqquuGjev9fXzJPb8oiWPTK6wzCddzheQ,2125
 xpk/core/jobset.py,sha256=PJ4Fd8TNNLuYKNOMehoMYRIUEXyc5jsbHctJGqfW_8Y,4037
-xpk/core/kueue_manager.py,sha256=qpz4Df7tfWNKzBFlTbMUfsHnXl15SdI7r_mHlCFRYdc,19998
-xpk/core/kueue_manager_test.py,sha256=iJZFQE-fhQAI8MVXe66zUJpSbU5HHUZmNFnnCPCXNZs,22042
+xpk/core/kubectl_common.py,sha256=mVuHJERdRXVEigU9Fxcmq7SHABkWDUmfYTJFvkHEbLs,2353
+xpk/core/kubectl_common_test.py,sha256=2PTTcr5d8pSifuki2L7uA8-UWOqcIpv6PJQ1hlFruJQ,4469
+xpk/core/kueue_manager.py,sha256=8VVOEvkh6Cif3s9ii0HfUR27HuxEJWUwnt4x4YM6ydc,20031
+xpk/core/kueue_manager_test.py,sha256=Xb-w_eF11doqkpfNPb5EYm9BneEo3dtOxHuH-jtfjAI,23001
 xpk/core/monitoring.py,sha256=__bzTq_DIDAK8yIaN4F3MJh-yjYw5X1OlxmRgYOpf1g,4332
 xpk/core/mtc.py,sha256=pO7p3l-EzLFdTE8MdwWV8i0Zu-7epGql_kPoksVofIU,6259
 xpk/core/nap.py,sha256=gBxXu8Png1-BlAHbxLWZgbSXeLMGVixufkQVMR0fmvk,12963
 xpk/core/network.py,sha256=Oulb7U69lWkpOKxOC1C7ekJDpC51TLwd7XdZA3NQ7E0,10505
-xpk/core/nodepool.py,sha256=FX2ljKvwMsG3fXfn_CDCRwiKH4UAArQeDiFLq3XK9F0,25495
-xpk/core/nodepool_test.py,sha256=9xSFpn-1j9Vd0J8KFzbq8ywS_Ibsbx4CgR1er68mRnw,17542
+xpk/core/nodepool.py,sha256=dgE4l-HETNvnVj6WfCNWZkegzqrLiNg61RotGTMPDd0,26575
+xpk/core/nodepool_test.py,sha256=Zf1NW9Hoj0D517NxSsxcj6p8NxeEbFmRpdIACz9NPdM,20207
 xpk/core/pathways.py,sha256=9w_VrpLLjQSSdNd8HJLWWtIYzA0NpR7t70knRSVLK0w,11574
 xpk/core/pathways_test.py,sha256=UeuSo_g9BNI27to-wflQwc6dJFVSA5-kOK_cjmY5qgU,1809
 xpk/core/ray.py,sha256=JWhc_ToRHpF4_URGnuE_47FMgamaRsA4KVUMpqThWzw,6145
@@ -142,9 +144,9 @@ xpk/utils/validation.py,sha256=rE9LTkXJT7jIesodFb9pONL7ixhLqiQleyoaz7N39Dw,2765
 xpk/utils/validation_test.py,sha256=PEDSMUqZdt_Lx1FSR-LOTXKKtsJ47JH1fxugM0Gfz6Y,1168
 xpk/utils/versions.py,sha256=_Ep68W70a9605XjiaOOpBa9Is9jXlsoOiwL8v5Xt-WA,897
 xpk/utils/yaml.py,sha256=j8xuAJ9yAAwnQi6ozwZ-nMnDyDnc3xWkeBZMtSuP4RU,844
-xpk-1.1.1.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-xpk-1.1.1.dist-info/METADATA,sha256=CqGGSACNDUVw7uPODTokMvMqpBMc1gDbY99AUmvQ68Q,10013
-xpk-1.1.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-xpk-1.1.1.dist-info/entry_points.txt,sha256=mzEtiIesFkT1kmcTUVDA1o3uOhiniX6tIz2wmOlMu1M,38
-xpk-1.1.1.dist-info/top_level.txt,sha256=aDe4N0jicmuWExx_6w0TxWQJaEuPSs9BnLU-3aF1GLo,4
-xpk-1.1.1.dist-info/RECORD,,
+xpk-1.2.0.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+xpk-1.2.0.dist-info/METADATA,sha256=Jkgz_1jDxS-vJnbqKI_Kdd5LRH9mikljvnEFXTk8lWA,10013
+xpk-1.2.0.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+xpk-1.2.0.dist-info/entry_points.txt,sha256=mzEtiIesFkT1kmcTUVDA1o3uOhiniX6tIz2wmOlMu1M,38
+xpk-1.2.0.dist-info/top_level.txt,sha256=aDe4N0jicmuWExx_6w0TxWQJaEuPSs9BnLU-3aF1GLo,4
+xpk-1.2.0.dist-info/RECORD,,

{xpk-1.1.1.dist-info → xpk-1.2.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.9.0)
+Generator: setuptools (80.10.2)
 Root-Is-Purelib: true
 Tag: py3-none-any