Designing a Simple Class that caches files from AWS

Question

Background

I have a pipeline that I run and it spins up multiple EC2 instances that process many jobs. I noticed that in many parts of my pipeline i am downloading files from amazon s3 multiple times. I can save lot of time if i just cache them since after a certain stage in the pipeline those files are not updated but just downloaded multiple times.

In that light I came up with a simple design that caches files from amazon s3 and also ensures the cached files are valid.

The idea is simple i maintain a dictionary that maps the file name with its respective etag. Etag is provided amazon s3 it is basically like md5 checksum for a file.

My Code

import os
import s3
import pickle


class Cache:
    def _init_(self):
        self.s3_files_etag_map = {}

    def file_exists(self, local_path):
        return os.path.exists(local_path)

    def is_file_valid(self, local_path, file_name, s3_path):
        with open(f"{local_path}/s3_files_etag_map.p", "br") as f:
            self.s3_files_etag_map = pickle.load(f)
        local_etag_value = self.s3_files_etag_map[file_name]
        s3_etag_value = self.get_s3_file_etag_value(s3_path, file_name)
        return s3_etag_value == local_etag_value

    def download_and_cache_file(self, s3_path, local_path, file_name):
        s3.copy_with_python_retry(
            from_path=f"{s3_path}/{file_name}", to_path=f"{local_path}/{file_name}"
        )
        s3_etag_value = self.get_s3_file_etag_value(s3_path, file_name)
        self.s3_files_etag_map[file_name] = s3_etag_value
        with open(f"{local_path}/s3_files_etag_map.p", "bw") as f:
            pickle.dump(self.s3_files_etag_map, f)

    def get_s3_file_etag_value(self, s3_path, file_name):
        s3_bucket_and_key = s3.deconstruct_s3_path(f"{s3_path}/{file_name}")
        s3_etag_value = s3.get_object_etag(s3_bucket_and_key[0], s3_bucket_and_key[1])
        return s3_etag_value

    def cache_file(self, s3_path, local_path, file_name):
        if self.file_exists(local_path):
            if self.is_file_valid(local_path, file_name, s3_path):
                with open(f"{local_path}/s3_files_etag_map.p", "bw") as f:
                    pickle.dump(self.s3_files_etag_map, f)
            else:
                self.download_and_cache_file(s3_path, local_path, file_name)
        else:
            os.makedirs(local_path)
            self.download_and_cache_file(s3_path, local_path, file_name)

My Thought Process

Here i will attempt to explain my thought process and how i came up with the code above. The idea for me is simple. I need to create a caching system that will last the lifetime of my EC2 instance which could be a day or two. So In my case when my pipeline downloads some file from s3 it will cache the file instead.

cache_device = Cache()
cache_device.cache_file(s3_path, local_path, file_name)

The users of the caching service will have knowledge of what the s3_path to the file is, the name of the file and the local_path where they want to cache the file.

Inside the Cache class the idea is simple. We first check if the path exists. If it does then we check if the etag value of the local file matches the etag of it's s3 counter part. This will let us know if our cached file is valid or not. If the file is valid then we dump the dictionary as a pickle file. (Basically convert python dictionary to a file format that is saved for easy access by other objects)

If the path does not exist then we create the path and download the file. When we download the file we also store the file name and its etag value in a dictionary and dump it as a pickle file.

Assumptions

You can assume I have a s3.py file that provides the service needed in the code

Caller Code

To keep things simple. Let just say for different ids i am running a script multiple times. So below i am looping over a 1000 times and each time i am calling the cache_file function where i can check if the file is cached and if it is i just get it from the local path and do the needful. if it is not cached then it will get cached via the code in caching class.

local_path = some_local_path
s3_path = some_s3_path
file_name = some_file_name
for i in range(1,1000):
    c = Cache()
    c.cache_file(local_path, s3_path, file_name)
    #read file from local path and do stuff with it.

Objective

I would love a code review on this as I am not good with python and I am open to any suggestions that can make code or overall design or anything else cleaner.

Answer 1

Constructor

def _init_(self):

is not a constructor; this needs to be __init__.

Pickle?

For a simple flat dictionary with a string key and a string value, pickle is an odd choice. It will be more difficult for you to debug, and unless you're storing millions of files I doubt that the performance difference to JSON will be visible. You should probably just use JSON.

Pathlib

Instead of

with open(f"{local_path}/s3_files_etag_map.p", "br") as f:

consider

local_path = Path(local_path_str)
with (local_path / 's3_files_etag_map.p').open('rb') as f:

Concurrency

Apparently you do not need multiple processes to all access this cache index file at the same time. As such, there is no need to read it on every single request. Read it once, keep the dict in memory, and write it out whenever it changes.

In other words: move your file-reading code to your __init__; keep the dictionary as a member variable on the class; and any time that you previously read from the file, just use the dictionary that's already been loaded. Any time that you modify the dictionary, make sure to modify the member variable, and write to the file.

Unpacking

s3_bucket_and_key = s3.deconstruct_s3_path(f"{s3_path}/{file_name}")
s3_etag_value = s3.get_object_etag(s3_bucket_and_key[0], s3_bucket_and_key[1])

should be

bucket, key = s3.deconstruct_s3_path(f"{s3_path}/{file_name}")
s3_etag_value = s3.get_object_etag(bucket, key)

Flattening logic

    if self.file_exists(local_path):
        if self.is_file_valid(local_path, file_name, s3_path):
            with open(f"{local_path}/s3_files_etag_map.p", "bw") as f:
                pickle.dump(self.s3_files_etag_map, f)
        else:
            self.download_and_cache_file(s3_path, local_path, file_name)
    else:
        os.makedirs(local_path)
        self.download_and_cache_file(s3_path, local_path, file_name)

can be

if not self.file_exists(local_path):
    os.makedirs(local_path)
    self.download_and_cache_file(s3_path, local_path, file_name)
elif self.is_file_valid(local_path, file_name, s3_path):
    with open(f"{local_path}/s3_files_etag_map.p", "bw") as f:
        pickle.dump(self.s3_files_etag_map, f)
else:
    self.download_and_cache_file(s3_path, local_path, file_name)

Testing

Any kind of meaningful unit testing on this thing will require mocking. This is a fairly detailed topic; read the official documentation here:

https://docs.python.org/3/library/unittest.mock.html

You need to mock out most of s3 to prevent it from connecting to the internet.