diff --git a/RELEASES.md b/RELEASES.md index 063701229..10782aa72 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -15,8 +15,9 @@ #### 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 #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) 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