feat(nemotron-parse): richer inference metrics, engine_kwargs passthrough, fanout stage spec

benchmarking/scripts/nemotron_parse_pdf_benchmark.py

@@ -28,7 +28,7 @@
 from typing import Any
 
 from loguru import logger
-from utils import collect_parquet_output_metrics, setup_executor, write_benchmark_results
+from utils import setup_executor, write_benchmark_results
 
 REPO_ROOT = Path(__file__).parent.parent.parent
 sys.path.insert(0, str(REPO_ROOT / "tutorials" / "interleaved" / "nemotron_parse_pdf"))
@@ -38,6 +38,26 @@
     create_nemotron_parse_pdf_pipeline,
 )
 
+from nemo_curator.tasks.utils import TaskPerfUtils  # noqa: E402
+
+
+def _safe_div(numerator: float, denominator: float) -> float:
+    return numerator / denominator if denominator else 0.0
+
+
+def _compute_pdf_parse_metrics(output_tasks: list, run_time_taken: float) -> dict[str, float]:
+    """Compute benchmark-level throughput metrics from additive task stats."""
+    task_metrics = TaskPerfUtils.aggregate_task_metrics(output_tasks, prefix="task")
+    metric_prefix = "task_nemotron_parse_inference_custom"
+
+    num_valid_pages = task_metrics.get(f"{metric_prefix}.num_valid_pages_sum", 0.0)
+    total_output_tokens = task_metrics.get(f"{metric_prefix}.total_output_tokens_sum", 0.0)
+
+    return {
+        "throughput_pages_per_sec": _safe_div(num_valid_pages, run_time_taken),
+        "throughput_output_tokens_per_sec": _safe_div(total_output_tokens, run_time_taken),
+    }
+
 
 def run_nemotron_parse_pdf_benchmark(args: argparse.Namespace) -> dict[str, Any]:
     """Run the Nemotron-Parse PDF benchmark and collect metrics."""
@@ -72,27 +92,13 @@ def run_nemotron_parse_pdf_benchmark(args: argparse.Namespace) -> dict[str, Any]
                     unique_samples.update(task.data.column("sample_id").to_pylist())
 
         num_pdfs_processed = len(unique_samples)
-        parquet_metrics = collect_parquet_output_metrics(output_dir)
-
-        stage_perf: dict[str, list[float]] = {}
-        for task in output_tasks:
-            for perf in task._stage_perf:
-                stage_perf.setdefault(perf.stage_name, []).append(perf.process_time)
-
-        stage_summary = {}
-        for stage_name, times in stage_perf.items():
-            stage_summary[stage_name] = {
-                "count": len(times),
-                "total_s": sum(times),
-                "mean_s": sum(times) / len(times) if times else 0,
-                "min_s": min(times) if times else 0,
-                "max_s": max(times) if times else 0,
-            }
+        pdf_parse_metrics = _compute_pdf_parse_metrics(output_tasks, run_time_taken)
 
         logger.success(f"Benchmark completed in {run_time_taken:.2f}s")
         logger.success(f"Processed {num_pdfs_processed} PDFs")
+        logger.success(f"Page throughput: {pdf_parse_metrics['throughput_pages_per_sec']:.2f} pages/s")
         logger.success(
-            f"Output: {parquet_metrics.get('num_rows', 0)} rows in {parquet_metrics.get('num_output_files', 0)} files"
+            f"Output token throughput: {pdf_parse_metrics['throughput_output_tokens_per_sec']:.2f} tokens/s"
         )
         success = True
 
@@ -102,8 +108,7 @@ def run_nemotron_parse_pdf_benchmark(args: argparse.Namespace) -> dict[str, Any]
         logger.debug(f"Full traceback:\n{error_traceback}")
         run_time_taken = time.perf_counter() - run_start_time
         num_pdfs_processed = 0
-        parquet_metrics = {}
-        stage_summary = {}
+        pdf_parse_metrics = {}
 
     return {
         "params": {
@@ -128,8 +133,7 @@ def run_nemotron_parse_pdf_benchmark(args: argparse.Namespace) -> dict[str, Any]
             "num_pdfs_processed": num_pdfs_processed,
             "num_output_tasks": len(output_tasks),
             "throughput_pdfs_per_sec": num_pdfs_processed / run_time_taken if run_time_taken > 0 else 0,
-            **parquet_metrics,
-            "stage_performance": stage_summary,
+            **pdf_parse_metrics,
         },
         "tasks": output_tasks,
     }

nemo_curator/stages/interleaved/pdf/nemotron_parse/inference.py

@@ -18,7 +18,9 @@
 
 import contextlib
 import io
+import time
 from dataclasses import dataclass, field
+from typing import Any
 
 import pyarrow as pa
 import torch
@@ -69,6 +71,9 @@ class NemotronParseInferenceStage(ProcessingStage[InterleavedBatch, InterleavedB
         Pages per GPU forward pass (HF backend only).
     max_num_seqs
         Maximum concurrent sequences (vLLM backend only).
+    engine_kwargs
+        Extra keyword arguments forwarded to the vLLM engine (e.g.
+        ``gpu_memory_utilization``, ``max_num_batched_tokens``). vLLM backend only.
     """
 
     model_path: str = DEFAULT_MODEL_PATH
@@ -78,6 +83,7 @@ class NemotronParseInferenceStage(ProcessingStage[InterleavedBatch, InterleavedB
     inference_batch_size: int = 4
     max_num_seqs: int = 64
     enforce_eager: bool = False
+    engine_kwargs: dict[str, Any] | None = None
     name: str = "nemotron_parse_inference"
     resources: Resources = field(default_factory=lambda: Resources(cpus=4.0, gpus=1.0))
 
@@ -134,11 +140,12 @@ def _setup_vllm(self) -> None:
         from nemo_curator.utils.vllm_utils import create_vllm_llm, resolve_local_model_path
 
         resolved_path = resolve_local_model_path(self.model_path)
-        self._llm = create_vllm_llm(
-            resolved_path,
-            max_num_seqs=self.max_num_seqs,
-            enforce_eager=self.enforce_eager,
-        )
+        engine_kwargs = {
+            "max_num_seqs": self.max_num_seqs,
+            "enforce_eager": self.enforce_eager,
+            **(self.engine_kwargs or {}),
+        }
+        self._llm = create_vllm_llm(resolved_path, **engine_kwargs)
         self._sampling_params = SamplingParams(
             temperature=0,
             top_k=1,
@@ -188,23 +195,80 @@ def _reset_vllm(self) -> None:
             torch.cuda.empty_cache()
         self._setup_vllm()
 
-    def _infer_vllm(self, images: list[Image.Image]) -> list[str]:
+    @staticmethod
+    def _vllm_metrics_from_outputs(  # noqa: PLR0913
+        outputs: list[Any],
+        *,
+        inference_time_s: float,
+        num_input_pages: int,
+        num_valid_pages: int,
+        num_skipped_pages: int,
+        vllm_retries: int = 0,
+    ) -> dict[str, float]:
+        """Build additive per-task vLLM metrics for TaskPerfUtils aggregation."""
+        total_prompt_tokens = 0
+        total_output_tokens = 0
+        total_output_chars = 0
+        num_length_truncated = 0
+        num_empty_outputs = 0
+
+        for req_out in outputs:
+            prompt_ids = getattr(req_out, "prompt_token_ids", None)
+            if prompt_ids is not None:
+                total_prompt_tokens += len(prompt_ids)
+
+            if not req_out.outputs:
+                num_empty_outputs += 1
+                continue
+
+            completion = req_out.outputs[0]
+            token_ids = getattr(completion, "token_ids", None)
+            if token_ids is not None:
+                total_output_tokens += len(token_ids)
+
+            text = getattr(completion, "text", "") or ""
+            total_output_chars += len(text)
+            if not text.strip():
+                num_empty_outputs += 1
+
+            if getattr(completion, "finish_reason", None) == "length":
+                num_length_truncated += 1
+
+        return {
+            "vllm_inference_time": inference_time_s,
+            "num_input_pages": float(num_input_pages),
+            "num_valid_pages": float(num_valid_pages),
+            "num_skipped_pages": float(num_skipped_pages),
+            "total_prompt_tokens": float(total_prompt_tokens),
+            "total_output_tokens": float(total_output_tokens),
+            "total_output_chars": float(total_output_chars),
+            "num_output_length_truncated": float(num_length_truncated),
+            "num_empty_outputs": float(num_empty_outputs),
+            "vllm_retries": float(vllm_retries),
+        }
+
+    def _infer_vllm(self, images: list[Image.Image]) -> tuple[list[str], list[Any], int]:
         if not images:
-            return []
+            return [], [], 0
         prompts = [{"prompt": self.task_prompt, "multi_modal_data": {"image": img}} for img in images]
 
         max_retries = 3
+        vllm_retries = 0
         for attempt in range(1, max_retries + 1):
             try:
                 outputs = self._llm.generate(prompts, self._sampling_params)
-                return [output.outputs[0].text for output in outputs]
             except Exception as e:
                 logger.warning(f"[vLLM] Inference failed (attempt {attempt}/{max_retries}): {e}")
                 if attempt < max_retries:
+                    vllm_retries += 1
                     self._reset_vllm()
                 else:
                     raise
-        return []
+            else:
+                texts = [output.outputs[0].text if output.outputs else "" for output in outputs]
+                return texts, outputs, vllm_retries
+        msg = "unreachable"
+        raise RuntimeError(msg)
 
     def _infer_hf(self, images: list[Image.Image]) -> list[str]:
         all_outputs: list[str] = []
@@ -233,19 +297,42 @@ def _infer_hf_single_fallback(self, images: list[Image.Image]) -> list[str]:
     def process(self, task: InterleavedBatch) -> InterleavedBatch | None:
         task_df = task.to_pandas()
         images = []
+        image_t0 = time.perf_counter()
         for idx, b in enumerate(task_df["binary_content"]):
             try:
                 images.append(Image.open(io.BytesIO(b)))
             except Exception as e:  # noqa: BLE001
                 logger.warning(f"Skipping page {idx} in {task.task_id}: {e}")
                 images.append(None)
-
+        self._log_metrics({"image_load_time": time.perf_counter() - image_t0})
         valid_mask = [img is not None for img in images]
         valid_images = [img for img in images if img is not None]
         if not valid_images:
             return None
 
-        valid_outputs = self._infer_vllm(valid_images) if self.backend == "vllm" else self._infer_hf(valid_images)
+        if self.backend == "vllm":
+            t0 = time.perf_counter()
+            valid_outputs, raw_outputs, vllm_retries = self._infer_vllm(valid_images)
+            inference_time_s = time.perf_counter() - t0
+            self._log_metrics(
+                self._vllm_metrics_from_outputs(
+                    raw_outputs,
+                    inference_time_s=inference_time_s,
+                    num_input_pages=len(images),
+                    num_valid_pages=len(valid_images),
+                    num_skipped_pages=len(images) - len(valid_images),
+                    vllm_retries=vllm_retries,
+                )
+            )
+        else:
+            valid_outputs = self._infer_hf(valid_images)
+            self._log_metrics(
+                {
+                    "num_input_pages": float(len(images)),
+                    "num_valid_pages": float(len(valid_images)),
+                    "num_skipped_pages": float(len(images) - len(valid_images)),
+                }
+            )
 
         all_outputs = []
         valid_iter = iter(valid_outputs)

nemo_curator/stages/interleaved/pdf/nemotron_parse/partitioning.py

@@ -22,6 +22,7 @@
 
 from loguru import logger
 
+from nemo_curator.backends.utils import RayStageSpecKeys
 from nemo_curator.stages.base import ProcessingStage
 from nemo_curator.stages.resources import Resources
 from nemo_curator.tasks import EmptyTask, FileGroupTask
@@ -80,6 +81,9 @@ def inputs(self) -> tuple[list[str], list[str]]:
     def outputs(self) -> tuple[list[str], list[str]]:
         return [], []
 
+    def ray_stage_spec(self) -> dict[str, Any]:
+        return {RayStageSpecKeys.IS_FANOUT_STAGE: True}
+
     def xenna_stage_spec(self) -> dict[str, Any]:
         return {"num_workers_per_node": 1}