fix: detect a worker that exits while steps are still assigned to it

mloda/core/runtime/compute_framework_executor.py

@@ -306,3 +306,6 @@ def multi_execute_step(self, step: Any) -> None:
             process, command_queue, result_queue = existing
 
         self.worker_manager.send_command(cfw_uuid, step)
+        # Record the assignment so a worker that exits still owing these results is
+        # detectable, including the clean exit the data-drop path produces.
+        self.worker_manager.record_assignment(cfw_uuid, set(step.get_uuids()))

mloda/core/runtime/run.py

@@ -154,6 +154,15 @@ def _check_for_error(self) -> bool:
             dead = self.worker_manager.find_dead_workers()
             if dead:
                 raise MlodaRunError(f"Worker process(es) died unexpectedly: {dead}")
+            # A clean exit is invisible above, so check assignments too: a worker that is
+            # gone will never answer, and waiting on it hangs the run instead of failing it.
+            orphaned = self.worker_manager.find_orphaned_steps()
+            if orphaned:
+                detail = "; ".join(
+                    f"cfw {cfw_uuid} exited with code {exitcode} owing step(s) {', '.join(str(step) for step in steps)}"
+                    for cfw_uuid, exitcode, steps in orphaned
+                )
+                raise MlodaRunError(f"Worker process(es) exited with steps still assigned: {detail}")
             return False
         return True
 

mloda/core/runtime/worker_manager.py

@@ -23,6 +23,10 @@ def __init__(self) -> None:
         self.process_register: dict[UUID, tuple[Any, Any, Any]] = {}
         self.result_queues_collection: set[Any] = set()
         self.result_uuids_collection: set[UUID] = set()
+        # cfw_uuid -> step uuids dispatched to that worker. Needed because a worker that
+        # exits cleanly is invisible to find_dead_workers, so the only way to notice the
+        # loss is that steps were assigned to it and no result ever arrived.
+        self.assigned_steps: dict[UUID, set[UUID]] = {}
 
     def add_thread_task(self, task: threading.Thread) -> None:
         """Add task to list and call task.start()."""
@@ -72,6 +76,10 @@ def poll_result_queues(self) -> None:
             if isinstance(msg, str):
                 self.result_uuids_collection.add(UUID(msg))
 
+    def record_assignment(self, cfw_uuid: UUID, step_uuids: set[UUID]) -> None:
+        """Remember that these steps were dispatched to this worker."""
+        self.assigned_steps.setdefault(cfw_uuid, set()).update(step_uuids)
+
     def find_dead_workers(self) -> list[tuple[UUID, int]]:
         """Return (cfw_uuid, exitcode) for workers that died abnormally (exitcode not in {None, 0})."""
         dead: list[tuple[UUID, int]] = []
@@ -81,6 +89,29 @@ def find_dead_workers(self) -> list[tuple[UUID, int]]:
                 dead.append((cfw_uuid, exitcode))
         return dead
 
+    def find_orphaned_steps(self) -> list[tuple[UUID, int, list[UUID]]]:
+        """Return (cfw_uuid, exitcode, orphaned step uuids) per exited worker still owing results.
+
+        Complements ``find_dead_workers``, which only reports a non-zero exitcode. A worker
+        that takes the data-drop path breaks its own loop and exits with code 0, so it is
+        invisible there while the steps dispatched to it stay in ``currently_running_steps``
+        forever and the run loop waits on a process that is gone.
+
+        Any exitcode counts here, including 0: once a process has exited it will never
+        produce a result, so an assigned step with no result is lost whatever the code.
+        Results are checked against ``result_uuids_collection``, so a step whose result
+        arrived before the exit is not reported.
+        """
+        orphaned: list[tuple[UUID, int, list[UUID]]] = []
+        for cfw_uuid, (process, _, _) in self.process_register.items():
+            exitcode = process.exitcode
+            if exitcode is None:
+                continue
+            pending = self.assigned_steps.get(cfw_uuid, set()) - self.result_uuids_collection
+            if pending:
+                orphaned.append((cfw_uuid, exitcode, sorted(pending, key=str)))
+        return orphaned
+
     def is_step_done(self, step_uuid: UUID) -> bool:
         """Return step_uuid in result_uuids_collection."""
         return step_uuid in self.result_uuids_collection

tests/test_core/test_runtime/test_compute_framework_executor.py

@@ -843,6 +843,9 @@ def test_prepares_step_with_multiprocessing_mode(self) -> None:
         step.children_if_root = []
         step.compute_framework = Mock()
         step.compute_framework.get_class_name.return_value = "TestCFW"
+        # multi_execute_step records the dispatch, so the step must answer get_uuids()
+        # with real uuids like a FeatureGroupStep does.
+        step.get_uuids.return_value = {uuid4()}
 
         cfw_uuid = uuid4()
         cfw_register.get_cfw_uuid.return_value = cfw_uuid
@@ -871,6 +874,9 @@ def test_creates_worker_process_if_not_exists(self, mock_worker: Any) -> None:
         step.children_if_root = []
         step.compute_framework = Mock()
         step.compute_framework.get_class_name.return_value = "TestCFW"
+        # multi_execute_step records the dispatch, so the step must answer get_uuids()
+        # with real uuids like a FeatureGroupStep does.
+        step.get_uuids.return_value = {uuid4()}
 
         cfw_uuid = uuid4()
         cfw_register.get_cfw_uuid.return_value = cfw_uuid
@@ -904,6 +910,9 @@ def test_uses_existing_worker_process_if_exists(self) -> None:
         step.children_if_root = []
         step.compute_framework = Mock()
         step.compute_framework.get_class_name.return_value = "TestCFW"
+        # multi_execute_step records the dispatch, so the step must answer get_uuids()
+        # with real uuids like a FeatureGroupStep does.
+        step.get_uuids.return_value = {uuid4()}
 
         cfw_uuid = uuid4()
         cfw_register.get_cfw_uuid.return_value = cfw_uuid
@@ -932,6 +941,9 @@ def test_sends_command_to_worker(self) -> None:
         step.children_if_root = []
         step.compute_framework = Mock()
         step.compute_framework.get_class_name.return_value = "TestCFW"
+        # multi_execute_step records the dispatch, so the step must answer get_uuids()
+        # with real uuids like a FeatureGroupStep does.
+        step.get_uuids.return_value = {uuid4()}
 
         cfw_uuid = uuid4()
         cfw_register.get_cfw_uuid.return_value = cfw_uuid
@@ -956,6 +968,9 @@ def test_prepares_from_cfw_for_transform_framework_step(self) -> None:
         step.from_framework = Mock()
         step.from_framework.get_class_name.return_value = "FromCFW"
         step.required_uuids = [uuid4()]
+        # multi_execute_step records the dispatch, so the step must answer get_uuids()
+        # with real uuids like a TransformFrameworkStep does.
+        step.get_uuids.return_value = {uuid4()}
 
         from_cfw_uuid = uuid4()
         cfw_register.get_cfw_uuid.side_effect = [from_cfw_uuid, from_cfw_uuid]

tests/test_core/test_runtime/test_worker_manager.py

@@ -408,6 +408,106 @@ def test_find_dead_workers_ignores_alive_and_clean_workers(self) -> None:
         assert manager.find_dead_workers() == []
 
 
+class TestWorkerManagerOrphanedStepDetection:
+    """A worker that exits while steps are still assigned to it must be detectable.
+
+    find_dead_workers only reports a non-zero exitcode. The data-drop path breaks the
+    worker loop and the process exits with code 0, so it is invisible there while the steps
+    dispatched to it stay in currently_running_steps with no result ever arriving.
+    """
+
+    @staticmethod
+    def _exited(manager: WorkerManager, exitcode: int) -> UUID:
+        """Register a worker whose process has already exited with *exitcode*."""
+        cfw_uuid = uuid4()
+        process = MagicMock()
+        process.exitcode = exitcode
+        process.is_alive.return_value = False
+        manager.process_register[cfw_uuid] = (process, MagicMock(), MagicMock())
+        return cfw_uuid
+
+    def test_a_clean_exit_owing_a_step_is_reported(self) -> None:
+        """Exitcode 0 is the case find_dead_workers cannot see, so it is the case that matters."""
+        manager = WorkerManager()
+        cfw_uuid = self._exited(manager, 0)
+        step_uuid = uuid4()
+        manager.record_assignment(cfw_uuid, {step_uuid})
+
+        # Premise: the existing check stays silent, which is why this one has to exist.
+        assert manager.find_dead_workers() == []
+
+        assert manager.find_orphaned_steps() == [(cfw_uuid, 0, [step_uuid])]
+
+    def test_an_abnormal_exit_owing_a_step_is_reported_too(self) -> None:
+        """Any exitcode counts: an exited process will never answer, whatever the code."""
+        manager = WorkerManager()
+        cfw_uuid = self._exited(manager, -9)
+        step_uuid = uuid4()
+        manager.record_assignment(cfw_uuid, {step_uuid})
+
+        assert manager.find_orphaned_steps() == [(cfw_uuid, -9, [step_uuid])]
+
+    def test_a_step_whose_result_already_arrived_is_not_reported(self) -> None:
+        """The worker finished its work and then exited; nothing was lost."""
+        manager = WorkerManager()
+        cfw_uuid = self._exited(manager, 0)
+        step_uuid = uuid4()
+        manager.record_assignment(cfw_uuid, {step_uuid})
+        manager.result_uuids_collection.add(step_uuid)
+
+        assert manager.find_orphaned_steps() == []
+
+    def test_only_the_steps_still_owed_are_named(self) -> None:
+        """A partially-drained worker reports the remainder, not everything it was sent."""
+        manager = WorkerManager()
+        cfw_uuid = self._exited(manager, 0)
+        done, pending = uuid4(), uuid4()
+        manager.record_assignment(cfw_uuid, {done, pending})
+        manager.result_uuids_collection.add(done)
+
+        assert manager.find_orphaned_steps() == [(cfw_uuid, 0, [pending])]
+
+    def test_a_live_worker_is_never_orphaned(self) -> None:
+        """A running worker still owes results; that is work in progress, not a loss."""
+        manager = WorkerManager()
+        cfw_uuid = uuid4()
+        process = MagicMock()
+        process.exitcode = None
+        process.is_alive.return_value = True
+        manager.process_register[cfw_uuid] = (process, MagicMock(), MagicMock())
+        manager.record_assignment(cfw_uuid, {uuid4()})
+
+        assert manager.find_orphaned_steps() == []
+
+    def test_an_exited_worker_with_no_assignments_is_not_reported(self) -> None:
+        """A worker stopped after draining its queue exits owing nothing."""
+        manager = WorkerManager()
+        self._exited(manager, 0)
+
+        assert manager.find_orphaned_steps() == []
+
+    def test_record_assignment_accumulates_across_dispatches(self) -> None:
+        """multi_execute_step runs once per step, so assignments must union, not replace."""
+        manager = WorkerManager()
+        cfw_uuid = self._exited(manager, 0)
+        first, second = uuid4(), uuid4()
+        manager.record_assignment(cfw_uuid, {first})
+        manager.record_assignment(cfw_uuid, {second})
+
+        assert manager.assigned_steps[cfw_uuid] == {first, second}
+        assert manager.find_orphaned_steps() == [(cfw_uuid, 0, sorted([first, second], key=str))]
+
+    def test_each_exited_worker_is_reported_separately(self) -> None:
+        """The run loop names every affected cfw, not only the first one found."""
+        manager = WorkerManager()
+        first_cfw = self._exited(manager, 0)
+        second_cfw = self._exited(manager, 0)
+        manager.record_assignment(first_cfw, {uuid4()})
+        manager.record_assignment(second_cfw, {uuid4()})
+
+        assert {entry[0] for entry in manager.find_orphaned_steps()} == {first_cfw, second_cfw}
+
+
 class TestWorkerManagerDropCompletion:
     """Test waiting for drop completion messages."""