Repository navigation
New async & distributed ML backend (aka "PSv2") #515
Description
Status update (2026-06-27): PSv2 is now the platform default. #1353 merged & deployed —
async_pipeline_workersdefaults toTrueand all existing projects were migrated to the async/distributed backend. The async path (JobDispatchMode.ASYNC_API— NATS queue, worker pull, result ingestion, cancellation, monitoring) is the common path. Remaining PSv2 items are non-blocking hardening, performance, and feature-completeness follow-ups — see the open PSv2 issues and PSv2 PRs.
@kaviecos and @mihow have designed & written the specifications for a new ML backend that orchestrates multiple types of models by different research teams, across multiple stages of processing, and is horizontally scalable. This expands on the current ML backend API defined here https://ml.dev.insectai.org/ by adding asynchronous processing, a controller & queue system, auth and many other production features.
The initial spec and notes are here, but are being re-written in the Aarhus GitLab wiki as the backend is developed.
https://docs.google.com/document/d/1caKxxfZhWhRi9Jfv9fy5fVeoM9bvhYPJ/
Docs in progress:
https://gitlab.au.dk/ecos/hoyelab/ami/ami-pipeline-controller/-/wikis/Getting-Started
https://gitlab.au.dk/ecos/hoyelab/ami/ami-pipeline-controller/-/wikis/Pipeline-Stages
https://gitlab.au.dk/ecos/hoyelab/ami/ami-pipeline-controller/-/wikis/Architecture-Overview
Known remaining tasks:
- Fix truncated requests for images with many insects #516
- Add more validation responses instead of 500 errors (specifically for the model APIs)
- Add docs for the data flow between stages
- Add docs for registering new consumers for new & existing models (regional classifiers) #517
- Install NewRelic infrastructure monitoring on all servers (get token from @mihow)
- Consider installing NewRelic into actual codebase
- Test initiating requests from AMI application (when Update platform to use upcoming asynchronous ML backend API #505 is complete)
- Implement re-processing of detections
- Status endpoint on the controller, for checking the status of a request #540
- Document possible error responses
- CNN features are stored to a generic "URI" field (could reference a remote storage, or s3:// or vector database or file)
- Add dummy/mockup detector that returns many detections, and no detections
- Consider seeing how many large requests the controller can take and send back (to callback) (load test)
- Pipeline requests should have a cancellation method #539
- For detectionRequests: Change single source image to array of images, even though we are using one
- Add "stageParams": {} to the controller request, which are then passed via URL to the stage implementation requests.
- Update the registry to support a "*" projectID that any project can use
- Think about an endpoint for checking what stages are registered. Use the projectID to determine if it's a private or shared stage implementation.
- Stage Implementation Monitoring #538
- Implement automatic retries when a task fails.

Activity
mihow commentedon Aug 15, 2024
Controller can be tested here:
https://preview.ami.ecoscience.dk/swagger-ui/index.html#/pipeline-controller/createRequest
Example request:
{ "projectId": "ecos", "jobId": "ea12ac70-288c-11ef-9ca5-00155d926c42", "sourceImages": [ { "id": "NScxODE3NzEyMwo=", "url": "https://anon.erda.au.dk/share_redirect/DSWDMAO70L/ias/denmark/DK1/2023_07_05/20230705000135-00-07.jpg", "eventId": "1234" } ], "pipelineConfig": { "stages": [ { "stage": "OBJECT_DETECTION", "stageImplementation": "flatbug" }, { "stage": "CLASSIFICATION", "stageImplementation": "mcc24" } ] }, "callback": { "callbackUrl": "http://127.0.0.1:8080/example/callback", "callbackToken": "1234" } }The callback can be inspected using Ngrok locally:
ngrok http 2222callbackUrlin the sample request to the generated ngrok URLImage that causes a truncated request
https://static.dev.insectai.org/ami-trapdata/Panama/E43B615A/20231113004009-snapshot.jpg
kaviecos commentedon Aug 15, 2024
There are a few Refactorings I'm considering in the PipelineController.
Stagein the PipelineConfig is currently unused by thecontroller. I originally included it because I thought different stages might be handled differently but I think the current solution is cleaner - every stage use the same interface. One problem with handling different stages differently is that one stage implementation may perform multiple operations on the image.ImageCrop- Originally I thought the pipeline asObjectDetection->ImageCrop->Classification. But I realized that we don't need to crop in order to do classification. Actually it performs a lot better if the entire sourceimage is used with bounding boxes. So maybe cropping is post-pipeline step? It can also be added as a stage and set thecropUrlon the detection.mihow commentedon Aug 20, 2024
@kaviecos Have you considered adding a status endpoint on the controller? It would be nice if a callback is missed, or if the job is taking a long time, for the client to be able to request the status of a request. PROCESSING, FAILED, WAITING, 2/6 complete, etc.
Also will you document the types of failures and what those will look like in a callback?
I just added these as subtasks as well.
mihow commentedon Aug 20, 2024
NOTES from call with @mihow and @kaviecos 2024-08-20
docs for OpenAI
https://platform.openai.com/docs/guides/batch/getting-started
mihow commentedon Aug 20, 2024
@kaviecos Let's also keep in mind that scenario where we have stage implementations running behind a firewall and cannot have a publicly accessible endpoint. Our major providers (Compute Canada and the UK's JASMIN) both have way more compute available to us, but not on persistent VMs like we are using now. For a big job like Biodiversa+'s data, we can request a bunch of GPUs and provide the docker container to run via a SLURM job scheduler. Each job can access the internet, so can pull from the queue and send back results. But they won't have a publicly accessible endpoint (unless we can use a tunnel?)
kaviecos commentedon Aug 22, 2024
@mihow The current implementation actually allows for consumers hosted on other servers. The requirements are that they can make an outbound connection to the RabbitMQ server (port 5672). Then all communication with the controller can go through RabbitMQ. Of course they also need to be able to access the source images.
When it comes to monitoring this also means that we need to monitor both the consumers and the stage implementations. And if we cannot make inbound http requests then we need to think that into the monitoring solution. Maybe a push-based solution (last_seen as you mentioned). It would be nice to know more about the restrictions - like, is it even possible to connect to RabbitMQ?
3 remaining items
vanessavmac commentedon Jul 8, 2025
Goal: Have a draft of the a working celery worker by the end of the month (July)
mihow commentedon Jul 8, 2025
Here are some notes from a conversation in Slack
Antenna processes results in the results queue (converts raw results to Occurrences, Taxa, etc.)
mihow commentedon Jul 8, 2025
We can use our existing PipelineRequest & PipelineResponse objects, but use them to pass serialized data to & from the celery workers. Something like:
mihow commentedon Jul 17, 2025
process_pipeline_request(PipelineRequest) -> PipelineRequestin Antennaprocess_images()in antenna to submit PipelineRequest to queue. Callprocess_pipeline_request.apply_async(args=[pipeline_request], queue=queue_name, routing_key=routing_key)(queue name or routing key should be specific to each pipeline)routing_keycan be dynamic? (use the pipeline slug)process_pipeline_requestswith it's dummy task signatureCurrent implementation
New implementation
v1
keep API as-is
add celery to processing service
Antenna calls
process_pipeline_request.apply_async()processing service calls
save_results.apply_async()Antenna saves the results as normal
v2
Processing service can operate without public API
mihow commentedon Jul 17, 2025
v1

v2

vanessavmac commentedon Jul 21, 2025
rename default queue to antenna queue
save results can go to the antenna queue
is a separate queue required to ensure save results is high priority?
Next steps:
processing.pywhich contains the process logic; don't make http request, just call a functionMLJob.run()shouldn't endlessly poll for status; we should be able to callcheck_ml_job_status()anytime to quickly check status oftask_ids--> See comment# TLDR; MLJob.run() just submits processing tasks + schedules celery beat; celery beat responsible for everything aftera. Add a new field for storing all task ids (
process_pipeline_requestandsave_results) to the MLJob model. An array field of strings calledtask_idsb. Create a separate function (
check_ml_job_status) which celery beat runs periodicallyc.
check_ml_job_statusshould still watch the status ofprocess_pipeline_requesttasks; but only use the status to update the UI; shouldn't do anything with the results.d. The processing service should produce the
save_resultstasks, queue it into the antenna queue (maybe change this to ensure higher priority?), and then the celerybeat watches for these tasks. Celerybeat will watch for any tasks with a parent id equal totasks_ids(celery might already attach parent ids).GOAL for next meeting: start a job, process images, and see the job status change to success (only after all sub-tasks
process_pipeline_requestandsave_resultsall complete)vanessavmac commentedon Jul 30, 2025
Next Tasks (08/04)
Edge cases to consider:
[-]New async & distributed ML backend [/-][+]New async & distributed ML backend (aka "PSv2")[/+]mihow commentedon Jun 16, 2026
Claude says: (posted on Michael's behalf — consolidated status map of the PSv2 effort, assembled from the full
PSv2-labeled issue/PR set.)PSv2 — implementation history and current status
This is a consolidated map of the async + distributed ML backend ("PSv2") work: how it was built, what is fixed, what was attempted and abandoned, what is still open, and the root issues that look most critical before async processing can become the default path. It was assembled by reading every issue and PR carrying the new
PSv2label (~120 items) and following the cross-references between them. Where a claim is an interpretation rather than a settled fact, it is flagged as such.All PSv2 work is now tagged with the
PSv2label for easier tracking.How we got here (the arc)
Phase 0 — controller/push era (spec). PSv2 began as an external pipeline controller designed at Aarhus (this issue, plus #505, #516, #517, #538, #539, #540). It was push-based and lived outside Antenna. Cancellation (#539) and stage monitoring (#538) were specified here and are still open — they resurface years later.
Phase 1 — in-Antenna push baseline. #632 landed V1: register a processing service, pick the lowest-latency online one, push work to it. #1060 raised the RabbitMQ ack timeout so long sync jobs stopped dying at 30 minutes. This is the synchronous baseline PSv2 replaces.
Phase 2 — the pull pivot (design). A deliberate move from push to pull: #968 (pull API spec), #969 (reference impl), #970 (implement), #971 (scheduler). Rationale, recorded in #987: per-task tracking for resilience, horizontal scale, and — critically — no public endpoint required, so workers can sit behind HPC/university firewalls. #1046 landed the API scaffold with the
tasks()/result()endpoints stubbed.Phase 3 — core implementation. #987 is the keystone: NATS JetStream chosen (after rejecting RabbitMQ/Beanstalkd for lacking the disconnected pull/ack semantics we need), task state + queue managers, real queue/pull/ack/result. Shipped disabled by default behind the
async_pipeline_workersflag for extended testing. #1109 added the NATS host alias for deployment.Phase 4 — registration, state decoupling, cleanup. Pipeline registration for pull mode (#1086 → #1076,
endpoint_urlnullable). Decoupling the Celery task state from the job state — the recurring hard problem — seeded by #1062, framed in #1084, solved in #1114 (a job is SUCCESS only when all stages complete, not when the queueing task returns), hardened in #1125. Routing via the newJob.dispatch_modeenum (#1111 → #1118). Resource cleanup on completion (#1083 → #1113).Phase 5 — progress, hardening, cancellation. Progress parity with the sync path (#1121). Integration-test hardening after real worker flooding (#1135, #1142, with test infra #1129/#1143). Service "live" status groundwork (#1122 → #1146). Cancellation was attempted in #1144 but closed unmerged and reworked later.
What is fixed and merged
dispatch_moderouting (fix: PSv2: Workers should not try to fetch tasks from v1 jobs #1118), resource cleanup on completion (PSv2: Implement queue clean-up upon job completion #1113), Celery-vs-job state decoupling (fix: Properly handle async job state with celery tasks #1114/PSv2 cleanup: use is_complete() and dispatch_mode in job progress handler #1125).async_apijobs killed by transient Redis errors duringupdate_state—RedisErrorand "state actually missing" are conflated into a single fatal path #1219).What was attempted and then abandoned, reverted, or superseded
allow_reconnect=False) route instead.A pattern worth naming: several fixes introduced the next bug. #1162 (cancellation + fail-on-missing-Redis) created the "transient Redis error kills the job" path that #1219/#1231 then had to repair, and whose residue is still open (#1232, #1241). #1234 dropped
max_deliverfrom 5 to 2, which sharpened the lost-image symptom that #1244 then had to reconcile. The slice has repeatedly fixed a symptom and surfaced the next one.What is still open
The never-closed roots (each has had multiple PRs chip at sub-cases without closing the general guarantee):
max_deliverwithout results. Five PRs mitigate it (Support for NATS dead letter queue #1175/feat(jobs): schedule periodic stale-job check #1227/feat(jobs): drain zombie NATS streams in periodic health check #1239/fix(jobs): mark failed or lost images as failed so job can be marked complete #1244/feat(jobs): show worker availability in async job logs #1240); the detection mechanism as specified and the root cause (Decouple ML result ingestion from the HTTP POST: large bodies, no durable handoff, no retry-record path #1223, the large-result-POST → 503 → redelivery cascade) are unbuilt.The newly-named concurrency root:
progress/status. fix(jobs): fixes for concurrent ML processing jobs #1261 droppedselect_for_updateto kill the row-lock contention that was serializing every result handler (feat(job): refactor job logging so it isn't a bottleneck #1256) and causing redelivery storms. That fixed contention but reintroduced a classic lost-update race: two handlers read the same snapshot, both write the whole JSONB blob back, one update is lost. This is the hinge — we traded a contention problem for a correctness problem — and it is the common cause behind the stuck-STARTED display bugs, the false REVOKEs of completed jobs (fix(jobs): fix dangling jobs from going to revoked #1276 mitigates), the SUCCESS→STARTED resurrection, and the scheduling crossover below.The crossover that makes "just a display bug" unsafe: status is also a scheduling input, not only a UI output. A job wrongly left non-terminal stays claimable, so the worker
/nextendpoint keeps handing it out and starves newer jobs. #1282 (CANCELLED jobs leaking through/next) is exactly this with a different trigger. The/nextreshape in #1265 must agree with #1337 on which states are claimable, or the crossover reopens.Open PRs gating the default flip: #1276 (reaper reads Redis instead of clobbered progress — mergeable), #1279 (propagate pipeline config to pull-mode tasks), #1312 (captures marked processed with zero detections — currently conflicting), #1324 (dispatch/cancel hardening — see caution below), plus #1287 (savepoint around the log insert) and the not-yet-PR items #1265 and #1288.
Other open blockers and debt: #1263 (non-superuser worker tokens 403 on
/tasks//result, currently masked by making the worker a superuser), #1302 (Postgres connection-slot exhaustion under load — infra), #1283 (NATS stream/consumer not cleaned on REVOKE/restart → duplicate tasks), #1247 (one worker batch stranded after long jobs), the/tasksendpoint cluster (#1141/#1152/#1153/#1182), and the residual Redis-state handling (#1232/#1241).A caution on #1324: it is reasonable in intent but stacks two risky behaviors. Its
acks_late+reject_on_worker_lostredelivery-safety relies on readingjob.statusin an early guard — the very field #1337 says is unreliable under the clobber race. AndCELERY_WORKER_POOL_OPTIMIZATION = "fair"is not read by Celery from settings (fair scheduling is the-O fairCLI flag), so as written that half is a no-op while the cross-container head-of-line lever the originating incident hit is left commented out. Worth gating behind the #1337 fix (at least its conditional-status-write piece) rather than merging as standalone hardening. Also note its RabbitMQacks_lateinteraction withconsumer_timeoutdiffers by environment.Most critical remaining root issues, ranked
/next, starves new jobs. A concrete, deterministic production-impacting instance of Saving job progress concurrently is the root of multiple issues related to incorrect job statuses #1337's crossover; small fix, but must align with Add GET /jobs/next action; deprecate list(ids_only=1) claim-semantics #1265.A separate root-cause issue (#1337) now tracks item 1 in detail with proposed fix directions; #1285 tracks the diagnostic triangulation across the DB/Redis/NATS sources of truth that would make this whole class easier to debug.
mihow commentedon Jun 16, 2026
Claude says: A proposed order for the remaining PSv2 root work, before async becomes the default path. The principle behind the ordering: fix the things that corrupt or lose state before the things that re-architect or optimize, and land each root fix on prod + demo + staging before relying on it. Sequence, not a hard dependency graph — several items can run in parallel where noted. Framed as a proposal to discuss, not a settled plan.
Stage 0 — stop the bleeding (correctness, small, parallelizable)
These are deterministic, low-risk, and each removes a way the system currently loses or corrupts state. They gate the flip.
/nextstarvation crossover. Small PR, ready. Everything else that readsjob.status(including Improve celery task dispatch and cancellation to prevent stuck jobs #1324's redelivery guard) is safer once this is in./next. The concrete production instance of the crossover. Tiny fix (status__in+ add CANCELLED/terminal to the claim filter). Must be designed together with Add GET /jobs/next action; deprecate list(ids_only=1) claim-semantics #1265 so the claimable-state set is defined once.valid()semantic change, plus the cleanup-command predicate patch.Stage 1 — define the claim contract (the endpoint the workers pull from)
GET /jobs/next/, removing the three temporary hacks fix(jobs): fixes for concurrent ML processing jobs #1261 left inviews.py. This is where the "which states are claimable" decision lives; Saving job progress concurrently is the root of multiple issues related to incorrect job statuses #1337 Layer 1 and CANCELLED jobs leak through /next filter, starve newer async_api jobs #1282 must agree with it. Doing this after Stage 0 means the claim filter is written once, against the corrected state machine, rather than patched twice.Stage 2 — the structural state fix (needs design first)
Stage 3 — durability and the hang/stranding class
max_deliverexhaustion. This is the architectural fix under PSv2: Async jobs hang forever when NATS tasks exhaust max_deliver without posting results #1168; it should come after the state store is settled because a durable result-handoff path interacts with how per-task state is recorded.Cross-cutting (run alongside, not blocking)
test_ml_job_e2e+ chaos), zombie-stream sweep clean under that load, and the broker/Redis config confirmed per environment. Idle-green is not load-green.A note on #1324
Reasonable in intent but should be gated behind #1337 Layer 1 (its
acks_lateredelivery guard readsjob.status), and its fair-scheduler setting is a no-op as written (fair scheduling is the-O fairCLI flag, not a settings value). TheJob.cancel()rewrite in it is sound and could be split out and merged early. Itsacks_lateinteraction with RabbitMQconsumer_timeoutalso differs by environment and needs the timeout raised where it currently uses the 30-minute default.