def CreateStruct(
self,
request: executor_pb2.CreateStructRequest,
context: grpc.ServicerContext,
) -> executor_pb2.CreateStructResponse:
"""Creates a struct embedded in the executor."""
py_typecheck.check_type(request, executor_pb2.CreateStructRequest)
try:
with self._lock:
elem_futures = [
self._values[e.value_ref.id] for e in request.element
]
elem_names = [
str(elem.name) if elem.name else None
for elem in request.element
]
async def _processing():
elem_values = await asyncio.gather(
*[asyncio.wrap_future(v) for v in elem_futures])
elements = list(zip(elem_names, elem_values))
struct = structure.Struct(elements)
return await self.executor.create_struct(struct)
result_fut = self._run_coro_threadsafe_with_tracing(_processing())
result_id = str(uuid.uuid4())
with self._lock:
self._values[result_id] = result_fut
return executor_pb2.CreateStructResponse(
value_ref=executor_pb2.ValueRef(id=result_id))
except (ValueError, TypeError) as err:
_set_invalid_arg_err(context, err)
return executor_pb2.CreateStructResponse()
def test_create_struct_returns_value(self, mock_executor_grpc_stub):
response = executor_pb2.CreateStructResponse()
instance = mock_executor_grpc_stub.return_value
instance.CreateStruct = mock.Mock(side_effect=[response])
stub = create_stub()
result = stub.create_struct(request=executor_pb2.CreateStructRequest())
instance.CreateStruct.assert_called_once()
self.assertEqual(result, response)
def test_create_struct_returns_remote_value(self, mock_stub):
mock_stub.create_struct.return_value = executor_pb2.CreateStructResponse(
)
executor = remote_executor.RemoteExecutor(mock_stub)
_set_cardinalities_with_mock(executor, mock_stub)
type_signature = computation_types.TensorType(tf.int32)
value_1 = remote_executor.RemoteValue(executor_pb2.ValueRef(),
type_signature, executor)
value_2 = remote_executor.RemoteValue(executor_pb2.ValueRef(),
type_signature, executor)
result = asyncio.run(executor.create_struct([value_1, value_2]))
mock_stub.create_struct.assert_called_once()
self.assertIsInstance(result, remote_executor.RemoteValue)
def test_create_struct_returns_remote_value(self, mock_stub):
response = executor_pb2.CreateStructResponse()
instance = mock_stub.return_value
instance.CreateStruct = mock.Mock(side_effect=[response])
loop = asyncio.get_event_loop()
executor = create_remote_executor()
type_signature = computation_types.TensorType(tf.int32)
value_1 = remote_executor.RemoteValue(executor_pb2.ValueRef(),
type_signature, executor)
value_2 = remote_executor.RemoteValue(executor_pb2.ValueRef(),
type_signature, executor)
result = loop.run_until_complete(executor.create_struct([value_1, value_2]))
instance.CreateStruct.assert_called_once()
self.assertIsInstance(result, remote_executor.RemoteValue)
def test_create_struct_returns_remote_value(self, mock_stub):
response = executor_pb2.ExecuteResponse(
create_struct=executor_pb2.CreateStructResponse())
executor = _setup_mock_streaming_executor(mock_stub, response)
loop = asyncio.get_event_loop()
type_signature = computation_types.TensorType(tf.int32)
value_1 = remote_executor.RemoteValue(executor_pb2.ValueRef(),
type_signature, executor)
value_2 = remote_executor.RemoteValue(executor_pb2.ValueRef(),
type_signature, executor)
result = loop.run_until_complete(
executor.create_struct([value_1, value_2]))
self.assertIsInstance(result, remote_executor.RemoteValue)
