pyrox.dev
Merge pull request #309779 from GaetanLepage/torchsnapshot
python311Packages.torchsnapshot: init at 0.1.0
authored by Pol Dellaiera and committed by GitHub e9778401 56a37107
+71
pkgs/development/python-modules/torchsnapshot/default.nix
+71
pkgs/development/python-modules/torchsnapshot/default.nix
···
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
+
22
+
buildPythonPackage 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
+
}
+2
pkgs/top-level/python-packages.nix
+2
pkgs/top-level/python-packages.nix
···
15223
15223
openai-triton = self.openai-triton-bin;
15224
15224
};
15225
15225
15226
+
torchsnapshot = callPackage ../development/python-modules/torchsnapshot { };
15227
+
15226
15228
torchWithCuda = self.torch.override {
15227
15229
openai-triton = self.openai-triton-cuda;
15228
15230
cudaSupport = true;