xpk 0.14.4__py3-none-any.whl → 0.16.0__py3-none-any.whl

xpk/core/kueue_manager.py

@@ -20,17 +20,17 @@ from dataclasses import dataclass
 from typing import Optional, List, Dict, Any
 import json
 from jinja2 import Environment, FileSystemLoader
-from ..utils.execution_context import is_dry_run
-from ..utils.kueue import is_queued_cluster
 
-from .capacity import B200_DEVICE_TYPE, H100_MEGA_DEVICE_TYPE, H200_DEVICE_TYPE
-from .scheduling import (
-    create_accelerator_label,
-    create_machine_label,
-)
+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
 from .system_characteristics import (
+    SUB_SLICING_TOPOLOGIES,
+    AcceleratorType,
     AcceleratorTypeToAcceleratorCharacteristics,
     SystemCharacteristics,
+    create_accelerator_label,
+    create_machine_label,
 )
 from ..core.commands import (
     run_command_for_value,
@@ -38,10 +38,12 @@ from ..core.commands import (
     run_command_with_updates_retry,
 )
 from ..utils.file import write_tmp_file
-from ..utils.console import xpk_print, xpk_exit
+from ..utils.console import xpk_print, xpk_exit, ask_for_user_consent
 from ..utils.templates import TEMPLATE_PATH, get_templates_absolute_path
 from packaging.version import Version
 
+KUEUE_VERSION = Version("v0.14.3")
+LATEST_BREAKING_VERSION = Version("v0.14.0")
 WAIT_FOR_KUEUE_TIMEOUT = "10m"
 CLUSTER_QUEUE_NAME = "cluster-queue"
 LOCAL_QUEUE_NAME = "multislice-queue"
@@ -52,10 +54,9 @@ KUEUE_CONTROLLER_MANAGER_JINJA_FILE = "kueue_controller_manager.yaml.j2"
 KUEUE_SUB_SLICING_TOPOLOGY_JINJA_FILE = "kueue_sub_slicing_topology.yaml.j2"
 MEMORY_SIZE_PER_VM = 1.2
 MIN_MEMORY_LIMIT_SIZE = 4096
-KUEUE_VERSION = Version("v0.12.2")
 
 
-@dataclass
+@dataclass(frozen=True)
 class KueueConfig:
   system: SystemCharacteristics
   total_chips: int
@@ -68,7 +69,7 @@ class KueueConfig:
   num_slices: int = 1
 
 
-@dataclass
+@dataclass(frozen=True)
 class _NameAndYaml:
   name: str
   yaml: str
@@ -79,9 +80,13 @@ class KueueManager:
 
   def __init__(
       self,
+      project: str,
+      zone: str,
       kueue_version: Version = KUEUE_VERSION,
       template_path=TEMPLATE_PATH,
   ):
+    self.project = project
+    self.zone = zone
     self.kueue_version = kueue_version
 
     self.template_env = Environment(
@@ -102,10 +107,10 @@ class KueueManager:
     Args:
         tolerations: An optional list of tolerations to apply to the kueue-controller-manager.
     """
-    return_code, installed_version = self.get_installed_kueue_version()
+    return_code, installed_version = get_installed_kueue_version()
 
-    if return_code == 0:
-      if installed_version and installed_version > self.kueue_version:
+    if return_code == 0 and installed_version:
+      if installed_version > self.kueue_version:
         xpk_print(
             f"Cluster has a newer Kueue version, {installed_version}. Skipping"
             " installation."
@@ -113,6 +118,10 @@ class KueueManager:
         return 0
       else:
         xpk_print(f"Upgrading Kueue to version v{self.kueue_version}...")
+        assert installed_version
+        prepare_code = self.__prepare_for_upgrade(installed_version)
+        if prepare_code != 0:
+          return prepare_code
     else:
       xpk_print(f"Installing Kueue version v{self.kueue_version}...")
 
@@ -122,24 +131,6 @@ class KueueManager:
 
     return self.__configure(kueue_config)
 
-  def get_installed_kueue_version(self) -> tuple[int, Version | None]:
-    command = (
-        "kubectl get deployment kueue-controller-manager -n kueue-system -o"
-        " jsonpath='{.spec.template.spec.containers[0].image}'"
-    )
-    task = "Get kueue version on server"
-    return_code, val = run_command_for_value(
-        command,
-        task,
-        dry_run_return_val="",
-    )
-    if return_code != 0:
-      return return_code, None
-    version_tag = val.split(":")
-    if len(version_tag) == 1:
-      return 1, None
-    return return_code, Version(version_tag[-1])
-
   def __install(
       self,
       tolerations: Optional[List[Dict[str, Any]]] = None,
@@ -161,6 +152,60 @@ class KueueManager:
 
     return self.__wait_for_kueue_available()
 
+  def __prepare_for_upgrade(self, installed_version: Version) -> int:
+    if installed_version >= LATEST_BREAKING_VERSION:
+      return 0
+
+    xpk_print(
+        f"Currently installed Kueue version v{installed_version} is"
+        f" incompatible with the newer v{self.kueue_version}."
+    )
+
+    changelog_link = f"https://github.com/kubernetes-sigs/kueue/blob/main/CHANGELOG/CHANGELOG-{self.kueue_version.major}.{self.kueue_version.minor}.md"
+    agreed = ask_for_user_consent(
+        "Do you want to allow XPK to update Kueue automatically? This will"
+        " delete all existing Kueue resources and create new ones. If you"
+        " decline, you will need to upgrade the Kueue manually (see"
+        f" {changelog_link} for help)."
+    )
+    if not agreed:
+      return 1
+
+    return self.__delete_all_kueue_resources()
+
+  def __delete_all_kueue_resources(self) -> int:
+    return_code, kueue_crds_string = run_command_for_value(
+        "kubectl get crd -o name | grep .kueue.x-k8s.io", "Get Kueue CRDs"
+    )
+    if return_code != 0:
+      return return_code
+
+    kueue_crds = [
+        line.strip().removeprefix(
+            "customresourcedefinition.apiextensions.k8s.io/"
+        )
+        for line in kueue_crds_string.strip().split("\n")
+    ]
+
+    for crd in kueue_crds:
+      return_code = run_command_with_updates(
+          f"kubectl delete {crd} --all", f"Delete all resources of type {crd}"
+      )
+      if return_code != 0:
+        return return_code
+
+    for crd in kueue_crds:
+      return_code = run_command_with_updates(
+          f"kubectl delete crd {crd}", f"Delete CRD {crd}"
+      )
+      if return_code != 0:
+        return return_code
+
+    return run_command_with_updates(
+        "kubectl delete deployment kueue-controller-manager -n kueue-system",
+        "Delete Kueue Controller Manager deployment",
+    )
+
   def __install_kueue_crs(self) -> int:
     manifest_url = f"https://github.com/kubernetes-sigs/kueue/releases/download/v{self.kueue_version}/manifests.yaml"
     install_command = (
@@ -228,6 +273,7 @@ class KueueManager:
     topology_name = (
         topology_name_and_yaml.name if topology_name_and_yaml else None
     )
+    cpu_limit, memory_limit = self.__autocorrect_resource_limits(kueue_config)
 
     # The manager builds the context internally based on its opinionated logic
     context = self.__build_template_context(
@@ -237,8 +283,8 @@ class KueueManager:
         autoprovisioning=kueue_config.autoprovisioning_enabled,
         flex=kueue_config.flex,
         num_slices=kueue_config.num_slices,
-        cpu_limit=kueue_config.cpu_limit,
-        memory_limit=kueue_config.memory_limit,
+        cpu_limit=cpu_limit,
+        memory_limit=memory_limit,
         topology_name=topology_name,
     )
 
@@ -273,19 +319,16 @@ class KueueManager:
     main_flavor_name = f"{num_slices}x{device_type_str}"
 
     node_labels_dict = {}
-    accelerator_label = create_accelerator_label(
-        system.accelerator_type, system
-    )
+    accelerator_label = create_accelerator_label(system)
     if accelerator_label:
       key, value = accelerator_label.split(":", 1)
       node_labels_dict[key] = value.strip()
 
-    machine_label = create_machine_label(
-        system.accelerator_type, system, autoprovisioning
-    )
-    if machine_label:
-      key, value = machine_label.split(":", 1)
-      node_labels_dict[key] = value.strip()
+    if not autoprovisioning:
+      machine_label = create_machine_label(system)
+      if machine_label:
+        key, value = machine_label.split(":", 1)
+        node_labels_dict[key] = value.strip()
 
     topology_label = f"topologyName: {topology_name}" if topology_name else ""
 
@@ -352,11 +395,10 @@ class KueueManager:
   def __get_topology_name_and_yaml(
       self, system: SystemCharacteristics, configure_sub_slicing: bool
   ) -> _NameAndYaml | None:
-    if system.device_type in [
-        H100_MEGA_DEVICE_TYPE,
-        H200_DEVICE_TYPE,
-        B200_DEVICE_TYPE,
-    ]:
+    if (
+        system.accelerator_type == AcceleratorType["GPU"]
+        and system.gpu_requires_topology
+    ):
       return _NameAndYaml(
           name="gke-default",
           yaml=self.template_env.get_template(
@@ -364,12 +406,25 @@ class KueueManager:
           ).render(),
       )
     elif configure_sub_slicing:
+      sorted_topologies = sorted(
+          SUB_SLICING_TOPOLOGIES, key=get_topology_product, reverse=True
+      )
+      levels = [
+          get_slice_topology_level(topology)
+          for topology in sorted_topologies
+          if is_topology_contained(
+              contained=topology, container=system.topology
+          )
+      ]
+      levels.append("kubernetes.io/hostname")
+
       return _NameAndYaml(
           name=SUB_SLICE_TOPOLOGY_NAME,
           yaml=self.template_env.get_template(
               KUEUE_SUB_SLICING_TOPOLOGY_JINJA_FILE
           ).render({
               "sub_slice_topology_name": SUB_SLICE_TOPOLOGY_NAME,
+              "levels": levels,
           }),
       )
     else:
@@ -377,8 +432,6 @@ class KueueManager:
 
   def __apply_manifest(self, manifest: str) -> int:
     task = "Applying Kueue Custom Resources"
-    if is_dry_run():
-      xpk_print(f"Applying following Kueue resources:{manifest}")
     tmp_file = write_tmp_file(manifest)
     command = f"kubectl apply -f {tmp_file}"
     return run_command_with_updates(command, task)
@@ -422,13 +475,114 @@ class KueueManager:
       xpk_print(f"{task} returned ERROR {return_code}")
     return return_code
 
+  def __autocorrect_resource_limits(
+      self, kueue_config: KueueConfig
+  ) -> tuple[int, str]:
+    """Verify specified CPU and memory limits against machine type."""
+
+    cpu_limit = kueue_config.cpu_limit
+    memory_limit_str = kueue_config.memory_limit
+    if not cpu_limit and not memory_limit_str:
+      return cpu_limit, memory_limit_str
+
+    # Get CPU and memory capacity from machine type
+    command = (
+        "gcloud compute machine-types describe"
+        f" {kueue_config.system.gce_machine_type} "
+        f" --project={self.project} --zone={self.zone}"
+        " --format='value(guestCpus,memoryMb)'"
+    )
+    return_code, out = run_command_for_value(
+        command,
+        "Get vCPU and memory capacity for machine type",
+        dry_run_return_val="10 10",
+    )
+    if return_code != 0:
+      xpk_print(
+          "Unable to verify vCPU and memory capacity for machine type."
+          " XPK will proceed with using user-defined  limits."
+      )
+      return cpu_limit, memory_limit_str
+
+    cpu_capacity_str, memory_capacity_MB_str = out.split()
+    if cpu_limit:
+      cpu_limit = _autocorrect_cpu_limit(cpu_limit, int(cpu_capacity_str))
+    if memory_limit_str:
+      memory_limit_str = _autocorrect_memory_limit(
+          memory_limit_str, memory_capacity_MB_str
+      )
+    return cpu_limit, memory_limit_str
+
+
+def get_installed_kueue_version(
+    dry_run_version: Version | None = None,
+) -> tuple[int, Version | None]:
+  command = (
+      "kubectl get deployment kueue-controller-manager -n kueue-system -o"
+      " jsonpath='{.spec.template.spec.containers[0].image}'"
+  )
+  task = "Get kueue version on server"
+  return_code, val = run_command_for_value(
+      command,
+      task,
+      dry_run_return_val=(
+          f"registry.k8s.io/kueue/kueue:v{dry_run_version}"
+          if dry_run_version
+          else ""
+      ),
+  )
+  if return_code != 0:
+    return return_code, None
+  version_tag = val.split(":")
+  if len(version_tag) == 1:
+    return 1, None
+  return return_code, Version(version_tag[-1])
+
 
 def has_sub_slicing_enabled() -> tuple[int, bool | None]:
   return_code, value = run_command_for_value(
-      command="kubectl get topology", task="Get defined topologies"
+      command="kubectl get topology",
+      task="Get defined topologies",
+      dry_run_return_val=SUB_SLICE_TOPOLOGY_NAME,
   )
 
   if return_code != 0:
     return return_code, None
 
   return return_code, SUB_SLICE_TOPOLOGY_NAME in value
+
+
+def _autocorrect_cpu_limit(cpu_limit: int, cpu_capacity: int) -> int:
+  if cpu_limit > cpu_capacity:
+    xpk_print(
+        "The CPU limit is above the available capacity."
+        f" We will set CPU limit to {cpu_capacity}."
+    )
+  elif cpu_limit < cpu_capacity:
+    xpk_print(
+        "The CPU limit is below the available capacity, which would lead"
+        f" to underutilization. We will set CPU limit to {cpu_capacity}."
+    )
+  return cpu_capacity
+
+
+def _autocorrect_memory_limit(
+    memory_limit_str: str, memory_capacity_MB_str: str
+) -> str:
+  memory_limit_bytes = parse_quantity(memory_limit_str)
+  memory_capacity_bytes = int(memory_capacity_MB_str) << 20
+  if memory_limit_bytes == memory_capacity_bytes:
+    return memory_limit_str
+  memory_limit_str = memory_capacity_MB_str + "Mi"
+  if memory_limit_bytes > memory_capacity_bytes:
+    xpk_print(
+        "The memory limit is above the available capacity. We will set"
+        f" memory limit to {memory_limit_str}."
+    )
+  else:
+    xpk_print(
+        "The memory limit is below the available capacity, which would"
+        " lead to underutilization. We will set the memory limit to"
+        f" {memory_limit_str}."
+    )
+  return memory_limit_str

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 GpuConfig, DockerPlatform, AcceleratorType, SystemCharacteristics, UserFacingNameToSystemCharacteristics
 from xpk.core.testing.commands_tester import CommandsTester
 from packaging.version import Version
 
@@ -35,6 +35,7 @@ TPU_SYSTEM: SystemCharacteristics = SystemCharacteristics(
     accelerator_type=AcceleratorType.TPU,
     device_type="v5p-8",
     supports_sub_slicing=False,
+    docker_platform=DockerPlatform.ARM,
 )
 
 KUEUE_CONFIG: KueueConfig = KueueConfig(
@@ -61,6 +62,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 +86,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 +110,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 +123,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 +281,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 +325,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,
@@ -287,6 +406,8 @@ def test_configure_generates_correct_manifest_with_gke_default_topology(
           accelerator_type=AcceleratorType.GPU,
           device_type="h100-mega-80gb-8",
           supports_sub_slicing=False,
+          docker_platform=DockerPlatform.ARM,
+          gpu_config=GpuConfig(requires_topology=True),
       ),
   )
 
@@ -317,6 +438,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 +451,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")
@@ -373,6 +504,29 @@ def test_configure_generates_correct_manifest_with_pathways(
   assert pathways_rg["flavors"][0]["resources"][1]["nominalQuota"] == "2000G"
 
 
+@patch("xpk.core.kueue_manager.write_tmp_file")
+def test_configure_generates_correct_manifest_for_a4x(
+    write_tmp_file_mock: MagicMock,
+    mock_commands: CommandsTester,
+    kueue_manager: KueueManager,
+):
+  """Test that __configure generates correct manifest for a4x GPUs."""
+  set_installed_kueue_version(mock_commands, None)
+  kueue_config = dataclasses.replace(
+      KUEUE_CONFIG,
+      system=UserFacingNameToSystemCharacteristics["gb200-4"],
+  )
+
+  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))
+
+  # Check that the gke-default topology is present for a4x.
+  topology = _first(doc for doc in manifest_docs if doc["kind"] == "Topology")
+  assert topology["metadata"]["name"] == "gke-default"
+
+
 def test_has_sub_slicing_enabled_returns_exit_code_when_command_fails(
     mock_commands: CommandsTester,
 ):

xpk/core/nap.py

@@ -30,9 +30,8 @@ from .commands import run_command_with_updates, run_commands
 from .gcloud_context import get_cluster_location
 from .nodepool import get_all_nodepools_programmatic
 from .resources import (
-    CLUSTER_METADATA_CONFIGMAP,
-    CLUSTER_RESOURCES_CONFIGMAP,
     AutoprovisioningConfig,
+    ConfigMapType,
     get_cluster_configmap,
 )
 from .scheduling import get_total_chips_requested_from_args
@@ -266,14 +265,12 @@ def is_autoprovisioning_enabled(
     int of 0 if successful and 1 otherwise.
   """
 
-  resources_configmap_name = f'{args.cluster}-{CLUSTER_RESOURCES_CONFIGMAP}'
-  cluster_config_map = get_cluster_configmap(resources_configmap_name)
+  cluster_config_map = get_cluster_configmap(
+      args.cluster, ConfigMapType.RESOURCES
+  )
 
   if cluster_config_map is None:
-    xpk_print(
-        f'Unable to find config map: {resources_configmap_name}.'
-        ' Autoprovisioning is not enabled.'
-    )
+    xpk_print('Unable to find config map. Autoprovisioning is not enabled.')
     return False, 0
 
   return_code, autoprovisioning_value = get_value_from_map(
@@ -281,8 +278,8 @@ def is_autoprovisioning_enabled(
   )
   if return_code != 0:
     xpk_print(
-        'gke_accelerator type not found in config map:'
-        f' {resources_configmap_name}. Autoprovisioning is not enabled.'
+        'gke_accelerator type not found in config map. Autoprovisioning is not'
+        ' enabled.'
     )
     return False, 0
 
@@ -319,8 +316,9 @@ def get_autoprovisioning_node_selector_args(args) -> tuple[str, int]:
 
   if capacity_type_str == CapacityType.UNKNOWN.name:
     # Use default settings from cluster creation.
-    metadata_configmap_name = f'{args.cluster}-{CLUSTER_METADATA_CONFIGMAP}'
-    cluster_config_map = get_cluster_configmap(metadata_configmap_name)
+    cluster_config_map = get_cluster_configmap(
+        args.cluster, ConfigMapType.METADATA
+    )
 
     # Error out if the metadata config map doesn't exist, and is attempting to use
     # autoprovisioning.
@@ -363,8 +361,9 @@ def get_autoprovisioning_node_selector_args(args) -> tuple[str, int]:
 
 
 def get_cluster_provisioner(args) -> str:
-  metadata_configmap_name = f'{args.cluster}-{CLUSTER_METADATA_CONFIGMAP}'
-  cluster_config_map = get_cluster_configmap(metadata_configmap_name)
+  cluster_config_map = get_cluster_configmap(
+      args.cluster, ConfigMapType.METADATA
+  )
   cluster_provisioner = 'gcloud'
   if not cluster_config_map is None:
     provisioner = cluster_config_map.get('provisioner')