pyrox.dev
1{
2 lib,
3 buildPythonPackage,
4 pythonOlder,
5 fetchFromGitHub,
6 setuptools,
7 wheel,
8 aiofiles,
9 aiohttp,
10 importlib-metadata,
11 nest-asyncio,
12 psutil,
13 pyyaml,
14 torch,
15 typing-extensions,
16 pytest-asyncio,
17 pytestCheckHook,
18 pythonAtLeast,
19 stdenv,
20}:
21
22buildPythonPackage rec {
23 pname = "torchsnapshot";
24 version = "0.1.0";
25 pyproject = true;
26
27 disabled = pythonOlder "3.7";
28
29 src = fetchFromGitHub {
30 owner = "pytorch";
31 repo = "torchsnapshot";
32 rev = "refs/tags/${version}";
33 hash = "sha256-F8OaxLH8BL6MPNLFv1hBuVmeEdnEQ5w2Qny6by1wP6k=";
34 };
35
36 build-system = [
37 setuptools
38 wheel
39 ];
40
41 dependencies = [
42 aiofiles
43 aiohttp
44 importlib-metadata
45 nest-asyncio
46 psutil
47 pyyaml
48 torch
49 typing-extensions
50 ];
51
52 pythonImportsCheck = [ "torchsnapshot" ];
53
54 nativeCheckInputs = [
55 pytest-asyncio
56 pytestCheckHook
57 ];
58
59 meta = with lib; {
60 description = "A performant, memory-efficient checkpointing library for PyTorch applications, designed with large, complex distributed workloads in mind";
61 homepage = "https://github.com/pytorch/torchsnapshot/";
62 changelog = "https://github.com/pytorch/torchsnapshot/releases/tag/${version}";
63 license = licenses.bsd3;
64 maintainers = with maintainers; [ GaetanLepage ];
65 broken =
66 # https://github.com/pytorch/torchsnapshot/issues/175
67 pythonAtLeast "3.12"
68 # ModuleNotFoundError: No module named 'torch._C._distributed_c10d'; 'torch._C' is not a package
69 || stdenv.isDarwin;
70 };
71}