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 = "13.0.1";
15 pyproject = true;
16
17 src = fetchFromGitHub {
18 owner = "gpuopenanalytics";
19 repo = "pynvml";
20 tag = version;
21 hash = "sha256-Jwj3cm0l7qR/q5jzwKbD52L7ePYCdzXrYFOceMA776M=";
22 };
23
24 build-system = [
25 setuptools
26 ];
27
28 pythonRelaxDeps = [
29 "nvidia-ml-py"
30 ];
31
32 dependencies = [ nvidia-ml-py ];
33
34 pythonImportsCheck = [
35 "pynvml_utils"
36 ];
37
38 nativeCheckInputs = [ pytestCheckHook ];
39
40 doCheck = false;
41
42 passthru.tests.tester-nvmlInit = cudaPackages.writeGpuTestPython { libraries = [ pynvml ]; } ''
43 from pynvml_utils import nvidia_smi # noqa: F401
44 nvsmi = nvidia_smi.getInstance()
45 print(nvsmi.DeviceQuery('memory.free, memory.total'))
46 '';
47
48 meta = {
49 description = "Unofficial Python bindings for the NVIDIA Management Library";
50 homepage = "https://github.com/gpuopenanalytics/pynvml";
51 changelog = "https://github.com/gpuopenanalytics/pynvml?tab=readme-ov-file#release-notes";
52 license = lib.licenses.bsd3;
53 maintainers = [ lib.maintainers.bcdarwin ];
54 };
55}