QuLab 2.1.0__cp310-cp310-macosx_10_9_universal2.whl → 2.1.2__cp310-cp310-macosx_10_9_universal2.whl

qulab/scan/recorder.py

@@ -1,18 +1,11 @@
 import asyncio
-import itertools
 import os
 import pickle
-import sys
 import time
-import uuid
-from collections import defaultdict
 from pathlib import Path
-from threading import Lock
-from types import EllipsisType
 
 import click
 import dill
-import numpy as np
 import zmq
 from loguru import logger
 
@@ -21,8 +14,7 @@ from qulab.sys.rpc.zmq_socket import ZMQContextManager
 from .curd import query_record, remove_tags, tag, update_tags
 from .models import Record as RecordInDB
 from .models import Session, create_engine, create_tables, sessionmaker, utcnow
-
-_notgiven = object()
+from .record import Record
 
 try:
     default_record_port = int(os.getenv('QULAB_RECORD_PORT', 6789))
@@ -36,487 +28,7 @@ else:
 datapath.mkdir(parents=True, exist_ok=True)
 
 record_cache = {}
-
-
-def random_path(base):
-    while True:
-        s = uuid.uuid4().hex
-        path = base / s[:2] / s[2:4] / s[4:6] / s[6:]
-        if not path.exists():
-            return path
-
-
-def index_in_slice(slice_obj: slice | int, index: int):
-    if isinstance(slice_obj, int):
-        return slice_obj == index
-    start, stop, step = slice_obj.start, slice_obj.stop, slice_obj.step
-    if start is None:
-        start = 0
-    if step is None:
-        step = 1
-    if stop is None:
-        stop = sys.maxsize
-
-    if step > 0:
-        return start <= index < stop and (index - start) % step == 0
-    else:
-        return stop < index <= start and (index - start) % step == 0
-
-
-class BufferList():
-
-    def __init__(self, file=None, slice=None):
-        self._list = []
-        self.lu = ()
-        self.rd = ()
-        self.inner_shape = ()
-        self.file = file
-        self._slice = slice
-        self._lock = Lock()
-        self._data_id = None
-
-    def __repr__(self):
-        return f"<BufferList: shape={self.shape}, lu={self.lu}, rd={self.rd}, slice={self._slice}>"
-
-    def __getstate__(self):
-        self.flush()
-        if isinstance(self.file, Path):
-            file = '/'.join(self.file.parts[-4:])
-        else:
-            file = self.file
-        return {
-            'file': file,
-            'lu': self.lu,
-            'rd': self.rd,
-            'inner_shape': self.inner_shape,
-        }
-
-    def __setstate__(self, state):
-        self.file = state['file']
-        self.lu = state['lu']
-        self.rd = state['rd']
-        self.inner_shape = state['inner_shape']
-        self._list = []
-        self._slice = None
-        self._lock = Lock()
-        self._data_id = None
-
-    @property
-    def shape(self):
-        return tuple([i - j
-                      for i, j in zip(self.rd, self.lu)]) + self.inner_shape
-
-    def flush(self):
-        if not self._list:
-            return
-        if isinstance(self.file, Path):
-            with self._lock:
-                with open(self.file, 'ab') as f:
-                    for item in self._list:
-                        dill.dump(item, f)
-                self._list.clear()
-
-    def append(self, pos, value, dims=None):
-        if dims is not None:
-            if any([p != 0 for i, p in enumerate(pos) if i not in dims]):
-                return
-            pos = tuple([pos[i] for i in dims])
-        self.lu = tuple([min(i, j) for i, j in zip(pos, self.lu)])
-        self.rd = tuple([max(i + 1, j) for i, j in zip(pos, self.rd)])
-        if hasattr(value, 'shape'):
-            if self.inner_shape is None:
-                self.inner_shape = value.shape
-            elif self.inner_shape != value.shape:
-                self.inner_shape = ()
-        self._list.append((pos, value))
-        if len(self._list) > 1000:
-            self.flush()
-
-    def _iter_file(self):
-        if isinstance(self.file, Path) and self.file.exists():
-            with self._lock:
-                with open(self.file, 'rb') as f:
-                    while True:
-                        try:
-                            pos, value = dill.load(f)
-                            yield pos, value
-                        except EOFError:
-                            break
-
-    def iter(self):
-        if self._data_id is None:
-            for pos, value in itertools.chain(self._iter_file(), self._list):
-                if not self._slice:
-                    yield pos, value
-                elif all(
-                    [index_in_slice(s, i) for s, i in zip(self._slice, pos)]):
-                    if self.inner_shape:
-                        yield pos, value[self._slice[len(pos):]]
-                    else:
-                        yield pos, value
-        else:
-            server, record_id, key = self._data_id
-            with ZMQContextManager(zmq.DEALER, connect=server) as socket:
-                socket.send_pyobj({
-                    'method': 'bufferlist_slice',
-                    'record_id': record_id,
-                    'key': key,
-                    'slice': self._slice
-                })
-                ret = socket.recv_pyobj()
-                yield from ret
-
-    def value(self):
-        d = []
-        for pos, value in self.iter():
-            d.append(value)
-        return d
-
-    def pos(self):
-        p = []
-        for pos, value in self.iter():
-            p.append(pos)
-        return p
-
-    def items(self):
-        p, d = [], []
-        for pos, value in self.iter():
-            p.append(pos)
-            d.append(value)
-        return p, d
-
-    def array(self):
-        pos, data = self.items()
-        if self._slice:
-            pos = np.asarray(pos)
-            lu = tuple(np.min(pos, axis=0))
-            rd = tuple(np.max(pos, axis=0) + 1)
-            pos = np.asarray(pos) - np.asarray(lu)
-            shape = []
-            for k, (s, i, j) in enumerate(zip(self._slice, rd, lu)):
-                if s.step is not None:
-                    pos[:, k] = pos[:, k] / s.step
-                    shape.append(round(np.ceil((i - j) / s.step)))
-                else:
-                    shape.append(i - j)
-            shape = tuple(shape)
-        else:
-            shape = tuple([i - j for i, j in zip(self.rd, self.lu)])
-            pos = np.asarray(pos) - np.asarray(self.lu)
-        data = np.asarray(data)
-        inner_shape = data.shape[1:]
-        x = np.full(shape + inner_shape, np.nan, dtype=data[0].dtype)
-        x.__setitem__(tuple(pos.T), data)
-        return x
-
-    def _full_slice(self, slice_tuple: slice
-                    | tuple[slice | int | EllipsisType, ...]):
-        ndim = len(self.lu)
-        if self.inner_shape:
-            ndim += len(self.inner_shape)
-
-        if isinstance(slice_tuple, slice):
-            slice_tuple = (
-                slice_tuple, ) + (slice(0, sys.maxsize, 1), ) * (ndim - 1)
-        if slice_tuple is Ellipsis:
-            slice_tuple = (slice(0, sys.maxsize, 1), ) * ndim
-        else:
-            head, tail = (), ()
-            for i, s in enumerate(slice_tuple):
-                if s is Ellipsis:
-                    head = slice_tuple[:i]
-                    tail = slice_tuple[i + 1:]
-                    break
-            else:
-                head = slice_tuple
-                tail = ()
-            slice_tuple = head + (slice(
-                0, sys.maxsize, 1), ) * (ndim - len(head) - len(tail)) + tail
-        slice_list = []
-        contract = []
-        reversed = []
-        for i, s in enumerate(slice_tuple):
-            if isinstance(s, int):
-                if s >= 0:
-                    slice_list.append(slice(s, s + 1, 1))
-                elif i < len(self.lu):
-                    s = self.rd[i] + s
-                    slice_list.append(slice(s, s + 1, 1))
-                else:
-                    slice_list.append(slice(s, s - 1, -1))
-                contract.append(i)
-            else:
-                start, stop, step = s.start, s.stop, s.step
-                if step is None:
-                    step = 1
-                if step < 0 and i < len(self.lu):
-                    step = -step
-                    reversed.append(i)
-                    if start is None and stop is None:
-                        start, stop = 0, sys.maxsize
-                    elif start is None:
-                        start, stop = self.lu[i], sys.maxsize
-                    elif stop is None:
-                        start, stop = 0, start + self.lu[i]
-                    else:
-                        start, stop = stop + self.lu[i] + 1, start + self.lu[
-                            i] + 1
-
-                if start is None:
-                    start = 0
-                elif start < 0 and i < len(self.lu):
-                    start = self.rd[i] + start
-                if step is None:
-                    step = 1
-                if stop is None:
-                    stop = sys.maxsize
-                elif stop < 0 and i < len(self.lu):
-                    stop = self.rd[i] + stop
-
-                slice_list.append(slice(start, stop, step))
-        return tuple(slice_list), contract, reversed
-
-    def __getitem__(self, slice_tuple: slice | EllipsisType
-                    | tuple[slice | int | EllipsisType, ...]):
-        self._slice, contract, reversed = self._full_slice(slice_tuple)
-        ret = self.array()
-        slices = []
-        for i, s in enumerate(self._slice):
-            if i in contract:
-                slices.append(0)
-            elif isinstance(s, slice):
-                if i in reversed:
-                    slices.append(slice(None, None, -1))
-                else:
-                    slices.append(slice(None, None, 1))
-        ret = ret.__getitem__(tuple(slices))
-        self._slice = None
-        return ret
-
-
-class Record():
-
-    def __init__(self, id, database, description=None):
-        from .scan import OptimizeSpace
-
-        self.id = id
-        self.database = database
-        self.description = description
-        self._keys = set()
-        self._items = {}
-        self._index = []
-        self._pos = []
-        self._last_vars = set()
-        self._file = None
-        self.independent_variables = {}
-        self.constants = {}
-        self.dims = {}
-
-        for name, value in self.description['consts'].items():
-            if name not in self._items:
-                self._items[name] = value
-            self.constants[name] = value
-            self.dims[name] = ()
-        for level, range_list in self.description['loops'].items():
-            for name, iterable in range_list:
-                if isinstance(iterable, OptimizeSpace):
-                    self.dims[name] = tuple(range(level + 1))
-                    continue
-                elif isinstance(iterable, (np.ndarray, list, tuple, range)):
-                    self._items[name] = iterable
-                    self.independent_variables[name] = iterable
-                self.dims[name] = (level, )
-
-        for level, group in self.description['order'].items():
-            for names in group:
-                for name in names:
-                    if name not in self.description['dependents']:
-                        if name not in self.dims:
-                            self.dims[name] = (level, )
-                    else:
-                        d = set()
-                        for n in self.description['dependents'][name]:
-                            d.update(self.dims[n])
-                        if name not in self.dims:
-                            self.dims[name] = tuple(sorted(d))
-                        else:
-                            self.dims[name] = tuple(
-                                sorted(set(self.dims[name]) | d))
-
-        if self.is_local_record():
-            self.database = Path(self.database)
-            self._file = random_path(self.database / 'objects')
-            self._file.parent.mkdir(parents=True, exist_ok=True)
-
-    def __getstate__(self) -> dict:
-        return {
-            'id': self.id,
-            'database': self.database,
-            'description': self.description,
-            '_keys': self._keys,
-            '_items': self._items,
-            '_index': self._index,
-            '_pos': self._pos,
-            '_last_vars': self._last_vars,
-            'independent_variables': self.independent_variables,
-            'constants': self.constants,
-            'dims': self.dims,
-        }
-
-    def __setstate__(self, state: dict):
-        self.id = state['id']
-        self.database = state['database']
-        self.description = state['description']
-        self._keys = state['_keys']
-        self._items = state['_items']
-        self._index = state['_index']
-        self._pos = state['_pos']
-        self._last_vars = state['_last_vars']
-        self.independent_variables = state['independent_variables']
-        self.constants = state['constants']
-        self.dims = state['dims']
-        self._file = None
-
-    def is_local_record(self):
-        return not self.is_cache_record() and not self.is_remote_record()
-
-    def is_cache_record(self):
-        return self.database is None
-
-    def is_remote_record(self):
-        return isinstance(self.database,
-                          str) and self.database.startswith("tcp://")
-
-    def __del__(self):
-        self.flush()
-
-    def __getitem__(self, key):
-        return self.get(key, buffer_to_array=True)
-
-    def get(self, key, default=_notgiven, buffer_to_array=False, slice=None):
-        if self.is_remote_record():
-            with ZMQContextManager(zmq.DEALER,
-                                   connect=self.database) as socket:
-                socket.send_pyobj({
-                    'method': 'record_getitem',
-                    'record_id': self.id,
-                    'key': key
-                })
-                ret = socket.recv_pyobj()
-                if isinstance(ret, BufferList):
-                    if buffer_to_array:
-                        socket.send_pyobj({
-                            'method': 'bufferlist_slice',
-                            'record_id': self.id,
-                            'key': key,
-                            'slice': slice
-                        })
-                        lst = socket.recv_pyobj()
-                        ret._list = lst
-                        ret._slice = slice
-                        return ret.array()
-                    else:
-                        ret._data_id = self.database, self.id, key
-                        return ret
-                else:
-                    return ret
-        else:
-            if default is _notgiven:
-                d = self._items.get(key)
-            else:
-                d = self._items.get(key, default)
-            if isinstance(d, BufferList):
-                if isinstance(d.file, str):
-                    d.file = self._file.parent.parent.parent.parent / d.file
-                d._slice = slice
-                if buffer_to_array:
-                    return d.array()
-                else:
-                    return d
-            else:
-                return d
-
-    def keys(self):
-        if self.is_remote_record():
-            with ZMQContextManager(zmq.DEALER,
-                                   connect=self.database) as socket:
-                socket.send_pyobj({
-                    'method': 'record_keys',
-                    'record_id': self.id
-                })
-                return socket.recv_pyobj()
-        else:
-            return list(self._keys)
-
-    def append(self, level, step, position, variables):
-        if level < 0:
-            self.flush()
-            return
-
-        for key in set(variables.keys()) - self._last_vars:
-            if key not in self.dims:
-                self.dims[key] = tuple(range(level + 1))
-
-        self._last_vars = set(variables.keys())
-        self._keys.update(variables.keys())
-
-        if level >= len(self._pos):
-            l = level + 1 - len(self._pos)
-            self._index.extend(([0] * (l - 1)) + [step])
-            self._pos.extend(([0] * (l - 1)) + [position])
-            pos = tuple(self._pos)
-        elif level == len(self._pos) - 1:
-            self._index[-1] = step
-            self._pos[-1] = position
-            pos = tuple(self._pos)
-        else:
-            self._index = self._index[:level + 1]
-            self._pos = self._pos[:level + 1]
-            self._index[-1] = step + 1
-            self._pos[-1] = position
-            pos = tuple(self._pos)
-            self._pos[-1] += 1
-
-        for key, value in variables.items():
-            if self.dims[key] == ():
-                if key not in self._items:
-                    self._items[key] = value
-            elif level == self.dims[key][-1]:
-                if key not in self._items:
-                    if self.is_local_record():
-                        bufferlist_file = random_path(self.database /
-                                                      'objects')
-                        bufferlist_file.parent.mkdir(parents=True,
-                                                     exist_ok=True)
-                        self._items[key] = BufferList(bufferlist_file)
-                    else:
-                        self._items[key] = BufferList()
-                    self._items[key].lu = pos
-                    self._items[key].rd = tuple([i + 1 for i in pos])
-                    self._items[key].append(pos, value, self.dims[key])
-                elif isinstance(self._items[key], BufferList):
-                    self._items[key].append(pos, value, self.dims[key])
-
-    def flush(self):
-        if self.is_remote_record() or self.is_cache_record():
-            return
-
-        for key, value in self._items.items():
-            if isinstance(value, BufferList):
-                value.flush()
-
-        with open(self._file, 'wb') as f:
-            dill.dump(self, f)
-
-    def __repr__(self):
-        return f"<Record: id={self.id} app={self.description['app']}, keys={self.keys()}>"
-
-    # def _repr_html_(self):
-    #     return f"""
-    #     <h3>Record: id={self.id}, app={self.description['app']}</h3>
-    #     <p>keys={self.keys()}</p>
-    #     <p>dims={self.dims}</p>
-    #     """
+CACHE_SIZE = 1024
 
 
 class Request():
@@ -534,11 +46,11 @@ async def reply(req, resp):
 
 
 def clear_cache():
-    if len(record_cache) < 1024:
+    if len(record_cache) < CACHE_SIZE:
         return
 
     for k, (t, _) in zip(sorted(record_cache.items(), key=lambda x: x[1][0]),
-                         range(len(record_cache) - 1024)):
+                         range(len(record_cache) - CACHE_SIZE)):
         del record_cache[k]
 
 
@@ -547,14 +59,20 @@ def flush_cache():
         r.flush()
 
 
+def get_local_record(session: Session, id: int, datapath: Path) -> Record:
+    record_in_db = session.get(RecordInDB, id)
+    record_in_db.atime = utcnow()
+    path = datapath / 'objects' / record_in_db.file
+    with open(path, 'rb') as f:
+        record = dill.load(f)
+        record.database = datapath
+        record._file = path
+    return record
+
+
 def get_record(session: Session, id: int, datapath: Path) -> Record:
     if id not in record_cache:
-        record_in_db = session.get(RecordInDB, id)
-        record_in_db.atime = utcnow()
-        path = datapath / 'objects' / record_in_db.file
-        with open(path, 'rb') as f:
-            record = dill.load(f)
-            record._file = path
+        record = get_local_record(session, id, datapath)
     else:
         record = record_cache[id][1]
     clear_cache()
@@ -570,6 +88,7 @@ def record_create(session: Session, description: dict, datapath: Path) -> int:
     if 'tags' in description:
         record_in_db.tags = [tag(session, t) for t in description['tags']]
     record_in_db.file = '/'.join(record._file.parts[-4:])
+    record._file = datapath / 'objects' / record_in_db.file
     session.add(record_in_db)
     try:
         session.commit()
@@ -618,7 +137,7 @@ async def handle(session: Session, request: Request, datapath: Path):
                           msg['position'], msg['variables'], datapath)
         case 'record_description':
             record = get_record(session, msg['record_id'], datapath)
-            await reply(request, dill.dumps(record.description))
+            await reply(request, dill.dumps(record))
         case 'record_getitem':
             record = get_record(session, msg['record_id'], datapath)
             await reply(request, record.get(msg['key'], buffer_to_array=False))