xla icon indicating copy to clipboard operation
xla copied to clipboard

Published 20 hours ago • verified publisher icon openxla

Ported profiler's code from pybind11 to nanobind + threadsafety in FT

Open vfdev-5 opened this issue 7 months ago • 0 comments

Description:

  • Ported profiler's code from pybind11 to nanobind
  • Added nanobind mutexes to protect PythonHooks and PythonHookContext under free-threading
    • Race seen in JAX: https://github.com/jax-ml/jax/actions/runs/14651237353/job/41117294221?pr=28245#step:18:1561
#8 xla::profiler::PythonHookContext::ProfileFast(_frame*, int, _object*) /proc/self/cwd/external/xla/xla/python/profiler/internal/python_hooks.cc:292:25

and

#6 xla::profiler::PythonHookContext::ProfileFast(_frame*, int, _object*) /proc/self/cwd/external/xla/xla/python/profiler/internal/python_hooks.cc:292:25

cc @hawkinsp

vfdev-5 avatar Apr 25 '25 10:04 vfdev-5