Dispatch registrations table is published whole under a lock (#76 B6)

tests/api/test_threadstate.py

@@ -130,6 +130,57 @@ def body():
         f"thread A's op was dropped after another thread ended its capture: {out}"
 
 
+_REGISTRATIONS_CHILD = r"""
+import json, sys, threading
+import torch
+assert hasattr(torch._C, "_aten_implemented"), "the child imported upstream, not the shim"
+sys.setswitchinterval(1e-6)
+key = "CompositeImplicitAutograd"
+n = 8
+barrier = threading.Barrier(n)
+seen = [None] * n
+
+def body(i):
+    barrier.wait()
+    seen[i] = torch._C._dispatch_get_registrations_for_dispatch_key(key)
+
+threads = [threading.Thread(target=body, args=(i,)) for i in range(n)]
+for t in threads:
+    t.start()
+for t in threads:
+    t.join(60)
+final = torch._C._dispatch_get_registrations_for_dispatch_key(key)
+print(json.dumps({"seen": [len(s) if s is not None else None for s in seen],
+                  "same": [s == final for s in seen],
+                  "final": len(final), "unique": len(set(final))}))
+"""
+
+
+def test_dispatch_registrations_built_by_eight_threads_at_once_are_whole():
+    # #76 B6: the table was filled lazily into the shared dict, and the
+    # `if _DISPATCH_REGISTRATIONS: return` check is true as soon as the keys
+    # are inserted, before any name is. A second thread got back a partial
+    # table, and two threads filling at once appended every name twice. Each
+    # run is a fresh subprocess because the table is built once per process.
+    import json
+    import subprocess
+    if not hasattr(torch._C, "_aten_implemented"):
+        raise _skip.Skip("upstream torch: no file-built registration table")
+    env = dict(os.environ)
+    env["PYTHONPATH"] = os.pathsep.join(p for p in sys.path[:2])
+    for attempt in range(3):
+        proc = subprocess.run([sys.executable, "-c", _REGISTRATIONS_CHILD],
+                              capture_output=True, text=True, timeout=600, env=env)
+        lines = [l for l in proc.stdout.splitlines() if l.startswith("{")]
+        assert proc.returncode == 0 and lines, proc.stderr[-2000:]
+        out = json.loads(lines[-1])
+        assert out["final"] > 0, out
+        assert out["unique"] == out["final"], (
+            f"run {attempt}: the table holds duplicate names, two threads filled it at once: {out}")
+        assert all(out["same"]), (
+            f"run {attempt}: a thread got back a different table than the finished one: {out}")
+
+
 def test_backward_on_another_thread_than_the_forward():
     # #78: upstream's graph hangs off the tensors, so a forward recorded in one
     # thread can be differentiated from another. The shim's eager tape was per

torchnative/rust/torch_c/src/bootstrap.py

@@ -1711,7 +1711,9 @@ def _aten_tags(module, qualname: str, overload: str) -> list:
 
 
 #: `<dispatch key>` -> `["aten::<name>[.<overload>]", ...]`, built on first use.
-_DISPATCH_REGISTRATIONS: dict = {}
+#: `None` until a whole table is published; see `_scan_dispatch_registrations`.
+_DISPATCH_REGISTRATIONS: dict | None = None
+_DISPATCH_REGISTRATIONS_LOCK = threading.Lock()
 
 #: The keys `native_functions.yaml` is authoritative about. These four are
 #: *alias* keys: they say how an op is put together, not which backend runs it,
@@ -2011,18 +2013,35 @@ def _scan_dispatch_registrations() -> dict:
     builtin -- the same 176-name gap `_is_absent_inplace_variant` documents,
     seen from another side.
     """
-    if _DISPATCH_REGISTRATIONS:
+    global _DISPATCH_REGISTRATIONS
+    table = _DISPATCH_REGISTRATIONS
+    if table is not None:
+        return table
+    with _DISPATCH_REGISTRATIONS_LOCK:
+        if _DISPATCH_REGISTRATIONS is None:
+            _DISPATCH_REGISTRATIONS = _build_dispatch_registrations()
         return _DISPATCH_REGISTRATIONS
-    for key in _FILE_DECLARED_DISPATCH_KEYS:
-        _DISPATCH_REGISTRATIONS[key] = []
+
+
+def _build_dispatch_registrations() -> dict:
+    """The scan `_scan_dispatch_registrations` publishes, into a fresh dict.
+
+    #76 B6: this used to fill the shared dict in place, and the "already
+    built" check was true as soon as the keys were inserted, before any name
+    was. Eight threads asking at once got back seven empty lists, and two
+    threads filling at once appended every name twice. Built here, then
+    published whole under `_DISPATCH_REGISTRATIONS_LOCK`, a reader sees either
+    no table or the finished one.
+    """
+    table = {key: [] for key in _FILE_DECLARED_DISPATCH_KEYS}
     source = _native_functions_source()
     if not os.path.isabs(source or ""):
-        return _DISPATCH_REGISTRATIONS
+        return table
     try:
         with open(source, encoding="utf-8") as handle:
             lines = handle.read().splitlines()
     except OSError:
-        return _DISPATCH_REGISTRATIONS
+        return table
 
     state = {"name": None, "keys": set(), "dispatch": False,
              "structured": False, "delegate": False, "in_block": False}
@@ -2035,8 +2054,8 @@ def flush():
         if not state["dispatch"] and not state["structured"] and not state["delegate"]:
             keys.add("CompositeImplicitAutograd")
         for key in keys:
-            if key in _DISPATCH_REGISTRATIONS:
-                _DISPATCH_REGISTRATIONS[key].append(f"aten::{name}")
+            if key in table:
+                table[key].append(f"aten::{name}")
 
     for line in lines:
         if line.startswith("- func:"):
@@ -2072,7 +2091,7 @@ def flush():
             for key in stripped.split(":", 1)[0].split(","):
                 state["keys"].add(key.strip())
     flush()
-    return _DISPATCH_REGISTRATIONS
+    return table
 
 
 def _dispatch_registrations(key: str) -> list: