nv-ingest-api 2025.4.18.dev20250418__py3-none-any.whl → 2025.4.20.dev20250420__py3-none-any.whl

nv_ingest_api/util/message_brokers/simple_message_broker/broker.py

@@ -1,465 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
-# All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-import uuid
-import socket
-import socketserver
-import json
-import logging
-import threading
-from typing import Optional
-
-from pydantic import ValidationError
-
-from nv_ingest_api.internal.schemas.message_brokers.request_schema import (
-    PushRequestSchema,
-    PopRequestSchema,
-    SizeRequestSchema,
-)
-from nv_ingest_api.internal.schemas.message_brokers.response_schema import ResponseSchema
-from nv_ingest_api.util.message_brokers.simple_message_broker.ordered_message_queue import OrderedMessageQueue
-
-logger = logging.getLogger(__name__)
-
-
-class SimpleMessageBrokerHandler(socketserver.BaseRequestHandler):
-    """
-    Handles incoming client requests for the SimpleMessageBroker server, processes commands such as
-    PUSH, POP, SIZE, and PING, and manages message queues with thread-safe operations.
-    """
-
-    def handle(self):
-        """
-        Handles incoming client requests, validates the request data, and dispatches to the appropriate
-        command handler.
-
-        Raises
-        ------
-        ValidationError
-            If the incoming request data fails schema validation.
-        Exception
-            If there is an unexpected error while processing the request.
-        """
-
-        client_address = self.client_address
-
-        data_bytes = None
-        try:
-            data_length_bytes = self._recv_exact(8)
-            if not data_length_bytes:
-                logger.debug("No data length received. Closing connection.")
-                return
-            data_length = int.from_bytes(data_length_bytes, "big")
-
-            data_bytes = self._recv_exact(data_length)
-            if not data_bytes:
-                logger.debug("No data received. Closing connection.")
-                return
-
-            data = data_bytes.decode("utf-8").strip()
-            request_data = json.loads(data)
-
-            command = request_data.get("command")
-            if not command:
-                response = ResponseSchema(response_code=1, response_reason="No command specified")
-                self._send_response(response)
-                return
-
-            # Handle the PING command directly
-            if command == "PING":
-                self._handle_ping()
-                return
-
-            # Validate and extract common fields
-            queue_name = request_data.get("queue_name")
-
-            # Initialize the queue and its lock if necessary
-            if queue_name:
-                self.server._initialize_queue(queue_name)
-                queue_lock = self.server.queue_locks[queue_name]
-                queue = self.server.queues[queue_name]
-            else:
-                queue_lock = None
-                queue = None  # For commands that don't require a queue
-
-            # Dispatch to the appropriate handler
-            if command == "PUSH":
-                validated_data = PushRequestSchema(**request_data)
-                self._handle_push(validated_data, transaction_id=str(uuid.uuid4()), queue=queue, queue_lock=queue_lock)
-            elif command == "PUSH_FOR_NV_INGEST":
-                validated_data = PushRequestSchema(**request_data)
-                self._handle_push_for_nv_ingest(validated_data, queue=queue, queue_lock=queue_lock)
-            elif command == "POP":
-                validated_data = PopRequestSchema(**request_data)
-                self._handle_pop(validated_data, queue=queue, queue_lock=queue_lock)
-            elif command == "SIZE":
-                validated_data = SizeRequestSchema(**request_data)
-                response = self._size_of_queue(validated_data, queue, queue_lock)
-                self._send_response(response)
-            else:
-                response = ResponseSchema(response_code=1, response_reason="Unknown command")
-                self._send_response(response)
-
-        except ValidationError as ve:
-            response = ResponseSchema(response_code=1, response_reason=str(ve))
-            self._send_response(response)
-        except Exception as e:
-            logger.error(f"Error processing command from {client_address}: {e}\n{data_bytes}")
-            response = ResponseSchema(response_code=1, response_reason=str(e))
-            try:
-                self._send_response(response)
-            except BrokenPipeError:
-                logger.error("Cannot send error response; client connection closed.")
-
-    def _handle_ping(self):
-        """
-        Responds to a PING command with a PONG response.
-        """
-
-        response = ResponseSchema(response_code=0, response="PONG")
-        self._send_response(response)
-
-    def _handle_push(
-        self, data: PushRequestSchema, transaction_id: str, queue: OrderedMessageQueue, queue_lock: threading.Lock
-    ):
-        """
-        Handles a PUSH command to add a message to the specified queue.
-
-        Parameters
-        ----------
-        data : PushRequestSchema
-            The validated data for the PUSH command.
-        transaction_id : str
-            The unique transaction ID for the operation.
-        queue : OrderedMessageQueue
-            The queue where the message will be pushed.
-        queue_lock : threading.Lock
-            The lock object to ensure thread-safe access to the queue.
-        """
-
-        timeout = data.timeout
-
-        with queue_lock:
-            if queue.full():
-                # Return failure response immediately
-                response = ResponseSchema(response_code=1, response_reason="Queue is full")
-                self._send_response(response)
-                return
-
-        # Proceed with the 3-way handshake
-        initial_response = ResponseSchema(
-            response_code=0, response="Transaction initiated. Waiting for ACK.", transaction_id=transaction_id
-        )
-        self._send_response(initial_response)
-
-        # Wait for ACK
-        if not self._wait_for_ack(transaction_id, timeout):
-            logger.debug(f"Transaction {transaction_id}: ACK not received. Discarding data.")
-            final_response = ResponseSchema(
-                response_code=1, response_reason="ACK not received.", transaction_id=transaction_id
-            )
-        else:
-            # Perform the PUSH operation after ACK
-            with queue_lock:
-                queue.push(data.message)
-            final_response = ResponseSchema(response_code=0, response="Data stored.", transaction_id=transaction_id)
-
-        # Send final response
-        self._send_response(final_response)
-
-    def _handle_push_for_nv_ingest(
-        self, data: PushRequestSchema, queue: OrderedMessageQueue, queue_lock: threading.Lock
-    ):
-        """
-        Handles a PUSH_FOR_NV_INGEST command, which includes generating a unique job ID and
-        updating the message payload accordingly.
-
-        Parameters
-        ----------
-        data : PushRequestSchema
-            The validated data for the PUSH_FOR_NV_INGEST command.
-        queue : OrderedMessageQueue
-            The queue where the message will be pushed.
-        queue_lock : threading.Lock
-            The lock object to ensure thread-safe access to the queue.
-        """
-
-        timeout = data.timeout
-
-        # Deserialize the message
-        try:
-            message_dict = json.loads(data.message)
-        except json.JSONDecodeError:
-            response = ResponseSchema(response_code=1, response_reason="Invalid JSON message")
-            self._send_response(response)
-            return
-
-        # Generate a UUID for 'job_id' and use it as transaction_id
-        transaction_id = str(uuid.uuid4())
-        message_dict["job_id"] = transaction_id
-
-        # Re-serialize the message
-        updated_message = json.dumps(message_dict)
-
-        with queue_lock:
-            if queue.full():
-                # Return failure response immediately
-                response = ResponseSchema(response_code=1, response_reason="Queue is full")
-                self._send_response(response)
-                return
-
-        # Proceed with the 3-way handshake
-        initial_response = ResponseSchema(
-            response_code=0, response="Transaction initiated. Waiting for ACK.", transaction_id=transaction_id
-        )
-        self._send_response(initial_response)
-
-        # Wait for ACK
-        if not self._wait_for_ack(transaction_id, timeout):
-            logger.debug(f"Transaction {transaction_id}: ACK not received. Discarding data.")
-            final_response = ResponseSchema(
-                response_code=1, response_reason="ACK not received.", transaction_id=transaction_id
-            )
-        else:
-            # Perform the PUSH operation after ACK
-            with queue_lock:
-                queue.push(updated_message)
-            final_response = ResponseSchema(response_code=0, response="Data stored.", transaction_id=transaction_id)
-
-        # Send final response
-        self._send_response(final_response)
-
-    def _handle_pop(self, data: PopRequestSchema, queue: OrderedMessageQueue, queue_lock: threading.Lock):
-        """
-        Handles a POP command to retrieve a message from the specified queue.
-
-        Parameters
-        ----------
-        data : PopRequestSchema
-            The validated data for the POP command.
-        queue : OrderedMessageQueue
-            The queue from which the message will be retrieved.
-        queue_lock : threading.Lock
-            The lock object to ensure thread-safe access to the queue.
-        """
-
-        timeout = data.timeout
-        transaction_id = str(uuid.uuid4())
-
-        with queue_lock:
-            if queue.empty():
-                # Return failure response immediately
-                response = ResponseSchema(response_code=1, response_reason="Queue is empty")
-                self._send_response(response)
-                return
-            # Pop the message from the queue
-            message = queue.pop(transaction_id)
-
-        # Proceed with the 3-way handshake
-        initial_response = ResponseSchema(response_code=0, response=message, transaction_id=transaction_id)
-        self._send_response(initial_response)
-
-        # Wait for ACK
-        if not self._wait_for_ack(transaction_id, timeout):
-            logger.debug(f"Transaction {transaction_id}: ACK not received. Returning data to queue.")
-            with queue_lock:
-                queue.return_message(transaction_id)
-            final_response = ResponseSchema(
-                response_code=1, response_reason="ACK not received.", transaction_id=transaction_id
-            )
-        else:
-            with queue_lock:
-                queue.acknowledge(transaction_id)
-            final_response = ResponseSchema(response_code=0, response="Data processed.", transaction_id=transaction_id)
-
-        # Send final response
-        self._send_response(final_response)
-
-    def _size_of_queue(
-        self, data: SizeRequestSchema, queue: OrderedMessageQueue, queue_lock: threading.Lock
-    ) -> ResponseSchema:
-        """
-        Retrieves the size of the specified queue.
-
-        Parameters
-        ----------
-        data : SizeRequestSchema
-            The validated data for the SIZE command.
-        queue : OrderedMessageQueue
-            The queue whose size will be queried.
-        queue_lock : threading.Lock
-            The lock object to ensure thread-safe access to the queue.
-
-        Returns
-        -------
-        ResponseSchema
-            A response containing the size of the queue.
-        """
-
-        with queue_lock:
-            size = queue.qsize()
-        return ResponseSchema(response_code=0, response=str(size))
-
-    def _wait_for_ack(self, transaction_id: str, timeout: Optional[float]) -> bool:
-        """
-        Waits for an acknowledgment (ACK) from the client for a specific transaction.
-
-        Parameters
-        ----------
-        transaction_id : str
-            The unique transaction ID for the operation.
-        timeout : float, optional
-            The timeout period for waiting for the ACK.
-
-        Returns
-        -------
-        bool
-            True if the ACK is received, False otherwise.
-        """
-
-        try:
-            self.request.settimeout(timeout)
-            ack_length_bytes = self._recv_exact(8)
-            if not ack_length_bytes or len(ack_length_bytes) < 8:
-                return False
-            ack_length = int.from_bytes(ack_length_bytes, "big")
-            ack_data_bytes = self._recv_exact(ack_length)
-            if not ack_data_bytes or len(ack_data_bytes) < ack_length:
-                return False
-            ack_data = ack_data_bytes.decode("utf-8")
-            ack_response = json.loads(ack_data)
-            return ack_response.get("transaction_id") == transaction_id and ack_response.get("ack") is True
-        except (socket.timeout, json.JSONDecodeError, ConnectionResetError) as e:
-            logger.error(f"Error waiting for ACK: {e}")
-            return False
-        finally:
-            self.request.settimeout(None)
-
-    def _send_response(self, response: ResponseSchema):
-        """
-        Sends a response back to the client.
-
-        Parameters
-        ----------
-        response : ResponseSchema
-            The response to send to the client.
-
-        Raises
-        ------
-        Exception
-            If there is an error while sending the response.
-        """
-
-        try:
-            response_json = response.model_dump_json().encode("utf-8")
-            total_length = len(response_json)
-            self.request.sendall(total_length.to_bytes(8, "big"))
-            self.request.sendall(response_json)
-        except BrokenPipeError as e:
-            logger.error(f"BrokenPipeError while sending response: {e}")
-            # Handle the broken pipe gracefully
-        except Exception as e:
-            logger.error(f"Unexpected error while sending response: {e}")
-
-    def _recv_exact(self, num_bytes: int) -> Optional[bytes]:
-        """
-        Receives an exact number of bytes from the client connection.
-
-        Parameters
-        ----------
-        num_bytes : int
-            The number of bytes to receive.
-
-        Returns
-        -------
-        Optional[bytes]
-            The received bytes, or None if the connection is closed.
-        """
-
-        data = bytearray()
-        while len(data) < num_bytes:
-            packet = self.request.recv(num_bytes - len(data))
-            if not packet:
-                return None
-            data.extend(packet)
-        return bytes(data)
-
-
-class SimpleMessageBroker(socketserver.ThreadingMixIn, socketserver.TCPServer):
-    """
-    A thread-safe message broker server that manages multiple message queues and supports commands
-    such as PUSH, POP, SIZE, and PING.
-    """
-
-    allow_reuse_address = True
-    _instances = {}
-    _instances_lock = threading.Lock()
-
-    def __new__(cls, host: str, port: int, max_queue_size: int):
-        """
-        Ensures that only one instance of SimpleMessageBroker is created per host and port combination.
-
-        Parameters
-        ----------
-        host : str
-            The hostname or IP address for the server.
-        port : int
-            The port number for the server.
-        max_queue_size : int
-            The maximum size of each message queue.
-
-        Returns
-        -------
-        SimpleMessageBroker
-            The singleton instance of the server.
-        """
-
-        key = (host, port)
-        with cls._instances_lock:
-            if key not in cls._instances:
-                # Create a new instance and store it in the instances dictionary
-                instance = super(SimpleMessageBroker, cls).__new__(cls)
-                cls._instances[key] = instance
-            else:
-                instance = cls._instances[key]
-        return instance
-
-    def __init__(self, host: str, port: int, max_queue_size: int):
-        """
-        Initializes the SimpleMessageBroker server, setting up message queues and locks.
-
-        Parameters
-        ----------
-        host : str
-            The hostname or IP address for the server.
-        port : int
-            The port number for the server.
-        max_queue_size : int
-            The maximum size of each message queue.
-        """
-
-        # Prevent __init__ from running multiple times on the same instance
-        if hasattr(self, "_initialized") and self._initialized:
-            return
-        super().__init__((host, port), SimpleMessageBrokerHandler)
-        self.max_queue_size = max_queue_size
-        self.queues = {}
-        self.queue_locks = {}  # Dictionary to hold locks for each queue
-        self.lock = threading.Lock()  # Global lock to protect access to queues and locks
-        self._initialized = True  # Flag to indicate initialization is complete
-
-    def _initialize_queue(self, queue_name: str):
-        """
-        Initializes a new message queue with the specified name if it doesn't already exist.
-
-        Parameters
-        ----------
-        queue_name : str
-            The name of the queue to initialize.
-        """
-
-        with self.lock:
-            if queue_name not in self.queues:
-                self.queues[queue_name] = OrderedMessageQueue(maxsize=self.max_queue_size)
-                self.queue_locks[queue_name] = threading.Lock()

nv_ingest_api/util/message_brokers/simple_message_broker/ordered_message_queue.py

@@ -1,71 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
-# All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-
-import threading
-import heapq
-
-
-class OrderedMessageQueue:
-    def __init__(self, maxsize=0):
-        self.queue = []  # List of (index, message) tuples
-        self.maxsize = maxsize
-        self.next_index = 0  # Monotonically increasing message index
-        self.in_flight = {}  # Mapping of transaction_id to (index, message)
-        self.lock = threading.Lock()
-        self.not_empty = threading.Condition(self.lock)
-        self.not_full = threading.Condition(self.lock)
-
-    def can_push(self):
-        """Check if the queue can accept more messages."""
-        with self.lock:
-            return self.maxsize == 0 or (len(self.queue) + len(self.in_flight)) < self.maxsize
-
-    def push(self, message):
-        """Add a message to the queue after it has been acknowledged."""
-        with self.lock:
-            index = self.next_index
-            self.next_index += 1
-            heapq.heappush(self.queue, (index, message))
-            self.not_empty.notify()
-            return index
-
-    def pop(self, transaction_id):
-        """Pop a message from the queue and mark it as in-flight."""
-        with self.lock:
-            while not self.queue:
-                self.not_empty.wait()
-            index, message = heapq.heappop(self.queue)
-            self.in_flight[transaction_id] = (index, message)
-            self.not_full.notify()
-            return message
-
-    def acknowledge(self, transaction_id):
-        """Acknowledge that a message has been processed."""
-        with self.lock:
-            if transaction_id in self.in_flight:
-                del self.in_flight[transaction_id]
-
-    def return_message(self, transaction_id):
-        """Return an unacknowledged message back to the queue."""
-        with self.lock:
-            if transaction_id in self.in_flight:
-                index, message = self.in_flight.pop(transaction_id)
-                heapq.heappush(self.queue, (index, message))
-                self.not_empty.notify()
-
-    def qsize(self):
-        """Get the number of messages currently in the queue."""
-        with self.lock:
-            return len(self.queue)
-
-    def empty(self):
-        """Check if the queue is empty."""
-        with self.lock:
-            return not self.queue
-
-    def full(self):
-        """Check if the queue is full."""
-        with self.lock:
-            return self.maxsize > 0 and (len(self.queue) + len(self.in_flight)) >= self.maxsize