"""Save-time evaluation policy for retaining one best training checkpoint. The regular :class:`transformers.Trainer` save schedule is deliberately kept as the clock for this feature. ``CheckpointEvaluationTriggerCallback`` turns each native save request into an evaluation request; the trainer then calls ``CheckpointEvaluationPolicy.observe`` with the score produced by its in-memory evaluator. Only an improving score is allowed to turn saving back on. Keeping the decision policy independent from Trainer makes it small, easy to resume, and testable without loading a model or importing an accelerator. """ from __future__ import annotations import json import math import os from dataclasses import asdict, dataclass from pathlib import Path from typing import Any, Callable, Mapping from uuid import uuid4 from opsd_utils.checkpoint_eval_paths import update_final_checkpoint_link try: # Allow the lightweight policy tests to run without the training stack. from transformers import TrainerCallback from transformers.trainer_callback import ExportableState except ImportError: # pragma: no cover - exercised only in minimal tooling envs class TrainerCallback: # type: ignore[no-redef] """Minimal compatibility base used when transformers is unavailable.""" class ExportableState: # type: ignore[no-redef] """Compatibility marker for environments that only test this module.""" @dataclass class CheckpointEvaluationState: """Durable state required to make save decisions after a resume.""" best_score: float | None = None best_step: int | None = None # Kept for checkpoint compatibility; with tie_policy="count" this is the # consecutive non-improvement streak (both lower and equal scores). lower_score_streak: int = 0 evaluation_count: int = 0 # A terminal policy result must survive a crash before the trainer exits. # Otherwise a resumed job could perform another long interval of training # before noticing that the third lower score had already occurred. stop_requested: bool = False # ``True`` means best_score is the run's step-zero reference rather than # a checkpoint produced by this run. It participates in comparisons but # must never be treated as ``output_dir/checkpoint-``. best_is_external: bool = False @dataclass(frozen=True) class CheckpointEvaluationDecision: """The result of comparing one evaluation score with the retained best.""" score: float step: int should_save: bool should_stop: bool reason: str best_score: float best_step: int lower_score_streak: int evaluation_count: int best_is_external: bool = False @property def improved(self) -> bool: """Whether this result replaced the best checkpoint.""" return self.should_save def to_dict(self) -> dict[str, Any]: return asdict(self) @classmethod def from_dict(cls, payload: Mapping[str, Any]) -> "CheckpointEvaluationDecision": """Restore a decision sent over a distributed object broadcast.""" if not isinstance(payload, Mapping): raise ValueError("checkpoint evaluation decision must be a mapping") required = { "score", "step", "should_save", "should_stop", "reason", "best_score", "best_step", "lower_score_streak", "evaluation_count", } missing = required.difference(payload) if missing: raise ValueError(f"checkpoint evaluation decision missing fields: {sorted(missing)}") return cls( score=float(payload["score"]), step=int(payload["step"]), should_save=bool(payload["should_save"]), should_stop=bool(payload["should_stop"]), reason=str(payload["reason"]), best_score=float(payload["best_score"]), best_step=int(payload["best_step"]), lower_score_streak=int(payload["lower_score_streak"]), evaluation_count=int(payload["evaluation_count"]), best_is_external=bool(payload.get("best_is_external", False)), ) class CheckpointEvaluationPolicy: """Compare checkpoint-time scores and implement best-only retention. A first finite score always establishes the initial recoverable checkpoint. Later scores save only when strictly greater than the best. Strictly lower scores increment ``lower_score_streak`` and stop at ``patience``. Equal scores are controlled by ``tie_policy``: * ``reset`` (the training default): reset the low-score streak; * ``ignore``: leave the streak unchanged; * ``stop``: stop immediately, without saving. * ``count``: count as a non-improvement toward ``patience``. """ STATE_VERSION = 1 def __init__( self, *, patience: int = 3, tie_policy: str = "reset", patience_start_step: int = 0, state: CheckpointEvaluationState | None = None, ) -> None: try: patience = int(patience) except (TypeError, ValueError) as exc: raise ValueError("patience must be a positive integer") from exc if patience <= 0: raise ValueError(f"patience must be a positive integer, got {patience}") tie_policy = str(tie_policy).strip().lower() if tie_policy not in {"reset", "ignore", "stop", "count"}: raise ValueError("tie_policy must be one of: reset, ignore, stop, count") try: patience_start_step = int(patience_start_step) except (TypeError, ValueError) as exc: raise ValueError("patience_start_step must be a non-negative integer") from exc if patience_start_step < 0: raise ValueError("patience_start_step must be a non-negative integer") self.patience = patience self.tie_policy = tie_policy self.patience_start_step = patience_start_step self.state = state if state is not None else CheckpointEvaluationState() self._validate_state(self.state) @staticmethod def _coerce_finite_score(score: Any) -> float: try: value = float(score) except (TypeError, ValueError) as exc: raise ValueError(f"checkpoint evaluation score must be numeric, got {score!r}") from exc if not math.isfinite(value): raise ValueError(f"checkpoint evaluation score must be finite, got {score!r}") return value @staticmethod def _coerce_step(step: Any) -> int: # bool is an int subclass but never a meaningful Trainer global step. if isinstance(step, bool): raise ValueError("checkpoint evaluation step must be a non-negative integer") try: result = int(step) except (TypeError, ValueError) as exc: raise ValueError("checkpoint evaluation step must be a non-negative integer") from exc if result < 0 or result != step: raise ValueError("checkpoint evaluation step must be a non-negative integer") return result @classmethod def _validate_state(cls, state: CheckpointEvaluationState) -> None: if state.best_score is not None: cls._coerce_finite_score(state.best_score) if state.best_step is not None: cls._coerce_step(state.best_step) if state.best_score is None and state.best_step is not None: raise ValueError("checkpoint evaluation state has best_step without best_score") if state.best_score is not None and state.best_step is None: raise ValueError("checkpoint evaluation state has best_score without best_step") if int(state.lower_score_streak) < 0 or int(state.evaluation_count) < 0: raise ValueError("checkpoint evaluation counters cannot be negative") if int(state.evaluation_count) > 0 and state.best_score is None: raise ValueError( "checkpoint evaluation state has completed evaluations without a best score" ) if state.stop_requested and state.best_score is None: raise ValueError("checkpoint evaluation stop state has no best score") if state.best_is_external and state.best_score is None: raise ValueError("external checkpoint evaluation baseline has no score") def observe(self, score: Any, *, step: Any) -> CheckpointEvaluationDecision: """Record a score and return the corresponding save/stop decision. This mutates only policy state. The caller remains responsible for saving the model on every process when ``should_save`` is true. """ score_value = self._coerce_finite_score(score) step_value = self._coerce_step(step) state = self.state if state.stop_requested: raise RuntimeError( "checkpoint evaluation policy already requested training stop; " "create a new run instead of observing more scores" ) state.evaluation_count += 1 if state.best_score is None: state.best_score = score_value state.best_step = step_value state.lower_score_streak = 0 state.stop_requested = False state.best_is_external = False return self._decision(score_value, step_value, True, False, "initial") if score_value > state.best_score: state.best_score = score_value state.best_step = step_value state.lower_score_streak = 0 state.stop_requested = False state.best_is_external = False return self._decision(score_value, step_value, True, False, "improved") # Early validation points are health monitors. They may still replace # the best on a strict improvement (handled above), but a transient # plateau/drop before the configured boundary never consumes patience. if step_value < self.patience_start_step: state.lower_score_streak = 0 state.stop_requested = False relation = "lower" if score_value < state.best_score else "tie" return self._decision( score_value, step_value, False, False, f"patience_warmup_{relation}", ) if score_value < state.best_score: state.lower_score_streak += 1 should_stop = state.lower_score_streak >= self.patience state.stop_requested = should_stop return self._decision( score_value, step_value, False, should_stop, "patience_exhausted" if should_stop else "lower", ) # Equality must not overwrite the old checkpoint. The explicit # policy makes repeated plateaus predictable for long-running jobs. if self.tie_policy == "reset": state.lower_score_streak = 0 state.stop_requested = False return self._decision(score_value, step_value, False, False, "tie_reset") if self.tie_policy == "stop": state.stop_requested = True return self._decision(score_value, step_value, False, True, "tie_stop") if self.tie_policy == "count": state.lower_score_streak += 1 should_stop = state.lower_score_streak >= self.patience state.stop_requested = should_stop return self._decision( score_value, step_value, False, should_stop, "patience_exhausted" if should_stop else "tie_counted", ) state.stop_requested = False return self._decision(score_value, step_value, False, False, "tie_ignored") def establish_external_baseline(self, score: Any, *, step: Any = 0) -> CheckpointEvaluationState: """Record a run-local reference score without treating it as an evaluation. Step-zero validation exists to establish what this *specific* student starts from. It must not be a minimum score, consume patience, create a checkpoint, or prevent optimization. Once a resume has already restored a policy state, keep that durable reference unchanged. """ score_value = self._coerce_finite_score(score) step_value = self._coerce_step(step) state = self.state if state.best_score is not None: return state state.best_score = score_value state.best_step = step_value state.lower_score_streak = 0 state.stop_requested = False state.best_is_external = True self._validate_state(state) return state def _decision( self, score: float, step: int, should_save: bool, should_stop: bool, reason: str, ) -> CheckpointEvaluationDecision: # best_score/best_step are guaranteed after the first branch of observe. assert self.state.best_score is not None assert self.state.best_step is not None return CheckpointEvaluationDecision( score=score, step=step, should_save=should_save, should_stop=should_stop, reason=reason, best_score=self.state.best_score, best_step=self.state.best_step, lower_score_streak=self.state.lower_score_streak, evaluation_count=self.state.evaluation_count, best_is_external=self.state.best_is_external, ) def state_dict(self) -> dict[str, Any]: """Return JSON-safe policy state suitable for Trainer checkpoint state.""" return { "version": self.STATE_VERSION, "patience": self.patience, "tie_policy": self.tie_policy, "patience_start_step": self.patience_start_step, "state": asdict(self.state), } @property def checkpoint_eval_state(self) -> dict[str, Any]: """Alias used by Trainer/callback checkpoint serialization glue.""" return self.state_dict() @checkpoint_eval_state.setter def checkpoint_eval_state(self, payload: Mapping[str, Any]) -> None: self.load_state_dict(payload) def load_state_dict(self, payload: Mapping[str, Any]) -> None: """Restore a state emitted by :meth:`state_dict`. Policy settings are included to prevent a resumed run from silently changing its early-stop semantics. A caller may construct with the same settings (the normal path), or explicitly restore an older state whose settings match. """ if not isinstance(payload, Mapping): raise ValueError("checkpoint evaluation state must be a mapping") version = payload.get("version", self.STATE_VERSION) if version != self.STATE_VERSION: raise ValueError(f"unsupported checkpoint evaluation state version: {version!r}") if int(payload.get("patience", self.patience)) != self.patience: raise ValueError("resumed checkpoint evaluation patience does not match configuration") if str(payload.get("tie_policy", self.tie_policy)).strip().lower() != self.tie_policy: raise ValueError("resumed checkpoint evaluation tie_policy does not match configuration") if int(payload.get("patience_start_step", self.patience_start_step)) != self.patience_start_step: raise ValueError( "resumed checkpoint evaluation patience_start_step does not match configuration" ) raw_state = payload.get("state", payload) if not isinstance(raw_state, Mapping): raise ValueError("checkpoint evaluation state payload must contain a state mapping") state = CheckpointEvaluationState( best_score=raw_state.get("best_score"), best_step=raw_state.get("best_step"), lower_score_streak=int(raw_state.get("lower_score_streak", 0)), evaluation_count=int(raw_state.get("evaluation_count", 0)), stop_requested=bool(raw_state.get("stop_requested", False)), best_is_external=bool(raw_state.get("best_is_external", False)), ) self._validate_state(state) self.state = state @classmethod def from_state_dict(cls, payload: Mapping[str, Any]) -> "CheckpointEvaluationPolicy": """Construct a policy with the settings and state in ``payload``.""" if not isinstance(payload, Mapping): raise ValueError("checkpoint evaluation state must be a mapping") policy = cls( patience=payload.get("patience", 3), tie_policy=payload.get("tie_policy", "reset"), patience_start_step=payload.get("patience_start_step", 0), ) policy.load_state_dict(payload) return policy def apply_checkpoint_evaluation_decision(control: Any, decision: CheckpointEvaluationDecision | Mapping[str, Any]) -> Any: """Apply a policy decision to a Trainer control object. This tiny helper centralizes the required veto: non-improving evaluations must clear the original save event, while an exhausted patience counter requests a clean training-loop stop. It accepts a mapping so rank-zero decisions can be broadcast as ordinary serializable objects. """ if isinstance(decision, Mapping): decision = CheckpointEvaluationDecision.from_dict(decision) if not isinstance(decision, CheckpointEvaluationDecision): raise TypeError("decision must be CheckpointEvaluationDecision or a mapping") control.should_save = bool(decision.should_save) # Custom in-memory evaluation does not run Trainer.evaluate's normal # callback path unless its caller does so explicitly. Clear this flag # here as a final lifecycle guard: otherwise the epoch-end # _maybe_log_save_evaluate invocation can evaluate the same save event a # second time (especially on the final optimizer step). if hasattr(control, "should_evaluate"): control.should_evaluate = False if decision.should_stop: control.should_training_stop = True return control class CheckpointEvaluationTriggerCallback(TrainerCallback, ExportableState): """Redirect native Trainer save events to the in-memory evaluator. ``DefaultFlowCallback`` sets ``control.should_save`` at configured save steps/epochs. This callback is appended after it, clears that flag and requests evaluation instead. The trainer consumes the pending flag during its custom ``evaluate`` implementation and uses the policy decision to set ``should_save`` back to true only for an improved score. """ def __init__( self, *, enabled: bool = True, patience: int = 3, tie_policy: str = "reset", patience_start_step: int = 0, gate_steps: list[int] | tuple[int, ...] | None = None, regular_steps: int | None = None, save_total_limit: int = 1, output_dir: str | os.PathLike[str] | None = None, ) -> None: self.enabled = bool(enabled) self.policy = CheckpointEvaluationPolicy( patience=patience, tie_policy=tie_policy, patience_start_step=patience_start_step, ) try: self.save_total_limit = int(save_total_limit) except (TypeError, ValueError) as exc: raise ValueError("checkpoint evaluation save_total_limit must be positive") from exc if self.save_total_limit <= 0: raise ValueError("checkpoint evaluation save_total_limit must be positive") normalized_gate_steps: set[int] = set() for raw_step in gate_steps or (): step = int(raw_step) if step <= 0: raise ValueError("checkpoint evaluation gate_steps must contain positive integers") normalized_gate_steps.add(step) self.gate_steps = tuple(sorted(normalized_gate_steps)) if regular_steps is not None: regular_steps = int(regular_steps) if regular_steps <= 0: raise ValueError("checkpoint evaluation regular_steps must be positive") self.regular_steps = regular_steps self.output_dir = os.fspath(output_dir) if output_dir is not None else None self.pending_checkpoint_evaluation = False self.trigger_count = 0 # The in-flight save flag is process-local: a checkpoint is serialized # before ``on_save``, so persisting it would make a restored callback # believe a successful save was still in flight. The terminal-stop # flag is mirrored in policy state and intentionally survives resume. self._pending_checkpoint_save_step: int | None = None self._policy_stop_requested = False # ``persist`` can be called by the trainer immediately after a policy # decision. Cache the process role from on_train_begin so all ranks # may invoke it safely without racing on the metadata file. self._is_world_process_zero: bool | None = None @property def state_path(self) -> Path | None: """Path of the external policy metadata, or ``None`` when disabled.""" return Path(self.output_dir) / "checkpoint_eval_state.json" if self.output_dir else None @staticmethod def _world_process_zero(state: Any) -> bool: """Use Trainer's global-rank flag, defaulting to single-process true.""" return bool(getattr(state, "is_world_process_zero", True)) def _set_output_dir_from_training_args(self, args: Any) -> None: if self.output_dir is None: output_dir = getattr(args, "output_dir", None) if output_dir: self.output_dir = os.fspath(output_dir) def _best_checkpoint_dir(self, *, policy: CheckpointEvaluationPolicy | None = None) -> Path | None: """Return the policy's retained checkpoint directory when it exists. The sidecar is allowed to be newer than the retained Trainer checkpoint after one or more lower-scoring evaluations. It is *not* allowed to nominate a best step whose checkpoint was never committed: that can happen if a process dies between an evaluation and native Trainer checkpoint writing. Checking this path is the commit barrier for sidecar recovery. """ active_policy = self.policy if policy is None else policy step = active_policy.state.best_step if step is None or self.output_dir is None: return None return Path(self.output_dir) / f"checkpoint-{int(step)}" def _sidecar_policy_is_committed(self, policy: CheckpointEvaluationPolicy) -> bool: """Whether a sidecar policy can safely override checkpoint state.""" # ``output_dir`` is always set by on_train_begin in real Trainer runs. # Keep direct, lightweight uses of ``load`` backwards compatible when # a caller intentionally has no filesystem location. if self.output_dir is None: return True checkpoint_dir = self._best_checkpoint_dir(policy=policy) if policy.state.best_is_external: return True return checkpoint_dir is not None and checkpoint_dir.is_dir() def persist(self, *, is_world_process_zero: bool | None = None) -> Path | None: """Atomically persist policy metadata outside model checkpoints. This is intentionally a sidecar under ``output_dir`` rather than an edit to ``checkpoint-*/trainer_state.json``. A lower-scoring result has no checkpoint to write, but it still has to survive a crash so the three-strike early-stop counter remains correct after resuming. The trainer may call this on every rank. Only global rank zero writes; all other ranks return without touching the filesystem. """ is_main = self._is_world_process_zero if is_world_process_zero is None else bool(is_world_process_zero) if is_main is False: return None path = self.state_path if path is None: return None path.parent.mkdir(parents=True, exist_ok=True) payload = { "version": CheckpointEvaluationPolicy.STATE_VERSION, "checkpoint_eval_state": self.checkpoint_eval_state, } temporary = path.with_name(f".{path.name}.tmp-{os.getpid()}-{uuid4().hex}") try: with temporary.open("w", encoding="utf-8") as handle: json.dump(payload, handle, indent=2, sort_keys=True) handle.write("\n") handle.flush() os.fsync(handle.fileno()) os.replace(temporary, path) finally: if temporary.exists(): temporary.unlink() return path def load(self, *, is_world_process_zero: bool | None = None) -> bool: """Load sidecar policy metadata, returning false when it is absent. Invalid or incompatible metadata fails loudly instead of silently resetting patience: silently doing so could retain a lower-scoring model after a resume. Like :meth:`persist`, only rank zero reads it; the trainer's normal evaluation decision broadcast synchronizes the restored policy to other ranks before the next save decision. """ is_main = self._is_world_process_zero if is_world_process_zero is None else bool(is_world_process_zero) if is_main is False: return False path = self.state_path if path is None or not path.exists(): return False try: with path.open("r", encoding="utf-8") as handle: payload = json.load(handle) except (OSError, json.JSONDecodeError) as exc: raise ValueError(f"Cannot load checkpoint evaluation metadata from {path}") from exc if not isinstance(payload, Mapping): raise ValueError(f"Checkpoint evaluation metadata at {path} must be a mapping") if payload.get("version") != CheckpointEvaluationPolicy.STATE_VERSION: raise ValueError( f"Unsupported checkpoint evaluation metadata version at {path}: " f"{payload.get('version')!r}" ) policy_state = payload.get("checkpoint_eval_state") if not isinstance(policy_state, Mapping): raise ValueError(f"Checkpoint evaluation metadata at {path} has no policy state") # A retained Trainer checkpoint is the durable model/optimizer source # of truth. A sidecar may add later *lower* evaluations, but must # never roll the callback backwards or point it to a best model that # was not actually written. Construct a temporary policy first so # incompatible/corrupt metadata still fails loudly without mutating # the live callback. candidate = CheckpointEvaluationPolicy( patience=self.policy.patience, tie_policy=self.policy.tie_policy, patience_start_step=self.policy.patience_start_step, ) candidate.load_state_dict(policy_state) if candidate.state.evaluation_count <= self.policy.state.evaluation_count: return False if not self._sidecar_policy_is_committed(candidate): return False self.policy.load_state_dict(candidate.state_dict()) return True @property def checkpoint_eval_state(self) -> dict[str, Any]: """Serializable policy state restored by Transformers callback state.""" return self.policy.state_dict() @checkpoint_eval_state.setter def checkpoint_eval_state(self, payload: Mapping[str, Any]) -> None: self.policy.load_state_dict(payload) def record_checkpoint_evaluation_decision( self, decision: CheckpointEvaluationDecision | Mapping[str, Any] ) -> None: """Record transient lifecycle state for the trainer's policy result. A saving decision becomes durable only from :meth:`on_save`, which is invoked after ``Trainer._save_checkpoint`` has written and rotated the checkpoint. Non-saving decisions are persisted immediately by the trainer because otherwise their lower-score streak has no native checkpoint in which to live. """ if isinstance(decision, Mapping): decision = CheckpointEvaluationDecision.from_dict(decision) if not isinstance(decision, CheckpointEvaluationDecision): raise TypeError("decision must be CheckpointEvaluationDecision or a mapping") if decision.should_save: if ( self._pending_checkpoint_save_step is not None and self._pending_checkpoint_save_step != decision.step ): raise RuntimeError( "checkpoint evaluation attempted to schedule a second save before " "the prior best checkpoint committed" ) self._pending_checkpoint_save_step = decision.step if decision.should_stop: self._policy_stop_requested = True def _commit_saved_checkpoint(self, *, step: int) -> None: """Publish an already-written best checkpoint to sidecar/link users.""" is_main = self._is_world_process_zero if is_main is None: is_main = True if not is_main: return best_step = self.policy.state.best_step if best_step != step: raise RuntimeError( "checkpoint evaluation save commit does not match the policy best step: " f"saved_step={step}, best_step={best_step}" ) checkpoint_dir = self._best_checkpoint_dir() if checkpoint_dir is None or not checkpoint_dir.is_dir(): raise RuntimeError( "checkpoint evaluation expected native Trainer to write the new best " f"checkpoint before on_save: {checkpoint_dir}" ) # Keep the compatibility final_checkpoint pointer valid throughout a # long run, not merely after ``train()`` returns. Native rotation may # have just deleted the old target, so this is intentionally after # _save_checkpoint rather than at evaluation time. if self.save_total_limit == 1: update_final_checkpoint_link(self.output_dir, checkpoint_dir) else: update_final_checkpoint_link( self.output_dir, checkpoint_dir, max_checkpoints=self.save_total_limit, ) self.persist(is_world_process_zero=True) def _recover_committed_checkpoint_artifacts(self) -> None: """Repair sidecar/link state after resuming an already-written best.""" is_main = self._is_world_process_zero if is_main is None: is_main = True if ( not is_main or self.policy.state.best_step is None or self.policy.state.best_is_external ): return checkpoint_dir = self._best_checkpoint_dir() if checkpoint_dir is None or not checkpoint_dir.is_dir(): return if self.save_total_limit == 1: update_final_checkpoint_link(self.output_dir, checkpoint_dir) else: update_final_checkpoint_link( self.output_dir, checkpoint_dir, max_checkpoints=self.save_total_limit, ) self.persist(is_world_process_zero=True) @staticmethod def _wait_for_native_checkpoint_write() -> None: """Wait for rank-local checkpoint shards before publishing the link. Trainer invokes ``on_save`` on every rank after its own checkpoint writes. A small distributed barrier here prevents rank zero from publishing ``final_checkpoint`` while another rank is still writing optimizer/RNG or sharded model data. Single-process and non-torch backends intentionally take the no-op path. """ try: import torch.distributed as dist if dist.is_available() and dist.is_initialized(): dist.barrier() except (ImportError, RuntimeError): # Trainer itself supports non-torch distributed backends. Do not # make the policy callback less portable than the native saver. return def _broadcast_rank_zero_outcome( self, *, operation: str, success_payload: Mapping[str, Any] | None = None, rank_zero_error: Exception | None = None, ) -> dict[str, Any] | None: """Share a rank-zero operation's outcome before any rank can proceed. A sidecar is intentionally read and written only by global rank zero. That normally avoids filesystem races, but it has an important failure mode: if rank zero raises before a following collective, workers block forever at that collective. Put both successful values and failures in one object broadcast so every torch-distributed rank either proceeds with the exact same result or raises a contextual error. ``None`` is returned only for an apparent nonzero rank when torch distributed is not initialized. That case is not a real distributed run and preserves the callback's lightweight/direct-use behavior. """ is_main = self._is_world_process_zero is not False outcome: list[Any] = [None] if is_main: if rank_zero_error is None: outcome[0] = {"ok": True, **dict(success_payload or {})} else: outcome[0] = { "ok": False, "error_type": type(rank_zero_error).__name__, "error": str(rank_zero_error), } try: import torch.distributed as dist except ImportError: # pragma: no cover - torch is present in training distributed = False else: distributed = bool(dist.is_available() and dist.is_initialized()) if distributed: dist.broadcast_object_list(outcome, src=0) if not distributed and not is_main: return None received = outcome[0] if not isinstance(received, Mapping) or not isinstance(received.get("ok"), bool): raise RuntimeError( f"checkpoint evaluation {operation} did not receive a valid rank-zero outcome" ) if not received["ok"]: error_type = received.get("error_type", "RuntimeError") detail = received.get("error", "unknown error") message = ( f"checkpoint evaluation {operation} failed on global rank zero " f"({error_type}): {detail}" ) if rank_zero_error is not None: raise RuntimeError(message) from rank_zero_error raise RuntimeError(message) return dict(received) def _run_on_rank_zero_and_broadcast_failure( self, *, operation: str, action: Callable[[], Any] ) -> None: """Run a rank-zero-only side effect without stranding other ranks.""" is_main = self._is_world_process_zero is not False rank_zero_error: Exception | None = None if is_main: try: action() except Exception as exc: # Do not re-raise until every worker has received this result. rank_zero_error = exc self._broadcast_rank_zero_outcome( operation=operation, rank_zero_error=rank_zero_error, ) def _load_and_broadcast_policy_state_from_world_zero(self) -> bool: """Load sidecar state on rank zero and share success *or* failure. ``load`` performs validation and may raise for corrupt or incompatible metadata. It must therefore be inside the same error-envelope collective as the resulting policy state; calling ``load`` followed by a raw policy broadcast would otherwise leave every worker waiting when rank zero rejects a sidecar. """ is_main = self._is_world_process_zero is not False loaded = False rank_zero_error: Exception | None = None success_payload: Mapping[str, Any] | None = None if is_main: try: loaded = self.load(is_world_process_zero=True) success_payload = { "loaded": bool(loaded), "policy_state": self.policy.state_dict(), } except Exception as exc: # Keep the original exception for rank zero's exception chain, # while workers receive its serializable description. rank_zero_error = exc received = self._broadcast_rank_zero_outcome( operation="sidecar policy restore", success_payload=success_payload, rank_zero_error=rank_zero_error, ) if received is None: return False policy_state = received.get("policy_state") if not isinstance(policy_state, Mapping): raise RuntimeError( "checkpoint evaluation sidecar policy restore received no valid policy state" ) self.policy.load_state_dict(policy_state) return bool(received.get("loaded", False)) def _broadcast_policy_state_from_world_zero(self) -> None: """Synchronize sidecar-restored policy state to every torch rank. Callback state restored from ``trainer_state.json`` is normally the same on every rank. The sidecar can be newer, however, because lower scores do not create a native checkpoint. Only global rank zero reads that file, so immediately broadcast the resulting state before a resumed training loop can make a save or early-stop decision. Keep this helper local to the callback rather than importing ``accelerate`` here: the policy module is intentionally usable by its lightweight tests and by tooling which does not import the full training stack. """ try: import torch.distributed as dist except ImportError: # pragma: no cover - torch is present in training return if not dist.is_available() or not dist.is_initialized(): return payload: list[Any] = [self.policy.state_dict() if self._is_world_process_zero else None] dist.broadcast_object_list(payload, src=0) policy_state = payload[0] if not isinstance(policy_state, Mapping): raise RuntimeError( "checkpoint evaluation did not receive valid policy state from global rank zero" ) self.policy.load_state_dict(policy_state) def _redirect_save_to_evaluate(self, control: Any) -> Any: # HF still invokes on_epoch_end after a policy-requested early stop. # DefaultFlowCallback can set should_save there, which would otherwise # cause one extra ChartQA evaluation after the third lower score. Do # not use control.should_training_stop here: at max_steps it is also # true and the final scheduled checkpoint must still be evaluated. if self._policy_stop_requested: control.should_save = False control.should_evaluate = False self.pending_checkpoint_evaluation = False return control if not self.enabled or not getattr(control, "should_save", False): return control control.should_save = False control.should_evaluate = True self.pending_checkpoint_evaluation = True self.trigger_count += 1 return control def _step_is_scheduled(self, state: Any) -> bool: step = int(getattr(state, "global_step", 0)) return bool( step in self.gate_steps or (self.regular_steps is not None and step > 0 and step % self.regular_steps == 0) ) def on_step_end(self, args: Any, state: Any, control: Any, **kwargs: Any) -> Any: if self.enabled and self._step_is_scheduled(state): # Promote irregular warm-up gates (10/25/50) to the same native # save-event lifecycle used by the regular checkpoint cadence. control.should_save = True return self._redirect_save_to_evaluate(control) def on_epoch_end(self, args: Any, state: Any, control: Any, **kwargs: Any) -> Any: return self._redirect_save_to_evaluate(control) def on_train_begin(self, args: Any, state: Any, control: Any, **kwargs: Any) -> Any: """Restore sidecar metadata on rank zero before any save-time eval.""" self._set_output_dir_from_training_args(args) self._is_world_process_zero = self._world_process_zero(state) self._pending_checkpoint_save_step = None self._policy_stop_requested = bool(self.policy.state.stop_requested) # Do not call ``load`` followed by a separate raw broadcast: an invalid # sidecar would make rank zero raise before workers enter that latter # collective. This helper broadcasts either the restored state or the # rank-zero error, so the whole job fails coherently. self._load_and_broadcast_policy_state_from_world_zero() self._policy_stop_requested = bool(self.policy.state.stop_requested) if self._policy_stop_requested: # ``CallbackHandler.on_train_begin`` resets this field before # invoking callbacks. Restore the terminal policy decision here # so a resumed job exits without taking another optimizer step. control.should_training_stop = True # If a process died after native checkpoint writing but before on_save, # callback state in trainer_state.json is still valid. Re-publish the # sidecar and final_checkpoint link before continuing. self._run_on_rank_zero_and_broadcast_failure( operation="committed checkpoint artifact recovery", action=self._recover_committed_checkpoint_artifacts, ) return control def on_save(self, args: Any, state: Any, control: Any, **kwargs: Any) -> Any: """Commit sidecar/link metadata only after native checkpoint writing.""" expected_step = self._pending_checkpoint_save_step if expected_step is None: return control self._pending_checkpoint_save_step = None saved_step = int(getattr(state, "global_step", expected_step)) if saved_step != expected_step: raise RuntimeError( "checkpoint evaluation received an unexpected native save event: " f"expected_step={expected_step}, saved_step={saved_step}" ) self._wait_for_native_checkpoint_write() # Native shard writes have completed on every rank. The final link # and sidecar are rank-zero-only publication work, so share a failure # from that phase before workers can advance into a later collective. self._run_on_rank_zero_and_broadcast_failure( operation="saved checkpoint publication", action=lambda: self._commit_saved_checkpoint(step=saved_step), ) return control def consume_checkpoint_evaluation(self) -> bool: """Return and clear whether the next evaluation is a save-time one.""" pending = self.pending_checkpoint_evaluation self.pending_checkpoint_evaluation = False return pending def state(self) -> dict[str, Any]: """Expose callback progress through Transformers' callback state API.""" return { "args": { "enabled": self.enabled, "patience": self.policy.patience, "tie_policy": self.policy.tie_policy, "patience_start_step": self.policy.patience_start_step, "gate_steps": list(self.gate_steps), "regular_steps": self.regular_steps, "save_total_limit": self.save_total_limit, "output_dir": self.output_dir, }, "attributes": { "pending_checkpoint_evaluation": self.pending_checkpoint_evaluation, "trigger_count": self.trigger_count, "checkpoint_eval_state": self.checkpoint_eval_state, }, } @classmethod def from_state(cls, state: Mapping[str, Any]) -> "CheckpointEvaluationTriggerCallback": args = state.get("args", {}) if isinstance(state, Mapping) else {} attrs = state.get("attributes", {}) if isinstance(state, Mapping) else {} callback = cls( enabled=bool(args.get("enabled", True)), patience=args.get("patience", 3), tie_policy=args.get("tie_policy", "reset"), patience_start_step=args.get("patience_start_step", 0), gate_steps=args.get("gate_steps"), regular_steps=args.get("regular_steps"), save_total_limit=args.get("save_total_limit", 1), output_dir=args.get("output_dir"), ) callback.pending_checkpoint_evaluation = bool(attrs.get("pending_checkpoint_evaluation", False)) callback.trigger_count = int(attrs.get("trigger_count", 0)) if "checkpoint_eval_state" in attrs: callback.checkpoint_eval_state = attrs["checkpoint_eval_state"] callback._policy_stop_requested = bool(callback.policy.state.stop_requested) return callback __all__ = [ "CheckpointEvaluationDecision", "CheckpointEvaluationPolicy", "CheckpointEvaluationState", "CheckpointEvaluationTriggerCallback", "apply_checkpoint_evaluation_decision", ]