Spaces:
Sleeping
Sleeping
File size: 1,557 Bytes
c42fe7e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
import platform
import re
import traceback
from torch.multiprocessing import Manager, Process, current_process, get_context
is_main_process = not bool(re.match(r'((.*Process)|(SyncManager)|(.*PoolWorker))-\d+', current_process().name))
def main_process_print(self, *args, sep=' ', end='\n', file=None):
if is_main_process:
print(self, *args, sep=sep, end=end, file=file)
def chunked_worker_run(map_func, args, results_queue=None):
for a in args:
# noinspection PyBroadException
try:
res = map_func(*a)
results_queue.put(res)
except KeyboardInterrupt:
break
except Exception:
traceback.print_exc()
results_queue.put(None)
def chunked_multiprocess_run(map_func, args, num_workers, q_max_size=1000):
num_jobs = len(args)
if num_jobs < num_workers:
num_workers = num_jobs
queues = [Manager().Queue(maxsize=q_max_size // num_workers) for _ in range(num_workers)]
if platform.system().lower() != 'windows':
process_creation_func = get_context('spawn').Process
else:
process_creation_func = Process
workers = []
for i in range(num_workers):
worker = process_creation_func(
target=chunked_worker_run, args=(map_func, args[i::num_workers], queues[i]), daemon=True
)
workers.append(worker)
worker.start()
for i in range(num_jobs):
yield queues[i % num_workers].get()
for worker in workers:
worker.join()
worker.close()
|