Mentors4EDU's picture
Publish training pipeline under training/; update repository notes
db3982a verified
Raw History Blame Contribute Delete
3.84 kB
from __future__ import annotations
import argparse
import json
import logging
import subprocess
import sys
import time
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from riemann.paths import CHECKPOINT_ROOT
logger = logging.getLogger("riemann.supervise")
def current_step(checkpoint_dir: Path) -> int:
p = checkpoint_dir / "training_state.json"
if not p.exists():
return 0
try:
return int(json.loads(p.read_text(encoding="utf-8")).get("global_step", 0))
except Exception:
return 0
def main_gpu_present() -> bool:
try:
import torch_directml
if torch_directml.device_count() == 0:
return False
name = torch_directml.device_name(0)
return "880M" not in name
except Exception:
try:
import torch
return torch.cuda.is_available()
except Exception:
return False
def wait_for_gpu(poll_seconds: int) -> None:
waited = 0
while not main_gpu_present():
if waited % (poll_seconds * 6) == 0:
logger.warning(
"Main GPU not present; waiting for it to reconnect (%ds so far)", waited
)
time.sleep(poll_seconds)
waited += poll_seconds
if waited:
logger.info("Main GPU is back after %ds; resuming", waited)
time.sleep(10)
def main() -> None:
logging.basicConfig(
level=logging.INFO, format="%(asctime)s %(levelname)s [supervise] %(message)s"
)
parser = argparse.ArgumentParser()
parser.add_argument("--stage", required=True)
parser.add_argument("--max-steps", type=int, required=True)
parser.add_argument("--max-restarts", type=int, default=200)
parser.add_argument("--gpu-poll-seconds", type=int, default=15)
parser.add_argument("--checkpoint-root", default=str(CHECKPOINT_ROOT))
parser.add_argument("--log-file", default=None)
parser.add_argument("pretrain_args", nargs=argparse.REMAINDER)
args = parser.parse_args()
checkpoint_dir = Path(args.checkpoint_root) / args.stage
script = Path(__file__).resolve().parent / "pretrain.py"
extra = [a for a in args.pretrain_args if a != "--"]
restarts = 0
while True:
step = current_step(checkpoint_dir)
if step >= args.max_steps:
logger.info("Target reached: %s at step %d", args.stage, step)
return
wait_for_gpu(args.gpu_poll_seconds)
cmd = [
sys.executable,
"-u",
str(script),
"--stage",
args.stage,
"--resume",
"--max-steps",
str(args.max_steps),
"--checkpoint-root",
args.checkpoint_root,
*extra,
]
logger.info("Launch %d from step %d: %s", restarts + 1, step, " ".join(cmd[1:]))
if args.log_file:
with open(args.log_file, "a", encoding="utf-8") as fh:
rc = subprocess.call(cmd, stdout=fh, stderr=subprocess.STDOUT)
else:
rc = subprocess.call(cmd)
new_step = current_step(checkpoint_dir)
if rc == 0 and new_step >= args.max_steps:
logger.info("Finished: %s at step %d", args.stage, new_step)
return
restarts += 1
logger.warning(
"pretrain.py exited with code %d at step %d (restart %d/%d)",
rc,
new_step,
restarts,
args.max_restarts,
)
if restarts >= args.max_restarts:
logger.error("Too many restarts; giving up at step %d", new_step)
sys.exit(1)
time.sleep(20)
if __name__ == "__main__":
main()