207 lines
7.5 KiB
Python
207 lines
7.5 KiB
Python
#!/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()
|