pkgs/development/python-modules/pytorch-lightning/default.nix

   1 {
   2   lib,
   3   buildPythonPackage,
   4   fetchFromGitHub,
   5
   6   # build-system
   7   setuptools,
   8
   9   # dependencies
  10   fsspec,
  11   lightning-utilities,
  12   numpy,
  13   packaging,
  14   pyyaml,
  15   tensorboardx,
  16   torch,
  17   torchmetrics,
  18   tqdm,
  19   traitlets,
  20
  21   # tests
  22   psutil,
  23   pytestCheckHook,
  24 }:
  25
  26 buildPythonPackage rec {
  27   pname = "pytorch-lightning";
  28   version = "2.5.0.post0";
  29   pyproject = true;
  30
  31   src = fetchFromGitHub {
  32     owner = "Lightning-AI";
  33     repo = "pytorch-lightning";
  34     tag = version;
  35     hash = "sha256-TkwDncyfv1VoV/IErUgF4p0Or5PJbwKoABqo1xXGLVg=";
  36   };
  37
  38   preConfigure = ''
  39     export PACKAGE_NAME=pytorch
  40   '';
  41
  42   build-system = [ setuptools ];
  43
  44   dependencies = [
  45     fsspec
  46     lightning-utilities
  47     numpy
  48     packaging
  49     pyyaml
  50     tensorboardx
  51     torch
  52     torchmetrics
  53     tqdm
  54     traitlets
  55   ] ++ fsspec.optional-dependencies.http;
  56
  57   nativeCheckInputs = [
  58     psutil
  59     pytestCheckHook
  60   ];
  61
  62   # Some packages are not in NixPkgs; other tests try to build distributed
  63   # models, which doesn't work in the sandbox.
  64   doCheck = false;
  65
  66   pythonImportsCheck = [ "pytorch_lightning" ];
  67
  68   meta = {
  69     description = "Lightweight PyTorch wrapper for machine learning researchers";
  70     homepage = "https://github.com/Lightning-AI/pytorch-lightning";
  71     changelog = "https://github.com/Lightning-AI/pytorch-lightning/releases/tag/${src.tag}";
  72     license = lib.licenses.asl20;
  73     maintainers = with lib.maintainers; [ tbenst ];
  74   };
  75 }