Run workflows through managed hook sessions

docs/source/start.rst

@@ -51,6 +51,41 @@ Gradients reads ``("baseline", "input")`` and writes
 ``("attr", "input")``. Here, ``attributions.shape`` is ``(1, 4)``, matching
 the input.
 
+Interactive operations across a workflow
+----------------------------------------
+
+Use :meth:`~tdhook.workflow.Workflow.session` when captures or interventions
+must wrap every model execution in a declared workflow. The managed object
+exposes the normal :class:`~tdhook.session.HookSession` operations and runs the
+bound workflow without requiring a second, manually nested context:
+
+.. code-block:: python
+
+   from tdhook.latent import ActivationCaching
+   from tdhook.targets import Target
+   from tdhook.workflow import Workflow
+
+   workflow = Workflow(
+       ActivationCaching("0", cache_key=("activations", "first")),
+       ActivationCaching("2", cache_key=("activations", "second")),
+   )
+
+   with workflow.session(model) as session:
+       observed = session.capture(Target("0", "activation", -1, (0,)))
+       execution = session.run(data)
+
+   execution.plan       # the exact WorkflowPlan that ran
+   execution.program    # the HookProgram applied around that plan
+   observed.values      # one entry per matching model execution, in call order
+
+Session operations wrap the complete run; they are not workflow steps and do
+not change planning or co-execution decisions. If
+:meth:`~tdhook.session.HookSession.stop` reaches its target, it aborts the
+workflow run, so no complete workflow result is returned. Its
+:class:`~tdhook.session.EarlyStopResult` still exposes the partial module
+output, and the surrounding managed context restores workflow hooks, session
+hooks, and temporary model state.
+
 Where to go next
 ----------------
 

src/tdhook/session.py

@@ -2,8 +2,8 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass
-from typing import Literal
+from dataclasses import dataclass, field
+from typing import Literal, Self
 import weakref
 
 import torch
@@ -19,9 +19,19 @@
 
 @dataclass
 class CapturedTarget:
-    """Mutable result populated when a capture hook observes its target."""
+    """Mutable result populated whenever a capture hook observes its target.
+
+    ``value`` retains the most recent observation for compatibility with a
+    single model execution. ``values`` preserves every observation in call
+    order, so repeated executions remain distinguishable.
+    """
 
     value: Tensor | None = None
+    values: list[Tensor] = field(default_factory=list)
+
+    def _record(self, value: Tensor) -> None:
+        self.value = value
+        self.values.append(value)
 
 
 @dataclass
@@ -60,7 +70,7 @@ def program(self) -> HookProgram:
 
         return self._builder.program if self._builder is not None else self._program
 
-    def __enter__(self) -> "HookSession":
+    def __enter__(self) -> Self:
         if self._builder is not None:
             raise RuntimeError("Cannot enter a HookSession twice")
         self._model()
@@ -95,15 +105,15 @@ def capture(self, target: Target, *, prepend: bool = False) -> CapturedTarget:
 
         if target.kind == "parameter":
             parameter = module.get_parameter(target.parameter)  # type: ignore[arg-type]
-            captured.value = target._select(parameter).detach().clone()
+            captured._record(target._select(parameter).detach().clone())
             builder.record(spec)
         else:
 
             def forward_hook(_module: nn.Module, _args: tuple[object, ...], value: object):
-                captured.value = target.select_output(value).detach().clone()
+                captured._record(target.select_output(value).detach().clone())
 
             def gradient_hook(_module: nn.Module, values: tuple[Tensor | None, ...]):
-                captured.value = target.select_output(values).detach().clone()
+                captured._record(target.select_output(values).detach().clone())
 
             builder.register(
                 module,

src/tdhook/workflow.py

@@ -16,6 +16,7 @@
 from tdhook.execution import AutogradLifetime, ExecutionSpec, GradientMode
 from tdhook.descriptions import ConfiguredStepDescription
 from tdhook.runtime import HookProgram
+from tdhook.session import HookSession
 
 
 @runtime_checkable
@@ -100,11 +101,13 @@ class WorkflowResult:
 
     Use :meth:`Workflow.run_with_plan` when consumers need execution metadata.
     :meth:`Workflow.run` continues to return only the TensorDict for existing
-    callers.
+    callers. ``program`` identifies the imperative operations that wrapped a
+    managed :class:`WorkflowSession` run and is ``None`` for direct execution.
     """
 
     data: TensorDictBase
     plan: WorkflowPlan
+    program: HookProgram | None = None
 
 
 @dataclass(frozen=True)
@@ -581,10 +584,42 @@ def run(self, model: nn.Module, data: TensorDictBase) -> TensorDictBase:
 
         return self.run_with_plan(model, data).data
 
+    def session(self, model: nn.Module) -> "WorkflowSession":
+        """Bind this workflow and ``model`` to one managed hook session.
+
+        Operations registered on the returned session wrap the complete
+        workflow run without participating in planning or co-execution.
+        """
+
+        return WorkflowSession(self, model)
+
     def __call__(self, model: nn.Module, data: TensorDictBase) -> TensorDictBase:
         return self.run(model, data)
 
 
+class WorkflowSession(HookSession):
+    """A :class:`HookSession` whose operations wrap a complete workflow run."""
+
+    def __init__(self, workflow: Workflow, model: nn.Module):
+        super().__init__(model)
+        self._workflow = workflow
+
+    def run(self, data: TensorDictBase) -> WorkflowResult:
+        """Execute the bound workflow and associate its plan with this program.
+
+        The session must be active. Managed early stopping aborts the workflow
+        run and exits the surrounding context, leaving partial results on the
+        corresponding :class:`~tdhook.session.EarlyStopResult`.
+        """
+
+        self._active_state()
+        result = self._workflow.run_with_plan(self._model(), data)
+        return WorkflowResult(data=result.data, plan=result.plan, program=self.program)
+
+    def __call__(self, data: TensorDictBase) -> WorkflowResult:
+        return self.run(data)
+
+
 __all__ = [
     "CompatibilityDecision",
     "MethodBinding",
@@ -593,5 +628,6 @@ def __call__(self, model: nn.Module, data: TensorDictBase) -> TensorDictBase:
     "WorkflowMethod",
     "WorkflowPlan",
     "WorkflowResult",
+    "WorkflowSession",
     "WorkflowUpdate",
 ]

tests/test_session.py

@@ -22,6 +22,7 @@ def test_session_captures_and_replaces_an_activation(default_test_model):
         program = session.program
 
     assert captured.value is not None
+    assert captured.values == [captured.value]
     assert captured.value.shape == (3, 1)
     assert not torch.allclose(modified, baseline)
     assert torch.allclose(default_test_model(x), baseline)
@@ -33,6 +34,20 @@ def test_session_captures_and_replaces_an_activation(default_test_model):
     )
 
 
+def test_session_preserves_repeated_captures_in_call_order(default_test_model):
+    target = Target("linear1", "activation", -1, (0,))
+    inputs = (torch.zeros(1, 10), torch.ones(1, 10))
+
+    with HookSession(default_test_model) as session:
+        captured = session.capture(target)
+        for value in inputs:
+            default_test_model(value)
+
+    assert len(captured.values) == 2
+    assert captured.value is captured.values[-1]
+    assert not torch.equal(captured.values[0], captured.values[1])
+
+
 def test_session_preserves_single_tensor_tuple_output():
     class TupleModule(nn.Module):
         def forward(self, x):

tests/test_workflow.py

@@ -10,9 +10,16 @@
 from tdhook.execution import AutogradLifetime, ExecutionSpec, GradientMode
 from tdhook.latent import ActivationCaching, Probing
 from tdhook.modules import HookedModule
-from tdhook.runtime import BoundHookProgram, HookProgramBuilder, HookSpec
+from tdhook.runtime import BoundHookProgram, HookProgram, HookProgramBuilder, HookSpec
 from tdhook.targets import Target
-from tdhook.workflow import PlannedExecution, Workflow, WorkflowResult, WorkflowUpdate, _DeferredAutogradCleanup
+from tdhook.workflow import (
+    PlannedExecution,
+    Workflow,
+    WorkflowResult,
+    WorkflowSession,
+    WorkflowUpdate,
+    _DeferredAutogradCleanup,
+)
 
 
 class CountingModel(nn.Module):
@@ -163,10 +170,77 @@ def test_workflow_returns_the_plan_used_by_execution(default_test_model):
 
     assert isinstance(result, WorkflowResult)
     assert result.plan.model_passes == 1
+    assert result.program is None
     assert result.data["output"].shape == (2, 5)
     assert workflow.run(default_test_model, data).shape == data.shape
 
 
+def test_managed_workflow_session_applies_operations_to_every_model_execution(default_test_model):
+    workflow = Workflow(CaptureOutput(), ReplaceOutput())
+    data = TensorDict({"input": torch.ones(2, 10)}, batch_size=[2])
+    target = Target("", "activation", -1, (0,))
+
+    with workflow.session(default_test_model) as session:
+        assert isinstance(session, WorkflowSession)
+        session.replace(target, 0)
+        captured = session.capture(target)
+        result = session.run(data)
+
+    assert result.plan.model_passes == 2
+    assert result.program == HookProgram(
+        (
+            HookSpec("", "replace", "fwd", target=target),
+            HookSpec("", "capture", "fwd", target=target),
+        )
+    )
+    assert len(captured.values) == 2
+    assert all(torch.equal(value, torch.zeros(2, 1)) for value in captured.values)
+    assert all(not module._forward_hooks for module in default_test_model.modules())
+
+
+def test_managed_workflow_session_early_stop_aborts_the_run_and_restores_hooks(default_test_model):
+    method = CaptureOutput()
+    workflow = Workflow(method)
+    data = TensorDict({"input": torch.ones(2, 10)}, batch_size=[2])
+
+    with workflow.session(default_test_model) as session:
+        stopped = session.stop("linear1")
+        session.run(data)
+        pytest.fail("managed early stopping must abort the workflow run")
+
+    assert stopped.reached
+    assert stopped.output is not None
+    assert method.values == []
+    assert session.program == HookProgram((HookSpec("linear1", "stop", "fwd"),), stopped_at="linear1")
+    assert all(not module._forward_hooks for module in default_test_model.modules())
+
+
+def test_managed_workflow_session_restores_state_after_workflow_failure(default_test_model):
+    def fail(value):
+        raise RuntimeError("workflow failed")
+
+    failing = TensorDictModule(fail, in_keys=["output"], out_keys=["unused"])
+    workflow = Workflow(CaptureOutput(), failing)
+    data = TensorDict({"input": torch.ones(2, 10)}, batch_size=[2])
+    target = Target("linear1", "parameter", 0, (0,), parameter="weight")
+    original = default_test_model.linear1.weight.detach().clone()
+
+    with pytest.raises(RuntimeError, match="workflow failed"):
+        with workflow.session(default_test_model) as session:
+            session.replace(target, -3)
+            session.run(data)
+
+    assert torch.equal(default_test_model.linear1.weight, original)
+    assert all(not module._forward_hooks for module in default_test_model.modules())
+
+
+def test_managed_workflow_session_requires_its_context(default_test_model):
+    session = Workflow().session(default_test_model)
+
+    with pytest.raises(RuntimeError, match="active context"):
+        session.run(TensorDict())
+
+
 def test_executed_plan_is_not_stale_after_a_preflight_plan(default_test_model):
     class ChangesAfterPreflight(HookingContextFactory):
         def __init__(self):