xpk 0.14.3__py3-none-any.whl → 0.15.0__py3-none-any.whl

xpk/core/kueue_manager_test.py

@@ -22,7 +22,7 @@ import yaml
 from unittest.mock import MagicMock, patch
 
 from xpk.core.kueue_manager import KueueConfig, KueueManager, has_sub_slicing_enabled
-from xpk.core.system_characteristics import AcceleratorType, SystemCharacteristics
+from xpk.core.system_characteristics import AcceleratorType, SystemCharacteristics, UserFacingNameToSystemCharacteristics
 from xpk.core.testing.commands_tester import CommandsTester
 from packaging.version import Version
 
@@ -61,6 +61,13 @@ def set_installed_kueue_version(
   )
 
 
+@pytest.fixture(autouse=True)
+def mock_ask_for_user_consent(mocker: MockerFixture) -> MagicMock:
+  return mocker.patch(
+      "xpk.core.kueue_manager.ask_for_user_consent", return_value=True
+  )
+
+
 @pytest.fixture(autouse=True)
 def mock_commands(mocker: MockerFixture) -> CommandsTester:
   return CommandsTester(
@@ -78,7 +85,7 @@ def mock_commands(mocker: MockerFixture) -> CommandsTester:
 @pytest.fixture(autouse=True)
 @patch("jinja2.Environment", return_value=MagicMock())
 def kueue_manager(mock_env: MagicMock) -> KueueManager:
-  return KueueManager()
+  return KueueManager("test-project", "test-zone")
 
 
 def test_install_or_upgrade_when_newer_version_already_installed(
@@ -102,7 +109,7 @@ def test_install_or_upgrade_when_outdated(
   result = kueue_manager.install_or_upgrade(KUEUE_CONFIG)
 
   assert result == 0
-  mock_commands.assert_command_run("kubectl apply", "v0.12.2/manifests.yaml")
+  mock_commands.assert_command_run("kubectl apply", "v0.14.3/manifests.yaml")
   mock_commands.assert_command_run("kubectl apply -f", "/tmp/")
 
 
@@ -115,10 +122,84 @@ def test_install_or_upgrade_when_not_installed(
   result = kueue_manager.install_or_upgrade(KUEUE_CONFIG)
 
   assert result == 0
-  mock_commands.assert_command_run("kubectl apply", "v0.12.2/manifests.yaml")
+  mock_commands.assert_command_run("kubectl apply", "v0.14.3/manifests.yaml")
   mock_commands.assert_command_run("kubectl apply -f", "/tmp/")
 
 
+def test_upgrade_when_no_breaking_changes_between_versions_no_preparation_needed(
+    mock_commands: CommandsTester,
+    kueue_manager: KueueManager,
+    mock_ask_for_user_consent: MagicMock,
+):
+  set_installed_kueue_version(mock_commands, Version("0.14.0"))
+
+  kueue_manager.install_or_upgrade(KUEUE_CONFIG)
+
+  mock_ask_for_user_consent.assert_not_called()
+
+
+def test_upgrade_with_breaking_changes_between_versions_runs_preparation(
+    mock_commands: CommandsTester,
+    kueue_manager: KueueManager,
+    mock_ask_for_user_consent: MagicMock,
+):
+  set_installed_kueue_version(mock_commands, Version("0.11.0"))
+  fake_crds = (
+      "customresourcedefinition.apiextensions.k8s.io/kueue-crd-1.kueue.x-k8s.io\n"
+      "customresourcedefinition.apiextensions.k8s.io/kueue-crd-2.kueue.x-k8s.io"
+  )
+  mock_commands.set_result_for_command(
+      (0, fake_crds), "kubectl get crd -o name"
+  )
+  mock_ask_for_user_consent.return_value = True
+
+  result = kueue_manager.install_or_upgrade(KUEUE_CONFIG)
+
+  assert result == 0
+  mock_ask_for_user_consent.assert_called_once()
+  assert (
+      "CHANGELOG/CHANGELOG-0.14.md"
+      in mock_ask_for_user_consent.mock_calls[0].args[0]
+  )
+  mock_commands.assert_command_run(
+      "kubectl delete kueue-crd-1.kueue.x-k8s.io --all"
+  )
+  mock_commands.assert_command_run(
+      "kubectl delete kueue-crd-2.kueue.x-k8s.io --all"
+  )
+  mock_commands.assert_command_run(
+      "kubectl delete crd kueue-crd-1.kueue.x-k8s.io"
+  )
+  mock_commands.assert_command_run(
+      "kubectl delete crd kueue-crd-2.kueue.x-k8s.io"
+  )
+  mock_commands.assert_command_run(
+      "kubectl delete deployment kueue-controller-manager"
+  )
+
+
+def test_upgrade_with_breaking_changes_between_versions_does_not_run_preparation_without_consent(
+    mock_commands: CommandsTester,
+    kueue_manager: KueueManager,
+    mock_ask_for_user_consent: MagicMock,
+):
+  set_installed_kueue_version(mock_commands, Version("0.11.0"))
+  mock_commands.set_result_for_command(
+      (
+          0,
+          "customresourcedefinition.apiextensions.k8s.io/kueue-crd-1.kueue.x-k8s.io",
+      ),
+      "kubectl get crd -o name",
+  )
+  mock_ask_for_user_consent.return_value = False
+
+  result = kueue_manager.install_or_upgrade(KUEUE_CONFIG)
+
+  assert result == 1
+  # Assert there was no command run for the Kueue crd:
+  mock_commands.assert_command_not_run("kueue-crd-1.kueue.x-k8s.io")
+
+
 def test_installation_with_tolerations(
     mock_commands: CommandsTester, kueue_manager: KueueManager
 ):
@@ -199,6 +280,10 @@ def test_configure_generates_correct_manifest_for_tpu(
 ):
   """Test that __configure generates the correct manifest content for TPUs."""
   set_installed_kueue_version(mock_commands, None)
+  mock_commands.set_result_for_command(
+      (0, "100 102400"), "gcloud compute machine-types describe"
+  )
+
   tpu_kueue_config = dataclasses.replace(
       KUEUE_CONFIG, system=TPU_SYSTEM, num_slices=2
   )
@@ -239,6 +324,39 @@ def test_configure_generates_correct_manifest_for_tpu(
   )
 
 
+@patch("xpk.core.kueue_manager.write_tmp_file")
+def test_install_autocorrects_resource_limits(
+    write_tmp_file_mock: MagicMock,
+    mock_commands: CommandsTester,
+    kueue_manager: KueueManager,
+):
+  """Test that installation auto-corrects the specified resource limits."""
+  set_installed_kueue_version(mock_commands, None)
+  # set 50 vCPU, 200Gi memory
+  mock_commands.set_result_for_command(
+      (0, "50 204800"), "gcloud compute machine-types describe"
+  )
+
+  kueue_config = dataclasses.replace(
+      KUEUE_CONFIG, cpu_limit=100, memory_limit="100Gi"
+  )
+
+  kueue_manager.install_or_upgrade(kueue_config)
+
+  rendered_manifest: str = write_tmp_file_mock.call_args[0][0]
+  manifest_docs = list(yaml.safe_load_all(rendered_manifest))
+  cluster_queue = _first(
+      doc for doc in manifest_docs if doc["kind"] == "ClusterQueue"
+  )
+  resources = cluster_queue["spec"]["resourceGroups"][0]["flavors"][0][
+      "resources"
+  ]
+  cpu_resource = _first(r for r in resources if r["name"] == "cpu")
+  memory_resource = _first(r for r in resources if r["name"] == "memory")
+  assert cpu_resource["nominalQuota"] == 50
+  assert memory_resource["nominalQuota"] == "204800Mi"
+
+
 @patch("xpk.core.kueue_manager.write_tmp_file")
 def test_configure_generates_manifest_with_admission_checks_for_flex_single_slice(
     write_tmp_file_mock: MagicMock,
@@ -317,6 +435,7 @@ def test_configure_generates_correct_manifest_with_sub_slicing(
   kueue_config = dataclasses.replace(
       KUEUE_CONFIG,
       configure_sub_slicing=True,
+      system=UserFacingNameToSystemCharacteristics["v6e-8x8"],
   )
 
   kueue_manager.install_or_upgrade(kueue_config)
@@ -329,6 +448,15 @@ def test_configure_generates_correct_manifest_with_sub_slicing(
   assert resource_flavor["spec"]["topologyName"] == "sub-slice-topology"
   topology = _first(doc for doc in manifest_docs if doc["kind"] == "Topology")
   assert topology["metadata"]["name"] == "sub-slice-topology"
+  expected_levels = [
+      "cloud.google.com/gke-tpu-slice-8x8-id",
+      "cloud.google.com/gke-tpu-slice-4x8-id",
+      "cloud.google.com/gke-tpu-slice-4x4-id",
+      "cloud.google.com/gke-tpu-slice-2x4-id",
+      "kubernetes.io/hostname",
+  ]
+  actual_levels = [level["nodeLabel"] for level in topology["spec"]["levels"]]
+  assert actual_levels == expected_levels
 
 
 @patch("xpk.core.kueue_manager.write_tmp_file")

xpk/core/nodepool.py

@@ -15,8 +15,8 @@ limitations under the License.
 """
 
 from typing import List
-from ..utils.console import get_user_input, xpk_print
-from ..utils.topology import get_topology_product, is_topology_valid
+from ..utils.console import ask_for_user_consent, xpk_print
+from .scheduling import get_placement_policy_name, is_placement_policy_supported
 from .capacity import (
     AUTOPROVISIONING_CONFIG_VALUE,
     H100_MEGA_DEVICE_TYPE,
@@ -110,6 +110,7 @@ def run_gke_node_pool_create_command(
       existing_node_pool_names, args.cluster, desired_node_pool_count
   )
 
+  node_pools_to_delete = []
   node_pools_to_remain = []
   delete_commands = []
   delete_task_names = []
@@ -186,14 +187,10 @@ def run_gke_node_pool_create_command(
   # when cluster is getting updated from 'x' device_type/gke_accelerator to 'y' device_type/gke_accelerator.
   # In that case, '{args.cluster}-np-i' nodepool will be re-created for 'y' device_type/gke_accelerator.
   if delete_commands:
-    will_delete = True
-    if node_pools_to_delete and not args.force:
-      will_delete = get_user_input(
-          f'Planning to delete {len(node_pools_to_delete)} node pools including'
-          f' {node_pools_to_delete}. \nDo you wish to delete: y (yes) / n'
-          ' (no):\n'
-      )
-    if not will_delete:
+    if node_pools_to_delete and not ask_for_user_consent(
+        f'Planning to delete {len(node_pools_to_delete)} node pools including'
+        f' {node_pools_to_delete}. \nDo you wish to delete?'
+    ):
       xpk_print(
           'You have requested to not delete the existing nodepools in the'
           ' cluster. There will be no change to the cluster.'
@@ -215,18 +212,15 @@ def run_gke_node_pool_create_command(
 
   # Enable Workload Identity on existing Nodepools
   if update_WI_commands:
-    will_update_WI = True
-    if node_pools_to_update_WI and not args.force:
-      will_update_WI = get_user_input(
-          'Planning to enable Workload Identity Federation on'
-          f' {len(node_pools_to_update_WI)} existing node pools including'
-          f' {node_pools_to_update_WI}.This immediately enables Workload'
-          ' Identity Federation for GKE for any workloads running in the node'
-          ' pool. Also, xpk does not support disabling Workload Identity on'
-          ' clusters that have it enabled already \nDo you wish to update: y'
-          ' (yes) / n (no):\n'
-      )
-    if not will_update_WI:
+    will_update_WI = not node_pools_to_update_WI or ask_for_user_consent(
+        'Planning to enable Workload Identity Federation on'
+        f' {len(node_pools_to_update_WI)} existing node pools including'
+        f' {node_pools_to_update_WI}. This immediately enables Workload'
+        ' Identity Federation for GKE for any workloads running in the node'
+        ' pool. Also, xpk does not support disabling Workload Identity on'
+        ' clusters that have it enabled already \nDo you wish to update?'
+    )
+    if will_update_WI:
       for i, command in enumerate(update_WI_commands):
         xpk_print(
             f'To complete {update_WI_task_names[i]} we are executing {command}'
@@ -264,10 +258,8 @@ def run_gke_node_pool_create_command(
         return 1
 
   placement_args = ''
-  if system.requires_workload_policy and is_topology_valid(system.topology):
-    placement_policy = (
-        f'{system.device_type}-{system.topology}-placement-policy'
-    )
+  if is_placement_policy_supported(system):
+    placement_policy = get_placement_policy_name(system)
     ensure_resource_policy_exists(placement_policy, args, system.topology)
     placement_args = f' --placement-policy={placement_policy}'
 
@@ -290,16 +282,16 @@ def run_gke_node_pool_create_command(
     )
     if system.accelerator_type == AcceleratorType.TPU:
       command += f' --node-version={gke_node_pool_version}'
-      topology_product = get_topology_product(system.topology)
       if capacity_type == CapacityType.FLEX_START:
         command += ' --num-nodes=0'
-      elif topology_product > 1:
+      else:
         command += f' --num-nodes={system.vms_per_slice}'
       command += (
           f' --scopes=storage-full,gke-default,{CLOUD_PLATFORM_AUTH_SCOPE_URL}'
       )
 
-      if topology_product > 1:
+      # --tpu-topology should not be set for single-host node pools
+      if system.vms_per_slice > 1:
         # --placement-type=COMPACT enables group placement policy which
         # is mutually exclusive with workload policy, --tpu-topology should
         # also not be passed when workload policy is used

xpk/core/nodepool_test.py

@@ -145,22 +145,24 @@ def mock_nodepool_dependencies(mocker):
       "xpk.core.nodepool.get_cluster_location", return_value="us-central1"
   )
   mocker.patch("xpk.core.nodepool.run_commands", return_value=0)
-  mocker.patch("xpk.core.nodepool.get_user_input", return_value=True)
-  mock_is_topology_valid = mocker.patch("xpk.core.nodepool.is_topology_valid")
+  mocker.patch("xpk.core.nodepool.ask_for_user_consent", return_value=True)
+  mock_is_placement_policy_supported = mocker.patch(
+      "xpk.core.nodepool.is_placement_policy_supported"
+  )
   mock_ensure_resource_policy = mocker.patch(
       "xpk.core.nodepool.ensure_resource_policy_exists"
   )
-  return mock_is_topology_valid, mock_ensure_resource_policy
+  return mock_is_placement_policy_supported, mock_ensure_resource_policy
 
 
 def test_placement_policy_created_for_gpu_with_valid_topology(
     mocker, mock_nodepool_dependencies
 ):
   """Tests that placement policy is created for GPUs with a valid topology."""
-  mock_is_topology_valid, mock_ensure_resource_policy = (
+  mock_is_placement_policy_supported, mock_ensure_resource_policy = (
       mock_nodepool_dependencies
   )
-  mock_is_topology_valid.return_value = True
+  mock_is_placement_policy_supported.return_value = True
   args = mocker.Mock(
       tpu_type=None,
       device_type="h100-80gb-8",
@@ -170,7 +172,7 @@ def test_placement_policy_created_for_gpu_with_valid_topology(
   )
   system = SystemCharacteristics(
       topology="N/A",
-      vms_per_slice=1,
+      vms_per_slice=2,
       gke_accelerator="nvidia-h100-80gb",
       gce_machine_type="a3-highgpu-8g",
       chips_per_vm=8,
@@ -188,10 +190,10 @@ def test_placement_policy_not_created_for_gpu_with_invalid_topology(
     mocker, mock_nodepool_dependencies
 ):
   """Tests that placement policy is not created for GPUs with an invalid topology."""
-  mock_is_topology_valid, mock_ensure_resource_policy = (
+  mock_is_placement_policy_supported, mock_ensure_resource_policy = (
       mock_nodepool_dependencies
   )
-  mock_is_topology_valid.return_value = False
+  mock_is_placement_policy_supported.return_value = False
   args = mocker.Mock(
       tpu_type=None,
       device_type="h100-80gb-8",
@@ -200,7 +202,7 @@ def test_placement_policy_not_created_for_gpu_with_invalid_topology(
   )
   system = SystemCharacteristics(
       topology="N/A",
-      vms_per_slice=1,
+      vms_per_slice=2,
       gke_accelerator="nvidia-h100-80gb",
       gce_machine_type="a3-highgpu-8g",
       chips_per_vm=8,
@@ -218,10 +220,10 @@ def test_placement_policy_created_for_tpu7x_with_valid_topology(
     mocker, mock_nodepool_dependencies
 ):
   """Tests that placement policy is created for tpu7x with a valid topology."""
-  mock_is_topology_valid, mock_ensure_resource_policy = (
+  mock_is_placement_policy_supported, mock_ensure_resource_policy = (
       mock_nodepool_dependencies
   )
-  mock_is_topology_valid.return_value = True
+  mock_is_placement_policy_supported.return_value = True
   args = mocker.Mock(
       tpu_type="tpu7x-8",
       device_type=None,
@@ -232,7 +234,7 @@ def test_placement_policy_created_for_tpu7x_with_valid_topology(
   )
   system = SystemCharacteristics(
       topology="2x2x1",
-      vms_per_slice=1,
+      vms_per_slice=2,
       gke_accelerator="tpu7x",
       gce_machine_type="tpu7x-standard-4t",
       chips_per_vm=4,
@@ -251,14 +253,14 @@ def test_placement_policy_not_created_for_non7x_tpu(
     mocker, mock_nodepool_dependencies
 ):
   """Tests that placement policy is not created for non-tpu7x TPUs."""
-  mock_is_topology_valid, mock_ensure_resource_policy = (
+  mock_is_placement_policy_supported, mock_ensure_resource_policy = (
       mock_nodepool_dependencies
   )
-  mock_is_topology_valid.return_value = True
+  mock_is_placement_policy_supported.return_value = False
   args = mocker.Mock(
       tpu_type="v6e",
       device_type=None,
-      num_slices=1,
+      num_slices=2,
       cluster="test-cluster",
       project="test-project",
       zone="us-central1-a",

xpk/core/scheduling.py

@@ -14,7 +14,9 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+from ..utils.topology import get_slice_topology_level
 from ..utils.console import xpk_print
+from ..utils.topology import is_topology_valid
 from ..utils.execution_context import is_dry_run
 from .capacity import AUTOPROVISIONING_CONFIG_MAXIMUM_KEY, AUTOPROVISIONING_CONFIG_VALUE
 from .resources import CLUSTER_RESOURCES_CONFIGMAP, get_cluster_configmap
@@ -299,7 +301,20 @@ def create_sub_slicing_annotations(sub_slicing_topology: str) -> list[str]:
   return [
       (
           'kueue.x-k8s.io/podset-required-topology:'
-          f' "google.com/gke-tpu-slice-{sub_slicing_topology}-id"'
+          f' "{get_slice_topology_level(sub_slicing_topology)}"'
       ),
       f'cloud.google.com/gke-tpu-slice-topology: {sub_slicing_topology}',
   ]
+
+
+def create_placement_policy_label(system: SystemCharacteristics) -> str:
+  name = get_placement_policy_name(system)
+  return f'cloud.google.com/placement-policy-name: {name}'
+
+
+def get_placement_policy_name(system: SystemCharacteristics) -> str:
+  return f'{system.device_type}-{system.topology}-placement-policy'
+
+
+def is_placement_policy_supported(system: SystemCharacteristics) -> bool:
+  return system.requires_workload_policy and is_topology_valid(system.topology)

xpk/core/scheduling_test.py

@@ -14,18 +14,97 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-from .scheduling import create_sub_slicing_annotations
+from .scheduling import create_sub_slicing_annotations, create_placement_policy_label, get_placement_policy_name, is_placement_policy_supported
+from .system_characteristics import SystemCharacteristics, AcceleratorType
 
 
 def test_create_sub_slicing_annotations_returns_valid_annotations():
-  subslicing_topology = '2x2'
-
-  result = create_sub_slicing_annotations(subslicing_topology)
+  result = create_sub_slicing_annotations(sub_slicing_topology='2x4')
 
   assert result == [
       (
           'kueue.x-k8s.io/podset-required-topology:'
-          ' "google.com/gke-tpu-slice-2x2-id"'
+          ' "cloud.google.com/gke-tpu-slice-2x4-id"'
       ),
-      'cloud.google.com/gke-tpu-slice-topology: 2x2',
+      'cloud.google.com/gke-tpu-slice-topology: 2x4',
   ]
+
+
+def test_create_placement_policy_label_returns_valid_label():
+  system_characteristics = SystemCharacteristics(
+      chips_per_vm=1,
+      gce_machine_type='tpu7x-standard-1t',
+      gke_accelerator='tpu7x',
+      requires_workload_policy=False,
+      topology='1x1x1',
+      vms_per_slice=1,
+      device_type='tpu7x',
+      accelerator_type=AcceleratorType.TPU,
+      supports_sub_slicing=False,
+  )
+  label = create_placement_policy_label(system_characteristics)
+  assert (
+      label
+      == 'cloud.google.com/placement-policy-name: tpu7x-1x1x1-placement-policy'
+  )
+
+
+def test_get_placement_policy_name_returns_valid_name():
+  system_characteristics = SystemCharacteristics(
+      chips_per_vm=1,
+      gce_machine_type='tpu7x-standard-1t',
+      gke_accelerator='tpu7x',
+      requires_workload_policy=False,
+      topology='1x1x1',
+      vms_per_slice=1,
+      device_type='tpu7x',
+      accelerator_type=AcceleratorType.TPU,
+      supports_sub_slicing=False,
+  )
+  name = get_placement_policy_name(system_characteristics)
+  assert name == 'tpu7x-1x1x1-placement-policy'
+
+
+def test_is_placement_policy_supported_returns_true_for_system_characteristics_supporting_workload_policy_and_having_valid_topology():
+  system_characteristics = SystemCharacteristics(
+      chips_per_vm=1,
+      gce_machine_type='tpu7x-standard-1t',
+      gke_accelerator='tpu7x',
+      requires_workload_policy=True,
+      topology='1x1x1',
+      vms_per_slice=1,
+      device_type='tpu7x',
+      accelerator_type=AcceleratorType.TPU,
+      supports_sub_slicing=False,
+  )
+  assert is_placement_policy_supported(system_characteristics) is True
+
+
+def test_is_placement_policy_supported_returns_false_for_system_characteristics_not_supporting_workload_policy_and_having_valid_topology():
+  system_characteristics = SystemCharacteristics(
+      chips_per_vm=1,
+      gce_machine_type='tpu7x-standard-1t',
+      gke_accelerator='tpu7x',
+      requires_workload_policy=False,
+      topology='1x1x1',
+      vms_per_slice=1,
+      device_type='tpu7x',
+      accelerator_type=AcceleratorType.TPU,
+      supports_sub_slicing=False,
+  )
+  assert is_placement_policy_supported(system_characteristics) is False
+
+
+def test_is_placement_policy_supported_returns_false_for_system_characteristics_supporting_workload_policy_and_having_invalid_topology():
+  system_characteristics = SystemCharacteristics(
+      chips_per_vm=1,
+      gce_machine_type='tpu7x-standard-1t',
+      gke_accelerator='tpu7x',
+      requires_workload_policy=True,
+      topology='aaa',
+      vms_per_slice=1,
+      device_type='tpu7x',
+      accelerator_type=AcceleratorType.TPU,
+      supports_sub_slicing=False,
+  )
+  assert is_placement_policy_supported(system_characteristics) is False