Fail wait_for_job if a Pathways worker pod failed

kinetic/backend/pathways_client.py

@@ -162,6 +162,39 @@ def _raise_with_details(base_msg, core_v1, job_name, namespace):
   raise RuntimeError(msg)
 
 
+def _failed_worker_pod_names(core_v1, job_name, namespace):
+  """Return names of non-leader pods in the Failed phase."""
+  leader_pod_name = _get_leader_pod_name(job_name)
+  try:
+    pods = k8s_utils.list_job_pods(core_v1, job_name, namespace)
+  except ApiException:
+    return []
+  return [
+    pod.metadata.name
+    for pod in pods
+    if pod.metadata.name != leader_pod_name and pod.status.phase == "Failed"
+  ]
+
+
+def _raise_if_worker_failed(core_v1, job_name, namespace, success_msg):
+  """Raise if any worker pod failed, otherwise log ``success_msg``.
+
+  LWS reports the leader's status, but a worker pod can fail while the
+  leader still terminates cleanly. Treat any failed worker as a failed
+  job so callers see the real outcome instead of a false success.
+  """
+  failed = _failed_worker_pod_names(core_v1, job_name, namespace)
+  if failed:
+    names = ", ".join(failed)
+    _raise_with_details(
+      f"Pathways job {job_name} failed: worker pod(s) {names} failed",
+      core_v1,
+      job_name,
+      namespace,
+    )
+  logging.info(success_msg)
+
+
 def wait_for_job(job_id, namespace="default", timeout=3600, poll_interval=10):
   """Wait for Pathways Job (LeaderWorkerSet) to complete."""
   core_v1 = k8s_utils.core_v1()
@@ -189,7 +222,12 @@ def wait_for_job(job_id, namespace="default", timeout=3600, poll_interval=10):
           logged_running = True
 
         if pod.status.phase == "Succeeded":
-          logging.info(f"[REMOTE] Job {job_name} completed successfully")
+          _raise_if_worker_failed(
+            core_v1,
+            job_name,
+            namespace,
+            f"[REMOTE] Job {job_name} completed successfully",
+          )
           return "success"
 
         if pod.status.phase == "Failed":
@@ -224,7 +262,12 @@ def wait_for_job(job_id, namespace="default", timeout=3600, poll_interval=10):
         # Check current state
         if container_status.state.terminated:
           if container_status.state.terminated.exit_code == 0:
-            logging.info(f"[REMOTE] Job {job_name} completed successfully")
+            _raise_if_worker_failed(
+              core_v1,
+              job_name,
+              namespace,
+              f"[REMOTE] Job {job_name} completed successfully",
+            )
             return "success"
           else:
             _raise_with_details(
@@ -237,8 +280,11 @@ def wait_for_job(job_id, namespace="default", timeout=3600, poll_interval=10):
         # Check last state (in case it restarted)
         if container_status.last_state.terminated:
           if container_status.last_state.terminated.exit_code == 0:
-            logging.info(
-              f"[REMOTE] Job {job_name} completed successfully (restarted)"
+            _raise_if_worker_failed(
+              core_v1,
+              job_name,
+              namespace,
+              f"[REMOTE] Job {job_name} completed successfully (restarted)",
             )
             return "success"
           else:

kinetic/backend/pathways_client_test.py

@@ -457,6 +457,17 @@ def setUp(self):
     self.mock_core = self.enterContext(
       mock.patch("kinetic.backend.k8s_utils.core_v1")
     ).return_value
+    # Default, no workers visible to the failure scan, so successful
+    # leader status returns success.
+    self.mock_core.list_namespaced_pod.return_value.items = []
+    # Stub failure-details collection so MagicMock pod fields don't leak
+    # into string formatting inside k8s_utils.
+    self.enterContext(
+      mock.patch(
+        "kinetic.backend.k8s_utils.collect_pod_failure_details",
+        return_value="",
+      )
+    )
 
     self.mock_streamer = MagicMock()
     self.enterContext(
@@ -465,12 +476,18 @@ def setUp(self):
     self.mock_streamer.__enter__ = MagicMock(return_value=self.mock_streamer)
     self.mock_streamer.__exit__ = MagicMock(return_value=False)
 
-  def _make_pod(self, phase, container_statuses=None):
+  def _make_pod(self, phase, container_statuses=None, name=None):
     pod = MagicMock()
     pod.status.phase = phase
     pod.status.container_statuses = container_statuses
+    if name is not None:
+      pod.metadata.name = name
     return pod
 
+  def _set_worker_pods(self, *pods):
+    """Set the pods returned by list_namespaced_pod (failure scan)."""
+    self.mock_core.list_namespaced_pod.return_value.items = list(pods)
+
   def _make_container_status(
     self, state_terminated=None, last_state_terminated=None
   ):
@@ -593,6 +610,65 @@ def test_no_streaming_when_pod_pending(self):
     self.assertEqual(result, "success")
     self.mock_streamer.start.assert_not_called()
 
+  def test_worker_failed_overrides_leader_success_phase(self):
+    """Regression for #239: leader Succeeded but a worker pod failed."""
+    self.mock_core.read_namespaced_pod.return_value = self._make_pod(
+      "Succeeded", name="keras-pathways-j1-0"
+    )
+    leader = self._make_pod("Succeeded", name="keras-pathways-j1-0")
+    worker = self._make_pod("Failed", name="keras-pathways-j1-1")
+    self._set_worker_pods(leader, worker)
+    with self.assertRaisesRegex(RuntimeError, "keras-pathways-j1-1"):
+      wait_for_job("j1")
+
+  def test_worker_failed_overrides_container_exit_zero(self):
+    """Container exit 0 must not mask a failed worker pod."""
+    cs = self._make_container_status(state_terminated=self._make_terminated(0))
+    pod = self._make_pod(
+      "Running", container_statuses=[cs], name="keras-pathways-j1-0"
+    )
+    self.mock_core.read_namespaced_pod.return_value = pod
+    worker = self._make_pod("Failed", name="keras-pathways-j1-2")
+    self._set_worker_pods(pod, worker)
+    with self.assertRaisesRegex(RuntimeError, "keras-pathways-j1-2"):
+      wait_for_job("j1")
+
+  def test_worker_failed_overrides_last_state_exit_zero(self):
+    """Restarted leader exit 0 must not mask a failed worker pod."""
+    cs = self._make_container_status(
+      state_terminated=None,
+      last_state_terminated=self._make_terminated(0),
+    )
+    pod = self._make_pod(
+      "Running", container_statuses=[cs], name="keras-pathways-j1-0"
+    )
+    self.mock_core.read_namespaced_pod.return_value = pod
+    worker = self._make_pod("Failed", name="keras-pathways-j1-1")
+    self._set_worker_pods(pod, worker)
+    with self.assertRaisesRegex(RuntimeError, "keras-pathways-j1-1"):
+      wait_for_job("j1")
+
+  def test_success_when_all_workers_healthy(self):
+    """Leader Succeeded with all workers Succeeded should return success."""
+    self.mock_core.read_namespaced_pod.return_value = self._make_pod(
+      "Succeeded", name="keras-pathways-j1-0"
+    )
+    leader = self._make_pod("Succeeded", name="keras-pathways-j1-0")
+    w1 = self._make_pod("Succeeded", name="keras-pathways-j1-1")
+    w2 = self._make_pod("Succeeded", name="keras-pathways-j1-2")
+    self._set_worker_pods(leader, w1, w2)
+    self.assertEqual(wait_for_job("j1"), "success")
+
+  def test_worker_failure_scan_api_error_does_not_block_success(self):
+    """If listing pods fails, fall back to leader status (no false failure)."""
+    self.mock_core.read_namespaced_pod.return_value = self._make_pod(
+      "Succeeded", name="keras-pathways-j1-0"
+    )
+    self.mock_core.list_namespaced_pod.side_effect = ApiException(
+      status=500, reason="Server Error"
+    )
+    self.assertEqual(wait_for_job("j1"), "success")
+
 
 class TestCleanupJob(absltest.TestCase):
   def setUp(self):