[GR-55037] Fixes for tokenizers

PullRequest: graalpython/3390

Author: msimacek

graalpython/com.oracle.graal.python.cext/src/capi.c

@@ -853,7 +853,7 @@ void initialize_hashes();
 void _PyFloat_InitState(PyInterpreterState* state);
 
 Py_LOCAL_SYMBOL TruffleContext* TRUFFLE_CONTEXT;
-Py_LOCAL_SYMBOL int graalpy_finalizing;
+Py_LOCAL_SYMBOL int32_t graalpy_finalizing;
 
 PyAPI_FUNC(void) initialize_graal_capi(TruffleEnv* env, void **builtin_closures) {
     clock_t t = clock();
@@ -880,58 +880,12 @@ PyAPI_FUNC(void) initialize_graal_capi(TruffleEnv* env, void **builtin_closures)
     PyTruffle_Log(PY_TRUFFLE_LOG_FINE, "initialize_graal_capi: %fs", ((double) (clock() - t)) / CLOCKS_PER_SEC);
 }
 
-/*
-This is a workaround for C++ modules, namely PyTorch, that declare global/static variables with destructors that call
-_Py_DECREF. The destructors get called by libc during exit during which we cannot make upcalls as that would segfault.
-So we rebind them to no-ops when exiting.
-*/
-Py_ssize_t nop_GraalPy_get_PyObject_ob_refcnt(PyObject* obj) {
-    return IMMORTAL_REFCNT; // large dummy refcount
-}
-
-void nop_GraalPy_set_PyObject_ob_refcnt(PyObject* obj, Py_ssize_t refcnt) {
-    // do nothing
-}
-
-void nop_GraalPyTruffle_NotifyRefCount(PyObject* obj, Py_ssize_t refcnt) {
-    // do nothing
-}
-
-void nop_GraalPyTruffle_BulkNotifyRefCount(void* ptrs, int count) {
-    // do nothing
-}
-
-/*
- * This array contains pairs of variable address and "reset value".
- * The variable location is usually the address of a function pointer variable
- * and the reset value is a new value to set at VM shutdown.
- * For further explanation why this is required, see Java method
- * 'CApiContext.ensureCapiWasLoaded'.
- *
- * Array format: [ var_addr, reset_val, var_addr1, reset_val1, ..., NULL ]
- *
- * ATTENTION: If the structure of the array's content is changed, method
- *            'CApiContext.addNativeFinalizer' *MUST BE* adopted.
- *
- * ATTENTION: the array is expected to be NULL-terminated !
- *
- */
-static int64_t reset_ptrs[] = {
-        &graalpy_finalizing, 1,
-        &GraalPy_get_PyObject_ob_refcnt, nop_GraalPy_get_PyObject_ob_refcnt,
-        &GraalPy_set_PyObject_ob_refcnt, nop_GraalPy_set_PyObject_ob_refcnt,
-        &GraalPyTruffle_NotifyRefCount, nop_GraalPyTruffle_NotifyRefCount,
-        &GraalPyTruffle_BulkNotifyRefCount, nop_GraalPyTruffle_NotifyRefCount,
-        /* sentinel (required) */
-        NULL
-};
-
 /*
  * This function is called from Java during C API initialization to get the
- * pointer to array 'reset_pts'.
+ * pointer `graalpy_finalizing`.
  */
-PyAPI_FUNC(int64_t *) GraalPy_get_finalize_capi_pointer_array() {
-    return reset_ptrs;
+PyAPI_FUNC(int32_t *) GraalPy_get_finalize_capi_pointer() {
+    return &graalpy_finalizing;
 }
 
 #if ((__linux__ && __GNU_LIBRARY__) || __APPLE__)

graalpython/com.oracle.graal.python.cext/src/object.c

@@ -169,9 +169,15 @@ void
 _Py_DecRef(PyObject *o)
 {
     // GraalPy change: different implementation
+    if (UNLIKELY(graalpy_finalizing)) {
+        /*
+         * Many native libraries in C++ do decrefs in global object destructors that run
+         * after the VM has exitted and we have deallocated the native wrappers.
+         */
+        return;
+    }
     const Py_ssize_t refcnt = Py_REFCNT(o);
-    if (refcnt != IMMORTAL_REFCNT)
-    {
+    if (refcnt != IMMORTAL_REFCNT) {
         const Py_ssize_t updated_refcnt = refcnt - 1;
         Py_SET_REFCNT(o, updated_refcnt);
         if (updated_refcnt != 0) {

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/cext/capi/CApiContext.java

@@ -834,11 +834,11 @@ public static CApiContext ensureCapiWasLoaded(Node node, PythonContext context,
                      * is terminated by a signal, the context exit is skipped. For that case we set
                      * up the shutdown hook.
                      */
-                    Object finalizeFunction = U.readMember(capiLibrary, "GraalPy_get_finalize_capi_pointer_array");
+                    Object finalizeFunction = U.readMember(capiLibrary, "GraalPy_get_finalize_capi_pointer");
                     Object finalizeSignature = env.parseInternal(Source.newBuilder(J_NFI_LANGUAGE, "():POINTER", "exec").build()).call();
-                    Object resetFunctionPointerArray = SignatureLibrary.getUncached().call(finalizeSignature, finalizeFunction);
+                    Object finalizingPointer = SignatureLibrary.getUncached().call(finalizeSignature, finalizeFunction);
                     try {
-                        cApiContext.addNativeFinalizer(env, resetFunctionPointerArray);
+                        cApiContext.addNativeFinalizer(env, finalizingPointer);
                     } catch (RuntimeException e) {
                         // This can happen when other languages restrict multithreading
                         LOGGER.warning(() -> "didn't register a native finalizer due to: " + e.getMessage());
@@ -865,33 +865,18 @@ public static CApiContext ensureCapiWasLoaded(Node node, PythonContext context,
     }
 
     /**
-     * Registers a VM shutdown hook, that resets certain C API function to NOP functions to avoid
-     * segfaults due to improper Python context shutdown. In particular, the shutdown hook reads
-     * pairs of {@code variable address} and {@code reset value} provided in a native array and
-     * writes the {@code reset value} to the {@code variable address}. Currently, this is used to
-     * change the incref and decref upcall functions which would crash if they are called after the
-     * Python context was closed. The format of the array is:
-     *
-     * <pre>
-     *     [ var_addr, reset_val, var_addr1, reset_val1, ..., NULL ]
-     * </pre>
+     * Registers a VM shutdown hook, that sets {@code graalpy_finalizing} variable to let the C side
+     * know that it's not safe to do upcalls and that native wrappers might have been deallocated.
+     * We need to do it in a VM shutdown hook to make sure C atexit won't crash even if our context
+     * finalization didn't run.
      */
-    private void addNativeFinalizer(Env env, Object resetFunctionPointerArray) {
+    private void addNativeFinalizer(Env env, Object finalizingPointerObj) {
         final Unsafe unsafe = getContext().getUnsafe();
-        InteropLibrary lib = InteropLibrary.getUncached(resetFunctionPointerArray);
-        if (!lib.isNull(resetFunctionPointerArray) && lib.isPointer(resetFunctionPointerArray)) {
+        InteropLibrary lib = InteropLibrary.getUncached(finalizingPointerObj);
+        if (!lib.isNull(finalizingPointerObj) && lib.isPointer(finalizingPointerObj)) {
             try {
-                long lresetFunctionPointerArray = lib.asPointer(resetFunctionPointerArray);
-                nativeFinalizerRunnable = () -> {
-                    long resetFunctionPointerLocation;
-                    long curElemAddr = lresetFunctionPointerArray;
-                    while ((resetFunctionPointerLocation = unsafe.getLong(curElemAddr)) != 0) {
-                        curElemAddr += Long.BYTES;
-                        long replacingFunctionPointer = unsafe.getLong(curElemAddr);
-                        unsafe.putAddress(resetFunctionPointerLocation, replacingFunctionPointer);
-                        curElemAddr += Long.BYTES;
-                    }
-                };
+                long finalizingPointer = lib.asPointer(finalizingPointerObj);
+                nativeFinalizerRunnable = () -> unsafe.putInt(finalizingPointer, 1);
                 nativeFinalizerShutdownHook = env.newTruffleThreadBuilder(nativeFinalizerRunnable).build();
                 Runtime.getRuntime().addShutdownHook(nativeFinalizerShutdownHook);
             } catch (UnsupportedMessageException e) {

graalpython/lib-graalpython/patches/tokenizers/tokenizers-0.13.3.patch

@@ -22,6 +22,34 @@ index 6282c31..47e6b12 100644
 
  [features]
  default = ["pyo3/extension-module"]
+diff --git a/src/lib.rs b/src/lib.rs
+index bf1adbd..89c34fe 100644
+--- a/src/lib.rs
++++ b/src/lib.rs
+@@ -48,14 +48,16 @@ extern "C" fn child_after_fork() {
+ pub fn tokenizers(_py: Python, m: &PyModule) -> PyResult<()> {
+     let _ = env_logger::try_init_from_env("TOKENIZERS_LOG");
+ 
++    // GraalPy change: Disable the atfork warning. This triggers a ton of false positives when
++    // jline calls stty and we don't support fork anyway
+     // Register the fork callback
+-    #[cfg(target_family = "unix")]
+-    unsafe {
+-        if !REGISTERED_FORK_CALLBACK {
+-            libc::pthread_atfork(None, None, Some(child_after_fork));
+-            REGISTERED_FORK_CALLBACK = true;
+-        }
+-    }
++    // #[cfg(target_family = "unix")]
++    // unsafe {
++    //     if !REGISTERED_FORK_CALLBACK {
++    //         libc::pthread_atfork(None, None, Some(child_after_fork));
++    //         REGISTERED_FORK_CALLBACK = true;
++    //     }
++    // }
+ 
+     m.add_class::<tokenizer::PyTokenizer>()?;
+     m.add_class::<tokenizer::PyAddedToken>()?;
 diff --git a/tokenizers-lib/src/models/bpe/trainer.rs b/tokenizers-lib/src/models/bpe/trainer.rs
 index 43ab848..55f95f8 100644
 --- a/tokenizers-lib/src/models/bpe/trainer.rs

graalpython/lib-graalpython/patches/tokenizers/tokenizers-0.15.0.patch

@@ -22,3 +22,31 @@ index 51ac712..a7d1e9b 100644
 
  [features]
  defaut = ["pyo3/extension-module"]
+diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs
+index 8625944..a8ac19d 100644
+--- a/bindings/python/src/lib.rs
++++ b/bindings/python/src/lib.rs
+@@ -50,14 +50,16 @@ extern "C" fn child_after_fork() {
+ pub fn tokenizers(_py: Python, m: &PyModule) -> PyResult<()> {
+     let _ = env_logger::try_init_from_env("TOKENIZERS_LOG");
+ 
++    // GraalPy change: Disable the atfork warning. This triggers a ton of false positives when
++    // jline calls stty and we don't support fork anyway
+     // Register the fork callback
+-    #[cfg(target_family = "unix")]
+-    unsafe {
+-        if !REGISTERED_FORK_CALLBACK {
+-            libc::pthread_atfork(None, None, Some(child_after_fork));
+-            REGISTERED_FORK_CALLBACK = true;
+-        }
+-    }
++    // #[cfg(target_family = "unix")]
++    // unsafe {
++    //     if !REGISTERED_FORK_CALLBACK {
++    //         libc::pthread_atfork(None, None, Some(child_after_fork));
++    //         REGISTERED_FORK_CALLBACK = true;
++    //     }
++    // }
+ 
+     m.add_class::<tokenizer::PyTokenizer>()?;
+     m.add_class::<tokenizer::PyAddedToken>()?;