Memory leak when iterating a Dataset
NONE
Is someone assigned to this issue?
2022-08-30T04:32:42Z
https://github.com/huggingface/datasets/issues/4528#issuecomment-1231128248
https://api.github.com/repos/huggingface/datasets/issues/4528
IC_kwDODunzps5JYYa4
https://api.github.com/repos/huggingface/datasets/issues/comments/1231128248/reactions
2022-08-30T04:32:42Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1231128248
https://avatars.githubusercontent.com/u/4252994?v=4
https://api.github.com/users/OhadRubin/events{/privacy}
https://api.github.com/users/OhadRubin/followers
https://api.github.com/users/OhadRubin/following{/other_user}
https://api.github.com/users/OhadRubin/gists{/gist_id}
https://github.com/OhadRubin
OhadRubin
MDQ6VXNlcjQyNTI5OTQ=
https://api.github.com/users/OhadRubin/orgs
https://api.github.com/users/OhadRubin/received_events
https://api.github.com/users/OhadRubin/repos
https://api.github.com/users/OhadRubin/starred{/owner}{/repo}
https://api.github.com/users/OhadRubin/subscriptions
User
https://api.github.com/users/OhadRubin
public
CONTRIBUTOR
The same issue is being debugged here: https://github.com/huggingface/datasets/issues/4883

2022-09-09T04:35:12Z
https://github.com/huggingface/datasets/issues/4528#issuecomment-1241488981
https://api.github.com/repos/huggingface/datasets/issues/4528
IC_kwDODunzps5J_55V
https://api.github.com/repos/huggingface/datasets/issues/comments/1241488981/reactions
2022-09-09T04:35:12Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1241488981
https://avatars.githubusercontent.com/u/10676103?v=4
https://api.github.com/users/stas00/events{/privacy}
https://api.github.com/users/stas00/followers
https://api.github.com/users/stas00/following{/other_user}
https://api.github.com/users/stas00/gists{/gist_id}
https://github.com/stas00
stas00
MDQ6VXNlcjEwNjc2MTAz
https://api.github.com/users/stas00/orgs
https://api.github.com/users/stas00/received_events
https://api.github.com/users/stas00/repos
https://api.github.com/users/stas00/starred{/owner}{/repo}
https://api.github.com/users/stas00/subscriptions
User
https://api.github.com/users/stas00
public
CONTRIBUTOR
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)
2022-09-09T04:56:36Z
https://github.com/huggingface/datasets/issues/4528#issuecomment-1241499052
https://api.github.com/repos/huggingface/datasets/issues/4528
IC_kwDODunzps5J_8Ws
https://api.github.com/repos/huggingface/datasets/issues/comments/1241499052/reactions
2022-09-09T05:00:15Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1241499052
https://avatars.githubusercontent.com/u/10676103?v=4
https://api.github.com/users/stas00/events{/privacy}
https://api.github.com/users/stas00/followers
https://api.github.com/users/stas00/following{/other_user}
https://api.github.com/users/stas00/gists{/gist_id}
https://github.com/stas00
stas00
MDQ6VXNlcjEwNjc2MTAz
https://api.github.com/users/stas00/orgs
https://api.github.com/users/stas00/received_events
https://api.github.com/users/stas00/repos
https://api.github.com/users/stas00/starred{/owner}{/repo}
https://api.github.com/users/stas00/subscriptions
User
https://api.github.com/users/stas00
public
CONTRIBUTOR
@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
2022-09-10T02:59:48Z
https://github.com/huggingface/datasets/issues/4528#issuecomment-1242603126
https://api.github.com/repos/huggingface/datasets/issues/4528
IC_kwDODunzps5KEJ52
https://api.github.com/repos/huggingface/datasets/issues/comments/1242603126/reactions
2022-09-10T02:59:48Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1242603126
https://avatars.githubusercontent.com/u/10676103?v=4
https://api.github.com/users/stas00/events{/privacy}
https://api.github.com/users/stas00/followers
https://api.github.com/users/stas00/following{/other_user}
https://api.github.com/users/stas00/gists{/gist_id}
https://github.com/stas00
stas00
MDQ6VXNlcjEwNjc2MTAz
https://api.github.com/users/stas00/orgs
https://api.github.com/users/stas00/received_events
https://api.github.com/users/stas00/repos
https://api.github.com/users/stas00/starred{/owner}{/repo}
https://api.github.com/users/stas00/subscriptions
User
https://api.github.com/users/stas00
public
MEMBER
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
2022-09-12T08:51:12Z
https://github.com/huggingface/datasets/issues/4528#issuecomment-1243416999
https://api.github.com/repos/huggingface/datasets/issues/4528
IC_kwDODunzps5KHQmn
https://api.github.com/repos/huggingface/datasets/issues/comments/1243416999/reactions
2022-09-12T08:51:12Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1243416999
https://avatars.githubusercontent.com/u/29777165?v=4
https://api.github.com/users/NouamaneTazi/events{/privacy}
https://api.github.com/users/NouamaneTazi/followers
https://api.github.com/users/NouamaneTazi/following{/other_user}
https://api.github.com/users/NouamaneTazi/gists{/gist_id}
https://github.com/NouamaneTazi
NouamaneTazi
MDQ6VXNlcjI5Nzc3MTY1
https://api.github.com/users/NouamaneTazi/orgs
https://api.github.com/users/NouamaneTazi/received_events
https://api.github.com/users/NouamaneTazi/repos
https://api.github.com/users/NouamaneTazi/starred{/owner}{/repo}
https://api.github.com/users/NouamaneTazi/subscriptions
User
https://api.github.com/users/NouamaneTazi
public
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
