diff --git a/docs/docs/extraction/nemo-retriever-api-reference.md b/docs/docs/extraction/nemo-retriever-api-reference.md index 68a1c688ff..1397ca9aa3 100644 --- a/docs/docs/extraction/nemo-retriever-api-reference.md +++ b/docs/docs/extraction/nemo-retriever-api-reference.md @@ -140,6 +140,79 @@ For a support-oriented mapping of extraction paths, error signals, corrective actions, and escalation criteria, refer to [Python API error triage](troubleshoot.md#python-api-error-triage). +## Capture remote NIM requests { #capture-remote-nim-requests } + +Use `InferenceCaptureConfig` to persist the final outbound request payload for +remote inference. This is useful when you want to replay representative +extraction or retrieval requests against a compatible self-hosted NIM. + +The feature is disabled by default. It captures supported HTTP JSON NIM requests +from the generic NIM client, embedding, and reranking paths, plus Triton gRPC +requests. It works with self-hosted and NVIDIA-hosted endpoints. It does not +capture inputs sent to local Transformers or vLLM models because those execution +modes do not make a NIM network request. + +The following example enables capture for an in-process ingestion pipeline. + +```python +from nemo_retriever import create_ingestor +from nemo_retriever.common.inference_capture import InferenceCaptureConfig + +capture = InferenceCaptureConfig( + storage_uri="/var/lib/nemo-retriever/nim-captures", + failure_mode="required", + operations=("ingest",), +) + +pipeline = create_ingestor( + run_mode="inprocess", + inference_capture=capture, +) +``` + +Pass the same configuration to `GraphIngestor` or `Retriever` when you create +those objects directly. Use `operations=("query",)` to capture query embedding +and reranking requests, or omit `operations` to capture both ingestion and +query requests. Use `stages` to restrict capture to named inference stages. + +Each captured attempt has its own directory. HTTP JSON captures contain +`manifest.json` and `request.json`. Triton gRPC captures contain `manifest.json` +and `request.bin`, a NumPy `savez_compressed` archive of the input tensors. Its +manifest has `protocol: "grpc"` and records the input names, data types, requested +output names, and inference parameters. All manifests record the operation, +stage, sanitized endpoint, model when available, timestamp, and retry attempt. +Request inputs are persisted for replay. Embedding HTTP captures also add an +optional replay-metadata sidecar in `manifest.json` at `metadata.replay`. It does +not change `request.json`, so the captured request remains the exact payload +sent to the NIM. + +`metadata.replay.replay_version` is `1`. Its `records` list is aligned with the +request `input` positions. Each item includes `input_index`, an `input_sha256`, +and the source `record`. The source record preserves available row identity and +context, including VectorDB fields or application query-identifying fields when +they are present. It omits `_content` and any `metadata.embedding` value. Do not +assume a normalized VectorDB identity or a dedicated query-ID field. + +The manifest `operation` labels ingestion as `ingest` and retrieval embedding as +`query`. Use the operation, input position, and replay record together to map a +captured embedding request back to its source row. The replay record can contain +document metadata, retrieved content, and query identifiers. Treat it as +sensitive data. The recorder does not persist authorization headers, API keys, +or endpoint query strings. + +By default, `failure_mode="best_effort"` logs a capture write failure and still +sends the inference request. Set `failure_mode="required"` when you generate +replay fixtures and want a capture write failure to stop the request before it +is sent. `storage_uri` accepts a local directory or an +[fsspec-compatible](https://filesystem-spec.readthedocs.io/) URI. + +!!! warning "Treat capture artifacts as sensitive data" + + Request payloads can contain document text, user queries, and base64-encoded + page images. Store captures only in an approved location. Configure access + controls, encryption, and retention policies before enabling capture in a + production environment. + !!! note "Version-specific behavior" This reference describes the current NeMo Retriever Library. Older diff --git a/nemo_retriever/dev/compose/service-mode.compose.yaml b/nemo_retriever/dev/compose/service-mode.compose.yaml index 8985afd914..c64ad265f6 100644 --- a/nemo_retriever/dev/compose/service-mode.compose.yaml +++ b/nemo_retriever/dev/compose/service-mode.compose.yaml @@ -99,11 +99,11 @@ services: restart: unless-stopped # Page-elements and table-structure are separate Compose services (matching - # Helm) but both run the combined nemotron-object-detection:2.0.0 image. + # Helm) but both run the combined nemotron-object-detection:2.0.1 image. nim-page-elements: <<: *nim-service profiles: [nims-core] - image: ${NIM_PAGE_ELEMENTS_IMAGE:-nvcr.io/nim/nvidia/nemotron-object-detection}:${NIM_PAGE_ELEMENTS_TAG:-2.0.0} + image: ${NIM_PAGE_ELEMENTS_IMAGE:-nvcr.io/nim/nvidia/nemotron-object-detection}:${NIM_PAGE_ELEMENTS_TAG:-2.0.1} ports: ["${NIM_PAGE_ELEMENTS_HOST_PORT:-8001}:8000"] environment: <<: *nim-environment @@ -125,7 +125,7 @@ services: nim-table-structure: <<: *nim-service profiles: [nims-core] - image: ${NIM_TABLE_STRUCTURE_IMAGE:-nvcr.io/nim/nvidia/nemotron-object-detection}:${NIM_TABLE_STRUCTURE_TAG:-2.0.0} + image: ${NIM_TABLE_STRUCTURE_IMAGE:-nvcr.io/nim/nvidia/nemotron-object-detection}:${NIM_TABLE_STRUCTURE_TAG:-2.0.1} ports: ["${NIM_TABLE_STRUCTURE_HOST_PORT:-8002}:8000"] environment: <<: *nim-environment @@ -147,7 +147,7 @@ services: nim-ocr: <<: *nim-service profiles: [nims-core] - image: ${NIM_OCR_IMAGE:-nvcr.io/nim/nvidia/nemotron-ocr-v2}:${NIM_OCR_TAG:-2.0.0} + image: ${NIM_OCR_IMAGE:-nvcr.io/nim/nvidia/nemotron-ocr-v2}:${NIM_OCR_TAG:-2.0.1} ports: ["${NIM_OCR_HOST_PORT:-8003}:8000"] environment: <<: *nim-environment @@ -170,7 +170,7 @@ services: nim-embedding: <<: *nim-service profiles: [nims-core] - image: ${NIM_EMBED_IMAGE:-nvcr.io/nim/nvidia/llama-nemotron-embed-vl-1b-v2}:${NIM_EMBED_TAG:-1.12.0} + image: ${NIM_EMBED_IMAGE:-nvcr.io/nim/nvidia/llama-nemotron-embed-vl-1b-v2}:${NIM_EMBED_TAG:-2.3.0} ports: ["${NIM_EMBED_HOST_PORT:-8004}:8000"] environment: <<: *nim-environment diff --git a/nemo_retriever/helm/README.md b/nemo_retriever/helm/README.md index d8c501cf99..dc6b0f7fd3 100644 --- a/nemo_retriever/helm/README.md +++ b/nemo_retriever/helm/README.md @@ -182,6 +182,53 @@ helm install retriever ./nemo_retriever/helm \ via `optional: true` `secretKeyRef`, so the install still succeeds when the secret is absent (useful for fully local NIM endpoints). +### Capture outbound NIM requests + +The retriever service can persist final outbound request payloads for supported +remote HTTP JSON and Triton gRPC NIM calls during service ingestion, plus +VectorDB query embedding calls. This is useful when you want to collect replay +fixtures for a compatible self-hosted NIM. It does not capture other query-phase +traffic. + +The capture location is administrator-owned service configuration. It is not a +field that an ingest request can override. Mount or otherwise make an approved +location writable by both the retriever service and VectorDB query deployment, +then configure the chart: + +```yaml +serviceConfig: + inferenceCapture: + enabled: true + storageUri: /var/lib/nemo-retriever/nim-captures + failureMode: best_effort + # Optional filters. Empty lists capture all supported service-ingestion stages. + operations: [] + stages: [] +``` + +`storageUri` is required when capture is enabled. Use `failureMode: required` +when every request must be captured before the service sends it to the NIM. The +default, `best_effort`, logs capture write failures and continues ingestion. +Each HTTP JSON attempt writes `manifest.json` and `request.json` in its own +capture directory. A Triton gRPC attempt writes `manifest.json` and `request.bin`, +a NumPy `savez_compressed` archive of input tensors; its manifest records the +input/output names, data types, and inference parameters. Embedding HTTP +captures also include an optional replay-metadata sidecar in `manifest.json` at +`metadata.replay`; it does not alter the raw `request.json` payload. Version 1 +metadata has `records` aligned with request input positions. Each record includes +the input index, a SHA-256 hash of the input, and its source row. The source row +preserves available VectorDB identity, document context, or application query +identity, but does not provide a normalized VDB schema or a dedicated query-ID +field. `manifest.json` labels service ingestion as `operation: "ingest"` and +VectorDB query embedding as `operation: "query"`. The recorder does not persist +authorization headers, API keys, or endpoint query strings. + +Capture artifacts can include document text, base64-encoded page images, source +row metadata, retrieved content, and query-identifying fields. Use an approved +destination with suitable access controls, encryption, and retention policies. +This feature is separate from pipeline `.store()`, which persists ingest output +artifacts rather than outbound NIM requests. + ### 3. Install with the NIM Operator (in-cluster NIMs) Install the [NIM Operator](https://docs.nvidia.com/nim-operator/) first so @@ -324,6 +371,11 @@ The retriever service picks up the in-cluster ASR endpoint when `nimOperator.aud | `serviceConfig.pipeline.realtimeWorkers` | `24` | Per-pod realtime worker count. | | `serviceConfig.pipeline.batchWorkers` | `48` | Per-pod batch worker count. Refer to [Timeouts and alleviating ingest failures](#timeouts-and-alleviating-ingest-failures) if embed or pool errors appear under load. | | `serviceConfig.resources.maxUploadBytes` | `500000000` | Maximum upload file size in bytes; requests exceeding the limit are rejected before buffering. | +| `serviceConfig.inferenceCapture.enabled` | `false` | Enables administrator-owned capture for supported outbound remote-NIM requests during ingestion and VectorDB query embedding. | +| `serviceConfig.inferenceCapture.storageUri` | `""` | Local directory or fsspec-compatible URI. Required when capture is enabled. | +| `serviceConfig.inferenceCapture.failureMode` | `best_effort` | Use `required` to fail before a request is sent when it cannot be captured. | +| `serviceConfig.inferenceCapture.operations` | `[]` | Optional operation filter. | +| `serviceConfig.inferenceCapture.stages` | `[]` | Optional inference-stage filter. | | `serviceConfig.nimEndpoints.*InvokeUrl` | `""` | Override the auto-resolved NIM Operator URL. Available knobs: `pageElementsInvokeUrl`, `tableStructureInvokeUrl`, `ocrInvokeUrl`, `embedInvokeUrl`, and `captionInvokeUrl` (refer to [Image captioning (Omni 30B)](#image-captioning-omni-30b)). | | `serviceConfig.nimEndpoints.captionModelName` | `""` | Model id sent to the remote VLM. Auto-set to `nvidia/nemotron-3-nano-omni-30b-a3b-reasoning` whenever a caption URL is resolved. | | `serviceConfig.llm.enabled` | `false` | Enables `POST /v1/answer`. Auto-flips to true when `nimOperator.answer_llm` is enabled and the operator URL resolves. | diff --git a/nemo_retriever/helm/templates/configmap.yaml b/nemo_retriever/helm/templates/configmap.yaml index 4a170089b2..cfa61199e0 100644 --- a/nemo_retriever/helm/templates/configmap.yaml +++ b/nemo_retriever/helm/templates/configmap.yaml @@ -96,6 +96,13 @@ logging: file: {{ .Values.serviceConfig.logging.file | quote }} format: {{ .Values.serviceConfig.logging.format | quote }} +inference_capture: + enabled: {{ .Values.serviceConfig.inferenceCapture.enabled }} + storage_uri: {{ if .Values.serviceConfig.inferenceCapture.storageUri }}{{ .Values.serviceConfig.inferenceCapture.storageUri | quote }}{{ else }}null{{ end }} + failure_mode: {{ .Values.serviceConfig.inferenceCapture.failureMode | quote }} + operations: {{ .Values.serviceConfig.inferenceCapture.operations | toJson }} + stages: {{ .Values.serviceConfig.inferenceCapture.stages | toJson }} + nim_endpoints: page_elements_invoke_url: {{ .pageElementsURL | quote }} table_structure_invoke_url: {{ .tableStructureURL | quote }} diff --git a/nemo_retriever/helm/templates/deployment-vectordb.yaml b/nemo_retriever/helm/templates/deployment-vectordb.yaml index 1df0a86a3a..8d89f851e5 100644 --- a/nemo_retriever/helm/templates/deployment-vectordb.yaml +++ b/nemo_retriever/helm/templates/deployment-vectordb.yaml @@ -137,6 +137,14 @@ spec: containerPort: {{ $vdb.port }} protocol: TCP env: + {{- if .Values.serviceConfig.inferenceCapture.enabled }} + - name: NEMO_RETRIEVER_INFERENCE_CAPTURE_URI + value: {{ .Values.serviceConfig.inferenceCapture.storageUri | quote }} + - name: NEMO_RETRIEVER_INFERENCE_CAPTURE_FAILURE_MODE + value: {{ .Values.serviceConfig.inferenceCapture.failureMode | quote }} + - name: NEMO_RETRIEVER_INFERENCE_CAPTURE_OPERATION + value: query + {{- end }} {{- if $internalAuth.enabled }} - name: NRL_INTERNAL_VDB_TOKEN valueFrom: diff --git a/nemo_retriever/helm/values.yaml b/nemo_retriever/helm/values.yaml index d26a4419a1..d30a7198d2 100644 --- a/nemo_retriever/helm/values.yaml +++ b/nemo_retriever/helm/values.yaml @@ -538,6 +538,15 @@ serviceConfig: file: "/var/lib/nemo-retriever/retriever-service.log" format: "%(asctime)s | %(levelname)s | %(name)s | %(message)s" + # Administrator-owned persistence of outbound remote-NIM request bodies. + # Artifacts can contain document content; use an approved, writable location. + inferenceCapture: + enabled: false + storageUri: "" + failureMode: best_effort + operations: [] + stages: [] + # External NIM endpoints. Used as-is when the operator NIM is disabled # or when the NIM Operator CRDs are absent. When the corresponding # `nimOperator..enabled` is true (and the CRDs exist), the diff --git a/nemo_retriever/src/nemo_retriever/common/inference_capture.py b/nemo_retriever/src/nemo_retriever/common/inference_capture.py new file mode 100644 index 0000000000..af80274c5c --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/common/inference_capture.py @@ -0,0 +1,251 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +"""Opt-in persistence of outbound remote-inference requests. + +The recorder is deliberately transport and model-schema agnostic. It records +the final JSON value handed to an HTTP client, not application-level inputs, +so artifacts can be replayed against a compatible NIM endpoint. +""" + +from __future__ import annotations + +import contextlib +import contextvars +import json +import logging +import os +import uuid +from dataclasses import dataclass +from datetime import UTC, datetime +from pathlib import Path +from typing import Any, Iterator, Literal +from urllib.parse import urlsplit, urlunsplit + +logger = logging.getLogger(__name__) + +CaptureFailureMode = Literal["best_effort", "required"] + + +@dataclass(frozen=True) +class InferenceCaptureConfig: + """Configuration for recording remote model requests. + + ``storage_uri`` accepts a local directory or an fsspec-compatible URI. + The default mode never changes inference behavior when the capture sink is + unavailable; callers that generate replay fixtures can use ``required``. + """ + + storage_uri: str + failure_mode: CaptureFailureMode = "best_effort" + operations: tuple[str, ...] = () + stages: tuple[str, ...] = () + + def __post_init__(self) -> None: + if not str(self.storage_uri).strip(): + raise ValueError("inference capture storage_uri must not be empty") + if self.failure_mode not in {"best_effort", "required"}: + raise ValueError("inference capture failure_mode must be 'best_effort' or 'required'") + + @classmethod + def from_value(cls, value: "InferenceCaptureConfig | dict[str, Any] | None") -> "InferenceCaptureConfig | None": + if value is None or isinstance(value, cls): + return value + values = dict(value) + for name in ("operations", "stages"): + if name in values and values[name] is not None: + values[name] = tuple(values[name]) + return cls(**values) + + +_capture_config: contextvars.ContextVar[InferenceCaptureConfig | None] = contextvars.ContextVar( + "nemo_retriever_inference_capture", default=None +) +_capture_operation: contextvars.ContextVar[str | None] = contextvars.ContextVar( + "nemo_retriever_inference_capture_operation", default=None +) +# ThreadPool workers used by NIM clients do not inherit ContextVars. Keep a +# process fallback for the active operation; context-local values still win. +_active_config: InferenceCaptureConfig | None = None +_active_operation: str | None = None + + +@contextlib.contextmanager +def activate_inference_capture( + config: InferenceCaptureConfig | dict[str, Any] | None, + *, + operation: str | None = None, +) -> Iterator[None]: + """Activate capture for synchronous work in the current execution context.""" + + global _active_config, _active_operation + parsed = InferenceCaptureConfig.from_value(config) + prior_config, prior_operation = _active_config, _active_operation + _active_config, _active_operation = parsed, operation + config_token = _capture_config.set(parsed) + operation_token = _capture_operation.set(operation) + try: + yield + finally: + _capture_operation.reset(operation_token) + _capture_config.reset(config_token) + _active_config, _active_operation = prior_config, prior_operation + + +def _safe_endpoint(endpoint: str) -> str: + parts = urlsplit(str(endpoint)) + if not parts.scheme or not parts.netloc: + return str(endpoint).split("?", 1)[0].split("#", 1)[0] + return urlunsplit((parts.scheme, parts.netloc, parts.path, "", "")) + + +def _environment_capture_config() -> InferenceCaptureConfig | None: + uri = os.environ.get("NEMO_RETRIEVER_INFERENCE_CAPTURE_URI", "").strip() + if not uri: + return None + return InferenceCaptureConfig( + storage_uri=uri, + failure_mode=os.environ.get("NEMO_RETRIEVER_INFERENCE_CAPTURE_FAILURE_MODE", "best_effort"), + ) + + +def _matches(config: InferenceCaptureConfig, operation: str | None, stage: str) -> bool: + return (not config.operations or (operation or "") in config.operations) and ( + not config.stages or stage in config.stages + ) + + +def stage_from_endpoint(endpoint: str, *, default: str = "remote_nim") -> str: + """Return a stable, schema-independent stage label from an invoke URL.""" + path = urlsplit(str(endpoint)).path.rstrip("/") + if not path: + return default + return path.rsplit("/", 1)[-1].replace("-", "_") or default + + +def _write_local(directory: Path, manifest: dict[str, Any], body: bytes, suffix: str) -> None: + directory.mkdir(parents=True, exist_ok=False) + try: + manifest_tmp = directory / ".manifest.tmp" + request_tmp = directory / ".request.tmp" + manifest_tmp.write_text(json.dumps(manifest, indent=2, sort_keys=True), encoding="utf-8") + request_tmp.write_bytes(body) + os.replace(manifest_tmp, directory / "manifest.json") + os.replace(request_tmp, directory / f"request.{suffix}") + except Exception: + # A partially populated capture is never useful as a replay fixture. + for item in directory.glob("*"): + item.unlink(missing_ok=True) + directory.rmdir() + raise + + +def _write_capture( + config: InferenceCaptureConfig, capture_id: str, manifest: dict[str, Any], body: bytes, suffix: str +) -> None: + uri = str(config.storage_uri) + if "://" not in uri: + _write_local(Path(uri).expanduser().resolve() / capture_id, manifest, body, suffix) + return + + import fsspec # noqa: PLC0415 + + base = uri.rstrip("/") + "/" + capture_id + with fsspec.open(base + "/manifest.json", "wt") as handle: + json.dump(manifest, handle, indent=2, sort_keys=True) + with fsspec.open(base + f"/request.{suffix}", "wb") as handle: + handle.write(body) + + +def record_binary_request( + *, + stage: str, + endpoint: str, + payload: bytes, + protocol: str, + model: str | None = None, + attempt: int = 0, + metadata: dict[str, Any] | None = None, + operation: str | None = None, +) -> None: + """Persist one opaque transport payload, for example Triton gRPC tensors.""" + config = _capture_config.get() or _active_config or _environment_capture_config() + selected_operation = ( + operation + or _capture_operation.get() + or _active_operation + or os.environ.get("NEMO_RETRIEVER_INFERENCE_CAPTURE_OPERATION") + ) + if config is None or not _matches(config, selected_operation, stage): + return + try: + capture_id = f"{datetime.now(UTC).strftime('%Y%m%dT%H%M%S%fZ')}-{uuid.uuid4().hex}" + manifest = { + "capture_version": 1, + "capture_id": capture_id, + "timestamp": datetime.now(UTC).isoformat(), + "operation": selected_operation, + "stage": stage, + "protocol": protocol, + "endpoint": _safe_endpoint(endpoint), + "model": model, + "attempt": int(attempt), + "content_type": "application/octet-stream", + "metadata": metadata or {}, + } + _write_capture(config, capture_id, manifest, payload, "bin") + except Exception as exc: + if config.failure_mode == "required": + raise RuntimeError(f"Failed to persist inference capture for {stage}: {exc}") from exc + logger.warning("Failed to persist inference capture for %s: %s", stage, exc) + + +def record_json_request( + *, + stage: str, + endpoint: str, + payload: Any, + method: str = "POST", + model: str | None = None, + attempt: int = 0, + operation: str | None = None, + metadata: dict[str, Any] | None = None, +) -> None: + """Persist one final JSON request when capture is active. + + Credentials are intentionally not accepted by this function. Endpoint + query strings are omitted because they can contain credentials. + """ + + config = _capture_config.get() or _active_config or _environment_capture_config() + selected_operation = ( + operation + or _capture_operation.get() + or _active_operation + or os.environ.get("NEMO_RETRIEVER_INFERENCE_CAPTURE_OPERATION") + ) + if config is None or not _matches(config, selected_operation, stage): + return + + try: + body = json.dumps(payload, ensure_ascii=False, separators=(",", ":"), default=str).encode("utf-8") + capture_id = f"{datetime.now(UTC).strftime('%Y%m%dT%H%M%S%fZ')}-{uuid.uuid4().hex}" + manifest = { + "capture_version": 1, + "capture_id": capture_id, + "timestamp": datetime.now(UTC).isoformat(), + "operation": selected_operation, + "stage": stage, + "protocol": "http", + "method": method, + "endpoint": _safe_endpoint(endpoint), + "model": model, + "attempt": int(attempt), + "content_type": "application/json", + "metadata": metadata or {}, + } + _write_capture(config, capture_id, manifest, body, "json") + except Exception as exc: + if config.failure_mode == "required": + raise RuntimeError(f"Failed to persist inference capture for {stage}: {exc}") from exc + logger.warning("Failed to persist inference capture for %s: %s", stage, exc) diff --git a/nemo_retriever/src/nemo_retriever/common/params/models.py b/nemo_retriever/src/nemo_retriever/common/params/models.py index be5cae75e8..1752d5d5f6 100644 --- a/nemo_retriever/src/nemo_retriever/common/params/models.py +++ b/nemo_retriever/src/nemo_retriever/common/params/models.py @@ -316,6 +316,7 @@ class IngestorCreateParams(_ParamsModel): node_overrides: Optional[dict[str, dict[str, Any]]] = None api_key: Optional[str] = None error_policy: Literal["raise", "collect"] = "raise" + inference_capture: Any = None # service run mode: maximum number of concurrent page uploads. Lower # values (e.g. 2-4) reduce burst pressure on Kubernetes NodePort / # kube-proxy paths that otherwise reset connections under heavy load. diff --git a/nemo_retriever/src/nemo_retriever/graph/retriever.py b/nemo_retriever/src/nemo_retriever/graph/retriever.py index d342523a87..6b3572bd7b 100644 --- a/nemo_retriever/src/nemo_retriever/graph/retriever.py +++ b/nemo_retriever/src/nemo_retriever/graph/retriever.py @@ -91,6 +91,7 @@ class Retriever: embed_kwargs: dict[str, Any] = field(default_factory=dict) vdb_kwargs: dict[str, Any] = field(default_factory=dict) rerank_kwargs: dict[str, Any] = field(default_factory=dict) + inference_capture: Any = None _cached_graph: Any = field(default=None, init=False, repr=False, compare=False) _cache_key: Any = field(default=None, init=False, repr=False, compare=False) @@ -480,13 +481,16 @@ def queries( if self.graph is None: embed_kwargs = self._resolve_embed_kwargs(index_model, embed_kwargs, index_revision) - raw_hits = self._execute_queries_graph( - query_texts, - effective_top_k=candidate_top_k, - retrieval_top_k=retrieval_top_k, - vdb_call_kwargs=vdb_call_kwargs, - embed_extra=embed_kwargs, - ) + from nemo_retriever.common.inference_capture import activate_inference_capture + + with activate_inference_capture(self.inference_capture, operation="query"): + raw_hits = self._execute_queries_graph( + query_texts, + effective_top_k=candidate_top_k, + retrieval_top_k=retrieval_top_k, + vdb_call_kwargs=vdb_call_kwargs, + embed_extra=embed_kwargs, + ) return [ shape_query_hits( hits, diff --git a/nemo_retriever/src/nemo_retriever/ingestor/core.py b/nemo_retriever/src/nemo_retriever/ingestor/core.py index c300089cab..5af50bfd2e 100644 --- a/nemo_retriever/src/nemo_retriever/ingestor/core.py +++ b/nemo_retriever/src/nemo_retriever/ingestor/core.py @@ -81,6 +81,7 @@ def create_ingestor( allow_no_gpu=parsed.allow_no_gpu, node_overrides=parsed.node_overrides, error_policy=parsed.error_policy, + inference_capture=parsed.inference_capture, ) diff --git a/nemo_retriever/src/nemo_retriever/ingestor/graph_ingestor.py b/nemo_retriever/src/nemo_retriever/ingestor/graph_ingestor.py index 953397c738..7ba4742837 100644 --- a/nemo_retriever/src/nemo_retriever/ingestor/graph_ingestor.py +++ b/nemo_retriever/src/nemo_retriever/ingestor/graph_ingestor.py @@ -471,6 +471,7 @@ def __init__( node_overrides: Optional[Dict[str, Dict[str, Any]]] = None, show_progress: bool = True, error_policy: str = "raise", + inference_capture: Any = None, ) -> None: super().__init__(documents=documents) if run_mode not in {"batch", "inprocess"}: @@ -488,6 +489,7 @@ def __init__( self._node_overrides: Dict[str, Dict[str, Any]] = node_overrides or {} self._show_progress = show_progress self._error_policy = error_policy + self._inference_capture = inference_capture self._rd_dataset: Any = None self._buffers: list[tuple[str, BytesIO]] = [] self._inline_texts: list[str] | None = None @@ -916,13 +918,22 @@ def _execute_single_graph_inprocess( executor = InprocessExecutor(graph, show_progress=self._show_progress) self._rd_dataset = None if self._inline_texts: - return executor.ingest(self._inline_text_dataframe()) + from nemo_retriever.common.inference_capture import activate_inference_capture + + with activate_inference_capture(self._inference_capture, operation="ingest"): + return executor.ingest(self._inline_text_dataframe()) if self._buffers: import pandas as pd df = pd.DataFrame([{"bytes": buf.getvalue(), "path": name} for name, buf in self._buffers]) - return executor.ingest(df) - return executor.ingest(self._documents) + from nemo_retriever.common.inference_capture import activate_inference_capture + + with activate_inference_capture(self._inference_capture, operation="ingest"): + return executor.ingest(df) + from nemo_retriever.common.inference_capture import activate_inference_capture + + with activate_inference_capture(self._inference_capture, operation="ingest"): + return executor.ingest(self._documents) def _execute_extraction_branches( self, diff --git a/nemo_retriever/src/nemo_retriever/models/inference/main_text_embed.py b/nemo_retriever/src/nemo_retriever/models/inference/main_text_embed.py index 66a628f64c..7508d4399f 100644 --- a/nemo_retriever/src/nemo_retriever/models/inference/main_text_embed.py +++ b/nemo_retriever/src/nemo_retriever/models/inference/main_text_embed.py @@ -15,6 +15,7 @@ ```python import pandas as pd + from nemo_retriever.models.inference.main_text_embed import create_text_embeddings_for_df # df must have a `text` column (recommended) and may have `metadata` dicts. @@ -36,12 +37,16 @@ def local_embedder(texts): from __future__ import annotations +import hashlib +import json import logging from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple import pandas as pd + +from nemo_retriever.common.inference_capture import record_json_request from nemo_retriever.common.api.util.string_processing import ( ensure_openai_embeddings_http_url, prepend_model_provider_prefix, @@ -88,6 +93,32 @@ class TextEmbeddingConfig: nim_http_max_concurrent: int = 32 +def _json_safe(value: Any) -> Any: + """Return a JSON-compatible copy for capture sidecars.""" + return json.loads(json.dumps(value, default=str)) + + +def _embedding_replay_records(rows: list[pd.Series], inputs: list[str]) -> list[dict[str, Any]]: + """Return source rows aligned with a final outbound embedding payload.""" + records: list[dict[str, Any]] = [] + for index, (row, input_value) in enumerate(zip(rows, inputs)): + record = row.to_dict() + metadata = record.get("metadata") + if isinstance(metadata, dict): + metadata = dict(metadata) + metadata.pop("embedding", None) + record["metadata"] = metadata + record.pop("_content", None) + records.append( + { + "input_index": index, + "input_sha256": hashlib.sha256(input_value.encode("utf-8")).hexdigest(), + "record": _json_safe(record), + } + ) + return records + + # ------------------------------------------------------------------------------ # Batch processing utilities # ------------------------------------------------------------------------------ @@ -293,6 +324,8 @@ def _http_embed_openai_compat( model_provider_prefix: Optional[str] = None, dimensions: Optional[int] = None, timeout_s: float = 600.0, + replay_records: Optional[List[dict[str, Any]]] = None, + capture_operation: Optional[str] = None, ) -> List[Optional[List[float]]]: """ Best-effort HTTP embeddings call using an OpenAI-compatible schema. @@ -324,6 +357,10 @@ def _http_embed_openai_compat( if dimensions is not None: payload["dimensions"] = int(dimensions) + record_json_request( + stage="embed", endpoint=url, payload=payload, model=model_name, operation=capture_operation, + metadata={"replay": {"replay_version": 1, "records": replay_records or []}}, + ) with httpx.Client(timeout=float(timeout_s)) as client: resp = client.post(url, headers=headers, json=payload) resp.raise_for_status() @@ -360,6 +397,8 @@ def _make_async_request( modalities: Optional[List[str]] = None, dimensions: Optional[int] = None, timeout_s: float = 600.0, + replay_records: Optional[List[dict[str, Any]]] = None, + capture_operation: Optional[str] = None, ) -> dict: """ Send an HTTP OpenAI-compatible embedding request. @@ -383,6 +422,8 @@ def _make_async_request( truncate=str(truncate), dimensions=dimensions, timeout_s=timeout_s, + replay_records=replay_records, + capture_operation=capture_operation, ) response["embedding"] = vecs response["info_msg"] = None @@ -409,10 +450,15 @@ def _async_request_handler( dimensions: Optional[int] = None, max_concurrent: Optional[int] = None, timeout_s: float = 600.0, + replay_records: Optional[List[List[dict[str, Any]]]] = None, + capture_operation: Optional[str] = None, ) -> List[dict]: if modalities is None: modalities = [None] * len(prompts) # type: ignore[assignment] + if replay_records is None: + replay_records = [None] * len(prompts) # type: ignore[assignment] + pool_size = max_concurrent if max_concurrent and max_concurrent > 0 else None with ThreadPoolExecutor(max_workers=pool_size) as executor: futures = [ @@ -430,8 +476,10 @@ def _async_request_handler( modalities=modality_batch, # type: ignore[arg-type] dimensions=dimensions, timeout_s=timeout_s, + replay_records=replay_batch, + capture_operation=capture_operation, ) - for prompt_batch, modality_batch in zip(prompts, modalities) + for prompt_batch, modality_batch, replay_batch in zip(prompts, modalities, replay_records) ] results = [future.result() for future in futures] @@ -452,6 +500,8 @@ def _async_runner( dimensions: Optional[int] = None, max_concurrent: Optional[int] = None, timeout_s: float = 600.0, + replay_records: Optional[List[List[dict[str, Any]]]] = None, + capture_operation: Optional[str] = None, ) -> dict: results = _async_request_handler( prompts, @@ -467,6 +517,8 @@ def _async_runner( dimensions=dimensions, max_concurrent=max_concurrent, timeout_s=timeout_s, + replay_records=replay_records, + capture_operation=capture_operation, ) flat_results = {"embeddings": [], "info_msgs": []} @@ -611,6 +663,9 @@ def create_text_embeddings_for_df( if timeout_raw is None: timeout_raw = getattr(transform_config, "request_timeout_s", 600.0) request_timeout_s = float(timeout_raw) + capture_operation = str(task_config.get("inference_capture_operation") or ( + "query" if str(transform_config.input_type).strip().lower() == "query" else "ingest" + )) if df_transform_ledger.empty: return df_transform_ledger, {"trace_info": execution_trace_log} @@ -643,6 +698,13 @@ def _text_image_content(r: pd.Series) -> Optional[str]: df_content["_content"] = extracted_content valid_content_mask = df_content["_content"].notna() + valid_rows = [row for _, row in df_content.loc[valid_content_mask].iterrows()] + + def _capture_batches(inputs: list[str]) -> list[list[dict[str, Any]]]: + return _generate_batches( + _embedding_replay_records(valid_rows, inputs), batch_size=int(transform_config.batch_size) + ) + if valid_content_mask.any(): if embed_modality in IMAGE_MODALITIES and multimodal_embedder is not None: # Local multimodal path: use _multimodal_callable_runner @@ -676,6 +738,7 @@ def _text_image_content(r: pd.Series) -> Optional[str]: filtered_content_batches = _generate_batches( filtered_content_list, batch_size=int(transform_config.batch_size) ) + capture_batches = _capture_batches(filtered_content_list) content_embeddings = _async_runner( filtered_content_batches, api_key, @@ -690,6 +753,8 @@ def _text_image_content(r: pd.Series) -> Optional[str]: dimensions=dimensions, max_concurrent=nim_http_max_concurrent, timeout_s=request_timeout_s, + replay_records=capture_batches, + capture_operation=capture_operation, ) else: # Text-only path (default) @@ -697,6 +762,7 @@ def _text_image_content(r: pd.Series) -> Optional[str]: filtered_content_batches = _generate_batches( filtered_content_list, batch_size=int(transform_config.batch_size) ) + capture_batches = _capture_batches(filtered_content_list) if endpoint_url: content_embeddings = _async_runner( @@ -713,6 +779,8 @@ def _text_image_content(r: pd.Series) -> Optional[str]: dimensions=dimensions, max_concurrent=nim_http_max_concurrent, timeout_s=request_timeout_s, + replay_records=capture_batches, + capture_operation=capture_operation, ) elif callable(embedder): content_embeddings = _callable_runner( diff --git a/nemo_retriever/src/nemo_retriever/models/nim/nim.py b/nemo_retriever/src/nemo_retriever/models/nim/nim.py index 6613b50a60..0361b958c9 100644 --- a/nemo_retriever/src/nemo_retriever/models/nim/nim.py +++ b/nemo_retriever/src/nemo_retriever/models/nim/nim.py @@ -13,6 +13,8 @@ import requests +from nemo_retriever.common.inference_capture import record_json_request, stage_from_endpoint + logger = logging.getLogger(__name__) @@ -170,6 +172,9 @@ def _post_with_retries( except Exception as exc: logger.warning("OpenTelemetry trace propagation failed for NIM request: %s", exc) try: + record_json_request( + stage=stage_from_endpoint(invoke_url), endpoint=invoke_url, payload=payload, attempt=attempt + ) response = requests.post( invoke_url, headers=request_headers, json=payload, timeout=float(timeout_s) ) diff --git a/nemo_retriever/src/nemo_retriever/models/nim/primitives/nim_client.py b/nemo_retriever/src/nemo_retriever/models/nim/primitives/nim_client.py index 27b5bed25f..096a5df300 100644 --- a/nemo_retriever/src/nemo_retriever/models/nim/primitives/nim_client.py +++ b/nemo_retriever/src/nemo_retriever/models/nim/primitives/nim_client.py @@ -5,6 +5,7 @@ from __future__ import annotations import hashlib +import io import inspect import json import logging @@ -22,6 +23,8 @@ import numpy as np import requests +from nemo_retriever.common.inference_capture import record_binary_request, record_json_request + from nemo_retriever.common.api.internal.primitives.tracing.tagging import traceable_func from nemo_retriever.common.api.util.string_processing import generate_url @@ -462,6 +465,24 @@ def _grpc_infer( while attempt < self.max_retries: try: + capture_buffer = io.BytesIO() + np.savez_compressed( + capture_buffer, **{str(name): value for name, value in zip(input_names, formatted_input)} + ) + record_binary_request( + stage=str(self.model_interface.name()).replace("-", "_"), + endpoint=self._grpc_endpoint, + payload=capture_buffer.getvalue(), + protocol="grpc", + model=model_name, + attempt=attempt, + metadata={ + "input_names": list(input_names), + "dtypes": list(dtypes), + "output_names": list(output_names), + "parameters": parameters, + }, + ) response = _call_infer_with_optional_headers( self.client, headers=_inject_trace_headers(), @@ -600,6 +621,13 @@ def _http_infer(self, formatted_input: dict) -> dict: except Exception as exc: logger.warning("OpenTelemetry trace propagation failed for NimClient HTTP request: %s", exc) + record_json_request( + stage=str(self.model_interface.name()).replace("-", "_"), + endpoint=self.endpoint_url, + payload=formatted_input, + model=str(self.model_interface.name()), + attempt=attempt, + ) response = requests.post( self.endpoint_url, json=formatted_input, headers=request_headers, timeout=self.timeout ) diff --git a/nemo_retriever/src/nemo_retriever/operators/rerank.py b/nemo_retriever/src/nemo_retriever/operators/rerank.py index 48cce39cdd..b65781c786 100644 --- a/nemo_retriever/src/nemo_retriever/operators/rerank.py +++ b/nemo_retriever/src/nemo_retriever/operators/rerank.py @@ -219,6 +219,7 @@ def _rerank_via_endpoint( Documents not returned by ``top_n`` truncation receive ``-inf``. """ import requests + from nemo_retriever.common.inference_capture import record_json_request url = _normalize_rerank_endpoint(endpoint) use_v1_rerank_schema = _uses_v1_rerank_schema(url) @@ -229,6 +230,7 @@ def _rerank_via_endpoint( payload = _v1_rerank_payload(query, documents, model_name, images_b64) else: payload = _nim_ranking_payload(query, documents, model_name, images_b64) + record_json_request(stage="rerank", endpoint=url, payload=payload, model=model_name) response = requests.post(url, json=payload, headers=headers) response.raise_for_status() data = response.json() diff --git a/nemo_retriever/src/nemo_retriever/service/config.py b/nemo_retriever/src/nemo_retriever/service/config.py index e68a72c265..a1c7488697 100644 --- a/nemo_retriever/src/nemo_retriever/service/config.py +++ b/nemo_retriever/src/nemo_retriever/service/config.py @@ -544,6 +544,34 @@ def to_policy(self, *, caption_enabled: bool = False) -> "PipelineOverridesPolic ) +class InferenceCaptureServiceConfig(RichModel): + """Administrator-owned remote inference request capture settings.""" + + model_config = ConfigDict(extra="forbid") + + enabled: bool = False + storage_uri: str | None = None + failure_mode: Literal["best_effort", "required"] = "best_effort" + operations: list[str] = Field(default_factory=list) + stages: list[str] = Field(default_factory=list) + + @model_validator(mode="after") + def _validate_storage_uri(self) -> "InferenceCaptureServiceConfig": + if self.enabled and not (self.storage_uri or "").strip(): + raise ValueError("inference_capture.storage_uri is required when inference_capture.enabled is true") + return self + + def to_capture_config(self) -> dict[str, Any] | None: + if not self.enabled: + return None + return { + "storage_uri": str(self.storage_uri), + "failure_mode": self.failure_mode, + "operations": tuple(self.operations), + "stages": tuple(self.stages), + } + + class ServiceConfig(RichModel): """Top-level configuration for the retriever service mode. @@ -575,6 +603,7 @@ class ServiceConfig(RichModel): work_queue: WorkQueueConfig = Field(default_factory=WorkQueueConfig) vectordb: VectorDbConfig = Field(default_factory=VectorDbConfig) pipeline_overrides: PipelineOverridesConfig = Field(default_factory=PipelineOverridesConfig) + inference_capture: InferenceCaptureServiceConfig = Field(default_factory=InferenceCaptureServiceConfig) @model_validator(mode="after") def _cap_process_pool_workers_for_local_models(self) -> "ServiceConfig": diff --git a/nemo_retriever/src/nemo_retriever/service/services/pipeline_executor.py b/nemo_retriever/src/nemo_retriever/service/services/pipeline_executor.py index ccd9d27927..4c9c0daea0 100644 --- a/nemo_retriever/src/nemo_retriever/service/services/pipeline_executor.py +++ b/nemo_retriever/src/nemo_retriever/service/services/pipeline_executor.py @@ -559,6 +559,7 @@ def _build_graph_ingestor_from_spec( spec: dict[str, Any] | None, base_caption: dict[str, Any] | None = None, base_asr: dict[str, Any] | None = None, + inference_capture: dict[str, Any] | None = None, ) -> "tuple[Any, str, bool]": """Construct a :class:`GraphIngestor` reflecting the per-request *spec*. @@ -611,7 +612,7 @@ def _build_graph_ingestor_from_spec( asr_params = ASRParams(**base_asr) if base_asr else None - ingestor = GraphIngestor(run_mode="inprocess", show_progress=False) + ingestor = GraphIngestor(run_mode="inprocess", show_progress=False, inference_capture=inference_capture) ingestor = ingestor.buffers([(filename, BytesIO(payload))]) if extraction_mode == "video": @@ -741,6 +742,7 @@ def _run_pipeline_in_process( write_context: DocumentWriteContext | None = None, job_id: str | None = None, internal_api_token: str | None = None, + inference_capture: dict[str, Any] | None = None, ) -> tuple[int, list[dict[str, Any]], float]: """Execute one pipeline run inside a child process. @@ -780,6 +782,7 @@ def _run_pipeline_in_process( pipeline_spec, caption_params_dict, asr_params_dict, + inference_capture, ) result_df = ingestor.ingest() @@ -1140,6 +1143,7 @@ async def _work(item: WorkItem) -> tuple[int, list[dict[str, Any]]]: write_context, item.job_id, config.vectordb.internal_api_token, + getattr(getattr(config, "inference_capture", None), "to_capture_config", lambda: None)(), ) except BrokenProcessPool: logger.error( diff --git a/nemo_retriever/tests/test_inference_capture.py b/nemo_retriever/tests/test_inference_capture.py new file mode 100644 index 0000000000..000c1f67e0 --- /dev/null +++ b/nemo_retriever/tests/test_inference_capture.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +import json +from pathlib import Path + +import pandas as pd +import pytest + +from nemo_retriever.models.inference.main_text_embed import _embedding_replay_records + +from nemo_retriever.common.inference_capture import ( + InferenceCaptureConfig, + activate_inference_capture, + record_binary_request, + record_json_request, +) + + +def test_records_sanitized_replay_artifact(tmp_path: Path) -> None: + with activate_inference_capture(InferenceCaptureConfig(str(tmp_path), failure_mode="required"), operation="query"): + record_json_request( + stage="rerank", + endpoint="https://nim.example/v1/rerank?api_key=secret", + payload={"query": "hello", "model": "test"}, + model="test", + ) + + capture_dir = next(tmp_path.iterdir()) + manifest = json.loads((capture_dir / "manifest.json").read_text()) + assert manifest["operation"] == "query" + assert manifest["stage"] == "rerank" + assert manifest["endpoint"] == "https://nim.example/v1/rerank" + assert "secret" not in (capture_dir / "manifest.json").read_text() + assert json.loads((capture_dir / "request.json").read_text()) == {"query": "hello", "model": "test"} + + +def test_best_effort_does_not_interrupt_inference(tmp_path: Path) -> None: + target = tmp_path / "file-not-directory" + target.write_text("not a directory") + with activate_inference_capture(InferenceCaptureConfig(str(target)), operation="ingest"): + record_json_request(stage="ocr", endpoint="http://nim/v1/ocr", payload={"image": "x"}) + + +def test_required_capture_failure_raises(tmp_path: Path) -> None: + target = tmp_path / "file-not-directory" + target.write_text("not a directory") + with activate_inference_capture(InferenceCaptureConfig(str(target), failure_mode="required"), operation="ingest"): + with pytest.raises(RuntimeError, match="Failed to persist inference capture"): + record_json_request(stage="ocr", endpoint="http://nim/v1/ocr", payload={"image": "x"}) + + +def test_records_binary_transport_artifact(tmp_path: Path) -> None: + with activate_inference_capture(InferenceCaptureConfig(str(tmp_path), failure_mode="required"), operation="ingest"): + record_binary_request( + stage="asr", + endpoint="asr.example:50051", + payload=b"grpc-input", + protocol="grpc", + model="asr-model", + metadata={"input_names": ["audio"]}, + ) + capture_dir = next(tmp_path.iterdir()) + manifest = json.loads((capture_dir / "manifest.json").read_text()) + assert manifest["protocol"] == "grpc" + assert manifest["metadata"]["input_names"] == ["audio"] + assert (capture_dir / "request.bin").read_bytes() == b"grpc-input" + + +def test_environment_config_captures_query_requests(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("NEMO_RETRIEVER_INFERENCE_CAPTURE_URI", str(tmp_path)) + monkeypatch.setenv("NEMO_RETRIEVER_INFERENCE_CAPTURE_FAILURE_MODE", "required") + monkeypatch.setenv("NEMO_RETRIEVER_INFERENCE_CAPTURE_OPERATION", "query") + + record_json_request(stage="embed", endpoint="http://nim/v1/embeddings", payload={"input": ["q"]}) + + capture_dir = next(tmp_path.iterdir()) + manifest = json.loads((capture_dir / "manifest.json").read_text()) + assert manifest["operation"] == "query" + assert manifest["stage"] == "embed" + + +def test_embedding_replay_metadata_is_aligned_and_does_not_change_payload(tmp_path: Path) -> None: + row = pd.Series({"id": "chunk-1", "text": "hello", "metadata": {"source_id": "doc-1", "embedding": [1.0]}}) + replay_records = _embedding_replay_records([row], ["hello"]) + + with activate_inference_capture(InferenceCaptureConfig(str(tmp_path), failure_mode="required"), operation="ingest"): + record_json_request( + stage="embed", + endpoint="http://nim/v1/embeddings", + payload={"input": ["hello"], "input_type": "passage"}, + metadata={"replay": {"replay_version": 1, "records": replay_records}}, + ) + + capture_dir = next(tmp_path.iterdir()) + manifest = json.loads((capture_dir / "manifest.json").read_text()) + assert json.loads((capture_dir / "request.json").read_text()) == {"input": ["hello"], "input_type": "passage"} + record = manifest["metadata"]["replay"]["records"][0] + assert record["input_index"] == 0 + assert record["record"]["id"] == "chunk-1" + assert record["record"]["metadata"] == {"source_id": "doc-1"} diff --git a/nemo_retriever/tests/test_nim_tracing.py b/nemo_retriever/tests/test_nim_tracing.py index c2b0dda8d8..5e9ca4ec50 100644 --- a/nemo_retriever/tests/test_nim_tracing.py +++ b/nemo_retriever/tests/test_nim_tracing.py @@ -6,6 +6,8 @@ from __future__ import annotations +import json +from pathlib import Path from typing import Any import numpy as np @@ -196,6 +198,7 @@ def _post(*args: Any, headers: dict[str, str], **kwargs: Any) -> _Response: def test_internal_nim_client_grpc_passes_trace_context_headers_when_supported( monkeypatch: pytest.MonkeyPatch, exported_spans: list[Any], + tmp_path: Path, ) -> None: captured_headers: list[dict[str, str]] = [] @@ -259,13 +262,23 @@ def close(self) -> None: max_429_retries=1, ) - with tracing.start_span("test.parent"): - parent_trace_id = tracing.current_trace_id_hex() - parsed_output, batch_data = client._process_batch( - np.array([[1.0]], dtype=np.float32), - batch_data={"batch": 2}, - model_name="detector-grpc", - ) + from nemo_retriever.common.inference_capture import InferenceCaptureConfig, activate_inference_capture + + with activate_inference_capture(InferenceCaptureConfig(str(tmp_path), failure_mode="required"), operation="ingest"): + with tracing.start_span("test.parent"): + parent_trace_id = tracing.current_trace_id_hex() + parsed_output, batch_data = client._process_batch( + np.array([[1.0]], dtype=np.float32), + batch_data={"batch": 2}, + model_name="detector-grpc", + ) + + capture_dir = next(tmp_path.iterdir()) + manifest = json.loads((capture_dir / "manifest.json").read_text()) + assert manifest["protocol"] == "grpc" + assert manifest["model"] == "detector-grpc" + with np.load(capture_dir / "request.bin") as captured_inputs: + assert np.array_equal(captured_inputs["input"], np.array([[1.0]], dtype=np.float32)) assert isinstance(parsed_output, np.ndarray) assert batch_data == {"batch": 2}