Fix SIGSEGV: release prediction input arrays with the GIL held

[kasper0406]

Summary

MLModel.predict can corrupt the Python heap and segfault, because Core ML releases the prediction's input MLMultiArrays on one of its own dispatch queues, and PybindCompatibleArray's compiler-generated .cxx_destruct then runs Py_DECREF on the wrapped numpy array without the GIL (and without any Python thread state).

PybindCompatibleArray keeps the caller's numpy buffer alive by storing a py::array ivar inside an Objective-C object whose lifetime belongs to Core ML. Core ML does not necessarily release input feature values before -predictionFromFeatures: returns: the MLE5 engine leaves its input ports bound and unbinds them ~1 second later, from -[MLE5ExecutionStream resetAfterLingering:] on com.apple.coreml.MLE5ExecutionStream.resetQueue. When the ObjC wrapper is the sole surviving owner of the numpy object — the common case, since predict() builds its input dict locally and _update_float16_multiarray_input_to_float32 substitutes fresh temporaries — that background unbind drops the last reference off-thread. The result is a race against the interpreter: a crash in _PyObject_Free, either immediately or during some later, unrelated prediction.

Thread 3   com.apple.coreml.MLE5ExecutionStream.resetQueue
  _PyObject_Free                                     libpython3.12.dylib
  -[PybindCompatibleArray .cxx_destruct]             libcoremlpython.so
  object_cxxDestructFromClass                        libobjc.A.dylib
  -[MLFeatureValue dealloc]                          CoreML
  -[MLE5InputPortBinder reset]
  -[MLE5InputPort reset] / -[MLE5ExecutionStreamOperation reset]
  -[MLE5ExecutionStream _reset]
  __43-[MLE5ExecutionStream resetAfterLingering:]_block_invoke
EXC_BAD_ACCESS (SIGSEGV), KERN_INVALID_ADDRESS

In practice this bites applications that alternate between several MLModel instances (e.g. per-function instances of a multifunction model, sharing an MLState): each instance's execution stream reliably idles past the linger timeout while other instances predict. That is how we found it — it initially looked like an MLState bug.

Fix

Give PybindCompatibleArray a -dealloc that drops the Python reference with the GIL held (PyGILState_Ensure/Py_DECREF/PyGILState_Release, with a Py_IsInitialized guard for interpreter shutdown), leaving .cxx_destruct a no-op. The file builds with ARC, so the early returns are safe. This is the only place a Python reference is stored in an ObjC-lifetime object; the __block py::object uses in CoreMLPythonUtils.mm are synchronous handler blocks and are unaffected.

Deadlock consideration: Model::predict currently holds the GIL for the duration of the prediction, so Core ML's reset queue now waits in PyGILState_Ensure until the predict returns. Stress-tested (33 multi-second predictions across 3 lingering instances, ~280 s): no deadlock — the prediction path never waits on the reset queue. Releasing the GIL around -predictionFromFeatures: would remove even that stall and is a natural follow-up, left out here to keep the change minimal.

Reproduction

No MLState, no multifunction model, no large weights: a 1.8 KiB single-add mlprogram, one instance, one predict(), drop the input array, idle one second. Two variants — detect (an ndarray subclass reports its __del__ running 1.06 s after predict on a _DummyThread; deterministic) and crash (plain numpy + allocation churn on the main thread; segfaults unpatched). On this machine (macOS 26.4, Apple Silicon, Python 3.12, coremltools 9.0 wheel and this repo's main): unpatched → detect reports the off-main-thread free and crash dies with SIGSEGV in round 0; patched → both pass, and the free happens on the main thread with the GIL.

repro script (self-contained)

"""Minimal upstream repro: coremltools SIGSEGVs because Core ML releases input
MLMultiArrays on a background queue, and PybindCompatibleArray's Py_DECREF then
runs without the GIL.

Nothing exotic is required -- no MLState, no multifunction model, no large
weights. All that is needed is:

  1. an mlprogram model,
  2. an input numpy array whose only surviving owner after predict() returns is
     the Objective-C MLMultiArray that Core ML kept bound to its input port,
  3. ~1 second of idle time, so the MLE5 execution stream "lingers" and then
     resets itself on com.apple.coreml.MLE5ExecutionStream.resetQueue,
  4. the main thread executing Python (allocating) while that happens.

Step 3 fires -[MLE5InputPortBinder reset] -> -[MLFeatureValue dealloc] ->
-[PybindCompatibleArray .cxx_destruct] -> ~py::array() -> Py_DECREF, from a
thread that holds neither the GIL nor a Python thread state.

Two variants:
  detect  -- ndarray subclass whose __del__ reports the freeing thread.
             Deterministic; proves the cross-thread release exists.
  crash   -- plain numpy arrays plus Python allocation churn. Races pymalloc
             and segfaults on an unpatched coremltools.
"""
import subprocess
import sys
import tempfile
import shutil
from pathlib import Path

SHAPE = (1, 64, 512)

BUILD = r'''
import numpy as np, coremltools as ct
from coremltools.converters.mil import Builder as mb
from coremltools.converters.mil.mil import types

@mb.program(input_specs=[mb.TensorSpec(%r, dtype=types.fp32)],
            opset_version=ct.target.iOS18)
def prog(x):
    return mb.add(x=x, y=np.float32(1.0))

ct.convert(prog, convert_to="mlprogram",
           minimum_deployment_target=ct.target.iOS18,
           skip_model_load=True).save(%r)
print("BUILT")
'''

DETECT = r'''
import gc, threading, time
import numpy as np, coremltools as ct

P, SHAPE = %r, %r
MAIN = threading.main_thread().ident
EVENTS = []
T0 = time.time()

class Tracked(np.ndarray):
    def __del__(self):
        t = threading.current_thread()
        EVENTS.append((time.time() - T0, t.ident != MAIN, type(t).__name__))

m = ct.models.MLModel(P, compute_units=ct.ComputeUnit.CPU_ONLY)
T0 = time.time()
m.predict({"x": np.zeros(SHAPE, dtype=np.float32).view(Tracked)})
gc.collect()
print("immediately after predict(): freed=%%d" %% len(EVENTS), flush=True)

waited = 0.0
while waited < 10.0 and not EVENTS:
    time.sleep(0.25); waited += 0.25

if EVENTS:
    dt, off, cls = EVENTS[0]
    print("input array freed %%.2fs after predict, off_main_thread=%%s (%%s)"
          %% (dt, off, cls), flush=True)
    print("VERDICT: %%s" %% ("CROSS-THREAD RELEASE CONFIRMED" if off else "released on main thread"),
          flush=True)
else:
    print("VERDICT: never freed within 10s", flush=True)
'''

CRASH = r'''
import faulthandler; faulthandler.enable()
import time
import numpy as np, coremltools as ct

P, SHAPE, ROUNDS = %r, %r, %d
m = ct.models.MLModel(P, compute_units=ct.ComputeUnit.CPU_ONLY)

for i in range(ROUNDS):
    # The dict and the array die when predict() returns; the ObjC MLMultiArray
    # bound to Core ML's input port becomes the sole owner of the numpy object.
    m.predict({"x": np.zeros(SHAPE, dtype=np.float32)})
    # Churn the Python heap for longer than the ~1s linger timeout, so the
    # GIL-less Py_DECREF on the reset queue collides with pymalloc.
    deadline = time.time() + 1.5
    while time.time() < deadline:
        junk = [object() for _ in range(512)]
        del junk
    print("round %%d ok" %% i, flush=True)
print("ALL OK", flush=True)
'''


def main():
    tmp = tempfile.mkdtemp(prefix="ctrepro-")
    try:
        pkg = str(Path(tmp) / "tiny.mlpackage")
        rb = subprocess.run([sys.executable, "-c", BUILD % (SHAPE, pkg)],
                            capture_output=True, text=True, timeout=1800)
        if "BUILT" not in rb.stdout:
            print("BUILD FAILED", rb.returncode)
            print(rb.stderr[-2000:])
            return
        size = sum(f.stat().st_size for f in Path(pkg).rglob("*") if f.is_file())
        print(f"tiny model: {size/1024:.1f} KiB, one fp32 input {SHAPE}, one add op\n")

        print("== variant 'detect' ==")
        r = subprocess.run([sys.executable, "-c", DETECT % (pkg, SHAPE)],
                           capture_output=True, text=True, timeout=1800)
        for l in r.stdout.splitlines():
            if l.startswith(("immediately", "input array", "VERDICT")):
                print("   ", l)
        print(f"    rc={r.returncode}")

        print("\n== variant 'crash' ==")
        r = subprocess.run([sys.executable, "-c", CRASH % (pkg, SHAPE, 8)],
                           capture_output=True, text=True, timeout=1800)
        done = [l for l in r.stdout.splitlines() if l.strip()]
        print(f"    rc={r.returncode} (rc=-11 => SIGSEGV reproduced)")
        print(f"    last stdout: {done[-1] if done else '(none)'}")
        for l in r.stderr.splitlines():
            if "MLE5" in l or "Segmentation" in l or "Fatal Python" in l:
                print("    stderr:", l)
    finally:
        shutil.rmtree(tmp, ignore_errors=True)


main()

Verification

• The repro above: both variants pass patched, crash segfaults unpatched.
• A real workload that previously segfaulted in 4 different instance-alternation sequences (multifunction stateful Gemma export, MLState shared across per-size function instances) now runs all sequences clean, and state values verifiably carry (warm-vs-fresh logits diff matches the Swift reference behavior).
• Existing suites on the patched build: test_api_examples.py stateful tests, test_modelpackage.py (49 passed, 1 skipped), test_model.py + test/blob (44 tests) — all green.

One related observation (not addressed here): Core ML keeps the raw mutable_data() pointer of the input array bound to its port for that same ~1 s window, so a caller mutating an input array in place immediately after predict() returns is racing Core ML. The retained reference makes it memory-safe, but it may deserve a note in the docs.

🤖 Generated with Claude Code