Coverage for transformer_lens/model_bridge/sources/inspect/driver.py: 94%
134 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-10 05:10 +0000
1"""InspectDriver — torch-free consumer of an ``inspect_ai`` provider's output.
3Talks to a provider through the inspect_ai ``ModelOutput`` envelope. Everything
4provider-specific — the request schema, which hooks are served, full vs last-token
5logits, intervention translation — lives in a :mod:`profiles` Profile; the driver just
6drives it. Stays numpy-only (``to_torch`` runs at the bridge boundary), so this file
7imports zero torch symbols (enforced by a unit test). ``inspect_ai`` is imported lazily.
8"""
9from __future__ import annotations
11import asyncio
12import logging
13import os
14import threading
15import warnings
16from concurrent.futures import TimeoutError as FutureTimeout
17from typing import Any, Mapping
19import numpy as np
21from transformer_lens.model_bridge.driver_protocol import (
22 ForwardResult,
23 Intervention,
24 TensorLike,
25)
26from transformer_lens.model_bridge.sources._driver_base import DriverBase
28from . import hooks, wire
29from .masks import normalize_attention_mask
30from .profiles import TLBridgeProfile
32# Cap a single provider call so a hung remote/provider forward unblocks the sync caller
33# instead of stalling it forever. Generous by default; override for slow remote backends.
34_PROVIDER_TIMEOUT_S = float(os.environ.get("TL_INSPECT_TIMEOUT_S", "300"))
37class InspectDriver(DriverBase):
38 """Driver wrapping an ``inspect_ai`` model; capture + interventions via a Profile."""
40 # Remote provider — no torch weight/grad surface.
41 _supported_features = frozenset()
43 def __init__(self, model: Any, adapter: Any, tokenizer: Any, profile: Any = None) -> None:
44 super().__init__(adapter.cfg, tokenizer)
45 self._model = model
46 self._profile = profile if profile is not None else TLBridgeProfile()
47 self._n_layers = int(self.bridge_config.n_layers)
48 self._d_vocab = int(self.bridge_config.d_vocab)
49 # Provider-specific: loss/both allowed only if the provider returns full logits.
50 self.provides_sequence_logits = self._profile.provides_sequence_logits
51 self.supported_hook_points = self._profile.supported_hooks(self._n_layers)
52 # Everything the registry could serve (boundaries + head-split) that this
53 # provider/model doesn't, plus the never-fireable set (embed, ln_final, scores).
54 universe = hooks.all_hook_points(self._n_layers)
55 self.non_fireable_hook_points = hooks.nonfireable_hook_points(self._n_layers) | (
56 universe - self.supported_hook_points
57 )
58 # Background event loop, created lazily on first forward.
59 self._loop: asyncio.AbstractEventLoop | None = None
60 self._loop_thread: threading.Thread | None = None
61 self._loop_lock = threading.Lock() # guards lazy creation / abandonment
62 self._warned_missing: set[str] = set() # hooks we've already warned were absent
64 def forward(
65 self,
66 input_ids: TensorLike | None = None,
67 *,
68 capture: tuple[str, ...] = (),
69 intervene: Mapping[str, Intervention] | None = None,
70 max_new_tokens: int = 1,
71 return_logits: bool = True,
72 attention_mask: TensorLike | None = None,
73 **kwargs: Any,
74 ) -> ForwardResult:
75 if self._model is None:
76 raise RuntimeError("InspectDriver is closed.")
77 if input_ids is None: 77 ↛ 78line 77 didn't jump to line 78 because the condition on line 77 was never true
78 raise ValueError("InspectDriver requires input_ids")
79 if int(max_new_tokens) != 1:
80 raise NotImplementedError(
81 "InspectDriver supports max_new_tokens=1 only (single-forward capture)."
82 )
83 ids = self._normalize_input_ids(input_ids)
84 mask = None
85 if attention_mask is not None:
86 if not getattr(self._profile, "supports_attention_mask", False):
87 raise NotImplementedError("This Inspect provider does not support attention_mask.")
88 mask = normalize_attention_mask(attention_mask, len(ids))
89 # capture is authoritative: the bridge passes exactly the hooks with handlers,
90 # so () means "capture nothing" (logits only), not "capture everything".
91 names = list(capture)
92 wire_keys = self._wire_keys(names)
93 interventions = self._profile.translate_interventions(
94 intervene or {}, self.supported_hook_points
95 )
96 prompt, extra_args = self._profile.build_request(
97 ids, wire_keys, interventions, return_logits, self.tokenizer
98 )
99 if mask is not None:
100 extra_args["attention_mask"] = mask
102 output = self._run_coro(self._generate(prompt, extra_args))
104 captured = self._assemble_captures(output, names)
105 logits = (
106 self._profile.decode_logits(output, len(ids), self._d_vocab, self.tokenizer)
107 if return_logits
108 else None
109 )
110 return ForwardResult(logits=logits, captured=captured, raw_output=output)
112 async def _generate(self, prompt: Any, extra_args: dict[str, Any]) -> Any:
113 from inspect_ai.model import GenerateConfig
115 config = GenerateConfig(
116 temperature=0.0, max_tokens=1, extra_body={"extra_args": extra_args}
117 )
118 return await self._model.generate(prompt, config=config)
120 def _assemble_captures(self, output: Any, names: list[str]) -> dict[str, np.ndarray]:
121 """Decode the requested hooks → ``{hook_name: (1, ...)}`` (batch dim added onto the
122 provider's batchless array — rank 2 for boundaries, 3 for head-split/pattern);
123 names the provider didn't return are skipped (and warned once)."""
124 metadata = getattr(output, "metadata", None) or {}
125 decoded = wire.decode_activations(metadata, self._wire_keys(names))
126 captured: dict[str, np.ndarray] = {}
127 missing: list[str] = []
128 for name in names:
129 resolved = hooks.resolve(name)
130 if resolved is None: 130 ↛ 131line 130 didn't jump to line 131 because the condition on line 130 was never true
131 continue
132 arr = decoded.get(hooks.wire_key(*resolved))
133 if arr is None:
134 missing.append(name)
135 continue
136 batchless = hooks.WIRE_BATCHLESS_NDIM.get(resolved[1], 2)
137 captured[name] = arr[np.newaxis, ...] if arr.ndim == batchless else arr
138 self._warn_missing(missing)
139 return captured
141 def _warn_missing(self, missing: list[str]) -> None:
142 """Warn once per hook the provider was asked for but didn't return — else its
143 cache entry is silently absent and surfaces only as a later KeyError."""
144 new = [name for name in missing if name not in self._warned_missing]
145 if new:
146 self._warned_missing.update(new)
147 warnings.warn(
148 f"InspectDriver: provider returned no activation for {sorted(new)} "
149 "(requested and in supported_hook_points); those cache keys will be absent.",
150 UserWarning,
151 stacklevel=2,
152 )
154 def close(self) -> None:
155 log = logging.getLogger("transformer_lens.inspect")
156 with self._loop_lock:
157 loop, self._loop = self._loop, None
158 thread, self._loop_thread = self._loop_thread, None
159 if loop is not None:
160 try:
161 loop.call_soon_threadsafe(loop.stop)
162 if thread is not None: 162 ↛ 164line 162 didn't jump to line 164 because the condition on line 162 was always true
163 thread.join(timeout=5)
164 loop.close()
165 except Exception as e:
166 log.debug("event-loop teardown failed during close(): %s", e)
167 self._model = None # drop the provider reference; the server owns its own lifecycle
169 # ---- helpers ----
171 def _ensure_loop(self) -> asyncio.AbstractEventLoop:
172 """A private loop on a daemon thread — works even when the caller is already
173 inside a running loop (Jupyter), unlike asyncio.run(). Locked: unlocked
174 check-then-create would leak loops under concurrent first calls."""
175 loop = self._loop
176 if loop is not None: # GIL-safe fast path; the lock only guards create/abandon
177 return loop
178 with self._loop_lock:
179 if self._loop is None:
180 self._loop = asyncio.new_event_loop()
181 self._loop_thread = threading.Thread(
182 target=self._loop.run_forever, daemon=True, name="inspect-driver-loop"
183 )
184 self._loop_thread.start()
185 return self._loop
187 def _abandon_loop(self) -> None:
188 """After a timeout the loop thread is still occupied by the hung forward (cancel
189 can't interrupt sync work), so every later call would queue behind it and time out.
190 Abandon it — the daemon thread stops once the hung call returns — and rebuild
191 lazily on the next forward."""
192 with self._loop_lock:
193 loop, self._loop, self._loop_thread = self._loop, None, None
194 if loop is not None: 194 ↛ 199line 194 didn't jump to line 199 because the condition on line 194 was always true
195 try:
196 loop.call_soon_threadsafe(loop.stop) # takes effect when the hung call ends
197 except Exception:
198 pass
199 warnings.warn(
200 "InspectDriver: abandoning the wedged event-loop thread after a provider "
201 "timeout; a fresh loop will be created on the next forward.",
202 UserWarning,
203 stacklevel=3,
204 )
206 def _run_coro(self, coro: Any) -> Any:
207 future = asyncio.run_coroutine_threadsafe(coro, self._ensure_loop())
208 try:
209 return future.result(timeout=_PROVIDER_TIMEOUT_S) # re-raises provider errors
210 except FutureTimeout:
211 future.cancel() # best-effort; can't interrupt a sync forward on the loop thread
212 self._abandon_loop()
213 raise TimeoutError(
214 f"Inspect provider call exceeded {_PROVIDER_TIMEOUT_S:.0f}s "
215 "(set TL_INSPECT_TIMEOUT_S to change) — the remote/provider forward looks "
216 "hung; unblocking the caller rather than stalling indefinitely."
217 ) from None
219 def _wire_keys(self, names: list[str]) -> list[str]:
220 """Unique ``<layer>:<kind>`` keys for the requested hook names (aliases collapse)."""
221 keys = {hooks.wire_key(*r) for r in (hooks.resolve(n) for n in names) if r is not None}
222 return sorted(keys)
224 @staticmethod
225 def _normalize_input_ids(input_ids: Any) -> list[int]:
226 """Coerce to a flat list[int] (batch_size=1 only); numpy/list/tensor, no torch import."""
227 # Duck-type a torch tensor onto CPU first — np.asarray can't read CUDA memory,
228 # and a torch-free driver can't import torch to .cpu() it.
229 if hasattr(input_ids, "detach") and hasattr(input_ids, "cpu"):
230 input_ids = input_ids.detach().cpu()
231 arr = np.asarray(input_ids)
232 if arr.ndim == 2:
233 if arr.shape[0] != 1:
234 raise NotImplementedError("InspectDriver supports batch_size=1 only.")
235 arr = arr[0]
236 return [int(x) for x in arr.tolist()]