nv-ingest-api 2025.4.15.dev20250415__py3-none-any.whl → 2025.4.17.dev20250417__py3-none-any.whl

nv_ingest_api/util/message_brokers/simple_message_broker/simple_client.py

@@ -0,0 +1,435 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# NOTE: This code is duplicated from the ingest service:
+# src/nv_ingest/util/message_brokers/simple_message_broker/simple_client.py
+# Eventually we should move all client wrappers for the message broker into a shared library that both the ingest
+# service and the client can use.
+
+import socket
+import json
+import time
+import logging
+from typing import Optional
+
+from nv_ingest_api.internal.schemas.message_brokers.response_schema import ResponseSchema
+from nv_ingest_api.util.service_clients.client_base import MessageBrokerClientBase
+
+logger = logging.getLogger(__name__)
+
+
+class SimpleClient(MessageBrokerClientBase):
+    """
+    A client for interfacing with SimpleMessageBroker, creating a new socket connection per request
+    to ensure thread safety and robustness. Respects timeouts for all operations.
+    """
+
+    def __init__(
+        self,
+        host: str,
+        port: int,
+        db: int = 0,
+        max_retries: int = 3,
+        max_backoff: int = 32,
+        connection_timeout: int = 300,
+        max_pool_size: int = 128,
+        use_ssl: bool = False,
+    ):
+        """
+        Initialize the SimpleClient with configuration parameters.
+
+        Parameters
+        ----------
+        host : str
+            The hostname or IP address of the broker.
+        port : int
+            The port number of the broker.
+        db : int, optional
+            The database index (default: 0).
+        max_retries : int, optional
+            Maximum number of retries for operations (default: 3).
+        max_backoff : int, optional
+            Maximum backoff time in seconds for retries (default: 32).
+        connection_timeout : int, optional
+            Timeout in seconds for establishing a connection (default: 300).
+        max_pool_size : int, optional
+            Maximum pool size for connections (default: 128).
+        use_ssl : bool, optional
+            Whether to use SSL for connections (default: False).
+        """
+
+        self._host = host
+        self._port = port
+        self._db = db
+        self._max_retries = max_retries
+        self._max_backoff = max_backoff
+        self._max_pool_size = max_pool_size
+        self._connection_timeout = connection_timeout
+        self._use_ssl = use_ssl
+
+    def get_client(self):
+        """
+        Retrieve the current client instance.
+
+        Returns
+        -------
+        SimpleClient
+            The current client instance.
+        """
+        return self
+
+    def submit_message(
+        self, queue_name: str, message: str, timeout: Optional[float] = None, for_nv_ingest: bool = False
+    ) -> ResponseSchema:
+        """
+        Submit a message to the specified queue.
+
+        Parameters
+        ----------
+        queue_name : str
+            The name of the queue.
+        message : str
+            The message to be submitted.
+        timeout : float, optional
+            Timeout in seconds for the operation.
+        for_nv_ingest : bool, optional
+            Indicates whether the message is for NV ingest operations.
+
+        Returns
+        -------
+        ResponseSchema
+            The response from the broker.
+        """
+        return self._handle_push(queue_name, message, timeout, for_nv_ingest)
+
+    def fetch_message(self, queue_name: str, timeout: Optional[float] = None) -> ResponseSchema:
+        """
+        Fetch a message from the specified queue.
+
+        Parameters
+        ----------
+        queue_name : str
+            The name of the queue.
+        timeout : float, optional
+            Timeout in seconds for the operation.
+
+        Returns
+        -------
+        ResponseSchema
+            The response containing the fetched message.
+        """
+        return self._handle_pop(queue_name, timeout)
+
+    def ping(self) -> ResponseSchema:
+        """
+        Ping the broker to check connectivity.
+
+        Returns
+        -------
+        ResponseSchema
+            The response indicating the success of the ping operation.
+        """
+        command = {"command": "PING"}
+        return self._execute_simple_command(command)
+
+    def size(self, queue_name: str) -> ResponseSchema:
+        """
+        Fetch the size of the specified queue.
+
+        Parameters
+        ----------
+        queue_name : str
+            The name of the queue.
+
+        Returns
+        -------
+        ResponseSchema
+            The response containing the queue size.
+        """
+        command = {"command": "SIZE", "queue_name": queue_name}
+        return self._execute_simple_command(command)
+
+    def _handle_push(
+        self, queue_name: str, message: str, timeout: Optional[float], for_nv_ingest: bool
+    ) -> ResponseSchema:
+        """
+        Push a message to the queue with optional timeout.
+
+        Parameters
+        ----------
+        queue_name : str
+            The name of the queue.
+        message : str
+            The message to push.
+        timeout : float, optional
+            Timeout in seconds for the operation.
+        for_nv_ingest : bool
+            Indicates whether the message is for NV ingest operations.
+
+        Returns
+        -------
+        ResponseSchema
+            The response from the broker.
+        """
+
+        if not queue_name or not isinstance(queue_name, str):
+            return ResponseSchema(response_code=1, response_reason="Invalid queue name.")
+        if not message or not isinstance(message, str):
+            return ResponseSchema(response_code=1, response_reason="Invalid message.")
+
+        if for_nv_ingest:
+            command = {"command": "PUSH_FOR_NV_INGEST", "queue_name": queue_name, "message": message}
+        else:
+            command = {"command": "PUSH", "queue_name": queue_name, "message": message}
+
+        if timeout is not None:
+            command["timeout"] = timeout
+
+        start_time = time.time()
+        while True:
+            elapsed = time.time() - start_time
+            remaining_timeout = (timeout - elapsed) if (timeout is not None) else None
+            if (remaining_timeout is not None) and (remaining_timeout <= 0):
+                return ResponseSchema(response_code=1, response_reason="PUSH operation timed out.")
+
+            try:
+                with socket.create_connection((self._host, self._port), timeout=self._connection_timeout) as sock:
+                    self._send(sock, json.dumps(command).encode("utf-8"))
+                    # Receive initial response with transaction ID
+                    response_data = self._recv(sock)
+                    response = json.loads(response_data)
+
+                    if response.get("response_code") != 0:
+                        if (
+                            response.get("response_reason") == "Queue is full"
+                            or response.get("response_reason") == "Queue is not available"
+                        ):
+                            time.sleep(0.5)
+                            continue
+                        else:
+                            return ResponseSchema(**response)
+
+                    if "transaction_id" not in response:
+                        error_msg = "No transaction_id in response."
+                        logger.error(error_msg)
+
+                        return ResponseSchema(response_code=1, response_reason=error_msg)
+
+                    transaction_id = response["transaction_id"]
+
+                    # Send ACK
+                    ack_data = json.dumps({"transaction_id": transaction_id, "ack": True}).encode("utf-8")
+                    self._send(sock, ack_data)
+
+                    # Receive final response
+                    final_response_data = self._recv(sock)
+                    final_response = json.loads(final_response_data)
+
+                    return ResponseSchema(**final_response)
+
+            except (ConnectionError, socket.error, BrokenPipeError):
+                pass
+            except json.JSONDecodeError:
+                return ResponseSchema(response_code=1, response_reason="Invalid JSON response from server.")
+            except Exception as e:
+                return ResponseSchema(response_code=1, response_reason=str(e))
+
+            time.sleep(0.5)  # Backoff delay before retry
+
+    def _handle_pop(self, queue_name: str, timeout: Optional[float]) -> ResponseSchema:
+        """
+        Pop a message from the queue with optional timeout.
+
+        Parameters
+        ----------
+        queue_name : str
+            The name of the queue.
+        timeout : float, optional
+            Timeout in seconds for the operation.
+
+        Returns
+        -------
+        ResponseSchema
+            The response containing the popped message.
+        """
+
+        if not queue_name or not isinstance(queue_name, str):
+            return ResponseSchema(response_code=1, response_reason="Invalid queue name.")
+
+        command = {"command": "POP", "queue_name": queue_name}
+        if timeout is not None:
+            command["timeout"] = timeout
+
+        start_time = time.time()
+        while True:
+            elapsed = time.time() - start_time
+            remaining_timeout = timeout - elapsed if timeout else None
+            if remaining_timeout is not None and remaining_timeout <= 0:
+                return ResponseSchema(response_code=1, response_reason="POP operation timed out.")
+
+            try:
+                with socket.create_connection((self._host, self._port), timeout=self._connection_timeout) as sock:
+                    self._send(sock, json.dumps(command).encode("utf-8"))
+                    # Receive initial response with transaction ID and message
+                    response_data = self._recv(sock)
+                    response = json.loads(response_data)
+
+                    if response.get("response_code") != 0:
+                        if response.get("response_reason") == "Queue is empty":
+                            time.sleep(0.1)
+                            continue
+                        else:
+                            return ResponseSchema(**response)
+
+                    if "transaction_id" not in response:
+                        error_msg = "No transaction_id in response."
+
+                        return ResponseSchema(response_code=1, response_reason=error_msg)
+
+                    transaction_id = response["transaction_id"]
+                    message = response.get("response")
+
+                    # Send ACK
+                    ack_data = json.dumps({"transaction_id": transaction_id, "ack": True}).encode("utf-8")
+                    self._send(sock, ack_data)
+
+                    # Receive final response
+                    final_response_data = self._recv(sock)
+                    final_response = json.loads(final_response_data)
+
+                    if final_response.get("response_code") == 0:
+                        return ResponseSchema(response_code=0, response=message, transaction_id=transaction_id)
+                    else:
+                        return ResponseSchema(**final_response)
+
+            except (ConnectionError, socket.error, BrokenPipeError):
+                pass
+            except json.JSONDecodeError:
+                return ResponseSchema(response_code=1, response_reason="Invalid JSON response from server.")
+            except Exception as e:
+                return ResponseSchema(response_code=1, response_reason=str(e))
+
+            time.sleep(0.1)  # Backoff delay before retry
+
+    def _execute_simple_command(self, command: dict) -> ResponseSchema:
+        """
+        Execute a simple command on the broker and process the response.
+
+        Parameters
+        ----------
+        command : dict
+            The command to execute.
+
+        Returns
+        -------
+        ResponseSchema
+            The response from the broker.
+        """
+
+        if isinstance(command, dict):
+            data = json.dumps(command).encode("utf-8")
+        elif isinstance(command, str):
+            data = command.encode("utf-8")
+
+        try:
+            with socket.create_connection((self._host, self._port), timeout=self._connection_timeout) as sock:
+                self._send(sock, data)
+                response_data = self._recv(sock)
+                response = json.loads(response_data)
+                return ResponseSchema(**response)
+        except (ConnectionError, socket.error, BrokenPipeError) as e:
+            return ResponseSchema(response_code=1, response_reason=f"Connection error: {e}")
+        except json.JSONDecodeError:
+            return ResponseSchema(response_code=1, response_reason="Invalid JSON response from server.")
+        except Exception as e:
+            return ResponseSchema(response_code=1, response_reason=str(e))
+
+    def _send(self, sock: socket.socket, data: bytes) -> None:
+        """
+        Send data over a socket connection with a length header.
+
+        Parameters
+        ----------
+        sock : socket.socket
+            The socket connection.
+        data : bytes
+            The data to send.
+
+        Raises
+        ------
+        ConnectionError
+            If sending data fails.
+        """
+
+        total_length = len(data)
+        if total_length == 0:
+            raise ValueError("Cannot send an empty message.")
+
+        try:
+            sock.sendall(total_length.to_bytes(8, "big"))
+            sock.sendall(data)
+        except (socket.error, BrokenPipeError):
+            raise ConnectionError("Failed to send data.")
+
+    def _recv(self, sock: socket.socket) -> str:
+        """
+        Receive data from a socket connection based on a length header.
+
+        Parameters
+        ----------
+        sock : socket.socket
+            The socket connection.
+
+        Returns
+        -------
+        str
+            The received data as a string.
+
+        Raises
+        ------
+        ConnectionError
+            If receiving data fails.
+        """
+
+        try:
+            length_header = self._recv_exact(sock, 8)
+            if not length_header:
+                raise ConnectionError("Incomplete length header received.")
+            total_length = int.from_bytes(length_header, "big")
+            data_bytes = self._recv_exact(sock, total_length)
+            if not data_bytes:
+                raise ConnectionError("Incomplete message received.")
+            return data_bytes.decode("utf-8")
+        except (socket.error, BrokenPipeError, ConnectionError):
+            raise ConnectionError("Failed to receive data.")
+
+    def _recv_exact(self, sock: socket.socket, num_bytes: int) -> Optional[bytes]:
+        """
+        Receive an exact number of bytes from a socket connection.
+
+        Parameters
+        ----------
+        sock : socket.socket
+            The socket connection.
+        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:
+            try:
+                packet = sock.recv(num_bytes - len(data))
+                if not packet:
+                    return None
+                data.extend(packet)
+            except socket.timeout:
+                return None
+            except Exception:
+                return None
+        return bytes(data)

nv_ingest_api/util/metadata/__init__.py

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