feat(profiling): add semantic batch NVTX ranges

[jioffe502]

Description

Add a stable, decorator-based NVTX hierarchy for batch profiling:

• nrl.batch::pipeline.ingest covers one complete RayDataExecutor run (_ingest), including return_results=False summary mode.

• nrl.batch::... ranges cover Ray actor startup, actor batches (PDF split, PDF extract, Page Elements, Table Structure, OCR, embedding), materialization, terminal streaming, and retained-result concatenation. Local GPU actors and endpoint-backed CPU actors use the same stage names.

• Local model operations are annotated with @nvtx.annotate directly on existing high-level methods. No methods exist only to hold a decorator:

  • nrl.model::ocr_v2.invoke, nrl.model::page_elements.invoke, nrl.model::table_structure.invoke
  • nrl.model::llama_nemotron_embed_vl.{embed,embed_queries,embed_images,embed_text_image} on both embedder backends: the HF class and LlamaNemotronEmbedVL1BV2VLLMEmbedder, which is what vLLM-backend library runs (e.g. NRB) use
  • nrl.model::vlm_captioner.caption_batch, nrl.model::parakeet_ctc.decode_batch

  These ranges cover the whole operation, including CPU input preparation and postprocessing. Correlate them with the CUDA rows in Nsight Systems to see the device portion.

• nrl.remote::nim.request covers the shared NIM HTTP request, including retries and backoff.

Breaking changes for profiling users:

• Remove nemo_retriever.common.nvtx (gpu_inference_range). It duplicated native NVTX behavior and coupled annotations to torch.cuda.nvtx. All in-repo callers are migrated.
• Remove the gpu_inference range and the dynamic Model | bs=N labels. gpu_inference was misleading because it was not GPU-only time. Commands using --capture-range=nvtx --nvtx-capture=gpu_inference should trigger on nrl.batch::pipeline.ingest instead. No in-repo code or docs used that capture filter.
• NemotronOCRV1 no longer emits model ranges. Live ingestion uses NemotronOCRV2.

Scope: profiling structure and labels only. Scheduling, resource allocation, model inputs and outputs, retry policy, and pipeline results are unchanged.

Validation:

• Full bo767 Nsight Systems run (H200, 26.08-core-single-gpu-library-main, at c732f89): 767 files / 54,730 pages, nDCG@10 0.777, exit 0. nrl.batch::pipeline.ingest appears on the driver and actor ranges appear in 24 Ray worker processes. All OCR and Page Elements model ranges are nested inside their stage's batch range. Neither gpu_inference nor nrl.remote::nim.request appears; every endpoint was local. In an earlier JP20 profile, the old gpu_inference range started and ended with the model range nested inside it, and the GPU was busy for only ~25–27% of it.
• The summary-mode pipeline fix and the vLLM embedder decorators landed after that run and are covered by unit tests.
• Targeted tests pass: 275 model/actor/executor/NIM/NVTX tests at c732f89, and 107 NVTX and embedder tests at the vLLM change.
• Decorator contract tests cover successful calls, original exceptions, the pipeline boundary in both result modes, the vLLM embedder, and the shared remote request boundary.
• Pre-commit and git diff --check pass.
• MkDocs is not installed in this environment, so the docs build was not run.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.

Link to Devin session: https://nvidia-cloud.devinenterprise.com/sessions/a2c7aa2145764e96a43c536ab8b8fbc5
Open in Devin Desktop: https://nvidia-cloud.devinenterprise.com/desktop/session/a2c7aa2145764e96a43c536ab8b8fbc5?variant=devin

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge; the previously reported profiling-boundary, import, documentation, licensing, and constructor-signature concerns are resolved or withdrawn.

Summary

This PR replaces the legacy PyTorch-coupled profiling helper with native nvtx annotations and introduces a stable semantic hierarchy across batch execution, Ray actors, local model operations, materialization, streaming, and remote NIM requests.

• Adds nvtx as a core dependency so CPU-only and torch-free installations can emit annotations.
• Covers the complete _ingest execution boundary in both retained-result and summary modes.
• Uses consistent batch-stage names across local GPU and endpoint-backed CPU actors.
• Documents the profiling commands, range meanings, concurrency caveats, and intentional removal of gpu_inference.
• Adds contract tests for range boundaries, exceptions, torch-free imports, vLLM embedding, and pipeline result modes.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  I["nrl.batch::pipeline.ingest"] --> P["PDF split / extraction batches"]
  P --> A["Actor startup ranges"]
  A --> B["Actor batch ranges"]
  B --> M["Local nrl.model::* operations"]
  B --> R["Remote nrl.remote::nim.request"]
  B --> D["nrl.batch::ray.materialize"]
  B --> S["nrl.batch::pipeline.terminal_stream"]
  S --> C["nrl.batch::result.concat when retained frames exist"]

Loading

_{Reviews (22) · Last reviewed commit: "Merge branch 'main' into jioffe502/batch..." · Reviewed by Greptile}

[greptile-apps]

Comments Outside Diff

These findings could not be posted inline.

•  Compatibility import removed nemo_retriever/src/nemo_retriever/common/nvtx.py:1 ▶

  Deleting nemo_retriever.common.nvtx means existing integrations that import gpu_inference_range now fail with ModuleNotFoundError instead of receiving the documented deprecation warning. The previous implementation and tests preserved this compatibility entry point and verified its labels, results, and exception behavior. This violates the repository's backward-compatibility directive; retain the module for a deprecation cycle and keep its migration documentation.

•  Public NVTX API removed nemo_retriever/src/nemo_retriever/common/nvtx.py:46 ▶

  Deleting nemo_retriever.common.nvtx also removes its documented gpu_inference_range function. Existing integrations that import this helper will now fail with ModuleNotFoundError. This violates the repository requirement that public API removals use a deprecation cycle with migration documentation.

[jioffe502]

/devin add nvtx.annotate ranges for pdf_split, pdf_extract, and table_structure (CPU+GPU) and document them in the performance guide

[devin-ai-integration]

On it. Plan: nrl.batch::pdf_split.batch and nrl.batch::pdf_extract.batch on the CPU actors (these stages have no GPU variant, and their constructors only store kwargs, so no startup range), plus nrl.batch::table_structure.startup / .batch on both TableStructureActor (GPU) and TableStructureCPUActor. The performance guide table will list them in pipeline order, and the subprocess range test will cover each one.

Written by Devin

[devin-ai-integration]

Done in 3bad715:

• PDFSplitCPUActor.process → nrl.batch::pdf_split.batch
• PDFExtractionCPUActor.process → nrl.batch::pdf_extract.batch
• TableStructureActor (GPU) and TableStructureCPUActor: __init__ → nrl.batch::table_structure.startup, process → nrl.batch::table_structure.batch

The archetype wrappers are left unannotated so the ranges don't double up. The performance guide table lists the new ranges in pipeline order. test_batch_actors_emit_stage_ranges checks each one, including remote resolution to TableStructureCPUActor and the GPU table actor built with endpoint URLs. A Node/cloudpickle test confirms the annotated keyword-only TableStructureCPUActor constructor still round-trips.

Written by Devin

[devin-ai-integration]

Devin Review found 2 new potential issues.