134 lines
4.7 KiB
Python
134 lines
4.7 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
fetch_train.py v4 — 蒸馏+GPU实时进度同步 (FIXED: phase/epoch/step detection)
|
|
"""
|
|
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():
|
|
# Use tail -300 to ensure we catch epoch print + loss lines
|
|
raw = ssh(['tail','-300','/root/autodl-tmp/distill_mother.log'])
|
|
if not raw:
|
|
return {}, 'no_log'
|
|
lines = raw.replace('\r', '\n').split('\n')
|
|
|
|
# Phase detection: look for tqdm progress bar pattern EpX: X%|
|
|
if 'ALL DONE' in raw or 'ALL VERIFIED' in raw:
|
|
return {'phase': 'completed'}, 'completed'
|
|
|
|
# Check if tqdm is running (EpX: XX%| pattern)
|
|
has_progress = False
|
|
for line in lines:
|
|
if re.search(r'Ep\d+:\s+\d+%\|', line):
|
|
has_progress = True
|
|
break
|
|
|
|
if has_progress:
|
|
phase = 'training'
|
|
elif 'Train' in raw or 'Training' in raw:
|
|
phase = 'training'
|
|
elif 'Load models' in raw or 'Loading weights' in raw:
|
|
phase = 'loading_models'
|
|
elif 'Tokenize' in raw:
|
|
phase = 'tokenizing'
|
|
elif 'Save final' in raw or 'Upload' in raw:
|
|
phase = 'saving'
|
|
else:
|
|
phase = 'running'
|
|
|
|
# Epoch: try "EpX" from tqdm first, then "Epoch X/Y" from print
|
|
epoch = None; total_epoch = 3
|
|
for line in reversed(lines):
|
|
m = re.search(r'Ep(\d+):', line)
|
|
if m:
|
|
epoch = int(m.group(1))
|
|
break
|
|
if epoch is None:
|
|
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 & total: from tqdm progress bar step/15781
|
|
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: from "step=X loss=Y.ZZZZ" lines
|
|
loss = '--'
|
|
for line in reversed(lines):
|
|
m = re.search(r"loss=([0-9.]+)", line)
|
|
if m:
|
|
loss = m.group(1)
|
|
break
|
|
|
|
# ETA: from tqdm progress <HH:MM:SS>
|
|
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') in ('completed', 'done'):
|
|
info['display_step'] = '完成'
|
|
info['display_pct'] = 100.0
|
|
elif info.get('step') and info.get('total_steps'):
|
|
# Calculate overall progress: (epoch-1)*100/3 + step/total*100/3
|
|
ep = info.get('epoch') or 1
|
|
total_ep = info.get('total_epoch') or 3
|
|
ep_progress = (ep - 1) * 100 / total_ep
|
|
ep_step = info['step'] / info['total_steps'] * 100 / total_ep
|
|
info['display_pct'] = round(ep_progress + ep_step, 1)
|
|
if info.get('epoch'):
|
|
info['display_step'] = f"Ep{info['epoch']}/{info['total_epoch']} · {info['step']}/{info['total_steps']}"
|
|
else:
|
|
info['display_step'] = f"{info['step']}/{info['total_steps']}"
|
|
else:
|
|
info['display_step'] = phase if phase != 'no_log' else '等待数据...'
|
|
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, ensure_ascii=False)
|
|
print(json.dumps(info, ensure_ascii=False, indent=2))
|
|
|
|
if __name__ == '__main__':
|
|
main()
|