Move array_garbage_collection_guard logic from PyArray_tp_clear to PyArray_tp_finalize.

PiperOrigin-RevId: 859832848

Author: ICGog

docs/jax.rst

@@ -42,6 +42,7 @@ Configuration
 .. autosummary::
    :toctree: _autosummary
 
+   array_garbage_collection_guard
    config
    check_tracer_leaks
    checking_leaks

jax/__init__.py

@@ -68,6 +68,7 @@
   jax2tf_associative_scan_reductions as jax2tf_associative_scan_reductions,
   legacy_prng_key as legacy_prng_key,
   threefry_partitionable as threefry_partitionable,
+  array_garbage_collection_guard as array_garbage_collection_guard,
   transfer_guard as transfer_guard,
   transfer_guard_host_to_device as transfer_guard_host_to_device,
   transfer_guard_device_to_device as transfer_guard_device_to_device,

jax/_src/test_util.py

@@ -227,6 +227,7 @@ def get_output() -> str:
       f.seek(0)
       captured = f.read()
       os.dup2(original_fd, fp.fileno())
+      os.close(original_fd)
 
 
 capture_stdout = partial(_capture_output, sys.stdout)

jaxlib/py_array.cc

@@ -332,8 +332,11 @@ extern "C" int PyArray_tp_traverse(PyObject* self, visitproc visit, void* arg) {
   return 0;
 }
 
-// dynamic_attr: Allow the GC to clear the dictionary.
-extern "C" int PyArray_tp_clear(PyObject* self) {
+extern "C" void PyArray_tp_finalize(PyObject* self) {
+  // This method assumes that `PyObject_CallFinalizerFromDealloc` is not called
+  // from `PyArray_tp_dealloc`. If this assumption is violated, then the garbage
+  // collector guard would trigger for an array deallocated via reference
+  // counting.
   switch (auto guard_level = GetGarbageCollectArrayGuard(); guard_level) {
     case GarbageCollectionGuardLevel::kAllow:
       break;
@@ -361,13 +364,28 @@ extern "C" int PyArray_tp_clear(PyObject* self) {
       if (guard_level == GarbageCollectionGuardLevel::kFatal) {
         Py_FatalError(error_msg.c_str());
       } else {
+#if PY_VERSION_HEX < 0x030C0000
+        PyObject *err_type, *err_value, *err_traceback;
+        PyErr_Fetch(&err_type, &err_value, &err_traceback);
+#else
+        PyObject* exc = PyErr_GetRaisedException();
+#endif
         PyErr_SetString(PyExc_RuntimeError, error_msg.c_str());
         PyErr_Print();
         PyErr_Clear();
+#if PY_VERSION_HEX < 0x030C0000
+        PyErr_Restore(err_type, err_value, err_traceback);
+#else
+        PyErr_SetRaisedException(exc);
+#endif
       }
       break;
     }
   }
+}
+
+// dynamic_attr: Allow the GC to clear the dictionary.
+extern "C" int PyArray_tp_clear(PyObject* self) {
 #if PY_VERSION_HEX < 0x030C0000
   PyObject*& dict = *_PyObject_GetDictPtr(self);
   Py_CLEAR(dict);
@@ -2079,6 +2097,7 @@ PyMemberDef array_impl_members[] = {
 
 PyType_Slot array_impl_slots[] = {
     {Py_tp_new, reinterpret_cast<void*>(PyArray_tp_new)},
+    {Py_tp_finalize, reinterpret_cast<void*>(PyArray_tp_finalize)},
     {Py_tp_dealloc, reinterpret_cast<void*>(PyArray_tp_dealloc)},
     {Py_tp_members, reinterpret_cast<void*>(array_impl_members)},
     {Py_tp_traverse, reinterpret_cast<void*>(PyArray_tp_traverse)},

tests/garbage_collection_guard_test.py

@@ -12,17 +12,32 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import dataclasses
 import gc
 import weakref
 
 from absl.testing import absltest
 import jax
 from jax._src import config
+from jax._src.lib import jaxlib_extension_version
 import jax._src.test_util as jtu
 import jax.numpy as jnp
 
 jax.config.parse_flags_with_absl()
 
+_GC_ERROR_MESSAGE = "`jax.Array` was deleted by the Python garbage collector"
+
+
+@dataclasses.dataclass()
+class _A:
+  a: jax.Array
+  ref: "_B"
+
+
+@dataclasses.dataclass()
+class _B:
+  ref: _A | None = None
+
 
 def _create_array_cycle():
   """Creates a reference cycle of two jax.Arrays."""
@@ -45,11 +60,9 @@ def test_gced_array_is_not_logged_by_default(self):
       self.assertIsNone(ref())  # Cycle collected.
     # Check that no error message is logged because
     # `array_garbage_collection_guard` defaults to `allow`.
-    self.assertNotIn(
-        "`jax.Array` was deleted by the Python garbage collector", stderr(),
-    )
+    self.assertNotIn(_GC_ERROR_MESSAGE, stderr())
 
-  def test_gced_array_is_logged(self):
+  def test_array_part_of_cycle_is_logged(self):
     with config.array_garbage_collection_guard("log"):
       with jtu.capture_stderr() as stderr:
         # Create a reference cycle of two jax.Arrays.
@@ -59,9 +72,27 @@ def test_gced_array_is_logged(self):
         self.assertIsNone(ref())  # Cycle collected.
     # Verify that an error message is logged because two jax.Arrays were garbage
     # collected.
-    self.assertIn(
-        "`jax.Array` was deleted by the Python garbage collector", stderr()
-    )
+    self.assertIn(_GC_ERROR_MESSAGE, stderr())
+
+  def test_array_reachable_from_cycle_is_logged(self):
+    if jaxlib_extension_version < 401:
+      self.skipTest("This functionality is not yet supported in jaxlib.")
+    with config.array_garbage_collection_guard("log"):
+      with jtu.capture_stderr() as stderr:
+        b = _B()
+        a = _A(a=jnp.array([1, 2, 3]), ref=b)
+        b.ref = a
+        del a
+        del b
+        gc.collect()
+    self.assertIn(_GC_ERROR_MESSAGE, stderr())
+
+  def test_no_error_is_logged_when_array_is_not_gced(self):
+    with config.array_garbage_collection_guard("log"):
+      with jtu.capture_stderr() as stderr:
+        a = jnp.array([1, 2, 3])
+        del a
+    self.assertNotIn(_GC_ERROR_MESSAGE, stderr())
 
 
 if __name__ == "__main__":