xpk 0.14.4__py3-none-any.whl → 0.16.0__py3-none-any.whl

xpk/main.py

@@ -32,11 +32,14 @@ Next Steps:
 """
 
 import argparse
+import argcomplete
 import sys
 
 from .parser.core import set_parser
 from .core.updates import print_xpk_hello
-from .utils.console import xpk_print
+from .core.config import set_config, FileSystemConfig
+from .core.telemetry import MetricsCollector, send_clearcut_payload, should_send_telemetry
+from .utils.console import xpk_print, exit_code_to_int
 from .utils.execution_context import set_context
 ################### Compatibility Check ###################
 # Check that the user runs the below version or greater.
@@ -59,19 +62,36 @@ if (
 
 
 def main() -> None:
-  # Create top level parser for xpk command.
-  parser = argparse.ArgumentParser(description='xpk command', prog='xpk')
-  set_parser(parser=parser)
+  try:
+    # Create top level parser for xpk command.
+    parser = argparse.ArgumentParser(description='xpk command', prog='xpk')
+    set_parser(parser=parser)
+    argcomplete.autocomplete(parser)
 
-  main_args = parser.parse_args()
-  main_args.enable_ray_cluster = False
-  set_context(
-      dry_run_value='dry_run' in main_args and main_args.dry_run,
-      quiet_value='quiet' in main_args and main_args.quiet,
-  )
-  print_xpk_hello()
-  main_args.func(main_args)
-  xpk_print('XPK Done.', flush=True)
+    main_args = parser.parse_args()
+    main_args.enable_ray_cluster = False
+    set_config(FileSystemConfig())
+    set_context(
+        dry_run_value='dry_run' in main_args and main_args.dry_run,
+        quiet_value=(
+            ('quiet' in main_args and main_args.quiet)
+            or ('force' in main_args and main_args.force)
+        ),
+    )
+    MetricsCollector.log_start(main_args.xpk_subcommands)
+    print_xpk_hello()
+    main_args.func(main_args)
+    xpk_print('XPK Done.', flush=True)
+    MetricsCollector.log_complete(0)
+  except SystemExit as e:
+    MetricsCollector.log_complete(exit_code_to_int(e.code))
+    raise
+  except:
+    MetricsCollector.log_complete(-1)
+    raise
+  finally:
+    if should_send_telemetry():
+      send_clearcut_payload(MetricsCollector.flush())
 
 
 if __name__ == '__main__':

xpk/parser/cluster.py

@@ -26,10 +26,10 @@ from ..commands.cluster import (
     cluster_describe,
     cluster_list,
 )
-from ..commands.config import xpk_cfg
+from ..core.config import get_config
 from ..core.config import CFG_BUCKET_KEY
 from ..core.vertex import DEFAULT_VERTEX_TENSORBOARD_NAME
-from .common import add_shared_arguments, ParserOrArgumentGroup
+from .common import add_shared_arguments, ParserOrArgumentGroup, add_tpu_type_argument, add_tpu_and_device_type_arguments
 from .validators import name_type
 from ..utils.feature_flags import FeatureFlags
 
@@ -98,21 +98,7 @@ def set_cluster_create_parser(cluster_create_parser: ArgumentParser):
           required=True
       )
   )
-  cluster_device_group.add_argument(
-      '--tpu-type',
-      type=str,
-      default=None,
-      help='The tpu type to use, v5litepod-16, etc.',
-  )
-  cluster_device_group.add_argument(
-      '--device-type',
-      type=str,
-      default=None,
-      help=(
-          'The device type to use (can be tpu or gpu or cpu), v5litepod-16,'
-          ' h100-80gb-8, n2-standard-32-4 etc.'
-      ),
-  )
+  add_tpu_and_device_type_arguments(cluster_device_group)
 
   ### Optional arguments specific to "cluster create"
   cluster_create_optional_arguments = cluster_create_parser.add_argument_group(
@@ -124,7 +110,7 @@ def set_cluster_create_parser(cluster_create_parser: ArgumentParser):
   cluster_create_optional_arguments.add_argument(
       '--cluster-state-gcs-bucket',
       type=str,
-      default=xpk_cfg.get(CFG_BUCKET_KEY),
+      default=get_config().get(CFG_BUCKET_KEY),
       help='The name of the bucket to store cluster state.',
       required=False,
   )
@@ -143,12 +129,9 @@ 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',
-    )
+    add_cluster_create_sub_slicing_arguments(cluster_create_optional_arguments)
 
   autoprovisioning_arguments = cluster_create_parser.add_argument_group(
       'Autoprovisioning Arguments',
@@ -204,11 +187,8 @@ def set_cluster_create_pathways_parser(
   add_shared_cluster_create_required_arguments(
       cluster_create_pathways_required_arguments
   )
-  cluster_create_pathways_required_arguments.add_argument(
-      '--tpu-type',
-      type=str,
-      default=None,
-      help='The tpu type to use, v5litepod-16, etc.',
+  add_tpu_type_argument(
+      cluster_create_pathways_required_arguments, required=True
   )
 
   ### Optional arguments specific to "cluster create-pathways"
@@ -221,6 +201,10 @@ def set_cluster_create_pathways_parser(
   add_shared_cluster_create_optional_arguments(
       cluster_create_pathways_optional_arguments
   )
+  if FeatureFlags.SUB_SLICING_ENABLED:
+    add_cluster_create_sub_slicing_arguments(
+        cluster_create_pathways_optional_arguments
+    )
 
   autoprovisioning_arguments = (
       cluster_create_pathways_parser.add_argument_group(
@@ -281,13 +265,8 @@ def set_cluster_create_ray_parser(cluster_create_ray_parser: ArgumentParser):
   add_shared_cluster_create_required_arguments(
       cluster_create_ray_required_arguments
   )
-  cluster_create_ray_required_arguments.add_argument(
-      '--tpu-type',
-      type=str,
-      default=None,
-      help='The tpu type to use, v5litepod-16, etc.',
-      required=True,
-  )
+  add_tpu_type_argument(cluster_create_ray_required_arguments, required=True)
+
   # TODO(bzmarke): Add --device-type to support GPU/CPU
   cluster_create_ray_required_arguments.add_argument(
       '--ray-version',
@@ -350,7 +329,9 @@ def set_cluster_create_ray_parser(cluster_create_ray_parser: ArgumentParser):
   )
   add_resource_limits(cluster_create_resource_limits)
 
-  cluster_create_ray_parser.set_defaults(func=cluster_create_ray_cluster)
+  cluster_create_ray_parser.set_defaults(
+      func=cluster_create_ray_cluster, sub_slicing=False
+  )
 
 
 def set_cluster_delete_parser(cluster_delete_parser: ArgumentParser):
@@ -375,7 +356,7 @@ def set_cluster_delete_parser(cluster_delete_parser: ArgumentParser):
   cluster_delete_optional_arguments.add_argument(
       '--cluster-state-gcs-bucket',
       type=str,
-      default=xpk_cfg.get(CFG_BUCKET_KEY),
+      default=get_config().get(CFG_BUCKET_KEY),
       help='The name of the bucket to store cluster state.',
       required=False,
   )
@@ -404,21 +385,7 @@ def set_cluster_cacheimage_parser(cluster_cacheimage_parser: ArgumentParser):
   )
 
   ### Device Type Argument
-  cluster_cacheimage_group.add_argument(
-      '--tpu-type',
-      type=str,
-      default=None,
-      help='The tpu type to cache images on, v5litepod-16, etc.',
-  )
-  cluster_cacheimage_group.add_argument(
-      '--device-type',
-      type=str,
-      default=None,
-      help=(
-          'The device type to cache images on (can be tpu or gpu),'
-          ' v5litepod-16, h100-80gb-8, etc.'
-      ),
-  )
+  add_tpu_and_device_type_arguments(cluster_cacheimage_group)
 
   ### Required arguments
   cluster_cacheimage_required_arguments.add_argument(
@@ -503,21 +470,7 @@ def set_cluster_adapt_parser(cluster_adapt_parser: ArgumentParser):
           required=True
       )
   )
-  cluster_adapt_device_group.add_argument(
-      '--tpu-type',
-      type=str,
-      default=None,
-      help='The tpu type used on cluster, v5litepod-16, etc.',
-  )
-  cluster_adapt_device_group.add_argument(
-      '--device-type',
-      type=str,
-      default=None,
-      help=(
-          'The device type used on cluster (can be tpu or gpu or cpu), eg.'
-          ' h100-80gb-8, n2-standard-32-4 etc.'
-      ),
-  )
+  add_tpu_and_device_type_arguments(cluster_adapt_device_group)
 
   cluster_adapt_optional_arguments = cluster_adapt_parser.add_argument_group(
       'Optional Arguments',
@@ -662,6 +615,11 @@ def add_shared_cluster_create_optional_arguments(
           ' regional clusters, all zones must support the machine type.'
       ),
   )
+  parser_or_group.add_argument(
+      '--managed-mldiagnostics',
+      action='store_true',
+      help='Enables the installation of required ML Diagnostics components.',
+  )
   parser_or_group.add_argument(
       '--cluster-cpu-machine-type',
       type=str,
@@ -790,6 +748,11 @@ def add_driver_arguments(parser_or_group: ParserOrArgumentGroup):
       action='store_true',
       help='Enable Lustre CSI driver on the cluster.',
   )
+  parser_or_group.add_argument(
+      '--enable-legacy-lustre-port',
+      action='store_true',
+      help='Enable legacy port for Lustre CSI driver on the cluster.',
+  )
 
 
 def add_shared_cluster_create_tensorboard_arguments(
@@ -937,3 +900,13 @@ def add_resource_limits(parser_or_group: ParserOrArgumentGroup):
       default=None,
       help='The CPU limit for the Kueue controller manager.',
   )
+
+
+def add_cluster_create_sub_slicing_arguments(
+    parser_or_group: ParserOrArgumentGroup,
+):
+  parser_or_group.add_argument(
+      '--sub-slicing',
+      action='store_true',
+      help='Whether to set up cluster to support sub-slicing',
+  )

xpk/parser/cluster_test.py

@@ -15,7 +15,7 @@ limitations under the License.
 """
 
 import argparse
-from xpk.parser.cluster import set_cluster_create_parser
+from xpk.parser.cluster import set_cluster_create_parser, set_cluster_create_pathways_parser, set_cluster_create_ray_parser
 import pytest
 from ..utils.feature_flags import FeatureFlags
 
@@ -49,7 +49,7 @@ def test_cluster_create_sub_slicing_is_false_by_default():
 
   set_cluster_create_parser(parser)
   args = parser.parse_args(
-      ["--cluster", "test-cluster", "--tpu-type", "test-tpu"]
+      ["--cluster", "test-cluster", "--tpu-type", "tpu7x-2"]
   )
 
   assert args.sub_slicing is False
@@ -60,7 +60,87 @@ def test_cluster_create_sub_slicing_can_be_set():
 
   set_cluster_create_parser(parser)
   args = parser.parse_args(
-      ["--cluster", "test-cluster", "--tpu-type", "test-tpu", "--sub-slicing"]
+      ["--cluster", "test-cluster", "--tpu-type", "tpu7x-2", "--sub-slicing"]
   )
 
   assert args.sub_slicing is True
+
+
+def test_cluster_create_pathways_sub_slicing_is_hidden_with_flag_off():
+  FeatureFlags.SUB_SLICING_ENABLED = False
+  parser = argparse.ArgumentParser()
+
+  set_cluster_create_pathways_parser(parser)
+  help_str = parser.format_help()
+
+  assert "--sub-slicing" not in help_str
+
+
+def test_cluster_create_pathways_sub_slicing_can_be_set():
+  parser = argparse.ArgumentParser()
+
+  set_cluster_create_pathways_parser(parser)
+  args = parser.parse_args(
+      ["--cluster", "test-cluster", "--tpu-type", "tpu7x-2", "--sub-slicing"]
+  )
+
+  assert args.sub_slicing is True
+
+
+def test_cluster_create_ray_sub_slicing_is_hidden_but_set_to_false():
+  parser = argparse.ArgumentParser()
+
+  set_cluster_create_ray_parser(parser)
+  args = parser.parse_args([
+      "--cluster",
+      "test-cluster",
+      "--tpu-type",
+      "tpu7x-2",
+      "--ray-version",
+      "19.32.0",
+  ])
+  help_str = parser.format_help()
+
+  assert args.sub_slicing is False
+  assert "--sub-slicing" not in help_str
+
+
+def test_cluster_create_managed_mldiagnostics():
+  parser = argparse.ArgumentParser()
+
+  set_cluster_create_parser(parser)
+  args = parser.parse_args([
+      "--cluster",
+      "test-cluster",
+      "--tpu-type",
+      "v5p-8",
+      "--managed-mldiagnostics",
+  ])
+
+  assert args.managed_mldiagnostics is True
+
+
+def test_cluster_create_enable_lustre_legacy_port_is_false_by_default():
+  parser = argparse.ArgumentParser()
+
+  set_cluster_create_parser(parser)
+  args = parser.parse_args(
+      ["--cluster", "test-cluster", "--tpu-type", "tpu7x-2"]
+  )
+
+  assert args.enable_legacy_lustre_port is False
+
+
+def test_cluster_create_enable_lustre_legacy_port_can_be_set():
+  parser = argparse.ArgumentParser()
+
+  set_cluster_create_parser(parser)
+  args = parser.parse_args([
+      "--cluster",
+      "test-cluster",
+      "--tpu-type",
+      "tpu7x-2",
+      "--enable-legacy-lustre-port",
+  ])
+
+  assert args.enable_legacy_lustre_port is True

xpk/parser/common.py

@@ -16,6 +16,10 @@ limitations under the License.
 
 import argparse
 from typing import Protocol, Any
+from ..core.system_characteristics import get_system_characteristics_keys_by_accelerator_type, AcceleratorType
+import difflib
+from argcomplete import ChoicesCompleter
+from argparse import Action, ArgumentError
 
 
 class ParserOrArgumentGroup(Protocol):
@@ -24,6 +28,46 @@ class ParserOrArgumentGroup(Protocol):
     ...
 
 
+class ManyChoicesAction(Action):
+  """An action class to output better error message for arguments with large lists of choices."""
+
+  def __init__(self, *args, large_choice_list, **kwargs):
+    self.large_list_of_choices = large_choice_list
+    super().__init__(*args, **kwargs)
+
+  def __call__(self, parser, namespace, value, option_string=None):
+    if value not in self.large_list_of_choices:
+      close_matches = difflib.get_close_matches(
+          value, self.large_list_of_choices, n=5, cutoff=0
+      )
+      msg = (
+          f"invalid choice: '{value}' (closest matches:"
+          f" {', '.join(close_matches)})"
+      )
+      raise ArgumentError(self, msg)
+    setattr(namespace, self.dest, value)
+
+
+def add_many_choices_argument(
+    parserOrGroup: ParserOrArgumentGroup,
+    flag_name,
+    choices: list[str],
+    metavar: str,
+    help_msg: str,
+    required: bool = False,
+) -> None:
+  parserOrGroup.add_argument(
+      flag_name,
+      action=ManyChoicesAction,
+      large_choice_list=choices,
+      type=str,
+      metavar=metavar,
+      help=help_msg,
+      required=required,
+      default=None,
+  ).completer = ChoicesCompleter(choices)
+
+
 def add_shared_arguments(
     custom_parser_or_group: ParserOrArgumentGroup, required=False
 ) -> None:
@@ -285,3 +329,43 @@ def add_slurm_arguments(custom_parser_or_group: ParserOrArgumentGroup):
           ' `very-high`. Defaults to `medium`.'
       ),
   )
+
+
+def add_tpu_type_argument(
+    custom_parser_or_group: ParserOrArgumentGroup,
+    required: bool = False,
+) -> None:
+  add_many_choices_argument(
+      custom_parser_or_group,
+      '--tpu-type',
+      choices=get_system_characteristics_keys_by_accelerator_type(
+          [AcceleratorType.TPU]
+      ),
+      metavar='TPU_TYPE',
+      help_msg='The tpu type to use, v5litepod-16, etc.',
+      required=required,
+  )
+
+
+def add_device_type_argument(
+    custom_parser_or_group: ParserOrArgumentGroup,
+    required: bool = False,
+) -> None:
+  add_many_choices_argument(
+      custom_parser_or_group,
+      '--device-type',
+      choices=get_system_characteristics_keys_by_accelerator_type(),
+      metavar='DEVICE_TYPE',
+      help_msg=(
+          'The device type to use (can be tpu or gpu or cpu), v5litepod-16,'
+          ' h100-80gb-8, n2-standard-32-4 etc.'
+      ),
+      required=required,
+  )
+
+
+def add_tpu_and_device_type_arguments(
+    custom_parser_or_group: ParserOrArgumentGroup,
+) -> None:
+  add_tpu_type_argument(custom_parser_or_group)
+  add_device_type_argument(custom_parser_or_group)

xpk/parser/storage.py

@@ -104,6 +104,16 @@ def add_storage_attach_parser(
       help='If true workloads can only read from storage',
   )
 
+  lustre_args = storage_attach_parser.add_argument_group(
+      'Lustre arguments',
+      'Arguments used when --type=lustre',
+  )
+  lustre_args.add_argument(
+      '--enable-legacy-lustre-port',
+      action='store_true',
+      help='Enable legacy port for Lustre CSI driver on the cluster.',
+  )
+
   gcsfuse_args = storage_attach_parser.add_argument_group(
       'FUSE arguments',
       'Arguments used when --type=gcsfuse',

xpk/parser/storage_test.py

@@ -0,0 +1,47 @@
+"""
+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.storage import set_storage_parser
+
+DEFAULT_ATTACH_ARGUMENTS = (
+    "attach test-storage --cluster test-cluster --zone test-zone"
+    " --project test-project --mount-point test-mount-point"
+    " --readonly false --auto-mount true"
+)
+
+DEFAULT_LUSTRE_ATTACH_ARGUMENTS = (
+    DEFAULT_ATTACH_ARGUMENTS + " --type lustre --manifest test-manifest"
+)
+
+
+def test_cluster_create_enable_lustre_legacy_port_is_false_by_default():
+  parser = argparse.ArgumentParser()
+
+  set_storage_parser(parser)
+  args = parser.parse_args(DEFAULT_LUSTRE_ATTACH_ARGUMENTS.split())
+
+  assert args.enable_legacy_lustre_port is False
+
+
+def test_cluster_create_enable_lustre_legacy_port_can_be_set():
+  parser = argparse.ArgumentParser()
+  set_storage_parser(parser)
+  args = parser.parse_args(
+      DEFAULT_LUSTRE_ATTACH_ARGUMENTS.split() + ["--enable-legacy-lustre-port"]
+  )
+
+  assert args.enable_legacy_lustre_port is True

xpk/parser/workload.py

@@ -22,9 +22,8 @@ from ..commands.workload import (
     workload_list,
 )
 from ..core.docker_image import DEFAULT_DOCKER_IMAGE, DEFAULT_SCRIPT_DIR
-from .common import add_shared_arguments
+from .common import add_shared_arguments, add_tpu_type_argument, add_tpu_and_device_type_arguments
 from .validators import directory_path_type, name_type
-from ..utils.feature_flags import FeatureFlags
 
 
 def set_workload_parsers(workload_parser: ArgumentParser):
@@ -118,21 +117,7 @@ def set_workload_create_parser(workload_create_parser: ArgumentParser):
           required=True
       )
   )
-  workload_device_group.add_argument(
-      '--tpu-type',
-      type=str,
-      default=None,
-      help='The tpu type to use, v5litepod-16, etc.',
-  )
-  workload_device_group.add_argument(
-      '--device-type',
-      type=str,
-      default=None,
-      help=(
-          'The device type to use (can be tpu or gpu or cpu), v5litepod-16,'
-          ' h100-80gb-8, n2-standard-32-4 etc.'
-      ),
-  )
+  add_tpu_and_device_type_arguments(workload_device_group)
 
   workload_create_parser_optional_arguments.add_argument(
       '--storage',
@@ -280,11 +265,8 @@ def set_workload_create_pathways_parser(
       )
   )
   ### "workload create-pathways" Required arguments, specific to Pathways
-  workload_create_pathways_parser_required_arguments.add_argument(
-      '--tpu-type',
-      type=str,
-      default=None,
-      help='The tpu type to use, v5litepod-16, etc.',
+  add_tpu_type_argument(
+      workload_create_pathways_parser_required_arguments, required=True
   )
 
   ### "workload create-pathways" Optional arguments, specific to Pathways
@@ -601,6 +583,16 @@ def add_shared_workload_create_optional_arguments(args_parsers):
             ' `jax-tpu`.'
         ),
     )
+    custom_parser.add_argument(
+        '--output-manifest-file',
+        type=str,
+        default=None,
+        help=(
+            'If you want to see the generated manifest, provide a file path'
+            ' here. This will write the manifest to the file. If used with'
+            ' --dry-run, it will skip the actual deployment and cluster checks.'
+        ),
+    )
     custom_parser.add_argument(
         '--num-slices',
         type=int,
@@ -659,13 +651,6 @@ 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

@@ -16,53 +16,9 @@ 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():
+def test_workload_create_parses():
   parser = argparse.ArgumentParser()
 
   set_workload_create_parser(parser)
@@ -74,9 +30,7 @@ def test_workload_create_sub_slicing_topology_can_be_set():
       "--workload",
       "test",
       "--tpu-type",
-      "test-tpu",
-      "--sub-slicing-topology",
-      "2x2",
+      "tpu7x-2",
   ])
 
-  assert args.sub_slicing_topology is "2x2"
+  assert args