Possible Memory leak in the `convert` module

[MikeLippincott]

I noticed that while converting multiple large ~30GiB .sqlite files to .parquet files the memory in use never dropped from the spike of conveting.

A code example of use:

for file in list_of_sqlite_files:
    convert(
        source_path=source_path,
        dest_path=dest_path,
        dest_datatype=dest_datatype,
        preset=preset,
    )

After each iteration the memory is use was ~2x what the file convert size was and never reset.
So for 30GiB files after the first iteration the memory spike was 60GiB and on iteration 2 I maxed out my 128GiB of memory for the second 30GiB file.

I am not sure if this as python thing on my end or a memory leak on CytoTable end. To test this I have an example with no loop to test on below.

I have a subset 23.8MiB .sqlite file I can send for a reproducible issue with profiled code:

# Notebook to call CytoTable
# import
import pathlib
from cytotable import convert
from memory_profiler import profile
from time import sleep

# set args
sqlite_file = 'PBMC_batch_1_subset'
source_path = str(pathlib.Path('PBMC_batch_1_subset.sqlite'))
dest_path = str(pathlib.Path('PBMC_batch_1_subset.parquet'))
dest_datatype = "parquet"
preset = "cellprofiler_sqlite_pycytominer"

print(f"Performing merge single cells and conversion on {sqlite_file}!")

@profile 
def my_func(source_path,dest_path,dest_datatype,preset):
    convert(
        source_path=source_path,
        dest_path=dest_path,
        dest_datatype=dest_datatype,
        preset=preset,
    )

# call profiled function
my_func(source_path,dest_path,dest_datatype,preset)

# to show that memory spike persists while script is running post convert call
sleep(45)

Below is the memory-time graph from the profiling

[d33bs]

Great investigation on this and thank you for raising the issue! I wanted to follow up after our conversation just now to highlight possible technical causes. In addition to what we talked about I wanted to highlight a few things:

1. Python garbage collection mechanisms can sometimes cause delays in memory deallocation. The way this can appear sometimes is similar to what you highlighted where there are spikes and consistent lines afterwards. There are ways around this, but it's often more maintainable to let Python handle these things on their own. When there are multiple procedural needs (as in the case of a loop) the memory should be reallocated as necessary as references to objects are no longer needed (so long as the program and dependencies allow for this to happen).
2. If it remains the same after each loop iteration, the use of dest_path in your loop might be causing CytoTable to attempt to perform concatenation or join work on each previous result dataset, potentially becoming a cause for the increased memory usage you mentioned. CytoTable is designed to reference multiple source datasets when providing a directory as a source. It might be that you could see performance improvements by referencing the source data as a directory with a single call so long as the source data are similar in their schema (for example, all generated from CellProfiler, etc).

Other technical aspects internal to CytoTable which could be cause for the memory issues:

1. Parsl, which is used as a workflow engine within CytoTable, may not be fully deallocating memory or object references on completions.
2. PyArrow, which is used to facilitate Parquet concatenation, may somehow be experiencing issues with memory bloat as the files are written.
3. DuckDB, which is used to perform data extraction and join work, might for some reason be holding onto memory after completing this work. Calls here are generally designed to directly export to Parquet, so this seems unlikely, but there could be something happening.

I plan to take a look into this to find out more and will circle back with you after I know more. Please don't hesitate to let me know if you have any questions.

[d33bs]

Following up here in addition to our discussion in person to make sure this issue receives similar information for visibility. I'll go through things below in a way that I hope builds a foundation of understanding and the basis for further work.

Python resident memory with time.sleep()

When using memory_profiler with default arguments I believe the chart you shared demonstrates resident memory usage. Python's memory manager generally won't aggressively deallocate resident memory already used by a Python process unless it's told to explicitly do so (this process is automatically performed). Even when we do tell Python's garbage collection to run, we sometimes see mixed results. I bring this up to help illustrate that despite our efforts within CytoTable, Python may not be designed to yield resident memory without more nuanced changes. See below's examples and charts (using Python 3.9 on MacOS with memray).

Without garbage collection

"""
Example demonstrating what happens when we use sleep after
allocating some memory in Python.
"""

import random
import time

# create randomized number data
example = [[random.random() for _ in range(100)] for _ in range(1000)]

# sleep for memory visibility
time.sleep(10)

---

With garbage collection

"""
Example demonstrating what happens when we use sleep after
allocating some memory in Python.
"""

import gc
import time

import random

# create randomized number data
example = [[random.random() for _ in range(100)] for _ in range(1000)]

# sleep for memory visibility
time.sleep(10)

# delete reference to example data
del example

# invoke garbage collection
gc.collect()

# sleep for memory visibility
time.sleep(10)

With the above two examples we can observe that even when we invoke garbage collection after deleting a reference to data within Python that we only receive a part of it back.

CytoTable resident and heap memory usage

While exploring the nature of this issue I tested CytoTable with memray to gain insights. In so doing I was able to reproduce what you illustrated in created this issue and also found that specifically heap memory seemed misaligned.

Heap and resident memory

My understanding is that Python uses heap memory for object data storage and resident memory for executable code, loaded libraries, variable references, and the execution environment itself. Unnecessary growth or spikes in either of these memory types may be signs of memory leakage (memory which is not deallocated). We're also contending with automated garbage collection so this can sometimes "hide" memory challenges.

CytoTable heap and resident memory in a loop

The above chart demonstrates the resident and heap memory usage when reading a ~4MB SQLite dataset. We can see large spikes in the heap usage, with the second being larger than the other. This heap usage is well out of alignment with the overall size of the dataset. These spikes I thought might be interrelated to the resident memory usage increases somehow, so I opted to explore these more deeply below. Generally the timing of these seem to align with when DuckDB reads data from parquet files and CytoTable uses arrow data structures to reach the results.

DuckDB, Closed Connections, and Heap Usage

Digging into this further using the previously referenced list of possibilities I happened to find a DuckDB Discord conversation which sounded like (though didn't exactly match) what was occurring here. I referenced the mentions of explicit connect() and close() usage as a way to further test and see what was occurring.

DuckDB parquet read without explicit close()

"""
An example to test duckdb execution without closing connections
using parquet files.
"""

import duckdb

def _duckdb_reader() -> duckdb.DuckDBPyConnection:
    """
    Creates a DuckDB connection with the
    sqlite_scanner installed and loaded.

    Returns:
        duckdb.DuckDBPyConnection
    """

    return duckdb.connect()


for parquet_file in [
    "./examples/data/random_number_data.parquet",
    "./examples/data/random_number_data-copy.parquet",
]:
    sql_stmt = f"""
    SELECT * from '{parquet_file}';
    """
    _duckdb_reader().execute(sql_stmt)

---

DuckDB parquet read with explicit close() (and arrow objects)

"""
An example to test duckdb execution with closed connections
using parquet files.
"""

import duckdb

def _duckdb_reader(sql_stmt: str) -> duckdb.DuckDBPyConnection:
    """
    Creates a DuckDB connection with the
    sqlite_scanner installed and loaded.

    Returns:
        duckdb.DuckDBPyConnection
    """

    # setup a duckdb client with various configuration options
    duckdb_client = duckdb.connect()

    # gather the result from the duckdb connection
    result = duckdb_client.execute(sql_stmt).arrow()

    # explicitly close the connection before returning result
    duckdb_client.close()

    return result


for parquet_file in [
    "./examples/data/random_number_data.parquet",
    "./examples/data/random_number_data-copy.parquet",
]:
    sql_stmt = f"""
    SELECT * from '{parquet_file}';
    """
    _duckdb_reader(sql_stmt)

The above two examples reference a very small random number dataset using DuckDB both with and without connection closes. In this case, the connection close appears to demonstrate a similar pattern to what occurs in CytoTable. Doing a bit more investigation I came across this issue duckdb/duckdb#7856 which appears to indicate that Arrow may be the cause. In the example of explicitly closing the connection we create an arrow data structure which might be closer to the root cause.

Next steps: I plan to take a further look into PyArrow and how Arrow tables allocate + deallocate memory. I will follow up once I know more here.

[d33bs]

Posting another follow up here based on recent findings:

PyArrow Memory Allocation

PyArrow provides several ways to allocate memory using malloc, mimalloc, and jemalloc. My understanding is that each of these may be preferred based on the system where PyArrow is used. Many of these functions and objects may be controlled through PyArrow memory and buffer methods.

There are some existing PyArrow issues which might be related to or could help with this that can be found below:

• [Python][Docs] Provide guidelines for expectations about and how to investigate memory usage apache/arrow#36378
• Memory leak using pyarrow.Table and to_pandas method apache/arrow#36101

I tested each of these memory allocators using PyArrow v12.0.1 on my MacOS device to see how they might differ when reading two identical parquet files within a loop. This work can be found in a fork which I hope to eventually merge as part of resolving this issue: https://github.com/d33bs/CytoTable-benchmarks/blob/explore-loop-memory-usage/notebooks/pyarrow_looped_memory_usage.ipynb.

malloc performance

mimalloc performance

jemalloc performance

In terms of heap memory, the above showed malloc performing the best, jemalloc shortly behind, and mimalloc with the worst performance (around 5x more heap memory consumed than the other two allocators). My very brief testing also showed that mimalloc is used as a default (at least on my device), so this might be linked to your findings as well. A quick way to determine which memory allocator is being used is to run the following code snippet:

import pyarrow as pa
print(pa.default_memory_pool().backend_name)

Based on the above, I felt it might indicate we should use malloc as the primary memory allocator for PyArrow. This likely comes with many caveats; making a memory allocator selection to override the PyArrow defaults could mean we make this package unusable. Just the same, I tested performance and unfortunately, I found there were still large spikes of heap memory consumption despite selecting another memory allocator.

PyArrow parquet.read_table Memory Mapping

PyArrow offers the ability to use memory mapping when reading parquet files through pyarrow.parquet.read_table(..., memory_map=True). For equivalency with the above, I tested to see how each of the memory allocators performed using memory mapped techniques.

malloc performance with memory mapping

mimalloc performance with memory mapping

jemalloc performance with memory mapping

The above showed improved performance for all allocators with malloc and jemalloc able to remain under resident memory consumption. In general, this parameter seems like a good candidate to add to CytoTable in order to marginally increase read performance (regardless of allocator selection). However, when testing this, as before, the large spike in heap memory use remained.

Next steps

Given the above, I think we should implement memory mapping for parquet file reads within CytoTable, but I don't consider this to be a fix for the large heap memory consumption.

My next target for exploration with this issue will be _join_source_chunk to see if we can avoid in-memory PyArrow table duplication when modifying columns included and sorting (PyArrow tables are read-only, and as a result, they may be duplicating with each modification after DuckDB extraction before they can be deallocated). Please don't hesitate to let me know if you have any questions or concerns in the meantime.

[d33bs]

Quick update: as a result of work in this issue I've opened apache/arrow#37361 for further details on Arrow memory allocator selection in case it may assist here (or elsewhere).

[d33bs]

Another update:

Based on the feedback and some further consideration on the Arrow memory allocators, my feeling is that we should likely stick with the defaults selected. Generally I've noticed with CytoTable that jemalloc tends to be slightly more aggressive in deallocating memory than mimalloc, with both being better about dellocations than malloc/system.

Lately I've been working on using gc.set_debug(gc.DEBUG_LEAK) to find and estimate references, size, and contents of leaking objects. The gist of this can be see below:

import gc
import sys

import cytotable
import pandas as pd

gc.set_debug(gc.DEBUG_LEAK)

cytotable.convert(
    source_path="./examples/data/all_cellprofiler.sqlite",
    dest_path="./examples/data/test-result.parquet",
    dest_datatype="parquet",
    preset="cellprofiler_sqlite_pycytominer",
)

collected = gc.collect()
if gc.garbage:
    print(f"Memory leak detected: {len(gc.garbage)} objects")
    df = pd.DataFrame(
        [
            {
                "id": id(obj),
                "type": type(obj),
                "refcount": sys.getrefcount(obj),
                "repr": repr(obj),
                "size": sys.getsizeof(obj),
            }
            for obj in gc.garbage
        ]
    )

df.head()

I noticed many of the leaking objects seemingly come from dictionaries of parameters sent to Parsl python_app's (which act as "tasks" in the context of larger join_app "workflows". I attempted to nest some of the larger parameterized data work (columns, schema operations) inside the Parsl apps where they occur to see how this effected things. I also increased the number of times CytoTable runs with testing to 8 in order to see if there was a noticeable difference.

It does seem that the memory allocator begins to curb the use of memory around the 6th iteration, but I don't feel that we've avoided the growth entirely. Generally the changes are fairly complex and don't seem to indicate a significant amount of performance, so I'd feel hesitant to pull these in. Note: there were times where I'd witness the memory allocator dropping to a maximum, seemingly appearing to operate without leaky behavior. I couldn't reproduce this consistently, so it may be a flaw in how I'm running tests or perhaps something to do with "cold starts".

Below are the results of 8 looped runs of CytoTable using mimalloc:

As-is:

With nested functionality to avoid large parameter objects:

Based on my current understanding, it might be that the ThreadPoolExecutor (the default executor for Parsl used by CytoTable now) functions this way by default and/or is optimized for different kinds of work than what we seek in this package. I feel that we might be able to move to the HighThroughputExecutor with positive effects.

In testing the same looped runs as above using the HighThroughputExecutor, I saw the following. These results were faster and also better performing than those found using ThreadPoolExecutor.