def CreateCall(
self,
request: executor_pb2.CreateCallRequest,
context: grpc.ServicerContext,
) -> executor_pb2.CreateCallResponse:
"""Creates a call embedded in the executor."""
py_typecheck.check_type(request, executor_pb2.CreateCallRequest)
try:
function_id = str(request.function_ref.id)
argument_id = str(request.argument_ref.id)
with self._lock:
function_val = self._values[function_id]
argument_val = self._values[
argument_id] if argument_id else None
async def _processing():
function = await asyncio.wrap_future(function_val)
argument = await asyncio.wrap_future(
argument_val) if argument_val is not None else None
return await self.executor.create_call(function, argument)
coro = _processing()
result_fut = self._run_coro_threadsafe_with_tracing(coro)
result_id = str(uuid.uuid4())
with self._lock:
self._values[result_id] = result_fut
return executor_pb2.CreateCallResponse(
value_ref=executor_pb2.ValueRef(id=result_id))
except (ValueError, TypeError) as err:
_set_invalid_arg_err(context, err)
return executor_pb2.CreateCallResponse()
def CreateCall(self, request, context):
"""Creates a call embedded in the executor.
Args:
request: An instance of `executor_pb2.CreateCallRequest`.
context: An instance of `grpc.ServicerContext`.
Returns:
An instance of `executor_pb2.CreateCallResponse`.
"""
py_typecheck.check_type(request, executor_pb2.CreateCallRequest)
try:
function_id = str(request.function_ref.id)
argument_id = str(request.argument_ref.id)
with self._lock:
function_val = self._values[function_id]
argument_val = self._values[
argument_id] if argument_id else None
function = function_val.result()
argument = argument_val.result(
) if argument_val is not None else None
result_val = asyncio.run_coroutine_threadsafe(
self._executor.create_call(function, argument),
self._event_loop)
result_id = str(uuid.uuid4())
with self._lock:
self._values[result_id] = result_val
return executor_pb2.CreateCallResponse(
value_ref=executor_pb2.ValueRef(id=result_id))
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.CreateCallResponse()
def test_create_call_returns_remote_value(self, mock_executor_grpc_stub):
response = executor_pb2.CreateCallResponse()
instance = mock_executor_grpc_stub.return_value
instance.CreateCall = mock.Mock(side_effect=[response])
stub = create_stub()
result = stub.create_call(request=executor_pb2.CreateCallRequest())
instance.CreateCall.assert_called_once()
self.assertEqual(result, response)
def test_create_call_returns_remote_value(self, mock_stub):
mock_stub.create_call.return_value = executor_pb2.CreateCallResponse()
executor = remote_executor.RemoteExecutor(mock_stub)
_set_cardinalities_with_mock(executor, mock_stub)
type_signature = computation_types.FunctionType(None, tf.int32)
fn = remote_executor.RemoteValue(executor_pb2.ValueRef(),
type_signature, executor)
result = asyncio.run(executor.create_call(fn, None))
mock_stub.create_call.assert_called_once()
self.assertIsInstance(result, remote_executor.RemoteValue)
def test_create_call_returns_remote_value(self, mock_stub):
response = executor_pb2.ExecuteResponse(
create_call=executor_pb2.CreateCallResponse())
executor = _setup_mock_streaming_executor(mock_stub, response)
loop = asyncio.get_event_loop()
type_signature = computation_types.FunctionType(None, tf.int32)
fn = remote_executor.RemoteValue(executor_pb2.ValueRef(),
type_signature, executor)
result = loop.run_until_complete(executor.create_call(fn, None))
self.assertIsInstance(result, remote_executor.RemoteValue)
def test_create_call_returns_remote_value(self, mock_stub):
response = executor_pb2.CreateCallResponse()
instance = mock_stub.return_value
instance.CreateCall = mock.Mock(side_effect=[response])
loop = asyncio.get_event_loop()
executor = create_remote_executor()
type_signature = computation_types.FunctionType(None, tf.int32)
fn = remote_executor.RemoteValue(executor_pb2.ValueRef(),
type_signature, executor)
result = loop.run_until_complete(executor.create_call(fn, None))
instance.CreateCall.assert_called_once()
self.assertIsInstance(result, remote_executor.RemoteValue)
