103 lines
3.2 KiB
Python
103 lines
3.2 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import Callable
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
import torch.nn as nn
|
|
|
|
|
|
class BaseState:
|
|
"""Base class for hook state containers."""
|
|
|
|
def reset(self) -> None: # pragma: no cover - default is no-op
|
|
pass
|
|
|
|
|
|
class StateManager:
|
|
"""Manage per-context hook state instances."""
|
|
|
|
def __init__(self, state_cls: Callable[[], BaseState]):
|
|
self._state_cls = state_cls
|
|
self._states: dict[str, BaseState] = {}
|
|
self._context: str = "default"
|
|
|
|
def set_context(self, name: str) -> None:
|
|
self._context = name or "default"
|
|
|
|
def get_state(self) -> BaseState:
|
|
if self._context not in self._states:
|
|
self._states[self._context] = self._state_cls()
|
|
return self._states[self._context]
|
|
|
|
def reset(self) -> None:
|
|
self._states.clear()
|
|
|
|
|
|
class ModelHook:
|
|
"""Base class for model hooks that can override a module's forward."""
|
|
|
|
def initialize_hook(self, module: nn.Module) -> nn.Module:
|
|
raise NotImplementedError
|
|
|
|
def new_forward(self, module: nn.Module, *args: Any, **kwargs: Any):
|
|
raise NotImplementedError
|
|
|
|
def reset_state(self, module: nn.Module) -> nn.Module:
|
|
raise NotImplementedError
|
|
|
|
|
|
@dataclass
|
|
class _WrappedForward:
|
|
module: nn.Module
|
|
|
|
def __call__(self, *args: Any, **kwargs: Any):
|
|
registry: HookRegistry | None = getattr(self.module, "_hook_registry", None)
|
|
if registry is None or not registry._hooks:
|
|
return self.module._original_forward(*args, **kwargs)
|
|
return registry.dispatch(*args, **kwargs)
|
|
|
|
|
|
class HookRegistry:
|
|
"""Registry of hooks attached to a module."""
|
|
|
|
def __init__(self, module: nn.Module):
|
|
self.module = module
|
|
self._hooks: dict[str, ModelHook] = {}
|
|
|
|
@classmethod
|
|
def get_or_create(cls, module: nn.Module) -> HookRegistry:
|
|
registry: HookRegistry | None = getattr(module, "_hook_registry", None)
|
|
if registry is None:
|
|
registry = cls(module)
|
|
setattr(module, "_hook_registry", registry)
|
|
|
|
# Wrap module.forward once so hooks can intercept calls.
|
|
if not hasattr(module, "_original_forward"):
|
|
module._original_forward = module.forward # type: ignore[attr-defined]
|
|
module.forward = _WrappedForward(module) # type: ignore[assignment]
|
|
|
|
return registry
|
|
|
|
def register_hook(self, name: str, hook: ModelHook) -> None:
|
|
hook.initialize_hook(self.module)
|
|
self._hooks[name] = hook
|
|
|
|
def get_hook(self, name: str) -> ModelHook | None:
|
|
return self._hooks.get(name)
|
|
|
|
def dispatch(self, *args: Any, **kwargs: Any):
|
|
# For now we support a single active hook and call it directly.
|
|
# This can be extended to a chain if needed.
|
|
if not self._hooks:
|
|
return self.module._original_forward(*args, **kwargs) # type: ignore[attr-defined]
|
|
# Deterministic order: sort by name.
|
|
name = sorted(self._hooks.keys())[0]
|
|
hook = self._hooks[name]
|
|
return hook.new_forward(self.module, *args, **kwargs)
|
|
|
|
def reset_hook(self, name: str) -> None:
|
|
hook = self._hooks.get(name)
|
|
if hook is not None:
|
|
hook.reset_state(self.module)
|