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_multiple_inputs(self):
ex = sizing_executor.SizingExecutor(eager_tf_executor.EagerTFExecutor())
int_type = computation_types.TensorType(tf.int32, 10)
float_type = computation_types.TensorType(tf.float64, 10)
@computations.tf_computation(float_type, int_type)
def add(x, y):
x = tf.cast(x, tf.int64)
y = tf.cast(y, tf.int64)
return x + y
async def _make():
v1 = await ex.create_value(add)
v2 = await ex.create_value([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], float_type)
v3 = await ex.create_value([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], int_type)
v4 = await ex.create_struct(
anonymous_tuple.AnonymousTuple([(None, v2), (None, v3)]))
v5 = await ex.create_call(v1, v4)
return await v5.compute()
asyncio.get_event_loop().run_until_complete(_make())
self.assertCountEqual(ex.broadcast_history,
[[10, tf.int32], [10, tf.float64]])
self.assertCountEqual(ex.aggregate_history, [[10, tf.int64]])
def test_nested_tuple(self):
ex = sizing_executor.SizingExecutor(eager_tf_executor.EagerTFExecutor())
a = computation_types.TensorType(tf.int32, [4])
b = computation_types.TensorType(tf.bool, [2])
c = computation_types.TensorType(tf.int64, [2, 3])
inner_type = computation_types.NamedTupleType([('a', a), ('b', b),
('c', c)])
outer_type = computation_types.NamedTupleType([('a', inner_type),
('b', inner_type)])
inner_type_val = collections.OrderedDict()
inner_type_val['a'] = [0, 1, 2, 3]
inner_type_val['b'] = [True, False]
inner_type_val['c'] = [[1, 2, 3], [4, 5, 6]]
outer_type_val = collections.OrderedDict()
outer_type_val['a'] = inner_type_val
outer_type_val['b'] = inner_type_val
async def _make():
v1 = await ex.create_value(outer_type_val, outer_type)
return await v1.compute()
asyncio.get_event_loop().run_until_complete(_make())
self.assertCountEqual(ex.broadcast_history,
[[4, tf.int32], [2, tf.bool], [6, tf.int64],
[4, tf.int32], [2, tf.bool], [6, tf.int64]])
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 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_empty_tuple(self):
ex = sizing_executor.SizingExecutor(eager_tf_executor.EagerTFExecutor())
tup = computation_types.NamedTupleType([])
empty_dict = collections.OrderedDict()
async def _make():
v1 = await ex.create_value(empty_dict, tup)
return await v1.compute()
asyncio.get_event_loop().run_until_complete(_make())
self.assertCountEqual(ex.broadcast_history, [])
def test_unnamed_tuple(self):
ex = sizing_executor.SizingExecutor(eager_tf_executor.EagerTFExecutor())
type_spec = computation_types.NamedTupleType([tf.int32, tf.int32])
value = anonymous_tuple.AnonymousTuple([(None, 0), (None, 1)])
async def _make():
v1 = await ex.create_value(value, type_spec)
return await v1.compute()
asyncio.get_event_loop().run_until_complete(_make())
self.assertCountEqual(ex.broadcast_history, [[1, tf.int32], [1, tf.int32]])
self.assertCountEqual(ex.aggregate_history, [[1, tf.int32], [1, tf.int32]])
def test_string(self):
ex = sizing_executor.SizingExecutor(eager_tf_executor.EagerTFExecutor())
tensor_type = computation_types.TensorType(tf.string, [4])
strings = ['hi', 'bye', 'tensor', 'federated']
total_string_length = sum([len(s) for s in strings])
async def _make():
v1 = await ex.create_value(strings, tensor_type)
return await v1.compute()
asyncio.get_event_loop().run_until_complete(_make())
self.assertCountEqual(ex.broadcast_history,
[[total_string_length, tf.string]])
self.assertCountEqual(ex.aggregate_history,
[[total_string_length, tf.string]])
def _create_sizing_stack(num_clients, clients_per_thread):
"""Constructs local stack with sizing wired in at each client."""
sizing_stacks = [
sizing_executor.SizingExecutor(_create_bottom_stack())
for _ in range(num_clients // clients_per_thread)
]
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(federated_executor.FederatedExecutor(executor_dict))
def test_unnamed_tuple(self):
ex = sizing_executor.SizingExecutor(
eager_tf_executor.EagerTFExecutor())
type_spec = computation_types.StructType([tf.int32, tf.int32])
value = structure.Struct([(None, 0), (None, 1)])
async def _make():
v1 = await ex.create_value(value, type_spec)
return await v1.compute()
asyncio.run(_make())
self.assertCountEqual(ex.broadcast_history,
[[1, tf.int32], [1, tf.int32]])
self.assertCountEqual(ex.aggregate_history,
[[1, tf.int32], [1, tf.int32]])
def test_ordered_dict(self):
a = computation_types.TensorType(tf.string, [4])
b = computation_types.TensorType(tf.int64, [2, 3])
tup = computation_types.NamedTupleType([('a', a), ('b', b)])
ex = sizing_executor.SizingExecutor(eager_tf_executor.EagerTFExecutor())
od = collections.OrderedDict()
od['a'] = ['some', 'arbitrary', 'string', 'here']
od['b'] = [[3, 4, 1], [6, 8, -5]]
total_string_length = sum([len(s) for s in od['a']])
async def _make():
v1 = await ex.create_value(od, tup)
return await v1.compute()
asyncio.get_event_loop().run_until_complete(_make())
self.assertCountEqual(ex.broadcast_history,
[[total_string_length, tf.string], [6, tf.int64]])
def test_different_input_output(self):
ex = sizing_executor.SizingExecutor(eager_tf_executor.EagerTFExecutor())
tensor_type = computation_types.TensorType(tf.int32, 10)
@computations.tf_computation(tensor_type)
def return_constant(x):
del x
return tf.constant(0, tf.int32)
async def _make():
v1 = await ex.create_value(return_constant)
v2 = await ex.create_value([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], tensor_type)
v3 = await ex.create_call(v1, v2)
return await v3.compute()
asyncio.get_event_loop().run_until_complete(_make())
self.assertCountEqual(ex.broadcast_history, [[10, tf.int32]])
self.assertCountEqual(ex.aggregate_history, [[1, tf.int32]])
def test_simple(self):
ex = sizing_executor.SizingExecutor(eager_tf_executor.EagerTFExecutor())
tensor_type = computation_types.TensorType(tf.int32, 10)
@computations.tf_computation(tensor_type)
def add_one(x):
return tf.add(x, 1)
async def _make():
v1 = await ex.create_value(add_one)
v2 = await ex.create_value(
tf.constant([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], tf.int32), tensor_type)
v3 = await ex.create_call(v1, v2)
v4 = await ex.create_struct(anonymous_tuple.AnonymousTuple([('foo', v3)]))
v5 = await ex.create_selection(v4, name='foo')
return await v5.compute()
asyncio.get_event_loop().run_until_complete(_make())
self.assertCountEqual(ex.broadcast_history, [[10, tf.int32]])
self.assertCountEqual(ex.aggregate_history, [[10, tf.int32]])
def _create_sizing_stack(num_clients: int, num_client_executors: int,
unplaced_ex_factory: UnplacedExecutorFactory):
"""Constructs local stack with sizing wired in at each client."""
sizing_stacks = []
for _ in range(num_client_executors):
sizing_ex = sizing_executor.SizingExecutor(
unplaced_ex_factory.create_executor(
placement=placement_literals.CLIENTS))
sizing_stacks.append(sizing_ex)
executor_dict = {
placement_literals.CLIENTS:
[sizing_stacks[k % len(sizing_stacks)] for k in range(num_clients)],
placement_literals.SERVER:
unplaced_ex_factory.create_executor(
placement=placement_literals.SERVER),
None:
unplaced_ex_factory.create_executor(
placement=placement_literals.SERVER),
}
return _wrap_executor_in_threading_stack(
federating_executor.FederatingExecutor(executor_dict)), sizing_stacks
