def test_deserialize_federated_value_promotes_types(self): x = [10] smaller_type = computation_types.StructType([ (None, computation_types.to_type(tf.int32)) ]) smaller_type_member_proto, _ = executor_serialization.serialize_value( x, smaller_type) larger_type = computation_types.StructType([ ('a', computation_types.to_type(tf.int32)) ]) larger_type_member_proto, _ = executor_serialization.serialize_value( x, larger_type) type_at_clients = type_serialization.serialize_type( computation_types.at_clients(tf.int32)) unspecified_member_federated_type = computation_pb2.FederatedType( placement=type_at_clients.federated.placement, all_equal=False) federated_proto = executor_pb2.Value.Federated( type=unspecified_member_federated_type, value=[larger_type_member_proto, smaller_type_member_proto]) federated_value_proto = executor_pb2.Value(federated=federated_proto) self.assertIsInstance(smaller_type_member_proto, executor_pb2.Value) self.assertIsInstance(larger_type_member_proto, executor_pb2.Value) self.assertIsInstance(federated_value_proto, executor_pb2.Value) _, deserialized_type_spec = executor_serialization.deserialize_value( federated_value_proto) self.assert_types_identical(deserialized_type_spec, computation_types.at_clients(larger_type))
def test_executor_service_create_one_arg_computation_value_and_call(self): ex_factory = executor_stacks.ResourceManagingExecutorFactory( lambda _: eager_tf_executor.EagerTFExecutor()) env = TestEnv(ex_factory) @computations.tf_computation(tf.int32) def comp(x): return tf.add(x, 1) value_proto, _ = executor_serialization.serialize_value(comp) response = env.stub.CreateValue( executor_pb2.CreateValueRequest(value=value_proto)) self.assertIsInstance(response, executor_pb2.CreateValueResponse) comp_ref = response.value_ref value_proto, _ = executor_serialization.serialize_value(10, tf.int32) response = env.stub.CreateValue( executor_pb2.CreateValueRequest(value=value_proto)) self.assertIsInstance(response, executor_pb2.CreateValueResponse) arg_ref = response.value_ref response = env.stub.CreateCall( executor_pb2.CreateCallRequest( function_ref=comp_ref, argument_ref=arg_ref)) self.assertIsInstance(response, executor_pb2.CreateCallResponse) value_id = str(response.value_ref.id) value = env.get_value(value_id) self.assertEqual(value, 11) del env
def test_deserialize_federated_value_with_incompatible_member_types_raises( self): x = 10 x_type = computation_types.to_type(tf.int32) int_member_proto, _ = executor_serialization.serialize_value(x, x_type) y = 10. y_type = computation_types.to_type(tf.float32) float_member_proto, _ = executor_serialization.serialize_value(y, y_type) fully_specified_type_at_clients = type_serialization.serialize_type( computation_types.at_clients(tf.int32)) unspecified_member_federated_type = computation_pb2.FederatedType( placement=fully_specified_type_at_clients.federated.placement, all_equal=False) federated_proto = executor_pb2.Value.Federated( type=unspecified_member_federated_type, value=[int_member_proto, float_member_proto]) federated_value_proto = executor_pb2.Value(federated=federated_proto) self.assertIsInstance(int_member_proto, executor_pb2.Value) self.assertIsInstance(float_member_proto, executor_pb2.Value) self.assertIsInstance(federated_value_proto, executor_pb2.Value) with self.assertRaises(TypeError): executor_serialization.deserialize_value(federated_value_proto)
def test_dispose_does_not_trigger_cleanup(self): class MockFactory(executor_factory.ExecutorFactory, mock.MagicMock): def create_executor(self, *args, **kwargs): return mock.MagicMock() def clean_up_executors(self): return ex_factory = MockFactory() ex_factory.clean_up_executors = mock.MagicMock() env = TestEnv(ex_factory) value_proto, _ = executor_serialization.serialize_value( tf.constant(10.0).numpy(), tf.float32) # Create the value response = env.stub.CreateValue( executor_pb2.CreateValueRequest(value=value_proto)) self.assertIsInstance(response, executor_pb2.CreateValueResponse) value_id = str(response.value_ref.id) # Check that the value appears in the _values map env.get_value_future_directly(value_id) # Dispose of the value dispose_request = executor_pb2.DisposeRequest() dispose_request.value_ref.append(response.value_ref) response = env.stub.Dispose(dispose_request) # We shouldn't be propagating close down the executor stack on Dispose--this # would close the bidi stream and cause a hang in the streaming case with # intermediate aggregation. Python GC takes care of pushing Dispose requests # from the aggregators to the workers. ex_factory.clean_up_executors.assert_not_called()
def test_serialize_sequence_bad_element_type(self): x = tf.data.Dataset.range(5).map(lambda x: x * 2) with self.assertRaisesRegex( TypeError, r'Cannot serialize dataset .* int64\* .* float32\*.*'): _ = executor_serialization.serialize_value( x, computation_types.SequenceType(tf.float32))
def test_serialize_deserialize_sequence_of_namedtuples(self): test_tuple_type = collections.namedtuple('TestTuple', ['a', 'b', 'c']) def make_test_tuple(x): return test_tuple_type(a=x * 2, b=tf.cast(x, tf.int32), c=tf.cast(x - 1, tf.float32)) ds = tf.data.Dataset.range(5).map(make_test_tuple) element_type = computation_types.StructType([ ('a', tf.int64), ('b', tf.int32), ('c', tf.float32), ]) sequence_type = computation_types.SequenceType(element=element_type) value_proto, value_type = executor_serialization.serialize_value( ds, sequence_type) self.assertIsInstance(value_proto, executor_pb2.Value) self.assertEqual(value_type, sequence_type) y, type_spec = executor_serialization.deserialize_value(value_proto) self.assertEqual(type_spec, sequence_type) actual_values = self.evaluate(list(y)) expected_values = [ test_tuple_type(a=x * 2, b=x, c=x - 1.) for x in range(5) ] for actual, expected in zip(actual_values, expected_values): self.assertAllClose(actual, expected)
def test_serialize_deserialize_tensor_value_with_different_dtype(self): x = tf.constant(10.0).numpy() value_proto, value_type = (executor_serialization.serialize_value( x, computation_types.TensorType(tf.int32))) self.assertIsInstance(value_proto, executor_pb2.Value) self.assertEqual(str(value_type), 'int32') y, type_spec = executor_serialization.deserialize_value(value_proto) self.assertEqual(str(type_spec), 'int32') self.assertEqual(y, 10)
def test_serialize_deserialize_nested_tuple_value_without_names(self): x = (10, 20) x_type = computation_types.to_type((tf.int32, tf.int32)) value_proto, value_type = executor_serialization.serialize_value(x, x_type) self.assertIsInstance(value_proto, executor_pb2.Value) self.assert_types_identical(value_type, x_type) y, type_spec = executor_serialization.deserialize_value(value_proto) self.assert_types_equivalent(type_spec, x_type) self.assertEqual(y, structure.from_container((10, 20)))
def test_serialize_deserialize_sequence_of_scalars(self): ds = tf.data.Dataset.range(5).map(lambda x: x * 2) value_proto, value_type = executor_serialization.serialize_value( ds, computation_types.SequenceType(tf.int64)) self.assertIsInstance(value_proto, executor_pb2.Value) self.assertEqual(str(value_type), 'int64*') y, type_spec = executor_serialization.deserialize_value(value_proto) self.assertEqual(str(type_spec), 'int64*') self.assertAllEqual(list(y), [x * 2 for x in range(5)])
def test_serialize_deserialize_tensor_value_with_nontrivial_shape(self): x = tf.constant([10, 20, 30]).numpy() value_proto, value_type = executor_serialization.serialize_value( x, computation_types.TensorType(tf.int32, [3])) self.assertIsInstance(value_proto, executor_pb2.Value) self.assertEqual(str(value_type), 'int32[3]') y, type_spec = executor_serialization.deserialize_value(value_proto) self.assertEqual(str(type_spec), 'int32[3]') self.assertTrue(np.array_equal(x, y))
def test_serialize_deserialize_tensor_value(self): x = tf.constant(10.0).numpy() type_spec = computation_types.TensorType(tf.as_dtype(x.dtype), x.shape) value_proto, value_type = executor_serialization.serialize_value( x, type_spec) self.assertIsInstance(value_proto, executor_pb2.Value) self.assertEqual(str(value_type), 'float32') y, type_spec = executor_serialization.deserialize_value(value_proto) self.assertEqual(str(type_spec), 'float32') self.assertTrue(np.array_equal(x, y))
def test_serialize_deserialize_federated_at_server(self): x = 10 x_type = computation_types.at_server(tf.int32) value_proto, value_type = executor_serialization.serialize_value( x, x_type) self.assertIsInstance(value_proto, executor_pb2.Value) self.assertEqual(str(value_type), 'int32@SERVER') y, type_spec = executor_serialization.deserialize_value(value_proto) self.assertEqual(str(type_spec), str(x_type)) self.assertEqual(y, 10)
def test_serialize_deserialize_federated_at_clients(self): x = [10, 20] x_type = computation_types.at_clients(tf.int32) value_proto, value_type = executor_serialization.serialize_value( x, x_type) self.assertIsInstance(value_proto, executor_pb2.Value) self.assertEqual(str(value_type), '{int32}@CLIENTS') y, type_spec = executor_serialization.deserialize_value(value_proto) self.assertEqual(str(type_spec), str(x_type)) self.assertEqual(y, [10, 20])
def test_serialize_deserialize_nested_tuple_value_without_names(self): x = tuple([10, 20]) x_type = computation_types.to_type(tuple([tf.int32, tf.int32])) value_proto, value_type = executor_serialization.serialize_value( x, x_type) self.assertIsInstance(value_proto, executor_pb2.Value) self.assertEqual(str(value_type), '<int32,int32>') y, type_spec = executor_serialization.deserialize_value(value_proto) self.assertEqual(str(type_spec), str(x_type)) self.assertCountEqual(y, (10, 20))
def test_serialize_deserialize_federated_at_server(self): x = 10 x_type = computation_types.at_server(tf.int32) value_proto, value_type = executor_serialization.serialize_value(x, x_type) self.assertIsInstance(value_proto, executor_pb2.Value) self.assert_types_identical(value_type, computation_types.at_server(tf.int32)) y, type_spec = executor_serialization.deserialize_value(value_proto) self.assert_types_identical(type_spec, x_type) self.assertEqual(y, 10)
def test_executor_service_raises_after_cleanup_without_configuration(self): ex_factory = executor_stacks.ResourceManagingExecutorFactory( lambda _: eager_tf_executor.EagerTFExecutor()) env = TestEnv(ex_factory) env.stub.ClearExecutor(executor_pb2.ClearExecutorRequest()) value_proto, _ = executor_serialization.serialize_value( tf.constant(10.0).numpy(), tf.float32) with self.assertRaises(grpc.RpcError): env.stub.CreateValue( executor_pb2.CreateValueRequest(value=value_proto))
def _iterator(): @computations.tf_computation(tf.int32) def comp(x): return tf.add(x, 1) value_proto, _ = executor_serialization.serialize_value(comp) request = executor_pb2.ExecuteRequest( create_value=executor_pb2.CreateValueRequest( value=value_proto)) yield request
def test_executor_service_create_and_select_from_tuple(self): ex_factory = executor_stacks.ResourceManagingExecutorFactory( lambda _: eager_tf_executor.EagerTFExecutor()) env = TestEnv(ex_factory) value_proto, _ = executor_serialization.serialize_value(10, tf.int32) response = env.stub.CreateValue( executor_pb2.CreateValueRequest(value=value_proto)) self.assertIsInstance(response, executor_pb2.CreateValueResponse) ten_ref = response.value_ref self.assertEqual(env.get_value(ten_ref.id), 10) value_proto, _ = executor_serialization.serialize_value(20, tf.int32) response = env.stub.CreateValue( executor_pb2.CreateValueRequest(value=value_proto)) self.assertIsInstance(response, executor_pb2.CreateValueResponse) twenty_ref = response.value_ref self.assertEqual(env.get_value(twenty_ref.id), 20) response = env.stub.CreateStruct( executor_pb2.CreateStructRequest(element=[ executor_pb2.CreateStructRequest.Element(name='a', value_ref=ten_ref), executor_pb2.CreateStructRequest.Element(name='b', value_ref=twenty_ref) ])) self.assertIsInstance(response, executor_pb2.CreateStructResponse) tuple_ref = response.value_ref self.assertEqual(str(env.get_value(tuple_ref.id)), '<a=10,b=20>') for arg_name, arg_val, result_val in [('name', 'a', 10), ('name', 'b', 20), ('index', 0, 10), ('index', 1, 20)]: response = env.stub.CreateSelection( executor_pb2.CreateSelectionRequest(source_ref=tuple_ref, **{arg_name: arg_val})) self.assertIsInstance(response, executor_pb2.CreateSelectionResponse) selection_ref = response.value_ref self.assertEqual(env.get_value(selection_ref.id), result_val) del env
def test_serialize_deserialize_tensor_value_with_nontrivial_shape(self): x = tf.constant([10, 20, 30]) value_proto, value_type = executor_serialization.serialize_value( x, computation_types.TensorType(tf.int32, [3])) self.assertIsInstance(value_proto, executor_pb2.Value) self.assert_types_identical( value_type, computation_types.TensorType(tf.int32, [3])) y, type_spec = executor_serialization.deserialize_value(value_proto) self.assert_types_identical( type_spec, computation_types.TensorType(tf.int32, [3])) self.assertAllEqual(x, y)
def test_serialize_deserialize_tensor_value_with_different_dtype(self): x = tf.constant(10.0) value_proto, value_type = executor_serialization.serialize_value( x, computation_types.TensorType(tf.int32)) self.assertIsInstance(value_proto, executor_pb2.Value) self.assert_types_identical(value_type, computation_types.TensorType(tf.int32)) y, type_spec = executor_serialization.deserialize_value(value_proto) self.assert_types_identical(type_spec, computation_types.TensorType(tf.int32)) self.assertEqual(y, 10)
def test_executor_service_slowly_create_tensor_value(self): class SlowExecutorValue(executor_value_base.ExecutorValue): def __init__(self, v, t): self._v = v self._t = t @property def type_signature(self): return self._t async def compute(self): return self._v class SlowExecutor(executor_base.Executor): def __init__(self): self.status = 'idle' self.busy = threading.Event() self.done = threading.Event() async def create_value(self, value, type_spec=None): self.status = 'busy' self.busy.set() self.done.wait() self.status = 'done' return SlowExecutorValue(value, type_spec) async def create_call(self, comp, arg=None): raise NotImplementedError async def create_struct(self, elements): raise NotImplementedError async def create_selection(self, source, index=None, name=None): raise NotImplementedError def close(self): pass ex = SlowExecutor() ex_factory = executor_stacks.ResourceManagingExecutorFactory(lambda _: ex) env = TestEnv(ex_factory) self.assertEqual(ex.status, 'idle') value_proto, _ = executor_serialization.serialize_value(10, tf.int32) response = env.stub.CreateValue( executor_pb2.CreateValueRequest(value=value_proto)) ex.busy.wait() self.assertEqual(ex.status, 'busy') ex.done.set() value = env.get_value(response.value_ref.id) self.assertEqual(ex.status, 'done') self.assertEqual(value, 10)
def test_serialize_deserialize_computation_value(self): @computations.tf_computation def comp(): return tf.constant(10) value_proto, value_type = executor_serialization.serialize_value(comp) self.assertEqual(value_proto.WhichOneof('value'), 'computation') self.assertEqual(str(value_type), '( -> int32)') comp, type_spec = executor_serialization.deserialize_value(value_proto) self.assertIsInstance(comp, computation_pb2.Computation) self.assertEqual(str(type_spec), '( -> int32)')
def test_serialize_deserialize_sequence_of_nested_structures(self): test_tuple_type = collections.namedtuple('TestTuple', ['u', 'v']) def _make_nested_tf_structure(x): return collections.OrderedDict([ ('b', tf.cast(x, tf.int32)), ('a', tuple([ x, test_tuple_type(x * 2, x * 3), collections.OrderedDict([('x', x**2), ('y', x**3)]) ])), ]) ds = tf.data.Dataset.range(5).map(_make_nested_tf_structure) element_type = computation_types.StructType([ ('b', tf.int32), ('a', computation_types.StructType([ (None, tf.int64), (None, test_tuple_type(tf.int64, tf.int64)), (None, computation_types.StructType([('x', tf.int64), ('y', tf.int64)])), ])), ]) sequence_type = computation_types.SequenceType(element=element_type) value_proto, value_type = executor_serialization.serialize_value( ds, sequence_type) self.assertIsInstance(value_proto, executor_pb2.Value) self.assertEqual(value_type, sequence_type) y, type_spec = executor_serialization.deserialize_value(value_proto) # These aren't the same because ser/de destroys the PyContainer type_spec.check_equivalent_to(sequence_type) def _build_expected_structure(x): return collections.OrderedDict([ ('b', x), ('a', tuple([ x, test_tuple_type(x * 2, x * 3), collections.OrderedDict([('x', x**2), ('y', x**3)]) ])), ]) actual_values = self.evaluate(list(y)) expected_values = [_build_expected_structure(x) for x in range(5)] for actual, expected in zip(actual_values, expected_values): self.assertEqual(type(actual), type(expected)) self.assertAllClose(actual, expected)
def test_serialize_deserialize_tensor_value(self): x = tf.constant(10.0) type_spec = computation_types.TensorType(tf.as_dtype(x.dtype), x.shape) value_proto, value_type = executor_serialization.serialize_value( x, type_spec) self.assertIsInstance(value_proto, executor_pb2.Value) self.assert_types_identical(value_type, computation_types.TensorType(tf.float32)) y, type_spec = executor_serialization.deserialize_value(value_proto) self.assert_types_identical(type_spec, computation_types.TensorType(tf.float32)) self.assertAllEqual(x, y)
def test_executor_service_create_tensor_value(self): ex_factory = executor_stacks.ResourceManagingExecutorFactory( lambda _: eager_tf_executor.EagerTFExecutor()) env = TestEnv(ex_factory) value_proto, _ = executor_serialization.serialize_value( tf.constant(10.0).numpy(), tf.float32) response = env.stub.CreateValue( executor_pb2.CreateValueRequest(value=value_proto)) self.assertIsInstance(response, executor_pb2.CreateValueResponse) value_id = str(response.value_ref.id) value = env.get_value(value_id) self.assertEqual(value, 10.0) del env
def test_serialize_deserialize_computation_value(self): @computations.tf_computation def comp(): return tf.constant(10) value_proto, value_type = executor_serialization.serialize_value(comp) self.assertEqual(value_proto.WhichOneof('value'), 'computation') self.assert_types_identical( value_type, computation_types.FunctionType(parameter=None, result=tf.int32)) comp, type_spec = executor_serialization.deserialize_value(value_proto) self.assertIsInstance(comp, computation_pb2.Computation) self.assert_types_identical( type_spec, computation_types.FunctionType(parameter=None, result=tf.int32))
def test_serialize_deserialize_sequence_of_tuples(self): ds = tf.data.Dataset.range(5).map(lambda x: (x * 2, tf.cast( x, tf.int32), tf.cast(x - 1, tf.float32))) value_proto, value_type = executor_serialization.serialize_value( ds, computation_types.SequenceType(element=(tf.int64, tf.int32, tf.float32))) self.assertIsInstance(value_proto, executor_pb2.Value) self.assertEqual(str(value_type), '<int64,int32,float32>*') y, type_spec = executor_serialization.deserialize_value(value_proto) self.assertEqual(str(type_spec), '<int64,int32,float32>*') self.assertAllEqual(self.evaluate(list(y)), [(x * 2, x, x - 1.) for x in range(5)])
def test_serialize_deserialize_nested_tuple_value_with_names(self): x = collections.OrderedDict( a=10, b=[20, 30], c=collections.OrderedDict(d=40)) x_type = computation_types.to_type( collections.OrderedDict( a=tf.int32, b=[tf.int32, tf.int32], c=collections.OrderedDict(d=tf.int32))) value_proto, value_type = executor_serialization.serialize_value(x, x_type) self.assertIsInstance(value_proto, executor_pb2.Value) self.assert_types_identical(value_type, x_type) y, type_spec = executor_serialization.deserialize_value(value_proto) # Don't assert on the Python container since it is lost in serialization. self.assert_types_equivalent(type_spec, x_type) self.assertEqual(y, structure.from_container(x, recursive=True))
def test_serialize_deserialize_sequence_of_tuples(self): ds = tf.data.Dataset.range(5).map(lambda x: (x * 2, tf.cast( x, tf.int32), tf.cast(x - 1, tf.float32))) value_proto, value_type = executor_serialization.serialize_value( ds, computation_types.SequenceType(element=(tf.int64, tf.int32, tf.float32))) expected_type = computation_types.SequenceType( (tf.int64, tf.int32, tf.float32)) self.assertIsInstance(value_proto, executor_pb2.Value) self.assert_types_identical(value_type, expected_type) y, type_spec = executor_serialization.deserialize_value(value_proto) # Only checking for equivalence, we don't have the Python container # after deserialization. self.assert_types_equivalent(type_spec, expected_type) self.assertAllEqual(list(y), [(x * 2, x, x - 1.) for x in range(5)])
def test_serialize_deserialize_nested_tuple_value_with_names(self): x = collections.OrderedDict([('a', 10), ('b', [20, 30]), ('c', collections.OrderedDict([('d', 40) ]))]) x_type = computation_types.to_type( collections.OrderedDict([ ('a', tf.int32), ('b', [tf.int32, tf.int32]), ('c', collections.OrderedDict([('d', tf.int32)])) ])) value_proto, value_type = executor_serialization.serialize_value( x, x_type) self.assertIsInstance(value_proto, executor_pb2.Value) self.assertEqual(str(value_type), '<a=int32,b=<int32,int32>,c=<d=int32>>') y, type_spec = executor_serialization.deserialize_value(value_proto) self.assertEqual(str(type_spec), str(x_type)) self.assertTrue(str(y), '<a=10,b=<20,30>,c=<d=40>>')