opengris-scaler 1.12.37__cp38-cp38-musllinux_1_2_x86_64.whl

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)

scaler/io/ymq/_ymq.pyi

@@ -0,0 +1,95 @@
+# NOTE: NOT IMPLEMENTATION, TYPE INFORMATION ONLY
+# This file contains type stubs for the Ymq Python C Extension module
+from enum import IntEnum
+from typing import Callable, Optional, SupportsBytes, Union
+
+try:
+    from collections.abc import Buffer  # type: ignore[attr-defined]
+except ImportError:
+    from typing_extensions import Buffer
+
+class Bytes(Buffer):
+    data: bytes | None
+    len: int
+
+    def __init__(self, data: Buffer | None = None) -> None: ...
+    def __repr__(self) -> str: ...
+    def __len__(self) -> int: ...
+
+    # this type signature is not 100% accurate because it's implemented in C
+    # but this satisfies the type check and is good enough
+    def __buffer__(self, flags: int, /) -> memoryview: ...
+
+class Message:
+    address: Bytes | None
+    payload: Bytes
+
+    def __init__(
+        self, address: Bytes | bytes | SupportsBytes | None, payload: Bytes | bytes | SupportsBytes
+    ) -> None: ...
+    def __repr__(self) -> str: ...
+    def __str__(self) -> str: ...
+
+class IOSocketType(IntEnum):
+    Uninit = 0
+    Binder = 1
+    Connector = 2
+    Unicast = 3
+    Multicast = 4
+
+class BaseIOContext:
+    num_threads: int
+
+    def __init__(self, num_threads: int = 1) -> None: ...
+    def __repr__(self) -> str: ...
+    def createIOSocket(
+        self, callback: Callable[[Union[BaseIOSocket, Exception]], None], identity: str, socket_type: IOSocketType
+    ) -> None:
+        """Create an io socket with an identity and socket type"""
+
+class BaseIOSocket:
+    identity: str
+    socket_type: IOSocketType
+
+    def __repr__(self) -> str: ...
+    def send(self, callback: Callable[[Optional[Exception]], None], message: Message) -> None:
+        """Send a message to one of the socket's peers"""
+
+    def recv(self, callback: Callable[[Union[Message, Exception]], None]) -> None:
+        """Receive a message from one of the socket's peers"""
+
+    def bind(self, callback: Callable[[Optional[Exception]], None], address: str) -> None:
+        """Bind the socket to an address and listen for incoming connections"""
+
+    def connect(self, callback: Callable[[Optional[Exception]], None], address: str) -> None:
+        """Connect to a remote socket"""
+
+class ErrorCode(IntEnum):
+    Uninit = 0
+    InvalidPortFormat = 1
+    InvalidAddressFormat = 2
+    ConfigurationError = 3
+    SignalNotSupported = 4
+    CoreBug = 5
+    RepetetiveIOSocketIdentity = 6
+    RedundantIOSocketRefCount = 7
+    MultipleConnectToNotSupported = 8
+    MultipleBindToNotSupported = 9
+    InitialConnectFailedWithInProgress = 10
+    SendMessageRequestCouldNotComplete = 11
+    SetSockOptNonFatalFailure = 12
+    IPv6NotSupported = 13
+    RemoteEndDisconnectedOnSocketWithoutGuaranteedDelivery = 14
+    ConnectorSocketClosedByRemoteEnd = 15
+    IOSocketStopRequested = 16
+    BinderSendMessageWithNoAddress = 17
+
+    def explanation(self) -> str: ...
+
+class YMQException(Exception):
+    code: ErrorCode
+    message: str
+
+    def __init__(self, /, code: ErrorCode, message: str) -> None: ...
+    def __repr__(self) -> str: ...
+    def __str__(self) -> str: ...

scaler/io/ymq/ymq.py

@@ -0,0 +1,138 @@
+# This file wraps the interface exported by the C implementation of the module
+# and provides a more ergonomic interface supporting both asynchronous and synchronous execution
+
+__all__ = ["IOSocket", "IOContext", "Message", "IOSocketType", "YMQException", "Bytes", "ErrorCode"]
+
+import asyncio
+import concurrent.futures
+from typing import Any, Callable, Optional, TypeVar, Union
+
+try:
+    from typing import Concatenate, ParamSpec  # type: ignore[attr-defined]
+except ImportError:
+    from typing_extensions import ParamSpec, Concatenate  # type: ignore[assignment]
+
+from scaler.io.ymq._ymq import BaseIOContext, BaseIOSocket, Bytes, ErrorCode, IOSocketType, Message, YMQException
+
+
+class IOSocket:
+    _base: BaseIOSocket
+
+    def __init__(self, base: BaseIOSocket) -> None:
+        self._base = base
+
+    @property
+    def socket_type(self) -> IOSocketType:
+        return self._base.socket_type
+
+    @property
+    def identity(self) -> str:
+        return self._base.identity
+
+    async def bind(self, address: str) -> None:
+        """Bind the socket to an address and listen for incoming connections"""
+        await call_async(self._base.bind, address)
+
+    def bind_sync(self, address: str, /, timeout: Optional[float] = None) -> None:
+        """Bind the socket to an address and listen for incoming connections"""
+        call_sync(self._base.bind, address, timeout=timeout)
+
+    async def connect(self, address: str) -> None:
+        """Connect to a remote socket"""
+        await call_async(self._base.connect, address)
+
+    def connect_sync(self, address: str, /, timeout: Optional[float] = None) -> None:
+        """Connect to a remote socket"""
+        call_sync(self._base.connect, address, timeout=timeout)
+
+    async def send(self, message: Message) -> None:
+        """Send a message to one of the socket's peers"""
+        await call_async(self._base.send, message)
+
+    def send_sync(self, message: Message, /, timeout: Optional[float] = None) -> None:
+        """Send a message to one of the socket's peers"""
+        call_sync(self._base.send, message, timeout=timeout)
+
+    async def recv(self) -> Message:
+        """Receive a message from one of the socket's peers"""
+        return await call_async(self._base.recv)
+
+    def recv_sync(self, /, timeout: Optional[float] = None) -> Message:
+        """Receive a message from one of the socket's peers"""
+        return call_sync(self._base.recv, timeout=timeout)
+
+
+class IOContext:
+    _base: BaseIOContext
+
+    def __init__(self, num_threads: int = 1) -> None:
+        self._base = BaseIOContext(num_threads)
+
+    @property
+    def num_threads(self) -> int:
+        return self._base.num_threads
+
+    async def createIOSocket(self, identity: str, socket_type: IOSocketType) -> IOSocket:
+        """Create an io socket with an identity and socket type"""
+        return IOSocket(await call_async(self._base.createIOSocket, identity, socket_type))
+
+    def createIOSocket_sync(self, identity: str, socket_type: IOSocketType) -> IOSocket:
+        """Create an io socket with an identity and socket type"""
+        return IOSocket(call_sync(self._base.createIOSocket, identity, socket_type))
+
+
+P = ParamSpec("P")
+T = TypeVar("T")
+
+
+async def call_async(
+    func: Callable[Concatenate[Callable[[Union[T, BaseException]], None], P], None],  # type: ignore
+    *args: P.args,  # type: ignore
+    **kwargs: P.kwargs,  # type: ignore
+) -> T:
+    loop = asyncio.get_event_loop()
+    future = loop.create_future()
+
+    def callback(result: Union[T, BaseException]):
+        if isinstance(result, BaseException):
+            loop.call_soon_threadsafe(_safe_set_exception, future, result)
+        else:
+            loop.call_soon_threadsafe(_safe_set_result, future, result)
+
+    func(callback, *args, **kwargs)
+    return await future
+
+
+# about the ignore directives: mypy cannot properly handle typing extension's ParamSpec and Concatenate in python <=3.9
+# these type hints are correctly understood in Python 3.10+
+def call_sync(  # type: ignore[valid-type]
+    func: Callable[Concatenate[Callable[[Union[T, BaseException]], None], P], None],  # type: ignore
+    *args: P.args,  # type: ignore
+    timeout: Optional[float] = None,
+    **kwargs: P.kwargs,  # type: ignore
+) -> T:  # type: ignore
+    future: concurrent.futures.Future = concurrent.futures.Future()
+
+    def callback(result: Union[T, BaseException]):
+        if future.done():
+            return
+
+        if isinstance(result, BaseException):
+            future.set_exception(result)
+        else:
+            future.set_result(result)
+
+    func(callback, *args, **kwargs)
+    return future.result(timeout)
+
+
+def _safe_set_result(future: asyncio.Future, result: Any) -> None:
+    if future.done():
+        return
+    future.set_result(result)
+
+
+def _safe_set_exception(future: asyncio.Future, exc: BaseException) -> None:
+    if future.done():
+        return
+    future.set_exception(exc)

scaler/io/ymq_async_object_storage_connector.py

@@ -0,0 +1,184 @@
+import asyncio
+import logging
+import socket
+import uuid
+from typing import Dict, Optional, Tuple
+
+from scaler.io.mixins import AsyncObjectStorageConnector
+from scaler.io.ymq.ymq import IOContext, IOSocketType, Message, YMQException
+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
+
+
+class PyYMQAsyncObjectStorageConnector(AsyncObjectStorageConnector):
+    """An asyncio connector that uses YMQ to connect to a Scaler's object storage instance."""
+
+    def __init__(self):
+        self._host: Optional[str] = None
+        self._port: Optional[int] = None
+
+        self._connected_event = asyncio.Event()
+
+        self._next_request_id = 0
+        self._pending_get_requests: Dict[ObjectID, asyncio.Future] = {}
+
+        self._lock = asyncio.Lock()
+        self._identity: str = f"{self.__class__.__name__}|{socket.gethostname().split('.')[0]}|{uuid.uuid4()}"
+        self._io_context: IOContext = IOContext()
+        self._io_socket = self._io_context.createIOSocket_sync(self._identity, IOSocketType.Connector)
+
+    def __del__(self):
+        if not self.is_connected():
+            return
+        self._io_socket = None
+
+    async def connect(self, host: str, port: int):
+        self._host = host
+        self._port = port
+
+        if self.is_connected():
+            raise ObjectStorageException("connector is already connected.")
+        await self._io_socket.connect(self.address)
+        self._connected_event.set()
+
+    async def wait_until_connected(self):
+        await self._connected_event.wait()
+
+    def is_connected(self) -> bool:
+        return self._connected_event.is_set()
+
+    async def destroy(self):
+        if not self.is_connected():
+            return
+        self._io_socket = None
+
+    @property
+    def address(self) -> str:
+        return f"tcp://{self._host}:{self._port}"
+
+    async def routine(self):
+        await self.wait_until_connected()
+
+        response = await self.__receive_response()
+        if response is None:
+            return
+
+        header, payload = response
+
+        if header.response_type != ObjectResponseHeader.ObjectResponseType.GetOK:
+            return
+
+        pending_get_future = self._pending_get_requests.pop(header.object_id, None)
+
+        if pending_get_future is None:
+            logging.warning(f"unknown get-ok response for unrequested object_id={repr(header.object_id)}.")
+            return
+
+        pending_get_future.set_result(payload)
+
+    async def set_object(self, object_id: ObjectID, payload: bytes) -> None:
+        await self.__send_request(object_id, len(payload), ObjectRequestHeader.ObjectRequestType.SetObject, payload)
+
+    async def get_object(self, object_id: ObjectID, max_payload_length: int = 2**64 - 1) -> bytes:
+        pending_get_future = self._pending_get_requests.get(object_id)
+
+        if pending_get_future is None:
+            pending_get_future = asyncio.Future()
+            self._pending_get_requests[object_id] = pending_get_future
+
+            await self.__send_request(
+                object_id, max_payload_length, ObjectRequestHeader.ObjectRequestType.GetObject, None
+            )
+
+        return await pending_get_future
+
+    async def delete_object(self, object_id: ObjectID) -> None:
+        await self.__send_request(object_id, 0, ObjectRequestHeader.ObjectRequestType.DeleteObject, None)
+
+    async def duplicate_object_id(self, object_id: ObjectID, new_object_id: ObjectID) -> None:
+        object_id_payload = to_capnp_object_id(object_id).to_bytes()
+
+        await self.__send_request(
+            new_object_id,
+            len(object_id_payload),
+            ObjectRequestHeader.ObjectRequestType.DuplicateObjectID,
+            object_id_payload,
+        )
+
+    def __ensure_is_connected(self):
+        if self._io_socket is None:
+            raise ObjectStorageException("connector is not connected.")
+
+    async def __send_request(
+        self,
+        object_id: ObjectID,
+        payload_length: int,
+        request_type: ObjectRequestHeader.ObjectRequestType,
+        payload: Optional[bytes],
+    ):
+        self.__ensure_is_connected()
+
+        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)
+
+        try:
+            async with self._lock:
+                await self.__write_request_header(header)
+
+                if payload is not None:
+                    await self.__write_request_payload(payload)
+
+        except YMQException:
+            self._io_socket = None
+            self.__raise_connection_failure()
+
+    async def __write_request_header(self, header: ObjectRequestHeader):
+        assert self._io_socket is not None
+        await self._io_socket.send(Message(address=None, payload=header.get_message().to_bytes()))
+
+    async def __write_request_payload(self, payload: bytes):
+        assert self._io_socket is not None
+        await self._io_socket.send(Message(address=None, payload=payload))
+
+    async def __receive_response(self) -> Optional[Tuple[ObjectResponseHeader, bytes]]:
+        if self._io_socket is None:
+            return None
+
+        try:
+            header = await self.__read_response_header()
+            payload = await self.__read_response_payload(header)
+        except YMQException:
+            self._io_socket = None
+            self.__raise_connection_failure()
+
+        return header, payload
+
+    async def __read_response_header(self) -> ObjectResponseHeader:
+        assert self._io_socket is not None
+
+        msg = await self._io_socket.recv()
+        header_data = msg.payload.data
+        assert len(header_data) == ObjectResponseHeader.MESSAGE_LENGTH
+
+        with _object_storage.ObjectResponseHeader.from_bytes(header_data) as header_message:
+            return ObjectResponseHeader(header_message)
+
+    async def __read_response_payload(self, header: ObjectResponseHeader) -> bytes:
+        assert self._io_socket is not None
+        # assert self._reader is not None
+
+        if header.payload_length > 0:
+            res = await self._io_socket.recv()
+            assert len(res.payload) == header.payload_length
+            return res.payload.data
+        else:
+            return b""
+
+    @staticmethod
+    def __raise_connection_failure():
+        raise ObjectStorageException("connection failure to object storage server.")