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