diff --git a/mloda/core/runtime/compute_framework_executor.py b/mloda/core/runtime/compute_framework_executor.py index 029b857d..e0f7e1a1 100644 --- a/mloda/core/runtime/compute_framework_executor.py +++ b/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())) diff --git a/mloda/core/runtime/run.py b/mloda/core/runtime/run.py index 64de06b9..3644413e 100644 --- a/mloda/core/runtime/run.py +++ b/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 diff --git a/mloda/core/runtime/worker_manager.py b/mloda/core/runtime/worker_manager.py index 892605b6..1c109c8a 100644 --- a/mloda/core/runtime/worker_manager.py +++ b/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 diff --git a/tests/test_core/test_runtime/test_compute_framework_executor.py b/tests/test_core/test_runtime/test_compute_framework_executor.py index 795e9dcf..fe431c1a 100644 --- a/tests/test_core/test_runtime/test_compute_framework_executor.py +++ b/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] diff --git a/tests/test_core/test_runtime/test_worker_manager.py b/tests/test_core/test_runtime/test_worker_manager.py index 5ebefe18..81df5a6f 100644 --- a/tests/test_core/test_runtime/test_worker_manager.py +++ b/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."""