Custom hook dataclasses with postponed annotations fail during module loading

[Linxiushen]

Setting agentlightning.hooks to a Python file that defines a dataclass with postponed annotations fails before the hook can run. load_hooks executes the file without first registering its module in sys.modules; the dataclass decorator needs that registration to resolve string annotations.

Reproduction on current main (ff9457587fb6ec900e16e93be9ad2d77409afa08), with the locked development dependencies and Python 3.12.13:

Save this as custom_hooks.py:

from __future__ import annotations

from dataclasses import dataclass

from agentlightning.hooks import RolloutHooks


@dataclass
class Settings:
    label: str = "training"


class Hooks(RolloutHooks):
    settings = Settings()

Then run:

from agentlightning.hooks import load_hooks

hooks = load_hooks("custom_hooks.py")

Expected: a Hooks instance is returned. Actual: AttributeError: 'NoneType' object has no attribute '__dict__' from dataclasses._is_type, at sys.modules.get(cls.__module__).__dict__.

The trainer's _ensure_hooks uses this same loader for agentlightning.hooks. A hook file without postponed annotations loads successfully. Classes from loaded hook files also cannot currently be pickled because their _agl_hooks module cannot be imported.

I have a focused fix that registers the module before execution, gives different resolved file paths separate module identities, and restores the previous registration if loading fails. Regression tests exercise real on-disk hook files and RolloutCreate transformations. AI assistance was used to investigate and prepare the fix.

[hzy46]

Hi 林SO (@Linxiushen) , is there any specific reason that you want to use hooks?