diff --git a/python/mujoco/callbacks.cc b/python/mujoco/callbacks.cc index ea565127..82ef1032 100644 --- a/python/mujoco/callbacks.cc +++ b/python/mujoco/callbacks.cc @@ -22,6 +22,7 @@ #include #include "errors.h" +#include "gil.h" #include "structs.h" #include "raw.h" #include @@ -32,6 +33,7 @@ namespace mujoco::python { namespace { namespace py = ::pybind11; + [[noreturn]] static void EscapeWithPythonException() { mju_error("Python exception raised"); std::terminate(); // not actually reachable, mju_error doesn't return @@ -121,6 +123,10 @@ static const py::handle MjWrapperLookup(const Raw* ptr) { return MjWrapperLookup(const_cast(ptr)); } +// CallPyCallback takes ownership of py_callback: it will Py_XDECREF it before +// returning or escaping. This is necessary because EscapeWithPythonException() +// calls mju_error which uses longjmp, bypassing C++ stack unwinding and any +// trailing Py_XDECREF in callers. template static Return CallPyCallback(const char* name, PyObject* py_callback, Args... args) { @@ -135,9 +141,12 @@ CallPyCallback(const char* name, PyObject* py_callback, Args... args) { try { if constexpr (std::is_void_v) { callback(args...); + Py_XDECREF(py_callback); return; } else { - return callback(args...).template cast(); + auto result = callback(args...).template cast(); + Py_XDECREF(py_callback); + return result; } } catch (py::error_already_set& e) { e.restore(); @@ -154,13 +163,23 @@ CallPyCallback(const char* name, PyObject* py_callback, Args... args) { PyErr_SetString(PyExc_TypeError, msg.str().c_str()); } } + // Error path: DECREF before escaping (longjmp won't unwind the stack). + Py_XDECREF(py_callback); } EscapeWithPythonException(); } static PyObject* py_mju_user_warning = nullptr; static void PyMjuUserWarning(const char* msg) { - CallPyCallback("mju_user_warning", py_mju_user_warning, msg); + PyObject* cb; + { + py::gil_scoped_acquire gil; + MutexLockIfGilDisabled lock(GetCallbackMutex()); + cb = py_mju_user_warning; + Py_XINCREF(cb); + } + // CallPyCallback takes ownership of cb (will XDECREF it). + CallPyCallback("mju_user_warning", cb, msg); } // We only support ctypes function pointers for these. @@ -171,20 +190,41 @@ static PyObject* py_mju_user_free = nullptr; static PyObject* py_mjcb_passive = nullptr; static void PyMjcbPassive(const raw::MjModel* m, raw::MjData* d) { - CallPyCallback("mjcb_passive", py_mjcb_passive, + PyObject* cb; + { + py::gil_scoped_acquire gil; + MutexLockIfGilDisabled lock(GetCallbackMutex()); + cb = py_mjcb_passive; + Py_XINCREF(cb); + } + CallPyCallback("mjcb_passive", cb, MjWrapperLookup(m), MjWrapperLookup(d)); } static PyObject* py_mjcb_control = nullptr; static void PyMjcbControl(const raw::MjModel* m, raw::MjData* d) { - CallPyCallback("mjcb_control", py_mjcb_control, + PyObject* cb; + { + py::gil_scoped_acquire gil; + MutexLockIfGilDisabled lock(GetCallbackMutex()); + cb = py_mjcb_control; + Py_XINCREF(cb); + } + CallPyCallback("mjcb_control", cb, MjWrapperLookup(m), MjWrapperLookup(d)); } static PyObject* py_mjcb_contactfilter = nullptr; static int PyMjcbContactfilter( const raw::MjModel* m, raw::MjData* d, int geom1, int geom2) { - return CallPyCallback("mjcb_contactfilter", py_mjcb_contactfilter, + PyObject* cb; + { + py::gil_scoped_acquire gil; + MutexLockIfGilDisabled lock(GetCallbackMutex()); + cb = py_mjcb_contactfilter; + Py_XINCREF(cb); + } + return CallPyCallback("mjcb_contactfilter", cb, MjWrapperLookup(m), MjWrapperLookup(d), geom1, geom2); } @@ -192,34 +232,72 @@ static int PyMjcbContactfilter( static PyObject* py_mjcb_sensor = nullptr; static void PyMjcbSensor(const raw::MjModel* m, raw::MjData* d, int stage) { - CallPyCallback("mjcb_sensor", py_mjcb_sensor, + PyObject* cb; + { + py::gil_scoped_acquire gil; + MutexLockIfGilDisabled lock(GetCallbackMutex()); + cb = py_mjcb_sensor; + Py_XINCREF(cb); + } + CallPyCallback("mjcb_sensor", cb, MjWrapperLookup(m), MjWrapperLookup(d), stage); } static PyObject* py_mjcb_time = nullptr; static mjtNum PyMjcbTime() { - return CallPyCallback("mjcb_time", py_mjcb_time); + PyObject* cb; + { + py::gil_scoped_acquire gil; + MutexLockIfGilDisabled lock(GetCallbackMutex()); + cb = py_mjcb_time; + Py_XINCREF(cb); + } + return CallPyCallback("mjcb_time", cb); } static PyObject* py_mjcb_act_dyn = nullptr; static mjtNum PyMjcbActDyn(const raw::MjModel* m, const raw::MjData* d, int id) { - return CallPyCallback("mjcb_act_dyn", py_mjcb_act_dyn, - MjWrapperLookup(m), MjWrapperLookup(d), id); + PyObject* cb; + { + py::gil_scoped_acquire gil; + MutexLockIfGilDisabled lock(GetCallbackMutex()); + cb = py_mjcb_act_dyn; + Py_XINCREF(cb); + } + return CallPyCallback("mjcb_act_dyn", cb, + MjWrapperLookup(m), + MjWrapperLookup(d), id); } static PyObject* py_mjcb_act_gain = nullptr; static mjtNum PyMjcbActGain(const raw::MjModel* m, const raw::MjData* d, int id) { - return CallPyCallback("mjcb_act_gain", py_mjcb_act_gain, - MjWrapperLookup(m), MjWrapperLookup(d), id); + PyObject* cb; + { + py::gil_scoped_acquire gil; + MutexLockIfGilDisabled lock(GetCallbackMutex()); + cb = py_mjcb_act_gain; + Py_XINCREF(cb); + } + return CallPyCallback("mjcb_act_gain", cb, + MjWrapperLookup(m), + MjWrapperLookup(d), id); } static PyObject* py_mjcb_act_bias = nullptr; static mjtNum PyMjcbActBias(const raw::MjModel* m, const raw::MjData* d, int id) { - return CallPyCallback("mjcb_act_bias", py_mjcb_act_bias, - MjWrapperLookup(m), MjWrapperLookup(d), id); + PyObject* cb; + { + py::gil_scoped_acquire gil; + MutexLockIfGilDisabled lock(GetCallbackMutex()); + cb = py_mjcb_act_bias; + Py_XINCREF(cb); + } + return CallPyCallback("mjcb_act_bias", cb, + MjWrapperLookup(m), + MjWrapperLookup(d), id); } // If the Python object is a ctypes function pointer, returns the corresponding @@ -287,66 +365,96 @@ template void SetCallback(py::handle h, CFuncPtr py_trampoline, PyObject** py_callback, CFuncPtr* mj_callback) { CFuncPtr cfuncptr = GetCFuncPtr(h); + PyObject* old = nullptr; if (h.is_none()) { - Py_XDECREF(*py_callback); - *py_callback = nullptr; - *mj_callback = nullptr; + { + MutexLockIfGilDisabled lock(GetCallbackMutex()); + old = *py_callback; + *py_callback = nullptr; + *mj_callback = nullptr; + } } else if (cfuncptr) { - Py_XDECREF(*py_callback); Py_INCREF(h.ptr()); - *py_callback = h.ptr(); - *mj_callback = cfuncptr; + { + MutexLockIfGilDisabled lock(GetCallbackMutex()); + old = *py_callback; + *py_callback = h.ptr(); + *mj_callback = cfuncptr; + } } else if (IsCallable(h)) { - Py_XDECREF(*py_callback); Py_INCREF(h.ptr()); - *py_callback = h.ptr(); - *mj_callback = py_trampoline; + { + MutexLockIfGilDisabled lock(GetCallbackMutex()); + old = *py_callback; + *py_callback = h.ptr(); + *mj_callback = py_trampoline; + } } else { throw py::type_error("callback is not an Optional[Callable]"); } + // XDECREF outside mutex: __del__ may re-enter callback setters/getters. + Py_XDECREF(old); } -py::object GetCallback(PyObject* py_callback) { - if (!py_callback) { +py::object GetCallback(PyObject** py_callback) { + MutexLockIfGilDisabled lock(GetCallbackMutex()); + if (!*py_callback) { return py::none(); } - return py::reinterpret_borrow(py_callback); + return py::reinterpret_borrow(*py_callback); } -PYBIND11_MODULE(_callbacks, pymodule) { +PYBIND11_MODULE(_callbacks, pymodule, pybind11::mod_gil_not_used()) { // Setters pymodule.def("set_mju_user_warning", [](py::handle h) { - SetCallback(h, PyMjuUserWarning, &py_mju_user_warning, &::mju_user_warning); + SetCallback(h, PyMjuUserWarning, &py_mju_user_warning, + &::mju_user_warning); }); pymodule.def("set_mju_user_malloc", [](py::handle h) { + PyObject* old = nullptr; if (h.is_none()) { - Py_XDECREF(py_mju_user_malloc); - py_mju_user_malloc = nullptr; + { + MutexLockIfGilDisabled lock(GetCallbackMutex()); + old = py_mju_user_malloc; + py_mju_user_malloc = nullptr; + } } else { auto* cfuncptr = GetCFuncPtr(h); if (!cfuncptr) { throw py::type_error("mju_user_malloc must be a C function pointer"); } - Py_XDECREF(py_mju_user_malloc); Py_XINCREF(h.ptr()); - py_mju_user_malloc = h.ptr(); - ::mju_user_malloc = cfuncptr; + { + MutexLockIfGilDisabled lock(GetCallbackMutex()); + old = py_mju_user_malloc; + py_mju_user_malloc = h.ptr(); + ::mju_user_malloc = cfuncptr; + } } + Py_XDECREF(old); }); pymodule.def("set_mju_user_free", [](py::handle h) { + PyObject* old = nullptr; if (h.is_none()) { - Py_XDECREF(py_mju_user_free); - py_mju_user_free = nullptr; + { + MutexLockIfGilDisabled lock(GetCallbackMutex()); + old = py_mju_user_free; + py_mju_user_free = nullptr; + } } else { auto* cfuncptr = GetCFuncPtr(h); if (!cfuncptr) { throw py::type_error("mju_user_free must be a C function pointer"); } - Py_XDECREF(py_mju_user_free); Py_XINCREF(h.ptr()); - py_mju_user_free = h.ptr(); - ::mju_user_free = cfuncptr; + { + MutexLockIfGilDisabled lock(GetCallbackMutex()); + old = py_mju_user_free; + py_mju_user_free = h.ptr(); + ::mju_user_free = cfuncptr; + } } + Py_XDECREF(old); }); pymodule.def("set_mjcb_passive", [](py::handle h) { SetCallback(h, PyMjcbPassive, &py_mjcb_passive, &::mjcb_passive); @@ -376,37 +484,37 @@ PYBIND11_MODULE(_callbacks, pymodule) { // Getters pymodule.def("get_mju_user_warning", []() { - return GetCallback(py_mju_user_warning); + return GetCallback(&py_mju_user_warning); }); pymodule.def("get_mju_user_malloc", []() { - return GetCallback(py_mju_user_malloc); + return GetCallback(&py_mju_user_malloc); }); pymodule.def("get_mju_user_free", []() { - return GetCallback(py_mju_user_free); + return GetCallback(&py_mju_user_free); }); pymodule.def("get_mjcb_passive", []() { - return GetCallback(py_mjcb_passive); + return GetCallback(&py_mjcb_passive); }); pymodule.def("get_mjcb_control", []() { - return GetCallback(py_mjcb_control); + return GetCallback(&py_mjcb_control); }); pymodule.def("get_mjcb_contactfilter", []() { - return GetCallback(py_mjcb_contactfilter); + return GetCallback(&py_mjcb_contactfilter); }); pymodule.def("get_mjcb_sensor", []() { - return GetCallback(py_mjcb_sensor); + return GetCallback(&py_mjcb_sensor); }); pymodule.def("get_mjcb_time", []() { - return GetCallback(py_mjcb_time); + return GetCallback(&py_mjcb_time); }); pymodule.def("get_mjcb_act_dyn", []() { - return GetCallback(py_mjcb_act_dyn); + return GetCallback(&py_mjcb_act_dyn); }); pymodule.def("get_mjcb_act_gain", []() { - return GetCallback(py_mjcb_act_gain); + return GetCallback(&py_mjcb_act_gain); }); pymodule.def("get_mjcb_act_bias", []() { - return GetCallback(py_mjcb_act_bias); + return GetCallback(&py_mjcb_act_bias); }); } // PYBIND11_MODULE } // namespace diff --git a/python/mujoco/constants.cc b/python/mujoco/constants.cc index 023b78cb..a07bc61e 100644 --- a/python/mujoco/constants.cc +++ b/python/mujoco/constants.cc @@ -48,7 +48,7 @@ py::tuple MakeTuple(const char* (&strings)[N][3]) { return std::move(result); } -PYBIND11_MODULE(_constants, pymodule) { +PYBIND11_MODULE(_constants, pymodule, pybind11::mod_gil_not_used()) { #define X(var) pymodule.attr(#var) = var // from mjmodel.h diff --git a/python/mujoco/enums.cc b/python/mujoco/enums.cc index 5227fe50..23a81010 100644 --- a/python/mujoco/enums.cc +++ b/python/mujoco/enums.cc @@ -144,7 +144,7 @@ void DefAllEnums(py::module_& m, Tuple&& tuple) { } } -PYBIND11_MODULE(_enums, pymodule) { +PYBIND11_MODULE(_enums, pymodule, pybind11::mod_gil_not_used()) { DefAllEnums(pymodule, python_traits::kAllEnums); } } // namespace diff --git a/python/mujoco/errors.cc b/python/mujoco/errors.cc index 4f9d0a43..1cbc411c 100644 --- a/python/mujoco/errors.cc +++ b/python/mujoco/errors.cc @@ -17,7 +17,7 @@ namespace mujoco::python { namespace { -PYBIND11_MODULE(_errors, m) { +PYBIND11_MODULE(_errors, m, pybind11::mod_gil_not_used()) { m.attr("FatalError") = FatalError::GetPyExc(); m.attr("UnexpectedError") = UnexpectedError::GetPyExc(); } diff --git a/python/mujoco/experimental/studio/native_viewer.cc b/python/mujoco/experimental/studio/native_viewer.cc index 0146db89..de907eaf 100644 --- a/python/mujoco/experimental/studio/native_viewer.cc +++ b/python/mujoco/experimental/studio/native_viewer.cc @@ -201,7 +201,7 @@ class Viewer { std::vector pixels_; }; -PYBIND11_MODULE(native_viewer_cc, m) { +PYBIND11_MODULE(native_viewer_cc, m, pybind11::mod_gil_not_used()) { pybind11::module_::import("mujoco._structs"); pybind11::class_(m, "Viewer") .def(pybind11::init()) diff --git a/python/mujoco/experimental/studio/parser.cc b/python/mujoco/experimental/studio/parser.cc index 6316f521..36f7c2ea 100644 --- a/python/mujoco/experimental/studio/parser.cc +++ b/python/mujoco/experimental/studio/parser.cc @@ -45,7 +45,7 @@ py::object Parse(std::string_view filepath) { } // namespace mujoco::python -PYBIND11_MODULE(parser, m) { +PYBIND11_MODULE(parser, m, pybind11::mod_gil_not_used()) { pybind11::module_::import("mujoco._structs"); m.def("parse", &mujoco::python::Parse, pybind11::return_value_policy::take_ownership); diff --git a/python/mujoco/experimental/studio/renderer.cc b/python/mujoco/experimental/studio/renderer.cc index 06186176..05b7c211 100644 --- a/python/mujoco/experimental/studio/renderer.cc +++ b/python/mujoco/experimental/studio/renderer.cc @@ -75,7 +75,7 @@ class Renderer { } // namespace mujoco::python -PYBIND11_MODULE(renderer, m) { +PYBIND11_MODULE(renderer, m, pybind11::mod_gil_not_used()) { pybind11::module_::import("mujoco._structs"); pybind11::class_(m, "Renderer") .def(pybind11::init()) diff --git a/python/mujoco/experimental/studio/sim.cc b/python/mujoco/experimental/studio/sim.cc index 810461a5..0dd285ed 100644 --- a/python/mujoco/experimental/studio/sim.cc +++ b/python/mujoco/experimental/studio/sim.cc @@ -25,7 +25,7 @@ namespace py = pybind11; using StepControl = mujoco::platform::StepControl; -PYBIND11_MODULE(sim, m) { +PYBIND11_MODULE(sim, m, pybind11::mod_gil_not_used()) { py::module_::import("mujoco._structs"); m.doc() = "MuJoCo platform simulation bindings for Link."; diff --git a/python/mujoco/experimental/studio/ux.cc b/python/mujoco/experimental/studio/ux.cc index 6e5d5a88..08d8750d 100644 --- a/python/mujoco/experimental/studio/ux.cc +++ b/python/mujoco/experimental/studio/ux.cc @@ -51,7 +51,7 @@ struct RenderFlags { std::array flags = {0}; }; -PYBIND11_MODULE(ux, m) { +PYBIND11_MODULE(ux, m, pybind11::mod_gil_not_used()) { py::module_::import("mujoco._structs"); py::class_(m, "RenderFlags") .def(py::init<>()) diff --git a/python/mujoco/functions.cc b/python/mujoco/functions.cc index 97c7fb79..e56643b0 100644 --- a/python/mujoco/functions.cc +++ b/python/mujoco/functions.cc @@ -37,7 +37,7 @@ namespace mujoco::python { namespace { -PYBIND11_MODULE(_functions, pymodule) { +PYBIND11_MODULE(_functions, pymodule, pybind11::mod_gil_not_used()) { namespace py = ::pybind11; namespace traits = python_traits; diff --git a/python/mujoco/gil.h b/python/mujoco/gil.h new file mode 100644 index 00000000..b368f5f1 --- /dev/null +++ b/python/mujoco/gil.h @@ -0,0 +1,62 @@ +// Copyright 2026 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MUJOCO_PYTHON_GIL_H_ +#define MUJOCO_PYTHON_GIL_H_ + +#include + +namespace mujoco::python { + +// A scoped lock that acquires a std::mutex in free-threaded Python builds +// (Py_GIL_DISABLED) and does nothing in standard GIL builds. +// +// In standard builds the GIL already serializes access to Python state, so +// adding a C++ mutex would introduce unnecessary overhead and deadlock risk +// (two locks held simultaneously). In free-threaded builds the GIL is absent, +// so an explicit mutex is required to protect shared mutable C++ state. +// +// Usage: +// static std::mutex my_mutex; +// { +// MutexLockIfGilDisabled lock(my_mutex); +// // ... access shared state ... +// } +#ifdef Py_GIL_DISABLED +class MutexLockIfGilDisabled { + public: + explicit MutexLockIfGilDisabled(std::mutex& mtx) : lock_(mtx) {} + MutexLockIfGilDisabled(const MutexLockIfGilDisabled&) = delete; + MutexLockIfGilDisabled& operator=(const MutexLockIfGilDisabled&) = delete; + + private: + std::lock_guard lock_; +}; +#else +class MutexLockIfGilDisabled { + public: + explicit MutexLockIfGilDisabled(std::mutex& /*unused*/) {} + MutexLockIfGilDisabled(const MutexLockIfGilDisabled&) = delete; + MutexLockIfGilDisabled& operator=(const MutexLockIfGilDisabled&) = delete; +}; +#endif + +inline std::mutex& GetCallbackMutex() { + static std::mutex mtx; + return mtx; +} + +} // namespace mujoco::python + +#endif // MUJOCO_PYTHON_GIL_H_ diff --git a/python/mujoco/indexers.cc b/python/mujoco/indexers.cc index 3bcfb44e..17135421 100644 --- a/python/mujoco/indexers.cc +++ b/python/mujoco/indexers.cc @@ -19,6 +19,7 @@ #include #include +#include "gil.h" #include "indexers.h" #include "raw.h" #include "util/crossplatform.h" @@ -184,7 +185,7 @@ MjModelIndexer::MjModelIndexer(raw::MjModel* m, py::handle owner) name_to_id_(*m), id_to_name_(*m) #define XGROUP(MjModelFieldGroupedViews, field, nfield, FIELD_XMACROS) \ - , field##_(m->nfield, std::nullopt) + , field##_(m->nfield) MJMODEL_VIEW_GROUPS #undef XGROUP {} @@ -194,6 +195,7 @@ MjModelIndexer::MjModelIndexer(raw::MjModel* m, py::handle owner) if (i >= field##_.size() || i < 0) { \ throw py::index_error(IndexErrorMessage(i, field##_.size())); \ } \ + MutexLockIfGilDisabled lock(lazy_init_mutex_); \ auto& indexer = field##_[i]; \ if (!indexer.has_value()) { \ const std::string& name = id_to_name_.field[i]; \ @@ -224,7 +226,7 @@ MjDataIndexer::MjDataIndexer(raw::MjData* d, const raw::MjModel* m, name_to_id_(*m), id_to_name_(*m) #define XGROUP(MjDataGroupedViews, field, nfield, FIELD_XMACROS) \ - , field##_(m->nfield, std::nullopt) + , field##_(m->nfield) MJDATA_VIEW_GROUPS #undef XGROUP {} @@ -234,6 +236,7 @@ MjDataIndexer::MjDataIndexer(raw::MjData* d, const raw::MjModel* m, if (i >= field##_.size() || i < 0) { \ throw py::index_error(IndexErrorMessage(i, field##_.size())); \ } \ + MutexLockIfGilDisabled lock(lazy_init_mutex_); \ auto& indexer = field##_[i]; \ if (!indexer.has_value()) { \ const std::string& name = id_to_name_.field[i]; \ @@ -270,6 +273,7 @@ MJDATA_VIEW_GROUPS #define MJ_M(n) m_->n #define X(type, prefix, var, dim0, dim1) \ py::array_t XGROUP::var() { \ + MutexLockIfGilDisabled lock(lazy_init_mutex_); \ if (!var##_.has_value()) { \ var##_.emplace(MakeArray<&raw::MjModel::dim0>( \ m_->prefix##var, index_, MAKE_SHAPE(dim1), *m_, owner_)); \ @@ -361,6 +365,7 @@ MJMODEL_KEYFRAME #define MJ_M(n) m_->n #define X(type, prefix, var, dim0, dim1) \ py::array_t XGROUP::var() { \ + MutexLockIfGilDisabled lock(lazy_init_mutex_); \ if (!var##_.has_value()) { \ var##_.emplace(MakeArray<&raw::MjModel::dim0>( \ d_->prefix##var, index_, MAKE_SHAPE(dim1), *m_, owner_)); \ diff --git a/python/mujoco/indexers.h b/python/mujoco/indexers.h index 2b9defad..279b241c 100644 --- a/python/mujoco/indexers.h +++ b/python/mujoco/indexers.h @@ -15,6 +15,7 @@ #ifndef MUJOCO_PYTHON_INDEXERS_H_ #define MUJOCO_PYTHON_INDEXERS_H_ +#include #include #include #include @@ -104,6 +105,7 @@ class MjModelGroupedViewsBase { std::string name_; raw::MjModel* m_; pybind11::handle owner_; + mutable std::mutex lazy_init_mutex_; }; #define XGROUP(MjModelGroupedViews, field, nfield, FIELD_XMACROS) \ @@ -142,6 +144,7 @@ class MjModelIndexer { pybind11::handle owner_; NameToIDMappings name_to_id_; IDToNameMappings id_to_name_; + mutable std::mutex lazy_init_mutex_; // Lazily instantiate a grouped views object when accessed from Python, but // cache it once made so that we can return the same one if requested again. @@ -172,6 +175,7 @@ class MjDataGroupedViewsBase { raw::MjData* d_; const raw::MjModel* m_; pybind11::handle owner_; + mutable std::mutex lazy_init_mutex_; }; #define XGROUP(MjDataGroupedViews, field, nfield, FIELD_XMACROS) \ @@ -212,6 +216,7 @@ class MjDataIndexer { pybind11::handle owner_; NameToIDMappings name_to_id_; IDToNameMappings id_to_name_; + mutable std::mutex lazy_init_mutex_; // Lazily instantiate a grouped views object when accessed from Python, but // cache it once made so that we can return the same one if requested again. diff --git a/python/mujoco/render.cc b/python/mujoco/render.cc index d16af680..57c633ec 100644 --- a/python/mujoco/render.cc +++ b/python/mujoco/render.cc @@ -147,7 +147,7 @@ void MjrContextWrapper::Free() { } // namespace _impl namespace { -PYBIND11_MODULE(_render, pymodule) { +PYBIND11_MODULE(_render, pymodule, pybind11::mod_gil_not_used()) { namespace py = ::pybind11; namespace traits = python_traits; diff --git a/python/mujoco/rollout.cc b/python/mujoco/rollout.cc index 0928e09b..d0e34eaa 100644 --- a/python/mujoco/rollout.cc +++ b/python/mujoco/rollout.cc @@ -330,7 +330,7 @@ class Rollout { std::shared_ptr pool_; }; -PYBIND11_MODULE(_rollout, pymodule) { +PYBIND11_MODULE(_rollout, pymodule, pybind11::mod_gil_not_used()) { namespace py = ::pybind11; py::class_(pymodule, "Rollout") diff --git a/python/mujoco/simulate.cc b/python/mujoco/simulate.cc index 29f4dddd..0e7aa99e 100644 --- a/python/mujoco/simulate.cc +++ b/python/mujoco/simulate.cc @@ -323,7 +323,7 @@ inline auto SetIfNotNull(T mujoco::Simulate::* member) { }; } -PYBIND11_MODULE(_simulate, pymodule) { +PYBIND11_MODULE(_simulate, pymodule, pybind11::mod_gil_not_used()) { py::class_(pymodule, "Mutex") .def( "__enter__", [](SimulateMutex& mtx) { mtx.lock(); }, diff --git a/python/mujoco/specs.cc b/python/mujoco/specs.cc index a11fd950..45bc46d6 100644 --- a/python/mujoco/specs.cc +++ b/python/mujoco/specs.cc @@ -225,7 +225,7 @@ static std::string addPrefixAndSuffix(const std::string& original_path, return addSuffixBeforeExtension(prefixed_path, suffix_to_add); } -PYBIND11_MODULE(_specs, m) { +PYBIND11_MODULE(_specs, m, pybind11::mod_gil_not_used()) { auto structs_m = py::module::import("mujoco._structs"); py::function mjmodel_from_raw_ptr = structs_m.attr("MjModel").attr("_from_model_ptr"); diff --git a/python/mujoco/structs.cc b/python/mujoco/structs.cc index 6b35b9c0..359b14eb 100644 --- a/python/mujoco/structs.cc +++ b/python/mujoco/structs.cc @@ -101,7 +101,7 @@ py::tuple RecompileSpec(raw::MjSpec* spec, const MjModelWrapper& old_m, } // namespace -PYBIND11_MODULE(_structs, m) { +PYBIND11_MODULE(_structs, m, pybind11::mod_gil_not_used()) { py::module_::import("mujoco._enums"); // ==================== MJOPTION ============================================= diff --git a/python/mujoco/structs.h b/python/mujoco/structs.h index 2db39ff2..f8b72049 100644 --- a/python/mujoco/structs.h +++ b/python/mujoco/structs.h @@ -16,6 +16,7 @@ #define MUJOCO_PYTHON_STRUCTS_H_ #include +#include #include #include #include @@ -35,6 +36,7 @@ #include #include #include +#include "gil.h" #include "indexers.h" #include "raw.h" #include @@ -154,7 +156,12 @@ class StructListBase { } StructListBase(const StructListBase& other) = delete; - StructListBase(StructListBase&& other) = default; + StructListBase(StructListBase&& other) + : ptr_(other.ptr_), + num_(other.num_), + owner_(std::move(other.owner_)), + wrappers_(std::move(other.wrappers_)) {} + // populate_mutex_ is default-constructed (std::mutex is not movable) virtual ~StructListBase() = default; @@ -175,6 +182,8 @@ class StructListBase { protected: void PopulateUpTo(int n) { + MutexLockIfGilDisabled lock(populate_mutex_); + wrappers_.reserve(n + 1); while (wrappers_.size() <= n) { wrappers_.push_back( std::make_shared>(&ptr_[wrappers_.size()], owner_)); @@ -202,6 +211,7 @@ class StructListBase { // Using shared_ptr here so that we get identical Python objects when slicing. std::vector>> wrappers_; + mutable std::mutex populate_mutex_; }; template diff --git a/python/mujoco/structs_wrappers.cc b/python/mujoco/structs_wrappers.cc index 3444ea6d..60bd3bff 100644 --- a/python/mujoco/structs_wrappers.cc +++ b/python/mujoco/structs_wrappers.cc @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -35,6 +36,7 @@ #include #include #include "errors.h" +#include "gil.h" #include "private.h" #include "raw.h" #include "serialization.h" @@ -50,6 +52,9 @@ namespace mujoco::python::_impl { +using ::mujoco::python::GetCallbackMutex; +using ::mujoco::python::MutexLockIfGilDisabled; + namespace py = ::pybind11; namespace { @@ -221,11 +226,17 @@ MjModelRawPointerMap() { return *hash_map; } +static std::mutex& MjModelMapMutex() { + static auto* mtx = new std::mutex; + return *mtx; +} + MjModelWrapper* MjModelWrapper::FromRawPointer(raw::MjModel* m) noexcept { try { auto& map = MjModelRawPointerMap(); { py::gil_scoped_acquire gil; + MutexLockIfGilDisabled lock(MjModelMapMutex()); auto found = map.find(m); return found != map.end() ? found->second : nullptr; } @@ -250,6 +261,7 @@ MjModelWrapper::MjWrapper(raw::MjModel* ptr) bool is_newly_inserted = false; { py::gil_scoped_acquire gil; + MutexLockIfGilDisabled lock(MjModelMapMutex()); is_newly_inserted = MjModelRawPointerMap().insert({ptr_, this}).second; } if (!is_newly_inserted) { @@ -271,6 +283,7 @@ MjModelWrapper::MjWrapper(MjModelWrapper&& other) bool is_newly_inserted = false; { py::gil_scoped_acquire gil; + MutexLockIfGilDisabled lock(MjModelMapMutex()); is_newly_inserted = MjModelRawPointerMap().insert_or_assign(ptr_, this).second; } @@ -294,6 +307,7 @@ MjModelWrapper::~MjWrapper() { bool erased = false; { py::gil_scoped_acquire gil; + MutexLockIfGilDisabled lock(MjModelMapMutex()); erased = MjModelRawPointerMap().erase(ptr_); } if (!erased) { @@ -611,11 +625,17 @@ absl::flat_hash_map& MjDataRawPointerMap() { return *hash_map; } +static std::mutex& MjDataMapMutex() { + static auto* mtx = new std::mutex; + return *mtx; +} + MjDataWrapper* MjDataWrapper::FromRawPointer(raw::MjData* m) noexcept { try { auto& map = MjDataRawPointerMap(); { py::gil_scoped_acquire gil; + MutexLockIfGilDisabled lock(MjDataMapMutex()); auto found = map.find(m); return found != map.end() ? found->second : nullptr; } @@ -657,6 +677,7 @@ MjDataWrapper::MjWrapper(MjModelWrapper* model) bool is_newly_inserted = false; { py::gil_scoped_acquire gil; + MutexLockIfGilDisabled lock(MjDataMapMutex()); is_newly_inserted = MjDataRawPointerMap().insert({ptr_, this}).second; } if (!is_newly_inserted) { @@ -666,7 +687,7 @@ MjDataWrapper::MjWrapper(MjModelWrapper* model) // install default timer if not already installed { - py::gil_scoped_acquire gil; + MutexLockIfGilDisabled lock(GetCallbackMutex()); if (!mjcb_time) { mjcb_time = GetTime; } @@ -727,6 +748,7 @@ MjDataWrapper::MjWrapper(MjDataWrapper&& other) bool is_newly_inserted = false; { py::gil_scoped_acquire gil; + MutexLockIfGilDisabled lock(MjDataMapMutex()); is_newly_inserted = MjDataRawPointerMap().insert_or_assign(ptr_, this).second; } @@ -760,6 +782,7 @@ MjDataWrapper::MjWrapper(const MjDataWrapper& other, MjModelWrapper* model) bool is_newly_inserted = false; { py::gil_scoped_acquire gil; + MutexLockIfGilDisabled lock(MjDataMapMutex()); is_newly_inserted = MjDataRawPointerMap().insert({ptr_, this}).second; } if (!is_newly_inserted) { @@ -790,6 +813,7 @@ MjDataWrapper::MjWrapper(MjModelWrapper* model, raw::MjData* d) bool is_newly_inserted = false; { py::gil_scoped_acquire gil; + MutexLockIfGilDisabled lock(MjDataMapMutex()); is_newly_inserted = MjDataRawPointerMap().insert({ptr_, this}).second; } if (!is_newly_inserted) { @@ -803,6 +827,7 @@ MjDataWrapper::~MjWrapper() { bool erased = false; { py::gil_scoped_acquire gil; + MutexLockIfGilDisabled lock(MjDataMapMutex()); erased = MjDataRawPointerMap().erase(ptr_); } if (!erased) { diff --git a/python/mujoco/thread_safety_test.py b/python/mujoco/thread_safety_test.py new file mode 100644 index 00000000..1140f496 --- /dev/null +++ b/python/mujoco/thread_safety_test.py @@ -0,0 +1,208 @@ +# Copyright 2026 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Thread-safety tests for MuJoCo Python bindings. + +Validates that the free-threading mutex guards in the bindings work correctly: +- Concurrent access to lazy-initialized indexers and struct lists. +- Concurrent callback setter/getter operations. +- Concurrent MjModel/MjData creation and destruction. +- Re-entrancy safety (no deadlock when __del__ re-enters callback setters). +""" + +import threading + +from absl.testing import absltest +import mujoco + + +# A model with contacts so data.contact is non-empty after stepping. +_CONTACT_XML = r""" + + + + + + + + + +""" + +_NUM_THREADS = 8 +_ITERS_PER_THREAD = 50 + + +class ConcurrentStressTest(absltest.TestCase): + """Tests concurrent access to bindings under multiple threads.""" + + def test_concurrent_indexer_access(self): + """Concurrent first-access of lazy indexers must not corrupt state.""" + model = mujoco.MjModel.from_xml_string(_CONTACT_XML) + data = mujoco.MjData(model) + mujoco.mj_step(model, data) + errors = [] + + def access_indexers(): + try: + for _ in range(_ITERS_PER_THREAD): + # Each of these triggers lazy init on first access. + _ = model.geom(0) + _ = data.qpos + _ = data.xpos + except Exception as e: # pylint: disable=broad-except + errors.append(e) + + threads = [threading.Thread(target=access_indexers) + for _ in range(_NUM_THREADS)] + for t in threads: + t.start() + for t in threads: + t.join() + self.assertEmpty(errors, f"Errors in concurrent indexer access: {errors}") + + def test_concurrent_struct_list_access(self): + """Concurrent access to contact struct lists must not crash.""" + model = mujoco.MjModel.from_xml_string(_CONTACT_XML) + data = mujoco.MjData(model) + # Step to generate contacts. + mujoco.mj_step(model, data) + errors = [] + + def access_contacts(): + try: + for _ in range(_ITERS_PER_THREAD): + ncon = data.ncon + if ncon > 0: + contacts = data.contact[:ncon] + _ = len(contacts) + except Exception as e: # pylint: disable=broad-except + errors.append(e) + + threads = [threading.Thread(target=access_contacts) + for _ in range(_NUM_THREADS)] + for t in threads: + t.start() + for t in threads: + t.join() + self.assertEmpty( + errors, f"Errors in concurrent struct list access: {errors}") + + def test_concurrent_callback_set_get(self): + """Concurrent callback setter/getter calls must not crash or deadlock.""" + errors = [] + + def toggle_callback(): + try: + for i in range(_ITERS_PER_THREAD): + if i % 2 == 0: + mujoco.set_mjcb_passive(lambda m, d: None) + else: + mujoco.set_mjcb_passive(None) + _ = mujoco.get_mjcb_passive() + except Exception as e: # pylint: disable=broad-except + errors.append(e) + + threads = [threading.Thread(target=toggle_callback) + for _ in range(_NUM_THREADS)] + for t in threads: + t.start() + for t in threads: + t.join() + # Clean up. + mujoco.set_mjcb_passive(None) + self.assertEmpty( + errors, f"Errors in concurrent callback set/get: {errors}") + + def test_concurrent_model_data_lifecycle(self): + """Concurrent MjModel/MjData creation and destruction must not crash.""" + errors = [] + + def create_destroy(): + try: + for _ in range(_ITERS_PER_THREAD): + m = mujoco.MjModel.from_xml_string(_CONTACT_XML) + d = mujoco.MjData(m) + mujoco.mj_step(m, d) + del d + del m + except Exception as e: # pylint: disable=broad-except + errors.append(e) + + threads = [threading.Thread(target=create_destroy) + for _ in range(_NUM_THREADS)] + for t in threads: + t.start() + for t in threads: + t.join() + self.assertEmpty( + errors, f"Errors in concurrent model/data lifecycle: {errors}") + + +class ReentrancyTest(absltest.TestCase): + """Tests that mutexes don't deadlock when __del__ re-enters setters.""" + + def test_callback_setter_from_del(self): + """Setting a callback whose __del__ re-enters the setter must not hang.""" + + class ReentrantCallback: + """A callable whose destructor re-enters the callback setter.""" + + def __call__(self, m, d): + pass + + def __del__(self): + # When this object is destroyed by set_mjcb_passive(None) or by + # being replaced, __del__ will try to set the callback again. + # This must not deadlock. + try: + mujoco.set_mjcb_passive(None) + except Exception: # pylint: disable=broad-except + pass # Swallow — we just want to verify no deadlock. + + mujoco.set_mjcb_passive(ReentrantCallback()) + # This replaces the callback → drops last ref → triggers __del__ → + # re-enters set_mjcb_passive. Must complete without hanging. + mujoco.set_mjcb_passive(None) + + # If we get here, no deadlock occurred. + self.assertIsNone(mujoco.get_mjcb_passive()) + + def test_callback_replacement_from_del(self): + """Replacing a callback whose __del__ sets a new callback must not hang.""" + + class ChainedCallback: + """A callable whose destructor sets a different callback.""" + + def __call__(self, m, d): + pass + + def __del__(self): + try: + mujoco.set_mjcb_passive(lambda m, d: None) + except Exception: # pylint: disable=broad-except + pass + + mujoco.set_mjcb_passive(ChainedCallback()) + # Replace with a plain lambda — old ChainedCallback.__del__ fires. + mujoco.set_mjcb_passive(lambda m, d: None) + + # Clean up. + cb = mujoco.get_mjcb_passive() + self.assertIsNotNone(cb) + mujoco.set_mjcb_passive(None) + + +if __name__ == "__main__": + absltest.main()