New inline_array kwarg for open_dataset

[TomNicholas]

Exposes the inline_array kwarg of dask.array.from_array in xr.open_dataset, and ds/da/variable.chunk.

What setting this to True does is inline the array into the opening/chunking task, which avoids an an extra array object at the start of the task graph. That's useful because the presence of that single common task connecting otherwise independent parts of the graph can confuse the graph optimizer.

With open_dataset(..., inline_array=False):

With open_dataset(..., inline_array=True):

In our case (xGCM) this is important because once inlined the optimizer understands that all the remaining parts of the graph are embarrasingly-parallel, and realizes that it can fuze all our chunk-wise padding tasks into one padding task per chunk.

I think this option could help in any case where someone is opening data from a Zarr store (the reason we had this opener task) or a netCDF file.

The value of the kwarg should be kept optional because in theory inlining is a tradeoff between fewer tasks and more memory use, but I think there might be a case for setting the default to be True?

Questions:

1. How should I test this?
2. Should it default to False or True?
3. inline_array or inline? (inline_array doesn't really make sense for open_dataset, which creates multiple arrays)

• ClosesAvoid Adapters in task graphs? #1895
• Tests added
• User visible changes (including notable bug fixes) are documented in whats-new.rst

@rabernat@jbusecke

[rabernat]

Exposing this options seems like a great idea IMO.

I'm not sure the best way to test this. I think the most basic test is just to make sure the inline=True option gets invoked in the test suite. Going further, one could examine the dask graph to make sure inlining is actually happening, but that sounds fragile and maybe also not xarray's responsibility. Let's just make sure it gets to dask.

[rabernat]

What is the point of this function if it doesn't do anything?

[dcherian]

It means that Dataset.chunk doesn't have to specifically deal with IndexVariable (convenient!) but is the cause of #6204

[TomNicholas]

I think the test failure might be because our minimum dependencies CI uses dask_core=2.30, but the inline_array kwarg was added in dask_core=2021.01.0. That's actually only a few versions afterwards, so I'll try bumping the dependency in this PR.

[TomNicholas]

There is some discussion on the dask PR that added this feature about what the default value for the flag should be. They suggest that at least for datasets opened from zarr it might always be better to inline_array=True. I guess we could change the default in open_zarr, if not in open_dataset?

Tagging @shoyer because he had opinions in that dask PR discussion.

[dcherian]

that's actually only a few versions afterwards, so I'll try bumping the dependency in this PR.

See #6559 getting the env to work took some effort

[TomNicholas]

@dcherian thanks for the heads-up. #6559 has a recent enough version of dask, so if that gets merged then I can just pull it into this PR and avoid messing about with versions here.

[dcherian]

@TomNicholas can you review and check that I didn't miss anything?

[TomNicholas]

We discussed this in the team meeting today.

Questions:

1. How should I test this?

I've added a test which simply counts the number of nodes in the dask graph and checks that it is smaller when inline_array is True.

2. Should it default to False or True?

We decided False for now, and maybe switch it in a future PR

3. inline_array or inline? (inline_array doesn't really make sense for open_dataset, which . creates multiple arrays)

I'll just leave it as inline_array for now.

I think this can be merged?

[dcherian]

LGTM. Thanks @TomNicholas

[TomNicholas]

For some reason counting the number of tasks in the dask graph via len(ds.__dask_graph__())raises an Error on Windows.

> assert num_graph_nodes(inlined) < num_graph_nodes(not_inlined)

...

> os.unlink(fullname)
E PermissionError: [WinError 32] The process cannot access the file because it is being used by another process: 'C:\\Users\\RUNNER~1\\AppData\\Local\\Temp\\tmpnmm87jlx\\temp-415.nc'

I only need to do this in the test, so unless someone knows a more robust way to check the number of tasks in the task graph, I'll just add a skipif windows to the test.

[dcherian]

Lets skip windows for now.

@crusaderky this looks weird:

For some reason counting the number of tasks in the dask graph via len(ds.dask_graph()) raises an Error on Windows.

[shoyer]

For whatever reason, Windows seems to be much stricter about requiring file handles to be explicitly closed. So my guess is that this could be solved by using open_dataset() as a context manager.

[crusaderky]

Lets skip windows for now.

@crusaderky this looks weird:

For some reason counting the number of tasks in the dask graph via len(ds.dask_graph()) raises an Error on Windows.

I think that's the context manager teardown, not the task counting