QuLab 2.10.10__cp313-cp313-macosx_10_13_universal2.whl

qulab/scan/scan.py

@@ -0,0 +1,1166 @@
+import asyncio
+import contextlib
+import copy
+import inspect
+import itertools
+import lzma
+import os
+import pickle
+import re
+import sys
+import uuid
+from concurrent.futures import ProcessPoolExecutor
+from graphlib import TopologicalSorter
+from pathlib import Path
+from typing import Any, Awaitable, Callable, Iterable
+
+import dill
+import numpy as np
+import zmq
+
+from ..expression import Env, Expression, Symbol
+from ..sys.rpc.zmq_socket import ZMQContextManager
+from .optimize import NgOptimizer
+from .record import Record
+from .server import default_record_port
+from .space import Optimizer, OptimizeSpace, Space
+from .utils import (async_zip, call_function, dump_dict, dump_globals,
+                    get_installed_packages, get_system_info, yapf_reformat)
+
+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()
+__notebook_id = None
+
+if os.getenv('QULAB_SERVER'):
+    default_server = os.getenv('QULAB_SERVER')
+else:
+    default_server = f'tcp://127.0.0.1:{default_record_port}'
+if os.getenv('QULAB_EXECUTOR'):
+    default_executor = os.getenv('QULAB_EXECUTOR')
+else:
+    default_executor = default_server
+
+
+class Promise():
+    __slots__ = ['task', 'key', 'attr']
+
+    def __init__(self, task, key=None, attr=None):
+        self.task = task
+        self.key = key
+        self.attr = attr
+
+    def __await__(self):
+
+        async def _getitem(task, key):
+            return (await task)[key]
+
+        async def _getattr(task, attr):
+            return getattr(await task, attr)
+
+        if self.key is not None:
+            return _getitem(self.task, self.key).__await__()
+        elif self.attr is not None:
+            return _getattr(self.task, self.attr).__await__()
+        else:
+            return self.task.__await__()
+
+    def __getitem__(self, key):
+        return Promise(self.task, key, None)
+
+    def __getattr__(self, attr):
+        return Promise(self.task, None, attr)
+
+
+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()
+
+
+def task_uuid():
+    return uuid.uuid3(__process_uuid, str(next(__task_counter)))
+
+
+def _get_depends(func: Callable):
+    try:
+        sig = inspect.signature(func)
+    except:
+        return []
+
+    args = []
+    for name, param in sig.parameters.items():
+        if param.kind in [
+                param.POSITIONAL_ONLY, param.POSITIONAL_OR_KEYWORD,
+                param.KEYWORD_ONLY
+        ]:
+            args.append(name)
+        elif param.kind == param.VAR_KEYWORD:
+            pass
+        elif param.kind == param.VAR_POSITIONAL:
+            raise ValueError('not support VAR_POSITIONAL')
+    return args
+
+
+def _run_function_in_process(buf):
+    func, args, kwds = dill.loads(buf)
+    return func(*args, **kwds)
+
+
+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):
+                coros.append(coro)
+        variables[name] = value
+    if coros:
+        await asyncio.gather(*coros)
+
+
+async def _iter_level(variables,
+                      iters: list[tuple[str, Iterable | Expression | Callable
+                                        | OptimizeSpace]],
+                      order: list[list[str]],
+                      functions: dict[str, Callable | Expression],
+                      optimizers: dict[str, Optimizer],
+                      setters: dict[str, Callable] = {},
+                      getters: dict[str, Callable] = {}):
+    iters_d = {}
+    env = Env()
+    env.variables = variables
+    opts = {}
+
+    for name, iter in iters:
+        if isinstance(iter, OptimizeSpace):
+            if iter.optimizer.name not in opts:
+                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:
+            iters_d[name] = iter
+
+    maxiter = 0xffffffff
+    for name, opt in opts.items():
+        opt_cfg = optimizers[name]
+        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)
+        for name, opt in opts.items():
+            args = opt.ask()
+            opt_cfg = optimizers[name]
+            await update_variables(variables, {
+                n: v
+                for n, v in zip(opt_cfg.dimensions.keys(), args)
+            }, setters)
+
+        await update_variables(
+            variables, await call_many_functions(order, functions, variables),
+            setters)
+
+        yield variables
+
+        variables.update(await call_many_functions(order, getters, variables))
+
+        if opts:
+            for key in list(variables.keys()):
+                if key.startswith('*') or ',' in key:
+                    await _unpack(key, variables)
+
+        for name, opt in opts.items():
+            opt_cfg = optimizers[name]
+            args = [variables[n] for n in opt_cfg.dimensions.keys()]
+
+            if name not in variables:
+                raise ValueError(f'{name} not in variables.')
+            fun = variables[name]
+            if inspect.isawaitable(fun):
+                fun = await fun
+            if opt_cfg.minimize:
+                opt.tell(args, fun)
+            else:
+                opt.tell(args, -fun)
+
+    if opts:
+        for name, opt in opts.items():
+            opt_cfg = optimizers[name]
+            result = opt.get_result()
+            await update_variables(
+                variables, {
+                    name: value
+                    for name, value in zip(opt_cfg.dimensions.keys(), result.x)
+                }, setters)
+
+        yield variables
+
+        variables.update(await call_many_functions(order, getters, variables))
+
+        for key in list(variables.keys()):
+            if key.startswith('*') or ',' in key:
+                await _unpack(key, 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
+
+
+async def _unpack(key, variables):
+    x = variables[key]
+    if inspect.isawaitable(x):
+        x = await x
+    if key.startswith('**'):
+        assert isinstance(
+            x, dict), f"Should promise a dict for `**` symbol. {key}"
+        if "{key}" in key:
+            for k, v in x.items():
+                variables[key[2:].format(key=k)] = v
+        else:
+            variables.update(x)
+    elif key.startswith('*'):
+        assert isinstance(
+            x, (list, tuple,
+                np.ndarray)), f"Should promise a list for `*` symbol. {key}"
+        for i, v in enumerate(x):
+            k = key[1:].format(i=i)
+            variables[k] = v
+    elif ',' in key:
+        keys1, keys2 = [], []
+        args = None
+        for k in key.split(','):
+            if k.startswith('*'):
+                if args is None:
+                    args = k
+                else:
+                    raise ValueError(f'Only one `*` symbol is allowed. {key}')
+            elif args is None:
+                keys1.append(k)
+            else:
+                keys2.append(k)
+        assert isinstance(
+            x,
+            (list, tuple,
+             np.ndarray)), f"Should promise a list for multiple symbols. {key}"
+        if args is None:
+            assert len(keys1) == len(
+                x), f"Length of keys and values should be equal. {key}"
+            for k, v in zip(keys1, x):
+                variables[k] = v
+        else:
+            assert len(keys1) + len(keys2) <= len(
+                x), f"Too many values for unpacking. {key}"
+            for k, v in zip(keys1, x[:len(keys1)]):
+                variables[k] = v
+            end = -len(keys2) if keys2 else None
+            for i, v in enumerate(x[len(keys1):end]):
+                k = args[1:].format(i=i)
+                variables[k] = v
+            if keys2:
+                for k, v in zip(keys2, x[end:]):
+                    variables[k] = v
+    else:
+        return
+    del variables[key]
+
+
+class Scan():
+
+    def __new__(cls, *args, mixin=None, **kwds):
+        if mixin is None:
+            return super().__new__(cls)
+        for k in dir(mixin):
+            if not hasattr(cls, k):
+                try:
+                    setattr(cls, k, getattr(mixin, k))
+                except:
+                    pass
+        return super().__new__(cls)
+
+    def __init__(self,
+                 app: str = 'task',
+                 tags: tuple[str] = (),
+                 database: str | Path
+                 | None = default_server,
+                 dump_globals: bool = False,
+                 max_workers: int = 4,
+                 max_promise: int = 100,
+                 max_message: int = 1000,
+                 config: dict | None = None,
+                 mixin=None):
+        self.id = task_uuid()
+        self.record = None
+        self.config = {} if config is None else copy.deepcopy(config)
+        self._raw_config_copy = copy.deepcopy(self.config)
+        self.description = {
+            'app': app,
+            'tags': tags,
+            'config': None,
+            'loops': {},
+            'intrinsic_loops': {},
+            'consts': {},
+            'functions': {},
+            'getters': {},
+            'setters': {},
+            'optimizers': {},
+            'namespace': {} if dump_globals else None,
+            'actions': {},
+            'dependents': {},
+            'order': {},
+            'axis': {},
+            'independent_variables': set(),
+            'filters': {},
+            'total': {},
+            'database': database,
+            'hiden': ['self', 'config', r'^__.*', r'.*__$', r'^#.*'],
+            'entry': {
+                'system': get_system_info(),
+                'env': {},
+                'shell': '',
+                'cmds': [],
+                'scripts': []
+            },
+        }
+        self._current_level = 0
+        self._variables = {}
+        self._main_task = None
+        self._background_tasks = ()
+        self._sock = None
+        self._sem = asyncio.Semaphore(max_promise + 1)
+        self._bar: dict[int, tqdm] = {}
+        self._hide_pattern_re = re.compile('|'.join(self.description['hiden']))
+        self._msg_queue = asyncio.Queue(max_message)
+        self._prm_queue = asyncio.Queue()
+        self._single_step = True
+        self._max_workers = max_workers
+        self._max_promise = max_promise
+        self._max_message = max_message
+        self._executors = ProcessPoolExecutor(max_workers=max_workers)
+
+    def __del__(self):
+        try:
+            self._main_task.cancel()
+        except:
+            pass
+
+    def __getstate__(self) -> dict:
+        state = self.__dict__.copy()
+        del state['record']
+        del state['_sock']
+        del state['_main_task']
+        del state['_background_tasks']
+        del state['_bar']
+        del state['_msg_queue']
+        del state['_prm_queue']
+        del state['_sem']
+        del state['_executors']
+        return state
+
+    def __setstate__(self, state: dict) -> None:
+        self.__dict__.update(state)
+        self.record = None
+        self._sock = None
+        self._main_task = None
+        self._background_tasks = ()
+        self._bar = {}
+        self._prm_queue = asyncio.Queue()
+        self._msg_queue = asyncio.Queue(self._max_message)
+        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
+        try:
+            self._executors.shutdown()
+        except:
+            pass
+
+    @property
+    def current_level(self):
+        return self._current_level
+
+    @property
+    def variables(self) -> dict[str, Any]:
+        return self._variables
+
+    async def _emit(self, current_level, step, position, variables: dict[str,
+                                                                         Any]):
+        for key, value in list(variables.items()):
+            if key.startswith('*') or ',' in key:
+                await _unpack(key, variables)
+            elif inspect.isawaitable(value) and not self.hiden(key):
+                variables[key] = await value
+
+        if self.record is None:
+            self.record = await self.create_record()
+        if self._sock is not None:
+            await self._sock.send_pyobj({
+                'task': self.id,
+                'method': 'record_append',
+                'record_id': self.record.id,
+                'level': current_level,
+                'step': step,
+                'position': position,
+                'variables': {
+                    k: v
+                    for k, v in variables.items() if not self.hiden(k)
+                }
+            })
+        else:
+            self.record.append(current_level, step, position, {
+                k: v
+                for k, v in variables.items() if not self.hiden(k)
+            })
+
+    async def emit(self, current_level, step, position, variables: dict[str,
+                                                                        Any]):
+        await self._msg_queue.put(
+            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']))
+
+    def hiden(self, name: str) -> bool:
+        return bool(self._hide_pattern_re.match(name)) or name.startswith(
+            '*') or ',' in name
+
+    async def _filter(self, variables: dict[str, Any], level: int = 0):
+        try:
+            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
+
+    async def create_record(self):
+        if self._sock is None:
+            return Record(None, self.description['database'], self.description)
+
+        if self.config:
+            self.description['config'] = await create_config(
+                self._raw_config_copy, 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
+
+        await self._sock.send_pyobj({
+            'task':
+            self.id,
+            'method':
+            'record_create',
+            'description':
+            dump_dict(self.description,
+                      keys=[
+                          'intrinsic_loops', 'app', 'tags', 'loops',
+                          'independent_variables', 'axis', 'config', 'entry'
+                      ])
+        })
+
+        record_id = await self._sock.recv_pyobj()
+        return Record(record_id, self.description['database'],
+                      self.description)
+
+    def get(self, name: str):
+        if name in self.description['consts']:
+            return self.description['consts'][name]
+        else:
+            try:
+                return self._query_config(name)
+            except:
+                return Symbol(name)
+
+    def _add_search_space(self, name: str, level: int, space):
+        if level not in self.description['loops']:
+            self.description['loops'][level] = []
+        self.description['loops'][level].append((name, space))
+
+    def add_depends(self, name: str, depends: list[str]):
+        if isinstance(depends, str):
+            depends = [depends]
+        if 'self' in depends:
+            depends.append('config')
+        if name not in self.description['dependents']:
+            self.description['dependents'][name] = set()
+        self.description['dependents'][name].update(depends)
+
+    def add_filter(self, func: Callable, level: int = -1):
+        """
+        Add a filter function to the scan.
+
+        Args:
+            func: A callable object or an instance of Expression.
+            level: The level of the scan to add the filter. -1 means any level.
+        """
+        if level not in self.description['filters']:
+            self.description['filters'][level] = []
+        self.description['filters'][level].append(func)
+
+    def set(self,
+            name: str,
+            value,
+            depends: Iterable[str] | None = None,
+            setter: Callable | None = None):
+        try:
+            dill.dumps(value)
+        except:
+            raise ValueError('value is not serializable.')
+        if isinstance(value, Expression):
+            self.add_depends(name, value.symbols())
+            self.description['functions'][name] = value
+        elif callable(value):
+            if depends:
+                self.add_depends(name, depends)
+                s = ','.join(depends)
+                self.description['functions'][f'#{name}'] = value
+                self.description['functions'][name] = eval(
+                    f"lambda self, {s}: self.description['functions']['#{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 '.' in name:
+            self.add_depends('config', [name])
+
+        if ',' in name:
+            for key in name.split(','):
+                if not key.startswith('*'):
+                    self.add_depends(key, [name])
+        if setter:
+            self.description['setters'][name] = setter
+
+    def search(self,
+               name: str,
+               space: Iterable | Expression | Callable | OptimizeSpace,
+               level: int | None = None,
+               setter: Callable | None = None,
+               intrinsic: bool = False):
+        if level is not None:
+            if not intrinsic:
+                assert level >= 0, 'level must be greater than or equal to 0.'
+        if intrinsic:
+            assert isinstance(space, (np.ndarray, list, tuple, range, Space)), \
+                'space must be an instance of np.ndarray, list, tuple, range or Space.'
+            self.description['intrinsic_loops'][name] = level
+            self.set(name, space)
+        elif isinstance(space, OptimizeSpace):
+            space.name = name
+            space.optimizer.dimensions[name] = space.space
+            if space.suggestion is not None:
+                space.optimizer.suggestion[name] = space.suggestion
+            self._add_search_space(name, space.optimizer.level, space)
+            self.add_depends(space.optimizer.name, [name])
+        else:
+            if level is None:
+                raise ValueError('level must be provided.')
+            try:
+                space = Space.fromarray(space)
+            except:
+                pass
+            self._add_search_space(name, level, space)
+            if isinstance(space, Expression) or callable(space):
+                self.add_depends(name, space.symbols())
+        if setter:
+            self.description['setters'][name] = setter
+        if '.' in name:
+            self.add_depends('config', [name])
+
+    def trace(self,
+              name: str,
+              depends: list[str],
+              getter: Callable | None = None):
+        self.add_depends(name, depends)
+        if getter:
+            self.description['getters'][name] = getter
+
+    def minimize(self,
+                 name: str,
+                 level: int,
+                 method=NgOptimizer,
+                 maxiter=100,
+                 getter: Callable | None = None,
+                 **kwds) -> Optimizer:
+        assert level >= 0, 'level must be greater than or equal to 0.'
+        opt = Optimizer(self,
+                        name,
+                        level,
+                        method,
+                        maxiter,
+                        minimize=True,
+                        **kwds)
+        self.description['optimizers'][name] = opt
+        if getter:
+            self.description['getters'][name] = getter
+        return opt
+
+    def maximize(self,
+                 name: str,
+                 level: int,
+                 method=NgOptimizer,
+                 maxiter=100,
+                 getter: Callable | None = None,
+                 **kwds) -> Optimizer:
+        assert level >= 0, 'level must be greater than or equal to 0.'
+        opt = Optimizer(self,
+                        name,
+                        level,
+                        method,
+                        maxiter,
+                        minimize=False,
+                        **kwds)
+        self.description['optimizers'][name] = opt
+        if getter:
+            self.description['getters'][name] = getter
+        return opt
+
+    def _synchronize_config(self):
+        for key, value in self.variables.items():
+            if '.' in key:
+                d = self.config
+                ks = key.split('.')
+                if not ks:
+                    continue
+                for k in ks[:-1]:
+                    if k in d:
+                        d = d[k]
+                    else:
+                        d[k] = {}
+                        d = d[k]
+                d[ks[-1]] = value
+        return self.config
+
+    def _query_config(self, key):
+        d = self.config
+        for k in key.split('.'):
+            d = d[k]
+        return d
+
+    async def _update_progress(self):
+        while True:
+            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()
+
+    @contextlib.asynccontextmanager
+    async def _send_msg_and_update_bar(self):
+        send_msg_task = asyncio.create_task(self._send_msg())
+        update_progress_task = asyncio.create_task(self._update_progress())
+        try:
+            yield (send_msg_task, update_progress_task)
+        finally:
+            update_progress_task.cancel()
+            send_msg_task.cancel()
+            while True:
+                try:
+                    task = self._prm_queue.get_nowait()
+                except:
+                    break
+                try:
+                    task.cancel()
+                except:
+                    pass
+
+    async def _check_background_tasks(self):
+        for task in self._background_tasks:
+            if task.done():
+                await task
+
+    async def run(self):
+        assymbly(self.description)
+        self._background_tasks = ()
+
+        if isinstance(
+                self.description['database'],
+                str) and self.description['database'].startswith("tcp://"):
+            async with ZMQContextManager(zmq.DEALER,
+                                         connect=self.description['database'],
+                                         socket=self._sock) as socket:
+                self._sock = socket
+                async with self._send_msg_and_update_bar() as background_tasks:
+                    self._background_tasks = background_tasks
+                    await self._run()
+        else:
+            if self.config:
+                self.description['config'] = self._raw_config_copy
+            async with self._send_msg_and_update_bar() as background_tasks:
+                self._background_tasks = background_tasks
+                await self._run()
+
+    async def _run(self):
+        self._variables = {'self': self, 'config': self.config}
+
+        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)
+
+        updates = await call_many_functions(
+            self.description['order'].get(-1, []),
+            self.description['functions'], self.variables)
+        await update_variables(self.variables, updates,
+                               self.description['setters'])
+        await self._check_background_tasks()
+        await self.work()
+        for level, bar in self._bar.items():
+            bar.close()
+        await self._check_background_tasks()
+        if self._single_step:
+            self.variables.update(await call_many_functions(
+                self.description['order'].get(-1, []),
+                self.description['getters'], self.variables))
+
+            await self.emit(0, 0, 0, self.variables)
+            await self.emit(-1, 0, 0, {})
+        await self._check_background_tasks()
+        await self._prm_queue.join()
+        await self._msg_queue.join()
+        return self.variables
+
+    async def done(self):
+        if self._main_task is not None:
+            try:
+                await self._main_task
+            except asyncio.CancelledError:
+                pass
+
+    def finished(self):
+        return self._main_task.done()
+
+    def start(self):
+        import asyncio
+        self._main_task = asyncio.create_task(self.run())
+
+    async def submit(self, server=default_executor):
+        assymbly(self.description)
+        async with ZMQContextManager(zmq.DEALER,
+                                     connect=server,
+                                     socket=self._sock) as socket:
+            await socket.send_pyobj({
+                'method': 'task_submit',
+                'description': dill.dumps(self.description)
+            })
+            self.id = await socket.recv_pyobj()
+            await socket.send_pyobj({
+                'method': 'task_get_record_id',
+                'id': self.id
+            })
+            record_id = await socket.recv_pyobj()
+            self.record = Record(record_id, self.description['database'],
+                                 self.description)
+
+    def cancel(self):
+        if self._main_task is not None:
+            self._main_task.cancel()
+
+    async def _reset_progress_bar(self, level):
+        if level in self._bar:
+            self._bar[level].reset()
+
+    async def _update_progress_bar(self, level, n: int):
+        if level in self._bar:
+            self._bar[level].update(n)
+
+    async def iter(self, **kwds):
+        if self.current_level >= len(self.description['loops']):
+            return
+        step = 0
+        position = 0
+        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, []),
+                self.description['order'].get(self.current_level, []),
+                self.description['functions']
+                | {'config': self._synchronize_config},
+                self.description['optimizers'], self.description['setters'],
+                self.description['getters']):
+            await self._check_background_tasks()
+            self._current_level += 1
+            if await self._filter(variables, self.current_level - 1):
+                yield variables
+                self._single_step = False
+                await self.emit(self.current_level - 1, step, position,
+                                variables)
+                step += 1
+            position += 1
+            self._current_level -= 1
+            self._prm_queue.put_nowait(
+                self._update_progress_bar(self.current_level, 1))
+            await self._check_background_tasks()
+        if self.current_level == 0:
+            await self.emit(self.current_level - 1, 0, 0, {})
+            for name, value in self.variables.items():
+                if inspect.isawaitable(value):
+                    self.variables[name] = await value
+            await self._check_background_tasks()
+            await self._prm_queue.join()
+
+    async def work(self, **kwds):
+        if self.current_level in self.description['actions']:
+            action = self.description['actions'][self.current_level]
+            coro = action(self, **kwds)
+            if inspect.isawaitable(coro):
+                await coro
+        else:
+            async for variables in self.iter(**kwds):
+                await self.do_something(**kwds)
+
+    async def do_something(self, **kwds):
+        await self.work(**kwds)
+
+    def mount(self, action: Callable, level: int):
+        """
+        Mount a action to the scan.
+
+        Args:
+            action: A callable object.
+            level: The level of the scan to mount the action.
+        """
+        self.description['actions'][level] = action
+
+    async def promise(self, awaitable: Awaitable | Callable, *args,
+                      **kwds) -> Promise:
+        """
+        Promise to calculate asynchronous function and return the result in future.
+
+        Args:
+            awaitable: An awaitable object.
+
+        Returns:
+            Promise: A promise object.
+        """
+        if inspect.isawaitable(awaitable):
+            async with self._sem:
+                task = asyncio.create_task(self._await(awaitable))
+                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)
+            except:
+                return awaitable(*args, **kwds)
+            self._prm_queue.put_nowait(task)
+            return Promise(task)
+        else:
+            return awaitable
+
+    async def _await(self, awaitable: Awaitable):
+        async with self._sem:
+            return await awaitable
+
+
+def _get_environment(description):
+    import __main__
+    from IPython import get_ipython
+
+    if isinstance(description['namespace'], dict):
+        description['namespace'] = dump_globals()
+
+    ipy = get_ipython()
+    if ipy is not None:
+        description['entry']['shell'] = 'ipython'
+        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
+
+    description['entry']['env'] = {k: v for k, v in os.environ.items()}
+
+    return description
+
+
+def _mapping_levels(description):
+    mapping = {
+        label: level
+        for level, label in enumerate(
+            sorted(
+                set(description['loops'].keys())
+                | {k
+                   for k in description['actions'].keys() if k >= 0}))
+    }
+
+    if -1 in description['actions']:
+        mapping[-1] = max(mapping.values()) + 1
+
+    levels = sorted(mapping.values())
+    for k in description['actions'].keys():
+        if k < -1:
+            mapping[k] = levels[k]
+
+    description['loops'] = dict(
+        sorted([(mapping[k], v) for k, v in description['loops'].items()]))
+    description['actions'] = {
+        mapping[k]: v
+        for k, v in description['actions'].items()
+    }
+
+    for level, loops in description['loops'].items():
+        description['total'][level] = np.inf
+        for name, space in loops:
+            try:
+                description['total'][level] = min(description['total'][level],
+                                                  len(space))
+            except:
+                pass
+    return levels
+
+
+def _get_independent_variables(description):
+    independent_variables = set(description['intrinsic_loops'].keys())
+    for level, loops in description['loops'].items():
+        for name, iterable in loops:
+            if isinstance(iterable, (np.ndarray, list, tuple, range, Space)):
+                independent_variables.add(name)
+    return independent_variables
+
+
+def _build_dependents(description, levels, independent_variables):
+    dependents = copy.deepcopy(description['dependents'])
+    all_nodes = set(description['dependents'].keys())
+    for key, deps in dependents.items():
+        all_nodes.update(deps)
+
+    for level in levels:
+        range_list = description['loops'].get(level, [])
+        if level > 0:
+            if f'#__loop_{level}' not in description['dependents']:
+                dependents[f'#__loop_{level}'] = set()
+            dependents[f'#__loop_{level}'].add(f'#__loop_{level-1}')
+        for name, _ in range_list:
+            if name not in description['dependents']:
+                dependents[name] = set()
+            dependents[name].add(f'#__loop_{level}')
+
+    after_yield = set()
+    for key in all_nodes:
+        if key not in independent_variables and key not in description[
+                'consts']:
+            if key not in dependents:
+                after_yield.add(key)
+
+    def _get_all_depends(key, graph):
+        ret = set()
+        if key not in graph:
+            return ret
+
+        for e in graph[key]:
+            ret.update(_get_all_depends(e, graph))
+        ret.update(graph[key])
+        return ret
+
+    full_depends = {}
+    for key in dependents:
+        full_depends[key] = _get_all_depends(key, dependents)
+        if full_depends[key] & after_yield:
+            after_yield.add(key)
+
+    return dependents, full_depends, after_yield
+
+
+def _build_order(description, levels, dependents, full_depends):
+    levels = {}
+    passed = set()
+    all_keys = set(description['consts'].keys())
+    for key in dependents:
+        all_keys.add(key)
+        all_keys.update(dependents[key])
+    for level in reversed(description['loops'].keys()):
+        tag = f'#__loop_{level}'
+        for key, deps in full_depends.items():
+            all_keys.update(deps)
+            all_keys.add(key)
+            if key.startswith('#__loop_'):
+                continue
+            if tag in deps:
+                if level not in levels:
+                    levels[level] = set()
+                if key not in passed:
+                    passed.add(key)
+                    levels[level].add(key)
+    levels[-1] = {
+        key
+        for key in all_keys - passed if not key.startswith('#__loop_')
+    }
+
+    order = []
+    ts = TopologicalSorter(dependents)
+    ts.prepare()
+    while ts.is_active():
+        ready = ts.get_ready()
+        order.append(ready)
+        for k in ready:
+            ts.done(k)
+
+    description['order'] = {}
+
+    for level in sorted(levels):
+        keys = set(levels[level])
+        description['order'][level] = []
+        for ready in order:
+            ready = list(keys & set(ready))
+            if ready:
+                description['order'][level].append(ready)
+                keys -= set(ready)
+
+
+def _make_axis(description):
+    axis = {}
+
+    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
+            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'] = {
+        k: tuple([x for x in v if x >= 0])
+        for k, v in axis.items()
+    }
+
+
+def assymbly(description):
+    _get_environment(description)
+    levels = _mapping_levels(description)
+    independent_variables = _get_independent_variables(description)
+    description['independent_variables'] = independent_variables
+
+    dependents, full_depends, after_yield = _build_dependents(
+        description, levels, independent_variables)
+
+    _build_order(description, levels, dependents, full_depends)
+    _make_axis(description)
+
+    return description