Refactor pathways workload scheduling to use Jinja template (#1116)

* Refactor pathways workload scheduling to use Jinja template

* Refactor jinja injection and cleanup workload.py

Author: jamOne-

recipes/Workload_create_pathways.md

@@ -42,7 +42,7 @@ docker buildx build --platform=linux/amd64 -f 4b6736a12db8ea0f78ce793fd0d4ee0c94
 docker tag dry-run-runner gcr.io/golden-project/dry-run-runner:prefix-current
 [XPK] Task: `Upload Docker Image` is implemented by the following command not running since it is a dry run. 
 docker push gcr.io/golden-project/dry-run-runner:prefix-current
-[XPK] Temp file (5394ec8d9ca40eb8e048844d5622be8da9ac27cc9565535366c11a379ea35f58) content: 
+[XPK] Temp file (8e311bca9f9f54ee09e88dbb12e7d20d536478aa69d990a59142ee6fb70da079) content: 
 apiVersion: jobset.x-k8s.io/v1alpha2
 kind: JobSet
 metadata:
@@ -76,7 +76,6 @@ spec:
             dnsPolicy: ClusterFirstWithHostNet
             nodeSelector:
               cloud.google.com/gke-nodepool: cpu-np
-              
             initContainers:
               - name: pathways-proxy
                 image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest
@@ -132,9 +131,7 @@ spec:
                     cpu: "8"
                     memory: 32G
                 restartPolicy: Always
-
             containers:
-
               - name: jax-tpu
                 image: gcr.io/golden-project/dry-run-runner:prefix-current
                 imagePullPolicy: Always
@@ -209,8 +206,6 @@ spec:
             nodeSelector:
               cloud.google.com/gke-tpu-accelerator: tpu-v5p-slice
               cloud.google.com/gke-tpu-topology: 2x2x1
-              
-              
             containers:
               - name: pathways-worker
                 image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest
@@ -282,7 +277,7 @@ spec:
   suspend: false
 
 [XPK] Task: `Creating Workload` is implemented by the following command not running since it is a dry run. 
-kubectl apply -f 5394ec8d9ca40eb8e048844d5622be8da9ac27cc9565535366c11a379ea35f58
+kubectl apply -f 8e311bca9f9f54ee09e88dbb12e7d20d536478aa69d990a59142ee6fb70da079
 [XPK] Task: `GKE Dashboard List` is implemented by the following command not running since it is a dry run. 
 gcloud monitoring dashboards list --project=golden-project --filter="displayName:'GKE - TPU Monitoring Dashboard'" --format="value(name)" --verbosity=error
 [XPK] Check statistics and outlier mode of GKE metrics here: https://console.cloud.google.com/monitoring/dashboards/builder/0?project=golden-project&f.rlabel.cluster_name.ClusterName=golden-cluster. To view the metric data for your workload, select golden-workload from the JobName filter on the dashboard.

src/xpk/commands/workload.py

@@ -15,6 +15,8 @@
 """
 
 import urllib
+import argparse
+from ..core.system_characteristics import SystemCharacteristics
 from ..core.blueprint.blueprint_generator import (
     a3high_device_type,
     a4x_device_types,
@@ -41,14 +43,9 @@
 )
 from ..core.network import get_cluster_subnetworks
 from ..core.pathways import (
-    append_custom_colocated_python_sidecar,
-    append_custom_pathways_proxy_server,
-    append_custom_pathways_server,
-    append_custom_pathways_worker,
     check_if_pathways_job_is_installed,
     ensure_pathways_workload_prerequisites,
     get_pathways_unified_query_link,
-    get_user_workload_for_pathways,
     try_to_delete_pathwaysjob_first,
 )
 from ..core.resources import get_cluster_capacity_type, get_cluster_system_characteristics_from_config_map
@@ -58,9 +55,7 @@
     ONE_TO_ONE_REPLICA_NODE_POOL_ASSIGNMENT_ANNOTATION,
     WorkloadScheduling,
     check_if_workload_can_schedule,
-    create_tpu_machine_type,
     create_tpu_slice_topology_annotation,
-    create_tpu_topology,
     get_cpu_affinity,
     get_gpu_scheduler,
     create_sub_slicing_annotations,
@@ -106,6 +101,8 @@
 from jinja2 import Environment, FileSystemLoader
 from ..utils.templates import get_templates_absolute_path
 
+_PATHWAYS_WORKLOAD_TEMPLATE = 'pathways_workload_create.yaml.j2'
+
 _SUPER_SLICING_WORKLOAD_NAME_LIMIT = 28
 """Maximum safe workload name length to avoid exceeding GCE's 63-character limit.
 
@@ -263,89 +260,90 @@
               containers:
               {container}
 """
-# The indentation of PW_WORKLOAD_CREATE_YAML is intentional to allow reusing the user workload container YAML.
-PW_WORKLOAD_CREATE_YAML = """apiVersion: jobset.x-k8s.io/v1alpha2
-kind: JobSet
-metadata:
-  name: {args.workload}
-  labels:
-    kueue.x-k8s.io/queue-name: {local_queue_name}  # Name of the LocalQueue
-    xpk.google.com/workload: {args.workload}
-spec:
-  coordinator:
-    replicatedJob: pathways-head
-  network:
-    enableDNSHostnames: true
-    publishNotReadyAddresses: true
-  failurePolicy:
-    restartStrategy: Recreate
-  replicatedJobs:
-  - name: pathways-head
-    replicas: 1
-    template:
-      spec:
-        backoffLimit: 0
-        completionMode: Indexed
-        completions: 1
-        parallelism: 1
-        template:
-          metadata:
-            annotations:
-              alpha.jobset.sigs.k8s.io/exclusive-topology: kubernetes.io/hostname
-          spec:
-            hostNetwork: true
-            dnsPolicy: ClusterFirstWithHostNet
-            nodeSelector:
-              cloud.google.com/gke-nodepool: cpu-np
-              {autoprovisioning_args}
-{pathways_head_containers}
-            restartPolicy: Never
-            volumes:
-            - hostPath:
-                path: /tmp
-                type: DirectoryOrCreate
-              name: shared-tmp
-  - name: worker
-    replicas: {args.num_slices}
-    template:
-      spec:
-        backoffLimit: {worker_backoff_limit}
-        completionMode: Indexed
-        completions: {vms_per_slice}
-        parallelism: {vms_per_slice}
-        template:
-          metadata:
-            labels:
-              xpk.google.com/workload: {args.workload}
-            annotations:
-              alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
-          spec:
-            hostNetwork: true
-            dnsPolicy: ClusterFirstWithHostNet
-            terminationGracePeriodSeconds: {args.termination_grace_period_seconds}
-            priorityClassName: {args.priority}
-            nodeSelector:
-              {accelerator_label}
-              {node_selector_machine_label}
-              {placement_policy_label}
-              {autoprovisioning_args}
-            containers:
-              {custom_pathways_worker}
-            restartPolicy: OnFailure
-            volumes:
-            - hostPath:
-                path: /tmp
-                type: DirectoryOrCreate
-              name: shared-tmp
-  startupPolicy:
-    startupPolicyOrder: InOrder
-  {success_policy}
-  suspend: false
-"""
 
 ARM_GPU_WORKLOAD_CREATE_JINJA_FILE = 'arm_gpu_workload_crate.yaml.j2'
 
 
+def _generate_pathways_workload_yaml(
+    args: argparse.Namespace,
+    workload_system: SystemCharacteristics,
+    parallel_containers: int,
+    placement_policy_label: str,
+    autoprovisioning_args: str | None,
+) -> str:
+  worker_backoff_limit = (
+      (args.max_slice_restarts * workload_system.vms_per_slice)
+      if getattr(args, 'elastic_slices', 0) > 0
+      else (workload_system.vms_per_slice * 4)
+  )
+
+  proxy_server_image = (
+      getattr(args, 'proxy_server_image', None)
+      or 'us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest'
+  )
+  server_image = (
+      getattr(args, 'server_image', None)
+      or 'us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest'
+  )
+  worker_image = getattr(args, 'worker_image', None) or server_image
+  instance_type = (
+      f'{workload_system.pathways_tpu_version}:{workload_system.topology}'
+      if workload_system.pathways_tpu_version
+      else workload_system.gce_machine_type
+  )
+  if args.headless:
+    user_workload_container = ''
+    user_workload_env_vars = []
+  else:
+    user_workload_container, _ = get_user_workload_container(
+        args, workload_system, parallel_containers
+    )
+
+    user_workload_env_vars = [
+        {
+            'name': 'PATHWAYS_HEAD',
+            'valueFrom': "metadata.labels['jobset.sigs.k8s.io/coordinator']",
+        },
+        {
+            'name': 'JAX_PLATFORMS',
+            'value': 'proxy',
+        },
+        {
+            'name': 'XCLOUD_ENVIRONMENT',
+            'value': 'GCP',
+        },
+        {
+            'name': 'JAX_BACKEND_TARGET',
+            'value': 'grpc://$(PATHWAYS_HEAD):29000',
+        },
+    ]
+
+  template_env = Environment(
+      loader=FileSystemLoader(searchpath=get_templates_absolute_path()),
+      trim_blocks=True,
+      lstrip_blocks=True,
+      keep_trailing_newline=True,
+  )
+  workload_create_yaml = template_env.get_template(_PATHWAYS_WORKLOAD_TEMPLATE)
+  return workload_create_yaml.render(
+      args=args,
+      local_queue_name=LOCAL_QUEUE_NAME,
+      proxy_server_image=proxy_server_image,
+      server_image=server_image,
+      instance_type=instance_type,
+      user_workload_container=user_workload_container,
+      user_workload_env_vars=user_workload_env_vars,
+      worker_backoff_limit=worker_backoff_limit,
+      vms_per_slice=workload_system.vms_per_slice,
+      workload_system=workload_system,
+      accelerator_label=create_accelerator_label(workload_system),
+      node_selector_machine_label=create_machine_label(workload_system),
+      placement_policy_label=placement_policy_label,
+      autoprovisioning_args=autoprovisioning_args,
+      worker_image=worker_image,
+  )
+
+
 def workload_create_pathways(args) -> None:
   """Run jobset apply command for a file, specifically for Pathways.
 
@@ -695,46 +693,12 @@ def workload_create(args) -> None:
   elif args.use_pathways and ensure_pathways_workload_prerequisites(
       args, workload_system
   ):
-    if args.headless:
-      pathways_head_containers = f"""            containers:
-{append_custom_pathways_proxy_server(args)}
-{append_custom_pathways_server(args, workload_system)}
-{append_custom_colocated_python_sidecar(args)}"""
-      success_policy = ''
-    else:
-      pathways_head_containers = f"""            initContainers:
-{append_custom_pathways_proxy_server(args)}
-{append_custom_pathways_server(args, workload_system)}
-{append_custom_colocated_python_sidecar(args)}
-            containers:
-{get_user_workload_for_pathways(args, workload_system, parallel_containers)}"""
-      success_policy = """successPolicy:
-    operator: All
-    targetReplicatedJobs:
-    - pathways-head"""
-
-    worker_backoff_limit = (
-        (args.max_slice_restarts * workload_system.vms_per_slice)
-        if getattr(args, 'elastic_slices', 0) > 0
-        else (workload_system.vms_per_slice * 4)
-    )
-
-    yml_string = PW_WORKLOAD_CREATE_YAML.format(
+    yml_string = _generate_pathways_workload_yaml(
         args=args,
-        topology=create_tpu_topology(workload_system),
-        machine_type=create_tpu_machine_type(workload_system),
-        pathways_head_containers=pathways_head_containers,
-        custom_pathways_worker=append_custom_pathways_worker(
-            args, workload_system
-        ),
-        worker_backoff_limit=worker_backoff_limit,
-        success_policy=success_policy,
-        local_queue_name=LOCAL_QUEUE_NAME,
-        autoprovisioning_args=autoprovisioning_args,
+        workload_system=workload_system,
+        parallel_containers=parallel_containers,
         placement_policy_label=placement_policy_label,
-        vms_per_slice=workload_system.vms_per_slice,
-        accelerator_label=create_accelerator_label(workload_system),
-        node_selector_machine_label=create_machine_label(workload_system),
+        autoprovisioning_args=autoprovisioning_args,
     )
   else:
     if use_sub_slicing:

src/xpk/commands/workload_test.py

@@ -329,7 +329,7 @@ def test_workload_create_pathways_jobset_yaml(mocker):
   args.use_vertex_tensorboard = False
   args.headless = False
   args.num_slices = 2
-  args.elastic_slices = 0
+  args.elastic_slices = 2
   args.max_restarts = 1
   args.max_slice_restarts = 1
   args.termination_grace_period_seconds = 30
@@ -382,16 +382,25 @@ def test_workload_create_pathways_jobset_yaml(mocker):
       return_value=True,
   )
   mocker.patch(
-      'xpk.core.pathways.get_user_workload_container',
-      return_value=('- name: test-docker\n  image: test-image', '123'),
-  )
-  mocker.patch('xpk.commands.workload.create_tpu_topology', return_value='4x4')
-  mocker.patch(
-      'xpk.commands.workload.create_tpu_machine_type',
-      return_value='ct4p-hightpu-4t',
+      'xpk.commands.workload.get_user_workload_container',
+      return_value=(
+          (
+              '- name: test-docker\n  image: test-image\n  env:\n    - name:'
+              ' FOO\n      value: BAR'
+          ),
+          '123',
+      ),
   )
 
-  mock_write_file = mocker.patch('builtins.open', mocker.mock_open())
+  real_open = open
+  m_open = mocker.mock_open()
+
+  def custom_open(file, *args, **kwargs):
+    if str(file) == 'pw_manifest.yaml':
+      return m_open(file, *args, **kwargs)
+    return real_open(file, *args, **kwargs)
+
+  mocker.patch('builtins.open', side_effect=custom_open)
 
   mocker.patch(
       'xpk.commands.workload.write_tmp_file', return_value='/tmp/test.yaml'
@@ -405,10 +414,8 @@ def test_workload_create_pathways_jobset_yaml(mocker):
 
   workload_create(args)
 
-  mock_write_file.assert_called_once_with(
-      'pw_manifest.yaml', 'w', encoding='utf-8'
-  )
-  written_content = mock_write_file.return_value.write.call_args[0][0]
+  m_open.assert_called_once_with('pw_manifest.yaml', 'w', encoding='utf-8')
+  written_content = m_open.return_value.write.call_args[0][0]
 
   assert 'apiVersion: jobset.x-k8s.io/v1alpha2' in written_content
   assert 'kind: JobSet' in written_content
@@ -422,6 +429,14 @@ def test_workload_create_pathways_jobset_yaml(mocker):
   assert '- name: pathways-worker' in written_content
   assert f'replicas: {args.num_slices}' in written_content  # worker replicas
 
+  # Assert custom arguments are correctly injected
+  assert '- --custom_proxy_arg' in written_content
+  assert '- --custom_server_arg' in written_content
+  assert '- --custom_worker_arg' in written_content
+
+  # Assert elastic_slices is rendered
+  assert '- --num_elastic_slices=2' in written_content
+
   # Assert newly migrated JobSet specifics
   assert 'coordinator:' in written_content
   assert 'replicatedJob: pathways-head' in written_content
@@ -431,4 +446,14 @@ def test_workload_create_pathways_jobset_yaml(mocker):
   assert 'completionMode: Indexed' in written_content
   assert 'startupPolicyOrder: InOrder' in written_content
   assert 'operator: All' in written_content
-  assert f'backoffLimit: {workload_system.vms_per_slice * 4}' in written_content
+  assert (
+      f'backoffLimit: {args.max_slice_restarts * workload_system.vms_per_slice}'
+      in written_content
+  )
+  assert f'image: {args.proxy_server_image}' in written_content
+  assert f'image: {args.server_image}' in written_content
+  assert f'image: {args.colocated_python_sidecar_image}' in written_content
+  assert f'image: {args.worker_image}' in written_content
+  assert (
+      f'--gcs_scratch_location={args.pathways_gcs_location}' in written_content
+  )