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