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