This is an automated email from the ASF dual-hosted git repository. ryankert01 pushed a commit to branch qdp-unify-gpu-detection in repository https://gitbox.apache.org/repos/asf/mahout.git
commit 8306a5adda1b1afdd272efebdbce60c43532af57 Author: Hsien-Cheng Huang <[email protected]> AuthorDate: Fri Jun 26 19:09:50 2026 +0000 feat(qdp): expose native cuda_available() and gate GPU tests on it Follow-up to #1321 (addresses #1414). After #1321, `_qdp` can build and import without the CUDA toolkit (stub runtime), so "extension importable" no longer implies a usable GPU. The test suite gated GPU tests on `torch.cuda.is_available()`, a proxy that is wrong on this PR's headline scenario -- a GPU host with PyTorch but no toolkit, where `_qdp` is a stub build yet torch still reports a device. - qdp-core: add `cuda_runtime_available()`, which queries `cudaGetDeviceCount` (false in a stub build via the existing 999 sentinel stub, and on hosts with no device). Re-exported from the crate root. - _qdp: expose it as `_qdp.cuda_available()`. - qumat_qdp: add `is_cuda_available()`, mirroring `is_triton_amd_available()`, as the single Python source of truth. - testing/conftest: gate the `@pytest.mark.gpu` auto-skip on the native signal (falling back to torch only if the helper is absent). - test_fallback: coverage that runs on a stub build too, guarding that querying availability returns a bool without aborting. Verified on GPU (tests run and pass) and with CUDA hidden (tests skip); full Rust suite, clippy --all-features, ruff, and ty all clean. --- qdp/qdp-core/src/gpu/cuda_ffi.rs | 24 ++++++++++++++++++++++++ qdp/qdp-core/src/lib.rs | 1 + qdp/qdp-python/qumat_qdp/__init__.py | 22 ++++++++++++++++++++++ qdp/qdp-python/src/lib.rs | 13 +++++++++++++ testing/conftest.py | 13 +++++++++++-- testing/qdp_python/test_fallback.py | 23 +++++++++++++++++++++++ 6 files changed, 94 insertions(+), 2 deletions(-) diff --git a/qdp/qdp-core/src/gpu/cuda_ffi.rs b/qdp/qdp-core/src/gpu/cuda_ffi.rs index dc8930584..dbbe844db 100644 --- a/qdp/qdp-core/src/gpu/cuda_ffi.rs +++ b/qdp/qdp-core/src/gpu/cuda_ffi.rs @@ -69,6 +69,10 @@ unsafe extern "C" { pub(crate) fn cudaMemGetInfo(free: *mut usize, total: *mut usize) -> i32; + /// Number of CUDA-capable devices. Used to detect whether a GPU is + /// actually present at runtime (see `cuda_runtime_available`). + pub(crate) fn cudaGetDeviceCount(count: *mut i32) -> i32; + pub(crate) fn cudaMemcpyAsync( dst: *mut c_void, src: *const c_void, @@ -153,6 +157,10 @@ mod no_cuda_stubs { QDP_CUDA_UNAVAILABLE } + pub(crate) unsafe fn cudaGetDeviceCount(_count: *mut i32) -> i32 { + QDP_CUDA_UNAVAILABLE + } + pub(crate) unsafe fn cudaMemcpyAsync( _dst: *mut c_void, _src: *const c_void, @@ -216,3 +224,19 @@ mod no_cuda_stubs { #[cfg(qdp_no_cuda)] pub(crate) use no_cuda_stubs::*; + +/// Returns `true` if a usable CUDA device is present at runtime. +/// +/// This is deliberately distinct from "the extension is importable": a +/// no-toolkit (`qdp_no_cuda`) build links CUDA stubs that return a non-zero +/// sentinel, and a host with the toolkit but no GPU reports zero devices. Both +/// yield `false` here, so callers (and the Python test suite) can gate GPU work +/// on an accurate signal rather than mere importability. +pub fn cuda_runtime_available() -> bool { + let mut count: i32 = 0; + // SAFETY: `cudaGetDeviceCount` writes a single `i32` through `count` and + // accesses no other memory. In a `qdp_no_cuda` build this resolves to the + // stub above, which returns the unavailable sentinel without dereferencing. + let rc = unsafe { cudaGetDeviceCount(&mut count) }; + rc == CUDA_SUCCESS && count > 0 +} diff --git a/qdp/qdp-core/src/lib.rs b/qdp/qdp-core/src/lib.rs index 822fba96e..95259793c 100644 --- a/qdp/qdp-core/src/lib.rs +++ b/qdp/qdp-core/src/lib.rs @@ -36,6 +36,7 @@ pub mod types; mod profiling; pub use error::{MahoutError, Result, cuda_error_to_string}; +pub use gpu::cuda_ffi::cuda_runtime_available; pub use gpu::memory::Precision; pub use reader::{FloatElem, NullHandling, handle_float32_nulls, handle_float64_nulls}; pub use types::{Dtype, Encoding}; diff --git a/qdp/qdp-python/qumat_qdp/__init__.py b/qdp/qdp-python/qumat_qdp/__init__.py index 2b7566141..81ffd6634 100644 --- a/qdp/qdp-python/qumat_qdp/__init__.py +++ b/qdp/qdp-python/qumat_qdp/__init__.py @@ -44,6 +44,27 @@ else: BACKEND = get_backend() + +def is_cuda_available() -> bool: + """Return whether a usable CUDA device is available to the native engine. + + Unlike the importability of the ``_qdp`` extension -- which only means it + was built, possibly against CUDA stubs on a host without the toolkit -- this + reflects whether GPU work can actually run: ``False`` for a stub build or a + host with no CUDA device, ``True`` only when the native runtime reports at + least one device. + """ + if _qdp_mod is None: + return False + probe = getattr(_qdp_mod, "cuda_available", None) + if probe is None: + return False + try: + return bool(probe()) + except Exception: + return False + + from qumat_qdp.api import ( LatencyResult, QdpBenchmark, @@ -68,6 +89,7 @@ __all__ = [ "ThroughputResult", "TritonAmdEngine", "force_backend", + "is_cuda_available", "is_triton_amd_available", "run_throughput_pipeline_py", ] diff --git a/qdp/qdp-python/src/lib.rs b/qdp/qdp-python/src/lib.rs index 04d772a90..0bfd77189 100644 --- a/qdp/qdp-python/src/lib.rs +++ b/qdp/qdp-python/src/lib.rs @@ -68,6 +68,18 @@ fn run_throughput_pipeline_py( )) } +/// Returns ``True`` if a usable CUDA device is available to the native engine. +/// +/// This reflects whether GPU work can actually run -- it is ``False`` for a +/// stub build (the extension built without the CUDA toolkit) or a host with no +/// CUDA device, and ``True`` only when the native runtime reports a device. +/// Importability of this module only means it was built, which a stub build +/// makes possible without a GPU. +#[pyfunction] +fn cuda_available() -> bool { + qdp_core::cuda_runtime_available() +} + /// Quantum Data Plane (QDP) Python module /// /// GPU-accelerated quantum data encoding with DLPack integration. @@ -78,6 +90,7 @@ fn _qdp(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::<QdpEngine>()?; m.add_class::<QuantumTensor>()?; + m.add_function(wrap_pyfunction!(cuda_available, m)?)?; #[cfg(target_os = "linux")] m.add_class::<PyQuantumLoader>()?; #[cfg(target_os = "linux")] diff --git a/testing/conftest.py b/testing/conftest.py index 95009cc3c..010b7230a 100644 --- a/testing/conftest.py +++ b/testing/conftest.py @@ -47,9 +47,18 @@ def _gpu_available() -> bool: toolkit -- it links stub CUDA Runtime symbols (see qdp-core ``build.rs`` / ``cuda_ffi.rs``). So importing ``_qdp`` no longer implies a working GPU, and ``@pytest.mark.gpu`` tests must additionally check for a real device. - This mirrors the ``torch.cuda.is_available()`` guard used by the inline - skips throughout the QDP test modules. + + Prefers the native engine's own signal (``qumat_qdp.is_cuda_available()``), + which is correct even for a stub build: a host with a GPU + PyTorch but no + CUDA toolkit builds ``_qdp`` against stubs, where ``torch.cuda.is_available()`` + would wrongly report True. Falls back to torch when the helper is absent. """ + try: + from qumat_qdp import is_cuda_available + + return bool(is_cuda_available()) + except Exception: + pass try: import torch diff --git a/testing/qdp_python/test_fallback.py b/testing/qdp_python/test_fallback.py index 9b34ff12b..ac9e7739d 100644 --- a/testing/qdp_python/test_fallback.py +++ b/testing/qdp_python/test_fallback.py @@ -85,6 +85,29 @@ class TestBackendDetection: t = get_torch() assert t is not None # torch is available in test env + def test_is_cuda_available_returns_bool(self): + """is_cuda_available() returns a plain bool without crashing. + + On a stub build (the extension built without the CUDA toolkit) or a host + with no device it must be False -- and querying it must NOT abort the + process, which is the regression this guards (a stub build previously + aborted when GPU code ran). + """ + from qumat_qdp import is_cuda_available + + assert isinstance(is_cuda_available(), bool) + + def test_is_cuda_available_matches_extension(self): + """The helper mirrors the native ``_qdp.cuda_available()`` signal.""" + from qumat_qdp import is_cuda_available + + try: + import _qdp + except ImportError: + assert is_cuda_available() is False + else: + assert is_cuda_available() == bool(_qdp.cuda_available()) + # --------------------------------------------------------------------------- # Loader with explicit PyTorch backend
