[stf] Enable free-threaded Python (3.14t) support for the cuda.stf bindings

[caugonnet]

Summary

Enable the cuda.stf Python bindings on free-threaded CPython (3.13t/3.14t). The bindings were correct under the GIL; free-threading is a new requirement that removes the implicit serialization the wrapper objects relied on for lifetime management. This PR adds the synchronization the wrappers need to own, declares the extension freethreading_compatible (so importing cuda.stf no longer re-enables the GIL), documents the threading contract, and adds multi-threaded contract tests.

Motivation: concurrent task submission on a shared context is supported and internally synchronized by the C++ runtime (contract documented and tested in NVIDIA#10892), and free-threaded submitters measured beneficial for host-side-bound workloads. Today, importing the extension on 3.14t re-enables the GIL (RuntimeWarning: The global interpreter lock (GIL) has been enabled to load module ...), which forfeits that benefit.

What is in here (one commit per concern)

1. Per-context lifetime lock (_AliveFlag + PyThread_type_lock): serializes child __dealloc__ check-then-destroy against finalize() flip-and-free, makes racing finalize() calls resolve to exactly one teardown, and guards the stackable open-scope counter. PyThread_type_lock is used because it survives nogil regions (a critical section would not) and exists on every supported CPython. Contended acquisition releases the GIL while blocking, and the locked regions are kept free of Python allocation so the non-reentrant lock cannot be re-entered by a GC-triggered __dealloc__ on the owning thread.
2. _PrimaryContextPin.release: atomic test-and-set so the driver's primary-context refcount is decremented exactly once (an over-decrement could tear down a primary context other libraries still use).
3. LaunchableGraph: atomic claim of the C-level shared handle so duplicate reset()/__dealloc__ cannot double-free it, and the owner scope closes exactly once. The release itself remains thread-confined (the stackable scope structure is per-thread in the runtime) and this is now documented.
4. exec_place affinity scopes: per-thread scope stacks (threading.local) instead of a single slot on the shared place object; also makes single-thread nesting of with place: work.
5. check_errors(): EAFP pop so each pending host-callback error is surfaced to exactly one caller.
6. TaskGraph: a lock around the record-once/launch-many state machine; launch() checks state under the lock but launches outside it.
7. Clean errors on finalized contexts: creating tasks/logical_data/tokens on a finalized context now raises RuntimeError instead of passing a null handle to the C layer (the C-side null checks are _CCCL_ASSERTs, compiled out of release builds).
8. freethreading_compatible=True on the extension module -- deliberately ordered after the fixes.
9. Documented threading contract (README "Thread safety" section + docstrings): safe from multiple threads -- task submission on a shared context, wrapper teardown from any thread, idempotent finalize, per-thread affinity scopes; requires external quiescence -- finalize()/fence()/wait() as single-caller phase operations, thread-confined stackable scope recording, place configuration before submission. This mirrors the C++ contract.
10. Tests (tests/stf/test_free_threading.py): contract assertions, not luck -- N submitter threads on a shared context, concurrent logical_data create/destroy, wrapper references dropped concurrently with finalize(), finalize-vs-finalize, submission-after-finalize raising, LaunchableGraph reset idempotency + concurrent no-op resets, per-thread scopes, multi-threaded pytorch_task, and (on free-threaded builds) sys._is_gil_enabled() staying False after import. Thread/iteration counts scale via STF_TEST_NUM_THREADS / STF_TEST_ITERS.

Design note: no locks were added to the per-task submission hot path. The runtime already synchronizes submission on the C++ side; the Python-side fixes are strictly about wrapper-object lifetime, which lives on cold paths (creation guards, teardown, finalize, scope open/close).

Testing (GB300, CUDA 13.4, aarch64)

Build (editable, no build isolation), on both interpreters:

uv venv /tmp/stf-312 --python 3.12 && uv venv /tmp/stf-314t --python 3.14t  # cpython-3.14.7+freethreaded
uv pip install --python <venv> scikit-build-core setuptools_scm 'cython>=3.1' 'cmake>=3.30' ninja
cd python/cuda_stf && uv pip install --python <venv> --no-build-isolation -e '.[test-sysctk13]'

• Python 3.12 (GIL) regression: full existing suite pytest tests/ -- 265 passed, 23 skipped (skips are optional-dependency gates), no failures.
• Python 3.14.7 free-threaded: new test file 12/12 passed (including the pytorch interop test with the free-threaded torch 2.13.0+cu130 wheel and cupy-cuda13x installed); sys._is_gil_enabled() is False after import and stays False.
• Stress campaign on 3.14t: STF_TEST_NUM_THREADS=16 STF_TEST_ITERS=100, 5/5 repetitions green.
• Before/after on the unpatched base (same node): importing the extension re-enables the GIL with the documented RuntimeWarning; forcing it off (PYTHON_GIL=0, i.e. what shipping the directive alone would mean) segfaults deterministically (3/3) on the submission-after-finalize path fixed by commit 7. The narrower lifetime races (e.g. __dealloc__-vs-finalize) did not fire in 40 targeted repetitions on the unpatched build -- their windows are a few instructions wide -- so those fixes rest on the interleaving analysis in the commit messages plus the contract tests, which pass 5/5 under stress on the patched build.

Base

Branch stf-python-freethreading-base is a snapshot of upstream main (2f03020); the PR is opened fork-internal for review before submitting upstream.

🤖 Generated with Claude Code