xpk 0.11.0__py3-none-any.whl → 0.13.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
 
 
@@ -174,6 +174,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)
 
 
@@ -208,6 +215,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(
@@ -237,6 +252,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)
 
 
@@ -312,6 +336,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)
 
 
@@ -529,15 +560,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 +577,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 +588,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 +608,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 +633,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 +655,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 +666,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 +676,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 +687,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 +698,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 +706,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 +717,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 +730,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 +745,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 +763,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 +811,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 +824,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 +841,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 +850,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 +859,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 +870,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 +892,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 +901,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,
@@ -869,3 +910,23 @@ def add_shared_cluster_create_mtc_arguments(parser: ArgumentParser):
           ' 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/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',