fix: progress-based job drain at acquisition end + JobRunner death detection

software/control/core/job_processing.py

@@ -6,7 +6,7 @@
 import json
 from datetime import datetime
 from contextlib import contextmanager
-from typing import ClassVar, Dict, Generic, List, Optional, Set, Tuple, TypeVar, Union
+from typing import Callable, ClassVar, Dict, Generic, List, Optional, Sequence, Set, Tuple, TypeVar, Union
 from uuid import uuid4
 
 from dataclasses import dataclass, field
@@ -859,8 +859,18 @@ def output_queue(self) -> multiprocessing.Queue:
         return self._output_queue
 
     def has_pending(self):
-        with self._pending_count.get_lock():
-            return self._pending_count.value > 0
+        return self.pending_count() > 0
+
+    def pending_count(self) -> int:
+        """Number of dispatched jobs the subprocess has not completed yet.
+
+        Returns 0 after shutdown() has released the shared counter.
+        """
+        pending_count = self._pending_count
+        if pending_count is None:
+            return 0
+        with pending_count.get_lock():
+            return pending_count.value
 
     def wait_ready(self, timeout_s: float = 5.0) -> bool:
         """Wait for the subprocess to signal it's ready to process jobs.
@@ -1034,3 +1044,96 @@ def run(self):
         log_memory("WORKER_SHUTDOWN", include_children=False)
         stop_worker_monitoring()
         self._log.info("Shutdown request received, exiting run.")
+
+
+@dataclass
+class DrainResult:
+    """Outcome of drain_runners(): which pending jobs had to be abandoned, and why."""
+
+    # job class name -> number of pending jobs abandoned for that runner
+    abandoned: Dict[str, int] = field(default_factory=dict)
+    # job class names whose runner subprocess was found dead
+    dead: List[str] = field(default_factory=list)
+
+    @property
+    def total_abandoned(self) -> int:
+        return sum(self.abandoned.values())
+
+
+def find_dead_runners(
+    runners: Sequence[Tuple[type, Optional[JobRunner]]],
+) -> List[Tuple[type, JobRunner]]:
+    """Return runners that have pending jobs but whose subprocess is no longer alive.
+
+    Pending jobs on a dead runner can never complete: the pending/backpressure
+    counters are only decremented by the subprocess, so a dead runner presents as a
+    permanently-full job queue to the acquisition loop.
+    """
+    return [
+        (job_class, runner)
+        for job_class, runner in runners
+        if runner is not None and runner.pending_count() > 0 and not runner.is_alive()
+    ]
+
+
+def drain_runners(
+    runners: Sequence[Tuple[type, Optional[JobRunner]]],
+    stall_timeout_s: float = 10.0,
+    poll_fn: Optional[Callable[[], None]] = None,
+    poll_interval_s: float = 0.1,
+) -> DrainResult:
+    """Wait for all pending jobs on the given runners to complete.
+
+    The deadline is progress-based: every completed job resets the stall clock, so a
+    full-but-steadily-draining queue gets as long as it needs. Pending jobs are
+    abandoned only when nothing completes for stall_timeout_s (the stalled runner is
+    killed so later shutdown cannot hang on it), or immediately when a runner
+    subprocess has died.
+
+    Args:
+        runners: (job_class, runner) pairs; None runners are skipped.
+        stall_timeout_s: max time with zero completed jobs before giving up.
+        poll_fn: optional callback invoked every poll iteration (e.g. to drain
+            result queues while waiting).
+        poll_interval_s: sleep between polls.
+    """
+    result = DrainResult()
+    waiting = [(job_class, runner) for job_class, runner in runners if runner is not None]
+    stall_deadline = time.monotonic() + stall_timeout_s
+    last_total = None
+
+    while True:
+        if poll_fn is not None:
+            poll_fn()
+
+        still_waiting = []
+        total = 0
+        for job_class, runner in waiting:
+            pending = runner.pending_count()
+            if pending == 0:
+                continue
+            if not runner.is_alive():
+                # Nothing can complete these jobs anymore - abandon immediately.
+                result.abandoned[job_class.__name__] = pending
+                result.dead.append(job_class.__name__)
+                continue
+            still_waiting.append((job_class, runner))
+            total += pending
+        waiting = still_waiting
+
+        if total == 0:
+            return result
+
+        if last_total is None or total < last_total:
+            # Progress since the last look: reset the stall clock.
+            stall_deadline = time.monotonic() + stall_timeout_s
+            last_total = total
+        elif time.monotonic() > stall_deadline:
+            for job_class, runner in waiting:
+                pending = runner.pending_count()
+                if pending > 0:
+                    result.abandoned[job_class.__name__] = pending
+                    runner.kill()
+            return result
+
+        time.sleep(poll_interval_s)

software/control/core/multi_point_worker.py

@@ -45,6 +45,8 @@
     JobImage,
     JobRunner,
     JobResult,
+    drain_runners,
+    find_dead_runners,
 )
 from control.core.mosaic_utils import (
     calculate_overlap_pixels,
@@ -438,6 +440,24 @@ def _abort_due_to_error(self) -> None:
         self._abort_cause = "error"
         self.request_abort_fn()
 
+    def _abort_if_job_runners_dead(self) -> bool:
+        """Abort the acquisition if any job runner subprocess died with jobs pending.
+
+        A dead runner can never complete its jobs or release backpressure capacity:
+        the acquisition would otherwise crawl at one throttle timeout per frame while
+        losing every image dispatched to that runner. Returns True if aborted.
+        """
+        dead = find_dead_runners(self._job_runners)
+        if not dead:
+            return False
+        names = ", ".join(job_class.__name__ for job_class, _ in dead)
+        self._log.error(
+            f"Job runner subprocess died with jobs pending ({names}); those jobs can never "
+            f"complete. Aborting acquisition."
+        )
+        self._abort_due_to_error()
+        return True
+
     def _run_state_beat(self) -> None:
         self._run_state.beat(
             {
@@ -652,7 +672,7 @@ def _wait_for_outstanding_callback_images(self):
         self._ready_for_next_trigger.set()
         self._image_callback_idle.set()
 
-    def _finish_jobs(self, timeout_s=10):
+    def _finish_jobs(self, stall_timeout_s=10):
         # Drain and summarize all currently available job results before waiting for completion
         self._summarize_runner_outputs(drain_all=True)
 
@@ -661,34 +681,28 @@ def _finish_jobs(self, timeout_s=10):
         ]
 
         self._log.info(f"Waiting for jobs to finish on {len(active_runners)} job runners before shutting them down...")
-        timeout_time = time.time() + timeout_s
-
-        def timed_out():
-            return time.time() > timeout_time
-
-        def time_left():
-            return max(timeout_time - time.time(), 0)
 
-        # Wait for all pending jobs across all runners (round-robin to avoid blocking on one)
-        while not timed_out():
-            any_pending = False
-            for job_class, job_runner in active_runners:
-                if job_runner.has_pending():
-                    any_pending = True
-                    break
-            if not any_pending:
-                break
-            # Process any available results while waiting
-            self._summarize_runner_outputs(drain_all=True)
-            time.sleep(0.1)
-        else:
-            # Timed out - kill any runners that still have pending jobs
-            for job_class, job_runner in active_runners:
-                if job_runner.has_pending():
-                    self._log.error(
-                        f"Timed out after {timeout_s} [s] waiting for jobs to finish. Pending jobs for {job_class.__name__} abandoned!!!"
-                    )
-                    job_runner.kill()
+        # Progress-based drain: a full-but-steadily-draining queue gets as long as it
+        # needs (under backpressure saturation the pending count equals the job limit
+        # by construction, so any fixed deadline would abandon the acquisition tail).
+        # Jobs are abandoned only when nothing completes for stall_timeout_s or the
+        # runner subprocess died.
+        drain_result = drain_runners(
+            active_runners,
+            stall_timeout_s=stall_timeout_s,
+            poll_fn=lambda: self._summarize_runner_outputs(drain_all=True),
+        )
+        for job_class_name, abandoned_count in drain_result.abandoned.items():
+            cause = (
+                "its runner subprocess died"
+                if job_class_name in drain_result.dead
+                else f"no jobs completed for {stall_timeout_s} [s]"
+            )
+            self._log.error(
+                f"Abandoned {abandoned_count} pending {job_class_name} job(s) because {cause}. "
+                f"Data for these jobs is lost!"
+            )
+        self._acquisition_error_count += drain_result.total_abandoned
 
         # Drain results before shutdown
         self._summarize_runner_outputs(drain_all=True)
@@ -708,9 +722,11 @@ def shutdown_runner(job_runner, timeout):
                 log.error(f"Error shutting down job runner in background: {e}")
 
         self._log.info("Shutting down job runners (non-blocking)...")
-        remaining_time = time_left()
+        # Runners are idle after a clean drain (and already killed if they stalled or
+        # died), so a short join timeout before terminate() suffices.
+        shutdown_timeout_s = 2.0
         for job_class, job_runner in active_runners:
-            t = threading.Thread(target=shutdown_runner, args=(job_runner, remaining_time), daemon=True)
+            t = threading.Thread(target=shutdown_runner, args=(job_runner, shutdown_timeout_s), daemon=True)
             t.start()
 
         # Final drain of all output queues (should be empty, but check anyway)
@@ -1471,12 +1487,18 @@ def acquire_camera_image(
         # Backpressure check AFTER previous frame dispatched, BEFORE next trigger
         # This is when we know the previous image's jobs have been dispatched (and counters incremented)
         if self._backpressure.should_throttle():
+            # A dead runner subprocess can never release capacity - detect it up front
+            # instead of burning the full throttle timeout on every frame.
+            if self._abort_if_job_runners_dead():
+                return
             with self._timing.get_timer("backpressure.wait_for_capacity"):
                 got_capacity = self._backpressure.wait_for_capacity()
-                if not got_capacity:
-                    self._log.error(
-                        f"Backpressure timeout - disk I/O cannot keep up. Stats: {self._backpressure.get_stats()}"
-                    )
+            if not got_capacity:
+                if self._abort_if_job_runners_dead():
+                    return
+                self._log.error(
+                    f"Backpressure timeout - disk I/O cannot keep up. Stats: {self._backpressure.get_stats()}"
+                )
 
         with self._timing.get_timer("get_ready_for_trigger re-check"):
             # This should be a noop - we have the frame already.  Still, check!