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