QuLab 2.0.9__cp311-cp311-macosx_10_9_universal2.whl → 2.1.1__cp311-cp311-macosx_10_9_universal2.whl

qulab/scan/scan.py

@@ -1,34 +1,52 @@
 import asyncio
 import copy
-import datetime
 import inspect
 import itertools
 import os
 import re
 import sys
 import uuid
-import warnings
+from concurrent.futures import ProcessPoolExecutor
 from graphlib import TopologicalSorter
 from pathlib import Path
-from types import MethodType
-from typing import Any, Awaitable, Callable, Iterable, Type
+from typing import Any, Awaitable, Callable, Iterable
 
 import dill
 import numpy as np
-import skopt
 import zmq
-from skopt.space import Categorical, Integer, Real
-from tqdm.notebook import tqdm
 
 from ..sys.rpc.zmq_socket import ZMQContextManager
 from .expression import Env, Expression, Symbol
 from .optimize import NgOptimizer
-from .recorder import Record, default_record_port
-from .utils import async_zip, call_function
+from .record import Record
+from .recorder import default_record_port
+from .space import Optimizer, OptimizeSpace, Space
+from .utils import async_zip, call_function, dump_globals
+
+try:
+    from tqdm.notebook import tqdm
+except:
+
+    class tqdm():
+
+        def update(self, n):
+            pass
+
+        def close(self):
+            pass
+
+        def reset(self):
+            pass
+
 
 __process_uuid = uuid.uuid1()
 __task_counter = itertools.count()
 
+if os.getenv('QULAB_SERVER'):
+    default_server = os.getenv('QULAB_SERVER')
+else:
+    default_server = f'tcp://127.0.0.1:{default_record_port}'
+
 
 def task_uuid():
     return uuid.uuid3(__process_uuid, str(next(__task_counter)))
@@ -54,79 +72,6 @@ def _get_depends(func: Callable):
     return args
 
 
-class OptimizeSpace():
-
-    def __init__(self, optimizer: 'Optimizer', space):
-        self.optimizer = optimizer
-        self.space = space
-        self.name = None
-
-    def __len__(self):
-        return self.optimizer.maxiter
-
-
-class Optimizer():
-
-    def __init__(self,
-                 scanner: 'Scan',
-                 name: str,
-                 level: int,
-                 method: str | Type = skopt.Optimizer,
-                 maxiter: int = 1000,
-                 minimize: bool = True,
-                 **kwds):
-        self.scanner = scanner
-        self.method = method
-        self.maxiter = maxiter
-        self.dimensions = {}
-        self.name = name
-        self.level = level
-        self.kwds = kwds
-        self.minimize = minimize
-
-    def create(self):
-        return self.method(list(self.dimensions.values()), **self.kwds)
-
-    def Categorical(self,
-                    categories,
-                    prior=None,
-                    transform=None,
-                    name=None) -> OptimizeSpace:
-        return OptimizeSpace(self,
-                             Categorical(categories, prior, transform, name))
-
-    def Integer(self,
-                low,
-                high,
-                prior="uniform",
-                base=10,
-                transform=None,
-                name=None,
-                dtype=np.int64) -> OptimizeSpace:
-        return OptimizeSpace(
-            self, Integer(low, high, prior, base, transform, name, dtype))
-
-    def Real(self,
-             low,
-             high,
-             prior="uniform",
-             base=10,
-             transform=None,
-             name=None,
-             dtype=float) -> OptimizeSpace:
-        return OptimizeSpace(
-            self, Real(low, high, prior, base, transform, name, dtype))
-
-    def __getstate__(self) -> dict:
-        state = self.__dict__.copy()
-        del state['scanner']
-        return state
-
-    def __setstate__(self, state: dict) -> None:
-        self.__dict__.update(state)
-        self.scanner = None
-
-
 class Promise():
     __slots__ = ['task', 'key', 'attr']
 
@@ -157,6 +102,11 @@ class Promise():
         return Promise(self.task, None, attr)
 
 
+def _run_function_in_process(buf):
+    func, args, kwds = dill.loads(buf)
+    return func(*args, **kwds)
+
+
 class Scan():
 
     def __new__(cls, *args, mixin=None, **kwds):
@@ -174,8 +124,10 @@ class Scan():
                  app: str = 'task',
                  tags: tuple[str] = (),
                  database: str | Path
-                 | None = f'tcp://127.0.0.1:{default_record_port}',
+                 | None = default_server,
                  dump_globals: bool = False,
+                 max_workers: int = 4,
+                 max_promise: int = 100,
                  mixin=None):
         self.id = task_uuid()
         self.record = None
@@ -193,6 +145,8 @@ class Scan():
             'actions': {},
             'dependents': {},
             'order': {},
+            'axis': {},
+            'independent_variables': set(),
             'filters': {},
             'total': {},
             'database': database,
@@ -207,23 +161,21 @@ class Scan():
         self._variables = {}
         self._main_task = None
         self._sock = None
-        self._sem = asyncio.Semaphore(100)
+        self._sem = asyncio.Semaphore(max_promise + 1)
         self._bar: dict[int, tqdm] = {}
         self._hide_pattern_re = re.compile('|'.join(self.description['hiden']))
-        self._task_queue = asyncio.Queue()
-        self._task_pool = []
+        self._msg_queue = asyncio.Queue()
+        self._prm_queue = asyncio.Queue()
         self._single_step = True
+        self._max_workers = max_workers
+        self._max_promise = max_promise
+        self._executors = ProcessPoolExecutor(max_workers=max_workers)
 
     def __del__(self):
         try:
             self._main_task.cancel()
         except:
             pass
-        for task in self._task_pool:
-            try:
-                task.cancel()
-            except:
-                pass
 
     def __getstate__(self) -> dict:
         state = self.__dict__.copy()
@@ -231,9 +183,10 @@ class Scan():
         del state['_sock']
         del state['_main_task']
         del state['_bar']
-        del state['_task_queue']
-        del state['_task_pool']
+        del state['_msg_queue']
+        del state['_prm_queue']
         del state['_sem']
+        del state['_executors']
         return state
 
     def __setstate__(self, state: dict) -> None:
@@ -242,12 +195,20 @@ class Scan():
         self._sock = None
         self._main_task = None
         self._bar = {}
-        self._task_queue = asyncio.Queue()
-        self._task_pool = []
-        self._sem = asyncio.Semaphore(100)
+        self._prm_queue = asyncio.Queue()
+        self._msg_queue = asyncio.Queue()
+        self._sem = asyncio.Semaphore(self._max_promise + 1)
+        self._executors = ProcessPoolExecutor(max_workers=self._max_workers)
         for opt in self.description['optimizers'].values():
             opt.scanner = self
 
+    def __del__(self):
+        try:
+            self._main_task.cancel()
+        except:
+            pass
+        self._executors.shutdown()
+
     @property
     def current_level(self):
         return self._current_level
@@ -256,8 +217,8 @@ class Scan():
     def variables(self) -> dict[str, Any]:
         return self._variables
 
-    async def emit(self, current_level, step, position, variables: dict[str,
-                                                                        Any]):
+    async def _emit(self, current_level, step, position, variables: dict[str,
+                                                                         Any]):
         for key, value in list(variables.items()):
             if inspect.isawaitable(value) and not self.hiden(key):
                 variables[key] = await value
@@ -280,6 +241,11 @@ class Scan():
                 for k, v in variables.items() if not self.hiden(k)
             })
 
+    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())))
+
     def hide(self, name: str):
         self.description['hiden'].append(name)
         self._hide_pattern_re = re.compile('|'.join(self.description['hiden']))
@@ -289,11 +255,11 @@ class Scan():
 
     async def _filter(self, variables: dict[str, Any], level: int = 0):
         try:
-            return all([
-                await call_function(fun, variables) for fun in itertools.chain(
+            return all(await asyncio.gather(*[
+                call_function(fun, variables) for fun in itertools.chain(
                     self.description['filters'].get(level, []),
                     self.description['filters'].get(-1, []))
-            ])
+            ]))
         except:
             return True
 
@@ -345,7 +311,11 @@ class Scan():
             self.description['filters'][level] = []
         self.description['filters'][level].append(func)
 
-    def set(self, name: str, value, setter: Callable | None = None):
+    def set(self,
+            name: str,
+            value,
+            depends: Iterable[str] | None = None,
+            setter: Callable | None = None):
         try:
             dill.dumps(value)
         except:
@@ -354,9 +324,21 @@ class Scan():
             self.add_depends(name, value.symbols())
             self.description['functions'][name] = value
         elif callable(value):
-            self.add_depends(name, _get_depends(value))
-            self.description['functions'][name] = value
+            if depends:
+                self.add_depends(name, depends)
+                s = ','.join(depends)
+                self.description['functions'][f'_tmp_{name}'] = value
+                self.description['functions'][name] = eval(
+                    f"lambda self, {s}: self.description['functions']['_tmp_{name}']({s})"
+                )
+            else:
+                self.add_depends(name, _get_depends(value))
+                self.description['functions'][name] = value
         else:
+            try:
+                value = Space.fromarray(value)
+            except:
+                pass
             self.description['consts'][name] = value
         if setter:
             self.description['setters'][name] = setter
@@ -376,6 +358,10 @@ class Scan():
         else:
             if level is None:
                 raise ValueError('level must be provided.')
+            try:
+                range = Space.fromarray(range)
+            except:
+                pass
             self._add_loop_var(name, level, range)
             if isinstance(range, Expression) or callable(range):
                 self.add_depends(name, range.symbols())
@@ -432,11 +418,15 @@ class Scan():
 
     async def _update_progress(self):
         while True:
-            task = await self._task_queue.get()
-            if isinstance(task, asyncio.Event):
-                task.set()
-            elif inspect.isawaitable(task):
-                await task
+            task = await self._prm_queue.get()
+            await task
+            self._prm_queue.task_done()
+
+    async def _send_msg(self):
+        while True:
+            task = await self._msg_queue.get()
+            await task
+            self._msg_queue.task_done()
 
     async def run(self):
         assymbly(self.description)
@@ -452,44 +442,48 @@ class Scan():
             await self._run()
 
     async def _run(self):
-        task = asyncio.create_task(self._update_progress())
-        self._task_pool.append(task)
+        send_msg = asyncio.create_task(self._send_msg())
+        update_progress_task = asyncio.create_task(self._update_progress())
+
         self._variables = {'self': self}
-        self._variables.update(self.description['consts'].copy())
+
+        consts = {}
+        for k, v in self.description['consts'].items():
+            if isinstance(v, Space):
+                consts[k] = v.toarray()
+            else:
+                consts[k] = v
+
+        await update_variables(self._variables, consts,
+                               self.description['setters'])
         for level, total in self.description['total'].items():
             if total == np.inf:
                 total = None
             self._bar[level] = tqdm(total=total)
-        for group in self.description['order'].get(-1, []):
-            for name in group:
-                if name in self.description['functions']:
-                    self.variables[name] = await call_function(
-                        self.description['functions'][name], self.variables)
-                if name in self.description['setters']:
-                    coro = self.description['setters'][name](
-                        self.variables[name])
-                    if inspect.isawaitable(coro):
-                        await coro
+
+        updates = await call_many_functions(
+            self.description['order'].get(-1, []),
+            self.description['functions'], self.variables)
+        await update_variables(self.variables, updates,
+                               self.description['setters'])
+
         self.record = await self.create_record()
         await self.work()
         for level, bar in self._bar.items():
             bar.close()
 
-        while not self._task_queue.empty():
-            evt = self._task_queue.get_nowait()
-            if isinstance(evt, asyncio.Event):
-                evt.set()
-            elif inspect.isawaitable(evt):
-                await evt
         if self._single_step:
-            for group in self.description['order'].get(-1, []):
-                for name in group:
-                    if name in self.description['getters']:
-                        self.variables[name] = await call_function(
-                            self.description['getters'][name], self.variables)
-            await self.emit(0, 0, 0, self.variables.copy())
-            await self.emit(-1, 0, 0, {})
-        task.cancel()
+            self.variables.update(await call_many_functions(
+                self.description['order'].get(-1, []),
+                self.description['getters'], self.variables))
+
+            self.emit(0, 0, 0, self.variables)
+            self.emit(-1, 0, 0, {})
+
+        await self._prm_queue.join()
+        update_progress_task.cancel()
+        await self._msg_queue.join()
+        send_msg.cancel()
         return self.variables
 
     async def done(self):
@@ -536,8 +530,8 @@ class Scan():
             return
         step = 0
         position = 0
-        self._task_queue.put_nowait(
-            self._reset_progress_bar(self.current_level))
+        self._prm_queue.put_nowait(self._reset_progress_bar(
+            self.current_level))
         async for variables in _iter_level(
                 self.variables,
                 self.description['loops'].get(self.current_level, []),
@@ -548,23 +542,18 @@ class Scan():
             if await self._filter(variables, self.current_level - 1):
                 yield variables
                 self._single_step = False
-                asyncio.create_task(
-                    self.emit(self.current_level - 1, step, position,
-                              variables.copy()))
+                self.emit(self.current_level - 1, step, position, variables)
                 step += 1
             position += 1
             self._current_level -= 1
-            self._task_queue.put_nowait(
+            self._prm_queue.put_nowait(
                 self._update_progress_bar(self.current_level, 1))
         if self.current_level == 0:
-            await self.emit(self.current_level - 1, 0, 0, {})
+            self.emit(self.current_level - 1, 0, 0, {})
             for name, value in self.variables.items():
                 if inspect.isawaitable(value):
                     self.variables[name] = await value
-            while not self._task_queue.empty():
-                task = self._task_queue.get_nowait()
-                if inspect.isawaitable(task):
-                    await task
+            await self._prm_queue.join()
 
     async def work(self, **kwds):
         if self.current_level in self.description['actions']:
@@ -589,7 +578,8 @@ class Scan():
         """
         self.description['actions'][level] = action
 
-    async def promise(self, awaitable: Awaitable) -> Promise:
+    async def promise(self, awaitable: Awaitable | Callable, *args,
+                      **kwds) -> Promise:
         """
         Promise to calculate asynchronous function and return the result in future.
 
@@ -602,8 +592,19 @@ class Scan():
         if inspect.isawaitable(awaitable):
             async with self._sem:
                 task = asyncio.create_task(self._await(awaitable))
-                self._task_queue.put_nowait(task)
+                self._prm_queue.put_nowait(task)
                 return Promise(task)
+        elif inspect.iscoroutinefunction(awaitable):
+            return await self.promise(awaitable(*args, **kwds))
+        elif callable(awaitable):
+            try:
+                buf = dill.dumps((awaitable, args, kwds))
+                task = asyncio.get_running_loop().run_in_executor(
+                    self._executors, _run_function_in_process, buf)
+                self._prm_queue.put_nowait(task)
+                return Promise(task)
+            except:
+                return awaitable(*args, **kwds)
         else:
             return awaitable
 
@@ -612,51 +613,6 @@ class Scan():
             return await awaitable
 
 
-class Unpicklable:
-
-    def __init__(self, obj):
-        self.type = str(type(obj))
-        self.id = id(obj)
-
-    def __repr__(self):
-        return f'<Unpicklable: {self.type} at 0x{id(self):x}>'
-
-
-class TooLarge:
-
-    def __init__(self, obj):
-        self.type = str(type(obj))
-        self.id = id(obj)
-
-    def __repr__(self):
-        return f'<TooLarge: {self.type} at 0x{id(self):x}>'
-
-
-def dump_globals(ns=None, *, size_limit=10 * 1024 * 1024, warn=False):
-    import __main__
-
-    if ns is None:
-        ns = __main__.__dict__
-
-    namespace = {}
-
-    for name, value in ns.items():
-        try:
-            buf = dill.dumps(value)
-        except:
-            namespace[name] = Unpicklable(value)
-            if warn:
-                warnings.warn(f'Unpicklable: {name} {type(value)}')
-        if len(buf) > size_limit:
-            namespace[name] = TooLarge(value)
-            if warn:
-                warnings.warn(f'TooLarge: {name} {type(value)}')
-        else:
-            namespace[name] = buf
-
-    return namespace
-
-
 def assymbly(description):
     import __main__
     from IPython import get_ipython
@@ -782,16 +738,52 @@ def assymbly(description):
             if ready:
                 description['order'][level].append(ready)
                 keys -= set(ready)
+
+    axis = {}
+    independent_variables = set()
+
+    for name in description['consts']:
+        axis[name] = ()
+    for level, range_list in description['loops'].items():
+        for name, iterable in range_list:
+            if isinstance(iterable, OptimizeSpace):
+                axis[name] = tuple(range(level + 1))
+                continue
+            elif isinstance(iterable, (np.ndarray, list, tuple, range, Space)):
+                independent_variables.add(name)
+            axis[name] = (level, )
+
+    for level, group in description['order'].items():
+        for names in group:
+            for name in names:
+                if name not in description['dependents']:
+                    if name not in axis:
+                        axis[name] = (level, )
+                else:
+                    d = set()
+                    for n in description['dependents'][name]:
+                        d.update(axis[n])
+                    if name not in axis:
+                        axis[name] = tuple(sorted(d))
+                    else:
+                        axis[name] = tuple(sorted(set(axis[name]) | d))
+    description['axis'] = axis
+    description['independent_variables'] = independent_variables
+
     return description
 
 
-async def _update_variables(variables, updates, setters):
+async def update_variables(variables: dict[str, Any], updates: dict[str, Any],
+                           setters: dict[str, Callable]):
+    coros = []
     for name, value in updates.items():
         if name in setters:
             coro = setters[name](value)
             if inspect.isawaitable(coro):
-                await coro
+                coros.append(coro)
         variables[name] = value
+    if coros:
+        await asyncio.gather(*coros)
 
 
 async def _iter_level(variables,
@@ -813,6 +805,8 @@ async def _iter_level(variables,
                 opts[iter.optimizer.name] = iter.optimizer.create()
         elif isinstance(iter, Expression):
             iters_d[name] = iter.eval(env)
+        elif isinstance(iter, Space):
+            iters_d[name] = iter.toarray()
         elif callable(iter):
             iters_d[name] = await call_function(iter, variables)
         else:
@@ -824,31 +818,23 @@ async def _iter_level(variables,
         maxiter = min(maxiter, opt_cfg.maxiter)
 
     async for args in async_zip(*iters_d.values(), range(maxiter)):
-        await _update_variables(variables, dict(zip(iters_d.keys(),
-                                                    args[:-1])), setters)
+        await update_variables(variables, dict(zip(iters_d.keys(), args[:-1])),
+                               setters)
         for name, opt in opts.items():
             args = opt.ask()
             opt_cfg = optimizers[name]
-            await _update_variables(variables, {
+            await update_variables(variables, {
                 n: v
                 for n, v in zip(opt_cfg.dimensions.keys(), args)
             }, setters)
 
-        for group in order:
-            for name in group:
-                if name in functions:
-                    await _update_variables(variables, {
-                        name:
-                        await call_function(functions[name], variables)
-                    }, setters)
+        await update_variables(
+            variables, await call_many_functions(order, functions, variables),
+            setters)
 
         yield variables
 
-        for group in order:
-            for name in group:
-                if name in getters:
-                    variables[name] = await call_function(
-                        getters[name], variables)
+        variables.update(await call_many_functions(order, getters, variables))
 
         for name, opt in opts.items():
             opt_cfg = optimizers[name]
@@ -866,10 +852,28 @@ async def _iter_level(variables,
     for name, opt in opts.items():
         opt_cfg = optimizers[name]
         result = opt.get_result()
-        variables.update({
-            n: v
-            for n, v in zip(opt_cfg.dimensions.keys(), result.x)
-        })
+        await update_variables(
+            variables, {
+                name: value
+                for name, value in zip(opt_cfg.dimensions.keys(), result.x)
+            }, setters)
         variables[name] = result.fun
     if opts:
         yield variables
+
+
+async def call_many_functions(order: list[list[str]],
+                              functions: dict[str, Callable],
+                              variables: dict[str, Any]) -> dict[str, Any]:
+    ret = {}
+    for group in order:
+        waited = []
+        coros = []
+        for name in group:
+            if name in functions:
+                waited.append(name)
+                coros.append(call_function(functions[name], variables | ret))
+        if coros:
+            results = await asyncio.gather(*coros)
+            ret.update(dict(zip(waited, results)))
+    return ret

qulab/scan/server.py

@@ -1,20 +1,18 @@
 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
 
+from .scan import Scan
+
 pool = {}
 
+
 class Request():
     __slots__ = ['sock', 'identity', 'msg', 'method']