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)