diff --git a/.env.example b/.env.example index 12a099f..3ac89ce 100644 --- a/.env.example +++ b/.env.example @@ -15,3 +15,20 @@ AMI_ANTENNA_API_BASE_URL=http://localhost:8000/api/v2 AMI_ANTENNA_API_AUTH_TOKEN=your_antenna_auth_token_here AMI_ANTENNA_API_BATCH_SIZE=4 AMI_ANTENNA_SERVICE_NAME="AMI Data Companion" + +# Worker process recycling. The worker exits cleanly at the next batch +# boundary when it receives SIGUSR1 (e.g. `supervisorctl signal USR1 `), +# after processing AMI_WORKER_MAX_JOBS jobs, or when its resident memory, +# sampled between jobs, reaches or exceeds AMI_WORKER_MAX_RSS_MB (in MiB). +# 0 disables either cap. +# +# Recycling requires a process manager configured to restart the worker after +# a *successful* exit, because a drain exits with status 0. The usual defaults +# do not: supervisord's `autorestart=unexpected` and systemd's +# `Restart=on-failure` both treat status 0 as a reason to stay stopped, as +# does Docker's `on-failure` restart policy. Set `autorestart=true` +# (supervisord), `Restart=always` (systemd), or `--restart always` / +# `unless-stopped` (Docker). Without it the worker recycles once and stays +# down. +# AMI_WORKER_MAX_JOBS=10 +# AMI_WORKER_MAX_RSS_MB=8192 diff --git a/trapdata/antenna/tests/test_worker_drain.py b/trapdata/antenna/tests/test_worker_drain.py new file mode 100644 index 0000000..05f9f5f --- /dev/null +++ b/trapdata/antenna/tests/test_worker_drain.py @@ -0,0 +1,388 @@ +"""Unit tests for the worker's drain-and-exit machinery. + +The worker exits at a batch boundary — never mid-batch — when an operator +sends SIGUSR1, after ``AMI_WORKER_MAX_JOBS`` jobs, or when its resident +memory, sampled between jobs, exceeds ``AMI_WORKER_MAX_RSS_MB``. These tests +cover the drain state and signal handler, the between-jobs recycle checks +and their per-job memory log line, that ``_process_job`` stops cleanly at a +batch boundary while flushing pending result posts, and that the polling +loop exits after a drain request. No models are loaded; inference is mocked +out. +""" + +import multiprocessing +import os +import signal +import unittest +from types import SimpleNamespace +from unittest import TestCase +from unittest.mock import MagicMock, patch + +from trapdata.antenna.worker import ( + _after_job_check, + _current_rss_bytes, + _DrainRequest, + _install_drain_handler, + _install_parent_drain_handler, + _process_job, + _worker_loop, +) + +_HAS_SIGUSR1 = hasattr(signal, "SIGUSR1") + + +def _spawn_event(): + """An event of the kind the parent shares with spawned worker instances.""" + return multiprocessing.get_context("spawn").Event() + + +class TestDrainRequest(TestCase): + def test_starts_unrequested(self): + drain = _DrainRequest() + assert not drain.requested + assert drain.reason is None + + def test_signal_handler_sets_requested(self): + drain = _DrainRequest() + # The handler ignores its (signum, frame) arguments. + drain.handle_signal(None, None) + assert drain.requested + assert drain.reason == "SIGUSR1" + + def test_first_reason_is_kept(self): + drain = _DrainRequest() + drain.request("first") + drain.request("second") + assert drain.reason == "first" + + @unittest.skipUnless(_HAS_SIGUSR1, "platform has no SIGUSR1") + def test_installed_handler_receives_a_real_signal(self): + drain = _DrainRequest() + previous = signal.getsignal(signal.SIGUSR1) + try: + _install_drain_handler(drain) + os.kill(os.getpid(), signal.SIGUSR1) + assert drain.requested + finally: + signal.signal(signal.SIGUSR1, previous) + + +class TestParentDrainPropagation(TestCase): + """A drain reaches a worker instance even if it arrives before that + instance is ready to handle a signal. + + Worker instances are started with the spawn method, so each begins life + with SIGUSR1 at its default action of terminating the process. The parent + therefore never signals them: it sets a shared event that each instance + reads at its own safe points. The guarantee under test is that an event + already set when the instance starts is honoured, which is exactly the + case a signal could not survive. + """ + + def test_event_set_before_start_is_observed(self): + event = _spawn_event() + event.set() + drain = _DrainRequest(event) + assert drain.requested + assert "parent" in drain.reason + + def test_event_set_after_start_is_observed(self): + event = _spawn_event() + drain = _DrainRequest(event) + assert not drain.requested + event.set() + assert drain.requested + + def test_unset_event_does_not_request_a_drain(self): + # Without this, an always-true `requested` would pass the positive + # cases above while draining every worker immediately. + drain = _DrainRequest(_spawn_event()) + assert not drain.requested + assert drain.reason is None + + def test_own_reason_survives_a_parent_event(self): + event = _spawn_event() + drain = _DrainRequest(event) + drain.request("SIGUSR1") + event.set() + assert drain.reason == "SIGUSR1" + + @unittest.skipUnless(_HAS_SIGUSR1, "platform has no SIGUSR1") + def test_parent_handler_publishes_to_the_event(self): + event = _spawn_event() + previous = signal.getsignal(signal.SIGUSR1) + try: + _install_parent_drain_handler(event) + assert not event.is_set() + os.kill(os.getpid(), signal.SIGUSR1) + assert event.is_set() + finally: + signal.signal(signal.SIGUSR1, previous) + + +def _recycle_settings(max_rss_mb: int = 0, max_jobs: int = 0) -> SimpleNamespace: + return SimpleNamespace(worker_max_rss_mb=max_rss_mb, worker_max_jobs=max_jobs) + + +class TestAfterJobCheck(TestCase): + """Each recycle cap must trigger only past its threshold and stay off at 0. + + Both directions matter: a cap that never fires leaves memory growth + unbounded, and a cap that fires with the triggers disabled would recycle + every deployment that did not opt in. + """ + + @patch("trapdata.antenna.worker._current_rss_bytes", return_value=2 * 1024**3) + def test_default_caps_never_request_drain(self, _rss): + drain = _DrainRequest() + _after_job_check(drain, _recycle_settings(), jobs_processed=100) + assert not drain.requested + + @patch("trapdata.antenna.worker._current_rss_bytes", return_value=2 * 1024**3) + def test_rss_per_job_line_is_always_logged(self, _rss): + # The log line is the instrument for answering whether memory climbs + # job after job in a deployment, so it must not depend on any cap + # being enabled. + with patch("trapdata.antenna.worker.logger") as mock_logger: + _after_job_check(_DrainRequest(), _recycle_settings(), jobs_processed=3) + logged = " ".join(str(c.args[0]) for c in mock_logger.info.call_args_list) + assert "Resident memory after job 3: 2048 MiB" in logged + + @patch("trapdata.antenna.worker._current_rss_bytes", return_value=9 * 1024**3) + def test_rss_over_cap_requests_drain(self, _rss): + drain = _DrainRequest() + _after_job_check(drain, _recycle_settings(max_rss_mb=8192), jobs_processed=1) + assert drain.requested + # 9 GiB = 9216 MiB, over an 8192 MiB cap + assert "9216" in drain.reason + + @patch("trapdata.antenna.worker._current_rss_bytes", return_value=8 * 1024**3) + def test_rss_exactly_at_cap_requests_drain(self, _rss): + # The cap is inclusive, which is what the setting's documentation + # promises operators. + drain = _DrainRequest() + _after_job_check(drain, _recycle_settings(max_rss_mb=8192), jobs_processed=1) + assert drain.requested + + @patch("trapdata.antenna.worker._current_rss_bytes", return_value=2 * 1024**3) + def test_rss_under_cap_does_nothing(self, _rss): + drain = _DrainRequest() + _after_job_check(drain, _recycle_settings(max_rss_mb=8192), jobs_processed=1) + assert not drain.requested + + @patch("trapdata.antenna.worker._current_rss_bytes", return_value=None) + def test_unreadable_rss_skips_the_memory_cap(self, _rss): + drain = _DrainRequest() + _after_job_check(drain, _recycle_settings(max_rss_mb=8192), jobs_processed=1) + assert not drain.requested + + @patch("trapdata.antenna.worker._current_rss_bytes", return_value=None) + def test_job_cap_works_without_rss(self, _rss): + drain = _DrainRequest() + _after_job_check(drain, _recycle_settings(max_jobs=5), jobs_processed=5) + assert drain.requested + assert "cap 5" in drain.reason + + @patch("trapdata.antenna.worker._current_rss_bytes", return_value=2 * 1024**3) + def test_job_cap_not_reached_does_nothing(self, _rss): + drain = _DrainRequest() + _after_job_check(drain, _recycle_settings(max_jobs=5), jobs_processed=4) + assert not drain.requested + + @patch("trapdata.antenna.worker._current_rss_bytes", return_value=2 * 1024**3) + def test_existing_drain_reason_is_kept(self, _rss): + drain = _DrainRequest() + drain.request("SIGUSR1") + _after_job_check(drain, _recycle_settings(max_jobs=1), jobs_processed=1) + assert drain.reason == "SIGUSR1" + + def test_current_rss_reads_a_positive_value_where_proc_exists(self): + rss = _current_rss_bytes() + if rss is None: + self.skipTest("/proc/self/status not available on this platform") + assert rss > 0 + + +class TestRecycleCapValidation(TestCase): + """Only 0 disables a recycle cap. + + A negative value would otherwise be accepted and read as "disabled" by + the threshold checks, quietly turning off a cap an operator meant to set. + """ + + def test_negative_caps_are_rejected(self): + import pydantic + + from trapdata.settings import Settings + + for field in ("worker_max_rss_mb", "worker_max_jobs"): + with self.subTest(field=field): + with self.assertRaises(pydantic.ValidationError): + Settings(**{field: -1}) + + def test_zero_and_positive_caps_are_accepted(self): + from trapdata.settings import Settings + + for value in (0, 4096): + settings = Settings(worker_max_rss_mb=value, worker_max_jobs=value) + assert settings.worker_max_rss_mb == value + assert settings.worker_max_jobs == value + + +def _fake_batches(n: int) -> list[dict]: + """Minimal truthy batches; contents are irrelevant with inference mocked.""" + return [{"images": [object()], "image_ids": [f"img_{i}"]} for i in range(n)] + + +@patch("trapdata.antenna.worker.ResultPoster") +@patch("trapdata.antenna.worker.APIMothDetector") +@patch("trapdata.antenna.worker._process_batch", return_value=(1, 0, [], 0.0, 0.0)) +@patch("trapdata.antenna.worker.should_filter_detections", return_value=False) +@patch.dict( + "trapdata.antenna.worker.CLASSIFIER_CHOICES", + {"fake_pipeline": MagicMock()}, +) +@patch("trapdata.antenna.worker.get_rest_dataloader") +class TestProcessJobStopsAtBatchBoundary(TestCase): + """``should_stop`` must stop a job between batches, never lose posted work. + + The guard has to hold in both directions: with a stop requested after the + first batch, later batches must not run; without one, every batch must + run — otherwise a should_stop that always trips would pass the first + assertion while silently truncating every job. + """ + + def _settings(self) -> MagicMock: + settings = MagicMock() + settings.antenna_api_base_url = "http://testserver/api/v2" + settings.antenna_api_auth_token = "test-token" + return settings + + def _give_poster_real_metrics(self, mock_poster_cls: MagicMock) -> None: + # The end-of-job summary formats these with numeric format specs, + # which a bare MagicMock attribute cannot satisfy. + mock_poster_cls.return_value.get_metrics.return_value = SimpleNamespace( + total_posts=1, + successful_posts=1, + failed_posts=0, + success_rate=100.0, + total_post_time=0.1, + max_queue_size=1, + ) + + def test_stops_after_current_batch( + self, + mock_loader, + mock_should_filter, + mock_process_batch, + mock_detector, + mock_poster_cls, + ): + mock_loader.return_value = _fake_batches(3) + self._give_poster_real_metrics(mock_poster_cls) + + stop_flag = {"stop": False} + + def on_batch(batch_num: int, items: int): + # Simulates a drain request arriving while batch 1 is in flight. + stop_flag["stop"] = True + + result = _process_job( + "fake_pipeline", + job_id=123, + settings=self._settings(), + on_batch_complete=on_batch, + should_stop=lambda: stop_flag["stop"], + ) + + assert result is True + # Batch 1 ran; batches 2 and 3 were left for the next worker. + assert mock_process_batch.call_count == 1 + # Pending posts were flushed and the poster shut down cleanly. + poster = mock_poster_cls.return_value + poster.wait_for_all_posts.assert_called_once() + poster.shutdown.assert_called_once() + + def test_without_stop_request_all_batches_run( + self, + mock_loader, + mock_should_filter, + mock_process_batch, + mock_detector, + mock_poster_cls, + ): + mock_loader.return_value = _fake_batches(3) + self._give_poster_real_metrics(mock_poster_cls) + + result = _process_job( + "fake_pipeline", + job_id=123, + settings=self._settings(), + should_stop=lambda: False, + ) + + assert result is True + assert mock_process_batch.call_count == 3 + + +class TestWorkerLoopExitsOnDrain(TestCase): + @unittest.skipUnless(_HAS_SIGUSR1, "platform has no SIGUSR1") + @patch("trapdata.antenna.worker.get_jobs") + @patch("trapdata.antenna.worker.read_settings") + def test_loop_returns_after_signal(self, mock_read_settings, mock_get_jobs): + settings = MagicMock() + settings.antenna_service_name = "test-worker" + settings.worker_max_rss_mb = 0 + settings.worker_max_jobs = 0 + mock_read_settings.return_value = settings + + polls: list[int] = [] + + def fake_get_jobs(**kwargs): + polls.append(1) + if len(polls) > 1: + raise AssertionError("worker loop kept polling after the drain request") + # The signal is delivered to this process before get_jobs returns, + # like an operator signalling mid-poll. + os.kill(os.getpid(), signal.SIGUSR1) + return [] + + mock_get_jobs.side_effect = fake_get_jobs + + previous = signal.getsignal(signal.SIGUSR1) + try: + _worker_loop(0, ["fake_pipeline"]) + finally: + signal.signal(signal.SIGUSR1, previous) + + assert polls == [1] + + @patch("trapdata.antenna.worker.get_jobs") + @patch("trapdata.antenna.worker.read_settings") + def test_loop_honours_a_drain_requested_before_it_started( + self, mock_read_settings, mock_get_jobs + ): + """A worker instance started after the parent already drained exits + without claiming a job. + + This is the startup case a forwarded signal cannot cover: the parent + may drain while an instance is still starting, before that instance + can handle SIGUSR1. + """ + settings = MagicMock() + settings.antenna_service_name = "test-worker" + settings.worker_max_rss_mb = 0 + settings.worker_max_jobs = 0 + mock_read_settings.return_value = settings + + event = _spawn_event() + event.set() + + previous = signal.getsignal(signal.SIGUSR1) if _HAS_SIGUSR1 else None + try: + _worker_loop(0, ["fake_pipeline"], event) + finally: + if _HAS_SIGUSR1: + signal.signal(signal.SIGUSR1, previous) + + mock_get_jobs.assert_not_called() diff --git a/trapdata/antenna/worker.py b/trapdata/antenna/worker.py index 2b7e1db..41a2dd4 100644 --- a/trapdata/antenna/worker.py +++ b/trapdata/antenna/worker.py @@ -3,8 +3,11 @@ from __future__ import annotations import datetime +import signal import time from collections.abc import Callable +from multiprocessing.synchronize import Event as EventType +from types import FrameType import numpy as np import torch @@ -29,6 +32,141 @@ MAX_PENDING_POSTS = 5 # Maximum number of concurrent result posts before blocking SLEEP_TIME_SECONDS = 5 +_MIB = 1024 * 1024 + + +class _DrainRequest: + """A request for the worker process to exit at the next safe point. + + Set by SIGUSR1 (an operator or process supervisor asking for a recycle), + by the worker itself when a between-jobs recycle cap is hit (see + ``_after_job_check``), or by the parent of a multi-GPU worker group + through ``parent_event``. A safe point is a batch boundary: the batch in + flight finishes and its results are posted before the process exits with + status 0, so a process manager configured to restart the worker after a + clean exit starts a fresh one (fresh memory, file descriptors, and + DataLoader state) without losing in-flight work. + """ + + def __init__(self, parent_event: EventType | None = None) -> None: + self.reason: str | None = None + self._parent_event = parent_event + + @property + def requested(self) -> bool: + # Reading the parent's event here means every safe point already in + # the worker — the polling loop and the batch-boundary + # ``should_stop`` — observes a drain published by the parent without + # having to check for it separately. + if self._parent_event is not None and self._parent_event.is_set(): + self.request("parent process requested a drain") + return self.reason is not None + + def request(self, reason: str) -> None: + # The first reason wins; later triggers change nothing. + if self.reason is None: + self.reason = reason + + def handle_signal(self, signum: int, frame: FrameType | None) -> None: + logger.info( + "SIGUSR1 received: will finish the batch in flight, post its " + "results, and exit cleanly" + ) + self.request("SIGUSR1") + + +def _install_drain_handler(drain: _DrainRequest) -> None: + """Route SIGUSR1 to ``drain`` on platforms that have the signal.""" + if hasattr(signal, "SIGUSR1"): + signal.signal(signal.SIGUSR1, drain.handle_signal) + + +def _install_parent_drain_handler(drain_event: EventType) -> None: + """Let the parent of a worker group publish drains without signalling children. + + The parent owns SIGUSR1 on behalf of the whole group, because signalling + the children directly is unsafe during their startup: they are started + with the spawn method, so each begins life with SIGUSR1 at its default + action of terminating the process and only installs its own handler once + ``_worker_loop`` runs. A signal landing in that window kills the child, + and torch responds by terminating its siblings — the opposite of a clean + drain. Publishing the request through an event the children read at their + own safe points removes that window rather than narrowing it. + + Installing this before any child is spawned also closes the matching + window in the parent, which is otherwise still at SIGUSR1's default + action while its children are already running, so a drain arriving then + would kill the parent and orphan workers holding GPU memory. + """ + if not hasattr(signal, "SIGUSR1"): + return + + def _publish_drain(signum: int, frame: FrameType | None) -> None: + logger.info( + "SIGUSR1 received: asking every worker instance to finish the " + "batch in flight, post its results, and exit cleanly" + ) + drain_event.set() + + signal.signal(signal.SIGUSR1, _publish_drain) + + +def _current_rss_bytes() -> int | None: + """Resident set size of this process, or None where /proc is unavailable.""" + try: + with open("/proc/self/status") as status: + for line in status: + if line.startswith("VmRSS:"): + return int(line.split()[1]) * 1024 + except (OSError, ValueError): + return None + return None + + +def _after_job_check( + drain: _DrainRequest, settings: Settings, jobs_processed: int +) -> None: + """Log per-job memory and request a drain when a recycle cap is hit. + + Runs between jobs, when the working set is at its idle baseline, so the + logged reading — and the ``worker_max_rss_mb`` comparison — reflects + memory retained across jobs rather than a job's transient working set. + The log line lets operators answer "does resident memory keep climbing + job after job, or plateau after the first model load?" from ordinary + worker logs, without instrumenting the host. + + Two independent triggers, each disabled at 0 (the default): + + - ``worker_max_jobs``: drain after this many jobs. Deterministic bound + for retention that scales with jobs processed, useful even before the + retained memory's size or source is known. + - ``worker_max_rss_mb``: drain when resident memory reaches this many + MiB. Catches whatever the job cap does not predict. + """ + rss_bytes = _current_rss_bytes() + rss_mb = rss_bytes // _MIB if rss_bytes is not None else None + if rss_mb is not None: + logger.info(f"Resident memory after job {jobs_processed}: {rss_mb} MiB") + if drain.requested: + return + max_jobs = settings.worker_max_jobs + if max_jobs > 0 and jobs_processed >= max_jobs: + logger.info( + f"Processed {jobs_processed} jobs, reaching the " + f"{max_jobs}-job cap (AMI_WORKER_MAX_JOBS); exiting cleanly so " + "the process supervisor restarts a fresh worker" + ) + drain.request(f"{jobs_processed} jobs processed, cap {max_jobs}") + return + max_rss_mb = settings.worker_max_rss_mb + if max_rss_mb > 0 and rss_mb is not None and rss_mb >= max_rss_mb: + logger.warning( + f"Resident memory {rss_mb} MiB reached the {max_rss_mb} MiB cap " + "(AMI_WORKER_MAX_RSS_MB); exiting cleanly so the process " + "supervisor restarts a fresh worker" + ) + drain.request(f"resident memory {rss_mb} MiB reached the {max_rss_mb} MiB cap") + def run_worker(pipelines: list[str]): """Run the worker to process images from the REST API queue. @@ -55,14 +193,25 @@ def run_worker(pipelines: list[str]): gpu_count = torch.cuda.device_count() if gpu_count > 1: logger.info(f"Found {gpu_count} GPUs, spawning one AMI worker instance per GPU") + # A process manager delivers signals to this parent process, which + # relays them to the group through this event. Both the event and its + # handler are set up before the first child exists, so a drain + # arriving at any point during startup is recorded rather than lost. + # The event has to come from the spawn context to be shareable with + # processes started by mp.spawn. + drain_event = mp.get_context("spawn").Event() + _install_parent_drain_handler(drain_event) # Don't pass settings through mp.spawn — Settings contains enums that # can't be pickled. Each child process calls read_settings() itself. - mp.spawn( + context = mp.spawn( _worker_loop, - args=(pipelines,), + args=(pipelines, drain_event), nprocs=gpu_count, - join=True, + join=False, ) + assert context is not None + while not context.join(): + pass else: if gpu_count == 1: logger.info(f"Found 1 GPU: {torch.cuda.get_device_name(0)}") @@ -71,13 +220,28 @@ def run_worker(pipelines: list[str]): _worker_loop(0, pipelines) -def _worker_loop(gpu_id: int, pipelines: list[str]): +def _worker_loop( + gpu_id: int, pipelines: list[str], drain_event: EventType | None = None +): """Main polling loop for a single AMI worker instance, pinned to a specific GPU. + The loop runs until a drain is requested — by SIGUSR1, by a between-jobs + recycle cap (see ``_after_job_check``), or by the parent of a multi-GPU + group — and then returns so the process exits with status 0 for its + process manager to restart. + Args: gpu_id: GPU index to pin this AMI worker instance to (0 for CPU-only). pipelines: List of pipeline slugs to poll for jobs. + drain_event: Shared event the parent of a multi-GPU worker group sets + to ask this instance to drain. None when the worker runs + in-process and owns SIGUSR1 itself. """ + # Set up drain handling before any slower startup work, so a drain + # arriving early is recorded rather than acted on by SIGUSR1's default + # action of terminating the process. + drain = _DrainRequest(drain_event) + _install_drain_handler(drain) settings = read_settings() device = torch.device(f"cuda:{gpu_id}" if torch.cuda.is_available() else "cpu") if torch.cuda.is_available() and torch.cuda.device_count() > 0: @@ -90,7 +254,8 @@ def _worker_loop(gpu_id: int, pipelines: list[str]): full_service_name = get_full_service_name(settings.antenna_service_name) logger.info(f"Running worker as: {full_service_name}") - while True: + jobs_processed = 0 + while not drain.requested: # TODO CGJS: Support pulling and prioritizing single image tasks, which are used in interactive testing # These should probably come from a dedicated endpoint and should preempt batch jobs under the assumption that they # would run on the same GPU. @@ -104,30 +269,48 @@ def _worker_loop(gpu_id: int, pipelines: list[str]): pipeline_slugs=pipelines, ) for job_id, pipeline in jobs: + if drain.requested: + break logger.info( f"[GPU {gpu_id}] Processing job {job_id} with pipeline {pipeline}" ) + # A job that raised mid-processing still counts toward the + # recycle caps: it may have loaded models and run batches, so + # it retains memory like a completed job. Claims that yielded + # no work do not count. + work_attempted = False try: any_work_done = _process_job( pipeline=pipeline, job_id=job_id, settings=settings, device=device, + should_stop=lambda: drain.requested, ) any_jobs = any_jobs or any_work_done + work_attempted = any_work_done except Exception as e: logger.error( f"[GPU {gpu_id}] Failed to process job {job_id} with pipeline {pipeline}: {e}", exc_info=True, ) # Continue to next job rather than crashing the worker + work_attempted = True + if work_attempted: + jobs_processed += 1 + _after_job_check(drain, settings, jobs_processed) - if not any_jobs: + if not any_jobs and not drain.requested: logger.info( f"[GPU {gpu_id}] No jobs found, sleeping for {SLEEP_TIME_SECONDS} seconds" ) time.sleep(SLEEP_TIME_SECONDS) + logger.info( + f"[GPU {gpu_id}] Drain complete ({drain.reason}); exiting cleanly for " + "the process supervisor to restart a fresh worker" + ) + def _apply_binary_classification( binary_filter: "MothClassifierBinary", @@ -403,6 +586,7 @@ def _process_job( settings: Settings, device: torch.device | None = None, on_batch_complete: Callable | None = None, + should_stop: Callable[[], bool] | None = None, ) -> bool: """Run the worker to process images from the REST API queue. @@ -413,6 +597,10 @@ def _process_job( device: The device to use for processing. Auto-detected if None. on_batch_complete: Optional callback invoked after each batch, with kwargs batch_num (int) and items (int, cumulative items processed so far). + should_stop: Optional callable checked at every batch boundary. When + it returns True the job stops before starting another batch: + results for completed batches are still posted, and the job's + remaining tasks stay queued for the next worker to claim. Returns: True if any work was done, False otherwise """ @@ -453,6 +641,13 @@ def _process_job( _, t_total = log_time() try: for i, batch in enumerate(batch_source): + if should_stop and should_stop(): + logger.info( + f"Stop requested; leaving job {job_id} at a batch " + "boundary. Remaining tasks stay queued for the next " + "worker to claim." + ) + break cls_time = 0.0 det_time = 0.0 load_time, t = t() @@ -506,7 +701,7 @@ def _process_job( ) batch_total, t_total = t_total() logger.info( - f"Batch {i + 1}: {batch_total/max(n_items, 1):.2f}s/image, " + f"Batch {i + 1}: {batch_total / max(n_items, 1):.2f}s/image, " f"Classification time: {cls_time:.2f}s, Detection time: {det_time:.2f}s, " f"Load time: {load_time:.2f}s" ) diff --git a/trapdata/settings.py b/trapdata/settings.py index b07e043..d9ee07f 100644 --- a/trapdata/settings.py +++ b/trapdata/settings.py @@ -42,6 +42,14 @@ class Settings(BaseSettings): antenna_api_auth_token: str = "" antenna_service_name: str = "AMI Data Companion" antenna_api_batch_size: int = 24 + # The worker exits cleanly (for its process manager to restart) when + # resident memory, sampled between jobs, reaches or exceeds this many MiB. + # Only 0 disables the cap, so a negative value is rejected rather than + # silently turning the cap off. + worker_max_rss_mb: int = Field(default=0, ge=0) + # The worker exits cleanly after processing this many jobs. Bounds + # memory retained per job even when its size is not known. 0 disables. + worker_max_jobs: int = Field(default=0, ge=0) @pydantic.field_validator("image_base_path", "user_data_path") def validate_path(cls, v):