Add sparse tensors constructed via legacy constructor to _sparse_tensors_to_validate (pytorch#147759)

This is a redo of pytorch#147408 which added validation at the end of the legacy constructor calls.

The reason why I didn't land that was because in `legacy_load`, constructor would be called before storages of indices/values are set. So the tensor would not actually be validated.

Technically, torch.sparse.{Foo}Tensor should not even be called by our rebuild process since afaict this was the first PR that added support for sparse tensor serialization pytorch#27062 and it already uses `_rebuild_sparse_tensor` (which would add the rebuilt tensor to the list to validate), but torch.sparse.FooTensor is allowlisted

This PR adds tensors constructed as such to the list to validate at the end of torch.load.

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

Author: mikaylagawarecki

test/test_serialization.py

@@ -443,6 +443,40 @@ def __reduce_ex__(self, proto):
                         "size is inconsistent with indices"):
                     y = torch.load(f, weights_only=weights_only)
 
+    def test_serialization_sparse_invalid_legacy_ctor(self):
+        # This is set in test class setup but would not be check when running user code
+        prev_invariant_check_enabled = torch.sparse.check_sparse_tensor_invariants.is_enabled()
+        try:
+            torch.sparse.check_sparse_tensor_invariants.disable()
+            x = torch.zeros(3, 3)
+            x[1][1] = 1
+            x = x.to_sparse()
+            x_legacy_ctor = torch.sparse.FloatTensor(x.indices(), x.values())
+
+            # technically legacy ctor will still always be rebuilt with _rebuild_sparse_tensor
+            # this is to test that legacy ctor in data.pkl will be validated by weights_only unpickler
+            class LegacyCtorSerializationSpoofer:
+                def __init__(self, tensor):
+                    self.tensor = tensor
+
+                def __reduce_ex__(self, proto):
+                    indices = self.tensor._indices()
+                    indices[0][0] = 3
+                    return (torch.sparse.FloatTensor, (indices, self.tensor._values(), self.tensor.size()))
+
+            with tempfile.NamedTemporaryFile() as f:
+                sd = {"spoofed_legacy_ctor": LegacyCtorSerializationSpoofer(x_legacy_ctor)}
+                torch.save(sd, f)
+                for weights_only in (True,):
+                    f.seek(0)
+                    with self.assertRaisesRegex(
+                            RuntimeError,
+                            "size is inconsistent with indices"):
+                        y = torch.load(f, weights_only=weights_only)
+        finally:
+            if prev_invariant_check_enabled:
+                torch.sparse.check_sparse_tensor_invariants.enable()
+
     def _test_serialization_sparse_compressed_invalid(self,
                                                       conversion,
                                                       get_compressed_indices,

torch/_weights_only_unpickler.py

@@ -71,7 +71,7 @@
 from typing import Any, Callable, Union
 
 import torch
-from torch._utils import IMPORT_MAPPING, NAME_MAPPING
+from torch._utils import _sparse_tensors_to_validate, IMPORT_MAPPING, NAME_MAPPING
 
 
 # modules in this list are never allowed, even if the user attempts to allowlist
@@ -387,7 +387,10 @@ def load(self):
                     cls in _get_user_allowed_globals().values()
                     or cls in _get_allowed_globals().values()
                 ):
-                    self.append(cls.__new__(cls, *args))
+                    result = cls.__new__(cls, *args)
+                    if cls in torch._tensor_classes and "sparse" in cls.__module__:
+                        _sparse_tensors_to_validate.append(result)
+                    self.append(result)
                 else:
                     raise UnpicklingError(
                         "Can only create new object for nn.Parameter or classes allowlisted "
@@ -403,7 +406,10 @@ def load(self):
                     raise UnpicklingError(
                         f"Trying to call reduce for unrecognized function {func}"
                     )
-                self.stack[-1] = func(*args)
+                result = func(*args)
+                if func in torch._tensor_classes and "sparse" in func.__module__:
+                    _sparse_tensors_to_validate.append(result)
+                self.stack[-1] = result
             elif key[0] == BUILD[0]:
                 state = self.stack.pop()
                 inst = self.stack[-1]