-
Notifications
You must be signed in to change notification settings - Fork 123
BitGenerator support #499
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
BitGenerator support #499
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This should be removed
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
sure, will do when I’m done. I like working on multiple machines, and I don’t like re-doing settings for individual projects
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I still don't love the drop impl, but with the way to manually release it with a Python
token, it may be acceptable. Maybe @davidhewitt has an idea and/or comments about the appoach. Otherwise I only have a few minor remarks.
Thanks for the comments, I’ll address them! The main issue is that I think I’m triggering UB somehow and I don’t know how: when running all tests, often some unrelated test run after this one crashes …
I didn’t forget about this either, will look! /edit: the C API for random is there since 1.19: https://numpy.org/doc/1.26/reference/random/c-api.html |
//! # use pyo3::prelude::*; | ||
//! use rand::Rng as _; | ||
//! # use numpy::random::{PyBitGenerator, PyBitGeneratorMethods as _}; | ||
//! # // TODO: reuse function definition from above? |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It feels like there should be a convenient way to get this. I'm thinking about something like
impl PyBitGenerator {
fn new(py: Python<'_>) -> PyResult<Bound<..>>;
}
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
there are many implementations, we’d have to cover all of them.
I’d rather leave this minimal until this PR is mostly done.
.getattr(intern!(py, "capsule"))? | ||
.downcast_into::<PyCapsule>()?; | ||
let lock = self.getattr(intern!(py, "lock"))?; | ||
// we’re holding the GIL, so there’s no race condition checking the lock and acquiring it later. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This may not be true under free-threaded Python. Is the lock known to be threadsafe and acquire
simply fails if the lock is already acquired? If not we may need to guard the whole module under cfg(not(Py_GIL_DISABLED))
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
it doesn’t fail, it hangs, but that’s configurable with a timeout or by making it non-blocking: https://docs.python.org/3/library/threading.html#threading.Lock.acquire
and it’s a threading.Lock
!
OK, with the |
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
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:
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. |
See
Fixes #498
The idea is to have a safe wrapper around the
npy_bitgen
struct that implementsrand::RngCore
. That way pyo3 functions could be passed anp.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:
downcast
anp.random.BitGenerator
instance into anumpy::random::PyBitGenerator
..lock()
on it to get anumpy::random::PyBitGeneratorGuard
.TODO:
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
andPyBitGeneratorGuard
, allowing to choose if someone wants toPyBitGenerator
’srandom_*
methods directly on that object while holding the GIL and without locking itnp.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.