Created
October 12, 2018 00:00
-
-
Save JevinJ/3d468dd947e5121a0ae8e9b515305a4f to your computer and use it in GitHub Desktop.
Interruptible Python multiprocessing pool.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import sys | |
| import multiprocessing | |
| import os | |
| import time | |
| class Pool(): | |
| def __init__(self): | |
| self.cpu_count = os.cpu_count() | |
| self.pool = [] | |
| self.pool_lock = multiprocessing.Lock() | |
| def map(self, target, iterable): | |
| try: | |
| while iterable: | |
| while iterable and len(self.pool) < os.cpu_count(): | |
| with self.pool_lock: | |
| try: | |
| p = multiprocessing.Process(target=target, args=(iterable.pop(),), daemon=True) | |
| p.start() | |
| self.pool.append(p) | |
| except IndexError: | |
| break | |
| self.join() | |
| except KeyboardInterrupt: | |
| self.join() | |
| def join(self): | |
| with self.pool_lock: | |
| for p in self.pool: | |
| p.join() | |
| self.pool = [] | |
| #Example processing function | |
| def run(num): | |
| try: | |
| print(num) | |
| time.sleep(1) | |
| except KeyboardInterrupt: | |
| raise | |
| def fetch(num): | |
| try: | |
| run(num) | |
| except: | |
| sys.exit() | |
| if __name__ == '__main__': | |
| nums = [i for i in range(100)] | |
| pool = Pool() | |
| pool.map(fetch, nums) | |
| print('done') |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment