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

xpk/commands/workload_test.py

@@ -17,8 +17,9 @@ limitations under the License.
 import dataclasses
 from unittest.mock import MagicMock, patch
 import pytest
-from ..core.system_characteristics import SystemCharacteristics
-from .workload import _validate_sub_slicing_topology
+from ..core.system_characteristics import SystemCharacteristics, AcceleratorType
+from .workload import _validate_sub_slicing_topology, _validate_sub_slicing_availability
+from packaging.version import Version
 
 
 SYSTEM_CHARACTERISTICS = SystemCharacteristics(
@@ -27,7 +28,7 @@ SYSTEM_CHARACTERISTICS = SystemCharacteristics(
     gke_accelerator='nvidia-l4',
     gce_machine_type='g2-standard-12',
     chips_per_vm=1,
-    accelerator_type=1,
+    accelerator_type=AcceleratorType.TPU,
     device_type='l4-1',
     supports_sub_slicing=True,
     requires_workload_policy=False,
@@ -40,7 +41,7 @@ def xpk_print(mocker):
 
 
 def test_validate_sub_slicing_topology_exits_for_unsupported_topology(
-    xpk_print,
+    xpk_print: MagicMock,
 ):
   with pytest.raises(SystemExit):
     _validate_sub_slicing_topology(SYSTEM_CHARACTERISTICS, '2x1')
@@ -50,7 +51,9 @@ def test_validate_sub_slicing_topology_exits_for_unsupported_topology(
   )
 
 
-def test_validate_sub_slicing_topology_exits_for_too_large_topology(xpk_print):
+def test_validate_sub_slicing_topology_exits_for_too_large_topology(
+    xpk_print: MagicMock,
+):
   with pytest.raises(SystemExit):
     _validate_sub_slicing_topology(SYSTEM_CHARACTERISTICS, '16x16')
 
@@ -64,6 +67,86 @@ def test_validate_sub_slicing_topology_does_nothing_for_supported_topology():
   _validate_sub_slicing_topology(SYSTEM_CHARACTERISTICS, '4x4')
 
 
+def test_validate_sub_slicing_availability_exits_when_getting_topologies_fails(
+    xpk_print: MagicMock, mocker
+):
+  mocker.patch(
+      'xpk.commands.workload.has_sub_slicing_enabled',
+      return_value=(1, None),
+  )
+  with pytest.raises(SystemExit):
+    _validate_sub_slicing_availability()
+
+  assert (
+      'Unable to validate sub-slicing support'
+      in xpk_print.mock_calls[0].args[0]
+  )
+
+
+def test_validate_sub_slicing_availability_exits_when_subslicing_topology_is_not_defined(
+    xpk_print: MagicMock, mocker
+):
+  mocker.patch(
+      'xpk.commands.workload.has_sub_slicing_enabled',
+      return_value=(0, False),
+  )
+  with pytest.raises(SystemExit):
+    _validate_sub_slicing_availability()
+
+  assert (
+      'Cluster has not been not set up for Sub-slicing.'
+      in xpk_print.mock_calls[0].args[0]
+  )
+
+
+def test_validate_sub_slicing_availability_exits_when_kueue_version_cannot_be_determined(
+    xpk_print: MagicMock, mocker
+):
+  mocker.patch(
+      'xpk.commands.workload.has_sub_slicing_enabled',
+      return_value=(0, True),
+  )
+  mocker.patch(
+      'xpk.commands.workload.KueueManager.get_installed_kueue_version',
+      return_value=(1, None),
+  )
+  with pytest.raises(SystemExit):
+    _validate_sub_slicing_availability()
+
+  assert 'Unable to validate sub-slicing' in xpk_print.mock_calls[0].args[0]
+
+
+def test_validate_sub_slicing_availability_exits_when_kueue_version_does_not_meet_minimum_requirements(
+    xpk_print: MagicMock, mocker
+):
+  mocker.patch(
+      'xpk.commands.workload.has_sub_slicing_enabled',
+      return_value=(0, True),
+  )
+  mocker.patch(
+      'xpk.commands.workload.KueueManager.get_installed_kueue_version',
+      return_value=(0, Version('0.0.0')),
+  )
+  with pytest.raises(SystemExit):
+    _validate_sub_slicing_availability()
+
+  assert 'The minimal required version is' in xpk_print.mock_calls[0].args[0]
+
+
+def test_validate_sub_slicing_availability_does_nothing_when_cluster_is_correctly_configured_for_subslicing(
+    mocker,
+):
+  mocker.patch(
+      'xpk.commands.workload.has_sub_slicing_enabled',
+      return_value=(0, True),
+  )
+  mocker.patch(
+      'xpk.commands.workload.KueueManager.get_installed_kueue_version',
+      return_value=(0, Version('0.13.0')),
+  )
+  _validate_sub_slicing_availability()
+
+
 @patch('xpk.commands.common.xpk_print')
 def test_validate_sub_slicing_topology_fails_for_unsupported_system(
     common_xpk_print: MagicMock,

xpk/core/blueprint/blueprint_definitions.py

@@ -36,6 +36,21 @@ class DeploymentModule:
   settings: Optional[dict[str, Any]] = None
   use: Optional[list[str]] = None
 
+  def update_settings(self, additionalSettings: dict[str, Any]) -> None:
+    if self.settings is None:
+      self.settings = dict()
+    self.settings.update(additionalSettings)
+
+  def set_setting(self, key: str, value: Any) -> None:
+    if self.settings is None:
+      self.settings = dict()
+    self.settings[key] = value
+
+  def append_use(self, use: str) -> None:
+    if self.use is None:
+      self.use = list()
+    self.use.append(use)
+
 
 @dataclass
 class DeploymentGroup:
@@ -59,6 +74,6 @@ class Blueprint:
   blueprint_name: Optional[str]
   toolkit_modules_url: str
   toolkit_modules_version: str
-  vars: dict[str, str | list[str]] | None
+  vars: dict[str, str | list[str] | dict[str, str]] | None
   terraform_providers: Optional[dict[str, Any]] = None
   validators: Optional[list[Any]] = None

xpk/core/blueprint/blueprint_generator.py

@@ -211,9 +211,9 @@ class BlueprintGenerator:
         outputs=["instructions"],
     )
     if capacity_type == CapacityType.FLEX_START:
-      a3_megagpu_pool_0.settings.update(self.get_dws_flex_start())
+      a3_megagpu_pool_0.update_settings(self.get_dws_flex_start())
     else:
-      a3_megagpu_pool_0.settings.update({"static_node_count": num_nodes})
+      a3_megagpu_pool_0.update_settings({"static_node_count": num_nodes})
 
     set_placement_policy = capacity_type != CapacityType.SPOT
     workload = DeploymentModule(
@@ -252,8 +252,8 @@ class BlueprintGenerator:
 
     print(reservation_placement_policy)
     if reservation_placement_policy is not None:
-      a3_megagpu_pool_0.settings["placement_policy"] = (
-          reservation_placement_policy
+      a3_megagpu_pool_0.set_setting(
+          "placement_policy", reservation_placement_policy
       )
 
     primary_group = DeploymentGroup(
@@ -268,7 +268,7 @@ class BlueprintGenerator:
         ],
     )
     if set_placement_policy and reservation_placement_policy is None:
-      a3_megagpu_pool_0.use.append(group_placement_0.id)
+      a3_megagpu_pool_0.append_use(group_placement_0.id)
       primary_group.modules.append(group_placement_0)
     a3_mega_blueprint = Blueprint(
         terraform_backend_defaults=self._getblock_terraform_backend(
@@ -580,9 +580,9 @@ class BlueprintGenerator:
         outputs=["instructions"],
     )
     if capacity_type == CapacityType.FLEX_START:
-      gpu_pool.settings.update(self.get_dws_flex_start())
+      gpu_pool.update_settings(self.get_dws_flex_start())
     else:
-      gpu_pool.settings.update({"static_node_count": num_nodes})
+      gpu_pool.update_settings({"static_node_count": num_nodes})
 
     workload_manager_install_id = "workload-manager-install"
     workload_manager_install = DeploymentModule(
@@ -855,9 +855,9 @@ class BlueprintGenerator:
         outputs=["instructions"],
     )
     if capacity_type == CapacityType.FLEX_START:
-      gpu_pool.settings.update(self.get_dws_flex_start())
+      gpu_pool.update_settings(self.get_dws_flex_start())
     else:
-      gpu_pool.settings.update({"static_node_count": num_nodes})
+      gpu_pool.update_settings({"static_node_count": num_nodes})
 
     workload_manager_install_id = "workload-manager-install"
     workload_manager_install = DeploymentModule(
@@ -956,7 +956,7 @@ class BlueprintGenerator:
     )
 
   def _getblock_terraform_backend(
-      self, gcs_bucket: str, cluster_name: str, prefix: str = ""
+      self, gcs_bucket: str | None, cluster_name: str, prefix: str = ""
   ) -> dict | None:
     if gcs_bucket is None:
       return None
@@ -986,7 +986,7 @@ class BlueprintGenerator:
       yaml_parser.dump(xpk_blueprint, blueprint_file)
     return blueprint_path
 
-  def _get_blueprint_path(self, blueprint_name, prefix: str = ""):
+  def _get_blueprint_path(self, blueprint_name, prefix: str = "") -> str:
     blueprint_path = os.path.join(
         self._get_storage_path(prefix), f"{blueprint_name}.yaml"
     )

xpk/core/capacity.py

@@ -152,6 +152,23 @@ def get_reservation_placement_policy(
   return output.strip()
 
 
+def get_reservation_deployment_type(
+    reservation: str, zone: str, project: str
+) -> str:
+  """Get reservation deployment type."""
+  command = (
+      f'gcloud beta compute reservations describe {reservation}'
+      f' --project={project} --zone={zone} --format="value(deploymentType)"'
+  )
+  return_code, output = run_command_for_value(
+      command, 'Get reservation deployment type', dry_run_return_val='DENSE'
+  )
+  if return_code != 0:
+    xpk_print(f'Get reservation deployment type ERROR {return_code}')
+    xpk_exit(1)
+  return output.strip()
+
+
 def verify_reservation_exists(args) -> int:
   """Verify the reservation exists.
 

xpk/core/capacity_test.py

@@ -0,0 +1,50 @@
+"""
+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, patch
+from .capacity import get_reservation_deployment_type
+
+
+@patch('xpk.core.capacity.xpk_print')
+def test_get_reservation_deployment_type_exits_with_command_fails(
+    xpk_print: MagicMock, mocker
+):
+  mocker.patch(
+      target='xpk.core.capacity.run_command_for_value', return_value=(1, '')
+  )
+  with pytest.raises(SystemExit):
+    get_reservation_deployment_type(
+        reservation='reservation', zone='zone', project='project'
+    )
+
+  assert (
+      'Get reservation deployment type ERROR 1'
+      in xpk_print.mock_calls[0].args[0]
+  )
+
+
+def test_get_reservation_deployment_type_returns_deployment_type_when_command_succeeds(
+    mocker,
+):
+  mocker.patch(
+      target='xpk.core.capacity.run_command_for_value',
+      return_value=(0, 'DENSE'),
+  )
+  result = get_reservation_deployment_type(
+      reservation='reservation', zone='zone', project='project'
+  )
+  assert result == 'DENSE'

xpk/core/config.py

@@ -22,7 +22,7 @@ from ..utils import file
 from ..utils.console import xpk_print
 
 # This is the version for XPK PyPI package
-__version__ = 'v0.14.2'
+__version__ = 'v0.14.3'
 XPK_CURRENT_VERSION = __version__
 XPK_CONFIG_FILE = os.path.expanduser('~/.config/xpk/config.yaml')
 

xpk/core/docker_container.py

@@ -97,7 +97,7 @@ def get_main_container(args, system, docker_image, resource_type) -> str:
     )
 
   gpu_workload_terminate_command = ''
-  if system.accelerator_type == AcceleratorType['GPU']:
+  if system.accelerator_type == AcceleratorType.GPU:
     gpu_workload_terminate_command = (
         'echo Main app is done > /usr/share/workload/workload_terminated; '
     )
@@ -105,7 +105,7 @@ def get_main_container(args, system, docker_image, resource_type) -> str:
   tpu_stacktrace_terminate_command = ''
   if (
       not args.use_pathways
-      and system.accelerator_type == AcceleratorType['TPU']
+      and system.accelerator_type == AcceleratorType.TPU
       and args.deploy_stacktrace_sidecar
   ):
     tpu_stacktrace_terminate_command = (
@@ -193,7 +193,7 @@ def get_user_workload_container(args, system: SystemCharacteristics):
   ].resource_type
   if (
       not args.use_pathways
-      and system.accelerator_type == AcceleratorType['TPU']
+      and system.accelerator_type == AcceleratorType.TPU
       and args.deploy_stacktrace_sidecar
   ):
     xpk_print(
@@ -219,7 +219,7 @@ def get_main_container_docker_image(args, system: SystemCharacteristics) -> str:
       Workload docker image as a YAML string
   """
 
-  if system.accelerator_type == AcceleratorType['GPU']:
+  if system.accelerator_type == AcceleratorType.GPU:
     return 'gpu-image'
 
   return f'{args.docker_name}'

xpk/core/docker_resources.py

@@ -43,10 +43,10 @@ def get_main_container_resources(
     return resources_yaml
 
   gpu_resources_yaml = """nvidia.com/gpu: {system.chips_per_vm}"""
-  if system.accelerator_type == AcceleratorType['GPU']:
+  if system.accelerator_type == AcceleratorType.GPU:
     return gpu_resources_yaml.format(system=system)
 
-  if system.accelerator_type == AcceleratorType['CPU']:
+  if system.accelerator_type == AcceleratorType.CPU:
     # CPUs don't have chips, but have a subresource called vCPUs.
     # system.chips_per_vm is used as a proxy for vCPUs.
     # Some vCPUs get used in hosting system pods of the workloads,
@@ -67,10 +67,10 @@ def get_env_container(args, system: SystemCharacteristics) -> str:
     str:
       YAML with the env config for the main container, as a YAML string.
   """
-  if system.accelerator_type == AcceleratorType['GPU']:
+  if system.accelerator_type == AcceleratorType.GPU:
     return get_gpu_env(args, system)
 
-  if system.accelerator_type == AcceleratorType['CPU']:
+  if system.accelerator_type == AcceleratorType.CPU:
     return get_cpu_env(args, system)
 
   return format_env_dict(args.env, system)
@@ -176,7 +176,7 @@ def get_cpu_env(args, system) -> str:
 
 
 def format_env_dict(env, system: SystemCharacteristics) -> str:
-  if system.accelerator_type == AcceleratorType['GPU']:
+  if system.accelerator_type == AcceleratorType.GPU:
     # For GPUs, it has two more spaces ahead of name and value respectively
     env_format = '''
                   - name: {key}
@@ -265,7 +265,7 @@ def get_volumes(args, system: SystemCharacteristics) -> str:
                   driver: {driver}"""
 
   if (
-      system.accelerator_type == AcceleratorType['TPU']
+      system.accelerator_type == AcceleratorType.TPU
       and args.deploy_stacktrace_sidecar
   ):
     volumes += """
@@ -317,7 +317,7 @@ def get_volume_mounts(args, system: SystemCharacteristics) -> str:
                   name: shared-tmp
                 """
   elif (
-      system.accelerator_type == AcceleratorType['TPU']
+      system.accelerator_type == AcceleratorType.TPU
       and args.deploy_stacktrace_sidecar
   ):
     volume_mount_yaml += """- name: tpu-stack-trace
@@ -325,7 +325,7 @@ def get_volume_mounts(args, system: SystemCharacteristics) -> str:
                 - name: shared-data
                   mountPath: /shared-volume
                 """
-  elif system.accelerator_type == AcceleratorType['GPU']:
+  elif system.accelerator_type == AcceleratorType.GPU:
     volume_mount_yaml = ''
 
   storages: list[Storage] = (
@@ -379,7 +379,7 @@ def add_container_ports(args, system: SystemCharacteristics) -> str:
     return ''
 
   gpu_port_yaml = """- containerPort: 6002"""
-  if system.accelerator_type == AcceleratorType['GPU']:
+  if system.accelerator_type == AcceleratorType.GPU:
     return gpu_port_yaml
   return port_yaml
 
@@ -394,7 +394,7 @@ def add_jax_coordinator_port(system) -> str:
     str:
       jax coordinator port as a YAML string
   """
-  if system.accelerator_type == AcceleratorType['CPU']:
+  if system.accelerator_type == AcceleratorType.CPU:
     return '- containerPort: 1234'
   return ''
 
@@ -411,6 +411,6 @@ def add_image_pull_policy_for_pw_or_gpu(args, system: SystemCharacteristics):
   """
   yaml = """imagePullPolicy: Always"""
 
-  if args.use_pathways or system.accelerator_type == AcceleratorType['GPU']:
+  if args.use_pathways or system.accelerator_type == AcceleratorType.GPU:
     return yaml.format(args=args)
   return ''

xpk/core/kjob.py

@@ -296,15 +296,13 @@ def create_job_template_instance(
     working_directory = JobTemplateDefaults.WORKING_DIRECTORY.value
   resources = (
       job_resources_template.format(gpu_per_node=system.chips_per_vm)
-      if system is not None
-      and system.accelerator_type == AcceleratorType["GPU"]
+      if system is not None and system.accelerator_type == AcceleratorType.GPU
       else ""
   )
 
   node_selector = (
       job_node_selector_template.format(gpu_name=system.gke_accelerator)
-      if system is not None
-      and system.accelerator_type == AcceleratorType["GPU"]
+      if system is not None and system.accelerator_type == AcceleratorType.GPU
       else ""
   )
   yml_string = job_template_yaml.format(
@@ -319,7 +317,7 @@ def create_job_template_instance(
       priority=args.priority if hasattr(args, "priority") else "medium",
       service_account=service_account,
   )
-  if system is not None and system.accelerator_type == AcceleratorType["GPU"]:
+  if system is not None and system.accelerator_type == AcceleratorType.GPU:
     yml_string = decorate_job_template_with_gpu(yml_string, system.device_type)
 
   return run_kubectl_apply(

xpk/core/kueue_manager.py

@@ -40,6 +40,7 @@ from ..core.commands import (
 from ..utils.file import write_tmp_file
 from ..utils.console import xpk_print, xpk_exit
 from ..utils.templates import TEMPLATE_PATH, get_templates_absolute_path
+from packaging.version import Version
 
 WAIT_FOR_KUEUE_TIMEOUT = "10m"
 CLUSTER_QUEUE_NAME = "cluster-queue"
@@ -51,7 +52,7 @@ 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 = "v0.12.2"
+KUEUE_VERSION = Version("v0.12.2")
 
 
 @dataclass
@@ -78,7 +79,7 @@ class KueueManager:
 
   def __init__(
       self,
-      kueue_version: str = KUEUE_VERSION,
+      kueue_version: Version = KUEUE_VERSION,
       template_path=TEMPLATE_PATH,
   ):
     self.kueue_version = kueue_version
@@ -111,9 +112,9 @@ class KueueManager:
         )
         return 0
       else:
-        xpk_print(f"Upgrading Kueue to version {self.kueue_version}...")
+        xpk_print(f"Upgrading Kueue to version v{self.kueue_version}...")
     else:
-      xpk_print(f"Installing Kueue version {self.kueue_version}...")
+      xpk_print(f"Installing Kueue version v{self.kueue_version}...")
 
     install_return_code = self.__install(tolerations)
     if install_return_code != 0:
@@ -121,7 +122,7 @@ class KueueManager:
 
     return self.__configure(kueue_config)
 
-  def get_installed_kueue_version(self) -> tuple[int, str | None]:
+  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}'"
@@ -130,15 +131,14 @@ class KueueManager:
     return_code, val = run_command_for_value(
         command,
         task,
-        dry_run_return_val="""
-        v0.12.1""",
+        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_tag[-1]
+    return return_code, Version(version_tag[-1])
 
   def __install(
       self,
@@ -162,7 +162,7 @@ class KueueManager:
     return self.__wait_for_kueue_available()
 
   def __install_kueue_crs(self) -> int:
-    manifest_url = f"https://github.com/kubernetes-sigs/kueue/releases/download/{self.kueue_version}/manifests.yaml"
+    manifest_url = f"https://github.com/kubernetes-sigs/kueue/releases/download/v{self.kueue_version}/manifests.yaml"
     install_command = (
         f"kubectl apply --server-side --force-conflicts -f {manifest_url}"
     )
@@ -199,7 +199,7 @@ class KueueManager:
       0 if successful and 1 otherwise.
     """
     command = (
-        "kubectl wait deploy/kueue-controller-manager -nkueue-system"
+        "kubectl wait deploy/kueue-controller-manager -n kueue-system"
         f" --for=condition=available --timeout={WAIT_FOR_KUEUE_TIMEOUT}"
     )
     task = "Wait for Kueue to be available"
@@ -421,3 +421,14 @@ class KueueManager:
     if return_code != 0:
       xpk_print(f"{task} returned ERROR {return_code}")
     return return_code
+
+
+def has_sub_slicing_enabled() -> tuple[int, bool | None]:
+  return_code, value = run_command_for_value(
+      command="kubectl get topology", task="Get defined topologies"
+  )
+
+  if return_code != 0:
+    return return_code, None
+
+  return return_code, SUB_SLICE_TOPOLOGY_NAME in value