pyrox.dev
1{
2 lib,
3 buildPythonPackage,
4 fetchFromGitHub,
5 pythonOlder,
6 setuptools,
7 # build inputs
8 torch,
9 numpy,
10 ninja,
11 # check inputs
12 pytestCheckHook,
13 parameterized,
14 pytest-cov-stub,
15 pytest-timeout,
16 remote-pdb,
17}:
18let
19 pname = "fairscale";
20 version = "0.4.13";
21in
22buildPythonPackage {
23 inherit pname version;
24 format = "pyproject";
25
26 disabled = pythonOlder "3.10";
27
28 src = fetchFromGitHub {
29 owner = "facebookresearch";
30 repo = "fairscale";
31 tag = "v${version}";
32 hash = "sha256-L2Rl/qL6l0OLAofygzJBGQdp/2ZrgDFarwZRjyAR3dw=";
33 };
34
35 # setup.py depends on ninja python dependency, but we have the binary in nixpkgs
36 postPatch = ''
37 substituteInPlace setup.py \
38 --replace 'setup_requires=["ninja"]' 'setup_requires=[]'
39 '';
40
41 nativeBuildInputs = [
42 ninja
43 setuptools
44 ];
45
46 propagatedBuildInputs = [
47 torch
48 numpy
49 ];
50
51 nativeCheckInputs = [
52 pytestCheckHook
53 parameterized
54 pytest-cov-stub
55 pytest-timeout
56 remote-pdb
57 ];
58
59 # Some tests try to build distributed models, which doesn't work in the sandbox.
60 doCheck = false;
61
62 pythonImportsCheck = [ "fairscale" ];
63
64 meta = with lib; {
65 description = "PyTorch extensions for high performance and large scale training";
66 mainProgram = "wgit";
67 homepage = "https://github.com/facebookresearch/fairscale";
68 changelog = "https://github.com/facebookresearch/fairscale/releases/tag/v${version}";
69 license = with licenses; [
70 mit
71 asl20
72 bsd3
73 ];
74 maintainers = with maintainers; [ happysalada ];
75 };
76}