Operationalization of ML Pipelines on Apache Mesos and Hadoop using Airflow
By: Mihail Vieru
Tags: Mesos, Airflow, Hadoop, Machine Learning

Operationalization of ML Pipelines on Apache Mesos and Hadoop using Airflow

An architecture for bringing ML models into production at NEW YORKER

This blog post originally appeared on the NEW YORKER Tech Blog.

Photo by Simon Migaj on Unsplash

One of the highest impact use cases for machine learning in retail companies is markdown pricing. As one of Europe’s largest fashion retailers, NEW YORKER sets an initial price for each item sold in our 1100 stores in 45 countries. The price is gradually marked down in a series of price reductions until the item is sold out. The goal is to predict when and how to mark down the price, so that the whole inventory is sold at the highest possible price. Thus, maximizing the company’s revenue.

Solving this use case involves two tasks: creating a model which delivers accurate predictions, and operationalizing it into a production pipeline on our current infrastructure.

In this post we focus on the latter task, and present an architecture for orchestrating machine learning pipelines in production on Apache Mesos and Hadoop using Airflow. We start with the requirements for the machine learning model along with the ones posed by our current infrastructure and data engineering architecture. We then present an architecture which fulfills those, including containerization, pipeline orchestration with Airflow, performance considerations and the simplified DAG code. We conclude with a summary and an outlook on future work.

Requirements

Let’s start with the requirements of the model we wish to operationalize, along with the ones stemming from our current infrastructure and data engineering architecture.

The training data for the model are stored as Parquet files on HDFS in our Hadoop cluster. They are the result of extensive prior batch data processing with Spark, including data cleaning, aggregation and imputation where necessary. This legacy ETL is not orchestrated with Airflow yet and uses Luigi for the workflow management and Cron Jobs for the scheduling.

The model is written in Python and the machine learning framework employed should be up to the choice of our Data Scientists. Currently LightGBM is used. The training and inference are done on a per country basis — a characteristic which we can leverage later for parallelization. GPU support is important, as our Data Scientists are also considering deep learning approaches, which would profit from it.

The resulting predictions of the model need to be written to our ERP system’s persistence layer — a relational database. The predictions serve as a guide to our Pricing department, which checks them via the ERP’s UI for plausibility and applies the suggested price reductions where and when it deems them appropriate. As the predictions change very slowly, a refresh is required once at the beginning of each week.

Besides the aforementioned Hadoop cluster, we employ a DC/OS cluster in our infrastructure. It provides us with container orchestration and resource scheduling based on Apache Mesos. GPUs can be used in containers via Marathon application and Metronome job definitions for long and short running applications. We have a containerized Airflow installation running on DC/OS for orchestrating most of our data pipelines.

Architecture

For fulfilling the requirements stated above, we have come up with the following architecture.

Containerization

We’ve decided to containerize the model training and inference computation and use DC/OS to run it. This has two advantages. First, it preserves the independence of the machine learning framework chosen by our Data Scientists. Second, we can leverage existing GPUs on the DC/OS cluster to speed up future deep learning approaches.

If we would have chosen to use SparkML or parallelize another ML framework with Spark on Hadoop, we would have had to migrate the existing LightGBM model to this approach. We would have then essentially locked in the ML framework employed. Furthermore, to facilitate GPU usage, we would have needed to install GPUs on the Hadoop cluster or move them from the DC/OS cluster to it. A Hadoop version upgrade to 3.x from the present 2.x would have also been required. All this would have been non trivial and very time consuming.

The container itself has the following workflow:

At the beginning, it downloads the training data from HDFS. Then it trains the model on it and computes the predictions for a given country. Finally, the predictions are refreshed in the ERP’s underlying relational database.

The refreshing of the predictions is done in a single transaction, i.e. it deletes old predictions for all items for a single country and it writes new ones. We don’t use updates here, as the set of predictions can vary across time. In this way, the ERP’s users will always see predictions in the UI for the selected set of items with no downtime.

We’ve parameterized the container, so that it can calculate and refresh predictions not only for a single country but for a specified group of countries.

Pipeline Orchestration

Having a parameterized container for the model training and inference in place, we can now proceed to deploy and orchestrate multiple instances of it to form a proper pipeline with the help of Airflow:

For short-running applications, i.e. jobs, DC/OS provides Metronome as a task scheduler on top of Mesos. As our container is ephemeral in nature, we can employ Metronome to deploy it on the DC/OS cluster. We do this from Airflow using the MetronomeOperator, which we’ll introduce more in-depth shortly.

We start multiple instances of the MetronomeOperator for groups of countries so that all 45 countries are covered, thus effectively parallelizing the work.

As the legacy ETL — which processes the needed training data on Hadoop — is not orchestrated with Airflow yet, we employ a modified HDFS sensor in order to trigger the start of the MetronomeOperator instances. The original Airflow HDFS sensor was only functioning with Python 2. We’ve made some small changes so that it works with the Python 3 version that we use in our containerized Airflow installation.

The code for it can be found here:

https://github.com/NewYorkerData/ny-public-airflow-operators

We schedule the pipeline to start simultaneously with the legacy ETL on Hadoop. The HDFS sensor will continuously poll for the success file and trigger the remaining operators in the DAG when the file is found.

Note: As part of the production pipeline we also write the predictions to HDFS for future analysis. During the computation thereof we log all steps to Elasticsearch. We have Kibana dashboards in place and currently use elastalert to monitor and alert on anomalies, e.g. the number of predictions decreased by 10% from last run.

Before promoting changes to production, we use Neptune as part of the test pipeline to keep track and evaluate the impact of the various changes to the model itself.

Metronome Airflow Operator

No Metronome operator existed for Airflow to the best of our knowledge at the inception of the project. So we created one.

It takes a Metronome job definition as a JSON parameter and works as depicted in the figure below:

First, we get an authorization token from the IAM REST API to be able to deploy containers on the DC/OS cluster. Second, we check whether the job exists via the Metronome REST API. Depending on the response, we upsert the job definition. Then we start the job. Finally, we poll for the job status and, depending on the response, change the Airflow operator’s status to SUCCESSFUL or FAILED.

The code for the operator can be found here:

https://github.com/NewYorkerData/ny-public-airflow-operators

Performance Considerations

In the Metronome job definition we can specify CPU, RAM and GPU usage for the container. Therefore we can scale vertically. Regarding the horizontal scaling, we do this by parallelizing the model training and inference for the groups of countries. The degree of parallelism is however not unbounded and is limited by the available cluster resources. In order not to use up all resources, we limit the number of concurrently running containers using the Airflow resource pool feature. The resource pool name is set in all MetronomeOperator instances, as we’ll see in the DAG code below. The size of the resource pool is thus equal to the degree of parallelism.

Implementation as Airflow DAG

Below you can find the simplified version of the Airflow DAG:

default_args = {

"owner": "ny-data-science",

"start_date": datetime(2020, 1, 1),

"provide_context": True

} dag = DAG(

"markdown-pricing-dag",

schedule_interval="0 8 * * 1", # every Monday at 8am

dagrun_timeout=timedelta(days=1),

default_args=default_args,

max_active_runs=1,

catchup=True

) # Example list of countries for the model training and inference

COUNTRY_LIST = ["DEU", "AUT", "NLD", "FRA", "ESP", "ITA"] # Job template we will use in the MetronomeOperator

METRONOME_JOB_TEMPLATE = """

{{

"id": "{}",

"description": "Markdown Pricing ML Job",

"run": {{

"cpus": 4,

"mem": 32768,

"gpus": 0,

"disk": 0,

"ucr": {{

"image": {{

"id": "registry-dns-address:1234/markdown-pricing:latest-master",

"forcePull": true

}}

}},

"env": {{

"COUNTRY_GROUP": "{}",

"DB_USER": "{}",

"DB_PASSWORD": "{}"

}}

}}

}}

""" # Creates a Metronome job on DC/OS by instantiating the MetronomeOperator with the job template above

# and setting the country group and other environment variables

def create_metronome_job_for_country_group(country_group, index):

return MetronomeOperator(

task_id=f"metronome_operator_{index}",

metronome_job_json=

METRONOME_JOB_TEMPLATE.format(f"markdown-job-{index}",

country_group,

Variable.get("markdown_erp_db_user"), Variable.get("markdown_erp_db_pwd")),

dcos_http_conn_id="dcos_master",

dcos_robot_user_name=Variable.get("robot_user_name_dcos"),

dcos_robot_user_pwd=Variable.get("robot_user_pwd_dcos"),

dag=dag,

pool="markdown_metronome_job_pool",

retries=3

) # Get the resource pool size (in slots) for the MetronomeOperator instances from Airflow configuration

metronome_job_pool_size = get_pool_size(pool_name="markdown_metronome_job_pool") # Split the country list into groups into N parts of approximately equal length for parallelization purposes.

# N is here the size of the Metronome job pool.

# Given the COUNTRY_LIST defined above and N = 3, the function will return: [["DEU","AUT"], ["NLD","FRA"], ["ESP","ITA"]]

country_groups = split_country_list(COUNTRY_LIST, metronome_job_pool_size) # Iterates through the country groups and creates a Metronome job for each of those groups

metronome_countries_jobs = [create_metronome_job_for_country_group(country_group=country_group, index=index) for

index, country_group in enumerate(country_groups)] # HDFS sensor on the latest training data

training_data_sensor = NYHDFSSensor(

task_id="training_data_sensor",

filepaths=[f"/data/production/markdown_training_data/{get_current_date()}/_SUCCESS"],

hdfs_conn_id="hdfs_conn_default",

retries=1440,

retry_delay=timedelta(minutes=1),

timeout=0,

dag=dag) # Create DAG

training_data_sensor >> metronome_countries_jobs

Conclusion and Future Steps

We have shown an architecture which allows us to easily orchestrate machine learning pipelines in a mixed Mesos and Hadoop cluster environment using Airflow. We leverage Mesos’ container orchestration and resource scheduling to scale the model training and inference both horizontally and vertically, while being able to take advantage of available hardware acceleration provided by GPUs. We employ Airflow’s powerful features, such as sensors and dynamic DAGs, to manage the whole workflow effectively across the clusters.

There are still many improvement areas we can work on in the future.

For instance, migrating the legacy ETL on Hadoop from Luigi/Cron to Airflow will allow us to deprecate the HDFS sensor and simplify the DAG by using Airflow’s pipeline interdependency feature.

We’re also aware that some computation is common across all countries. This step could be extracted in a separate stage, on which all further stages would depend on, possibly improving the performance.

Finally, we’re in the process of preparing the migration of our clusters from Mesos to Kubernetes. In the short-term we will need to adjust all our machine learning pipelines in Airflow. This should be in theory relatively easy to accomplish, e.g. use the KubernetesOperator instead of the MetronomeOperator in the DAGs. In the longer term, Kubernetes opens up a new world of possibilities in terms of machine learning operationalization approaches, such as Kubeflow.