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

xpk/core/config_test.py

@@ -14,7 +14,9 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-from xpk.core.config import XpkConfig, CFG_BUCKET_KEY, CLUSTER_NAME_KEY, PROJECT_KEY, ZONE_KEY
+from xpk.core.config import FileSystemConfig, InMemoryXpkConfig, CFG_BUCKET_KEY, CLUSTER_NAME_KEY, PROJECT_KEY, ZONE_KEY, _get_version
+from unittest.mock import patch
+from importlib.metadata import PackageNotFoundError
 
 import os
 import pytest
@@ -31,15 +33,60 @@ def _():
     os.remove(config_tmp_path)
 
 
-def test_config(_):
-  cfg = XpkConfig(config_tmp_path)
+@patch('os.getenv', return_value='10.0.0')
+def test_get_version_returns_overriden_value_when_it_is_overriden(_):
+  assert _get_version() == '10.0.0'
+
+
+@patch('os.getenv', return_value='')
+@patch('xpk.core.config.setuptools_get_version', return_value='10.0.0')
+def test_get_version_returns_value_from_setuptools_scm_when_there_is_no_override(
+    *_,
+):
+  assert _get_version() == '10.0.0'
+
+
+@patch('os.getenv', return_value='')
+@patch(
+    'xpk.core.config.setuptools_get_version',
+    side_effect=LookupError('unable to find git version'),
+)
+@patch('xpk.core.config.version', return_value='10.0.0')
+def test_get_version_returns_value_from_pip_when_there_is_no_setuptools_could_be_resolved(
+    *_,
+):
+  assert _get_version() == '10.0.0'
+
+
+@patch('os.getenv', return_value='')
+@patch(
+    'xpk.core.config.setuptools_get_version',
+    side_effect=LookupError('unable to find git version'),
+)
+@patch(
+    'xpk.core.config.version',
+    side_effect=PackageNotFoundError('unable to locate package'),
+)
+def test_get_version_returns_none_when_no_version_could_be_resolved(*_):
+  with pytest.raises(LookupError):
+    _get_version()
+
+
+@pytest.mark.parametrize(
+    argnames='cfg',
+    argvalues=[(FileSystemConfig(config_tmp_path)), (InMemoryXpkConfig())],
+)
+def test_config(_, cfg):
   cfg.set('project-id', 'foo')
   project_id = cfg.get('project-id')
   assert project_id == 'foo'
 
 
-def test_config_get_all(_):
-  cfg = XpkConfig(config_tmp_path)
+@pytest.mark.parametrize(
+    argnames='cfg',
+    argvalues=[(FileSystemConfig(config_tmp_path)), (InMemoryXpkConfig())],
+)
+def test_config_get_all(_, cfg):
   cfg.set(PROJECT_KEY, 'foo')
   cfg.set(CLUSTER_NAME_KEY, 'bar')
   cfg.set(ZONE_KEY, 'europe-west1-a')
@@ -52,20 +99,29 @@ def test_config_get_all(_):
   assert cfg_all[CFG_BUCKET_KEY] == 'cfg-bucket'
 
 
-def test_config_get_empty(_):
-  cfg = XpkConfig(config_tmp_path)
+@pytest.mark.parametrize(
+    argnames='cfg',
+    argvalues=[(FileSystemConfig(config_tmp_path)), (InMemoryXpkConfig())],
+)
+def test_config_get_empty(_, cfg):
   val = cfg.get(PROJECT_KEY)
   assert val is None
 
 
-def test_config_get_all_empty(_):
-  cfg = XpkConfig(config_tmp_path)
+@pytest.mark.parametrize(
+    argnames='cfg',
+    argvalues=[(FileSystemConfig(config_tmp_path)), (InMemoryXpkConfig())],
+)
+def test_config_get_all_empty(_, cfg):
   val = cfg.get_all()
   assert not val
 
 
-def test_config_set_incorrect(_):
-  cfg = XpkConfig(config_tmp_path)
+@pytest.mark.parametrize(
+    argnames='cfg',
+    argvalues=[(FileSystemConfig(config_tmp_path)), (InMemoryXpkConfig())],
+)
+def test_config_set_incorrect(cfg, _):
   cfg.set('foo', 'bar')
   cfg_all = cfg.get_all()
   assert not cfg_all

xpk/core/docker_container.py

@@ -182,7 +182,9 @@ def get_user_workload_container(args, system: SystemCharacteristics):
       debugging_dashboard_id: id of the GKE dashboard
   """
 
-  setup_docker_image_code, docker_image = setup_docker_image(args)
+  setup_docker_image_code, docker_image = setup_docker_image(
+      args, system.docker_platform
+  )
   if setup_docker_image_code != 0:
     xpk_exit(setup_docker_image_code)
 

xpk/core/docker_image.py

@@ -19,6 +19,7 @@ import os
 import random
 import string
 
+from .system_characteristics import DockerPlatform
 from ..utils.console import xpk_exit, xpk_print
 from ..utils.file import write_tmp_file
 from ..utils.execution_context import is_dry_run
@@ -26,7 +27,6 @@ from .commands import run_command_with_updates
 
 DEFAULT_DOCKER_IMAGE = 'python:3.10'
 DEFAULT_SCRIPT_DIR = os.getcwd()
-PLATFORM = 'linux/amd64'
 
 
 def validate_docker_image(docker_image, args) -> int:
@@ -63,7 +63,9 @@ def validate_docker_image(docker_image, args) -> int:
     return 0
 
 
-def build_docker_image_from_base_image(args, verbose=True) -> tuple[int, str]:
+def build_docker_image_from_base_image(
+    args, docker_platform: DockerPlatform, verbose=True
+) -> tuple[int, str]:
   """Adds script dir to the base docker image and uploads the image.
 
   Args:
@@ -97,8 +99,8 @@ def build_docker_image_from_base_image(args, verbose=True) -> tuple[int, str]:
   )
   tmp = write_tmp_file(docker_file)
   docker_build_command = (
-      f'docker buildx build --platform={PLATFORM} -f {str(tmp)} -t'
-      f' {docker_name} {args.script_dir}'
+      f'docker buildx build --platform={docker_platform.value} -f'
+      f' {str(tmp)} -t {docker_name} {args.script_dir}'
   )
   xpk_print(f'Building {args.script_dir} into docker image.')
   return_code = run_command_with_updates(
@@ -158,7 +160,9 @@ def build_docker_image_from_base_image(args, verbose=True) -> tuple[int, str]:
   return return_code, cloud_docker_image
 
 
-def setup_docker_image(args) -> tuple[int, str]:
+def setup_docker_image(
+    args, docker_platform: DockerPlatform
+) -> tuple[int, str]:
   """Does steps to verify docker args, check image, and build image (if asked).
 
   Args:
@@ -177,7 +181,7 @@ def setup_docker_image(args) -> tuple[int, str]:
     if validate_docker_image_code != 0:
       xpk_exit(validate_docker_image_code)
     build_docker_image_code, docker_image = build_docker_image_from_base_image(
-        args
+        args, docker_platform
     )
     if build_docker_image_code != 0:
       xpk_exit(build_docker_image_code)

xpk/core/docker_resources.py

@@ -16,7 +16,6 @@ limitations under the License.
 
 import os
 import re
-from .capacity import H100_DEVICE_TYPE, H100_MEGA_DEVICE_TYPE, H200_DEVICE_TYPE
 from .cluster import setup_k8s_env
 from .storage import GCS_FUSE_TYPE, GCP_FILESTORE_TYPE, PARALLELSTORE_TYPE, GCE_PD_TYPE, LUSTRE_TYPE, Storage, get_storages_to_mount
 from .system_characteristics import AcceleratorType, SystemCharacteristics
@@ -109,14 +108,6 @@ def get_gpu_env(args, system) -> str:
                     value: "{args.command}"
                   {custom_envs}"""
 
-  gpu_direct_name = 'fastrak'
-  if args.device_type == H100_DEVICE_TYPE:
-    gpu_direct_name = 'tcpx'
-  elif args.device_type == H100_MEGA_DEVICE_TYPE:
-    gpu_direct_name = 'tcpxo'
-  elif args.device_type == H200_DEVICE_TYPE:
-    gpu_direct_name = 'rdma'
-
   gpu_env_dic = {
       'JAX_COORDINATOR_PORT': '6002',
       'JAX_COORDINATOR_ADDRESS': (
@@ -129,7 +120,7 @@ def get_gpu_env(args, system) -> str:
   return gpu_env_yaml.format(
       args=args,
       chips_per_vm=system.chips_per_vm,
-      gpu_direct_name=gpu_direct_name,
+      gpu_direct_name=system.gpu_config.gpu_direct_name,
       custom_envs=format_env_dict(args.env, system),
   )
 

xpk/core/kjob.py

@@ -25,7 +25,6 @@ from kubernetes.client.rest import ApiException
 from ..utils import templates
 from ..utils.execution_context import is_dry_run
 from ..utils.console import xpk_exit, xpk_print
-from .capacity import H100_DEVICE_TYPE, H100_MEGA_DEVICE_TYPE, H200_DEVICE_TYPE
 from .cluster import DEFAULT_NAMESPACE, XPK_SA, setup_k8s_env
 from .commands import (
     run_command_for_value,
@@ -38,7 +37,7 @@ from .config import (
     KJOB_SHELL_IMAGE,
     KJOB_SHELL_INTERACTIVE_COMMAND,
     KJOB_SHELL_WORKING_DIRECTORY,
-    xpk_config,
+    get_config,
 )
 from .network import get_cluster_subnetworks
 from .system_characteristics import AcceleratorType, SystemCharacteristics
@@ -52,7 +51,6 @@ from .storage import (
 )
 from .workload_decorators import (
     rdma_decorator,
-    tcpx_decorator,
     tcpxo_decorator,
 )
 from .workload_decorators.tcpxo_decorator import get_tcpxo_deamon_entry
@@ -234,7 +232,7 @@ def get_pod_template_interactive_command() -> str:
   Returns:
     str - PodTemplate's interactive command
   """
-  pod_command = xpk_config.get(KJOB_SHELL_INTERACTIVE_COMMAND)
+  pod_command = get_config().get(KJOB_SHELL_INTERACTIVE_COMMAND)
   if pod_command is None or len(pod_command) == 0:
     pod_command = PodTemplateDefaults.INTERACTIVE_COMMAND.value
 
@@ -260,14 +258,17 @@ def create_app_profile_instance(volume_bundles: list[str]) -> int:
   )
 
 
-def decorate_job_template_with_gpu(yml_string: str, gpu_type: str) -> str:
+def decorate_job_template_with_gpu(
+    yml_string: str, system: SystemCharacteristics
+) -> str:
   job_spec = yaml.safe_load(yml_string)["template"]
-  if gpu_type == H100_DEVICE_TYPE:
-    job_spec = tcpx_decorator.decorate_kjob_template(job_spec)
-  if gpu_type == H100_MEGA_DEVICE_TYPE:
-    job_spec = tcpxo_decorator.decorate_kjob_template(job_spec)
-  if gpu_type == H200_DEVICE_TYPE:
-    job_spec = rdma_decorator.decorate_kjob_template(job_spec)
+  kjob_decorator = (
+      system.gpu_config.kjob_decorator_fn
+      if system.gpu_config and system.gpu_config.kjob_decorator_fn
+      else None
+  )
+  if kjob_decorator:
+    job_spec = kjob_decorator(job_spec)
   job_template_dict = yaml.safe_load(yml_string)
   job_template_dict["template"] = job_spec
   yaml_result: str = yaml.dump(job_template_dict, sort_keys=False)
@@ -286,10 +287,10 @@ def create_job_template_instance(
   Returns:
     exit_code > 0 if creating JobTemplate fails, 0 otherwise
   """
-  job_image = xpk_config.get(KJOB_BATCH_IMAGE)
+  job_image = get_config().get(KJOB_BATCH_IMAGE)
   if job_image is None or len(job_image) == 0:
     job_image = JobTemplateDefaults.IMAGE.value
-  working_directory = xpk_config.get(KJOB_BATCH_WORKING_DIRECTORY)
+  working_directory = get_config().get(KJOB_BATCH_WORKING_DIRECTORY)
   if working_directory is None or len(working_directory) == 0:
     working_directory = JobTemplateDefaults.WORKING_DIRECTORY.value
   resources = (
@@ -316,7 +317,7 @@ def create_job_template_instance(
       service_account=service_account,
   )
   if system is not None and system.accelerator_type == AcceleratorType.GPU:
-    yml_string = decorate_job_template_with_gpu(yml_string, system.device_type)
+    yml_string = decorate_job_template_with_gpu(yml_string, system)
 
   return run_kubectl_apply(
       yml_string,
@@ -330,10 +331,10 @@ def create_pod_template_instance(service_account: str) -> int:
   Returns:
     exit_code > 0 if creating PodTemplate fails, 0 otherwise
   """
-  pod_image = xpk_config.get(KJOB_SHELL_IMAGE)
+  pod_image = get_config().get(KJOB_SHELL_IMAGE)
   if pod_image is None or len(pod_image) == 0:
     pod_image = PodTemplateDefaults.IMAGE.value
-  working_directory = xpk_config.get(KJOB_SHELL_WORKING_DIRECTORY)
+  working_directory = get_config().get(KJOB_SHELL_WORKING_DIRECTORY)
   if working_directory is None or len(working_directory) == 0:
     working_directory = PodTemplateDefaults.WORKING_DIRECTORY.value
 

xpk/core/kueue_manager.py

@@ -24,15 +24,13 @@ from jinja2 import Environment, FileSystemLoader
 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 .capacity import B200_DEVICE_TYPE, H100_MEGA_DEVICE_TYPE, H200_DEVICE_TYPE
-from .scheduling import (
-    create_accelerator_label,
-    create_machine_label,
-)
 from .system_characteristics import (
     SUB_SLICING_TOPOLOGIES,
+    AcceleratorType,
     AcceleratorTypeToAcceleratorCharacteristics,
     SystemCharacteristics,
+    create_accelerator_label,
+    create_machine_label,
 )
 from ..core.commands import (
     run_command_for_value,
@@ -321,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 ""
 
@@ -400,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(

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, UserFacingNameToSystemCharacteristics
+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(
@@ -405,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),
       ),
   )
 
@@ -501,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')

xpk/core/nodepool.py

@@ -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
 
@@ -247,20 +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 is_placement_policy_supported(system):
     placement_policy = get_placement_policy_name(system)
-    ensure_resource_policy_exists(placement_policy, args, system.topology)
+    ensure_resource_policy_exists(
+        resource_policy_name=placement_policy,
+        project=args.project,
+        zone=args.zone,
+        topology=system.topology,
+    )
     placement_args = f' --placement-policy={placement_policy}'
 
   create_commands = []
@@ -311,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):
@@ -587,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',
   )
@@ -605,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
@@ -179,6 +194,8 @@ def test_placement_policy_created_for_gpu_with_valid_topology(
       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")
@@ -209,6 +226,8 @@ def test_placement_policy_not_created_for_gpu_with_invalid_topology(
       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")
@@ -242,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")
@@ -274,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()