Add Workaround for Megascale Num Slices Issue on PW (#375)

* Support custom Pathways args. (#364)

* Remove Pathways specific args from workload create flow. (#365)

* Remove Pathways specific optional args from workload create.

* Deprecate --use-pathways from workload create, set it in create-pathways.

* Remove --enable-pathways from cluster create flow.

* Adding --use-pathways and --enable-pathways args - they are used to determine Pathways flows.

* Add Workaround for Megascale Num Slices Issue on PW

---------

Co-authored-by: Roshani Narasimhan <roshanin@google.com>

Author: SujeethJinesh

src/xpk/commands/workload.py

@@ -298,6 +298,28 @@
                 - mountPath: /tmp
                   name: shared-tmp
                 {storage_volume_mounts}
+                env:
+                  # Workaround for v6e
+                  - name: MEGASCALE_GRPC_ENABLE_XOR_TRACER
+                    value: "false"
+                  - name: MEGASCALE_NUM_SLICES
+                    valueFrom:
+                      fieldRef:
+                        fieldPath: "metadata.labels['jobset.sigs.k8s.io/replicatedjob-replicas']"
+                  - name: JOBSET_NAME
+                    valueFrom:
+                      fieldRef:
+                        fieldPath: metadata.annotations['jobset.sigs.k8s.io/jobset-name']
+                  - name: REPLICATED_JOB_NAME
+                    valueFrom:
+                      fieldRef:
+                        fieldPath: metadata.annotations['jobset.sigs.k8s.io/replicatedjob-name']
+                  - name: MEGASCALE_SLICE_ID
+                    valueFrom:
+                      fieldRef:
+                        fieldPath: "metadata.labels['jobset.sigs.k8s.io/job-index']"
+                  - name: MEGASCALE_COORDINATOR_ADDRESS
+                    value: "$(JOBSET_NAME)-$(REPLICATED_JOB_NAME)-$(MEGASCALE_SLICE_ID)-0.$(JOBSET_NAME)"
               {pathways_sidecar_container}
               nodeSelector:
                 {accelerator_label}
@@ -397,6 +419,13 @@ def workload_create_pathways(args) -> None:
     0 if successful and 1 otherwise.
   """
   args.use_pathways = True
+  if args.headless:
+    xpk_print(
+        'Please use kubectl port forwarding to connect to the Pathways proxy.'
+        ' kubectl get pods kubectl port-forward <proxy-pod-name> 29000:29000'
+        ' JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 python'
+        " -c 'import pathwaysutils; import jax; print(jax.devices())'"
+    )
   workload_create(args)