From 39e0b042a2de975139cff070a592fbcefbf60c16 Mon Sep 17 00:00:00 2001
From: Sujeeth Jinesh <sujinesh@google.com>
Date: Tue, 14 Jan 2025 17:17:47 +0000
Subject: [PATCH] Add Remote Python Sidecar

---
 src/xpk/commands/workload.py | 20 ++++++++++++++++++++
 src/xpk/parser/workload.py   |  8 ++++++++
 2 files changed, 28 insertions(+)

diff --git a/src/xpk/commands/workload.py b/src/xpk/commands/workload.py
index bc26e4e8..75d6c0b1 100644
--- a/src/xpk/commands/workload.py
+++ b/src/xpk/commands/workload.py
@@ -254,6 +254,26 @@
               volumeMounts:
               - mountPath: /tmp
                 name: shared-tmp
+            initContainers:
+            - args:
+              name: remote-python-sidecar
+              image: {args.remote_python_sidecar_image}
+              imagePullPolicy: Always
+              command:
+              - "bash"
+              - "-c"
+              - |
+                python start_remote_python.py
+              securityContext:
+                privileged: true
+              volumeMounts:
+              - mountPath: /tmp
+                name: shared-tmp
+              ports:
+                - containerPort: 50051
+                env:
+                - name: GRPC_SERVER_ADDRESS
+                  value: "0.0.0.0:50051"
             nodeSelector:
               {accelerator_label}
               {machine_label}
diff --git a/src/xpk/parser/workload.py b/src/xpk/parser/workload.py
index f12b4224..00726ccf 100644
--- a/src/xpk/parser/workload.py
+++ b/src/xpk/parser/workload.py
@@ -503,6 +503,14 @@ def add_shared_workload_create_optional_arguments(args_parsers):
             ' event or deletion request.Defaults to 30 seconds.'
         ),
     )
+    custom_parser.add_argument(
+        '--remote-python-sidecar-image',
+        type=str,
+        default='TempImageLocation',
+        help=(
+            'Remote Python sidecar image to use.'
+        ),
+    )
     custom_parser.add_argument(
         '--enable-debug-logs',
         action='store_true',