Skip to content
Open
Show file tree
Hide file tree
Changes from 8 commits
Commits
Show all changes
20 commits
Select commit Hold shift + click to select a range
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 21 additions & 0 deletions nemo_retriever/helm/templates/deployment.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -272,6 +278,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: POD_IP
valueFrom:
fieldRef:
Expand Down Expand Up @@ -316,6 +328,10 @@ spec:
- name: config
mountPath: /etc/nemo-retriever
readOnly: true
{{- if $.Values.retrieverResults.enabled }}
- name: retriever-results
mountPath: {{ $.Values.retrieverResults.mountPath }}
{{- end }}
- name: tmp
mountPath: /tmp
{{- if $svc.startupProbe.enabled }}
Expand All @@ -336,5 +352,10 @@ spec:
name: {{ include "nemo-retriever.role.configMapName" (dict "context" $ "role" $role) }}
- name: tmp
emptyDir: {}
{{- if $.Values.retrieverResults.enabled }}
- name: retriever-results
persistentVolumeClaim:
claimName: {{ include "nemo-retriever.retrieverResultsPvcName" $ }}
{{- end }}
{{- end }}
{{- end }}
8 changes: 8 additions & 0 deletions nemo_retriever/helm/values.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -679,10 +679,18 @@ serviceConfig:
# 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.
# In split mode, gateway, realtime, and batch pods all mount this claim so
# result lookup is deterministic regardless of which worker processed a
# document.
#
# Use local/host-path storage for results to avoid network I/O overhead.
# A ReadWriteOnce claim can be shared only by pods on the same node. Use an
# RWX-capable storage class when split roles can schedule across multiple nodes.
retrieverResults:
enabled: true
# Unconsumed results and interrupted temporary/claim files are removed
# opportunistically after this age. Covers the job tracker stale and terminal retention windows.
ttlSeconds: 28800
existingClaim: ""
storageClass: ""
accessModes:
Expand Down
11 changes: 8 additions & 3 deletions nemo_retriever/src/nemo_retriever/service/routers/ingest.py
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@
INGEST_REQUESTS_TOTAL,
)
from nemo_retriever.service.services.proxy import get_proxy
from nemo_retriever.service.services.worker_result_store import consume_result_data
from nemo_retriever.service.utils.file_type import (
FileCategory,
FileClassifier,
Expand Down Expand Up @@ -253,7 +254,11 @@ async def _enqueue_or_reject(pool_type: PoolType, item: WorkItem) -> None:

Comment thread
greptile-apps[bot] marked this conversation as resolved.

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*."""
"""Consume shared rows, falling back to legacy worker Service lookup."""
rows = consume_result_data(document_id)
if rows is not None:
return rows

proxy = get_proxy()
if proxy is None:
return None
Expand Down Expand Up @@ -776,6 +781,8 @@ async def get_job_document(
)
is_terminal = rec.status in (DocumentStatus.COMPLETED, DocumentStatus.FAILED)
result_data = tracker.consume_result_data(document_id) if is_terminal else None
if is_terminal and result_data is None and rec.result_rows 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)

Expand Down Expand Up @@ -1672,8 +1679,6 @@ 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)
if rows is None:
raise HTTPException(
Expand Down
Original file line number Diff line number Diff line change
@@ -1,38 +1,209 @@
# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.
# SPDX-License-Identifier: Apache-2.0

"""Ephemeral per-document result rows on worker pods (split topology).
"""Per-document result rows for split-topology workers.

Worker → gateway completion callbacks intentionally omit ``result_data``
to keep POST bodies small. Rows are held here until the gateway (or a
client polling through the gateway) fetches them via
``GET /v1/internal/document-result/{document_id}``.
to keep POST bodies small. When ``NEMO_RETRIEVER_RESULTS_DIR`` is set, rows
are written atomically to that shared directory so any gateway or worker pod
can consume them by document ID. The in-memory store remains as a fallback
for local and non-Helm deployments that do not configure shared storage.
"""

from __future__ import annotations

import hashlib
import json
import logging
import math
import os
import threading
import time
import uuid
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)

_lock = threading.Lock()
_store: dict[str, list[dict[str, Any]]] = {}
_RESULTS_DIR_ENV = "NEMO_RETRIEVER_RESULTS_DIR"
_RESULTS_TTL_S_ENV = "NEMO_RETRIEVER_RESULTS_TTL_SECONDS"
_DEFAULT_RESULTS_TTL_S = 8 * 3600 # stale-job window plus terminal-job retention
_SWEEP_INTERVAL_S = 60
_last_sweep_dir: Path | None = None
_last_sweep_at = 0.0


def _results_dir() -> Path | None:
value = os.environ.get(_RESULTS_DIR_ENV, "").strip()
return Path(value) if value else None


def _result_path(results_dir: Path, document_id: str) -> Path:
"""Return a traversal-safe, deterministic path for *document_id*."""
digest = hashlib.sha256(document_id.encode("utf-8")).hexdigest()
return results_dir / f"{digest}.json"


def _results_ttl_s() -> float:
value = os.environ.get(_RESULTS_TTL_S_ENV, "").strip()
if not value:
return _DEFAULT_RESULTS_TTL_S
try:
ttl_s = float(value)
except ValueError:
ttl_s = 0
if not math.isfinite(ttl_s) or ttl_s <= 0:
logger.warning("Ignoring invalid %s=%r; using %s seconds", _RESULTS_TTL_S_ENV, value, _DEFAULT_RESULTS_TTL_S)
return _DEFAULT_RESULTS_TTL_S
return ttl_s


def _is_result_store_file(path: Path) -> bool:
"""Return whether *path* is a result or an intermediate result file."""
name = path.name
if name.endswith(".json"):
digest = name.removesuffix(".json")
return len(digest) == 64 and all(character in "0123456789abcdef" for character in digest)
if not name.startswith(".") or not name.endswith((".tmp", ".claim", ".cleanup")):
return False
parts = name[1:].split(".")
return (
len(parts) == 4
and len(parts[0]) == 64
and all(character in "0123456789abcdef" for character in parts[0])
and parts[1] == "json"
and len(parts[2]) == 32
and all(character in "0123456789abcdef" for character in parts[2])
)


def _remove_expired_result(path: Path, *, cutoff: float) -> bool:
"""Atomically claim and remove an expired result without deleting a replacement."""
try:
if path.stat().st_mtime > cutoff:
return False
if not path.name.endswith(".json"):
path.unlink(missing_ok=True)
return True

claimed = path.with_name(f".{path.name}.{uuid.uuid4().hex}.cleanup")
# Restoring a concurrently replaced result requires an atomic,
# no-overwrite hard link. Verify that operation is supported before
# moving the canonical result out of the way; some RWX filesystems do
# not support hard links, so expiry is best-effort there.
os.link(path, claimed)
claimed.unlink(missing_ok=True)

os.replace(path, claimed)
if claimed.stat().st_mtime <= cutoff:
claimed.unlink(missing_ok=True)
return True

# A writer replaced the stale file between stat() and os.replace().
# Restore the fresh inode only if no newer result now owns the path.
try:
os.link(claimed, path)
except FileExistsError:
pass
claimed.unlink(missing_ok=True)
Comment thread
charlesbluca marked this conversation as resolved.
Outdated
except FileNotFoundError:
pass # Another pod consumed, replaced, or swept it first.
except OSError:
logger.debug("Unable to remove expired shared result file %s", path, exc_info=True)
return False


def _sweep_expired_files(results_dir: Path, *, now: float, ttl_s: float) -> None:
"""Best-effort removal of expired files owned by this result store."""
try:
paths = list(results_dir.iterdir())
except FileNotFoundError:
return
except OSError:
logger.warning("Unable to scan shared result directory %s", results_dir, exc_info=True)
return

cutoff = now - ttl_s
removed = sum(_remove_expired_result(path, cutoff=cutoff) for path in paths if _is_result_store_file(path))
if removed:
logger.info("Removed %d expired shared result file(s) from %s", removed, results_dir)


def _maybe_sweep_expired_files(results_dir: Path) -> None:
"""Sweep at most once per interval in this process; other pods may also sweep."""
global _last_sweep_at, _last_sweep_dir

monotonic_now = time.monotonic()
with _lock:
if _last_sweep_dir == results_dir and monotonic_now - _last_sweep_at < _SWEEP_INTERVAL_S:
return
_last_sweep_dir = results_dir
_last_sweep_at = monotonic_now
_sweep_expired_files(results_dir, now=time.time(), ttl_s=_results_ttl_s())


def _store_on_filesystem(results_dir: Path, document_id: str, result_data: list[dict[str, Any]]) -> None:
results_dir.mkdir(parents=True, exist_ok=True)
_maybe_sweep_expired_files(results_dir)
target = _result_path(results_dir, document_id)
temporary = results_dir / f".{target.name}.{uuid.uuid4().hex}.tmp"
try:
with temporary.open("x", encoding="utf-8") as stream:
json.dump(result_data, stream, ensure_ascii=False, separators=(",", ":"))
stream.flush()
os.fsync(stream.fileno())
os.replace(temporary, target)
finally:
temporary.unlink(missing_ok=True)
Comment thread
charlesbluca marked this conversation as resolved.
Outdated


def _consume_from_filesystem(results_dir: Path, document_id: str) -> list[dict[str, Any]] | None:
_maybe_sweep_expired_files(results_dir)
target = _result_path(results_dir, document_id)
claimed = results_dir / f".{target.name}.{uuid.uuid4().hex}.claim"
try:
os.replace(target, claimed)
except FileNotFoundError:
return None

try:
with claimed.open(encoding="utf-8") as stream:
rows = json.load(stream)
if not isinstance(rows, list):
raise ValueError(f"Shared result payload for {document_id!r} is not a JSON list")
return rows
except ValueError:
logger.warning("Unable to decode shared result payload for %r", document_id, exc_info=True)
return None
finally:
claimed.unlink(missing_ok=True)
Comment thread
charlesbluca marked this conversation as resolved.
Outdated


def store_result_data(document_id: str, result_data: list[dict[str, Any]] | None) -> None:
"""Retain *result_data* for a completed document on this worker pod."""
"""Retain *result_data* for a completed document."""
if not document_id or not result_data:
return
if results_dir := _results_dir():
_store_on_filesystem(results_dir, document_id, result_data)
return
with _lock:
_store[document_id] = result_data


def consume_result_data(document_id: str) -> list[dict[str, Any]] | None:
"""Return stored rows for *document_id* and remove them from memory."""
"""Return stored rows for *document_id* and remove them from the store."""
if results_dir := _results_dir():
return _consume_from_filesystem(results_dir, document_id)
with _lock:
return _store.pop(document_id, None)


def clear_for_tests() -> None:
"""Test helper — drop all cached rows."""
global _last_sweep_at, _last_sweep_dir
with _lock:
_store.clear()
_last_sweep_dir = None
_last_sweep_at = 0.0
Loading
Loading