| """Durable monitor and automatic packaging/upload after successful training.""" |
| import json |
| import os |
| import re |
| import shutil |
| import signal |
| import subprocess |
| import sys |
| import time |
| from pathlib import Path |
|
|
| ROOT = Path('/data/realstyle2') |
| pid = int(sys.argv[1]) |
| log = Path(sys.argv[2]) |
| started = time.time() |
| target_epochs = 40 |
| eta_reduced = False |
| while True: |
| data = log.read_text(errors='replace') |
| matches = re.findall(r'steps:\s+[^\r\n]*?\b(\d+)/(\d+)\s+\[([^\]]+)\](?:[^\r\n]*)', data) |
| epochs = re.findall(r'epoch (\d+)/(\d+)', data) |
| progress = matches[-1] if matches else None |
| projected_seconds = None |
| if progress and int(progress[0]) >= 100: |
| rate = re.search(r'([\d.]+)s/it', progress[2]) |
| if rate: |
| projected_seconds = float(rate[1]) * int(progress[1]) * target_epochs / 40 |
| if projected_seconds > 4 * 3600 and target_epochs > 30: |
| target_epochs = 35 if float(rate[1]) * int(progress[1]) * 35 / 40 <= 4 * 3600 else 30 |
| eta_reduced = True |
| (ROOT / 'metadata/eta_reduction.json').write_text(json.dumps({ |
| 'target_epochs': target_epochs, 'estimated_full_40_seconds': float(rate[1]) * int(progress[1]), |
| 'at_step': int(progress[0]), 'reason': 'User requested a reduction when projected training exceeds 4 hours.'}, indent=2)) |
| stopped_at_limit = False |
| if eta_reduced and epochs and int(epochs[-1][0]) > target_epochs: |
| |
| checkpoint = ROOT / 'output' / f'realstyle2_anima-{target_epochs:06d}.safetensors' |
| assert checkpoint.exists() |
| os.kill(pid, signal.SIGTERM) |
| for _ in range(30): |
| try: |
| if Path(f'/proc/{pid}/stat').read_text().split()[2] == 'Z': |
| break |
| except FileNotFoundError: |
| break |
| time.sleep(1) |
| shutil.copyfile(checkpoint, ROOT / 'output/realstyle2_anima.safetensors') |
| import toml |
| config_path = ROOT / 'configs/train.toml' |
| config = toml.load(config_path) |
| config['max_train_epochs'] = target_epochs |
| config_path.write_text(toml.dumps(config)) |
| stopped_at_limit = True |
| gpu = subprocess.check_output(['nvidia-smi', '--query-gpu=memory.used,utilization.gpu', '--format=csv,noheader'], text=True).strip() |
| status = {'pid': pid, 'epoch': epochs[-1] if epochs else None, 'progress': progress, |
| 'gpu': gpu, 'updated_at': time.time(), 'elapsed_monitor_seconds': time.time()-started, |
| 'target_epochs': target_epochs, 'projected_training_seconds': projected_seconds, |
| 'saved_checkpoints': len(list((ROOT / 'output').glob('*.safetensors')))} |
| (ROOT / 'status.json').write_text(json.dumps(status, indent=2)) |
| print(json.dumps(status), flush=True) |
| try: |
| os.kill(pid, 0) |
| process_state = Path(f'/proc/{pid}/stat').read_text().split()[2] |
| alive = process_state != 'Z' |
| except (ProcessLookupError, FileNotFoundError): |
| alive = False |
| if stopped_at_limit: |
| break |
| if not alive: |
| if 'model saved.' not in data or not (ROOT / 'output/realstyle2_anima.safetensors').is_file(): |
| raise RuntimeError('Training process exited before successful final save; inspect training log') |
| break |
| time.sleep(30) |
| print('Training complete. Validating, packaging, uploading.', flush=True) |
| with (ROOT / 'logs/upload.log').open('w') as out: |
| subprocess.run(['/data/venv-training/bin/python', str(ROOT / 'scripts/package_upload.py')], stdout=out, stderr=subprocess.STDOUT, check=True) |
| print('UPLOAD COMPLETE', flush=True) |
|
|
