# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: Apache-2.0
# pylint: skip-file
import asyncio
from io import BytesIO
from typing import Any, Dict, List, Optional, Tuple
import base64
import json
import logging
import os
import time
import uuid
import random
from pathlib import Path
import fsspec
from fastapi import APIRouter, Request, Response
from fastapi import HTTPException
from fastapi.responses import StreamingResponse
from redis import RedisError
from nv_ingest.framework.schemas.framework_message_wrapper_schema import MessageWrapper
from nv_ingest_api.util.service_clients.client_base import FetchMode
from nv_ingest_api.util.dataloader.dataloader import DataLoader
from nv_ingest_api.internal.schemas.meta.ingest_job_schema import DocumentTypeEnum
# For PDF splitting
import pypdfium2 as pdfium
# Reuse V1 state constants and helper functions
from ..v1.ingest import (
trace_id_to_uuid,
INGEST_SERVICE_T,
STATE_RETRIEVED_DESTRUCTIVE,
STATE_RETRIEVED_NON_DESTRUCTIVE,
STATE_RETRIEVED_CACHED,
STATE_FAILED,
STATE_SUBMITTED,
INTERMEDIATE_STATES,
)
from .. import traced_endpoint
from opentelemetry import trace
logger = logging.getLogger("uvicorn")
router = APIRouter()
DEFAULT_PDF_SPLIT_PAGE_COUNT = 32
# Default QoS thresholds (pages). Tunable via environment variables:
# QOS_MAX_PAGES_MICRO, QOS_MAX_PAGES_SMALL, QOS_MAX_PAGES_MEDIUM
_QOS_DEFAULTS = {
"micro": 8,
"small": 64,
"medium": 256,
}
[docs]
def get_qos_tier_for_page_count(page_count: int) -> str:
"""
Select QoS tier for a document based on its total page count.
Tiers: 'micro', 'small', 'medium', 'large', 'default'
Thresholds can be tuned via environment variables:
- QOS_MAX_PAGES_MICRO (default: 4)
- QOS_MAX_PAGES_SMALL (default: 16)
- QOS_MAX_PAGES_MEDIUM (default: 64)
Anything above MEDIUM is 'large'. Non-positive page_count returns 'default'.
"""
try:
micro_max = int(os.getenv("QOS_MAX_PAGES_MICRO", str(_QOS_DEFAULTS["micro"])))
small_max = int(os.getenv("QOS_MAX_PAGES_SMALL", str(_QOS_DEFAULTS["small"])))
medium_max = int(os.getenv("QOS_MAX_PAGES_MEDIUM", str(_QOS_DEFAULTS["medium"])))
except ValueError:
micro_max, small_max, medium_max = _QOS_DEFAULTS["micro"], _QOS_DEFAULTS["small"], _QOS_DEFAULTS["medium"]
if page_count <= 0:
return "default"
if page_count <= micro_max:
return "micro"
if page_count <= small_max:
return "small"
if page_count <= medium_max:
return "medium"
return "large"
[docs]
def get_pdf_split_page_count(client_override: Optional[int] = None) -> int:
"""
Resolve the page chunk size for PDF splitting with client override support.
Priority: client_override (clamped) > env var > default (32)
Enforces boundaries: min=1, max=128
"""
MIN_PAGES = 1
MAX_PAGES = 128
# Client override takes precedence if provided
if client_override is not None:
clamped = max(MIN_PAGES, min(client_override, MAX_PAGES))
if clamped != client_override:
logger.warning(
"Client requested split_page_count=%s; clamped to %s (min=%s, max=%s)",
client_override,
clamped,
MIN_PAGES,
MAX_PAGES,
)
return clamped
# Fall back to environment variable
raw_value = os.environ.get("PDF_SPLIT_PAGE_COUNT")
if raw_value is None:
return DEFAULT_PDF_SPLIT_PAGE_COUNT
try:
parsed = int(raw_value)
except ValueError:
logger.warning(
"Invalid PDF_SPLIT_PAGE_COUNT '%s'; falling back to default %s", raw_value, DEFAULT_PDF_SPLIT_PAGE_COUNT
)
return DEFAULT_PDF_SPLIT_PAGE_COUNT
clamped = max(MIN_PAGES, min(parsed, MAX_PAGES))
if clamped != parsed:
logger.warning(
"Env PDF_SPLIT_PAGE_COUNT=%s clamped to %s (min=%s, max=%s)",
parsed,
clamped,
MIN_PAGES,
MAX_PAGES,
)
return clamped
[docs]
def split_pdf_to_chunks(pdf_content: bytes, pages_per_chunk: int) -> List[Dict[str, Any]]:
"""
Split a PDF into multi-page chunks using pypdfium2.
Returns a list of dictionaries containing the chunk bytes and page range metadata.
Note: this currently buffers each chunk in-memory; consider streaming in future upgrades.
"""
chunks: List[Dict[str, Any]] = []
if pages_per_chunk <= 0:
pages_per_chunk = 1
pdf = pdfium.PdfDocument(pdf_content)
total_pages = len(pdf)
try:
for chunk_index, start_zero in enumerate(range(0, total_pages, pages_per_chunk)):
end_zero = min(start_zero + pages_per_chunk, total_pages)
page_indices = list(range(start_zero, end_zero))
new_pdf = pdfium.PdfDocument.new()
try:
new_pdf.import_pages(pdf, page_indices)
buffer = BytesIO()
try:
new_pdf.save(buffer)
chunk_bytes = buffer.getvalue()
finally:
buffer.close()
finally:
new_pdf.close()
start_page = start_zero + 1
end_page = end_zero
chunk_info: Dict[str, Any] = {
"bytes": chunk_bytes,
"chunk_index": chunk_index,
"start_page": start_page,
"end_page": end_page,
"page_count": end_page - start_page + 1,
}
chunks.append(chunk_info)
finally:
pdf.close()
return chunks
[docs]
def get_pdf_page_count(pdf_content: bytes) -> int:
"""Get the number of pages in a PDF using pypdfium2."""
try:
pdf = pdfium.PdfDocument(pdf_content)
page_count = len(pdf)
pdf.close()
return page_count
except Exception as e:
logger.warning(f"Failed to get PDF page count: {e}")
return 1 # Assume single page on error
def _create_subjob_dict(
job_id: str,
job_payload: Dict[str, Any],
job_spec_template: Dict[str, Any],
current_trace_id: int,
parent_job_id: str,
start_key: Dict[str, Any],
) -> Dict[str, Any]:
job_spec = {
key: value
for key, value in job_spec_template.items()
if key not in {"job_payload", "job_id", "tracing_options"}
}
job_spec["job_payload"] = job_payload
job_spec["job_id"] = job_id
base_tracing_options = job_spec_template.get("tracing_options") or {}
tracing_options = dict(base_tracing_options)
tracing_options.setdefault("trace", True)
tracing_options["trace_id"] = str(current_trace_id)
tracing_options["ts_send"] = int(time.time() * 1000)
tracing_options["parent_job_id"] = parent_job_id
for key, value in start_key.items():
tracing_options[key] = value
job_spec["tracing_options"] = tracing_options
return job_spec
def _create_payload_dict(
job_spec_template: Dict[str, Any],
content: str,
source_id: str,
source_name: str,
document_type: str,
) -> Dict[str, Any]:
subjob_payload_template = job_spec_template.get("job_payload", {})
subjob_payload = {
key: value
for key, value in subjob_payload_template.items()
if key not in {"content", "source_id", "source_name"}
}
subjob_payload["content"] = [content]
subjob_payload["source_id"] = [source_id]
subjob_payload["source_name"] = [source_name]
subjob_payload["document_type"] = [document_type]
return subjob_payload
def _prepare_chunk_submission(
job_spec_template: Dict[str, Any],
chunk: Dict[str, Any],
*,
parent_uuid: uuid.UUID,
parent_job_id: str,
current_trace_id: int,
source_id: str,
source_name: str,
document_type: str,
) -> Tuple[str, MessageWrapper]:
"""Create a subjob MessageWrapper for a PDF chunk and return its identifier."""
chunk_number = chunk["chunk_index"] + 1
subjob_uuid = uuid.uuid5(parent_uuid, f"chunk-{chunk_number}")
subjob_id = str(subjob_uuid)
subjob_payload_template = job_spec_template.get("job_payload", {})
chunk_bytes = base64.b64encode(chunk["bytes"]).decode("utf-8")
subjob_payload = _create_payload_dict(subjob_payload_template, chunk_bytes, source_id, source_name, document_type)
start = chunk["start_page"] if "start_page" in chunk else chunk["start"]
subjob_spec = _create_subjob_dict(
subjob_id, subjob_payload, job_spec_template, current_trace_id, parent_job_id, {"page_num": start}
)
return subjob_id, MessageWrapper(payload=json.dumps(subjob_spec))
# ============================================================================
# Helper Functions for Fetch Job Aggregation
# ============================================================================
async def _gather_in_batches(coroutines: List, batch_size: int, return_exceptions: bool = False) -> List[Any]:
"""
Execute coroutines in batches to respect concurrency limits.
Parameters
----------
coroutines : List
List of coroutines to execute
batch_size : int
Maximum number of coroutines to execute concurrently
return_exceptions : bool
Whether to return exceptions as results (passed to asyncio.gather)
Returns
-------
List[Any]
Results from all coroutines in original order
"""
results: List[Any] = []
for offset in range(0, len(coroutines), batch_size):
batch = coroutines[offset : offset + batch_size]
batch_results = await asyncio.gather(*batch, return_exceptions=return_exceptions)
results.extend(batch_results)
return results
async def _update_job_state_after_fetch(job_id: str, ingest_service: INGEST_SERVICE_T) -> None:
"""
Update job state after successful fetch based on configured fetch mode.
Parameters
----------
job_id : str
The job identifier
ingest_service : IngestServiceMeta
The ingest service instance
"""
try:
current_fetch_mode = await ingest_service.get_fetch_mode()
if current_fetch_mode == FetchMode.DESTRUCTIVE:
target_state = STATE_RETRIEVED_DESTRUCTIVE
elif current_fetch_mode == FetchMode.NON_DESTRUCTIVE:
target_state = STATE_RETRIEVED_NON_DESTRUCTIVE
else:
target_state = STATE_RETRIEVED_CACHED
await ingest_service.set_job_state(job_id, target_state)
logger.debug(f"Updated job {job_id} state to {target_state}")
except Exception as e:
logger.error(f"Failed to update job state for {job_id}: {e}")
def _stream_json_response(data: Dict[str, Any]) -> StreamingResponse:
"""
Create a StreamingResponse for JSON data.
Parameters
----------
data : Dict[str, Any]
The data to serialize and stream
Returns
-------
StreamingResponse
FastAPI streaming response with JSON content
"""
json_bytes = json.dumps(data).encode("utf-8")
return StreamingResponse(iter([json_bytes]), media_type="application/json", status_code=200)
async def _check_all_subjob_states(
ordered_descriptors: List[Dict[str, Any]], max_parallel_ops: int, ingest_service: INGEST_SERVICE_T
) -> Tuple[List[Optional[str]], List[Dict[str, object]]]:
"""
Check the state of all subjobs in parallel batches.
Parameters
----------
ordered_descriptors : List[Dict[str, Any]]
List of subjob descriptors with job_id and chunk_index
max_parallel_ops : int
Maximum number of parallel operations
ingest_service : IngestServiceMeta
The ingest service instance
Returns
-------
Tuple[List[Optional[str]], List[Dict[str, object]]]
Tuple of (subjob_states, failed_subjobs_list)
Raises
------
HTTPException
If any subjob is still processing (202)
"""
# Gather all subjob states in parallel batches
state_coroutines = [ingest_service.get_job_state(descriptor.get("job_id")) for descriptor in ordered_descriptors]
subjob_states = await _gather_in_batches(state_coroutines, max_parallel_ops)
# Check for failures and pending work
failed_subjobs: List[Dict[str, object]] = []
for page_index, (descriptor, subjob_state) in enumerate(zip(ordered_descriptors, subjob_states), start=1):
subjob_id = descriptor.get("job_id")
if subjob_state == STATE_FAILED:
logger.warning(f"Subjob {subjob_id} failed")
failed_subjobs.append({"subjob_id": subjob_id, "chunk_index": page_index})
elif subjob_state in INTERMEDIATE_STATES:
raise HTTPException(status_code=202, detail="Parent job still processing. Some pages not complete.")
return subjob_states, failed_subjobs
async def _fetch_all_subjob_results(
ordered_descriptors: List[Dict[str, Any]],
subjob_states: List[Optional[str]],
failed_subjobs: List[Dict[str, object]],
max_parallel_ops: int,
ingest_service: INGEST_SERVICE_T,
) -> List[Optional[Dict[str, Any]]]:
"""
Fetch results for all completed subjobs in parallel batches.
Parameters
----------
ordered_descriptors : List[Dict[str, Any]]
List of subjob descriptors
subjob_states : List[Optional[str]]
States of all subjobs (from _check_all_subjob_states)
failed_subjobs : List[Dict[str, object]]
List to append failed fetch attempts to (modified in place)
max_parallel_ops : int
Maximum number of parallel operations
ingest_service : IngestServiceMeta
The ingest service instance
Returns
-------
List[Optional[Dict[str, Any]]]
Results for each subjob (None for failed ones)
Raises
------
HTTPException
If any subjob is not ready yet (202)
"""
# Initialize results array with None placeholders
subjob_results: List[Optional[Dict[str, Any]]] = [None] * len(ordered_descriptors)
# Build list of fetch tasks (only for non-failed subjobs)
fetch_coroutines = []
fetch_targets: List[Dict[str, Any]] = []
for list_index, (page_index, descriptor, subjob_state) in enumerate(
zip(range(1, len(ordered_descriptors) + 1), ordered_descriptors, subjob_states)
):
subjob_id = descriptor.get("job_id")
# Skip failed subjobs (already recorded in failed_subjobs)
if subjob_state == STATE_FAILED:
continue
# Skip intermediate states (should have been caught earlier, but defensive)
if subjob_state in INTERMEDIATE_STATES:
continue
# Queue this subjob for fetching
fetch_coroutines.append(ingest_service.fetch_job(subjob_id))
fetch_targets.append(
{
"list_index": list_index,
"page_index": page_index,
"subjob_id": subjob_id,
}
)
# Fetch all results in parallel batches
if fetch_coroutines:
fetch_results = await _gather_in_batches(fetch_coroutines, max_parallel_ops, return_exceptions=True)
# Process results and handle errors
for target, fetch_result in zip(fetch_targets, fetch_results):
subjob_id = target["subjob_id"]
page_index = target["page_index"]
list_index = target["list_index"]
if isinstance(fetch_result, TimeoutError):
logger.debug(f"Subjob {subjob_id} not ready yet; deferring aggregation")
raise HTTPException(status_code=202, detail="Parent job still processing. Some pages not complete.")
if isinstance(fetch_result, Exception):
logger.error(f"Failed to fetch subjob {subjob_id}: {fetch_result}")
failed_subjobs.append(
{
"subjob_id": subjob_id,
"chunk_index": page_index,
"error": str(fetch_result),
}
)
continue
subjob_results[list_index] = fetch_result
return subjob_results
def _extract_ray_telemetry(result: Dict[str, Any]) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]:
"""Return the trace and annotation dictionaries emitted by the sink stage."""
if not isinstance(result, dict):
return None, None
trace = result.get("trace")
annotations = result.get("annotations")
trace_dict = trace if isinstance(trace, dict) else None
annotations_dict = annotations if isinstance(annotations, dict) else None
return trace_dict, annotations_dict
def _normalize_chunk_records(
records: Optional[List[Any]],
descriptor: Dict[str, Any],
parent_metadata: Dict[str, Any],
) -> List[Any]:
"""Re-map chunk-local metadata to document-level context for aggregation."""
if not isinstance(records, list):
return []
total_pages = parent_metadata.get("total_pages")
original_source_id = parent_metadata.get("original_source_id")
original_source_name = parent_metadata.get("original_source_name")
start_page = descriptor.get("start_page")
page_offset = start_page - 1 if isinstance(start_page, int) and start_page > 0 else 0
normalized_entries: List[Any] = []
for entry in records:
if not isinstance(entry, dict):
normalized_entries.append(entry)
continue
normalized_entry = entry.copy()
original_metadata = entry.get("metadata")
if isinstance(original_metadata, dict):
normalized_metadata = original_metadata.copy()
normalized_entry["metadata"] = normalized_metadata
original_source_meta = original_metadata.get("source_metadata")
if isinstance(original_source_meta, dict):
normalized_source_meta = original_source_meta.copy()
normalized_metadata["source_metadata"] = normalized_source_meta
if original_source_id:
normalized_source_meta["source_id"] = original_source_id
if original_source_name:
normalized_source_meta["source_name"] = original_source_name
original_content_meta = original_metadata.get("content_metadata")
if isinstance(original_content_meta, dict):
normalized_content_meta = original_content_meta.copy()
normalized_metadata["content_metadata"] = normalized_content_meta
page_number = normalized_content_meta.get("page_number")
if isinstance(page_number, int) and page_number >= 0:
normalized_content_meta["page_number"] = page_number + page_offset
if isinstance(total_pages, int) and isinstance(normalized_content_meta.get("page_count"), int):
# Ensure optional per-record page count reflects the full document
normalized_content_meta["page_count"] = total_pages
original_hierarchy = original_content_meta.get("hierarchy")
if isinstance(original_hierarchy, dict):
normalized_hierarchy = original_hierarchy.copy()
normalized_content_meta["hierarchy"] = normalized_hierarchy
hierarchy_page = normalized_hierarchy.get("page")
if isinstance(hierarchy_page, int) and hierarchy_page >= 0:
normalized_hierarchy["page"] = hierarchy_page + page_offset
if isinstance(total_pages, int):
normalized_hierarchy["page_count"] = total_pages
normalized_entries.append(normalized_entry)
return normalized_entries
def _aggregate_parent_traces(chunk_traces: Dict[str, Any]) -> Dict[str, Any]:
"""
Aggregate chunk-level traces into parent-level metrics.
For each stage found in chunk traces:
- trace::entry::<stage> = min(all chunk entries) - earliest start
- trace::exit::<stage> = max(all chunk exits) - latest finish
- trace::resident_time::<stage> = sum(chunk durations) - total compute
Parameters
----------
chunk_traces : Dict[str, Any]
Trace dict with chunk-prefixed keys (chunk_N::trace::entry::stage_name)
Returns
-------
Dict[str, Any]
Parent-level aggregated traces (trace::entry::stage_name, etc.)
"""
# Group by stage: {stage_name: {chunk_idx: {entry: float, exit: float}}}
stage_data: Dict[str, Dict[int, Dict[str, Any]]] = {}
for key, value in chunk_traces.items():
if not key.startswith("chunk_"):
continue
parts = key.split("::")
if len(parts) < 4: # Minimum: chunk_N::trace::entry/exit::stage_name
continue
if parts[1] != "trace": # Ensure it's a trace key
continue
chunk_idx_str = parts[0].split("_")[1] # "chunk_1" -> "1"
try:
chunk_idx = int(chunk_idx_str)
except ValueError:
continue
event_type = parts[2] # "entry" or "exit"
# Stage name is everything after trace::entry:: or trace::exit::
# Handles both simple (pdf_extractor) and nested (pdf_extractor::pdf_extraction::pdfium_0)
stage_name = "::".join(parts[3:]) # Join remaining parts
if event_type not in ("entry", "exit"):
continue
if stage_name not in stage_data:
stage_data[stage_name] = {}
if chunk_idx not in stage_data[stage_name]:
stage_data[stage_name][chunk_idx] = {}
stage_data[stage_name][chunk_idx][event_type] = value
# Compute aggregated metrics
parent_traces: Dict[str, Any] = {}
for stage_name, chunks in stage_data.items():
entries = []
exits = []
durations = []
for chunk_data in chunks.values():
entry = chunk_data.get("entry")
exit_time = chunk_data.get("exit")
# Both entry and exit must exist for valid pair
if entry is not None and exit_time is not None:
entries.append(entry)
exits.append(exit_time)
durations.append(exit_time - entry)
# Only add parent traces if we have valid data
if entries and exits:
parent_traces[f"trace::entry::{stage_name}"] = min(entries)
parent_traces[f"trace::exit::{stage_name}"] = max(exits)
parent_traces[f"trace::resident_time::{stage_name}"] = sum(durations)
return parent_traces
def _build_aggregated_response(
parent_job_id: str,
subjob_results: List[Optional[Dict[str, Any]]],
failed_subjobs: List[Dict[str, object]],
ordered_descriptors: List[Dict[str, Any]],
metadata: Dict[str, Any],
) -> Dict[str, Any]:
"""
Build the aggregated response from subjob results.
Parameters
----------
parent_job_id : str
The parent job identifier
subjob_results : List[Optional[Dict[str, Any]]]
Results from all subjobs (None for failed ones)
failed_subjobs : List[Dict[str, object]]
List of failed subjob information
ordered_descriptors : List[Dict[str, Any]]
Subjob descriptors in original order
metadata : Dict[str, Any]
Parent job metadata
Returns
-------
Dict[str, Any]
Aggregated response with combined data and metadata
"""
any_failed = len(failed_subjobs) > 0
subjob_ids = [desc.get("job_id") for desc in ordered_descriptors]
aggregated_result = {
"data": [],
"status": "failed" if any_failed else "success",
"description": (
"One or more subjobs failed to complete" if any_failed else "Aggregated result composed from subjob outputs"
),
# Top-level trace/annotations for V1 compatibility
"trace": {},
"annotations": {},
"metadata": {
"parent_job_id": parent_job_id,
"total_pages": metadata.get("total_pages", len(subjob_ids)),
"pages_per_chunk": metadata.get("pages_per_chunk"),
"original_source_id": metadata.get("original_source_id"),
"original_source_name": metadata.get("original_source_name"),
"subjobs_failed": sum(1 for r in subjob_results if r is None),
"failed_subjobs": failed_subjobs,
"subjob_ids": subjob_ids,
"chunks": [],
"trace_segments": [],
"annotation_segments": [],
},
}
# Aggregate subjob data in page order
for page_num, (result, descriptor) in enumerate(zip(subjob_results, ordered_descriptors), 1):
if result is not None:
# Add page data to aggregated result
if "data" in result:
normalized_records = _normalize_chunk_records(result.get("data"), descriptor, metadata)
aggregated_result["data"].extend(normalized_records)
chunk_entry = dict(descriptor)
aggregated_result["metadata"]["chunks"].append(chunk_entry)
trace_data, annotation_data = _extract_ray_telemetry(result)
start_page = descriptor.get("start_page")
end_page = descriptor.get("end_page")
if trace_data:
# Add to trace_segments (detailed, per-chunk view)
aggregated_result["metadata"]["trace_segments"].append(
{
"job_id": descriptor.get("job_id"),
"chunk_index": descriptor.get("chunk_index"),
"start_page": start_page,
"end_page": end_page,
"trace": trace_data,
}
)
# Chunk traces stay in metadata.trace_segments only (not in top-level)
if annotation_data:
# Add to annotation_segments (detailed, per-chunk view)
aggregated_result["metadata"]["annotation_segments"].append(
{
"job_id": descriptor.get("job_id"),
"chunk_index": descriptor.get("chunk_index"),
"start_page": start_page,
"end_page": end_page,
"annotations": annotation_data,
}
)
# Merge into top-level annotations (annotations have unique UUIDs, safe to merge)
aggregated_result["annotations"].update(annotation_data)
else:
# Note failed page
logger.warning(f"Page {page_num} failed or missing")
# Compute parent-level trace aggregations from trace_segments
trace_segments = aggregated_result["metadata"]["trace_segments"]
if trace_segments:
# Build a temporary chunk trace dict for aggregation
temp_chunk_traces = {}
for segment in trace_segments:
chunk_idx = segment.get("chunk_index")
chunk_trace = segment.get("trace", {})
for trace_key, trace_value in chunk_trace.items():
prefixed_key = f"chunk_{chunk_idx}::{trace_key}"
temp_chunk_traces[prefixed_key] = trace_value
# Aggregate and set as top-level trace (only parent traces, no chunk traces)
parent_level_traces = _aggregate_parent_traces(temp_chunk_traces)
aggregated_result["trace"] = parent_level_traces
return aggregated_result
# ---------------------------------------------------------------------------
# Bursty submission helpers (fairness without long-lived in-flight tasks)
# ---------------------------------------------------------------------------
def _get_submit_burst_params() -> Tuple[int, int, int]:
"""
Returns (burst_size, pause_ms, jitter_ms) from environment with sane defaults.
- V2_SUBMIT_BURST_SIZE (default: 16)
- V2_SUBMIT_BURST_PAUSE_MS (default: 25)
- V2_SUBMIT_BURST_JITTER_MS (default: 10)
"""
burst_size = int(os.getenv("V2_SUBMIT_BURST_SIZE", "16"))
pause_ms = int(os.getenv("V2_SUBMIT_BURST_PAUSE_MS", "50"))
jitter_ms = int(os.getenv("V2_SUBMIT_BURST_JITTER_MS", "15"))
return max(1, burst_size), max(0, pause_ms), max(0, jitter_ms)
async def _submit_subjobs_in_bursts(
items: List[Tuple[str, MessageWrapper]],
ingest_service: "INGEST_SERVICE_T",
*,
burst_size: int,
pause_ms: int,
jitter_ms: int,
) -> None:
"""
Submit subjobs in sequential bursts and await each burst to completion.
This avoids keeping a large number of pending tasks in the REST handler
and allows other concurrent requests to interleave enqueue work between bursts.
"""
for offset in range(0, len(items), burst_size):
burst = items[offset : offset + burst_size]
tasks = [ingest_service.submit_job(wrapper, subjob_id) for (subjob_id, wrapper) in burst]
# Propagate any errors from this burst
await asyncio.gather(*tasks)
# Pause with jitter to yield to other request handlers before next burst
if offset + burst_size < len(items):
delay_ms = pause_ms + (random.randint(0, jitter_ms) if jitter_ms > 0 else 0)
if delay_ms > 0:
await asyncio.sleep(delay_ms / 1000.0)
# POST /v2/submit_job
[docs]
@router.post(
"/submit_job",
responses={
200: {"description": "Jobs were successfully submitted"},
500: {"description": "Error encountered while submitting jobs."},
503: {"description": "Service unavailable."},
},
tags=["Ingestion"],
summary="submit jobs to the core nv ingestion service for processing with PDF splitting",
operation_id="submit_job_v2",
)
@traced_endpoint("http-submit-job-v2")
async def submit_job_v2(
request: Request, response: Response, job_spec: MessageWrapper, ingest_service: INGEST_SERVICE_T
):
span = trace.get_current_span()
source_id = None
document_type = None
try:
span.add_event("Submitting file for processing (V2)")
current_trace_id = span.get_span_context().trace_id
parent_job_id = trace_id_to_uuid(current_trace_id)
# Parse job spec
job_spec_dict = json.loads(job_spec.payload)
# Extract PDF configuration if provided by client
pdf_config = job_spec_dict.get("pdf_config", {})
client_split_page_count = pdf_config.get("split_page_count") if pdf_config else None
# Extract document type and payload from the proper structure
job_payload = job_spec_dict.get("job_payload", {})
document_types = job_payload.get("document_type", [])
payloads = job_payload.get("content", [])
# Resolve original source metadata up front for logging / subjob naming
source_ids = job_payload.get("source_id", ["unknown_source.pdf"])
source_names = job_payload.get("source_name", ["unknown_source.pdf"])
original_source_id = source_ids[0] if source_ids else "unknown_source.pdf"
original_source_name = source_names[0] if source_names else "unknown_source.pdf"
# Track page count for all PDFs (used for both splitting logic and metadata)
pdf_page_count_cache = None
submission_items: List[Tuple[str, MessageWrapper]] = []
subjob_ids: List[str] = []
subjob_descriptors: List[Dict[str, Any]] = []
parent_metadata: Dict[str, Any] = {}
submission_items: List[Tuple[str, MessageWrapper]] = []
try:
parent_uuid = uuid.UUID(parent_job_id)
except ValueError:
logger.warning(
"Parent job id %s is not a valid UUID; generating fallback namespace for subjobs",
parent_job_id,
)
parent_uuid = uuid.uuid4()
# Check if this is a PDF that needs splitting
if document_types and payloads and document_types[0].lower() == "pdf":
# Decode the payload to check page count
pdf_content = base64.b64decode(payloads[0])
page_count = get_pdf_page_count(pdf_content)
pdf_page_count_cache = page_count # Cache for later use
qos_tier = get_qos_tier_for_page_count(page_count)
pages_per_chunk = get_pdf_split_page_count(client_override=client_split_page_count)
document_type = DocumentTypeEnum.PDF
# Split if the document has more pages than our chunk size
if page_count > pages_per_chunk:
logger.warning(
"Splitting PDF %s into %s-page chunks (total pages: %s) -> (qos_tier: %s)",
original_source_name,
pages_per_chunk,
page_count,
qos_tier,
)
chunks = split_pdf_to_chunks(pdf_content, pages_per_chunk)
subjob_ids: List[str] = []
subjob_descriptors: List[Dict[str, Any]] = []
submission_items: List[Tuple[str, MessageWrapper]] = []
try:
parent_uuid = uuid.UUID(parent_job_id)
except ValueError:
logger.warning(
"Parent job id %s is not a valid UUID; generating fallback namespace for subjobs",
parent_job_id,
)
parent_uuid = uuid.uuid4()
for chunk in chunks:
start = chunk["start_page"]
end = chunk["end_page"]
page_suffix = f"page_{start}" if start == end else f"pages_{start}-{end}"
source_id = f"{original_source_id}#{page_suffix}"
source_name = f"{original_source_name}#{page_suffix}"
subjob_id, subjob_wrapper = _prepare_chunk_submission(
job_spec_dict,
chunk,
document_type=DocumentTypeEnum.PDF,
parent_uuid=parent_uuid,
parent_job_id=parent_job_id,
current_trace_id=current_trace_id,
source_id=source_id,
source_name=source_name,
)
# Inject QoS routing hint into subjob routing_options (keeps API and service loosely coupled)
try:
sub_spec = json.loads(subjob_wrapper.payload)
routing_opts = sub_spec.get("routing_options") or {}
routing_opts["queue_hint"] = qos_tier
sub_spec["routing_options"] = routing_opts
subjob_wrapper = MessageWrapper(payload=json.dumps(sub_spec))
except Exception:
# Best-effort; if we cannot inject, fall back to default routing
pass
submission_items.append((subjob_id, subjob_wrapper))
subjob_ids.append(subjob_id)
subjob_descriptors.append(
{
"job_id": subjob_id,
"chunk_index": len(subjob_descriptors) + 1,
"start_page": chunk.get("start_page"),
"end_page": chunk.get("end_page"),
"page_count": chunk.get("page_count"),
}
)
parent_metadata.update(
{
"total_pages": page_count,
"pages_per_chunk": pages_per_chunk,
"original_source_id": original_source_id,
"original_source_name": original_source_name,
"document_type": document_types[0] if document_types else "pdf",
"subjob_order": subjob_ids,
}
)
elif document_types and payloads and document_types[0].lower() in ["mp4", "mov", "avi", "mp3", "wav", "mkv"]:
document_type = document_types[0]
upload_path = f"./{Path(original_source_id).name}"
# dump the payload to a file, just came from client
with fsspec.open(upload_path, "wb") as f:
f.write(base64.b64decode(payloads[0]))
dataloader = DataLoader(
path=upload_path, output_dir="./audio_chunks/", audio_only=True, split_interval=50000000
)
document_type = DocumentTypeEnum.MP3
parent_uuid = uuid.UUID(parent_job_id)
for task in job_spec_dict["tasks"]:
if "task_properties" in task and "document_type" in task["task_properties"]:
task["task_properties"]["document_type"] = document_type
end = 0
for idx, (file_path, duration) in enumerate(dataloader.files_completed):
start = end
end = int(start + duration)
chunk = {
"bytes": file_path.encode("utf-8"),
"chunk_index": idx,
"start": start,
"end": end,
}
subjob_id, subjob_wrapper = _prepare_chunk_submission(
job_spec_dict,
chunk,
parent_uuid=parent_uuid,
parent_job_id=parent_job_id,
current_trace_id=current_trace_id,
source_id=file_path,
source_name=upload_path,
document_type=document_type,
)
submission_items.append((subjob_id, subjob_wrapper))
subjob_ids.append(subjob_id)
subjob_descriptors.append(
{
"job_id": subjob_id,
"chunk_index": idx + 1,
"start_page": chunk.get("start"),
"end_page": chunk.get("end"),
"page_count": chunk.get("page_count", 0),
}
)
logger.debug(f"Removing uploaded file {upload_path}")
os.remove(upload_path)
if submission_items:
burst_size, pause_ms, jitter_ms = _get_submit_burst_params()
await _submit_subjobs_in_bursts(
submission_items,
ingest_service,
burst_size=burst_size,
pause_ms=pause_ms,
jitter_ms=jitter_ms,
)
parent_metadata.update(
{
"original_source_id": original_source_id,
"original_source_name": original_source_name,
"document_type": document_type,
"subjob_order": subjob_ids,
}
)
# raise ValueError(f"Setting parent job mapping for {parent_job_id} with {len(subjob_ids)} subjobs")
await ingest_service.set_parent_job_mapping(
parent_job_id,
subjob_ids,
parent_metadata,
subjob_descriptors=subjob_descriptors,
)
await ingest_service.set_job_state(parent_job_id, STATE_SUBMITTED)
span.add_event(f"Split into {len(subjob_ids)} subjobs")
response.headers["x-trace-id"] = trace.format_trace_id(current_trace_id)
return parent_job_id
# For non-PDFs or cases where splitting is not required, submit as normal
if "tracing_options" not in job_spec_dict:
job_spec_dict["tracing_options"] = {"trace": True}
job_spec_dict["tracing_options"]["trace_id"] = str(current_trace_id)
# If this was a PDF and we computed page_count, route the single job using the same QoS tier
try:
if (
document_types
and document_types[0].lower() == "pdf"
and "queue_hint" not in (job_spec_dict.get("routing_options") or {})
):
job_spec_dict.setdefault("routing_options", {})["queue_hint"] = qos_tier
except Exception:
pass
updated_job_spec = MessageWrapper(payload=json.dumps(job_spec_dict))
span.add_event("Submitting as single job (no split needed)")
# Submit the job to the pipeline task queue
await ingest_service.submit_job(updated_job_spec, parent_job_id)
await ingest_service.set_job_state(parent_job_id, STATE_SUBMITTED)
# If this was a PDF (even if not split), store page count metadata for tracking
if pdf_page_count_cache is not None:
try:
# Use cached page count from earlier check to avoid re-decoding
# Store minimal metadata for non-split PDFs (consistent with split PDFs)
single_pdf_metadata: Dict[str, Any] = {
"total_pages": pdf_page_count_cache,
"pages_per_chunk": pdf_page_count_cache, # Single chunk = entire document
"original_source_id": original_source_id,
"original_source_name": original_source_name,
"document_type": document_types[0],
"subjob_order": [], # No subjobs for non-split PDFs
}
# Store as parent job metadata with empty subjob list for consistency
await ingest_service.set_parent_job_mapping(
parent_job_id,
[], # Empty subjob list
single_pdf_metadata,
subjob_descriptors=[],
)
logger.debug(
f"Stored page count metadata for non-split PDF {original_source_name}: {pdf_page_count_cache} pages"
)
except Exception as metadata_err:
# Don't fail the job if metadata storage fails
logger.warning(f"Failed to store page count metadata for {parent_job_id}: {metadata_err}")
response.headers["x-trace-id"] = trace.format_trace_id(current_trace_id)
return parent_job_id
except Exception as ex:
logger.exception(f"Error submitting job: {str(ex)}, {source_id}")
raise HTTPException(status_code=500, detail=f"Nv-Ingest Internal Server Error: {str(ex)}, for: \n{source_id}")
# GET /v2/fetch_job
[docs]
@router.get(
"/fetch_job/{job_id}",
responses={
200: {"description": "Job result successfully retrieved."},
202: {"description": "Job is processing or result not yet available. Retry later."},
404: {"description": "Job ID not found or associated state has expired."},
410: {"description": "Job result existed but is now gone (expired or retrieved destructively/cached)."},
500: {"description": "Internal server error during fetch processing."},
503: {"description": "Job processing failed, or backend service temporarily unavailable preventing fetch."},
},
tags=["Ingestion"],
summary="Fetch the result of a previously submitted job by its job_id (V2 with aggregation)",
operation_id="fetch_job_v2",
)
async def fetch_job_v2(job_id: str, ingest_service: INGEST_SERVICE_T):
"""
V2 fetch that handles parent job aggregation.
"""
try:
# Check if this is a parent job with subjobs
subjob_info = await ingest_service.get_parent_job_info(job_id)
if subjob_info is None:
# Not a parent job, fetch identical to V1
current_state = await ingest_service.get_job_state(job_id)
logger.debug(f"Initial state check for job {job_id}: {current_state}")
if current_state is None:
logger.warning(f"Job {job_id} not found or expired. Returning 404.")
raise HTTPException(status_code=404, detail="Job ID not found or state has expired.")
if current_state == STATE_FAILED:
logger.error(f"Job {job_id} failed. Returning 503.")
raise HTTPException(status_code=503, detail="Job processing failed.")
if current_state == STATE_RETRIEVED_DESTRUCTIVE:
logger.warning(f"Job {job_id} was destructively retrieved. Returning 410.")
raise HTTPException(status_code=410, detail="Job result is gone (destructive read).")
if current_state in INTERMEDIATE_STATES or current_state in {
STATE_RETRIEVED_NON_DESTRUCTIVE,
STATE_RETRIEVED_CACHED,
}:
logger.debug(f"Attempting fetch for job {job_id} in state {current_state}.")
try:
job_response = await ingest_service.fetch_job(job_id)
logger.debug(f"Fetched result for job {job_id}.")
try:
current_fetch_mode = await ingest_service.get_fetch_mode()
if current_fetch_mode == FetchMode.DESTRUCTIVE:
target_state = STATE_RETRIEVED_DESTRUCTIVE
elif current_fetch_mode == FetchMode.NON_DESTRUCTIVE:
target_state = STATE_RETRIEVED_NON_DESTRUCTIVE
elif current_fetch_mode == FetchMode.CACHE_BEFORE_DELETE:
target_state = STATE_RETRIEVED_CACHED
else:
target_state = "RETRIEVED_UNKNOWN"
if target_state != "RETRIEVED_UNKNOWN":
await ingest_service.set_job_state(job_id, target_state)
logger.debug(f"Updated job {job_id} state to {target_state}.")
except Exception as state_err:
logger.error(f"Failed to set job state for {job_id} after fetch: {state_err}")
try:
json_bytes = json.dumps(job_response).encode("utf-8")
return StreamingResponse(iter([json_bytes]), media_type="application/json", status_code=200)
except TypeError as json_err:
logger.exception(f"Serialization error for job {job_id}: {json_err}")
raise HTTPException(
status_code=500, detail="Internal server error: Failed to serialize result."
)
except (TimeoutError, RedisError, ConnectionError) as fetch_err:
# Handle timeout/error cases same as V1
fetch_err_type = type(fetch_err).__name__
if isinstance(fetch_err, TimeoutError):
logger.debug(
f"Job {job_id} still processing (state: {current_state}), fetch attempt timed out cleanly."
)
else:
logger.warning(
f"Backend error ({fetch_err_type}) during fetch attempt for job {job_id} "
f"(state: {current_state}): {fetch_err}"
)
if current_state == STATE_RETRIEVED_NON_DESTRUCTIVE:
if isinstance(fetch_err, TimeoutError):
raise HTTPException(status_code=410, detail="Job result is gone (TTL expired).")
else:
raise HTTPException(
status_code=503, detail="Backend service unavailable preventing access to job result."
)
elif current_state == STATE_RETRIEVED_CACHED:
raise HTTPException(
status_code=410, detail="Job result is gone (previously cached, fetch failed)."
)
elif current_state in INTERMEDIATE_STATES:
if isinstance(fetch_err, TimeoutError):
raise HTTPException(
status_code=202, detail=f"Job is processing (state: {current_state}). Retry later."
)
else:
raise HTTPException(
status_code=503, detail="Backend service unavailable preventing fetch of job result."
)
else:
logger.error(f"Unexpected state '{current_state}' for job {job_id} after fetch failure.")
raise HTTPException(
status_code=500, detail="Internal server error: Unexpected job state after fetch failure."
)
else:
logger.error(f"Unknown job state '{current_state}' for job {job_id}.")
raise HTTPException(
status_code=500, detail=f"Internal server error: Unknown job state '{current_state}'."
)
else:
# This is a parent job - orchestrate aggregation using declarative helpers
subjob_ids = subjob_info.get("subjob_ids", [])
metadata = subjob_info.get("metadata", {})
logger.debug(f"Parent job {job_id} has {len(subjob_ids)} subjobs")
# Special case: Non-split PDFs have metadata but no subjobs
# Fetch the result directly and augment with page count metadata
if len(subjob_ids) == 0:
logger.debug(f"Job {job_id} is a non-split PDF, fetching result directly")
try:
job_response = await ingest_service.fetch_job(job_id)
# Augment response with page count metadata
if isinstance(job_response, dict):
if "metadata" not in job_response:
job_response["metadata"] = {}
job_response["metadata"]["total_pages"] = metadata.get("total_pages")
job_response["metadata"]["original_source_id"] = metadata.get("original_source_id")
job_response["metadata"]["original_source_name"] = metadata.get("original_source_name")
# Update job state after successful fetch
await _update_job_state_after_fetch(job_id, ingest_service)
return _stream_json_response(job_response)
except (TimeoutError, RedisError, ConnectionError):
logger.debug(f"Job {job_id} (non-split PDF) not ready yet")
raise HTTPException(status_code=202, detail="Job is processing. Retry later.")
except Exception as e:
logger.exception(f"Error fetching non-split PDF job {job_id}: {e}")
raise HTTPException(status_code=500, detail="Internal server error during job fetch.")
# Build ordered descriptors for subjobs
stored_descriptors = subjob_info.get("subjob_descriptors") or []
descriptor_lookup = {entry.get("job_id"): entry for entry in stored_descriptors if isinstance(entry, dict)}
ordered_descriptors: List[Dict[str, Any]] = []
for idx, subjob_id in enumerate(subjob_ids, 1):
descriptor = descriptor_lookup.get(subjob_id, {})
ordered_descriptors.append(
{
"job_id": subjob_id,
"chunk_index": descriptor.get("chunk_index", idx),
"start_page": descriptor.get("start_page"),
"end_page": descriptor.get("end_page"),
"page_count": descriptor.get("page_count"),
}
)
# Calculate max parallel operations (stay within Redis connection pool)
max_parallel_ops = max(
1, min(len(ordered_descriptors), getattr(ingest_service, "_concurrency_level", 10) // 2)
)
# Check all subjob states (raises 202 if any still processing)
subjob_states, failed_subjobs = await _check_all_subjob_states(
ordered_descriptors, max_parallel_ops, ingest_service
)
# Fetch all subjob results (raises 202 if any not ready)
subjob_results = await _fetch_all_subjob_results(
ordered_descriptors, subjob_states, failed_subjobs, max_parallel_ops, ingest_service
)
# Build aggregated response from all subjob results
aggregated_result = _build_aggregated_response(
job_id, subjob_results, failed_subjobs, ordered_descriptors, metadata
)
# Update parent job state after successful aggregation
await _update_job_state_after_fetch(job_id, ingest_service)
# Return aggregated result as streaming response
return _stream_json_response(aggregated_result)
except HTTPException:
raise
except Exception as e:
logger.exception(f"Unexpected error in fetch_job_v2: {e}")
raise HTTPException(status_code=500, detail="Internal server error during job fetch.")
