xpk 0.14.0__py3-none-any.whl → 0.14.1__py3-none-any.whl

xpk/core/kueue_manager_test.py

@@ -76,9 +76,7 @@ class KueueManagerTest(unittest.TestCase):
       mock_install.assert_called_once()
       mock_configure.assert_called_once()
 
-  @patch(
-      "xpk.core.kueue_manager.KueueManager._KueueManager__get_installed_kueue_version"
-  )
+  @patch("xpk.core.kueue_manager.KueueManager.get_installed_kueue_version")
   @patch("xpk.core.kueue_manager.KueueManager._KueueManager__install")
   @patch("xpk.core.kueue_manager.KueueManager._KueueManager__configure")
   def test_install_or_upgrade_when_newer_version_already_installed(
@@ -95,9 +93,7 @@ class KueueManagerTest(unittest.TestCase):
     mock_install.assert_not_called()
     mock_configure.assert_not_called()
 
-  @patch(
-      "xpk.core.kueue_manager.KueueManager._KueueManager__get_installed_kueue_version"
-  )
+  @patch("xpk.core.kueue_manager.KueueManager.get_installed_kueue_version")
   def test_install_or_upgrade_when_outdated(
       self,
       mock_get_version,
@@ -121,9 +117,7 @@ class KueueManagerTest(unittest.TestCase):
     mock_install.assert_called_once()
     mock_configure.assert_called_once()
 
-  @patch(
-      "xpk.core.kueue_manager.KueueManager._KueueManager__get_installed_kueue_version"
-  )
+  @patch("xpk.core.kueue_manager.KueueManager.get_installed_kueue_version")
   def test_install_or_upgrade_when_not_installed(
       self,
       mock_get_version,
@@ -155,7 +149,7 @@ class KueueManagerTest(unittest.TestCase):
             return_value=0,
         ) as mock_run_retry,
         patch(
-            "xpk.core.kueue_manager.KueueManager._KueueManager__get_installed_kueue_version",
+            "xpk.core.kueue_manager.KueueManager.get_installed_kueue_version",
             return_value=(1, None),
         ),
         patch(
@@ -199,7 +193,7 @@ class KueueManagerTest(unittest.TestCase):
             return_value=0,
         ) as mock_run_retry,
         patch(
-            "xpk.core.kueue_manager.KueueManager._KueueManager__get_installed_kueue_version",
+            "xpk.core.kueue_manager.KueueManager.get_installed_kueue_version",
             return_value=(1, None),
         ),
         patch(
@@ -224,9 +218,7 @@ class KueueManagerTest(unittest.TestCase):
       self.assertEqual(result, 0)
       self.assertEqual(mock_run_retry.call_count, 0)
 
-  @patch(
-      "xpk.core.kueue_manager.KueueManager._KueueManager__get_installed_kueue_version"
-  )
+  @patch("xpk.core.kueue_manager.KueueManager.get_installed_kueue_version")
   @patch("xpk.core.kueue_manager.KueueManager._KueueManager__apply_manifest")
   def test_configuration_updates_resources(
       self, mock_apply_manifest, mock_get_version
@@ -240,6 +232,7 @@ class KueueManagerTest(unittest.TestCase):
         total_chips=8,
         cpu_limit=100,
         memory_limit="100Gi",
+        configure_sub_slicing=False,
     )
 
     with (
@@ -265,6 +258,7 @@ class KueueManagerTest(unittest.TestCase):
         total_chips=8,
         cpu_limit=100,
         memory_limit="100Gi",
+        configure_sub_slicing=False,
     )
 
     with (
@@ -274,7 +268,7 @@ class KueueManagerTest(unittest.TestCase):
         ),
         patch.object(
             self.kueue_manager,
-            "_KueueManager__get_installed_kueue_version",
+            "get_installed_kueue_version",
             return_value=(1, None),
         ),
         patch.object(
@@ -307,6 +301,7 @@ class KueueManagerTest(unittest.TestCase):
         total_chips=8,
         cpu_limit=100,
         memory_limit="100Gi",
+        configure_sub_slicing=False,
     )
 
     with (
@@ -316,7 +311,7 @@ class KueueManagerTest(unittest.TestCase):
         ),
         patch.object(
             self.kueue_manager,
-            "_KueueManager__get_installed_kueue_version",
+            "get_installed_kueue_version",
             return_value=(1, None),
         ),
         patch.object(
@@ -344,7 +339,7 @@ class KueueManagerTest(unittest.TestCase):
   @patch(
       "xpk.core.kueue_manager.KueueManager._KueueManager__update_kueue_resources_if_necessary"
   )
-  def test_configure_generates_correct_manifest(
+  def test_configure_generates_correct_manifest_for_tpu(
       self, mock_update_resources, mock_install
   ):
     """Test that __configure generates the correct manifest content for TPUs."""
@@ -357,6 +352,7 @@ class KueueManagerTest(unittest.TestCase):
         memory_limit="100Gi",
         autoprovisioning_enabled=False,
         num_slices=2,
+        configure_sub_slicing=False,
     )
 
     rendered_manifest = self._trigger_installation(kueue_config)
@@ -413,6 +409,7 @@ class KueueManagerTest(unittest.TestCase):
         autoprovisioning_enabled=False,
         num_slices=1,
         flex=True,
+        configure_sub_slicing=False,
     )
 
     rendered_manifest = self._trigger_installation(kueue_config)
@@ -432,7 +429,7 @@ class KueueManagerTest(unittest.TestCase):
   @patch(
       "xpk.core.kueue_manager.KueueManager._KueueManager__update_kueue_resources_if_necessary"
   )
-  def test_configure_generates_correct_manifest_with_topology(
+  def test_configure_generates_correct_manifest_with_gke_default_topology(
       self, mock_update_resources, mock_install
   ):
     """Test that __configure generates correct manifest for GPUs."""
@@ -444,11 +441,11 @@ class KueueManagerTest(unittest.TestCase):
         cpu_limit=100,
         memory_limit="100Gi",
         num_slices=2,
+        configure_sub_slicing=False,
     )
 
     rendered_manifest = self._trigger_installation(kueue_config)
 
-    self.assertIn("kind: Topology", rendered_manifest)
     manifest_docs = list(yaml.safe_load_all(rendered_manifest))
     resource_flavor = _first(
         doc for doc in manifest_docs if doc["kind"] == "ResourceFlavor"
@@ -459,6 +456,40 @@ class KueueManagerTest(unittest.TestCase):
         ],
         "h100-mega-80gb-8",
     )
+    self.assertEqual(resource_flavor["spec"]["topologyName"], "gke-default")
+    topology = _first(doc for doc in manifest_docs if doc["kind"] == "Topology")
+    self.assertEqual(topology["metadata"]["name"], "gke-default")
+
+  @patch("xpk.core.kueue_manager.KueueManager._KueueManager__install")
+  @patch(
+      "xpk.core.kueue_manager.KueueManager._KueueManager__update_kueue_resources_if_necessary"
+  )
+  def test_configure_generates_correct_manifest_with_sub_slicing(
+      self, mock_update_resources, mock_install
+  ):
+    """Test that __configure generates correct manifest with sub-slicing topology."""
+    mock_install.return_value = 0
+    mock_update_resources.return_value = 0
+    kueue_config = KueueConfig(
+        system=self.mock_system_chars,
+        total_chips=16,
+        cpu_limit=100,
+        memory_limit="100Gi",
+        num_slices=2,
+        configure_sub_slicing=True,
+    )
+
+    rendered_manifest = self._trigger_installation(kueue_config)
+
+    manifest_docs = list(yaml.safe_load_all(rendered_manifest))
+    resource_flavor = _first(
+        doc for doc in manifest_docs if doc["kind"] == "ResourceFlavor"
+    )
+    self.assertEqual(
+        resource_flavor["spec"]["topologyName"], "sub-slice-topology"
+    )
+    topology = _first(doc for doc in manifest_docs if doc["kind"] == "Topology")
+    self.assertEqual(topology["metadata"]["name"], "sub-slice-topology")
 
   @patch("xpk.core.kueue_manager.KueueManager._KueueManager__install")
   @patch(
@@ -477,6 +508,7 @@ class KueueManagerTest(unittest.TestCase):
         memory_limit="100Gi",
         is_pathways_cluster=True,
         num_slices=2,
+        configure_sub_slicing=False,
     )
 
     rendered_manifest = self._trigger_installation(kueue_config)
@@ -513,7 +545,7 @@ class KueueManagerTest(unittest.TestCase):
     """Calls Kueue installation and returns the rendered manifest."""
     with (
         patch.object(
-            self.kueue_manager, "_KueueManager__get_installed_kueue_version"
+            self.kueue_manager, "get_installed_kueue_version"
         ) as mock_get_version,
         patch.object(
             self.kueue_manager, "_KueueManager__apply_manifest"

xpk/core/system_characteristics.py

@@ -135,10 +135,9 @@ def get_tpu_system_characteristics_map(
 ) -> dict[str, SystemCharacteristics]:
   system_characteristics_map = {}
   for topology in supported_topologies:
-    total_chips = get_topology_product(topology)
-    num_tensorcores = total_chips * tensorcores_per_chip
-    chips_per_vm = 1 if total_chips == 1 else 4
-    vms_per_slice = total_chips // chips_per_vm
+    chips_per_vm = compute_chips_per_vm(topology)
+    vms_per_slice = compute_vms_per_slice(topology)
+    num_tensorcores = compute_num_tensorcores(tensorcores_per_chip, topology)
     system = SystemCharacteristics(
         topology=topology,
         vms_per_slice=vms_per_slice,
@@ -156,6 +155,19 @@ def get_tpu_system_characteristics_map(
   return system_characteristics_map
 
 
+def compute_chips_per_vm(topology: str) -> int:
+  return 1 if get_topology_product(topology) == 1 else 4
+
+
+def compute_num_tensorcores(tensorcores_per_chip: int, topology: str) -> int:
+  return get_topology_product(topology) * tensorcores_per_chip
+
+
+def compute_vms_per_slice(topology: str) -> int:
+  chips_per_vm = compute_chips_per_vm(topology)
+  return get_topology_product(topology) // chips_per_vm
+
+
 ################### Subcommand Helper Functions #############################
 """ !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 IF YOU MODIFY THE BELOW UserFacingNameToSystemCharacteristics MAP YOU SHOULD

xpk/core/system_characteristics_test.py

@@ -0,0 +1,73 @@
+"""
+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.
+"""
+
+from .system_characteristics import get_tpu_system_characteristics_map, SystemCharacteristics
+
+
+def test_get_tpu_system_characteristics_map_returns_correct_values_for_1x1_topology():
+  result = get_tpu_system_characteristics_map(
+      prefix="test",
+      tensorcores_per_chip=1,
+      gke_accelerator="test",
+      machine_type="test",
+      supported_topologies=["1x1"],
+      supports_sub_slicing=False,
+      requires_workload_policy=True,
+  )
+
+  expected_system_characteristics = SystemCharacteristics(
+      topology="1x1",
+      vms_per_slice=1,
+      gke_accelerator="test",
+      gce_machine_type="test",
+      chips_per_vm=1,
+      accelerator_type=1,
+      device_type="test-1",
+      supports_sub_slicing=False,
+      requires_workload_policy=True,
+  )
+  assert result == {
+      "test-1": expected_system_characteristics,
+      "test-1x1": expected_system_characteristics,
+  }
+
+
+def test_get_tpu_system_characteristics_map_returns_correct_values_for_2x2_topology():
+  result = get_tpu_system_characteristics_map(
+      prefix="test",
+      tensorcores_per_chip=2,
+      gke_accelerator="test",
+      machine_type="test",
+      supported_topologies=["2x2"],
+      supports_sub_slicing=False,
+      requires_workload_policy=True,
+  )
+
+  expected_system_characteristics = SystemCharacteristics(
+      topology="2x2",
+      vms_per_slice=1,
+      gke_accelerator="test",
+      gce_machine_type="test",
+      chips_per_vm=4,
+      accelerator_type=1,
+      device_type="test-8",
+      supports_sub_slicing=False,
+      requires_workload_policy=True,
+  )
+  assert result == {
+      "test-8": expected_system_characteristics,
+      "test-2x2": expected_system_characteristics,
+  }

xpk/templates/cluster_preheat.yaml.j2

@@ -0,0 +1,31 @@
+apiVersion: apps/v1
+kind: DaemonSet
+metadata:
+  name: {{ cachekey }}
+  labels:
+    k8s-app: {{ cachekey }}
+spec:
+  selector:
+    matchLabels:
+      k8s-app: {{ cachekey }}
+  updateStrategy:
+    type: RollingUpdate
+  template:
+    metadata:
+      labels:
+        name: {{ cachekey }}
+        k8s-app: {{ cachekey }}
+    spec:
+      affinity:
+        nodeAffinity:
+          requiredDuringSchedulingIgnoredDuringExecution:
+            nodeSelectorTerms:
+            - matchExpressions:
+              - key: {{ nodeSelectorKey }}
+                operator: Exists
+      tolerations:
+      - operator: "Exists"
+      containers:
+      - image: {{ image_name }}
+        name: {{ cachekey }}
+        command: [ "sleep", "inf" ]

xpk/templates/filestore-pv.yaml

@@ -0,0 +1,17 @@
+apiVersion: v1
+kind: PersistentVolume
+metadata:
+  name: xpk-filestore-pv
+spec:
+  storageClassName:
+  capacity:
+    storage: 
+  accessModes:
+  persistentVolumeReclaimPolicy: Retain
+  volumeMode: Filesystem
+  csi:
+    driver: filestore.csi.storage.gke.io
+    volumeHandle:
+    volumeAttributes:
+      ip:
+      volume:

xpk/templates/filestore-pvc.yaml

@@ -0,0 +1,11 @@
+kind: PersistentVolumeClaim
+apiVersion: v1
+metadata:
+  name: 
+spec:
+  accessModes:
+  storageClassName:
+  volumeName: 
+  resources:
+    requests:
+      storage: 

xpk/templates/filestore-sc.yaml

@@ -0,0 +1,10 @@
+apiVersion: storage.k8s.io/v1
+kind: StorageClass
+metadata:
+  name: 
+provisioner: filestore.csi.storage.gke.io
+volumeBindingMode: Immediate
+allowVolumeExpansion: true
+parameters:
+  tier: standard
+  network: default

xpk/templates/fuse-pv.yaml

@@ -0,0 +1,17 @@
+apiVersion: v1
+kind: PersistentVolume
+metadata:
+  name:
+spec:
+  accessModes:
+  - ReadWriteMany
+  capacity:
+    storage:
+  storageClassName: example-storage-class
+  mountOptions:
+    - implicit-dirs
+  csi:
+    driver: gcsfuse.csi.storage.gke.io
+    volumeHandle:
+    volumeAttributes:
+      gcsfuseLoggingSeverity: warning

xpk/templates/fuse-pvc.yaml

@@ -0,0 +1,13 @@
+apiVersion: v1
+kind: PersistentVolumeClaim
+metadata:
+  name:
+  namespace: default
+spec:
+  accessModes:
+  - ReadWriteMany
+  resources:
+    requests:
+      storage:
+  volumeName:
+  storageClassName: example-storage-class

xpk/templates/kueue_config.yaml.j2

@@ -0,0 +1,95 @@
+{% for flavor in flavors %}
+apiVersion: kueue.x-k8s.io/v1beta1
+kind: ResourceFlavor
+metadata:
+  name: "{{ flavor.name }}"
+spec:
+  nodeLabels: {{ flavor.nodeLabels | tojson }}
+  {% if flavor.topologyLabel %}
+  {{ flavor.topologyLabel }}
+  {% endif %}
+---
+{% endfor %}
+apiVersion: kueue.x-k8s.io/v1beta1
+kind: AdmissionCheck
+metadata:
+  name: dws-prov
+spec:
+  controllerName: kueue.x-k8s.io/provisioning-request
+  parameters:
+    apiGroup: kueue.x-k8s.io
+    kind: ProvisioningRequestConfig
+    name: dws-config
+---
+apiVersion: kueue.x-k8s.io/v1beta1
+kind: ProvisioningRequestConfig
+metadata:
+  name: dws-config
+spec:
+  provisioningClassName: queued-provisioning.gke.io
+  podSetUpdates:
+    nodeSelector:
+    - key: autoscaling.gke.io/provisioning-request
+      valueFromProvisioningClassDetail: ResizeRequestName
+  managedResources:
+  - {{ managed_resource }}
+---
+apiVersion: kueue.x-k8s.io/v1beta1
+kind: ClusterQueue
+metadata:
+  name: "{{ cluster_queue_name }}"
+spec:
+  preemption:
+    reclaimWithinCohort: Never # Don't preempt other queues in the cohort.
+    withinClusterQueue: LowerPriority
+  namespaceSelector: {} # match all.
+  resourceGroups: {{ resource_groups }}
+  {{ admission_checks | indent(2) }}
+---
+apiVersion: kueue.x-k8s.io/v1beta1
+kind: LocalQueue
+metadata:
+  namespace: default
+  name: {{ local_queue_name }}
+spec:
+  clusterQueue: {{ cluster_queue_name }}
+---
+apiVersion: scheduling.k8s.io/v1
+kind: PriorityClass
+metadata:
+  name: very-low
+value: 100
+globalDefault: false
+description: "Very Low"
+---
+apiVersion: scheduling.k8s.io/v1
+kind: PriorityClass
+metadata:
+  name: low
+value: 250
+globalDefault: false
+description: "Low"
+---
+apiVersion: scheduling.k8s.io/v1
+kind: PriorityClass
+metadata:
+  name: medium
+value: 500
+globalDefault: false
+description: "Medium"
+---
+apiVersion: scheduling.k8s.io/v1
+kind: PriorityClass
+metadata:
+  name: high
+value: 750
+globalDefault: false
+description: "High"
+---
+apiVersion: scheduling.k8s.io/v1
+kind: PriorityClass
+metadata:
+  name: very-high
+value: 1000
+globalDefault: false
+description: "Very High"

xpk/templates/kueue_gke_default_topology.yaml.j2

@@ -0,0 +1,10 @@
+apiVersion: kueue.x-k8s.io/v1beta1
+kind: Topology
+metadata:
+  name: "gke-default"
+spec:
+  levels:
+  - nodeLabel: "cloud.google.com/gce-topology-block"
+  - nodeLabel: "cloud.google.com/gce-topology-subblock"
+  - nodeLabel: "cloud.google.com/gce-topology-host"
+  - nodeLabel: "kubernetes.io/hostname"

xpk/templates/kueue_sub_slicing_topology.yaml.j2

@@ -0,0 +1,14 @@
+apiVersion: kueue.x-k8s.io/v1beta1
+kind: Topology
+metadata:
+  name: {{ sub_slice_topology_name }}
+spec:
+  levels:
+  - nodeLabel: "cloud.google.com/gke-tpu-slice-16x16-id"
+  - nodeLabel: "cloud.google.com/gke-tpu-slice-8x16-id"
+  - nodeLabel: "cloud.google.com/gke-tpu-slice-8x8-id"
+  - nodeLabel: "cloud.google.com/gke-tpu-slice-4x8-id"
+  - nodeLabel: "cloud.google.com/gke-tpu-slice-4x4-id"
+  - nodeLabel: "cloud.google.com/gke-tpu-slice-2x4-id"
+  - nodeLabel: "cloud.google.com/gke-tpu-slice-2x2-id"
+  - nodeLabel: "kubernetes.io/hostname"

xpk/templates/mtc-cpc.yaml

@@ -0,0 +1,15 @@
+apiVersion: checkpointing.gke.io/v1
+kind: CheckpointConfiguration
+metadata:
+  name: my-checkpointconfiguration
+spec:
+  cloudStorageBucketName:
+  # This field is optional
+  nodeSelector:
+    node.kubernetes.io/instance-type:
+  # This field is optional
+  tolerations:
+  - key:
+    operator: Exists
+    effect: NoSchedule
+  inMemoryVolumeSize:

xpk/templates/volume_bundle.yaml

@@ -0,0 +1,7 @@
+apiVersion: kjobctl.x-k8s.io/v1alpha1
+kind: VolumeBundle
+metadata:
+  name: $NAME
+spec:
+  volumes: []
+  containerVolumeMounts: []

xpk/utils/templates.py

@@ -18,7 +18,7 @@ import os
 
 import ruamel.yaml
 
-TEMPLATE_PATH = "src/xpk/templates/"
+TEMPLATE_PATH = "templates"
 
 yaml = ruamel.yaml.YAML()
 
@@ -28,3 +28,16 @@ def load(path: str) -> dict:
   with open(template_path, "r", encoding="utf-8") as file:
     data: dict = yaml.load(file)
   return data
+
+
+def get_templates_absolute_path(templates_path: str = TEMPLATE_PATH) -> str:
+  """
+  Return the absolute path to the templates folder
+
+  Args:
+    templates_path: The path to the templates folder relative to the src/xpk directory
+  """
+  current_file_path = os.path.abspath(__file__)
+  current_dir = os.path.dirname(current_file_path)
+  xpk_package_dir = os.path.dirname(current_dir)
+  return os.path.join(xpk_package_dir, templates_path)

xpk/utils/topology.py

@@ -35,3 +35,12 @@ def parse_topology(topology: str) -> list[int]:
     raise ValueError("Topology is an empty string")
 
   return [int(el) for el in topology.lower().split("x")]
+
+
+def is_topology_contained(contained: str, container: str) -> bool:
+  contained_parsed = parse_topology(contained)
+  container_parsed = parse_topology(container)
+  return len(contained_parsed) == len(container_parsed) and all(
+      contained <= container
+      for contained, container in zip(contained_parsed, container_parsed)
+  )

xpk/utils/topology_test.py

@@ -15,7 +15,7 @@ limitations under the License.
 """
 
 import pytest
-from .topology import is_topology_valid, get_topology_product, parse_topology
+from .topology import is_topology_valid, get_topology_product, parse_topology, is_topology_contained
 
 
 def test_is_topology_valid_with_invalid_topology():
@@ -41,3 +41,23 @@ def test_parse_topology_with_empty_input():
 def test_get_topology_product():
   result = get_topology_product("1x2x3")
   assert result == 6
+
+
+def test_is_topology_contained_with_container_smaller_than_contained_returns_false():
+  result = is_topology_contained(contained="3x3x3", container="2x2x2")
+  assert result is False
+
+
+def test_is_topology_contained_with_container_larger_than_contained_returns_true():
+  result = is_topology_contained(contained="1x1x1", container="2x2x2")
+  assert result is True
+
+
+def test_is_topology_contained_with_container_equal_to_contained_returns_true():
+  result = is_topology_contained(contained="2x2x2", container="2x2x2")
+  assert result is True
+
+
+def test_is_topology_contained_with_different_topologies_dimensions_returns_false():
+  result = is_topology_contained(contained="2x2", container="2x2x2")
+  assert result is False

{xpk-0.14.0.dist-info → xpk-0.14.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xpk
-Version: 0.14.0
+Version: 0.14.1
 Summary: xpk helps Cloud developers to orchestrate training jobs on accelerators on GKE.
 Author-email: XPK team <xpk-code-reviewers@google.com>
 License: Apache-2.0