QuLab 2.10.10__cp313-cp313-macosx_10_13_universal2.whl

qulab/sys/rpc/msgpack.pyi

@@ -0,0 +1,41 @@
+from typing import Any
+
+__version__: str
+
+version: tuple[int, int, int]
+
+def pack(obj, fp, **options) -> None: ...
+def packb(obj, **options) -> bytes: ...
+def dump(obj, fp, **options) -> None: ...
+def dumps(obj, **options) -> bytes: ...
+
+def unpackb(s: bytes | bytearray, **options) -> Any: ...
+def unpack(fp, **options) -> Any: ...
+def loads(s: bytes | bytearray, **options) -> Any: ...
+def load(fp, **options) -> Any: ...
+
+class Ext:
+    type: int
+    data: bytes
+    def __init__(self, type: int, data: bytes) -> None: ...
+    def __eq__(self, other) -> bool: ...
+    def __ne__(self, other) -> bool: ...
+    def __hash__(self) -> int: ...
+
+class InvalidString(bytes): ...
+
+def ext_serializable(ext_type: int): ...
+
+class PackException(Exception): ...
+class UnpackException(Exception): ...
+class UnsupportedTypeException(PackException): ...
+class InsufficientDataException(UnpackException): ...
+class InvalidStringException(UnpackException): ...
+class UnsupportedTimestampException(UnpackException): ...
+class ReservedCodeException(UnpackException): ...
+class UnhashableKeyException(UnpackException): ...
+class DuplicateKeyException(UnpackException): ...
+KeyNotPrimitiveException = UnhashableKeyException
+KeyDuplicateException = DuplicateKeyException
+
+compatibility: bool

qulab/sys/rpc/router.py

@@ -0,0 +1,35 @@
+import asyncio
+import random
+
+import zmq
+import zmq.asyncio
+
+
+async def handle_client(socket, identity, message):
+    print(f"Received request from {identity}: {message.decode()}")
+    # 随机延时 0 到 3 秒
+    await asyncio.sleep(random.uniform(0, 3))
+    task_id = random.randint(1000, 9999)  # 随机生成一个任务 ID
+    await socket.send_multipart([identity, f"Task ID: {task_id}".encode()])
+    print(f"Sent Task ID {task_id} to {identity}")
+
+
+async def server():
+    context = zmq.asyncio.Context()
+    socket = context.socket(zmq.ROUTER)
+    socket.bind("tcp://*:5555")
+
+    while True:
+        try:
+            identity, message = await socket.recv_multipart()
+            asyncio.create_task(handle_client(socket, identity, message))
+        except Exception as e:
+            print(f"An error occurred: {e}")
+            break
+
+    socket.close()
+    context.term()
+
+
+if __name__ == "__main__":
+    asyncio.run(server())

qulab/sys/rpc/rpc.py

@@ -0,0 +1,412 @@
+import asyncio
+import functools
+import inspect
+import logging
+import platform
+import struct
+from abc import ABC, abstractmethod
+from collections.abc import Awaitable
+
+from ...version import __version__
+from .exceptions import RPCError, RPCServerError, RPCTimeout
+from .serialize import pack, unpack
+from .utils import acceptArg
+
+if platform.system() == "Windows":
+    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
+
+log = logging.getLogger(__name__)  # pylint: disable=invalid-name
+
+msgIDFormat = struct.Struct("!IIQ")
+
+__msgIndex = 1024
+
+
+def nextMsgID(clientID, sessionID=0):
+    global __msgIndex
+    __msgIndex += 1
+    return msgIDFormat.pack(clientID, sessionID, __msgIndex)
+
+
+def parseMsgID(msgID):
+    """
+    return: (clientID, sessionID, msgIndex)
+    """
+    return msgIDFormat.unpack(msgID)
+
+
+# message type
+
+RPC_REQUEST = b'\x01'
+RPC_RESPONSE = b'\x02'
+RPC_PING = b'\x03'
+RPC_PONG = b'\x04'
+RPC_CANCEL = b'\x05'
+RPC_SHUTDOWN = b'\x06'
+RPC_CONNECT = b'\x07'
+RPC_WELCOME = b'\x08'
+
+RPC_MSGIDSIZE = msgIDFormat.size
+
+
+class RPCMixin(ABC):
+    info = {}
+
+    def start(self):
+        pass
+
+    def stop(self):
+        pass
+
+    def close(self):
+        pass
+
+    @property
+    @abstractmethod
+    def loop(self):
+        """
+        Event loop.
+        """
+
+    @abstractmethod
+    async def sendto(self, data, address):
+        """
+        Send message to address.
+        """
+
+    __rpc_handlers = {
+        RPC_CONNECT: 'on_connect',
+        RPC_WELCOME: 'on_welcome',
+        RPC_PING: 'on_ping',
+        RPC_PONG: 'on_pong',
+        RPC_REQUEST: 'on_request',
+        RPC_RESPONSE: 'on_response',
+        RPC_CANCEL: 'on_cancel',
+        RPC_SHUTDOWN: 'on_shutdown',
+    }
+
+    def parseData(self, data):
+        msg_type, msg = data[:1], data[1:]
+        if msg_type in [RPC_PING, RPC_PONG, RPC_CONNECT, RPC_WELCOME]:
+            return msg_type, msg
+        elif msg_type in [RPC_REQUEST, RPC_RESPONSE, RPC_CANCEL, RPC_SHUTDOWN]:
+            msgID, msg = msg[:RPC_MSGIDSIZE], msg[RPC_MSGIDSIZE:]
+            return msg_type, msgID, msg
+        # elif msg_type in [RPC_LONGREQUEST, RPC_LONGRESPONSE]:
+        #     msgID, sessionID, msg = msg[:20], msg[20:40], msg[40:]
+        #     return msg_type, msgID, sessionID, msg
+        else:
+            raise RPCError(f'Unkown message type {msg_type}.')
+
+    def handle(self, source, data):
+        """
+        Handle received data.
+
+        Should be called whenever received data from outside.
+        """
+        msg_type, *args = self.parseData(data)
+        log.debug(f'received request {msg_type} from {source}')
+        handler = self.__rpc_handlers.get(msg_type, None)
+        if handler is not None:
+            getattr(self, handler)(source, *args)
+        else:
+            log.error(f'No handler found for request {msg_type} from {source}')
+
+    async def pong(self, addr):
+        await self.sendto(RPC_PONG, addr)
+
+    def on_ping(self, source, msg):
+        log.debug(f"received ping from {source}")
+        asyncio.ensure_future(self.pong(source), loop=self.loop)
+
+
+class RPCClientMixin(RPCMixin):
+    _client_defualt_timeout = 10
+    __pending = None
+    __clientID = 1
+
+    @property
+    def pending(self):
+        if self.__pending is None:
+            self.__pending = {}
+        return self.__pending
+
+    @property
+    def clientID(self):
+        return self.__clientID
+
+    def createPending(self, addr, msgID, timeout=1, cancelRemote=True):
+        """
+        Create a future for request, wait response before timeout.
+        """
+        fut = self.loop.create_future()
+        self.pending[msgID] = (fut,
+                               self.loop.call_later(timeout,
+                                                    self.cancelPending, addr,
+                                                    msgID, cancelRemote))
+
+        def clean(fut, msgID=msgID):
+            if msgID in self.pending:
+                del self.pending[msgID]
+
+        fut.add_done_callback(clean)
+
+        return fut
+
+    def cancelPending(self, addr, msgID, cancelRemote):
+        """
+        Give up when request timeout and try to cancel remote task.
+        """
+        if msgID in self.pending:
+            fut, timeout = self.pending[msgID]
+            if cancelRemote:
+                self.cancelRemoteTask(addr, msgID)
+            if not fut.done():
+                fut.set_exception(
+                    RPCTimeout(
+                        f'Node({self.info}): Wait response from {addr} timeout.'
+                    ))
+
+    def cancelRemoteTask(self, addr, msgID):
+        """
+        Try to cancel remote task.
+        """
+        asyncio.ensure_future(self.sendto(RPC_CANCEL + msgID, addr),
+                              loop=self.loop)
+
+    def close(self):
+        self.stop()
+        for fut, timeout in list(self.pending.values()):
+            fut.cancel()
+            timeout.cancel()
+        self.pending.clear()
+
+    def setTimeout(self, timeout=10):
+        self._client_defualt_timeout = timeout
+
+    def remoteCall(self, addr, methodNane, sessionID=0, args=(), kw={}):
+        if 'timeout' in kw:
+            timeout = kw['timeout']
+        else:
+            timeout = self._client_defualt_timeout
+        msg = pack((methodNane, args, kw))
+        msgID = nextMsgID(self.clientID, sessionID)
+        asyncio.ensure_future(self.request(addr, msgID, msg), loop=self.loop)
+        return self.createPending(addr, msgID, timeout)
+
+    async def connect(self, addr, authkey=b"", timeout=None):
+        if timeout is None:
+            timeout = self._client_defualt_timeout
+        await self.sendto(RPC_CONNECT + authkey, addr)
+        fut = self.createPending(addr, addr, timeout, False)
+        msgID = await fut
+        clientID, *_ = parseMsgID(msgID)
+        if clientID < 1024:
+            raise RPCError(f'Connect {addr} fail')
+        self.__clientID = clientID
+
+    async def ping(self, addr, timeout=1):
+        await self.sendto(RPC_PING, addr)
+        fut = self.createPending(addr, addr, timeout, False)
+        try:
+            return await fut
+        except RPCTimeout:
+            return False
+
+    async def request(self, address, msgID, msg):
+        log.debug(f'send request {address}, {msgID.hex()}, {msg}')
+        await self.sendto(RPC_REQUEST + msgID + msg, address)
+
+    async def shutdown(self, address, roleAuth):
+        await self.sendto(RPC_SHUTDOWN + nextMsgID(self.clientID) + roleAuth,
+                          address)
+
+    def on_welcome(self, source, msg):
+        if source in self.pending:
+            fut, timeout = self.pending[source]
+            timeout.cancel()
+            if not fut.done():
+                fut.set_result(msg)
+
+    def on_pong(self, source, msg):
+        log.debug(f"received pong from {source}")
+        if source in self.pending:
+            fut, timeout = self.pending[source]
+            timeout.cancel()
+            if not fut.done():
+                fut.set_result(True)
+
+    def on_response(self, source, msgID, msg):
+        """
+        Client side.
+        """
+        if msgID not in self.pending:
+            return
+        fut, timeout = self.pending[msgID]
+        timeout.cancel()
+        try:
+            result = unpack(msg)
+        except Exception as e:
+            fut.set_exception(e)
+            return
+        if not fut.done():
+            if isinstance(result, Exception):
+                fut.set_exception(result)
+            else:
+                fut.set_result(result)
+
+
+class RPCServerMixin(RPCMixin):
+    __tasks = None
+    __sessions = None
+    __nextClientID = 1024
+    __nextSessionID = 1024
+
+    def info(self):
+        return {'version': __version__}
+
+    @property
+    def nextClientID(self):
+        self.__nextClientID += 1
+        return self.__nextClientID
+
+    @property
+    def nextSessionID(self):
+        self.__nextSessionID += 1
+        return self.__nextSessionID
+
+    @property
+    def tasks(self):
+        if self.__tasks is None:
+            self.__tasks = {}
+        return self.__tasks
+
+    @property
+    def sessions(self):
+        if self.__sessions is None:
+            self.__sessions = {}
+        return self.__sessions
+
+    def createTask(self, msgID, coro, timeout=0):
+        """
+        Create a new task for msgID.
+        """
+        if timeout > 0:
+            coro = asyncio.wait_for(coro, timeout)
+        task = asyncio.ensure_future(coro, loop=self.loop)
+        self.tasks[msgID] = task
+
+        def clean(fut, msgID=msgID):
+            if msgID in self.tasks:
+                del self.tasks[msgID]
+
+        task.add_done_callback(clean)
+
+    def cancelTask(self, msgID):
+        """
+        Cancel the task for msgID.
+        """
+        if msgID in self.tasks:
+            self.tasks[msgID].cancel()
+
+    def createSession(self, clientID, obj):
+        sessionID = self.nextSessionID
+        self.sessions[(clientID, sessionID)] = obj
+        return sessionID
+
+    def removeSession(self, clientID, sessionID):
+        del self.sessions[(clientID, sessionID)]
+
+    def close(self):
+        self.stop()
+        for task in list(self.tasks.values()):
+            task.cancel()
+        self.tasks.clear()
+
+    def _unpack_request(self, msg):
+        try:
+            method, args, kw = unpack(msg)
+        except:
+            raise RPCError("Could not read packet: %r" % msg)
+        return method, args, kw
+
+    @property
+    def executor(self):
+        return None
+
+    @abstractmethod
+    def getRequestHandler(self, methodNane, source, msgID, args=(), kw={}):
+        """
+        Get suitable handler for request.
+
+        You should implement this method yourself.
+        """
+
+    def processResult(self, result, method, msgID):
+        return result
+
+    async def handle_request(self, source, msgID, method, args, kw):
+        """
+        Handle a request from source.
+        """
+        try:
+            func = self.getRequestHandler(method, source=source, msgID=msgID)
+            result = await self.callMethod(func, *args, **kw)
+            result = self.processResult(result, method, msgID)
+        except (RPCError, RPCServerError) as e:
+            result = e
+        except Exception as e:
+            result = RPCServerError.make(e)
+        msg = pack(result)
+        await self.response(source, msgID, msg)
+
+    async def callMethod(self, func, *args, **kw):
+        if 'timeout' in kw and not acceptArg(func, 'timeout'):
+            del kw['timeout']
+        if inspect.iscoroutinefunction(func):
+            result = await func(*args, **kw)
+        else:
+            result = await self.loop.run_in_executor(
+                self.executor, functools.partial(func, *args, **kw))
+            if isinstance(result, Awaitable):
+                result = await result
+        return result
+
+    async def response(self, address, msgID, msg):
+        log.debug(f'send response {address}, {msgID.hex()}, {msg}')
+        await self.sendto(RPC_RESPONSE + msgID + msg, address)
+
+    def on_connect(self, source, msg):
+        log.debug(f"connect from {source}")
+        if self.auth(source, msg):
+            clientID = self.nextClientID
+        else:
+            clientID = 0
+        msg = pack(self.info())
+        asyncio.ensure_future(
+            self.sendto(
+                RPC_WELCOME + nextMsgID(clientID),  # + msg,
+                source),
+            loop=self.loop)
+
+    def on_request(self, source, msgID, msg):
+        """
+        Received a request from source.
+        """
+        method, args, kw = self._unpack_request(msg)
+        self.createTask(msgID,
+                        self.handle_request(source, msgID, method, args, kw),
+                        timeout=kw.get('timeout', 0))
+
+    def on_shutdown(self, source, msgID, roleAuth):
+        if self.is_admin(source, roleAuth):
+            raise SystemExit(0)
+
+    def auth(self, source, authkey):
+        return True
+
+    def is_admin(self, source, roleAuth):
+        return True
+
+    def on_cancel(self, source, msgID, msg):
+        self.cancelTask(msgID)

qulab/sys/rpc/serialize.py

@@ -0,0 +1,139 @@
+import pickle
+import zlib
+from typing import Any, Callable, TypeVar
+
+import msgpack
+
+__index = 0
+__pack_handlers = {}
+__unpack_handlers = {}
+__compress_level = zlib.Z_NO_COMPRESSION
+
+cls = TypeVar('cls')
+
+
+def compress_level(level: int) -> None:
+    """
+    Set compress level
+
+    Args:
+        level: int
+            An integer from 0 to 9 controlling the level of compression;
+            1 is fastest and produces the least compression, 9 is slowest
+            and produces the most. 0 is no compression. The default value
+            is 0.
+    """
+    global __compress_level
+    __compress_level = level
+
+
+def register(cls: type,
+             encode: Callable[[cls], bytes] = pickle.dumps,
+             decode: Callable[[bytes], cls] = pickle.loads) -> None:
+    """
+    Register a serializable type
+
+    Args:
+        cls: type
+        encode: Callable
+            translate an object of type `cls` into `bytes`
+            default: pickle.dumps
+        decode: Callable
+            translate `bytes` to an object of type `cls`
+            default: pickle.loads
+    """
+    global __index
+    __index += 1
+    t = __index
+    __pack_handlers[cls] = (t, encode)
+    __unpack_handlers[t] = decode
+
+
+def default(obj: Any) -> msgpack.ExtType:
+    for cls, (t, encode) in __pack_handlers.items():
+        if isinstance(obj, cls):
+            return msgpack.ExtType(t, encode(obj))
+    else:
+        raise TypeError("Unknown type: %r" % (obj, ))
+
+
+def ext_hook(code: int, data: bytes) -> msgpack.ExtType:
+    for c, decode in __unpack_handlers.items():
+        if code == c:
+            return decode(data)
+    else:
+        return msgpack.ExtType(code, data)
+
+
+def pack(obj: Any) -> bytes:
+    """
+    Serialize
+    """
+    return msgpack.packb(obj, default=default, use_bin_type=True)
+
+
+def unpack(buff: bytes) -> Any:
+    """
+    Unserialize
+    """
+    return msgpack.unpackb(buff, ext_hook=ext_hook, raw=False)
+
+
+def packz(obj: Any) -> bytes:
+    """
+    Serialize and compress.
+    """
+    return zlib.compress(pack(obj), level=__compress_level)
+
+
+def unpackz(buff: bytes) -> Any:
+    """
+    Decompress and unserialize.
+    """
+    return unpack(zlib.decompress(buff))
+
+
+def encode_excepion(e: Exception) -> bytes:
+    e.__traceback__ = None
+    return pickle.dumps(e)
+
+
+register(Exception, encode_excepion)
+
+
+try:
+    import numpy as np
+
+    dtypes = [
+        np.bool_, np.int8, np.int16, np.int32, np.int64, np.uint8, np.uint16,
+        np.uint32, np.uint64, np.float16, np.float32, np.float64, np.complex64,
+        np.complex128
+    ]
+
+    _dtype_map1 = {t: i for i, t in enumerate(dtypes)}
+    _dtype_map2 = {i: t for i, t in enumerate(dtypes)}
+
+    def encode_ndarray(x: np.ndarray) -> bytes:
+        dtype = x.dtype
+        if isinstance(dtype, np.dtype):
+            dtype = dtype.type
+        if x.flags['F_CONTIGUOUS']:
+            x = x.T
+            T = True
+        else:
+            T = False
+        return pack((_dtype_map1[dtype], x.shape, T, x.data))
+
+    def decode_ndarray(buff: bytes) -> np.ndarray:
+        t, shape, T, buff = unpack(buff)
+        x = np.ndarray(shape, dtype=_dtype_map2[t], buffer=buff, order='C')
+        if T:
+            x = x.T
+        return x
+
+    register(np.ndarray, encode_ndarray, decode_ndarray)
+
+except:
+    pass
+
+__all__ = ['compress_level', 'register', 'pack', 'unpack', 'packz', 'unpackz']

qulab/sys/rpc/server.py

@@ -0,0 +1,29 @@
+from ..net.kcp import listen
+
+
+class Peer:
+
+    def __init__(self, addr, handler, **kwargs):
+        self.addr = addr
+        self.handler = handler
+        self.kwargs = kwargs
+
+    async def run(self):
+        conv = self.kwargs.get('conv', 0)
+        ttl = self.kwargs.get('ttl', 60)
+        async with listen(self.addr, conv, ttl, **self.kwargs) as server:
+            async for conn in server:
+                await self.handler(conn)
+
+    def __call__(self, handler):
+        self.handler = handler
+        return self
+
+    def __await__(self):
+        return self.run().__await__()
+
+    def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        pass

qulab/sys/rpc/socket.py

@@ -0,0 +1,29 @@
+from ..net.kcp import listen
+
+
+class Peer:
+
+    def __init__(self, addr, handler, **kwargs):
+        self.addr = addr
+        self.handler = handler
+        self.kwargs = kwargs
+
+    async def run(self):
+        conv = self.kwargs.get('conv', 0)
+        ttl = self.kwargs.get('ttl', 60)
+        async with listen(self.addr, conv, ttl, **self.kwargs) as server:
+            async for conn in server:
+                await self.handler(conn)
+
+    def __call__(self, handler):
+        self.handler = handler
+        return self
+
+    def __await__(self):
+        return self.run().__await__()
+
+    def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        pass

qulab/sys/rpc/utils.py

@@ -0,0 +1,25 @@
+import inspect
+import os
+from hashlib import sha1
+
+
+def acceptArg(f, name, keyword=True):
+    """
+    Test if argument is acceptable by function.
+
+    Args:
+        f: callable
+            function
+        name: str
+            argument name
+    """
+    sig = inspect.signature(f)
+    for param in sig.parameters.values():
+        if param.name == name and param.kind != param.VAR_POSITIONAL:
+            return True
+        elif param.kind == param.VAR_KEYWORD:
+            return True
+        elif param.kind == param.VAR_POSITIONAL and not keyword:
+            return True
+    return False
+