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

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()]
 

xpk/core/scheduling.py

@@ -14,61 +14,63 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+from enum import Enum
+
+from .kueue_manager import get_installed_kueue_version, has_sub_slicing_enabled
+from ..utils.feature_flags import FeatureFlags
 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
 from .system_characteristics import (
+    SUB_SLICING_TOPOLOGIES,
     AcceleratorType,
-    AcceleratorTypeToAcceleratorCharacteristics,
     SystemCharacteristics,
+    create_accelerator_label,
+    create_machine_label,
 )
+from packaging.version import Version
 
+_SUB_SLICING_MINIMUM_KUEUE_VERSION = Version('0.13.0')
 
-def check_if_workload_can_schedule(args, system: SystemCharacteristics) -> bool:
-  """Check if workload can schedule based on the cluster resources (tpu_type and maximum VM in cluster).
 
-  Args:
-    args: user provided arguments for running the command.
-    system: system characteristics
+class WorkloadScheduling(Enum):
+  UNAVAILABLE = 0
+  AVAILABLE = 1
+  SUB_SLICING_AVAILABLE = 2
+
+
+def check_if_workload_can_schedule(
+    args,
+    workload_system: SystemCharacteristics,
+    cluster_system: SystemCharacteristics | None,
+    resources_config_map: dict[str, str] | None,
+) -> WorkloadScheduling:
+  """Check if workload can schedule based on the cluster resources (tpu_type and maximum VM in cluster).
 
   Returns:
-    returns true if workload can schedule, otherwise returns false.
+    returns WorkloadScheduling describing scheduling option.
   """
-  resources_configmap_name = f'{args.cluster}-{CLUSTER_RESOURCES_CONFIGMAP}'
-  cluster_config_map = get_cluster_configmap(resources_configmap_name)
+  if is_dry_run() and not cluster_system:
+    xpk_print('Skipping workload scheduling validation in dry run.')
+    return WorkloadScheduling.AVAILABLE
 
-  # Prevents workload creation failure for existing clusters with no ConfigMap
-  if cluster_config_map is None:
+  if resources_config_map is None:
     xpk_print(
-        'No ConfigMap exist for cluster with the name'
-        f' {resources_configmap_name}.'
+        "Skipping workload scheduling validation, because there's no Resources"
+        ' ConfigMap in the cluster.'
     )
-    return True
-
-  if is_dry_run():
-    return True
+    return WorkloadScheduling.AVAILABLE
 
-  # Check for gke accelerator type:
-  missing_gke_accelerator_type = False
-  if not cluster_config_map.get(system.gke_accelerator):
-    xpk_print(
-        f'GKE Accelerator Type Check: {args.workload} is requesting'
-        f' {system.gke_accelerator} but cluster only contains'
-        f' {cluster_config_map.keys()}. '
-    )
-    missing_gke_accelerator_type = True
-  elif (
-      cluster_config_map[system.gke_accelerator]
-      == AUTOPROVISIONING_CONFIG_VALUE
-  ):
+  if _is_cluster_set_up_for_nap(workload_system, resources_config_map):
     # Run total chip check when in autoprovisioning mode.
     max_chips_in_cluster = int(
-        cluster_config_map[AUTOPROVISIONING_CONFIG_MAXIMUM_KEY]
+        resources_config_map[AUTOPROVISIONING_CONFIG_MAXIMUM_KEY]
+    )
+    num_chips_in_workload = get_total_chips_requested_from_args(
+        args, workload_system
     )
-    num_chips_in_workload = get_total_chips_requested_from_args(args, system)
 
     if num_chips_in_workload > max_chips_in_cluster:
       xpk_print(
@@ -77,44 +79,100 @@ def check_if_workload_can_schedule(args, system: SystemCharacteristics) -> bool:
           '  Resize the cluster to support more chips with'
           ' `xpk cluster create --autoprovisioning-max-chips=X ...`'
       )
-      return False
-    return True
+      return WorkloadScheduling.UNAVAILABLE
+    return WorkloadScheduling.AVAILABLE
+
+  if workload_system.device_type in resources_config_map:
+    if _check_workload_size_fits(
+        args,
+        workload_system,
+        max_vm_in_cluster=int(
+            resources_config_map[workload_system.device_type]
+        ),
+    ):
+      return WorkloadScheduling.AVAILABLE
+    else:
+      return WorkloadScheduling.UNAVAILABLE
+
+  if _check_sub_slicing_availability(
+      workload_system=workload_system, cluster_system=cluster_system
+  ):
+    assert cluster_system
+    if _check_workload_size_fits(
+        args,
+        workload_system,
+        max_vm_in_cluster=int(resources_config_map[cluster_system.device_type]),
+    ):
+      return WorkloadScheduling.SUB_SLICING_AVAILABLE
+    else:
+      return WorkloadScheduling.UNAVAILABLE
+
+  xpk_print(
+      'Workload scheduling validation failed. XPK will not create the workload'
+      f' {args.workload}.'
+  )
+  return WorkloadScheduling.UNAVAILABLE
+
+
+def _is_cluster_set_up_for_nap(
+    workload_system: SystemCharacteristics, resources_config_map: dict[str, str]
+) -> bool:
+  return (
+      resources_config_map.get(workload_system.gke_accelerator, None)
+      == AUTOPROVISIONING_CONFIG_VALUE
+  )
 
-  # Check for device type
-  missing_device_type = False
-  device_type = system.device_type
-  if device_type not in cluster_config_map:
-    xpk_print(
-        f'Device Type Check: {args.workload} is requesting {device_type} but '
-        f'cluster only contains {cluster_config_map.keys()}. '
-    )
-    missing_device_type = True
 
-  if missing_device_type and missing_gke_accelerator_type:
+def _check_workload_size_fits(
+    args,
+    workload_system: SystemCharacteristics,
+    max_vm_in_cluster: int,
+) -> bool:
+  if workload_system.accelerator_type == AcceleratorType.GPU:
+    vm_required_by_workload = args.num_nodes
+  else:
+    vm_required_by_workload = args.num_slices * workload_system.vms_per_slice
+
+  if vm_required_by_workload > max_vm_in_cluster:
     xpk_print(
-        'Both Device Type and GKE Accelerator Type checks failed.'
-        f' XPK will not create the workload {args.workload}.'
+        f'{args.workload} is requesting {args.num_slices} slice/slices of'
+        f' {workload_system.device_type}, which is'
+        f' {vm_required_by_workload} VMs, but the cluster only contains'
+        f' {max_vm_in_cluster} VMs of {workload_system.device_type}. XPK will'
+        ' not create this workload.'
     )
     return False
-  else:
-    # Check if the size of the workload will fit in the cluster.
-    max_vm_in_cluster = int(cluster_config_map[device_type])
-    if system.accelerator_type == AcceleratorType.GPU:
-      vm_required_by_workload = args.num_nodes
-    else:
-      vm_required_by_workload = args.num_slices * system.vms_per_slice
-    if vm_required_by_workload > max_vm_in_cluster:
-      xpk_print(
-          f'{args.workload} is requesting {args.num_slices} slice/slices of'
-          f' {device_type}, which is {vm_required_by_workload} VMs, but the'
-          f' cluster only contains {max_vm_in_cluster} VMs of {device_type}.'
-          ' XPK will not create this workload.'
-      )
-      return False
-
   return True
 
 
+def _check_sub_slicing_availability(
+    workload_system: SystemCharacteristics,
+    cluster_system: SystemCharacteristics | None,
+) -> bool:
+  if (
+      (not FeatureFlags.SUB_SLICING_ENABLED)
+      or (not cluster_system)
+      or (workload_system.gke_accelerator != cluster_system.gke_accelerator)
+      or (not cluster_system.supports_sub_slicing)
+      or (workload_system.topology not in SUB_SLICING_TOPOLOGIES)
+  ):
+    return False
+
+  return_code, sub_slicing_enabled = has_sub_slicing_enabled()
+  if return_code != 0 or not sub_slicing_enabled:
+    return False
+
+  return_code, current_version = get_installed_kueue_version(
+      dry_run_version=Version('0.13')
+  )
+
+  return (
+      return_code == 0
+      and current_version is not None
+      and current_version >= _SUB_SLICING_MINIMUM_KUEUE_VERSION
+  )
+
+
 def get_total_chips_requested_from_args(
     args, system: SystemCharacteristics
 ) -> int:
@@ -135,7 +193,7 @@ def get_total_chips_requested_from_args(
   return int(num_chips)
 
 
-def get_cpu_affinity(accelerator_type) -> str:
+def get_cpu_affinity(accelerator_type: AcceleratorType) -> str:
   """Generate affinity rules for CPU nodepools, so that workload pods are
   not scheduled on the default pool machines.
   Args:
@@ -199,10 +257,8 @@ def get_gpu_scheduler(
               """
     gpu_scheduler = gpu_scheduler_yaml.format(
         scheduler_name=args.scheduler,
-        accelerator_label=create_accelerator_label(
-            system.accelerator_type, system
-        ),
-        machine_label=create_machine_label(system.accelerator_type, system),
+        accelerator_label=create_accelerator_label(system),
+        machine_label=create_machine_label(system),
         node_pool_name=f'{args.cluster}-np-0',
         autoprovisioning_args=autoprovisioning_args,
     )
@@ -217,74 +273,14 @@ def get_gpu_scheduler(
   return gpu_scheduler, return_code
 
 
-def create_accelerator_label(accelerator_type, system) -> str:
-  """Generates accelerator label.
-
-  Args:
-    accelerator_type: type of accelerator.
-    system: system characteristics.
-
-  Returns:
-    The accelerator label.
-  """
-  if accelerator_type == AcceleratorType.CPU:
-    return ''
-  return (
-      f'{AcceleratorTypeToAcceleratorCharacteristics[accelerator_type].accelerator_label}:'
-      f' {system.gke_accelerator}'
-  )
-
-
-def create_tpu_machine_type(accelerator_type, system) -> str:
-  """Generates TPU machine type..
-
-  Args:
-    accelerator_type: type of accelerator.
-    system: system characteristics.
-
-  Returns:
-    The accelerator label.
-  """
-  if accelerator_type == AcceleratorType.TPU:
+def create_tpu_machine_type(system: SystemCharacteristics) -> str:
+  if system.accelerator_type == AcceleratorType.TPU:
     return f'{system.gce_machine_type}'
   return ''
 
 
-def create_machine_label(
-    accelerator_type, system, autoprovisioning_enabled: bool = False
-) -> str:
-  """Generates machine label.
-
-  Args:
-    accelerator_type: type of accelerator.
-    system: system characteristics.
-    autoprovisioning_enabled: describes autoprovisioning enablement.
-
-  Returns:
-    The machine label.
-  """
-  if accelerator_type == AcceleratorType.TPU and not autoprovisioning_enabled:
-    return (
-        f'{AcceleratorTypeToAcceleratorCharacteristics[accelerator_type].machine_label}:'
-        f' {system.topology}'
-    )
-  return ''
-
-
-def create_tpu_topology(
-    accelerator_type, system, autoprovisioning_enabled: bool = False
-) -> str:
-  """Generates TPU topology.
-
-  Args:
-    accelerator_type: type of accelerator.
-    system: system characteristics.
-    autoprovisioning_enabled: describes autoprovisioning enablement.
-
-  Returns:
-    The machine label.
-  """
-  if accelerator_type == AcceleratorType.TPU and not autoprovisioning_enabled:
+def create_tpu_topology(system: SystemCharacteristics) -> str:
+  if system.accelerator_type == AcceleratorType.TPU:
     return f'{system.topology}'
   return ''