# Resume all interrupted trainings in yolov5/ dir including DDP trainings

# Usage: $ python utils/aws/resume.py

import os

import sys

from pathlib import Path

import torch

import yaml

FILE = Path(__file__).resolve()

ROOT = FILE.parents[2]  # YOLOv5 root directory

if str(ROOT) not in sys.path:

    sys.path.append(str(ROOT))  # add ROOT to PATH

port = 0  # --master_port

path = Path('').resolve()

for last in path.rglob('*/**/last.pt'):

    ckpt = torch.load(last)

    if ckpt['optimizer'] is None:

        continue

    # Load opt.yaml

    with open(last.parent.parent / 'opt.yaml', errors='ignore') as f:

        opt = yaml.safe_load(f)

    # Get device count

    d = opt['device'].split(',')  # devices

    nd = len(d)  # number of devices

    ddp = nd > 1 or (nd == 0 and torch.cuda.device_count() > 1)  # distributed data parallel

    if ddp:  # multi-GPU

        port += 1

        cmd = f'python -m torch.distributed.run --nproc_per_node {nd} --master_port {port} train.py --resume {last}'

    else:  # single-GPU

        cmd = f'python train.py --resume {last}'

    cmd += ' > /dev/null 2>&1 &'  # redirect output to dev/null and run in daemon thread

    print(cmd)

    os.system(cmd)