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))