nv-ingest-client 2025.10.18.dev20251018__tar.gz → 2025.12.14.dev20251214__tar.gz

{nv_ingest_client-2025.10.18.dev20251018/src/nv_ingest_client.egg-info → nv_ingest_client-2025.12.14.dev20251214}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nv-ingest-client
-Version: 2025.10.18.dev20251018
+Version: 2025.12.14.dev20251214
 Summary: Python client for the nv-ingest service
 Author-email: Jeremy Dyer <jdyer@nvidia.com>
 License:                                  Apache License
@@ -223,6 +223,7 @@ Requires-Dist: pydantic-settings>2.0.0
 Requires-Dist: requests>=2.28.2
 Requires-Dist: setuptools>=78.1.1
 Requires-Dist: tqdm>=4.67.1
+Requires-Dist: lancedb>=0.25.3
 Provides-Extra: milvus
 Requires-Dist: pymilvus==2.5.10; extra == "milvus"
 Requires-Dist: pymilvus[bulk_writer,model]; extra == "milvus"

{nv_ingest_client-2025.10.18.dev20251018 → nv_ingest_client-2025.12.14.dev20251214}/pyproject.toml

@@ -30,6 +30,7 @@ dependencies = [
     "requests>=2.28.2",
     "setuptools>=78.1.1",
     "tqdm>=4.67.1",
+    "lancedb>=0.25.3",
 ]
 
 [project.optional-dependencies]

{nv_ingest_client-2025.10.18.dev20251018 → nv_ingest_client-2025.12.14.dev20251214}/src/nv_ingest_client/client/client.py

@@ -44,6 +44,50 @@ from nv_ingest_client.util.util import (
 logger = logging.getLogger(__name__)
 
 
+def _compute_resident_times(trace_dict: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Compute resident_time entries from entry/exit pairs if not already present.
+
+    This ensures consistency between split jobs (where server computes resident_time)
+    and non-split jobs (where we compute it client-side).
+
+    Parameters
+    ----------
+    trace_dict : Dict[str, Any]
+        Trace dictionary with entry/exit pairs
+
+    Returns
+    -------
+    Dict[str, Any]
+        Trace dictionary with resident_time entries added
+    """
+    if not trace_dict or not isinstance(trace_dict, dict):
+        return trace_dict
+
+    # Check if resident_time already exists (server-computed for split jobs)
+    has_resident = any(k.startswith("trace::resident_time::") for k in trace_dict.keys())
+    if has_resident:
+        return trace_dict  # Already computed by server
+
+    # Compute resident_time from entry/exit pairs
+    result = dict(trace_dict)
+    stages = set()
+
+    # Find all unique stages
+    for key in trace_dict:
+        if key.startswith("trace::entry::"):
+            stages.add(key.replace("trace::entry::", ""))
+
+    # Compute resident_time for each stage
+    for stage in stages:
+        entry_key = f"trace::entry::{stage}"
+        exit_key = f"trace::exit::{stage}"
+        if entry_key in trace_dict and exit_key in trace_dict:
+            result[f"trace::resident_time::{stage}"] = trace_dict[exit_key] - trace_dict[entry_key]
+
+    return result
+
+
 class DataDecodeException(Exception):
     """
     Exception raised for errors in decoding data.
@@ -87,6 +131,7 @@ class _ConcurrentProcessor:
         stream_to_callback_only: bool,
         return_full_response: bool,
         verbose: bool = False,
+        return_traces: bool = False,
     ):
         """
         Initializes the concurrent processor.
@@ -120,6 +165,8 @@ class _ConcurrentProcessor:
             initiating job submission or fetching fails for a batch.
         verbose : bool, optional
             If True, enables detailed debug logging. Default is False.
+        return_traces : bool, optional
+            If True, parent-level trace data for each completed job is stored.
 
         Raises
         ------
@@ -142,17 +189,26 @@ class _ConcurrentProcessor:
         self.stream_to_callback_only = stream_to_callback_only
         self.return_full_response = return_full_response
         self.verbose = verbose
+        self.return_traces = return_traces
 
         # State variables managed across batch cycles
         self.retry_job_ids: List[str] = []
         self.retry_counts: Dict[str, int] = defaultdict(int)
         self.results: List[Dict[str, Any]] = []  # Stores successful results (full dicts)
         self.failures: List[Tuple[str, str]] = []  # (job_index, error_message)
+        self.traces: List[Optional[Dict[str, Any]]] = []
 
         # --- Initial Checks ---
         if not self.job_queue_id:
             logger.warning("job_queue_id is not set; submission of new jobs will fail.")
 
+        # Executor check required for run_async
+        if not hasattr(client, "_worker_pool"):
+            raise AttributeError("Client object is missing the '_worker_pool' attribute, required for run_async.")
+        if not isinstance(client._worker_pool, ThreadPoolExecutor):
+            raise TypeError("Client's '_worker_pool' must be a ThreadPoolExecutor for run_async.")
+        self._executor = client._worker_pool
+
     # --------------------------------------------------------------------------
     # Private Methods
     # --------------------------------------------------------------------------
@@ -197,7 +253,7 @@ class _ConcurrentProcessor:
         # Attempt to mark state as FAILED locally in the client (best effort)
         try:
             # Use a method assumed to safely get the state object
-            job_state = self.client._get_job_state_object(job_index)
+            job_state = self.client._get_and_check_job_state(job_index)
             # Check state exists and is not already terminal before updating
             if (
                 job_state and hasattr(job_state, "state") and job_state.state not in ["FAILED", "COMPLETED"]
@@ -247,6 +303,14 @@ class _ConcurrentProcessor:
             # When requested, return the full response envelope (includes 'trace' and 'annotations')
             self.results.append(result_data if self.return_full_response else result_data.get("data"))
 
+        # Extract trace data for all successful (non-failed) jobs
+        if self.return_traces and not is_failed:
+            trace_payload = result_data.get("trace") if result_data else None
+            # Compute resident_time if not already present (for consistency)
+            if trace_payload:
+                trace_payload = _compute_resident_times(trace_payload)
+            self.traces.append(trace_payload if trace_payload else None)
+
         # Cleanup retry count if it exists
         if job_index in self.retry_counts:
             del self.retry_counts[job_index]
@@ -438,7 +502,10 @@ class _ConcurrentProcessor:
 
         return batch_futures_dict, normalized_job_indices
 
-    def run(self) -> Tuple[List[Dict[str, Any]], List[Tuple[str, str]]]:
+    # --------------------------------------------------------------------------
+    # Core Processing Logic
+    # --------------------------------------------------------------------------
+    def _process_all_jobs(self) -> Tuple[List[Any], List[Tuple[str, str]], List[Optional[Dict[str, Any]]]]:
         """
         Executes the main processing loop in batches.
 
@@ -581,7 +648,45 @@ class _ConcurrentProcessor:
         # --- Final Logging ---
         self._log_final_status(total_jobs)
 
-        return self.results, self.failures
+        return self.results, self.failures, self.traces if self.return_traces else []
+
+    # --------------------------------------------------------------------------
+    # Public Methods
+    # --------------------------------------------------------------------------
+
+    def run(self) -> Tuple[List[Dict[str, Any]], List[Tuple[str, str]], List[Optional[Dict[str, Any]]]]:
+        """
+        Executes the main processing loop synchronously.
+
+        This method orchestrates the job processing by maintaining a constant
+        pool of in-flight jobs, handling submissions, fetches, and retries until
+        all jobs are complete. It blocks until all jobs are processed.
+
+        Returns
+        -------
+        Tuple[List[Any], List[Tuple[str, str]], List[Optional[Dict[str, Any]]]]
+            A tuple containing:
+            1. A list of successfully fetched job results.
+            2. A list of tuples for failed jobs (job_index, error_message).
+            3. A list of trace dictionaries if `return_traces` was True.
+        """
+        return self._process_all_jobs()
+
+    def run_async(self) -> Future[Tuple[List[Any], List[Tuple[str, str]], List[Optional[Dict[str, Any]]]]]:
+        """
+        Executes the main processing loop asynchronously.
+
+        Submits the entire processing logic to the client's background
+        thread pool and returns a Future that resolves with the final
+        results, failures, and traces once all jobs are complete.
+
+        Returns
+        -------
+        Future
+            A future representing the asynchronous execution. Its result()
+            will be a tuple containing (results, failures, traces).
+        """
+        return self._executor.submit(self._process_all_jobs)
 
 
 class NvIngestClient:
@@ -1212,7 +1317,12 @@ class NvIngestClient:
         stream_to_callback_only: bool = False,
         return_full_response: bool = False,
         verbose: bool = False,
-    ) -> Union[List[Any], Tuple[List[Any], List[Tuple[str, str]]]]:
+        return_traces: bool = False,
+    ) -> Union[
+        List[Any],
+        Tuple[List[Any], List[Tuple[str, str]]],
+        Tuple[List[Any], List[Tuple[str, str]], List[Optional[Dict[str, Any]]]],
+    ]:
         """
         Submit and fetch multiple jobs concurrently.
 
@@ -1247,6 +1357,8 @@ class NvIngestClient:
             Ignored when stream_to_callback_only=True. Default is False.
         verbose : bool, optional
             If True, enable debug logging. Default is False.
+        return_traces : bool, optional
+            If True, parent-level aggregated trace metrics are extracted and returned. Default is False.
 
         Returns
         -------
@@ -1254,6 +1366,9 @@ class NvIngestClient:
             List of successful job results when `return_failures` is False.
         results, failures : tuple
             Tuple of (successful results, failure tuples) when `return_failures` is True.
+        results, failures, traces : tuple
+            Tuple of (successful results, failure tuples, trace dicts) when both
+            `return_failures` and `return_traces` are True.
 
         Raises
         ------
@@ -1266,7 +1381,12 @@ class NvIngestClient:
 
         # Handle empty input
         if not job_indices:
-            return ([], []) if return_failures else []
+            if return_failures and return_traces:
+                return [], [], []
+            elif return_failures:
+                return [], []
+            else:
+                return []
 
         # Validate and set batch_size
         validated_batch_size = self._validate_batch_size(batch_size)
@@ -1289,17 +1409,84 @@ class NvIngestClient:
             stream_to_callback_only=stream_to_callback_only,
             return_full_response=return_full_response,
             verbose=verbose,
+            return_traces=return_traces,
         )
 
-        results, failures = processor.run()
+        results, failures, traces = processor.run()
 
-        if return_failures:
+        if return_failures and return_traces:
+            return results, failures, traces
+        elif return_failures:
             return results, failures
+        elif return_traces:
+            return results, traces
 
         if failures:
             logger.warning(f"{len(failures)} job(s) failed during concurrent processing." " Check logs for details.")
         return results
 
+    def process_jobs_concurrently_async(
+        self,
+        job_indices: Union[str, List[str]],
+        job_queue_id: Optional[str] = None,
+        batch_size: Optional[int] = None,
+        timeout: int = 100,
+        max_job_retries: Optional[int] = None,
+        retry_delay: float = 0.5,
+        initial_fetch_delay: float = 0.3,
+        fail_on_submit_error: bool = False,
+        completion_callback: Optional[Callable[[Any, str], None]] = None,
+        stream_to_callback_only: bool = False,
+        return_full_response: bool = False,
+        verbose: bool = False,
+        return_traces: bool = False,
+    ) -> Future[Tuple[List[Any], List[Tuple[str, str]], List[Optional[Dict[str, Any]]]]]:
+        """
+        Submit and fetch multiple jobs concurrently and asynchronously.
+
+        This method initializes the processing and returns a Future immediately. The Future
+        will resolve with a fixed 3-part tuple `(results, failures, traces)` once all
+        jobs have completed.
+
+        Parameters are identical to `process_jobs_concurrently`.
+
+        Returns
+        -------
+        Future[Tuple[List[Any], List[Tuple[str, str]], List[Optional[Dict[str, Any]]]]]
+            A future that completes when all jobs are done. Its result is a tuple
+            containing (successful_results, failures, traces).
+        """
+        if isinstance(job_indices, str):
+            job_indices = [job_indices]
+
+        if not job_indices:
+            immediate_future: Future = Future()
+            immediate_future.set_result(([], [], []))
+            return immediate_future
+
+        validated_batch_size = self._validate_batch_size(batch_size)
+        effective_timeout: Tuple[int, Optional[float]] = (int(timeout), None)
+
+        processor = _ConcurrentProcessor(
+            client=self,
+            batch_size=validated_batch_size,
+            job_indices=job_indices,
+            job_queue_id=job_queue_id,
+            timeout=effective_timeout,
+            max_job_retries=max_job_retries,
+            retry_delay=retry_delay,
+            initial_fetch_delay=initial_fetch_delay,
+            completion_callback=completion_callback,
+            fail_on_submit_error=fail_on_submit_error,
+            stream_to_callback_only=stream_to_callback_only,
+            return_full_response=return_full_response,
+            verbose=verbose,
+            return_traces=return_traces,
+        )
+
+        # Asynchronous call
+        return processor.run_async()
+
     def _ensure_submitted(self, job_ids: Union[str, List[str]]) -> None:
         """
         Block until all specified jobs have been marked submitted.
@@ -1628,9 +1815,6 @@ class NvIngestClient:
                     )
                     logger.error(error_msg)
                     failures.append((self._job_index_to_job_spec[job_id].source_id, str(e)))
-                finally:
-                    # Clean up the job spec mapping
-                    del self._job_index_to_job_spec[job_id]
 
         if return_failures:
             return results, failures

{nv_ingest_client-2025.10.18.dev20251018 → nv_ingest_client-2025.12.14.dev20251214}/src/nv_ingest_client/client/ingest_job_handler.py

@@ -323,18 +323,40 @@ class IngestJobHandler:
 
                 futures_dict: Dict[Any, str] = self.client.fetch_job_result_async(self._job_ids_batch, data_only=False)
                 for future in as_completed(futures_dict.keys()):
+                    pages_per_sec = None
                     try:
                         # Block as each future completes; this mirrors CLI behavior
                         future_response, trace_id = self._handle_future_result(future)
                         job_id: str = futures_dict[future]
                         trace_ids[job_id_map[job_id]] = trace_id
 
-                        first_page_metadata = future_response["data"][0]["metadata"]
-                        file_page_counts: Dict[str, int] = {
-                            first_page_metadata["source_metadata"]["source_name"]: first_page_metadata[
-                                "content_metadata"
-                            ]["hierarchy"]["page_count"]
-                        }
+                        # Extract page count: prefer V2 metadata location, fall back to V1
+                        page_count = None
+                        source_name = None
+
+                        # Try V2 metadata location first (top-level metadata.total_pages)
+                        if "metadata" in future_response and future_response["metadata"]:
+                            response_metadata = future_response["metadata"]
+                            page_count = response_metadata.get("total_pages")
+                            source_name = response_metadata.get("original_source_name")
+
+                        # Fall back to V1 location (first data element's hierarchy.page_count)
+                        if page_count is None and future_response.get("data"):
+                            try:
+                                first_page_metadata = future_response["data"][0]["metadata"]
+                                page_count = first_page_metadata["content_metadata"]["hierarchy"]["page_count"]
+                                source_name = first_page_metadata["source_metadata"]["source_name"]
+                            except (KeyError, IndexError, TypeError):
+                                # If we can't extract from V1 location, use defaults
+                                pass
+
+                        # Use extracted values or defaults
+                        if page_count is None:
+                            page_count = 0  # Default if not found
+                        if source_name is None:
+                            source_name = "unknown_source"
+
+                        file_page_counts: Dict[str, int] = {source_name: page_count}
 
                         if self.output_directory:
                             self._save_response_data(