From 4c089c2472ad789e205813e5c3e7614ad1d90b08 Mon Sep 17 00:00:00 2001 From: Michael Bunsen Date: Wed, 27 May 2026 09:19:09 -0700 Subject: [PATCH 01/11] perf(ml): bulk-query filter_processed_images to fix Collect-stage stall filter_processed_images() iterated images one at a time, issuing 3-5 ORM round trips per image against Detection / Classification. On a ~147k-image collection that produces ~500k single-row queries inside one Celery task, which silences the AMQP heartbeat long enough that the reaper flips the job from STARTED to REVOKED before any forward progress is reported. Rewrite to chunk the input via itertools.islice and run at most two bulk queries per chunk: one over Detection rows, one over Classification rows for the real detections in that chunk. Semantics for all five branches (no detections, null-only, unclassified real, partial classifier coverage, fully classified) are preserved. Adds tests covering empty input, mixed batches across all branches, an assertion that query count scales with batch count rather than image count, and ordering preservation across batch boundaries. Fixes #1321. Co-Authored-By: Claude --- ami/ml/models/pipeline.py | 111 +++++++++++++++++++++++++------------- ami/ml/tests.py | 89 ++++++++++++++++++++++++++++++ 2 files changed, 164 insertions(+), 36 deletions(-) diff --git a/ami/ml/models/pipeline.py b/ami/ml/models/pipeline.py index c259e4aea..415127846 100644 --- a/ami/ml/models/pipeline.py +++ b/ami/ml/models/pipeline.py @@ -8,6 +8,7 @@ import collections import dataclasses +import itertools import logging import time import typing @@ -23,7 +24,6 @@ from ami.base.models import BaseModel, BaseQuerySet from ami.base.schemas import ConfigurableStage, default_stages from ami.main.models import ( - NULL_DETECTIONS_FILTER, Classification, Deployment, Detection, @@ -57,10 +57,14 @@ logger = logging.getLogger(__name__) +FILTER_PROCESSED_BATCH_SIZE = 1000 + + def filter_processed_images( images: typing.Iterable[SourceImage], pipeline: Pipeline, task_logger: logging.Logger = logger, + batch_size: int = FILTER_PROCESSED_BATCH_SIZE, ) -> typing.Iterable[SourceImage]: """ Return only images that need to be processed by a given pipeline. @@ -79,56 +83,91 @@ def filter_processed_images( Null detections are sentinels that mark an image as "processed, nothing found." They are excluded from classification checks so they don't trigger reprocessing. + + Input images are processed in batches so the work scales as O(image_count / + batch_size) database round-trips instead of O(image_count). This keeps the + Collect stage well under the Celery broker's AMQP heartbeat window for the + large-collection case (see issue #1321). """ - pipeline_algorithms = pipeline.algorithms.all() + pipeline_algorithms = list(pipeline.algorithms.all()) + pipeline_algorithm_ids = [a.id for a in pipeline_algorithms] - detection_type_keys = Algorithm.detection_task_types - detection_algorithms = pipeline_algorithms.filter(task_type__in=detection_type_keys) - if not detection_algorithms.exists(): + detection_type_keys = set(Algorithm.detection_task_types) + has_detection_algorithm = any(a.task_type in detection_type_keys for a in pipeline_algorithms) + if not has_detection_algorithm: task_logger.warning(f"Pipeline {pipeline} has no detection algorithms saved. Will reprocess all images.") - classification_algorithms = pipeline_algorithms.exclude(task_type__in=detection_type_keys) - if not classification_algorithms.exists(): + pipeline_classifier_ids = {a.id for a in pipeline_algorithms if a.task_type not in detection_type_keys} + if not pipeline_classifier_ids: task_logger.warning(f"Pipeline {pipeline} has no classification algorithms saved. Will reprocess all images.") - for image in images: - existing_detections = image.detections.filter(detection_algorithm__in=pipeline_algorithms) - if not existing_detections.exists(): - task_logger.debug(f"Image {image} needs processing: has no existing detections from pipeline's detector") - # If there are no existing detections from this pipeline, send the image - yield image - elif not existing_detections.exclude(NULL_DETECTIONS_FILTER).exists(): # type: ignore - # All detections for this image are null (processed but nothing found) — skip - task_logger.debug(f"Image {image} has only null detections from pipeline {pipeline}, skipping!") - continue - elif existing_detections.exclude(NULL_DETECTIONS_FILTER).filter(classifications__isnull=True).exists(): - # Check if any real detections (non-null) have no classifications - task_logger.debug( - f"Image {image} needs processing: has existing detections with no classifications " - "from pipeline {pipeline}" - ) - yield image - else: - # If there are existing detections with classifications, - # Compare their classification algorithms to the current pipeline's algorithms - pipeline_algorithm_ids = set(classification_algorithms.values_list("id", flat=True)) - detection_algorithm_ids = set(existing_detections.values_list("classifications__algorithm_id", flat=True)) + image_iter = iter(images) + while True: + batch = list(itertools.islice(image_iter, batch_size)) + if not batch: + return + + batch_ids = [image.pk for image in batch] + + images_with_any_detection: set[int] = set() + real_detections_per_image: dict[int, set[int]] = collections.defaultdict(set) + detection_rows = Detection.objects.filter( + source_image_id__in=batch_ids, + detection_algorithm_id__in=pipeline_algorithm_ids, + ).values_list("source_image_id", "id", "bbox") + for source_image_id, detection_id, bbox in detection_rows: + images_with_any_detection.add(source_image_id) + if bbox not in (None, []): + real_detections_per_image[source_image_id].add(detection_id) + + classifier_ids_per_detection: dict[int, set[int]] = collections.defaultdict(set) + real_detection_ids = [d for dets in real_detections_per_image.values() for d in dets] + if real_detection_ids: + classification_rows = Classification.objects.filter( + detection_id__in=real_detection_ids, + ).values_list("detection_id", "algorithm_id") + for detection_id, algorithm_id in classification_rows: + classifier_ids_per_detection[detection_id].add(algorithm_id) + + for image in batch: + image_id = image.pk + + if image_id not in images_with_any_detection: + task_logger.debug( + f"Image {image} needs processing: has no existing detections from pipeline's detector" + ) + yield image + continue - if not pipeline_algorithm_ids.issubset(detection_algorithm_ids): + real_det_ids = real_detections_per_image.get(image_id) + if not real_det_ids: + task_logger.debug(f"Image {image} has only null detections from pipeline {pipeline}, skipping!") + continue + + if any(detection_id not in classifier_ids_per_detection for detection_id in real_det_ids): task_logger.debug( - f"Image {image} has existing detections that haven't been classified by the pipeline: {pipeline}:" - f" {detection_algorithm_ids} vs {pipeline_algorithm_ids}" + f"Image {image} needs processing: has existing detections with no classifications " + f"from pipeline {pipeline}" + ) + yield image + continue + + observed_classifier_ids: set[int] = set() + for detection_id in real_det_ids: + observed_classifier_ids.update(classifier_ids_per_detection[detection_id]) + + if not pipeline_classifier_ids.issubset(observed_classifier_ids): + missing_algos = pipeline_classifier_ids - observed_classifier_ids + task_logger.debug( + f"Image {image} has existing detections that haven't been classified by the pipeline: {pipeline}: " + f"{observed_classifier_ids} vs {pipeline_classifier_ids}. " f"Since we do yet have a mechanism to reclassify detections, processing the image from scratch." ) - # log all algorithms that are in the pipeline but not in the detection - missing_algos = pipeline_algorithm_ids - detection_algorithm_ids task_logger.debug(f"Image #{image.pk} needs classification by pipeline's algorithms: {missing_algos}") yield image else: - # If all detections have been classified by the pipeline, skip the image task_logger.debug( f"Image {image} has existing detections classified by the pipeline: {pipeline}, skipping!" ) - continue def collect_images( diff --git a/ami/ml/tests.py b/ami/ml/tests.py index 0eaa0a237..3148116f2 100644 --- a/ami/ml/tests.py +++ b/ami/ml/tests.py @@ -962,6 +962,95 @@ def test_filter_processed_images_skips_null_and_fully_classified(self): result = list(filter_processed_images([image], self.pipeline)) self.assertEqual(result, [], "Fully classified image with null detection should be skipped") + def test_filter_processed_images_empty_input(self): + """An empty iterable should yield nothing and run no per-image queries.""" + from ami.ml.models.pipeline import filter_processed_images + + with self.assertNumQueries(1): # one query: pipeline.algorithms.all() + result = list(filter_processed_images([], self.pipeline)) + self.assertEqual(result, []) + + def test_filter_processed_images_mixed_batch(self): + """ + A mixed batch of images covering all five branches should yield only + the ones that need processing, in input order. + """ + from ami.ml.models.pipeline import filter_processed_images + + detector = self.algorithms["random-detector"] + binary = self.algorithms["random-binary-classifier"] + species = self.algorithms["random-species-classifier"] + + unprocessed = SourceImage.objects.create(path="unprocessed.jpg") + null_only = SourceImage.objects.create(path="null_only.jpg") + unclassified = SourceImage.objects.create(path="unclassified.jpg") + fully_classified = SourceImage.objects.create(path="fully_classified.jpg") + + Detection.objects.create(source_image=null_only, detection_algorithm=detector, bbox=None) + + Detection.objects.create(source_image=unclassified, detection_algorithm=detector, bbox=[0.1, 0.2, 0.3, 0.4]) + + real_det = Detection.objects.create( + source_image=fully_classified, detection_algorithm=detector, bbox=[0.1, 0.2, 0.3, 0.4] + ) + taxon = Taxon.objects.create(name="Test Mixed Batch Taxon") + Classification.objects.create( + detection=real_det, taxon=taxon, algorithm=binary, score=0.9, timestamp=datetime.datetime.now() + ) + Classification.objects.create( + detection=real_det, taxon=taxon, algorithm=species, score=0.8, timestamp=datetime.datetime.now() + ) + + images = [unprocessed, null_only, unclassified, fully_classified] + result = list(filter_processed_images(images, self.pipeline)) + self.assertEqual(result, [unprocessed, unclassified]) + + def test_filter_processed_images_query_count_is_bounded_per_batch(self): + """ + With N images and batch_size=B, the query count should scale as + O(N / B), not O(N). Locks in the bulk-query rewrite from issue #1321. + + Each batch issues at most: + - 1 Detection bulk select + - 1 Classification bulk select (only when real detections exist) + Plus one initial pipeline.algorithms.all() that's shared across batches. + """ + from ami.ml.models.pipeline import filter_processed_images + + detector = self.algorithms["random-detector"] + binary = self.algorithms["random-binary-classifier"] + species = self.algorithms["random-species-classifier"] + taxon = Taxon.objects.create(name="Bounded Query Test Taxon") + + # 10 images. Batch 1: 5 fully-classified (triggers classification query). + # Batch 2: 5 unprocessed (no detections, classification query skipped). + images = [SourceImage.objects.create(path=f"bulk-{i}.jpg") for i in range(10)] + for image in images[:5]: + real_det = Detection.objects.create( + source_image=image, detection_algorithm=detector, bbox=[0.1, 0.2, 0.3, 0.4] + ) + Classification.objects.create( + detection=real_det, taxon=taxon, algorithm=binary, score=0.9, timestamp=datetime.datetime.now() + ) + Classification.objects.create( + detection=real_det, taxon=taxon, algorithm=species, score=0.8, timestamp=datetime.datetime.now() + ) + + # Expected: 1 (pipeline) + 2 (detection × 2 batches) + 1 (classification, batch 1 only) = 4. + with self.assertNumQueries(4): + result = list(filter_processed_images(images, self.pipeline, batch_size=5)) + + # First 5 fully classified → skipped. Last 5 fresh → yielded. + self.assertEqual(result, images[5:]) + + def test_filter_processed_images_preserves_input_order_across_batches(self): + """Output should preserve input ordering even when batched.""" + from ami.ml.models.pipeline import filter_processed_images + + images = [SourceImage.objects.create(path=f"order-{i}.jpg") for i in range(7)] + result = list(filter_processed_images(images, self.pipeline, batch_size=3)) + self.assertEqual(result, images) + def test_null_detections_are_algorithm_specific(self): """ Null detections from different pipelines/algorithms should not be shared. From 25c3be30fdc3305bdf4b550e36f86b6b9e7611ee Mon Sep 17 00:00:00 2001 From: Michael Bunsen Date: Wed, 27 May 2026 09:30:37 -0700 Subject: [PATCH 02/11] fix(ml): strip default ordering, dedupe null-bbox check, rename batch dicts Takeaway-review followups on the bulk filter_processed_images rewrite: - Strip Detection's Meta.ordering on the bulk select via .order_by(). The per-batch query was sorting all matching rows on (frame_num, timestamp) before returning, wasted work since results land in a dict. - Extract bbox_is_null() helper in ami.main.models next to NULL_DETECTIONS_FILTER so the SQL filter and the in-memory check share a single source of truth. - Rename images_with_any_detection -> images_with_pipeline_detection and real_detections_per_image -> real_pipeline_detections_per_image to make clear the dicts only cover detections from this pipeline's algorithms. Co-Authored-By: Claude --- ami/main/models.py | 5 +++++ ami/ml/models/pipeline.py | 39 +++++++++++++++++++++++---------------- 2 files changed, 28 insertions(+), 16 deletions(-) diff --git a/ami/main/models.py b/ami/main/models.py index 3d21c07cd..1de9a01ee 100644 --- a/ami/main/models.py +++ b/ami/main/models.py @@ -98,6 +98,11 @@ class TaxonRank(OrderedEnum): NULL_DETECTIONS_FILTER = Q(bbox__isnull=True) | Q(bbox=[]) +def bbox_is_null(bbox) -> bool: + """In-memory equivalent of NULL_DETECTIONS_FILTER for already-fetched bbox values.""" + return bbox is None or bbox == [] + + def get_media_url(path: str) -> str: """ If path is a full URL, return it as-is. diff --git a/ami/ml/models/pipeline.py b/ami/ml/models/pipeline.py index 415127846..f4a675a45 100644 --- a/ami/ml/models/pipeline.py +++ b/ami/ml/models/pipeline.py @@ -34,6 +34,7 @@ TaxaList, Taxon, TaxonRank, + bbox_is_null, update_calculated_fields_for_events, update_occurrence_determination, ) @@ -108,37 +109,43 @@ def filter_processed_images( batch_ids = [image.pk for image in batch] - images_with_any_detection: set[int] = set() - real_detections_per_image: dict[int, set[int]] = collections.defaultdict(set) - detection_rows = Detection.objects.filter( - source_image_id__in=batch_ids, - detection_algorithm_id__in=pipeline_algorithm_ids, - ).values_list("source_image_id", "id", "bbox") + images_with_pipeline_detection: set[int] = set() + real_pipeline_detections_per_image: dict[int, set[int]] = collections.defaultdict(set) + detection_rows = ( + Detection.objects.filter( + source_image_id__in=batch_ids, + detection_algorithm_id__in=pipeline_algorithm_ids, + ) + .order_by() + .values_list("source_image_id", "id", "bbox") + ) for source_image_id, detection_id, bbox in detection_rows: - images_with_any_detection.add(source_image_id) - if bbox not in (None, []): - real_detections_per_image[source_image_id].add(detection_id) + images_with_pipeline_detection.add(source_image_id) + if not bbox_is_null(bbox): + real_pipeline_detections_per_image[source_image_id].add(detection_id) classifier_ids_per_detection: dict[int, set[int]] = collections.defaultdict(set) - real_detection_ids = [d for dets in real_detections_per_image.values() for d in dets] - if real_detection_ids: - classification_rows = Classification.objects.filter( - detection_id__in=real_detection_ids, - ).values_list("detection_id", "algorithm_id") + real_pipeline_detection_ids = [d for dets in real_pipeline_detections_per_image.values() for d in dets] + if real_pipeline_detection_ids: + classification_rows = ( + Classification.objects.filter(detection_id__in=real_pipeline_detection_ids) + .order_by() + .values_list("detection_id", "algorithm_id") + ) for detection_id, algorithm_id in classification_rows: classifier_ids_per_detection[detection_id].add(algorithm_id) for image in batch: image_id = image.pk - if image_id not in images_with_any_detection: + if image_id not in images_with_pipeline_detection: task_logger.debug( f"Image {image} needs processing: has no existing detections from pipeline's detector" ) yield image continue - real_det_ids = real_detections_per_image.get(image_id) + real_det_ids = real_pipeline_detections_per_image.get(image_id) if not real_det_ids: task_logger.debug(f"Image {image} has only null detections from pipeline {pipeline}, skipping!") continue From 2f6ed070723535ea86ab91df3073292f68263919 Mon Sep 17 00:00:00 2001 From: Michael Bunsen Date: Wed, 27 May 2026 13:57:46 -0700 Subject: [PATCH 03/11] perf(ml): prefetch deployment + data_source in collect_images MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit collect_images() returned a SourceImage queryset without select_related on the deployment/data_source FK chain. Downstream queue_images_to_nats() calls image.url() per image, which traverses self.deployment.data_source — two FK lookups per call. On a ~97k-image collection that's ~4.6ms × 97k = ~7 minutes of pure ORM round-trips before NATS publishes start, well inside the 10-minute reaper window (see issue #1321 follow-up comment). Add .select_related("deployment__data_source") to both queryset branches in collect_images so the joins ride along on the initial SourceImage fetch and image.url() in the queue loop stays a pure-Python operation. Test asserts that accessing .deployment.data_source on the returned rows triggers zero additional queries. Co-Authored-By: Claude --- ami/ml/models/pipeline.py | 8 +++++--- ami/ml/tests.py | 37 +++++++++++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 3 deletions(-) diff --git a/ami/ml/models/pipeline.py b/ami/ml/models/pipeline.py index f4a675a45..8d992833c 100644 --- a/ami/ml/models/pipeline.py +++ b/ami/ml/models/pipeline.py @@ -197,13 +197,15 @@ def collect_images( else: job = None - # Set source to first argument that is not None + # Set source to first argument that is not None. Always prefetch the + # deployment + data_source joins so later calls to image.url() in + # queue_images_to_nats don't trigger N+1 lookups (see issue #1321). if collection: - images = collection.images.all() + images = collection.images.select_related("deployment__data_source") elif source_images: images = source_images elif deployment: - images = SourceImage.objects.filter(deployment=deployment) + images = SourceImage.objects.filter(deployment=deployment).select_related("deployment__data_source") else: raise ValueError("Must specify a collection, deployment or a list of images") diff --git a/ami/ml/tests.py b/ami/ml/tests.py index 3148116f2..f1caba0e6 100644 --- a/ami/ml/tests.py +++ b/ami/ml/tests.py @@ -554,6 +554,43 @@ def test_collect_images(self): images = list(collect_images(collection=self.image_collection, pipeline=self.pipeline)) assert len(images) == 2 + def test_collect_images_prefetches_deployment_and_data_source(self): + """ + collect_images() must hand back SourceImage rows with deployment and + deployment.data_source already joined, so the downstream queue_images_to_nats + loop doesn't trigger N+1 FK lookups inside image.url() (see issue #1321). + """ + from ami.main.models import S3StorageSource + + data_source = S3StorageSource.objects.create( + name="ds-prefetch-test", + bucket="prefetch-bucket", + access_key="x", + secret_key="y", + public_base_url="https://example.invalid/", + project=self.project, + ) + deployment = Deployment.objects.create( + name="prefetch-deployment", project=self.project, data_source=data_source + ) + images = [ + SourceImage.objects.create(path=f"prefetch-{i}.jpg", deployment=deployment, project=self.project) + for i in range(3) + ] + collection = SourceImageCollection.objects.create(project=self.project, name="prefetch-collection") + collection.images.set(images) + + collected = list(collect_images(collection=collection, pipeline=self.pipeline)) + self.assertEqual(len(collected), 3) + + # Accessing deployment.data_source on the returned images should + # require zero extra queries because select_related joined both. + with self.assertNumQueries(0): + for image in collected: + self.assertEqual(image.deployment_id, deployment.pk) + self.assertEqual(image.deployment.data_source_id, data_source.pk) + self.assertEqual(image.deployment.data_source.public_base_url, "https://example.invalid/") + def fake_pipeline_results( self, source_images: list[SourceImage], From 528247e357392df34f7b95582b164485d6a70879 Mon Sep 17 00:00:00 2001 From: Michael Bunsen Date: Wed, 27 May 2026 13:57:54 -0700 Subject: [PATCH 04/11] perf(ml): call image.url() once per image in queue_images_to_nats MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The pre-NATS prep loop in queue_images_to_nats() called image.url() twice per iteration — once in the truthy conditional, once for assignment. image.url() goes through SourceImage.public_url() which touches deployment and data_source; doubling it doubles the FK chain cost. On ~97k-image collections the redundant call alone adds ~7 minutes (see issue #1321 follow-up). Cache the call in a local once and drop the now-redundant hasattr/url double-check. The companion select_related in collect_images keeps the single remaining call cheap. Co-Authored-By: Claude --- ami/ml/orchestration/jobs.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/ami/ml/orchestration/jobs.py b/ami/ml/orchestration/jobs.py index 79f787e28..47fee8c95 100644 --- a/ami/ml/orchestration/jobs.py +++ b/ami/ml/orchestration/jobs.py @@ -91,7 +91,11 @@ def queue_images_to_nats(job: "Job", images: list[SourceImage]): skipped_count = 0 for image in images: image_id = str(image.pk) - image_url = image.url() if hasattr(image, "url") and image.url() else "" + # Call image.url() exactly once per iteration — the implementation + # touches deployment + data_source and the call cost adds up across + # large collections. The upstream queryset in collect_images() also + # prefetches those joins so this stays cheap (see issue #1321). + image_url = image.url() if hasattr(image, "url") else None if not image_url: job.logger.warning(f"Image {image.pk} has no URL, skipping queuing to NATS for job '{job.pk}'") skipped_count += 1 From 59d0561932c32ac5ef30d887f2d0a1605087e288 Mon Sep 17 00:00:00 2001 From: Michael Bunsen Date: Wed, 27 May 2026 16:07:25 -0700 Subject: [PATCH 05/11] perf(ml): fanout NATS publishes in queue_images_to_nats MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Each publish_task awaits a JetStream ack round-trip (~1.3ms measured on Serbia 2026-05-27 for the 105k-image incident-shape collection). Sequential awaits stack linearly, so 105k images took 139s and 700k would push past the 10-min reaper threshold on its own. Switch the inner loop to chunked asyncio.gather batches of 200 so the client can pipeline ack roundtrips. _ensure_stream / _ensure_consumer are hoisted out of the loop and pre-warmed once on entry — they were already cache-skipped after first call (nats_queue.py:310,358) but each call still serialised inside its publish coroutine. JetStream still assigns ack.seq atomically server-side, so messages remain ordered in the stream by arrival; only the order *within* a fanout chunk is non-deterministic from the caller's POV. ADC consumers fetch in seq order and treat each image-task independently — no per-image ordering dependency. Adds 5 unit tests in ami/ml/orchestration/tests/test_jobs.py mocking TaskQueueManager: warm-up count, chunk-boundary publish count, partial publish failure path, stream-setup failure short-circuit, and a sentinel test locking the chunk constant. Co-Authored-By: Claude --- ami/ml/orchestration/jobs.py | 55 +++++++--- ami/ml/orchestration/tests/test_jobs.py | 135 ++++++++++++++++++++++++ 2 files changed, 175 insertions(+), 15 deletions(-) create mode 100644 ami/ml/orchestration/tests/test_jobs.py diff --git a/ami/ml/orchestration/jobs.py b/ami/ml/orchestration/jobs.py index 47fee8c95..868295245 100644 --- a/ami/ml/orchestration/jobs.py +++ b/ami/ml/orchestration/jobs.py @@ -1,3 +1,4 @@ +import asyncio import logging from asgiref.sync import async_to_sync @@ -10,6 +11,14 @@ logger = logging.getLogger(__name__) +# Number of concurrent JetStream publishes per fanout chunk. The bottleneck on +# large-collection jobs is the per-message ack round-trip (~1.3ms each on +# Serbia 2026-05-27, sequential), so awaiting publishes one at a time scales +# linearly with image count and pushes >450k-image jobs past the reaper +# threshold. Issuing ~200 publishes per gather() lets NATS pipeline the acks +# back to us; the chunk boundary keeps memory and concurrent-task counts bounded. +NATS_PUBLISH_FANOUT_CHUNK_SIZE = 200 + def cleanup_async_job_resources(job_id: int) -> bool: """ @@ -124,28 +133,44 @@ async def queue_all_images(): # 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: + # Warm the stream + consumer caches once so per-publish calls skip + # the cached-noop branches in publish_task -> _ensure_stream / + # _ensure_consumer. Even though those branches are O(1) after the + # first call, each one still runs inside the publish coroutine and + # serialises with the gather below. + try: + await manager._ensure_stream(job.pk) + await manager._ensure_consumer(job.pk) + except Exception as e: + await manager.log_async( + logging.ERROR, + f"Failed to set up NATS stream/consumer for job '{job.pk}': {e}", + exc_info=True, + ) + return 0, len(tasks) + + async def publish_one(image_pk: int, task: PipelineProcessingTask) -> bool: try: - 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, - ) + return await manager.publish_task(job_id=job.pk, data=task) except Exception as 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 - - if success: - successful_queues += 1 - else: - failed_queues += 1 + return False + + for chunk_start in range(0, len(tasks), NATS_PUBLISH_FANOUT_CHUNK_SIZE): + chunk = tasks[chunk_start : chunk_start + NATS_PUBLISH_FANOUT_CHUNK_SIZE] + results = await asyncio.gather( + *(publish_one(image_pk, task) for image_pk, task in chunk), + return_exceptions=False, + ) + for success in results: + if success: + successful_queues += 1 + else: + failed_queues += 1 return successful_queues, failed_queues diff --git a/ami/ml/orchestration/tests/test_jobs.py b/ami/ml/orchestration/tests/test_jobs.py new file mode 100644 index 000000000..9d14698f9 --- /dev/null +++ b/ami/ml/orchestration/tests/test_jobs.py @@ -0,0 +1,135 @@ +from unittest.mock import AsyncMock, patch + +from django.test import TestCase + +from ami.jobs.models import Job, JobDispatchMode, MLJob +from ami.main.models import Deployment, Project, S3StorageSource, SourceImage, SourceImageCollection +from ami.ml.models import Pipeline +from ami.ml.orchestration import jobs as orchestration_jobs +from ami.ml.orchestration.jobs import NATS_PUBLISH_FANOUT_CHUNK_SIZE, queue_images_to_nats + + +def _stub_manager(mock_manager_cls, publish_results=None): + """Return an AsyncMock TaskQueueManager instance with publish_task pre-stubbed. + + publish_results is an optional iterable of bools matching publish_task call + order. Defaults to always-True. + """ + instance = mock_manager_cls.return_value + instance.__aenter__ = AsyncMock(return_value=instance) + instance.__aexit__ = AsyncMock(return_value=False) + instance._ensure_stream = AsyncMock() + instance._ensure_consumer = AsyncMock() + instance.log_async = AsyncMock() + if publish_results is None: + instance.publish_task = AsyncMock(return_value=True) + else: + instance.publish_task = AsyncMock(side_effect=list(publish_results)) + return instance + + +@patch("ami.ml.orchestration.jobs.AsyncJobStateManager") +@patch("ami.ml.orchestration.jobs.TaskQueueManager") +class TestQueueImagesToNatsFanout(TestCase): + """Unit-test the chunk+gather behaviour of queue_images_to_nats. + + The integration test in test_cleanup.py uses a real NATS server. These + tests instead mock TaskQueueManager so they run without infrastructure and + can assert on call counts. + """ + + def setUp(self): + self.project = Project.objects.create(name="fanout test") + self.data_source = S3StorageSource.objects.create( + name="ds", + bucket="b", + access_key="x", + secret_key="y", + public_base_url="https://example.invalid/", + project=self.project, + ) + self.deployment = Deployment.objects.create(name="d", project=self.project, data_source=self.data_source) + self.pipeline = Pipeline.objects.create(name="p", slug="p", version="1") + self.collection = SourceImageCollection.objects.create(name="c", project=self.project) + + def _make_images(self, n: int) -> list[SourceImage]: + # public_base_url is set so SourceImage.url() returns a truthy URL — + # otherwise queue_images_to_nats skips the image and the inner loop + # we're trying to test never runs. + rows = [ + SourceImage( + path=f"fanout/{i}.jpg", + deployment=self.deployment, + project=self.project, + public_base_url="https://example.invalid/", + ) + for i in range(n) + ] + SourceImage.objects.bulk_create(rows) + return list(SourceImage.objects.filter(project=self.project).order_by("pk")) + + def _make_job(self) -> Job: + return Job.objects.create( + name="fanout", + job_type_key=MLJob.key, + project=self.project, + pipeline=self.pipeline, + source_image_collection=self.collection, + dispatch_mode=JobDispatchMode.ASYNC_API, + ) + + def test_warms_stream_and_consumer_once_before_publishing(self, mock_mgr_cls, _state_cls): + instance = _stub_manager(mock_mgr_cls) + images = self._make_images(5) + job = self._make_job() + + queue_images_to_nats(job, images) + + self.assertEqual(instance._ensure_stream.await_count, 1) + self.assertEqual(instance._ensure_consumer.await_count, 1) + self.assertEqual(instance.publish_task.await_count, 5) + + def test_chunks_across_fanout_boundary(self, mock_mgr_cls, _state_cls): + instance = _stub_manager(mock_mgr_cls) + # One full chunk + one short chunk: forces the inner loop to iterate twice. + n = NATS_PUBLISH_FANOUT_CHUNK_SIZE + 7 + images = self._make_images(n) + job = self._make_job() + + queue_images_to_nats(job, images) + + self.assertEqual(instance.publish_task.await_count, n) + # Warm-up still only runs once even when we span multiple chunks. + self.assertEqual(instance._ensure_stream.await_count, 1) + self.assertEqual(instance._ensure_consumer.await_count, 1) + + def test_partial_failure_counted_as_failed_not_raised(self, mock_mgr_cls, _state_cls): + # 3 images, middle one fails. Method should return False (failure summary) + # but not raise — and the success count should reflect the two that worked. + instance = _stub_manager(mock_mgr_cls, publish_results=[True, False, True]) + images = self._make_images(3) + job = self._make_job() + + result = queue_images_to_nats(job, images) + + self.assertFalse(result) + self.assertEqual(instance.publish_task.await_count, 3) + + def test_stream_setup_failure_short_circuits_publishing(self, mock_mgr_cls, _state_cls): + # If _ensure_stream raises, we should not attempt any publish_task calls — + # the whole batch is marked failed in one shot. + instance = _stub_manager(mock_mgr_cls) + instance._ensure_stream = AsyncMock(side_effect=RuntimeError("nats down")) + images = self._make_images(4) + job = self._make_job() + + result = queue_images_to_nats(job, images) + + self.assertFalse(result) + self.assertEqual(instance.publish_task.await_count, 0) + + def test_default_chunk_size_is_explicit(self, _mgr_cls, _state_cls): + # Lock the constant so a future change to the chunk size is a deliberate + # decision visible in the diff. The value itself is not a tuned magic + # number — it's an empirical floor that keeps gather() bounded. + self.assertEqual(orchestration_jobs.NATS_PUBLISH_FANOUT_CHUNK_SIZE, 200) From 9ae8999eeed7e100a234f0af2fdc187dc8724591 Mon Sep 17 00:00:00 2001 From: Michael Bunsen Date: Wed, 27 May 2026 16:34:45 -0700 Subject: [PATCH 06/11] feat(ml): emit throttled Collect-stage progress in filter_processed_images MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit filter_processed_images now accepts optional `job` and `total` kwargs and emits a fractional `collect` progress update at most once per COLLECT_PROGRESS_SAVE_INTERVAL_SECONDS (5s wall time), capped at COLLECT_PROGRESS_MAX_FRACTION (0.99). collect_images wires both through. Reaper at ami/jobs/tasks.py:929 triggers on (status in running_states) AND (updated_at < cutoff, default 10min). Before this commit, the collect stage held a worker for several seconds to ~minutes on large collections without ever saving the Job row, so updated_at went stale and the reaper revoked otherwise-healthy jobs. The throttled save now keeps updated_at fresh. Cap at 0.99 so the caller's terminal status=SUCCESS, progress=1 transition still owns the final flip. Legacy callers without `job` see zero job.save() calls — the throttle block is fully gated on both new kwargs. Two new tests in ami/ml/tests.py: - mocked time.monotonic verifies the >=5s gate fires exactly the expected number of times across a multi-chunk run. - omitting `job` keeps the legacy code path silent. Co-Authored-By: Claude --- ami/ml/models/pipeline.py | 42 +++++++++++++++++++++++- ami/ml/tests.py | 67 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 108 insertions(+), 1 deletion(-) diff --git a/ami/ml/models/pipeline.py b/ami/ml/models/pipeline.py index 8d992833c..4028f578a 100644 --- a/ami/ml/models/pipeline.py +++ b/ami/ml/models/pipeline.py @@ -59,6 +59,13 @@ FILTER_PROCESSED_BATCH_SIZE = 1000 +# Minimum wall-time (seconds) between in-loop Job.progress writes inside +# filter_processed_images. The reaper's "no forward progress" heuristic keys off +# job.updated_at, so we need to tick at least once per ~10min; 5s gives ample +# headroom while still amortising the JSONB UPDATE across many chunks. Caps at +# 0.99 so the caller still owns the final status=SUCCESS, progress=1 flip. +COLLECT_PROGRESS_SAVE_INTERVAL_SECONDS = 5.0 +COLLECT_PROGRESS_MAX_FRACTION = 0.99 def filter_processed_images( @@ -66,6 +73,8 @@ def filter_processed_images( pipeline: Pipeline, task_logger: logging.Logger = logger, batch_size: int = FILTER_PROCESSED_BATCH_SIZE, + job: Job | None = None, + total: int | None = None, ) -> typing.Iterable[SourceImage]: """ Return only images that need to be processed by a given pipeline. @@ -102,6 +111,11 @@ def filter_processed_images( task_logger.warning(f"Pipeline {pipeline} has no classification algorithms saved. Will reprocess all images.") image_iter = iter(images) + # Track how many of the input images we've inspected so far so we can emit + # a fractional `collect` progress to the Job row. Only used when both + # `job` and `total` are passed by the caller; legacy callers stay silent. + processed_count = 0 + last_progress_save_monotonic = time.monotonic() while True: batch = list(itertools.islice(image_iter, batch_size)) if not batch: @@ -176,6 +190,20 @@ def filter_processed_images( f"Image {image} has existing detections classified by the pipeline: {pipeline}, skipping!" ) + # Throttled progress emit. Save only when both `job` and `total` are + # provided, the total is non-zero, and at least + # COLLECT_PROGRESS_SAVE_INTERVAL_SECONDS of wall time have passed since + # the last save. Capped at COLLECT_PROGRESS_MAX_FRACTION so the caller's + # final status=SUCCESS, progress=1 flip still owns the terminal value. + processed_count += len(batch) + if job is not None and total: + now_monotonic = time.monotonic() + if now_monotonic - last_progress_save_monotonic >= COLLECT_PROGRESS_SAVE_INTERVAL_SECONDS: + fraction = min(processed_count / total, COLLECT_PROGRESS_MAX_FRACTION) + job.progress.update_stage("collect", progress=fraction) + job.save(update_fields=["progress"]) + last_progress_save_monotonic = now_monotonic + def collect_images( collection: SourceImageCollection | None = None, @@ -213,7 +241,19 @@ def collect_images( if pipeline and not reprocess_all_images: msg = f"Filtering images that have already been processed by pipeline {pipeline}" task_logger.info(msg) - images = list(filter_processed_images(images, pipeline, task_logger=task_logger)) + # Pass `job` and `total` so filter_processed_images can emit throttled + # `collect` progress updates as it chews through the input — keeps the + # reaper "no forward progress" heuristic happy on large collections + # (issue #1321 follow-up). + images = list( + filter_processed_images( + images, + pipeline, + task_logger=task_logger, + job=job, + total=total_images, + ) + ) else: msg = "NOT filtering images that have already been processed" task_logger.info(msg) diff --git a/ami/ml/tests.py b/ami/ml/tests.py index f1caba0e6..cf2a5a0c8 100644 --- a/ami/ml/tests.py +++ b/ami/ml/tests.py @@ -1088,6 +1088,73 @@ def test_filter_processed_images_preserves_input_order_across_batches(self): result = list(filter_processed_images(images, self.pipeline, batch_size=3)) self.assertEqual(result, images) + def test_filter_processed_images_emits_throttled_collect_progress(self): + """ + When `job` and `total` are passed, filter_processed_images should call + job.save(update_fields=["progress"]) at most once per + COLLECT_PROGRESS_SAVE_INTERVAL_SECONDS of wall time, capped at + COLLECT_PROGRESS_MAX_FRACTION. Keeps the reaper's "no forward progress" + heuristic happy on multi-minute Collect stages without hot-saving the + Job row on every chunk (issue #1321 follow-up). + """ + from unittest.mock import patch + + from ami.jobs.models import Job, MLJob + from ami.ml.models.pipeline import COLLECT_PROGRESS_MAX_FRACTION, filter_processed_images + + job = Job.objects.create( + project=self.project, + name="collect progress cadence test", + pipeline=self.pipeline, + job_type_key=MLJob.key, + ) + # First save triggered MLJob.setup → "collect" stage exists. + job.progress.get_stage("collect") + + images = [SourceImage.objects.create(path=f"cadence-{i}.jpg") for i in range(10)] + + # batch_size=3 over 10 images → 4 batches. + # Per-batch time.monotonic() return values: + # init=0, after batch1=1 (gap 1, no save), + # after batch2=6 (gap 6, SAVE → last=6), + # after batch3=7 (gap 1, no save), + # after batch4=12 (gap 6, SAVE → last=12). + # Expected: exactly 2 saves. + monotonic_values = iter([0.0, 1.0, 6.0, 7.0, 12.0]) + with patch("ami.ml.models.pipeline.time.monotonic", side_effect=lambda: next(monotonic_values)): + with patch.object(Job, "save", autospec=True) as mock_save: + list(filter_processed_images(images, self.pipeline, batch_size=3, job=job, total=10)) + + self.assertEqual(mock_save.call_count, 2, "Expected throttle to allow exactly 2 saves") + for call in mock_save.call_args_list: + self.assertEqual( + call.kwargs.get("update_fields"), + ["progress"], + "Throttled saves should write only the progress column", + ) + + # Final emitted fraction comes from the second save (batch 4): processed=10, + # total=10 → raw 1.0, capped at COLLECT_PROGRESS_MAX_FRACTION. + collect_stage = job.progress.get_stage("collect") + self.assertEqual(collect_stage.progress, COLLECT_PROGRESS_MAX_FRACTION) + + def test_filter_processed_images_skips_progress_emission_without_job(self): + """ + Legacy callers that omit `job` (the only callers before this change) + should see zero job.save() calls — the throttle block is fully gated + on both `job` and `total` being passed. + """ + from unittest.mock import patch + + from ami.jobs.models import Job + from ami.ml.models.pipeline import filter_processed_images + + images = [SourceImage.objects.create(path=f"nojob-{i}.jpg") for i in range(5)] + with patch.object(Job, "save", autospec=True) as mock_save: + list(filter_processed_images(images, self.pipeline, batch_size=2)) + + self.assertEqual(mock_save.call_count, 0) + def test_null_detections_are_algorithm_specific(self): """ Null detections from different pipelines/algorithms should not be shared. From abf1fa7ff69a33826964dcfe2dd1465d98b20249 Mon Sep 17 00:00:00 2001 From: Michael Bunsen Date: Wed, 27 May 2026 16:51:06 -0700 Subject: [PATCH 07/11] fix(ml): address Copilot review nits in filter_processed_images MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three small fixes flagged on PR #1322: - Typo in the debug log message ("do yet have" → "do not yet have"). - Reword the `collect_images` source-selection comment so it accurately describes when the `deployment__data_source` join is applied — the `source_images` branch hands the caller's iterable through unchanged. - Fail fast on non-positive `batch_size` instead of silently filtering out every image when an `islice` of 0/negative returns an empty batch. Co-Authored-By: Claude --- ami/ml/models/pipeline.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/ami/ml/models/pipeline.py b/ami/ml/models/pipeline.py index 4028f578a..845ba895b 100644 --- a/ami/ml/models/pipeline.py +++ b/ami/ml/models/pipeline.py @@ -99,6 +99,9 @@ def filter_processed_images( Collect stage well under the Celery broker's AMQP heartbeat window for the large-collection case (see issue #1321). """ + if batch_size <= 0: + raise ValueError(f"batch_size must be > 0, got {batch_size}") + pipeline_algorithms = list(pipeline.algorithms.all()) pipeline_algorithm_ids = [a.id for a in pipeline_algorithms] @@ -181,7 +184,8 @@ def filter_processed_images( task_logger.debug( f"Image {image} has existing detections that haven't been classified by the pipeline: {pipeline}: " f"{observed_classifier_ids} vs {pipeline_classifier_ids}. " - f"Since we do yet have a mechanism to reclassify detections, processing the image from scratch." + f"Since we do not yet have a mechanism to reclassify detections, " + f"processing the image from scratch." ) task_logger.debug(f"Image #{image.pk} needs classification by pipeline's algorithms: {missing_algos}") yield image @@ -225,9 +229,12 @@ def collect_images( else: job = None - # Set source to first argument that is not None. Always prefetch the - # deployment + data_source joins so later calls to image.url() in - # queue_images_to_nats don't trigger N+1 lookups (see issue #1321). + # Set source to first argument that is not None. The collection- and + # deployment-backed branches prefetch the deployment + data_source joins so + # later calls to image.url() in queue_images_to_nats don't trigger N+1 + # lookups (see issue #1321). The source_images branch passes the caller's + # iterable through unchanged — callers handing us a pre-built list are + # responsible for any prefetching they need. if collection: images = collection.images.select_related("deployment__data_source") elif source_images: From 8410a166c68a2a9a8b70c515c324645be02ca9da Mon Sep 17 00:00:00 2001 From: Michael Bunsen Date: Wed, 27 May 2026 16:51:15 -0700 Subject: [PATCH 08/11] refactor(ml): expose ensure_job_resources public wrapper on TaskQueueManager Per Copilot review on PR #1322: calling _ensure_stream / _ensure_consumer directly from the orchestrator coupled queue_images_to_nats to private TaskQueueManager internals. Adds a thin public wrapper that calls both, keeps the per-instance "already-warmed" caches behind the manager's API, and gives future refactors a single seam to evolve. Switches the publish-loop warm-up in queue_images_to_nats to the new public method and updates the fanout unit tests accordingly. Co-Authored-By: Claude --- ami/ml/orchestration/jobs.py | 3 +-- ami/ml/orchestration/nats_queue.py | 14 ++++++++++++++ ami/ml/orchestration/tests/test_jobs.py | 13 ++++++------- 3 files changed, 21 insertions(+), 9 deletions(-) diff --git a/ami/ml/orchestration/jobs.py b/ami/ml/orchestration/jobs.py index 868295245..872a8559f 100644 --- a/ami/ml/orchestration/jobs.py +++ b/ami/ml/orchestration/jobs.py @@ -139,8 +139,7 @@ async def queue_all_images(): # first call, each one still runs inside the publish coroutine and # serialises with the gather below. try: - await manager._ensure_stream(job.pk) - await manager._ensure_consumer(job.pk) + await manager.ensure_job_resources(job.pk) except Exception as e: await manager.log_async( logging.ERROR, diff --git a/ami/ml/orchestration/nats_queue.py b/ami/ml/orchestration/nats_queue.py index 1338e05ff..7f85c1ad3 100644 --- a/ami/ml/orchestration/nats_queue.py +++ b/ami/ml/orchestration/nats_queue.py @@ -290,6 +290,20 @@ async def _stream_exists(self, stream_name: str) -> bool: except nats.js.errors.NotFoundError: return False + async def ensure_job_resources(self, job_id: int) -> None: + """Ensure both the stream and consumer for the given job exist. + + Public wrapper for the two ``_ensure_*`` calls. Callers that want to + pre-warm both before a hot loop (e.g. a publish fanout) should use this + rather than touching the private methods directly — keeps the + TaskQueueManager's stream/consumer setup story in one place. + + Subsequent calls in the same manager session skip the NATS round-trip + via the per-instance ``_streams_logged`` / ``_consumers_logged`` sets. + """ + await self._ensure_stream(job_id) + await self._ensure_consumer(job_id) + async def _ensure_stream(self, job_id: int): """Ensure stream exists for the given job. diff --git a/ami/ml/orchestration/tests/test_jobs.py b/ami/ml/orchestration/tests/test_jobs.py index 9d14698f9..f5f23de49 100644 --- a/ami/ml/orchestration/tests/test_jobs.py +++ b/ami/ml/orchestration/tests/test_jobs.py @@ -18,6 +18,7 @@ def _stub_manager(mock_manager_cls, publish_results=None): instance = mock_manager_cls.return_value instance.__aenter__ = AsyncMock(return_value=instance) instance.__aexit__ = AsyncMock(return_value=False) + instance.ensure_job_resources = AsyncMock() instance._ensure_stream = AsyncMock() instance._ensure_consumer = AsyncMock() instance.log_async = AsyncMock() @@ -85,8 +86,7 @@ def test_warms_stream_and_consumer_once_before_publishing(self, mock_mgr_cls, _s queue_images_to_nats(job, images) - self.assertEqual(instance._ensure_stream.await_count, 1) - self.assertEqual(instance._ensure_consumer.await_count, 1) + self.assertEqual(instance.ensure_job_resources.await_count, 1) self.assertEqual(instance.publish_task.await_count, 5) def test_chunks_across_fanout_boundary(self, mock_mgr_cls, _state_cls): @@ -100,8 +100,7 @@ def test_chunks_across_fanout_boundary(self, mock_mgr_cls, _state_cls): self.assertEqual(instance.publish_task.await_count, n) # Warm-up still only runs once even when we span multiple chunks. - self.assertEqual(instance._ensure_stream.await_count, 1) - self.assertEqual(instance._ensure_consumer.await_count, 1) + self.assertEqual(instance.ensure_job_resources.await_count, 1) def test_partial_failure_counted_as_failed_not_raised(self, mock_mgr_cls, _state_cls): # 3 images, middle one fails. Method should return False (failure summary) @@ -116,10 +115,10 @@ def test_partial_failure_counted_as_failed_not_raised(self, mock_mgr_cls, _state self.assertEqual(instance.publish_task.await_count, 3) def test_stream_setup_failure_short_circuits_publishing(self, mock_mgr_cls, _state_cls): - # If _ensure_stream raises, we should not attempt any publish_task calls — - # the whole batch is marked failed in one shot. + # If ensure_job_resources raises, we should not attempt any publish_task + # calls — the whole batch is marked failed in one shot. instance = _stub_manager(mock_mgr_cls) - instance._ensure_stream = AsyncMock(side_effect=RuntimeError("nats down")) + instance.ensure_job_resources = AsyncMock(side_effect=RuntimeError("nats down")) images = self._make_images(4) job = self._make_job() From 4133e6679d97588348d76ca4d19309b8a90d6517 Mon Sep 17 00:00:00 2001 From: Michael Bunsen Date: Wed, 27 May 2026 18:01:38 -0700 Subject: [PATCH 09/11] test(ml): drop two redundant tests, harden throttle clock stub MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to the takeaway review on PR #1322: - Drop `test_filter_processed_images_preserves_input_order_across_batches` — `test_filter_processed_images_mixed_batch` already asserts the output order with `[unprocessed, unclassified]`. There is no shuffle/reorder path in the function for a cross-batch ordering test to exercise. - Drop `test_default_chunk_size_is_explicit` — locking the constant value catches no behavior regression; the explanatory comment at `ami/ml/orchestration/jobs.py:14-19` is the real documentation. - Replace the `iter([...])` + `next(...)` clock stub in `test_filter_processed_images_emits_throttled_collect_progress` with a counter-based stub. The previous pattern would crash with `StopIteration` if a future change in `filter_processed_images` added another `time.monotonic()` call; the counter form fails with a clear cadence-mismatch assertion instead. No coverage lost. Co-Authored-By: Claude --- ami/ml/orchestration/tests/test_jobs.py | 7 ----- ami/ml/tests.py | 34 ++++++++++++------------- 2 files changed, 17 insertions(+), 24 deletions(-) diff --git a/ami/ml/orchestration/tests/test_jobs.py b/ami/ml/orchestration/tests/test_jobs.py index f5f23de49..b257bb62d 100644 --- a/ami/ml/orchestration/tests/test_jobs.py +++ b/ami/ml/orchestration/tests/test_jobs.py @@ -5,7 +5,6 @@ from ami.jobs.models import Job, JobDispatchMode, MLJob from ami.main.models import Deployment, Project, S3StorageSource, SourceImage, SourceImageCollection from ami.ml.models import Pipeline -from ami.ml.orchestration import jobs as orchestration_jobs from ami.ml.orchestration.jobs import NATS_PUBLISH_FANOUT_CHUNK_SIZE, queue_images_to_nats @@ -126,9 +125,3 @@ def test_stream_setup_failure_short_circuits_publishing(self, mock_mgr_cls, _sta self.assertFalse(result) self.assertEqual(instance.publish_task.await_count, 0) - - def test_default_chunk_size_is_explicit(self, _mgr_cls, _state_cls): - # Lock the constant so a future change to the chunk size is a deliberate - # decision visible in the diff. The value itself is not a tuned magic - # number — it's an empirical floor that keeps gather() bounded. - self.assertEqual(orchestration_jobs.NATS_PUBLISH_FANOUT_CHUNK_SIZE, 200) diff --git a/ami/ml/tests.py b/ami/ml/tests.py index cf2a5a0c8..ed1e72bf2 100644 --- a/ami/ml/tests.py +++ b/ami/ml/tests.py @@ -1080,14 +1080,6 @@ def test_filter_processed_images_query_count_is_bounded_per_batch(self): # First 5 fully classified → skipped. Last 5 fresh → yielded. self.assertEqual(result, images[5:]) - def test_filter_processed_images_preserves_input_order_across_batches(self): - """Output should preserve input ordering even when batched.""" - from ami.ml.models.pipeline import filter_processed_images - - images = [SourceImage.objects.create(path=f"order-{i}.jpg") for i in range(7)] - result = list(filter_processed_images(images, self.pipeline, batch_size=3)) - self.assertEqual(result, images) - def test_filter_processed_images_emits_throttled_collect_progress(self): """ When `job` and `total` are passed, filter_processed_images should call @@ -1113,15 +1105,23 @@ def test_filter_processed_images_emits_throttled_collect_progress(self): images = [SourceImage.objects.create(path=f"cadence-{i}.jpg") for i in range(10)] - # batch_size=3 over 10 images → 4 batches. - # Per-batch time.monotonic() return values: - # init=0, after batch1=1 (gap 1, no save), - # after batch2=6 (gap 6, SAVE → last=6), - # after batch3=7 (gap 1, no save), - # after batch4=12 (gap 6, SAVE → last=12). - # Expected: exactly 2 saves. - monotonic_values = iter([0.0, 1.0, 6.0, 7.0, 12.0]) - with patch("ami.ml.models.pipeline.time.monotonic", side_effect=lambda: next(monotonic_values)): + # batch_size=3 over 10 images → 4 batches. Each monotonic() call + # advances the clock by 3s. Expected sequence: init=0, batch1=3 + # (gap 3, no save), batch2=6 (gap 6, SAVE → last=6), batch3=9 + # (gap 3, no save), batch4=12 (gap 6, SAVE → last=12). Two saves. + # + # Counter-based stub (vs an iter/next sequence) means extra + # monotonic() calls added to filter_processed_images later will + # advance time faster and the assertion will fail with a clear + # cadence mismatch, not StopIteration. + clock = {"t": 0.0} + + def fake_monotonic(): + t = clock["t"] + clock["t"] += 3.0 + return t + + with patch("ami.ml.models.pipeline.time.monotonic", side_effect=fake_monotonic): with patch.object(Job, "save", autospec=True) as mock_save: list(filter_processed_images(images, self.pipeline, batch_size=3, job=job, total=10)) From 51a7fff55d30fee4a10dc8691db8d9424b1ae84e Mon Sep 17 00:00:00 2001 From: Michael Bunsen Date: Wed, 27 May 2026 18:24:21 -0700 Subject: [PATCH 10/11] fix(ml): address CodeRabbit review on filter_processed_images MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three correctness fixes and one nit cleanup from the CodeRabbit review on PR #1322: 1. **`Job.save` heartbeat now bumps `updated_at`.** The throttled progress save in `filter_processed_images` was calling `job.save(update_fields=["progress"])`. Django's `auto_now=True` pre_save hook only fires for fields listed in `update_fields`, so `Job.updated_at` was never refreshed. The reaper at `ami/jobs/tasks.py:929-944` keys off `Job.updated_at < cutoff`, so the heartbeat did not actually defeat the reaper as intended. Including `updated_at` in `update_fields` restores the auto_now behavior. Test assertion updated to enforce the contract. 2. **Short-circuit the "no classifier" path.** When a pipeline has zero classifier algorithms registered, `pipeline_classifier_ids` is the empty set and `set().issubset(anything)` is vacuously True — so the subset check below would skip every image with existing detections, directly contradicting the warning "Will reprocess all images." Now `yield from images; return` after the warning so the warning's stated behavior actually holds. 3. **Drop the `hasattr(image, "url")` guard in `queue_images_to_nats`.** `images` is typed `list[SourceImage]` and `SourceImage` always defines `url()`. The guard masked nothing real and confused readers. 4. **`# noqa: S106` on the fixture `secret_key="y"` values** in two test files so future runs under Ruff's bandit ruleset stay clean. These are clearly fixture values, never real credentials. Co-Authored-By: Claude --- ami/ml/models/pipeline.py | 13 ++++++++++++- ami/ml/orchestration/jobs.py | 2 +- ami/ml/orchestration/tests/test_jobs.py | 2 +- ami/ml/tests.py | 7 ++++--- 4 files changed, 18 insertions(+), 6 deletions(-) diff --git a/ami/ml/models/pipeline.py b/ami/ml/models/pipeline.py index 845ba895b..7cb54d14f 100644 --- a/ami/ml/models/pipeline.py +++ b/ami/ml/models/pipeline.py @@ -112,6 +112,11 @@ def filter_processed_images( pipeline_classifier_ids = {a.id for a in pipeline_algorithms if a.task_type not in detection_type_keys} if not pipeline_classifier_ids: task_logger.warning(f"Pipeline {pipeline} has no classification algorithms saved. Will reprocess all images.") + # set().issubset(anything) is vacuously True, so without this short-circuit + # every image with existing detections gets skipped by the subset check + # below — contradicting the warning above. Yield everything and exit. + yield from images + return image_iter = iter(images) # Track how many of the input images we've inspected so far so we can emit @@ -199,13 +204,19 @@ def filter_processed_images( # COLLECT_PROGRESS_SAVE_INTERVAL_SECONDS of wall time have passed since # the last save. Capped at COLLECT_PROGRESS_MAX_FRACTION so the caller's # final status=SUCCESS, progress=1 flip still owns the terminal value. + # + # `updated_at` is included in update_fields explicitly: Django only fires + # auto_now's pre_save hook for fields listed in update_fields, so without + # it the reaper's `Job.updated_at < cutoff` heuristic + # (`ami/jobs/tasks.py:929-944`) would not see this heartbeat and could + # still revoke the job mid-Collect. processed_count += len(batch) if job is not None and total: now_monotonic = time.monotonic() if now_monotonic - last_progress_save_monotonic >= COLLECT_PROGRESS_SAVE_INTERVAL_SECONDS: fraction = min(processed_count / total, COLLECT_PROGRESS_MAX_FRACTION) job.progress.update_stage("collect", progress=fraction) - job.save(update_fields=["progress"]) + job.save(update_fields=["progress", "updated_at"]) last_progress_save_monotonic = now_monotonic diff --git a/ami/ml/orchestration/jobs.py b/ami/ml/orchestration/jobs.py index 872a8559f..a965ed7c9 100644 --- a/ami/ml/orchestration/jobs.py +++ b/ami/ml/orchestration/jobs.py @@ -104,7 +104,7 @@ def queue_images_to_nats(job: "Job", images: list[SourceImage]): # touches deployment + data_source and the call cost adds up across # large collections. The upstream queryset in collect_images() also # prefetches those joins so this stays cheap (see issue #1321). - image_url = image.url() if hasattr(image, "url") else None + image_url = image.url() if not image_url: job.logger.warning(f"Image {image.pk} has no URL, skipping queuing to NATS for job '{job.pk}'") skipped_count += 1 diff --git a/ami/ml/orchestration/tests/test_jobs.py b/ami/ml/orchestration/tests/test_jobs.py index b257bb62d..04aaf30ed 100644 --- a/ami/ml/orchestration/tests/test_jobs.py +++ b/ami/ml/orchestration/tests/test_jobs.py @@ -44,7 +44,7 @@ def setUp(self): name="ds", bucket="b", access_key="x", - secret_key="y", + secret_key="y", # noqa: S106 - fixture value, never used as a real credential public_base_url="https://example.invalid/", project=self.project, ) diff --git a/ami/ml/tests.py b/ami/ml/tests.py index ed1e72bf2..b2f5a34ed 100644 --- a/ami/ml/tests.py +++ b/ami/ml/tests.py @@ -566,7 +566,7 @@ def test_collect_images_prefetches_deployment_and_data_source(self): name="ds-prefetch-test", bucket="prefetch-bucket", access_key="x", - secret_key="y", + secret_key="y", # noqa: S106 - fixture value, never used as a real credential public_base_url="https://example.invalid/", project=self.project, ) @@ -1129,8 +1129,9 @@ def fake_monotonic(): for call in mock_save.call_args_list: self.assertEqual( call.kwargs.get("update_fields"), - ["progress"], - "Throttled saves should write only the progress column", + ["progress", "updated_at"], + "Throttled saves must include `updated_at` so Django's auto_now fires " + "and the reaper's stale-job heuristic sees forward motion.", ) # Final emitted fraction comes from the second save (batch 4): processed=10, From 373b16eeeb30a998f58f9c79d2a490311add22b5 Mon Sep 17 00:00:00 2001 From: Michael Bunsen Date: Wed, 27 May 2026 18:37:00 -0700 Subject: [PATCH 11/11] test(ml): cover no-classifier short-circuit in filter_processed_images Add a behavioural test that pins the fix from 51a7fff5: a pipeline with no classifier algorithms registered must yield every input image, matching the "Will reprocess all images" warning. Without the short-circuit added in 51a7fff5, the empty `pipeline_classifier_ids` set makes the `set().issubset(observed)` check vacuously True and images with detections get silently skipped. Co-Authored-By: Claude --- ami/ml/tests.py | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/ami/ml/tests.py b/ami/ml/tests.py index b2f5a34ed..ea375135e 100644 --- a/ami/ml/tests.py +++ b/ami/ml/tests.py @@ -1007,6 +1007,33 @@ def test_filter_processed_images_empty_input(self): result = list(filter_processed_images([], self.pipeline)) self.assertEqual(result, []) + def test_filter_processed_images_yields_all_when_pipeline_has_no_classifiers(self): + """ + When a pipeline has no classifier algorithms registered, filter_processed_images + must yield every image (matching the "Will reprocess all images" warning). + Without the short-circuit, the empty `pipeline_classifier_ids` set makes + `set().issubset(observed) == True` and every image with existing detections + is silently skipped — directly contradicting the warning. + """ + from ami.ml.models.pipeline import filter_processed_images + + detector_only_pipeline = Pipeline.objects.create(name="Detector Only Pipeline") + detector_only_pipeline.algorithms.set([self.algorithms["random-detector"]]) + + # Image with a real, fully-processed-looking detection from the detector. + # Pre-short-circuit this would be skipped because the empty pipeline classifier + # set is vacuously a subset of any observed-classifier set. + image_with_detection = SourceImage.objects.create(path="no-classifier-with-det.jpg") + Detection.objects.create( + source_image=image_with_detection, + detection_algorithm=self.algorithms["random-detector"], + bbox=[0.1, 0.2, 0.3, 0.4], + ) + image_unprocessed = SourceImage.objects.create(path="no-classifier-unprocessed.jpg") + + result = list(filter_processed_images([image_with_detection, image_unprocessed], detector_only_pipeline)) + self.assertEqual(result, [image_with_detection, image_unprocessed]) + def test_filter_processed_images_mixed_batch(self): """ A mixed batch of images covering all five branches should yield only