1# Copyright 2021 The Chromium Authors 2# Use of this source code is governed by a BSD-style license that can be 3# found in the LICENSE file. 4 5from __future__ import print_function 6 7import contextlib 8from typing import Generator, Optional 9 10from pathos import pools 11 12 13def GetProcessPool(nodes: Optional[int] = None) -> pools.ProcessPool: 14 """Returns a pathos.pools.ProcessPool instance. 15 16 Split out for ease of unittesting since pathos can still run into pickling 17 issues with MagicMocks used in tests. 18 19 Args: 20 nodes: How many processes processes to spawn in the process pool. 21 22 Returns: 23 A pathos.pools.ProcessPool instance. 24 """ 25 return pools.ProcessPool(nodes=nodes) 26 27 28@contextlib.contextmanager 29def GetProcessPoolContext( 30 nodes: Optional[int] = None) -> Generator[pools.ProcessPool, None, None]: 31 try: 32 pool = GetProcessPool(nodes) 33 yield pool 34 finally: 35 pool.close() 36 pool.join() 37