#

# Licensed to the Apache Software Foundation (ASF) under one or more

# contributor license agreements. See the NOTICE file distributed with

# this work for additional information regarding copyright ownership.

# The ASF licenses this file to You under the Apache License, Version 2.0

# (the "License"); you may not use this file except in compliance with

# the License. You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

#

"""Tests for apache_beam.ml.base."""

import math

import multiprocessing

import os

import pickle

import random

import sys

import tempfile

import time

import unittest

import unittest.mock

from collections.abc import Iterable

from collections.abc import Mapping

from collections.abc import Sequence

from typing import Any

from typing import Optional

from typing import Union

import pytest

import apache_beam as beam

from apache_beam.examples.inference import run_inference_side_inputs

from apache_beam.metrics.metric import MetricsFilter

from apache_beam.ml.inference import base

from apache_beam.options.pipeline_options import StandardOptions

from apache_beam.testing.test_pipeline import TestPipeline

from apache_beam.testing.util import assert_that

from apache_beam.testing.util import equal_to

from apache_beam.transforms import trigger

from apache_beam.transforms import window

from apache_beam.transforms.periodicsequence import TimestampedValue

from apache_beam.utils import multi_process_shared

class FakeModel:

def predict(self, example: int) -> int:

return example + 1

class FakeStatefulModel:

def __init__(self, state: int):

if state == 100:

raise Exception('Oh no')

self._state = state

def predict(self, example: int) -> int:

return self._state

def increment_state(self, amount: int):

self._state += amount

class FakeSlowModel:

def __init__(self, sleep_on_load_seconds=0, file_path_write_on_del=None):

self._file_path_write_on_del = file_path_write_on_del

time.sleep(sleep_on_load_seconds)

def predict(self, example: int) -> int:

time.sleep(example)

return example

def __del__(self):

if self._file_path_write_on_del is not None:

with open(self._file_path_write_on_del, 'a') as myfile:

myfile.write('Deleted FakeSlowModel')

class FakeIncrementingModel:

def __init__(self):

self._state = 0

def predict(self, example: int) -> int:

self._state += 1

return self._state

class FakeSlowModelHandler(base.ModelHandler[int, int, FakeModel]):

def __init__(

self,

sleep_on_load: int,

multi_process_shared=False,

file_path_write_on_del=None):

self._sleep_on_load = sleep_on_load

self._multi_process_shared = multi_process_shared

self._file_path_write_on_del = file_path_write_on_del

def load_model(self):

return FakeSlowModel(self._sleep_on_load, self._file_path_write_on_del)

def run_inference(

self,

batch: Sequence[int],

model: FakeModel,

inference_args=None) -> Iterable[int]:

for example in batch:

yield model.predict(example)

def share_model_across_processes(self) -> bool:

return self._multi_process_shared

def batch_elements_kwargs(self):

return {'min_batch_size': 1, 'max_batch_size': 1}

class FakeModelHandler(base.ModelHandler[int, int, FakeModel]):

def __init__(

self,

clock=None,

min_batch_size=1,

max_batch_size=9999,

multi_process_shared=False,

state=None,

incrementing=False,

max_copies=1,

num_bytes_per_element=None,

**kwargs):

self._fake_clock = clock

self._min_batch_size = min_batch_size

self._max_batch_size = max_batch_size

self._env_vars = kwargs.get('env_vars', {})

self._multi_process_shared = multi_process_shared

self._state = state

self._incrementing = incrementing

self._max_copies = max_copies

self._num_bytes_per_element = num_bytes_per_element

def load_model(self):

assert (not self._incrementing or self._state is None)

if self._fake_clock:

self._fake_clock.current_time_ns += 500_000_000 # 500ms

if self._incrementing:

return FakeIncrementingModel()

if self._state is not None:

return FakeStatefulModel(self._state)

return FakeModel()

def run_inference(

self,

batch: Sequence[int],

model: FakeModel,

inference_args=None) -> Iterable[int]:

multi_process_shared_loaded = "multi_process_shared" in str(type(model))

if self._multi_process_shared != multi_process_shared_loaded:

raise Exception(

f'Loaded model of type {type(model)}, was' +

f'{"" if self._multi_process_shared else " not"} ' +

'expecting multi_process_shared_model')

if self._fake_clock:

self._fake_clock.current_time_ns += 3_000_000 # 3 milliseconds

for example in batch:

yield model.predict(example)

def update_model_path(self, model_path: Optional[str] = None):

pass

def batch_elements_kwargs(self):

return {

'min_batch_size': self._min_batch_size,

'max_batch_size': self._max_batch_size

}

def share_model_across_processes(self):

return self._multi_process_shared

def model_copies(self):

return self._max_copies

def get_num_bytes(self, batch: Sequence[int]) -> int:

if self._num_bytes_per_element:

return self._num_bytes_per_element * len(batch)

return super().get_num_bytes(batch)

class FakeModelHandlerReturnsPredictionResult(

base.ModelHandler[int, base.PredictionResult, FakeModel]):

def __init__(

self,

clock=None,

model_id='fake_model_id_default',

multi_process_shared=False,

state=None):

self.model_id = model_id

self._fake_clock = clock

self._env_vars = {}

self._multi_process_shared = multi_process_shared

self._state = state

def load_model(self):

if self._state is not None:

return FakeStatefulModel(0)

return FakeModel()

def run_inference(

self,

batch: Sequence[int],

model: Union[FakeModel, FakeStatefulModel],

inference_args=None) -> Iterable[base.PredictionResult]:

multi_process_shared_loaded = "multi_process_shared" in str(type(model))

if self._multi_process_shared != multi_process_shared_loaded:

raise Exception(

f'Loaded model of type {type(model)}, was' +

f'{"" if self._multi_process_shared else " not"} ' +

'expecting multi_process_shared_model')

for example in batch:

yield base.PredictionResult(

model_id=self.model_id,

example=example,

inference=model.predict(example))

if self._state is not None:

model.increment_state(1) # type: ignore[union-attr]

def update_model_path(self, model_path: Optional[str] = None):

self.model_id = model_path if model_path else self.model_id

def share_model_across_processes(self):

return self._multi_process_shared

class FakeModelHandlerNoEnvVars(base.ModelHandler[int, int, FakeModel]):

def __init__(

self, clock=None, min_batch_size=1, max_batch_size=9999, **kwargs):

self._fake_clock = clock

self._min_batch_size = min_batch_size

self._max_batch_size = max_batch_size

def load_model(self):

if self._fake_clock:

self._fake_clock.current_time_ns += 500_000_000 # 500ms

return FakeModel()

def run_inference(

self,

batch: Sequence[int],

model: FakeModel,

inference_args=None) -> Iterable[int]:

if self._fake_clock:

self._fake_clock.current_time_ns += 3_000_000 # 3 milliseconds

for example in batch:

yield model.predict(example)

def update_model_path(self, model_path: Optional[str] = None):

pass

def batch_elements_kwargs(self):

return {

'min_batch_size': self._min_batch_size,

'max_batch_size': self._max_batch_size

}

class FakeClock:

def __init__(self):

# Start at 10 seconds.

self.current_time_ns = 10_000_000_000

def time_ns(self) -> int:

return self.current_time_ns

class ExtractInferences(beam.DoFn):

def process(self, prediction_result):

yield prediction_result.inference

class FakeModelHandlerNeedsBigBatch(FakeModelHandler):

def run_inference(self, batch, unused_model, inference_args=None):

if len(batch) < 100:

raise ValueError('Unexpectedly small batch')

return batch

def batch_elements_kwargs(self):

return {'min_batch_size': 9999}

class FakeModelHandlerFailsOnInferenceArgs(FakeModelHandler):

def run_inference(self, batch, unused_model, inference_args=None):

raise ValueError(

'run_inference should not be called because error should already be '

'thrown from the validate_inference_args check.')

def validate_inference_args(self, inference_args: Optional[dict[str, Any]]):

if inference_args:

raise ValueError(

'inference_args were provided, but should be None because this '

'framework does not expect extra arguments on inferences.')

class FakeModelHandlerExpectedInferenceArgs(FakeModelHandler):

def run_inference(self, batch, unused_model, inference_args=None):

if not inference_args:

raise ValueError('inference_args should exist')

return batch

def validate_inference_args(self, inference_args):

pass

class RunInferenceBaseTest(unittest.TestCase):

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_run_inference_impl_simple_examples_multi_process_shared(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(multi_process_shared=True))

assert_that(actual, equal_to(expected), label='assert:inferences')

def test_run_inference_impl_simple_examples_multi_process_shared_multi_copy(

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(multi_process_shared=True, max_copies=4))

assert_that(actual, equal_to(expected), label='assert:inferences')

def test_run_inference_impl_multi_process_shared_incrementing_multi_copy(

self):

with TestPipeline() as pipeline:

examples = [1, 5, 3, 10, 1, 5, 3, 10, 1, 5, 3, 10, 1, 5, 3, 10]

expected = [1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4]

pcoll = pipeline | 'start' >> beam.Create(examples)

actual = pcoll | base.RunInference(

FakeModelHandler(

multi_process_shared=True,

max_copies=4,

incrementing=True,

max_batch_size=1))

assert_that(actual, equal_to(expected), label='assert:inferences')

def test_run_inference_impl_mps_nobatch_incrementing_multi_copy(self):

with TestPipeline() as pipeline:

examples = [1, 5, 3, 10, 1, 5, 3, 10, 1, 5, 3, 10, 1, 5, 3, 10]

expected = [1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4]

batched_examples = [[example] for example in examples]

pcoll = pipeline | 'start' >> beam.Create(batched_examples)

actual = pcoll | base.RunInference(

FakeModelHandler(

multi_process_shared=True, max_copies=4,

incrementing=True).with_no_batching())

assert_that(actual, equal_to(expected), label='assert:inferences')

def test_run_inference_impl_keyed_mps_incrementing_multi_copy(self):

with TestPipeline() as pipeline:

examples = [1, 5, 3, 10, 1, 5, 3, 10, 1, 5, 3, 10, 1, 5, 3, 10]

keyed_examples = [('abc', example) for example in examples]

expected = [1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4]

keyed_expected = [('abc', val) for val in expected]

pcoll = pipeline | 'start' >> beam.Create(keyed_examples)

actual = pcoll | base.RunInference(

base.KeyedModelHandler(

FakeModelHandler(

multi_process_shared=True,

max_copies=4,

incrementing=True,

max_batch_size=1)))

assert_that(actual, equal_to(keyed_expected), label='assert:inferences')

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(

base.KeyedModelHandler(FakeModelHandler()))

assert_that(actual, equal_to(expected), label='assert:inferences')

def test_run_inference_impl_with_keyed_examples_many_model_handlers(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)]

expected[0] = (0, 200)

pcoll = pipeline | 'start' >> beam.Create(keyed_examples)

mhs = [

base.KeyModelMapping([0],

FakeModelHandler(

state=200, multi_process_shared=True)),

base.KeyModelMapping([1, 2, 3],

FakeModelHandler(multi_process_shared=True))

]

actual = pcoll | base.RunInference(base.KeyedModelHandler(mhs))

assert_that(actual, equal_to(expected), label='assert:inferences')

def test_run_inference_impl_with_keyed_examples_many_model_handlers_metrics(

self):

pipeline = TestPipeline()

examples = [1, 5, 3, 10]

metrics_namespace = 'test_namespace'

keyed_examples = [(i, example) for i, example in enumerate(examples)]

pcoll = pipeline | 'start' >> beam.Create(keyed_examples)

mhs = [

base.KeyModelMapping([0],

FakeModelHandler(

state=200, multi_process_shared=True)),

base.KeyModelMapping([1, 2, 3],

FakeModelHandler(multi_process_shared=True))

]

_ = pcoll | base.RunInference(

base.KeyedModelHandler(mhs), metrics_namespace=metrics_namespace)

result = pipeline.run()

result.wait_until_finish()

metrics_filter = MetricsFilter().with_namespace(namespace=metrics_namespace)

metrics = result.metrics().query(metrics_filter)

assert len(metrics['counters']) != 0

assert len(metrics['distributions']) != 0

metrics_filter = MetricsFilter().with_name('0-_num_inferences')

metrics = result.metrics().query(metrics_filter)

num_inferences_counter_key_0 = metrics['counters'][0]

self.assertEqual(num_inferences_counter_key_0.committed, 1)

metrics_filter = MetricsFilter().with_name('1-_num_inferences')

metrics = result.metrics().query(metrics_filter)

num_inferences_counter_key_1 = metrics['counters'][0]

self.assertEqual(num_inferences_counter_key_1.committed, 3)

metrics_filter = MetricsFilter().with_name('num_inferences')

metrics = result.metrics().query(metrics_filter)

num_inferences_counter_aggregate = metrics['counters'][0]

self.assertEqual(num_inferences_counter_aggregate.committed, 4)

metrics_filter = MetricsFilter().with_name('0-_failed_batches_counter')

metrics = result.metrics().query(metrics_filter)

failed_batches_counter_key_0 = metrics['counters']

self.assertEqual(len(failed_batches_counter_key_0), 0)

metrics_filter = MetricsFilter().with_name('failed_batches_counter')

metrics = result.metrics().query(metrics_filter)

failed_batches_counter_aggregate = metrics['counters']

self.assertEqual(len(failed_batches_counter_aggregate), 0)

metrics_filter = MetricsFilter().with_name(

'0-_load_model_latency_milli_secs')

metrics = result.metrics().query(metrics_filter)

load_latency_dist_key_0 = metrics['distributions'][0]

self.assertEqual(load_latency_dist_key_0.committed.count, 1)

metrics_filter = MetricsFilter().with_name('load_model_latency_milli_secs')

metrics = result.metrics().query(metrics_filter)

load_latency_dist_aggregate = metrics['distributions'][0]

self.assertEqual(load_latency_dist_aggregate.committed.count, 2)

def test_run_inference_impl_with_keyed_examples_many_mhs_max_models_hint(

self):

pipeline = TestPipeline()

examples = [1, 5, 3, 10, 2, 4, 6, 8, 9, 7, 1, 5, 3, 10, 2, 4, 6, 8, 9, 7]

metrics_namespace = 'test_namespace'

keyed_examples = [(i, example) for i, example in enumerate(examples)]

pcoll = pipeline | 'start' >> beam.Create(keyed_examples)

mhs = [

base.KeyModelMapping([0, 2, 4, 6, 8],

FakeModelHandler(

state=200, multi_process_shared=True)),

base.KeyModelMapping(

[1, 3, 5, 7, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19],

FakeModelHandler(multi_process_shared=True))

]

_ = pcoll | base.RunInference(

base.KeyedModelHandler(mhs, max_models_per_worker_hint=1),

metrics_namespace=metrics_namespace)

result = pipeline.run()

result.wait_until_finish()

metrics_filter = MetricsFilter().with_namespace(namespace=metrics_namespace)

metrics = result.metrics().query(metrics_filter)

assert len(metrics['counters']) != 0

assert len(metrics['distributions']) != 0

metrics_filter = MetricsFilter().with_name('load_model_latency_milli_secs')

metrics = result.metrics().query(metrics_filter)

load_latency_dist_aggregate = metrics['distributions'][0]

# We should flip back and forth between models a bit since

# max_models_per_worker_hint=1, but we shouldn't thrash forever

# since most examples belong to the second ModelMapping

self.assertGreater(load_latency_dist_aggregate.committed.count, 2)

self.assertLess(load_latency_dist_aggregate.committed.count, 12)

def test_keyed_many_model_handlers_validation(self):

def mult_two(example: str) -> int:

return int(example) * 2

mhs = [

base.KeyModelMapping(

[0],

FakeModelHandler(

state=200,

multi_process_shared=True).with_preprocess_fn(mult_two)),

base.KeyModelMapping([1, 2, 3],

FakeModelHandler(multi_process_shared=True))

]

with self.assertRaises(ValueError):

base.KeyedModelHandler(mhs)

mhs = [

base.KeyModelMapping(

[0],

FakeModelHandler(

state=200,

multi_process_shared=True).with_postprocess_fn(mult_two)),

base.KeyModelMapping([1, 2, 3],

FakeModelHandler(multi_process_shared=True))

]

with self.assertRaises(ValueError):

base.KeyedModelHandler(mhs)

mhs = [

base.KeyModelMapping([0],

FakeModelHandler(

state=200, multi_process_shared=True)),

base.KeyModelMapping([0, 1, 2, 3],

FakeModelHandler(multi_process_shared=True))

]

with self.assertRaises(ValueError):

base.KeyedModelHandler(mhs)

mhs = [

base.KeyModelMapping([],

FakeModelHandler(

state=200, multi_process_shared=True)),

base.KeyModelMapping([0, 1, 2, 3],

FakeModelHandler(multi_process_shared=True))

]

with self.assertRaises(ValueError):

base.KeyedModelHandler(mhs)

def test_keyed_model_handler_get_num_bytes(self):

mh = base.KeyedModelHandler(FakeModelHandler(num_bytes_per_element=10))

batch = [('key1', 1), ('key2', 2), ('key1', 3)]

expected = len(pickle.dumps(('key1', 'key2', 'key1'))) + 30

actual = mh.get_num_bytes(batch)

self.assertEqual(expected, actual)

def test_keyed_model_handler_multiple_models_get_num_bytes(self):

mhs = [

base.KeyModelMapping(['key1'],

FakeModelHandler(num_bytes_per_element=10)),

base.KeyModelMapping(['key2'],

FakeModelHandler(num_bytes_per_element=20))

]

mh = base.KeyedModelHandler(mhs)

batch = [('key1', 1), ('key2', 2), ('key1', 3)]

expected = len(pickle.dumps(('key1', 'key2', 'key1'))) + 40

actual = mh.get_num_bytes(batch)

self.assertEqual(expected, actual)

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_run_inference_impl_with_keyed_examples_multi_process_shared(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(

base.KeyedModelHandler(FakeModelHandler(multi_process_shared=True)))

assert_that(actual, equal_to(expected), label='assert:inferences')

def test_run_inference_impl_with_maybe_keyed_examples_multi_process_shared(

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(multi_process_shared=True))

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_run_inference_preprocessing(self):

def mult_two(example: str) -> int:

return int(example) * 2

with TestPipeline() as pipeline:

examples = ["1", "5", "3", "10"]

expected = [int(example) * 2 + 1 for example in examples]

pcoll = pipeline | 'start' >> beam.Create(examples)

actual = pcoll | base.RunInference(

FakeModelHandler().with_preprocess_fn(mult_two))

assert_that(actual, equal_to(expected), label='assert:inferences')

def test_run_inference_prebatched(self):

with TestPipeline() as pipeline:

examples = [[1, 5], [3, 10]]

expected = [int(example) + 1 for batch in examples for example in batch]

pcoll = pipeline | 'start' >> beam.Create(examples)

actual = pcoll | base.RunInference(FakeModelHandler().with_no_batching())

assert_that(actual, equal_to(expected), label='assert:inferences')

def test_run_inference_preprocessing_multiple_fns(self):

def add_one(example: str) -> int:

return int(example) + 1

def mult_two(example: int) -> int:

return example * 2

with TestPipeline() as pipeline:

examples = ["1", "5", "3", "10"]

expected = [(int(example) + 1) * 2 + 1 for example in examples]

pcoll = pipeline | 'start' >> beam.Create(examples)

actual = pcoll | base.RunInference(

FakeModelHandler().with_preprocess_fn(mult_two).with_preprocess_fn(

add_one))

assert_that(actual, equal_to(expected), label='assert:inferences')

def test_run_inference_postprocessing(self):

def mult_two(example: int) -> str:

return str(example * 2)

with TestPipeline() as pipeline:

examples = [1, 5, 3, 10]

expected = [str((example + 1) * 2) for example in examples]

pcoll = pipeline | 'start' >> beam.Create(examples)

actual = pcoll | base.RunInference(

FakeModelHandler().with_postprocess_fn(mult_two))

assert_that(actual, equal_to(expected), label='assert:inferences')

def test_run_inference_postprocessing_multiple_fns(self):

def add_one(example: int) -> str:

return str(int(example) + 1)

def mult_two(example: int) -> int:

return example * 2

with TestPipeline() as pipeline:

examples = [1, 5, 3, 10]

expected = [str(((example + 1) * 2) + 1) for example in examples]

pcoll = pipeline | 'start' >> beam.Create(examples)

actual = pcoll | base.RunInference(

FakeModelHandler().with_postprocess_fn(mult_two).with_postprocess_fn(

add_one))

assert_that(actual, equal_to(expected), label='assert:inferences')

def test_run_inference_preprocessing_dlq(self):

def mult_two(example: str) -> int:

if example == "5":

raise Exception("TEST")

return int(example) * 2

with TestPipeline() as pipeline:

examples = ["1", "5", "3", "10"]

expected = [3, 7, 21]

expected_bad = ["5"]

pcoll = pipeline | 'start' >> beam.Create(examples)

main, other = pcoll | base.RunInference(

FakeModelHandler().with_preprocess_fn(mult_two)

).with_exception_handling()

assert_that(main, equal_to(expected), label='assert:inferences')

assert_that(

other.failed_inferences, equal_to([]), label='assert:bad_infer')

# bad will be in form [element, error]. Just pull out bad element.

bad_without_error = other.failed_preprocessing[0] | beam.Map(

lambda x: x[0])

assert_that(

bad_without_error, equal_to(expected_bad), label='assert:failures')

def test_run_inference_postprocessing_dlq(self):

def mult_two(example: int) -> str:

if example == 6:

raise Exception("TEST")

return str(example * 2)

with TestPipeline() as pipeline:

examples = [1, 5, 3, 10]

expected = ["4", "8", "22"]

expected_bad = [6]

pcoll = pipeline | 'start' >> beam.Create(examples)

main, other = pcoll | base.RunInference(

FakeModelHandler().with_postprocess_fn(mult_two)

).with_exception_handling()

assert_that(main, equal_to(expected), label='assert:inferences')

assert_that(

other.failed_inferences, equal_to([]), label='assert:bad_infer')

# bad will be in form [element, error]. Just pull out bad element.

bad_without_error = other.failed_postprocessing[0] | beam.Map(

lambda x: x[0])

assert_that(

bad_without_error, equal_to(expected_bad), label='assert:failures')

def test_run_inference_pre_and_post_processing_dlq(self):

def mult_two_pre(example: str) -> int:

if example == "5":

raise Exception("TEST")

return int(example) * 2

def mult_two_post(example: int) -> str:

if example == 7:

raise Exception("TEST")

return str(example * 2)

with TestPipeline() as pipeline:

examples = ["1", "5", "3", "10"]

expected = ["6", "42"]

expected_bad_pre = ["5"]

expected_bad_post = [7]

pcoll = pipeline | 'start' >> beam.Create(examples)

main, other = pcoll | base.RunInference(

FakeModelHandler().with_preprocess_fn(

mult_two_pre

).with_postprocess_fn(

mult_two_post

)).with_exception_handling()

assert_that(main, equal_to(expected), label='assert:inferences')

assert_that(

other.failed_inferences, equal_to([]), label='assert:bad_infer')

# bad will be in form [elements, error]. Just pull out bad element.

bad_without_error_pre = other.failed_preprocessing[0] | beam.Map(

lambda x: x[0])

assert_that(

bad_without_error_pre,

equal_to(expected_bad_pre),

label='assert:failures_pre')

# bad will be in form [elements, error]. Just pull out bad element.

bad_without_error_post = other.failed_postprocessing[0] | beam.Map(

lambda x: x[0])

assert_that(

bad_without_error_post,

equal_to(expected_bad_post),

label='assert:failures_post')

def test_run_inference_keyed_pre_and_post_processing(self):

def mult_two(element):

return (element[0], element[1] * 2)

with TestPipeline() as pipeline:

examples = [1, 5, 3, 10]

keyed_examples = [(i, example) for i, example in enumerate(examples)]

expected = [(i, ((example * 2) + 1) * 2)

for i, example in enumerate(examples)]

pcoll = pipeline | 'start' >> beam.Create(keyed_examples)

actual = pcoll | base.RunInference(

base.KeyedModelHandler(FakeModelHandler()).with_preprocess_fn(

mult_two).with_postprocess_fn(mult_two))

assert_that(actual, equal_to(expected), label='assert:inferences')

def test_run_inference_maybe_keyed_pre_and_post_processing(self):

def mult_two(element):

return element * 2

def mult_two_keyed(element):

return (element[0], element[1] * 2)

with TestPipeline() as pipeline:

examples = [1, 5, 3, 10]

keyed_examples = [(i, example) for i, example in enumerate(examples)]

expected = [((2 * example) + 1) * 2 for example in examples]

keyed_expected = [(i, ((2 * example) + 1) * 2)

for i, example in enumerate(examples)]

model_handler = base.MaybeKeyedModelHandler(FakeModelHandler())

pcoll = pipeline | 'Unkeyed' >> beam.Create(examples)

actual = pcoll | 'RunUnkeyed' >> base.RunInference(

model_handler.with_preprocess_fn(mult_two).with_postprocess_fn(

mult_two))

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.with_preprocess_fn(mult_two_keyed).with_postprocess_fn(

mult_two_keyed))

assert_that(keyed_actual, equal_to(keyed_expected), label='CheckKeyed')

def test_run_inference_impl_dlq(self):

with TestPipeline() as pipeline:

examples = [1, 'TEST', 3, 10, 'TEST2']

expected_good = [2, 4, 11]

expected_bad = ['TEST', 'TEST2']

pcoll = pipeline | 'start' >> beam.Create(examples)

main, other = pcoll | base.RunInference(

FakeModelHandler(

min_batch_size=1,

max_batch_size=1

)).with_exception_handling()

assert_that(main, equal_to(expected_good), label='assert:inferences')

# bad.failed_inferences will be in form [batch[elements], error].

# Just pull out bad element.

bad_without_error = other.failed_inferences | beam.Map(lambda x: x[0][0])

assert_that(

bad_without_error, equal_to(expected_bad), label='assert:failures')

def test_run_inference_timeout_on_load_dlq(self):

with TestPipeline() as pipeline:

examples = [1, 2]

expected_good = []

expected_bad = [1, 2]

pcoll = pipeline | 'start' >> beam.Create(examples)

main, other = pcoll | base.RunInference(

FakeSlowModelHandler(10)).with_exception_handling(timeout=2)

assert_that(main, equal_to(expected_good), label='assert:inferences')

# bad.failed_inferences will be in form [batch[elements], error].

# Just pull out bad element.

bad_without_error = other.failed_inferences | beam.Map(lambda x: x[0][0])

assert_that(

bad_without_error, equal_to(expected_bad), label='assert:failures')

def test_run_inference_timeout_on_inference_dlq(self):

with TestPipeline() as pipeline:

examples = [10, 11]

expected_good = []

expected_bad = [10, 11]

pcoll = pipeline | 'start' >> beam.Create(examples)

main, other = pcoll | base.RunInference(

FakeSlowModelHandler(0)).with_exception_handling(timeout=5)

assert_that(main, equal_to(expected_good), label='assert:inferences')

# bad.failed_inferences will be in form [batch[elements], error].

# Just pull out bad element.

bad_without_error = other.failed_inferences | beam.Map(lambda x: x[0][0])

assert_that(

bad_without_error, equal_to(expected_bad), label='assert:failures')

def test_run_inference_timeout_not_hit(self):

with TestPipeline() as pipeline:

examples = [1, 2]

expected_good = [1, 2]

expected_bad = []

pcoll = pipeline | 'start' >> beam.Create(examples)

main, other = pcoll | base.RunInference(

FakeSlowModelHandler(3)).with_exception_handling(timeout=500)

assert_that(main, equal_to(expected_good), label='assert:inferences')

# bad.failed_inferences will be in form [batch[elements], error].

# Just pull out bad element.

bad_without_error = other.failed_inferences | beam.Map(lambda x: x[0][0])

assert_that(

bad_without_error, equal_to(expected_bad), label='assert:failures')

@unittest.skipIf(

sys.platform == "win32" or sys.version_info < (3, 11),

"This test relies on the __del__ lifecycle method, but __del__ does " +

"not get invoked in the same way on older versions of Python or on " +

"windows, breaking this test. See " +

"github.com/python/cpython/issues/87950#issuecomment-1807570983 " +

"for example.")

def test_run_inference_timeout_does_garbage_collection(self):

with tempfile.TemporaryDirectory() as tmp_dirname:

tmp_path = os.path.join(tmp_dirname, 'tmp_filename')

expected_file_contents = 'Deleted FakeSlowModel'

with TestPipeline() as pipeline:

# Start with bad example which gets timed out.

# Then provide plenty of time for GC to happen.

examples = [20] + [1] * 15

expected_good = [1] * 15

expected_bad = [20]

pcoll = pipeline | 'start' >> beam.Create(examples)

main, other = pcoll | base.RunInference(

FakeSlowModelHandler(

0, True, tmp_path)).with_exception_handling(timeout=5)

assert_that(main, equal_to(expected_good), label='assert:inferences')

# # bad.failed_inferences will be in form [batch[elements], error].

# # Just pull out bad element.

bad_without_error = other.failed_inferences | beam.Map(

lambda x: x[0][0])

assert_that(

bad_without_error, equal_to(expected_bad), label='assert:failures')

with open(tmp_path) as f:

s = f.read()

self.assertEqual(s, expected_file_contents)

def test_run_inference_impl_inference_args(self):

with TestPipeline() as pipeline:

examples = [1, 5, 3, 10]

pcoll = pipeline | 'start' >> beam.Create(examples)

inference_args = {'key': True}

actual = pcoll | base.RunInference(

FakeModelHandlerExpectedInferenceArgs(),

inference_args=inference_args)

assert_that(actual, equal_to(examples), label='assert:inferences')

def test_run_inference_metrics_with_custom_namespace(self):

metrics_namespace = 'my_custom_namespace'

pipeline = TestPipeline()

examples = [1, 5, 3, 10]

pcoll = pipeline | 'start' >> beam.Create(examples)

_ = pcoll | base.RunInference(

FakeModelHandler(), metrics_namespace=metrics_namespace)

result = pipeline.run()

result.wait_until_finish()

metrics_filter = MetricsFilter().with_namespace(namespace=metrics_namespace)

metrics = result.metrics().query(metrics_filter)

assert len(metrics['counters']) != 0

assert len(metrics['distributions']) != 0

metrics_filter = MetricsFilter().with_namespace(namespace='fake_namespace')

metrics = result.metrics().query(metrics_filter)

assert len(metrics['counters']) == len(metrics['distributions']) == 0

def test_unexpected_inference_args_passed(self):

with self.assertRaisesRegex(ValueError, r'inference_args were provided'):

with TestPipeline() as pipeline:

examples = [1, 5, 3, 10]

pcoll = pipeline | 'start' >> beam.Create(examples)

inference_args = {'key': True}

_ = pcoll | base.RunInference(

FakeModelHandlerFailsOnInferenceArgs(),

inference_args=inference_args)

def test_increment_failed_batches_counter(self):

with self.assertRaises(ValueError):

# TODO(https://github.com/apache/beam/issues/34549): This test relies on

# metrics filtering which doesn't work on Prism yet because Prism renames

# steps (e.g. "Do" becomes "ref_AppliedPTransform_Do_7").

# https://github.com/apache/beam/blob/5f9cd73b7c9a2f37f83971ace3a399d633201dd1/sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner.py#L1590

with TestPipeline('FnApiRunner') as pipeline:

examples = [7]

pcoll = pipeline | 'start' >> beam.Create(examples)

_ = pcoll | base.RunInference(FakeModelHandlerExpectedInferenceArgs())

run_result = pipeline.run()

run_result.wait_until_finish()

metric_results = (

run_result.metrics().query(

MetricsFilter().with_name('failed_batches_counter')))

num_failed_batches_counter = metric_results['counters'][0]

self.assertEqual(num_failed_batches_counter.committed, 3)

# !!!: The above will need to be updated if retry behavior changes

def test_failed_batches_counter_no_failures(self):

pipeline = TestPipeline()

examples = [7]

pcoll = pipeline | 'start' >> beam.Create(examples)

inference_args = {'key': True}

_ = pcoll | base.RunInference(

FakeModelHandlerExpectedInferenceArgs(), inference_args=inference_args)

run_result = pipeline.run()

run_result.wait_until_finish()

metric_results = (

run_result.metrics().query(

MetricsFilter().with_name('failed_batches_counter')))

self.assertEqual(len(metric_results['counters']), 0)

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()