Industrialization of a Machine Learning model using Apache Airflow and Apache BEAM
By: Massy Bourennani
Tags: Data Engineering, Apache Beam, Google Cloud Platform, Machine Learning, Apache Airflow

Figure-4 is a representation of the DAG that schedules the computation of predictions. The frequency is defined by the DAG’s schedule_interval which takes the value 0 1 * * * to specify a daily execution at 01:00 AM.

Each step represents an operator. start and end are dummy operators, the most interesting ones are described in the following sections.

github_bq_query

This operator is responsible of executing a SQL query, located on Github. It is identified by a path in the repository and a tag which allows us to execute a particular version of that query file. The SQL can be templated in Airflow.

Internally the operator is a specialization of BigQueryOperator where the self.sql parameter is retrieved from GitHub (instead of a raw string or local file). The render_template_fields of BaseOperator is used as a hook to know when to exactly get the query and save it back to self.sql .

The result of the query is then saved in a table by specifying the destination_dataset_table let it be my_project.my_datset.churn for our example.

bq_to_gcs

The content of my_project.my_dataset.churn is then saved on GCS in new line delimited JSON (NDJSON) format.

start_beam_job

In this operator an Apache BEAM pipeline is started. It reads NDJSON files one line at a time, transforms each one to a dict, applies the model and saves back the result to GCS after converting it to string. We show the pipeline below

Figure-5: Apache BEAM pipeline

Below is the code of the main.py reflecting the pipeline.

The most interesting part resides in the ExecuteModelDoFn where the model.pkl is taken from GCS, deserialized and later used to compute the score of each data point.

model.pkl is copied from Github to GCS. It is identified by two things the path to the model and the tag, like what is described in Figure-2 which allows us to execute a specific version of the model. This operation is done in a previous step of the Airflow DAG.

Figure-5 shows that the model is applied to one item at a time which is not very efficient. In order to decrease latency of the predictions we can batch in the DoFn and compute the predictions one batch at a time.

The code below details the implementation of ExecuteModelDoFn

The setup() method is executed only once by each worker on Dataflow. In that method model.pkl is downloaded from GCS, deserialized and saved in the the instance attribute self.model .

The from google.cloud import storage inside the method may seem peculiar at first sight. We must not forget that beam_pipeline.py is run in Airflow and the pipeline object is serialized and sent to Dataflow along with the DoFn . Only the __init__ method is executed in the context (think pyhon packages) of Airflow. Putting the import in the beginning of the file constrains us to have google-cloud-storage installed in Airflow workers (and scheduler).

Thanks to Apache BEAM and Dataflow we have the scalability we need to execute the model on millions of rows. In addition to that, the computation of the score of one observation is independent from the others. Consequently, the process() method is executed by many Dataflow workers on different elements which allows us to reduce prediction’s latency of the whole dataset.

What about flexibility ? We must allow data scientist to use the same libraries used to train the model and some custom packages too. Knowing that the execution context of process() is Dataflow worker, we need to provide them with the right dependencies. This is done using the setup.py

gcs_to_bq

The predictions are then sent from GCS to BigQuery using GoogleCloudStorageToBigQueryOperator.

Conclusion

We presented the automation of the prediction phase of ML models based on python. In order to orchestrate the workflow we used Apache Airflow. BEAM and Dataflow were used to apply the model on the whole dataset allowing to scale to millions of data points. Flexibility is also achieved by providing Dataflow workers with the right dependencies using the pipeline’s setup.py . A particular version of the model is retrieved from Github using tags. BigQuery and GCS are used as storage components which allows us to separate data and compute.

Monitoring is achieved using Airflow’s web UI and the alerting (sending a slack message, email …etc) is done by specifying on_failure_callback in the DAG’s default_args .

References

[1] https://cloud.google.com/ai-platform/docs/ml-solutions-overview