Repository: dask/dask
Pull request: Add config for handling obj normalization with missing __dask_tokenize__
- [x] Closes #6555
  - [x] Added `tokenize.allow-random` config, accepting `True` or `False`.
- [x] Tests added / passed
  - [x] Added `test_normalize_object` to `dask/tests/test_base.py`.
- [x] Passes `black dask` / `flake8 dask`

Problem statement:
Flag for raising an error when `normalize_object` doesn't find `__dask_tokenize__`
I ran into a bit of a nasty surprised when I realized some of my Delayed objects resulted in a different token when I called `tokenize()`. This is because in my data I had some objects which didn't have `__dask_tokenize__`. I fixed those errors, but I'm afraid they might come up again silently. It would be amazing if I could set a flag that would raise an error any time that I tried to tokenize something which contained anything that couldn't be tokenized deterministically. I'm thinking something in `dask.base.normalize_object` that raises an error instead of returning `uuid.uuid4().hex`.

Thanks!

Hints from issue discussion:
Thanks for raising an issue @STOpandthink! I can imagine this was tedious to track down. Adding a config option to raise an error or warn when falling back to a random token (which is `False` by default) seems like a nice addition for certain use cases. Is this something you're interested in working on? (no obligation though)
Hi @jrbourbeau, could you please elaborate on the "config option" part?
I suppose, based on the wording, you're referring to something other than a flag that's passed to a particular function or set of functions.
Dask has it's own [configuration system](https://docs.dask.org/en/latest/configuration.html) that let's user customize some of Dask's behavior (e.g. logging verbosity). The proposal here is to add a new configuration option that controls if users are informed when [`dask.base.normalize_object`](https://github.com/dask/dask/blob/1767fedd2f0662ba2fd6a9caebf2d3595ca8046c/dask/base.py#L838) isn't able to return a deterministic result and falls back to returning a [random `uuid`](https://github.com/dask/dask/blob/1767fedd2f0662ba2fd6a9caebf2d3595ca8046c/dask/base.py#L842).

Commit message: Add config for handling obj normalization with missing __dask_tokenize__ (#6555)

Patch to review:
diff --git a/dask/base.py b/dask/base.py
index 1d5ad89b8bd..4ea837f451d 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -839,7 +839,42 @@ def normalize_object(o):
     method = getattr(o, "__dask_tokenize__", None)
     if method is not None:
         return method()
-    return normalize_function(o) if callable(o) else uuid.uuid4().hex
+
+    if callable(o):
+        return normalize_function(o)
+
+    handle_missing_dask_tokenize_method = config.get(
+        "normalize_token.handle_missing_dask_tokenize_method",
+        "ignore",
+    )
+
+    if handle_missing_dask_tokenize_method == "ignore":
+        return uuid.uuid4().hex
+
+    missing_dask_tokenize_method_msg = (
+        f"Object {str(o)} doesn't have a `__dask_tokenize__` method and "
+        "cannot be normalized deterministically."
+    )
+
+    if handle_missing_dask_tokenize_method == "warn":
+        normalized = uuid.uuid4().hex
+
+        import warnings
+
+        warnings.warn(
+            missing_dask_tokenize_method_msg
+            + f' Normalizing non-deterministically to "{normalized}".'
+        )
+
+        return normalized
+
+    if handle_missing_dask_tokenize_method == "raise":
+        raise TypeError(missing_dask_tokenize_method_msg)
+
+    raise ValueError(
+        f"Invalid value for config `normalize_token.handle_missing_dask_tokenize_method`. "
+        f'Expected one of ["ignore", "warn", "raise"]; got "{handle_missing_dask_tokenize_method}".'
+    )
 
 
 function_cache = {}
diff --git a/dask/tests/test_base.py b/dask/tests/test_base.py
index f1681b41228..16d03e7e5d3 100644
--- a/dask/tests/test_base.py
+++ b/dask/tests/test_base.py
@@ -8,6 +8,8 @@
 
 from tlz import merge, partial, compose, curry
 
+from unittest import mock
+
 import dask
 import dask.bag as db
 from dask import delayed
@@ -224,6 +226,86 @@ def test_normalize_base():
         assert normalize_token(i) is i
 
 
+def test_normalize_object():
+    import warnings
+
+    warnings.simplefilter("error")
+
+    o = object()
+
+    # Baseline config behavior: no explicitly set config value.
+    # Default behavior is equivalent to "ignore", i.e., proceed with non-
+    # deterministic normalization.
+    with mock.patch("uuid.uuid4") as mock_uuid4:
+        mock_uuid4.return_value = mock.Mock(hex="123")
+        assert normalize_token(o) == "123"
+
+    # Explicitly set config value to "ignore".
+    with dask.config.set(
+        {"normalize_token.handle_missing_dask_tokenize_method": "ignore"}
+    ):
+        with mock.patch("uuid.uuid4") as mock_uuid4:
+            mock_uuid4.return_value = mock.Mock(hex="123")
+            assert normalize_token(o) == "123"
+
+    # Explicitly set config value to "warn".
+    with dask.config.set(
+        {"normalize_token.handle_missing_dask_tokenize_method": "warn"}
+    ):
+        with mock.patch("uuid.uuid4") as mock_uuid4:
+            mock_uuid4.return_value = mock.Mock(hex="123")
+
+            with pytest.raises(UserWarning) as excinfo:
+                assert normalize_token(o) == "123"
+
+            assert "123" in str(excinfo.value)
+
+    # Explicitly set config value to "raise".
+    with dask.config.set(
+        {"normalize_token.handle_missing_dask_tokenize_method": "raise"}
+    ):
+        with mock.patch("uuid.uuid4") as mock_uuid4:
+            mock_uuid4.return_value = mock.Mock(hex="123")
+
+            with pytest.raises(TypeError) as excinfo:
+                normalize_token(o)
+
+            assert "__dask_tokenize__" in str(excinfo.value)
+
+    # Explicitly set config value to some unsupported value.
+    with dask.config.set(
+        {"normalize_token.handle_missing_dask_tokenize_method": "unsupported"}
+    ):
+        with pytest.raises(ValueError) as excinfo:
+            normalize_token(o)
+
+        assert "unsupported" in str(excinfo.value)
+
+    # Assert that objects, which have a corresponding `__dask_tokenize__` are
+    # unaffected by the config, and get normalized as per that method.
+    o2 = mock.Mock()
+    o2.__dask_tokenize__ = mock.Mock(return_value="abc")
+
+    for h in ["ignore", "warn", "raise"]:
+        with dask.config.set(
+            {"normalize_token.handle_missing_dask_tokenize_method": h}
+        ):
+            assert normalize_token(o2) == "abc"
+
+    # Assert that handling callables is unaffected by the config as it's logically distinct.
+    # Only one type of callables is tested here.
+    c = lambda: None
+
+    with mock.patch("cloudpickle.dumps") as mock_cloudpickle_dumps:
+        mock_cloudpickle_dumps.return_value = "678"
+
+        for h in ["ignore", "warn", "raise"]:
+            with dask.config.set(
+                {"normalize_token.handle_missing_dask_tokenize_method": h}
+            ):
+                assert normalize_token(c) == "678"
+
+
 @pytest.mark.skipif("not pd")
 def test_tokenize_pandas():
     a = pd.DataFrame({"x": [1, 2, 3], "y": ["4", "asd", None]}, index=[1, 2, 3])


Code review comments:
diff_hunk: @@ -839,7 +839,42 @@ def normalize_object(o):
     method = getattr(o, "__dask_tokenize__", None)
     if method is not None:
         return method()
-    return normalize_function(o) if callable(o) else uuid.uuid4().hex
+
+    if callable(o):
+        return normalize_function(o)
+
+    handle_missing_dask_tokenize_method = config.get(
+        "normalize_token.handle_missing_dask_tokenize_method",
+        "ignore",
+    )
+
+    if handle_missing_dask_tokenize_method == "ignore":
+        return uuid.uuid4().hex
+
+    missing_dask_tokenize_method_msg = (
+        f"Object {str(o)} doesn't have a `__dask_tokenize__` method and "
+        "cannot be normalized deterministically."
+    )
+
+    if handle_missing_dask_tokenize_method == "warn":
+        normalized = uuid.uuid4().hex
+
+        import warnings
+
+        warnings.warn(
path: dask/base.py
text: @author:
This implicitly displays a `UserWarning`. Would a different type of warning be more suitable here (e.g., `RuntimeWarning`)?

diff_hunk: @@ -839,7 +839,42 @@ def normalize_object(o):
     method = getattr(o, "__dask_tokenize__", None)
     if method is not None:
         return method()
-    return normalize_function(o) if callable(o) else uuid.uuid4().hex
+
+    if callable(o):
+        return normalize_function(o)
+
+    handle_missing_dask_tokenize_method = config.get(
+        "normalize_token.handle_missing_dask_tokenize_method",
+        "ignore",
+    )
+
+    if handle_missing_dask_tokenize_method == "ignore":
+        return uuid.uuid4().hex
+
+    missing_dask_tokenize_method_msg = (
+        f"Object {str(o)} doesn't have a `__dask_tokenize__` method and "
+        "cannot be normalized deterministically."
+    )
+
+    if handle_missing_dask_tokenize_method == "warn":
+        normalized = uuid.uuid4().hex
+
+        import warnings
+
+        warnings.warn(
+            missing_dask_tokenize_method_msg
+            + f' Normalizing non-deterministically to "{normalized}".'
+        )
+
+        return normalized
+
+    if handle_missing_dask_tokenize_method == "raise":
+        raise TypeError(missing_dask_tokenize_method_msg)
path: dask/base.py
text: @author:
I've opted for raising a `TypeError` here. Please, let me know if you believe some other type would be more suitable for this circumstance (e.g., `ValueError`).