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:
@@ -337,24 +383,28 @@ def get_job_status(job_name, namespace="default") -> JobStatus:
       )
     raise RuntimeError(f"Failed to read leader pod status: {e.reason}") from e
 
-  if pod.status.phase == "Succeeded":
+  # A leader Succeeded result must be downgraded to FAILED if any worker
+  # pod failed, otherwise async callers see the same false success that
+  # wait_for_job used to return.
+  def _succeeded_or_worker_failed() -> JobStatus:
+    if _failed_worker_pod_names(core_v1, job_name, namespace):
+      return JobStatus.FAILED
     return JobStatus.SUCCEEDED
+
+  if pod.status.phase == "Succeeded":
+    return _succeeded_or_worker_failed()
   if pod.status.phase == "Failed":
     return JobStatus.FAILED
   if pod.status.container_statuses:
     container_status = pod.status.container_statuses[0]
     if container_status.state.terminated:
-      return (
-        JobStatus.SUCCEEDED
-        if container_status.state.terminated.exit_code == 0
-        else JobStatus.FAILED
-      )
+      if container_status.state.terminated.exit_code == 0:
+        return _succeeded_or_worker_failed()
+      return JobStatus.FAILED
     if container_status.last_state.terminated:
-      return (
-        JobStatus.SUCCEEDED
-        if container_status.last_state.terminated.exit_code == 0
-        else JobStatus.FAILED
-      )
+      if container_status.last_state.terminated.exit_code == 0:
+        return _succeeded_or_worker_failed()
+      return JobStatus.FAILED
   if pod.status.phase == "Running":
     return JobStatus.RUNNING
   return JobStatus.PENDING

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):
@@ -640,10 +716,13 @@ def setUp(self):
     self.mock_custom_api = self.enterContext(
       mock.patch(f"{_MODULE}._custom_api")
     ).return_value
+    # Default, no worker pods visible to the failure scan.
+    self.mock_core.list_namespaced_pod.return_value.items = []
 
-  def _make_pod(self, phase, exit_code=None):
+  def _make_pod(self, phase, exit_code=None, name="keras-pathways-job-1-0"):
     pod = MagicMock()
     pod.status.phase = phase
+    pod.metadata.name = name
     if exit_code is None:
       pod.status.container_statuses = None
     else:
@@ -653,6 +732,9 @@ def _make_pod(self, phase, exit_code=None):
       pod.status.container_statuses = [container_status]
     return pod
 
+  def _set_worker_pods(self, *pods):
+    self.mock_core.list_namespaced_pod.return_value.items = list(pods)
+
   def test_get_job_status_running(self):
     self.mock_core.read_namespaced_pod.return_value = self._make_pod("Running")
 
@@ -732,6 +814,7 @@ def test_get_job_status_failed_from_last_state(self):
   def test_get_job_status_succeeded_from_last_state(self):
     pod = MagicMock()
     pod.status.phase = "Running"
+    pod.metadata.name = "keras-pathways-job-1-0"
     container_status = MagicMock()
     container_status.state.terminated = None
     container_status.last_state.terminated = MagicMock(exit_code=0)
@@ -742,6 +825,42 @@ def test_get_job_status_succeeded_from_last_state(self):
 
     self.assertEqual(status, JobStatus.SUCCEEDED)
 
+  def test_get_job_status_failed_when_worker_failed_phase_succeeded(self):
+    """Regression for #239 in async path, leader Succeeded but worker failed."""
+    self.mock_core.read_namespaced_pod.return_value = self._make_pod(
+      "Succeeded"
+    )
+    leader = self._make_pod("Succeeded", name="keras-pathways-job-1-0")
+    worker = self._make_pod("Failed", name="keras-pathways-job-1-1")
+    self._set_worker_pods(leader, worker)
+
+    self.assertEqual(get_job_status("keras-pathways-job-1"), JobStatus.FAILED)
+
+  def test_get_job_status_failed_when_worker_failed_container_exit_zero(self):
+    """Container exit 0 must not mask a failed worker pod for async callers."""
+    self.mock_core.read_namespaced_pod.return_value = self._make_pod(
+      "Running", exit_code=0
+    )
+    worker = self._make_pod("Failed", name="keras-pathways-job-1-2")
+    self._set_worker_pods(worker)
+
+    self.assertEqual(get_job_status("keras-pathways-job-1"), JobStatus.FAILED)
+
+  def test_get_job_status_failed_when_worker_failed_last_state_exit_zero(self):
+    """Restarted leader exit 0 must not mask a failed worker pod."""
+    pod = MagicMock()
+    pod.status.phase = "Running"
+    pod.metadata.name = "keras-pathways-job-1-0"
+    container_status = MagicMock()
+    container_status.state.terminated = None
+    container_status.last_state.terminated = MagicMock(exit_code=0)
+    pod.status.container_statuses = [container_status]
+    self.mock_core.read_namespaced_pod.return_value = pod
+    worker = self._make_pod("Failed", name="keras-pathways-job-1-1")
+    self._set_worker_pods(worker)
+
+    self.assertEqual(get_job_status("keras-pathways-job-1"), JobStatus.FAILED)
+
   def test_get_job_status_api_error_raises(self):
     self.mock_core.read_namespaced_pod.side_effect = ApiException(
       status=500, reason="Internal Server Error"