Skip to content

Commit 3b45e02

Browse files
yutingyemeta-codesync[bot]
authored andcommitted
Fix Linux wheel C++ runtime loading (#1566)
Summary: Pull Request resolved: #1566 Linux PyPI wheels bundle a renamed private libstdc++ because conda-forge runtime dependencies require newer GLIBCXX symbols than manylinux_2_28's system libstdc++ provides. The original import failure was caused by loading that private C++ runtime before PyTorch under Python's default RTLD_LOCAL mode: import pymomentum.geometry import torch That lets a second libstdc++ instance exist in the process before torch._C is loaded, and torch can segfault during import. A scoped RTLD_GLOBAL loader makes that order pass, but it is not a correct fix: it exposes PyMomentum's private libstdc++ symbols to the whole process. In the opposite order, that can poison later native library loads. In particular, Triton's compiler fails in ConvertTritonGPUToLLVM after this order: import torch import pymomentum.geometry import pymomentum.quaternion # backend=\"triton\" Fix this by changing the Linux wheel repair shim to import torch, when torch is available, immediately before loading any PyMomentum native extension. The extension still loads with Python's normal RTLD_LOCAL behavior, so PyMomentum's private libstdc++ remains hidden from later native extensions such as Triton, while torch owns the process-global C++ runtime state before PyMomentum's private runtime is loaded. The generated pymomentum/__init__.py is added only during Linux wheel repair. This keeps macOS and Windows behavior unchanged for now because the known bug is in Linux ELF/dlopen/libstdc++ interaction. If we later want one package shape on all platforms, we can move a no-op-on-non-Linux __init__.py into the source tree, but that should be a separate deliberate package-shape change. Also expand wheel smoke tests so every wheel variant installs torch and checks the pymomentum.geometry-before-torch order. The strict GPU wheel test continues to require CUDA and Triton. Reviewed By: cstollmeta Differential Revision: D110929696 fbshipit-source-id: c46695726d1ecd76e04ff10cff5c1080b82e22b6
1 parent 31d0933 commit 3b45e02

2 files changed

Lines changed: 145 additions & 29 deletions

File tree

‎pyproject-pypi.toml.j2‎

Lines changed: 95 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,10 @@
1919
# The pixi environment's pybind11 uses pybind11NewTools which requires python3_add_library
2020
# from CMake's FindPython, but manylinux Python only has Interpreter component
2121
"-DMOMENTUM_USE_SYSTEM_PYBIND11=OFF",
22-
# Use manylinux container's compiler (gcc-toolset-12) instead of pixi's GCC 14.3
23-
# to ensure manylinux_2_28 compatibility (requires GLIBCXX <= 3.4.24, CXXABI <= 1.3.11)
24-
"-DCMAKE_CXX_FLAGS=-static-libstdc++ -static-libgcc",
25-
"-DCMAKE_C_FLAGS=-static-libgcc",
22+
# Do not static-link libstdc++ into extension modules: duplicated static C++
23+
# runtime symbols can put later extension imports, notably torch._C, into an
24+
# incompatible state under Python's default RTLD_LOCAL mode. The Linux wheel
25+
# repair step below gives libstdc++ a private SONAME instead.
2626
{%- endmacro %}
2727

2828
{#- macOS-specific args -#}
@@ -364,8 +364,10 @@ repair-wheel-command = """
364364
# - manylinux_2_28 policy rejects our libs (too-recent GLIBCXX symbols)
365365
# - linux_x86_64 policy is not available inside the manylinux container
366366
#
367-
# We manually bundle dependencies and rename libstdc++ to avoid the dynamic
368-
# linker collision where Python/torch load the system libstdc++.so.6 first.
367+
# We manually bundle dependencies, including libstdc++, because conda-forge
368+
# dependencies require newer GLIBCXX symbols than manylinux_2_28's system
369+
# libstdc++ provides. libstdc++ is renamed to a private SONAME so torch-first
370+
# imports cannot pin the older system libstdc++.so.6 for PyMomentum's DSOs.
369371

370372
set -e
371373

@@ -377,6 +379,80 @@ echo '=== Unpacking wheel ==='
377379
cd $WORK
378380
unzip -q {wheel}
379381

382+
echo '=== Adding Linux native loader shim ==='
383+
cat > pymomentum/__init__.py <<'PYMOMENTUM_INIT'
384+
# Copyright (c) Meta Platforms, Inc. and affiliates.
385+
#
386+
# This source code is licensed under the MIT license found in the
387+
# LICENSE file in the root directory of this source tree.
388+
389+
from __future__ import annotations
390+
391+
import importlib.abc
392+
import importlib.machinery
393+
import sys
394+
395+
396+
_NATIVE_MODULES = {
397+
f"{__name__}.{module}"
398+
for module in (
399+
"axel",
400+
"camera",
401+
"diff_geometry",
402+
"geometry",
403+
"marker_tracking",
404+
"renderer",
405+
"solver",
406+
"solver2",
407+
)
408+
}
409+
410+
411+
class _TorchFirstExtensionLoader(importlib.abc.Loader):
412+
def __init__(self, loader: importlib.abc.Loader) -> None:
413+
self._loader = loader
414+
415+
def __getattr__(self, name: str):
416+
return getattr(self._loader, name)
417+
418+
def create_module(self, spec):
419+
_import_torch_if_available()
420+
if not hasattr(self._loader, "create_module"):
421+
return None
422+
return self._loader.create_module(spec)
423+
424+
def exec_module(self, module) -> None:
425+
_import_torch_if_available()
426+
self._loader.exec_module(module)
427+
428+
429+
class _TorchFirstExtensionFinder(importlib.abc.MetaPathFinder):
430+
def find_spec(self, fullname: str, path=None, target=None):
431+
if fullname not in _NATIVE_MODULES:
432+
return None
433+
spec = importlib.machinery.PathFinder.find_spec(fullname, path)
434+
if spec is None or not isinstance(
435+
spec.loader, importlib.machinery.ExtensionFileLoader
436+
):
437+
return spec
438+
if not isinstance(spec.loader, _TorchFirstExtensionLoader):
439+
spec.loader = _TorchFirstExtensionLoader(spec.loader)
440+
return spec
441+
442+
443+
def _import_torch_if_available() -> None:
444+
try:
445+
import torch # noqa: F401
446+
except ImportError:
447+
pass
448+
449+
450+
if sys.platform.startswith("linux") and not any(
451+
isinstance(finder, _TorchFirstExtensionFinder) for finder in sys.meta_path
452+
):
453+
sys.meta_path.insert(0, _TorchFirstExtensionFinder())
454+
PYMOMENTUM_INIT
455+
380456
# Find or create the .libs directory for this distribution.
381457
LIBS_DIR="{{ wheel_libs_dir }}"
382458
mkdir -p "$LIBS_DIR"
@@ -392,7 +468,6 @@ is_skip_lib() {
392468
libresolv*|libnss*|ld-linux*) return 0 ;;
393469
libgcc_s*) return 0 ;;
394470
libpython*) return 0 ;;
395-
libstdc++.so*) return 0 ;;
396471
*) return 1 ;;
397472
esac
398473
}
@@ -462,35 +537,30 @@ done
462537
echo "Bundled libraries:"
463538
ls -la "$LIBS_DIR/"
464539

465-
# Special handling: rename libstdc++ to avoid dynamic linker collision
540+
# Keep PyMomentum's C++ runtime dependency private. Python package init imports
541+
# torch, when available, before loading PyMomentum extension modules so torch
542+
# owns its process-global C++ runtime state without exposing PyMomentum's private
543+
# libstdc++ symbols to later native extensions such as Triton.
466544
echo '=== Renaming libstdc++ ==='
467545
RENAMED_LIBSTDCXX="libstdc++-pymomentum.so.6"
468-
469-
# Find pixi's libstdc++ (the one with GLIBCXX_3.4.29)
470-
PIXI_LIBSTDCXX=$(find "$PIXI_ENV/lib" -name 'libstdc++.so.6.*' -not -type l | sort -V | tail -1)
471-
if [ -z "$PIXI_LIBSTDCXX" ]; then
472-
PIXI_LIBSTDCXX="$PIXI_ENV/lib/libstdc++.so.6"
546+
if [ -f "$LIBS_DIR/libstdc++.so.6" ]; then
547+
mv "$LIBS_DIR/libstdc++.so.6" "$LIBS_DIR/$RENAMED_LIBSTDCXX"
548+
else
549+
PIXI_LIBSTDCXX=$(find "$PIXI_ENV/lib" -name 'libstdc++.so.6*' -not -type l | sort -V | tail -1)
550+
if [ -z "$PIXI_LIBSTDCXX" ]; then
551+
PIXI_LIBSTDCXX="$PIXI_ENV/lib/libstdc++.so.6"
552+
fi
553+
cp "$(readlink -f "$PIXI_LIBSTDCXX")" "$LIBS_DIR/$RENAMED_LIBSTDCXX"
473554
fi
474-
PIXI_LIBSTDCXX=$(readlink -f "$PIXI_LIBSTDCXX")
475-
echo "Pixi libstdc++: $PIXI_LIBSTDCXX"
476-
477-
# Copy with renamed name
478-
cp "$PIXI_LIBSTDCXX" "$LIBS_DIR/$RENAMED_LIBSTDCXX"
479555
chmod 755 "$LIBS_DIR/$RENAMED_LIBSTDCXX"
480-
# Remove the original-named copy if bundled above
481-
rm -f "$LIBS_DIR/libstdc++.so.6"
482-
483-
# Set the SONAME of the renamed library
484556
patchelf --set-soname "$RENAMED_LIBSTDCXX" "$LIBS_DIR/$RENAMED_LIBSTDCXX"
485557

486-
# Patch DT_NEEDED in all .so files
487558
echo '=== Patching DT_NEEDED and RPATH ==='
488559
find . -name '*.so*' -type f | while read so_file; do
489560
case "$so_file" in
490-
*libstdc++-pymomentum*) continue ;;
561+
*"$RENAMED_LIBSTDCXX") continue ;;
491562
esac
492563

493-
# Replace libstdc++.so.6 reference with renamed version
494564
if patchelf --print-needed "$so_file" 2>/dev/null | grep -q '^libstdc++[.]so[.]6$'; then
495565
echo " Patching DT_NEEDED: $so_file"
496566
patchelf --replace-needed libstdc++.so.6 "$RENAMED_LIBSTDCXX" "$so_file"

‎scripts/test_wheel.py‎

Lines changed: 50 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -131,8 +131,6 @@ def get_import_test_script(
131131
if torch_runtime:
132132
lines.extend(
133133
[
134-
" # Import torch before PyMomentum's torch-backed Python modules",
135-
" # or any PyMomentum extension so PyTorch sets up native paths first.",
136134
" import torch",
137135
' print(f" PyTorch {torch.__version__}, CUDA={torch.version.cuda}")',
138136
]
@@ -287,6 +285,23 @@ def get_import_test_script(
287285
return "\n".join(lines) + "\n"
288286

289287

288+
def get_torch_after_pymomentum_script() -> str:
289+
"""Return a regression test for PyMomentum's native loader state."""
290+
return "\n".join(
291+
[
292+
"import faulthandler",
293+
"faulthandler.enable()",
294+
'print("Testing pymomentum.geometry before torch...")',
295+
"import pymomentum.geometry",
296+
'print(" pymomentum.geometry imported")',
297+
"import torch",
298+
'print(f" PyTorch {torch.__version__} imported")',
299+
'print("[PASS] pymomentum.geometry before torch")',
300+
"",
301+
]
302+
)
303+
304+
290305
def run_import_tests(
291306
python_exe: str,
292307
wheel_type: str,
@@ -305,6 +320,21 @@ def run_import_tests(
305320
env.pop(name, None)
306321
# Keep the checkout's namespace-package directories out of wheel import checks.
307322
with tempfile.TemporaryDirectory() as tmpdir:
323+
if has_torch_runtime(wheel_type):
324+
result = subprocess.run(
325+
[
326+
python_exe,
327+
"-I",
328+
"-c",
329+
get_torch_after_pymomentum_script(),
330+
],
331+
capture_output=False,
332+
cwd=tmpdir,
333+
env=env,
334+
)
335+
if result.returncode != 0:
336+
return result.returncode
337+
308338
result = subprocess.run(
309339
[
310340
python_exe,
@@ -374,7 +404,7 @@ def test_in_container(
374404
"$PYTHON -m uv venv --python $PYTHON /tmp/test_venv",
375405
"",
376406
]
377-
if has_torch_extensions(wheel_type):
407+
if has_torch_runtime(wheel_type):
378408
torch_index = get_torch_index(wheel_type)
379409
container_lines.extend(
380410
[
@@ -385,6 +415,7 @@ def test_in_container(
385415
]
386416
)
387417

418+
torch_after_pymomentum_script = get_torch_after_pymomentum_script()
388419
import_test_script = get_import_test_script(
389420
wheel_type,
390421
require_cuda,
@@ -395,6 +426,21 @@ def test_in_container(
395426
'echo "Installing wheel..."',
396427
f"$PYTHON -m uv pip install --python /tmp/test_venv/bin/python /wheel/{wheel_file.name}",
397428
"",
429+
]
430+
)
431+
if has_torch_runtime(wheel_type):
432+
container_lines.extend(
433+
[
434+
'echo "Testing pymomentum.geometry before torch..."',
435+
"cat >/tmp/test_torch_after_pymomentum.py <<'PYTEST'",
436+
*torch_after_pymomentum_script.rstrip().splitlines(),
437+
"PYTEST",
438+
"/tmp/test_venv/bin/python -I /tmp/test_torch_after_pymomentum.py",
439+
"",
440+
]
441+
)
442+
container_lines.extend(
443+
[
398444
'echo "Testing imports..."',
399445
"cat >/tmp/test_wheel_imports.py <<'PYTEST'",
400446
*import_test_script.rstrip().splitlines(),
@@ -465,7 +511,7 @@ def test_locally_with_uv(
465511
print(f"[FAIL] Failed to create venv: {result.stderr}", file=sys.stderr)
466512
return result.returncode
467513

468-
if has_torch_extensions(wheel_type):
514+
if has_torch_runtime(wheel_type):
469515
torch_index = get_torch_index(wheel_type)
470516
print(f"Installing torch from {torch_index}...")
471517
result = subprocess.run(

0 commit comments

Comments
 (0)