def get_value(self, value_id): response = self._stub.Compute( executor_pb2.ComputeRequest(value_ref=executor_pb2.ValueRef( id=value_id))) py_typecheck.check_type(response, executor_pb2.ComputeResponse) value, _ = executor_service_utils.deserialize_value(response.value) return value
async def _compute(self, value_ref): py_typecheck.check_type(value_ref, executor_pb2.ValueRef) request = executor_pb2.ComputeRequest(value_ref=value_ref) response = _request(self._stub.Compute, request) py_typecheck.check_type(response, executor_pb2.ComputeResponse) value, _ = executor_serialization.deserialize_value(response.value) return value
def get_value(self, value_id: str): """Retrieves a value using the `Compute` endpoint.""" response = self._stub.Compute( executor_pb2.ComputeRequest( value_ref=executor_pb2.ValueRef(id=value_id))) py_typecheck.check_type(response, executor_pb2.ComputeResponse) value, _ = executor_serialization.deserialize_value(response.value) return value
async def _compute(self, value_ref, type_spec): self._check_has_executor_id() py_typecheck.check_type(value_ref, executor_pb2.ValueRef) request = executor_pb2.ComputeRequest( executor=self._executor_id, value_ref=value_ref) response = self._stub.compute(request) py_typecheck.check_type(response, executor_pb2.ComputeResponse) value, _ = value_serialization.deserialize_value(response.value, type_spec) return value
def test_compute_raises_retryable_error_on_grpc_error_unavailable( self, mock_executor_grpc_stub): instance = mock_executor_grpc_stub.return_value instance.Compute = mock.Mock(side_effect=_raise_grpc_error_unavailable) stub = create_stub() with self.assertRaises(executors_errors.RetryableError): stub.compute( executor_pb2.ComputeRequest(value_ref=executor_pb2.ValueRef()))
async def _compute(self, value_ref): py_typecheck.check_type(value_ref, executor_pb2.ValueRef) request = executor_pb2.ComputeRequest(value_ref=value_ref) if self._bidi_stream is None: response = _request(self._stub.Compute, request) else: response = (await self._bidi_stream.send_request( executor_pb2.ExecuteRequest(compute=request))).compute py_typecheck.check_type(response, executor_pb2.ComputeResponse) value, _ = executor_serialization.deserialize_value(response.value) return value
async def _compute(self, value_ref): py_typecheck.check_type(value_ref, executor_pb2.ValueRef) request = executor_pb2.ComputeRequest(value_ref=value_ref) if not self._bidi_stream: response = self._stub.Compute(request) else: response = self._bidi_stream.send_request( executor_pb2.ExecuteRequest(compute=request)).compute py_typecheck.check_type(response, executor_pb2.ComputeResponse) value, _ = executor_service_utils.deserialize_value(response.value) return value
async def _compute(self, value_ref): py_typecheck.check_type(value_ref, executor_pb2.ValueRef) request = executor_pb2.ComputeRequest(value_ref=value_ref) if not self._bidi_stream: try: response = self._stub.Compute(request) except grpc.RpcError as e: self._handle_grpc_error(e) else: response = (await self._bidi_stream.send_request( executor_pb2.ExecuteRequest(compute=request))).compute py_typecheck.check_type(response, executor_pb2.ComputeResponse) value, _ = executor_service_utils.deserialize_value(response.value) return value
def test_compute_reraises_grpc_error(self, grpc_stub): instance = grpc_stub.return_value instance.Compute = mock.Mock( side_effect=_raise_non_retryable_grpc_error) stub = create_stub() request = executor_pb2.ComputeRequest( executor=executor_pb2.ExecutorId(), value_ref=executor_pb2.ValueRef()) with self.assertRaises(grpc.RpcError) as context: stub.compute(request) self.assertEqual(context.exception.code(), grpc.StatusCode.ABORTED)
def iterator(self): @computations.tf_computation() def comp(): return 1 value_proto, _ = executor_serialization.serialize_value(comp) request = executor_pb2.ExecuteRequest( create_value=executor_pb2.CreateValueRequest(value=value_proto)) yield request response = self.queue.get() create_call_proto = executor_pb2.CreateCallRequest( function_ref=response.create_value.value_ref, argument_ref=None) request = executor_pb2.ExecuteRequest(create_call=create_call_proto) yield request response = self.queue.get() compute_proto = executor_pb2.ComputeRequest( value_ref=response.create_call.value_ref) request = executor_pb2.ExecuteRequest(compute=compute_proto) yield request
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)