QuLab 2.0.2__cp310-cp310-macosx_10_9_universal2.whl

qulab/scan/utils.py

@@ -0,0 +1,37 @@
+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):]
+    }
+    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)
+    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

qulab/storage/__main__.py

@@ -0,0 +1,51 @@
+import click
+
+
+@click.command()
+@click.option('--executor', default='', help='Executor address to use.')
+@click.option('--port', default=8080, help='Port to run the server on.')
+@click.option('--host', default='127.0.0.1', help='Host to run the server on.')
+@click.option('--db-url', default=None, help='Database URL to use.')
+@click.option('--data-path',
+              default='waveforms/data',
+              help='Path to the data directory.')
+@click.option('--debug', is_flag=True, help='Run in debug mode.')
+@click.option('--workers',
+              default=1,
+              help='Number of workers to run the server with.')
+@click.option('--timeout', default=60, help='Timeout for requests.')
+@click.option('--log-level',
+              default='INFO',
+              help='Log level to run the server with.')
+@click.option('--log-file',
+              default='/var/log/waveforms/server.log',
+              help='Log file to run the server with.')
+@click.option('--log-format',
+              default='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+              help='Log format to run the server with.')
+def main(executor, port, host, db_url, data_path, debug, workers, timeout,
+         log_level, log_file, log_format):
+    """
+    Main entry point for the server.
+    """
+    from waveforms.server import create_app
+
+    app = create_app(
+        executor=executor,
+        port=port,
+        host=host,
+        db_url=db_url,
+        data_path=data_path,
+        debug=debug,
+        workers=workers,
+        timeout=timeout,
+        log_level=log_level,
+        log_file=log_file,
+        log_format=log_format,
+    )
+    app.run()
+
+
+if __name__ == '__main__':
+    main()
+    

qulab/storage/backend/redis.py

@@ -0,0 +1,204 @@
+import dill
+import numpy as np
+import redis
+
+
+def try_dumps(value):
+    try:
+        value_b = dill.dumps(value)
+    except:
+        value_b = dill.dumps(str(value))
+    finally:
+        return value_b
+
+
+class redisClient(object):
+
+    def __init__(self,
+                 name,
+                 server=None,
+                 addr='redis://localhost:6379/0',
+                 expire_time=172800):
+        self._r = redis.Redis.from_url(addr) if server is None else server
+        self.name = name
+        self.expire_time = int(expire_time)  # 默认数据两天过期，单位秒
+
+    def delete(self):
+        self._r.delete(self.name)
+
+
+class redisString(redisClient):
+    '''Redis String client'''
+
+    def set(self, value):
+        if value is not None:
+            value_b = try_dumps(value)
+            self._r.set(self.name, value_b, ex=self.expire_time)
+
+    def get(self):
+        v_b = self._r.get(self.name)
+        value = dill.loads(v_b) if v_b is not None else None
+        return value
+
+    @property
+    def data(self):
+        return self.get()
+
+
+class redisList(redisClient):
+    '''Redis List client'''
+
+    def add(self, *arg):
+        arg_b = [dill.dumps(i) for i in arg]
+        self._r.rpush(self.name, *arg_b)
+        self._r.expire(self.name, self.expire_time)
+
+    def read(self, start=0, end=-1):
+        data_b = self._r.lrange(self.name, start, end)
+        data = [dill.loads(i) for i in data_b]
+        return data
+
+    @property
+    def size(self):
+        return self._r.llen(self.name)
+
+    @property
+    def data(self):
+        return self.read()
+
+
+class redisSet(redisClient):
+    '''Redis Set client'''
+
+    def add(self, *arg):
+        arg_b = {dill.dumps(i) for i in arg}
+        self._r.sadd(self.name, *arg_b)
+        self._r.expire(self.name, self.expire_time)
+
+    def read(self):
+        data_b = self._r.smembers(self.name)
+        data = {dill.loads(i) for i in data_b}
+        return data
+
+    @property
+    def size(self):
+        return self._r.scard(self.name)
+
+    @property
+    def data(self):
+        return self.read()
+
+
+class redisZSet(redisClient):
+    '''有序集合'''
+
+    def __init__(self,
+                 name,
+                 server=None,
+                 addr='redis://localhost:6379/0',
+                 expire_time=172800):
+        super().__init__(name, server, addr, expire_time)
+        self.__score = 0
+
+    def delete(self):
+        super().delete()
+        self.__score = 0
+
+    def add(self, *elements):
+        mapping = {}
+        for ele in elements:
+            ele_b = dill.dumps(ele)
+            self.__score += 1
+            mapping.update({ele_b: self.__score})
+        self._r.zadd(self.name, mapping, nx=True)  # 只添加新元素
+        self._r.expire(self.name, self.expire_time)
+
+    def read(self, start=0, end=-1):
+        data_b = self._r.zrange(self.name, start, end, withscores=False)
+        data = [dill.loads(i) for i in data_b]
+        return data
+
+    @property
+    def size(self):
+        return self._r.zcard(self.name)
+
+    @property
+    def data(self):
+        return self.read()
+
+
+class redisHash(redisClient):
+    '''Redis Hash client'''
+
+    def add(self, **kw):
+        kw_b = {k: try_dumps(v) for k, v in kw.items()}
+        self._r.hmset(self.name, kw_b)
+        self._r.expire(self.name, self.expire_time)
+
+    def read(self):
+        data_b = self._r.hgetall(self.name)
+        data = {k_b.decode(): dill.loads(v_b) for k_b, v_b in data_b.items()}
+        return data
+
+    def get(self, key):
+        '''读取Hash中的一个key'''
+        v_b = self._r.hget(self.name, key)
+        value = dill.loads(v_b) if v_b is not None else None
+        return value
+
+    @property
+    def size(self):
+        return self._r.hlen(self.name)
+
+    @property
+    def data(self):
+        return self.read()
+
+
+class redisArray(redisClient):
+    '''Redis np.array client'''
+
+    def __init__(self,
+                 name,
+                 server=None,
+                 addr='redis://localhost:6379/0',
+                 expire_time=172800,
+                 dtype='complex128'):
+        super().__init__(name, server, addr, expire_time)
+        _r_dtype = self._r.get(f'{name}.dtype')
+        if _r_dtype is None:
+            self._r.set(f'{name}.dtype', dtype, ex=self.expire_time)
+            self.dtype = dtype
+        else:
+            self.dtype = _r_dtype
+
+    def delete(self):
+        self._r.delete(self.name)
+        self._r.delete(f'{self.name}.dtype')
+
+    def add(self, *args):
+        for arg in args:
+            buf = np.asarray(arg).astype(self.dtype).tobytes()
+            # self._r.append(self.name, buf)
+            self._r.rpush(self.name, buf)
+        self._r.expire(self.name, self.expire_time)
+
+    def read(self):
+        # buf = self._r.get(self.name)
+        buf_list = self._r.lrange(self.name, 0, -1)
+        buf = b''.join(buf_list)
+        data = np.frombuffer(buf,
+                             dtype=self.dtype) if buf is not None else None
+        return data
+
+    @property
+    def size(self):
+        array = self.data
+        if array is None:
+            return 0
+        else:
+            return array.size
+
+    @property
+    def data(self):
+        return self.read()

qulab/storage/base_dataset.py

@@ -0,0 +1,352 @@
+import bisect
+from concurrent.futures import Future
+from datetime import datetime
+from itertools import chain
+from multiprocessing import Lock
+from queue import Queue
+from typing import Any, Sequence
+
+from ..scan.base import StepStatus, Tracker, _get_all_dependence
+
+_NODEFAULT = object()
+
+
+class BaseDataset(Tracker):
+    """
+    A tracker that stores the results of the steps.
+
+    Parameters
+    ----------
+    data : dict
+        The data of the results.
+    shape : tuple
+        The shape of the results.
+    ctime : datetime.datetime
+        The creation time of the tracker.
+    mtime : datetime.datetime
+        The modification time of the tracker.
+    """
+
+    def __init__(
+            self,
+            data: dict = None,
+            shape: tuple = (),
+            save_kwds: bool | Sequence[str] = True,
+            frozen_keys: tuple = (),
+            ignores: tuple = (),
+    ):
+        self.ctime = datetime.utcnow()
+        self.mtime = datetime.utcnow()
+        self.data = data if data is not None else {}
+        self.cache = {}
+        self.pos = {}
+        self.timestamps = {}
+        self.iteration = {}
+        self._init_keys = list(self.data.keys())
+        self._frozen_keys = frozen_keys
+        self._ignores = ignores
+        self._key_levels = ()
+        self.depends = {}
+        self.dims = {}
+        self.vars_dims = {}
+        self.shape = shape
+        self.count = 0
+        self.save_kwds = save_kwds
+        self.queue = Queue()
+        self._queue_buffer = None
+        self._lock = Lock()
+
+    def init(self, loops: dict, functions: dict, constants: dict, graph: dict,
+             order: list):
+        """
+        Initialize the tracker.
+
+        Parameters
+        ----------
+        loops : dict
+            The map of iterables.
+        functions : dict
+            The map of functions.
+        constants : dict
+            The map of constants.
+        graph : dict
+            The dependence graph.
+        order : list
+            The order of the dependence graph.
+        """
+        from numpy import ndarray
+
+        self.depends = graph
+
+        for level, (keys, iters) in enumerate(loops.items()):
+            self._key_levels = self._key_levels + ((keys, level), )
+            if isinstance(keys, str):
+                keys = (keys, )
+                iters = (iters, )
+            if (len(keys) > 1 and len(iters) == 1
+                    and isinstance(iters[0], ndarray) and iters[0].ndim == 2
+                    and iters[0].shape[1] == len(keys)):
+                iters = iters[0]
+                for i, key in enumerate(keys):
+                    self.data[key] = iters[:, i]
+                    self._frozen_keys = self._frozen_keys + (key, )
+                    self._init_keys.append(key)
+                continue
+            if not isinstance(iters, tuple) or len(keys) != len(iters):
+                continue
+            for key, iter in zip(keys, iters):
+                if self.depends.get(key, set()):
+                    dims = set()
+                    for dep in self.depends[key]:
+                        if dep in self.vars_dims:
+                            dims.update(set(self.vars_dims[dep]))
+                    dims.add(level)
+                    self.vars_dims[key] = tuple(sorted(dims))
+                else:
+                    self.vars_dims[key] = (level, )
+                if level not in self.dims:
+                    self.dims[level] = ()
+                self.dims[level] = self.dims[level] + (key, )
+                if key not in self.data and isinstance(iter,
+                                                       (list, range, ndarray)):
+                    if key.startswith('__'):
+                        continue
+                    self.data[key] = iter
+                    self._frozen_keys = self._frozen_keys + (key, )
+                    self._init_keys.append(key)
+
+        for key, value in constants.items():
+            if key.startswith('__'):
+                continue
+            self.data[key] = value
+            self._init_keys.append(key)
+            self.vars_dims[key] = ()
+
+        for ready in order:
+            for key in ready:
+                if key in functions:
+                    deps = _get_all_dependence(key, graph)
+                    dims = set()
+                    for k in deps:
+                        dims.update(set(self.vars_dims.get(k, ())))
+                    self.vars_dims[key] = tuple(sorted(dims))
+
+        for k, v in self.vars_dims.items():
+            if len(v) == 1:
+                if v[0] in self.dims and k not in self.dims[v[0]]:
+                    self.dims[v[0]] = self.dims[v[0]] + (k, )
+                elif v[0] not in self.dims:
+                    self.dims[v[0]] = (k, )
+
+    def feed(self,
+             step: StepStatus,
+             dataframe: dict | Future,
+             store=False,
+             **options):
+        """
+        Feed the results of the step to the dataset.
+
+        Parameters
+        ----------
+        step : StepStatus
+            The step.
+        dataframe : dict
+            The results of the step.
+        """
+        import numpy as np
+
+        if not store:
+            return
+        self.mtime = datetime.utcnow()
+        if not self.shape:
+            self.shape = tuple([i + 1 for i in step.pos])
+        else:
+            self.shape = tuple(
+                [max(i + 1, j) for i, j in zip(step.pos, self.shape)])
+        if self.save_kwds:
+            if isinstance(self.save_kwds, bool):
+                kwds = step.kwds
+            else:
+                kwds = {
+                    key: step.kwds.get(key, np.nan)
+                    for key in self.save_kwds
+                }
+        else:
+            kwds = {}
+
+        if isinstance(dataframe, dict):
+            dataframe = self._prune(dataframe)
+        self.queue.put_nowait(
+            (step.iteration, step.pos, dataframe, kwds, self.mtime))
+        self.flush()
+
+    def _prune(self, dataframe: dict[str, Any]) -> dict[str, Any]:
+        return {
+            k: v
+            for k, v in dataframe.items() if k not in self._ignores
+            and k not in self._frozen_keys and not k.startswith('__')
+        }
+
+    def _append(self, iteration, pos, dataframe, kwds, now):
+        for k, v in chain(kwds.items(), dataframe.items()):
+            if k in self._frozen_keys or k in self._ignores:
+                continue
+            if k.startswith('__'):
+                continue
+            if self.vars_dims.get(k, ()) == () and k not in dataframe:
+                continue
+            self.count += 1
+            if k not in self.data:
+                self.data[k] = [v]
+                if k in self.vars_dims:
+                    self.pos[k] = tuple([pos[i]] for i in self.vars_dims[k])
+                else:
+                    self.pos[k] = tuple([i] for i in pos)
+                self.timestamps[k] = [now.timestamp()]
+                self.iteration[k] = [iteration]
+            else:
+                if k in self.vars_dims:
+                    pos_k = tuple(pos[i] for i in self.vars_dims[k])
+                    if k not in dataframe and pos_k in zip(*self.pos[k]):
+                        continue
+                    for i, l in zip(pos_k, self.pos[k]):
+                        l.append(i)
+                else:
+                    for i, l in zip(pos, self.pos[k]):
+                        l.append(i)
+                self.timestamps[k].append(now.timestamp())
+                self.iteration[k].append(iteration)
+                self.data[k].append(v)
+
+    def flush(self, block=False):
+        with self._lock:
+            self._flush(block=block)
+
+    def _dataframe_done(self, dataframe: Future | dict) -> bool:
+        if isinstance(dataframe, Future):
+            return dataframe.done()
+        else:
+            return all(x.done() for x in dataframe.values()
+                       if isinstance(x, Future))
+
+    def _dataframe_result(self, dataframe: Future | dict) -> dict:
+        if isinstance(dataframe, Future):
+            return dataframe.result()
+        else:
+            return {
+                k: v.result() if isinstance(v, Future) else v
+                for k, v in dataframe.items()
+            }
+
+    def _flush(self, block=False):
+        if self._queue_buffer is not None:
+            iteration, pos, dataframe, kwds, now = self._queue_buffer
+            if self._dataframe_done(dataframe) or block:
+                self._append(iteration, pos,
+                             self._prune(self._dataframe_result(dataframe)),
+                             kwds, now)
+                self._queue_buffer = None
+            else:
+                return
+        while not self.queue.empty():
+            iteration, pos, dataframe, kwds, now = self.queue.get()
+            if not self._dataframe_done(dataframe) and not block:
+                self._queue_buffer = (iteration, pos, dataframe, kwds, now)
+                return
+            else:
+                self._append(iteration, pos,
+                             self._prune(self._dataframe_result(dataframe)),
+                             kwds, now)
+
+    def _get_array(self, key, shape, count):
+        import numpy as np
+
+        if key in self.vars_dims:
+            shape = tuple([shape[i] for i in self.vars_dims[key]])
+
+        data, data_shape, data_count = self.cache.get(key, (None, (), 0))
+        if (data_shape, data_count) == (shape, count):
+            return data
+        try:
+            tmp = np.asarray(self.data[key])
+            if data_shape != shape:
+                data = np.full(shape + tmp.shape[1:], np.nan, dtype=tmp.dtype)
+        except:
+            tmp = self.data[key]
+            if data_shape != shape:
+                data = np.full(shape, np.nan, dtype=object)
+        try:
+            data[self.pos[key]] = tmp
+        except:
+            raise
+        self.cache[key] = (data, shape, count)
+        return data
+
+    def _get_part(self, key, skip):
+        i = bisect.bisect_left(self.iteration[key], skip)
+        pos = tuple(p[i:] for p in self.pos[key])
+        iteration = self.iteration[key][i:]
+        data = self.data[key][i:]
+        return data, iteration, pos
+
+    def keys(self):
+        """
+        Get the keys of the dataset.
+        """
+        self.flush()
+        return self.data.keys()
+
+    def values(self):
+        """
+        Get the values of the dataset.
+        """
+        self.flush()
+        return [self[k] for k in self.data]
+
+    def items(self):
+        """
+        Get the items of the dataset.
+        """
+        self.flush()
+        return list(zip(self.keys(), self.values()))
+
+    def get(self, key, default=_NODEFAULT, skip=None, block=False):
+        """
+        Get the value of the dataset.
+        """
+        self.flush(block)
+        if key in self._init_keys:
+            return self.data[key]
+        elif key in self.data:
+            if skip is None:
+                return self._get_array(key, self.shape, self.count)
+            else:
+                return self._get_part(key, skip)
+        elif default is _NODEFAULT:
+            raise KeyError(key)
+        else:
+            return default
+
+    def __getitem__(self, key):
+        return self.get(key)
+
+    def __getstate__(self):
+        self.flush()
+        data = dict(self.items())
+        return {
+            'data': data,
+            'pos': self.pos,
+            'timestamps': self.timestamps,
+            'iteration': self.iteration,
+            'depends': self.depends,
+            'shape': self.shape,
+            'dims': self.dims,
+            'vars_dims': self.vars_dims,
+            'ctime': self.ctime,
+            'mtime': self.mtime,
+            '_init_keys': self._init_keys,
+            '_frozen_keys': self._frozen_keys,
+            '_ignores': self._ignores,
+            '_key_levels': self._key_levels,
+            'save_kwds': self.save_kwds,
+        }

qulab/storage/chunk.py

@@ -0,0 +1,60 @@
+import hashlib
+import zlib
+from pathlib import Path
+
+DATAPATH = Path.home() / 'data'
+CHUNKSIZE = 1024 * 1024 * 4  # 4 MB
+
+
+def set_data_path(base_path: str) -> None:
+    global DATAPATH
+    DATAPATH = Path(base_path)
+
+
+def get_data_path() -> Path:
+    return DATAPATH
+
+
+def save_chunk(data: bytes, compressed: bool = False) -> tuple[str, str]:
+    if compressed:
+        data = zlib.compress(data)
+    hashstr = hashlib.sha1(data).hexdigest()
+    file = get_data_path(
+    ) / 'chunks' / hashstr[:2] / hashstr[2:4] / hashstr[4:]
+    file.parent.mkdir(parents=True, exist_ok=True)
+    with open(file, 'wb') as f:
+        f.write(data)
+    return str('/'.join(file.parts[-4:])), len(data)
+
+
+def load_chunk(file: str, compressed: bool = False) -> bytes:
+    if file.startswith('chunks/'):
+        with open(get_data_path() / file, 'rb') as f:
+            data = f.read()
+    elif file.startswith('packs/'):
+        *filepath, start, size = file.split('/')
+        filepath = '/'.join(filepath)
+        with open(get_data_path() / filepath, 'rb') as f:
+            f.seek(int(start))
+            data = f.read(int(size))
+    else:
+        raise ValueError('Invalid file path: ' + file)
+    if compressed:
+        data = zlib.decompress(data)
+    return data
+
+
+def pack_chunk(pack: str, chunkfile: str) -> str:
+    pack = get_data_path() / 'packs' / pack
+    pack.parent.mkdir(parents=True, exist_ok=True)
+    with open(pack, 'ab') as f:
+        buf = load_chunk(chunkfile)
+        start = f.tell()
+        size = len(buf)
+        f.write(buf)
+    return str('/'.join(pack.parts[-2:])) + '/' + str(start) + '/' + str(size)
+
+
+def delete_chunk(file: str):
+    file = get_data_path() / file
+    file.unlink()