Per-thread grad mode and device-mode stack (#76, part 1)

.github/scripts/gate_suites.py

@@ -83,6 +83,7 @@
 SUITES = {
     "_support/test_shim.py": BOTH,
     "api/test_import.py": BOTH,
+    "api/test_threadstate.py": BOTH,
     "api/test_tnnamespace.py": BOTH,
     "api/test_tntransformers.py": BOTH,
     "bindings/test_aliasinc.py": BOTH,

tests/api/test_threadstate.py

@@ -0,0 +1,138 @@
+"""State that upstream keeps per thread must not leak between threads (#76).
+
+Upstream keeps grad mode (`c10/core/GradMode.h`) and the `torch.device(...)`
+mode stack per thread. The shim held both process-wide, so `torch.no_grad()`
+in one thread silently stopped another thread's `grad_fn` recording. That is
+wrong with the GIL on, not only on a free-threaded interpreter: the GIL
+serialises bytecode, it does not make a global per thread.
+
+Each test parks one thread inside a context manager (a no-grad region, a device
+mode) and asserts from another thread that the state it sees is its own. The
+same tests pass against upstream torch, which is how they were checked before
+the fix.
+
+What this cannot see: a race that leaves the right answer. These tests check
+isolation, not memory safety; the RNG stream free race of #76 needs the
+free-threaded hammer in tests/api/test_freethreaded.py.
+"""
+
+import os
+import sys
+import threading
+
+# The shim imports upstream extension libraries that resolve only with global
+# symbols, as every other suite that imports torch sets.
+os.environ.setdefault("TORCH_USE_RTLD_GLOBAL", "1")
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "_support"))
+sys.path.insert(0, os.path.join(str(next(p for p in __import__("pathlib").Path(__file__).resolve().parents if (p / "AGENTS.md").is_file())), "torchnative", "python"))
+
+import torch  # noqa: E402
+import _skip  # noqa: E402
+
+
+def _park(enter, seen):
+    """Run `enter()` in a thread, record what it sees inside, hold it open."""
+    inside, release = threading.Event(), threading.Event()
+
+    def body():
+        with enter():
+            seen["inside"] = torch.is_grad_enabled()
+            inside.set()
+            release.wait(10)
+        seen["after"] = torch.is_grad_enabled()
+
+    t = threading.Thread(target=body)
+    t.start()
+    assert inside.wait(10), "parked thread never entered"
+    return t, release
+
+
+def test_no_grad_in_another_thread_does_not_stop_recording_here():
+    seen = {}
+    t, release = _park(torch.no_grad, seen)
+    try:
+        x = torch.ones(2, requires_grad=True)
+        y = x * 2
+        here = torch.is_grad_enabled()
+    finally:
+        release.set()
+        t.join(10)
+    assert seen["inside"] is False, seen
+    assert here is True, "grad mode leaked from another thread's no_grad"
+    assert y.grad_fn is not None, "grad_fn not recorded: another thread's no_grad applied here"
+
+
+def test_overlapping_no_grad_regions_restore_each_threads_own_mode():
+    # Thread A enters no_grad, then the main thread enters and leaves its own.
+    # With one global flag, the main thread's exit restores A's False (or A's
+    # exit restores the main thread's), and someone is left with the wrong mode.
+    seen = {}
+    t, release = _park(torch.no_grad, seen)
+    try:
+        with torch.no_grad():
+            assert torch.is_grad_enabled() is False
+        after_main = torch.is_grad_enabled()
+    finally:
+        release.set()
+        t.join(10)
+    assert after_main is True, "main thread left with another thread's mode"
+    assert seen["after"] is True, seen
+
+
+def test_a_device_mode_in_another_thread_does_not_apply_here():
+    inside, release = threading.Event(), threading.Event()
+
+    def body():
+        with torch.device("meta"):
+            inside.set()
+            release.wait(10)
+
+    t = threading.Thread(target=body)
+    t.start()
+    try:
+        assert inside.wait(10)
+        dev = torch.empty(2).device
+    finally:
+        release.set()
+        t.join(10)
+    assert dev.type == "cpu", f"another thread's torch.device('meta') applied here: {dev}"
+
+
+def test_ending_a_capture_on_one_thread_does_not_stop_another_threads():
+    # #76: the "is anyone recording" flag was one bool, cleared by whichever
+    # thread ended its capture first, while the recorders are per thread. The
+    # other thread's later ops then skipped recording and vanished from its
+    # trace. Here the main thread captures and ends while thread A is mid-capture.
+    C = torch._C
+    if not hasattr(C, "_capture_begin"):
+        raise _skip.Skip("no capture API (upstream torch)")
+    began, release, out = threading.Event(), threading.Event(), {}
+
+    def body():
+        x = torch.ones(3)
+        C._capture_begin([x])
+        began.set()
+        release.wait(10)
+        y = x.neg()
+        out["ops"] = [n["op"] for n in C._capture_end(y).nodes]
+
+    t = threading.Thread(target=body)
+    t.start()
+    try:
+        assert began.wait(10)
+        a = torch.ones(2)
+        C._capture_begin([a])
+        C._capture_end(a.exp())
+    finally:
+        release.set()
+        t.join(10)
+    assert any("neg" in op for op in out.get("ops", [])), \
+        f"thread A's op was dropped after another thread ended its capture: {out}"
+
+
+if __name__ == "__main__":
+    assert hasattr(torch._C, "_aten_implemented") or os.environ.get("TORCHNATIVE_ORACLE") == "1", \
+        "imported upstream torch, not the shim"
+    raise SystemExit(1 if _skip.run_tests(
+        [(n, f) for n, f in list(globals().items()) if n.startswith("test_") and callable(f)],
+        "api/test_threadstate") else 0)

torchnative/rust/torch_c/src/bootstrap.py

@@ -4969,13 +4969,40 @@ def _get_tensor_metadata(tensor):
 # `DeviceContext`; `torch/__init__.py` has `set_default_device` and
 # `get_default_device`. All five `_C` names they bottom out in are below.
 #
-# **The stack is process-wide where upstream's is thread-local.**
-# `PythonTorchFunctionTLS` is per-thread; this is one list. Recorded rather than
-# fixed: a mode entered on one thread would be seen by another, which upstream
-# would not do. Nothing in this shim's measured paths is multi-threaded, and the
-# fix (a `threading.local`) is a two-line change in this block if that stops
-# being true.
-_MODE_STACK: list = []
+# **The stack is per thread**, as upstream's `PythonTorchFunctionTLS` is. It was
+# one process-wide list, so `with torch.device("meta")` on one thread made
+# another thread's `torch.empty` land on meta (#76). The proxy keeps the
+# list-shaped spelling every caller below uses, over a `threading.local`.
+class _ThreadModeStack:
+    __slots__ = ("_tls",)
+
+    def __init__(self):
+        self._tls = threading.local()
+
+    def _list(self) -> list:
+        try:
+            return self._tls.stack
+        except AttributeError:
+            self._tls.stack = []
+            return self._tls.stack
+
+    def append(self, mode):
+        self._list().append(mode)
+
+    def pop(self):
+        return self._list().pop()
+
+    def __len__(self):
+        return len(self._list())
+
+    def __bool__(self):
+        return bool(self._list())
+
+    def __getitem__(self, index):
+        return self._list()[index]
+
+
+_MODE_STACK = _ThreadModeStack()
 
 
 def _through_torch_function_modes(func, args, kwargs):
@@ -7181,11 +7208,10 @@ def _install_grad_mode(module, varfns) -> None:
     `torch.is_grad_enabled()` and `torch._C.is_grad_enabled()` disagree.
     """
     state = {
-        "grad": True,
         # `torch/autograd/grad_mode.py:340` and `:393`. Both are real
         # backend-configuration switches (thread pool, layout enforcement) with
         # nothing behind them here, and both are context managers that restore
-        # what they read -- so, like `grad`, they have to round-trip.
+        # what they read -- so, like grad mode, they have to round-trip.
         "multithreading": True,
         "layout_enforcement": False,
         # `torch.is_inference_mode_enabled()`. Upstream spells this ONLY at
@@ -7194,7 +7220,7 @@ def _install_grad_mode(module, varfns) -> None:
         # as the table-less refusal made `fake_tensor.py:1801` stop the whole of
         # `torch.export` on a predicate that has an obvious answer.
         #
-        # It lives HERE, in the same dict as `grad`, rather than as a constant,
+        # It lives HERE, in this dict rather than as a constant,
         # because docs/graph/EXPORT.md §2.2 is about exactly the other choice: a
         # constant `False` would make `with torch.inference_mode():` a block
         # that enters, reports itself absent, and changes nothing. The setter
@@ -7203,8 +7229,12 @@ def _install_grad_mode(module, varfns) -> None:
         "inference": False,
     }
 
+    # Grad mode is per thread, as upstream's `GradMode` is (#76); the other
+    # entries in `state` stay process-wide, as their upstream switches are.
+    grad_tls = threading.local()
+
     def is_grad_enabled():
-        return state["grad"]
+        return getattr(grad_tls, "grad", True)
 
     def _set_grad_enabled(mode):
         # Two writes, one source of truth. `state` stays the value
@@ -7214,7 +7244,7 @@ def _set_grad_enabled(mode):
         # differentiate. See `tensor.rs`'s `GRAD_ENABLED`, and
         # `test_no_grad_gates_grad_fn` for the two being checked against each
         # other rather than assumed equal.
-        state["grad"] = bool(mode)
+        grad_tls.grad = bool(mode)
         module._shim_set_grad_enabled_flag(bool(mode))
 
     def is_inference_mode_enabled():

torchnative/rust/torch_c/src/capture.rs

@@ -61,7 +61,13 @@ use crate::tensor::PyTensorBase;
 /// Global rather than thread-local because reading a `thread_local!` costs a
 /// TLS lookup, and the recorder itself is thread-local anyway -- a second
 /// thread that takes the branch finds no recorder and falls straight back out.
-static CAPTURING: AtomicBool = AtomicBool::new(false);
+///
+/// A **count** of open regions, not a bool. As a bool it was cleared by
+/// whichever thread ended its capture first, and every later op of a thread
+/// still mid-capture skipped recording and vanished from its trace (#76).
+/// `capture_begin` refuses a second region on the same thread, so each thread
+/// contributes at most one, and only a thread that took its recorder decrements.
+static CAPTURING: AtomicUsize = AtomicUsize::new(0);
 
 thread_local! {
     static RECORDER: RefCell<Option<Recorder>> = const { RefCell::new(None) };
@@ -76,7 +82,7 @@ thread_local! {
 
 #[inline(always)]
 pub fn is_active() -> bool {
-    CAPTURING.load(Ordering::Relaxed)
+    CAPTURING.load(Ordering::Relaxed) != 0
 }
 
 /// Whether the eager tape records at all.
@@ -1738,7 +1744,7 @@ pub fn capture_begin(py: Python<'_>, inputs: &Bound<'_, PyAny>) -> PyResult<()>
     let _ = py;
 
     RECORDER.with(|cell| *cell.borrow_mut() = Some(rec));
-    CAPTURING.store(true, Ordering::Relaxed);
+    CAPTURING.fetch_add(1, Ordering::Relaxed);
     Ok(())
 }
 
@@ -1753,7 +1759,9 @@ pub fn capture_abandon() -> PyResult<()> {
 
 fn take_recorder() -> PyResult<Recorder> {
     let taken = RECORDER.with(|cell| cell.borrow_mut().take());
-    CAPTURING.store(false, Ordering::Relaxed);
+    if taken.is_some() {
+        CAPTURING.fetch_sub(1, Ordering::Relaxed);
+    }
     taken.ok_or_else(not_recording)
 }