Fix dynamic rollout hook module registration

agentlightning/hooks.py

@@ -4,6 +4,8 @@
 
 from __future__ import annotations
 
+import sys
+from hashlib import sha256
 from typing import TYPE_CHECKING, Any, Protocol
 
 from agentlightning.schemas import RolloutCreate
@@ -43,10 +45,20 @@ def load_hooks(path: str) -> RolloutHooks:
     if not module_path.exists():
         raise FileNotFoundError(f"Hooks module not found: {module_path}")
 
-    spec = importlib.util.spec_from_file_location("_agl_hooks", str(module_path))
+    module_name = f"_agl_hooks_{sha256(str(module_path).encode()).hexdigest()[:16]}"
+    spec = importlib.util.spec_from_file_location(module_name, str(module_path))
     assert spec is not None and spec.loader is not None
     module = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(module)
+    previous_module = sys.modules.get(module_name)
+    sys.modules[module_name] = module
+    try:
+        spec.loader.exec_module(module)
+    except BaseException:
+        if previous_module is None:
+            sys.modules.pop(module_name, None)
+        else:
+            sys.modules[module_name] = previous_module
+        raise
 
     hook_classes = [
         obj

tests/test_hooks.py

@@ -0,0 +1,43 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Regression coverage for dynamically loaded rollout hooks."""
+
+import pickle
+import sys
+
+from agentlightning.hooks import load_hooks
+from agentlightning.schemas import RolloutCreate
+
+
+def test_load_hooks_registers_dataclass_hook_module(tmp_path) -> None:
+    hooks_path = tmp_path / "custom_hooks.py"
+    hooks_path.write_text(
+        """
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from agentlightning.hooks import RolloutHooks
+from agentlightning.schemas import RolloutCreate
+
+
+@dataclass
+class HookConfig:
+    prefix: str = "checked:"
+
+
+class ConfiguredHooks(RolloutHooks):
+    def __init__(self) -> None:
+        self.config = HookConfig()
+
+    def on_enqueue(self, request: RolloutCreate) -> RolloutCreate:
+        return request.model_copy(update={"input": f"{self.config.prefix}{request.input}"})
+""",
+        encoding="utf-8",
+    )
+
+    hooks = load_hooks(str(hooks_path))
+
+    assert hooks.on_enqueue(RolloutCreate(input="task")).input == "checked:task"
+    assert hooks.__class__.__module__ in sys.modules
+    assert pickle.loads(pickle.dumps(hooks)).config.prefix == "checked:"