google / google/tensorstore

Memory question

Open
#223 16 comments 1 reaction 0 assignees View on GitHub
Dominant language
C++
Stars
1.5k
Forks
145
PR merge metrics
No merged PRs in 30d

Description

I am writing an in-memory numpy array to tensorstore using the zarr driver. When I write the array, the memory increases by about the size of the array. I am not sure why this is, as I would expect references to the numpy array to be made by tensorstore, and then it is just writting to disk. Even if copies are being made in memory by tensorstore, I am not sure why the memory is not released after writing has finished. Here is an example. I show that even after `main` has returned, the numpy array is still being held in memory. This doesn't happen if I don't write to tensorstore.

```
import os
import tempfile
from pathlib import Path

import numpy as np
import psutil
import tensorstore

process = psutil.Process(os.getpid())

def main():
x = np.random.randint(low=0, high=255, size=(30000, 30000, 3), dtype='uint8')

with tempfile.TemporaryDirectory() as tmpdir:
store = tensorstore.open(
spec={
'driver': 'zarr',
'kvstore': {
'driver': 'file',
'path': str(Path(tmpdir) / 'z.zarr')
}
},
shape=x.shape,
dtype=x.dtype,
chunk_layout=tensorstore.ChunkLayout(
chunk=tensorstore.ChunkLayout.Grid(shape=(4096, 4096, 3))
),
create=True,
).result()

print('before write: ', process.memory_info().rss)
store[:] = x
print('after write: ', process.memory_info().rss)

if __name__ == '__main__':
print('before main: ', process.memory_info().rss)
main()
print('after main: ', process.memory_info().rss)
```

output:

```
before main: 86323200
before write: 2787667968
after write: 5585666048
after main: 2885693440
```

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.