AI/ML pipelines

Beam <3 machine learning. Being productive and successful as a machine learning practitioner is often dependent on your ability to efficiently leverage large volumes of data in a way that is uniquely tailored to your resources, requirements, and budget. Whether starting your next AI/ML project or upscaling an existing project, consider adding Apache Beam to your project.

AI/ML workloads

Let’s take a look at the different building blocks that we need to create an end-to-end AI/ML use case and where Apache Beam can help.

Overview of AI/ML building blocks and where Apache Beam can be used

  1. Data ingestion: Incoming new data is stored in your file system or database, or it’s published to a messaging queue.
  2. Data validation: After you receieve your data, check the quality of your data. For example, you might want to detect outliers and calculate standard deviations and class distributions.
  3. Data preprocessing: After you validate your data, transform the data so that it is ready to use to train your model.
  4. Model training: When your data is ready, you can start training your AI/ML model. This step is typically repeated multiple times, depending on the quality of your trained model.
  5. Model validation: Before you deploy your new model, validate its performance and accuracy.
  6. Model deployment: Deploy your model, using it to run inference on new or existing data.

To keep your model up to date and performing well as your data grows and evolves, run these steps multiple times. In addition, you can apply MLOps to your project to automate the AI/ML workflows throughout the model and data lifecycle. Use orchestrators to automate this flow and to handle the transition between the different building blocks in your project.

You can use Apache Beam for data validation, data preprocessing, and model deployment/inference. The next section examines these building blocks in more detail and explores how they can be orchestrated.

Data processing

You can use Apache Beam for data validation and preprocessing by setting up data pipelines that transform your data and output metrics computed from your data. Beam has a rich set of I/O connectors for ingesting and writing data, which allows you to integrate it with your existing file system, database, or messaging queue.

When developing your ML model, you can also first explore your data with the Beam DataFrame API. The DataFrom API lets you identify and implement the required preprocessing steps, making it easier for you to move your pipeline to production.

Steps executed during preprocessing often also need to be applied before running inference, in which case you can use the same Beam implementation twice. Lastly, when you need to do postprocessing after running inference, Apache Beam allows you to incoporate the postprocessing into your model inference pipeline.

Further reading:


Beam provides different ways to implement inference as part of your pipeline. You can run your ML model directly in your pipeline and apply it on big scale datasets, both in batch and streaming pipelines.


The recommended way to implement inference is by using the RunInference API. RunInference takes advantage of existing Apache Beam concepts, such as the BatchElements transform and the Shared class, to enable you to use models in your pipelines to create transforms optimized for machine learning inferences. The ability to create arbitrarily complex workflow graphs also allows you to build multi-model pipelines.

You can integrate your model in your pipeline by using the corresponding model handlers. A ModelHandler is an object that wraps the underlying model and allows you to configure its parameters. Model handlers are available for PyTorch, scikit-learn, and TensorFlow. Examples of how to use RunInference for PyTorch, scikit-learn, and TensorFlow are shown in this notebook.

Because they can process multiple computations simultaneously, GPUs are optimized for training artificial intelligence and deep learning models. RunInference also allows you to use GPUs for significant inference speedup. An example of how to use RunInference with GPUs is demonstrated on the RunInference metrics page.

Custom Inference

The RunInference API doesn’t currently support making remote inference calls using, for example, the Natural Language API or the Cloud Vision API. Therefore, in order to use these remote APIs with Apache Beam, you need to write custom inference calls. The Remote inference in Apache Beam notebook shows how to implement a custom remote inference call using beam.DoFn. When you implement a remote inference for real life projects, consider the following factors:


In order to automate and track the AI/ML workflows throughout your project, you can use orchestrators such as Kubeflow pipelines (KFP) or TensorFlow Extended (TFX). These orchestrators automate your different building blocks and handle the transitions between them.

When you use Apache Beam as one of the building blocks in your project, these orchestrators are able to launch your Apache Beam job and to keep track of the input and output of your pipeline. These tasks are essential when moving your AI/ML solution into production, because they allow you to handle your model and data over time and improve the quality and reproducibility of results.


You can find examples of end-to-end AI/ML pipelines for several use cases: