From fb74cb5e7cb9f233d275a63ad2d7ed9c3d1d0162 Mon Sep 17 00:00:00 2001 From: Michael Bunsen Date: Fri, 10 Apr 2026 18:55:02 -0700 Subject: [PATCH 01/14] feat(nats): surface queue lifecycle events on the per-job logger MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit TaskQueueManager now accepts an optional job_logger. When set, lifecycle events (stream/consumer create+reuse with state/config snapshot, publish failures, cleanup deletions, and a forensic consumer-stats line before deletion) are mirrored to the per-job logger in addition to the module logger. The UI job log now reflects what the NATS layer is actually doing for that specific job instead of a silent gap. Per-message and per-poll paths (publish_task success, reserve_tasks, acknowledge_task) intentionally stay on the module logger only — a 10k-image job would otherwise drown its own log. Lifecycle log lines are deduped per manager session so a loop over N images still only emits a single "Created NATS stream" line per job. cleanup_async_job_resources and queue_images_to_nats pass job.logger through to TaskQueueManager so real async_api jobs pick up the new logging without further caller changes. Closes RolnickLab/antenna#1220 --- ami/ml/orchestration/jobs.py | 24 +- ami/ml/orchestration/nats_queue.py | 196 +++++++++++--- ami/ml/orchestration/tests/test_nats_queue.py | 248 ++++++++++++++++++ 3 files changed, 421 insertions(+), 47 deletions(-) diff --git a/ami/ml/orchestration/jobs.py b/ami/ml/orchestration/jobs.py index 95c763b1b..af27a74e1 100644 --- a/ami/ml/orchestration/jobs.py +++ b/ami/ml/orchestration/jobs.py @@ -22,8 +22,13 @@ def cleanup_async_job_resources(job_id: int, _logger: logging.Logger) -> bool: Cleanup failures are logged but don't fail the job - data is already saved. Args: - job_id: The Job ID (integer primary key) - _logger: Logger to use for logging cleanup results + job_id: The Job ID (integer primary key). For ASYNC_API jobs this should + be called with the per-job logger (``job.logger``) so the UI log shows + cleanup events and the forensic consumer-stats snapshot that + ``TaskQueueManager.cleanup_job_resources`` emits before deletion. + _logger: Logger to use for logging cleanup results. Passed through to + TaskQueueManager so lifecycle events land on both the module logger + and the per-job logger. Returns: bool: True if both cleanups succeeded, False otherwise """ @@ -39,9 +44,10 @@ def cleanup_async_job_resources(job_id: int, _logger: logging.Logger) -> bool: except Exception as e: _logger.error(f"Error cleaning up Redis state for job {job_id}: {e}") - # Cleanup NATS resources + # Cleanup NATS resources. Pass _logger through so TaskQueueManager can + # log final consumer stats and deletion events against the per-job logger. async def cleanup(): - async with TaskQueueManager() as manager: + async with TaskQueueManager(job_logger=_logger) as manager: return await manager.cleanup_job_resources(job_id) try: @@ -97,16 +103,20 @@ async def queue_all_images(): successful_queues = 0 failed_queues = 0 - async with TaskQueueManager() as manager: + # Pass job.logger so stream/consumer setup and any publish failures + # appear in the UI job log (not just the module logger). Per-image + # success logs stay at module level so a 10k-image job doesn't drown + # the job log. + async with TaskQueueManager(job_logger=job.logger) as manager: for image_pk, task in tasks: try: - logger.info(f"Queueing image {image_pk} to stream for job '{job.pk}': {task.image_url}") + logger.debug(f"Queueing image {image_pk} to stream for job '{job.pk}': {task.image_url}") success = await manager.publish_task( job_id=job.pk, data=task, ) except Exception as e: - logger.error(f"Failed to queue image {image_pk} to stream for job '{job.pk}': {e}") + job.logger.error(f"Failed to queue image {image_pk} to stream for job '{job.pk}': {e}") success = False if success: diff --git a/ami/ml/orchestration/nats_queue.py b/ami/ml/orchestration/nats_queue.py index b6e9af254..d186bb765 100644 --- a/ami/ml/orchestration/nats_queue.py +++ b/ami/ml/orchestration/nats_queue.py @@ -54,21 +54,64 @@ class TaskQueueManager: nats_url: NATS server URL. Falls back to settings.NATS_URL, then "nats://nats:4222". max_ack_pending: Max unacknowledged messages per consumer. Falls back to settings.NATS_MAX_ACK_PENDING, then 1000. + job_logger: Optional per-job logger. When set, lifecycle events (stream / + consumer create or reuse, cleanup stats, publish failures) are mirrored + to this logger in addition to the module logger, so they appear in the + job's own log stream as seen from the UI. Per-message and per-poll + events stay on the module logger only to avoid drowning large jobs. Use as an async context manager: - async with TaskQueueManager() as manager: + async with TaskQueueManager(job_logger=job.logger) as manager: await manager.publish_task(123, {'data': 'value'}) tasks = await manager.reserve_tasks(123, count=64) await manager.acknowledge_task(tasks[0].reply_subject) """ - def __init__(self, nats_url: str | None = None, max_ack_pending: int | None = None): + def __init__( + self, + nats_url: str | None = None, + max_ack_pending: int | None = None, + job_logger: logging.Logger | None = None, + ): self.nats_url = nats_url or getattr(settings, "NATS_URL", "nats://nats:4222") self.max_ack_pending = ( max_ack_pending if max_ack_pending is not None else getattr(settings, "NATS_MAX_ACK_PENDING", 1000) ) + self.job_logger = job_logger self.nc: nats.NATS | None = None self.js: JetStreamContext | None = None + # Dedupe lifecycle log lines per manager session so a job that publishes + # hundreds of tasks doesn't emit hundreds of "reusing stream" messages. + self._streams_logged: set[int] = set() + self._consumers_logged: set[int] = set() + + def _log(self, level: int, msg: str) -> None: + """Log to both the module logger and the job logger (if set). + + Module logger still fires so ops dashboards (stdout/New Relic) aren't + affected. Job logger fires so the UI job log reflects what NATS is doing. + """ + logger.log(level, msg) + if self.job_logger is not None and self.job_logger is not logger: + self.job_logger.log(level, msg) + + @staticmethod + def _format_consumer_stats(info) -> str: + """Format ConsumerInfo into a compact stats string. + + All nats-py ConsumerInfo fields are Optional, so defensive access is + required: this method renders missing values as '?'. Used for both + reuse-announcements and forensic cleanup lines. + """ + delivered = info.delivered.consumer_seq if info.delivered is not None else "?" + ack_floor = info.ack_floor.consumer_seq if info.ack_floor is not None else "?" + return ( + f"delivered={delivered} " + f"ack_floor={ack_floor} " + f"num_pending={info.num_pending if info.num_pending is not None else '?'} " + f"num_ack_pending={info.num_ack_pending if info.num_ack_pending is not None else '?'} " + f"num_redelivered={info.num_redelivered if info.num_redelivered is not None else '?'}" + ) async def __aenter__(self): """Create connection on enter.""" @@ -127,27 +170,52 @@ async def _stream_exists(self, stream_name: str) -> bool: return False async def _ensure_stream(self, job_id: int): - """Ensure stream exists for the given job.""" + """Ensure stream exists for the given job. + + Logs a lifecycle line to both the module and job logger the first time it + sees a given job in this manager session (creation or reuse). Subsequent + calls in the same session are silent, so a job publishing N images doesn't + emit N log lines. + """ if self.js is None: raise RuntimeError("Connection is not open. Use TaskQueueManager as an async context manager.") - if not await self._job_stream_exists(job_id): - stream_name = self._get_stream_name(job_id) - subject = self._get_subject(job_id) - logger.warning(f"Stream {stream_name} does not exist") - # Stream doesn't exist, create it - await asyncio.wait_for( - self.js.add_stream( - name=stream_name, - subjects=[subject], - max_age=86400, # 24 hours retention - ), - timeout=NATS_JETSTREAM_TIMEOUT, - ) - logger.info(f"Created stream {stream_name}") + stream_name = self._get_stream_name(job_id) + subject = self._get_subject(job_id) + + try: + info = await asyncio.wait_for(self.js.stream_info(stream_name), timeout=NATS_JETSTREAM_TIMEOUT) + if job_id not in self._streams_logged: + self._log( + logging.INFO, + f"Reusing NATS stream {stream_name} " + f"(messages={info.state.messages}, last_seq={info.state.last_seq})", + ) + self._streams_logged.add(job_id) + return + except nats.js.errors.NotFoundError: + pass + + await asyncio.wait_for( + self.js.add_stream( + name=stream_name, + subjects=[subject], + max_age=86400, # 24 hours retention + ), + timeout=NATS_JETSTREAM_TIMEOUT, + ) + self._log(logging.INFO, f"Created NATS stream {stream_name}") + self._streams_logged.add(job_id) async def _ensure_consumer(self, job_id: int): - """Ensure consumer exists for the given job.""" + """Ensure consumer exists for the given job. + + On first sight in this manager session (creation or reuse), emits a line + to both the module and job logger. On creation the line includes the + config snapshot (max_deliver, ack_wait, max_ack_pending, deliver_policy, + ack_policy) so forensic readers can see exactly what delivery semantics + were in effect. + """ if self.js is None: raise RuntimeError("Connection is not open. Use TaskQueueManager as an async context manager.") @@ -160,27 +228,43 @@ async def _ensure_consumer(self, job_id: int): self.js.consumer_info(stream_name, consumer_name), timeout=NATS_JETSTREAM_TIMEOUT, ) - logger.debug(f"Consumer {consumer_name} already exists: {info}") + if job_id not in self._consumers_logged: + self._log( + logging.INFO, + f"Reusing NATS consumer {consumer_name} ({self._format_consumer_stats(info)})", + ) + self._consumers_logged.add(job_id) + return except asyncio.TimeoutError: raise # NATS unreachable — let caller handle it except Exception: - # Consumer doesn't exist, create it - await asyncio.wait_for( - self.js.add_consumer( - stream=stream_name, - config=ConsumerConfig( - durable_name=consumer_name, - ack_policy=AckPolicy.EXPLICIT, - ack_wait=TASK_TTR, # Visibility timeout (TTR) - max_deliver=5, # Max retry attempts - deliver_policy=DeliverPolicy.ALL, - max_ack_pending=self.max_ack_pending, - filter_subject=subject, - ), + # Consumer doesn't exist, fall through to create it. + pass + + await asyncio.wait_for( + self.js.add_consumer( + stream=stream_name, + config=ConsumerConfig( + durable_name=consumer_name, + ack_policy=AckPolicy.EXPLICIT, + ack_wait=TASK_TTR, # Visibility timeout (TTR) + max_deliver=5, # Max retry attempts + deliver_policy=DeliverPolicy.ALL, + max_ack_pending=self.max_ack_pending, + filter_subject=subject, ), - timeout=NATS_JETSTREAM_TIMEOUT, - ) - logger.info(f"Created consumer {consumer_name}") + ), + timeout=NATS_JETSTREAM_TIMEOUT, + ) + self._log( + logging.INFO, + f"Created NATS consumer {consumer_name} " + f"(max_deliver=5, ack_wait={TASK_TTR}s, " + f"max_ack_pending={self.max_ack_pending}, " + f"deliver_policy={DeliverPolicy.ALL.value}, " + f"ack_policy={AckPolicy.EXPLICIT.value})", + ) + self._consumers_logged.add(job_id) async def publish_task(self, job_id: int, data: PipelineProcessingTask) -> bool: """ @@ -212,7 +296,10 @@ async def publish_task(self, job_id: int, data: PipelineProcessingTask) -> bool: return True except Exception as e: - logger.error(f"Failed to publish task to stream for job '{job_id}': {e}") + # Per-message success logs stay at module level (noise in 10k-image + # jobs), but a failure on even a single publish deserves to surface + # in the job log — otherwise the failure path is invisible to users. + self._log(logging.ERROR, f"Failed to publish task to stream for job '{job_id}': {e}") return False async def reserve_tasks(self, job_id: int, count: int, timeout: float = 5) -> list[PipelineProcessingTask]: @@ -292,6 +379,31 @@ async def acknowledge_task(self, reply_subject: str) -> bool: logger.error(f"Failed to acknowledge task: {e}") return False + async def _log_final_consumer_stats(self, job_id: int) -> None: + """Log one forensic line about the consumer state before deletion. + + This is the single most useful line in a post-mortem: it tells you how + many messages were delivered, how many were acked, and how many were + redelivered before the consumer vanished. Failures here must NOT block + cleanup — if the consumer or stream is already gone, just skip it. + """ + if self.js is None: + return + stream_name = self._get_stream_name(job_id) + consumer_name = self._get_consumer_name(job_id) + try: + info = await asyncio.wait_for( + self.js.consumer_info(stream_name, consumer_name), + timeout=NATS_JETSTREAM_TIMEOUT, + ) + except Exception as e: + logger.debug(f"Could not fetch consumer info for {consumer_name} before deletion: {e}") + return + self._log( + logging.INFO, + f"Finalizing NATS consumer {consumer_name} before deletion " f"({self._format_consumer_stats(info)})", + ) + async def delete_consumer(self, job_id: int) -> bool: """ Delete the consumer for a job. @@ -313,10 +425,10 @@ async def delete_consumer(self, job_id: int) -> bool: self.js.delete_consumer(stream_name, consumer_name), timeout=NATS_JETSTREAM_TIMEOUT, ) - logger.info(f"Deleted consumer {consumer_name} for job '{job_id}'") + self._log(logging.INFO, f"Deleted NATS consumer {consumer_name} for job '{job_id}'") return True except Exception as e: - logger.error(f"Failed to delete consumer for job '{job_id}': {e}") + self._log(logging.ERROR, f"Failed to delete NATS consumer for job '{job_id}': {e}") return False async def delete_stream(self, job_id: int) -> bool: @@ -339,10 +451,10 @@ async def delete_stream(self, job_id: int) -> bool: self.js.delete_stream(stream_name), timeout=NATS_JETSTREAM_TIMEOUT, ) - logger.info(f"Deleted stream {stream_name} for job '{job_id}'") + self._log(logging.INFO, f"Deleted NATS stream {stream_name} for job '{job_id}'") return True except Exception as e: - logger.error(f"Failed to delete stream for job '{job_id}': {e}") + self._log(logging.ERROR, f"Failed to delete NATS stream for job '{job_id}': {e}") return False async def _setup_advisory_stream(self): @@ -482,6 +594,10 @@ async def cleanup_job_resources(self, job_id: int) -> bool: Returns: bool: True if successful, False otherwise """ + # Log a forensic snapshot of the consumer state BEFORE we destroy it. + # This is the highest-leverage line for post-mortem investigations. + await self._log_final_consumer_stats(job_id) + # Delete consumer first, then stream, then the durable DLQ advisory consumer consumer_deleted = await self.delete_consumer(job_id) stream_deleted = await self.delete_stream(job_id) diff --git a/ami/ml/orchestration/tests/test_nats_queue.py b/ami/ml/orchestration/tests/test_nats_queue.py index da47f3429..86144de9d 100644 --- a/ami/ml/orchestration/tests/test_nats_queue.py +++ b/ami/ml/orchestration/tests/test_nats_queue.py @@ -1,6 +1,7 @@ """Unit tests for TaskQueueManager.""" import json +import logging import unittest from unittest.mock import AsyncMock, MagicMock, patch @@ -235,3 +236,250 @@ async def test_get_dead_letter_image_ids_no_messages(self): self.assertEqual(result, []) mock_psub.unsubscribe.assert_called_once() + + +class TestTaskQueueManagerJobLogger(unittest.IsolatedAsyncioTestCase): + """Tests covering the job_logger lifecycle-mirroring behavior (#1220).""" + + def _create_sample_task(self): + return PipelineProcessingTask( + id="task-1", + image_id="img-1", + image_url="https://example.com/image.jpg", + ) + + def _create_mock_nats_connection(self): + """Duplicate of the sibling helper — kept local so the two test classes + can evolve independently.""" + nc = MagicMock() + nc.is_closed = False + nc.close = AsyncMock() + nc.flush = AsyncMock() + + js = MagicMock() + js.stream_info = AsyncMock() + js.add_stream = AsyncMock() + js.add_consumer = AsyncMock() + js.consumer_info = AsyncMock() + js.publish = AsyncMock(return_value=MagicMock(seq=1)) + js.pull_subscribe = AsyncMock() + js.delete_consumer = AsyncMock() + js.delete_stream = AsyncMock() + + return nc, js + + def _make_consumer_info(self, delivered=10, ack_floor=8, num_pending=2, num_ack_pending=2, num_redelivered=1): + """Build a ConsumerInfo-like MagicMock with nested SequenceInfo stubs.""" + info = MagicMock() + info.delivered = MagicMock(consumer_seq=delivered) + info.ack_floor = MagicMock(consumer_seq=ack_floor) + info.num_pending = num_pending + info.num_ack_pending = num_ack_pending + info.num_redelivered = num_redelivered + return info + + def _make_stream_info(self, messages=5, last_seq=5): + info = MagicMock() + info.state = MagicMock(messages=messages, last_seq=last_seq) + return info + + def _make_captured_logger(self) -> logging.Logger: + """A real Logger that captures to a list — better than MagicMock.log + because it exercises the actual `logger.log(level, msg)` dispatch and + surfaces any type surprises in the call site.""" + log_logger = logging.getLogger(f"test.job_logger.{id(self)}") + log_logger.handlers.clear() + log_logger.setLevel(logging.DEBUG) + + captured = [] + + class CaptureHandler(logging.Handler): + def emit(self, record): + captured.append((record.levelno, record.getMessage())) + + log_logger.addHandler(CaptureHandler()) + log_logger._captured = captured # type: ignore[attr-defined] + return log_logger + + async def test_create_stream_and_consumer_logs_to_job_logger(self): + """First publish on a brand-new job should log stream/consumer creation + to both the module logger and the passed-in job_logger.""" + nc, js = self._create_mock_nats_connection() + js.stream_info.side_effect = nats.js.errors.NotFoundError() + js.consumer_info.side_effect = nats.js.errors.NotFoundError() + + job_logger = self._make_captured_logger() + captured = job_logger._captured # type: ignore[attr-defined] + + with patch("ami.ml.orchestration.nats_queue.get_connection", AsyncMock(return_value=(nc, js))): + async with TaskQueueManager(job_logger=job_logger) as manager: + await manager.publish_task(42, self._create_sample_task()) + + messages = [m for _, m in captured] + self.assertTrue( + any("Created NATS stream job_42" in m for m in messages), + f"expected stream-create log on job_logger, got {messages}", + ) + self.assertTrue( + any("Created NATS consumer job-42-consumer" in m for m in messages), + f"expected consumer-create log on job_logger, got {messages}", + ) + # Config snapshot should appear on the creation line. + self.assertTrue( + any("max_deliver=5" in m and "ack_policy=" in m for m in messages), + f"expected consumer config snapshot in log, got {messages}", + ) + + async def test_publish_success_does_not_spam_job_logger(self): + """After the first publish, subsequent publishes in the same session + must NOT emit new setup lines — per-message logging is forbidden for + 10k-image jobs.""" + nc, js = self._create_mock_nats_connection() + # First call hits NotFound (create path), subsequent calls succeed (reuse path) + js.stream_info.side_effect = [nats.js.errors.NotFoundError(), self._make_stream_info()] + js.consumer_info.side_effect = [nats.js.errors.NotFoundError(), self._make_consumer_info()] + + job_logger = self._make_captured_logger() + captured = job_logger._captured # type: ignore[attr-defined] + + with patch("ami.ml.orchestration.nats_queue.get_connection", AsyncMock(return_value=(nc, js))): + async with TaskQueueManager(job_logger=job_logger) as manager: + await manager.publish_task(42, self._create_sample_task()) + captured_after_first = list(captured) + await manager.publish_task(42, self._create_sample_task()) + + new_messages = captured[len(captured_after_first) :] + # The second publish should not add any lifecycle log lines — dedup set + # should swallow them after the first publish for this job_id. + lifecycle_terms = ("Created NATS", "Reusing NATS") + for _, m in new_messages: + self.assertFalse( + any(term in m for term in lifecycle_terms), + f"unexpected lifecycle log on second publish: {m}", + ) + + async def test_reuse_stream_and_consumer_logs_with_stats(self): + """When stream and consumer already exist, the reuse line should include + a summary of current consumer state so forensic readers can tell whether + the queue is empty, backed up, or mid-redelivery.""" + nc, js = self._create_mock_nats_connection() + js.stream_info.return_value = self._make_stream_info(messages=17, last_seq=17) + js.consumer_info.return_value = self._make_consumer_info( + delivered=12, ack_floor=10, num_pending=5, num_ack_pending=2, num_redelivered=3 + ) + + job_logger = self._make_captured_logger() + captured = job_logger._captured # type: ignore[attr-defined] + + with patch("ami.ml.orchestration.nats_queue.get_connection", AsyncMock(return_value=(nc, js))): + async with TaskQueueManager(job_logger=job_logger) as manager: + await manager.publish_task(99, self._create_sample_task()) + + messages = [m for _, m in captured] + self.assertTrue( + any("Reusing NATS stream job_99" in m and "messages=17" in m and "last_seq=17" in m for m in messages), + f"expected reuse-stream log with state, got {messages}", + ) + self.assertTrue( + any( + "Reusing NATS consumer job-99-consumer" in m + and "delivered=12" in m + and "ack_floor=10" in m + and "num_pending=5" in m + and "num_redelivered=3" in m + for m in messages + ), + f"expected reuse-consumer log with stats, got {messages}", + ) + + async def test_cleanup_logs_final_consumer_stats_before_delete(self): + """cleanup_job_resources must emit a forensic snapshot of the consumer + state BEFORE the delete calls land. This is the single most useful line + for a post-mortem — without it, the consumer is already gone by the + time anyone investigates.""" + nc, js = self._create_mock_nats_connection() + final_info = self._make_consumer_info( + delivered=434, ack_floor=420, num_pending=0, num_ack_pending=14, num_redelivered=5 + ) + js.consumer_info.return_value = final_info + + job_logger = self._make_captured_logger() + captured = job_logger._captured # type: ignore[attr-defined] + + with patch("ami.ml.orchestration.nats_queue.get_connection", AsyncMock(return_value=(nc, js))): + async with TaskQueueManager(job_logger=job_logger) as manager: + await manager.cleanup_job_resources(123) + + messages = [m for _, m in captured] + finalizing_idx = None + delete_idx = None + for i, m in enumerate(messages): + if "Finalizing NATS consumer job-123-consumer" in m: + finalizing_idx = i + if delete_idx is None and "Deleted NATS consumer job-123-consumer" in m: + delete_idx = i + + self.assertIsNotNone(finalizing_idx, f"expected forensic finalize-log, got {messages}") + self.assertIsNotNone(delete_idx, f"expected delete-log, got {messages}") + self.assertLess( + finalizing_idx, # type: ignore[arg-type] + delete_idx, # type: ignore[arg-type] + "finalize snapshot must log BEFORE the delete", + ) + # The stats themselves should make it into the line. + final_line = messages[finalizing_idx] # type: ignore[index] + for expected in ("delivered=434", "ack_floor=420", "num_redelivered=5"): + self.assertIn(expected, final_line) + + async def test_cleanup_tolerates_missing_consumer(self): + """If the consumer is already gone when cleanup runs, the pre-delete + stats call must NOT raise or block — cleanup is called in failure + paths where the consumer may have already been deleted.""" + nc, js = self._create_mock_nats_connection() + js.consumer_info.side_effect = nats.js.errors.NotFoundError() + + job_logger = self._make_captured_logger() + + with patch("ami.ml.orchestration.nats_queue.get_connection", AsyncMock(return_value=(nc, js))): + async with TaskQueueManager(job_logger=job_logger) as manager: + # Must not raise. + result = await manager.cleanup_job_resources(77) + + # delete_consumer / delete_stream are still called on the mock and + # return truthy, so overall cleanup is reported successful. + self.assertTrue(result) + + async def test_publish_failure_surfaces_on_job_logger(self): + """A failed publish (which today only logs to the module logger) must + now also land on the job_logger so users see the failure in the UI.""" + nc, js = self._create_mock_nats_connection() + js.stream_info.return_value = self._make_stream_info() + js.consumer_info.return_value = self._make_consumer_info() + js.publish = AsyncMock(side_effect=RuntimeError("simulated nats outage")) + + job_logger = self._make_captured_logger() + captured = job_logger._captured # type: ignore[attr-defined] + + with patch("ami.ml.orchestration.nats_queue.get_connection", AsyncMock(return_value=(nc, js))): + async with TaskQueueManager(job_logger=job_logger) as manager: + result = await manager.publish_task(55, self._create_sample_task()) + + self.assertFalse(result) + messages = [m for level, m in captured if level >= logging.ERROR] + self.assertTrue( + any("Failed to publish task" in m and "simulated nats outage" in m for m in messages), + f"expected publish failure on job_logger, got {messages}", + ) + + async def test_no_job_logger_falls_back_to_module_logger_only(self): + """When job_logger is None (e.g., module-level uses like advisory + listener), lifecycle logs must still be emitted to the module logger + without crashing on a None attribute access.""" + nc, js = self._create_mock_nats_connection() + js.stream_info.side_effect = nats.js.errors.NotFoundError() + js.consumer_info.side_effect = nats.js.errors.NotFoundError() + + with patch("ami.ml.orchestration.nats_queue.get_connection", AsyncMock(return_value=(nc, js))): + async with TaskQueueManager() as manager: # no job_logger passed + # Must not raise. + await manager.publish_task(1, self._create_sample_task()) From 7e8ce284c78a52f98a90d076fb09fbc0d392a0cf Mon Sep 17 00:00:00 2001 From: Michael Bunsen Date: Sat, 11 Apr 2026 14:01:55 -0700 Subject: [PATCH 02/14] fix(nats): bridge job-logger mirroring through sync_to_async MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit TaskQueueManager._log() was calling job_logger.log() synchronously from inside async methods (_ensure_stream, _ensure_consumer, cleanup). That triggers JobLogHandler.emit() which does a Django ORM refresh_from_db + save — forbidden from an event loop. Every lifecycle line was silently dropped with "Failed to save logs for job #N: You cannot call this from an async context", defeating the point of the original change. Convert _log to async and await sync_to_async(job_logger.log)(...) so the ORM work runs in a thread. Update all call sites to await. Apply the same fix to the publish-failure path in queue_images_to_nats. Verified on ami-demo with job #74: lifecycle lines fired on module logger but "Failed to save logs" errors swallowed the job-logger mirror. --- ami/ml/orchestration/jobs.py | 10 ++++++-- ami/ml/orchestration/nats_queue.py | 40 +++++++++++++++++++----------- 2 files changed, 34 insertions(+), 16 deletions(-) diff --git a/ami/ml/orchestration/jobs.py b/ami/ml/orchestration/jobs.py index af27a74e1..cd53c7157 100644 --- a/ami/ml/orchestration/jobs.py +++ b/ami/ml/orchestration/jobs.py @@ -1,6 +1,6 @@ import logging -from asgiref.sync import async_to_sync +from asgiref.sync import async_to_sync, sync_to_async from ami.jobs.models import Job, JobState from ami.main.models import SourceImage @@ -116,7 +116,13 @@ async def queue_all_images(): data=task, ) except Exception as e: - job.logger.error(f"Failed to queue image {image_pk} to stream for job '{job.pk}': {e}") + # job.logger.error triggers a sync Django ORM save inside + # JobLogHandler.emit, which raises SynchronousOnlyOperation + # when called directly from the event loop. Bridge it so + # the line actually lands in job.logs.stdout. + await sync_to_async(job.logger.error)( + f"Failed to queue image {image_pk} to stream for job '{job.pk}': {e}" + ) success = False if success: diff --git a/ami/ml/orchestration/nats_queue.py b/ami/ml/orchestration/nats_queue.py index d186bb765..4eb59f356 100644 --- a/ami/ml/orchestration/nats_queue.py +++ b/ami/ml/orchestration/nats_queue.py @@ -15,6 +15,7 @@ import logging import nats +from asgiref.sync import sync_to_async from django.conf import settings from nats.js import JetStreamContext from nats.js.api import AckPolicy, ConsumerConfig, DeliverPolicy @@ -85,15 +86,26 @@ def __init__( self._streams_logged: set[int] = set() self._consumers_logged: set[int] = set() - def _log(self, level: int, msg: str) -> None: + async def _log(self, level: int, msg: str) -> None: """Log to both the module logger and the job logger (if set). - Module logger still fires so ops dashboards (stdout/New Relic) aren't - affected. Job logger fires so the UI job log reflects what NATS is doing. + Module logger fires synchronously (ops dashboards / stdout / New Relic + are unaffected). The job logger call is bridged through + ``sync_to_async`` because Django's ``JobLogHandler`` does an ORM + ``refresh_from_db`` + ``save`` on every emit — calling that directly + from the event loop raises ``SynchronousOnlyOperation`` and the log + line is silently dropped. The bridge offloads the handler work to a + thread so the line actually lands in ``job.logs.stdout``. + + Exceptions from the job logger are swallowed so logging a lifecycle + event never breaks the actual NATS operation. """ logger.log(level, msg) if self.job_logger is not None and self.job_logger is not logger: - self.job_logger.log(level, msg) + try: + await sync_to_async(self.job_logger.log)(level, msg) + except Exception as e: + logger.warning(f"Failed to mirror log to job logger: {e}") @staticmethod def _format_consumer_stats(info) -> str: @@ -186,7 +198,7 @@ async def _ensure_stream(self, job_id: int): try: info = await asyncio.wait_for(self.js.stream_info(stream_name), timeout=NATS_JETSTREAM_TIMEOUT) if job_id not in self._streams_logged: - self._log( + await self._log( logging.INFO, f"Reusing NATS stream {stream_name} " f"(messages={info.state.messages}, last_seq={info.state.last_seq})", @@ -204,7 +216,7 @@ async def _ensure_stream(self, job_id: int): ), timeout=NATS_JETSTREAM_TIMEOUT, ) - self._log(logging.INFO, f"Created NATS stream {stream_name}") + await self._log(logging.INFO, f"Created NATS stream {stream_name}") self._streams_logged.add(job_id) async def _ensure_consumer(self, job_id: int): @@ -229,7 +241,7 @@ async def _ensure_consumer(self, job_id: int): timeout=NATS_JETSTREAM_TIMEOUT, ) if job_id not in self._consumers_logged: - self._log( + await self._log( logging.INFO, f"Reusing NATS consumer {consumer_name} ({self._format_consumer_stats(info)})", ) @@ -256,7 +268,7 @@ async def _ensure_consumer(self, job_id: int): ), timeout=NATS_JETSTREAM_TIMEOUT, ) - self._log( + await self._log( logging.INFO, f"Created NATS consumer {consumer_name} " f"(max_deliver=5, ack_wait={TASK_TTR}s, " @@ -299,7 +311,7 @@ async def publish_task(self, job_id: int, data: PipelineProcessingTask) -> bool: # Per-message success logs stay at module level (noise in 10k-image # jobs), but a failure on even a single publish deserves to surface # in the job log — otherwise the failure path is invisible to users. - self._log(logging.ERROR, f"Failed to publish task to stream for job '{job_id}': {e}") + await self._log(logging.ERROR, f"Failed to publish task to stream for job '{job_id}': {e}") return False async def reserve_tasks(self, job_id: int, count: int, timeout: float = 5) -> list[PipelineProcessingTask]: @@ -399,7 +411,7 @@ async def _log_final_consumer_stats(self, job_id: int) -> None: except Exception as e: logger.debug(f"Could not fetch consumer info for {consumer_name} before deletion: {e}") return - self._log( + await self._log( logging.INFO, f"Finalizing NATS consumer {consumer_name} before deletion " f"({self._format_consumer_stats(info)})", ) @@ -425,10 +437,10 @@ async def delete_consumer(self, job_id: int) -> bool: self.js.delete_consumer(stream_name, consumer_name), timeout=NATS_JETSTREAM_TIMEOUT, ) - self._log(logging.INFO, f"Deleted NATS consumer {consumer_name} for job '{job_id}'") + await self._log(logging.INFO, f"Deleted NATS consumer {consumer_name} for job '{job_id}'") return True except Exception as e: - self._log(logging.ERROR, f"Failed to delete NATS consumer for job '{job_id}': {e}") + await self._log(logging.ERROR, f"Failed to delete NATS consumer for job '{job_id}': {e}") return False async def delete_stream(self, job_id: int) -> bool: @@ -451,10 +463,10 @@ async def delete_stream(self, job_id: int) -> bool: self.js.delete_stream(stream_name), timeout=NATS_JETSTREAM_TIMEOUT, ) - self._log(logging.INFO, f"Deleted NATS stream {stream_name} for job '{job_id}'") + await self._log(logging.INFO, f"Deleted NATS stream {stream_name} for job '{job_id}'") return True except Exception as e: - self._log(logging.ERROR, f"Failed to delete NATS stream for job '{job_id}': {e}") + await self._log(logging.ERROR, f"Failed to delete NATS stream for job '{job_id}': {e}") return False async def _setup_advisory_stream(self): From 439300ab2e4f26dc1fd7a3bab4366dd64d2363b9 Mon Sep 17 00:00:00 2001 From: Michael Bunsen Date: Sat, 11 Apr 2026 18:44:49 -0700 Subject: [PATCH 03/14] fix(nats): narrow consumer_info exception + stop forwarding module logger to TaskQueueManager in cleanup Addresses CodeRabbit review on #1222. 1. _ensure_consumer caught broad Exception when consumer_info() failed, masking auth/API/transient JetStream errors as "consumer missing" and emitting misleading creation logs. Narrowed to NotFoundError to match the pattern already used in _ensure_stream (line 208). 2. cleanup_async_job_resources forwarded its `_logger` argument into TaskQueueManager as `job_logger`. One caller (_fail_job on the Job.DoesNotExist path in ami/jobs/tasks.py:198) passes a plain module logger, which would then have cleanup lifecycle lines mirrored into an unrelated logger via sync_to_async. Added a separate `job_logger` parameter, defaulted to None, and updated the two callers that have real job context (_fail_job happy path, cleanup_async_job_if_needed) to pass `job.logger` explicitly. The DoesNotExist path leaves job_logger=None, so TaskQueueManager falls through to the module logger only. Tests: 18/18 in test_nats_queue.py pass, pre-commit clean. Co-Authored-By: Claude Opus 4.6 (1M context) --- ami/jobs/tasks.py | 7 +++++-- ami/ml/orchestration/jobs.py | 33 ++++++++++++++++++++---------- ami/ml/orchestration/nats_queue.py | 7 +++++-- 3 files changed, 32 insertions(+), 15 deletions(-) diff --git a/ami/jobs/tasks.py b/ami/jobs/tasks.py index d8ff89d39..45f0126e1 100644 --- a/ami/jobs/tasks.py +++ b/ami/jobs/tasks.py @@ -192,9 +192,12 @@ def _fail_job(job_id: int, reason: str) -> None: job.save(update_fields=["status", "progress", "finished_at"]) job.logger.error(f"Job {job_id} marked as FAILURE: {reason}") - cleanup_async_job_resources(job.pk, job.logger) + cleanup_async_job_resources(job.pk, job.logger, job_logger=job.logger) except Job.DoesNotExist: logger.error(f"Cannot fail job {job_id}: not found") + # No job_logger here — the job row is gone, so cleanup lifecycle lines + # have nowhere to be mirrored to. TaskQueueManager falls through to + # the module logger. cleanup_async_job_resources(job_id, logger) @@ -423,7 +426,7 @@ def cleanup_async_job_if_needed(job) -> None: # import here to avoid circular imports from ami.ml.orchestration.jobs import cleanup_async_job_resources - cleanup_async_job_resources(job.pk, job.logger) + cleanup_async_job_resources(job.pk, job.logger, job_logger=job.logger) @task_prerun.connect(sender=run_job) diff --git a/ami/ml/orchestration/jobs.py b/ami/ml/orchestration/jobs.py index cd53c7157..2aabae02e 100644 --- a/ami/ml/orchestration/jobs.py +++ b/ami/ml/orchestration/jobs.py @@ -11,7 +11,11 @@ logger = logging.getLogger(__name__) -def cleanup_async_job_resources(job_id: int, _logger: logging.Logger) -> bool: +def cleanup_async_job_resources( + job_id: int, + _logger: logging.Logger, + job_logger: logging.Logger | None = None, +) -> bool: """ Clean up NATS JetStream and Redis resources for a completed job. @@ -22,13 +26,17 @@ def cleanup_async_job_resources(job_id: int, _logger: logging.Logger) -> bool: Cleanup failures are logged but don't fail the job - data is already saved. Args: - job_id: The Job ID (integer primary key). For ASYNC_API jobs this should - be called with the per-job logger (``job.logger``) so the UI log shows - cleanup events and the forensic consumer-stats snapshot that - ``TaskQueueManager.cleanup_job_resources`` emits before deletion. - _logger: Logger to use for logging cleanup results. Passed through to - TaskQueueManager so lifecycle events land on both the module logger - and the per-job logger. + job_id: The Job ID (integer primary key). + _logger: Logger to use for the local Redis/NATS outcome lines emitted + by this function itself. May be a plain module logger when the + caller has no job context (e.g. the ``Job.DoesNotExist`` path in + ``_fail_job``). + job_logger: Optional per-job logger (``job.logger``) to forward to + ``TaskQueueManager`` so the UI job log sees the forensic + consumer-stats snapshot and the stream/consumer delete lines. + Must only be set when the caller actually has a ``job.logger`` — + otherwise cleanup lifecycle lines would be mirrored into an + unrelated module logger. Returns: bool: True if both cleanups succeeded, False otherwise """ @@ -44,10 +52,13 @@ def cleanup_async_job_resources(job_id: int, _logger: logging.Logger) -> bool: except Exception as e: _logger.error(f"Error cleaning up Redis state for job {job_id}: {e}") - # Cleanup NATS resources. Pass _logger through so TaskQueueManager can - # log final consumer stats and deletion events against the per-job logger. + # Cleanup NATS resources. Forward the per-job logger (if any) so the + # forensic pre-delete consumer-stats snapshot and the delete lifecycle + # lines land in the UI job log. When job_logger is None (e.g. the + # Job.DoesNotExist fallback), TaskQueueManager falls back to the module + # logger only. async def cleanup(): - async with TaskQueueManager(job_logger=_logger) as manager: + async with TaskQueueManager(job_logger=job_logger) as manager: return await manager.cleanup_job_resources(job_id) try: diff --git a/ami/ml/orchestration/nats_queue.py b/ami/ml/orchestration/nats_queue.py index 4eb59f356..74cb7e0f8 100644 --- a/ami/ml/orchestration/nats_queue.py +++ b/ami/ml/orchestration/nats_queue.py @@ -249,8 +249,11 @@ async def _ensure_consumer(self, job_id: int): return except asyncio.TimeoutError: raise # NATS unreachable — let caller handle it - except Exception: - # Consumer doesn't exist, fall through to create it. + except nats.js.errors.NotFoundError: + # Consumer doesn't exist, fall through to create it. Other + # JetStream errors (auth, API, transient) must propagate so we + # don't mask them as "missing consumer" and emit misleading + # creation logs. pass await asyncio.wait_for( From 90071971e7316eace6f2b803cb331112ec4d0cc2 Mon Sep 17 00:00:00 2001 From: Michael Bunsen Date: Sun, 12 Apr 2026 07:22:12 -0700 Subject: [PATCH 04/14] refactor(jobs): simplify cleanup_async_job_resources to match codebase job_logger pattern MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Drop the separate `_logger` parameter that was redundant with `job_logger`. Follow the existing convention (e.g. save_results in ami/jobs/tasks.py): `_log = job_logger or logger` — use per-job logger when available, module logger otherwise. Callers now read cleanly: cleanup_async_job_resources(job.pk, job_logger=job.logger) # has job cleanup_async_job_resources(job_id) # no job Co-Authored-By: Claude Opus 4.6 (1M context) --- ami/jobs/tasks.py | 9 +++------ ami/ml/orchestration/jobs.py | 39 ++++++++++++++---------------------- 2 files changed, 18 insertions(+), 30 deletions(-) diff --git a/ami/jobs/tasks.py b/ami/jobs/tasks.py index 45f0126e1..e176dd08c 100644 --- a/ami/jobs/tasks.py +++ b/ami/jobs/tasks.py @@ -192,13 +192,10 @@ def _fail_job(job_id: int, reason: str) -> None: job.save(update_fields=["status", "progress", "finished_at"]) job.logger.error(f"Job {job_id} marked as FAILURE: {reason}") - cleanup_async_job_resources(job.pk, job.logger, job_logger=job.logger) + cleanup_async_job_resources(job.pk, job_logger=job.logger) except Job.DoesNotExist: logger.error(f"Cannot fail job {job_id}: not found") - # No job_logger here — the job row is gone, so cleanup lifecycle lines - # have nowhere to be mirrored to. TaskQueueManager falls through to - # the module logger. - cleanup_async_job_resources(job_id, logger) + cleanup_async_job_resources(job_id) def _ack_task_via_nats(reply_subject: str, job_logger: logging.Logger) -> None: @@ -426,7 +423,7 @@ def cleanup_async_job_if_needed(job) -> None: # import here to avoid circular imports from ami.ml.orchestration.jobs import cleanup_async_job_resources - cleanup_async_job_resources(job.pk, job.logger, job_logger=job.logger) + cleanup_async_job_resources(job.pk, job_logger=job.logger) @task_prerun.connect(sender=run_job) diff --git a/ami/ml/orchestration/jobs.py b/ami/ml/orchestration/jobs.py index 2aabae02e..7a4c314cf 100644 --- a/ami/ml/orchestration/jobs.py +++ b/ami/ml/orchestration/jobs.py @@ -11,11 +11,7 @@ logger = logging.getLogger(__name__) -def cleanup_async_job_resources( - job_id: int, - _logger: logging.Logger, - job_logger: logging.Logger | None = None, -) -> bool: +def cleanup_async_job_resources(job_id: int, job_logger: logging.Logger | None = None) -> bool: """ Clean up NATS JetStream and Redis resources for a completed job. @@ -27,19 +23,15 @@ def cleanup_async_job_resources( Args: job_id: The Job ID (integer primary key). - _logger: Logger to use for the local Redis/NATS outcome lines emitted - by this function itself. May be a plain module logger when the - caller has no job context (e.g. the ``Job.DoesNotExist`` path in - ``_fail_job``). - job_logger: Optional per-job logger (``job.logger``) to forward to - ``TaskQueueManager`` so the UI job log sees the forensic - consumer-stats snapshot and the stream/consumer delete lines. - Must only be set when the caller actually has a ``job.logger`` — - otherwise cleanup lifecycle lines would be mirrored into an - unrelated module logger. + job_logger: Per-job logger (``job.logger``) when the caller has a job + context. Falls back to the module logger when None (e.g. the + ``Job.DoesNotExist`` path in ``_fail_job``). Also forwarded to + ``TaskQueueManager`` so the forensic consumer-stats snapshot and + stream/consumer delete lines land in the UI job log. Returns: bool: True if both cleanups succeeded, False otherwise """ + _log = job_logger or logger redis_success = False nats_success = False @@ -47,15 +39,14 @@ def cleanup_async_job_resources( try: state_manager = AsyncJobStateManager(job_id) state_manager.cleanup() - _logger.info(f"Cleaned up Redis state for job {job_id}") + _log.info(f"Cleaned up Redis state for job {job_id}") redis_success = True except Exception as e: - _logger.error(f"Error cleaning up Redis state for job {job_id}: {e}") + _log.error(f"Error cleaning up Redis state for job {job_id}: {e}") - # Cleanup NATS resources. Forward the per-job logger (if any) so the - # forensic pre-delete consumer-stats snapshot and the delete lifecycle - # lines land in the UI job log. When job_logger is None (e.g. the - # Job.DoesNotExist fallback), TaskQueueManager falls back to the module + # Cleanup NATS resources. Forward job_logger to TaskQueueManager so the + # forensic pre-delete consumer-stats snapshot lands in the UI job log. + # When job_logger is None, TaskQueueManager falls back to the module # logger only. async def cleanup(): async with TaskQueueManager(job_logger=job_logger) as manager: @@ -64,11 +55,11 @@ async def cleanup(): try: nats_success = async_to_sync(cleanup)() if nats_success: - _logger.info(f"Cleaned up NATS resources for job {job_id}") + _log.info(f"Cleaned up NATS resources for job {job_id}") else: - _logger.warning(f"Failed to clean up NATS resources for job {job_id}") + _log.warning(f"Failed to clean up NATS resources for job {job_id}") except Exception as e: - _logger.error(f"Error cleaning up NATS resources for job {job_id}: {e}") + _log.error(f"Error cleaning up NATS resources for job {job_id}: {e}") return redis_success and nats_success From f8d09ac644c183bea844e835db14a80767564a26 Mon Sep 17 00:00:00 2001 From: Michael Bunsen Date: Sun, 12 Apr 2026 07:41:24 -0700 Subject: [PATCH 05/14] refactor(nats): read consumer config from ConsumerInfo instead of hardcoding in creation log The "Created NATS consumer" log line was hardcoding max_deliver=5 and interpolating TASK_TTR for ack_wait. Now reads from the ConsumerInfo returned by add_consumer(), so the log always reflects what the server accepted. Added _format_consumer_config() alongside the existing _format_consumer_stats() for the two different log contexts (creation vs runtime stats). Updated test mocks to return ConsumerInfo-like objects with a config sub-object from add_consumer. Co-Authored-By: Claude Opus 4.6 (1M context) --- ami/ml/orchestration/nats_queue.py | 29 ++++++++++++++----- ami/ml/orchestration/tests/test_nats_queue.py | 27 +++++++++++++++-- 2 files changed, 47 insertions(+), 9 deletions(-) diff --git a/ami/ml/orchestration/nats_queue.py b/ami/ml/orchestration/nats_queue.py index 74cb7e0f8..addffa393 100644 --- a/ami/ml/orchestration/nats_queue.py +++ b/ami/ml/orchestration/nats_queue.py @@ -107,9 +107,28 @@ async def _log(self, level: int, msg: str) -> None: except Exception as e: logger.warning(f"Failed to mirror log to job logger: {e}") + @staticmethod + def _format_consumer_config(info) -> str: + """Format ConsumerInfo config into a compact creation-time string. + + Reads the actual config from the ConsumerInfo returned by + ``add_consumer`` or ``consumer_info``, so the log always reflects + what the server accepted rather than what we requested. + """ + cfg = info.config if info.config is not None else None + if cfg is None: + return "config=?" + return ( + f"max_deliver={cfg.max_deliver if cfg.max_deliver is not None else '?'}, " + f"ack_wait={cfg.ack_wait if cfg.ack_wait is not None else '?'}s, " + f"max_ack_pending={cfg.max_ack_pending if cfg.max_ack_pending is not None else '?'}, " + f"deliver_policy={cfg.deliver_policy if cfg.deliver_policy is not None else '?'}, " + f"ack_policy={cfg.ack_policy if cfg.ack_policy is not None else '?'}" + ) + @staticmethod def _format_consumer_stats(info) -> str: - """Format ConsumerInfo into a compact stats string. + """Format ConsumerInfo into a compact runtime stats string. All nats-py ConsumerInfo fields are Optional, so defensive access is required: this method renders missing values as '?'. Used for both @@ -256,7 +275,7 @@ async def _ensure_consumer(self, job_id: int): # creation logs. pass - await asyncio.wait_for( + info = await asyncio.wait_for( self.js.add_consumer( stream=stream_name, config=ConsumerConfig( @@ -273,11 +292,7 @@ async def _ensure_consumer(self, job_id: int): ) await self._log( logging.INFO, - f"Created NATS consumer {consumer_name} " - f"(max_deliver=5, ack_wait={TASK_TTR}s, " - f"max_ack_pending={self.max_ack_pending}, " - f"deliver_policy={DeliverPolicy.ALL.value}, " - f"ack_policy={AckPolicy.EXPLICIT.value})", + f"Created NATS consumer {consumer_name} ({self._format_consumer_config(info)})", ) self._consumers_logged.add(job_id) diff --git a/ami/ml/orchestration/tests/test_nats_queue.py b/ami/ml/orchestration/tests/test_nats_queue.py index 86144de9d..d1d651450 100644 --- a/ami/ml/orchestration/tests/test_nats_queue.py +++ b/ami/ml/orchestration/tests/test_nats_queue.py @@ -268,14 +268,34 @@ def _create_mock_nats_connection(self): return nc, js - def _make_consumer_info(self, delivered=10, ack_floor=8, num_pending=2, num_ack_pending=2, num_redelivered=1): - """Build a ConsumerInfo-like MagicMock with nested SequenceInfo stubs.""" + def _make_consumer_info( + self, + delivered=10, + ack_floor=8, + num_pending=2, + num_ack_pending=2, + num_redelivered=1, + max_deliver=5, + ack_wait=30, + max_ack_pending=1000, + deliver_policy="all", + ack_policy="explicit", + ): + """Build a ConsumerInfo-like MagicMock with nested SequenceInfo stubs + and a config sub-object for creation-time logging.""" info = MagicMock() info.delivered = MagicMock(consumer_seq=delivered) info.ack_floor = MagicMock(consumer_seq=ack_floor) info.num_pending = num_pending info.num_ack_pending = num_ack_pending info.num_redelivered = num_redelivered + info.config = MagicMock( + max_deliver=max_deliver, + ack_wait=ack_wait, + max_ack_pending=max_ack_pending, + deliver_policy=deliver_policy, + ack_policy=ack_policy, + ) return info def _make_stream_info(self, messages=5, last_seq=5): @@ -307,6 +327,7 @@ async def test_create_stream_and_consumer_logs_to_job_logger(self): nc, js = self._create_mock_nats_connection() js.stream_info.side_effect = nats.js.errors.NotFoundError() js.consumer_info.side_effect = nats.js.errors.NotFoundError() + js.add_consumer = AsyncMock(return_value=self._make_consumer_info(delivered=0, ack_floor=0)) job_logger = self._make_captured_logger() captured = job_logger._captured # type: ignore[attr-defined] @@ -338,6 +359,7 @@ async def test_publish_success_does_not_spam_job_logger(self): # First call hits NotFound (create path), subsequent calls succeed (reuse path) js.stream_info.side_effect = [nats.js.errors.NotFoundError(), self._make_stream_info()] js.consumer_info.side_effect = [nats.js.errors.NotFoundError(), self._make_consumer_info()] + js.add_consumer = AsyncMock(return_value=self._make_consumer_info(delivered=0, ack_floor=0)) job_logger = self._make_captured_logger() captured = job_logger._captured # type: ignore[attr-defined] @@ -478,6 +500,7 @@ async def test_no_job_logger_falls_back_to_module_logger_only(self): nc, js = self._create_mock_nats_connection() js.stream_info.side_effect = nats.js.errors.NotFoundError() js.consumer_info.side_effect = nats.js.errors.NotFoundError() + js.add_consumer = AsyncMock(return_value=self._make_consumer_info(delivered=0, ack_floor=0)) with patch("ami.ml.orchestration.nats_queue.get_connection", AsyncMock(return_value=(nc, js))): async with TaskQueueManager() as manager: # no job_logger passed From 4c62d346324bb8987d450f812cdec8e67e6e825f Mon Sep 17 00:00:00 2001 From: Michael Bunsen Date: Sun, 12 Apr 2026 08:14:26 -0700 Subject: [PATCH 06/14] refactor(nats): skip redundant NATS calls after first ensure, fix enum rendering in config log MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three cleanup items found on self-review: 1. _ensure_stream and _ensure_consumer were calling stream_info/consumer_info on every publish_task (once per image). Since the stream and consumer are never deleted mid-flight (cleanup uses a separate manager session), added an early return when job_id is already in the logged set. Saves 2 NATS round-trips per image after the first. 2. _format_consumer_config was rendering enum fields as "DeliverPolicy.ALL" instead of "all" — added a _val() helper that unwraps .value when present. 3. Removed now-redundant inner dedup checks (the early return makes them always-true). Co-Authored-By: Claude Opus 4.6 (1M context) --- ami/ml/orchestration/nats_queue.py | 49 +++++++++++++++++------------- 1 file changed, 28 insertions(+), 21 deletions(-) diff --git a/ami/ml/orchestration/nats_queue.py b/ami/ml/orchestration/nats_queue.py index addffa393..655dbb41b 100644 --- a/ami/ml/orchestration/nats_queue.py +++ b/ami/ml/orchestration/nats_queue.py @@ -118,12 +118,17 @@ def _format_consumer_config(info) -> str: cfg = info.config if info.config is not None else None if cfg is None: return "config=?" + + def _val(v): + """Unwrap enum .value if present, pass through scalars.""" + return v.value if hasattr(v, "value") else v + return ( - f"max_deliver={cfg.max_deliver if cfg.max_deliver is not None else '?'}, " - f"ack_wait={cfg.ack_wait if cfg.ack_wait is not None else '?'}s, " - f"max_ack_pending={cfg.max_ack_pending if cfg.max_ack_pending is not None else '?'}, " - f"deliver_policy={cfg.deliver_policy if cfg.deliver_policy is not None else '?'}, " - f"ack_policy={cfg.ack_policy if cfg.ack_policy is not None else '?'}" + f"max_deliver={_val(cfg.max_deliver) if cfg.max_deliver is not None else '?'}, " + f"ack_wait={_val(cfg.ack_wait) if cfg.ack_wait is not None else '?'}s, " + f"max_ack_pending={_val(cfg.max_ack_pending) if cfg.max_ack_pending is not None else '?'}, " + f"deliver_policy={_val(cfg.deliver_policy) if cfg.deliver_policy is not None else '?'}, " + f"ack_policy={_val(cfg.ack_policy) if cfg.ack_policy is not None else '?'}" ) @staticmethod @@ -205,9 +210,11 @@ async def _ensure_stream(self, job_id: int): Logs a lifecycle line to both the module and job logger the first time it sees a given job in this manager session (creation or reuse). Subsequent - calls in the same session are silent, so a job publishing N images doesn't - emit N log lines. + calls in the same session skip the NATS round-trip entirely — the stream + won't be deleted mid-flight (cleanup uses a separate manager session). """ + if job_id in self._streams_logged: + return if self.js is None: raise RuntimeError("Connection is not open. Use TaskQueueManager as an async context manager.") @@ -216,13 +223,12 @@ async def _ensure_stream(self, job_id: int): try: info = await asyncio.wait_for(self.js.stream_info(stream_name), timeout=NATS_JETSTREAM_TIMEOUT) - if job_id not in self._streams_logged: - await self._log( - logging.INFO, - f"Reusing NATS stream {stream_name} " - f"(messages={info.state.messages}, last_seq={info.state.last_seq})", - ) - self._streams_logged.add(job_id) + await self._log( + logging.INFO, + f"Reusing NATS stream {stream_name} " + f"(messages={info.state.messages}, last_seq={info.state.last_seq})", + ) + self._streams_logged.add(job_id) return except nats.js.errors.NotFoundError: pass @@ -245,8 +251,10 @@ async def _ensure_consumer(self, job_id: int): to both the module and job logger. On creation the line includes the config snapshot (max_deliver, ack_wait, max_ack_pending, deliver_policy, ack_policy) so forensic readers can see exactly what delivery semantics - were in effect. + were in effect. Subsequent calls skip the NATS round-trip. """ + if job_id in self._consumers_logged: + return if self.js is None: raise RuntimeError("Connection is not open. Use TaskQueueManager as an async context manager.") @@ -259,12 +267,11 @@ async def _ensure_consumer(self, job_id: int): self.js.consumer_info(stream_name, consumer_name), timeout=NATS_JETSTREAM_TIMEOUT, ) - if job_id not in self._consumers_logged: - await self._log( - logging.INFO, - f"Reusing NATS consumer {consumer_name} ({self._format_consumer_stats(info)})", - ) - self._consumers_logged.add(job_id) + await self._log( + logging.INFO, + f"Reusing NATS consumer {consumer_name} ({self._format_consumer_stats(info)})", + ) + self._consumers_logged.add(job_id) return except asyncio.TimeoutError: raise # NATS unreachable — let caller handle it From 0d3160c81e5a1b5d46f39e4533b8ff759e4db8c3 Mon Sep 17 00:00:00 2001 From: Michael Bunsen Date: Mon, 13 Apr 2026 15:03:05 -0700 Subject: [PATCH 07/14] fix(nats): preserve traceback on publish failure + defensive stream state access MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses two CodeRabbit findings: 1. queue_images_to_nats publish exception path only wrote a stringified error to the job log. Added logger.exception() on the module logger so ops dashboards keep the full traceback. The job.logger.error bridge still writes the user-facing message. 2. _ensure_stream accessed info.state.messages / info.state.last_seq without defending against None, unlike _format_consumer_stats which already does. Match the defensive pattern so the reuse line doesn't blow up if the server ever returns a StreamInfo with no state. Also pushing back on the "ack_wait nanoseconds" finding in the review thread — nats-py does convert it to seconds in from_response via _convert_nanoseconds (source: val / _NANOSECOND where _NANOSECOND = 1e9). Co-Authored-By: Claude Opus 4.6 (1M context) --- ami/ml/orchestration/jobs.py | 2 ++ ami/ml/orchestration/nats_queue.py | 6 ++++-- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/ami/ml/orchestration/jobs.py b/ami/ml/orchestration/jobs.py index 7a4c314cf..cee9bb38d 100644 --- a/ami/ml/orchestration/jobs.py +++ b/ami/ml/orchestration/jobs.py @@ -118,6 +118,8 @@ async def queue_all_images(): data=task, ) except Exception as e: + # Module logger gets the full traceback for ops dashboards. + logger.exception("Failed to queue image %s to stream for job '%s'", image_pk, job.pk) # job.logger.error triggers a sync Django ORM save inside # JobLogHandler.emit, which raises SynchronousOnlyOperation # when called directly from the event loop. Bridge it so diff --git a/ami/ml/orchestration/nats_queue.py b/ami/ml/orchestration/nats_queue.py index 655dbb41b..d39cf14fc 100644 --- a/ami/ml/orchestration/nats_queue.py +++ b/ami/ml/orchestration/nats_queue.py @@ -223,10 +223,12 @@ async def _ensure_stream(self, job_id: int): try: info = await asyncio.wait_for(self.js.stream_info(stream_name), timeout=NATS_JETSTREAM_TIMEOUT) + state = info.state + messages = state.messages if state is not None else "?" + last_seq = state.last_seq if state is not None else "?" await self._log( logging.INFO, - f"Reusing NATS stream {stream_name} " - f"(messages={info.state.messages}, last_seq={info.state.last_seq})", + f"Reusing NATS stream {stream_name} (messages={messages}, last_seq={last_seq})", ) self._streams_logged.add(job_id) return From b06ec28e3379ef4b7f3d130642a5e165529a3465 Mon Sep 17 00:00:00 2001 From: Michael Bunsen Date: Mon, 13 Apr 2026 15:07:56 -0700 Subject: [PATCH 08/14] refactor(jobs): resolve job logger internally in cleanup_async_job_resources MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Match the pattern used by save_results in ami/jobs/tasks.py: take only job_id, resolve job.logger internally, fall back to the module logger when the Job row is missing. Keeps call sites consistent across the codebase — cleanup_async_job_resources(job.pk) everywhere — and makes the job object available inside the function for future use. Co-Authored-By: Claude Opus 4.6 (1M context) --- ami/jobs/tasks.py | 4 ++-- ami/ml/orchestration/jobs.py | 20 ++++++++++++++------ 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/ami/jobs/tasks.py b/ami/jobs/tasks.py index e176dd08c..ad3e18ca8 100644 --- a/ami/jobs/tasks.py +++ b/ami/jobs/tasks.py @@ -192,7 +192,7 @@ def _fail_job(job_id: int, reason: str) -> None: job.save(update_fields=["status", "progress", "finished_at"]) job.logger.error(f"Job {job_id} marked as FAILURE: {reason}") - cleanup_async_job_resources(job.pk, job_logger=job.logger) + cleanup_async_job_resources(job.pk) except Job.DoesNotExist: logger.error(f"Cannot fail job {job_id}: not found") cleanup_async_job_resources(job_id) @@ -423,7 +423,7 @@ def cleanup_async_job_if_needed(job) -> None: # import here to avoid circular imports from ami.ml.orchestration.jobs import cleanup_async_job_resources - cleanup_async_job_resources(job.pk, job_logger=job.logger) + cleanup_async_job_resources(job.pk) @task_prerun.connect(sender=run_job) diff --git a/ami/ml/orchestration/jobs.py b/ami/ml/orchestration/jobs.py index cee9bb38d..e7af3bb1f 100644 --- a/ami/ml/orchestration/jobs.py +++ b/ami/ml/orchestration/jobs.py @@ -11,7 +11,7 @@ logger = logging.getLogger(__name__) -def cleanup_async_job_resources(job_id: int, job_logger: logging.Logger | None = None) -> bool: +def cleanup_async_job_resources(job_id: int) -> bool: """ Clean up NATS JetStream and Redis resources for a completed job. @@ -21,17 +21,25 @@ def cleanup_async_job_resources(job_id: int, job_logger: logging.Logger | None = Cleanup failures are logged but don't fail the job - data is already saved. + Resolves the job (and its per-job logger) internally so callers only need + to pass the ``job_id`` — matches the pattern used by ``save_results`` in + ``ami/jobs/tasks.py``. If the ``Job`` row is gone (e.g. the + ``Job.DoesNotExist`` path in ``_fail_job``), the function falls back to + the module logger and TaskQueueManager's module-logger path. + Args: job_id: The Job ID (integer primary key). - job_logger: Per-job logger (``job.logger``) when the caller has a job - context. Falls back to the module logger when None (e.g. the - ``Job.DoesNotExist`` path in ``_fail_job``). Also forwarded to - ``TaskQueueManager`` so the forensic consumer-stats snapshot and - stream/consumer delete lines land in the UI job log. Returns: bool: True if both cleanups succeeded, False otherwise """ + job_logger: logging.Logger | None = None + try: + job = Job.objects.get(pk=job_id) + job_logger = job.logger + except Job.DoesNotExist: + pass _log = job_logger or logger + redis_success = False nats_success = False From 3257c778420ef9f5cbc51ed39862c7bf455197f6 Mon Sep 17 00:00:00 2001 From: Michael Bunsen Date: Mon, 13 Apr 2026 15:13:40 -0700 Subject: [PATCH 09/14] refactor(jobs): rename _log to job_logger to avoid name collision with TaskQueueManager._log MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two _log identifiers in the same module had incompatible semantics: - nats_queue.TaskQueueManager._log(level, msg) — async coroutine that fans out to module + job loggers with sync_to_async bridging. - jobs.cleanup_async_job_resources local _log = job_logger or logger — a plain Logger instance called as _log.info(...) / .error(...). Rename the local to job_logger and assign directly from `job.logger if job else logger`, matching the save_results pattern. Co-Authored-By: Claude Opus 4.6 (1M context) --- ami/ml/orchestration/jobs.py | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/ami/ml/orchestration/jobs.py b/ami/ml/orchestration/jobs.py index e7af3bb1f..8d5ec7911 100644 --- a/ami/ml/orchestration/jobs.py +++ b/ami/ml/orchestration/jobs.py @@ -32,13 +32,14 @@ def cleanup_async_job_resources(job_id: int) -> bool: Returns: bool: True if both cleanups succeeded, False otherwise """ - job_logger: logging.Logger | None = None + # Resolve the logger up front: job.logger when the Job exists, module + # logger otherwise. Matches the pattern used by save_results. + job: Job | None = None try: job = Job.objects.get(pk=job_id) - job_logger = job.logger except Job.DoesNotExist: pass - _log = job_logger or logger + job_logger: logging.Logger = job.logger if job else logger redis_success = False nats_success = False @@ -47,27 +48,26 @@ def cleanup_async_job_resources(job_id: int) -> bool: try: state_manager = AsyncJobStateManager(job_id) state_manager.cleanup() - _log.info(f"Cleaned up Redis state for job {job_id}") + job_logger.info(f"Cleaned up Redis state for job {job_id}") redis_success = True except Exception as e: - _log.error(f"Error cleaning up Redis state for job {job_id}: {e}") + job_logger.error(f"Error cleaning up Redis state for job {job_id}: {e}") - # Cleanup NATS resources. Forward job_logger to TaskQueueManager so the - # forensic pre-delete consumer-stats snapshot lands in the UI job log. - # When job_logger is None, TaskQueueManager falls back to the module - # logger only. + # Cleanup NATS resources. Only forward a real per-job logger to + # TaskQueueManager — passing the module logger would mirror cleanup + # lifecycle lines into an unrelated logger. async def cleanup(): - async with TaskQueueManager(job_logger=job_logger) as manager: + async with TaskQueueManager(job_logger=job.logger if job else None) as manager: return await manager.cleanup_job_resources(job_id) try: nats_success = async_to_sync(cleanup)() if nats_success: - _log.info(f"Cleaned up NATS resources for job {job_id}") + job_logger.info(f"Cleaned up NATS resources for job {job_id}") else: - _log.warning(f"Failed to clean up NATS resources for job {job_id}") + job_logger.warning(f"Failed to clean up NATS resources for job {job_id}") except Exception as e: - _log.error(f"Error cleaning up NATS resources for job {job_id}: {e}") + job_logger.error(f"Error cleaning up NATS resources for job {job_id}: {e}") return redis_success and nats_success From ae51c21d83ac6cdccd20bdc673a370fcd137a47d Mon Sep 17 00:00:00 2001 From: Michael Bunsen Date: Mon, 13 Apr 2026 15:18:11 -0700 Subject: [PATCH 10/14] chore(nats): tidy up minor code smells from self-review MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Drop redundant ternary: cfg = info.config, not `info.config if info.config is not None else None` - Merge two-line f-string in _log_final_consumer_stats - Drop redundant `except asyncio.TimeoutError: raise` in _ensure_consumer (now that the catch is narrowed to NotFoundError, TimeoutError propagates naturally — matches _ensure_stream style) - Explicit comment on the intentionally-broad `except Exception` in _log_final_consumer_stats clarifying it's different from _ensure_consumer Co-Authored-By: Claude Opus 4.6 (1M context) --- ami/ml/orchestration/nats_queue.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/ami/ml/orchestration/nats_queue.py b/ami/ml/orchestration/nats_queue.py index d39cf14fc..f65690ee5 100644 --- a/ami/ml/orchestration/nats_queue.py +++ b/ami/ml/orchestration/nats_queue.py @@ -115,7 +115,7 @@ def _format_consumer_config(info) -> str: ``add_consumer`` or ``consumer_info``, so the log always reflects what the server accepted rather than what we requested. """ - cfg = info.config if info.config is not None else None + cfg = info.config if cfg is None: return "config=?" @@ -275,13 +275,11 @@ async def _ensure_consumer(self, job_id: int): ) self._consumers_logged.add(job_id) return - except asyncio.TimeoutError: - raise # NATS unreachable — let caller handle it except nats.js.errors.NotFoundError: # Consumer doesn't exist, fall through to create it. Other - # JetStream errors (auth, API, transient) must propagate so we - # don't mask them as "missing consumer" and emit misleading - # creation logs. + # JetStream errors (auth, API, transient) and asyncio.TimeoutError + # propagate naturally — we don't want to mask them as "missing + # consumer" and emit misleading creation logs. pass info = await asyncio.wait_for( @@ -436,11 +434,15 @@ async def _log_final_consumer_stats(self, job_id: int) -> None: timeout=NATS_JETSTREAM_TIMEOUT, ) except Exception as e: + # Broad catch is intentional here (unlike _ensure_consumer): at + # cleanup time we tolerate any failure — stream gone, consumer + # already deleted, auth, timeout — so the delete calls below + # still get a chance to run. logger.debug(f"Could not fetch consumer info for {consumer_name} before deletion: {e}") return await self._log( logging.INFO, - f"Finalizing NATS consumer {consumer_name} before deletion " f"({self._format_consumer_stats(info)})", + f"Finalizing NATS consumer {consumer_name} before deletion ({self._format_consumer_stats(info)})", ) async def delete_consumer(self, job_id: int) -> bool: From ed7bf566550faf72f80b9eaf3c496837a6b4534a Mon Sep 17 00:00:00 2001 From: Michael Bunsen Date: Mon, 13 Apr 2026 15:34:02 -0700 Subject: [PATCH 11/14] refactor(nats): unify log fan-out behind manager.log_async MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rename TaskQueueManager._log → log_async and route every log line in the queue_images_to_nats async block through it (debug, error + traceback). Drops the ad-hoc sync_to_async(job.logger.error) bridge and the separate logger.exception call — one consistent API, one place that knows how to bridge JobLogHandler's ORM save through sync_to_async. log_async also now accepts exc_info=True so callers don't need to pair a module-only logger.exception with a job-logger error call. Co-Authored-By: Claude Opus 4.6 (1M context) --- ami/ml/orchestration/jobs.py | 29 +++++++++++++------------ ami/ml/orchestration/nats_queue.py | 34 ++++++++++++++++++------------ 2 files changed, 36 insertions(+), 27 deletions(-) diff --git a/ami/ml/orchestration/jobs.py b/ami/ml/orchestration/jobs.py index 8d5ec7911..79f787e28 100644 --- a/ami/ml/orchestration/jobs.py +++ b/ami/ml/orchestration/jobs.py @@ -1,6 +1,6 @@ import logging -from asgiref.sync import async_to_sync, sync_to_async +from asgiref.sync import async_to_sync from ami.jobs.models import Job, JobState from ami.main.models import SourceImage @@ -113,27 +113,28 @@ async def queue_all_images(): successful_queues = 0 failed_queues = 0 - # Pass job.logger so stream/consumer setup and any publish failures - # appear in the UI job log (not just the module logger). Per-image - # success logs stay at module level so a 10k-image job doesn't drown - # the job log. + # Pass job.logger so stream/consumer setup, per-image debug lines, and + # publish failures all appear in the UI job log (not just the module + # logger). All log calls inside this block go through manager.log_async + # so module + job logger stay in sync with one consistent API — and + # the sync_to_async bridge for JobLogHandler's ORM save lives in one + # place instead of being re-implemented at every call site. async with TaskQueueManager(job_logger=job.logger) as manager: for image_pk, task in tasks: try: - logger.debug(f"Queueing image {image_pk} to stream for job '{job.pk}': {task.image_url}") + await manager.log_async( + logging.DEBUG, + f"Queueing image {image_pk} to stream for job '{job.pk}': {task.image_url}", + ) success = await manager.publish_task( job_id=job.pk, data=task, ) except Exception as e: - # Module logger gets the full traceback for ops dashboards. - logger.exception("Failed to queue image %s to stream for job '%s'", image_pk, job.pk) - # job.logger.error triggers a sync Django ORM save inside - # JobLogHandler.emit, which raises SynchronousOnlyOperation - # when called directly from the event loop. Bridge it so - # the line actually lands in job.logs.stdout. - await sync_to_async(job.logger.error)( - f"Failed to queue image {image_pk} to stream for job '{job.pk}': {e}" + await manager.log_async( + logging.ERROR, + f"Failed to queue image {image_pk} to stream for job '{job.pk}': {e}", + exc_info=True, ) success = False diff --git a/ami/ml/orchestration/nats_queue.py b/ami/ml/orchestration/nats_queue.py index f65690ee5..5228b610f 100644 --- a/ami/ml/orchestration/nats_queue.py +++ b/ami/ml/orchestration/nats_queue.py @@ -86,9 +86,14 @@ def __init__( self._streams_logged: set[int] = set() self._consumers_logged: set[int] = set() - async def _log(self, level: int, msg: str) -> None: + async def log_async(self, level: int, msg: str, *, exc_info: bool = False) -> None: """Log to both the module logger and the job logger (if set). + Named ``log_async`` (not ``log``) to flag at every call site that this + is the async fan-out helper, distinct from stdlib ``Logger.log`` — + callers must ``await`` it. Use this from any async context where the + line should appear in both ops dashboards and the job's UI log. + Module logger fires synchronously (ops dashboards / stdout / New Relic are unaffected). The job logger call is bridged through ``sync_to_async`` because Django's ``JobLogHandler`` does an ORM @@ -97,13 +102,16 @@ async def _log(self, level: int, msg: str) -> None: line is silently dropped. The bridge offloads the handler work to a thread so the line actually lands in ``job.logs.stdout``. + Pass ``exc_info=True`` inside an ``except`` block to capture the + traceback on both loggers (same semantics as stdlib ``Logger.log``). + Exceptions from the job logger are swallowed so logging a lifecycle event never breaks the actual NATS operation. """ - logger.log(level, msg) + logger.log(level, msg, exc_info=exc_info) if self.job_logger is not None and self.job_logger is not logger: try: - await sync_to_async(self.job_logger.log)(level, msg) + await sync_to_async(self.job_logger.log)(level, msg, exc_info=exc_info) except Exception as e: logger.warning(f"Failed to mirror log to job logger: {e}") @@ -226,7 +234,7 @@ async def _ensure_stream(self, job_id: int): state = info.state messages = state.messages if state is not None else "?" last_seq = state.last_seq if state is not None else "?" - await self._log( + await self.log_async( logging.INFO, f"Reusing NATS stream {stream_name} (messages={messages}, last_seq={last_seq})", ) @@ -243,7 +251,7 @@ async def _ensure_stream(self, job_id: int): ), timeout=NATS_JETSTREAM_TIMEOUT, ) - await self._log(logging.INFO, f"Created NATS stream {stream_name}") + await self.log_async(logging.INFO, f"Created NATS stream {stream_name}") self._streams_logged.add(job_id) async def _ensure_consumer(self, job_id: int): @@ -269,7 +277,7 @@ async def _ensure_consumer(self, job_id: int): self.js.consumer_info(stream_name, consumer_name), timeout=NATS_JETSTREAM_TIMEOUT, ) - await self._log( + await self.log_async( logging.INFO, f"Reusing NATS consumer {consumer_name} ({self._format_consumer_stats(info)})", ) @@ -297,7 +305,7 @@ async def _ensure_consumer(self, job_id: int): ), timeout=NATS_JETSTREAM_TIMEOUT, ) - await self._log( + await self.log_async( logging.INFO, f"Created NATS consumer {consumer_name} ({self._format_consumer_config(info)})", ) @@ -336,7 +344,7 @@ async def publish_task(self, job_id: int, data: PipelineProcessingTask) -> bool: # Per-message success logs stay at module level (noise in 10k-image # jobs), but a failure on even a single publish deserves to surface # in the job log — otherwise the failure path is invisible to users. - await self._log(logging.ERROR, f"Failed to publish task to stream for job '{job_id}': {e}") + await self.log_async(logging.ERROR, f"Failed to publish task to stream for job '{job_id}': {e}") return False async def reserve_tasks(self, job_id: int, count: int, timeout: float = 5) -> list[PipelineProcessingTask]: @@ -440,7 +448,7 @@ async def _log_final_consumer_stats(self, job_id: int) -> None: # still get a chance to run. logger.debug(f"Could not fetch consumer info for {consumer_name} before deletion: {e}") return - await self._log( + await self.log_async( logging.INFO, f"Finalizing NATS consumer {consumer_name} before deletion ({self._format_consumer_stats(info)})", ) @@ -466,10 +474,10 @@ async def delete_consumer(self, job_id: int) -> bool: self.js.delete_consumer(stream_name, consumer_name), timeout=NATS_JETSTREAM_TIMEOUT, ) - await self._log(logging.INFO, f"Deleted NATS consumer {consumer_name} for job '{job_id}'") + await self.log_async(logging.INFO, f"Deleted NATS consumer {consumer_name} for job '{job_id}'") return True except Exception as e: - await self._log(logging.ERROR, f"Failed to delete NATS consumer for job '{job_id}': {e}") + await self.log_async(logging.ERROR, f"Failed to delete NATS consumer for job '{job_id}': {e}") return False async def delete_stream(self, job_id: int) -> bool: @@ -492,10 +500,10 @@ async def delete_stream(self, job_id: int) -> bool: self.js.delete_stream(stream_name), timeout=NATS_JETSTREAM_TIMEOUT, ) - await self._log(logging.INFO, f"Deleted NATS stream {stream_name} for job '{job_id}'") + await self.log_async(logging.INFO, f"Deleted NATS stream {stream_name} for job '{job_id}'") return True except Exception as e: - await self._log(logging.ERROR, f"Failed to delete NATS stream for job '{job_id}': {e}") + await self.log_async(logging.ERROR, f"Failed to delete NATS stream for job '{job_id}': {e}") return False async def _setup_advisory_stream(self): From ebfdbb7cb0c8169086754606f67af87ef1bbd9eb Mon Sep 17 00:00:00 2001 From: Michael Bunsen Date: Mon, 13 Apr 2026 16:02:19 -0700 Subject: [PATCH 12/14] docs(nats): note future intent to route, not mirror, lifecycle logs log_async currently mirrors granular per-job lifecycle to both the module and job loggers. Document why this is intentional for now (async ML processing still stabilizing, stdout visibility helping us debug) and the eventual target shape (route to job logger only at INFO/DEBUG, mirror at WARNING+). Matches the pattern in ami.jobs.tasks.save_results, where job.logger.propagate=False keeps granular per-job state out of ops logs. Co-Authored-By: Claude Opus 4.6 (1M context) --- ami/ml/orchestration/nats_queue.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/ami/ml/orchestration/nats_queue.py b/ami/ml/orchestration/nats_queue.py index 5228b610f..742ba7771 100644 --- a/ami/ml/orchestration/nats_queue.py +++ b/ami/ml/orchestration/nats_queue.py @@ -107,6 +107,20 @@ async def log_async(self, level: int, msg: str, *, exc_info: bool = False) -> No Exceptions from the job logger are swallowed so logging a lifecycle event never breaks the actual NATS operation. + + FUTURE: this currently mirrors granular per-job lifecycle (stream / + consumer create+reuse, per-image debug, forensic stats) to BOTH the + module logger and the job logger. The longer-term preference is to + route — granular lifecycle stays on ``job.logger`` only (matching + ``ami.jobs.tasks.save_results`` and friends, where ``job.logger`` has + ``propagate=False`` and never reaches stdout / NR), with the module + logger reserved for true ops signals (connection failures, NATS-side + errors). Kept symmetric for now because async ML processing is still + being stabilized and the extra stdout visibility is helping us + debug. Once we trust the per-job UI log as the canonical place to + inspect a job, switch ``log_async`` to route-not-mirror at INFO/DEBUG + and only auto-mirror at WARNING+ (so true error signals still always + reach ops dashboards). """ logger.log(level, msg, exc_info=exc_info) if self.job_logger is not None and self.job_logger is not logger: From b5e9a875c9af93040ad9a820ec85eb1560d2947e Mon Sep 17 00:00:00 2001 From: Michael Bunsen Date: Mon, 13 Apr 2026 18:19:30 -0700 Subject: [PATCH 13/14] docs(nats): correct safety claim for _ensure_stream/_ensure_consumer dedup MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The original docstring claimed the stream/consumer won't be deleted mid-flight because cleanup uses a separate manager session. That's incomplete — Job.cancel() runs cleanup_async_job_resources in the request thread while queue_images_to_nats is still in its publish loop in the Celery worker. So a concurrent delete across manager sessions is possible. The early-return is still safe in that scenario, but for a different reason than the original claim: downstream publish_task fails loudly (returns False, logs ERROR) when the stream is gone, rather than silently recreating an orphan stream without a consumer (which is what the non-deduped baseline would do). Updates _ensure_stream and _ensure_consumer docstrings to describe the actual safety argument accurately. Co-Authored-By: Claude Opus 4.6 (1M context) --- ami/ml/orchestration/nats_queue.py | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/ami/ml/orchestration/nats_queue.py b/ami/ml/orchestration/nats_queue.py index 742ba7771..3c36f1ab5 100644 --- a/ami/ml/orchestration/nats_queue.py +++ b/ami/ml/orchestration/nats_queue.py @@ -232,8 +232,17 @@ async def _ensure_stream(self, job_id: int): Logs a lifecycle line to both the module and job logger the first time it sees a given job in this manager session (creation or reuse). Subsequent - calls in the same session skip the NATS round-trip entirely — the stream - won't be deleted mid-flight (cleanup uses a separate manager session). + calls in the same session skip the NATS round-trip entirely via the + ``_streams_logged`` set. + + Concurrency note: ``Job.cancel()`` can trigger ``cleanup_async_job_resources`` + in the request thread while this manager is still in its publish loop in + the Celery worker, so the stream *can* be deleted mid-flight from a + different manager session. The early-return is still safe in that case — + subsequent ``publish_task`` calls will fail loudly (``self.js.publish`` + returns an error, caught and logged by ``publish_task``) rather than + silently recreating the stream without a consumer. Failing loud on a + cancel race is the correct behavior. """ if job_id in self._streams_logged: return @@ -275,7 +284,13 @@ async def _ensure_consumer(self, job_id: int): to both the module and job logger. On creation the line includes the config snapshot (max_deliver, ack_wait, max_ack_pending, deliver_policy, ack_policy) so forensic readers can see exactly what delivery semantics - were in effect. Subsequent calls skip the NATS round-trip. + were in effect. Subsequent calls skip the NATS round-trip via the + ``_consumers_logged`` set. + + Same concurrency caveat as ``_ensure_stream``: a concurrent cancel can + delete the consumer mid-flight. The early-return stays safe because + downstream ``publish_task`` fails loudly rather than silently recreating + an orphan consumer. """ if job_id in self._consumers_logged: return From 766137c1504cc294ba5e4896155d0b12756aae5e Mon Sep 17 00:00:00 2001 From: Michael Bunsen Date: Mon, 13 Apr 2026 18:25:51 -0700 Subject: [PATCH 14/14] perf(nats): gate log_async on isEnabledFor before sync_to_async mirror MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Without this gate, every log_async call fires the job-logger mirror through sync_to_async — a ThreadPoolExecutor submit per call — regardless of whether the effective level would drop the record. For a 10k-image queue this amounts to 10k unnecessary thread-pool submissions when DEBUG is off. stdlib Logger.log does the same isEnabledFor check internally before formatting. We need to do it explicitly here because the mirror goes through sync_to_async, bypassing the in-logger short-circuit. No behavior change when at least one logger is enabled for the level; pure short-circuit when both are gated out. Co-Authored-By: Claude Opus 4.6 (1M context) --- ami/ml/orchestration/nats_queue.py | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/ami/ml/orchestration/nats_queue.py b/ami/ml/orchestration/nats_queue.py index 3c36f1ab5..43d9d65e5 100644 --- a/ami/ml/orchestration/nats_queue.py +++ b/ami/ml/orchestration/nats_queue.py @@ -108,6 +108,15 @@ async def log_async(self, level: int, msg: str, *, exc_info: bool = False) -> No Exceptions from the job logger are swallowed so logging a lifecycle event never breaks the actual NATS operation. + Gated by ``isEnabledFor`` up front so a disabled level returns + immediately without paying for the ``sync_to_async`` round-trip. + Matters most at DEBUG during large queues — stdlib ``Logger.log`` + does the same level check internally before formatting a message; + we have to do it explicitly here because the job-logger mirror + happens through ``sync_to_async`` (ThreadPoolExecutor submit), which + would otherwise fire once per image even when the handler is about + to drop the record. + FUTURE: this currently mirrors granular per-job lifecycle (stream / consumer create+reuse, per-image debug, forensic stats) to BOTH the module logger and the job logger. The longer-term preference is to @@ -122,8 +131,15 @@ async def log_async(self, level: int, msg: str, *, exc_info: bool = False) -> No and only auto-mirror at WARNING+ (so true error signals still always reach ops dashboards). """ - logger.log(level, msg, exc_info=exc_info) - if self.job_logger is not None and self.job_logger is not logger: + module_enabled = logger.isEnabledFor(level) + job_enabled = ( + self.job_logger is not None and self.job_logger is not logger and self.job_logger.isEnabledFor(level) + ) + if not module_enabled and not job_enabled: + return + if module_enabled: + logger.log(level, msg, exc_info=exc_info) + if job_enabled: try: await sync_to_async(self.job_logger.log)(level, msg, exc_info=exc_info) except Exception as e: