From f250591fb9a03725a496456dc90a61bfbcfc6795 Mon Sep 17 00:00:00 2001 From: Ahmed Eldeeb <62363199+deeb01@users.noreply.github.com> Date: Wed, 26 Aug 2026 01:29:06 +0300 Subject: [PATCH] Load triton before TensorFlow in ot.backend (closes #816) triton and TensorFlow both ship a statically linked LLVM. Loading triton's libtriton.so into a process that has already imported TensorFlow segfaults inside dlopen, during the static initializers of a 461 MB shared object. torch imports triton lazily rather than at import time: torch.optim's Optimizer.__init__ calls add_param_group, which is decorated with @torch._disable_dynamo, whose wrapper imports torch._dynamo on first call, which calls torch.utils._triton.has_triton_package(), which imports triton. So the first construction of any torch optimizer loads libtriton, and in a process that imported ot.backend that happens after TensorFlow is resident. This is why the crash only appeared in the jobs that install every backend into one process, and why the individual tests all passed on their own. It is not specific to torch 2.12; that was simply when torch started shipping a triton version that collides. ot.backend already imports torch before TensorFlow, so probing for triton at the torch import site loads it while the process is still clean. The probe is guarded on TensorFlow actually being importable, so installations without TensorFlow do not pay the cost of loading libtriton. Reproduced on Linux with torch 2.13.0, triton 3.7.1 and TensorFlow 2.21.0: the crash is a SIGSEGV at triton/knobs.py:15 inside create_module, matching the traceback in the issue. Verified that importing triton after TensorFlow still crashes, so it is the ordering and not the import itself that matters. Adds a non-regression test that builds a torch optimizer after importing ot.backend in a subprocess, so the segfault is reported as a test failure instead of killing the test runner. Removes the torch<2.12 pin from .github/requirements_doctests.txt and docs/requirements.txt. --- .github/requirements_doctests.txt | 2 +- RELEASES.md | 6 ++++++ docs/requirements.txt | 2 +- ot/backend.py | 16 ++++++++++++++++ test/test_backend.py | 28 ++++++++++++++++++++++++++++ 5 files changed, 52 insertions(+), 2 deletions(-) diff --git a/.github/requirements_doctests.txt b/.github/requirements_doctests.txt index cae6a079d..7c1aa4990 100644 --- a/.github/requirements_doctests.txt +++ b/.github/requirements_doctests.txt @@ -5,7 +5,7 @@ autograd pymanopt cvxopt scikit-learn -torch<2.12 +torch jax jaxlib tensorflow; python_version < '3.14' diff --git a/RELEASES.md b/RELEASES.md index f91b42bc0..23a5490eb 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -1,5 +1,11 @@ # Releases +## 0.9.8dev + +#### Closed issues + +- Load triton before TensorFlow in `ot.backend` so that building a torch optimizer no longer segfaults the interpreter, and remove the `torch<2.12` pin from the doctest and documentation requirements (PR #839, Issue #816) + ## 0.9.7.post1 This release is identical to 0.9.7 but will allow the upload of a source distribution to PyPI and release on conda-forge (that requires a source distribution). diff --git a/docs/requirements.txt b/docs/requirements.txt index beb66b6bc..021e833ca 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -5,7 +5,7 @@ autograd pymanopt cvxopt scikit-learn -torch<2.12 +torch pytest torch_geometric cvxpy diff --git a/ot/backend.py b/ot/backend.py index 7749c948d..af622734d 100644 --- a/ot/backend.py +++ b/ot/backend.py @@ -86,6 +86,7 @@ # # License: MIT License +import importlib.util import os import time import warnings @@ -108,6 +109,21 @@ import torch torch_type = torch.Tensor + + # Load triton before TensorFlow is imported below. Both triton and + # TensorFlow ship a statically linked LLVM, and loading triton's + # libtriton.so into a process that has already imported TensorFlow + # segfaults inside dlopen. torch only imports triton lazily, on the + # first use of a feature that needs it (constructing an optimizer is + # enough), which would otherwise happen after TensorFlow is loaded. + # See https://github.com/PythonOT/POT/issues/816 + if not os.environ.get(DISABLE_TF_KEY, False) and ( + importlib.util.find_spec("tensorflow") is not None + ): + try: + import triton # noqa: F401 + except ImportError: + pass except ImportError: torch = False torch_type = float diff --git a/test/test_backend.py b/test/test_backend.py index fe6af9c67..4df918140 100644 --- a/test/test_backend.py +++ b/test/test_backend.py @@ -6,6 +6,10 @@ # # License: MIT License +import importlib.util +import subprocess +import sys + import numpy as np import pytest from numpy.testing import assert_array_almost_equal_nulp @@ -914,3 +918,27 @@ def test_get_backend_none(): assert str(nx) == "numpy" with pytest.raises(ValueError): get_backend(None, None) + + +@pytest.mark.skipif( + not torch or not tf or importlib.util.find_spec("triton") is None, + reason="Requires torch, tensorflow and triton installed together", +) +def test_torch_optimizer_after_tensorflow_import(): + """Non-regression test for issue #816. + + Building a torch optimizer makes torch import triton lazily. If TensorFlow + was imported first, loading libtriton.so segfaults the interpreter, so this + has to run in a subprocess. + """ + code = ( + "import ot.backend\n" + "import torch\n" + "x = torch.zeros(3, requires_grad=True)\n" + "torch.optim.SGD([x], lr=0.1)\n" + ) + result = subprocess.run([sys.executable, "-c", code], capture_output=True) + assert result.returncode == 0, ( + f"interpreter died with returncode {result.returncode}: " + f"{result.stderr.decode(errors='replace')[-2000:]}" + )