[MRG] Do not let the triton probe break import ot (follow-up to #839)

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)

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

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