nv-ingest-api 26.1.0rc4__py3-none-any.whl

nv_ingest_api/internal/primitives/control_message_task.py

@@ -0,0 +1,16 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from uuid import UUID
+
+from pydantic import BaseModel, Field, ConfigDict
+from typing import Any, Dict, Union
+
+
+class ControlMessageTask(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    type: str
+    id: Union[str, UUID]
+    properties: Dict[str, Any] = Field(default_factory=dict)

nv_ingest_api/internal/primitives/ingest_control_message.py

@@ -0,0 +1,307 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import copy
+import re
+from datetime import datetime
+from collections import defaultdict
+from typing import Any, Dict, Generator, List, Optional, Union
+
+import logging
+import pandas as pd
+
+from nv_ingest_api.internal.primitives.control_message_task import ControlMessageTask
+
+logger = logging.getLogger(__name__)
+
+
+def remove_task_by_type(ctrl_msg, task: str):
+    """
+    Remove a task from the control message by matching its type.
+
+    This function iterates over the tasks in the control message, and if it finds a task
+    whose type matches the provided task string, it removes that task (using its unique id)
+    and returns the task's properties.
+
+    Parameters
+    ----------
+    ctrl_msg : IngestControlMessage
+        The control message from which to remove the task.
+    task : str
+        The task type to remove.
+
+    Returns
+    -------
+    dict
+        The properties of the removed task.
+
+    Raises
+    ------
+    ValueError
+        If no task with the given type is found.
+    """
+    task_obj = None
+    for t in ctrl_msg.get_tasks():
+        if t.type == task:
+            task_obj = t
+            break
+
+    if task_obj is None:
+        err_msg = f"process_control_message: Task '{task}' not found in control message."
+        logger.error(err_msg)
+        raise ValueError(err_msg)
+
+    removed_task = ctrl_msg.remove_task(task_obj.id)
+    return removed_task.properties
+
+
+def remove_all_tasks_by_type(ctrl_msg, task: str):
+    """
+    Remove all tasks from the control message by matching their type.
+
+    This function iterates over the tasks in the control message, finds all tasks
+    whose type matches the provided task string, removes them, and returns their
+    properties as a list.
+
+    Parameters
+    ----------
+    ctrl_msg : IngestControlMessage
+        The control message from which to remove the tasks.
+    task : str
+        The task type to remove.
+
+    Returns
+    -------
+    list[dict]
+        A list of dictionaries of properties for all removed tasks.
+
+    Raises
+    ------
+    ValueError
+        If no tasks with the given type are found.
+    """
+    matching_tasks = []
+
+    # Find all tasks with matching type
+    for t in ctrl_msg.get_tasks():
+        if t.type == task:
+            matching_tasks.append(t)
+
+    if not matching_tasks:
+        err_msg = f"process_control_message: No tasks of type '{task}' found in control message."
+        logger.error(err_msg)
+        raise ValueError(err_msg)
+
+    # Remove all matching tasks and collect their properties
+    removed_task_properties = []
+    for task_obj in matching_tasks:
+        removed_task = ctrl_msg.remove_task(task_obj.id)
+        removed_task_properties.append(removed_task.properties)
+
+    return removed_task_properties
+
+
+class IngestControlMessage:
+    """
+    A control message class for ingesting tasks and managing associated metadata,
+    timestamps, configuration, and payload.
+    """
+
+    def __init__(self):
+        """
+        Initialize a new IngestControlMessage instance.
+        """
+        self._tasks: Dict[str, List[ControlMessageTask]] = defaultdict(list)
+        self._metadata: Dict[str, Any] = {}
+        self._timestamps: Dict[str, datetime] = {}
+        self._payload: Optional[pd.DataFrame] = None
+        self._config: Dict[str, Any] = {}
+
+    def add_task(self, task: ControlMessageTask):
+        """
+        Add a task to the control message. Multiple tasks with the same ID are supported.
+        """
+        self._tasks[task.id].append(task)
+
+    def get_tasks(self) -> Generator[ControlMessageTask, None, None]:
+        """
+        Return all tasks as a generator.
+        """
+        for task_list in self._tasks.values():
+            yield from task_list
+
+    def has_task(self, task_id: str) -> bool:
+        """
+        Check if any tasks with the given ID exist.
+        """
+        return task_id in self._tasks and len(self._tasks[task_id]) > 0
+
+    def remove_task(self, task_id: str) -> ControlMessageTask:
+        """
+        Remove the first task with the given ID. Warns if no task exists.
+        """
+        if task_id in self._tasks and self._tasks[task_id]:
+            task = self._tasks[task_id].pop(0)
+            # Clean up empty lists
+            if not self._tasks[task_id]:
+                del self._tasks[task_id]
+            return task
+        else:
+            raise RuntimeError(f"Attempted to remove non-existent task with id: {task_id}")
+
+    def config(self, config: Dict[str, Any] = None) -> Dict[str, Any]:
+        """
+        Get or update the control message configuration.
+
+        If 'config' is provided, it must be a dictionary. The configuration is updated with the
+        provided values. If no argument is provided, returns a copy of the current configuration.
+
+        Raises
+        ------
+        ValueError
+            If the provided configuration is not a dictionary.
+        """
+        if config is None:
+            return self._config.copy()
+
+        if not isinstance(config, dict):
+            raise ValueError("Configuration must be provided as a dictionary.")
+
+        self._config.update(config)
+        return self._config.copy()
+
+    def copy(self) -> "IngestControlMessage":
+        """
+        Create a deep copy of this control message.
+        """
+        return copy.deepcopy(self)
+
+    def get_metadata(self, key: Union[str, re.Pattern] = None, default_value: Any = None) -> Any:
+        """
+        Retrieve metadata. If 'key' is None, returns a copy of all metadata.
+
+        Parameters
+        ----------
+        key : str or re.Pattern, optional
+            If a string is provided, returns the value for that exact key.
+            If a regex pattern is provided, returns a dictionary of all metadata key-value pairs
+            where the key matches the regex. If no matches are found, returns default_value.
+        default_value : Any, optional
+            The value to return if the key is not found or no regex matches.
+
+        Returns
+        -------
+        Any
+            The metadata value for an exact string key, or a dict of matching metadata if a regex is provided.
+        """
+        if key is None:
+            return self._metadata.copy()
+
+        # If key is a regex pattern (i.e. has a search method), perform pattern matching.
+        if hasattr(key, "search"):
+            matches = {k: v for k, v in self._metadata.items() if key.search(k)}
+            return matches if matches else default_value
+
+        # Otherwise, perform an exact lookup.
+        return self._metadata.get(key, default_value)
+
+    def has_metadata(self, key: Union[str, re.Pattern]) -> bool:
+        """
+        Check if a metadata key exists.
+
+        Parameters
+        ----------
+        key : str or re.Pattern
+            If a string is provided, checks for the exact key.
+            If a regex pattern is provided, returns True if any metadata key matches the regex.
+
+        Returns
+        -------
+        bool
+            True if the key (or any matching key, in case of a regex) exists, False otherwise.
+        """
+        if hasattr(key, "search"):
+            return any(key.search(k) for k in self._metadata)
+        return key in self._metadata
+
+    def list_metadata(self) -> list:
+        """
+        List all metadata keys.
+        """
+        return list(self._metadata.keys())
+
+    def set_metadata(self, key: str, value: Any) -> None:
+        """
+        Set a metadata key-value pair.
+        """
+        self._metadata[key] = value
+
+    def filter_timestamp(self, regex_filter: str) -> Dict[str, datetime]:
+        """
+        Retrieve timestamps whose keys match the regex filter.
+        """
+        pattern = re.compile(regex_filter)
+        timestamps_snapshot = self._timestamps.copy()
+        return {key: ts for key, ts in timestamps_snapshot.items() if pattern.search(key)}
+
+    def get_timestamp(self, key: str, fail_if_nonexist: bool = False) -> datetime:
+        """
+        Retrieve a timestamp for a given key.
+
+        Raises
+        ------
+        KeyError
+            If the key is not found and 'fail_if_nonexist' is True.
+        """
+        if key in self._timestamps:
+            return self._timestamps[key]
+        if fail_if_nonexist:
+            raise KeyError(f"Timestamp for key '{key}' does not exist.")
+        return None
+
+    def get_timestamps(self) -> Dict[str, datetime]:
+        """
+        Retrieve all timestamps.
+        """
+        return self._timestamps.copy()
+
+    def set_timestamp(self, key: str, timestamp: Any) -> None:
+        """
+        Set a timestamp for a given key. Accepts either a datetime object or an ISO format string.
+
+        Raises
+        ------
+        ValueError
+            If the provided timestamp is neither a datetime object nor a valid ISO format string.
+        """
+        if isinstance(timestamp, datetime):
+            self._timestamps[key] = timestamp
+
+        elif isinstance(timestamp, str):
+            try:
+                dt = datetime.fromisoformat(timestamp)
+                self._timestamps[key] = dt
+            except ValueError as e:
+                raise ValueError(f"Invalid timestamp format: {timestamp}") from e
+
+        else:
+            raise ValueError("timestamp must be a datetime object or ISO format string")
+
+    def payload(self, payload: pd.DataFrame = None) -> pd.DataFrame:
+        """
+        Get or set the payload DataFrame.
+
+        Raises
+        ------
+        ValueError
+            If the provided payload is not a pandas DataFrame.
+        """
+        if payload is None:
+            return self._payload
+
+        if not isinstance(payload, pd.DataFrame):
+            raise ValueError("Payload must be a pandas DataFrame")
+
+        self._payload = payload
+        return self._payload

nv_ingest_api/internal/primitives/nim/__init__.py

@@ -0,0 +1,9 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from .nim_client import NimClient
+from .nim_client import get_nim_client_manager
+from .nim_model_interface import ModelInterface
+
+__all__ = ["NimClient", "ModelInterface", "get_nim_client_manager"]

nv_ingest_api/internal/primitives/nim/default_values.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# Copyright (c) 2024, NVIDIA CORPORATION.
+
+
+YOLOX_MAX_BATCH_SIZE = 8
+YOLOX_MAX_WIDTH = 1536
+YOLOX_MAX_HEIGHT = 1536
+YOLOX_CONF_THRESHOLD = 0.01
+YOLOX_IOU_THRESHOLD = 0.5
+YOLOX_MIN_SCORE = 0.1
+YOLOX_FINAL_SCORE = 0.48

nv_ingest_api/internal/primitives/nim/model_interface/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0

nv_ingest_api/internal/primitives/nim/model_interface/cached.py

@@ -0,0 +1,274 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+
+import base64
+import io
+import logging
+import PIL.Image as Image
+from typing import Any, Dict, Optional, List
+
+import numpy as np
+
+from nv_ingest_api.internal.primitives.nim import ModelInterface
+from nv_ingest_api.util.image_processing.transforms import base64_to_numpy
+
+logger = logging.getLogger(__name__)
+
+
+class CachedModelInterface(ModelInterface):
+    """
+    An interface for handling inference with a Cached model, supporting both gRPC and HTTP
+    protocols, including batched input.
+    """
+
+    def name(self) -> str:
+        """
+        Get the name of the model interface.
+
+        Returns
+        -------
+        str
+            The name of the model interface ("Cached").
+        """
+        return "Cached"
+
+    def prepare_data_for_inference(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Decode base64-encoded images into NumPy arrays, storing them in `data["image_arrays"]`.
+
+        Parameters
+        ----------
+        data : dict of str -> Any
+            The input data containing either:
+             - "base64_image": a single base64-encoded image, or
+             - "base64_images": a list of base64-encoded images.
+
+        Returns
+        -------
+        dict of str -> Any
+            The updated data dictionary with decoded image arrays stored in
+            "image_arrays", where each array has shape (H, W, C).
+
+        Raises
+        ------
+        KeyError
+            If neither 'base64_image' nor 'base64_images' is provided.
+        ValueError
+            If 'base64_images' is provided but is not a list.
+        """
+        if "base64_images" in data:
+            base64_list = data["base64_images"]
+            if not isinstance(base64_list, list):
+                raise ValueError("The 'base64_images' key must contain a list of base64-encoded strings.")
+            data["image_arrays"] = [base64_to_numpy(img) for img in base64_list]
+
+        elif "base64_image" in data:
+            # Fallback to single image case; wrap it in a list to keep the interface consistent
+            data["image_arrays"] = [base64_to_numpy(data["base64_image"])]
+
+        else:
+            raise KeyError("Input data must include 'base64_image' or 'base64_images' with base64-encoded images.")
+
+        return data
+
+    def format_input(self, data: Dict[str, Any], protocol: str, max_batch_size: int, **kwargs) -> Any:
+        """
+        Format input data for the specified protocol ("grpc" or "http"), handling batched images.
+        Additionally, returns batched data that coalesces the original image arrays and their dimensions
+        in the same order as provided.
+
+        Parameters
+        ----------
+        data : dict of str -> Any
+            The input data dictionary, expected to contain "image_arrays" (a list of np.ndarray).
+        protocol : str
+            The protocol to use, "grpc" or "http".
+        max_batch_size : int
+            The maximum number of images per batch.
+
+        Returns
+        -------
+        tuple
+            A tuple (formatted_batches, formatted_batch_data) where:
+              - For gRPC: formatted_batches is a list of NumPy arrays, each of shape (B, H, W, C)
+                with B <= max_batch_size.
+              - For HTTP: formatted_batches is a list of JSON-serializable dict payloads.
+              - In both cases, formatted_batch_data is a list of dicts with the keys:
+                    "image_arrays": the list of original np.ndarray images for that batch, and
+                    "image_dims":  a list of (height, width) tuples for each image in the batch.
+
+        Raises
+        ------
+        KeyError
+            If "image_arrays" is missing in the data dictionary.
+        ValueError
+            If the protocol is invalid, or if no valid images are found.
+        """
+        if "image_arrays" not in data:
+            raise KeyError("Expected 'image_arrays' in data. Make sure prepare_data_for_inference was called.")
+
+        image_arrays = data["image_arrays"]
+        # Compute dimensions for each image.
+        image_dims = [(img.shape[0], img.shape[1]) for img in image_arrays]
+
+        # Helper: chunk a list into sublists of length up to chunk_size.
+        def chunk_list(lst: list, chunk_size: int) -> List[list]:
+            return [lst[i : i + chunk_size] for i in range(0, len(lst), chunk_size)]
+
+        if protocol == "grpc":
+            logger.debug("Formatting input for gRPC Cached model (batched).")
+            batched_images = []
+            for arr in image_arrays:
+                # Expand from (H, W, C) to (1, H, W, C) if needed
+                if arr.ndim == 3:
+                    arr = np.expand_dims(arr, axis=0)
+                batched_images.append(arr.astype(np.float32))
+
+            if not batched_images:
+                raise ValueError("No valid images found for gRPC formatting.")
+
+            # Chunk the processed images, original arrays, and dimensions.
+            batched_image_chunks = chunk_list(batched_images, max_batch_size)
+            orig_chunks = chunk_list(image_arrays, max_batch_size)
+            dims_chunks = chunk_list(image_dims, max_batch_size)
+
+            batched_inputs = []
+            formatted_batch_data = []
+            for proc_chunk, orig_chunk, dims_chunk in zip(batched_image_chunks, orig_chunks, dims_chunks):
+                # Concatenate along the batch dimension => shape (B, H, W, C)
+                batched_input = np.concatenate(proc_chunk, axis=0)
+                batched_inputs.append(batched_input)
+                formatted_batch_data.append({"image_arrays": orig_chunk, "image_dims": dims_chunk})
+            return batched_inputs, formatted_batch_data
+
+        elif protocol == "http":
+            logger.debug("Formatting input for HTTP Cached model (batched).")
+            content_list: List[Dict[str, Any]] = []
+            for arr in image_arrays:
+                # Convert to uint8 if needed, then to PIL Image and base64-encode it.
+                if arr.dtype != np.uint8:
+                    arr = (arr * 255).astype(np.uint8)
+                image_pil = Image.fromarray(arr)
+                buffered = io.BytesIO()
+                image_pil.save(buffered, format="PNG")
+                base64_img = base64.b64encode(buffered.getvalue()).decode("utf-8")
+                image_item = {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{base64_img}"}}
+                content_list.append(image_item)
+
+            # Chunk the content list, original arrays, and dimensions.
+            content_chunks = chunk_list(content_list, max_batch_size)
+            orig_chunks = chunk_list(image_arrays, max_batch_size)
+            dims_chunks = chunk_list(image_dims, max_batch_size)
+
+            payload_batches = []
+            formatted_batch_data = []
+            for chunk, orig_chunk, dims_chunk in zip(content_chunks, orig_chunks, dims_chunks):
+                message = {"content": chunk}
+                payload = {"messages": [message]}
+                payload_batches.append(payload)
+                formatted_batch_data.append({"image_arrays": orig_chunk, "image_dims": dims_chunk})
+            return payload_batches, formatted_batch_data
+
+        else:
+            raise ValueError("Invalid protocol specified. Must be 'grpc' or 'http'.")
+
+    def parse_output(self, response: Any, protocol: str, data: Optional[Dict[str, Any]] = None, **kwargs: Any) -> Any:
+        """
+        Parse the output from the Cached model's inference response.
+
+        Parameters
+        ----------
+        response : Any
+            The raw response from the model inference.
+        protocol : str
+            The protocol used ("grpc" or "http").
+        data : dict of str -> Any, optional
+            Additional input data (unused here, but available for consistency).
+        **kwargs : Any
+            Additional keyword arguments for future compatibility.
+
+        Returns
+        -------
+        Any
+            The parsed output data (e.g., list of strings), depending on the protocol.
+
+        Raises
+        ------
+        ValueError
+            If the protocol is invalid.
+        RuntimeError
+            If the HTTP response is not as expected (missing 'data' key).
+        """
+        if protocol == "grpc":
+            logger.debug("Parsing output from gRPC Cached model (batched).")
+            parsed: List[str] = []
+            # Assume `response` is iterable, each element a list/array of byte strings
+            for single_output in response:
+                joined_str = " ".join(o.decode("utf-8") for o in single_output)
+                parsed.append(joined_str)
+            return parsed
+
+        elif protocol == "http":
+            logger.debug("Parsing output from HTTP Cached model (batched).")
+            if not isinstance(response, dict):
+                raise RuntimeError("Expected JSON/dict response for HTTP, got something else.")
+            if "data" not in response or not response["data"]:
+                raise RuntimeError("Unexpected response format: 'data' key missing or empty.")
+
+            contents: List[str] = []
+            for item in response["data"]:
+                # Each "item" might have a "content" key
+                content = item.get("content", "")
+                contents.append(content)
+
+            return contents
+
+        else:
+            raise ValueError("Invalid protocol specified. Must be 'grpc' or 'http'.")
+
+    def process_inference_results(self, output: Any, protocol: str, **kwargs: Any) -> Any:
+        """
+        Process inference results for the Cached model.
+
+        Parameters
+        ----------
+        output : Any
+            The raw output from the model.
+        protocol : str
+            The inference protocol used ("grpc" or "http").
+        **kwargs : Any
+            Additional parameters for post-processing (not used here).
+
+        Returns
+        -------
+        Any
+            The processed inference results, which here is simply returned as-is.
+        """
+        # For Cached model, we simply return what we parsed (e.g., a list of strings or a single string)
+        return output
+
+    def _extract_content_from_nim_response(self, json_response: Dict[str, Any]) -> Any:
+        """
+        Extract content from the JSON response of a NIM (HTTP) API request.
+
+        Parameters
+        ----------
+        json_response : dict of str -> Any
+            The JSON response from the NIM API.
+
+        Returns
+        -------
+        Any
+            The extracted content from the response.
+
+        Raises
+        ------
+        RuntimeError
+            If the response format is unexpected (missing 'data' or empty).
+        """
+        if "data" not in json_response or not json_response["data"]:
+            raise RuntimeError("Unexpected response format: 'data' key is missing or empty.")
+
+        return json_response["data"][0]["content"]

nv_ingest_api/internal/primitives/nim/model_interface/decorators.py

@@ -0,0 +1,56 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from functools import wraps
+from multiprocessing import Lock
+from multiprocessing import Manager
+
+logger = logging.getLogger(__name__)
+
+# Create a shared manager and lock for thread-safe access
+manager = Manager()
+global_cache = manager.dict()
+lock = Lock()
+
+
+def multiprocessing_cache(max_calls):
+    """
+    A decorator that creates a global cache shared between multiple processes.
+    The cache is invalidated after `max_calls` number of accesses.
+
+    Args:
+        max_calls (int): The number of calls after which the cache is cleared.
+
+    Returns:
+        function: The decorated function with global cache and invalidation logic.
+    """
+
+    def decorator(func):
+        call_count = manager.Value("i", 0)  # Shared integer for call counting
+
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            key = (func.__name__, args, frozenset(kwargs.items()))
+
+            with lock:
+                call_count.value += 1
+
+                if call_count.value > max_calls:
+                    global_cache.clear()
+                    call_count.value = 0
+
+                if key in global_cache:
+                    return global_cache[key]
+
+            result = func(*args, **kwargs)
+
+            with lock:
+                global_cache[key] = result
+
+            return result
+
+        return wrapper
+
+    return decorator