tjh.dev
1{
2 lib,
3 fetchPypi,
4 buildPythonPackage,
5 substituteAll,
6 addDriverRunpath,
7 setuptools,
8 cudaPackages,
9 nvidia-ml-py,
10}:
11
12buildPythonPackage rec {
13 pname = "nvidia-ml-py";
14 version = "12.560.30";
15
16 pyproject = true;
17
18 src = fetchPypi {
19 inherit pname version;
20 extension = "tar.gz";
21 hash = "sha256-8CVNx0AGR2gKBy7gJQm/1GECtgvf7KMhV21NSBfn/pc=";
22 };
23
24 patches = [
25 (substituteAll {
26 src = ./0001-locate-libnvidia-ml.so.1-on-NixOS.patch;
27 inherit (addDriverRunpath) driverLink;
28 })
29 ];
30
31 build-system = [
32 setuptools
33 ];
34
35 # no tests
36 doCheck = false;
37
38 pythonImportsCheck = [ "pynvml" ];
39
40 passthru.tests.tester-nvmlInit =
41 cudaPackages.writeGpuTestPython { libraries = [ nvidia-ml-py ]; }
42 ''
43 from pynvml import (
44 nvmlInit,
45 nvmlSystemGetDriverVersion,
46 nvmlDeviceGetCount,
47 nvmlDeviceGetHandleByIndex,
48 nvmlDeviceGetName,
49 )
50
51 nvmlInit()
52 print(f"Driver Version: {nvmlSystemGetDriverVersion()}")
53
54 for i in range(nvmlDeviceGetCount()):
55 handle = nvmlDeviceGetHandleByIndex(i)
56 print(f"Device {i} : {nvmlDeviceGetName(handle)}")
57 '';
58
59 meta = {
60 description = "Python Bindings for the NVIDIA Management Library";
61 homepage = "https://pypi.org/project/nvidia-ml-py";
62 license = lib.licenses.bsd3;
63 maintainers = with lib.maintainers; [ GaetanLepage ];
64 };
65}