Research worker frameworks for processing service v3 bundled in Antenna

[mihow]

Plain-language summary

Antenna currently distributes ML work to processing workers via NATS JetStream pull-subscriptions. Workers are reference-implemented in ami-data-companion (a standalone Python process). Researchers can also run their own workers — same protocol, same auth, no public endpoint required on their side.

The reference worker has accumulated a long tail of subprocess/memory hygiene issues (RolnickLab/ami-data-companion#140, #145, #144, #147). PR RolnickLab/ami-data-companion#148 ships a near-term fix that addresses the immediate OOM family, but the broader pattern is that we're maintaining custom process-management, GPU memory hygiene, and worker-recycling code that mature OSS handles for free.

This ticket is research and discovery for a "processing service v3" — a worker SDK/runtime bundled in the Antenna repo, designed so the same protocol works for:

• a Jupyter notebook function (researcher iterates on a new model),
• a single-GPU lab workstation,
• a SLURM array job on a university HPC cluster (no inbound network access),
• a production multi-GPU service.

Outcome: a recommendation document + proof-of-concept on one or two candidate frameworks. No implementation in this ticket — that lands in follow-up PRs once a direction is picked.

Why now

• PR ami-data-companion#148 fixes the immediate fire but doesn't change the architecture. We're going to keep filing variants of the same bug if we don't migrate the worker internals off bespoke code.
• The "researchers run their own workers" requirement is genuinely uncommon in commercial ML serving (Triton, TorchServe, Ray Serve, KServe, BentoML are all push-based) — but it's exactly what Antenna's collaborative-research model needs.
• The NATS-based federation boundary is the right choice and not what we want to change. The question is what runs inside the worker.

Requirements the current design solves (don't break these)

1. Federation: anybody can show up with their own worker, no Antenna-controlled fleet.
2. Firewall-friendly: workers initiate outbound connection only.
3. Heterogeneous shapes: same protocol from Jupyter to HPC.
4. Open ML problem: workers ship custom models, custom preprocessing, custom collation. Antenna provides the platform, researchers contribute the science.
5. Strong API contracts (OpenAPI + pydantic): keeps worker/server loose-coupled across model evolution.

What needs deciding

What goes inside the worker boundary, behind the NATS protocol? Three layers to evaluate:

1. Process management + recycling: kills the leak treadmill (Bump sqlalchemy from 2.0.8 to 2.0.12 in /backend #140, Bump ra-data-simple-rest from 4.9.0 to 4.10.1 in /frontend #145, Add favicons #147)

   • multiprocessing.Pool with maxtasksperchild (stdlib only, cheapest)
   • concurrent.futures.ProcessPoolExecutor
   • Celery worker with --max-tasks-per-child (custom NATS broker required)
   • Ray actors with max_restarts / max_calls_per_actor

2. GPU memory + batching:

   • Manual (current path)
   • NVIDIA Triton Inference Server (in-process, Python backend)
   • Ray Serve replicas

3. Pipeline orchestration (detect → classify → species-classify):

   • Sequential Python (current)
   • Ray DAGs
   • Triton ensembles

Candidate frameworks to evaluate

For each: federation fit, researcher onboarding cost, GPU memory handling, recycling story, integration cost with NATS, license, community health.

• stdlib multiprocessing.Pool — simplest, no new deps, solves Add favicons #147 in ~20 LoC. Doesn't help GPU memory.
• Ray + Ray Serve — Python-native actor model, GPU pinning, automatic restart. Researchers write @ray.remote classes. Adds Ray as a heavy dep.
• NVIDIA Triton (in-process Python backend) — best-in-class GPU mem mgmt, dynamic batching, instance pools. Researchers ship Triton model repos (pbtxt config). Steepest learning curve, highest ceiling.
• Celery with custom NATS transport — maximum reuse of Antenna's existing Celery stack. Custom broker transport is a maintenance commitment.
• Temporal — durable workflow engine with pull workers. Bigger lift; only worth it if we want long-running multi-stage workflows.

Discovery scope

• Write a comparison matrix (federation fit, researcher onboarding, GPU mem handling, recycling, NATS integration cost, license, ecosystem health).
• Build a thin proof-of-concept worker for the top one or two candidates that pulls a real NATS task, runs quebec_vermont_moths_2023, posts results.
• Measure: 1000-image batch RSS curve, recovery from a simulated OOM, time-to-first-result for a "Jupyter notebook" shape.
• Write recommendation doc with phased migration path.

Out of scope (for this ticket)

• Picking one and migrating.
• Rewriting ami-data-companion.
• API contract changes between worker and Antenna server.
• Changes to NATS broker setup.

Pluggable-backend SDK pattern (worth evaluating during discovery)

A direction worth testing in the proof-of-concept: a single SDK with backends picked per deployment:

# Jupyter / quick research
@antenna.worker(pipeline="my_model_v2", backend="inline")
def process(images): return run_my_model(images)

# Lab workstation
@antenna.worker(pipeline="my_model_v2", backend="pool", max_tasks_per_child=50)
def process(images): return run_my_model(images)

# Production cluster
@antenna.worker(pipeline="my_model_v2", backend="ray", num_replicas=4)
class MyModel:
    def __init__(self): self.model = load_my_model()
    def __call__(self, images): return self.model(images)

All backends speak NATS upward; Antenna server is unaware which is in use.

References

• DataLoader subprocess deadlock: MainThread hangs forever on queue.get after shared-memory unlink failure ami-data-companion#140 — DataLoader deadlock
• DataLoader pipe FD leak: ~10k pipes/day accumulate, worker hits EMFILE silently after ~6 days uptime ami-data-companion#145 — pipe FD leak
• LRU eviction for idle pipeline models + audit residual RSS growth after #118 ami-data-companion#144 — residual RSS audit
• Add SIGUSR1 drain-and-exit handler to worker; enable safer auto-recycle ami-data-companion#147 — drain-and-exit recycler
• fix(dataloader): switch to forkserver context + add timeout + explicit teardown (closes #140, partial #145) ami-data-companion#148 — current OOM fix (does not change architecture)