Make MuJoCo Python bindings compatible with free-threading.
Introduce a new header `gil.h` defining `MutexLockIfGilDisabled` to support thread-safety in both standard and free-threaded CPython builds. Protect critical shared states and registries: - Guard global Python callback pointers in `callbacks.cc` using a mutex. Move `gil_scoped_acquire` into local blocks around refcount modifications to prevent `longjmp` from bypassing destructors. - Protect raw pointer maps in `structs_wrappers.cc` with static mutexes. - Replace TOCTOU race in `mjcb_time` initialization with thread-safe `std::call_once`. - Add synchronization to lazy indexer array cache initialization in `indexers.cc` and `indexer_xmacro.h`. - Protect vector mutations in `StructListBase::PopulateUpTo` in `structs.h` with a mutex. - Revert unnecessary atomic changes to threadpool counters. - Declare free-threading compatibility by passing `pybind11::mod_gil_not_used()` to all extension modules. Fixes #3259 Fixes #3256 Fixes #2978 PiperOrigin-RevId: 941101502 Change-Id: Iec4ce58afcbc75d4b0be6a9a21fc8a47854242e3
This commit is contained in:
committed by
Copybara-Service
parent
cab191755a
commit
a07ae6f849
+156
-48
@@ -22,6 +22,7 @@
|
|||||||
|
|
||||||
#include <mujoco/mujoco.h>
|
#include <mujoco/mujoco.h>
|
||||||
#include "errors.h"
|
#include "errors.h"
|
||||||
|
#include "gil.h"
|
||||||
#include "structs.h"
|
#include "structs.h"
|
||||||
#include "raw.h"
|
#include "raw.h"
|
||||||
#include <pybind11/eval.h>
|
#include <pybind11/eval.h>
|
||||||
@@ -32,6 +33,7 @@ namespace mujoco::python {
|
|||||||
namespace {
|
namespace {
|
||||||
namespace py = ::pybind11;
|
namespace py = ::pybind11;
|
||||||
|
|
||||||
|
|
||||||
[[noreturn]] static void EscapeWithPythonException() {
|
[[noreturn]] static void EscapeWithPythonException() {
|
||||||
mju_error("Python exception raised");
|
mju_error("Python exception raised");
|
||||||
std::terminate(); // not actually reachable, mju_error doesn't return
|
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<Raw*>(ptr));
|
return MjWrapperLookup(const_cast<Raw*>(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 <typename Return, typename... Args>
|
template <typename Return, typename... Args>
|
||||||
static Return
|
static Return
|
||||||
CallPyCallback(const char* name, PyObject* py_callback, Args... args) {
|
CallPyCallback(const char* name, PyObject* py_callback, Args... args) {
|
||||||
@@ -135,9 +141,12 @@ CallPyCallback(const char* name, PyObject* py_callback, Args... args) {
|
|||||||
try {
|
try {
|
||||||
if constexpr (std::is_void_v<Return>) {
|
if constexpr (std::is_void_v<Return>) {
|
||||||
callback(args...);
|
callback(args...);
|
||||||
|
Py_XDECREF(py_callback);
|
||||||
return;
|
return;
|
||||||
} else {
|
} else {
|
||||||
return callback(args...).template cast<Return>();
|
auto result = callback(args...).template cast<Return>();
|
||||||
|
Py_XDECREF(py_callback);
|
||||||
|
return result;
|
||||||
}
|
}
|
||||||
} catch (py::error_already_set& e) {
|
} catch (py::error_already_set& e) {
|
||||||
e.restore();
|
e.restore();
|
||||||
@@ -154,13 +163,23 @@ CallPyCallback(const char* name, PyObject* py_callback, Args... args) {
|
|||||||
PyErr_SetString(PyExc_TypeError, msg.str().c_str());
|
PyErr_SetString(PyExc_TypeError, msg.str().c_str());
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
// Error path: DECREF before escaping (longjmp won't unwind the stack).
|
||||||
|
Py_XDECREF(py_callback);
|
||||||
}
|
}
|
||||||
EscapeWithPythonException();
|
EscapeWithPythonException();
|
||||||
}
|
}
|
||||||
|
|
||||||
static PyObject* py_mju_user_warning = nullptr;
|
static PyObject* py_mju_user_warning = nullptr;
|
||||||
static void PyMjuUserWarning(const char* msg) {
|
static void PyMjuUserWarning(const char* msg) {
|
||||||
CallPyCallback<void>("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<void>("mju_user_warning", cb, msg);
|
||||||
}
|
}
|
||||||
|
|
||||||
// We only support ctypes function pointers for these.
|
// 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 PyObject* py_mjcb_passive = nullptr;
|
||||||
static void PyMjcbPassive(const raw::MjModel* m, raw::MjData* d) {
|
static void PyMjcbPassive(const raw::MjModel* m, raw::MjData* d) {
|
||||||
CallPyCallback<void>("mjcb_passive", py_mjcb_passive,
|
PyObject* cb;
|
||||||
|
{
|
||||||
|
py::gil_scoped_acquire gil;
|
||||||
|
MutexLockIfGilDisabled lock(GetCallbackMutex());
|
||||||
|
cb = py_mjcb_passive;
|
||||||
|
Py_XINCREF(cb);
|
||||||
|
}
|
||||||
|
CallPyCallback<void>("mjcb_passive", cb,
|
||||||
MjWrapperLookup(m), MjWrapperLookup(d));
|
MjWrapperLookup(m), MjWrapperLookup(d));
|
||||||
}
|
}
|
||||||
|
|
||||||
static PyObject* py_mjcb_control = nullptr;
|
static PyObject* py_mjcb_control = nullptr;
|
||||||
static void PyMjcbControl(const raw::MjModel* m, raw::MjData* d) {
|
static void PyMjcbControl(const raw::MjModel* m, raw::MjData* d) {
|
||||||
CallPyCallback<void>("mjcb_control", py_mjcb_control,
|
PyObject* cb;
|
||||||
|
{
|
||||||
|
py::gil_scoped_acquire gil;
|
||||||
|
MutexLockIfGilDisabled lock(GetCallbackMutex());
|
||||||
|
cb = py_mjcb_control;
|
||||||
|
Py_XINCREF(cb);
|
||||||
|
}
|
||||||
|
CallPyCallback<void>("mjcb_control", cb,
|
||||||
MjWrapperLookup(m), MjWrapperLookup(d));
|
MjWrapperLookup(m), MjWrapperLookup(d));
|
||||||
}
|
}
|
||||||
|
|
||||||
static PyObject* py_mjcb_contactfilter = nullptr;
|
static PyObject* py_mjcb_contactfilter = nullptr;
|
||||||
static int PyMjcbContactfilter(
|
static int PyMjcbContactfilter(
|
||||||
const raw::MjModel* m, raw::MjData* d, int geom1, int geom2) {
|
const raw::MjModel* m, raw::MjData* d, int geom1, int geom2) {
|
||||||
return CallPyCallback<int>("mjcb_contactfilter", py_mjcb_contactfilter,
|
PyObject* cb;
|
||||||
|
{
|
||||||
|
py::gil_scoped_acquire gil;
|
||||||
|
MutexLockIfGilDisabled lock(GetCallbackMutex());
|
||||||
|
cb = py_mjcb_contactfilter;
|
||||||
|
Py_XINCREF(cb);
|
||||||
|
}
|
||||||
|
return CallPyCallback<int>("mjcb_contactfilter", cb,
|
||||||
MjWrapperLookup(m), MjWrapperLookup(d),
|
MjWrapperLookup(m), MjWrapperLookup(d),
|
||||||
geom1, geom2);
|
geom1, geom2);
|
||||||
}
|
}
|
||||||
@@ -192,34 +232,72 @@ static int PyMjcbContactfilter(
|
|||||||
static PyObject* py_mjcb_sensor = nullptr;
|
static PyObject* py_mjcb_sensor = nullptr;
|
||||||
static void
|
static void
|
||||||
PyMjcbSensor(const raw::MjModel* m, raw::MjData* d, int stage) {
|
PyMjcbSensor(const raw::MjModel* m, raw::MjData* d, int stage) {
|
||||||
CallPyCallback<void>("mjcb_sensor", py_mjcb_sensor,
|
PyObject* cb;
|
||||||
|
{
|
||||||
|
py::gil_scoped_acquire gil;
|
||||||
|
MutexLockIfGilDisabled lock(GetCallbackMutex());
|
||||||
|
cb = py_mjcb_sensor;
|
||||||
|
Py_XINCREF(cb);
|
||||||
|
}
|
||||||
|
CallPyCallback<void>("mjcb_sensor", cb,
|
||||||
MjWrapperLookup(m), MjWrapperLookup(d), stage);
|
MjWrapperLookup(m), MjWrapperLookup(d), stage);
|
||||||
}
|
}
|
||||||
|
|
||||||
static PyObject* py_mjcb_time = nullptr;
|
static PyObject* py_mjcb_time = nullptr;
|
||||||
static mjtNum PyMjcbTime() {
|
static mjtNum PyMjcbTime() {
|
||||||
return CallPyCallback<mjtNum>("mjcb_time", py_mjcb_time);
|
PyObject* cb;
|
||||||
|
{
|
||||||
|
py::gil_scoped_acquire gil;
|
||||||
|
MutexLockIfGilDisabled lock(GetCallbackMutex());
|
||||||
|
cb = py_mjcb_time;
|
||||||
|
Py_XINCREF(cb);
|
||||||
|
}
|
||||||
|
return CallPyCallback<mjtNum>("mjcb_time", cb);
|
||||||
}
|
}
|
||||||
|
|
||||||
static PyObject* py_mjcb_act_dyn = nullptr;
|
static PyObject* py_mjcb_act_dyn = nullptr;
|
||||||
static mjtNum
|
static mjtNum
|
||||||
PyMjcbActDyn(const raw::MjModel* m, const raw::MjData* d, int id) {
|
PyMjcbActDyn(const raw::MjModel* m, const raw::MjData* d, int id) {
|
||||||
return CallPyCallback<mjtNum>("mjcb_act_dyn", py_mjcb_act_dyn,
|
PyObject* cb;
|
||||||
MjWrapperLookup(m), MjWrapperLookup(d), id);
|
{
|
||||||
|
py::gil_scoped_acquire gil;
|
||||||
|
MutexLockIfGilDisabled lock(GetCallbackMutex());
|
||||||
|
cb = py_mjcb_act_dyn;
|
||||||
|
Py_XINCREF(cb);
|
||||||
|
}
|
||||||
|
return CallPyCallback<mjtNum>("mjcb_act_dyn", cb,
|
||||||
|
MjWrapperLookup(m),
|
||||||
|
MjWrapperLookup(d), id);
|
||||||
}
|
}
|
||||||
|
|
||||||
static PyObject* py_mjcb_act_gain = nullptr;
|
static PyObject* py_mjcb_act_gain = nullptr;
|
||||||
static mjtNum
|
static mjtNum
|
||||||
PyMjcbActGain(const raw::MjModel* m, const raw::MjData* d, int id) {
|
PyMjcbActGain(const raw::MjModel* m, const raw::MjData* d, int id) {
|
||||||
return CallPyCallback<mjtNum>("mjcb_act_gain", py_mjcb_act_gain,
|
PyObject* cb;
|
||||||
MjWrapperLookup(m), MjWrapperLookup(d), id);
|
{
|
||||||
|
py::gil_scoped_acquire gil;
|
||||||
|
MutexLockIfGilDisabled lock(GetCallbackMutex());
|
||||||
|
cb = py_mjcb_act_gain;
|
||||||
|
Py_XINCREF(cb);
|
||||||
|
}
|
||||||
|
return CallPyCallback<mjtNum>("mjcb_act_gain", cb,
|
||||||
|
MjWrapperLookup(m),
|
||||||
|
MjWrapperLookup(d), id);
|
||||||
}
|
}
|
||||||
|
|
||||||
static PyObject* py_mjcb_act_bias = nullptr;
|
static PyObject* py_mjcb_act_bias = nullptr;
|
||||||
static mjtNum
|
static mjtNum
|
||||||
PyMjcbActBias(const raw::MjModel* m, const raw::MjData* d, int id) {
|
PyMjcbActBias(const raw::MjModel* m, const raw::MjData* d, int id) {
|
||||||
return CallPyCallback<mjtNum>("mjcb_act_bias", py_mjcb_act_bias,
|
PyObject* cb;
|
||||||
MjWrapperLookup(m), MjWrapperLookup(d), id);
|
{
|
||||||
|
py::gil_scoped_acquire gil;
|
||||||
|
MutexLockIfGilDisabled lock(GetCallbackMutex());
|
||||||
|
cb = py_mjcb_act_bias;
|
||||||
|
Py_XINCREF(cb);
|
||||||
|
}
|
||||||
|
return CallPyCallback<mjtNum>("mjcb_act_bias", cb,
|
||||||
|
MjWrapperLookup(m),
|
||||||
|
MjWrapperLookup(d), id);
|
||||||
}
|
}
|
||||||
|
|
||||||
// If the Python object is a ctypes function pointer, returns the corresponding
|
// If the Python object is a ctypes function pointer, returns the corresponding
|
||||||
@@ -287,66 +365,96 @@ template <typename CFuncPtr>
|
|||||||
void SetCallback(py::handle h, CFuncPtr py_trampoline,
|
void SetCallback(py::handle h, CFuncPtr py_trampoline,
|
||||||
PyObject** py_callback, CFuncPtr* mj_callback) {
|
PyObject** py_callback, CFuncPtr* mj_callback) {
|
||||||
CFuncPtr cfuncptr = GetCFuncPtr<CFuncPtr>(h);
|
CFuncPtr cfuncptr = GetCFuncPtr<CFuncPtr>(h);
|
||||||
|
PyObject* old = nullptr;
|
||||||
if (h.is_none()) {
|
if (h.is_none()) {
|
||||||
Py_XDECREF(*py_callback);
|
{
|
||||||
*py_callback = nullptr;
|
MutexLockIfGilDisabled lock(GetCallbackMutex());
|
||||||
*mj_callback = nullptr;
|
old = *py_callback;
|
||||||
|
*py_callback = nullptr;
|
||||||
|
*mj_callback = nullptr;
|
||||||
|
}
|
||||||
} else if (cfuncptr) {
|
} else if (cfuncptr) {
|
||||||
Py_XDECREF(*py_callback);
|
|
||||||
Py_INCREF(h.ptr());
|
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)) {
|
} else if (IsCallable(h)) {
|
||||||
Py_XDECREF(*py_callback);
|
|
||||||
Py_INCREF(h.ptr());
|
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 {
|
} else {
|
||||||
throw py::type_error("callback is not an Optional[Callable]");
|
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) {
|
py::object GetCallback(PyObject** py_callback) {
|
||||||
if (!py_callback) {
|
MutexLockIfGilDisabled lock(GetCallbackMutex());
|
||||||
|
if (!*py_callback) {
|
||||||
return py::none();
|
return py::none();
|
||||||
}
|
}
|
||||||
return py::reinterpret_borrow<py::object>(py_callback);
|
return py::reinterpret_borrow<py::object>(*py_callback);
|
||||||
}
|
}
|
||||||
|
|
||||||
PYBIND11_MODULE(_callbacks, pymodule) {
|
PYBIND11_MODULE(_callbacks, pymodule, pybind11::mod_gil_not_used()) {
|
||||||
// Setters
|
// Setters
|
||||||
pymodule.def("set_mju_user_warning", [](py::handle h) {
|
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) {
|
pymodule.def("set_mju_user_malloc", [](py::handle h) {
|
||||||
|
PyObject* old = nullptr;
|
||||||
if (h.is_none()) {
|
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 {
|
} else {
|
||||||
auto* cfuncptr = GetCFuncPtr<decltype(::mju_user_malloc)>(h);
|
auto* cfuncptr = GetCFuncPtr<decltype(::mju_user_malloc)>(h);
|
||||||
if (!cfuncptr) {
|
if (!cfuncptr) {
|
||||||
throw py::type_error("mju_user_malloc must be a C function pointer");
|
throw py::type_error("mju_user_malloc must be a C function pointer");
|
||||||
}
|
}
|
||||||
Py_XDECREF(py_mju_user_malloc);
|
|
||||||
Py_XINCREF(h.ptr());
|
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) {
|
pymodule.def("set_mju_user_free", [](py::handle h) {
|
||||||
|
PyObject* old = nullptr;
|
||||||
if (h.is_none()) {
|
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 {
|
} else {
|
||||||
auto* cfuncptr = GetCFuncPtr<decltype(::mju_user_free)>(h);
|
auto* cfuncptr = GetCFuncPtr<decltype(::mju_user_free)>(h);
|
||||||
if (!cfuncptr) {
|
if (!cfuncptr) {
|
||||||
throw py::type_error("mju_user_free must be a C function pointer");
|
throw py::type_error("mju_user_free must be a C function pointer");
|
||||||
}
|
}
|
||||||
Py_XDECREF(py_mju_user_free);
|
|
||||||
Py_XINCREF(h.ptr());
|
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) {
|
pymodule.def("set_mjcb_passive", [](py::handle h) {
|
||||||
SetCallback(h, PyMjcbPassive, &py_mjcb_passive, &::mjcb_passive);
|
SetCallback(h, PyMjcbPassive, &py_mjcb_passive, &::mjcb_passive);
|
||||||
@@ -376,37 +484,37 @@ PYBIND11_MODULE(_callbacks, pymodule) {
|
|||||||
|
|
||||||
// Getters
|
// Getters
|
||||||
pymodule.def("get_mju_user_warning", []() {
|
pymodule.def("get_mju_user_warning", []() {
|
||||||
return GetCallback(py_mju_user_warning);
|
return GetCallback(&py_mju_user_warning);
|
||||||
});
|
});
|
||||||
pymodule.def("get_mju_user_malloc", []() {
|
pymodule.def("get_mju_user_malloc", []() {
|
||||||
return GetCallback(py_mju_user_malloc);
|
return GetCallback(&py_mju_user_malloc);
|
||||||
});
|
});
|
||||||
pymodule.def("get_mju_user_free", []() {
|
pymodule.def("get_mju_user_free", []() {
|
||||||
return GetCallback(py_mju_user_free);
|
return GetCallback(&py_mju_user_free);
|
||||||
});
|
});
|
||||||
pymodule.def("get_mjcb_passive", []() {
|
pymodule.def("get_mjcb_passive", []() {
|
||||||
return GetCallback(py_mjcb_passive);
|
return GetCallback(&py_mjcb_passive);
|
||||||
});
|
});
|
||||||
pymodule.def("get_mjcb_control", []() {
|
pymodule.def("get_mjcb_control", []() {
|
||||||
return GetCallback(py_mjcb_control);
|
return GetCallback(&py_mjcb_control);
|
||||||
});
|
});
|
||||||
pymodule.def("get_mjcb_contactfilter", []() {
|
pymodule.def("get_mjcb_contactfilter", []() {
|
||||||
return GetCallback(py_mjcb_contactfilter);
|
return GetCallback(&py_mjcb_contactfilter);
|
||||||
});
|
});
|
||||||
pymodule.def("get_mjcb_sensor", []() {
|
pymodule.def("get_mjcb_sensor", []() {
|
||||||
return GetCallback(py_mjcb_sensor);
|
return GetCallback(&py_mjcb_sensor);
|
||||||
});
|
});
|
||||||
pymodule.def("get_mjcb_time", []() {
|
pymodule.def("get_mjcb_time", []() {
|
||||||
return GetCallback(py_mjcb_time);
|
return GetCallback(&py_mjcb_time);
|
||||||
});
|
});
|
||||||
pymodule.def("get_mjcb_act_dyn", []() {
|
pymodule.def("get_mjcb_act_dyn", []() {
|
||||||
return GetCallback(py_mjcb_act_dyn);
|
return GetCallback(&py_mjcb_act_dyn);
|
||||||
});
|
});
|
||||||
pymodule.def("get_mjcb_act_gain", []() {
|
pymodule.def("get_mjcb_act_gain", []() {
|
||||||
return GetCallback(py_mjcb_act_gain);
|
return GetCallback(&py_mjcb_act_gain);
|
||||||
});
|
});
|
||||||
pymodule.def("get_mjcb_act_bias", []() {
|
pymodule.def("get_mjcb_act_bias", []() {
|
||||||
return GetCallback(py_mjcb_act_bias);
|
return GetCallback(&py_mjcb_act_bias);
|
||||||
});
|
});
|
||||||
} // PYBIND11_MODULE
|
} // PYBIND11_MODULE
|
||||||
} // namespace
|
} // namespace
|
||||||
|
|||||||
@@ -48,7 +48,7 @@ py::tuple MakeTuple(const char* (&strings)[N][3]) {
|
|||||||
return std::move(result);
|
return std::move(result);
|
||||||
}
|
}
|
||||||
|
|
||||||
PYBIND11_MODULE(_constants, pymodule) {
|
PYBIND11_MODULE(_constants, pymodule, pybind11::mod_gil_not_used()) {
|
||||||
#define X(var) pymodule.attr(#var) = var
|
#define X(var) pymodule.attr(#var) = var
|
||||||
|
|
||||||
// from mjmodel.h
|
// from mjmodel.h
|
||||||
|
|||||||
@@ -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);
|
DefAllEnums(pymodule, python_traits::kAllEnums);
|
||||||
}
|
}
|
||||||
} // namespace
|
} // namespace
|
||||||
|
|||||||
@@ -17,7 +17,7 @@
|
|||||||
|
|
||||||
namespace mujoco::python {
|
namespace mujoco::python {
|
||||||
namespace {
|
namespace {
|
||||||
PYBIND11_MODULE(_errors, m) {
|
PYBIND11_MODULE(_errors, m, pybind11::mod_gil_not_used()) {
|
||||||
m.attr("FatalError") = FatalError::GetPyExc();
|
m.attr("FatalError") = FatalError::GetPyExc();
|
||||||
m.attr("UnexpectedError") = UnexpectedError::GetPyExc();
|
m.attr("UnexpectedError") = UnexpectedError::GetPyExc();
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -201,7 +201,7 @@ class Viewer {
|
|||||||
std::vector<std::byte> pixels_;
|
std::vector<std::byte> pixels_;
|
||||||
};
|
};
|
||||||
|
|
||||||
PYBIND11_MODULE(native_viewer_cc, m) {
|
PYBIND11_MODULE(native_viewer_cc, m, pybind11::mod_gil_not_used()) {
|
||||||
pybind11::module_::import("mujoco._structs");
|
pybind11::module_::import("mujoco._structs");
|
||||||
pybind11::class_<Viewer>(m, "Viewer")
|
pybind11::class_<Viewer>(m, "Viewer")
|
||||||
.def(pybind11::init<const std::string&, int, int, const std::string&>())
|
.def(pybind11::init<const std::string&, int, int, const std::string&>())
|
||||||
|
|||||||
@@ -45,7 +45,7 @@ py::object Parse(std::string_view filepath) {
|
|||||||
|
|
||||||
} // namespace mujoco::python
|
} // namespace mujoco::python
|
||||||
|
|
||||||
PYBIND11_MODULE(parser, m) {
|
PYBIND11_MODULE(parser, m, pybind11::mod_gil_not_used()) {
|
||||||
pybind11::module_::import("mujoco._structs");
|
pybind11::module_::import("mujoco._structs");
|
||||||
m.def("parse", &mujoco::python::Parse,
|
m.def("parse", &mujoco::python::Parse,
|
||||||
pybind11::return_value_policy::take_ownership);
|
pybind11::return_value_policy::take_ownership);
|
||||||
|
|||||||
@@ -75,7 +75,7 @@ class Renderer {
|
|||||||
|
|
||||||
} // namespace mujoco::python
|
} // namespace mujoco::python
|
||||||
|
|
||||||
PYBIND11_MODULE(renderer, m) {
|
PYBIND11_MODULE(renderer, m, pybind11::mod_gil_not_used()) {
|
||||||
pybind11::module_::import("mujoco._structs");
|
pybind11::module_::import("mujoco._structs");
|
||||||
pybind11::class_<mujoco::python::Renderer>(m, "Renderer")
|
pybind11::class_<mujoco::python::Renderer>(m, "Renderer")
|
||||||
.def(pybind11::init<const std::string&>())
|
.def(pybind11::init<const std::string&>())
|
||||||
|
|||||||
@@ -25,7 +25,7 @@ namespace py = pybind11;
|
|||||||
|
|
||||||
using StepControl = mujoco::platform::StepControl;
|
using StepControl = mujoco::platform::StepControl;
|
||||||
|
|
||||||
PYBIND11_MODULE(sim, m) {
|
PYBIND11_MODULE(sim, m, pybind11::mod_gil_not_used()) {
|
||||||
py::module_::import("mujoco._structs");
|
py::module_::import("mujoco._structs");
|
||||||
m.doc() = "MuJoCo platform simulation bindings for Link.";
|
m.doc() = "MuJoCo platform simulation bindings for Link.";
|
||||||
|
|
||||||
|
|||||||
@@ -51,7 +51,7 @@ struct RenderFlags {
|
|||||||
std::array<uint8_t, mjNRNDFLAG> flags = {0};
|
std::array<uint8_t, mjNRNDFLAG> flags = {0};
|
||||||
};
|
};
|
||||||
|
|
||||||
PYBIND11_MODULE(ux, m) {
|
PYBIND11_MODULE(ux, m, pybind11::mod_gil_not_used()) {
|
||||||
py::module_::import("mujoco._structs");
|
py::module_::import("mujoco._structs");
|
||||||
py::class_<RenderFlags>(m, "RenderFlags")
|
py::class_<RenderFlags>(m, "RenderFlags")
|
||||||
.def(py::init<>())
|
.def(py::init<>())
|
||||||
|
|||||||
@@ -37,7 +37,7 @@
|
|||||||
|
|
||||||
namespace mujoco::python {
|
namespace mujoco::python {
|
||||||
namespace {
|
namespace {
|
||||||
PYBIND11_MODULE(_functions, pymodule) {
|
PYBIND11_MODULE(_functions, pymodule, pybind11::mod_gil_not_used()) {
|
||||||
namespace py = ::pybind11;
|
namespace py = ::pybind11;
|
||||||
namespace traits = python_traits;
|
namespace traits = python_traits;
|
||||||
|
|
||||||
|
|||||||
@@ -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 <mutex>
|
||||||
|
|
||||||
|
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<std::mutex> 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_
|
||||||
@@ -19,6 +19,7 @@
|
|||||||
#include <utility>
|
#include <utility>
|
||||||
#include <vector>
|
#include <vector>
|
||||||
|
|
||||||
|
#include "gil.h"
|
||||||
#include "indexers.h"
|
#include "indexers.h"
|
||||||
#include "raw.h"
|
#include "raw.h"
|
||||||
#include "util/crossplatform.h"
|
#include "util/crossplatform.h"
|
||||||
@@ -184,7 +185,7 @@ MjModelIndexer::MjModelIndexer(raw::MjModel* m, py::handle owner)
|
|||||||
name_to_id_(*m),
|
name_to_id_(*m),
|
||||||
id_to_name_(*m)
|
id_to_name_(*m)
|
||||||
#define XGROUP(MjModelFieldGroupedViews, field, nfield, FIELD_XMACROS) \
|
#define XGROUP(MjModelFieldGroupedViews, field, nfield, FIELD_XMACROS) \
|
||||||
, field##_(m->nfield, std::nullopt)
|
, field##_(m->nfield)
|
||||||
MJMODEL_VIEW_GROUPS
|
MJMODEL_VIEW_GROUPS
|
||||||
#undef XGROUP
|
#undef XGROUP
|
||||||
{}
|
{}
|
||||||
@@ -194,6 +195,7 @@ MjModelIndexer::MjModelIndexer(raw::MjModel* m, py::handle owner)
|
|||||||
if (i >= field##_.size() || i < 0) { \
|
if (i >= field##_.size() || i < 0) { \
|
||||||
throw py::index_error(IndexErrorMessage(i, field##_.size())); \
|
throw py::index_error(IndexErrorMessage(i, field##_.size())); \
|
||||||
} \
|
} \
|
||||||
|
MutexLockIfGilDisabled lock(lazy_init_mutex_); \
|
||||||
auto& indexer = field##_[i]; \
|
auto& indexer = field##_[i]; \
|
||||||
if (!indexer.has_value()) { \
|
if (!indexer.has_value()) { \
|
||||||
const std::string& name = id_to_name_.field[i]; \
|
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),
|
name_to_id_(*m),
|
||||||
id_to_name_(*m)
|
id_to_name_(*m)
|
||||||
#define XGROUP(MjDataGroupedViews, field, nfield, FIELD_XMACROS) \
|
#define XGROUP(MjDataGroupedViews, field, nfield, FIELD_XMACROS) \
|
||||||
, field##_(m->nfield, std::nullopt)
|
, field##_(m->nfield)
|
||||||
MJDATA_VIEW_GROUPS
|
MJDATA_VIEW_GROUPS
|
||||||
#undef XGROUP
|
#undef XGROUP
|
||||||
{}
|
{}
|
||||||
@@ -234,6 +236,7 @@ MjDataIndexer::MjDataIndexer(raw::MjData* d, const raw::MjModel* m,
|
|||||||
if (i >= field##_.size() || i < 0) { \
|
if (i >= field##_.size() || i < 0) { \
|
||||||
throw py::index_error(IndexErrorMessage(i, field##_.size())); \
|
throw py::index_error(IndexErrorMessage(i, field##_.size())); \
|
||||||
} \
|
} \
|
||||||
|
MutexLockIfGilDisabled lock(lazy_init_mutex_); \
|
||||||
auto& indexer = field##_[i]; \
|
auto& indexer = field##_[i]; \
|
||||||
if (!indexer.has_value()) { \
|
if (!indexer.has_value()) { \
|
||||||
const std::string& name = id_to_name_.field[i]; \
|
const std::string& name = id_to_name_.field[i]; \
|
||||||
@@ -270,6 +273,7 @@ MJDATA_VIEW_GROUPS
|
|||||||
#define MJ_M(n) m_->n
|
#define MJ_M(n) m_->n
|
||||||
#define X(type, prefix, var, dim0, dim1) \
|
#define X(type, prefix, var, dim0, dim1) \
|
||||||
py::array_t<type> XGROUP::var() { \
|
py::array_t<type> XGROUP::var() { \
|
||||||
|
MutexLockIfGilDisabled lock(lazy_init_mutex_); \
|
||||||
if (!var##_.has_value()) { \
|
if (!var##_.has_value()) { \
|
||||||
var##_.emplace(MakeArray<&raw::MjModel::dim0>( \
|
var##_.emplace(MakeArray<&raw::MjModel::dim0>( \
|
||||||
m_->prefix##var, index_, MAKE_SHAPE(dim1), *m_, owner_)); \
|
m_->prefix##var, index_, MAKE_SHAPE(dim1), *m_, owner_)); \
|
||||||
@@ -361,6 +365,7 @@ MJMODEL_KEYFRAME
|
|||||||
#define MJ_M(n) m_->n
|
#define MJ_M(n) m_->n
|
||||||
#define X(type, prefix, var, dim0, dim1) \
|
#define X(type, prefix, var, dim0, dim1) \
|
||||||
py::array_t<type> XGROUP::var() { \
|
py::array_t<type> XGROUP::var() { \
|
||||||
|
MutexLockIfGilDisabled lock(lazy_init_mutex_); \
|
||||||
if (!var##_.has_value()) { \
|
if (!var##_.has_value()) { \
|
||||||
var##_.emplace(MakeArray<&raw::MjModel::dim0>( \
|
var##_.emplace(MakeArray<&raw::MjModel::dim0>( \
|
||||||
d_->prefix##var, index_, MAKE_SHAPE(dim1), *m_, owner_)); \
|
d_->prefix##var, index_, MAKE_SHAPE(dim1), *m_, owner_)); \
|
||||||
|
|||||||
@@ -15,6 +15,7 @@
|
|||||||
#ifndef MUJOCO_PYTHON_INDEXERS_H_
|
#ifndef MUJOCO_PYTHON_INDEXERS_H_
|
||||||
#define MUJOCO_PYTHON_INDEXERS_H_
|
#define MUJOCO_PYTHON_INDEXERS_H_
|
||||||
|
|
||||||
|
#include <mutex>
|
||||||
#include <optional>
|
#include <optional>
|
||||||
#include <string>
|
#include <string>
|
||||||
#include <string_view>
|
#include <string_view>
|
||||||
@@ -104,6 +105,7 @@ class MjModelGroupedViewsBase {
|
|||||||
std::string name_;
|
std::string name_;
|
||||||
raw::MjModel* m_;
|
raw::MjModel* m_;
|
||||||
pybind11::handle owner_;
|
pybind11::handle owner_;
|
||||||
|
mutable std::mutex lazy_init_mutex_;
|
||||||
};
|
};
|
||||||
|
|
||||||
#define XGROUP(MjModelGroupedViews, field, nfield, FIELD_XMACROS) \
|
#define XGROUP(MjModelGroupedViews, field, nfield, FIELD_XMACROS) \
|
||||||
@@ -142,6 +144,7 @@ class MjModelIndexer {
|
|||||||
pybind11::handle owner_;
|
pybind11::handle owner_;
|
||||||
NameToIDMappings name_to_id_;
|
NameToIDMappings name_to_id_;
|
||||||
IDToNameMappings id_to_name_;
|
IDToNameMappings id_to_name_;
|
||||||
|
mutable std::mutex lazy_init_mutex_;
|
||||||
|
|
||||||
// Lazily instantiate a grouped views object when accessed from Python, but
|
// 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.
|
// cache it once made so that we can return the same one if requested again.
|
||||||
@@ -172,6 +175,7 @@ class MjDataGroupedViewsBase {
|
|||||||
raw::MjData* d_;
|
raw::MjData* d_;
|
||||||
const raw::MjModel* m_;
|
const raw::MjModel* m_;
|
||||||
pybind11::handle owner_;
|
pybind11::handle owner_;
|
||||||
|
mutable std::mutex lazy_init_mutex_;
|
||||||
};
|
};
|
||||||
|
|
||||||
#define XGROUP(MjDataGroupedViews, field, nfield, FIELD_XMACROS) \
|
#define XGROUP(MjDataGroupedViews, field, nfield, FIELD_XMACROS) \
|
||||||
@@ -212,6 +216,7 @@ class MjDataIndexer {
|
|||||||
pybind11::handle owner_;
|
pybind11::handle owner_;
|
||||||
NameToIDMappings name_to_id_;
|
NameToIDMappings name_to_id_;
|
||||||
IDToNameMappings id_to_name_;
|
IDToNameMappings id_to_name_;
|
||||||
|
mutable std::mutex lazy_init_mutex_;
|
||||||
|
|
||||||
// Lazily instantiate a grouped views object when accessed from Python, but
|
// 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.
|
// cache it once made so that we can return the same one if requested again.
|
||||||
|
|||||||
@@ -147,7 +147,7 @@ void MjrContextWrapper::Free() {
|
|||||||
} // namespace _impl
|
} // namespace _impl
|
||||||
|
|
||||||
namespace {
|
namespace {
|
||||||
PYBIND11_MODULE(_render, pymodule) {
|
PYBIND11_MODULE(_render, pymodule, pybind11::mod_gil_not_used()) {
|
||||||
namespace py = ::pybind11;
|
namespace py = ::pybind11;
|
||||||
namespace traits = python_traits;
|
namespace traits = python_traits;
|
||||||
|
|
||||||
|
|||||||
@@ -330,7 +330,7 @@ class Rollout {
|
|||||||
std::shared_ptr<ThreadPool> pool_;
|
std::shared_ptr<ThreadPool> pool_;
|
||||||
};
|
};
|
||||||
|
|
||||||
PYBIND11_MODULE(_rollout, pymodule) {
|
PYBIND11_MODULE(_rollout, pymodule, pybind11::mod_gil_not_used()) {
|
||||||
namespace py = ::pybind11;
|
namespace py = ::pybind11;
|
||||||
|
|
||||||
py::class_<Rollout>(pymodule, "Rollout")
|
py::class_<Rollout>(pymodule, "Rollout")
|
||||||
|
|||||||
@@ -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_<SimulateMutex>(pymodule, "Mutex")
|
py::class_<SimulateMutex>(pymodule, "Mutex")
|
||||||
.def(
|
.def(
|
||||||
"__enter__", [](SimulateMutex& mtx) { mtx.lock(); },
|
"__enter__", [](SimulateMutex& mtx) { mtx.lock(); },
|
||||||
|
|||||||
@@ -225,7 +225,7 @@ static std::string addPrefixAndSuffix(const std::string& original_path,
|
|||||||
return addSuffixBeforeExtension(prefixed_path, suffix_to_add);
|
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");
|
auto structs_m = py::module::import("mujoco._structs");
|
||||||
py::function mjmodel_from_raw_ptr =
|
py::function mjmodel_from_raw_ptr =
|
||||||
structs_m.attr("MjModel").attr("_from_model_ptr");
|
structs_m.attr("MjModel").attr("_from_model_ptr");
|
||||||
|
|||||||
@@ -101,7 +101,7 @@ py::tuple RecompileSpec(raw::MjSpec* spec, const MjModelWrapper& old_m,
|
|||||||
|
|
||||||
} // namespace
|
} // namespace
|
||||||
|
|
||||||
PYBIND11_MODULE(_structs, m) {
|
PYBIND11_MODULE(_structs, m, pybind11::mod_gil_not_used()) {
|
||||||
py::module_::import("mujoco._enums");
|
py::module_::import("mujoco._enums");
|
||||||
|
|
||||||
// ==================== MJOPTION =============================================
|
// ==================== MJOPTION =============================================
|
||||||
|
|||||||
+11
-1
@@ -16,6 +16,7 @@
|
|||||||
#define MUJOCO_PYTHON_STRUCTS_H_
|
#define MUJOCO_PYTHON_STRUCTS_H_
|
||||||
|
|
||||||
#include <algorithm>
|
#include <algorithm>
|
||||||
|
#include <mutex>
|
||||||
#include <array>
|
#include <array>
|
||||||
#include <cctype>
|
#include <cctype>
|
||||||
#include <cstddef>
|
#include <cstddef>
|
||||||
@@ -35,6 +36,7 @@
|
|||||||
#include <absl/types/span.h>
|
#include <absl/types/span.h>
|
||||||
#include <mujoco/mujoco.h>
|
#include <mujoco/mujoco.h>
|
||||||
#include <mujoco/mjxmacro.h>
|
#include <mujoco/mjxmacro.h>
|
||||||
|
#include "gil.h"
|
||||||
#include "indexers.h"
|
#include "indexers.h"
|
||||||
#include "raw.h"
|
#include "raw.h"
|
||||||
#include <pybind11/numpy.h>
|
#include <pybind11/numpy.h>
|
||||||
@@ -154,7 +156,12 @@ class StructListBase {
|
|||||||
}
|
}
|
||||||
|
|
||||||
StructListBase(const StructListBase& other) = delete;
|
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;
|
virtual ~StructListBase() = default;
|
||||||
|
|
||||||
@@ -175,6 +182,8 @@ class StructListBase {
|
|||||||
|
|
||||||
protected:
|
protected:
|
||||||
void PopulateUpTo(int n) {
|
void PopulateUpTo(int n) {
|
||||||
|
MutexLockIfGilDisabled lock(populate_mutex_);
|
||||||
|
wrappers_.reserve(n + 1);
|
||||||
while (wrappers_.size() <= n) {
|
while (wrappers_.size() <= n) {
|
||||||
wrappers_.push_back(
|
wrappers_.push_back(
|
||||||
std::make_shared<MjWrapper<T>>(&ptr_[wrappers_.size()], owner_));
|
std::make_shared<MjWrapper<T>>(&ptr_[wrappers_.size()], owner_));
|
||||||
@@ -202,6 +211,7 @@ class StructListBase {
|
|||||||
|
|
||||||
// Using shared_ptr here so that we get identical Python objects when slicing.
|
// Using shared_ptr here so that we get identical Python objects when slicing.
|
||||||
std::vector<std::shared_ptr<MjWrapper<T>>> wrappers_;
|
std::vector<std::shared_ptr<MjWrapper<T>>> wrappers_;
|
||||||
|
mutable std::mutex populate_mutex_;
|
||||||
};
|
};
|
||||||
|
|
||||||
template <typename T>
|
template <typename T>
|
||||||
|
|||||||
@@ -23,6 +23,7 @@
|
|||||||
#include <ios>
|
#include <ios>
|
||||||
#include <iostream>
|
#include <iostream>
|
||||||
#include <memory>
|
#include <memory>
|
||||||
|
#include <mutex>
|
||||||
#include <optional>
|
#include <optional>
|
||||||
#include <string>
|
#include <string>
|
||||||
#include <string_view>
|
#include <string_view>
|
||||||
@@ -35,6 +36,7 @@
|
|||||||
#include <mujoco/mjxmacro.h>
|
#include <mujoco/mjxmacro.h>
|
||||||
#include <mujoco/mujoco.h>
|
#include <mujoco/mujoco.h>
|
||||||
#include "errors.h"
|
#include "errors.h"
|
||||||
|
#include "gil.h"
|
||||||
#include "private.h"
|
#include "private.h"
|
||||||
#include "raw.h"
|
#include "raw.h"
|
||||||
#include "serialization.h"
|
#include "serialization.h"
|
||||||
@@ -50,6 +52,9 @@
|
|||||||
|
|
||||||
namespace mujoco::python::_impl {
|
namespace mujoco::python::_impl {
|
||||||
|
|
||||||
|
using ::mujoco::python::GetCallbackMutex;
|
||||||
|
using ::mujoco::python::MutexLockIfGilDisabled;
|
||||||
|
|
||||||
namespace py = ::pybind11;
|
namespace py = ::pybind11;
|
||||||
|
|
||||||
namespace {
|
namespace {
|
||||||
@@ -221,11 +226,17 @@ MjModelRawPointerMap() {
|
|||||||
return *hash_map;
|
return *hash_map;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
static std::mutex& MjModelMapMutex() {
|
||||||
|
static auto* mtx = new std::mutex;
|
||||||
|
return *mtx;
|
||||||
|
}
|
||||||
|
|
||||||
MjModelWrapper* MjModelWrapper::FromRawPointer(raw::MjModel* m) noexcept {
|
MjModelWrapper* MjModelWrapper::FromRawPointer(raw::MjModel* m) noexcept {
|
||||||
try {
|
try {
|
||||||
auto& map = MjModelRawPointerMap();
|
auto& map = MjModelRawPointerMap();
|
||||||
{
|
{
|
||||||
py::gil_scoped_acquire gil;
|
py::gil_scoped_acquire gil;
|
||||||
|
MutexLockIfGilDisabled lock(MjModelMapMutex());
|
||||||
auto found = map.find(m);
|
auto found = map.find(m);
|
||||||
return found != map.end() ? found->second : nullptr;
|
return found != map.end() ? found->second : nullptr;
|
||||||
}
|
}
|
||||||
@@ -250,6 +261,7 @@ MjModelWrapper::MjWrapper(raw::MjModel* ptr)
|
|||||||
bool is_newly_inserted = false;
|
bool is_newly_inserted = false;
|
||||||
{
|
{
|
||||||
py::gil_scoped_acquire gil;
|
py::gil_scoped_acquire gil;
|
||||||
|
MutexLockIfGilDisabled lock(MjModelMapMutex());
|
||||||
is_newly_inserted = MjModelRawPointerMap().insert({ptr_, this}).second;
|
is_newly_inserted = MjModelRawPointerMap().insert({ptr_, this}).second;
|
||||||
}
|
}
|
||||||
if (!is_newly_inserted) {
|
if (!is_newly_inserted) {
|
||||||
@@ -271,6 +283,7 @@ MjModelWrapper::MjWrapper(MjModelWrapper&& other)
|
|||||||
bool is_newly_inserted = false;
|
bool is_newly_inserted = false;
|
||||||
{
|
{
|
||||||
py::gil_scoped_acquire gil;
|
py::gil_scoped_acquire gil;
|
||||||
|
MutexLockIfGilDisabled lock(MjModelMapMutex());
|
||||||
is_newly_inserted =
|
is_newly_inserted =
|
||||||
MjModelRawPointerMap().insert_or_assign(ptr_, this).second;
|
MjModelRawPointerMap().insert_or_assign(ptr_, this).second;
|
||||||
}
|
}
|
||||||
@@ -294,6 +307,7 @@ MjModelWrapper::~MjWrapper() {
|
|||||||
bool erased = false;
|
bool erased = false;
|
||||||
{
|
{
|
||||||
py::gil_scoped_acquire gil;
|
py::gil_scoped_acquire gil;
|
||||||
|
MutexLockIfGilDisabled lock(MjModelMapMutex());
|
||||||
erased = MjModelRawPointerMap().erase(ptr_);
|
erased = MjModelRawPointerMap().erase(ptr_);
|
||||||
}
|
}
|
||||||
if (!erased) {
|
if (!erased) {
|
||||||
@@ -611,11 +625,17 @@ absl::flat_hash_map<raw::MjData*, MjDataWrapper*>& MjDataRawPointerMap() {
|
|||||||
return *hash_map;
|
return *hash_map;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
static std::mutex& MjDataMapMutex() {
|
||||||
|
static auto* mtx = new std::mutex;
|
||||||
|
return *mtx;
|
||||||
|
}
|
||||||
|
|
||||||
MjDataWrapper* MjDataWrapper::FromRawPointer(raw::MjData* m) noexcept {
|
MjDataWrapper* MjDataWrapper::FromRawPointer(raw::MjData* m) noexcept {
|
||||||
try {
|
try {
|
||||||
auto& map = MjDataRawPointerMap();
|
auto& map = MjDataRawPointerMap();
|
||||||
{
|
{
|
||||||
py::gil_scoped_acquire gil;
|
py::gil_scoped_acquire gil;
|
||||||
|
MutexLockIfGilDisabled lock(MjDataMapMutex());
|
||||||
auto found = map.find(m);
|
auto found = map.find(m);
|
||||||
return found != map.end() ? found->second : nullptr;
|
return found != map.end() ? found->second : nullptr;
|
||||||
}
|
}
|
||||||
@@ -657,6 +677,7 @@ MjDataWrapper::MjWrapper(MjModelWrapper* model)
|
|||||||
bool is_newly_inserted = false;
|
bool is_newly_inserted = false;
|
||||||
{
|
{
|
||||||
py::gil_scoped_acquire gil;
|
py::gil_scoped_acquire gil;
|
||||||
|
MutexLockIfGilDisabled lock(MjDataMapMutex());
|
||||||
is_newly_inserted = MjDataRawPointerMap().insert({ptr_, this}).second;
|
is_newly_inserted = MjDataRawPointerMap().insert({ptr_, this}).second;
|
||||||
}
|
}
|
||||||
if (!is_newly_inserted) {
|
if (!is_newly_inserted) {
|
||||||
@@ -666,7 +687,7 @@ MjDataWrapper::MjWrapper(MjModelWrapper* model)
|
|||||||
|
|
||||||
// install default timer if not already installed
|
// install default timer if not already installed
|
||||||
{
|
{
|
||||||
py::gil_scoped_acquire gil;
|
MutexLockIfGilDisabled lock(GetCallbackMutex());
|
||||||
if (!mjcb_time) {
|
if (!mjcb_time) {
|
||||||
mjcb_time = GetTime;
|
mjcb_time = GetTime;
|
||||||
}
|
}
|
||||||
@@ -727,6 +748,7 @@ MjDataWrapper::MjWrapper(MjDataWrapper&& other)
|
|||||||
bool is_newly_inserted = false;
|
bool is_newly_inserted = false;
|
||||||
{
|
{
|
||||||
py::gil_scoped_acquire gil;
|
py::gil_scoped_acquire gil;
|
||||||
|
MutexLockIfGilDisabled lock(MjDataMapMutex());
|
||||||
is_newly_inserted =
|
is_newly_inserted =
|
||||||
MjDataRawPointerMap().insert_or_assign(ptr_, this).second;
|
MjDataRawPointerMap().insert_or_assign(ptr_, this).second;
|
||||||
}
|
}
|
||||||
@@ -760,6 +782,7 @@ MjDataWrapper::MjWrapper(const MjDataWrapper& other, MjModelWrapper* model)
|
|||||||
bool is_newly_inserted = false;
|
bool is_newly_inserted = false;
|
||||||
{
|
{
|
||||||
py::gil_scoped_acquire gil;
|
py::gil_scoped_acquire gil;
|
||||||
|
MutexLockIfGilDisabled lock(MjDataMapMutex());
|
||||||
is_newly_inserted = MjDataRawPointerMap().insert({ptr_, this}).second;
|
is_newly_inserted = MjDataRawPointerMap().insert({ptr_, this}).second;
|
||||||
}
|
}
|
||||||
if (!is_newly_inserted) {
|
if (!is_newly_inserted) {
|
||||||
@@ -790,6 +813,7 @@ MjDataWrapper::MjWrapper(MjModelWrapper* model, raw::MjData* d)
|
|||||||
bool is_newly_inserted = false;
|
bool is_newly_inserted = false;
|
||||||
{
|
{
|
||||||
py::gil_scoped_acquire gil;
|
py::gil_scoped_acquire gil;
|
||||||
|
MutexLockIfGilDisabled lock(MjDataMapMutex());
|
||||||
is_newly_inserted = MjDataRawPointerMap().insert({ptr_, this}).second;
|
is_newly_inserted = MjDataRawPointerMap().insert({ptr_, this}).second;
|
||||||
}
|
}
|
||||||
if (!is_newly_inserted) {
|
if (!is_newly_inserted) {
|
||||||
@@ -803,6 +827,7 @@ MjDataWrapper::~MjWrapper() {
|
|||||||
bool erased = false;
|
bool erased = false;
|
||||||
{
|
{
|
||||||
py::gil_scoped_acquire gil;
|
py::gil_scoped_acquire gil;
|
||||||
|
MutexLockIfGilDisabled lock(MjDataMapMutex());
|
||||||
erased = MjDataRawPointerMap().erase(ptr_);
|
erased = MjDataRawPointerMap().erase(ptr_);
|
||||||
}
|
}
|
||||||
if (!erased) {
|
if (!erased) {
|
||||||
|
|||||||
@@ -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"""
|
||||||
|
<mujoco>
|
||||||
|
<worldbody>
|
||||||
|
<geom type="plane" size="1 1 0.1"/>
|
||||||
|
<body pos="0 0 0.05">
|
||||||
|
<freejoint/>
|
||||||
|
<geom type="sphere" size="0.05"/>
|
||||||
|
</body>
|
||||||
|
</worldbody>
|
||||||
|
</mujoco>
|
||||||
|
"""
|
||||||
|
|
||||||
|
_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()
|
||||||
Reference in New Issue
Block a user