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

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/core/vertex.py

@@ -66,7 +66,7 @@ def create_vertex_experiment(args) -> dict | None:
   )
 
   metadata_configmap_name = f'{args.cluster}-{CLUSTER_METADATA_CONFIGMAP}'
-  cluster_config_map = get_cluster_configmap(args, metadata_configmap_name)
+  cluster_config_map = get_cluster_configmap(metadata_configmap_name)
 
   if cluster_config_map is None or 'tensorboard_name' not in cluster_config_map:
     xpk_print(

xpk/core/workload.py

@@ -17,7 +17,7 @@ limitations under the License.
 import re
 from ..utils.console import xpk_exit, xpk_print
 from .commands import run_command_for_value
-from .gcloud_context import zone_to_region
+from .gcloud_context import get_cluster_location
 
 
 def workload_list_awk_command(filter_key) -> str:
@@ -131,7 +131,7 @@ def get_workload_list(args) -> tuple[int, str]:
   if hasattr(args, 'filter_by_job'):
     task += f' with filter-by-job={args.filter_by_job}'
 
-  return_code, return_value = run_command_for_value(command, task, args)
+  return_code, return_value = run_command_for_value(command, task)
   return return_code, return_value
 
 
@@ -152,7 +152,7 @@ def check_if_workload_exists(args) -> bool:
 
   command = f"kubectl get workloads -o=custom-columns='{s}'"
   return_code, return_msg = run_command_for_value(
-      command, 'Check if Workload Already Exists', args
+      command, 'Check if Workload Already Exists'
   )
 
   if return_code != 0:
@@ -186,7 +186,7 @@ def wait_for_job_completion(args) -> int:
   # Get the full workload name
   get_workload_name_cmd = f'kubectl get workloads | grep jobset-{args.workload}'
   return_code, return_value = run_command_for_value(
-      get_workload_name_cmd, 'Get full workload name', args
+      get_workload_name_cmd, 'Get full workload name'
   )
   if return_code != 0:
     xpk_print(f'Get full workload name request returned ERROR {return_code}')
@@ -205,7 +205,6 @@ def wait_for_job_completion(args) -> int:
   return_code, return_value = run_command_for_value(
       wait_cmd,
       f'Wait for workload to finish with timeout of {timeout_msg}',
-      args,
       print_timer=True,
   )
   if return_code != 0:
@@ -214,7 +213,7 @@ def wait_for_job_completion(args) -> int:
           f'Timed out waiting for your workload after {timeout_msg}, see your'
           ' workload here:'
           # pylint: disable=line-too-long
-          f' https://console.cloud.google.com/kubernetes/service/{zone_to_region(args.zone)}/{args.cluster}/default/{args.workload}/details?project={args.project}'
+          f' https://console.cloud.google.com/kubernetes/service/{get_cluster_location(args.project, args.cluster, args.zone)}/{args.cluster}/default/{args.workload}/details?project={args.project}'
       )
       return 124
     else:
@@ -224,14 +223,14 @@ def wait_for_job_completion(args) -> int:
   xpk_print(
       'Finished waiting for your workload, see your workload here:'
       # pylint: disable=line-too-long
-      f' https://console.cloud.google.com/kubernetes/service/{zone_to_region(args.zone)}/{args.cluster}/default/{args.workload}/details?project={args.project}'
+      f' https://console.cloud.google.com/kubernetes/service/{get_cluster_location(args.project, args.cluster, args.zone)}/{args.cluster}/default/{args.workload}/details?project={args.project}'
   )
   status_cmd = (
       f'kubectl get jobset {args.workload} -o'
       " jsonpath='{.status.conditions[-1].type}'"
   )
   return_code, return_value = run_command_for_value(
-      status_cmd, 'Get jobset status', args
+      status_cmd, 'Get jobset status'
   )
   if return_code != 0:
     xpk_print(f'Get workload status request returned ERROR {return_code}')

xpk/main.py

@@ -36,7 +36,6 @@ import sys
 
 from .parser.core import set_parser
 from .utils.console import xpk_print
-from .utils.validation import validate_dependencies
 from .utils.execution_context import set_dry_run
 ################### Compatibility Check ###################
 # Check that the user runs the below version or greater.
@@ -66,9 +65,8 @@ def main() -> None:
   xpk_print('Starting xpk', flush=True)
   main_args = parser.parse_args()
   main_args.enable_ray_cluster = False
-  set_dry_run('dry_run' in main_args and main_args.dry_run)
-  if not main_args.dry_run:
-    validate_dependencies()
+  dry_run = 'dry_run' in main_args and main_args.dry_run
+  set_dry_run(dry_run)
   main_args.func(main_args)
   xpk_print('XPK Done.', flush=True)
 

xpk/parser/cluster.py

@@ -31,6 +31,7 @@ from ..core.config import CFG_BUCKET_KEY
 from ..core.vertex import DEFAULT_VERTEX_TENSORBOARD_NAME
 from .common import add_shared_arguments, ParserOrArgumentGroup
 from .validators import name_type
+from ..utils.feature_flags import FeatureFlags
 
 
 def set_cluster_parser(cluster_parser: ArgumentParser):
@@ -142,6 +143,12 @@ def set_cluster_create_parser(cluster_create_parser: ArgumentParser):
           ' enable cluster to accept Pathways workloads.'
       ),
   )
+  if FeatureFlags.SUB_SLICING_ENABLED:
+    cluster_create_optional_arguments.add_argument(
+        '--sub-slicing',
+        action='store_true',
+        help='Whether to set up cluster to support sub-slicing',
+    )
 
   autoprovisioning_arguments = cluster_create_parser.add_argument_group(
       'Autoprovisioning Arguments',

xpk/parser/cluster_test.py

@@ -0,0 +1,66 @@
+"""
+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 argparse
+from xpk.parser.cluster import set_cluster_create_parser
+import pytest
+from ..utils.feature_flags import FeatureFlags
+
+
+@pytest.fixture(autouse=True)
+def with_sub_slicing_enabled():
+  FeatureFlags.SUB_SLICING_ENABLED = True
+
+
+def test_cluster_create_sub_slicing_is_hidden_with_flag_off():
+  FeatureFlags.SUB_SLICING_ENABLED = False
+  parser = argparse.ArgumentParser()
+
+  set_cluster_create_parser(parser)
+  help_str = parser.format_help()
+
+  assert "--sub-slicing" not in help_str
+
+
+def test_cluster_create_sub_slicing_is_shown_with_flag_on():
+  parser = argparse.ArgumentParser()
+
+  set_cluster_create_parser(parser)
+  help_str = parser.format_help()
+
+  assert "--sub-slicing" in help_str
+
+
+def test_cluster_create_sub_slicing_is_false_by_default():
+  parser = argparse.ArgumentParser()
+
+  set_cluster_create_parser(parser)
+  args = parser.parse_args(
+      ["--cluster", "test-cluster", "--tpu-type", "test-tpu"]
+  )
+
+  assert args.sub_slicing is False
+
+
+def test_cluster_create_sub_slicing_can_be_set():
+  parser = argparse.ArgumentParser()
+
+  set_cluster_create_parser(parser)
+  args = parser.parse_args(
+      ["--cluster", "test-cluster", "--tpu-type", "test-tpu", "--sub-slicing"]
+  )
+
+  assert args.sub_slicing is True

xpk/parser/common.py

@@ -62,6 +62,17 @@ def add_shared_arguments(
       ),
       required=required,
   )
+  custom_parser_or_group.add_argument(
+      '--skip-validation',
+      type=bool,
+      action=argparse.BooleanOptionalAction,
+      default=False,
+      help=(
+          'Skip dependency validation checks (kubectl, gcloud, docker, etc). '
+          'Independent of --dry-run.'
+      ),
+      required=required,
+  )
 
 
 def add_cluster_arguments(

xpk/parser/workload.py

@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+from argparse import ArgumentParser
 from ..commands.workload import (
     workload_create,
     workload_create_pathways,
@@ -23,9 +24,10 @@ from ..commands.workload import (
 from ..core.docker_image import DEFAULT_DOCKER_IMAGE, DEFAULT_SCRIPT_DIR
 from .common import add_shared_arguments
 from .validators import directory_path_type, name_type
+from ..utils.feature_flags import FeatureFlags
 
 
-def set_workload_parsers(workload_parser):
+def set_workload_parsers(workload_parser: ArgumentParser):
   workload_subcommands = workload_parser.add_subparsers(
       title='workload subcommands',
       dest='xpk_workload_subcommands',
@@ -39,6 +41,28 @@ def set_workload_parsers(workload_parser):
   workload_create_parser = workload_subcommands.add_parser(
       'create', help='Create a new job.'
   )
+  set_workload_create_parser(workload_create_parser)
+
+  # "workload create-pathways" command parser.
+  workload_create_pathways_parser = workload_subcommands.add_parser(
+      'create-pathways', help='Create a new job.'
+  )
+  set_workload_create_pathways_parser(workload_create_pathways_parser)
+
+  # "workload delete" command parser.
+  workload_delete_parser = workload_subcommands.add_parser(
+      'delete', help='Delete job.'
+  )
+  set_workload_delete_parser(workload_delete_parser)
+
+  # "workload list" command parser.
+  workload_list_parser = workload_subcommands.add_parser(
+      'list', help='List jobs.'
+  )
+  set_workload_list_parser(workload_list_parser)
+
+
+def set_workload_create_parser(workload_create_parser: ArgumentParser):
   workload_create_parser_required_arguments = (
       workload_create_parser.add_argument_group(
           'Workload Built-in Arguments',
@@ -193,10 +217,33 @@ def set_workload_parsers(workload_parser):
       ),
   )
 
-  # "workload create-pathways" command parser.
-  workload_create_pathways_parser = workload_subcommands.add_parser(
-      'create-pathways', help='Create a new job.'
-  )
+  add_shared_workload_create_required_arguments([
+      workload_create_parser_required_arguments,
+  ])
+  add_shared_workload_create_optional_arguments([
+      workload_create_parser_optional_arguments,
+  ])
+  add_shared_workload_create_env_arguments([
+      workload_create_parser_optional_arguments,
+  ])
+  add_shared_workload_base_docker_image_arguments([
+      workload_base_docker_image_arguments,
+  ])
+  add_shared_workload_docker_image_arguments([
+      workload_docker_image_arguments,
+  ])
+  add_shared_workload_create_tensorboard_arguments([
+      workload_vertex_tensorboard_arguments,
+  ])
+  add_shared_workload_create_autoprovisioning_arguments([
+      workload_create_autoprovisioning_arguments,
+  ])
+  workload_create_parser.set_defaults(func=workload_create)
+
+
+def set_workload_create_pathways_parser(
+    workload_create_pathways_parser: ArgumentParser,
+):
   workload_create_pathways_parser_required_arguments = (
       workload_create_pathways_parser.add_argument_group(
           'Workload create-pathways Built-in Arguments',
@@ -232,7 +279,6 @@ def set_workload_parsers(workload_parser):
           'Arguments for creating Vertex AI Experiment in workload create.',
       )
   )
-
   ### "workload create-pathways" Required arguments, specific to Pathways
   workload_create_pathways_parser_required_arguments.add_argument(
       '--tpu-type',
@@ -353,42 +399,30 @@ def set_workload_parsers(workload_parser):
   )
 
   add_shared_workload_create_required_arguments([
-      workload_create_parser_required_arguments,
       workload_create_pathways_parser_required_arguments,
   ])
   add_shared_workload_create_optional_arguments([
-      workload_create_parser_optional_arguments,
       workload_create_pathways_parser_optional_arguments,
   ])
   add_shared_workload_create_env_arguments([
-      workload_create_parser_optional_arguments,
       workload_create_pathways_parser_optional_arguments,
   ])
   add_shared_workload_base_docker_image_arguments([
-      workload_base_docker_image_arguments,
       workload_create_pathways_base_docker_image_arguments,
   ])
   add_shared_workload_docker_image_arguments([
-      workload_docker_image_arguments,
       workload_create_pathways_docker_image_arguments,
   ])
   add_shared_workload_create_tensorboard_arguments([
-      workload_vertex_tensorboard_arguments,
       workload_create_pathways_vertex_tensorboard_arguments,
   ])
   add_shared_workload_create_autoprovisioning_arguments([
-      workload_create_autoprovisioning_arguments,
       workload_create_pathways_autoprovisioning_arguments,
   ])
-
-  # Set defaults for both workload create and workload create-pathways after adding all shared args.
-  workload_create_parser.set_defaults(func=workload_create)
   workload_create_pathways_parser.set_defaults(func=workload_create_pathways)
 
-  # "workload delete" command parser.
-  workload_delete_parser = workload_subcommands.add_parser(
-      'delete', help='Delete job.'
-  )
+
+def set_workload_delete_parser(workload_delete_parser: ArgumentParser):
   workload_delete_parser_required_arguments = (
       workload_delete_parser.add_argument_group(
           'Required Arguments',
@@ -454,14 +488,10 @@ def set_workload_parsers(workload_parser):
           'Forces workload deletion command to run without additional approval.'
       ),
   )
-
   workload_delete_parser.set_defaults(func=workload_delete)
 
-  # "workload list" command parser.
-  workload_list_parser = workload_subcommands.add_parser(
-      'list', help='List jobs.'
-  )
 
+def set_workload_list_parser(workload_list_parser: ArgumentParser):
   workload_list_parser.add_argument(
       '--cluster',
       type=name_type,
@@ -629,6 +659,13 @@ def add_shared_workload_create_optional_arguments(args_parsers):
             ' the workload.'
         ),
     )
+    if FeatureFlags.SUB_SLICING_ENABLED:
+      custom_parser.add_argument(
+          '--sub-slicing-topology',
+          type=str,
+          help='Sub-slicing topology to use.',
+          required=False,
+      )
 
 
 def add_shared_workload_create_env_arguments(args_parsers):

xpk/parser/workload_test.py

@@ -0,0 +1,82 @@
+"""
+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 argparse
+from xpk.parser.workload import set_workload_create_parser
+from ..utils.feature_flags import FeatureFlags
+import pytest
+
+
+@pytest.fixture(autouse=True)
+def with_sub_slicing_enabled():
+  FeatureFlags.SUB_SLICING_ENABLED = True
+
+
+def test_workload_create_sub_slicing_topology_is_hidden_with_flag_off():
+  FeatureFlags.SUB_SLICING_ENABLED = False
+  parser = argparse.ArgumentParser()
+
+  set_workload_create_parser(parser)
+  help_str = parser.format_help()
+
+  assert "--sub-slicing" not in help_str
+
+
+def test_workload_create_sub_slicing_topology_is_shown_with_flag_on():
+  parser = argparse.ArgumentParser()
+
+  set_workload_create_parser(parser)
+  help_str = parser.format_help()
+
+  assert "--sub-slicing" in help_str
+
+
+def test_workload_create_sub_slicing_topology_is_none_by_default():
+  parser = argparse.ArgumentParser()
+
+  set_workload_create_parser(parser)
+  args = parser.parse_args([
+      "--cluster",
+      "test-cluster",
+      "--command",
+      "python3",
+      "--workload",
+      "test",
+      "--tpu-type",
+      "test-tpu",
+  ])
+
+  assert args.sub_slicing_topology is None
+
+
+def test_workload_create_sub_slicing_topology_can_be_set():
+  parser = argparse.ArgumentParser()
+
+  set_workload_create_parser(parser)
+  args = parser.parse_args([
+      "--cluster",
+      "test-cluster",
+      "--command",
+      "python3",
+      "--workload",
+      "test",
+      "--tpu-type",
+      "test-tpu",
+      "--sub-slicing-topology",
+      "2x2",
+  ])
+
+  assert args.sub_slicing_topology is "2x2"

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