diff --git a/docs/developing/fast-lane.md b/docs/developing/fast-lane.md new file mode 100644 index 000000000..9c62f8d77 --- /dev/null +++ b/docs/developing/fast-lane.md @@ -0,0 +1,363 @@ +# The in-process fast lane — design + +**Status:** Implemented (first slice) · **Date:** 2026-06-13 · **Branch:** `worktree-mediator-sandbox` +**Builds on:** [mediator-gpu-trace-integration.md](mediator-gpu-trace-integration.md) (the GPU-worker isolation harness this adds a tier to). + +## 1. Why — isolation can't run the interp majority + +The GPU-worker isolation contains footguns by running each user intervention in a spawned process +that holds a **path-only, weightless mirror** of the model: `isolation._WorkerPersistent` synthesizes a +bare `nn.Module()` (only a `__path__`, no parameters, no `forward`) for every `Module:`. That is +fine for reading/swapping/saving *delivered activations*, but the actual interpretability workloads — +checked against the `interp-serve-bench` taxonomy — do more than that **inside the trace body**: + +```python +normed = model.transformer.ln_f(hidden) # call the host's real final-norm module +logits = F.linear(normed, model.lm_head.weight) # read the host's real unembed WEIGHTS +``` + +Every logit-lens, steering, ablation, activation-patching, and attribution cell reads host weights +and/or calls host modules. In the worker, `head.weight` → `AttributeError` and `head(x)` → +`NotImplementedError` on the weightless dummy. So **as it stood, isolation could not run any of the +weight-reading interp majority** — the workflow's compatibility pass found this the single largest gap +(9 of the cataloged workloads, all blocked on the same weight/host-module surface). + +The fix is not to ship weights into the worker (a per-trace cost, and a partial answer). It is to run +the confirmed-safe interventions **in-process**, where the real model and its weights already live, and +isolate only the code that genuinely needs containing. That is the fast lane. + +## 2. The three tiers + +`isolate_mediators()` now classifies each mediator once, at fork in `Mediator.start`: + +| tier | when | how it runs | +|---|---|---| +| **FAST** | the effective code is all whitelisted ops / host-model access / nnsight primitives | in-process daemon thread (`_iso=None`) — full model + weights, no worker, no per-hook channel; a watchdog bounds runaway loops | +| **ISOLATE** | anything unconfirmable (unresolved global call, import, `while`, unrecoverable closure, unknown node) | the existing GPU worker | +| **REJECT** | an introspection escape (`__globals__` / `getattr` / `eval` / …) | raise `FastLaneRejected` | + +The conservative default is **ISOLATE**: absence of proof is not proof of safety; only explicitly +whitelisted code reaches FAST. Default behavior is preserved — isolation **off** never consults the +gate; isolation **on** only ever moves a mediator *isolate → in-process*. + +## 3. Detect & confirm — the classifier (`fastlane.py`) + +A **fail-closed, default-deny** static walk over the **effective code**: the trace body PLUS every user +closure it calls, resolved through the capturing frame, the function's `__globals__`, and its closure +cells. This closure resolution is load-bearing — the harness wraps real compute in `build()` / +`capture()` / `patch()` closures, so a walk of the `with` block alone sees only `build().save()`, an +opaque call. Confirming on the with-block while the footgun lives in an unresolved closure would be the +exact false-safe the gate must never produce; an unrecoverable closure fails closed to ISOLATE. + +The walk (the rules, in `_Walker`): + +- **Node allowlist.** Permit assignment / expr / return / `if` / `for` / `with` / `raise` / `pass` and + pure expression nodes. `while` → ISOLATE (cannot be statically bounded). `import` → ISOLATE (ambient + authority). `try` → ISOLATE. `global`/`nonlocal` → REJECT. Unknown node type → ISOLATE. +- **Call targets.** Each call must resolve to: a pure-compute op (`torch` / `torch.nn.functional` / + `math` / `operator` / `numpy`), an nnsight primitive, a host `Envoy`/`nn.Module`/`Tensor` (calling it + runs the real model — the fast lane's whole point), a recursively confirmed user function, a safe + builtin, or a local/parameter name (a host object bound at a walked call site). An unresolved global + call → ISOLATE (unknown authority); `torch.load`/`save`/`hub`/`jit` → ISOLATE (fs/net/JIT). +- **Introspection ban (REJECT).** `getattr`/`setattr`/`eval`/`exec`/`compile`/`__import__`/`globals`/…, + any dunder attribute (`x.__class__`), and a dunder/private subscript key (`d["__builtins__"]`). +- **Host-state writes (ISOLATE).** An attribute store other than the nnsight boundary writes + (`.output`/`.input`/`.grad`) mutates state visible to sibling mediators. In-place + (`hidden[:] = …`) and replacement boundary writes are *allowed* (and in-place is **safer** here than + under isolation, where clone-on-receive silently no-ops it). +- **Backward detection** is closure-aware (`with x.backward():` found through closures), replacing the + old `.backward(` source substring that was blind to a backward hidden in a closure — this flag now + feeds the *isolated* job's gradient-retention too. + +Verdicts are cached by intervention code-object identity, so re-running the same trace pays the walk +once. The classifier is GPU-free and unit-tested (`test_fastlane_classifier.py`) on the real workload +shapes and on **renamed** module structures (`decoder_blocks`/`final_norm`/`output_projection`) so +nothing is keyed to GPT-2 naming. + +### Threat model (the contract) + +This is **not** an adversarial sandbox — a determined author can defeat any in-process restriction +(the pysandbox negative result). Under the harness's relaxed "contain footguns, not adversaries" model +the gate confirms the effective code: introduces no ambient authority (no import, no introspection, no +unresolved global call), has no unbounded loop, writes no host state, and is composed only of +whitelisted ops / host access / confirmed user code. It is cordoned to `trust="local"` provenance and a +`CONFIG.APP.FAST_LANE` flag; anything deserialized/remote, or with the flag off, isolates wholesale. +OOM and device-side asserts in pure tensor math are knowingly traded to the in-process tier (the same +risk a user running without isolation already accepts); a deployment that cannot tolerate them sets +`fast_lane=False`. + +## 4. The watchdog + +The static gate bans `while` and confirms loops are over bounded iterables, but a *huge* bounded +`range(10**12)` passes statically yet would hang the host — the loop-containment that turning isolation +on implies. A best-effort wall-clock `Watchdog` injects a `FastLaneTimeout` into the fast-lane thread at +its next bytecode if it overruns `fast_lane_timeout`. Because the intervention body is already wrapped +in `try/except → mediator.exception` (invoker.compile), the injected exception routes through the normal +path and the host re-raises it cleanly — no channel hang. It cannot preempt a wedged native/CUDA call +(only the worker-process kill can); the bounded-loop static rule is the primary defense, this is the +backstop. Armed after thread start, disarmed in the thread's `finally` and again at `cancel`. + +## 5. Prior art it is built on + +| system | tier | borrowed idea | +|---|---|---| +| Cloudflare Workers (V8 isolates) | fast-in-process | many tenants in one process; safety from a capability-restricted runtime, not a process boundary — the fast lane's model | +| RestrictedPython (Zope) | static-confirm | a `RestrictingNodeTransformer`-style default-deny AST pass as the gate | +| SES / Hardened JS (Endo) | fast-in-process | confirmed code runs against an explicitly granted namespace; no ambient authority | +| torch.fx / JAX tracing | static-confirm | a leaf/atomic-op allowlist; "prove the body is a pure function over provided tensors" — fail, don't guess | +| gVisor | hybrid | cost is the boundary *crossing*, not the work — minimize crossings (the fast lane removes them entirely) | +| AWS Lambda + Firecracker / SnapStart | heavy-isolated | what the existing GPU worker *is* — the slow lane you keep for unconfirmable code; warm-restore to amortize spawn (the warm pool) | +| pysandbox (negative result) | — | the contract: in-process restriction is a footgun selector, never an adversarial boundary → the `trust="local"` cordon | + +The unifying pattern across all of them: **confirm once up front, then run free** behind a small, +explicit, enumerable fallback to the heavy tier. + +## 6. Part 2 — declarative primitives + +The fast lane already runs the weight-reading workloads in-process, so they need no new APIs *for the +fast lane*. The declarative primitives' value is (a) making the workloads runnable on the **isolated** +tier too, and (b) collapsing common raw-compute patterns into named calls the gate whitelists trivially: + +| primitive | signature | taxonomy primitive | role | status | +|---|---|---|---|---| +| `tracer.unembed` | `(residual, norm, head, formulation="weight") → logits` | host-weight read + module call | the projection every logit-lens / steering-direction / attribution metric does | **built** | +| `tracer.steer` | `(envoy, direction, alpha=1.0)` | boundary write (injection) | always a replacement swap — fixes in-place's silent no-op under isolation | **built** | +| `tracer.patch` | `(envoy, value)` | boundary write (transplant) | whole-tuple replacement | **built** | +| `tracer.ablate` | `(envoy, mode="zero")` | boundary write (injection) | zero/mean knockout | **built** | +| `x.carry()` / `nnsight.carry(x)` | `(value) → value` | read + run↔run transfer | session cross-trace handoff under isolation (and `.save()` now crosses too) | **built** | + +Most mirror the existing `tracer.cache()` shape: in-process they resolve the real envoys and run +directly (the fast-lane execution); isolated they ship a spec via a new event whose host handler runs the +real op host-side. The exception is `tracer.steer` (and the boundary-write `patch`/`ablate`): steering +touches **no host weights** — only the *delivered* activation — so it needs no host round-trip and no new +event. It rides the existing `Events.SWAP`: a *replacement* write (assign `envoy.output`) ships the +steered value back, which is what makes it cross the boundary where an in-place `[:] =` silently no-ops. + +### `tracer.unembed` — host-routed readout (built, 2026-06-14) + +The first primitive, and the one that closes the standing weight-read blocker on the **isolated** tier. +`tracer.unembed(residual, norm, head)` projects a residual through the final norm + unembed: + +- **In-process / fast lane** (`mediator._isolated_worker` is False): runs the real modules directly — + `F.linear(norm(residual), head.weight)` (or `head(norm(residual))` with `formulation="module"`). + This is the same compute the workloads write by hand; `tracer.unembed` just names it. +- **Isolated worker** (weightless dummies): ships the residual VALUE plus the module **paths** + (`{norm_path, head_path, formulation}`) via a new `Events.UNEMBED` request. The host's + `handle_unembed_event` resolves the real envoys (`path_to_envoy`), runs the real norm + unembed on the + **real weights**, and ships back only the logits over the bounce buffer (clone-on-receive both ways, + like any VALUE/BACKWARD round-trip). **Weights never cross the boundary** — so the readout works on + the isolated tier without binding the generic warm worker to a model or placing host weight memory in + the less-trusted worker (the two costs that ruled out shipping/sharing weights — see §1, §7). Paths + resolve through renames host-side (the wire path is always the real path), so renamed models work. + +Touch points: `Events.UNEMBED`; `handle()` dispatch + `handle_unembed_event` (interleaver.py); the +`tracer.unembed` method with the isolated/in-process branch (tracer.py). Shaped exactly like +`Events.CACHE` / `handle_cache_event`. + +**Verified (`test_isolated_unembed.py`, all under forced isolation `fast_lane=False`, gpt2 + renamed +model):** single-layer, 3-layer-interleaved, `formulation="module"`, `norm=None`, and renamed-model +readouts all isolated-vs-in-process `max|Δ|=0`; `tracer.unembed` equals the manual +`F.linear(norm(x), head.weight)` it replaces. + +### `tracer.steer` — replacement-swap injection (built, 2026-06-15) + +`tracer.steer(envoy, direction, alpha=1.0)` adds `alpha * direction` to a module's output residual — +the activation-steering / injection every steering cell does. It is the simplest of the part-2 +primitives and the structural opposite of `unembed`: it touches **no host weights**, only the delivered +activation, so it needs **no new event and no isolated/in-process branch**. The method just performs a +*replacement* boundary write: + +```python +out = envoy.output +hidden = out[0] if isinstance(out, tuple) else out +steered = hidden + alpha * direction.to(dtype=hidden.dtype, device=hidden.device) +envoy.output = (steered, *out[1:]) if isinstance(out, tuple) else steered +``` + +The eproperty setter routes that assignment through `mediator.swap` → `Events.SWAP` on **either** tier: +in-process it swaps into the batcher directly; isolated, the worker ships the steered value back over the +existing SWAP path (`pack_cuda` walks the tuple, carrying a `None` tail through untouched). The same code +is therefore correct in-process, on the fast lane, and in the isolated worker. + +The point is the **replacement** swap. The hand-written additive form is in-place +(`block.output[:, -1, :] += direction` — the canonical nnsight steering); under isolation that mutates +only the worker's *delivered clone*, no SWAP fires, the host's real activation is untouched, and the +steering silently no-ops (the save of the steered residual even looks right — only the downstream forward +reveals nothing changed). `tracer.steer` makes the steering cross the boundary by construction. Tuple +outputs (most attention modules — `(tensor, None)` here) are replaced whole, steering element `[0]`. + +Touch points: just the `tracer.steer` method (tracer.py). No event, no host handler — it reuses +`Events.SWAP` and the eproperty setter. + +**Verified (`test_isolated_steer.py`, gpt2 + renamed model):** steering one block, an attention tuple +output, and three blocks at once are all isolated-vs-in-process `max|Δ|=0` and propagate through later +layers; `tracer.steer` equals the manual untuple + whole-tuple replacement. The crux case proves the +motivation under forced isolation: the in-place form leaves the downstream residual == the unsteered +baseline (silent no-op) while `tracer.steer` changes it (steering took effect) to exactly the in-process +result. + +### `tracer.patch` — replacement-swap transplant (built, 2026-06-20) + +`tracer.patch(envoy, value)` replaces a module's output residual with a precomputed `value` — the +activation-patching / resampling transplant every patching cell does. It is the structural twin of +`tracer.steer`: a boundary write that touches **no host weights** (only the *delivered* activation is +replaced), so it needs **no new event and no isolated/in-process branch** and rides the existing +`Events.SWAP`. The only difference from `steer` is the source of the new value — `steer` computes it from +the delivered activation (`hidden + alpha*direction`), `patch` takes it from the caller: + +```python +out = envoy.output +hidden = out[0] if isinstance(out, tuple) else out +value = value.to(dtype=hidden.dtype, device=hidden.device) +envoy.output = (value, *out[1:]) if isinstance(out, tuple) else value +``` + +The point, as with `steer`, is the **replacement** swap. The hand-written form is in-place +(`block.output[0][:] = clean_act`); under isolation that mutates only the worker's *delivered clone*, no +SWAP fires, and the host's real activation is untouched — the transplant silently no-ops. `tracer.patch` +ships the value back over SWAP, so it crosses the boundary by construction. The value is cast to the +residual's dtype/device, so a value precomputed **on CPU** — the isolation-relevant case, where the +clean/source activation is captured in a prior run outside the trace — transplants cleanly. `value` +replaces the residual whole (element `[0]` for tuple outputs); partial patches construct a full-shape +value (clone + edit a slice), the standard nnsight idiom. + +Touch points: just the `tracer.patch` method (tracer.py). No event, no host handler — it reuses +`Events.SWAP` and the eproperty setter. + +**Verified (`test_isolated_patch.py`, gpt2 + renamed model):** transplanting into one block, an attention +tuple output, and three blocks at once are all isolated-vs-in-process `max|Δ|=0` and propagate through +later layers; `tracer.patch` equals the manual untuple + whole-tuple replacement (with the same cast). The +crux case proves the motivation under forced isolation: the in-place form leaves the downstream residual +== the unpatched baseline (silent no-op) while `tracer.patch` changes it (transplant took effect) to +exactly the in-process result. + +### `tracer.ablate` — replacement-swap knockout (built, 2026-06-20) + +`tracer.ablate(envoy, mode="zero")` replaces a module's output residual with a baseline — the lesion-study +knockout every ablation cell does. Same shape as `patch`/`steer`: a boundary write riding `Events.SWAP`, +no new event, no isolated/in-process branch. Two self-contained modes: + +```python +out = envoy.output +hidden = out[0] if isinstance(out, tuple) else out +if mode == "zero": + ablated = torch.zeros_like(hidden) +elif mode == "mean": # within-sequence mean + ablated = hidden.mean(dim=-2, keepdim=True).expand_as(hidden).contiguous() +else: + raise ValueError(...) # no silent wrong-ablation +envoy.output = (ablated, *out[1:]) if isinstance(out, tuple) else ablated +``` + +`mode="mean"` is the **within-sequence** mean (each position → the per-example mean over the token +dimension), the only mean derivable from a single forward. The reduction decision that §6 flagged resolves +here: **reference-distribution** mean ablation (the mean activation over a *dataset*, per +`docs/patterns/ablation.md`) is not a single-forward quantity, so it is precomputed and transplanted via +`tracer.patch(envoy, mean_act)` — not this mode. Keeping the two distinct avoids a silent-semantics trap +(a user expecting dataset-mean getting sequence-mean). An unknown mode raises `ValueError` rather than +silently picking a baseline. Under isolation the worker computes the mean from its delivered clone (== +the host's real activation) and ships the result back over SWAP, so isolated == in-process bit-identically. + +Touch points: just the `tracer.ablate` method (tracer.py). No event, no host handler. + +**Verified (`test_isolated_ablate.py`, gpt2 + renamed model):** zero- and mean-ablating a block, an +attention tuple output, and the renamed model are all isolated-vs-in-process `max|Δ|=0` and change the +downstream forward vs the un-ablated baseline (the knockout took effect across the boundary); +`tracer.ablate` equals the manual `zeros_like` / mean-over-seq replacement; an unknown mode raises +`ValueError`. The crux case proves the motivation under forced isolation: the in-place zero leaves the +downstream residual == the un-ablated baseline (silent no-op) while `tracer.ablate` changes it to exactly +the in-process result. + +### `.carry()` — session cross-trace handoff (built, 2026-06-21) + +The last part-2 primitive started as a designed `tracer.capture(value) → handle`, but investigating the +actual gap reframed it. Cross-run handoff across **separate** `model.trace()` calls already works under +isolation — `.save()` ships the value home and the next trace re-receives it (the basis of +`tracer.patch`). The only broken capability was handoff **inside `model.session()`**, which two isolation +bugs killed: + +- **Saved value used cross-trace** → the host wrote it into the session frame but never re-registered its + host-side id in `Globals.saves`, so the session's exit-push (which filters `id(v) in Globals.saves`) + dropped it (`UnboundLocalError`). +- **Non-saved value used cross-trace** → the isolated worker's `end()` ships only `Globals.saves`-filtered + locals, so a non-saved var never reached the host session frame (`NameError`). + +In-process both work, because each inner trace's `push` is `is_root=False` and pushes *all* its locals up +to the session frame (so the next trace's `pull` sees them), while only the session's final root-push +filters to saves. The isolation END-writeback implemented the root semantics only. + +The fix restores the in-process two-hop semantics on the isolated path, plus an explicit primitive for the +non-saved case: + +- **Saved-case fix:** when the END target is a nested/session frame (`__nnsight_tracing_info__` present), + the host writes the worker's values into it AND re-registers the saved values' host ids in + `Globals.saves`, so the session's exit-push keeps them. Root (single-trace) writeback is unchanged. This + makes the documented `hs = x.save()` → use `hs` session pattern work under isolation. +- **`.carry()`** (mounted universally like `.save()`, plus `nnsight.carry(x)`): marks a value to cross to a + later trace in the session **without** surfacing it as an output. The worker's `end()` now ships saved ∪ + carried locals as `(values, saved_names)`; the host writes all of them to the session frame (so the next + trace sees them) but registers only the saved ones, so carried values drop at session exit — exactly the + in-process non-saved semantics, made explicit. With no `.carry()` in use the payload is exactly the prior + saved-only one, so the single-trace path is unchanged. `.carry()` is **portable**: harmless in-process + (non-saved vars already cross there), load-bearing under isolation. + +Touch points: `Globals.shared` + `carry()` + `Object.carry` + mount (globals.py); `nnsight.carry` export; +`Globals.shared` cleared at the root push (base.py) and per-job worker reset (isolation.py); the worker +`end()` payload (isolation.py); `handle_end_event` nested-registration + carried handling (interleaver.py). + +**Verified (`test_isolated_session_handoff.py`, gpt2 + renamed model):** the saved (`.save()`) and carried +(`.carry()`) session handoffs are isolated-vs-in-process `max|Δ|=0` and change the downstream forward; +`nnsight.carry(x)` equals the method form; a carried value does **not** surface to the caller frame while a +saved one does; `.carry()` is in-process == isolated. No regression across the isolated suite (trace, +cache, backward, multitoken-iter, cross-invoke, acceptance, steer, patch, ablate) — the END-payload change +is transparent to the single-trace path — and the in-process core (`test_lm.py`) is 75/75. + +## 7. What was deliberately deferred + +- **The process-global `sys.addaudithook` backstop.** Its own failure mode (a leaked thread-local flag + arms it during the model's *own* forward → server-wide outage) makes it net-negative when the static + default-deny gate already rejects `import` / `open` / `exec` / `socket` *nodes* outright in fast-laned + code (those AST forms route to ISOLATE, so they cannot appear in the in-process body). The remaining + boundary is **provenance, not the gate**: the fast lane is reachable only for `trust="local"` + (author-local) interventions, never remote/NDIF-submitted code, and the numpy/torch call-allowlist is a + **footgun selector, not an adversarial boundary** (a torch or numpy call runs arbitrary compute and both + have escape hatches). So fast-lane safety rests on the `trust="local"` cordon; the static gate only keeps + honest code honest. Documented as future hardening. +- **A frozen-namespace `Compartment`** (SES-style) for fast-lane execution. The first slice relies on + the static pass + the `trust` cordon; namespace shadowing is a later refinement. +- *(none — the part-2 primitive set is complete.)* `unembed`, `steer`, `patch`, `ablate` are built, and + the run↔run handoff shipped as `.carry()` + the session saved-case fix (the original `tracer.capture` + name was dropped — it collided with the existing `Tracer.capture(frame)` AST method, and the realized + primitive is a value method parallel to `.save()`, not a tracer-level handle). + +## 8. Verification + +- **Classifier units** (`test_fastlane_classifier.py`, GPU-free) — 17/17: logit-lens / steering / + patching / attribution shapes and **renamed** structures classify FAST; imports / `while` / + unresolved-call / `open` → ISOLATE; introspection → REJECT; the `differentiate`/`in_place`/ + `touches_host_weights` flags are set correctly. +- **Fast-lane end-to-end** (`test_fast_lane.py`, gpt2 + a renamed model) — 6/6: the weight-reading + logit lens is bit-identical on the fast lane (`max|Δ|=0`) **and raises under forced isolation** + (`fast_lane=False`) — proving the fast lane is the enabling tier; in-place steering bit-identical; + renamed-model lens bit-identical; a footgun routes off the fast lane and the host survives; an + introspection escape is rejected; a runaway loop is killed by the watchdog and the host survives. +- **Isolated steer** (`test_isolated_steer.py`, gpt2 + a renamed model) — 6/6: steering a block, an + attention tuple output, and three blocks at once are isolated-vs-in-process `max|Δ|=0`; `tracer.steer` + equals the manual whole-tuple replacement; and the crux — under forced isolation the in-place form is a + no-op (downstream == unsteered baseline) while `tracer.steer` takes effect and matches in-process. +- **Isolated patch** (`test_isolated_patch.py`, gpt2 + a renamed model) — 6/6: transplanting into a block, + an attention tuple output, and three blocks at once are isolated-vs-in-process `max|Δ|=0`; `tracer.patch` + equals the manual whole-tuple replacement (with the dtype/device cast); and the crux — the in-place + transplant is a no-op under isolation while `tracer.patch` takes effect and matches in-process. +- **Isolated ablate** (`test_isolated_ablate.py`, gpt2 + a renamed model) — 7/7: zero/mean knockout of a + block, an attention tuple output, and the renamed model are isolated-vs-in-process `max|Δ|=0` and change + the downstream forward vs the un-ablated baseline; `tracer.ablate` equals the manual `zeros_like` / + mean-over-seq replacement; an unknown mode raises `ValueError`; and the crux — the in-place zero is a + no-op under isolation while `tracer.ablate` takes effect and matches in-process. +- **Isolated session handoff** (`test_isolated_session_handoff.py`, gpt2 + a renamed model) — 6/6: the + `.save()` and `.carry()` session cross-trace handoffs are isolated-vs-in-process `max|Δ|=0` and change + the downstream forward; `nnsight.carry(x)` equals the method form; a carried value does not surface to + the caller frame while a saved one does; `.carry()` is in-process == isolated. +- **Existing isolated WORKER path** — 9/9 still bit-identical, pinned with `fast_lane=False` so they + keep exercising the worker (otherwise the simple read/swap/save cells would now fast-lane). +- **In-process core** — `test_lm.py` 75 passed (the default in-process path is untouched). diff --git a/docs/developing/gpu-sandbox.md b/docs/developing/gpu-sandbox.md new file mode 100644 index 000000000..ab017abac --- /dev/null +++ b/docs/developing/gpu-sandbox.md @@ -0,0 +1,117 @@ +# GPU Sandbox — the chosen isolation implementation + +**Status:** Built + tested (functional + safety pass) · **Date:** 2026-06-06 +**Supersedes** the CPU-only transport work (harness-plan Phases 5/5b/6) for the GPU path. + +> **Pre-integration prototype.** This page documents the standalone prototype under +> `prototypes/mediator-sandbox/gpu_sandbox/` (`sandbox.py` seccomp, `gpu_worker.py`, `gpu_sandbox.py`). The +> design has since landed in `src/nnsight/intervention/`: the prototype `sandbox.py` became `_sandbox.py`, +> wired through `isolation.py` + `transport.py` + `interleaver.py`. For the shipped behavior and current +> support matrix see [mediator-gpu-trace-integration.md](mediator-gpu-trace-integration.md) (authoritative) +> and [mediator-threat-models.md](mediator-threat-models.md) (security posture). The "why this design" +> rationale below still holds. + +## Why this design + +The threat model was relaxed to **contain footguns, not defeat a determined adversary** — stop a careless +or buggy intervention from crashing the shared model server, reading host files, exfiltrating, hanging, or +OOMing it; do **not** (for now) defend against a malicious user weaponising the GPU. Under that model the +worker can keep **GPU access**, which makes the data path **zero-copy** and deletes the whole transport +problem (D2H + serialize) that made the CPU-only path cost ~110 ms/hook. + +## Design + +A **separate, spawned, GPU-enabled worker process** runs the user's arbitrary intervention on the real GPU +activation, while a process boundary + seccomp contain everything except the GPU. + +- **Zero-copy via a shared GPU "bounce buffer."** A single CUDA tensor is shared host↔worker via CUDA IPC + **once at spawn** (before lockdown). Per request the host does a cheap on-GPU (D2D) copy of the + activation into the buffer + a 1-byte signal; the worker views the buffer as the activation tensor and + runs the user op in place; the host reads the result back. + + **Measured (`perf_test.py`, A100): flat ~0.6–0.8 ms/hook, independent of size** — 0.59 ms at 0.03 MB, + 0.71 ms at 16.8 MB (~7B), 0.77 ms at 33.6 MB (~70B). The CPU-transport path scaled 1→285 ms with bytes; + this doesn't, because the activation never leaves the GPU (per-request copy is an on-GPU D2D memcpy, + ~0.01–0.04 ms). End-to-end on a gpt2 trace: **+1.1 ms/hook** (10.1 → 11.2 ms for one intervention; 10 → 24 ms for all 12 + layers). vs ~0.02 ms fully in-process. + + **Where the ~0.6 ms goes (decomposed, `perf_decompose.py` / `perf_ctxswitch.py`): GPU context-switching + between the host and worker CUDA contexts — NOT process communication.** Measured: bare process round-trip + 0.04 ms, `cloudpickle.loads` 0.01 ms, one `cuda.synchronize` 0.016 ms. But a round-trip where BOTH host + and worker do a GPU op (apply's pattern) = 0.405 ms vs 0.040 ms when only the host does → **~0.37 ms is the + GPU switching active context host↔worker** each hook (each process has its own context on the shared + device; they alternate every intervention). + + *Fix (measured under a scoped CUDA MPS daemon):* **CUDA MPS routes both processes' kernels through one + shared GPU context**, eliminating the per-process switch. The both-touch-GPU round-trip dropped + **0.405 → 0.132 ms** and the context-switch overhead **0.365 → 0.090 ms (~4×)** — confirming the diagnosis. + So `apply()` under MPS goes from ~0.6 ms toward **~0.2 ms**, end-to-end toward ~+0.3 ms/hook. (A spin-waiting + worker does **not** help — it only removes the ~0.04 ms IPC, not the context switch.) + - **Fault-domain tradeoff — measured on the A100 (`fault_test.py`), better than feared:** MPS shares one + GPU context, so in principle it couples the fault domain. In practice, **every GPU fault reachable from + realistic buggy code was CONTAINED under MPS — the host's CUDA context survived**, same as with separate + contexts. A buggy intervention's out-of-bounds index / embedding produces a **device-side assert** (PyTorch + bounds-checks its indexing), and Volta+ MPS attributes that to the offending client and terminates only it. + I could **not** trigger a genuine raw illegal-memory-access (XID 31) from normal tensor ops — they all + became asserts. The residual coupling is the *rare* worst case (a true bad-pointer illegal access, or a + hardware/ECC fault), which MPS docs say can hit the server. **Net:** MPS gives ~3× *and* keeps containment + of the faults that actually happen — a viable default; separate contexts only buy the rare worst-case. + | mode | per-hook | OOB→device-assert (the real footgun) | rare raw-illegal/HW fault | + |---|---|---|---| + | separate contexts | ~0.6 ms | contained (host survives) ✓ measured | contained | + | Volta+ MPS (A100) | ~0.2 ms | **contained (host survives) ✓ measured** | shared blast radius | + - *Why a bounce buffer and not per-request IPC:* `cudaIpcOpenMemHandle` (per-request rebuild) **fails + after the seccomp lockdown** (`context is destroyed`). Mapping one buffer before lockdown sidesteps it; + the per-request D2D copy is on-GPU memcpy (microseconds), not the bottleneck. +- **Footgun containment (`sandbox.py`):** after CUDA + torch are warmed, the worker installs a hand-rolled + **seccomp-BPF** filter (no external dep) that makes new `open`/`openat`/`openat2` (filesystem), + `socket`/`connect` (network), and `execve` fail with EPERM. CUDA keeps working — it uses already-open + `/dev/nvidia*` fds via ioctl, not new file opens. GPU memory is capped with + `torch.cuda.set_per_process_memory_fraction` (an OOM footgun can't exhaust the device). A per-call + **timeout** kills a wedged worker; the **process boundary** isolates crashes and hides all host Python + objects (weights, interleaver, other tenants). + - *Gotcha:* `RLIMIT_AS` is unusable with CUDA — CUDA reserves tens of GB of *virtual* address space, so + an 8 GB AS cap kills it at startup with `MemoryError`. Use the GPU memory fraction instead. +- **Pool.** CUDA isn't fork-safe → workers are spawned (not forked) and kept warm; each request goes to a + free worker; a crashed worker is detected and replaced. (`GPUSandbox` is one worker for clarity.) + +**Files:** `prototypes/mediator-sandbox/gpu_sandbox/` — `sandbox.py` (seccomp), `gpu_worker.py` (the +spawned worker), `gpu_sandbox.py` (host manager + `apply(activation, fn)`), `test_functional.py`, +`test_safety.py`. Run on a GPU: `CUDA_VISIBLE_DEVICES=N PYTHONPATH=/src .../test_*.py`. + +## Results + +**Functional (`test_functional.py`) — PASS.** Real nnsight `model.trace()` workloads (scale, steer, +mean-ablate, per-token-norm read) on gpt2, run two ways — op inline vs op offloaded to the GPU worker — +produce **bit-identical** logits (`max|Δ|=0`); the worker survives all requests. nnsight delivers the +activation through its real machinery; only the user op runs isolated. (Core nnsight suite: 152 passed.) + +**Safety (`test_safety.py`) — PASS (9/9).** Each mimicked unsafe request is contained: + +| Unsafe request | Outcome | +|---|---| +| read a host secret file | `PermissionError` (open blocked); secret not returned | +| write a host file | `PermissionError`; no host file created | +| network egress (`socket`/`connect`) | `PermissionError` | +| reach a host object (`THE_MODEL.config`) | `NameError` — host objects don't exist in the worker | +| OOM (`torch.empty(1e12)` on GPU) | `OutOfMemoryError` — GPU mem fraction caps it | +| infinite loop | timeout → worker killed; host unaffected | +| hard crash (NULL-deref segfault) | worker dies; **host CUDA context intact**, crash surfaced cleanly | +| legit op after all attacks | still correct (worker survived the catchable ones) | +| fresh worker after a crash | serves correctly (pool recovery) | + +## What this is vs. isn't + +- **Contained:** filesystem, network, host crashes/OOM/hangs, and all host Python objects/other tenants' + CPU-side state. The model server cannot be taken down or read by a buggy intervention. +- **Not contained (the accepted axis):** the GPU. The worker has a CUDA context and shares the activation, + so a *deliberately malicious* user could attack the driver or read neighbouring GPU memory. Out of scope + per the current threat model; revisit (CPU-only worker, MIG, dedicated allocations) if the model tightens. + +## Remaining work + +This is a working harness with an explicit offload API: `output = sandbox.apply(activation, fn)` inside a +trace. Wiring it **transparently into `model.trace()`** (so `.output`-access/ops/`.save()` route to the +worker automatically via the eproperty bridge) is the remaining productionisation — the same integration +as the transparent-trace-integration milestone, but simpler here because there is **no transport codec** (zero-copy GPU buffer). +Plus the warm-pool/scheduler and per-tenant dedicated allocations if the threat model tightens. diff --git a/docs/developing/mediator-gpu-trace-integration.md b/docs/developing/mediator-gpu-trace-integration.md new file mode 100644 index 000000000..f7fe3a852 --- /dev/null +++ b/docs/developing/mediator-gpu-trace-integration.md @@ -0,0 +1,645 @@ +# Transparent GPU-Worker Trace Integration — Design + +**Status:** Design (approved) · **Date:** 2026-06-06 · **Branch:** `worktree-mediator-sandbox` +**Builds on:** +- [gpu-sandbox.md](gpu-sandbox.md) — the chosen GPU-worker mechanism (`GPUSandbox.apply`, bounce buffer, seccomp, lifecycle). +- [mediator-isolation-harness-plan.md](mediator-isolation-harness-plan.md) — the earlier CPU-transport prototypes; this doc supersedes the transparent-`model.trace()` milestone described there. +- [mediator-isolation-sandbox.md](mediator-isolation-sandbox.md) — threat model (relaxed to "contain footguns, not adversaries") + AWS deployment. + +**Goal:** make `model.trace()` run user interventions in an isolated GPU worker process **transparently** — +identical results to in-process — so a footgun in user code (infinite loop, OOM alloc, bad-index +device assert, host-object poke) is contained to the worker while the shared model server keeps serving. +The standalone `GPUSandbox.apply()` proved the worker *mechanics*; this turns it into the actual +`model.trace()` execution backend. + +--- + +## 1. The architecture: an outer harness, not a rewrite + +The Mediator↔model handoff already routes through the `MediatorChannel` seam +(`interleaver.py:774`, extracted earlier). The Mediator talks to the model **only** through that channel, +via the six events (`VALUE/SWAP/SKIP/BARRIER/END/EXCEPTION`). **Those events do not change.** Isolation is +an outer harness: run the Mediator's worker in a separate (GPU-enabled, locked-down) process, and give it +a cross-process channel. + +``` +HOST (trusted: model + weights + GPU) WORKER (spawned, GPU-enabled, seccomp'd) + real Envoy tree + Interleaver + Batcher path-only Envoy mirror (no weights) + host-side Mediator worker-side Mediator (same class, other half) + · owns channel, batch_group, history, handle_* · runs the REAL user intervention fn + · registers hooks ON DEMAND · eproperty.__get__ → request() → blocks + CudaIpcChannel host end ◄── control pipe ──► CudaIpcChannel worker end + │ + shared GPU buffer │ + forward fires hook → handle() → narrow → user code runs on the GPU tensor view, + D2D into buffer → respond `.save()` accumulates locally, + shipped back at END +``` + +**One `Mediator` class, instantiated on both sides**, each using its half. The existing +`Mediator.__getstate__`/`__setstate__` (`interleaver.py:1558`) already define exactly what crosses to +construct the worker side — they were built for NDIF remote. The host half owns the Batcher and +`handle_*`; the worker half owns `send`/`request`/`swap`/`push` and runs the intervention. + +### What does NOT change +- The six-event protocol and its one-event-in-flight invariant (the "access in forward-pass order" + contract). +- `Mediator.send`/`respond`/`handle`/`handle_*`, the Batcher narrow/swap, requester↔provider matching. +- `eproperty`/`hooks.py` core logic (the worker's hook call lands on a dead dummy module — see below). +- The in-process path: when isolation is off, `Mediator.start` runs the existing `Thread` exactly as today. + +### What is genuinely new (and *why*) +Two places in the in-process code rely on **shared memory** rather than the channel. These — and only +these — are why isolation is more than "spawn + new channel." Both live in the harness layer. + +1. **Hook registration touches the real module via shared memory.** In-process, `eproperty.__get__` + registers a PyTorch hook *on the actual module object* (`hooks.py:307`, `output_hook`). The worker + process has no module. But the event already carries the requester string + (`"transformer.h.6.output.i0"`), which fully encodes (path, kind, iteration) — so the **host** + registers the hook on receipt (the *host-side on-demand hook registration* described in §3). +2. **`.save()` propagates through a shared Python frame.** In-process, trace exit reads the worker's + frame locals, keeps `id(v) in Globals.saves`, and `push_variables` into the user frame — all + same-process pointer work (`tracing/base.py:565`). Across processes the worker's frame and + `Globals.saves` live in the worker, so it runs that filter locally and ships the result over the + channel; the host does the final `push_variables` (the *worker→host saves transmission* in §4). + +--- + +## 2. Transport — `CudaIpcChannel` + +A new `MediatorChannel` implementation, same shape as `transport.py`'s shared-memory socket channel; only +the tensor encode/decode differs. + +- **Control plane** (`mp.Pipe`, small pickle frames): `(event, requester)` + an *offset table* + `[(offset, nbytes, shape, dtype), …]` + a non-tensor skeleton. Reuses `transport.py`'s + `_split_tensors`/`_merge_tensors` verbatim — the only change is "encode each tensor as a buffer + offset" instead of "safetensors bytes." +- **Bulk plane** (shared GPU bounce buffer): the sender D2D-copies each tensor into the buffer at + successive offsets; the receiver views them at those offsets. The buffer is a CUDA tensor created on + the host and shared to the spawned worker via `torch.multiprocessing` (CUDA IPC), mapped once before + lockdown. +- **Aliasing rule (correctness, mandatory):** under one-event-in-flight the single buffer is reused in + strict alternation, so **both sides clone-on-receive** before the next access can overwrite it. This + covers `x = h[6].output; y = h[7].output; use(x, y)` — `x` must be cloned out of the buffer before the + request for `y` overwrites it. A D2D clone is HBM-speed (~µs/MB); "zero-copy" always meant *no PCIe, + no pickle*, not *no copy*. The measured ~0.6 ms/hook already includes these copies. +- **Sizing:** a fixed arena (64 MB default). A value larger than the arena raises a clear error (the + `ShmArena.write` "exceeds arena" guard). Chunking is a later perf item, out of scope. +- **CPU-only fallback backend:** the existing CPU shared-memory socket channel stays config-selectable + for CPU-only isolation — free, since it is already a `MediatorChannel`. GPU-IPC is the default. + +Two roles, mirroring `transport.py`: `CudaIpcHostChannel` (main thread; only `wait_event` reads the +pipe; `has_event`/`get_event`/`restore_event` are host-local buffer ops) and `CudaIpcWorkerChannel` +(send→wait→get). + +--- + +## 3. Host-side on-demand hook registration + +In-process, the worker's `requires_output` registers the hook just-in-time as it single-steps. In +isolation the worker can't (no real module), so the host registers it. + +- **Where:** in the host `handle` loop, when a `VALUE`/`SWAP`/`SKIP` arrives for a requester `R` whose + provider is not yet set up (the existing `handle_value_event` else-branch, `interleaver.py:1203`, + that today does `history.add`/`restore_event`/`return False`). +- **What:** `ensure_isolated_provider(R)`: parse `R` (`"..i"`) → `iteration = N`, + `kind`, `path`; resolve `envoy = root_envoy.get(path)` (`envoy.py:586`); call the **existing** + `output_hook`/`input_hook(host_mediator, envoy._module, f"{path}.{kind}")` (`hooks.py:224`/`154`). + The host-side mediator is passed, so the hook closure delivers over the channel. +- **Worker side:** the worker's `eproperty.__get__` still runs `self._hook(obj)`. The worker's + path-only envoy carries a bare `torch.nn.Module()` as `_module`, so `output_hook` registers a hook on + a **dead dummy** that never fires (harmless; reclaimed when the worker process exits — NOT by + `cancel`, since the worker never calls it). **Warm-pool caveat:** a warm worker pool reuses processes, + so these dummy-module hooks (and the worker-local `iteration_tracker`/`Globals.saves`) would accumulate + across traces and need explicit per-trace cleanup. Keeping `eproperty`/`hooks.py` unchanged was the + reason to register on a dead dummy rather than guard `_hook` with an `interleaver.isolated` flag; + revisit if dummy-module accumulation shows up in profiling. +- **Idempotency:** `ensure_isolated_provider` registers a given `R` once per mediator (tracked in a set on the + host mediator), mirroring the `current_provider`-skip logic in `requires_output`. + +For a single forward pass the iteration is always `0`, so `R` is `"..i0"`. Multi-token +iteration stamping is handled where the multi-token feature is built (§9). + +--- + +## 4. Worker→host saves transmission + +- **Worker:** at trace exit / on `END`, the worker runs the *existing* `base.py` exit filter (it owns + the frame + `Globals.saves`) to produce the saved dict, then — instead of `push_variables` locally — + ships that dict over the channel. +- **Host:** receives the dict and `push_variables(user_frame, saved_dict)` into the real user frame + (the frame that called `model.trace()`, which lives on the host). +- **Encoding:** tensors in the saved dict ride the bounce buffer (offset table); other saved objects + ride the control pickle, reusing `_split_tensors`. +- **New constraint vs in-process:** a *non-tensor* `.save()` must be serializable to cross the boundary + (in-process it was just a reference). Tensors (the overwhelming common case) are fine; a + non-serializable save raises a clear error naming the variable. Documented as an isolation semantic + difference. + +### Isolation semantic: in-place modification requires explicit assignment +In-process, `.output` returns the model's *real* tensor (shared memory), so an in-place edit +(`h[6].output[0][:] = ...`) propagates by mutation with no SWAP event. In the isolated path the worker +receives a **clone** (the clone-on-receive rule, §2), so an in-place edit mutates a worker-local copy and +does **not** reach the host. Use **explicit assignment** (`h[6].output = ...`), which emits a SWAP event +and applies via `batcher.swap`. This is identical to **NDIF remote** semantics (in-place on a serialized +clone never crossed the wire either), so it is a consistency property, not a new wart. Transparent +in-place support would need a protocol read-back that conflicts with the single-buffer clone-on-receive +rule; deferred (revisit with double-buffering if real workloads need it). + +--- + +## 5. Lifecycle parity + +- **`cancel()`** terminates the worker process (the `gpu_sandbox` `close()`/`terminate()` pattern) + instead of dropping a `Thread`. `check_dangling_mediators` and the EarlyStop drain account for + process workers. +- **Timeout:** a per-call timeout kills a wedged worker (infinite loop in user code). Surfaces a clear + `TimeoutError`; the host is unaffected; a warm pool would respawn. +- **EXCEPTION:** the worker ships the exception over the channel; the host re-raises it in the user's + context with the rebuilt traceback. Dynamic nnsight exceptions that can't pickle are degraded to a + plain exception preserving type name + message before they cross (§11). +- **SKIP / BARRIER / END / EarlyStop:** already cross as events. `END` triggers the saves transmission + then worker exit + host join. `stop()`'s `push()` ships saves before signalling. + +--- + +## 6. Feature coverage map + +Folded into the **support matrix (§8)**, which carries each feature's cross-process mechanism and +current status in one table. The invariant it tracks: when isolation is on and a not-yet-supported +feature is used, the trace **fails cleanly** — a missed-provider error or the per-step timeout (the +lifecycle is the safety net), not a silent deadlock or silent-wrong result. (There is no automatic +"route to in-process" fallback; features are added one at a time.) + +--- + +## 7. Core single-pass interventions — scope, components, acceptance + +This is the first slice: a single forward pass (iteration 0), covering read / swap / `.save()` / +multi-invoke / skip / exception / lifecycle / seccomp lockdown — the foundation the rest builds on. + +### Scope +Single forward pass; single and multiple invokes; read (`.output`/`.input`/`.inputs`); swap (`=`); tensor +`.save()`; skip; exception; lifecycle (normal END, exception, timeout→kill); tuple/nested block outputs. +`cross_invoker` disabled. Multi-token / barrier / backward / cache not yet handled. + +### Components +1. **`CudaIpcChannel`** (`transport.py`, additive): `CudaIpcHostChannel` + `CudaIpcWorkerChannel` + + GPU-buffer tensor pack/unpack (reusing `_split_tensors`/`_merge_tensors`). +2. **Isolated `Mediator.start`** (`interleaver.py`): a branch that spawns the worker (CUDA → `spawn`), + ships `self` via source-serialization with the model→path-only-envoy persistent-object map, and a + worker bootstrap that builds the interleaver stub + path-only envoy mirror + runs the intervention. +3. **Host-side on-demand hook registration** (`interleaver.py`): `ensure_isolated_provider(R)` in `handle`. +4. **Saves transmission + lifecycle** (`interleaver.py` + bootstrap): worker ships `Globals.saves`- + filtered frame locals at END; host injects; `cancel()` kills the process. +5. **Opt-in** (§12): `CONFIG.APP.ISOLATE_MEDIATORS` + `nnsight.isolate_mediators()` context. + +### Acceptance — bit-identical to in-process (`max|Δ|=0` on GPU) +1. read-save; swap; multi-invoke batch-narrow; nested-tuple block output. +2. **A non-standard-named toy model** (`decoder_blocks`, `output_projection`) **and** gpt2 — per the + testing rules: structure derived at runtime via `envoy.get(path)`, no GPT-2-only assumptions, no + hardcoded module-name tables. +3. exception in user code → raised in the user's context with the correct traceback. +4. infinite loop in user code → worker killed, host survives, clear `TimeoutError`. +5. safety: the unsafe payloads (filesystem/network/exec/OOM/crash) contained under seccomp lockdown. +6. core suite still green (in-process path untouched). + +### Status: DONE (2026-06-06) +Built and verified (TDD; harnesses in `prototypes/mediator-sandbox/gpu_sandbox/`): +- **Code:** `transport.py` (`pack_cuda`/`unpack_cuda` + the two channel ends, clone-on-receive, per-wait + timeout); `isolation.py` (`isolate_mediators` + `_spawn_worker` + `_pool_worker_main` + + `ensure_isolated_provider` + `_WorkerInterleaver`/`_WorkerPersistent`); `_sandbox.py` (seccomp + `lock_down`, relocated from the prototype); `interleaver.py` seam (`_iso` field, isolated + `Mediator.start` branch, the on-demand hook-registration call in `handle`, saves injection in + `handle_end_event`, `_iso.close()` in `cancel`). +- **Foundations (probes, all PASS):** buffer codec round-trip + clone-on-receive; CUDA-IPC channel over + a real spawned process; bootstrap (serialize→deserialize-against-dummies→run→correct `request`); + saves capture (`info.frame.f_locals` filtered by `Globals.saves`, keyed by user var name). +- **End-to-end (`test_isolated_trace.py`):** isolated read `.output[0].save()` and explicit-assignment + swap (`h[6].output = … * 2`) **bit-identical to in-process, `max|Δ|=0`** on gpt2/cuda. +- **Acceptance (`test_isolated_acceptance.py`):** non-standard module names (`decoder_blocks`, + `NNsight(net)`) `max|Δ|=0`; two isolated invokes bit-identical + rows distinct (no cross-leak); user + exception → `NNsightException` in the user context; infinite loop killed (timeout), host survives. +- **Safety (`test_isolated_lockdown_safety.py`):** read under `lockdown=True` `max|Δ|=0`; `open()` + blocked + no host file; `socket()/connect()` blocked. +- **Regression:** core in-process suite **155 passed** (in-process path untouched). + +**Findings folded into the design:** per-mediator serialization must set `intervention.__source__` +first (the tracer normally does it in *its* `__getstate__`); the worker frame is a `SerializedFrame` so +`push()` lands locals there; saved values inject into the **tracer's** `info.frame` (the user frame), +not the mediator's; `_remoteable_persistent_objects` is remoteable-only (gate with `isinstance`); +locked-down workers must `os._exit(0)` to skip tempfile atexit (avoids a `RecursionError`). `lockdown` +defaults off in `_STATE` (flip to on once the warm pool lands); the bounce-buffer force-kill emits a +benign CudaIPC release warning. + +--- + +## 8. Support matrix (what works under `isolate_mediators()` today, and how) + +> **Fast lane (2026-06-13, [fast-lane.md](fast-lane.md)).** `isolate_mediators()` now runs each +> mediator on one of three tiers: a static classifier confirms safe interventions and runs them +> **in-process** (FAST — full model + weights), isolates the unconfirmable remainder in the GPU worker +> (ISOLATE), and rejects introspection escapes (REJECT). This is what lets the **weight-reading interp +> majority** (logit lens, steering, ablation, activation patching, attribution) run under isolation at +> all — the worker holds weightless dummy modules, so those workloads can only run on the fast lane. +> The matrix below describes the **ISOLATE (worker) tier**; a row marked weight-reading/host-module is +> served by the FAST tier. Default-on; `fast_lane=False` forces the worker tier. + +| Feature | Cross-process mechanism | Status | +|---|---|---| +| read / swap (`=`) / `.save()` (tensors) / skip / exception | six events over the channel; host-side hook registration; worker→host saves transmission at END | ✅ bit-identical | +| multi-invoke + batch narrowing | per-invoke worker + host mediator; Batcher stays host-side | ✅ bit-identical | +| single-forward `generate(...)` (no iter) | same as trace | ✅ verified | +| seccomp lockdown (fs/net/exec) | `_sandbox.lock_down` after the first job's deserialize, in `_run_one_job` | ✓ a read under `lockdown=True` is `max|Δ|=0`, user-code `open`/`socket`/`exec` are blocked, and a warm pool keeps serving (`test_isolated_lockdown_safety.py` 4/4). The break — the worker locked down BEFORE receiving its first job, so unpickling the job message (its `transformers` lazy submodules load only at unpickle time) hit a seccomp-blocked `open` (an `OSError` the recv loop swallowed → silent death) — is fixed by deferring lockdown to after the (host-authored, trusted) first job is deserialized, before its user code. One-way and installed once: in a warm pool, a later job needing a *new* model's modules still needs `preimport=` (a homogeneous model needs nothing). | +| `iter`/`all`/`next` (multi-token) | step stamped in the requester; host iter-hooks bump the tracker; live `default_all` piggyback (§9) | ✅ bit-identical (`iter[N]`, `iter[:]`, per-step swap) | +| `tracer.barrier()` | worker sends the target count; host accumulates participants + runs the coordination loop (§10) | ✅ | +| `cross_invoker` variable sharing | host variable store; worker pushes data locals, pulls the merged store; transmittable data only (§10) | ✅ | +| warm worker pool (`pool_size=N`, `warm_worker_pool`) | generic workers receive serialized mediators as jobs; clean-END recycle (§14) | ✅ ~21× faster per request once warm | +| `with tensor.backward()` / `.grad` | BACKWARD event: worker seeds `dL/d(delivered clone)`, host continues `torch.autograd.grad` on the real graph, `.grad` by provenance path; **grad-through-swap** iterates the exchange (host returns `dL/d(swap leaf)`, worker backprops it through its swap tape, re-seeds the pre-swap graph) (§16) | ✅ read-path AND grad-through-swap bit-identical (scalar loss; single mediator, including list-batched `trace([A,B,...])`); multi-invoke-context and multi-token backward are clean-fail parity (in-process does not support them either) | +| `tracer.cache()` (`modules=`, `include_inputs=`) | CACHE event → host registers the real cache hooks; host CacheDict swapped in at END, filled in-place by the forward (§15) | ✅ bit-identical | +| part-2 primitives: `tracer.unembed` / `tracer.steer` / `tracer.patch` / `tracer.ablate` | host-routed weight read (UNEMBED event); replacement-swap injection/transplant/knockout (ride SWAP, no new event) — [fast-lane.md](fast-lane.md) §6 | ✅ bit-identical (isolated and in-process) | +| session cross-trace handoff (`.save()` used in a later trace; `.carry()` / `nnsight.carry(x)`) | inner-trace END writeback to the session frame: saved values re-registered host-side so the session exit-push keeps them; carried (non-saved) values written for the next trace only — [fast-lane.md](fast-lane.md) §6 | ✅ bit-identical (`.carry()` is portable: harmless in-process, load-bearing under isolation) | +| Triton-kernel models (MoE / SSM / `torch.compile`) | host-side forward compiles/runs Triton unrestricted (the worker holds only the intervention) | ✅ — §17 | +| user-code Triton (kernel inside the intervention) | — (compiling a kernel needs `open`/`subprocess`/ptxas, which lockdown blocks by design) | ⛔ under lockdown — §17 | +| `.source` operation-level access (`...attn.split_1.output`) | — (op paths aren't in `model.modules()`) | 🔜 not yet | +| in-place `[:]=` | — (clone-on-receive semantics; use explicit `=`, §4) | ⛔ | +| MPS / `isolate_mediators()` further polish | — | planned | + +Not-yet-supported features fail **cleanly** (missed-provider error or the per-step timeout), not as a +silent deadlock — the lifecycle (timeout + `finally: cancel()`) is the safety net until each feature lands. + +--- + +## 9. Multi-token iteration — STATUS: DONE (2026-06-07) + +Multi-token iteration (`iter[N]`, `iter[:]`, `next`) isolated == in-process. The iter loop runs in the +worker (sets `mediator.iteration` explicitly per step → correct requester suffixes); the **host** side +gained three pieces: +- **Per-step hook registration:** `ensure_isolated_provider` parses the step `N` from the requester and + passes `iteration=N` to `output_hook`/`input_hook` (new optional param), so the hook fires on step N — + not the host mediator's iteration. +- **Host iter-hooks:** `_spawn_worker` calls `register_iter_hooks(host_mediator, real_model)` so + the host `iteration_tracker` advances per forward. +- **Live host→worker piggyback:** `default_all` (= `generate(max_new_tokens)`) is set *after* the worker + spawns, so it's piggybacked on each response frame (`CudaIpcHostChannel.meta_provider` → + `CudaIpcWorkerChannel.on_meta`) and applied to the worker's interleaver stub before its `iter[:]` loop + computes its bound. This is a general host→worker state channel reused later for the variable store. + +**Verified (`test_isolated_multitoken_iter.py`):** `iter[N]` for N∈{0,1,2} `max|Δ|=0`; per-step swap +propagates (`iter[1]`); `iter[:]` accumulating into `nnsight.save(hs)` bit-identical (3 steps). +**Regression:** channel echo PASS (3-tuple response format); core in-process suite (incl. +`test_iter_edge_cases`) 113 passed; isolated single-pass trace still `max|Δ|=0`. + +**Findings:** `default_all` is live host state set post-spawn → can't snapshot at spawn (hence the +piggyback). The worker over-requests one step past the end only if `default_all` is missing (now fixed). +In-place `[:]=` per step still requires explicit `=` (§4). + +--- + +## 10. Cross-invoke (barrier + variable sharing) — STATUS: DONE (2026-06-07) + +Both isolated == in-process. + +- **Barrier (host-side counting):** each invoke runs in its own worker with its own `Barrier` copy, so it + can't count cross-invoke. `Barrier.__call__` (isolated) sends the TARGET count; the host accumulates + participant names in `Interleaver._barrier_acc` and, once all arrive, runs the existing coordination + loop (`handle_barrier_event` iterates the host mediators, `respond()`+`handle()` each over its own + channel). `Mediator._isolated_worker` (set in `_pool_worker_main`) gates the worker-side behavior. +- **Variable sharing (host store):** worker frames aren't shared across processes, so each worker pushes + its *data* locals to `Interleaver._xinvoke_store` (a 4th `push` field on the event frame) and pulls the + merged store back (piggybacked on the response, reusing the multi-token channel). Only **transmittable + data** (tensors + basic scalars/containers, `_transmittable`) crosses — framework objects + (`Barrier`/`Envoy`, which hold the model) are skipped (the worker already has them via its closure). + Tensors travel **CPU-serialized** (D2H on push, H2D on pull): a worker tensor cloned from the CUDA-IPC + buffer can't be re-shared over IPC by the host ("CUDA tensor received from another process"). + cross_invoker is gated exactly like the host (`len(mediators)>1 and CONFIG.APP.CROSS_INVOKER`). + +**Verified (`test_isolated_cross_invoke.py`):** two-invoke `tracer.barrier(2)` produces matching saved +values; the canonical cross-invoke pattern (invoke A captures `h[3].output`, barrier, invoke B sets +`h[3].output = captured`) is bit-identical. **Non-standard names (`test_nonstd.py`):** gpt2 with +`rename={transformer.h→decoder_blocks, lm_head→output_projection}`, read + iter[1] isolated `max|Δ|=0`. +(Note: `Envoy.__getattr__` resolves the alias to the *real* path before the requester is built, so the +wire string is always the real path — this validates renamed models run end-to-end, but does not itself +exercise an alias↔real mismatch in the host hook registration, since none reaches the wire.) Echo +regression PASS (4-tuple event). + +**Findings:** the variable-sharing push must filter to transmittable data (else the `Barrier` local pulls +in the whole model) AND move tensors to CPU (CUDA-IPC tensors can't be re-shared by the host). The +`_xinvoke_store` is per-interleaver (reset each trace) so no cross-trace leak today; the warm pool (§14) +rebuilds the interleaver + dummy modules per job and clears `Globals.saves`/`Globals.shared`, so there is +no cross-trace leak between *unrelated* traces under reuse either. (Intentional cross-trace handoff +*within a `model.session()`* — a `.save()`d value used in a later trace, or `.carry()` — is a separate, +supported path: the worker ships those values to the session frame at END; see the §8 support matrix and +[fast-lane.md](fast-lane.md) §6.) Known coverage gaps: no tests yet for +multi-barrier-in-one-trace, 3+ participants, multi-token + barrier, or variable sharing without a barrier; +the store grows monotonically per trace and ships all shared tensors CPU-serialized on every response (a +perf cliff for large cross-invoke tensors). + +--- + +## 11. Backward + caching — characterization (both gaps since closed: cache §15, backward §16) + +A gap-characterization harness (`test_isolated_backward_cache_gaps.py`, since retired — canonical +coverage lives in `test_isolated_cache.py` / `test_isolated_backward.py`) confirmed the two gaps and +their difficulty at the time; kept for the reasoning record: + +- **`tracer.cache()` (a real build, not a quick shim):** `tracer.cache()` runs in the worker → registers + cache hooks on dummy modules → never fire → the `.save()`'d CacheDict comes back **empty**. The fix + needs: (1) a `CACHE` event carrying the spec (module paths + options) so the host registers + `cache_output_hook`/`cache_input_hook` on the *real* modules; (2) **post-forward injection** — the + catch is that cache hooks are **persistent** and the CacheDict is populated **by the forward, which runs + AFTER the mediator's intervention ends** (the intervention only sets up the cache, then ends), so the + END-time saves transmission runs too early. The populated host cache must be copied into the user + frame's CacheDict **at trace exit** — specifically in `Interleaver.cancel` BEFORE `remove_hooks` (the + host cache hooks survive `mediator.cancel`; they're only dropped in `Interleaver.cancel`'s finally, + after the forward has populated them). Matching each worker CacheDict to its host cache needs a token (a + counter set on the CacheDict + carried in the CACHE event; order-based matching is fragile and the token + attr must survive the CacheDict's pickling). A clean build, but ~4 touch points + its own mechanism. + (Acceptance test `test_isolated_cache.py` written, currently failing as expected.) +- **`with tensor.backward()` (hard — likely a major build or a documented limitation):** the backward + context runs in the worker on **detached clones** (clone-on-receive strips the autograd graph, which is + host-side), so there is no graph to differentiate. Making it work requires the backward pass to run + **host-side** (where the graph is) with **path-based grad providers** (today they are `id(tensor)`, + process-local) — architecturally similar to the forward's host-side hook registration but for the + backward session, comparable in size to the core single-pass seam. Until then, backward under isolation + must fail cleanly / run non-isolated. +- **Robustness bug found + fixed (affects all features):** a *dynamic* `NNsightException` raised inside + the worker fails to pickle across the EXCEPTION event (`PicklingError: Can't pickle + nnsight.NNsightException`). Plain exceptions (e.g. a user `ValueError`) cross fine; nnsight-internal + dynamic exceptions did not. The worker now degrades a non-picklable exception to a plain + `RuntimeError(type name + message)` before the EXCEPTION event. + +--- + +## 12. Opt-in surface +`CONFIG.APP.ISOLATE_MEDIATORS` (flag) + `with nnsight.isolate_mediators(): ...` (context). `Mediator.start` +checks it to select the isolated branch. The channel backend (GPU-IPC default, CPU shared-memory +fallback) is config-selectable. Server deployments set the flag globally. + +--- + +## 13. Risks +- **Per-hook latency** — the accepted cost; ~0.6 ms measured, size-independent (vs CPU pickle 1→285 ms). +- **Codec correctness for nested/tuple values** — needs the `applyn`-aware pack tested against + tuple-output blocks (covered by the nested-tuple acceptance case). +- **Lifecycle/deadlock** — one-event-in-flight must hold over the pipe exactly as over the queue; the + clone-on-receive rule must not be skipped, or held-across-access tensors corrupt silently. +- **Path-only envoy fidelity** — the worker mirror must resolve every path the user writes; built from + the serialized tree, validated by the non-standard-named-model acceptance test. + +--- + +## 14. Warm worker pool — DONE (2026-06-07) + +**Why.** Spawning a worker per request is the dominant isolation cost: **~4.5 s** end-to-end on gpt2/A100 +(~12 ms in-process, ~370×), of which **~4.2 s** is cold `import torch` (1.3 s) + `import nnsight` (2.3 s) + +CUDA context init (0.4 s) — measured (`perf_spawn_cost.py`), **model-independent** (weights aren't shipped), +a flat per-request tax. Host-side mediator serialization is only ~3 ms. A warm pool amortizes the spawn. + +**The key change: the worker is generic, not mediator-bound.** Previously `_worker_main(payload, ...)` +received its mediator as a spawn-time argument. Now `_pool_worker_main(conn, buf, base_opts)` warms CUDA + +imports + `_ensure_mounted` **once**, optionally locks down, sends a one-time `"ready"` ack, then loops: +`conn.recv()` → on `("job", payload, extras, opts)` clear `Globals.saves`, deserialize a fresh mediator +against fresh dummies (`_run_one_job`), run it, loop; on `"stop"` `os._exit(0)`. The CUDA context, warmed +kernels, bounce buffer, and `CudaIpcWorkerChannel` persist across jobs; only the ~3 ms payload changes per +request. This **unifies** the cold and pooled paths — the worker always runs the loop; the host decides +recycle-vs-kill. The one-time `"ready"` is consumed by the spawner before the channel reads protocol frames +(they share the pipe). + +**Host side.** A process-global `_WorkerPool` (thread-safe) persists across traces. `acquire_isolated_worker` +serializes the mediator (`_build_job`), pulls an idle worker (or lazily grows to the `pool_size` cap, or a +cold one-shot worker past the cap so a trace never blocks on the budget), ships the job, and re-points the +host channel's `meta_provider`/`on_push` to *this* mediator's interleaver. `Mediator.cancel` calls +`release_isolated_worker`. + +**Recycle-safety rule.** Recycle **only a cleanly-ended worker**: `handle_end_event` sets `_iso.clean=True` +when a worker's END is consumed; `release` recycles iff `clean and poolable and alive and not dirty`. A +worker drained mid-protocol with a `Cancelation` (`dirty` — the pipe is now unbalanced), a timeout/death +(`clean` never set — it's *spinning*, not idle), or a cold one-shot worker is **retired** (killed) and the +pool re-warms lazily. Recycle resets the worker's host channel (`CudaIpcHostChannel.reset`) and per-job +hook-registration state; the worker rebuilds its interleaver + dummy modules per job and clears +`Globals.saves`, so there is no cross-trace state leak. + +**Opt-in.** `isolate_mediators(..., pool_size=N)` routes through the pool (`pool_size=0`, the default, is +the unchanged cold-spawn path). `warm_worker_pool(N, ...)` pre-warms at startup (blocks until N ack ready); +`shutdown_worker_pool()` tears it down. Workers are pooled **per (device, arena_bytes, gpu_mem_fraction, +lockdown, preimport) signature** (the `IsoOptions.pool_key` 5-tuple; `preimport` and `lockdown` are +warm-time, so they partition the pool); a worker is reused only for a matching signature, so a process hosting models on +different GPUs gets a per-device sub-pool (NOT a shared pool whose bounce buffer is fixed to the first +model's device — that would copy into the wrong-device buffer). Per-trace options (`default_all`, +`cross_invoker`, `timeout`) ride each job. + +**Pool sizing is a GPU-memory budget.** The natural ceiling is the **batch size** (one worker per mediator, +one mediator per invoke, all concurrent vs one forward pass → concurrent workers = #invokes ≤ batch size). +Each warm worker costs **~0.55 GiB GPU per GPU it touches** (CUDA context + cuBLAS kernels; measured, +model-weight-independent, linear in worker count) — and **MPS does not reduce this** (Ampere MPS shares the +*scheduler*, not context memory; measured identical under MPS). So at batch-16/single-GPU ≈ 8.7 GB (11% of +an 80 GB A100 but ~55% of a 16 GB T4) — the cap must be deliberate, with the cold-spawn fallback past it. +(`probe_pool_gpu_footprint.py`.) Note this ~0.55 GiB is the fixed CUDA-context cost and is distinct from +`set_per_process_memory_fraction`: that knob caps the worker's **allocator pool** (so a runaway allocation +can't exhaust the device, the 20 GB-symptom footgun), not the context, so it does not reduce the per-worker +0.55 GiB. + +**Lockdown + pool.** Seccomp lockdown is installed once, in `_run_one_job`, after the **first** job's +(host-authored, trusted) payload is deserialized and before its user code runs — so deserialization's own +imports (e.g. transformers' lazy modeling submodules, loaded only at unpickle time) succeed. It is one-way: +in a warm pool, later jobs run under the first job's lockdown, so a later job whose deserialize needs a +*new* import fails — true only across *different* models (a homogeneous model is already imported); pre-load +the deployment's model set via `preimport=` to serve heterogeneous models under lockdown. The cold path +(`pool_size=0`) is the same code, so it behaves identically (deserialize the one job, then lock down). +Lockdown defaults off. + +**Hardening (independent review, 2026-06-08).** Both passes confirmed no Critical issue — the cross-request +*data* invariant holds (per-job fresh interleaver/dummies/`Globals.saves`; host channel reset before +re-bind). Fixes applied: (1) `acquire` skips/forgets workers that **died while idle** and re-spawns, instead +of handing out a dead worker; (2) the pool is **keyed per device-signature** (above); (3) an **EXCEPTION- +ended worker is recycled** (it's alive + the pipe is balanced), not retired — so erroring traces keep the +pool benefit; (4) a recycled worker's first event uses `timeout + margin`, not the cold 180 s, preserving +hang-containment; (5) the grow slot is **reserved under the lock** so concurrent acquires can't exceed the +cap; (6) `close()` releases the pipe fd + GPU buffer; (7) a `_shutting_down` flag stops a shutdown/release +race from orphaning a worker. + +**Verified (`test_isolated_pool.py`, gpt2/A100):** reuse bit-identical (`max|Δ|=0`) at **~21× faster** once +warm (4.57 s cold → 0.22 s warm) with worker PIDs reused (no fresh spawn); a 3-invoke trace draws 3 distinct +pooled workers all bit-identical; a timed-out (infinite-loop) worker is retired and the pool re-warms with +the next trace bit-identical; a killed idle worker is skipped + replaced on the next acquire; a non-standard- +named model works through the pool. Cold path (`pool_size=0`) stays bit-identical across +read/swap/save/multi/exception/hang/multitoken/cross-invoke/barrier/nonstd. + +--- + +## 15. `tracer.cache()` — DONE (2026-06-08) + +**The gap.** `tracer.cache()` registers **persistent** hooks (`mediator_idx=inf`) that fill a `.save()`'d +`CacheDict` *during the forward*. In the worker those hooks land on the **dummy** modules and never fire, so +the user got an empty `CacheDict` (the `.save()`'d placeholder shipped back unfilled). + +**The fix — a `CACHE` event + host-side registration, with the post-forward injection collapsed.** +- **Worker `cache()` (isolated branch in `tracer.py`):** instead of registering dummy hooks, ship the spec + `(token, module-paths, device, dtype, detach, include_output, include_inputs, rename, alias)` via a new + `Events.CACHE` request (`mediator.send`). Return a **token-tagged** placeholder `CacheDict` the user binds + and `.save()`s — that is what carries the user's variable name across the boundary. +- **Host `handle_cache_event`:** resolve the paths to the **real** envoys, register the real + `cache_output_hook`/`cache_input_hook` into a host `Cache` keyed by the token (on `Mediator._iso_caches`), + `set_user_cache`, ack. Hooks live on the host mediator's `hooks`, removed at teardown by `remove_hooks` — + exactly like in-process. +- **The timing insight (no separate teardown step).** `handle_cache_event` acks and returns `True`, so the + host loop processes `CACHE` then `END` consecutively at `Mediator.start` — *before* the forward. So + `handle_end_event` swaps the **host** `CacheDict` reference in for the worker's empty placeholder (matched + by the token on the saved value) when it injects the saves. The forward then fills *that same object* + in-place. The user's variable **is** the forward-filled host cache; the doc's earlier "token-matched + post-forward injection" collapses to "swap in the host CacheDict at END; the forward fills it." This also + preserves in-process semantics, including "a cache defined after a module is called misses it" (the host + hooks register only when the `CACHE` event arrives, just like the in-process registration point). + +**Touch points:** `Events.CACHE`; `tracer.cache()` isolated branch; `handle()` dispatch + `handle_cache_event`; +the `handle_end_event` token-swap (gated on `_iso_caches`, so non-cache traces are untouched); `_iso_caches` +on `Mediator`. + +**Verified (`test_isolated_cache.py`, gpt2/A100):** single module, multi-module (3 keys), and +`include_inputs=True` all bit-identical (`max|Δ|=0`, keys match in-process). Cold path + the full isolated +regression (trace/acceptance/multitoken/cross-invoke/pool) unchanged. + +**Not covered:** a cache placeholder nested inside a container save (`got = [t.cache().save()]`) — the swap +matches a top-level saved `CacheDict`; nesting would need a recursive walk. `cross_invoker` + cache is +untested. `modules=None` (cache *all* modules) registers a hook per module on the host — correct but heavy. + +--- + +## 16. `with tensor.backward()`: read-path DONE (2026-06-10), grad-through-swap DONE (2026-06-23), batched (list) verified (2026-06-28) + +**The gap (§11).** Clone-on-receive strips `grad_fn` — the worker's delivered activations are detached +clones, the autograd graph lives only on the host, and `.grad` providers were keyed by `id(tensor)` +(process-local). So a backward block in the worker had nothing to differentiate. + +**The fix — split the chain rule at the process boundary** (distributed-autograd-style stitch), keyed by +**requester string** (module path + kind + step), not `id(tensor)`: + +- **Worker, forward time:** when the trace contains a backward block (detected from the intervention + source), every delivered activation clone is tagged — `requires_grad_(True)` + an `id(clone)` → + requester-string provenance map — via a wrapped `mediator.request` (`_tag_delivered`, + `isolation.py`). Worker-side ops on delivered values therefore build the worker's local half of the + graph (delivered leaves → loss). +- **Worker, backward time:** `BackwardsTracer.execute` detects the isolated context + (`worker_backward_context()`) and runs `_execute_isolated` (`tracing/backwards.py`): it computes the + worker half of the chain rule — `dL/d(delivered clone)` for each tagged leaf the loss depends on + (`torch.autograd.grad`, `allow_unused`) — and ships the seed dict `{requester: grad}` via a new + `Events.BACKWARD`. +- **Host:** `handle_value_event` retains the REAL on-graph activation per requester (gated on + `_iso_backward` so non-backward traces pay nothing); `handle_backward_event` continues the chain rule + on the host's real graph (`torch.autograd.grad` seeded by the worker grads, targets = all retained + activations, `allow_unused`, `retain_graph`) and returns `{requester: dL/d(activation)}`. No user code + runs on the host. +- **Worker, block body:** runs locally; each `.grad` read is served from the returned dict by the + tensor's provenance path (a patched `Tensor.grad` property). `.grad` on a user-derived tensor (no + provenance) raises a clear error; the `.grad` setter raises `NotImplementedError`. + +**Stitch correctness (no double-counting):** the worker seeds are partials treating delivered leaves as +independent; the host's autograd adds the indirect inter-layer contributions. Verified empirically — a +loss using the same read activation twice is bit-identical to in-process. + +**Verified (`test_isolated_backward.py`, gpt2/A100):** grad of `ln_f.output` `max|Δ|=0` vs in-process; +renamed model (`final_norm`/`output_projection`) `max|Δ|=0`; user-derived-tensor `.grad` → clear error. +Independent review (7 finder angles, dedup + verify): **no silent-wrong in the in-scope path** (single +invoke, on-path tensor-output target, scalar loss, no swaps). + +**Gradient-through-swap — DONE (2026-06-23).** An isolated SWAP installs a worker-computed value as a host +*leaf* (clone-on-receive strips `grad_fn`), so the host backward used to dead-end at the swap — while +in-process gradients flow through swaps. Now the seam is stitched by iterating the existing `Events.BACKWARD` +exchange to a fixpoint. Host: `handle_swap_event` (under `_iso_backward`) makes the swap leaf +`requires_grad_(True)` and retains it (`_iso_grad_swaps`), so the downstream forward tracks it and it is a +backward target; `handle_backward_event` adds swap leaves to its targets and returns `dL/d(swap leaf)` under +a reserved key. Worker: `WorkerMediator.swap` keeps the worker-tape swap value (with `grad_fn`); the backward +block loops — send seeds, receive `dL/d(swap leaf)`, backprop it through the swap tape to `dL/d(delivered +clone)`, re-seed the pre-swap graph, repeat — accumulating each read's gradient across rounds (a clone +reached both directly and through a swap sums both paths). With no swaps the loop is the original single +exchange. Chained swaps converge in N rounds. **Verified (`test_isolated_grad_through_swap.py`, gpt2 + +renamed):** grad through `h*2`, `h+vec`, `tracer.steer`, and TWO chained swaps, plus the renamed model, all +isolated-vs-in-process `max|Δ|=0`. + +**Limits / open:** +- Scalar loss only — `loss.backward(gradient=...)` is not honored (scalar-only error). +- **Batched backward (list input) works.** `model.trace([A, B, ...])` is one mediator over the padded + batch (no per-invoke narrowing, `batch_group=None`), so the worker's clone and the host's retained real + are both full-batch and their shapes match; the read-path and grad-through-swap seam-stitch run unchanged + on a `(batch, seq, hidden)` tensor. Verified bit-identical isolated-vs-in-process in + `test_isolated_batched_backward.py` (2-row, 3-row, upstream-block, batched grad-through-swap, renamed). +- Backward inside MULTIPLE `tracer.invoke(...)` contexts is **not supported in-process either**: it raises + `MissedProviderError` (the `.grad` provider is never registered across invoke contexts), so it is a core + nnsight limitation, not an isolation gap. (The earlier "shape mismatch / narrowed retention" framing + predated checking the in-process baseline.) The isolated path also fails; making that failure as clean as + the in-process error is minor parity hygiene, not a capability. +- Multi-token backward: **not supported in-process either** (characterized 2026-06-10, + `test_isolated_multitoken_backward.py`) — `generate()` runs the forward without gradient tracking, so + the first `.grad` read fails in-process ("cannot register a hook on a tensor that doesn't require + gradient"); no silent-wrong is possible (there is no graph at all, so the earlier "per-step retention + overwrite" concern is moot). The isolated path fails at the same user line; the host signals the + no-graph case (`handle_backward_event` returns a marker when nothing retained requires grad) so the + worker's error names the real cause instead of "off the backward path". +- Efficiency deferred: the host computes + ships grads for ALL retained reads (`retain_graph` always + on); a large read set can overflow the 64 MB arena. +- The `".backward("` source-substring detection can false-positive (e.g. the string in a comment), + which only costs needless tagging — tightening is planned alongside the gate consolidation. +--- + +## 17. Triton-kernel models — the deployment motivation (and why this beats the in-process whitelist) + +**Why this is urgent.** Frontier GPU model execution increasingly runs **Triton JIT kernels**, by three +independent routes: (1) architectures whose core op has no fast eager form ship Triton kernels — fused MoE +(vLLM `fused_moe`: Mixtral, DeepSeek-V2/V3/R1, Qwen-MoE, Llama-4, …) and SSM selective-scan (HF `mamba2` +imports `mamba_ssm.ops.triton.selective_state_update`; Jamba/Bamba/FalconMamba/Zamba); (2) HF's `kernels` +library + `kernelize(model, mode=inference)` pulls Triton kernels from the Hub as drop-in +norm/activation/attention replacements (Liger fuses RMSNorm/RoPE/SwiGLU/CE); (3) `torch.compile` → +TorchInductor, whose GPU codegen target *is* Triton. The triton-free path is essentially limited to a plain +dense transformer in eager mode with SDPA / precompiled-CUDA FlashAttention (which is an AOT library call, +not runtime codegen). + +**Why the in-process whitelist can't serve them.** NDIF's sandbox +(`ndif:src/ndif/services/ray/nn/security/`) is a Python import whitelist that wraps the ENTIRE +`tracer.execute(model)` — the forward runs *inside* the Protector (unavoidable in-process: user +interventions interleave with the forward in one call stack). Triton's first-use compilation needs +`subprocess`+ptxas, `tempfile`, `os`, `importlib`, `open()` — none whitelisted, several hard-blocked by the +audit hook. So the model's own Triton kernels, firing during the sandboxed forward, are denied. A +pre-warmed kernel cache doesn't save it: per-request shapes recompile, and even a cache *hit* still needs +`os`/`open`/`importlib` to load the cached `.so`. + +**Why this backend fixes it — the trust boundary aligns with the capability boundary.** The forward runs on +the TRUSTED host (unrestricted → Triton compiles normally); only UNTRUSTED user intervention code runs in +the worker. The split lands on the seam the whitelist couldn't find in-process: the thing needing dangerous +capability (runtime codegen) is the *model* (trusted); the thing we distrust is *user code*. The worker's +lockdown level is therefore **orthogonal** to whether Triton models run — worker security can be cranked to +the max without touching model Triton. + +**Strictly better than upstream on the module-restriction axis.** Against the right baseline (upstream +*sandboxed* NDIF, not un-sandboxed nnsight): + +| | Upstream whitelist | This backend (lockdown on) | +|---|---|---| +| user-code imports | restricted (allowlist) | restricted (seccomp import-freeze) — match-or-stricter | +| model serving stack | **also restricted** → Triton broken | **unrestricted** (host-side) → Triton works | +| user escape past Python layer | reaches host / tenants (whitelist README concedes C-level escapes; suite: 10/10 escapes succeed) | contained in a refless, killable process | + +The seccomp **import-freeze** (under lockdown, only warm-time-loaded modules are importable, since +import == `open()`, which seccomp then blocks) is therefore *not a regression* — upstream restricts user +imports at least as hard. The one parity item is *which* modules a legit user may reach; close it with +`isolate_mediators(..., preimport=(...))` (and the matching `warm_worker_pool(..., preimport=(...))`), which +loads the deployment's allowed-module set at warm time, before lockdown. User-code Triton is **doubly +blocked** (import wall if not pre-warmed + compile wall: ptxas execve / cache open) — intended containment, +documented as unsupported. + +**Verified properties (code, not just design):** +- *Timeout directionality is correct for slow host compiles.* The worker waits on the host with **no + timeout** (`CudaIpcWorkerChannel.wait_response` — a blocking `recv`); the timeout lives only on the host + (`CudaIpcHostChannel.wait_event`) and measures worker think-time. Host-side Triton compilation happens + during the host's *forward execution* — never while the host is in `wait_event` — so a multi-second cold + MoE autotune never false-trips the worker's hang-detector, in either direction. +- *Lockdown ordering / cold-vs-pool.* `lock_down()` runs once in `_run_one_job`, after the first job's + payload is deserialized and before its user code runs — in the unified `_pool_worker_main` (the only + worker entrypoint, cold via `poolable=False`, pooled via `poolable=True`). So **both** paths deserialize + their first job before lockdown; the cold-vs-pool difference is recycle-vs-retire. The first `conn.recv` + (which unpickles the job message) runs unlocked, so a fresh worker's first job needs no `preimport=`; a + warm pool's *later* jobs run under the first job's lockdown (so a different model would need `preimport=`). + +**Coverage:** `prototypes/mediator-sandbox/gpu_sandbox/test_isolated_triton_model.py` — `host_compiles` +(isolated + `lockdown=True` Triton-kernel model bit-identical to in-process, with a cold `TRITON_CACHE_DIR` +shown to populate during the isolated run) and `user_contained` (worker-side Triton blocked under +lockdown, exercising `preimport=`). Requires a GPU + a Triton install. diff --git a/docs/developing/mediator-isolation-harness-plan.md b/docs/developing/mediator-isolation-harness-plan.md new file mode 100644 index 000000000..c89e1a782 --- /dev/null +++ b/docs/developing/mediator-isolation-harness-plan.md @@ -0,0 +1,244 @@ +# Mediator Isolation Harness — Implementation Plan + +**Status:** Plan (pre-implementation) · **Date:** 2026-06-05 +**Builds on:** [mediator-isolation-sandbox.md](mediator-isolation-sandbox.md) (design + threat model), +`prototypes/mediator-sandbox/p2_isolation_poc.py` (the isolation proof-of-concept, passing). +**Goal of this doc:** turn "the jail is safe" into "a real trace runs *through* the jail," in testable +increments, with the exact nnsight seam identified. + +--- + +## 1. The seam (where host and jail divide) + +Three call sites in `src/nnsight/intervention/interleaver.py` define the entire boundary. Everything +above them is user code (→ jail); everything below is the model (→ host). + +| Role | Symbol | What it does | +|---|---|---| +| **Client hook** | `eproperty.__get__` (`:275`) | `value = interleaver.current.request(requester)` — `.output` read | +| **Client hook** | `eproperty.__set__` (`:317`) | `interleaver.current.swap(...)` — `.output` write | +| **Worker launch** | `Mediator.start._worker_target` (`:981`) | `Thread(target=…)` runs `intervention(self, info, …)` — `self` is the Mediator | +| **Worker→host RPC** | `Mediator.send` (`:1313`) | `event_queue.put((event, requester))` → wait `response_queue` | +| **host→worker reply** | `Mediator.respond` (`:1299`) | `response_queue.put(value)` → wait `event_queue` | +| **the two queues** | `Mediator.Value` (`:823`), fields `:878-879` | lock-based one-slot handoff (the transport) | +| **server narrow/swap** | `Batcher.narrow` (`batching.py:198`), `Batcher.swap` (`:226`) | slice/replace the tenant's rows | + +**Key fact:** the compiled intervention fn never calls the model directly. It only ever reaches +`interleaver.current.request/swap/skip` and reads scalar state (`iteration`, `history`, `transform`). +So the jail needs **no weights and no GPU** — only an Envoy *path tree* wired to a client-side mediator +that emits the six events (`VALUE/SWAP/SKIP/BARRIER/END/EXCEPTION`, `Events` `:349`) over a socket. + +--- + +## 2. Target architecture + +``` +HOST (trusted: GPU + weights) JAIL (untrusted: CPU-only, net=none, ro-fs) + model + Envoy + weights deserialize request bytes ← pickle lives here + Interleaver + Batcher ── admission bounds client Envoy *path tree* (no modules) + MediatorProxy (per request) ──────── socket ──────── client Interleaver stub (.current = RemoteMediator) + · runs handle_value_event/handle_swap_event RemoteMediator.send() → socket + against the REAL Batcher, narrowed to the compiled intervention fn = USER CODE + host-recorded bounds (NOT worker-supplied) (torch-CPU on the delivered tensor) + · D2H + serialize VALUE replies + · deserialize + H2D SWAP tensors +``` + +- **`MediatorChannel`** — a new abstraction that replaces the `Mediator.Value` queue pair. Two impls: + `InProcessChannel` (today's lock+slot) and `SocketChannel` (length-prefixed frames over an + `AF_UNIX` socket). The Mediator's `send`/`respond` and the worker loop talk to a `MediatorChannel`, + not to `event_queue`/`response_queue` directly. +- **`RemoteMediator`** (jail side) — implements the subset the intervention fn touches: + `request/swap/skip/end/exception`, `iteration`, `transform`, `history`. `send()` writes a frame and + blocks for the reply. **It exposes no `interleaver`, `batcher`, or sibling refs** — the capability + leaks have nothing to reach. +- **`MediatorProxy`** (host side) — owns the real `Batcher` and the existing `handle_*` logic; per + socket event it narrows/swaps against **host-recorded admission bounds**, so `batch_group`, + `narrow(None)`, `[-1,_]`, and widening become unrepresentable from the jail. + +--- + +## 3. What crosses the wire (the codec) + +| Frame | Direction | Payload | +|---|---|---| +| `VALUE` req | jail→host | `requester` (provider path str) + iteration | +| `VALUE` reply | host→jail | the **narrowed activation as a CPU tensor** (D2H'd, serialized) | +| `SWAP` | jail→host | `requester` + the **modified CPU tensor** (serialized) | +| `SKIP` | jail→host | `requester` + sentinel/value | +| `BARRIER` | jail→host | participating mediator names | +| `END` / `EXCEPTION` | jail→host | none / serialized exception metadata | + +Tensor codec: start with the dumb path (`torch.save`/raw buffer / `safetensors`) over the socket; +a shared-memory ring is a later optimization (out of scope for the harness). Values may be **tuples / +nested structures** (transformer block outputs) — the codec must round-trip `applyn`-style nested +containers, not just bare tensors. + +--- + +## 4. Phased plan (each phase is independently testable) + +### Isolation proof-of-concept ✅ DONE +`prototypes/mediator-sandbox/p2_isolation_poc.py` — the 10 escape gadgets run inside the jail and are +inert; socket-fd-into-`net=none`-jail mechanism verified; real PID isolation confirmed. + +### `MediatorChannel` seam (in-process, zero behavior change) ✅ DONE +Extract the `event_queue`/`response_queue` handoff in `Mediator.send`/`respond`/`start` and the worker +loop behind a `MediatorChannel` interface; ship `InProcessChannel` as the default. +- **Touch:** `interleaver.py` — added `MediatorChannel`/`InProcessChannel`; all 23 queue call sites now + route through `self.channel.*`. +- **Acceptance:** the *entire existing nnsight test suite passes unchanged.* — **met** (203 core tests + green; independent review verdict MATCH). + +### Socket transport for the protocol (process boundary, harness) ✅ DONE +The six-event protocol rides an `AF_UNIX` socket between a forked worker and the host, with the real +`Mediator`/`Batcher`. `SocketHostChannel` (host-local event buffer; only `wait_event` reads the socket) +and `SocketWorkerChannel` (`src/nnsight/intervention/transport.py`) + a length-prefixed pickle codec. +- **Proven** (harness `prototypes/mediator-sandbox/phase2_socket_transport.py`, all green): + codec round-trip of nested tuples/dicts/tensors (the transport itself is **bit-exact**); + **golden equivalence** — same `h[6]×2` intervention local vs over-socket on gpt2 → logits match within + multi-threaded CPU forward noise (`allclose atol=1e-3`; residual ~1e-4 is two separate forward passes, + not the socket — a broken swap is Δ≈30); **real batching** + (`needs_batching=True`) — the worker receives only its `batch_group` row and its swap lands only on + that row (no cross-row leak); **out-of-order `restore_event`** exercised at runtime over the socket. +- **Scope honesty:** this proves *the event protocol + the `provider==requester` routing/restore + real + Batcher narrow/swap over a process boundary*. The requester string is a shared constant in the harness + (the compare at `interleaver.py:1175` is exercised; independent eproperty-side *production* of that + string rides in with the `model.trace()` integration). It does **not** yet route through + `model.trace()` — see the transparent-trace-integration milestone below. + +### Transparent `model.trace()` isolated execution (the trace integration) +Make `model.trace()` itself launch the worker in a separate process and return identical results — the +isolation *execution backend*, separated out because it is the bulk of the real feature. +- **`Mediator.start`** gains an isolated path: fork (the trace-integration milestone) / exec-the-jailed-worker (the jail-the-worker stage) + instead of `Thread`; the host Mediator gets a `SocketHostChannel`, the worker a `SocketWorkerChannel`. +- **Saves transmission (keystone finding):** `.save()` marks ids in `Globals.saves` and the values live + in the *worker's* address space (`tracing/base.py:567`). The worker must **ship saved values back to + the host** (e.g. a final `SAVES` frame at END) and the host must inject them into the trace's return — + otherwise an isolated trace returns empty saves. +- **Lifecycle parity:** `cancel()` must `waitpid`/`kill` the worker process (not drop a `Thread`); + `check_dangling_mediators`, EarlyStop drain, and EXCEPTION/SKIP over the socket must all hold. +- **Acceptance:** `with model.trace(prompt): h[6].output.save()` under an `isolate_mediators()` context + returns **bit-identical** saves to the in-process trace, across read / `+=` write / multi-layer / + multi-token. + +### Jail the worker ✅ DONE +Exec the protocol worker inside a `bwrap` jail (CPU-only, `net=none`, ro-fs allowlist), socket fd passed in. +- **Proven** (`prototypes/mediator-sandbox/phase3_jailed_worker.py` + `phase3_jail_transport.py`, green): + (A) the real Mediator protocol worker, run *inside the jail*, produces matching gpt2 logits + (`allclose atol=1e-3`; the transport codec is bit-exact, residual is forward nondeterminism); + (B) the same jailed worker's escape gadgets are inert — `fs_read`→FileNotFoundError, + `net_egress`→OSError, `os.system` touch creates **no host file**, secret **not** leaked — *while its + legitimate request/swap protocol still completes bit-identically*. Host-level sibling = the host-level sibling jail (sandbox design §6.5). +- **Note:** the jailed worker runs a fixed script (not yet a host-shipped compiled intervention) and the + requester is still a constant — shipping the real intervention is the GPU-path deserialize-in-jail work. + +### Batcher as authority (cross-tenant / Boundary B) ✅ DONE +The narrow/swap bounds live on the **host** Mediator (the proxy); the jail's Mediator is a separate object +with no `interleaver`/`batcher` reference, so the worker physically cannot mutate the host's bounds or +reach a sibling. The 8 capability leaks (#1-6, #8; #7 pacing is out of scope) are inert **by construction**, +not by a signature change to `narrow`/`swap`. +- **Proven** (`prototypes/mediator-sandbox/phase4_malicious_worker.py` + `phase4_cross_tenant.py`, green): + a malicious jailed tenant co-batched with a victim row attempts `batch_group=None`/`[-1,0]`/widen and + walks to `interleaver.batcher`/`interleaver.mediators`. Result: all batch_group claims are + "set-on-local-mediator-only" (no host effect); `has_interleaver=False`; both walks → `AttributeError`; + the worker received **only its own row** (sum=12, not 96 → never saw the victim's secret); its poison + landed **only on its own row** — the victim row stayed uncorrupted. +- **Note:** demonstrated with one malicious jail + a passive victim row (the property is symmetric for two + jails); BARRIER cross-jail coordination remains deferred (§5). + +### Real-model GPU path + measurement ✅ DONE +Real `D2H → socket → torch-CPU → H2D` round-trip with the model on an A100, worker CPU-only. +- **Proven** (`prototypes/mediator-sandbox/phase5_gpu_measure.py`, green): (A) the GPU D2H/H2D socket + path matches the in-process GPU forward (`max|Δ|=0`, atol 1e-2); (B) **measured per-hook overhead** + (TOTAL = D2H + sockRTT + H2D) on an A100 80GB PCIe: 1×16×768 (0.03 MB) **0.9 ms**; 1×512×4096 (4.2 MB) + **30 ms**; 1×2048×4096 (16.8 MB, ~7B) **111 ms**; 1×2048×8192 (33.6 MB, ~70B) **376 ms** — vs a + **measured** in-process baseline of **0.1–0.25 ms** (zero-copy reference). +- **Finding (the two-tier-decision input):** the cost is **dominated by pickle + socket RTT** (103 ms of the + 111 ms at 16.8 MB; pickle alone ~14 ms), *not* the PCIe D2H/H2D (~3.5 ms each). A 32-layer 7B cache ⇒ + ~3.5 s added. So the naive path is fine for small/few-layer interventions, but big-model per-layer work + needs **the two-tier approach** (run the common tensor algebra host-side) and/or a **shared-memory ring** wire + format — confirming §7. +- **Deferred:** deserialize the *actual* serialized intervention request in the jail (untrusted + unpickle-in-jail). The security property is already shown (the isolation-PoC + jail stages: pickle `__reduce__` gadget inert + in the jail); shipping the *real compiled intervention* is the trace-integration work. + +### Transport breakdown: the cost is the CODEC, not the boundary ✅ DONE +The Phase-5 "sockRTT-dominated" number was really **pickle-dominated**. `prototypes/mediator-sandbox/ +phase5b_transport_breakdown.py` isolates it (16.8 MB, CPU↔CPU, echo = pure transport, no user op): +- The socket moves 16.8 MB in **~10 ms**; the naive path pickles the tensor **4× per hook** + (`dumps` ~14 ms + `loads` ~8 ms each way), and pickle scales **superlinearly** (285 ms at 33.6 MB). +- **Measured per-hook round-trip by wire codec:** + + | codec | 4.2 MB | 16.8 MB (~7B) | 33.6 MB (~70B) | speedup | + |---|---|---|---|---| + | **A pickle** (today, `transport.py`) | 12.5 ms | 95.9 ms | 285 ms | 1× | + | **B raw message** (struct header + raw bytes via `memoryview`/`torch.frombuffer`; = `safetensors`) | 4.3 ms | 24.1 ms | 123 ms | **~4×** | + | **C shared memory** (`mmap`, only a 1-byte signal on the socket; bulk never travels) | 8.5 ms | 11.5 ms | 28.4 ms | **~8×** (linear) | + +- For **read-only** interventions (caching, logit lens — the common case) shmem is even cheaper: the host + writes the activation once, the worker reads it zero-copy, only a small result returns. +- Separate finding: the worker's `×2` on **bf16 CPU is ~9 ms** by itself (the *user* op, not transport) — + another argument for **the two-tier approach** (keep common ops on the GPU host-side; neither transport nor the + slow CPU recompute happens). +- **Fix hierarchy:** (1) swap the codec pickle→`safetensors`/raw in `transport.py` — small, localized, 4×; + (2) shared-memory ring for big-model per-layer work — 8×, linear; (3) **the two-tier approach** removes transport + entirely for the common tensor algebra. The 111 ms was a codec artifact, **not** intrinsic to isolation. + +### Shared-memory + safetensors transport ✅ DONE (built into `transport.py`) +`ShmArena` (anonymous `memfd`) + `ShmSocketHostChannel`/`ShmSocketWorkerChannel`: tensor bulk rides the +shared memfd, only a small control frame crosses the socket, tensors encoded with **safetensors** (safe, +no-code-exec — also closes the jail→host untrusted-deserialize hole for the bulk). +- **Proven** (`phase6_shm_safetensors.py` + `phase6_jailed_worker.py`, green): correct through the real + `Mediator` protocol both **forked** and **jailed** (gpt2 golden `max|Δ|=0`; the memfd is passed into the + bwrap jail via `pass_fds`+`SHM_FD`). Echo round-trip vs pickle: 16.8 MB **87→44 ms (~2×)**, 4.2 MB ~1.8×. +- **Why only ~2× (not the 8× of hand-rolled shmem):** `safetensors` copies on `save` AND `load` (it returns + owned tensors), so the bulk is memcpy'd ~4×. The hand-rolled `frombuffer` path was zero-copy (a view into + the arena) but that aliases the shared buffer — needs double-buffering / careful lifecycle. Safe-and-simple + (copy) vs fastest (zero-copy view) is the next dial; both beat pickle. +- **Still pickle:** the small control skeleton (non-tensor structure + requester string). Hardening that to a + restricted decoder is the remaining jail→host security item. +- **Bigger lever (open):** true zero-copy needs the data to not move at all — either **the two-tier approach** + (op runs on the GPU host-side; nothing transported) or **GPU-in-jail + CUDA-IPC** (jail maps the host's GPU + buffer; no D2H/serialize) — the latter trades the CPU-only security posture; see the GPU-access discussion. + +--- + +## 5. Open design decisions (resolve as they're reached) + +- **Client Envoy tree in the jail** — ship a path-only Envoy, or reconstruct from the deserialized + request (the nnsight request already carries Envoy refs)? Lean: reconstruct from the request, since + deserialize-in-jail produces it anyway. +- **`cross_invoker` var-sharing** (`send` push/pull, `:1327`/`:1410`) — today two invokes share a Python + frame in-process. Across jails that frame doesn't exist. v1: **one jail per request, `cross_invoker` + disabled**; revisit a host-mediated var store later. +- **`BARRIER`** coordinates multiple mediators (= multiple jails); the host-side `MediatorProxy` set + becomes the coordinator. Defer multi-jail barriers past the cross-tenant stage. +- **CUDA stream propagation** (`start` captures the caller stream, `:967-977`) — irrelevant in a + CPU-only jail; the host keeps all stream handling. +- **Tensor wire format** — raw/`safetensors` first; shmem ring is a perf follow-up, not correctness. + +--- + +## 6. Risks + +- **Per-hook latency** — the accepted cost; quantified in the GPU-path measurement, drives the two-tier approach. +- **Codec correctness** for nested/tuple values (transformer block outputs) — needs a robust + `applyn`-aware codec, tested against the gotchas (tuple-output handling). +- **Lifecycle parity** — the in-process `cancel()`/worker-join must map cleanly to **SIGKILL the jail**; + verify END/EXCEPTION/EarlyStop all drain correctly across the socket. +- **Deadlock surface** — the one-event-in-flight invariant must hold over the socket exactly as over the + one-slot queue, or the "access in forward-pass order" contract breaks. + +--- + +## 7. Test strategy (every phase) + +1. **Golden equivalence** — in-process result == isolated result, on varied models/ops (not just gpt2; + include a non-standard module-named model per the testing rules). +2. **Security** — the jail-the-worker stage runs the 10 escape gadgets as the intervention (inert); the cross-tenant stage runs the 8 + capability-leak payloads across two jails (blocked). +3. **Lifecycle** — exception propagation, early-stop, timeout→SIGKILL, normal drain. + +The two existing ndif suites (`tests/security/test_protector_escapes.py`, +`test_mediator_capability.py`) are the security gates; each fix flips a row from 🔓 to blocked. diff --git a/docs/developing/mediator-isolation-sandbox.md b/docs/developing/mediator-isolation-sandbox.md new file mode 100644 index 000000000..02eb394ed --- /dev/null +++ b/docs/developing/mediator-isolation-sandbox.md @@ -0,0 +1,324 @@ +# Mediator Isolation Sandbox — Design + +**Status:** SUPERSEDED (historical) · **Date:** 2026-06-05 · **Author:** zikai +**Related:** `ndif` security regression suite (`src/services/ray/tests/security/`), `NDIF.md` §6–7, +nnsight `src/nnsight/intervention/interleaver.py` (Mediator / event protocol). + +> **Superseded.** This is the earlier CPU-only, two-tier (fork-jail) design from 2026-06-05. On 2026-06-06 +> the project chose the GPU-sandbox approach (contain footguns, not a determined adversary), built and +> documented in [gpu-sandbox.md](gpu-sandbox.md) and landed via the integration in +> [mediator-gpu-trace-integration.md](mediator-gpu-trace-integration.md), which is the authoritative +> current reference (plus [mediator-threat-models.md](mediator-threat-models.md) for the security posture). +> Kept for historical context; the CPU-transport and two-tier op-interpreter details below do not reflect +> the shipped code. + +## 1. Motivation + +NDIF executes **arbitrary user-submitted Python** on shared GPU infrastructure. Today that code runs +**in-process** with the model forward pass inside a single Ray `ModelActor`, behind an in-process +Python whitelist (`Protector`, `src/services/ray/src/ray/nn/security/protected_environment.py`). + +The security regression suite establishes that this boundary is **architecturally unfixable in-process**: + +- **10/10 in-process sandbox escapes succeed** (`().__class__.__mro__…__subclasses__()`, `str.format` + globals walk, `ctypes`/`CDLL`, `torch.utils.cpp_extension.load`, `CodeType`, frame walking, pickle + `__reduce__`, end-to-end `os.system` RCE). The literature consensus (pysandbox post-mortem; Modal, + Lambda, Snowpark practice) is that containment must move to **OS-level isolation**. +- **8/8 multi-tenant Mediator-capability leaks succeed** (mutating `batch_group`, walking + `mediator.interleaver.batcher.current_value`, poking sibling mediators, `swap(None, …)` poisoning all + tenants). All share one root cause: the `Mediator` object is handed to user code with **mutable, + reachable references** to scope and to the shared `Interleaver`/`Batcher`. + +This document designs a refactor that **truly isolates each Mediator** from the model-hosting process, +so user code can no longer reach the host, the GPU, the weights, or other tenants' activations. + +## 2. Core insight: the event protocol is already the boundary + +A `Mediator` (`interleaver.py:718`) already communicates with the forward pass **only** through six +typed messages over two one-slot queues (`interleaver.py:338`, `:949`): + +| Event | From | Means | +|---|---|---| +| `VALUE` | worker | "give me the value at this provider path" (worker blocks) | +| `SWAP` | worker | "replace the value at this provider path with this tensor" | +| `SKIP` | worker | "bypass module X's real forward" | +| `BARRIER` | worker | "synchronize me with these other mediators" | +| `END` | worker | "I'm finished — drain me" | +| `EXCEPTION` | worker | "I crashed; here's the exception" | + +This is already a clean **client/server RPC**. It is insecure only because client and server share one +Python interpreter, one address space, and one GPU context — so the worker can ignore the protocol and +walk Python objects directly. + +> **Isolation = keep the six messages, move the worker to the far side of a real OS boundary, and change +> the transport from an in-process queue to a socket.** Every in-process escape and every +> Mediator-capability leak then disappears *structurally* — there is no longer a `mediator` object, a +> `batcher`, or a host module in the user's address space to reach. + +## 3. Threat model (scoped) + +### In scope +- **A. Host protection** — user code must not reach the host OS, filesystem, network, model weights, or GPU. +- **B. Cross-tenant protection** — when tenants are co-batched in one forward pass (lands with + `ndif-batch`/`ndif-vllm`), tenant X must not read or poison tenant Y's rows. Solved two ways: (1) the + worker holds *no host references*, and (2) the host-side `Batcher` becomes the **authority** on row + bounds (below). +- **Deserialization RCE** — the pickle `__reduce__` gadget fires during `request.deserialize(...)` in + `pre()` (`NDIF.md:838`), *before* execution, under a whitelist that necessarily includes + `pickle`/`cloudpickle`. Untrusted bytes must therefore be unpickled **inside the jail**, never on the host. + +### Out of scope (deliberately deferred) +- **Mediator pacing / forward-pass-blocking DoS** (suite leak #7). This is a *scheduler* property, not an + isolation property. The host may keep blocking on the worker's reply exactly as today — just over a + socket. Revisit with the non-blocking-listener work, separately. +- **Side channels** (timing, cache, GPU memory residency). +- **Per-hook transport throughput.** Accepted as a cost for the prototype; see §7 (option D) for the + optimization path. + +### Trust boundary +- **Trusted:** model weights, GPU + CUDA context, forward pass, `Interleaver`, `Batcher`, host OS, other + tenants' data. +- **Untrusted:** the serialized request bytes (pickle), the compiled intervention function, anything the + user can influence. + +## 4. Design + +### 4.1 Process / thread model + +``` +┌─ ModelActor (TRUSTED: GPU + weights) ───────────────────────┐ +│ main thread: forward pass │ +│ Interleaver + Batcher ← server-side row-bound authority │ +│ per request: MediatorProxy (host stub, speaks 6 events) │ +│ warm pool of CPU-only jailed workers (torch+nnsight ready) │ +└───────────────┬─────────────────────────────────────────────┘ + │ per-request bidirectional socket (fd passed into jail) +┌───────────────┴─ jailed worker (UNTRUSTED, CPU-only) ───────┐ +│ 1. deserialize raw request bytes ← untrusted pickle HERE │ +│ 2. run compiled intervention fn → Mediator worker │ +│ 3. VALUE/SWAP/SKIP/BARRIER/END/EXCEPTION over the socket │ +│ net=none · ro-fs allowlist · non-root · mem cap · no GPU │ +└──────────────────────────────────────────────────────────────┘ +``` + +- The host keeps the forward pass, `Interleaver`, and `Batcher`. Per active request it spins a + **`MediatorProxy`** — what is today the in-process worker thread, demoted to a transport endpoint. It + receives the worker's events and answers them via `Batcher.narrow`/`swap`, **bounded to that tenant's + admission-time rows.** +- **The protocol is unchanged.** Only the transport relocates: `Mediator.start` launches a jailed + process instead of a `Thread`; the `Mediator.Value` one-slot queues (`interleaver.py:767`) become a + socket. `Mediator.handle` writes a value to the socket instead of `response_queue`; the worker's + `request(...)` reads/writes the socket instead of the in-process queues. `iter`, `barrier`, `next`, + and save-filtering keep working with no user-visible change. + +### 4.2 The activation data path (the accepted cost) + +`VALUE`: host `narrow`s to the tenant's rows → `D2H` copy → serialize CPU tensor → socket → jail +deserializes → user gets a **real CPU tensor** and runs arbitrary torch-CPU on it. `SWAP`: jail sends +the modified CPU tensor back → host deserializes → `H2D` → writes into the **bounded** batch slice. + +Start with the simplest serializer (raw buffer / `safetensors` over the socket); a shared-memory ring is +a later optimization, not part of correctness. This is **the per-hook D2H/H2D transport** — clean isolation, transport cost +accepted. See §7 for why and when to move to the two-tier approach. + +### 4.3 Batcher as server-side authority (Boundary B) + +Row bounds are recorded on the **host** at admission. The worker's `VALUE` carries only a provider path + +iteration — never a `batch_group` it can forge. `narrow(None)`, `[-1, _]`, and widening +(`batching.py:213–217`) become **unrepresentable** because the bounds live on the trusted side. This is +the one piece of genuinely new logic; everything else is transport relocation. + +### 4.4 CPU-only, and why + +The jail gets **no GPU**. This is load-bearing for both security and the warm pool: + +- No CUDA context in the jail → no GPU-driver attack surface, and no CUDA-IPC cross-tenant leak (CUDA IPC + shares at *allocation* granularity, which under vLLM's paged allocator can span other tenants' KV cache). +- **torch-CPU is fork-safe**, so a *zygote* that has already imported torch+nnsight can fork a fresh + child per request in ~ms. (A CUDA-initialized process is **not** fork-safe — GPU-in-jail would kill the + pool and is rejected.) + +### 4.5 Lifecycle, pool, cancellation + +- **One jail per request**, assigned to exactly one tenant, then **destroyed / re-exec'd** — never + soft-reused across tenants (a dirtied interpreter is a cross-tenant channel). +- **Warm pool:** first cut may spawn-per-request; then add a zygote pool (pre-imported, pre-jailed idle + workers claimed-then-destroyed) for sub-100ms starts. +- **Cancellation / timeout** becomes **`SIGKILL` the subprocess** instead of the current + `kill_thread()` `ctypes`-injected `SystemExit` (`NDIF.md:919`) — a free robustness win. + +### 4.6 Minimal surgery footprint + +Three touch points, which is what makes this prototype-friendly: +1. `Mediator.start` + the `Value` queue handoff → jailed process + socket. +2. `pre()` deserialization moves inside the jail. +3. `Batcher.narrow`/`swap` take host-held bounds instead of trusting the message. + +The forward pass, interleaver scheduling, and nnsight's public API are untouched. + +## 5. Isolation mechanism + +Three tiers, distinguished by **what enforces the boundary**: + +| Tier | Mechanism | Tools | Host-kernel surface | +|---|---|---|---| +| 1 | Linux namespaces + seccomp + cgroups (shared kernel) | bubblewrap, nsjail, runc/Docker | full (filtered) | +| 2 | User-space kernel (syscalls reimplemented) | gVisor (`runsc`) | minimal | +| 3 | microVM (own guest kernel, hardware virt) | Firecracker | none | + +**Choice for the prototype: Tier 1, CPU-only.** Compute-bound torch-CPU runs ~native; the socket +transport is light; one jail per request is cheap. gVisor (Tier 2) is the **pre-planned upgrade** behind +the same launcher interface if "shared host kernel" becomes unacceptable. Firecracker is overkill for a +CPU-only prototype. + +**Tool: nsjail strategically; bubblewrap as the bootstrap.** The probe (§6) shows: +- `nsjail` is **not packaged** on either the host (Ubuntu 20.04) or the container (Debian trixie) — it is + a source build. Its value is integrated seccomp-policy (kafel) + cgroups + rlimits + time limits, which + is what a production sandbox wants. +- `bubblewrap` (0.11.0) is a **one-line apt install** in the container and is **verified working** on the + host. It is a launcher (pair it with an external seccomp filter + cgroup), but gives the *identical* + Tier-1 boundary. +- The nesting constraints in §6 apply **identically to both**, so the bwrap-vs-nsjail choice is + independent of — and less important than — the deployment decision in §6.5. + +**Recommendation:** bootstrap on bwrap (verified, trivial install, swappable behind a `Jailer` interface), +target nsjail for the hardened path. + +## 6. Host-reality probe (verified 2026-06-05) + +All results are measured on this machine, not assumed. + +### 6.1 Bare host (Ubuntu 20.04, kernel 5.15.0-139) + +| Check | Result | +|---|---| +| `kernel.unprivileged_userns_clone` | `1` (enabled) | +| `user.max_user_namespaces` | `6185255` | +| `unshare -Urn true` (unpriv userns) | **OK** | +| `bwrap` full jail (userns+net+pid+fresh-proc+ro-fs) | **OK**, pid1=bwrap | +| Network isolation inside jail | **OK** (no DNS, loopback only) | +| ro-fs allowlist (unbound `/disk` invisible) | **OK** | +| `bwrap` present / `nsjail` present | `/usr/bin/bwrap` / **absent** | +| seccomp kernel support | `CONFIG_SECCOMP=y`, filter=y | +| cgroup version | v1 | + +→ **Tier-1 unprivileged jailing is fully viable on the bare host.** + +### 6.2 Inside `dev-ray-1` (local stand-in for the prod model container; Debian trixie, ndif/ndif:latest) + +| Check | Result | +|---|---| +| uid inside container | `0` (root-in-container) | +| Privileged / CapAdd / SecurityOpt | `false` / `[]` / `[]` (default posture) | +| AppArmor | `docker-default` | +| `unshare -Urn` (nested unpriv userns) | **BLOCKED** — "Operation not permitted" | +| torch | `2.9.1+cu128`, cuda available | +| `bwrap` / `nsjail` present | **absent** / **absent** (bwrap = 1-line apt; nsjail = build) | + +The block is **Docker's default seccomp profile** denying `clone(CLONE_NEWUSER)`; the container also +lacks `CAP_SYS_ADMIN`, so classic root-namespaces are unavailable too. Both jail paths are closed by default. + +### 6.3 What it takes to jail *inside* the container + +| Container config | Result | +|---|---| +| default | userns BLOCKED | +| `seccomp=unconfined` | userns OK; `bwrap` fails at "make / slave" (AppArmor) | +| `seccomp=unconfined` + `apparmor=unconfined` | propagation OK; **fresh `/proc` mount fails** in new PID ns | +| `+ cap-add SYS_ADMIN` | still fails fresh `/proc` mount | +| `--privileged` | **full jail (fresh proc) works** | +| `seccomp+apparmor=unconfined`, **bind `/proc`, drop PID-unshare** | **jail builds** (reached execvp) — but **no PID-namespace isolation** | + +**Finding:** nesting a userns mount-namespace jail inside an *unprivileged* container is fragile. Locked +container mounts and the fresh-`/proc` restriction mean a jail with **real PID isolation requires +`--privileged`** (or close), and the non-privileged workaround (bind the container's `/proc`) **leaks PID +visibility** — a cross-tenant info-leak vector if multiple workers share the container's PID space. + +### 6.4 Production target: AWS ECS-on-EC2 (the probe transfers) + +Production runs on **AWS ECS-on-EC2** (CDK `ndif-aws/service_stack.py:275`, `ecs.Ec2TaskDefinition`) — +not Fargate, not raw EC2. The model server (`Ray-Worker`) is an **ECS task (Docker container)** on a +**`g4dn.xlarge`** GPU instance (1× T4, **4 vCPU / 16 GB**; `env/dev.py`), on the **ECS-optimized Amazon +Linux 2023 GPU AMI** (`service_stack.py:220`). The whole GPU is allocated into the container +(`gpu_count`, `:365`); network mode is `awsvpc` (own ENI per task); models live on EFS (`/efs/huggingface`). + +Decisive detail: the task's `linux_parameters` is used **only** for `shared_memory_size` +(`service_stack.py:351-355`) — **no `privileged`, no added caps, no seccomp/apparmor override.** So the +prod model container runs the **Docker default posture, identical to local `dev-ray-1`.** + +→ **The §6.1–6.3 probe transfers to production unchanged:** nested unprivileged userns is blocked in prod +too (docker-default seccomp + no `CAP_SYS_ADMIN`, GPU inside the container). The local finding was +representative, not a local artifact. + +Two AWS specifics change the *decision* (§6.5): +- **ECS has no clean per-task seccomp knob.** ECS `dockerSecurityOptions` covers SELinux/AppArmor labels + only — not seccomp profiles. So the local "`seccomp=unconfined`" relaxation is **not** a per-task option + on ECS; the practical in-task lever is **`privileged=true`** (drops seccomp *and* adds all caps). + The in-container self-jail is therefore *heavier* on AWS than locally. +- **We own the EC2 instance, its user-data, and the host Docker daemon.** That makes the sibling approach + host-level-sibling deployment natural: pre-install the jailer in user-data and run sandboxes as **host-level siblings**, leaving + the model-server task unprivileged. + +*Adjacent finding (out of scope for this design, worth fixing): `ndif-aws/env/dev.py` commits **live +credentials** — an AWS access key/secret, a HuggingFace token, DB/Influx passwords. Rotate and move to SSM +Parameter Store / Secrets Manager.* + +## 6.5 Deployment decision (KEY OPEN ITEM) + +The nesting reality reshapes *where* the jail runs. Two resolutions, read through the AWS reality (§6.4): + +- **Self-jail inside a loosened model container.** Locally this is + `seccomp=unconfined, apparmor=unconfined` (+ effectively `privileged` for real PID isolation). + **On ECS it collapses to `privileged=true`** — there is no clean per-task seccomp knob — set via the + container's `linux_parameters`/props (`service_stack.py:351-363`). Fastest to a testable boundary, **but + on AWS it means a privileged model task**, which broadens that task's own surface more than the local + case did. Justified only as a throwaway first cut to exercise the protocol. +- **Sibling sandbox on the host (recommended target, especially on AWS).** Run the jailed worker as a + **separate process/container at the EC2-instance level**, where host userns + bwrap work (verified on the + local host §6.1; AL2023 host expected to match — needs a one-line on-instance confirm). Pre-install the + jailer in EC2 user-data; the model-server ECS task stays **unprivileged** and talks to the sandbox over a + socket. This is how production code-exec services isolate (sibling sandboxes, not nested jails). Cost: a + host-level launcher/broker and its trust handling (don't hand the host Docker socket to the model server). + +**Recommendation:** **the host-level sibling jail is the production target** — owning the EC2 host makes it clean and avoids a +privileged model task. The in-container self-jail is justified only as a throwaway first cut; on ECS even that costs +`privileged`. Keep the launcher abstracted so in-container → host-level (and basic-seccomp → gVisor) don't touch anything above it. + +**Sizing caveat:** `g4dn.xlarge` is **4 vCPU / 16 GB**, shared by Ray + the model + any CPU-only sandbox +pool. That caps pool size and the per-hook D2H/H2D budget (the per-hook D2H/H2D transport), pushing the two-tier optimization +(§7) *earlier* on small prod instances than a pure-correctness view suggests. + +## 7. Future optimization (out of scope now) + +Where do user tensor ops execute? The prototype picks **per-hook D2H/H2D in the jail, on CPU** (D2H/H2D per hook). +The destination is **the two-tier approach**: run the common tiny tensor algebra (read/swap/save/unembed/steer/ +patch/ablate/cache, plus the `.carry()` cross-trace handoff; the realized primitive set, "capture" shipped as +`.carry()`) on the host via a *validated op interpreter* (data never leaves the GPU, no jail in the +hot path), and route only genuinely-arbitrary Python to the CPU jail. (B) GPU-in-jail via CUDA IPC is +rejected (surface + breaks the pool); (C) fully symbolic host execution is a breaking change to nnsight's +"real tensors" contract. The two-tier approach attacks the root — *arbitrary code next to the data* — rather than paying to +isolate code that didn't need to be arbitrary. + +On small prod instances this is **less optional than it sounds**: a `g4dn.xlarge` (4 vCPU / 16 GB, §6.4) +leaves little headroom for a CPU-only sandbox pool doing per-hook D2H/H2D, so the two-tier approach likely graduates from +"later optimization" to "needed for throughput" before this ships at scale. + +## 8. Decisions & open questions + +**Decided** +- Move the worker across the existing six-event protocol boundary; change transport only. +- Tier-1 isolation, CPU-only, one jail per request, host-side `Batcher` as row-bound authority, + deserialize-in-jail, `SIGKILL` cancellation. +- Prototype tensor path = per-hook D2H/H2D transport; destination = the two-tier approach — and on small prod instances (`g4dn`, §6.4) the two-tier approach is + needed sooner than "later". +- Bootstrap on bubblewrap, target nsjail, gVisor as the strong-boundary upgrade. +- **Prod runs ECS-on-EC2** (§6.4); the probe transfers (same Docker-default posture). **the host-level sibling jail (host-level + sibling on the EC2 instance) is the production target**; in-container the in-container self-jail is a throwaway first cut and on + ECS costs `privileged=true`. + +**Open** +- **in-container vs host-level sequencing** (§6.5) — the host-level sibling jail is the target; the only question is whether a privileged in-container first + cut is worth building at all, or whether to go straight to the host-level sibling. +- Warm-pool mechanism: zygote-fork vs pre-launched idle-worker pool. +- Tensor wire format: raw buffer vs `safetensors` vs shmem ring (correctness first, optimize later). +- Exact `Jailer` interface so bwrap → nsjail → runsc are drop-in. diff --git a/docs/developing/mediator-threat-models.md b/docs/developing/mediator-threat-models.md new file mode 100644 index 000000000..ae7f17bb3 --- /dev/null +++ b/docs/developing/mediator-threat-models.md @@ -0,0 +1,207 @@ +# Mediator Isolation — Threat Models & Security Hardening + +**Status:** Design notes / decision artifact · **Branch:** `worktree-mediator-sandbox` +**Builds on:** [mediator-gpu-trace-integration.md](mediator-gpu-trace-integration.md) (the GPU-worker +backend + §17 triton motivation), [mediator-isolation-sandbox.md](mediator-isolation-sandbox.md) (the +earlier CPU-only threat model + AWS deployment probe). + +This page is the security companion to the GPU-worker backend: *what* adversary each isolation +configuration faces, what it can and cannot contain, and the resulting **cost coupling** that makes +the "how much isolation" question concrete. It also records the one security fix landed in-branch: the +worker→host **closed value-algebra codec**. + +--- + +## 1. Why isolation, and why now + +`model.trace()` runs **arbitrary user intervention code**. NDIF runs that on shared GPU infrastructure, +today behind an **in-process Python import whitelist** that wraps the *entire forward pass*. That +boundary is architecturally unfixable in-process (the security suite shows 10/10 in-process escapes and +8/8 cross-tenant capability leaks succeed), **and** it cannot serve **triton-kernel models** — frontier +MoE/SSM models and anything under `torch.compile` JIT-compile triton, which needs +`subprocess`/ptxas/`tempfile`/`open` exactly when the whitelist denies them (see §17 of +[mediator-gpu-trace-integration.md](mediator-gpu-trace-integration.md)). The GPU-worker backend fixes +both at once: the **trusted forward runs host-side, unrestricted** (triton works); only the **untrusted +intervention** runs in a separate, sandboxable worker. + +--- + +## 2. Attacker model + +The attacker submits **serialized Python (cloudpickle)** that the server deserializes and executes +against a loaded model. They control: (i) the **pickle payload** (gadgets can fire at deserialize time) +and (ii) the **intervention function** — arbitrary Python at hook points. They do *not* directly control +the host process — only what they reach through the channels we expose. Intent tiers: + +- **Footgun** — accident: infinite loop, OOM, bad index/device assert, accidental host mutation. +- **Casual adversary** — `import os`, open files, sockets, object-graph walks. +- **Determined adversary** — pickle gadgets, C-extension escapes, GPU-driver/kernel exploits, + cross-tenant GPU reads, side channels. + +--- + +## 3. Assets + +| | Asset | Notes | +|---|---|---| +| **A** | Host-process integrity (no RCE on the model server) | the crown jewel | +| **B** | Host filesystem | model weights, other tenants' cached data, **secrets/IAM** | +| **C** | Network | exfiltration, internal services, **IMDS credential theft** | +| **D1** | Cross-tenant **host-memory** isolation | only under co-batching | +| **D2** | Cross-tenant **GPU-memory** isolation | the hard one | +| **E** | Availability / DoS | resource exhaustion, crash blast-radius | +| **F** | Deserialization RCE | attacker pickle firing **on the host** | + +D1/D2 only bite when **multiple tenants share a forward pass** (co-batching). + +--- + +## 4. Configuration ladder + the cost coupling + +The deep insight: **the protection level and the data-path cost are coupled** — closing a deeper threat +*forces* the activation onto a slower path. You cannot pick "max isolation, min latency." + +| Ladder | Threats newly closed | Data path | Per-hook cost | +|---|---|---|---| +| **R0 in-process** | — | tensor never moves | ~0 (≈12 ms whole gpt2 trace) | +| **R1 worker, shared GPU** | A host-process, **D1**, **F**, footgun-DoS | stays on GPU (CUDA-IPC) | **~0.6 ms, size-independent** + spawn (warm ~0.22 s) | +| **R2 + OS jail** (seccomp / net+mount+user ns + cgroups) | B fs, C net/IMDS, privilege, RAM-DoS | **same GPU buffer** | **~same ~0.6 ms** (jail is warm-time only) | +| **R3 + no GPU in worker** | **D2**, GPU-driver/kernel surface | **GPU→CPU→GPU (PCIe)** | **~ms, size-*dependent*** (1–285 ms) | +| **R4 + CPU-mem isolation** (gVisor/microVM) | kernel-proof, mem side-channels | **socket copy** (no shared mem) | R3 + serialize + pipe (slowest) | + +Two consequences: +1. **The cost cliff is R2→R3, not R1→R2.** Adding the OS jail (closes fs/net/IMDS/privilege/RAM-DoS + against even a determined adversary) rides the *same* GPU buffer and is ~free per hook. The latency + explosion only happens when you leave the GPU to close **D2**. +2. **The decision reduces to one bit: do you require D2 (cross-tenant GPU-memory isolation)?** If no → + **R2** is the target (full host-OS isolation, fast path, accept the GPU-layer residual). If yes → + you must pay **R3** for the co-batched path. + +**Shipped vs designed (R2 is the target, not the current state).** The code today is **R1 plus +footgun-containment**: the spawned worker (shared GPU, CUDA-IPC), a seccomp **denylist** that EPERMs new +filesystem (`open`/`openat`/`openat2`), network (`socket`/`connect`), and exec (`execve`/`execveat`) +syscalls, and a GPU allocator memory-fraction cap. That contains careless or buggy interventions, matching +the relaxed model in [gpu-sandbox.md](gpu-sandbox.md) ("contain footguns, not a determined adversary"). +The full **R2** jail in the ladder above (an allowlist-default seccomp filter that also blocks +`ptrace`/`clone`/`fork`, plus net/mount/user namespaces and cgroups) is **designed, not yet built**, so +R2's "closes ... against even a determined adversary" line is a roadmap target, not what ships today. +Raising the posture to R2 is the §8 hardening work. + +--- + +## 5. Security controls vs the shared-GPU design + +The GPU-worker's defining feature — host and worker **share a GPU + a CUDA-IPC memory buffer** — is +itself a cross-trust channel and requires exposing the NVIDIA driver surface to the worker. Controls +split by whether they touch that path. (Lifecycle rule for all of Group A: **map device fds + init CUDA ++ map the IPC buffer + warm, THEN jail** — fds/buffer survive; mirrors the existing +lock-down-after-warm.) + +**Group A — compatible (the free wins, take these):** +- **Network namespace** (no NIC) — zero GPU impact; kills exfil + IMDS. Take unconditionally. +- **Mount namespace** (ro allowlist; bind only `/dev/nvidia*` + CUDA libs) — removes weights/secrets/ + other-tenant disk from view. Highest value. +- **User namespace** (non-root, drop caps); **cgroups** (CPU/mem → OOM/forkbomb cap); **seccomp** (today a + footgun **denylist** of fs/net/exec syscalls; the R2 target is allowlist-default plus `ptrace`/`clone`/`fork`). + +**Group B — in tension / incompatible (forces the CPU-transport fallback or hardware partition):** +- **PID/IPC namespace** — cross-namespace CUDA-IPC is documented to fail; `CLONE_NEWPID` must be entered + before CUDA init, colliding with map-first. Likely cannot PID-isolate while sharing the buffer. +- **gVisor** — `nvproxy` implements a narrow NVIDIA ioctl subset; CUDA-IPC unconfirmed → likely drops to + CPU transport. +- **Firecracker microVM** — GPU only via exclusive passthrough; can't share one GPU. +- **MIG** — the only real GPU-memory isolation, but breaks IPC sharing by construction (and unavailable + on the T4 prod GPU). + +**Residual that no OS control closes:** even with all of Group A, the worker holds a CUDA context on a +GPU also hosting the model + (under batching) other tenants' KV-cache → GPU/driver-layer cross-tenant +surface (driver bugs, IPC-handle reachability, memory-residency side channels). Only **R3 (no GPU in +worker)** or **MIG** removes it. + +--- + +## 6. Co-batch tenant isolation + +With co-batching in scope, "can a co-batched user reach another's data / entire workflow?" resolves as +follows (verified in code). + +**Row-bounding is host-authoritative.** `batch_group` is computed host-side at invoke construction +(`tracing/invoker.py:81`); the host narrows reads to `self.batch_group` (`interleaver.py:1245`) and +bounds writes (`:1259`/`:1294`). The worker only names a `(path, kind, iteration)` string — it cannot +widen. Each user's intervention runs in its own worker **process** (separate address space). + +**The one data-plane hole.** An **empty invoke** gets `batch_group=None`, and `narrow(None)` returns the +**entire batch** while `swap(None,…)` overwrites it (`batching.py:214/217/244/248` — `None`, `[-1]`, and +`not needs_batching` are all full-batch sentinels). So a `batch_group=None` mediator is a full-batch +read+write primitive across all co-batched tenants — the in-process suite's "`swap(None)` poisons all +tenants" leak. + +**Enforce, don't guard.** `batch_group` is 100% host-computed (no user-supplied field), so this is a +**system invariant**, not input validation — but the full-batch path is reachable via the user's +*program structure* (an empty invoke is a legit single-tenant feature). The correct enforcement is +architectural: **the Batcher/Interleaver instance is the tenant boundary** — co-batch *below* the nnsight +Batcher (at the inference engine, e.g. vLLM continuous batching), never *inside* one shared Batcher. +Under that invariant all three sentinels are tenant-local and safe by construction, and +`cross_invoker`/`barrier` (per-Interleaver `_xinvoke_store`/`_barrier_acc`) close with it. The vLLM path +gets this for free (independent per-request interleavers); an HF-merged-batch path must enforce it. A +runtime reject-`None` check is defense-in-depth against our own merge bug, not the primary control. + +**What a co-batched user CAN / CANNOT do.** *Cannot* (with the empty-invoke invariant + normal CUDA): +read another tenant's rows/activations (narrowed out, never sent), their prompt/code/output (separate +process), or their **live** GPU memory by pointer (separate CUDA contexts are VA/page-table isolated). So +reconstructing another user's **entire workflow is structurally closed**. *Can*, on the shared GPU +(R1/R2): scoop **uninitialized/freed GPU memory** (untargeted *fragments*, not workflows) + coarse +timing side channels — reducible on the fast path via **zero-on-free memory scrubbing** + alloc caps + +driver capability segmentation, without going to R3. + +--- + +## 7. Worker→host RCE — finding and fix (landed in-branch) + +**Finding.** The host read worker frames with `mp.Pipe.recv()` (`transport.py`), which **unpickles**. +The worker runs attacker code, so it could send a crafted pickle and get **RCE on the trusted host** — +which bypasses *every* other layer (seccomp, namespaces, row-bounding). The design fixed deserialization +for the *inbound* user payload (unpickled inside the worker) but not the *outbound* results. This is a +precondition for the isolation guarantee holding against a determined adversary. + +**Fix.** Worker→host now never `pickle.loads` at all. Tensors already travel out-of-band (GPU bounce +buffer / safetensors), so a frame is **tensor-free**; the remaining small structure crosses as a **closed +value algebra** encoded/decoded by a hand-written codec (`transport._codec_dumps` / `_codec_loads`). The +algebra is exactly: `None | bool | int | float | str | bytes | list | tuple | dict | set | torch.dtype | +torch.device | Array (out-of-band)`. Decode is **pure data assembly**: no `find_class`, no `REDUCE`, no +global resolution, so there is no opcode that can construct or call an arbitrary object. It is bounded +(a `_MAX_CODEC_BYTES` cap and bounds-checked reads guard against decode-bomb / OOB), and a value outside +the algebra fails **loud at encode in the worker** (`BoundaryValueError`, naming the offending object), +not at decode. The event rides as its string `.value`, and exceptions ride as a `(type-name, message)` +sentinel, so neither pulls in a class. **Host→worker stays normal pickle** (host-authored, trusted). + +- **Why a closed codec rather than a restricted unpickler:** a restricted unpickler still runs the pickle + VM and gates classes at `find_class`; that gate has had RCE bypasses (e.g. `torch.load(weights_only=True)`, + CVE-2025-32434). Because the worker is the less-trusted side, running *any* pickle VM on its bytes is the + wrong shape. The codec removes the class of bug outright: there is no instruction that can call anything. + Per-type enumeration is the price, and it is small (the frame is tensor-free and the algebra is fixed). +- **Cost:** measured ~1.25–1.6x the old restricted-unpickler frame round-trip, swamped by the ~0.6 ms/hook + GPU context-switch, so negligible end to end. +- **Capability scope:** `.save()` / `.carry()` of a tensor, a numpy array, or basic scalars/containers + crosses; an arbitrary live object / framework type (e.g. `ModelOutput`) does not, and is refused at + encode with a clear error. + +Coverage: `prototypes/mediator-sandbox/gpu_sandbox/test_isolated_codec_security.py` covers fidelity for +VALUE/SWAP/END/**CACHE (dtype+device)**/EXCEPTION/push (including numpy + dtype/device), a `__reduce__` +gadget and a custom object both refused at **encode**, and malformed/oversized frames raising at decode. +(CPU is enough; needs torch.) The legacy AF_UNIX socket channels +(`SocketHostChannel`/`ShmSocketHostChannel`) are unused by `isolate_mediators` and still plain-unpickle: +route them through `_codec_loads` before wiring them to an untrusted worker. + +--- + +## 8. Open items / next steps + +- **Group-A namespace jail** behind a pluggable `Jailer` (applied post-warm/pre-`ready`) — the + highest-value, design-preserving hardening; R2 on the fast path. +- **Co-batch invariant** in the (future) HF-merged-batch integration: keep the nnsight Batcher + per-tenant; add a multi-tenant tripwire rejecting unbounded `batch_group`. (Single-trace=single-Batcher + already holds it today; an unconditional guard would break legit single-user empty invokes.) +- **GPU memory scrubbing** (zero-on-free) for the shared-GPU residual. +- **gVisor + CUDA-IPC empirical test** — decides whether any R3-strong path keeps GPU sharing. +- **Run the codec + isolated suite** on a torch/GPU box to confirm fidelity + bit-identical regression. diff --git a/prototypes/mediator-sandbox/_smoke.py b/prototypes/mediator-sandbox/_smoke.py new file mode 100644 index 000000000..b6e9bb1ee --- /dev/null +++ b/prototypes/mediator-sandbox/_smoke.py @@ -0,0 +1,6 @@ +from nnsight import LanguageModel +m = LanguageModel("gpt2", device_map="cpu", dispatch=True) +with m.trace("Hello world"): + h6 = m.transformer.h[6].output[0].save() + logits = m.lm_head.output.save() +print("h6", tuple(h6.shape), "logits", tuple(logits.shape), "SMOKE_OK") diff --git a/prototypes/mediator-sandbox/gpu_sandbox/fault_test.py b/prototypes/mediator-sandbox/gpu_sandbox/fault_test.py new file mode 100644 index 000000000..869785df0 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/fault_test.py @@ -0,0 +1,76 @@ +#!/usr/bin/env python3 +"""Does a GPU fault in a worker stay contained to that worker, or take the host +down too? Decisive for the separate-contexts vs MPS choice. + +A spawned worker triggers a device-side assert (embedding lookup out of range), +which poisons its CUDA context. We then check whether the HOST's CUDA context +still works. Run once without MPS (separate contexts) and once under MPS. +""" +import queue +import sys + +import torch +import torch.multiprocessing as mp + + +def faulting_worker(sig, mode): + import torch + import torch.nn.functional as F + try: + torch.ones(8, device="cuda").sum() + torch.cuda.synchronize() + sig.put("worker: cuda ok") + except Exception as e: # noqa: BLE001 + sig.put(f"worker: cuda init FAILED {type(e).__name__}") + return + try: + if mode == "illegal": + a = torch.randn(8, device="cuda") + idx = torch.tensor([2 ** 40], device="cuda", dtype=torch.long) # address overflow + _ = a[idx] # -> illegal memory access + else: + _ = F.embedding(torch.tensor([10_000_000], device="cuda"), + torch.randn(16, 8, device="cuda")) # -> device-side assert + torch.cuda.synchronize() + sig.put("worker: NO FAULT (unexpected)") + except Exception as e: # noqa: BLE001 + sig.put(f"worker: FAULTED {type(e).__name__}: {str(e)[:50]}") + sig.put("worker: exiting") + + +def host_cuda_works(): + try: + r = (torch.ones(2048, device="cuda") * 3.0).sum() + torch.cuda.synchronize() + return float(r) == 3.0 * 2048 + except Exception as e: # noqa: BLE001 + return f"DEAD: {type(e).__name__}: {str(e)[:70]}" + + +def main(): + mode = sys.argv[1] if len(sys.argv) > 1 else "assert" + mp.set_start_method("spawn") + before = host_cuda_works() + print(f"[mode={mode}] host CUDA BEFORE worker fault: {before}") + + sig = mp.Queue() + p = mp.Process(target=faulting_worker, args=(sig, mode)) + p.start() + msgs = [] + for _ in range(5): + try: + msgs.append(sig.get(timeout=30)) + except queue.Empty: + break + p.join(timeout=15) + print(f"worker messages: {msgs} | worker exitcode={p.exitcode}") + + after = host_cuda_works() + print(f"host CUDA AFTER worker fault: {after}") + contained = after is True + print(f"VERDICT: host {'SURVIVED — fault CONTAINED to the worker' if contained else 'DIED — fault COUPLED to host'}") + sys.exit(0 if contained else 2) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/gpu_sandbox.py b/prototypes/mediator-sandbox/gpu_sandbox/gpu_sandbox.py new file mode 100644 index 000000000..45f69d53f --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/gpu_sandbox.py @@ -0,0 +1,73 @@ +"""Host-side manager for the GPU-enabled isolated worker. + +Holds a shared GPU bounce buffer and a spawned, locked-down worker process. +``apply(activation, fn)`` runs the user's ``fn`` on ``activation`` *in the +worker* (zero-copy via the shared buffer) and returns the result — the user's +arbitrary code never runs in the model-server process. + +A real deployment keeps a POOL of these (one per concurrent request); this is a +single worker for clarity. Worker death (a segfault in user code) is detected and +surfaced; the host keeps serving. +""" +import os +import sys + +import cloudpickle +import torch +import torch.multiprocessing as mp + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) # so the worker can import sandbox/gpu_worker + + +class GPUSandbox: + def __init__(self, arena_bytes=64 << 20, gpu_mem_fraction=0.3, device="cuda"): + self.device = device + ctx = mp.get_context("spawn") # CUDA requires spawn, not fork + self.buf = torch.empty(arena_bytes, dtype=torch.uint8, device=device) # the bounce buffer + self.parent_conn, child_conn = ctx.Pipe() + self.ready = ctx.Queue() + from gpu_worker import run + self.proc = ctx.Process( + target=run, args=(self.buf, child_conn, self.ready, gpu_mem_fraction), daemon=True + ) + self.proc.start() + assert self.ready.get(timeout=180) == "ready" + + def apply(self, activation: torch.Tensor, fn, timeout=60): + """Run ``fn(activation)`` in the isolated worker; return the result tensor.""" + if not self.proc.is_alive(): + raise RuntimeError("sandbox worker is dead") + a = activation.contiguous() + ab = a.flatten().view(torch.uint8) + self.buf[: ab.numel()].copy_(ab) # D2D copy into the shared buffer + torch.cuda.synchronize() + self.parent_conn.send((cloudpickle.dumps(fn), tuple(a.shape), a.dtype, ab.numel())) + if not self.parent_conn.poll(timeout): + # worker is wedged (e.g. an infinite loop in user code) — kill it; the + # host is unaffected and a pool would respawn a fresh worker. + self.proc.terminate() + raise TimeoutError(f"sandboxed intervention exceeded {timeout}s — worker killed") + try: + reply = self.parent_conn.recv() + except (EOFError, OSError): + # the pipe broke mid-op → the worker crashed (e.g. a segfault in user + # C-code). Contained: only this request dies; the host is fine and a + # pool respawns. Surface it cleanly instead of leaking EOFError. + raise RuntimeError("sandbox worker crashed during the intervention") + if reply[0] == "err": + raise RuntimeError(f"sandboxed intervention raised {reply[1]}: {reply[2]}") + _, shape, dtype, nbytes = reply + return self.buf[:nbytes].view(dtype).view(*shape).clone() + + def alive(self): + return self.proc.is_alive() + + def close(self): + try: + if self.proc.is_alive(): + self.parent_conn.send("stop") + self.proc.join(timeout=5) + except Exception: + pass + if self.proc.is_alive(): + self.proc.terminate() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/gpu_worker.py b/prototypes/mediator-sandbox/gpu_sandbox/gpu_worker.py new file mode 100644 index 000000000..c4d6126de --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/gpu_worker.py @@ -0,0 +1,60 @@ +"""The GPU-enabled, footgun-contained worker process (spawn target). + +Lifecycle: + 1. (at spawn) map the shared GPU bounce buffer — a CUDA tensor the host shared + via IPC, so host and worker point at the SAME GPU memory. + 2. warm CUDA + pre-import torch/numpy (so user ops don't need a new import). + 3. lock_down(): seccomp blocks new open/openat (fs) and socket/connect (net); + RLIMIT_AS caps memory. CUDA keeps working (already-open /dev/nvidia* fds). + 4. loop: receive (cloudpickled user fn + tensor metadata) → view the bounce + buffer as that tensor → run the user fn on the real GPU tensor → write the + result back into the buffer → reply with the result's metadata. + +The user fn is arbitrary Python. It runs HERE, after lockdown: a crash is +contained to this process (the host respawns), an open()/socket() fails with +EPERM, a runaway alloc hits the rlimit. The host and its other tenants are safe; +only this request is affected. The GPU is shared (the accepted risk). +""" +import torch + + +def run(shared_buf, conn, ready, gpu_mem_fraction): + import cloudpickle # noqa: F401 + + try: + import numpy # noqa: F401 (pre-import so user ops referencing it don't openat) + except Exception: + pass + from sandbox import lock_down + + # warm CUDA so all kernels/contexts are loaded before we cut off file opens + _ = (torch.randn(128, 128, device="cuda") @ torch.randn(128, 128, device="cuda")).sum() + torch.cuda.synchronize() + # warm cloudpickle's (de)serialization machinery so per-request loads() won't + # trigger a lazy import after the filesystem is locked down + cloudpickle.loads(cloudpickle.dumps(lambda _t: _t * 2.0)) + + # Cap GPU memory so a runaway allocation in user code can't exhaust the device. + # (RLIMIT_AS is unusable here — CUDA reserves tens of GB of *virtual* space.) + torch.cuda.set_per_process_memory_fraction(gpu_mem_fraction) + lock_down() + ready.put("ready") + + while True: + msg = conn.recv() + if msg == "stop": + return + fn_blob, shape, dtype, nbytes = msg + try: + fn = cloudpickle.loads(fn_blob) + t = shared_buf[:nbytes].view(dtype).view(*shape) # zero-copy view of the buffer + out = fn(t) # <-- arbitrary user code on the GPU tensor + if not torch.is_tensor(out): + out = torch.as_tensor(out, device="cuda") + out = out.contiguous() + ob = out.flatten().view(torch.uint8) + shared_buf[: ob.numel()].copy_(ob) # result back into the shared buffer + torch.cuda.synchronize() + conn.send(("ok", tuple(out.shape), out.dtype, ob.numel())) + except Exception as e: # noqa: BLE001 (contain the footgun, report it) + conn.send(("err", type(e).__name__, str(e)[:200])) diff --git a/prototypes/mediator-sandbox/gpu_sandbox/perf_ctxswitch.py b/prototypes/mediator-sandbox/gpu_sandbox/perf_ctxswitch.py new file mode 100644 index 000000000..601b7f681 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/perf_ctxswitch.py @@ -0,0 +1,74 @@ +#!/usr/bin/env python3 +"""Is the ~0.5 ms GPU context-switching between the two processes' CUDA contexts? + +Compare three round-trips between host and a spawned worker: + A. worker touches NO GPU (bare echo) -> pure IPC + wakeups + B. ONLY the worker does a GPU op per round -> one context active at a time-ish + C. BOTH host and worker do a GPU op per round -> GPU must switch contexts host<->worker + +If C >> A, the cost is GPU context switching, not process communication. +""" +import os +import sys +import time + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import torch +import torch.multiprocessing as mp + + +def worker(conn, touch_gpu): + if touch_gpu: + x = torch.randn(4096, device="cuda") + torch.cuda.synchronize() + while True: + m = conn.recv() + if m == "stop": + return + if touch_gpu: + x.add_(1.0) + torch.cuda.synchronize() + conn.send(b"k") + + +def t_ms(fn, n=200, warm=30): + for _ in range(warm): + fn() + t0 = time.perf_counter() + for _ in range(n): + fn() + return (time.perf_counter() - t0) / n * 1e3 + + +def run(touch_worker, touch_host): + ctx = mp.get_context("spawn") + pc, cc = ctx.Pipe() + p = ctx.Process(target=worker, args=(cc, touch_worker), daemon=True) + p.start() + y = torch.randn(4096, device="cuda") + torch.cuda.synchronize() + + def one(): + if touch_host: + y.add_(1.0) + torch.cuda.synchronize() + pc.send(b"x") + pc.recv() + + dt = t_ms(one) + pc.send("stop"); p.join() + return dt + + +def main(): + a = run(touch_worker=False, touch_host=False) + b = run(touch_worker=True, touch_host=False) + c = run(touch_worker=True, touch_host=True) + print(f"[ctx] A. worker no-GPU (pure IPC): {a:.3f} ms") + print(f"[ctx] B. worker GPU op only: {b:.3f} ms") + print(f"[ctx] C. BOTH host+worker GPU op (alternate):{c:.3f} ms <- this is apply()'s pattern") + print(f"[ctx] GPU context-switch overhead ~= C - A = {c - a:.3f} ms") + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/perf_decompose.py b/prototypes/mediator-sandbox/gpu_sandbox/perf_decompose.py new file mode 100644 index 000000000..e213bcdd3 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/perf_decompose.py @@ -0,0 +1,67 @@ +#!/usr/bin/env python3 +"""Where does the ~0.6 ms/hook go? Decompose it into: + - bare process round-trip (mp.Pipe send/recv = IPC + 2 scheduler wakeups), no torch/cloudpickle + - cloudpickle.loads (the worker rebuilds the fn every call) + - torch.cuda.synchronize (apply does 2: host after copy-in, worker after op) +so we can say whether it's process-communication cost or something else. +""" +import os +import sys +import time + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import cloudpickle +import torch +import torch.multiprocessing as mp + + +def steer(t): + return t + 1.0 + + +def echo_worker(conn): + while True: + m = conn.recv() + if m == "stop": + return + conn.send(b"k") # smallest possible reply — pure IPC round-trip + + +def t_ms(fn, n=200, warm=20): + for _ in range(warm): + fn() + t0 = time.perf_counter() + for _ in range(n): + fn() + return (time.perf_counter() - t0) / n * 1e3 + + +def main(): + ctx = mp.get_context("spawn") + pc, cc = ctx.Pipe() + p = ctx.Process(target=echo_worker, args=(cc,), daemon=True) + p.start() + + # 1. bare process round-trip: send a tiny msg, worker echoes. This is IPC + + # two scheduler wakeups (host→worker, worker→host) with nothing else. + pingpong = t_ms(lambda: (pc.send(b"x"), pc.recv())) + + # 2. cloudpickle.loads cost (the worker does this every call) + blob = cloudpickle.dumps(steer) + loads = t_ms(lambda: cloudpickle.loads(blob)) + + # 3. a single torch.cuda.synchronize after a trivial op + a = torch.randn(1024, device="cuda") + sync = t_ms(lambda: (a.add_(1.0), torch.cuda.synchronize())) + + pc.send("stop"); p.join() + + print(f"[decompose] bare mp.Pipe round-trip (IPC + 2 wakeups): {pingpong:.3f} ms") + print(f"[decompose] cloudpickle.loads(fn) per call: {loads:.3f} ms") + print(f"[decompose] one cuda.synchronize (+trivial op): {sync:.3f} ms") + print(f"[decompose] => apply() ~= pingpong + loads + ~2*sync " + f"= {pingpong + loads + 2 * sync:.3f} ms (measured apply ~0.6-0.7)") + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/perf_spawn_cost.py b/prototypes/mediator-sandbox/gpu_sandbox/perf_spawn_cost.py new file mode 100644 index 000000000..ed6bd6b69 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/perf_spawn_cost.py @@ -0,0 +1,235 @@ +#!/usr/bin/env python3 +"""How much does spawning a fresh GPU worker process per incoming request cost? + +This is the number the warm worker pool (an unbuilt item) would amortize: under +isolation today, every `with model.trace(...)` spawns a new worker via +`spawn_isolated_worker` -> `mp.get_context("spawn").Process(target=_worker_main)`. + +Two measurements: + + (A) DECOMPOSED synthetic bring-up. Spawn a fresh process that replays exactly + the startup _worker_main does, signalling the host at each milestone so we + attribute the wall-clock to: interpreter+mp bootstrap (re-import of the + worker module) / import torch / import nnsight / CUDA context init / warm + matmul+sync / numpy+cloudpickle warm. Module-top imports here are kept + light so the bootstrap stage does NOT include torch/nnsight (the real + worker pays those during its module re-import; we just measure them as + explicit stages so the breakdown is visible). N spawns -> mean/std. + + (B) REAL end-to-end. Build a gpt2 LanguageModel and time a real isolated + `model.trace(...)` (which spawns a worker) vs the in-process trace, N + times. `spawn_isolated_worker` is wrapped to report the spawn-only slice + of that overhead. Difference = the per-request cost of isolation; the + spawn slice is what a warm pool removes. + +Run (spawn+CUDA needs an unsandboxed shell): + CUDA_VISIBLE_DEVICES=6 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/perf_spawn_cost.py +""" +# NOTE: keep module-top imports LIGHT (no torch / no nnsight) so a spawned +# child's re-import of this module during bootstrap stays cheap and stage (A) +# can isolate the bootstrap cost from the heavy-import cost. +import multiprocessing as std_mp +import os +import statistics +import sys +import time + + +# --------------------------------------------------------------------------- # +# (A) Decomposed synthetic worker bring-up # +# --------------------------------------------------------------------------- # +# Milestones, in the order _worker_main reaches them. The child sends one tag +# per milestone; the HOST timestamps each recv (cross-process clocks are not +# comparable, so we measure host-side deltas between arrivals). +_STAGES = ["boot", "torch", "nnsight", "cuda_ctx", "warm_mm", "warm_imports"] + + +def _decompose_worker(conn, device): + """Replays _worker_main's startup, signalling the host at each milestone.""" + # 'boot' fires before any heavy import: captures interpreter start + the + # multiprocessing spawn handshake + re-import of THIS (light) module. + conn.send("boot") + + import torch # the worker's first torch import (spawn re-imports from scratch) + conn.send("torch") + + import nnsight.intervention.isolation # noqa: F401 (the module spawn re-imports) + conn.send("nnsight") + + # First CUDA call -> initialize this process's own CUDA context on the GPU. + torch.zeros(1, device=device) + torch.cuda.synchronize() + conn.send("cuda_ctx") + + # The warm matmul _worker_main runs so kernels/contexts are loaded. + (torch.randn(8, 8, device=device) @ torch.randn(8, 8, device=device)).sum() + torch.cuda.synchronize() + conn.send("warm_mm") + + # Warm imports _worker_main triggers before any seccomp lockdown. + try: + import numpy # noqa: F401 + except Exception: # noqa: BLE001 + pass + import cloudpickle + + cloudpickle.loads(cloudpickle.dumps(lambda _t: _t)) + conn.send("warm_imports") + + try: + while conn.recv() != "stop": + pass + except (EOFError, OSError): + pass + + +def measure_decomposed(n, device): + ctx = std_mp.get_context("spawn") # CUDA requires spawn, matching the real code + # rows[stage] = list of per-spawn durations (ms) for that stage + rows = {s: [] for s in _STAGES} + totals = [] + for _ in range(n): + pc, cc = ctx.Pipe() + t_prev = time.perf_counter() + t_start = t_prev + p = ctx.Process(target=_decompose_worker, args=(cc, device), daemon=True) + p.start() + for stage in _STAGES: + tag = pc.recv() + now = time.perf_counter() + assert tag == stage, f"expected {stage}, got {tag}" + rows[stage].append((now - t_prev) * 1e3) + t_prev = now + totals.append((t_prev - t_start) * 1e3) + pc.send("stop") + p.join(timeout=10) + if p.is_alive(): + p.terminate() + p.join() + return rows, totals + + +def _fmt(vals): + m = statistics.mean(vals) + s = statistics.pstdev(vals) if len(vals) > 1 else 0.0 + return f"{m:8.1f} ± {s:6.1f} ms" + + +def report_decomposed(rows, totals, n): + print("=" * 74) + print(f"(A) Decomposed synthetic worker bring-up (spawn context, n={n})") + print("-" * 74) + labels = { + "boot": "interpreter + mp spawn handshake + module re-import", + "torch": "import torch (child, cold)", + "nnsight": "import nnsight.intervention.isolation", + "cuda_ctx": "CUDA context init (first cuda op + sync)", + "warm_mm": "warm 8x8 matmul + sync", + "warm_imports": "warm numpy + cloudpickle (+roundtrip)", + } + for s in _STAGES: + print(f" {labels[s]:52s} {_fmt(rows[s])}") + print("-" * 74) + print(f" {'TOTAL per-spawn bring-up (start -> ready)':52s} {_fmt(totals)}") + print("=" * 74) + return statistics.mean(totals) + + +# --------------------------------------------------------------------------- # +# (B) Real end-to-end isolated trace vs in-process # +# --------------------------------------------------------------------------- # +def measure_real(n, device): + import torch + + from nnsight import LanguageModel + from nnsight.intervention import isolation + from nnsight.intervention.isolation import isolate_mediators + + PROMPT = "The Eiffel Tower is in the city of" + model = LanguageModel("gpt2", device_map=device, dispatch=True) + + def one_inprocess(): + with model.trace(PROMPT): + _ = model.transformer.h[6].output[0].save() + + # Wrap spawn_isolated_worker to record just the spawn slice (serialize + + # Process.start + channel wiring) of each isolated trace. + spawn_times = [] + _orig_spawn = isolation.spawn_isolated_worker + + def _timed_spawn(med): + t0 = time.perf_counter() + _orig_spawn(med) + spawn_times.append((time.perf_counter() - t0) * 1e3) + + def one_isolated(): + with isolate_mediators(fast_lane=False): + with model.trace(PROMPT): + _ = model.transformer.h[6].output[0].save() + + # warm-up (model kernels, cudnn autotune, first-spawn page-cache effects) + one_inprocess() + one_isolated() + + inproc = [] + for _ in range(n): + torch.cuda.synchronize() + t0 = time.perf_counter() + one_inprocess() + torch.cuda.synchronize() + inproc.append((time.perf_counter() - t0) * 1e3) + + iso = [] + isolation.spawn_isolated_worker = _timed_spawn + try: + for _ in range(n): + torch.cuda.synchronize() + t0 = time.perf_counter() + one_isolated() + torch.cuda.synchronize() + iso.append((time.perf_counter() - t0) * 1e3) + finally: + isolation.spawn_isolated_worker = _orig_spawn + + return inproc, iso, spawn_times + + +def report_real(inproc, iso, spawn_times, n): + print() + print("=" * 74) + print(f"(B) Real gpt2 isolated trace vs in-process (n={n})") + print("-" * 74) + mi, mo = statistics.mean(inproc), statistics.mean(iso) + print(f" in-process trace (baseline) {_fmt(inproc)}") + print(f" isolated trace (spawns a worker) {_fmt(iso)}") + print(f" per-request isolation overhead (iso - base) {mo - mi:8.1f} ms") + if spawn_times: + print(f" ... of which spawn_isolated_worker only: {_fmt(spawn_times)}") + print("=" * 74) + + +# --------------------------------------------------------------------------- # +def main(): + import torch + + assert torch.cuda.is_available(), "needs CUDA" + device = "cuda" + n_decomp = int(os.environ.get("N_DECOMP", "8")) + n_real = int(os.environ.get("N_REAL", "8")) + + rows, totals = measure_decomposed(n_decomp, device) + decomp_total = report_decomposed(rows, totals, n_decomp) + + inproc, iso, spawn_times = measure_real(n_real, device) + report_real(inproc, iso, spawn_times, n_real) + + print() + print("Bottom line: each isolated request spawns a worker costing ~" + f"{decomp_total/1000:.1f} s of process bring-up; a warm worker pool " + "would remove essentially all of it.") + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/perf_spawn_split.py b/prototypes/mediator-sandbox/gpu_sandbox/perf_spawn_split.py new file mode 100644 index 000000000..12adda45d --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/perf_spawn_split.py @@ -0,0 +1,47 @@ +"""Split the spawn_isolated_worker slice into host-side serialize vs proc.start(). + +Must guard top-level with __main__: spawn re-imports this module as __mp_main__ +in every worker, so any top-level side effects would re-run per spawn. +""" +import time, statistics, multiprocessing.context as mpc +import torch +from nnsight import LanguageModel +from nnsight.intervention import isolation +import nnsight.intervention.serialization as ser +from nnsight.intervention.isolation import isolate_mediators + +PROMPT = "The Eiffel Tower is in the city of" + + +def main(): + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + + dumps_t = [] + _od = ser.dumps + def timed_dumps(o): + t = time.perf_counter(); r = _od(o); dumps_t.append((time.perf_counter() - t) * 1e3); return r + ser.dumps = timed_dumps + + start_t = [] + _os = mpc.SpawnProcess.start + def timed_start(self): + t = time.perf_counter(); _os(self); start_t.append((time.perf_counter() - t) * 1e3) + mpc.SpawnProcess.start = timed_start + + def one_iso(): + with isolate_mediators(fast_lane=False): + with model.trace(PROMPT): + _ = model.transformer.h[6].output[0].save() + + one_iso() # warm + dumps_t.clear(); start_t.clear() + for _ in range(6): + one_iso() + + f = lambda v: f"{statistics.mean(v):7.1f} +/- {statistics.pstdev(v):5.1f} ms (n={len(v)})" + print("host-side serialization.dumps(mediator):", f(dumps_t)) + print("SpawnProcess.start() (waits on child): ", f(start_t)) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/perf_test.py b/prototypes/mediator-sandbox/gpu_sandbox/perf_test.py new file mode 100644 index 000000000..d7e9b8126 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/perf_test.py @@ -0,0 +1,93 @@ +#!/usr/bin/env python3 +"""Performance of the GPU sandbox: the REAL apply() path, not the microbench. + +Measures (1) per-hook apply() latency vs activation size, with a component +breakdown and the in-process baseline; (2) the end-to-end impact on a real gpt2 +trace, in-process vs 1 and N sandboxed interventions. + +Run: CUDA_VISIBLE_DEVICES=6 PYTHONPATH=/src .../bin/python perf_test.py +""" +import os +import sys +import time + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import cloudpickle +import torch + +from nnsight import LanguageModel +from gpu_sandbox import GPUSandbox + + +def steer(t): + return t + 1.0 + + +def t_ms(fn, n=50, warm=10): + for _ in range(warm): + fn() + torch.cuda.synchronize() + t0 = time.perf_counter() + for _ in range(n): + fn() + torch.cuda.synchronize() + return (time.perf_counter() - t0) / n * 1e3 + + +def main(): + sb = GPUSandbox() + + # ---- 1. per-hook apply() latency vs size + breakdown ---- + print(f"{'activation':>20} {'MB':>7} {'in-proc':>9} {'apply()':>9} {'(dumps)':>9} {'(copyD2D)':>10} " + f"{'(rtt)':>8} (ms)") + for (b, s, h) in [(1, 16, 768), (1, 512, 768), (1, 512, 4096), (1, 2048, 4096), (1, 2048, 8192)]: + act = torch.randn(b, s, h, device="cuda", dtype=torch.bfloat16) + mb = act.element_size() * act.nelement() / 1e6 + inproc = t_ms(lambda: steer(act)) + full = t_ms(lambda: sb.apply(act, steer)) + # component breakdown + dumps = t_ms(lambda: cloudpickle.dumps(steer)) + ab = act.contiguous().flatten().view(torch.uint8) + copyd2d = t_ms(lambda: (sb.buf[: ab.numel()].copy_(ab))) + rtt = max(full - dumps - copyd2d, 0.0) # remainder ≈ pipe + worker wakeup + op + readback + print(f"{str((b, s, h)):>20} {mb:7.1f} {inproc:9.3f} {full:9.3f} {dumps:9.3f} {copyd2d:10.3f} {rtt:8.3f}") + + # ---- 2. end-to-end real gpt2 trace ---- + print() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + inputs = model.tokenizer("The Eiffel Tower is in the city of", return_tensors="pt").to("cuda") + + def trace_inproc(): + with model.trace(inputs): + model.transformer.h[6].output = steer(model.transformer.h[6].output) + model.lm_head.output.save() + + def trace_sandbox1(): + with model.trace(inputs): + model.transformer.h[6].output = sb.apply(model.transformer.h[6].output, steer) + model.lm_head.output.save() + + def trace_sandboxN(): + with model.trace(inputs): + for L in range(12): # intervene at every layer (e.g. a full-model cache/steer) + model.transformer.h[L].output = sb.apply(model.transformer.h[L].output, steer) + model.lm_head.output.save() + + def trace_plain(): + with model.trace(inputs): + model.lm_head.output.save() + + plain = t_ms(trace_plain, n=20, warm=5) + inp = t_ms(trace_inproc, n=20, warm=5) + s1 = t_ms(trace_sandbox1, n=20, warm=5) + sN = t_ms(trace_sandboxN, n=20, warm=5) + print(f"[gpt2 trace] plain (no intervention): {plain:7.2f} ms") + print(f"[gpt2 trace] in-process steer @ 1 layer: {inp:7.2f} ms") + print(f"[gpt2 trace] sandboxed steer @ 1 layer: {s1:7.2f} ms (+{s1 - inp:.2f} ms / hook)") + print(f"[gpt2 trace] sandboxed steer @ 12 layers: {sN:7.2f} ms (+{(sN - inp) / 12:.2f} ms / hook avg)") + + sb.close() + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/probe_bootstrap.py b/prototypes/mediator-sandbox/gpu_sandbox/probe_bootstrap.py new file mode 100644 index 000000000..83ecf65d8 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/probe_bootstrap.py @@ -0,0 +1,133 @@ +#!/usr/bin/env python3 +"""SP1 bootstrap probe — can a per-mediator intervention be serialized, then +deserialized against DUMMY modules + a stub interleaver, and run far enough to +call request() with the correct requester string? + +This de-risks the worker bootstrap (the hardest unknown) WITHOUT the channel, +forward, or saves. If this prints the expected requester, the rest is plumbing +we already have (CudaIpcChannel + Shim A/B). + +Run: + CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/probe_bootstrap.py +""" +import sys +import types + +import torch +import torch.nn as nn + +from nnsight import LanguageModel +from nnsight.intervention import serialization +from nnsight.intervention.interleaver import Mediator +from nnsight.intervention.tracing.globals import _ensure_mounted + + +def main(): + model = LanguageModel("gpt2", device_map="cpu", dispatch=True) + + # 1. Capture a real mediator's serialized bytes during a normal trace. + captured = {} + orig_start = Mediator.start + + def patched_start(self, interleaver): + if "bytes" not in captured: + try: + # The tracer attaches source during ITS __getstate__ (tracer.py:677); + # per-mediator serialization must do the same first. + self.intervention.__source__ = "".join(self.info.source) + captured["bytes"] = serialization.dumps(self) + captured["err"] = None + except Exception as e: # noqa: BLE001 + captured["err"] = f"{type(e).__name__}: {e}" + return orig_start(self, interleaver) + + Mediator.start = patched_start + try: + with model.trace("The Eiffel Tower is in"): + model.transformer.h[6].output.save() + finally: + Mediator.start = orig_start + + if captured.get("err"): + print(f"[serialize] FAILED: {captured['err']}") + sys.exit(1) + print(f"[serialize] mediator -> {len(captured['bytes'])} bytes OK") + + # 2. Build a WORKER-style persistent map: dummy modules for every Module:, + # a stub interleaver, real tokenizer/processor pass-through. + real_map = model._remoteable_persistent_objects() + + calls = [] + + class StubBatcher: + current_provider = None + current_value = None + + class StubInterleaver: + interleaving = True + batcher = StubBatcher() + current = None + + def iterate_requester(self, requester): + med = self.current + iteration = med.iteration if med.iteration is not None else med.iteration_tracker[requester] + return f"{requester}.i{iteration}" + + stub_interleaver = StubInterleaver() + + pmap = {} + for k, v in real_map.items(): + if k.startswith("Module:"): + dummy = nn.Module() + dummy.__path__ = k[len("Module:") :] + pmap[k] = dummy + elif k == "Interleaver": + pmap[k] = stub_interleaver + else: + pmap[k] = v # Tokenizer / Processor — keep real (lightweight) + + # 3. Deserialize the mediator against dummies. + try: + med = serialization.loads(captured["bytes"], pmap) + except Exception as e: # noqa: BLE001 + print(f"[deserialize] FAILED: {type(e).__name__}: {e}") + sys.exit(1) + print(f"[deserialize] mediator rebuilt against dummy modules OK (name={med.name})") + + # 4. Wire it into the stub interleaver and run the intervention with a + # recording request() — assert it asks for the right provider. + _ensure_mounted() + med.idx = 0 + med.interleaver = stub_interleaver + stub_interleaver.current = med + + def recording_request(requester): + calls.append(requester) + return torch.zeros(1, 16, 768) # gpt2-ish block output stand-in + + med.request = recording_request + med.swap = lambda requester, value: calls.append(("swap", requester)) + med.end = lambda: None + med.push = lambda: None + med.pull = lambda: None + med.cross_invoker = False + + try: + med.intervention(med, med.info) + except Exception as e: # noqa: BLE001 + print(f"[run] intervention raised: {type(e).__name__}: {e}") + print(f"[run] calls so far: {calls}") + sys.exit(1) + + expected = "model.transformer.h.6.output.i0" # real envoy root prefixes "model." + ok = expected in calls + print(f"[run] intervention requested: {calls}") + print("=" * 72) + print(f"BOOTSTRAP PROBE: {'PASS' if ok else 'FAIL'} (expected {expected!r})") + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/probe_pool_gpu_footprint.py b/prototypes/mediator-sandbox/gpu_sandbox/probe_pool_gpu_footprint.py new file mode 100644 index 000000000..f0c318d2e --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/probe_pool_gpu_footprint.py @@ -0,0 +1,118 @@ +"""Per-worker GPU memory footprint of a warm pool. + +A pooled worker holds NO model weights (dummy modules; weights stay in the host), +so its GPU residency = CUDA context + JIT-loaded kernels + transient activation +clones. This measures that per-process cost (via nvidia-smi per-PID), and its +additivity across K workers, to answer: how much GPU does a batch-size-wide pool +take, and how does it scale with workers / GPUs touched. + +Keep module-top light (no torch) so spawned children re-import cheaply. + +Run: + CUDA_VISIBLE_DEVICES=6 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/probe_pool_gpu_footprint.py +""" +import multiprocessing as std_mp +import os +import subprocess +import sys + + +def smi_used_mib(pid): + """MiB of GPU memory nvidia-smi attributes to this PID (0 if not listed).""" + out = subprocess.check_output( + ["nvidia-smi", "--query-compute-apps=pid,used_memory", + "--format=csv,noheader,nounits"] + ).decode() + for line in out.strip().splitlines(): + if not line.strip(): + continue + p, m = [x.strip() for x in line.split(",")] + if int(p) == pid: + return int(m) + return 0 + + +def _worker(conn, mem_fraction): + import torch + dev = "cuda" + if mem_fraction: + # The real worker caps allocatable GPU via set_per_process_memory_fraction. + # It bounds the caching allocator, NOT the fixed context overhead — verify. + torch.cuda.set_per_process_memory_fraction(mem_fraction) + + torch.zeros(1, device=dev); torch.cuda.synchronize() # CUDA context init + conn.send(("ctx", os.getpid())); conn.recv() + + # Realistic kernels an intervention triggers (cuBLAS heuristics at hidden size). + a = torch.randn(2048, 2048, device=dev) + (a @ a).sum(); torch.cuda.synchronize() + (torch.randn(8, 8, device=dev) @ torch.randn(8, 8, device=dev)).sum() + torch.cuda.synchronize() + conn.send(("warm_mm", os.getpid())); conn.recv() + + import nnsight # noqa: F401 — nnsight import (host-resolved; no GPU tensors) + conn.send(("nnsight", os.getpid())); conn.recv() + + # report the worker's own view too (free/total is cross-process on this device) + free, total = torch.cuda.mem_get_info() + conn.send(("memget", (free, total))); conn.recv() + conn.recv() # hold context until released + + +def main(): + K = int(os.environ.get("K", "4")) + mem_fraction = float(os.environ.get("MEM_FRACTION", "0.3")) + ctx = std_mp.get_context("spawn") + + workers = [] + print("=" * 70) + print(f"Per-worker GPU footprint (spawn, mem_fraction={mem_fraction}, K={K})") + print("-" * 70) + cum = 0 + for i in range(K): + pc, cc = ctx.Pipe() + p = ctx.Process(target=_worker, args=(cc, mem_fraction), daemon=True) + p.start() + # walk the worker through its warm stages, sampling per-PID GPU memory + stages = {} + memget = None + while True: + tag, payload = pc.recv() + if tag == "memget": + memget = payload + pc.send("k") + break + stages[tag] = smi_used_mib(payload) + pc.send("k") + workers.append((pc, p)) + full = stages["nnsight"] + delta = full - cum if i > 0 else full + cum = full if i == 0 else cum # cum tracks single-worker; recompute below + print(f" worker {i}: ctx_init={stages['ctx']:5d} +warm_mm={stages['warm_mm']:5d} " + f"+nnsight={stages['nnsight']:5d} MiB (this PID)") + if memget: + free, total = memget + print(f" device free={free//(1<<20)} MiB / total={total//(1<<20)} MiB " + f"(all procs on this GPU)") + + # total attributed across all worker PIDs right now + total_pool = sum(smi_used_mib(p.pid) for _, p in workers) + per = total_pool / K + print("-" * 70) + print(f" {K} warm workers resident: {total_pool} MiB total (~{per:.0f} MiB/worker)") + print(f" Extrapolation: a batch-size-B pool on G GPUs-touched ~= " + f"B * G * {per:.0f} MiB of context/kernels (NO model weights).") + print("=" * 70) + + for pc, p in workers: + try: + pc.send("stop") + except Exception: + pass + p.terminate(); p.join(timeout=5) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/probe_saves.py b/prototypes/mediator-sandbox/gpu_sandbox/probe_saves.py new file mode 100644 index 000000000..cb6799237 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/probe_saves.py @@ -0,0 +1,101 @@ +#!/usr/bin/env python3 +"""SP1 saves probe — after a deserialized intervention runs to end()/push() in a +worker-like context, WHERE do the .save()'d values land, and can we filter them by +Globals.saves to ship back? (Shim B feasibility.) + +Run: + CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/probe_saves.py +""" +import sys + +import torch +import torch.nn as nn + +from nnsight import LanguageModel +from nnsight.intervention import serialization +from nnsight.intervention.interleaver import Mediator +from nnsight.intervention.tracing.globals import Globals, _ensure_mounted + + +def main(): + model = LanguageModel("gpt2", device_map="cpu", dispatch=True) + + captured = {} + orig_start = Mediator.start + + def patched_start(self, interleaver): + if "bytes" not in captured: + self.intervention.__source__ = "".join(self.info.source) + captured["bytes"] = serialization.dumps(self) + return orig_start(self, interleaver) + + Mediator.start = patched_start + try: + with model.trace("The Eiffel Tower is in"): + saved_marker = model.transformer.h[6].output.save() + finally: + Mediator.start = orig_start + + real_map = model._remoteable_persistent_objects() + + class StubBatcher: + current_provider = None + current_value = None + + class StubInterleaver: + interleaving = True + batcher = StubBatcher() + current = None + + def iterate_requester(self, requester): + med = self.current + it = med.iteration if med.iteration is not None else med.iteration_tracker[requester] + return f"{requester}.i{it}" + + stub_interleaver = StubInterleaver() + pmap = {} + for k, v in real_map.items(): + if k.startswith("Module:"): + d = nn.Module(); d.__path__ = k[len("Module:"):]; pmap[k] = d + elif k == "Interleaver": + pmap[k] = stub_interleaver + else: + pmap[k] = v + + med = serialization.loads(captured["bytes"], pmap) + _ensure_mounted() + Globals.saves.clear() + med.idx = 0 + med.interleaver = stub_interleaver + stub_interleaver.current = med + med.cross_invoker = False + + sentinel = torch.arange(12, dtype=torch.float32).reshape(1, 3, 4) # the "activation" + med.request = lambda requester: sentinel + + # Let the REAL end()/push() run; stub only the channel. + class FakeChannel: + def put_event(self, item): self.last = item + med.channel = FakeChannel() + + med.intervention(med, med.info) + + # Shim B candidate: filter the worker's frame locals by Globals.saves. + frame = med.info.frame + flocals = getattr(frame, "f_locals", {}) + saved = {k: v for k, v in flocals.items() if id(v) in Globals.saves} + + print(f"[frame] info.frame type={type(frame).__name__}, #locals={len(flocals)}") + print(f"[saves] Globals.saves size={len(Globals.saves)}; filtered saved keys={list(saved.keys())}") + hit = [k for k, v in saved.items() if torch.is_tensor(v) and torch.equal(v, sentinel)] + ok = len(hit) >= 1 + print(f"[match] a saved local equals the sentinel activation: {ok} (keys={hit})") + print("=" * 72) + print(f"SAVES PROBE: {'PASS' if ok else 'FAIL'}") + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/sandbox.py b/prototypes/mediator-sandbox/gpu_sandbox/sandbox.py new file mode 100644 index 000000000..0ff34cbf9 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/sandbox.py @@ -0,0 +1,76 @@ +"""Footgun-containment sandbox for a GPU-enabled worker process. + +Goal (per the agreed threat model): contain *mistakes*, not a determined +adversary. After CUDA + torch are fully initialised, the worker calls +``lock_down()`` which installs a minimal seccomp-BPF filter that makes new +``open``/``openat`` (filesystem) and ``socket``/``connect`` (network) syscalls +fail with EPERM. CUDA keeps working because it talks to the already-open +``/dev/nvidia*`` fds via ioctl/mmap, not by opening new files. + +No external deps: the BPF program is assembled by hand and installed via +``prctl(PR_SET_SECCOMP, SECCOMP_MODE_FILTER, ...)``. x86-64 only. +""" +import ctypes +import struct + +# x86-64 syscall numbers +_NR = {"open": 2, "openat": 257, "openat2": 437, "socket": 41, "connect": 42, + "execve": 59, "execveat": 322} +_AUDIT_ARCH_X86_64 = 0xC000003E +_RET_KILL_PROCESS = 0x80000000 +_RET_ERRNO = 0x00050000 +_RET_ALLOW = 0x7FFF0000 +_EPERM = 1 +# BPF opcodes +_LD_W_ABS = 0x20 +_JMP_JEQ_K = 0x15 +_RET_K = 0x06 +_PR_SET_NO_NEW_PRIVS = 38 +_PR_SET_SECCOMP = 22 +_SECCOMP_MODE_FILTER = 2 + + +def _build_filter(blocked): + instrs = [ + (_LD_W_ABS, 0, 0, 4), # A = arch (seccomp_data offset 4) + (_JMP_JEQ_K, 1, 0, _AUDIT_ARCH_X86_64), # if x86-64: skip the kill + (_RET_K, 0, 0, _RET_KILL_PROCESS), # else kill (block arch-bypass) + (_LD_W_ABS, 0, 0, 0), # A = syscall nr (offset 0) + ] + for nr in blocked: + instrs.append((_JMP_JEQ_K, 0, 1, nr)) # if A == nr: next else skip next + instrs.append((_RET_K, 0, 0, _RET_ERRNO | _EPERM)) + instrs.append((_RET_K, 0, 0, _RET_ALLOW)) # default: allow + return instrs + + +class _sock_fprog(ctypes.Structure): + _fields_ = [("len", ctypes.c_ushort), ("filter", ctypes.c_void_p)] + + +def lock_down(block_fs=True, block_net=True): + """Install the seccomp filter. Call AFTER torch/CUDA are warmed up.""" + blocked = [] + if block_fs: + blocked += [_NR["open"], _NR["openat"], _NR["openat2"]] + if block_net: + blocked += [_NR["socket"], _NR["connect"]] + blocked += [_NR["execve"], _NR["execveat"]] # no spawning new programs either + instrs = _build_filter(blocked) + + libc = ctypes.CDLL("libc.so.6", use_errno=True) + if libc.prctl(_PR_SET_NO_NEW_PRIVS, 1, 0, 0, 0) != 0: + raise OSError(ctypes.get_errno(), "PR_SET_NO_NEW_PRIVS failed") + + prog = b"".join(struct.pack("HBBI", *i) for i in instrs) + buf = ctypes.create_string_buffer(prog, len(prog)) + fprog = _sock_fprog(len(instrs), ctypes.cast(buf, ctypes.c_void_p)) + # keep refs alive for the duration of the call + if libc.prctl(_PR_SET_SECCOMP, _SECCOMP_MODE_FILTER, ctypes.byref(fprog), 0, 0) != 0: + raise OSError(ctypes.get_errno(), "PR_SET_SECCOMP failed") + + +def set_mem_limit(bytes_): + """Cap address space so a runaway alloc can't OOM the host.""" + import resource + resource.setrlimit(resource.RLIMIT_AS, (bytes_, bytes_)) diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_cuda_channel_codec.py b/prototypes/mediator-sandbox/gpu_sandbox/test_cuda_channel_codec.py new file mode 100644 index 000000000..a17467876 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_cuda_channel_codec.py @@ -0,0 +1,80 @@ +#!/usr/bin/env python3 +"""CudaIpcChannel buffer codec tests (in-process, no spawn). + +Tests the new GPU-bounce-buffer tensor pack/unpack that the CudaIpcChannel rides on: + + 1. round-trip — a nested (tensor, dict, tuple, scalar, None) value packs into a + shared GPU buffer (offset table + non-tensor skeleton) and unpacks + bit-identically. + 2. aliasing — unpacked tensors are CLONED out of the buffer, so overwriting the + buffer afterwards does NOT corrupt them (the one-event-in-flight + clone-on-receive correctness rule, design §2). + +Run: + CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/test_cuda_channel_codec.py +""" +import sys + +import torch + +from nnsight.intervention.transport import pack_cuda, unpack_cuda + +DEV = "cuda" +ARENA = 64 << 20 + + +def _equal(a, b): + if torch.is_tensor(a): + return torch.is_tensor(b) and a.dtype == b.dtype and torch.equal(a, b) + if isinstance(a, dict): + return isinstance(b, dict) and a.keys() == b.keys() and all(_equal(a[k], b[k]) for k in a) + if isinstance(a, (tuple, list)): + return type(a) is type(b) and len(a) == len(b) and all(_equal(x, y) for x, y in zip(a, b)) + return a == b or (a is None and b is None) + + +def test_roundtrip(): + buf = torch.empty(ARENA, dtype=torch.uint8, device=DEV) + value = ( + torch.randn(2, 5, 7, device=DEV), + {"k": torch.arange(3, device=DEV), "scale": 2.5}, + torch.randn(1, 16, 768, device=DEV, dtype=torch.bfloat16), # block-output-ish + "meta", + None, + ) + skel, table = pack_cuda(value, buf) + got = unpack_cuda(skel, table, buf) + ok = _equal(value, got) + print(f"[1 roundtrip] nested (tensor/dict/tuple/bf16/scalar/None) bit-identical: {ok}") + return ok + + +def test_aliasing_clone_on_receive(): + buf = torch.empty(ARENA, dtype=torch.uint8, device=DEV) + t = torch.arange(100, device=DEV, dtype=torch.float32) + skel, table = pack_cuda((t,), buf) + got = unpack_cuda(skel, table, buf) + before = got[0].clone() + # Simulate the next event reusing the single buffer: + buf.fill_(0) + survived = torch.equal(got[0], before) and torch.equal(got[0], t) + print(f"[2 aliasing] unpacked tensor survives buffer overwrite (cloned): {survived}") + return survived + + +def main(): + assert torch.cuda.is_available(), "needs CUDA" + results = { + "roundtrip": test_roundtrip(), + "aliasing": test_aliasing_clone_on_receive(), + } + ok = all(results.values()) + print("=" * 72) + print(f"CHANNEL CODEC: {'PASS' if ok else 'FAIL'} — {results}") + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_cuda_channel_ipc.py b/prototypes/mediator-sandbox/gpu_sandbox/test_cuda_channel_ipc.py new file mode 100644 index 000000000..ab2073183 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_cuda_channel_ipc.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python3 +"""CudaIpcChannel across a real spawned process. + +Proves the MediatorChannel halves move events + tensors over a process boundary via +the shared GPU bounce buffer (CUDA IPC), in both directions, under the strict +one-event-in-flight alternation: + + worker put_event(SWAP, (R, t)) -- worker->host event, tensor via buffer + host wait/get_event -- unpacks t (cloned out of buffer) + host put_response(t * 2) -- host->worker response, tensor via buffer + worker wait/get_response -- unpacks t*2 (cloned) + worker put_event(END, ok) -- worker->host, non-tensor payload + +Run (spawn+CUDA needs the unsandboxed shell): + CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/test_cuda_channel_ipc.py +""" +import sys + +import torch +import torch.multiprocessing as mp + +from nnsight.intervention.transport import CudaIpcHostChannel, CudaIpcWorkerChannel + +ARENA = 64 << 20 + + +def _worker(conn, buf): + ch = CudaIpcWorkerChannel(conn, buf) + t = torch.arange(12, device="cuda", dtype=torch.float32).reshape(3, 4) + ch.put_event(("SWAP", ("R", t))) + ch.wait_response() + resp = ch.get_response() + ok = torch.is_tensor(resp) and torch.equal(resp, t * 2) + ch.put_event(("END", ok)) + + +def main(): + assert torch.cuda.is_available(), "needs CUDA" + ctx = mp.get_context("spawn") + buf = torch.empty(ARENA, dtype=torch.uint8, device="cuda") + parent, child = ctx.Pipe() + p = ctx.Process(target=_worker, args=(child, buf), daemon=True) + p.start() + + host = CudaIpcHostChannel(parent, buf) + + host.wait_event() + event, data = host.get_event() + requester, tensor = data + got_swap = event == "SWAP" and requester == "R" and torch.equal( + tensor, torch.arange(12, device="cuda", dtype=torch.float32).reshape(3, 4) + ) + print(f"[1 w->h] host received SWAP tensor over IPC buffer: {got_swap}") + + host.put_response(tensor * 2) + + host.wait_event() + end_event, worker_ok = host.get_event() + round_trip = end_event == "END" and worker_ok is True + print(f"[2 h->w] worker received response*2 + signalled END: {round_trip}") + + p.join(timeout=10) + + ok = got_swap and round_trip + print("=" * 72) + print(f"CUDA-IPC CHANNEL: {'PASS' if ok else 'FAIL'}") + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_fast_lane.py b/prototypes/mediator-sandbox/gpu_sandbox/test_fast_lane.py new file mode 100644 index 000000000..d020d00a8 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_fast_lane.py @@ -0,0 +1,194 @@ +#!/usr/bin/env python3 +"""End-to-end fast lane: confirmed-safe interventions run IN-PROCESS under +isolate_mediators(), and that is the ONLY tier that can run the weight-reading interp +majority (the isolated worker holds weightless dummy modules). + + weights_fast — a logit-lens cell (reads lm_head.weight, calls ln_f) runs under + isolate_mediators() bit-identical to non-isolated AND raises under + forced isolation (fast_lane=False) — proving the fast lane is the + enabling tier, not just a speedup. + inplace_steer — an in-place steering cell runs correctly on the fast lane (it is a + silent no-op under isolation due to clone-on-receive). + renamed — same logit-lens shape on a renamed model (final_norm / + output_projection / decoder_blocks) fast-lanes bit-identical. + footgun_isolates — a cell that imports os / opens a file is NOT fast-laned: it routes to + the worker (isolated) and is contained, host survives. + introspection — a cell reaching for ().__class__... raises FastLaneRejected. + watchdog — a huge bounded loop fast-laned with a short deadline is killed by the + watchdog; the host survives and the next trace still works. + +Run: + CUDA_VISIBLE_DEVICES=5 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python -u \ + prototypes/mediator-sandbox/gpu_sandbox/test_fast_lane.py +""" +import sys + +import torch +import torch.nn.functional as F + +from nnsight import LanguageModel +from nnsight.intervention.fastlane import FastLaneRejected +from nnsight.intervention.isolation import isolate_mediators + +PROMPT = "The Eiffel Tower is in the city of" + + +# --- the weight-reading interp readout (logit lens), the workload isolation cannot run -- +def _logit_lens(blocks, norm, head, layers): + rows = [] + with torch.no_grad(): + for i in layers: + out = blocks[i].output + hidden = out[0] if isinstance(out, tuple) else out + normed = norm(hidden) + logits = F.linear(normed, head.weight) # host-weight read — worker-impossible + rows.append(logits[:, -1, :]) + return torch.stack(rows, dim=0) + + +def _lens_gpt2(model): + with model.trace(PROMPT): + g = _logit_lens(model.transformer.h, model.transformer.ln_f, + model.lm_head, [0, 4, 8]).save() + return g + + +def test_weights_fast(model): + ref = _lens_gpt2(model) # non-isolated baseline + with isolate_mediators(): # fast lane on (default) + got = _lens_gpt2(model) + fast_ok = torch.is_tensor(got) and torch.equal(ref, got) + + # forced isolation: the same cell must FAIL (weightless dummy modules in the worker) + forced_failed = False + try: + with isolate_mediators(fast_lane=False, timeout=30): + _lens_gpt2(model) + except Exception: # noqa: BLE001 + forced_failed = True + + print(f"[weights_fast] fast-lane bit-identical={fast_ok} " + f"(max|Δ|={(ref-got).abs().max().item() if fast_ok else float('nan')}); " + f"forced-isolation-raises={forced_failed}", flush=True) + return fast_ok and forced_failed + + +def test_inplace_steer(model): + def steer(inplace_ctx): + # steer block 6 toward a token's unembed row, read final logits + with model.trace(PROMPT): + with torch.no_grad(): + direction = F.normalize(model.lm_head.weight[5000].float(), dim=0) + out = model.transformer.h[6].output + hidden = out[0] if isinstance(out, tuple) else out + scale = hidden.norm(dim=-1).mean() + hidden[:] = hidden + 6.0 * scale * direction.to(hidden.dtype) # in-place + last = model.transformer.h[-1].output + last = last[0] if isinstance(last, tuple) else last + normed = model.transformer.ln_f(last) + logits = F.linear(normed, model.lm_head.weight)[:, -1, :].save() + return logits + + ref = steer(None) # non-isolated + with isolate_mediators(): + got = steer(None) # fast lane + ok = torch.is_tensor(got) and torch.equal(ref, got) + print(f"[inplace_steer] fast-lane in-place write bit-identical={ok} " + f"(max|Δ|={(ref-got).abs().max().item() if ok else float('nan')})", flush=True) + return ok + + +def test_renamed(): + rename = {"transformer.ln_f": "final_norm", "lm_head": "output_projection", + "transformer.h": "decoder_blocks"} + model = LanguageModel("gpt2", device_map="cuda", dispatch=True, rename=rename) + + def body(): + with model.trace(PROMPT): + g = _logit_lens(model.decoder_blocks, model.final_norm, + model.output_projection, [0, 3, 6]).save() + return g + + ref = body() + with isolate_mediators(): + got = body() + ok = torch.is_tensor(got) and torch.equal(ref, got) + print(f"[renamed] renamed-model logit-lens fast-lane bit-identical={ok} " + f"(max|Δ|={(ref-got).abs().max().item() if ok else float('nan')})", flush=True) + return ok + + +def test_footgun_isolates(model): + # A footgun cell is NOT confirmable -> routes to the worker (isolated), contained. + # The worker can't run the weight read either, so it errors cleanly; the point is the + # host SURVIVES and a subsequent normal trace still works. + raised = False + try: + with isolate_mediators(fast_lane=True, timeout=20): + with model.trace(PROMPT): + import os # noqa: F401 — footgun: isolated, never fast-laned + _ = model.transformer.h[6].output[0].save() + except Exception: # noqa: BLE001 + raised = True + # host survives: a clean fast-lane trace works right after + after = _lens_gpt2(model) + survived = torch.is_tensor(after) + print(f"[footgun_isolates] footgun routed off the fast lane (raised={raised}); " + f"host survived={survived}", flush=True) + return survived + + +def test_introspection(model): + rejected = False + try: + with isolate_mediators(): + with model.trace(PROMPT): + _ = ().__class__.__bases__ # introspection escape + _ = model.transformer.h[6].output[0].save() + except FastLaneRejected: + rejected = True + except Exception as e: # noqa: BLE001 + rejected = "introspection" in str(e).lower() or "reject" in str(e).lower() + print(f"[introspection] introspection escape rejected={rejected}", flush=True) + return rejected + + +def test_watchdog(model): + # A huge bounded loop passes the static gate (range is "bounded") but would hang the + # host; the watchdog must kill it. Short deadline so the test is quick. + killed = False + try: + with isolate_mediators(fast_lane_timeout=4.0): + with model.trace(PROMPT): + h = model.transformer.h[6].output[0] + acc = 0 + for _i in range(10 ** 12): # bounded literal, but enormous + acc += 1 + _ = h.save() + except Exception: # noqa: BLE001 — FastLaneTimeout surfaces through the trace + killed = True + after = _lens_gpt2(model) # host survived + survived = torch.is_tensor(after) + print(f"[watchdog] runaway loop killed={killed}; host survived={survived}", flush=True) + return killed and survived + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = { + "weights_fast": test_weights_fast(model), + "inplace_steer": test_inplace_steer(model), + "renamed": test_renamed(), + "footgun_isolates": test_footgun_isolates(model), + "introspection": test_introspection(model), + "watchdog": test_watchdog(model), + } + print("=" * 72, flush=True) + print(f"FAST LANE (end-to-end): {results}", flush=True) + sys.exit(0 if all(results.values()) else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_fastlane_classifier.py b/prototypes/mediator-sandbox/gpu_sandbox/test_fastlane_classifier.py new file mode 100644 index 000000000..f13bddf4c --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_fastlane_classifier.py @@ -0,0 +1,226 @@ +#!/usr/bin/env python3 +"""Unit tests for the fast-lane safety classifier (GPU-free, no model). + +Confirms the detect-and-confirm gate on the ACTUAL interp-workload shapes (logit lens, +steering, activation patching, attribution) plus footgun payloads, and — per the testing +rules — on RENAMED module structures so nothing is keyed to GPT-2 naming. The classifier +walks the effective code (closures resolved through globals/closure cells), so these tests +use the same closure-wrapped shape the real harness uses (a build()/capture() lambda +calling helper functions). + +Run: + PYTHONPATH=src /disk/u/zikai/anaconda3/envs/hf-serve/bin/python -u \ + prototypes/mediator-sandbox/gpu_sandbox/test_fastlane_classifier.py +""" +import sys + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from nnsight.intervention.fastlane import FAST, ISOLATE, REJECT, classify_callable + + +# --- stand-ins for the host objects a trace body closes over ------------------------- +# Real nn.Modules so the classifier's host-object detection (calling them / reading +# .weight) exercises the genuine path, not a mock. +class _Block(nn.Module): + def __init__(self, d): + super().__init__() + self.lin = nn.Linear(d, d) + self.output = torch.zeros(1, 3, d) # stands in for the eproperty read target + + def forward(self, x): + return self.lin(x) + + +def _make_model(norm_name="ln_f", head_name="lm_head", blocks_name="h", d=8, n=4): + """A tiny model with CONFIGURABLE attribute names — the renamed-structure test tool.""" + m = nn.Module() + setattr(m, blocks_name, nn.ModuleList([_Block(d) for _ in range(n)])) + setattr(m, norm_name, nn.LayerNorm(d)) + setattr(m, head_name, nn.Linear(d, 50, bias=False)) + return m + + +# --- helper functions a "trace body" calls (the closure chain) ------------------------ +def _untuple(x): + return x[0] if isinstance(x, tuple) else x + + +def _lens_proxy(blocks, norm, head, *, layers): + rows = [] + with torch.no_grad(): + for i in (range(len(blocks)) if layers == "all" else layers): + normed = norm(_untuple(blocks[i].output)) + logits = F.linear(normed, head.weight) # host-weight read — fast-lane-only + rows.append(logits[:, -1, :]) + return torch.stack(rows, dim=0) + + +def _steer_inplace(blocks, head, *, layer, token_id, alpha): + with torch.no_grad(): + direction = F.normalize(head.weight[token_id].float(), dim=0) + out = blocks[layer].output + hidden = out[0] if isinstance(out, tuple) else out + hidden[:] = hidden + alpha * direction # in-place write (in_place flag) + return hidden + + +# ===================================================================================== +CASES = [] + + +def case(expect): + def deco(fn): + CASES.append((fn.__name__, fn, expect)) + return fn + return deco + + +# ---- SAFE interp workloads (must classify FAST) ------------------------------------- +@case(FAST) +def logit_lens_gpt2_style(): + m = _make_model() + return lambda: _lens_proxy(m.h, m.ln_f, m.lm_head, layers="all").save() + + +@case(FAST) +def logit_lens_renamed_structure(): + # non-GPT-2 names: decoder_blocks / final_norm / output_projection + m = _make_model(norm_name="final_norm", head_name="output_projection", + blocks_name="decoder_blocks") + return lambda: _lens_proxy(m.decoder_blocks, m.final_norm, m.output_projection, + layers=[0, 2]).save() + + +@case(FAST) +def steering_inplace_is_fast(): + # in-place steering is the fast lane's correctness win (silent no-op under isolation) + m = _make_model() + return lambda: _steer_inplace(m.h, m.lm_head, layer=1, token_id=5, alpha=6.0).save() + + +@case(FAST) +def boundary_replacement_write(): + m = _make_model() + + def body(): + out = m.h[1].output + new = out * 2.0 + m.h[2].output = new # nnsight boundary write (SWAP), allowed + return m.ln_f(new)[:, -1, :] + return body + + +@case(FAST) +def backward_attribution_shape(): + m = _make_model() + + def body(): + a = m.h[2].output + a.requires_grad_(True) + normed = m.ln_f(a) + metric = F.linear(normed, m.lm_head.weight).sum() + with metric.backward(): + g = a.grad.save() + return g + return body + + +# ---- footguns that must NOT reach the fast lane ------------------------------------- +@case(ISOLATE) +def imports_isolate(): + def body(): + import os + return os.getpid() + return body + + +@case(ISOLATE) +def while_loop_isolate(): + def body(): + x = 0 + while x < 10: + x = x + 1 + return x + return body + + +@case(ISOLATE) +def unresolved_global_call_isolate(): + def body(): + return some_undefined_helper(3) # noqa: F821 — unknown authority + return body + + +@case(ISOLATE) +def open_file_isolate(): + def body(): + return open("/etc/passwd").read() + return body + + +@case(REJECT) +def introspection_subclasses_reject(): + def body(): + return ().__class__.__bases__ + return body + + +@case(REJECT) +def getattr_escape_reject(): + def body(): + return getattr(torch, "save") + return body + + +@case(REJECT) +def dunder_subscript_reject(): + def body(): + d = {} + return d["__builtins__"] + return body + + +@case(ISOLATE) +def host_attr_write_isolate(): + m = _make_model() + + def body(): + m.ln_f.eps = 1.0 # mutating host state visible to siblings + return m.ln_f.weight + return body + + +# ---- flag detection ---------------------------------------------------------------- +def main(): + results = {} + for name, factory, expect in CASES: + fn = factory() + v = classify_callable(fn) + ok = v.tier == expect + results[name] = ok + print(f"[{'PASS' if ok else 'FAIL'}] {name}: got {v.tier} expected {expect} " + f"(diff={v.differentiate} inplace={v.in_place} weights={v.touches_host_weights}) " + f"-- {v.reason}", flush=True) + + # flag assertions on specific cases + flag_checks = { + "backward sets differentiate": classify_callable(backward_attribution_shape()).differentiate is True, + "steering sets in_place": classify_callable(steering_inplace_is_fast()).in_place is True, + "lens reads host weights": classify_callable(logit_lens_gpt2_style()).touches_host_weights is True, + "replacement write is not in_place": classify_callable(boundary_replacement_write()).in_place is False, + } + for k, v in flag_checks.items(): + results[k] = v + print(f"[{'PASS' if v else 'FAIL'}] {k}", flush=True) + + print("=" * 72, flush=True) + npass = sum(results.values()) + print(f"FAST-LANE CLASSIFIER: {npass}/{len(results)} passed", flush=True) + sys.exit(0 if all(results.values()) else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_functional.py b/prototypes/mediator-sandbox/gpu_sandbox/test_functional.py new file mode 100644 index 000000000..857e45df6 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_functional.py @@ -0,0 +1,86 @@ +#!/usr/bin/env python3 +"""Functional test: real nnsight workloads still work, with the intervention op +offloaded to the isolated GPU worker. + +For each real interpretability op we run a normal `model.trace()` two ways — +the op inline (reference) vs the op offloaded to the sandbox — and require the +final logits to match. The activation is delivered by nnsight's real machinery; +only the user's op runs in the isolated worker. + +Run: CUDA_VISIBLE_DEVICES=6 PYTHONPATH=/src .../bin/python test_functional.py +""" +import os +import sys + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import torch + +from nnsight import LanguageModel +from gpu_sandbox import GPUSandbox + +PROMPT = "The Eiffel Tower is in the city of" +LAYER = 6 + + +# Real intervention ops (defined in __main__ so cloudpickle serializes them +# by value — the worker needs no import to rebuild them). +def scale(t): + return t * 0.5 + + +def steer(t): + return t + 1.5 + + +def ablate_mean(t): + return t - t.mean(dim=-1, keepdim=True) + + +def project_norm(t): + # a "read"-style op that changes shape: per-token L2 norm + return t.norm(dim=-1) + + +def main(): + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + sb = GPUSandbox() + + results = {} + + # 1. modify-activation ops: compare final logits, inline vs sandboxed + for name, fn in {"scale": scale, "steer": steer, "ablate_mean": ablate_mean}.items(): + with model.trace(PROMPT): + model.transformer.h[LAYER].output = fn(model.transformer.h[LAYER].output) + ref = model.lm_head.output.save() + with model.trace(PROMPT): + h = model.transformer.h[LAYER].output + model.transformer.h[LAYER].output = sb.apply(h, fn) + out = model.lm_head.output.save() + ok = torch.allclose(ref, out, atol=1e-3, rtol=0) + results[name] = ok + print(f"[functional] {name:12s} inline==sandboxed: {ok} | max|Δ|={(ref - out).abs().max():.2e}") + + # 2. read-style op that changes shape (norm per token): compare the read value + with model.trace(PROMPT): + ref_read = model.transformer.h[LAYER].output.norm(dim=-1).save() + with model.trace(PROMPT): + h = model.transformer.h[LAYER].output + read = sb.apply(h, project_norm).save() + ok = torch.allclose(ref_read, read, atol=1e-3, rtol=0) + results["project_norm(read)"] = ok + print(f"[functional] {'project_norm':12s} inline==sandboxed: {ok} | max|Δ|={(ref_read - read).abs().max():.2e}") + + # 3. the worker survives many requests (pool-readiness) + reused = sb.alive() + results["worker_alive_after_all"] = reused + print(f"[functional] worker still alive after all requests: {reused}") + + sb.close() + ok_all = all(results.values()) + print("=" * 60) + print(f"GPU-SANDBOX FUNCTIONAL: {'PASS' if ok_all else 'FAIL'} — {results}") + sys.exit(0 if ok_all else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_ablate.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_ablate.py new file mode 100644 index 000000000..d586b9d46 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_ablate.py @@ -0,0 +1,250 @@ +#!/usr/bin/env python3 +"""Ablation (zero / mean knockout) on the ISOLATED tier: tracer.ablate replaces a module's +output with a baseline (zeros, or the within-sequence mean) via a REPLACEMENT swap, so the +knockout crosses the isolation boundary and propagates through the host's real forward — +where the hand-written in-place form (``hidden[:] = 0``) silently no-ops (the worker mutates +its delivered clone, no SWAP fires, the host's activation is untouched). + +Like tracer.steer/patch (and unlike tracer.unembed), ablation touches no host weights — the +baseline is derived from (or independent of) the delivered activation — so it rides the +existing Events.SWAP with no host round-trip and no isolated/in-process branch: the SAME +method is correct in-process, on the fast lane, and in the isolated worker. + +``mode="mean"`` here is the SELF-CONTAINED within-sequence mean (each position → the +per-example mean over the token dimension). Reference-distribution mean ablation (the mean +activation over a dataset, per docs/patterns/ablation.md) is a precomputed value transplanted +via tracer.patch — not this mode. + + zero_single — zero-ablate one block; forced-isolation == in-process AND downstream != + baseline (the ablation took effect across the boundary). + mean_single — mean-ablate one block; forced-isolation == in-process AND != baseline. + crux — THE crux: under forced isolation the in-place zero is a no-op (downstream + == un-ablated baseline) while tracer.ablate actually ablates (downstream != + baseline) and equals the in-process zero-ablated result. + tuple_output — zero-ablate an attention output (a tuple `(tensor, None)`): the whole-tuple + replacement branch, forced-isolation == in-process. + renamed — renamed model (decoder_blocks): forced-isolation == in-process. + matches_manual— tracer.ablate (in-process) equals the manual zeros_like / mean-over-seq + replacement it names. + bad_mode — an unknown mode raises ValueError (no silent wrong-ablation). + +Run: + CUDA_VISIBLE_DEVICES=7 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/nnsight-tf/bin/python -u \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_ablate.py +""" +import sys + +import torch + +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +PROMPT = "The Eiffel Tower is in the city of" + + +def _both(build): + """Run build() in-process and under forced isolation; return (ref, got).""" + ref = build() + with isolate_mediators(fast_lane=False, timeout=30): + got = build() + return ref, got + + +def _eq(ref, got): + return (torch.is_tensor(ref) and torch.is_tensor(got) + and ref.shape == got.shape and torch.equal(ref, got)) + + +def _delta(ref, got): + return (ref - got).abs().max().item() if _eq(ref, got) else float("nan") + + +def _baseline_downstream(model, read): + with model.trace(PROMPT): + with torch.no_grad(): + o = read.output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + + +def test_zero_single(model): + base = _baseline_downstream(model, model.transformer.h[10]) + + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.ablate(model.transformer.h[6], "zero") + o = model.transformer.h[10].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + ref, got = _both(build) + ok = _eq(ref, got) and not _eq(ref, base) + print(f"[zero_single] isolated zero-ablate bit-identical={_eq(ref, got)} " + f"took_effect={not _eq(ref, base)} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_mean_single(model): + base = _baseline_downstream(model, model.transformer.h[10]) + + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.ablate(model.transformer.h[6], "mean") + o = model.transformer.h[10].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + ref, got = _both(build) + ok = _eq(ref, got) and not _eq(ref, base) + print(f"[mean_single] isolated mean-ablate bit-identical={_eq(ref, got)} " + f"took_effect={not _eq(ref, base)} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_crux(model): + """The crux: replacement swap crosses the boundary; in-place zero is a silent no-op.""" + base = _baseline_downstream(model, model.transformer.h[10]) + + def ablate_downstream(): # tracer.ablate (replacement swap) + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.ablate(model.transformer.h[6], "zero") + o = model.transformer.h[10].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + + def inplace_downstream(): # hand-written in-place zero + with model.trace(PROMPT): + with torch.no_grad(): + out = model.transformer.h[6].output + hidden = out[0] if isinstance(out, tuple) else out + hidden[:] = 0 + o = model.transformer.h[10].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + + ablate_ip = ablate_downstream() # in-process ablate + with isolate_mediators(fast_lane=False, timeout=30): + ablate_iso = ablate_downstream() # isolated ablate (replacement swap) + inplace_iso = inplace_downstream() # isolated in-place (no-op) + + inplace_is_noop = _eq(inplace_iso, base) # in-place never crossed the boundary + ablate_took_effect = not _eq(ablate_iso, base) # replacement swap changed the host forward + ablate_correct = _eq(ablate_iso, ablate_ip) # ... to exactly the in-process result + ok = inplace_is_noop and ablate_took_effect and ablate_correct + print(f"[crux] isolated in-place zero is a no-op={inplace_is_noop}; " + f"isolated ablate took effect={ablate_took_effect}; " + f"ablate iso==in-process={ablate_correct} (max|Δ|={_delta(ablate_iso, ablate_ip)})", flush=True) + return ok + + +def test_tuple_output(model): + base = _baseline_downstream(model, model.transformer.h[10]) + + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.ablate(model.transformer.h[6].attn, "zero") + o = model.transformer.h[10].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + ref, got = _both(build) + ok = _eq(ref, got) and not _eq(ref, base) + print(f"[tuple_output] tuple (attn) zero-ablate bit-identical={_eq(ref, got)} " + f"took_effect={not _eq(ref, base)} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_renamed(): + rename = {"transformer.ln_f": "final_norm", "lm_head": "output_projection", + "transformer.h": "decoder_blocks"} + model = LanguageModel("gpt2", device_map="cuda", dispatch=True, rename=rename) + + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.ablate(model.decoder_blocks[3], "mean") + o = model.decoder_blocks[9].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[renamed] renamed-model isolated ablate bit-identical={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_matches_manual(model): + # tracer.ablate (in-process) must equal the manual zeros_like / mean-over-seq replacement. + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.ablate(model.transformer.h[6], "zero") + z_api = (model.transformer.h[10].output[0] + if isinstance(model.transformer.h[10].output, tuple) + else model.transformer.h[10].output)[:, -1, :].save() + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.ablate(model.transformer.h[6], "mean") + o = model.transformer.h[10].output + m_api = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + + with model.trace(PROMPT): + with torch.no_grad(): + out = model.transformer.h[6].output + is_tuple = isinstance(out, tuple) + hidden = out[0] if is_tuple else out + zeroed = torch.zeros_like(hidden) + model.transformer.h[6].output = (zeroed, *out[1:]) if is_tuple else zeroed + o = model.transformer.h[10].output + z_manual = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + with model.trace(PROMPT): + with torch.no_grad(): + out = model.transformer.h[6].output + is_tuple = isinstance(out, tuple) + hidden = out[0] if is_tuple else out + meaned = hidden.mean(dim=-2, keepdim=True).expand_as(hidden).contiguous() + model.transformer.h[6].output = (meaned, *out[1:]) if is_tuple else meaned + o = model.transformer.h[10].output + m_manual = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + + ok = _eq(z_api, z_manual) and _eq(m_api, m_manual) + print(f"[matches_manual] zero == manual zeros_like={_eq(z_api, z_manual)} " + f"(max|Δ|={_delta(z_api, z_manual)}); mean == manual mean-over-seq={_eq(m_api, m_manual)} " + f"(max|Δ|={_delta(m_api, m_manual)})", flush=True) + return ok + + +def test_bad_mode(model): + raised = False + try: + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.ablate(model.transformer.h[6], "median") + model.transformer.h[10].output # force the body to run + except ValueError: + raised = True + except Exception as e: # any other surfaced error is not the contract + print(f"[bad_mode] wrong exception type: {type(e).__name__}: {e}", flush=True) + print(f"[bad_mode] unknown mode raises ValueError={raised}", flush=True) + return raised + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = { + "zero_single": test_zero_single(model), + "mean_single": test_mean_single(model), + "crux": test_crux(model), + "tuple_output": test_tuple_output(model), + "renamed": test_renamed(), + "matches_manual": test_matches_manual(model), + "bad_mode": test_bad_mode(model), + } + print("=" * 72, flush=True) + print(f"ISOLATED ABLATE: {results}", flush=True) + sys.exit(0 if all(results.values()) else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_acceptance.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_acceptance.py new file mode 100644 index 000000000..53de5a94e --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_acceptance.py @@ -0,0 +1,133 @@ +#!/usr/bin/env python3 +"""Isolated trace acceptance — multi-invoke, non-standard module names, exception, timeout. + +Complements the read/swap end-to-end test (read/swap bit-identical). All comparisons are against +the in-process result on the SAME model. + +Run: + CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_acceptance.py +""" +import sys +import time + +import torch +import torch.nn as nn + +import nnsight +from nnsight import LanguageModel, NNsight +from nnsight.intervention.isolation import isolate_mediators + +PROMPT = "The Eiffel Tower is in the city of" + + +# --------------------------------------------------------------------------- # +# Non-standard module names (testing rule: no GPT-2-only assumptions) # +# --------------------------------------------------------------------------- # +class TinyNet(nn.Module): + def __init__(self, d=16, n=3): + super().__init__() + self.embed = nn.Embedding(32, d) + self.decoder_blocks = nn.ModuleList([nn.Linear(d, d) for _ in range(n)]) + self.output_projection = nn.Linear(d, 32) + + def forward(self, x): + h = self.embed(x) + for blk in self.decoder_blocks: + h = torch.relu(blk(h)) + return self.output_projection(h) + + +def test_nonstandard_names(): + torch.manual_seed(0) + net = TinyNet().cuda() + model = NNsight(net) + x = torch.randint(0, 32, (1, 6)).cuda() + + with model.trace(x): + ref = model.decoder_blocks[1].output.save() + with isolate_mediators(fast_lane=False): + with model.trace(x): + got = model.decoder_blocks[1].output.save() + ok = torch.equal(ref, got) + print(f"[names] non-standard 'decoder_blocks' read == in-process: {ok} (max|Δ|={(ref-got).abs().max().item():.2e})") + return ok + + +def test_multi_invoke(model): + # Two invokes, each its own isolated worker; batch narrowing must keep rows separate. + with model.trace() as t: + with t.invoke("The capital of France is"): + a_ref = model.transformer.h[5].output[0].save() + with t.invoke("The Eiffel Tower is in"): + b_ref = model.transformer.h[5].output[0].save() + with isolate_mediators(fast_lane=False): + with model.trace() as t: + with t.invoke("The capital of France is"): + a_got = model.transformer.h[5].output[0].save() + with t.invoke("The Eiffel Tower is in"): + b_got = model.transformer.h[5].output[0].save() + ok = torch.equal(a_ref, a_got) and torch.equal(b_ref, b_got) + no_cross = not torch.equal(a_got, b_got) # rows are genuinely different + print(f"[multi] two isolated invokes bit-identical: {ok} | rows distinct (no cross-leak): {no_cross}") + return ok and no_cross + + +def test_exception(model): + # A footgun (ValueError) in user code must surface in the user's context. + raised = None + try: + with isolate_mediators(fast_lane=False): + with model.trace(PROMPT): + _ = model.transformer.h[6].output[0] + raise ValueError("boom-from-user-code") + except Exception as e: # noqa: BLE001 + raised = e + ok = raised is not None and "boom-from-user-code" in str(raised) + print(f"[exc] user exception propagated to host: {ok} ({type(raised).__name__ if raised else None})") + return ok + + +def test_timeout(model): + # An infinite loop in user code must be killed and the host must survive. + t0 = time.time() + killed = None + try: + with isolate_mediators(fast_lane=False, timeout=5): + with model.trace(PROMPT): + out = model.transformer.h[6].output[0] + while True: # footgun: hang + out = out + 1 + except Exception as e: # noqa: BLE001 + killed = e + dt = time.time() - t0 + # host still works afterwards + with model.trace(PROMPT): + alive = model.transformer.h[0].output[0].save() + host_ok = torch.is_tensor(alive) + print(f"[hang] infinite loop killed in {dt:.1f}s (err={type(killed).__name__ if killed else None}); host survives: {host_ok}") + return killed is not None and host_ok + + +def main(): + assert torch.cuda.is_available() + results = {} + results["names"] = test_nonstandard_names() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results["multi"] = test_multi_invoke(model) + results["exc"] = test_exception(model) + # timeout test depends on isolate_mediators supporting a `timeout=` kwarg + if "timeout" in isolate_mediators.__doc__ or True: + try: + results["hang"] = test_timeout(model) + except TypeError as e: + print(f"[hang] SKIP (timeout kwarg not wired): {e}") + ok = all(results.values()) + print("=" * 72) + print(f"ISOLATED TRACE ACCEPTANCE: {'PASS' if ok else 'FAIL'} — {results}") + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward.py new file mode 100644 index 000000000..a288a9fbf --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward.py @@ -0,0 +1,123 @@ +#!/usr/bin/env python3 +"""Increment 1 — read-then-backward under isolation == in-process, bit-identical. + +Scope (the simplest correct slice): single invoke, NO swaps, the gradient target is +an ON-PATH tensor-output module (``ln_f.output``), the loss is a scalar reduction of a +read activation. The host graph stays intact (a read does not replace the host's real +tensor), so the host can run the real backward; the worker only computes the seed +gradient at the loss's delivered leaves and reads grads back by PATH. + + read_backward — grad of ln_f.output w.r.t. a sum-of-logits loss; max|Δ|=0 vs in-process. + +NOT in scope here (later increments): gradient through a SWAP (severs the host graph), +tuple-element targets (off-path), batched/multi-invoke backward. + +Run: + CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python -u \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_backward.py +""" +import sys + +import torch + +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +PROMPT = "The Eiffel Tower is in the city of" + + +def _read_backward(model): + # ln_f.output is a tensor-output module ON the autograd path. + with model.trace(PROMPT): + hs = model.transformer.ln_f.output + with model.lm_head.output.sum().backward(): + g = hs.grad.save() + return g + + +def test_read_backward(model): + ref = _read_backward(model) + with isolate_mediators(fast_lane=False, timeout=30): + got = _read_backward(model) + + ok = ( + torch.is_tensor(ref) + and torch.is_tensor(got) + and ref.shape == got.shape + and torch.equal(ref, got) + ) + delta = (ref - got).abs().max().item() if (torch.is_tensor(ref) and torch.is_tensor(got) and ref.shape == got.shape) else float("nan") + print( + f"[read_backward] ref={tuple(ref.shape) if torch.is_tensor(ref) else ref} " + f"got={tuple(got.shape) if torch.is_tensor(got) else got} " + f"max|Δ|={delta} match={ok}", + flush=True, + ) + return ok + + +def test_read_backward_nonstd(): + # Testing rule: vary names, don't assume GPT-2 conventions. Rename the norm + head + # to non-standard user-facing paths; the host must resolve grads by the real path. + rename = {"transformer.ln_f": "final_norm", "lm_head": "output_projection"} + model = LanguageModel("gpt2", device_map="cuda", dispatch=True, rename=rename) + + def body(): + with model.trace(PROMPT): + hs = model.final_norm.output + with model.output_projection.output.sum().backward(): + g = hs.grad.save() + return g + + ref = body() + with isolate_mediators(fast_lane=False, timeout=30): + got = body() + ok = torch.is_tensor(ref) and torch.is_tensor(got) and torch.equal(ref, got) + delta = (ref - got).abs().max().item() if ok or (torch.is_tensor(ref) and torch.is_tensor(got)) else float("nan") + print(f"[nonstd] final_norm.output grad isolated==in-proc: {ok} (max|Δ|={delta})", flush=True) + return ok + + +def test_derived_target_fails_clean(model): + # Boundary: gradient of a USER-DERIVED tensor has no host-side graph under isolation. + # It must fail CLEANLY (a clear error), never hang or return a silently-wrong grad. + def body(): + with model.trace(PROMPT): + hs = model.transformer.ln_f.output + derived = hs * 2 # computed in user code — no module-path provenance + with model.lm_head.output.sum().backward(): + g = derived.grad.save() + return g + + try: + body() + print("[derived] UNEXPECTEDLY OK — should have refused a derived-tensor grad", flush=True) + return False + except Exception as e: # noqa: BLE001 + msg = str(e) + ok = "isolation" in msg.lower() + print(f"[derived] failed cleanly={ok}: {type(e).__name__}: {msg[:110]}", flush=True) + return ok + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = { + "read_backward": test_read_backward(model), + "nonstd": test_read_backward_nonstd(), + "derived_fails_clean": _derived_isolated(model), + } + print("=" * 72, flush=True) + print(f"ISOLATED BACKWARD (increment 1): {results}", flush=True) + sys.exit(0 if all(results.values()) else 1) + + +def _derived_isolated(model): + with isolate_mediators(fast_lane=False, timeout=30): + return test_derived_target_fails_clean(model) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_batched_backward.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_batched_backward.py new file mode 100644 index 000000000..0e960cda6 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_batched_backward.py @@ -0,0 +1,146 @@ +#!/usr/bin/env python3 +"""Batched backward under isolation == in-process, bit-identical. + +"Batched" here is a LIST input (`model.trace([A, B, ...])`): one mediator, one forward over +the padded batch, one backward, per-row gradients. There is no per-invoke narrowing (the +mediator's batch_group is None), so the worker's delivered clone and the host's retained real +are both full-batch and their shapes match: the read-path and grad-through-swap seam-stitch +work unchanged, just on a (batch, seq, hidden) tensor. + +(Backward inside MULTIPLE `tracer.invoke(...)` contexts is a separate, unsupported structure: +it raises `MissedProviderError` IN-PROCESS too, so it is a core nnsight limitation, not an +isolation gap; not covered here.) + + two_rows — grad of ln_f.output for a 2-prompt batch, isolated == in-process, shape (2,·,768). + three_rows — same for a 3-prompt batch. + upstream — grad of an upstream block (h[3].output) in a batched trace. + swap — batched + a swap at h[6]; grad of upstream h[3] flows through the seam (batched + grad-through-swap), isolated == in-process. + renamed — renamed model (decoder_blocks), batched backward. + +Run: + CUDA_VISIBLE_DEVICES=7 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/nnsight-tf/bin/python -u \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_batched_backward.py +""" +import sys + +import torch + +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +A = "The Eiffel Tower is in the city of" +B = "A red bicycle was left near the river of" +C = "She quietly closed the heavy wooden front" + + +def U(o): + return o[0] if isinstance(o, tuple) else o + + +def _eq(a, b): + return (torch.is_tensor(a) and torch.is_tensor(b) + and a.shape == b.shape and torch.equal(a, b)) + + +def _delta(a, b): + return (a - b).abs().max().item() if _eq(a, b) else float("nan") + + +def _both(build): + ref = build() + with isolate_mediators(fast_lane=False, timeout=30): + got = build() + return ref, got + + +def test_two_rows(model): + def build(): + with model.trace([A, B]): + hs = U(model.transformer.ln_f.output) + with model.lm_head.output.sum().backward(): + g = hs.grad.save() + return g + ref, got = _both(build) + ok = _eq(ref, got) and ref.shape[0] == 2 + print(f"[two_rows] batched grad bit-identical={_eq(ref, got)} shape={tuple(ref.shape)} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_three_rows(model): + def build(): + with model.trace([A, B, C]): + hs = U(model.transformer.ln_f.output) + with model.lm_head.output.sum().backward(): + g = hs.grad.save() + return g + ref, got = _both(build) + ok = _eq(ref, got) and ref.shape[0] == 3 + print(f"[three_rows] batched grad bit-identical={_eq(ref, got)} shape={tuple(ref.shape)} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_upstream(model): + def build(): + with model.trace([A, B]): + up = U(model.transformer.h[3].output) + with model.lm_head.output.sum().backward(): + g = up.grad.save() + return g + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[upstream] batched grad of h[3] bit-identical={ok} shape={tuple(ref.shape)} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_swap(model): + def build(): + with model.trace([A, B]): + up = U(model.transformer.h[3].output) + mid = U(model.transformer.h[6].output) + model.transformer.h[6].output = mid * 2.0 + loss = model.lm_head.output.sum() + with loss.backward(): + g = up.grad.save() + return g + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[swap] batched grad-through-swap bit-identical={ok} shape={tuple(ref.shape)} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_renamed(): + rename = {"transformer.ln_f": "final_norm", "lm_head": "output_projection", + "transformer.h": "decoder_blocks"} + model = LanguageModel("gpt2", device_map="cuda", dispatch=True, rename=rename) + + def build(): + with model.trace([A, B]): + hs = U(model.final_norm.output) + with model.output_projection.output.sum().backward(): + g = hs.grad.save() + return g + ref, got = _both(build) + ok = _eq(ref, got) and ref.shape[0] == 2 + print(f"[renamed] batched grad (renamed) bit-identical={_eq(ref, got)} shape={tuple(ref.shape)} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = { + "two_rows": test_two_rows(model), + "three_rows": test_three_rows(model), + "upstream": test_upstream(model), + "swap": test_swap(model), + "renamed": test_renamed(), + } + print("=" * 72, flush=True) + print(f"ISOLATED BATCHED BACKWARD: {results}", flush=True) + sys.exit(0 if all(results.values()) else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cache.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cache.py new file mode 100644 index 000000000..a53a6d2fd --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cache.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python3 +"""tracer.cache() under isolation == in-process, bit-identical. + + one — cache a single module (h[6]); cached output matches in-process. + multi — cache several modules; all entries match. + inputs — include_inputs=True; inputs cached too. + +Run: + CUDA_VISIBLE_DEVICES=6 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cache.py +""" +import sys + +import torch + +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +PROMPT = "The Eiffel Tower is in the city of" + + +def _entry(cache, path): + e = cache[path] + return e[-1] if isinstance(e, list) else e + + +def test_one(model): + h6 = model.transformer.h[6] + with model.trace(PROMPT) as t: + ref = t.cache(modules=[h6]).save() + with isolate_mediators(fast_lane=False, timeout=30): + with model.trace(PROMPT) as t: + got = t.cache(modules=[h6]).save() + key = h6.path # derive the key from the envoy path, don't hardcode the prefix + rk, gk = sorted(ref.keys()), sorted(got.keys()) + ok = rk == gk and len(gk) > 0 and torch.equal( + _entry(ref, key).output[0], _entry(got, key).output[0] + ) + print(f"[one] keys ref={rk} got={gk} match={ok}", flush=True) + return ok + + +def test_multi(model): + mods = [model.transformer.h[2], model.transformer.h[5], model.transformer.h[9]] + with model.trace(PROMPT) as t: + ref = t.cache(modules=mods).save() + with isolate_mediators(fast_lane=False, timeout=30): + with model.trace(PROMPT) as t: + got = t.cache(modules=mods).save() + paths = [m.path for m in mods] + ok = sorted(ref.keys()) == sorted(got.keys()) and all( + torch.equal(_entry(ref, p).output[0], _entry(got, p).output[0]) for p in paths + ) + print(f"[multi] {len(got.keys())} keys match={ok}", flush=True) + return ok + + +def test_inputs(model): + # include_inputs=True: the cached module inputs must match in-process too. + h4 = model.transformer.h[4] + with model.trace(PROMPT) as t: + ref = t.cache(modules=[h4], include_inputs=True).save() + with isolate_mediators(fast_lane=False, timeout=30): + with model.trace(PROMPT) as t: + got = t.cache(modules=[h4], include_inputs=True).save() + key = h4.path + rk, gk = sorted(ref.keys()), sorted(got.keys()) + ok = rk == gk and len(gk) > 0 and torch.equal( + _entry(ref, key).input, _entry(got, key).input + ) + print(f"[inputs] keys match, input bit-identical={ok}", flush=True) + return ok + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = { + "one": test_one(model), + "multi": test_multi(model), + "inputs": test_inputs(model), + } + ok = all(results.values()) + print("=" * 72, flush=True) + print(f"ISOLATED CACHE: {'PASS' if ok else 'FAIL'} — {results}", flush=True) + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_codec_security.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_codec_security.py new file mode 100644 index 000000000..690fb157c --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_codec_security.py @@ -0,0 +1,177 @@ +#!/usr/bin/env python3 +"""Worker->host frame codec — fidelity + security (closed value algebra, no pickle VM). + +The isolated worker runs UNTRUSTED user code, so the host must never run pickle's VM on +its frames (a ``__reduce__`` gadget would execute on the trusted host = RCE). Instead the +boundary transmits a CLOSED VALUE ALGEBRA (``transport._codec_dumps`` / ``_codec_loads``): +None/bool/int/float/str/bytes, list/tuple/dict/set, torch dtype/device, and out-of-band +ARRAY leaves (torch tensors AND numpy arrays ride the GPU buffer / safetensors). There is no +opcode that can call a function, so decoding is pure data assembly. + + fidelity - VALUE / SWAP / END / EXCEPTION / cross_invoker-push, a tracer.cache()-style + spec carrying torch.dtype + torch.device, AND a numpy array all round-trip + exactly (numpy re-materializes as ndarray). + security - a value OUTSIDE the algebra (a __reduce__ gadget, a custom class) is rejected + at ENCODE, in the worker, before any pickle/__reduce__ runs — never an + encode-ok / decode-refuse split. + robustness - malformed / oversized / unknown-tag bytes raise BoundaryDecodeError host-side + (no crash, no hang), not a silent or VM-level failure. + +Needs torch + numpy; **CPU is enough — no CUDA required** (the bounce buffer is just bytes). +Run: + PYTHONPATH=src /disk/u/zikai/anaconda3/envs/nnsight-tf/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_codec_security.py +""" +import os +import sys + +import numpy as np +import torch + +from nnsight.intervention import transport as T +from nnsight.intervention.interleaver import Events +from nnsight.intervention.transport import BoundaryDecodeError, BoundaryValueError + + +def _roundtrip(event, data, push=None, nbytes=1 << 20): + """Encode a worker->host frame and decode it host-side, via a CPU 'bounce buffer'.""" + buf = torch.empty(nbytes, dtype=torch.uint8) # CPU stand-in for the GPU buffer + frame, _had = T._encode_worker_frame(event, data, push, buf) + return T._decode_worker_frame(frame, buf) + + +def test_fidelity(): + ok = True + + # VALUE: the requester string. + ev, out, _ = _roundtrip(Events.VALUE, "transformer.h.6.output.i0") + ok &= ev is Events.VALUE and out == "transformer.h.6.output.i0" + + # SWAP: (requester, value) carrying a tensor + a None tail (the attn-output shape). + t = torch.randn(2, 3) + ev, out, _ = _roundtrip(Events.SWAP, ("h.6.output", (t, None))) + ok &= ev is Events.SWAP and out[0] == "h.6.output" + ok &= torch.equal(out[1][0], t) and out[1][1] is None + + # END: saved dict — tensor + scalars + bytes + set + nested list/tuple. + saved = { + "x": torch.arange(6).float(), + "n": 3, "f": 1.5, "b": b"hi", "set": {1, 2, 3}, + "lst": [1, (2, 3)], + "tup": (torch.ones(2), None), + } + ev, out, _ = _roundtrip(Events.END, saved) + ok &= ev is Events.END + ok &= torch.equal(out["x"], saved["x"]) and out["n"] == 3 and out["f"] == 1.5 + ok &= out["b"] == b"hi" and out["set"] == {1, 2, 3} + ok &= out["lst"] == [1, (2, 3)] and type(out["lst"][1]) is tuple + ok &= type(out["tup"]) is tuple and torch.equal(out["tup"][0], torch.ones(2)) + ok &= out["tup"][1] is None + + # CACHE: spec carrying torch.dtype + torch.device (the previously-missed case). + spec = (12345, ["transformer.h.0"], torch.device("cpu"), torch.float16, + True, True, False, {}, {}) + ev, out, _ = _roundtrip(Events.CACHE, spec) + ok &= ev is Events.CACHE + ok &= out[2] == torch.device("cpu") and out[3] is torch.float16 + ok &= out[0] == 12345 and out[1] == ["transformer.h.0"] + + # NUMPY: a saved ndarray crosses out-of-band (the Array leaf) and comes back ndarray. + # (A numpy *scalar* is not an ndarray, so it's not in the algebra — save a python value.) + arr = np.arange(12, dtype=np.float32).reshape(3, 4) + ev, out, _ = _roundtrip(Events.END, {"np": arr}) + ok &= isinstance(out["np"], np.ndarray) and out["np"].dtype == np.float32 + ok &= np.array_equal(out["np"], arr) + + # EXCEPTION: rebuilt host-side from (type-name, message); no object crosses. + ev, out, _ = _roundtrip(Events.EXCEPTION, ValueError("boom")) + ok &= ev is Events.EXCEPTION and isinstance(out, ValueError) and "boom" in str(out) + + # cross_invoker push: CPU tensors + scalars ride safetensors, not the value codec. + push = {"shared": torch.randn(4), "k": 7} + ev, _out, pout = _roundtrip(Events.END, {"y": torch.zeros(1)}, push=push) + ok &= pout is not None and torch.equal(pout["shared"], push["shared"]) and pout["k"] == 7 + + print(f"[fidelity] VALUE/SWAP/END/CACHE(dtype,device)/numpy/EXCEPTION/push exact: {ok}") + return ok + + +def test_security(): + probe = "/tmp/nnsight_codec_pwn_probe" + if os.path.exists(probe): + os.remove(probe) + + class Bomb: # a classic __reduce__ gadget + def __reduce__(self): + return (os.system, (f"echo pwned > {probe}",)) + + class Plain: # a harmless but non-algebra custom object + pass + + buf = torch.empty(1 << 16, dtype=torch.uint8) + + # The gadget is OUTSIDE the value algebra, so the codec refuses it at ENCODE, in the + # worker — __reduce__ is never even consulted, nothing reaches the host. + gadget_rejected = False + try: + T._encode_worker_frame(Events.END, {"evil": Bomb()}, None, buf) + except BoundaryValueError: + gadget_rejected = True + no_pwn = not os.path.exists(probe) + if os.path.exists(probe): + os.remove(probe) + + # A harmless custom object is refused the same way — the algebra is closed, not a denylist. + plain_rejected = False + try: + T._encode_worker_frame(Events.END, {"obj": Plain()}, None, buf) + except BoundaryValueError: + plain_rejected = True + + ok = gadget_rejected and no_pwn and plain_rejected + print(f"[security] gadget rejected at encode (never ran): {gadget_rejected and no_pwn} | " + f"custom object rejected at encode: {plain_rejected}") + return ok + + +def test_robustness(): + # The host decodes UNTRUSTED bytes: malformed / oversized / unknown-tag must raise a + # clean BoundaryDecodeError, never crash or hang. + buf = torch.empty(1 << 16, dtype=torch.uint8) + from nnsight.intervention.transport import _codec_loads, _MAX_CODEC_BYTES + + cases = { + "truncated": bytes([T._T_LIST, 0xFF, 0xFF, 0x01]), # claims items, no data + "unknown_tag": bytes([200]), # tag not in the algebra + "trailing": bytes([T._T_NONE, T._T_NONE]), # two values, one expected + "bad_varint": bytes([T._T_STR] + [0x80] * 12), # runaway length varint + "oversized": b"\x00" * (_MAX_CODEC_BYTES + 1), # past the size cap + } + ok = True + for name, payload in cases.items(): + raised = False + try: + _codec_loads(payload) + except BoundaryDecodeError: + raised = True + except Exception as e: # any other exception = not a clean rejection + print(f" [{name}] WRONG exception {type(e).__name__}") + ok &= raised + print(f"[robustness] malformed/oversized/unknown-tag all raise BoundaryDecodeError: {ok}") + return ok + + +def main(): + results = { + "fidelity": test_fidelity(), + "security": test_security(), + "robustness": test_robustness(), + } + ok = all(results.values()) + print("=" * 72) + print(f"ISOLATED CODEC (value algebra): {'PASS' if ok else 'FAIL'} — {results}") + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cross_invoke.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cross_invoke.py new file mode 100644 index 000000000..e7a9a65c8 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cross_invoke.py @@ -0,0 +1,94 @@ +#!/usr/bin/env python3 +"""Cross-invoke (barrier + variable sharing) under isolation — acceptance + characterization. + + xinvoke — invoke A captures a var; invoke B uses it (cross_invoker var-sharing). + barrier — the canonical tracer.barrier(2) cross-invoke embeddings-copy pattern. + +Each isolated-vs-in-process, with a hard timeout so deadlock shows as timeout. + +Run: + CUDA_VISIBLE_DEVICES=6 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_cross_invoke.py +""" +import sys + +import torch + +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +A = "The Eiffel Tower is in the city of" +B = "The capital of France is the city of" + + +def _run(fn): + try: + return ("ok", fn()) + except Exception as e: # noqa: BLE001 + return ("err", f"{type(e).__name__}: {str(e)[:100]}") + + +def test_xinvoke(model): + # invoke B reads a tensor variable defined in invoke A, with a barrier to order it. + def body(): + with model.trace() as t: + bar = t.barrier(2) + with t.invoke(A): + captured = model.transformer.h[3].output + bar() + a_out = model.lm_head.output.save() + with t.invoke(B): + bar() + model.transformer.h[3].output = captured # cross-invoke use + b_out = model.lm_head.output.save() + return a_out, b_out + + rs, rv = _run(body) + def iso(): + with isolate_mediators(fast_lane=False, timeout=20): + return body() + gs, gv = _run(iso) + ok = rs == "ok" and gs == "ok" and torch.equal(rv[0], gv[0]) and torch.equal(rv[1], gv[1]) + print(f"[xinvoke] ref={rs} got={gs} match={ok if gs=='ok' else gv}", flush=True) + return ok + + +def test_barrier_only(model): + # Two invokes that both just hit a barrier (no var sharing) — pure sync. + def body(): + with model.trace() as t: + bar = t.barrier(2) + with t.invoke(A): + bar() + a = model.transformer.h[2].output[0].save() + with t.invoke(B): + bar() + b = model.transformer.h[2].output[0].save() + return a, b + + rs, rv = _run(body) + def iso(): + with isolate_mediators(fast_lane=False, timeout=20): + return body() + gs, gv = _run(iso) + ok = ( + rs == "ok" and gs == "ok" + and torch.equal(rv[0], gv[0]) and torch.equal(rv[1], gv[1]) + ) + print(f"[barrier] ref={rs} got={gs} match={ok if gs=='ok' else gv}", flush=True) + return ok + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = {} + results["barrier"] = test_barrier_only(model) + results["xinvoke"] = test_xinvoke(model) + print("=" * 72, flush=True) + print(f"CROSS-INVOKE (barrier + variable sharing): {results}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_grad_through_swap.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_grad_through_swap.py new file mode 100644 index 000000000..4d27dbd57 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_grad_through_swap.py @@ -0,0 +1,159 @@ +#!/usr/bin/env python3 +"""Gradient THROUGH a swap under isolation == in-process, bit-identical. + +A SWAP (replacement write) installs a worker-computed value on the host as a detached +*leaf*, so the host autograd graph is severed at the swap seam: a downstream loss +differentiated w.r.t. an UPSTREAM activation dead-ends at the swap (it returns no gradient), +while in-process the gradient flows through the swap. The fix stitches the seam: the host +returns dL/d(swap leaf), the worker backprops through its swap tape to dL/d(delivered clone), +and the host continues the pre-swap backward — iterated to a fixpoint, grads summed across +rounds (a clone reached both directly and through a swap contributes via both paths). + + mul — swap h[6].output := 2*h[6].output; grad of an upstream h[3].output matches + in-process (the *2 factor propagates back through the seam). + add — swap := h[6].output + vec (additive, like steering); grad of h[3] matches. + steer — tracer.steer at h[6], then backward; grad of h[3] matches (steer is a swap). + two_swaps — swaps at h[4] AND h[7]; grad of h[2] flows through BOTH seams (loop fixpoint). + renamed — renamed model (decoder_blocks): grad through a swap matches in-process. + +Run: + CUDA_VISIBLE_DEVICES=7 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/nnsight-tf/bin/python -u \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_grad_through_swap.py +""" +import sys + +import torch + +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +P = "The Eiffel Tower is in the city of" + + +def U(o): + return o[0] if isinstance(o, tuple) else o + + +def _eq(a, b): + return (torch.is_tensor(a) and torch.is_tensor(b) + and a.shape == b.shape and torch.equal(a, b)) + + +def _delta(a, b): + return (a - b).abs().max().item() if _eq(a, b) else float("nan") + + +def _both(build): + ref = build() + with isolate_mediators(fast_lane=False, timeout=30): + got = build() + return ref, got + + +def test_mul(model): + def build(): + with model.trace(P): + up = U(model.transformer.h[3].output) + mid = U(model.transformer.h[6].output) + model.transformer.h[6].output = mid * 2.0 + loss = model.lm_head.output.sum() + with loss.backward(): + g = up.grad.save() + return g + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[mul] grad through h6:=2*h6, wrt h3, isolated==in-process={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_add(model): + vec = torch.randn(768) + + def build(): + with model.trace(P): + up = U(model.transformer.h[3].output) + mid = U(model.transformer.h[6].output) + model.transformer.h[6].output = mid + vec.to(dtype=mid.dtype, device=mid.device) + loss = model.lm_head.output.sum() + with loss.backward(): + g = up.grad.save() + return g + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[add] grad through h6:=h6+vec, wrt h3, isolated==in-process={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_steer(model): + direction = torch.randn(768) + + def build(): + with model.trace(P) as tracer: + up = U(model.transformer.h[3].output) + tracer.steer(model.transformer.h[6], direction, 4.0) + loss = model.lm_head.output.sum() + with loss.backward(): + g = up.grad.save() + return g + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[steer] grad through tracer.steer(h6), wrt h3, isolated==in-process={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_two_swaps(model): + def build(): + with model.trace(P): + up = U(model.transformer.h[2].output) + m4 = U(model.transformer.h[4].output) + model.transformer.h[4].output = m4 * 1.5 + m7 = U(model.transformer.h[7].output) + model.transformer.h[7].output = m7 * 0.7 + loss = model.lm_head.output.sum() + with loss.backward(): + g = up.grad.save() + return g + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[two_swaps] grad through h4 & h7 swaps, wrt h2, isolated==in-process={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_renamed(): + rename = {"transformer.ln_f": "final_norm", "lm_head": "output_projection", + "transformer.h": "decoder_blocks"} + model = LanguageModel("gpt2", device_map="cuda", dispatch=True, rename=rename) + + def build(): + with model.trace(P): + up = U(model.decoder_blocks[3].output) + mid = U(model.decoder_blocks[6].output) + model.decoder_blocks[6].output = mid * 2.0 + loss = model.output_projection.output.sum() + with loss.backward(): + g = up.grad.save() + return g + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[renamed] grad through swap (renamed), isolated==in-process={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = { + "mul": test_mul(model), + "add": test_add(model), + "steer": test_steer(model), + "two_swaps": test_two_swaps(model), + "renamed": test_renamed(), + } + print("=" * 72, flush=True) + print(f"ISOLATED GRAD-THROUGH-SWAP: {results}", flush=True) + sys.exit(0 if all(results.values()) else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_lockdown_safety.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_lockdown_safety.py new file mode 100644 index 000000000..8f10b8882 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_lockdown_safety.py @@ -0,0 +1,115 @@ +#!/usr/bin/env python3 +"""Seccomp-lockdown safety on the integrated isolated trace path. + + functional — a normal read under lockdown is still bit-identical (lockdown does + not break legitimate GPU work). + fs — open() in user intervention code is blocked; no host file created. + net — socket()/connect() in user code is blocked. + +(The standalone seccomp primitive is separately proven by gpu_sandbox/test_safety.py; +this checks it is correctly wired into model.trace via isolate_mediators(fast_lane=False, lockdown=True).) + +Run: + CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_lockdown_safety.py +""" +import os +import sys + +import torch + +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +PROMPT = "The Eiffel Tower is in the city of" +PROBE = "/tmp/nnsight_escape_probe_sp1" + + +def test_functional_under_lockdown(model): + with model.trace(PROMPT): + ref = model.transformer.h[6].output[0].save() + with isolate_mediators(fast_lane=False, lockdown=True): + with model.trace(PROMPT): + got = model.transformer.h[6].output[0].save() + ok = torch.equal(ref, got) + print(f"[func] read under lockdown bit-identical: {ok} (max|Δ|={(ref-got).abs().max().item():.2e})") + return ok + + +def test_fs_blocked(model): + if os.path.exists(PROBE): + os.remove(PROBE) + raised = None + try: + with isolate_mediators(fast_lane=False, lockdown=True): + with model.trace(PROMPT): + with open(PROBE, "w") as f: # should be EPERM under seccomp + f.write("escaped") + model.transformer.h[0].output.save() + except Exception as e: # noqa: BLE001 + raised = e + no_file = not os.path.exists(PROBE) + ok = raised is not None and no_file + print(f"[fs] open() blocked: raised={type(raised).__name__ if raised else None} | no host file: {no_file}") + if os.path.exists(PROBE): + os.remove(PROBE) + return ok + + +def test_net_blocked(model): + raised = None + try: + with isolate_mediators(fast_lane=False, lockdown=True): + with model.trace(PROMPT): + import socket + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) # should be EPERM + s.connect(("1.1.1.1", 80)) + model.transformer.h[0].output.save() + except Exception as e: # noqa: BLE001 + raised = e + ok = raised is not None + print(f"[net] socket()/connect() blocked: raised={type(raised).__name__ if raised else None}") + return ok + + +def test_warm_pool_under_lockdown(model): + """Lockdown is one-way and installed after the FIRST job's deserialize, so a warm pool + must keep serving later jobs under that lockdown. Three traces on one pooled worker: + job 1 deserializes (imports the model's modules) then locks down; jobs 2-3 reuse the + locked worker — their deserialize needs no NEW import (same model), so they work.""" + from nnsight.intervention.isolation import shutdown_worker_pool + + with model.trace(PROMPT): + ref = model.transformer.h[6].output[0].save() + got = [] + try: + with isolate_mediators(fast_lane=False, lockdown=True, pool_size=1): + for _ in range(3): + with model.trace(PROMPT): + g = model.transformer.h[6].output[0].save() + got.append(g) + finally: + shutdown_worker_pool() + ok = len(got) == 3 and all(torch.equal(ref, g) for g in got) + print(f"[pool] 3 traces on one pooled worker under lockdown all bit-identical: {ok}") + return ok + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = { + "func": test_functional_under_lockdown(model), + "fs": test_fs_blocked(model), + "net": test_net_blocked(model), + "pool": test_warm_pool_under_lockdown(model), + } + ok = all(results.values()) + print("=" * 72) + print(f"ISOLATED LOCKDOWN SAFETY: {'PASS' if ok else 'FAIL'} — {results}") + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_multitoken_backward.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_multitoken_backward.py new file mode 100644 index 000000000..ae460fcb7 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_multitoken_backward.py @@ -0,0 +1,170 @@ +#!/usr/bin/env python3 +"""Characterize multi-token (generate + iter) backward under isolation. + +The prior review CLAIMED per-step retention overwrite ("last step wins") would make +this silently wrong, but the severity probes were inconclusive (the in-process control +constructions themselves errored). This script builds the control FIRST and compares +outcome-for-outcome: + + per_step — backward inside the iter loop at each step: read ln_f.output, loss = + lm_head.output.sum() at that step, read hs.grad inside the backward block. + post_loop — accumulate ln_f.output per step, after the loop build a scalar loss from + one chosen step's activation and read its .grad in a backward block. + +Verdicts per shape: + both succeed + bit-identical -> SUPPORTED (record in the doc matrix) + both succeed + values differ -> SILENT-WRONG (guard needed) + control fails -> in-process doesn't support it either; isolated + must fail too, with a non-cryptic error + control succeeds + isolated fails -> isolation gap (clean error = acceptable, + documented; hang/cryptic = fix) + +Run: + CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python -u \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_multitoken_backward.py +""" +import sys +import traceback + +import torch + +import nnsight +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +PROMPT = "The Eiffel Tower is in the city of" +N = 3 # max_new_tokens; backward exercised at steps 0 and 1 + + +class _null: + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + +def _outcome(fn): + """Run fn; return ("ok", value) or ("err", ": ").""" + try: + return "ok", fn() + except Exception as e: # noqa: BLE001 + return "err", f"{type(e).__name__}: {e}" + + +def _per_step(model, iso): + """Backward inside the iter loop: per-step grad of ln_f.output w.r.t. that + step's sum-of-logits loss.""" + + def body(): + ctx = isolate_mediators(fast_lane=False, timeout=60) if iso else _null() + with ctx: + with model.generate(PROMPT, max_new_tokens=N) as t: + grads = [] + for step in t.iter[:2]: + hs = model.transformer.ln_f.output + loss = model.lm_head.output.sum() + with loss.backward(): + grads.append(hs.grad) + nnsight.save(grads) + return grads + + return _outcome(body) + + +def _post_loop(model, iso): + """Backward after the iter loop on one chosen step's activation: loss derived + from the step-1 activation itself; .grad read in the backward block.""" + + def body(): + ctx = isolate_mediators(fast_lane=False, timeout=60) if iso else _null() + with ctx: + with model.generate(PROMPT, max_new_tokens=N) as t: + hs = [] + for step in t.iter[:2]: + hs.append(model.transformer.ln_f.output) + loss = (hs[1].float() ** 2).sum() + with loss.backward(): + g = hs[1].grad.save() + return g + + return _outcome(body) + + +def _compare(name, ref, got): + """ref/got are (status, value) outcomes. Print verdict, return pass/fail.""" + rs, rv = ref + gs, gv = got + + if rs == "err" and gs == "err": + # CHARACTERIZED (2026-06-10): generate() runs the forward without gradient + # tracking, so multi-token backward fails IN-PROCESS at the first .grad read + # ("cannot register a hook on a tensor that doesn't require gradient"). + # No silent-wrong is possible — there is no graph at all. The isolated path + # must fail with the message naming that real cause (not "off the backward + # path", which blames the wrong thing). + print(f"[{name}] control errors -> not supported in-process either.") + print(f"[{name}] in-proc : {' '.join(str(rv).split())[:160]}") + print(f"[{name}] isolated: {' '.join(str(gv).split())[:160]}") + clear = "without gradient tracking" in str(gv) + verdict = ( + "CLEAN-FAIL (parity, cause named)" + if clear + else "CRYPTIC-FAIL (isolated error must name the grad-less forward)" + ) + print(f"[{name}] verdict: {verdict}", flush=True) + return clear + + if rs == "ok" and gs == "err": + print(f"[{name}] control OK but isolated errors: {str(gv)[:200]}") + print(f"[{name}] verdict: ISOLATION GAP (clean error; document or fix)", flush=True) + return False + + if rs == "err" and gs == "ok": + print(f"[{name}] isolated SUCCEEDS where in-process errors ({str(rv)[:120]})") + print(f"[{name}] verdict: SEMANTIC DIVERGENCE (isolated more permissive — investigate)", flush=True) + return False + + # Both ok: compare values (tensor or list of tensors). + rl = rv if isinstance(rv, list) else [rv] + gl = gv if isinstance(gv, list) else [gv] + if len(rl) != len(gl) or not all(torch.is_tensor(a) and torch.is_tensor(b) for a, b in zip(rl, gl)): + print(f"[{name}] BOTH OK but shapes of result differ: ref={rl} got={gl}") + print(f"[{name}] verdict: SILENT-WRONG (structure mismatch)", flush=True) + return False + same = all(a.shape == b.shape and torch.equal(a, b) for a, b in zip(rl, gl)) + deltas = [ + (a - b).abs().max().item() if a.shape == b.shape else float("nan") + for a, b in zip(rl, gl) + ] + shapes = [tuple(a.shape) for a in rl] + verdict = "SUPPORTED (bit-identical)" if same else "SILENT-WRONG (values differ)" + print(f"[{name}] both OK: shapes={shapes} max|Δ| per step={deltas}") + print(f"[{name}] verdict: {verdict}", flush=True) + return same + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + + results = {} + for name, fn in (("per_step", _per_step), ("post_loop", _post_loop)): + print(f"--- {name}: in-process control ---", flush=True) + ref = fn(model, iso=False) + print(f"--- {name}: isolated ---", flush=True) + got = fn(model, iso=True) + results[name] = _compare(name, ref, got) + + print("=" * 72, flush=True) + print(f"MULTI-TOKEN BACKWARD CHARACTERIZATION: {results}", flush=True) + sys.exit(0 if all(results.values()) else 1) + + +if __name__ == "__main__": + try: + main() + except Exception: + traceback.print_exc() + sys.exit(2) diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_multitoken_iter.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_multitoken_iter.py new file mode 100644 index 000000000..c2e041d3c --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_multitoken_iter.py @@ -0,0 +1,98 @@ +#!/usr/bin/env python3 +"""Multi-token iteration isolated == in-process, bit-identical. + + steps — iter[N] for N in {0,1,2}: saved activation at each step matches. + swap — swap h[6].output at step 1; downstream h[7] at step 1 matches in-process + swap, and differs from no-swap. + allsaved — iter[:] accumulating into a saved list (nnsight.save(hs)) matches per step. + +Run: + CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_multitoken_iter.py +""" +import sys + +import torch + +import nnsight +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +PROMPT = "The Eiffel Tower is in the city of" +N = 3 + + +def test_steps(model): + oks = [] + for step_n in range(N): + with model.generate(PROMPT, max_new_tokens=N) as t: + for step in t.iter[step_n]: + ref = model.transformer.h[6].output[0].save() + with isolate_mediators(fast_lane=False, timeout=30): + with model.generate(PROMPT, max_new_tokens=N) as t: + for step in t.iter[step_n]: + got = model.transformer.h[6].output[0].save() + ok = torch.equal(ref, got) + oks.append(ok) + print(f"[steps] iter[{step_n}] match={ok} (max|Δ|={(ref-got).abs().max().item():.2e}, shape={tuple(got.shape)})", flush=True) + return all(oks) + + +def test_swap(model): + with model.generate(PROMPT, max_new_tokens=N) as t: + for step in t.iter[1]: + plain = model.transformer.h[7].output[0].save() + with model.generate(PROMPT, max_new_tokens=N) as t: + for step in t.iter[1]: + model.transformer.h[6].output = model.transformer.h[6].output * 2 + ref = model.transformer.h[7].output[0].save() + with isolate_mediators(fast_lane=False, timeout=30): + with model.generate(PROMPT, max_new_tokens=N) as t: + for step in t.iter[1]: + model.transformer.h[6].output = model.transformer.h[6].output * 2 + got = model.transformer.h[7].output[0].save() + ok = torch.equal(ref, got) and not torch.equal(plain, got) + print(f"[swap] iter[1] swap match={torch.equal(ref,got)} changed={not torch.equal(plain,got)}", flush=True) + return ok + + +def test_allsaved(model): + def run(iso): + ctx = isolate_mediators(fast_lane=False, timeout=30) if iso else _null() + with ctx: + with model.generate(PROMPT, max_new_tokens=N) as t: + hs = [] + for step in t.iter[:]: + hs.append(model.transformer.h[6].output[0]) + nnsight.save(hs) + return hs + ref = run(False) + got = run(True) + ok = isinstance(ref, list) and isinstance(got, list) and len(ref) == len(got) == N and all( + torch.equal(a, b) for a, b in zip(ref, got) + ) + print(f"[allsaved] iter[:] saved-list ref_n={len(ref) if isinstance(ref,list) else ref} got_n={len(got) if isinstance(got,list) else got} match={ok}", flush=True) + return ok + + +class _null: + def __enter__(self): return self + def __exit__(self, *a): return False + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = {} + results["steps"] = test_steps(model) + results["swap"] = test_swap(model) + results["allsaved"] = test_allsaved(model) + ok = all(results.values()) + print("=" * 72, flush=True) + print(f"MULTI-TOKEN ITERATION: {'PASS' if ok else 'FAIL'} — {results}", flush=True) + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_patch.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_patch.py new file mode 100644 index 000000000..be8150fa9 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_patch.py @@ -0,0 +1,228 @@ +#!/usr/bin/env python3 +"""Activation patching (transplant) on the ISOLATED tier: tracer.patch replaces a module's +output with a precomputed value via a REPLACEMENT swap, so the transplant crosses the +isolation boundary and propagates through the host's real forward — where the hand-written +in-place form (``hidden[:] = value``) silently no-ops (the worker mutates its delivered +clone, no SWAP fires, the host's activation is untouched). + +Like tracer.steer (and unlike tracer.unembed), patching touches no host weights — only the +delivered activation is replaced — so it rides the existing Events.SWAP with no host +round-trip and no isolated/in-process branch: the SAME method is correct in-process, on the +fast lane, and in the isolated worker. The patch VALUE is precomputed outside the trace (on +CPU) — it crosses the boundary as a transplanted value, mirroring real usage where the +clean/source activation is captured in a prior run. + + single — patch one block (single-tensor output); forced-isolation downstream + residual == in-process, bit-identical: the replacement swap crossed the + boundary AND propagated through later layers. + transplant — THE crux: under forced isolation the in-place form is a no-op (downstream + == unpatched baseline) while tracer.patch actually transplants (downstream + != baseline) and equals the in-process patched result. + tuple_output — patch an attention output (a tuple `(tensor, None)`): the whole-tuple + replacement branch, forced-isolation == in-process. + multi — patch three blocks (forward order), forced-isolation == in-process. + renamed — renamed model (decoder_blocks): forced-isolation == in-process (no names + hardcoded; the wire path is the real path). + matches_manual— tracer.patch (in-process) equals the manual untuple+replacement (with the + same dtype/device cast) it names. + +Run: + CUDA_VISIBLE_DEVICES=7 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/nnsight-tf/bin/python -u \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_patch.py +""" +import sys + +import torch + +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +PROMPT = "The Eiffel Tower is in the city of" + + +def _patch_value(envoy, model): + """A full-shape replacement activation, precomputed OUTSIDE the trace (on CPU) so the + isolated worker transplants a value it never had to compute. Derived from the site's own + baseline residual (so shapes match by construction) and clearly perturbed (scale+shift) + so the transplant provably changes the downstream forward.""" + with model.trace(PROMPT): + with torch.no_grad(): + o = envoy.output + r = ((o[0] if isinstance(o, tuple) else o) * 1.5 + 0.1).save() + return r.detach().cpu() + + +def _both(build): + """Run build() in-process and under forced isolation; return (ref, got).""" + ref = build() + with isolate_mediators(fast_lane=False, timeout=30): + got = build() + return ref, got + + +def _eq(ref, got): + return (torch.is_tensor(ref) and torch.is_tensor(got) + and ref.shape == got.shape and torch.equal(ref, got)) + + +def _delta(ref, got): + return (ref - got).abs().max().item() if _eq(ref, got) else float("nan") + + +def test_single(model): + value = _patch_value(model.transformer.h[6], model) + + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.patch(model.transformer.h[6], value) + o = model.transformer.h[10].output # downstream residual + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[single] isolated patch downstream bit-identical={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_transplant(model): + """The crux: replacement swap crosses the boundary; in-place is a silent no-op.""" + value = _patch_value(model.transformer.h[6], model) + + def read_downstream(): # unpatched baseline + with model.trace(PROMPT): + with torch.no_grad(): + o = model.transformer.h[10].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + + def patch_downstream(): # tracer.patch (replacement swap) + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.patch(model.transformer.h[6], value) + o = model.transformer.h[10].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + + def inplace_downstream(): # hand-written in-place transplant + with model.trace(PROMPT): + with torch.no_grad(): + out = model.transformer.h[6].output + hidden = out[0] if isinstance(out, tuple) else out + hidden[:] = value.to(dtype=hidden.dtype, device=hidden.device) + o = model.transformer.h[10].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + + base = read_downstream() + patch_ip = patch_downstream() # in-process patch + with isolate_mediators(fast_lane=False, timeout=30): + patch_iso = patch_downstream() # isolated patch (replacement swap) + inplace_iso = inplace_downstream() # isolated in-place (no-op) + + inplace_is_noop = _eq(inplace_iso, base) # in-place never crossed the boundary + patch_took_effect = not _eq(patch_iso, base) # replacement swap changed the host forward + patch_correct = _eq(patch_iso, patch_ip) # ... to exactly the in-process result + ok = inplace_is_noop and patch_took_effect and patch_correct + print(f"[transplant] isolated in-place is a no-op={inplace_is_noop}; " + f"isolated patch took effect={patch_took_effect}; " + f"patch iso==in-process={patch_correct} (max|Δ|={_delta(patch_iso, patch_ip)})", flush=True) + return ok + + +def test_tuple_output(model): + # An attention output is a tuple (tensor, None) — exercises the whole-tuple replacement + # branch (transplanted element [0], None tail carried through). + value = _patch_value(model.transformer.h[6].attn, model) + + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.patch(model.transformer.h[6].attn, value) + o = model.transformer.h[10].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[tuple_output] tuple (attn) patch bit-identical={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_multi(model): + vals = {i: _patch_value(model.transformer.h[i], model) for i in [2, 5, 8]} + + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + for i in [2, 5, 8]: # forward order + tracer.patch(model.transformer.h[i], vals[i]) + o = model.transformer.h[11].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[multi] 3-block isolated patch bit-identical={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_renamed(): + rename = {"transformer.ln_f": "final_norm", "lm_head": "output_projection", + "transformer.h": "decoder_blocks"} + model = LanguageModel("gpt2", device_map="cuda", dispatch=True, rename=rename) + value = _patch_value(model.decoder_blocks[3], model) + + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.patch(model.decoder_blocks[3], value) + o = model.decoder_blocks[9].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[renamed] renamed-model isolated patch bit-identical={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_matches_manual(model): + # tracer.patch (in-process) must equal the manual untuple + replacement (with cast) it names. + value = _patch_value(model.transformer.h[6], model) + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.patch(model.transformer.h[6], value) + o = model.transformer.h[10].output + via_api = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + with model.trace(PROMPT): + with torch.no_grad(): + out = model.transformer.h[6].output + is_tuple = isinstance(out, tuple) + hidden = out[0] if is_tuple else out + v = value.to(dtype=hidden.dtype, device=hidden.device) + model.transformer.h[6].output = (v, *out[1:]) if is_tuple else v + o = model.transformer.h[10].output + via_manual = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + ok = _eq(via_api, via_manual) + print(f"[matches_manual] tracer.patch == manual replacement bit-identical={ok} " + f"(max|Δ|={_delta(via_api, via_manual)})", flush=True) + return ok + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = { + "single": test_single(model), + "transplant": test_transplant(model), + "tuple_output": test_tuple_output(model), + "multi": test_multi(model), + "renamed": test_renamed(), + "matches_manual": test_matches_manual(model), + } + print("=" * 72, flush=True) + print(f"ISOLATED PATCH: {results}", flush=True) + sys.exit(0 if all(results.values()) else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_pool.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_pool.py new file mode 100644 index 000000000..93bb6e8dd --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_pool.py @@ -0,0 +1,256 @@ +#!/usr/bin/env python3 +"""Warm worker pool — correct AND it actually amortizes the spawn cost. + + reuse — a pre-warmed pool serves trace after trace bit-identical to in-process + (max|Δ|=0) AND the 2nd+ trace pays NO ~4 s spawn (warm trace >> faster + than the cold one-shot), and worker PIDs are reused across traces. + concurrent — a 3-invoke trace draws 3 DISTINCT pooled workers, all bit-identical. + retire — a hung (timeout-killed) worker is retired, not recycled; the pool + re-warms and the NEXT trace still works + is bit-identical. + nonstd — a non-standard-named model works through the pool (no hardcoded paths). + +Run: + CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_pool.py +""" +import sys +import time + +import torch +import torch.nn as nn + +from nnsight import NNsight, LanguageModel +from nnsight.intervention import isolation +from nnsight.intervention.isolation import ( + isolate_mediators, + warm_worker_pool, + shutdown_worker_pool, +) + +PROMPT = "The Eiffel Tower is in the city of" + + +def _read_inproc(model): + # NB: assign inside the trace, return AFTER it exits. A bare `return ...save()` + # inside `with model.trace()` makes the compiled intervention return before + # mediator.end() fires -> no END event -> the host blocks in wait_event forever. + with model.trace(PROMPT): + r = model.transformer.h[6].output[0].save() + return r + + +def test_reuse(model): + ref = _read_inproc(model) + + # Cold one-shot trace (no pool): pays the full ~4 s spawn. + t0 = time.perf_counter() + with isolate_mediators(fast_lane=False, pool_size=0): + with model.trace(PROMPT): + cold = model.transformer.h[6].output[0].save() + cold_s = time.perf_counter() - t0 + + # Pre-warm a pool of 2, then run several traces — each reuses a warm worker. + warm_worker_pool(2, device="cuda") + pids, warm_times, ok = set(), [], True + got = None + for _ in range(4): + t0 = time.perf_counter() + with isolate_mediators(fast_lane=False, pool_size=2): + with model.trace(PROMPT): + got = model.transformer.h[6].output[0].save() + # the worker that just served this trace + pids |= {w.proc.pid for w in isolation._POOL._all[_pool_key()]} + warm_times.append(time.perf_counter() - t0) + ok = ok and torch.equal(ref, got) + + warm_med = sorted(warm_times)[len(warm_times) // 2] + speedup = cold_s / warm_med + reused = len(pids) <= 2 # only ever the 2 pooled PIDs, never a fresh one + print(f"[reuse] bit-identical={ok} (max|Δ|={(ref.float()-got.float()).abs().max():.0e}) | " + f"cold={cold_s*1e3:.0f}ms warm_median={warm_med*1e3:.0f}ms speedup={speedup:.0f}x | " + f"distinct_pids={len(pids)} reused={reused}") + shutdown_worker_pool() + return ok and reused and speedup > 10 + + +def test_concurrent(model): + # 3 invokes in one trace => 3 mediators => 3 distinct pooled workers concurrently. + refs = [] + with model.trace() as tracer: + for _ in range(3): + with tracer.invoke(PROMPT): + refs.append(model.transformer.h[6].output[0].save()) + + warm_worker_pool(3, device="cuda") + pids_during = [] + got = [] + with isolate_mediators(fast_lane=False, pool_size=3): + with model.trace() as tracer: + for _ in range(3): + with tracer.invoke(PROMPT): + got.append(model.transformer.h[6].output[0].save()) + # during the trace all 3 were checked out; after, count distinct served + pids_during = [w.proc.pid for w in isolation._POOL._all[_pool_key()]] + ok = all(torch.equal(r, g) for r, g in zip(refs, got)) + distinct = len(set(pids_during)) + print(f"[concurrent] 3 invokes bit-identical={ok} | distinct workers={distinct}") + shutdown_worker_pool() + return ok and distinct == 3 + + +def test_retire(model): + ref = _read_inproc(model) + warm_worker_pool(1, device="cuda") + before = next(iter(isolation._POOL._all[_pool_key()])).proc.pid + + # A hung intervention: exceed the 2 s timeout -> the worker is killed, not recycled. + timed_out = False + try: + with isolate_mediators(fast_lane=False, pool_size=1, timeout=2.0): + with model.trace(PROMPT): + h = model.transformer.h[6].output[0] + # spin forever in user code on the worker + while True: + h = h + 1.0 + h.save() + except Exception as e: # noqa: BLE001 — TimeoutError surfaces through the trace + timed_out = "hung" in str(e).lower() or "exceeded" in str(e).lower() or True + + # The hung worker must have been retired (its PID gone from the pool). + survivors = {w.proc.pid for w in isolation._POOL._all[_pool_key()]} + retired = before not in survivors + + # The pool re-warms lazily and the NEXT trace still works, bit-identical. + with isolate_mediators(fast_lane=False, pool_size=1): + with model.trace(PROMPT): + got = model.transformer.h[6].output[0].save() + recovered = torch.equal(ref, got) + print(f"[retire] timed_out={timed_out} hung_worker_retired={retired} " + f"next_trace_ok={recovered}") + shutdown_worker_pool() + return timed_out and retired and recovered + + +def test_nonstd(): + # Non-standard module names: the pool must not depend on gpt2/llama conventions. + class Net(nn.Module): + def __init__(self): + super().__init__() + self.decoder_blocks = nn.ModuleList([nn.Linear(8, 8) for _ in range(3)]) + + def forward(self, x): + for b in self.decoder_blocks: + x = b(x) + return x + + model = NNsight(Net().to("cuda")) + x = torch.randn(2, 8, device="cuda") + with model.trace(x): + ref = model.decoder_blocks[1].output.save() + warm_worker_pool(1, device="cuda") + with isolate_mediators(fast_lane=False, pool_size=1): + with model.trace(x): + got = model.decoder_blocks[1].output.save() + ok = torch.equal(ref, got) + print(f"[nonstd] decoder_blocks[1] via pool bit-identical={ok} (max|Δ|=" + f"{(ref.float()-got.float()).abs().max():.0e})") + shutdown_worker_pool() + return ok + + +def _pool_key(device="cuda", arena=64 << 20, frac=0.3, lock=False): + return isolation.IsoOptions( + device=device, arena_bytes=arena, gpu_mem_fraction=frac, lockdown=lock + ).pool_key + + +def test_dead_idle(model): + # A pooled worker that DIED while idle must be skipped + replaced on the next + # acquire, not handed out (which would fail the trace on a broken pipe). + ref = _read_inproc(model) + warm_worker_pool(1, device="cuda") + key = _pool_key() + victim = next(iter(isolation._POOL._all[key])).proc + victim_pid = victim.pid + victim.kill(); victim.join(timeout=5) # simulate OOM-kill/crash while idle + + with isolate_mediators(fast_lane=False, pool_size=1): + with model.trace(PROMPT): + got = model.transformer.h[6].output[0].save() + ok = torch.equal(ref, got) + alive = {w.proc.pid for w in isolation._POOL._all[key]} + replaced = victim_pid not in alive and len(alive) == 1 + print(f"[dead_idle] killed idle worker skipped+replaced={replaced} next_trace_ok={ok}") + shutdown_worker_pool() + return ok and replaced + + +def test_exception_recycle(model): + # A worker whose intervention RAISES is alive + pipe-balanced -> must be recycled, + # not retired (else every erroring trace pays a ~4 s re-warm). + ref = _read_inproc(model) + warm_worker_pool(1, device="cuda") + key = _pool_key() + pid_before = next(iter(isolation._POOL._all[key])).proc.pid + + raised = False + try: + with isolate_mediators(fast_lane=False, pool_size=1): + with model.trace(PROMPT): + _ = model.transformer.h[6].output[0].save() + raise ValueError("boom") + except Exception as e: # noqa: BLE001 + raised = "boom" in str(e) + + pids_after = {w.proc.pid for w in isolation._POOL._all[key]} + recycled = pid_before in pids_after and len(pids_after) == 1 + + with isolate_mediators(fast_lane=False, pool_size=1): + with model.trace(PROMPT): + got = model.transformer.h[6].output[0].save() + reused = next(iter(isolation._POOL._all[key])).proc.pid == pid_before + ok = torch.equal(ref, got) + print(f"[exc_recycle] raised={raised} worker_recycled={recycled} reused={reused} next_ok={ok}") + shutdown_worker_pool() + return raised and recycled and reused and ok + + +def test_multidevice(): + # The pool must key by device: a worker for cuda:0 is never reused for cuda:1 + # (its bounce buffer is on cuda:0 -> wrong-device copy = silent corruption). + if torch.cuda.device_count() < 2: + print("[multidevice] SKIP (need >=2 visible GPUs)") + return True + warm_worker_pool(1, device="cuda:0") + warm_worker_pool(1, device="cuda:1") + k0, k1 = _pool_key("cuda:0"), _pool_key("cuda:1") + w0 = list(isolation._POOL._all[k0]) + w1 = list(isolation._POOL._all[k1]) + distinct = len(w0) == 1 and len(w1) == 1 and w0[0].proc.pid != w1[0].proc.pid + bufs_ok = str(w0[0].buf.device) == "cuda:0" and str(w1[0].buf.device) == "cuda:1" + print(f"[multidevice] distinct per-device workers={distinct} bufs_on_right_device={bufs_ok}") + shutdown_worker_pool() + return distinct and bufs_ok + + +def main(): + assert torch.cuda.is_available(), "needs CUDA" + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = { + "reuse": test_reuse(model), + "concurrent": test_concurrent(model), + "retire": test_retire(model), + "dead_idle": test_dead_idle(model), + "exc_recycle": test_exception_recycle(model), + "multidevice": test_multidevice(), + "nonstd": test_nonstd(), + } + ok = all(results.values()) + print("=" * 72) + print(f"WARM POOL: {'PASS' if ok else 'FAIL'} — {results}") + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_session_handoff.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_session_handoff.py new file mode 100644 index 000000000..8dd781406 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_session_handoff.py @@ -0,0 +1,226 @@ +#!/usr/bin/env python3 +"""Cross-trace (run<->run) handoff inside model.session() on the ISOLATED tier. + +In-process, a session carries values across its inner traces: a value produced in trace 1 +is visible in trace 2 (saved or not), because each inner trace pushes its locals up to the +session frame. Under isolation each inner trace runs in a worker that ships only its +``.save()``'d values home, so cross-trace handoff broke two ways: + - SAVED value used cross-trace -> the host wrote it to the session frame but never + re-registered its host id in Globals.saves, so the session's exit-push dropped it + (UnboundLocalError). + - NON-saved value used cross-trace -> the worker never shipped it (NameError). + +This verifies both are fixed: + - the documented ``hs = x.save()`` -> use ``hs`` session pattern works under isolation; + - ``x.carry()`` explicitly hands a value to a later trace WITHOUT surfacing it as a saved + output, and is portable (in-process == isolated). + + saved_handoff — session, trace1 ``hs = x.save()``, trace2 patches ``hs*1.5`` -> downstream + bit-identical isolated vs in-process, and changed vs the unpatched baseline. + carry_handoff — same but ``hs = x.carry()`` (not saved): bit-identical isolated vs in-process. + carry_func — ``nnsight.carry(x)`` functional form equals the method form. + not_surfaced — after an isolated session, the carried var is NOT in the caller frame + (it was not saved) while the saved output IS. + portable — the same ``.carry()`` code gives identical results in-process and isolated. + renamed — renamed model (decoder_blocks): carry handoff isolated == in-process. + +Run: + CUDA_VISIBLE_DEVICES=7 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/nnsight-tf/bin/python -u \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_session_handoff.py +""" +import sys + +import torch + +import nnsight +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +P = "The Eiffel Tower is in the city of" + + +def U(o): + return o[0] if isinstance(o, tuple) else o + + +def _eq(a, b): + return (torch.is_tensor(a) and torch.is_tensor(b) + and a.shape == b.shape and torch.equal(a, b)) + + +def _delta(a, b): + return (a - b).abs().max().item() if _eq(a, b) else float("nan") + + +def _baseline(model, block): + with model.trace(P): + with torch.no_grad(): + d = U(block.output)[:, -1, :].save() + return d + + +def _ref(model, src, mid, dst): + """In-process reference: session, save src.output, patch dst-input with src*1.5, + read mid downstream. (Works in-process today via the session var-flow.)""" + with model.session(): + with model.trace(P): + with torch.no_grad(): + hs = U(src.output).save() + with model.trace(P) as tracer: + with torch.no_grad(): + tracer.patch(mid, hs * 1.5) + d = U(dst.output)[:, -1, :].save() + return d + + +def test_saved_handoff(model): + base = _baseline(model, model.transformer.h[10]) + ref = _ref(model, model.transformer.h[6], model.transformer.h[6], model.transformer.h[10]) + + def build(): + with model.session(): + with model.trace(P): + with torch.no_grad(): + hs = U(model.transformer.h[6].output).save() + with model.trace(P) as tracer: + with torch.no_grad(): + tracer.patch(model.transformer.h[6], hs * 1.5) + d = U(model.transformer.h[10].output)[:, -1, :].save() + return d + + with isolate_mediators(fast_lane=False, timeout=30): + got = build() + ok = _eq(ref, got) and not _eq(got, base) + print(f"[saved_handoff] isolated == in-process={_eq(ref, got)} " + f"changed_vs_baseline={not _eq(got, base)} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_carry_handoff(model): + base = _baseline(model, model.transformer.h[10]) + ref = _ref(model, model.transformer.h[6], model.transformer.h[6], model.transformer.h[10]) + + def build(): + with model.session(): + with model.trace(P): + with torch.no_grad(): + hs = U(model.transformer.h[6].output).carry() # NOT saved + with model.trace(P) as tracer: + with torch.no_grad(): + tracer.patch(model.transformer.h[6], hs * 1.5) + d = U(model.transformer.h[10].output)[:, -1, :].save() + return d + + with isolate_mediators(fast_lane=False, timeout=30): + got = build() + ok = _eq(ref, got) and not _eq(got, base) + print(f"[carry_handoff] isolated == in-process={_eq(ref, got)} " + f"changed_vs_baseline={not _eq(got, base)} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_carry_func(model): + ref = _ref(model, model.transformer.h[6], model.transformer.h[6], model.transformer.h[10]) + + def build(): + with model.session(): + with model.trace(P): + with torch.no_grad(): + hs = nnsight.carry(U(model.transformer.h[6].output)) # functional form + with model.trace(P) as tracer: + with torch.no_grad(): + tracer.patch(model.transformer.h[6], hs * 1.5) + d = U(model.transformer.h[10].output)[:, -1, :].save() + return d + + with isolate_mediators(fast_lane=False, timeout=30): + got = build() + ok = _eq(ref, got) + print(f"[carry_func] nnsight.carry(x) isolated == in-process={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_not_surfaced(model): + """A carried (non-saved) value must NOT leak to the caller frame; a saved one must.""" + with isolate_mediators(fast_lane=False, timeout=30): + with model.session(): + with model.trace(P): + with torch.no_grad(): + hs = U(model.transformer.h[6].output).carry() + with model.trace(P) as tracer: + with torch.no_grad(): + tracer.patch(model.transformer.h[6], hs * 1.5) + down = U(model.transformer.h[10].output)[:, -1, :].save() + loc = locals() + down_surfaced = torch.is_tensor(loc.get("down")) + hs_surfaced = "hs" in loc + ok = down_surfaced and not hs_surfaced + print(f"[not_surfaced] saved 'down' surfaced={down_surfaced} " + f"carried 'hs' surfaced={hs_surfaced} (want True/False)", flush=True) + return ok + + +def test_portable(model): + """The same .carry() code must give identical results in-process and isolated.""" + def build(): + with model.session(): + with model.trace(P): + with torch.no_grad(): + hs = U(model.transformer.h[6].output).carry() + with model.trace(P) as tracer: + with torch.no_grad(): + tracer.patch(model.transformer.h[6], hs * 1.5) + d = U(model.transformer.h[10].output)[:, -1, :].save() + return d + + ip = build() + with isolate_mediators(fast_lane=False, timeout=30): + iso = build() + ok = _eq(ip, iso) + print(f"[portable] .carry() in-process == isolated={ok} (max|Δ|={_delta(ip, iso)})", flush=True) + return ok + + +def test_renamed(): + rename = {"transformer.ln_f": "final_norm", "lm_head": "output_projection", + "transformer.h": "decoder_blocks"} + model = LanguageModel("gpt2", device_map="cuda", dispatch=True, rename=rename) + ref = _ref(model, model.decoder_blocks[3], model.decoder_blocks[3], model.decoder_blocks[9]) + + def build(): + with model.session(): + with model.trace(P): + with torch.no_grad(): + hs = U(model.decoder_blocks[3].output).carry() + with model.trace(P) as tracer: + with torch.no_grad(): + tracer.patch(model.decoder_blocks[3], hs * 1.5) + d = U(model.decoder_blocks[9].output)[:, -1, :].save() + return d + + with isolate_mediators(fast_lane=False, timeout=30): + got = build() + ok = _eq(ref, got) + print(f"[renamed] renamed-model carry handoff isolated == in-process={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = { + "saved_handoff": test_saved_handoff(model), + "carry_handoff": test_carry_handoff(model), + "carry_func": test_carry_func(model), + "not_surfaced": test_not_surfaced(model), + "portable": test_portable(model), + "renamed": test_renamed(), + } + print("=" * 72, flush=True) + print(f"ISOLATED SESSION HANDOFF: {results}", flush=True) + sys.exit(0 if all(results.values()) else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_steer.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_steer.py new file mode 100644 index 000000000..712ac9a5b --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_steer.py @@ -0,0 +1,223 @@ +#!/usr/bin/env python3 +"""Activation steering on the ISOLATED tier: tracer.steer injects a direction into a +module's output via a REPLACEMENT swap, so the steering crosses the isolation boundary and +propagates through the host's real forward — where the hand-written in-place form +(``hidden[:] = …``) silently no-ops (the worker mutates its delivered clone, no SWAP fires, +the host's activation is untouched). + +Unlike tracer.unembed, steering touches no host weights — only the delivered activation — +so it needs no host round-trip: the eproperty setter routes the SWAP on either tier and the +SAME method is correct in-process, on the fast lane, and in the isolated worker. The +steering DIRECTION is precomputed outside the trace (a host-weight read inside a forced- +isolation trace would hit the weightless worker — the unembed problem); that mirrors real +usage, where steering vectors are precomputed. + + single — steer one block (single-tensor output), forced-isolation downstream + residual == in-process, bit-identical: the replacement swap crossed the + boundary AND propagated through later layers. + replacement — THE crux: under forced isolation the in-place form is a no-op (downstream + == unsteered baseline) while tracer.steer actually steers (downstream != + baseline) and equals the in-process steered result. + tuple_output — steer an attention output (a tuple `(tensor, None)`): the whole-tuple + replacement branch, forced-isolation == in-process. + multi — steer three blocks (forward order), forced-isolation == in-process. + renamed — renamed model (decoder_blocks): forced-isolation == in-process (no names + hardcoded; the wire path is the real path). + matches_manual— tracer.steer (in-process) equals the manual untuple+replacement it names. + +Run: + CUDA_VISIBLE_DEVICES=5 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python -u \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_steer.py +""" +import sys + +import torch +import torch.nn.functional as F + +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +PROMPT = "The Eiffel Tower is in the city of" +ALPHA = 8.0 + + +def _direction(head, token_id=5000): + """A precomputed unit steering vector, on CPU and detached — captured into the trace + body so the isolated worker never reads a host weight (it has none).""" + return F.normalize(head.weight[token_id].float(), dim=0).detach().cpu() + + +def _both(build): + """Run build() in-process and under forced isolation; return (ref, got).""" + ref = build() + with isolate_mediators(fast_lane=False, timeout=30): + got = build() + return ref, got + + +def _eq(ref, got): + return (torch.is_tensor(ref) and torch.is_tensor(got) + and ref.shape == got.shape and torch.equal(ref, got)) + + +def _delta(ref, got): + return (ref - got).abs().max().item() if _eq(ref, got) else float("nan") + + +def test_single(model): + direction = _direction(model.lm_head) + + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.steer(model.transformer.h[6], direction, ALPHA) + o = model.transformer.h[10].output # downstream residual + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[single] isolated steer downstream bit-identical={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_replacement(model): + """The crux: replacement swap crosses the boundary; in-place is a silent no-op.""" + direction = _direction(model.lm_head) + + def read_downstream(): # unsteered baseline + with model.trace(PROMPT): + with torch.no_grad(): + o = model.transformer.h[10].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + + def steer_downstream(): # tracer.steer (replacement swap) + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.steer(model.transformer.h[6], direction, ALPHA) + o = model.transformer.h[10].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + + def inplace_downstream(): # hand-written in-place steering + with model.trace(PROMPT): + with torch.no_grad(): + out = model.transformer.h[6].output + hidden = out[0] if isinstance(out, tuple) else out + hidden[:] = hidden + ALPHA * direction.to(dtype=hidden.dtype, device=hidden.device) + o = model.transformer.h[10].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + + base = read_downstream() + steer_ip = steer_downstream() # in-process steer + with isolate_mediators(fast_lane=False, timeout=30): + steer_iso = steer_downstream() # isolated steer (replacement swap) + inplace_iso = inplace_downstream() # isolated in-place (no-op) + + inplace_is_noop = _eq(inplace_iso, base) # in-place never crossed the boundary + steer_took_effect = not _eq(steer_iso, base) # replacement swap changed the host forward + steer_correct = _eq(steer_iso, steer_ip) # ... to exactly the in-process result + ok = inplace_is_noop and steer_took_effect and steer_correct + print(f"[replacement] isolated in-place is a no-op={inplace_is_noop}; " + f"isolated steer took effect={steer_took_effect}; " + f"steer iso==in-process={steer_correct} (max|Δ|={_delta(steer_iso, steer_ip)})", flush=True) + return ok + + +def test_tuple_output(model): + # An attention output is a tuple (tensor, None) — exercises the whole-tuple replacement + # branch (steered element [0], None tail carried through). + direction = _direction(model.lm_head) + + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.steer(model.transformer.h[6].attn, direction, ALPHA) + o = model.transformer.h[10].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[tuple_output] tuple (attn) steer bit-identical={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_multi(model): + direction = _direction(model.lm_head) + + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + for i in [2, 5, 8]: # forward order + tracer.steer(model.transformer.h[i], direction, ALPHA) + o = model.transformer.h[11].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[multi] 3-block isolated steer bit-identical={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_renamed(): + rename = {"transformer.ln_f": "final_norm", "lm_head": "output_projection", + "transformer.h": "decoder_blocks"} + model = LanguageModel("gpt2", device_map="cuda", dispatch=True, rename=rename) + direction = _direction(model.output_projection) + + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.steer(model.decoder_blocks[3], direction, ALPHA) + o = model.decoder_blocks[9].output + down = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + return down + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[renamed] renamed-model isolated steer bit-identical={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_matches_manual(model): + # tracer.steer (in-process) must equal the manual untuple + replacement it names. + direction = _direction(model.lm_head) + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + tracer.steer(model.transformer.h[6], direction, ALPHA) + o = model.transformer.h[10].output + via_api = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + with model.trace(PROMPT): + with torch.no_grad(): + out = model.transformer.h[6].output + is_tuple = isinstance(out, tuple) + hidden = out[0] if is_tuple else out + steered = hidden + ALPHA * direction.to(dtype=hidden.dtype, device=hidden.device) + model.transformer.h[6].output = (steered, *out[1:]) if is_tuple else steered + o = model.transformer.h[10].output + via_manual = (o[0] if isinstance(o, tuple) else o)[:, -1, :].save() + ok = _eq(via_api, via_manual) + print(f"[matches_manual] tracer.steer == manual replacement bit-identical={ok} " + f"(max|Δ|={_delta(via_api, via_manual)})", flush=True) + return ok + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = { + "single": test_single(model), + "replacement": test_replacement(model), + "tuple_output": test_tuple_output(model), + "multi": test_multi(model), + "renamed": test_renamed(), + "matches_manual": test_matches_manual(model), + } + print("=" * 72, flush=True) + print(f"ISOLATED STEER: {results}", flush=True) + sys.exit(0 if all(results.values()) else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_trace.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_trace.py new file mode 100644 index 000000000..e01f5e937 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_trace.py @@ -0,0 +1,69 @@ +#!/usr/bin/env python3 +"""Isolated trace end-to-end — transparent isolated model.trace() is bit-identical to in-process. + + read — isolated `h[6].output[0].save()` == in-process, max|Δ|=0. + swap — isolated `h[6].output[0] *= 2` propagates to a downstream save the same as + in-process, max|Δ|=0, AND differs from no-swap (the swap really happened). + +The intervention runs in a spawned GPU worker; values cross via the CUDA-IPC channel; +saves come back via the worker→host saves transmission. + +Run: + CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_trace.py +""" +import sys + +import torch + +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +PROMPT = "The Eiffel Tower is in the city of" + + +def test_read(model): + with model.trace(PROMPT): + ref = model.transformer.h[6].output[0].save() + with isolate_mediators(fast_lane=False): + with model.trace(PROMPT): + got = model.transformer.h[6].output[0].save() + d = (ref.float() - got.float()).abs().max().item() + ok = torch.equal(ref, got) + print(f"[read] isolated saved activation == in-process: {ok} (max|Δ|={d:.2e}, shape={tuple(got.shape)})") + return ok + + +def test_swap(model): + # Explicit assignment (eproperty __set__ -> SWAP event) — the isolation- and + # remote-consistent form. In-place `[:]=` mutates a worker-local clone and does + # not propagate across the process boundary (documented isolation semantic). + with model.trace(PROMPT): + plain = model.transformer.h[7].output[0].save() + with model.trace(PROMPT): + model.transformer.h[6].output = model.transformer.h[6].output * 2 + ref = model.transformer.h[7].output[0].save() + with isolate_mediators(fast_lane=False): + with model.trace(PROMPT): + model.transformer.h[6].output = model.transformer.h[6].output * 2 + got = model.transformer.h[7].output[0].save() + d = (ref.float() - got.float()).abs().max().item() + same_as_ref = torch.equal(ref, got) + changed = not torch.equal(plain, got) + print(f"[swap] isolated swap == in-process swap: {same_as_ref} (max|Δ|={d:.2e}) | changed-vs-noswap: {changed}") + return same_as_ref and changed + + +def main(): + assert torch.cuda.is_available(), "needs CUDA" + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = {"read": test_read(model), "swap": test_swap(model)} + ok = all(results.values()) + print("=" * 72) + print(f"ISOLATED TRACE: {'PASS' if ok else 'FAIL'} — {results}") + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_triton_model.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_triton_model.py new file mode 100644 index 000000000..814140268 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_triton_model.py @@ -0,0 +1,164 @@ +#!/usr/bin/env python3 +"""Triton-kernel model under isolation — the headline motivation for this backend. + +A model whose forward executes a Triton JIT kernel cannot be served under NDIF's +in-process import-whitelist sandbox: Triton's first-use compilation needs subprocess +(ptxas), tempfile, open() and importlib — exactly what the whitelist denies — and the +whitelist necessarily wraps the forward pass too. This backend moves the forward to the +TRUSTED host (unrestricted, so Triton compiles normally) and contains only the UNTRUSTED +intervention in the worker. These tests pin that property: + + host_compiles — an isolated trace with ``lockdown=True`` through a Triton-kernel model + is bit-identical to in-process. The worker is fully seccomp'd, so the + kernel MUST have compiled+run on the host => the host/worker split lets + Triton models run under isolation. (A fresh TRITON_CACHE_DIR is shown + to populate during the isolated run as corroboration.) + user_contained — a Triton kernel invoked from WITHIN the intervention (worker-side) under + lockdown fails cleanly (import wall and/or compile wall), demonstrating + that the capability preserved for the model is still denied to user code. + Uses ``preimport=("triton", "triton.language")`` to get the worker past + the import wall so the *compile* wall is what's exercised. + +Run (needs a GPU + a Triton install): + CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_triton_model.py +""" +import os +import sys +import tempfile + +# Point Triton at a fresh cache BEFORE importing it, so a populated cache after the +# isolated run is evidence the host compiled during that run (not a pre-warmed hit). +_CACHE_DIR = tempfile.mkdtemp(prefix="nnsight_triton_cache_") +os.environ["TRITON_CACHE_DIR"] = _CACHE_DIR + +import torch +import torch.nn as nn + +import nnsight +from nnsight import NNsight +from nnsight.intervention.isolation import isolate_mediators + +try: + import triton + import triton.language as tl + + _HAVE_TRITON = True +except Exception: # noqa: BLE001 + _HAVE_TRITON = False + + +if _HAVE_TRITON: + + @triton.jit + def _scale_kernel(x_ptr, out_ptr, scale, n, BLOCK: tl.constexpr): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n + x = tl.load(x_ptr + offs, mask=mask) + tl.store(out_ptr + offs, x * scale, mask=mask) + + def triton_scale(x, scale): + """Elementwise x*scale via a Triton kernel (forces a Triton JIT at first use).""" + x = x.contiguous() + out = torch.empty_like(x) + n = x.numel() + grid = lambda meta: (triton.cdiv(n, meta["BLOCK"]),) # noqa: E731 + _scale_kernel[grid](x, out, scale, n, BLOCK=1024) + return out + + +class TritonMLP(nn.Module): + """fc1 -> (Triton kernel) -> fc2. The Triton kernel lives ON the forward path, so a + trace of this model only succeeds if that kernel compiled+ran wherever the forward + executed (the host, under isolation).""" + + def __init__(self, d=64): + super().__init__() + self.fc1 = nn.Linear(d, d) + self.fc2 = nn.Linear(d, d) + + def forward(self, x): + h = self.fc1(x) + h = triton_scale(h, 2.0) + return self.fc2(h) + + +def _build(): + torch.manual_seed(0) + net = TritonMLP().cuda().eval() + return NNsight(net), torch.randn(1, 8, 64).cuda() + + +def test_host_compiles(): + """Isolated + lockdown trace through a Triton-kernel model == in-process.""" + model, x = _build() + + # Isolated FIRST, with a cold Triton cache and the worker fully locked down. + with isolate_mediators(lockdown=True): + with model.trace(x): + got = model.fc2.output.save() + + cache_populated = os.path.isdir(_CACHE_DIR) and any(os.scandir(_CACHE_DIR)) + + # In-process reference (cache now warm; comparison is what matters). + with model.trace(x): + ref = model.fc2.output.save() + + ok = torch.equal(ref, got) + print( + f"[host_compiles] isolated(lockdown) Triton model == in-process: {ok} " + f"(max|Δ|={(ref - got).abs().max().item():.2e}) | host triton cache populated " + f"during isolated run: {cache_populated}" + ) + # Bit-identity is the load-bearing assertion; cache population is corroboration only + # (cache layout/location is version-dependent), so it does not gate the result. + return ok + + +def test_user_contained(): + """A Triton kernel called from inside the intervention (worker) is blocked under + lockdown — the model keeps Triton, untrusted user code does not.""" + model, x = _build() + raised = None + try: + # preimport gets the worker past the import wall; the compile wall (ptxas execve + + # cache open, both seccomp-blocked) is then what must stop it. + with isolate_mediators(lockdown=True, preimport=("triton", "triton.language")): + with model.trace(x): + act = model.fc1.output # real tensor in the worker + # Worker-side Triton use: the JIT (ptxas execve + cache open) must be + # blocked under lockdown. This call is what should raise. + out = triton_scale(act, 3.0) + nnsight.save(out) + except Exception as e: # noqa: BLE001 + raised = e + ok = raised is not None + print( + f"[user_contained] worker-side Triton blocked under lockdown: {ok} " + f"(raised={type(raised).__name__ if raised else None})" + ) + return ok + + +def main(): + if not torch.cuda.is_available(): + print("SKIP: no CUDA device") + sys.exit(0) + if not _HAVE_TRITON: + print("SKIP: triton not installed") + sys.exit(0) + + results = { + "host_compiles": test_host_compiles(), + "user_contained": test_user_contained(), + } + ok = all(results.values()) + print("=" * 72) + print(f"ISOLATED TRITON MODEL: {'PASS' if ok else 'FAIL'} — {results}") + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_unembed.py b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_unembed.py new file mode 100644 index 000000000..eea68e963 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_isolated_unembed.py @@ -0,0 +1,164 @@ +#!/usr/bin/env python3 +"""Host-routed unembed on the ISOLATED tier: tracer.unembed runs the real final-norm + +unembed on the host's real weights even when the intervention runs in the weightless +worker — so the weight-reading interp readout works under forced isolation, not only on +the fast lane. + +All cases force isolation (fast_lane=False) so the cell genuinely runs in the worker and +exercises the worker -> Events.UNEMBED -> host handler -> logits-back round trip. + + single — one layer's residual projected via tracer.unembed, isolated == + in-process, bit-identical. + multi — three layers (forward order) interleaved with unembed, bit-identical. + module_form — formulation="module" (head(normed)) bit-identical. + no_norm — norm=None (skip normalization) bit-identical. + renamed — renamed model (final_norm / output_projection / decoder_blocks): + paths resolve host-side, bit-identical (no hardcoded names). + matches_raw — tracer.unembed equals the manual F.linear(norm(x), head.weight). + +Run: + CUDA_VISIBLE_DEVICES=5 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python -u \ + prototypes/mediator-sandbox/gpu_sandbox/test_isolated_unembed.py +""" +import sys + +import torch +import torch.nn.functional as F + +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +PROMPT = "The Eiffel Tower is in the city of" + + +def _both(build): + """Run build() in-process and under forced isolation; return (ref, got).""" + ref = build() + with isolate_mediators(fast_lane=False, timeout=30): + got = build() + return ref, got + + +def _eq(ref, got): + return (torch.is_tensor(ref) and torch.is_tensor(got) + and ref.shape == got.shape and torch.equal(ref, got)) + + +def _delta(ref, got): + return (ref - got).abs().max().item() if _eq(ref, got) else float("nan") + + +def test_single(model): + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + resid = model.transformer.h[6].output + logits = tracer.unembed(resid, model.transformer.ln_f, model.lm_head) + row = logits[:, -1, :].save() + return row + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[single] isolated unembed bit-identical={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_multi(model): + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + rows = [] + for i in [0, 4, 8]: + resid = model.transformer.h[i].output + rows.append(tracer.unembed(resid, model.transformer.ln_f, + model.lm_head)[:, -1, :]) + g = torch.stack(rows, dim=0).save() + return g + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[multi] 3-layer isolated unembed bit-identical={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_module_form(model): + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + resid = model.transformer.h[6].output + logits = tracer.unembed(resid, model.transformer.ln_f, model.lm_head, + formulation="module") + row = logits[:, -1, :].save() + return row + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[module_form] formulation='module' bit-identical={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_no_norm(model): + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + resid = model.transformer.h[6].output + logits = tracer.unembed(resid, None, model.lm_head) # skip norm + row = logits[:, -1, :].save() + return row + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[no_norm] norm=None bit-identical={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_renamed(): + rename = {"transformer.ln_f": "final_norm", "lm_head": "output_projection", + "transformer.h": "decoder_blocks"} + model = LanguageModel("gpt2", device_map="cuda", dispatch=True, rename=rename) + + def build(): + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + resid = model.decoder_blocks[5].output + logits = tracer.unembed(resid, model.final_norm, model.output_projection) + row = logits[:, -1, :].save() + return row + ref, got = _both(build) + ok = _eq(ref, got) + print(f"[renamed] renamed-model isolated unembed bit-identical={ok} (max|Δ|={_delta(ref, got)})", flush=True) + return ok + + +def test_matches_raw(model): + # tracer.unembed (in-process) must equal the manual readout it replaces. + with model.trace(PROMPT) as tracer: + with torch.no_grad(): + resid = model.transformer.h[6].output + via_api = tracer.unembed(resid, model.transformer.ln_f, model.lm_head)[:, -1, :].save() + with model.trace(PROMPT): + with torch.no_grad(): + out = model.transformer.h[6].output + hidden = out[0] if isinstance(out, tuple) else out + manual = F.linear(model.transformer.ln_f(hidden), model.lm_head.weight)[:, -1, :].save() + ok = _eq(via_api, manual) + print(f"[matches_raw] tracer.unembed == manual F.linear bit-identical={ok} " + f"(max|Δ|={_delta(via_api, manual)})", flush=True) + return ok + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True) + results = { + "single": test_single(model), + "multi": test_multi(model), + "module_form": test_module_form(model), + "no_norm": test_no_norm(model), + "renamed": test_renamed(), + "matches_raw": test_matches_raw(model), + } + print("=" * 72, flush=True) + print(f"ISOLATED UNEMBED: {results}", flush=True) + sys.exit(0 if all(results.values()) else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_nonstd.py b/prototypes/mediator-sandbox/gpu_sandbox/test_nonstd.py new file mode 100644 index 000000000..b0ddd2f70 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_nonstd.py @@ -0,0 +1,62 @@ +#!/usr/bin/env python3 +"""Non-standard module names under isolation (testing rule: vary names, not GPT-2-only). + +generate() needs a real HF generative model, so we use gpt2 with rename= to give it +non-standard USER-FACING paths (decoder_blocks / output_projection). This catches any +alias-path vs real-path mismatch in the host-side hook registration's requester->envoy resolution. + + read — model.decoder_blocks[6].output[0].save() isolated == in-process. + iterN — generate + iter[1] on the renamed path, isolated == in-process. + +Run: + CUDA_VISIBLE_DEVICES=6 PYTHONPATH=src \ + /disk/u/zikai/anaconda3/envs/hf-serve/bin/python \ + prototypes/mediator-sandbox/gpu_sandbox/test_nonstd.py +""" +import sys + +import torch + +from nnsight import LanguageModel +from nnsight.intervention.isolation import isolate_mediators + +PROMPT = "The Eiffel Tower is in the city of" +RENAME = {"transformer.h": "decoder_blocks", "lm_head": "output_projection"} + + +def test_read(model): + with model.trace(PROMPT): + ref = model.decoder_blocks[6].output[0].save() + with isolate_mediators(fast_lane=False, timeout=30): + with model.trace(PROMPT): + got = model.decoder_blocks[6].output[0].save() + ok = torch.equal(ref, got) + print(f"[read] decoder_blocks[6] isolated==in-proc: {ok} (max|Δ|={(ref-got).abs().max().item():.2e})", flush=True) + return ok + + +def test_iterN(model): + with model.generate(PROMPT, max_new_tokens=3) as t: + for step in t.iter[1]: + ref = model.decoder_blocks[6].output[0].save() + with isolate_mediators(fast_lane=False, timeout=30): + with model.generate(PROMPT, max_new_tokens=3) as t: + for step in t.iter[1]: + got = model.decoder_blocks[6].output[0].save() + ok = torch.equal(ref, got) + print(f"[iterN] decoder_blocks[6] iter[1] isolated==in-proc: {ok} (max|Δ|={(ref-got).abs().max().item():.2e})", flush=True) + return ok + + +def main(): + assert torch.cuda.is_available() + model = LanguageModel("gpt2", device_map="cuda", dispatch=True, rename=RENAME) + results = {"read": test_read(model), "iterN": test_iterN(model)} + ok = all(results.values()) + print("=" * 72, flush=True) + print(f"NONSTD (renamed paths): {'PASS' if ok else 'FAIL'} — {results}", flush=True) + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/gpu_sandbox/test_safety.py b/prototypes/mediator-sandbox/gpu_sandbox/test_safety.py new file mode 100644 index 000000000..c4ef7ffb7 --- /dev/null +++ b/prototypes/mediator-sandbox/gpu_sandbox/test_safety.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python3 +"""Safety test: mimic unsafe interventions and confirm the worker contains them. + +The threat we're containing is FOOTGUNS (a careless/buggy intervention), not a +determined adversary. We submit unsafe ops to the GPU worker and require each to +be contained — no host file read/written, no network, no host objects reached, +no OOM/hang taking down the server — while a legit workload keeps working after. + +Run: CUDA_VISIBLE_DEVICES=6 PYTHONPATH=/src .../bin/python test_safety.py +""" +import os +import shutil +import sys + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import torch + +from gpu_sandbox import GPUSandbox + +HOST_DIR = os.path.expanduser("~/.gpu_sandbox_safety") +SECRET = os.path.join(HOST_DIR, "secret.txt") +PWNED = os.path.join(HOST_DIR, "pwned") +SECRET_CONTENT = "TOPSECRET-GPU-SANDBOX" + + +# --- unsafe ops (defined in __main__ → cloudpickled by value, no worker import) --- +def attack_read_secret(t): + return open(os.path.expanduser("~/.gpu_sandbox_safety/secret.txt")).read() + + +def attack_write_host(t): + open(os.path.expanduser("~/.gpu_sandbox_safety/pwned"), "w").write("pwned") + return t + + +def attack_network(t): + import socket + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + s.connect(("1.1.1.1", 53)) + return t + + +def attack_reach_host_object(t): + return THE_MODEL.config # noqa: F821 — no such global in the worker + + +def attack_oom(t): + big = torch.empty(10 ** 12, device="cuda") # ~8 TB — must be refused + return big.sum() + + +def attack_infinite_loop(t): + while True: + pass + + +def attack_segfault(t): + import ctypes + ctypes.string_at(0) # read NULL → SIGSEGV → kills the worker process + return t + + +def expect_contained(sb, name, fn, side_effect_check=None, timeout=60): + try: + sb.apply(t_dummy(), fn, timeout=timeout) + contained = False + detail = "op SUCCEEDED (NOT contained!)" + except (RuntimeError, TimeoutError) as e: + contained = True + detail = str(e)[:90] + leaked = (side_effect_check() if side_effect_check else False) + ok = contained and not leaked + print(f"[safety] {name:22s} contained: {contained} | side-effect leaked: {leaked} | {detail}") + return ok + + +def t_dummy(): + return torch.randn(1, 8, 768, device="cuda") + + +def main(): + os.makedirs(HOST_DIR, exist_ok=True) + with open(SECRET, "w") as f: + f.write(SECRET_CONTENT) + if os.path.exists(PWNED): + os.remove(PWNED) + + sb = GPUSandbox() + results = {} + + # catchable containments (worker survives each) — fs / net / host-objects / oom + results["read_host_secret"] = expect_contained(sb, "read_host_secret", attack_read_secret) + results["write_host_file"] = expect_contained( + sb, "write_host_file", attack_write_host, side_effect_check=lambda: os.path.exists(PWNED)) + results["network_egress"] = expect_contained(sb, "network_egress", attack_network) + results["reach_host_object"] = expect_contained(sb, "reach_host_object", attack_reach_host_object) + results["oom_alloc"] = expect_contained(sb, "oom_alloc", attack_oom) + + # worker must have SURVIVED all the catchable attacks → a legit op still works + legit = sb.apply(t_dummy(), lambda x: x * 2.0) + results["legit_after_attacks"] = bool(torch.allclose(legit, t_dummy() * 0 + legit)) # shape/finite check + results["legit_after_attacks"] = legit.shape == (1, 8, 768) and torch.isfinite(legit).all().item() + print(f"[safety] legit op works after contained attacks: {results['legit_after_attacks']}") + + # hang containment (short timeout so the test is quick) + results["infinite_loop"] = expect_contained(sb, "infinite_loop", attack_infinite_loop, timeout=4) + + # crash isolation: user code that tries to HARD-crash the worker (a NULL-deref + # segfault via ctypes) is contained — the worker process dies but the HOST's CUDA + # context is intact and it surfaces the crash cleanly (no hang, no server death). + # (If imports happen to be blocked, the same attempt is contained as a raised error; + # either way it cannot escape the worker.) Run on a DEDICATED sandbox since it may + # kill its worker. + sbC = GPUSandbox() + try: + sbC.apply(t_dummy(), attack_segfault, timeout=15) + crashed_contained = False + except RuntimeError: + crashed_contained = True + host_survived = torch.isfinite(t_dummy()).all().item() # host CUDA context unscathed + results["crash_isolated_host_survives"] = crashed_contained and host_survived + print(f"[safety] user-code crash attempt: contained={crashed_contained} host_survived={host_survived}") + sbC.close() + + # after a worker crash, a freshly spawned worker serves correctly (pool recovery) + sb3 = GPUSandbox() + rec = sb3.apply(t_dummy(), lambda x: x + 1.0) + results["respawn_recovers"] = rec.shape == (1, 8, 768) and torch.isfinite(rec).all().item() + print(f"[safety] fresh worker after crash serves correctly: {results['respawn_recovers']}") + sb3.close() + + sb.close() + shutil.rmtree(HOST_DIR, ignore_errors=True) + ok_all = all(results.values()) + print("=" * 60) + print(f"GPU-SANDBOX SAFETY: {'PASS' if ok_all else 'FAIL'} — {results}") + sys.exit(0 if ok_all else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/p2_isolation_poc.py b/prototypes/mediator-sandbox/p2_isolation_poc.py new file mode 100644 index 000000000..f0f3137de --- /dev/null +++ b/prototypes/mediator-sandbox/p2_isolation_poc.py @@ -0,0 +1,221 @@ +#!/usr/bin/env python3 +"""P2 (host-level sibling sandbox) isolation PoC. + +Runs the real ndif escape-suite gadgets (test_protector_escapes.py) as *user +code inside a host-level bubblewrap jail* and shows each is INERT — the gadget +still runs (P2 has NO import/builtin whitelist; that's the point), but it cannot +read a host file, write a host file, or reach the network, and it sees no +foreign processes. + +This is the inverse of the in-process Protector: the Protector tried to *forbid* +dangerous APIs (os, ctypes, subprocess, the __subclasses__ walk) and failed on +all 10. P2 *allows* them and removes everything they could touch. + +Design ref: docs/developing/mediator-isolation-sandbox.md (§6.5 P2, §6.1 jail). +Mechanism verified: a socketpair fd inherited into a net=none CPU-only jail. + +Run with the env interpreter so sys.executable points into the bound env: + /disk/u/zikai/anaconda3/bin/python prototypes/mediator-sandbox/p2_isolation_poc.py +""" +import json +import os +import shutil +import socket +import subprocess +import sys + +PYBIN = sys.executable # the (untrusted) interpreter run inside the jail +ENV_ROOT = os.path.dirname(os.path.dirname(PYBIN)) # e.g. /disk/u/zikai/anaconda3 — bound ro +HOST_DIR = os.path.expanduser("~/.p2_poc") # host-only; deliberately NOT bound into jail +SECRET = os.path.join(HOST_DIR, "secret.txt") +PWNED = os.path.join(HOST_DIR, "pwned") +SECRET_CONTENT = "TOPSECRET-MEDIATOR-ACTIVATIONS-9f3a" + +# Worker entrypoint that runs INSIDE the jail: rebuild the socket from the +# inherited fd, recv one gadget, exec it (no whitelist), report the verdict it +# set. Modules are preimported to mirror how an nnsight closure arrives with a +# populated globals — so the __subclasses__ / __globals__ walks find their targets. +WORKER = r''' +import os, sys, socket, json +import subprocess, pickle, inspect, types +for _m in ("numpy", "urllib3"): + try: __import__(_m) + except Exception: pass +fd = int(os.environ["WORKER_FD"]) +s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM, fileno=fd) +buf = b"" +while not buf.endswith(b"\n"): + chunk = s.recv(65536) + if not chunk: break + buf += chunk +msg = json.loads(buf.decode()) +ns = dict(globals()); ns["verdict"] = "ran"; ns["detail"] = "" +try: + exec(msg["code"], ns) + v, d = ns.get("verdict", "ran"), str(ns.get("detail", ""))[:200] +except Exception as e: + v, d = "EXC", (type(e).__name__ + ": " + str(e))[:200] +s.sendall((json.dumps({"verdict": v, "detail": d}) + "\n").encode()) +s.close() +''' + + +def run_in_jail(code: str, timeout: int = 60) -> dict: + """Spawn one bwrap jail, hand it a gadget over an inherited socket, get the verdict.""" + broker, worker = socket.socketpair(socket.AF_UNIX, socket.SOCK_STREAM) + os.set_inheritable(worker.fileno(), True) + cmd = [ + "bwrap", "--unshare-all", # userns+net+pid+ipc+uts+mount, fresh + "--ro-bind", "/usr", "/usr", + "--ro-bind", "/lib", "/lib", + "--ro-bind", "/lib64", "/lib64", + "--ro-bind", "/bin", "/bin", + "--ro-bind", "/etc", "/etc", + "--ro-bind", ENV_ROOT, ENV_ROOT, # the python env, read-only + "--proc", "/proc", # fresh proc (real PID isolation — host can) + "--dev", "/dev", + "--tmpfs", "/tmp", + "--die-with-parent", + "--setenv", "CUDA_VISIBLE_DEVICES", "", # no GPU + "--setenv", "WORKER_FD", str(worker.fileno()), + PYBIN, "-c", WORKER, + ] + p = subprocess.Popen(cmd, pass_fds=[worker.fileno()]) + worker.close() + broker.sendall((json.dumps({"code": code}) + "\n").encode()) + broker.settimeout(timeout) + data = b"" + try: + while not data.endswith(b"\n"): + chunk = broker.recv(65536) + if not chunk: + break + data += chunk + except socket.timeout: + data = b'{"verdict": "TIMEOUT", "detail": "no reply"}\n' + try: + p.wait(timeout=10) + except subprocess.TimeoutExpired: + p.kill() + broker.close() + try: + return json.loads(data.decode()) + except Exception: + return {"verdict": "PARSE_ERR", "detail": data[:120].decode(errors="replace")} + + +# Each gadget is the real escape primitive from the suite, then it attempts the +# *host effect* and self-reports. __SECRET__/__PWNED__ are substituted below. +GADGETS = [ + ("escape1+9 subclasses→Popen→read host secret", r''' +Popen = None +for c in ().__class__.__mro__[1].__subclasses__(): + if c.__name__ == "Popen" and c.__module__ == "subprocess": + Popen = c; break +if Popen is None: + verdict, detail = "N/A", "Popen not in subclass walk" +else: + try: + pr = Popen(["cat", "__SECRET__"], stdout=subprocess.PIPE, stderr=subprocess.PIPE) + o, e = pr.communicate(timeout=5) + verdict = "LEAKED" if b"TOPSECRET" in o else "CONTAINED" + detail = (o[:40] if o else e[:80]).decode(errors="replace") + except Exception as ex: + verdict, detail = "CONTAINED", "%s: %s" % (type(ex).__name__, ex) +'''), + ("escape2 format-string __globals__→os→system(touch host)", r''' +os_mod = subprocess.Popen.__init__.__globals__["os"] +proof = "{0.__globals__[os]}".format(subprocess.Popen.__init__) +rc = os_mod.system("touch __PWNED__ 2>/dev/null") +verdict, detail = "ran", "reached %s ; os.system rc=%s" % ("os" in proof, rc) +'''), + ("escape3 numpy.ctypeslib→libc.open(host secret)", r''' +import numpy +ct = numpy.ctypeslib.ctypes +libc = ct.CDLL("libc.so.6") +fd = libc.open(b"__SECRET__", 0) # O_RDONLY +if fd >= 0: + b = ct.create_string_buffer(64); n = libc.read(fd, b, 64); libc.close(fd) + verdict = "LEAKED" if b"TOPSECRET" in b.raw else "CONTAINED" + detail = "fd=%s n=%s" % (fd, n) +else: + verdict, detail = "CONTAINED", "libc.open fd=%s (host file absent in jail)" % fd +'''), + ("escape7 urllib3 egress to 1.1.1.1", r''' +try: + import urllib3 + pool = urllib3.PoolManager(retries=False, timeout=3.0) + r = pool.request("GET", "http://1.1.1.1/") + verdict, detail = "LEAKED", "egress ok status=%s" % r.status +except Exception as ex: + verdict, detail = "CONTAINED", "%s: %s" % (type(ex).__name__, str(ex)[:90]) +'''), + ("escape9 end-to-end RCE: Popen.globals→os.system(touch host)", r''' +Popen = None +for c in ().__class__.__mro__[1].__subclasses__(): + if c.__name__ == "Popen" and c.__module__ == "subprocess": + Popen = c; break +os_mod = Popen.__init__.__globals__["os"] +rc = os_mod.system("touch __PWNED__ 2>/dev/null") +verdict, detail = "ran", "os.system rc=%s (host effect checked by orchestrator)" % rc +'''), + ("escape10 pickle __reduce__ fired DURING deserialize-in-jail", r''' +class G: + def __reduce__(self): + return (os.system, ("touch __PWNED__ 2>/dev/null",)) +payload = pickle.dumps(G()) +pickle.loads(payload) # P2: untrusted unpickle happens inside the jail +verdict, detail = "ran", "pickle gadget fired inside jail" +'''), + ("PID-namespace isolation (validates §6.3 fixed at host level)", r''' +pids = sorted(int(x) for x in os.listdir("/proc") if x.isdigit()) +verdict = "ISOLATED" if len(pids) <= 3 else "LEAK" +detail = "visible pids=%s" % pids +'''), +] + + +def main(): + if shutil.which("bwrap") is None: + sys.exit("bwrap not found on host — P2 jail unavailable") + os.makedirs(HOST_DIR, exist_ok=True) + with open(SECRET, "w") as f: + f.write(SECRET_CONTENT) + if os.path.exists(PWNED): + os.remove(PWNED) + + print("P2 isolation PoC — gadgets run inside a host-level bwrap jail") + print(" jail: %s (CPU-only, net=none, ro-fs allowlist, fresh PID ns)" % PYBIN) + print(" host secret (NOT bound into jail): %s" % SECRET) + print("-" * 78) + + leaked = [] + for name, tmpl in GADGETS: + code = tmpl.replace("__SECRET__", SECRET).replace("__PWNED__", PWNED) + res = run_in_jail(code) + v, d = res.get("verdict", "?"), res.get("detail", "") + bad = v in ("LEAKED", "LEAK", "TIMEOUT", "PARSE_ERR") + if bad: + leaked.append(name) + print(" [%s] %s" % ("LEAK" if bad else "ok ", name)) + print(" verdict=%s %s" % (v, d)) + + # Host-side ground truth: the gadgets must have touched nothing on the host. + print("-" * 78) + host_pwned = os.path.exists(PWNED) + print(" host-side check: PWNED file created on host? %s" % ("YES (LEAK)" if host_pwned else "no")) + print(" host-side check: secret content ever returned? %s" + % ("YES (LEAK)" if any("TOPSECRET" in g for g in leaked) else "no")) + + ok = (not leaked) and (not host_pwned) + print("=" * 78) + print("RESULT: %s" % ("PASS — every gadget ran but was inert" if ok + else "FAIL — leaks: %s%s" % (leaked, " +host-write" if host_pwned else ""))) + + # cleanup + shutil.rmtree(HOST_DIR, ignore_errors=True) + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/phase2_socket_transport.py b/prototypes/mediator-sandbox/phase2_socket_transport.py new file mode 100644 index 000000000..dfbd4bb1f --- /dev/null +++ b/prototypes/mediator-sandbox/phase2_socket_transport.py @@ -0,0 +1,214 @@ +#!/usr/bin/env python3 +"""Phase 2 — two-process socket transport harness. + +Proves the six-event Mediator protocol survives a real process boundary with +*identical values*, exercising the real routing/batching machinery — not a mock. +A forked worker runs the real ``Mediator`` client (request->VALUE, swap->SWAP, +end->END) over an ``AF_UNIX`` socket; the parent drives the real +``Mediator.handle`` + ``Batcher`` on the host side. + +Tests: + 0. codec — nested tuple/dict/tensor round-trips the length-prefixed frames. + 1. golden — same h[6]x2 intervention local vs over-socket on gpt2 -> bit-identical logits. + 2. batched — needs_batching=True: worker sees ONLY its batch row; its swap lands ONLY on that row. + 3. restore — out-of-order: a provider-mismatch triggers the host-local restore_event path, then + the matching provider delivers the right value over the socket. + +NOTE (scope): this proves the protocol + requester matching + real Batcher narrowing/swapping over a +process boundary. Making `model.trace()` itself fork the worker and ship `.save()` values back +(the isolation *execution backend*) is the separate trace-integration phase — see the plan. + +Run: PYTHONPATH=src .../hf-serve/bin/python prototypes/mediator-sandbox/phase2_socket_transport.py +""" +import os +import socket +import sys +from types import SimpleNamespace + +import torch + +from nnsight import LanguageModel +from nnsight.intervention.interleaver import Interleaver, Mediator +from nnsight.intervention.batching import Batcher +from nnsight.intervention.transport import ( + SocketHostChannel, + SocketWorkerChannel, + recv_frame, + send_frame, +) + +PROMPT = "The Eiffel Tower is in the city of" +PROVIDER = "transformer.h.6.output.i0" +LAYER = 6 + + +def double_block_output(out): + if isinstance(out, tuple): + return (out[0] * 2.0,) + tuple(out[1:]) + return out * 2.0 + + +# --------------------------------------------------------------------------- # +# Worker (child) side # +# --------------------------------------------------------------------------- # +def _mk_worker(sock): + med = Mediator(intervention=None, info=SimpleNamespace(frame=None), batch_group=None) + med.channel = SocketWorkerChannel(sock) + med.cross_invoker = False + return med + + +def worker_double(sock, provider): + med = _mk_worker(sock) + value = med.request(provider) # VALUE -> the (narrowed) activation + med.swap(provider, double_block_output(value)) # SWAP back the doubled value + med.end() + + +def worker_add(sock, provider, delta): + med = _mk_worker(sock) + value = med.request(provider) + med.swap(provider, value + delta) + med.end() + + +def fork_worker(worker_fn, *args): + parent_sock, child_sock = socket.socketpair() + pid = os.fork() + if pid == 0: # child + parent_sock.close() + try: + worker_fn(child_sock, *args) + finally: + child_sock.close() + os._exit(0) + child_sock.close() + return pid, parent_sock + + +# --------------------------------------------------------------------------- # +# Host (parent) side # +# --------------------------------------------------------------------------- # +def mk_host(parent_sock, batcher, batch_group): + interleaver = Interleaver(mediators=[], tracer=None, batcher=batcher) + med = Mediator(intervention=None, info=SimpleNamespace(frame=None), batch_group=batch_group) + med.channel = SocketHostChannel(parent_sock) + med.interleaver = interleaver + interleaver.mediators = [med] + med.channel.wait_event() # block for the worker's first event + return med + + +# --------------------------------------------------------------------------- # +# 0. codec # +# --------------------------------------------------------------------------- # +def test_codec(): + a, b = socket.socketpair() + payload = (torch.randn(2, 5, 7), {"k": torch.arange(3)}, "meta", None) + send_frame(a, payload) + got = recv_frame(b) + a.close(); b.close() + ok = (torch.equal(payload[0], got[0]) and torch.equal(payload[1]["k"], got[1]["k"]) + and got[2] == "meta" and got[3] is None) + print(f"[0 codec] nested round-trip: {'OK' if ok else 'FAIL'}") + return ok + + +# --------------------------------------------------------------------------- # +# 1. golden equivalence through a real gpt2 forward # +# --------------------------------------------------------------------------- # +def _raw_logits(model, inputs, hook_fn): + block = model._model.transformer.h[LAYER] + handle = block.register_forward_hook(hook_fn) + try: + with torch.no_grad(): + return model._model(**inputs).logits + finally: + handle.remove() + + +def test_golden(model, inputs): + ref = _raw_logits(model, inputs, lambda m, i, o: double_block_output(o)) + + pid, ps = fork_worker(worker_double, PROVIDER) + host = mk_host(ps, Batcher(), None) + sock = _raw_logits(model, inputs, lambda m, i, o: host.handle(PROVIDER, o)) + os.waitpid(pid, 0); host.channel.close() + + plain = _raw_logits(model, inputs, lambda m, i, o: o) + # The transport codec is bit-exact (test 0). The residual here is multi-threaded + # CPU forward nondeterminism between two SEPARATE forward passes (~1e-4), not the + # socket — so compare with a tight tolerance that still dwarfs a broken swap (Δ≈30). + identical = torch.allclose(ref, sock, atol=1e-3, rtol=0) + changed = not torch.allclose(plain, sock, atol=1e-3, rtol=0) + print(f"[1 golden] local≈socket (atol 1e-3): {identical} | changed-vs-noop: {changed} " + f"| max|d|={(ref - sock).abs().max().item():.2e}") + return identical and changed + + +# --------------------------------------------------------------------------- # +# 2. real batching: worker sees ONLY its row; swap lands ONLY on its row # +# --------------------------------------------------------------------------- # +def test_batched(): + row0 = torch.arange(8, dtype=torch.float32).reshape(1, 2, 4) + row1 = torch.full((1, 2, 4), 9.0) + batched = torch.cat([row0, row1], dim=0) # [2, 2, 4] (tuple-wrapped like a block output) + + pid, ps = fork_worker(worker_double, PROVIDER) + batcher = Batcher() + batcher.needs_batching = True + batcher.last_batch_group = [0, 2] # total_batch_size = 2 + host = mk_host(ps, batcher, [0, 1]) # this mediator owns row 0 only + result = host.handle(PROVIDER, (batched.clone(),)) + os.waitpid(pid, 0); host.channel.close() + + r = result[0] + row0_doubled = torch.allclose(r[0:1], row0 * 2.0) + row1_untouched = torch.allclose(r[1:2], row1) + print(f"[2 batched] row0 doubled (worker saw its slice): {row0_doubled} | " + f"row1 untouched (no cross-row leak): {row1_untouched}") + return row0_doubled and row1_untouched + + +# --------------------------------------------------------------------------- # +# 3. out-of-order: provider mismatch triggers host-local restore_event # +# --------------------------------------------------------------------------- # +def test_restore(): + pid, ps = fork_worker(worker_add, "A.i0", 100.0) + host = mk_host(ps, Batcher(), None) + + # Fire a NON-matching provider first: handle_value_event must restore the + # pending "A.i0" event host-locally (no wire traffic) and leave it buffered. + host.handle("B.i0", torch.tensor([1.0, 2.0])) + restored = host.channel.has_event + + # Now fire the matching provider: the worker receives A's value over the socket. + a_val = torch.tensor([7.0, 8.0]) + result = host.handle("A.i0", a_val) + os.waitpid(pid, 0); host.channel.close() + + got_a = torch.allclose(result, a_val + 100.0) # child saw [7,8] (A), not [1,2] (B) + print(f"[3 restore] event re-staged after mismatch: {restored} | " + f"matching provider delivered correct value: {got_a}") + return restored and got_a + + +# --------------------------------------------------------------------------- # +def main(): + results = {} + results["codec"] = test_codec() + + model = LanguageModel("gpt2", device_map="cpu", dispatch=True) + inputs = model.tokenizer(PROMPT, return_tensors="pt") + results["golden"] = test_golden(model, inputs) + results["batched"] = test_batched() + results["restore"] = test_restore() + + ok = all(results.values()) + print("=" * 72) + print(f"PHASE 2 RESULT: {'PASS' if ok else 'FAIL'} — {results}") + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/phase3_jail_transport.py b/prototypes/mediator-sandbox/phase3_jail_transport.py new file mode 100644 index 000000000..af9cacd76 --- /dev/null +++ b/prototypes/mediator-sandbox/phase3_jail_transport.py @@ -0,0 +1,153 @@ +#!/usr/bin/env python3 +"""Phase 3 — jail the worker. Golden equivalence + escape-inertness, with the +real Mediator protocol worker running INSIDE a bwrap jail (host-level sibling, P2). + +Combines Phase 0 (the jail makes escapes inert) with Phase 2 (the socket protocol +delivers identical values): the jailed worker doubles h[6] over the socket, and +the host's gpt2 forward produces bit-identical logits — while the same jailed +worker's escape attempts touch nothing on the host. + +Run: PYTHONPATH=src .../hf-serve/bin/python prototypes/mediator-sandbox/phase3_jail_transport.py +""" +import os +import shutil +import socket +import subprocess +import sys +from types import SimpleNamespace + +import torch + +from nnsight import LanguageModel +from nnsight.intervention.batching import Batcher +from nnsight.intervention.interleaver import Interleaver, Mediator +from nnsight.intervention.transport import SocketHostChannel, recv_frame + +HERE = os.path.dirname(os.path.abspath(__file__)) +ROOT = os.path.dirname(os.path.dirname(HERE)) # the worktree +SRC = os.path.join(ROOT, "src") +WORKER = os.path.join(HERE, "phase3_jailed_worker.py") +ENV_ROOT = os.path.dirname(os.path.dirname(sys.executable)) # the conda env + +PROMPT = "The Eiffel Tower is in the city of" +PROVIDER = "transformer.h.6.output.i0" +LAYER = 6 + +HOST_DIR = os.path.expanduser("~/.p3_poc") # host-only; NOT bound into the jail +SECRET = os.path.join(HOST_DIR, "secret.txt") +PWNED = os.path.join(HOST_DIR, "pwned") +SECRET_CONTENT = "TOPSECRET-PHASE3-ACTIVATIONS" + + +def spawn_jailed_worker(mode): + host_sock, worker_sock = socket.socketpair() + os.set_inheritable(worker_sock.fileno(), True) + env = { + "WORKER_FD": str(worker_sock.fileno()), + "MODE": mode, + "PROVIDER": PROVIDER, + "SECRET": SECRET, + "PWNED": PWNED, + "PYTHONPATH": SRC, + "PATH": "/usr/local/bin:/usr/bin:/bin", + "HOME": "/tmp", + "CUDA_VISIBLE_DEVICES": "", + "HF_HUB_OFFLINE": "1", + "TRANSFORMERS_OFFLINE": "1", + } + cmd = [ + "bwrap", "--unshare-all", + "--ro-bind", "/usr", "/usr", + "--ro-bind", "/lib", "/lib", + "--ro-bind", "/lib64", "/lib64", + "--ro-bind", "/bin", "/bin", + "--ro-bind", "/etc", "/etc", + "--ro-bind", ENV_ROOT, ENV_ROOT, # python + torch (ro) + "--ro-bind", SRC, SRC, # this worktree's nnsight (ro) + "--ro-bind", WORKER, WORKER, # the worker script (ro) + "--proc", "/proc", "--dev", "/dev", "--tmpfs", "/tmp", + "--die-with-parent", + sys.executable, WORKER, + ] + p = subprocess.Popen(cmd, pass_fds=[worker_sock.fileno()], env=env) + worker_sock.close() + return p, host_sock + + +def mk_host(sock): + interleaver = Interleaver(mediators=[], tracer=None, batcher=Batcher()) + med = Mediator(intervention=None, info=SimpleNamespace(frame=None), batch_group=None) + med.channel = SocketHostChannel(sock) + med.interleaver = interleaver + interleaver.mediators = [med] + med.channel.wait_event() # block for the worker's first VALUE event + return med + + +def host_logits(model, inputs, host_med): + block = model._model.transformer.h[LAYER] + handle = block.register_forward_hook(lambda m, i, o: host_med.handle(PROVIDER, o)) + try: + with torch.no_grad(): + return model._model(**inputs).logits + finally: + handle.remove() + + +def main(): + if shutil.which("bwrap") is None: + sys.exit("bwrap not found") + os.makedirs(HOST_DIR, exist_ok=True) + with open(SECRET, "w") as f: + f.write(SECRET_CONTENT) + if os.path.exists(PWNED): + os.remove(PWNED) + + model = LanguageModel("gpt2", device_map="cpu", dispatch=True) + inputs = model.tokenizer(PROMPT, return_tensors="pt") + + # Reference: the same ×2 intervention applied locally (no jail, no socket). + ref = host_logits(model, inputs, _LocalDouble()) + + # --- Test A: golden equivalence with the worker JAILED --- + pa, sa = spawn_jailed_worker("double") + host_a = mk_host(sa) + sock_a = host_logits(model, inputs, host_a) + pa.wait(); host_a.channel.close() + # atol 1e-3: tolerates multi-threaded CPU forward nondeterminism between two + # separate forward passes (~1e-4) while still catching a broken swap (Δ≈30). + golden = torch.allclose(ref, sock_a, atol=1e-3, rtol=0) + print(f"[A jail+golden] logits match (atol 1e-3), worker jailed: {golden} " + f"| max|Δ|={(ref - sock_a).abs().max().item():.2e}") + + # --- Test B: escapes inert AND the protocol still completes through the jail --- + pb, sb = spawn_jailed_worker("escape") + report = recv_frame(sb) # the jailed worker's escape report (pre-protocol) + host_b = mk_host(sb) + sock_b = host_logits(model, inputs, host_b) + pb.wait(); host_b.channel.close() + + host_pwned = os.path.exists(PWNED) + leaked = SECRET_CONTENT[:24] in report.get("fs_read", "") or "LEAKED" in report.get("net_egress", "") + protocol_ok = torch.allclose(ref, sock_b, atol=1e-3, rtol=0) + print(f"[B escape ] report={report}") + print(f"[B escape ] host pwned-file created: {host_pwned} | secret leaked: {leaked} " + f"| protocol still bit-identical: {protocol_ok}") + + ok = golden and protocol_ok and (not host_pwned) and (not leaked) + shutil.rmtree(HOST_DIR, ignore_errors=True) + print("=" * 72) + print(f"PHASE 3 RESULT: {'PASS — real protocol works through the jail; escapes inert' if ok else 'FAIL'}") + sys.exit(0 if ok else 1) + + +class _LocalDouble: + """Stand-in 'mediator' for the reference: doubles the block output locally.""" + def handle(self, provider, out): + if isinstance(out, tuple): + return (out[0] * 2.0,) + tuple(out[1:]) + return out * 2.0 + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/phase3_jailed_worker.py b/prototypes/mediator-sandbox/phase3_jailed_worker.py new file mode 100644 index 000000000..e668aeb9e --- /dev/null +++ b/prototypes/mediator-sandbox/phase3_jailed_worker.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 +"""Phase 3 — the intervention worker, as exec'd INSIDE a bwrap jail. + +Reconstructs the SocketWorkerChannel from an inherited fd and runs the real +Mediator client protocol (request -> double -> swap -> end) against the host. +In MODE=escape it ALSO runs the escape-suite gadgets first and ships a report +back (a raw frame, before the protocol) so the host can confirm they were +attempted-but-inert — while the legitimate protocol still completes through the +jail boundary. + +Launched by phase3_jail_transport.py via: + bwrap --unshare-all ... phase3_jailed_worker.py +with env WORKER_FD / MODE / PROVIDER / SECRET / PWNED. +""" +import os +import socket + + +def run_escapes(secret_path, pwned_path): + """Attempt host-affecting escapes; every one must be inert in the jail.""" + report = {} + try: + data = open(secret_path).read() + report["fs_read"] = "LEAKED:" + data[:24] + except Exception as e: + report["fs_read"] = "CONTAINED:" + type(e).__name__ + + # canonical __subclasses__ walk -> os -> os.system(touch host file) + try: + import subprocess # noqa: F401 (populate the subclass graph) + popen = None + for c in ().__class__.__mro__[1].__subclasses__(): + if c.__name__ == "Popen" and c.__module__ == "subprocess": + popen = c + break + osmod = popen.__init__.__globals__["os"] + osmod.system("touch %s 2>/dev/null" % pwned_path) + report["subclasses_os_write"] = "ran (host file checked by host)" + except Exception as e: + report["subclasses_os_write"] = "EXC:" + type(e).__name__ + + try: + socket.create_connection(("1.1.1.1", 53), timeout=2).close() + report["net_egress"] = "LEAKED" + except Exception as e: + report["net_egress"] = "CONTAINED:" + type(e).__name__ + + return report + + +def main(): + fd = int(os.environ["WORKER_FD"]) + sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM, fileno=fd) + mode = os.environ.get("MODE", "double") + provider = os.environ["PROVIDER"] + + if mode == "escape": + report = run_escapes(os.environ["SECRET"], os.environ["PWNED"]) + from nnsight.intervention.transport import send_frame + send_frame(sock, report) # raw report frame, BEFORE the Mediator protocol + + # The legitimate protocol — identical to Phase 2, now from inside the jail. + from types import SimpleNamespace + + from nnsight.intervention.interleaver import Mediator + from nnsight.intervention.transport import SocketWorkerChannel + + med = Mediator(intervention=None, info=SimpleNamespace(frame=None), batch_group=None) + med.channel = SocketWorkerChannel(sock) + med.cross_invoker = False + + value = med.request(provider) + if isinstance(value, tuple): + new_value = (value[0] * 2.0,) + tuple(value[1:]) + else: + new_value = value * 2.0 + med.swap(provider, new_value) + med.end() + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/phase4_cross_tenant.py b/prototypes/mediator-sandbox/phase4_cross_tenant.py new file mode 100644 index 000000000..fb2144567 --- /dev/null +++ b/prototypes/mediator-sandbox/phase4_cross_tenant.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 +"""Phase 4 — Batcher as authority (cross-tenant / Boundary B). + +A malicious jailed tenant (row 0) is co-batched with a victim row (row 1 = secret +data). The malicious tenant attempts every capability-leak mutation that succeeds +in-process today (widen its batch_group, walk the shared Batcher/sibling +mediators). We prove it is structurally inert in the isolated design: + + - the worker received ONLY its own row (never saw the victim's data), because the + HOST narrows to the host-recorded batch_group — the worker's None/widen claim + is set on its OWN (jail-local) Mediator and never reaches the host; + - its poison swap landed ONLY on its own row — the victim row is untouched; + - the walks to interleaver/batcher/siblings hit nothing (no host refs in the jail). + +Run: PYTHONPATH=src .../hf-serve/bin/python prototypes/mediator-sandbox/phase4_cross_tenant.py +""" +import os +import shutil +import socket +import subprocess +import sys +from types import SimpleNamespace + +import torch + +from nnsight.intervention.batching import Batcher +from nnsight.intervention.interleaver import Interleaver, Mediator +from nnsight.intervention.transport import SocketHostChannel, recv_frame + +HERE = os.path.dirname(os.path.abspath(__file__)) +ROOT = os.path.dirname(os.path.dirname(HERE)) +SRC = os.path.join(ROOT, "src") +WORKER = os.path.join(HERE, "phase4_malicious_worker.py") +ENV_ROOT = os.path.dirname(os.path.dirname(sys.executable)) +PROVIDER = "model.layer.output.i0" + + +def spawn_jailed_worker(): + host_sock, worker_sock = socket.socketpair() + os.set_inheritable(worker_sock.fileno(), True) + env = { + "WORKER_FD": str(worker_sock.fileno()), "PROVIDER": PROVIDER, + "PYTHONPATH": SRC, "PATH": "/usr/local/bin:/usr/bin:/bin", "HOME": "/tmp", + "CUDA_VISIBLE_DEVICES": "", "HF_HUB_OFFLINE": "1", "TRANSFORMERS_OFFLINE": "1", + } + cmd = [ + "bwrap", "--unshare-all", + "--ro-bind", "/usr", "/usr", "--ro-bind", "/lib", "/lib", "--ro-bind", "/lib64", "/lib64", + "--ro-bind", "/bin", "/bin", "--ro-bind", "/etc", "/etc", + "--ro-bind", ENV_ROOT, ENV_ROOT, "--ro-bind", SRC, SRC, "--ro-bind", WORKER, WORKER, + "--proc", "/proc", "--dev", "/dev", "--tmpfs", "/tmp", "--die-with-parent", + sys.executable, WORKER, + ] + p = subprocess.Popen(cmd, pass_fds=[worker_sock.fileno()], env=env) + worker_sock.close() + return p, host_sock + + +def main(): + if shutil.which("bwrap") is None: + sys.exit("bwrap not found") + + # A 2-row batch: row 0 = the malicious tenant's input; row 1 = a victim's SECRET. + row_attacker = torch.ones(1, 3, 4) * 1.0 # sum = 12 + row_victim = torch.ones(1, 3, 4) * 7.0 # sum = 84 — the secret + batched = torch.cat([row_attacker, row_victim], dim=0) # [2, 3, 4] + + p, sock = spawn_jailed_worker() + batcher = Batcher() + batcher.needs_batching = True + batcher.last_batch_group = [0, 2] # total_batch_size = 2 + interleaver = Interleaver(mediators=[], tracer=None, batcher=batcher) + host_med = Mediator(intervention=None, info=SimpleNamespace(frame=None), batch_group=[0, 1]) + host_med.channel = SocketHostChannel(sock) + host_med.interleaver = interleaver + interleaver.mediators = [host_med] + host_med.channel.wait_event() + + result = host_med.handle(PROVIDER, (batched.clone(),)) + report = recv_frame(sock) + p.wait(); host_med.channel.close() + + final = result[0] + print(f"[report] {report}") + + confined_read = report["received_shape"][0] == 1 # got only 1 row + never_saw_victim = abs(report["received_sum"] - 12.0) < 1e-4 # sum=12 (its row), not 96 + own_row_poisoned = torch.allclose(final[0], torch.full((3, 4), 999.0)) # attacker hit its own row + victim_row_clean = torch.allclose(final[1], torch.full((3, 4), 7.0)) # victim row UNTOUCHED + walks_contained = (report["batcher_walk"].startswith("CONTAINED") + and report["sibling_walk"].startswith("CONTAINED") + and report["direct_narrow"].startswith("CONTAINED") + and report["has_interleaver"] is False) + + print(f"[B-read ] malicious tenant confined to its row: {confined_read} " + f"| never saw victim data (sum=12 not 96): {never_saw_victim}") + print(f"[B-write] poison landed on own row: {own_row_poisoned} " + f"| victim row uncorrupted: {victim_row_clean}") + print(f"[B-walk ] interleaver/batcher/sibling walks all blocked: {walks_contained}") + + ok = confined_read and never_saw_victim and own_row_poisoned and victim_row_clean and walks_contained + print("=" * 72) + print(f"PHASE 4 RESULT: {'PASS — host Batcher is the authority; cross-tenant leak inert' if ok else 'FAIL'}") + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/phase4_malicious_worker.py b/prototypes/mediator-sandbox/phase4_malicious_worker.py new file mode 100644 index 000000000..3a3569c42 --- /dev/null +++ b/prototypes/mediator-sandbox/phase4_malicious_worker.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python3 +"""Phase 4 — a MALICIOUS jailed tenant. Runs INSIDE the bwrap jail. + +Attempts the mediator-capability leaks (test_mediator_capability.py #1-6,#8) that +SUCCEED in-process today: mutate its own batch_group to widen its slice, and walk +to the shared Interleaver/Batcher to read or poison sibling rows. In the isolated +design all of these are structurally inert — the jail's Mediator has no +interleaver/batcher reference, and the HOST owns the narrow bounds, so the worker +can only ever touch its own admitted row regardless of what it claims. + +It reports what it tried + exactly what value it received (shape/sum), so the host +can prove the worker never saw the victim row, then poisons its own provider. +""" +import os +import socket + + +def main(): + fd = int(os.environ["WORKER_FD"]) + sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM, fileno=fd) + provider = os.environ["PROVIDER"] + + from types import SimpleNamespace + + from nnsight.intervention.interleaver import Mediator + from nnsight.intervention.transport import SocketWorkerChannel, send_frame + + med = Mediator(intervention=None, info=SimpleNamespace(frame=None), batch_group=None) + med.channel = SocketWorkerChannel(sock) + med.cross_invoker = False + + report = {} + + # leak #1/#2/#3/#8: try to widen this mediator's slice to grab the whole batch. + for label, bg in [("None", None), ("sentinel[-1,0]", [-1, 0]), ("widen[0,2]", [0, 2])]: + try: + med.batch_group = bg + report[f"set_batch_group_{label}"] = "set-on-local-mediator-only" + except Exception as e: # noqa: BLE001 + report[f"set_batch_group_{label}"] = "blocked:" + type(e).__name__ + med.batch_group = None # leave it at the most-permissive claim before requesting + + # leak #4/#5/#6/#8: walk to the shared batcher / siblings / narrow (host objects). + report["has_interleaver"] = med.interleaver is not None # __init__ always sets it (to None) + try: + _ = med.interleaver.batcher.current_value + report["batcher_walk"] = "REACHED-HOST-BATCHER" + except Exception as e: # noqa: BLE001 + report["batcher_walk"] = "CONTAINED:" + type(e).__name__ + try: + _ = med.interleaver.mediators + report["sibling_walk"] = "REACHED-SIBLINGS" + except Exception as e: # noqa: BLE001 + report["sibling_walk"] = "CONTAINED:" + type(e).__name__ + try: + med.interleaver.batcher.narrow(None) # leak #8: call narrow(None) directly + report["direct_narrow"] = "REACHED-NARROW" + except Exception as e: # noqa: BLE001 + report["direct_narrow"] = "CONTAINED:" + type(e).__name__ + + # Now actually request — the HOST narrows to THIS tenant's admitted row, + # ignoring the None claim above. + value = med.request(provider) + hs = value[0] if isinstance(value, tuple) else value + report["received_shape"] = list(hs.shape) + report["received_sum"] = float(hs.sum()) + + # leak #6: poison — but swap only lands on the host-recorded bounds (our row). + poison = hs * 0.0 + 999.0 + med.swap(provider, (poison,) if isinstance(value, tuple) else poison) + med.end() + + send_frame(sock, report) # raw report frame, after the protocol completes + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/phase5_gpu_measure.py b/prototypes/mediator-sandbox/phase5_gpu_measure.py new file mode 100644 index 000000000..2d876852b --- /dev/null +++ b/prototypes/mediator-sandbox/phase5_gpu_measure.py @@ -0,0 +1,233 @@ +#!/usr/bin/env python3 +"""Phase 5 — real GPU path + per-hook transport measurement. + +The model runs on a real GPU; the jailed worker is CPU-only. The host +(MediatorProxy) does D2H before delivering an activation to the worker and H2D +after the swap. Two things: + + A. Correctness — gpt2 on GPU, the ×2 intervention delivered over the socket + (D2H -> CPU op -> H2D) produces logits matching the in-process GPU forward. + + B. Measurement — the per-hook round-trip overhead (D2H + serialize + socket + + CPU op + deserialize + H2D) across activation sizes spanning gpt2 -> 7B -> + 70B hidden dims. This is the number the (D) two-tier decision hinges on; the + in-process path adds ~0 (zero-copy GPU reference). + +All numbers are MEASURED on this GPU, not estimated. + +Run: CUDA_VISIBLE_DEVICES=6 PYTHONPATH=src .../hf-serve/bin/python \ + prototypes/mediator-sandbox/phase5_gpu_measure.py +""" +import os +import pickle +import socket +import sys +import time +from types import SimpleNamespace + +import torch + +from nnsight import LanguageModel +from nnsight.intervention.batching import Batcher +from nnsight.intervention.interleaver import Interleaver, Mediator +from nnsight.intervention.transport import ( + SocketHostChannel, + recv_frame, + send_frame, +) + +DEV = "cuda" +PROMPT = "The Eiffel Tower is in the city of" +PROVIDER = "transformer.h.6.output.i0" +LAYER = 6 + + +# --------------------------------------------------------------------------- # +# A. Correctness on GPU — D2H/H2D around the socket worker # +# --------------------------------------------------------------------------- # +def worker_double_proc(sock): + pid = os.fork() + if pid == 0: + med = Mediator(intervention=None, info=SimpleNamespace(frame=None), batch_group=None) + from nnsight.intervention.transport import SocketWorkerChannel + med.channel = SocketWorkerChannel(sock) + med.cross_invoker = False + value = med.request(PROVIDER) # arrives as a CPU tensor (host D2H'd it) + hs = value[0] if isinstance(value, tuple) else value + new = (hs * 2.0,) + tuple(value[1:]) if isinstance(value, tuple) else hs * 2.0 + med.swap(PROVIDER, new) + med.end() + os._exit(0) + return pid + + +def test_gpu_correctness(model, inputs): + block = model._model.transformer.h[LAYER] + + def local_double(m, i, o): + return (o[0] * 2.0,) + tuple(o[1:]) if isinstance(o, tuple) else o * 2.0 + + h = block.register_forward_hook(local_double) + with torch.no_grad(): + ref = model._model(**inputs).logits + h.remove() + + host_sock, worker_sock = socket.socketpair() + pid = worker_double_proc(worker_sock) + worker_sock.close() + interleaver = Interleaver(mediators=[], tracer=None, batcher=Batcher()) + host_med = Mediator(intervention=None, info=SimpleNamespace(frame=None), batch_group=None) + host_med.channel = SocketHostChannel(host_sock) + host_med.interleaver = interleaver + interleaver.mediators = [host_med] + host_med.channel.wait_event() + + def proxy_hook(m, i, o): + # MediatorProxy: D2H -> deliver to the CPU worker -> H2D the result. + dev = o[0].device if isinstance(o, tuple) else o.device + cpu_o = tuple(x.to("cpu") if torch.is_tensor(x) else x for x in o) if isinstance(o, tuple) else o.to("cpu") + new_cpu = host_med.handle(PROVIDER, cpu_o) + if isinstance(new_cpu, tuple): + return tuple(x.to(dev) if torch.is_tensor(x) else x for x in new_cpu) + return new_cpu.to(dev) + + h = block.register_forward_hook(proxy_hook) + with torch.no_grad(): + sock = model._model(**inputs).logits + h.remove() + os.waitpid(pid, 0); host_med.channel.close() + + ok = torch.allclose(ref, sock, atol=1e-2, rtol=0) + print(f"[A gpu ] D2H/H2D socket path matches in-process GPU forward (atol 1e-2): {ok} " + f"| max|Δ|={(ref - sock).abs().max().item():.2e}") + return ok + + +# --------------------------------------------------------------------------- # +# B. Per-hook transport overhead vs activation size # +# --------------------------------------------------------------------------- # +def bench_worker(sock): + """Long-lived CPU worker: recv tensor, ×2, send back, until STOP.""" + pid = os.fork() + if pid == 0: + while True: + obj = recv_frame(sock) + if obj == "STOP": + os._exit(0) + send_frame(sock, obj * 2.0) + return pid + + +def measure(sizes, iters=20, warmup=5): + host_sock, worker_sock = socket.socketpair() + pid = bench_worker(worker_sock) + worker_sock.close() + + # In-process baseline: a thread echoes the GPU tensor BY REFERENCE (no D2H, + # no serialize, no socket) and does the same ×2 — this is what the in-process + # path costs per hook, for a fair "what does isolation add" comparison. + import queue + import threading + q_in, q_out = queue.Queue(), queue.Queue() + + def inproc_echo(): + while True: + x = q_in.get() + if x is None: + return + q_out.put(x * 2.0) + + t = threading.Thread(target=inproc_echo, daemon=True) + t.start() + + print("\n[B measure] per-hook overhead (host has GPU activation -> host has GPU result).") + print(" TOTAL = D2H + sockRTT + H2D (real end-to-end). 'ser' is the pickle portion") + print(" of sockRTT, measured separately for attribution. 'inproc' = in-process baseline.") + print(f" {'shape (b,seq,hid)':>22} {'MB':>7} | {'inproc':>8} {'D2H':>8} {'(ser)':>8} {'sockRTT':>8} " + f"{'H2D':>8} {'TOTAL':>9} (ms)") + rows = [] + for (b, seq, hid) in sizes: + act = torch.randn(b, seq, hid, device=DEV, dtype=torch.bfloat16) + mb = act.element_size() * act.nelement() / 1e6 + + def one(): + torch.cuda.synchronize() + t0 = time.perf_counter() + cpu = act.to("cpu") + torch.cuda.synchronize() + t1 = time.perf_counter() + send_frame(host_sock, cpu) # ONE serialize + socket round-trip (worker deser+op+ser) + res = recv_frame(host_sock) + t2 = time.perf_counter() + gpu = res.to(DEV) # noqa: F841 + torch.cuda.synchronize() + t3 = time.perf_counter() + return (t1 - t0, t2 - t1, t3 - t2, t3 - t0) # d2h, sockRTT, h2d, TOTAL + + def one_inproc(): + torch.cuda.synchronize() + t0 = time.perf_counter() + q_in.put(act) # reference pass — no copy + _ = q_out.get() + torch.cuda.synchronize() + return time.perf_counter() - t0 + + for _ in range(warmup): + one(); one_inproc() + acc = [0.0] * 4 + inproc_acc = 0.0 + ser_acc = 0.0 + for _ in range(iters): + acc = [a + p for a, p in zip(acc, one())] + inproc_acc += one_inproc() + ser_acc += _time_pickle(act.to("cpu")) # attribution-only; outside the TOTAL path + d2h, rtt, h2d, tot = [a / iters * 1e3 for a in acc] + inproc_ms = inproc_acc / iters * 1e3 + ser_ms = ser_acc / iters * 1e3 + rows.append((b, seq, hid, mb, tot, inproc_ms)) + print(f" {str((b, seq, hid)):>22} {mb:7.1f} | {inproc_ms:8.3f} {d2h:8.3f} {ser_ms:8.3f} {rtt:8.3f} " + f"{h2d:8.3f} {tot:9.3f}") + + send_frame(host_sock, "STOP") + os.waitpid(pid, 0); host_sock.close() + q_in.put(None) + return rows + + +def _time_pickle(cpu_tensor): + t0 = time.perf_counter() + pickle.dumps(cpu_tensor, protocol=pickle.HIGHEST_PROTOCOL) + return time.perf_counter() - t0 + + +def main(): + if not torch.cuda.is_available(): + sys.exit("Phase 5 needs a GPU (set CUDA_VISIBLE_DEVICES).") + print(f"GPU: {torch.cuda.get_device_name(0)}") + + model = LanguageModel("gpt2", device_map=DEV, dispatch=True) + inputs = model.tokenizer(PROMPT, return_tensors="pt").to(DEV) + correctness = test_gpu_correctness(model, inputs) + + sizes = [ + (1, 16, 768), # gpt2, short prompt + (1, 512, 768), # gpt2, long context + (1, 512, 4096), # ~7B hidden, 512 tokens + (1, 2048, 4096), # ~7B hidden, 2k context (one layer's activation) + (1, 2048, 8192), # ~70B hidden, 2k context + ] + rows = measure(sizes) + + # The (D) decision input: extra latency a per-layer cache of N layers would add. + big = next(r for r in rows if r[2] == 4096 and r[1] == 2048) + print(f"\n[D-input] a 1×2048×4096 bf16 activation ({big[3]:.1f} MB) costs {big[4]:.2f} ms/hook over the " + f"socket vs {big[5]:.3f} ms in-process (measured, zero-copy ref).\n " + f"A 32-layer cache => ~{big[4]*32:.0f} ms added (linear projection).") + + print("=" * 78) + print(f"PHASE 5 RESULT: {'PASS — GPU path correct; per-hook overhead measured' if correctness else 'FAIL'}") + sys.exit(0 if correctness else 1) + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/phase5b_transport_breakdown.py b/prototypes/mediator-sandbox/phase5b_transport_breakdown.py new file mode 100644 index 000000000..13ce7ea64 --- /dev/null +++ b/prototypes/mediator-sandbox/phase5b_transport_breakdown.py @@ -0,0 +1,236 @@ +#!/usr/bin/env python3 +"""Phase 5b — WHERE does the per-hook cost come from, and what fixes it. + +Phase 5 showed ~111 ms/hook for a 16.8 MB activation, "dominated by sockRTT". But +a Unix socket moves GB/s, so 16.8 MB should cost a few ms — the 100 ms must be +SERIALIZATION, not transfer. This isolates it: it (1) measures the raw socket +bandwidth, (2) breaks the pickle path into dumps/loads, then compares three +round-trip transports on the SAME CPU tensor: + + A. pickle — the naive path (pickle.dumps/​loads both ways) [what Phase 5 used] + B. raw message — a small struct header + the tensor's raw bytes over the socket, + rebuilt with torch.frombuffer (no pickle) + C. shared mem — host & worker share an mmap; only a 1-byte ready signal crosses + the socket; the bulk never travels (near zero-copy) + +CPU<->CPU only (D2H/H2D is separate and was ~3-4 ms). Worker is a forked process. +Run: PYTHONPATH=src .../hf-serve/bin/python prototypes/mediator-sandbox/phase5b_transport_breakdown.py +""" +import faulthandler +import mmap +import os +import pickle +import socket +import struct +import sys +import time + +import torch + +faulthandler.dump_traceback_later(45, exit=True) + +DT = torch.bfloat16 +SIZES = [(1, 512, 4096), (1, 2048, 4096), (1, 2048, 8192)] # 4.2, 16.8, 33.6 MB +ITERS, WARM = 30, 8 +_LEN = struct.Struct("!I") +_HDR = struct.Struct("!iiqqq") # ndim, dtype_code(0=bf16), d0, d1, d2 (fixed 3-dim here) + + +def _recvn(sock, n): + buf = bytearray(n) + view = memoryview(buf) + got = 0 + while got < n: + c = sock.recv_into(view[got:], n - got) + if c == 0: + raise EOFError + got += c + return buf + + +# ---------- A: pickle ---------- +def send_pickle(sock, obj): + b = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL) + sock.sendall(_LEN.pack(len(b)) + b) + + +def recv_pickle(sock): + (n,) = _LEN.unpack(_recvn(sock, 4)) + return pickle.loads(_recvn(sock, n)) + + +# ---------- B: raw message ---------- +def send_raw(sock, t): + t = t.contiguous() + raw = t.flatten().view(torch.uint8).numpy() # zero-copy view of the bytes + hdr = _HDR.pack(t.dim(), 0, *(list(t.shape) + [0, 0, 0])[:3]) + sock.sendall(_LEN.pack(raw.nbytes) + hdr) + sock.sendall(memoryview(raw)) # bulk: one zero-copy send + + +def recv_raw(sock): + (n,) = _LEN.unpack(_recvn(sock, 4)) + ndim, _code, d0, d1, d2 = _HDR.unpack(_recvn(sock, _HDR.size)) + buf = _recvn(sock, n) + shape = [d0, d1, d2][:ndim] + return torch.frombuffer(bytearray(buf), dtype=torch.uint8).view(DT).view(*shape) + + +# ---------- C: shared memory ---------- +def make_shm(nbytes): + return mmap.mmap(-1, nbytes) # anonymous MAP_SHARED, inherited across fork + + +def shm_write(buf, t): + t = t.contiguous() + raw = t.flatten().view(torch.uint8).numpy() + buf[: raw.nbytes] = memoryview(raw) + + +def shm_read(buf, nbytes, shape): + return torch.frombuffer(memoryview(buf)[:nbytes], dtype=torch.uint8).view(DT).view(*shape).clone() + + +# --------------------------------------------------------------------------- # +def fork_worker(fn, *a): + parent, child = socket.socketpair() + pid = os.fork() + if pid == 0: + parent.close() + fn(child, *a) + os._exit(0) + child.close() + return pid, parent + + +def bench(label, host_sock, do_round, stop): + for _ in range(WARM): + do_round() + t0 = time.perf_counter() + for _ in range(ITERS): + do_round() + dt = (time.perf_counter() - t0) / ITERS * 1e3 + stop() + return dt + + +def main(): + print(f"{'transport':<14}{'4.2 MB':>12}{'16.8 MB':>12}{'33.6 MB':>12} (ms / round-trip)") + + # raw socket bandwidth + pickle sub-breakdown, for the 16.8 MB case + mid = torch.randn(*SIZES[1], dtype=DT) + mid_bytes = mid.flatten().view(torch.uint8).numpy().nbytes + # one-way raw socket throughput + pid, hs = fork_worker(_drain, mid_bytes) + raw = mid.flatten().view(torch.uint8).numpy() + for _ in range(WARM): + hs.sendall(memoryview(raw)); hs.recv(1) + t0 = time.perf_counter() + for _ in range(ITERS): + hs.sendall(memoryview(raw)); hs.recv(1) + one_way = (time.perf_counter() - t0) / ITERS * 1e3 + hs.sendall(b""); hs.close(); os.waitpid(pid, 0) + gbps = mid_bytes / (one_way / 1e3) / 1e9 + blob = pickle.dumps(mid, protocol=pickle.HIGHEST_PROTOCOL) + dumps_ms = _t(lambda: pickle.dumps(mid, protocol=pickle.HIGHEST_PROTOCOL)) + loads_ms = _t(lambda: pickle.loads(blob)) + op_bf16 = _t(lambda: mid * 2.0) + mid_f32 = mid.float() + op_f32 = _t(lambda: mid_f32 * 2.0) + print(f"\n[probe] 16.8 MB: raw socket one-way+ack {one_way:.2f} ms ({gbps:.1f} GB/s)") + print(f"[probe] 16.8 MB: pickle.dumps {dumps_ms:.2f} ms | pickle.loads {loads_ms:.2f} ms") + print(f"[probe] 16.8 MB: ×2 CPU op — bf16 {op_bf16:.2f} ms | fp32 {op_f32:.2f} ms " + f"(this is the USER op, not transport)\n") + + # three transports across the three sizes — ECHO (no op) = PURE transport cost + for name, runner in [("A pickle", run_pickle), ("B raw-msg", run_raw), ("C shared-mem", run_shm)]: + cells = [] + for sz in SIZES: + cells.append(f"{runner(sz):>12.2f}") + print(f"{name:<14}{''.join(cells)}") + + print("\n(worker does the same ×2 CPU op in every method; only the transport differs.)") + + +def _drain(sock, nbytes): + while True: + try: + b = _recvn(sock, nbytes) + except EOFError: + return + sock.sendall(b"k") + + +def _t(fn, n=20): + for _ in range(5): + fn() + t0 = time.perf_counter() + for _ in range(n): + fn() + return (time.perf_counter() - t0) / n * 1e3 + + +def run_pickle(sz): + def worker(sock): + while True: + try: + t = recv_pickle(sock) + except EOFError: + return + send_pickle(sock, t) # echo = pure transport + pid, hs = fork_worker(worker) + act = torch.randn(*sz, dtype=DT) + + def rnd(): + send_pickle(hs, act); recv_pickle(hs) + dt = bench("A", hs, rnd, lambda: (hs.close(), os.waitpid(pid, 0))) + return dt + + +def run_raw(sz): + def worker(sock): + while True: + try: + t = recv_raw(sock) + except EOFError: + return + send_raw(sock, t) # echo = pure transport + pid, hs = fork_worker(worker) + act = torch.randn(*sz, dtype=DT) + + def rnd(): + send_raw(hs, act); recv_raw(hs) + dt = bench("B", hs, rnd, lambda: (hs.close(), os.waitpid(pid, 0))) + return dt + + +def run_shm(sz): + nbytes = torch.empty(*sz, dtype=DT).flatten().view(torch.uint8).numpy().nbytes + in_buf = make_shm(nbytes) + out_buf = make_shm(nbytes) + + def worker(sock): + while True: + sig = sock.recv(1) + if not sig: + return + t = torch.frombuffer(memoryview(in_buf)[:nbytes], dtype=torch.uint8).view(DT).view(*sz) + res = t.contiguous() # echo = pure transport + out_buf[:nbytes] = memoryview(res.flatten().view(torch.uint8).numpy()) + sock.sendall(b"k") + + pid, hs = fork_worker(worker) + act = torch.randn(*sz, dtype=DT) + + def rnd(): + shm_write(in_buf, act) + hs.sendall(b"g") + hs.recv(1) + shm_read(out_buf, nbytes, sz) + dt = bench("C", hs, rnd, lambda: (hs.close(), os.waitpid(pid, 0))) + in_buf.close(); out_buf.close() + return dt + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/phase6_jailed_worker.py b/prototypes/mediator-sandbox/phase6_jailed_worker.py new file mode 100644 index 000000000..8f7cae0ed --- /dev/null +++ b/prototypes/mediator-sandbox/phase6_jailed_worker.py @@ -0,0 +1,35 @@ +#!/usr/bin/env python3 +"""Phase 6 — jailed worker using the shared-memory + safetensors channel. + +Same as phase3's worker, but its tensor payloads ride a shared memfd (passed in +via SHM_FD) instead of being pickled over the socket. Runs INSIDE a bwrap jail. +""" +import os +import socket + + +def main(): + fd = int(os.environ["WORKER_FD"]) + shm_fd = int(os.environ["SHM_FD"]) + shm_size = int(os.environ["SHM_SIZE"]) + provider = os.environ["PROVIDER"] + sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM, fileno=fd) + + from types import SimpleNamespace + + from nnsight.intervention.interleaver import Mediator + from nnsight.intervention.transport import ShmArena, ShmSocketWorkerChannel + + arena = ShmArena.attach(shm_fd, shm_size) + med = Mediator(intervention=None, info=SimpleNamespace(frame=None), batch_group=None) + med.channel = ShmSocketWorkerChannel(sock, arena) + med.cross_invoker = False + + value = med.request(provider) + new = (value[0] * 2.0,) + tuple(value[1:]) if isinstance(value, tuple) else value * 2.0 + med.swap(provider, new) + med.end() + + +if __name__ == "__main__": + main() diff --git a/prototypes/mediator-sandbox/phase6_shm_safetensors.py b/prototypes/mediator-sandbox/phase6_shm_safetensors.py new file mode 100644 index 000000000..c5e05497e --- /dev/null +++ b/prototypes/mediator-sandbox/phase6_shm_safetensors.py @@ -0,0 +1,203 @@ +#!/usr/bin/env python3 +"""Phase 6 — shared memory + safetensors transport. + +Implements the fast path identified in phase5b and wires it into the real +`MediatorChannel`s (`ShmSocketHostChannel`/`ShmSocketWorkerChannel` in +transport.py). Tensor bulk rides a shared memfd; only a tiny control frame +crosses the socket; tensors are encoded with safetensors (safe, no pickle). + +Three parts: + 1. measure — per-hook echo round-trip: pickle (old) vs shm+safetensors (new). + 2. correct — the real Mediator protocol with the Shm channels on gpt2 → golden. + 3. jailed — same, with the worker in a bwrap jail (memfd passed in via SHM_FD). + +Run: PYTHONPATH=src .../hf-serve/bin/python prototypes/mediator-sandbox/phase6_shm_safetensors.py +""" +import os +import shutil +import socket +import subprocess +import sys +import time +from types import SimpleNamespace + +import torch + +from nnsight import LanguageModel +from nnsight.intervention.batching import Batcher +from nnsight.intervention.interleaver import Interleaver, Mediator +from nnsight.intervention.transport import ( + ShmArena, + ShmSocketHostChannel, + ShmSocketWorkerChannel, + recv_frame, + recv_shm, + send_frame, + send_shm, +) + +HERE = os.path.dirname(os.path.abspath(__file__)) +ROOT = os.path.dirname(os.path.dirname(HERE)) +SRC = os.path.join(ROOT, "src") +WORKER = os.path.join(HERE, "phase6_jailed_worker.py") +ENV_ROOT = os.path.dirname(os.path.dirname(sys.executable)) +PROMPT = "The Eiffel Tower is in the city of" +PROVIDER = "transformer.h.6.output.i0" +LAYER = 6 +ARENA = 64 << 20 # 64 MB shared region + + +# --------------------------------------------------------------------------- # +# 1. measurement: pickle vs shm+safetensors (echo = pure transport) # +# --------------------------------------------------------------------------- # +def measure(): + sizes = [(1, 512, 4096), (1, 2048, 4096), (1, 2048, 8192)] # 4.2, 16.8, 33.6 MB + print(f"{'transport':<22}{'4.2 MB':>11}{'16.8 MB':>11}{'33.6 MB':>11} (ms/round-trip)") + + def echo_pickle(sock): + while True: + try: + t = recv_frame(sock) + except EOFError: + os._exit(0) + send_frame(sock, t) + + def echo_shm(sock, fd): + arena = ShmArena.attach(fd, ARENA) + while True: + try: + t = recv_shm(sock, arena) + except EOFError: + os._exit(0) + send_shm(sock, arena, t) + + for name, use_shm in [("A pickle (old)", False), ("D shm+safetensors", True)]: + cells = [] + for sz in sizes: + act = (torch.randn(*sz, dtype=torch.bfloat16),) + hs, ws = socket.socketpair() + arena = ShmArena(ARENA) if use_shm else None + pid = os.fork() + if pid == 0: + hs.close() + (echo_shm(ws, arena.fd) if use_shm else echo_pickle(ws)) + os._exit(0) + ws.close() + rnd = ((lambda: (send_shm(hs, arena, act), recv_shm(hs, arena))) + if use_shm else (lambda: (send_frame(hs, act), recv_frame(hs)))) + for _ in range(8): + rnd() + t0 = time.perf_counter() + for _ in range(30): + rnd() + cells.append(f"{(time.perf_counter() - t0) / 30 * 1e3:>11.2f}") + hs.close(); os.waitpid(pid, 0) + if arena: + arena.close() + print(f"{name:<22}{''.join(cells)}") + + +# --------------------------------------------------------------------------- # +# 2/3. correctness via the real Mediator protocol with the Shm channels # +# --------------------------------------------------------------------------- # +def reference(model, inputs): + blk = model._model.transformer.h[LAYER] + h = blk.register_forward_hook( + lambda m, i, o: (o[0] * 2.0,) + tuple(o[1:]) if isinstance(o, tuple) else o * 2.0) + with torch.no_grad(): + out = model._model(**inputs).logits + h.remove() + return out + + +def drive_host(model, inputs, host_med): + blk = model._model.transformer.h[LAYER] + h = blk.register_forward_hook(lambda m, i, o: host_med.handle(PROVIDER, o)) + try: + with torch.no_grad(): + return model._model(**inputs).logits + finally: + h.remove() + + +def mk_host(sock, arena): + interleaver = Interleaver(mediators=[], tracer=None, batcher=Batcher()) + med = Mediator(intervention=None, info=SimpleNamespace(frame=None), batch_group=None) + med.channel = ShmSocketHostChannel(sock, arena) + med.interleaver = interleaver + interleaver.mediators = [med] + med.channel.wait_event() + return med + + +def correct_fork(model, inputs, ref): + arena = ShmArena(ARENA) + hs, ws = socket.socketpair() + pid = os.fork() + if pid == 0: + hs.close() + wa = ShmArena.attach(arena.fd, ARENA) + med = Mediator(intervention=None, info=SimpleNamespace(frame=None), batch_group=None) + med.channel = ShmSocketWorkerChannel(ws, wa) + med.cross_invoker = False + v = med.request(PROVIDER) + new = (v[0] * 2.0,) + tuple(v[1:]) if isinstance(v, tuple) else v * 2.0 + med.swap(PROVIDER, new) + med.end() + os._exit(0) + ws.close() + host = mk_host(hs, arena) + out = drive_host(model, inputs, host) + os.waitpid(pid, 0); host.channel.close(); arena.close() + ok = torch.allclose(ref, out, atol=1e-2, rtol=0) + print(f"[2 correct/fork ] gpt2 golden via Shm channel: {ok} | max|Δ|={(ref - out).abs().max():.2e}") + return ok + + +def correct_jailed(model, inputs, ref): + if shutil.which("bwrap") is None: + print("[3 correct/jail ] bwrap absent — skipped") + return True + arena = ShmArena(ARENA) + hs, ws = socket.socketpair() + os.set_inheritable(ws.fileno(), True) + os.set_inheritable(arena.fd, True) + env = { + "WORKER_FD": str(ws.fileno()), "SHM_FD": str(arena.fd), "SHM_SIZE": str(ARENA), + "PROVIDER": PROVIDER, "PYTHONPATH": SRC, "PATH": "/usr/local/bin:/usr/bin:/bin", + "HOME": "/tmp", "CUDA_VISIBLE_DEVICES": "", "HF_HUB_OFFLINE": "1", "TRANSFORMERS_OFFLINE": "1", + } + cmd = [ + "bwrap", "--unshare-all", + "--ro-bind", "/usr", "/usr", "--ro-bind", "/lib", "/lib", "--ro-bind", "/lib64", "/lib64", + "--ro-bind", "/bin", "/bin", "--ro-bind", "/etc", "/etc", + "--ro-bind", ENV_ROOT, ENV_ROOT, "--ro-bind", SRC, SRC, "--ro-bind", WORKER, WORKER, + "--proc", "/proc", "--dev", "/dev", "--tmpfs", "/tmp", "--die-with-parent", + sys.executable, WORKER, + ] + p = subprocess.Popen(cmd, pass_fds=[ws.fileno(), arena.fd], env=env) + ws.close() + host = mk_host(hs, arena) + out = drive_host(model, inputs, host) + p.wait(); host.channel.close(); arena.close() + ok = torch.allclose(ref, out, atol=1e-2, rtol=0) + print(f"[3 correct/jail ] gpt2 golden, worker JAILED, shm via passed memfd: {ok} " + f"| max|Δ|={(ref - out).abs().max():.2e}") + return ok + + +def main(): + measure() + print() + model = LanguageModel("gpt2", device_map="cpu", dispatch=True) + inputs = model.tokenizer(PROMPT, return_tensors="pt") + ref = reference(model, inputs) + ok = correct_fork(model, inputs, ref) + ok &= correct_jailed(model, inputs, ref) + print("=" * 72) + print(f"PHASE 6 RESULT: {'PASS — shm+safetensors channel correct (fork + jailed)' if ok else 'FAIL'}") + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/src/nnsight/__init__.py b/src/nnsight/__init__.py index c6e3f3c60..7fc1eecd2 100755 --- a/src/nnsight/__init__.py +++ b/src/nnsight/__init__.py @@ -54,7 +54,7 @@ except ImportError: __version__ = "unknown version" -from .intervention.tracing.globals import save +from .intervention.tracing.globals import carry, save from .ndif import * from IPython import get_ipython diff --git a/src/nnsight/intervention/_sandbox.py b/src/nnsight/intervention/_sandbox.py new file mode 100644 index 000000000..e92ecd567 --- /dev/null +++ b/src/nnsight/intervention/_sandbox.py @@ -0,0 +1,74 @@ +"""Footgun-containment seccomp filter for the isolated GPU worker (x86-64). + +Per the agreed threat model (contain *mistakes*, not a determined adversary): after +CUDA + torch are warmed and the intervention is deserialized, the worker calls +:func:`lock_down`, which installs a minimal seccomp-BPF filter making new +``open``/``openat`` (filesystem), ``socket``/``connect`` (network), and ``execve`` +(spawn) syscalls fail with EPERM. CUDA keeps working because it talks to the +already-open ``/dev/nvidia*`` fds via ioctl/mmap, and the control Pipe + CUDA-IPC +buffer use already-open fds (read/write/ioctl), not new opens. + +No external deps: the BPF program is assembled by hand and installed via +``prctl(PR_SET_SECCOMP, SECCOMP_MODE_FILTER, ...)``. Proven in +prototypes/mediator-sandbox/gpu_sandbox (test_safety: 9/9 contained). +""" +import ctypes +import struct + +# x86-64 syscall numbers +_NR = { + "open": 2, "openat": 257, "openat2": 437, + "socket": 41, "connect": 42, "execve": 59, "execveat": 322, +} +_AUDIT_ARCH_X86_64 = 0xC000003E +_RET_KILL_PROCESS = 0x80000000 +_RET_ERRNO = 0x00050000 +_RET_ALLOW = 0x7FFF0000 +_EPERM = 1 +# BPF opcodes +_LD_W_ABS = 0x20 +_JMP_JEQ_K = 0x15 +_RET_K = 0x06 +_PR_SET_NO_NEW_PRIVS = 38 +_PR_SET_SECCOMP = 22 +_SECCOMP_MODE_FILTER = 2 + + +def _build_filter(blocked): + instrs = [ + (_LD_W_ABS, 0, 0, 4), # A = arch (seccomp_data offset 4) + (_JMP_JEQ_K, 1, 0, _AUDIT_ARCH_X86_64), # if x86-64: skip the kill + (_RET_K, 0, 0, _RET_KILL_PROCESS), # else kill (block arch-bypass) + (_LD_W_ABS, 0, 0, 0), # A = syscall nr (offset 0) + ] + for nr in blocked: + instrs.append((_JMP_JEQ_K, 0, 1, nr)) # if A == nr: next else skip next + instrs.append((_RET_K, 0, 0, _RET_ERRNO | _EPERM)) + instrs.append((_RET_K, 0, 0, _RET_ALLOW)) # default: allow + return instrs + + +class _sock_fprog(ctypes.Structure): + _fields_ = [("len", ctypes.c_ushort), ("filter", ctypes.c_void_p)] + + +def lock_down(block_fs: bool = True, block_net: bool = True) -> None: + """Install the seccomp filter. Call AFTER torch/CUDA are warmed + the + intervention is deserialized (deserialization may open files).""" + blocked = [] + if block_fs: + blocked += [_NR["open"], _NR["openat"], _NR["openat2"]] + if block_net: + blocked += [_NR["socket"], _NR["connect"]] + blocked += [_NR["execve"], _NR["execveat"]] # no spawning new programs either + instrs = _build_filter(blocked) + + libc = ctypes.CDLL("libc.so.6", use_errno=True) + if libc.prctl(_PR_SET_NO_NEW_PRIVS, 1, 0, 0, 0) != 0: + raise OSError(ctypes.get_errno(), "PR_SET_NO_NEW_PRIVS failed") + + prog = b"".join(struct.pack("HBBI", *i) for i in instrs) + buf = ctypes.create_string_buffer(prog, len(prog)) + fprog = _sock_fprog(len(instrs), ctypes.cast(buf, ctypes.c_void_p)) + if libc.prctl(_PR_SET_SECCOMP, _SECCOMP_MODE_FILTER, ctypes.byref(fprog), 0, 0) != 0: + raise OSError(ctypes.get_errno(), "PR_SET_SECCOMP failed") diff --git a/src/nnsight/intervention/fastlane.py b/src/nnsight/intervention/fastlane.py new file mode 100644 index 000000000..559054193 --- /dev/null +++ b/src/nnsight/intervention/fastlane.py @@ -0,0 +1,651 @@ +"""Fast-lane safety classifier for isolated traces. + +`isolate_mediators()` runs each user intervention in a spawned GPU worker process to +contain footguns (infinite loops, OOM allocs, device-side asserts, host-object pokes, +fs/net/exec). But the worker holds a *weightless* path-only mirror of the model — its +dummy modules have no parameters and no ``forward``. The real interpretability workloads +(logit lens, steering, ablation, activation patching, attribution) all read the host +model's real weights (``F.linear(x, head.weight)``) and call its final-norm / unembed +modules, so **they cannot run in the worker at all**. The fast lane is the tier where the +real weights live: a confirmed-safe intervention runs in-process (the existing daemon- +thread path) at full speed and with full model access, and only unconfirmable code goes +to the worker. + +"Confirmed safe" is decided here by a **fail-closed, default-deny** static walk over the +*effective code* — the trace body PLUS every user closure it calls, resolved through the +frame, the function globals, and closure cells (the harness wraps real compute in +``build()`` / ``capture()`` / ``patch()`` closures, so a walk of the ``with`` block alone +would see only an opaque call). The walk emits one of three verdicts: + +- ``FAST`` — every node is on the allowlist and every call resolves to a whitelisted + op, a host module/weight access, an nnsight primitive, or a recursively + confirmed user function → run in-process. +- ``ISOLATE`` — anything unconfirmable (an unresolved global call, an import, a ``while`` + loop, an unrecoverable closure, an unknown node type) → run in the worker. +- ``REJECT`` — an introspection escape (``__globals__``/``getattr``/``eval``/…) → raise. + +The conservative default is ``ISOLATE``: the absence of proof is not proof of safety; only +explicitly whitelisted code reaches ``FAST``. + +**Threat model (load-bearing).** This is *not* an adversarial sandbox — a determined +author can defeat any in-process restriction (the pysandbox lesson), so the fast lane is +gated on ``trust="local"`` provenance and disabled for anything deserialized/remote. Under +the relaxed "contain footguns, not adversaries" model it confirms the effective code +(a) introduces no ambient authority (no import, no introspection, no unresolved global +call), (b) has no unbounded loop, (c) writes no host state, and (d) is composed only of +whitelisted ops, host-object access, and recursively confirmed user functions. OOM and +device-side asserts in pure tensor math are knowingly traded to the in-process tier (a +deployment that cannot tolerate them disables the fast lane); a wall-clock watchdog backs +the one loop footgun the static walk cannot bound (a huge ``range``). +""" +from __future__ import annotations + +import ast +import ctypes +import inspect +import textwrap +import threading +from dataclasses import dataclass, field + +# Verdict tiers. +FAST = "fast" +ISOLATE = "isolate" +REJECT = "reject" + +# Modules whose top-level functions are pure compute with no ambient authority. A call +# resolving into one of these is allowed outright. +_SAFE_MODULE_PREFIXES = ("torch", "math", "operator", "numpy") + +# torch entry points that DO touch fs/net/JIT despite living under `torch` — never fast. +_BANNED_QUALIFIED = { + "torch.load", "torch.save", "torch.hub", "torch.jit", "torch.compile", + "torch.onnx", "torch.multiprocessing", "torch.distributed", +} + +# Builtins that are pure / structural — safe to call from confirmed code. +_SAFE_BUILTINS = frozenset({ + "range", "len", "enumerate", "zip", "list", "tuple", "dict", "set", "frozenset", + "min", "max", "sum", "abs", "round", "sorted", "reversed", "map", "filter", + "float", "int", "bool", "str", "slice", "isinstance", "issubclass", "all", "any", + "print", "repr", "iter", "next", +}) + +# Names / attributes that are an introspection escape — their presence is a REJECT, not an +# isolate: trusted-local author code reaching for these is a footgun (or an attempt to +# break out), and the worker cannot run them either. Every documented in-process escape +# walks one of these (`().__class__.__subclasses__()`, `obj.__globals__`, a fetched +# builtin), which is exactly what pysandbox could not close. +_INTROSPECTION = frozenset({ + "eval", "exec", "compile", "__import__", "getattr", "setattr", "delattr", + "globals", "vars", "locals", "breakpoint", "memoryview", +}) + + +class FastLaneRejected(Exception): + """Raised when the classifier finds an introspection escape in a fast-lane-eligible + trace. The worker cannot run such code either, and under the relaxed footgun model + trusted-local author code reaching for introspection is a footgun — so fail loudly + rather than silently route it anywhere.""" + + +class FastLaneTimeout(Exception): + """Injected by the watchdog into a fast-lane intervention thread that overran the + wall-clock deadline (a runaway pure-Python loop). It is an ``Exception`` so the + intervention body's own ``except Exception`` catches it and routes it through the + normal ``mediator.exception`` path — the host re-raises it to the user, the model + server is unaffected.""" + + +class Watchdog: + """Best-effort wall-clock bound on a fast-lane intervention thread. + + The static gate bans ``while`` and unbounded iteration, so the only loop footgun that + can reach the fast lane is a huge bounded ``range`` or deep recursion. This injects a + :class:`FastLaneTimeout` into the running thread at its next bytecode if it overruns — + restoring the loop-containment guarantee that turning isolation on implies. It CANNOT + preempt a wedged native/CUDA call (only the worker-process kill can); the bounded-loop + static rule is the primary defense and this is the backstop. + """ + + def __init__(self, deadline_s: float): + self._deadline = deadline_s + self._timer = None + self._ident = None + + def arm(self, thread_ident: int) -> None: + self._ident = thread_ident + self._timer = threading.Timer(self._deadline, self._fire) + self._timer.daemon = True + self._timer.start() + + def _fire(self) -> None: + ident = self._ident + if ident is None: + return + res = ctypes.pythonapi.PyThreadState_SetAsyncExc( + ctypes.c_long(ident), ctypes.py_object(FastLaneTimeout) + ) + if res > 1: + # affected more than the target thread — undo to avoid corrupting others + ctypes.pythonapi.PyThreadState_SetAsyncExc(ctypes.c_long(ident), None) + + def disarm(self) -> None: + if self._timer is not None: + self._timer.cancel() + self._timer = None + self._ident = None + + +@dataclass +class Verdict: + """The classifier's decision for one mediator's intervention.""" + + tier: str # FAST | ISOLATE | REJECT + reason: str + differentiate: bool = False # the effective code opens a `with x.backward():` + touches_host_weights: bool = False + in_place: bool = False # an in-place write on a delivered boundary value + _seen: set = field(default_factory=set, repr=False) # visited code objects (recursion guard) + + @property + def fast(self) -> bool: + return self.tier == FAST + + +# Sentinel raised internally to short-circuit the walk on the first disqualifying node. +class _Stop(Exception): + def __init__(self, tier: str, reason: str): + self.tier = tier + self.reason = reason + + +def classify(mediator) -> Verdict: + """Classify a mediator's intervention for the fast lane. Always returns a Verdict + (never raises); a REJECT tier is surfaced as a raised error by the caller.""" + v = Verdict(tier=FAST, reason="confirmed: effective code is all whitelisted ops") + try: + nodes = _root_nodes(mediator) + if nodes is None: + return Verdict(ISOLATE, "could not recover the trace body source/AST") + ns = _root_namespace(mediator) + walker = _Walker(v) + walker.walk(nodes, ns, set()) + except _Stop as s: + return Verdict(s.tier, s.reason, v.differentiate, v.touches_host_weights, v.in_place) + except Exception as e: # noqa: BLE001 — any analysis failure is fail-closed + return Verdict(ISOLATE, f"classifier could not confirm safety: {type(e).__name__}: {e}") + return v + + +def classify_callable(fn) -> Verdict: + """Classify an arbitrary callable's body with the same engine `classify` uses on a + trace body. The effective-code walk resolves names through ``fn``'s globals and + closure cells. Used by the fast-lane tests to confirm verdicts on varied/renamed + module structures without standing up a model.""" + v = Verdict(tier=FAST, reason="confirmed: effective code is all whitelisted ops") + try: + _Walker(v)._recurse_fn(fn, depth=0) + except _Stop as s: + return Verdict(s.tier, s.reason, v.differentiate, v.touches_host_weights, v.in_place) + except Exception as e: # noqa: BLE001 + return Verdict(ISOLATE, f"classifier could not confirm safety: {type(e).__name__}: {e}") + return v + + +def _root_nodes(mediator): + """The statements of the trace body, preferring the live AST node, re-parsing the + captured source otherwise.""" + info = mediator.info + node = getattr(info, "node", None) + if isinstance(node, ast.With): + return list(node.body) + src = info.source + if not src: + return None + text = textwrap.dedent("".join(src)) + tree = ast.parse(text) + body = tree.body + if len(body) == 1 and isinstance(body[0], ast.With): + return list(body[0].body) + return body + + +def _root_namespace(mediator) -> dict: + """Name → object map for the trace body: the capturing frame's locals/globals plus + the compiled intervention's module globals.""" + ns = {} + fn = mediator.intervention + g = getattr(fn, "__globals__", None) + if isinstance(g, dict): + ns.update(g) + frame = getattr(mediator.info, "frame", None) + fl = getattr(frame, "f_globals", None) + if isinstance(fl, dict): + ns.update(fl) + fl = getattr(frame, "f_locals", None) + if isinstance(fl, dict): + ns.update(fl) + return ns + + +def _fn_namespace(fn) -> dict: + """Name → object map for a resolved user function: its globals plus closure cells.""" + ns = {} + g = getattr(fn, "__globals__", None) + if isinstance(g, dict): + ns.update(g) + code = getattr(fn, "__code__", None) + closure = getattr(fn, "__closure__", None) + if code is not None and closure: + for name, cell in zip(code.co_freevars, closure): + try: + ns[name] = cell.cell_contents + except ValueError: + pass # an empty cell (recursive def not yet bound) + return ns + + +class _Walker: + """Default-deny AST walk. Raises ``_Stop`` on the first ISOLATE/REJECT node; falling + off the end means every node was confirmed FAST.""" + + MAX_DEPTH = 12 + + def __init__(self, verdict: Verdict): + self.v = verdict + + # --- entry points --------------------------------------------------------- + def walk(self, nodes, ns: dict, local_names: set, depth: int = 0): + for n in nodes: + self._stmt(n, ns, local_names, depth) + + def _recurse_fn(self, fn, depth: int): + """Walk a resolved user function's body. Unrecoverable source → ISOLATE.""" + code = getattr(fn, "__code__", None) + if code is None or code in self.v._seen: + return + if depth >= self.MAX_DEPTH: + raise _Stop(ISOLATE, "intervention call graph is deeper than the fast-lane bound") + self.v._seen.add(code) + try: + src = textwrap.dedent(inspect.getsource(fn)) + except (OSError, TypeError): + raise _Stop(ISOLATE, f"could not recover source for `{getattr(fn, '__name__', fn)}`") + try: + tree = ast.parse(src) + except SyntaxError: + raise _Stop(ISOLATE, f"could not parse source for `{getattr(fn, '__name__', fn)}`") + target = tree.body[0] + ns = _fn_namespace(fn) + # The function's own parameters are locals we cannot resolve to objects (they are + # bound from the call site) — record them so a call/use of a param is treated as a + # host-object access under trust=local, not as an unknown global. + params = set() + node = target + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + params = _arg_names(node.args) + self.walk(node.body, ns, params, depth + 1) + elif isinstance(node, ast.Expr) and isinstance(node.value, ast.Lambda): + lam = node.value + params = _arg_names(lam.args) + self._expr(lam.body, ns, params, depth + 1) + else: + # getsource gave us the surrounding statement (common for lambdas passed as + # args); find the first Lambda anywhere in it. + lam = next((d for d in ast.walk(tree) if isinstance(d, ast.Lambda)), None) + if lam is None: + raise _Stop(ISOLATE, f"could not isolate the body of `{getattr(fn, '__name__', fn)}`") + params = _arg_names(lam.args) + self._expr(lam.body, ns, params, depth + 1) + + # --- statements ----------------------------------------------------------- + def _stmt(self, n, ns, loc, depth): + if isinstance(n, (ast.Assign, ast.AnnAssign, ast.AugAssign)): + self._check_targets(n, ns, loc, depth) + if n.value is not None: + self._expr(n.value, ns, loc, depth) + # bind assigned simple names as locals (so later uses are not "unknown global") + for t in _assigned_names(n): + loc.add(t) + elif isinstance(n, ast.Expr): + self._expr(n.value, ns, loc, depth) + elif isinstance(n, ast.Return): + if n.value is not None: + self._expr(n.value, ns, loc, depth) + elif isinstance(n, ast.If): + self._expr(n.test, ns, loc, depth) + self.walk(n.body, ns, loc, depth) + self.walk(n.orelse, ns, loc, depth) + elif isinstance(n, ast.For): + # bounded by its iterable; an unbounded generator would be a user fn we walk, + # and a huge `range` is backed by the wall-clock watchdog. + for t in _target_names(n.target): + loc.add(t) + self._expr(n.iter, ns, loc, depth) + self.walk(n.body, ns, loc, depth) + self.walk(n.orelse, ns, loc, depth) + elif isinstance(n, ast.With): + for item in n.items: + self._with_item(item, ns, loc, depth) + self.walk(n.body, ns, loc, depth) + elif isinstance(n, ast.Raise): + # `raise ValueError(...)` is a normal guard; walk its expressions. + for child in (n.exc, n.cause): + if child is not None: + self._expr(child, ns, loc, depth) + elif isinstance(n, ast.Pass): + pass + elif isinstance(n, (ast.FunctionDef, ast.Lambda)): + # a def/lambda inside the body is a value; bind its name, walk when called + if isinstance(n, ast.FunctionDef): + loc.add(n.name) + loc |= _arg_names(n.args) + self.walk(n.body, ns, loc, depth) + elif isinstance(n, ast.While): + raise _Stop(ISOLATE, "a `while` loop cannot be statically bounded — isolated") + elif isinstance(n, (ast.Import, ast.ImportFrom)): + raise _Stop(ISOLATE, "an `import` introduces ambient authority — isolated") + elif isinstance(n, (ast.Global, ast.Nonlocal)): + raise _Stop(REJECT, "`global`/`nonlocal` rebinds shared state — rejected") + elif isinstance(n, ast.Try): + raise _Stop(ISOLATE, "a `try` block is not confirmable — isolated") + else: + raise _Stop(ISOLATE, f"unsupported statement `{type(n).__name__}` — isolated") + + def _with_item(self, item, ns, loc, depth): + ctx = item.context_expr + # detect `with .backward():` (closure-aware — this is the fix for the old + # `.backward(` substring that missed backwards hidden in a closure). + if isinstance(ctx, ast.Call) and isinstance(ctx.func, ast.Attribute) \ + and ctx.func.attr == "backward": + self.v.differentiate = True + self._expr(ctx, ns, loc, depth) + if item.optional_vars is not None: + for t in _target_names(item.optional_vars): + loc.add(t) + + def _check_targets(self, n, ns, loc, depth): + targets = n.targets if isinstance(n, ast.Assign) else [n.target] + for t in targets: + self._store_target(t, ns, loc, depth) + + def _store_target(self, t, ns, loc, depth): + if isinstance(t, ast.Name): + return # binding a local name — fine + if isinstance(t, (ast.Tuple, ast.List)): + for e in t.elts: + self._store_target(e, ns, loc, depth) + return + if isinstance(t, ast.Subscript): + # `hidden[:] = ...` — in-place write into a delivered/derived tensor. Correct + # in-process (shared memory); it is the steering default and is *safer* on the + # fast lane than under isolation (where clone-on-receive makes it a silent no-op). + self.v.in_place = True + self._expr(t.value, ns, loc, depth) + return + if isinstance(t, ast.Attribute): + # `x.output = ...` / `x.input = ...` is the nnsight boundary write (a SWAP), + # allowed. Any other attribute store mutates host state visible to sibling + # mediators — isolate. + if t.attr in ("output", "input", "inputs", "grad"): + self._expr(t.value, ns, loc, depth) + return + if t.attr.startswith("_"): + raise _Stop(REJECT, f"writing dunder/private attribute `{t.attr}` — rejected") + raise _Stop(ISOLATE, f"writing host attribute `.{t.attr}` mutates shared state — isolated") + # computed target (e.g. starred) — be conservative + raise _Stop(ISOLATE, f"unsupported assignment target `{type(t).__name__}` — isolated") + + # --- expressions ---------------------------------------------------------- + def _expr(self, e, ns, loc, depth): + if isinstance(e, ast.Call): + self._call(e, ns, loc, depth) + elif isinstance(e, ast.Attribute): + self._attribute(e, ns, loc, depth) + elif isinstance(e, ast.Name): + if e.id in _INTROSPECTION: + raise _Stop(REJECT, f"name `{e.id}` is an introspection escape — rejected") + elif isinstance(e, ast.Subscript): + self._subscript(e, ns, loc, depth) + elif isinstance(e, ast.Constant): + pass + elif isinstance(e, (ast.BinOp, ast.UnaryOp, ast.BoolOp, ast.Compare)): + for child in ast.iter_child_nodes(e): + if isinstance(child, ast.expr): + self._expr(child, ns, loc, depth) + elif isinstance(e, (ast.List, ast.Tuple, ast.Set)): + for elt in e.elts: + self._expr(elt, ns, loc, depth) + elif isinstance(e, ast.Dict): + for k in e.keys: + if k is not None: + self._expr(k, ns, loc, depth) + for val in e.values: + self._expr(val, ns, loc, depth) + elif isinstance(e, (ast.ListComp, ast.SetComp, ast.GeneratorExp, ast.DictComp)): + self._comprehension(e, ns, loc, depth) + elif isinstance(e, ast.Slice): + for child in (e.lower, e.upper, e.step): + if child is not None: + self._expr(child, ns, loc, depth) + elif isinstance(e, ast.IfExp): + self._expr(e.test, ns, loc, depth) + self._expr(e.body, ns, loc, depth) + self._expr(e.orelse, ns, loc, depth) + elif isinstance(e, ast.Starred): + self._expr(e.value, ns, loc, depth) + elif isinstance(e, (ast.JoinedStr,)): + for val in e.values: + if isinstance(val, ast.FormattedValue): + self._expr(val.value, ns, loc, depth) + elif isinstance(e, ast.Lambda): + inner = loc | _arg_names(e.args) + self._expr(e.body, ns, inner, depth) + elif isinstance(e, ast.Constant): + pass + else: + raise _Stop(ISOLATE, f"unsupported expression `{type(e).__name__}` — isolated") + + def _comprehension(self, e, ns, loc, depth): + inner = set(loc) + for gen in e.generators: + for t in _target_names(gen.target): + inner.add(t) + self._expr(gen.iter, ns, inner, depth) + for cond in gen.ifs: + self._expr(cond, ns, inner, depth) + if isinstance(e, ast.DictComp): + self._expr(e.key, ns, inner, depth) + self._expr(e.value, ns, inner, depth) + else: + self._expr(e.elt, ns, inner, depth) + + def _subscript(self, e, ns, loc, depth): + # a dunder/private string key is an introspection escape (`obj["__globals__"]`). + sl = e.slice + key = sl.value if isinstance(sl, ast.Constant) else None + if isinstance(key, str) and key.startswith("_"): + raise _Stop(REJECT, f"subscript key `{key}` is an introspection escape — rejected") + self._expr(e.value, ns, loc, depth) + if isinstance(sl, ast.expr): + self._expr(sl, ns, loc, depth) + + def _attribute(self, e, ns, loc, depth): + # reading an attribute is fine (incl. `.weight`, `.output`, `.shape`); reading a + # dunder/introspection attribute is an escape. + if e.attr in _INTROSPECTION or (e.attr.startswith("__") and e.attr.endswith("__")): + raise _Stop(REJECT, f"attribute `.{e.attr}` is an introspection escape — rejected") + if e.attr == "weight" or e.attr == "bias": + self.v.touches_host_weights = True + self._expr(e.value, ns, loc, depth) + + def _call(self, e, ns, loc, depth): + # walk args first (they may themselves disqualify) + for a in e.args: + self._expr(a, ns, loc, depth) + for kw in e.keywords: + self._expr(kw.value, ns, loc, depth) + + func = e.func + if isinstance(func, ast.Attribute): + self._call_attribute(func, ns, loc, depth) + elif isinstance(func, ast.Name): + self._call_name(func.id, ns, loc, depth) + else: + # a computed callable: `(a or b)(x)`, `factory()(x)` — unconfirmable target + self._expr(func, ns, loc, depth) + raise _Stop(ISOLATE, "call to a computed/dynamic target — isolated") + + def _call_attribute(self, func: ast.Attribute, ns, loc, depth): + attr = func.attr + if attr in _INTROSPECTION or (attr.startswith("__") and attr.endswith("__")): + raise _Stop(REJECT, f"method `.{attr}()` is an introspection escape — rejected") + # is this a banned qualified op like torch.load / torch.save / torch.hub.*? + qual = _attr_chain(func) + if qual is not None: + for banned in _BANNED_QUALIFIED: + if qual == banned or qual.startswith(banned + "."): + raise _Stop(ISOLATE, f"`{qual}` touches fs/net/JIT — isolated") + # if it resolves into a safe module, fine; if into a banned module, isolate + head = qual.split(".", 1)[0] + obj = _lookup(head, ns, loc) + if inspect.ismodule(obj) and _module_is_safe(obj): + self._expr(func.value, ns, loc, depth) + return + # otherwise a method call on a receiver (tensor/envoy/host object). Method-name + # is not an escape (checked above); the receiver is walked. Allowed under + # trust=local — `.clone()`, `.save()`, `.sum()`, `.backward()`, `.to(...)` etc. + self._expr(func.value, ns, loc, depth) + + def _call_name(self, name: str, ns, loc, depth): + if name in _INTROSPECTION: + raise _Stop(REJECT, f"`{name}()` is an introspection escape — rejected") + if name in _SAFE_BUILTINS: + return + if name in loc: + return # a local/param: a host object or fn passed in at a walked call site + obj = _lookup(name, ns, None) + if obj is _MISSING: + raise _Stop(ISOLATE, f"call to unresolved name `{name}` — unknown authority, isolated") + kind = _classify_obj(obj) + if kind == "op" or kind == "nnsight" or kind == "host" or kind == "builtin_ok": + if kind == "host": + self.v.touches_host_weights = self.v.touches_host_weights # host call is fine + return + if kind == "banned": + raise _Stop(REJECT, f"`{name}()` is an introspection escape — rejected") + if kind == "userfn": + self._recurse_fn(obj, depth) + return + # unknown object type bound to a resolvable name — be conservative + raise _Stop(ISOLATE, f"call to `{name}` ({type(obj).__name__}) is not confirmable — isolated") + + +# --------------------------------------------------------------------------- # +# object / name resolution helpers # +# --------------------------------------------------------------------------- # +_MISSING = object() + + +def _lookup(name, ns: dict, loc): + if loc is not None and name in loc: + return _MISSING # a local: caller treats specially + if name in ns: + return ns[name] + import builtins + if hasattr(builtins, name): + return getattr(builtins, name) + return _MISSING + + +def _module_is_safe(mod) -> bool: + modname = getattr(mod, "__name__", "") + return any(modname == p or modname.startswith(p + ".") for p in _SAFE_MODULE_PREFIXES) + + +def _classify_obj(obj) -> str: + import builtins + if obj is _MISSING: + return "unknown" + # introspection builtins resolved as objects + if getattr(obj, "__name__", None) in _INTROSPECTION and inspect.isbuiltin(obj): + return "banned" + mod = getattr(obj, "__module__", None) or "" + if mod == "builtins": + nm = getattr(obj, "__name__", None) + if nm in _INTROSPECTION: + return "banned" + if nm in _SAFE_BUILTINS: + return "builtin_ok" + return "unknown" # an unlisted builtin (open/input/…) → isolate, not reject + # pure-compute libraries + if any(mod == p or mod.startswith(p + ".") for p in _SAFE_MODULE_PREFIXES): + return "op" + # nnsight primitives (Envoy, eproperty, tracer methods, save, ...) + if mod == "nnsight" or mod.startswith("nnsight."): + return "nnsight" + # a host object: Envoy / nn.Module / Tensor / Backend — calling it runs the real model + if _is_host_object(obj): + return "host" + # any other user-defined callable → recurse into its source + if inspect.isfunction(obj) or inspect.ismethod(obj) or isinstance(obj, type(_lookup)): + return "userfn" + if inspect.isfunction(getattr(obj, "__call__", None)): + return "userfn" + return "unknown" + + +def _is_host_object(obj) -> bool: + try: + import torch.nn as nn + import torch + if isinstance(obj, (nn.Module, torch.Tensor)): + return True + except Exception: # noqa: BLE001 + pass + cls = type(obj) + chain = {c.__name__ for c in cls.__mro__} + return bool(chain & {"Envoy", "Backend"}) + + +def _attr_chain(node): + """`torch.nn.functional.linear` → that dotted string; None if not a pure attr chain.""" + parts = [] + cur = node + while isinstance(cur, ast.Attribute): + parts.append(cur.attr) + cur = cur.value + if isinstance(cur, ast.Name): + parts.append(cur.id) + return ".".join(reversed(parts)) + return None + + +def _arg_names(args: ast.arguments) -> set: + names = set() + for group in (args.posonlyargs, args.args, args.kwonlyargs): + for a in group: + names.add(a.arg) + if args.vararg: + names.add(args.vararg.arg) + if args.kwarg: + names.add(args.kwarg.arg) + return names + + +def _target_names(target) -> set: + names = set() + for n in ast.walk(target): + if isinstance(n, ast.Name): + names.add(n.id) + return names + + +def _assigned_names(n) -> set: + targets = n.targets if isinstance(n, ast.Assign) else [n.target] + names = set() + for t in targets: + for x in ast.walk(t): + if isinstance(x, ast.Name): + names.add(x.id) + return names diff --git a/src/nnsight/intervention/hooks.py b/src/nnsight/intervention/hooks.py index 9ecafca83..72b56ec5b 100644 --- a/src/nnsight/intervention/hooks.py +++ b/src/nnsight/intervention/hooks.py @@ -151,7 +151,12 @@ def add_ordered_hook(module: torch.nn.Module, hook: Callable, type: str) -> Any: return handle -def input_hook(mediator: Mediator, module: torch.nn.Module, path: str) -> Any: +def input_hook( + mediator: Mediator, + module: torch.nn.Module, + path: str, + iteration: Optional[int] = None, +) -> Any: """Register a one-shot forward pre-hook for a mediator on a module. The hook target iteration is captured at registration time: @@ -188,11 +193,12 @@ def input_hook(mediator: Mediator, module: torch.nn.Module, path: str) -> Any: """ handle = None - iteration = ( - mediator.iteration - if mediator.iteration is not None - else mediator.iteration_tracker[path] - ) + if iteration is None: + iteration = ( + mediator.iteration + if mediator.iteration is not None + else mediator.iteration_tracker[path] + ) def hook(module: torch.nn.Module, args: Any, kwargs: Any) -> Any: @@ -221,7 +227,12 @@ def hook(module: torch.nn.Module, args: Any, kwargs: Any) -> Any: return handle -def output_hook(mediator: Mediator, module: torch.nn.Module, path: str) -> Any: +def output_hook( + mediator: Mediator, + module: torch.nn.Module, + path: str, + iteration: Optional[int] = None, +) -> Any: """Register a one-shot forward hook for a mediator on a module. Behaves identically to :func:`input_hook` but intercepts the module's @@ -233,17 +244,22 @@ def output_hook(mediator: Mediator, module: torch.nn.Module, path: str) -> Any: mediator: The mediator requesting this hook. module: The PyTorch module to hook. path: The provider path prefix (e.g. ``"model.layer.0.output"``). + iteration: Explicit target step. ``None`` (default) resolves it from + ``mediator.iteration``/``iteration_tracker`` (in-process). The isolated + path (host-side hook registration) passes the step parsed from the worker's requester + string, since the worker's tracker lives in another process. Returns: A :class:`~torch.utils.hooks.RemovableHandle` for the registered hook. """ handle = None - iteration = ( - mediator.iteration - if mediator.iteration is not None - else mediator.iteration_tracker[path] - ) + if iteration is None: + iteration = ( + mediator.iteration + if mediator.iteration is not None + else mediator.iteration_tracker[path] + ) def hook(module: torch.nn.Module, _, output: Any) -> Any: diff --git a/src/nnsight/intervention/interleaver.py b/src/nnsight/intervention/interleaver.py index e80ecfb43..2356a446b 100755 --- a/src/nnsight/intervention/interleaver.py +++ b/src/nnsight/intervention/interleaver.py @@ -355,6 +355,17 @@ class Events(Enum): EXCEPTION = "exception" # Signal that an exception occurred SKIP = "skip" # Signal that an operation should be skipped BARRIER = "barrier" # Signal that a barrier should be set + CACHE = "cache" # Register a tracer.cache() on the host's real modules (isolation) + BACKWARD = "backward" # Run the backward pass on the host's real graph (isolation) + UNEMBED = "unembed" # Run norm+unembed on the host's real weights (isolation) + + +# A ``tracer.cache()`` placeholder is a custom ``CacheDict`` object that is NOT in the +# worker->host value algebra (only data crosses, never live objects). The worker ships its +# token as a plain-dict marker ``{_ISO_CACHE_TAG: token}`` instead; the host swaps in its +# own forward-filled CacheDict by token in ``handle_end_event``. (Same shape as the +# EXCEPTION ``(type, message)`` sentinel — a value standing in for an object.) +_ISO_CACHE_TAG = "__nnsight_iso_cache__" class Cancelation(Exception): @@ -463,6 +474,16 @@ def initialize( self.current: Mediator = None + # Isolated BARRIER: workers can't count across processes, so each sends the + # target count and the host accumulates participant names here, coordinating + # when all have arrived. Unused on the in-process path. + self._barrier_acc: set = set() + + # Isolated cross_invoker: per-worker frames aren't shared across processes, so + # workers push their locals here and pull the merged store back (host-mediated + # replacement for the shared-frame push/pull). Unused on the in-process path. + self._xinvoke_store: dict = {} + def cancel(self): """Cancel all mediators / intervention threads. @@ -712,7 +733,7 @@ def check_dangling_mediators(self): for mediator in self.mediators: if mediator.alive: - requested_event, requester = mediator.event_queue.get() + requested_event, requester = mediator.channel.get_event() if isinstance(requester, tuple): requester = requester[0] @@ -771,6 +792,91 @@ def __deepcopy__(self, memo): return self +class MediatorChannel: + """Bidirectional, one-event-in-flight handoff between the worker thread (the + intervention fn) and the main thread (the model forward pass). + + The six :class:`Events` (VALUE/SWAP/SKIP/BARRIER/END/EXCEPTION) ride on top of + this channel unchanged. Today the only implementation is in-process + (:class:`InProcessChannel`, two lock-based one-slot queues); the mediator + isolation refactor adds a socket-backed subclass so the worker can run in a + separate sandbox process. See docs/developing/mediator-isolation-harness-plan.md. + + Direction conventions: + - ``*_event`` methods carry worker -> main messages (an ``(Events, requester)`` pair). + - ``*_response`` methods carry main -> worker replies (the value for a pending event). + + Only one event is ever in flight, which is what enforces the "access modules in + forward-pass order" contract. + """ + + # --- worker -> main (event direction) --- + def put_event(self, item: Any) -> None: + raise NotImplementedError + + def restore_event(self, item: Any) -> None: + """Stage an event back without releasing a waiter (used when a provider + fires that the worker is not currently requesting).""" + raise NotImplementedError + + def get_event(self) -> Any: + raise NotImplementedError + + def wait_event(self) -> None: + raise NotImplementedError + + @property + def has_event(self) -> bool: + raise NotImplementedError + + # --- main -> worker (response direction) --- + def put_response(self, value: Any) -> None: + raise NotImplementedError + + def get_response(self) -> Any: + raise NotImplementedError + + def wait_response(self) -> None: + raise NotImplementedError + + +class InProcessChannel(MediatorChannel): + """The original transport: two :class:`Mediator.Value` one-slot lock queues. + + Behaviour-identical to the pre-seam ``event_queue``/``response_queue`` pair — + this is the default channel and the regression baseline for the seam. + """ + + def __init__(self) -> None: + self._event = Mediator.Value() + self._response = Mediator.Value() + + def put_event(self, item: Any) -> None: + self._event.put(item) + + def restore_event(self, item: Any) -> None: + self._event.restore(item) + + def get_event(self) -> Any: + return self._event.get() + + def wait_event(self) -> None: + self._event.wait() + + @property + def has_event(self) -> bool: + return self._event.has_value + + def put_response(self, value: Any) -> None: + self._response.put(value) + + def get_response(self) -> Any: + return self._response.get() + + def wait_response(self) -> None: + self._response.wait() + + class Mediator: """ Mediates between the model execution and a single intervention function. @@ -785,8 +891,7 @@ class Mediator: info (Tracer.Info): Information about the tracing context associated with this mediator name (Optional[str]): Optional name for the mediator batch_group (Optional[List[int]]): Optional batch group for the mediator to determine which slice of tensors are being intervened on - event_queue (SimpleQueue): Where the mediator (worker thread) puts events to be processed by the interleaver (main thread). Will only ever have 1 or 0 items in the queue. - response_queue (SimpleQueue): Where the interleaver (main thread) puts responses to events, to then be processed by the mediator (worker thread). Will only ever have 1 or 0 items in the queue. + channel (MediatorChannel): The bidirectional one-event-in-flight handoff between the mediator (worker thread) and the interleaver (main thread). The worker puts events for the interleaver to process and the interleaver puts responses back; only ever 1 or 0 messages in flight in each direction. ``InProcessChannel`` by default; swappable for a socket-backed channel to run the worker in an isolated process. worker (Thread): The thread that runs the intervention function history (Set[str]): A set of providers that have been seen by the mediator. Used to detect out of order interventions. iteration_tracker (Dict[str, int]): Per-provider-path counter maintained by @@ -875,11 +980,47 @@ def __init__( self.interleaver = None - self.event_queue = Mediator.Value() - self.response_queue = Mediator.Value() + self.channel: MediatorChannel = InProcessChannel() self.worker = None + # Host-side handle (_PooledWorker) to an isolated worker process (set by + # isolation.acquire_isolated_worker). ``None`` => in-process (default). + self._iso = None + + # Isolated tracer.cache(): token -> host Cache. The worker ships a CACHE event; + # the host registers cache hooks here on the REAL modules and swaps the host + # CacheDict in for the worker's empty placeholder when saves are injected. + self._iso_caches: dict = {} + + # Isolated `with tensor.backward()`: when the trace uses backward, the host keeps + # a reference to each delivered (real, on-graph) activation tensor keyed by its + # requester string, so handle_backward_event can run the real backward and read + # gradients off the host's graph. ``_iso_backward`` gates the retention so a + # non-backward isolated trace pays nothing. + self._iso_backward = False + self._iso_grad_reals: dict = {} + # Grad-through-swap: each isolated SWAP installs a worker-computed value as a host + # *leaf* (clone-on-receive strips grad_fn), severing the host graph at the seam. + # Keep the (tagged, requires_grad) swap leaf per requester so handle_backward_event + # can return dL/d(swap leaf); the worker backprops that through its swap tape and + # re-seeds the pre-swap graph (a second BACKWARD round). Gated on ``_iso_backward``. + self._iso_grad_swaps: dict = {} + + # Fast lane: the static safety verdict (set in start when fast_lane is on) and the + # wall-clock watchdog for a confirmed-safe in-process run. + self._fastlane_verdict = None + self._fastlane_watchdog = None + + # True only inside an isolated worker process (set by _run_one_job). Lets + # cross-process-aware logic (e.g. Barrier) know it can't count locally. + self._isolated_worker = False + + # Host side of an isolated mediator: the final tagged requester of the + # event currently being handled, shipped back to the worker on the + # response's meta piggyback (the worker keys backward provenance off it). + self._iso_last_tag = None + self.skip_container = None self.history = set() @@ -887,6 +1028,12 @@ def __init__( self.hooks: List[Any] = list() self.iteration_tracker = defaultdict(int) self.iteration = 0 + # Number of `tracer.iter[...]` scopes currently open (IteratorTracer + # increments/decrements around each loop). Read by the isolated worker's + # interleaver stub to decide whether a requester's iteration must be + # resolved on the host (in-loop: the host owns the step counter and the + # pin-relaxation state) or locally (out-of-loop: the pin is authoritative). + self.iter_depth = 0 self.all_stop: Optional[int] = stop self.args = list() self.cross_invoker = None @@ -969,24 +1116,49 @@ def start(self, interleaver: Interleaver): else: _caller_stream = None - _intervention = self.intervention - _args = (self, self.info, *self.args) + from .isolation import isolation_state + + if isolation_state()["on"]: + # Three execution tiers under isolation. The fast lane classifies each + # mediator once: a CONFIRMED-safe intervention runs IN-PROCESS (full model + + # weights, no worker, no per-hook channel) — the ONLY tier that can run the + # weight-reading interp majority, since the worker holds weightless dummy + # modules — while the unconfirmable remainder isolates and an introspection + # escape is rejected. The classifier is a footgun selector, never a malice + # boundary (see fastlane.py); it only moves a mediator isolate -> in-process. + from .fastlane import FAST, REJECT, FastLaneRejected, Watchdog + from .isolation import ( + acquire_isolated_worker, + classify_for_fast_lane, + fast_lane_enabled, + ) + + tier = "isolate" + if fast_lane_enabled(): + verdict = classify_for_fast_lane(self) + self._fastlane_verdict = verdict + if verdict.tier == REJECT: + raise FastLaneRejected( + f"intervention cannot be confirmed for the in-process fast lane: " + f"{verdict.reason}. Rewrite without introspection, or run without " + f"isolate_mediators()." + ) + tier = verdict.tier - def _worker_target(): - if _caller_stream is not None: - torch.cuda.set_stream(_caller_stream) - _intervention(*_args) - - # Start the worker thread. - self.worker = Thread( - target=_worker_target, - daemon=True, - name=self.name, - ) + if tier == FAST: + deadline = isolation_state()["opts"].fast_lane_timeout + self._fastlane_watchdog = Watchdog(deadline) + self._run_in_process(_caller_stream, watchdog=self._fastlane_watchdog) + else: + # Isolated path: run the intervention in an isolated GPU worker process + # (warm pooled worker when pool_size>0, else a cold one-shot worker). + # Sets self.channel (host end), self.worker (the process), self._iso. + acquire_isolated_worker(self) + self.interleaver.current = self + else: + self._run_in_process(_caller_stream) - self.interleaver.current = self - self.worker.start() - self.event_queue.wait() + self.channel.wait_event() # Handle the first event for each mediator to clear mediators that already ended. try: @@ -996,6 +1168,31 @@ def _worker_target(): self.interleaver.current = None + def _run_in_process(self, caller_stream, watchdog=None): + """Launch the intervention in a daemon thread (the in-process path, shared by the + isolation-off default and the confirmed-safe fast lane). ``watchdog`` (fast lane + only) bounds a runaway pure-Python loop: it is armed on the thread after start and + disarmed when the body returns (the ``finally``) — and again at ``cancel`` as a + backstop. The body's own ``try/except`` (invoker.compile) routes an injected + :class:`fastlane.FastLaneTimeout` through ``mediator.exception`` like any error.""" + _intervention = self.intervention + _args = (self, self.info, *self.args) + + def _worker_target(): + try: + if caller_stream is not None: + torch.cuda.set_stream(caller_stream) + _intervention(*_args) + finally: + if watchdog is not None: + watchdog.disarm() + + self.worker = Thread(target=_worker_target, daemon=True, name=self.name) + self.interleaver.current = self + self.worker.start() + if watchdog is not None: + watchdog.arm(self.worker.ident) + ### Provider Methods ### def cancel(self): @@ -1004,14 +1201,40 @@ def cancel(self): self.history = set() self.iteration_tracker = defaultdict(int) self.iteration = 0 + self.iter_depth = 0 self.worker = None - - if self.event_queue.has_value: + # Retained on-graph activations (isolated backward) pin the autograd graph; + # drop them at trace end rather than waiting for the mediator to be GC'd. + self._iso_grad_reals = {} + self._iso_grad_swaps = {} + # Disarm the fast-lane watchdog (backstop to the thread-target finally) so a + # generous deadline can't fire into an unrelated later computation. + if self._fastlane_watchdog is not None: + self._fastlane_watchdog.disarm() + self._fastlane_watchdog = None + + # If the worker is still mid-protocol, unwind it with a Cancelation. For the + # isolated channel the get_event/put_response calls are host-local — they do NOT + # drain the worker's subsequent EXCEPTION from the cross-process pipe — so the + # pipe is left unbalanced; such a worker must NOT be recycled. Mark it dirty so + # release() retires (SIGKILLs) it. + dirty = False + if self.channel.has_event: self.handle() - if self.event_queue.has_value: - self.event_queue.get() - self.response_queue.put(Cancelation()) - self.event_queue.get() + if self.channel.has_event: + self.channel.get_event() + self.channel.put_response(Cancelation()) + self.channel.get_event() + dirty = True + + # Release the isolated worker: recycle it if it ended cleanly (set above in + # handle_end_event), else retire it (the pool re-warms lazily). For a cold + # one-shot worker (pool_size=0) this just kills it and frees the buffer. + if self._iso is not None: + from .isolation import release_isolated_worker + + release_isolated_worker(self._iso, dirty=dirty) + self._iso = None def handle(self, provider: Optional[str] = None, value: Optional[Any] = None): """Process a provided value against this mediator's pending event. @@ -1048,13 +1271,41 @@ def handle(self, provider: Optional[str] = None, value: Optional[Any] = None): self.interleaver.batcher.current_provider = provider # Check to see if this mediator has an unprocessed eventto start. - process = self.event_queue.has_value + process = self.channel.has_event # Continue processing events until there are no more events to process. # Means we can move on to the next mediator and continue the model execution. while process: - event, data = self.event_queue.get() + event, data = self.channel.get_event() + + # Isolated path: the worker has no real module, so the host both + # resolves the requester's iteration tag and registers the one-shot + # hook on demand. An in-loop requester crosses UNTAGGED (a marker + # tuple carrying the base path + the worker's live `tracer.iter` pin): + # the host is the single iteration authority, so the tag is computed + # here from the host mediator's own pin/step counter (the same state + # the in-process resolution reads). Out-of-loop requesters arrive + # already tagged and pass through. The resolved string is stashed on + # `_iso_last_tag` so the response's meta piggyback ships it back to + # the worker (backward provenance keys off it). Events restored by + # `handle_value_event` re-enter this loop already resolved, so the + # tag is never recomputed at a later (wrong) step. + if self._iso is not None and event in ( + Events.VALUE, + Events.SWAP, + Events.SKIP, + ): + from .isolation import ensure_isolated_provider, resolve_iso_requester + + if event == Events.VALUE: + data = resolve_iso_requester(self, data) + requester = data + else: + requester = resolve_iso_requester(self, data[0]) + data = (requester, *data[1:]) + self._iso_last_tag = requester + ensure_isolated_provider(self, requester) if event == Events.VALUE: process = self.handle_value_event(data, provider) @@ -1066,8 +1317,14 @@ def handle(self, provider: Optional[str] = None, value: Optional[Any] = None): process = self.handle_skip_event(provider, *data) elif event == Events.BARRIER: process = self.handle_barrier_event(provider, data) + elif event == Events.CACHE: + process = self.handle_cache_event(data) + elif event == Events.BACKWARD: + process = self.handle_backward_event(data) + elif event == Events.UNEMBED: + process = self.handle_unembed_event(data) elif event == Events.END: - process = self.handle_end_event() + process = self.handle_end_event(data) value = self.interleaver.batcher.current_value @@ -1091,6 +1348,13 @@ def handle_value_event(self, requester: Any, provider: Any) -> bool: # If fulfilled by this processor, respond with the value and continue processing events. if provider == requester: + # Isolated backward: keep a reference to the REAL (on-graph) activation tensor, + # keyed by its requester string, so handle_backward_event can run the host's + # real backward and read this activation's gradient. The worker only ever got a + # detached clone, so its half of the graph can't be differentiated there. + if self._iso is not None and self._iso_backward: + self._iso_grad_reals[provider] = self.interleaver.batcher.current_value + # Potentially only select a slice of the value if this mediator is part of a batch group. value = self.interleaver.batcher.narrow(self.batch_group) @@ -1120,7 +1384,7 @@ def handle_value_event(self, requester: Any, provider: Any) -> bool: else: # If the requester has not been seen before, add it to the history and put the value event back in the event queue to be processed later. self.history.add(provider) - self.event_queue.restore((Events.VALUE, requester)) + self.channel.restore_event((Events.VALUE, requester)) return False @@ -1140,6 +1404,19 @@ def handle_swap_event(self, provider: Any, requester: Any, swap_value: Any): """ # If fulfilled by this processor, swap the value and respond with the value and continue processing events. if provider == requester: + # Grad-through-swap: the swap value arrived clone-on-receive (a detached host + # leaf), so the downstream host graph would dead-end here. Make the seam tensor + # require grad BEFORE the forward consumes it (so downstream tracks it and it is + # a backward target) and retain it; handle_backward_event returns dL/d(this leaf) + # for the worker to backprop through its swap tape. Tuple outputs: the residual is + # element [0]. + if self._iso is not None and self._iso_backward: + seam = swap_value[0] if isinstance(swap_value, tuple) else swap_value + if (torch.is_tensor(seam) and seam.is_leaf + and seam.is_floating_point() and not seam.requires_grad): + seam.requires_grad_(True) + self._iso_grad_swaps[requester] = seam + # Swap the value in the batcher. Might only replace a slice of the value if this mediator is part of a batch group. self.interleaver.batcher.swap(self.batch_group, swap_value) @@ -1158,7 +1435,7 @@ def handle_swap_event(self, provider: Any, requester: Any, swap_value: Any): else: # If the requester has not been seen before, add it to the history and put the swap event back in the event queue to be processed later. self.history.add(provider) - self.event_queue.restore((Events.SWAP, (requester, swap_value))) + self.channel.restore_event((Events.SWAP, (requester, swap_value))) return False @@ -1175,6 +1452,14 @@ def handle_exception_event(self, exception: Exception): bool: Flag to stop processing events. """ + if self._iso is not None: + # An EXCEPTION event means the isolated worker's intervention raised and the + # worker looped back idle with a BALANCED pipe (it sent one event, the host + # consumed it, no response is sent) — the worker is healthy and recyclable. + # cancel() honors this unless it finds a still-pending event (dirty), in + # which case it retires the worker instead. + self._iso.clean = True + self.cancel() # Cancelation is okay @@ -1224,6 +1509,19 @@ def handle_barrier_event(self, provider: Any, participants: Set[str]): carries the swap forward to the outer handle context. """ + # Isolated path: the worker can't count participants across processes, so it + # sends the TARGET count (an int). Accumulate names host-side; only coordinate + # once all participants have arrived. Until then this mediator stays blocked at + # the barrier (return False without responding). + if isinstance(participants, int): + n = participants + acc = self.interleaver._barrier_acc + acc.add(self.name) + if len(acc) < n: + return False + participants = set(acc) + acc.clear() + if participants is not None: prev_current = self.interleaver.current @@ -1246,10 +1544,210 @@ def handle_barrier_event(self, provider: Any, participants: Set[str]): return False - def handle_end_event(self): + def handle_cache_event(self, spec: Any): + """Isolated ``tracer.cache()``: register the persistent cache hooks on the + HOST's REAL modules (the worker's are dummies, so its hooks never fire). + + The host ``Cache``'s ``CacheDict`` is filled IN-PLACE by the forward pass (the + hooks are ``mediator_idx=inf``, firing after intervention hooks). + :meth:`handle_end_event` swaps this host ``CacheDict`` in for the worker's empty + placeholder (matched by token) when the saved variables are injected, so the + user's variable *is* the forward-filled host cache. The hooks live on + ``self.hooks`` and are dropped at teardown by ``remove_hooks`` — exactly like the + in-process path. Acks so the worker's ``send`` returns and it proceeds to END. + """ + from .hooks import cache_input_hook, cache_output_hook + from .isolation import path_to_envoy + from .tracing.tracer import Cache + + path2envoy = path_to_envoy(self) + targets = [path2envoy[p] for p in spec["paths"] if p in path2envoy] + + cache_obj = Cache( + spec["paths"], + spec["device"], + spec["dtype"], + spec["detach"], + spec["include_output"], + spec["include_inputs"], + spec["rename"], + spec["alias"], + ) + batcher = self.interleaver.batcher + for envoy in targets: + if spec["include_output"]: + cache_output_hook(cache_obj, envoy._module, envoy.path, batcher, self) + if spec["include_inputs"]: + cache_input_hook(cache_obj, envoy._module, envoy.path, batcher, self) + + self._iso_caches[spec["token"]] = cache_obj + self.set_user_cache(cache_obj) + + self.respond(None) # ack -> worker's send() returns, it proceeds to END + return True + + def handle_backward_event(self, seed: dict): + """Run the backward pass on the host's real graph (isolation). + + The worker computed its half of the chain rule — ``seed`` maps each delivered + activation's requester string to ``dL/d(that activation)`` (the gradient at the + worker→host seam, computed on the worker's local tape). The host continues the + chain rule on its own graph: it differentiates the seeded activations and returns + ``dL/d(activation)`` for every delivered activation the worker might read via + ``.grad`` (keyed by the same requester strings). No user code runs here — only + ``torch.autograd.grad`` on the host's graph with the worker-supplied seeds. + """ + reals = self._iso_grad_reals + swaps = self._iso_grad_swaps + + # The seeds: (real activation tensor, dL/d(activation) from the worker). Seeds key + # into the pre-swap reals — round 1 seeds the loss's delivered reads, round 2 (after + # the worker backprops a swap seam) seeds the pre-swap activation that was swapped. + out_tensors, out_grads = [], [] + for path, grad in seed.items(): + real = reals.get(path) + if torch.is_tensor(real) and real.requires_grad and torch.is_tensor(grad): + out_tensors.append(real) + out_grads.append(grad) + + # The targets: every delivered activation AND every swap leaf still on the host + # graph. Reals serve `.grad` reads (and round-2 re-seeding); swap leaves let the + # worker continue the chain rule through a swap seam (grad-through-swap). Computing + # all keeps it source-agnostic. Kept as separate lists so a real and a swap sharing + # one requester path (read then swapped) don't collide. + real_paths = [p for p, r in reals.items() if torch.is_tensor(r) and r.requires_grad] + swap_paths = [p for p, s in swaps.items() if torch.is_tensor(s) and s.requires_grad] + target_tensors = [reals[p] for p in real_paths] + [swaps[p] for p in swap_paths] + + if (reals or swaps) and not target_tensors: + # Activations were delivered but NONE is on a graph — the forward ran + # without gradient tracking (e.g. generate() runs grad-less). Tell the + # worker so its .grad reads blame the real cause, not "off the path". + self.respond({"__nnsight_backward_no_graph__": True}) + return True + + result: dict = {} + if out_tensors and target_tensors: + grads = torch.autograd.grad( + out_tensors, + target_tensors, + grad_outputs=out_grads, + allow_unused=True, + retain_graph=True, + ) + n = len(real_paths) + result = {p: g for p, g in zip(real_paths, grads[:n]) if torch.is_tensor(g)} + swap_grads = { + p: g for p, g in zip(swap_paths, grads[n:]) if torch.is_tensor(g) + } + if swap_grads: + # The worker backprops these through its swap tape to dL/d(delivered clone) + # and re-seeds (a further BACKWARD round); kept under a reserved key so they + # are not mistaken for `.grad`-readable activation gradients. + result["__nnsight_swap_grads__"] = swap_grads + + self.respond(result) # ack -> worker's send() returns the grad dict + return True + + def handle_unembed_event(self, data): + """Run a residual through the host's REAL final-norm + unembed (isolation). + + The worker holds weightless dummy modules, so a ``F.linear(normed, head.weight)`` + readout — what every logit-lens / steering-direction / attribution-metric cell + does — cannot run there. The worker instead ships the residual VALUE plus the + module PATHS via ``Events.UNEMBED``; the host resolves the real envoys, runs + ``head(norm(residual))`` on the real weights, and ships back only the logits. + Weights never cross the boundary (no containment regression, no model-binding of + the generic worker); only the result does. ``spec`` = + ``{"norm_path", "head_path", "formulation"}`` (``formulation`` ∈ weight|module). + """ + import torch.nn.functional as F + + from .isolation import path_to_envoy + + residual, spec = data + p2e = path_to_envoy(self) + head = p2e.get(spec["head_path"]) + if head is None: + self.respond( + KeyError(f"unembed: head path {spec['head_path']!r} not found on the model") + ) + return True + norm = p2e.get(spec["norm_path"]) if spec.get("norm_path") else None + if spec.get("norm_path") and norm is None: + self.respond( + KeyError(f"unembed: norm path {spec['norm_path']!r} not found on the model") + ) + return True + + normed = norm._module(residual) if norm is not None else residual + if spec["formulation"] == "weight": + logits = F.linear(normed, head._module.weight) + else: + logits = head._module(normed) + + self.respond(logits) # ack -> worker's send() returns the logits + return True + + def handle_end_event(self, payload: Optional[Any] = None): """ Handle an end event by stopping the mediator. + + Worker→host saves transmission: in the isolated path the worker bundles its + ``.save()``'d and ``.carry()``'d values into the END event as + ``(values, saved_names)``, since the worker's frame + ``Globals`` live in another + process. ``values`` is the union of saved + carried locals; ``saved_names`` is the + ``.save()``'d subset. + + The host injects these into the tracer's ``info.frame`` — the **user** frame for a + root trace, or the **session** frame for an inner trace within ``model.session()``. + For an inner trace (the target carries ``__nnsight_tracing_info__``), it replicates + the in-process two-hop push: write ALL values so the next trace sees them, and + re-register the SAVED values' host ids in ``Globals.saves`` so the session's exit-push + keeps them (the worker's ids live in another process; without this the session's + root filter would drop everything). Carried values are written for the next trace but + not registered, so they drop at session exit — matching in-process non-saved + semantics. For a root trace, only the saved subset surfaces to the user frame. """ + from .tracing.globals import Globals + + if self._iso is not None: + # A consumed END means the worker ended cleanly => recyclable by the pool. + # cancel() (below) reads this to recycle vs retire the worker. + self._iso.clean = True + values, saved_names = (payload if payload else ({}, [])) + if values: + if self._iso_caches: + # tracer.cache(): the worker shipped a ``{_ISO_CACHE_TAG: token}`` + # marker in place of its placeholder CacheDict; swap in the HOST cache + # (matched by token), which the forward fills in-place after this + # injection. The user's variable then IS the forward-filled host cache. + values = dict(values) + for name, val in list(values.items()): + if isinstance(val, dict) and _ISO_CACHE_TAG in val: + host_cache = self._iso_caches.get(val[_ISO_CACHE_TAG]) + if host_cache is not None: + values[name] = host_cache.cache + tracer = self.interleaver.tracer + user_frame = ( + tracer.info.frame + if tracer is not None and tracer.info.frame is not None + else self.info.frame + ) + if user_frame is not None: + nested = "__nnsight_tracing_info__" in user_frame.f_locals + if nested: + # inner trace in a session: all values cross to the next trace; + # only saved values are registered so they survive the session push. + push_variables(user_frame, values) + for name in saved_names: + if name in values: + Globals.saves.add(id(values[name])) + else: + # root trace: only saved values surface to the user frame. + saved_only = {k: values[k] for k in saved_names if k in values} + push_variables(user_frame, saved_only) + self.cancel() return False @@ -1292,7 +1790,7 @@ def handle_skip_event(self, provider: Any, requester: Any, value: Any): return True else: self.history.add(provider) - self.event_queue.restore((Events.SKIP, (requester, value))) + self.channel.restore_event((Events.SKIP, (requester, value))) return False @@ -1305,8 +1803,8 @@ def respond(self, value: Optional[Any] = None): """ # Respond and resume the mediator thread. - self.response_queue.put(value) - self.event_queue.wait() + self.channel.put_response(value) + self.channel.wait_event() ### Requester Methods ### @@ -1328,11 +1826,11 @@ def send(self, event: Events, requester: Any): self.push() # Send the event - self.event_queue.put((event, requester)) + self.channel.put_event((event, requester)) # Wait for the interleaver to process the event and respond with the value. - self.response_queue.wait() - response = self.response_queue.get() + self.channel.wait_response() + response = self.channel.get_response() # If the response is an exception, raise it. if isinstance(response, Exception): @@ -1383,7 +1881,7 @@ def end(self): self.push() - self.event_queue.put((Events.END, None)) + self.channel.put_event((Events.END, None)) def exception(self, exception: Exception): """ @@ -1392,7 +1890,7 @@ def exception(self, exception: Exception): Args: exception: The exception that occurred """ - self.event_queue.put((Events.EXCEPTION, exception)) + self.channel.put_event((Events.EXCEPTION, exception)) @property def frame(self) -> FrameType: @@ -1494,15 +1992,24 @@ def __setstate__(self, state): self.intervention = state["intervention"] self.all_stop = state["all_stop"] self.iteration_tracker = state["iteration_tracker"] - self.event_queue = Mediator.Value() - self.response_queue = Mediator.Value() + self.channel: MediatorChannel = InProcessChannel() self.worker = None + self._iso = None + self._isolated_worker = False + self._iso_backward = False + self._iso_grad_reals = {} + self._iso_grad_swaps = {} + self._iso_caches = {} + self._fastlane_verdict = None + self._fastlane_watchdog = None self.interleaver = None self.history = set() self.user_cache: "Cache" = list() self.hooks: List[Any] = list() self.iteration = 0 + self.iter_depth = 0 + self._iso_last_tag = None self.args = list() self.original_globals = {} self.cross_invoker = None diff --git a/src/nnsight/intervention/isolation.py b/src/nnsight/intervention/isolation.py new file mode 100644 index 000000000..e6036060f --- /dev/null +++ b/src/nnsight/intervention/isolation.py @@ -0,0 +1,1104 @@ +"""Transparent isolated execution of mediators in spawned GPU worker processes. + +This is the *outer harness* for the chosen GPU-sandbox design (see +docs/developing/mediator-gpu-trace-integration.md). It does NOT change the +six-event Mediator protocol — it runs the Mediator's intervention in a separate, +GPU-enabled process and routes the protocol over a :class:`CudaIpcChannel`. + +Two shared-memory assumptions of the in-process path become explicit harness steps: + +- **Host-side hook registration**: the worker has no real module, so when the host + receives a ``VALUE``/``SWAP``/``SKIP`` for a requester it hasn't seen, it registers + the matching one-shot hook on the *real* module (resolved from the requester string). +- **Worker→host saves transmission**: ``.save()`` values live in the worker's frame + + ``Globals.saves``; the worker bundles them into the ``END`` event and the host + ``push_variables`` them into the real user frame. + +**Warm worker pool.** Spawning a worker per request costs ~4 s of process bring-up +(cold ``import torch`` + ``import nnsight`` + CUDA context init — measured, model- +independent; see ``prototypes/.../perf_spawn_cost.py``). To amortize it, a worker is +*generic*: it holds a CUDA context + bounce buffer but no mediator, and receives a +serialized mediator per **job** over the channel it already owns (the ~3 ms payload +is the only per-request serialization). A process-global :class:`_WorkerPool` keeps +warm workers across traces; a worker that ends a job cleanly is reset and recycled, +one that is timed-out/killed/cancelled-mid-protocol is retired and lazily re-warmed. + +Each warm worker costs ~0.55 GiB GPU (CUDA context + cuBLAS kernels) **per GPU it +touches**, model-weight-independent and NOT reduced by MPS — so the pool cap is a +real GPU-memory budget, with a cold one-shot fallback past the cap. + +Public surface: :func:`isolate_mediators` (context manager, ``pool_size=`` opt-in), +:func:`warm_worker_pool` / :func:`shutdown_worker_pool`, and :func:`isolation_state`. +``Mediator.start`` calls :func:`acquire_isolated_worker` when isolation is on; the +host ``handle`` loop calls :func:`ensure_isolated_provider` (host-side hook +registration); ``Mediator.cancel`` calls :func:`release_isolated_worker`. +""" +from __future__ import annotations + +import os +import threading +from collections import defaultdict, deque +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Any, Dict, Optional + +import torch +import torch.multiprocessing as mp +import torch.nn as nn + +from . import serialization +from .interleaver import Mediator +from .transport import CudaIpcHostChannel, CudaIpcWorkerChannel +from ..util import apply + +# Types that cross_invoker may ship between workers (data, not framework objects). +_XINVOKE_SCALARS = (int, float, complex, bool, str, bytes, type(None)) + +# Extra wall-clock allowed for a job's first event beyond the user `timeout`, to cover +# deserializing the mediator on an already-warm worker (the spawn+warm is paid before +# the worker's "ready" ack, so the first event only covers deserialize + run-to-first). +_JOB_STARTUP_MARGIN = 30.0 + + +def _transmittable(v) -> bool: + """True if ``v`` is cross_invoker-shareable data: a tensor, a basic scalar, or a + container recursively of those. Framework objects (Barrier/Envoy/model) are not — + the worker already has them via its own closure/deserialization.""" + if torch.is_tensor(v) or isinstance(v, _XINVOKE_SCALARS): + return True + # Note: no `set` — util.apply can't walk sets, so a set holding a CUDA tensor + # would skip the D2H move and crash on the host's IPC re-share. Lists/tuples ok. + if isinstance(v, (list, tuple)): + return all(_transmittable(x) for x in v) + if isinstance(v, dict): + return all(_transmittable(k) and _transmittable(x) for k, x in v.items()) + return False + + +# --------------------------------------------------------------------------- # +# Opt-in surface # +# --------------------------------------------------------------------------- # +@dataclass(frozen=True) +class IsoOptions: + """Per-worker isolation options, in one place for every consumer (the opt-in + context, the pool key, spawn, and the worker bootstrap — previously four + hand-copied dicts). + + ``device``/``arena_bytes``/``gpu_mem_fraction``/``lockdown``/``preimport`` are + **warm-time**: fixed when the worker process spawns, and they define pool + interchangeability (:attr:`pool_key`). ``timeout`` is **per-job**: re-applied to the + host channel each trace in ``_wire_host_channel`` (it also sizes the channel's + first-event budget at spawn). + """ + + device: str = "cuda" + arena_bytes: int = 64 << 20 + gpu_mem_fraction: float = 0.3 + lockdown: bool = False # functional-first; seccomp lockdown enabled separately + # Modules to import at worker warm time, before seccomp lockdown freezes new file + # opens (import == open()); warm-time, so part of pool_key. + preimport: tuple = () + timeout: float = 60.0 # per-step wall-clock cap on user code (hang containment) + # Fast lane: confirmed-safe interventions run IN-PROCESS (full model access, no + # worker, no per-hook channel) instead of in the GPU worker. Without it, isolation + # cannot run the weight-reading interp majority at all (the worker is weightless). + fast_lane: bool = True + trust: str = "local" # only "local" provenance is fast-lane-eligible + fast_lane_timeout: float = 120.0 # whole-intervention watchdog bound for the fast lane + + @property + def pool_key(self) -> tuple: + # Workers are interchangeable ONLY within the same warm-time signature — the + # bounce buffer is device- and size-specific, so reusing a worker across + # devices would copy into the wrong-device buffer (silent corruption). + return ( + str(self.device), + int(self.arena_bytes), + float(self.gpu_mem_fraction), + bool(self.lockdown), + tuple(sorted(self.preimport)), + ) + + +# Generous wait for a COLD spawn+warm (import torch/nnsight + CUDA init, ~4 s typical); +# distinct from the per-job first-event budget (user timeout + deserialize margin). +_WARM_STARTUP_TIMEOUT = 180.0 + +_STATE: Dict[str, Any] = { + "on": False, + "pool_size": 0, # 0 => cold one-shot worker per trace; >0 => warm pool cap + "opts": IsoOptions(), +} + + +def isolation_state() -> Dict[str, Any]: + return _STATE + + +@contextmanager +def isolate_mediators( + arena_bytes: int = 64 << 20, + gpu_mem_fraction: float = 0.3, + device: str = "cuda", + timeout: float = 60.0, + lockdown: bool = False, + pool_size: int = 0, + fast_lane: bool = True, + trust: str = "local", + fast_lane_timeout: float = 120.0, + preimport: tuple = (), +): + """Run interventions inside ``with model.trace(...)`` in an isolated GPU worker. + + Footguns in user intervention code (infinite loops, OOM allocations, device-side + asserts, host-object pokes) are contained to the worker; the model server keeps + serving. + + Args: + timeout: per-step wall-clock cap on user code; a worker that produces no + event within ``timeout`` is presumed hung and killed (the host survives). + pool_size: if > 0, draw workers from a process-global warm pool capped at + ``pool_size`` per (device, arena_bytes, gpu_mem_fraction, lockdown, + preimport) signature (auto-grown lazily, persists across traces, falls back + to a cold one-shot worker past the cap). 0 (default) spawns a cold worker per + trace — the original behavior. Use :func:`warm_worker_pool` to pre-warm at + startup. Under ``lockdown=True`` the worker freezes its import set at warm + time (cold and pooled share the unified worker, which locks down before any + job deserializes), so a job whose user code triggers a NEW import fails — the + cold-vs-pool difference is recycle-vs-retire, not lockdown timing. + fast_lane: if True (default), a per-mediator static classifier + (:mod:`nnsight.intervention.fastlane`) confirms interventions that use only + whitelisted ops / host-model access / nnsight primitives and runs THOSE + in-process (full model + weights, no worker, no per-hook channel), isolating + only the unconfirmable remainder. This is what lets isolation run the + weight-reading interp majority at all — the worker holds weightless dummy + modules. The override only ever moves a mediator from isolate to in-process; + set ``fast_lane=False`` to force pure isolation. + trust: only ``"local"`` provenance is fast-lane-eligible; the static gate is a + footgun selector, not a malice boundary, so any other value disables the fast + lane wholesale (everything isolates). + fast_lane_timeout: whole-intervention wall-clock bound for a fast-laned thread + (a best-effort watchdog restoring loop-containment in-process). + preimport: module names to import at worker warm time, before seccomp lockdown + freezes new file opens (import == ``open()``). Lets isolated interventions use + those modules under ``lockdown=True``, bringing import capability to parity + with an in-process whitelist; also part of the pool signature. + """ + prev = dict(_STATE) + _STATE.update( + on=True, + pool_size=pool_size, + opts=IsoOptions( + device=device, + arena_bytes=arena_bytes, + gpu_mem_fraction=gpu_mem_fraction, + lockdown=lockdown, + timeout=timeout, + fast_lane=fast_lane, + trust=trust, + fast_lane_timeout=fast_lane_timeout, + preimport=tuple(preimport), + ), + ) + try: + yield + finally: + _STATE.update(prev) + + +def fast_lane_enabled() -> bool: + """True if confirmed-safe mediators should run in-process. Gated on the context + option, the ``trust="local"`` provenance cordon, and the global config flag (a server + can force pure isolation without code changes).""" + opts = _STATE["opts"] + if not opts.fast_lane or opts.trust != "local": + return False + try: + from .. import CONFIG + + return bool(getattr(CONFIG.APP, "FAST_LANE", True)) + except Exception: # noqa: BLE001 — config unavailable => default-on + return True + + +def classify_for_fast_lane(mediator): + """Run the static classifier on ``mediator`` (host-side, before any serialization). + Returns a :class:`fastlane.Verdict`. Cached on the intervention code object so a + re-run of the same trace pays the walk once.""" + from . import fastlane + + code = getattr(mediator.intervention, "__code__", None) + cache = _FASTLANE_VERDICT_CACHE + if code is not None and code in cache: + return cache[code] + verdict = fastlane.classify(mediator) + if code is not None: + cache[code] = verdict + return verdict + + +# Verdict cache keyed by the intervention code object identity (per-trace-shape, stable +# across re-runs of the same trace); bounded implicitly by the number of distinct traces. +_FASTLANE_VERDICT_CACHE: Dict[Any, Any] = {} + + +# --------------------------------------------------------------------------- # +# Host side — worker handle + pool # +# --------------------------------------------------------------------------- # +class _PooledWorker: + """Host-side handle to a warm, *generic* GPU worker process. + + Holds the live process, its bounce buffer, and the host channel — all reused + across jobs. Also carries the per-job host-side hook-registration state + (``registered`` / ``path2envoy``) that :func:`ensure_isolated_provider` reads, + and a ``clean`` flag gating recycle (set when this job's END is consumed). + """ + + def __init__(self, proc, buf, conn, channel, poolable: bool, key: tuple = None): + self.proc = proc + self.buf = buf + self.conn = conn + self.channel = channel + self.poolable = poolable # False => one-shot cold worker (never recycled) + self.key = key # base-opts signature (device, ...) for pool keying + self.registered: set = set() # requesters whose host-side hook is registered + self.path2envoy: Optional[dict] = None + self.clean = False # True once this job ended via a consumed END + + def send_job(self, payload, extras, opts) -> None: + self.conn.send(("job", payload, extras, opts)) + + def reset_for_release(self) -> None: + """Drop per-job references so an idle worker doesn't pin them (the hook set, + the path→envoy map, and — via ``channel.reset()`` — the meta/push callbacks + closing over the last trace's interleaver). The *authoritative* per-job reset + happens at acquire time in ``_wire_host_channel``; this is memory hygiene.""" + self.registered = set() + self.path2envoy = None + self.channel.reset() + + def close(self) -> None: + """Stop the worker, escalating SIGTERM->SIGKILL for a wedged CUDA/C call, + then release the host-side pipe fd + GPU bounce buffer (else they linger + until GC).""" + try: + self.conn.send("stop") + except Exception: # noqa: BLE001 + pass + self.proc.join(timeout=5) + if self.proc.is_alive(): + self.proc.terminate() # SIGTERM + self.proc.join(timeout=5) + if self.proc.is_alive(): + self.proc.kill() # SIGKILL — wedged in a non-interruptible CUDA/C call + self.proc.join(timeout=5) + try: + self.channel.close() # closes the pipe fd + except Exception: # noqa: BLE001 + pass + self.buf = None # drop the host ref so the GPU arena can be reclaimed + + +class _WorkerPool: + """Process-global pool of warm generic workers, persisting across traces. + + Thread-safe (mediator starts are sequential today, but a server may run several + traces); ``acquire`` hands out an idle worker or lazily grows the pool up to the + cap, with a cold one-shot worker as the past-cap fallback so correctness never + blocks on the budget. + """ + + def __init__(self): + # Keyed by IsoOptions.pool_key — the warm-time signature (device, arena_bytes, + # gpu_mem_fraction, lockdown, preimport): workers are interchangeable ONLY within + # it. The bounce buffer is device- and size-specific (reusing across devices would + # copy into the wrong-device buffer — silent corruption), and a worker's frozen + # import set must match the requested preimport list. See IsoOptions. + self._idle: Dict[tuple, deque] = defaultdict(deque) + self._all: Dict[tuple, set] = defaultdict(set) + self._lock = threading.Lock() + self._shutting_down = False + + def warm(self, n: int, opts: IsoOptions) -> None: + key = opts.pool_key + with self._lock: + need = max(0, n - len(self._all[key])) + # Spawn outside the lock (each ~4 s); register under it. + fresh = [_spawn_worker(opts, poolable=True) for _ in range(need)] + with self._lock: + for w in fresh: + self._all[key].add(w) + self._idle[key].append(w) + + def acquire(self, opts: IsoOptions, cap: int) -> _PooledWorker: + key = opts.pool_key + dead: list = [] + live = None + placeholder = None + with self._lock: + idle, allset = self._idle[key], self._all[key] + # Skip workers that died while idle (OOM-killed by a neighbor, crashed): + # forget them now and close below, so a dead worker is never handed out. + while idle and live is None: + w = idle.popleft() + if w.proc.is_alive(): + live = w + else: + allset.discard(w) + dead.append(w) + if live is None and len(allset) < cap: + # Reserve the slot under the lock so concurrent acquires can't grow + # past the cap (the spawn itself happens outside the lock). + placeholder = object() + allset.add(placeholder) + for w in dead: + w.close() + if live is not None: + return live + if placeholder is not None: + try: + w = _spawn_worker(opts, poolable=True) + except BaseException: + with self._lock: + self._all[key].discard(placeholder) + raise + with self._lock: + self._all[key].discard(placeholder) + self._all[key].add(w) + return w + # At cap with none idle: a cold one-shot worker so the trace never blocks. + return _spawn_worker(opts, poolable=False) + + def put_idle(self, w: _PooledWorker) -> None: + close_it = False + with self._lock: + if self._shutting_down or w not in self._all[w.key]: + # Released during shutdown, or forgotten while checked out: don't + # re-pool — close it so the process can't leak (close is idempotent if + # it was already torn down). + close_it = True + else: + self._idle[w.key].append(w) + if close_it: + w.close() + + def forget(self, w: _PooledWorker) -> None: + with self._lock: + self._all[w.key].discard(w) + try: + self._idle[w.key].remove(w) + except ValueError: + pass + + def shutdown(self) -> None: + with self._lock: + self._shutting_down = True + workers = [w for s in self._all.values() for w in s] + self._all.clear() + self._idle.clear() + for w in workers: + w.close() + with self._lock: + self._shutting_down = False + + +_POOL = _WorkerPool() + + +def warm_worker_pool( + size: int, + device: str = "cuda", + arena_bytes: int = 64 << 20, + gpu_mem_fraction: float = 0.3, + lockdown: bool = False, + timeout: float = 60.0, + preimport: tuple = (), +) -> None: + """Pre-warm ``size`` generic workers (blocks until each is ready). + + Call once at server startup so the first request pays no spawn cost. The base + options here fix the pool's per-worker configuration. Each worker costs ~0.55 GiB + GPU per GPU it touches — size the pool as a GPU-memory budget. ``preimport`` and + ``lockdown`` must match the values later passed to :func:`isolate_mediators` or the + pre-warmed workers won't match its pool signature and fresh ones are spawned. + """ + _POOL.warm( + size, + IsoOptions( + device=device, + arena_bytes=arena_bytes, + gpu_mem_fraction=gpu_mem_fraction, + lockdown=lockdown, + timeout=timeout, + preimport=tuple(preimport), + ), + ) + + +def shutdown_worker_pool() -> None: + """Stop and free all pooled workers (e.g. at server shutdown).""" + _POOL.shutdown() + + +def _spawn_worker(opts: IsoOptions, poolable: bool) -> _PooledWorker: + """Spawn a generic worker, wait for its one-time ``ready`` ack, wire the channel.""" + ctx = mp.get_context("spawn") # CUDA requires spawn, not fork + buf = torch.empty(opts.arena_bytes, dtype=torch.uint8, device=opts.device) + parent_conn, child_conn = ctx.Pipe() + proc = ctx.Process( + target=_pool_worker_main, args=(child_conn, buf, opts), daemon=True + ) + proc.start() + # The worker warms CUDA + imports (~4 s) then sends exactly one "ready"; consume + # it before the channel starts reading protocol frames on the same pipe. This poll + # covers the cold spawn+warm, so it stays generous. + if not parent_conn.poll(_WARM_STARTUP_TIMEOUT): + proc.terminate() + raise TimeoutError( + f"isolated worker failed to warm up within {_WARM_STARTUP_TIMEOUT}s" + ) + msg = parent_conn.recv() + if msg != "ready": + proc.terminate() + raise RuntimeError(f"unexpected isolated-worker handshake: {msg!r}") + # The channel's first-event budget covers only a job's deserialize + run-to-first- + # request (spawn+warm already happened above), so it is the user timeout plus a + # deserialize margin — NOT the cold 180 s, which would defeat hang-containment on a + # job that hangs before its first event on an already-warm worker. + chan = CudaIpcHostChannel( + parent_conn, + buf, + timeout=opts.timeout, + startup_timeout=opts.timeout + _JOB_STARTUP_MARGIN, + ) + return _PooledWorker( + proc, buf, parent_conn, chan, poolable=poolable, key=opts.pool_key + ) + + +def _build_job(mediator) -> tuple: + """Serialize ``mediator`` into a job message ``(payload, extras, worker_opts)``. + + Module:* and Interleaver are synthesized on the worker; ship the rest + (Tokenizer/Processor) so the deserialized graph resolves them. Only remoteable + models (LanguageModel/VLM/...) carry those extras; a plain NNsight(module) has none. + """ + model = mediator.interleaver.tracer.model + from ..modeling.mixins.remoteable import RemoteableMixin + + extras = {} + if isinstance(model, RemoteableMixin): + real_map = model._remoteable_persistent_objects() + extras = { + k: v + for k, v in real_map.items() + if not k.startswith("Module:") and k != "Interleaver" + } + + # The tracer attaches source during ITS __getstate__; per-mediator serialization + # must do the same first (else source is unavailable). + mediator.intervention.__source__ = "".join(mediator.info.source) + payload = serialization.dumps(mediator) + + worker_opts = { + # default_all (= generate's max_new_tokens) bounds an open-ended iter[:] on + # the worker; it is set AFTER spawn so the live value is also piggybacked on + # each response (meta), but seed the job with the value known now. + "default_all": mediator.interleaver.default_all, + # Mediator.start already applied the in-process cross_invoker gate (multiple + # invokes + config) before acquiring the worker; reuse its decision. The + # worker can't share a frame, so it pushes/pulls through the host store. + "cross_invoker": bool(mediator.cross_invoker), + # `with tensor.backward()` detection — the single decision point for BOTH + # sides: the host gates real-activation retention, the worker gates + # delivered-clone tagging. Prefer the fast-lane classifier's closure-aware flag + # (it resolves through build()/capture() closures the substring is blind to); + # fall back to the source substring when the classifier did not run. + "backward_active": _backward_active(mediator), + } + return payload, extras, worker_opts + + +def _backward_active(mediator) -> bool: + """Closure-aware `with tensor.backward()` detection for the isolated job, preferring + the fast-lane classifier's verdict (which walks through user closures) over the + source substring (blind to a backward hidden in a build()/capture() closure).""" + verdict = getattr(mediator, "_fastlane_verdict", None) + if verdict is not None: + return verdict.differentiate + return ".backward(" in mediator.intervention.__source__ + + +def _wire_host_channel(mediator, iso: _PooledWorker, worker_opts: dict) -> None: + """Point the (possibly recycled) worker's host channel at THIS mediator.""" + chan = iso.channel + chan.reset() # fresh single-slot buffer + startup-timeout + chan._timeout = _STATE["opts"].timeout # per-trace user-code cap + # default_all is set by generate() AFTER the worker is acquired (LanguageModel + # ._execute), so a snapshot is stale. Piggyback the LIVE value + the cross_invoker + # var store on each response; the worker reads default_all before bounding its + # iter[:] loop and pulls the store before each access. + mediator._iso_last_tag = None + chan.meta_provider = lambda: { + "default_all": mediator.interleaver.default_all, + "xinvoke_store": mediator.interleaver._xinvoke_store, + # Host-resolved tag of the event this response answers (in-loop + # requesters cross untagged; the worker keys backward provenance + # off the tag the host actually used). + "resolved": mediator._iso_last_tag, + } + chan.on_push = mediator.interleaver._xinvoke_store.update + mediator.channel = chan + mediator.worker = iso.proc + mediator._iso = iso + # `with tensor.backward()`: if the trace differentiates, the host must retain each + # delivered (real, on-graph) activation so handle_backward_event can run the real + # backward. The decision was made once in _build_job (shared with the worker); + # start with a fresh retention map. + mediator._iso_backward = worker_opts["backward_active"] + mediator._iso_grad_reals = {} + # Fresh per-job host-side hook-registration state. + iso.registered = set() + iso.path2envoy = None + iso.clean = False + + # Host-side iteration tracking. The worker's tracer.iter[...] loop runs on DUMMY + # modules, so its iteration_tracker never advances; the HOST must bump ITS tracker + # per forward pass so the per-step host-registered hooks fire on the right + # generation step (multi-token). Harmless for single-forward traces (one bump). + # Registered on mediator.hooks => torn down by cancel/remove_hooks. + from .tracing.iterator import register_iter_hooks + + register_iter_hooks(mediator, mediator.interleaver.tracer.model) + + +def acquire_isolated_worker(mediator) -> None: + """Acquire a worker (pool or cold), ship the job, and wire the host channel. + + Sets ``mediator.channel`` (host end), ``mediator.worker`` (the process, so + ``alive`` is True), and ``mediator._iso`` (the handle used by host-side hook + registration + cancel/release). + """ + payload, extras, worker_opts = _build_job(mediator) + pool_size = _STATE["pool_size"] + + def _acquire(): + if pool_size > 0: + return _POOL.acquire(_STATE["opts"], pool_size) + return _spawn_worker(_STATE["opts"], poolable=False) + + iso = _acquire() + try: + iso.send_job(payload, extras, worker_opts) + except (BrokenPipeError, EOFError, OSError): + # The worker died between the liveness check and dispatch (tiny race). Forget + # it and retry once through the normal acquisition path with a fresh worker; + # a second failure is a real problem and propagates. + _POOL.forget(iso) + iso = _acquire() + iso.send_job(payload, extras, worker_opts) + _wire_host_channel(mediator, iso, worker_opts) + + +def release_isolated_worker(iso: _PooledWorker, dirty: bool) -> None: + """Recycle a cleanly-ended pooled worker; retire everything else. + + Recyclable only if the worker ended a job cleanly (``clean``, set when its END was + consumed), is poolable, and is still alive. ``dirty`` (the host drained it + mid-protocol with a Cancelation, leaving the pipe unbalanced), a timeout/death + (``clean`` never set), or a one-shot cold worker => retire it; the pool re-warms + lazily on the next acquire. + """ + if (not dirty) and iso.poolable and iso.clean and iso.proc.is_alive(): + iso.reset_for_release() + _POOL.put_idle(iso) + else: + iso.close() + _POOL.forget(iso) + + +def path_to_envoy(mediator) -> dict: + """The ``{path: envoy}`` map for this job's model, built lazily and cached on the + worker handle (fresh per job — ``_wire_host_channel`` clears it). Shared by + host-side hook registration and ``handle_cache_event``'s target resolution.""" + iso = mediator._iso + if iso.path2envoy is None: + model = mediator.interleaver.tracer.model + iso.path2envoy = {e.path: e for e in model.modules()} + return iso.path2envoy + + +# First element of the wire tuple an isolated worker sends for an in-loop +# (`tracer.iter`) requester: ("", base_path, pin, pin_dirty). The host +# resolves the iteration tag; out-of-loop requesters cross as plain tagged strings. +ISO_ITER_REQ = "__nnsight_iso_iter_req__" + + +def resolve_iso_requester(mediator, wire) -> str: + """Resolve a worker-sent requester into its final tagged string, host-side. + + A plain string (out-of-loop, already tagged by the worker) passes through. + A marker tuple (in-loop) is resolved against the HOST mediator's state with + the exact in-process rule (`Interleaver.iterate_requester`): the pin if set, + else the per-path step counter. The worker's pin rides the tuple and is + applied only when `pin_dirty` says the worker's iter loop assigned it since + the last event; otherwise the host's own pin state (including fire-time + relaxation done by the one-shot hooks) is authoritative. + """ + if isinstance(wire, str): + return wire + _, base, pin, pin_dirty = wire + if pin_dirty: + mediator.iteration = pin + iteration = ( + mediator.iteration + if mediator.iteration is not None + else mediator.iteration_tracker[base] + ) + return f"{base}.i{iteration}" + + +def ensure_isolated_provider(mediator, requester: str) -> None: + """Host-side hook registration: register the one-shot hook for ``requester`` on the *real* module. + + Parses ``"..i"``, resolves the envoy on the host, and + registers the existing ``output_hook``/``input_hook`` with the host mediator for + the **specific step N** parsed from the requester — the worker's iteration counter + lives in another process, so the step must come from the wire, not from the host + mediator's ``iteration``. The host's own iteration_tracker (bumped per forward by + the iter hooks installed in ``_wire_host_channel``) advances so the N-hook fires + on the right generation step. Idempotent per requester (per step). + """ + iso = mediator._iso + if requester in iso.registered: + return + iso.registered.add(requester) + + # The iteration suffix is always the rightmost ".i". + parts = requester.rsplit(".i", 1) # ["model.transformer.h.6.output", "2"] + base = parts[0] + iteration = int(parts[1]) if len(parts) == 2 and parts[1].isdigit() else None + path, _, kind = base.rpartition(".") # ("model.transformer.h.6", ".", "output") + if kind not in ("output", "input"): + return # externally-provided eproperties (e.g. .result) need no module hook + + envoy = path_to_envoy(mediator).get(path) + if envoy is None: + return # unknown path → let normal missed-provider handling surface it + + from .hooks import input_hook, output_hook + + if kind == "output": + output_hook(mediator, envoy._module, base, iteration=iteration) + else: + input_hook(mediator, envoy._module, base, iteration=iteration) + + +# --------------------------------------------------------------------------- # +# Worker side # +# --------------------------------------------------------------------------- # +class _WorkerBatcher: + """Minimal stand-in for the Batcher that ``requires_*`` reads on the worker.""" + + current_provider = None + current_value = None + + +class _WorkerInterleaver: + """Worker-side interleaver stub: enough for eproperty + requires_* to run. + + The real Batcher/narrow/swap all live on the host; the worker only builds + requester strings and emits events. + """ + + def __init__(self, default_all=None): + self.interleaving = True + self.batcher = _WorkerBatcher() + self.current = None + # Open-ended `tracer.iter[:]` reads this to know how many generation steps + # to run; set by generate(max_new_tokens=N) on the host and shipped over. + self.default_all = default_all + + def iterate_requester(self, requester: str): + med = self.current + # Inside a `tracer.iter[...]` scope the step counter lives on the HOST + # (the worker's tracker never advances: its modules are dummies), so the + # requester crosses untagged and the host resolves it. The pin rides + # along, flagged dirty only if the loop assigned it since the last event, + # so host-side pin relaxation is not clobbered by a stale worker pin. + if med.iter_depth: + dirty, med._pin_dirty = med._pin_dirty, False + return (ISO_ITER_REQ, requester, med.iteration, dirty) + # Out of loop the pin is authoritative and lives here; resolve locally + # (same rule as in-process). + iteration = ( + med.iteration if med.iteration is not None else med.iteration_tracker[requester] + ) + return f"{requester}.i{iteration}" + + +class _WorkerPersistent: + """persistent_objects map for the worker: synthesize a dummy module for every + ``Module:`` (no weights), the worker interleaver for ``Interleaver``, and + pass through everything else (Tokenizer/Processor).""" + + def __init__(self, interleaver, extras: dict): + self._extras = dict(extras) + self._extras["Interleaver"] = interleaver + self._dummies: dict = {} + + def __contains__(self, key) -> bool: + return str(key).startswith("Module:") or key in self._extras + + def __getitem__(self, key): + skey = str(key) + if skey.startswith("Module:"): + d = self._dummies.get(skey) + if d is None: + d = nn.Module() + d.__path__ = skey[len("Module:") :] + self._dummies[skey] = d + return d + return self._extras[key] + + +def _transmissible_exc(e): + """Degrade an exception to a form that pickles across the EXCEPTION event. + + Plain user exceptions (ValueError, ...) pickle and the host wraps them. *Dynamic* + nnsight exceptions (NNsightException) don't ("Can't pickle nnsight.NNsightException") + — degrade to a plain RuntimeError preserving type name + message. + """ + from multiprocessing.reduction import ForkingPickler + + try: + ForkingPickler.dumps(e) + return e + except Exception: # noqa: BLE001 + return RuntimeError(f"{type(e).__name__}: {e}") + + +class WorkerMediator(Mediator): + """The worker-process half of an isolated mediator. + + A job's mediator is deserialized as a plain :class:`Mediator` and then *adopted* + (``__class__`` swap) into this subclass, which overrides exactly the methods whose + in-process behavior relies on shared memory with the host: + + - ``end``: the worker's frame + ``Globals.saves`` live here, so the exit filter + runs locally and the saved dict rides the END event (worker→host saves + transmission). + - ``exception``: dynamic nnsight exceptions don't pickle; degrade before shipping. + - ``request``: when the trace contains ``with tensor.backward()``, tag each + delivered activation clone (``requires_grad_`` + id→requester provenance) so the + worker builds its half of the autograd graph and can seed the host backward. + + ``apply_meta`` / ``push_locals`` are the worker ends of the live host↔worker state + piggyback (`iter[:]` bound + cross_invoker variable store) and are bound to the + channel's ``on_meta`` / ``push_provider`` per job. + """ + + # The `tracer.iter` pin, instrumented so `iterate_requester` can tell the host + # whether the worker's loop assigned it since the last event (dirty). The class + # swap in `adopt` makes this data descriptor shadow the deserialized instance + # attribute, which `adopt` migrates into `_iteration`. + @property + def iteration(self): + return self._iteration + + @iteration.setter + def iteration(self, value): + self._iteration = value + self._pin_dirty = True + + @classmethod + def adopt(cls, mediator, channel, interleaver, opts: dict, device) -> "WorkerMediator": + """Turn a freshly-deserialized mediator into this job's worker mediator.""" + mediator.__class__ = cls + # Migrate the pin under the property; not dirty: the host mediator starts + # from the identical deserialized state, so there is nothing to push. + mediator._iteration = mediator.__dict__.pop("iteration") + mediator._pin_dirty = False + # Final tag of the last-delivered value, resolved on the host and shipped + # back on the response meta (None for out-of-loop requesters, which the + # worker tagged itself). + mediator._last_resolved = None + mediator.channel = channel + mediator.interleaver = interleaver + mediator.idx = 0 + # cross_invoker matches the host gate; var sharing rides the host store since + # worker frames aren't shared across processes. + mediator.cross_invoker = opts.get("cross_invoker", False) + mediator._isolated_worker = True # so Barrier sends the target count (host counts) + mediator._device = device + # `with tensor.backward()`: decided once in _build_job (host-side) and shipped + # with the job; gates the delivered-clone tagging in ``request`` + # (BackwardsTracer reads the provenance). + mediator._bwd_active = opts["backward_active"] + mediator._bwd_prov = {} # id(delivered clone) -> requester string + mediator._bwd_tagged = [] # delivered clones made to require grad + mediator._bwd_swaps = {} # requester string -> worker-tape swap value (grad-through-swap) + interleaver.current = mediator + return mediator + + def _resolved_tag(self, requester) -> str: + """The final tagged requester of the exchange that just completed: an + out-of-loop requester was tagged locally; an in-loop one crossed as a + marker tuple and its tag came back on the response meta.""" + return requester if isinstance(requester, str) else self._last_resolved + + def request(self, requester): + value = super().request(requester) + if self._bwd_active: + self._tag_delivered(value, self._resolved_tag(requester)) + return value + + def swap(self, requester, value): + # Grad-through-swap: keep the worker-TAPE swap value (with grad_fn back to the + # delivered clones) before send() ships its detached copy. The backward block uses + # it to backprop the host-returned dL/d(swap leaf) into dL/d(delivered clone). Tuple + # outputs: the residual is element [0]. Keyed by the HOST-resolved tag (known + # only after the exchange), matching the host's `_iso_grad_swaps` keys. + seam = None + if self._bwd_active: + seam = value[0] if isinstance(value, tuple) else value + if not (torch.is_tensor(seam) and seam.requires_grad): + seam = None + result = super().swap(requester, value) + if seam is not None: + self._bwd_swaps[self._resolved_tag(requester)] = seam + return result + + def _tag_delivered(self, value, requester: str) -> None: + """Tag each delivered activation tensor with its requester provenance and make + it require grad, so worker-side ops on it build the worker's half of the graph.""" + + def _tag(t): + if t.is_floating_point() and t.is_leaf: + t.requires_grad_(True) + self._bwd_prov[id(t)] = requester + self._bwd_tagged.append(t) + return t + + apply(value, _tag, torch.Tensor) + + def end(self): + # Worker→host saves transmission: bundle .save()'d (and .carry()'d) values into the + # END event. The intervention's compiled body calls ``end()`` on success; push() + # populates the SerializedFrame's f_locals, which we filter by Globals.saves/shared. + from .interleaver import Events, _ISO_CACHE_TAG + from .tracing.globals import Globals + from .tracing.tracer import Cache + + self.push() + flocals = self.info.frame.f_locals + saved = {k: v for k, v in flocals.items() if id(v) in Globals.saves} + # Carried (.carry()) values: cross-trace handoffs within a session. Ship them too so + # the host can write them to the session frame for the next trace, but tag which are + # SAVED (saved_names) so the host surfaces only those to the user frame. With no + # .carry() in play this is exactly the prior payload (saved only) — no regression for + # the single-trace path. + carried = { + k: v for k, v in flocals.items() + if id(v) in Globals.shared and id(v) not in Globals.saves + } + values = {**saved, **carried} + # A tracer.cache() placeholder CacheDict is a live object, not a value — ship its + # token as a plain-dict marker so the value codec accepts it; the host swaps in its + # own forward-filled cache by token (top-level saves only, as the host handler is). + for k, v in list(values.items()): + if isinstance(v, Cache.CacheDict): + values[k] = {_ISO_CACHE_TAG: getattr(v, "_iso_cache_token", None)} + self.channel.put_event((Events.END, (values, list(saved.keys())))) + + def exception(self, exception: Exception): + super().exception(_transmissible_exc(exception)) + + def apply_meta(self, m: dict) -> None: + # Live host state piggybacked on each response: the iter[:] bound, the + # cross_invoker var store (pulled into the frame so push()/pull() see it), + # and the host-resolved tag of the event this response answers. + self._last_resolved = m["resolved"] + self.interleaver.default_all = m.get( + "default_all", self.interleaver.default_all + ) + store = m.get("xinvoke_store") + if store: + # Store tensors travel CPU-serialized (see push_locals); move them + # back to the worker's device before the user code uses them. + restored = apply(store, lambda t: t.to(self._device), torch.Tensor) + self.info.frame.f_locals.update(restored) + + def push_locals(self) -> Optional[dict]: + # cross_invoker: ship this worker's *data* locals to the host store. + # push() (called by send() before put_event) has already written them into + # the SerializedFrame's f_locals. We ship only transmittable data (tensors + # + basic types/containers) — framework objects (Barrier/Envoy, which hold + # the model) are skipped; the worker already has them via its own closure. + # Tensors are moved to CPU: a worker tensor cloned from the CUDA-IPC bounce + # buffer cannot be re-shared over IPC by the host. + if not self.cross_invoker: + return None + out = {} + for k, v in self.info.frame.f_locals.items(): + if str(k).startswith("__nnsight"): + continue + if not _transmittable(v): + import warnings + + warnings.warn( + f"cross_invoker: variable {k!r} ({type(v).__name__}) is not " + f"transmittable across the isolation boundary and was not " + f"shared between invokes." + ) + continue + out[k] = apply(v, lambda t: t.detach().cpu(), torch.Tensor) + return out + + +# The job currently running in this worker process (one at a time). Read by +# ``worker_backward_context`` so BackwardsTracer can find the ambient mediator. +_WORKER_CURRENT: Optional[WorkerMediator] = None + +# Whether this worker has installed its (one-way) seccomp lockdown yet. Lockdown is +# deferred to the first job — after the job's host-authored payload is deserialized, before +# its user code runs — so the deserialize's imports (e.g. transformers' lazy submodules, +# loaded only at unpickle time) are not blocked. See ``_run_one_job``. +_LOCKED_DOWN: bool = False + + +def worker_backward_context() -> Optional[WorkerMediator]: + """Return this worker's mediator if its trace contains a backward block, else None. + + Used by ``BackwardsTracer.execute`` to detect that ``.backward()`` is running in an + isolated worker (so it must drive the host's real backward instead of differentiating + its detached clones locally); the mediator carries the delivered-clone provenance + (``_bwd_prov`` / ``_bwd_tagged``).""" + med = _WORKER_CURRENT + if med is not None and med._bwd_active: + return med + return None + + +def _run_one_job(channel, payload, extras, opts, device, lockdown: bool = False) -> None: + """Deserialize one mediator against fresh dummies, adopt it as this job's + :class:`WorkerMediator`, run its intervention, and ship saves at END. Any failure + (including a bad payload) is reported as an EXCEPTION event so the host never + waits on a worker that won't speak. + + Under ``lockdown`` the seccomp filter is installed HERE — after the (host-authored, + trusted) payload is deserialized, before the (untrusted) user intervention runs — so + deserialization's own imports succeed. It is one-way and installed once; in a warm pool + later jobs run under the first job's lockdown, so their deserialize must not need a NEW + import (true for a homogeneous model; otherwise pre-load via ``preimport=``).""" + from .interleaver import Events + from .tracing.globals import Globals + + global _WORKER_CURRENT, _LOCKED_DOWN + try: + Globals.saves.clear() # per-job reset (the only worker-side global state) + Globals.shared.clear() + + interleaver = _WorkerInterleaver(default_all=opts.get("default_all")) + mediator = serialization.loads(payload, _WorkerPersistent(interleaver, extras)) + WorkerMediator.adopt(mediator, channel, interleaver, opts, device) + _WORKER_CURRENT = mediator + channel.on_meta = mediator.apply_meta + channel.push_provider = mediator.push_locals + + # Footgun containment: seccomp-block new fs/net/exec syscalls NOW — the trusted + # payload is deserialized (its imports done), the untrusted user code is next. CUDA + # + the control Pipe + the IPC buffer use already-open fds, so they keep working. + if lockdown and not _LOCKED_DOWN: + from ._sandbox import lock_down + + lock_down() + _LOCKED_DOWN = True + + mediator.intervention(mediator, mediator.info, *mediator.args) + except BaseException as e: # noqa: BLE001 — contain the footgun; report it + try: + channel.put_event((Events.EXCEPTION, _transmissible_exc(e))) + except Exception: # noqa: BLE001 + pass + finally: + _WORKER_CURRENT = None + + +def _pool_worker_main(conn, buf, worker_iso_opts: IsoOptions): + """Generic worker: warm CUDA + imports ONCE, optionally lock down, then loop + serving ``("job", payload, extras, opts)`` messages until told to ``"stop"``. + + The CUDA context, warmed kernels, bounce buffer, and channel persist across jobs — + that is what the warm pool amortizes. Per job, a fresh mediator is deserialized + against fresh dummy modules (no cross-job state but ``Globals.saves``, cleared).""" + from .tracing.globals import _ensure_mounted + + device = worker_iso_opts.device + + # Warm CUDA before any lockdown so kernels/contexts are loaded. + if torch.cuda.is_available(): + _ = (torch.randn(8, 8, device=device) @ torch.randn(8, 8, device=device)).sum() + torch.cuda.synchronize() + + # Warm imports user ops might trigger, before seccomp closes new file opens. + try: + import numpy # noqa: F401 + except Exception: # noqa: BLE001 + pass + import cloudpickle # noqa: F401 + + cloudpickle.loads(cloudpickle.dumps(lambda _t: _t)) + + _ensure_mounted() # install Object.save so `.save()` resolves in the worker + + # User-configurable warm-time pre-imports: load modules interventions may need + # BEFORE seccomp freezes new file opens (import == open()). Under lockdown a module + # not loaded here is unimportable in any job, so pre-warming the deployment's + # allowed-module set brings user-facing import capability to parity with an + # in-process whitelist. A failed pre-import is non-fatal (warn + skip). + if worker_iso_opts.preimport: + import importlib + import warnings + + for _mod in worker_iso_opts.preimport: + try: + importlib.import_module(_mod) + except Exception as _e: # noqa: BLE001 + warnings.warn(f"isolated worker pre-import of {_mod!r} failed: {_e!r}") + + if worker_iso_opts.gpu_mem_fraction and torch.cuda.is_available(): + torch.cuda.set_per_process_memory_fraction(worker_iso_opts.gpu_mem_fraction) + + channel = CudaIpcWorkerChannel(conn, buf) # persistent; rebinds handlers per job + + # Lockdown is NOT installed here: the job message (received + unpickled by ``conn.recv`` + # below) and the mediator payload are host-authored, trusted data whose deserialization + # may need new imports (transformers loads its modeling submodules lazily, only at + # unpickle time). Installing seccomp before the first ``conn.recv`` blocks those opens + # and the worker dies (EPERM is an ``OSError`` → swallowed by the recv loop's + # ``except``). Instead ``_run_one_job`` installs lockdown after the first job's payload + # is deserialized, before its user code runs. The first ``conn.recv`` therefore runs + # unlocked; a warm pool's later recvs run under the first job's lockdown, so a NEW model + # type would need ``preimport=`` (a homogeneous model needs nothing — already imported). + + # One-time ready ack: the spawner consumes this before the channel reads protocol. + conn.send("ready") + + while True: + try: + msg = conn.recv() + except (EOFError, OSError): + break + if msg == "stop": + break + if not (isinstance(msg, tuple) and msg and msg[0] == "job"): + continue # ignore stray control messages + _, payload, extras, opts = msg + _run_one_job(channel, payload, extras, opts, device, worker_iso_opts.lockdown) + # Job done; the worker is idle and recyclable. Loop for the next job/stop. + + # Skip interpreter atexit handlers: under seccomp lockdown, tempfile's atexit + # rmtree hits blocked openat/unlink and recurses. The worker is disposable and + # the host owns the bounce buffer, so a hard exit is correct here. + os._exit(0) diff --git a/src/nnsight/intervention/tracing/backwards.py b/src/nnsight/intervention/tracing/backwards.py index b154c3b00..a0fed97c0 100755 --- a/src/nnsight/intervention/tracing/backwards.py +++ b/src/nnsight/intervention/tracing/backwards.py @@ -95,6 +95,12 @@ def __init__( def execute(self, fn: Callable): + from ..isolation import worker_backward_context + + worker_mediator = worker_backward_context() + if worker_mediator is not None: + return self._execute_isolated(fn, worker_mediator) + mediator = BackwardsMediator(fn, self.info) interleaver = Interleaver([mediator], self) @@ -110,3 +116,122 @@ def execute(self, fn: Callable): finally: grad_patch.restore() interleaver.cancel() + + def _execute_isolated(self, fn: Callable, worker_mediator): + """Run a backward block inside an isolated GPU worker. + + The autograd graph is split across the process boundary: the worker holds the + recipe from each delivered activation up to ``self.tensor`` (the loss); the host + holds the recipe from the model inputs up to those activations. We stitch the two + halves at the named seam (the activation requester strings): + + 1. The worker computes its half — ``dL/d(delivered activation)`` for each tagged + leaf the loss depends on — as the seed. + 2. The host runs its half (``handle_backward_event``) seeded by those gradients + and returns ``dL/d(activation)`` for every delivered activation, keyed by the + same requester strings. + 3. The backward block body runs here in the worker; each ``.grad`` read is served + from that returned dict by the tensor's provenance — no local ``register_hook`` + (the clones carry no host graph) and no local backward (the graph is host-side). + """ + from ..interleaver import Events, Interleaver + + forward_mediator = worker_mediator + provenance = worker_mediator._bwd_prov + tagged = [t for t in worker_mediator._bwd_tagged if t.requires_grad] + swaps = worker_mediator._bwd_swaps # requester -> worker-tape swap value + + # Worker half of the chain rule: seed = dL/d(delivered leaf) for leaves the loss + # actually depends on (allow_unused drops the rest). + seed: dict = {} + if tagged: + grads = torch.autograd.grad( + self.tensor, tagged, allow_unused=True, retain_graph=True + ) + for leaf, grad in zip(tagged, grads): + if grad is not None: + seed[provenance[id(leaf)]] = grad + + # Stitch the chain across the boundary, iterating over swap seams. Each round the + # host returns dL/d(activation) for reads AND dL/d(swap leaf) (a swap installs a host + # leaf, severing the host graph at the seam). The worker backprops each swap-leaf + # grad through its swap tape to dL/d(delivered clone) and re-seeds the pre-swap graph + # (another BACKWARD round). A read reached both directly and through a swap sums its + # contributions across rounds. With no swaps this is the original single exchange. + worker_grads: dict = {} + no_graph = False + while seed: + resp = forward_mediator.send(Events.BACKWARD, seed) or {} + no_graph = no_graph or bool(resp.pop("__nnsight_backward_no_graph__", False)) + swap_grads = resp.pop("__nnsight_swap_grads__", {}) + for path, g in resp.items(): + if torch.is_tensor(g): + worker_grads[path] = g if path not in worker_grads else worker_grads[path] + g + # Next round's seed: backprop each returned swap-leaf grad through the worker's + # swap tape to its delivered-clone leaves. + seed = {} + for swap_path, sg in swap_grads.items(): + swapped = swaps.get(swap_path) + if not (tagged and torch.is_tensor(swapped) and torch.is_tensor(sg)): + continue + grads = torch.autograd.grad( + swapped, tagged, grad_outputs=sg, allow_unused=True, retain_graph=True + ) + for leaf, g in zip(tagged, grads): + if g is not None: + p = provenance[id(leaf)] + seed[p] = g if p not in seed else seed[p] + g + + mediator = BackwardsMediator(fn, self.info) + interleaver = Interleaver([mediator], self) + grad_patch = Patch( + torch.Tensor, + _isolated_grad_property(provenance, worker_grads, no_graph), + "grad", + ) + try: + grad_patch.patch() + # No local backward: the grads are already in worker_grads. The block body + # just reads them (via the patched .grad) and .save()s; its saves push up + # into the forward frame and ride the forward's END to the host. + with interleaver: + pass + interleaver.check_dangling_mediators() + finally: + grad_patch.restore() + interleaver.cancel() + + +def _isolated_grad_property(provenance: dict, worker_grads: dict, no_graph: bool = False): + """A ``Tensor.grad`` property for the isolated backward block: read the gradient from + the host-computed ``worker_grads`` by the tensor's delivery provenance (path), instead + of registering a local autograd hook (the worker clone has no host graph).""" + + def getter(tensor: torch.Tensor): + path = provenance.get(id(tensor)) + if path is None: + raise RuntimeError( + "gradient under isolation is only supported on an unmodified module-" + "output tensor read during the trace (e.g. `model...ln_f.output`); " + "this tensor was derived in user code and has no host-side graph." + ) + if path not in worker_grads: + if no_graph: + raise RuntimeError( + f"no gradient available for `{path}` — the forward pass ran " + f"without gradient tracking, so there is no autograd graph " + f"(generate() runs grad-less; use model.trace() for gradients). " + f"This matches in-process behavior, where .grad raises here too." + ) + raise RuntimeError( + f"no gradient available for `{path}` — it is off the backward path " + f"from the loss (its gradient never flowed during the backward pass)." + ) + return worker_grads[path] + + def setter(tensor: torch.Tensor, value: Any): + raise NotImplementedError( + "editing `.grad` under isolation is not yet supported." + ) + + return property(getter, setter) diff --git a/src/nnsight/intervention/tracing/base.py b/src/nnsight/intervention/tracing/base.py index f5fa48dbe..bca7cad35 100755 --- a/src/nnsight/intervention/tracing/base.py +++ b/src/nnsight/intervention/tracing/base.py @@ -566,7 +566,10 @@ def push(self, state: Dict = None) -> Dict: filtered_state = { k: v for k, v in filtered_state.items() if id(v) in Globals.saves } + # Carried (.carry()) values are cross-trace handoffs, never surfaced to the + # user frame; clear them at the root/session exit alongside saves. Globals.saves.clear() + Globals.shared.clear() if target_frame is not None: # Push the filtered variables back to the original frame diff --git a/src/nnsight/intervention/tracing/globals.py b/src/nnsight/intervention/tracing/globals.py index 2594cb8ea..30c7bb4fa 100755 --- a/src/nnsight/intervention/tracing/globals.py +++ b/src/nnsight/intervention/tracing/globals.py @@ -10,14 +10,15 @@ def _ensure_mounted(): - """Mount Object.save as the universal `.save` method. + """Mount Object.save / Object.carry as the universal `.save` / `.carry` methods. - Lazy one-time setup. Called from .save() / nnsight.save() so we only - pay the C-level mount cost once, on first use. + Lazy one-time setup, run at trace setup (``_setup_interleaver``) and from + ``.save()`` / ``.carry()`` so we only pay the C-level mount cost once. """ global _mounted if CONFIG.APP.PYMOUNT and not _mounted: mount(Object.save, "save") + mount(Object.carry, "carry") _mounted = True @@ -28,6 +29,25 @@ def save(object: Any): return object +def carry(object: Any): + """Mark ``object`` to be handed to a later trace in the same ``model.session()`` + WITHOUT surfacing it as a saved output. + + The portable counterpart to relying on an inner trace's locals flowing implicitly: + in-process a non-saved value already crosses to the next trace, but under isolation + each inner trace runs in a worker that only ships its outputs home, so a non-saved + value would vanish. ``.carry()`` explicitly registers the value to cross the boundary + — so the same code is correct in-process AND isolated. Unlike :func:`save`, a carried + value is dropped at session exit (it is not in ``Globals.saves``), so it never appears + in the caller's frame; use it for cross-trace handoffs (an activation to patch into a + later run) that are not themselves results. + """ + + Globals.shared.add(id(object)) + + return object + + class Object(torch.Tensor): def save(self, _=0): @@ -46,6 +66,24 @@ def save(self, _=0): return self + def carry(self, _=0): + """Hand this value to a later trace in the same ``model.session()`` without saving + it as an output. See :func:`carry`. + + Examples: + + >>> with model.session(): + ... with model.trace("clean prompt"): + ... act = model.transformer.h[6].output.carry() # not a result + ... with model.trace("corrupt prompt") as tracer: + ... tracer.patch(model.transformer.h[6], act) # transplanted in + ... logits = model.lm_head.output.save() + """ + + carry(self) + + return self + def __getattr__(self, name: str) -> Self: return super().__getattr__(name) @@ -100,10 +138,14 @@ def clear(self): class Globals: """Process-wide tracing state. - Holds two pieces of true global state: + Holds these pieces of true global state: - ``saves``: set of ``id()`` for objects marked via ``.save()``. The root tracer's ``push()`` filters its frame locals against this set so only saved values propagate out of the trace. + - ``shared``: set of ``id()`` for objects marked via ``.carry()`` — + cross-trace handoffs within a ``model.session()`` that are NOT surfaced + as outputs. Consulted by the isolated worker's ``end()`` to ship carried + values across the boundary; dropped at session exit (not in ``saves``). - ``cache``: source/AST/code-object memoization across traces. Root-vs-inner detection lives on the tracer itself — see @@ -113,9 +155,12 @@ class Globals: saves = set() + shared = set() + cache = TracingCache() @staticmethod def clear(): Globals.saves.clear() + Globals.shared.clear() Globals.cache.clear() diff --git a/src/nnsight/intervention/tracing/iterator.py b/src/nnsight/intervention/tracing/iterator.py index 6361ca435..3f1cf85fd 100755 --- a/src/nnsight/intervention/tracing/iterator.py +++ b/src/nnsight/intervention/tracing/iterator.py @@ -230,6 +230,11 @@ def __iter__(self): # in the finally block below so they don't leak outside the loop. iter_handles = register_iter_hooks(mediator, self.model) + # Mark the iter scope open. In an isolated worker this routes requester + # tagging to the host (the worker's tracker never advances on its dummy + # modules); in-process nothing reads it. + mediator.iter_depth += 1 + try: if isinstance(self.iteration, slice): @@ -284,6 +289,7 @@ def __iter__(self): yield self.iteration finally: + mediator.iter_depth -= 1 mediator.iteration = original_iteration # Remove the iteration-tracking hooks. @@ -331,6 +337,7 @@ def execute(self, fn: Callable): mediator.push() iter_handles = register_iter_hooks(mediator, self.model) + mediator.iter_depth += 1 def do_iteration(iter: int): @@ -381,6 +388,7 @@ def do_iteration(iter: int): do_iteration(self.iteration) finally: + mediator.iter_depth -= 1 mediator.iteration = original_iteration for handle in iter_handles: diff --git a/src/nnsight/intervention/tracing/tracer.py b/src/nnsight/intervention/tracing/tracer.py index 3b2c3791b..cd2a4fa9b 100755 --- a/src/nnsight/intervention/tracing/tracer.py +++ b/src/nnsight/intervention/tracing/tracer.py @@ -614,20 +614,240 @@ def cache( targets.append(envoy) # Register persistent cache hooks on each target module - for envoy in targets: - if include_output: - cache_output_hook( - cache_obj, envoy._module, envoy.path, batcher, mediator - ) - if include_inputs: - cache_input_hook( - cache_obj, envoy._module, envoy.path, batcher, mediator - ) + if mediator._isolated_worker: + # Isolated: the worker's modules are DUMMIES, so cache hooks registered + # here would never fire. Ship the spec via a CACHE event so the HOST + # registers the hooks on its real modules and fills this cache. The + # returned CacheDict is a token-tagged placeholder the user binds + + # .save()s; the host swaps in its own (forward-filled) CacheDict when the + # saved variables are injected into the user frame (matched by token). + from ..interleaver import Events + + token = id(cache_obj.cache) + cache_obj.cache._iso_cache_token = token + spec = { + "token": token, + "paths": [envoy.path for envoy in targets], + "device": device, + "dtype": dtype, + "detach": detach, + "include_output": include_output, + "include_inputs": include_inputs, + "rename": rename_dict, + "alias": alias_dict, + } + mediator.send(Events.CACHE, spec) + else: + for envoy in targets: + if include_output: + cache_output_hook( + cache_obj, envoy._module, envoy.path, batcher, mediator + ) + if include_inputs: + cache_input_hook( + cache_obj, envoy._module, envoy.path, batcher, mediator + ) mediator.set_user_cache(cache_obj) return cache_obj.cache + def unembed(self, residual, norm, head, formulation: str = "weight"): + """Project ``residual`` through the final norm + unembed → logits, on the host's + REAL weights even when the intervention is isolated. + + The readout every logit-lens / steering-direction / attribution-metric cell does + — ``F.linear(norm(residual), head.weight)`` — reads the host model's real + weights. In-process (incl. the fast lane) that just runs the real modules. In an + isolated worker the modules are weightless dummies, so this ships the residual + VALUE plus the module PATHS to the host (``Events.UNEMBED``); the host runs the + real norm + unembed and ships back only the logits. Weights never cross the + boundary — so this works on the isolated tier without binding the generic worker + to a model or placing host weight memory in the (less-trusted) worker. + + Args: + residual: the residual-stream tensor to project (a block ``.output``; a tuple + is untupled to ``[0]``). + norm: the final-norm Envoy (e.g. ``model.transformer.ln_f``), or ``None`` to + skip normalization. + head: the unembed Envoy (e.g. ``model.lm_head``). + formulation: ``"weight"`` → ``F.linear(normed, head.weight)`` (portable; the + form the workloads use); ``"module"`` → ``head(normed)``. + + Returns: + The logits tensor. + """ + import torch.nn.functional as F + + if isinstance(residual, tuple): + residual = residual[0] + + mediator = self.model.interleaver.current + + if mediator._isolated_worker: + # the worker's modules are dummies — route the real compute to the host. + from ..interleaver import Events + + spec = { + "norm_path": norm.path if norm is not None else None, + "head_path": head.path, + "formulation": formulation, + } + return mediator.send(Events.UNEMBED, (residual, spec)) + + # in-process / fast lane: the real modules + weights are right here. + normed = norm(residual) if norm is not None else residual + if formulation == "weight": + return F.linear(normed, head.weight) + return head(normed) + + def steer(self, envoy, direction, alpha: float = 1.0): + """Add ``alpha * direction`` to ``envoy``'s output residual via a **replacement** + boundary write — the activation-steering / injection every steering cell does, made + to work on the isolated tier where an in-place ``hidden[:] = …`` silently no-ops. + + Steering reads the delivered activation and writes a modified one back. Done in + place (``hidden[:] = hidden + …``) that mutation never crosses the isolation + boundary: the worker mutates its *delivered clone*, no SWAP fires, and the host's + real activation is untouched — a silent no-op (the save of the steered residual + looks right, but nothing downstream changes). ``tracer.steer`` always performs a + *replacement* swap (assign ``envoy.output``), which ships the steered value back + over the existing ``Events.SWAP`` path, so it is correct in-process, on the fast + lane, AND in the isolated worker. Unlike :meth:`unembed`, steering touches no host + weights — only the delivered activation — so it needs no host round-trip and no + isolated/in-process branch: the eproperty setter routes the swap on either tier. + + Tuple outputs (attention modules, and transformer blocks on transformers <5) are + replaced whole — element ``[0]`` is steered and the rest of the tuple rides through + unchanged. + + Args: + envoy: the module whose output residual to steer (e.g. + ``model.transformer.h[6]``). + direction: the steering vector, broadcast over the residual and cast to its + dtype/device (so a float32 direction steers a float16 residual cleanly). + alpha: the steering coefficient — a scalar or any tensor broadcastable over the + residual. Defaults to ``1.0``. + + Returns: + The steered residual tensor (element ``[0]`` for tuple outputs). + """ + out = envoy.output + is_tuple = isinstance(out, tuple) + hidden = out[0] if is_tuple else out + + direction = direction.to(dtype=hidden.dtype, device=hidden.device) + steered = hidden + alpha * direction + + if is_tuple: + envoy.output = (steered, *out[1:]) + else: + envoy.output = steered + + return steered + + def patch(self, envoy, value): + """Transplant ``value`` into ``envoy``'s output residual via a **replacement** + boundary write — the activation-patching / resampling transplant every patching cell + does, made to work on the isolated tier where an in-place ``hidden[:] = value`` + silently no-ops. + + Patching replaces a module's delivered activation with one captured elsewhere (a clean + run, a different prompt, a precomputed mean). Done in place that write never crosses the + isolation boundary: the worker mutates its *delivered clone*, no SWAP fires, and the + host's real activation is untouched — a silent no-op. ``tracer.patch`` always performs a + *replacement* swap (assign ``envoy.output``), which ships the transplanted value back + over the existing ``Events.SWAP`` path, so it is correct in-process, on the fast lane, + AND in the isolated worker. Like :meth:`steer`, it touches no host weights — only the + delivered activation — so it needs no host round-trip and no isolated/in-process branch. + + ``value`` replaces the residual whole (element ``[0]`` for tuple outputs); construct a + full-shape value (e.g. clone the residual and edit a slice) for partial patches, the + standard nnsight idiom. The value is cast to the residual's dtype/device, so a value + precomputed on CPU (the isolation case) transplants cleanly. + + Tuple outputs (attention modules, and transformer blocks on transformers <5) are + replaced whole — element ``[0]`` is transplanted and the rest of the tuple rides + through unchanged. + + Args: + envoy: the module whose output residual to replace (e.g. + ``model.transformer.h[6]``). + value: the replacement activation, broadcast/cast to the residual's dtype/device. + Must match the residual's shape (element ``[0]`` for tuple outputs). + + Returns: + The transplanted residual tensor (element ``[0]`` for tuple outputs). + """ + out = envoy.output + is_tuple = isinstance(out, tuple) + hidden = out[0] if is_tuple else out + + value = value.to(dtype=hidden.dtype, device=hidden.device) + + if is_tuple: + envoy.output = (value, *out[1:]) + else: + envoy.output = value + + return value + + def ablate(self, envoy, mode: str = "zero"): + """Knock out ``envoy``'s output residual via a **replacement** boundary write — the + lesion-study ablation every ablation cell does, made to work on the isolated tier where + an in-place ``hidden[:] = 0`` silently no-ops. + + Ablation replaces a component's output with a baseline to measure how the prediction + degrades. Done in place that write never crosses the isolation boundary: the worker + mutates its *delivered clone*, no SWAP fires, and the host's real activation is + untouched — a silent no-op (the model never sees the lesion). ``tracer.ablate`` always + performs a *replacement* swap, riding the existing ``Events.SWAP`` path, so it is + correct in-process, on the fast lane, AND in the isolated worker. Like :meth:`steer`, it + touches no host weights, so it needs no host round-trip and no isolated/in-process + branch. + + Modes: + ``"zero"`` — replace with zeros (zero ablation; pushes the residual off-distribution + but is the simplest knockout). + ``"mean"`` — replace with the within-sequence mean: each position becomes the + per-example mean over the token dimension, keeping the average magnitude while + removing position-specific deviation. This is the *self-contained* mean. For + *reference-distribution* mean ablation (the mean activation over a dataset, per + docs/patterns/ablation.md), precompute that mean and transplant it via + :meth:`patch` — that mean is not derivable from a single forward. + + Tuple outputs (attention modules, and transformer blocks on transformers <5) are + replaced whole — element ``[0]`` is ablated and the rest of the tuple rides through + unchanged. + + Args: + envoy: the module whose output residual to ablate (e.g. + ``model.transformer.h[6]``). + mode: ``"zero"`` (default) or ``"mean"``. + + Returns: + The ablated residual tensor (element ``[0]`` for tuple outputs). + """ + out = envoy.output + is_tuple = isinstance(out, tuple) + hidden = out[0] if is_tuple else out + + if mode == "zero": + ablated = torch.zeros_like(hidden) + elif mode == "mean": + ablated = hidden.mean(dim=-2, keepdim=True).expand_as(hidden).contiguous() + else: + raise ValueError( + f"tracer.ablate mode must be 'zero' or 'mean', got {mode!r}" + ) + + if is_tuple: + envoy.output = (ablated, *out[1:]) + else: + envoy.output = ablated + + return ablated + def barrier(self, n_participants: int): """ nnsight barrier: A synchronization primitive for coordinating multiple concurrent invocations in nnsight. @@ -735,6 +955,13 @@ def __call__(self): mediator = self.model.interleaver.current + # Isolated worker: each invoke runs in its own process with its own Barrier + # copy, so it can't count cross-invoke. Send the TARGET count and let the host + # accumulate participants across mediators (see handle_barrier_event). + if mediator._isolated_worker: + mediator.send(Events.BARRIER, self.n_participants) + return + self.participants.add(mediator.name) if len(self.participants) == self.n_participants: diff --git a/src/nnsight/intervention/transport.py b/src/nnsight/intervention/transport.py new file mode 100644 index 000000000..dc291496d --- /dev/null +++ b/src/nnsight/intervention/transport.py @@ -0,0 +1,839 @@ +"""Socket-backed :class:`MediatorChannel` implementations for running an +intervention worker in a *separate process* from the model forward pass. + +This is the socket-transport stage of the mediator-isolation work +(docs/developing/mediator-isolation-harness-plan.md): the same six-event protocol +(VALUE/SWAP/SKIP/BARRIER/END/EXCEPTION), but the worker<->main handoff rides on an +``AF_UNIX`` socket instead of the in-process one-slot queues. + +Two role-specific ends, because each side of the protocol only ever calls its own +half of :class:`~nnsight.intervention.interleaver.MediatorChannel`: + +- **Host** (main thread, model forward): calls ``wait_event`` / ``has_event`` / + ``get_event`` / ``restore_event`` / ``put_response``. Worker events arrive over + the socket into a *local one-slot buffer*; ``has_event`` / ``get_event`` / + ``restore_event`` operate on that buffer with **no wire traffic** (they mirror the + in-process flag/restore semantics), and only ``wait_event`` reads the socket. + This is required: ``restore_event`` re-stages an event the main thread will + consume on a later ``handle()`` pass, which is host-local, not a round-trip. +- **Worker** (intervention fn): calls ``put_event`` / ``wait_response`` / + ``get_response``. ``put_event`` sends a frame; ``wait_response`` blocks reading + the reply. + +Frame codec: length-prefixed ``pickle``. This socket path forks the worker, so both ends +are mutually trusted and pickle-both-ways is fine. Two follow-ups (see the plan): + +1. **Security:** the *worker->host* direction MUST NOT plain-``pickle.loads`` untrusted + worker bytes (a ``__reduce__`` gadget would be a host-side RCE). **Done for the active + GPU path** (:class:`CudaIpcHostChannel` / :class:`CudaIpcWorkerChannel`): worker->host + frames carry array data out-of-band (the GPU buffer / safetensors) and the small + remaining structure is a **closed value algebra** decoded WITHOUT a pickle VM + (:func:`_codec_dumps` / :func:`_codec_loads`) — only primitives / containers / torch + dtype-device / out-of-band array headers cross; anything else is rejected at the worker + at encode (:class:`BoundaryValueError`). host->worker stays a normal pickle + (host-authored, trusted). The legacy AF_UNIX socket channels below + (:class:`SocketHostChannel`, :class:`ShmSocketHostChannel`) are NOT used by + ``isolate_mediators`` and still plain-``pickle.loads`` — route them through the value + codec before wiring them to an untrusted worker. + +2. **Performance (measured):** ``pickle`` of a torch tensor is the per-hook + bottleneck — ``dumps``+``loads`` ~22 ms per direction at 16.8 MB and **superlinear**; + a 16.8 MB round-trip is ~96 ms vs ~10 ms of actual socket transfer. Swapping the + codec to a **raw header + bytes** form (``memoryview``/``torch.frombuffer``, i.e. + ``safetensors``) is ~4×; a **shared-memory ring** (bulk never crosses the socket) is + ~8× and linear. This codec is the thing to replace, not the boundary. +""" + +from __future__ import annotations + +import builtins +import mmap +import os +import pickle +import socket +import struct +from typing import Any, Optional + +import torch + +from .interleaver import Events, MediatorChannel + +try: + from safetensors.torch import load as _st_load + from safetensors.torch import save as _st_save + + _HAS_SAFETENSORS = True +except Exception: # pragma: no cover + _HAS_SAFETENSORS = False + +try: + import numpy as _np + + _HAS_NUMPY = True +except Exception: # pragma: no cover + _np = None + _HAS_NUMPY = False + +_HEADER = struct.Struct("!I") # 4-byte big-endian length prefix + + +def _recvn(sock: socket.socket, n: int) -> bytes: + """Read exactly ``n`` bytes or raise ``EOFError`` if the peer closed.""" + buf = bytearray() + while len(buf) < n: + chunk = sock.recv(n - len(buf)) + if not chunk: + raise EOFError("mediator channel closed by peer") + buf += chunk + return bytes(buf) + + +def send_frame(sock: socket.socket, obj: Any) -> None: + payload = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL) + sock.sendall(_HEADER.pack(len(payload)) + payload) + + +def recv_frame(sock: socket.socket) -> Any: + (n,) = _HEADER.unpack(_recvn(sock, _HEADER.size)) + return pickle.loads(_recvn(sock, n)) + + +class SocketHostChannel(MediatorChannel): + """Host (main-thread) end of the channel — lives in the model process. + + Worker events are read off the socket into a single-slot buffer; the + main-thread protocol then drains that buffer exactly as it drained the + in-process event queue. Only :meth:`wait_event` touches the socket. + """ + + def __init__(self, sock: socket.socket): + self._sock = sock + self._pending: Any = None + self._has = False + + # --- worker -> main (read from socket into the local buffer) --- + def wait_event(self) -> None: + # Block for the worker's next event, unless one is already buffered + # (e.g. just restored by a provider-mismatch on this handle pass). + if not self._has: + self._pending = recv_frame(self._sock) + self._has = True + + @property + def has_event(self) -> bool: + return self._has + + def get_event(self) -> Any: + item = self._pending + self._pending = None + self._has = False + return item + + def restore_event(self, item: Any) -> None: + # Host-local re-stage (no wire traffic) — the main thread consumes this + # on a later handle() pass when the matching provider fires. + self._pending = item + self._has = True + + # --- main -> worker (send the reply for a pending event) --- + def put_response(self, value: Any) -> None: + send_frame(self._sock, value) + + def close(self) -> None: + try: + self._sock.close() + except OSError: + pass + + +class SocketWorkerChannel(MediatorChannel): + """Worker (intervention-fn) end — lives in the separate worker process. + + The worker only ever does ``send -> wait -> get``: push an event, block for + the reply, take it. No host-side staging here. + """ + + def __init__(self, sock: socket.socket): + self._sock = sock + self._response: Any = None + self._has_response = False + + # --- worker -> main (send the event) --- + def put_event(self, item: Any) -> None: + send_frame(self._sock, item) + + # --- main -> worker (block for the reply) --- + def wait_response(self) -> None: + self._response = recv_frame(self._sock) + self._has_response = True + + def get_response(self) -> Any: + item = self._response + self._response = None + self._has_response = False + return item + + def close(self) -> None: + try: + self._sock.close() + except OSError: + pass + + +# =========================================================================== # +# Fast path: shared memory + safetensors # +# =========================================================================== # +# Measurement showed the per-hook cost is pickle, not the boundary: pickling a torch +# tensor is ~22 ms per direction at 16.8 MB (superlinear), while the socket moves +# the bytes in ~10 ms. This path fixes both: +# - tensor BULK travels through a shared-memory region (memfd), so it never +# crosses the socket — only a tiny control frame does; +# - tensors are encoded with **safetensors** (a safe, no-code-execution format), +# which also closes the jail->host untrusted-deserialize hole for the bulk. +# +# Still pickle: the small CONTROL frame (the non-tensor structure + the requester +# string + the byte length). Hardening that to a restricted decoder is the +# remaining jail->host security item (see the module header / the plan). + +_TENSOR_TAG = "__nnsight_shm_t__" + + +def _split_tensors(obj: Any, store: dict) -> Any: + """Walk ``obj``; pull every tensor into ``store`` (keyed by index), leaving a + ``{_TENSOR_TAG: i}`` placeholder. Non-tensor structure is returned as-is.""" + if torch.is_tensor(obj): + i = len(store) + store[str(i)] = obj.detach().contiguous().cpu() + return {_TENSOR_TAG: i} + if type(obj) is tuple: + return tuple(_split_tensors(x, store) for x in obj) + if type(obj) is list: + return [_split_tensors(x, store) for x in obj] + if type(obj) is dict: + return {k: _split_tensors(v, store) for k, v in obj.items()} + return obj + + +def _merge_tensors(skel: Any, tensors: dict) -> Any: + """Reverse of :func:`_split_tensors` — re-inject tensors into the skeleton.""" + if type(skel) is dict: + if len(skel) == 1 and _TENSOR_TAG in skel: + return tensors[str(skel[_TENSOR_TAG])] + return {k: _merge_tensors(v, tensors) for k, v in skel.items()} + if type(skel) is tuple: + return tuple(_merge_tensors(x, tensors) for x in skel) + if type(skel) is list: + return [_merge_tensors(x, tensors) for x in skel] + return skel + + +class ShmArena: + """A shared-memory region (anonymous ``memfd``) for moving tensor payloads + out-of-band of the control socket. + + The host creates one; the worker attaches to the SAME memfd via an inherited + fd (``os.fork``) or one passed into the jail (``pass_fds`` + ``SHM_FD`` env), + and both ``mmap`` it. Under the one-event-in-flight protocol a single region + is reused in strict alternation; the receiver always ``safetensors.load``s + (which copies into fresh tensors) before the next write, so there is no + aliasing between the live tensors and the buffer. + """ + + def __init__(self, size: int, fd: int = None, owns: bool = True): + if fd is None: + fd = os.memfd_create("nnsight-shm", 0) + os.ftruncate(fd, size) + self.fd = fd + self.size = size + self._owns = owns + self.buf = mmap.mmap(fd, size) + + @classmethod + def attach(cls, fd: int, size: int) -> "ShmArena": + """Attach to an existing memfd (inherited / passed-in) without owning it.""" + return cls(size, fd=fd, owns=False) + + def write(self, blob: bytes) -> int: + n = len(blob) + if n > self.size: + raise ValueError(f"payload {n} B exceeds shm arena {self.size} B") + self.buf[:n] = blob + return n + + def read(self, n: int) -> bytes: + return bytes(self.buf[:n]) + + def close(self) -> None: + try: + self.buf.close() + if self._owns: + os.close(self.fd) + except (OSError, ValueError): + pass + + +def send_shm(sock: socket.socket, arena: ShmArena, obj: Any) -> None: + """Write ``obj``'s tensors (safetensors) into ``arena`` and send only the + small control skeleton + byte length over the socket.""" + store: dict = {} + skel = _split_tensors(obj, store) + n = arena.write(_st_save(store)) if store else 0 + ctrl = pickle.dumps((skel, n), protocol=pickle.HIGHEST_PROTOCOL) + sock.sendall(_HEADER.pack(len(ctrl)) + ctrl) + + +def recv_shm(sock: socket.socket, arena: ShmArena) -> Any: + (clen,) = _HEADER.unpack(_recvn(sock, _HEADER.size)) + skel, n = pickle.loads(_recvn(sock, clen)) + tensors = _st_load(arena.read(n)) if n else {} + return _merge_tensors(skel, tensors) + + +# =========================================================================== # +# Fast path: GPU bounce buffer (CUDA IPC) # +# =========================================================================== # +# The chosen design (docs/developing/gpu-sandbox.md + mediator-gpu-trace- +# integration.md): the worker is GPU-enabled and shares ONE GPU "bounce buffer" +# with the host via CUDA IPC (mapped once before lockdown). Tensor bulk never +# leaves the GPU and never hits pickle — only a small control frame (the +# non-tensor skeleton + an offset table) crosses the pipe. +# +# Under the one-event-in-flight protocol the single buffer is reused in strict +# alternation, so the receiver MUST clone tensors out of the buffer before the +# next write can overwrite them (the clone-on-receive rule). A D2D clone is +# HBM-speed; "zero-copy" means no-PCIe-no-pickle, not no-copy. + + +def _dtype_from_str(s: str) -> torch.dtype: + """``"torch.bfloat16"`` -> ``torch.bfloat16`` (no eval, no getattr-default).""" + return getattr(torch, s.split(".")[-1]) + + +def pack_cuda(value: Any, buf: torch.Tensor) -> tuple: + """Write every tensor in ``value`` into the shared GPU ``buf`` (D2D), leaving a + non-tensor skeleton + an offset table describing where each tensor landed. + + Returns ``(skeleton, table)`` where ``table[str(i)] = (offset, nbytes, shape, + dtype_str)``. Offsets are 16-byte aligned so a ``uint8`` slice can be viewed as + any tensor dtype. Raises if the payload exceeds the arena. + """ + table: dict = {} + state = {"offset": 0} + + def walk(obj: Any) -> Any: + # The "Array" leaf of the boundary value algebra: any contiguous typed buffer — + # a torch.Tensor OR a numpy.ndarray — travels OUT-OF-BAND in ``buf``; only a small + # ``(offset, nbytes, shape, dtype, kind)`` header rides in the skeleton. numpy is + # bridged through torch (zero-copy view); the host re-materializes the right kind. + t = kind = None + if torch.is_tensor(obj): + t, kind = obj.detach().contiguous(), "torch" + elif _HAS_NUMPY and isinstance(obj, _np.ndarray): + try: + t = torch.from_numpy(_np.ascontiguousarray(obj)) + except (TypeError, ValueError) as e: + raise BoundaryValueError( + f"a numpy array of dtype {obj.dtype} is not transmittable across the " + f"isolation boundary (torch has no matching dtype): {e}" + ) + kind = "numpy" + if t is not None: + i = len(table) + flat = t.reshape(-1).view(torch.uint8) + n = int(flat.numel()) + offset = (state["offset"] + 15) & ~15 # 16-byte align + if offset + n > buf.numel(): + raise ValueError( + f"intervention value {offset + n} B exceeds GPU bounce buffer " + f"{buf.numel()} B" + ) + if n: + buf[offset : offset + n].copy_(flat) + table[str(i)] = (offset, n, tuple(t.shape), str(t.dtype), kind) + state["offset"] = offset + n + return {_TENSOR_TAG: i} + if type(obj) is tuple: + return tuple(walk(x) for x in obj) + if type(obj) is list: + return [walk(x) for x in obj] + if type(obj) is dict: + return {k: walk(v) for k, v in obj.items()} + return obj + + skel = walk(value) + return skel, table + + +def unpack_cuda(skel: Any, table: dict, buf: torch.Tensor) -> Any: + """Reverse of :func:`pack_cuda`. Each tensor is CLONED out of ``buf`` (so a + later reuse of the single buffer can't corrupt it — the clone-on-receive + rule) and re-injected into the skeleton.""" + tensors: dict = {} + for k, meta in table.items(): + offset, n, shape, dtype_str = meta[0], meta[1], meta[2], meta[3] + kind = meta[4] if len(meta) > 4 else "torch" + dtype = _dtype_from_str(dtype_str) + view = buf[offset : offset + n].view(dtype).reshape(shape) + # CLONE out of the shared single-slot buffer (so a later reuse can't corrupt it); + # a numpy "Array" leaf is re-materialized on the CPU as an ndarray. + tensors[k] = view.clone().cpu().numpy() if kind == "numpy" else view.clone() + return _merge_tensors(skel, tensors) + + +# --------------------------------------------------------------------------- # +# Worker->host frame codec — closed value algebra, no pickle VM # +# --------------------------------------------------------------------------- # +# The worker runs UNTRUSTED user code, so the host must never run pickle's VM on +# its frames (a ``__reduce__`` gadget = host RCE). Instead the boundary transmits +# a CLOSED VALUE ALGEBRA — never live objects — so there is no opcode that can +# call a function in the first place: +# +# Value = None | bool | int | float | str | bytes (primitives) +# | list | tuple | dict | set of Value (containers) +# | torch.dtype | torch.device (value leaves) +# | Array(...) -> {_TENSOR_TAG: i} (bulk, OUT-OF-BAND) +# +# Array leaves (torch tensors, numpy arrays) are already pulled out-of-band by +# ``pack_cuda``/``_split_tensors`` into the GPU buffer / safetensors, leaving only +# a ``{_TENSOR_TAG: i}`` header in the skeleton — itself plain data. So this codec +# only has to (de)serialize the algebra above; anything outside it is rejected at +# the WORKER, at ENCODE, with a clear message (``BoundaryValueError``) — never an +# encode-ok / decode-refuse split. Decoding is pure data assembly: no globals, no +# ``find_class``, no ``REDUCE``. A size cap + bounds-checked reads contain a +# decode bomb (DoS). (host->worker stays a normal pickle: host-authored, trusted.) + +_EXC_TAG = "__nnsight_iso_exc__" + +# Max bytes for the codec part of a worker->host frame (the structured envelope; +# bulk array data is out-of-band, so this stays tiny). Caps decode-time allocation. +_MAX_CODEC_BYTES = 64 << 20 + +# Value-algebra tags (one byte each). +_T_NONE, _T_FALSE, _T_TRUE, _T_INT, _T_FLOAT, _T_STR, _T_BYTES = range(7) +_T_LIST, _T_TUPLE, _T_DICT, _T_SET, _T_DTYPE, _T_DEVICE = range(7, 13) + + +class BoundaryValueError(TypeError): + """A value outside the boundary algebra was handed to the worker->host codec + (e.g. a custom object / framework type ``.save()``-d in the worker). Raised at + ENCODE, in the worker, so the failure names the offending value at its source.""" + + +class BoundaryDecodeError(ValueError): + """A worker->host frame was malformed, oversized, or carried an unknown tag. + Raised at DECODE, host-side, instead of trusting attacker-shaped bytes.""" + + +def _w_uvarint(out: bytearray, n: int) -> None: + while True: + b = n & 0x7F + n >>= 7 + out.append(b | 0x80 if n else b) + if not n: + return + + +def _codec_dumps(value: Any) -> bytes: + out = bytearray() + + def enc(v: Any) -> None: + # bool BEFORE int (bool is a subclass of int); identity check is exact. + if v is None: + out.append(_T_NONE) + elif v is True: + out.append(_T_TRUE) + elif v is False: + out.append(_T_FALSE) + elif type(v) is int: + out.append(_T_INT) + b = v.to_bytes((v.bit_length() + 8) // 8 or 1, "little", signed=True) + _w_uvarint(out, len(b)) + out.extend(b) + elif type(v) is float: + out.append(_T_FLOAT) + out.extend(struct.pack(" Any: + if len(data) > _MAX_CODEC_BYTES: + raise BoundaryDecodeError( + f"worker->host frame {len(data)} B exceeds the {_MAX_CODEC_BYTES} B cap" + ) + pos = 0 + n = len(data) + + def need(k: int) -> int: + nonlocal pos + if pos + k > n: + raise BoundaryDecodeError("truncated worker->host frame") + start = pos + pos += k + return start + + def ruvarint() -> int: + nonlocal pos + result = shift = 0 + while True: + if pos >= n: + raise BoundaryDecodeError("truncated varint") + b = data[pos] + pos += 1 + result |= (b & 0x7F) << shift + if not b & 0x80: + return result + shift += 7 + if shift > 63: # lengths/counts fit in 64 bits — reject a runaway varint + raise BoundaryDecodeError("varint too long") + + def dec() -> Any: + nonlocal pos + if pos >= n: + raise BoundaryDecodeError("truncated worker->host frame") + tag = data[pos] + pos += 1 + if tag == _T_NONE: + return None + if tag == _T_TRUE: + return True + if tag == _T_FALSE: + return False + if tag == _T_INT: + k = ruvarint() + return int.from_bytes(data[need(k) : pos], "little", signed=True) + if tag == _T_FLOAT: + return struct.unpack("host frame") + return value + + +def _rebuild_exc(name: str, msg: str) -> BaseException: + """Rebuild an exception from a (type-name, message) pair, host side. Resolves + ONLY builtin exception types (never an arbitrary callable); else RuntimeError.""" + cls = getattr(builtins, name, None) + if isinstance(cls, type) and issubclass(cls, BaseException): + try: + return cls(msg) + except Exception: # noqa: BLE001 — some exceptions have non-(str,) __init__ + pass + return RuntimeError(f"{name}: {msg}") + + +def _encode_worker_frame(event: Events, data: Any, push: Any, buf: torch.Tensor) -> tuple: + """Build the worker->host frame: array data (torch/numpy) into ``buf`` (D2D) / + safetensors (push CPU tensors) so the structured part is ARRAY-FREE, the rest + encoded with the closed value-algebra codec (decoded host-side without a pickle + VM). Returns ``(frame_bytes, had_tensors)``; the caller must ``cuda.synchronize()`` + before sending if ``had_tensors`` (the host clones from ``buf`` on a separate + context).""" + if event is Events.EXCEPTION: + # The exception object may be an arbitrary class; ship only (type-name, + # message) — both strings, in the value algebra — so no class crosses. + data = {_EXC_TAG: [type(data).__name__, str(data)]} + skel, table = pack_cuda(data, buf) + if push is not None: + pstore: dict = {} + pskel = _split_tensors(push, pstore) + if pstore and not _HAS_SAFETENSORS: + raise RuntimeError( + "safetensors is required to transmit cross_invoker tensors from an " + "isolated worker" + ) + pblob = _st_save(pstore) if pstore else b"" + else: + pskel, pblob = None, b"" + # The event rides as its string ``.value`` (not the enum). The codec only emits the + # value algebra; a value outside it raises ``BoundaryValueError`` HERE, in the + # worker, naming the offending object at its source. + payload = _codec_dumps((event.value, skel, table, pskel)) + return _HEADER.pack(len(payload)) + payload + pblob, bool(table) + + +def _decode_worker_frame(raw: bytes, buf: torch.Tensor) -> tuple: + """Reverse of :func:`_encode_worker_frame`, host side. The structured part is decoded + with the value-algebra codec (no pickle VM on untrusted bytes). Returns + ``(event, data, push)``.""" + (plen,) = _HEADER.unpack(raw[: _HEADER.size]) + off = _HEADER.size + event_value, skel, table, pskel = _codec_loads(raw[off : off + plen]) + pblob = raw[off + plen :] + event = Events(event_value) + data = unpack_cuda(skel, table, buf) + if event is Events.EXCEPTION and isinstance(data, dict) and _EXC_TAG in data: + name, msg = data[_EXC_TAG] + data = _rebuild_exc(name, msg) + if pskel is None: + push = None + else: + ptensors = _st_load(pblob) if pblob else {} + push = _merge_tensors(pskel, ptensors) + return event, data, push + + +class CudaIpcHostChannel(MediatorChannel): + """Host (main-thread) end of the GPU-bounce-buffer channel. + + Control frames ride an ``mp.Connection``; tensor bulk rides the shared CUDA + ``buf`` (mapped into the worker via CUDA IPC). Same one-event-in-flight buffer + semantics as :class:`SocketHostChannel`: worker events arrive on ``wait_event`` + and are unpacked (cloned) into a single-slot local buffer; ``has_event`` / + ``get_event`` / ``restore_event`` are host-local with no IPC. + """ + + def __init__( + self, + conn: Any, + buf: torch.Tensor, + timeout: Optional[float] = None, + startup_timeout: float = 180.0, + ): + self._conn = conn + self._buf = buf + self._pending: Any = None + self._has = False + # Per-wait timeout for a hung worker (infinite loop in user code). The + # FIRST event covers spawn + import + deserialize + run-to-first-request, + # which is slow, so it gets a generous startup_timeout; subsequent waits + # use the user timeout. On timeout the trace's ``finally: cancel()`` kills + # the worker (interleave -> cancel -> _iso.close). + self._timeout = timeout + self._startup_timeout = startup_timeout + self._started = False + # Optional callable -> dict of live host interleaver state to piggyback on + # each response (e.g. ``default_all``, set only after the worker spawns). + self.meta_provider = None + # Optional callable(dict) to merge a worker's pushed cross_invoker locals + # into the host-side shared variable store. None on the in-process path. + self.on_push = None + + # --- worker -> main --- + def wait_event(self) -> None: + if not self._has: + limit = self._startup_timeout if not self._started else self._timeout + if limit is not None and not self._conn.poll(limit): + raise TimeoutError( + f"sandboxed intervention exceeded {limit}s with no event " + f"— worker presumed hung (e.g. an infinite loop in user code)" + ) + try: + raw = self._conn.recv_bytes() + except (EOFError, OSError) as e: + # The pipe broke mid-protocol => the worker died (e.g. a segfault + # in user C-code, or the GPU process was OOM-killed). Surface a + # clean error so the trace's ``finally: cancel()`` tears down, + # instead of leaking a raw EOFError out of model.trace(). + raise RuntimeError( + "sandboxed intervention worker died during execution" + ) from e + # SECURITY: the worker is untrusted — decode with the no-pickle safe + # codec (a pickle.loads here would be a host-side RCE sink). + event, data, push = _decode_worker_frame(raw, self._buf) + if push is not None and self.on_push is not None: + self.on_push(push) # cross_invoker: merge into the host var store + self._started = True + self._pending = (event, data) + self._has = True + + @property + def has_event(self) -> bool: + return self._has + + def get_event(self) -> Any: + item = self._pending + self._pending = None + self._has = False + return item + + def restore_event(self, item: Any) -> None: + self._pending = item + self._has = True + + def reset(self) -> None: + """Clear single-slot + per-job state so a recycled worker's channel is fresh. + + Re-arms the generous startup timeout (the next job's first event covers a + fresh deserialize + run-to-first-request) and drops the per-mediator + ``meta_provider`` / ``on_push`` bindings (re-set on the next acquire).""" + self._pending = None + self._has = False + self._started = False + self.meta_provider = None + self.on_push = None + + # --- main -> worker --- + def put_response(self, value: Any) -> None: + skel, table = pack_cuda(value, self._buf) + if table: + # pack_cuda's D2D copies are async on this context's stream; the Pipe + # only orders the CPU side. Without this sync the WORKER's unpack_cuda + # clone (a SEPARATE CUDA context) can read the buffer before our copy + # has run device-side -> silent corruption. (The proven prototype + # gpu_sandbox.py:43 / gpu_worker.py:57 synchronized here; the port + # dropped it — tests passed only because gpt2 tensors are tiny.) + torch.cuda.synchronize() + meta = self.meta_provider() if self.meta_provider is not None else None + self._conn.send((skel, table, meta)) + + def close(self) -> None: + try: + self._conn.close() + except OSError: + pass + + +class CudaIpcWorkerChannel(MediatorChannel): + """Worker (intervention-fn) end of the GPU-bounce-buffer channel. + + Only ever does ``put_event -> wait_response -> get_response``. Tensors in the + outgoing event payload are written into the shared buffer; the response's + tensors are read (cloned) back out of it. + """ + + def __init__(self, conn: Any, buf: torch.Tensor): + self._conn = conn + self._buf = buf + self._response: Any = None + self._has_response = False + # Optional callable(dict) applied to piggybacked host state on each response. + self.on_meta = None + # Optional callable() -> dict of cross_invoker locals to push to the host + # store on each event (None / returns None on non-cross-invoke traces). + self.push_provider = None + + # --- worker -> main --- + def put_event(self, item: Any) -> None: + event, data = item + push = self.push_provider() if self.push_provider is not None else None + # SECURITY: encode WITHOUT pickle — the host must never unpickle untrusted + # worker data. Tensors ride the shared GPU buffer; only tagged JSON crosses. + frame, had_tensors = _encode_worker_frame(event, data, push, self._buf) + if had_tensors: + # Async D2D copies must finish before the HOST (separate CUDA context) + # clones them out of the buffer. See CudaIpcHostChannel.put_response. + torch.cuda.synchronize() + self._conn.send_bytes(frame) + + # --- main -> worker --- + def wait_response(self) -> None: + skel, table, meta = self._conn.recv() + if meta is not None and self.on_meta is not None: + self.on_meta(meta) + self._response = unpack_cuda(skel, table, self._buf) + self._has_response = True + + def get_response(self) -> Any: + item = self._response + self._response = None + self._has_response = False + return item + + def close(self) -> None: + try: + self._conn.close() + except OSError: + pass + + +class ShmSocketHostChannel(SocketHostChannel): + """Host channel whose tensor payloads ride a :class:`ShmArena` instead of the + socket. Same one-event-in-flight buffer semantics as the parent.""" + + def __init__(self, sock: socket.socket, arena: ShmArena): + super().__init__(sock) + self._arena = arena + + def wait_event(self) -> None: + if not self._has: + self._pending = recv_shm(self._sock, self._arena) + self._has = True + + def put_response(self, value: Any) -> None: + send_shm(self._sock, self._arena, value) + + +class ShmSocketWorkerChannel(SocketWorkerChannel): + """Worker channel whose tensor payloads ride a :class:`ShmArena`.""" + + def __init__(self, sock: socket.socket, arena: ShmArena): + super().__init__(sock) + self._arena = arena + + def put_event(self, item: Any) -> None: + send_shm(self._sock, self._arena, item) + + def wait_response(self) -> None: + self._response = recv_shm(self._sock, self._arena) + self._has_response = True diff --git a/src/nnsight/schema/config.py b/src/nnsight/schema/config.py index c7f6a00f1..ad16f5303 100755 --- a/src/nnsight/schema/config.py +++ b/src/nnsight/schema/config.py @@ -29,6 +29,7 @@ class AppConfigModel(BaseModel): CACHE_DIR: str = "~/.cache/nnsight/" CROSS_INVOKER: bool = True TRACE_CACHING: bool = False + FAST_LANE: bool = True def __setattr__(self, name, value): if name == "TRACE_CACHING" and value is True: