def testEstimatorModelPredict(self):
example_path = self._get_output_data_dir('examples')
self._prepare_predict_examples(example_path)
model_path = self._get_output_data_dir('model')
self._build_predict_model(model_path)
prediction_log_path = self._get_output_data_dir('predictions')
self._run_inference_with_beam(
example_path,
model_spec_pb2.InferenceSpecType(
saved_model_spec=model_spec_pb2.SavedModelSpec(
model_path=model_path)), prediction_log_path)
results = self._get_results(prediction_log_path)
self.assertLen(results, 2)
self.assertEqual(
results[0].predict_log.request.inputs[
run_inference._DEFAULT_INPUT_KEY].string_val[0],
self._predict_examples[0].SerializeToString())
self.assertEqual(results[0].predict_log.response.outputs['y'].dtype,
tf.float32)
self.assertLen(
results[0].predict_log.response.outputs['y'].tensor_shape.dim, 2)
self.assertEqual(
results[0].predict_log.response.outputs['y'].tensor_shape.dim[0].
size, 1)
self.assertEqual(
results[0].predict_log.response.outputs['y'].tensor_shape.dim[1].
size, 1)
def _get_inference_spec(
self, model_path: Text, model_version: Text,
ai_platform_serving_args: Dict[Text, Any]
) -> model_spec_pb2.InferenceSpecType:
if 'project_id' not in ai_platform_serving_args:
raise ValueError(
'\'project_id\' is missing in \'ai_platform_serving_args\'')
project_id = ai_platform_serving_args['project_id']
if 'model_name' not in ai_platform_serving_args:
raise ValueError(
'\'model_name\' is missing in \'ai_platform_serving_args\'')
model_name = ai_platform_serving_args['model_name']
ai_platform_prediction_model_spec = (
model_spec_pb2.AIPlatformPredictionModelSpec(
project_id=project_id,
model_name=model_name,
version_name=model_version))
model_signature = self._get_model_signature(model_path)
if (len(model_signature.inputs) == 1 and list(
model_signature.inputs.values())[0].dtype == tf.string.as_datatype_enum
):
ai_platform_prediction_model_spec.use_serialization_config = True
logging.info(
'Using hosted model on Cloud AI platform, model_name: %s,'
'model_version: %s.', model_name, model_version)
result = model_spec_pb2.InferenceSpecType()
result.ai_platform_prediction_model_spec.CopyFrom(
ai_platform_prediction_model_spec)
return result
def testClassifyModel(self):
example_path = self._get_output_data_dir('examples')
self._prepare_multihead_examples(example_path)
model_path = self._get_output_data_dir('model')
self._build_multihead_model(model_path)
prediction_log_path = self._get_output_data_dir('predictions')
self._run_inference_with_beam(
example_path,
model_spec_pb2.InferenceSpecType(
saved_model_spec=model_spec_pb2.SavedModelSpec(
model_path=model_path, signature_name=['classify_sum'])),
prediction_log_path)
results = self._get_results(prediction_log_path)
self.assertLen(results, 2)
classify_log = results[0].classify_log
self.assertLen(classify_log.request.input.example_list.examples, 1)
self.assertEqual(classify_log.request.input.example_list.examples[0],
self._multihead_examples[0])
self.assertLen(classify_log.response.result.classifications, 1)
self.assertLen(classify_log.response.result.classifications[0].classes,
1)
self.assertAlmostEqual(
classify_log.response.result.classifications[0].classes[0].score,
1.0)
def run(args, pipeline_args):
"""
Run inference pipeline using data generated from streaming pipeline.
"""
pipeline_options = PipelineOptions(pipeline_args,
save_main_session=True,
streaming=True)
with beam.Pipeline(options=pipeline_options) as pipeline:
_ = (
pipeline
| 'ReadTFExample' >> beam.io.gcp.pubsub.ReadStringsFromPubSub(
subscription=args.pubsub_subscription)
|
'ParseExamples' >> beam.Map(lambda x: Parse(x, tf.train.Example()))
| RunInference(
model_spec_pb2.InferenceSpecType(
saved_model_spec=model_spec_pb2.SavedModelSpec(
signature_name=['serving_default'],
model_path=args.saved_model_location)))
| beam.ParDo(
process_encdec_inf_rtn.ProcessReturn(
config={
'tf_transform_graph_dir': args.tf_transform_graph_dir,
'model_config': config.MODEL_CONFIG
}))
| beam.ParDo(process_encdec_inf_rtn.CheckAnomalous(threshold=0.7))
| beam.ParDo(print))
def testMultiInferenceModel(self):
example_path = self._get_output_data_dir('examples')
self._prepare_multihead_examples(example_path)
model_path = self._get_output_data_dir('model')
self._build_multihead_model(model_path)
prediction_log_path = self._get_output_data_dir('predictions')
self._run_inference_with_beam(
example_path,
model_spec_pb2.InferenceSpecType(
saved_model_spec=model_spec_pb2.SavedModelSpec(
model_path=model_path,
signature_name=['regress_diff', 'classify_sum'])),
prediction_log_path)
results = self._get_results(prediction_log_path)
self.assertLen(results, 2)
multi_inference_log = results[0].multi_inference_log
self.assertLen(multi_inference_log.request.input.example_list.examples, 1)
self.assertEqual(multi_inference_log.request.input.example_list.examples[0],
self._multihead_examples[0])
self.assertLen(multi_inference_log.response.results, 2)
signature_names = []
for result in multi_inference_log.response.results:
signature_names.append(result.model_spec.signature_name)
self.assertIn('regress_diff', signature_names)
self.assertIn('classify_sum', signature_names)
result = multi_inference_log.response.results[0]
self.assertEqual(result.model_spec.signature_name, 'regress_diff')
self.assertLen(result.regression_result.regressions, 1)
self.assertAlmostEqual(result.regression_result.regressions[0].value, 0.6)
result = multi_inference_log.response.results[1]
self.assertEqual(result.model_spec.signature_name, 'classify_sum')
self.assertLen(result.classification_result.classifications, 1)
self.assertLen(result.classification_result.classifications[0].classes, 1)
self.assertAlmostEqual(
result.classification_result.classifications[0].classes[0].score, 1.0)
def testDoSkippedModelCreation(self, mock_runner, mock_run_model_inference,
_):
input_dict = {
'examples': [self._examples],
'model': [self._model],
'model_blessing': [self._model_blessing],
}
output_dict = {
'inference_result': [self._inference_result],
}
ai_platform_serving_args = {
'model_name': 'model_name',
'project_id': 'project_id'
}
# Create exe properties.
exec_properties = {
'data_spec':
proto_utils.proto_to_json(bulk_inferrer_pb2.DataSpec()),
'custom_config':
json_utils.dumps(
{executor.SERVING_ARGS_KEY: ai_platform_serving_args}),
}
mock_runner.get_service_name_and_api_version.return_value = ('ml', 'v1')
mock_runner.create_model_for_aip_prediction_if_not_exist.return_value = False
# Run executor.
bulk_inferrer = executor.Executor(self._context)
bulk_inferrer.Do(input_dict, output_dict, exec_properties)
ai_platform_prediction_model_spec = (
model_spec_pb2.AIPlatformPredictionModelSpec(
project_id='project_id',
model_name='model_name',
version_name=self._model_version))
ai_platform_prediction_model_spec.use_serialization_config = True
inference_endpoint = model_spec_pb2.InferenceSpecType()
inference_endpoint.ai_platform_prediction_model_spec.CopyFrom(
ai_platform_prediction_model_spec)
mock_run_model_inference.assert_called_once_with(mock.ANY, mock.ANY,
mock.ANY, mock.ANY,
mock.ANY,
inference_endpoint)
executor_class_path = '%s.%s' % (bulk_inferrer.__class__.__module__,
bulk_inferrer.__class__.__name__)
with telemetry_utils.scoped_labels(
{telemetry_utils.LABEL_TFX_EXECUTOR: executor_class_path}):
job_labels = telemetry_utils.make_labels_dict()
mock_runner.deploy_model_for_aip_prediction.assert_called_once_with(
serving_path=path_utils.serving_model_path(self._model.uri),
model_version_name=mock.ANY,
ai_platform_serving_args=ai_platform_serving_args,
labels=job_labels,
api=mock.ANY,
skip_model_endpoint_creation=True,
set_default=False)
mock_runner.delete_model_from_aip_if_exists.assert_called_once_with(
model_version_name=mock.ANY,
ai_platform_serving_args=ai_platform_serving_args,
api=mock.ANY,
delete_model_endpoint=False)
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_request_body_with_binary_data(self):
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())
inference_spec_type = model_spec_pb2.InferenceSpecType(
ai_platform_prediction_model_spec=model_spec_pb2.
AIPlatformPredictionModelSpec(project_id='test_project',
model_name='test_model',
version_name='test_version'))
remote_predict = run_inference._RemotePredictDoFn(
inference_spec_type, None)
result = list(remote_predict._prepare_instances([example]))
self.assertEqual(result, [
{
'x_bytes': {
'b64': 'QVNhOGFzZGY='
},
'x': 'JLK7ljk3',
'y': [1, 2],
'z': [4.5, 5, 5.5]
},
])
def test_request_serialized_example(self):
examples = [
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] }}}
}
""", tf.train.Example()),
text_format.Parse(
"""
context {
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] }}}
}
""", tf.train.SequenceExample())
]
serialized_examples = [e.SerializeToString() for e in examples]
inference_spec_type = model_spec_pb2.InferenceSpecType(
ai_platform_prediction_model_spec=model_spec_pb2.
AIPlatformPredictionModelSpec(project_id='test_project',
model_name='test_model',
version_name='test_version',
use_serialization_config=True))
remote_predict = run_inference._BatchRemotePredictDoFn(
inference_spec_type, None)
result = remote_predict._make_instances(examples, serialized_examples)
self.assertEqual(result, [{
'b64': base64.b64encode(se).decode()
} for se in serialized_examples])
def run(args, pipeline_args):
"""
Run inference pipeline using data generated from streaming pipeline.
"""
pipeline_options = PipelineOptions(pipeline_args,
save_main_session=True,
streaming=False)
with beam.Pipeline(options=pipeline_options) as pipeline:
_ = (pipeline
| 'ReadTFExample' >> beam.io.tfrecordio.ReadFromTFRecord(
file_pattern=args.tfrecord_folder)
| 'ParseExamples' >> beam.Map(tf.train.Example.FromString)
| RunInference(
model_spec_pb2.InferenceSpecType(
saved_model_spec=model_spec_pb2.SavedModelSpec(
signature_name=['serving_default'],
model_path=args.saved_model_location)))
| beam.ParDo(
process_encdec_inf_rtn.ProcessReturn(
config={
'tf_transform_graph_dir': args.tf_transform_graph_dir,
'model_config': config.MODEL_CONFIG
}))
| beam.ParDo(
process_encdec_inf_rtn.CheckAnomalous(threshold=0.15))
| beam.ParDo(print))
def test_exception_raised_when_project_id_is_empty(self):
inference_spec_type = model_spec_pb2.InferenceSpecType(
ai_platform_prediction_model_spec=model_spec_pb2.
AIPlatformPredictionModelSpec(model_name='test-model', ))
with self.assertRaises(ValueError):
self._set_up_pipeline(inference_spec_type)
self._run_inference_with_beam()
def testModelPathInvalid(self):
example_path = self._get_output_data_dir('examples')
self._prepare_predict_examples(example_path)
prediction_log_path = self._get_output_data_dir('predictions')
with self.assertRaisesRegexp(IOError, 'SavedModel file does not exist.*'):
self._run_inference_with_beam(
example_path,
model_spec_pb2.InferenceSpecType(
saved_model_spec=model_spec_pb2.SavedModelSpec(
model_path=self._get_output_data_dir())), prediction_log_path)
def testInfersElementType(self, input_element, output_type):
# TODO(zwestrick): Skip building the model, which is not actually used, or
# stop using parameterized tests if performance becomes an issue.
model_path = self._get_output_data_dir('model')
self._build_predict_model(model_path)
spec = model_spec_pb2.InferenceSpecType(
saved_model_spec=model_spec_pb2.SavedModelSpec(model_path=model_path))
inference_transform = run_inference.RunInferenceImpl(spec)
with beam.Pipeline() as p:
inference = (p | beam.Create([input_element]) | inference_transform)
self.assertEqual(inference.element_type, output_type)
def _run_model_inference(self, model_path: Text,
example_uris: Mapping[Text,
Text], output_path: Text,
model_spec: bulk_inferrer_pb2.ModelSpec) -> None:
"""Runs model inference on given example data.
Args:
model_path: Path to model.
example_uris: Mapping of example split name to example uri.
output_path: Path to output generated prediction logs.
model_spec: bulk_inferrer_pb2.ModelSpec instance.
Returns:
None
"""
try:
from tfx_bsl.public.beam import run_inference
from tfx_bsl.public.proto import model_spec_pb2
except ImportError:
# TODO(b/151468119): Remove this branch after next release.
run_inference = importlib.import_module(
'tfx_bsl.beam.run_inference')
model_spec_pb2 = importlib.import_module(
'tfx_bsl.proto.model_spec_pb2')
saved_model_spec = model_spec_pb2.SavedModelSpec(
model_path=model_path,
tag=model_spec.tag,
signature_name=model_spec.model_signature_name)
# TODO(b/151468119): Remove this branch after next release.
if getattr(model_spec_pb2, 'InferenceEndpoint', False):
inference_endpoint = getattr(model_spec_pb2, 'InferenceEndpoint')()
else:
inference_endpoint = model_spec_pb2.InferenceSpecType()
inference_endpoint.saved_model_spec.CopyFrom(saved_model_spec)
with self._make_beam_pipeline() as pipeline:
data_list = []
for split, example_uri in example_uris.items():
data = (
pipeline
| 'ReadData[{}]'.format(split) >> beam.io.ReadFromTFRecord(
file_pattern=io_utils.all_files_pattern(example_uri)))
data_list.append(data)
_ = ([data for data in data_list]
| 'FlattenExamples' >> beam.Flatten(pipeline=pipeline)
| 'ParseExamples' >> beam.Map(tf.train.Example.FromString)
| 'RunInference' >>
run_inference.RunInference(inference_endpoint)
| 'WritePredictionLogs' >> beam.io.WriteToTFRecord(
output_path,
file_name_suffix='.gz',
coder=beam.coders.ProtoCoder(
prediction_log_pb2.PredictionLog)))
logging.info('Inference result written to %s.', output_path)
def _get_inference_spec(
self, model_path: Text,
exec_properties: Dict[Text, Any]) -> model_spec_pb2.InferenceSpecType:
model_spec = bulk_inferrer_pb2.ModelSpec()
proto_utils.json_to_proto(exec_properties['model_spec'], model_spec)
saved_model_spec = model_spec_pb2.SavedModelSpec(
model_path=model_path,
tag=model_spec.tag,
signature_name=model_spec.model_signature_name)
result = model_spec_pb2.InferenceSpecType()
result.saved_model_spec.CopyFrom(saved_model_spec)
return result
def _get_inference_spec(
self, model_path: str,
exec_properties: Dict[str, Any]) -> model_spec_pb2.InferenceSpecType:
model_spec = bulk_inferrer_pb2.ModelSpec()
proto_utils.json_to_proto(
exec_properties[standard_component_specs.MODEL_SPEC_KEY], model_spec)
saved_model_spec = model_spec_pb2.SavedModelSpec(
model_path=model_path,
tag=model_spec.tag,
signature_name=model_spec.model_signature_name)
result = model_spec_pb2.InferenceSpecType()
result.saved_model_spec.CopyFrom(saved_model_spec)
return result
def testKerasModelPredict(self):
inputs = tf.keras.Input(shape=(1, ), name='input1')
output1 = tf.keras.layers.Dense(1,
activation=tf.nn.sigmoid,
name='output1')(inputs)
output2 = tf.keras.layers.Dense(1,
activation=tf.nn.sigmoid,
name='output2')(inputs)
inference_model = tf.keras.models.Model(inputs, [output1, output2])
class TestKerasModel(tf.keras.Model):
def __init__(self, inference_model):
super(TestKerasModel, self).__init__(name='test_keras_model')
self.inference_model = inference_model
@tf.function(input_signature=[
tf.TensorSpec(shape=[None], dtype=tf.string, name='inputs')
])
def call(self, serialized_example):
features = {
'input1':
tf.compat.v1.io.FixedLenFeature([1],
dtype=tf.float32,
default_value=0)
}
input_tensor_dict = tf.io.parse_example(
serialized_example, features)
return inference_model(input_tensor_dict['input1'])
model = TestKerasModel(inference_model)
model.compile(optimizer=tf.keras.optimizers.Adam(lr=.001),
loss=tf.keras.losses.binary_crossentropy,
metrics=['accuracy'])
model_path = self._get_output_data_dir('model')
tf.compat.v1.keras.experimental.export_saved_model(model,
model_path,
serving_only=True)
example_path = self._get_output_data_dir('examples')
self._prepare_predict_examples(example_path)
prediction_log_path = self._get_output_data_dir('predictions')
self._run_inference_with_beam(
example_path,
model_spec_pb2.InferenceSpecType(
saved_model_spec=model_spec_pb2.SavedModelSpec(
model_path=model_path)), prediction_log_path)
results = self._get_results(prediction_log_path)
self.assertLen(results, 2)
def testTelemetry(self):
example_path = self._get_output_data_dir('examples')
self._prepare_multihead_examples(example_path)
model_path = self._get_output_data_dir('model')
self._build_multihead_model(model_path)
inference_spec_type = model_spec_pb2.InferenceSpecType(
saved_model_spec=model_spec_pb2.SavedModelSpec(
model_path=model_path, signature_name=['classify_sum']))
pipeline = beam.Pipeline()
_ = (pipeline
| 'ReadExamples' >> beam.io.ReadFromTFRecord(example_path)
| 'MaybeDecode' >> beam.Map(lambda x: x if _randbool() else tf.
train.Example.FromString(x))
| 'RunInference' >>
run_inference.RunInferenceImpl(inference_spec_type))
run_result = pipeline.run()
run_result.wait_until_finish()
num_inferences = run_result.metrics().query(
MetricsFilter().with_name('num_inferences'))
self.assertTrue(num_inferences['counters'])
self.assertEqual(num_inferences['counters'][0].result, 2)
num_instances = run_result.metrics().query(
MetricsFilter().with_name('num_instances'))
self.assertTrue(num_instances['counters'])
self.assertEqual(num_instances['counters'][0].result, 2)
inference_request_batch_size = run_result.metrics().query(
MetricsFilter().with_name('inference_request_batch_size'))
self.assertTrue(inference_request_batch_size['distributions'])
self.assertEqual(
inference_request_batch_size['distributions'][0].result.sum, 2)
inference_request_batch_byte_size = run_result.metrics().query(
MetricsFilter().with_name('inference_request_batch_byte_size'))
self.assertTrue(inference_request_batch_byte_size['distributions'])
self.assertEqual(
inference_request_batch_byte_size['distributions'][0].result.sum,
sum(element.ByteSize() for element in self._multihead_examples))
inference_batch_latency_micro_secs = run_result.metrics().query(
MetricsFilter().with_name('inference_batch_latency_micro_secs'))
self.assertTrue(inference_batch_latency_micro_secs['distributions'])
self.assertGreaterEqual(
inference_batch_latency_micro_secs['distributions'][0].result.sum,
0)
load_model_latency_milli_secs = run_result.metrics().query(
MetricsFilter().with_name('load_model_latency_milli_secs'))
self.assertTrue(load_model_latency_milli_secs['distributions'])
self.assertGreaterEqual(
load_model_latency_milli_secs['distributions'][0].result.sum, 0)
def test_basic(self):
examples = [
text_format.Parse(
"""
features {
feature { key: "x" value { float_list { value: 0 }}}
}
""", tf.train.Example()),
text_format.Parse(
"""
features {
feature { key: "x" value { float_list { value: 1 }}}
}
""", tf.train.Example())
]
model_paths = [self._get_output_data_dir(m) for m in ('model1', 'model2')]
for model_path in model_paths:
self._build_predict_model(model_path)
specs = [
model_spec_pb2.InferenceSpecType(
saved_model_spec=model_spec_pb2.SavedModelSpec(model_path=p))
for p in model_paths
]
with self._make_beam_pipeline() as pipeline:
predictions = (
pipeline
| beam.Create(examples)
| run_inference.RunInferencePerModelImpl(specs)
| beam.MapTuple(
lambda _, p2: p2.predict_log.response.outputs['y'].float_val[0]))
assert_that(predictions, equal_to([0.0, 2.0]))
predictions_table = (
pipeline
|
'CreateTable' >> beam.Create([(i, e) for i, e in enumerate(examples)])
| 'RunInferencePerModelTable' >>
run_inference.RunInferencePerModelImpl(specs)
| beam.MapTuple(lambda k, v: # pylint: disable=g-long-lambda
(k, v[1].predict_log.response.outputs['y'].float_val[
0])))
assert_that(
predictions_table,
equal_to([(0, 0.0), (1, 2.0)]),
label='AssertTable')
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()
def testRegressModel(self, keyed_input: bool, decode_examples: bool):
example_path = self._get_output_data_dir('examples')
self._prepare_multihead_examples(example_path)
model_path = self._get_output_data_dir('model')
self._build_multihead_model(model_path)
prediction_log_path = self._get_output_data_dir('predictions')
self._run_inference_with_beam(
example_path,
model_spec_pb2.InferenceSpecType(
saved_model_spec=model_spec_pb2.SavedModelSpec(
model_path=model_path, signature_name=['regress_diff'])),
prediction_log_path, keyed_input, decode_examples)
results = self._get_results(prediction_log_path)
self.assertLen(results, 2)
regress_log = results[0].regress_log
self.assertLen(regress_log.request.input.example_list.examples, 1)
self.assertEqual(regress_log.request.input.example_list.examples[0],
self._multihead_examples[0])
self.assertLen(regress_log.response.result.regressions, 1)
self.assertAlmostEqual(regress_log.response.result.regressions[0].value,
0.6)
def run(args, pipeline_args):
"""
Run inference pipeline using data generated from streaming pipeline.
"""
pipeline_options = PipelineOptions(pipeline_args,
save_main_session=True,
streaming=False)
with beam.Pipeline(options=pipeline_options) as pipeline:
_ = (pipeline
| 'ReadTFExample' >> beam.io.tfrecordio.ReadFromTFRecord(
file_pattern=args.tfrecord_folder)
| 'ParseExamples' >> beam.Map(tf.train.Example.FromString)
| beam.ParDo(filter_examples.FilterExamples())
| RunInference(
model_spec_pb2.InferenceSpecType(
saved_model_spec=model_spec_pb2.SavedModelSpec(
signature_name=['serving_default'],
model_path=args.saved_model_location)))
| beam.ParDo(process_inference_return.ProcessReturn())
| beam.ParDo(process_inference_return.CheckAnomalous())
| beam.ParDo(print))
def Do(self, input_dict: Dict[Text, List[types.Artifact]],
output_dict: Dict[Text, List[types.Artifact]],
exec_properties: Dict[Text, Any]) -> None:
"""Runs batch inference on a given model with given input examples.
Args:
input_dict: Input dict from input key to a list of Artifacts.
- examples: examples for inference.
- model: exported model.
- model_blessing: model blessing result, optional.
output_dict: Output dict from output key to a list of Artifacts.
- output: bulk inference results.
exec_properties: A dict of execution properties.
- model_spec: JSON string of bulk_inferrer_pb2.ModelSpec instance.
- data_spec: JSON string of bulk_inferrer_pb2.DataSpec instance.
Returns:
None
"""
self._log_startup(input_dict, output_dict, exec_properties)
source = exec_properties[StepKeys.SOURCE]
args = exec_properties[StepKeys.ARGS]
c = source_utils.load_source_path_class(source)
inferrer_step: BaseInferrer = c(**args)
output_examples = artifact_utils.get_single_instance(
output_dict[PREDICTIONS])
if EXAMPLES not in input_dict:
raise ValueError('\'examples\' is missing in input dict.')
if MODEL not in input_dict:
raise ValueError('Input models are not valid, model '
'need to be specified.')
if MODEL_BLESSING in input_dict:
model_blessing = artifact_utils.get_single_instance(
input_dict['model_blessing'])
if not model_utils.is_model_blessed(model_blessing):
logging.info('Model on %s was not blessed', model_blessing.uri)
return
else:
logging.info(
'Model blessing is not provided, exported model will be '
'used.')
model = artifact_utils.get_single_instance(input_dict[MODEL])
model_path = path_utils.serving_model_path(model.uri)
logging.info('Use exported model from %s.', model_path)
output_example_spec = bulk_inferrer_pb2.OutputExampleSpec(
output_columns_spec=[
bulk_inferrer_pb2.OutputColumnsSpec(
predict_output=bulk_inferrer_pb2.PredictOutput(
output_columns=[
bulk_inferrer_pb2.PredictOutputCol(
output_key=x,
output_column=f'{x}_label',
) for x in inferrer_step.get_labels()
]))
])
model_spec = bulk_inferrer_pb2.ModelSpec()
saved_model_spec = model_spec_pb2.SavedModelSpec(
model_path=model_path,
tag=model_spec.tag,
signature_name=model_spec.model_signature_name)
inference_spec = model_spec_pb2.InferenceSpecType()
inference_spec.saved_model_spec.CopyFrom(saved_model_spec)
self._run_model_inference(output_example_spec, input_dict[EXAMPLES],
output_examples, inference_spec,
inferrer_step)
