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)
