#

# 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.

#

# pytype: skip-file

import os

import secrets

import shutil

import tempfile

import time

import unittest

from collections.abc import Sequence

from dataclasses import dataclass

from typing import Any

from typing import Optional

import numpy as np

from parameterized import param

from parameterized import parameterized

import apache_beam as beam

from apache_beam.metrics.metric import MetricsFilter

from apache_beam.ml.inference.base import ModelHandler

from apache_beam.ml.inference.base import RunInference

from apache_beam.ml.transforms import base

from apache_beam.testing.util import assert_that

from apache_beam.testing.util import equal_to

# pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports

try:

from apache_beam.ml.transforms import tft

from apache_beam.ml.transforms.handlers import TFTProcessHandler

from apache_beam.ml.transforms.tft import TFTOperation

except ImportError:

tft = None # type: ignore

try:

import PIL

from PIL.Image import Image as PIL_Image

except ImportError:

PIL = None # type: ignore[assignment]

PIL_Image = Any # type: ignore[misc, assignment]

try:

class _FakeOperation(TFTOperation):

def __init__(self, name, *args, **kwargs):

super().__init__(*args, **kwargs)

self.name = name

def apply_transform(self, inputs, output_column_name, **kwargs):

return {output_column_name: inputs}

except: # pylint: disable=bare-except

pass

try:

from apache_beam.runners.dataflow.internal import apiclient

except ImportError:

apiclient = None # type: ignore

class BaseMLTransformTest(unittest.TestCase):

def setUp(self) -> None:

self.artifact_location = tempfile.mkdtemp()

def tearDown(self):

shutil.rmtree(self.artifact_location)

def test_ml_transform_no_read_or_write_artifact_lcoation(self):

with self.assertRaises(ValueError):

_ = base.MLTransform(transforms=[])

@unittest.skipIf(tft is None, 'tft module is not installed.')

def test_ml_transform_appends_transforms_to_process_handler_correctly(self):

fake_fn_1 = _FakeOperation(name='fake_fn_1', columns=['x'])

transforms = [fake_fn_1]

ml_transform = base.MLTransform(

transforms=transforms, write_artifact_location=self.artifact_location)

ml_transform = ml_transform.with_transform(

transform=_FakeOperation(name='fake_fn_2', columns=['x']))

self.assertEqual(len(ml_transform.transforms), 2)

self.assertEqual(ml_transform.transforms[0].name, 'fake_fn_1')

self.assertEqual(ml_transform.transforms[1].name, 'fake_fn_2')

@unittest.skipIf(tft is None, 'tft module is not installed.')

def test_ml_transform_on_dict(self):

transforms = [tft.ScaleTo01(columns=['x'])]

data = [{'x': 1}, {'x': 2}]

with beam.Pipeline() as p:

result = (

p

| beam.Create(data)

| base.MLTransform(

write_artifact_location=self.artifact_location,

transforms=transforms))

expected_output = [

np.array([0.0], dtype=np.float32),

np.array([1.0], dtype=np.float32),

]

actual_output = result | beam.Map(lambda x: x.x)

assert_that(

actual_output, equal_to(expected_output, equals_fn=np.array_equal))

@unittest.skipIf(tft is None, 'tft module is not installed.')

def test_ml_transform_on_list_dict(self):

transforms = [tft.ScaleTo01(columns=['x'])]

data = [{'x': [1, 2, 3]}, {'x': [4, 5, 6]}]

with beam.Pipeline() as p:

result = (

p

| beam.Create(data)

| base.MLTransform(

transforms=transforms,

write_artifact_location=self.artifact_location))

expected_output = [

np.array([0, 0.2, 0.4], dtype=np.float32),

np.array([0.6, 0.8, 1], dtype=np.float32),

]

actual_output = result | beam.Map(lambda x: x.x)

assert_that(

actual_output, equal_to(expected_output, equals_fn=np.array_equal))

@parameterized.expand([

param(

input_data=[{

'x': 1,

'y': 2.0,

}],

input_types={

'x': int, 'y': float

},

expected_dtype={

'x': Sequence[np.float32],

'y': Sequence[np.float32],

},

),

param(

input_data=[{

'x': np.array([1], dtype=np.int64),

'y': np.array([2.0], dtype=np.float32),

}],

input_types={

'x': np.int32, 'y': np.float32

},

expected_dtype={

'x': Sequence[np.float32],

'y': Sequence[np.float32],

},

),

param(

input_data=[{

'x': [1, 2, 3], 'y': [2.0, 3.0, 4.0]

}],

input_types={

'x': list[int], 'y': list[float]

},

expected_dtype={

'x': Sequence[np.float32],

'y': Sequence[np.float32],

},

),

param(

input_data=[{

'x': [1, 2, 3], 'y': [2.0, 3.0, 4.0]

}],

input_types={

'x': Sequence[int],

'y': Sequence[float],

},

expected_dtype={

'x': Sequence[np.float32],

'y': Sequence[np.float32],

},

),

])

@unittest.skipIf(tft is None, 'tft module is not installed.')

def test_ml_transform_dict_output_pcoll_schema(

self, input_data, input_types, expected_dtype):

transforms = [tft.ScaleTo01(columns=['x'])]

with beam.Pipeline() as p:

schema_data = (

p

| beam.Create(input_data)

| beam.Map(lambda x: beam.Row(**x)).with_output_types(

beam.row_type.RowTypeConstraint.from_fields(

list(input_types.items()))))

transformed_data = schema_data | base.MLTransform(

write_artifact_location=self.artifact_location, transforms=transforms)

for name, typ in transformed_data.element_type._fields:

if name in expected_dtype:

self.assertEqual(expected_dtype[name], typ)

@unittest.skipIf(tft is None, 'tft module is not installed.')

def test_ml_transform_fail_for_non_global_windows_in_produce_mode(self):

transforms = [tft.ScaleTo01(columns=['x'])]

with beam.Pipeline() as p:

with self.assertRaises(RuntimeError):

_ = (

p

| beam.Create([{

'x': 1, 'y': 2.0

}])

| beam.WindowInto(beam.window.FixedWindows(1))

| base.MLTransform(

transforms=transforms,

write_artifact_location=self.artifact_location,

))

@unittest.skipIf(tft is None, 'tft module is not installed.')

def test_ml_transform_on_multiple_columns_single_transform(self):

transforms = [tft.ScaleTo01(columns=['x', 'y'])]

data = [{'x': [1, 2, 3], 'y': [1.0, 10.0, 20.0]}]

with beam.Pipeline() as p:

result = (

p

| beam.Create(data)

| base.MLTransform(

transforms=transforms,

write_artifact_location=self.artifact_location))

expected_output_x = [

np.array([0, 0.5, 1], dtype=np.float32),

]

expected_output_y = [np.array([0, 0.47368422, 1], dtype=np.float32)]

actual_output_x = result | beam.Map(lambda x: x.x)

actual_output_y = result | beam.Map(lambda x: x.y)

assert_that(

actual_output_x,

equal_to(expected_output_x, equals_fn=np.array_equal))

assert_that(

actual_output_y,

equal_to(expected_output_y, equals_fn=np.array_equal),

label='y')

@unittest.skipIf(tft is None, 'tft module is not installed.')

def test_ml_transforms_on_multiple_columns_multiple_transforms(self):

transforms = [

tft.ScaleTo01(columns=['x']),

tft.ComputeAndApplyVocabulary(columns=['y'])

]

data = [{'x': [1, 2, 3], 'y': ['a', 'b', 'c']}]

with beam.Pipeline() as p:

result = (

p

| beam.Create(data)

| base.MLTransform(

transforms=transforms,

write_artifact_location=self.artifact_location))

expected_output_x = [

np.array([0, 0.5, 1], dtype=np.float32),

]

expected_output_y = [np.array([2, 1, 0])]

actual_output_x = result | beam.Map(lambda x: x.x)

actual_output_y = result | beam.Map(lambda x: x.y)

assert_that(

actual_output_x,

equal_to(expected_output_x, equals_fn=np.array_equal))

assert_that(

actual_output_y,

equal_to(expected_output_y, equals_fn=np.array_equal),

label='actual_output_y')

@unittest.skipIf(tft is None, 'tft module is not installed.')

def test_mltransform_with_counter(self):

transforms = [

tft.ComputeAndApplyVocabulary(columns=['y']),

tft.ScaleTo01(columns=['x'])

]

data = [{'x': [1, 2, 3], 'y': ['a', 'b', 'c']}]

with beam.Pipeline() as p:

_ = (

p | beam.Create(data)

| base.MLTransform(

transforms=transforms,

write_artifact_location=self.artifact_location))

scale_to_01_counter = MetricsFilter().with_name('BeamML_ScaleTo01')

vocab_counter = MetricsFilter().with_name(

'BeamML_ComputeAndApplyVocabulary')

mltransform_counter = MetricsFilter().with_name('BeamML_MLTransform')

result = p.result

self.assertEqual(

result.metrics().query(scale_to_01_counter)['counters'][0].result, 1)

self.assertEqual(

result.metrics().query(vocab_counter)['counters'][0].result, 1)

self.assertEqual(

result.metrics().query(mltransform_counter)['counters'][0].result, 1)

def test_non_ptransfrom_provider_class_to_mltransform(self):

class Add:

def __call__(self, x):

return x + 1

with self.assertRaisesRegex(TypeError, 'transform must be a subclass of'):

with beam.Pipeline() as p:

_ = (

p

| beam.Create([{

'x': 1

}])

| base.MLTransform(

write_artifact_location=self.artifact_location).with_transform(

Add()))

def test_read_mode_with_transforms(self):

with self.assertRaises(ValueError):

_ = base.MLTransform(

# fake callable

transforms=[lambda x: x],

read_artifact_location=self.artifact_location)

class FakeModel:

def __call__(self, example: list[str]) -> list[str]:

for i in range(len(example)):

if not isinstance(example[i], str):

raise TypeError('Input must be a string')

example[i] = example[i][::-1]

return example

class FakeModelHandler(ModelHandler):

def run_inference(

self,

batch: Sequence[str],

model: Any,

inference_args: Optional[dict[str, Any]] = None):

return model(batch)

def load_model(self):

return FakeModel()

class FakeEmbeddingsManager(base.EmbeddingsManager):

def __init__(self, columns, **kwargs):

super().__init__(columns=columns, **kwargs)

def get_model_handler(self) -> ModelHandler:

FakeModelHandler.__repr__ = lambda x: 'FakeEmbeddingsManager' # type: ignore[method-assign]

return FakeModelHandler()

def get_ptransform_for_processing(self, **kwargs) -> beam.PTransform:

return (RunInference(model_handler=base._TextEmbeddingHandler(self)))

def __repr__(self):

return 'FakeEmbeddingsManager'

class InvalidEmbeddingsManager(base.EmbeddingsManager):

def __init__(self, **kwargs):

super().__init__(**kwargs)

def get_model_handler(self) -> ModelHandler:

InvalidEmbeddingsManager.__repr__ = lambda x: 'InvalidEmbeddingsManager' # type: ignore[method-assign]

return FakeModelHandler()

def get_ptransform_for_processing(self, **kwargs) -> beam.PTransform:

return (RunInference(model_handler=base._TextEmbeddingHandler(self)))

def __repr__(self):

return 'InvalidEmbeddingsManager'

class TextEmbeddingHandlerTest(unittest.TestCase):

def setUp(self) -> None:

self.embedding_conig = FakeEmbeddingsManager(columns=['x'])

self.artifact_location = tempfile.mkdtemp()

def tearDown(self) -> None:

shutil.rmtree(self.artifact_location)

def test_no_columns_or_type_adapter(self):

with self.assertRaises(ValueError):

_ = InvalidEmbeddingsManager()

def test_handler_with_incompatible_datatype(self):

text_handler = base._TextEmbeddingHandler(

embeddings_manager=self.embedding_conig)

data = [

('x', 1),

('x', 2),

('x', 3),

]

with self.assertRaises(TypeError):

text_handler.run_inference(data, None, None)

def test_handler_with_dict_inputs(self):

data = [

{

'x': "Hello world"

},

{

'x': "Apache Beam"

},

]

expected_data = [{

key: value[::-1]

for key, value in d.items()

} for d in data]

with beam.Pipeline() as p:

result = (

p

| beam.Create(data)

| base.MLTransform(

write_artifact_location=self.artifact_location).with_transform(

self.embedding_conig))

assert_that(

result,

equal_to(expected_data),

)

def test_handler_with_batch_sizes(self):

self.embedding_conig.max_batch_size = 100

self.embedding_conig.min_batch_size = 10

data = [

{

'x': "Hello world"

},

{

'x': "Apache Beam"

},

] * 100

expected_data = [{

key: value[::-1]

for key, value in d.items()

} for d in data]

with beam.Pipeline() as p:

result = (

p

| beam.Create(data)

| base.MLTransform(

write_artifact_location=self.artifact_location).with_transform(

self.embedding_conig))

assert_that(

result,

equal_to(expected_data),

)

def test_handler_on_multiple_columns(self):

data = [

{

'x': "Hello world", 'y': "Apache Beam", 'z': 'unchanged'

},

{

'x': "Apache Beam", 'y': "Hello world", 'z': 'unchanged'

},

]

embedding_config = FakeEmbeddingsManager(columns=['x', 'y'])

expected_data = [{

key: (value[::-1] if key in embedding_config.columns else value)

for key, value in d.items()

} for d in data]

with beam.Pipeline() as p:

result = (

p

| beam.Create(data)

| base.MLTransform(write_artifact_location=self.artifact_location).

with_transform(embedding_config))

assert_that(

result,

equal_to(expected_data),

)

def test_handler_on_columns_not_exist_in_input_data(self):

data = [

{

'x': "Hello world", 'y': "Apache Beam"

},

{

'x': "Apache Beam", 'y': "Hello world"

},

]

embedding_config = FakeEmbeddingsManager(columns=['x', 'y', 'a'])

with self.assertRaises(RuntimeError):

with beam.Pipeline() as p:

_ = (

p

| beam.Create(data)

| base.MLTransform(write_artifact_location=self.artifact_location).

with_transform(embedding_config))

def test_handler_with_list_data(self):

data = [{

'x': ['Hello world', 'Apache Beam'],

}, {

'x': ['Apache Beam', 'Hello world'],

}]

with self.assertRaisesRegex(Exception, "Embeddings can only be generated"):

with beam.Pipeline() as p:

_ = (

p

| beam.Create(data)

| base.MLTransform(

write_artifact_location=self.artifact_location).with_transform(

self.embedding_conig))

def test_handler_with_inconsistent_keys(self):

data = [

{

'x': 'foo', 'y': 'bar', 'z': 'baz'

},

{

'x': 'foo2', 'y': 'bar2'

},

{

'x': 'foo3', 'y': 'bar3', 'z': 'baz3'

},

]

self.embedding_conig.min_batch_size = 2

with self.assertRaises(RuntimeError):

with beam.Pipeline() as p:

_ = (

p

| beam.Create(data)

| base.MLTransform(

write_artifact_location=self.artifact_location).with_transform(

self.embedding_conig))

class FakeImageModel:

def __call__(self, example: list[PIL_Image]) -> list[PIL_Image]:

for i in range(len(example)):

if not isinstance(example[i], PIL_Image):

raise TypeError('Input must be an Image')

return example

class FakeImageModelHandler(ModelHandler):

def run_inference(

self,

batch: Sequence[PIL_Image],

model: Any,

inference_args: Optional[dict[str, Any]] = None):

return model(batch)

def load_model(self):

return FakeImageModel()

class FakeImageEmbeddingsManager(base.EmbeddingsManager):

def __init__(self, columns, **kwargs):

super().__init__(columns=columns, **kwargs)

def get_model_handler(self) -> ModelHandler:

FakeModelHandler.__repr__ = lambda x: 'FakeImageEmbeddingsManager' # type: ignore[method-assign]

return FakeImageModelHandler()

def get_ptransform_for_processing(self, **kwargs) -> beam.PTransform:

return (RunInference(model_handler=base._ImageEmbeddingHandler(self)))

def __repr__(self):

return 'FakeImageEmbeddingsManager'

class TestImageEmbeddingHandler(unittest.TestCase):

def setUp(self) -> None:

self.embedding_config = FakeImageEmbeddingsManager(columns=['x'])

self.artifact_location = tempfile.mkdtemp()

def tearDown(self) -> None:

shutil.rmtree(self.artifact_location)

@unittest.skipIf(PIL is None, 'PIL module is not installed.')

def test_handler_with_non_dict_datatype(self):

image_handler = base._ImageEmbeddingHandler(

embeddings_manager=self.embedding_config)

data = [

('x', 'hi there'),

('x', 'not an image'),

('x', 'image_path.jpg'),

]

with self.assertRaises(TypeError):

image_handler.run_inference(data, None, None)

@unittest.skipIf(PIL is None, 'PIL module is not installed.')

def test_handler_with_non_image_datatype(self):

image_handler = base._ImageEmbeddingHandler(

embeddings_manager=self.embedding_config)

data = [

{

'x': 'hi there'

},

{

'x': 'not an image'

},

{

'x': 'image_path.jpg'

},

]

with self.assertRaises(TypeError):

image_handler.run_inference(data, None, None)

@unittest.skipIf(PIL is None, 'PIL module is not installed.')

def test_handler_with_dict_inputs(self):

img_one = PIL.Image.new(mode='RGB', size=(1, 1))

img_two = PIL.Image.new(mode='RGB', size=(1, 1))

data = [

{

'x': img_one

},

{

'x': img_two

},

]

expected_data = [{key: value for key, value in d.items()} for d in data]

with beam.Pipeline() as p:

result = (

p

| beam.Create(data)

| base.MLTransform(

write_artifact_location=self.artifact_location).with_transform(

self.embedding_config))

assert_that(

result,

equal_to(expected_data),

)

@dataclass

class FakeMultiModalInput:

image: Optional[PIL_Image] = None

video: Optional[Any] = None

text: Optional[str] = None

class FakeMultiModalModel:

def __call__(self,

example: list[FakeMultiModalInput]) -> list[FakeMultiModalInput]:

for i in range(len(example)):

if not isinstance(example[i], FakeMultiModalInput):

raise TypeError('Input must be a MultiModalInput')

return example

class FakeMultiModalModelHandler(ModelHandler):

def run_inference(

self,

batch: Sequence[FakeMultiModalInput],

model: Any,

inference_args: Optional[dict[str, Any]] = None):

return model(batch)

def load_model(self):

return FakeMultiModalModel()

class FakeMultiModalEmbeddingsManager(base.EmbeddingsManager):

def __init__(self, columns, **kwargs):

super().__init__(columns=columns, **kwargs)

def get_model_handler(self) -> ModelHandler:

FakeModelHandler.__repr__ = lambda x: 'FakeMultiModalEmbeddingsManager' # type: ignore[method-assign]

return FakeMultiModalModelHandler()

def get_ptransform_for_processing(self, **kwargs) -> beam.PTransform:

return (RunInference(model_handler=base._MultiModalEmbeddingHandler(self)))

def __repr__(self):

return 'FakeMultiModalEmbeddingsManager'

class TestMultiModalEmbeddingHandler(unittest.TestCase):

def setUp(self) -> None:

self.embedding_config = FakeMultiModalEmbeddingsManager(columns=['x'])

self.artifact_location = tempfile.mkdtemp()

def tearDown(self) -> None:

shutil.rmtree(self.artifact_location)

@unittest.skipIf(PIL is None, 'PIL module is not installed.')

def test_handler_with_non_dict_datatype(self):

image_handler = base._MultiModalEmbeddingHandler(

embeddings_manager=self.embedding_config)

data = [

('x', 'hi there'),

('x', 'not an image'),

('x', 'image_path.jpg'),

]

with self.assertRaises(TypeError):

image_handler.run_inference(data, None, None)

@unittest.skipIf(PIL is None, 'PIL module is not installed.')

def test_handler_with_incorrect_datatype(self):

image_handler = base._MultiModalEmbeddingHandler(

embeddings_manager=self.embedding_config)

data = [

{

'x': 'hi there'

},

{

'x': 'not an image'

},

{

'x': 'image_path.jpg'

},

]

with self.assertRaises(TypeError):

image_handler.run_inference(data, None, None)

@unittest.skipIf(PIL is None, 'PIL module is not installed.')

def test_handler_with_dict_inputs(self):

input_one = FakeMultiModalInput(

image=PIL.Image.new(mode='RGB', size=(1, 1)), text="test image one")

input_two = FakeMultiModalInput(

image=PIL.Image.new(mode='RGB', size=(1, 1)), text="test image two")

input_three = FakeMultiModalInput(

image=PIL.Image.new(mode='RGB', size=(1, 1)),

video=bytes.fromhex('2Ef0 F1f2 '),

text="test image three with video")

data = [

{

'x': input_one

},

{

'x': input_two

},

{

'x': input_three

},

]

expected_data = [{key: value for key, value in d.items()} for d in data]

with beam.Pipeline() as p:

result = (

p

| beam.Create(data)

| base.MLTransform(

write_artifact_location=self.artifact_location).with_transform(

self.embedding_config))

assert_that(

result,

equal_to(expected_data),

)

class TestUtilFunctions(unittest.TestCase):

def test_dict_input_fn_normal(self):

input_list = [{'a': 1, 'b': 2}, {'a': 3, 'b': 4}]

columns = ['a', 'b']

expected_output = [1, 2, 3, 4]

self.assertEqual(base._dict_input_fn(columns, input_list), expected_output)

def test_dict_output_fn_normal(self):

input_list = [{'a': 1, 'b': 2}, {'a': 3, 'b': 4}]

columns = ['a', 'b']

embeddings = [1.1, 2.2, 3.3, 4.4]

expected_output = [{'a': 1.1, 'b': 2.2}, {'a': 3.3, 'b': 4.4}]

self.assertEqual(

base._dict_output_fn(columns, input_list, embeddings), expected_output)

def test_dict_input_fn_on_list_inputs(self):

input_list = [{'a': [1, 2, 10], 'b': 3}, {'a': [1], 'b': 5}]

columns = ['a', 'b']

expected_output = [[1, 2, 10], 3, [1], 5]

self.assertEqual(base._dict_input_fn(columns, input_list), expected_output)

def test_dict_output_fn_on_list_inputs(self):

input_list = [{'a': [1, 2, 10], 'b': 3}, {'a': [1], 'b': 5}]

columns = ['a', 'b']

embeddings = [1.1, 2.2, 3.3, 4.4]

expected_output = [{'a': 1.1, 'b': 2.2}, {'a': 3.3, 'b': 4.4}]

self.assertEqual(

base._dict_output_fn(columns, input_list, embeddings), expected_output)

class TestJsonPickleTransformAttributeManager(unittest.TestCase):

def setUp(self):

self.attribute_manager = base._transform_attribute_manager

self.artifact_location = tempfile.mkdtemp()

def tearDown(self):

shutil.rmtree(self.artifact_location)

@unittest.skipIf(tft is None, 'tft module is not installed.')

def test_save_tft_process_handler(self):

transforms = [

tft.ScaleTo01(columns=['x']),

tft.ComputeAndApplyVocabulary(columns=['y'])

]

process_handler = TFTProcessHandler(

transforms=transforms,

artifact_location=self.artifact_location,

)

self.attribute_manager.save_attributes(

ptransform_list=[process_handler],

artifact_location=self.artifact_location,

)

files = os.listdir(self.artifact_location)

self.assertTrue(len(files) == 1)

self.assertTrue(files[0] == base._ATTRIBUTE_FILE_NAME)

def test_save_run_inference(self):

self.attribute_manager.save_attributes(

ptransform_list=[RunInference(model_handler=FakeModelHandler())],

artifact_location=self.artifact_location,

)

files = os.listdir(self.artifact_location)

self.assertTrue(len(files) == 1)

self.assertTrue(files[0] == base._ATTRIBUTE_FILE_NAME)

def test_save_and_load_run_inference(self):

ptransform_list = [RunInference(model_handler=FakeModelHandler())]

self.attribute_manager.save_attributes(

ptransform_list=ptransform_list,

artifact_location=self.artifact_location,

)

loaded_ptransform_list = self.attribute_manager.load_attributes(

artifact_location=self.artifact_location,

)

self.assertTrue(len(loaded_ptransform_list) == len(ptransform_list))

self.assertListEqual(

list(loaded_ptransform_list[0].__dict__.keys()),

list(ptransform_list[0].__dict__.keys()))

get_keys = lambda x: list(x.__dict__.keys())

for i, transform in enumerate(ptransform_list):

self.assertListEqual(

get_keys(transform), get_keys(loaded_ptransform_list[i]))

if hasattr(transform, 'model_handler'):

model_handler = transform.model_handler

loaded_model_handler = loaded_ptransform_list[i].model_handler

self.assertListEqual(

get_keys(model_handler), get_keys(loaded_model_handler))

def test_mltransform_to_ptransform_wrapper(self):

transforms = [

FakeEmbeddingsManager(columns=['x']),

FakeEmbeddingsManager(columns=['y', 'z']),

]

ptransform_mapper = base._MLTransformToPTransformMapper(

transforms=transforms,

artifact_location=self.artifact_location,

artifact_mode=None)

ptransform_list = ptransform_mapper.create_ptransform_list()

self.assertTrue(len(ptransform_list) == 2)

self.assertEqual(type(ptransform_list[0]), RunInference)

expected_columns = [['x'], ['y', 'z']]

for i in range(len(ptransform_list)):

self.assertEqual(type(ptransform_list[i]), RunInference)

self.assertEqual(

type(ptransform_list[i]._model_handler), base._TextEmbeddingHandler)

self.assertEqual(

ptransform_list[i]._model_handler.columns, expected_columns[i])

@unittest.skipIf(apiclient is None, 'apache_beam[gcp] is not installed.')

def test_with_gcs_location_with_none_options(self):

path = f"gs://fake_path_{secrets.token_hex(3)}_{int(time.time())}"

with self.assertRaises(RuntimeError):

self.attribute_manager.save_attributes(

ptransform_list=[], artifact_location=path, options=None)

with self.assertRaises(RuntimeError):

self.attribute_manager.save_attributes(

ptransform_list=[], artifact_location=path)

def test_with_same_local_artifact_location(self):

artifact_location = self.artifact_location

attribute_manager = base._JsonPickleTransformAttributeManager()

ptransform_list = [RunInference(model_handler=FakeModelHandler())]

attribute_manager.save_attributes(

ptransform_list, artifact_location=artifact_location)

with self.assertRaises(FileExistsError):

attribute_manager.save_attributes([lambda x: x],

artifact_location=artifact_location)

class MLTransformDLQTest(unittest.TestCase):

def setUp(self) -> None:

self.artifact_location = tempfile.mkdtemp()

def tearDown(self):

shutil.rmtree(self.artifact_location)

def test_dlq_with_embeddings(self):

with beam.Pipeline() as p:

good, bad = (

p

| beam.Create([{

'x': 1

},

{

'x': 3,

},

{

'x': 'Hello'

}

],

)

| base.MLTransform(

write_artifact_location=self.artifact_location).with_transform(

FakeEmbeddingsManager(

columns=['x'])).with_exception_handling())

good_expected_elements = [{'x': 'olleH'}]

assert_that(

good,

equal_to(good_expected_elements),

label='good',

)

# batching happens in RunInference hence elements

# are in lists in the bad pcoll.

bad_expected_elements = [[{'x': 1}], [{'x': 3}]]

assert_that(

bad | beam.Map(lambda x: x.element),

equal_to(bad_expected_elements),

label='bad',

)

def test_mltransform_with_dlq_and_extract_tranform_name(self):

with beam.Pipeline() as p:

good, bad = (

p

| beam.Create([{

'x': 1

},

{

'x': 3,

},

{

'x': 'Hello'

}

],

)

| base.MLTransform(

write_artifact_location=self.artifact_location).with_transform(

FakeEmbeddingsManager(

columns=['x'])).with_exception_handling())

good_expected_elements = [{'x': 'olleH'}]

assert_that(

good,

equal_to(good_expected_elements),

label='good',

)

bad_expected_transform_name = [

'FakeEmbeddingsManager', 'FakeEmbeddingsManager'

]

assert_that(

bad | beam.Map(lambda x: x.transform_name),

equal_to(bad_expected_transform_name),

)

if __name__ == '__main__':

unittest.main()