Question / possible bug leaking python object in binding

[awni]

For some reason I see leak warnings at program exit when passing a certain Python objects to a C++ class with Nanobind bindings. For example for x = 3 the object y below is destroyed and there are no leaks. But when x is a function or a user defined class y leaks.

Is this expected? Is there something wrong with the C++ code or bindings?

The error message is the following:

nanobind: leaked 1 instances!
 - leaked instance 0x100fe1f48 of type "my_ext.Test"
nanobind: leaked 1 types!
 - leaked type "my_ext.Test"
nanobind: leaked 1 functions!
 - leaked function "__init__"
nanobind: this is likely caused by a reference counting issue in the binding code.

Here is the extension and the Python code.

#include <iostream>
#include <nanobind/nanobind.h>

namespace nb = nanobind;
using namespace nb::literals;

class Test {
 public:
  Test(nb::object x) : x(x) { }
  nb::object x;
};

NB_MODULE(my_ext, m) {
  nb::class_<Test>(m, "Test").def(nb::init<nb::object>(), "f"_a);
}

import my_ext

# Doesn't leak
# x = 3

# Leaks
def x():
    pass

y = my_ext.Test(x)

Edit:

I have read through the relevant section in the docs.

• I am using python 3.12 with nanobind 2.4

[wjakob]

I think this is an expected reference and unrelated to nanobind, i.e.: you would have a leak here with other binding frameworks as well, they just wouldn't warn you about it.

Your x function implicitly depends on globals(), which in turn contains y, which references x. To inform Python that this is a collectable cycle, you'll have to implement tp_traverse for your Test type.

This is actually a nice and concise way to illustrate a problem that people often aren't aware of, I think I'll integrate it into the FAQ.

[awni]

I guess it makes sense if the globals hold a reference to the other object. So we’ll have to implement the cycle breaking. Thanks!

[wjakob]

There were so many questions about this over time that I thought it's best to make an exhaustive documentation section that explains the topic: https://nanobind.readthedocs.io/en/latest/refleaks.html

[awni]

That's awesome!! 🙏

[awni]

We have a few cases that boil down to this pattern:

  m.def(
      "test",
      [](nb::callable fun) {
        return nb::cpp_function([fun](){ return 0; });
      },
      "fun"_a);

Just wondering if you know of a clean way to deal with that? We can always make a custom object, bind it, and install the cycle breaking functions, but that seems pretty tedious for every time we want to return a lambda that can enclose Python objects.

[awni]

If I make a nb::cpp_function from a C++ functor with a Python binding and type slots installed it still leaks.

E.g.

  nb::class_<Test>(m, "_Test", nb::type_slots(slots));

  m.def(
      "test",
      [](nb::callable fun) {
        return nb::cpp_function(Test{fun});
      },
      "fun"_a);

It would be nice if there was a way to use nb::cpp_function this way as we use the Extra parameters to set the name, docstring etc which is quite handy for wrapped functions.

Maybe there is a better way that I am missing, not sure.

Thanks for your help with this!

[awni]

@wjakob here's a complete example which shows the problem using nb::cpp_function to wrap another function. Oddly the function itself is not reported as leaking (I assume nanobind is just not detecting it..?):

It feels ike nb::cpp_function should have the cycle clearing slots installed or somehow expose a way to let the user do so. Does that make sense?

#include <nanobind/nanobind.h>
#include <nanobind/stl/function.h>

namespace nb = nanobind;
using namespace nb::literals;

nb::object func_ret(nb::callable f) {
    return nb::cpp_function([f](int i) { return i+1; }, nb::arg("number"));
}

class Test {
 public:
  Test() { };
};

NB_MODULE(my_ext, m) {
  m.def("func_ret", &func_ret);
  nb::class_<Test>(m, "Test").def(nb::init<>());
}

import my_ext

def x():
    pass

t = my_ext.Test()
y = my_ext.func_ret(x)

[awni]

The output from running the python code:

nanobind: leaked 1 instances!
 - leaked instance 0x102c71f48 of type "my_ext.Test"
nanobind: leaked 1 types!
 - leaked type "my_ext.Test"
nanobind: leaked 2 functions!
 - leaked function "__init__"
 - leaked function ""
nanobind: this is likely caused by a reference counting issue in the binding code.

[wjakob]

Although it would in principle be possible to extend the nanobind function object with user-defined garbage collection hooks, this comes at a cost (2 pointers per function) and would be difficult to use in general. Normally, cpp_function is constructed from a lambda function, potentially with a closure object.

When that closure object is the cause of a leak, we'd have to insert two extra functions that operate on the same closure object (to traverse and clear), but something like that is not possible in C++. So all of this would only work with a named functor like in your example, but that is a much more specific use case. For that reason, I'm not eager to modify nanobind to add such a feature.

However: i's possible to create a thin wrapper around nb_func that adds extra functionality, such as garbage collection hooks. You will usually want to implement such low level features at the CPython API level. Here is an example.

I did not test this thoroughly, and there are probably still issues with this code. But it should be a good starting point. In case anything is missing, you can also take a look at src/nb_func.cpp and src/nb_internals.cpp in the nanobind repository for information on how the nanobind function object itself is implemented.

#include <nanobind/nanobind.h>
#include <nanobind/stl/function.h>

namespace nb = nanobind;
using namespace nb::literals;

struct gc_func {
    PyObject_HEAD
    // Vector call implementation that forwards calls to nanobind
    PyObject* (*vectorcall)(PyObject *, PyObject * const*, size_t, PyObject *);
    // The function itself
    PyObject *func;
    // A non-owning reference to something that is owned by 'func'
    PyObject *dep;
};

int gc_func_tp_traverse(PyObject *self, visitproc visit, void *arg) {
    gc_func *w = (gc_func *) self;
    Py_VISIT(w->func);
    Py_VISIT(w->dep);
    #if PY_VERSION_HEX >= 0x03090000
        Py_VISIT(Py_TYPE(self));
    #endif
    return 0;
};

int gc_func_tp_clear(PyObject *self) {
    gc_func *w = (gc_func *) self;
    Py_CLEAR(w->func);
    return 0;
}

PyObject *gc_func_get_doc(PyObject *self, void *) {
    return PyObject_GetAttrString(((gc_func *) self)->func, "__doc__");
}

PyObject* gc_func_vectorcall(PyObject *self, PyObject * const* args, size_t nargs, PyObject *kwnames) {
    return PyObject_Vectorcall(((gc_func *) self)->func, args, nargs, kwnames);
}

void gc_func_dealloc(PyObject *self) {
    PyObject_GC_UnTrack(self);
    Py_XDECREF(((gc_func *) self)->func);
    PyObject_GC_Del(self);
}


static PyMemberDef gc_func_members[] = {
    { "__vectorcalloffset__", T_PYSSIZET,
      (Py_ssize_t) offsetof(gc_func, vectorcall), READONLY, nullptr },
    { nullptr, 0, 0, 0, nullptr }
};

static PyGetSetDef gc_func_getset[] = {
    { "__doc__", gc_func_get_doc, nullptr, nullptr, nullptr },
    { nullptr, nullptr, nullptr, nullptr, nullptr }
};

// Table of custom type slots we want to install
PyType_Slot gc_func_slots[] = {
    { Py_tp_traverse, (void *) gc_func_tp_traverse },
    { Py_tp_clear, (void *) gc_func_tp_clear },
    { Py_tp_getset, (void *) gc_func_getset },
    { Py_tp_members, (void *) gc_func_members },
    { Py_tp_call, (void *) PyVectorcall_Call },
    { Py_tp_dealloc, (void *) gc_func_dealloc },
    { 0, 0 }
};

static PyType_Spec gc_func_spec = {
    /* .name = */ "mlx.gc_func",
    /* .basicsize = */ (int) sizeof(gc_func),
    /* .itemsize = */ 0,
    /* .flags = */ Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | NB_HAVE_VECTORCALL,
    /* .slots = */ gc_func_slots
};

static PyTypeObject *gc_func_tp = nullptr;

nb::callable wrap_it(nb::object func, nb::handle dep) {
    gc_func *r = (gc_func *) PyType_GenericAlloc(gc_func_tp, 0);
    r->func = func.inc_ref().ptr();
    r->dep = dep.ptr();
    r->vectorcall = gc_func_vectorcall;
    return nb::steal<nb::callable>((PyObject *) r);
}

nb::object func_ret(nb::callable f) {
    nb::object func =
        nb::cpp_function([f](int i) { return i + 1; }, nb::arg("number"));
    return wrap_it(std::move(func), f);
}


class Test {
public:
    Test(){};
};

NB_MODULE(test_classes_ext, m) {
    gc_func_tp = (PyTypeObject *) PyType_FromSpec(&gc_func_spec);
    if (!gc_func_tp)
        nb::raise("Could not register GC function type!");

    m.def("func_ret", &func_ret);
    nb::class_<Test>(m, "Test").def(nb::init<>());
}

[awni]

@wjakob understood and thanks a ton for the detailed response with the full code example. That is extremely useful 🙏