test(serialization): make the serialization suite runnable and add it to CI

[Hotragn]

Summary

The four serialization test files are not in the CI list, and two of them had stopped working entirely without anyone noticing. This makes the suite runnable, removes the order-dependence in it, corrects a comment in serialization.py that claims a check which isn't there, and adds all four files to CI.

Found while checking which test files CI does not cover.

What was broken

1. test_whitelist_serialization.py did not collect at all.

ImportError: cannot import name 'SERVER_MODULES_WHITELIST' from
'nnsight.intervention.serialization'

SERVER_MODULES_WHITELIST / _is_whitelisted_module were imported at module scope, so their removal took all 12 tests in the file down — including the nine that test source serialization and still pass. Seven tests in test_serialization_edge_cases.py failed the same way on PicklingProhibitedError.

Both are now resolved per-test and skip with a stated reason rather than erroring at import. The rest of each file runs, and these come back automatically if the APIs return — which felt better than deleting the maintainers' intent.

2. Four tests were order-dependent and passed by accident.

test_no_register_needed, test_serialization_includes_module_functions, and the two *_serialized_by_source tests all assert b"def normalize" in dumps(normalize). Serializing by source is opt-in:

before register: 47 bytes, source embedded = False
after  register: 887 bytes, source embedded = True

register mutates a process-global set in cloudpickle (_PICKLE_BY_VALUE_MODULES), and every trace calls get_local_env(), which registers local modules. So these four passed in a whole-file run — and every one of them failed when run alone:

test_non_whitelisted_function_serialized_by_source  -> 1 failed
test_serialization_includes_module_functions        -> 1 failed
test_no_register_needed                             -> 1 failed
test_non_whitelisted_class_serialized_by_source     -> 1 failed

They now use a fixture that registers and unregisters around each test. A second fixture lifts the entries back out of the global set so the by-reference half of the contract can be asserted wherever in the file it runs. Every test in the file now passes (or skips) in its own process.

3. test_no_register_needed asserted the opposite of current behaviour.

Its premise — that register() is unnecessary — was true under the whitelist design that has since been removed. Without registration, dumps emits a GLOBAL the server cannot import, which is the ModuleNotFoundError that docs/remote/register-local-modules.md exists to prevent. Replaced with two tests that pin the real contract in both directions:

• test_register_switches_to_serialization_by_value — registration is what makes source travel with the job
• test_unregistered_module_is_serialized_by_reference — without it, the payload is a bare GLOBAL mymethods.stateful normalize

4. serialization.py claimed a security check it no longer performs.

# SECURITY CHECK: Check function globals for prohibited modules/functions.
# This catches dangerous patterns like `import os; os.getcwd()` or
# `from subprocess import run; run(...)` where the prohibited object
# is captured in the function's globals.
captured_globals = slotstate.get("__globals__", {})

# CLOSURE: Extract closure values and names.

captured_globals is assigned and never read — it is the only remaining trace of the lint-time blacklist added in ccb1702. Left as-is, this reads to anyone auditing the remote path as though input is being validated there. Replaced with a note stating the check is absent, where it would go, and that its tests are already written.

I have not restored the blacklist itself — see #695 for that question. It would start rejecting user code that references pathlib, io, glob and friends, which is a policy call rather than something to slip into a test-coverage PR.

Result

before:  248 tests in CI
after:   323 passed, 12 skipped, 1 xpassed

75 previously-unrun tests, covering the path that ships user code to NDIF.

Note

The CI run will stay red until #693 lands. test_lm.py::TestGradients::test_backward_with_multiple_invokers is already failing on dev and main and is unrelated to this change. Happy to rebase this on top of #693, or to drop the workflow change into a follow-up if you'd rather land the test repairs first.

_{Run with transformers 5.15.1 / torch 2.9.1 / cloudpickle 3.1.2 on CPU.}

[JadenFiotto-Kaufman]

Hey @Hotragn ! Thank you so much for the recent PRs!

Right now im actively developing off the 0.8 branch here: https://github.com/ndif-team/nnsight/tree/0.8
Could you base your work off this instead? I'll still respond in your PRs as they are for now.

[Hotragn]

Thanks — moved over to 0.8, and closing this one: almost everything it fixes is already handled there.

Checked each piece against 0.8:

• Stale imports — gone. test_whitelist_serialization.py and test_serialization_edge_cases.py are consolidated into tests/test_serialization.py, with no SERVER_MODULES_WHITELIST / PicklingProhibitedError references left.
• The misleading # SECURITY CHECK comment and dead captured_globals — gone from 0.8's serialization.py.
• The CI gap — gone. 0.8 runs pytest tests/ -q --ignore=tests/vllm --ignore=tests/tp, so the whole directory is covered rather than an enumerated list. That's a better fix than the one I proposed.
• The order-dependent registration tests — no longer present in that form.

So the whole PR is redundant on 0.8. Nothing for you to action here.

Two things did come out of running the 0.8 suite, though:

1. fix(serialization): find the function's symtable child by type, not by count (Python 3.14) #698 — tests/test_serialization.py::TestScopeFiltering has two real failures on Python 3.14 (which pyproject.toml lists as supported). _function_referenced_names picks the function's symtable child by counting children, and PEP 649 adds an __annotate__ table to every def, so the guard trips and the scope filter is inert — reinstating the closed-file capture its docstring describes. CI runs 3.12, where PEP 649 isn't active, so it doesn't show.
2. feat(backward): explain the freed autograd graph in terms of invokes, and cover batched-invoke gradients #699 — the cross-invoke .backward() diagnostic from feat(backward): explain the freed autograd graph in terms of invokes #694, which still reproduces on 0.8, plus the batched-invoke gradient coverage 0.8 is missing.

Also updated #695 — most of what I raised there is resolved on 0.8; only the policy question remains, and it's a much smaller one now.

[Hotragn]

Will do — everything is on 0.8 now. Thanks for the pointer; it changed the picture quite a bit, so here's where things stand.

I re-checked all four PRs against 0.8 rather than porting them blind, and three of them turned out to be already fixed there:

• fix(util): traverse container subclasses so ModelOutput values are not silently skipped #692 (util.apply skipping container subclasses) — 0.8's apply is a full rewrite that handles all of it, and applyn is gone. Verified ModelOutput narrowing works on 0.8 (rel_err 2.1e-07, distinct objects per invoke). Closed.
• test(lm): fix the batched-gradient assertion that has kept CI red for a month #693 (the batched-gradient assertion) — that test doesn't exist on 0.8. Closed.
• test(serialization): make the serialization suite runnable and add it to CI #696 (serialization tests + CI) — stale imports gone, misleading comment gone, and pytest tests/ covers everything. Closed.

Two are still live, both verified against 0.8 first:

• fix(serialization): find the function's symtable child by type, not by count (Python 3.14) #698 — TestScopeFiltering has two genuine failures on Python 3.14. _function_referenced_names finds the function's symtable child by counting children, and PEP 649 gives every def an extra __annotate__ table, so the guard trips and the scope filter silently falls back to the block rule — bringing back the closed-file capture its docstring warns about. pyproject.toml lists 3.14 as supported; CI runs 3.12, where PEP 649 isn't active, so it doesn't surface. One-line fix plus regression tests. Full suite goes to 856 passed / 0 failed.
• feat(backward): explain the freed autograd graph in terms of invokes, and cover batched-invoke gradients #699 — the cross-invoke .backward() diagnostic (still reproduces on 0.8), plus the batched-invoke gradient coverage the suite is missing: test_backward.py is single-invoke only and test_batching.py is all @torch.no_grad(), so nothing asserted an invoke's gradient is its own rows. Built on the existing MLP + _BatchEnvoy pattern, so it's exact and adds no model download.

#699 will show the two TestScopeFiltering failures until #698 lands, since it's based on clean 0.8. Happy to stack it on #698 instead if you'd rather.

One note in case it's useful: pyproject.toml advertises 3.13 and 3.14 but CI only runs 3.12. #698 is exactly the kind of thing that gap hides — PEP 695 generics would trip the same symtable code path on 3.13. Glad to open a PR adding 3.13/3.14 to the test matrix if you want the coverage.