Fix get_unsafe_globals_in_checkpoint to account for user allowed globals per docstring

mikaylagawarecki · mikaylagawarecki · commit dddee440fafc · 2024-11-14T11:37:28.000-08:00
[ghstack-poisoned]
diff --git a/test/test_serialization.py b/test/test_serialization.py
@@ -4413,6 +4413,10 @@ def test_get_unsafe_globals_in_checkpoint(self):
             unsafe_globals = torch.serialization.get_unsafe_globals_in_checkpoint(f)
             self.assertEqual(set(unsafe_globals), expected_unsafe_global_strs)
             f.seek(0)
+            with torch.serialization.safe_globals([TwoTensor]):
+                unsafe_globals = torch.serialization.get_unsafe_globals_in_checkpoint(f)
+                self.assertEqual(set(unsafe_globals), set())
+            f.seek(0)
             try:
                 old_get_allowed_globals = torch._weights_only_unpickler._get_allowed_globals
                 torch._weights_only_unpickler._get_allowed_globals = lambda: dict()  # noqa: PIE807
diff --git a/torch/serialization.py b/torch/serialization.py
@@ -342,7 +342,13 @@ def get_unsafe_globals_in_checkpoint(f: FILE_LIKE) -> List[str]:
     Returns:
         A list of strings of pickle GLOBALs in the checkpoint that are not allowlisted for ``weights_only``.
     """
-    safe_global_strings = set(_weights_only_unpickler._get_allowed_globals().keys())
+    default_safe_globals_strings = set(
+        _weights_only_unpickler._get_allowed_globals().keys()
+    )
+    user_safe_global_strings = set(
+        _weights_only_unpickler._get_user_allowed_globals().keys()
+    )
+    safe_global_strings = default_safe_globals_strings.union(user_safe_global_strings)
 
     with _open_file_like(f, "rb") as opened_file:
         if not _is_zipfile(opened_file):
