Fix tests in test/test_serialization that were failing if run individually (pytorch#141300)

pytorch#140739 and pytorch#140740 made it such that `get_safe_globals` no longer return an empty List by default

This caused some tests that check the content of `get_safe_globals` to fail, in particular when run individually (they didn't fail in test suite as other tests ran before them called `clear_safe_globals`) but will fail when tests are run individually [T208186010](https://www.internalfb.com/intern/tasks/?t=208186010)

test_safe_globals_for_weights_only
test_safe_globals_context_manager_weights_only

This PR fixes that and also makes most tests calling `clear_safe_globals` use the `safe_globals` context manager rather than try: finally

Pull Request resolved: pytorch#141300
Approved by: https://github.com/awgu

Author: mikaylagawarecki

test/test_serialization.py

@@ -1099,12 +1099,9 @@ def __reduce__(self):
             # Safe load should assert
             with self.assertRaisesRegex(pickle.UnpicklingError, "Unsupported global: GLOBAL builtins.print"):
                 torch.load(f, weights_only=True)
-            try:
-                torch.serialization.add_safe_globals([print])
+            with torch.serialization.safe_globals([print]):
                 f.seek(0)
                 torch.load(f, weights_only=True)
-            finally:
-                torch.serialization.clear_safe_globals()
 
     def test_weights_only_safe_globals_newobj(self):
         # This will use NEWOBJ
@@ -1116,12 +1113,9 @@ def test_weights_only_safe_globals_newobj(self):
                                         "GLOBAL __main__.Point was not an allowed global by default"):
                 torch.load(f, weights_only=True)
             f.seek(0)
-            try:
-                torch.serialization.add_safe_globals([Point])
+            with torch.serialization.safe_globals([Point]):
                 loaded_p = torch.load(f, weights_only=True)
                 self.assertEqual(loaded_p, p)
-            finally:
-                torch.serialization.clear_safe_globals()
 
     def test_weights_only_safe_globals_build(self):
         counter = 0
@@ -1138,21 +1132,20 @@ def fake_set_state(obj, *args):
                                         "GLOBAL __main__.ClassThatUsesBuildInstruction was not an allowed global by default"):
                 torch.load(f, weights_only=True)
             try:
-                torch.serialization.add_safe_globals([ClassThatUsesBuildInstruction])
-                # Test dict update path
-                f.seek(0)
-                loaded_c = torch.load(f, weights_only=True)
-                self.assertEqual(loaded_c.num, 2)
-                self.assertEqual(loaded_c.foo, 'bar')
-                # Test setstate path
-                ClassThatUsesBuildInstruction.__setstate__ = fake_set_state
-                f.seek(0)
-                loaded_c = torch.load(f, weights_only=True)
-                self.assertEqual(loaded_c.num, 2)
-                self.assertEqual(counter, 1)
-                self.assertFalse(hasattr(loaded_c, 'foo'))
+                with torch.serialization.safe_globals([ClassThatUsesBuildInstruction]):
+                    # Test dict update path
+                    f.seek(0)
+                    loaded_c = torch.load(f, weights_only=True)
+                    self.assertEqual(loaded_c.num, 2)
+                    self.assertEqual(loaded_c.foo, 'bar')
+                    # Test setstate path
+                    ClassThatUsesBuildInstruction.__setstate__ = fake_set_state
+                    f.seek(0)
+                    loaded_c = torch.load(f, weights_only=True)
+                    self.assertEqual(loaded_c.num, 2)
+                    self.assertEqual(counter, 1)
+                    self.assertFalse(hasattr(loaded_c, 'foo'))
             finally:
-                torch.serialization.clear_safe_globals()
                 ClassThatUsesBuildInstruction.__setstate__ = None
 
     @parametrize("slots", ['some', 'all'])
@@ -4629,10 +4622,12 @@ def test_safe_globals_for_weights_only(self):
                 sd = torch.load(f, weights_only=True)
 
             # Loading tensor subclass should work if the class is marked safe
+            safe_globals_before = torch.serialization.get_safe_globals()
             f.seek(0)
             try:
                 torch.serialization.add_safe_globals([TwoTensor])
-                self.assertTrue(torch.serialization.get_safe_globals() == [TwoTensor])
+                expected_safe_globals = set(safe_globals_before + [TwoTensor])
+                self.assertEqual(set(torch.serialization.get_safe_globals()), expected_safe_globals)
                 sd = torch.load(f, weights_only=True)
                 self.assertEqual(sd['t'], t)
                 self.assertEqual(sd['p'], p)
@@ -4645,6 +4640,7 @@ def test_safe_globals_for_weights_only(self):
                     torch.load(f, weights_only=True)
             finally:
                 torch.serialization.clear_safe_globals()
+                torch.serialization.add_safe_globals(safe_globals_before)
 
     def test_safe_globals_context_manager_weights_only(self):
         '''
@@ -4654,20 +4650,23 @@ def test_safe_globals_context_manager_weights_only(self):
         p = torch.nn.Parameter(t)
         sd = OrderedDict([('t', t), ('p', p)])
 
+        safe_globals_before = torch.serialization.get_safe_globals()
         try:
             torch.serialization.add_safe_globals([TestEmptySubclass])
             with tempfile.NamedTemporaryFile() as f:
                 torch.save(sd, f)
                 with safe_globals([TwoTensor]):
                     f.seek(0)
                     torch.load(f, weights_only=True)
-                self.assertTrue(torch.serialization.get_safe_globals() == [TestEmptySubclass])
+                expected_safe_globals = set(safe_globals_before + [TestEmptySubclass])
+                self.assertEqual(set(torch.serialization.get_safe_globals()), expected_safe_globals)
                 f.seek(0)
                 with self.assertRaisesRegex(pickle.UnpicklingError,
                                             "Unsupported global: GLOBAL torch.testing._internal.two_tensor.TwoTensor"):
                     torch.load(f, weights_only=True)
         finally:
             torch.serialization.clear_safe_globals()
+            torch.serialization.add_safe_globals(safe_globals_before)
 
     def test_sets_are_loadable_with_weights_only(self):
         s = {1, 2, 3}