Add TCP readiness probes to Pathways head containers and scope custom libtpu to Pathways containers.

pathwaysutils/experimental/gke/jobset.py

@@ -195,7 +195,8 @@ def _apply_container_kwargs(
     container.env = merged_env
 
   for field, value in kwargs.items():
-    if getattr(container, field) is not None:
+    existing = getattr(container, field)
+    if existing is not None and existing != value:
       _logger.warning(
           "Replacing the default %s of the %s container.",
           field,
@@ -314,6 +315,10 @@ def _build_rm_container(
           ),
           client.V1ContainerPort(container_port=29002, protocol="TCP"),
       ],
+      readiness_probe=client.V1Probe(
+          tcp_socket=client.V1TCPSocketAction(port=PATHWAYS_RM_PORT),
+          period_seconds=1,
+      ),
       resources=client.V1ResourceRequirements(
           limits={"cpu": "8", "memory": "32G"}
       ),
@@ -375,6 +380,10 @@ def _build_proxy_container(
               container_port=PATHWAYS_PROXY_PORT, protocol="TCP"
           )
       ],
+      readiness_probe=client.V1Probe(
+          tcp_socket=client.V1TCPSocketAction(port=PATHWAYS_PROXY_PORT),
+          period_seconds=1,
+      ),
       resources=client.V1ResourceRequirements(
           limits={"cpu": "16", "memory": "100G"}
       ),
@@ -1153,12 +1162,20 @@ def with_custom_libtpu(
       )
     path = posixpath.normpath(stripped)
 
+    target_containers: str | Sequence[str] = (
+        ("pathways-rm", "pathways-proxy", "pathways-worker")
+        if containers == "all"
+        else containers
+    )
+
     for job_template in (self._head_job_template, self._worker_job_template):
       pod_spec = job_template.spec.template.spec
       all_containers = (pod_spec.containers or []) + (
           pod_spec.init_containers or []
       )
-      matching = self._filter_matching_containers(containers, all_containers)
+      matching = self._filter_matching_containers(
+          target_containers, all_containers
+      )
       if not matching:
         continue
 

pathwaysutils/test/experimental/gke/jobset_test.py

@@ -169,12 +169,26 @@ def test_headless_head_job_containers(self):
         rm_container["image"],
         "us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest",
     )
+    self.assertEqual(
+        rm_container["readinessProbe"],
+        {
+            "periodSeconds": 1,
+            "tcpSocket": {"port": jobset.PATHWAYS_RM_PORT},
+        },
+    )
 
     proxy_container = helper.containers["pathways-head"]["pathways-proxy"]
     self.assertEqual(
         proxy_container["image"],
         "us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest",
     )
+    self.assertEqual(
+        proxy_container["readinessProbe"],
+        {
+            "periodSeconds": 1,
+            "tcpSocket": {"port": jobset.PATHWAYS_PROXY_PORT},
+        },
+    )
     self.assertIn("--num_elastic_slices=2", proxy_container["args"])
     self.assertTrue(
         any(
@@ -1429,9 +1443,17 @@ def test_container_kwargs_warns_when_replacing_default_field(self):
   def test_container_kwargs_no_warning_for_unset_field(self):
     with self.assertNoLogs(jobset._logger, level="WARNING"):
       self._create_jobset(
+          proxy_container_kwargs={
+              "readiness_probe": client.V1Probe(
+                  tcp_socket=client.V1TCPSocketAction(
+                      port=jobset.PATHWAYS_PROXY_PORT
+                  ),
+                  period_seconds=1,
+              ),
+          },
           worker_container_kwargs={
               "security_context": client.V1SecurityContext(privileged=True)
-          }
+          },
       )
 
   def test_worker_pod_spec_volumes_keep_shared_tmp(self):
@@ -1744,6 +1766,48 @@ def test_with_custom_libtpu_replaces_existing_env(self):
     self.assertIn("PATHWAYS_HEAD", [e["name"] for e in worker["env"]])
     self.assertIn("MEGASCALE_SLICE_ID", [e["name"] for e in worker["env"]])
 
+  def test_with_custom_libtpu_all_skips_non_pathways_containers(self):
+    stage_init = client.V1Container(name="stage-libtpu", image="cloud-sdk:slim")
+    pw_jobset = self._create_jobset(
+        topology="2x2",
+        num_slices=1,
+        head_pod_spec_kwargs={"init_containers": [stage_init]},
+        worker_pod_spec_kwargs={
+            "init_containers": [
+                client.V1Container(name="stage-libtpu", image="cloud-sdk:slim")
+            ]
+        },
+    )
+    pw_jobset.add_colocated_python(
+        image="gcr.io/my-project/colocated-python:latest"
+    )
+    pw_jobset.add_user_workload(
+        image="us-docker.pkg.dev/my-project/test:v1",
+        command="python3 test.py",
+    )
+    pw_jobset.with_custom_libtpu("/tmp/custom_libtpu/libtpu.so")
+
+    helper = JobSetManifestHelper(pw_jobset.to_dict())
+    for container_name in ("pathways-rm", "pathways-proxy", "pathways-worker"):
+      matches = helper.get_all_containers_by_name(container_name)
+      self.assertLen(matches, 1, container_name)
+      _, container = matches[0]
+      self.assertIn(
+          {"name": "TPU_LIBRARY_PATH", "value": "/tmp/custom_libtpu/libtpu.so"},
+          container["env"],
+      )
+
+    for non_pw_name in (
+        "stage-libtpu",
+        "user-workload",
+        "colocated-python-sidecar",
+    ):
+      matches = helper.get_all_containers_by_name(non_pw_name)
+      self.assertNotEmpty(matches, non_pw_name)
+      for _, container in matches:
+        env_names = [e["name"] for e in container.get("env") or []]
+        self.assertNotIn("TPU_LIBRARY_PATH", env_names, non_pw_name)
+
   def test_add_gcsfuse_replaces_duplicate_mount_path(self):
     pw_jobset = self._create_jobset(topology="2x2", num_slices=1)
     pw_jobset.add_gcsfuse(