diff --git a/.github/workflows/config/.secrets.baseline b/.github/workflows/config/.secrets.baseline index aeeb9a66f4..bd24ddd864 100644 --- a/.github/workflows/config/.secrets.baseline +++ b/.github/workflows/config/.secrets.baseline @@ -160,7 +160,7 @@ "filename": "fern/versions/main/pages/curate-text/synthetic/inference-server.mdx", "hashed_secret": "ce7501007f04a6529e650f1f1b3fc0586d1d94eb", "is_verified": false, - "line_number": 420 + "line_number": 422 } ], "fern/versions/main/pages/curate-text/synthetic/llm-client.mdx": [ @@ -534,5 +534,5 @@ } ] }, - "generated_at": "2026-09-02T18:31:29Z" + "generated_at": "2026-09-09T19:37:37Z" } diff --git a/benchmarking/4xGB200-64CPU.yaml b/benchmarking/4xGB200-64CPU.yaml index 4242767b9a..6e157f97e2 100644 --- a/benchmarking/4xGB200-64CPU.yaml +++ b/benchmarking/4xGB200-64CPU.yaml @@ -234,6 +234,12 @@ entries: - name: nemotron_parse_pdf_raydata timeout_s: 2400 + - name: nemotron_parse_pdf_inference_server_ray_serve + timeout_s: 2400 + + - name: nemotron_parse_pdf_inference_server_dynamo + timeout_s: 2400 + - name: alm_pipeline_xenna timeout_s: 1200 diff --git a/benchmarking/benchmarks.yaml b/benchmarking/benchmarks.yaml index 76411e79cb..ac5d463041 100644 --- a/benchmarking/benchmarks.yaml +++ b/benchmarking/benchmarks.yaml @@ -1666,17 +1666,6 @@ entries: --max-pages=10 --max-num-seqs=64 timeout_s: 1200 - sink_data: - - name: slack - additional_metrics: - - num_pdfs_processed - - num_pages_processed - - num_output_tokens - - num_output_tasks - - throughput_pages_per_sec - - throughput_output_tokens_per_sec - ping_on_failure: - - U082G3E46R0 # Abhinav Garg requirements: # Exact counts: the dataset is fixed and every file in it renders, so a # correct run reproduces these precisely. Verified identical across three @@ -1728,17 +1717,6 @@ entries: --max-pages=10 --max-num-seqs=64 timeout_s: 1200 - sink_data: - - name: slack - additional_metrics: - - num_pdfs_processed - - num_pages_processed - - num_output_tokens - - num_output_tasks - - throughput_pages_per_sec - - throughput_output_tokens_per_sec - ping_on_failure: - - U082G3E46R0 # Abhinav Garg requirements: # Exact counts: the dataset is fixed and every file in it renders, so a # correct run reproduces these precisely. Verified identical across three @@ -1773,6 +1751,71 @@ entries: min_value: 4630000 max_value: 5115000 + - name: nemotron_parse_pdf_inference_server_ray_serve + enabled: true + script: nemotron_parse_pdf_benchmark.py + args: >- + --benchmark-results-path={session_entry_dir} + --executor=ray_data + --manifest={dataset:nemotron_parse_pdf,manifest} + --pdf-dir={dataset:nemotron_parse_pdf,pdf_dir} + --output-dir={session_entry_dir}/scratch/output + --model-path={dataset:nemotron_parse_pdf_model,files} + --model-id=nvidia/NVIDIA-Nemotron-Parse-v1.2 + --backend=vllm + --inference-server-type=ray-serve + --inference-server-client-workers-per-replica=4 + --inference-batch-size=32 + --enforce-eager + --pdfs-per-task=20 + --max-pdfs=1278 + --max-pages=10 + timeout_s: 1800 + requirements: + - metric: num_pdfs_processed + exact_value: 1278 + - metric: num_pages_processed + exact_value: 5794 + - metric: throughput_pages_per_sec + min_value: 7.0 + - metric: throughput_output_tokens_per_sec + min_value: 6000 + - metric: num_output_tokens + min_value: 4630000 + max_value: 5115000 + + - name: nemotron_parse_pdf_inference_server_dynamo + enabled: true + script: nemotron_parse_pdf_benchmark.py + args: >- + --benchmark-results-path={session_entry_dir} + --executor=ray_data + --manifest={dataset:nemotron_parse_pdf,manifest} + --pdf-dir={dataset:nemotron_parse_pdf,pdf_dir} + --output-dir={session_entry_dir}/scratch/output + --model-path={dataset:nemotron_parse_pdf_model,files} + --model-id=nvidia/NVIDIA-Nemotron-Parse-v1.2 + --backend=vllm + --inference-server-type=dynamo + --inference-server-client-workers-per-replica=4 + --inference-batch-size=32 + --enforce-eager + --pdfs-per-task=20 + --max-pdfs=1278 + --max-pages=10 + timeout_s: 1800 + requirements: + - metric: num_pdfs_processed + exact_value: 1278 + - metric: num_pages_processed + exact_value: 5794 + - metric: throughput_pages_per_sec + min_value: 7.0 + - metric: throughput_output_tokens_per_sec + min_value: 6000 + - metric: num_output_tokens + min_value: 4630000 + max_value: 5115000 # CPU-only; the 8-GPU runtime target does not apply. - name: alm_pipeline_xenna enabled: true diff --git a/benchmarking/scripts/inference_server_utils.py b/benchmarking/scripts/inference_server_utils.py new file mode 100644 index 0000000000..ffb20ddb67 --- /dev/null +++ b/benchmarking/scripts/inference_server_utils.py @@ -0,0 +1,136 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# ruff: noqa: PLR0913 + +"""Shared helpers for inference servers used by benchmark scripts.""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING, Any, Literal + +if TYPE_CHECKING: + from nemo_curator.core.serve import InferenceServer + +InferenceServerBackend = Literal["ray-serve", "dynamo"] + + +def parse_json_object(value: str | None, *, argument: str) -> dict[str, Any]: + """Parse an optional command-line JSON object.""" + if value is None: + return {} + try: + parsed = json.loads(value) + except json.JSONDecodeError as error: + msg = f"{argument} must be valid JSON: {error}" + raise ValueError(msg) from error + if not isinstance(parsed, dict): + msg = f"{argument} must decode to a JSON object" + raise TypeError(msg) + return parsed + + +def static_num_replicas(autoscaling_config: dict[str, Any] | None) -> int: + """Resolve a fixed replica count from a Ray-style autoscaling config.""" + if not autoscaling_config: + return 1 + min_replicas = int(autoscaling_config.get("min_replicas", 1)) + max_replicas = int(autoscaling_config.get("max_replicas", min_replicas)) + if min_replicas != max_replicas: + msg = ( + "Dynamo does not support autoscaling in benchmarks; " + f"min_replicas ({min_replicas}) must equal max_replicas ({max_replicas})." + ) + raise ValueError(msg) + if min_replicas < 1: + msg = f"num_replicas must be at least 1, got {min_replicas}" + raise ValueError(msg) + return min_replicas + + +def start_inference_server( + *, + backend: InferenceServerBackend, + model_id: str, + num_replicas: int, + engine_kwargs: dict[str, Any] | None = None, + model_path: str | None = None, + model_runtime_env: dict[str, Any] | None = None, + dynamo_kwargs: dict[str, Any] | None = None, + dynamo_router_kwargs: dict[str, Any] | None = None, + dynamo_subprocess_env: dict[str, str] | None = None, + ray_serve_deployment_config: dict[str, Any] | None = None, + health_check_timeout_s: int = 900, +) -> InferenceServer: + """Build, start, and return an inference server. + + If ``model_path`` is set, the server loads weights from that local path + while exposing ``model_id`` as the served model name. + + ``model_runtime_env`` is passed to Ray Serve replicas or Dynamo workers. + For gpt-oss, set ``TIKTOKEN_RS_CACHE_DIR`` there to read the Harmony + encoding from a pre-populated local cache instead of downloading it from + Azure at startup (see https://github.com/openai/harmony/issues/101). + + ``health_check_timeout_s`` controls how long server startup waits for the + model to register at ``/v1/models``. + """ + from nemo_curator.core.serve import InferenceServer + + if num_replicas < 1: + msg = f"num_replicas must be at least 1, got {num_replicas}" + raise ValueError(msg) + if backend == "dynamo": + from nemo_curator.core.serve import DynamoRouterConfig, DynamoServerConfig, DynamoVLLMModelConfig + + model = DynamoVLLMModelConfig( + model_identifier=model_path or model_id, + model_name=model_id if model_path else None, + engine_kwargs=engine_kwargs or {}, + num_replicas=num_replicas, + dynamo_kwargs=dynamo_kwargs or {}, + runtime_env=model_runtime_env or {}, + ) + server = InferenceServer( + models=[model], + backend=DynamoServerConfig( + request_plane="tcp", + router=DynamoRouterConfig(router_kwargs=dynamo_router_kwargs or {}), + subprocess_env=dynamo_subprocess_env or {}, + ), + health_check_timeout_s=health_check_timeout_s, + ) + else: + if backend != "ray-serve": + msg = f"Unsupported inference server backend: {backend}" + raise ValueError(msg) + + from nemo_curator.core.serve import RayServeModelConfig + + model = RayServeModelConfig( + model_identifier=model_path or model_id, + model_name=model_id if model_path else None, + deployment_config=( + ray_serve_deployment_config + if ray_serve_deployment_config is not None + else {"num_replicas": num_replicas} + ), + engine_kwargs=engine_kwargs or {}, + runtime_env=model_runtime_env or {}, + ) + server = InferenceServer(models=[model], health_check_timeout_s=health_check_timeout_s) + + server.start() + return server diff --git a/benchmarking/scripts/ndd_benchmark.py b/benchmarking/scripts/ndd_benchmark.py index b773c73cf6..7a8b472d85 100644 --- a/benchmarking/scripts/ndd_benchmark.py +++ b/benchmarking/scripts/ndd_benchmark.py @@ -37,12 +37,12 @@ """ import argparse -import json import os import time from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import Any +from inference_server_utils import parse_json_object, start_inference_server, static_num_replicas from loguru import logger from utils import load_dataset_files, setup_executor, write_benchmark_results @@ -52,109 +52,6 @@ from nemo_curator.stages.text.io.writer.jsonl import JsonlWriter from nemo_curator.tasks.utils import TaskPerfUtils -if TYPE_CHECKING: - from nemo_curator.core.serve import InferenceServer - - -def _start_ray_serve_inference_server( - model_id: str, - engine_kwargs: dict[str, Any] | None = None, - autoscaling_config: dict[str, Any] | None = None, - model_path: str | None = None, - tiktoken_cache_dir: str | None = None, - health_check_timeout_s: int | None = None, -) -> "InferenceServer": - """Start a local Ray Serve-backed InferenceServer and return it. - - If ``model_path`` is set, vLLM loads weights from that local path while - ``model_id`` is used as the served name in ``/v1/models``. - - ``tiktoken_cache_dir``, if set, is passed to replicas as ``TIKTOKEN_RS_CACHE_DIR`` - so gpt-oss's harmony encoding is read from a pre-populated local cache instead of - being downloaded from Azure blob storage at startup (see openai/harmony#101). - - ``health_check_timeout_s`` controls how long ``_wait_for_models`` waits for the - model to register at ``/v1/models`` before raising ``SubprocessError``. - """ - from nemo_curator.core.serve import InferenceServer, RayServeModelConfig - from nemo_curator.core.serve.constants import DEFAULT_SERVE_HEALTH_TIMEOUT_S - - engine_kwargs = engine_kwargs or {} - autoscaling_config = autoscaling_config or {"min_replicas": 1, "max_replicas": 1} - runtime_env = {"env_vars": {"TIKTOKEN_RS_CACHE_DIR": tiktoken_cache_dir}} if tiktoken_cache_dir else {} - - server_config = RayServeModelConfig( - model_identifier=model_path or model_id, - model_name=model_id if model_path else None, - deployment_config={"autoscaling_config": autoscaling_config}, - engine_kwargs=engine_kwargs, - runtime_env=runtime_env, - ) - - server = InferenceServer( - models=[server_config], - health_check_timeout_s=health_check_timeout_s or DEFAULT_SERVE_HEALTH_TIMEOUT_S, - ) - server.start() - return server - - -def _start_dynamo_inference_server( - model_id: str, - engine_kwargs: dict[str, Any] | None = None, - autoscaling_config: dict[str, Any] | None = None, - model_path: str | None = None, - tiktoken_cache_dir: str | None = None, - health_check_timeout_s: int | None = None, -) -> "InferenceServer": - """Start a local Dynamo-backed InferenceServer and return it. - - Dynamo has no autoscaling — ``min_replicas`` and ``max_replicas`` (when - supplied) must match and are used as a static ``num_replicas``. - If ``model_path`` is set, vLLM loads weights from that local path while - ``model_id`` is used as the served name in ``/v1/models``. - - ``tiktoken_cache_dir``, if set, is passed to workers as ``TIKTOKEN_RS_CACHE_DIR`` - so gpt-oss's harmony encoding is read from a pre-populated local cache instead of - being downloaded from Azure blob storage at startup (see openai/harmony#101). - - ``health_check_timeout_s`` controls how long ``_wait_for_models`` waits for the - model to register at ``/v1/models`` before raising ``SubprocessError``. - """ - from nemo_curator.core.serve import DynamoServerConfig, DynamoVLLMModelConfig, InferenceServer - from nemo_curator.core.serve.constants import DEFAULT_SERVE_HEALTH_TIMEOUT_S - - engine_kwargs = engine_kwargs or {} - num_replicas = 1 - if autoscaling_config: - min_r = autoscaling_config.get("min_replicas", 1) - max_r = autoscaling_config.get("max_replicas", min_r) - if min_r != max_r: - msg = ( - f"Dynamo backend does not support autoscaling; min_replicas ({min_r}) " - f"must equal max_replicas ({max_r})." - ) - raise ValueError(msg) - num_replicas = min_r - - runtime_env = {"env_vars": {"TIKTOKEN_RS_CACHE_DIR": tiktoken_cache_dir}} if tiktoken_cache_dir else {} - - model_config = DynamoVLLMModelConfig( - model_identifier=model_path or model_id, - model_name=model_id if model_path else None, - engine_kwargs=engine_kwargs, - num_replicas=num_replicas, - runtime_env=runtime_env, - ) - - server = InferenceServer( - models=[model_config], - backend=DynamoServerConfig(), - health_check_timeout_s=health_check_timeout_s or DEFAULT_SERVE_HEALTH_TIMEOUT_S, - ) - server.start() - return server - def run_nemotron_cc_sdg_benchmark( # noqa: PLR0915 inference_server_type: str, @@ -193,18 +90,23 @@ def run_nemotron_cc_sdg_benchmark( # noqa: PLR0915 if inference_server_type in ("ray-serve", "dynamo"): logger.info(f"Starting local {inference_server_type} InferenceServer with engine_kwargs={engine_kwargs}") serve_start = time.perf_counter() - starter = ( - _start_ray_serve_inference_server + num_replicas = static_num_replicas(autoscaling_config) if inference_server_type == "dynamo" else 1 + ray_serve_deployment_config = ( + {"autoscaling_config": autoscaling_config or {"min_replicas": 1, "max_replicas": 1}} if inference_server_type == "ray-serve" - else _start_dynamo_inference_server + else None ) - inference_server = starter( - model_id, - engine_kwargs, - autoscaling_config, + inference_server = start_inference_server( + backend=inference_server_type, + model_id=model_id, model_path=model_path, - tiktoken_cache_dir=tiktoken_cache_dir, - health_check_timeout_s=health_check_timeout_s, + num_replicas=num_replicas, + engine_kwargs=engine_kwargs, + model_runtime_env=( + {"env_vars": {"TIKTOKEN_RS_CACHE_DIR": tiktoken_cache_dir}} if tiktoken_cache_dir else None + ), + ray_serve_deployment_config=ray_serve_deployment_config, + health_check_timeout_s=health_check_timeout_s or 300, ) serve_startup_s = time.perf_counter() - serve_start logger.info(f"InferenceServer ready at {inference_server.endpoint} (startup: {serve_startup_s:.1f}s)") @@ -372,9 +274,8 @@ def main() -> int: logger.info("=== Nemotron-CC SDG Benchmark Starting ===") logger.info(f"Arguments: {vars(args)}") - # Parse JSON string args - engine_kwargs = json.loads(args.engine_kwargs) if args.engine_kwargs else None - autoscaling_config = json.loads(args.autoscaling_config) if args.autoscaling_config else None + engine_kwargs = parse_json_object(args.engine_kwargs, argument="--engine-kwargs") + autoscaling_config = parse_json_object(args.autoscaling_config, argument="--autoscaling-config") success_code = 1 result_dict: dict[str, Any] = { diff --git a/benchmarking/scripts/nemotron_parse_pdf_benchmark.py b/benchmarking/scripts/nemotron_parse_pdf_benchmark.py index 0bccaa364a..6158ccb323 100644 --- a/benchmarking/scripts/nemotron_parse_pdf_benchmark.py +++ b/benchmarking/scripts/nemotron_parse_pdf_benchmark.py @@ -12,14 +12,17 @@ # See the License for the specific language governing permissions and # limitations under the License. +# ruff: noqa: ANN401, E402, PLR0915 + """Nemotron-Parse PDF pipeline benchmarking script. Reuses the pipeline and argparser from -tutorials/interleaved/nemotron_parse_pdf/main.py with comprehensive +tutorials/interleaved/nemotron_parse_pdf/pipeline_utils.py with comprehensive metrics collection. """ import argparse +import contextlib import json import sys import time @@ -27,25 +30,44 @@ from pathlib import Path from typing import Any +from inference_server_utils import InferenceServerBackend, parse_json_object from loguru import logger 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")) -from main import ( # noqa: E402 +from pipeline_utils import ( create_nemotron_parse_pdf_argparser, create_nemotron_parse_pdf_pipeline, ) -from nemo_curator.tasks.utils import TaskPerfUtils # noqa: E402 +from nemo_curator.backends.utils import get_available_cpu_gpu_resources +from nemo_curator.stages.interleaved.pdf.nemotron_parse import create_nemotron_parse_inference_server +from nemo_curator.tasks.utils import TaskPerfUtils def _safe_div(numerator: float, denominator: float) -> float: return numerator / denominator if denominator else 0.0 -def _sample_ids_from_table(data: Any) -> set[str]: # noqa: ANN401 +def _resolve_num_replicas(configured_num_replicas: int | None) -> int: + num_replicas = int(configured_num_replicas) if configured_num_replicas is not None else _available_gpu_count() + if num_replicas < 1: + msg = f"--num-replicas must be at least 1, got {num_replicas}." + raise ValueError(msg) + return num_replicas + + +def _available_gpu_count() -> int: + num_gpus = int(get_available_cpu_gpu_resources(init_and_shutdown=True)[1]) + if num_gpus < 1: + msg = f"Nemotron-Parse inference needs at least one GPU, found {num_gpus}." + raise RuntimeError(msg) + return num_gpus + + +def _sample_ids_from_table(data: Any) -> set[str]: """Pull sample ids from an in-memory Arrow table, if that is what the task carries.""" if data is None or not hasattr(data, "column"): return set() @@ -55,7 +77,7 @@ def _sample_ids_from_table(data: Any) -> set[str]: # noqa: ANN401 return set() -def _sample_ids_from_metadata(task: Any) -> set[str]: # noqa: ANN401 +def _sample_ids_from_metadata(task: Any) -> set[str]: """Derive sample ids from the manifest entries recorded in task metadata. The parquet writer returns a ``FileGroupTask`` whose ``data`` is a list of @@ -97,14 +119,25 @@ def _count_unique_pdfs(output_tasks: list) -> int: return len(unique) -def _compute_pdf_parse_metrics(output_tasks: list, run_time_taken: float) -> dict[str, float]: +def _compute_pdf_parse_metrics( + output_tasks: list, + run_time_taken: float, + num_inference_gpus: int, + inference_stage_parallelism: int, +) -> 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_input_tokens = task_metrics.get(f"{metric_prefix}.total_prompt_tokens_sum", 0.0) total_output_tokens = task_metrics.get(f"{metric_prefix}.total_output_tokens_sum", 0.0) + inference_stage_process_time_sum_s = task_metrics.get("task_nemotron_parse_inference_process_time_sum", 0.0) + 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) + inference_stage_active_time_s = _safe_div(inference_stage_process_time_sum_s, inference_stage_parallelism) + inference_stage_gpu_time_s = inference_stage_active_time_s * num_inference_gpus return { # Surfaced as first-class metrics (not just throughput denominators) so # entries can assert on work actually completed rather than on wall-clock @@ -113,8 +146,17 @@ def _compute_pdf_parse_metrics(output_tasks: list, run_time_taken: float) -> dic # model emits EOS), so it is asserted as a band rather than an exact value. "num_pages_processed": num_valid_pages, "num_output_tokens": total_output_tokens, - "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), + # Stage process time excludes model/server setup. Normalizing the sum + # across the stage's concurrent workers estimates the active inference + # wall time for both in-process and HTTP inference. These intermediate + # values are intentionally not exposed as top-level metrics. + "inference_stage_pages_per_sec_per_gpu": _safe_div(num_valid_pages, inference_stage_gpu_time_s), + "inference_stage_input_tokens_per_sec_per_gpu": _safe_div(total_input_tokens, inference_stage_gpu_time_s), + "inference_stage_output_tokens_per_sec_per_gpu": _safe_div(total_output_tokens, inference_stage_gpu_time_s), + "throughput_pages_per_sec": throughput_pages_per_sec, + "throughput_output_tokens_per_sec": throughput_output_tokens_per_sec, + "throughput_pages_per_sec_per_gpu": _safe_div(throughput_pages_per_sec, num_inference_gpus), + "throughput_output_tokens_per_sec_per_gpu": _safe_div(throughput_output_tokens_per_sec, num_inference_gpus), } @@ -125,19 +167,72 @@ def run_nemotron_parse_pdf_benchmark(args: argparse.Namespace) -> dict[str, Any] output_dir = Path(args.output_dir).absolute() output_dir.mkdir(parents=True, exist_ok=True) + inference_server = None + inference_server_startup_s = 0.0 + num_replicas = 0 + num_inference_gpus = 0 + inference_stage_parallelism = 0 + server_engine_kwargs: dict[str, Any] | None = None + server_type: InferenceServerBackend | None = args.inference_server_type + logger.info(f"Manifest: {args.manifest}") logger.info(f"PDF source: zip_base_dir={args.zip_base_dir}, pdf_dir={args.pdf_dir}") logger.info(f"Output: {output_dir}") logger.info(f"Model: {args.model_path}, backend={args.backend}") logger.info(f"PDFs per task: {args.pdfs_per_task}, max PDFs: {args.max_pdfs}") - pipeline = create_nemotron_parse_pdf_pipeline(args) - run_start_time = time.perf_counter() success = False output_tasks: list = [] try: + if server_type is not None: + if args.backend != "vllm": + msg = f"--inference-server-type requires --backend=vllm, got {args.backend!r}." + raise ValueError(msg) # noqa: TRY301 + if args.inference_server_client_workers_per_replica < 1: + msg = "--inference-server-client-workers-per-replica must be at least 1" + raise ValueError(msg) # noqa: TRY301 + num_replicas = _resolve_num_replicas(args.num_replicas) + num_inference_gpus = num_replicas + inference_stage_parallelism = args.inference_server_client_workers_per_replica * num_replicas + model_name = args.model_id or args.model_path + server_engine_kwargs = parse_json_object(args.engine_kwargs, argument="--engine-kwargs") + if args.enforce_eager: + server_engine_kwargs["enforce_eager"] = True + logger.info( + f"Starting {server_type} inference server with {num_replicas} replicas; " + f"PDF client stage workers={inference_stage_parallelism}" + ) + server_start = time.perf_counter() + inference_server = create_nemotron_parse_inference_server( + backend=server_type, + model_path=args.model_path, + model_name=model_name, + num_replicas=num_replicas, + engine_kwargs=server_engine_kwargs, + request_timeout_s=args.inference_server_request_timeout_s, + health_check_timeout_s=args.inference_server_health_timeout_s, + ) + server_engine_kwargs = inference_server.models[0].engine_kwargs + inference_server.start() + inference_server_startup_s = time.perf_counter() - server_start + pipeline = create_nemotron_parse_pdf_pipeline( + args, + inprocess_backend=args.backend, + inference_server_endpoint=inference_server.endpoint, + inference_server_model_name=model_name, + inference_server_client_num_workers=inference_stage_parallelism, + ) + logger.info( + f"Inference server ready at {inference_server.endpoint} after {inference_server_startup_s:.2f}s" + ) + else: + num_inference_gpus = _available_gpu_count() + inference_stage_parallelism = num_inference_gpus + pipeline = create_nemotron_parse_pdf_pipeline(args, inprocess_backend=args.backend) + + run_start_time = time.perf_counter() logger.info("Running Nemotron-Parse PDF pipeline...") logger.info(f"Pipeline description:\n{pipeline.describe()}") @@ -145,7 +240,12 @@ def run_nemotron_parse_pdf_benchmark(args: argparse.Namespace) -> dict[str, Any] run_time_taken = time.perf_counter() - run_start_time num_pdfs_processed = _count_unique_pdfs(output_tasks) - pdf_parse_metrics = _compute_pdf_parse_metrics(output_tasks, run_time_taken) + pdf_parse_metrics = _compute_pdf_parse_metrics( + output_tasks, + run_time_taken, + num_inference_gpus, + inference_stage_parallelism, + ) logger.success(f"Benchmark completed in {run_time_taken:.2f}s") logger.success(f"Processed {num_pdfs_processed} PDFs") @@ -153,7 +253,16 @@ def run_nemotron_parse_pdf_benchmark(args: argparse.Namespace) -> dict[str, Any] logger.success( f"Output token throughput: {pdf_parse_metrics['throughput_output_tokens_per_sec']:.2f} tokens/s" ) - success = True + logger.success( + "Inference-stage per-GPU throughput: " + f"{pdf_parse_metrics['inference_stage_pages_per_sec_per_gpu']:.2f} pages/s/GPU, " + f"{pdf_parse_metrics['inference_stage_output_tokens_per_sec_per_gpu']:.2f} output tokens/s/GPU, " + f"{pdf_parse_metrics['inference_stage_input_tokens_per_sec_per_gpu']:.2f} input tokens/s/GPU" + ) + if not num_pdfs_processed or not pdf_parse_metrics["num_pages_processed"]: + logger.error("Benchmark produced no PDFs or pages") + else: + success = True except Exception as e: error_traceback = traceback.format_exc() @@ -166,10 +275,20 @@ def run_nemotron_parse_pdf_benchmark(args: argparse.Namespace) -> dict[str, Any] pdf_parse_metrics = { "num_pages_processed": 0.0, "num_output_tokens": 0.0, + "inference_stage_pages_per_sec_per_gpu": 0.0, + "inference_stage_input_tokens_per_sec_per_gpu": 0.0, + "inference_stage_output_tokens_per_sec_per_gpu": 0.0, "throughput_pages_per_sec": 0.0, "throughput_output_tokens_per_sec": 0.0, + "throughput_pages_per_sec_per_gpu": 0.0, + "throughput_output_tokens_per_sec_per_gpu": 0.0, } + finally: + if inference_server is not None: + with contextlib.suppress(Exception): + inference_server.stop() + return { "params": { "executor": args.executor, @@ -180,16 +299,28 @@ def run_nemotron_parse_pdf_benchmark(args: argparse.Namespace) -> dict[str, Any] "benchmark_results_path": str(args.benchmark_results_path), "model_path": args.model_path, "backend": args.backend, + "inference_server_type": server_type, + "num_replicas": num_replicas, + "num_inference_gpus": num_inference_gpus, + "inference_server_client_workers_per_replica": ( + args.inference_server_client_workers_per_replica if server_type is not None else None + ), "pdfs_per_task": args.pdfs_per_task, "max_pdfs": args.max_pdfs, "dpi": args.dpi, "max_pages": args.max_pages, "inference_batch_size": args.inference_batch_size, - "max_num_seqs": args.max_num_seqs, + "max_num_seqs": ( + args.max_num_seqs if server_type is None else (server_engine_kwargs or {}).get("max_num_seqs") + ), + "max_tokens": args.max_tokens, + "enforce_eager": args.enforce_eager, + "server_engine_kwargs": server_engine_kwargs, }, "metrics": { "is_success": success, "time_taken_s": run_time_taken, + "inference_server_startup_s": inference_server_startup_s, "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, @@ -214,7 +345,58 @@ def main() -> int: choices=["xenna", "ray_data"], help="Executor to use for pipeline execution", ) - + parser.add_argument( + "--backend", + default="vllm", + choices=["hf", "vllm"], + help="In-process inference backend; inference-server runs require vllm", + ) + parser.add_argument( + "--inference-server-type", + choices=["ray-serve", "dynamo"], + default=None, + help="Run PDF inference through a managed vLLM Ray Serve or Dynamo server; requires --backend=vllm", + ) + parser.add_argument( + "--num-replicas", + type=int, + default=None, + help="Inference-server replicas; defaults to the GPU count reported by Ray", + ) + parser.add_argument( + "--inference-server-client-workers-per-replica", + type=int, + default=4, + help="Parallel HTTP client stage workers per inference-server replica", + ) + parser.add_argument( + "--model-id", + default=None, + help="Served model name; defaults to --model-path", + ) + parser.add_argument( + "--engine-kwargs", + default=None, + help="JSON object of additional vLLM engine arguments for the inference server", + ) + parser.add_argument( + "--inference-server-health-timeout-s", + type=int, + default=900, + help="Seconds to wait for the inference server to become healthy", + ) + parser.add_argument( + "--inference-server-request-timeout-s", + type=float, + default=300.0, + help="Timeout for each page inference request", + ) + parser.add_argument( + "--inference-server-max-retries", + type=int, + default=3, + help="Retries after a failed page inference request", + ) args = parser.parse_args() logger.info("=== Nemotron-Parse PDF Pipeline Benchmark Starting ===") diff --git a/fern/versions/main/pages/curate-text/load-data/nemotron-parse-pdf.mdx b/fern/versions/main/pages/curate-text/load-data/nemotron-parse-pdf.mdx index 713905619c..a19ba1b59f 100644 --- a/fern/versions/main/pages/curate-text/load-data/nemotron-parse-pdf.mdx +++ b/fern/versions/main/pages/curate-text/load-data/nemotron-parse-pdf.mdx @@ -1,5 +1,5 @@ --- -description: "Convert PDF datasets into interleaved Parquet output using NVIDIA's Nemotron-Parse VLM with the four-stage NemotronParsePDFReader composite pipeline" +description: "Convert PDF datasets into interleaved Parquet output with Nemotron-Parse and the recommended Dynamo InferenceServer deployment" categories: ["how-to-guides"] tags: ["pdf", "nemotron-parse", "interleaved", "vllm", "parsing"] personas: ["data-scientist-focused", "mle-focused"] @@ -18,7 +18,7 @@ Convert PDF datasets into interleaved Parquet output using NVIDIA's [Nemotron-Pa 1. **`PDFPartitioningStage`** — reads a JSONL manifest of PDF entries and packs them into `FileGroupTask` objects. 2. **`PDFPreprocessStage`** — extracts PDF bytes from the configured source, renders pages to images with scale-to-fit safeguarding against OOM on large pages. -3. **`NemotronParseInferenceStage`** — runs Nemotron-Parse via vLLM (recommended) or Hugging Face Transformers, with `text_in_pic` and `enforce_eager` flags and free-port retry on collisions. +3. **`NemotronParseHTTPClientStage`** (recommended) or **`NemotronParseInferenceStage`** — calls an OpenAI-compatible `InferenceServer`, or runs vLLM/Hugging Face Transformers in process. 4. **`NemotronParsePostprocessStage`** — parses model output, aligns images and captions, crops images, and emits the final interleaved rows. The output is interleaved Parquet ready to be filtered with [Interleaved Filters](/curate-text/process-data/interleaved/filters) and written to MINT-1T-style WebDataset shards. @@ -27,8 +27,9 @@ The output is interleaved Parquet ready to be filtered with [Interleaved Filters Choose your PDF source and confirm the prerequisites: -- **GPU**: Required. Nemotron-Parse runs on GPU via vLLM (recommended) or Hugging Face Transformers. -- **vLLM**: Strongly recommended for throughput. Falls back to HF Transformers if `backend="hf"` is set. +- **GPU**: Required. Use one inference-server model replica per GPU. +- **Inference server**: Recommended for production. It lets the serving layer batch page requests across pipeline tasks while rendering and postprocessing scale independently. Internal PDF-pipeline comparisons found this to be the better default; benchmark your corpus before further tuning. +- **In-process inference**: Retained for local validation and debugging. Use vLLM when possible, or set `backend="hf"` when vLLM is unavailable. - **`pypdfium2`**: Required Python dependency for PDF rendering. Installed automatically with the `interleaved_cpu` or `interleaved_cuda12` extras (e.g., `uv sync --extra interleaved_cuda12`). - **Manifest**: A JSONL file listing the PDFs to process. Each line should specify the PDF location relative to the source directory you choose. @@ -42,50 +43,99 @@ Pass exactly one of `pdf_dir`, `zip_base_dir`, or `jsonl_base_dir` so the prepro | `zip_base_dir` | A `CC-MAIN-2021-31-PDF-UNTRUNCATED` zip hierarchy | Common Crawl PDF dumps | | `jsonl_base_dir` | JSONL-encoded PDF datasets where each line carries the PDF bytes | GitHub-hosted PDF datasets, custom JSONL collections | -### Backend Selection +### Inference Topology -| Backend | When to Use | +| Topology | When to Use | |---------|-------------| -| `vllm` (recommended) | High-throughput GPU inference with batching. Set `enforce_eager=True` if you hit compilation issues. | -| `hf` | Hugging Face Transformers fallback when vLLM is unavailable or for debugging. | +| Dynamo `InferenceServer` + `NemotronParseHTTPClientStage` (recommended) | Production pipelines. Dynamo owns GPU replicas and batches requests across pipeline tasks. | +| In-process `vllm` | Small local runs or debugging without a separate server lifecycle. | +| In-process `hf` | Compatibility fallback when vLLM is unavailable. | -The inference stage retries transient port collisions when starting vLLM. Non-retryable startup failures, such as invalid configuration or GPU out-of-memory errors, fail immediately. +For the recommended topology, use a fixed HTTP stage pool of `4 * num_gpus` workers and start with `inference_batch_size=32` concurrent page requests per worker. We validated both 32 and 64 on 8 H100 GPUs and retained 32 because 64 improved aggregate page throughput only slightly while increasing individual request latency. Because the best value depends on the GPU and corpus, benchmark 64 on the target workload and keep it only when it improves throughput without request failures or out-of-memory errors. The in-process stage retries transient port collisions when starting vLLM; non-retryable startup failures fail immediately. --- ## Usage -A minimal end-to-end pipeline that reads PDFs from a directory and writes interleaved Parquet: +A runnable tutorial entry point starts Dynamo, waits for its OpenAI-compatible +endpoint, runs the PDF pipeline, and stops the server. From a NeMo Curator +checkout, install both dependency groups: + +```bash +uv sync --extra interleaved_cuda12 --extra inference_server +``` + +Dynamo starts local `etcd` and `nats-server` processes. The NeMo Curator +container includes both binaries; for a source environment outside the +container, install them with +[`docker/common/install_etcd_nats.sh`](https://github.com/NVIDIA-NeMo/Curator/blob/main/docker/common/install_etcd_nats.sh) +before running the entry point. + +Create a manifest with one PDF filename per line: + +```jsonl +{"file_name": "document.pdf"} +``` + +Start Dynamo and run the pipeline with one command: + +```bash +uv run python tutorials/interleaved/nemotron_parse_pdf/main.py \ + --manifest ./pdfs.jsonl \ + --pdf-dir /data/pdfs \ + --output-dir ./parsed_pdfs \ + --model-path nvidia/NVIDIA-Nemotron-Parse-v1.2 \ + --inference-batch-size 32 +``` + +The entry point detects the Ray-visible GPUs and configures one +`DynamoVLLMModelConfig` replica per GPU. It passes the resulting +`server.endpoint` to `create_nemotron_parse_pdf_pipeline`, fixes the HTTP stage +pool at `4 * num_gpus` workers, and uses Ray Data for the pipeline. Use +`CUDA_VISIBLE_DEVICES` or your Ray cluster resources to select the GPUs. The +entry point defaults to 32 concurrent requests per HTTP worker; pass +`--inference-batch-size 64` to compare the higher concurrency on your corpus. + +The server configuration uses Dynamo's TCP request plane, enables multimodal +vLLM, and sets the frontend and model-worker options needed by Nemotron-Parse. +The shared `create_nemotron_parse_inference_server` helper owns those +model-specific settings so tutorials and benchmarks use the same configuration. +It does not use `dyn_chat_processor`. See the runnable +[`main.py`](https://github.com/NVIDIA-NeMo/Curator/blob/main/tutorials/interleaved/nemotron_parse_pdf/main.py) +source for the complete lifecycle and [Inference Server](/curate-text/synthetic/inference-server) +for the underlying configuration objects. For executor options, refer to +[Execution Backends](/reference/infra/execution-backends). + +`main.py` always starts Dynamo with vLLM and does not expose a backend option. +For local validation without an inference server, run `inprocess.py` with +`--backend vllm` or `--backend hf`. + +For programmatic use, the helper returns an unstarted `InferenceServer`; use it +as a context manager to start it, wait for health, and stop it: ```python -from nemo_curator.pipeline import Pipeline -from nemo_curator.backends.xenna import XennaExecutor -from nemo_curator.stages.interleaved.pdf.nemotron_parse import NemotronParsePDFReader -from nemo_curator.stages.interleaved.io.writers.tabular import InterleavedParquetWriter +from nemo_curator.backends.utils import get_available_cpu_gpu_resources +from nemo_curator.stages.interleaved.pdf.nemotron_parse import ( + NemotronParsePDFReader, + create_nemotron_parse_inference_server, +) -pipeline = Pipeline(name="pdf_to_interleaved") +_, available_gpus = get_available_cpu_gpu_resources(init_and_shutdown=True) +num_gpus = int(available_gpus) -# 1. Parse PDFs into interleaved rows -pipeline.add_stage( - NemotronParsePDFReader( +with create_nemotron_parse_inference_server( + model_path="nvidia/NVIDIA-Nemotron-Parse-v1.2", + num_replicas=num_gpus, +) as server: + reader = NemotronParsePDFReader( manifest_path="./pdfs.jsonl", pdf_dir="/data/pdfs", - backend="vllm", - pdfs_per_task=10, - max_pages=50, - inference_batch_size=4, + inference_server_endpoint=server.endpoint, + inference_server_client_num_workers=4 * num_gpus, + inference_batch_size=32, ) -) - -# 2. Write interleaved Parquet -pipeline.add_stage(InterleavedParquetWriter(output_dir="./parsed_pdfs")) - -executor = XennaExecutor() -pipeline.run(executor) ``` -For executor options and configuration, refer to [Execution Backends](/reference/infra/execution-backends). - ### Example: CC-MAIN PDF Dump Parse a Common Crawl PDF dump from its zip hierarchy: @@ -94,9 +144,12 @@ Parse a Common Crawl PDF dump from its zip hierarchy: NemotronParsePDFReader( manifest_path="./cc_pdfs.jsonl", zip_base_dir="/data/CC-MAIN-2021-31-PDF-UNTRUNCATED", - backend="vllm", file_names_field="cc_pdf_file_names", pdfs_per_task=20, + inference_server_endpoint=server.endpoint, + inference_server_model_name=model_name, + inference_server_client_num_workers=4 * num_gpus, + inference_batch_size=32, ) ``` @@ -108,7 +161,10 @@ Parse a JSONL-encoded dataset (e.g., GitHub-hosted PDFs where each line contains NemotronParsePDFReader( manifest_path="./github_pdfs.jsonl", jsonl_base_dir="/data/github_pdfs", - backend="vllm", + inference_server_endpoint=server.endpoint, + inference_server_model_name=model_name, + inference_server_client_num_workers=4 * num_gpus, + inference_batch_size=32, ) ``` @@ -121,13 +177,14 @@ NemotronParsePDFReader( | `zip_base_dir` | str \| None | `None` | Root directory of CC-MAIN PDF zip hierarchy. | | `jsonl_base_dir` | str \| None | `None` | Root directory of JSONL-encoded PDF datasets. | | `model_path` | str | `"nvidia/NVIDIA-Nemotron-Parse-v1.2"` | Local path or HF repo ID for the Nemotron-Parse weights. | -| `backend` | str | `"vllm"` | Inference backend (`vllm` or `hf`). | +| `backend` | str | `"vllm"` | In-process inference backend (`vllm` or `hf`); ignored by the HTTP stage. | | `pdfs_per_task` | int | `10` | Number of PDFs grouped into each `FileGroupTask`. | | `max_pdfs` | int \| None | `None` | Hard cap on total PDFs processed (debug aid). | | `dpi` | int | `300` | Render DPI for PDF pages. | | `max_pages` | int | `50` | Maximum pages rendered per PDF; longer PDFs are truncated. | -| `inference_batch_size` | int | `4` | vLLM/HF batch size. | +| `inference_batch_size` | int | `4` | Pages per HF pass or concurrent page requests per HTTP worker. The HTTP stage warns below the validated starting point of 32. | | `max_num_seqs` | int | `64` | Maximum concurrent vLLM sequences. | +| `max_tokens` | int | `8192` | Maximum output tokens generated per page. | | `text_in_pic` | bool | `False` | When `True`, treat embedded text within rendered images as part of the text content. | | `enforce_eager` | bool | `False` | Disable vLLM compilation for compatibility with restricted environments. | | `min_crop_px` | int | `10` | Minimum dimension (pixels) for cropped image regions. | @@ -135,8 +192,27 @@ NemotronParsePDFReader( | `file_name_field` | str | `"file_name"` | Manifest field naming a single PDF file. | | `file_names_field` | str | `"cc_pdf_file_names"` | Manifest field naming a list of PDF files (CC-MAIN layout). | | `url_field` | str | `"url"` | Manifest field for the source URL passthrough. | +| `inference_server_endpoint` | str \| None | `None` | OpenAI-compatible endpoint. Set this to use the recommended HTTP stage instead of in-process inference. | +| `inference_server_model_name` | str \| None | `None` | Served model name; defaults to `model_path`. | +| `inference_server_client_num_workers` | int | `4` | Fixed HTTP stage workers. Set to `4 * num_gpus` for production. | +| `inference_server_request_timeout_s` | float | `300.0` | Per-page HTTP request timeout. | +| `inference_server_max_retries` | int | `3` | Retries for transient HTTP failures. | + +### Failure Behavior + +PDF read and render failures happen in `PDFPreprocessStage`, before inference, +so they behave identically for in-process and inference-server runs: the bad PDF +is skipped, and a task with no renderable PDFs returns no output. After vLLM or +HTTP retries—or the Hugging Face single-page fallback—are exhausted, inference +failures raise the task instead of emitting a partially parsed document. The +executor's failure policy determines whether the overall run stops or records +that task as failed. + +### Tune In-Process vLLM -### Tune the vLLM Engine +Use this section only for the in-process fallback. For the recommended +inference-server topology, configure vLLM through `DynamoVLLMModelConfig` as shown in the +[Inference Server guide](/curate-text/synthetic/inference-server). `NemotronParseInferenceStage.engine_kwargs` is `None` by default. It passes additional settings to NeMo Curator's shared vLLM initializer: vLLM engine settings are forwarded to `vllm.LLM`, while helper settings such as `max_port_retries` control initialization itself. Use this field when you compose the pipeline from individual stages and need controls that are not exposed by `NemotronParsePDFReader`: @@ -247,7 +323,7 @@ The output is directly compatible with [Interleaved IO](/curate-text/process-dat ## Inspect Inference Metrics -`NemotronParseInferenceStage` records additive custom metrics on each output task. Aggregate the final pipeline results with `TaskPerfUtils`: +Both Nemotron-Parse inference stages record additive custom metrics on each output task. Aggregate the final pipeline results with `TaskPerfUtils`: ```python import time @@ -278,29 +354,31 @@ print(f"{output_tokens_per_second:.2f} output tokens/s") | Custom metric | Backends | Description | | --- | --- | --- | | `image_load_time` | vLLM, HF | Seconds spent decoding page-image bytes for the task. | -| `num_input_pages` | vLLM, HF | Page rows presented to the inference stage. | -| `num_valid_pages` | vLLM, HF | Pages successfully decoded and sent to the model. | -| `num_skipped_pages` | vLLM, HF | Pages skipped because their image bytes could not be decoded. | +| `num_input_pages` | vLLM, HF, HTTP | Page rows presented to the inference stage. | +| `num_valid_pages` | vLLM, HF, HTTP | Pages sent to the model. | +| `num_skipped_pages` | vLLM, HF, HTTP | Pages skipped because image data was unavailable or invalid. | | `vllm_inference_time` | vLLM | Seconds spent in vLLM generation, including retried inference attempts. | -| `total_prompt_tokens` | vLLM | Prompt tokens reported by vLLM across valid pages. | -| `total_output_tokens` | vLLM | Generated token count across valid pages. | -| `total_output_chars` | vLLM | Characters in generated text across valid pages. | -| `num_output_length_truncated` | vLLM | Completions whose vLLM finish reason was `length`. | -| `num_empty_outputs` | vLLM | Requests with no completion or blank completion text. | +| `inference_server_request_time` | HTTP | Seconds spent waiting for page requests in the HTTP stage. | +| `total_prompt_tokens` | vLLM, HTTP | Prompt tokens reported across valid pages. | +| `total_output_tokens` | vLLM, HTTP | Generated token count across valid pages. | +| `total_output_chars` | vLLM, HTTP | Characters in generated text across valid pages. | +| `num_output_length_truncated` | vLLM, HTTP | Completions whose finish reason was `length`. | +| `num_empty_outputs` | vLLM, HTTP | Requests with no completion or blank completion text. | | `vllm_retries` | vLLM | Inference-engine resets after generation failures. This does not count startup port-collision retries. | -The HF backend records image-loading and page-count metrics, but it does not expose vLLM token, character, truncation, or retry metrics. +The HF backend records image-loading and page-count metrics, but it does not expose token, character, truncation, or retry metrics. The HTTP stage reports server-request time rather than in-process image-loading or vLLM-engine time. -Use the quality signals alongside throughput. A high `num_output_length_truncated` value means outputs are reaching the stage's built-in 9,000-token generation limit and warrants inspection of those pages; `num_empty_outputs` and `num_skipped_pages` identify model-output and image-decoding failures that raw pages-per-second figures can hide. +Use the quality signals alongside throughput. A high `num_output_length_truncated` value means outputs are reaching the configured `max_tokens` limit (8,192 by default) and warrants inspection of those pages; `num_empty_outputs` and `num_skipped_pages` identify model-output and image failures that raw pages-per-second figures can hide. HTTP request failures are not counted as partial successes: after retries are exhausted, the task raises. -## vLLM Retry Behavior +## Retry Behavior -There are two separate retry paths: +There are three separate retry paths: 1. **Engine startup:** `create_vllm_llm()` chooses a new `MASTER_PORT` and retries up to `max_port_retries=3` times for direct address-in-use errors or vLLM v1's wrapped `Engine core initialization failed` error. Retries wait two to five seconds with jitter. Known non-retryable failures, including out-of-memory, device-side assertion, and invalid configuration errors, are raised immediately. 2. **Inference:** vLLM generation is attempted up to three times. After a failed attempt, the stage resets the engine before retrying. Successful retries contribute to the `vllm_retries` task metric; the final exception is raised after the third failed attempt. +3. **HTTP requests:** `NemotronParseHTTPClientStage` uses the existing `AsyncOpenAIClient` exponential-backoff retry behavior. Configure the retry count with `inference_server_max_retries`; a request that still fails raises the task so the pipeline cannot silently emit a partial document. If startup retries are exhausted, first check the worker log for the original failure. For repeated port collisions, reduce the number of vLLM replicas starting simultaneously or set a larger stage-level `max_port_retries` through `engine_kwargs`. Do not mask CUDA out-of-memory or invalid-model errors by increasing retries; tune memory-related engine settings or correct the model configuration instead. @@ -312,11 +390,11 @@ You don't need to configure this — it works automatically. If you find legitim ## Benchmarking -A standalone benchmark script ships at `benchmarking/scripts/nemotron_parse_pdf_benchmark.py`. It uses the same `TaskPerfUtils` aggregation shown above and reports end-to-end pages per second and output tokens per second. Use a representative manifest and the public tutorial arguments to compare configurations before scaling to your full corpus; the repository's nightly benchmark orchestration is not required. +A standalone benchmark script ships at `benchmarking/scripts/nemotron_parse_pdf_benchmark.py`. It uses `TaskPerfUtils` to report stage-normalized per-GPU pages, input tokens, and output tokens per second for both in-process and inference-server runs. The benchmark harness still records full-script `exec_time_s`, including inference-server startup; `time_taken_s` covers `pipeline.run()`, and `inference_server_startup_s` reports server startup separately. Use a representative manifest and the public tutorial arguments to compare configurations before scaling to your full corpus; the repository's nightly benchmark orchestration is not required. ## Best Practices -- **Use vLLM unless you can't**: the `vllm` backend is substantially faster than `hf`. Only fall back to `hf` for debugging or in environments where vLLM is unavailable. +- **Use Dynamo serving for production**: keep one model replica per GPU and four HTTP workers per GPU. Start at 32 concurrent requests per worker, then benchmark any change on the target GPU and corpus. Retain in-process vLLM and HF for small validation runs or debugging. - **Cap `max_pages` for outliers**: very long PDFs (1000+ pages) can dominate runtime. The default 50 pages handles most academic papers and articles; raise to 200+ for book-length sources. - **Tune `pdfs_per_task` for parallelism**: smaller values (5–10) parallelize better across many GPUs; larger values (20–50) reduce per-task overhead on smaller clusters. - **Set `enforce_eager=True` in restricted environments**: vLLM's torch.compile path can fail on certain hosts. Disabling compilation trades throughput for compatibility. @@ -324,6 +402,7 @@ A standalone benchmark script ships at `benchmarking/scripts/nemotron_parse_pdf_ ## Related Topics +- **[Inference Server](/curate-text/synthetic/inference-server)** — Dynamo lifecycle, model configuration, and troubleshooting. - **[Interleaved IO](/curate-text/process-data/interleaved/io)** — readers and writers that consume the Parquet output of this pipeline. - **[Interleaved Filters](/curate-text/process-data/interleaved/filters)** — sample-level filters to apply after parsing. - **[Common Crawl](/curate-text/load-data/common-crawl)** — companion source for web-scale PDF input via CC-MAIN dumps. diff --git a/fern/versions/main/pages/curate-text/synthetic/inference-server.mdx b/fern/versions/main/pages/curate-text/synthetic/inference-server.mdx index c8e6fd3466..04a87ced17 100644 --- a/fern/versions/main/pages/curate-text/synthetic/inference-server.mdx +++ b/fern/versions/main/pages/curate-text/synthetic/inference-server.mdx @@ -47,6 +47,11 @@ the NVIDIA package index. Local GPU serving through this extra supports **x86_64 Linux** and **aarch64 Linux**. The vLLM, ai-dynamo, and NIXL dependencies are not installed on macOS. +The NeMo Curator container also includes the `etcd` and `nats-server` binaries +that Dynamo starts. When using a source environment outside the container, +install them with +[`docker/common/install_etcd_nats.sh`](https://github.com/NVIDIA-NeMo/Curator/blob/main/docker/common/install_etcd_nats.sh), +or configure existing endpoints. The stack uses: @@ -314,7 +319,6 @@ backend = DynamoServerConfig( router=DynamoRouterConfig( mode="kv", kv_events=False, - router_kwargs={"dyn_chat_processor": "vllm"}, ), ) ``` @@ -331,15 +335,13 @@ backend = DynamoServerConfig( Additional `router_kwargs` are forwarded to the Dynamo frontend. Do not put `router_mode` or `router_kv_events` in that dictionary; use the typed fields instead. Boolean false values are emitted as `--no-*` flags. -For multimodal OpenAI content arrays, the Dynamo frontend can require: +Multimodal workers need their model-specific settings, such as +`limit_mm_per_prompt` in `engine_kwargs` and `enable_multimodal` in +`dynamo_kwargs`. Use Dynamo's default frontend processing path. -```python -DynamoRouterConfig( - router_kwargs={"dyn_chat_processor": "vllm"}, -) -``` - -The worker also needs its model-specific multimodal settings, such as `limit_mm_per_prompt` in `engine_kwargs` and `enable_multimodal` in `dynamo_kwargs`. + +The default frontend path avoids the slower compatibility processor used by older examples. + ## Runtime Environments and Subprocess Variables @@ -437,6 +439,18 @@ with InferenceServer(models=[model]) as server: CPU-only pipelines can use any executor while the server is active. + +For Nemotron-Parse PDF pipelines, use the Dynamo backend with +`NemotronParseHTTPClientStage`. Start with one model replica per GPU, a fixed +HTTP stage pool of `4 * num_gpus` workers, and 32 concurrent requests per HTTP +worker. This starting point was validated on 8 H100 GPUs. Because request +concurrency is hardware- and corpus-dependent, benchmark 64 on the target +workload and keep it only if throughput improves without request failures or +out-of-memory errors. The +[Nemotron-Parse PDF guide](/curate-text/load-data/nemotron-parse-pdf) provides +the exact install and run commands. + + ## Migrate from `InferenceModelConfig` Before: diff --git a/nemo_curator/models/client/llm_client.py b/nemo_curator/models/client/llm_client.py index d406cbed84..0f14b2b4f9 100644 --- a/nemo_curator/models/client/llm_client.py +++ b/nemo_curator/models/client/llm_client.py @@ -15,11 +15,14 @@ import asyncio import secrets from abc import ABC, abstractmethod -from collections.abc import Iterable +from collections.abc import Awaitable, Callable, Iterable from dataclasses import dataclass +from typing import TypeVar from loguru import logger +_T = TypeVar("_T") + class ConversationFormatter(ABC): """ @@ -116,7 +119,7 @@ async def _query_model_impl( msg = "Subclass of AsyncLLMClient must implement '_query_model_impl'" raise NotImplementedError(msg) - async def query_model( # noqa: C901, PLR0912 + async def query_model( self, *, messages: Iterable, @@ -133,7 +136,17 @@ async def query_model( # noqa: C901, PLR0912 elif isinstance(generation_config, dict): generation_config = GenerationConfig(**generation_config) - # Initialize semaphore if not already done or if we're in a different event loop + return await self._execute_with_retries( + lambda: self._query_model_impl( + messages=messages, + model=model, + conversation_formatter=conversation_formatter, + generation_config=generation_config, + ) + ) + + async def _execute_with_retries(self, request: Callable[[], Awaitable[_T]]) -> _T: # noqa: C901, PLR0912 + """Run an async request with the client's concurrency and retry policy.""" current_loop = asyncio.get_running_loop() if self._semaphore is None or self._semaphore_loop != current_loop: self._semaphore = asyncio.Semaphore(self.max_concurrent_requests) @@ -179,12 +192,7 @@ async def query_model( # noqa: C901, PLR0912 # Attempt the query try: - return await self._query_model_impl( - messages=messages, - model=model, - conversation_formatter=conversation_formatter, - generation_config=generation_config, - ) + return await request() except Exception as e: last_exception = e # If this is the last attempt, provide helpful error message @@ -207,8 +215,5 @@ async def query_model( # noqa: C901, PLR0912 if last_exception: raise last_exception - # This should never be reached, but add explicit return for linter - logger.warning( - "Unexpected code path: AsyncLLMClient.query_model completed without returning a result or raising an exception" - ) - return [] + msg = "Async request completed without returning a result or raising an exception" + raise RuntimeError(msg) diff --git a/nemo_curator/models/client/openai_client.py b/nemo_curator/models/client/openai_client.py index 3ca232fa1e..24d89f1af1 100644 --- a/nemo_curator/models/client/openai_client.py +++ b/nemo_curator/models/client/openai_client.py @@ -17,6 +17,7 @@ from loguru import logger from openai import AsyncOpenAI, OpenAI +from openai.types.chat import ChatCompletion from nemo_curator.models.client.llm_client import AsyncLLMClient, ConversationFormatter, GenerationConfig, LLMClient @@ -122,6 +123,41 @@ async def _query_model_impl( """ Internal implementation of query_model without retry/concurrency logic. """ + response = await self._query_model_response_impl( + messages=messages, + model=model, + conversation_formatter=conversation_formatter, + generation_config=generation_config, + ) + + return [choice.message.content for choice in response.choices] + + async def query_model_response( + self, + *, + messages: Iterable, + model: str, + conversation_formatter: ConversationFormatter | None = None, + generation_config: GenerationConfig | dict | None = None, + ) -> ChatCompletion: + """Query a model and return its raw response with retry and concurrency control.""" + return await self._execute_with_retries( + lambda: self._query_model_response_impl( + messages=messages, + model=model, + conversation_formatter=conversation_formatter, + generation_config=generation_config, + ) + ) + + async def _query_model_response_impl( + self, + *, + messages: Iterable, + model: str, + conversation_formatter: ConversationFormatter | None = None, + generation_config: GenerationConfig | dict | None = None, + ) -> ChatCompletion: if conversation_formatter is not None: warnings.warn("conversation_formatter is not used in an AsyncOpenAIClient", stacklevel=2) @@ -155,6 +191,4 @@ async def _query_model_impl( if not hasattr(self, "client"): self.setup() - response = await self.client.chat.completions.create(**create_kwargs) - - return [choice.message.content for choice in response.choices] + return await self.client.chat.completions.create(**create_kwargs) diff --git a/nemo_curator/stages/interleaved/pdf/nemotron_parse/__init__.py b/nemo_curator/stages/interleaved/pdf/nemotron_parse/__init__.py index b27ead0b32..815923cb58 100644 --- a/nemo_curator/stages/interleaved/pdf/nemotron_parse/__init__.py +++ b/nemo_curator/stages/interleaved/pdf/nemotron_parse/__init__.py @@ -13,15 +13,21 @@ # limitations under the License. from nemo_curator.stages.interleaved.pdf.nemotron_parse.composite import NemotronParsePDFReader -from nemo_curator.stages.interleaved.pdf.nemotron_parse.inference import NemotronParseInferenceStage +from nemo_curator.stages.interleaved.pdf.nemotron_parse.inference import ( + NemotronParseHTTPClientStage, + NemotronParseInferenceStage, +) from nemo_curator.stages.interleaved.pdf.nemotron_parse.partitioning import PDFPartitioningStage from nemo_curator.stages.interleaved.pdf.nemotron_parse.postprocess import NemotronParsePostprocessStage from nemo_curator.stages.interleaved.pdf.nemotron_parse.preprocess import PDFPreprocessStage +from nemo_curator.stages.interleaved.pdf.nemotron_parse.server import create_nemotron_parse_inference_server __all__ = [ + "NemotronParseHTTPClientStage", "NemotronParseInferenceStage", "NemotronParsePDFReader", "NemotronParsePostprocessStage", "PDFPartitioningStage", "PDFPreprocessStage", + "create_nemotron_parse_inference_server", ] diff --git a/nemo_curator/stages/interleaved/pdf/nemotron_parse/composite.py b/nemo_curator/stages/interleaved/pdf/nemotron_parse/composite.py index 29da17fe86..9476cc2903 100644 --- a/nemo_curator/stages/interleaved/pdf/nemotron_parse/composite.py +++ b/nemo_curator/stages/interleaved/pdf/nemotron_parse/composite.py @@ -20,7 +20,9 @@ from nemo_curator.stages.base import CompositeStage, ProcessingStage from nemo_curator.stages.interleaved.pdf.nemotron_parse.inference import ( + DEFAULT_MAX_TOKENS, DEFAULT_MODEL_PATH, + NemotronParseHTTPClientStage, NemotronParseInferenceStage, ) from nemo_curator.stages.interleaved.pdf.nemotron_parse.partitioning import PDFPartitioningStage @@ -37,7 +39,7 @@ class NemotronParsePDFReader(CompositeStage[EmptyTask, InterleavedBatch]): 1. :class:`PDFPartitioningStage` — read manifest, create FileGroupTasks 2. :class:`PDFPreprocessStage` — extract PDFs, render pages to images - 3. :class:`NemotronParseInferenceStage` — GPU model inference + 3. :class:`NemotronParseInferenceStage` or :class:`NemotronParseHTTPClientStage` — model inference 4. :class:`NemotronParsePostprocessStage` — parse output, align, crop Parameters @@ -63,7 +65,9 @@ class NemotronParsePDFReader(CompositeStage[EmptyTask, InterleavedBatch]): max_pages Maximum pages to render per PDF. inference_batch_size - Pages per GPU forward pass (HF only). + Pages per GPU forward pass for HF, or maximum concurrent page requests + from each HTTP client worker for an inference server. Start with 32; + tune on the target hardware and corpus. max_num_seqs Maximum concurrent sequences (vLLM only). text_in_pic @@ -78,6 +82,13 @@ class NemotronParsePDFReader(CompositeStage[EmptyTask, InterleavedBatch]): JSONL field containing a list of PDF filenames (CC-MAIN style). url_field JSONL field containing the source URL. + inference_server_endpoint + OpenAI-compatible inference server endpoint. A Dynamo-backed server is + the recommended production topology; when omitted, inference runs in + process. + inference_server_client_num_workers + Fixed number of concurrent HTTP client stage workers. Use four times + the number of inference GPUs. """ manifest_path: str | None = None @@ -92,6 +103,7 @@ class NemotronParsePDFReader(CompositeStage[EmptyTask, InterleavedBatch]): max_pages: int = 50 inference_batch_size: int = 4 max_num_seqs: int = 64 + max_tokens: int = DEFAULT_MAX_TOKENS text_in_pic: bool = False enforce_eager: bool = False min_crop_px: int = 10 @@ -99,6 +111,11 @@ class NemotronParsePDFReader(CompositeStage[EmptyTask, InterleavedBatch]): file_name_field: str = "file_name" file_names_field: str = "cc_pdf_file_names" url_field: str = "url" + inference_server_endpoint: str | None = None + inference_server_model_name: str | None = None + inference_server_client_num_workers: int = 4 + inference_server_request_timeout_s: float = 300.0 + inference_server_max_retries: int = 3 def __post_init__(self) -> None: super().__init__() @@ -121,14 +138,30 @@ def __post_init__(self) -> None: dpi=self.dpi, max_pages=self.max_pages, ) - self._inference = NemotronParseInferenceStage( - model_path=self.model_path, - text_in_pic=self.text_in_pic, - backend=self.backend, - inference_batch_size=self.inference_batch_size, - max_num_seqs=self.max_num_seqs, - enforce_eager=self.enforce_eager, - ) + if self.inference_server_endpoint is None: + self._inference = NemotronParseInferenceStage( + model_path=self.model_path, + text_in_pic=self.text_in_pic, + backend=self.backend, + inference_batch_size=self.inference_batch_size, + max_num_seqs=self.max_num_seqs, + max_tokens=self.max_tokens, + enforce_eager=self.enforce_eager, + ) + else: + if self.inference_server_client_num_workers < 1: + msg = "inference_server_client_num_workers must be at least 1" + raise ValueError(msg) + self._inference = NemotronParseHTTPClientStage( + endpoint=self.inference_server_endpoint, + model_name=self.inference_server_model_name or self.model_path, + model_path=self.model_path, + text_in_pic=self.text_in_pic, + request_timeout_s=self.inference_server_request_timeout_s, + max_retries=self.inference_server_max_retries, + inference_batch_size=self.inference_batch_size, + max_tokens=self.max_tokens, + ).with_(num_workers=self.inference_server_client_num_workers) self._postprocessor = NemotronParsePostprocessStage( min_crop_px=self.min_crop_px, ) diff --git a/nemo_curator/stages/interleaved/pdf/nemotron_parse/inference.py b/nemo_curator/stages/interleaved/pdf/nemotron_parse/inference.py index 9ee70cb19c..d01752c11f 100644 --- a/nemo_curator/stages/interleaved/pdf/nemotron_parse/inference.py +++ b/nemo_curator/stages/interleaved/pdf/nemotron_parse/inference.py @@ -12,10 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""GPU inference stage for Nemotron-Parse.""" +"""In-process and HTTP inference stages for Nemotron-Parse.""" from __future__ import annotations +import asyncio +import base64 import contextlib import io import time @@ -27,12 +29,38 @@ from loguru import logger from PIL import Image +from nemo_curator.models.client.llm_client import GenerationConfig +from nemo_curator.models.client.openai_client import AsyncOpenAIClient from nemo_curator.stages.base import ProcessingStage from nemo_curator.stages.resources import Resources from nemo_curator.tasks import InterleavedBatch DEFAULT_MODEL_PATH = "nvidia/NVIDIA-Nemotron-Parse-v1.2" PROMPT_BASE = "" +DEFAULT_MAX_TOKENS = 8192 +_RECOMMENDED_HTTP_CONCURRENCY = 32 + +_NEMOTRON_PARSE_SAMPLING_PARAMS: dict[str, Any] = { + "temperature": 0, + "top_p": 1.0, + "top_k": 1, + "repetition_penalty": 1.1, + "max_tokens": DEFAULT_MAX_TOKENS, + "skip_special_tokens": False, + "seed": None, +} +_OPENAI_GENERATION_CONFIG_FIELDS = {"max_tokens", "seed", "temperature", "top_p"} + + +def _nemotron_parse_sampling_params(max_tokens: int) -> dict[str, Any]: + return {**_NEMOTRON_PARSE_SAMPLING_PARAMS, "max_tokens": max_tokens} + + +def _nemotron_parse_server_generation_config(max_tokens: int) -> GenerationConfig: + sampling_params = _nemotron_parse_sampling_params(max_tokens) + config_params = {key: value for key, value in sampling_params.items() if key in _OPENAI_GENERATION_CONFIG_FIELDS} + extra_body = {key: value for key, value in sampling_params.items() if key not in config_params} + return GenerationConfig(**config_params, extra_kwargs={"extra_body": extra_body}) def build_task_prompt(*, text_in_pic: bool = False) -> str: @@ -71,6 +99,8 @@ class NemotronParseInferenceStage(ProcessingStage[InterleavedBatch, InterleavedB Pages per GPU forward pass (HF backend only). max_num_seqs Maximum concurrent sequences (vLLM backend only). + max_tokens + Maximum number of generated tokens (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. @@ -82,6 +112,7 @@ class NemotronParseInferenceStage(ProcessingStage[InterleavedBatch, InterleavedB backend: str = "vllm" inference_batch_size: int = 4 max_num_seqs: int = 64 + max_tokens: int = DEFAULT_MAX_TOKENS enforce_eager: bool = False engine_kwargs: dict[str, Any] | None = None name: str = "nemotron_parse_inference" @@ -146,13 +177,7 @@ def _setup_vllm(self) -> None: **(self.engine_kwargs or {}), } self._llm = create_vllm_llm(resolved_path, **engine_kwargs) - self._sampling_params = SamplingParams( - temperature=0, - top_k=1, - repetition_penalty=1.1, - max_tokens=9000, - skip_special_tokens=False, - ) + self._sampling_params = SamplingParams(**_nemotron_parse_sampling_params(self.max_tokens)) from transformers import AutoProcessor processor = AutoProcessor.from_pretrained(resolved_path, trust_remote_code=True) @@ -289,7 +314,7 @@ def _infer_hf_single_fallback(self, images: list[Image.Image]) -> list[str]: results.extend(self._infer_batch_hf([img])) except (RuntimeError, ValueError, TypeError) as e: logger.warning(f"Single page fallback failed: {e}") - results.append("") + raise return results # -- process -- @@ -351,3 +376,182 @@ def process(self, task: InterleavedBatch) -> InterleavedBatch | None: _metadata=metadata, _stage_perf=task._stage_perf, ) + + +@dataclass +class _HTTPPageResult: + text: str = "" + prompt_tokens: int = 0 + output_tokens: int = 0 + finish_reason: str | None = None + + +@dataclass +class NemotronParseHTTPClientStage(ProcessingStage[InterleavedBatch, InterleavedBatch]): + """Call Nemotron-Parse through an OpenAI-compatible HTTP endpoint. + + ``model_name`` is the served name used in requests. ``model_path`` is the + underlying model identifier recorded for postprocessing and defaults to the + served name. We validated ``inference_batch_size=32`` on 8 H100 GPUs as a + starting point; tune the concurrency on the target hardware and corpus. + ``proc_size`` must match the served model's image processor. + A page request that still fails after client retries raises the whole task, + matching in-process inference and preventing partial document output. + """ + + endpoint: str + model_name: str + model_path: str | None = None + text_in_pic: bool = False + task_prompt: str | None = None + request_timeout_s: float = 300.0 + max_retries: int = 3 + retry_base_delay_s: float = 1.0 + inference_batch_size: int = 4 + max_tokens: int = DEFAULT_MAX_TOKENS + proc_size: tuple[int, int] = (2048, 1664) + name: str = "nemotron_parse_inference" + resources: Resources = field(default_factory=lambda: Resources(cpus=1.0)) + + def __post_init__(self) -> None: + if self.task_prompt is None: + self.task_prompt = build_task_prompt(text_in_pic=self.text_in_pic) + self.max_retries = max(0, int(self.max_retries)) + self.inference_batch_size = int(self.inference_batch_size) + if self.inference_batch_size < 1: + msg = "inference_batch_size must be at least 1" + raise ValueError(msg) + if self.inference_batch_size < _RECOMMENDED_HTTP_CONCURRENCY: + logger.warning( + "NemotronParseHTTPClientStage inference_batch_size={} may underutilize the server; " + "start with 32 concurrent requests per HTTP worker and tune on the target hardware and corpus", + self.inference_batch_size, + ) + self._generation_config = _nemotron_parse_server_generation_config(self.max_tokens) + + def inputs(self) -> tuple[list[str], list[str]]: + return ["data"], [] + + def outputs(self) -> tuple[list[str], list[str]]: + return ["data"], [] + + def ray_stage_spec(self) -> dict[str, Any]: + return {"is_actor_stage": False} + + async def _query_page( + self, + client: AsyncOpenAIClient, + image_bytes: bytes, + content_type: str, + ) -> _HTTPPageResult: + image_url = f"data:{content_type};base64,{base64.b64encode(image_bytes).decode('ascii')}" + response = await client.query_model_response( + model=self.model_name, + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": self.task_prompt or ""}, + {"type": "image_url", "image_url": {"url": image_url}}, + ], + } + ], + generation_config=self._generation_config, + ) + + choice = response.choices[0] if response.choices else None + usage = response.usage + return _HTTPPageResult( + text=str(getattr(getattr(choice, "message", None), "content", "") or ""), + prompt_tokens=int(getattr(usage, "prompt_tokens", 0) or 0), + output_tokens=int(getattr(usage, "completion_tokens", 0) or 0), + finish_reason=getattr(choice, "finish_reason", None), + ) + + async def _query_pages(self, images: list[tuple[bytes, str]]) -> list[_HTTPPageResult]: + client = AsyncOpenAIClient( + max_concurrent_requests=self.inference_batch_size, + max_retries=self.max_retries, + base_delay=self.retry_base_delay_s, + api_key="unused", # pragma: allowlist secret + base_url=self.endpoint.rstrip("/"), + timeout=self.request_timeout_s, + ) + try: + return list( + await asyncio.gather( + *(self._query_page(client, image_bytes, content_type) for image_bytes, content_type in images) + ) + ) + finally: + if hasattr(client, "client"): + with contextlib.suppress(Exception): + await client.client.close() + + def _build_metrics( + self, + results: list[_HTTPPageResult], + *, + request_time_s: float, + num_input_pages: int, + num_valid_pages: int, + ) -> dict[str, float]: + total_output_tokens = float(sum(result.output_tokens for result in results)) + total_output_chars = float(sum(len(result.text) for result in results)) + return { + "inference_server_request_time": request_time_s, + "num_input_pages": float(num_input_pages), + "num_valid_pages": float(num_valid_pages), + "num_skipped_pages": float(num_input_pages - num_valid_pages), + "total_prompt_tokens": float(sum(result.prompt_tokens for result in results)), + "total_output_tokens": total_output_tokens, + "total_output_chars": total_output_chars, + "num_output_length_truncated": float(sum(result.finish_reason == "length" for result in results)), + "num_empty_outputs": float(sum(not result.text.strip() for result in results)), + } + + def process(self, task: InterleavedBatch) -> InterleavedBatch | None: + task_df = task.to_pandas() + valid_mask: list[bool] = [] + images: list[tuple[bytes, str]] = [] + + for _, row in task_df.iterrows(): + raw_bytes = row.get("binary_content") + try: + image_bytes = bytes(raw_bytes) if raw_bytes is not None else b"" + except TypeError: + image_bytes = b"" + is_valid = bool(image_bytes) + valid_mask.append(is_valid) + if is_valid: + images.append((image_bytes, str(row.get("content_type") or "image/png"))) + + if not images: + return None + + request_start = time.perf_counter() + results = asyncio.run(self._query_pages(images)) + request_time_s = time.perf_counter() - request_start + self._log_metrics( + self._build_metrics( + results, + request_time_s=request_time_s, + num_input_pages=len(valid_mask), + num_valid_pages=len(images), + ) + ) + + result_iter = iter(result.text for result in results) + task_df["text_content"] = [next(result_iter) if is_valid else "" for is_valid in valid_mask] + + metadata = dict(task._metadata) + metadata["proc_size"] = list(self.proc_size) + metadata["model_path"] = self.model_path or self.model_name + metadata["inference_server_endpoint"] = self.endpoint + + return InterleavedBatch( + dataset_name=task.dataset_name, + data=pa.Table.from_pandas(task_df, preserve_index=False), + _metadata=metadata, + _stage_perf=task._stage_perf, + ) diff --git a/nemo_curator/stages/interleaved/pdf/nemotron_parse/server.py b/nemo_curator/stages/interleaved/pdf/nemotron_parse/server.py new file mode 100644 index 0000000000..089a30afac --- /dev/null +++ b/nemo_curator/stages/interleaved/pdf/nemotron_parse/server.py @@ -0,0 +1,95 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Inference-server configuration for Nemotron-Parse.""" + +from __future__ import annotations + +from typing import Any, Literal + +from nemo_curator.core.serve import ( + DynamoRouterConfig, + DynamoServerConfig, + DynamoVLLMModelConfig, + InferenceServer, + RayServeModelConfig, +) +from nemo_curator.stages.interleaved.pdf.nemotron_parse.inference import DEFAULT_MODEL_PATH + +NemotronParseServerBackend = Literal["ray-serve", "dynamo"] + +_DEFAULT_ENGINE_KWARGS: dict[str, Any] = { + "trust_remote_code": True, + "dtype": "bfloat16", + "limit_mm_per_prompt": {"image": 1}, + "enable_prefix_caching": False, + "disable_hybrid_kv_cache_manager": False, +} + + +def create_nemotron_parse_inference_server( # noqa: PLR0913 + *, + model_path: str = DEFAULT_MODEL_PATH, + model_name: str | None = None, + backend: NemotronParseServerBackend = "dynamo", + num_replicas: int = 1, + engine_kwargs: dict[str, Any] | None = None, + request_timeout_s: float = 300.0, + health_check_timeout_s: int = 900, +) -> InferenceServer: + """Return an inference server configured for Nemotron-Parse PDFs. + + The returned server is not started. Use it as a context manager or call + :meth:`InferenceServer.start` and :meth:`InferenceServer.stop` explicitly. + """ + if num_replicas < 1: + msg = f"num_replicas must be at least 1, got {num_replicas}" + raise ValueError(msg) + if request_timeout_s < 1: + msg = f"request_timeout_s must be at least 1, got {request_timeout_s}" + raise ValueError(msg) + + resolved_engine_kwargs = {**_DEFAULT_ENGINE_KWARGS, **(engine_kwargs or {})} + model_kwargs = { + "model_identifier": model_path, + "model_name": model_name, + "engine_kwargs": resolved_engine_kwargs, + "runtime_env": {"uv": {"packages": ["albumentations==2.0.8"]}}, + } + + if backend == "dynamo": + model = DynamoVLLMModelConfig( + **model_kwargs, + num_replicas=num_replicas, + dynamo_kwargs={"enable_multimodal": True}, + ) + server_config = DynamoServerConfig( + request_plane="tcp", + router=DynamoRouterConfig(router_kwargs={"trust_remote_code": True}), + subprocess_env={"DYN_TCP_REQUEST_TIMEOUT": str(int(request_timeout_s))}, + ) + return InferenceServer( + models=[model], + backend=server_config, + health_check_timeout_s=health_check_timeout_s, + ) + if backend == "ray-serve": + model = RayServeModelConfig( + **model_kwargs, + deployment_config={"num_replicas": num_replicas}, + ) + return InferenceServer(models=[model], health_check_timeout_s=health_check_timeout_s) + + msg = f"Unsupported inference server backend: {backend}" + raise ValueError(msg) diff --git a/tests/models/client/test_openai_client.py b/tests/models/client/test_openai_client.py index 884d4e7fdf..650c13ca5f 100644 --- a/tests/models/client/test_openai_client.py +++ b/tests/models/client/test_openai_client.py @@ -460,6 +460,40 @@ def side_effect(*_args: object, **_kwargs: object) -> Mock: assert result == ["Success after retry"] assert call_count == 3 # Should have tried 3 times + @pytest.mark.asyncio + @patch("nemo_curator.models.client.openai_client.AsyncOpenAI") + async def test_query_model_response_preserves_metadata_and_uses_parent_retries( + self, mock_async_openai: Mock + ) -> None: + call_count = 0 + expected_response = Mock() + expected_response.choices = [Mock()] + expected_response.usage = Mock(prompt_tokens=5, completion_tokens=7) + + async def create_response(**_kwargs: object) -> Mock: + nonlocal call_count + call_count += 1 + if call_count == 1: + error_msg = "429 Rate limit exceeded" + raise RuntimeError(error_msg) + return expected_response + + mock_client = AsyncMock() + mock_client.chat.completions.create.side_effect = create_response + mock_async_openai.return_value = mock_client + client = AsyncOpenAIClient(max_retries=1, base_delay=0) + + with patch("nemo_curator.models.client.llm_client.asyncio.sleep", new_callable=AsyncMock): + response = await client.query_model_response( + messages=[{"role": "user", "content": "test"}], + model="gpt-4", + ) + + assert response is expected_response + assert response.usage.prompt_tokens == 5 + assert response.usage.completion_tokens == 7 + assert call_count == 2 + @pytest.mark.asyncio @patch("nemo_curator.models.client.openai_client.AsyncOpenAI") async def test_concurrent_request_limiting(self, mock_async_openai: Mock) -> None: diff --git a/tests/stages/interleaved/pdf/nemotron_parse/test_server.py b/tests/stages/interleaved/pdf/nemotron_parse/test_server.py new file mode 100644 index 0000000000..c1060c1b1f --- /dev/null +++ b/tests/stages/interleaved/pdf/nemotron_parse/test_server.py @@ -0,0 +1,59 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for Nemotron-Parse inference-server configuration.""" + +import pytest + +from nemo_curator.core.serve import DynamoServerConfig, DynamoVLLMModelConfig, RayServeModelConfig +from nemo_curator.stages.interleaved.pdf.nemotron_parse.server import create_nemotron_parse_inference_server + + +def test_dynamo_server_has_pdf_defaults_and_overrides() -> None: + server = create_nemotron_parse_inference_server( + model_path="/models/nemotron-parse", + model_name="nemotron-parse", + num_replicas=2, + engine_kwargs={"enforce_eager": True}, + request_timeout_s=123, + ) + + model = server.models[0] + assert isinstance(model, DynamoVLLMModelConfig) + assert model.model_identifier == "/models/nemotron-parse" + assert model.model_name == "nemotron-parse" + assert model.num_replicas == 2 + assert model.engine_kwargs["limit_mm_per_prompt"] == {"image": 1} + assert model.engine_kwargs["enforce_eager"] is True + assert model.dynamo_kwargs == {"enable_multimodal": True} + assert model.runtime_env == {"uv": {"packages": ["albumentations==2.0.8"]}} + assert isinstance(server.backend, DynamoServerConfig) + assert server.backend.request_plane == "tcp" + assert server.backend.router.router_kwargs == {"trust_remote_code": True} + assert server.backend.subprocess_env == {"DYN_TCP_REQUEST_TIMEOUT": "123"} + + +def test_ray_serve_server_has_pdf_defaults() -> None: + server = create_nemotron_parse_inference_server(backend="ray-serve", num_replicas=3) + + model = server.models[0] + assert isinstance(model, RayServeModelConfig) + assert model.deployment_config == {"num_replicas": 3} + assert model.engine_kwargs["limit_mm_per_prompt"] == {"image": 1} + + +@pytest.mark.parametrize("num_replicas", [0, -1]) +def test_rejects_non_positive_replica_count(num_replicas: int) -> None: + with pytest.raises(ValueError, match="num_replicas must be at least 1"): + create_nemotron_parse_inference_server(num_replicas=num_replicas) diff --git a/tests/stages/interleaved/pdf/nemotron_parse/test_stages.py b/tests/stages/interleaved/pdf/nemotron_parse/test_stages.py index 36c1e285bf..925ea9c8b8 100644 --- a/tests/stages/interleaved/pdf/nemotron_parse/test_stages.py +++ b/tests/stages/interleaved/pdf/nemotron_parse/test_stages.py @@ -16,13 +16,15 @@ from __future__ import annotations +import asyncio import base64 +import contextlib import io import json import zipfile from types import SimpleNamespace from typing import TYPE_CHECKING -from unittest.mock import patch +from unittest.mock import AsyncMock, patch import pytest from PIL import Image @@ -496,6 +498,39 @@ def fake_create_vllm_llm(model_path: str, **kwargs) -> object: assert captured_kwargs["gpu_memory_utilization"] == 0.9 assert stage._proc_size == (100, 100) + def test_in_process_and_http_client_sampling_parameters_match(self) -> None: + from nemo_curator.stages.interleaved.pdf.nemotron_parse.inference import ( + _nemotron_parse_sampling_params, + _nemotron_parse_server_generation_config, + ) + + http_client_config = _nemotron_parse_server_generation_config(1234) + http_client_params = { + "temperature": http_client_config.temperature, + "top_p": http_client_config.top_p, + "max_tokens": http_client_config.max_tokens, + "seed": http_client_config.seed, + **http_client_config.extra_kwargs["extra_body"], + } + + assert http_client_params == _nemotron_parse_sampling_params(1234) + + def test_in_process_and_http_client_default_to_8192_max_tokens(self) -> None: + from nemo_curator.stages.interleaved.pdf.nemotron_parse.inference import ( + NemotronParseHTTPClientStage, + NemotronParseInferenceStage, + ) + + in_process_stage = NemotronParseInferenceStage() + http_client_stage = NemotronParseHTTPClientStage( + endpoint="http://localhost:8000/v1", + model_name="nemotron-parse", + ) + + assert in_process_stage.max_tokens == 8192 + assert http_client_stage.max_tokens == 8192 + assert http_client_stage._generation_config.max_tokens == 8192 + def test_infer_vllm_empty_outputs_produces_empty_string(self) -> None: """RequestOutput with no completions should yield '' rather than IndexError.""" from nemo_curator.stages.interleaved.pdf.nemotron_parse.inference import NemotronParseInferenceStage @@ -512,6 +547,17 @@ def test_infer_vllm_empty_outputs_produces_empty_string(self) -> None: assert raw == [empty_req_output] assert retries == 0 + def test_hf_page_failure_raises_after_batch_fallback(self) -> None: + from nemo_curator.stages.interleaved.pdf.nemotron_parse.inference import NemotronParseInferenceStage + + stage = NemotronParseInferenceStage(backend="hf") + + with ( + patch.object(stage, "_infer_batch_hf", side_effect=RuntimeError("inference failed")), + pytest.raises(RuntimeError, match="inference failed"), + ): + stage._infer_hf([Image.new("RGB", (10, 10))]) + def test_infer_vllm_unreachable_loop_path_raises(self) -> None: from nemo_curator.stages.interleaved.pdf.nemotron_parse.inference import NemotronParseInferenceStage @@ -522,3 +568,141 @@ def test_infer_vllm_unreachable_loop_path_raises(self) -> None: with patch("builtins.range", return_value=()), pytest.raises(RuntimeError, match="unreachable"): stage._infer_vllm([image]) + + +class TestNemotronParseHTTPClientStage: + def test_rejects_non_positive_inference_batch_size(self) -> None: + from nemo_curator.stages.interleaved.pdf.nemotron_parse.inference import ( + NemotronParseHTTPClientStage, + ) + + with pytest.raises(ValueError, match="inference_batch_size must be at least 1"): + NemotronParseHTTPClientStage( + endpoint="http://localhost:8000/v1", + model_name="nemotron-parse", + inference_batch_size=0, + ) + + def test_warns_when_request_concurrency_is_below_recommended_starting_point(self) -> None: + from nemo_curator.stages.interleaved.pdf.nemotron_parse.inference import ( + NemotronParseHTTPClientStage, + ) + + with patch("nemo_curator.stages.interleaved.pdf.nemotron_parse.inference.logger.warning") as warning: + NemotronParseHTTPClientStage( + endpoint="http://localhost:8000/v1", + model_name="nemotron-parse", + inference_batch_size=4, + ) + + warning.assert_called_once() + + def test_query_pages_runs_up_to_inference_batch_size_concurrently(self) -> None: + from nemo_curator.stages.interleaved.pdf.nemotron_parse.inference import ( + NemotronParseHTTPClientStage, + ) + + active_requests = 0 + max_active_requests = 0 + two_requests_started = asyncio.Event() + + async def create_response(**create_kwargs: object) -> SimpleNamespace: + nonlocal active_requests, max_active_requests + active_requests += 1 + max_active_requests = max(max_active_requests, active_requests) + if active_requests == 2: + two_requests_started.set() + with contextlib.suppress(TimeoutError): + await asyncio.wait_for(two_requests_started.wait(), timeout=0.05) + + messages = create_kwargs["messages"] + image_url = messages[0]["content"][1]["image_url"]["url"] # type: ignore[index] + page_text = base64.b64decode(image_url.split(",", 1)[1]).decode() + active_requests -= 1 + return SimpleNamespace( + choices=[SimpleNamespace(message=SimpleNamespace(content=page_text), finish_reason="stop")], + usage=SimpleNamespace(prompt_tokens=1, completion_tokens=1), + ) + + sdk_client = SimpleNamespace( + chat=SimpleNamespace(completions=SimpleNamespace(create=create_response)), + close=AsyncMock(), + ) + stage = NemotronParseHTTPClientStage( + endpoint="http://localhost:8000/v1", + model_name="nemotron-parse", + ) + stage.inference_batch_size = 2 + + with patch("nemo_curator.models.client.openai_client.AsyncOpenAI", return_value=sdk_client): + results = asyncio.run(stage._query_pages([(f"page-{index}".encode(), "image/png") for index in range(3)])) + + assert max_active_requests == 2 + assert [result.text for result in results] == ["page-0", "page-1", "page-2"] + + def test_terminal_request_failure_raises(self) -> None: + from nemo_curator.stages.interleaved.pdf.nemotron_parse.inference import ( + NemotronParseHTTPClientStage, + ) + + stage = NemotronParseHTTPClientStage( + endpoint="http://localhost:8000/v1", + model_name="nemotron-parse", + ) + client = SimpleNamespace(query_model_response=AsyncMock(side_effect=RuntimeError("request failed"))) + + with pytest.raises(RuntimeError, match="request failed"): + asyncio.run(stage._query_page(client, b"png-bytes", "image/png")) + + def test_process_uses_openai_client_response_and_records_usage(self) -> None: + import pandas as pd + import pyarrow as pa + + from nemo_curator.stages.interleaved.pdf.nemotron_parse.inference import ( + NemotronParseHTTPClientStage, + ) + from nemo_curator.tasks import InterleavedBatch + + task = InterleavedBatch( + dataset_name="test", + data=pa.Table.from_pandas( + pd.DataFrame( + [ + { + "sample_id": "s1", + "position": 0, + "modality": "page_image", + "content_type": "image/png", + "text_content": None, + "binary_content": b"png-bytes", + "source_ref": None, + } + ] + ) + ), + ) + response = SimpleNamespace( + choices=[SimpleNamespace(message=SimpleNamespace(content="parsed page"), finish_reason="stop")], + usage=SimpleNamespace(prompt_tokens=5, completion_tokens=7), + ) + stage = NemotronParseHTTPClientStage( + endpoint="http://localhost:8000/v1", + model_name="nemotron-parse", + model_path="/models/NVIDIA-Nemotron-Parse-v1.1", + ) + + with patch( + "nemo_curator.stages.interleaved.pdf.nemotron_parse.inference.AsyncOpenAIClient.query_model_response", + new_callable=AsyncMock, + return_value=response, + ): + result = stage.process(task) + + assert result is not None + assert result.to_pandas().iloc[0]["text_content"] == "parsed page" + assert result._metadata["inference_server_endpoint"] == "http://localhost:8000/v1" + assert result._metadata["model_path"] == "/models/NVIDIA-Nemotron-Parse-v1.1" + assert result._metadata["proc_size"] == [2048, 1664] + assert stage._custom_metrics["total_prompt_tokens"] == 5.0 + assert stage._custom_metrics["total_output_tokens"] == 7.0 + assert "vllm_inference_time" not in stage._custom_metrics diff --git a/tutorials/interleaved/README.md b/tutorials/interleaved/README.md index 3c0f1feea6..9bb1190d7d 100644 --- a/tutorials/interleaved/README.md +++ b/tutorials/interleaved/README.md @@ -7,7 +7,7 @@ Hands-on tutorials for curating **interleaved multimodal data** — documents th | Tutorial | Description | Files | |----------|-------------|-------| | **[Getting Started](getting-started/)** | Load, explore, filter, and save interleaved data from MINT-1T PDF shards | `interleaved_data_quickstart.ipynb`, `interleaved_pipeline.py` | -| **[PDF Extraction Pipeline (Nemotron-Parse)](nemotron_parse_pdf/)** | Convert PDFs into structured interleaved Parquet using Nemotron-Parse v1.2 | `main.py` | +| **[PDF Extraction Pipeline (Nemotron-Parse)](nemotron_parse_pdf/)** | Convert PDFs into structured interleaved Parquet using the recommended Dynamo inference-server path | `main.py`, `inprocess.py` | ## Quick Start diff --git a/tutorials/interleaved/nemotron_parse_pdf/README.md b/tutorials/interleaved/nemotron_parse_pdf/README.md index 033397de57..0ce029001f 100644 --- a/tutorials/interleaved/nemotron_parse_pdf/README.md +++ b/tutorials/interleaved/nemotron_parse_pdf/README.md @@ -2,16 +2,37 @@ Convert PDFs into structured, interleaved parquet — text blocks, tables, images, and captions in reading order — using **Nemotron-Parse v1.2**. +We recommend using NeMo Curator's Dynamo-backed +`InferenceServer` with the HTTP client stage instead of loading vLLM inside the +pipeline stage. Internal comparisons found this to be the better default +because the serving layer can batch requests across pipeline tasks and keep +model replicas fed while PDF rendering and postprocessing scale independently. + +We tested the tutorial on 8 H100 GPUs with request concurrency values of 32 and +64. Use this starting configuration: + +- One inference-server replica per GPU. +- A fixed HTTP stage pool of `4 * num_gpus` workers. +- `inference_batch_size=32`, which is the maximum number of concurrent page + requests sent by each HTTP client worker. Because the best value depends on + the GPU and corpus, benchmark 64 on the target workload and keep it only if it + improves throughput without request failures or out-of-memory errors. + ## Setup ```bash git clone https://github.com/NVIDIA-NeMo/Curator.git cd Curator pip install uv -uv sync --extra interleaved_cuda12 +uv sync --extra interleaved_cuda12 --extra inference_server ``` -## Quickstart +The NeMo Curator container includes the `etcd` and `nats-server` binaries that +Dynamo starts. For a source environment outside the container, install them +with [`docker/common/install_etcd_nats.sh`](https://github.com/NVIDIA-NeMo/Curator/blob/main/docker/common/install_etcd_nats.sh) +before running the recommended entry point. + +## Run the tutorial **Step 1 — Create a manifest listing your PDFs:** @@ -22,59 +43,45 @@ for f in /path/to/pdfs/*.pdf; do done ``` -**Step 2 — Run the pipeline:** +**Step 2 — Start Dynamo and run the pipeline (recommended):** ```bash -python tutorials/interleaved/nemotron_parse_pdf/main.py \ +uv run python tutorials/interleaved/nemotron_parse_pdf/main.py \ --manifest manifest.jsonl \ --pdf-dir /path/to/pdfs \ --output-dir /path/to/output \ - --backend vllm \ - --enforce-eager + --model-path nvidia/NVIDIA-Nemotron-Parse-v1.2 \ + --inference-batch-size 32 ``` -## Dynamo serving example +`main.py` detects the Ray-visible GPUs, starts one Dynamo model replica per GPU, +waits for the OpenAI-compatible endpoint to become healthy, runs the pipeline, +and stops Dynamo. No separately managed server process is required. It fixes +the HTTP stage pool at `4 * num_gpus` workers. Use `CUDA_VISIBLE_DEVICES` or +your Ray cluster resources to control which GPUs are used. -When serving Nemotron-Parse through `InferenceServer` with the Dynamo backend, -set the vLLM chat processor explicitly so multimodal OpenAI content arrays are -flattened correctly, enable multimodal handling on the worker, and pass any -runtime-specific Dynamo environment variables through `subprocess_env`: +`main.py` always starts Dynamo with vLLM and does not expose a `--backend` +option. -```python -from nemo_curator.core.serve import ( - DynamoRouterConfig, - DynamoServerConfig, - DynamoVLLMModelConfig, - InferenceServer, -) - -server = InferenceServer( - models=[ - DynamoVLLMModelConfig( - model_identifier="/path/to/NVIDIA-Nemotron-Parse-v1.2", - engine_kwargs={ - "trust_remote_code": True, - "dtype": "bfloat16", - "limit_mm_per_prompt": {"image": 1}, - "enable_prefix_caching": False, - "disable_hybrid_kv_cache_manager": False, - }, - dynamo_kwargs={"enable_multimodal": True}, - ) - ], - backend=DynamoServerConfig( - request_plane="tcp", - router=DynamoRouterConfig( - router_kwargs={ - "dyn_chat_processor": "vllm", - } - ), - subprocess_env={"DYN_TCP_REQUEST_TIMEOUT": "180"}, - ), -) -server.start() +**Alternative — Run inference in process:** + +```bash +uv run python tutorials/interleaved/nemotron_parse_pdf/inprocess.py \ + --manifest manifest.jsonl \ + --pdf-dir /path/to/pdfs \ + --output-dir /path/to/output \ + --backend vllm \ + --enforce-eager ``` +Use `inprocess.py` for local validation and debugging when you do not want a +separate serving topology. + +The entry point uses `create_nemotron_parse_inference_server`, which keeps the +Nemotron-Parse vLLM, Dynamo, and runtime-environment settings shared with the +benchmark. See the [Inference Server guide](https://docs.nvidia.com/nemo/curator/latest/curate-text/synthetic/inference-server) +for details about the underlying configuration objects. + ## Input formats The pipeline supports three input formats selected by a mutually exclusive flag: @@ -121,9 +128,10 @@ images = [Image.open(io.BytesIO(b)) for b in df[df["modality"] == "image"]["bina | Flag | Default | Description | |------|---------|-------------| -| `--backend` | `vllm` | Inference backend (`vllm` or `hf`) | +| `--backend` | `vllm` | In-process engine (`inprocess.py` only); also supports `hf`. | | `--enforce-eager` | off | Skip vLLM CUDA graph capture (~35 min savings on first run) | | `--max-num-seqs` | 64 | Max concurrent sequences for vLLM | +| `--inference-batch-size` | 32 (`main.py`), 4 (`inprocess.py`) | Concurrent requests per HTTP worker, or pages per in-process HF pass | | `--pdfs-per-task` | 10 | PDFs batched per processing task | | `--max-pdfs` | — | Cap total PDFs (for testing) | | `--dpi` | 300 | PDF rendering resolution | diff --git a/tutorials/interleaved/nemotron_parse_pdf/inprocess.py b/tutorials/interleaved/nemotron_parse_pdf/inprocess.py new file mode 100644 index 0000000000..c0b0030d77 --- /dev/null +++ b/tutorials/interleaved/nemotron_parse_pdf/inprocess.py @@ -0,0 +1,215 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tutorial: Process PDFs through Nemotron-Parse with in-process inference. + +This pipeline reads PDFs (from a directory or CC-MAIN-style zip archives), +renders each page to an image, runs Nemotron-Parse for structured extraction +(text, tables, images), and writes interleaved parquet output. + +Pipeline stages:: + + 1. PDFPartitioningStage (EmptyTask -> FileGroupTask) [CPU] + Reads a JSONL manifest and packs PDF entries into FileGroupTasks. + + 2. PDFPreprocessStage (FileGroupTask -> InterleavedBatch) [CPU] + Extracts PDF bytes (from directory or zip), renders pages to images. + + 3. NemotronParseInferenceStage or NemotronParseHTTPClientStage + (InterleavedBatch -> InterleavedBatch) + Runs Nemotron-Parse in process or calls an OpenAI-compatible inference + server. The inference-server path is recommended for production. + + 4. NemotronParsePostprocessStage (InterleavedBatch -> InterleavedBatch) [CPU] + Parses model output, aligns images/captions, crops, builds rows. + + 5. InterleavedParquetWriterStage (InterleavedBatch -> FileGroupTask) + Writes final interleaved parquet output. + +Supported data sources: + +- **PDF directory**: Set ``--pdf-dir`` to a directory containing ``.pdf`` files. + Create a simple manifest with:: + + for f in /path/to/pdfs/*.pdf; do + echo "{\"file_name\": \"$(basename $f)\"}" >> manifest.jsonl + done + +- **CC-MAIN zip archives**: Set ``--zip-base-dir`` to the root of the + CC-MAIN-2021-31-PDF-UNTRUNCATED zip hierarchy. The manifest should use + ``cc_pdf_file_names`` (list) or ``file_name`` fields. + See: https://github.com/tballison/CC-MAIN-2021-31-PDF-UNTRUNCATED + +Usage:: + + # From a PDF directory (3 PDFs for testing) + python inprocess.py --pdf-dir /path/to/pdfs --manifest manifest.jsonl \\ + --output-dir ./output --max-pdfs 3 + + # From CC-MAIN zip archives + python inprocess.py --zip-base-dir /path/to/zipfiles --manifest manifest.jsonl \\ + --output-dir ./output + + # Small in-process vLLM run + python inprocess.py --pdf-dir /path/to/pdfs --manifest manifest.jsonl \\ + --output-dir ./output --backend vllm + +For production, run ``main.py``. It starts a Dynamo ``InferenceServer`` and +calls ``create_nemotron_parse_pdf_pipeline`` with four HTTP stage workers per +inference GPU and ``--inference-batch-size`` set to 32 by default. +""" + +from __future__ import annotations + +import json +import os +import time +from dataclasses import dataclass + +from loguru import logger +from pipeline_utils import create_nemotron_parse_pdf_argparser, create_nemotron_parse_pdf_pipeline + +from nemo_curator.backends.xenna import XennaExecutor +from nemo_curator.core.client import RayClient +from nemo_curator.stages.base import ProcessingStage +from nemo_curator.tasks import FileGroupTask + + +@dataclass +class PerfLoggingStage(ProcessingStage[FileGroupTask, FileGroupTask]): + """Append per-task stage perf stats to a JSONL file as each task completes. + + Placed after the writer stage so perf data is flushed to disk + incrementally — survives job kills from Slurm time limits. + """ + + output_dir: str + name: str = "perf_logging" + + def inputs(self) -> tuple[list[str], list[str]]: + return ["data"], [] + + def outputs(self) -> tuple[list[str], list[str]]: + return ["data"], [] + + def process(self, task: FileGroupTask) -> FileGroupTask: + perf_path = os.path.join(self.output_dir, f"_perf_stats_{os.getpid()}.jsonl") + record = { + "task_id": task.task_id, + "stages": [ + { + "stage_name": p.stage_name, + "process_time_s": p.process_time, + "actor_idle_time_s": p.actor_idle_time, + "num_items_processed": p.num_items_processed, + **{f"custom_{k}": v for k, v in p.custom_metrics.items()}, + } + for p in task._stage_perf + ], + } + with open(perf_path, "a") as f: + f.write(json.dumps(record) + "\n") + return task + + +def _write_perf_summary(results: list, output_dir: str, wall_time: float) -> None: + """Write per-task stage timings to a parquet file and log aggregate stats.""" + import pandas as pd + + valid_results = [r for r in results if r is not None] if results else [] + if not valid_results: + logger.warning("No results to write perf summary for") + return + + if len(valid_results) < len(results): + logger.warning(f"{len(results) - len(valid_results)} tasks returned None (failed)") + + rows = [] + for task in valid_results: + for perf in task._stage_perf: + row = { + "task_id": task.task_id, + "stage_name": perf.stage_name, + "process_time_s": perf.process_time, + "actor_idle_time_s": perf.actor_idle_time, + "num_items_processed": perf.num_items_processed, + } + for k, v in perf.custom_metrics.items(): + row[f"custom_{k}"] = v + rows.append(row) + + df = pd.DataFrame(rows) + job_id = os.environ.get("SLURM_JOB_ID", f"local_{int(time.time())}") + perf_path = os.path.join(output_dir, f"_perf_stats_{job_id}.parquet") + df.to_parquet(perf_path, index=False) + logger.info(f"Wrote {len(df)} perf records ({len(valid_results)} tasks) to {perf_path}") + + n_tasks = len(valid_results) + logger.info(f"\n{'=' * 70}\n PERFORMANCE SUMMARY (wall_time={wall_time:.1f}s, tasks={n_tasks})\n{'=' * 70}") + for stage_name, group in df.groupby("stage_name", sort=False): + avg_t = group["process_time_s"].mean() + sum_t = group["process_time_s"].sum() + p50 = group["process_time_s"].median() + p95 = group["process_time_s"].quantile(0.95) + total_items = group["num_items_processed"].sum() + logger.info( + f" {stage_name:40s} avg={avg_t:8.2f}s p50={p50:8.2f}s p95={p95:8.2f}s " + f"sum={sum_t:10.1f}s items={total_items}" + ) + logger.info(f"{'=' * 70}\n") + + +def main() -> None: + parser = create_nemotron_parse_pdf_argparser() + parser.add_argument("--backend", default="vllm", choices=["hf", "vllm"], help="In-process inference backend") + args = parser.parse_args() + + args.output_dir = os.path.abspath(args.output_dir) + os.makedirs(args.output_dir, exist_ok=True) + + if os.environ.get("SLURM_JOB_ID"): + from nemo_curator.core.client import SlurmRayClient + + ray_client = SlurmRayClient() + else: + ray_client = RayClient() + ray_client.start() + + try: + pipeline = create_nemotron_parse_pdf_pipeline(args, inprocess_backend=args.backend) + logger.info(f"\n{pipeline.describe()}") + + executor = XennaExecutor( + config={ + "execution_mode": "streaming", + "ignore_failures": True, + "failures_return_nones": True, + "reset_workers_on_failure": True, + } + ) + + t0 = time.perf_counter() + results = pipeline.run(executor=executor) + wall_time = time.perf_counter() - t0 + + n_valid = sum(1 for r in results if r is not None) + n_failed = len(results) - n_valid + logger.info(f"Pipeline finished in {wall_time:.1f}s, {n_valid} output tasks ({n_failed} failed)") + _write_perf_summary(results, args.output_dir, wall_time) + finally: + ray_client.stop() + + +if __name__ == "__main__": + main() diff --git a/tutorials/interleaved/nemotron_parse_pdf/main.py b/tutorials/interleaved/nemotron_parse_pdf/main.py index 1e98b8f426..13c1a373eb 100644 --- a/tutorials/interleaved/nemotron_parse_pdf/main.py +++ b/tutorials/interleaved/nemotron_parse_pdf/main.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,301 +12,49 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Tutorial: Process PDFs through Nemotron-Parse into interleaved parquet. - -This pipeline reads PDFs (from a directory or CC-MAIN-style zip archives), -renders each page to an image, runs Nemotron-Parse for structured extraction -(text, tables, images), and writes interleaved parquet output. - -Pipeline stages:: - - 1. PDFPartitioningStage (EmptyTask -> FileGroupTask) [CPU] - Reads a JSONL manifest and packs PDF entries into FileGroupTasks. - - 2. PDFPreprocessStage (FileGroupTask -> InterleavedBatch) [CPU] - Extracts PDF bytes (from directory or zip), renders pages to images. - - 3. NemotronParseInferenceStage (InterleavedBatch -> InterleavedBatch) [GPU] - Runs Nemotron-Parse model inference on page images. - - 4. NemotronParsePostprocessStage (InterleavedBatch -> InterleavedBatch) [CPU] - Parses model output, aligns images/captions, crops, builds rows. - - 5. InterleavedParquetWriterStage (InterleavedBatch -> FileGroupTask) - Writes final interleaved parquet output. - -Supported data sources: - -- **PDF directory**: Set ``--pdf-dir`` to a directory containing ``.pdf`` files. - Create a simple manifest with:: - - for f in /path/to/pdfs/*.pdf; do - echo "{\"file_name\": \"$(basename $f)\"}" >> manifest.jsonl - done - -- **CC-MAIN zip archives**: Set ``--zip-base-dir`` to the root of the - CC-MAIN-2021-31-PDF-UNTRUNCATED zip hierarchy. The manifest should use - ``cc_pdf_file_names`` (list) or ``file_name`` fields. - See: https://github.com/tballison/CC-MAIN-2021-31-PDF-UNTRUNCATED - -Usage:: - - # From a PDF directory (3 PDFs for testing) - python main.py --pdf-dir /path/to/pdfs --manifest manifest.jsonl \\ - --output-dir ./output --max-pdfs 3 - - # From CC-MAIN zip archives - python main.py --zip-base-dir /path/to/zipfiles --manifest manifest.jsonl \\ - --output-dir ./output - - # With vLLM backend (recommended for throughput) - python main.py --pdf-dir /path/to/pdfs --manifest manifest.jsonl \\ - --output-dir ./output --backend vllm -""" +"""Run the Nemotron-Parse PDF tutorial with NVIDIA Dynamo serving.""" from __future__ import annotations -import argparse -import json import os -import time -from dataclasses import dataclass -from loguru import logger +from pipeline_utils import create_nemotron_parse_pdf_argparser, create_nemotron_parse_pdf_pipeline -from nemo_curator.backends.xenna import XennaExecutor -from nemo_curator.core.client import RayClient -from nemo_curator.pipeline import Pipeline -from nemo_curator.stages.base import ProcessingStage -from nemo_curator.stages.interleaved.io import InterleavedParquetWriterStage -from nemo_curator.stages.interleaved.pdf.nemotron_parse import NemotronParsePDFReader -from nemo_curator.tasks import FileGroupTask - - -@dataclass -class PerfLoggingStage(ProcessingStage[FileGroupTask, FileGroupTask]): - """Append per-task stage perf stats to a JSONL file as each task completes. - - Placed after the writer stage so perf data is flushed to disk - incrementally — survives job kills from Slurm time limits. - """ - - output_dir: str - name: str = "perf_logging" - - def inputs(self) -> tuple[list[str], list[str]]: - return ["data"], [] - - def outputs(self) -> tuple[list[str], list[str]]: - return ["data"], [] - - def process(self, task: FileGroupTask) -> FileGroupTask: - perf_path = os.path.join(self.output_dir, f"_perf_stats_{os.getpid()}.jsonl") - record = { - "task_id": task.task_id, - "stages": [ - { - "stage_name": p.stage_name, - "process_time_s": p.process_time, - "actor_idle_time_s": p.actor_idle_time, - "num_items_processed": p.num_items_processed, - **{f"custom_{k}": v for k, v in p.custom_metrics.items()}, - } - for p in task._stage_perf - ], - } - with open(perf_path, "a") as f: - f.write(json.dumps(record) + "\n") - return task - - -def create_nemotron_parse_pdf_argparser() -> argparse.ArgumentParser: - """Create the argument parser for the Nemotron-Parse PDF pipeline.""" - parser = argparse.ArgumentParser(description="Process PDFs through Nemotron-Parse into interleaved parquet") - - # Data source - parser.add_argument("--manifest", required=True, help="Path to JSONL manifest listing PDFs") - source = parser.add_mutually_exclusive_group(required=True) - source.add_argument("--pdf-dir", help="Directory containing PDF files") - source.add_argument("--zip-base-dir", help="Root of CC-MAIN zip archive hierarchy") - source.add_argument("--jsonl-base-dir", help="Root of JSONL-based PDF dataset (e.g. GitHub PDFs)") - - # Output - parser.add_argument("--output-dir", required=True, help="Output directory for parquet files") - parser.add_argument("--dataset-name", default="pdf_dataset", help="Dataset name for output tasks") - - # Model - parser.add_argument( - "--model-path", - default="nvidia/NVIDIA-Nemotron-Parse-v1.2", - help="HuggingFace model ID or local path", - ) - parser.add_argument("--backend", default="vllm", choices=["hf", "vllm"], help="Inference backend") - - # Processing - parser.add_argument("--pdfs-per-task", type=int, default=10, help="PDFs per processing task") - parser.add_argument("--max-pdfs", type=int, default=None, help="Limit total PDFs (for testing)") - parser.add_argument("--dpi", type=int, default=300, help="PDF rendering resolution") - parser.add_argument("--max-pages", type=int, default=50, help="Max pages per PDF") - parser.add_argument("--min-crop-size", type=int, default=10, help="Min pixel dimension for image crops") - parser.add_argument( - "--text-in-pic", - action="store_true", - help="Predict text inside pictures (v1.2+ only). Default: no text in pictures.", - ) - - # Inference - parser.add_argument("--inference-batch-size", type=int, default=4, help="Pages per GPU pass (HF only)") - parser.add_argument("--max-num-seqs", type=int, default=64, help="Max concurrent sequences (vLLM only)") - parser.add_argument( - "--enforce-eager", - action="store_true", - help="Disable vLLM CUDA graph capture (enforce_eager=True). Eliminates ~35min compilation " - "idle at startup; slight throughput reduction. Recommended on clusters with GPU " - "utilization enforcement.", - ) - - # Executor - parser.add_argument( - "--execution-mode", - default="streaming", - choices=["streaming", "batch"], - help="XennaExecutor execution mode", - ) - - # Manifest field names - parser.add_argument("--file-name-field", default="file_name", help="JSONL field for single PDF filename") - parser.add_argument( - "--file-names-field", default="cc_pdf_file_names", help="JSONL field for list of PDF filenames" - ) - parser.add_argument("--url-field", default="url", help="JSONL field for source URL") - - return parser - - -def create_nemotron_parse_pdf_pipeline(args: argparse.Namespace) -> Pipeline: - """Build the Nemotron-Parse PDF processing pipeline from parsed arguments.""" - pipeline = Pipeline( - name="nemotron_parse_pdf", - description="PDF -> Nemotron-Parse -> Interleaved Parquet", - ) - pipeline.add_stage( - NemotronParsePDFReader( - manifest_path=args.manifest, - zip_base_dir=args.zip_base_dir, - pdf_dir=args.pdf_dir, - jsonl_base_dir=args.jsonl_base_dir, - model_path=args.model_path, - backend=args.backend, - pdfs_per_task=args.pdfs_per_task, - max_pdfs=args.max_pdfs, - dpi=args.dpi, - max_pages=args.max_pages, - inference_batch_size=args.inference_batch_size, - max_num_seqs=args.max_num_seqs, - text_in_pic=args.text_in_pic, - enforce_eager=args.enforce_eager, - min_crop_px=args.min_crop_size, - dataset_name=args.dataset_name, - file_name_field=args.file_name_field, - file_names_field=args.file_names_field, - url_field=args.url_field, - ) - ) - pipeline.add_stage( - InterleavedParquetWriterStage( - path=args.output_dir, - materialize_on_write=False, - ) - ) - return pipeline - - -def _write_perf_summary(results: list, output_dir: str, wall_time: float) -> None: - """Write per-task stage timings to a parquet file and log aggregate stats.""" - import pandas as pd - - valid_results = [r for r in results if r is not None] if results else [] - if not valid_results: - logger.warning("No results to write perf summary for") - return - - if len(valid_results) < len(results): - logger.warning(f"{len(results) - len(valid_results)} tasks returned None (failed)") - - rows = [] - for task in valid_results: - for perf in task._stage_perf: - row = { - "task_id": task.task_id, - "stage_name": perf.stage_name, - "process_time_s": perf.process_time, - "actor_idle_time_s": perf.actor_idle_time, - "num_items_processed": perf.num_items_processed, - } - for k, v in perf.custom_metrics.items(): - row[f"custom_{k}"] = v - rows.append(row) - - df = pd.DataFrame(rows) - job_id = os.environ.get("SLURM_JOB_ID", f"local_{int(time.time())}") - perf_path = os.path.join(output_dir, f"_perf_stats_{job_id}.parquet") - df.to_parquet(perf_path, index=False) - logger.info(f"Wrote {len(df)} perf records ({len(valid_results)} tasks) to {perf_path}") - - n_tasks = len(valid_results) - logger.info(f"\n{'=' * 70}\n PERFORMANCE SUMMARY (wall_time={wall_time:.1f}s, tasks={n_tasks})\n{'=' * 70}") - for stage_name, group in df.groupby("stage_name", sort=False): - avg_t = group["process_time_s"].mean() - sum_t = group["process_time_s"].sum() - p50 = group["process_time_s"].median() - p95 = group["process_time_s"].quantile(0.95) - total_items = group["num_items_processed"].sum() - logger.info( - f" {stage_name:40s} avg={avg_t:8.2f}s p50={p50:8.2f}s p95={p95:8.2f}s " - f"sum={sum_t:10.1f}s items={total_items}" - ) - logger.info(f"{'=' * 70}\n") +from nemo_curator.backends.ray_data import RayDataExecutor +from nemo_curator.backends.utils import get_available_cpu_gpu_resources +from nemo_curator.stages.interleaved.pdf.nemotron_parse import create_nemotron_parse_inference_server def main() -> None: parser = create_nemotron_parse_pdf_argparser() + parser.set_defaults(inference_batch_size=32) args = parser.parse_args() args.output_dir = os.path.abspath(args.output_dir) os.makedirs(args.output_dir, exist_ok=True) - if os.environ.get("SLURM_JOB_ID"): - from nemo_curator.core.client import SlurmRayClient - - ray_client = SlurmRayClient() - else: - ray_client = RayClient() - ray_client.start() - - try: - pipeline = create_nemotron_parse_pdf_pipeline(args) - logger.info(f"\n{pipeline.describe()}") - - executor = XennaExecutor( - config={ - "execution_mode": args.execution_mode, - "ignore_failures": True, - "failures_return_nones": True, - "reset_workers_on_failure": True, - } + _, available_gpus = get_available_cpu_gpu_resources(init_and_shutdown=True) + num_gpus = int(available_gpus) + if num_gpus < 1: + parser.error("Dynamo serving requires at least one Ray-visible GPU") + + model_name = args.model_path + server = create_nemotron_parse_inference_server( + model_path=args.model_path, + model_name=model_name, + backend="dynamo", + num_replicas=num_gpus, + engine_kwargs={"enforce_eager": True} if args.enforce_eager else None, + ) + + with server: + pipeline = create_nemotron_parse_pdf_pipeline( + args, + inference_server_endpoint=server.endpoint, + inference_server_model_name=model_name, + inference_server_client_num_workers=4 * num_gpus, ) - - t0 = time.perf_counter() - results = pipeline.run(executor=executor) - wall_time = time.perf_counter() - t0 - - n_valid = sum(1 for r in results if r is not None) - n_failed = len(results) - n_valid - logger.info(f"Pipeline finished in {wall_time:.1f}s, {n_valid} output tasks ({n_failed} failed)") - _write_perf_summary(results, args.output_dir, wall_time) - finally: - ray_client.stop() + pipeline.run(RayDataExecutor()) if __name__ == "__main__": diff --git a/tutorials/interleaved/nemotron_parse_pdf/pipeline_utils.py b/tutorials/interleaved/nemotron_parse_pdf/pipeline_utils.py new file mode 100644 index 0000000000..b5b42e0677 --- /dev/null +++ b/tutorials/interleaved/nemotron_parse_pdf/pipeline_utils.py @@ -0,0 +1,134 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared argument and pipeline builders for the Nemotron-Parse PDF tutorial.""" + +from __future__ import annotations + +import argparse + +from nemo_curator.pipeline import Pipeline +from nemo_curator.stages.interleaved.io import InterleavedParquetWriterStage +from nemo_curator.stages.interleaved.pdf.nemotron_parse import NemotronParsePDFReader +from nemo_curator.stages.interleaved.pdf.nemotron_parse.inference import DEFAULT_MAX_TOKENS + + +def create_nemotron_parse_pdf_argparser() -> argparse.ArgumentParser: + """Create the argument parser for the Nemotron-Parse PDF pipeline.""" + parser = argparse.ArgumentParser(description="Process PDFs through Nemotron-Parse into interleaved parquet") + + parser.add_argument("--manifest", required=True, help="Path to JSONL manifest listing PDFs") + source = parser.add_mutually_exclusive_group(required=True) + source.add_argument("--pdf-dir", help="Directory containing PDF files") + source.add_argument("--zip-base-dir", help="Root of CC-MAIN zip archive hierarchy") + source.add_argument("--jsonl-base-dir", help="Root of JSONL-based PDF dataset (e.g. GitHub PDFs)") + + parser.add_argument("--output-dir", required=True, help="Output directory for parquet files") + parser.add_argument("--dataset-name", default="pdf_dataset", help="Dataset name for output tasks") + + parser.add_argument( + "--model-path", + default="nvidia/NVIDIA-Nemotron-Parse-v1.2", + help="HuggingFace model ID or local path", + ) + parser.add_argument("--pdfs-per-task", type=int, default=10, help="PDFs per processing task") + parser.add_argument("--max-pdfs", type=int, default=None, help="Limit total PDFs (for testing)") + parser.add_argument("--dpi", type=int, default=300, help="PDF rendering resolution") + parser.add_argument("--max-pages", type=int, default=50, help="Max pages per PDF") + parser.add_argument("--min-crop-size", type=int, default=10, help="Min pixel dimension for image crops") + parser.add_argument( + "--text-in-pic", + action="store_true", + help="Predict text inside pictures (v1.2+ only). Default: no text in pictures.", + ) + + parser.add_argument( + "--inference-batch-size", + type=int, + default=4, + help="Pages per HF GPU pass or maximum concurrent inference-server page requests", + ) + parser.add_argument("--max-num-seqs", type=int, default=64, help="Max concurrent sequences (vLLM only)") + parser.add_argument("--max-tokens", type=int, default=DEFAULT_MAX_TOKENS, help="Maximum output tokens per page") + parser.add_argument( + "--enforce-eager", + action="store_true", + help="Disable vLLM CUDA graph capture (enforce_eager=True). Eliminates ~35min compilation " + "idle at startup; slight throughput reduction. Recommended on clusters with GPU " + "utilization enforcement.", + ) + + parser.add_argument("--file-name-field", default="file_name", help="JSONL field for single PDF filename") + parser.add_argument( + "--file-names-field", default="cc_pdf_file_names", help="JSONL field for list of PDF filenames" + ) + parser.add_argument("--url-field", default="url", help="JSONL field for source URL") + + return parser + + +def create_nemotron_parse_pdf_pipeline( + args: argparse.Namespace, + *, + inprocess_backend: str = "vllm", + inference_server_endpoint: str | None = None, + inference_server_model_name: str | None = None, + inference_server_client_num_workers: int = 4, +) -> Pipeline: + """Build the PDF pipeline, optionally using an inference server. + + For an inference-server deployment, use four HTTP client workers per + serving GPU and start with ``args.inference_batch_size=32``. Tune request + concurrency on the target hardware and corpus. + """ + pipeline = Pipeline( + name="nemotron_parse_pdf", + description="PDF -> Nemotron-Parse -> Interleaved Parquet", + ) + pipeline.add_stage( + NemotronParsePDFReader( + manifest_path=args.manifest, + zip_base_dir=args.zip_base_dir, + pdf_dir=args.pdf_dir, + jsonl_base_dir=args.jsonl_base_dir, + model_path=args.model_path, + backend=inprocess_backend, + pdfs_per_task=args.pdfs_per_task, + max_pdfs=args.max_pdfs, + dpi=args.dpi, + max_pages=args.max_pages, + inference_batch_size=args.inference_batch_size, + max_num_seqs=args.max_num_seqs, + max_tokens=args.max_tokens, + text_in_pic=args.text_in_pic, + enforce_eager=args.enforce_eager, + min_crop_px=args.min_crop_size, + dataset_name=args.dataset_name, + file_name_field=args.file_name_field, + file_names_field=args.file_names_field, + url_field=args.url_field, + inference_server_endpoint=inference_server_endpoint, + inference_server_model_name=inference_server_model_name, + inference_server_client_num_workers=inference_server_client_num_workers, + inference_server_request_timeout_s=getattr(args, "inference_server_request_timeout_s", 300.0), + inference_server_max_retries=getattr(args, "inference_server_max_retries", 3), + ) + ) + pipeline.add_stage( + InterleavedParquetWriterStage( + path=args.output_dir, + materialize_on_write=False, + ) + ) + return pipeline