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

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/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 = "src/xpk/templates/"
+
 yaml = ruamel.yaml.YAML()
 
 

xpk/utils/topology.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.
+"""
+
+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")]

xpk/utils/topology_test.py

@@ -0,0 +1,43 @@
+"""
+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
+
+
+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