Add timeout check for smartstacks

banzai/dbs.py

@@ -664,6 +664,19 @@ def get_subframes(db_address, moluid):
     with get_session(db_address, site_deploy=True) as session:
         return session.query(Subframe).filter(
             Subframe.moluid == moluid
+        ).order_by(Subframe.stack_num).all()
+
+
+def get_active_subframes_for_camera(db_address, camera):
+    """Get active subframe records for a camera ordered by stack and arrival time."""
+    with get_session(db_address, site_deploy=True) as session:
+        return session.query(Subframe).filter(
+            Subframe.camera == camera,
+            Subframe.status == 'active',
+        ).order_by(
+            Subframe.moluid,
+            Subframe.created_at,
+            Subframe.stack_num,
         ).all()
 
 

banzai/stacking.py

@@ -1,8 +1,10 @@
 """Smart stacking: worker, supervisor, and helper functions."""
 import argparse
+import datetime
 import multiprocessing
 import signal
 import time
+from collections import defaultdict
 
 import redis as redis_lib
 
@@ -16,6 +18,8 @@
 # Helpers
 # ---------------------------------------------------------------------------
 
+# Time that may elapse after the latest reduced subframe before a stack times out.
+STACK_TIMEOUT_SECONDS = 15 * 60
 REQUIRED_MESSAGE_FIELDS = ('fits_file', 'last_frame', 'instrument_enqueue_timestamp')
 
 
@@ -36,6 +40,22 @@ def check_stack_complete(subframes, frmtotal):
     return bool(subframes) and (all_arrived or has_last)
 
 
+def stack_has_timed_out(subframes, now=None):
+    """Return True when no newer reduced subframe has arrived before the timeout."""
+    latest_reduced_time = max(
+        (
+            subframe.created_at
+            for subframe in subframes
+            if subframe.created_at is not None
+        ),
+        default=None,
+    )
+    if latest_reduced_time is None:
+        return False
+
+    now = now or datetime.datetime.utcnow()
+    return (now - latest_reduced_time).total_seconds() > STACK_TIMEOUT_SECONDS
+
 # ---------------------------------------------------------------------------
 # Notifications
 # ---------------------------------------------------------------------------
@@ -91,6 +111,7 @@ def run_worker_loop(camera, db_address, redis_url, retention_days=30, poll_inter
     while True:
         try:
             process_notifications(db_address, redis_client, camera)
+            check_timeouts(db_address, camera)
             dbs.cleanup_old_subframes(db_address, retention_days)
             time.sleep(poll_interval)
         except Exception as e:
@@ -122,6 +143,27 @@ def finalize_stack(db_address, moluid, status='complete'):
     logger.debug(f'Mock ingester upload for {moluid}', extra_tags={'moluid': moluid})
 
 
+def check_timeouts(db_address, camera, now=None):
+    """Finalize active stacks that are complete or have exceeded the fixed timeout."""
+    active_subframes = dbs.get_active_subframes_for_camera(db_address, camera)
+    # The DB scan is camera-wide, so regroup rows into independent stack states
+    # before making terminal-status decisions.
+    subframes_by_moluid = defaultdict(list)
+    for subframe in active_subframes:
+        subframes_by_moluid[subframe.moluid].append(subframe)
+
+    for moluid, subframes in subframes_by_moluid.items():
+        frmtotal = subframes[0].frmtotal
+        # Completion wins over timeout so a fully reduced stack cannot be marked
+        # partial just because its cadence also exceeded the fixed threshold.
+        if check_stack_complete(subframes, frmtotal):
+            finalize_stack(db_address, moluid, status='complete')
+        # Incomplete active stacks are finalized only when reduced-row timing
+        # indicates that the next expected frame is no longer arriving normally.
+        elif stack_has_timed_out(subframes, now=now):
+            finalize_stack(db_address, moluid, status='timeout')
+
+
 # ---------------------------------------------------------------------------
 # Supervisor
 # ---------------------------------------------------------------------------

banzai/tests/site_e2e/test_site_e2e.py

@@ -5,6 +5,7 @@
 import subprocess
 import time
 import json
+import datetime
 from pathlib import Path
 
 import pytest
@@ -349,23 +350,80 @@ def check():
                     dbs.Subframe.moluid == 'mol-e2e-test'
                 ).all()
                 if subframes and all(s.status == 'complete' for s in subframes):
-                    return [s.filepath for s in subframes]
+                    return subframes
                 return None
 
-        filepaths = poll_until(check, timeout=300)
-        assert filepaths, "Subframe stack did not complete within timeout"
+        subframes = poll_until(check, timeout=300)
+        assert subframes, "Subframe stack did not complete within timeout"
 
         # Verify the reduced output file exists on disk.
         # Paths in the DB are now absolute host paths, so check directly.
-        assert filepaths[0], "Subframe has no filepath after completion"
-        expected_path = Path(filepaths[0])
+        assert subframes[0].filepath, "Subframe has no filepath after completion"
+        expected_path = Path(subframes[0].filepath)
 
         found = poll_until(
             lambda p=expected_path: p.exists() and p.stat().st_size > 0,
             timeout=60, interval=5
         )
         assert found, f"Reduced subframe output not found: {expected_path}"
 
+    @pytest.mark.e2e_site_startup
+    def test_13_subframe_stack_times_out(self, site_deployment):
+        """Verify the deployed stacking supervisor can mark an incomplete stack timeout."""
+
+        raw_dir = RAW_DIR
+        src_path = raw_dir / RAW_FRAME_FILENAME
+        subframe_path = raw_dir / 'subframe_timeout_test.fits.fz'
+        moluid = 'mol-e2e-timeout-test'
+
+        assert src_path.exists(), f"Raw frame not found: {src_path}"
+        shutil.copy2(str(src_path), str(subframe_path))
+
+        with fits.open(str(subframe_path), mode='update') as hdul:
+            hdul['SCI'].header['MOLUID'] = moluid
+            hdul['SCI'].header['MOLFRNUM'] = 1
+            hdul['SCI'].header['FRMTOTAL'] = 2
+            hdul['SCI'].header['STACK'] = 'T'
+
+        body = json.dumps({
+            'fits_file': str(subframe_path),
+            'last_frame': False,
+            'instrument_enqueue_timestamp': int(time.time() * 1000),
+        })
+        stack_queue = os.environ.get('STACK_QUEUE_NAME', 'banzai_stack_queue')
+        publish_raw_string_to_queue(stack_queue, body)
+
+        def row_exists():
+            with dbs.get_session(LOCAL_DB_ADDRESS, site_deploy=True) as session:
+                return session.query(dbs.Subframe).filter(
+                    dbs.Subframe.moluid == moluid
+                ).first()
+
+        subframe = poll_until(row_exists, timeout=120, interval=2)
+        assert subframe, \
+            "Timed-out subframe row was not inserted"
+        assert subframe.filepath, \
+            "Timed-out subframe row was inserted without a reduced filepath"
+
+        stale_created_at = datetime.datetime.utcnow() - datetime.timedelta(hours=2)
+        with dbs.get_session(LOCAL_DB_ADDRESS, site_deploy=True) as session:
+            session.execute(
+                text("UPDATE subframes SET created_at = :created_at WHERE moluid = :moluid"),
+                {'created_at': stale_created_at, 'moluid': moluid},
+            )
+
+        def timed_out():
+            with dbs.get_session(LOCAL_DB_ADDRESS, site_deploy=True) as session:
+                subframes = session.query(dbs.Subframe).filter(
+                    dbs.Subframe.moluid == moluid
+                ).all()
+                if subframes and all(s.status == 'timeout' for s in subframes):
+                    return subframes
+                return None
+
+        assert poll_until(timed_out, timeout=120, interval=2), \
+            "Incomplete subframe stack was not marked timeout"
+
     @pytest.mark.e2e_site_startup
     def test_14_cache_init_reuses_existing_slot(self, site_deployment):
         """Verify cache-init succeeds when the replication slot already exists on the publisher."""

banzai/tests/test_smart_stacking.py

@@ -13,8 +13,9 @@
 from banzai.dbs import insert_subframe, get_subframes, mark_stack_complete, cleanup_old_subframes
 from banzai.stacking import (validate_message, check_stack_complete,
                               push_notification, pop_notification, REDIS_KEY_PREFIX,
-                              process_notifications, finalize_stack, run_worker_loop,
-                              StackingSupervisor)
+                              process_notifications, finalize_stack, check_timeouts,
+                              stack_has_timed_out, STACK_TIMEOUT_SECONDS,
+                              run_worker_loop, StackingSupervisor)
 from banzai.scheduling import process_subframe
 from banzai.main import SubframeListener
 
@@ -96,6 +97,18 @@ def test_insert_subframe_and_get_subframes(self, db_address):
         assert subframe.is_last is False
         assert subframe.dateobs == dateobs
 
+    def test_get_subframes_orders_by_stack_num(self, db_address):
+        dateobs = datetime.datetime(2024, 6, 15, 12, 0, 0)
+        for stack_num in (3, 1, 2):
+            insert_subframe(
+                db_address, moluid='mol-ordered', stack_num=stack_num, frmtotal=3,
+                camera='cam1', filepath=f'/data/frame{stack_num}.fits', is_last=False, dateobs=dateobs,
+            )
+
+        subframes = get_subframes(db_address, moluid='mol-ordered')
+
+        assert [subframe.stack_num for subframe in subframes] == [1, 2, 3]
+
     def test_insert_subframe_duplicate_resets_state_for_retry(self, db_address):
         dateobs = datetime.datetime(2024, 6, 15, 12, 0, 0)
         insert_subframe(
@@ -163,6 +176,128 @@ def test_mark_stack_complete_sets_timeout(self, db_address):
             assert s.completed_at is not None
 
 
+# ---------------------------------------------------------------------------
+# Fixed reduced-subframe timeout
+# ---------------------------------------------------------------------------
+
+class TestFixedReducedSubframeTimeout:
+
+    @staticmethod
+    def _created(base, seconds):
+        return base + datetime.timedelta(seconds=seconds)
+
+    @staticmethod
+    def _insert_frame(db_address, moluid, stack_num, created_at, camera='cam1',
+                      filepath='/data/reduced.fits', frmtotal=5, is_last=False):
+        dateobs = datetime.datetime(2024, 6, 15, 12, 0, 0)
+        insert_subframe(
+            db_address, moluid=moluid, stack_num=stack_num, frmtotal=frmtotal,
+            camera=camera, filepath=filepath, is_last=is_last, dateobs=dateobs,
+        )
+        with dbs.get_session(db_address, site_deploy=True) as session:
+            session.execute(
+                text("UPDATE subframes SET created_at = :created_at WHERE moluid = :moluid AND stack_num = :stack_num"),
+                {'created_at': created_at, 'moluid': moluid, 'stack_num': stack_num},
+            )
+
+    def test_one_reduced_row_times_out_only_after_fixed_threshold(self, db_address):
+        start = datetime.datetime(2024, 1, 1, 0, 0, 0)
+        self._insert_frame(db_address, 'mol-one', 1, start)
+
+        check_timeouts(db_address, 'cam1', now=self._created(start, STACK_TIMEOUT_SECONDS))
+
+        assert {s.status for s in get_subframes(db_address, 'mol-one')} == {'active'}
+
+        check_timeouts(db_address, 'cam1', now=self._created(start, STACK_TIMEOUT_SECONDS + 1))
+
+        assert {s.status for s in get_subframes(db_address, 'mol-one')} == {'timeout'}
+
+    def test_newer_reduced_row_resets_timeout(self, db_address):
+        start = datetime.datetime(2024, 1, 1, 0, 0, 0)
+        second_arrival = self._created(start, STACK_TIMEOUT_SECONDS + 1)
+        self._insert_frame(db_address, 'mol-reset', 1, start)
+        self._insert_frame(db_address, 'mol-reset', 2, second_arrival)
+
+        check_timeouts(
+            db_address,
+            'cam1',
+            now=self._created(second_arrival, STACK_TIMEOUT_SECONDS),
+        )
+
+        subframes = get_subframes(db_address, 'mol-reset')
+        assert stack_has_timed_out(subframes, now=self._created(second_arrival, STACK_TIMEOUT_SECONDS)) is False
+        assert {s.status for s in subframes} == {'active'}
+
+    def test_latest_reduced_row_over_threshold_times_out(self, db_address):
+        start = datetime.datetime(2024, 1, 1, 0, 0, 0)
+        second_arrival = self._created(start, 60)
+        self._insert_frame(db_address, 'mol-latest-timeout', 1, start)
+        self._insert_frame(db_address, 'mol-latest-timeout', 2, second_arrival)
+
+        check_timeouts(db_address, 'cam1', now=self._created(second_arrival, STACK_TIMEOUT_SECONDS + 1))
+
+        assert {s.status for s in get_subframes(db_address, 'mol-latest-timeout')} == {'timeout'}
+
+    def test_later_newer_reduced_row_resets_timeout(self, db_address):
+        start = datetime.datetime(2024, 1, 1, 0, 0, 0)
+        third_arrival = self._created(start, 60 + STACK_TIMEOUT_SECONDS + 1)
+        self._insert_frame(db_address, 'mol-late-reset', 1, start)
+        self._insert_frame(db_address, 'mol-late-reset', 2, self._created(start, 60))
+        self._insert_frame(
+            db_address,
+            'mol-late-reset',
+            3,
+            third_arrival,
+        )
+
+        subframes = get_subframes(db_address, 'mol-late-reset')
+        assert stack_has_timed_out(subframes, now=self._created(third_arrival, 1)) is False
+
+    def test_incomplete_stack_inside_timeout_window_remains_active(self, db_address):
+        start = datetime.datetime(2024, 1, 1, 0, 0, 0)
+        self._insert_frame(db_address, 'mol-active', 1, start)
+        self._insert_frame(db_address, 'mol-active', 2, self._created(start, 120))
+
+        check_timeouts(db_address, 'cam1', now=self._created(start, 200))
+
+        subframes = get_subframes(db_address, 'mol-active')
+        assert {s.status for s in subframes} == {'active'}
+
+    def test_check_timeouts_marks_timeout_for_worker_camera_only(self, db_address):
+        start = datetime.datetime(2024, 1, 1, 0, 0, 0)
+        self._insert_frame(db_address, 'mol-timeout', 1, start, camera='cam1')
+        self._insert_frame(db_address, 'mol-other-camera', 1, start, camera='cam2')
+
+        check_timeouts(db_address, 'cam1', now=self._created(start, STACK_TIMEOUT_SECONDS + 1))
+
+        assert {s.status for s in get_subframes(db_address, 'mol-timeout')} == {'timeout'}
+        assert {s.status for s in get_subframes(db_address, 'mol-other-camera')} == {'active'}
+
+    def test_check_timeouts_marks_complete_before_timeout(self, db_address):
+        start = datetime.datetime(2024, 1, 1, 0, 0, 0)
+        for stack_num in range(1, 4):
+            self._insert_frame(
+                db_address,
+                'mol-complete-wins',
+                stack_num,
+                self._created(start, (stack_num - 1) * (STACK_TIMEOUT_SECONDS + 1)),
+                frmtotal=3, is_last=(stack_num == 3),
+            )
+
+        check_timeouts(db_address, 'cam1', now=self._created(start, 3 * (STACK_TIMEOUT_SECONDS + 1)))
+
+        assert {s.status for s in get_subframes(db_address, 'mol-complete-wins')} == {'complete'}
+
+    def test_check_timeouts_ignores_terminal_rows(self, db_address):
+        start = datetime.datetime(2024, 1, 1, 0, 0, 0)
+        self._insert_frame(db_address, 'mol-terminal', 1, start)
+        mark_stack_complete(db_address, 'mol-terminal', 'complete')
+
+        check_timeouts(db_address, 'cam1', now=self._created(start, STACK_TIMEOUT_SECONDS + 1))
+
+        assert {s.status for s in get_subframes(db_address, 'mol-terminal')} == {'complete'}
+
+
 # ---------------------------------------------------------------------------
 # Redis notifications
 # ---------------------------------------------------------------------------
@@ -656,14 +791,37 @@ class TestWorkerLoopResilience:
 
     @patch('banzai.stacking.redis_lib.Redis.from_url')
     @patch('banzai.stacking.time.sleep')
+    @patch('banzai.stacking.check_timeouts')
+    @patch('banzai.stacking.dbs.cleanup_old_subframes')
     @patch('banzai.stacking.process_notifications')
-    def test_run_worker_loop_continues_after_exception(self, mock_process, mock_sleep, mock_redis_cls):
+    def test_run_worker_loop_continues_after_exception(
+        self, mock_process, mock_cleanup, mock_check_timeouts, mock_sleep, mock_redis_cls
+    ):
         """run_worker_loop must not crash when process_notifications raises; it should log and continue."""
         # First call raises, second call raises KeyboardInterrupt to escape the infinite loop.
         mock_process.side_effect = [Exception('boom'), KeyboardInterrupt]
         with pytest.raises(KeyboardInterrupt):
             run_worker_loop('cam1', 'sqlite:///fake.db', 'redis://localhost:6379', poll_interval=0)
         # process_notifications was called twice: once raised Exception, once raised KeyboardInterrupt.
         assert mock_process.call_count == 2
+        mock_check_timeouts.assert_not_called()
+        mock_cleanup.assert_not_called()
         # sleep should have been called once (after the transient Exception, before continuing).
         mock_sleep.assert_called_once_with(0)
+
+    @patch('banzai.stacking.redis_lib.Redis.from_url')
+    @patch('banzai.stacking.time.sleep')
+    @patch('banzai.stacking.check_timeouts')
+    @patch('banzai.stacking.dbs.cleanup_old_subframes')
+    @patch('banzai.stacking.process_notifications')
+    def test_run_worker_loop_invokes_timeout_sweep(
+        self, mock_process, mock_cleanup, mock_check_timeouts, mock_sleep, mock_redis_cls
+    ):
+        """run_worker_loop should sweep fixed timeouts after notification processing."""
+        mock_process.side_effect = [None, KeyboardInterrupt]
+        with pytest.raises(KeyboardInterrupt):
+            run_worker_loop('cam1', 'sqlite:///fake.db', 'redis://localhost:6379', poll_interval=0)
+
+        mock_check_timeouts.assert_called_once_with('sqlite:///fake.db', 'cam1')
+        mock_cleanup.assert_called_once_with('sqlite:///fake.db', 30)
+        mock_sleep.assert_called_once_with(0)