search

nicholasjng / nanobind-bazel

Bazel defs and rules for building Python projects with nanobind extensions.
Apache License 2.0
6 stars 0 forks source link

Add `nanobind_shared_library` rule #31

Closed nicholasjng closed 3 weeks ago

nicholasjng commented 1 month ago

Equivalent to a native cc_shared_library with a static nanobind dependency.

This target enables linking multiple extensions with a single shared nanobind target, producing smaller binding sizes.

Closes #25.

Still TODO: Avoid auto-injection of @nanobind into the nanobind_extension deps, so that we don't link nanobind multiple times.

Reviews and opinions welcome.

nicholasjng commented 1 month ago

More precisely, this is the proposed way of achieving project organization in Bazel laid out in this issue comment: https://github.com/nicholasjng/nanobind-bazel/pull/7#issuecomment-1989303343

tensorflow/__init__.py
tensorflow/libnanobind-tensorflow.so
tensorflow/_tf_cuda.cpython-312-x86_64-linux-gnu.so
tensorflow/_tf_tpu.cpython-312-x86_64-linux-gnu.so
.. (potentially many more)

These binary extension files (tf*.so) share the libnanobind-tensorflow.so component instead of linking a static library into each one. But note that this libnanobind.so variant is local to the tensorflow package. In other words, it isn't copied into a central lib directory and shared with other Python packages since this would be really fragile.

With Bazel, we get the locality for free, and after removing the unconditional dependency on @nanobind, we can substitute a nanobind_shared_library no problem. (It would be nice to catch the case where the user forgets to specify any nanobind dep, but that will require checking the deps by hand)

Note also that this library was compiled with the NB_DOMAIN set to "tensorflow". This isolates the package from random nanobind settings/bindings that other packages may install, and it also gives it a unique name which is important for some platforms. (Windows, e.g., will not load two .dll files with the same name)

This is where I'm still unsure. cc_shared_library takes only lib dependencies, and thus no defines, so it's not possible (or useful) to pass an NB_DOMAIN. We can still produce the same name by tweaking the shared library's name argument (as described in the docstring), but I'm not sure if that's enough to replicate the domain setting.

wjakob commented 1 month ago

Is the plan that TF moves to nanobind? 😮

nicholasjng commented 1 month ago

Is the plan that TF moves to nanobind? 😮

¯\_(ツ)_/¯ No idea, I'm not a TF developer - but fingers crossed!

This was more to recap the motivation for having a target like this.