Skip to content

Build internal libraries as shared libraries in Python wheels - #390

Merged
scal444 merged 3 commits into
NVIDIA-BioNeMo:mainfrom
scal444:fix/wheel-shared-libs-main
Oct 30, 2026
Merged

scal444 merged 3 commits into
NVIDIA-BioNeMo:mainfrom
scal444:fix/wheel-shared-libs-main

Conversation

@scal444

@scal444 scal444 commented Oct 30, 2026

Copy link
Copy Markdown
Collaborator

Static internal libraries copied their objects, including each CUDA fatbin, into every Python module that linked them. The BFGS minimizer kernels alone were duplicated across four modules, which pushed wheels past 150 MB.

Build the internal libraries as shared libraries installed next to the modules, link the shared CUDA runtime, and add $ORIGIN to the rpath that repair_wheel.sh writes. Link etkdg_stages to etkdg_impl, which provides the symbols it uses, instead of etkdg; the old link formed a cycle that CMake only permits among static libraries.

Static internal libraries copied their objects, including each CUDA fatbin,
into every Python module that linked them. The BFGS minimizer kernels alone
were duplicated across four modules, which pushed wheels past 150 MB.

Build the internal libraries as shared libraries installed next to the
modules, link the shared CUDA runtime, and add $ORIGIN to the rpath that
repair_wheel.sh writes. Link etkdg_stages to etkdg_impl, which provides the
symbols it uses, instead of etkdg; the old link formed a cycle that CMake
only permits among static libraries.
@greptile-apps

greptile-apps Bot commented Oct 30, 2026 •

Copy link
Copy Markdown
Contributor

RetriggerView in GreptileConfidence Score: 5/5

[Medium impact] The PR appears safe to merge; no new blocking issues were found.

Summary

The PR builds internal libraries as shared libraries beside the Python modules to avoid copying CUDA code into multiple modules.

  • Python wheels put shared CUDA libraries beside the modules that use them.
  • Force-field energy calculations use one shared reduction launcher.
  • ETKDG stages link to the library that supplies their implementation.

Changes since the previous review only clarify comments in repair_wheel.sh. No new issues or repository-rule violations were found.

Reviews (3) · Last reviewed commit: "Describe the shared-library rpath in rep..." · Reviewed by Greptile

Comment thread cmake/python_install.cmake Resolved
Shared libraries resolve their own dependencies, so:
- Launch reduceEnergiesKernel through FFKernelUtils::launchReduceEnergiesKernel
  in the library that defines it. CUDA 13 gives kernel host stubs hidden
  visibility, so another library cannot launch the kernel directly. This
  replaces the identical MMFF and DistGeom launchers.
- Declare the libraries each test links instead of relying on static
  libraries to forward their private dependencies.
- Build conformer_checkers only with tests or benchmarks so wheels do not
  install it.
@scal444
scal444 requested a review from evasnow1992 October 30, 2026 16:18
Comment thread admin/distribute/repair_wheel.sh Resolved

@evasnow1992 evasnow1992 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just one minor comment. The changes look good to me overall! Thanks for addressing the PyPI artifact size limit issue!

The header still described patching only the _*.so modules without the
$ORIGIN entry for the internal shared libraries.
@scal444
scal444 merged commit d568fd0 into NVIDIA-BioNeMo:main Oct 30, 2026
14 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants