xpk 0.7.2__py3-none-any.whl → 0.9.0__py3-none-any.whl

xpk/parser/storage.py

@@ -70,10 +70,10 @@ def add_storage_attach_parser(
       '--type',
       type=str,
       help=(
-          'The type of storage. Currently supported types: ["gcsfuse",'
-          ' "gcpfilestore"]'
+          'The type of storage. Currently supported types: "gcsfuse",'
+          ' "gcpfilestore", "parallelstore", "pd"'
       ),
-      choices=['gcsfuse', 'gcpfilestore'],
+      choices=['gcsfuse', 'gcpfilestore', 'parallelstore', 'pd'],
       required=True,
   )
   add_cluster_arguments(req_args, required=True)
@@ -114,6 +114,15 @@ def add_storage_attach_parser(
           ' is infered as a bucket name.'
       ),
   )
+  gcsfuse_args.add_argument(
+      '--prefetch-metadata',
+      action=argparse.BooleanOptionalAction,
+      default=True,
+      help=(
+          '(optional) Enables metadata pre-population when'
+          ' mounting the volume. True by default.'
+      ),
+  )
 
   gcpfilestore_args = storage_attach_parser.add_argument_group(
       'Filestore arguments',
@@ -146,13 +155,19 @@ def add_storage_attach_parser(
 
   opt_args = storage_attach_parser.add_argument_group(
       'Optional Arguments',
-      'Optional arguments for storage create.',
+      'Optional arguments for storage attach.',
   )
   opt_args.add_argument(
       '--manifest',
       type=str,
       help='Path to manifest file containing volume definitions',
   )
+  opt_args.add_argument(
+      '--mount-options',
+      type=str,
+      help='Comma-separated list of mountOptions for PersistentVolume',
+      default='implicit-dirs',
+  )
   add_kind_cluster_arguments(opt_args)
 
 
@@ -184,7 +199,6 @@ def add_storage_create_parser(
       ),
       required=True,
   )
-
   req_args.add_argument(
       '--type',
       type=str,
@@ -248,6 +262,12 @@ def add_storage_create_parser(
       type=str,
       help='Path to manifest file containing volume definitions',
   )
+  opt_args.add_argument(
+      '--mount-options',
+      type=str,
+      help='Comma-separated list of mountOptions for PersistentVolume',
+      default='',
+  )
 
   add_kind_cluster_arguments(opt_args)
 

xpk/parser/workload.py

@@ -134,6 +134,24 @@ def set_workload_parsers(workload_parser):
           ' to use `gke.io/topology-aware-auto`.'
       ),
   )
+  workload_create_parser_optional_arguments.add_argument(
+      '--ramdisk-directory',
+      type=str,
+      default='',
+      help=(
+          'The directory of the locally mounted RAM disk. This is only to'
+          ' be used with the CSI driver provided by GKE.'
+      ),
+  )
+  workload_create_parser_optional_arguments.add_argument(
+      '--mtc-enabled',
+      action='store_true',
+      help=(
+          'The workload can use multi-tier checkpointing controllers when the'
+          ' --ramdisk-directory argument is used with this additional'
+          ' argument.'
+      ),
+  )
   workload_create_parser_optional_arguments.add_argument(
       '--debug-dump-gcs',
       type=str,
@@ -161,6 +179,19 @@ def set_workload_parsers(workload_parser):
           ' create Pathways workloads.'
       ),
   )
+  workload_create_parser_optional_arguments.add_argument(
+      '--restart-on-exit-codes',
+      type=str,
+      default=None,
+      help=(
+          'Adding this argument specifies additional user-defined exit codes'
+          ' that allow restarting the workload when --max-restarts is set to'
+          ' a value greater than 0. By default, workloads restart on exit'
+          ' codes 42 and 127-255. Any exit codes provided through this flag'
+          ' will be included alongside the default codes for restarting'
+          ' conditions.'
+      ),
+  )
 
   # Autoprovisioning workload arguments
   workload_create_autoprovisioning_arguments.add_argument(
@@ -244,9 +275,7 @@ def set_workload_parsers(workload_parser):
   workload_create_pathways_parser_optional_arguments.add_argument(
       '--proxy-server-image',
       type=str,
-      default=(
-          'us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest'
-      ),
+      default='',
       help=(
           'Please provide the proxy server image for Pathways. This arg can'
           ' only be used in `xpk workload create-pathways`.'
@@ -255,7 +284,7 @@ def set_workload_parsers(workload_parser):
   workload_create_pathways_parser_optional_arguments.add_argument(
       '--server-image',
       type=str,
-      default='us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest',
+      default='',
       help=(
           'Please provide the server image for Pathways. This arg can only be'
           ' used in `xpk workload create-pathways`.'
@@ -293,7 +322,7 @@ def set_workload_parsers(workload_parser):
   workload_create_pathways_parser_optional_arguments.add_argument(
       '--custom-pathways-server-args',
       type=str,
-      default=None,
+      default='',
       help=(
           'Provide custom Pathways server args as follows -'
           " --custom-pathways-server-args='--arg_1=xxx --arg2=yyy'"
@@ -304,7 +333,7 @@ def set_workload_parsers(workload_parser):
   workload_create_pathways_parser_optional_arguments.add_argument(
       '--custom-pathways-proxy-server-args',
       type=str,
-      default=None,
+      default='',
       help=(
           'Provide custom Pathways proxy server args as follows -'
           " --custom-pathways-proxy-server-args='--arg_1=xxx --arg2=yyy'"
@@ -315,7 +344,7 @@ def set_workload_parsers(workload_parser):
   workload_create_pathways_parser_optional_arguments.add_argument(
       '--custom-pathways-worker-args',
       type=str,
-      default=None,
+      default='',
       help=(
           'Provide custom Pathways worker args as follows -'
           " --custom-pathways-worker-args='--arg_1=xxx --arg2=yyy'"
@@ -323,6 +352,27 @@ def set_workload_parsers(workload_parser):
       required=False,
   )
 
+  workload_create_pathways_parser_optional_arguments.add_argument(
+      '--elastic-slices',
+      type=int,
+      default=0,
+      help=(
+          'Enable elastic slices in Pathways and specify'
+          ' the number of slices the workload could lose.'
+      ),
+      required=False,
+  )
+  workload_create_pathways_parser_optional_arguments.add_argument(
+      '--max-slice-restarts',
+      type=int,
+      default=1,
+      help=(
+          'Specify the maximum times the workers in a slice can be'
+          ' restarted. Used with --elastic-slices for Pathways workloads.'
+      ),
+      required=False,
+  )
+
   add_shared_workload_create_required_arguments([
       workload_create_parser_required_arguments,
       workload_create_pathways_parser_required_arguments,
@@ -583,9 +633,9 @@ def add_shared_workload_create_optional_arguments(args_parsers):
         ),
     )
     custom_parser.add_argument(
-        '--remote-python-sidecar-image',
+        '--colocated-python-sidecar-image',
         type=str,
-        default=None,
+        default='',
         help='Remote Python sidecar server image.',
     )
     custom_parser.add_argument(
@@ -596,28 +646,6 @@ def add_shared_workload_create_optional_arguments(args_parsers):
             ' the workload.'
         ),
     )
-    custom_parser.add_argument(
-        '--restart-on-exit-codes',
-        type=str,
-        default=None,
-        help=(
-            'Adding this argument specifies additional user-defined exit codes'
-            ' that allow restarting the workload when --max-restarts is set to'
-            ' a value greater than 0. By default, workloads restart on exit'
-            ' codes 42 and 127-255. Any exit codes provided through this flag'
-            ' will be included alongside the default codes for restarting'
-            ' conditions.'
-        ),
-    )
-    custom_parser.add_argument(
-        '--ramdisk-directory',
-        type=str,
-        default='',
-        help=(
-            'The directory of the locally mounted RAM disk. This is only to'
-            ' be used with the CSI driver provided by GKE.'
-        ),
-    )
 
 
 def add_shared_workload_create_env_arguments(args_parsers):

xpk/utils/kubectl.py

@@ -20,10 +20,11 @@ from kubernetes.dynamic import DynamicClient
 from .console import xpk_print
 
 
-def apply_kubectl_manifest(client, manifest):
+def apply_kubectl_manifest(client, manifest) -> int:
   xpk_print('Applying manifest')
   dynamic_client = DynamicClient(client)
 
+  status_code = 0
   for obj in manifest:
     api_version = obj['apiVersion']
     kind = obj['kind']
@@ -55,3 +56,5 @@ def apply_kubectl_manifest(client, manifest):
         )
       else:
         xpk_print(f'Error applying {kind}: {e}')
+        status_code = 1
+  return status_code