from __future__ import print_function
import contextlib
from typing import Generator, Optional
from pathos import pools
def GetProcessPool(nodes: Optional[int] = None) -> pools.ProcessPool:
"""Returns a pathos.pools.ProcessPool instance.
Split out for ease of unittesting since pathos can still run into pickling
issues with MagicMocks used in tests.
Args:
nodes: How many processes processes to spawn in the process pool.
Returns:
A pathos.pools.ProcessPool instance.
"""
return pools.ProcessPool(nodes=nodes)
@contextlib.contextmanager
def GetProcessPoolContext(
nodes: Optional[int] = None) -> Generator[pools.ProcessPool, None, None]:
try:
pool = GetProcessPool(nodes)
yield pool
finally:
pool.close()
pool.join()
