From 25f18f2ce28f1e05ec2763c51d37da39cb03bdaf Mon Sep 17 00:00:00 2001 From: Tom White Date: Wed, 14 Aug 2024 13:08:21 +0100 Subject: [PATCH 1/4] Make chunk manager an option in `set_options` --- xarray/core/options.py | 3 +++ xarray/namedarray/parallelcompat.py | 5 +++-- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/xarray/core/options.py b/xarray/core/options.py index f5614104357..25bde13bb80 100644 --- a/xarray/core/options.py +++ b/xarray/core/options.py @@ -10,6 +10,7 @@ Options = Literal[ "arithmetic_join", + "chunk_manager", "cmap_divergent", "cmap_sequential", "display_max_rows", @@ -36,6 +37,7 @@ class T_Options(TypedDict): arithmetic_broadcast: bool arithmetic_join: Literal["inner", "outer", "left", "right", "exact"] + chunk_manager: str cmap_divergent: str | Colormap cmap_sequential: str | Colormap display_max_rows: int @@ -62,6 +64,7 @@ class T_Options(TypedDict): OPTIONS: T_Options = { "arithmetic_broadcast": True, "arithmetic_join": "inner", + "chunk_manager": "dask", "cmap_divergent": "RdBu_r", "cmap_sequential": "viridis", "display_max_rows": 12, diff --git a/xarray/namedarray/parallelcompat.py b/xarray/namedarray/parallelcompat.py index 3394aca9a5e..b90e0f99782 100644 --- a/xarray/namedarray/parallelcompat.py +++ b/xarray/namedarray/parallelcompat.py @@ -14,6 +14,7 @@ import numpy as np +from xarray.core.options import OPTIONS from xarray.core.utils import emit_user_level_warning from xarray.namedarray.pycompat import is_chunked_array @@ -101,8 +102,8 @@ def guess_chunkmanager( # use the only option available manager = next(iter(chunkmanagers.keys())) else: - # default to trying to use dask - manager = "dask" + # use the one in options (default dask) + manager = OPTIONS["chunk_manager"] if isinstance(manager, str): if manager not in chunkmanagers: From 1315c4a61744b990161d913b67aeec85a7a765b6 Mon Sep 17 00:00:00 2001 From: Tom White Date: Thu, 15 Aug 2024 09:37:18 +0100 Subject: [PATCH 2/4] Add docs --- doc/internals/chunked-arrays.rst | 3 ++- xarray/core/options.py | 4 +++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/doc/internals/chunked-arrays.rst b/doc/internals/chunked-arrays.rst index ba7ce72c834..eb0a8b197e7 100644 --- a/doc/internals/chunked-arrays.rst +++ b/doc/internals/chunked-arrays.rst @@ -91,7 +91,8 @@ Once the chunkmanager subclass has been registered, xarray objects wrapping the The latter two methods ultimately call the chunkmanager's implementation of ``.from_array``, to which they pass the ``from_array_kwargs`` dict. The ``chunked_array_type`` kwarg selects which registered chunkmanager subclass to dispatch to. It defaults to ``'dask'`` if Dask is installed, otherwise it defaults to whichever chunkmanager is registered if only one is registered. -If multiple chunkmanagers are registered it will raise an error by default. +If multiple chunkmanagers are registered, the ``chunk_manager`` configuration option (which can be set using :py:func:`set_options`) +will be used to determine which chunkmanager to use, defaulting to ``'dask'``. Parallel processing without chunks ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/xarray/core/options.py b/xarray/core/options.py index 25bde13bb80..f31413a2a1a 100644 --- a/xarray/core/options.py +++ b/xarray/core/options.py @@ -175,7 +175,9 @@ class set_options: - "override": if indexes are of same size, rewrite indexes to be those of the first object with that dimension. Indexes for the same dimension must have the same size in all objects. - + chunk_manager : str, default: "dask" + Chunk manager to use for chunked array computations when multiple + options are installed. cmap_divergent : str or matplotlib.colors.Colormap, default: "RdBu_r" Colormap to use for divergent data plots. If string, must be matplotlib built-in colormap. Can also be a Colormap object From dfc408cfb027855578647e5f3bfcd4abbfaa42b5 Mon Sep 17 00:00:00 2001 From: Tom White Date: Thu, 15 Aug 2024 09:53:08 +0100 Subject: [PATCH 3/4] Add unit test --- xarray/tests/test_parallelcompat.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/xarray/tests/test_parallelcompat.py b/xarray/tests/test_parallelcompat.py index dbe40be710c..67c68aac534 100644 --- a/xarray/tests/test_parallelcompat.py +++ b/xarray/tests/test_parallelcompat.py @@ -6,6 +6,7 @@ import numpy as np import pytest +from xarray import set_options from xarray.core.types import T_Chunks, T_DuckArray, T_NormalizedChunks from xarray.namedarray._typing import _Chunks from xarray.namedarray.daskmanager import DaskManager @@ -152,6 +153,11 @@ def test_get_chunkmanger(self, register_dummy_chunkmanager) -> None: chunkmanager = guess_chunkmanager("dummy") assert isinstance(chunkmanager, DummyChunkManager) + def test_get_chunkmanger_via_set_options(self, register_dummy_chunkmanager) -> None: + with set_options(chunk_manager="dummy"): + chunkmanager = guess_chunkmanager(None) + assert isinstance(chunkmanager, DummyChunkManager) + def test_fail_on_nonexistent_chunkmanager(self) -> None: with pytest.raises(ValueError, match="unrecognized chunk manager foo"): guess_chunkmanager("foo") From c9a95728d4dcf0dc8ac0b6567b7a93094a1fc18e Mon Sep 17 00:00:00 2001 From: Tom White Date: Fri, 16 Aug 2024 10:41:01 +0100 Subject: [PATCH 4/4] Add what's new entry --- doc/whats-new.rst | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 67e58e52619..c8f3a40e87f 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -22,7 +22,8 @@ v2024.07.1 (unreleased) New Features ~~~~~~~~~~~~ - +- Make chunk manager an option in ``set_options`` (:pull:`9362`). + By `Tom White `_. Breaking changes ~~~~~~~~~~~~~~~~