Add pathways + nap support (#623)

* feat: add pathways + nap support

* feat: add support for nap arguments in pathways workload create

* style: remove unecessary comments

Author: scaliby

src/xpk/commands/cluster.py

@@ -141,8 +141,6 @@ def cluster_adapt(args) -> None:
     if not tensorboard_config:
       xpk_exit(1)
 
-  # Provision node pools dynamically based on incoming workloads:
-  # Currently autoprovisioning is not supported with Pathways.
   autoprovisioning_config = None
   if args.enable_autoprovisioning:
     xpk_print('Enabling Autoprovisioning')
@@ -294,7 +292,7 @@ def cluster_create(args) -> None:
   # Provision node pools dynamically based on incoming workloads:
   # Currently autoprovisioning is not supported with Pathways.
   autoprovisioning_config = None
-  if not args.enable_pathways and args.enable_autoprovisioning:
+  if args.enable_autoprovisioning:
     xpk_print('Enabling Autoprovisioning')
     autoprovisioning_config, return_code = enable_autoprovisioning_on_cluster(
         args, system

src/xpk/core/nap.py

@@ -269,9 +269,6 @@ def is_autoprovisioning_enabled(
     bool is true if autoprovisioning is enabled, false otherwise.
     int of 0 if successful and 1 otherwise.
   """
-  # Currently autoprovisioning is not enabled for Pathways workloads. b/360898087
-  if args.use_pathways:
-    return False, 0
 
   resources_configmap_name = f'{args.cluster}-{CLUSTER_RESOURCES_CONFIGMAP}'
   cluster_config_map = get_cluster_configmap(args, resources_configmap_name)

src/xpk/core/resources.py

@@ -108,13 +108,7 @@ def create_cluster_configmaps(
   device_type = system.device_type
   if system.accelerator_type == AcceleratorType['GPU']:
     resources_data = f'{device_type}: "{int(args.num_nodes)}"'
-  elif (
-      not args.enable_pathways
-      and args.enable_autoprovisioning
-      and autoprovisioning_config
-  ):
-    # Currently autoprovisioning is not supported with Pathways.
-    # Auto provisioning will have variable topologies for a gke accelerator type.
+  elif args.enable_autoprovisioning and autoprovisioning_config:
     resources_data = (
         f'{system.gke_accelerator}: {AUTOPROVISIONING_CONFIG_VALUE}'
     )

src/xpk/parser/cluster.py

@@ -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(

src/xpk/parser/workload.py

@@ -193,43 +193,6 @@ def set_workload_parsers(workload_parser):
       ),
   )
 
-  # Autoprovisioning workload arguments
-  workload_create_autoprovisioning_arguments.add_argument(
-      '--on-demand',
-      action='store_true',
-      help=(
-          'Sets autoprovisioning to use on-demand resources for the workload'
-          ' request. See `--reservation` or `--spot` for other capacity types.'
-      ),
-  )
-  workload_create_autoprovisioning_arguments.add_argument(
-      '--reservation',
-      type=str,
-      help=(
-          'Sets autoprovisioning to use reservation resources for the workload'
-          ' request. This will attempt to find the provided reservation. See'
-          ' `--spot`, `--flex` or `--on-demand` for other capacity types.'
-      ),
-  )
-  workload_create_autoprovisioning_arguments.add_argument(
-      '--spot',
-      action='store_true',
-      help=(
-          'Sets autoprovisioning to use spot resources. See `--reservation`,'
-          ' `--flex` or `--on-demand` for other capacity types.'
-      ),
-  )
-
-  workload_create_autoprovisioning_arguments.add_argument(
-      '--flex',
-      action='store_true',
-      help=(
-          'Sets autoprovisioning to use flex-start resources. See'
-          ' `--reservation`, `--spot` or `--on-demand` for other capacity'
-          ' types.'
-      ),
-  )
-
   # "workload create-pathways" command parser.
   workload_create_pathways_parser = workload_subcommands.add_parser(
       'create-pathways', help='Create a new job.'
@@ -257,6 +220,12 @@ def set_workload_parsers(workload_parser):
       '`--base-docker-image` is used by default. Set this argument if the'
       ' user wants the docker image to be used directly by the xpk workload.',
   )
+  workload_create_pathways_autoprovisioning_arguments = (
+      workload_create_pathways_parser.add_argument_group(
+          'Optional Autoprovisioning Arguments',
+          'Arguments for configuring autoprovisioning.',
+      )
+  )
   workload_create_pathways_vertex_tensorboard_arguments = (
       workload_create_pathways_parser.add_argument_group(
           'Vertex Tensorboard Arguments',
@@ -407,6 +376,10 @@ def set_workload_parsers(workload_parser):
       workload_vertex_tensorboard_arguments,
       workload_create_pathways_vertex_tensorboard_arguments,
   ])
+  add_shared_workload_create_autoprovisioning_arguments([
+      workload_create_autoprovisioning_arguments,
+      workload_create_pathways_autoprovisioning_arguments,
+  ])
 
   # Set defaults for both workload create and workload create-pathways after adding all shared args.
   workload_create_parser.set_defaults(func=workload_create)
@@ -770,3 +743,48 @@ def add_shared_workload_create_tensorboard_arguments(args_parsers):
             '<cluster>-<workload> will be created.'
         ),
     )
+
+
+def add_shared_workload_create_autoprovisioning_arguments(args_parsers):
+  """Add shared autoprovisioning arguments
+
+  Args:
+      List of workload create optional arguments parsers
+  """
+  for custom_parser in args_parsers:
+    custom_parser.add_argument(
+        '--on-demand',
+        action='store_true',
+        help=(
+            'Sets autoprovisioning to use on-demand resources for the workload'
+            ' request. See `--reservation` or `--spot` for other capacity'
+            ' types.'
+        ),
+    )
+    custom_parser.add_argument(
+        '--reservation',
+        type=str,
+        help=(
+            'Sets autoprovisioning to use reservation resources for the'
+            ' workload request. This will attempt to find the provided'
+            ' reservation. See `--spot`, `--flex` or `--on-demand` for other'
+            ' capacity types.'
+        ),
+    )
+    custom_parser.add_argument(
+        '--spot',
+        action='store_true',
+        help=(
+            'Sets autoprovisioning to use spot resources. See `--reservation`,'
+            ' `--flex` or `--on-demand` for other capacity types.'
+        ),
+    )
+    custom_parser.add_argument(
+        '--flex',
+        action='store_true',
+        help=(
+            'Sets autoprovisioning to use flex-start resources. See'
+            ' `--reservation`, `--spot` or `--on-demand` for other capacity'
+            ' types.'
+        ),
+    )