Memory leak when looping through data variables of a dataset loaded from a VRT #774

Activity

1. amaissen

   added

   bugSomething isn't workingSomething isn't working

   on May 3, 2024
2. snowman2

   added

   questionFurther information is requestedFurther information is requested

   and removed

   bugSomething isn't workingSomething isn't working

   on May 3, 2024
3. snowman2 commented on May 3, 2024

   

   snowman2

   on May 3, 2024

   Member

   More actions

   Add these kwargs to open_rasterio to disable caching:

   lock=False,  # disable internal caching
   cache=False,  # don't keep data loaded in memory. pull from disk every time
   
4. amaissen commented on May 3, 2024

   

   amaissen

   on May 3, 2024

   Last edited by amaissen

   Author

   More actions

   @snowman2 , thanks for pointing to these options. I tired the options you suggested but they did not help to release the memory.

   However, when I store the entire raster to zarr storage with to_zarr(), and load with raster = xarray.open_zarr(...), I don't see any memory leaks when iterating through the data variables. This would look like

   import rioxarray as rxr
   import xarray as xr
   import gc
   
   PATH = "path_to_multi_band_vrt.vrt"
   
   def no_memory_leak():
     # Read from VRT and save to zarr (one chunk per band)
     rxr.open_rasterio(PATH, band_as_variable=True, chunks={"x": -1, "y": -1}).to_zarr(some_temp_dataset)
     
     # Open zarr and iterate over data vars.
     raster = xr.open_zarr(some_temp_dataset, chunks={"x": -1, "y": -1})
     bands = list(raster.data_vars)
     
     for band in bands:
       data = raster[band].copy(deep=True).load()
     
       del data
       gc.collect()
5. GregoryPetrochenkov-NOAA commented on Jun 11, 2024

   

   GregoryPetrochenkov-NOAA

   on Jun 11, 2024

   Last edited by GregoryPetrochenkov-NOAA

   More actions

   I have experienced similar issues with memory leaks. I did an experiment with rioxarray loading GeoTIFFs and similarly with xarray loading NetCDF files.

   The tests use:

   • python: 3.10.14
   • memray: 1.12.0
   • rasterio: 1.3.9
   • xarray: 2024.3.0
   • rioxarray: 0.15.3
   • netCDF4: 1.6.5

   I would run the operation 5 times and memory profile with memray as follows (run in a Jupyter Notebook):

   import time
   import subprocess
   import os
   import gc
   from functools import partial
   
   import memray
   import rioxarray as rxr
   import xarray as xr
   
   %load_ext memray
   
   def run_test(func):
       """
       Driver to run memory accumulation test by running 5 times simulating a batch process
       """
   
       for x in range(5):
           func()
   
       time.sleep(1)

   I also set up these arguments in advance:

   xarray_kwargs = {
      "cand_file": "./subsample_benchmark_mean.nc",
      "bench_file": "./subsample_candidate_mean.nc",
       "cache": False,
       "lock": False
   }
   
   rio_kwargs = {
       "cand_file": "./c_uint8.tif",
       "bench_file": "./b_uint8.tif",
       "cache": False,
       "lock": False
   }

   The first snippet shows rioxarray loading the two GeoTIFFs in a context wrapper which automatically calls the close method when finished:

   %%memray_flamegraph --temporal
   
   def run_context_load(cand_file, bench_file, cache=False, lock=False):
       """
       Loads with context wrappers
       """
   
       with (rxr.open_rasterio(cand_file, cache=cache, lock=lock) as ds,
               rxr.open_rasterio(bench_file, cache=cache, lock=lock) as ds2):
   
           # Pure load
           ds.load()
           ds2.load()
                 
   run_test(partial(run_context_load, **rio_kwargs))

   

   As can be seen the memory is never released. In practice of larger iterations it leads to very large memory consumption due to accumulation of memory.

   The same method is done with xarray and NetCDF files as follows:

   %%memray_flamegraph --temporal
   
   def run_context_load_xr(cand_file, bench_file, cache=False, lock=False):
       """
       Loads with context wrappers
       """
   
       with (xr.open_dataset(cand_file, cache=cache, lock=lock) as ds,
                xr.open_dataset(bench_file, cache=cache, lock=lock) as ds2):
   
           # Pure load
           ds.load()
           ds2.load()
                 
   run_test(partial(run_context_load_xr, **xarray_kwargs))

   

   As can be seen, as expected all memory is released by the end of the operation.

   I tried changing the cache and lock arguments to no avail, I could not get rioxarray to behave similarly. The only way to fully release memory is to directly delete the objects and garbage collect:

   %%memray_flamegraph --temporal
   
   def run_context_load_delete_gc(cand_file, bench_file, cache=False, lock=False):
       """
       Loads with context wrappers and deletes objects
       """
   
       with (rxr.open_rasterio(cand_file, cache=cache, lock=lock) as ds,
               rxr.open_rasterio(bench_file, cache=cache, lock=lock) as ds2):
   
           # Pure load
           ds.load()
           ds2.load()
   
       del ds, ds2
       gc.collect()
       
   run_test(partial(run_context_load_delete_gc, **rio_kwargs))

   

   While this does work, it is not a clean solution and would necessitate prescribing users to do so. I would suggest to relabel this issue as a bug because this takes extra work for a user to diagnose. A user would not expect this behavior when loading GeoTIFFs in rioxarray. This is a dependency in a package I am supporting and the memory accumulation caused issues for workflows as can be seen in this memory profiling example:

   

   Also, I apologize for not being able to simply paste the data as would be preferable but If you like I can provide my notebooks and data which total in about 100 MBs.

   👍React with 👍1Reacted by Fernando Aristizabal

   ❤️React with ❤️4Reacted by Fernando Aristizabal, Papadopoulos Nikolas, Jacob-Shay and nkwonctrees

   🚀React with 🚀2Reacted by Fernando Aristizabal and Jacob-Shay
6. snowman2 commented on Jun 21, 2024

   

   snowman2

   on Jun 21, 2024

   Member

   More actions

   The GDAL cache settings may be worth looking into: https://gdal.org/user/configoptions.html#performance-and-caching
7. snowman2

   added

   bugSomething isn't workingSomething isn't working

   on Jun 21, 2024