show_pickle/model_dump: Handle invalid UTF-8 in pickles (pytorch#57661)

Summary:
Pull Request resolved: pytorch#57661

Thie Pickle "specification" (pickletools.py) states that the argument to
a BINUNICODE opcode must be UTF-8 encoded.  However, if a PyTorch custom
class returns a non-UTF-8 std::string from its pickle method the
libtorch Pickler will write it to the output pickle without complaining.
Python's _Unpickler (the Python implementation of Unpickler) always
throws an exception when trying to deserialize these invalid pickles.

We still want to be able to dump these pickle files.  Update
DumpUnpickler to create its own opcode dispatch table (initialized as a
clone of the _Unpickler dispatch table) and patch in a custom function
for the BINUNICODE op.  We try to emulate the default behavior, but any
UnicodeDecodeError is caught and replaced with a dummy object.  This
could violate the assumptions of a user that expects a str in that
position, so we disable this behavior by default.

Update model_dump to recognize this special object and allow it to be
rendered.

Test Plan: Dumped and viewed a model with an invalid string in an object state.

Reviewed By: malfet

Differential Revision: D28531392

Pulled By: dreiss

fbshipit-source-id: ab5aea20975a0ef53ef52a880deaa2c5a626e4a2

Author: dreiss

torch/utils/model_dump/__init__.py

@@ -162,6 +162,17 @@ def hierarchical_pickle(data):
             assert isinstance(name, str)
             # Just forget that it was a device and return the name.
             return name
+        if typename == "builtin.UnicodeDecodeError":
+            assert data.state is None
+            msg, = data.args
+            assert isinstance(msg, str)
+            # Hack: Pretend this is a module so we don't need custom serialization.
+            # Hack: Wrap the message in a tuple so it looks like a nice state object.
+            # TODO: Undo at least that second hack.  We should support string states.
+            return {
+                "__module_type__": typename,
+                "state": hierarchical_pickle((msg,)),
+            }
         raise Exception(f"Can't prepare fake object of type for JS: {typename}")
     raise Exception(f"Can't prepare data of type for JS: {type(data)}")
 
@@ -210,7 +221,7 @@ def get_model_info(
         version = zf.read(path_prefix + "/version").decode("utf-8").strip()
 
         with zf.open(path_prefix + "/data.pkl") as handle:
-            raw_model_data = torch.utils.show_pickle.DumpUnpickler(handle).load()
+            raw_model_data = torch.utils.show_pickle.DumpUnpickler(handle, catch_invalid_utf8=True).load()
             model_data = hierarchical_pickle(raw_model_data)
 
         # Intern strings that are likely to be re-used.
@@ -287,7 +298,7 @@ def ist(s):
                 # TODO: handle errors here and just ignore the file?
                 # NOTE: For a lot of these files (like bytecode),
                 # we could get away with just unpickling, but this should be safer.
-                obj = torch.utils.show_pickle.DumpUnpickler(handle).load()
+                obj = torch.utils.show_pickle.DumpUnpickler(handle, catch_invalid_utf8=True).load()
             buf = io.StringIO()
             pprint.pprint(obj, buf)
             contents = buf.getvalue()

torch/utils/model_dump/code.js

@@ -276,6 +276,11 @@ class ModelData extends Component {
         }
       } else if (mstate.__tuple_values__) {
         parts.push(html`<br/><${ModelData} prefix="" indent=${new_indent} data=${mstate} />`);
+      } else if (mstate.__module_type__) {
+        // We normally wouldn't have the state of a module be another module,
+        // but we use "modules" to encode special values (like Unicode decode
+        // errors) that might be valid states.  Just go with it.
+        parts.push(html`<br/><${ModelData} prefix="" indent=${new_indent} data=${mstate} />`);
       } else {
         throw new Error("Bad module state");
       }

torch/utils/show_pickle.py

@@ -1,10 +1,11 @@
 #!/usr/bin/env python3
 import sys
 import pickle
+import struct
 import pprint
 import zipfile
 import fnmatch
-from typing import IO, BinaryIO, Union
+from typing import Any, IO, BinaryIO, Union
 
 
 class FakeObject(object):
@@ -58,12 +59,42 @@ def fake_new(self, *args):
 
 
 class DumpUnpickler(pickle._Unpickler):  # type: ignore[name-defined]
+    def __init__(
+            self,
+            file,
+            *,
+            catch_invalid_utf8=False,
+            **kwargs):
+        super().__init__(file, **kwargs)
+        self.catch_invalid_utf8 = catch_invalid_utf8
+
     def find_class(self, module, name):
         return FakeClass(module, name)
 
     def persistent_load(self, pid):
         return FakeObject("pers", "obj", (pid,))
 
+    dispatch = dict(pickle._Unpickler.dispatch)  # type: ignore[attr-defined]
+
+    # Custom objects in TorchScript are able to return invalid UTF-8 strings
+    # from their pickle (__getstate__) functions.  Install a custom loader
+    # for strings that catches the decode exception and replaces it with
+    # a sentinel object.
+    def load_binunicode(self):
+        strlen, = struct.unpack("<I", self.read(4))
+        if strlen > sys.maxsize:
+            raise Exception("String too long.")
+        str_bytes = self.read(strlen)
+        obj: Any
+        try:
+            obj = str(str_bytes, "utf-8", "surrogatepass")
+        except UnicodeDecodeError as exn:
+            if not self.catch_invalid_utf8:
+                raise
+            obj = FakeObject("builtin", "UnicodeDecodeError", (str(exn),))
+        self.append(obj)
+    dispatch[pickle.BINUNICODE[0]] = load_binunicode
+
     @classmethod
     def dump(cls, in_stream, out_stream):
         value = cls(in_stream).load()