#
# 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.
import abc
import functools
import logging
import os
import tempfile
import uuid
from collections.abc import Callable
from collections.abc import Mapping
from collections.abc import Sequence
from dataclasses import dataclass
from typing import Any
from typing import Dict
from typing import Generic
from typing import Iterable
from typing import List
from typing import Optional
from typing import TypeVar
from typing import Union
from typing import cast
import jsonpickle
import numpy as np
import apache_beam as beam
from apache_beam.io.filesystems import FileSystems
from apache_beam.metrics.metric import Metrics
from apache_beam.ml.inference.base import ModelHandler
from apache_beam.ml.inference.base import ModelT
from apache_beam.ml.inference.base import RunInferenceDLQ
from apache_beam.options.pipeline_options import PipelineOptions
_LOGGER = logging.getLogger(__name__)
_ATTRIBUTE_FILE_NAME = 'attributes.json'
__all__ = [
'MLTransform',
'ProcessHandler',
'MLTransformProvider',
'BaseOperation',
'EmbeddingsManager'
]
TransformedDatasetT = TypeVar('TransformedDatasetT')
TransformedMetadataT = TypeVar('TransformedMetadataT')
# Input/Output types to the MLTransform.
MLTransformOutputT = TypeVar('MLTransformOutputT')
ExampleT = TypeVar('ExampleT')
# Input to the apply() method of BaseOperation.
OperationInputT = TypeVar('OperationInputT')
# Output of the apply() method of BaseOperation.
OperationOutputT = TypeVar('OperationOutputT')
# Input to the EmbeddingTypeAdapter input_fn
EmbeddingTypeAdapterInputT = TypeVar(
'EmbeddingTypeAdapterInputT') # e.g., Chunk
# Output of the EmbeddingTypeAdapter output_fn
EmbeddingTypeAdapterOutputT = TypeVar(
'EmbeddingTypeAdapterOutputT') # e.g., Embedding
@dataclass
class EmbeddingTypeAdapter(Generic[EmbeddingTypeAdapterInputT,
EmbeddingTypeAdapterOutputT]):
"""Adapts input types to text for embedding and converts output embeddings.
Args:
input_fn: Function to extract text for embedding from input type
output_fn: Function to create output type from input and embeddings
"""
input_fn: Callable[[Sequence[EmbeddingTypeAdapterInputT]], List[str]]
output_fn: Callable[[Sequence[EmbeddingTypeAdapterInputT], Sequence[Any]],
List[EmbeddingTypeAdapterOutputT]]
def __reduce__(self):
"""Custom serialization that preserves type information during
jsonpickle."""
return (self.__class__, (self.input_fn, self.output_fn))
def _map_errors_to_beam_row(element, cls_name=None):
row_elements = {
'element': element[0],
'msg': str(element[1][1]),
'stack': str(element[1][2]),
}
if cls_name is not None:
row_elements['transform_name'] = cls_name
return beam.Row(**row_elements)
class ArtifactMode(object):
PRODUCE = 'produce'
CONSUME = 'consume'
[docs]
class BaseOperation(Generic[OperationInputT, OperationOutputT],
MLTransformProvider,
abc.ABC):
def __init__(self, columns: list[str]) -> None:
"""
Base Opertation class data processing transformations.
Args:
columns: List of column names to apply the transformation.
"""
self.columns = columns
def __call__(self, data: OperationInputT,
output_column_name: str) -> dict[str, OperationOutputT]:
"""
This method is called when the instance of the class is called.
This method will invoke the apply() method of the class.
"""
transformed_data = self.apply_transform(data, output_column_name)
return transformed_data
[docs]
class ProcessHandler(
beam.PTransform[beam.PCollection[ExampleT],
Union[beam.PCollection[MLTransformOutputT],
tuple[beam.PCollection[MLTransformOutputT],
beam.PCollection[beam.Row]]]],
abc.ABC):
"""
Only for internal use. No backwards compatibility guarantees.
"""
def _dict_input_fn(
columns: Sequence[str], batch: Sequence[Union[Dict[str, Any],
beam.Row]]) -> List[str]:
"""Extract text from specified columns in batch."""
if batch and hasattr(batch[0], '_asdict'):
batch = [row._asdict() if hasattr(row, '_asdict') else row for row in batch]
if not batch or not isinstance(batch[0], dict):
raise TypeError(
'Expected data to be dicts, got '
f'{type(batch[0])} instead.')
result = []
expected_keys = set(batch[0].keys())
expected_columns = set(columns)
# Process one batch item at a time
for item in batch:
item_keys = item.keys() if isinstance(item, dict) else set()
if set(item_keys) != expected_keys:
extra_keys = item_keys - expected_keys
missing_keys = expected_keys - item_keys
raise RuntimeError(
f'All dicts in batch must have the same keys. '
f'extra keys: {extra_keys}, '
f'missing keys: {missing_keys}')
missing_columns = expected_columns - item_keys
if (missing_columns):
raise RuntimeError(
f'Data does not contain the following columns '
f': {missing_columns}.')
# Get all columns for this item
for col in columns:
if isinstance(item, dict):
result.append(item[col])
return result
def _dict_output_fn(
columns: Sequence[str],
batch: Sequence[Union[Dict[str, Any], beam.Row]],
embeddings: Sequence[Any]) -> list[Union[dict[str, Any], beam.Row]]:
"""Map embeddings back to columns in batch."""
is_beam_row = False
if batch and hasattr(batch[0], '_asdict'):
is_beam_row = True
batch = [row._asdict() if hasattr(row, '_asdict') else row for row in batch]
result = []
for batch_idx, item in enumerate(batch):
for col_idx, col in enumerate(columns):
embedding_idx = batch_idx * len(columns) + col_idx
if isinstance(item, dict):
item[col] = embeddings[embedding_idx]
result.append(item)
if is_beam_row:
result = [beam.Row(**item) for item in result if isinstance(item, dict)]
return result
def _create_dict_adapter(
columns: List[str]) -> EmbeddingTypeAdapter[Dict[str, Any], Dict[str, Any]]:
"""Create adapter for dict-based processing."""
return EmbeddingTypeAdapter[Dict[str, Any], Dict[str, Any]](
input_fn=cast(
Callable[[Sequence[Dict[str, Any]]], List[str]],
functools.partial(_dict_input_fn, columns)),
output_fn=cast(
Callable[[Sequence[Dict[str, Any]], Sequence[Any]],
List[Dict[str, Any]]],
functools.partial(_dict_output_fn, columns)))
# TODO:https://github.com/apache/beam/issues/29356
# Add support for inference_fn
[docs]
class EmbeddingsManager(MLTransformProvider):
def __init__(
self,
*,
columns: Optional[list[str]] = None,
type_adapter: Optional[EmbeddingTypeAdapter] = None,
# common args for all ModelHandlers.
load_model_args: Optional[dict[str, Any]] = None,
min_batch_size: Optional[int] = None,
max_batch_size: Optional[int] = None,
large_model: bool = False,
**kwargs):
self.load_model_args = load_model_args or {}
self.min_batch_size = min_batch_size
self.max_batch_size = max_batch_size
self.large_model = large_model
self.columns = columns
if columns is not None:
self.type_adapter = _create_dict_adapter(columns)
elif type_adapter is not None:
self.type_adapter = type_adapter
else:
raise ValueError("Either columns or type_adapter must be specified")
self.inference_args = kwargs.pop('inference_args', {})
if kwargs:
_LOGGER.warning("Ignoring the following arguments: %s", kwargs.keys())
# TODO:https://github.com/apache/beam/pull/29564 add set_model_handler method
[docs]
@abc.abstractmethod
def get_model_handler(self) -> ModelHandler:
"""
Return framework specific model handler.
"""
[docs]
def get_columns_to_apply(self):
return self.columns
class MLTransformMetricsUsage(beam.PTransform):
def __init__(self, ml_transform: MLTransform):
self._ml_transform = ml_transform
self._ml_transform._counter.inc()
def expand(self, pipeline):
def _increment_counters():
# increment for MLTransform.
self._ml_transform._counter.inc()
# increment if data processing transforms are passed.
transforms = self._ml_transform.transforms
if transforms:
for transform in transforms:
transform.get_counter().inc()
_ = (
pipeline
| beam.Create([None])
| beam.Map(lambda _: _increment_counters()))
class _TransformAttributeManager:
"""
Base class used for saving and loading the attributes.
"""
@staticmethod
def save_attributes(artifact_location):
"""
Save the attributes to json file using stdlib json.
"""
raise NotImplementedError
@staticmethod
def load_attributes(artifact_location):
"""
Load the attributes from json file.
"""
raise NotImplementedError
class _JsonPickleTransformAttributeManager(_TransformAttributeManager):
"""
Use Jsonpickle to save and load the attributes. Here the attributes refer
to the list of PTransforms that are used to process the data.
jsonpickle is used to serialize the PTransforms and save it to a json file and
is compatible across python versions.
"""
@staticmethod
def _is_remote_path(path):
is_gcs = path.find('gs://') != -1
# TODO:https://github.com/apache/beam/issues/29356
# Add support for other remote paths.
if not is_gcs and path.find('://') != -1:
raise RuntimeError(
"Artifact locations are currently supported for only available for "
"local paths and GCS paths. Got: %s" % path)
return is_gcs
@staticmethod
def save_attributes(
ptransform_list,
artifact_location,
**kwargs,
):
# if an artifact location is present, instead of overwriting the
# existing file, raise an error since the same artifact location
# can be used by multiple beam jobs and this could result in undesired
# behavior.
if FileSystems.exists(FileSystems.join(artifact_location,
_ATTRIBUTE_FILE_NAME)):
raise FileExistsError(
"The artifact location %s already exists and contains %s. Please "
"specify a different location." %
(artifact_location, _ATTRIBUTE_FILE_NAME))
if _JsonPickleTransformAttributeManager._is_remote_path(artifact_location):
temp_dir = tempfile.mkdtemp()
temp_json_file = os.path.join(temp_dir, _ATTRIBUTE_FILE_NAME)
with open(temp_json_file, 'w+') as f:
f.write(jsonpickle.encode(ptransform_list))
with open(temp_json_file, 'rb') as f:
from apache_beam.runners.dataflow.internal import apiclient
_LOGGER.info('Creating artifact location: %s', artifact_location)
# pipeline options required to for the client to configure project.
options = kwargs.get('options')
try:
apiclient.DataflowApplicationClient(options=options).stage_file(
gcs_or_local_path=artifact_location,
file_name=_ATTRIBUTE_FILE_NAME,
stream=f,
mime_type='application/json')
except Exception as exc:
if not options:
raise RuntimeError(
"Failed to create Dataflow client. "
"Pipeline options are required to save the attributes."
"in the artifact location %s" % artifact_location) from exc
raise
else:
if not FileSystems.exists(artifact_location):
FileSystems.mkdirs(artifact_location)
# FileSystems.open() fails if the file does not exist.
with open(os.path.join(artifact_location, _ATTRIBUTE_FILE_NAME),
'w+') as f:
f.write(jsonpickle.encode(ptransform_list))
@staticmethod
def load_attributes(artifact_location):
with FileSystems.open(os.path.join(artifact_location, _ATTRIBUTE_FILE_NAME),
'rb') as f:
return jsonpickle.decode(f.read())
_transform_attribute_manager = _JsonPickleTransformAttributeManager
class _MLTransformToPTransformMapper:
"""
This class takes in a list of data processing transforms compatible to be
wrapped around MLTransform and returns a list of PTransforms that are used to
run the data processing transforms.
The _MLTransformToPTransformMapper is responsible for loading and saving the
PTransforms or attributes of PTransforms to the artifact location to seal
the gap between the training and inference pipelines.
"""
def __init__(
self,
transforms: list[MLTransformProvider],
artifact_location: str,
artifact_mode: str = ArtifactMode.PRODUCE,
pipeline_options: Optional[PipelineOptions] = None,
):
self.transforms = transforms
self._parent_artifact_location = artifact_location
self.artifact_mode = artifact_mode
self.pipeline_options = pipeline_options
def create_and_save_ptransform_list(self):
ptransform_list = self.create_ptransform_list()
self.save_transforms_in_artifact_location(ptransform_list)
return ptransform_list
def create_ptransform_list(self):
previous_ptransform_type = None
current_ptransform = None
ptransform_list = []
for transform in self.transforms:
if not isinstance(transform, MLTransformProvider):
raise RuntimeError(
'Transforms must be instances of MLTransformProvider and '
'implement get_ptransform_for_processing() method.')
# for each instance of PTransform, create a new artifact location
current_ptransform = transform.get_ptransform_for_processing(
artifact_location=os.path.join(
self._parent_artifact_location, uuid.uuid4().hex[:6]),
artifact_mode=self.artifact_mode)
append_transform = hasattr(current_ptransform, 'append_transform')
if (type(current_ptransform)
!= previous_ptransform_type) or not append_transform:
ptransform_list.append(current_ptransform)
previous_ptransform_type = type(current_ptransform)
# If different PTransform is appended to the list and the PTransform
# supports append_transform, append the transform to the PTransform.
if append_transform:
ptransform_list[-1].append_transform(transform)
return ptransform_list
def save_transforms_in_artifact_location(self, ptransform_list):
"""
Save the ptransform references to json file.
"""
_transform_attribute_manager.save_attributes(
ptransform_list=ptransform_list,
artifact_location=self._parent_artifact_location,
options=self.pipeline_options)
@staticmethod
def load_transforms_from_artifact_location(artifact_location):
return _transform_attribute_manager.load_attributes(artifact_location)
class _EmbeddingHandler(ModelHandler):
"""
A ModelHandler intended to be work on list[dict[str, Any]] inputs.
The inputs to the model handler are expected to be a list of dicts.
For example, if the original mode is used with RunInference to take a
PCollection[E] to a PCollection[P], this ModelHandler would take a
PCollection[dict[str, E]] to a PCollection[dict[str, P]].
_EmbeddingHandler will accept an EmbeddingsManager instance, which
contains the details of the model to be loaded and the inference_fn to be
used. The purpose of _EmbeddingHandler is to generate embeddings for
general inputs using the EmbeddingsManager instance.
This is an internal class and offers no backwards compatibility guarantees.
Args:
embeddings_manager: An EmbeddingsManager instance.
"""
def __init__(self, embeddings_manager: EmbeddingsManager):
self.embedding_config = embeddings_manager
self._underlying = self.embedding_config.get_model_handler()
self.columns = self.embedding_config.get_columns_to_apply()
def load_model(self):
model = self._underlying.load_model()
return model
def _validate_column_data(self, batch):
pass
def run_inference(
self,
batch: Sequence[dict[str, list[str]]],
model: ModelT,
inference_args: Optional[dict[str, Any]] = None,
) -> list[dict[str, Union[list[float], list[str]]]]:
"""
Runs inference on a batch of text inputs. The inputs are expected to be
a list of dicts. Each dict should have the same keys, and the shape
should be of the same size for a single key across the batch.
"""
embedding_input = self.embedding_config.type_adapter.input_fn(batch)
self._validate_column_data(batch=embedding_input)
prediction = self._underlying.run_inference(
embedding_input, model, inference_args)
# Convert prediction to Sequence[Any]
if isinstance(prediction, np.ndarray):
prediction_seq = prediction.tolist()
elif isinstance(prediction, Iterable) and not isinstance(prediction,
(str, bytes)):
prediction_seq = list(prediction)
else:
prediction_seq = [prediction]
return self.embedding_config.type_adapter.output_fn(batch, prediction_seq)
def get_metrics_namespace(self) -> str:
return (
self._underlying.get_metrics_namespace() or 'BeamML_EmbeddingHandler')
def batch_elements_kwargs(self) -> Mapping[str, Any]:
batch_sizes_map = {}
if self.embedding_config.max_batch_size:
batch_sizes_map['max_batch_size'] = self.embedding_config.max_batch_size
if self.embedding_config.min_batch_size:
batch_sizes_map['min_batch_size'] = self.embedding_config.min_batch_size
return (self._underlying.batch_elements_kwargs() or batch_sizes_map)
def __repr__(self):
return self._underlying.__repr__()
def validate_inference_args(self, _):
pass
class _TextEmbeddingHandler(_EmbeddingHandler):
"""
A ModelHandler intended to be work on list[dict[str, str]] inputs.
The inputs to the model handler are expected to be a list of dicts.
For example, if the original mode is used with RunInference to take a
PCollection[E] to a PCollection[P], this ModelHandler would take a
PCollection[dict[str, E]] to a PCollection[dict[str, P]].
_TextEmbeddingHandler will accept an EmbeddingsManager instance, which
contains the details of the model to be loaded and the inference_fn to be
used. The purpose of _TextEmbeddingHandler is to generate embeddings for
text inputs using the EmbeddingsManager instance.
If the input is not a text column, a RuntimeError will be raised.
This is an internal class and offers no backwards compatibility guarantees.
Args:
embeddings_manager: An EmbeddingsManager instance.
"""
def _validate_column_data(self, batch):
if not isinstance(batch[0], (str, bytes)):
raise TypeError(
'Embeddings can only be generated on dict[str, str].'
f'Got dict[str, {type(batch[0])}] instead.')
def get_metrics_namespace(self) -> str:
return (
self._underlying.get_metrics_namespace() or
'BeamML_TextEmbeddingHandler')
class _ImageEmbeddingHandler(_EmbeddingHandler):
"""
A ModelHandler intended to be work on list[dict[str, Image]] inputs.
The inputs to the model handler are expected to be a list of dicts.
For example, if the original mode is used with RunInference to take a
PCollection[E] to a PCollection[P], this ModelHandler would take a
PCollection[dict[str, E]] to a PCollection[dict[str, P]].
_ImageEmbeddingHandler will accept an EmbeddingsManager instance, which
contains the details of the model to be loaded and the inference_fn to be
used. The purpose of _ImageEmbeddingHandler is to generate embeddings for
image inputs using the EmbeddingsManager instance.
If the input is not an Image representation column, a RuntimeError will be
raised.
This is an internal class and offers no backwards compatibility guarantees.
Args:
embeddings_manager: An EmbeddingsManager instance.
"""
def _validate_column_data(self, batch):
# Don't want to require framework-specific imports
# here, so just catch columns of primatives for now.
if isinstance(batch[0], (int, str, float, bool)):
raise TypeError(
'Embeddings can only be generated on dict[str, Image].'
f'Got dict[str, {type(batch[0])}] instead.')
def get_metrics_namespace(self) -> str:
return (
self._underlying.get_metrics_namespace() or
'BeamML_ImageEmbeddingHandler')