`xarray` data objects not being accounted in the worker's managed memory
- Dominant language
- Python
- Stars
- 1.7k
- Forks
- 778
- Avg merge
- 2h 50m
- Merged PRs (30d)
- 3
Description
**What happened**: When we use `xarray` data objects in the dask computations, they are not being accounted in the worker's managed memory. We see only the size of reference tables of `xarray` is being shown in the managed memory.
**What you expected to happen**: As `xarray` is one of the standard library, we would like to see these data objects being accounted in the managed memory. As workers can only spill the managed memory to disk, we faced number of issues when running out of memory in our workflows.
**Minimal Complete Verifiable Example**:
```
import time
import xarray as xr
from dask import delayed
from distributed import Client, wait
from dask.sizeof import sizeof
# Initialise sample data
def init_data(n):
return xr.Dataset( {"testdata_1": (["a", "b", "y", "x"],
np.random.rand(1, 1, n, n))} , {"testdata_2": (["a", "b", "y", "x"],
np.random.rand(1, 1, n, n))} )
# Get size of objects using dask sizeof
def get_size():
data_obj = init_data(1024)
obj_size = sizeof(data_obj)
print("xarray data object size is {} MiB".format(obj_size / (1024 * 1024)))
# Sample operations
def sample_operations(ds):
time.sleep(30)
ds = ds * 2 + 1.5
return ds
if __name__ == "__main__":
get_size()
print("Creating scheduler and 4 workers")
client = Client(n_workers=4, threads_per_worker=1)
data = delayed(init_data)(1024)
graph = delayed(sample_operations)(data)
future = client.compute(graph, sync=True)
wait(future)
```
**Anything else we need to know?**: The origin of this issue is in the dask's `sizeof()` implementation, `xarray` is not included. If we register `xarray` in the dispatcher function, we will fix this issue.
```
@sizeof.register_lazy("xarray")
def register_xarray():
import xarray
@sizeof.register(xarray.core.dataset.Dataset)
def sizeof_xarray_dataset(x):
return int(x.nbytes)
```
**Environment**:
- Dask version: 2021.9.0
- Python version: 3.8.10
- Operating System: Debian GNU/Linux 10
- Install method (conda, pip, source): pip
Contributor guide
Assessment
This issue has not been assessed yet.