Run Spark Jobs on Azure Batch using Azure Container Registry and Blob storage
By: Kris Peeters
Tags: Docker, Azure, Spark, Data Engineering

This is another one of those “how to” blogs that can hopefully help people get up-and-running quickly because it took me a while to figure it out. Here is the full code if you want to try this yourself: https://github.com/datamindedbe/spark_on_azure_batch_demo. Obviously, this is all just demo code, with a lot of hardcoded strings and shortcuts. This is not production-grade. So take what you can use, and adapt it to your context.

Situation: Give me cheap Spark data pipelines!

You want to run spark jobs on Azure. But cost is a concern. So you’re looking into Azure Batch because you can use Low-Prio VMs, which are not available on HDInsights or Azure Databricks Cloud. See pricing here: https://azure.microsoft.com/en-us/pricing/details/batch/. You quickly realise that you can save up to 80%, which is significant. We already use a lot of SPOT instances on AWS (the equivalent of low prio VMs), even in Production workloads, because we aim to write our jobs in such a way, that they can easily recover from failure, and simply run again. With Compute typically being the biggest cost in big data workloads, this will have a very positive impact on the bottom line.

Current pricing for Low Prio VMs on Azure. From $0.16 to $0.03 for a 16GB VM.

Here’s another piece of the puzzle: you have to run a bunch of daily jobs, which individually, are not really that big. But you have to run it for 100 days. You can look into AZTK which can run large spark clusters on Azure Batch. But I don’t really need one big Spark cluster. Just give me 100 single node Spark engines.

This is where Azure Batch comes in. You simply wrap your Spark code in a Docker container, and schedule 100 of those containers on a pool of 20 Azure Batch nodes, and let them execute the 100 jobs in parallel. You get to use Low Prio VMs out of the box, and you only pay for what you need. No long-running Spark or Hadoop clusters.

Required setup to start

In order to run this demo, you need to set up:

An Azure Storage Account to read and write your data

An Azure Container Registry to store the Docker container

An Azure Batch Account to run your batch jobs in.

Chances are you already have these running in Azure. If not, it’s simple click and install on the Azure Portal. For each of these 3 services, you need to get the Keys and add them to a config.py:

STORAGE_ACCOUNT_NAME = "datamindeddemodata"

STORAGE_ACCOUNT_KEY = "XYZ"



BATCH_ACCOUNT_NAME = "dataminteddemobatch"

BATCH_ACCOUNT_KEY = "XYZ"

BATCH_ACCOUNT_URL = "https://dataminteddemobatch.westeurope.batch.azure.com"



ACR_LOGINSERVER = "datamindeddemo.azurecr.io"

ACR_USERNAME = "datamindeddemo"

ACR_PASSWORD = "XYZ"

Write the Spark job

For this demo, we need a simple spark job. Any job will do, really, The only thing you need to be aware of, is that it has to be able to read from Blob Storage. So we need some extra libraries when we initialise Spark. In this case, I’m doing a simple aggregation on airline data, which you can find here.

import argparse



from pyspark.sql import SparkSession



import config





def get_azure_spark_connection(storage_account_name, storage_account_key):

spark = (

SparkSession.builder

.config('spark.jars.packages', 'org.apache.hadoop:hadoop-azure:2.7.3')

.config('spark.hadoop.fs.azure', "org.apache.hadoop.fs.azure.NativeAzureFileSystem")

.config("spark.hadoop.fs.azure.account.key." + storage_account_name + ".blob.core.windows.net",

storage_account_key)

.appName("AzureSparkDemo")

.getOrCreate())



(spark.sparkContext._jsc.hadoopConfiguration().set("fs.wasbs.impl",

"org.apache.hadoop.fs.azure.NativeAzureFileSystem"))

return spark





if __name__ == '__main__':

parser = argparse.ArgumentParser()

parser.add_argument("-i", "--input", help="input file to parse", type=str)

parser.add_argument("-o", "--output", help="result file to write", type=str)

args = parser.parse_args()

spark = get_azure_spark_connection(config.STORAGE_ACCOUNT_NAME, config.STORAGE_ACCOUNT_KEY)

df = (spark.read.option("header", "true")

.option("delimiter", ",")

.option("inferSchema", "true")

.csv(args.input))

df.registerTempTable("airlines")

result = spark.sql("""

select Year, Month, DayofMonth, avg(ArrDelay) as avg_ArrDelay, avg(DepDelay) as avg_DepDelay

from airlines

group by Year, Month, DayofMonth

""")

result.repartition(1).write.mode("overwrite").parquet(args.output)

What I’m doing here, is taking a full year of airline data, grouping it by year, month, day, and then calculate the average arrival and departure delay. A very basic example. To get this to work, you need a few azure and pyspark dependencies, which you typically add to requirements.txt and then do pip install -r requirements.txt , ideally in a virtual environment (which is out-of-scope for this demo):

azure

azure-storage

azure-storage-blob

pyspark==2.4.0

Dockerize it

The next step is that we wrap this job in a Docker container and push it to the Azure Container Registry. A simple Dockerfile can be as follows:

FROM python:3.6



# Install OpenJDK 8, and monitoring tooling

RUN \

apt-get update && \

apt-get install -y openjdk-8-jdk htop bmon && \

rm -rf /var/lib/apt/lists/*



RUN pip install --upgrade setuptools

COPY requirements.txt /

RUN pip install -r /requirements.txt



ENV PATH=$PATH:/src

ENV PYTHONPATH /src



ADD ./ /src

WORKDIR /src/

It starts from a python3.6 base image, installs openjdk and some monitoring tooling, and then all the python dependencies as listed in requirements.txt . Finally, it copies your code to a src folder. You can build it with the command:

docker build -t dataminded/spark_on_azure_batch_demo .

Now you could already run this script locally as follows:

This simply says to docker: please run the above docker image, and within that image, execute the python command with input a CSV file and output a Parquet file on Azure Blob storage. We’ll need this later to schedule the same command on Azure Batch.

Next, you need to log in to the Azure Container Registry, with datamindeddemo replaced by your own registry of course. Enter the admin username you see on the Azure portal:

docker login datamindeddemo.azurecr.io

Finally you can tag your docker container:

docker tag dataminded/spark_on_azure_batch_demo:latest datamindeddemo.azurecr.io/datamindeddemo/spark_on_azure_batch_demo:latest

And upload it to ACR:

docker push datamindeddemo.azurecr.io/datamindeddemo/spark_on_azure_batch_demo:latest

Again, replacing this with your own ACR path and credentials. Note that the first time you push this Docker container, it will be huge, about 1GB: it has an ubuntu image, a bunch of spark libraries, a JDK, … The next push you do of this container, will only push your latest code changes.

Run it on Azure Batch

Ok, so we’ve written a sample Spark application, with the right Azure libraries. We’ve dockerized it and uploaded it to Azure Container Registry. Now we’re finally ready to run our code on Azure Batch. This is going to be quite a long python script, but it does a lot of things as well:

It connects to your container registry and uses those docker images

Create a Pool if it doesn’t exist yet. Here, you can configure which kind of VMs and how many of them you want in your pool. And more importantly, you can specify that it are Low Prio VMs, which are cheap.

Create a Job within the Pool

Create a separate task to process each year of data. In a real-life situation, you would have a task for each day of data.

import datetime

import sys

import time



import azure.batch.batch_auth as batch_auth

import azure.batch.batch_service_client as batch

import azure.batch.models as batchmodels



import config



sys.path.append('.')

sys.path.append('..')



IMAGE_NAME = 'datamindeddemo.azurecr.io/datamindeddemo/spark_on_azure_batch_demo'

IMAGE_VERSION = 'latest'





def create_pool(batch_service_client, container_registry, image_name, pool_id, pool_vm_size, pool_node_count,

skip_if_exists=True):

print('Creating pool [{}]...'.format(pool_id))



container_conf = batch.models.ContainerConfiguration(

container_image_names=[image_name],

container_registries=[container_registry]

)



image_ref_to_use = batch.models.ImageReference(

publisher='microsoft-azure-batch',

offer='ubuntu-server-container',

sku='16-04-lts',

version='latest')



new_pool = batch.models.PoolAddParameter(

id=pool_id,

virtual_machine_configuration=batch.models.VirtualMachineConfiguration(

image_reference=image_ref_to_use,

container_configuration=container_conf,

node_agent_sku_id='batch.node.ubuntu 16.04'),

vm_size=pool_vm_size,

target_low_priority_nodes=pool_node_count)



if not skip_if_exists or not batch_service_client.pool.exists(pool_id):

batch_service_client.pool.add(new_pool)





def create_job(batch_service_client, job_id, pool_id):

print('Creating job [{}]...'.format(job_id))



job = batch.models.JobAddParameter(

id=job_id,

pool_info=batch.models.PoolInformation(pool_id=pool_id))



batch_service_client.job.add(job)





def add_task(batch_service_client, image_name, image_version, job_id, command, name):

user = batchmodels.UserIdentity(

auto_user=batchmodels.AutoUserSpecification(

elevation_level=batchmodels.ElevationLevel.admin,

scope=batchmodels.AutoUserScope.task))



task_id = name

task_container_settings = batch.models.TaskContainerSettings(

image_name=image_name + ':' + image_version,

container_run_options='--rm -p 4040:4040')

task = batch.models.TaskAddParameter(

id=task_id,

command_line=command,

container_settings=task_container_settings,

user_identity=user

)

print("running " + command)



batch_service_client.task.add(job_id, task)





def wait_for_tasks_to_complete(batch_service_client, job_id, timeout):

timeout_expiration = datetime.datetime.now() + timeout



print("Monitoring all tasks for 'Completed' state, timeout in {}..."

.format(timeout), end='')



while datetime.datetime.now() < timeout_expiration:

print('.', end='')

sys.stdout.flush()

tasks = batch_service_client.task.list(job_id)



incomplete_tasks = [task for task in tasks if

task.state != batchmodels.TaskState.completed]

if not incomplete_tasks:

print()

return True

else:

time.sleep(1)



print()

raise RuntimeError("ERROR: Tasks did not reach 'Completed' state within "

"timeout period of " + str(timeout))





def daterange(start_date, end_date):

for n in range(int((end_date - start_date).days)):

yield start_date + datetime.timedelta(n)





if __name__ == '__main__':



start_time = datetime.datetime.now().replace(microsecond=0)

print('Start batch job {}'.format(start_time))

print()



image_name = IMAGE_NAME

image_version = IMAGE_VERSION



# Create the blob client, for use in obtaining references to

# blob storage containers and uploading files to containers.



credentials = batch_auth.SharedKeyCredentials(config.BATCH_ACCOUNT_NAME,

config.BATCH_ACCOUNT_KEY)



batch_client = batch.BatchServiceClient(credentials, base_url=config.BATCH_ACCOUNT_URL)



job_id = 'Job-' + '{:%Y-%m-%d-%H-%M-%S}'.format(datetime.datetime.now())

pool_id = 'Airlines'



try:

container_registry = batch.models.ContainerRegistry(

registry_server=config.ACR_LOGINSERVER,

user_name=config.ACR_USERNAME,

password=config.ACR_PASSWORD)



create_pool(

batch_service_client=batch_client,

pool_id=pool_id,

container_registry=container_registry,

image_name=image_name,

pool_node_count=3,

pool_vm_size='Standard_E2s_v3',

skip_if_exists=True)



# Create the job that will run the tasks.

create_job(batch_client, job_id, pool_id)



for year in range(2001, 2009):

command = "python /src/airline_analytics.py --input wasbs://demo@datamindeddata.blob.core.windows.net/raw/airlines/{0}.csv.bz2 --output wasbs://demo@datamindeddata.blob.core.windows.net/aggregated/airlines/{1}.parquet" \

.format(year, year)

add_task(

batch_service_client=batch_client,

image_name=image_name,

image_version=image_version,

job_id=job_id,

command=command,

name='airlines{0}'.format(year))



# Pause execution until tasks reach Completed state.

wait_for_tasks_to_complete(batch_client, job_id, datetime.timedelta(hours=2))



print(" Success! All tasks reached the 'Completed' state within the specified timeout period.")



except batchmodels.BatchErrorException as err:

print(err)

raise



# Print out some timing info

end_time = datetime.datetime.now().replace(microsecond=0)

print()

print('Sample end: {}'.format(end_time))

print('Elapsed time: {}'.format(end_time - start_time))

print()

print()

And finally, you will see your jobs run on Azure Batch:

The pool with 3 Low Prio VMs, costing 0.10$ in total

The jobs are running year by year. Because I have 3 nodes, I have 3 jobs running in parallel. If I had 20 nodes, I would be able to process 20 jobs in parallel

Note that all nodes in a pool are completely separate. This is not a one big Spark cluster concept. You can log in to those individual nodes with a handy connect button:

The connect button on Azure Batch Explorer

By the way, Azure Batch Explorer is downloadable from github: https://azure.github.io/BatchExplorer/. It makes things slightly more convenient than looking in the azure Portal. Once connected, you can see how those nodes are performing using a tool such as htop.

You can find all code to reproduce this here: https://github.com/datamindedbe/spark_on_azure_batch_demo

Next steps

This solves 90% of the needs of a project we’re working on. However, there are always next steps: