xpk 0.10.1__py3-none-any.whl → 0.12.0__py3-none-any.whl

xpk/parser/cluster.py

@@ -29,7 +29,7 @@ from ..commands.cluster import (
 from ..commands.config import xpk_cfg
 from ..core.config import CFG_BUCKET_KEY
 from ..core.vertex import DEFAULT_VERTEX_TENSORBOARD_NAME
-from .common import add_shared_arguments
+from .common import add_shared_arguments, ParserOrArgumentGroup
 from .validators import name_type
 
 
@@ -208,6 +208,14 @@ def set_cluster_create_pathways_parser(
       cluster_create_pathways_optional_arguments
   )
 
+  autoprovisioning_arguments = (
+      cluster_create_pathways_parser.add_argument_group(
+          'Autoprovisioning Arguments',
+          'Optional arguments for enabling autoprovisioning.',
+      )
+  )
+  add_autoprovisioning_arguments(autoprovisioning_arguments)
+
   ### Capacity arguments specific to "cluster create-pathways"
   cluster_create_pathways_capacity_arguments = (
       cluster_create_pathways_parser.add_argument_group(
@@ -529,15 +537,15 @@ def set_cluster_adapt_parser(cluster_adapt_parser: ArgumentParser):
   cluster_adapt_parser.set_defaults(func=cluster_adapt)
 
 
-def add_autoprovisioning_arguments(parser: ArgumentParser):
-  parser.add_argument(
+def add_autoprovisioning_arguments(parser_or_group: ParserOrArgumentGroup):
+  parser_or_group.add_argument(
       '--enable-autoprovisioning',
       action='store_true',
       help=(
           'Enable GKE features for autoprovisioning node pools in GKE clusters.'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--autoprovisioning-min-chips',
       type=int,
       help=(
@@ -546,7 +554,7 @@ def add_autoprovisioning_arguments(parser: ArgumentParser):
           ' resources in the cluster as the minimum, and maximum.'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--autoprovisioning-max-chips',
       type=int,
       help=(
@@ -557,13 +565,15 @@ def add_autoprovisioning_arguments(parser: ArgumentParser):
   )
 
 
-def add_shared_cluster_create_required_arguments(parser: ArgumentParser):
+def add_shared_cluster_create_required_arguments(
+    parser_or_group: ParserOrArgumentGroup,
+):
   """Add shared required arguments in cluster create and Pathways cluster create.
 
   Args:
-    parser: cluster create argument parser or argument group
+    parser_or_group: cluster create argument parser or argument group
   """
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--cluster',
       type=name_type,
       default=None,
@@ -575,21 +585,23 @@ def add_shared_cluster_create_required_arguments(parser: ArgumentParser):
   )
 
 
-def add_shared_cluster_create_optional_arguments(parser: ArgumentParser):
+def add_shared_cluster_create_optional_arguments(
+    parser_or_group: ParserOrArgumentGroup,
+):
   """Add shared optional arguments in cluster create and Pathways cluster create.
 
   Args:
-    parser: cluster create argument parser or argument group
+    parser_or_group: cluster create argument parser or argument group
   """
-  add_shared_arguments(parser)
-  parser.add_argument(
+  add_shared_arguments(parser_or_group)
+  parser_or_group.add_argument(
       '--host-maintenance-interval',
       type=str,
       choices=['AS_NEEDED', 'PERIODIC'],
       default='AS_NEEDED',
       help='The maintenance policy of the cluster and respective clusters.',
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--gke-version',
       type=str,
       help=(
@@ -598,20 +610,20 @@ def add_shared_cluster_create_optional_arguments(parser: ArgumentParser):
           ' recommended version.'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--num-slices',
       type=int,
       default=1,
       help='The number of slices to run the job on, defaults to 1.',
       required=False,
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--pathways-gce-machine-type',
       type=str,
       default='n2-standard-64',
       help='The CPU type for Pathways CPU nodepools',
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--default-pool-cpu-machine-type',
       type=str,
       default='e2-standard-16',
@@ -620,7 +632,7 @@ def add_shared_cluster_create_optional_arguments(parser: ArgumentParser):
           ' regional clusters, all zones must support the machine type.'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--cluster-cpu-machine-type',
       type=str,
       default='',
@@ -631,7 +643,7 @@ def add_shared_cluster_create_optional_arguments(parser: ArgumentParser):
           ' cpu nodepools using --device-type.'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--default-pool-cpu-num-nodes',
       type=int,
       default=6,
@@ -641,7 +653,7 @@ def add_shared_cluster_create_optional_arguments(parser: ArgumentParser):
           ' over time.'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--custom-cluster-arguments',
       type=str,
       default='',
@@ -652,7 +664,7 @@ def add_shared_cluster_create_optional_arguments(parser: ArgumentParser):
           " --custom-cluster-arguments='--network=mtu9k --subnetwork=mtu9k'"
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--custom-nodepool-arguments',
       type=str,
       default='',
@@ -663,7 +675,7 @@ def add_shared_cluster_create_optional_arguments(parser: ArgumentParser):
           ' --custom-nodepool-arguments="--disk-size=300"'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--force',
       action='store_true',
       help=(
@@ -671,7 +683,7 @@ def add_shared_cluster_create_optional_arguments(parser: ArgumentParser):
           ' additional approval.'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--custom-tpu-nodepool-arguments',
       type=str,
       default='',
@@ -682,7 +694,7 @@ def add_shared_cluster_create_optional_arguments(parser: ArgumentParser):
           ' --custom-tpu-nodepool-arguments="--enable-ip-alias"'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--private',
       action='store_true',
       help=(
@@ -695,7 +707,7 @@ def add_shared_cluster_create_optional_arguments(parser: ArgumentParser):
           ' clusters.'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--authorized-networks',
       action='extend',
       nargs='+',
@@ -710,16 +722,16 @@ def add_shared_cluster_create_optional_arguments(parser: ArgumentParser):
           ' Example usage: --authorized-networks 1.2.3.0/24 1.2.4.5/32'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--enable-workload-identity',
       action='store_true',
       help='Enable Workload Identity Federation on the cluster and node-pools.',
   )
-  add_driver_arguments(parser)
+  add_driver_arguments(parser_or_group)
 
 
-def add_driver_arguments(parser: ArgumentParser):
-  parser.add_argument(
+def add_driver_arguments(parser_or_group: ParserOrArgumentGroup):
+  parser_or_group.add_argument(
       '--enable-gcsfuse-csi-driver',
       action='store_true',
       help=(
@@ -728,42 +740,44 @@ def add_driver_arguments(parser: ArgumentParser):
           ' Identity is enabled by default.'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--enable-gcpfilestore-csi-driver',
       action='store_true',
       help='Enable GCPFilestore driver on the cluster.',
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--enable-parallelstore-csi-driver',
       action='store_true',
       help='Enable Parallelstore CSI driver on the cluster.',
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--enable-pd-csi-driver',
       action='store_true',
       help='Enable PersistentDisk CSI driver on the cluster.',
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--enable-lustre-csi-driver',
       action='store_true',
       help='Enable Lustre CSI driver on the cluster.',
   )
 
 
-def add_shared_cluster_create_tensorboard_arguments(parser: ArgumentParser):
+def add_shared_cluster_create_tensorboard_arguments(
+    parser_or_group: ParserOrArgumentGroup,
+):
   """Add shared tensorboard arguments in cluster create and Pathways cluster create.
   Note that this feature enables non-Pathways workloads to use tensorboard arguments
   on a Pathways cluster.
 
   Args:
-    parser: cluster create argument parser or argument group
+    parser_or_group: cluster create argument parser or argument group
   """
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--create-vertex-tensorboard',
       action='store_true',
       help='Set this flag to create a Tensorboard instance in Vertex AI.',
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--tensorboard-region',
       type=str,
       default='us-central1',
@@ -774,7 +788,7 @@ def add_shared_cluster_create_tensorboard_arguments(parser: ArgumentParser):
           ' instance will be created in us-central1.'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--tensorboard-name',
       type=str,
       required=False,
@@ -787,13 +801,15 @@ def add_shared_cluster_create_tensorboard_arguments(parser: ArgumentParser):
   )
 
 
-def add_shared_cluster_create_capacity_arguments(parser: ArgumentParser):
+def add_shared_cluster_create_capacity_arguments(
+    parser_or_group: ParserOrArgumentGroup,
+):
   """Add shared capacity arguments in cluster create and Pathways cluster create.
 
   Args:
-    parser: cluster create argument parser or argument group
+    parser_or_group: cluster create argument parser or argument group
   """
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--on-demand',
       action='store_true',
       help=(
@@ -802,7 +818,7 @@ def add_shared_cluster_create_capacity_arguments(parser: ArgumentParser):
           ' types.'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--reservation',
       type=str,
       help=(
@@ -811,7 +827,7 @@ def add_shared_cluster_create_capacity_arguments(parser: ArgumentParser):
           ' `--flex` or `--on-demand` for other capacity types.'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--spot',
       action='store_true',
       help=(
@@ -820,7 +836,7 @@ def add_shared_cluster_create_capacity_arguments(parser: ArgumentParser):
           ' capacity types.'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--flex',
       action='store_true',
       help=(
@@ -831,18 +847,20 @@ def add_shared_cluster_create_capacity_arguments(parser: ArgumentParser):
   )
 
 
-def add_shared_cluster_create_mtc_arguments(parser: ArgumentParser):
+def add_shared_cluster_create_mtc_arguments(
+    parser_or_group: ParserOrArgumentGroup,
+):
   """Add shared Multi-tier Checkpointing arguments in cluster create and Pathways cluster create.
 
   Args:
-      List of cluster create MTC arguments parsers
+      List of cluster create MTC arguments parsers or group
   """
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--enable-mtc',
       action='store_true',
       help='Enable MTC on the cluster.',
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--mtc-ramdisk-size',
       type=str,
       default=None,
@@ -851,7 +869,7 @@ def add_shared_cluster_create_mtc_arguments(parser: ArgumentParser):
           ' used for multi-tier checkpointing. e.g. "64Mi" '
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--mtc-gcs-bucket',
       type=str,
       default=None,
@@ -860,7 +878,7 @@ def add_shared_cluster_create_mtc_arguments(parser: ArgumentParser):
           ' multi-tier checkpointing.'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--mtc-toleration-key',
       type=str,
       default=None,

xpk/parser/common.py

@@ -15,24 +15,31 @@ limitations under the License.
 """
 
 import argparse
+from typing import Protocol, Any
+
+
+class ParserOrArgumentGroup(Protocol):
+
+  def add_argument(self, *args, **kwargs) -> Any:
+    ...
 
 
 def add_shared_arguments(
-    custom_parser: argparse.ArgumentParser, required=False
+    custom_parser_or_group: ParserOrArgumentGroup, required=False
 ) -> None:
-  """Add shared arguments to the parser.
+  """Add shared arguments to the parser or argument group.
 
   Args:
-    custom_parser: parser to add shared arguments to.
+    custom_parser_or_group: parser or argument group to add shared arguments to.
   """
-  custom_parser.add_argument(
+  custom_parser_or_group.add_argument(
       '--project',
       type=str,
       default=None,
       help='GCE project name, defaults to "gcloud config project."',
       required=required,
   )
-  custom_parser.add_argument(
+  custom_parser_or_group.add_argument(
       '--zone',
       type=str,
       default=None,
@@ -43,7 +50,7 @@ def add_shared_arguments(
       ),
       required=required,
   )
-  custom_parser.add_argument(
+  custom_parser_or_group.add_argument(
       '--dry-run',
       type=bool,
       action=argparse.BooleanOptionalAction,
@@ -58,14 +65,14 @@ def add_shared_arguments(
 
 
 def add_cluster_arguments(
-    custom_parser: argparse.ArgumentParser, required=False
+    custom_parser_or_group: ParserOrArgumentGroup, required=False
 ) -> None:
-  """Add cluster argument to the parser.
+  """Add cluster argument to the parser or argument group.
 
   Args:
-    custom_parser: parser to add shared arguments to.
+    custom_parser_or_group: parser or argument group to add shared arguments to.
   """
-  custom_parser.add_argument(
+  custom_parser_or_group.add_argument(
       '--cluster',
       type=str,
       default=None,
@@ -74,13 +81,15 @@ def add_cluster_arguments(
   )
 
 
-def add_kind_cluster_arguments(custom_parser: argparse.ArgumentParser) -> None:
-  """Add kind cluster arguments to the parser.
+def add_kind_cluster_arguments(
+    custom_parser_or_group: ParserOrArgumentGroup,
+) -> None:
+  """Add kind cluster arguments to the parser or argument group.
 
   Args:
-    custom_parser: parser to add shared arguments to.
+    custom_parser_or_group: parser or argument group to add shared arguments to.
   """
-  custom_parser.add_argument(
+  custom_parser_or_group.add_argument(
       '--kind-cluster',
       type=bool,
       action=argparse.BooleanOptionalAction,
@@ -89,13 +98,13 @@ def add_kind_cluster_arguments(custom_parser: argparse.ArgumentParser) -> None:
   )
 
 
-def add_global_arguments(custom_parser: argparse.ArgumentParser):
+def add_global_arguments(custom_parser_or_group: ParserOrArgumentGroup):
   """Add global - no cloud dependent -  arguments to the parser.
 
   Args:
-    custom_parser: parser to add global arguments to.
+    custom_parser_or_group: parser or argument group to add global arguments to.
   """
-  custom_parser.add_argument(
+  custom_parser_or_group.add_argument(
       '--dry-run',
       type=bool,
       action=argparse.BooleanOptionalAction,
@@ -108,20 +117,20 @@ def add_global_arguments(custom_parser: argparse.ArgumentParser):
   )
 
 
-def add_slurm_arguments(custom_parser: argparse.ArgumentParser):
+def add_slurm_arguments(custom_parser_or_group: ParserOrArgumentGroup):
   """Add Slurm job arguments to the parser.
 
   Args:
-    custom_parser: parser to add global arguments to.
+    custom_parser_or_group: parser or argument group to add global arguments to.
   """
-  custom_parser.add_argument(
+  custom_parser_or_group.add_argument(
       '--ignore-unknown-flags',
       type=bool,
       action=argparse.BooleanOptionalAction,
       default=False,
       help='Ignore all the unsupported flags in the bash script.',
   )
-  custom_parser.add_argument(
+  custom_parser_or_group.add_argument(
       '-a',
       '--array',
       type=str,
@@ -137,32 +146,32 @@ def add_slurm_arguments(custom_parser: argparse.ArgumentParser):
           ' 0. The maximum index value is 2147483647.'
       ),
   )
-  custom_parser.add_argument(
+  custom_parser_or_group.add_argument(
       '-c',
       '--cpus-per-task',
       type=str,
       default=None,
       help='How much cpus a container inside a pod requires.',
   )
-  custom_parser.add_argument(
+  custom_parser_or_group.add_argument(
       '--gpus-per-task',
       type=str,
       default=None,
       help='How much gpus a container inside a pod requires.',
   )
-  custom_parser.add_argument(
+  custom_parser_or_group.add_argument(
       '--mem',
       type=str,
       default=None,
       help='How much memory a pod requires.',
   )
-  custom_parser.add_argument(
+  custom_parser_or_group.add_argument(
       '--mem-per-task',
       type=str,
       default=None,
       help='How much memory a container requires.',
   )
-  custom_parser.add_argument(
+  custom_parser_or_group.add_argument(
       '--mem-per-cpu',
       type=str,
       default=None,
@@ -171,7 +180,7 @@ def add_slurm_arguments(custom_parser: argparse.ArgumentParser):
           'of requested cpus per task by mem-per-cpu.'
       ),
   )
-  custom_parser.add_argument(
+  custom_parser_or_group.add_argument(
       '--mem-per-gpu',
       type=str,
       default=None,
@@ -180,21 +189,21 @@ def add_slurm_arguments(custom_parser: argparse.ArgumentParser):
           'of requested gpus per task by mem-per-gpu.'
       ),
   )
-  custom_parser.add_argument(
+  custom_parser_or_group.add_argument(
       '-N',
       '--nodes',
       type=int,
       default=None,
       help='Number of pods to be used at a time.',
   )
-  custom_parser.add_argument(
+  custom_parser_or_group.add_argument(
       '-n',
       '--ntasks',
       type=int,
       default=None,
       help='Number of identical containers inside of a pod, usually 1.',
   )
-  custom_parser.add_argument(
+  custom_parser_or_group.add_argument(
       '-o',
       '--output',
       type=str,
@@ -204,7 +213,7 @@ def add_slurm_arguments(custom_parser: argparse.ArgumentParser):
           ' passed it proceeds to stdout, and is available via kubectl logs.'
       ),
   )
-  custom_parser.add_argument(
+  custom_parser_or_group.add_argument(
       '-e',
       '--error',
       type=str,
@@ -214,27 +223,27 @@ def add_slurm_arguments(custom_parser: argparse.ArgumentParser):
           ' proceeds to stdout, and is available via kubectl logs.'
       ),
   )
-  custom_parser.add_argument(
+  custom_parser_or_group.add_argument(
       '--input',
       type=str,
       default=None,
       help='What to pipe into the script.',
   )
-  custom_parser.add_argument(
+  custom_parser_or_group.add_argument(
       '-J',
       '--job-name',
       type=str,
       default=None,
       help='What is the job name.',
   )
-  custom_parser.add_argument(
+  custom_parser_or_group.add_argument(
       '-D',
       '--chdir',
       type=str,
       default=None,
       help='Change directory before executing the script.',
   )
-  custom_parser.add_argument(
+  custom_parser_or_group.add_argument(
       '-t',
       '--time',
       type=str,
@@ -247,7 +256,7 @@ def add_slurm_arguments(custom_parser: argparse.ArgumentParser):
           'and "days-hours:minutes:seconds".'
       ),
   )
-  custom_parser.add_argument(
+  custom_parser_or_group.add_argument(
       '--priority',
       type=str,
       default='medium',

xpk/parser/storage.py

@@ -28,6 +28,13 @@ from .common import (
     add_kind_cluster_arguments,
     add_shared_arguments,
 )
+from typing import Protocol, Any
+
+
+class Subcommands(Protocol):
+
+  def add_parser(self, *args, **kwargs) -> Any:
+    ...
 
 
 def set_storage_parser(storage_parser: argparse.ArgumentParser) -> None:
@@ -47,7 +54,7 @@ def set_storage_parser(storage_parser: argparse.ArgumentParser) -> None:
 
 
 def add_storage_attach_parser(
-    storage_subcommands_parser: argparse.ArgumentParser,
+    storage_subcommands_parser: Subcommands,
 ) -> None:
 
   storage_attach_parser: argparse.ArgumentParser = (
@@ -171,9 +178,7 @@ def add_storage_attach_parser(
   add_kind_cluster_arguments(opt_args)
 
 
-def add_storage_create_parser(
-    storage_subcommands_parser: argparse.ArgumentParser,
-) -> None:
+def add_storage_create_parser(storage_subcommands_parser: Subcommands) -> None:
   storage_create_parser: argparse.ArgumentParser = (
       storage_subcommands_parser.add_parser(
           'create', help='create XPK Storage.'
@@ -272,9 +277,7 @@ def add_storage_create_parser(
   add_kind_cluster_arguments(opt_args)
 
 
-def add_storage_list_parser(
-    storage_subcommands_parser: argparse.ArgumentParser,
-):
+def add_storage_list_parser(storage_subcommands_parser: Subcommands) -> None:
   storage_list_parser: argparse.ArgumentParser = (
       storage_subcommands_parser.add_parser('list', help='List XPK Storages.')
   )
@@ -290,9 +293,7 @@ def add_storage_list_parser(
   )
 
 
-def add_storage_detach_parser(
-    storage_subcommands_parser: argparse.ArgumentParser,
-):
+def add_storage_detach_parser(storage_subcommands_parser: Subcommands) -> None:
   storage_detach_parser: argparse.ArgumentParser = (
       storage_subcommands_parser.add_parser(
           'detach', help='Detach XPK Storage.'
@@ -315,9 +316,7 @@ def add_storage_detach_parser(
   add_kind_cluster_arguments(opt_args)
 
 
-def add_storage_delete_parser(
-    storage_subcommands_parser: argparse.ArgumentParser,
-):
+def add_storage_delete_parser(storage_subcommands_parser: Subcommands) -> None:
   storage_delete_parser: argparse.ArgumentParser = (
       storage_subcommands_parser.add_parser(
           'delete', help='Delete XPK Storage.'