meerschaum.utils.pool

 1#! /usr/bin/env python
 2# -*- coding: utf-8 -*-
 3# vim:fenc=utf-8
 4
 5"""
 6Global pools that are joined on exit.
 7"""
 8
 9from __future__ import annotations
10from meerschaum.utils.typing import Optional, Callable, List, Any
11from meerschaum.utils.threading import Lock, RLock
12import signal
13
14pools = {}
15_locks = {
16    'pools': Lock(),
17}
18
19def _initializer():
20    """Ignore keyboard interrupt in workers."""
21    signal.signal(signal.SIGINT, signal.SIG_IGN)
22
23def get_pool(
24        pool_class_name: str = 'ThreadPool',
25        workers: Optional[int] = None,
26        initializer: Optional[Callable[[None], None]] = None,
27        initargs: Optional[List[Any]] = None,
28    ):
29    """If the requested pool does not exist, instantiate it here.
30    Pools are joined and closed on exit."""
31    from multiprocessing import cpu_count
32    if workers is None:
33        workers = cpu_count()
34    pool_key = pool_class_name + f'-{workers}'
35
36    def build_pool(workers):
37        from meerschaum.utils.warnings import warn
38        from meerschaum.utils.packages import attempt_import
39        import importlib
40        try:
41            Pool = getattr(
42                importlib.import_module('multiprocessing.pool'),
43                pool_class_name
44            )
45        except Exception as e:
46            warn(e, stacklevel=3)
47            Pool = getattr(
48                importlib.import_module('multiprocessing.pool'),
49                'ThreadPool'
50            )
51
52        try:
53            pool = Pool(workers, initializer=initializer, initargs=initargs)
54        except Exception as e:
55            print(e)
56            pool = None
57
58        with _locks['pools']:
59            pools[pool_key] = pool
60
61    if pools.get(pool_key, None) is None:
62        build_pool(workers)
63
64    if (
65        pools[pool_key] is not None
66        and pools[pool_key]._state not in ('RUN', 0)
67    ):
68        try:
69            pools[pool_key].close()
70            pools[pool_key].terminate()
71        except Exception as e:
72            pass
73        del pools[pool_key]
74        build_pool(workers)
75
76    return pools[pool_key]
77
78
79def get_pools():
80    """Return the global pools dictionary."""
81    global pools
82    if pools is None:
83        with _locks['pools']:
84            pools = {}
85    return pools
86
87
88def get_pool_executor(workers: Optional[int] = None):
89    """ Return a new `ThreadPoolExecutor`. """
90    try:
91        from multiprocessing import cpu_count
92        from concurrent.futures import ThreadPoolExecutor
93        workers = cpu_count() if workers is None else workers
94    except Exception as e:
95        return None
96
97    return ThreadPoolExecutor(max_workers=workers) if ThreadPoolExecutor is not None else None