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()
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_serialization.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_compute_returns_result(self, mock_stub):
tensor_proto = tf.make_tensor_proto(1)
any_pb = any_pb2.Any()
any_pb.Pack(tensor_proto)
value = executor_pb2.Value(tensor=any_pb)
response = executor_pb2.ComputeResponse(value=value)
instance = mock_stub.return_value
instance.Compute = mock.Mock(side_effect=[response])
loop = asyncio.get_event_loop()
executor = create_remote_executor()
type_signature = computation_types.FunctionType(None, tf.int32)
comp = remote_executor.RemoteValue(executor_pb2.ValueRef(),
type_signature, executor)
result = loop.run_until_complete(comp.compute())
instance.Compute.assert_called_once()
self.assertEqual(result, 1)
def test_compute_returns_result(self, mock_stub):
tensor_proto = tf.make_tensor_proto(1)
any_pb = any_pb2.Any()
any_pb.Pack(tensor_proto)
value = executor_pb2.Value(tensor=any_pb)
mock_stub.compute.return_value = executor_pb2.ComputeResponse(
value=value)
executor = remote_executor.RemoteExecutor(mock_stub)
_set_cardinalities_with_mock(executor, mock_stub)
executor.set_cardinalities({placements.CLIENTS: 3})
type_signature = computation_types.FunctionType(None, tf.int32)
comp = remote_executor.RemoteValue(executor_pb2.ValueRef(),
type_signature, executor)
result = asyncio.run(comp.compute())
mock_stub.compute.assert_called_once()
self.assertEqual(result, 1)
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)
with self._try_handle_request_context(request, context,
executor_pb2.ComputeResponse):
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, _ = value_serialization.serialize_value(
result_val, val_type)
return executor_pb2.ComputeResponse(value=value_proto)
def test_compute_returns_result(self, mock_executor_grpc_stub):
tensor_proto = tf.make_tensor_proto(1)
any_pb = any_pb2.Any()
any_pb.Pack(tensor_proto)
value = executor_pb2.Value(tensor=any_pb)
response = executor_pb2.ComputeResponse(value=value)
instance = mock_executor_grpc_stub.return_value
instance.Compute = mock.Mock(side_effect=[response])
request = executor_pb2.ComputeRequest(
executor=executor_pb2.ExecutorId(),
value_ref=executor_pb2.ValueRef())
stub = create_stub()
result = stub.compute(request)
instance.Compute.assert_called_once()
value, _ = value_serialization.deserialize_value(result.value)
self.assertEqual(value, 1)
