nv-ingest-api 2025.10.30.dev20251030__tar.gz → 2025.10.31.dev20251031__tar.gz

{nv_ingest_api-2025.10.30.dev20251030/src/nv_ingest_api.egg-info → nv_ingest_api-2025.10.31.dev20251031}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nv-ingest-api
-Version: 2025.10.30.dev20251030
+Version: 2025.10.31.dev20251031
 Summary: Python module with core document ingestion functions.
 Author-email: Jeremy Dyer <jdyer@nvidia.com>
 License:                                  Apache License

{nv_ingest_api-2025.10.30.dev20251030 → nv_ingest_api-2025.10.31.dev20251031}/src/nv_ingest_api/internal/schemas/meta/ingest_job_schema.py

@@ -43,6 +43,24 @@ class PdfConfigSchema(BaseModelNoExt):
     split_page_count: Annotated[int, Field(ge=1)] = 32
 
 
+class RoutingOptionsSchema(BaseModelNoExt):
+    # Queue routing hint for QoS scheduler
+    queue_hint: Optional[str] = None
+
+    @field_validator("queue_hint")
+    @classmethod
+    def validate_queue_hint(cls, v):
+        if v is None:
+            return v
+        if not isinstance(v, str):
+            raise ValueError("queue_hint must be a string")
+        s = v.lower()
+        allowed = {"default", "immediate", "micro", "small", "medium", "large"}
+        if s not in allowed:
+            raise ValueError("queue_hint must be one of: default, immediate, micro, small, medium, large")
+        return s
+
+
 # Ingest Task Schemas
 
 
@@ -283,8 +301,27 @@ class IngestJobSchema(BaseModelNoExt):
     job_id: Union[str, int]
     tasks: List[IngestTaskSchema]
     tracing_options: Optional[TracingOptionsSchema] = None
+    routing_options: Optional[RoutingOptionsSchema] = None
     pdf_config: Optional[PdfConfigSchema] = None
 
+    @model_validator(mode="before")
+    @classmethod
+    def migrate_queue_hint(cls, values):
+        """
+        Backward-compatibility shim: if a legacy client sends
+        tracing_options.queue_hint, move it into routing_options.queue_hint.
+        """
+        try:
+            topt = values.get("tracing_options") or {}
+            ropt = values.get("routing_options") or {}
+            if isinstance(topt, dict) and "queue_hint" in topt and "queue_hint" not in ropt:
+                ropt["queue_hint"] = topt.pop("queue_hint")
+                values["routing_options"] = ropt
+                values["tracing_options"] = topt
+        except Exception:
+            pass
+        return values
+
 
 # ------------------------------------------------------------------------------
 # Utility Functions

nv_ingest_api-2025.10.31.dev20251031/src/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-2025.10.30.dev20251030 → nv_ingest_api-2025.10.31.dev20251031}/src/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-2025.10.30.dev20251030 → nv_ingest_api-2025.10.31.dev20251031}/src/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.30.dev20251030 → nv_ingest_api-2025.10.31.dev20251031/src/nv_ingest_api.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nv-ingest-api
-Version: 2025.10.30.dev20251030
+Version: 2025.10.31.dev20251031
 Summary: Python module with core document ingestion functions.
 Author-email: Jeremy Dyer <jdyer@nvidia.com>
 License:                                  Apache License

{nv_ingest_api-2025.10.30.dev20251030 → nv_ingest_api-2025.10.31.dev20251031}/src/nv_ingest_api.egg-info/SOURCES.txt

@@ -150,6 +150,7 @@ src/nv_ingest_api/util/logging/__init__.py
 src/nv_ingest_api/util/logging/configuration.py
 src/nv_ingest_api/util/logging/sanitize.py
 src/nv_ingest_api/util/message_brokers/__init__.py
+src/nv_ingest_api/util/message_brokers/qos_scheduler.py
 src/nv_ingest_api/util/message_brokers/simple_message_broker/__init__.py
 src/nv_ingest_api/util/message_brokers/simple_message_broker/broker.py
 src/nv_ingest_api/util/message_brokers/simple_message_broker/ordered_message_queue.py