Test with free-threaded Python (3.14t)

[agriyakhetarpal]

This PR tries to run our test suite against the free-threaded Python interpreter (Python 3.14t). I'm not testing against Python 3.13t anymore as previous iterations on this PR were doing, because many of the builtins and stdlib support are better in 3.14t, and 3.14t moved out of experimental support following PEP 779.

Such testing helps us check two aspects: (i) bugs in NumPy and SciPy related to free-threading that we should report upstream (I had a few with 3.13t, but fortunately none with 3.14t) and (ii) whether we have anything of concern in our implementation in case there is nothing to report upstream.

In light of (ii), I found a few issues, and I think that I have been able to fix them somewhat:

• The first one that I caught in the tests was to make the TraceStack use a ContextVar. This kind of builds from https://py-free-threading.github.io/porting/#converting-global-state-to-thread-local-state. TraceStack.top is a shared counter that every thread increments/decrements. As far as I can tell, this is the only piece of runtime-mutable global state (that I've been able to find). The idea is that concurrent grad/jacobian/forward-mode calls must not induce a race on a global top counter. This was previously using threading.local, but I switched to ContextVars based on the review below in Test with free-threaded Python (3.14t) #659 (review) (see also Recommend contextvars.ContextVar as a thread-safe and async-safe mechanism for global state Quansight-Labs/free-threaded-compatibility#337).
• There were a few cases of a module-level shared random state, which I think should be categorised as more of a thread safety issue in the test suite rather than anything functionality-related, as our test suite infrastructure is showing a bit of age due to the lack of modern patterns such as parameterisations (xref Do away with combo_check and switch to test parametrization #662). These tests now use a local RNG defined within the test. Note that there are still a few tests that are ill-conditioned/non-smooth operations, which doesn't help, but dropping the global npr.randn seems to have fixed quite a lot of issues. I've switched to local RNGs everywhere in the test code, which makes up most of this diff.
• The other change is to make check_grads deterministic by construction, because it draws cotangents from a global RNG. In that sense, I've added rng arguments to VSpace.randn and its ArrayVSpace and ComplexArrayVSpace counterparts. Now every check_grads call creates its own freshly seeded generator, and is more reproducible that way. This seems to help quite a bit with test_odeint and test_pinv, which are highly susceptible and numerically fragile in my testing. The gradient for the pinv is ill-conditioned for near-singular inputs that a racing global RNG would otherwise produce (and cause the test to be flaky).

I also added four explicitly multithreaded tests in test_thread_safety.py.

In this regard, we validate independent concurrent differentiation, i.e., each thread performs its own grad. What we do not support is sharing a single in-flight Box/trace across threads, i.e., differentiating between objects that another thread is mutating. That sounds unusual to me. Based on what I have been able to get from reading many of the docs around this stuff, I think NumPy has the same stance, which is to not mutate a shared ndarray across threads. I added some docs in the tutorial, and added a link to the NumPy thread-safety guide.

The rest of the changes are all plumbing-related. I added a nox session that runs on free-threaded Python, and another that runs on free-threaded Python with the pytest-run-parallel plugin. I tested with --parallel-threads 4 to match what the NumPy CI is doing for now, but I also tried things locally with --iterations 10. I am not setting that in CI, however.

I am more or less following the approach suggested at https://py-free-threading.github.io, which is thankfully written (and written well) by many of my amazing colleagues! <3

[agriyakhetarpal]

Pushed a few updates:

• Used ContextVar for TraceStack.top, so it should now be both thread-safe and async-safe. This turned out to be less invasive than I assumed.
• Added a TSan CI job with the Docker images. I mostly followed the NumPy CI reference. This is green and completes in 7 minutes for the full test suite, which is good because our PyPy jobs take even longer, so I am fine keeping it on PRs as such. Our test suite outside of TSan is fairly fast, so I am not surprised. It turns out that SciPy is not included as per the README, so I have left out those tests for now.
• I used Claude Code to write me a libcst-based script to do a bunch of these per-method RNG calls and it fared pretty well at this. In a few places I had to pitch in manually though. I checked all the files, and they all look correct to me, so no shared RNG state remains anywhere. There are a few cases of npr.seed() in the examples, which I am not too concerned about.

I'll slowly work through the rest of the review comments next!

[agriyakhetarpal]

I've added a lock for the const_graph cache and verified that the multithreaded test fails in its absence. I also spent a few hours investigating @pb01ka's review in depth. At this point, only the user-facing docs stuff remains (which can go in a separate PR because of the large diff already, but happy to put something here itself if you insist). Is there a chance both of you could re-review?

Edit: I discovered that const_graph is actually something we are not using yet as part of the functionality right now and there is a TODO comment about it, rather:

autograd/autograd/differential_operators.py

Lines 144 to 160 in 7c6a626

	# TODO(mattjj): update this function using make_jvp and const_graph
	def make_ggnvp(f, g=lambda x: 1.0 / 2 * np.sum(x**2, axis=-1), f_argnum=0):
	"""Builds a function for evaluating generalized-Gauss-Newton-vector products
	at a point. Slightly more expensive than mixed-mode."""
	@unary_to_nary
	def _make_ggnvp(f, x):
	f_vjp, f_x = _make_vjp(f, x)
	g_hvp, grad_g_x = _make_vjp(grad(g), f_x)
	f_jvp, _ = _make_vjp(f_vjp, vspace(grad_g_x).zeros())
	def ggnvp(v):
	return f_vjp(g_hvp(f_jvp(v)))
	return ggnvp
	return _make_ggnvp(f, f_argnum)

So the fix is not really needed at this time, but I've kept it for later if I ever come back to that part as it makes sense for it to be thread safe anyway.

[pb01ka]

Gave final suggestion in #659 (comment). I leave the decision of incorporating it upto you.

[ngoldbaum]

Just one comment inline. Also are you sure "beta" is right, given the level of work you've put in here? IMO this is more polished than "beta".

[ngoldbaum]

This works and is safe but it could be a real multithreaded scaling bottleneck if the graph is read from often enough.

Instead, for caches that are read from more often than they are written to, I tend to recommend the CoW pattern: https://py-free-threading.github.io/porting/#copy-on-write.

[agriyakhetarpal]

Hmm, in this case, the lock is double-checked. We construct the graph once and store it in the cache, so all subsequent reads from the cache skip the lock entirely.

I did some benchmarking:

First benchmark with cached reads on 1 thread with 300k calls

Benchmark 1: dcl  1thr
  Time (mean ± σ):      3.275 s ±  0.104 s    [User: 3.204 s, System: 0.027 s]
  Range (min … max):    3.151 s …  3.429 s    8 runs
 
Benchmark 2: cow  1thr
  Time (mean ± σ):      3.590 s ±  0.314 s    [User: 3.399 s, System: 0.057 s]
  Range (min … max):    3.293 s …  4.211 s    8 runs
 
Summary
  dcl  1thr ran
    1.10 ± 0.10 times faster than cow  1thr

Second benchmark with 8 threads and 100k calls each

Benchmark 1: dcl  8thr
  Time (mean ± σ):      6.447 s ±  0.142 s    [User: 41.107 s, System: 5.481 s]
  Range (min … max):    6.304 s …  6.752 s    8 runs
 
Benchmark 2: cow  8thr
  Time (mean ± σ):      6.430 s ±  0.096 s    [User: 41.557 s, System: 5.511 s]
  Range (min … max):    6.307 s …  6.589 s    8 runs
 
Summary
  cow  8thr ran
    1.00 ± 0.03 times faster than dcl  8thr

That suggests there isn't a major bottleneck in reads so far. With CoW, I faced a different sort of problem – it redundantly retraces the function quite a lot (5.6x on average) across multiple runs of the benchmark, so the CPU cost is wasted due to the races to create the tracing cache.

import autograd.numpy as np
from autograd.misc import const_graph as dcl
from impls import const_graph_cow as cow

# Count how many times the function body is traced when n threads hit the
# first call concurrently
def make(counter):
    def loss(x):
        counter[0] += 1
        return np.sum(np.sin(x) ** 2 + np.exp(x) + np.dot(x, x))
    return loss

def measure(cg, n_threads=8, trials=40):
    redundant = 0; total_traces = 0
    for _ in range(trials):
        counter = [0]
        f = cg(make(counter))
        x = np.ones(8)
        b = threading.Barrier(n_threads)
        def w():
            b.wait(); f(x)
        ts = [threading.Thread(target=w) for _ in range(n_threads)]
        [t.start() for t in ts]; [t.join() for t in ts]
        total_traces += counter[0]
        if counter[0] > 1: redundant += 1
    return total_traces, trials, redundant

for name, cg in [("DCL", dcl), ("CoW", cow)]:
    total, trials, redundant = measure(cg)
    print(f"{name}: {total} total traces over {trials} trials x 8 threads "
          f"(ideal={trials}). Trials with redundant tracing: {redundant}/{trials}")

DCL: 40 total traces over 40 trials x8 threads (ideal=40). Trials with redundant tracing: 0/40
CoW: 223 total traces over 40 trials x8 threads (ideal=40). Trials with redundant tracing: 40/40

---

(Also, this function isn't currently used directly in the functionality either, so that's a TODO I will be taking a look at after this to implement)

[ngoldbaum]

Benchmarks beat my idle speculation 😀

[agriyakhetarpal]

Just one comment inline. Also are you sure "beta" is right, given the level of work you've put in here? IMO this is more polished than "beta".

Haha, I would be happy to bump it up to stable if you insist. But at this time we should note it's still in untested waters, aside from me playing around with the multithreaded tests and my own experiments.

Thanks a lot for reviewing!