Threaded cached implementation of dask.get Parameters ---------- dsk: dict A dask dictionary specifying a workflow keys: key or list of keys Keys corresponding to desired data num_workers: integer of thread count The number of threads to use in the Threa
(
dsk: Mapping,
keys: Sequence[Key] | Key,
cache=None,
num_workers=None,
pool=None,
**kwargs,
)
| 36 | |
| 37 | |
| 38 | def get( |
| 39 | dsk: Mapping, |
| 40 | keys: Sequence[Key] | Key, |
| 41 | cache=None, |
| 42 | num_workers=None, |
| 43 | pool=None, |
| 44 | **kwargs, |
| 45 | ): |
| 46 | """Threaded cached implementation of dask.get |
| 47 | |
| 48 | Parameters |
| 49 | ---------- |
| 50 | |
| 51 | dsk: dict |
| 52 | A dask dictionary specifying a workflow |
| 53 | keys: key or list of keys |
| 54 | Keys corresponding to desired data |
| 55 | num_workers: integer of thread count |
| 56 | The number of threads to use in the ThreadPool that will actually execute tasks |
| 57 | cache: dict-like (optional) |
| 58 | Temporary storage of results |
| 59 | |
| 60 | Examples |
| 61 | -------- |
| 62 | >>> inc = lambda x: x + 1 |
| 63 | >>> add = lambda x, y: x + y |
| 64 | >>> dsk = {'x': 1, 'y': 2, 'z': (inc, 'x'), 'w': (add, 'z', 'y')} |
| 65 | >>> get(dsk, 'w') |
| 66 | 4 |
| 67 | >>> get(dsk, ['w', 'y']) |
| 68 | (4, 2) |
| 69 | """ |
| 70 | global default_pool |
| 71 | pool = pool or config.get("pool", None) |
| 72 | num_workers = num_workers or config.get("num_workers", None) |
| 73 | thread = current_thread() |
| 74 | |
| 75 | with pools_lock: |
| 76 | if pool is None: |
| 77 | if num_workers is None and thread is main_thread: |
| 78 | if default_pool is None: |
| 79 | default_pool = ThreadPoolExecutor(CPU_COUNT) |
| 80 | atexit.register(default_pool.shutdown) |
| 81 | pool = default_pool |
| 82 | elif thread in pools and num_workers in pools[thread]: |
| 83 | pool = pools[thread][num_workers] |
| 84 | else: |
| 85 | pool = ThreadPoolExecutor(num_workers) |
| 86 | atexit.register(pool.shutdown) |
| 87 | pools[thread][num_workers] = pool |
| 88 | elif isinstance(pool, multiprocessing.pool.Pool): |
| 89 | pool = MultiprocessingPoolExecutor(pool) |
| 90 | |
| 91 | results = get_async( |
| 92 | pool.submit, |
| 93 | pool._max_workers, |
| 94 | dsk, |
| 95 | keys, |