pkgs/development/python-modules/nvidia-ml-py/default.nix

   1 {
   2   lib,
   3   fetchPypi,
   4   buildPythonPackage,
   5   substituteAll,
   6   addDriverRunpath,
   7   setuptools,
   8   cudaPackages,
   9   nvidia-ml-py,
  10 }:
  11
  12 buildPythonPackage 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 }