def tensor_db(): """Prepare tensor db.""" db = TensorDB() array_1 = np.array([0, 1, 2, 3, 4]) tensor_key_1 = TensorKey('tensor_name', 'agg', 0, False, ('col1', )) array_2 = np.array([2, 3, 4, 5, 6]) tensor_key_2 = TensorKey('tensor_name', 'agg', 0, False, ('col2', )) db.cache_tensor({tensor_key_1: array_1, tensor_key_2: array_2}) return db
def test_get_aggregated_tensor_error_aggregation_function(tensor_db): """Test that get_aggregated_tensor raise error if aggregation function is not callable.""" collaborator_weight_dict = {'col1': 0.1, 'col2': 0.9} tensor_key = TensorKey('tensor_name', 'agg', 0, False, ()) with pytest.raises(KeyError): tensor_db.get_aggregated_tensor(tensor_key, collaborator_weight_dict, 'fake_agg_function')
def tensor_key(collaborator_mock, named_tensor): """Initialize the tensor_key mock.""" tensor_key = TensorKey(named_tensor.name, collaborator_mock.collaborator_name, named_tensor.round_number, named_tensor.report, tuple(named_tensor.tags)) return tensor_key
def test_get_aggregated_tensor_new_aggregation_function(tensor_db): """Test that get_aggregated_tensor works correctly with a given agg function.""" collaborator_weight_dict = {'col1': 0.1, 'col2': 0.9} tensor_key = TensorKey('tensor_name', 'agg', 0, False, ()) agg_nparray, agg_metadata_dict = tensor_db.get_aggregated_tensor( tensor_key, collaborator_weight_dict, ['sum']) assert np.array_equal(agg_nparray, np.array([2, 4, 6, 8, 10]))
def tensor_key_trained(collaborator_mock, named_tensor): """Initialize the tensor_key_trained mock.""" named_tensor.tags.append('trained') named_tensor.tags.remove('model') tensor_key = TensorKey(named_tensor.name, collaborator_mock.collaborator_name, named_tensor.round_number, named_tensor.report, tuple(named_tensor.tags)) return tensor_key
def tensor_key(named_tensor): """Initialize the tensor_key mock.""" tensor_key = TensorKey( named_tensor.name, 'col1', named_tensor.round_number, named_tensor.report, tuple(named_tensor.tags) ) return tensor_key
def test_find_dependencies_with_zero_round(tensor_key): """Test that find_dependencies returns empty list when round number is 0.""" tensor_codec = TensorCodec(NoCompressionPipeline()) tensor_name, origin, round_number, report, tags = tensor_key tensor_key = TensorKey( tensor_name, origin, round_number, report, ('model',) ) tensor_key_dependencies = tensor_codec.find_dependencies(tensor_key, True) assert len(tensor_key_dependencies) == 0
def test_find_dependencies_without_send_model_deltas(tensor_key): """Test that find_dependencies returns empty list when send_model_deltas = False.""" tensor_codec = TensorCodec(NoCompressionPipeline()) tensor_name, origin, round_number, report, tags = tensor_key tensor_key = TensorKey( tensor_name, origin, 5, report, ('model',) ) tensor_key_dependencies = tensor_codec.find_dependencies(tensor_key, False) assert len(tensor_key_dependencies) == 0
def test_get_aggregated_tensor_multiple_aggregation_functions(tensor_db): """Test that get_aggregated_tensor works correctly with multiple agg functions.""" collaborator_weight_dict = {'col1': 0.1, 'col2': 0.9} tensor_key = TensorKey('tensor_name', 'agg', 0, False, ()) agg_nparray, agg_metadata_dict = tensor_db.get_aggregated_tensor( tensor_key, collaborator_weight_dict, ['sum', 'mean']) assert np.array_equal(agg_nparray, np.array([2, 4, 6, 8, 10])) assert 'mean' in agg_metadata_dict assert np.array_equal(agg_metadata_dict['mean'], np.array([1., 2., 3., 4., 5.]))
def test_get_aggregated_tensor_directly(nparray, tensor_key): """Test that get_aggregated_tensor returns tensors directly.""" db = TensorDB() db.cache_tensor({tensor_key: nparray}) tensor_name, origin, round_number, report, tags = tensor_key tensor_key = TensorKey(tensor_name, 'col2', round_number, report, ('model', )) db.cache_tensor({tensor_key: nparray}) agg_nparray, agg_metadata_dict = db.get_aggregated_tensor( tensor_key, {}, None) assert np.array_equal(nparray, agg_nparray)
def test_get_aggregated_tensor_only_col(nparray, tensor_key): """Test that get_aggregated_tensor returns None if data presents for only collaborator.""" db = TensorDB() db.cache_tensor({tensor_key: nparray}) tensor_name, origin, round_number, report, tags = tensor_key tensor_key = TensorKey(tensor_name, 'col2', round_number, report, ('model', )) collaborator_weight_dict = {'col1': 0.5, 'col2': 0.5} agg_nparray, agg_metadata_dict = db.get_aggregated_tensor( tensor_key, collaborator_weight_dict, None) assert agg_nparray is None
def test_get_aggregated_tensor(nparray, tensor_key): """Test that get_aggregated_tensor returns tensors directly.""" db = TensorDB() db.cache_tensor({tensor_key: nparray}) tensor_name, origin, round_number, report, tags = tensor_key tensor_key = TensorKey(tensor_name, 'col2', round_number, report, ('model', )) db.cache_tensor({tensor_key: nparray}) collaborator_weight_dict = {'col1': 0.5, 'col2': 0.5} agg_nparray, agg_metadata_dict = db.get_aggregated_tensor( tensor_key, collaborator_weight_dict, WeightedAverage()) assert np.array_equal(nparray, agg_nparray)
def test_get_aggregated_tensor_weights(tensor_db): """Test that get_aggregated_tensor calculates correctly.""" collaborator_weight_dict = {'col1': 0.1, 'col2': 0.9} tensor_key = TensorKey('tensor_name', 'agg', 0, False, ()) agg_nparray, agg_metadata_dict = tensor_db.get_aggregated_tensor( tensor_key, collaborator_weight_dict, None) control_nparray = np.average( [np.array([0, 1, 2, 3, 4]), np.array([2, 3, 4, 5, 6])], weights=np.array(list(collaborator_weight_dict.values())), axis=0) assert np.array_equal(agg_nparray, control_nparray)
def test_decompress_require_lossless_no_compressed_in_tags(tensor_key, named_tensor): """Test that decompress raises error when require_lossless is True and is no compressed tag.""" tensor_codec = TensorCodec(NoCompressionPipeline()) tensor_name, origin, round_number, report, tags = tensor_key tensor_key = TensorKey( tensor_name, origin, round_number, report, ('lossy_compressed',) ) metadata = [{'int_to_float': proto.int_to_float, 'int_list': proto.int_list, 'bool_list': proto.bool_list } for proto in named_tensor.transformer_metadata] with pytest.raises(AssertionError): tensor_codec.decompress( tensor_key, named_tensor.data_bytes, metadata, require_lossless=True )
def test_find_dependencies(tensor_key): """Test that find_dependencies works correctly.""" tensor_codec = TensorCodec(NoCompressionPipeline()) tensor_name, origin, round_number, report, tags = tensor_key round_number = 2 tensor_key = TensorKey( tensor_name, origin, round_number, report, ('model',) ) tensor_key_dependencies = tensor_codec.find_dependencies(tensor_key, True) assert len(tensor_key_dependencies) == 2 tensor_key_dependency_0, tensor_key_dependency_1 = tensor_key_dependencies assert tensor_key_dependency_0.round_number == round_number - 1 assert tensor_key_dependency_0.tags == tensor_key.tags assert tensor_key_dependency_1.tags == ('aggregated', 'delta', 'compressed')
def test_decompress_compressed_in_tags(tensor_key, named_tensor): """Test that decompress works correctly when there is compressed tag.""" tensor_codec = TensorCodec(NoCompressionPipeline()) tensor_name, origin, round_number, report, tags = tensor_key tensor_key = TensorKey( tensor_name, origin, round_number, report, ('compressed',) ) metadata = [{'int_to_float': proto.int_to_float, 'int_list': proto.int_list, 'bool_list': proto.bool_list } for proto in named_tensor.transformer_metadata] decompressed_tensor_key, decompressed_nparray = tensor_codec.decompress( tensor_key, named_tensor.data_bytes, metadata ) assert 'compressed' not in decompressed_tensor_key.tags
def test_get_aggregated_tensor_new_aggregation_function(tensor_db): """Test that get_aggregated_tensor works correctly with a given agg function.""" collaborator_weight_dict = {'col1': 0.1, 'col2': 0.9} class Sum(AggregationFunctionInterface): def call(self, local_tensors, *_): tensors = [local_tensor.tensor for local_tensor in local_tensors] return np.sum(tensors, axis=0) tensor_key = TensorKey('tensor_name', 'agg', 0, False, ()) agg_nparray = tensor_db.get_aggregated_tensor(tensor_key, collaborator_weight_dict, Sum()) assert np.array_equal(agg_nparray, np.array([2, 4, 6, 8, 10]))
def test_decompress_call_lossless_pipeline_with_require_lossless(tensor_key, named_tensor): """Test that decompress calls lossless pipeline when require_lossless is True.""" tensor_codec = TensorCodec(NoCompressionPipeline()) tensor_name, origin, round_number, report, tags = tensor_key tensor_key = TensorKey( tensor_name, origin, round_number, report, ('compressed',) ) metadata = [{'int_to_float': proto.int_to_float, 'int_list': proto.int_list, 'bool_list': proto.bool_list } for proto in named_tensor.transformer_metadata] tensor_codec.lossless_pipeline = mock.Mock() tensor_codec.decompress( tensor_key, named_tensor.data_bytes, metadata, require_lossless=True ) tensor_codec.lossless_pipeline.backward.assert_called_with( named_tensor.data_bytes, metadata)
def test_decompress_call_compression_pipeline(tensor_key, named_tensor): """Test that decompress calls compression pipeline when there is no compressed tag.""" tensor_codec = TensorCodec(NoCompressionPipeline()) tensor_name, origin, round_number, report, tags = tensor_key tensor_key = TensorKey( tensor_name, origin, round_number, report, ('lossy_compressed',) ) metadata = [{'int_to_float': proto.int_to_float, 'int_list': proto.int_list, 'bool_list': proto.bool_list } for proto in named_tensor.transformer_metadata] tensor_codec.compression_pipeline = mock.Mock() tensor_codec.decompress( tensor_key, named_tensor.data_bytes, metadata ) tensor_codec.compression_pipeline.backward.assert_called_with( named_tensor.data_bytes, metadata)
def test_generate_delta_assert_model_in_tags(tensor_key, named_tensor): """Test that generate_delta raises exception when there is model tag.""" tensor_codec = TensorCodec(NoCompressionPipeline()) tensor_name, origin, round_number, report, tags = tensor_key tensor_key = TensorKey( tensor_name, origin, round_number, report, ('model',) ) metadata = [{'int_to_float': proto.int_to_float, 'int_list': proto.int_list, 'bool_list': proto.bool_list } for proto in named_tensor.transformer_metadata] array_shape = tuple(metadata[0]['int_list']) flat_array = np.frombuffer(named_tensor.data_bytes, dtype=np.float32) nparray = np.reshape(flat_array, newshape=array_shape, order='C') with pytest.raises(AssertionError): tensor_codec.generate_delta(tensor_key, nparray, nparray)
def test_apply_delta_agg(tensor_key, named_tensor): """Test that apply_delta works for aggregator tensor_key.""" tensor_codec = TensorCodec(NoCompressionPipeline()) tensor_name, origin, round_number, report, tags = tensor_key tensor_key = TensorKey( tensor_name, 'aggregator_1', round_number, report, ('delta',) ) metadata = [{'int_to_float': proto.int_to_float, 'int_list': proto.int_list, 'bool_list': proto.bool_list } for proto in named_tensor.transformer_metadata] array_shape = tuple(metadata[0]['int_list']) flat_array = np.frombuffer(named_tensor.data_bytes, dtype=np.float32) nparray = np.reshape(flat_array, newshape=array_shape, order='C') new_model_tensor_key, nparray_with_delta = tensor_codec.apply_delta( tensor_key, nparray, nparray) assert 'delta' not in new_model_tensor_key.tags assert np.array_equal(nparray_with_delta, nparray + nparray)