From 625a419967ce01f56a09bfdf4ad32cdcd9d1445a Mon Sep 17 00:00:00 2001 From: Ahmed Eldeeb <62363199+deeb01@users.noreply.github.com> Date: Fri, 18 Sep 2026 14:28:56 -0700 Subject: [PATCH 1/3] Do not let the triton probe break `import ot` (follow-up to #839) PR #839 added an eager `import triton` to ot.backend so that triton loads before TensorFlow. It caught only ImportError, but a broken native install raises other things, typically OSError for a missing shared object. That exception escaped both the probe and the enclosing torch block, so `import ot` failed outright, and a failure in find_spec("tensorflow") would have been caught by the torch block and silently disabled the torch backend instead. The probe is a best-effort optimisation and must never be fatal, so it now catches Exception and covers the find_spec call as well. Adds a regression test that puts a triton raising OSError on the path and asserts `import ot` still succeeds. --- RELEASES.md | 6 ++++++ ot/backend.py | 15 +++++++++------ test/test_backend.py | 36 ++++++++++++++++++++++++++++++++++++ 3 files changed, 51 insertions(+), 6 deletions(-) diff --git a/RELEASES.md b/RELEASES.md index 063701229..6e9dd3ae1 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -2,6 +2,12 @@ ## 0.9.8dev +#### Closed issues + +- Probing for triton in `ot.backend` no longer propagates errors other than `ImportError`, so a broken native triton install can no longer stop `import ot` or silently disable the torch backend (PR #859, follow-up to PR #839 and Issue #816) + +## 0.9.8dev + #### New features - Use `ot.utils.check_marginal` (and shape-tuple support in `ot.utils.unif`) to fill and validate default marginals consistently across solvers (Gromov, low-rank, stochastic, barycenter, factored) (PR #856) diff --git a/ot/backend.py b/ot/backend.py index fc087495c..b19d892cb 100644 --- a/ot/backend.py +++ b/ot/backend.py @@ -117,13 +117,16 @@ # 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: + # Probing must never be fatal: a broken or partial triton install must + # not stop `import ot`, and must not silently disable the torch backend + # either, so this catches more than ImportError. + try: + if not os.environ.get(DISABLE_TF_KEY, False) and ( + importlib.util.find_spec("tensorflow") is not None + ): import triton # noqa: F401 - except ImportError: - pass + except Exception: # pragma: no cover - depends on the installation + pass except ImportError: torch = False torch_type = float diff --git a/test/test_backend.py b/test/test_backend.py index c88ee5052..e49d8955c 100644 --- a/test/test_backend.py +++ b/test/test_backend.py @@ -7,6 +7,7 @@ # License: MIT License import importlib.util +import os import subprocess import sys @@ -982,3 +983,38 @@ def test_no_cuda_context_for_cpu_only_work(): f"interpreter exited with returncode {result.returncode}: " f"{result.stderr.decode(errors='replace')[-2000:]}" ) + + +@pytest.mark.skipif(not torch, reason="Requires torch") +def test_broken_triton_does_not_break_import(tmp_path): + """Probing for triton must never be fatal (see issue #816). + + ot.backend imports triton eagerly so that it loads before TensorFlow. A + broken native install raises something other than ImportError, which must + not stop `import ot` nor disable the torch backend. + """ + (tmp_path / "triton.py").write_text( + 'raise OSError("libtriton.so: cannot open shared object file")\n' + ) + # a stub is enough: the probe only needs find_spec("tensorflow") to succeed + tf_stub = tmp_path / "tensorflow" + tf_stub.mkdir() + (tf_stub / "__init__.py").write_text("class Tensor:\n pass\n") + (tf_stub / "experimental").mkdir() + (tf_stub / "experimental" / "__init__.py").write_text('raise ImportError("stub")\n') + + env = dict(os.environ) + env["PYTHONPATH"] = str(tmp_path) + os.pathsep + env.get("PYTHONPATH", "") + result = subprocess.run( + [ + sys.executable, + "-c", + "import ot.backend as b\n" + "assert b.torch is not False, 'torch backend was disabled by the probe'\n" + "print('ok')\n", + ], + capture_output=True, + env=env, + ) + assert result.returncode == 0, result.stderr.decode(errors="replace")[-2000:] + assert b"ok" in result.stdout From e887b80d071155df32f798e8ea3e85c9fc893e32 Mon Sep 17 00:00:00 2001 From: Ahmed Eldeeb <62363199+deeb01@users.noreply.github.com> Date: Fri, 18 Sep 2026 14:51:56 -0700 Subject: [PATCH 2/3] Correct PR number in RELEASES.md --- RELEASES.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/RELEASES.md b/RELEASES.md index 6e9dd3ae1..f4a8ccb50 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -4,7 +4,7 @@ #### Closed issues -- Probing for triton in `ot.backend` no longer propagates errors other than `ImportError`, so a broken native triton install can no longer stop `import ot` or silently disable the torch backend (PR #859, follow-up to PR #839 and Issue #816) +- Probing for triton in `ot.backend` no longer propagates errors other than `ImportError`, so a broken native triton install can no longer stop `import ot` or silently disable the torch backend (PR #865, follow-up to PR #839 and Issue #816) ## 0.9.8dev @@ -22,7 +22,7 @@ #### Closed issues - Remove a leftover debug `print` from `ot.utils.projection_sparse_simplex` with `axis=1`, and make the `ot.datasets.make_gauss_hd` docstring a raw string so importing `ot` no longer emits a `SyntaxWarning` (PR #860) -- Fix `ot.dist` ignoring the weights `w` for `metric="cityblock"`, which returned the unweighted distance although the weights are documented for this metric (PR #859) +- Fix `ot.dist` ignoring the weights `w` for `metric="cityblock"`, which returned the unweighted distance although the weights are documented for this metric (PR #865) - Fix swapped arguments to `div_to_product` in `ot.gromov.fused_unbalanced_across_spaces_cost`: with `reg_type="independent"` (UCOOT) the entropic terms used the plan marginals as the reference measures and vice versa (PR #855, Issue #854) - Fix device placement in `ot.batch.bregman_projection_batch` so `ot.solve_batch(..., method="sinkhorn")` no longer crashes on GPU when the torch default device is CPU (PR #851) - Preserve input dtype and device for expected sliced plans, avoid materializing dense distance matrices for sparse plans, and fix weighted sparse-distance ordering (PR #846, Issue #845) From 9628dd6a31a9461111d16a1c30b625ed66e9fa5f Mon Sep 17 00:00:00 2001 From: Ahmed Eldeeb <62363199+deeb01@users.noreply.github.com> Date: Fri, 18 Sep 2026 14:52:26 -0700 Subject: [PATCH 3/3] Fold RELEASES entry into the existing 0.9.8dev section --- RELEASES.md | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/RELEASES.md b/RELEASES.md index f4a8ccb50..10782aa72 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -2,12 +2,6 @@ ## 0.9.8dev -#### Closed issues - -- Probing for triton in `ot.backend` no longer propagates errors other than `ImportError`, so a broken native triton install can no longer stop `import ot` or silently disable the torch backend (PR #865, follow-up to PR #839 and Issue #816) - -## 0.9.8dev - #### New features - Use `ot.utils.check_marginal` (and shape-tuple support in `ot.utils.unif`) to fill and validate default marginals consistently across solvers (Gromov, low-rank, stochastic, barycenter, factored) (PR #856) @@ -21,6 +15,7 @@ #### Closed issues +- Probing for triton in `ot.backend` no longer propagates errors other than `ImportError`, so a broken native triton install can no longer stop `import ot` or silently disable the torch backend (PR #865, follow-up to PR #839 and Issue #816) - Remove a leftover debug `print` from `ot.utils.projection_sparse_simplex` with `axis=1`, and make the `ot.datasets.make_gauss_hd` docstring a raw string so importing `ot` no longer emits a `SyntaxWarning` (PR #860) - Fix `ot.dist` ignoring the weights `w` for `metric="cityblock"`, which returned the unweighted distance although the weights are documented for this metric (PR #865) - Fix swapped arguments to `div_to_product` in `ot.gromov.fused_unbalanced_across_spaces_cost`: with `reg_type="independent"` (UCOOT) the entropic terms used the plan marginals as the reference measures and vice versa (PR #855, Issue #854)