huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-12 04:44:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Capability as usizes

Browse Source
This commit is contained in:
Daniël de Kok 2024-09-27 11:42:17 +00:00
parent a29636ee0a
commit 3eb68a371d
1 changed files with 7 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
launcher/src/gpu.rs
View File

@ -1,9 +1,9 @@
use std::sync::LazyLock; use std::sync::LazyLock;
pub static COMPUTE_CAPABILITY: LazyLock<Option<(isize, isize)>> = pub static COMPUTE_CAPABILITY: LazyLock<Option<(usize, usize)>> =
LazyLock::new(get_cuda_capability); LazyLock::new(get_cuda_capability);
fn get_cuda_capability() -> Option<(isize, isize)> { fn get_cuda_capability() -> Option<(usize, usize)> {
use pyo3::prelude::*; use pyo3::prelude::*;
let py_get_capability = |py: Python| -> PyResult<(isize, isize)> { let py_get_capability = |py: Python| -> PyResult<(isize, isize)> {
@ -13,7 +13,11 @@ fn get_cuda_capability() -> Option<(isize, isize)> {
}; };
match pyo3::Python::with_gil(py_get_capability) { match pyo3::Python::with_gil(py_get_capability) {
Ok(capability) => Some(capability), Ok((major, minor)) if major < 0 || minor < 0 => {
tracing::warn!("Ignoring negative GPU compute capabilities: {major}.{minor}");
None
}
Ok((major, minor)) => Some((major as usize, minor as usize)),
Err(err) => { Err(err) => {
tracing::warn!("Cannot determine GPU compute capability: {}", err); tracing::warn!("Cannot determine GPU compute capability: {}", err);
None None