104 lines
3.5 KiB
Python
104 lines
3.5 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""
|
||
|
|
fetch_train.py v3 — 蒸馏+GPU实时进度同步
|
||
|
|
"""
|
||
|
|
import subprocess, json, re
|
||
|
|
from datetime import datetime, timezone
|
||
|
|
|
||
|
|
DATA_FILE = '/opt/guanghulab-repo/homepage/training-status.json'
|
||
|
|
SSH_BASE = ['sshpass','-p','HkM43lFVUIsc','ssh','-o','StrictHostKeyChecking=no',
|
||
|
|
'-o','ConnectTimeout=10','-p','23647','root@connect.westd.seetacloud.com']
|
||
|
|
|
||
|
|
def ssh(cmd_list):
|
||
|
|
try:
|
||
|
|
r = subprocess.run(SSH_BASE + cmd_list, capture_output=True, text=True, timeout=30)
|
||
|
|
return r.stdout.strip()
|
||
|
|
except:
|
||
|
|
return ''
|
||
|
|
|
||
|
|
def get_distill_status():
|
||
|
|
raw = ssh(['tail','-50','/root/autodl-tmp/distill_mother.log'])
|
||
|
|
if not raw:
|
||
|
|
return {}, 'no_log'
|
||
|
|
lines = raw.replace('\r', '\n').split('\n')
|
||
|
|
|
||
|
|
if 'DONE' in raw:
|
||
|
|
return {'phase': 'completed'}, 'done'
|
||
|
|
if 'Train' in raw:
|
||
|
|
phase = 'training'
|
||
|
|
elif 'Load models' in raw or 'Loading weights' in raw:
|
||
|
|
phase = 'loading_models'
|
||
|
|
elif 'Tokenize' in raw:
|
||
|
|
phase = 'tokenizing'
|
||
|
|
else:
|
||
|
|
phase = 'unknown'
|
||
|
|
|
||
|
|
epoch = None; total_epoch = 3
|
||
|
|
for line in reversed(lines):
|
||
|
|
m = re.search(r'Epoch\s+(\d+)/(\d+)', line)
|
||
|
|
if m:
|
||
|
|
epoch, total_epoch = int(m.group(1)), int(m.group(2))
|
||
|
|
break
|
||
|
|
|
||
|
|
step = None; total_steps = None
|
||
|
|
for line in reversed(lines):
|
||
|
|
m = re.search(r'(\d+)/(\d+)\s+\[', line)
|
||
|
|
if m:
|
||
|
|
step, total_steps = int(m.group(1)), int(m.group(2))
|
||
|
|
break
|
||
|
|
|
||
|
|
loss = '--'
|
||
|
|
for line in reversed(lines):
|
||
|
|
m = re.search(r"loss[=:]?\s*'?([0-9.]+)'?", line)
|
||
|
|
if m:
|
||
|
|
loss = m.group(1)
|
||
|
|
break
|
||
|
|
|
||
|
|
eta = '--'
|
||
|
|
for line in reversed(lines):
|
||
|
|
m = re.search(r'<([0-9]+:[0-9]+:[0-9]+)', line)
|
||
|
|
if m:
|
||
|
|
eta = m.group(1)
|
||
|
|
break
|
||
|
|
|
||
|
|
return {'phase': phase, 'epoch': epoch, 'total_epoch': total_epoch,
|
||
|
|
'step': step, 'total_steps': total_steps, 'loss': loss, 'eta': eta}, phase
|
||
|
|
|
||
|
|
def get_gpu_status():
|
||
|
|
raw = ssh(['nvidia-smi','--query-gpu=temperature.gpu,memory.used,memory.total,utilization.gpu,utilization.memory',
|
||
|
|
'--format=csv,noheader'])
|
||
|
|
parts = raw.split(', ')
|
||
|
|
if len(parts) >= 5:
|
||
|
|
return {'temp_c': parts[0].strip(),
|
||
|
|
'mem_used_gb': round(int(parts[1].strip().split()[0])/1024, 1),
|
||
|
|
'mem_total_gb': round(int(parts[2].strip().split()[0])/1024, 1),
|
||
|
|
'gpu_util_pct': parts[3].strip().split()[0],
|
||
|
|
'mem_util_pct': parts[4].strip().split()[0]}
|
||
|
|
return {}
|
||
|
|
|
||
|
|
def main():
|
||
|
|
info = {'mode': 'distill_shuangyan'}
|
||
|
|
status, phase = get_distill_status()
|
||
|
|
info.update(status)
|
||
|
|
info['gpu'] = get_gpu_status()
|
||
|
|
|
||
|
|
if info.get('phase') == 'completed':
|
||
|
|
info['display_step'] = '完成'
|
||
|
|
info['display_pct'] = 100.0
|
||
|
|
elif info.get('step') and info.get('total_steps') and info.get('epoch'):
|
||
|
|
ep_progress = (info['epoch'] - 1) * 100 / info['total_epoch']
|
||
|
|
ep_step = info['step'] / info['total_steps'] * 100 / info['total_epoch']
|
||
|
|
info['display_pct'] = round(ep_progress + ep_step, 1)
|
||
|
|
info['display_step'] = f"Ep{info['epoch']}/{info['total_epoch']} · {info['step']}/{info['total_steps']}"
|
||
|
|
else:
|
||
|
|
info['display_step'] = phase
|
||
|
|
info['display_pct'] = 0
|
||
|
|
|
||
|
|
info['updated'] = datetime.now(timezone.utc).strftime('%Y-%m-%dT%H:%M:%S.000000+00:00Z')
|
||
|
|
with open(DATA_FILE, 'w') as f:
|
||
|
|
json.dump(info, f)
|
||
|
|
print(json.dumps(info, ensure_ascii=False, indent=2))
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
main()
|