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)