diff --git a/conda/environments/all_cuda-118_arch-x86_64.yaml b/conda/environments/all_cuda-118_arch-x86_64.yaml index d21497c4def..1ec002d3ec6 100644 --- a/conda/environments/all_cuda-118_arch-x86_64.yaml +++ b/conda/environments/all_cuda-118_arch-x86_64.yaml @@ -71,6 +71,7 @@ dependencies: - ptxcompiler - pyarrow>=14.0.0,<19.0.0a0 - pydata-sphinx-theme!=0.14.2 +- pynvml>=11.4.1,<12.0.0a0 - pytest-benchmark - pytest-cases>=3.8.2 - pytest-cov diff --git a/conda/environments/all_cuda-125_arch-x86_64.yaml b/conda/environments/all_cuda-125_arch-x86_64.yaml index 400c1195e00..b6d1cf75721 100644 --- a/conda/environments/all_cuda-125_arch-x86_64.yaml +++ b/conda/environments/all_cuda-125_arch-x86_64.yaml @@ -69,6 +69,7 @@ dependencies: - pyarrow>=14.0.0,<19.0.0a0 - pydata-sphinx-theme!=0.14.2 - pynvjitlink>=0.0.0a0 +- pynvml>=11.4.1,<12.0.0a0 - pytest-benchmark - pytest-cases>=3.8.2 - pytest-cov diff --git a/conda/recipes/dask-cudf/meta.yaml b/conda/recipes/dask-cudf/meta.yaml index 1e6c0a35a09..74ecded8ead 100644 --- a/conda/recipes/dask-cudf/meta.yaml +++ b/conda/recipes/dask-cudf/meta.yaml @@ -43,6 +43,7 @@ requirements: run: - python - cudf ={{ version }} + - pynvml >=11.4.1,<12.0.0a0 - rapids-dask-dependency ={{ minor_version }} - {{ pin_compatible('cuda-version', max_pin='x', min_pin='x') }} diff --git a/dependencies.yaml b/dependencies.yaml index 682aaa612b4..6c38d1c290a 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -758,6 +758,7 @@ dependencies: common: - output_types: [conda, requirements, pyproject] packages: + - pynvml>=11.4.1,<12.0.0a0 - rapids-dask-dependency==24.12.*,>=0.0.0a0 run_custreamz: common: diff --git a/python/dask_cudf/dask_cudf/io/parquet.py b/python/dask_cudf/dask_cudf/io/parquet.py index bf8fae552c2..bbedd046760 100644 --- a/python/dask_cudf/dask_cudf/io/parquet.py +++ b/python/dask_cudf/dask_cudf/io/parquet.py @@ -55,7 +55,7 @@ def _get_device_size(): handle = pynvml.nvmlDeviceGetHandleByIndex(int(index)) return pynvml.nvmlDeviceGetMemoryInfo(handle).total - except (ImportError, ValueError): + except ValueError: # Fall back to a conservative 8GiB default return 8 * 1024**3 diff --git a/python/dask_cudf/pyproject.toml b/python/dask_cudf/pyproject.toml index 07d9143db36..5dac70cc295 100644 --- a/python/dask_cudf/pyproject.toml +++ b/python/dask_cudf/pyproject.toml @@ -24,6 +24,7 @@ dependencies = [ "fsspec>=0.6.0", "numpy>=1.23,<3.0a0", "pandas>=2.0,<2.2.4dev0", + "pynvml>=11.4.1,<12.0.0a0", "rapids-dask-dependency==24.12.*,>=0.0.0a0", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. classifiers = [