Пример #1
0
class ProcessPool:
    def __init__(self, processors: int = None, **kwargs):
        if __debug__ and processors is None:
            processors = 1
        self._pool = Pool(processors, **kwargs)
        self._results = list()  # type: List[AsyncResult]

    def queue_task(self, task: Callable, *args, **kwargs) -> AsyncResult:
        res = self._pool.apply_async(task, args=args, kwds=kwargs)
        self._results.append(res)
        return res

    def get_async_results(self) -> List[AsyncResult]:
        slot = self._results
        self._results = type(self._results)()
        return slot

    def get_results(self) -> Iterator[Any]:
        return (i.get() for i in self.get_async_results())

    def check_exceptions(self) -> None:
        list(self.get_results())

    def __enter__(self):
        self._pool.__enter__()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        return self._pool.__exit__(exc_type, exc_val, exc_tb)

    def close(self):
        return self._pool.close()

    def join(self):
        return self._pool.join()
Пример #2
0
def run_cmds(cmds, cache_path, jittor_path, msg="run_cmds"):
    global pool_size, p
    bk = mp.current_process()._config.get('daemon')
    mp.current_process()._config['daemon'] = False
    if pool_size == 0:
        try:
            mem_bytes = os.sysconf('SC_PAGE_SIZE') * os.sysconf(
                'SC_PHYS_PAGES')
            mem_gib = mem_bytes / (1024.**3)
            pool_size = min(16, max(int(mem_gib // 3), 1))
            LOG.i(
                f"Total mem: {mem_gib:.2f}GB, using {pool_size} procs for compiling."
            )
        except ValueError:
            # On macOS, python with version lower than 3.9 do not support SC_PHYS_PAGES.
            # Use hard coded pool size instead.
            pool_size = 4
            LOG.i(f"using {pool_size} procs for compiling.")

        p = Pool(pool_size, initializer=pool_initializer)
        p.__enter__()
        import atexit
        atexit.register(pool_cleanup)
    cmds = [[cmd, cache_path, jittor_path] for cmd in cmds]
    try:
        n = len(cmds)
        dp = DelayProgress(msg, n)
        for i, _ in enumerate(p.imap_unordered(do_compile, cmds)):
            dp.update(i)
    finally:
        mp.current_process()._config['daemon'] = bk
Пример #3
0
class OptionalPool:
    def __init__(self, threads=1):
        self.pool = Pool(threads) if threads > 1 else None

    def map(self, func, iterable, chunksize=None):
        return self.pool.map(func, iterable,
                             chunksize) if self.pool is not None else map(
                                 func, iterable)

    def starmap(self, func, iterable, chunksize=None):
        return self.pool.starmap(func, iterable, chunksize) if self.pool is not None \
            else [func(*params) for params in iterable]

    def __enter__(self):
        if self.pool is not None:
            self.pool.__enter__()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.pool is not None:
            self.pool.__exit__(exc_type, exc_val, exc_tb)

    def close(self):
        if self.pool is not None:
            self.pool.close()
Пример #4
0
class LazyPool:
    """
    A version of :class:`multiprocessing.Pool` that limits the number of inputs
    being queued at the same time.
    At the beginning, ``processes * item_per_process`` inputs are sent
    to the pool. Further inputs are only sent if the previous ones have
    already been processed. This prevents outputs from being overgenerated
    and filling up the memory when the main processes cannot keep up.

    Note: ATM only :meth:`imap` is implemented.

    Note also that :class:`multiprocessing.Pool` cannot (easily?) be subclassed.
    So :class:`LazyPool` is not a subclass of :class:`multiprocessing.Pool`, just
    uses the latter. After Python 3.5, :class:`multiprocessing.pool.Pool` is
    exposed (or simply exists?), so it's a bit easier.
    """
    def __init__(self,
                 processes=None,
                 item_per_process=5,
                 initializer=None,
                 initargs=(),
                 maxtasksperchild=None):
        self.item_per_process = item_per_process
        self.max_items = item_per_process * processes
        self.pool = Pool(processes, initializer, initargs, maxtasksperchild)

    def __enter__(self):
        self.pool.__enter__()
        return self

    def __exit__(self, *args):
        return self.pool.__exit__(*args)

    def imap(self, func: Callable,
             iterable: Iterable) -> Generator[Tuple[Any, Any], None, None]:
        outputs = deque(maxlen=self.max_items)
        it = iter(iterable)
        # Add the first batch of inputs at the same time
        for inp in islice(it, self.max_items):
            outputs.append(self.pool.apply_async(func, (inp, )))
        # After that, add new input whenever one has been processed. We
        # wait for the first result to keep the order of the input and the
        # output corpus the same.
        for inp in it:
            outputs[0].wait()
            yield outputs.popleft().get()
            outputs.append(self.pool.apply_async(func, (inp, )))
        # Just consume the rest
        while outputs:
            outputs[0].wait()
            yield outputs.popleft().get()

    def __getattr__(self, name):
        return getattr(self.pool, name)
Пример #5
0
def run_cmds(cmds, cache_path, jittor_path, msg="run_cmds"):
    global pool_size, p
    bk = mp.current_process()._config.get('daemon')
    mp.current_process()._config['daemon'] = False
    if pool_size == 0:
        mem_bytes = os.sysconf('SC_PAGE_SIZE') * os.sysconf('SC_PHYS_PAGES')
        mem_gib = mem_bytes/(1024.**3)
        pool_size = min(16,max(int(mem_gib // 3), 1))
        LOG.i(f"Total mem: {mem_gib:.2f}GB, using {pool_size} procs for compiling.")
        p = Pool(pool_size)
        p.__enter__()
        import atexit
        atexit.register(pool_cleanup)
    cmds = [ [cmd, cache_path, jittor_path] for cmd in cmds ]
    try:
        n = len(cmds)
        dp = DelayProgress(msg, n)
        for i,_ in enumerate(p.imap_unordered(do_compile, cmds)):
            dp.update(i)
    finally:
        mp.current_process()._config['daemon'] = bk
Пример #6
0
def run_cmds(cmds, cache_path, jittor_path, msg="run_cmds"):
    global pool_size, p
    bk = mp.current_process()._config.get('daemon')
    mp.current_process()._config['daemon'] = False
    if pool_size == 0:
        try:
            mem_bytes = get_total_mem()
            mem_gib = mem_bytes / (1024.**3)
            pool_size = min(16, max(int(mem_gib // 3), 1))
            LOG.i(
                f"Total mem: {mem_gib:.2f}GB, using {pool_size} procs for compiling."
            )
        except ValueError:
            # On macOS, python with version lower than 3.9 do not support SC_PHYS_PAGES.
            # Use hard coded pool size instead.
            pool_size = 4
            LOG.i(f"using {pool_size} procs for compiling.")
        if os.name == 'nt':
            # a hack way to by pass windows
            # multiprocess spawn init_main_from_path.
            # check spawn.py:get_preparation_data
            spec_bk = sys.modules['__main__'].__spec__
            tmp = lambda x: x
            tmp.name = '__main__'
            sys.modules['__main__'].__spec__ = tmp
        p = Pool(pool_size, initializer=pool_initializer)
        p.__enter__()
        if os.name == 'nt':
            sys.modules['__main__'].__spec__ = spec_bk
        import atexit
        atexit.register(pool_cleanup)
    cmds = [[cmd, cache_path, jittor_path] for cmd in cmds]
    try:
        n = len(cmds)
        dp = DelayProgress(msg, n)
        for i, _ in enumerate(p.imap_unordered(do_compile, cmds)):
            dp.update(i)
    finally:
        mp.current_process()._config['daemon'] = bk