refactor(minimal-worker): address PR review — typed schemas, env file, simpler loop

Apply feedback from PR #1252 review:

- Consolidate schemas into `processing_services/minimal/api/schemas.py`. v1
  push and v2 worker now share a single source of truth; the separate
  `worker/schemas.py` mirror is removed. Client/runner/register/loop use real
  Pydantic types in signatures (PipelineProcessingTask, PipelineTaskResult,
  ProcessingServiceClientInfo, AsyncPipelineRegistrationRequest).

- Move defaults into `processing_services/minimal/.env.dev`. Strip the
  hard-coded `os.environ.get(..., "<default>")` fallbacks in register.py and
  worker_main.py; replace with `os.environ[...]`. The inline environment
  block in docker-compose.yml becomes `env_file:`.

- Replace the loop's job→slug reverse-lookup with per-slug iteration: the
  outer loop variable IS the slug, so `_slug_for_job` is gone.

- Fix register.py docstring: the PS is identified by the Authorization
  header's user (main) or the API key (#1194), not by
  `processing_service_name`. That field just labels the DB row
  `get_or_create(name=...)` lands on.

- Import pipelines directly in register.py instead of HTTP-GETting
  /info from the co-located FastAPI. Removes the FastAPI readiness wait
  entirely, and means register.py works in MODE=worker (where FastAPI isn't
  running at all).

Co-Authored-By: Claude <noreply@anthropic.com>

Author: mihow

processing_services/docker-compose.yml

@@ -10,17 +10,11 @@ services:
       - minio:host-gateway
     networks:
       - antenna_network
-    # MODE=api+worker runs both the v1 FastAPI service and the v2 worker
-    # poll loop in one container. start.sh handles sequencing: FastAPI first
-    # (for register.py to read /info), then register.py, then worker loop.
-    # For pure v1 testing, set MODE=api. For pure v2, set MODE=worker.
-    environment:
-      MODE: api+worker
-      ANTENNA_API_URL: http://django:8000
-      ANTENNA_DEFAULT_PROJECT_NAME: Default Project
-      ANTENNA_USER: antenna@insectai.org
-      ANTENNA_PASSWORD: localadmin
-      ANTENNA_SERVICE_NAME: minimal-worker-dev
+    # start.sh dispatches on MODE: api | worker | api+worker. For api+worker,
+    # start.sh runs FastAPI first (for register.py to read /info), then
+    # register.py, then the worker poll loop.
+    env_file:
+      - ./minimal/.env.dev
 
   ml_backend_example:
     build:

processing_services/minimal/.env.dev

@@ -0,0 +1,29 @@
+# Dev defaults for the minimal processing service container.
+#
+# Loaded via `env_file:` in processing_services/docker-compose.yml when you run
+# the stub locally against a fresh Antenna stack. These are the same creds
+# seeded by ami/main/management/commands/ensure_default_project.py and
+# .envs/.local/.django — intentionally not secret.
+#
+# Copy to .env.minimal-worker and edit to point at a different Antenna
+# deployment. CI loads this same file.
+
+# ── Container mode ──────────────────────────────────────────────────────────
+# api          FastAPI only (default in Dockerfile; used by docker-compose.ci.yml)
+# worker       poll loop only
+# api+worker   both, plus register.py on boot — used by local dev compose below
+MODE=api+worker
+
+# ── Antenna target ─────────────────────────────────────────────────────────
+ANTENNA_API_URL=http://django:8000
+ANTENNA_DEFAULT_PROJECT_NAME=Default Project
+ANTENNA_SERVICE_NAME=minimal-worker-dev
+
+# ── Auth (user/password fallback; overridden by ANTENNA_API_KEY / _AUTH_TOKEN if set) ──
+ANTENNA_USER=antenna@insectai.org
+ANTENNA_PASSWORD=localadmin
+
+# ── Worker tuning ──────────────────────────────────────────────────────────
+WORKER_POLL_INTERVAL_SECONDS=2.0
+WORKER_BATCH_SIZE=4
+WORKER_REQUEST_TIMEOUT_SECONDS=30

processing_services/minimal/api/schemas.py

@@ -291,3 +291,67 @@ class ProcessingServiceInfoResponse(pydantic.BaseModel):
     #    default=list,
     #    examples=[RANDOM_BINARY_CLASSIFIER],
     # )
+
+
+# -----------------------------------------------------------------------------
+# v2 async worker schemas
+#
+# Mirror of the relevant classes in ami/ml/schemas.py. Kept in this package so
+# both the v1 FastAPI side (push / `/process`) and the v2 worker side
+# (pull / polling) share a single source of truth. When Antenna evolves the
+# canonical schemas, keep these in sync — field-for-field parity matters for
+# correct JSON round-trips.
+# -----------------------------------------------------------------------------
+
+
+class PipelineResultsError(pydantic.BaseModel):
+    """Error result when pipeline processing fails for a single task."""
+
+    error: str
+    image_id: str | None = None
+
+
+class PipelineProcessingTask(pydantic.BaseModel):
+    """A single image task reserved from the async job queue.
+
+    `reply_subject` is the NATS subject Antenna ACKs on when the result comes
+    back — the worker must round-trip it verbatim in the matching
+    PipelineTaskResult.
+    """
+
+    id: str
+    image_id: str
+    image_url: str
+    reply_subject: str | None = None
+
+
+class TasksResponse(pydantic.BaseModel):
+    """Response body of `POST /api/v2/jobs/{id}/tasks/`."""
+
+    tasks: list[PipelineProcessingTask] = []
+
+
+class PipelineTaskResult(pydantic.BaseModel):
+    """Result of processing a single PipelineProcessingTask."""
+
+    reply_subject: str
+    result: PipelineResultsResponse | PipelineResultsError
+
+
+class ProcessingServiceClientInfo(pydantic.BaseModel):
+    """Identity metadata sent by a processing service worker.
+
+    A ProcessingService in the DB may have multiple physical workers running
+    simultaneously; this lets the server distinguish them. Fields are
+    intentionally open — processing services can send any useful key/value
+    pairs (hostname, software version, pod name, etc).
+    """
+
+    model_config = pydantic.ConfigDict(extra="allow")
+
+
+class AsyncPipelineRegistrationRequest(pydantic.BaseModel):
+    """Body for `POST /api/v2/projects/{id}/pipelines/` from an async processing service."""
+
+    processing_service_name: str
+    pipelines: list[PipelineConfigResponse] = []

processing_services/minimal/register.py

@@ -1,70 +1,72 @@
 """
 Self-register this processing service's pipelines with Antenna.
 
-Modeled on the PR #1194 version (which targets API-key auth once merged). For
-now this targets main, which still uses user-token auth and requires
-`processing_service_name` in the registration body.
-
-Auth priority:
-  1. ANTENNA_API_KEY set → use `Api-Key <key>` (TODO: activate once PR #1194 merges)
-  2. ANTENNA_API_AUTH_TOKEN set → use `Token <token>`
-  3. ANTENNA_USER / ANTENNA_PASSWORD → log in, use the returned token.
-
-Self-provisioning (option 3) matches the local-dev and CI defaults baked into
-.envs/.local/.django (antenna@insectai.org / localadmin).
-
-Environment:
-  ANTENNA_API_URL          Base URL (e.g. http://django:8000)
-  ANTENNA_PROJECT_ID       Project PK OR ANTENNA_DEFAULT_PROJECT_NAME to resolve to one.
-  ANTENNA_DEFAULT_PROJECT_NAME  Fallback lookup by name (default: "Default Project")
-  ANTENNA_SERVICE_NAME     ProcessingService name (default: minimal-worker-<hostname>)
-  ANTENNA_API_KEY          Optional (future, PR #1194 path)
-  ANTENNA_API_AUTH_TOKEN   Optional static token (skips login)
-  ANTENNA_USER             Fallback login email (default: antenna@insectai.org)
-  ANTENNA_PASSWORD         Fallback login password (default: localadmin)
+What this does, in order:
+  1. Resolve an Authorization header (env var or fallback login).
+  2. Resolve the target project id (env var, else look up by name).
+  3. Fetch our own /info to get the list of pipelines this container serves.
+  4. POST that list to `/api/v2/projects/{id}/pipelines/` so Antenna knows which
+     async pipelines this ProcessingService can handle.
+
+About identity: on main, the server looks up / creates a `ProcessingService`
+record by the `processing_service_name` field in the request body, and grants
+write access based on the Authorization header's user. PR #1194 changes that
+to use API keys — the PS record is derived from the key itself and
+`processing_service_name` is no longer sent. We tolerate both by sending
+`processing_service_name` now; #1194-enabled Antenna will ignore the field and
+pick the PS from the key.
+
+Env vars are read via `os.environ[...]` without fallbacks — the .env file is
+expected to provide them. See `processing_services/.env.example`.
 """
 
 import logging
 import os
-import platform
-import socket
 import sys
 import time
 
 import requests
+from api.api import pipelines as pipeline_classes  # type: ignore[import-not-found]
+from api.schemas import (  # type: ignore[import-not-found]
+    AsyncPipelineRegistrationRequest,
+    PipelineConfigResponse,
+    ProcessingServiceClientInfo,
+)
 
 logger = logging.getLogger(__name__)
 
 MAX_RETRIES = 20
 RETRY_DELAY = 3  # seconds
 
-DEFAULT_USER = "antenna@insectai.org"
-DEFAULT_PASSWORD = "localadmin"
-DEFAULT_PROJECT_NAME = "Default Project"
-LOCAL_INFO_URL = "http://localhost:2000/info"
-LOCAL_LIVEZ_URL = "http://localhost:2000/livez"
+CACHED_AUTH_HEADER_PATH = "/tmp/antenna_auth_header"
 
 
-def get_client_info() -> dict:
-    """Identity metadata sent to Antenna.
+def get_client_info() -> ProcessingServiceClientInfo:
+    """Identity metadata sent to Antenna in the registration body.
 
-    Extra keys are allowed by the ProcessingServiceClientInfo schema (Config.extra = "allow"),
-    so it's fine to add more here; main's registration endpoint currently ignores this field
-    and PR #1194 reads it.
+    `ProcessingServiceClientInfo` has `extra="allow"`, so any keys here are
+    forwarded verbatim. On main the registration serializer ignores unknown
+    fields; PR #1194 consumes this field.
     """
-    return {
-        "hostname": socket.gethostname(),
-        "software": "antenna-minimal-worker",
-        "version": "0.1.0",
-        "platform": platform.platform(),
-    }
+    import platform
+    import socket
+
+    return ProcessingServiceClientInfo.model_validate(
+        {
+            "hostname": socket.gethostname(),
+            "software": "antenna-minimal-worker",
+            "version": "0.1.0",
+            "platform": platform.platform(),
+        }
+    )
 
 
 def auth_header() -> dict[str, str] | None:
     """Pick an auth header based on what env vars are set, or None to trigger login flow."""
     api_key = os.environ.get("ANTENNA_API_KEY")
     if api_key:
-        # TODO(PR #1194): Api-Key auth is enabled on Antenna once #1194 merges.
+        # PR #1194 path. Harmless on main — main ignores unknown auth schemes
+        # and falls through to the next header, which we don't send.
         return {"Authorization": f"Api-Key {api_key}"}
 
     token = os.environ.get("ANTENNA_API_AUTH_TOKEN")
@@ -93,7 +95,7 @@ def resolve_project_id(api_url: str, headers: dict[str, str]) -> str:
     if explicit:
         return explicit
 
-    name = os.environ.get("ANTENNA_DEFAULT_PROJECT_NAME", DEFAULT_PROJECT_NAME)
+    name = os.environ["ANTENNA_DEFAULT_PROJECT_NAME"]
     resp = requests.get(f"{api_url}/api/v2/projects/", headers=headers, timeout=10)
     resp.raise_for_status()
     for project in resp.json().get("results", []):
@@ -104,57 +106,54 @@ def resolve_project_id(api_url: str, headers: dict[str, str]) -> str:
     raise RuntimeError(f"No project found with name '{name}' — ensure_default_project should have created it")
 
 
-def fetch_own_pipelines() -> list[dict]:
-    resp = requests.get(LOCAL_INFO_URL, timeout=5)
-    resp.raise_for_status()
-    return resp.json().get("pipelines", [])
-
+def fetch_own_pipelines() -> list[PipelineConfigResponse]:
+    """Return the pipeline configs this container serves.
 
-def wait_for_local_server() -> None:
-    for attempt in range(MAX_RETRIES):
-        try:
-            r = requests.get(LOCAL_LIVEZ_URL, timeout=2)
-            if r.status_code == 200:
-                return
-        except (requests.ConnectionError, requests.Timeout):
-            pass
-        logger.info("Waiting for local FastAPI server (%d/%d)", attempt + 1, MAX_RETRIES)
-        time.sleep(RETRY_DELAY)
-    raise RuntimeError("Local FastAPI server did not come up in time")
+    Imported directly from the api module rather than fetched over HTTP from
+    the co-located FastAPI service — register.py runs in the same container,
+    and importing avoids having to wait for FastAPI to be up (which it isn't
+    in MODE=worker).
+    """
+    return [p.config for p in pipeline_classes]
 
 
-def register(api_url: str, project_id: str, headers: dict[str, str], pipelines: list[dict]) -> None:
+def register(
+    api_url: str,
+    project_id: str,
+    headers: dict[str, str],
+    pipelines: list[PipelineConfigResponse],
+) -> None:
     """POST pipelines to the project registration endpoint.
 
-    Body shape for main:
-        {"processing_service_name": str, "pipelines": [...], "client_info": {...}}
-    PR #1194 drops `processing_service_name`. Sending both now is safe because
-    main's serializer ignores unknown fields and #1194's serializer ignores
-    `processing_service_name`.
+    Sends the schema-defined `AsyncPipelineRegistrationRequest` body. We also
+    attach a `client_info` field, which is ignored on main (unknown field) and
+    read by PR #1194.
     """
-    service_name = os.environ.get("ANTENNA_SERVICE_NAME", f"minimal-worker-{socket.gethostname()}")
-    payload = {
-        "processing_service_name": service_name,
-        "pipelines": pipelines,
-        "client_info": get_client_info(),
-    }
+    service_name = os.environ["ANTENNA_SERVICE_NAME"]
+    body = AsyncPipelineRegistrationRequest(
+        processing_service_name=service_name,
+        pipelines=pipelines,
+    ).model_dump(mode="json")
+    body["client_info"] = get_client_info().model_dump(mode="json")
+
     url = f"{api_url}/api/v2/projects/{project_id}/pipelines/"
-    resp = requests.post(url, json=payload, headers=headers, timeout=30)
+    resp = requests.post(url, json=body, headers=headers, timeout=30)
     if resp.status_code in (200, 201):
-        logger.info("Registered %d pipelines as '%s' (project=%s)", len(pipelines), service_name, project_id)
+        logger.info(
+            "Registered %d pipelines as '%s' (project=%s)",
+            len(pipelines),
+            service_name,
+            project_id,
+        )
         return
     raise RuntimeError(f"Registration failed: {resp.status_code} {resp.text}")
 
 
 def main() -> int:
     logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
 
-    api_url = os.environ.get("ANTENNA_API_URL")
-    if not api_url:
-        logger.error("ANTENNA_API_URL not set; skipping registration")
-        return 0
+    api_url = os.environ["ANTENNA_API_URL"]
 
-    wait_for_local_server()
     pipelines = fetch_own_pipelines()
     if not pipelines:
         logger.warning("No pipelines found from local /info; nothing to register")
@@ -163,8 +162,8 @@ def main() -> int:
     # Auth: use explicit header if provided, else log in.
     headers = auth_header()
     if headers is None:
-        email = os.environ.get("ANTENNA_USER", DEFAULT_USER)
-        password = os.environ.get("ANTENNA_PASSWORD", DEFAULT_PASSWORD)
+        email = os.environ["ANTENNA_USER"]
+        password = os.environ["ANTENNA_PASSWORD"]
         # Retry login so we tolerate "Django not up yet".
         for attempt in range(MAX_RETRIES):
             try:
@@ -193,7 +192,7 @@ def main() -> int:
 
     # Cache the resolved id and auth header for worker_main.py to reuse.
     os.environ["ANTENNA_PROJECT_ID"] = project_id
-    with open("/tmp/antenna_auth_header", "w") as f:
+    with open(CACHED_AUTH_HEADER_PATH, "w") as f:
         f.write(next(iter(headers.values())))
 
     for attempt in range(MAX_RETRIES):

processing_services/minimal/start.sh

@@ -28,13 +28,11 @@ start_api() {
 }
 
 start_register() {
-    # register.py self-provisions a ProcessingService and registers this
-    # container's pipelines with Antenna. It has its own retry loop for
-    # "Antenna not up yet", so we don't need to poll for readiness here.
+    # register.py imports pipelines from the api module directly (no HTTP to
+    # self), self-provisions a ProcessingService, and registers pipelines
+    # with Antenna. Its retry loop handles "Antenna not up yet".
     # Skip if registration env vars aren't set — the container still works
     # as a pure v1 push service without them.
-    # Need ANTENNA_API_URL and either an explicit ANTENNA_PROJECT_ID or a
-    # project name (ANTENNA_DEFAULT_PROJECT_NAME) that register.py can look up.
     if [ -z "${ANTENNA_API_URL:-}" ]; then
         echo "[start.sh] Skipping registration (ANTENNA_API_URL not set)"
         return
@@ -62,10 +60,11 @@ case "$MODE" in
         start_worker
         ;;
     api+worker)
+        # FastAPI first so /process is available even while register.py is
+        # still retrying against a not-yet-ready Antenna. register.py doesn't
+        # depend on FastAPI (imports pipeline configs directly), so there's
+        # no startup race here.
         start_api
-        # Give the FastAPI side a beat to bind its port before register.py
-        # tries to GET /info from it. Cheap; register.py also retries.
-        sleep 2
         start_register
         start_worker
         ;;