def test_executor_service_create_one_arg_computation_value_and_call(self):
env = TestEnv(eager_executor.EagerExecutor())
@computations.tf_computation(tf.int32)
def comp(x):
return tf.add(x, 1)
value_proto, _ = executor_service_utils.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_service_utils.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 Compute(self, request, context):
"""Computes a value embedded in the executor.
Args:
request: An instance of `executor_pb2.ComputeRequest`.
context: An instance of `grpc.ServicerContext`.
Returns:
An instance of `executor_pb2.ComputeResponse`.
"""
py_typecheck.check_type(request, executor_pb2.ComputeRequest)
try:
value_id = str(request.value_ref.id)
with self._lock:
future_val = self._values[value_id]
val = future_val.result()
py_typecheck.check_type(val, executor_value_base.ExecutorValue)
result = asyncio.run_coroutine_threadsafe(val.compute(),
self._event_loop)
result_val = result.result()
val_type = val.type_signature
value_proto, _ = executor_service_utils.serialize_value(
result_val, val_type)
return executor_pb2.ComputeResponse(value=value_proto)
except (ValueError, TypeError) as err:
logging.error(traceback.format_exc())
context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
context.set_details(str(err))
return executor_pb2.ComputeResponse()
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.NamedTupleType([
('a', tf.int64),
('b', tf.int32),
('c', tf.float32),
])
sequence_type = computation_types.SequenceType(element=element_type)
value_proto, value_type = executor_service_utils.serialize_value(
ds, sequence_type)
self.assertIsInstance(value_proto, executor_pb2.Value)
self.assertEqual(value_type, sequence_type)
y, type_spec = executor_service_utils.deserialize_value(value_proto)
self.assertEqual(type_spec, sequence_type)
actual_values = self.evaluate([y_val for y_val in 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_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_service_utils.serialize_value(
x, computation_types.SequenceType(tf.float32))
async def create_value(self, value, type_spec=None):
value_proto, type_spec = (executor_service_utils.serialize_value(
value, type_spec))
response = self._stub.CreateValue(
executor_pb2.CreateValueRequest(value=value_proto))
py_typecheck.check_type(response, executor_pb2.CreateValueResponse)
return RemoteValue(response.value_ref, type_spec, self)
def test_executor_service_create_tensor_value(self):
value_proto = executor_service_utils.serialize_value(
tf.constant(10.0).numpy(), tf.float32)
request = executor_pb2.CreateValueRequest(value=value_proto)
response = self._stub.CreateValue(request)
self.assertIsInstance(response, executor_pb2.CreateValueResponse)
value_id = str(response.value_ref.id)
value = self._extract_value_from_service(value_id)
self.assertEqual(value.internal_representation.numpy(), 10.0)
def test_serialize_deserialize_tensor_value_with_nontrivial_shape(self):
x = tf.constant([10, 20, 30]).numpy()
value_proto, value_type = executor_service_utils.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_service_utils.deserialize_value(value_proto)
self.assertEqual(str(type_spec), 'int32[3]')
self.assertTrue(np.array_equal(x, y))
def test_serialize_deserialize_tensor_value_with_different_dtype(self):
x = tf.constant(10.0).numpy()
value_proto, value_type = (executor_service_utils.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_service_utils.deserialize_value(value_proto)
self.assertEqual(str(type_spec), 'int32')
self.assertEqual(y, 10)
def test_serialize_deserialize_sequence_of_scalars(self):
ds = tf.data.Dataset.range(5).map(lambda x: x * 2)
value_proto, value_type = executor_service_utils.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_service_utils.deserialize_value(value_proto)
self.assertEqual(str(type_spec), 'int64*')
self.assertAllEqual([y_val for y_val in y], [x * 2 for x in range(5)])
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_service_utils.serialize_value(x, x_type) self.assertIsInstance(value_proto, executor_pb2.Value) self.assertEqual(str(value_type), '<int32,int32>') y, type_spec = executor_service_utils.deserialize_value(value_proto) self.assertEqual(str(type_spec), str(x_type)) self.assertCountEqual(y, (10, 20))
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_service_utils.serialize_value(
x, type_spec)
self.assertIsInstance(value_proto, executor_pb2.Value)
self.assertEqual(str(value_type), 'float32')
y, type_spec = executor_service_utils.deserialize_value(value_proto)
self.assertEqual(str(type_spec), 'float32')
self.assertTrue(np.array_equal(x, y))
def test_serialize_deserialize_federated_at_clients(self):
x = [10, 20]
x_type = type_factory.at_clients(tf.int32)
value_proto, value_type = executor_service_utils.serialize_value(
x, x_type)
self.assertIsInstance(value_proto, executor_pb2.Value)
self.assertEqual(str(value_type), '{int32}@CLIENTS')
y, type_spec = executor_service_utils.deserialize_value(value_proto)
self.assertEqual(str(type_spec), str(x_type))
self.assertEqual(y, [10, 20])
def test_serialize_deserialize_federated_at_server(self):
x = 10
x_type = type_factory.at_server(tf.int32)
value_proto, value_type = executor_service_utils.serialize_value(
x, x_type)
self.assertIsInstance(value_proto, executor_pb2.Value)
self.assertEqual(str(value_type), 'int32@SERVER')
y, type_spec = executor_service_utils.deserialize_value(value_proto)
self.assertEqual(str(type_spec), str(x_type))
self.assertEqual(y, 10)
def test_serialize_deserialize_computation_value(self):
@computations.tf_computation
def comp():
return tf.constant(10)
value_proto, value_type = executor_service_utils.serialize_value(comp)
self.assertEqual(value_proto.WhichOneof('value'), 'computation')
self.assertEqual(str(value_type), '( -> int32)')
comp, type_spec = executor_service_utils.deserialize_value(value_proto)
self.assertIsInstance(comp, computation_pb2.Computation)
self.assertEqual(str(type_spec), '( -> int32)')
def test_executor_service_create_tensor_value(self):
env = TestEnv(eager_executor.EagerExecutor())
value_proto, _ = executor_service_utils.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_executor_service_create_and_select_from_tuple(self):
env = TestEnv(eager_executor.EagerExecutor())
value_proto, _ = executor_service_utils.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_service_utils.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.CreateTuple(
executor_pb2.CreateTupleRequest(element=[
executor_pb2.CreateTupleRequest.Element(name='a',
value_ref=ten_ref),
executor_pb2.CreateTupleRequest.Element(name='b',
value_ref=twenty_ref)
]))
self.assertIsInstance(response, executor_pb2.CreateTupleResponse)
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
async def create_value(self, value, type_spec=None):
value_proto, type_spec = (
executor_service_utils.serialize_value(value, type_spec))
create_value_request = executor_pb2.CreateValueRequest(value=value_proto)
if not self._bidi_stream:
response = self._stub.CreateValue(create_value_request)
else:
response = (await self._bidi_stream.send_request(
executor_pb2.ExecuteRequest(create_value=create_value_request)
)).create_value
py_typecheck.check_type(response, executor_pb2.CreateValueResponse)
return RemoteValue(response.value_ref, type_spec, self)
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.NamedTupleType([
('b', tf.int32),
('a',
computation_types.NamedTupleType([
(None, tf.int64),
(None, test_tuple_type(tf.int64, tf.int64)),
(None,
computation_types.NamedTupleType([('x', tf.int64),
('y', tf.int64)])),
])),
])
sequence_type = computation_types.SequenceType(element=element_type)
value_proto, value_type = executor_service_utils.serialize_value(
ds, sequence_type)
self.assertIsInstance(value_proto, executor_pb2.Value)
self.assertEqual(value_type, sequence_type)
y, type_spec = executor_service_utils.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([y_val for y_val in 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_executor_service_create_computation_value(self):
@computations.tf_computation
def comp():
return tf.constant(10)
value_proto = executor_service_utils.serialize_value(comp)
request = executor_pb2.CreateValueRequest(value=value_proto)
response = self._stub.CreateValue(request)
self.assertIsInstance(response, executor_pb2.CreateValueResponse)
value_id = str(response.value_ref.id)
value = self._extract_value_from_service(value_id)
self.assertTrue(callable(value.internal_representation))
self.assertEqual(value.internal_representation().numpy(), 10.0)
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_service_utils.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_service_utils.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_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_tuple(self, elements):
raise NotImplementedError
async def create_selection(self, source, index=None, name=None):
raise NotImplementedError
def close(self):
pass
ex = SlowExecutor()
env = TestEnv(ex)
self.assertEqual(ex.status, 'idle')
value_proto, _ = executor_service_utils.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_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_service_utils.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_service_utils.deserialize_value(value_proto)
self.assertEqual(str(type_spec), '<int64,int32,float32>*')
self.assertAllEqual(self.evaluate([y_val for y_val in y]),
[(x * 2, x, x - 1.) for x in range(5)])
def test_executor_service_create_no_arg_computation_value_and_call(self):
env = TestEnv(eager_executor.EagerExecutor())
@computations.tf_computation
def comp():
return tf.constant(10)
value_proto, _ = executor_service_utils.serialize_value(comp)
response = env.stub.CreateValue(
executor_pb2.CreateValueRequest(value=value_proto))
self.assertIsInstance(response, executor_pb2.CreateValueResponse)
response = env.stub.CreateCall(
executor_pb2.CreateCallRequest(function_ref=response.value_ref))
self.assertIsInstance(response, executor_pb2.CreateCallResponse)
value_id = str(response.value_ref.id)
value = env.get_value(value_id)
self.assertEqual(value, 10)
del env
async def _Compute(
self,
request: executor_pb2.ComputeRequest,
context: grpc.ServicerContext,
) -> executor_pb2.ComputeResponse:
"""Asynchronous implemention of `Compute`."""
py_typecheck.check_type(request, executor_pb2.ComputeRequest)
try:
value_id = str(request.value_ref.id)
with self._lock:
future_val = asyncio.wrap_future(self._values[value_id])
val = await future_val
result_val = await val.compute()
val_type = val.type_signature
value_proto, _ = executor_service_utils.serialize_value(
result_val, val_type)
return executor_pb2.ComputeResponse(value=value_proto)
except (ValueError, TypeError) as err:
_set_invalid_arg_err(context, err)
return executor_pb2.ComputeResponse()
def test_executor_service_value_unavailable_after_dispose(self):
env = TestEnv(eager_tf_executor.EagerTFExecutor())
value_proto, _ = executor_service_utils.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)
self.assertIsInstance(response, executor_pb2.DisposeResponse)
# Check that the value is gone from the _values map
# get_value_future_directly is used here so that we can catch the
# exception rather than having it occur on the GRPC thread.
with self.assertRaises(KeyError):
env.get_value_future_directly(value_id)
def test_serialize_sequence_not_a_dataset(self):
with self.assertRaisesRegex(
TypeError,
r'Cannot serialize Python type int as .* float32\*'):
_ = executor_service_utils.serialize_value(
5, computation_types.SequenceType(tf.float32))
def test_serialize_deserialize_tensor_value_with_bad_shape(self):
x = tf.constant([10, 20, 30]).numpy()
with self.assertRaises(TypeError):
executor_service_utils.serialize_value(x, tf.int32)
def serialize_value():
return executor_service_utils.serialize_value(value, type_spec)
