Make the is_picklable function more robust (#17270)

Author: leng-yue

src/lightning/pytorch/CHANGELOG.md

@@ -4,6 +4,13 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
+## [2.0.2] - 2023-04-12
+
+### Fixed
+
+- Make the `is_picklable` function more robust ([#17270](https://github.com/Lightning-AI/lightning/pull/17270))
+
+
 ## [2.0.1] - 2023-03-30
 
 ### Changed
@@ -12,6 +19,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Generalized `Optimizer` validation to accommodate both FSDP 1.x and 2.x ([#16733](https://github.com/Lightning-AI/lightning/pull/16733))
 - Disable `torch.inference_mode` with `torch.compile` in PyTorch 2.0 ([#17215](https://github.com/Lightning-AI/lightning/pull/17215))
 
+
 ### Fixed
 
 - Fixed issue where pickling the module instance would fail with a DataLoader error ([#17130](https://github.com/Lightning-AI/lightning/pull/17130))

src/lightning/pytorch/utilities/parsing.py

@@ -32,7 +32,7 @@ def is_picklable(obj: object) -> bool:
     try:
         pickle.dumps(obj)
         return True
-    except (pickle.PickleError, AttributeError, RuntimeError):
+    except (pickle.PickleError, AttributeError, RuntimeError, TypeError):
         return False
 
 

tests/tests_pytorch/utilities/test_parsing.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import inspect
+import threading
 
 import pytest
 from torch.jit import ScriptModule
@@ -169,7 +170,7 @@ class UnpicklableClass:
         pass
 
     true_cases = [None, True, 123, "str", (123, "str"), max]
-    false_cases = [unpicklable_function, UnpicklableClass, ScriptModule()]
+    false_cases = [unpicklable_function, UnpicklableClass, ScriptModule(), threading.Lock()]
 
     for case in true_cases:
         assert is_picklable(case) is True