Developing with Apache Beam notebooks

Using the Apache Beam interactive runner with JupyterLab notebooks lets you iteratively develop pipelines, inspect your pipeline graph, and parse individual PCollections in a read-eval-print-loop (REPL) workflow. These Apache Beam notebooks are made available through AI Platform Notebooks, a managed service that hosts notebook virtual machines pre-installed with the latest data science and machine learning frameworks.

This guide focuses on the functionality introduced by Apache Beam notebooks, but does not show how to build one. For more information on Apache Beam, see the Apache Beam programming guide.

Note: Apache Beam notebooks currently only support Python. Apache Beam pipeline segments running in these notebooks are run in a test environment, and not against a production Apache Beam runner; however, users can export pipelines created in an Apache Beam notebook and launch them on the Dataflow service. For more details, see Launching Dataflow jobs from your notebook.

Before you begin

1. Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

2. In the Google Cloud Console, on the project selector page, select or create a Google Cloud project.

   Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

   Go to project selector

3. Make sure that billing is enabled for your Cloud project. Learn how to confirm that billing is enabled for your project.

4. Enable the Compute Engine, AI Platform Notebooks APIs.

   Enable the APIs

Note: Enable additional APIs for pipelines that use other services, such as Pub/Sub, before creating your Apache Beam notebook instance. If not specified, the notebook instance is executed by the default compute engine service account with the IAM project editor role. If the project explicitly limits the roles of the service account, make sure it still has enough authorization to run your notebooks. For example, reading from a Pub/Sub topic implicitly creates a subscription, and your service account needs an IAM Pub/Sub editor role. By contrast, reading from a Pub/Sub subscription only requires an IAM Pub/Sub subscriber role.

When you finish this guide, you can avoid continued billing by deleting the resources you created. For more details, see Cleaning up.

Launching an Apache Beam notebooks instance

1. In the Google Cloud Console, on the project selector page, select or create a Google Cloud project.
2. Navigate to Dataflow in the side panel and click Notebooks.
3. In the toolbar, click New Instance.
4. Select Apache Beam.
5. On the New notebook instance page, select a network for the notebook VM and click Create.
6. (Optional) If you want to set up a custom notebook instance, click Customize. For more information on customizing instance properties, see Create an AI Platform Notebooks instance with specific properties.
7. Click Open JupyterLab when the link becomes active. AI Platform Notebooks creates a new Apache Beam notebook instance.

Installing dependencies (Optional)

Apache Beam notebooks already come with Apache Beam and Google Cloud connector dependencies installed. If your pipeline contains custom connectors or custom PTransforms that depend on third-party libraries, you can install them after you create a notebook instance. For more information, see Installing Dependencies in the AI Platform Notebooks documentation.

Getting started with Apache Beam notebooks

After opening an AI Platform Notebooks instance, example notebooks are available in the Examples folder. The following are currently available:

• Word Count
• Streaming Word Count
• Streaming NYC Taxi Ride Data
• Dataflow Word Count

These notebooks include explanatory text and commented code blocks to help you understand Apache Beam concepts and API usage.

Note: Example code from the Streaming Word Count notebook is used in the following sections. There might be some minor discrepancies between the code snippets in this guide and what is found in the Streaming Word Count notebook.

Creating a notebook instance

Navigate to File > New > Notebook and select a kernel that is Apache Beam 2.22 or later.

Note: Apache Beam notebooks are built against the master branch of the Apache Beam SDK. This means that the latest version of the kernel shown in the notebooks UI might be ahead of the most recently released version of the SDK.

Apache Beam is installed on your notebook instance, so include the interactive_runner and interactive_beam modules in your notebook.

import apache_beam as beam
from apache_beam.runners.interactive.interactive_runner import InteractiveRunner
import apache_beam.runners.interactive.interactive_beam as ib

If your notebook uses other Google APIs, add the following import statements:

from apache_beam.options import pipeline_options
from apache_beam.options.pipeline_options import GoogleCloudOptions
import google.auth

Setting interactivity options

The following sets the amount of time the InteractiveRunner records data from an unbounded source. In this example, the duration is set to 10 minutes.

ib.options.recording_duration = '10m'

You can also change the recording size limit (in bytes) for an unbounded source through the recording_size_limit property.

# Set the recording size limit to 1 GB.
ib.options.recording_size_limit = 1e9

Note: When changing either the recording duration or size, call stop, clear, and record after changing values to allow new data to be recorded for the updated duration or size limit.

For additional interactive options, see the interactive_beam.options class.

Creating your pipeline

Initialize the pipeline using an InteractiveRunner object.

options = pipeline_options.PipelineOptions()

# Set the pipeline mode to stream the data from Pub/Sub.
options.view_as(pipeline_options.StandardOptions).streaming = True

p = beam.Pipeline(InteractiveRunner(), options=options)

Reading and visualizing the data

The following example shows an Apache Beam pipeline that creates a subscription to the given Pub/Sub topic and reads from the subscription.

words = p 
    | "read" >> beam.io.ReadFromPubSub(topic="projects/pubsub-public-data/topics/shakespeare-kinglear")

The pipeline counts the words by windows from the source. It creates fixed windowing with each window being 10 seconds in duration.

windowed_words = (words
   | "window" >> beam.WindowInto(beam.window.FixedWindows(10)))

After the data is windowed, the words are counted by window.

windowed_words_counts = (windowed_words
   | "count" >> beam.combiners.Count.PerElement())

The show() method visualizes the resulting PCollection in the notebook.

ib.show(windowed_word_counts, include_window_info=True)

You can scope the result set back from show() by setting two optional parameters: n and duration. Setting n limits the result set to show at most n number of elements, such as 20. If n is not set, the default behavior is to list the most recent elements captured until the source recording is over. Setting duration limits the result set to a specified number of seconds worth of data starting from the beginning of the source recording. If duration is not set, the default behavior is to list all elements until the recording is over.

If both optional parameters are set, show() stops whenever either threshold is met. In the following example, show() returns at most 20 elements that are computed based on the first 30 seconds worth of data from the recorded sources.

ib.show(windowed_word_counts, include_window_info=True, n=20, duration=30)

To display visualizations of your data, pass visualize_data=True into the show() method. You can apply multiple filters to your visualizations. The following visualization allows you to filter by label and axis:

Note: To ensure replayability while prototyping streaming pipelines, the show() method calls reuse the captured data by default. By setting interactive_beam.options.enable_capture_replay = False, you can change this behavior and have the show() method always fetch new data. Also, if you add a second unbounded source to your notebook, the data from the previous unbounded source is discarded.

Another useful visualization in Apache Beam notebooks is a Pandas DataFrame. The following example first converts the words to lowercase and then computes the frequency of each word.

windowed_lower_word_counts = (windowed_words
   | beam.Map(lambda word: word.lower())
   | "count" >> beam.combiners.Count.PerElement())

The collect() method provides the output in a Pandas DataFrame.

ib.collect(windowed_lower_word_counts, include_window_info=True)

Note: Editing and re-executing a cell is a common practice in notebook development. When you edit and re-execute a cell in a Apache Beam notebook, the cell does not undo the intended action of the code in the original cell. For example, if a cell adds a PTransform to a pipeline, re-executing that cell would add an additional PTransform to the pipeline. If you want to clear the state, restart the kernel and rerun the cells.

Understanding a pipeline's recording status

In addition to visualizations, you can also inspect the recording status for one or all pipelines in your notebook instance by calling describe.

# Return the recording status of a specific pipeline. Leave the parameter list empty to return
# the recording status of all pipelines.
ib.recordings.describe(p)

The describe() method provides the following details:

• Total size (in bytes) of all of the recordings for the pipeline on disk
• Start time of when the background recording job started (in seconds from Unix epoch)
• Current pipeline status of the background recording job
• Python variable for the pipeline

Launching Dataflow jobs from a pipeline created in your notebook

1. (Optional) Before using your notebook to run Dataflow jobs, restart the kernel, rerun all cells, and verify the output. If you skip this step, hidden states in the notebook might affect the job graph in the pipeline object.
2. Enable the Dataflow API.

3. Add the following import statement:

   from apache_beam.runners import DataflowRunner
   

4. Pass in your pipeline options.

   # Set up Apache Beam pipeline options.
   options = pipeline_options.PipelineOptions()
   
   # Set the project to the default project in your current Google Cloud
   # environment.
   _, options.view_as(GoogleCloudOptions).project = google.auth.default()
   
   # Set the Google Cloud region to run Dataflow.
   options.view_as(GoogleCloudOptions).region = 'us-central1'
   
   # Choose a Cloud Storage location.
   dataflow_gcs_location = 'gs://<change me>/dataflow'
   
   # Set the staging location. This location is used to stage the
   # Dataflow pipeline and SDK binary.
   options.view_as(GoogleCloudOptions).staging_location = '%s/staging' % dataflow_gcs_location
   
   # Set the temporary location. This location is used to store temporary files
   # or intermediate results before outputting to the sink.
   options.view_as(GoogleCloudOptions).temp_location = '%s/temp' % dataflow_gcs_location
   
   # Set the SDK location. This is used by Dataflow to locate the
   # SDK needed to run the pipeline.
   options.view_as(pipeline_options.SetupOptions).sdk_location = (
       '/root/apache-beam-custom/packages/beam/sdks/python/dist/apache-beam-%s0.tar.gz' %
       beam.version.__version__)
   

   You can adjust the parameter values. For example, you can change the region value from us-central1.

5. Run the pipeline with DataflowRunner. This runs your job on the Dataflow service.

   runner = DataflowRunner()
   runner.run_pipeline(p, options=options)
   

   p is a pipeline object from Creating your pipeline.

For an example on how to perform this conversion on an interactive notebook, see the Dataflow Word Count notebook in your notebook instance.

Alternatively, you can export your notebook as an executable script, modify the generated .py file using the previous steps, and then deploy your pipeline to the Dataflow service.

Saving your notebook

Notebooks you create are saved locally in your running notebook instance. If you reset or shut down the notebook instance during development, those new notebooks are deleted. To keep your notebooks for future use, download them locally to your workstation, save them to GitHub, or export them to a different file format.

Cleaning up

After you've finished using your Apache Beam notebook instance, clean up the resources you created on Google Cloud by shutting down the notebook instance.