Worker implementation for processing images in the Antenna task queue (PSv2)

[carlosgjs]

Implements a worker service (ami worker) that processes images queued by users in the Antenna web platform. When users upload images to Antenna and request them to be processed, this worker pulls tasks from the queue via the jobs API, runs them through the local ML pipeline (detection + classification), and posts results back. This allows processing to be run behind a firewall (for example in a university HPCC) and for any number of workers to process images in parallel.

This is the counterpart to RolnickLab/antenna#987 which adds the job queue API to Antenna.

Usage

# Configure via environment variables or .env file
export AMI_ANTENNA_API_AUTH_TOKEN=<token>
export AMI_ANTENNA_API_BASE_URL=http://localhost:8000/api/v2  # default

# Register pipelines with projects (optional)
ami worker register "My Worker" --project 1 --project 2

# Start processing jobs - all pipelines by default
ami worker

# Or specify specific pipeline(s)
ami worker --pipeline uk_denmark_moths_2023
ami worker --pipeline moth_binary --pipeline panama_moths_2024

Or configure in settings:

• antenna_api_base_url - Antenna API endpoint (default: http://localhost:8000/api/v2)
• antenna_api_auth_token - Authentication token for Antenna project
• antenna_api_batch_size - Number of tasks per batch (default: 4)

Architecture

The worker functionality is implemented in a dedicated trapdata/antenna/ module for separation of concerns and future portability:

• trapdata/antenna/client.py - API client for fetching jobs and posting results
• trapdata/antenna/worker.py - Worker loop and job processing logic
• trapdata/antenna/registration.py - Pipeline registration with Antenna projects
• trapdata/antenna/schemas.py - Pydantic models for Antenna API requests/responses
• trapdata/antenna/datasets.py - RESTDataset for streaming tasks from the API
• trapdata/antenna/tests/ - Integration tests with mock Antenna API server
• trapdata/cli/worker.py - Thin CLI wrapper (~75 lines) that delegates to the antenna module

Changes

Worker implementation (trapdata/antenna/worker.py)

• Polls /jobs endpoint for available jobs by pipeline slug
• Fetches tasks from /jobs/{id}/tasks
• Processes images through localization + classification pipeline
• Posts results back to /jobs/{id}/result/
• Continuous loop with configurable polling interval

API client (trapdata/antenna/client.py)

• get_jobs() - Fetches job IDs for a given pipeline
• post_batch_results() - Posts processed results back to Antenna

Pipeline registration (trapdata/antenna/registration.py)

• register_pipelines() - Registers available pipelines with Antenna projects
• get_user_projects() - Fetches accessible projects from Antenna API
• Reads pipeline configurations from CLASSIFIER_CHOICES

Session handling (trapdata/api/utils.py)

• get_http_session() - Creates HTTP session with connection pooling
• Automatic retries with exponential backoff via urllib3 adapter (hardcoded: 3 retries, 0.5s backoff)
• Only retries on server errors (5XX) and network failures, not client errors (4XX)

Schemas (trapdata/antenna/schemas.py)

• Pydantic models for Antenna API: AntennaJobsListResponse, AntennaTasksListResponse, AntennaTaskResult
• AsyncPipelineRegistrationRequest and AsyncPipelineRegistrationResponse for pipeline registration

Dataset (trapdata/antenna/datasets.py)

• RESTDataset - IterableDataset that streams tasks from Antenna API
• Downloads images on-the-fly with error handling
• Reports download failures back to Antenna

Settings (trapdata/settings.py)

• Added antenna_api_base_url, antenna_api_auth_token, antenna_api_batch_size with Kivy UI integration

CLI (trapdata/cli/worker.py)

• ami worker - Default command to run the worker (processes all pipelines)
• ami worker --pipeline <slug> - Process specific pipeline(s) (repeatable flag)
• ami worker register <name> - Register pipelines with Antenna projects
• Uses @cli.callback(invoke_without_command=True) to make worker the default action

Tests (trapdata/antenna/tests/)

• Integration tests with real ML inference (detector + classifier)
• Mock Antenna API server for testing worker interactions
• Tests for registration client functions
• End-to-end test: register → get jobs → process → post results

Credits

• @carlosgjs - Original worker and pipeline registration implementation
• @mihow - Integration tests, session handling, Settings pattern refactoring, module extraction

Test plan

• Run ami test all - all tests pass
• Manual test with local Antenna instance
• Integration tests with real ML inference
• Worker handles network errors gracefully (via retry logic)
• Test with GPU and CPU-only environments

Summary by CodeRabbit

Release Notes

• New Features

  • Added Antenna Worker with pipeline execution and registration capabilities
  • Introduced REST API dataset integration for streaming data processing
  • Added pipeline registration across multiple projects
  • Extended CLI with new worker subcommand group

• Bug Fixes

  • Fixed conditional CUDA cache clearing for systems without GPU support

• Configuration

  • Added Antenna API configuration environment variables
  • Updated pre-commit hooks to latest versions
  • Enhanced pytest test discovery paths

• Documentation

  • Added comprehensive Antenna Worker setup and usage guide

• Tests

  • Added mock Antenna API server for integration testing
  • Implemented comprehensive worker integration and unit test coverage

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[Copilot]

Pull request overview

Adds an end-to-end “pulling (v2)” worker implementation for PyTorch inference that polls an Antenna service for queued tasks, runs detection/classification, and posts results back.

Changes:

• Added ami worker CLI command and worker runtime to poll /api/v2/jobs and post /result/ payloads.
• Implemented REST-backed IterableDataset + DataLoader plumbing for streaming tasks and image downloads.
• Added small API/model utilities (service-info caching, reset helpers, task schema) to support the worker flow.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 7 comments.

Show a summary per file

File	Description
trapdata/common/utils.py	Adds a small timing/logging helper used by the worker.
trapdata/cli/worker.py	New worker loop: fetch jobs, process batches, post results.
trapdata/cli/base.py	Registers worker CLI command and validates pipeline args.
trapdata/api/schemas.py	Adds PipelineProcessingTask schema for REST task payloads.
trapdata/api/models/localization.py	Adds detector reset helper; simplifies result saving loop.
trapdata/api/models/classification.py	Adds classifier reset + shared “update detection classification” helper.
trapdata/api/datasets.py	Adds REST task IterableDataset, collate fn, and DataLoader factory.
trapdata/api/api.py	Caches /info response via FastAPI lifespan initialization.
.vscode/launch.json	Adds VS Code launch configs for debugging worker and API.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[mihow]

Things to follow up on:

• Removing many of the worker tests
• Move the code in cli/worker.py somewhere. Keep the CLI minimal, import the worker functions.

[coderabbitai]

📝 Walkthrough

Walkthrough

This pull request introduces comprehensive Antenna API integration, adding a new worker module for distributed processing, REST-based dataset streaming, pipeline registration orchestration, and API client utilities. It includes configuration management, CLI commands, lifecycle improvements to the FastAPI app, and extensive integration tests alongside development tooling updates.

Changes

Cohort / File(s)	Summary
Configuration & Environment Setup .env.example, .gitignore, .pre-commit-config.yaml, pyproject.toml	Added Antenna API environment variables (base URL, auth token, batch size); expanded .gitignore for test artifacts; updated flake8 hook to v4.0.0; configured pytest test discovery paths for nested test directories.
Development Tools .vscode/launch.json, CLAUDE.md	New VSCode launch configuration with three debug profiles (Python debugger, worker, and API). Updated development guidelines to emphasize interactive staging (git add -p), TDD practice, and simplicity.
Documentation README.md	Added "Running the Antenna Worker" section detailing setup steps, configuration, runtime commands for single/multiple pipelines, parallel worker notes, token binding, and shutdown behavior with link to Antenna platform documentation.
API Server & Service Infrastructure trapdata/api/api.py, trapdata/api/utils.py, trapdata/api/models/classification.py, trapdata/api/models/localization.py, trapdata/api/service.conf	Introduced async FastAPI lifecycle management with service-info caching at startup. Added HTTP session utility with retry strategy and optional auth. Extended APIMothClassifier with reset method and new update_detection_classification method. Simplified APIMothDetector's save_results by replacing threaded processing with synchronous flow. Expanded API supervisord configuration with autostart, autorestart, and log settings.
API Testing trapdata/api/tests/test_api.py, trapdata/api/tests/utils.py	Refactored test_api.py to delegate to shared test utilities. New utils.py module provides helpers for test image sourcing, pipeline class retrieval, and Antenna API request patching via TestClient.
Settings & Configuration trapdata/settings.py	Added three Antenna API settings (antenna_api_base_url, antenna_api_auth_token, antenna_api_batch_size) with corresponding Kivy UI metadata in the Config.fields dictionary.
CLI & Worker Interface trapdata/cli/base.py, trapdata/cli/test.py, trapdata/cli/worker.py	Added new "worker" CLI subcommand group with run and register commands for job processing and pipeline registration. Improved error handling with pretty_exceptions_show_locals=False. Updated test.py to invoke pytest via Python interpreter. Added worker.py module with validated pipeline selection and project registration.
Common Utilities trapdata/common/utils.py	Introduced log_time function for elapsed time measurement with optional logging and chained timing support via functools.partial.
Database Models trapdata/db/models/detections.py	Replaced list comprehensions with no-op assignments in save_detected_objects and save_classified_objects (no functional change).
ML Models trapdata/ml/models/base.py	Made CUDA cache clearing conditional on torch.cuda.is_available() instead of unconditional clearing.
Antenna API Integration (New Module) trapdata/antenna/__init__.py, trapdata/antenna/client.py, trapdata/antenna/datasets.py, trapdata/antenna/registration.py, trapdata/antenna/schemas.py, trapdata/antenna/worker.py	New comprehensive module for Antenna API integration: client.py provides REST API helpers (fetch jobs, post results, get projects); datasets.py implements RESTDataset for streaming task batches with image loading and custom collate; registration.py orchestrates pipeline registration across projects; schemas.py defines Pydantic models for API request/response payloads; worker.py implements the main job processing loop with lazy component initialization, detection, classification, and result aggregation.
Antenna Testing & Service Config trapdata/antenna/tests/antenna_api_server.py, trapdata/antenna/tests/test_worker.py, trapdata/antenna/service.conf	New mock FastAPI Antenna API server for integration testing with endpoints simulating job/task fetching and result posting. Comprehensive integration and end-to-end tests covering rest_collate_fn, RESTDataset, get_jobs, _process_job, and full worker workflows with real inference. New supervisord configuration for the Antenna worker service.

Sequence Diagram(s)

sequenceDiagram
    participant Worker as run_worker()
    participant Client as get_jobs()
    participant API as Antenna API
    participant Dataset as RESTDataset
    participant Detector as APIMothDetector
    participant Classifier as APIMothClassifier
    participant PostResults as post_batch_results()

    Worker->>Client: get_jobs(base_url, token, pipeline)
    Client->>API: GET /jobs?pipeline_slug=...
    API-->>Client: {results: [job_ids]}
    Client-->>Worker: [job_ids]
    
    Worker->>Worker: for each job_id
    Worker->>Dataset: __init__(base_url, token, job_id)
    Worker->>Dataset: __iter__()
    Dataset->>API: GET /jobs/{job_id}/tasks?batch_size=...
    API-->>Dataset: {tasks: [AntennaPipelineProcessingTask]}
    Dataset->>Dataset: load_image(url) for each task
    Dataset-->>Worker: {images, image_ids, reply_subjects}
    
    Worker->>Detector: run(images)
    Detector-->>Worker: detections
    
    Worker->>Classifier: run(image_crops, detection_indices)
    Classifier-->>Worker: predictions
    
    Worker->>PostResults: post_batch_results(base_url, token, job_id, results)
    PostResults->>API: POST /jobs/{job_id}/result/
    API-->>PostResults: 200 OK
    PostResults-->>Worker: True

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~50 minutes

Poem

🐰 Hoppy times ahead, I declare!
New Antenna workers in the air,
REST datasets streaming with care,
Registration pipelines everywhere,
Inference flowing—a distributed affair! 🚀

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 58.52% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately summarizes the primary change: implementing a worker service for processing images from the Antenna task queue, which is the main focus of this substantial PR.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 5

🤖 Fix all issues with AI agents

In `@trapdata/api/datasets.py`:
- Around line 232-290: The exception handler can reference worker_id before it
is assigned; initialize worker_id (and num_workers) before the try so they’re
always defined. In the method shown, set default values like worker_id = 0 and
num_workers = 1 above the try that calls torch.utils.data.get_worker_info(),
then keep the existing logic that overrides them from worker_info; this ensures
logger.error(f"Worker {worker_id}: Exception in iterator: {e}") won’t raise
UnboundLocalError.

In `@trapdata/api/models/classification.py`:
- Around line 156-158: Replace the bare assert that checks
detection.source_image_id == image_id with an explicit validation that raises a
clear exception (e.g., ValueError) when the invariant fails; locate the block
that accesses self.detections[detection_idx] (variables: self.detections,
detection_idx, detection, image_id) and after retrieving detection, if
detection.source_image_id != image_id raise an error containing detection_idx,
detection.source_image_id and the expected image_id to prevent silent failures
when Python assertions are disabled.

In `@trapdata/api/utils.py`:
- Around line 41-90: The current get_http_session function retries POST by
default; change its signature to accept an allowed_methods parameter (e.g.,
allowed_methods: tuple[str, ...] = ("GET",)) and use that value when
constructing the urllib3.Retry (pass allowed_methods=list(allowed_methods) or
equivalent), update the docstring to document the new parameter and default, and
leave existing callers to opt-in to POST retries (e.g., post_batch_results
should explicitly pass allowed_methods=("GET","POST") only if the endpoint is
idempotent or uses idempotency keys). Ensure the Retry call uses the provided
allowed_methods and that session mounting and auth handling remain unchanged.

In `@trapdata/cli/test.py`:
- Around line 48-50: Remove the unnecessary inline noqa directive from the
event_day annotation; specifically edit the Annotated declaration referencing
event_day and typer.Argument(formats=["%Y-%m-%d"]) to delete the trailing "#
noqa: F722" token so the annotation remains exactly as
Annotated[datetime.datetime, typer.Argument(formats=["%Y-%m-%d"])]; no other
changes needed.

In `@trapdata/cli/worker.py`:
- Around line 425-433: The CLI currently falls back to ANTENNA_API_* env vars;
change the registration defaults to use the app settings (read_settings()) and
AMI_ prefixed vars instead: in the function handling registration (where
base_url and auth_token are set), call read_settings() to obtain
settings.api_base_url / settings.api_token (or their AMI_ equivalents) and use
those values as fallbacks before finally checking os.environ for AMI_*
(AMI_API_BASE_URL / AMI_API_TOKEN); then update the logger message to reference
AMI_API_TOKEN when auth_token remains empty. Ensure you update references to
base_url and auth_token in that block only (the section that currently checks
ANTENNA_API_BASE_URL/ANTENNA_API_TOKEN).

🧹 Nitpick comments (13)

trapdata/db/models/detections.py (2)

326-326: Remove the no‑op list construction.
images is already loaded; building and discarding [img.id for img in images] adds overhead without side effects. If the goal was eager-loading, this doesn’t do it.

♻️ Proposed change

-        _ = [img.id for img in images]
+        # No-op removed; images are already loaded by .all()

---

407-407: Avoid the unused query before bulk_update_mappings.
This extra SELECT adds a DB round‑trip and the results are discarded. If existence checks are needed, handle them explicitly; otherwise remove the query.

♻️ Proposed change

-        _ = sesh.query(DetectedObject).filter(DetectedObject.id.in_(object_ids)).all()
+        # Removed unused prefetch; bulk_update_mappings does not require it

trapdata/common/utils.py (2)

126-149: LGTM with a minor style suggestion.

The timing utility is well-designed with a clean chaining API. The docstring clearly explains usage.

Consider updating the return type annotation to use the built-in tuple syntax for consistency with the modern style used elsewhere (str | None):

✨ Style consistency suggestion

-from typing import Any, Callable, Tuple, Union
+from typing import Any, Callable, Union

-def log_time(start: float = 0, msg: str | None = None) -> Tuple[float, Callable]:
+def log_time(start: float = 0, msg: str | None = None) -> tuple[float, Callable]:

---

8-8: Optional: Simplify Union usage.

Since the codebase targets Python 3.10+, you can replace Union[pathlib.Path, str] (lines 79-80) and Union[pathlib.Path, None] with the modern | syntax throughout this file for consistency.

.env.example (1)

12-16: Optional: reorder keys to satisfy dotenv-linter warning.

🔧 Proposed reorder

-AMI_ANTENNA_API_BASE_URL=http://localhost:8000/api/v2
-AMI_ANTENNA_API_AUTH_TOKEN=your_antenna_auth_token_here
+AMI_ANTENNA_API_AUTH_TOKEN=your_antenna_auth_token_here
+AMI_ANTENNA_API_BASE_URL=http://localhost:8000/api/v2
 AMI_ANTENNA_API_BATCH_SIZE=4

trapdata/api/models/localization.py (1)

49-52: Consider zip(..., strict=True) to surface length mismatches.

🔧 Proposed fix

-        for image_id, image_output in zip(item_ids, batch_output):
+        for image_id, image_output in zip(item_ids, batch_output, strict=True):

trapdata/api/api.py (1)

359-374: Startup builds full model configs—consider lazy or metadata-only construction.

Instantiating detectors/classifiers here may load weights and slow startup; a lightweight metadata registry (or lazy load on first /info) can reduce cold-start latency and memory pressure.

trapdata/api/tests/utils.py (1)

15-37: Stabilize test image ordering.

glob() ordering is filesystem-dependent; slicing [:num] can lead to flaky tests. Sort before slicing.

🛠️ Proposed change

-    source_image_urls = [
-        file_server.get_url(f.relative_to(test_images_dir))
-        for f in images_dir.glob("*.jpg")
-    ][:num]
+    source_image_urls = [
+        file_server.get_url(f.relative_to(test_images_dir))
+        for f in sorted(images_dir.glob("*.jpg"))
+    ][:num]

trapdata/api/tests/antenna_api_server.py (2)

10-19: Remove unused import AntennaTaskResults.

The AntennaTaskResults import is not used anywhere in this module.

🧹 Proposed fix

 from trapdata.api.schemas import (
     AntennaJobListItem,
     AntennaJobsListResponse,
     AntennaPipelineProcessingTask,
     AntennaTaskResult,
-    AntennaTaskResults,
     AntennaTasksListResponse,
     AsyncPipelineRegistrationRequest,
     AsyncPipelineRegistrationResponse,
 )

---

30-45: Consider prefixing unused parameters with underscore.

The pipeline__slug, ids_only, and incomplete_only parameters are part of the API contract but unused in this mock. Prefixing with underscores documents this intentionally and silences linter warnings.

🧹 Proposed fix

 `@app.get`("/api/v2/jobs")
-def get_jobs(pipeline__slug: str, ids_only: int, incomplete_only: int):
+def get_jobs(_pipeline__slug: str, _ids_only: int, _incomplete_only: int):
     """Return available job IDs.

trapdata/api/tests/test_worker.py (3)

40-128: Minor style inconsistency: TestRestCollateFn doesn't inherit from TestCase.

All other test classes in this file inherit from unittest.TestCase, but TestRestCollateFn uses pytest-style class syntax. This works fine with pytest but creates an inconsistent pattern. Consider inheriting from TestCase for consistency, or keep as-is if you prefer the lighter pytest style for unit tests.

The test coverage itself is excellent — covering all-successful, all-failed, mixed, and single-item edge cases.

---

285-294: Test test_query_params_sent doesn't actually validate query parameters.

The test comment claims to validate query params, but since the mock server ignores pipeline__slug, ids_only, and incomplete_only, this test only verifies the function returns a list. Consider either:

1. Enhancing the mock to capture and expose the received query params for assertion
2. Renaming/documenting this as a basic smoke test rather than param validation

---

319-329: Duplicate _make_settings helper in two test classes.

_make_settings() is duplicated identically in both TestProcessJobIntegration and TestWorkerEndToEnd. Consider extracting to a shared fixture or utility function in the test module or utils.py.

♻️ Example extraction

# At module level or in utils.py
def make_test_settings():
    """Create mock settings for worker tests."""
    settings = MagicMock()
    settings.antenna_api_base_url = "http://testserver/api/v2"
    settings.antenna_api_auth_token = "test-token"
    settings.antenna_api_batch_size = 2
    settings.antenna_api_retry_max = 3
    settings.antenna_api_retry_backoff = 0.5
    settings.num_workers = 0
    settings.localization_batch_size = 2
    return settings

Also applies to: 469-478

[coderabbitai]

Actionable comments posted: 2

🤖 Fix all issues with AI agents

In `@trapdata/cli/worker.py`:
- Around line 239-245: Clamp bbox coordinates from dresp.bbox to the image
tensor bounds before slicing: retrieve image shape from
image_tensors[dresp.source_image_id], convert bbox coords to ints and clamp
x1,x2 to [0,width] and y1,y2 to [0,height]; verify x2>x1 and y2>y1
(skip/continue for invalid boxes) and only then perform the crop and
unsqueeze—update the crop logic around image_tensors, dresp.bbox, and crop to
include these checks.
- Around line 198-203: Validate lengths before performing any zip operations to
avoid silent truncation: check that len(image_ids) equals len(images) before
zipping image_ids and images (the zip(image_ids, images) call) and check that
len(image_ids) == len(images) == len(reply_subjects) == len(image_urls) before
zipping image_ids, reply_subjects, image_urls, images; if lengths differ, raise
a ValueError with a clear message. Also update the zip calls to use zip(...,
strict=True) (Python 3.10+) so mismatches fail fast; reference the variables
image_ids, images, reply_subjects, image_urls and the zip usages in the worker
function when making these changes.

[coderabbitai]

Actionable comments posted: 2

🤖 Fix all issues with AI agents

In `@trapdata/api/tests/antenna_api_server.py`:
- Around line 10-19: The import list in trapdata/api/tests/antenna_api_server.py
includes an unused symbol AntennaTaskResults; remove AntennaTaskResults from the
from trapdata.api.schemas import (...) statement so the module only imports the
actually used names (e.g., AntennaJobListItem, AntennaJobsListResponse,
AntennaPipelineProcessingTask, AntennaTaskResult, AntennaTasksListResponse,
AsyncPipelineRegistrationRequest, AsyncPipelineRegistrationResponse).

In `@trapdata/api/tests/test_worker.py`:
- Around line 40-41: The test class TestRestCollateFn is missing
unittest.TestCase inheritance so unittest discovery won't run it; update the
class definition to inherit from unittest.TestCase (i.e., class
TestRestCollateFn(unittest.TestCase):) and add an import for unittest if absent,
and while here convert any bare assert statements inside TestRestCollateFn to
unittest assertions (self.assertEqual / self.assertRaises, etc.) so the tests
run correctly under the project's unittest runner.

🧹 Nitpick comments (1)

trapdata/api/tests/antenna_api_server.py (1)

30-45: Consider using filter parameters for more realistic mock behavior.

The pipeline__slug, ids_only, and incomplete_only parameters are accepted but ignored. While acceptable for a basic mock, using these parameters would enable tests to verify that the worker sends correct query params. For instance, you could filter _jobs_queue by a stored pipeline slug or return full job objects when ids_only=0.

If this level of fidelity isn't needed now, consider adding a # noqa: ARG001 comment to suppress the linter warnings.

[coderabbitai]

Actionable comments posted: 1

🤖 Fix all issues with AI agents

In `@trapdata/api/datasets.py`:
- Around line 193-218: In _load_image, fix orientation by applying EXIF-based
transpose immediately after opening the image and before converting mode or
applying image_transforms: after Image.open(BytesIO(response.content)) call
ImageOps.exif_transpose (or equivalent EXIF transpose helper) to correct
rotation, then continue with the existing mode check and conversion to "RGB",
and then call image_transforms; keep existing error handling and use the same
symbols (image_fetch_session, BytesIO, Image, image_transforms).

[coderabbitai]

Actionable comments posted: 3

🤖 Fix all issues with AI agents

In `@trapdata/api/tests/antenna_api_server.py`:
- Around line 29-44: The get_jobs handler currently receives FastAPI query
params pipeline__slug, ids_only, and incomplete_only but doesn't use them; to
satisfy linters without renaming the parameters, explicitly mark them as
intentionally unused inside get_jobs (e.g., assign them to a throwaway variable
or use typing.cast/# pragma if preferred) so the function signature stays
unchanged; update the body of get_jobs to reference pipeline__slug, ids_only,
and incomplete_only in a no-op way (e.g., _ = pipeline__slug; _ = ids_only; _ =
incomplete_only) before computing job_ids.

In `@trapdata/cli/worker.py`:
- Line 171: Replace unguarded calls to torch.cuda.empty_cache() with a guarded
check using torch.cuda.is_available(): wrap the call as if
torch.cuda.is_available(): torch.cuda.empty_cache(). Update the call in the
worker module where torch.cuda.empty_cache() is invoked and make the same change
in the model base module (trapdata.ml.models.base) so CPU-only builds skip CUDA
cache clearing and device selection logic auto-detects GPU availability.
- Around line 310-312: post_batch_results(...) return value is ignored causing
silent data loss when posting fails; change the call in the worker loop to
capture its boolean result (e.g., success = post_batch_results(settings, job_id,
batch_results)), only add its timing to total_save_time when success, and on
failure either requeue the batch (so RESTDataset tasks aren’t lost) or raise/log
a clear error so the job can be retried/inspected; update the code around
post_batch_results, job_id, batch_results and total_save_time to implement this
check and appropriate error handling.

[coderabbitai]

Actionable comments posted: 9

🤖 Fix all issues with AI agents

In `@docs/claude/planning/antenna-module-refactor.md`:
- Around line 108-116: Update the documented test paths and validation checklist
to use the new module layout: replace occurrences of
"trapdata/api/tests/test_worker.py" with "trapdata/antenna/tests/test_worker.py"
(including the pytest command in the example code block and any checklist
entries noted around lines referenced, e.g., the validation checklist and the
commands block showing "pytest ..."). Search the file for any other mentions of
the old path and update them to the new "trapdata/antenna/tests/test_worker.py"
so the example commands and checklist accurately reflect the current structure.

In `@docs/claude/planning/simplify-worker-tests.md`:
- Around line 15-26: The markdown tables containing rows like
`TestRESTDatasetIntegration.test_empty_queue`,
`TestGetJobsIntegration.test_empty_queue`,
`TestProcessJobIntegration.test_empty_queue`,
`TestRESTDatasetIntegration.test_multiple_batches`, and
`TestWorkerEndToEnd.test_multiple_batches_processed` are flagged by markdownlint
MD060 for pipe spacing; fix by normalizing pipe spacing and column alignment
(ensure a single space after and before each pipe or run a markdown table
formatter) so each row and header aligns consistently, then re-run the linter/CI
to confirm the MD060 warning is resolved.

In `@docs/claude/planning/worker-integration-tests.md`:
- Around line 336-345: The markdown table starting with the header "| Test Class
| Tests | Type | What It Tests |" is triggering MD060 due to inconsistent pipe
spacing; fix by normalizing column pipe alignment and spacing (e.g., ensure a
single space on both sides of each pipe for every row) or run the project's
markdown formatter (prettier/markdownlint) to reflow the table, and if alignment
cannot be preserved, add an inline markdownlint disable comment for MD060 above
the table to suppress the warning; target the table block and the similar table
at the later occurrence (the block beginning with the same header).

In `@scripts/validate_dwc_export.py`:
- Around line 69-74: The subspecies count uses a misspelled field 'tqaxonRank'
which will KeyError; update the generator to check the correct field name
'taxonRank' (same as used for species) when computing subspecies from rows
(i.e., fix the expression that computes subspecies to reference row['taxonRank']
instead of row['tqaxonRank']), leaving the rest of the logic (variables species,
subspecies, rows and the print statements) unchanged.
- Around line 36-38: There is a typo in the loop over rank_counts: the computed
percentage variable is named pct but the print uses qpct causing a NameError;
update the print statement inside the loop that iterates "for rank, count in
sorted(rank_counts.items())" to reference pct instead of qpct so it prints f" 
{rank}: {count} ({pct:.1f}%)".

In `@trapdata/antenna/client.py`:
- Around line 3-56: The broad except in get_jobs should be narrowed: replace the
generic "except Exception" that wraps resp.json() and
AntennaJobsListResponse.model_validate() with two specific handlers catching
json.JSONDecodeError (from calling resp.json()) and pydantic.ValidationError
(from AntennaJobsListResponse.model_validate()). Import json and
pydantic.ValidationError (or from pydantic import ValidationError), log distinct
error messages (e.g., "Failed to decode JSON" and "Failed to validate jobs
response") referencing base_url and the exception, and return [] in each
handler; keep the existing requests.RequestException handler for request errors.

In `@trapdata/antenna/schemas.py`:
- Around line 5-9: The import list in the module currently includes an unused
symbol ProcessingServiceInfoResponse; remove ProcessingServiceInfoResponse from
the from-import tuple (the line importing PipelineConfigResponse,
PipelineResultsResponse, ProcessingServiceInfoResponse) so only the actually
used symbols (PipelineConfigResponse and PipelineResultsResponse) are imported
to satisfy flake8.

In `@trapdata/cli/base.py`:
- Around line 2-8: The file imports unused symbols Annotated and
CLASSIFIER_CHOICES which will fail linting; remove Annotated from the typing
import (leave Optional) and delete the import of CLASSIFIER_CHOICES from
trapdata.api.api so only used symbols remain, then re-run flake8 to verify no
unused-import errors.

In `@trapdata/cli/test.py`:
- Line 27: The subprocess call using the bare "pytest" is PATH-dependent; change
the invocation of subprocess.call that sets return_code (the call using
subprocess.call(["pytest", "-v"])) to use the current Python interpreter by
replacing the args with [sys.executable, "-m", "pytest", "-v"] and add an import
sys at the top of the module so the active virtualenv's pytest module is used
reliably.

🧹 Nitpick comments (6)

scripts/validate_dwc_export.py (1)

21-23: Consider handling empty files to avoid division by zero.

If the TSV file contains no data rows, total_taxa will be 0, causing ZeroDivisionError on subsequent percentage calculations (lines 29, 37, 53, 60, 66).

🛡️ Proposed early exit for empty files

     total_taxa = len(rows)
+    if total_taxa == 0:
+        print("No taxa found in file.")
+        return
     print(f"Total Taxa: {total_taxa}")

trapdata/api/api.py (1)

358-373: Avoid eager model initialization when building pipeline configs.
initialize_service_info() instantiates detector/classifier objects for every pipeline; if constructors load weights, startup and CLI registration can become heavy. Consider a lightweight metadata path (e.g., classmethod/constant config) and defer model loading until inference.

trapdata/antenna/schemas.py (1)

64-71: Prefer default_factory for list fields for consistency with the rest of the file.
Other list fields in this file (lines 57, 79–88) use pydantic.Field(default_factory=list). This pattern is also recommended in Pydantic v2 to make intent explicit.

♻️ Suggested change

 class AsyncPipelineRegistrationRequest(pydantic.BaseModel):
     """
     Request to register pipelines from an async processing service
     """
 
     processing_service_name: str
-    pipelines: list[PipelineConfigResponse] = []
+    pipelines: list[PipelineConfigResponse] = pydantic.Field(default_factory=list)

trapdata/antenna/datasets.py (2)

122-137: Consider handling EXIF orientation for rotated images.

The code converts to RGB correctly, but per coding guidelines images should also handle EXIF orientation. Some camera trap images may have EXIF rotation metadata that could cause the model to receive incorrectly oriented images.

The broad Exception catch is acceptable here since external image loading can fail in many unpredictable ways (network errors, invalid formats, corrupted data, etc.).

♻️ Proposed fix to handle EXIF orientation

             response = self.image_fetch_session.get(image_url, timeout=30)
             response.raise_for_status()
             image = Image.open(BytesIO(response.content))
 
+            # Handle EXIF orientation
+            from PIL import ImageOps
+            image = ImageOps.exif_transpose(image)
+
             # Convert to RGB if necessary
             if image.mode != "RGB":
                 image = image.convert("RGB")

As per coding guidelines: "Handle EXIF orientation when preprocessing images; ensure models receive RGB format"

---

179-205: Minor: Redundant conditional and commented debug code.

1. Line 204: The ternary if errors else None is redundant since we're already inside an if errors: block - errors is guaranteed truthy here.

2. Lines 182, 186: Commented-out log_time calls should be removed or uncommented if still needed for debugging.

♻️ Proposed cleanup

                 for task in tasks:
                     errors = []
                     # Load the image
-                    # _, t = log_time()
                     image_tensor = (
                         self._load_image(task.image_url) if task.image_url else None
                     )
-                    # _, t = t(f"Loaded image from {image_url}")
 
                     if image_tensor is None:
                         errors.append("failed to load image")
 
                     if errors:
                         logger.warning(
                             f"Worker {worker_id}: Errors in task for image '{task.image_id}': {', '.join(errors)}"
                         )
 
                     # Yield the data row
                     row = {
                         "image": image_tensor,
                         "reply_subject": task.reply_subject,
                         "image_id": task.image_id,
                         "image_url": task.image_url,
                     }
                     if errors:
-                        row["error"] = "; ".join(errors) if errors else None
+                        row["error"] = "; ".join(errors)
                     yield row

trapdata/antenna/tests/test_worker.py (1)

374-383: Consider extracting shared _make_settings helper to reduce duplication.

The _make_settings method is duplicated between TestProcessJobIntegration and TestWorkerEndToEnd classes. Consider extracting it to a module-level helper or a shared test base class.

♻️ Proposed refactor

# At module level or in a shared test utilities module
def make_test_settings():
    """Create mock settings for worker tests."""
    settings = MagicMock()
    settings.antenna_api_base_url = "http://testserver/api/v2"
    settings.antenna_api_auth_token = "test-token"
    settings.antenna_api_batch_size = 2
    settings.antenna_api_retry_max = 3
    settings.antenna_api_retry_backoff = 0.5
    settings.num_workers = 0
    settings.localization_batch_size = 2
    return settings

Then use in both test classes:

def _make_settings(self):
    return make_test_settings()

Also applies to: 224-234

[coderabbitai]

Actionable comments posted: 1

🤖 Fix all issues with AI agents

In `@trapdata/antenna/tests/test_worker.py`:
- Around line 139-146: Replace hardcoded auth token strings in tests with an
environment-backed test constant: create a single constant (e.g.,
TEST_AUTH_TOKEN = os.getenv("AMI_TEST_AUTH_TOKEN", "test-token")) and use it
wherever RESTDataset is instantiated (refer to _make_dataset and other test
helpers that call RESTDataset), so calls like RESTDataset(...,
auth_token="test-token") become auth_token=TEST_AUTH_TOKEN; ensure you import os
and define the constant once at the top of the test module to keep tests
deterministic and satisfy Ruff S105/S106.

🧹 Nitpick comments (2)

trapdata/antenna/datasets.py (1)

96-121: Consider handling EXIF orientation when loading images.

The _load_image method converts images to RGB but doesn't handle EXIF orientation metadata. Images from cameras may appear rotated incorrectly. As per coding guidelines: "Handle EXIF orientation when preprocessing images."

♻️ Proposed fix to handle EXIF orientation

+from PIL import ImageOps
 ...
             image = Image.open(BytesIO(response.content))
 
+            # Handle EXIF orientation before any other processing
+            image = ImageOps.exif_transpose(image)
+
             # Convert to RGB if necessary
             if image.mode != "RGB":
                 image = image.convert("RGB")

trapdata/antenna/tests/test_worker.py (1)

207-487: Consider gating the real-inference integration tests behind an opt‑in flag.
These tests run real ML inference (not mocked) and local file servers; if included in default CI runs, they can be slow or fail on CPU‑only or model‑less environments.

♻️ Suggested opt‑in gating (unittest)

+import os
-from unittest import TestCase
+from unittest import TestCase, skipUnless
@@
+RUN_INTEGRATION = os.getenv("RUN_INTEGRATION_TESTS") == "1"
@@
+@skipUnless(RUN_INTEGRATION, "integration tests are opt-in")
 class TestProcessJobIntegration(TestCase):
@@
+@skipUnless(RUN_INTEGRATION, "integration tests are opt-in")
 class TestWorkerEndToEnd(TestCase):