nv-ingest-api 2025.10.22.dev20251022__py3-none-any.whl → 2025.11.2.dev20251102__py3-none-any.whl

nv_ingest_api/internal/transform/embed_text.py

@@ -7,6 +7,7 @@ from concurrent.futures import ThreadPoolExecutor
 from functools import partial
 from typing import Any, Dict, Tuple, Optional, Iterable, List
 
+import glom
 import pandas as pd
 from openai import OpenAI
 
@@ -282,6 +283,33 @@ def _add_embeddings(row, embeddings, info_msgs):
     return row
 
 
+def _add_custom_embeddings(row, embeddings, result_target_field):
+    """
+    Updates a DataFrame row with embedding data and associated error info
+    based on a user supplied custom content field.
+
+    Parameters
+    ----------
+    row : pandas.Series
+        A row of the DataFrame.
+    embeddings : dict
+        Dictionary mapping row indices to embeddings.
+    result_target_field: str
+        The field in custom_content to output the embeddings to
+
+    Returns
+    -------
+    pandas.Series
+        The updated row
+    """
+    embedding = embeddings.get(row.name, None)
+
+    if embedding is not None:
+        row["metadata"] = glom.assign(row["metadata"], "custom_content." + result_target_field, embedding, missing=dict)
+
+    return row
+
+
 def _format_image_input_string(image_b64: Optional[str]) -> str:
     if not image_b64:
         return
@@ -381,6 +409,20 @@ def _get_pandas_audio_content(row, modality="text"):
     return row.get("audio_metadata", {}).get("audio_transcript")
 
 
+def _get_pandas_custom_content(row, custom_content_field):
+    custom_content = row.get("custom_content", {})
+    content = glom.glom(custom_content, custom_content_field, default=None)
+    if content is None:
+        logger.warning(f"Custom content field: {custom_content_field} not found")
+        return None
+
+    try:
+        return str(content)
+    except (TypeError, ValueError):
+        logger.warning(f"Cannot convert custom content field: {custom_content_field} to string")
+        return None
+
+
 # ------------------------------------------------------------------------------
 # Batch Processing Utilities
 # ------------------------------------------------------------------------------
@@ -519,6 +561,7 @@ def transform_create_text_embeddings_internal(
     api_key = task_config.get("api_key") or transform_config.api_key
     endpoint_url = task_config.get("endpoint_url") or transform_config.embedding_nim_endpoint
     model_name = task_config.get("model_name") or transform_config.embedding_model
+    custom_content_field = task_config.get("custom_content_field") or transform_config.custom_content_field
 
     if execution_trace_log is None:
         execution_trace_log = {}
@@ -612,4 +655,43 @@ def transform_create_text_embeddings_internal(
         content_masks.append(content_mask)
 
     combined_df = _concatenate_extractions_pandas(df_transform_ledger, embedding_dataframes, content_masks)
+
+    # Embed custom content
+    if custom_content_field is not None:
+        result_target_field = task_config.get("result_target_field") or custom_content_field + "_embedding"
+
+        extracted_custom_content = (
+            combined_df["metadata"]
+            .apply(partial(_get_pandas_custom_content, custom_content_field=custom_content_field))
+            .apply(lambda x: x.strip() if isinstance(x, str) and x.strip() else None)
+        )
+
+        valid_custom_content_mask = extracted_custom_content.notna()
+        if valid_custom_content_mask.any():
+            custom_content_list = extracted_custom_content[valid_custom_content_mask].to_list()
+            custom_content_batches = _generate_batches(custom_content_list, batch_size=transform_config.batch_size)
+
+            custom_content_embeddings = _async_runner(
+                custom_content_batches,
+                api_key,
+                endpoint_url,
+                model_name,
+                transform_config.encoding_format,
+                transform_config.input_type,
+                transform_config.truncate,
+                False,
+            )
+            custom_embeddings_dict = dict(
+                zip(
+                    extracted_custom_content.loc[valid_custom_content_mask].index,
+                    custom_content_embeddings.get("embeddings", []),
+                )
+            )
+        else:
+            custom_embeddings_dict = {}
+
+        combined_df = combined_df.apply(
+            _add_custom_embeddings, embeddings=custom_embeddings_dict, result_target_field=result_target_field, axis=1
+        )
+
     return combined_df, {"trace_info": execution_trace_log}

nv_ingest_api/util/dataloader/dataloader.py

@@ -254,22 +254,29 @@ else:
         file = None
         try:
             for file in paths:
+                if thread_stop.is_set():
+                    return
                 if isinstance(file, tuple):
                     video_file, audio_file = file
+                    if thread_stop.is_set():
+                        return
                     with open(video_file, "rb") as f:
                         video = f.read()
+                    if thread_stop.is_set():
+                        return
                     with open(audio_file, "rb") as f:
                         audio = f.read()
                     queue.put((video, audio))
                 else:
-                    if thread_stop:
+                    if thread_stop.is_set():
                         return
                     with open(file, "rb") as f:
                         queue.put(f.read())
         except Exception as e:
             logging.error(f"Error processing file {file}: {e}")
             queue.put(RuntimeError(f"Error processing file {file}: {e}"))
-        queue.put(StopIteration)
+        finally:
+            queue.put(StopIteration)
 
     class DataLoader:
         """
@@ -290,7 +297,7 @@ else:
         ):
             interface = interface if interface else MediaInterface()
             self.thread = None
-            self.thread_stop = False
+            self.thread_stop = threading.Event()
             self.queue = queue.Queue(size)
             self.path = Path(path)
             self.output_dir = output_dir
@@ -323,16 +330,20 @@ else:
             Reset itertor by stopping the thread and clearing the queue.
             """
             if self.thread:
-                self.thread_stop = True
+                self.thread_stop.set()
                 self.thread.join()
-            self.thread_stop = False
-            while self.queue.qsize() != 0:
-                with self.queue.mutex:
-                    self.queue.queue.clear()
+                self.thread = None
+            try:
+                while True:
+                    self.queue.get_nowait()
+            except Exception:
+                pass
+            finally:
+                self.thread_stop.clear()
 
         def __iter__(self):
             self.stop()
-            self.thread_stop = False
+            self.thread_stop.clear()
             self.thread = threading.Thread(
                 target=load_data,
                 args=(

nv_ingest_api/util/message_brokers/qos_scheduler.py

@@ -0,0 +1,283 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import Dict, Optional
+import logging
+import time
+import random
+
+
+class _SchedulingStrategy:
+    """
+    Base scheduling strategy interface. Implementations must provide a non-blocking
+    single-sweep attempt over non-immediate queues and return a job or None.
+    """
+
+    def try_once(self, client, queues: Dict[str, str], order: list[str]) -> Optional[dict]:
+        raise NotImplementedError
+
+
+class _LotteryStrategy(_SchedulingStrategy):
+    """
+    Lottery scheduling with fixed weights.
+    Weights: micro=4, small=2, large=1, medium=1, default=1
+    """
+
+    def __init__(self, prioritize_immediate: bool = True) -> None:
+        self._weights: Dict[str, int] = {
+            "micro": 4,
+            "small": 2,
+            "large": 1,
+            "medium": 1,
+            "default": 1,
+        }
+        self._prioritize_immediate: bool = bool(prioritize_immediate)
+
+    def try_once(self, client, queues: Dict[str, str], order: list[str]) -> Optional[dict]:
+        # Immediate-first if enabled (non-blocking)
+        if self._prioritize_immediate:
+            try:
+                job = client.fetch_message(queues["immediate"], 0)
+                if job is not None:
+                    return job
+            except TimeoutError:
+                pass
+        candidates = list(order)
+        weights = [self._weights[q] for q in candidates]
+        while candidates:
+            try:
+                chosen = random.choices(candidates, weights=weights, k=1)[0]
+                job = client.fetch_message(queues[chosen], 0)
+                if job is not None:
+                    return job
+            except TimeoutError:
+                pass
+            finally:
+                idx = candidates.index(chosen)
+                del candidates[idx]
+                del weights[idx]
+        return None
+
+
+class _SimpleStrategy(_SchedulingStrategy):
+    """
+    Simple strategy placeholder. Actual simple-mode handling is done in QosScheduler.fetch_next
+    to directly fetch from the base 'default' queue using the provided timeout.
+    """
+
+    def try_once(self, client, queues: Dict[str, str], order: list[str]) -> Optional[dict]:
+        # Block up to 30s on the base/default queue and return first available job
+        try:
+            return client.fetch_message(queues["default"], 30.0)
+        except TimeoutError:
+            return None
+
+
+class _RoundRobinStrategy(_SchedulingStrategy):
+    """
+    Simple round-robin over non-immediate queues. Maintains rotation across calls.
+    """
+
+    def __init__(self, order: list[str], prioritize_immediate: bool = True) -> None:
+        self._order = list(order)
+        self._len = len(self._order)
+        self._idx = 0
+        self._prioritize_immediate: bool = bool(prioritize_immediate)
+
+    def try_once(self, client, queues: Dict[str, str], order: list[str]) -> Optional[dict]:
+        # Immediate-first if enabled (non-blocking)
+        if self._prioritize_immediate:
+            try:
+                job = client.fetch_message(queues["immediate"], 0)
+                if job is not None:
+                    return job
+            except TimeoutError:
+                pass
+        start_idx = self._idx
+        for step in range(self._len):
+            i = (start_idx + step) % self._len
+            qname = self._order[i]
+            try:
+                job = client.fetch_message(queues[qname], 0)
+                if job is not None:
+                    # advance rotation to the position after the chosen one
+                    self._idx = (i + 1) % self._len
+                    return job
+            except TimeoutError:
+                continue
+        return None
+
+
+class _WeightedRoundRobinStrategy(_SchedulingStrategy):
+    """
+    Smooth Weighted Round Robin (SWRR) using weights micro=4, small=2, large=1, medium=1, default=1.
+    Maintains current weights across calls.
+    """
+
+    def __init__(self, prioritize_immediate: bool = True) -> None:
+        self._weights: Dict[str, int] = {
+            "micro": 4,
+            "small": 2,
+            "large": 1,
+            "medium": 1,
+            "default": 1,
+        }
+        self._current: Dict[str, int] = {k: 0 for k in self._weights.keys()}
+        self._total: int = sum(self._weights.values())
+        self._prioritize_immediate: bool = bool(prioritize_immediate)
+
+    def try_once(self, client, queues: Dict[str, str], order: list[str]) -> Optional[dict]:
+        # Immediate-first if enabled (non-blocking)
+        if self._prioritize_immediate:
+            try:
+                job = client.fetch_message(queues["immediate"], 0)
+                if job is not None:
+                    return job
+            except TimeoutError:
+                pass
+        # Attempt up to len(order) selections per sweep, excluding queues that prove empty
+        active = list(order)
+        for _ in range(len(order)):
+            if not active:
+                break
+            for q in active:
+                self._current[q] += self._weights[q]
+            chosen = max(active, key=lambda q: self._current[q])
+            self._current[chosen] -= self._total
+            try:
+                job = client.fetch_message(queues[chosen], 0)
+                if job is not None:
+                    return job
+            except TimeoutError:
+                job = None
+            # If no job available from chosen, exclude it for the remainder of this sweep
+            if job is None and chosen in active:
+                active.remove(chosen)
+        # Fallback: single non-blocking attempt for each queue in order
+        for q in order:
+            try:
+                job = client.fetch_message(queues[q], 0)
+                if job is not None:
+                    return job
+            except TimeoutError:
+                continue
+        return None
+
+
+class QosScheduler:
+    """
+    Simplified scheduler that fetches jobs from the default queue only.
+    Uses the provided timeout value when polling the broker.
+    """
+
+    def __init__(
+        self,
+        base_queue: str,
+        total_buffer_capacity: int = 1,
+        num_prefetch_threads: int = 0,
+        prefetch_poll_interval: float = 0.0,
+        prefetch_non_immediate: bool = False,
+        strategy: str = "lottery",
+        prioritize_immediate: bool = True,
+    ) -> None:
+        self.base_queue = base_queue
+
+        # Define all derived queues; default behavior still uses only "default"
+        self.queues: Dict[str, str] = {
+            "default": f"{base_queue}",
+            "immediate": f"{base_queue}_immediate",
+            "micro": f"{base_queue}_micro",
+            "small": f"{base_queue}_small",
+            "medium": f"{base_queue}_medium",
+            "large": f"{base_queue}_large",
+        }
+
+        # Priority order for multi-queue fetching; "immediate" always first
+        self._priority_order = [
+            "immediate",
+            "micro",
+            "small",
+            "medium",
+            "large",
+            "default",
+        ]
+
+        # Non-immediate queue order reference
+        self._non_immediate_order = ["micro", "small", "large", "medium", "default"]
+
+        # Logger
+        self._logger = logging.getLogger(__name__)
+
+        # No prefetching - just direct calls
+        self._total_buffer_capacity: int = int(total_buffer_capacity)
+        self._num_prefetch_threads: int = int(num_prefetch_threads)
+        self._prefetch_poll_interval: float = float(prefetch_poll_interval)
+        self._prefetch_non_immediate: bool = bool(prefetch_non_immediate)
+
+        # Strategy selection
+        self._simple_mode: bool = False
+        if strategy == "simple":
+            self._strategy_impl: _SchedulingStrategy = _SimpleStrategy()
+            self._simple_mode = True
+        elif strategy == "round_robin":
+            self._strategy_impl = _RoundRobinStrategy(self._non_immediate_order, prioritize_immediate)
+        elif strategy == "weighted_round_robin":
+            self._strategy_impl = _WeightedRoundRobinStrategy(prioritize_immediate)
+        else:
+            self._strategy_impl = _LotteryStrategy(prioritize_immediate)
+
+    # Context manager helpers for clean shutdown
+    def __enter__(self) -> "QosScheduler":
+        return self
+
+    def __exit__(self, exc_type, exc, tb) -> None:
+        self.close()
+
+    # ---------------------------- Public API ----------------------------
+    def close(self) -> None:
+        """
+        Cleanly close the scheduler. No-op for the current implementation
+        since we do not spin background threads.
+        """
+        return None
+
+    def fetch_next(self, client, timeout: float = 0.0) -> Optional[dict]:
+        """
+        Immediate-first, then strategy-based scheduling among non-immediate queues.
+
+        Behavior:
+        - Always check 'immediate' first (non-blocking). If present, return immediately.
+        - If not, select using the configured strategy (lottery, round_robin, weighted_round_robin).
+        - If no job is found in a full pass:
+          - If timeout <= 0: return None.
+          - Else: sleep in 0.5s increments and retry until accumulated elapsed time >= timeout.
+        """
+        # Simple mode: delegate to the strategy (blocks up to 30s on base queue)
+        if getattr(self, "_simple_mode", False):
+            return self._strategy_impl.try_once(client, self.queues, self._non_immediate_order)
+
+        start = time.monotonic()
+        while True:
+            # Strategy-based attempt (strategy may include immediate priority internally)
+            job = self._strategy_impl.try_once(client, self.queues, self._non_immediate_order)
+            if job is not None:
+                return job
+
+            # No job found in this sweep
+            if timeout <= 0:
+                return None
+
+            elapsed = time.monotonic() - start
+            if elapsed >= timeout:
+                return None
+
+            # Sleep up to 0.5s, but not beyond remaining timeout
+            remaining = timeout - elapsed
+            sleep_time = 0.5 if remaining > 0.5 else remaining
+            if sleep_time > 0:
+                time.sleep(sleep_time)
+            else:
+                return None

nv_ingest_api/util/message_brokers/simple_message_broker/simple_client.py

@@ -35,6 +35,7 @@ class SimpleClient(MessageBrokerClientBase):
         connection_timeout: int = 300,
         max_pool_size: int = 128,
         use_ssl: bool = False,
+        api_version: str = "v1",
     ):
         """
         Initialize the SimpleClient with configuration parameters.

nv_ingest_api/util/multi_processing/mp_pool_singleton.py

@@ -5,8 +5,9 @@
 
 import logging
 import math
-import multiprocessing as mp
 import os
+import sys
+import multiprocessing as mp
 from threading import Lock
 from typing import Any, Callable, Optional
 
@@ -103,7 +104,12 @@ class ProcessWorkerPoolSingleton:
             The total number of worker processes to start.
         """
         self._total_workers = total_max_workers
-        self._context: mp.context.ForkContext = mp.get_context("fork")
+
+        start_method = "fork"
+        if sys.platform.lower() == "darwin":
+            start_method = "spawn"
+        self._context: mp.context.ForkContext = mp.get_context(start_method)
+
         # Bounded task queue: maximum tasks queued = 2 * total_max_workers.
         self._task_queue: mp.Queue = self._context.Queue(maxsize=2 * total_max_workers)
         self._next_task_id: int = 0

nv_ingest_api/util/service_clients/redis/redis_client.py

@@ -650,6 +650,22 @@ class RedisClient(MessageBrokerClientBase):
             except Exception as e:
                 logger.exception(f"{log_prefix}: Cache read error: {e}. Trying Redis.")
 
+        # If caller requests non-blocking behavior (timeout <= 0), attempt immediate pop.
+        if timeout is not None and timeout <= 0:
+            try:
+                client = self.get_client()
+                popped = client.lpop(channel_name)
+                if popped is None:
+                    return None
+                try:
+                    return json.loads(popped)
+                except json.JSONDecodeError as e:
+                    logger.error(f"Failed to decode JSON from non-blocking LPOP on '{channel_name}': {e}")
+                    return None
+            except Exception as e:
+                logger.warning(f"Non-blocking LPOP failed for '{channel_name}': {e}")
+                return None
+
         while True:
             try:
                 fetch_result: Union[Dict[str, Any], List[Dict[str, Any]]]
@@ -711,6 +727,150 @@ class RedisClient(MessageBrokerClientBase):
                 logger.exception(f"{log_prefix}: Unexpected error during fetch: {e}")
                 raise ValueError(f"Unexpected error during fetch: {e}") from e
 
+    def fetch_message_from_any(self, channel_names: List[str], timeout: float = 0) -> Optional[Dict[str, Any]]:
+        """
+        Attempt to fetch a message from the first non-empty list among the provided channel names
+        using Redis BLPOP. If the popped item represents a fragmented message, this method will
+        continue popping from the same channel to reconstruct the full message.
+
+        Parameters
+        ----------
+        channel_names : List[str]
+            Ordered list of Redis list keys to attempt in priority order.
+        timeout : float, optional
+            Timeout in seconds to wait for any item across the provided lists. Redis supports
+            integer-second timeouts; sub-second values will be truncated.
+
+        Returns
+        -------
+        dict or None
+            The reconstructed message dictionary if an item was fetched; otherwise None on timeout.
+        """
+        if not channel_names:
+            return None
+
+        client = self.get_client()
+        blpop_timeout = int(max(0, timeout))
+        try:
+            res = client.blpop(channel_names, timeout=blpop_timeout)
+        except (redis.RedisError, ConnectionError) as e:
+            logger.debug(f"BLPOP error on {channel_names}: {e}")
+            return None
+
+        if res is None:
+            return None
+
+        list_key, first_bytes = res
+        if isinstance(list_key, bytes):
+            try:
+                list_key = list_key.decode("utf-8")
+            except Exception:
+                list_key = str(list_key)
+        # Decode first element
+        try:
+            first_msg = json.loads(first_bytes)
+        except json.JSONDecodeError as e:
+            logger.error(f"Failed to decode JSON popped from '{list_key}': {e}")
+            return None
+
+        expected_count: int = int(first_msg.get("fragment_count", 1))
+        if expected_count <= 1:
+            return first_msg
+
+        # Collect remaining fragments from the same list key
+        fragments: List[Dict[str, Any]] = [first_msg]
+        accumulated = 0.0
+        start_time = time.monotonic()
+        for i in range(1, expected_count):
+            remaining = max(0, timeout - accumulated)
+            per_frag_timeout = int(max(1, remaining)) if timeout else 1
+            try:
+                frag_res = client.blpop([list_key], timeout=per_frag_timeout)
+            except (redis.RedisError, ConnectionError) as e:
+                logger.error(f"BLPOP error while collecting fragments from '{list_key}': {e}")
+                return None
+            if frag_res is None:
+                logger.error(f"Timeout while collecting fragment {i}/{expected_count-1} from '{list_key}'")
+                return None
+            _, frag_key_bytes_or_val = frag_res
+            # Redis returns (key, value); we don't need the key here
+            frag_bytes = frag_key_bytes_or_val
+            try:
+                frag_msg = json.loads(frag_bytes)
+                fragments.append(frag_msg)
+            except json.JSONDecodeError as e:
+                logger.error(f"Failed to decode fragment JSON from '{list_key}': {e}")
+                return None
+            accumulated = time.monotonic() - start_time
+
+        # Combine and return
+        try:
+            return self._combine_fragments(fragments)
+        except Exception as e:
+            logger.error(f"Error combining fragments from '{list_key}': {e}")
+            return None
+
+    def fetch_message_from_any_with_key(
+        self, channel_names: List[str], timeout: float = 0
+    ) -> Optional[Tuple[str, Dict[str, Any]]]:
+        """
+        Like fetch_message_from_any(), but returns the Redis list key together with the message.
+        This is useful for higher-level schedulers that need to apply per-category quotas.
+        """
+        if not channel_names:
+            return None
+
+        client = self.get_client()
+        blpop_timeout = int(max(0, timeout))
+        try:
+            res = client.blpop(channel_names, timeout=blpop_timeout)
+        except (redis.RedisError, ConnectionError) as e:
+            logger.debug(f"BLPOP error on {channel_names}: {e}")
+            return None
+
+        if res is None:
+            return None
+
+        list_key, first_bytes = res
+        try:
+            first_msg = json.loads(first_bytes)
+        except json.JSONDecodeError as e:
+            logger.error(f"Failed to decode JSON popped from '{list_key}': {e}")
+            return None
+
+        expected_count: int = int(first_msg.get("fragment_count", 1))
+        if expected_count <= 1:
+            return list_key, first_msg
+
+        fragments: List[Dict[str, Any]] = [first_msg]
+        accumulated = 0.0
+        start_time = time.monotonic()
+        for i in range(1, expected_count):
+            remaining = max(0, timeout - accumulated)
+            per_frag_timeout = int(max(1, remaining)) if timeout else 1
+            try:
+                frag_res = client.blpop([list_key], timeout=per_frag_timeout)
+            except (redis.RedisError, ConnectionError) as e:
+                logger.error(f"BLPOP error while collecting fragments from '{list_key}': {e}")
+                return None
+            if frag_res is None:
+                logger.error(f"Timeout while collecting fragment {i}/{expected_count-1} from '{list_key}'")
+                return None
+            _, frag_bytes = frag_res
+            try:
+                frag_msg = json.loads(frag_bytes)
+                fragments.append(frag_msg)
+            except json.JSONDecodeError as e:
+                logger.error(f"Failed to decode fragment JSON from '{list_key}': {e}")
+                return None
+            accumulated = time.monotonic() - start_time
+
+        try:
+            return list_key, self._combine_fragments(fragments)
+        except Exception as e:
+            logger.error(f"Error combining fragments from '{list_key}': {e}")
+            return None
+
     @staticmethod
     def _combine_fragments(fragments: List[Dict[str, Any]]) -> Dict[str, Any]:
         """

{nv_ingest_api-2025.10.22.dev20251022.dist-info → nv_ingest_api-2025.11.2.dev20251102.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nv-ingest-api
-Version: 2025.10.22.dev20251022
+Version: 2025.11.2.dev20251102
 Summary: Python module with core document ingestion functions.
 Author-email: Jeremy Dyer <jdyer@nvidia.com>
 License:                                  Apache License
@@ -222,6 +222,7 @@ Requires-Dist: fsspec>=2025.5.1
 Requires-Dist: universal_pathlib>=0.2.6
 Requires-Dist: ffmpeg-python==0.2.0
 Requires-Dist: tritonclient
+Requires-Dist: glom
 Dynamic: license-file
 
 # nv-ingest-api