Optimize rollout-related ut execution time

tests/ray/test_evaluator.py

@@ -22,10 +22,12 @@ class TestEvaluator(unittest.TestCase):
     @classmethod
     def setUpClass(cls) -> None:
         os.environ["XTUNER_USE_FA3"] = "1"
-
+        os.environ["LMD_SKIP_WARMUP"] = "1"
+
     @classmethod
     def tearDownClass(cls) -> None:
         del os.environ["XTUNER_USE_FA3"]
+        del os.environ["LMD_SKIP_WARMUP"]
 
     def init_config(self):
         self.resources_cfg = AcceleratorResourcesConfig(

tests/ray/test_mock_rollout.py

@@ -1,4 +1,5 @@
 import os
+import asyncio
 import unittest
 import ray
 from transformers import AutoTokenizer
@@ -47,7 +48,6 @@ class TestMockRollout(unittest.TestCase):
     def setUpClass(cls):
         os.environ["XTUNER_USE_FA3"] = "1"
 
-
     @classmethod
     def tearDownClass(cls):
         del os.environ["XTUNER_USE_FA3"]
@@ -61,15 +61,6 @@ def setUp(self):
         self.max_retry_times = 3
         self.temp_dir = tempfile.TemporaryDirectory()
         self.worker_log_dir = os.path.join(self.temp_dir.name, "work_dirs")
-
-        self.resources_cfg = AcceleratorResourcesConfig(
-            accelerator=resource_map[torch.accelerator.current_accelerator().type],
-            num_workers=8,
-            num_cpus_per_worker=8,
-            cpu_memory_per_worker=16 * 1024**3,  # 16 GB
-        )
-        self.pg = AutoAcceleratorWorkers.build_placement_group(self.resources_cfg)
-
         self.rollout_cfg = RolloutConfig(
             env="test_mock_rollout",
             model_path=MODEL_PATH,
@@ -108,33 +99,40 @@ def tearDown(self):
         ray.shutdown()
         self.temp_dir.cleanup()
 
-    def _run_mock_test(self, mock_controller_cls, error_name: str):
-        rollout_controller = mock_controller_cls.remote(self.rollout_cfg, self.pg)
-        self.test_env = SingleTurnEnvironment.remote("env", self.pg, self.rollout_cfg, rollout_controller=rollout_controller)
+    async def _run_mock_test(self, mock_controller_cls, error_name, pg):
+        rollout_controller = mock_controller_cls.remote(self.rollout_cfg, pg)
+        self.test_env = SingleTurnEnvironment.remote("env", pg, self.rollout_cfg, rollout_controller=rollout_controller)
         self.test_dataflow = DataFlow.remote("dataflow", self.dataflow_cfg, self.replay_buffer_cfg, self.test_env)
 
-        completed_rollouts = ray.get(self.test_dataflow.run.remote(num=3))["data_groups"]
-        status = ray.get(self.test_dataflow.get_replaybuffer_status.remote())
+        result = await self.test_dataflow.run.remote(num=3)
+        completed_rollouts = result["data_groups"]
+        status = await self.test_dataflow.get_replaybuffer_status.remote()
         self.assertEqual(len(completed_rollouts), 0, f"[{error_name}] Expected no rollouts to complete successfully.")
         self.assertEqual(status["remain_completed_samples_count"], 0, f"[{error_name}] Completed count in buffer should be 0.")
         self.assertEqual(status["remain_aborted_samples_count"], 0, f"[{error_name}] Expected no rollouts to be interrupted.")
-        ray.get(self.test_env.shutdown.remote())
+        await self.test_env.shutdown.remote()
 
     @unittest.skipIf(os.environ.get("XTUNER_USE_LMDEPLOY", "0") == "0", "lmdeploy backend is not enabled")
-    def test_rollout_with_timeout_mock(self):
-        self._run_mock_test(MockTimeoutRolloutController, "timeout")
+    def test_parallel_mock_rollout(self):
+        async def run_parallel():
+            res_cfg_small = AcceleratorResourcesConfig(
+                accelerator=resource_map[torch.accelerator.current_accelerator().type],
+                num_workers=2,
+                num_cpus_per_worker=2,
+            )
+
+            pgs = [AutoAcceleratorWorkers.build_placement_group(res_cfg_small, name=f"pg_{i}") for i in range(4)]
+            await asyncio.gather(*[pg.ready() for pg in pgs])
 
-    @unittest.skipIf(os.environ.get("XTUNER_USE_LMDEPLOY", "0") == "0", "lmdeploy backend is not enabled")  
-    def test_rollout_with_request_error_mock(self):
-        self._run_mock_test(MockRequestErrorRolloutController, "request error")
-
-    @unittest.skipIf(os.environ.get("XTUNER_USE_LMDEPLOY", "0") == "0", "lmdeploy backend is not enabled")
-    def test_rollout_with_client_error_mock(self):
-        self._run_mock_test(MockClientErrorRolloutController, "client error")
-
-    @unittest.skipIf(os.environ.get("XTUNER_USE_LMDEPLOY", "0") == "0", "lmdeploy backend is not enabled")
-    def test_rollout_with_server_error_mock(self):
-        self._run_mock_test(MockServerErrorRolloutController, "server error")
+            tasks = [
+                self._run_mock_test(MockTimeoutRolloutController, "timeout", pgs[0]),
+                self._run_mock_test(MockRequestErrorRolloutController, "request_error", pgs[1]),
+                self._run_mock_test(MockClientErrorRolloutController, "client_error", pgs[2]),
+                self._run_mock_test(MockServerErrorRolloutController, "server_error", pgs[3]),
+            ]
+            await asyncio.gather(*tasks)
+
+        asyncio.run(run_parallel())
 
 if __name__ == "__main__":
     unittest.main()

tests/ray/test_rl_trainer.py

@@ -38,11 +38,13 @@ class TestRLTrainer(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
         os.environ["XTUNER_USE_FA3"] = "1"
+        os.environ["LMD_SKIP_WARMUP"] = "1"
 
     @classmethod
     def tearDownClass(cls):
         del os.environ["XTUNER_USE_FA3"]
-
+        del os.environ["LMD_SKIP_WARMUP"]
+
     def init_traine_worker_config(self, train_optimizer_steps, pack_max_length):
         model_cfg = get_model_config_from_hf(Path(MODEL_PATH))
         optim_cfg = AdamWConfig(lr=1e-6, betas=(0.9, 0.999), max_grad_norm=1.0, weight_decay=0.1, foreach=False)