diff --git a/src/xpk/commands/workload.py b/src/xpk/commands/workload.py index 81648d1eb..f2f744a77 100644 --- 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 index 71ed27ea1..875b3e188 100644 --- 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 index 023c2edf4..efbadc13e 100644 --- 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',