Allow user to manually pass module.name associated with global in {add}_safe_global (pytorch#142153)

Fixes pytorch#142144

A global x is saved in checkpoint as `GLOBAL x.__module__ x.__name__`. So , after allowlisting a GLOBAL it is expected to match any GLOBAL instruction of the form `GLOBAL x.__module__ x.__name__`  but there are edge cases when for the same API from the same module, what `__module__` gives changes between versions which prevents users from allowlisting the global.

In this case, in numpy < 2.1

```
torch.save("bla", np_array)
# checkpoint has GLOBAL "np.core.multiarray" "_reconstruct"
```
In np version 2.1

```
with safe_globals([np.core.multiarray._reconstruct]):
    torch.load("bla")
```
np.core.multiarray._reconstruct.__module__ gives "np._core.multiarray" (note the extra _ before core) and see what was done [here](https://github.com/numpy/numpy/blob/main/numpy/core/multiarray.py)

Since the dictionary to access safe globals is keyed on "{foo.__module__}.{foo.__name__}", __module__, __name__ will no longer match that in the checkpoint so "np.core.multiarray._reconstruct" can no longer be properly allowlisted (instead np._core.multiarray._reconstruct is a key in the dict).

We allow `add_safe_globals/safe_globals` to optionally take tuples of (global, str of module.name) to workaround such (odd/edge case) situations.

Pull Request resolved: pytorch#142153
Approved by: https://github.com/albanD

Author: mikaylagawarecki

test/test_cpp_extensions_open_device_registration.py

@@ -539,10 +539,6 @@ def test_open_device_tensorlist_type_fallback(self):
         np.__version__ < "1.25",
         "versions < 1.25 serialize dtypes differently from how it's serialized in data_legacy_numpy",
     )
-    @unittest.skipIf(
-        np.__version__ >= "2.1",
-        "weights_only failure on numpy >= 2.1",
-    )
     def test_open_device_numpy_serialization(self):
         """
         This tests the legacy _rebuild_device_tensor_from_numpy serialization path
@@ -601,7 +597,9 @@ def test_open_device_numpy_serialization(self):
 
         with safe_globals(
             [
-                np.core.multiarray._reconstruct,
+                (np.core.multiarray._reconstruct, "numpy.core.multiarray._reconstruct")
+                if np.__version__ >= "2.1"
+                else np.core.multiarray._reconstruct,
                 np.ndarray,
                 np.dtype,
                 _codecs.encode,

torch/_weights_only_unpickler.py

@@ -68,7 +68,7 @@
 )
 from struct import unpack
 from sys import maxsize
-from typing import Any, Callable, Dict, List, Set, Tuple
+from typing import Any, Callable, Dict, List, Set, Tuple, Union
 
 import torch
 from torch._utils import IMPORT_MAPPING, NAME_MAPPING
@@ -83,15 +83,15 @@
     "nt",
 ]
 
-_marked_safe_globals_set: Set[Any] = set()
+_marked_safe_globals_set: Set[Union[Callable, Tuple[Callable, str]]] = set()
 
 
-def _add_safe_globals(safe_globals: List[Any]):
+def _add_safe_globals(safe_globals: List[Union[Callable, Tuple[Callable, str]]]):
     global _marked_safe_globals_set
     _marked_safe_globals_set = _marked_safe_globals_set.union(set(safe_globals))
 
 
-def _get_safe_globals() -> List[Any]:
+def _get_safe_globals() -> List[Union[Callable, Tuple[Callable, str]]]:
     global _marked_safe_globals_set
     return list(_marked_safe_globals_set)
 
@@ -101,13 +101,15 @@ def _clear_safe_globals():
     _marked_safe_globals_set = set()
 
 
-def _remove_safe_globals(globals_to_remove: List[Any]):
+def _remove_safe_globals(
+    globals_to_remove: List[Union[Callable, Tuple[Callable, str]]],
+):
     global _marked_safe_globals_set
     _marked_safe_globals_set = _marked_safe_globals_set - set(globals_to_remove)
 
 
 class _safe_globals:
-    def __init__(self, safe_globals: List[Any]):
+    def __init__(self, safe_globals: List[Union[Callable, Tuple[Callable, str]]]):
         self.safe_globals = safe_globals
 
     def __enter__(self):
@@ -127,8 +129,20 @@ def __exit__(self, type, value, tb):
 def _get_user_allowed_globals():
     rc: Dict[str, Any] = {}
     for f in _marked_safe_globals_set:
-        module, name = f.__module__, f.__name__
-        rc[f"{module}.{name}"] = f
+        if isinstance(f, tuple):
+            if len(f) != 2:
+                raise ValueError(
+                    f"Expected tuple of length 2 (global, str of callable full path), but got tuple of length: {len(f)}"
+                )
+            if type(f[1]) is not str:
+                raise TypeError(
+                    f"Expected second item in tuple to be str of callable full path, but got: {type(f[1])}"
+                )
+            f, name = f
+            rc[name] = f
+        else:
+            module, name = f.__module__, f.__name__
+            rc[f"{module}.{name}"] = f
     return rc
 
 

torch/serialization.py

@@ -261,21 +261,29 @@ def clear_safe_globals() -> None:
     _weights_only_unpickler._clear_safe_globals()
 
 
-def get_safe_globals() -> List[Any]:
+def get_safe_globals() -> List[Union[Callable, Tuple[Callable, str]]]:
     """
     Returns the list of user-added globals that are safe for ``weights_only`` load.
     """
     return _weights_only_unpickler._get_safe_globals()
 
 
-def add_safe_globals(safe_globals: List[Any]) -> None:
+def add_safe_globals(safe_globals: List[Union[Callable, Tuple[Callable, str]]]) -> None:
     """
     Marks the given globals as safe for ``weights_only`` load. For example, functions
     added to this list can be called during unpickling, classes could be instantiated
     and have state set.
 
+    Each item in the list can either be a function/class or a tuple of the form
+    (function/class, string) where string is the full path of the function/class.
+
+    Within the serialized format, each function is identified with its full
+    path as ``{__module__}.{__name__}``. When calling this API, you can provide this
+    full path that should match the one in the checkpoint otherwise the default
+    ``{fn.__module__}.{fn.__name__}`` will be used.
+
     Args:
-        safe_globals (List[Any]): list of globals to mark as safe
+        safe_globals (List[Union[Callable, Tuple[Callable, str]]]): list of globals to mark as safe
 
     Example:
         >>> # xdoctest: +SKIP("Can't torch.save(t, ...) as doctest thinks MyTensor is defined on torch.serialization")