Support custom Pathways args. (#364)

Author: RoshaniN

src/xpk/core/pathways.py

@@ -46,6 +46,8 @@ def get_pathways_worker_args(args) -> str:
                 - --resource_manager_address={rm_address}
                 - --gcs_scratch_location={args.pathways_gcs_location}"""
   if args.use_pathways:
+    if args.custom_pathways_worker_args:
+      yaml = append_custom_pathways_args(yaml, args.custom_pathways_worker_args)
     return yaml.format(args=args, rm_address=get_rm_address(args))
   else:
     return ''
@@ -64,6 +66,10 @@ def get_pathways_proxy_args(args) -> str:
                 - --gcs_scratch_location={args.pathways_gcs_location}"""
 
   if args.use_pathways:
+    if args.custom_pathways_proxy_server_args:
+      yaml = append_custom_pathways_args(
+          yaml, args.custom_pathways_proxy_server_args
+      )
     return yaml.format(args=args, rm_address=get_rm_address(args))
   else:
     return ''
@@ -233,6 +239,8 @@ def get_pathways_rm_args(args, system: SystemCharacteristics) -> str:
                 - --instance_count={instance_count}
                 - --instance_type={instance_type}"""
   if args.use_pathways:
+    if args.custom_pathways_server_args:
+      yaml = append_custom_pathways_args(yaml, args.custom_pathways_server_args)
     return yaml.format(
         args=args,
         instance_count=args.num_slices,
@@ -242,6 +250,28 @@ def get_pathways_rm_args(args, system: SystemCharacteristics) -> str:
     return ''
 
 
+def append_custom_pathways_args(yaml, custom_args) -> str:
+  """Append custom Pathways args to the YAML with proper indentation.
+
+  Args:
+      yaml (string): existing yaml containing args
+
+  Returns:
+      yaml (string): yaml with additional args appended.
+  """
+  second_line = yaml.split('\n')[1]
+  if (
+      not second_line
+  ):  # to cover edge case if only one arg remains, we would have to look at the entire YAML in this case.
+    return yaml
+  # Calculate the indentation based on the second line of existing YAML.
+  indentation = ' ' * (len(second_line) - len(second_line.lstrip()))
+  custom_args = custom_args.split(' ')
+  for arg in custom_args:
+    yaml += '\n' + indentation + '- ' + arg
+  return yaml
+
+
 def get_user_workload_for_pathways(
     args,
     system: SystemCharacteristics,

src/xpk/parser/workload.py

@@ -256,6 +256,39 @@ def set_workload_parsers(workload_parser):
       help='Names of storages the workload uses',
   )
 
+  workload_create_pathways_parser_optional_arguments.add_argument(
+      '--custom-pathways-server-args',
+      type=str,
+      default=None,
+      help=(
+          'Provide custom Pathways server args as follows -'
+          " --custom-pathways-server-args='--arg_1=xxx --arg2=yyy'"
+      ),
+      required=False,
+  )
+
+  workload_create_pathways_parser_optional_arguments.add_argument(
+      '--custom-pathways-proxy-server-args',
+      type=str,
+      default=None,
+      help=(
+          'Provide custom Pathways proxy server args as follows -'
+          " --custom-pathways-proxy-server-args='--arg_1=xxx --arg2=yyy'"
+      ),
+      required=False,
+  )
+
+  workload_create_pathways_parser_optional_arguments.add_argument(
+      '--custom-pathways-worker-args',
+      type=str,
+      default=None,
+      help=(
+          'Provide custom Pathways worker args as follows -'
+          " --custom-pathways-worker-args='--arg_1=xxx --arg2=yyy'"
+      ),
+      required=False,
+  )
+
   add_shared_workload_create_required_arguments([
       workload_create_parser_required_arguments,
       workload_create_pathways_parser_required_arguments,