Code Clinic case study

What happens in an ICCS Climate Code Clinic? Principal Research Software Engineer Tom Meltzer gives an example.

Background

I was approached by a researcher. They had a ML training pipeline which rapidly consumed memory before crashing their system. To summarise, we used a programme called memray to profile their code and reduce memory consumption by a factor of 7.

Approach 

The researcher gave me a brief description of the code, along with their suggestions for where they thought it might be going wrong. I like a guesstimate as much as the next person, but computers are often tricky and misleading... So I suggested we dive straight into profiling.

I normally work on compiled languages like C, C++ and Fortran. The researcher's code was written in Python, but, similar concepts apply. So after a quick Google search I came across Bloomberg’s memray project.

We had two requirements:

1. We needed a tool that could track memory allocations over time; and
2. We needed to know which lines of code corresponded to these allocations.

Thankfully, memray provided everything we needed.

Method

We had a relatively simple setup. A single file compiler.py which we need to profile. With memray this can be done as follows:

$ memray run compiler.py

After running, it produces a binary file called something like

$ ls memray*.bin
memray-compiler.py.8689.bin

Initially we wanted to see some overall behaviour of the program. This is possible using memray’s flamegraph tool.

$ memray flamegraph memray-compiler.py.8689.bin

This will then create a .html file which we can inspect using a browser.

$ firefox memray-compiler.py.8689.html

From this we can view a flamegraph, but also a memory timeline by clicking “Memory Graph”. Below is a figure that shows the memory usage over time.

So we have confirmed there is indeed a problem, the memory keeps on increasing. The issue is where does this occur. For this, it is easiest to view using memray’s tree utility.

$ memray tree memray-compiler.py.8689.bin

This should look something like the following:

From the tree we can see that there is roughly 8GB being consumed on a single line compiler.py:187. The specific line in question looked something like the following:

pressure_np = pressure_array.compute().astype("float32").values

The compiler.py code uses a library called dask so that it can process datasets that are too large to fit in memory. dask represents computations as task graphs and computes values only when necessary. This helps to batch the dataset and reduce the amount of memory required to process the full dataset. The issue is that the .compute() method forces dask to compute the full dataset too early, which forces the full dataset into RAM causing a memory crash.

Interestingly, this is not the area of code where we expected the memory error. Thankfully, we made a few changes later on in the code which enable us to use dask arrays throughout, which removed the need for the premature .compute() method.

Profiling the code a second time, we now get:

The code takes slightly longer to run, but the memory pressure is reduced significantly. We now spike around 2GB instead of 14GB like before. This means we can now run significantly larger datasets, and prepare our training data locally without needing big HPC.

This small example shows the benefit of using a profiler to optimise code, and how an ICCS Code Clinic can work.