diff --git a/scripts/auto_pipeline.py b/scripts/auto_pipeline.py new file mode 100644 index 0000000..8399842 --- /dev/null +++ b/scripts/auto_pipeline.py @@ -0,0 +1,121 @@ +#!/usr/bin/env python3 +"""自动训练流水线 - 铸渊自管理系统 +母模型训练完成 → 上传COS → 启动代码模型训练 → 上传COS → 关实例 + +使用方法: + nohup python3 auto_pipeline.py > pipeline.log 2>&1 & +""" + +import os, json, time, sys, subprocess +sys.stdout.reconfigure(line_buffering=True) + +MOTHER_LOG = "/root/autodl-tmp/train_mother.log" +MOTHER_OUT = "/root/autodl-tmp/output/qwen25-7b-sft/final" + +COS_BUCKET = "sy-finetune-corpus-1317346199" +COS_REGION = "ap-guangzhou" +COS_SECRET_ID = "AKIDkQuBQhoiS2OYXWebXLwMbdT7cvAScbbU" +COS_SECRET_KEY = "nPoZKArgUJBA4nJenjSxJSQBj5FCj3A4" + +CODE_MODEL = "Qwen/Qwen2.5-Coder-7B" +CODE_CACHE = "/root/autodl-tmp/cache" +CODE_OUT = "/root/autodl-tmp/output/qwen25-coder-7b-sft" +CODE_SCRIPT = "/root/autodl-tmp/train_coder.py" +DATA_FILE = "/root/autodl-tmp/data/sft.jsonl" +POLL_INTERVAL = 300 + +def run_cmd(cmd): + result = subprocess.run(cmd, shell=True, capture_output=True, text=True) + return result.stdout.strip(), result.returncode + +def log(msg): + t = time.strftime("%Y-%m-%d %H:%M:%S") + print(f"[{t}] {msg}") + sys.stdout.flush() + +def upload_to_cos(local_path, cos_key): + log(f"Uploading {local_path} → cos://{COS_BUCKET}/{cos_key}") + script = f""" +from qcloud_cos import CosConfig, CosS3Client +import os +config = CosConfig(Region='{COS_REGION}', SecretId='{COS_SECRET_ID}', SecretKey='{COS_SECRET_KEY}') +client = CosS3Client(config) +if os.path.isdir('{local_path}'): + import glob + for f in glob.glob('{local_path}/**/*', recursive=True): + if os.path.isfile(f): + rel = os.path.relpath(f, '{local_path}') + client.upload_file(Bucket='{COS_BUCKET}', Key='{cos_key}/' + rel, LocalFilePath=f) +else: + client.upload_file(Bucket='{COS_BUCKET}', Key='{cos_key}', LocalFilePath='{local_path}') +print('Upload OK') +""" + with open("/tmp/cos_upload.py", "w") as fp: + fp.write(script) + out, rc = run_cmd(f"export PATH=/root/miniconda3/bin:$PATH && python3 /tmp/cos_upload.py") + log(out) + return rc == 0 + +def check_mother_done(): + if not os.path.exists(MOTHER_LOG): + return False + out, _ = run_cmd(f"grep -a 'DONE' {MOTHER_LOG} | tail -1") + return "DONE" in out + +def check_mother_output(): + return os.path.isdir(MOTHER_OUT) and any(f.endswith('.safetensors') for f in os.listdir(MOTHER_OUT)) + +def generate_coder_script(): + return open('/root/autodl-tmp/train_mother.py').read().replace( + 'Qwen2.5-7B', 'Qwen2.5-Coder-7B' + ).replace( + 'qwen25-7b-sft', 'qwen25-coder-7b-sft' + ) + +def main(): + log("=== 铸渊自动训练流水线启动 ===") + + # Phase 1: Wait for mother model + log("等待母模型训练完成...") + while True: + if check_mother_done() or check_mother_output(): + log("✅ 母模型训练完成!") + break + log(f" 母模型仍在训练中... (每{POLL_INTERVAL}秒检查)") + time.sleep(POLL_INTERVAL) + + if not check_mother_output(): + log("❌ 母模型输出不存在") + return + + # Phase 2: Upload to COS + log("=== 上传母模型到COS ===") + upload_to_cos(MOTHER_OUT, "models/qwen25-7b-sft/final") + + # Phase 3: Start code model training + log("=== 准备代码模型训练 ===") + with open(CODE_SCRIPT, 'w') as f: + f.write(generate_coder_script()) + + cmd = f"export PATH=/root/miniconda3/bin:$PATH && cd /root/autodl-tmp && nohup python3 -u {CODE_SCRIPT} > train_coder.log 2>&1 &" + run_cmd(cmd) + log("代码模型训练已启动") + + # Phase 4: Wait for code model + while True: + out, _ = run_cmd(f"grep -a 'DONE' /root/autodl-tmp/train_coder.log | tail -1") + if "DONE" in out: + log("✅ 代码模型训练完成!") + break + time.sleep(POLL_INTERVAL) + + # Phase 5: Upload code model to COS + upload_to_cos(f"{CODE_OUT}/final", "models/qwen25-coder-7b-sft/final") + + log("=" * 50) + log("U0001F389 全部训练任务完成!") + log("⚠️ GPU实例仍在运行,建议手动关停") + log("=" * 50) + +if __name__ == "__main__": + main()