pyrox.dev
1{
2 lib,
3 buildPythonPackage,
4 cudaPackages,
5 fetchFromGitHub,
6 setuptools,
7 pytestCheckHook,
8 nvidia-ml-py,
9 pynvml,
10}:
11
12buildPythonPackage rec {
13 pname = "pynvml";
14 version = "12.0.0";
15 pyproject = true;
16
17 src = fetchFromGitHub {
18 owner = "gpuopenanalytics";
19 repo = "pynvml";
20 tag = version;
21 hash = "sha256-bfES6QqL9DO7rmQ3btXxVzW0KlL0eHyUvYJVeijymBk=";
22 };
23
24 build-system = [
25 setuptools
26 ];
27
28 dependencies = [ nvidia-ml-py ];
29
30 pythonImportsCheck = [
31 "pynvml_utils"
32 ];
33
34 nativeCheckInputs = [ pytestCheckHook ];
35
36 doCheck = false;
37
38 passthru.tests.tester-nvmlInit = cudaPackages.writeGpuTestPython { libraries = [ pynvml ]; } ''
39 from pynvml_utils import nvidia_smi # noqa: F401
40 nvsmi = nvidia_smi.getInstance()
41 print(nvsmi.DeviceQuery('memory.free, memory.total'))
42 '';
43
44 meta = {
45 description = "Unofficial Python bindings for the NVIDIA Management Library";
46 homepage = "https://github.com/gpuopenanalytics/pynvml";
47 changelog = "https://github.com/gpuopenanalytics/pynvml?tab=readme-ov-file#release-notes";
48 license = lib.licenses.bsd3;
49 maintainers = [ lib.maintainers.bcdarwin ];
50 };
51}