Project: Beam (App Definition and Development / Streaming & Messaging)
File: multi-model-pipelines.md

title: "Multi-model pipelines" <!-- Licensed 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. --> Apache Beam allows you to develop multi-model pipelines. This example demonstrates how to ingest and transform input data, run it through a model, and then pass the outcome of your first model into a second model. This page explains how multi-model pipelines work and gives an overview of what you need to know to build one. Before reading this section, it is recommended that you become familiar with the information in the . A typical machine learning workflow involves a series of data transformation steps, such as data ingestion, data processing tasks, inference, and post-processing. Apache Beam enables you to orchestrate all of those steps together by encapsulating them in a single Apache Beam Directed Acyclic Graph (DAG), which allows you to build resilient and scalable end-to-end machine learning systems. To deploy your machine learning model in an Apache Beam pipeline, use the , which facilitates the integration of your model as a `PTransform` step in your DAG. Composing multiple `RunInference` transforms within a single DAG makes it possible to build a pipeline that consists of multiple ML models. In this way, Apache Beam supports the development of complex ML systems. You can use different patterns to build multi-model pipelines in Apache Beam. This page explores A/B patterns and cascade patterns. The A/B pattern describes a framework multiple where ML models are running in parallel. One application for this pattern is to test the performance of different machine learning models and decide whether a new model is an improvement over an existing one. This is also known as the Champion/Challenger method. Typically, you define a business metric to compare the performance of a control model with the current model. An example could be recommendation engine models where you have an existing model that recommends ads based on the users preferences and activity history. When deciding to deploy a new model, you could split the incoming user traffic into two branches, where half of the users are exposed to the new model and the other half to the current one. After, you can measure the average click-through rate (CTR) of ads for both sets of users over a defined period of time to determine if the new model is performing better than the existing one. ``` import apache_beam as beam with beam.Pipeline() as pipeline: usersetatraffic, usersetbtraffic = (pipeline | 'ReadFromStream' >> beam.ReadFromStream('stream_source') | Partition >> beam.partition(split_dataset, 2, ratio=[5, 5]) ) modelapredictions = usersetatraffic | RunInference(<modelhandlerA>) modelbpredictions = usersetbtraffic | RunInference(<modelhandlerB>) ``` Where

is used to split the data source into 50/50 split partitions. For more information about data partitioning, see . The Cascade pattern is used when the solution to a problem involves a series of ML models. In this scenario, the output of a model is typically transformed to a suitable format using a `PTransform` before passing it to another model. ``` with pipeline as p: data = p | 'Read' >> beam.ReadFromSource('a_source') modelapredictions = data | RunInference(<modelhandlerA>) modelbpredictions = modelapredictions | beam.ParDo(postprocessing()) | RunInference(<modelhandler_B>) ``` The notebook shows an end-to-end example of a cascade pipeline used to generate and rank image captions. The solution consists of two open-source models: A caption generation model () that generates candidate image captions from an input image. A caption ranking model () that uses the image and candidate captions to rank the captions in the order in which they best describe the image. You can use a `KeyedModelHandler` to load several different models into the `RunInference` transform. Use the associated key to determine which model to use with which data. The following example loads a model by using `config1`. That model is used for inference for all examples associated with `key1`. It loads a second model by using `config2`. That model is used for all examples associated with `key2` and `key3`. ``` from apache_beam.ml.inference.base import KeyedModelHandler keyedmodelhandler = KeyedModelHandler([ KeyModelMapping(['key1'], PytorchModelHandlerTensor(<config1>)), KeyModelMapping(['key2', 'key3'], PytorchModelHandlerTensor(<config2>)) ]) with pipeline as p: data = p | beam.Create([ ('key1', torch.tensor([[1,2,3],[4,5,6],...])), ('key2', torch.tensor([[1,2,3],[4,5,6],...])), ('key3', torch.tensor([[1,2,3],[4,5,6],...])), ]) predictions = data | RunInference(keyedmodelhandler) ``` For a more detailed example, see the notebook . Loading multiple models at the same times increases the risk of out of memory errors (OOMs). By default, `KeyedModelHandler` doesn't limit the number of models loaded into memory at the same time. If the models don't all fit into memory, your pipeline might fail with an out of memory error. To avoid this issue, use the `maxmodelsperworkerhint` parameter to set the maximum number of models that can be loaded into memory at the same time. The following example loads at most two models per SDK worker process at a time. It unloads models that aren't currently in use. ``` mhs = [ KeyModelMapping(['key1'], PytorchModelHandlerTensor(<config1>)), KeyModelMapping(['key2', 'key3'], PytorchModelHandlerTensor(<config2>)), KeyModelMapping(['key4'], PytorchModelHandlerTensor(<config3>)), KeyModelMapping(['key5', 'key6', 'key7'], PytorchModelHandlerTensor(<config4>)), ] keyedmodelhandler = KeyedModelHandler(mhs, maxmodelsperworkerhint=2) ``` Runners that have multiple SDK worker processes on a given machine load at most `maxmodelsperworkerhint*<num worker processes>` models onto the machine. Leave enough space for the models and any additional memory needs from other transforms. Because the memory might not be released immediately after a model is offloaded, leaving an additional buffer is recommended. Note: Having many models but a small `maxmodelsperworkerhint` can cause memory thrashing, where a large amount of execution time is used to swap models in and out of memory. To reduce the likelihood and impact of memory thrashing, if you're using a distributed runner, insert a transform before your inference step. The `GroupByKey` transform reduces thrashing by ensuring that elements with the same key and model are collocated on the same worker. For more information, see .