Fix get_unsafe_globals_in_checkpoint to account for user allowed globals per docstring (pytorch#140738)

bugfix: this function did not account for the user allowed globals :( Differential Revision: [D65960696](https://our.internmc.facebook.com/intern/diff/D65960696) Pull Request resolved: pytorch#140738 Approved by: https://github.com/malfet

Author: mikaylagawarecki

test/test_serialization.py

@@ -4422,6 +4422,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

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):