class TensorBoardReleaseManagerReleaseTest(parameterized.TestCase,
unittest.IsolatedAsyncioTestCase,
tf.test.TestCase):
# pyformat: disable
@parameterized.named_parameters(
# materialized values
('bool', True, computation_types.TensorType(tf.bool), [('', True)]),
('int', 1, computation_types.TensorType(tf.int32), [('', 1)]),
('tensor_int', tf.constant(1), computation_types.TensorType(
tf.int32), [('', tf.constant(1))]),
('numpy_int', np.int32(1), computation_types.TensorType(
tf.int32), [('', np.int32(1))]),
# value references
('materializable_value_reference_tensor',
program_test_utils.TestMaterializableValueReference(1),
computation_types.TensorType(tf.int32), [('', 1)]),
# structures
('list', [
True,
program_test_utils.TestMaterializableValueReference(1), 'a'
], computation_types.SequenceType([tf.bool, tf.int32, tf.string
]), [('0', True), ('1', 1)]),
('list_nested', [
[True,
program_test_utils.TestMaterializableValueReference(1)], ['a']
], computation_types.SequenceType([[tf.bool, tf.int32], [tf.string]
]), [('0/0', True), ('0/1', 1)]),
('dict', {
'a': True,
'b': program_test_utils.TestMaterializableValueReference(1),
'c': 'a'
},
computation_types.SequenceType([('a', tf.bool), ('b', tf.int32),
('c', tf.string)]), [('a', True),
('b', 1)]),
('dict_nested', {
'x': {
'a': True,
'b': program_test_utils.TestMaterializableValueReference(1)
},
'y': {
'c': 'a'
}
},
computation_types.SequenceType(
[('x', [('a', tf.bool), ('b', tf.int32)]),
('y', [('c', tf.string)])]), [('x/a', True), ('x/b', 1)]),
('attr',
program_test_utils.TestAttrObject2(
True, program_test_utils.TestMaterializableValueReference(1)),
computation_types.SequenceType([('a', tf.bool),
('b', tf.int32)]), [('a', True),
('b', 1)]),
('attr_nested',
program_test_utils.TestAttrObject2(
program_test_utils.TestAttrObject2(
True, program_test_utils.TestMaterializableValueReference(1)),
program_test_utils.TestAttrObject1('a')),
computation_types.SequenceType(
[('x', [('a', tf.bool), ('b', tf.int32)]),
('y', [('c', tf.string)])]), [('a/a', True), ('a/b', 1)]),
)
# pyformat: enable
async def test_writes_value_scalar(self, value, type_signature,
expected_calls):
summary_dir = self.create_tempdir()
release_mngr = tensorboard_release_manager.TensorBoardReleaseManager(
summary_dir=summary_dir)
with mock.patch.object(tf.summary, 'scalar') as mock_scalar:
await release_mngr.release(value, type_signature, 1)
self.assertLen(mock_scalar.mock_calls, len(expected_calls))
for call, expected_args in zip(mock_scalar.mock_calls,
expected_calls):
_, args, kwargs = call
actual_name, actual_value = args
expected_name, expected_value = expected_args
self.assertEqual(actual_name, expected_name)
self.assertEqual(actual_value, expected_value)
self.assertEqual(kwargs, {'step': 1})
# pyformat: disable
@parameterized.named_parameters(
# materialized values
('tensor_array', tf.ones(
[3], tf.int32), computation_types.TensorType(
tf.int32, [3]), [('', tf.ones([3], tf.int32))]),
('numpy_array', np.ones(
[3], np.int32), computation_types.TensorType(
tf.int32, [3]), [('', np.ones([3], np.int32))]),
# value references
('materializable_value_reference_sequence',
program_test_utils.TestMaterializableValueReference(
tf.data.Dataset.from_tensor_slices([1, 2, 3])),
computation_types.SequenceType([tf.int32] * 3), [('', [1, 2, 3])]),
)
# pyformat: enable
async def test_writes_value_histogram(self, value, type_signature,
expected_calls):
summary_dir = self.create_tempdir()
release_mngr = tensorboard_release_manager.TensorBoardReleaseManager(
summary_dir=summary_dir)
with mock.patch.object(tf.summary, 'histogram') as mock_histogram:
await release_mngr.release(value, type_signature, 1)
self.assertLen(mock_histogram.mock_calls, len(expected_calls))
for call, expected_args in zip(mock_histogram.mock_calls,
expected_calls):
_, args, kwargs = call
actual_name, actual_value = args
expected_name, expected_value = expected_args
self.assertEqual(actual_name, expected_name)
self.assertAllEqual(actual_value, expected_value)
self.assertEqual(kwargs, {'step': 1})
async def test_writes_value_scalar_and_histogram(self):
summary_dir = self.create_tempdir()
release_mngr = tensorboard_release_manager.TensorBoardReleaseManager(
summary_dir=summary_dir)
value = [1, tf.ones([3], tf.int32)]
type_signature = computation_types.StructType([
tf.int32,
computation_types.TensorType(tf.float32, [3]),
])
patched_scalar = mock.patch.object(tf.summary, 'scalar')
patched_histogram = mock.patch.object(tf.summary, 'histogram')
with patched_scalar as mock_scalar, patched_histogram as mock_histogram:
await release_mngr.release(value, type_signature, 1)
mock_scalar.assert_called_once_with('0', 1, step=1)
self.assertLen(mock_histogram.mock_calls, 1)
call = mock_histogram.mock_calls[0]
_, args, kwargs = call
actual_name, actual_value = args
expected_name, expected_value = '1', tf.ones([3], tf.int32)
self.assertEqual(actual_name, expected_name)
self.assertAllEqual(actual_value, expected_value)
self.assertEqual(kwargs, {'step': 1})
@parameterized.named_parameters(
# materialized values
('none', None, computation_types.StructType([])),
('str', 'a', computation_types.TensorType(tf.string)),
('tensor_str', tf.constant('a'), computation_types.TensorType(
tf.string)),
# structures
('list_empty', [], computation_types.StructType(())),
('dict_empty', {}, computation_types.StructType(())),
)
async def test_does_not_write_value(self, value, type_signature):
summary_dir = self.create_tempdir()
release_mngr = tensorboard_release_manager.TensorBoardReleaseManager(
summary_dir=summary_dir)
patch_scalar = mock.patch.object(tf.summary, 'scalar')
patch_histogram = mock.patch.object(tf.summary, 'histogram')
with patch_scalar as mock_scalar, patch_histogram as mock_histogram:
await release_mngr.release(value, type_signature, 1)
mock_scalar.assert_not_called()
mock_histogram.assert_not_called()
@parameterized.named_parameters(
('negative_1', -1),
('0', 0),
('1', 1),
)
async def test_does_not_raise_type_error_with_key(self, key):
summary_dir = self.create_tempdir()
release_mngr = tensorboard_release_manager.TensorBoardReleaseManager(
summary_dir=summary_dir)
value = 1
type_signature = computation_types.TensorType(tf.int32)
try:
await release_mngr.release(value, type_signature, key)
except TypeError:
self.fail('Raised TypeError unexpectedly.')
@parameterized.named_parameters(
('none', None),
('str', 'a'),
('list', []),
)
async def test_raises_type_error_with_key(self, key):
summary_dir = self.create_tempdir()
release_mngr = tensorboard_release_manager.TensorBoardReleaseManager(
summary_dir=summary_dir)
value = 1
type_signature = computation_types.TensorType(tf.int32)
with self.assertRaises(TypeError):
await release_mngr.release(value, type_signature, key)
class FileProgramStateManagerLoadTest(parameterized.TestCase,
unittest.IsolatedAsyncioTestCase,
tf.test.TestCase):
# pyformat: disable
@parameterized.named_parameters(
# materialized values
('none', None, None),
('bool', True, tf.constant(True)),
('int', 1, tf.constant(1)),
('str', 'a', tf.constant('a')),
('tensor_int', tf.constant(1), tf.constant(1)),
('tensor_str', tf.constant('a'), tf.constant('a')),
('tensor_array', tf.ones([3], tf.int32), tf.ones([3], tf.int32)),
('numpy_int', np.int32(1), tf.constant(1)),
('numpy_array', np.ones([3], int), tf.ones([3], tf.int32)),
# value references
('materializable_value_reference_tensor',
program_test_utils.TestMaterializableValueReference(1), tf.constant(1)
),
('materializable_value_reference_sequence',
program_test_utils.TestMaterializableValueReference(
tf.data.Dataset.from_tensor_slices(
[1, 2, 3])), tf.data.Dataset.from_tensor_slices([1, 2, 3])),
# structures
('list',
[True,
program_test_utils.TestMaterializableValueReference(1), 'a'],
[tf.constant(True),
tf.constant(1), tf.constant('a')]),
('list_empty', [], []),
('list_nested', [[
True, program_test_utils.TestMaterializableValueReference(1)
], ['a']], [[tf.constant(True), tf.constant(1)], [tf.constant('a')]]),
('dict', {
'a': True,
'b': program_test_utils.TestMaterializableValueReference(1),
'c': 'a'
}, {
'a': tf.constant(True),
'b': tf.constant(1),
'c': tf.constant('a')
}),
('dict_empty', {}, {}),
('dict_nested', {
'x': {
'a': True,
'b': program_test_utils.TestMaterializableValueReference(1)
},
'y': {
'c': 'a'
}
}, {
'x': {
'a': tf.constant(True),
'b': tf.constant(1)
},
'y': {
'c': tf.constant('a')
}
}),
('attr',
program_test_utils.TestAttrObject2(
True, program_test_utils.TestMaterializableValueReference(1)),
program_test_utils.TestAttrObject2(tf.constant(True),
tf.constant(1))),
('attr_nested',
program_test_utils.TestAttrObject2(
program_test_utils.TestAttrObject2(
True, program_test_utils.TestMaterializableValueReference(1)),
program_test_utils.TestAttrObject1('a')),
program_test_utils.TestAttrObject2(
program_test_utils.TestAttrObject2(tf.constant(True),
tf.constant(1)),
program_test_utils.TestAttrObject1(tf.constant('a')))),
)
# pyformat: enable
async def test_returns_saved_program_state(self, program_state,
expected_program_state):
root_dir = self.create_tempdir()
program_state_mngr = file_program_state_manager.FileProgramStateManager(
root_dir=root_dir, prefix='a_')
await program_state_mngr.save(program_state, 1)
structure = program_state
actual_program_state = await program_state_mngr.load(1, structure)
if isinstance(actual_program_state, tf.data.Dataset):
actual_program_state = list(actual_program_state)
if isinstance(expected_program_state, tf.data.Dataset):
expected_program_state = list(expected_program_state)
self.assertAllEqual(actual_program_state, expected_program_state)
@parameterized.named_parameters(
('0', 0),
('1', 1),
('2', 2),
)
async def test_returns_saved_program_state_with_version(self, version):
root_dir = self.create_tempdir()
program_state_mngr = file_program_state_manager.FileProgramStateManager(
root_dir=root_dir, prefix='a_')
for i in range(3):
await program_state_mngr.save(f'state_{i}', i)
structure = 'state'
actual_program_state = await program_state_mngr.load(
version, structure)
expected_program_state = f'state_{version}'
self.assertEqual(actual_program_state, expected_program_state)
async def test_raises_version_not_found_error_with_no_saved_program_state(
self):
root_dir = self.create_tempdir()
program_state_mngr = file_program_state_manager.FileProgramStateManager(
root_dir=root_dir, prefix='a_')
with self.assertRaises(
program_state_manager.ProgramStateManagerStateNotFoundError):
await program_state_mngr.load(0, None)
async def test_raises_version_not_found_error_with_unknown_version(self):
root_dir = self.create_tempdir()
program_state_mngr = file_program_state_manager.FileProgramStateManager(
root_dir=root_dir, prefix='a_')
await program_state_mngr.save('state_1', 1)
structure = 'state'
with self.assertRaises(
program_state_manager.ProgramStateManagerStateNotFoundError):
await program_state_mngr.load(10, structure)
async def test_raises_structure_error(self):
root_dir = self.create_tempdir()
program_state_mngr = file_program_state_manager.FileProgramStateManager(
root_dir=root_dir, prefix='a_')
await program_state_mngr.save('state_1', 1)
structure = []
with self.assertRaises(
program_state_manager.ProgramStateManagerStructureError):
await program_state_mngr.load(1, structure)
@parameterized.named_parameters(
('none', None),
('str', 'a'),
('list', []),
)
async def test_raises_type_error_with_version(self, version):
root_dir = self.create_tempdir()
program_state_mngr = file_program_state_manager.FileProgramStateManager(
root_dir=root_dir, prefix='a_')
with self.assertRaises(TypeError):
await program_state_mngr.load(version, None)
class FileProgramStateManagerSaveTest(parameterized.TestCase,
unittest.IsolatedAsyncioTestCase,
tf.test.TestCase):
# pyformat: disable
@parameterized.named_parameters(
# materialized values
('none', None, [None]),
('bool', True, [True]),
('int', 1, [1]),
('str', 'a', ['a']),
('tensor_int', tf.constant(1), [tf.constant(1)]),
('tensor_str', tf.constant('a'), [tf.constant('a')]),
('tensor_array', tf.ones([3], tf.int32), [tf.ones([3], tf.int32)]),
('numpy_int', np.int32(1), [np.int32(1)]),
('numpy_array', np.ones([3], int), [np.ones([3], int)]),
# value references
('materializable_value_reference_tensor',
program_test_utils.TestMaterializableValueReference(1), [1]),
('materializable_value_reference_sequence',
program_test_utils.TestMaterializableValueReference(
tf.data.Dataset.from_tensor_slices(
[1, 2, 3])), [tf.data.Dataset.from_tensor_slices([1, 2, 3])]),
# structures
('list', [
True,
program_test_utils.TestMaterializableValueReference(1), 'a'
], [True, 1, 'a']),
('list_empty', [], []),
('list_nested', [[
True, program_test_utils.TestMaterializableValueReference(1)
], ['a']], [True, 1, 'a']),
('dict', {
'a': True,
'b': program_test_utils.TestMaterializableValueReference(1),
'c': 'a'
}, [True, 1, 'a']),
('dict_empty', {}, []),
('dict_nested', {
'x': {
'a': True,
'b': program_test_utils.TestMaterializableValueReference(1)
},
'y': {
'c': 'a'
}
}, [True, 1, 'a']),
('attr',
program_test_utils.TestAttrObject2(
True, program_test_utils.TestMaterializableValueReference(1)),
[True, 1]),
('attr_nested',
program_test_utils.TestAttrObject2(
program_test_utils.TestAttrObject2(
True, program_test_utils.TestMaterializableValueReference(1)),
program_test_utils.TestAttrObject1('a')), [True, 1, 'a']),
)
# pyformat: enable
async def test_writes_program_state(self, program_state, expected_value):
root_dir = self.create_tempdir()
program_state_mngr = file_program_state_manager.FileProgramStateManager(
root_dir=root_dir, prefix='a_', keep_total=0)
with mock.patch.object(file_utils,
'write_saved_model') as mock_write_saved_model:
await program_state_mngr.save(program_state, 1)
mock_write_saved_model.assert_called_once()
call = mock_write_saved_model.mock_calls[0]
_, args, kwargs = call
actual_value, actual_path = args
def _normalize(value: Any) -> Any:
if isinstance(value, tf.data.Dataset):
return list(value)
return value
actual_value = tree.map_structure(_normalize, actual_value)
expected_value = tree.map_structure(_normalize, expected_value)
self.assertAllEqual(actual_value, expected_value)
expected_path = os.path.join(root_dir, 'a_1')
self.assertEqual(actual_path, expected_path)
self.assertEqual(kwargs, {})
async def test_removes_saved_program_state(self):
root_dir = self.create_tempdir()
program_state_mngr = file_program_state_manager.FileProgramStateManager(
root_dir=root_dir, prefix='a_')
with mock.patch.object(
program_state_mngr,
'_remove_old_program_state') as mock_remove_old_program_state:
await program_state_mngr.save('state_1', 1)
mock_remove_old_program_state.assert_called_once()
async def test_raises_version_already_exists_error_with_existing_version(
self):
root_dir = self.create_tempdir()
program_state_mngr = file_program_state_manager.FileProgramStateManager(
root_dir=root_dir, prefix='a_')
await program_state_mngr.save('state_1', 1)
with self.assertRaises(program_state_manager.
ProgramStateManagerStateAlreadyExistsError):
await program_state_mngr.save('state_1', 1)
@parameterized.named_parameters(
('none', None),
('str', 'a'),
('list', []),
)
async def test_raises_type_error_with_version(self, version):
root_dir = self.create_tempdir()
program_state_mngr = file_program_state_manager.FileProgramStateManager(
root_dir=root_dir, prefix='a_')
with self.assertRaises(TypeError):
await program_state_mngr.save('state', version)
class FlattenWithNameTest(parameterized.TestCase, tf.test.TestCase):
# pyformat: disable
@parameterized.named_parameters(
# materialized values
('none', None, [('', None)]),
('bool', True, [('', True)]),
('int', 1, [('', 1)]),
('str', 'a', [('', 'a')]),
('tensor_int', tf.constant(1), [('', tf.constant(1))]),
('tensor_str', tf.constant('a'), [('', tf.constant('a'))]),
('tensor_array', tf.ones([3]), [('', tf.ones([3]))]),
('numpy_int', np.int32(1), [('', np.int32(1))]),
('numpy_array', np.ones([3]), [('', np.ones([3]))]),
# value references
('materializable_value_reference_tensor',
program_test_utils.TestMaterializableValueReference(1), [
('', program_test_utils.TestMaterializableValueReference(1))
]),
('materializable_value_reference_sequence',
program_test_utils.TestMaterializableValueReference(
tf.data.Dataset.from_tensor_slices([1, 2, 3])), [
('',
program_test_utils.TestMaterializableValueReference(
tf.data.Dataset.from_tensor_slices([1, 2, 3])))
]),
# structures
('list', [
True,
program_test_utils.TestMaterializableValueReference(1), 'a'
], [('0', True),
('1', program_test_utils.TestMaterializableValueReference(1)),
('2', 'a')]),
('list_empty', [], []),
('list_nested', [[
True, program_test_utils.TestMaterializableValueReference(1)
], ['a']], [
('0/0', True),
('0/1', program_test_utils.TestMaterializableValueReference(1)),
('1/0', 'a')
]),
('dict', {
'a': True,
'b': program_test_utils.TestMaterializableValueReference(1),
'c': 'a'
}, [('a', True),
('b', program_test_utils.TestMaterializableValueReference(1)),
('c', 'a')]),
('dict_empty', {}, []),
('dict_nested', {
'x': {
'a': True,
'b': program_test_utils.TestMaterializableValueReference(1)
},
'y': {
'c': 'a'
}
}, [('x/a', True),
('x/b', program_test_utils.TestMaterializableValueReference(1)),
('y/c', 'a')]),
('attr',
program_test_utils.TestAttrObject2(
True, program_test_utils.TestMaterializableValueReference(1)), [
('a', True),
('b', program_test_utils.TestMaterializableValueReference(1))
]),
('attr_nested',
program_test_utils.TestAttrObject2(
program_test_utils.TestAttrObject2(
True, program_test_utils.TestMaterializableValueReference(1)),
program_test_utils.TestAttrObject1('a')),
[('a/a', True),
('a/b', program_test_utils.TestMaterializableValueReference(1)),
('b/a', 'a')]),
)
# pyformat: enable
def test_returns_result(self, structure, expected_result):
actual_result = structure_utils.flatten_with_name(structure)
for actual_item, expected_item in zip(actual_result, expected_result):
actual_path, actual_value = actual_item
expected_path, expected_value = expected_item
self.assertEqual(actual_path, expected_path)
self.assertAllEqual(actual_value, expected_value)
class MemoryReleaseManagerTest(parameterized.TestCase,
unittest.IsolatedAsyncioTestCase,
tf.test.TestCase):
# pyformat: disable
@parameterized.named_parameters(
# materialized values
('none', None, computation_types.StructType([]), None),
('bool', True, computation_types.TensorType(tf.bool), True),
('int', 1, computation_types.TensorType(tf.int32), 1),
('str', 'a', computation_types.TensorType(tf.string), 'a'),
('tensor_int', tf.constant(1), computation_types.TensorType(
tf.int32), tf.constant(1)),
('tensor_str', tf.constant('a'), computation_types.TensorType(
tf.string), tf.constant('a')),
('tensor_array', tf.ones([3], tf.int32),
computation_types.TensorType(tf.int32, [3]), tf.ones([3], tf.int32)),
('numpy_int', np.int32(1), computation_types.TensorType(
tf.int32), np.int32(1)),
('numpy_array', np.ones([3], int),
computation_types.TensorType(tf.int32, [3]), np.ones([3], int)),
# value references
('value_reference_tensor',
program_test_utils.TestMaterializableValueReference(1),
computation_types.TensorType(tf.int32), 1),
('value_reference_sequence',
program_test_utils.TestMaterializableValueReference(
tf.data.Dataset.from_tensor_slices(
[1, 2, 3])), computation_types.SequenceType(
tf.int32), tf.data.Dataset.from_tensor_slices([1, 2, 3])),
# structures
('list', [
True,
program_test_utils.TestMaterializableValueReference(1), 'a'
], computation_types.SequenceType([tf.bool, tf.int32, tf.string
]), [True, 1, 'a']),
('list_empty', [], computation_types.SequenceType([]), []),
('list_nested', [
[True,
program_test_utils.TestMaterializableValueReference(1)], ['a']
], computation_types.SequenceType([[tf.bool, tf.int32], [tf.string]
]), [[True, 1], ['a']]),
('dict', {
'a': True,
'b': program_test_utils.TestMaterializableValueReference(1),
'c': 'a'
},
computation_types.SequenceType([('a', tf.bool), ('b', tf.int32),
('c', tf.string)]), {
'a': True,
'b': 1,
'c': 'a'
}),
('dict_empty', {}, computation_types.SequenceType([]), {}),
('dict_nested', {
'x': {
'a': True,
'b': program_test_utils.TestMaterializableValueReference(1)
},
'y': {
'c': 'a'
}
},
computation_types.SequenceType([('x', [('a', tf.bool),
('b', tf.int32)]),
('y', [('c', tf.string)])]), {
'x': {
'a': True,
'b': 1
},
'y': {
'c': 'a'
}
}),
('attr',
program_test_utils.TestAttrObject2(
True, program_test_utils.TestMaterializableValueReference(1)),
computation_types.SequenceType([
('a', tf.bool), ('b', tf.int32)
]), program_test_utils.TestAttrObject2(True, 1)),
('attr_nested',
program_test_utils.TestAttrObject2(
program_test_utils.TestAttrObject2(
True, program_test_utils.TestMaterializableValueReference(1)),
program_test_utils.TestAttrObject1('a')),
computation_types.SequenceType([('a', [('a', tf.bool),
('b', tf.int32)]),
('b', [('c', tf.string)])]),
program_test_utils.TestAttrObject2(
program_test_utils.TestAttrObject2(True, 1),
program_test_utils.TestAttrObject1('a'))),
)
# pyformat: enable
async def test_release_saves_value(self, value, type_signature,
expected_value):
release_mngr = memory_release_manager.MemoryReleaseManager()
await release_mngr.release(value, type_signature, 1)
self.assertLen(release_mngr._values, 1)
actual_value = release_mngr._values[1]
if isinstance(actual_value, tf.data.Dataset):
actual_value = list(actual_value)
if isinstance(expected_value, tf.data.Dataset):
expected_value = list(expected_value)
self.assertAllEqual(actual_value, expected_value)
@parameterized.named_parameters(
('none', None),
('bool', True),
('int', 1),
('str', 'a'),
)
async def test_release_does_not_raise_type_error_with_key(self, key):
release_mngr = memory_release_manager.MemoryReleaseManager()
value = 1
type_signature = computation_types.TensorType(tf.int32)
try:
await release_mngr.release(value, type_signature, key)
except TypeError:
self.fail('Raised TypeError unexpectedly.')
@parameterized.named_parameters(
('list', []), )
async def test_release_raises_type_error_with_key(self, key):
release_mngr = memory_release_manager.MemoryReleaseManager()
value = 1
type_signature = computation_types.TensorType(tf.int32)
with self.assertRaises(TypeError):
await release_mngr.release(value, type_signature, key)
@parameterized.named_parameters(
('0', 0),
('1', 1),
('10', 10),
)
def test_values_returns_values(self, count):
release_mngr = memory_release_manager.MemoryReleaseManager()
for i in range(count):
release_mngr._values[i] = i * 10
values = release_mngr.values()
self.assertEqual(values, {i: i * 10 for i in range(count)})
def test_values_returns_copy(self):
release_mngr = memory_release_manager.MemoryReleaseManager()
values_1 = release_mngr.values()
values_2 = release_mngr.values()
self.assertIsNot(values_1, values_2)
class MaterializeValueTest(parameterized.TestCase,
unittest.IsolatedAsyncioTestCase, tf.test.TestCase):
# pyformat: disable
@parameterized.named_parameters(
# materialized values
('none', None, None),
('bool', True, True),
('int', 1, 1),
('str', 'a', 'a'),
('tensor_int', tf.constant(1), tf.constant(1)),
('tensor_str', tf.constant('a'), tf.constant('a')),
('tensor_array', tf.ones([3]), tf.ones([3])),
('numpy_int', np.int32(1), np.int32(1)),
('numpy_array', np.ones([3]), tf.ones([3])),
# value references
('materializable_value_reference_tensor',
program_test_utils.TestMaterializableValueReference(1), 1),
('materializable_value_reference_sequence',
program_test_utils.TestMaterializableValueReference(
tf.data.Dataset.from_tensor_slices(
[1, 2, 3])), tf.data.Dataset.from_tensor_slices([1, 2, 3])),
# structures
('list', [
True,
program_test_utils.TestMaterializableValueReference(1), 'a'
], [True, 1, 'a']),
('list_empty', [], []),
('list_nested', [[
True, program_test_utils.TestMaterializableValueReference(1)
], ['a']], [[True, 1], ['a']]),
('dict', {
'a': True,
'b': program_test_utils.TestMaterializableValueReference(1),
'c': 'a'
}, {
'a': True,
'b': 1,
'c': 'a'
}),
('dict_empty', {}, {}),
('dict_nested', {
'x': {
'a': True,
'b': program_test_utils.TestMaterializableValueReference(1)
},
'y': {
'c': 'a'
}
}, {
'x': {
'a': True,
'b': 1
},
'y': {
'c': 'a'
}
}),
('attr',
program_test_utils.TestAttrObject2(
True, program_test_utils.TestMaterializableValueReference(1)),
program_test_utils.TestAttrObject2(True, 1)),
('attr_nested',
program_test_utils.TestAttrObject2(
program_test_utils.TestAttrObject2(
True, program_test_utils.TestMaterializableValueReference(1)),
program_test_utils.TestAttrObject1('a')),
program_test_utils.TestAttrObject2(
program_test_utils.TestAttrObject2(True, 1),
program_test_utils.TestAttrObject1('a'))),
)
# pyformat: enable
async def test_returns_value(self, value, expected_value):
actual_value = await value_reference.materialize_value(value)
if isinstance(actual_value, tf.data.Dataset):
actual_value = list(actual_value)
if isinstance(expected_value, tf.data.Dataset):
expected_value = list(expected_value)
self.assertAllEqual(actual_value, expected_value)
class SavedModelFileReleaseManagerReleaseTest(parameterized.TestCase,
unittest.IsolatedAsyncioTestCase,
tf.test.TestCase):
# pyformat: disable
@parameterized.named_parameters(
# materialized values
('none', None, computation_types.StructType([]), [None]),
('bool', True, computation_types.TensorType(tf.bool), [True]),
('int', 1, computation_types.TensorType(tf.int32), [1]),
('str', 'a', computation_types.TensorType(tf.string), ['a']),
('tensor_int', tf.constant(1), computation_types.TensorType(
tf.int32), [tf.constant(1)]),
('tensor_str', tf.constant('a'), computation_types.TensorType(
tf.string), [tf.constant('a')]),
('tensor_array', tf.ones([3], tf.int32),
computation_types.TensorType(tf.int32, [3]), [tf.ones([3])]),
('numpy_int', np.int32(1), computation_types.TensorType(
tf.int32), [np.int32(1)]),
('numpy_array', np.ones([3], int),
computation_types.TensorType(tf.int32, [3]), [np.ones([3], int)]),
# value references
('materializable_value_reference_tensor',
program_test_utils.TestMaterializableValueReference(1),
computation_types.TensorType(tf.int32), [1]),
('materializable_value_reference_sequence',
program_test_utils.TestMaterializableValueReference(
tf.data.Dataset.from_tensor_slices([1, 2, 3])),
computation_types.SequenceType(
tf.int32), [tf.data.Dataset.from_tensor_slices([1, 2, 3])]),
# structures
('list', [
True,
program_test_utils.TestMaterializableValueReference(1), 'a'
], computation_types.SequenceType([tf.bool, tf.int32, tf.string
]), [True, 1, 'a']),
('list_empty', [], computation_types.SequenceType([]), []),
('list_nested', [
[True,
program_test_utils.TestMaterializableValueReference(1)], ['a']
], computation_types.SequenceType([[tf.bool, tf.int32], [tf.string]
]), [True, 1, 'a']),
('dict', {
'a': True,
'b': program_test_utils.TestMaterializableValueReference(1),
'c': 'a'
},
computation_types.SequenceType([('a', tf.bool), ('b', tf.int32),
('c', tf.string)]), [True, 1, 'a']),
('dict_empty', {}, computation_types.SequenceType([]), []),
('dict_nested', {
'x': {
'a': True,
'b': program_test_utils.TestMaterializableValueReference(1)
},
'y': {
'c': 'a'
}
},
computation_types.SequenceType(
[('x', [('a', tf.bool), ('b', tf.int32)]),
('y', [('c', tf.string)])]), [True, 1, 'a']),
('attr',
program_test_utils.TestAttrObject2(
True, program_test_utils.TestMaterializableValueReference(1)),
computation_types.SequenceType([('a', tf.bool),
('b', tf.int32)]), [True, 1]),
('attr_nested',
program_test_utils.TestAttrObject2(
program_test_utils.TestAttrObject2(
True, program_test_utils.TestMaterializableValueReference(1)),
program_test_utils.TestAttrObject1('a')),
computation_types.SequenceType(
[('a', [('a', tf.bool), ('b', tf.int32)]),
('b', [('c', tf.string)])]), [True, 1, 'a']),
)
# pyformat: enable
async def test_writes_value(self, value, type_signature, expected_value):
root_dir = self.create_tempdir()
release_mngr = file_release_manager.SavedModelFileReleaseManager(
root_dir=root_dir, prefix='a_')
with mock.patch.object(file_utils,
'write_saved_model') as mock_write_saved_model:
await release_mngr.release(value, type_signature, 1)
mock_write_saved_model.assert_called_once()
call = mock_write_saved_model.mock_calls[0]
_, args, kwargs = call
actual_value, actual_path = args
def _normalize(value: Any) -> Any:
if isinstance(value, tf.data.Dataset):
return list(value)
return value
actual_value = tree.map_structure(_normalize, actual_value)
expected_value = tree.map_structure(_normalize, expected_value)
self.assertAllEqual(actual_value, expected_value)
expected_path = os.path.join(root_dir, 'a_1')
self.assertEqual(actual_path, expected_path)
self.assertEqual(kwargs, {'overwrite': True})
@parameterized.named_parameters(
('negative_1', -1),
('0', 0),
('1', 1),
)
async def test_does_not_raise_type_error_with_key(self, key):
root_dir = self.create_tempdir()
release_mngr = file_release_manager.SavedModelFileReleaseManager(
root_dir=root_dir, prefix='a_')
value = 1
type_signature = computation_types.TensorType(tf.int32)
try:
await release_mngr.release(value, type_signature, key)
except TypeError:
self.fail('Raised TypeError unexpectedly.')
@parameterized.named_parameters(
('none', None),
('str', 'a'),
('list', []),
)
async def test_raises_type_error_with_key(self, key):
root_dir = self.create_tempdir()
release_mngr = file_release_manager.SavedModelFileReleaseManager(
root_dir=root_dir, prefix='a_')
value = 1
type_signature = computation_types.TensorType(tf.int32)
with self.assertRaises(TypeError):
await release_mngr.release(value, type_signature, key)
class CSVFileReleaseManagerReleaseTest(parameterized.TestCase,
unittest.IsolatedAsyncioTestCase):
async def test_calls_remove_values_greater_than_with_empty_file(self):
file_path = self.create_tempfile()
os.remove(file_path)
release_mngr = file_release_manager.CSVFileReleaseManager(
file_path=file_path)
value = {'a': 11, 'b': 21}
type_signature = computation_types.StructType([
('a', tf.int32),
('b', tf.int32),
])
with mock.patch.object(release_mngr, '_remove_values_greater_than'
) as mock_remove_values_greater_than:
await release_mngr.release(value, type_signature, 1)
mock_remove_values_greater_than.assert_called_once_with(0)
self.assertEqual(release_mngr._latest_key, 1)
async def test_calls_remove_values_greater_than_with_existing_file(self):
file_path = self.create_tempfile()
_write_values_to_csv(file_path=file_path,
fieldnames=['key', 'a', 'b'],
values=[{
'key': 1,
'a': 10,
'b': 20
}])
release_mngr = file_release_manager.CSVFileReleaseManager(
file_path=file_path)
value = {'a': 11, 'b': 21}
type_signature = computation_types.StructType([
('a', tf.int32),
('b', tf.int32),
])
with mock.patch.object(release_mngr, '_remove_values_greater_than'
) as mock_remove_values_greater_than:
await release_mngr.release(value, type_signature, 1)
mock_remove_values_greater_than.assert_called_once_with(0)
self.assertEqual(release_mngr._latest_key, 1)
async def test_calls_append_value(self):
file_path = self.create_tempfile()
os.remove(file_path)
release_mngr = file_release_manager.CSVFileReleaseManager(
file_path=file_path,
save_mode=file_release_manager.CSVSaveMode.APPEND)
value = {'a': 11, 'b': 21}
type_signature = computation_types.StructType([
('a', tf.int32),
('b', tf.int32),
])
with mock.patch.object(release_mngr,
'_append_value') as mock_append_value:
await release_mngr.release(value, type_signature, 1)
mock_append_value.assert_called_once_with({
'key': 1,
'a': 11,
'b': 21
})
self.assertEqual(release_mngr._latest_key, 1)
async def test_calls_write_value(self):
file_path = self.create_tempfile()
os.remove(file_path)
release_mngr = file_release_manager.CSVFileReleaseManager(
file_path=file_path,
save_mode=file_release_manager.CSVSaveMode.WRITE)
value = {'a': 11, 'b': 21}
type_signature = computation_types.StructType([
('a', tf.int32),
('b', tf.int32),
])
with mock.patch.object(release_mngr,
'_write_value') as mock_write_value:
await release_mngr.release(value, type_signature, 1)
mock_write_value.assert_called_once_with({
'key': 1,
'a': 11,
'b': 21
})
self.assertEqual(release_mngr._latest_key, 1)
# pyformat: disable
@parameterized.named_parameters(
# materialized values
('none', None, computation_types.StructType([]), [{
'key': '1',
'': ''
}]),
('bool', True, computation_types.TensorType(tf.bool), [{
'key': '1',
'': 'True'
}]),
('int', 1, computation_types.TensorType(tf.int32), [{
'key': '1',
'': '1'
}]),
('str', 'a', computation_types.TensorType(tf.string), [{
'key': '1',
'': 'a'
}]),
('tensor_int', tf.constant(1), computation_types.TensorType(
tf.int32), [{
'key': '1',
'': '1'
}]),
('tensor_str', tf.constant('a'), computation_types.TensorType(
tf.string), [{
'key': '1',
'': 'b\'a\''
}]),
('tensor_array', tf.ones([3], tf.int32),
computation_types.TensorType(tf.int32, [3]), [{
'key': '1',
'': '[1, 1, 1]'
}]),
('numpy_int', np.int32(1), computation_types.TensorType(
tf.int32), [{
'key': '1',
'': '1'
}]),
('numpy_array', np.ones([3], int),
computation_types.TensorType(tf.int32, [3]), [{
'key': '1',
'': '[1, 1, 1]'
}]),
# value references
('materializable_value_reference_tensor',
program_test_utils.TestMaterializableValueReference(1),
computation_types.TensorType(tf.int32), [{
'key': '1',
'': '1'
}]),
('materializable_value_reference_sequence',
program_test_utils.TestMaterializableValueReference(
tf.data.Dataset.from_tensor_slices([1, 2, 3])),
computation_types.SequenceType(tf.int32), [{
'key': '1',
'': '[1, 2, 3]'
}]),
# structures
('list', [
True,
program_test_utils.TestMaterializableValueReference(1), 'a'
], computation_types.SequenceType([tf.bool, tf.int32, tf.string
]), [{
'key': '1',
'0': 'True',
'1': '1',
'2': 'a'
}]),
('list_empty', [], computation_types.SequenceType([]), [{
'key': '1'
}]),
('list_nested', [
[True,
program_test_utils.TestMaterializableValueReference(1)], ['a']
], computation_types.SequenceType([[tf.bool, tf.int32], [tf.string]
]), [{
'key': '1',
'0/0': 'True',
'0/1': '1',
'1/0': 'a'
}]),
('dict', {
'a': True,
'b': program_test_utils.TestMaterializableValueReference(1),
'c': 'a'
},
computation_types.SequenceType([('a', tf.bool), ('b', tf.int32),
('c', tf.string)]), [{
'key': '1',
'a': 'True',
'b': '1',
'c': 'a'
}]),
('dict_empty', {}, computation_types.SequenceType([]), [{
'key': '1'
}]),
('dict_nested', {
'x': {
'a': True,
'b': program_test_utils.TestMaterializableValueReference(1)
},
'y': {
'c': 'a'
}
},
computation_types.SequenceType(
[('x', [('a', tf.bool), ('b', tf.int32)]),
('y', [('c', tf.string)])]), [{
'key': '1',
'x/a': 'True',
'x/b': '1',
'y/c': 'a'
}]),
('attr',
program_test_utils.TestAttrObject2(
True, program_test_utils.TestMaterializableValueReference(1)),
computation_types.SequenceType([('a', tf.bool),
('b', tf.int32)]), [{
'key': '1',
'a': 'True',
'b': '1'
}]),
('attr_nested',
program_test_utils.TestAttrObject2(
program_test_utils.TestAttrObject2(
True, program_test_utils.TestMaterializableValueReference(1)),
program_test_utils.TestAttrObject1('a')),
computation_types.SequenceType(
[('a', [('a', tf.bool), ('b', tf.int32)]),
('b', [('c', tf.string)])]), [{
'key': '1',
'a/a': 'True',
'a/b': '1',
'b/a': 'a'
}]),
)
# pyformat: enable
async def test_writes_value(self, value, type_signature, expected_value):
file_path = self.create_tempfile()
os.remove(file_path)
release_mngr = file_release_manager.CSVFileReleaseManager(
file_path=file_path)
await release_mngr.release(value, type_signature, 1)
_, actual_value = _read_values_from_csv(file_path)
self.assertEqual(actual_value, expected_value)
@parameterized.named_parameters(
('negative_1', -1),
('0', 0),
('1', 1),
)
async def test_does_not_raise_type_error_with_key(self, key):
file_path = self.create_tempfile()
os.remove(file_path)
release_mngr = file_release_manager.CSVFileReleaseManager(
file_path=file_path)
value = 1
type_signature = computation_types.TensorType(tf.int32)
try:
await release_mngr.release(value, type_signature, key)
except TypeError:
self.fail('Raised TypeError unexpectedly.')
@parameterized.named_parameters(
('none', None),
('str', 'a'),
('list', []),
)
async def test_raises_type_error_with_key(self, key):
file_path = self.create_tempfile()
os.remove(file_path)
release_mngr = file_release_manager.CSVFileReleaseManager(
file_path=file_path)
value = 1
type_signature = computation_types.TensorType(tf.int32)
with self.assertRaises(TypeError):
await release_mngr.release(value, type_signature, key)