def test_close_then_use_executor(self): ex = concurrent_executor.ConcurrentExecutor( eager_executor.EagerExecutor()) ex.close() result = self.use_executor(ex) self.assertIsInstance(result, eager_executor.EagerValue) self.assertEqual(result.internal_representation.numpy(), 11)
def test_multiple_computations_with_same_executor(self): @computations.tf_computation(tf.int32) def add_one(x): return tf.add(x, 1) ex = concurrent_executor.ConcurrentExecutor( eager_executor.EagerExecutor()) async def compute(): return await ex.create_selection(await ex.create_tuple( collections.OrderedDict([ ('a', await ex.create_call(await ex.create_value(add_one), await ex.create_value(10, tf.int32))) ])), name='a') result = asyncio.get_event_loop().run_until_complete(compute()) self.assertIsInstance(result, eager_executor.EagerValue) self.assertEqual(result.internal_representation.numpy(), 11) # After this call, the ConcurrentExecutor has been closed, and needs # to be re-initialized. ex.close() result = asyncio.get_event_loop().run_until_complete(compute()) self.assertIsInstance(result, eager_executor.EagerValue) self.assertEqual(result.internal_representation.numpy(), 11)
def test_close_then_use_executor_with_cache(self): # Integration that use after close is compatible with the combined # concurrent executors and cached executors. This was broken in # the past due to interactions between closing, caching, and the # concurrent executor. See b/148288711 for context. ex = concurrent_executor.ConcurrentExecutor( caching_executor.CachingExecutor(eager_executor.EagerExecutor())) self.use_executor(ex) ex.close() self.use_executor(ex)
def make_output(): test_ex = FakeExecutor() executors = [ concurrent_executor.ConcurrentExecutor(test_ex) for _ in range(10) ] loop = asyncio.get_event_loop() vals = [ex.create_value(idx) for idx, ex in enumerate(executors)] results = loop.run_until_complete(asyncio.gather(*vals)) self.assertCountEqual(list(results), list(range(10))) del executors return test_ex.output
def test_end_to_end(self): @computations.tf_computation(tf.int32) def add_one(x): return tf.add(x, 1) ex = concurrent_executor.ConcurrentExecutor( eager_executor.EagerExecutor()) set_default_executor.set_default_executor(ex) self.assertEqual(add_one(7), 8) # After this invocation, the ConcurrentExecutor has been closed, and needs # to be re-initialized. self.assertEqual(add_one(8), 9) set_default_executor.set_default_executor()
def _complete_stack(ex): return lambda_executor.LambdaExecutor( caching_executor.CachingExecutor( concurrent_executor.ConcurrentExecutor(ex)))
def _create_bottom_stack(): return lambda_executor.LambdaExecutor( caching_executor.CachingExecutor( concurrent_executor.ConcurrentExecutor( eager_executor.EagerExecutor())))
def test_something(self): # TODO(b/134543154): Actually test something. concurrent_executor.ConcurrentExecutor([])
def _create_single_worker_stack(): ex = eager_executor.EagerExecutor() ex = concurrent_executor.ConcurrentExecutor(ex) ex = caching_executor.CachingExecutor(ex) return lambda_executor.LambdaExecutor(ex)
def _make(n): return [concurrent_executor.ConcurrentExecutor(bottom_ex) for _ in range(n)]