def test_exception_raised_when_response_body_contains_error_entry(self):
error_msg = 'Base64 decode failed.'
builder = http.RequestMockBuilder({
'ml.projects.predict':
(None, self._make_response_body(error_msg, successful=False))
})
resource = discovery.build('ml',
'v1',
http=http.HttpMock(
self._discovery_testdata_dir,
{'status': http_client.OK}),
requestBuilder=builder)
with mock.patch('googleapiclient.discovery.' 'build') as response_mock:
response_mock.side_effect = lambda service, version: resource
inference_spec_type = model_spec_pb2.InferenceSpecType(
ai_platform_prediction_model_spec=model_spec_pb2.
AIPlatformPredictionModelSpec(
project_id='test-project',
model_name='test-model',
))
try:
self._set_up_pipeline(inference_spec_type)
self._run_inference_with_beam()
except ValueError as exc:
actual_error_msg = str(exc)
self.assertTrue(actual_error_msg.startswith(error_msg))
else:
self.fail('Test was expected to throw ValueError exception')
def test_model_predict(self):
predictions = [{'output_1': [0.901], 'output_2': [0.997]}]
builder = http.RequestMockBuilder({
'ml.projects.predict':
(None, self._make_response_body(predictions, successful=True))
})
resource = discovery.build('ml',
'v1',
http=http.HttpMock(
self._discovery_testdata_dir,
{'status': http_client.OK}),
requestBuilder=builder)
with mock.patch('googleapiclient.discovery.' 'build') as response_mock:
response_mock.side_effect = lambda service, version: resource
inference_spec_type = model_spec_pb2.InferenceSpecType(
ai_platform_prediction_model_spec=model_spec_pb2.
AIPlatformPredictionModelSpec(
project_id='test-project',
model_name='test-model',
))
prediction_log = prediction_log_pb2.PredictionLog()
prediction_log.predict_log.response.outputs['output_1'].CopyFrom(
tf.make_tensor_proto(values=[0.901],
dtype=tf.double,
shape=(1, 1)))
prediction_log.predict_log.response.outputs['output_2'].CopyFrom(
tf.make_tensor_proto(values=[0.997],
dtype=tf.double,
shape=(1, 1)))
self._set_up_pipeline(inference_spec_type)
assert_that(self.pcoll, equal_to([prediction_log]))
self._run_inference_with_beam()
def test_can_format_requests(self):
predictions = [{
'output_1': [0.901],
'output_2': [0.997]
}] * len(self._predict_examples)
builder = http.RequestMockBuilder({
'ml.projects.predict':
(None, self._make_response_body(predictions, successful=True))
})
resource = discovery.build(
'ml',
'v1',
http=http.HttpMock(self._discovery_testdata_dir,
{'status': http_client.OK}),
requestBuilder=builder)
with mock.patch('googleapiclient.discovery.' 'build') as response_mock:
response_mock.side_effect = lambda service, version: resource
inference_spec_type = model_spec_pb2.InferenceSpecType(
ai_platform_prediction_model_spec=model_spec_pb2
.AIPlatformPredictionModelSpec(
project_id='test-project',
model_name='test-model',
))
example = text_format.Parse(
"""
features {
feature { key: "x_bytes" value { bytes_list { value: ["ASa8asdf"] }}}
feature { key: "x" value { bytes_list { value: "JLK7ljk3" }}}
feature { key: "y" value { int64_list { value: [1, 2] }}}
feature { key: "z" value { float_list { value: [4.5, 5, 5.5] }}}
}
""", tf.train.Example())
self.pipeline = self._make_beam_pipeline()
self.pcoll = (
self.pipeline
| 'CreateExamples' >> beam.Create([example])
| 'RunInference'
>> run_inference.RunInferenceImpl(inference_spec_type))
self._run_inference_with_beam()
def test_model_predict(self, keyed_input: bool):
predictions = [{
'output_1': [0.901],
'output_2': [0.997]
}] * len(self._predict_examples)
builder = http.RequestMockBuilder({
'ml.projects.predict':
(None, self._make_response_body(predictions, successful=True))
})
resource = discovery.build(
'ml',
'v1',
http=http.HttpMock(self._discovery_testdata_dir,
{'status': http_client.OK}),
requestBuilder=builder)
with mock.patch('googleapiclient.discovery.' 'build') as response_mock:
response_mock.side_effect = lambda service, version: resource
inference_spec_type = model_spec_pb2.InferenceSpecType(
ai_platform_prediction_model_spec=model_spec_pb2
.AIPlatformPredictionModelSpec(
project_id='test-project',
model_name='test-model',
))
expected = []
for example in self._predict_examples:
prediction_log = prediction_log_pb2.PredictionLog()
predict_log = prediction_log.predict_log
input_tensor_proto = predict_log.request.inputs['inputs']
input_tensor_proto.dtype = tf.string.as_datatype_enum
input_tensor_proto.tensor_shape.dim.add().size = 1
input_tensor_proto.string_val.append(example.SerializeToString())
predict_log.response.outputs['output_1'].CopyFrom(
tf.make_tensor_proto(values=[0.901], dtype=tf.double, shape=(1, 1)))
predict_log.response.outputs['output_2'].CopyFrom(
tf.make_tensor_proto(values=[0.997], dtype=tf.double, shape=(1, 1)))
expected.append(prediction_log)
self._set_up_pipeline(inference_spec_type, keyed_input)
assert_that(self.pcoll, equal_to(expected))
self._run_inference_with_beam()