From 4fb0356993c8a2a73016be14e4e9eee17209b393 Mon Sep 17 00:00:00 2001
From: Sadi Kneipp <ksadi@google.com>
Date: Fri, 15 Nov 2024 08:16:27 +0000
Subject: [PATCH 1/2] add custom pw proxy args support

---
 src/xpk/core/pathways.py   |  4 ++++
 src/xpk/parser/workload.py | 10 ++++++++++
 2 files changed, 14 insertions(+)

diff --git a/src/xpk/core/pathways.py b/src/xpk/core/pathways.py
index 71ed27ea..af0b8a70 100644
--- a/src/xpk/core/pathways.py
+++ b/src/xpk/core/pathways.py
@@ -14,6 +14,8 @@
 limitations under the License.
 """
 
+import logging
+
 from ..utils.console import xpk_exit, xpk_print
 from .core import (
     AcceleratorType,
@@ -60,6 +62,8 @@ def get_pathways_proxy_args(args) -> str:
   yaml = """- --server_port=29000
               - --resource_manager_address={rm_address}
               - --gcs_scratch_location={args.pathways_gcs_location}"""
+  if args.additional_pw_proxy_args:
+    yaml += args.additional_pw_proxy_args.replace(' ', '\n' + ' '*14 +'- ')
 
   if args.use_pathways:
     return yaml.format(args=args, rm_address=get_rm_address(args))
diff --git a/src/xpk/parser/workload.py b/src/xpk/parser/workload.py
index fecebef9..adf1d189 100644
--- a/src/xpk/parser/workload.py
+++ b/src/xpk/parser/workload.py
@@ -244,6 +244,16 @@ def set_workload_parsers(workload_parser):
       required=False,
   )
 
+  workload_pathways_workload_arguments.add_argument(
+      '--additional_pw_proxy_args',
+      type=str,
+      default='',
+      help=(
+          'Additional arguments to pass to the proxy server.'
+      ),
+      required=False,
+  )
+
   add_shared_workload_create_required_arguments([
       workload_create_parser_required_arguments,
       workload_create_pathways_parser_required_arguments,

From 827113b57b4d86b22e22012f713f9bb0fc16db43 Mon Sep 17 00:00:00 2001
From: Sadi Kneipp <ksadi@google.com>
Date: Fri, 15 Nov 2024 08:27:10 +0000
Subject: [PATCH 2/2] remove unused import

---
 src/xpk/core/pathways.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/src/xpk/core/pathways.py b/src/xpk/core/pathways.py
index af0b8a70..c17b4c8f 100644
--- a/src/xpk/core/pathways.py
+++ b/src/xpk/core/pathways.py
@@ -14,8 +14,6 @@
 limitations under the License.
 """
 
-import logging
-
 from ..utils.console import xpk_exit, xpk_print
 from .core import (
     AcceleratorType,