add some env vars to jobset

GitOrigin-RevId: 67dd26c

Author: muyangyuapple

axlearn/cloud/gcp/jobset_utils.py

@@ -502,6 +502,29 @@ def _build_container(self) -> Nested[Any]:
         k8s_env_vars.append(
             {"name": "NODE_NAME", "valueFrom": {"fieldRef": {"fieldPath": "spec.nodeName"}}}
         )
+        # pylint: disable=line-too-long
+        k8s_env_vars.append(
+            {
+                "name": "NUM_REPLICAS",
+                "valueFrom": {
+                    "fieldRef": {
+                        "fieldPath": "metadata.annotations['jobset.sigs.k8s.io/replicatedjob-replicas']"
+                    }
+                },
+            }
+        )
+        # pylint: enable=line-too-long
+
+        k8s_env_vars.append(
+            {
+                "name": "REPLICA_ID",
+                "valueFrom": {
+                    "fieldRef": {
+                        "fieldPath": "metadata.annotations['jobset.sigs.k8s.io/job-index']"
+                    }
+                },
+            }
+        )
 
         return dict(
             name=cfg.name,

axlearn/cloud/gcp/jobset_utils_test.py

@@ -359,6 +359,18 @@ def test_build_pod(
                 container_env["NODE_IP"]["valueFrom"]["fieldRef"]["fieldPath"],
             )
 
+            # Verify NUM_REPLICAS in container env.
+            self.assertEqual(
+                "metadata.annotations['jobset.sigs.k8s.io/replicatedjob-replicas']",
+                container_env["NUM_REPLICAS"]["valueFrom"]["fieldRef"]["fieldPath"],
+            )
+
+            # Verify REPLICA_ID in container env.
+            self.assertEqual(
+                "metadata.annotations['jobset.sigs.k8s.io/job-index']",
+                container_env["REPLICA_ID"]["valueFrom"]["fieldRef"]["fieldPath"],
+            )
+
             # Verify uploader container specs
             self.assertEqual(len(pod_spec["initContainers"]), 1)