def _send_histograms(self, name, hists): msg = tree_pb2.Histograms() for hist in hists: ciphertext = _encode_encrypted_numbers(hist) msg.hists.append( tree_pb2.EncryptedNumbers(ciphertext=ciphertext)) self._bridge.send_proto(self._bridge.current_iter_id, name, msg)
def _receive_and_decrypt_histogram(self, name): msg = tree_pb2.Histograms() self._bridge.receive_proto(self._bridge.current_iter_id, name).Unpack(msg) if not self._pool: return _decrypt_histogram_helper( (0, self._public_key, self._private_key, msg.hists)) job_size = (len(msg.hists) + self._num_parallel - 1) // self._num_parallel args = [(i * job_size, self._public_key, self._private_key, msg.hists[i * job_size:(i + 1) * job_size]) for i in range(self._num_parallel)] hists = self._pool.map(_decrypt_histogram_helper, args) return sum(hists, [])