#!/usr/bin/env python3
"""自动训练流水线 - 铸渊自管理系统
母模型完成 → 上传COS → 启动代码模型训练 → 上传COS

使用方法：
  export ZY_OSS_KEY=<SecretId>
  export ZY_OSS_SECRET=<SecretKey>
  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 = os.environ.get("ZY_OSS_KEY") or os.environ.get("COS_SECRET_ID")
COS_SECRET_KEY = os.environ.get("ZY_OSS_SECRET") or os.environ.get("COS_SECRET_KEY")
if not COS_SECRET_ID or not COS_SECRET_KEY:
    print("❌ 需要环境变量：export ZY_OSS_KEY=... ZY_OSS_SECRET=...")
    sys.exit(1)

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):
    r = subprocess.run(cmd, shell=True, capture_output=True, text=True)
    return r.stdout.strip(), r.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}")
    s = "from qcloud_cos import CosConfig, CosS3Client; "
    s += f"c=CosS3Client(CosConfig(Region='{COS_REGION}',SecretId='{COS_SECRET_ID}',SecretKey='{COS_SECRET_KEY}')); "
    s += f"c.upload_file(Bucket='{COS_BUCKET}',Key='{cos_key}',LocalFilePath='{local_path}')"
    out, rc = run_cmd(f"python3 -c \"{s}\"")
    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():
    s = open("/root/autodl-tmp/train_mother.py").read()
    s = s.replace("Qwen2___5-7B", "Qwen2___5-Coder-7B")
    s = s.replace("qwen25-7b-sft", "qwen25-coder-7b-sft")
    return s

def main():
    log("=== 自动训练流水线启动 ===")
    while True:
        if check_mother_done() or check_mother_output():
            break
        time.sleep(POLL_INTERVAL)
    log("✅ 母模型训练完成")
    if upload_to_cos(MOTHER_OUT, "models/qwen25-7b-sft/final"):
        log("✅ 母模型上传COS完成")
    with open(CODE_SCRIPT, 'w') as f: f.write(generate_coder_script())
    run_cmd(f"nohup python3 -u {CODE_SCRIPT} > train_coder.log 2>&1 &")
    log("✅ 代码模型训练已启动")
    while True:
        out, _ = run_cmd("grep -a 'DONE' /root/autodl-tmp/train_coder.log | tail -1")
        if "DONE" in out:
            log("✅ 代码模型训练完成"); break
        time.sleep(POLL_INTERVAL)
    if upload_to_cos(f"{CODE_OUT}/final", "models/qwen25-coder-7b-sft/final"):
        log("✅ 代码模型上传COS完成")
    log("🎉 全部训练完成！")

if __name__ == "__main__":
    main()