Skip to content

Memory usage (memory leak ?).

Hello,

I’m developing an application that uses simplestac to read and process STAC collections. During testing, I’ve encountered significant memory consumption issues when loading and manipulating large STAC datasets with Dask (especially when stacking into xarray).

However, when I use odc-stac to perform the same operation (converting a STAC collection into an xarray.Dataset), the memory usage is much more stable, and the pipeline runs smoothly even with large datasets (e.g., Sentinel-2).

Observations

  • simplestac → stackstac → xarray

    high memory usage, sometimes Dask graph explosion on large collections.

  • simplestac → odc-stac.load() → xarray

    much more efficient, lower memory overhead, better handling of chunking and parallelism.

Proposal

Would it be possible to:

  1. add an optional backend based on odc-stac for generating xarray.Dataset/DataArray ?
  2. Or replace the stackstac step with odc-stac if it fits within the current API?

Example

This script sets up a dask cluster (via SLURM), loads a stac collection into an xarray.DataArray, and applies a simple processing function (predict_block) using xr.map_blocks to simulate a typical block-wise computation workflow (e.g., inference, raster processing, etc.).

The MemorySampler from dask diagnostics is used to monitor and plot memory consumption during the zarr export. In practice, this setup shows how quickly memory usage increases when the data are stacked with stackstac, whereas the same workflow using odc-stac.load() remains stable and efficient.

The shape of the incoming data is (5,12,10980,10980) : actually 5 Sentinel-2 acquisition. The next step is to use Sentinel-2 over a whole year.

import xarray as xr
from dask.distributed import Client
from dask_jobqueue import SLURMCluster
import dask.array as da
import numpy as np
from distributed.diagnostics import MemorySampler

from simplestac.utils import ItemCollection

def predict_block(block):
    H, W = block.sizes["y"], block.sizes["x"]
    return xr.DataArray(
        np.ones((H, W)),
        dims=("y", "x"),
        coords={
            "y": block.coords["y"],
            "x": block.coords["x"],
        },
    )

cluster = SLURMCluster(
    cores=1,  
    processes=1,
    memory="16GB",
    walltime="10:00:00",
    job_extra_directives=[
        "--cpus-per-task=1", 
        "--mem=16GB",
        "--account=theia_oso"
    ],
    interface="ib0"
)
cluster.scale(jobs=10)
client = Client(cluster)
print(client.dashboard_link)

collection_path = "collection_samples.json"
coll = ItemCollection.from_file(collection_path)

chunk_size_y = 5
chunk_size_x = 10980

xarray_data = coll.to_xarray(resolution=10, chunksize=(-1, -1, chunk_size_y, chunk_size_x))

_, _, y_size, x_size = xarray_data.shape

T, B = xarray_data.sizes["time"], xarray_data.sizes["band"]

template = xr.DataArray(
    da.zeros((y_size, x_size), dtype=np.uint8),
    dims=("y", "x"),
    coords={
        "y": xarray_data.coords["y"][:y_size],
        "x": xarray_data.coords["x"][:x_size],
    },
).chunk({"y": chunk_size_y, "x": chunk_size_x})

pred_da = xr.map_blocks(
    predict_block,
    xarray_data,
    template=template,
)

prediction_file = "prediction.zarr"
ms = MemorySampler()
with ms.sample("memory_usage_stackstac"):
    pred_da.to_dataset(name="prediction").to_zarr(
            prediction_file,
            mode="w",
            consolidated=True)
ms.plot(align=True)
client.close()

which produce the memory usage

image

If I replace

xarray_data = coll.to_xarray(resolution=10, chunksize=(-1, -1, chunk_size_y, chunk_size_x))

by

xarray_data = stac_load(coll, chunks={"x":chunk_size_x, "y":chunk_size_y, "band":-1, "time":-1}, resolution=10, resampling="bilinear", dtype="float64") # float64 to mimic stackstac
xarray_data = xarray_data.to_array(dim="band")

the memory usage become

image

Conclusion

odc-stac is faster with less memory than stackstac.

Edited by Arthur Vincent