Add support for free-threading

[lysnikolaou]

Implement thread-safety using Cython 3.1+ critical sections to protect Pool operations (alloc, free, realloc) from race conditions. Update dependencies and CI to support free-threaded builds.

[lysnikolaou]

Unfortunately, because the Pool public methods are only exposed in Cython, it's not easy to test them in CI. I used the following to test multi-threaded usage locally:

# cython: freethreading_compatible=True

from collections import deque
import threading

from cymem.cymem cimport Pool


def test_concurrent_pool_operations():
    cdef Pool pool = Pool()
    cdef int num_threads = 10
    cdef int ops_per_thread = 100
    errors = []

    def worker(thread_id):
        cdef Pool p = pool
        cdef void* ptr
        cdef void* old_ptr
        cdef void* new_ptr
        cdef int i
        cdef size_t addr

        try:
            ptrs = deque()
            for _ in range(5):
                ptr = p.alloc(1, 64)
                ptrs.append(<size_t>ptr)

            for i in range(ops_per_thread):
                op = i % 3

                if op == 0:  # alloc
                    ptr = p.alloc(1, 64)
                    ptrs.append(<size_t>ptr)
                elif op == 1 and ptrs:  # free
                    addr = <size_t>ptrs.popleft()
                    ptr = <void*>addr
                    p.free(ptr)
                elif op == 2 and ptrs:  # realloc
                    addr = <size_t>ptrs.popleft()
                    old_ptr = <void*>addr
                    new_ptr = p.realloc(old_ptr, 128)
                    ptrs.append(<size_t>new_ptr)

            # Cleanup remaining allocations
            for addr_obj in ptrs:
                addr = <size_t>addr_obj
                ptr = <void*>addr
                p.free(ptr)
        except Exception as e:
            errors.append((thread_id, e))

    threads = [threading.Thread(target=worker, args=(i,)) for i in range(num_threads)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()

    assert not errors, f"Thread errors: {errors}"
    assert pool.size == 0, f"Pool size should be 0, got {pool.size}"
    assert len(pool.addresses) == 0, f"Pool addresses should be empty, got {len(pool.addresses)}"

And then in a command line:

cythonize -i test_threadsafety.pyx
python -c "from cymem.tests.test_threadsafety import test_concurrent_pool_operations; test_concurrent_pool_operations()"

[ngoldbaum]

If it's part of the public API for these Pool objects to access mutable state, then I think you should add new reader methods that read the state under a critical section and document that third party libraries accessing the C state directly without holding a critical section on the Pool isn't thread-safe.

It'd also be nice to package your test script as a git repo with a meson build configuration. That will also make it easier to run the test under TSan.

[ngoldbaum]

I just went over preshed which depends on cymem. It looks like preshed only uses alloc and free and doesn't access any internal state. So for what that's worth it looks like what you have here is more than sufficient from preshed's perspective.

Let's hold off on adding any new API until we find that we need it.

Once this PR is merged I'll go ahead and make my fork of preshed depend on your fork of cymem. From there we'll need to do a breadth-search reverse search up the spacy dependency tree, setting up CI running against forks that have had a thread safety pass. I don't think we need to go to the trouble of doing pypi releases and actually forking the packages, we'll just muddle through by hacking the CI.

[ngoldbaum]

Here's my WIP work for preshed: ngoldbaum/preshed#1

[lysnikolaou]

Thanks for the reviews, folks! I've addressed all of the feedback.

Once this PR is merged I'll go ahead and make my fork of preshed depend on your fork of cymem.

I'd propose to not merge these PRs on our local forks and just depend on the branch in CI. That makes it easier to upstream the PRs.

It'd also be nice to package your test script as a git repo with a meson build configuration. That will also make it easier to run the test under TSan.

Here you go!

[ngoldbaum]

I'd propose to not merge these PRs on our local forks and just depend on the branch in CI. That makes it easier to upstream the PRs.

What do you think about working in a free-threading or similar branch in our forks instead of the master branch? I'd still like to work in reviewable chunks in PRs and not in one huge PR. I think that works out to the same thing here, where you're doing everything in one chunk so it's all one PR. Either way I'll go ahead and work on making my fork of preshed depend on this branch of cymem and we'll leave this branch as-is until we're ready to make an upstream PR.

[crusaderky]

What do you think about working in a free-threading or similar branch in our forks instead of the master branch? I'd still like to work in reviewable chunks in PRs and not in one huge PR.

Strongly in favour

[ngoldbaum]

Here you go!

Awesome! I think this actually solves most of the "testing Cython is hard" worries the cymem maintainers had about testing Cython code. This would probably be a nice pattern to follow to add more tests. When we eventually speak with the Explosion folks I might propose adding your test as a blueprint for more tests. I probably would have skipped including spin but that's just about the only comment I have. Thanks for doing that work, it's much nicer to work with pip-installable Cython code.

I guess this is an issue with the cymem packaging, but when I try to run your test code, I get a build error:

FAILED: [code=1] test_pool.cpython-314t-darwin.so.p/test_pool.pyx.c
cython -M --fast-fail -3 -Xfreethreading_compatible=true /Users/goldbaum/Documents/test-cymem-threadsafety/test_pool.pyx -o test_pool.cpython-314t-darwin.so.p/test_pool.pyx.c

Error compiling Cython file:
------------------------------------------------------------
...
from collections import deque
import threading

from cymem.cymem cimport Pool
^
------------------------------------------------------------

/Users/goldbaum/Documents/test-cymem-threadsafety/test_pool.pyx:4:0: 'cymem/cymem.pxd' not found

It looks like you did specify your fork of cymem as a build dependency, not sure what's going wrong there.

If I install your branch of cymem and then build your test project with --no-build-isolation, I get past the missing cython header problem if I make the following patch (maybe different meson or meson-python versions?):

diff --git a/meson.build b/meson.build
index 7d7b64b..32061f8 100644
--- a/meson.build
+++ b/meson.build
@@ -4,7 +4,7 @@ project(
   version: '0.1.0',
 )

-py = import('python').find_installation()
+py = import('python').find_installation(pure: false)
 dep_py = py.dependency()

 cy = meson.get_compiler('cython')

That avoids this error:

      meson-python: error: The test_cymem_threadsafety package is split between
purelib and platlib: 'purelib/test_cymem_threadsafety/tests/tests' and 'platlib/test_cymem_threadsafety/test_pool.cpython-314t-darwin.so', a "pure: false" argument may be missing in meson.build. It is recommended to set it in "import('python').find_installation()"
      [end of output]

After doing that I'm able to build and install your test using a TSAN build of Python 3.14-dev and don't see any race reports. Woohoo!

[lysnikolaou]

Folks, I've addressed all feedback here and I think this is now in a good state. If you could give it another review, that'd be great!

[crusaderky]

Just a minor performance tweak remaining, reducing the number of dict accesses from resize() from 4 to 2.

[crusaderky]

Suggested change

	if <size_t>p not in self.addresses:
	self.pyfree.free(new_ptr)
	raise ValueError("Pointer %d not found in Pool %s" % (<size_t>p, self.addresses))
	old_size = self.addresses[<size_t>p]
	old_size = self.addresses.pop(<size_t>p, 0)
	if old_size == 0:
	self.pyfree.free(new_ptr)
	raise ValueError("Pointer %d not found in Pool %s" % (<size_t>p, self.addresses))

[lysnikolaou]

I did a

try:
    self.addresses.pop(...)
except KeyError:
   raise ValueError()

Theoretically, malloc(0) can succeed, which means that 0 can also be a valid value in the dict. I know we could instead default to -1 or something, but the try/except feels clearer here.

[crusaderky]

PyFree could, in theory, release the critical section.
In practice, this is an exception condition that most times in user code corresponds to a bug, so I'm not overly concerned.

[ngoldbaum]

It looks like everything is working over in preshed if I make it depend on this branch of cymem and the master branch of murmurhash:

ngoldbaum/preshed#2

Commenting about that mostly for visibility to you two, I'm going to self-merge it. The thread-safety work over in preshed looks similar to what happened here (although thankfully it looks like state is stored outside the interpreter so it might be simpler). I'll probably ping both of you for review, or just Lysandros if it slips into next week.

[crusaderky]

LGTM

[lysnikolaou]

Closing this since upstream PR is now open: explosion#53