opengris-scaler 1.12.28__cp313-cp313-musllinux_1_2_x86_64.whl

scaler/io/mixins.py

@@ -0,0 +1,154 @@
+import abc
+from typing import Awaitable, Callable, Optional
+
+from scaler.protocol.python.mixins import Message
+from scaler.protocol.python.status import BinderStatus
+from scaler.utility.identifiers import ObjectID
+from scaler.utility.mixins import Looper, Reporter
+
+
+class AsyncBinder(Looper, Reporter, metaclass=abc.ABCMeta):
+    @property
+    @abc.abstractmethod
+    def identity(self):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def destroy(self):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def register(self, callback: Callable[[bytes, Message], Awaitable[None]]):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def send(self, to: bytes, message: Message):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def get_status(self) -> BinderStatus:
+        raise NotImplementedError()
+
+
+class AsyncConnector(Looper, metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def destroy(self):
+        raise NotImplementedError()
+
+    @property
+    @abc.abstractmethod
+    def identity(self) -> bytes:
+        raise NotImplementedError()
+
+    @property
+    @abc.abstractmethod
+    def address(self) -> str:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def send(self, message: Message):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def receive(self) -> Optional[Message]:
+        raise NotImplementedError()
+
+
+class SyncConnector(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def destroy(self):
+        raise NotImplementedError()
+
+    @property
+    @abc.abstractmethod
+    def identity(self) -> bytes:
+        raise NotImplementedError()
+
+    @property
+    @abc.abstractmethod
+    def address(self) -> str:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def send(self, message: Message):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def receive(self) -> Optional[Message]:
+        raise NotImplementedError()
+
+
+class AsyncObjectStorageConnector(Looper, metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    async def connect(self, host: str, port: int):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def wait_until_connected(self):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def is_connected(self) -> bool:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def destroy(self):
+        raise NotImplementedError()
+
+    @property
+    @abc.abstractmethod
+    def address(self) -> str:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def set_object(self, object_id: ObjectID, payload: bytes) -> None:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def get_object(self, object_id: ObjectID, max_payload_length: int = 2**64 - 1) -> bytes:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def delete_object(self, object_id: ObjectID) -> None:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def duplicate_object_id(self, object_id: ObjectID, new_object_id: ObjectID) -> None:
+        raise NotImplementedError()
+
+
+class SyncObjectStorageConnector(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def destroy(self):
+        raise NotImplementedError()
+
+    @property
+    @abc.abstractmethod
+    def address(self) -> str:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def set_object(self, object_id: ObjectID, payload: bytes):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def get_object(self, object_id: ObjectID, max_payload_length: int = 2**64 - 1) -> bytearray:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def delete_object(self, object_id: ObjectID) -> bool:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def duplicate_object_id(self, object_id: ObjectID, new_object_id: ObjectID) -> None:
+        raise NotImplementedError()
+
+
+class SyncSubscriber(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def destroy(self):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def run(self) -> None:
+        raise NotImplementedError()

scaler/io/sync_connector.py

@@ -0,0 +1,68 @@
+import logging
+import os
+import socket
+import threading
+import uuid
+from typing import Optional
+
+import zmq
+
+from scaler.config.types.zmq import ZMQConfig
+from scaler.io.mixins import SyncConnector
+from scaler.io.utility import deserialize, serialize
+from scaler.protocol.python.mixins import Message
+
+
+class ZMQSyncConnector(SyncConnector):
+    def __init__(self, context: zmq.Context, socket_type: int, address: ZMQConfig, identity: Optional[bytes]):
+        self._address = address
+
+        self._context = context
+        self._socket = self._context.socket(socket_type)
+
+        self._identity: bytes = (
+            f"{os.getpid()}|{socket.gethostname().split('.')[0]}|{uuid.uuid4()}".encode()
+            if identity is None
+            else identity
+        )
+
+        # set socket option
+        self._socket.setsockopt(zmq.IDENTITY, self._identity)
+        self._socket.setsockopt(zmq.SNDHWM, 0)
+        self._socket.setsockopt(zmq.RCVHWM, 0)
+
+        self._socket.connect(self._address.to_address())
+
+        self._lock = threading.Lock()
+
+    def destroy(self):
+        self._socket.close()
+
+    @property
+    def address(self) -> str:
+        return self._address.to_address()
+
+    @property
+    def identity(self) -> bytes:
+        return self._identity
+
+    def send(self, message: Message):
+        with self._lock:
+            self._socket.send(serialize(message), copy=False)
+
+    def receive(self) -> Optional[Message]:
+        with self._lock:
+            payload = self._socket.recv(copy=False)
+
+        return self.__compose_message(payload.bytes)
+
+    def __compose_message(self, payload: bytes) -> Optional[Message]:
+        result: Optional[Message] = deserialize(payload)
+        if result is None:
+            logging.error(f"{self.__get_prefix()}: received unknown message: {payload!r}")
+            return None
+
+        return result
+
+    def __get_prefix(self):
+        return f"{self.__class__.__name__}[{self._identity.decode()}]:"

scaler/io/sync_object_storage_connector.py

@@ -0,0 +1,247 @@
+import collections
+import socket
+import struct
+import uuid
+from threading import Lock
+from typing import Iterable, List, Optional, Tuple
+
+from scaler.io.mixins import SyncObjectStorageConnector
+from scaler.protocol.capnp._python import _object_storage  # noqa
+from scaler.protocol.python.object_storage import ObjectRequestHeader, ObjectResponseHeader, to_capnp_object_id
+from scaler.utility.exceptions import ObjectStorageException
+from scaler.utility.identifiers import ObjectID
+
+# Some OSes raise an OSError when sending buffers too large with send() or sendmsg().
+MAX_CHUNK_SIZE = 128 * 1024 * 1024
+
+
+class PySyncObjectStorageConnector(SyncObjectStorageConnector):
+    """An synchronous connector that uses an raw TCP socket to connect to a Scaler's object storage instance."""
+
+    def __init__(self, host: str, port: int):
+        self._host = host
+        self._port = port
+
+        self._identity: bytes = (
+            f"{self.__class__.__name__}|{socket.gethostname().split('.')[0]}|{uuid.uuid4()}".encode()
+        )
+
+        self._socket: Optional[socket.socket] = socket.create_connection((self._host, self._port))
+        self._socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+
+        self._next_request_id = 0
+
+        self._socket_lock = Lock()
+
+        self.__send_buffers([struct.pack("<Q", len(self._identity)), self._identity])
+        self.__read_framed_message()  # receive server identity
+
+    def __del__(self):
+        self.destroy()
+
+    def destroy(self):
+        with self._socket_lock:
+            if self._socket is not None:
+                self._socket.close()
+                self._socket = None
+
+    @property
+    def address(self) -> str:
+        return f"tcp://{self._host}:{self._port}"
+
+    def set_object(self, object_id: ObjectID, payload: bytes):
+        """
+        Sets the object's payload on the object storage server.
+        """
+
+        with self._socket_lock:
+            self.__send_request(object_id, len(payload), ObjectRequestHeader.ObjectRequestType.SetObject, payload)
+            response_header, response_payload = self.__receive_response()
+
+        self.__ensure_response_type(response_header, [ObjectResponseHeader.ObjectResponseType.SetOK])
+        self.__ensure_empty_payload(response_payload)
+
+    def get_object(self, object_id: ObjectID, max_payload_length: int = 2**64 - 1) -> bytearray:
+        """
+        Returns the object's payload from the object storage server.
+
+        Will block until the object is available.
+        """
+
+        with self._socket_lock:
+            self.__send_request(object_id, max_payload_length, ObjectRequestHeader.ObjectRequestType.GetObject)
+            response_header, response_payload = self.__receive_response()
+
+        self.__ensure_response_type(response_header, [ObjectResponseHeader.ObjectResponseType.GetOK])
+
+        return response_payload
+
+    def delete_object(self, object_id: ObjectID) -> bool:
+        """
+        Removes the object from the object storage server.
+
+        Returns `False` if the object wasn't found in the server. Otherwise returns `True`.
+        """
+
+        with self._socket_lock:
+            self.__send_request(object_id, 0, ObjectRequestHeader.ObjectRequestType.DeleteObject)
+            response_header, response_payload = self.__receive_response()
+
+        self.__ensure_response_type(
+            response_header,
+            [ObjectResponseHeader.ObjectResponseType.DelOK, ObjectResponseHeader.ObjectResponseType.DelNotExists],
+        )
+        self.__ensure_empty_payload(response_payload)
+
+        return response_header.response_type == ObjectResponseHeader.ObjectResponseType.DelOK
+
+    def duplicate_object_id(self, object_id: ObjectID, new_object_id: ObjectID) -> None:
+        """
+        Link an object's content to a new object ID on the object storage server.
+        """
+
+        object_id_payload = to_capnp_object_id(object_id).to_bytes()
+
+        with self._socket_lock:
+            self.__send_request(
+                new_object_id,
+                len(object_id_payload),
+                ObjectRequestHeader.ObjectRequestType.DuplicateObjectID,
+                object_id_payload,
+            )
+            response_header, response_payload = self.__receive_response()
+
+        self.__ensure_response_type(response_header, [ObjectResponseHeader.ObjectResponseType.DuplicateOK])
+        self.__ensure_empty_payload(response_payload)
+
+    def __ensure_is_connected(self):
+        if self._socket is None:
+            raise ObjectStorageException("connector is closed.")
+
+    def __ensure_response_type(
+        self, header: ObjectResponseHeader, valid_response_types: Iterable[ObjectResponseHeader.ObjectResponseType]
+    ):
+        if header.response_type not in valid_response_types:
+            raise RuntimeError(f"unexpected object storage response_type={header.response_type}.")
+
+    def __ensure_empty_payload(self, payload: bytearray):
+        if len(payload) != 0:
+            raise RuntimeError(f"unexpected response payload_length={len(payload)}, expected 0.")
+
+    def __send_request(
+        self,
+        object_id: ObjectID,
+        payload_length: int,
+        request_type: ObjectRequestHeader.ObjectRequestType,
+        payload: Optional[bytes] = None,
+    ):
+        self.__ensure_is_connected()
+        assert self._socket is not None
+
+        request_id = self._next_request_id
+        self._next_request_id += 1
+        self._next_request_id %= 2**64 - 1  # UINT64_MAX
+
+        header = ObjectRequestHeader.new_msg(object_id, payload_length, request_id, request_type)
+        header_bytes = header.get_message().to_bytes()
+
+        if payload is not None:
+            self.__send_buffers(
+                [struct.pack("<Q", len(header_bytes)), header_bytes, struct.pack("<Q", len(payload)), payload]
+            )
+        else:
+            self.__send_buffers([struct.pack("<Q", len(header_bytes)), header_bytes])
+
+    def __send_buffers(self, buffers: List[bytes]) -> None:
+        if len(buffers) < 1:
+            return
+
+        total_size = sum(len(buffer) for buffer in buffers)
+
+        # If the message is small enough, first try to send it at once with sendmsg(). This would ensure the message can
+        # be transmitted within a single TCP segment.
+        if total_size < MAX_CHUNK_SIZE:
+            sent = self._socket.sendmsg(buffers)
+
+            if sent <= 0:
+                self.__raise_connection_failure()
+
+            remaining_buffers = collections.deque(buffers)
+            while sent > len(remaining_buffers[0]):
+                removed_buffer = remaining_buffers.popleft()
+                sent -= len(removed_buffer)
+
+            if sent > 0:
+                # Truncate the first partially sent buffer
+                remaining_buffers[0] = memoryview(remaining_buffers[0])[sent:]
+
+            buffers = list(remaining_buffers)
+
+        # Send the remaining buffers sequentially
+        for buffer in buffers:
+            self.__send_buffer(buffer)
+
+    def __send_buffer(self, buffer: bytes) -> None:
+        buffer_view = memoryview(buffer)
+
+        total_sent = 0
+        while total_sent < len(buffer):
+            sent = self._socket.send(buffer_view[total_sent : MAX_CHUNK_SIZE + total_sent])
+
+            if sent <= 0:
+                self.__raise_connection_failure()
+
+            total_sent += sent
+
+    def __receive_response(self) -> Tuple[ObjectResponseHeader, bytearray]:
+        assert self._socket is not None
+
+        header = self.__read_response_header()
+        payload = self.__read_response_payload(header)
+
+        return header, payload
+
+    def __read_response_header(self) -> ObjectResponseHeader:
+        assert self._socket is not None
+
+        header_bytearray = self.__read_framed_message()
+
+        # pycapnp does not like to read from a bytearray object. This look like an not-yet-resolved issue.
+        # That's is annoying because it leads to an unnecessary copy of the header's buffer.
+        # See https://github.com/capnproto/pycapnp/issues/153
+        header_bytes = bytes(header_bytearray)
+
+        with _object_storage.ObjectResponseHeader.from_bytes(header_bytes) as header_message:
+            return ObjectResponseHeader(header_message)
+
+    def __read_response_payload(self, header: ObjectResponseHeader) -> bytearray:
+        if header.payload_length > 0:
+            res = self.__read_framed_message()
+            assert len(res) == header.payload_length
+            return res
+        else:
+            return bytearray()
+
+    def __read_exactly(self, length: int) -> bytearray:
+        buffer = bytearray(length)
+
+        total_received = 0
+        while total_received < length:
+            chunk_size = min(MAX_CHUNK_SIZE, length - total_received)
+            received = self._socket.recv_into(memoryview(buffer)[total_received:], chunk_size)
+
+            if received <= 0:
+                self.__raise_connection_failure()
+
+            total_received += received
+
+        return buffer
+
+    def __read_framed_message(self) -> bytearray:
+        length_bytes = self.__read_exactly(8)
+        (payload_length,) = struct.unpack("<Q", length_bytes)
+        return self.__read_exactly(payload_length) if payload_length > 0 else bytearray()
+
+    @staticmethod
+    def __raise_connection_failure():
+        raise ObjectStorageException("connection failure to object storage server.")

scaler/io/sync_subscriber.py

@@ -0,0 +1,83 @@
+import logging
+import threading
+from typing import Callable, Optional
+
+import zmq
+
+from scaler.config.types.zmq import ZMQConfig
+from scaler.io.mixins import SyncSubscriber
+from scaler.io.utility import deserialize
+from scaler.protocol.python.mixins import Message
+
+
+class ZMQSyncSubscriber(SyncSubscriber, threading.Thread):
+    def __init__(
+        self,
+        address: ZMQConfig,
+        callback: Callable[[Message], None],
+        topic: bytes,
+        exit_callback: Optional[Callable[[], None]] = None,
+        stop_event: threading.Event = threading.Event(),
+        daemonic: bool = False,
+        timeout_seconds: int = -1,
+    ):
+        threading.Thread.__init__(self)
+
+        self._stop_event = stop_event
+        self._address = address
+        self._callback = callback
+        self._exit_callback = exit_callback
+        self._topic = topic
+        self.daemon = bool(daemonic)
+        self._timeout_seconds = timeout_seconds
+
+        self._context: Optional[zmq.Context] = None
+        self._socket: Optional[zmq.Socket] = None
+
+    def __close(self):
+        self._socket.close()
+
+    def __stop_polling(self):
+        self._stop_event.set()
+
+    def destroy(self):
+        self.__stop_polling()
+
+    def run(self) -> None:
+        self.__initialize()
+
+        while not self._stop_event.is_set():
+            self.__routine_polling()
+
+        if self._exit_callback is not None:
+            self._exit_callback()
+
+        self.__close()
+
+    def __initialize(self):
+        self._context = zmq.Context.instance()
+        self._socket = self._context.socket(zmq.SUB)
+        self._socket.setsockopt(zmq.RCVHWM, 0)
+
+        if self._timeout_seconds == -1:
+            self._socket.setsockopt(zmq.RCVTIMEO, self._timeout_seconds)
+        else:
+            self._socket.setsockopt(zmq.RCVTIMEO, self._timeout_seconds * 1000)
+
+        self._socket.subscribe(self._topic)
+        self._socket.connect(self._address.to_address())
+        self._socket.connect(self._address.to_address())
+
+    def __routine_polling(self):
+        try:
+            self.__routine_receive(self._socket.recv(copy=False).bytes)
+        except zmq.Again:
+            raise TimeoutError(f"Cannot connect to {self._address.to_address()} in {self._timeout_seconds} seconds")
+
+    def __routine_receive(self, payload: bytes):
+        result: Optional[Message] = deserialize(payload)
+        if result is None:
+            logging.error(f"received unknown message: {payload!r}")
+            return None
+
+        self._callback(result)

scaler/io/utility.py

@@ -0,0 +1,80 @@
+import logging
+import os
+from typing import List, Optional
+
+from scaler.config.defaults import CAPNP_DATA_SIZE_LIMIT, CAPNP_MESSAGE_SIZE_LIMIT
+from scaler.config.types.network_backend import NetworkBackend
+from scaler.io.async_object_storage_connector import PyAsyncObjectStorageConnector
+from scaler.io.mixins import AsyncObjectStorageConnector, SyncObjectStorageConnector
+from scaler.io.sync_object_storage_connector import PySyncObjectStorageConnector
+from scaler.protocol.capnp._python import _message  # noqa
+from scaler.protocol.python.message import PROTOCOL
+from scaler.protocol.python.mixins import Message
+
+try:
+    from collections.abc import Buffer  # type: ignore[attr-defined]
+except ImportError:
+    from typing_extensions import Buffer
+
+
+def get_scaler_network_backend_from_env():
+    backend_str = os.environ.get("SCALER_NETWORK_BACKEND", "tcp_zmq")  # Default to tcp_zmq
+    try:
+        return NetworkBackend[backend_str]
+    except KeyError:
+        return None
+
+
+def create_async_object_storage_connector(*args, **kwargs) -> AsyncObjectStorageConnector:
+    connector_type = get_scaler_network_backend_from_env()
+    if connector_type == NetworkBackend.ymq:
+        from scaler.io.ymq_async_object_storage_connector import PyYMQAsyncObjectStorageConnector
+
+        return PyYMQAsyncObjectStorageConnector(*args, **kwargs)
+
+    elif connector_type == NetworkBackend.tcp_zmq:
+        return PyAsyncObjectStorageConnector(*args, **kwargs)
+
+    else:
+        raise ValueError(
+            f"Invalid SCALER_NETWORK_BACKEND value." f"Expected one of: {[e.name for e in NetworkBackend]}"
+        )
+
+
+def create_sync_object_storage_connector(*args, **kwargs) -> SyncObjectStorageConnector:
+    connector_type = get_scaler_network_backend_from_env()
+    if connector_type == NetworkBackend.ymq:
+        from scaler.io.ymq_sync_object_storage_connector import PyYMQSyncObjectStorageConnector
+
+        return PyYMQSyncObjectStorageConnector(*args, **kwargs)
+
+    elif connector_type == NetworkBackend.tcp_zmq:
+        return PySyncObjectStorageConnector(*args, **kwargs)
+    else:
+        raise ValueError(
+            f"Invalid SCALER_NETWORK_BACKEND value." f"Expected one of: {[e.name for e in NetworkBackend]}"
+        )
+
+
+def deserialize(data: Buffer) -> Optional[Message]:
+    with _message.Message.from_bytes(data, traversal_limit_in_words=CAPNP_MESSAGE_SIZE_LIMIT) as payload:
+        if not hasattr(payload, payload.which()):
+            logging.error(f"unknown message type: {payload.which()}")
+            return None
+
+        message = getattr(payload, payload.which())
+        return PROTOCOL[payload.which()](message)
+
+
+def serialize(message: Message) -> bytes:
+    payload = _message.Message(**{PROTOCOL.inverse[type(message)]: message.get_message()})
+    return payload.to_bytes()
+
+
+def chunk_to_list_of_bytes(data: bytes) -> List[bytes]:
+    # TODO: change to list of memoryview when capnp can support memoryview
+    return [data[i : i + CAPNP_DATA_SIZE_LIMIT] for i in range(0, len(data), CAPNP_DATA_SIZE_LIMIT)]
+
+
+def concat_list_of_bytes(data: List[bytes]) -> bytes:
+    return bytearray().join(data)