diff --git a/banzai/dbs.py b/banzai/dbs.py index b4c6dc69..2c969e0d 100755 --- a/banzai/dbs.py +++ b/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() diff --git a/banzai/stacking.py b/banzai/stacking.py index 28a5af41..12b44061 100644 --- a/banzai/stacking.py +++ b/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 # --------------------------------------------------------------------------- diff --git a/banzai/tests/site_e2e/test_site_e2e.py b/banzai/tests/site_e2e/test_site_e2e.py index 15326081..d5bf9492 100644 --- a/banzai/tests/site_e2e/test_site_e2e.py +++ b/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,16 +350,16 @@ 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, @@ -366,6 +367,63 @@ def check(): ) 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.""" diff --git a/banzai/tests/test_smart_stacking.py b/banzai/tests/test_smart_stacking.py index 938581bb..140aa35f 100644 --- a/banzai/tests/test_smart_stacking.py +++ b/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,8 +791,12 @@ 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] @@ -665,5 +804,24 @@ def test_run_worker_loop_continues_after_exception(self, mock_process, mock_slee 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) diff --git a/docs/smartstacking_architecture.md b/docs/smartstacking_architecture.md index f80635df..23ba553f 100644 --- a/docs/smartstacking_architecture.md +++ b/docs/smartstacking_architecture.md @@ -140,14 +140,22 @@ database level. 4. Checks the completion rule when at least one row is active. 5. Marks a complete group as `complete`. 6. Repeats from step 1 until the sorted set is empty. -7. Deletes old terminal rows according to the retention setting. -8. Sleeps for the polling interval. +7. Sweeps active rows for its camera to recover missed notifications and apply timeouts. +8. Deletes old terminal rows according to the retention setting. +9. Sleeps for the polling interval. + +After the notification set is empty, the worker scans active rows for its +camera. This fallback sweep lets complete stacks reach terminal state even if +a notification was missed. It also marks an incomplete stack as `timeout` when +no newer reduced row has arrived within 15 minutes of the latest row's +`created_at`. A newer reduced row resets that window. Completion is checked +before timeout, so complete stacks become `complete` rather than `timeout`. The worker does not build stack state from Redis messages. Ten notifications for the same pending `MOLUID` become one sorted-set member and one database query. Multiple `MOLUID`s may be active for one camera. None exclusively owns -the worker: an incomplete group waits for another notification while older -pending work for other groups continues in FIFO order. +the worker: an incomplete group waits for another notification or the fallback +sweep while older pending work for other groups continues in FIFO order. ```mermaid flowchart TD @@ -158,6 +166,9 @@ flowchart TD active{"Any active rows?"} complete{"Completion rule true?"} mark["Mark every row complete"] + sweep["Scan active rows
for this camera"] + eligible{"Complete or timed out?"} + finalize["Mark complete or timeout"] cleanup["Delete old terminal rows"] sleep["Sleep"] @@ -167,7 +178,9 @@ flowchart TD active -->|"no"| pop complete -->|"yes"| mark --> pop complete -->|"no"| pop - notified -->|"no"| cleanup + notified -->|"no"| sweep --> eligible + eligible -->|"yes"| finalize --> cleanup + eligible -->|"no"| cleanup cleanup --> sleep --> tick ``` @@ -208,7 +221,7 @@ one or more raw e00 subframes | `banzai-subframe-listener` | `banzai.main:run_subframe_worker` | Consume and validate RabbitMQ messages; dispatch Celery tasks. Despite the CLI name, this is the listener process. | | `banzai-subframe-worker` | Celery worker consuming `SUBFRAME_TASK_QUEUE_NAME` | Reduce the FITS file, upsert the reduced row, and send a notification. | | `banzai-stacking-supervisor` | `banzai.stacking:run_supervisor` | Discover cameras at startup, start one child per camera, and restart children that exit. | -| `stacking-worker-{camera}` | `banzai.stacking:run_worker_loop` | Process unique notifications in FIFO order, query stack state, mark complete groups, and run retention cleanup. | +| `stacking-worker-{camera}` | `banzai.stacking:run_worker_loop` | Process unique notifications in FIFO order, sweep active stacks for completion or timeout, and run retention cleanup. | Camera discovery uses `dbs.get_instruments_at_site(site_id, db_address)` once when the supervisor starts. Adding a camera to the database does not create a @@ -264,7 +277,7 @@ state. | `is_last` | Copy of the queue message's `last_frame` flag. | | `status` | Expected values are `active`, `complete`, and `timeout`. | | `dateobs` | Observation time from `DATE-OBS`, when available. | -| `created_at` | Time the row was inserted or reset by an upsert. | +| `created_at` | Time the reduced row was inserted or reset; the fixed timeout is measured from the newest value in a stack. | | `completed_at` | Time the group was marked terminal. | Current transitions are: @@ -273,16 +286,17 @@ Current transitions are: stateDiagram-v2 [*] --> active: reduced row upserted active --> complete: completion rule passes + active --> timeout: no newer reduced row for 15 minutes complete --> active: same subframe is reprocessed timeout --> active: same subframe is reprocessed ``` -The model and database helper support `timeout`, but no production worker path -currently sets it. - Retention cleanup deletes rows whose status is not `active` and whose `completed_at` is older than `STACK_RETENTION_DAYS` (30 days by default). -Active rows are retained indefinitely. +Active rows remain until the completion rule passes or the fixed timeout +expires. `STACK_TIMEOUT_MINUTES` in `site-banzai-env.default` is a legacy, +unused setting; the worker timeout is the 15-minute constant in +`banzai.stacking`. ## Current limitations and failure behavior @@ -290,16 +304,6 @@ These details are important when debugging or extending the subsystem: - **No `e45` stack product is generated.** Completion currently means a database transition only; only the individual `e09` inputs exist. -- **No timeout finalization runs.** `STACK_TIMEOUT_MINUTES` remains in - `site-banzai-env.default`, but the current Python code does not read it. -- **There is no fallback database sweep.** An active stack is reconsidered only - after a Redis notification for its `MOLUID`. If notification is disabled, - lost, or missed, the stack remains active until another notification arrives. -- **The in-flight notification has a crash window.** `ZPOPMIN` removes one - `MOLUID` before its database work begins. A worker crash can lose that one - notification, but all other members remain ordered in the sorted set. - PostgreSQL remains correct, but there is currently no sweep to rediscover an - active stack whose final notification was lost. - **Cleanup is repeated per camera.** Every camera worker calls the same database-wide terminal-row cleanup each tick. This is correct but redundant. - **Camera discovery is startup-only.** Restart the supervisor after adding or