nv-ingest-api 2025.10.4.dev20251004__py3-none-any.whl → 2025.11.2.dev20251102__py3-none-any.whl

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/util/service_clients/rest/rest_client.py

@@ -103,6 +103,17 @@ class RestClient(MessageBrokerClientBase):
             Default timeout in seconds for waiting for data after connection. Default is None.
         http_allocator : Optional[Callable[[], Any]], optional
             A callable that returns an HTTP client instance. If None, `requests.Session()` is used.
+        **kwargs : dict
+            Additional keyword arguments. Supported keys:
+            - api_version : str, optional
+                API version to use ('v1' or 'v2'). Defaults to 'v1' if not specified.
+                Invalid versions will log a warning and fall back to 'v1'.
+            - base_url : str, optional
+                Override the generated base URL.
+            - headers : dict, optional
+                Additional headers to include in requests.
+            - auth : optional
+                Authentication configuration for requests.
 
         Returns
         -------
@@ -137,13 +148,30 @@ class RestClient(MessageBrokerClientBase):
                 )
                 self._client = requests.Session()
 
-        self._submit_endpoint: str = "/v1/submit_job"
-        self._fetch_endpoint: str = "/v1/fetch_job"
+        # Validate and normalize API version to prevent misconfiguration
+        # Default to v1 for backwards compatibility if not explicitly provided
+        VALID_API_VERSIONS = {"v1", "v2"}
+        raw_api_version = kwargs.get("api_version", "v1")
+        api_version = str(raw_api_version).strip().lower()
+
+        if api_version not in VALID_API_VERSIONS:
+            logger.warning(
+                f"Invalid API version '{raw_api_version}' specified. "
+                f"Valid versions are: {VALID_API_VERSIONS}. Falling back to 'v1'."
+            )
+            api_version = "v1"
+
+        self._api_version = api_version
+        self._submit_endpoint: str = f"/{api_version}/submit_job"
+        self._fetch_endpoint: str = f"/{api_version}/fetch_job"
         self._base_url: str = kwargs.get("base_url") or self._generate_url(self._host, self._port)
         self._headers = kwargs.get("headers", {})
         self._auth = kwargs.get("auth", None)
 
         logger.debug(f"RestClient base URL set to: {self._base_url}")
+        logger.info(
+            f"RestClient using API version: {api_version} (endpoints: {self._submit_endpoint}, {self._fetch_endpoint})"
+        )
 
     @staticmethod
     def _generate_url(host: str, port: int) -> str:
@@ -308,7 +336,18 @@ class RestClient(MessageBrokerClientBase):
 
         retries: int = 0
         url: str = f"{self._base_url}{self._fetch_endpoint}/{job_id}"
-        req_timeout: Tuple[float, Optional[float]] = self._timeout
+        # Derive per-call timeout if provided; otherwise use default
+        if timeout is None:
+            req_timeout: Tuple[float, Optional[float]] = self._timeout
+        else:
+            if isinstance(timeout, tuple):
+                # Expect (connect, read)
+                connect_t = float(timeout[0])
+                read_t = None if (len(timeout) < 2 or timeout[1] is None) else float(timeout[1])
+                req_timeout = (connect_t, read_t)
+            else:
+                # Single float means override read timeout, keep a small connect timeout
+                req_timeout = (min(self._default_connect_timeout, 5.0), float(timeout))
 
         while True:
             result: Optional[Any] = None

nv_ingest_api/util/string_processing/yaml.py

@@ -1,5 +1,6 @@
 import os
 import re
+from typing import Optional
 
 # This regex finds all forms of environment variables:
 # $VAR, ${VAR}, $VAR|default, and ${VAR|default}
@@ -20,12 +21,46 @@ def _replacer(match: re.Match) -> str:
     var_name = match.group("braced") or match.group("named")
     default_val = match.group("braced_default") or match.group("named_default")
 
-    # Get value from environment, or use default.
-    value = os.environ.get(var_name, default_val)
+    # First try the primary env var
+    value = os.environ.get(var_name)
+    if value is not None:
+        return value
 
-    if value is None:
+    # If primary is missing, try the default.
+    resolved_default = _resolve_default_with_single_fallback(default_val)
+
+    if resolved_default is None:
         return ""
-    return value
+
+    return resolved_default
+
+
+def _is_var_ref(token: str) -> Optional[str]:
+    """If token is a $VAR or ${VAR} reference, return VAR name; else None."""
+    if not token:
+        return None
+    if token.startswith("${") and token.endswith("}"):
+        inner = token[2:-1]
+        return inner if re.fullmatch(r"\w+", inner) else None
+    if token.startswith("$"):
+        inner = token[1:]
+        return inner if re.fullmatch(r"\w+", inner) else None
+    return None
+
+
+def _resolve_default_with_single_fallback(default_val: Optional[str]) -> Optional[str]:
+    """
+    Support a single-level fallback where the default itself can be another env var.
+    For example, in $A|$B or ${A|$B}, we try B if A missing.
+    """
+    if default_val is None:
+        return None
+
+    var = _is_var_ref(default_val)
+    if var is not None:
+        return os.environ.get(var, None)
+
+    return default_val
 
 
 def substitute_env_vars_in_yaml_content(raw_content: str) -> str:
@@ -35,6 +70,8 @@ def substitute_env_vars_in_yaml_content(raw_content: str) -> str:
     This function finds all occurrences of environment variable placeholders
     ($VAR, ${VAR}, $VAR|default, ${VAR|default}) in the input string
     and replaces them with their corresponding environment variable values.
+    Also supports a single fallback to another env var: $VAR|$OTHER, ${VAR|$OTHER}
+    Quoted defaults are preserved EXACTLY as written (e.g., 'a,b' keeps quotes).
 
     Args:
         raw_content: The raw string content of a YAML file.

{nv_ingest_api-2025.10.4.dev20251004.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.4.dev20251004
+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