Memory leak when iterating a Dataset
Opened by NouamaneTazi (closed)

e## Describe the bug
It seems that memory never gets freed after iterating a `Dataset` (using `.map()` or a simple `for` loop)

## Steps to reproduce the bug
```python
import gc
import logging
import time
import pyarrow
from datasets import load_dataset
from tqdm import trange
import os, psutil

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
process = psutil.Process(os.getpid())

print(process.memory_info().rss)  # output: 633507840 bytes

corpus = load_dataset("BeIR/msmarco", 'corpus', keep_in_memory=False, streaming=False)['corpus'] # or "BeIR/trec-covid" for a smaller dataset

print(process.memory_info().rss)  # output: 698601472 bytes

logger.info("Applying method to all examples in all splits")
for i in trange(0, len(corpus), 1000):
    batch = corpus[i:i+1000]
    data = pyarrow.total_allocated_bytes()
    if data > 0:
        logger.info(f"{i}/{len(corpus)}: {data}")

print(process.memory_info().rss)  # output: 3788247040 bytes

del batch
gc.collect()

print(process.memory_info().rss)  # output: 3788247040 bytes

logger.info("Done...")
time.sleep(100)
```

## Expected results
Limited memory usage, and memory to be freed after processing

## Actual results
Memory leak
![test](https://user-images.githubusercontent.com/29777165/174578276-f2c37e6c-b5d8-4985-b4d8-8413eb2b3241.png)
You can see how the memory allocation keeps increasing until it reaches a steady state when we hit the `time.sleep(100)`, which showcases that even the garbage collector couldn't free the allocated memory

## Environment info
<!-- You can run the command `datasets-cli env` and copy-and-paste its output below. -->
- `datasets` version: 2.3.2
- Platform: Linux-5.4.0-90-generic-x86_64-with-glibc2.31
- Python version: 3.9.7
- PyArrow version: 8.0.0
- Pandas version: 1.4.2


Comments:
Is someone assigned to this issue?

The same issue is being debugged here: https://github.com/huggingface/datasets/issues/4883


Here is a modified repro example that makes it easier to see the leak:

```
$ cat ds2.py
import gc, sys
import time
from datasets import load_dataset
import os, psutil

process = psutil.Process(os.getpid())

print(process.memory_info().rss/2**20)

corpus = load_dataset("BeIR/msmarco", 'corpus', keep_in_memory=False, streaming=False)['corpus']
corpus = corpus.select(range(200000))

print(process.memory_info().rss/2**20)

batch = None

mem_before_start = psutil.Process(os.getpid()).memory_info().rss / 2**20

step = 20000
for i in range(0, 10*step, step):
    mem_before = psutil.Process(os.getpid()).memory_info().rss / 2**20
    batch = corpus[i:i+step]
    import objgraph
    #objgraph.show_refs([batch])
    #objgraph.show_refs([corpus])
    #sys.exit()
    gc.collect()

    mem_after = psutil.Process(os.getpid()).memory_info().rss / 2**20
    print(f"{i:6d} {mem_after - mem_before:12.4f} {mem_after - mem_before_start:12.4f}")

```

Let's run:

```
$ python ds2.py
     0      36.5391      36.5391
 20000      10.4609      47.0000
 40000       5.9766      52.9766
 60000       7.8906      60.8672
 80000       6.0586      66.9258
100000       8.4453      75.3711
120000       6.7422      82.1133
140000       8.5664      90.6797
160000       5.7344      96.4141
180000       8.3398     104.7539
```

You can see the last column of total RSS memory keeps on growing in MBs. The mid column is by how much it was grown during a single iteration of the repro script (20000 items)

@NouamaneTazi, please check my analysis here https://github.com/huggingface/datasets/issues/4883#issuecomment-1242599722 so if you agree with my research this Issue can be closed as well.

I also made a suggestion at how to proceed to hunt for a real leak here https://github.com/huggingface/datasets/issues/4883#issuecomment-1242600626

you may find this one to be useful as well https://github.com/huggingface/datasets/issues/4883#issuecomment-1242597966

Amazing job! Thanks for taking time to debug this 🤗

For my side, I tried to do some more research as well, but to no avail. https://github.com/huggingface/datasets/issues/4883#issuecomment-1243415957