Xarray engine: chunks on a Dask cluster

This notebook demonstrates how to use chunking in computations when a GRIB fieldlist is converted to to Xarray with to_xarray(). Chunking can be used to handle data that does not fit into memory.

In this example the data processing runs on a Dask cluster and is distributed among several Dask workers.

First, set up a default Dask cluster.

[1]:

from dask.distributed import LocalCluster
cluster = LocalCluster()
client = cluster.get_client()
client

[1]:

Client

Client-f8ed0adc-840c-11f0-a338-4ea7d729840e

Connection method: Cluster object	Cluster type: distributed.LocalCluster
Dashboard: http://127.0.0.1:8787/status

Cluster Info

LocalCluster

cfef27f0

Dashboard: http://127.0.0.1:8787/status	Workers: 5
Total threads: 10	Total memory: 16.00 GiB
Status: running	Using processes: True

Scheduler Info

Scheduler

Scheduler-b5f14249-5b0b-4f78-8504-11d4adbd9de2

Comm: tcp://127.0.0.1:64818	Workers: 5
Dashboard: http://127.0.0.1:8787/status	Total threads: 10
Started: Just now	Total memory: 16.00 GiB

Workers

Worker: 0

Comm: tcp://127.0.0.1:64833	Total threads: 2
Dashboard: http://127.0.0.1:64837/status	Memory: 3.20 GiB
Nanny: tcp://127.0.0.1:64821
Local directory: /var/folders/93/w0p869rx17q98wxk83gn9ys40000gn/T/dask-scratch-space/worker-87cc8a2_

Worker: 1

Comm: tcp://127.0.0.1:64832	Total threads: 2
Dashboard: http://127.0.0.1:64840/status	Memory: 3.20 GiB
Nanny: tcp://127.0.0.1:64823
Local directory: /var/folders/93/w0p869rx17q98wxk83gn9ys40000gn/T/dask-scratch-space/worker-gifb9_fd

Worker: 2

Comm: tcp://127.0.0.1:64831	Total threads: 2
Dashboard: http://127.0.0.1:64839/status	Memory: 3.20 GiB
Nanny: tcp://127.0.0.1:64825
Local directory: /var/folders/93/w0p869rx17q98wxk83gn9ys40000gn/T/dask-scratch-space/worker-r956y1p8

Worker: 3

Comm: tcp://127.0.0.1:64835	Total threads: 2
Dashboard: http://127.0.0.1:64838/status	Memory: 3.20 GiB
Nanny: tcp://127.0.0.1:64827
Local directory: /var/folders/93/w0p869rx17q98wxk83gn9ys40000gn/T/dask-scratch-space/worker-1r57z07n

Worker: 4

Comm: tcp://127.0.0.1:64834	Total threads: 2
Dashboard: http://127.0.0.1:64836/status	Memory: 3.20 GiB
Nanny: tcp://127.0.0.1:64829
Local directory: /var/folders/93/w0p869rx17q98wxk83gn9ys40000gn/T/dask-scratch-space/worker-wrpxguiw

2025-08-28 13:47:30,196 - distributed.nanny - WARNING - Restarting worker
2025-08-28 13:47:30,199 - distributed.nanny - WARNING - Restarting worker
2025-08-28 13:47:30,226 - distributed.nanny - WARNING - Restarting worker
2025-08-28 13:47:30,229 - distributed.nanny - WARNING - Restarting worker
2025-08-28 13:47:30,331 - distributed.nanny - WARNING - Restarting worker
2025-08-28 13:47:30,332 - distributed.nanny - WARNING - Restarting worker
2025-08-28 13:47:30,636 - distributed.nanny - WARNING - Restarting worker
2025-08-28 13:47:30,638 - distributed.nanny - WARNING - Restarting worker
2025-08-28 13:47:31,643 - distributed.nanny - WARNING - Restarting worker
2025-08-28 13:47:31,646 - distributed.nanny - WARNING - Restarting worker

Next, we get 2m temperature data for a whole year on a low resolution regular latitude-longitude grid. It contains 2 fields per day (at 0 and 6 UTC). This data obviously fits into memory, so only used for demonstration purposes.

[2]:

import earthkit.data as ekd
ds_fl = ekd.from_source("sample", "t2_1_year_hourly.grib")
len(ds_fl)

[2]:

Now, we convert the GRIB Fieldlist to Xarray using the chunk size of 10 fields.

[3]:

ds = ds_fl.to_xarray(time_dim_mode="valid_time",
                     chunks={"valid_time": 10},
                     add_earthkit_attrs=False)
ds["2t"]

[3]:

<xarray.DataArray '2t' (valid_time: 732, latitude: 13, longitude: 24)> Size: 2MB
dask.array<open_dataset-2t, shape=(732, 13, 24), dtype=float64, chunksize=(10, 13, 24), chunktype=numpy.ndarray>
Coordinates:
  * valid_time  (valid_time) datetime64[ns] 6kB 2020-01-01 ... 2020-12-31T06:...
  * latitude    (latitude) float64 104B 90.0 75.0 60.0 ... -60.0 -75.0 -90.0
  * longitude   (longitude) float64 192B 0.0 15.0 30.0 ... 315.0 330.0 345.0
Attributes:
    standard_name:  air_temperature
    long_name:      2 metre temperature
    units:          K

Finally, we compute the mean along the temporal dimension. Xarray will load data in chunks for this computation keeping the memory usage low.

[4]:

m = ds["2t"].mean(dim="valid_time").load()
m

In the Dask dashboard you should see something like this:

[ ]: