Memory leak when streaming

[Jourdelune]

Describe the bug

I try to use a dataset with streaming=True, the issue I have is that the RAM usage becomes higher and higher until it is no longer sustainable.

I understand that huggingface store data in ram during the streaming, and more worker in dataloader there are, more a lot of shard will be stored in ram, but the issue I have is that the ram usage is not constant. So after each new shard loaded, the ram usage will be higher and higher.

Steps to reproduce the bug

You can run this code and see you ram usage, after each shard of 255 examples, your ram usage will be extended.

from datasets import load_dataset
from torch.utils.data import DataLoader

dataset = load_dataset("WaveGenAI/dataset", streaming=True)

dataloader = DataLoader(dataset["train"], num_workers=3)

for i, data in enumerate(dataloader):
    print(i, end="\r")

Expected behavior

The Ram usage should be always the same (just 3 shards loaded in the ram).

Environment info

• datasets version: 3.0.1
• Platform: Linux-6.10.5-arch1-1-x86_64-with-glibc2.40
• Python version: 3.12.4
• huggingface_hub version: 0.26.0
• PyArrow version: 17.0.0
• Pandas version: 2.2.3
• fsspec version: 2024.6.1

[enze5088]

I seem to have encountered the same problem when loading non streaming datasets. load_from_disk. Causing hundreds of GB of memory, but the dataset actually only has 50GB

[lhoestq]

FYI when streaming parquet data, only one row group per worker is loaded in memory at a time.

Btw for datasets of embeddings you can surely optimize your RAM by reading the data as torch tensors directly instead of the default python lists

from datasets import load_dataset
from torch.utils.data import DataLoader

dataset = load_dataset("WaveGenAI/dataset", streaming=True).with_format("torch")

dataloader = DataLoader(dataset["train"], num_workers=3)

[SuperSonnix71]

Im also, hitting this issue.....

 # This is what's causing the leak:
  batch_datasets = []
  for file_path in batch_files:
      dataset = load_dataset(..., streaming=True)
      shuffled_dataset = dataset.shuffle(seed=42, buffer_size=1000)  # 1000-item buffer
      batch_datasets.append(shuffled_dataset)  # Buffer persists

  interleaved_dataset = interleave_datasets(batch_datasets, seed=42)  

And, nothing helps

      del batch_datasets, interleaved_dataset
      gc.collect()  # This doesn't work for HuggingFace internal memory structures

so my guess is that they wrote this in RUST and forgot to clean up!!!

Now, if i remove the interleaving and process files sequentially... like this it still leaks

     # Process files one by one - no batching, no interleaving
     for file_idx, file_path in enumerate(file_paths):
         dataset = load_dataset("parquet", data_files=file_path, split="train", streaming=True)
         shuffled_dataset = dataset.shuffle(seed=42, buffer_size=1000) 
         
         for record in shuffled_dataset:
             # Process record immediately
             pass
         
         del dataset, shuffled_dataset
         gc.collect()

• File 1: 42.4% memory
• File 2: 42.5% memory
• File 3: 42.5% memory
• File 4: 48.4% memory (+6%)
• File 5: 52.7% memory (+4.3%)
• File 6: 56.7% memory (+4%)
• File 7: 59.6% memory (+2.9%)
• File 8: 62.0% memory (+2.4%)

I had to go back to sequential shuffling (NO Interleaving) and clean up like this

     dataset.cleanup_cache_files()                                                         
         del dataset, shuffled_dataset                                                         
         gc.collect()                                                                          
         pa.default_memory_pool().release_unused()                                             
         libc.malloc_trim(0)  # when available   

[bghira]

i have also observed these memory leaks inside the huggingface library when developing bghira/captionflow and had the same outcome of being unable to actually free anything when it occurs. i've worked around it by avoiding some of the more damaging parts of the library, but in doing so i've essentially restricted the compatibility levels of the project.

[lhoestq]

Could it be a leak from PyArrow which is used to stream the data from the Parquet files ?

[bghira]

i believe it's heavily involved yeah

[meetdoshi90]

Is there any update on this? I’m seeing the same issue when using multiple streaming datasets for long runs.

[bghira]

see https://github.com/bghira/webshart for a Rust-based implementation (relatively bare minimum for my needs, sorry - open to PRs) that helped me to workaround the problem.

[meetdoshi90]

Thank you ✌🏻

[lhoestq]

Btw for local datasets since a Dataset uses memory mapped Arrow files, iterating on the data does iterate on memory mapped files which loads pages in RAM progressively but discards them once they are not needed anymore. Generally, the OS does it automatically when memory is needed for something else. So you might see your RSS (memory) go up but it won't OOM since data are paged out automatically by the OS.

[meetdoshi90]

Hi @lhoestq
It does go OOM since most of the memory occupied does not free up.
Eg: For my case, the compressed data size is about 400 GB with over 100 datasets. Still, the occupied memory goes up to over a TB for longer runs. I think it's most likely related to the dataloader buffer; every time it is refreshed, it fails to free up the occupied memory by the previous buffer.
MAX MEM: 908 Gbytes; AVG MEM: 843 Gbytes;

[simonschoenhofen]

Hi @lhoestq ,

We also encounter the same issue. When investigating the memory leak, i noticed that we have RSS Growth and no Pyarrow memory issue.

I did some digging and i think i found the cause of it.

_generate_tables

Calls

with open(file, "rb") as f:
    parquet_fragment = parquet_file_format.make_fragment(f)

1. It creates a Python file handle (using fsspec/xopen for GCS/S3 or open for local).

2. It passes this Python object into the PyArrow C++ engine (parquet_file_format).

3. Every time PyArrow needs data (which happens constantly with small Row Groups, as per default in datasets), it must pause, call back into Python to execute f.read(), acquire the GIL, copy bytes from Python to C++, and release the GIL.

__
In theory this should be no issue but i assume what is happening (although i am not convinced yet if this is actually the underlying reason) is that:

With tiny Row Groups (e.g., size 10-100), this boundary is crossed thousands of times. The internal buffers created during these Python-to-C++ callbacks (specifically fsspec buffers and PyArrow's internal Python file wrappers) are accumulating faster than the Python Garbage Collector can identify and free them. This results in RSS growth even though the PyArrow memory pool remains empty.
__

However, we were able to fix it by monkey patching it and could confirm that this is the root cause of the issue. We are just not 100 % sure why 😓

__
Our internal current fix:
We confirmed that the leak disappears completely if we bypass the Python file handle and use native PyArrow FileSystems (C++) by passing the path string directly. However, simply switching to pq.ParquetFile(path).read_row_group(i) (our tested temp fix) caused a 5x performance regression because it lacks the async prefetching. But it solved the memory leak.

[bghira]

oh - that actually makes a lot of sense. when developing webshart in Rust i had to make some crucial decisions when handling native python inputs like file pointers. it stands to reason that the C++ layer of pyarrow has the same problem.

for my testing i found that we couldn't even link into glibc to bypass gc and clean things up, and the gc running manually did nothing to resolve the issue either. deleting handles on Python side doesn't help. you can completely remove all Python side holds and allocations and it will still be holding onto RSS in Arrow.

[bghira]

also the statement that "it won't OOM because they will be paged out" is inaccurate. even on a 128G unified memory macbook it will OOM when looping over a 5gb dataset.

[simonschoenhofen]

For reference, this was our test/patch code to confirm the memory issue and our quick fix (local and gcs streaming)
(as mentioned, this fix is way slower for gcs streaming at this stage):

import argparse
import time
import os
import psutil
import gc
import pyarrow as pa
import pyarrow.parquet as pq
import datasets
from torch.utils.data import DataLoader

# ==============================================================================
# 1. THE MONKEY PATCH
# ==============================================================================
def apply_monkey_patch():
    """
    Overwrites the internal Hugging Face Parquet loader.
    Fix: Uses pyarrow.dataset (FAST) but with C++ native filesystems (MEMORY SAFE).
    """
    try:
        from datasets.packaged_modules.parquet.parquet import Parquet
    except ImportError:
        import datasets.packaged_modules.parquet as parquet_module
        Parquet = parquet_module.Parquet

    import pyarrow.dataset as ds
    import pyarrow.fs as pa_fs
    import pyarrow as pa
    from datasets.utils.file_utils import xopen
    import itertools

    # 2. Initialize GCS C++ FileSystem
    try:
        gcs_fs = pa_fs.GcsFileSystem()
        print("[Fix] GCS Support Enabled (C++ Native)")
    except ImportError:
        gcs_fs = None
        print("[Fix] Warning: PyArrow GCS support missing.")

    # 3. Flattening Helper
    def flatten_files(files):
        if isinstance(files, str):
            yield files
            return
        try:
            for item in files:
                yield from flatten_files(item)
        except TypeError:
            yield str(files)
            
    # 4. Path Resolver (Handles FilesIterable wrappers)
    def resolve_path(file_obj):
        if hasattr(file_obj, "current"): return str(file_obj.current)
        if hasattr(file_obj, "path"): return str(file_obj.path)
        return str(file_obj)

    # 5. THE OPTIMIZED GENERATOR
    def optimized_generate_tables(self, files):
        # Setup the fast reader options
        filter_expr = None
        if self.config.filters:
             import pyarrow.parquet as pq
             filter_expr = pq.filters_to_expression(self.config.filters) if isinstance(self.config.filters, list) else self.config.filters

        parquet_file_format = ds.ParquetFileFormat(
            default_fragment_scan_options=self.config.fragment_scan_options
        )

        # Iterate over files
        file_iterator = flatten_files(files)
        
        for file_idx, file_obj in enumerate(file_iterator):
            file_path = resolve_path(file_obj)
            use_native = False
            filesystem = None
            clean_path = file_path

            if file_path.startswith("gs://") and gcs_fs:
                clean_path = file_path.replace("gs://", "", 1)
                filesystem = gcs_fs
                use_native = True
            elif not file_path.startswith("http") and "://" not in file_path:
                filesystem = None # Local FS is default
                use_native = True

            try:
                # --- CREATE FRAGMENT ---
                if use_native:
                    parquet_fragment = parquet_file_format.make_fragment(
                        clean_path, 
                        filesystem=filesystem
                    )
                    # We don't need a context manager here
                    ctx = None
                else:
                    # SLOW/LEAKY FALLBACK (For HTTP/S3)
                    f = xopen(file_path, "rb")
                    parquet_fragment = parquet_file_format.make_fragment(f)
                    ctx = f

                # --- READ BATCHES (FAST API) ---
                if parquet_fragment.row_groups:
                    # Use original efficient batching logic
                    batch_size = self.config.batch_size or parquet_fragment.row_groups[0].num_rows
                    
                    # .to_batches() is the high-performance async iterator
                    batch_iter = parquet_fragment.to_batches(
                        batch_size=batch_size,
                        columns=self.config.columns,
                        filter=filter_expr,
                        # Re-enable readahead for speed
                        batch_readahead=1, 
                        fragment_readahead=1
                    )

                    for batch_idx, record_batch in enumerate(batch_iter):
                        pa_table = pa.Table.from_batches([record_batch])
                        # Yield in the format HF expects
                        yield f"{file_idx}_{batch_idx}", self._cast_table(pa_table)
                        
            except Exception as e:
                print(f"Error reading {file_path}: {e}")
                raise e
            finally:
                if ctx:
                    ctx.close()

    # Apply the patch
    Parquet._generate_tables = optimized_generate_tables
    print("[Fix] Monkey patch applied: High-Performance C++ Mode")


# ==============================================================================
# 2. MONITORING & MAIN LOOP
# ==============================================================================
def print_memory(prefix=""):
    process = psutil.Process()
    mem = process.memory_info().rss / 1024 / 1024
    print(f"[{prefix}] RSS Memory: {mem:.1f} MB")

def run_test(args):
    # 1. Setup Data File
    files = args.data_files


    # 2. Apply Fix if requested
    if args.fix:
        apply_monkey_patch()
    else:
        print("[Baseline] Running STANDARD Hugging Face loading (Expect Leak)")

    # 3. Load Dataset
    print_memory("Start")
    
    # We deliberately use the standard HF API to prove the patch works transparently
    ds = datasets.load_dataset(
        "parquet", 
        data_files=files, 
        streaming=True,
        split="train"
    )

    # 4. Iterate
    dataloader = DataLoader(ds, batch_size=32)
    
    print("Starting iteration...")
    start_time = time.time()
    
    for i, batch in enumerate(dataloader):
        if i % 5 == 0:
            print_memory(f"Batch {i:4d}")
            
            # Simulate some "work"
            time.sleep(0.01) 
            
        if i >= args.max_batches:
            break

    print_memory("Final")
    print(f"Done in {time.time() - start_time:.2f}s")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--fix", action="store_true", help="Apply the memory leak fix")
    parser.add_argument("--data-files", type=str, default=None, help="Path to GCS or Local parquet file")
    parser.add_argument("--max-batches", type=int, default=500, help="Stop after N batches")
    args = parser.parse_args()

    run_test(args)