From 6a19642cb61ce542e6d3df834fa3f7b269911c40 Mon Sep 17 00:00:00 2001 From: Lindsey Gray Date: Sat, 7 Dec 2024 14:26:16 -0600 Subject: [PATCH] better check for dask._task_spec --- src/dask_awkward/layers/layers.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/dask_awkward/layers/layers.py b/src/dask_awkward/layers/layers.py index 24b33c99..b5c54523 100644 --- a/src/dask_awkward/layers/layers.py +++ b/src/dask_awkward/layers/layers.py @@ -6,7 +6,7 @@ import dask -_dask_uses_tasks = dask.__version__ >= "2024.12.0" +_dask_uses_tasks = hasattr(dask, "_task_spec") from dask.blockwise import Blockwise, BlockwiseDepDict, blockwise_token from dask.highlevelgraph import MaterializedLayer @@ -15,6 +15,9 @@ from dask_awkward.utils import LazyInputsDict +if _dask_uses_tasks: + from dask._task_spec import Task, TaskRef + if TYPE_CHECKING: from awkward import Array as AwkwardArray from awkward._nplikes.typetracer import TypeTracerReport @@ -173,8 +176,6 @@ def __init__( } if _dask_uses_tasks: - from dask._task_spec import Task, TaskRef - super_kwargs["task"] = Task(name, self.io_func, TaskRef(blockwise_token(0))) else: super_kwargs["dsk"] = {name: (self.io_func, blockwise_token(0))}