feat: fetch jobs for all pipelines in a single API request (#114)

* feat: fetch jobs for all pipelines in a single API request

Replace the per-pipeline loop with a single call using
pipeline__slug__in filter. This reduces N API requests per poll
cycle to 1, regardless of how many pipelines the worker handles.

The Antenna MinimalJobSerializer now returns pipeline_slug so the
worker knows which pipeline each job belongs to.

Co-Authored-By: Claude <[email protected]>

* fix: update tests for pipeline__slug__in and add empty guard

- Guard against empty pipeline_slugs list in get_jobs()
- Update mock server to accept pipeline__slug__in parameter
- Include pipeline_slug in AntennaJobListItem construction
- Update test assertions for new tuple return type

Co-Authored-By: Claude <[email protected]>

---------

Co-authored-by: Claude <[email protected]>

Author: mihow

trapdata/antenna/client.py

@@ -29,27 +29,29 @@ def get_full_service_name(service_name: str) -> str:
 def get_jobs(
     base_url: str,
     auth_token: str,
-    pipeline_slug: str,
+    pipeline_slugs: list[str],
     processing_service_name: str,
-) -> list[int]:
-    """Fetch job ids from the API for the given pipeline.
+) -> list[tuple[int, str]]:
+    """Fetch job ids from the API for the given pipelines in a single request.
 
-    Calls: GET {base_url}/jobs?pipeline__slug=<pipeline>&ids_only=1&processing_service_name=<name>
+    Calls: GET {base_url}/jobs?pipeline__slug__in=<slugs>&ids_only=1&processing_service_name=<name>
 
     Args:
         base_url: Antenna API base URL (e.g., "http://localhost:8000/api/v2")
         auth_token: API authentication token
-        pipeline_slug: Pipeline slug to filter jobs
+        pipeline_slugs: List of pipeline slugs to filter jobs
         processing_service_name: Name of the processing service
 
     Returns:
-        List of job ids (possibly empty) on success or error.
+        List of (job_id, pipeline_slug) tuples (possibly empty) on success or error.
     """
     with get_http_session(auth_token) as session:
         try:
+            if not pipeline_slugs:
+                return []
             url = f"{base_url.rstrip('/')}/jobs"
             params = {
-                "pipeline__slug": pipeline_slug,
+                "pipeline__slug__in": ",".join(pipeline_slugs),
                 "ids_only": 1,
                 "incomplete_only": 1,
                 "processing_service_name": processing_service_name,
@@ -61,7 +63,7 @@ def get_jobs(
 
             # Parse and validate response with Pydantic
             jobs_response = AntennaJobsListResponse.model_validate(resp.json())
-            return [job.id for job in jobs_response.results]
+            return [(job.id, job.pipeline_slug) for job in jobs_response.results]
         except requests.RequestException as e:
             logger.error(f"Failed to fetch jobs from {base_url}: {e}")
             return []

trapdata/antenna/schemas.py

@@ -25,6 +25,7 @@ class AntennaJobListItem(pydantic.BaseModel):
     """A single job item from the Antenna jobs list API response."""
 
     id: int
+    pipeline_slug: str
 
 
 class AntennaJobsListResponse(pydantic.BaseModel):

trapdata/antenna/tests/antenna_api_server.py

@@ -29,15 +29,15 @@
 
 @app.get("/api/v2/jobs")
 def get_jobs(
-    pipeline__slug: str,
-    ids_only: int,
-    incomplete_only: int,
+    pipeline__slug__in: str = "",
+    ids_only: int = 1,
+    incomplete_only: int = 1,
     processing_service_name: str = "",
 ):
     """Return available job IDs.
 
     Args:
-        pipeline__slug: Pipeline slug filter
+        pipeline__slug__in: Comma-separated pipeline slugs filter
         ids_only: If 1, return only job IDs
         incomplete_only: If 1, return only incomplete jobs
         processing_service_name: Name of the processing service making the request
@@ -48,9 +48,17 @@ def get_jobs(
     global _last_get_jobs_service_name
     _last_get_jobs_service_name = processing_service_name
 
+    # Determine pipeline slug for response (use first slug from filter)
+    slugs = (
+        [s for s in pipeline__slug__in.split(",") if s] if pipeline__slug__in else []
+    )
+    default_slug = slugs[0] if slugs else "test_pipeline"
+
     # Return all jobs in queue (for testing, we return all registered jobs)
     job_ids = list(_jobs_queue.keys())
-    results = [AntennaJobListItem(id=job_id) for job_id in job_ids]
+    results = [
+        AntennaJobListItem(id=job_id, pipeline_slug=default_slug) for job_id in job_ids
+    ]
     return AntennaJobsListResponse(results=results)
 
 

trapdata/antenna/tests/test_worker.py

@@ -220,10 +220,13 @@ def test_returns_job_ids(self):
 
         with patch_antenna_api_requests(self.antenna_client):
             result = get_jobs(
-                "http://testserver/api/v2", "test-token", "moths_2024", "Test Worker"
+                "http://testserver/api/v2",
+                "test-token",
+                ["moths_2024"],
+                "Test Worker",
             )
 
-        assert result == [10, 20, 30]
+        assert [job_id for job_id, _ in result] == [10, 20, 30]
         assert antenna_api_server.get_last_get_jobs_service_name() == "Test Worker"
 
 
@@ -460,9 +463,13 @@ def test_full_workflow_with_real_inference(self):
             assert success is True
 
             # Step 2: Get jobs
-            job_ids = get_jobs(
-                "http://testserver/api/v2", "test-token", pipeline_slug, "Test Worker"
+            jobs = get_jobs(
+                "http://testserver/api/v2",
+                "test-token",
+                [pipeline_slug],
+                "Test Worker",
             )
+            job_ids = [job_id for job_id, _ in jobs]
             assert 200 in job_ids
             assert antenna_api_server.get_last_get_jobs_service_name() == "Test Worker"
 

trapdata/antenna/worker.py

@@ -90,32 +90,33 @@ def _worker_loop(gpu_id: int, pipelines: list[str]):
         # These should probably come from a dedicated endpoint and should preempt batch jobs under the assumption that they
         # would run on the same GPU.
         any_jobs = False
-        for pipeline in pipelines:
-            logger.info(f"[GPU {gpu_id}] Checking for jobs for pipeline {pipeline}")
-            jobs = get_jobs(
-                base_url=settings.antenna_api_base_url,
-                auth_token=settings.antenna_api_auth_token,
-                pipeline_slug=pipeline,
-                processing_service_name=full_service_name,
+        logger.info(
+            f"[GPU {gpu_id}] Checking for jobs for pipelines: {', '.join(pipelines)}"
+        )
+        jobs = get_jobs(
+            base_url=settings.antenna_api_base_url,
+            auth_token=settings.antenna_api_auth_token,
+            pipeline_slugs=pipelines,
+            processing_service_name=full_service_name,
+        )
+        for job_id, pipeline in jobs:
+            logger.info(
+                f"[GPU {gpu_id}] Processing job {job_id} with pipeline {pipeline}"
             )
-            for job_id in jobs:
-                logger.info(
-                    f"[GPU {gpu_id}] Processing job {job_id} with pipeline {pipeline}"
+            try:
+                any_work_done = _process_job(
+                    pipeline=pipeline,
+                    job_id=job_id,
+                    settings=settings,
+                    processing_service_name=full_service_name,
                 )
-                try:
-                    any_work_done = _process_job(
-                        pipeline=pipeline,
-                        job_id=job_id,
-                        settings=settings,
-                        processing_service_name=full_service_name,
-                    )
-                    any_jobs = any_jobs or any_work_done
-                except Exception as e:
-                    logger.error(
-                        f"[GPU {gpu_id}] Failed to process job {job_id} with pipeline {pipeline}: {e}",
-                        exc_info=True,
-                    )
-                    # Continue to next job rather than crashing the worker
+                any_jobs = any_jobs or any_work_done
+            except Exception as e:
+                logger.error(
+                    f"[GPU {gpu_id}] Failed to process job {job_id} with pipeline {pipeline}: {e}",
+                    exc_info=True,
+                )
+                # Continue to next job rather than crashing the worker
 
         if not any_jobs:
             logger.info(