Dict memoization sorting before normalizing keys is unsafe

parsl/dataflow/memoization.py

@@ -5,7 +5,7 @@
 import pickle
 from parsl.dataflow.taskrecord import TaskRecord
 
-from typing import Dict, Any, List, Optional, TYPE_CHECKING
+from typing import Dict, Any, List, Optional, TYPE_CHECKING  # avoid circular imports
 
 if TYPE_CHECKING:
     from parsl import DataFlowKernel  # import loop at runtime - needed for typechecking - TODO turn into "if typing:"
@@ -14,7 +14,9 @@
 
 import types
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger(__name__)  # logger named name for logging purposes
+
+# memoization function with a single dispatch decorator
 
 
 @singledispatch
@@ -49,6 +51,8 @@ def id_for_memo(obj: object, output_ref: bool = False) -> bytes:
     logger.error("id_for_memo attempted on unknown type {}".format(type(obj)))
     raise ValueError("unknown type for memoization: {}".format(type(obj)))
 
+# type specific implementations - handle how each type should be serialized for memoization
+
 
 @id_for_memo.register(str)
 @id_for_memo.register(int)
@@ -94,10 +98,13 @@ def id_for_memo_dict(denormalized_dict: dict, output_ref: bool = False) -> bytes
     if type(denormalized_dict) is not dict:
         raise ValueError("id_for_memo_dict cannot work on subclasses of dict")
 
-    keys = sorted(denormalized_dict)
+    # keys = sorted(denormalized_dict)  Line that sirosen commented on
+    # Proposed solution was to normalize the keys and then sort them
+    keymap = {id_for_memo(k): k for k in denormalized_dict}
+    normed_keys = sorted(keymap.values())
 
     normalized_list = []
-    for k in keys:
+    for k in normed_keys:
         normalized_list.append(id_for_memo(k))
         normalized_list.append(id_for_memo(denormalized_dict[k], output_ref=output_ref))
     return pickle.dumps(normalized_list)

parsl/tests/test_python_apps/test_memoize_3.py

@@ -0,0 +1,29 @@
+import pytest
+import enum
+
+# Define an enum - collection of related consonants
+
+
+class Foo(enum.Enum):
+    x = enum.auto()
+    y = enum.auto()
+
+
+# Test function demonstrating the issue with unstable sorting when keys
+# are hashable but not comparable.
+
+
+def test_unstable_sorting():
+    # Functions
+    def foo():
+        return 1
+
+    def bar():
+        return 2
+
+    # Dictionary with problematic keys
+    d = {foo: 1, bar: 2}
+
+    # Sort the dictionary, it should raise a TypeError
+    with pytest.raises(TypeError):
+        sorted(d)