Add Remote Python Sidecar (#326)

* Add Remote Python Sidecar * Add comment * Add comment * Remote Python image addition * Update Testing Script * Update test script * Remove test * Changes * Clean up * Remove text * Fix comment * Update whitespace * import fix * Format * Remove command
AI-Hypercomputer · Jan 28, 2025 · 24ff1a2 · 24ff1a2
1 parent 6ac1fc3
commit 24ff1a2
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 0 deletions.
diff --git a/src/xpk/commands/workload.py b/src/xpk/commands/workload.py
@@ -52,6 +52,7 @@
     ensure_pathways_workload_prerequisites,
     get_pathways_proxy_args,
     get_pathways_rm_args,
+    get_pathways_sidecar_container,
     get_pathways_unified_query_link,
     get_pathways_worker_args,
     get_user_workload_for_pathways,
@@ -255,6 +256,7 @@
               volumeMounts:
               - mountPath: /tmp
                 name: shared-tmp
+            {pathways_sidecar_container}
             nodeSelector:
               {accelerator_label}
               {machine_label}
@@ -499,6 +501,7 @@ def workload_create(args) -> None:
         pathways_rm_args=get_pathways_rm_args(args, system),
         pathways_worker_args=get_pathways_worker_args(args),
         pathways_proxy_args=get_pathways_proxy_args(args),
+        pathways_sidecar_container=get_pathways_sidecar_container(args),
         user_workload=get_user_workload_for_pathways(args, system),
         resource_type=AcceleratorTypeToAcceleratorCharacteristics[
             system.accelerator_type

diff --git a/src/xpk/core/pathways.py b/src/xpk/core/pathways.py
@@ -67,6 +67,40 @@ def get_pathways_proxy_args(args) -> str:
     return ''
 
 
+def get_pathways_sidecar_container(args) -> str:
+  """This is a sidecar container that runs the remote python server.
+
+      It is a special case of the initContainer (designated by restartPolicy:
+      Always)
+      See https://kubernetes.io/docs/concepts/workloads/pods/sidecar-containers/
+      for more details.
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    str: yaml containing arguments for the Pathways sidecar container.
+  """
+  yaml = """initContainers:
+            - name: remote-python-sidecar
+              image: {args.remote_python_sidecar_image}
+              imagePullPolicy: Always
+              securityContext:
+                privileged: true
+              volumeMounts:
+              - mountPath: /tmp  # Shared volume mount with the main container.
+                name: shared-tmp
+              restartPolicy: Always
+              ports:
+              - containerPort: 50051
+              env:
+              - name: GRPC_SERVER_ADDRESS
+                value: '0.0.0.0:50051'"""
+  if args.use_pathways and args.remote_python_sidecar_image is not None:
+    return yaml.format(args=args)
+  else:
+    return ''
+
+
 def add_pw_resource_flavors(args):
   """Add resource flavors required for Pathways enabled clusters."""
   resource_flavor_yaml = """apiVersion: kueue.x-k8s.io/v1beta1

diff --git a/src/xpk/parser/workload.py b/src/xpk/parser/workload.py
@@ -503,6 +503,12 @@ 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=None,
+        help='Remote Python sidecar server image.',
+    )
     custom_parser.add_argument(
         '--enable-debug-logs',
         action='store_true',