pkgs/development/python-modules/fairscale/default.nix

   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,
  15   pytest-timeout,
  16   remote-pdb,
  17 }:
  18 let
  19   pname = "fairscale";
  20   version = "0.4.13";
  21 in
  22 buildPythonPackage {
  23   inherit pname version;
  24   format = "pyproject";
  25
  26   disabled = pythonOlder "3.10";
  27
  28   src = fetchFromGitHub {
  29     owner = "facebookresearch";
  30     repo = "fairscale";
  31     rev = "refs/tags/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
  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 }