ENH: spatial.KDTree: make thread-safe via immutability by ngoldbaum · Pull Request #24799 · scipy/scipy

ngoldbaum · 2026-03-12T20:04:01Z

Towards fixing #20669. Supersedes #20655.

Currently, scipy.spatial.KDTree is mostly thread-safe. There are two major issues right now that this PR helps to fix:

KDTree.tree is lazily-initialized, which can lead to data races.
Some attributes of the KDTree are exposed as mutable ndarrays, which can cause corruption if they are modified, with multithreading or no.

I've converted KDTree.tree to use a py_safe single-initialization API, as described here in the Cython docs. These are new in Cython 3.2, so I bumped the minimum Cython version. There's also a compiler warning I can't fix that I reported to Cython.

I also made the ndarray attributes of the KDTree class be read-only. While there are still ways to modify a read-only NumPy array, this at least makes it harder for users to shoot themselves in the foot. I clarified the docstring for the data argument to expand the warnings about modifying the data if no copy happens. I noticed that the boxsize attribute is undocumented, so I added entries. I also updated the thread safety docs.

Finally, I added a multithreaded query_ball_point test and a multithreaded KDTree.tree generation test. I can add multithreaded tests for more of the query functions if you'd like me to do that. I also added a test to validate that KDTrees cannot be modified from Python.

I also updated the multithreaded test helper to fix a possible deadlock and allow tests to get return values from worker threads without relying on global state.

#20655 has a lot of discussion about modifying the KDTree before __init__ or __cinit__ finishes. I'm pretty sure that's all a red herring. It's not possible to modify an object before the initializer finishes, and initializers run in an effectively single-threaded context. I think confusion around that point led to the choice of adding a lock. I don't think that's actually needed.

I ran the KDTree tests, including the new multithreaded tests, under thread sanitizer on my Mac and see no data races.

ngoldbaum · 2026-03-12T20:17:36Z

scipy/spatial/tests/test_kdtree.py

    assert_array_equal(t.maxes, np.amax(points, axis=0))
    assert_array_equal(t.mins, np.amin(points, axis=0))
-    assert t.data is points
+    assert t.data.base is points


This is due to chaging data to be a read-only view. I'm not sure if it counts as a breaking change.

I played with this feature branch a bit locally, and T.data does still provide access to the points of course, so seems like it would be a mild break at worst.

There might even be a bug fix argument I suppose since it prevents a footgun. Maybe we could ping Sturla before merging if we're worried about it.

That shouldn't count as a breaking change, more like a robustness improvement. The current docstring of data input argument already says: "This array is not copied unless this is necessary to produce a contiguous array of doubles, and so modifying this data will result in bogus results."

scipy/spatial/tests/test_kdtree.py

ngoldbaum · 2026-03-12T20:46:06Z

scipy/_lib/_testutils.py

+                # to avoid deadlocks if spawning failed for some reason
+                barrier.abort()
+
+    return [f.result() for f in futures]


The original implementation can deadlock if spawning one of the threads fails. I fixed that and also used a ThreadPoolExecutor, which allows me to return results from the worker functions.

ngoldbaum · 2026-03-12T21:00:09Z

None of the CI failed due to using a too-new Cython feature. Maybe it is OK to bump the minimum Cython version?

rgommers · 2026-03-13T09:47:10Z

These are new in Cython 3.2. I don't know whether it's OK to bump the minimum Cython version. It is currently set to Cython 3.0.8. If bumping the minimum Cython version is not OK, then I think I can vendor a version of py_safe_call_object_once inside of SciPy. There's also a compiler warning I can't fix that I reported to Cython.

It's okay to bump Cython to >=3.2. Ideally a Cython bug fix release with the compiler warning fixed would show up, but we can suppress the -Wunneeded-internal-declaration in this PR - can you add it to the checks here:

scipy/scipy/meson.build

Line 569 in a6e37c0

# C++ warning flags

and then add that to the extension module's cpp_args in spatial/meson.build?

rgommers · 2026-03-13T09:49:31Z

This looks pretty good, thanks for working on this @ngoldbaum! I'll have a closer look once it's out of draft.

Works around a Cython issue that will be addressed in the next bugfix release.

tylerjereddy · 2026-03-19T18:32:32Z

scipy/spatial/tests/test_kdtree.py

    assert_array_equal(t.maxes, np.amax(points, axis=0))
    assert_array_equal(t.mins, np.amin(points, axis=0))
-    assert t.data is points
+    assert t.data.base is points


I played with this feature branch a bit locally, and T.data does still provide access to the points of course, so seems like it would be a mild break at worst.

There might even be a bug fix argument I suppose since it prevents a footgun. Maybe we could ping Sturla before merging if we're worried about it.

tylerjereddy · 2026-03-19T18:37:04Z

scipy/spatial/_ckdtree.pyx

+        # read-only view so ban people modifying tree.data after the tree is
+        # constructed
+        data = data.view(np.ndarray)
+        data.flags.writeable = False


I did wonder when I was reading this if the need for this is milder with copy_data=True? I guess you still have problems, even if a little less serious.

It's still useful to set this read-only flag, as it insulates against mutating this .data attribute. It's milder indeed in that copy_data already insulates against the base array for the input data being mutated. But this seems like the right thing to do either way.

ngoldbaum · 2026-03-19T18:37:44Z

I think this is ready now. I bumped the minimum version to 3.2.0 because I wasn't sure downstream builders would necessarily have 3.2.4. Let me know if you'd prefer it to be 3.2.4.

ngoldbaum · 2026-03-19T19:18:34Z

The last commit fixes a possible race that David Woods pointed out to me: cython/cython#7581 (comment).

We could go back to having a boolean flag but I'd need to use c++ standard library atomics to do that. I don't think the overhead of creating the wrapper object on every call matters much.

ngoldbaum · 2026-03-19T21:25:03Z

but I'd need to use c++ standard library atomics to do that

Eh, the atomics aren't so bad, see the last example in cython/cython#7582. Should I just go ahead and do that, along with switching to a closure instead of a wrapper object?

da-woods · 2026-03-19T22:37:01Z

I think it's possible to avoid the closure and the atomics by using a global cdef function and casting the captured self to/from void. See cython/cython#7582 (comment) where it's written out slightly more fully.

scipy/spatial/_ckdtree.pyx

rgommers · 2026-03-23T06:09:59Z

@ngoldbaum given the Cython bug fix and discussion around avoiding wrapper, can you confirm whether the minimum Cython version this PR needs in its current form is 3.2.0 or 3.2.5?

da-woods · 2026-03-23T07:57:12Z

can you confirm whether the minimum Cython version this PR needs in its current form is 3.2.0 or 3.2.5?

The version here should work with 3.2.0 I believe

ngoldbaum · 2026-03-23T16:00:16Z

can you confirm whether the minimum Cython version this PR needs in its current form is 3.2.0 or 3.2.5?

I just double-checked: this branch builds in an environment with Cython 3.2.0 installed.

…7582) Based on my experience working on scipy/scipy#24799.

rgommers

I went over this again, and all looks good - time to get it in. Thanks @ngoldbaum & reviewers!

rgommers · 2026-03-25T08:54:24Z

scipy/spatial/_ckdtree.pyx

+        # read-only view so ban people modifying tree.data after the tree is
+        # constructed
+        data = data.view(np.ndarray)
+        data.flags.writeable = False


It's still useful to set this read-only flag, as it insulates against mutating this .data attribute. It's milder indeed in that copy_data already insulates against the base array for the input data being mutated. But this seems like the right thing to do either way.

lucascolley · 2026-04-09T14:38:32Z

ASan CI has been a bit shaky since this PR was merged, not sure if it's related at all but either way maybe somebody involved here would be well-placed to investigate #24980

github-actions bot added scipy._lib scipy.spatial Documentation Issues related to the SciPy documentation. Also check https://github.com/scipy/scipy.org Cython Issues with the internal Cython code base labels Mar 12, 2026

ngoldbaum mentioned this pull request Mar 12, 2026

ENH: spatial: ensure thread-safety for KDTree #20655

Closed

ngoldbaum commented Mar 12, 2026

View reviewed changes

scipy/spatial/tests/test_kdtree.py Outdated Show resolved Hide resolved

ngoldbaum commented Mar 12, 2026

View reviewed changes

ngoldbaum force-pushed the thread-safe-kdtree branch from 018a536 to e0ab9df Compare March 19, 2026 15:44

ngoldbaum added 6 commits March 19, 2026 11:37

BUG: Fix thread-unsafe initialization of cKDTree.tree property

55b8363

ENH: make ndarray attributes of cKDTree read-only

6a0bcdb

DOC: update docs and docstrings

ca13741

TST: add test for immutability of kdtree attributes

3bc5682

BLD: Ignore warnings in cKDTree build

ccf920d

Works around a Cython issue that will be addressed in the next bugfix release.

BLD: Bump minimum required Cython version

1e3fe8d

ngoldbaum force-pushed the thread-safe-kdtree branch from e0ab9df to 1e3fe8d Compare March 19, 2026 17:38

ngoldbaum mentioned this pull request Mar 19, 2026

[BUG] data race when calling a closure with py_safe_call_object_once cython/cython#7581

Closed

DOC: fix sphinx warning

fc0244a

tylerjereddy reviewed Mar 19, 2026

View reviewed changes

ngoldbaum marked this pull request as ready for review March 19, 2026 18:37

ngoldbaum requested review from peterbell10 and rgommers as code owners March 19, 2026 18:37

TST: fix buggy assert

ba5c690

ngoldbaum changed the title ~~Make KDTree thread-safe via immutability~~ ENH: Make KDTree thread-safe via immutability Mar 19, 2026

BUG: avoid racey branching and create wrapper unconditionally

3edfcf3

lucascolley added enhancement A new feature or improvement free-threading Items related to supporting free-threaded (a.k.a. "no-GIL") builds of CPython labels Mar 19, 2026

ngoldbaum mentioned this pull request Mar 19, 2026

Add thread-safe cache initialization section to free-threading docs cython/cython#7582

Merged

lucascolley changed the title ~~ENH: Make KDTree thread-safe via immutability~~ ENH: spatial. KDTree thread-safe via immutability Mar 20, 2026

lucascolley changed the title ~~ENH: spatial. KDTree thread-safe via immutability~~ ENH: spatial.KDTree: make thread-safe via immutability Mar 20, 2026

ngoldbaum mentioned this pull request Mar 20, 2026

Fix passing of extra arguments to py_safe_call_once cython/cython#7585

Merged

ngoldbaum commented Mar 20, 2026

View reviewed changes

scipy/spatial/_ckdtree.pyx Outdated Show resolved Hide resolved

MAINT: avoid creating a wrapper object

67b7af0

rgommers mentioned this pull request Mar 23, 2026

Tracker: Support free-threaded CPython builds #20669

Open

51 tasks

da-woods pushed a commit to cython/cython that referenced this pull request Mar 24, 2026

Add thread-safe cache initialization section to free-threading docs (#…

87d8efc

…7582) Based on my experience working on scipy/scipy#24799.

rgommers approved these changes Mar 25, 2026

View reviewed changes

rgommers merged commit 5ca1cdc into scipy:main Mar 25, 2026
57 of 58 checks passed

github-actions bot added the needs-release-note-decision A maintainer should decide if this PR requires a release note. label Mar 25, 2026

lucascolley added needs-release-note A maintainer should add a release note written by a reviewer/author to the wiki. and removed needs-release-note-decision A maintainer should decide if this PR requires a release note. labels Mar 27, 2026

Uh oh!

Conversation

ngoldbaum commented Mar 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ngoldbaum commented Mar 12, 2026

Uh oh!

rgommers commented Mar 13, 2026

Uh oh!

rgommers commented Mar 13, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ngoldbaum commented Mar 19, 2026

Uh oh!

ngoldbaum commented Mar 19, 2026

Uh oh!

ngoldbaum commented Mar 19, 2026

Uh oh!

da-woods commented Mar 19, 2026

Uh oh!

Uh oh!

rgommers commented Mar 23, 2026

Uh oh!

da-woods commented Mar 23, 2026

Uh oh!

ngoldbaum commented Mar 23, 2026

Uh oh!

rgommers left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

lucascolley commented Apr 9, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

ngoldbaum commented Mar 12, 2026 •

edited

Loading