Fix broken tests in torchsnapshot

Summary:
 ---
The test `test_read_sharded_tensor` was failing with `FileNotFoundError` because each of the 4 distributed worker processes was creating its own independent temporary directory via `tempfile.TemporaryDirectory()`. When rank 0 created a snapshot in its temporary directory, other ranks would try to read from their own (different) temporary directories, which didn't contain the snapshot files.

Fixed by implementing a proper distributed temporary directory pattern:
1. Rank 0 creates the temp directory using `tempfile.mkdtemp()`
2. The path is broadcast to all ranks using `dist.broadcast_object_list()`
3. All ranks use the same shared path for snapshot operations
4. Proper cleanup with a `dist.barrier()` before rank 0 removes the temp directory to prevent race conditions

 ---
test_dict_store was failing due to mismatched regex. fixing the regex fixed test

Reviewed By: galrotem

Differential Revision: D89585177

fbshipit-source-id: 3ee6c7c283d8ddc2cd539843931a052fa51c1933

Author: Vinti Pandey

tests/test_dist_store.py

@@ -78,13 +78,13 @@ def _test_linear_barrier_timeout() -> None:
         # should timeout in .arrive() and other non-leader ranks should timeout
         # in .depart().
         if rank == 0:
-            with tc.assertRaisesRegex(RuntimeError, "Socket Timeout"):
+            with tc.assertRaisesRegex(RuntimeError, "timeout"):
                 barrier.arrive(timeout=timedelta(seconds=5))
         elif rank == 1:
             pass
         else:
             barrier.arrive(timeout=timedelta(seconds=5))
-            with tc.assertRaisesRegex(RuntimeError, "Socket Timeout"):
+            with tc.assertRaisesRegex(RuntimeError, "timeout"):
                 barrier.depart(timeout=timedelta(seconds=5))
 
         dist.barrier()
@@ -102,7 +102,7 @@ def _test_linear_barrier_timeout() -> None:
             pass
         else:
             barrier.arrive(timeout=timedelta(seconds=5))
-            with tc.assertRaisesRegex(RuntimeError, "Socket Timeout"):
+            with tc.assertRaisesRegex(RuntimeError, "timeout"):
                 barrier.depart(timeout=timedelta(seconds=5))
 
         dist.barrier()

tests/test_read_object.py

@@ -64,24 +64,40 @@ def _test_read_sharded_tensor() -> None:
         for foo_shard, bar_shard in zip(foo.local_shards(), bar.local_shards()):
             tc.assertFalse(torch.allclose(foo_shard.tensor, bar_shard.tensor))
 
-        with tempfile.TemporaryDirectory() as path:
+        # Create temp directory on rank 0 and broadcast path to all ranks
+        if dist.get_rank() == 0:
+            path = tempfile.mkdtemp()
+        else:
+            path = ""
+        path_list = [path]
+        dist.broadcast_object_list(path_list, src=0)
+        path = path_list[0]
+
+        try:
             snapshot = torchsnapshot.Snapshot.take(
                 path=path, app_state={"state": torchsnapshot.StateDict(foo=foo)}
             )
             snapshot.read_object("0/state/foo", obj_out=bar)
             baz = snapshot.read_object("0/state/foo")
 
-        for foo_shard, bar_shard in zip(foo.local_shards(), bar.local_shards()):
-            tc.assertTrue(torch.allclose(foo_shard.tensor, bar_shard.tensor))
+            for foo_shard, bar_shard in zip(foo.local_shards(), bar.local_shards()):
+                tc.assertTrue(torch.allclose(foo_shard.tensor, bar_shard.tensor))
 
-        tc.assertEqual(baz.shape, torch.Size([20_000, 128]))
+            tc.assertEqual(baz.shape, torch.Size([20_000, 128]))
 
-        gathered_foo_tensor = torch.empty(20_000, 128)
-        if dist.get_rank() == 0:
-            foo.gather(dst=0, out=gathered_foo_tensor)
-            tc.assertTrue(torch.allclose(baz, gathered_foo_tensor))
-        else:
-            foo.gather(dst=0, out=None)
+            gathered_foo_tensor = torch.empty(20_000, 128)
+            if dist.get_rank() == 0:
+                foo.gather(dst=0, out=gathered_foo_tensor)
+                tc.assertTrue(torch.allclose(baz, gathered_foo_tensor))
+            else:
+                foo.gather(dst=0, out=None)
+        finally:
+            # Clean up the temp directory on rank 0
+            dist.barrier()
+            if dist.get_rank() == 0:
+                import shutil
+
+                shutil.rmtree(path, ignore_errors=True)
 
     def test_read_sharded_tensor(self) -> None:
         lc = get_pet_launch_config(nproc=4)