def test_raises_with_closure(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.federated_computation(tf.int32,
type_factory.at_server(tf.int32))
def foo(x, y):
@computations.federated_computation(tf.int32)
def bar(z):
del z
return x
return intrinsics.federated_map(bar, y)
v1 = loop.run_until_complete(ex.create_value(foo))
v2 = loop.run_until_complete(
ex.create_value([0, 0],
[tf.int32, type_factory.at_server(tf.int32)]))
with self.assertRaisesRegex(
RuntimeError,
'lambda passed to intrinsic contains references to captured variables'):
loop.run_until_complete(ex.create_call(v1, v2))
def test_with_federated_map(self):
eager_ex = eager_tf_executor.EagerTFExecutor()
federated_ex = federating_executor.FederatingExecutor({
None:
eager_ex,
placement_literals.SERVER:
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 create_executor(
self, cardinalities: executor_factory.CardinalitiesType
) -> executor_base.Executor:
"""Constructs a federated executor with requested cardinalities."""
num_clients = self._validate_requested_clients(cardinalities)
num_client_executors = math.ceil(num_clients /
self._clients_per_thread)
client_stacks = [
self._unplaced_executor_factory.create_executor(
cardinalities={}, placement=placements.CLIENTS)
for _ in range(num_client_executors)
]
if self._use_sizing:
client_stacks = [
sizing_executor.SizingExecutor(ex) for ex in client_stacks
]
self._sizing_executors.extend(client_stacks)
federating_strategy_factory = self._federated_strategy_factory(
{
placements.CLIENTS: [
client_stacks[k % len(client_stacks)]
for k in range(num_clients)
],
placements.SERVER:
self._unplaced_executor_factory.create_executor(
placement=placements.SERVER),
},
local_computation_factory=self._local_computation_factory)
unplaced_executor = self._unplaced_executor_factory.create_executor()
executor = federating_executor.FederatingExecutor(
federating_strategy_factory, unplaced_executor)
return _wrap_executor_in_threading_stack(executor)
def test_with_federated_map_and_broadcast(self):
eager_ex = eager_tf_executor.EagerTFExecutor()
factory = federated_resolving_strategy.FederatedResolvingStrategy.factory({
placement_literals.SERVER: eager_ex,
placement_literals.CLIENTS: [eager_ex for _ in range(3)]
})
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,
intrinsics.federated_broadcast(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.assertCountEqual([x.numpy() for x in result], [11, 11, 11])
def test_with_federated_map_and_broadcast(self):
eager_ex = eager_tf_executor.EagerTFExecutor()
factory = federated_resolving_strategy.FederatedResolvingStrategy.factory(
{
placements.SERVER: eager_ex,
placements.CLIENTS: [eager_ex for _ in range(3)]
})
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,
intrinsics.federated_broadcast(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.assertCountEqual([x.numpy() for x in result], [11, 11, 11])
def create_executor(
self, cardinalities: executor_factory.CardinalitiesType
) -> executor_base.Executor:
"""Constructs a federated executor with requested cardinalities."""
num_clients = self._validate_requested_clients(cardinalities)
client_stacks = [
self._unplaced_executor_factory.create_executor(
cardinalities={}, placement=placement_literals.CLIENTS)
for _ in range(self._num_client_executors)
]
if self._use_sizing:
client_stacks = [
sizing_executor.SizingExecutor(ex) for ex in client_stacks
]
self._sizing_executors.extend(client_stacks)
federating_strategy_factory = federated_resolving_strategy.FederatedResolvingStrategy.factory(
{
placement_literals.CLIENTS: [
client_stacks[k % len(client_stacks)]
for k in range(num_clients)
],
placement_literals.SERVER:
self._unplaced_executor_factory.create_executor(
cardinalities={}, placement=placement_literals.SERVER),
})
unplaced_executor = self._unplaced_executor_factory.create_executor(
cardinalities={})
executor = federating_executor.FederatingExecutor(
federating_strategy_factory, unplaced_executor)
return _wrap_executor_in_threading_stack(executor)
def test_raises_with_closure(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)
@federated_computation.federated_computation(
tf.int32, computation_types.at_server(tf.int32))
def foo(x, y):
@federated_computation.federated_computation(tf.int32)
def bar(z):
del z
return x
return intrinsics.federated_map(bar, y)
v1 = asyncio.run(ex.create_value(foo))
v2 = asyncio.run(
ex.create_value(structure.Struct([
('x', 0), ('y', 0)
]), [tf.int32, computation_types.at_server(tf.int32)]))
with self.assertRaisesRegex(
RuntimeError,
'lambda passed to intrinsic contains references to captured variables'
):
asyncio.run(ex.create_call(v1, v2))
def test_raises_type_error_with_no_target_executor_unplaced(self):
factory = federated_resolving_strategy.FederatedResolvingStrategy.factory({
placement_literals.SERVER: eager_tf_executor.EagerTFExecutor(),
placement_literals.CLIENTS: eager_tf_executor.EagerTFExecutor(),
})
with self.assertRaises(TypeError):
federating_executor.FederatingExecutor(factory, None)
def _create_worker_stack():
return federating_executor.FederatingExecutor({
placement_literals.SERVER:
_create_bottom_stack(),
placement_literals.CLIENTS: [_create_bottom_stack() for _ in range(2)],
None:
_create_bottom_stack()
})
def _create_worker_stack():
factory = federated_resolving_strategy.FederatedResovlingStrategy.factory({
placement_literals.SERVER:
_create_bottom_stack(),
placement_literals.CLIENTS: [_create_bottom_stack() for _ in range(2)],
})
return federating_executor.FederatingExecutor(factory,
_create_bottom_stack())
def test_raises_value_error_with_no_target_executor_unplaced(self):
executor = federating_executor.FederatingExecutor({
placement_literals.SERVER: eager_tf_executor.EagerTFExecutor(),
placement_literals.CLIENTS: eager_tf_executor.EagerTFExecutor(),
})
value, type_signature = executor_test_utils.create_dummy_value_unplaced()
with self.assertRaises(ValueError):
self.run_sync(executor.create_value(value, type_signature))
def test_raises_value_error_with_no_target_executor_clients(
self, value, type_signature):
executor = federating_executor.FederatingExecutor({
placement_literals.SERVER: eager_tf_executor.EagerTFExecutor(),
None: eager_tf_executor.EagerTFExecutor()
})
with self.assertRaises(ValueError):
self.run_sync(executor.create_value(value, type_signature))
def test_raises_value_error_with_no_target_executor_clients(
self, value, type_signature):
factory = federated_resolving_strategy.FederatedResolvingStrategy.factory({
placement_literals.SERVER: eager_tf_executor.EagerTFExecutor(),
})
executor = federating_executor.FederatingExecutor(
factory, eager_tf_executor.EagerTFExecutor())
with self.assertRaises(ValueError):
self.run_sync(executor.create_value(value, type_signature))
def create_worker_stack():
factroy = federated_resolving_strategy.FederatedResolvingStrategy.factory({
placements.SERVER:
create_bottom_stack(),
placements.CLIENTS: [
create_bottom_stack() for _ in range(clients_per_stack)
],
})
return federating_executor.FederatingExecutor(factroy,
create_bottom_stack())
def create_test_executor(
num_clients: int = 3) -> federating_executor.FederatingExecutor:
executor = eager_tf_executor.EagerTFExecutor()
# Note: A `ReferenceResolvingExecutor` is required to be below a
# `FederatingExecutor` because intrinsics often take `Lambda`s as arguments.
executor = reference_resolving_executor.ReferenceResolvingExecutor(executor)
return federating_executor.FederatingExecutor({
placement_literals.SERVER: executor,
placement_literals.CLIENTS: [executor] * num_clients,
None: executor
})
def test_raises_value_error_with_no_target_executor_server(
self, value, type_signature):
factory = federated_resolving_strategy.FederatedResolvingStrategy.factory({
placements.CLIENTS: eager_tf_executor.EagerTFExecutor(),
})
executor = federating_executor.FederatingExecutor(
factory, eager_tf_executor.EagerTFExecutor())
value, type_signature = executor_test_utils.create_whimsy_value_at_server()
with self.assertRaises(ValueError):
self.run_sync(executor.create_value(value, type_signature))
def create_test_federated_stack(
num_clients=3) -> federating_executor.FederatingExecutor:
def create_bottom_stack():
executor = eager_tf_executor.EagerTFExecutor()
return reference_resolving_executor.ReferenceResolvingExecutor(executor)
factory = federated_resolving_strategy.FederatedResolvingStrategy.factory({
placements.SERVER: create_bottom_stack(),
placements.CLIENTS: [create_bottom_stack() for _ in range(num_clients)],
})
return federating_executor.FederatingExecutor(factory, create_bottom_stack())
def _create_composing_stack(
self, *, server_executor: executor_base.Executor,
target_executors: Sequence[executor_base.Executor]
) -> executor_base.Executor:
composing_strategy_factory = federated_composing_strategy.FederatedComposingStrategy.factory(
server_executor, target_executors)
unplaced_executor = self._unplaced_ex_factory.create_executor()
composing_executor = federating_executor.FederatingExecutor(
composing_strategy_factory, unplaced_executor)
threaded_composing_executor = _wrap_executor_in_threading_stack(
composing_executor)
return threaded_composing_executor
def _create_composite_stack(
target_executors, unplaced_ex_factory: UnplacedExecutorFactory
) -> executor_base.Executor:
"""Creates a single composite stack."""
server_executor = unplaced_ex_factory.create_executor(
placement=placement_literals.SERVER)
federating_strategy_factory = federated_composing_strategy.FederatedComposingStrategy.factory(
server_executor, target_executors)
unplaced_executor = unplaced_ex_factory.create_executor(
placement=placement_literals.SERVER)
executor = federating_executor.FederatingExecutor(
federating_strategy_factory, unplaced_executor)
return _wrap_executor_in_threading_stack(executor)
def _make_test_executor(
num_clients=1,
use_reference_resolving_executor=False,
) -> federating_executor.FederatingExecutor:
bottom_ex = eager_tf_executor.EagerTFExecutor()
if use_reference_resolving_executor:
bottom_ex = reference_resolving_executor.ReferenceResolvingExecutor(
bottom_ex)
return federating_executor.FederatingExecutor({
placements.SERVER: bottom_ex,
placements.CLIENTS: [bottom_ex for _ in range(num_clients)],
None: bottom_ex
})
def _create_composing_stack(
self, *, target_executors: Sequence[executor_base.Executor]
) -> executor_base.Executor:
server_executor = self._unplaced_ex_factory.create_executor(
placement=placement_literals.SERVER)
composing_strategy_factory = federated_composing_strategy.FederatedComposingStrategy.factory(
server_executor, target_executors)
unplaced_executor = self._unplaced_ex_factory.create_executor()
composing_executor = federating_executor.FederatingExecutor(
composing_strategy_factory, unplaced_executor)
threaded_composing_executor = _wrap_executor_in_threading_stack(
composing_executor, can_resolve_references=False)
return threaded_composing_executor
def test_recovers_from_raising(self):
class _RaisingExecutor(eager_tf_executor.EagerTFExecutor):
"""An executor which can be configured to raise on `create_value`."""
def __init__(self):
self._should_raise = True
super().__init__()
def stop_raising(self):
self._should_raise = False
async def create_value(self, *args, **kwargs):
if self._should_raise:
raise AssertionError
return await super().create_value(*args, **kwargs)
raising_executors = [_RaisingExecutor() for _ in range(2)]
factory = federated_resolving_strategy.FederatedResolvingStrategy.factory({
placements.SERVER: _create_worker_stack(),
placements.CLIENTS: raising_executors,
})
federating_ex = federating_executor.FederatingExecutor(
factory, _create_worker_stack())
raising_stacks = [federating_ex for _ in range(3)]
executor = _create_middle_stack([
_create_middle_stack(raising_stacks),
_create_middle_stack(raising_stacks),
])
@computations.federated_computation(
computation_types.at_clients(tf.float32))
def comp(x):
return intrinsics.federated_mean(x)
# 2 clients per worker stack * 3 worker stacks * 2 middle stacks
num_clients = 12
arg = [float(x + 1) for x in range(num_clients)]
with self.assertRaises(AssertionError):
_invoke(executor, comp, arg)
for ex in raising_executors:
ex.stop_raising()
result = _invoke(executor, comp, arg)
self.assertEqual(result, 6.5)
def create_test_executor(
number_of_clients: int = 3) -> federating_executor.FederatingExecutor:
def create_bottom_stack():
executor = eager_tf_executor.EagerTFExecutor()
return reference_resolving_executor.ReferenceResolvingExecutor(executor)
return federating_executor.FederatingExecutor({
placement_literals.SERVER: create_bottom_stack(),
placement_literals.CLIENTS: [
create_bottom_stack() for _ in range(number_of_clients)
],
None: create_bottom_stack()
})
def create_executor(
self, cardinalities: executor_factory.CardinalitiesType
) -> executor_base.Executor:
"""Constructs a federated executor with requested cardinalities."""
num_clients = self._validate_requested_clients(cardinalities)
client_stacks = [
self._unplaced_executor_factory.create_executor(
cardinalities={}, placement=placement_literals.CLIENTS)
for _ in range(self._num_client_executors)
]
if self._use_sizing:
client_stacks = [
sizing_executor.SizingExecutor(ex) for ex in client_stacks
]
self._sizing_executors.extend(client_stacks)
paillier_stack = self._unplaced_executor_factory.create_executor(
cardinalities={}, placement=paillier_placement.AGGREGATOR)
if self._use_sizing:
paillier_stack = sizing_executor.SizingExecutor(paillier_stack)
# Set up secure channel between clients & Paillier executor
secure_channel_grid = channels.ChannelGrid({
(tff.CLIENTS, paillier_placement.AGGREGATOR):
channels.EasyBoxChannel,
(tff.CLIENTS, tff.SERVER):
channels.PlaintextChannel,
(paillier_placement.AGGREGATOR, tff.SERVER):
channels.PlaintextChannel
})
# Build a FederatingStrategy factory for Paillier aggregation with the secure channel setup
strategy_factory = paillier_strategy.PaillierAggregatingStrategy.factory(
{
placement_literals.CLIENTS: [
client_stacks[k % len(client_stacks)]
for k in range(num_clients)
],
placement_literals.SERVER:
self._unplaced_executor_factory.create_executor(
cardinalities={}, placement=placement_literals.SERVER),
paillier_placement.AGGREGATOR:
paillier_stack,
},
channel_grid=secure_channel_grid,
# NOTE: we let the server generate it's own key here, but for proper
# deployment we would want to supply a key verified by proper PKI
key_inputter=paillier_comp.make_keygen(modulus_bitlength=2048))
unplaced_executor = self._unplaced_executor_factory.create_executor(
cardinalities={})
executor = federating_executor.FederatingExecutor(
strategy_factory, unplaced_executor)
return executor_stacks._wrap_executor_in_threading_stack(executor)
def test_raises_value_error_with_no_target_executor_unplaced_type(self):
# A `ValueError` will be raised because `create_value` can not find a target
# executor of the appropriate placement, because the following
# `federating_executor.FederatingExecutor` was created without one.
executor = federating_executor.FederatingExecutor({
placement_literals.SERVER:
eager_tf_executor.EagerTFExecutor(),
placement_literals.CLIENTS:
eager_tf_executor.EagerTFExecutor(),
})
value, type_signature = create_dummy_unplaced_type()
with self.assertRaises(ValueError):
self.run_sync(executor.create_value(value, type_signature))
def _create_sizing_stack(num_clients, num_client_executors):
"""Constructs local stack with sizing wired in at each client."""
sizing_stacks = [
sizing_executor.SizingExecutor(_create_bottom_stack())
for _ in range(num_client_executors)
]
executor_dict = {
placement_literals.CLIENTS: [
sizing_stacks[k % len(sizing_stacks)] for k in range(num_clients)
],
placement_literals.SERVER: _create_bottom_stack(),
None: _create_bottom_stack()
}
return _complete_stack(
federating_executor.FederatingExecutor(executor_dict)), sizing_stacks
def create_test_executor(
num_clients=1,
use_reference_resolving_executor=False
) -> federating_executor.FederatingExecutor:
bottom_ex = eager_tf_executor.EagerTFExecutor()
if use_reference_resolving_executor:
bottom_ex = reference_resolving_executor.ReferenceResolvingExecutor(
bottom_ex)
return federating_executor.FederatingExecutor({
placement_literals.SERVER:
bottom_ex,
placement_literals.CLIENTS: [bottom_ex] * num_clients,
None:
bottom_ex
})
def _create_federated_stack(num_clients, clients_per_thread, device_scheduler):
"""Constructs local federated stack."""
bottom_stacks = [
_create_bottom_stack(device=device_scheduler.next_client_device())
for _ in range(num_clients // clients_per_thread)
]
executor_dict = {
placement_literals.CLIENTS:
[bottom_stacks[k % len(bottom_stacks)] for k in range(num_clients)],
placement_literals.SERVER:
_create_bottom_stack(device=device_scheduler.server_device()),
None:
_create_bottom_stack(device=device_scheduler.server_device())
}
return _complete_stack(
federating_executor.FederatingExecutor(executor_dict))
def test_executor_call_unsupported_intrinsic(self):
dummy_intrinsic = intrinsic_defs.IntrinsicDef(
'DUMMY_INTRINSIC', 'dummy_intrinsic',
computation_types.AbstractType('T'))
type_signature = computation_types.TensorType(tf.int32)
comp = pb.Computation(
type=type_serialization.serialize_type(type_signature),
intrinsic=pb.Intrinsic(uri='dummy_intrinsic'))
loop = asyncio.get_event_loop()
factory = federated_composing_strategy.FederatedComposingStrategy.factory(
_create_bottom_stack(), [_create_worker_stack()])
executor = federating_executor.FederatingExecutor(
factory, _create_bottom_stack())
v1 = loop.run_until_complete(executor.create_value(comp))
with self.assertRaises(NotImplementedError):
loop.run_until_complete(executor.create_call(v1))
def test_executor_call_unsupported_intrinsic(self):
# `whimsy_intrinsic` definition is needed to allow successful lookup.
whimsy_intrinsic = intrinsic_defs.IntrinsicDef(
'WHIMSY_INTRINSIC', 'whimsy_intrinsic',
computation_types.AbstractType('T'))
type_signature = computation_types.TensorType(tf.int32)
comp = pb.Computation(
type=type_serialization.serialize_type(type_signature),
intrinsic=pb.Intrinsic(uri='whimsy_intrinsic'))
del whimsy_intrinsic
factory = federated_composing_strategy.FederatedComposingStrategy.factory(
_create_bottom_stack(), [_create_worker_stack()])
executor = federating_executor.FederatingExecutor(
factory, _create_bottom_stack())
v1 = asyncio.run(executor.create_value(comp))
with self.assertRaises(NotImplementedError):
asyncio.run(executor.create_call(v1))
