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