xpk 0.6.0__py3-none-any.whl → 0.7.1__py3-none-any.whl

xpk/parser/batch.py

@@ -14,9 +14,12 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-import argparse
-
-from .common import add_shared_arguments
+from .common import (
+    add_shared_arguments,
+    add_slurm_arguments,
+    add_cluster_arguments,
+    add_kind_cluster_arguments,
+)
 from ..commands.batch import batch
 
 
@@ -32,153 +35,9 @@ def set_batch_parser(batch_parser):
   batch_required_arguments.add_argument(
       'script', help='script with batch task to run'
   )
-  batch_optional_arguments.add_argument(
-      '--cluster',
-      type=str,
-      default=None,
-      help='Cluster to which command applies.',
-  )
-  batch_optional_arguments.add_argument(
-      '--kind-cluster',
-      type=bool,
-      action=argparse.BooleanOptionalAction,
-      default=False,
-      help='Apply command to a local test cluster.',
-  )
-  add_shared_arguments(batch_optional_arguments)
 
+  add_cluster_arguments(batch_optional_arguments)
+  add_kind_cluster_arguments(batch_optional_arguments)
+  add_shared_arguments(batch_optional_arguments)
+  add_slurm_arguments(batch_optional_arguments)
   batch_parser.set_defaults(func=batch)
-
-  batch_optional_arguments.add_argument(
-      '--ignore-unknown-flags',
-      type=bool,
-      action=argparse.BooleanOptionalAction,
-      default=False,
-      help='Ignore all the unsupported flags in the bash script.',
-  )
-  batch_optional_arguments.add_argument(
-      '-a',
-      '--array',
-      type=str,
-      default=None,
-      help=(
-          'Submit a job array, multiple jobs to be executed with identical'
-          ' parameters. The indexes specification identifies what array index'
-          ' values should be used. For example, "--array=0-15" or'
-          ' "--array=0,6,16-32". Multiple values may be specified using a comma'
-          ' separated list and/or a range of values with a "-" separator. For'
-          ' example "--array=0-15%%4" will limit the number of simultaneously'
-          ' running tasks from this job array to 4. The minimum index value is'
-          ' 0. The maximum index value is 2147483647.'
-      ),
-  )
-  batch_optional_arguments.add_argument(
-      '-c',
-      '--cpus-per-task',
-      type=str,
-      default=None,
-      help='How much cpus a container inside a pod requires.',
-  )
-  batch_optional_arguments.add_argument(
-      '--gpus-per-task',
-      type=str,
-      default=None,
-      help='How much gpus a container inside a pod requires.',
-  )
-  batch_optional_arguments.add_argument(
-      '--mem',
-      type=str,
-      default=None,
-      help='How much memory a pod requires.',
-  )
-  batch_optional_arguments.add_argument(
-      '--mem-per-task',
-      type=str,
-      default=None,
-      help='How much memory a container requires.',
-  )
-  batch_optional_arguments.add_argument(
-      '--mem-per-cpu',
-      type=str,
-      default=None,
-      help=(
-          'How much memory a container requires, it multiplies the number '
-          'of requested cpus per task by mem-per-cpu.'
-      ),
-  )
-  batch_optional_arguments.add_argument(
-      '--mem-per-gpu',
-      type=str,
-      default=None,
-      help=(
-          'How much memory a container requires, it multiplies the number '
-          'of requested gpus per task by mem-per-gpu.'
-      ),
-  )
-  batch_optional_arguments.add_argument(
-      '-N',
-      '--nodes',
-      type=int,
-      default=None,
-      help='Number of pods to be used at a time.',
-  )
-  batch_optional_arguments.add_argument(
-      '-n',
-      '--ntasks',
-      type=int,
-      default=None,
-      help='Number of identical containers inside of a pod, usually 1.',
-  )
-  batch_optional_arguments.add_argument(
-      '-o',
-      '--output',
-      type=str,
-      default=None,
-      help=(
-          'Where to redirect the standard output stream of a task. If not'
-          ' passed it proceeds to stdout, and is available via kubectl logs.'
-      ),
-  )
-  batch_optional_arguments.add_argument(
-      '-e',
-      '--error',
-      type=str,
-      default=None,
-      help=(
-          'Where to redirect std error stream of a task. If not passed it'
-          ' proceeds to stdout, and is available via kubectl logs.'
-      ),
-  )
-  batch_optional_arguments.add_argument(
-      '--input',
-      type=str,
-      default=None,
-      help='What to pipe into the script.',
-  )
-  batch_optional_arguments.add_argument(
-      '-J',
-      '--job-name',
-      type=str,
-      default=None,
-      help='What is the job name.',
-  )
-  batch_optional_arguments.add_argument(
-      '-D',
-      '--chdir',
-      type=str,
-      default=None,
-      help='Change directory before executing the script.',
-  )
-  batch_optional_arguments.add_argument(
-      '-t',
-      '--time',
-      type=str,
-      default=None,
-      help=(
-          'Set a limit on the total run time of the job. '
-          'A time limit of zero requests that no time limit be imposed. '
-          'Acceptable time formats include "minutes", "minutes:seconds", '
-          '"hours:minutes:seconds", "days-hours", "days-hours:minutes" '
-          'and "days-hours:minutes:seconds".'
-      ),
-  )

xpk/parser/cluster.py

@@ -23,8 +23,11 @@ from ..commands.cluster import (
     cluster_describe,
     cluster_list,
 )
-from ..core.core import DEFAULT_VERTEX_TENSORBOARD_NAME
+from ..core.vertex import DEFAULT_VERTEX_TENSORBOARD_NAME
 from .common import add_shared_arguments
+from .validators import name_type
+from ..commands.config import xpk_cfg
+from ..core.config import CFG_BUCKET_KEY
 
 
 def set_cluster_parser(cluster_parser):
@@ -82,6 +85,13 @@ def set_cluster_parser(cluster_parser):
   )
 
   ### Optional arguments specific to "cluster create"
+  cluster_create_optional_arguments.add_argument(
+      '--cluster-state-gcs-bucket',
+      type=str,
+      default=xpk_cfg.get(CFG_BUCKET_KEY),
+      help='The name of the bucket to store cluster state.',
+      required=False,
+  )
   cluster_create_optional_arguments.add_argument(
       '--num-nodes',
       type=int,
@@ -93,11 +103,10 @@ def set_cluster_parser(cluster_parser):
       '--enable-pathways',
       action='store_true',
       help=(
-          'DEPRECATING SOON!!! Please use `xpk cluster create-pathways`.'
-          ' Enable cluster to accept Pathways workloads.'
+          'Please use `xpk cluster create-pathways` instead to'
+          ' enable cluster to accept Pathways workloads.'
       ),
   )
-
   ### Autoprovisioning arguments specific to "cluster create"
   cluster_create_autoprovisioning_arguments = (
       cluster_create_parser.add_argument_group(
@@ -269,13 +278,20 @@ def set_cluster_parser(cluster_parser):
   ### Required arguments
   cluster_delete_required_arguments.add_argument(
       '--cluster',
-      type=str,
+      type=name_type,
       default=None,
       help='The name of the cluster to be deleted.',
       required=True,
   )
 
   ### Optional Arguments
+  cluster_delete_optional_arguments.add_argument(
+      '--cluster-state-gcs-bucket',
+      type=str,
+      default=xpk_cfg.get(CFG_BUCKET_KEY),
+      help='The name of the bucket to store cluster state.',
+      required=False,
+  )
   add_shared_arguments(cluster_delete_optional_arguments)
   cluster_delete_parser.set_defaults(func=cluster_delete)
   cluster_delete_parser.add_argument(
@@ -326,7 +342,7 @@ def set_cluster_parser(cluster_parser):
   ### Required arguments
   cluster_cacheimage_required_arguments.add_argument(
       '--cluster',
-      type=str,
+      type=name_type,
       default=None,
       help='The name of the cluster to cache the image.',
       required=True,
@@ -370,7 +386,7 @@ def set_cluster_parser(cluster_parser):
   ### Required arguments
   cluster_describe_required_arguments.add_argument(
       '--cluster',
-      type=str,
+      type=name_type,
       default=None,
       help='The name of the cluster to be describe.',
       required=True,
@@ -402,7 +418,7 @@ def add_shared_cluster_create_required_arguments(args_parsers):
   for custom_parser in args_parsers:
     custom_parser.add_argument(
         '--cluster',
-        type=str,
+        type=name_type,
         default=None,
         help=(
             'The name of the cluster. Will be used as the prefix for internal'
@@ -548,6 +564,31 @@ def add_shared_cluster_create_optional_arguments(args_parsers):
             ' Example usage: --authorized-networks 1.2.3.0/24 1.2.4.5/32'
         ),
     )
+    custom_parser.add_argument(
+        '--enable-workload-identity',
+        action='store_true',
+        help=(
+            'Enable Workload Identity Federation on the cluster and node-pools.'
+        ),
+    )
+    custom_parser.add_argument(
+        '--enable-gcsfuse-csi-driver',
+        action='store_true',
+        help=(
+            'Enable GSCFuse driver on the cluster. This enables Workload'
+            ' Identity Federation. When using A3 ultra/A3 mega Workload'
+            ' Identity is enabled by default.'
+        ),
+    )
+
+    custom_parser.add_argument(
+        '--enable-gcpfilestore-csi-driver',
+        action='store_true',
+        help=(
+            'Enable GCPFilestore driver on the cluster. This enables Workload'
+            ' Identity Federation.'
+        ),
+    )
 
 
 def add_shared_cluster_create_tensorboard_arguments(args_parsers):

xpk/parser/common.py

@@ -17,7 +17,9 @@ limitations under the License.
 import argparse
 
 
-def add_shared_arguments(custom_parser: argparse.ArgumentParser):
+def add_shared_arguments(
+    custom_parser: argparse.ArgumentParser, required=False
+) -> None:
   """Add shared arguments to the parser.
 
   Args:
@@ -28,6 +30,7 @@ def add_shared_arguments(custom_parser: argparse.ArgumentParser):
       type=str,
       default=None,
       help='GCE project name, defaults to "gcloud config project."',
+      required=required,
   )
   custom_parser.add_argument(
       '--zone',
@@ -38,6 +41,7 @@ def add_shared_arguments(custom_parser: argparse.ArgumentParser):
           'compute/zone." Only one of --zone or --region is allowed in a '
           'command.'
       ),
+      required=required,
   )
   custom_parser.add_argument(
       '--dry-run',
@@ -49,6 +53,39 @@ def add_shared_arguments(custom_parser: argparse.ArgumentParser):
           ' but not run them. This is imperfect in cases where xpk might'
           ' branch based on the output of commands'
       ),
+      required=required,
+  )
+
+
+def add_cluster_arguments(
+    custom_parser: argparse.ArgumentParser, required=False
+) -> None:
+  """Add cluster argument to the parser.
+
+  Args:
+    custom_parser: parser to add shared arguments to.
+  """
+  custom_parser.add_argument(
+      '--cluster',
+      type=str,
+      default=None,
+      help='The name of the cluster.',
+      required=required,
+  )
+
+
+def add_kind_cluster_arguments(custom_parser: argparse.ArgumentParser) -> None:
+  """Add kind cluster arguments to the parser.
+
+  Args:
+    custom_parser: parser to add shared arguments to.
+  """
+  custom_parser.add_argument(
+      '--kind-cluster',
+      type=bool,
+      action=argparse.BooleanOptionalAction,
+      default=False,
+      help='Apply command to a local test cluster.',
   )
 
 
@@ -69,3 +106,154 @@ def add_global_arguments(custom_parser: argparse.ArgumentParser):
           ' branch based on the output of commands'
       ),
   )
+
+
+def add_slurm_arguments(custom_parser: argparse.ArgumentParser):
+  """Add Slurm job arguments to the parser.
+
+  Args:
+    custom_parser: parser to add global arguments to.
+  """
+  custom_parser.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(
+      '-a',
+      '--array',
+      type=str,
+      default=None,
+      help=(
+          'Submit a job array, multiple jobs to be executed with identical'
+          ' parameters. The indexes specification identifies what array index'
+          ' values should be used. For example, "--array=0-15" or'
+          ' "--array=0,6,16-32". Multiple values may be specified using a comma'
+          ' separated list and/or a range of values with a "-" separator. For'
+          ' example "--array=0-15%%4" will limit the number of simultaneously'
+          ' running tasks from this job array to 4. The minimum index value is'
+          ' 0. The maximum index value is 2147483647.'
+      ),
+  )
+  custom_parser.add_argument(
+      '-c',
+      '--cpus-per-task',
+      type=str,
+      default=None,
+      help='How much cpus a container inside a pod requires.',
+  )
+  custom_parser.add_argument(
+      '--gpus-per-task',
+      type=str,
+      default=None,
+      help='How much gpus a container inside a pod requires.',
+  )
+  custom_parser.add_argument(
+      '--mem',
+      type=str,
+      default=None,
+      help='How much memory a pod requires.',
+  )
+  custom_parser.add_argument(
+      '--mem-per-task',
+      type=str,
+      default=None,
+      help='How much memory a container requires.',
+  )
+  custom_parser.add_argument(
+      '--mem-per-cpu',
+      type=str,
+      default=None,
+      help=(
+          'How much memory a container requires, it multiplies the number '
+          'of requested cpus per task by mem-per-cpu.'
+      ),
+  )
+  custom_parser.add_argument(
+      '--mem-per-gpu',
+      type=str,
+      default=None,
+      help=(
+          'How much memory a container requires, it multiplies the number '
+          'of requested gpus per task by mem-per-gpu.'
+      ),
+  )
+  custom_parser.add_argument(
+      '-N',
+      '--nodes',
+      type=int,
+      default=None,
+      help='Number of pods to be used at a time.',
+  )
+  custom_parser.add_argument(
+      '-n',
+      '--ntasks',
+      type=int,
+      default=None,
+      help='Number of identical containers inside of a pod, usually 1.',
+  )
+  custom_parser.add_argument(
+      '-o',
+      '--output',
+      type=str,
+      default=None,
+      help=(
+          'Where to redirect the standard output stream of a task. If not'
+          ' passed it proceeds to stdout, and is available via kubectl logs.'
+      ),
+  )
+  custom_parser.add_argument(
+      '-e',
+      '--error',
+      type=str,
+      default=None,
+      help=(
+          'Where to redirect std error stream of a task. If not passed it'
+          ' proceeds to stdout, and is available via kubectl logs.'
+      ),
+  )
+  custom_parser.add_argument(
+      '--input',
+      type=str,
+      default=None,
+      help='What to pipe into the script.',
+  )
+  custom_parser.add_argument(
+      '-J',
+      '--job-name',
+      type=str,
+      default=None,
+      help='What is the job name.',
+  )
+  custom_parser.add_argument(
+      '-D',
+      '--chdir',
+      type=str,
+      default=None,
+      help='Change directory before executing the script.',
+  )
+  custom_parser.add_argument(
+      '-t',
+      '--time',
+      type=str,
+      default=None,
+      help=(
+          'Set a limit on the total run time of the job. '
+          'A time limit of zero requests that no time limit be imposed. '
+          'Acceptable time formats include "minutes", "minutes:seconds", '
+          '"hours:minutes:seconds", "days-hours", "days-hours:minutes" '
+          'and "days-hours:minutes:seconds".'
+      ),
+  )
+  custom_parser.add_argument(
+      '--priority',
+      type=str,
+      default='medium',
+      choices=['very-low', 'low', 'medium', 'high', 'very-high'],
+      help=(
+          'A priority, one of `very-low`, `low`, `medium`, `high` or'
+          ' `very-high`. Defaults to `medium`.'
+      ),
+  )

xpk/parser/config.py

@@ -0,0 +1,49 @@
+"""
+Copyright 2024 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 ..commands.config import get_config, set_config
+from ..core.config import DEFAULT_KEYS
+from .common import add_shared_arguments
+
+
+def set_config_parsers(config_parser):
+  add_shared_arguments(config_parser)
+
+  config_subcommands = config_parser.add_subparsers(
+      title='config subcommands', dest='xpk_config_subcommands'
+  )
+  config_set_parser = config_subcommands.add_parser(
+      'set', help='set config key'
+  )
+  config_get_parser = config_subcommands.add_parser(
+      'get', help='get config key'
+  )
+  config_set_parser.add_argument(
+      'set_config_args',
+      help=f"""Pair of (key, value) to be set in config. Allowed keys are: {DEFAULT_KEYS}.
+      Command usage: `xpk config set key value`""",
+      type=str,
+      nargs=2,
+  )
+  config_get_parser.add_argument(
+      'get_config_key',
+      help=f"""Get key value from config. Allowed keys are: {DEFAULT_KEYS} .
+      Command usage: `xpk config get key`""",
+      type=str,
+      nargs=1,
+  )
+  config_set_parser.set_defaults(func=set_config)
+  config_get_parser.set_defaults(func=get_config)

xpk/parser/core.py

@@ -16,9 +16,12 @@ limitations under the License.
 
 import argparse
 
+from .config import set_config_parsers
+
 from ..utils.console import xpk_print
 from .cluster import set_cluster_parser
 from .inspector import set_inspector_parser
+from .storage import set_storage_parser
 from .workload import set_workload_parsers
 from .batch import set_batch_parser
 from .job import set_job_parser
@@ -26,6 +29,7 @@ from .info import set_info_parser
 from .kind import set_kind_parser
 from .shell import set_shell_parser
 from .version import set_version_parser
+from .run import set_run_parser
 
 
 def set_parser(parser: argparse.ArgumentParser):
@@ -33,7 +37,10 @@ def set_parser(parser: argparse.ArgumentParser):
       title="xpk subcommands", dest="xpk_subcommands", help="Top level commands"
   )
   workload_parser = xpk_subcommands.add_parser(
-      "workload", help="Commands around workload management."
+      "workload", help="Commands around workload management"
+  )
+  storage_parser = xpk_subcommands.add_parser(
+      "storage", help="Commands around storage management"
   )
   cluster_parser = xpk_subcommands.add_parser(
       "cluster",
@@ -65,6 +72,15 @@ def set_parser(parser: argparse.ArgumentParser):
       "version", help="Command to get xpk version"
   )
 
+  config_parser = xpk_subcommands.add_parser(
+      "config", help="Commands to set and retrieve values from xpk config."
+  )
+
+  run_parser = xpk_subcommands.add_parser(
+      "run",
+      help="Command to run parallel jobs",
+  )
+
   def default_subcommand_function(
       _args,
   ) -> int:  # args is unused, so pylint: disable=invalid-name
@@ -86,7 +102,10 @@ def set_parser(parser: argparse.ArgumentParser):
     shell_parser.print_help()
     version_parser.print_help()
     kind_parser.print_help()
+    config_parser.print_help()
+    run_parser.print_help()
 
+    storage_parser.print_help()
     return 0
 
   parser.set_defaults(func=default_subcommand_function)
@@ -97,7 +116,11 @@ def set_parser(parser: argparse.ArgumentParser):
   job_parser.set_defaults(func=default_subcommand_function)
   kind_parser.set_defaults(func=default_subcommand_function)
   shell_parser.set_defaults(func=default_subcommand_function)
+  storage_parser.set_defaults(func=default_subcommand_function)
   version_parser.set_defaults(func=default_subcommand_function)
+  config_parser.set_defaults(func=default_subcommand_function)
+  run_parser.set_defaults(func=default_subcommand_function)
+
   set_workload_parsers(workload_parser=workload_parser)
   set_cluster_parser(cluster_parser=cluster_parser)
   set_inspector_parser(inspector_parser=inspector_parser)
@@ -106,4 +129,7 @@ def set_parser(parser: argparse.ArgumentParser):
   set_job_parser(job_parser=job_parser)
   set_kind_parser(kind_parser=kind_parser)
   set_shell_parser(shell_parser=shell_parser)
+  set_storage_parser(storage_parser=storage_parser)
   set_version_parser(version_parser=version_parser)
+  set_config_parsers(config_parser=config_parser)
+  set_run_parser(run_parser=run_parser)

xpk/parser/info.py

@@ -16,6 +16,7 @@ limitations under the License.
 
 from ..commands.info import info
 from .common import add_shared_arguments
+from .validators import name_type
 import argparse
 
 
@@ -29,7 +30,7 @@ def set_info_parser(info_parser: argparse.ArgumentParser) -> None:
 
   info_required_arguments.add_argument(
       '--cluster',
-      type=str,
+      type=name_type,
       default=None,
       help='Cluster to which command applies.',
       required=True,

xpk/parser/inspector.py

@@ -15,7 +15,7 @@ limitations under the License.
 """
 
 from ..commands.inspector import inspector
-from .validators import workload_name_type
+from .validators import name_type
 from .common import add_shared_arguments
 
 
@@ -37,7 +37,7 @@ def set_inspector_parser(inspector_parser):
 
   inspector_parser_required_arguments.add_argument(
       '--cluster',
-      type=str,
+      type=name_type,
       default=None,
       help='The name of the cluster to investigate.',
       required=True,
@@ -48,7 +48,7 @@ def set_inspector_parser(inspector_parser):
 
   inspector_parser_optional_arguments.add_argument(
       '--workload',
-      type=workload_name_type,
+      type=name_type,
       default=None,
       help='The name of the workload to investigate.',
   )