Repository: dask/dask (python)

Issue:
Stricter data type for dask keys
- Follow up to dask/distributed#8083 
- Blocks https://github.com/dask/distributed/issues/8130

dask/distributed#8083 recently redefined dask keys from any hashable to bytes, float, int, str, or tuples thereof.

While it is possible to run dask/dask on schedulers other than distributed, I don't think it's a healthy idea to have stricter restrictions in dask/distributed than in dask/dask.

# Actions
1. Create a new type alias,
```python
Key: TypeAlias = Union[str, bytes, int, float, tuple["Key", ...]]
```
2. Search the dask/dask codebase for all instances of `Hashable` and replace wherever applicable with `Key`
3. Update documentation, e.g. https://docs.dask.org/en/stable/spec.html#definitions
4. Move the new function `validate_key` from dask/distributed to dask/dask; add unit tests


Pull request:
Stricter data type for dask keys
Closes #10461

Patch:
diff --git a/.github/workflows/additional.yml b/.github/workflows/additional.yml
index 717717b81..dc72fffb2 100644
--- a/.github/workflows/additional.yml
+++ b/.github/workflows/additional.yml
@@ -13,7 +13,7 @@ jobs:
     timeout-minutes: 90
     steps:
       - name: Checkout source
-        uses: actions/checkout@v3.6.0
+        uses: actions/checkout@v4.0.0
 
       - name: Setup Conda Environment
         uses: conda-incubator/setup-miniconda@v2.2.0
@@ -42,7 +42,7 @@ jobs:
         python-version: ["3.9", "3.10", "3.11"]
     steps:
       - name: Checkout source
-        uses: actions/checkout@v3.6.0
+        uses: actions/checkout@v4.0.0
 
       - name: Setup Conda
         uses: conda-incubator/setup-miniconda@v2.2.0
diff --git a/.github/workflows/conda.yml b/.github/workflows/conda.yml
index 8a7ac3b97..4e9bbb0b2 100644
--- a/.github/workflows/conda.yml
+++ b/.github/workflows/conda.yml
@@ -25,7 +25,7 @@ jobs:
     name: Build (and upload)
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v3.6.0
+      - uses: actions/checkout@v4.0.0
         with:
           fetch-depth: 0
       - name: Set up Python
diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml
index 955fa92ce..007f4cf1e 100644
--- a/.github/workflows/pre-commit.yml
+++ b/.github/workflows/pre-commit.yml
@@ -11,7 +11,7 @@ jobs:
     name: pre-commit hooks
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v3.6.0
+      - uses: actions/checkout@v4.0.0
       - uses: actions/setup-python@v4
         with:
           python-version: '3.9'
diff --git a/.github/workflows/update-gpuci.yml b/.github/workflows/update-gpuci.yml
index e565efa6f..634945b93 100644
--- a/.github/workflows/update-gpuci.yml
+++ b/.github/workflows/update-gpuci.yml
@@ -11,7 +11,7 @@ jobs:
     if: github.repository == 'dask/dask'
 
     steps:
-      - uses: actions/checkout@v3.6.0
+      - uses: actions/checkout@v4.0.0
 
       - name: Parse current axis YAML
         id: rapids_current
diff --git a/.github/workflows/upstream.yml b/.github/workflows/upstream.yml
index 6e3f7a4d5..cf732d757 100644
--- a/.github/workflows/upstream.yml
+++ b/.github/workflows/upstream.yml
@@ -20,7 +20,7 @@ jobs:
     outputs:
       test-upstream: ${{ steps.detect-trigger.outputs.trigger-found }}
     steps:
-      - uses: actions/checkout@v3.6.0
+      - uses: actions/checkout@v4.0.0
         with:
           fetch-depth: 2
       - uses: xarray-contrib/ci-trigger@v1
@@ -47,7 +47,7 @@ jobs:
 
     steps:
       - name: Checkout source
-        uses: actions/checkout@v3.6.0
+        uses: actions/checkout@v4.0.0
 
       - name: Setup Conda Environment
         uses: conda-incubator/setup-miniconda@v2.2.0
diff --git a/dask/array/core.py b/dask/array/core.py
index 1c6622cac..1a1122d96 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -13,11 +13,11 @@ import warnings
 from bisect import bisect
 from collections.abc import (
     Collection,
-    Hashable,
     Iterable,
     Iterator,
     Mapping,
     MutableMapping,
+    Sequence,
 )
 from functools import partial, reduce, wraps
 from itertools import product, zip_longest
@@ -61,6 +61,7 @@ from dask.delayed import Delayed, delayed
 from dask.highlevelgraph import HighLevelGraph, MaterializedLayer
 from dask.layers import ArraySliceDep, reshapelist
 from dask.sizeof import sizeof
+from dask.typing import Graph, Key, NestedKeys
 from dask.utils import (
     IndexCallable,
     SerializableLock,
@@ -1387,13 +1388,13 @@ class Array(DaskMethodsMixin):
     def __reduce__(self):
         return (Array, (self.dask, self.name, self.chunks, self.dtype, self._meta))
 
-    def __dask_graph__(self):
+    def __dask_graph__(self) -> Graph:
         return self.dask
 
-    def __dask_layers__(self):
+    def __dask_layers__(self) -> Sequence[str]:
         return (self.name,)
 
-    def __dask_keys__(self):
+    def __dask_keys__(self) -> NestedKeys:
         if self._cached_keys is not None:
             return self._cached_keys
 
@@ -4475,7 +4476,7 @@ def insert_to_ooc(
 
 
 def retrieve_from_ooc(
-    keys: Collection[Hashable], dsk_pre: Mapping, dsk_post: Mapping
+    keys: Collection[Key], dsk_pre: Graph, dsk_post: Graph
 ) -> dict[tuple, Any]:
     """
     Creates a Dask graph for loading stored ``keys`` from ``dsk``.
diff --git a/dask/bag/core.py b/dask/bag/core.py
index f2c7b2e2a..bfc68478f 100644
--- a/dask/bag/core.py
+++ b/dask/bag/core.py
@@ -7,7 +7,7 @@ import operator
 import uuid
 import warnings
 from collections import defaultdict
-from collections.abc import Hashable, Iterable, Iterator, Mapping
+from collections.abc import Iterable, Iterator, Mapping, Sequence
 from functools import partial, reduce, wraps
 from random import Random
 from urllib.request import urlopen
@@ -54,6 +54,7 @@ from dask.delayed import Delayed, unpack_collections
 from dask.highlevelgraph import HighLevelGraph
 from dask.optimization import cull, fuse, inline
 from dask.sizeof import sizeof
+from dask.typing import Graph, NestedKeys
 from dask.utils import (
     apply,
     digit,
@@ -366,13 +367,13 @@ class Item(DaskMethodsMixin):
                 f"Layer {self._layer} not in the HighLevelGraph's layers: {list(dsk.layers)}"
             )
 
-    def __dask_graph__(self):
+    def __dask_graph__(self) -> Graph:
         return self.dask
 
-    def __dask_keys__(self):
+    def __dask_keys__(self) -> NestedKeys:
         return [self.key]
 
-    def __dask_layers__(self):
+    def __dask_layers__(self) -> Sequence[str]:
         return (self._layer,)
 
     def __dask_tokenize__(self):
@@ -476,13 +477,13 @@ class Bag(DaskMethodsMixin):
         self.name = name
         self.npartitions = npartitions
 
-    def __dask_graph__(self):
+    def __dask_graph__(self) -> Graph:
         return self.dask
 
-    def __dask_keys__(self) -> list[Hashable]:
+    def __dask_keys__(self) -> NestedKeys:
         return [(self.name, i) for i in range(self.npartitions)]
 
-    def __dask_layers__(self):
+    def __dask_layers__(self) -> Sequence[str]:
         return (self.name,)
 
     def __dask_tokenize__(self):
diff --git a/dask/base.py b/dask/base.py
index 9e95c8460..044beb7ac 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -11,7 +11,7 @@ import threading
 import uuid
 import warnings
 from collections import OrderedDict
-from collections.abc import Callable, Iterator, Mapping
+from collections.abc import Callable, Hashable, Iterator, Mapping
 from concurrent.futures import Executor
 from contextlib import contextmanager
 from contextvars import ContextVar
@@ -19,7 +19,7 @@ from enum import Enum
 from functools import partial
 from numbers import Integral, Number
 from operator import getitem
-from typing import TYPE_CHECKING, Any, Literal, Protocol
+from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeVar
 
 from tlz import curry, groupby, identity, merge
 from tlz.functoolz import Compose
@@ -31,7 +31,7 @@ from dask.core import get as simple_get
 from dask.core import literal, quote
 from dask.hashing import hash_buffer_hex
 from dask.system import CPU_COUNT
-from dask.typing import SchedulerGetCallable
+from dask.typing import Key, SchedulerGetCallable
 from dask.utils import (
     Dispatch,
     apply,
@@ -1498,7 +1498,7 @@ def get_collection_names(collection) -> set[str]:
     return {get_name_from_key(k) for k in flatten(collection.__dask_keys__())}
 
 
-def get_name_from_key(key) -> str:
+def get_name_from_key(key: Key) -> str:
     """Given a dask collection's key, extract the collection name.
 
     Parameters
@@ -1518,10 +1518,13 @@ def get_name_from_key(key) -> str:
         return key[0]
     if isinstance(key, str):
         return key
-    raise TypeError(f"Expected str or tuple[str, Hashable, ...]; got {key}")
+    raise TypeError(f"Expected str or a tuple starting with str; got {key!r}")
 
 
-def replace_name_in_key(key, rename: Mapping[str, str]):
+KeyOrStrT = TypeVar("KeyOrStrT", Key, str)
+
+
+def replace_name_in_key(key: KeyOrStrT, rename: Mapping[str, str]) -> KeyOrStrT:
     """Given a dask collection's key, replace the collection name with a new one.
 
     Parameters
@@ -1546,10 +1549,10 @@ def replace_name_in_key(key, rename: Mapping[str, str]):
         return (rename.get(key[0], key[0]),) + key[1:]
     if isinstance(key, str):
         return rename.get(key, key)
-    raise TypeError(f"Expected str or tuple[str, Hashable, ...]; got {key}")
+    raise TypeError(f"Expected str or a tuple starting with str; got {key!r}")
 
 
-def clone_key(key, seed):
+def clone_key(key: KeyOrStrT, seed: Hashable) -> KeyOrStrT:
     """Clone a key from a Dask collection, producing a new key with the same prefix and
     indices and a token which is a deterministic function of the previous key and seed.
 
@@ -1567,4 +1570,4 @@ def clone_key(key, seed):
     if isinstance(key, str):
         prefix = key_split(key)
         return prefix + "-" + tokenize(key, seed)
-    raise TypeError(f"Expected str or tuple[str, Hashable, ...]; got {key}")
+    raise TypeError(f"Expected str or a tuple starting with str; got {key!r}")
diff --git a/dask/blockwise.py b/dask/blockwise.py
index a978616ab..51d822418 100644
--- a/dask/blockwise.py
+++ b/dask/blockwise.py
@@ -14,6 +14,7 @@ from dask.base import clone_key, get_name_from_key, tokenize
 from dask.core import flatten, ishashable, keys_in_tasks, reverse_dict
 from dask.highlevelgraph import HighLevelGraph, Layer
 from dask.optimization import SubgraphCallable, fuse
+from dask.typing import Graph, Key
 from dask.utils import (
     _deprecated,
     apply,
@@ -377,7 +378,7 @@ class Blockwise(Layer):
 
     output: str
     output_indices: tuple[str, ...]
-    dsk: Mapping[str, tuple]
+    dsk: Graph
     indices: tuple[tuple[str, tuple[str, ...] | None], ...]
     numblocks: Mapping[str, Sequence[int]]
     concatenate: bool | None
@@ -389,7 +390,7 @@ class Blockwise(Layer):
         self,
         output: str,
         output_indices: Iterable[str],
-        dsk: Mapping[str, tuple],
+        dsk: Graph,
         indices: Iterable[tuple[str | BlockwiseDep, Iterable[str] | None]],
         numblocks: Mapping[str, Sequence[int]],
         concatenate: bool | None = None,
@@ -579,7 +580,7 @@ class Blockwise(Layer):
 
     def cull(
         self, keys: set, all_hlg_keys: Iterable
-    ) -> tuple[Layer, Mapping[Hashable, set]]:
+    ) -> tuple[Layer, Mapping[Key, set[Key]]]:
         # Culling is simple for Blockwise layers.  We can just
         # collect a set of required output blocks (tuples), and
         # only construct graph for these blocks in `make_blockwise_graph`
@@ -598,9 +599,9 @@ class Blockwise(Layer):
 
     def clone(
         self,
-        keys: set,
+        keys: set[Key],
         seed: Hashable,
-        bind_to: Hashable = None,
+        bind_to: Key | None = None,
     ) -> tuple[Layer, bool]:
         names = {get_name_from_key(k) for k in keys}
         # We assume that 'keys' will contain either all or none of the output keys of
@@ -616,6 +617,7 @@ class Blockwise(Layer):
         is_leaf = True
 
         indices = []
+        k: Key
         for k, idxv in self.indices:
             # Note: k may not be a key and thus not be hashable in the case where
             # one or more args of blockwise() are sequences of literals;
@@ -627,7 +629,7 @@ class Blockwise(Layer):
 
             indices.append((k, idxv))
 
-        numblocks = {}
+        numblocks: dict[str, Sequence[int]] = {}
         for k, nbv in self.numblocks.items():
             if k in names:
                 is_leaf = False
diff --git a/dask/core.py b/dask/core.py
index 7beec491e..b87340172 100644
--- a/dask/core.py
+++ b/dask/core.py
@@ -1,6 +1,10 @@
 from __future__ import annotations
 
 from collections import defaultdict
+from collections.abc import Collection, Iterable
+from typing import Any, cast
+
+from dask.typing import Key
 
 no_default = "__no_default__"
 
@@ -15,6 +19,10 @@ def ishashable(x):
     True
     >>> ishashable([1])
     False
+
+    See Also
+    --------
+    iskey
     """
     try:
         hash(x)
@@ -156,7 +164,7 @@ def get(dsk, out, cache=None):
     return result
 
 
-def keys_in_tasks(keys, tasks, as_list=False):
+def keys_in_tasks(keys: Collection[Key], tasks: Iterable[Any], as_list: bool = False):
     """Returns the keys in `keys` that are also in `tasks`
 
     Examples
@@ -193,32 +201,28 @@ def keys_in_tasks(keys, tasks, as_list=False):
     return ret if as_list else set(ret)
 
 
-def find_all_possible_keys(tasks) -> set:
-    """Returns all possible keys in `tasks` including hashable literals.
+def iskey(key: object) -> bool:
+    """Return True if the given object is a potential dask key; False otherwise.
 
-    The definition of a key in a Dask graph is any hashable object
-    that is not a task. This function returns all such objects in
-    `tasks` even if the object is in fact a literal.
+    The definition of a key in a Dask graph is any str, bytes, int, float, or tuple
+    thereof.
 
+    See Also
+    --------
+    ishashable
+    validate_key
+    dask.typing.Key
     """
-    ret = set()
-    while tasks:
-        work = []
-        for w in tasks:
-            typ = type(w)
-            if typ is tuple and w and callable(w[0]):  # istask(w)
-                work.extend(w[1:])
-            elif typ is list:
-                work.extend(w)
-            elif typ is dict:
-                work.extend(w.values())
-            else:
-                try:
-                    ret.add(w)
-                except TypeError:  # not hashable
-                    pass
-        tasks = work
-    return ret
+    typ = type(key)
+    if typ is tuple:
+        return all(iskey(i) for i in cast(tuple, key))
+    return typ in {bytes, int, float, str}
+
+
+def validate_key(key: object) -> None:
+    """Validate the format of a dask key."""
+    if not iskey(key):
+        raise TypeError(f"Unexpected key type {type(key)} (value: {key!r})")
 
 
 def get_dependencies(dsk, key=None, task=no_default, as_list=False):
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
index 742ad0bc0..8182fc6e7 100644
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import operator
 import warnings
-from collections.abc import Callable, Hashable, Iterator, Mapping, Sequence
+from collections.abc import Callable, Iterator, Mapping, Sequence
 from functools import partial, wraps
 from numbers import Integral, Number
 from operator import getitem
@@ -78,6 +78,7 @@ from dask.dataframe.utils import (
 from dask.delayed import Delayed, delayed, unpack_collections
 from dask.highlevelgraph import HighLevelGraph
 from dask.layers import DataFrameTreeReduction
+from dask.typing import Graph, NestedKeys
 from dask.utils import (
     IndexCallable,
     M,
@@ -220,16 +221,16 @@ class Scalar(DaskMethodsMixin, OperatorMethodMixin):
             )
         self._meta = meta
 
-    def __dask_graph__(self):
+    def __dask_graph__(self) -> Graph:
         return self.dask
 
-    def __dask_keys__(self):
+    def __dask_keys__(self) -> NestedKeys:
         return [self.key]
 
     def __dask_tokenize__(self):
         return self._name
 
-    def __dask_layers__(self):
+    def __dask_layers__(self) -> Sequence[str]:
         return (self._name,)
 
     __dask_optimize__ = globalmethod(
@@ -449,13 +450,13 @@ class _Frame(DaskMethodsMixin, OperatorMethodMixin):
                 self._meta = result._meta
                 self.divisions = result.divisions
 
-    def __dask_graph__(self):
+    def __dask_graph__(self) -> Graph:
         return self.dask
 
-    def __dask_keys__(self) -> list[Hashable]:
+    def __dask_keys__(self) -> NestedKeys:
         return [(self._name, i) for i in range(self.npartitions)]
 
-    def __dask_layers__(self):
+    def __dask_layers__(self) -> Sequence[str]:
         return (self._name,)
 
     def __dask_tokenize__(self):
diff --git a/dask/delayed.py b/dask/delayed.py
index 882d4c9b8..51886583f 100644
--- a/dask/delayed.py
+++ b/dask/delayed.py
@@ -4,7 +4,7 @@ import operator
 import types
 import uuid
 import warnings
-from collections.abc import Iterator
+from collections.abc import Iterator, Sequence
 from dataclasses import fields, is_dataclass, replace
 from functools import partial
 
@@ -22,6 +22,7 @@ from dask.base import tokenize as _tokenize
 from dask.context import globalmethod
 from dask.core import flatten, quote
 from dask.highlevelgraph import HighLevelGraph
+from dask.typing import Graph, NestedKeys
 from dask.utils import (
     OperatorMethodMixin,
     apply,
@@ -277,7 +278,7 @@ def delayed(obj, name=None, pure=None, nout=None, traverse=True):
     ----------
     obj : object
         The function or object to wrap
-    name : string or hashable, optional
+    name : Dask key, optional
         The key to use in the underlying graph for the wrapped object. Defaults
         to hashing content. Note that this only affects the name of the object
         wrapped by this call to delayed, and *not* the output of delayed
@@ -545,13 +546,13 @@ class Delayed(DaskMethodsMixin, OperatorMethodMixin):
     def dask(self):
         return self._dask
 
-    def __dask_graph__(self):
+    def __dask_graph__(self) -> Graph:
         return self.dask
 
-    def __dask_keys__(self):
+    def __dask_keys__(self) -> NestedKeys:
         return [self.key]
 
-    def __dask_layers__(self):
+    def __dask_layers__(self) -> Sequence[str]:
         return (self._layer,)
 
     def __dask_tokenize__(self):
diff --git a/dask/graph_manipulation.py b/dask/graph_manipulation.py
index a5f593181..a80fd8112 100644
--- a/dask/graph_manipulation.py
+++ b/dask/graph_manipulation.py
@@ -20,6 +20,7 @@ from dask.blockwise import blockwise
 from dask.core import flatten
 from dask.delayed import Delayed, delayed
 from dask.highlevelgraph import HighLevelGraph, Layer, MaterializedLayer
+from dask.typing import Key
 
 __all__ = ("bind", "checkpoint", "clone", "wait_on")
 
@@ -309,7 +310,7 @@ def _bind_one(
     child: T,
     blocker: Delayed | None,
     omit_layers: set[str],
-    omit_keys: set[Hashable],
+    omit_keys: set[Key],
     seed: Hashable,
 ) -> T:
     prev_coll_names = get_collection_names(child)
diff --git a/dask/highlevelgraph.py b/dask/highlevelgraph.py
index 2f9fd37d1..19aa3befe 100644
--- a/dask/highlevelgraph.py
+++ b/dask/highlevelgraph.py
@@ -3,7 +3,7 @@ from __future__ import annotations
 import abc
 import copy
 import html
-from collections.abc import Hashable, Iterable, KeysView, Mapping, Set
+from collections.abc import Collection, Hashable, Iterable, KeysView, Mapping, Set
 from typing import Any
 
 import tlz as toolz
@@ -12,6 +12,7 @@ import dask
 from dask import config
 from dask.base import clone_key, flatten, is_dask_collection
 from dask.core import keys_in_tasks, reverse_dict
+from dask.typing import Graph, Key
 from dask.utils import ensure_dict, import_required, key_split
 from dask.widgets import get_template
 
@@ -101,8 +102,8 @@ class Layer(Mapping):
         return self.keys()  # this implementation will materialize the graph
 
     def cull(
-        self, keys: set, all_hlg_keys: Iterable
-    ) -> tuple[Layer, Mapping[Hashable, set]]:
+        self, keys: set[Key], all_hlg_keys: Collection[Key]
+    ) -> tuple[Layer, Mapping[Key, set[Key]]]:
         """Remove unnecessary tasks from the layer
 
         In other words, return a new Layer with only the tasks required to
@@ -148,14 +149,14 @@ class Layer(Mapping):
 
         return MaterializedLayer(out, annotations=self.annotations), ret_deps
 
-    def get_dependencies(self, key: Hashable, all_hlg_keys: Iterable) -> set:
+    def get_dependencies(self, key: Key, all_hlg_keys: Collection[Key]) -> set:
         """Get dependencies of `key` in the layer
 
         Parameters
         ----------
-        key: Hashable
+        key:
             The key to find dependencies of
-        all_hlg_keys: Iterable
+        all_hlg_keys:
             All keys in the high level graph.
 
         Returns
@@ -169,7 +170,7 @@ class Layer(Mapping):
         self,
         keys: set,
         seed: Hashable,
-        bind_to: Hashable = None,
+        bind_to: Key | None = None,
     ) -> tuple[Layer, bool]:
         """Clone selected keys in the layer, as well as references to keys in other
         layers
@@ -347,7 +348,7 @@ class HighLevelGraph(Mapping):
         The subgraph layers, keyed by a unique name
     dependencies : Mapping[str, set[str]]
         The set of layers on which each layer depends
-    key_dependencies : Mapping[Hashable, set], optional
+    key_dependencies : dict[Key, set], optional
         Mapping (some) keys in the high level graph to their dependencies. If
         a key is missing, its dependencies will be calculated on-the-fly.
 
@@ -393,16 +394,16 @@ class HighLevelGraph(Mapping):
     """
 
     layers: Mapping[str, Layer]
-    dependencies: Mapping[str, Set]
-    key_dependencies: dict[Hashable, Set]
+    dependencies: Mapping[str, Set[str]]
+    key_dependencies: dict[Key, Set[Key]]
     _to_dict: dict
     _all_external_keys: set
 
     def __init__(
         self,
-        layers: Mapping[str, Mapping],
-        dependencies: Mapping[str, Set],
-        key_dependencies: dict[Hashable, Set] | None = None,
+        layers: Mapping[str, Graph],
+        dependencies: Mapping[str, Set[str]],
+        key_dependencies: dict[Key, Set[Key]] | None = None,
     ):
         self.dependencies = dependencies
         self.key_dependencies = key_dependencies or {}
@@ -565,7 +566,7 @@ class HighLevelGraph(Mapping):
     def values(self):
         return self.to_dict().values()
 
-    def get_all_dependencies(self) -> dict[Hashable, Set]:
+    def get_all_dependencies(self) -> dict[Key, Set[Key]]:
         """Get dependencies of all keys
 
         This will in most cases materialize all layers, which makes
diff --git a/dask/layers.py b/dask/layers.py
index c0bd0b7d0..9b7ee2756 100644
--- a/dask/layers.py
+++ b/dask/layers.py
@@ -214,8 +214,8 @@ def _expand_keys_around_center(k, dims, name=None, axes=None):
 
     Parameters
     ----------
-    k: tuple
-        They key around which to generate new keys
+    k: Key
+        The key around which to generate new keys
     dims: Sequence[int]
         The number of chunks in each dimension
     name: Option[str]
diff --git a/dask/local.py b/dask/local.py
index a0ac53dd1..0a5b724bb 100644
--- a/dask/local.py
+++ b/dask/local.py
@@ -109,7 +109,7 @@ See the function ``inline_functions`` for more information.
 from __future__ import annotations
 
 import os
-from collections.abc import Hashable, Mapping, Sequence
+from collections.abc import Mapping, Sequence
 from concurrent.futures import Executor, Future
 from functools import partial
 from queue import Empty, Queue
@@ -118,6 +118,7 @@ from dask import config
 from dask.callbacks import local_callbacks, unpack_callbacks
 from dask.core import _execute_task, flatten, get_dependencies, has_tasks, reverse_dict
 from dask.order import order
+from dask.typing import Key
 
 if os.name == "nt":
     # Python 3 windows Queue.get doesn't handle interrupts properly. To
@@ -548,7 +549,7 @@ class SynchronousExecutor(Executor):
 synchronous_executor = SynchronousExecutor()
 
 
-def get_sync(dsk: Mapping, keys: Sequence[Hashable] | Hashable, **kwargs):
+def get_sync(dsk: Mapping, keys: Sequence[Key] | Key, **kwargs):
     """A naive synchronous version of get_async
 
     Can be useful for debugging.
diff --git a/dask/multiprocessing.py b/dask/multiprocessing.py
index 137ae5c34..54401115b 100644
--- a/dask/multiprocessing.py
+++ b/dask/multiprocessing.py
@@ -7,7 +7,7 @@ import os
 import pickle
 import sys
 import traceback
-from collections.abc import Hashable, Mapping, Sequence
+from collections.abc import Mapping, Sequence
 from concurrent.futures import ProcessPoolExecutor
 from functools import partial
 from warnings import warn
@@ -18,6 +18,7 @@ from dask import config
 from dask.local import MultiprocessingPoolExecutor, get_async, reraise
 from dask.optimization import cull, fuse
 from dask.system import CPU_COUNT
+from dask.typing import Key
 from dask.utils import ensure_dict
 
 
@@ -145,7 +146,7 @@ def get_context():
 
 def get(
     dsk: Mapping,
-    keys: Sequence[Hashable] | Hashable,
+    keys: Sequence[Key] | Key,
     num_workers=None,
     func_loads=None,
     func_dumps=None,
diff --git a/dask/optimization.py b/dask/optimization.py
index caa11e9ce..91a91d827 100644
--- a/dask/optimization.py
+++ b/dask/optimization.py
@@ -954,7 +954,7 @@ class SubgraphCallable:
     ----------
     dsk : dict
         A dask graph
-    outkey : hashable
+    outkey : Dask key
         The output key from the graph
     inkeys : list
         A list of keys to be used as arguments to the callable.
diff --git a/dask/threaded.py b/dask/threaded.py
index b1fadf186..8ecc9eda1 100644
--- a/dask/threaded.py
+++ b/dask/threaded.py
@@ -10,13 +10,14 @@ import multiprocessing.pool
 import sys
 import threading
 from collections import defaultdict
-from collections.abc import Hashable, Mapping, Sequence
+from collections.abc import Mapping, Sequence
 from concurrent.futures import Executor, ThreadPoolExecutor
 from threading import Lock, current_thread
 
 from dask import config
 from dask.local import MultiprocessingPoolExecutor, get_async
 from dask.system import CPU_COUNT
+from dask.typing import Key
 
 
 def _thread_get_id():
@@ -35,7 +36,7 @@ def pack_exception(e, dumps):
 
 def get(
     dsk: Mapping,
-    keys: Sequence[Hashable] | Hashable,
+    keys: Sequence[Key] | Key,
     cache=None,
     num_workers=None,
     pool=None,
diff --git a/dask/typing.py b/dask/typing.py
index 839b8c280..b9c6beb17 100644
--- a/dask/typing.py
+++ b/dask/typing.py
@@ -2,36 +2,44 @@ from __future__ import annotations
 
 import abc
 from collections.abc import Callable, Hashable, Mapping, Sequence
-from typing import TYPE_CHECKING, Any, Protocol, TypeVar, runtime_checkable
+from typing import TYPE_CHECKING, Any, Protocol, TypeVar, Union, runtime_checkable
 
 if TYPE_CHECKING:
     # IPython import is relatively slow. Avoid if not necessary
     from IPython.display import DisplayObject
 
+    # TODO import from typing (requires Python >=3.10)
+    from typing_extensions import TypeAlias
 
 CollType = TypeVar("CollType", bound="DaskCollection")
 CollType_co = TypeVar("CollType_co", bound="DaskCollection", covariant=True)
 PostComputeCallable = Callable
 
 
+Key: TypeAlias = Union[str, bytes, int, float, tuple["Key", ...]]
+Graph: TypeAlias = Mapping[Key, Any]
+# Potentially nested list of Dask keys
+NestedKeys: TypeAlias = list[Union[Key, "NestedKeys"]]
+
+
 class SchedulerGetCallable(Protocol):
     """Protocol defining the signature of a ``__dask_scheduler__`` callable."""
 
     def __call__(
         self,
-        dsk: Mapping,
-        keys: Sequence[Hashable] | Hashable,
+        dsk: Graph,
+        keys: Sequence[Key] | Key,
         **kwargs: Any,
     ) -> Any:
         """Method called as the default scheduler for a collection.
 
         Parameters
         ----------
-        dsk : Mapping
+        dsk :
             The task graph.
-        keys : Sequence[Hashable] | Hashable
+        keys :
             Key(s) corresponding to the desired data.
-        **kwargs : Any
+        **kwargs :
             Additional arguments.
 
         Returns
@@ -48,7 +56,7 @@ class PostPersistCallable(Protocol[CollType_co]):
 
     def __call__(
         self,
-        dsk: Mapping,
+        dsk: Graph,
         *args: Any,
         rename: Mapping[str, str] | None = None,
     ) -> CollType_co:
@@ -88,7 +96,7 @@ class DaskCollection(Protocol):
     """Protocol defining the interface of a Dask collection."""
 
     @abc.abstractmethod
-    def __dask_graph__(self) -> Mapping:
+    def __dask_graph__(self) -> Graph:
         """The Dask task graph.
 
         The core Dask collections (Array, DataFrame, Bag, and Delayed)
@@ -110,7 +118,7 @@ class DaskCollection(Protocol):
         raise NotImplementedError("Inheriting class must implement this method.")
 
     @abc.abstractmethod
-    def __dask_keys__(self) -> list[Hashable]:
+    def __dask_keys__(self) -> NestedKeys:
         """The output keys of the task graph.
 
         Note that there are additional constraints on keys for a Dask
@@ -118,11 +126,11 @@ class DaskCollection(Protocol):
         specification documentation <spec>`. These additional
         constraints are described below.
 
-        All keys must either be non-empty strings or tuples where the
-        first element is a non-empty string, followed by zero or more
-        arbitrary hashables. The non-empty string is commonly known as
-        the *collection name*. All collections embedded in the dask
-        package have exactly one name, but this is not a requirement.
+        All keys must either be non-empty strings or tuples where the first element is a
+        non-empty string, followed by zero or more arbitrary str, bytes, int, float, or
+        tuples thereof. The non-empty string is commonly known as the *collection name*.
+        All collections embedded in the dask package have exactly one name, but this is
+        not a requirement.
 
         These are all valid outputs:
 
@@ -213,10 +221,10 @@ class DaskCollection(Protocol):
 
     Parameters
     ----------
-    dsk : MutableMapping
+    dsk : Graph
         The merged graphs from all collections sharing the same
         ``__dask_optimize__`` method.
-    keys : Sequence[Hashable]
+    keys : Sequence[Key]
         A list of the outputs from ``__dask_keys__`` from all
         collections sharing the same ``__dask_optimize__`` method.
     **kwargs : Any
@@ -263,7 +271,7 @@ class DaskCollection(Protocol):
             If True [default], the graph is optimized before
             computation. Otherwise the graph is run as is. This can be
             useful for debugging.
-        kwargs
+        kwargs :
             Extra keywords to forward to the scheduler function.
 
         Returns
diff --git a/docs/source/custom-collections.rst b/docs/source/custom-collections.rst
index df6acf2f4..0dd9ea875 100644
--- a/docs/source/custom-collections.rst
+++ b/docs/source/custom-collections.rst
@@ -430,7 +430,7 @@ Demonstrating this class:
 
     # The output keys for this graph.
     # The first element of each tuple must be the same across the whole collection;
-    # the remainder are arbitrary, unique hashables
+    # the remainder are arbitrary, unique str, bytes, int, or floats
     >>> keys = [("x", "k1"), ("x", 1), ("x", 2), ("x", 3)]
 
     >>> x = Tuple(dsk, keys)
diff --git a/docs/source/spec.rst b/docs/source/spec.rst
index 9811c6c0c..43ca94c72 100644
--- a/docs/source/spec.rst
+++ b/docs/source/spec.rst
@@ -20,7 +20,7 @@ A **Dask graph** is a dictionary mapping **keys** to **computations**:
     'w': (sum, ['x', 'y', 'z']),
     'v': [(sum, ['w', 'z']), 2]}
 
-A **key** is any hashable value that is not a **task**:
+A **key** is a str, bytes, int, float, or tuple thereof:
 
 .. code-block:: python
 


Test patch:
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index c6ea22e2b..92ebe3f00 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -55,7 +55,7 @@ jobs:
 
     steps:
       - name: Checkout source
-        uses: actions/checkout@v3.6.0
+        uses: actions/checkout@v4.0.0
         with:
           fetch-depth: 0  # Needed by codecov.io
       - name: Setup Java
diff --git a/dask/tests/test_base.py b/dask/tests/test_base.py
index 171189a14..776c360ad 100644
--- a/dask/tests/test_base.py
+++ b/dask/tests/test_base.py
@@ -53,6 +53,10 @@ np = import_or_none("numpy")
 sp = import_or_none("scipy.sparse")
 pd = import_or_none("pandas")
 
+# Arbitrary dask keys
+h1 = (1.2, "foo", (3,))
+h2 = "h2"
+
 
 def f1(a, b, c=1):
     pass
@@ -743,9 +747,6 @@ def test_get_collection_names():
 
     assert get_collection_names(DummyCollection({}, [])) == set()
 
-    # Arbitrary hashables
-    h1 = object()
-    h2 = object()
     # __dask_keys__() returns a nested list
     assert get_collection_names(
         DummyCollection(
@@ -756,10 +757,6 @@ def test_get_collection_names():
 
 
 def test_get_name_from_key():
-    # Arbitrary hashables
-    h1 = object()
-    h2 = object()
-
     assert get_name_from_key("foo") == "foo"
     assert get_name_from_key("foo-123"), "foo-123"
     assert get_name_from_key(("foo-123", h1, h2)) == "foo-123"
@@ -776,8 +773,6 @@ def test_replace_name_in_keys():
     assert replace_name_in_key("foo", {"bar": "baz"}) == "foo"
     assert replace_name_in_key("foo", {"foo": "bar", "baz": "asd"}) == "bar"
     assert replace_name_in_key("foo-123", {"foo-123": "bar-456"}) == "bar-456"
-    h1 = object()  # Arbitrary hashables
-    h2 = object()
     assert replace_name_in_key(("foo-123", h1, h2), {"foo-123": "bar"}) == (
         "bar",
         h1,
@@ -830,10 +825,6 @@ class Tuple(DaskMethodsMixin):
 
 
 def test_custom_collection():
-    # Arbitrary hashables
-    h1 = object()
-    h2 = object()
-
     dsk = {("x", h1): 1, ("x", h2): 2}
     dsk2 = {("y", h1): (add, ("x", h1), ("x", h2)), ("y", h2): (add, ("y", h1), 1)}
     dsk2.update(dsk)
@@ -1341,10 +1332,7 @@ def test_persist_delayed():
     assert x3.compute() == xx.compute()
 
 
-some_hashable = object()
-
-
-@pytest.mark.parametrize("key", ["a", ("a-123", some_hashable)])
+@pytest.mark.parametrize("key", ["a", ("a-123", h1)])
 def test_persist_delayed_custom_key(key):
     d = Delayed(key, {key: "b", "b": 1})
     assert d.compute() == 1
@@ -1360,7 +1348,7 @@ def test_persist_delayed_custom_key(key):
         ("a", {}, "a"),
         ("a", {"c": "d"}, "a"),
         ("a", {"a": "b"}, "b"),
-        (("a-123", some_hashable), {"a-123": "b-123"}, ("b-123", some_hashable)),
+        (("a-123", h1), {"a-123": "b-123"}, ("b-123", h1)),
     ],
 )
 def test_persist_delayed_rename(key, rename, new_key):
@@ -1596,14 +1584,13 @@ def test_optimizations_ctd():
 
 
 def test_clone_key():
-    h = object()  # arbitrary hashable
     assert clone_key("inc-1-2-3", 123) == "inc-4dfeea2f9300e67a75f30bf7d6182ea4"
     assert clone_key("x", 123) == "x-dc2b8d1c184c72c19faa81c797f8c6b0"
     assert clone_key("x", 456) == "x-b76f061b547b00d18b9c7a18ccc47e2d"
     assert clone_key(("x", 1), 456) == ("x-b76f061b547b00d18b9c7a18ccc47e2d", 1)
-    assert clone_key(("sum-1-2-3", h, 1), 123) == (
+    assert clone_key(("sum-1-2-3", h1, 1), 123) == (
         "sum-1efd41f02035dc802f4ebb9995d07e9d",
-        h,
+        h1,
         1,
     )
     with pytest.raises(TypeError):
diff --git a/dask/tests/test_core.py b/dask/tests/test_core.py
index 8ec5fbccd..a3d72380d 100644
--- a/dask/tests/test_core.py
+++ b/dask/tests/test_core.py
@@ -12,11 +12,14 @@ from dask.core import (
     get_deps,
     getcycle,
     has_tasks,
+    ishashable,
+    iskey,
     istask,
     literal,
     preorder_traversal,
     quote,
     subs,
+    validate_key,
 )
 from dask.utils_test import GetFunctionTestMixin, add, inc
 
@@ -32,6 +35,48 @@ def contains(a, b):
     return all(a.get(k) == v for k, v in b.items())
 
 
+def test_ishashable():
+    class C:
+        pass
+
+    assert ishashable("x")
+    assert ishashable(1)
+    assert ishashable(C())
+    assert ishashable((1, 2))
+    assert not ishashable([1, 2])
+    assert not ishashable({1: 2})
+
+
+def test_iskey():
+    class C:
+        pass
+
+    assert iskey("x")
+    assert iskey(1)
+    assert not iskey(C())  # Custom hashables can't be dask keys
+    assert not iskey((C(),))
+    assert iskey((1, 2))
+    assert iskey(())
+    assert iskey(("x",))
+    assert not iskey([1, 2])
+    assert not iskey({1, 2})
+    assert not iskey({1: 2})
+
+
+def test_iskey_numpy_types():
+    np = pytest.importorskip("numpy")
+    one = np.int64(1)
+    assert not iskey(one)
+    assert not iskey(("foo", one))
+
+
+def test_validate_key():
+    validate_key(1)
+    validate_key(("x", 1))
+    with pytest.raises(TypeError):
+        validate_key(["x", 1])
+
+
 def test_istask():
     assert istask((inc, 1))
     assert not istask(1)
@@ -230,18 +275,9 @@ def test_subs_with_surprisingly_friendly_eq():
         assert subs(df, "x", 1) is df
 
 
-def test_subs_unexpected_hashable_key():
-    class UnexpectedButHashable:
-        def __init__(self):
-            self.name = "a"
-
-        def __hash__(self):
-            return hash(self.name)
-
-        def __eq__(self, other):
-            return isinstance(other, UnexpectedButHashable)
-
-    assert subs((id, UnexpectedButHashable()), UnexpectedButHashable(), 1) == (id, 1)
+def test_subs_arbitrary_key():
+    key = (1.2, "foo", (3,))
+    assert subs((id, key), key, 1) == (id, 1)
 
 
 def test_quote():
diff --git a/dask/tests/test_graph_manipulation.py b/dask/tests/test_graph_manipulation.py
index fbd92205b..f3a6291b4 100644
--- a/dask/tests/test_graph_manipulation.py
+++ b/dask/tests/test_graph_manipulation.py
@@ -75,9 +75,9 @@ def assert_did_not_materialize(cloned, orig):
             assert not cv.is_materialized()
 
 
-# Generic hashables
-h1 = object()
-h2 = object()
+# Generic valid keys
+h1 = (1.2, "foo", ())
+h2 = "h2"
 
 
 def collections_with_node_counters():
diff --git a/dask/tests/test_typing.py b/dask/tests/test_typing.py
index 45802cf24..bec841ba5 100644
--- a/dask/tests/test_typing.py
+++ b/dask/tests/test_typing.py
@@ -12,7 +12,10 @@ from dask.context import globalmethod
 from dask.delayed import Delayed, delayed
 from dask.typing import (
     DaskCollection,
+    Graph,
     HLGDaskCollection,
+    Key,
+    NestedKeys,
     PostComputeCallable,
     PostPersistCallable,
 )
@@ -33,11 +36,11 @@ def finalize(x: Sequence[Any]) -> Any:
     return x[0]
 
 
-def get1(dsk: Mapping, keys: Sequence[Hashable] | Hashable, **kwargs: Any) -> Any:
+def get1(dsk: Mapping, keys: Sequence[Key] | Key, **kwargs: Any) -> Any:
     return dask.threaded.get(dsk, keys, **kwargs)
 
 
-def get2(dsk: Mapping, keys: Sequence[Hashable] | Hashable, **kwargs: Any) -> Any:
+def get2(dsk: Mapping, keys: Sequence[Key] | Key, **kwargs: Any) -> Any:
     return dask.get(dsk, keys, **kwargs)
 
 
@@ -45,10 +48,10 @@ class Inheriting(DaskCollection):
     def __init__(self, based_on: DaskCollection) -> None:
         self.based_on = based_on
 
-    def __dask_graph__(self) -> Mapping:
+    def __dask_graph__(self) -> Graph:
         return self.based_on.__dask_graph__()
 
-    def __dask_keys__(self) -> list[Hashable]:
+    def __dask_keys__(self) -> NestedKeys:
         return self.based_on.__dask_keys__()
 
     def __dask_postcompute__(self) -> tuple[PostComputeCallable, tuple]:
@@ -94,13 +97,13 @@ class HLGCollection(DaskMethodsMixin):
     def __init__(self, based_on: HLGDaskCollection) -> None:
         self.based_on = based_on
 
-    def __dask_graph__(self) -> Mapping:
+    def __dask_graph__(self) -> Graph:
         return self.based_on.__dask_graph__()
 
     def __dask_layers__(self) -> Sequence[str]:
         return self.based_on.__dask_layers__()
 
-    def __dask_keys__(self) -> list[Hashable]:
+    def __dask_keys__(self) -> NestedKeys:
         return self.based_on.__dask_keys__()
 
     def __dask_postcompute__(self) -> tuple[PostComputeCallable, tuple]:
@@ -125,10 +128,10 @@ class NotHLGCollection(DaskMethodsMixin):
     def __init__(self, based_on: DaskCollection) -> None:
         self.based_on = based_on
 
-    def __dask_graph__(self) -> Mapping:
+    def __dask_graph__(self) -> Graph:
         return self.based_on.__dask_graph__()
 
-    def __dask_keys__(self) -> list[Hashable]:
+    def __dask_keys__(self) -> NestedKeys:
         return self.based_on.__dask_keys__()
 
     def __dask_postcompute__(self) -> tuple[PostComputeCallable, tuple]: