xpk 0.12.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,7 @@ 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.
 
@@ -63,9 +63,10 @@ def main() -> None:
   set_parser(parser=parser)
 
   xpk_print('Starting xpk', flush=True)
-  validate_dependencies()
   main_args = parser.parse_args()
   main_args.enable_ray_cluster = False
+  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',
@@ -174,6 +181,13 @@ def set_cluster_create_parser(cluster_create_parser: ArgumentParser):
       'Arguments for configuring MTC in cluster create.',
   )
   add_shared_cluster_create_mtc_arguments(cluster_create_mtc_arguments)
+
+  cluster_create_resource_limits = cluster_create_parser.add_argument_group(
+      'Optional Resource Limits Arguments',
+      'Arguments for configuring resource limits in cluster create.',
+  )
+  add_resource_limits(cluster_create_resource_limits)
+
   cluster_create_parser.set_defaults(func=cluster_create)
 
 
@@ -245,6 +259,15 @@ def set_cluster_create_pathways_parser(
       )
   )
   add_shared_cluster_create_mtc_arguments(cluster_create_mtc_arguments)
+
+  cluster_create_resource_limits = (
+      cluster_create_pathways_parser.add_argument_group(
+          'Optional Resource Limits Arguments',
+          'Arguments for configuring resource limits in cluster create.',
+      )
+  )
+  add_resource_limits(cluster_create_resource_limits)
+
   cluster_create_pathways_parser.set_defaults(func=cluster_create_pathways)
 
 
@@ -320,6 +343,13 @@ def set_cluster_create_ray_parser(cluster_create_ray_parser: ArgumentParser):
       'Arguments for configuring MTC in cluster create.',
   )
   add_shared_cluster_create_mtc_arguments(cluster_create_mtc_arguments)
+
+  cluster_create_resource_limits = cluster_create_ray_parser.add_argument_group(
+      'Optional Resource Limits Arguments',
+      'Arguments for configuring resource limits in cluster create.',
+  )
+  add_resource_limits(cluster_create_resource_limits)
+
   cluster_create_ray_parser.set_defaults(func=cluster_create_ray_cluster)
 
 
@@ -887,3 +917,23 @@ def add_shared_cluster_create_mtc_arguments(
           ' checkpointing. By default, it is set to "google.com/tpu".'
       ),
   )
+
+
+def add_resource_limits(parser_or_group: ParserOrArgumentGroup):
+  """Add resource limits arguments in cluster create.
+
+  Args:
+      List of cluster create resource limits arguments parsers or group
+  """
+  parser_or_group.add_argument(
+      '--memory-limit',
+      type=str,
+      default=None,
+      help='The memory limit for the Kueue controller manager.',
+  )
+  parser_or_group.add_argument(
+      '--cpu-limit',
+      type=int,
+      default=None,
+      help='The CPU limit for the Kueue controller manager.',
+  )

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/execution_context.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.
+"""
+
+dry_run = False
+
+
+def set_dry_run(value: bool) -> None:
+  """Sets the dry_run flag."""
+  global dry_run
+  dry_run = value
+
+
+def is_dry_run() -> bool:
+  """Returns the current value of the dry_run flag."""
+  return dry_run

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

@@ -16,10 +16,11 @@ limitations under the License.
 
 import tempfile
 import os
-from .console import xpk_print
+import hashlib
+from .execution_context import is_dry_run
 
 
-def make_tmp_files(per_command_name):
+def make_tmp_files(per_command_name: list[str]) -> list[str]:
   """Make temporary files for each command.
 
   Args:
@@ -28,16 +29,19 @@ def make_tmp_files(per_command_name):
   Returns:
     A list of temporary files for each command.
   """
+  if is_dry_run():
+    return [_hash_filename(command) for command in per_command_name]
+
   # Supports removal of spaces from command names before converting to file name.
   return [
       tempfile.NamedTemporaryFile(
           delete=False, prefix=command.replace(' ', '-') + '-'
-      )
+      ).file.name
       for command in per_command_name
   ]
 
 
-def write_tmp_file(payload):
+def write_tmp_file(payload: str) -> str:
   """Writes `payload` to a temporary file.
 
   Args:
@@ -46,14 +50,17 @@ def write_tmp_file(payload):
   Returns:
     A file object that was written to.
   """
+  if is_dry_run():
+    return _hash_filename(payload)
+
   with tempfile.NamedTemporaryFile(delete=False) as tmp:
     with open(file=tmp.name, mode='w', encoding='utf=8') as f:
       f.write(payload)
       f.flush()
-    return tmp
+    return tmp.file.name
 
 
-def append_tmp_file(payload, file):
+def append_tmp_file(payload: str, file: str) -> str:
   """Appends `payload` to an already created file.
 
   Use `write_temporary_file` to create a file.
@@ -65,18 +72,26 @@ def append_tmp_file(payload, file):
   Returns:
     A file object that was written to.
   """
-  with open(file=file.name, mode='a', encoding='utf=8') as f:
+  if is_dry_run():
+    return file
+
+  with open(file=file, mode='a', encoding='utf=8') as f:
     f.write(payload)
     f.flush()
   return file
 
 
-def ensure_directory_exists(directory_path):
+def ensure_directory_exists(directory_path: str) -> None:
   """Checks if a directory exists and creates it if it doesn't.
 
   Args:
     directory_path: The path to the directory.
   """
-  if not os.path.exists(directory_path):
+  if not is_dry_run() and not os.path.exists(directory_path):
     os.makedirs(directory_path)
-    xpk_print(f"Directory '{directory_path}' created successfully.")
+
+
+def _hash_filename(seed: str) -> str:
+  m = hashlib.sha256()
+  m.update(seed.encode('utf-8'))
+  return m.hexdigest()

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