opengris-scaler 1.12.28__cp313-cp313-musllinux_1_2_x86_64.whl

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.")

scaler/io/ymq_sync_object_storage_connector.py

@@ -0,0 +1,184 @@
+import socket
+import uuid
+from threading import Lock
+from typing import Iterable, Optional
+
+from scaler.io.mixins import SyncObjectStorageConnector
+from scaler.io.ymq.ymq import IOContext, IOSocket, 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
+
+# Some OSes raise an OSError when sending buffers too large with send() or sendmsg().
+MAX_CHUNK_SIZE = 128 * 1024 * 1024
+
+
+class PyYMQSyncObjectStorageConnector(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._next_request_id = 0
+
+        self._socket_lock = Lock()
+
+        self._identity: str = f"{self.__class__.__name__}|{socket.gethostname().split('.')[0]}|{uuid.uuid4()}"
+        self._io_context = IOContext()
+        self._io_socket: IOSocket = self._io_context.createIOSocket_sync(self._identity, IOSocketType.Connector)
+        self._io_socket.connect_sync(self.address)
+
+    def __del__(self):
+        self.destroy()
+
+    def destroy(self):
+        with self._socket_lock:
+            if self._io_socket is not None:
+                self._io_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._io_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._io_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._io_socket.send_sync(Message(address=None, payload=header_bytes))
+            self._io_socket.send_sync(Message(address=None, payload=payload))
+        else:
+            self._io_socket.send_sync(Message(address=None, payload=header_bytes))
+
+    def __receive_response(self):
+        assert self._io_socket is not None
+
+        try:
+            header = self.__read_response_header()
+            payload = self.__read_response_payload(header)
+            return header, payload
+        except YMQException:
+            self.__raise_connection_failure()
+
+    def __read_response_header(self) -> ObjectResponseHeader:
+        assert self._io_socket is not None
+
+        header_bytes = self._io_socket.recv_sync().payload.data
+        if header_bytes is None:
+            self.__raise_connection_failure()
+
+        # 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._io_socket.recv_sync().payload.data
+            if res is None:
+                self.__raise_connection_failure()
+            assert len(res) == header.payload_length
+            return bytearray(res)
+        else:
+            return bytearray()
+
+    @staticmethod
+    def __raise_connection_failure():
+        raise ObjectStorageException("connection failure to object storage server.")

scaler/protocol/capnp/_python.py

@@ -0,0 +1,6 @@
+import capnp  # noqa
+
+import scaler.protocol.capnp.common_capnp as _common  # noqa
+import scaler.protocol.capnp.message_capnp as _message  # noqa
+import scaler.protocol.capnp.object_storage_capnp as _object_storage  # noqa
+import scaler.protocol.capnp.status_capnp as _status  # noqa

scaler/protocol/capnp/common.capnp

@@ -0,0 +1,68 @@
+@0xf57f79ac88fab620;
+
+enum TaskResultType {
+    success @0;           # if submit and task is done and get result
+    failed @1;            # if submit and task is failed on worker
+    failedWorkerDied @2;  # if submit and worker died (only happened when scheduler keep_task=False)
+}
+
+enum TaskCancelConfirmType {
+    canceled @0;               # if cancel success
+    cancelFailed @1;           # if cancel failed, this might happened if the task is in process
+    cancelNotFound @2;         # if cancel cannot find such task
+}
+
+enum TaskTransition {
+    hasCapacity @0;
+    taskResultSuccess @1;
+    taskResultFailed @2;
+    taskResultWorkerDied @3;
+    taskCancel @4;
+    taskCancelConfirmCanceled @5;
+    taskCancelConfirmFailed @6;
+    taskCancelConfirmNotFound @7;
+    balanceTaskCancel @8;
+    workerDisconnect @9;
+    schedulerHasTask @10;
+    schedulerHasNoTask @11;
+}
+
+enum TaskState {
+    inactive @0;
+    running @1;
+    canceling @2;
+    balanceCanceling @3;
+    success @4;
+    failed @5;
+    failedWorkerDied @6;
+    canceled @7;
+    canceledNotFound @8;
+    balanceCanceled @9;
+    workerDisconnecting @10;
+}
+
+enum WorkerState {
+    connected @0;
+    disconnected @1;
+}
+
+struct TaskCapability {
+    name @0 :Text;      # the name of the capability provided by the worker/required by the task (e.g. "gpu" or "linux")
+    value @1 :Int64;    # the quantity of the capability provided/required. Use -1 for quantity-less capabilities
+}
+
+struct ObjectMetadata {
+    objectIds @0 :List(Data);
+    objectTypes @1 :List(ObjectContentType);
+    objectNames @2 :List(Data);
+
+    enum ObjectContentType {
+        serializer @0;
+        object @1;
+    }
+}
+
+struct ObjectStorageAddress {
+    host @0 :Text;
+    port @1 :UInt16;
+}