Add multi-container support for McJax workload create

Author: zxhe-sean

goldens/Workload_create.txt

@@ -16,6 +16,7 @@ kubectl get configmap golden-cluster-resources-configmap -o=custom-columns="Conf
 [XPK] No gcp parallelstore instances to add detected.
 [XPK] No gce persistent disk instances to add detected.
 [XPK] No managed lustre instances to add detected.
+Namespace(xpk_subcommands='workload', func=<function workload_create at 0x7fbd15b03100>, xpk_workload_subcommands='create', command='bash hello', tpu_type='v5p-8', device_type=None, storage=[], num_nodes=1, multi_container=False, scheduler='default-scheduler', ramdisk_directory='', mtc_enabled=False, debug_dump_gcs=None, deploy_stacktrace_sidecar=False, use_pathways=False, restart_on_exit_codes=None, workload='golden-workload', cluster='golden-cluster', project='golden-project', zone='us-central1-a', dry_run=True, skip_validation=False, quiet=False, docker_name='jax-tpu', output_manifest_file=None, num_slices=1, priority='medium', max_restarts='0', ttl_seconds_after_finished=43200, termination_grace_period_seconds='30', colocated_python_sidecar_image='', enable_debug_logs=False, env_file=None, env={}, base_docker_image='python:3.10', script_dir='/tmp', docker_image=None, docker_image_pull_secret=None, use_vertex_tensorboard=False, experiment_name=None, on_demand=False, reservation=None, spot=False, flex=False, enable_ray_cluster=False)
 [XPK] Temp file (4b6736a12db8ea0f78ce793fd0d4ee0c94c652303f1dc0fecad085ea0993f688) content: 
 FROM python:3.10
 

goldens/Workload_create_pathways.txt

@@ -11,6 +11,7 @@ kubectl get configmap golden-cluster-metadata-configmap -o=custom-columns="Confi
 [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. 
 kubectl get configmap golden-cluster-resources-configmap -o=custom-columns="ConfigData:data" --no-headers=true
 [XPK] gke_accelerator type not found in config map. Autoprovisioning is not enabled.
+Namespace(xpk_subcommands='workload', func=<function workload_create_pathways at 0x7fa65b206fc0>, xpk_workload_subcommands='create-pathways', tpu_type='v5p-8', headless=False, proxy_server_image='', server_image='', pathways_gcs_location='gs://cloud-pathways-staging/tmp', command='bash hello', storage=[], custom_pathways_server_args='', custom_pathways_proxy_server_args='', custom_pathways_worker_args='', elastic_slices=0, max_slice_restarts=1, workload='golden-workload', cluster='golden-cluster', project='golden-project', zone='us-central1-a', dry_run=True, skip_validation=False, quiet=False, docker_name='jax-tpu', output_manifest_file=None, num_slices=1, priority='medium', max_restarts='0', ttl_seconds_after_finished=43200, termination_grace_period_seconds='30', colocated_python_sidecar_image='', enable_debug_logs=False, env_file=None, env={}, base_docker_image='python:3.10', script_dir='/tmp', docker_image=None, docker_image_pull_secret=None, use_vertex_tensorboard=False, experiment_name=None, on_demand=False, reservation=None, spot=False, flex=False, enable_ray_cluster=False, use_pathways=True)
 [XPK] Task: `Check if PathwaysJob is installed on golden-cluster` is implemented by the following command not running since it is a dry run. 
 kubectl get pods -n pathways-job-system --no-headers -o custom-columns=NAME:.metadata.name
 [XPK] check_if_pathways_job_is_installed 0 0

goldens/Workload_create_sub-slicing.txt

@@ -19,6 +19,7 @@ kubectl get configmap golden-cluster-resources-configmap -o=custom-columns="Conf
 [XPK] No gcp parallelstore instances to add detected.
 [XPK] No gce persistent disk instances to add detected.
 [XPK] No managed lustre instances to add detected.
+Namespace(xpk_subcommands='workload', func=<function workload_create at 0x7f750e2cf100>, xpk_workload_subcommands='create', command='bash hello', tpu_type='v6e-2x4', device_type=None, storage=[], num_nodes=1, multi_container=False, scheduler='default-scheduler', ramdisk_directory='', mtc_enabled=False, debug_dump_gcs=None, deploy_stacktrace_sidecar=False, use_pathways=False, restart_on_exit_codes=None, workload='golden-workload', cluster='golden-cluster', project='golden-project', zone='us-central1-a', dry_run=True, skip_validation=False, quiet=False, docker_name='jax-tpu', output_manifest_file=None, num_slices=1, priority='medium', max_restarts='0', ttl_seconds_after_finished=43200, termination_grace_period_seconds='30', colocated_python_sidecar_image='', enable_debug_logs=False, env_file=None, env={}, base_docker_image='python:3.10', script_dir='/tmp', docker_image=None, docker_image_pull_secret=None, use_vertex_tensorboard=False, experiment_name=None, on_demand=False, reservation=None, spot=False, flex=False, enable_ray_cluster=False)
 [XPK] Workload will be scheduled using the Sub-slicing feature.
 [XPK] Temp file (4b6736a12db8ea0f78ce793fd0d4ee0c94c652303f1dc0fecad085ea0993f688) content: 
 FROM python:3.10

goldens/Workload_create_super-slicing.txt

@@ -19,6 +19,7 @@ kubectl get configmap golden-cluster-resources-configmap -o=custom-columns="Conf
 [XPK] No gcp parallelstore instances to add detected.
 [XPK] No gce persistent disk instances to add detected.
 [XPK] No managed lustre instances to add detected.
+Namespace(xpk_subcommands='workload', func=<function workload_create at 0x7f1062083100>, xpk_workload_subcommands='create', command='bash hello', tpu_type='tpu7x-4x4x20', device_type=None, storage=[], num_nodes=1, multi_container=False, scheduler='default-scheduler', ramdisk_directory='', mtc_enabled=False, debug_dump_gcs=None, deploy_stacktrace_sidecar=False, use_pathways=False, restart_on_exit_codes=None, workload='golden-workload', cluster='golden-cluster', project='golden-project', zone='us-central1-a', dry_run=True, skip_validation=False, quiet=False, docker_name='jax-tpu', output_manifest_file=None, num_slices=1, priority='medium', max_restarts='0', ttl_seconds_after_finished=43200, termination_grace_period_seconds='30', colocated_python_sidecar_image='', enable_debug_logs=False, env_file=None, env={}, base_docker_image='python:3.10', script_dir='/tmp', docker_image=None, docker_image_pull_secret=None, use_vertex_tensorboard=False, experiment_name=None, on_demand=False, reservation=None, spot=False, flex=False, enable_ray_cluster=False)
 [XPK] Workload will be scheduled using the Super-slicing feature.
 [XPK] Temp file (4b6736a12db8ea0f78ce793fd0d4ee0c94c652303f1dc0fecad085ea0993f688) content: 
 FROM python:3.10

goldens/Workload_create_with_output-manifest-file.txt

@@ -16,6 +16,7 @@ kubectl get configmap golden-cluster-resources-configmap -o=custom-columns="Conf
 [XPK] No gcp parallelstore instances to add detected.
 [XPK] No gce persistent disk instances to add detected.
 [XPK] No managed lustre instances to add detected.
+Namespace(xpk_subcommands='workload', func=<function workload_create at 0x7f0fae8cf100>, xpk_workload_subcommands='create', command='bash hello', tpu_type='v5p-8', device_type=None, storage=[], num_nodes=1, multi_container=False, scheduler='default-scheduler', ramdisk_directory='', mtc_enabled=False, debug_dump_gcs=None, deploy_stacktrace_sidecar=False, use_pathways=False, restart_on_exit_codes=None, workload='golden-workload', cluster='golden-cluster', project='golden-project', zone='us-central1-a', dry_run=True, skip_validation=False, quiet=False, docker_name='jax-tpu', output_manifest_file='/var/tmp/manifest.yaml', num_slices=1, priority='medium', max_restarts='0', ttl_seconds_after_finished=43200, termination_grace_period_seconds='30', colocated_python_sidecar_image='', enable_debug_logs=False, env_file=None, env={}, base_docker_image='python:3.10', script_dir='/tmp', docker_image=None, docker_image_pull_secret=None, use_vertex_tensorboard=False, experiment_name=None, on_demand=False, reservation=None, spot=False, flex=False, enable_ray_cluster=False)
 [XPK] Temp file (4b6736a12db8ea0f78ce793fd0d4ee0c94c652303f1dc0fecad085ea0993f688) content: 
 FROM python:3.10
 

pathways-job

@@ -0,0 +1 @@
+Subproject commit 2880c34e7d71596664bafa1c3cecb5754a9991e7

src/xpk/commands/workload.py

@@ -482,14 +482,31 @@ def workload_create(args) -> None:
   # Currently failure policy rules are supported for Pathways workloads. b/408465881
   failure_policy_rules = ''
   pod_failure_policy = ''
+  print(args)
   if not args.use_pathways:
     failure_policy_rules = """rules:
       - action: FailJobSet
         onJobFailureReasons:
         - PodFailurePolicy"""
     restart_on_exit_codes_list = get_restart_exit_codes(args)
     restart_on_exit_codes = ','.join(map(str, restart_on_exit_codes_list))
-    pod_failure_policy = f"""
+    if args.multi_container:
+      pod_failure_policy = f"""
+          podFailurePolicy:
+            rules:
+            - action: FailJob
+              onExitCodes:
+                containerName: {get_main_container_docker_image(args, workload_system)}-1
+                operator: NotIn
+                values: [{restart_on_exit_codes}]
+            - action: FailJob
+              onExitCodes:
+                containerName: {get_main_container_docker_image(args, workload_system)}-2
+                operator: NotIn
+                values: [{restart_on_exit_codes}]"""
+
+    else:
+      pod_failure_policy = f"""
           podFailurePolicy:
             rules:
             - action: FailJob

src/xpk/core/docker_container.py

@@ -112,6 +112,65 @@ def get_main_container(args, system, docker_image, resource_type) -> str:
         'touch /shared-volume/stacktrace_signal; '
     )
 
+  if not args.use_pathways and args.multi_container:
+    containers = []
+    for i in range(2):
+      container_yaml = """
+              - name: {docker_name}
+                image: {docker_image}
+                {image_pull_policy}
+                env: {env}
+                securityContext:
+                  privileged: true
+                command:
+                - bash
+                - -c
+                - |
+                  echo XPK Start: $(date);
+                  _sigterm() (kill -SIGTERM $! 2>/dev/null;);
+                  trap _sigterm SIGTERM;
+                  {gsutil_test_command}
+                  ({command}) & PID=$!;
+                  while kill -0 $PID 2>/dev/null;
+                      do sleep 5;
+                  done;
+                  wait $PID;
+                  EXIT_CODE=$?;
+                  {xpk_internal_commands}
+                  echo XPK End: $(date);
+                  echo EXIT_CODE=$EXIT_CODE;
+                  {tpu_stacktrace_terminate_command}
+                  {gpu_workload_terminate_command}
+                  exit $EXIT_CODE
+                resources:
+                  limits:
+                    {resources}
+"""
+      volume_mounts = get_volume_mounts(args, system)
+      if volume_mounts != '':
+        container_yaml += """
+                volumeMounts:
+                {volume_mounts}
+"""
+      containers.append(
+          container_yaml.format(
+              args=args,
+              system=system,
+              image_pull_policy=add_image_pull_policy_for_pw_or_gpu(args, system),
+              env=get_env_container(args, system),
+              docker_name=f'jax-tpu-{i+1}',
+              docker_image=docker_image,
+              gsutil_test_command=gsutil_test_command,
+              command=command,
+              tpu_stacktrace_terminate_command=tpu_stacktrace_terminate_command,
+              gpu_workload_terminate_command=gpu_workload_terminate_command,
+              xpk_internal_commands=xpk_internal_commands,
+              resources=f'{resource_type}: {int(system.chips_per_vm / 2)}',
+              volume_mounts=volume_mounts,
+          )
+      )
+    return ''.join(containers)
+
   yaml = """- name: {docker_name}
                 image: {docker_image}
                 {image_pull_policy}

src/xpk/parser/workload.py

@@ -131,6 +131,11 @@ def set_workload_create_parser(workload_create_parser: ArgumentParser):
       default=1,
       help='The number of nodes to use, default=1.',
   )
+  workload_create_parser_optional_arguments.add_argument(
+      '--multi-container',
+      action='store_true',
+      help='Enable multi-container workload.',
+  )
   workload_create_parser_optional_arguments.add_argument(
       '--scheduler',
       type=str,