[Test] Fix sync buffer hook

Author: HAOCHENYE

tests/test_hooks/test_sync_buffers_hook.py

@@ -1,15 +1,14 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-import os
 from unittest.mock import MagicMock
 
 import torch
 import torch.distributed as torch_dist
 import torch.nn as nn
+from torch.testing._internal.common_distributed import DistributedTestBase
 
 from mmengine.dist import all_gather
 from mmengine.hooks import SyncBuffersHook
 from mmengine.registry import MODELS
-from mmengine.testing._internal import MultiProcessTestCase
 from mmengine.testing.runner_test_case import RunnerTestCase, ToyModel
 
 
@@ -23,22 +22,14 @@ def __init__(self, data_preprocessor=None):
     def init_weights(self):
         for buffer in self.buffers():
             buffer.fill_(
-                torch.tensor(int(os.environ['RANK']), dtype=torch.float32))
+                torch.tensor(torch_dist.get_rank(), dtype=torch.float32))
         return super().init_weights()
 
 
-class TestSyncBuffersHook(MultiProcessTestCase, RunnerTestCase):
-
-    def setUp(self) -> None:
-        super().setUp()
-        self._spawn_processes()
-
-    def prepare_subprocess(self):
-        MODELS.register_module(module=ToyModuleWithNorm, force=True)
-        super(MultiProcessTestCase, self).setUp()
+class TestSyncBuffersHook(DistributedTestBase, RunnerTestCase):
 
     def test_sync_buffers_hook(self):
-        self.setup_dist_env()
+        self.create_pg('cuda')
         runner = MagicMock()
         runner.model = ToyModuleWithNorm()
         runner.model.init_weights()
@@ -53,9 +44,12 @@ def test_sync_buffers_hook(self):
         for buffer in runner.model.buffers():
             buffer1, buffer2 = all_gather(buffer)
             self.assertTrue(torch.allclose(buffer1, buffer2))
+        torch_dist.destroy_process_group()
 
     def test_with_runner(self):
-        self.setup_dist_env()
+        MODELS.register_module(module=ToyModuleWithNorm, force=True)
+        self.create_pg('cuda')
+        RunnerTestCase.setUp(self)
         cfg = self.epoch_based_cfg
         cfg.model = dict(type='ToyModuleWithNorm')
         cfg.launch = 'pytorch'
@@ -67,8 +61,6 @@ def test_with_runner(self):
             buffer1, buffer2 = all_gather(buffer)
             self.assertTrue(torch.allclose(buffer1, buffer2))
 
-    def setup_dist_env(self):
-        super().setup_dist_env()
-        os.environ['RANK'] = str(self.rank)
-        torch_dist.init_process_group(
-            backend='gloo', rank=self.rank, world_size=self.world_size)
+    @property
+    def world_size(self) -> int:
+        return 2