tjh.dev
1{
2 lib,
3 buildPythonPackage,
4 cudaPackages,
5 fetchFromGitHub,
6 substituteAll,
7 pythonOlder,
8 addDriverRunpath,
9 setuptools,
10 pytestCheckHook,
11 versioneer,
12 pynvml,
13}:
14
15buildPythonPackage rec {
16 pname = "pynvml";
17 version = "11.5.0";
18 pyproject = true;
19 disabled = pythonOlder "3.6";
20
21 src = fetchFromGitHub {
22 owner = "gpuopenanalytics";
23 repo = "pynvml";
24 rev = "refs/tags/${version}";
25 hash = "sha256-K3ZENjgi+TVDxr55dRK1y8SwzfgVIzcnD4oEI+KHRa4=";
26 };
27
28 patches = [
29 (substituteAll {
30 src = ./0001-locate-libnvidia-ml.so.1-on-NixOS.patch;
31 inherit (addDriverRunpath) driverLink;
32 })
33 ];
34
35 # unvendor versioneer
36 postPatch = ''
37 rm versioneer.py
38 '';
39
40 nativeBuildInputs = [
41 setuptools
42 versioneer
43 ];
44
45 pythonImportsCheck = [
46 "pynvml"
47 "pynvml.smi"
48 ];
49
50 nativeCheckInputs = [ pytestCheckHook ];
51
52 # OSError: /run/opengl-driver/lib/libnvidia-ml.so.1: cannot open shared object file: No such file or directory
53 doCheck = false;
54
55 passthru.tests.tester-nvmlInit = cudaPackages.writeGpuTestPython { libraries = [ pynvml ]; } ''
56 import pynvml
57 from pynvml.smi import nvidia_smi # noqa: F401
58
59 print(f"{pynvml.nvmlInit()=}")
60 '';
61
62 meta = with lib; {
63 description = "Python bindings for the NVIDIA Management Library";
64 homepage = "https://github.com/gpuopenanalytics/pynvml";
65 license = licenses.bsd3;
66 maintainers = [ maintainers.bcdarwin ];
67 };
68}