Fix fragile override tests and restructure backend override logic

- Clear stray env vars in override tests to prevent false failures
- Branch on active backend first to handle both-overrides-set case
- Warn (instead of silently ignoring) the wrong override when both are set
- Reject overrides on unsupported backends (e.g. XPU)
- Add symmetric both-overrides-set tests for CUDA and ROCm

Author: sstamenk

bitsandbytes/cextension.py

@@ -33,32 +33,49 @@ def get_cuda_bnb_library_path(cuda_specs: CUDASpecs) -> Path:
     cuda_override_value = os.environ.get("BNB_CUDA_VERSION")
     rocm_override_value = os.environ.get("BNB_ROCM_VERSION")
 
-    if rocm_override_value:
-        library_name = re.sub(r"rocm\d+", f"rocm{rocm_override_value}", library_name, count=1)
-        if torch.version.cuda:
-            raise RuntimeError(
-                f"BNB_ROCM_VERSION={rocm_override_value} detected for CUDA!\n"
-                "Use BNB_CUDA_VERSION instead: export BNB_CUDA_VERSION=<version>\n"
-                "Clear the variable and retry: unset BNB_ROCM_VERSION\n"
+    if torch.version.hip:
+        if cuda_override_value:
+            if not rocm_override_value:
+                raise RuntimeError(
+                    f"BNB_CUDA_VERSION={cuda_override_value} detected but this is not a CUDA build!\n"
+                    "Use BNB_ROCM_VERSION instead: export BNB_ROCM_VERSION=<version>\n"
+                    "Clear the variable and retry: unset BNB_CUDA_VERSION\n"
+                )
+            logger.warning(
+                f"WARNING: BNB_CUDA_VERSION={cuda_override_value} is set but ignored on this ROCm build. "
+                "Clear the variable: unset BNB_CUDA_VERSION",
             )
-        logger.warning(
-            f"WARNING: BNB_ROCM_VERSION={rocm_override_value} environment variable detected; loading {library_name}.\n"
-            "This can be used to load a bitsandbytes version built with a ROCm version that is different from the PyTorch ROCm version.\n"
-            "If this was unintended clear the variable and retry: unset BNB_ROCM_VERSION\n"
-        )
-    elif cuda_override_value:
-        library_name = re.sub(r"cuda\d+", f"cuda{cuda_override_value}", library_name, count=1)
-        if torch.version.hip:
+        if rocm_override_value:
+            library_name = re.sub(r"rocm\d+", f"rocm{rocm_override_value}", library_name, count=1)
+            logger.warning(
+                f"WARNING: BNB_ROCM_VERSION={rocm_override_value} environment variable detected; loading {library_name}.\n"
+                "This can be used to load a bitsandbytes version built with a ROCm version that is different from the PyTorch ROCm version.\n"
+                "If this was unintended clear the variable and retry: unset BNB_ROCM_VERSION\n",
+            )
+    elif torch.version.cuda:
+        if rocm_override_value:
+            if not cuda_override_value:
+                raise RuntimeError(
+                    f"BNB_ROCM_VERSION={rocm_override_value} detected but this is not a ROCm build!\n"
+                    "Use BNB_CUDA_VERSION instead: export BNB_CUDA_VERSION=<version>\n"
+                    "Clear the variable and retry: unset BNB_ROCM_VERSION\n"
+                )
+            logger.warning(
+                f"WARNING: BNB_ROCM_VERSION={rocm_override_value} is set but ignored on this CUDA build. "
+                "Clear the variable: unset BNB_ROCM_VERSION",
+            )
+        if cuda_override_value:
+            library_name = re.sub(r"cuda\d+", f"cuda{cuda_override_value}", library_name, count=1)
+            logger.warning(
+                f"WARNING: BNB_CUDA_VERSION={cuda_override_value} environment variable detected; loading {library_name}.\n"
+                "This can be used to load a bitsandbytes version built with a CUDA version that is different from the PyTorch CUDA version.\n"
+                "If this was unintended clear the variable and retry: unset BNB_CUDA_VERSION\n",
+            )
+    else:
+        if rocm_override_value or cuda_override_value:
             raise RuntimeError(
-                f"BNB_CUDA_VERSION={cuda_override_value} detected for ROCm!\n"
-                f"Use BNB_ROCM_VERSION instead: export BNB_ROCM_VERSION=<version>\n"
-                f"Clear the variable and retry: unset BNB_CUDA_VERSION\n"
+                "BNB_ROCM_VERSION / BNB_CUDA_VERSION overrides are not supported on this backend.",
             )
-        logger.warning(
-            f"WARNING: BNB_CUDA_VERSION={cuda_override_value} environment variable detected; loading {library_name}.\n"
-            "This can be used to load a bitsandbytes version built with a CUDA version that is different from the PyTorch CUDA version.\n"
-            "If this was unintended clear the variable and retry: unset BNB_CUDA_VERSION\n"
-        )
 
     return PACKAGE_DIR / library_name
 

tests/test_cuda_setup_evaluator.py

@@ -25,20 +25,32 @@ def test_get_cuda_bnb_library_path(monkeypatch, cuda120_spec):
 @pytest.mark.skipif(BNB_BACKEND != "CUDA", reason="this test requires a CUDA backend")
 def test_get_cuda_bnb_library_path_override(monkeypatch, cuda120_spec, caplog):
     """BNB_CUDA_VERSION=110 overrides path selection to the CUDA 11.0 binary."""
+    monkeypatch.delenv("BNB_ROCM_VERSION", raising=False)
     monkeypatch.setenv("BNB_CUDA_VERSION", "110")
     assert get_cuda_bnb_library_path(cuda120_spec).stem == "libbitsandbytes_cuda110"
     assert "BNB_CUDA_VERSION" in caplog.text  # did we get the warning?
 
 
 @pytest.mark.skipif(BNB_BACKEND != "CUDA", reason="this test requires a CUDA backend")
 def test_get_cuda_bnb_library_path_rejects_rocm_override(monkeypatch, cuda120_spec):
-    """BNB_ROCM_VERSION should be rejected on CUDA with a helpful error."""
+    """BNB_ROCM_VERSION alone should be rejected on CUDA with a helpful error."""
     monkeypatch.delenv("BNB_CUDA_VERSION", raising=False)
     monkeypatch.setenv("BNB_ROCM_VERSION", "72")
-    with pytest.raises(RuntimeError, match=r"BNB_ROCM_VERSION.*detected for CUDA"):
+    with pytest.raises(RuntimeError, match=r"BNB_ROCM_VERSION.*not a ROCm build"):
         get_cuda_bnb_library_path(cuda120_spec)
 
 
+@pytest.mark.skipif(BNB_BACKEND != "CUDA", reason="this test requires a CUDA backend")
+def test_get_cuda_bnb_library_path_cuda_override_takes_priority(monkeypatch, cuda120_spec, caplog):
+    """When both overrides are set on CUDA, the CUDA override wins and the ROCm one is warned about."""
+    monkeypatch.setenv("BNB_CUDA_VERSION", "110")
+    monkeypatch.setenv("BNB_ROCM_VERSION", "72")
+    assert get_cuda_bnb_library_path(cuda120_spec).stem == "libbitsandbytes_cuda110"
+    assert "BNB_CUDA_VERSION" in caplog.text
+    assert "BNB_ROCM_VERSION" in caplog.text
+    assert "ignored on this CUDA build" in caplog.text
+
+
 @pytest.fixture
 def rocm70_spec() -> CUDASpecs:
     """Simulates torch+rocm7.0 (bundled ROCm) when the system ROCm is newer."""
@@ -60,15 +72,27 @@ def test_get_rocm_bnb_library_path(monkeypatch, rocm70_spec):
 @pytest.mark.skipif(BNB_BACKEND != "ROCm", reason="this test requires a ROCm backend")
 def test_get_rocm_bnb_library_path_override(monkeypatch, rocm70_spec, caplog):
     """BNB_ROCM_VERSION=72 overrides to load the ROCm 7.2 library instead of 7.0."""
+    monkeypatch.delenv("BNB_CUDA_VERSION", raising=False)
     monkeypatch.setenv("BNB_ROCM_VERSION", "72")
     assert get_cuda_bnb_library_path(rocm70_spec).stem == "libbitsandbytes_rocm72"
     assert "BNB_ROCM_VERSION" in caplog.text
 
 
+@pytest.mark.skipif(BNB_BACKEND != "ROCm", reason="this test requires a ROCm backend")
+def test_get_rocm_bnb_library_path_rocm_override_takes_priority(monkeypatch, rocm70_spec, caplog):
+    """When both overrides are set on ROCm, the ROCm override wins and the CUDA one is warned about."""
+    monkeypatch.setenv("BNB_ROCM_VERSION", "72")
+    monkeypatch.setenv("BNB_CUDA_VERSION", "110")
+    assert get_cuda_bnb_library_path(rocm70_spec).stem == "libbitsandbytes_rocm72"
+    assert "BNB_ROCM_VERSION" in caplog.text
+    assert "BNB_CUDA_VERSION" in caplog.text
+    assert "ignored on this ROCm build" in caplog.text
+
+
 @pytest.mark.skipif(BNB_BACKEND != "ROCm", reason="this test requires a ROCm backend")
 def test_get_rocm_bnb_library_path_rejects_cuda_override(monkeypatch, rocm70_spec):
-    """BNB_CUDA_VERSION should be rejected on ROCm with a helpful error."""
+    """BNB_CUDA_VERSION alone should be rejected on ROCm with a helpful error."""
     monkeypatch.delenv("BNB_ROCM_VERSION", raising=False)
-    monkeypatch.setenv("BNB_CUDA_VERSION", "120")
-    with pytest.raises(RuntimeError, match=r"BNB_CUDA_VERSION.*detected for ROCm"):
+    monkeypatch.setenv("BNB_CUDA_VERSION", "110")
+    with pytest.raises(RuntimeError, match=r"BNB_CUDA_VERSION.*not a CUDA build"):
         get_cuda_bnb_library_path(rocm70_spec)