def test_with_one_arg_tf_comp_in_no_arg_fed_comp(self):
ex = reference_resolving_executor.ReferenceResolvingExecutor(
eager_tf_executor.EagerTFExecutor())
loop = asyncio.get_event_loop()
@computations.tf_computation(tf.int32)
def add_one(x):
return x + 1
@computations.federated_computation
def comp():
return add_one(10)
v1 = loop.run_until_complete(ex.create_value(comp))
v2 = loop.run_until_complete(ex.create_call(v1))
result = loop.run_until_complete(v2.compute())
self.assertEqual(result.numpy(), 11)
def test_with_tuples(self):
ex = reference_resolving_executor.ReferenceResolvingExecutor(
eager_tf_executor.EagerTFExecutor())
loop = asyncio.get_event_loop()
@computations.tf_computation(tf.int32, tf.int32)
def add_numbers(x, y):
return x + y
@computations.federated_computation
def comp():
return add_numbers(10, 20)
v1 = loop.run_until_complete(ex.create_value(comp))
v2 = loop.run_until_complete(ex.create_call(v1))
result = loop.run_until_complete(v2.compute())
self.assertEqual(result.numpy(), 30)
def test_context(rpc_mode='REQUEST_REPLY'):
port = portpicker.pick_unused_port()
server_pool = logging_pool.pool(max_workers=1)
server = grpc.server(server_pool)
server.add_insecure_port('[::]:{}'.format(port))
target_factory = executor_stacks.local_executor_factory(num_clients=3)
tracers = []
def _tracer_fn(cardinalities):
tracer = executor_test_utils.TracingExecutor(
target_factory.create_executor(cardinalities))
tracers.append(tracer)
return tracer
service = executor_service.ExecutorService(
executor_stacks.ResourceManagingExecutorFactory(_tracer_fn))
executor_pb2_grpc.add_ExecutorServicer_to_server(service, server)
server.start()
channel = grpc.insecure_channel('localhost:{}'.format(port))
stub = executor_pb2_grpc.ExecutorStub(channel)
serialized_cards = executor_service_utils.serialize_cardinalities(
{placement_literals.CLIENTS: 3})
stub.SetCardinalities(
executor_pb2.SetCardinalitiesRequest(cardinalities=serialized_cards))
remote_exec = remote_executor.RemoteExecutor(channel, rpc_mode)
executor = reference_resolving_executor.ReferenceResolvingExecutor(
remote_exec)
try:
yield collections.namedtuple('_', 'executor tracers')(executor,
tracers)
finally:
executor.close()
for tracer in tracers:
tracer.close()
try:
channel.close()
except AttributeError:
pass # Public gRPC channel doesn't support close()
finally:
server.stop(None)
def test_with_nested_lambdas(self):
ex = reference_resolving_executor.ReferenceResolvingExecutor(
eager_tf_executor.EagerTFExecutor())
@tensorflow_computation.tf_computation(tf.int32, tf.int32)
def add_numbers(x, y):
return x + y
@federated_computation.federated_computation(tf.int32)
def comp(x):
@federated_computation.federated_computation(tf.int32)
def nested_comp(y):
return add_numbers(x, y)
return nested_comp(1)
v1 = asyncio.run(ex.create_value(comp))
v2 = asyncio.run(ex.create_value(10, tf.int32))
v3 = asyncio.run(ex.create_call(v1, v2))
result = asyncio.run(v3.compute())
self.assertEqual(result.numpy(), 11)
def test_with_functional_parameter(self):
ex = reference_resolving_executor.ReferenceResolvingExecutor(
eager_tf_executor.EagerTFExecutor())
@tensorflow_computation.tf_computation(tf.int32)
def add_one(x):
return x + 1
@federated_computation.federated_computation(
computation_types.FunctionType(tf.int32, tf.int32), tf.int32)
def comp(f, x):
return f(f(x))
v1 = asyncio.run(ex.create_value(comp))
v2 = asyncio.run(ex.create_value(add_one))
v3 = asyncio.run(ex.create_value(10, tf.int32))
v4 = asyncio.run(
ex.create_struct(structure.Struct([('f', v2), ('x', v3)])))
v5 = asyncio.run(ex.create_call(v1, v4))
result = asyncio.run(v5.compute())
self.assertEqual(result.numpy(), 12)
def test_context():
port = portpicker.pick_unused_port()
server_pool = logging_pool.pool(max_workers=1)
server = grpc.server(server_pool)
server.add_insecure_port('[::]:{}'.format(port))
target_factory = executor_test_utils.LocalTestExecutorFactory(
default_num_clients=3)
tracers = []
def _tracer_fn(cardinalities):
tracer = executor_test_utils.TracingExecutor(
target_factory.create_executor(cardinalities))
tracers.append(tracer)
return tracer
service = executor_service.ExecutorService(
executor_test_utils.BasicTestExFactory(_tracer_fn))
executor_pb2_grpc.add_ExecutorGroupServicer_to_server(service, server)
server.start()
channel = grpc.insecure_channel('localhost:{}'.format(port))
stub = remote_executor_grpc_stub.RemoteExecutorGrpcStub(channel)
remote_exec = remote_executor.RemoteExecutor(stub)
remote_exec.set_cardinalities({placements.CLIENTS: 3})
executor = reference_resolving_executor.ReferenceResolvingExecutor(
remote_exec)
try:
yield collections.namedtuple('_', 'executor tracers')(executor,
tracers)
finally:
executor.close()
for tracer in tracers:
tracer.close()
try:
channel.close()
except AttributeError:
pass # Public gRPC channel doesn't support close()
finally:
server.stop(None)
def create_executor(self, cardinalities):
device = tf.config.list_logical_devices(device_type='CPU')[0]
num_requested_clients = cardinalities.get(placements.CLIENTS)
if num_requested_clients is None:
num_clients = self._default_num_clients
else:
num_clients = num_requested_clients
client_stacks = [
self._leaf_executor_fn(device=device) for _ in range(num_clients)
]
federating_strategy_factory = federated_resolving_strategy.FederatedResolvingStrategy.factory(
{
placements.CLIENTS: client_stacks,
placements.SERVER: self._leaf_executor_fn(device=device),
},
local_computation_factory=tensorflow_computation_factory.
TensorFlowComputationFactory())
executor = federating_executor.FederatingExecutor(
federating_strategy_factory, self._leaf_executor_fn(device=device))
return reference_resolving_executor.ReferenceResolvingExecutor(
executor)
def test_with_one_arg_tf_comp_in_two_arg_fed_comp(self):
ex = reference_resolving_executor.ReferenceResolvingExecutor(
eager_tf_executor.EagerTFExecutor())
@tensorflow_computation.tf_computation(tf.int32, tf.int32)
def add_numbers(x, y):
return x + y
@federated_computation.federated_computation(tf.int32, tf.int32)
def comp(x, y):
return add_numbers(x, x), add_numbers(x, y), add_numbers(y, y)
v1 = asyncio.run(ex.create_value(comp))
v2 = asyncio.run(ex.create_value(10, tf.int32))
v3 = asyncio.run(ex.create_value(20, tf.int32))
v4 = asyncio.run(
ex.create_struct(structure.Struct([('x', v2), ('y', v3)])))
v5 = asyncio.run(ex.create_call(v1, v4))
result = asyncio.run(v5.compute())
self.assertEqual(
str(structure.map_structure(lambda x: x.numpy(), result)),
'<20,30,40>')
def test_with_federated_map(self):
eager_ex = eager_tf_executor.EagerTFExecutor()
factory = federated_resolving_strategy.FederatedResolvingStrategy.factory(
{placement_literals.SERVER: eager_ex})
federated_ex = federating_executor.FederatingExecutor(factory, eager_ex)
ex = reference_resolving_executor.ReferenceResolvingExecutor(federated_ex)
loop = asyncio.get_event_loop()
@computations.tf_computation(tf.int32)
def add_one(x):
return x + 1
@computations.federated_computation(type_factory.at_server(tf.int32))
def comp(x):
return intrinsics.federated_map(add_one, x)
v1 = loop.run_until_complete(ex.create_value(comp))
v2 = loop.run_until_complete(
ex.create_value(10, type_factory.at_server(tf.int32)))
v3 = loop.run_until_complete(ex.create_call(v1, v2))
result = loop.run_until_complete(v3.compute())
self.assertEqual(result.numpy(), 11)
def test_with_functional_parameter(self):
ex = reference_resolving_executor.ReferenceResolvingExecutor(
eager_tf_executor.EagerTFExecutor())
loop = asyncio.get_event_loop()
@computations.tf_computation(tf.int32)
def add_one(x):
return x + 1
@computations.federated_computation(
computation_types.FunctionType(tf.int32, tf.int32), tf.int32)
def comp(f, x):
return f(f(x))
v1 = loop.run_until_complete(ex.create_value(comp))
v2 = loop.run_until_complete(ex.create_value(add_one))
v3 = loop.run_until_complete(ex.create_value(10, tf.int32))
v4 = loop.run_until_complete(
ex.create_struct(structure.Struct([(None, v2), (None, v3)])))
v5 = loop.run_until_complete(ex.create_call(v1, v4))
result = loop.run_until_complete(v5.compute())
self.assertEqual(result.numpy(), 12)
def test_with_one_arg_tf_comp_in_two_arg_fed_comp(self):
ex = reference_resolving_executor.ReferenceResolvingExecutor(
eager_tf_executor.EagerTFExecutor())
loop = asyncio.get_event_loop()
@computations.tf_computation(tf.int32, tf.int32)
def add_numbers(x, y):
return x + y
@computations.federated_computation(tf.int32, tf.int32)
def comp(x, y):
return add_numbers(x, x), add_numbers(x, y), add_numbers(y, y)
v1 = loop.run_until_complete(ex.create_value(comp))
v2 = loop.run_until_complete(ex.create_value(10, tf.int32))
v3 = loop.run_until_complete(ex.create_value(20, tf.int32))
v4 = loop.run_until_complete(
ex.create_struct(structure.Struct([(None, v2), (None, v3)])))
v5 = loop.run_until_complete(ex.create_call(v1, v4))
result = loop.run_until_complete(v5.compute())
self.assertEqual(
str(structure.map_structure(lambda x: x.numpy(), result)),
'<20,30,40>')
def test_with_federated_map(self):
eager_ex = eager_tf_executor.EagerTFExecutor()
factory = federated_resolving_strategy.FederatedResolvingStrategy.factory(
{placements.SERVER: eager_ex})
federated_ex = federating_executor.FederatingExecutor(
factory, eager_ex)
ex = reference_resolving_executor.ReferenceResolvingExecutor(
federated_ex)
@tensorflow_computation.tf_computation(tf.int32)
def add_one(x):
return x + 1
@federated_computation.federated_computation(
computation_types.at_server(tf.int32))
def comp(x):
return intrinsics.federated_map(add_one, x)
v1 = asyncio.run(ex.create_value(comp))
v2 = asyncio.run(
ex.create_value(10, computation_types.at_server(tf.int32)))
v3 = asyncio.run(ex.create_call(v1, v2))
result = asyncio.run(v3.compute())
self.assertEqual(result.numpy(), 11)
def create_test_executor_factory():
executor = eager_tf_executor.EagerTFExecutor()
executor = caching_executor.CachingExecutor(executor)
executor = reference_resolving_executor.ReferenceResolvingExecutor(
executor)
return executor_factory.ExecutorFactoryImpl(lambda _: executor)
def create_bottom_stack():
executor = eager_tf_executor.EagerTFExecutor()
return reference_resolving_executor.ReferenceResolvingExecutor(
executor)
def test_runs_tf(self):
executor_test_utils.test_runs_tf(
self,
reference_resolving_executor.ReferenceResolvingExecutor(
eager_tf_executor.EagerTFExecutor()))
def create_test_executor_factory():
executor = eager_tf_executor.EagerTFExecutor()
executor = reference_resolving_executor.ReferenceResolvingExecutor(
executor)
return executor_test_utils.BasicTestExFactory(executor)
def _create_bottom_stack():
return reference_resolving_executor.ReferenceResolvingExecutor(
caching_executor.CachingExecutor(
thread_delegating_executor.ThreadDelegatingExecutor(
eager_tf_executor.EagerTFExecutor())))
def _create_middle_stack(children):
return reference_resolving_executor.ReferenceResolvingExecutor(
caching_executor.CachingExecutor(
composing_executor.ComposingExecutor(_create_bottom_stack(),
children)))
def create_test_executor_factory():
executor = eager_tf_executor.EagerTFExecutor()
executor = reference_resolving_executor.ReferenceResolvingExecutor(
executor)
return executor_stacks.ResourceManagingExecutorFactory(lambda _: executor)
def _complete_stack(ex):
return reference_resolving_executor.ReferenceResolvingExecutor(
caching_executor.CachingExecutor(
thread_delegating_executor.ThreadDelegatingExecutor(ex)))
def _create_middle_stack(children):
factory = federated_composing_strategy.FederatedComposingStrategy.factory(
_create_bottom_stack(), children)
executor = federating_executor.FederatingExecutor(factory,
_create_bottom_stack())
return reference_resolving_executor.ReferenceResolvingExecutor(executor)
