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

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,
@@ -28,10 +28,9 @@ from .capacity import (
 from .commands import run_command_for_value, run_commands
 from .gcloud_context import GkeServerConfig, get_cluster_location, zone_to_region
 from .resources import (
-    CLUSTER_CONFIGMAP_YAML,
-    CLUSTER_RESOURCES_CONFIGMAP,
+    ConfigMapType,
     check_cluster_resources,
-    create_or_update_cluster_configmap,
+    update_cluster_configmap,
 )
 from .system_characteristics import AcceleratorType
 
@@ -110,6 +109,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 +186,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 +211,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}'
@@ -253,22 +246,23 @@ def run_gke_node_pool_create_command(
         )
       else:
         resources_data = f'{device_type}: "0"'
-      resources_configmap_name = f'{args.cluster}-{CLUSTER_RESOURCES_CONFIGMAP}'
-      resources_yml = CLUSTER_CONFIGMAP_YAML.format(
-          args=args, name=resources_configmap_name, data=resources_data
+      return_code = update_cluster_configmap(
+          cluster_name=args.cluster,
+          config_map_type=ConfigMapType.RESOURCES,
+          data=resources_data,
       )
-      configmap_yml = {}
-      configmap_yml[resources_configmap_name] = resources_yml
-      return_code = create_or_update_cluster_configmap(configmap_yml)
       if return_code != 0:
         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(
+        resource_policy_name=placement_policy,
+        project=args.project,
+        zone=args.zone,
+        topology=system.topology,
     )
-    ensure_resource_policy_exists(placement_policy, args, system.topology)
     placement_args = f' --placement-policy={placement_policy}'
 
   create_commands = []
@@ -290,16 +284,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
@@ -319,7 +313,7 @@ def run_gke_node_pool_create_command(
       command += (
           ' --accelerator'
           f' type={system.gke_accelerator},count={str(system.chips_per_vm)},gpu-driver-version=latest'
-          f' --no-enable-autoupgrade --scopes={CLOUD_PLATFORM_AUTH_SCOPE_URL}'
+          f' --scopes={CLOUD_PLATFORM_AUTH_SCOPE_URL}'
       )
       if device_type == H100_MEGA_DEVICE_TYPE:
         for i in range(1, 9):
@@ -595,14 +589,14 @@ def get_desired_node_pool_names(
 
 
 def ensure_resource_policy_exists(
-    resource_policy_name: str, args, topology: str
+    resource_policy_name: str, project: str, zone: str, topology: str
 ) -> None:
   return_code, _ = run_command_for_value(
       (
           'gcloud compute resource-policies describe'
           f' {resource_policy_name} '
-          f'--project={args.project} '
-          f'--region={zone_to_region(args.zone)}'
+          f'--project={project} '
+          f'--region={zone_to_region(zone)}'
       ),
       'Retrieve resource policy',
   )
@@ -613,7 +607,7 @@ def ensure_resource_policy_exists(
   return_code, _ = run_command_for_value(
       (
           'gcloud compute resource-policies create workload-policy'
-          f' {resource_policy_name} --project={args.project} --region={zone_to_region(args.zone)} --type=HIGH_THROUGHPUT'
+          f' {resource_policy_name} --project={project} --region={zone_to_region(zone)} --type=HIGH_THROUGHPUT'
           f' --accelerator-topology={topology}'
       ),
       'Create resource policy',

xpk/core/nodepool_test.py

@@ -20,7 +20,7 @@ from xpk.core.nodepool import (
     get_desired_node_pool_names,
     run_gke_node_pool_create_command,
 )
-from xpk.core.system_characteristics import AcceleratorType, SystemCharacteristics
+from xpk.core.system_characteristics import AcceleratorType, SystemCharacteristics, DockerPlatform, GpuConfig
 
 CLUSTER_NAME = "running-cucumber"
 
@@ -96,7 +96,12 @@ def test_ensure_resource_policy_exists_with_existing_policy_retrieves_existing_p
   mock = mocker.patch(
       "xpk.core.nodepool.run_command_for_value", return_value=(0, "")
   )
-  ensure_resource_policy_exists("resource-policy", args, "2x2x1")
+  ensure_resource_policy_exists(
+      resource_policy_name="resource-policy",
+      project="test-project",
+      zone="us-central1-a",
+      topology="2x2x1",
+  )
   mock.assert_called_once()
 
 
@@ -108,7 +113,12 @@ def test_ensure_resource_policy_exists_without_existing_policy_creates_policy(
   mock = mocker.patch(
       "xpk.core.nodepool.run_command_for_value", side_effect=[(1, ""), (0, "")]
   )
-  ensure_resource_policy_exists("resource-policy", args, "2x2x1")
+  ensure_resource_policy_exists(
+      resource_policy_name="resource-policy",
+      project="test-project",
+      zone="us-central1-a",
+      topology="2x2x1",
+  )
   assert mock.call_count == 2
   assert mock.call_args_list[0].args[1] == "Retrieve resource policy"
 
@@ -125,7 +135,12 @@ def test_ensure_resource_policy_exits_without_existing_policy_throws_when_creati
         "xpk.core.nodepool.run_command_for_value",
         side_effect=[(1, ""), (1, "")],
     )
-    ensure_resource_policy_exists("resource-policy", args, "2x2x1")
+    ensure_resource_policy_exists(
+        resource_policy_name="resource-policy",
+        project="test-project",
+        zone="us-central1-a",
+        topology="2x2x1",
+    )
 
 
 @pytest.fixture
@@ -145,22 +160,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,13 +187,15 @@ 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,
       accelerator_type=AcceleratorType.GPU,
       device_type="h100-80gb-8",
       supports_sub_slicing=False,
+      docker_platform=DockerPlatform.ARM,
+      gpu_config=GpuConfig(requires_topology=True),
   )
 
   run_gke_node_pool_create_command(args, system, "1.2.3")
@@ -188,10 +207,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,13 +219,15 @@ 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,
       accelerator_type=AcceleratorType.GPU,
       device_type="h100-80gb-8",
       supports_sub_slicing=False,
+      docker_platform=DockerPlatform.ARM,
+      gpu_config=GpuConfig(requires_topology=True),
   )
 
   run_gke_node_pool_create_command(args, system, "1.2.3")
@@ -218,10 +239,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 +253,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,
@@ -240,6 +261,7 @@ def test_placement_policy_created_for_tpu7x_with_valid_topology(
       device_type="tpu7x-8",
       requires_workload_policy=True,
       supports_sub_slicing=False,
+      docker_platform=DockerPlatform.ARM,
   )
 
   run_gke_node_pool_create_command(args, system, "1.2.3")
@@ -251,14 +273,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",
@@ -272,6 +294,7 @@ def test_placement_policy_not_created_for_non7x_tpu(
       accelerator_type=AcceleratorType.TPU,
       device_type="v6e-4",
       supports_sub_slicing=True,
+      docker_platform=DockerPlatform.ARM,
   )
 
   run_gke_node_pool_create_command(args, system, "1.2.3")

xpk/core/pathways.py

@@ -333,3 +333,26 @@ def try_to_delete_pathwaysjob_first(args, workloads) -> bool:
     xpk_print(f'Delete Workload request returned ERROR {return_code}')
     return False
   return True
+
+
+def get_pathways_machine_types(
+    project: str, zone: str
+) -> tuple[int, list[str]]:
+  # Identify machine types with sufficient allocatable capacity to
+  # schedule the Pathways pod. This filter ensures the selected node
+  # is large enough to handle the control plane workload plus GKE
+  # system overhead.
+  min_memory_mb = 233 * 1024
+  command = (
+      'gcloud compute machine-types list --filter "guestCpus >= 49 AND memoryMb'
+      f' >= {min_memory_mb} AND zone = \'{zone}\'" --format="value(name)"'
+      f' --project={project}'
+  )
+  return_code, result = run_command_for_value(
+      command=command,
+      task='Retrieve available pathways machine types',
+      dry_run_return_val='n2-standard-64',
+  )
+  if return_code != 0:
+    return return_code, []
+  return 0, result.strip().splitlines()

xpk/core/pathways_test.py

@@ -0,0 +1,57 @@
+"""
+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.
+"""
+
+import pytest
+from unittest.mock import MagicMock
+from xpk.core.testing.commands_tester import CommandsTester
+from .pathways import get_pathways_machine_types
+
+
+@pytest.fixture(autouse=True)
+def commands_tester(mocker: MagicMock):
+  return CommandsTester(
+      mocker,
+      run_command_with_updates_path=(
+          "xpk.core.pathways.run_command_with_updates"
+      ),
+      run_command_for_value_path="xpk.core.pathways.run_command_for_value",
+  )
+
+
+def test_get_pathways_machine_types_when_command_fails_returns_failed_exit_code(
+    commands_tester: CommandsTester,
+):
+  commands_tester.set_result_for_command(
+      (1, ""), "gcloud compute machine-types list"
+  )
+  return_code, machine_types = get_pathways_machine_types(
+      project="gke-project", zone="us-central1-a"
+  )
+  assert return_code == 1
+  assert machine_types == []
+
+
+def test_get_pathways_machine_types_when_command_suceeds_returns_machine_types(
+    commands_tester: CommandsTester,
+):
+  commands_tester.set_result_for_command(
+      (0, "abc\ncba"), "gcloud compute machine-types list"
+  )
+  return_code, machine_types = get_pathways_machine_types(
+      project="gke-project", zone="us-central1-a"
+  )
+  assert return_code == 0
+  assert machine_types == ["abc", "cba"]

xpk/core/resources.py

@@ -15,6 +15,7 @@ limitations under the License.
 """
 
 from dataclasses import dataclass
+import os
 
 from ..utils.console import xpk_print
 from ..utils.file import write_tmp_file
@@ -30,9 +31,13 @@ from .capacity import (
 from .commands import run_command_for_value, run_commands
 from .config import XPK_CURRENT_VERSION
 from .system_characteristics import AcceleratorType, get_system_characteristics_by_device_type, SystemCharacteristics
+from enum import Enum
+
+
+class ConfigMapType(Enum):
+  RESOURCES = 'resources-configmap'
+  METADATA = 'metadata-configmap'
 
-CLUSTER_RESOURCES_CONFIGMAP = 'resources-configmap'
-CLUSTER_METADATA_CONFIGMAP = 'metadata-configmap'
 
 CLUSTER_CONFIGMAP_YAML = """kind: ConfigMap
 apiVersion: v1
@@ -50,7 +55,15 @@ class AutoprovisioningConfig:
   maximum_chips: int
 
 
-def get_cluster_configmap(configmap_name) -> dict[str, str] | None:
+def get_config_map_name(
+    cluster_name: str, config_map_type: ConfigMapType
+) -> str:
+  return f'{cluster_name}-{config_map_type.value}'
+
+
+def get_cluster_configmap(
+    cluster_name: str, config_map_type: ConfigMapType
+) -> dict[str, str] | None:
   """Run the Get GKE Cluster ConfigMap request.
 
   Args:
@@ -59,15 +72,17 @@ def get_cluster_configmap(configmap_name) -> dict[str, str] | None:
   Returns:
     key:value pairs stored in cluster ConfigMap.
   """
+  config_map_name = get_config_map_name(cluster_name, config_map_type)
   command = (
       'kubectl get configmap'
-      f' {configmap_name} -o=custom-columns="ConfigData:data" --no-headers=true'
+      f' {config_map_name} -o=custom-columns="ConfigData:data"'
+      ' --no-headers=true'
   )
 
   return_code, return_value = run_command_for_value(
       command,
       'GKE Cluster Get ConfigMap',
-      dry_run_return_val='map[]',
+      dry_run_return_val=_get_dry_run_config_map_value(config_map_type),
   )
   if return_code != 0:
     xpk_print(f'GKE Cluster Get ConfigMap request returned ERROR {return_code}')
@@ -89,9 +104,18 @@ def get_cluster_configmap(configmap_name) -> dict[str, str] | None:
   return config_map
 
 
+def _get_dry_run_config_map_value(config_map_type: ConfigMapType) -> str:
+  default_value = 'map[]'
+
+  if config_map_type == ConfigMapType.RESOURCES:
+    return os.getenv('DRY_RUN_RESOURCES_CONFIG_MAP', default_value)
+
+  return default_value
+
+
 def create_cluster_configmaps(
     args,
-    system,
+    system: SystemCharacteristics,
     tensorboard_config: dict,
     autoprovisioning_config: AutoprovisioningConfig | None,
 ) -> int:
@@ -127,9 +151,11 @@ def create_cluster_configmaps(
     resources_data = (
         f'{device_type}: "{int(args.num_slices) * system.vms_per_slice}"'
     )
-  resources_configmap_name = f'{args.cluster}-{CLUSTER_RESOURCES_CONFIGMAP}'
+  resources_configmap_name = get_config_map_name(
+      args.cluster, ConfigMapType.RESOURCES
+  )
   resources_yml = CLUSTER_CONFIGMAP_YAML.format(
-      args=args, name=resources_configmap_name, data=resources_data
+      name=resources_configmap_name, data=resources_data
   )
   configmap_yml[resources_configmap_name] = resources_yml
 
@@ -148,15 +174,17 @@ def create_cluster_configmaps(
   # Reservation ID if applicable.
   if capacity_type == CapacityType.RESERVATION:
     metadata += f'\n  {RESERVATION_CONFIG_KEY}: {args.reservation}'
-  metadata_configmap_name = f'{args.cluster}-{CLUSTER_METADATA_CONFIGMAP}'
+  metadata_configmap_name = get_config_map_name(
+      args.cluster, ConfigMapType.METADATA
+  )
   metadata_yml = CLUSTER_CONFIGMAP_YAML.format(
-      args=args, name=metadata_configmap_name, data=metadata
+      name=metadata_configmap_name, data=metadata
   )
   configmap_yml[metadata_configmap_name] = metadata_yml
-  return create_or_update_cluster_configmap(configmap_yml)
+  return _create_or_update_cluster_configmap(configmap_yml)
 
 
-def create_or_update_cluster_configmap(configmap_yml: dict) -> int:
+def _create_or_update_cluster_configmap(configmap_yml: dict[str, str]) -> int:
   """
   Args:
     configmap_yml: dict containing ConfigMap name and yml string.
@@ -187,7 +215,18 @@ def create_or_update_cluster_configmap(configmap_yml: dict) -> int:
   return 0
 
 
-def check_cluster_resources(args, system) -> tuple[bool, bool]:
+def update_cluster_configmap(
+    cluster_name: str, config_map_type: ConfigMapType, data: str
+) -> int:
+  config_map_name = get_config_map_name(cluster_name, config_map_type)
+  yaml = CLUSTER_CONFIGMAP_YAML.format(name=config_map_name, data=data)
+  config_map_dict = {config_map_name: yaml}
+  return _create_or_update_cluster_configmap(config_map_dict)
+
+
+def check_cluster_resources(
+    args, system: SystemCharacteristics
+) -> tuple[bool, bool]:
   """Check if cluster has resources of a specified device_type/gke_accelerator.
   This check will be skipped if <args.cluster>-<_CLUSTER_RESOURCES_CONFIGMAP> ConfigMap doesn't exist for the cluster.
 
@@ -200,8 +239,9 @@ def check_cluster_resources(args, system) -> tuple[bool, bool]:
     True if resources in the cluster should be checked, False otherwise.
     True if device_type/gke_accelerator exists in the cluster, False otherwise.
   """
-  resources_configmap_name = f'{args.cluster}-{CLUSTER_RESOURCES_CONFIGMAP}'
-  resources_config_map = get_cluster_configmap(resources_configmap_name)
+  resources_config_map = get_cluster_configmap(
+      args.cluster, ConfigMapType.RESOURCES
+  )
   if resources_config_map is None:
     xpk_print(
         f'No ConfigMap exist for cluster with the name {resources_config_map}.'
@@ -216,20 +256,35 @@ def check_cluster_resources(args, system) -> tuple[bool, bool]:
 
 
 def get_cluster_system_characteristics(args) -> SystemCharacteristics | None:
-  """Get systemCharcteristics based on the cluster resources configMap
+  """Get SystemCharcteristics based on the cluster resources configMap.
+
   Args:
     args: user provided arguments for running the command.
 
   Returns:
-    returns system characteristics
+    returns system characteristics, or None if not found.
+  """
+  resources_config_map = get_cluster_configmap(
+      args.cluster, ConfigMapType.RESOURCES
+  )
+  return get_cluster_system_characteristics_from_config_map(
+      resources_config_map
+  )
+
+
+def get_cluster_system_characteristics_from_config_map(
+    resources_config_map: dict[str, str] | None,
+) -> SystemCharacteristics | None:
+  """Get SystemCharcteristics based on the cluster resources configMap.
+
+  Returns:
+    returns system characteristics, or None if not found.
   """
-  resources_configmap_name = f'{args.cluster}-{CLUSTER_RESOURCES_CONFIGMAP}'
-  cluster_config_map = get_cluster_configmap(resources_configmap_name)
 
-  if cluster_config_map is None:
+  if resources_config_map is None:
     return None
 
-  for key in cluster_config_map:
+  for key in resources_config_map:
     system, result_code = get_system_characteristics_by_device_type(key)
     if result_code == 0:
       return system
@@ -238,20 +293,22 @@ def get_cluster_system_characteristics(args) -> SystemCharacteristics | None:
 
 
 def get_cluster_capacity_type(args) -> CapacityType | None:
-  """Get systemCharcteristics based on the cluster resources configMap
+  """Get CapacityType based on the cluster metadata configMap.
+
   Args:
     args: user provided arguments for running the command.
 
   Returns:
-    returns system characteristics
+    returns CapacityType, or None if not found.
   """
-  metadata_configmap_name = f'{args.cluster}-{CLUSTER_METADATA_CONFIGMAP}'
-  cluster_config_map = get_cluster_configmap(metadata_configmap_name)
+  metadata_configmap_name = get_cluster_configmap(
+      args.cluster, ConfigMapType.METADATA
+  )
 
-  if cluster_config_map is None:
+  if metadata_configmap_name is None:
     return None
 
-  capacityValue = cluster_config_map.get('capacity_type')
+  capacityValue = metadata_configmap_name.get('capacity_type')
   if capacityValue is not None:
     return CapacityType[capacityValue.upper()]