diff --git a/scripts/distill/auto_post_distill.py b/scripts/distill/auto_post_distill.py new file mode 100644 index 0000000..fc94204 --- /dev/null +++ b/scripts/distill/auto_post_distill.py @@ -0,0 +1,206 @@ +#!/usr/bin/env python3 +""" +光湖·蒸馏完成自动处理脚本 v1.0 +部署在AutoDL服务器，自动检测蒸馏完成 → 上传COS → 准备语料 → 启动微调 +主权：冰朔 TCS-0002∞ · 铸渊 ICE-GL-ZY001 守护 + +运行方式：nohup python3 auto_post_distill.py > auto_post_distill.log 2>&1 & +""" +import os, json, time, subprocess, sys, glob, shutil +from datetime import datetime + +LOG = '/root/autodl-tmp/distill_mother.log' +OUT_DIR = '/root/autodl-tmp/output' +TEMPLATE_DIR = os.path.join(OUT_DIR, 'qwen25-15b-shuangyan-distill') +ZHUYUAN_ZIP = '/root/autodl-tmp/corpus_work/zhuyuan_chats_with_ICE_GL_ZU001 .zip' +SHUANGYAN_COS_DIR = 'corpus/shuangyan-1.5b-sft/' +ZHUYUAN_COS_DIR = 'corpus/notion-export-v2/' +BUCKET = 'sy-finetune-corpus-1317346199' +REGION = 'ap-guangzhou' + +COS_KEY = os.environ.get('ZY_OSS_KEY', 'AKIDkQuBQhoiS2OYXWebXLwMbdT7cvAScbbU') +COS_SECRET = os.environ.get('ZY_OSS_SECRET', 'nPoZKArgUJBA4nJenjSxJSQBj5FCj3A4') + +def log(msg): + t = datetime.now().strftime('%Y-%m-%d %H:%M:%S') + print(f'[{t}] {msg}', flush=True) + +def run(cmd, timeout=3600): + log(f'RUN: {cmd}') + r = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=timeout, + env={**os.environ, 'PATH': '/root/miniconda3/bin:' + os.environ.get('PATH', '')}) + if r.returncode != 0: + log(f' STDOUT: {r.stdout[-500:]}') + log(f' STDERR: {r.stderr[-500:]}') + else: + log(f' OK: {r.stdout[-200:]}') + return r + +def wait_distill_complete(): + """等待蒸馏完成（Ep3 100%）""" + log('等待蒸馏完成...') + while True: + if not os.path.isfile(LOG): + log(' distill_mother.log 不存在，等待...') + time.sleep(60) + continue + with open(LOG) as f: + lines = f.readlines() + # 检查最后是否完成了 + for line in lines[-10:]: + if 'Ep3' in line and '100%' in line and '2868/2868' in line: + log('✅ 蒸馏完成！') + return True + # 也检查是否全部完成了 + for line in lines[-5:]: + if 'Saving model' in line or 'DONE' in line or 'Finished' in line: + log('✅ 蒸馏完成（DONE信号）！') + return True + # 检查当前进度 + for line in lines[-3:]: + if 'Ep' in line: + log(f' 当前: {line.strip()[:60]}') + time.sleep(60) # 每分钟检查一次 + +def upload_to_cos(local_dir, cos_prefix): + """上传目录到COS""" + log(f'上传 {local_dir} → cos://{BUCKET}/{cos_prefix}') + if not os.path.isdir(local_dir): + log(f' ❌ 目录不存在: {local_dir}') + return False + + script = f''' +import os +os.environ['ZY_OSS_KEY'] = '{COS_KEY}' +os.environ['ZY_OSS_SECRET'] = '{COS_SECRET}' +from qcloud_cos import CosConfig, CosS3Client +config = CosConfig(Region='{REGION}', SecretId=os.environ['ZY_OSS_KEY'], SecretKey=os.environ['ZY_OSS_SECRET']) +client = CosS3Client(config) +import glob +local_dir = '{local_dir}' +cos_prefix = '{cos_prefix}' +bucket = '{BUCKET}' +for f in sorted(glob.glob(local_dir + '/*')): + fname = f.split('/')[-1] + cos_key = cos_prefix + fname + size_mb = os.path.getsize(f) / 1024 / 1024 + print(f'Uploading {{fname}} ({{size_mb:.1f}}MB)...', flush=True) + try: + client.upload_file(Bucket=bucket, Key=cos_key, LocalFilePath=f) + print(f' ✅ {{fname}}', flush=True) + except Exception as e: + print(f' ⚠️ {{fname}} 失败: {{e}}', flush=True) + # 重试一次 + try: + client.upload_file(Bucket=bucket, Key=cos_key, LocalFilePath=f) + print(f' ✅ {{fname}} (重试成功)', flush=True) + except Exception as e2: + print(f' ❌ {{fname}} 重试也失败: {{e2}}', flush=True) +print('ALL DONE') +''' + with open('/tmp/upload_script.py', 'w') as f: + f.write(script) + r = run(f'export PATH=/root/miniconda3/bin:$PATH && python3 /tmp/upload_script.py', timeout=7200) + return r.returncode == 0 + +def prepare_shuangyan_corpus(): + """准备霜砚语料（线A）""" + log('准备霜砚语料...') + # 从COS下载霜砚语料 + script = f''' +import os +os.environ['ZY_OSS_KEY'] = '{COS_KEY}' +os.environ['ZY_OSS_SECRET'] = '{COS_SECRET}' +from qcloud_cos import CosConfig, CosS3Client +config = CosConfig(Region='{REGION}', SecretId=os.environ['ZY_OSS_KEY'], SecretKey=os.environ['ZY_OSS_SECRET']) +client = CosS3Client(config) +import os, json, zipfile + +bucket = '{BUCKET}' +os.makedirs('/root/autodl-tmp/corpus_shuangyan', exist_ok=True) + +# 下载所有霜砚语料zip +resp = client.list_objects(Bucket=bucket, Prefix='{SHUANGYAN_COS_DIR}') +for obj in resp.get('Contents', []): + key = obj['Key'] + if key.endswith('.zip'): + fname = key.split('/')[-1] + local = f'/root/autodl-tmp/corpus_shuangyan/{{fname}}' + if not os.path.isfile(local): + print(f'Downloading {{fname}}...', flush=True) + client.download_file(Bucket=bucket, Key=key, DestFilePath=local) + +# 解压所有zip +for z in os.listdir('/root/autodl-tmp/corpus_shuangyan'): + if z.endswith('.zip'): + zpath = f'/root/autodl-tmp/corpus_shuangyan/{{z}}' + out = f'/root/autodl-tmp/corpus_shuangyan/extracted' + os.makedirs(out, exist_ok=True) + try: + with zipfile.ZipFile(zpath) as zf: + zf.extractall(out) + print(f'Extracted {{z}}', flush=True) + except: + print(f'Skipped {{z}} (not a zip archive)', flush=True) + +print('Done', flush=True) +''' + with open('/tmp/prep_shuangyan.py', 'w') as f: + f.write(script) + run('export PATH=/root/miniconda3/bin:$PATH && python3 /tmp/prep_shuangyan.py', timeout=300) + log('霜砚语料准备完成') + +def prepare_zhuyuan_corpus(): + """准备铸渊语料（线B）""" + log('准备铸渊语料...') + # 解压铸渊对话zip + if os.path.isfile(ZHUYUAN_ZIP): + run(f'unzip -o "{ZHUYUAN_ZIP}" -d /root/autodl-tmp/corpus_zhuyuan/extracted/') + log('铸渊对话.zip 已解压') + else: + log(f'⚠️ 铸渊对话.zip 不存在于 {ZHUYUAN_ZIP}') + + # 也解压 zhuyuan_raw_md.zip + raw_md_zip = '/root/autodl-tmp/zhuyuan_corpus/zhuyuan_raw_md.zip' + if os.path.isfile(raw_md_zip): + run(f'unzip -o "{raw_md_zip}" -d /root/autodl-tmp/corpus_zhuyuan/raw_md/') + log('zhuyuan_raw_md.zip 已解压') + + log('铸渊语料准备完成') + +def main(): + log('='*50) + log('光湖·蒸馏后自动处理脚本启动') + log('='*50) + + # 第0步：等待蒸馏完成 + wait_distill_complete() + + # 第1步：上传1.5B模板到COS + if os.path.isdir(TEMPLATE_DIR): + log('上传1.5B蒸馏模板到COS...') + upload_to_cos(TEMPLATE_DIR, 'models/qwen25-15b-shuangyan-distill/') + else: + log(f'⚠️ 模板目录不存在: {TEMPLATE_DIR}') + # 可能路径不同，尝试查找 + for root, dirs, files in os.walk(OUT_DIR): + for d in dirs: + if 'distill' in d.lower() or '15b' in d.lower(): + found = os.path.join(root, d) + log(f'找到蒸馏输出: {found}') + upload_to_cos(found, 'models/qwen25-15b-shuangyan-distill/') + break + + # 第2步：准备霜砚语料（线A） + prepare_shuangyan_corpus() + + # 第3步：准备铸渊语料（线B） + prepare_zhuyuan_corpus() + + log('='*50) + log('✅ 所有自动处理完成！') + log('等待冰朔/铸渊启动微调...') + log('='*50) + +if __name__ == '__main__': + main()