mirror of https://github.com/razor-ai/soup.git
124 lines
4.3 KiB
Python
124 lines
4.3 KiB
Python
"""v0.53.6 #101 — Soup plugin TrainerCallback.
|
|
|
|
Bridges :mod:`soup_cli.plugins` registered hooks into the HF Trainer
|
|
callback surface. For every enabled plugin (via :func:`list_plugins` +
|
|
``spec.enabled``), discovers implemented hooks via :func:`discover_hooks`
|
|
and dispatches the matching trainer event.
|
|
|
|
Per-plugin hook exceptions are swallowed at WARNING level — one
|
|
misbehaving plugin must not crash a multi-hour training run. This
|
|
mirrors the v0.44.0 / v0.45.0 plugin loader policy.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from typing import Any
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _collect_active_hooks() -> list[tuple[str, dict[str, Any]]]:
|
|
"""Return ``[(plugin_name, hook_map), ...]`` for every enabled plugin.
|
|
|
|
Lazy-imports :mod:`soup_cli.plugins` so the callback module stays
|
|
cheap to import in CI / non-training contexts.
|
|
"""
|
|
from soup_cli.plugins import discover_hooks, list_plugins
|
|
|
|
out: list[tuple[str, dict[str, Any]]] = []
|
|
for name, spec in list_plugins().items():
|
|
if not spec.enabled:
|
|
continue
|
|
hooks = discover_hooks(spec.plugin)
|
|
if hooks:
|
|
out.append((name, hooks))
|
|
return out
|
|
|
|
|
|
def _safe_invoke(
|
|
plugin_name: str, hook_name: str, hook: Any, ctx: dict[str, Any]
|
|
) -> None:
|
|
try:
|
|
hook(ctx)
|
|
except Exception: # noqa: BLE001 — plugin failure must not crash training
|
|
logger.warning(
|
|
"Plugin %r hook %r raised; continuing",
|
|
plugin_name,
|
|
hook_name,
|
|
exc_info=True,
|
|
)
|
|
|
|
|
|
def _build_callback_class() -> type:
|
|
"""Construct the ``SoupPluginCallback`` class with transformers as parent.
|
|
|
|
Lazy-imports :mod:`transformers` so the wiring helper can be imported
|
|
in CI without the heavy dep installed.
|
|
"""
|
|
from transformers import TrainerCallback
|
|
|
|
class SoupPluginCallback(TrainerCallback):
|
|
"""Fans HF Trainer events out to every enabled Soup plugin."""
|
|
|
|
def __init__(
|
|
self, hooks: list[tuple[str, dict[str, Any]]] | None = None
|
|
) -> None:
|
|
super().__init__()
|
|
# Snapshot at construction time so a plugin registered MID-run
|
|
# does not silently start receiving hooks halfway through. The
|
|
# caller may pre-collect hooks (the ``build_plugin_callback``
|
|
# path) to avoid a redundant registry scan + close a tiny
|
|
# race-window between "is any plugin enabled?" and
|
|
# "snapshot the registry".
|
|
self._hooks = (
|
|
list(hooks) if hooks is not None else _collect_active_hooks()
|
|
)
|
|
|
|
def _dispatch(self, hook_name: str, context: dict[str, Any]) -> None:
|
|
for plugin_name, hooks in self._hooks:
|
|
hook = hooks.get(hook_name)
|
|
if hook is None:
|
|
continue
|
|
_safe_invoke(plugin_name, hook_name, hook, context)
|
|
|
|
def on_train_begin(self, args, state, control, **kwargs): # noqa: D401
|
|
self._dispatch(
|
|
"pre_train", {"args": args, "state": state, "control": control}
|
|
)
|
|
|
|
def on_train_end(self, args, state, control, **kwargs): # noqa: D401
|
|
self._dispatch(
|
|
"post_train", {"args": args, "state": state, "control": control}
|
|
)
|
|
|
|
def on_step_begin(self, args, state, control, **kwargs): # noqa: D401
|
|
self._dispatch(
|
|
"pre_step", {"args": args, "state": state, "control": control}
|
|
)
|
|
|
|
def on_step_end(self, args, state, control, **kwargs): # noqa: D401
|
|
self._dispatch(
|
|
"post_step", {"args": args, "state": state, "control": control}
|
|
)
|
|
|
|
return SoupPluginCallback
|
|
|
|
|
|
def build_plugin_callback() -> Any:
|
|
"""Return a new ``SoupPluginCallback`` instance, or ``None`` if no
|
|
enabled plugins implement any hook (no-op short-circuit).
|
|
|
|
Collects hooks ONCE and passes the snapshot to the callback so a
|
|
plugin registered between the "is empty?" check and the constructor
|
|
cannot silently slip into the active hook list — review fix.
|
|
"""
|
|
hooks = _collect_active_hooks()
|
|
if not hooks:
|
|
return None
|
|
callback_cls = _build_callback_class()
|
|
return callback_cls(hooks)
|
|
|
|
|
|
__all__ = ["build_plugin_callback"]
|