overbrainer 0.6.1

"""Axolotl plugin shipped by overbrainer.

Appends one JSON line per Trainer log to the file named by $OVERBRAINER_METRICS,
from the main process only. overbrainer tails that file for live metrics. While an
evaluation runs, an "eval" line says how far it got, at most every 2 seconds; once
the training loop ends, an "end" line says the trainer's final work began.

When the file named by $OVERBRAINER_SNAPSHOT appears (the snapshot request),
saves a checkpoint at the end of the current step, stops training, and writes
snapshot.json next to the request: the proof overbrainer reads once the job ends.
"""

import json
import math
import os
import time

from axolotl.integrations.base import BasePlugin
from transformers import TrainerCallback

try:
    import torch
    import torch.distributed as dist
except ImportError:
    torch = None

FIELDS = ("loss", "eval_loss", "learning_rate", "grad_norm")

EVAL_EVERY = 2.0

SNAPSHOT_FILE = "snapshot.json"
REASONS = ("requested", "deadline", "cost", "disk")


def _number(value):
    try:
        value = float(value)
    except (TypeError, ValueError):
        return None
    return value if math.isfinite(value) else None


def _length(loader):
    """The number of batches of `loader`, None when it cannot tell."""
    try:
        length = len(loader)
    except Exception:
        return None
    return length if isinstance(length, int) and length > 0 else None


class OverbrainerMetricsCallback(TrainerCallback):
    def __init__(self, path):
        self.path = path
        os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
        self.eval_step = 0
        self.eval_written = None

    def _write(self, record):
        with open(self.path, "a", encoding="utf-8") as file:
            file.write(json.dumps(record) + "\n")
            file.flush()

    def on_train_begin(self, args, state, control, **kwargs):
        if state.is_world_process_zero:
            self._write(
                {"event": "begin", "time": time.time(), "max_steps": state.max_steps}
            )

    def on_log(self, args, state, control, logs=None, **kwargs):
        if not state.is_world_process_zero or not logs:
            return
        record = {
            "event": "log",
            "time": time.time(),
            "step": state.global_step,
            "epoch": _number(state.epoch),
            "max_steps": state.max_steps,
        }
        for field in FIELDS:
            if field in logs:
                record[field] = _number(logs[field])
        self._write(record)

    # The progress lines below are a convenience: they never raise, so a
    # failure to write them never stops a training.

    def on_prediction_step(self, args, state, control, **kwargs):
        try:
            if not state.is_world_process_zero:
                return
            self.eval_step += 1
            now = time.time()
            if self.eval_written is not None and now - self.eval_written < EVAL_EVERY:
                return
            self.eval_written = now
            record = {"event": "eval", "time": now, "step": self.eval_step}
            total = _length(kwargs.get("eval_dataloader"))
            if total is not None:
                record["total"] = total
            self._write(record)
        except Exception:
            pass

    def _evaluated(self):
        self.eval_step = 0
        self.eval_written = None

    def on_evaluate(self, args, state, control, **kwargs):
        self._evaluated()

    def on_predict(self, args, state, control, **kwargs):
        self._evaluated()

    def on_train_end(self, args, state, control, **kwargs):
        try:
            if state.is_world_process_zero:
                self._write(
                    {"event": "end", "time": time.time(), "step": state.global_step}
                )
        except Exception:
            pass


def _reason(path):
    """The reason written in the request, `requested` when it names none."""
    try:
        with open(path, encoding="utf-8") as file:
            reason = file.read().strip()
    except OSError:
        return "requested"
    return reason if reason in REASONS else "requested"


class OverbrainerSnapshotCallback(TrainerCallback):
    """Saves a checkpoint and stops training once the request file exists.

    Rank 0 looks for the file after every step and, under distributed training,
    broadcasts what it found, so every rank stops at the same step. The trainer
    saves the checkpoint before it checks should_training_stop, then calls
    on_save, which writes snapshot.json atomically.
    """

    def __init__(self, request):
        self.request = request
        self.root = os.path.dirname(request) or "."
        self.reason = None

    def _requested(self, args, state):
        found = state.is_world_process_zero and os.path.exists(self.request)
        if torch is None or not (dist.is_available() and dist.is_initialized()):
            return found
        flag = torch.tensor([1 if found else 0], device=args.device)
        dist.broadcast(flag, src=0)
        return bool(flag.item())

    def on_step_end(self, args, state, control, **kwargs):
        # The last step ends the run anyway, with its model saved and merged:
        # a request then changes nothing. Every rank sees the same step.
        max_steps = getattr(state, "max_steps", 0) or 0
        if max_steps > 0 and state.global_step >= max_steps:
            return
        if self.reason is not None or not self._requested(args, state):
            return
        self.reason = _reason(self.request) if state.is_world_process_zero else "requested"
        control.should_save = True
        control.should_training_stop = True

    def on_save(self, args, state, control, **kwargs):
        if self.reason is None or not state.is_world_process_zero:
            return
        checkpoint = os.path.join(args.output_dir, f"checkpoint-{state.global_step}")
        record = {
            "checkpoint": os.path.relpath(checkpoint, self.root),
            "step": state.global_step,
            "time": time.time(),
            "reason": self.reason,
        }
        path = os.path.join(self.root, SNAPSHOT_FILE)
        with open(path + ".tmp", "w", encoding="utf-8") as file:
            file.write(json.dumps(record) + "\n")
            file.flush()
            os.fsync(file.fileno())
        os.replace(path + ".tmp", path)


class OverbrainerMetricsPlugin(BasePlugin):
    def add_callbacks_pre_trainer(self, cfg, model):
        callbacks = []
        path = os.environ.get("OVERBRAINER_METRICS")
        if path:
            callbacks.append(OverbrainerMetricsCallback(path))
        request = os.environ.get("OVERBRAINER_SNAPSHOT")
        if request:
            callbacks.append(OverbrainerSnapshotCallback(request))
        return callbacks