def _Predict( # pylint: disable=invalid-name
pcoll: beam.pvalue.PCollection,
inference_spec_type: model_spec_pb2.InferenceSpecType):
"""Performs predict PTransform."""
if _using_in_process_inference(inference_spec_type):
prediction_model_spec = _PredictModelSpec(inference_spec_type)
return pcoll | 'Predict' >> run_inference_base.RunInference(
prediction_model_spec)
else:
remote_prediction_model_spec = _RemotePrectictModelSpec(
inference_spec_type, pcoll.pipeline.options)
return pcoll | 'Predict' >> run_inference_base.RunInference(
remote_prediction_model_spec)
def test_counted_metrics(self):
pipeline = TestPipeline()
examples = [1, 5, 3, 10]
pcoll = pipeline | 'start' >> beam.Create(examples)
_ = pcoll | run_inference_base.RunInference(FakeModelLoader())
run_result = pipeline.run()
run_result.wait_until_finish()
metric_results = (
run_result.metrics().query(MetricsFilter().with_name('num_inferences')))
num_inferences_counter = metric_results['counters'][0]
self.assertEqual(num_inferences_counter.committed, 4)
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, 4)
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.assertGreaterEqual(
inference_request_batch_byte_size['distributions'][0].result.sum,
len(pickle.dumps(examples)))
inference_request_batch_byte_size = run_result.metrics().query(
MetricsFilter().with_name('model_byte_size'))
self.assertTrue(inference_request_batch_byte_size['distributions'])
def test_run_inference_impl_simple_examples(self):
with TestPipeline() as pipeline:
examples = [1, 5, 3, 10]
expected = [example + 1 for example in examples]
pcoll = pipeline | 'start' >> beam.Create(examples)
actual = pcoll | run_inference_base.RunInference(FakeModelLoader())
assert_that(actual, equal_to(expected), label='assert:inferences')
def _Regress( # pylint: disable=invalid-name
pcoll: beam.pvalue.PCollection,
inference_spec_type: model_spec_pb2.InferenceSpecType):
"""Performs regress PTransform."""
if not _using_in_process_inference(inference_spec_type):
raise NotImplementedError
return pcoll | 'Regress' >> run_inference_base.RunInference(
_RegressModelSpec(inference_spec_type))
def test_run_inference_impl_with_keyed_examples(self):
with TestPipeline() as pipeline:
examples = [1, 5, 3, 10]
keyed_examples = [(i, example) for i, example in enumerate(examples)]
expected = [(i, example + 1) for i, example in enumerate(examples)]
pcoll = pipeline | 'start' >> beam.Create(keyed_examples)
actual = pcoll | run_inference_base.RunInference(FakeModelLoader())
assert_that(actual, equal_to(expected), label='assert:inferences')
def test_timing_metrics(self):
pipeline = TestPipeline()
examples = [1, 5, 3, 10]
pcoll = pipeline | 'start' >> beam.Create(examples)
mock_clock = MockClock()
_ = pcoll | run_inference_base.RunInference(
FakeModelLoader(clock=mock_clock), clock=mock_clock)
res = pipeline.run()
res.wait_until_finish()
metric_results = (
res.metrics().query(
MetricsFilter().with_name('inference_batch_latency_micro_secs')))
batch_latency = metric_results['distributions'][0]
self.assertEqual(batch_latency.result.count, 3)
self.assertEqual(batch_latency.result.mean, 3000)
metric_results = (
res.metrics().query(
MetricsFilter().with_name('load_model_latency_milli_secs')))
load_model_latency = metric_results['distributions'][0]
self.assertEqual(load_model_latency.result.count, 1)
self.assertEqual(load_model_latency.result.mean, 50)