📜 蒸馏脚本：auto_post_distill.py
This commit is contained in:
bingshuo
parent
c2eb01512c
commit
4f86041191
206
scripts/distill/auto_post_distill.py
Normal file
206
scripts/distill/auto_post_distill.py
Normal file
@ -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()
|
||||
Loading…
x
Reference in New Issue
Block a user