QuLab 2.1.2__cp311-cp311-win_amd64.whl → 2.1.4__cp311-cp311-win_amd64.whl

{QuLab-2.1.2.dist-info → QuLab-2.1.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: QuLab
-Version: 2.1.2
+Version: 2.1.4
 Summary: contral instruments and manage data
 Author-email: feihoo87 <feihoo87@gmail.com>
 Maintainer-email: feihoo87 <feihoo87@gmail.com>

{QuLab-2.1.2.dist-info → QuLab-2.1.4.dist-info}/RECORD

@@ -1,7 +1,7 @@
 qulab/__init__.py,sha256=vkFybY8YSsQilYdThPRD83-btPAR41sy_WCXiM-6mME,141
-qulab/__main__.py,sha256=XN2wrhlmEkTIPq_ZeSaO8rWXfYgD2Czkm9DVFVoCw_U,515
-qulab/fun.cp311-win_amd64.pyd,sha256=pDbiPNpPRjooF6cAGawfNqWkl4mbMhiar3MdLOTP0q4,31232
-qulab/version.py,sha256=FM98Inv2t7vU9Dvod2ccoaPCOtP2LV-rqmhbddOPf9U,21
+qulab/__main__.py,sha256=V7iokU7awstgjCeiF_hoOdFyrqJwC_4QetiLe7cWvOQ,454
+qulab/fun.cp311-win_amd64.pyd,sha256=DejsYo6DVNrnZP6HToBN86HY5ICqbUqWNxGnC0q5t8A,31232
+qulab/version.py,sha256=HVJYMyJeujBoXcNkL2sXW9oxWN275q3pivrlKw3Hm3s,21
 qulab/monitor/__init__.py,sha256=xEVDkJF8issrsDeLqQmDsvtRmrf-UiViFcGTWuzdlFU,43
 qulab/monitor/__main__.py,sha256=k2H1H5Zf9LLXTDLISJkbikLH-z0f1e5i5i6wXXYPOrE,105
 qulab/monitor/config.py,sha256=y_5StMkdrbZO1ziyKBrvIkB7Jclp9RCPK1QbsOhCxnY,785
@@ -13,15 +13,14 @@ qulab/monitor/ploter.py,sha256=dg7W28XTwEbBxHVtdPkFV135OQgoQwTi-NJCZQF-HYU,3724
 qulab/monitor/qt_compat.py,sha256=Eq7zlA4_XstB92NhtAqebtWU_Btw4lcwFO30YxZ-TPE,804
 qulab/monitor/toolbar.py,sha256=HxqG6ywKFyQJM2Q1s7SnhuzjbyeROczAZKwxztD1WJ8,8213
 qulab/scan/__init__.py,sha256=RR_0NQcr8Mi3vpWdypydbijQ1rXA0D3DEidQ7xjNslM,133
-qulab/scan/curd.py,sha256=bEXtcmiaoAv5APToXx5O5tvmqAhE2_LkvdwLsI_8D5M,4662
+qulab/scan/curd.py,sha256=yaTglGiS6mlk0GqDHi_w8T02XGBMvDZtXSdML7zDywk,7117
 qulab/scan/expression.py,sha256=vwUM9E0OFQal4bljlUtLR3NJu4zGRyuWYrdyZSs3QTU,16199
-qulab/scan/models.py,sha256=TkiVHF_fUZzYHs4MsCTRh391thpf4Ozd3R_LAU0Gxkg,17657
+qulab/scan/models.py,sha256=ZvXkJEt5Yz3Sjx0JKzYka-q2Uo-w_iVzHgH8A6DbjF0,18236
 qulab/scan/optimize.py,sha256=MlT4y422CnP961IR384UKryyZh8riNvrPSd2z_MXLEg,2356
-qulab/scan/query.py,sha256=vT1OevzKImqG40plmiPbyvzHVTBzAETo_OxpJHa0UmU,11821
-qulab/scan/record.py,sha256=s1bY4F_2FxIeN9ER-FILslCqc7i0JRAFv9pk0InQmcg,17477
-qulab/scan/recorder.py,sha256=qlUlxksL-beolTU3EupBlKDBGViicrq3ZoxjH6ZVHrY,8964
-qulab/scan/scan.py,sha256=AJ2GxdgTRL165UkNT27qZwJp_ag5j_hevNAY1-fUh7U,30607
-qulab/scan/server.py,sha256=Q7lX9ms37WbnzuLhrv3YbNrusTsxQUNT9slFqon5VL0,2872
+qulab/scan/query.py,sha256=WZeTiEtycP4raqzIUEP6tNg-OMaQYfRCxcdyYP5Swyg,11819
+qulab/scan/record.py,sha256=K8h8psiT_xgHZem9HRwzEgmr9lE0dGYHjlfK11sHldw,19111
+qulab/scan/scan.py,sha256=hBl9BjKfH-PoofhtpaD-O6s37wxy88hBN59GUXAcAuo,34877
+qulab/scan/server.py,sha256=e7AqFmaN9924VYrbvyKwPrVBzqPcnXbT8YwexAN_ZjQ,11911
 qulab/scan/space.py,sha256=OPceWIIrb2KDDQaSxD3Vkzkf9NDSuqWuQoDnKiWqtAo,5381
 qulab/scan/utils.py,sha256=30qnYvyFyctwcWxOCUpCNxXgGysA7xdIDzYbjwxGUzA,3744
 qulab/storage/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -71,7 +70,7 @@ qulab/sys/rpc/server.py,sha256=W3bPwe8um1IeR_3HLx-ad6iCcbeuUQcSg11Ze4w6DJg,742
 qulab/sys/rpc/socket.py,sha256=W3bPwe8um1IeR_3HLx-ad6iCcbeuUQcSg11Ze4w6DJg,742
 qulab/sys/rpc/utils.py,sha256=BurIcqh8CS-Hsk1dYP6IiefK4qHivaEqD9_rBY083SA,619
 qulab/sys/rpc/worker.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-qulab/sys/rpc/zmq_socket.py,sha256=Ep4Uy36mV04tkzKtLsRSu5dQ0AT7cPeLhqZxRYgaxTQ,8135
+qulab/sys/rpc/zmq_socket.py,sha256=SHMCp5oawYglfawuvDDTJJrQtxFpErlBtvoEwd61OO0,8362
 qulab/visualization/__init__.py,sha256=Bkt9AK5c45d6HFLlT-f8cIppywXziHtJqhDtVxOoKKo,6317
 qulab/visualization/__main__.py,sha256=WduINFl21B-XMsi2rg2cVhIyU11hKKX3zOsjc56QLiQ,1710
 qulab/visualization/_autoplot.py,sha256=JFLR3e71bryPaUkOCWFzD9C4Nk-ulqIdBb1YN4QnZic,14562
@@ -79,9 +78,9 @@ qulab/visualization/plot_layout.py,sha256=yAnMONOms7_szCdng-8wPpUMPis5UnbaNNzV4K
 qulab/visualization/plot_seq.py,sha256=h9D0Yl_yO64IwlvBgzMu9EBKr9gg6y8QE55gu2PfTns,2783
 qulab/visualization/qdat.py,sha256=HubXFu4nfcA7iUzghJGle1C86G6221hicLR0b-GqhKQ,5887
 qulab/visualization/widgets.py,sha256=HcYwdhDtLreJiYaZuN3LfofjJmZcLwjMfP5aasebgDo,3266
-QuLab-2.1.2.dist-info/LICENSE,sha256=b4NRQ-GFVpJMT7RuExW3NwhfbrYsX7AcdB7Gudok-fs,1086
-QuLab-2.1.2.dist-info/METADATA,sha256=BawKmH85rDyYCZUqZ1sEDQKzkUnpZX4D-Iy8iiY7AhQ,3609
-QuLab-2.1.2.dist-info/WHEEL,sha256=nSybvzWlmdJnHiUQSY-d7V1ycwEVUTqXiTvr2eshg44,102
-QuLab-2.1.2.dist-info/entry_points.txt,sha256=ohBzutEnQimP_BZWiuXdSliu4QAYSHHcN0PZD8c7ZCY,46
-QuLab-2.1.2.dist-info/top_level.txt,sha256=3T886LbAsbvjonu_TDdmgxKYUn939BVTRPxPl9r4cEg,6
-QuLab-2.1.2.dist-info/RECORD,,
+QuLab-2.1.4.dist-info/LICENSE,sha256=b4NRQ-GFVpJMT7RuExW3NwhfbrYsX7AcdB7Gudok-fs,1086
+QuLab-2.1.4.dist-info/METADATA,sha256=viS-wETLJCqd-ugZW2Ac_Z_pl1_BGi8KbBXhrlxj2qs,3609
+QuLab-2.1.4.dist-info/WHEEL,sha256=nSybvzWlmdJnHiUQSY-d7V1ycwEVUTqXiTvr2eshg44,102
+QuLab-2.1.4.dist-info/entry_points.txt,sha256=ohBzutEnQimP_BZWiuXdSliu4QAYSHHcN0PZD8c7ZCY,46
+QuLab-2.1.4.dist-info/top_level.txt,sha256=3T886LbAsbvjonu_TDdmgxKYUn939BVTRPxPl9r4cEg,6
+QuLab-2.1.4.dist-info/RECORD,,

qulab/__main__.py

@@ -1,7 +1,6 @@
 import click
 
 from .monitor.__main__ import main as monitor
-from .scan.recorder import record
 from .scan.server import server
 from .sys.net.cli import dht
 from .visualization.__main__ import plot
@@ -21,7 +20,6 @@ def hello():
 main.add_command(monitor)
 main.add_command(plot)
 main.add_command(dht)
-main.add_command(record)
 main.add_command(server)
 
 if __name__ == '__main__':

qulab/scan/curd.py

@@ -1,4 +1,7 @@
+import lzma
+import pickle
 from datetime import date, datetime, timezone
+from pathlib import Path
 from typing import Sequence, Type, Union
 
 from sqlalchemy.orm import Query, Session, aliased
@@ -6,7 +9,8 @@ from sqlalchemy.orm.exc import NoResultFound
 from sqlalchemy.orm.session import Session
 from waveforms.dicttree import foldDict
 
-from .models import Comment, Record, Report, Sample, Tag
+from .models import (Cell, Comment, Config, InputText, Notebook, Record,
+                     Report, Sample, Tag, utcnow)
 
 
 def tag(session: Session, tag_text: str) -> Tag:
@@ -142,3 +146,76 @@ def remove_tags(session: Session, record_id: int, tags: Sequence[str]):
         session.rollback()
         return False
     return True
+
+
+def create_notebook(session: Session, notebook_name: str) -> Notebook:
+    """Create a notebook in the database."""
+    notebook = Notebook(name=notebook_name)
+    session.add(notebook)
+    return notebook
+
+
+def create_input_text(session: Session, input_text: str) -> InputText:
+    """Create an input text in the database."""
+    input = InputText()
+    input.text = input_text
+    try:
+        input = session.query(InputText).filter(
+            InputText.hash == input.hash,
+            InputText.text_field == input_text).one()
+    except NoResultFound:
+        session.add(input)
+    return input
+
+
+def create_cell(session: Session, notebook: Notebook, input_text: str) -> Cell:
+    """Create a cell in the database."""
+    cell = Cell()
+    cell.notebook = notebook
+    cell.input = create_input_text(session, input_text)
+    cell.index = len(notebook.cells) - 1
+    session.add(cell)
+    notebook.atime = cell.ctime
+    return cell
+
+
+def create_config(session: Session, config: dict | bytes, base: Path,
+                  filename: str) -> Config:
+    """Create a config in the database."""
+
+    if not isinstance(config, bytes):
+        buf = pickle.dumps(config)
+        buf = lzma.compress(buf)
+        content_type = 'application/pickle+lzma'
+    else:
+        buf = config
+        content_type = 'application/octet-stream'
+    config = Config(buf)
+    config.content_type = content_type
+    for cfg in session.query(Config).filter(Config.hash == config.hash).all():
+        with open(base / cfg.file, 'rb') as f:
+            if f.read() == buf:
+                cfg.atime = utcnow()
+                return cfg
+    else:
+        path = base / filename
+        path.parent.mkdir(parents=True, exist_ok=True)
+        with open(path, 'wb') as f:
+            f.write(buf)
+        config.file = filename
+        session.add(config)
+    return config
+
+
+def get_config(session: Session, config_id: int, base: Path):
+    config = session.get(Config, config_id)
+    if config is None:
+        return None
+    config.atime = utcnow()
+    path = base / config.file
+    with open(path, 'rb') as f:
+        buf = f.read()
+    if config.content_type == 'application/pickle+lzma':
+        buf = lzma.decompress(buf)
+        buf = pickle.loads(buf)
+    return buf

qulab/scan/models.py

@@ -1,11 +1,9 @@
 import hashlib
 import pickle
-import time
 from datetime import datetime, timezone
 from functools import singledispatchmethod
-from typing import Optional
 
-from sqlalchemy import (JSON, Column, DateTime, Float, ForeignKey, Integer,
+from sqlalchemy import (Column, DateTime, Float, ForeignKey, Integer,
                         LargeBinary, String, Table, Text, create_engine)
 from sqlalchemy.orm import (backref, declarative_base, relationship,
                             sessionmaker)
@@ -325,7 +323,7 @@ class InputText(Base):
     __tablename__ = 'inputs'
 
     id = Column(Integer, primary_key=True)
-    hash = Column(LargeBinary(20))
+    hash = Column(LargeBinary(20), index=True)
     text_field = Column(Text, unique=True)
 
     @property
@@ -432,6 +430,22 @@ class SampleTransfer(Base):
     comments = relationship("Comment", secondary=sample_transfer_comments)
 
 
+class Config(Base):
+    __tablename__ = 'configs'
+
+    id = Column(Integer, primary_key=True)
+    hash = Column(LargeBinary(20), index=True)
+    file = Column(String)
+    content_type = Column(String, default='application/pickle')
+    ctime = Column(DateTime, default=utcnow)
+    atime = Column(DateTime, default=utcnow)
+
+    records = relationship("Record", back_populates="config")
+
+    def __init__(self, data: bytes) -> None:
+        self.hash = hashlib.sha1(data).digest()
+
+
 class Record(Base):
     __tablename__ = 'records'
 
@@ -440,14 +454,14 @@ class Record(Base):
     mtime = Column(DateTime, default=utcnow)
     atime = Column(DateTime, default=utcnow)
     user_id = Column(Integer, ForeignKey('users.id'))
+    config_id = Column(Integer, ForeignKey('configs.id'))
     parent_id = Column(Integer, ForeignKey('records.id'))
     cell_id = Column(Integer, ForeignKey('cells.id'))
 
     app = Column(String)
     file = Column(String)
+    content_type = Column(String, default='application/pickle')
     key = Column(String)
-    config = Column(JSON)
-    task_hash = Column(LargeBinary(32))
 
     parent = relationship("Record",
                           remote_side=[id],
@@ -456,6 +470,7 @@ class Record(Base):
                             remote_side=[parent_id],
                             back_populates="parent")
 
+    config = relationship("Config", back_populates="records")
     user = relationship("User")
     samples = relationship("Sample",
                            secondary=sample_records,

qulab/scan/query.py

@@ -12,7 +12,7 @@ from sqlalchemy.orm import sessionmaker
 from qulab.sys.rpc.zmq_socket import ZMQContextManager
 
 from .record import Record
-from .recorder import get_local_record
+from .server import get_local_record
 from .scan import default_server
 
 

qulab/scan/record.py

@@ -95,6 +95,11 @@ class BufferList():
                         dill.dump(item, f)
                 self._list.clear()
 
+    def delete(self):
+        if isinstance(self.file, Path):
+            self.file.unlink()
+            self.file = None
+
     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]):
@@ -450,6 +455,20 @@ class Record():
         with open(self._file, 'wb') as f:
             dill.dump(self, f)
 
+    def delete(self):
+        if self.is_remote_record():
+            with ZMQContextManager(zmq.DEALER,
+                                   connect=self.database) as socket:
+                socket.send_pyobj({
+                    'method': 'record_delete',
+                    'record_id': self.id
+                })
+        elif self.is_local_record():
+            for key, value in self._items.items():
+                if isinstance(value, BufferList):
+                    value.delete()
+            self._file.unlink()
+
     def export(self, file):
         with zipfile.ZipFile(file,
                              'w',
@@ -470,8 +489,33 @@ class Record():
                 else:
                     items[key] = value
             with z.open('record.pkl', 'w') as f:
+                self.description['entry']['scripts'] = self.scripts()
                 dill.dump((self.description, items), f)
 
+    def scripts(self, session=None):
+        scripts = self.description['entry']['scripts']
+        if isinstance(scripts, list):
+            return scripts
+        else:
+            cell_id = scripts
+
+        if self.is_remote_record():
+            with ZMQContextManager(zmq.DEALER,
+                                   connect=self.database) as socket:
+                socket.send_pyobj({
+                    'method': 'notebook_history',
+                    'cell_id': cell_id
+                })
+                return socket.recv_pyobj()
+        elif self.is_local_record():
+            from .models import Cell
+            assert session is not None, "session is required for local record"
+            cell = session.get(Cell, cell_id)
+            return [
+                cell.input.text
+                for cell in cell.notebook.cells[1:cell.index + 2]
+            ]
+
     @classmethod
     def load(cls, file: str):
         with zipfile.ZipFile(file, 'r') as z:

qulab/scan/scan.py

@@ -2,8 +2,12 @@ import asyncio
 import copy
 import inspect
 import itertools
+import lzma
 import os
+import pickle
+import platform
 import re
+import subprocess
 import sys
 import uuid
 from concurrent.futures import ProcessPoolExecutor
@@ -19,7 +23,7 @@ from ..sys.rpc.zmq_socket import ZMQContextManager
 from .expression import Env, Expression, Symbol
 from .optimize import NgOptimizer
 from .record import Record
-from .recorder import default_record_port
+from .server import default_record_port
 from .space import Optimizer, OptimizeSpace, Space
 from .utils import async_zip, call_function, dump_globals
 
@@ -41,6 +45,7 @@ except:
 
 __process_uuid = uuid.uuid1()
 __task_counter = itertools.count()
+__notebook_id = None
 
 if os.getenv('QULAB_SERVER'):
     default_server = os.getenv('QULAB_SERVER')
@@ -52,6 +57,105 @@ else:
     default_executor = default_server
 
 
+def yapf_reformat(cell_text):
+    try:
+        import isort
+        import yapf.yapflib.yapf_api
+
+        fname = f"f{uuid.uuid1().hex}"
+
+        def wrap(source):
+            lines = [f"async def {fname}():"]
+            for line in source.split('\n'):
+                lines.append("    " + line)
+            return '\n'.join(lines)
+
+        def unwrap(source):
+            lines = []
+            for line in source.split('\n'):
+                if line.startswith(f"async def {fname}():"):
+                    continue
+                lines.append(line[4:])
+            return '\n'.join(lines)
+
+        cell_text = re.sub('^%', '#%#', cell_text, flags=re.M)
+        reformated_text = unwrap(
+            yapf.yapflib.yapf_api.FormatCode(wrap(isort.code(cell_text)))[0])
+        return re.sub('^#%#', '%', reformated_text, flags=re.M)
+    except:
+        return cell_text
+
+
+def get_installed_packages():
+    result = subprocess.run([sys.executable, '-m', 'pip', 'freeze'],
+                            stdout=subprocess.PIPE,
+                            text=True)
+
+    lines = result.stdout.split('\n')
+    packages = []
+    for line in lines:
+        if line:
+            packages.append(line)
+    return packages
+
+
+def get_system_info():
+    info = {
+        'OS': platform.uname()._asdict(),
+        'Python': sys.version,
+        'PythonExecutable': sys.executable,
+        'PythonPath': sys.path,
+        'packages': get_installed_packages()
+    }
+    return info
+
+
+def current_notebook():
+    return __notebook_id
+
+
+async def create_notebook(name: str, database=default_server, socket=None):
+    global __notebook_id
+
+    async with ZMQContextManager(zmq.DEALER, connect=database,
+                                 socket=socket) as socket:
+        await socket.send_pyobj({'method': 'notebook_create', 'name': name})
+        __notebook_id = await socket.recv_pyobj()
+
+
+async def save_input_cells(notebook_id,
+                           input_cells,
+                           database=default_server,
+                           socket=None):
+    async with ZMQContextManager(zmq.DEALER, connect=database,
+                                 socket=socket) as socket:
+        await socket.send_pyobj({
+            'method': 'notebook_extend',
+            'notebook_id': notebook_id,
+            'input_cells': input_cells
+        })
+        return await socket.recv_pyobj()
+
+
+async def create_config(config: dict, database=default_server, socket=None):
+    async with ZMQContextManager(zmq.DEALER, connect=database,
+                                 socket=socket) as socket:
+        buf = lzma.compress(pickle.dumps(config))
+        await socket.send_pyobj({'method': 'config_update', 'update': buf})
+        return await socket.recv_pyobj()
+
+
+async def get_config(config_id: int, database=default_server, socket=None):
+    async with ZMQContextManager(zmq.DEALER, connect=database,
+                                 socket=socket) as socket:
+        await socket.send_pyobj({
+            'method': 'config_get',
+            'config_id': config_id
+        })
+        buf = await socket.recv_pyobj()
+        return pickle.loads(lzma.decompress(buf))
+
+
 def task_uuid():
     return uuid.uuid3(__process_uuid, str(next(__task_counter)))
 
@@ -135,10 +239,11 @@ class Scan():
                  mixin=None):
         self.id = task_uuid()
         self.record = None
-        self.namespace = {}
+        self.config = None
         self.description = {
             'app': app,
             'tags': tags,
+            'config': None,
             'loops': {},
             'intrinsic_loops': {},
             'consts': {},
@@ -157,9 +262,11 @@ class Scan():
             'database': database,
             'hiden': ['self', r'^__.*', r'.*__$'],
             'entry': {
+                'system': get_system_info(),
                 'env': {},
                 'shell': '',
-                'cmds': []
+                'cmds': [],
+                'scripts': []
             },
         }
         self._current_level = 0
@@ -248,8 +355,7 @@ class Scan():
 
     def emit(self, current_level, step, position, variables: dict[str, Any]):
         self._msg_queue.put_nowait(
-            asyncio.create_task(
-                self._emit(current_level, step, position, variables.copy())))
+            self._emit(current_level, step, position, variables.copy()))
 
     def hide(self, name: str):
         self.description['hiden'].append(name)
@@ -287,8 +393,6 @@ class Scan():
     def get(self, name: str):
         if name in self.description['consts']:
             return self.description['consts'][name]
-        elif name in self.namespace:
-            return self.namespace.get(name)
         else:
             return Symbol(name)
 
@@ -442,6 +546,17 @@ class Scan():
 
     async def run(self):
         assymbly(self.description)
+        if self.config:
+            self.description['config'] = await create_config(
+                self.config, self.description['database'], self._sock)
+        if current_notebook() is None:
+            await create_notebook('untitle', self.description['database'],
+                                  self._sock)
+        cell_id = await save_input_cells(current_notebook(),
+                                         self.description['entry']['scripts'],
+                                         self.description['database'],
+                                         self._sock)
+        self.description['entry']['scripts'] = cell_id
         if isinstance(
                 self.description['database'],
                 str) and self.description['database'].startswith("tcp://"):
@@ -635,13 +750,21 @@ def assymbly(description):
     ipy = get_ipython()
     if ipy is not None:
         description['entry']['shell'] = 'ipython'
-        description['entry']['cmds'] = ipy.user_ns['In']
+        description['entry']['scripts'] = [
+            yapf_reformat(cell_text) for cell_text in ipy.user_ns['In']
+        ]
     else:
         try:
             description['entry']['shell'] = 'shell'
             description['entry']['cmds'] = [
                 sys.executable, __main__.__file__, *sys.argv[1:]
             ]
+            description['entry']['scripts'] = []
+            try:
+                with open(__main__.__file__) as f:
+                    description['entry']['scripts'].append(f.read())
+            except:
+                pass
         except:
             pass
 

qulab/scan/server.py

@@ -1,5 +1,8 @@
 import asyncio
+import os
 import pickle
+import time
+from pathlib import Path
 
 import click
 import dill
@@ -8,7 +11,26 @@ from loguru import logger
 
 from qulab.sys.rpc.zmq_socket import ZMQContextManager
 
-from .scan import Scan
+from .curd import (create_cell, create_config, create_notebook, get_config,
+                   query_record, remove_tags, tag, update_tags)
+from .models import Cell, Notebook
+from .models import Record as RecordInDB
+from .models import Session, create_engine, create_tables, sessionmaker, utcnow
+from .record import BufferList, Record, random_path
+
+try:
+    default_record_port = int(os.getenv('QULAB_RECORD_PORT', 6789))
+except:
+    default_record_port = 6789
+
+if os.getenv('QULAB_RECORD_PATH'):
+    datapath = Path(os.getenv('QULAB_RECORD_PATH'))
+else:
+    datapath = Path.home() / 'qulab' / 'data'
+datapath.mkdir(parents=True, exist_ok=True)
+
+record_cache = {}
+CACHE_SIZE = 1024
 
 pool = {}
 
@@ -27,15 +49,177 @@ async def reply(req, resp):
     await req.sock.send_multipart([req.identity, pickle.dumps(resp)])
 
 
+def clear_cache():
+    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) - CACHE_SIZE)):
+        del record_cache[k]
+
+
+def flush_cache():
+    for k, (t, r) in record_cache.items():
+        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()
+
+    if record_in_db.file.endswith('.zip'):
+        return Record.load(datapath / 'objects' / record_in_db.file)
+
+    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 = get_local_record(session, id, datapath)
+    else:
+        record = record_cache[id][1]
+    clear_cache()
+    record_cache[id] = time.time(), record
+    return record
+
+
+def record_create(session: Session, description: dict, datapath: Path) -> int:
+    record = Record(None, datapath, description)
+    record_in_db = RecordInDB()
+    if 'app' in description:
+        record_in_db.app = description['app']
+    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_in_db.config_id = description['config']
+    record._file = datapath / 'objects' / record_in_db.file
+    session.add(record_in_db)
+    try:
+        session.commit()
+        record.id = record_in_db.id
+        clear_cache()
+        record_cache[record.id] = time.time(), record
+        return record.id
+    except:
+        session.rollback()
+        raise
+
+
+def record_append(session: Session, record_id: int, level: int, step: int,
+                  position: int, variables: dict, datapath: Path):
+    record = get_record(session, record_id, datapath)
+    record.append(level, step, position, variables)
+    try:
+        record_in_db = session.get(RecordInDB, record_id)
+        record_in_db.mtime = utcnow()
+        record_in_db.atime = utcnow()
+        session.commit()
+    except:
+        session.rollback()
+        raise
+
+
+def record_delete(session: Session, record_id: int, datapath: Path):
+    record = get_local_record(session, record_id, datapath)
+    record.delete()
+    record_in_db = session.get(RecordInDB, record_id)
+    session.delete(record_in_db)
+    session.commit()
+
+
 @logger.catch
-async def handle(request: Request):
+async def handle(session: Session, request: Request, datapath: Path):
 
     msg = request.msg
 
     match request.method:
         case 'ping':
             await reply(request, 'pong')
+        case 'bufferlist_slice':
+            record = get_record(session, msg['record_id'], datapath)
+            bufferlist = record.get(msg['key'],
+                                    buffer_to_array=False,
+                                    slice=msg['slice'])
+            await reply(request, list(bufferlist.iter()))
+        case 'record_create':
+            description = dill.loads(msg['description'])
+            await reply(request, record_create(session, description, datapath))
+        case 'record_append':
+            record_append(session, msg['record_id'], msg['level'], msg['step'],
+                          msg['position'], msg['variables'], datapath)
+        case 'record_description':
+            record = get_record(session, msg['record_id'], datapath)
+            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))
+        case 'record_keys':
+            record = get_record(session, msg['record_id'], datapath)
+            await reply(request, record.keys())
+        case 'record_query':
+            total, apps, table = query_record(session,
+                                              offset=msg.get('offset', 0),
+                                              limit=msg.get('limit', 10),
+                                              app=msg.get('app', None),
+                                              tags=msg.get('tags', ()),
+                                              before=msg.get('before', None),
+                                              after=msg.get('after', None))
+            await reply(request, (total, apps, table))
+        case 'record_get_tags':
+            record_in_db = session.get(RecordInDB, msg['record_id'])
+            await reply(request, [t.name for t in record_in_db.tags])
+        case 'record_remove_tags':
+            remove_tags(session, msg['record_id'], msg['tags'])
+        case 'record_add_tags':
+            update_tags(session, msg['record_id'], msg['tags'], True)
+        case 'record_replace_tags':
+            update_tags(session, msg['record_id'], msg['tags'], False)
+        case 'notebook_create':
+            notebook = create_notebook(session, msg['name'])
+            session.commit()
+            await reply(request, notebook.id)
+        case 'notebook_extend':
+            notebook = session.get(Notebook, msg['notebook_id'])
+            inputCells = msg.get('input_cells', [""])
+            aready_saved = len(notebook.cells)
+            if len(inputCells) > aready_saved:
+                for cell in inputCells[aready_saved:]:
+                    cell = create_cell(session, notebook, cell)
+                session.commit()
+                await reply(request, cell.id)
+            else:
+                await reply(request, None)
+        case 'notebook_history':
+            cell = session.get(Cell, msg['cell_id'])
+            if cell:
+                await reply(request, [
+                    cell.input.text
+                    for cell in cell.notebook.cells[1:cell.index + 2]
+                ])
+            else:
+                await reply(request, None)
+        case 'config_get':
+            config = get_config(session,
+                                msg['config_id'],
+                                base=datapath / 'objects')
+            session.commit()
+            await reply(request, config)
+        case 'config_update':
+            config = create_config(session,
+                                   msg['update'],
+                                   base=datapath / 'objects',
+                                   filename='/'.join(
+                                       random_path(datapath /
+                                                   'objects').parts[-4:]))
+            session.commit()
+            await reply(request, config.id)
         case 'submit':
+            from .scan import Scan
             description = dill.loads(msg['description'])
             task = Scan()
             task.description = description
@@ -55,24 +239,42 @@ async def handle(request: Request):
             logger.error(f"Unknown method: {msg['method']}")
 
 
-async def _handle(request: Request):
+async def _handle(session: Session, request: Request, datapath: Path):
     try:
-        await handle(request)
+        await handle(session, request, datapath)
     except:
         await reply(request, 'error')
 
 
-async def serv(port):
+async def serv(port,
+               datapath,
+               url=None,
+               buffer_size=1024 * 1024 * 1024,
+               interval=60):
     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))
+        if url is None:
+            url = 'sqlite:///' + str(datapath / 'data.db')
+        engine = create_engine(url)
+        create_tables(engine)
+        Session = sessionmaker(engine)
+        with Session() as session:
+            logger.info('Server started.')
+            received = 0
+            last_flush_time = time.time()
+            while True:
+                identity, msg = await sock.recv_multipart()
+                received += len(msg)
+                req = Request(sock, identity, msg)
+                asyncio.create_task(_handle(session, req, datapath))
+                if received > buffer_size or time.time(
+                ) - last_flush_time > interval:
+                    flush_cache()
+                    received = 0
+                    last_flush_time = time.time()
 
 
-async def watch(port, timeout=1):
+async def watch(port, datapath, url=None, timeout=1, buffer=1024, interval=60):
     with ZMQContextManager(zmq.DEALER,
                            connect=f"tcp://127.0.0.1:{port}") as sock:
         sock.setsockopt(zmq.LINGER, 0)
@@ -84,20 +286,34 @@ async def watch(port, timeout=1):
                 else:
                     raise asyncio.TimeoutError()
             except (zmq.error.ZMQError, asyncio.TimeoutError):
-                return asyncio.create_task(serv(port))
+                return asyncio.create_task(
+                    serv(port, datapath, url, buffer * 1024 * 1024, interval))
             await asyncio.sleep(timeout)
 
 
-async def main(port, timeout=1):
-    task = await watch(port=port, timeout=timeout)
+async def main(port, datapath, url, timeout=1, buffer=1024, interval=60):
+    task = await watch(port=port,
+                       datapath=datapath,
+                       url=url,
+                       timeout=timeout,
+                       buffer=buffer,
+                       interval=interval)
     await task
 
 
 @click.command()
-@click.option('--port', default=6788, help='Port of the server.')
+@click.option('--port',
+              default=os.getenv('QULAB_RECORD_PORT', 6789),
+              help='Port of the server.')
+@click.option('--datapath', default=datapath, help='Path of the data.')
+@click.option('--url', default=None, help='URL of the database.')
 @click.option('--timeout', default=1, help='Timeout of ping.')
-def server(port, timeout):
-    asyncio.run(main(port, timeout))
+@click.option('--buffer', default=1024, help='Buffer size (MB).')
+@click.option('--interval',
+              default=60,
+              help='Interval of flush cache, in unit of second.')
+def server(port, datapath, url, timeout, buffer, interval):
+    asyncio.run(main(port, Path(datapath), url, timeout, buffer, interval))
 
 
 if __name__ == "__main__":

qulab/sys/rpc/zmq_socket.py

@@ -98,7 +98,8 @@ class ZMQContextManager:
                  public_keys_location: Optional[str] = None,
                  secret_key: Optional[bytes] = None,
                  public_key: Optional[bytes] = None,
-                 server_public_key: Optional[bytes] = None):
+                 server_public_key: Optional[bytes] = None,
+                 socket: Optional[zmq.Socket] = None):
         self.socket_type = socket_type
         if bind is None and connect is None:
             raise ValueError("Either 'bind' or 'connect' must be specified.")
@@ -129,6 +130,7 @@ class ZMQContextManager:
         self.auth = None
         self.context = None
         self.socket = None
+        self._external_socket = socket
 
     def _create_socket(self, asyncio=False) -> zmq.Socket:
         """
@@ -138,6 +140,8 @@ class ZMQContextManager:
         Returns:
             zmq.Socket: The configured ZeroMQ socket.
         """
+        if self._external_socket:
+            return self._external_socket
         if asyncio:
             self.context = zmq.asyncio.Context()
         else:
@@ -185,6 +189,8 @@ class ZMQContextManager:
         Closes the ZeroMQ socket and the context, and stops the authenticator
         if it was started.
         """
+        if self._external_socket:
+            return
         if self.observer:
             self.observer.stop()
             self.observer.join()

qulab/version.py

@@ -1 +1 @@
-__version__ = "2.1.2"
+__version__ = "2.1.4"

qulab/scan/recorder.py

@@ -1,247 +0,0 @@
-import asyncio
-import os
-import pickle
-import time
-from pathlib import Path
-
-import click
-import dill
-import zmq
-from loguru import logger
-
-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
-from .record import Record
-
-try:
-    default_record_port = int(os.getenv('QULAB_RECORD_PORT', 6789))
-except:
-    default_record_port = 6789
-
-if os.getenv('QULAB_RECORD_PATH'):
-    datapath = Path(os.getenv('QULAB_RECORD_PATH'))
-else:
-    datapath = Path.home() / 'qulab' / 'data'
-datapath.mkdir(parents=True, exist_ok=True)
-
-record_cache = {}
-CACHE_SIZE = 1024
-
-
-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)])
-
-
-def clear_cache():
-    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) - CACHE_SIZE)):
-        del record_cache[k]
-
-
-def flush_cache():
-    for k, (t, r) in record_cache.items():
-        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 = get_local_record(session, id, datapath)
-    else:
-        record = record_cache[id][1]
-    clear_cache()
-    record_cache[id] = time.time(), record
-    return record
-
-
-def record_create(session: Session, description: dict, datapath: Path) -> int:
-    record = Record(None, datapath, description)
-    record_in_db = RecordInDB()
-    if 'app' in description:
-        record_in_db.app = description['app']
-    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()
-        record.id = record_in_db.id
-        clear_cache()
-        record_cache[record.id] = time.time(), record
-        return record.id
-    except:
-        session.rollback()
-        raise
-
-
-def record_append(session: Session, record_id: int, level: int, step: int,
-                  position: int, variables: dict, datapath: Path):
-    record = get_record(session, record_id, datapath)
-    record.append(level, step, position, variables)
-    try:
-        record_in_db = session.get(RecordInDB, record_id)
-        record_in_db.mtime = utcnow()
-        record_in_db.atime = utcnow()
-        session.commit()
-    except:
-        session.rollback()
-        raise
-
-
-@logger.catch
-async def handle(session: Session, request: Request, datapath: Path):
-
-    msg = request.msg
-
-    match request.method:
-        case 'ping':
-            await reply(request, 'pong')
-        case 'bufferlist_slice':
-            record = get_record(session, msg['record_id'], datapath)
-            bufferlist = record.get(msg['key'],
-                                    buffer_to_array=False,
-                                    slice=msg['slice'])
-            await reply(request, list(bufferlist.iter()))
-        case 'record_create':
-            description = dill.loads(msg['description'])
-            await reply(request, record_create(session, description, datapath))
-        case 'record_append':
-            record_append(session, msg['record_id'], msg['level'], msg['step'],
-                          msg['position'], msg['variables'], datapath)
-        case 'record_description':
-            record = get_record(session, msg['record_id'], datapath)
-            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))
-        case 'record_keys':
-            record = get_record(session, msg['record_id'], datapath)
-            await reply(request, record.keys())
-        case 'record_query':
-            total, apps, table = query_record(session,
-                                              offset=msg.get('offset', 0),
-                                              limit=msg.get('limit', 10),
-                                              app=msg.get('app', None),
-                                              tags=msg.get('tags', ()),
-                                              before=msg.get('before', None),
-                                              after=msg.get('after', None))
-            await reply(request, (total, apps, table))
-        case 'record_get_tags':
-            record_in_db = session.get(RecordInDB, msg['record_id'])
-            await reply(request, [t.name for t in record_in_db.tags])
-        case 'record_remove_tags':
-            remove_tags(session, msg['record_id'], msg['tags'])
-        case 'record_add_tags':
-            update_tags(session, msg['record_id'], msg['tags'], True)
-        case 'record_replace_tags':
-            update_tags(session, msg['record_id'], msg['tags'], False)
-        case _:
-            logger.error(f"Unknown method: {msg['method']}")
-
-
-async def _handle(session: Session, request: Request, datapath: Path):
-    try:
-        await handle(session, request, datapath)
-    except:
-        await reply(request, 'error')
-
-
-async def serv(port,
-               datapath,
-               url=None,
-               buffer_size=1024 * 1024 * 1024,
-               interval=60):
-    logger.info('Server starting.')
-    async with ZMQContextManager(zmq.ROUTER, bind=f"tcp://*:{port}") as sock:
-        if url is None:
-            url = 'sqlite:///' + str(datapath / 'data.db')
-        engine = create_engine(url)
-        create_tables(engine)
-        Session = sessionmaker(engine)
-        with Session() as session:
-            logger.info('Server started.')
-            received = 0
-            last_flush_time = time.time()
-            while True:
-                identity, msg = await sock.recv_multipart()
-                received += len(msg)
-                req = Request(sock, identity, msg)
-                asyncio.create_task(_handle(session, req, datapath))
-                if received > buffer_size or time.time(
-                ) - last_flush_time > interval:
-                    flush_cache()
-                    received = 0
-                    last_flush_time = time.time()
-
-
-async def watch(port, datapath, url=None, timeout=1, buffer=1024, interval=60):
-    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, datapath, url, buffer * 1024 * 1024, interval))
-            await asyncio.sleep(timeout)
-
-
-async def main(port, datapath, url, timeout=1, buffer=1024, interval=60):
-    task = await watch(port=port,
-                       datapath=datapath,
-                       url=url,
-                       timeout=timeout,
-                       buffer=buffer,
-                       interval=interval)
-    await task
-
-
-@click.command()
-@click.option('--port',
-              default=os.getenv('QULAB_RECORD_PORT', 6789),
-              help='Port of the server.')
-@click.option('--datapath', default=datapath, help='Path of the data.')
-@click.option('--url', default=None, help='URL of the database.')
-@click.option('--timeout', default=1, help='Timeout of ping.')
-@click.option('--buffer', default=1024, help='Buffer size (MB).')
-@click.option('--interval',
-              default=60,
-              help='Interval of flush cache, in unit of second.')
-def record(port, datapath, url, timeout, buffer, interval):
-    asyncio.run(main(port, Path(datapath), url, timeout, buffer, interval))
-
-
-if __name__ == "__main__":
-    record()