[tune/core] serialization debugging utility (#12142)

Co-authored-by: SangBin Cho <rkooo567@gmail.com>
Co-authored-by: Kai Fricke <kai@anyscale.com>

python/ray/tune/experiment.py

@@ -286,21 +286,12 @@ class Experiment:
             try:
                 register_trainable(name, run_object)
             except (TypeError, PicklingError) as e:
-                msg = (
-                    f"{str(e)}. The trainable ({str(run_object)}) could not "
-                    "be serialized, which is needed for parallel execution. "
-                    "To diagnose the issue, try the following:\n\n"
-                    "\t- Run `tune.utils.diagnose_serialization(trainable)` "
-                    "to check if non-serializable variables are captured "
-                    "in scope.\n"
-                    "\t- Try reproducing the issue by calling "
-                    "`pickle.dumps(trainable)`.\n"
-                    "\t- If the error is typing-related, try removing "
-                    "the type annotations and try again.\n\n"
-                    "If you have any suggestions on how to improve "
-                    "this error message, please reach out to the "
-                    "Ray developers on github.com/ray-project/ray/issues/")
-                raise type(e)(msg) from None
+                extra_msg = (f"Other options: "
+                             "\n-Try reproducing the issue by calling "
+                             "`pickle.dumps(trainable)`. "
+                             "\n-If the error is typing-related, try removing "
+                             "the type annotations and try again.")
+                raise type(e)(str(e) + " " + extra_msg) from None
             return name
         else:
             raise TuneError("Improper 'run' - not string nor trainable.")

python/ray/tune/registry.py

@@ -119,7 +119,7 @@ class _Registry:
             from ray.tune import TuneError
             raise TuneError("Unknown category {} not among {}".format(
                 category, KNOWN_CATEGORIES))
-        self._to_flush[(category, key)] = pickle.dumps(value)
+        self._to_flush[(category, key)] = pickle.dumps_debug(value)
         if _internal_kv_initialized():
             self.flush_values()
 

python/ray/tune/tests/test_api.py

@@ -1038,8 +1038,6 @@ class TrainableFunctionApiTest(unittest.TestCase):
         self.assertLessEqual(status.get("PENDING", 0), 1)
 
     def testMetricCheckingEndToEnd(self):
-        from ray import tune
-
         def train(config):
             tune.report(val=4, second=8)
 
@@ -1130,6 +1128,50 @@ class TrainableFunctionApiTest(unittest.TestCase):
             self.assertFalse(found)
 
 
+class SerializabilityTest(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        ray.init(local_mode=True)
+
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown()
+
+    def tearDown(self):
+        if "RAY_PICKLE_VERBOSE_DEBUG" in os.environ:
+            del os.environ["RAY_PICKLE_VERBOSE_DEBUG"]
+
+    def testNotRaisesNonserializable(self):
+        import threading
+        lock = threading.Lock()
+
+        def train(config):
+            print(lock)
+            tune.report(val=4, second=8)
+
+        with self.assertRaisesRegex(TypeError, "RAY_PICKLE_VERBOSE_DEBUG"):
+            # The trial runner raises a ValueError, but the experiment fails
+            # with a TuneError
+            tune.run(train, metric="acc")
+
+    def testRaisesNonserializable(self):
+        os.environ["RAY_PICKLE_VERBOSE_DEBUG"] = "1"
+        import threading
+        lock = threading.Lock()
+
+        def train(config):
+            print(lock)
+            tune.report(val=4, second=8)
+
+        with self.assertRaises(TypeError) as cm:
+            # The trial runner raises a ValueError, but the experiment fails
+            # with a TuneError
+            tune.run(train, metric="acc")
+        msg = cm.exception.args[0]
+        assert "RAY_PICKLE_VERBOSE_DEBUG" not in msg
+        assert "thread.lock" in msg
+
+
 class ShimCreationTest(unittest.TestCase):
     def testCreateScheduler(self):
         kwargs = {"metric": "metric_foo", "mode": "min"}

python/ray/tune/utils/serialization.py

@@ -4,6 +4,7 @@ import types
 
 from ray import cloudpickle as cloudpickle
 from ray.utils import binary_to_hex, hex_to_binary
+from ray.util.debug import log_once
 
 logger = logging.getLogger(__name__)
 
@@ -15,7 +16,8 @@ class TuneFunctionEncoder(json.JSONEncoder):
         try:
             return super(TuneFunctionEncoder, self).default(obj)
         except Exception:
-            logger.debug("Unable to encode. Falling back to cloudpickle.")
+            if log_once(f"tune_func_encode:{str(obj)}"):
+                logger.debug("Unable to encode. Falling back to cloudpickle.")
             return self._to_cloudpickle(obj)
 
     def _to_cloudpickle(self, obj):