Coverage for transformer_lens/model_bridge/driver_protocol.py: 93%
81 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""Driver protocol: the contract every model-execution backend satisfies."""
2from __future__ import annotations
4from dataclasses import dataclass, field
5from typing import Any, Callable, Mapping, Protocol, Union, runtime_checkable
7import numpy as np
8import torch
10from transformer_lens.config import TransformerBridgeConfig
13@runtime_checkable
14class TensorLike(Protocol):
15 """Quacks like a tensor: ``__array__`` + ``shape`` + ``dtype``."""
17 # Attribute types stay loose — torch.Size, plain tuple, and numpy's shape
18 # don't share a Protocol-strict supertype.
19 @property
20 def shape(self) -> Any:
21 ...
23 @property
24 def dtype(self) -> Any:
25 ...
27 def __array__(self, dtype: Any = None) -> np.ndarray:
28 ...
31InterventionFn = Callable[[TensorLike], TensorLike]
32InterventionSpec = Mapping[str, Any]
33# Drivers that can dispatch Python at the engine boundary (HF) accept
34# InterventionFn; drivers that can't (vLLM under compile, remote APIs) accept
35# InterventionSpec only.
36Intervention = Union[InterventionFn, InterventionSpec]
39@dataclass(frozen=True)
40class ForwardResult:
41 """One forward call's outputs. Tensors are native to the driver's framework."""
43 logits: TensorLike | None = None
44 captured: Mapping[str, TensorLike] = field(default_factory=dict)
45 new_tokens: TensorLike | None = None
46 # Driver's native return value (HF CausalLMOutputWithPast, vLLM
47 # RequestOutput, ...). Bridge reads driver-specific extras here.
48 raw_output: Any = None
51# Feature strings ``supports()`` may be queried with. The bridge consumes
52# "parameters" (input-device placement); the rest are informational capability
53# declarations for callers. validate_driver rejects drivers declaring strings
54# outside this set.
55KNOWN_FEATURES = frozenset(
56 {"gradients", "parameters", "state_dict", "weight_access", "intervention_callbacks"}
57)
60@runtime_checkable
61class Driver(Protocol):
62 """The forward-pass contract. Hook installation is the driver's problem.
64 ``forward`` has two dialects:
66 - **Module-replacement drivers** (TransformersDriver): hooks fire via the
67 bridge's HookPoint system during the real torch forward, so ``capture``/
68 ``intervene``/``max_new_tokens`` are not served here — conforming drivers
69 raise ``NotImplementedError`` on them rather than silently ignore.
70 - **Spec drivers** (vLLM, Inspect): no local module, so ``capture`` names
71 hook points to record and ``intervene`` carries declarative edit specs;
72 results come back in ``ForwardResult.captured``.
73 """
75 architecture: str
76 bridge_config: TransformerBridgeConfig
77 tokenizer: Any
78 supported_hook_points: frozenset[str]
79 non_fireable_hook_points: frozenset[str]
80 # False when the driver returns logits for the final position only —
81 # the bridge then refuses return_type="loss"/"both" instead of NaN-ing.
82 provides_sequence_logits: bool
84 def forward(
85 self,
86 input_ids: TensorLike | None = None,
87 *,
88 capture: tuple[str, ...] = (),
89 intervene: Mapping[str, Intervention] | None = None,
90 max_new_tokens: int = 1,
91 return_logits: bool = True,
92 **kwargs: Any,
93 ) -> ForwardResult:
94 ...
96 def close(self) -> None:
97 ...
99 def supports(self, feature: str) -> bool:
100 """Capability flag over :data:`KNOWN_FEATURES`. The bridge consults
101 "parameters"; the others are caller-facing declarations."""
102 ...
104 # Note: torch-specific surface (parameters, named_parameters, state_dict,
105 # weight access) is NOT in the protocol. Drivers that can serve those
106 # methods provide them as implementation details, gated by supports("...").
109def to_torch(t: TensorLike, *, dtype: torch.dtype | None = None) -> torch.Tensor:
110 """Convert any TensorLike to torch.Tensor at the bridge boundary.
112 Order: torch passthrough → DLPack (jax/mlx/tf/cupy/numpy≥1.22, preserves
113 device) → ``__array__`` + ``from_numpy`` (CPU only).
114 """
115 if isinstance(t, torch.Tensor):
116 return t.to(dtype) if dtype is not None else t
118 if hasattr(t, "__dlpack__"): 118 ↛ 128line 118 didn't jump to line 128 because the condition on line 118 was always true
119 try:
120 out = torch.from_dlpack(t)
121 return out.to(dtype) if dtype is not None else out
122 except (BufferError, RuntimeError, ValueError, TypeError, AttributeError):
123 # Fall through on stream-sync, missing __dlpack_device__, or
124 # version-skew failures; the numpy path either succeeds or raises
125 # informatively.
126 pass
128 arr = np.asarray(t)
129 out = torch.from_numpy(arr)
130 return out.to(dtype) if dtype is not None else out
133# Parameter names a conforming driver's forward() must accept; missing names
134# get silently swallowed by **kwargs and break the contract.
135_DRIVER_FORWARD_REQUIRED_PARAMS = frozenset(
136 ("input_ids", "capture", "intervene", "max_new_tokens", "return_logits")
137)
140def validate_driver(driver: Any, *, after_bridge_construction: bool = False) -> None:
141 """Stronger than ``isinstance(driver, Driver)``: checks types, signatures,
142 and (optionally) post-construction state.
144 Args:
145 after_bridge_construction: when True, also requires at least one of
146 ``supported_hook_points`` / ``non_fireable_hook_points`` non-empty
147 (the bridge backfills the former, so empty-on-both means the
148 driver silently degrades to "supports nothing").
150 Raises:
151 TypeError: with a message naming the contract violation.
152 """
153 _expect_attr_type(driver, "architecture", str)
154 _expect_attr_type(driver, "bridge_config", TransformerBridgeConfig)
155 if not hasattr(driver, "tokenizer"): 155 ↛ 156line 155 didn't jump to line 156 because the condition on line 155 was never true
156 raise TypeError("Driver missing required attribute: 'tokenizer'")
157 _expect_attr_type(driver, "supported_hook_points", frozenset)
158 _expect_attr_type(driver, "non_fireable_hook_points", frozenset)
159 # getattr(..., True) defaults would silently pick the UNSAFE value for
160 # loss gating, so the attribute is mandatory.
161 _expect_attr_type(driver, "provides_sequence_logits", bool)
163 declared_features = getattr(driver, "_supported_features", None)
164 if declared_features is not None: 164 ↛ 172line 164 didn't jump to line 172 because the condition on line 164 was always true
165 unknown = frozenset(declared_features) - KNOWN_FEATURES
166 if unknown:
167 raise TypeError(
168 f"Driver._supported_features contains unknown feature strings "
169 f"{sorted(unknown)}; known features: {sorted(KNOWN_FEATURES)}."
170 )
172 overlap = driver.supported_hook_points & driver.non_fireable_hook_points
173 if overlap:
174 raise TypeError(
175 f"Driver.supported_hook_points and Driver.non_fireable_hook_points "
176 f"overlap on {sorted(overlap)[:3]}; a hook is either fireable or not."
177 )
178 for name in driver.supported_hook_points | driver.non_fireable_hook_points:
179 if not isinstance(name, str):
180 raise TypeError(f"Hook-point names must be str; got {type(name).__name__}: {name!r}")
182 forward = getattr(driver, "forward", None)
183 if not callable(forward): 183 ↛ 184line 183 didn't jump to line 184 because the condition on line 183 was never true
184 raise TypeError("Driver.forward must be callable")
185 import inspect
187 sig = inspect.signature(forward)
188 params = sig.parameters
189 has_var_keyword = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values())
190 missing = [p for p in _DRIVER_FORWARD_REQUIRED_PARAMS if p not in params]
191 if missing and not has_var_keyword:
192 raise TypeError(
193 f"Driver.forward must accept parameters {sorted(_DRIVER_FORWARD_REQUIRED_PARAMS)}; "
194 f"missing {sorted(missing)} (and no **kwargs to absorb them)."
195 )
197 if not callable(getattr(driver, "close", None)): 197 ↛ 198line 197 didn't jump to line 198 because the condition on line 197 was never true
198 raise TypeError("Driver.close must be callable")
200 if after_bridge_construction:
201 if not driver.supported_hook_points and not driver.non_fireable_hook_points:
202 raise TypeError(
203 "Driver has empty supported_hook_points AND non_fireable_hook_points "
204 "after bridge construction. Drivers must declare at least one — the "
205 "bridge backfills supported from registry minus non_fireable, but "
206 "empty-on-both means there's no contract for downstream code."
207 )
210def _expect_attr_type(obj: Any, name: str, expected: type) -> None:
211 if not hasattr(obj, name):
212 raise TypeError(f"Driver missing required attribute: {name!r}")
213 value = getattr(obj, name)
214 if not isinstance(value, expected):
215 raise TypeError(f"Driver.{name} must be {expected.__name__}; got {type(value).__name__}.")
218__all__ = [
219 "Driver",
220 "ForwardResult",
221 "KNOWN_FEATURES",
222 "Intervention",
223 "InterventionFn",
224 "InterventionSpec",
225 "TensorLike",
226 "to_torch",
227 "validate_driver",
228]