Added ScriptFunction pkl exception for issue pytorch#61210 pytorch#61381 (pytorch#67076)

Summary:
Fixes pytorch#61381, pytorch#61210

Pull Request resolved: pytorch#67076

Reviewed By: jbschlosser

Differential Revision: D32908175

Pulled By: suo

fbshipit-source-id: f6e175793243dc96cde5e44022d92f2623b934eb

Co-authored-by: LucaStubbe <[email protected]>
Co-authored-by: Kanon Tromp <[email protected]>

Author: 3 people

test/test_jit.py

@@ -383,6 +383,15 @@ def fn(x):
         files = zipfile.ZipFile(buf).filelist
         self.assertTrue(any(['archive/constants.pkl' == f.filename for f in files]))
 
+    def test_script_fn_pkl(self):
+        with self.assertRaisesRegex(pickle.PickleError, "ScriptFunction cannot be pickled"):
+
+            @torch.jit.script
+            def fn(x: torch.Tensor) -> torch.Tensor:
+                return x
+
+            pkl_fn = pickle.dumps(fn, protocol=0)
+
     def test_restore_device(self):
         class M(torch.jit.ScriptModule):
             def __init__(self, cpu_device_str):

torch/jit/_script.py

@@ -56,6 +56,13 @@
 """
 set_module(ScriptFunction, "torch.jit")
 
+# Throws an error if a jit function is pickled.
+# Helps to avoid Python crashes for Python versions 3.9.5 + when protocol 0 or 1 is given as an argument.
+def _reduce(cls):
+    raise pickle.PickleError("ScriptFunction cannot be pickled")
+
+ScriptFunction.__reduce__ = _reduce  # type: ignore[assignment]
+
 
 if _enabled:
     Attribute = collections.namedtuple("Attribute", ["value", "type"])