feat(profiling): add semantic batch NVTX ranges

docs/docs/extraction/performance_guide.md

@@ -77,6 +77,88 @@ Related batch-size, CPU, and GPU-per-actor flags are documented in the [CLI inge
 
 Use the Ray dashboard to verify the available-resource snapshot and the planned worker allocation when you tune throughput.
 
+## Profile batch ingestion with Nsight Systems
+
+Batch mode emits NVIDIA Tools Extension (NVTX) ranges from the driver and Ray
+actor processes through the `nvtx` Python package, which is installed with NeMo
+Retriever Library. These host annotations do not require PyTorch or a GPU.
+Use NVIDIA Nsight Systems to correlate these semantic stages
+with CUDA and operating system activity. Start with a representative, bounded
+input because profiler tracing can increase runtime and report size.
+
+Run a batch ingest under `nsys profile` as follows:
+
+```bash
+nsys profile \
+  --trace=nvtx,cuda,osrt \
+  --sample=none \
+  --wait=all \
+  --output=nrl-batch-profile \
+  -- retriever ingest batch /path/to/your/pdfs
+```
+
+The command writes `nrl-batch-profile.nsys-rep`. Open the report in the Nsight
+Systems graphical interface, or run `nsys-ui nrl-batch-profile.nsys-rep` on a
+workstation with the interface installed. Search the timeline for `nrl.batch::`,
+`nrl.model::`, and `nrl.remote::` to find the following ranges.
+
+| Range | Interpretation |
+| --- | --- |
+| `nrl.batch::pipeline.ingest` | One complete `RayDataExecutor` run on the driver, including `return_results=False` summary mode. Use this range for batch pipeline wall-clock time. Driver-side setup before the run, such as building and normalizing manifest branch datasets, is outside this range. |
+| `nrl.batch::pdf_split.batch` | One PDF split actor call that splits a batch of documents into single-page PDFs. |
+| `nrl.batch::pdf_extract.batch` | One PDF extraction actor call that extracts text and images from a batch of pages. |
+| `nrl.batch::page_elements.startup` | One Page Elements actor constructor. A backend that loads lazily can perform more startup work in the first batch. |
+| `nrl.batch::page_elements.batch` | One Page Elements actor processing call. |
+| `nrl.batch::table_structure.startup` | One Table Structure actor constructor. A backend that loads lazily can perform more startup work in the first batch. |
+| `nrl.batch::table_structure.batch` | One Table Structure actor processing call. |
+| `nrl.batch::ocr.startup` | One OCR actor constructor. A backend that loads lazily can perform more startup work in the first batch. |
+| `nrl.batch::ocr.batch` | One OCR actor processing call. |
+| `nrl.batch::embedding.startup` | One embedding actor constructor. A backend that loads lazily can perform more startup work in the first batch. |
+| `nrl.batch::embedding.batch` | One embedding actor processing call. |
+| `nrl.batch::ray.materialize` | Materialization of a terminal Ray dataset into a pandas DataFrame. This range can include lazy upstream execution. |
+| `nrl.batch::pipeline.terminal_stream` | Consumption of the lazy upstream pipeline and streaming of its records to the vector database backend. This range is not pure vector database time. |
+| `nrl.batch::result.concat` | Concatenation of retained frames when the graph ends at a streaming vector database sink and at least one batch was consumed. This range is absent when the sink has downstream nodes or the input has no batches. |
+
+Local and remote (endpoint-backed) Page Elements, Table Structure, OCR, and embedding actors emit the same ranges; for remote actors, the batch range includes network requests to the endpoint.
+
+Local model operations also emit stable, decorator-based ranges. These ranges
+wrap existing high-level methods, so they can include CPU input preparation,
+one or more backend calls, and output processing. Compare them with CUDA
+activity in the Nsight Systems timeline to isolate device execution.
+
+| Range | Interpretation |
+| --- | --- |
+| `nrl.model::llama_nemotron_embed_vl.embed` | One local document text embedding operation. |
+| `nrl.model::llama_nemotron_embed_vl.embed_queries` | One local query embedding operation. |
+| `nrl.model::llama_nemotron_embed_vl.embed_images` | One local image embedding operation. |
+| `nrl.model::llama_nemotron_embed_vl.embed_text_image` | One local paired text-image embedding operation. |
+| `nrl.model::ocr_v2.invoke` | One end-to-end Nemotron OCR v2 wrapper call, including input conversion and all backend calls. This is the model used by the current local OCR actor. |
+| `nrl.model::page_elements.invoke` | One local Page Elements model call, including the initialization check and model execution. Output scaling runs in the separate `postprocess` call, outside this range. |
+| `nrl.model::table_structure.invoke` | One local Table Structure operation, including input preparation, model execution, and output scaling. |
+| `nrl.model::vlm_captioner.caption_batch` | One local visual language model batch captioning operation. |
+| `nrl.model::parakeet_ctc.decode_batch` | One local Parakeet CTC chunk-batch operation, including feature preparation, model execution, and greedy decoding. |
+
+Remote NVIDIA Inference Microservice (NIM) calls emit
+`nrl.remote::nim.request`. This range covers one HTTP request and its retry and
+backoff lifecycle. Compare it with the enclosing actor batch range to separate
+remote service wait time from local preparation and response processing.
+
+Ranges from concurrent Ray actors can overlap. Their durations include CPU
+preparation, waits, and calls into inference backends. They do not represent
+CUDA kernel time by themselves. In the timeline, compare range boundaries with
+the CUDA activity rows to identify the stage that remains active near pipeline
+completion and to find idle gaps between batches.
+
+You can also print an aggregate NVTX summary:
+
+```bash
+nsys stats --report nvtx_sum nrl-batch-profile.nsys-rep
+```
+
+The summary adds durations from all matching ranges. With concurrent actors,
+the summed duration can exceed wall-clock time, so use the timeline to reason
+about overlap and the critical path.
+
 ## Tune remote OCR request batching
 
 Remote OCR batches cropped regions across the page rows supplied to one OCR actor call. This behavior applies to in-process, batch, and service ingestion with a remote OCR NIM. It preserves page and region output order.

nemo_retriever/pyproject.toml

@@ -57,6 +57,7 @@ classifiers = [
 dependencies = [
   # Core orchestration and framework
   "backoff",
+  "nvtx>=0.2.15",
   "ray[data,serve]>=2.56.1",
   "ffmpeg-python",
   "pandas>=2.0,<3",

nemo_retriever/src/nemo_retriever/graph/executor.py

@@ -11,6 +11,7 @@
 import math
 import time
 from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Set
+import nvtx
 import pandas as pd
 
 if TYPE_CHECKING:
@@ -133,6 +134,7 @@ def arrow_table_to_pandas(table: Any) -> pd.DataFrame:
     return _normalize_object_tensor_columns(_materialize_row_unsafe_columns(table, frame))
 
 
+@nvtx.annotate("nrl.batch::ray.materialize", color="blue")
 def ray_dataset_to_pandas(dataset: ray.data.Dataset) -> pd.DataFrame:
     """Materialize a Ray Dataset without returning malformed Arrow arrays.
 
@@ -672,6 +674,7 @@ def ingest(
             raise TypeError(f"RayDataExecutor.ingest() does not accept setting(s): {unsupported}")
         return self._ingest(data, return_results=return_results)
 
+    @nvtx.annotate("nrl.batch::pipeline.ingest", color="blue")
     def _ingest(
         self,
         data: Any,
@@ -741,7 +744,8 @@ def input_batches() -> Iterator[pd.DataFrame]:
             return ray_dataset_to_pandas(downstream)
 
         if terminal_frames:
-            return pd.concat(terminal_frames, ignore_index=True)
+            with nvtx.annotate("nrl.batch::result.concat", color="blue"):
+                return pd.concat(terminal_frames, ignore_index=True)
         schema = dataset.schema()
         names = getattr(schema, "names", None)
         return pd.DataFrame(columns=list(names) if names is not None else None)

nemo_retriever/src/nemo_retriever/models/local/llama_nemotron_embed_vl_1b_v2_embedder.py

@@ -8,11 +8,11 @@
 from dataclasses import dataclass, field
 from typing import Any, Optional, Sequence
 
+import nvtx
 import torch
 
 from nemo_retriever.models.hf_cache import configure_global_hf_cache_base
 from nemo_retriever.models.embed_model_spec import resolve_embed_model_revision
-from nemo_retriever.common.nvtx import gpu_inference_range
 
 
 def _l2_normalize(x: torch.Tensor, eps: float = 1e-12) -> torch.Tensor:
@@ -112,6 +112,7 @@ def _set_p_max_length(self, modality: str) -> None:
         if p is not None and hasattr(self._model, "processor"):
             self._model.processor.p_max_length = p
 
+    @nvtx.annotate("nrl.model::llama_nemotron_embed_vl.embed", color="blue")
     def embed(self, texts: Sequence[str], *, batch_size: int = 64) -> torch.Tensor:
         """Embed document texts. Returns CPU tensor ``[N, D]``."""
         self._ensure_loaded()
@@ -121,10 +122,10 @@ def embed(self, texts: Sequence[str], *, batch_size: int = 64) -> torch.Tensor:
         with torch.inference_mode(), warnings.catch_warnings():
             warnings.filterwarnings("ignore", message="`input_embeds` is deprecated", category=FutureWarning)
             self._set_p_max_length("text")
-            with gpu_inference_range("LlamaNemotronEmbedVL1B", batch_size=len(texts_list), mode="doc_text"):
-                out = self._model.encode_documents(texts=texts_list)
+            out = self._model.encode_documents(texts=texts_list)
         return _to_normalized_cpu(out)
 
+    @nvtx.annotate("nrl.model::llama_nemotron_embed_vl.embed_queries", color="blue")
     def embed_queries(self, texts: Sequence[str], *, batch_size: int = 64) -> torch.Tensor:
         """Embed query strings. Returns CPU tensor ``[N, D]``."""
         self._ensure_loaded()
@@ -133,10 +134,10 @@ def embed_queries(self, texts: Sequence[str], *, batch_size: int = 64) -> torch.
             return torch.empty((0, self.output_dimension), dtype=torch.float32)
         with torch.inference_mode(), warnings.catch_warnings():
             warnings.filterwarnings("ignore", message="`input_embeds` is deprecated", category=FutureWarning)
-            with gpu_inference_range("LlamaNemotronEmbedVL1B", batch_size=len(texts_list), mode="query"):
-                out = self._model.encode_queries(texts_list)
+            out = self._model.encode_queries(texts_list)
         return _to_normalized_cpu(out)
 
+    @nvtx.annotate("nrl.model::llama_nemotron_embed_vl.embed_images", color="blue")
     def embed_images(self, images_b64: Sequence[str], *, batch_size: int = 64) -> torch.Tensor:
         """Embed images (base64-encoded). Returns CPU tensor ``[N, D]``."""
         self._ensure_loaded()
@@ -146,10 +147,10 @@ def embed_images(self, images_b64: Sequence[str], *, batch_size: int = 64) -> to
         with torch.inference_mode(), warnings.catch_warnings():
             warnings.filterwarnings("ignore", message="`input_embeds` is deprecated", category=FutureWarning)
             self._set_p_max_length("image")
-            with gpu_inference_range("LlamaNemotronEmbedVL1B", batch_size=len(image_dicts), mode="doc_image"):
-                out = self._model.encode_documents(images=image_dicts)
+            out = self._model.encode_documents(images=image_dicts)
         return _to_normalized_cpu(out)
 
+    @nvtx.annotate("nrl.model::llama_nemotron_embed_vl.embed_text_image", color="blue")
     def embed_text_image(
         self, texts: Sequence[str], images_b64: Sequence[str], *, batch_size: int = 64
     ) -> torch.Tensor:
@@ -166,8 +167,7 @@ def embed_text_image(
         with torch.inference_mode(), warnings.catch_warnings():
             warnings.filterwarnings("ignore", message="`input_embeds` is deprecated", category=FutureWarning)
             self._set_p_max_length("text_image")
-            with gpu_inference_range("LlamaNemotronEmbedVL1B", batch_size=len(paired_images), mode="doc_text_image"):
-                out = self._model.encode_documents(texts=paired_texts, images=paired_images)
+            out = self._model.encode_documents(texts=paired_texts, images=paired_images)
         return _to_normalized_cpu(out)
 
     def unload(self) -> None:
@@ -259,6 +259,7 @@ def _finalize_vectors(self, vectors: Sequence[Sequence[float]]) -> torch.Tensor:
         tensor = torch.tensor(vectors, dtype=torch.float32)
         return _l2_normalize(tensor) if self.normalize else tensor
 
+    @nvtx.annotate("nrl.model::llama_nemotron_embed_vl.embed", color="blue")
     def embed(self, texts: Sequence[str], *, batch_size: int = 64) -> torch.Tensor:
         """Embed document texts. Returns CPU tensor ``[N, D]``."""
         self._ensure_loaded()
@@ -276,6 +277,7 @@ def embed(self, texts: Sequence[str], *, batch_size: int = 64) -> torch.Tensor:
         )
         return self._finalize_vectors(vectors)
 
+    @nvtx.annotate("nrl.model::llama_nemotron_embed_vl.embed_queries", color="blue")
     def embed_queries(self, texts: Sequence[str], *, batch_size: int = 64) -> torch.Tensor:
         """Embed query strings. Returns CPU tensor ``[N, D]``."""
         self._ensure_loaded()
@@ -293,6 +295,7 @@ def embed_queries(self, texts: Sequence[str], *, batch_size: int = 64) -> torch.
         )
         return self._finalize_vectors(vectors)
 
+    @nvtx.annotate("nrl.model::llama_nemotron_embed_vl.embed_images", color="blue")
     def embed_images(self, images_b64: Sequence[str], *, batch_size: int = 64) -> torch.Tensor:
         """Embed images (base64-encoded). Returns CPU tensor ``[N, D]``."""
         self._ensure_loaded()
@@ -314,6 +317,7 @@ def embed_images(self, images_b64: Sequence[str], *, batch_size: int = 64) -> to
         )
         return self._finalize_vectors(vectors)
 
+    @nvtx.annotate("nrl.model::llama_nemotron_embed_vl.embed_text_image", color="blue")
     def embed_text_image(
         self, texts: Sequence[str], images_b64: Sequence[str], *, batch_size: int = 64
     ) -> torch.Tensor:

nemo_retriever/src/nemo_retriever/models/local/nemotron_ocr_v1.py

@@ -13,7 +13,6 @@
 import torch
 from nemo_retriever.models.hf_cache import configure_global_hf_cache_base
 from nemo_retriever.models.hf_model_registry import install_pinned_hf_hub_download
-from nemo_retriever.common.nvtx import gpu_inference_range
 from nemo_retriever.models.model import BaseModel, ModelRunMode
 
 from PIL import Image
@@ -188,23 +187,19 @@ def invoke(
                 out: List[Any] = []
                 for i in range(int(input_data.shape[0])):
                     b64 = self._tensor_to_png_b64(input_data[i])
-                    with gpu_inference_range("NemotronOCRv1", batch_size=1):
-                        out.extend(self._model(b64.encode("utf-8"), merge_level=merge_level))
+                    out.extend(self._model(b64.encode("utf-8"), merge_level=merge_level))
                 return out
             if input_data.ndim == 3:
                 b64 = self._tensor_to_png_b64(input_data)
-                with gpu_inference_range("NemotronOCRv1", batch_size=1):
-                    return self._model(b64.encode("utf-8"), merge_level=merge_level)
+                return self._model(b64.encode("utf-8"), merge_level=merge_level)
             raise ValueError(f"Unsupported torch tensor shape for OCR: {tuple(input_data.shape)}")
 
         # Disambiguate str: existing file path vs base64 string.
         if isinstance(input_data, str):
-            with gpu_inference_range("NemotronOCRv1", batch_size=1):
-                return self._model(input_data.encode("utf-8"), merge_level=merge_level)
+            return self._model(input_data.encode("utf-8"), merge_level=merge_level)
 
         # bytes / ndarray / BytesIO are supported directly by nemotron_ocr.
-        with gpu_inference_range("NemotronOCRv1", batch_size=1):
-            return self._model(input_data, merge_level=merge_level)
+        return self._model(input_data, merge_level=merge_level)
 
     @property
     def model_name(self) -> str:

nemo_retriever/src/nemo_retriever/models/local/nemotron_ocr_v2.py

@@ -10,6 +10,7 @@
 from pathlib import Path  # noqa: F401
 
 import numpy as np
+import nvtx
 import torch
 from nemo_retriever.models.hf_cache import configure_global_hf_cache_base
 from nemo_retriever.models.hf_model_registry import install_pinned_hf_hub_download
@@ -177,6 +178,7 @@ def _extract_text(obj: Any) -> str:
                     return " ".join(parts).strip()
         return str(obj).strip()
 
+    @nvtx.annotate("nrl.model::ocr_v2.invoke", color="blue")
     def invoke(
         self,
         input_data: Union[torch.Tensor, str, bytes, np.ndarray, io.BytesIO],

nemo_retriever/src/nemo_retriever/models/local/nemotron_page_elements_v3.py

@@ -5,11 +5,11 @@
 from typing import Any, Dict, List, Sequence, Tuple, Union, cast  # noqa: F401
 
 from torch import nn
+import nvtx
 import torch
 import numpy as np
 from nemo_retriever.models.hf_cache import configure_global_hf_cache_base
 from nemo_retriever.models.hf_model_registry import install_pinned_hf_hub_download
-from nemo_retriever.common.nvtx import gpu_inference_range
 from nemo_retriever.models.model import HuggingFaceModel, ModelRunMode
 
 import nemotron_page_elements_v3.model as _page_elements_model
@@ -102,6 +102,7 @@ def __call__(
     ) -> Union[List[Dict[str, torch.Tensor]], Dict[str, torch.Tensor]]:
         return self.invoke(input_data, orig_shape)
 
+    @nvtx.annotate("nrl.model::page_elements.invoke", color="blue")
     def invoke(
         self, input_data: torch.Tensor, orig_shape: Union[Tuple[int, int], Sequence[Tuple[int, int]]]
     ) -> Union[List[Dict[str, torch.Tensor]], Dict[str, torch.Tensor]]:
@@ -110,9 +111,8 @@ def invoke(
 
         with torch.inference_mode():
             with torch.autocast(device_type="cuda"):
-                with gpu_inference_range("PageElementsV3", batch_size=input_data.shape[0]):
-                    out = self._model(input_data, orig_shape)
-                    return out
+                out = self._model(input_data, orig_shape)
+                return out
         # preds0: Any
         # if isinstance(out, (list, tuple)) and len(out) > 0:
         #     preds0 = out[0]