Use `tempfile` to handle temporary configuration files and ensure cleanup in training stages

training_pipeline/stages/training.py

@@ -8,6 +8,7 @@
 
 import json
 import os
+import tempfile
 from typing import Dict, Any, Generator
 
 import yaml
@@ -74,25 +75,27 @@ class SoVITSTrainStage(BaseStage):
         data["version"] = version_str
         
         # 写入临时配置
-        tmp_dir = os.path.join(os.getcwd(), "TEMP")
-        os.makedirs(tmp_dir, exist_ok=True)
-        tmp_config_path = f"{tmp_dir}/tmp_s2.json"
-        with open(tmp_config_path, "w") as f:
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
             f.write(json.dumps(data))
+            tmp_config_path = f.name
             
-        # 选择训练脚本
-        if version_str in ["v1", "v2", "v2Pro", "v2ProPlus"]:
-            cmd = f'PYTHONPATH=.:GPT_SoVITS "{cfg.python_exec}" -s GPT_SoVITS/s2_train.py --config "{tmp_config_path}"'
-        else:
-            cmd = f'PYTHONPATH=.:GPT_SoVITS "{cfg.python_exec}" -s GPT_SoVITS/s2_train_v3_lora.py --config "{tmp_config_path}"'
+        try:
+            # 选择训练脚本
+            if version_str in ["v1", "v2", "v2Pro", "v2ProPlus"]:
+                cmd = f'PYTHONPATH=.:GPT_SoVITS "{cfg.python_exec}" -s GPT_SoVITS/s2_train.py --config "{tmp_config_path}"'
+            else:
+                cmd = f'PYTHONPATH=.:GPT_SoVITS "{cfg.python_exec}" -s GPT_SoVITS/s2_train_v3_lora.py --config "{tmp_config_path}"'
+                
+            yield self._make_progress("SoVITS训练启动中...", 0.1)
             
-        yield self._make_progress("SoVITS训练启动中...", 0.1)
-        
-        self._process = self._run_command(cmd, wait=True)
-        self._process = None
-        
-        self._status = StageStatus.COMPLETED
-        yield self._make_progress("SoVITS训练完成", 1.0)
+            self._process = self._run_command(cmd, wait=True)
+            self._process = None
+            
+            self._status = StageStatus.COMPLETED
+            yield self._make_progress("SoVITS训练完成", 1.0)
+        finally:
+            if os.path.exists(tmp_config_path):
+                os.remove(tmp_config_path)
 
 
 class GPTTrainStage(BaseStage):
@@ -142,23 +145,25 @@ class GPTTrainStage(BaseStage):
         data["output_dir"] = f"{s1_dir}/logs_s1_{version_str}"
         
         # 写入临时配置
-        tmp_dir = os.path.join(os.getcwd(), "TEMP")
-        os.makedirs(tmp_dir, exist_ok=True)
-        tmp_config_path = f"{tmp_dir}/tmp_s1.yaml"
-        with open(tmp_config_path, "w") as f:
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
             f.write(yaml.dump(data, default_flow_style=False))
+            tmp_config_path = f.name
             
-        # 设置环境变量
-        os.environ["_CUDA_VISIBLE_DEVICES"] = cfg.gpu_numbers.replace("-", ",")
-        os.environ["hz"] = "25hz"
-        
-        cmd = f'PYTHONPATH=.:GPT_SoVITS "{cfg.python_exec}" -s GPT_SoVITS/s1_train.py --config_file "{tmp_config_path}"'
-        
-        yield self._make_progress("GPT训练启动中...", 0.1)
-        
-        self._process = self._run_command(cmd, wait=True)
-        self._process = None
-        
-        self._status = StageStatus.COMPLETED
-        yield self._make_progress("GPT训练完成", 1.0)
+        try:
+            # 设置环境变量
+            os.environ["_CUDA_VISIBLE_DEVICES"] = cfg.gpu_numbers.replace("-", ",")
+            os.environ["hz"] = "25hz"
+            
+            cmd = f'PYTHONPATH=.:GPT_SoVITS "{cfg.python_exec}" -s GPT_SoVITS/s1_train.py --config_file "{tmp_config_path}"'
+            
+            yield self._make_progress("GPT训练启动中...", 0.1)
+            
+            self._process = self._run_command(cmd, wait=True)
+            self._process = None
+            
+            self._status = StageStatus.COMPLETED
+            yield self._make_progress("GPT训练完成", 1.0)
+        finally:
+            if os.path.exists(tmp_config_path):
+                os.remove(tmp_config_path)