clearml/trains/utilities/async_manager.py
2019-07-20 23:55:14 +03:00

68 lines
2.4 KiB
Python

import os
import time
from threading import Lock
import six
class AsyncManagerMixin(object):
_async_results_lock = Lock()
# per pid (process) list of async jobs (support for sub-processes forking)
_async_results = {}
@classmethod
def _add_async_result(cls, result, wait_on_max_results=None, wait_time=30, wait_cb=None):
while True:
try:
cls._async_results_lock.acquire()
# discard completed results
pid = os.getpid()
cls._async_results[pid] = [r for r in cls._async_results.get(pid, []) if not r.ready()]
num_results = len(cls._async_results[pid])
if wait_on_max_results is not None and num_results >= wait_on_max_results:
# At least max_results results are still pending, wait
if wait_cb:
wait_cb(num_results)
if wait_time:
time.sleep(wait_time)
continue
# add result
if result and not result.ready():
if not cls._async_results.get(pid):
cls._async_results[pid] = []
cls._async_results[pid].append(result)
break
finally:
cls._async_results_lock.release()
@classmethod
def wait_for_results(cls, timeout=None, max_num_uploads=None):
remaining = timeout
count = 0
pid = os.getpid()
for r in cls._async_results.get(pid, []):
if r.ready():
continue
t = time.time()
# bugfix for python2.7 threading issues
if six.PY2 and not remaining:
while not r.ready():
r.wait(timeout=2.)
else:
r.wait(timeout=remaining)
count += 1
if max_num_uploads is not None and max_num_uploads - count <= 0:
break
if timeout is not None:
remaining = max(0., remaining - max(0., time.time() - t))
if not remaining:
break
@classmethod
def get_num_results(cls):
pid = os.getpid()
if cls._async_results.get(pid, []):
return len([r for r in cls._async_results.get(pid, []) if not r.ready()])
else:
return 0