BitGenerator support

[flying-sheep]

See

• https://numpy.org/doc/stable/reference/random/c-api.html
• https://numpy.org/doc/stable/reference/random/extending.html

Fixes #498

The idea is to have a safe wrapper around the npy_bitgen struct that implements rand::RngCore. That way pyo3 functions could be passed a np.random.Generator, get that wrapper from it, and pass it to Rust APIs, which could then call its methods repeatedly.

The way it’s implemented, the workflow would look like this:

1. acquire GIL
2. cast a np.random.BitGenerator instance into a numpy::random::PyBitGenerator.
3. call .lock() on it to get a numpy::random::PyBitGeneratorGuard.
4. release GIL
5. call functions on guard object without needing to hold the GIL

TODO:

• I see local crashes when running all tests, so there’s probably some UB, I’d appreciate help to fix it.

Safety

If somebody releases the threading lock of the BitGenerator while we’re using it, this isn’t safe 🤔

API design options

I could make this more complex by adding a new trait that is implemented by both PyBitGenerator and PyBitGeneratorGuard, allowing to choose if someone wants to

• use the PyBitGenerator’s random_* methods directly on that object while holding the GIL and without locking it
• use it like it’s used now, by locking the np.random.BitGenerator and returning a GIL-free object that can be used.

but for now I just implemented the use case that’s actually desired.

[Icxolu]

This looks like a useful addition! Thanks for working on it. I'm definitely not an expert here, but I left a few comment about things that stood out to me. Let me know what you think.
Also, are there any differences between numpy v1 and v2 that we need to consider?

[Icxolu]

This should be removed

[flying-sheep]

sure, will do when I’m done. I like working on multiple machines, and I don’t like re-doing settings for individual projects

[Icxolu]

I may have found a problem:

This fails as intended:

Python::with_gil(|py| {
    let obj = get_bit_generator(py)?;
    let a = obj.lock()?;
    let b = obj.lock()?;

    Ok::<_, PyErr>(())
})
.unwrap();

returning

called `Result::unwrap()` on an `Err` value: PyErr { type: <class 'RuntimeError'>, value: RuntimeError('BitGenerator is already locked'), traceback: None }

But this does not fail:

Python::with_gil(|py| {
    let a = get_bit_generator(py)?.lock()?;
    let b = get_bit_generator(py)?.lock()?;

    Ok::<_, PyErr>(())
})
.unwrap();

and crucially it gives the same pointers:

[src/random.rs:113:18] ptr = 0x00007b9f6be44cc0
[src/random.rs:113:18] *ptr = bitgen_t {
    state: 0x00007b9f6be44d08,
    next_uint64: 0x00007b9f6837d320,
    next_uint32: 0x00007b9f6837d370,
    next_double: 0x00007b9f6837d3f0,
    next_raw: 0x00007b9f6837d320,
}
[src/random.rs:113:18] ptr = 0x00007b9f6be44cc0
[src/random.rs:113:18] *ptr = bitgen_t {
    state: 0x00007b9f6be44d08,
    next_uint64: 0x00007b9f6837d320,
    next_uint32: 0x00007b9f6837d370,
    next_double: 0x00007b9f6837d3f0,
    next_raw: 0x00007b9f6837d320,
}

So when using multiple threads, for example multiple tests running in parallel, we have a data race on the state. I think we need a lock across all instances to make this work.

[mejrs]

Maybe we should skip the guard part and just lock and unlock within the RngCore implementation itself. Can you give an example for why you'd want this, and why the api has this form? Why would someone want to use this rather than the RngCore impl that rand ships with? Maybe we can come up with a better design.

[flying-sheep]

When implementing Python-facing APIs, having a rng: np.random.Generator parameter is common. I want to write code that actually respects that parameter and uses it instead of ignoring it or calling it once to seed the actual generator.

[juntyr]

Would it also possible to go the other way, i.e. provide a rand rng from Rust to Python as a numpy BitGenerator?

[flying-sheep]

Yes, that's part of numpy's API as well!

[flying-sheep]

OK, done.

UB fixed

turns out the BitGenerator itself has allocated memory, not just the _capsule, so letting that be freed was the bug. Now we bump its refcount and keep it around until we’re done.

double locking issue fixed

@Icxolu as you found in #499 (comment), the lock is re-entrant, so now we just keep a HashSet of locked bitgens so same-thread guards are unique.

PS: MSRV is 1.83, and in 1.84 ptr.addr() was stabilized, which would be a slightly better API than as usize here

[flying-sheep]

And done! Please take a look!

@Icxolu I followed your initial advice and converted it from a lock() -> Guard API into a lock(|guard| ...) API, getting rid of the silently-fallible Drop implementation.

I have gathered some experience using numpy’s generators and figured that using .spawn_one() or .spawn(n) to get more freedom here is totally acceptable.

outstanding questions & tasks

• naming: we have PyBitGenerator for the maybe-shared one and BitGenerator for the exclusive access one (&mut when locked, by value when owned by us). These names could be better I think.
• add new()/from_* functions. With the new API it/they would return a owned-by-us BitGenerator (BitGenerator support #499 (comment)) could maybe be combined with this, and change the BitGenerator so it can hold one of multiple backends
  • allow passing a rng_core object to numpy as BitGenerator (BitGenerator support #499 (comment))
• finally: remove .vscode stuff (BitGenerator support #499 (comment))

[Icxolu]

Thanks for picking this back up! It might take me a bit to get my head back into this, but I will try to give this a read as soon as I find some time.

[flying-sheep]

I’m quite happy with it now :D

Only caveat (apart from naming): There’s a lot of nit-picking possible with the from_numpy API which is why I initially didn’t implement any factory methods. If you want changes there, maybe I’ll just remove that part and you can implement it however you want, or you tell me exactly what you want and I do it.

[flying-sheep]

could alternatively be an Into implementation.

[Icxolu]

I gave this a first pass. For now I mostly looked at the technical stuff and did not really think hard about naming etc. Unfortunately I think there is still a potential soundness issue with lock, see that comment for more details. I believe spawn does not have the same issue, so even if we can't get lock working, this one should be possible.

[Icxolu]

I think the API as currently implemented is still unsound. f is allowed to call into Python, so I can do something like this

 bit_gen.lock(|b| {
     bit_gen.getattr("lock")?.call_method0("release")?;
     Ok::<_, PyErr>(b.next_double())
})

and release the lock early. We could prevent this by requiring the closure to be Ungil like detach, but I'm still not certain that arbitrary Python code could not be a problem as well.

[flying-sheep]

making it Ungil means we can no longer have these tests:

• lock_keeps_bit_generator_alive
• double_lock_*

But maybe we don’t need them?

I did this in 6c0f4c7.

Now the closure and its return value are required to be Ungil + Send, and we detach ourselves. This makes the examples actually nicer, although I’m not sure why the Send bound is necessary now and wasn’t before.

[Icxolu]

The Send bound is needed, because the closure to detach has to be Ungil as well. Ungil should be an auto trait, but given that these are still unstable, PyO3 piggy backs on Send as it matches the semantics quite closely most of the time. Given that implementing the Fn traits is also unstable, the only way to satisfy the Ungil bound is to rely on the blanket which requires the captures to be Send.

---

Unfortunately I found another issue with the API 😞 The mutable reference we pass into the closure allows users to smuggle out the generator with std::mem::swap:

Python::attach(|py| {
    // Possibly shared generator
    let py_generator = get_bit_generator(py)?;

    // Exclusive generator
    let mut exclusive_gen = BitGenerator::from_numpy(py, NumpyBitGenerator::PCG64)?;

    // This allows swapping the generator, allowing unsynchronized access the the shared generator below
    py_generator
        .lock(|shared_locked_gen| std::mem::swap(shared_locked_gen, &mut exclusive_gen))?;

    // This is a shared generator now, so unsynchronized access is unsound
    exclusive_gen.next_double();
    Ok(())
})

[Icxolu]

I think it would make more sense if this consumes the guard, so it can't be released twice.
We can use ManuallyDrop to prevent Drop from running as well in that case, no need for a field.

[Icxolu]

The Send bound is needed, because the closure to detach has to be Ungil as well. Ungil should be an auto trait, but given that these are still unstable, PyO3 piggy backs on Send as it matches the semantics quite closely most of the time. Given that implementing the Fn traits is also unstable, the only way to satisfy the Ungil bound is to rely on the blanket which requires the captures to be Send.

---

Unfortunately I found another issue with the API 😞 The mutable reference we pass into the closure allows users to smuggle out the generator with std::mem::swap:

Python::attach(|py| {
    // Possibly shared generator
    let py_generator = get_bit_generator(py)?;

    // Exclusive generator
    let mut exclusive_gen = BitGenerator::from_numpy(py, NumpyBitGenerator::PCG64)?;

    // This allows swapping the generator, allowing unsynchronized access the the shared generator below
    py_generator
        .lock(|shared_locked_gen| std::mem::swap(shared_locked_gen, &mut exclusive_gen))?;

    // This is a shared generator now, so unsynchronized access is unsound
    exclusive_gen.next_double();
    Ok(())
})

[Icxolu]

PyO3 actually dropped support for 3.8 PyO3/pyo3#6128 recently (lands in 0.30 I think), so rust-numpys next release will drop it as well.