How to run CPU-based Workloads for Deep Learning Using Thousands Of Spot Instances on AWS and GCP Without Getting a Headache
By: Assaf Pinhasi
Tags: Data Engineering, Aws Batch, Deep Learning, Pipeline, Mlops

How to run CPU-based Workloads for Deep Learning Using Thousands Of Spot Instances on AWS and GCP Without Getting a Headache Assaf Pinhasi Jul 26·10 min read

Photo by Andrey Sharpilo on Unsplash

Deep learning is notorious for consuming large amounts of GPU resources during training. However, there are multiple parts within the Deep Learning workflow that require large amounts of CPU resources :

Running large scale inference jobs Pre-processing input data — and materializing it on disk as

preparation for training

These workloads typically have the following properties:

The workload (job) is triggered periodically (vs. always-on processing) A job is comprised of many items. Each item can be processed independently from other items. Items are read, processed and written back to some storage (typically object storage or a centralized file system) Processing a single item may take anywhere from several seconds to minutes on a single CPU core Users care about the throughput of the job vs. the latency of processing a single item Single item size ranges from tens of KB’s to many MB’s or even GB of data The processing is stateless — i.e. if the processing of a single item failed, it can be safely retried without any side effect.

When you have enough items to process, and when the processing of a single item is heavy enough, these jobs can consume a lot of CPU resources.

How much is a lot?

Example workload — 3D rendering of training data

Let’s say we want to train a model with large amounts of images of real world objects taken in different environments. A great way to produce large amounts of examples is to use synthetic data — take 3D models of objects, and render views of these models from multiple angles, lighting conditions etc.

3D rendering is a notoriously CPU-intensive workload; rendering a single image can take minutes on a single CPU core. Rendering 250K images can consume 30K+ CPU hours.

How can we deal with such workloads?

First, let’s formalize our requirements.

System requirements

General

Scale-out to several thousand cores Minimize cloud costs by running on spot instances Minimize infrastructure setup and maintenance costs

API for running jobs

Configure the processing logic per job Configure the CPU +memory needed per item in the job Configure where to read/write data from Submit all the items of the job to be processed Cancel a job (i.e. do not process its items) Observe # of successful/remaining/failed items per job Observe CPU/memory consumption (to help tune resources) View logs per item (nice to have — search logs)

Logical System Design

The requirements are very suitable for a queue + workers design pattern.

Logical batch processing system design. Image by the author

Logical flow

Let’s walk through the diagram and understand what happens at each stage:

1 — The user pushes a container into a container registry.

The container includes the processing logic and dependencies needed.

2 — The user enqueues a message containing an item to process.

Each message contains:

The URI of an item in the storage

Optionally — configuration/metadata to control the processing logic

3 — If needed, the compute auto-scales.

When ready, a message is dequeued and handed to an available container.

4 — The container reads an item from the storage, processes it and writes the output back based on the instructions in the message

Lastly, the system automatically collects metrics and logs into a centralized location.

The hard way — build it yourself on Kubernetes

Here are a few of the steps you’ll need to take to achieve this with K8s:

First — you’ll need to set up K8s to be able to manage thousands of nodes — from redundancy to right-sizing the control plane etc. Create nodegroups that run on spots with a mixture of policies Deploy a message bus — either on the cluster or outside of it Figure out how to scale out :

should you let the user specify the scale in advance and deploy a replica set + find some way to scale down automatically by updating the size and using cluster autoscaler to kill idle nodes?

or use a dynamic scaler like KEDA? If the storage is a shared file system, you’ll need to create persistent volumes and mount them into the containers ourselves; if it’s object storage we may need to pass in the credentials as a secret Add a full observability stack with centralized logging and monitoring Turn this into infra-as-code with Terraform etc.

This is a no-starter for most teams as the effort is very significant — especially the need to spin and and manage hundreds or thousands of nodes in the cluster, even temporarily

Edit — Other cluster-based solutions

This section was added following a lively discussion on LinkedIn.

There are other solutions apart from K8s that help manage scale-out workloads — such as Ray, Dask, and even Spark.

What all of these tools have in common is that they have the notion of a cluster, built for workloads where machines need to exchange data between them — and at the very least require some form of communication between some of the machines (e.g. master to workers).

When you scale to 1000’s of nodes, this management and communication can introduce various edge cases that are not trivial to handle. If you deploy the cluster, you own it. Even in “managed” clusters like EKS, you still need to ensure it’s running properly.

In inference/pre-processing — there is no added value in running on a cluster where nodes communicate. These are share-nothing workloads. The cluster component is just a liability. Instead, what you want is to provide your code and a list of items to process, and get the compute to take care of it for you….

This leads us to the easier option…

Note: If your workload is written in something other than Python, or requires unusual OS-level libraries, not all of the cluster managers mentioned above can support it.

The easy way — leveraging managed solutions

AWS

AWS has built the perfect tool for this type of workload, aptly named

AWS Batch.

AWS Batch solution for large scale batch processing system. Image by the author

Let’s walk through the entities in the diagrams and explain how this works:

Compute environment

A collection of compute resources; can contain a mixture of node types, including spots or on-demand.

The compute environment specifies its mix/max size — where setting the minimum to 0 enables it to scale to zero when there is no work to perform.

When scaling out, the compute environment creates instances based on instance templates.