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

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/feature_flags.py

@@ -0,0 +1,28 @@
+"""
+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 os
+
+
+def _get_boolean_flag(flag: str, default: bool) -> bool:
+  return os.getenv(flag, str(default)).lower() == "true"
+
+
+class _FeatureFlags:
+  SUB_SLICING_ENABLED = _get_boolean_flag("SUB_SLICING_ENABLED", default=False)
+
+
+FeatureFlags = _FeatureFlags()

xpk/utils/kueue.py

@@ -0,0 +1,20 @@
+"""
+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.
+"""
+
+
+def is_queued_cluster(num_slices: int) -> bool:
+  """Determines if admission checks should be enabled and cluster queued."""
+  return num_slices <= 1

xpk/utils/templates.py

@@ -18,6 +18,8 @@ import os
 
 import ruamel.yaml
 
+TEMPLATE_PATH = "templates"
+
 yaml = ruamel.yaml.YAML()
 
 
@@ -26,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

@@ -0,0 +1,46 @@
+"""
+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 functools import reduce
+from operator import mul
+
+
+def is_topology_valid(topology: str) -> bool:
+  try:
+    parse_topology(topology)
+    return True
+  except ValueError:
+    return False
+
+
+def get_topology_product(topology: str) -> int:
+  return reduce(mul, parse_topology(topology), 1)
+
+
+def parse_topology(topology: str) -> list[int]:
+  if len(topology) <= 0:
+    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

@@ -0,0 +1,63 @@
+"""
+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 .topology import is_topology_valid, get_topology_product, parse_topology, is_topology_contained
+
+
+def test_is_topology_valid_with_invalid_topology():
+  result = is_topology_valid("N/A")
+  assert result is False
+
+
+def test_is_topology_valid_with_valid_topology():
+  result = is_topology_valid("1x1x1")
+  assert result is True
+
+
+def test_parse_topology_with_valid_topology():
+  result = parse_topology("1x2x3")
+  assert result == [1, 2, 3]
+
+
+def test_parse_topology_with_empty_input():
+  with pytest.raises(ValueError):
+    parse_topology("")
+
+
+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/utils/validation.py

@@ -15,66 +15,90 @@ limitations under the License.
 """
 
 from ..core.commands import run_command_for_value
+from ..core.config import __version__ as xpk_version
 from .console import xpk_exit, xpk_print
 from ..commands.config import xpk_cfg
 from ..core.config import DEPENDENCIES_KEY
-from ..commands.version import get_xpk_version
-
-
-validation_commands = {
-    'kubectl': {
-        'command': 'kubectl --help',
-        'message': (
-            '`kubectl` not installed. Please follow'
-            ' https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#prerequisites'
-            ' to install xpk prerequisites.'
-        ),
-    },
-    'kjob': {
-        'command': 'kubectl kjob --help',
-        'message': (
-            '`kjobctl` not installed. Please follow'
-            ' https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#prerequisites'
-            ' to install xpk prerequisites.'
-        ),
-    },
-    'gcloud': {
-        'command': 'gcloud version',
-        'message': (
-            '`gcloud not installed. Please follow'
-            ' https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#prerequisites'
-            ' to install xpk prerequisites.'
-        ),
-    },
-    'docker': {
-        'command': 'docker version',
-        'message': (
-            '`docker` not installed. Please follow'
-            ' https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#prerequisites'
-            ' to install xpk prerequisites.'
-        ),
-    },
-    'kueuectl': {
-        'command': 'kubectl kueue --help',
-        'message': (
-            '`kueuectl` not installed. Please follow'
-            ' https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#prerequisites'
-            ' to install xpk prerequisites.'
-        ),
-    },
-}
+from enum import Enum
+from dataclasses import dataclass
+
+
+@dataclass
+class _SystemDependency:
+  command: str
+  message: str
+
+
+class SystemDependency(Enum):
+  """Represents required system dependencies."""
+
+  KUBECTL = _SystemDependency(
+      command='kubectl --help',
+      message=(
+          '`kubectl` not installed. Please follow'
+          ' https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#prerequisites'
+          ' to install xpk prerequisites.'
+      ),
+  )
+  KJOB = _SystemDependency(
+      command='kubectl kjob --help',
+      message=(
+          '`kjobctl` not installed. Please follow'
+          ' https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#prerequisites'
+          ' to install xpk prerequisites.'
+      ),
+  )
+  GCLOUD = _SystemDependency(
+      command='gcloud version',
+      message=(
+          '`gcloud not installed. Please follow'
+          ' https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#prerequisites'
+          ' to install xpk prerequisites.'
+      ),
+  )
+  DOCKER = _SystemDependency(
+      command='docker version',
+      message=(
+          '`docker` not installed. Please follow'
+          ' https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#prerequisites'
+          ' to install xpk prerequisites.'
+      ),
+  )
+  KUEUECTL = _SystemDependency(
+      command='kubectl kueue --help',
+      message=(
+          '`kueuectl` not installed. Please follow'
+          ' https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#prerequisites'
+          ' to install xpk prerequisites.'
+      ),
+  )
+
+
+def should_validate_dependencies(args):
+  skip_validation = 'skip_validation' in args and args.skip_validation
+  dry_run = 'dry_run' in args and args.dry_run
+  return not skip_validation and not dry_run
 
 
 def validate_dependencies():
+  """Validates all system dependencies if validation has not been done with current XPK version."""
   deps_version = xpk_cfg.get(DEPENDENCIES_KEY)
-  xpk_version = get_xpk_version()
   if deps_version is None or deps_version != xpk_version:
-    for name, check in validation_commands.items():
-      cmd, message = check['command'], check['message']
-      code, _ = run_command_for_value(
-          cmd, f'Validate {name} installation.', None
-      )
-      if code != 0:
-        xpk_print(message)
-        xpk_exit(code)
-    xpk_cfg.set(DEPENDENCIES_KEY, get_xpk_version())
+    validate_dependencies_list(list(SystemDependency))
+    xpk_cfg.set(DEPENDENCIES_KEY, xpk_version)
+
+
+def validate_dependencies_list(dependencies: list[SystemDependency]):
+  """Validates a list of system dependencies and returns none or exits with error."""
+  for dependency in dependencies:
+    _validate_dependency(dependency)
+
+
+def _validate_dependency(dependency: SystemDependency) -> None:
+  """Validates system dependency and returns none or exits with error."""
+  name, value = dependency.name, dependency.value
+  cmd, message = value.command, value.message
+  code, _ = run_command_for_value(cmd, f'Validate {name} installation.')
+  if code != 0:
+    xpk_print(message)
+    xpk_exit(code)

xpk/utils/validation_test.py

@@ -0,0 +1,37 @@
+"""
+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 .validation import validate_dependencies_list, SystemDependency
+
+
+def test_validate_dependencies_list_returns_nothing_for_successful_validation(
+    mocker,
+):
+  mocker.patch(
+      'xpk.utils.validation.run_command_for_value', return_value=(0, '')
+  )
+  validate_dependencies_list([SystemDependency.DOCKER])
+
+
+def test_validate_dependencies_list_exits_with_error_for_failed_validation(
+    mocker,
+):
+  mocker.patch(
+      'xpk.utils.validation.run_command_for_value', return_value=(1, '')
+  )
+  with pytest.raises(SystemExit):
+    validate_dependencies_list([SystemDependency.DOCKER])

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xpk
-Version: 0.13.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
@@ -22,15 +22,18 @@ Requires-Dist: google-api-core==2.24.1
 Requires-Dist: packaging==24.2
 Requires-Dist: google-cloud-filestore==1.12.0
 Requires-Dist: google-cloud-storage
+Requires-Dist: Jinja2==3.1.6
 Provides-Extra: dev
 Requires-Dist: pyink==24.3.0; extra == "dev"
 Requires-Dist: pylint>=2.6.0; extra == "dev"
 Requires-Dist: pre-commit; extra == "dev"
 Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pytest-mock==3.15.1; extra == "dev"
 Requires-Dist: docker==7.1.0; extra == "dev"
 Requires-Dist: mypy~=1.17; extra == "dev"
 Requires-Dist: types-PyYAML==6.0.2; extra == "dev"
 Requires-Dist: types-docker~=7.1.0.0; extra == "dev"
+Requires-Dist: pylint-per-file-ignores==1.4.0; extra == "dev"
 Dynamic: license-file
 
 <!--
@@ -76,6 +79,7 @@ XPK supports the following TPU types:
 * v5e
 * v5p
 * Trillium (v6e)
+* Ironwood (tpu7x)
 
 and the following GPU types:
 * A100
@@ -83,6 +87,7 @@ and the following GPU types:
 * A3-Mega (h100-mega) - [Create cluster](#provisioning-a3-ultra-a3-mega-and-a4-clusters-gpu-machines), [Create workloads](#workloads-for-a3-ultra-a3-mega-and-a4-clusters-gpu-machines)
 * A3-Ultra (h200) - [Create cluster](#provisioning-a3-ultra-a3-mega-and-a4-clusters-gpu-machines), [Create workloads](#workloads-for-a3-ultra-a3-mega-and-a4-clusters-gpu-machines)
 * A4 (b200) - [Create cluster](#provisioning-a3-ultra-a3-mega-and-a4-clusters-gpu-machines), [Create workloads](#workloads-for-a3-ultra-a3-mega-and-a4-clusters-gpu-machines)
+* A4X (gb200)
 
 and the following CPU types:
 * n2-standard-32