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 | base.RunInference(FakeModelHandler()) assert_that(actual, equal_to(expected), label='assert:inferences')
def test_counted_metrics(self): pipeline = TestPipeline() examples = [1, 5, 3, 10] pcoll = pipeline | 'start' >> beam.Create(examples) _ = pcoll | base.RunInference(FakeModelHandler()) 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_unkeyed_examples_with_keyed_model_handler(self): pipeline = TestPipeline() with self.assertRaises(TypeError): examples = [1, 3, 5] model_handler = base.KeyedModelHandler(FakeModelHandler()) _ = (pipeline | 'Unkeyed' >> beam.Create(examples) | 'RunUnkeyed' >> base.RunInference(model_handler)) pipeline.run()
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 | base.RunInference(FakeModelLoader()) assert_that(actual, equal_to(expected), label='assert:inferences')
def test_run_inference_impl_kwargs(self): with TestPipeline() as pipeline: examples = [1, 5, 3, 10] pcoll = pipeline | 'start' >> beam.Create(examples) kwargs = {'key': True} actual = pcoll | base.RunInference(FakeModelHandlerWithKwargs(), ** kwargs) assert_that(actual, equal_to(examples), label='assert:inferences')
def test_bad_file_raises(self): with self.assertRaises(RuntimeError): with TestPipeline() as pipeline: examples = [numpy.array([0, 0])] pcoll = pipeline | 'start' >> beam.Create(examples) # TODO(BEAM-14305) Test against the public API. _ = pcoll | base.RunInference( SklearnModelLoader(model_uri='/var/bad_file_name')) pipeline.run()
def test_run_inference_keyed_examples_with_unkeyed_model_handler(self): pipeline = TestPipeline() examples = [1, 3, 5] keyed_examples = [(i, example) for i, example in enumerate(examples)] model_handler = FakeModelHandler() with self.assertRaises(TypeError): _ = (pipeline | 'keyed' >> beam.Create(keyed_examples) | 'RunKeyed' >> base.RunInference(model_handler)) pipeline.run()
def test_run_inference_impl_with_maybe_keyed_examples(self): with TestPipeline() as pipeline: examples = [1, 5, 3, 10] keyed_examples = [(i, example) for i, example in enumerate(examples)] expected = [example + 1 for example in examples] keyed_expected = [(i, example + 1) for i, example in enumerate(examples)] model_handler = base.MaybeKeyedModelHandler(FakeModelHandler()) pcoll = pipeline | 'Unkeyed' >> beam.Create(examples) actual = pcoll | 'RunUnkeyed' >> base.RunInference(model_handler) assert_that(actual, equal_to(expected), label='CheckUnkeyed') keyed_pcoll = pipeline | 'Keyed' >> beam.Create(keyed_examples) keyed_actual = keyed_pcoll | 'RunKeyed' >> base.RunInference( model_handler) assert_that(keyed_actual, equal_to(keyed_expected), label='CheckKeyed')
def test_pipeline_pickled(self): temp_file_name = self.tmpdir + os.sep + 'pickled_file' with open(temp_file_name, 'wb') as file: pickle.dump(build_model(), file) with TestPipeline() as pipeline: examples = [numpy.array([0, 0]), numpy.array([1, 1])] pcoll = pipeline | 'start' >> beam.Create(examples) #TODO(BEAM-14305) Test against the public API. actual = pcoll | base.RunInference( SklearnModelLoader(model_uri=temp_file_name)) expected = [ api.PredictionResult(numpy.array([0, 0]), 0), api.PredictionResult(numpy.array([1, 1]), 1) ] assert_that( actual, equal_to(expected, equals_fn=_compare_prediction_result))
def test_timing_metrics(self): pipeline = TestPipeline() examples = [1, 5, 3, 10] pcoll = pipeline | 'start' >> beam.Create(examples) fake_clock = FakeClock() _ = pcoll | base.RunInference(FakeModelHandler(clock=fake_clock), clock=fake_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, 500)
def test_forwards_batch_args(self): examples = list(range(100)) with TestPipeline() as pipeline: pcoll = pipeline | 'start' >> beam.Create(examples) actual = pcoll | base.RunInference(FakeModelHandlerNeedsBigBatch()) assert_that(actual, equal_to(examples), label='assert:inferences')
def expand(self, pcoll: beam.PCollection) -> beam.PCollection: return pcoll | base.RunInference(self._model_loader)