Skip to content

Commit a42f992

Browse files
authored
chore: Switch pynvml to nvidia-ml-py (#1650)
The following warning info can be hidden by switching the package `pynvml` to `nvidia-ml-py`. ``` .venv/lib/python3.12/site-packages/torch/cuda/__init__.py:61: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you. import pynvml # type: ignore[import] ``` The Usage remains the same. Ref: https://pypi.org/project/pynvml/ and https://pypi.org/project/nvidia-ml-py/ Relevant reproducible detail: ``` $ uv pip list | grep -E "(torch|flashinfer)" flashinfer-python 0.3.1 torch 2.7.1 torchaudio 2.7.1 torchvision 0.22.1 ```
1 parent 3c3b588 commit a42f992

File tree

2 files changed

+2
-2
lines changed

2 files changed

+2
-2
lines changed

docker/install/install_python_packages.sh

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,6 @@ set -u
2323
CUDA_VERSION=${1:-cu128}
2424

2525
pip3 install torch --index-url https://download.pytorch.org/whl/${CUDA_VERSION}
26-
pip3 install requests ninja pytest numpy scipy build pynvml cuda-python einops nvidia-nvshmem-cu12
26+
pip3 install requests ninja pytest numpy scipy build nvidia-ml-py cuda-python einops nvidia-nvshmem-cu12
2727
pip3 install nvidia-cutlass-dsl
2828
pip3 install 'nvidia-cudnn-frontend>=1.13.0' 'nvidia-cudnn-cu12>=9.11.0.98'

setup.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ def generate_build_meta(aot_build_meta: dict) -> None:
5858
"torch",
5959
"ninja",
6060
"requests",
61-
"pynvml",
61+
"nvidia-ml-py",
6262
"einops",
6363
"click",
6464
"tqdm",

0 commit comments

Comments
 (0)