Wrangling a 100TB Data Set for Machine Learning

Also — Ryan (our Founder) is giving a talk on this topic on June 11th, 2024. If you are in the Boston area, swing by!

The First Major Project

Our first major task at Fan Pier Labs was to process a massive 100-terabyte dataset (100 million images) and prepare it for training machine learning models. Handling such a large volume of data is challenging, but we successfully built a cloud infrastructure capable of processing all that data in just 12 hours per job. Not only could our infrastructure handle this enormous amount of data, but it was also designed to (hopefully) scale indefinitely!

Shameless plug: if you need a hand with managing a massive data set, contact us!

About the Client

We built this project for a YCombinator — backed startup that wanted to train their own model to generate images from user-entered prompts (much like MidJourney or DALL-E). In order to develop this product, they decided to train their own models instead of using a pre-trained base model.

The Starting Point

Our client had a large data set of over 100 million images and was storing them in an S3 bucket. They didn't have anything set up outside of the S3 bucket.

The Challenge

In order to prepare the images for training, we had to process each image multiple times to extract a number of additional pieces of information. We used some of this information to filter out images that we didn't want to train on, and other pieces we used in the training step itself.

Deciding on a Database

First, we had to pick a database to store all of the information about the images. There were four databases that stuck out as good candidates. These were Redshift, Postgres, Spanner and BigTable.

Redshift is well equipped to handle this project because it excels at storing large quantities of data, has a strong ecosystem, and is, of course, well supported by AWS. Just like Spanner and BigTable, Redshift can scale to petabytes of data into a single table. This is possible because Redshift can shard a single large table horizontally across a fleet of servers, instead of relying on vertical scaling like older databases do. Redshift is also a columnar database, which means that it can skip non-relevant data when running a query that processes just one or two columns. However, a columnar database tends to be more difficult to insert data into. Amazon recommends against using the INSERT command for large quantities of data, and instead, recommends uploading the data to AWS S3 and then using a special COPY SQL command to copy it into Redshift.

Postgres also stands out because it is a very powerful database that has a strong ecosystem. AWS has great support for running and managing Postgres instances inside of AWS Relational Database Service (RDS). Postgres, however, is a row-oriented database, which means it will always read an entire row of data from disk when you query any column in that row.

Spanner and Bigtable would also be able to handle this workload. However, most of the rest of our client's infrastructure was already located on AWS, so we decided to stick with a database that was also on AWS.

We decided to use Postgres for this project due to its popularity, robust ecosystem, ease of setup with AWS RDS, and excellent support within AWS. However, if we were to undertake this project again, using Redshift might be worth considering.

Setting up the Database

The first step was to add the S3 Object URLs to the database. We assigned a unique UUID to each image and renamed the images in the storage bucket to match the assigned UUID (plus the image extension).

Structuring the Database

We first created the table with only two columns — one to hold the UUID of each image and another to hold the S3 Object URL of each image. Later, we added more columns to keep track of the information derived from each image.

Example table:

uuid                             | s3_path
---------------------------------+-----------------------------------------------------
4283d3832db64599b619a464197b7eb6 | s3://my-bucket/4283d3832db64599b619a464197b7eb6.jpg
a04d043d837844838751d10ac8f9183a | s3://my-bucket/a04d043d837844838751d10ac8f9183a.jpg

let s3_path = "s3://my-bucket/" + uuid + ".jpg"

Which additional fields do we need?

We needed to calculate additional columns for various information about the images:

• Extract EXIF data (Camera make/model, flash, aperture, shutter speed, focal length, ISO, location)
• File Extension of each image
• OCR text from the image
• Height and Width of image
• Size of images in bytes
• Hash of image for de-duplication
• Number of people in the image
• AI-Aesthetic ranking (GPU required)
• AI-description of the image (GPU required)

Adding additional columns to the Database

ALTER TABLE my_table ADD COLUMN extension varchar(4);
ALTER TABLE my_table ADD COLUMN aesthetic_score int;
-- etc...

Populating the additional columns (the wrong way!)

The quick and dirty way is a Python script that processes rows one at a time:

import boto3
import pg

conn = pg.connect({...})
cursor = conn.cursor()
s3 = boto3.client('s3')

# DON'T RUN THIS CODE ON MULTIPLE COMPUTERS, IT DOESN'T SCALE
def main():
    sql = "SELECT uuid FROM my_table WHERE description is null limit 1000"
    rows = pg.run(sql)
    for row in rows:
        uuid = row['uuid']
        s3_path = "s3://my-bucket/" + uuid + ".jpg"
        response = s3.get_object(Bucket=s3_bucket_name, Key=s3_key)
        data = response['Body'].read().decode('utf-8')
        processed_data = process_data(data)
        sql = "UPDATE my_table SET description = %s WHERE uuid = %s"
        cursor.execute(sql, (processed_data, uuid))

Problem: it doesn't run fast enough and cannot scale. Running on multiple servers causes them to process the same images due to no coordination.

Enter Queues

Queues allow us to coordinate between servers. We considered Celery, AWS SQS, and RabbitMQ. We went with AWS SQS because it integrates well with the AWS ecosystem.

The only concerning limitation of SQS is that it can only retain each job in the queue for 14 days.

Breaking down the data into discrete buckets

We assign each image to a bucket based on the first few characters of each UUID. Using the first 4 hex characters gives us 65,536 buckets:

16^4 = 65,536 batches
100 million images / 65,536 batches = ~1,526 images per batch

import boto3

queue_name = 'your_queue_name'
sqs = boto3.client('sqs')
response = sqs.create_queue(QueueName=queue_name)
queue_url = response['QueueUrl']

def generate_hex_strings():
    for i in range(65536):
        hex_string = format(i, '04x')
        yield hex_string

for hex_string in generate_hex_strings():
    sqs.send_message(QueueUrl=queue_url, MessageBody=hex_string)

Optimizing the lookup queries

We added hash indexes for ultra-fast lookups:

CREATE INDEX my_table_bucket_4_hash ON my_table USING HASH (bucket_4);
CREATE INDEX my_table_bucket_5_hash ON my_table USING HASH (bucket_5);
CREATE INDEX my_table_bucket_6_hash ON my_table USING HASH (bucket_6);

Query:

SELECT uuid FROM my_table WHERE bucket_4 = 'ffff' AND job_column IS NULL;

The Main Processing Code

queue = sqs.get_queue_by_name(QueueName=SOURCE_QUEUE)

def process_one_bucket():
    message = json.loads(queue.receive_messages(MaxNumberOfMessages=1)[0].body)
    bucket = message['prefix']

    query = 'select uuid from my_table where bucket_4 = ' + bucket + ' and job_column is null'
    rows = runPostgresQuery(query)

    for row in rows:
        buffer = BytesIO()
        buffer = s3.download_fileobj("images-for-training", uuid, buffer)
        buffer.seek(0)
        new_column_info = process_image(buffer)

        cur = postgres.cursor()
        cur.execute('UPDATE my_table SET ' + column_name + ' = %s WHERE uuid = %s', (new_column_info, row.uuid))
        postgres.commit()

    message.delete()

Deploying on a fleet of servers

We used Docker and Kubernetes to deploy ~20 app servers:

apiVersion: apps/v1
kind: Deployment
metadata:
  name: processing-code
spec:
  replicas: 2
  selector:
    matchLabels:
      app.kubernetes.io/name: processing-code
  template:
    metadata:
      labels:
        app.kubernetes.io/name: processing-code
    spec:
      containers:
        - name: processing-code-2
          image: ghcr.io/my_organization/my_repo/worker
          resources:
            limits:
              cpu: 16
              memory: 86Gi
              nvidia.com/gpu: 1

Storing the secrets with Kubernetes

kubectl create secret generic my-secrets --from-literal=secrets=$(cat keys.json | base64)

Conclusion

Overall, the project was a great success. We hit all of the goals we set and were able to process 100 terabytes of data (100 million images) in 12 hours. It was interesting to build out all this infrastructure and very rewarding once we had it all working.

Stay tuned for part 2, where we discuss how to monitor this entire infrastructure deployment with Datadog and Grafana, and Part 3 on integrating NVIDIA GPUs.

We've got decades of experience working with venture backed startups and specialize in Web Development, AWS Infrastructure and Machine Learning. If you need a hand with managing a massive data set, contact us!