Skip to content

feat: save and load a trained map without pickle - #11

Merged
andremsouza merged 2 commits into
masterfrom
feat/provenance-and-artifacts
Jul 30, 2026
Merged

feat: save and load a trained map without pickle#11
andremsouza merged 2 commits into
masterfrom
feat/provenance-and-artifacts

Conversation

@andremsouza

Copy link
Copy Markdown
Owner

save_npz / load_npz with no pickle, plus TrainingReport and SOMConfig. Last feature PR for
0.4.0. Two commits: the analysis tooling is separate from the feature so each reviews on its own.

The problem this had to solve

A map holds four callables — the neighborhood, two decays and the distance — and a callable cannot
go into a file without pickle. Only the neighborhood had a name registry. So "save a map" had no
obvious answer, which is why this one got a design pass before any code.

What a map actually contains, measured rather than assumed:

kind serialisable without pickle?
shape, input_len, rates, radii, cyclic, seed yes, JSON
the models yes, npz
four callables only by name
np.random.Generator its bit_generator.state is JSON-able

Decisions

One file, not a file plus a sidecar. A JSON string stores inside an .npz and loads fine under
allow_pickle=False. Provenance that can be separated from its artifact will be.

Strategies resolve by name. DECAY_FUNCTIONS and DISTANCE_FUNCTIONS are added beside the
existing NEIGHBORHOOD_FUNCTIONS, keyed by each function's own name. A map built from shipped
functions round-trips completely; one built with your own function records the name for provenance and
refuses to load silently:

ArtifactError: this map was saved with custom function(s) that cannot be restored from a name
(distance_function='cosine'). A name is all that was saved, because a callable cannot be written to a
file without pickle. Pass them back explicitly: load_npz(path, distance_function=...)

Both halves are tested — the error and that its suggested fix works. An error whose suggestion
doesn't work is worse than no suggestion.

A functools.partial falls out correctly: no name, so it is treated as custom rather than silently
reloading partial(exponential_decay, factor=3.0) as factor=2.0.

Both the seed and the generator state are saved, and this is the part worth reviewing closely. The
seed says where the stream began; the state says where it has reached. Only the state resumes it.

The test that proves it works

uninterrupted:  train(200)
interrupted:    train(100) → saveloadtrain(100)
assert max|difference| == 0.0

Stepwise mode deliberately — it draws samples from the generator, so a lost state changes which
samples arrive. Batch mode would pass this with the state thrown away.

And a second test asserts that re-seeding really would diverge, so the first cannot pass for the
wrong reason. Without it, discarding the state entirely would still go green if re-seeding happened to
give the same stream.

Security

The load path is the reason this PR exists in the shape it does.

Cannot happen: code execution (allow_pickle=False passed explicitly); loading an arbitrary
function by name (names resolve only through the registries — nothing is imported or evaluated); a
partly-read file treated as valid (member list, JSON, format version, weight dimensions and weight
shape all checked before a map is built — every failure is ArtifactError).

Can still happen: resource exhaustion, because an .npz is a zip. The docs say so as plainly as
they say the rest. allow_pickle=False buys "cannot execute code", not "safe to load anything".

The headline test doesn't assert that — it demonstrates it. It builds a __reduce__ payload that
creates a file, proves the payload is live, then shows the loader refusing a file containing it
with nothing created. Liveness is proven through pickle directly rather than by asking NumPy to
unpickle, which keeps the enabled form of NumPy's pickle flag out of the repository entirely — the
architecture profile forbids it at the artifact boundary, and a test that had to switch that rule off
to run would be the worse test.

Reviewing the load path against the OWASP deserialisation rules turned up one real gap, now fixed: a
crafted generator state produced a raw NumPy ValueError ("state must be for a PCG64 RNG") instead of
ArtifactError, so a bad file did not fail consistently. Four malformed shapes are now tested.

Bugs found while building this

  • The seed and the state were collapsed under one rng key, so rng["state"] was the inner PCG
    state rather than the full bit_generator.state. Every round trip failed with "state must be for a
    PCG64 RNG".
  • supplied.get(key, resolve(name)) evaluated the fallback eagerly, so passing a custom callable
    back still tried to resolve its name and raised — on exactly the argument that exists to accept it.
  • The public-method surface test ignored classmethods. It used inspect.isfunction, which is
    False for a classmethod accessed through the class, so it silently skipped load_npz and would
    have skipped any classmethod added later. A surface pin with a hole in it is worse than none.
  • tqdm was left undefined when absent rather than bound to None, which reads to a type checker
    as possibly-unbound at every use site however carefully the use is guarded. Pre-existing; surfaced
    from IDE diagnostics.

Also

  • The require_provenance WARN that the architecture profile has carried since refactor: split the package into a functional core and a thin shell #7 is cleared. The
    count stays at 13 only by coincidence — the strategy protocols added one positional-arg WARN of
    their own.
  • New docs page save-and-load.md, in the nav, including a worked example of reading an artifact's
    metadata without this package installed, since it is plain JSON.

Verification

369 tests, 100% coverage, ruff / ruff format / mypy --strict / mkdocs build --strict clean,
twine check passes, architecture profile --strict passes with 0 ERROR, bandit reports no issues
across src/ and benchmarks/, pip-audit reports no known vulnerabilities.


Second commit: the analysis extra

bandit, pip-audit and pylint, pinned exactly. Nothing gates a merge — ruff and mypy stay the
enforced checks. CI now names the extras it needs rather than --all-extras, so tools no CI job runs
are not installed on every matrix entry.

Two candidates were rejected rather than added, which is the part worth reading:

  • osv-scanner is a slopsquat risk. The PyPI package of that name is version 0.0.1, registered
    2026-07-16 (two weeks ago), with no author, no home page, and the summary "Reserved name
    placeholder. No functionality."
    The real tool is a Go binary from Google. Not installed.
  • semgrep was added, then removed. It pins mcp==1.23.3 exactly, and that version carries
    PYSEC-2026-3481/3482/3483 (fixed in 1.27.2/1.28.1) — so the pin makes them unremediable from here,
    and 1.172.0 is already the latest semgrep. Against that: bandit finds nothing in src/, and this is
    a pure-numpy library with no network, auth, crypto, SQL or templating, which is most of what
    semgrep's rulesets target. Low coverage gained for three CVEs that cannot be patched.
  • guarddog is genuine (DataDog) but first released 2022-11-28, after this project, so it fails the
    rule that a dependency should predate what it is added to.

Static analysis for review passes, pinned exactly like the rest of the tooling. Nothing here gates a
merge: ruff and mypy remain the enforced checks. bandit reports no issues across src/ and benchmarks/;
pip-audit reports no known vulnerabilities.

bandit excludes tests/ rather than disabling B101, so `assert` is still flagged in src/, where it
would be wrong -- asserts vanish under `python -O`, so validation has to raise.

CI names the extras it needs instead of using --all-extras, so review tools no CI job runs are not
installed on every matrix entry. `examples` is named explicitly because tests/test_end_to_end.py
imports matplotlib.

Two candidates were rejected rather than added:

- `osv-scanner` on PyPI is version 0.0.1, registered 2026-07-16, with no author, no home page, and the
  summary "Reserved name placeholder. No functionality." The real tool is a Go binary from Google.
  Installing a code-free package squatting the expected name is the slopsquat case the dependency
  rules exist for.
- `semgrep` pins `mcp==1.23.3`, which carries PYSEC-2026-3481/3482/3483 (fixed in 1.27.2/1.28.1). The
  exact pin makes them unremediable from here and 1.172.0 is already the latest. Against that, bandit
  finds nothing in src/, and this is a pure-numpy library with no network, auth, crypto, SQL or
  templating -- most of what semgrep's rulesets target. Revisit if the pin is relaxed.
- `guarddog` is genuine but first released 2022-11-28, after this project, so it fails the rule that a
  dependency should predate what it is added to.
Until now the only way to keep a trained map was `pickle`, which is arbitrary code execution on load.
`save_npz` and `load_npz` write one `.npz` holding the models as an array and everything else as JSON
beside them: the configuration, the seed, the generator state and the last training report.

A reloaded map continues the same random stream, so training it further produces what never stopping
would have produced. Both the seed and the generator's current state are saved, because the seed alone
restarts the stream and a resumed run would then diverge from an uninterrupted one without saying so.
A test asserts the two are identical, and a second test asserts that re-seeding really would differ,
so the first cannot pass with the state discarded.

Strategies are saved by name, resolved through registries: DECAY_FUNCTIONS and DISTANCE_FUNCTIONS are
added alongside the existing NEIGHBORHOOD_FUNCTIONS, keyed by each function's own name. A map built
from the shipped functions round-trips completely. One built with a caller's own function records the
name for provenance and raises ArtifactError naming the argument to pass it back through, which the
loader accepts. A functools.partial has no name, so it is treated as custom rather than reloaded as
the wrapped function with different arguments.

The seed and the generator state are separate keys in the metadata. Collapsing them conflates "where
the stream began" with "where it has reached", and assigning an inner PCG state where a full
bit_generator.state belongs fails with "state must be for a PCG64 RNG".

Also adds SOMConfig and a frozen TrainingReport, exposed as `som.config()` and `som.last_report`.
`train()` still returns the quantization error. `final_learning_rate` is None for batch training,
which has no step size: Eq. (8) is a weighted mean, and reporting the unused initial value would read
as though it had been applied. The training loops now return the values they finished on rather than
having `train` recompute them, so the report states what ran.

On the load path: allow_pickle=False is passed explicitly, names resolve only through the registries
so nothing from the file is imported or evaluated, and the member list, metadata JSON, format version,
weight dimensions and weight shape are all validated before a map is constructed. Every failure,
including a malformed generator state that NumPy rejects, surfaces as ArtifactError. A test builds a
payload that really would create a file, proves it executes when unpickled, and shows the loader
refusing a file containing it with nothing created.

What that does not buy is stated in the docs as plainly as what it does: an `.npz` is a zip, so a
hostile file can still attempt resource exhaustion through decompression.

`tqdm` is now bound to None when absent rather than left undefined, so a type checker reads the
guarded use as reachable instead of possibly-unbound at every call site.

The public-method surface test now includes classmethods. It used `inspect.isfunction`, which is
False for a classmethod accessed through the class, so it silently ignored `load_npz` and would have
ignored any classmethod added later.
@andremsouza
andremsouza merged commit 0d48ab6 into master Jul 30, 2026
8 checks passed
@andremsouza
andremsouza deleted the feat/provenance-and-artifacts branch July 30, 2026 12:50
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant