Fix split-topology result retrieval

nemo_retriever/helm/templates/deployment.yaml

@@ -104,6 +104,12 @@ spec:
           env:
             - name: NEMO_RETRIEVER_SERVICE_CONFIG
               value: /etc/nemo-retriever/retriever-service.yaml
+            {{- if .Values.retrieverResults.enabled }}
+            - name: NEMO_RETRIEVER_RESULTS_DIR
+              value: {{ .Values.retrieverResults.mountPath | quote }}
+            - name: NEMO_RETRIEVER_RESULTS_TTL_SECONDS
+              value: {{ .Values.retrieverResults.ttlSeconds | quote }}
+            {{- end }}
             - name: NGC_API_KEY
               valueFrom:
                 secretKeyRef:
@@ -272,6 +278,14 @@ spec:
           env:
             - name: NEMO_RETRIEVER_SERVICE_CONFIG
               value: /etc/nemo-retriever/retriever-service.yaml
+            {{- if $.Values.retrieverResults.enabled }}
+            - name: NEMO_RETRIEVER_RESULTS_TTL_SECONDS
+              value: {{ $.Values.retrieverResults.ttlSeconds | quote }}
+            {{- if eq $role "gateway" }}
+            - name: NEMO_RETRIEVER_RESULTS_DIR
+              value: {{ $.Values.retrieverResults.mountPath | quote }}
+            {{- end }}
+            {{- end }}
             - name: POD_IP
               valueFrom:
                 fieldRef:
@@ -316,6 +330,10 @@ spec:
             - name: config
               mountPath: /etc/nemo-retriever
               readOnly: true
+            {{- if and $.Values.retrieverResults.enabled (eq $role "gateway") }}
+            - name: retriever-results
+              mountPath: {{ $.Values.retrieverResults.mountPath }}
+            {{- end }}
             - name: tmp
               mountPath: /tmp
           {{- if $svc.startupProbe.enabled }}
@@ -336,5 +354,10 @@ spec:
             name: {{ include "nemo-retriever.role.configMapName" (dict "context" $ "role" $role) }}
         - name: tmp
           emptyDir: {}
+        {{- if and $.Values.retrieverResults.enabled (eq $role "gateway") }}
+        - name: retriever-results
+          persistentVolumeClaim:
+            claimName: {{ include "nemo-retriever.retrieverResultsPvcName" $ }}
+        {{- end }}
 {{- end }}
 {{- end }}

nemo_retriever/helm/values.yaml

@@ -677,12 +677,22 @@ serviceConfig:
 # Retriever Results volume
 # =============================================================================
 # A PVC mounted at `mountPath` for storing per-job page results as JSON files.
-# The processing pool writes results here and the /v1/ingest/job/{id}/results
-# endpoint reads from this directory instead of SQLite for faster retrieval.
+# The gateway writes worker handoffs here and serves retained rows from this
+# directory instead of SQLite.
+# In split mode, only the gateway mounts this claim. Workers hand retained
+# rows back to the gateway after processing, so realtime and batch replicas
+# remain independently schedulable across nodes with ReadWriteOnce storage.
 #
 # Use local/host-path storage for results to avoid network I/O overhead.
+# The default ReadWriteOnce access mode supports the default single gateway;
+# configuring multiple gateway replicas requires storage and job-state
+# coordination beyond the current split-topology contract.
 retrieverResults:
   enabled: true
+  # Immutable result generations and interrupted temporary files are removed
+  # opportunistically after this age. Keep this at least as long as the job
+  # tracker stale and terminal retention windows.
+  ttlSeconds: 28800
   existingClaim: ""
   storageClass: ""
   accessModes:

nemo_retriever/src/nemo_retriever/service/app.py

@@ -131,6 +131,9 @@ async def _lifespan(app: FastAPI) -> AsyncIterator[None]:
     from nemo_retriever.service.services.pipeline_pool import init_pipeline_pool, shutdown_pipeline_pool
     from nemo_retriever.service.services.proxy import init_proxy, shutdown_proxy
     from nemo_retriever.service.services.sidecar_store import init_sidecar_store, shutdown_sidecar_store
+    from nemo_retriever.service.services.worker_result_store import validate_result_store
+
+    validate_result_store()
 
     if mode in ("gateway", "standalone"):
         app.state.metrics = init_metrics()

nemo_retriever/src/nemo_retriever/service/auth.py

@@ -35,6 +35,15 @@ def _strip_bearer(value: str) -> str:
     return value.strip()
 
 
+def auth_headers(config: AuthConfig) -> dict[str, str]:
+    """Build the configured service-auth header for internal HTTP calls."""
+    token = (config.api_token or "").strip()
+    if not token:
+        return {}
+    value = f"Bearer {token}" if config.header_name.lower() == "authorization" else token
+    return {config.header_name: value}
+
+
 class BearerAuthMiddleware(BaseHTTPMiddleware):
     """Reject requests that don't carry the configured token.
 

nemo_retriever/src/nemo_retriever/service/routers/ingest.py

@@ -18,13 +18,16 @@
 
 import asyncio
 import hashlib
+import ipaddress
 import json
 import logging
 import os
 import uuid
 from datetime import datetime, timezone
 from typing import Any
+from urllib.parse import quote
 
+import httpx
 from fastapi import APIRouter, File, Form, HTTPException, Request, UploadFile
 from fastapi.responses import JSONResponse, Response
 from pydantic import BaseModel, Field, model_validator
@@ -66,13 +69,19 @@
     INGEST_REQUESTS_TOTAL,
 )
 from nemo_retriever.service.services.proxy import get_proxy
+from nemo_retriever.service.services.worker_result_store import (
+    ResultStoreTemporarilyUnavailable,
+    get_result_data,
+    store_result_data,
+)
 from nemo_retriever.service.utils.file_type import (
     FileCategory,
     FileClassifier,
     enforce_media_dependencies,
 )
 
 _RETRY_AFTER_SECONDS = "5"
+_RESULT_RETRY_AFTER_SECONDS = 60
 _DRY_RUN_HEADER = "X-Nemo-Dry-Run"
 _GATEWAY_DOC_ID_HEADER = "X-Gateway-Document-Id"
 _GATEWAY_CALLBACK_HEADER = "X-Gateway-Callback-Url"
@@ -168,6 +177,13 @@ def _work_item_retain_results(request: Request, *, job_id: str | None) -> bool:
     return _job_retain_results(job_id)
 
 
+def _internal_auth_headers(request: Request) -> dict[str, str]:
+    """Return service credentials for pod-to-pod callback traffic."""
+    from nemo_retriever.service.auth import auth_headers
+
+    return auth_headers(request.app.state.config.auth)
+
+
 def _gateway_retain_results_headers(job_id: str) -> dict[str, str]:
     if _job_retain_results(job_id):
         return {_GATEWAY_RETAIN_RESULTS_HEADER: "true"}
@@ -253,36 +269,81 @@ async def _enqueue_or_reject(pool_type: PoolType, item: WorkItem) -> None:
 
 
 async def _fetch_result_data_from_workers(document_id: str) -> list[dict[str, Any]] | None:
-    """Pull cached rows from the batch/realtime pod that processed *document_id*."""
-    proxy = get_proxy()
-    if proxy is None:
-        return None
-    for pool_type in (PoolType.BATCH, PoolType.REALTIME):
-        client = proxy._client_for(pool_type)
-        try:
-            resp = await client.get(f"/v1/internal/document-result/{document_id}")
-        except Exception as exc:
-            logger.debug(
-                "Worker result fetch from %s failed for %s: %s",
-                pool_type.value,
-                document_id,
-                exc,
-            )
-            continue
-        if resp.status_code == 404:
-            continue
-        if resp.status_code != 200:
-            logger.warning(
-                "Worker result fetch from %s returned HTTP %d for %s",
-                pool_type.value,
-                resp.status_code,
-                document_id,
-            )
-            continue
-        rows = resp.json().get("result_data")
-        if rows is not None:
-            return rows
-    return None
+    """Read rows already handed off to this gateway's retained store."""
+    try:
+        rows = await asyncio.to_thread(get_result_data, document_id)
+    except ResultStoreTemporarilyUnavailable as exc:
+        raise HTTPException(
+            status_code=503,
+            detail=str(exc),
+            headers={"Retry-After": str(_RESULT_RETRY_AFTER_SECONDS)},
+        ) from exc
+    if rows is not None:
+        return rows
+    raise HTTPException(
+        status_code=503,
+        detail=f"Retained result data for {document_id!r} is temporarily unavailable",
+        headers={"Retry-After": str(_RESULT_RETRY_AFTER_SECONDS)},
+    )
+
+
+def _worker_result_url(request: Request, document_id: str, worker_ip_value: Any) -> str:
+    """Build a fixed-path owner URL from a validated worker pod IP."""
+    if not isinstance(worker_ip_value, str):
+        raise HTTPException(status_code=503, detail="Completion callback is missing result worker identity")
+    try:
+        worker_ip = ipaddress.ip_address(worker_ip_value)
+    except ValueError as exc:
+        raise HTTPException(status_code=400, detail="Completion callback has an invalid result worker IP") from exc
+    if worker_ip.is_unspecified or worker_ip.is_multicast or worker_ip.is_loopback or worker_ip.is_link_local:
+        raise HTTPException(status_code=400, detail="Completion callback has an unroutable result worker IP")
+
+    peer_value = request.client.host if request.client is not None else ""
+    try:
+        peer_ip = ipaddress.ip_address(peer_value)
+    except ValueError:
+        peer_ip = None
+    if peer_ip is not None and not peer_ip.is_loopback and peer_ip != worker_ip:
+        raise HTTPException(status_code=400, detail="Result worker IP does not match callback peer")
+
+    host = f"[{worker_ip}]" if worker_ip.version == 6 else str(worker_ip)
+    port = request.app.state.config.server.port
+    return f"http://{host}:{port}/v1/internal/document-result/{quote(document_id, safe='')}"
+
+
+async def _pull_and_store_worker_result(request: Request, document_id: str, worker_ip: Any) -> None:
+    """Copy rows from the exact completing worker into the gateway store."""
+    url = _worker_result_url(request, document_id, worker_ip)
+    try:
+        async with httpx.AsyncClient(timeout=10.0, headers=_internal_auth_headers(request)) as client:
+            response = await client.get(url)
+    except httpx.HTTPError as exc:
+        raise HTTPException(
+            status_code=503,
+            detail=f"Unable to fetch retained result for {document_id!r} from its worker",
+            headers={"Retry-After": "1"},
+        ) from exc
+    if response.status_code != 200:
+        raise HTTPException(
+            status_code=503,
+            detail=f"Worker returned HTTP {response.status_code} for retained result {document_id!r}",
+            headers={"Retry-After": "1"},
+        )
+    try:
+        payload = response.json()
+        rows = payload.get("result_data") if isinstance(payload, dict) else None
+    except ValueError as exc:
+        raise HTTPException(status_code=503, detail=f"Worker returned invalid result data for {document_id!r}") from exc
+    if not rows or not isinstance(rows, list) or not all(isinstance(row, dict) for row in rows):
+        raise HTTPException(status_code=503, detail=f"Worker returned invalid result data for {document_id!r}")
+    try:
+        await asyncio.to_thread(store_result_data, document_id, rows)
+    except (OSError, ValueError, TypeError) as exc:
+        raise HTTPException(
+            status_code=503,
+            detail=f"Unable to retain result data for {document_id!r} on the gateway",
+            headers={"Retry-After": "1"},
+        ) from exc
 
 
 def _build_callback_url(request: Request) -> str:
@@ -756,8 +817,8 @@ async def get_job_document(
 ) -> Response:
     """Single-document detail nested under the owning job.
 
-    Returns HTTP 200 when the document is terminal (consumes
-    ``result_data`` so memory is freed for the caller) and HTTP 202
+    Returns HTTP 200 when the document is terminal (including retained
+    ``result_data`` when requested for the job) and HTTP 202
     while still pending/processing. A 404 is returned if either the
     job is unknown or the document does not belong to this job — the
     latter prevents leaking document existence across tenants.
@@ -775,7 +836,15 @@ async def get_job_document(
             detail=f"Document {document_id!r} not found in job {job_id!r}",
         )
     is_terminal = rec.status in (DocumentStatus.COMPLETED, DocumentStatus.FAILED)
-    result_data = tracker.consume_result_data(document_id) if is_terminal else None
+    result_data = tracker.get_result_data(document_id) if is_terminal else None
+    if (
+        is_terminal
+        and result_data is None
+        and rec.result_rows
+        and _job_retain_results(rec.job_id)
+        and _is_gateway(request)
+    ):
+        result_data = await _fetch_result_data_from_workers(document_id)
     body = _document_to_response(rec, result_data=result_data).model_dump()
     return JSONResponse(content=body, status_code=200 if is_terminal else 202)
 
@@ -899,6 +968,7 @@ async def submit_document_to_job(
                 payload=file_bytes,
                 filename=file.filename,
                 callback_url=gw_callback_url,
+                callback_headers=_internal_auth_headers(request),
                 job_id=gw_job_id,
                 pipeline_spec=worker_spec.model_dump(mode="json") if worker_spec is not None else None,
                 retain_results=_work_item_retain_results(request, job_id=gw_job_id),
@@ -1032,6 +1102,7 @@ async def submit_page_to_job(
                     payload=file_bytes,
                     filename=file.filename,
                     callback_url=gw_callback_url,
+                    callback_headers=_internal_auth_headers(request),
                     job_id=gw_job_id,
                     retain_results=_work_item_retain_results(request, job_id=gw_job_id),
                 ),
@@ -1165,6 +1236,7 @@ async def submit_whole_document_to_job(
                     payload=file_bytes,
                     filename=file.filename,
                     callback_url=gw_callback_url,
+                    callback_headers=_internal_auth_headers(request),
                     job_id=gw_job_id,
                     pipeline_spec=worker_spec.model_dump(mode="json") if worker_spec is not None else None,
                     retain_results=_work_item_retain_results(request, job_id=gw_job_id),
@@ -1206,9 +1278,8 @@ async def _status_response(request: Request, item_id: str) -> JSONResponse:
     """Look up document status and return the appropriate HTTP code.
 
     Returns 200 for completed/failed, 202 for pending/processing, 404 if unknown.
-    When returning a terminal (200) response, result_data is consumed from the
-    tracker (or, in gateway mode, from the worker pod that ran the pipeline)
-    so memory is freed after the client has retrieved it.
+    Terminal result data is read idempotently from the tracker or, in gateway
+    mode, from the shared worker result store.
     """
     from nemo_retriever.service.services.job_tracker import DocumentStatus
 
@@ -1223,8 +1294,14 @@ async def _status_response(request: Request, item_id: str) -> JSONResponse:
         raise HTTPException(status_code=404, detail=f"No tracked document with id={item_id!r}")
 
     is_terminal = rec.status in (DocumentStatus.COMPLETED, DocumentStatus.FAILED)
-    result_data = tracker.consume_result_data(item_id) if is_terminal else None
-    if is_terminal and result_data is None and rec.result_rows and _is_gateway(request):
+    result_data = tracker.get_result_data(item_id) if is_terminal else None
+    if (
+        is_terminal
+        and result_data is None
+        and rec.result_rows
+        and _job_retain_results(rec.job_id)
+        and _is_gateway(request)
+    ):
         result_data = await _fetch_result_data_from_workers(item_id)
 
     body = JobStatusResponse(
@@ -1672,9 +1749,14 @@ async def query(request: Request) -> Response:
 )
 async def worker_document_result(document_id: str) -> JSONResponse:
     """Return rows stored by the worker pool after pipeline completion."""
-    from nemo_retriever.service.services.worker_result_store import consume_result_data
-
-    rows = consume_result_data(document_id)
+    try:
+        rows = get_result_data(document_id)
+    except ResultStoreTemporarilyUnavailable as exc:
+        raise HTTPException(
+            status_code=503,
+            detail=str(exc),
+            headers={"Retry-After": str(_RESULT_RETRY_AFTER_SECONDS)},
+        ) from exc
     if rows is None:
         raise HTTPException(
             status_code=404,
@@ -1731,9 +1813,30 @@ async def job_callback(request: Request) -> JSONResponse:
             elapsed_s=body.get("elapsed_s"),
         )
     else:
+        result_rows = body.get("result_rows", 0)
+        if pre_rec is None and result_rows and body.get("result_worker_ip"):
+            logger.warning(
+                "Permanently rejecting retained result handoff for unknown document %s",
+                item_id,
+            )
+            raise HTTPException(
+                status_code=410,
+                detail=f"Gateway has no tracked document {item_id!r} for retained result handoff",
+            )
+        if pre_rec is not None and result_rows and tracker.should_retain_results(pre_rec.job_id):
+            try:
+                retained_rows = await asyncio.to_thread(get_result_data, item_id)
+            except ResultStoreTemporarilyUnavailable as exc:
+                raise HTTPException(
+                    status_code=503,
+                    detail=str(exc),
+                    headers={"Retry-After": "1"},
+                ) from exc
+            if retained_rows is None:
+                await _pull_and_store_worker_result(request, item_id, body.get("result_worker_ip"))
         outcome = tracker.mark_completed(
             item_id,
-            result_rows=body.get("result_rows", 0),
+            result_rows=result_rows,
             elapsed_s=body.get("elapsed_s"),
         )