[tune] cleanup error messaging/diagnose_serialization helper (#10210)

python/ray/tune/utils/util.py

@@ -1,5 +1,6 @@
 import copy
 import logging
+import inspect
 import threading
 import time
 from collections import defaultdict, deque, Mapping, Sequence
@@ -269,6 +270,92 @@ def _from_pinnable(obj):
     return obj[0]
 
 
+def diagnose_serialization(trainable):
+    """Utility for detecting accidentally-scoped objects.
+
+    Args:
+        trainable (cls | func): The trainable object passed to
+            tune.run(trainable).
+
+    Returns:
+        bool | set of unserializable objects.
+
+    Example:
+
+    .. code-block::
+
+        import threading
+        # this is not serializable
+        e = threading.Event()
+
+        def test():
+            print(e)
+
+        diagnose_serialization(test)
+        # should help identify that 'e' should be moved into
+        # the `test` scope.
+
+        # correct implementation
+        def test():
+            e = threading.Event()
+            print(e)
+
+        assert diagnose_serialization(test) is True
+
+    """
+    from ray.tune.registry import register_trainable, check_serializability
+
+    def check_variables(objects, failure_set, printer):
+        for var_name, variable in objects.items():
+            msg = None
+            try:
+                check_serializability(var_name, variable)
+                status = "PASSED"
+            except Exception as e:
+                status = "FAILED"
+                msg = f"{e.__class__.__name__}: {str(e)}"
+                failure_set.add(var_name)
+            printer(f"{str(variable)}[name='{var_name}'']... {status}")
+            if msg:
+                printer(msg)
+
+    print(f"Trying to serialize {trainable}...")
+    try:
+        register_trainable("__test:" + str(trainable), trainable, warn=False)
+        print("Serialization succeeded!")
+        return True
+    except Exception as e:
+        print(f"Serialization failed: {e}")
+
+    print("Inspecting the scope of the trainable by running "
+          f"`inspect.getclosurevars({str(trainable)})`...")
+    closure = inspect.getclosurevars(trainable)
+    failure_set = set()
+    if closure.globals:
+        print(f"Detected {len(closure.globals)} global variables. "
+              "Checking serializability...")
+        check_variables(closure.globals, failure_set,
+                        lambda s: print("   " + s))
+
+    if closure.nonlocals:
+        print(f"Detected {len(closure.nonlocals)} nonlocal variables. "
+              "Checking serializability...")
+        check_variables(closure.nonlocals, failure_set,
+                        lambda s: print("   " + s))
+
+    if not failure_set:
+        print("Nothing was found to have failed the diagnostic test, though "
+              "serialization did not succeed. Feel free to raise an "
+              "issue on github.")
+        return failure_set
+    else:
+        print(f"Variable(s) {failure_set} was found to be non-serializable. "
+              "Consider either removing the instantiation/imports "
+              "of these objects or moving them into the scope of "
+              "the trainable. ")
+        return failure_set
+
+
 def validate_save_restore(trainable_cls,
                           config=None,
                           num_gpus=0,