diff --git a/software/control/core/job_processing.py b/software/control/core/job_processing.py index 189b1796f..bdbb92791 100644 --- a/software/control/core/job_processing.py +++ b/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) diff --git a/software/control/core/multi_point_worker.py b/software/control/core/multi_point_worker.py index 7fd973452..ec748e411 100644 --- a/software/control/core/multi_point_worker.py +++ b/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! diff --git a/software/tests/control/core/test_job_drain.py b/software/tests/control/core/test_job_drain.py new file mode 100644 index 000000000..94e723785 --- /dev/null +++ b/software/tests/control/core/test_job_drain.py @@ -0,0 +1,256 @@ +"""Tests for progress-based job draining and JobRunner subprocess death detection. + +These verify the fix for two failure modes found in the 2026-07-12 backpressure +investigation: + +1. ``_finish_jobs`` used a fixed timeout to wait for pending jobs at the end of an + acquisition (natural completion AND abort). Under saturation the pending count + equals the backpressure job limit by construction, so whenever + ``pending × per-job save time`` exceeded the fixed timeout the tail of the + acquisition was killed and its images silently lost — even though jobs were + completing steadily the whole time. The drain deadline must be progress-based: + abandon only when NO job completes for the stall timeout. + +2. A dead JobRunner subprocess was undetectable. Backpressure counters are only + decremented by the subprocess, so its death froze them at the limit and the + acquisition crawled at one throttle-timeout per FOV, presenting as "stuck with + jobs queue full". Dead runners with pending jobs must be detectable so the + worker can abort with a clear error instead. +""" + +import time +from dataclasses import dataclass + +import numpy as np +import pytest + +import squid.abc +from control.core.job_processing import ( + CaptureInfo, + Job, + JobImage, + JobRunner, + drain_runners, + find_dead_runners, +) +from control.models import AcquisitionChannel, CameraSettings, IlluminationSettings + + +def make_test_capture_info() -> CaptureInfo: + """Create a minimal CaptureInfo for testing.""" + return CaptureInfo( + position=squid.abc.Pos(x_mm=0.0, y_mm=0.0, z_mm=0.0, theta_rad=None), + z_index=0, + capture_time=time.time(), + configuration=AcquisitionChannel( + name="BF LED matrix full", + display_color="#FFFFFF", + camera=1, # v1.0: camera is int ID + illumination_settings=IlluminationSettings( + illumination_channel="BF LED matrix full", + intensity=50.0, + ), + camera_settings=CameraSettings( + exposure_time_ms=10.0, + gain_mode=1.0, + ), + z_offset_um=0.0, # v1.0: at channel level + ), + save_directory="/tmp/test", + file_id="test_0_0", + region_id="A1", + fov=0, + configuration_idx=0, + ) + + +def make_test_job_image() -> JobImage: + """Create a minimal JobImage for testing.""" + return JobImage(image_array=np.zeros((10, 10), dtype=np.uint16)) + + +@dataclass +class SlowJob(Job): + """A job that takes a configurable amount of time to run.""" + + duration_s: float = 0.1 + + def run(self): + time.sleep(self.duration_s) + return "done" + + +@dataclass +class HangingJob(Job): + """A job that runs much longer than any test timeout (simulates a stuck save).""" + + duration_s: float = 60.0 + + def run(self): + time.sleep(self.duration_s) + return "done" + + +def make_job(job_type, duration_s): + return job_type( + capture_info=make_test_capture_info(), + capture_image=make_test_job_image(), + duration_s=duration_s, + ) + + +def start_runner() -> JobRunner: + runner = JobRunner() + runner.daemon = True + runner.start() + assert runner.wait_ready(timeout_s=10.0), "JobRunner subprocess never became ready" + return runner + + +def wait_until(predicate, timeout_s, interval_s=0.05): + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(interval_s) + return predicate() + + +class TestPendingCount: + def test_pending_count_tracks_dispatch_and_completion(self): + """pending_count() reports dispatched-but-incomplete jobs, reaching 0 when drained.""" + runner = start_runner() + try: + for _ in range(3): + runner.dispatch(make_job(SlowJob, duration_s=0.2)) + assert runner.pending_count() == 3 + + assert wait_until(lambda: runner.pending_count() == 0, timeout_s=15.0) + finally: + runner.shutdown(timeout_s=1.0) + + def test_pending_count_zero_after_shutdown(self): + """pending_count() is safe to call after shutdown() clears the shared value.""" + runner = start_runner() + runner.shutdown(timeout_s=2.0) + assert runner.pending_count() == 0 + + +class TestDrainRunners: + def test_drains_fully_when_progress_is_slower_than_stall_timeout(self): + """THE regression case: total drain time exceeds the stall timeout, but jobs + complete steadily — nothing may be abandoned. + + With the old fixed-deadline behavior (timeout_s=1), the last jobs of this + queue would have been killed and lost. + """ + runner = start_runner() + try: + for _ in range(6): + runner.dispatch(make_job(SlowJob, duration_s=0.4)) + + start = time.monotonic() + result = drain_runners([(SlowJob, runner)], stall_timeout_s=1.0) + elapsed = time.monotonic() - start + + assert result.total_abandoned == 0 + assert result.abandoned == {} + assert result.dead == [] + assert runner.pending_count() == 0 + # It must have kept waiting well past the stall timeout (a fixed + # 1s deadline would have returned early and abandoned jobs). + assert elapsed > 1.0 + finally: + runner.shutdown(timeout_s=1.0) + + def test_abandons_pending_jobs_after_stall(self): + """A job making no progress for the stall timeout is abandoned and its runner killed.""" + runner = start_runner() + try: + runner.dispatch(make_job(HangingJob, duration_s=60.0)) + # Give the subprocess a moment to pick the job up. + time.sleep(0.3) + + start = time.monotonic() + result = drain_runners([(HangingJob, runner)], stall_timeout_s=0.7) + elapsed = time.monotonic() - start + + assert result.abandoned == {"HangingJob": 1} + assert result.total_abandoned == 1 + # Should return shortly after the stall timeout, not wait for the job. + assert elapsed < 10.0 + # The stalled runner is killed so shutdown cannot hang on it. + assert wait_until(lambda: not runner.is_alive(), timeout_s=5.0) + finally: + runner.shutdown(timeout_s=1.0) + + def test_abandons_dead_runner_immediately(self): + """A runner whose subprocess died is abandoned without waiting out the stall timeout.""" + runner = start_runner() + try: + for _ in range(3): + runner.dispatch(make_job(HangingJob, duration_s=60.0)) + time.sleep(0.3) + + runner.kill() + assert wait_until(lambda: not runner.is_alive(), timeout_s=5.0) + + start = time.monotonic() + result = drain_runners([(HangingJob, runner)], stall_timeout_s=30.0) + elapsed = time.monotonic() - start + + # Must not wait anywhere near the 30s stall timeout. + assert elapsed < 10.0 + assert result.abandoned == {"HangingJob": 3} + assert result.dead == ["HangingJob"] + finally: + runner.shutdown(timeout_s=1.0) + + def test_no_pending_jobs_returns_immediately(self): + """Draining runners with nothing pending is a no-op.""" + runner = start_runner() + try: + start = time.monotonic() + result = drain_runners([(SlowJob, runner)], stall_timeout_s=5.0) + elapsed = time.monotonic() - start + + assert result.total_abandoned == 0 + assert elapsed < 2.0 + finally: + runner.shutdown(timeout_s=1.0) + + +class TestFindDeadRunners: + def test_alive_runner_with_pending_jobs_is_not_reported(self): + runner = start_runner() + try: + runner.dispatch(make_job(SlowJob, duration_s=1.0)) + assert find_dead_runners([(SlowJob, runner)]) == [] + finally: + runner.shutdown(timeout_s=1.0) + + def test_dead_runner_with_pending_jobs_is_reported(self): + runner = start_runner() + try: + runner.dispatch(make_job(HangingJob, duration_s=60.0)) + time.sleep(0.2) + runner.kill() + assert wait_until(lambda: not runner.is_alive(), timeout_s=5.0) + + dead = find_dead_runners([(HangingJob, runner)]) + assert dead == [(HangingJob, runner)] + finally: + runner.shutdown(timeout_s=1.0) + + def test_dead_runner_without_pending_jobs_is_not_reported(self): + """A runner that exited with nothing pending is not an emergency.""" + runner = start_runner() + try: + runner.kill() + assert wait_until(lambda: not runner.is_alive(), timeout_s=5.0) + assert find_dead_runners([(SlowJob, runner)]) == [] + finally: + runner.shutdown(timeout_s=1.0) + + def test_none_runner_entries_are_skipped(self): + assert find_dead_runners([(SlowJob, None)]) == []