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_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_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))
async def _compute(self, value_ref):
py_typecheck.check_type(value_ref, executor_pb2.ValueRef)
request = executor_pb2.ComputeRequest(value_ref=value_ref)
response = _request(self._stub.Compute, request)
py_typecheck.check_type(response, executor_pb2.ComputeResponse)
value, _ = executor_serialization.deserialize_value(response.value)
return value
def get_value(self, value_id: str):
"""Retrieves a value using the `Compute` endpoint."""
response = self._stub.Compute(
executor_pb2.ComputeRequest(
value_ref=executor_pb2.ValueRef(id=value_id)))
py_typecheck.check_type(response, executor_pb2.ComputeResponse)
value, _ = executor_serialization.deserialize_value(response.value)
return value
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_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_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_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_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_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_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)')
async def _compute(self, value_ref):
py_typecheck.check_type(value_ref, executor_pb2.ValueRef)
request = executor_pb2.ComputeRequest(value_ref=value_ref)
if self._bidi_stream is None:
response = _request(self._stub.Compute, request)
else:
response = (await self._bidi_stream.send_request(
executor_pb2.ExecuteRequest(compute=request))).compute
py_typecheck.check_type(response, executor_pb2.ComputeResponse)
value, _ = executor_serialization.deserialize_value(response.value)
return value
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_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_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_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_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_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>>')
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.to_type(
collections.OrderedDict(b=tf.int32,
a=tuple([
tf.int64,
test_tuple_type(tf.int64, tf.int64),
collections.OrderedDict(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.assert_types_identical(value_type, sequence_type)
y, type_spec = executor_serialization.deserialize_value(value_proto)
# These aren't the same because ser/de destroys the PyContainer
self.assert_types_equivalent(type_spec, 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 = 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 CreateValue(
self,
request: executor_pb2.CreateValueRequest,
context: grpc.ServicerContext,
) -> executor_pb2.CreateValueResponse:
"""Creates a value embedded in the executor."""
py_typecheck.check_type(request, executor_pb2.CreateValueRequest)
try:
with tracing.span('ExecutorService.CreateValue', 'deserialize_value'):
value, value_type = (
executor_serialization.deserialize_value(request.value))
value_id = str(uuid.uuid4())
coro = self.executor.create_value(value, value_type)
future_val = self._run_coro_threadsafe_with_tracing(coro)
with self._lock:
self._values[value_id] = future_val
return executor_pb2.CreateValueResponse(
value_ref=executor_pb2.ValueRef(id=value_id))
except (ValueError, TypeError) as err:
_set_invalid_arg_err(context, err)
return executor_pb2.CreateValueResponse()
def test_deserialize_federated_value_with_unset_member_type(self):
x = 10
x_type = computation_types.to_type(tf.int32)
member_proto, _ = executor_serialization.serialize_value(x, x_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=fully_specified_type_at_clients.federated.all_equal)
federated_proto = executor_pb2.Value.Federated(
type=unspecified_member_federated_type, value=[member_proto])
federated_value_proto = executor_pb2.Value(federated=federated_proto)
self.assertIsInstance(member_proto, executor_pb2.Value)
self.assertIsInstance(federated_value_proto, executor_pb2.Value)
deserialized_federated_value, deserialized_type_spec = executor_serialization.deserialize_value(
federated_value_proto)
self.assert_types_identical(deserialized_type_spec,
computation_types.at_clients(tf.int32))
self.assertEqual(deserialized_federated_value, [10])
