Add multi-container support for McJax workload create

Author: zxhe-sean

pathways-job

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

src/xpk/commands/cluster_test.py

@@ -147,6 +147,7 @@ def construct_args(**kwargs: Any) -> Namespace:
       docker_image_pull_secret='',
       managed_mldiagnostics=False,
       output_manifest_file='',
+      multi_container=False,
   )
   args_dict.update(kwargs)
   return Namespace(**args_dict)

src/xpk/commands/workload.py

@@ -489,7 +489,23 @@ def workload_create(args) -> None:
         - 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/commands/workload_test.py

@@ -23,6 +23,7 @@
 from ..core.system_characteristics import DockerPlatform, SystemCharacteristics, AcceleratorType, UserFacingNameToSystemCharacteristics, GpuConfig
 from .workload import workload_create
 from .cluster_test import construct_args
+from ..core.docker_container import get_user_workload_container as real_get_user_workload_container
 
 
 SYSTEM_CHARACTERISTICS = SystemCharacteristics(
@@ -205,3 +206,76 @@ def test_workload_create_dry_run_with_output_file(mocker):
   written_content = mock_open.return_value.write.call_args[0][0]
   assert 'test-workload' in written_content
   assert 'cloud.google.com/gke-tpu-topology: 8x8' in written_content
+
+
+def test_workload_create_multi_container(
+    workload_create_mocks: _WorkloadCreateMocks,
+    mocker,
+):
+  """Tests that the generated YAML for a multi-container workload has correct pod failure policy and container structure."""
+
+  # Enable dry_run to prevent external calls like get_storages_to_mount -> gcloud
+  mocker.patch('xpk.utils.execution_context.dry_run', True)
+
+  # Mock dependencies required by get_user_workload_container -> get_main_container
+  mocker.patch(
+      'xpk.core.docker_container.setup_docker_image',
+      return_value=(0, 'dummy-image'),
+  )
+  mocker.patch(
+      'xpk.core.docker_container.get_gke_debugging_dashboard', return_value=None
+  )
+
+  # Use the real get_user_workload_container to test integration
+  workload_create_mocks.get_user_workload_container.side_effect = (
+      real_get_user_workload_container
+  )
+
+  # Use a system with chips_per_vm=4 to test resource splitting logic
+  system_chars = dataclasses.replace(SYSTEM_CHARACTERISTICS, chips_per_vm=4)
+
+  with patch(
+      'xpk.commands.workload.get_system_characteristics',
+      return_value=(system_chars, 0),
+  ):
+    args = construct_args(
+        workload='test-workload',
+        command='echo hello',
+        num_nodes=1,
+        restart_on_exit_codes=None,
+        multi_container=True,
+        docker_name='test-docker',
+        deploy_stacktrace_sidecar=False,
+        enable_debug_logs=False,
+        scheduler='default-scheduler',
+    )
+    workload_create(args)
+
+  assert workload_create_mocks.write_tmp_file.called
+  yaml_content = workload_create_mocks.write_tmp_file.call_args[0][0]
+  jobset = yaml.safe_load(yaml_content)
+
+  # Verify Pod Failure Policy
+  pod_failure_rules = jobset['spec']['replicatedJobs'][0]['template']['spec'][
+      'podFailurePolicy'
+  ]['rules']
+  # Should have 2 rules for multi_container
+  assert len(pod_failure_rules) == 2
+  assert pod_failure_rules[0]['onExitCodes']['containerName'].endswith('-1')
+  assert pod_failure_rules[1]['onExitCodes']['containerName'].endswith('-2')
+
+  # Verify Containers
+  # Navigate to the containers list in the YAML
+  containers = jobset['spec']['replicatedJobs'][0]['template']['spec'][
+      'template'
+  ]['spec']['containers']
+
+  assert len(containers) == 2
+  assert containers[0]['name'] == 'jax-tpu-1'
+  assert containers[0]['image'] == 'dummy-image'
+  # Check if resources are split correctly (4 chips / 2 containers = 2 chips)
+  assert containers[0]['resources']['limits']['google.com/tpu'] == 2
+
+  assert containers[1]['name'] == 'jax-tpu-2'
+  assert containers[1]['image'] == 'dummy-image'
+  assert containers[1]['resources']['limits']['google.com/tpu'] == 2

src/xpk/core/docker_container.py

@@ -112,6 +112,67 @@ 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,