QuLab 2.0.2__cp312-cp312-macosx_10_9_universal2.whl → 2.0.4__cp312-cp312-macosx_10_9_universal2.whl

qulab/scan/server.py

@@ -0,0 +1,106 @@
+import asyncio
+import pickle
+import sys
+import time
+import uuid
+from pathlib import Path
+from .scan import Scan
+import click
+import dill
+import numpy as np
+import zmq
+from loguru import logger
+
+from qulab.sys.rpc.zmq_socket import ZMQContextManager
+
+pool = {}
+
+class Request():
+    __slots__ = ['sock', 'identity', 'msg', 'method']
+
+    def __init__(self, sock, identity, msg):
+        self.sock = sock
+        self.identity = identity
+        self.msg = pickle.loads(msg)
+        self.method = self.msg.get('method', '')
+
+
+async def reply(req, resp):
+    await req.sock.send_multipart([req.identity, pickle.dumps(resp)])
+
+
+@logger.catch
+async def handle(request: Request):
+
+    msg = request.msg
+
+    match request.method:
+        case 'ping':
+            await reply(request, 'pong')
+        case 'submit':
+            description = dill.loads(msg['description'])
+            task = Scan()
+            task.description = description
+            task.start()
+            pool[task.id] = task
+            await reply(request, task.id)
+        case 'get_record_id':
+            task = pool.get(msg['id'])
+            for _ in range(10):
+                if task.record:
+                    await reply(request, task.record.id)
+                    break
+                await asyncio.sleep(1)
+            else:
+                await reply(request, None)
+        case _:
+            logger.error(f"Unknown method: {msg['method']}")
+
+
+async def _handle(request: Request):
+    try:
+        await handle(request)
+    except:
+        await reply(request, 'error')
+
+
+async def serv(port):
+    logger.info('Server starting.')
+    async with ZMQContextManager(zmq.ROUTER, bind=f"tcp://*:{port}") as sock:
+        logger.info('Server started.')
+        while True:
+            identity, msg = await sock.recv_multipart()
+            req = Request(sock, identity, msg)
+            asyncio.create_task(_handle(req))
+
+
+async def watch(port, timeout=1):
+    with ZMQContextManager(zmq.DEALER,
+                           connect=f"tcp://127.0.0.1:{port}") as sock:
+        sock.setsockopt(zmq.LINGER, 0)
+        while True:
+            try:
+                sock.send_pyobj({"method": "ping"})
+                if sock.poll(int(1000 * timeout)):
+                    sock.recv()
+                else:
+                    raise asyncio.TimeoutError()
+            except (zmq.error.ZMQError, asyncio.TimeoutError):
+                return asyncio.create_task(serv(port))
+            await asyncio.sleep(timeout)
+
+
+async def main(port, timeout=1):
+    task = await watch(port=port, timeout=timeout)
+    await task
+
+
+@click.command()
+@click.option('--port', default=6788, help='Port of the server.')
+@click.option('--timeout', default=1, help='Timeout of ping.')
+def server(port, timeout):
+    asyncio.run(main(port, timeout))
+
+
+if __name__ == "__main__":
+    server()

qulab/scan/utils.py

@@ -1,37 +1,83 @@
+import ast
+import asyncio
 import inspect
-from concurrent.futures import Future
-
-
-def call_func_with_kwds(func, args, kwds, log=None):
-    funcname = getattr(func, '__name__', repr(func))
-    sig = inspect.signature(func)
-    for p in sig.parameters.values():
-        if p.kind == p.VAR_KEYWORD:
-            return func(*args, **kwds)
-    kw = {
-        k: v
-        for k, v in kwds.items()
-        if k in list(sig.parameters.keys())[len(args):]
-    }
+from typing import Any, Callable
+
+from .expression import Env, Expression
+
+
+def is_valid_identifier(s: str) -> bool:
+    """
+    Check if a string is a valid identifier.
+    """
+    try:
+        ast.parse(f"f({s}=0)")
+        return True
+    except SyntaxError:
+        return False
+
+
+async def async_next(aiter):
+    try:
+        if hasattr(aiter, '__anext__'):
+            return await aiter.__anext__()
+        else:
+            return next(aiter)
+    except StopIteration:
+        raise StopAsyncIteration from None
+
+
+async def async_zip(*aiters):
+    aiters = [
+        ait.__aiter__() if hasattr(ait, '__aiter__') else iter(ait)
+        for ait in aiters
+    ]
+    try:
+        while True:
+            # 使用 asyncio.gather 等待所有异步生成器返回下一个元素
+            result = await asyncio.gather(*(async_next(ait) for ait in aiters))
+            yield tuple(result)
+    except StopAsyncIteration:
+        # 当任一异步生成器耗尽时停止迭代
+        return
+
+
+async def call_function(func: Callable | Expression, variables: dict[str,
+                                                                     Any]):
+    if isinstance(func, Expression):
+        env = Env()
+        for name in func.symbols():
+            if name in variables:
+                if inspect.isawaitable(variables[name]):
+                    variables[name] = await variables[name]
+                env.variables[name] = variables[name]
+            else:
+                raise ValueError(f'{name} is not provided.')
+        return func.eval(env)
+
     try:
-        args = [
-            arg.result() if isinstance(arg, Future) else arg for arg in args
-        ]
-        kw = {
-            k: v.result() if isinstance(v, Future) else v
-            for k, v in kw.items()
-        }
-        return func(*args, **kw)
+        sig = inspect.signature(func)
     except:
-        if log:
-            log.exception(f'Call {funcname} with {args} and {kw}')
-        raise
-    finally:
-        if log:
-            log.debug(f'Call {funcname} with {args} and {kw}')
-
-
-def try_to_call(x, args, kwds, log=None):
-    if callable(x):
-        return call_func_with_kwds(x, args, kwds, log)
-    return x
+        return func()
+    args = []
+    for name, param in sig.parameters.items():
+        if param.kind == param.POSITIONAL_OR_KEYWORD:
+            if name in variables:
+                if inspect.isawaitable(variables[name]):
+                    variables[name] = await variables[name]
+                args.append(variables[name])
+            elif param.default is not param.empty:
+                args.append(param.default)
+            else:
+                raise ValueError(f'parameter {name} is not provided.')
+        elif param.kind == param.VAR_POSITIONAL:
+            raise ValueError('not support VAR_POSITIONAL')
+        elif param.kind == param.VAR_KEYWORD:
+            ret = func(**variables)
+            if inspect.isawaitable(ret):
+                ret = await ret
+            return ret
+    ret = func(*args)
+    if inspect.isawaitable(ret):
+        ret = await ret
+    return ret

qulab/version.py

@@ -1 +1 @@
-__version__ = "2.0.2"
+__version__ = "2.0.4"