What is nb_type_0?

[zhqrbitee]

Hi,
I have a use case that needs to pickle some python classes derived from nanobind classes and send to other process. With the pickle package: https://github.com/uqfoundation/dill, I hit an issue when dump those classes (not instance).
For example:

#include "nanobind/nanobind.h"

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

struct Scoreboard {
  Scoreboard() = default;
};

NB_MODULE(base, m) {
  nb::class_<Scoreboard>(m, "Scoreboard", nanobind::dynamic_attr())
      .def(nanobind::init<>{});
}

import typing

from base import Scoreboard
from dataclasses import dataclass
import dill   # pip install dill
import nanobind

class ScoreboardV1(Scoreboard):
    pass


@dataclass
class Struct:
    scoreboard: typing.Type[Scoreboard]


struct = Struct(scoreboard=ScoreboardV1)

try:
    bytes1 = dill.dumps(struct)
    struct1 = dill.loads(bytes1)
except Exception as e:
    # error out here with msg:
    #  Can't pickle <class 'nanobind.nb_type_0'>: it's not found as nanobind.nb_type_0
    print(f"dill error:\n{e}")
else:
    print("dill success")

So what is the nb_type_0 class? I tried to search through the source codes, but didn't find it.

What is interesting is that if I manually set nb_type_0 on nanobind module, it seems workaround the issue:

import nanobind

enable_work_around = True
if enable_work_around:
    meta_cls = type(Scoreboard)
    setattr(nanobind, "nb_type_0", meta_cls)

[wjakob]

It is the nanobind metaclass. In other words, the class that constructs the Python type A when you call nb::class_<A>. It cannot be pickled, and I am not sure how it could be pickled. That's because we don't just have nb_type_0 -- it's part of a family of metaclasses nb_type_N that are created on demand based how how many extra bytes of storage N are required inside the type object. That extra storage could include arbitrary pointers that don't make sense across Python sessions. All of this is deep inside the C++ guts of nanobind and not accessible through normal Python code.

[wjakob]

Actually, thinking a bit more about this, I think that it may be possible to pickle nb_type_0 instances (i.e. with N=0, which means there isn't any supplemental data)`. This would require adding dummy __get/setstate__ methods to the metaclass.

[zhqrbitee]

@wjakob

Thanks for the explanation!
I think the pickle package just need to access to the nb_type_0 class from nanobind module, so it can use it as the metaclass to reconstruct the derived class, e.g.

types.new_class(name="ScoreboardV1", bases=(Scoreboard, ), kwds={"metaclass": nanobind.nb_type_0})

This is why I think after doing this hack, it seems work (at least I didn't find issue with my application):

meta_cls = type(Scoreboard)
setattr(nanobind, "nb_type_0", meta_cls)

Being said that, if I can guarantee that we compiled all our nanobind modules with the same version of nanobind (and same compiler), then would the nb_type_0 class be unique?

[wjakob]

I tried to reproduce your issue using classes that already exist in the nanobind test suite, and using the built-in pickle instead of dill. No luck, the following works fine (This is on 3.12.2).

import test_classes_ext as t
from dataclasses import dataclass
import typing

class StructWithAttrDerived(t.StructWithAttr):
    pass

@dataclass
class TypeWrapper:
    tp: typing.Type[t.Struct]

value = TypeWrapper(tp=StructWithAttrDerived)

import pickle
bytes1 = pickle.dumps(value, protocol=pickle.HIGHEST_PROTOCOL)
struct1 = pickle.loads(bytes1)

[zhqrbitee]

@wjakob

Thanks, I can confirm it works with build-in pickle.
The issue with build-in pickle is that you need to make sure load and dumps are called in the same scope/file (otherwise, you probably see some errors like Can't find 'xxx' object). This is why in my use case, we prefer third party dill (we need to dump all stuffs into bytes and send it to a subprocess and recover from it).

I know technically speaking, it is not a issue of nanobind, but to make it work with dill, do you see any potential issue for doing the hack below (assuming we compiled all our nanobind modules with the same version of nanobind and same compiler):

meta_cls = type(Scoreboard)
setattr(nanobind, "nb_type_0", meta_cls)

[wjakob]

I think that this hack is harmless if that's what you are asking? The C++ parts of nanobind will never actually do anything with the nanobind.* module. The fact that some modules are created in there is sort of an illusion. But that also means that you may not be able to even get the nanobind module if the package isn't installed on the target machine (i.e. an ImportError will be raised). Usually, projects only reference nanobind as a build-time dependency, but not a run-time one.