[serialization] Revert exception allowlist and list subclass APPEND from #189353 (#200956)

https://github.com/pytorch/pytorch/pull/189353 globally allowlisted all builtin exceptions on torch.distributed import, so every torch.load(weights_only=True) and get_unsafe_globals_in_checkpoint in the process treats them as safe, not just object collectives. It also let APPEND/APPENDS run on allowlisted list subclasses, which calls their overridden append/extend during load; this was only needed for traceback.StackSummary, which is no longer allowlisted.

Fix: remove the c10d exception allowlist, the Python 2 exceptions NAME_MAPPING that supported it, and the list subclass APPEND/APPENDS path. The object collective weights_only path is unchanged.

Test plan:

```
python test/test_serialization.py -k weights_only
python test/distributed/test_c10d_object_collectives.py
python test/distributed/checkpoint/test_utils.py
```

Authored with an AI assistant (Claude).
Pull Request resolved: https://github.com/pytorch/pytorch/pull/200956
Approved by: https://github.com/kapilsh, https://github.com/atalman
ghstack dependencies: #200797

test/test_serialization.py

@@ -919,6 +919,10 @@ class serialization_method:
 
 Point = namedtuple('Point', ['x', 'y'])
 
+class ListSubclass(list):
+    pass
+
+
 class ClassThatUsesBuildInstruction:
     def __init__(self, num):
         self.num = num
@@ -1373,6 +1377,15 @@ class TestSerialization(TestCase, SerializationMixin):
                 with self.assertRaisesRegex(pickle.UnpicklingError, error_msg):
                     torch.load(f, weights_only=True)
 
+    def test_weights_only_list_subclass_append(self):
+        # APPEND/APPENDS would call overridden methods of list subclasses.
+        with BytesIOContext() as f:
+            torch.save(ListSubclass([1, 2]), f)
+            f.seek(0)
+            with safe_globals([ListSubclass]):
+                with self.assertRaisesRegex(pickle.UnpicklingError, "Can only extend lists"):
+                    torch.load(f, weights_only=True)
+
     @parametrize("module", ["traceback", "linecache"])
     def test_weights_only_traceback_blocklist(self, module):
         # traceback.FrameSummary and linecache.getline read arbitrary files.

torch/_utils.py

@@ -1,5 +1,4 @@
 # mypy: allow-untyped-defs
-import _compat_pickle
 import copyreg
 import functools
 import importlib
@@ -1172,16 +1171,6 @@ NAME_MAPPING = {
     ("UserDict", "UserDict"): ("collections", "UserDict"),
 }
 
-# Protocol 2 pickle (torch.save's default) maps builtin exceptions to the
-# Python 2 "exceptions" module via REVERSE_NAME_MAPPING; map them back so
-# allowlisted exception types resolve under their builtins.* names.
-NAME_MAPPING.update(
-    {
-        ("exceptions", name): ("builtins", name)
-        for name in _compat_pickle.PYTHON2_EXCEPTIONS
-    }
-)
-
 
 def _chunk_or_narrow_cat(
     tensor: "torch.Tensor",

torch/_weights_only_unpickler.py

@@ -2,7 +2,7 @@
 # Unpickler restricted to loading only state dicts
 # Restrict constructing types to a list defined in _get_allowed_globals()
 # Restrict BUILD operation to `Tensor`, `Parameter` and `OrderedDict` types only
-# Restrict APPEND/APPENDS to `list` (and allowlisted list subclasses)
+# Restrict APPEND/APPENDS to `list`
 # In `GLOBALS` operation do not do class lookup by name, but rather rely on dictionary
 # defined by `_get_allowed_globals()` method, that contains:
 # - torch types (Storage, dtypes, Tensor, `torch.Size`),
@@ -467,12 +467,18 @@ class Unpickler:
             elif key[0] == APPEND[0]:
                 item = self.stack.pop()
                 list_obj = self.stack[-1]
-                self._check_append_target(list_obj)
+                if type(list_obj) is not list:
+                    raise UnpicklingError(
+                        f"Can only append to lists, but got {type(list_obj)}"
+                    )
                 list_obj.append(item)
             elif key[0] == APPENDS[0]:
                 items = self.pop_mark()
                 list_obj = self.stack[-1]
-                self._check_append_target(list_obj)
+                if type(list_obj) is not list:
+                    raise UnpicklingError(
+                        f"Can only extend lists, but got {type(list_obj)}"
+                    )
                 list_obj.extend(items)
             elif key[0] == SETITEM[0]:
                 (v, k) = (self.stack.pop(), self.stack.pop())
@@ -595,15 +601,6 @@ class Unpickler:
                 f"collections.Counter, but got {type(self.stack[-1])}"
             )
 
-    def _check_append_target(self, list_obj):
-        # list subclasses allowlisted via add_safe_globals can be appended to
-        # (e.g. traceback.StackSummary); plain list is always allowed
-        if type(list_obj) is not list and not (
-            isinstance(list_obj, list)
-            and type(list_obj) in _get_user_allowed_globals().values()
-        ):
-            raise UnpicklingError(f"Can only append to lists, but got {type(list_obj)}")
-
     def persistent_load(self, pid):
         raise UnpicklingError("unsupported persistent id encountered")
 

torch/distributed/distributed_c10d.py

@@ -1,7 +1,6 @@
 # mypy: allow-untyped-defs
 """Distributed Collective Communication (c10d)."""
 
-import builtins
 import collections.abc
 import contextlib
 import ctypes
@@ -319,17 +318,6 @@ def _pg_options_to_hints(pg_options: object) -> dict[str, str] | None:
 _pickler = pickle.Pickler
 _unpickler = pickle.Unpickler
 
-# Object collectives can serialize via torch.save/torch.load(weights_only=True)
-# by passing weights_only=True. Exceptions are commonly transmitted for error
-# propagation, so allowlist them.
-torch.serialization.add_safe_globals(
-    [
-        t
-        for t in vars(builtins).values()
-        if isinstance(t, type) and issubclass(t, BaseException)
-    ]
-)
-
 GroupName = NewType("GroupName", str)