Repository: dask/dask
Pull request #11794 (2025-02-27 13:02:21)

Refactor import tests
Those import test jobs have been flaky for weeks. I think with this we won't have any issues any more and should cover the same functionality

github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/dask/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

      9 files  ±  0        9 suites  ±0   3h 22m 25s ⏱️ - 5m 27s
 17 818 tests + 13   16 604 ✅ + 13   1 214 💤 ±0  0 ❌ ±0 
159 436 runs  +117  147 280 ✅ +117  12 156 💤 ±0  0 ❌ ±0 

Results for commit 622a543a. ± Comparison against base commit 5f61e423.

[test-results]:data:application/gzip;base64,H4sIAJtrwGcC/02MzQ6DIBAGX8Vw7mEXAbEv09AVElJ/GoST6bsXjd16nPl2ZxMhjn4V96a/NWItMTMMJbkcl7kiSlS6qjrmfcbOov3xYy1EuzQG1F++4vv8ZBdcHKsDFj6lJZ0mlflI61615mROq05aYMlpzYeX9MHXMi3TFHMFYaR0WrUOwKMJ1IEl1N54o8k/Q2gpoEUFg/h8AbiRm4wYAQAA


Patch:
diff --git a/.github/workflows/additional.yml b/.github/workflows/additional.yml
index ef19a5779..92157eeaf 100644
--- a/.github/workflows/additional.yml
+++ b/.github/workflows/additional.yml
@@ -30,29 +30,3 @@ jobs:
 
       - name: Run tests
         run: pytest -v --doctest-modules --ignore-glob='*/test_*.py' --ignore='dask/array/_array_expr/' dask
-
-  imports:
-    runs-on: "ubuntu-latest"
-    timeout-minutes: 90
-    strategy:
-      fail-fast: false
-      matrix:
-        python-version: ["3.10", "3.11", "3.12"]
-    steps:
-      - name: Checkout source
-        uses: actions/checkout@v4.1.5
-
-      - name: Setup Conda
-        uses: conda-incubator/setup-miniconda@v3.1.1
-        with:
-          miniforge-version: latest
-          use-mamba: true
-          channel-priority: strict
-          python-version: "3.10"
-          activate-environment: test-environment
-          auto-activate-base: false
-
-      - name: Run import tests
-        env:
-          PYTHON_VERSION: ${{ matrix.python-version }}
-        run: source continuous_integration/scripts/test_imports.sh
diff --git a/dask/array/blockwise.py b/dask/array/blockwise.py
index c1e7704b4..1d2a063b7 100644
--- a/dask/array/blockwise.py
+++ b/dask/array/blockwise.py
@@ -246,7 +246,18 @@ def blockwise(
     if not out:
         out = "{}-{}".format(
             token or utils.funcname(func).strip("_"),
-            base.tokenize(func, out_ind, argindsstr, dtype, **kwargs),
+            base.tokenize(
+                func,
+                out_ind,
+                argindsstr,
+                adjust_chunks,
+                new_axes,
+                align_arrays,
+                concatenate,
+                meta,
+                dtype,
+                **kwargs,
+            ),
         )
 
     graph = core_blockwise(
diff --git a/dask/array/core.py b/dask/array/core.py
index 41c22b802..c42c1536b 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -815,7 +815,7 @@ def map_blocks(
         )
         name = token
 
-    name = f"{name or funcname(func)}-{tokenize(func, dtype, chunks, drop_axis, new_axis, *args, **kwargs)}"
+    token = f"{name or funcname(func)}"
     new_axes = {}
 
     if isinstance(drop_axis, Number):
@@ -885,7 +885,7 @@ def map_blocks(
             *concat(argpairs),
             expected_ndim=len(out_ind),
             _func=func,
-            name=name,
+            token=token,
             new_axes=new_axes,
             dtype=dtype,
             concatenate=True,
@@ -899,7 +899,7 @@ def map_blocks(
             func,
             out_ind,
             *concat(argpairs),
-            name=name,
+            token=token,
             new_axes=new_axes,
             dtype=dtype,
             concatenate=True,
diff --git a/dask/base.py b/dask/base.py
index 3992ed3fb..b7448d2df 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -18,18 +18,12 @@ from tlz import groupby, merge
 
 from dask import config, local
 from dask._compatibility import EMSCRIPTEN
+from dask._task_spec import DataNode, Dict, List, Task, TaskRef
 from dask.core import flatten
 from dask.core import get as simple_get
-from dask.core import quote
 from dask.system import CPU_COUNT
 from dask.typing import Key, SchedulerGetCallable
-from dask.utils import (
-    apply,
-    ensure_dict,
-    is_namedtuple_instance,
-    key_split,
-    shorten_traceback,
-)
+from dask.utils import ensure_dict, is_namedtuple_instance, key_split, shorten_traceback
 
 _DistributedClient = None
 _get_distributed_client = None
@@ -493,47 +487,45 @@ def unpack_collections(*args, traverse=True):
         if is_dask_collection(expr):
             tok = tokenize(expr)
             if tok not in repack_dsk:
-                repack_dsk[tok] = (getitem, collections_token, len(collections))
+                repack_dsk[tok] = Task(
+                    tok, getitem, TaskRef(collections_token), len(collections)
+                )
                 collections.append(expr)
-            return tok
+            return TaskRef(tok)
 
         tok = uuid.uuid4().hex
+        tsk: DataNode | Task
         if not traverse:
-            tsk = quote(expr)
+            tsk = DataNode(None, expr)
         else:
             # Treat iterators like lists
             typ = list if isinstance(expr, Iterator) else type(expr)
             if typ in (list, tuple, set):
-                tsk = (typ, [_unpack(i) for i in expr])
+                tsk = Task(tok, typ, List(*[_unpack(i) for i in expr]))
             elif typ in (dict, OrderedDict):
-                tsk = (typ, [[_unpack(k), _unpack(v)] for k, v in expr.items()])
+                tsk = Task(
+                    tok, typ, Dict({_unpack(k): _unpack(v) for k, v in expr.items()})
+                )
             elif dataclasses.is_dataclass(expr) and not isinstance(expr, type):
-                tsk = (
-                    apply,
+                tsk = Task(
+                    tok,
                     typ,
-                    (),
-                    (
-                        dict,
-                        [
-                            [f.name, _unpack(getattr(expr, f.name))]
-                            for f in dataclasses.fields(expr)
-                        ],
-                    ),
+                    *[_unpack(getattr(expr, f.name)) for f in dataclasses.fields(expr)],
                 )
             elif is_namedtuple_instance(expr):
-                tsk = (typ, *[_unpack(i) for i in expr])
+                tsk = Task(tok, typ, *[_unpack(i) for i in expr])
             else:
                 return expr
 
         repack_dsk[tok] = tsk
-        return tok
+        return TaskRef(tok)
 
     out = uuid.uuid4().hex
-    repack_dsk[out] = (tuple, [_unpack(i) for i in args])
+    repack_dsk[out] = Task(out, tuple, List(*[_unpack(i) for i in args]))
 
     def repack(results):
         dsk = repack_dsk.copy()
-        dsk[collections_token] = quote(results)
+        dsk[collections_token] = DataNode(collections_token, results)
         return simple_get(dsk, out)
 
     # The original `collections` is kept alive by the closure


Test patch:
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index f086d4405..3f5ab4661 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -5,7 +5,7 @@ on:
   pull_request:
   workflow_dispatch:
   schedule:
-    - cron: "0 6,18 * * *"  
+    - cron: "0 6,18 * * *"
 
 # When this workflow is queued, automatically cancel any previous running
 # or pending jobs from the same branch
@@ -110,6 +110,12 @@ jobs:
           echo "DASK_ARRAY__QUERY_PLANNING: $DASK_ARRAY__QUERY_PLANNING"
           echo "DASK_ARRAY__QUERY_PLANNING=$DASK_ARRAY__QUERY_PLANNING" >> $GITHUB_ENV
 
+      - name: Run import tests
+        id: import_tests
+        env:
+          ARRAYEXPR: ${{ matrix.array-expr }}
+        run: pytest dask/tests/test_imports.py
+
       - name: Run tests
         id: run_tests
         env:
diff --git a/continuous_integration/scripts/run_tests.sh b/continuous_integration/scripts/run_tests.sh
index 0574fbc87..38d61d2f2 100755
--- a/continuous_integration/scripts/run_tests.sh
+++ b/continuous_integration/scripts/run_tests.sh
@@ -16,7 +16,7 @@ if [[ $COVERAGE == 'true' ]]; then
     export XTRATESTARGS="--cov=dask --cov-report=xml --junit-xml pytest.xml $XTRATESTARGS"
 fi
 
-echo "py.testdask --runslow $MARKERS $XTRATESTARGS"
+echo "py.test dask --runslow $MARKERS $XTRATESTARGS"
 py.test dask --runslow $MARKERS $XTRATESTARGS
 
 set +e
diff --git a/continuous_integration/scripts/test_imports.sh b/continuous_integration/scripts/test_imports.sh
deleted file mode 100644
index c23ed3c76..000000000
--- a/continuous_integration/scripts/test_imports.sh
+++ /dev/null
@@ -1,30 +0,0 @@
-#!/usr/bin/env bash
-set -o errexit
-
-
-test_import () {
-    echo "Create environment: python=$PYTHON_VERSION $1"
-    # Create an empty environment
-    mamba create -q -y -n test-imports -c conda-forge python=$PYTHON_VERSION packaging pyyaml fsspec toolz partd click cloudpickle importlib-metadata $1
-    conda activate test-imports
-    if [[ $1 =~ "distributed" ]]; then
-        # dask[distributed] depends on the latest version of distributed
-        python -m pip install git+https://github.com/dask/distributed
-    fi
-    python -m pip install -e .
-    mamba list
-    echo "python -c '$2'"
-    python -c "$2"
-    # Ensure that no non-deterministic objects are tokenized at init time,
-    # which can prevent the library from being imported at all.
-    echo "python -c '$2' (ensure deterministic)"
-    DASK_TOKENIZE__ENSURE_DETERMINISTIC=True python -c "$2"
-    conda deactivate
-    mamba env remove -n test-imports
-}
-
-test_import ""                                "import dask, dask.base, dask.multiprocessing, dask.threaded, dask.optimization, dask.bag, dask.delayed, dask.graph_manipulation, dask.layers"
-test_import "numpy"                           "import dask.array"
-test_import "pandas pyarrow"                  "import dask.dataframe"
-test_import "bokeh"                           "import dask.diagnostics"
-test_import "distributed"                     "import dask.distributed"
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index dbc79c970..c894ca695 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -1857,6 +1857,16 @@ def test_map_blocks_no_array_args():
     assert_eq(x, np.arange(8, dtype=np.float32))
 
 
+def test_map_blocks_unique_name_enforce_dim():
+    def func(some_3d, block_info=None):
+        return some_3d
+
+    input_arr = da.zeros((3, 4, 5), chunks=((3,), (4,), (5,)), dtype=np.float32)
+    x = da.map_blocks(func, input_arr, enforce_ndim=True, dtype=np.float32)
+    y = da.map_blocks(func, input_arr, enforce_ndim=False, dtype=np.float32)
+    assert x._name != y._name
+
+
 def test_map_blocks_unique_name_chunks_dtype():
     def func(block_info=None):
         loc = block_info[None]["array-location"]
diff --git a/dask/tests/test_imports.py b/dask/tests/test_imports.py
new file mode 100644
index 000000000..87b0456d4
--- /dev/null
+++ b/dask/tests/test_imports.py
@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+import importlib
+
+import pytest
+
+
+@pytest.fixture(autouse=True)
+def invalidate_caches():
+    importlib.invalidate_caches()
+
+
+@pytest.mark.parametrize(
+    "module",
+    [
+        "dask",
+        "dask.bag",
+        "dask.base",
+        "dask.delayed",
+        "dask.graph_manipulation",
+        "dask.layers",
+        "dask.multiprocessing",
+        "dask.optimization",
+        "dask.threaded",
+    ],
+)
+def test_defaults(module):
+    __import__(module)
+
+
+def test_array():
+    pytest.importorskip("numpy")
+    import dask.array  # noqa: F401
+
+
+def test_pandas_pyarrow():
+    pytest.importorskip("pandas")
+    pytest.importorskip("pyarrow")
+    import dask.dataframe  # noqa: F401
+
+
+def test_bokeh():
+    pytest.importorskip("bokeh")
+    import dask.diagnostics  # noqa: F401
+
+
+def test_distributed():
+    pytest.importorskip("distributed")
+    import dask.distributed  # noqa: F401