Enable system ptrace for pathways headnode

GitOrigin-RevId: 32af26e

Author: muyangyuapple

axlearn/cloud/gcp/pathways_utils.py

@@ -435,6 +435,9 @@ def _build_pathways_head_container(self) -> dict:
         }
         head_container["resources"] = resources
 
+        # This enables tracing tools like pyspy
+        head_container["securityContext"] = {"capabilities": {"add": ["SYS_PTRACE"]}}
+
         volume_mounts = head_container.get("volumeMounts", [])
         volume_mounts.append(dict(name="shared-memory", mountPath="/tmp/ifrt_proxy"))
         head_container["volumeMounts"] = volume_mounts
@@ -1238,6 +1241,8 @@ def _build_head_container(self) -> dict:
         container = dict(
             name=inner_cfg.name,
             image=inner_cfg.image_id or self._bundler.id(inner_cfg.name),
+            # This enables tracing tools like pyspy
+            securityContext={"capabilities": {"add": ["SYS_PTRACE"]}},
             command=["bash", "-c", command],
             env=[
                 {

axlearn/cloud/gcp/pathways_utils_test.py

@@ -155,6 +155,13 @@ def test_build_pathways_head_pod(self, instance_type, sidecars):
                 }.issubset(env_vars)
             )
 
+            # Check security context for SYS_PTRACE capability
+            self.assertIn("securityContext", head_container)
+            security_context = head_container["securityContext"]
+            self.assertIn("capabilities", security_context)
+            self.assertIn("add", security_context["capabilities"])
+            self.assertIn("SYS_PTRACE", security_context["capabilities"]["add"])
+
             # Check pathways-proxy container args for XLA flags.
             proxy_container = None
             rm_container = None