Skip to content

[MRG] Make BSP-OT random seeding reproducible across calls - #898

Open
xihaian251 wants to merge 3 commits into
PythonOT:masterfrom
xihaian251:bsp-seed-repro
Open

xihaian251 wants to merge 3 commits into
PythonOT:masterfrom
xihaian251:bsp-seed-repro

Conversation

@xihaian251

Copy link
Copy Markdown
Contributor

Why

ot.bsp.compute_bspot_bijection(..., seed=...) — and therefore random_state of
ot.solve_sample(..., method="bsp") — is documented as a reproducibility knob, but it only held for the
first call made on a thread. Reported with measurements in #897.

// ot/bsp/BSP-OT_header_only.h:547-553
inline std::uint64_t global_seed = 0;
inline std::mt19937& rng() { thread_local std::mt19937 eng(global_seed); return eng; }

// ot/bsp/bsp_wrapper.cpp:70-73
double BSPOT_wrap(..., int seed) { global_seed = seed; ... }

eng is a thread-local object whose constructor consumes global_seed the first time rng() is used on
that thread and never again, so on later calls global_seed = seed cannot reach it and every draw
continues the previous stream. Three identical calls in one process returned three different bijections
and three different costs (0.786970722 / 0.786334761 / 0.793428761 on the instance in #897).

The same lines also cost solution quality: the candidate matchings are produced by
#pragma omp parallel for (ot/bsp/bsp_wrapper.cpp:40-48) and each worker builds its engine from the
same global_seed, so with the default static schedule the iterations at equal positions in their
chunks consume identical subsequences and return identical matchings. Measured, n_plans=64:
64 distinct candidates at 1 thread, 32 at 2, 16 at 4, 8 at 8, 4 at 16 — i.e. ceil(n_plans / n_threads)
distinct plans, the rest being copies that MergePlansNoPar then walks over again.

What

Both changes are inside the loop that generates the candidate plans, in ot/bsp/bsp_wrapper.cpp:

  • re-seed the (still thread_local) engine at the start of each iteration with a value derived from
    (user_seed, plan index), and pass seed down as a parameter instead of reading the mutable global;
  • construct the BijectiveBSPMatching per iteration, so the loop body no longer writes state shared
    with other iterations (random_pivot was written by every thread of the team).

The derivation must not be additive: with user_seed + task, plan t of seed s is plan t-1 of seed
s+1, so changing the seed would leave nearly the whole candidate set unchanged — hence the SplitMix64
finalizer. The vendored amalgamated header, the Cython layer, the Python API and every default are
untouched.

Contract

For a fixed build, compute_bspot_bijection(X, Y, n_plans, p, initial_perm, gaussian_slicing, seed) is
now a function of its arguments: repeat-equality, equality after unrelated calls, equality with a first
call in a fresh interpreter, and equality across OMP_NUM_THREADS. Candidate plan i depends on
(seed, i) alone, which is what the new "seed and index" test checks by comparing the first rows of the
stacks returned for n_plans=4 and n_plans=37.

Two deliberate non-claims. Results are reproducible per build, not across standard libraries:
std::mt19937::seed is specified by the standard, but std::normal_distribution and
std::uniform_int_distribution are implementation-specified, so a guarantee across libstdc++/libc++/MSVC
would need self-contained distributions. And different seeds are not required to give different
bijections. For the same reason the tests do not assert that the candidate matchings are pairwise
distinct — independent streams can legitimately return the same bijection (for Y = 2X + t every BSP
candidate is the identity, and on a 20-point instance I measured 120 different seeds producing only 114
distinct matchings). Diversity of the candidates is the consequence that #897 reports as a quality
metric (4/64 → 64/64 here), not something a test should promise.

For a given seed the returned permutation does change; unavoidable, since the old sequence depended on
the call history and on the CPU count, so there was no stable behaviour to keep. Existing assertions on
fixed configurations (test_bsp_ot_exact_identity, test_bsp_ot_identity_null_cost,
test_bsp_ot_bijective, test_bsp_ot_torch_backend) hold for any stream and still pass.

Added cost: one mt19937 state initialisation (624 words) per candidate plan, i.e. O(1) in n
against the O(n log n) work of generating that plan.

Tests

Six tests added to test/test_bspot.py; no historical cost or permutation is hardcoded, so they do not
freeze the new random sequence. Two of them spawn child interpreters (cross-process reproducibility and
invariance to OMP_NUM_THREADS is otherwise untestable); the child also reports the SHA-256 of the
extension it loaded, and the parent asserts it matches its own, so a child importing a different build
cannot silently pass.

Measured on Windows 11, MSVC 14.44 with /openmp, 16 logical CPUs, Python 3.13.1, numpy 2.4.6,
torch 2.10.0+cpu, jax 0.12.0, in a fresh clone at 5cb82c9, built with
python setup.py build_ext --inplace:

  • master C++ + these tests: 5 failed, 7 passed (the structure test passes on master as well, which is
    correct — it is not the defect);
  • this branch: 12 passed, and also 12 passed at OMP_NUM_THREADS=1, 2, 4, 8, 16;
  • pytest test/test_solvers.py -k bsp: 6 passed;
  • a contract sweep over 39 configurations (sizes 50–500, dims 2/3/80, n_plans 1–256, both
    slicing modes, several seeds; each also compared against the exact ot.emd value) checks
    repeat-equality, history-equality, permutation structure, and consistency between the cost the solver
    reports for its own permutation and the cost of that permutation: 39 configurations pass, 0 fail.
    The same script against the master build fails every one of the 39.
  • ruff check on the changed test file reports the same 5 pre-existing findings as HEAD, so no new
    lint.

Known limitation, stated rather than hidden: on a build or machine where OpenMP is unavailable (POT's
setup.py drops the flag silently) or pinned to one thread, the thread-count test cannot distinguish
the old behaviour from the new one; and a change that simply deleted #pragma omp parallel for would
also pass these tests — staying parallel is a performance property, not something these assertions
cover.

Out of scope

n_plans=0 crashes MergePlansNoPar (it indexes an empty plan vector), and merge_bijections reads
plans.shape[1] as the number of plans, which is an out-of-bounds read when n_plans < n. Both are
separate bugs; I will report them separately instead of mixing them into this change.

Fixes #897

ot.bsp.compute_bspot_bijection documents `seed` as a reproducibility knob, but the
BSP-OT engines are `thread_local std::mt19937` objects constructed once per thread
from `global_seed` (BSP-OT_header_only.h:549-553), so the `global_seed = seed`
assignment in BSPOT_wrap reaches the engine of the first call made on that thread
only: a repeated call continued the previous stream and returned different plans,
different intermediary stacks and a different cost (Issue PythonOT#897).

The same lines also cost solution quality, because the candidate plans are
generated by `#pragma omp parallel for` and each worker built its engine from the
same `global_seed`: iterations at equal positions in their chunks consumed
identical subsequences and produced identical matchings, so `n_plans=64` yielded
ceil(n_plans / n_threads) distinct candidates on this machine and `MergePlansNoPar`
walked over the copies.

Re-seed the engine at the start of each iteration from a SplitMix64 mix of
(user_seed, plan index) and pass the seed down as a parameter, so that candidate
plan i is a function of the arguments alone. The mixing must not be additive: with
`user_seed + task`, plan t of seed s is plan t-1 of seed s+1, so changing the seed
would leave nearly the whole candidate set unchanged. The matching object is now
constructed per iteration, which removes the only mutable state shared between the
iterations of this loop, and the now dead `global_seed = seed` write is dropped;
those two thread-safety edits do not change the result for a given seed.
Repeated calls, interleaved history, a fresh-interpreter reference, invariance to
OMP_NUM_THREADS, and the property that candidate plan i is determined by
(seed, i), which is checked by comparing the first rows of the candidate stacks
returned for n_plans=4 and n_plans=37. Five of the six fail on current master.

The candidate matchings are deliberately not asserted to be pairwise distinct:
independent streams can legitimately return the same bijection -- for a similarity
transform every BSP candidate is the identity -- so distinctness is a property of
the instance and of the standard library's distributions rather than of the
seeding, and PythonOT#897 reports it as a quality metric.

No historical cost or permutation is hardcoded, so the tests do not freeze the new
random sequence. The children report the extension hash they loaded and the parent
asserts it matches, so a child importing another build cannot silently pass.
Required by .github/workflows/check_release.yml, which fails any pull request that
does not touch RELEASES.md.
@xihaian251

Copy link
Copy Markdown
Contributor Author

The red checks are an infrastructure failure, not a test failure: every failed or cancelled job died in
actions/checkout@v7 while cloning the deps/eigen submodule, e.g. (Lint job, annotations)

remote error: GitLab is currently unable to handle this request due to load (ID a4e936739ef26abb-MSP)
clone of 'https://gitlab.com/libeigen/eigen.git' into submodule path '/home/runner/work/POT/POT/deps/eigen' failed
The process '/usr/bin/git' failed with exit code 1

No job failed at an install, build, lint or test step. git ls-remote https://gitlab.com/libeigen/eigen.git
answers from outside CI now, so re-running the failed jobs should be enough — I cannot trigger a re-run
from a fork without write access.

For reference, the two jobs that did get past checkout are green on head 3173c82:

  • windows: 1541 passed, 123 skipped, 4 xfailed
  • macos: 1545 passed, 119 skipped, 4 xfailed

with test/test_bspot.py at 11 passed / 1 skipped (test_bsp_ot_torch_backend, torch is not installed in
those jobs) on both, including the six new BSP seeding tests. Check updated RELEASES.md file is green.

@xihaian251

Copy link
Copy Markdown
Contributor Author

Hi maintainers,

Could someone with write access please re-run the failed jobs for PR #898?

The failures occurred during actions/checkout because GitLab rejected the Eigen submodule clone due to server load. The affected jobs never reached compilation or testing.

Windows and macOS have already passed.

Affected runs:

A re-run would help verify whether there are any actual code or test failures.

Thanks for your help!

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

BSP-OT seed is not reproducible across repeated calls

1 participant