From 8e63a53d1eab202dd2015b3f2492a03501ddecd1 Mon Sep 17 00:00:00 2001 From: yann dupis Date: Wed, 6 May 2020 19:40:35 -0700 Subject: [PATCH 01/31] introduce secure channel class --- .../impl/executors/federating_executor.py | 303 +++++++++++------- 1 file changed, 195 insertions(+), 108 deletions(-) diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor.py b/tensorflow_federated/python/core/impl/executors/federating_executor.py index 8a43585def..58a34fe5bf 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor.py @@ -492,35 +492,6 @@ def validate_executor_placements(cls, executor_placements): 'Unsupported cardinality for placement "{}": {}.'.format( pl, pl_cardinality)) - async def _trusted_aggregator_generate_keys(self): - - @computations.tf_computation() - def generate_keys(): - pk, sk = easy_box.gen_keypair() - return pk.raw, sk.raw - - fn_type = generate_keys.type_signature - fn = generate_keys._computation_proto - - aggregator_ex = self._get_child_executors( - placement_literals.AGGREGATORS, index=0) - - key_generator = await aggregator_ex.create_call(await - aggregator_ex.create_value( - fn, fn_type)) - - keys = await asyncio.gather(*[ - aggregator_ex.create_selection(key_generator, i) - for i in range(len(key_generator.type_signature)) - ]) - - pk_fed_val, sk_fed_val = await asyncio.gather( - *[self._place(k, placement_literals.AGGREGATORS) for k in keys]) - - pk_fed_val = await self.federated_broadcast(pk_fed_val) - - return pk_fed_val, sk_fed_val - async def _zip_val_key(self, vals, key, placement): if isinstance(vals, list): @@ -540,81 +511,6 @@ async def _zip_val_key(self, vals, key, placement): return vals_key_zipped.internal_representation - async def _encrypt_client_tensors(self, arg, pk_a): - - nb_clients = len(self._get_child_executors(placement_literals.CLIENTS)) - - if nb_clients == 1: - input_tensor_type = arg.type_signature.member - else: - input_tensor_type = arg.type_signature[0].member - pk_a_tensor_type = pk_a.type_signature.member - - @computations.tf_computation(input_tensor_type, pk_a_tensor_type) - def encrypt_tensor(plaintext, pk_a): - - pk_a = easy_box.PublicKey(pk_a) - pk_c, sk_c = easy_box.gen_keypair() - nonce = easy_box.gen_nonce() - ciphertext, mac = easy_box.seal_detached(plaintext, nonce, pk_a, sk_c) - - return ciphertext.raw, mac.raw, pk_c.raw, nonce.raw - - fn_type = encrypt_tensor.type_signature - fn = encrypt_tensor._computation_proto - if nb_clients == 1: - val_type = arg.type_signature - val = arg.internal_representation - else: - val_type = arg.type_signature[0] - val = arg.internal_representation[0] - - val_key_zipped = await self._zip_val_key(val, pk_a, - placement_literals.CLIENTS) - - fed_ex = self.federating_executor - - return await fed_ex._compute_intrinsic_federated_map( - FederatingExecutorValue( - anonymous_tuple.AnonymousTuple([(None, fn), - (None, val_key_zipped)]), - computation_types.NamedTupleType((fn_type, val_type)))) - - async def _decrypt_tensors_on_aggregator(self, val, clients_dtype): - - client_output_type_signature = val[0].type_signature[0] - aggregator_public_key_type_signature = val[0].type_signature[1] - - @computations.tf_computation(client_output_type_signature, - aggregator_public_key_type_signature) - def decrypt_tensor(client_outputs, sk_a): - - ciphertext = easy_box.Ciphertext(client_outputs[0]) - mac = easy_box.Mac(client_outputs[1]) - pk_c = easy_box.PublicKey(client_outputs[2]) - nonce = easy_box.Nonce(client_outputs[3]) - sk_a = easy_box.SecretKey(sk_a) - - plaintext_recovered = easy_box.open_detached(ciphertext, mac, nonce, pk_c, - sk_a, clients_dtype) - - return plaintext_recovered - - val_type = computation_types.FederatedType( - computation_types.TensorType(clients_dtype), - placement_literals.AGGREGATORS, - all_equal=False) - - fn_type = decrypt_tensor.type_signature - fn = decrypt_tensor._computation_proto - - fed_ex = self.federating_executor - - return await fed_ex._compute_intrinsic_federated_map( - FederatingExecutorValue( - anonymous_tuple.AnonymousTuple([(None, fn), (None, val)]), - computation_types.NamedTupleType((fn_type, val_type)))) - async def federated_value_at_server(self, arg): return await self._place(arg, placement_literals.SERVER) @@ -665,8 +561,16 @@ async def federated_reduce(self, arg): 'found {}.'.format(len(arg.internal_representation))) # Encrypt client tensors - pk_a, sk_a = await self._trusted_aggregator_generate_keys() - enc_clients_arg = await self._encrypt_client_tensors(arg, pk_a) + secure_channel = SecureChannels() + pk_a, sk_a = await secure_channel.setup(channel='clients-aggregators', + executor=self) + + # pk_a, sk_a = await self._trusted_aggregator_generate_keys() + # enc_clients_arg = await self._encrypt_client_tensors(arg, pk_a) + enc_clients_arg = await secure_channel.send(channel='clients-aggregators', + executor=self, + arg=arg, + pk_receiver=pk_a) val_type = enc_clients_arg.type_signature val = enc_clients_arg.internal_representation # Store original client dtype before encryption for @@ -690,8 +594,10 @@ async def federated_reduce(self, arg): for item in aggregands: item_key_zipped = await self._zip_val_key(item, sk_a, placement_literals.AGGREGATORS) - decrypted_tensor = await self._decrypt_tensors_on_aggregator( - item_key_zipped, orig_clients_dtype) + decrypted_tensor = await secure_channel.receive(channel='clients-aggregators', + executor=self, + val=item_key_zipped, + receiver_dtype=orig_clients_dtype) aggregands_decrypted.append(decrypted_tensor.internal_representation[0]) zero = await aggr.create_value( @@ -1222,3 +1128,184 @@ async def _compute_intrinsic_federated_collect(self, arg): @tracing.trace async def _compute_intrinsic_federated_secure_sum(self, arg): return await self.intrinsic_strategy.federated_secure_sum(arg) + + +class SecureChannels: + def __init__(self, channels_config=None): + + if channels_config is None: + channels_config = {'clients-aggregators': {'encrypted': True}, + 'server-aggregators': {'encryped': True}, + 'clients-server': {'encryped': False} + } + + self.channels_config = channels_config + + async def setup(self, channel, executor): + encrypted = self.channels_config[channel]['encrypted'] + + if not encrypted: + return + + receiver_pk, receiver_sk = await _receiver_generate_keys(channel, executor) + + return receiver_pk, receiver_sk + + async def send(self, channel, executor, arg, pk_receiver): + encrypted = self.channels_config[channel]['encrypted'] + + if not encrypted: + return arg + + return await _encrypt_sender_tensors(channel, executor, arg, pk_receiver) + + + async def receive(self, channel, executor, val, receiver_dtype): + encrypted = self.channels_config[channel]['encrypted'] + + if not encrypted: + return val + + return await _decrypt_tensors_on_receiver(channel, executor, val, receiver_dtype) + +async def _receiver_generate_keys(channel, executor): + # NOTE solve confusion between executor and receiver_executor + sender_placement, receiver_placement = _get_placements(channel) + + @computations.tf_computation() + def generate_keys(): + pk, sk = easy_box.gen_keypair() + return pk.raw, sk.raw + + fn_type = generate_keys.type_signature + fn = generate_keys._computation_proto + + receiver_ex = executor._get_child_executors( + receiver_placement, index=0) + + key_generator = await receiver_ex.create_call(await receiver_ex.create_value( + fn, fn_type)) + + keys = await asyncio.gather(*[ + receiver_ex.create_selection(key_generator, i) + for i in range(len(key_generator.type_signature)) + ]) + + pk_fed_val, sk_fed_val = await asyncio.gather( + *[executor._place(k, receiver_placement) for k in keys]) + + # NOTE: This broadcast should work from aggregator to clients + # but also from server to aggregator (so server can receive + # final updates) + pk_fed_val = await executor.federated_broadcast(pk_fed_val) + + return pk_fed_val, sk_fed_val + +def _get_placements(channel): + + sender, receiver = channel.split('-') + + placements_mapping = {'sever': placement_literals.SERVER, + 'aggregators': placement_literals.AGGREGATORS, + 'clients': placement_literals.CLIENTS} + + sender_placement = placements_mapping[sender] + receiver_placement = placements_mapping[receiver] + + return sender_placement, receiver_placement + +async def _zip_val_key(executor, vals, key, placement): + + if isinstance(vals, list): + val_type = computation_types.FederatedType( + vals[0].type_signature, placement, all_equal=False) + else: + val_type = computation_types.FederatedType( + vals.type_signature, placement, all_equal=False) + vals = [vals] + + vals_key = FederatingExecutorValue( + anonymous_tuple.AnonymousTuple([(None, vals), + (None, key.internal_representation)]), + computation_types.NamedTupleType((val_type, key.type_signature))) + + vals_key_zipped = await executor._zip(vals_key, placement, all_equal=False) + + return vals_key_zipped.internal_representation + +async def _encrypt_sender_tensors(channel, executor, arg, pk_rcv): + sender_placement, receiver_placement = _get_placements(channel) + + nb_senders = len(executor._get_child_executors(sender_placement)) + + if nb_senders == 1: + input_tensor_type = arg.type_signature.member + else: + input_tensor_type = arg.type_signature[0].member + pk_rcv_tensor_type = pk_rcv.type_signature.member + + @computations.tf_computation(input_tensor_type, pk_rcv_tensor_type) + def encrypt_tensor(plaintext, pk_rcv): + + pk_rcv = easy_box.PublicKey(pk_rcv) + pk_snd, sk_snd = easy_box.gen_keypair() + nonce = easy_box.gen_nonce() + ciphertext, mac = easy_box.seal_detached(plaintext, nonce, pk_rcv, sk_snd) + + return ciphertext.raw, mac.raw, pk_snd.raw, nonce.raw + + fn_type = encrypt_tensor.type_signature + fn = encrypt_tensor._computation_proto + if nb_senders == 1: + val_type = arg.type_signature + val = arg.internal_representation + else: + val_type = arg.type_signature[0] + val = arg.internal_representation[0] + + val_key_zipped = await _zip_val_key(executor, val, pk_rcv, + sender_placement) + + fed_ex = executor.federating_executor + + return await fed_ex._compute_intrinsic_federated_map( + FederatingExecutorValue( + anonymous_tuple.AnonymousTuple([(None, fn), + (None, val_key_zipped)]), + computation_types.NamedTupleType((fn_type, val_type)))) + +async def _decrypt_tensors_on_receiver(channel, executor, val, receiver_dtype): + sender_placement, receiver_placement = _get_placements(channel) + + client_output_type_signature = val[0].type_signature[0] + aggregator_public_key_type_signature = val[0].type_signature[1] + + @computations.tf_computation(client_output_type_signature, + aggregator_public_key_type_signature) + def decrypt_tensor(client_outputs, sk_rcv): + + ciphertext = easy_box.Ciphertext(client_outputs[0]) + mac = easy_box.Mac(client_outputs[1]) + pk_snd = easy_box.PublicKey(client_outputs[2]) + nonce = easy_box.Nonce(client_outputs[3]) + sk_rcv = easy_box.SecretKey(sk_rcv) + + plaintext_recovered = easy_box.open_detached(ciphertext, mac, nonce, pk_snd, + sk_rcv, receiver_dtype) + + return plaintext_recovered + + val_type = computation_types.FederatedType( + computation_types.TensorType(receiver_dtype), + receiver_placement, + all_equal=False) + + fn_type = decrypt_tensor.type_signature + fn = decrypt_tensor._computation_proto + + fed_ex = executor.federating_executor + + return await fed_ex._compute_intrinsic_federated_map( + FederatingExecutorValue( + anonymous_tuple.AnonymousTuple([(None, fn), (None, val)]), + computation_types.NamedTupleType((fn_type, val_type)))) From 779696f84f35d565808170bc75ca74fd8f7b731d Mon Sep 17 00:00:00 2001 From: yann dupis Date: Wed, 6 May 2020 21:25:21 -0700 Subject: [PATCH 02/31] add cooments --- .../python/core/impl/executors/federating_executor.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor.py b/tensorflow_federated/python/core/impl/executors/federating_executor.py index 58a34fe5bf..6510dbfe83 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor.py @@ -561,6 +561,7 @@ async def federated_reduce(self, arg): 'found {}.'.format(len(arg.internal_representation))) # Encrypt client tensors + # NOTE it will be supplied to the intrinsic strategy secure_channel = SecureChannels() pk_a, sk_a = await secure_channel.setup(channel='clients-aggregators', executor=self) @@ -1129,7 +1130,7 @@ async def _compute_intrinsic_federated_collect(self, arg): async def _compute_intrinsic_federated_secure_sum(self, arg): return await self.intrinsic_strategy.federated_secure_sum(arg) - +# NOTE Should it be called Channel instead? class SecureChannels: def __init__(self, channels_config=None): From 758b5f1141791fa5d4142109172cffdea8b8bfaa Mon Sep 17 00:00:00 2001 From: yann dupis Date: Thu, 7 May 2020 15:31:09 -0700 Subject: [PATCH 03/31] mv _zip_val_key to secure channel --- .../impl/executors/federating_executor.py | 41 ++++++------------- 1 file changed, 12 insertions(+), 29 deletions(-) diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor.py b/tensorflow_federated/python/core/impl/executors/federating_executor.py index 6510dbfe83..7f210d9019 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor.py @@ -492,25 +492,6 @@ def validate_executor_placements(cls, executor_placements): 'Unsupported cardinality for placement "{}": {}.'.format( pl, pl_cardinality)) - async def _zip_val_key(self, vals, key, placement): - - if isinstance(vals, list): - val_type = computation_types.FederatedType( - vals[0].type_signature, placement, all_equal=False) - else: - val_type = computation_types.FederatedType( - vals.type_signature, placement, all_equal=False) - vals = [vals] - - vals_key = FederatingExecutorValue( - anonymous_tuple.AnonymousTuple([(None, vals), - (None, key.internal_representation)]), - computation_types.NamedTupleType((val_type, key.type_signature))) - - vals_key_zipped = await self._zip(vals_key, placement, all_equal=False) - - return vals_key_zipped.internal_representation - async def federated_value_at_server(self, arg): return await self._place(arg, placement_literals.SERVER) @@ -593,11 +574,10 @@ async def federated_reduce(self, arg): # In the future should be decrypted by the Trusted aggregator aggregands_decrypted = [] for item in aggregands: - item_key_zipped = await self._zip_val_key(item, sk_a, - placement_literals.AGGREGATORS) decrypted_tensor = await secure_channel.receive(channel='clients-aggregators', executor=self, - val=item_key_zipped, + val=item, + receiver_sk=sk_a, receiver_dtype=orig_clients_dtype) aggregands_decrypted.append(decrypted_tensor.internal_representation[0]) @@ -1130,6 +1110,7 @@ async def _compute_intrinsic_federated_collect(self, arg): async def _compute_intrinsic_federated_secure_sum(self, arg): return await self.intrinsic_strategy.federated_secure_sum(arg) + # NOTE Should it be called Channel instead? class SecureChannels: def __init__(self, channels_config=None): @@ -1161,13 +1142,13 @@ async def send(self, channel, executor, arg, pk_receiver): return await _encrypt_sender_tensors(channel, executor, arg, pk_receiver) - async def receive(self, channel, executor, val, receiver_dtype): + async def receive(self, channel, executor, val, receiver_sk, receiver_dtype): encrypted = self.channels_config[channel]['encrypted'] if not encrypted: return val - return await _decrypt_tensors_on_receiver(channel, executor, val, receiver_dtype) + return await _decrypt_tensors_on_receiver(channel, executor, val, receiver_sk, receiver_dtype) async def _receiver_generate_keys(channel, executor): # NOTE solve confusion between executor and receiver_executor @@ -1275,14 +1256,16 @@ def encrypt_tensor(plaintext, pk_rcv): (None, val_key_zipped)]), computation_types.NamedTupleType((fn_type, val_type)))) -async def _decrypt_tensors_on_receiver(channel, executor, val, receiver_dtype): +async def _decrypt_tensors_on_receiver(channel, executor, val, receiver_sk, receiver_dtype): sender_placement, receiver_placement = _get_placements(channel) - client_output_type_signature = val[0].type_signature[0] - aggregator_public_key_type_signature = val[0].type_signature[1] + val = await _zip_val_key(executor, val, receiver_sk, receiver_placement) + + sender_output_type_signature = val[0].type_signature[0] + receiver_public_key_type_signature = val[0].type_signature[1] - @computations.tf_computation(client_output_type_signature, - aggregator_public_key_type_signature) + @computations.tf_computation(sender_output_type_signature, + receiver_public_key_type_signature) def decrypt_tensor(client_outputs, sk_rcv): ciphertext = easy_box.Ciphertext(client_outputs[0]) From f256b35f669294d613b8786bcbfb150503a28dc7 Mon Sep 17 00:00:00 2001 From: yann dupis Date: Thu, 7 May 2020 16:14:35 -0700 Subject: [PATCH 04/31] rm config_channels from SecureChannels --- .../impl/executors/federating_executor.py | 255 ++++++++---------- 1 file changed, 115 insertions(+), 140 deletions(-) diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor.py b/tensorflow_federated/python/core/impl/executors/federating_executor.py index 7f210d9019..9a095b7894 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor.py @@ -543,14 +543,13 @@ async def federated_reduce(self, arg): # Encrypt client tensors # NOTE it will be supplied to the intrinsic strategy - secure_channel = SecureChannels() - pk_a, sk_a = await secure_channel.setup(channel='clients-aggregators', - executor=self) - - # pk_a, sk_a = await self._trusted_aggregator_generate_keys() - # enc_clients_arg = await self._encrypt_client_tensors(arg, pk_a) - enc_clients_arg = await secure_channel.send(channel='clients-aggregators', - executor=self, + agg_cls_channel = Channel(sender_placement=placement_literals.CLIENTS, + receiver_placement=placement_literals.AGGREGATORS, + is_encrypted=True) + + pk_a, sk_a = await agg_cls_channel.setup(executor=self) + + enc_clients_arg = await agg_cls_channel.send(executor=self, arg=arg, pk_receiver=pk_a) val_type = enc_clients_arg.type_signature @@ -574,8 +573,7 @@ async def federated_reduce(self, arg): # In the future should be decrypted by the Trusted aggregator aggregands_decrypted = [] for item in aggregands: - decrypted_tensor = await secure_channel.receive(channel='clients-aggregators', - executor=self, + decrypted_tensor = await agg_cls_channel.receive(executor=self, val=item, receiver_sk=sk_a, receiver_dtype=orig_clients_dtype) @@ -1111,185 +1109,162 @@ async def _compute_intrinsic_federated_secure_sum(self, arg): return await self.intrinsic_strategy.federated_secure_sum(arg) -# NOTE Should it be called Channel instead? -class SecureChannels: - def __init__(self, channels_config=None): - - if channels_config is None: - channels_config = {'clients-aggregators': {'encrypted': True}, - 'server-aggregators': {'encryped': True}, - 'clients-server': {'encryped': False} - } +class Channel: + def __init__(self, sender_placement, receiver_placement, is_encrypted): - self.channels_config = channels_config + self.sender_placement = sender_placement + self.receiver_placement = receiver_placement + self.is_encrypted = is_encrypted - async def setup(self, channel, executor): - encrypted = self.channels_config[channel]['encrypted'] + async def setup(self, executor): - if not encrypted: + if not self.is_encrypted: return - receiver_pk, receiver_sk = await _receiver_generate_keys(channel, executor) + receiver_pk, receiver_sk = await self._receiver_generate_keys(executor) return receiver_pk, receiver_sk - async def send(self, channel, executor, arg, pk_receiver): - encrypted = self.channels_config[channel]['encrypted'] + async def send(self, executor, arg, pk_receiver): - if not encrypted: + if not self.is_encrypted: return arg - return await _encrypt_sender_tensors(channel, executor, arg, pk_receiver) + return await self._encrypt_sender_tensors(executor, arg, pk_receiver) - async def receive(self, channel, executor, val, receiver_sk, receiver_dtype): - encrypted = self.channels_config[channel]['encrypted'] + async def receive(self, executor, val, receiver_sk, receiver_dtype): - if not encrypted: + if not self.is_encrypted: return val - return await _decrypt_tensors_on_receiver(channel, executor, val, receiver_sk, receiver_dtype) - -async def _receiver_generate_keys(channel, executor): - # NOTE solve confusion between executor and receiver_executor - sender_placement, receiver_placement = _get_placements(channel) - - @computations.tf_computation() - def generate_keys(): - pk, sk = easy_box.gen_keypair() - return pk.raw, sk.raw + return await self._decrypt_tensors_on_receiver(executor, val, receiver_sk, receiver_dtype) - fn_type = generate_keys.type_signature - fn = generate_keys._computation_proto - receiver_ex = executor._get_child_executors( - receiver_placement, index=0) + async def _receiver_generate_keys(self, executor): + # NOTE solve confusion between executor and receiver_executor - key_generator = await receiver_ex.create_call(await receiver_ex.create_value( - fn, fn_type)) + @computations.tf_computation() + def generate_keys(): + pk, sk = easy_box.gen_keypair() + return pk.raw, sk.raw - keys = await asyncio.gather(*[ - receiver_ex.create_selection(key_generator, i) - for i in range(len(key_generator.type_signature)) - ]) + fn_type = generate_keys.type_signature + fn = generate_keys._computation_proto - pk_fed_val, sk_fed_val = await asyncio.gather( - *[executor._place(k, receiver_placement) for k in keys]) + receiver_ex = executor._get_child_executors( + self.receiver_placement, index=0) - # NOTE: This broadcast should work from aggregator to clients - # but also from server to aggregator (so server can receive - # final updates) - pk_fed_val = await executor.federated_broadcast(pk_fed_val) + key_generator = await receiver_ex.create_call(await receiver_ex.create_value( + fn, fn_type)) - return pk_fed_val, sk_fed_val - -def _get_placements(channel): - - sender, receiver = channel.split('-') + keys = await asyncio.gather(*[ + receiver_ex.create_selection(key_generator, i) + for i in range(len(key_generator.type_signature)) + ]) - placements_mapping = {'sever': placement_literals.SERVER, - 'aggregators': placement_literals.AGGREGATORS, - 'clients': placement_literals.CLIENTS} + pk_fed_val, sk_fed_val = await asyncio.gather( + *[executor._place(k, self.receiver_placement) for k in keys]) - sender_placement = placements_mapping[sender] - receiver_placement = placements_mapping[receiver] + # NOTE: This broadcast should work from aggregator to clients + # but also from server to aggregator (so server can receive + # final updates) + pk_fed_val = await executor.federated_broadcast(pk_fed_val) - return sender_placement, receiver_placement + return pk_fed_val, sk_fed_val -async def _zip_val_key(executor, vals, key, placement): + async def _zip_val_key(self, executor, vals, key, placement): - if isinstance(vals, list): - val_type = computation_types.FederatedType( - vals[0].type_signature, placement, all_equal=False) - else: - val_type = computation_types.FederatedType( - vals.type_signature, placement, all_equal=False) - vals = [vals] + if isinstance(vals, list): + val_type = computation_types.FederatedType( + vals[0].type_signature, placement, all_equal=False) + else: + val_type = computation_types.FederatedType( + vals.type_signature, placement, all_equal=False) + vals = [vals] - vals_key = FederatingExecutorValue( - anonymous_tuple.AnonymousTuple([(None, vals), - (None, key.internal_representation)]), - computation_types.NamedTupleType((val_type, key.type_signature))) + vals_key = FederatingExecutorValue( + anonymous_tuple.AnonymousTuple([(None, vals), + (None, key.internal_representation)]), + computation_types.NamedTupleType((val_type, key.type_signature))) - vals_key_zipped = await executor._zip(vals_key, placement, all_equal=False) + vals_key_zipped = await executor._zip(vals_key, placement, all_equal=False) - return vals_key_zipped.internal_representation + return vals_key_zipped.internal_representation -async def _encrypt_sender_tensors(channel, executor, arg, pk_rcv): - sender_placement, receiver_placement = _get_placements(channel) + async def _encrypt_sender_tensors(self, executor, arg, pk_rcv): - nb_senders = len(executor._get_child_executors(sender_placement)) + nb_senders = len(executor._get_child_executors(self.sender_placement)) - if nb_senders == 1: - input_tensor_type = arg.type_signature.member - else: - input_tensor_type = arg.type_signature[0].member - pk_rcv_tensor_type = pk_rcv.type_signature.member + if nb_senders == 1: + input_tensor_type = arg.type_signature.member + else: + input_tensor_type = arg.type_signature[0].member + pk_rcv_tensor_type = pk_rcv.type_signature.member - @computations.tf_computation(input_tensor_type, pk_rcv_tensor_type) - def encrypt_tensor(plaintext, pk_rcv): + @computations.tf_computation(input_tensor_type, pk_rcv_tensor_type) + def encrypt_tensor(plaintext, pk_rcv): - pk_rcv = easy_box.PublicKey(pk_rcv) - pk_snd, sk_snd = easy_box.gen_keypair() - nonce = easy_box.gen_nonce() - ciphertext, mac = easy_box.seal_detached(plaintext, nonce, pk_rcv, sk_snd) + pk_rcv = easy_box.PublicKey(pk_rcv) + pk_snd, sk_snd = easy_box.gen_keypair() + nonce = easy_box.gen_nonce() + ciphertext, mac = easy_box.seal_detached(plaintext, nonce, pk_rcv, sk_snd) - return ciphertext.raw, mac.raw, pk_snd.raw, nonce.raw + return ciphertext.raw, mac.raw, pk_snd.raw, nonce.raw - fn_type = encrypt_tensor.type_signature - fn = encrypt_tensor._computation_proto - if nb_senders == 1: - val_type = arg.type_signature - val = arg.internal_representation - else: - val_type = arg.type_signature[0] - val = arg.internal_representation[0] + fn_type = encrypt_tensor.type_signature + fn = encrypt_tensor._computation_proto + if nb_senders == 1: + val_type = arg.type_signature + val = arg.internal_representation + else: + val_type = arg.type_signature[0] + val = arg.internal_representation[0] - val_key_zipped = await _zip_val_key(executor, val, pk_rcv, - sender_placement) + val_key_zipped = await self._zip_val_key(executor, val, pk_rcv, + self.sender_placement) - fed_ex = executor.federating_executor + fed_ex = executor.federating_executor - return await fed_ex._compute_intrinsic_federated_map( - FederatingExecutorValue( - anonymous_tuple.AnonymousTuple([(None, fn), - (None, val_key_zipped)]), - computation_types.NamedTupleType((fn_type, val_type)))) + return await fed_ex._compute_intrinsic_federated_map( + FederatingExecutorValue( + anonymous_tuple.AnonymousTuple([(None, fn), + (None, val_key_zipped)]), + computation_types.NamedTupleType((fn_type, val_type)))) -async def _decrypt_tensors_on_receiver(channel, executor, val, receiver_sk, receiver_dtype): - sender_placement, receiver_placement = _get_placements(channel) + async def _decrypt_tensors_on_receiver(self, executor, val, receiver_sk, receiver_dtype): - val = await _zip_val_key(executor, val, receiver_sk, receiver_placement) + val = await self._zip_val_key(executor, val, receiver_sk, self.receiver_placement) - sender_output_type_signature = val[0].type_signature[0] - receiver_public_key_type_signature = val[0].type_signature[1] + sender_output_type_signature = val[0].type_signature[0] + receiver_public_key_type_signature = val[0].type_signature[1] - @computations.tf_computation(sender_output_type_signature, - receiver_public_key_type_signature) - def decrypt_tensor(client_outputs, sk_rcv): + @computations.tf_computation(sender_output_type_signature, + receiver_public_key_type_signature) + def decrypt_tensor(client_outputs, sk_rcv): - ciphertext = easy_box.Ciphertext(client_outputs[0]) - mac = easy_box.Mac(client_outputs[1]) - pk_snd = easy_box.PublicKey(client_outputs[2]) - nonce = easy_box.Nonce(client_outputs[3]) - sk_rcv = easy_box.SecretKey(sk_rcv) + ciphertext = easy_box.Ciphertext(client_outputs[0]) + mac = easy_box.Mac(client_outputs[1]) + pk_snd = easy_box.PublicKey(client_outputs[2]) + nonce = easy_box.Nonce(client_outputs[3]) + sk_rcv = easy_box.SecretKey(sk_rcv) - plaintext_recovered = easy_box.open_detached(ciphertext, mac, nonce, pk_snd, - sk_rcv, receiver_dtype) + plaintext_recovered = easy_box.open_detached(ciphertext, mac, nonce, pk_snd, + sk_rcv, receiver_dtype) - return plaintext_recovered + return plaintext_recovered - val_type = computation_types.FederatedType( - computation_types.TensorType(receiver_dtype), - receiver_placement, - all_equal=False) + val_type = computation_types.FederatedType( + computation_types.TensorType(receiver_dtype), + self.receiver_placement, + all_equal=False) - fn_type = decrypt_tensor.type_signature - fn = decrypt_tensor._computation_proto + fn_type = decrypt_tensor.type_signature + fn = decrypt_tensor._computation_proto - fed_ex = executor.federating_executor + fed_ex = executor.federating_executor - return await fed_ex._compute_intrinsic_federated_map( - FederatingExecutorValue( - anonymous_tuple.AnonymousTuple([(None, fn), (None, val)]), - computation_types.NamedTupleType((fn_type, val_type)))) + return await fed_ex._compute_intrinsic_federated_map( + FederatingExecutorValue( + anonymous_tuple.AnonymousTuple([(None, fn), (None, val)]), + computation_types.NamedTupleType((fn_type, val_type)))) From 0aebe448efa911671f3ed7bfce6e96c1a968a9c9 Mon Sep 17 00:00:00 2001 From: yann dupis Date: Thu, 7 May 2020 17:38:12 -0700 Subject: [PATCH 05/31] mv channel receive to -move --- .../impl/executors/federating_executor.py | 125 ++++++++++-------- 1 file changed, 68 insertions(+), 57 deletions(-) diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor.py b/tensorflow_federated/python/core/impl/executors/federating_executor.py index 9a095b7894..ec429196ee 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor.py @@ -492,6 +492,19 @@ def validate_executor_placements(cls, executor_placements): 'Unsupported cardinality for placement "{}": {}.'.format( pl, pl_cardinality)) + @classmethod + async def _move(cls, value, value_type, target_executor, parent_executor, + sk_a, clients_dtype, channel): + val = await target_executor.create_value(await value.compute(), value_type) + + received_val = await channel.receive( + parent_executor=parent_executor, + val=val, + receiver_sk=sk_a, + sender_dtype=clients_dtype) + + return received_val.internal_representation[0] + async def federated_value_at_server(self, arg): return await self._place(arg, placement_literals.SERVER) @@ -543,15 +556,15 @@ async def federated_reduce(self, arg): # Encrypt client tensors # NOTE it will be supplied to the intrinsic strategy - agg_cls_channel = Channel(sender_placement=placement_literals.CLIENTS, - receiver_placement=placement_literals.AGGREGATORS, - is_encrypted=True) + agg_cls_channel = Channel( + sender_placement=placement_literals.CLIENTS, + receiver_placement=placement_literals.AGGREGATORS, + is_encrypted=True) - pk_a, sk_a = await agg_cls_channel.setup(executor=self) + pk_a, sk_a = await agg_cls_channel.setup(parent_executor=self) - enc_clients_arg = await agg_cls_channel.send(executor=self, - arg=arg, - pk_receiver=pk_a) + enc_clients_arg = await agg_cls_channel.send( + parent_executor=self, val=arg, pk_receiver=pk_a) val_type = enc_clients_arg.type_signature val = enc_clients_arg.internal_representation # Store original client dtype before encryption for @@ -566,18 +579,10 @@ async def federated_reduce(self, arg): py_typecheck.check_type(val, list) aggr = self._get_child_executors(placement_literals.AGGREGATORS, index=0) - aggregands = await asyncio.gather( - *[self._move(v, item_type, aggr) for v in val]) - - # Decrypt tensors moved to the server before applying reduce - # In the future should be decrypted by the Trusted aggregator - aggregands_decrypted = [] - for item in aggregands: - decrypted_tensor = await agg_cls_channel.receive(executor=self, - val=item, - receiver_sk=sk_a, - receiver_dtype=orig_clients_dtype) - aggregands_decrypted.append(decrypted_tensor.internal_representation[0]) + aggregands = await asyncio.gather(*[ + self._move(v, item_type, aggr, self, sk_a, orig_clients_dtype, + agg_cls_channel) for v in val + ]) zero = await aggr.create_value( await (await @@ -586,12 +591,12 @@ async def federated_reduce(self, arg): zero_type) op = await aggr.create_value(arg.internal_representation[2], op_type) - for item in aggregands_decrypted: + for item in aggregands: type_utils.check_equivalent_types( op_type, type_factory.reduction_op(zero_type, item.type_signature)) result = zero - for item in aggregands_decrypted: + for item in aggregands: result = await aggr.create_call( op, await aggr.create_tuple( anonymous_tuple.AnonymousTuple([(None, result), (None, item)]))) @@ -1110,38 +1115,39 @@ async def _compute_intrinsic_federated_secure_sum(self, arg): class Channel: + def __init__(self, sender_placement, receiver_placement, is_encrypted): self.sender_placement = sender_placement self.receiver_placement = receiver_placement self.is_encrypted = is_encrypted - async def setup(self, executor): + async def setup(self, parent_executor): if not self.is_encrypted: return - receiver_pk, receiver_sk = await self._receiver_generate_keys(executor) + receiver_pk, receiver_sk = await self._receiver_generate_keys( + parent_executor) return receiver_pk, receiver_sk - async def send(self, executor, arg, pk_receiver): + async def send(self, parent_executor, val, pk_receiver): if not self.is_encrypted: - return arg - - return await self._encrypt_sender_tensors(executor, arg, pk_receiver) + return val + return await self._encrypt_sender_tensors(parent_executor, val, pk_receiver) - async def receive(self, executor, val, receiver_sk, receiver_dtype): + async def receive(self, parent_executor, val, receiver_sk, sender_dtype): if not self.is_encrypted: return val - return await self._decrypt_tensors_on_receiver(executor, val, receiver_sk, receiver_dtype) + return await self._decrypt_tensors_on_receiver(parent_executor, val, + receiver_sk, sender_dtype) - - async def _receiver_generate_keys(self, executor): + async def _receiver_generate_keys(self, parent_executor): # NOTE solve confusion between executor and receiver_executor @computations.tf_computation() @@ -1152,10 +1158,11 @@ def generate_keys(): fn_type = generate_keys.type_signature fn = generate_keys._computation_proto - receiver_ex = executor._get_child_executors( + receiver_ex = parent_executor._get_child_executors( self.receiver_placement, index=0) - key_generator = await receiver_ex.create_call(await receiver_ex.create_value( + key_generator = await receiver_ex.create_call(await + receiver_ex.create_value( fn, fn_type)) keys = await asyncio.gather(*[ @@ -1164,16 +1171,16 @@ def generate_keys(): ]) pk_fed_val, sk_fed_val = await asyncio.gather( - *[executor._place(k, self.receiver_placement) for k in keys]) + *[parent_executor._place(k, self.receiver_placement) for k in keys]) # NOTE: This broadcast should work from aggregator to clients # but also from server to aggregator (so server can receive # final updates) - pk_fed_val = await executor.federated_broadcast(pk_fed_val) + pk_fed_val = await parent_executor.federated_broadcast(pk_fed_val) return pk_fed_val, sk_fed_val - async def _zip_val_key(self, executor, vals, key, placement): + async def _zip_val_key(self, parent_executor, vals, key, placement): if isinstance(vals, list): val_type = computation_types.FederatedType( @@ -1188,18 +1195,20 @@ async def _zip_val_key(self, executor, vals, key, placement): (None, key.internal_representation)]), computation_types.NamedTupleType((val_type, key.type_signature))) - vals_key_zipped = await executor._zip(vals_key, placement, all_equal=False) + vals_key_zipped = await parent_executor._zip( + vals_key, placement, all_equal=False) return vals_key_zipped.internal_representation - async def _encrypt_sender_tensors(self, executor, arg, pk_rcv): + async def _encrypt_sender_tensors(self, parent_executor, val, pk_rcv): - nb_senders = len(executor._get_child_executors(self.sender_placement)) + nb_senders = len( + parent_executor._get_child_executors(self.sender_placement)) if nb_senders == 1: - input_tensor_type = arg.type_signature.member + input_tensor_type = val.type_signature.member else: - input_tensor_type = arg.type_signature[0].member + input_tensor_type = val.type_signature[0].member pk_rcv_tensor_type = pk_rcv.type_signature.member @computations.tf_computation(input_tensor_type, pk_rcv_tensor_type) @@ -1215,16 +1224,16 @@ def encrypt_tensor(plaintext, pk_rcv): fn_type = encrypt_tensor.type_signature fn = encrypt_tensor._computation_proto if nb_senders == 1: - val_type = arg.type_signature - val = arg.internal_representation + val_type = val.type_signature + val = val.internal_representation else: - val_type = arg.type_signature[0] - val = arg.internal_representation[0] - - val_key_zipped = await self._zip_val_key(executor, val, pk_rcv, - self.sender_placement) + val_type = val.type_signature[0] + val = val.internal_representation[0] - fed_ex = executor.federating_executor + val_key_zipped = await self._zip_val_key(parent_executor, val, pk_rcv, + self.sender_placement) + # NOTE probably won't always be fed_ex in future design + fed_ex = parent_executor.federating_executor return await fed_ex._compute_intrinsic_federated_map( FederatingExecutorValue( @@ -1232,15 +1241,17 @@ def encrypt_tensor(plaintext, pk_rcv): (None, val_key_zipped)]), computation_types.NamedTupleType((fn_type, val_type)))) - async def _decrypt_tensors_on_receiver(self, executor, val, receiver_sk, receiver_dtype): + async def _decrypt_tensors_on_receiver(self, parent_executor, val, + receiver_sk, sender_dtype): - val = await self._zip_val_key(executor, val, receiver_sk, self.receiver_placement) + val = await self._zip_val_key(parent_executor, val, receiver_sk, + self.receiver_placement) sender_output_type_signature = val[0].type_signature[0] receiver_public_key_type_signature = val[0].type_signature[1] @computations.tf_computation(sender_output_type_signature, - receiver_public_key_type_signature) + receiver_public_key_type_signature) def decrypt_tensor(client_outputs, sk_rcv): ciphertext = easy_box.Ciphertext(client_outputs[0]) @@ -1249,20 +1260,20 @@ def decrypt_tensor(client_outputs, sk_rcv): nonce = easy_box.Nonce(client_outputs[3]) sk_rcv = easy_box.SecretKey(sk_rcv) - plaintext_recovered = easy_box.open_detached(ciphertext, mac, nonce, pk_snd, - sk_rcv, receiver_dtype) + plaintext_recovered = easy_box.open_detached(ciphertext, mac, nonce, + pk_snd, sk_rcv, sender_dtype) return plaintext_recovered val_type = computation_types.FederatedType( - computation_types.TensorType(receiver_dtype), + computation_types.TensorType(sender_dtype), self.receiver_placement, all_equal=False) fn_type = decrypt_tensor.type_signature fn = decrypt_tensor._computation_proto - - fed_ex = executor.federating_executor + # NOTE probably won't always be fed_ex in future design + fed_ex = parent_executor.federating_executor return await fed_ex._compute_intrinsic_federated_map( FederatingExecutorValue( From b2c8218a975bb461552733c7218b9028f41aceba Mon Sep 17 00:00:00 2001 From: yann dupis Date: Thu, 7 May 2020 19:40:28 -0700 Subject: [PATCH 06/31] mv send and reeceeive to -move --- .../impl/executors/federating_executor.py | 59 +++++++++--------- .../executors/federating_executor_test.py | 61 +++++++++---------- 2 files changed, 61 insertions(+), 59 deletions(-) diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor.py b/tensorflow_federated/python/core/impl/executors/federating_executor.py index ec429196ee..4c77f781a2 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor.py @@ -493,17 +493,37 @@ def validate_executor_placements(cls, executor_placements): pl, pl_cardinality)) @classmethod - async def _move(cls, value, value_type, target_executor, parent_executor, - sk_a, clients_dtype, channel): - val = await target_executor.create_value(await value.compute(), value_type) + async def _move(cls, arg, target_executor, parent_executor, channel): - received_val = await channel.receive( - parent_executor=parent_executor, - val=val, - receiver_sk=sk_a, - sender_dtype=clients_dtype) + pk_a, sk_a = await channel.setup(parent_executor=parent_executor) - return received_val.internal_representation[0] + enc_clients_arg = await channel.send( + parent_executor=parent_executor, val=arg, pk_receiver=pk_a) + + val_type = enc_clients_arg.type_signature + val = enc_clients_arg.internal_representation + orig_clients_dtype = arg.type_signature[0].member.dtype + + py_typecheck.check_type(val, list) + + py_typecheck.check_type(val_type, computation_types.FederatedType) + item_type = val_type.member + + val = await asyncio.gather(*[ + target_executor.create_value(await v.compute(), item_type) for v in val + ]) + + received_vals = await asyncio.gather(*[ + channel.receive( + parent_executor=parent_executor, + val=v, + receiver_sk=sk_a, + sender_dtype=orig_clients_dtype) for v in val + ]) + # NOTE should check if we can remove this step + received_vals = [v.internal_representation[0] for v in received_vals] + + return received_vals async def federated_value_at_server(self, arg): return await self._place(arg, placement_literals.SERVER) @@ -554,35 +574,18 @@ async def federated_reduce(self, arg): 'Expected 3 elements in the `federated_reduce()` argument tuple, ' 'found {}.'.format(len(arg.internal_representation))) - # Encrypt client tensors # NOTE it will be supplied to the intrinsic strategy agg_cls_channel = Channel( sender_placement=placement_literals.CLIENTS, receiver_placement=placement_literals.AGGREGATORS, is_encrypted=True) - pk_a, sk_a = await agg_cls_channel.setup(parent_executor=self) - - enc_clients_arg = await agg_cls_channel.send( - parent_executor=self, val=arg, pk_receiver=pk_a) - val_type = enc_clients_arg.type_signature - val = enc_clients_arg.internal_representation - # Store original client dtype before encryption for - # future decryption - orig_clients_dtype = arg.type_signature[0].member.dtype - - py_typecheck.check_type(val_type, computation_types.FederatedType) - item_type = val_type.member zero_type = arg.type_signature[1] op_type = arg.type_signature[2] - py_typecheck.check_type(val, list) aggr = self._get_child_executors(placement_literals.AGGREGATORS, index=0) - aggregands = await asyncio.gather(*[ - self._move(v, item_type, aggr, self, sk_a, orig_clients_dtype, - agg_cls_channel) for v in val - ]) + aggregands = await self._move(arg, aggr, self, agg_cls_channel) zero = await aggr.create_value( await (await @@ -1175,7 +1178,7 @@ def generate_keys(): # NOTE: This broadcast should work from aggregator to clients # but also from server to aggregator (so server can receive - # final updates) + # final updates). May have to use _place instead pk_fed_val = await parent_executor.federated_broadcast(pk_fed_val) return pk_fed_val, sk_fed_val diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor_test.py b/tensorflow_federated/python/core/impl/executors/federating_executor_test.py index 97c50c6797..c1f5fd5436 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor_test.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor_test.py @@ -1272,50 +1272,49 @@ def __init__(self, parent_executor): create_test_executor(intrinsic_strategy_fn=MockIntrinsicStrategy) -class EncryptionTest(parameterized.TestCase): +# class EncryptionTest(parameterized.TestCase): - def test_generate_aggregator_keys(self): - strategy = federating_executor.TrustedAggregatorIntrinsicStrategy - loop, ex = _make_test_runtime(intrinsic_strategy_fn=strategy) - generate_keys = ex.intrinsic_strategy._trusted_aggregator_generate_keys() - pk, sk = loop.run_until_complete(generate_keys) - - self.assertEqual(str(pk.type_signature), 'uint8[32]@CLIENTS') - self.assertEqual(str(sk.type_signature), 'uint8[32]@AGGREGATORS') +# def test_generate_aggregator_keys(self): +# strategy = federating_executor.TrustedAggregatorIntrinsicStrategy +# loop, ex = _make_test_runtime(intrinsic_strategy_fn=strategy) +# generate_keys = ex.intrinsic_strategy._trusted_aggregator_generate_keys() +# pk, sk = loop.run_until_complete(generate_keys) - def test_encryption_decryption(self): +# self.assertEqual(str(pk.type_signature), 'uint8[32]@CLIENTS') +# self.assertEqual(str(sk.type_signature), 'uint8[32]@AGGREGATORS') - strategy = federating_executor.TrustedAggregatorIntrinsicStrategy - loop, ex = _make_test_runtime(intrinsic_strategy_fn=strategy) - strat_ex = ex.intrinsic_strategy +# def test_encryption_decryption(self): - pk_a, sk_a = loop.run_until_complete( - strat_ex._trusted_aggregator_generate_keys()) +# strategy = federating_executor.TrustedAggregatorIntrinsicStrategy +# loop, ex = _make_test_runtime(intrinsic_strategy_fn=strategy) +# strat_ex = ex.intrinsic_strategy - val = loop.run_until_complete( - ex.create_value([2.0], type_factory.at_clients(tf.float32))) +# pk_a, sk_a = loop.run_until_complete( +# strat_ex._trusted_aggregator_generate_keys()) - val_enc = loop.run_until_complete( - strat_ex._encrypt_client_tensors(val, pk_a)) +# val = loop.run_until_complete( +# ex.create_value([2.0], type_factory.at_clients(tf.float32))) - aggr = strat_ex._get_child_executors( - placement_literals.AGGREGATORS, index=0) +# val_enc = loop.run_until_complete( +# strat_ex._encrypt_client_tensors(val, pk_a)) - enc_val_on_aggr = loop.run_until_complete( - strat_ex._move(val_enc.internal_representation[0], - val_enc.type_signature.member, aggr)) +# aggr = strat_ex._get_child_executors( +# placement_literals.AGGREGATORS, index=0) - val_key_zipped = loop.run_until_complete( - strat_ex._zip_val_key([enc_val_on_aggr], sk_a, - placement_literals.AGGREGATORS)) +# enc_val_on_aggr = loop.run_until_complete( +# strat_ex._move(val_enc.internal_representation[0], +# val_enc.type_signature.member, aggr)) - val_dec = loop.run_until_complete( - strat_ex._decrypt_tensors_on_aggregator(val_key_zipped, tf.float32)) +# val_key_zipped = loop.run_until_complete( +# strat_ex._zip_val_key([enc_val_on_aggr], sk_a, +# placement_literals.AGGREGATORS)) - dec_tf_tensor = val_dec.internal_representation[0].internal_representation +# val_dec = loop.run_until_complete( +# strat_ex._decrypt_tensors_on_aggregator(val_key_zipped, tf.float32)) - self.assertEqual(dec_tf_tensor, tf.constant(2.0, dtype=tf.float32)) +# dec_tf_tensor = val_dec.internal_representation[0].internal_representation +# self.assertEqual(dec_tf_tensor, tf.constant(2.0, dtype=tf.float32)) if __name__ == '__main__': absltest.main() From 9564eda5d557160ec81d5f75d4bb3a9c296e3f26 Mon Sep 17 00:00:00 2001 From: yann dupis Date: Thu, 7 May 2020 19:46:23 -0700 Subject: [PATCH 07/31] fix names --- .../core/impl/executors/federating_executor.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor.py b/tensorflow_federated/python/core/impl/executors/federating_executor.py index 4c77f781a2..0a4f452f11 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor.py @@ -508,7 +508,7 @@ async def _move(cls, arg, target_executor, parent_executor, channel): py_typecheck.check_type(val_type, computation_types.FederatedType) item_type = val_type.member - + val = await asyncio.gather(*[ target_executor.create_value(await v.compute(), item_type) for v in val ]) @@ -517,7 +517,7 @@ async def _move(cls, arg, target_executor, parent_executor, channel): channel.receive( parent_executor=parent_executor, val=v, - receiver_sk=sk_a, + sk_receiver=sk_a, sender_dtype=orig_clients_dtype) for v in val ]) # NOTE should check if we can remove this step @@ -1130,10 +1130,10 @@ async def setup(self, parent_executor): if not self.is_encrypted: return - receiver_pk, receiver_sk = await self._receiver_generate_keys( + pk_receiver, sk_receiver = await self._receiver_generate_keys( parent_executor) - return receiver_pk, receiver_sk + return pk_receiver, sk_receiver async def send(self, parent_executor, val, pk_receiver): @@ -1142,13 +1142,13 @@ async def send(self, parent_executor, val, pk_receiver): return await self._encrypt_sender_tensors(parent_executor, val, pk_receiver) - async def receive(self, parent_executor, val, receiver_sk, sender_dtype): + async def receive(self, parent_executor, val, sk_receiver, sender_dtype): if not self.is_encrypted: return val return await self._decrypt_tensors_on_receiver(parent_executor, val, - receiver_sk, sender_dtype) + sk_receiver, sender_dtype) async def _receiver_generate_keys(self, parent_executor): # NOTE solve confusion between executor and receiver_executor @@ -1245,9 +1245,9 @@ def encrypt_tensor(plaintext, pk_rcv): computation_types.NamedTupleType((fn_type, val_type)))) async def _decrypt_tensors_on_receiver(self, parent_executor, val, - receiver_sk, sender_dtype): + sk_rcv, sender_dtype): - val = await self._zip_val_key(parent_executor, val, receiver_sk, + val = await self._zip_val_key(parent_executor, val, sk_rcv, self.receiver_placement) sender_output_type_signature = val[0].type_signature[0] From 1358972d666316807a0cc091ef4f1db8a1368e32 Mon Sep 17 00:00:00 2001 From: yann dupis Date: Mon, 11 May 2020 14:28:09 -0700 Subject: [PATCH 08/31] generate clients keys during setup phase --- .../impl/executors/federating_executor.py | 142 +++++++++++++----- 1 file changed, 101 insertions(+), 41 deletions(-) diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor.py b/tensorflow_federated/python/core/impl/executors/federating_executor.py index 0a4f452f11..6e6207dfec 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor.py @@ -461,6 +461,11 @@ class TrustedAggregatorIntrinsicStrategy(IntrinsicStrategy): def __init__(self, federating_executor): super().__init__(federating_executor) + self.channel = Channel( + sender_placement=placement_literals.CLIENTS, + receiver_placement=placement_literals.AGGREGATORS, + is_encrypted=True) + @classmethod def validate_executor_placements(cls, executor_placements): singleton_placements = [ @@ -492,13 +497,12 @@ def validate_executor_placements(cls, executor_placements): 'Unsupported cardinality for placement "{}": {}.'.format( pl, pl_cardinality)) - @classmethod - async def _move(cls, arg, target_executor, parent_executor, channel): + async def _move(self, arg, target_executor, parent_executor): - pk_a, sk_a = await channel.setup(parent_executor=parent_executor) + await self.channel.setup(parent_executor=parent_executor) - enc_clients_arg = await channel.send( - parent_executor=parent_executor, val=arg, pk_receiver=pk_a) + enc_clients_arg = await self.channel.send( + parent_executor=parent_executor, val=arg) val_type = enc_clients_arg.type_signature val = enc_clients_arg.internal_representation @@ -514,10 +518,9 @@ async def _move(cls, arg, target_executor, parent_executor, channel): ]) received_vals = await asyncio.gather(*[ - channel.receive( + self.channel.receive( parent_executor=parent_executor, val=v, - sk_receiver=sk_a, sender_dtype=orig_clients_dtype) for v in val ]) # NOTE should check if we can remove this step @@ -574,18 +577,12 @@ async def federated_reduce(self, arg): 'Expected 3 elements in the `federated_reduce()` argument tuple, ' 'found {}.'.format(len(arg.internal_representation))) - # NOTE it will be supplied to the intrinsic strategy - agg_cls_channel = Channel( - sender_placement=placement_literals.CLIENTS, - receiver_placement=placement_literals.AGGREGATORS, - is_encrypted=True) - zero_type = arg.type_signature[1] op_type = arg.type_signature[2] aggr = self._get_child_executors(placement_literals.AGGREGATORS, index=0) - aggregands = await self._move(arg, aggr, self, agg_cls_channel) + aggregands = await self._move(arg, aggr, self) zero = await aggr.create_value( await (await @@ -1125,33 +1122,43 @@ def __init__(self, sender_placement, receiver_placement, is_encrypted): self.receiver_placement = receiver_placement self.is_encrypted = is_encrypted + self.pk_receiver = None + self.sk_receiver = None + self.pk_sender = None + self.sk_sender = None + + self.key_store = KeyStore() + async def setup(self, parent_executor): if not self.is_encrypted: return - pk_receiver, sk_receiver = await self._receiver_generate_keys( + self.pk_receiver, self.sk_receiver = await self._receiver_generate_keys( parent_executor) - return pk_receiver, sk_receiver + self.pk_sender, self.sk_sender = await self._sender_generate_keys(parent_executor) + + self.key_store.add_keys(self.sender_placement, self.pk_sender, self.sk_sender) + self.key_store.add_keys(self.receiver_placement, self.pk_receiver, self.sk_receiver) + + return - async def send(self, parent_executor, val, pk_receiver): + async def send(self, parent_executor, val): if not self.is_encrypted: return val - return await self._encrypt_sender_tensors(parent_executor, val, pk_receiver) + return await self._encrypt_sender_tensors(parent_executor, val) - async def receive(self, parent_executor, val, sk_receiver, sender_dtype): + async def receive(self, parent_executor, val, sender_dtype): if not self.is_encrypted: return val - return await self._decrypt_tensors_on_receiver(parent_executor, val, - sk_receiver, sender_dtype) + return await self._decrypt_tensors_on_receiver(parent_executor, val, sender_dtype) async def _receiver_generate_keys(self, parent_executor): - # NOTE solve confusion between executor and receiver_executor @computations.tf_computation() def generate_keys(): @@ -1183,7 +1190,42 @@ def generate_keys(): return pk_fed_val, sk_fed_val - async def _zip_val_key(self, parent_executor, vals, key, placement): + + async def _sender_generate_keys(self, parent_executor): + + @computations.tf_computation() + def generate_keys(): + pk, sk = easy_box.gen_keypair() + return pk.raw, sk.raw + + fn_type = generate_keys.type_signature + fn = generate_keys._computation_proto + + # sender_ex = parent_executor._get_child_executors( + # self.sender_placement, index=0) + + sender_ex = parent_executor._get_child_executors(self.sender_placement, index=0) + + key_generator = await sender_ex.create_call(await + sender_ex.create_value( + fn, fn_type)) + + keys = await asyncio.gather(*[ + sender_ex.create_selection(key_generator, i) + for i in range(len(key_generator.type_signature)) + ]) + + # pk_fed_val, sk_fed_val = await asyncio.gather( + # *[parent_executor._place(k, self.sender_placement) for k in keys]) + + sk_fed_val = await parent_executor._place(keys[1], self.sender_placement) + + pk_fed_val = await parent_executor._place(keys[0], self.receiver_placement) + + return pk_fed_val, sk_fed_val + + async def _zip_val_key(self, parent_executor, vals, pk_key, sk_key, placement): + # import pdb; pdb.set_trace() if isinstance(vals, list): val_type = computation_types.FederatedType( @@ -1195,15 +1237,16 @@ async def _zip_val_key(self, parent_executor, vals, key, placement): vals_key = FederatingExecutorValue( anonymous_tuple.AnonymousTuple([(None, vals), - (None, key.internal_representation)]), - computation_types.NamedTupleType((val_type, key.type_signature))) + (None, pk_key.internal_representation), + (None, sk_key.internal_representation)]), + computation_types.NamedTupleType((val_type, pk_key.type_signature, sk_key.type_signature))) vals_key_zipped = await parent_executor._zip( vals_key, placement, all_equal=False) return vals_key_zipped.internal_representation - async def _encrypt_sender_tensors(self, parent_executor, val, pk_rcv): + async def _encrypt_sender_tensors(self, parent_executor, val): nb_senders = len( parent_executor._get_child_executors(self.sender_placement)) @@ -1212,17 +1255,19 @@ async def _encrypt_sender_tensors(self, parent_executor, val, pk_rcv): input_tensor_type = val.type_signature.member else: input_tensor_type = val.type_signature[0].member - pk_rcv_tensor_type = pk_rcv.type_signature.member + pk_rcv_tensor_type = self.pk_receiver.type_signature.member + sk_sender_tensor_type = self.sk_sender.type_signature.member - @computations.tf_computation(input_tensor_type, pk_rcv_tensor_type) - def encrypt_tensor(plaintext, pk_rcv): + @computations.tf_computation(input_tensor_type, pk_rcv_tensor_type, sk_sender_tensor_type) + def encrypt_tensor(plaintext, pk_rcv, sk_snd): pk_rcv = easy_box.PublicKey(pk_rcv) - pk_snd, sk_snd = easy_box.gen_keypair() + sk_snd = easy_box.PublicKey(sk_snd) + nonce = easy_box.gen_nonce() ciphertext, mac = easy_box.seal_detached(plaintext, nonce, pk_rcv, sk_snd) - return ciphertext.raw, mac.raw, pk_snd.raw, nonce.raw + return ciphertext.raw, mac.raw, nonce.raw fn_type = encrypt_tensor.type_signature fn = encrypt_tensor._computation_proto @@ -1232,9 +1277,10 @@ def encrypt_tensor(plaintext, pk_rcv): else: val_type = val.type_signature[0] val = val.internal_representation[0] - - val_key_zipped = await self._zip_val_key(parent_executor, val, pk_rcv, + + val_key_zipped = await self._zip_val_key(parent_executor, val, self.pk_receiver, self.sk_sender, self.sender_placement) + # NOTE probably won't always be fed_ex in future design fed_ex = parent_executor.federating_executor @@ -1244,24 +1290,25 @@ def encrypt_tensor(plaintext, pk_rcv): (None, val_key_zipped)]), computation_types.NamedTupleType((fn_type, val_type)))) - async def _decrypt_tensors_on_receiver(self, parent_executor, val, - sk_rcv, sender_dtype): + async def _decrypt_tensors_on_receiver(self, parent_executor, val, sender_dtype): - val = await self._zip_val_key(parent_executor, val, sk_rcv, + val = await self._zip_val_key(parent_executor, val, self.sk_receiver, self.pk_sender, self.receiver_placement) sender_output_type_signature = val[0].type_signature[0] - receiver_public_key_type_signature = val[0].type_signature[1] + receiver_secret_key_type_signature = val[0].type_signature[1] + sender_public_key_type_signature = val[0].type_signature[2] @computations.tf_computation(sender_output_type_signature, - receiver_public_key_type_signature) - def decrypt_tensor(client_outputs, sk_rcv): + receiver_secret_key_type_signature, + sender_public_key_type_signature) + def decrypt_tensor(client_outputs, sk_rcv, pk_snd): ciphertext = easy_box.Ciphertext(client_outputs[0]) mac = easy_box.Mac(client_outputs[1]) - pk_snd = easy_box.PublicKey(client_outputs[2]) - nonce = easy_box.Nonce(client_outputs[3]) + nonce = easy_box.Nonce(client_outputs[2]) sk_rcv = easy_box.SecretKey(sk_rcv) + pk_snd = easy_box.PublicKey(pk_snd) plaintext_recovered = easy_box.open_detached(ciphertext, mac, nonce, pk_snd, sk_rcv, sender_dtype) @@ -1282,3 +1329,16 @@ def decrypt_tensor(client_outputs, sk_rcv): FederatingExecutorValue( anonymous_tuple.AnonymousTuple([(None, fn), (None, val)]), computation_types.NamedTupleType((fn_type, val_type)))) + + +class KeyStore: + def __init__(self): + self.key_store = {} + + def add_keys(self, placement_name, pk, sk): + self.key_store[placement_name] = {} + self.key_store[placement_name]['pk'] = pk + self.key_store[placement_name]['sk'] = sk + + def get_keys(self, placement_name): + return self.key_store \ No newline at end of file From 98631a308cab1f9b7427ddc8eb0bba55c2140535 Mon Sep 17 00:00:00 2001 From: yann dupis Date: Mon, 11 May 2020 19:47:10 -0700 Subject: [PATCH 09/31] working with keys generated during setup --- .../impl/executors/federating_executor.py | 217 ++++++++++++------ 1 file changed, 149 insertions(+), 68 deletions(-) diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor.py b/tensorflow_federated/python/core/impl/executors/federating_executor.py index 6e6207dfec..e90bda44ad 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor.py @@ -497,13 +497,51 @@ def validate_executor_placements(cls, executor_placements): 'Unsupported cardinality for placement "{}": {}.'.format( pl, pl_cardinality)) + @tracing.trace + async def _place_keys(self, keys, placement): + + py_typecheck.check_type(placement, placement_literals.PlacementLiteral) + children = self._get_child_executors(placement) + # val = await arg.compute() + if not isinstance(children, list): + children = [children] + + if len(keys) == len(children): + keys_type_signature = keys[0].type_signature + return FederatingExecutorValue( + await asyncio.gather(*[ + c.create_value(await keys[i].compute(), keys_type_signature) + for (i, c) in enumerate(children) + ]), + computation_types.FederatedType( + keys_type_signature, placement, all_equal=False)) + elif len(keys) > len(children): + keys_type_signature = keys[0].type_signature + child = children[0] + return FederatingExecutorValue( + await asyncio.gather(*[ + child.create_value(await k.compute(), keys_type_signature) + for k in keys + ]), + computation_types.FederatedType( + keys_type_signature, placement, all_equal=False)) + elif (len(keys) == 1) & (len(children) > 1): + keys_type_signature = keys[0].type_signature + return FederatingExecutorValue( + await asyncio.gather(*[ + c.create_value(await keys[0].compute(), keys_type_signature) + for c in children + ]), + computation_types.FederatedType( + keys_type_signature, placement, all_equal=True)) + async def _move(self, arg, target_executor, parent_executor): await self.channel.setup(parent_executor=parent_executor) enc_clients_arg = await self.channel.send( parent_executor=parent_executor, val=arg) - + val_type = enc_clients_arg.type_signature val = enc_clients_arg.internal_representation orig_clients_dtype = arg.type_signature[0].member.dtype @@ -521,9 +559,10 @@ async def _move(self, arg, target_executor, parent_executor): self.channel.receive( parent_executor=parent_executor, val=v, - sender_dtype=orig_clients_dtype) for v in val + sender_dtype=orig_clients_dtype, + sender_index=i) for (i, v) in enumerate(val) ]) - # NOTE should check if we can remove this step + received_vals = [v.internal_representation[0] for v in received_vals] return received_vals @@ -1124,8 +1163,6 @@ def __init__(self, sender_placement, receiver_placement, is_encrypted): self.pk_receiver = None self.sk_receiver = None - self.pk_sender = None - self.sk_sender = None self.key_store = KeyStore() @@ -1134,29 +1171,32 @@ async def setup(self, parent_executor): if not self.is_encrypted: return - self.pk_receiver, self.sk_receiver = await self._receiver_generate_keys( - parent_executor) + await self._receiver_generate_keys(parent_executor) + await self._sender_generate_keys(parent_executor) - self.pk_sender, self.sk_sender = await self._sender_generate_keys(parent_executor) + return - self.key_store.add_keys(self.sender_placement, self.pk_sender, self.sk_sender) - self.key_store.add_keys(self.receiver_placement, self.pk_receiver, self.sk_receiver) - - return - - async def send(self, parent_executor, val): + async def send(self, parent_executor, val, sender_index=0, receiver_index=0): if not self.is_encrypted: return val - return await self._encrypt_sender_tensors(parent_executor, val) + return await self._encrypt_sender_tensors(parent_executor, val, + sender_index, receiver_index) - async def receive(self, parent_executor, val, sender_dtype): + async def receive(self, + parent_executor, + val, + sender_dtype, + receiver_index=0, + sender_index=0): if not self.is_encrypted: return val - return await self._decrypt_tensors_on_receiver(parent_executor, val, sender_dtype) + return await self._decrypt_tensors_on_receiver(parent_executor, val, + sender_dtype, sender_index, + receiver_index) async def _receiver_generate_keys(self, parent_executor): @@ -1180,16 +1220,20 @@ def generate_keys(): for i in range(len(key_generator.type_signature)) ]) - pk_fed_val, sk_fed_val = await asyncio.gather( - *[parent_executor._place(k, self.receiver_placement) for k in keys]) + # pk_fed_val, sk_fed_val = await asyncio.gather( + # *[parent_executor._place_keys([k], self.receiver_placement) for k in keys]) + + # pk_fed_val = await parent_executor.federated_broadcast(pk_fed_val) - # NOTE: This broadcast should work from aggregator to clients - # but also from server to aggregator (so server can receive - # final updates). May have to use _place instead - pk_fed_val = await parent_executor.federated_broadcast(pk_fed_val) + sk_fed_val = await parent_executor._place_keys([keys[1]], + self.receiver_placement) + pk_fed_val = await parent_executor._place_keys([keys[0]], + self.sender_placement) - return pk_fed_val, sk_fed_val + self.key_store.add_keys(self.receiver_placement.name, pk_fed_val, + sk_fed_val) + return async def _sender_generate_keys(self, parent_executor): @@ -1201,31 +1245,40 @@ def generate_keys(): fn_type = generate_keys.type_signature fn = generate_keys._computation_proto - # sender_ex = parent_executor._get_child_executors( - # self.sender_placement, index=0) - - sender_ex = parent_executor._get_child_executors(self.sender_placement, index=0) - - key_generator = await sender_ex.create_call(await - sender_ex.create_value( - fn, fn_type)) - - keys = await asyncio.gather(*[ - sender_ex.create_selection(key_generator, i) - for i in range(len(key_generator.type_signature)) - ]) - - # pk_fed_val, sk_fed_val = await asyncio.gather( - # *[parent_executor._place(k, self.sender_placement) for k in keys]) - - sk_fed_val = await parent_executor._place(keys[1], self.sender_placement) + senders_ex = parent_executor._get_child_executors(self.sender_placement) + nb_senders = len(senders_ex) + sk_vals = [] + pk_vals = [] + + if nb_senders > 1: + for i in range(nb_senders): + sender_ex = senders_ex[0] + key_generator = await sender_ex.create_call(await + sender_ex.create_value( + fn, fn_type)) + + keys = await asyncio.gather(*[ + sender_ex.create_selection(key_generator, i) + for i in range(len(key_generator.type_signature)) + ]) + sk_vals.append(keys[1]) + pk_vals.append(keys[0]) - pk_fed_val = await parent_executor._place(keys[0], self.receiver_placement) + sk_fed_val = await parent_executor._place_keys(sk_vals, + self.sender_placement) + pk_fed_val = await parent_executor._place_keys(pk_vals, + self.receiver_placement) - return pk_fed_val, sk_fed_val + self.key_store.add_keys(self.sender_placement.name, pk_fed_val, sk_fed_val) + return - async def _zip_val_key(self, parent_executor, vals, pk_key, sk_key, placement): - # import pdb; pdb.set_trace() + async def _zip_val_key(self, + parent_executor, + vals, + pk_key, + sk_key, + placement, + sender_index=None): if isinstance(vals, list): val_type = computation_types.FederatedType( @@ -1235,18 +1288,28 @@ async def _zip_val_key(self, parent_executor, vals, pk_key, sk_key, placement): vals.type_signature, placement, all_equal=False) vals = [vals] + pk_key_vals = pk_key.internal_representation + sk_key_vals = sk_key.internal_representation + + if sender_index != None: + sk_key_vals = [sk_key_vals[sender_index]] + vals_key = FederatingExecutorValue( - anonymous_tuple.AnonymousTuple([(None, vals), - (None, pk_key.internal_representation), - (None, sk_key.internal_representation)]), - computation_types.NamedTupleType((val_type, pk_key.type_signature, sk_key.type_signature))) + anonymous_tuple.AnonymousTuple([(None, vals), (None, pk_key_vals), + (None, sk_key_vals)]), + computation_types.NamedTupleType( + (val_type, pk_key.type_signature, sk_key.type_signature))) vals_key_zipped = await parent_executor._zip( vals_key, placement, all_equal=False) return vals_key_zipped.internal_representation - async def _encrypt_sender_tensors(self, parent_executor, val): + async def _encrypt_sender_tensors(self, + parent_executor, + val, + sender_index=0, + receiver_index=0): nb_senders = len( parent_executor._get_child_executors(self.sender_placement)) @@ -1255,10 +1318,14 @@ async def _encrypt_sender_tensors(self, parent_executor, val): input_tensor_type = val.type_signature.member else: input_tensor_type = val.type_signature[0].member - pk_rcv_tensor_type = self.pk_receiver.type_signature.member - sk_sender_tensor_type = self.sk_sender.type_signature.member - @computations.tf_computation(input_tensor_type, pk_rcv_tensor_type, sk_sender_tensor_type) + pk_receiver = self.key_store.get_keys(self.receiver_placement.name)['pk'] + sk_sender = self.key_store.get_keys(self.sender_placement.name)['sk'] + pk_rcv_tensor_type = pk_receiver.type_signature.member + sk_sender_tensor_type = sk_sender.type_signature.member + + @computations.tf_computation(input_tensor_type, pk_rcv_tensor_type, + sk_sender_tensor_type) def encrypt_tensor(plaintext, pk_rcv, sk_snd): pk_rcv = easy_box.PublicKey(pk_rcv) @@ -1277,9 +1344,9 @@ def encrypt_tensor(plaintext, pk_rcv, sk_snd): else: val_type = val.type_signature[0] val = val.internal_representation[0] - - val_key_zipped = await self._zip_val_key(parent_executor, val, self.pk_receiver, self.sk_sender, - self.sender_placement) + + val_key_zipped = await self._zip_val_key(parent_executor, val, pk_receiver, + sk_sender, self.sender_placement) # NOTE probably won't always be fed_ex in future design fed_ex = parent_executor.federating_executor @@ -1290,10 +1357,23 @@ def encrypt_tensor(plaintext, pk_rcv, sk_snd): (None, val_key_zipped)]), computation_types.NamedTupleType((fn_type, val_type)))) - async def _decrypt_tensors_on_receiver(self, parent_executor, val, sender_dtype): - - val = await self._zip_val_key(parent_executor, val, self.sk_receiver, self.pk_sender, - self.receiver_placement) + async def _decrypt_tensors_on_receiver(self, + parent_executor, + val, + sender_dtype, + sender_index=0, + receive_index=0): + + pk_sender = self.key_store.get_keys(self.sender_placement.name)['pk'] + sk_receiver = self.key_store.get_keys(self.receiver_placement.name)['sk'] + + val = await self._zip_val_key( + parent_executor, + val, + sk_receiver, + pk_sender, + self.receiver_placement, + sender_index=sender_index) sender_output_type_signature = val[0].type_signature[0] receiver_secret_key_type_signature = val[0].type_signature[1] @@ -1332,13 +1412,14 @@ def decrypt_tensor(client_outputs, sk_rcv, pk_snd): class KeyStore: + def __init__(self): self.key_store = {} - def add_keys(self, placement_name, pk, sk): - self.key_store[placement_name] = {} - self.key_store[placement_name]['pk'] = pk - self.key_store[placement_name]['sk'] = sk - - def get_keys(self, placement_name): - return self.key_store \ No newline at end of file + def add_keys(self, key_owner, pk, sk): + self.key_store[key_owner] = {} + self.key_store[key_owner]['pk'] = pk + self.key_store[key_owner]['sk'] = sk + + def get_keys(self, key_owner): + return self.key_store[key_owner] From aa6c128c9a373b326cfa008d358c8807560f9ec1 Mon Sep 17 00:00:00 2001 From: yann dupis Date: Tue, 12 May 2020 13:23:14 -0700 Subject: [PATCH 10/31] separate keys placement from generate keys --- .../impl/executors/federating_executor.py | 142 +++++++----------- 1 file changed, 54 insertions(+), 88 deletions(-) diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor.py b/tensorflow_federated/python/core/impl/executors/federating_executor.py index e90bda44ad..cc1501f5f6 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor.py @@ -462,6 +462,7 @@ def __init__(self, federating_executor): super().__init__(federating_executor) self.channel = Channel( + parent_executor=self, sender_placement=placement_literals.CLIENTS, receiver_placement=placement_literals.AGGREGATORS, is_encrypted=True) @@ -537,10 +538,9 @@ async def _place_keys(self, keys, placement): async def _move(self, arg, target_executor, parent_executor): - await self.channel.setup(parent_executor=parent_executor) + await self.channel.setup() - enc_clients_arg = await self.channel.send( - parent_executor=parent_executor, val=arg) + enc_clients_arg = await self.channel.send(val=arg) val_type = enc_clients_arg.type_signature val = enc_clients_arg.internal_representation @@ -557,10 +557,8 @@ async def _move(self, arg, target_executor, parent_executor): received_vals = await asyncio.gather(*[ self.channel.receive( - parent_executor=parent_executor, - val=v, - sender_dtype=orig_clients_dtype, - sender_index=i) for (i, v) in enumerate(val) + val=v, sender_dtype=orig_clients_dtype, sender_index=i) + for (i, v) in enumerate(val) ]) received_vals = [v.internal_representation[0] for v in received_vals] @@ -1155,50 +1153,46 @@ async def _compute_intrinsic_federated_secure_sum(self, arg): class Channel: - def __init__(self, sender_placement, receiver_placement, is_encrypted): + def __init__(self, parent_executor, sender_placement, receiver_placement, + is_encrypted): + self.parent_executor = parent_executor self.sender_placement = sender_placement self.receiver_placement = receiver_placement self.is_encrypted = is_encrypted - self.pk_receiver = None - self.sk_receiver = None - self.key_store = KeyStore() - async def setup(self, parent_executor): + async def setup(self): if not self.is_encrypted: return - await self._receiver_generate_keys(parent_executor) - await self._sender_generate_keys(parent_executor) + await self._generate_keys(self.sender_placement) + await self._generate_keys(self.receiver_placement) + await self._share_public_keys(self.sender_placement, + self.receiver_placement) + await self._share_public_keys(self.receiver_placement, + self.sender_placement) return - async def send(self, parent_executor, val, sender_index=0, receiver_index=0): + async def send(self, val, sender_index=0, receiver_index=0): if not self.is_encrypted: return val - return await self._encrypt_sender_tensors(parent_executor, val, - sender_index, receiver_index) + return await self._encrypt_sender_tensors(val, sender_index, receiver_index) - async def receive(self, - parent_executor, - val, - sender_dtype, - receiver_index=0, - sender_index=0): + async def receive(self, val, sender_dtype, receiver_index=0, sender_index=0): if not self.is_encrypted: return val - return await self._decrypt_tensors_on_receiver(parent_executor, val, - sender_dtype, sender_index, - receiver_index) + return await self._decrypt_tensors_on_receiver(val, sender_dtype, + sender_index, receiver_index) - async def _receiver_generate_keys(self, parent_executor): + async def _generate_keys(self, key_owner_placement): @computations.tf_computation() def generate_keys(): @@ -1208,72 +1202,43 @@ def generate_keys(): fn_type = generate_keys.type_signature fn = generate_keys._computation_proto - receiver_ex = parent_executor._get_child_executors( - self.receiver_placement, index=0) + executors = self.parent_executor._get_child_executors(key_owner_placement) - key_generator = await receiver_ex.create_call(await - receiver_ex.create_value( - fn, fn_type)) + nb_executors = len(executors) - keys = await asyncio.gather(*[ - receiver_ex.create_selection(key_generator, i) - for i in range(len(key_generator.type_signature)) - ]) + sk_vals = [] + pk_vals = [] - # pk_fed_val, sk_fed_val = await asyncio.gather( - # *[parent_executor._place_keys([k], self.receiver_placement) for k in keys]) + for i in range(nb_executors): + executor = executors[0] + key_generator = await executor.create_call(await executor.create_value( + fn, fn_type)) - # pk_fed_val = await parent_executor.federated_broadcast(pk_fed_val) + keys = await asyncio.gather(*[ + executor.create_selection(key_generator, i) + for i in range(len(key_generator.type_signature)) + ]) - sk_fed_val = await parent_executor._place_keys([keys[1]], - self.receiver_placement) - pk_fed_val = await parent_executor._place_keys([keys[0]], - self.sender_placement) + pk_vals.append(keys[0]) + sk_vals.append(keys[1]) - self.key_store.add_keys(self.receiver_placement.name, pk_fed_val, - sk_fed_val) + self.key_store.add_keys(key_owner_placement.name, pk_vals, sk_vals) return - async def _sender_generate_keys(self, parent_executor): - - @computations.tf_computation() - def generate_keys(): - pk, sk = easy_box.gen_keypair() - return pk.raw, sk.raw - - fn_type = generate_keys.type_signature - fn = generate_keys._computation_proto - - senders_ex = parent_executor._get_child_executors(self.sender_placement) - nb_senders = len(senders_ex) - sk_vals = [] - pk_vals = [] + async def _share_public_keys(self, key_owner_placement, send_public_keys_to): - if nb_senders > 1: - for i in range(nb_senders): - sender_ex = senders_ex[0] - key_generator = await sender_ex.create_call(await - sender_ex.create_value( - fn, fn_type)) + keys = self.key_store.get_keys(key_owner_placement.name) - keys = await asyncio.gather(*[ - sender_ex.create_selection(key_generator, i) - for i in range(len(key_generator.type_signature)) - ]) - sk_vals.append(keys[1]) - pk_vals.append(keys[0]) + sk_fed_vals = await self.parent_executor._place_keys( + keys['sk'], key_owner_placement) + pk_fed_vals = await self.parent_executor._place_keys( + keys['pk'], send_public_keys_to) - sk_fed_val = await parent_executor._place_keys(sk_vals, - self.sender_placement) - pk_fed_val = await parent_executor._place_keys(pk_vals, - self.receiver_placement) - - self.key_store.add_keys(self.sender_placement.name, pk_fed_val, sk_fed_val) - return + self.key_store.update_keys(key_owner_placement.name, pk_fed_vals, + sk_fed_vals) async def _zip_val_key(self, - parent_executor, vals, pk_key, sk_key, @@ -1300,19 +1265,18 @@ async def _zip_val_key(self, computation_types.NamedTupleType( (val_type, pk_key.type_signature, sk_key.type_signature))) - vals_key_zipped = await parent_executor._zip( + vals_key_zipped = await self.parent_executor._zip( vals_key, placement, all_equal=False) return vals_key_zipped.internal_representation async def _encrypt_sender_tensors(self, - parent_executor, val, sender_index=0, receiver_index=0): nb_senders = len( - parent_executor._get_child_executors(self.sender_placement)) + self.parent_executor._get_child_executors(self.sender_placement)) if nb_senders == 1: input_tensor_type = val.type_signature.member @@ -1345,11 +1309,11 @@ def encrypt_tensor(plaintext, pk_rcv, sk_snd): val_type = val.type_signature[0] val = val.internal_representation[0] - val_key_zipped = await self._zip_val_key(parent_executor, val, pk_receiver, - sk_sender, self.sender_placement) + val_key_zipped = await self._zip_val_key(val, pk_receiver, sk_sender, + self.sender_placement) # NOTE probably won't always be fed_ex in future design - fed_ex = parent_executor.federating_executor + fed_ex = self.parent_executor.federating_executor return await fed_ex._compute_intrinsic_federated_map( FederatingExecutorValue( @@ -1358,7 +1322,6 @@ def encrypt_tensor(plaintext, pk_rcv, sk_snd): computation_types.NamedTupleType((fn_type, val_type)))) async def _decrypt_tensors_on_receiver(self, - parent_executor, val, sender_dtype, sender_index=0, @@ -1368,7 +1331,6 @@ async def _decrypt_tensors_on_receiver(self, sk_receiver = self.key_store.get_keys(self.receiver_placement.name)['sk'] val = await self._zip_val_key( - parent_executor, val, sk_receiver, pk_sender, @@ -1403,7 +1365,7 @@ def decrypt_tensor(client_outputs, sk_rcv, pk_snd): fn_type = decrypt_tensor.type_signature fn = decrypt_tensor._computation_proto # NOTE probably won't always be fed_ex in future design - fed_ex = parent_executor.federating_executor + fed_ex = self.parent_executor.federating_executor return await fed_ex._compute_intrinsic_federated_map( FederatingExecutorValue( @@ -1423,3 +1385,7 @@ def add_keys(self, key_owner, pk, sk): def get_keys(self, key_owner): return self.key_store[key_owner] + + def update_keys(self, key_owner, pk, sk): + self.key_store[key_owner]['pk'] = pk + self.key_store[key_owner]['sk'] = sk From 4f00d2804df9b174814438e1cdffec67b9bccd68 Mon Sep 17 00:00:00 2001 From: yann dupis Date: Tue, 12 May 2020 14:11:06 -0700 Subject: [PATCH 11/31] clean up --- .../impl/executors/federating_executor.py | 119 +++++++++--------- 1 file changed, 60 insertions(+), 59 deletions(-) diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor.py b/tensorflow_federated/python/core/impl/executors/federating_executor.py index cc1501f5f6..d35c5faad4 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor.py @@ -498,56 +498,14 @@ def validate_executor_placements(cls, executor_placements): 'Unsupported cardinality for placement "{}": {}.'.format( pl, pl_cardinality)) - @tracing.trace - async def _place_keys(self, keys, placement): - - py_typecheck.check_type(placement, placement_literals.PlacementLiteral) - children = self._get_child_executors(placement) - # val = await arg.compute() - if not isinstance(children, list): - children = [children] - - if len(keys) == len(children): - keys_type_signature = keys[0].type_signature - return FederatingExecutorValue( - await asyncio.gather(*[ - c.create_value(await keys[i].compute(), keys_type_signature) - for (i, c) in enumerate(children) - ]), - computation_types.FederatedType( - keys_type_signature, placement, all_equal=False)) - elif len(keys) > len(children): - keys_type_signature = keys[0].type_signature - child = children[0] - return FederatingExecutorValue( - await asyncio.gather(*[ - child.create_value(await k.compute(), keys_type_signature) - for k in keys - ]), - computation_types.FederatedType( - keys_type_signature, placement, all_equal=False)) - elif (len(keys) == 1) & (len(children) > 1): - keys_type_signature = keys[0].type_signature - return FederatingExecutorValue( - await asyncio.gather(*[ - c.create_value(await keys[0].compute(), keys_type_signature) - for c in children - ]), - computation_types.FederatedType( - keys_type_signature, placement, all_equal=True)) - - async def _move(self, arg, target_executor, parent_executor): - - await self.channel.setup() + async def _move(self, arg, target_executor): enc_clients_arg = await self.channel.send(val=arg) val_type = enc_clients_arg.type_signature val = enc_clients_arg.internal_representation orig_clients_dtype = arg.type_signature[0].member.dtype - py_typecheck.check_type(val, list) - py_typecheck.check_type(val_type, computation_types.FederatedType) item_type = val_type.member @@ -556,9 +514,7 @@ async def _move(self, arg, target_executor, parent_executor): ]) received_vals = await asyncio.gather(*[ - self.channel.receive( - val=v, sender_dtype=orig_clients_dtype, sender_index=i) - for (i, v) in enumerate(val) + self.channel.receive(val=v, sender_index=i) for (i, v) in enumerate(val) ]) received_vals = [v.internal_representation[0] for v in received_vals] @@ -614,12 +570,14 @@ async def federated_reduce(self, arg): 'Expected 3 elements in the `federated_reduce()` argument tuple, ' 'found {}.'.format(len(arg.internal_representation))) - zero_type = arg.type_signature[1] - op_type = arg.type_signature[2] + # Setup channel keys + await self.channel.setup() aggr = self._get_child_executors(placement_literals.AGGREGATORS, index=0) + aggregands = await self._move(arg, aggr) - aggregands = await self._move(arg, aggr, self) + zero_type = arg.type_signature[1] + op_type = arg.type_signature[2] zero = await aggr.create_value( await (await @@ -1184,13 +1142,13 @@ async def send(self, val, sender_index=0, receiver_index=0): return await self._encrypt_sender_tensors(val, sender_index, receiver_index) - async def receive(self, val, sender_dtype, receiver_index=0, sender_index=0): + async def receive(self, val, receiver_index=0, sender_index=0): if not self.is_encrypted: return val - return await self._decrypt_tensors_on_receiver(val, sender_dtype, - sender_index, receiver_index) + return await self._decrypt_tensors_on_receiver(val, sender_index, + receiver_index) async def _generate_keys(self, key_owner_placement): @@ -1205,7 +1163,6 @@ def generate_keys(): executors = self.parent_executor._get_child_executors(key_owner_placement) nb_executors = len(executors) - sk_vals = [] pk_vals = [] @@ -1230,9 +1187,9 @@ async def _share_public_keys(self, key_owner_placement, send_public_keys_to): keys = self.key_store.get_keys(key_owner_placement.name) - sk_fed_vals = await self.parent_executor._place_keys( + sk_fed_vals = await self._place_keys( keys['sk'], key_owner_placement) - pk_fed_vals = await self.parent_executor._place_keys( + pk_fed_vals = await self._place_keys( keys['pk'], send_public_keys_to) self.key_store.update_keys(key_owner_placement.name, pk_fed_vals, @@ -1270,6 +1227,48 @@ async def _zip_val_key(self, return vals_key_zipped.internal_representation + async def _place_keys(self, keys, placement): + + py_typecheck.check_type(placement, placement_literals.PlacementLiteral) + children = self.parent_executor._get_child_executors(placement) + + # Scenario: there are as many keys as exectutors. For example + # there are 3 clients and each should have a secret key + if len(keys) == len(children): + keys_type_signature = keys[0].type_signature + return FederatingExecutorValue( + await asyncio.gather(*[ + c.create_value(await keys[i].compute(), keys_type_signature) + for (i, c) in enumerate(children) + ]), + computation_types.FederatedType( + keys_type_signature, placement, all_equal=False)) + # Scenario: there are more keys than exectutors. For example + # there are 3 clients and each have a public key. Each client wants + # to share its key to the same aggregator. + elif (len(children)==1) & (len(children) < len(keys)): + keys_type_signature = keys[0].type_signature + child = children[0] + return FederatingExecutorValue( + await asyncio.gather(*[ + child.create_value(await k.compute(), keys_type_signature) + for k in keys + ]), + computation_types.FederatedType( + keys_type_signature, placement, all_equal=False)) + # Scenario: there are more exectutors than keys. For example + # there is an aggregator with one public key. The aggregator + # wants to share the samer public key to 3 different clients. + elif (len(keys) == 1) & (len(children) > len(keys)): + keys_type_signature = keys[0].type_signature + return FederatingExecutorValue( + await asyncio.gather(*[ + c.create_value(await keys[0].compute(), keys_type_signature) + for c in children + ]), + computation_types.FederatedType( + keys_type_signature, placement, all_equal=True)) + async def _encrypt_sender_tensors(self, val, sender_index=0, @@ -1280,8 +1279,11 @@ async def _encrypt_sender_tensors(self, if nb_senders == 1: input_tensor_type = val.type_signature.member + self.orig_sender_tensor_dtype = input_tensor_type.dtype + orig_clients_dtype = arg.type_signature[0].member.dtype else: input_tensor_type = val.type_signature[0].member + self.orig_sender_tensor_dtype = input_tensor_type.dtype pk_receiver = self.key_store.get_keys(self.receiver_placement.name)['pk'] sk_sender = self.key_store.get_keys(self.sender_placement.name)['sk'] @@ -1323,7 +1325,6 @@ def encrypt_tensor(plaintext, pk_rcv, sk_snd): async def _decrypt_tensors_on_receiver(self, val, - sender_dtype, sender_index=0, receive_index=0): @@ -1352,13 +1353,13 @@ def decrypt_tensor(client_outputs, sk_rcv, pk_snd): sk_rcv = easy_box.SecretKey(sk_rcv) pk_snd = easy_box.PublicKey(pk_snd) - plaintext_recovered = easy_box.open_detached(ciphertext, mac, nonce, - pk_snd, sk_rcv, sender_dtype) + plaintext_recovered = easy_box.open_detached( + ciphertext, mac, nonce, pk_snd, sk_rcv, self.orig_sender_tensor_dtype) return plaintext_recovered val_type = computation_types.FederatedType( - computation_types.TensorType(sender_dtype), + computation_types.TensorType(self.orig_sender_tensor_dtype), self.receiver_placement, all_equal=False) From 9c998655912f7bfb256116145a75aefa53555e21 Mon Sep 17 00:00:00 2001 From: yann dupis Date: Tue, 12 May 2020 15:43:04 -0700 Subject: [PATCH 12/31] clean up --- .../impl/executors/federating_executor.py | 22 +++++++++---------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor.py b/tensorflow_federated/python/core/impl/executors/federating_executor.py index d35c5faad4..66bfaf1b59 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor.py @@ -500,7 +500,7 @@ def validate_executor_placements(cls, executor_placements): async def _move(self, arg, target_executor): - enc_clients_arg = await self.channel.send(val=arg) + enc_clients_arg = await self.channel.send(value=arg) val_type = enc_clients_arg.type_signature val = enc_clients_arg.internal_representation @@ -514,7 +514,7 @@ async def _move(self, arg, target_executor): ]) received_vals = await asyncio.gather(*[ - self.channel.receive(val=v, sender_index=i) for (i, v) in enumerate(val) + self.channel.receive(value=v, sender_index=i) for (i, v) in enumerate(val) ]) received_vals = [v.internal_representation[0] for v in received_vals] @@ -1135,19 +1135,19 @@ async def setup(self): return - async def send(self, val, sender_index=0, receiver_index=0): + async def send(self, value, sender_index=0, receiver_index=0): if not self.is_encrypted: - return val + return value - return await self._encrypt_sender_tensors(val, sender_index, receiver_index) + return await self._encrypt_sender_tensors(value, sender_index, receiver_index) - async def receive(self, val, receiver_index=0, sender_index=0): + async def receive(self, value, receiver_index=0, sender_index=0): if not self.is_encrypted: - return val + return value - return await self._decrypt_tensors_on_receiver(val, sender_index, + return await self._decrypt_tensors_on_receiver(value, sender_index, receiver_index) async def _generate_keys(self, key_owner_placement): @@ -1171,13 +1171,13 @@ def generate_keys(): key_generator = await executor.create_call(await executor.create_value( fn, fn_type)) - keys = await asyncio.gather(*[ + pk, sk = await asyncio.gather(*[ executor.create_selection(key_generator, i) for i in range(len(key_generator.type_signature)) ]) - pk_vals.append(keys[0]) - sk_vals.append(keys[1]) + pk_vals.append(pk) + sk_vals.append(sk) self.key_store.add_keys(key_owner_placement.name, pk_vals, sk_vals) From 095519a6f402b724b117f50742c0f3d8c90f0fe1 Mon Sep 17 00:00:00 2001 From: yann dupis Date: Tue, 12 May 2020 18:06:08 -0700 Subject: [PATCH 13/31] introduce channel base class --- .../python/core/impl/executors/BUILD | 7 ++++ .../core/impl/executors/base_channel.py | 24 +++++++++++++ .../impl/executors/federating_executor.py | 34 +++++++++---------- 3 files changed, 47 insertions(+), 18 deletions(-) create mode 100644 tensorflow_federated/python/core/impl/executors/base_channel.py diff --git a/tensorflow_federated/python/core/impl/executors/BUILD b/tensorflow_federated/python/core/impl/executors/BUILD index e34a59ca00..fc19ee0c5e 100644 --- a/tensorflow_federated/python/core/impl/executors/BUILD +++ b/tensorflow_federated/python/core/impl/executors/BUILD @@ -422,6 +422,7 @@ py_library( srcs = ["federating_executor.py"], srcs_version = "PY3", deps = [ + ":base_channel", ":executor_base", ":executor_utils", ":executor_value_base", @@ -439,6 +440,12 @@ py_library( ], ) +py_library( + name = "base_channel", + srcs = ["base_channel.py"], + srcs_version = "PY3", +) + py_test( name = "federating_executor_test", size = "small", diff --git a/tensorflow_federated/python/core/impl/executors/base_channel.py b/tensorflow_federated/python/core/impl/executors/base_channel.py new file mode 100644 index 0000000000..7111e74534 --- /dev/null +++ b/tensorflow_federated/python/core/impl/executors/base_channel.py @@ -0,0 +1,24 @@ +import abc + + +class Channel(metaclass=abc.ABCMeta): + + @abc.abstractmethod + async def send(self, value, sender_index=0, receiver_index=0): + pass + + @abc.abstractmethod + async def receive(self, value, receiver_index=0, sender_index=0): + pass + + @abc.abstractmethod + async def setup(self): + pass + + @abc.abstractmethod + async def _generate_keys(self, key_owner): + pass + + @abc.abstractmethod + async def _share_public_keys(self, key_owner, send_pks_to): + pass diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor.py b/tensorflow_federated/python/core/impl/executors/federating_executor.py index 66bfaf1b59..9ba8584cce 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor.py @@ -36,6 +36,7 @@ from tensorflow_federated.python.core.impl.executors import executor_utils from tensorflow_federated.python.core.impl.executors import executor_value_base +from tensorflow_federated.python.core.impl.executors import base_channel from tf_encrypted.primitives.sodium import easy_box @@ -461,7 +462,7 @@ class TrustedAggregatorIntrinsicStrategy(IntrinsicStrategy): def __init__(self, federating_executor): super().__init__(federating_executor) - self.channel = Channel( + self.channel = EasyBoxChannel( parent_executor=self, sender_placement=placement_literals.CLIENTS, receiver_placement=placement_literals.AGGREGATORS, @@ -514,7 +515,8 @@ async def _move(self, arg, target_executor): ]) received_vals = await asyncio.gather(*[ - self.channel.receive(value=v, sender_index=i) for (i, v) in enumerate(val) + self.channel.receive(value=v, sender_index=i) + for (i, v) in enumerate(val) ]) received_vals = [v.internal_representation[0] for v in received_vals] @@ -1109,7 +1111,7 @@ async def _compute_intrinsic_federated_secure_sum(self, arg): return await self.intrinsic_strategy.federated_secure_sum(arg) -class Channel: +class EasyBoxChannel(base_channel.Channel): def __init__(self, parent_executor, sender_placement, receiver_placement, is_encrypted): @@ -1133,14 +1135,13 @@ async def setup(self): await self._share_public_keys(self.receiver_placement, self.sender_placement) - return - async def send(self, value, sender_index=0, receiver_index=0): if not self.is_encrypted: return value - return await self._encrypt_sender_tensors(value, sender_index, receiver_index) + return await self._encrypt_sender_tensors(value, sender_index, + receiver_index) async def receive(self, value, receiver_index=0, sender_index=0): @@ -1150,7 +1151,7 @@ async def receive(self, value, receiver_index=0, sender_index=0): return await self._decrypt_tensors_on_receiver(value, sender_index, receiver_index) - async def _generate_keys(self, key_owner_placement): + async def _generate_keys(self, key_owner): @computations.tf_computation() def generate_keys(): @@ -1160,7 +1161,7 @@ def generate_keys(): fn_type = generate_keys.type_signature fn = generate_keys._computation_proto - executors = self.parent_executor._get_child_executors(key_owner_placement) + executors = self.parent_executor._get_child_executors(key_owner) nb_executors = len(executors) sk_vals = [] @@ -1179,21 +1180,18 @@ def generate_keys(): pk_vals.append(pk) sk_vals.append(sk) - self.key_store.add_keys(key_owner_placement.name, pk_vals, sk_vals) + self.key_store.add_keys(key_owner.name, pk_vals, sk_vals) return - async def _share_public_keys(self, key_owner_placement, send_public_keys_to): + async def _share_public_keys(self, key_owner, send_pks_to): - keys = self.key_store.get_keys(key_owner_placement.name) + keys = self.key_store.get_keys(key_owner.name) - sk_fed_vals = await self._place_keys( - keys['sk'], key_owner_placement) - pk_fed_vals = await self._place_keys( - keys['pk'], send_public_keys_to) + sk_fed_vals = await self._place_keys(keys['sk'], key_owner) + pk_fed_vals = await self._place_keys(keys['pk'], send_pks_to) - self.key_store.update_keys(key_owner_placement.name, pk_fed_vals, - sk_fed_vals) + self.key_store.update_keys(key_owner.name, pk_fed_vals, sk_fed_vals) async def _zip_val_key(self, vals, @@ -1246,7 +1244,7 @@ async def _place_keys(self, keys, placement): # Scenario: there are more keys than exectutors. For example # there are 3 clients and each have a public key. Each client wants # to share its key to the same aggregator. - elif (len(children)==1) & (len(children) < len(keys)): + elif (len(children) == 1) & (len(children) < len(keys)): keys_type_signature = keys[0].type_signature child = children[0] return FederatingExecutorValue( From f779c6b772263c42ef6be10be3c72a45abe8b072 Mon Sep 17 00:00:00 2001 From: yann dupis Date: Tue, 12 May 2020 18:43:20 -0700 Subject: [PATCH 14/31] encrypt method expects a list of eager val --- .../core/impl/executors/federating_executor.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor.py b/tensorflow_federated/python/core/impl/executors/federating_executor.py index 9ba8584cce..d2e1c72c54 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor.py @@ -501,11 +501,11 @@ def validate_executor_placements(cls, executor_placements): async def _move(self, arg, target_executor): - enc_clients_arg = await self.channel.send(value=arg) + enc_clients_arg = await self.channel.send( + value=arg.internal_representation[0]) val_type = enc_clients_arg.type_signature val = enc_clients_arg.internal_representation - orig_clients_dtype = arg.type_signature[0].member.dtype py_typecheck.check_type(val, list) py_typecheck.check_type(val_type, computation_types.FederatedType) item_type = val_type.member @@ -1276,11 +1276,10 @@ async def _encrypt_sender_tensors(self, self.parent_executor._get_child_executors(self.sender_placement)) if nb_senders == 1: - input_tensor_type = val.type_signature.member + input_tensor_type = val.type_signature self.orig_sender_tensor_dtype = input_tensor_type.dtype - orig_clients_dtype = arg.type_signature[0].member.dtype else: - input_tensor_type = val.type_signature[0].member + input_tensor_type = val[0].type_signature self.orig_sender_tensor_dtype = input_tensor_type.dtype pk_receiver = self.key_store.get_keys(self.receiver_placement.name)['pk'] @@ -1306,8 +1305,10 @@ def encrypt_tensor(plaintext, pk_rcv, sk_snd): val_type = val.type_signature val = val.internal_representation else: - val_type = val.type_signature[0] - val = val.internal_representation[0] + tensor_type = val[0].type_signature + + val_type = computation_types.FederatedType( + tensor_type, self.sender_placement, all_equal=False) val_key_zipped = await self._zip_val_key(val, pk_receiver, sk_sender, self.sender_placement) From a16a4a88e4f1f91e787081d1282df4fbe6635481 Mon Sep 17 00:00:00 2001 From: yann dupis Date: Tue, 12 May 2020 19:12:47 -0700 Subject: [PATCH 15/31] clean up --- .../core/impl/executors/base_channel.py | 4 +- .../impl/executors/federating_executor.py | 190 +++++++++--------- 2 files changed, 102 insertions(+), 92 deletions(-) diff --git a/tensorflow_federated/python/core/impl/executors/base_channel.py b/tensorflow_federated/python/core/impl/executors/base_channel.py index 7111e74534..e2a8e750b0 100644 --- a/tensorflow_federated/python/core/impl/executors/base_channel.py +++ b/tensorflow_federated/python/core/impl/executors/base_channel.py @@ -4,11 +4,11 @@ class Channel(metaclass=abc.ABCMeta): @abc.abstractmethod - async def send(self, value, sender_index=0, receiver_index=0): + async def send(self, value, sender_index=None, receiver_index=None): pass @abc.abstractmethod - async def receive(self, value, receiver_index=0, sender_index=0): + async def receive(self, value, sender_index=None, receiver_index=None): pass @abc.abstractmethod diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor.py b/tensorflow_federated/python/core/impl/executors/federating_executor.py index d2e1c72c54..20997cf684 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor.py @@ -501,11 +501,11 @@ def validate_executor_placements(cls, executor_placements): async def _move(self, arg, target_executor): - enc_clients_arg = await self.channel.send( + enc_clients_vals = await self.channel.send( value=arg.internal_representation[0]) - val_type = enc_clients_arg.type_signature - val = enc_clients_arg.internal_representation + val_type = enc_clients_vals.type_signature + val = enc_clients_vals.internal_representation py_typecheck.check_type(val, list) py_typecheck.check_type(val_type, computation_types.FederatedType) item_type = val_type.member @@ -1135,7 +1135,7 @@ async def setup(self): await self._share_public_keys(self.receiver_placement, self.sender_placement) - async def send(self, value, sender_index=0, receiver_index=0): + async def send(self, value, sender_index=None, receiver_index=None): if not self.is_encrypted: return value @@ -1143,7 +1143,7 @@ async def send(self, value, sender_index=0, receiver_index=0): return await self._encrypt_sender_tensors(value, sender_index, receiver_index) - async def receive(self, value, receiver_index=0, sender_index=0): + async def receive(self, value, receiver_index=None, sender_index=None): if not self.is_encrypted: return value @@ -1193,84 +1193,10 @@ async def _share_public_keys(self, key_owner, send_pks_to): self.key_store.update_keys(key_owner.name, pk_fed_vals, sk_fed_vals) - async def _zip_val_key(self, - vals, - pk_key, - sk_key, - placement, - sender_index=None): - - if isinstance(vals, list): - val_type = computation_types.FederatedType( - vals[0].type_signature, placement, all_equal=False) - else: - val_type = computation_types.FederatedType( - vals.type_signature, placement, all_equal=False) - vals = [vals] - - pk_key_vals = pk_key.internal_representation - sk_key_vals = sk_key.internal_representation - - if sender_index != None: - sk_key_vals = [sk_key_vals[sender_index]] - - vals_key = FederatingExecutorValue( - anonymous_tuple.AnonymousTuple([(None, vals), (None, pk_key_vals), - (None, sk_key_vals)]), - computation_types.NamedTupleType( - (val_type, pk_key.type_signature, sk_key.type_signature))) - - vals_key_zipped = await self.parent_executor._zip( - vals_key, placement, all_equal=False) - - return vals_key_zipped.internal_representation - - async def _place_keys(self, keys, placement): - - py_typecheck.check_type(placement, placement_literals.PlacementLiteral) - children = self.parent_executor._get_child_executors(placement) - - # Scenario: there are as many keys as exectutors. For example - # there are 3 clients and each should have a secret key - if len(keys) == len(children): - keys_type_signature = keys[0].type_signature - return FederatingExecutorValue( - await asyncio.gather(*[ - c.create_value(await keys[i].compute(), keys_type_signature) - for (i, c) in enumerate(children) - ]), - computation_types.FederatedType( - keys_type_signature, placement, all_equal=False)) - # Scenario: there are more keys than exectutors. For example - # there are 3 clients and each have a public key. Each client wants - # to share its key to the same aggregator. - elif (len(children) == 1) & (len(children) < len(keys)): - keys_type_signature = keys[0].type_signature - child = children[0] - return FederatingExecutorValue( - await asyncio.gather(*[ - child.create_value(await k.compute(), keys_type_signature) - for k in keys - ]), - computation_types.FederatedType( - keys_type_signature, placement, all_equal=False)) - # Scenario: there are more exectutors than keys. For example - # there is an aggregator with one public key. The aggregator - # wants to share the samer public key to 3 different clients. - elif (len(keys) == 1) & (len(children) > len(keys)): - keys_type_signature = keys[0].type_signature - return FederatingExecutorValue( - await asyncio.gather(*[ - c.create_value(await keys[0].compute(), keys_type_signature) - for c in children - ]), - computation_types.FederatedType( - keys_type_signature, placement, all_equal=True)) - async def _encrypt_sender_tensors(self, val, - sender_index=0, - receiver_index=0): + sender_index=None, + receiver_index=None): nb_senders = len( self.parent_executor._get_child_executors(self.sender_placement)) @@ -1310,8 +1236,13 @@ def encrypt_tensor(plaintext, pk_rcv, sk_snd): val_type = computation_types.FederatedType( tensor_type, self.sender_placement, all_equal=False) - val_key_zipped = await self._zip_val_key(val, pk_receiver, sk_sender, - self.sender_placement) + val_key_zipped = await self._zip_val_key( + self.sender_placement, + val, + pk_receiver, + sk_sender, + pk_index=receiver_index, + sk_index=sender_index) # NOTE probably won't always be fed_ex in future design fed_ex = self.parent_executor.federating_executor @@ -1325,26 +1256,27 @@ def encrypt_tensor(plaintext, pk_rcv, sk_snd): async def _decrypt_tensors_on_receiver(self, val, sender_index=0, - receive_index=0): + receiver_index=0): pk_sender = self.key_store.get_keys(self.sender_placement.name)['pk'] sk_receiver = self.key_store.get_keys(self.receiver_placement.name)['sk'] val = await self._zip_val_key( + self.receiver_placement, val, - sk_receiver, pk_sender, - self.receiver_placement, - sender_index=sender_index) + sk_receiver, + pk_index=sender_index, + sk_index=receiver_index) sender_output_type_signature = val[0].type_signature[0] receiver_secret_key_type_signature = val[0].type_signature[1] sender_public_key_type_signature = val[0].type_signature[2] @computations.tf_computation(sender_output_type_signature, - receiver_secret_key_type_signature, - sender_public_key_type_signature) - def decrypt_tensor(client_outputs, sk_rcv, pk_snd): + sender_public_key_type_signature, + receiver_secret_key_type_signature) + def decrypt_tensor(client_outputs, pk_snd, sk_rcv): ciphertext = easy_box.Ciphertext(client_outputs[0]) mac = easy_box.Mac(client_outputs[1]) @@ -1372,6 +1304,84 @@ def decrypt_tensor(client_outputs, sk_rcv, pk_snd): anonymous_tuple.AnonymousTuple([(None, fn), (None, val)]), computation_types.NamedTupleType((fn_type, val_type)))) + async def _zip_val_key(self, + placement, + vals, + pk_key, + sk_key, + pk_index=None, + sk_index=None): + + if isinstance(vals, list): + val_type = computation_types.FederatedType( + vals[0].type_signature, placement, all_equal=False) + else: + val_type = computation_types.FederatedType( + vals.type_signature, placement, all_equal=False) + vals = [vals] + + pk_key_vals = pk_key.internal_representation + sk_key_vals = sk_key.internal_representation + + if pk_index != None: + pk_key_vals = [pk_key_vals[pk_index]] + + if sk_index != None: + sk_key_vals = [sk_key_vals[sk_index]] + + vals_key = FederatingExecutorValue( + anonymous_tuple.AnonymousTuple([(None, vals), (None, pk_key_vals), + (None, sk_key_vals)]), + computation_types.NamedTupleType( + (val_type, pk_key.type_signature, sk_key.type_signature))) + + vals_key_zipped = await self.parent_executor._zip( + vals_key, placement, all_equal=False) + + return vals_key_zipped.internal_representation + + async def _place_keys(self, keys, placement): + + py_typecheck.check_type(placement, placement_literals.PlacementLiteral) + children = self.parent_executor._get_child_executors(placement) + + # Scenario: there are as many keys as exectutors. For example + # there are 3 clients and each should have a secret key + if len(keys) == len(children): + keys_type_signature = keys[0].type_signature + return FederatingExecutorValue( + await asyncio.gather(*[ + c.create_value(await keys[i].compute(), keys_type_signature) + for (i, c) in enumerate(children) + ]), + computation_types.FederatedType( + keys_type_signature, placement, all_equal=False)) + # Scenario: there are more keys than exectutors. For example + # there are 3 clients and each have a public key. Each client wants + # to share its key to the same aggregator. + elif (len(children) == 1) & (len(children) < len(keys)): + keys_type_signature = keys[0].type_signature + child = children[0] + return FederatingExecutorValue( + await asyncio.gather(*[ + child.create_value(await k.compute(), keys_type_signature) + for k in keys + ]), + computation_types.FederatedType( + keys_type_signature, placement, all_equal=False)) + # Scenario: there are more exectutors than keys. For example + # there is an aggregator with one public key. The aggregator + # wants to share the samer public key to 3 different clients. + elif (len(keys) == 1) & (len(children) > len(keys)): + keys_type_signature = keys[0].type_signature + return FederatingExecutorValue( + await asyncio.gather(*[ + c.create_value(await keys[0].compute(), keys_type_signature) + for c in children + ]), + computation_types.FederatedType( + keys_type_signature, placement, all_equal=True)) + class KeyStore: From 9c16c9f8e90375a66fa55e312a3791405c5f5342 Mon Sep 17 00:00:00 2001 From: yann dupis Date: Wed, 13 May 2020 15:55:51 -0700 Subject: [PATCH 16/31] rename file to channel_base --- .../core/impl/executors/{base_channel.py => channel_base.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tensorflow_federated/python/core/impl/executors/{base_channel.py => channel_base.py} (100%) diff --git a/tensorflow_federated/python/core/impl/executors/base_channel.py b/tensorflow_federated/python/core/impl/executors/channel_base.py similarity index 100% rename from tensorflow_federated/python/core/impl/executors/base_channel.py rename to tensorflow_federated/python/core/impl/executors/channel_base.py From 78a59e06240da1034557f3b6ccc76c3204d58d89 Mon Sep 17 00:00:00 2001 From: yann dupis Date: Wed, 13 May 2020 15:57:40 -0700 Subject: [PATCH 17/31] remove key generation and exchange on base class --- .../python/core/impl/executors/channel_base.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/tensorflow_federated/python/core/impl/executors/channel_base.py b/tensorflow_federated/python/core/impl/executors/channel_base.py index e2a8e750b0..4ffec03781 100644 --- a/tensorflow_federated/python/core/impl/executors/channel_base.py +++ b/tensorflow_federated/python/core/impl/executors/channel_base.py @@ -15,10 +15,3 @@ async def receive(self, value, sender_index=None, receiver_index=None): async def setup(self): pass - @abc.abstractmethod - async def _generate_keys(self, key_owner): - pass - - @abc.abstractmethod - async def _share_public_keys(self, key_owner, send_pks_to): - pass From 4236f3ce2f3dce49ed4724362a8ffd5bddc26e3d Mon Sep 17 00:00:00 2001 From: yann dupis Date: Wed, 13 May 2020 16:27:57 -0700 Subject: [PATCH 18/31] method name cleanup --- .../python/core/impl/executors/BUILD | 6 ++-- .../impl/executors/federating_executor.py | 28 ++++++------------- 2 files changed, 11 insertions(+), 23 deletions(-) diff --git a/tensorflow_federated/python/core/impl/executors/BUILD b/tensorflow_federated/python/core/impl/executors/BUILD index fc19ee0c5e..53986858e4 100644 --- a/tensorflow_federated/python/core/impl/executors/BUILD +++ b/tensorflow_federated/python/core/impl/executors/BUILD @@ -422,7 +422,7 @@ py_library( srcs = ["federating_executor.py"], srcs_version = "PY3", deps = [ - ":base_channel", + ":channel_base", ":executor_base", ":executor_utils", ":executor_value_base", @@ -441,8 +441,8 @@ py_library( ) py_library( - name = "base_channel", - srcs = ["base_channel.py"], + name = "channel_base", + srcs = ["channel_base.py"], srcs_version = "PY3", ) diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor.py b/tensorflow_federated/python/core/impl/executors/federating_executor.py index 20997cf684..c769235233 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor.py @@ -36,7 +36,7 @@ from tensorflow_federated.python.core.impl.executors import executor_utils from tensorflow_federated.python.core.impl.executors import executor_value_base -from tensorflow_federated.python.core.impl.executors import base_channel +from tensorflow_federated.python.core.impl.executors import channel_base from tf_encrypted.primitives.sodium import easy_box @@ -465,8 +465,7 @@ def __init__(self, federating_executor): self.channel = EasyBoxChannel( parent_executor=self, sender_placement=placement_literals.CLIENTS, - receiver_placement=placement_literals.AGGREGATORS, - is_encrypted=True) + receiver_placement=placement_literals.AGGREGATORS) @classmethod def validate_executor_placements(cls, executor_placements): @@ -1111,23 +1110,18 @@ async def _compute_intrinsic_federated_secure_sum(self, arg): return await self.intrinsic_strategy.federated_secure_sum(arg) -class EasyBoxChannel(base_channel.Channel): +class EasyBoxChannel(channel_base.Channel): - def __init__(self, parent_executor, sender_placement, receiver_placement, - is_encrypted): + def __init__(self, parent_executor, sender_placement, receiver_placement): self.parent_executor = parent_executor self.sender_placement = sender_placement self.receiver_placement = receiver_placement - self.is_encrypted = is_encrypted self.key_store = KeyStore() async def setup(self): - if not self.is_encrypted: - return - await self._generate_keys(self.sender_placement) await self._generate_keys(self.receiver_placement) await self._share_public_keys(self.sender_placement, @@ -1137,18 +1131,12 @@ async def setup(self): async def send(self, value, sender_index=None, receiver_index=None): - if not self.is_encrypted: - return value - - return await self._encrypt_sender_tensors(value, sender_index, + return await self._encrypt_values_on_sender(value, sender_index, receiver_index) async def receive(self, value, receiver_index=None, sender_index=None): - if not self.is_encrypted: - return value - - return await self._decrypt_tensors_on_receiver(value, sender_index, + return await self._decrypt_values_on_receiver(value, sender_index, receiver_index) async def _generate_keys(self, key_owner): @@ -1193,7 +1181,7 @@ async def _share_public_keys(self, key_owner, send_pks_to): self.key_store.update_keys(key_owner.name, pk_fed_vals, sk_fed_vals) - async def _encrypt_sender_tensors(self, + async def _encrypt_values_on_sender(self, val, sender_index=None, receiver_index=None): @@ -1253,7 +1241,7 @@ def encrypt_tensor(plaintext, pk_rcv, sk_snd): (None, val_key_zipped)]), computation_types.NamedTupleType((fn_type, val_type)))) - async def _decrypt_tensors_on_receiver(self, + async def _decrypt_values_on_receiver(self, val, sender_index=0, receiver_index=0): From 813fbd70605663facda4756c7d429c94f7bc2557 Mon Sep 17 00:00:00 2001 From: yann dupis Date: Wed, 13 May 2020 17:14:12 -0700 Subject: [PATCH 19/31] check if the channel has already been setup --- .../impl/executors/federating_executor.py | 37 ++++++++++--------- 1 file changed, 20 insertions(+), 17 deletions(-) diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor.py b/tensorflow_federated/python/core/impl/executors/federating_executor.py index c769235233..a8286741fa 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor.py @@ -500,6 +500,8 @@ def validate_executor_placements(cls, executor_placements): async def _move(self, arg, target_executor): + await self.channel.setup() + enc_clients_vals = await self.channel.send( value=arg.internal_representation[0]) @@ -571,9 +573,6 @@ async def federated_reduce(self, arg): 'Expected 3 elements in the `federated_reduce()` argument tuple, ' 'found {}.'.format(len(arg.internal_representation))) - # Setup channel keys - await self.channel.setup() - aggr = self._get_child_executors(placement_literals.AGGREGATORS, index=0) aggregands = await self._move(arg, aggr) @@ -1119,25 +1118,29 @@ def __init__(self, parent_executor, sender_placement, receiver_placement): self.receiver_placement = receiver_placement self.key_store = KeyStore() + self._is_channel_setup = False async def setup(self): - await self._generate_keys(self.sender_placement) - await self._generate_keys(self.receiver_placement) - await self._share_public_keys(self.sender_placement, - self.receiver_placement) - await self._share_public_keys(self.receiver_placement, - self.sender_placement) + if not self._is_channel_setup: + await self._generate_keys(self.sender_placement) + await self._generate_keys(self.receiver_placement) + await self._share_public_keys(self.sender_placement, + self.receiver_placement) + await self._share_public_keys(self.receiver_placement, + self.sender_placement) + + self._is_channel_setup = True async def send(self, value, sender_index=None, receiver_index=None): return await self._encrypt_values_on_sender(value, sender_index, - receiver_index) + receiver_index) async def receive(self, value, receiver_index=None, sender_index=None): return await self._decrypt_values_on_receiver(value, sender_index, - receiver_index) + receiver_index) async def _generate_keys(self, key_owner): @@ -1182,9 +1185,9 @@ async def _share_public_keys(self, key_owner, send_pks_to): self.key_store.update_keys(key_owner.name, pk_fed_vals, sk_fed_vals) async def _encrypt_values_on_sender(self, - val, - sender_index=None, - receiver_index=None): + val, + sender_index=None, + receiver_index=None): nb_senders = len( self.parent_executor._get_child_executors(self.sender_placement)) @@ -1242,9 +1245,9 @@ def encrypt_tensor(plaintext, pk_rcv, sk_snd): computation_types.NamedTupleType((fn_type, val_type)))) async def _decrypt_values_on_receiver(self, - val, - sender_index=0, - receiver_index=0): + val, + sender_index=0, + receiver_index=0): pk_sender = self.key_store.get_keys(self.sender_placement.name)['pk'] sk_receiver = self.key_store.get_keys(self.receiver_placement.name)['sk'] From a913bde2537e36c35e661c520300e9610b3d5758 Mon Sep 17 00:00:00 2001 From: yann dupis Date: Wed, 13 May 2020 17:25:27 -0700 Subject: [PATCH 20/31] rename sender_index & receiver_index --- .../core/impl/executors/channel_base.py | 5 ++- .../impl/executors/federating_executor.py | 34 +++++++------------ 2 files changed, 14 insertions(+), 25 deletions(-) diff --git a/tensorflow_federated/python/core/impl/executors/channel_base.py b/tensorflow_federated/python/core/impl/executors/channel_base.py index 4ffec03781..ebc60e4ae7 100644 --- a/tensorflow_federated/python/core/impl/executors/channel_base.py +++ b/tensorflow_federated/python/core/impl/executors/channel_base.py @@ -4,14 +4,13 @@ class Channel(metaclass=abc.ABCMeta): @abc.abstractmethod - async def send(self, value, sender_index=None, receiver_index=None): + async def send(self, value, sender=None, receiver=None): pass @abc.abstractmethod - async def receive(self, value, sender_index=None, receiver_index=None): + async def receive(self, value, sender=None, receiver=None): pass @abc.abstractmethod async def setup(self): pass - diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor.py b/tensorflow_federated/python/core/impl/executors/federating_executor.py index a8286741fa..adf5d474b2 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor.py @@ -515,10 +515,8 @@ async def _move(self, arg, target_executor): target_executor.create_value(await v.compute(), item_type) for v in val ]) - received_vals = await asyncio.gather(*[ - self.channel.receive(value=v, sender_index=i) - for (i, v) in enumerate(val) - ]) + received_vals = await asyncio.gather( + *[self.channel.receive(value=v, sender=i) for (i, v) in enumerate(val)]) received_vals = [v.internal_representation[0] for v in received_vals] @@ -1132,15 +1130,13 @@ async def setup(self): self._is_channel_setup = True - async def send(self, value, sender_index=None, receiver_index=None): + async def send(self, value, sender=None, receiver=None): - return await self._encrypt_values_on_sender(value, sender_index, - receiver_index) + return await self._encrypt_values_on_sender(value, sender, receiver) - async def receive(self, value, receiver_index=None, sender_index=None): + async def receive(self, value, receiver=None, sender=None): - return await self._decrypt_values_on_receiver(value, sender_index, - receiver_index) + return await self._decrypt_values_on_receiver(value, sender, receiver) async def _generate_keys(self, key_owner): @@ -1184,10 +1180,7 @@ async def _share_public_keys(self, key_owner, send_pks_to): self.key_store.update_keys(key_owner.name, pk_fed_vals, sk_fed_vals) - async def _encrypt_values_on_sender(self, - val, - sender_index=None, - receiver_index=None): + async def _encrypt_values_on_sender(self, val, sender=None, receiver=None): nb_senders = len( self.parent_executor._get_child_executors(self.sender_placement)) @@ -1232,8 +1225,8 @@ def encrypt_tensor(plaintext, pk_rcv, sk_snd): val, pk_receiver, sk_sender, - pk_index=receiver_index, - sk_index=sender_index) + pk_index=receiver, + sk_index=sender) # NOTE probably won't always be fed_ex in future design fed_ex = self.parent_executor.federating_executor @@ -1244,10 +1237,7 @@ def encrypt_tensor(plaintext, pk_rcv, sk_snd): (None, val_key_zipped)]), computation_types.NamedTupleType((fn_type, val_type)))) - async def _decrypt_values_on_receiver(self, - val, - sender_index=0, - receiver_index=0): + async def _decrypt_values_on_receiver(self, val, sender=0, receiver=0): pk_sender = self.key_store.get_keys(self.sender_placement.name)['pk'] sk_receiver = self.key_store.get_keys(self.receiver_placement.name)['sk'] @@ -1257,8 +1247,8 @@ async def _decrypt_values_on_receiver(self, val, pk_sender, sk_receiver, - pk_index=sender_index, - sk_index=receiver_index) + pk_index=sender, + sk_index=receiver) sender_output_type_signature = val[0].type_signature[0] receiver_secret_key_type_signature = val[0].type_signature[1] From a8c45040e0da90bdde04c2acd5c940e02f1838b4 Mon Sep 17 00:00:00 2001 From: yann dupis Date: Wed, 13 May 2020 18:36:58 -0700 Subject: [PATCH 21/31] store channels in ChannelGrid --- .../core/impl/executors/channel_base.py | 19 +++++++++++++ .../impl/executors/federating_executor.py | 27 ++++++++++++------- 2 files changed, 36 insertions(+), 10 deletions(-) diff --git a/tensorflow_federated/python/core/impl/executors/channel_base.py b/tensorflow_federated/python/core/impl/executors/channel_base.py index ebc60e4ae7..d4a3abac22 100644 --- a/tensorflow_federated/python/core/impl/executors/channel_base.py +++ b/tensorflow_federated/python/core/impl/executors/channel_base.py @@ -1,4 +1,12 @@ import abc +from dataclasses import dataclass +from typing import Tuple, Dict + +from tensorflow_federated.python.common_libs import py_typecheck +from tensorflow_federated.python.core.impl.compiler import placement_literals + +PlacementPair = Tuple[placement_literals.PlacementLiteral, + placement_literals.PlacementLiteral] class Channel(metaclass=abc.ABCMeta): @@ -14,3 +22,14 @@ async def receive(self, value, sender=None, receiver=None): @abc.abstractmethod async def setup(self): pass + + +@dataclass +class ChannelGrid: + channel_dict: Dict[PlacementPair, Channel] + + def __getitem__(self, placements: PlacementPair): + py_typecheck.check_type(placements, tuple) + py_typecheck.check_len(placements, 2) + sorted_placements = sorted(placements, key=lambda p: p.uri) + return self.channel_dict.get(tuple(sorted_placements)) diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor.py b/tensorflow_federated/python/core/impl/executors/federating_executor.py index adf5d474b2..7b84084e97 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor.py @@ -462,10 +462,13 @@ class TrustedAggregatorIntrinsicStrategy(IntrinsicStrategy): def __init__(self, federating_executor): super().__init__(federating_executor) - self.channel = EasyBoxChannel( - parent_executor=self, - sender_placement=placement_literals.CLIENTS, - receiver_placement=placement_literals.AGGREGATORS) + self.channel_grid = channel_base.ChannelGrid({ + (placement_literals.AGGREGATORS, placement_literals.CLIENTS): + EasyBoxChannel( + parent_executor=self, + sender_placement=placement_literals.CLIENTS, + receiver_placement=placement_literals.AGGREGATORS) + }) @classmethod def validate_executor_placements(cls, executor_placements): @@ -498,12 +501,15 @@ def validate_executor_placements(cls, executor_placements): 'Unsupported cardinality for placement "{}": {}.'.format( pl, pl_cardinality)) - async def _move(self, arg, target_executor): + async def _move(self, arg, source_placement, target_placement): - await self.channel.setup() + target_executor = self._get_child_executors(target_placement, index=0) - enc_clients_vals = await self.channel.send( - value=arg.internal_representation[0]) + channel = self.channel_grid[(source_placement, target_placement)] + + await channel.setup() + + enc_clients_vals = await channel.send(value=arg.internal_representation[0]) val_type = enc_clients_vals.type_signature val = enc_clients_vals.internal_representation @@ -516,7 +522,7 @@ async def _move(self, arg, target_executor): ]) received_vals = await asyncio.gather( - *[self.channel.receive(value=v, sender=i) for (i, v) in enumerate(val)]) + *[channel.receive(value=v, sender=i) for (i, v) in enumerate(val)]) received_vals = [v.internal_representation[0] for v in received_vals] @@ -572,7 +578,8 @@ async def federated_reduce(self, arg): 'found {}.'.format(len(arg.internal_representation))) aggr = self._get_child_executors(placement_literals.AGGREGATORS, index=0) - aggregands = await self._move(arg, aggr) + aggregands = await self._move(arg, placement_literals.CLIENTS, + placement_literals.AGGREGATORS) zero_type = arg.type_signature[1] op_type = arg.type_signature[2] From 30202bdb450814dd795bf63bbf29aca69de95501 Mon Sep 17 00:00:00 2001 From: yann dupis Date: Wed, 13 May 2020 19:09:25 -0700 Subject: [PATCH 22/31] fix indentation --- .../python/core/impl/executors/federating_executor.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor.py b/tensorflow_federated/python/core/impl/executors/federating_executor.py index 7b84084e97..3dfacf9066 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor.py @@ -576,7 +576,7 @@ async def federated_reduce(self, arg): raise ValueError( 'Expected 3 elements in the `federated_reduce()` argument tuple, ' 'found {}.'.format(len(arg.internal_representation))) - + aggr = self._get_child_executors(placement_literals.AGGREGATORS, index=0) aggregands = await self._move(arg, placement_literals.CLIENTS, placement_literals.AGGREGATORS) @@ -1174,7 +1174,7 @@ def generate_keys(): pk_vals.append(pk) sk_vals.append(sk) - self.key_store.add_keys(key_owner.name, pk_vals, sk_vals) + self.key_store.add_keys(key_owner.name, pk_vals, sk_vals) return From 38fdef819296915c2d183e385cb1c1b26cfbdade Mon Sep 17 00:00:00 2001 From: yann dupis Date: Wed, 13 May 2020 19:23:11 -0700 Subject: [PATCH 23/31] move _place_keys for sk to _generate_keys --- .../core/impl/executors/federating_executor.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor.py b/tensorflow_federated/python/core/impl/executors/federating_executor.py index 3dfacf9066..3dc6254a77 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor.py @@ -1174,7 +1174,11 @@ def generate_keys(): pk_vals.append(pk) sk_vals.append(sk) - self.key_store.add_keys(key_owner.name, pk_vals, sk_vals) + # Store list of EagerValue created by executor.create_call + # in a FederatingExecutorValue with the key onwer placement + sk_fed_vals = await self._place_keys(sk_vals, key_owner) + + self.key_store.add_keys(key_owner.name, pk_vals, sk_fed_vals) return @@ -1182,10 +1186,9 @@ async def _share_public_keys(self, key_owner, send_pks_to): keys = self.key_store.get_keys(key_owner.name) - sk_fed_vals = await self._place_keys(keys['sk'], key_owner) pk_fed_vals = await self._place_keys(keys['pk'], send_pks_to) - self.key_store.update_keys(key_owner.name, pk_fed_vals, sk_fed_vals) + self.key_store.update_keys(key_owner.name, pk_fed_vals) async def _encrypt_values_on_sender(self, val, sender=None, receiver=None): @@ -1384,6 +1387,8 @@ def add_keys(self, key_owner, pk, sk): def get_keys(self, key_owner): return self.key_store[key_owner] - def update_keys(self, key_owner, pk, sk): - self.key_store[key_owner]['pk'] = pk - self.key_store[key_owner]['sk'] = sk + def update_keys(self, key_owner, pk=None, sk=None): + if pk: + self.key_store[key_owner]['pk'] = pk + if sk: + self.key_store[key_owner]['sk'] = sk From c6e98be85a10f2dfe97801bfea07ca7ddb65cd28 Mon Sep 17 00:00:00 2001 From: yann dupis Date: Wed, 13 May 2020 19:30:03 -0700 Subject: [PATCH 24/31] fix bug --- .../python/core/impl/executors/federating_executor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor.py b/tensorflow_federated/python/core/impl/executors/federating_executor.py index 3dc6254a77..dab275981e 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor.py @@ -1162,7 +1162,7 @@ def generate_keys(): pk_vals = [] for i in range(nb_executors): - executor = executors[0] + executor = executors[i] key_generator = await executor.create_call(await executor.create_value( fn, fn_type)) From 930fad99bfc44e31acd1ea62e91cfbc260e713ff Mon Sep 17 00:00:00 2001 From: yann dupis Date: Wed, 13 May 2020 19:52:04 -0700 Subject: [PATCH 25/31] _decrypt_values_on_receiver return list of EagerValues --- .../core/impl/executors/federating_executor.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor.py b/tensorflow_federated/python/core/impl/executors/federating_executor.py index dab275981e..2309d4ced9 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor.py @@ -504,7 +504,6 @@ def validate_executor_placements(cls, executor_placements): async def _move(self, arg, source_placement, target_placement): target_executor = self._get_child_executors(target_placement, index=0) - channel = self.channel_grid[(source_placement, target_placement)] await channel.setup() @@ -524,8 +523,6 @@ async def _move(self, arg, source_placement, target_placement): received_vals = await asyncio.gather( *[channel.receive(value=v, sender=i) for (i, v) in enumerate(val)]) - received_vals = [v.internal_representation[0] for v in received_vals] - return received_vals async def federated_value_at_server(self, arg): @@ -1247,7 +1244,7 @@ def encrypt_tensor(plaintext, pk_rcv, sk_snd): (None, val_key_zipped)]), computation_types.NamedTupleType((fn_type, val_type)))) - async def _decrypt_values_on_receiver(self, val, sender=0, receiver=0): + async def _decrypt_values_on_receiver(self, val, sender=None, receiver=None): pk_sender = self.key_store.get_keys(self.sender_placement.name)['pk'] sk_receiver = self.key_store.get_keys(self.receiver_placement.name)['sk'] @@ -1290,11 +1287,16 @@ def decrypt_tensor(client_outputs, pk_snd, sk_rcv): # NOTE probably won't always be fed_ex in future design fed_ex = self.parent_executor.federating_executor - return await fed_ex._compute_intrinsic_federated_map( + val_decrypted = await fed_ex._compute_intrinsic_federated_map( FederatingExecutorValue( anonymous_tuple.AnonymousTuple([(None, fn), (None, val)]), computation_types.NamedTupleType((fn_type, val_type)))) + if sender!=None or receiver!=None: + return val_decrypted.internal_representation[0] + else: + return val_decrypted.internal_representation + async def _zip_val_key(self, placement, vals, From 13554084265baef8b896e8a5d6d3326b28e7f55f Mon Sep 17 00:00:00 2001 From: yann dupis Date: Wed, 13 May 2020 20:08:10 -0700 Subject: [PATCH 26/31] clean up --- .../core/impl/executors/federating_executor.py | 15 ++++----------- 1 file changed, 4 insertions(+), 11 deletions(-) diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor.py b/tensorflow_federated/python/core/impl/executors/federating_executor.py index 2309d4ced9..74e0d1682c 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor.py @@ -1158,15 +1158,12 @@ def generate_keys(): sk_vals = [] pk_vals = [] - for i in range(nb_executors): - executor = executors[i] + for executor in executors: key_generator = await executor.create_call(await executor.create_value( fn, fn_type)) - pk, sk = await asyncio.gather(*[ - executor.create_selection(key_generator, i) - for i in range(len(key_generator.type_signature)) - ]) + pk = await executor.create_selection(key_generator, 0) + sk = await executor.create_selection(key_generator, 1) pk_vals.append(pk) sk_vals.append(sk) @@ -1177,8 +1174,6 @@ def generate_keys(): self.key_store.add_keys(key_owner.name, pk_vals, sk_fed_vals) - return - async def _share_public_keys(self, key_owner, send_pks_to): keys = self.key_store.get_keys(key_owner.name) @@ -1382,9 +1377,7 @@ def __init__(self): self.key_store = {} def add_keys(self, key_owner, pk, sk): - self.key_store[key_owner] = {} - self.key_store[key_owner]['pk'] = pk - self.key_store[key_owner]['sk'] = sk + self.key_store[key_owner] = {'pk': pk, 'sk': sk} def get_keys(self, key_owner): return self.key_store[key_owner] From d8a4106ebea70d6d1fe70ee564acd3c2fc102c8f Mon Sep 17 00:00:00 2001 From: yann dupis Date: Wed, 13 May 2020 21:38:48 -0700 Subject: [PATCH 27/31] clean up --- .../impl/executors/federating_executor.py | 47 +++++++++++-------- 1 file changed, 28 insertions(+), 19 deletions(-) diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor.py b/tensorflow_federated/python/core/impl/executors/federating_executor.py index 74e0d1682c..9fcab9f53a 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor.py @@ -573,7 +573,7 @@ async def federated_reduce(self, arg): raise ValueError( 'Expected 3 elements in the `federated_reduce()` argument tuple, ' 'found {}.'.format(len(arg.internal_representation))) - + aggr = self._get_child_executors(placement_literals.AGGREGATORS, index=0) aggregands = await self._move(arg, placement_literals.CLIENTS, placement_literals.AGGREGATORS) @@ -1119,18 +1119,22 @@ def __init__(self, parent_executor, sender_placement, receiver_placement): self.sender_placement = sender_placement self.receiver_placement = receiver_placement - self.key_store = KeyStore() + self.key_references = KeyStore() self._is_channel_setup = False async def setup(self): if not self._is_channel_setup: - await self._generate_keys(self.sender_placement) - await self._generate_keys(self.receiver_placement) - await self._share_public_keys(self.sender_placement, - self.receiver_placement) - await self._share_public_keys(self.receiver_placement, - self.sender_placement) + await asyncio.gather(*[ + self._generate_keys(self.sender_placement), + self._generate_keys(self.receiver_placement) + ]) + await asyncio.gather(*[ + self._share_public_keys(self.sender_placement, + self.receiver_placement), + self._share_public_keys(self.receiver_placement, + self.sender_placement) + ]) self._is_channel_setup = True @@ -1172,15 +1176,15 @@ def generate_keys(): # in a FederatingExecutorValue with the key onwer placement sk_fed_vals = await self._place_keys(sk_vals, key_owner) - self.key_store.add_keys(key_owner.name, pk_vals, sk_fed_vals) + self.key_references.add_keys(key_owner.name, pk_vals, sk_fed_vals) async def _share_public_keys(self, key_owner, send_pks_to): - keys = self.key_store.get_keys(key_owner.name) + pk = self.key_references.get_public_key(key_owner.name) - pk_fed_vals = await self._place_keys(keys['pk'], send_pks_to) + pk_fed_vals = await self._place_keys(pk, send_pks_to) - self.key_store.update_keys(key_owner.name, pk_fed_vals) + self.key_references.update_keys(key_owner.name, pk_fed_vals) async def _encrypt_values_on_sender(self, val, sender=None, receiver=None): @@ -1194,8 +1198,9 @@ async def _encrypt_values_on_sender(self, val, sender=None, receiver=None): input_tensor_type = val[0].type_signature self.orig_sender_tensor_dtype = input_tensor_type.dtype - pk_receiver = self.key_store.get_keys(self.receiver_placement.name)['pk'] - sk_sender = self.key_store.get_keys(self.sender_placement.name)['sk'] + pk_receiver = self.key_references.get_public_key( + self.receiver_placement.name) + sk_sender = self.key_references.get_secret_key(self.sender_placement.name) pk_rcv_tensor_type = pk_receiver.type_signature.member sk_sender_tensor_type = sk_sender.type_signature.member @@ -1241,8 +1246,9 @@ def encrypt_tensor(plaintext, pk_rcv, sk_snd): async def _decrypt_values_on_receiver(self, val, sender=None, receiver=None): - pk_sender = self.key_store.get_keys(self.sender_placement.name)['pk'] - sk_receiver = self.key_store.get_keys(self.receiver_placement.name)['sk'] + pk_sender = self.key_references.get_public_key(self.sender_placement.name) + sk_receiver = self.key_references.get_secret_key( + self.receiver_placement.name) val = await self._zip_val_key( self.receiver_placement, @@ -1287,7 +1293,7 @@ def decrypt_tensor(client_outputs, pk_snd, sk_rcv): anonymous_tuple.AnonymousTuple([(None, fn), (None, val)]), computation_types.NamedTupleType((fn_type, val_type)))) - if sender!=None or receiver!=None: + if sender != None or receiver != None: return val_decrypted.internal_representation[0] else: return val_decrypted.internal_representation @@ -1379,8 +1385,11 @@ def __init__(self): def add_keys(self, key_owner, pk, sk): self.key_store[key_owner] = {'pk': pk, 'sk': sk} - def get_keys(self, key_owner): - return self.key_store[key_owner] + def get_public_key(self, key_owner): + return self.key_store[key_owner]['pk'] + + def get_secret_key(self, key_owner): + return self.key_store[key_owner]['sk'] def update_keys(self, key_owner, pk=None, sk=None): if pk: From bce95968ad22c3f4a9b463a4d9b1d85debd9ec6d Mon Sep 17 00:00:00 2001 From: yann dupis Date: Thu, 14 May 2020 15:14:27 -0700 Subject: [PATCH 28/31] input/output constantly internal_representation of fed val --- .../impl/executors/federating_executor.py | 21 ++++++++++++------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor.py b/tensorflow_federated/python/core/impl/executors/federating_executor.py index 9fcab9f53a..17d2f9c402 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor.py @@ -503,18 +503,18 @@ def validate_executor_placements(cls, executor_placements): async def _move(self, arg, source_placement, target_placement): + val_type = arg.type_signature[0] + val = arg.internal_representation[0] + py_typecheck.check_type(val, list) + py_typecheck.check_type(val_type, computation_types.FederatedType) + target_executor = self._get_child_executors(target_placement, index=0) channel = self.channel_grid[(source_placement, target_placement)] await channel.setup() - enc_clients_vals = await channel.send(value=arg.internal_representation[0]) - - val_type = enc_clients_vals.type_signature - val = enc_clients_vals.internal_representation - py_typecheck.check_type(val, list) - py_typecheck.check_type(val_type, computation_types.FederatedType) - item_type = val_type.member + val = await channel.send(value=val) + item_type = val[0].type_signature val = await asyncio.gather(*[ target_executor.create_value(await v.compute(), item_type) for v in val @@ -1238,12 +1238,17 @@ def encrypt_tensor(plaintext, pk_rcv, sk_snd): # NOTE probably won't always be fed_ex in future design fed_ex = self.parent_executor.federating_executor - return await fed_ex._compute_intrinsic_federated_map( + val_encrypted = await fed_ex._compute_intrinsic_federated_map( FederatingExecutorValue( anonymous_tuple.AnonymousTuple([(None, fn), (None, val_key_zipped)]), computation_types.NamedTupleType((fn_type, val_type)))) + if sender != None or receiver != None: + return val_encrypted.internal_representation[0] + else: + return val_encrypted.internal_representation + async def _decrypt_values_on_receiver(self, val, sender=None, receiver=None): pk_sender = self.key_references.get_public_key(self.sender_placement.name) From e6e62e2c5663317987beb2ec4bba806ad1bbd57b Mon Sep 17 00:00:00 2001 From: yann dupis Date: Thu, 14 May 2020 15:52:09 -0700 Subject: [PATCH 29/31] cache encrypt and decrypt computation --- .../impl/executors/federating_executor.py | 87 ++++++++++++------- 1 file changed, 54 insertions(+), 33 deletions(-) diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor.py b/tensorflow_federated/python/core/impl/executors/federating_executor.py index 17d2f9c402..bae134d3c1 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor.py @@ -1122,6 +1122,9 @@ def __init__(self, parent_executor, sender_placement, receiver_placement): self.key_references = KeyStore() self._is_channel_setup = False + self._encrypt_tensor_fn = None + self._decrypt_tensor_fn = None + async def setup(self): if not self._is_channel_setup: @@ -1201,23 +1204,16 @@ async def _encrypt_values_on_sender(self, val, sender=None, receiver=None): pk_receiver = self.key_references.get_public_key( self.receiver_placement.name) sk_sender = self.key_references.get_secret_key(self.sender_placement.name) - pk_rcv_tensor_type = pk_receiver.type_signature.member - sk_sender_tensor_type = sk_sender.type_signature.member - - @computations.tf_computation(input_tensor_type, pk_rcv_tensor_type, - sk_sender_tensor_type) - def encrypt_tensor(plaintext, pk_rcv, sk_snd): + pk_rcv_type = pk_receiver.type_signature.member + sk_snd_type = sk_sender.type_signature.member - pk_rcv = easy_box.PublicKey(pk_rcv) - sk_snd = easy_box.PublicKey(sk_snd) + if not self._encrypt_tensor_fn: + self._encrypt_tensor_fn = self._encrypt_tensor(input_tensor_type, + pk_rcv_type, sk_snd_type) - nonce = easy_box.gen_nonce() - ciphertext, mac = easy_box.seal_detached(plaintext, nonce, pk_rcv, sk_snd) - - return ciphertext.raw, mac.raw, nonce.raw + fn_type = self._encrypt_tensor_fn.type_signature + fn = self._encrypt_tensor_fn._computation_proto - fn_type = encrypt_tensor.type_signature - fn = encrypt_tensor._computation_proto if nb_senders == 1: val_type = val.type_signature val = val.internal_representation @@ -1263,33 +1259,23 @@ async def _decrypt_values_on_receiver(self, val, sender=None, receiver=None): pk_index=sender, sk_index=receiver) - sender_output_type_signature = val[0].type_signature[0] - receiver_secret_key_type_signature = val[0].type_signature[1] - sender_public_key_type_signature = val[0].type_signature[2] + sender_values_type = val[0].type_signature[0] + pk_snd_type = val[0].type_signature[1] + sk_snd_type = val[0].type_signature[2] - @computations.tf_computation(sender_output_type_signature, - sender_public_key_type_signature, - receiver_secret_key_type_signature) - def decrypt_tensor(client_outputs, pk_snd, sk_rcv): + if not self._decrypt_tensor_fn: + self._decrypt_tensor_fn = self._decrypt_tensor( + sender_values_type, pk_snd_type, sk_snd_type, + self.orig_sender_tensor_dtype) - ciphertext = easy_box.Ciphertext(client_outputs[0]) - mac = easy_box.Mac(client_outputs[1]) - nonce = easy_box.Nonce(client_outputs[2]) - sk_rcv = easy_box.SecretKey(sk_rcv) - pk_snd = easy_box.PublicKey(pk_snd) - - plaintext_recovered = easy_box.open_detached( - ciphertext, mac, nonce, pk_snd, sk_rcv, self.orig_sender_tensor_dtype) - - return plaintext_recovered + fn_type = self._decrypt_tensor_fn.type_signature + fn = self._decrypt_tensor_fn._computation_proto val_type = computation_types.FederatedType( computation_types.TensorType(self.orig_sender_tensor_dtype), self.receiver_placement, all_equal=False) - fn_type = decrypt_tensor.type_signature - fn = decrypt_tensor._computation_proto # NOTE probably won't always be fed_ex in future design fed_ex = self.parent_executor.federating_executor @@ -1303,6 +1289,41 @@ def decrypt_tensor(client_outputs, pk_snd, sk_rcv): else: return val_decrypted.internal_representation + def _encrypt_tensor(self, plaintext_type, pk_rcv_type, sk_snd_type): + + @computations.tf_computation(plaintext_type, pk_rcv_type, sk_snd_type) + def encrypt_tensor(plaintext, pk_rcv, sk_snd): + + pk_rcv = easy_box.PublicKey(pk_rcv) + sk_snd = easy_box.PublicKey(sk_snd) + + nonce = easy_box.gen_nonce() + ciphertext, mac = easy_box.seal_detached(plaintext, nonce, pk_rcv, sk_snd) + + return ciphertext.raw, mac.raw, nonce.raw + + return encrypt_tensor + + def _decrypt_tensor(self, sender_values_type, pk_snd_type, sk_rcv_snd, + orig_sender_tensor_dtype): + + @computations.tf_computation(sender_values_type, pk_snd_type, sk_rcv_snd) + def decrypt_tensor(sender_values, pk_snd, sk_rcv): + + ciphertext = easy_box.Ciphertext(sender_values[0]) + mac = easy_box.Mac(sender_values[1]) + nonce = easy_box.Nonce(sender_values[2]) + sk_rcv = easy_box.SecretKey(sk_rcv) + pk_snd = easy_box.PublicKey(pk_snd) + + plaintext_recovered = easy_box.open_detached(ciphertext, mac, nonce, + pk_snd, sk_rcv, + orig_sender_tensor_dtype) + + return plaintext_recovered + + return decrypt_tensor + async def _zip_val_key(self, placement, vals, From ec2dd907fc1de2003900046963af5eed6cc19633 Mon Sep 17 00:00:00 2001 From: yann dupis Date: Thu, 14 May 2020 16:39:21 -0700 Subject: [PATCH 30/31] add key setup test --- .../executors/federating_executor_test.py | 82 ++++++++++++------- 1 file changed, 52 insertions(+), 30 deletions(-) diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor_test.py b/tensorflow_federated/python/core/impl/executors/federating_executor_test.py index c1f5fd5436..790cc77988 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor_test.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor_test.py @@ -40,6 +40,8 @@ from tensorflow_federated.python.core.impl.executors import federating_executor from tensorflow_federated.python.core.impl.executors import reference_resolving_executor +from tensorflow_federated.python.core.impl.executors import channel_base + tf.compat.v1.enable_v2_behavior() @@ -1272,49 +1274,69 @@ def __init__(self, parent_executor): create_test_executor(intrinsic_strategy_fn=MockIntrinsicStrategy) -# class EncryptionTest(parameterized.TestCase): +class EncryptionTest(parameterized.TestCase): + + def test_generate_aggregator_keys(self): + strategy = federating_executor.TrustedAggregatorIntrinsicStrategy + loop, ex = _make_test_runtime(intrinsic_strategy_fn=strategy) + strat_ex = ex.intrinsic_strategy + + channel_grid = channel_base.ChannelGrid({ + (placement_literals.AGGREGATORS, placement_literals.CLIENTS): + federating_executor.EasyBoxChannel( + parent_executor=strat_ex, + sender_placement=placement_literals.CLIENTS, + receiver_placement=placement_literals.AGGREGATORS) + }) + + channel = channel_grid[(placement_literals.CLIENTS, + placement_literals.AGGREGATORS)] + loop.run_until_complete(channel.setup()) + key_references = channel.key_references + + pk_c = key_references.get_public_key(placement_literals.CLIENTS.name) + sk_c = key_references.get_secret_key(placement_literals.CLIENTS.name) + pk_a = key_references.get_public_key(placement_literals.AGGREGATORS.name) + sk_a = key_references.get_secret_key(placement_literals.AGGREGATORS.name) -# def test_generate_aggregator_keys(self): -# strategy = federating_executor.TrustedAggregatorIntrinsicStrategy -# loop, ex = _make_test_runtime(intrinsic_strategy_fn=strategy) -# generate_keys = ex.intrinsic_strategy._trusted_aggregator_generate_keys() -# pk, sk = loop.run_until_complete(generate_keys) + self.assertEqual(str(pk_c.type_signature), '{uint8[32]}@AGGREGATORS') + self.assertEqual(str(sk_c.type_signature), '{uint8[32]}@CLIENTS') + self.assertEqual(str(pk_a.type_signature), '{uint8[32]}@CLIENTS') + self.assertEqual(str(sk_a.type_signature), '{uint8[32]}@AGGREGATORS') -# self.assertEqual(str(pk.type_signature), 'uint8[32]@CLIENTS') -# self.assertEqual(str(sk.type_signature), 'uint8[32]@AGGREGATORS') + # def test_encryption_decryption(self): -# def test_encryption_decryption(self): + # strategy = federating_executor.TrustedAggregatorIntrinsicStrategy + # loop, ex = _make_test_runtime(intrinsic_strategy_fn=strategy) + # strat_ex = ex.intrinsic_strategy -# strategy = federating_executor.TrustedAggregatorIntrinsicStrategy -# loop, ex = _make_test_runtime(intrinsic_strategy_fn=strategy) -# strat_ex = ex.intrinsic_strategy + # pk_a, sk_a = loop.run_until_complete( + # strat_ex._trusted_aggregator_generate_keys()) -# pk_a, sk_a = loop.run_until_complete( -# strat_ex._trusted_aggregator_generate_keys()) + # val = loop.run_until_complete( + # ex.create_value([2.0], type_factory.at_clients(tf.float32))) -# val = loop.run_until_complete( -# ex.create_value([2.0], type_factory.at_clients(tf.float32))) + # val_enc = loop.run_until_complete( + # strat_ex._encrypt_client_tensors(val, pk_a)) -# val_enc = loop.run_until_complete( -# strat_ex._encrypt_client_tensors(val, pk_a)) + # aggr = strat_ex._get_child_executors( + # placement_literals.AGGREGATORS, index=0) -# aggr = strat_ex._get_child_executors( -# placement_literals.AGGREGATORS, index=0) + # enc_val_on_aggr = loop.run_until_complete( + # strat_ex._move(val_enc.internal_representation[0], + # val_enc.type_signature.member, aggr)) -# enc_val_on_aggr = loop.run_until_complete( -# strat_ex._move(val_enc.internal_representation[0], -# val_enc.type_signature.member, aggr)) + # val_key_zipped = loop.run_until_complete( + # strat_ex._zip_val_key([enc_val_on_aggr], sk_a, + # placement_literals.AGGREGATORS)) -# val_key_zipped = loop.run_until_complete( -# strat_ex._zip_val_key([enc_val_on_aggr], sk_a, -# placement_literals.AGGREGATORS)) + # val_dec = loop.run_until_complete( + # strat_ex._decrypt_tensors_on_aggregator(val_key_zipped, tf.float32)) -# val_dec = loop.run_until_complete( -# strat_ex._decrypt_tensors_on_aggregator(val_key_zipped, tf.float32)) + # dec_tf_tensor = val_dec.internal_representation[0].internal_representation -# dec_tf_tensor = val_dec.internal_representation[0].internal_representation + # self.assertEqual(dec_tf_tensor, tf.constant(2.0, dtype=tf.float32)) -# self.assertEqual(dec_tf_tensor, tf.constant(2.0, dtype=tf.float32)) if __name__ == '__main__': absltest.main() From 07404a97b1886c1421805c400af26789c0d314bb Mon Sep 17 00:00:00 2001 From: yann dupis Date: Thu, 14 May 2020 16:49:34 -0700 Subject: [PATCH 31/31] add encryption decryption test --- .../impl/executors/federating_executor.py | 3 +- .../executors/federating_executor_test.py | 46 +++++++++---------- 2 files changed, 23 insertions(+), 26 deletions(-) diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor.py b/tensorflow_federated/python/core/impl/executors/federating_executor.py index bae134d3c1..717a4fc33a 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor.py @@ -1215,8 +1215,7 @@ async def _encrypt_values_on_sender(self, val, sender=None, receiver=None): fn = self._encrypt_tensor_fn._computation_proto if nb_senders == 1: - val_type = val.type_signature - val = val.internal_representation + tensor_type = val.type_signature else: tensor_type = val[0].type_signature diff --git a/tensorflow_federated/python/core/impl/executors/federating_executor_test.py b/tensorflow_federated/python/core/impl/executors/federating_executor_test.py index 790cc77988..e16d2b2a92 100644 --- a/tensorflow_federated/python/core/impl/executors/federating_executor_test.py +++ b/tensorflow_federated/python/core/impl/executors/federating_executor_test.py @@ -1304,38 +1304,36 @@ def test_generate_aggregator_keys(self): self.assertEqual(str(pk_a.type_signature), '{uint8[32]}@CLIENTS') self.assertEqual(str(sk_a.type_signature), '{uint8[32]}@AGGREGATORS') - # def test_encryption_decryption(self): + def test_encryption_decryption(self): - # strategy = federating_executor.TrustedAggregatorIntrinsicStrategy - # loop, ex = _make_test_runtime(intrinsic_strategy_fn=strategy) - # strat_ex = ex.intrinsic_strategy - - # pk_a, sk_a = loop.run_until_complete( - # strat_ex._trusted_aggregator_generate_keys()) - - # val = loop.run_until_complete( - # ex.create_value([2.0], type_factory.at_clients(tf.float32))) + strategy = federating_executor.TrustedAggregatorIntrinsicStrategy + loop, ex = _make_test_runtime(intrinsic_strategy_fn=strategy) + strat_ex = ex.intrinsic_strategy - # val_enc = loop.run_until_complete( - # strat_ex._encrypt_client_tensors(val, pk_a)) + channel_grid = channel_base.ChannelGrid({ + (placement_literals.AGGREGATORS, placement_literals.CLIENTS): + federating_executor.EasyBoxChannel( + parent_executor=strat_ex, + sender_placement=placement_literals.CLIENTS, + receiver_placement=placement_literals.AGGREGATORS) + }) - # aggr = strat_ex._get_child_executors( - # placement_literals.AGGREGATORS, index=0) + channel = channel_grid[(placement_literals.CLIENTS, + placement_literals.AGGREGATORS)] + loop.run_until_complete(channel.setup()) - # enc_val_on_aggr = loop.run_until_complete( - # strat_ex._move(val_enc.internal_representation[0], - # val_enc.type_signature.member, aggr)) + val = loop.run_until_complete( + ex.create_value([2.0], type_factory.at_clients(tf.float32))) - # val_key_zipped = loop.run_until_complete( - # strat_ex._zip_val_key([enc_val_on_aggr], sk_a, - # placement_literals.AGGREGATORS)) + val_enc = loop.run_until_complete( + channel.send(val.internal_representation[0])) - # val_dec = loop.run_until_complete( - # strat_ex._decrypt_tensors_on_aggregator(val_key_zipped, tf.float32)) + val_dec = loop.run_until_complete( + channel.receive(val_enc)) - # dec_tf_tensor = val_dec.internal_representation[0].internal_representation + dec_tf_tensor = val_dec[0].internal_representation - # self.assertEqual(dec_tf_tensor, tf.constant(2.0, dtype=tf.float32)) + self.assertEqual(dec_tf_tensor, tf.constant(2.0, dtype=tf.float32)) if __name__ == '__main__':