Bug fix in multi stage resume (#462)

Author: chenyushuo

tests/trainer/trainer_test.py

@@ -331,12 +331,15 @@ def test_trainer(self, mock_load):
             ),
         ]
         self.config.check_and_update()
+        old_taskset_path = self.config.stages[1].buffer.explorer_input.taskset.path
+        self.config.stages[1].buffer.explorer_input.taskset.path = "/invalid/path"
 
-        mock_load.return_value = self.config
+        mock_load.return_value = deepcopy(self.config)
 
-        run(config_path="dummy.yaml")
+        with self.assertRaises(Exception):
+            run(config_path="dummy.yaml")
 
-        stage_configs = [cfg.check_and_update() for cfg in self.config]
+        stage_configs = [cfg.check_and_update() for cfg in deepcopy(self.config)]
 
         # sft warmup stage
         sft_config = stage_configs[0]
@@ -351,6 +354,10 @@ def test_trainer(self, mock_load):
         self.assertEqual(parser.metric_min_step(response_metrics[0]), 1)
         self.assertEqual(parser.metric_max_step(response_metrics[0]), 3)
 
+        self.config.stages[1].buffer.explorer_input.taskset.path = old_taskset_path
+        mock_load.return_value = deepcopy(self.config)
+        run(config_path="dummy.yaml")
+
         # grpo stage
         grpo_config = stage_configs[1]
         parser = TensorBoardParser(os.path.join(grpo_config.monitor.cache_dir, "tensorboard"))

trinity/cli/launcher.py

@@ -191,6 +191,7 @@ def run(config_path: str, dlc: bool = False, plugin_dir: str = None):
                         f"> Skipping completed stage {i + 1}/{len(config.stages)}...\n"
                         "==========================================================="
                     )
+                    stage_config.check_and_update()
                 else:
                     logger.info(
                         "===========================================================\n"