def _create_closure(self):
@def_function.function()
def some_function():
return 1.0
return client.Closure(some_function)
def testProcessAtLeaseOnce(self):
closure_queue = client._CoordinatedClosureQueue(
MockCancellationManager())
labels = ['A', 'B', 'C', 'D', 'E']
processed_count = collections.defaultdict(int)
coord = coordinator.Coordinator(clean_stop_exception_types=[])
def process_queue():
with coord.stop_on_exception():
has_been_put_back = False
while True:
closure = closure_queue.get(timeout=30)
if closure is None:
break
if not has_been_put_back:
has_been_put_back = True
closure_queue.put_back(closure)
continue
closure._function()
closure_queue.mark_finished()
def get_func(label):
def func():
logging.info('Label: %s, before waiting 3 sec', label)
time.sleep(3)
processed_count[label] += 1
logging.info('Label: %s, after waiting 3 sec', label)
return func
for label in labels:
closure_queue.put(
client.Closure(get_func(label), MockCancellationManager()))
t1 = threading.Thread(target=process_queue, daemon=True)
t1.start()
t2 = threading.Thread(target=process_queue, daemon=True)
t2.start()
# Make sure multiple wait() calls are fine.
closure_queue.wait()
closure_queue.wait()
closure_queue.wait()
closure_queue.wait()
self.assertEqual(processed_count, collections.Counter(labels))
coord.join([t1, t2])
def testNotifyBeforeWait(self):
closure_queue = client._CoordinatedClosureQueue()
def func():
logging.info('func running')
coord = coordinator.Coordinator(clean_stop_exception_types=[])
def process_queue():
with coord.stop_on_exception():
closure = closure_queue.get()
closure_queue.mark_finished(closure)
closure_queue.put(client.Closure(func))
t = threading.Thread(target=process_queue)
t.start()
coord.join([t])
# This test asserts that waiting at the time the function has been processed
# doesn't time out.
closure_queue.wait()
def _create_closure(self, cancellation_mgr):
@def_function.function()
def some_function():
return 1.0
return client.Closure(some_function, cancellation_mgr)
def _create_closure(self):
@def_function.function()
def some_function():
return 1.0
return client.Closure(some_function, MockCancellationManager())
