diff --git a/pkgs/development/python-modules/fairscale/default.nix b/pkgs/development/python-modules/fairscale/default.nix new file mode 100644 index 000000000000..eff005bc5cae --- /dev/null +++ b/pkgs/development/python-modules/fairscale/default.nix @@ -0,0 +1,65 @@ +{ lib +, buildPythonPackage +, fetchFromGitHub +, pythonOlder +# build inputs +, torch +, numpy +, ninja +# check inputs +, pytestCheckHook +, parameterized +, pytest-cov +, pytest-timeout +, remote-pdb +}: +let + pname = "fairscale"; + version = "0.4.13"; +in +buildPythonPackage { + inherit pname version; + format = "setuptools"; + + disabled = pythonOlder "3.10"; + + src = fetchFromGitHub { + owner = "facebookresearch"; + repo = "fairscale"; + rev = "refs/tags/v${version}"; + hash = "sha256-L2Rl/qL6l0OLAofygzJBGQdp/2ZrgDFarwZRjyAR3dw="; + }; + + nativeBuildInputs = [ ninja ]; + dontUseNinjaBuild = true; + dontUseNinjaInstall = true; + dontUseNinjaCheck = true; + + propagatedBuildInputs = [ + torch + numpy + ]; + + nativeCheckInputs = [ + pytestCheckHook + parameterized + pytest-cov + pytest-timeout + remote-pdb + ]; + + # Some tests try to build distributed models, which doesn't work in the sandbox. + doCheck = false; + + pythonImportsCheck = [ + "fairscale" + ]; + + meta = with lib; { + description = "PyTorch extensions for high performance and large scale training"; + homepage = "https://github.com/facebookresearch/fairscale"; + changelog = "https://github.com/facebookresearch/fairscale/releases/tag/v${version}"; + license = with licenses; [ mit asl20 bsd3 ]; + maintainers = with maintainers; [ happysalada ]; + }; +} diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 082a8a8ec3b9..665ac989e32f 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -3519,6 +3519,8 @@ self: super: with self; { factory_boy = callPackage ../development/python-modules/factory_boy { }; + fairscale = callPackage ../development/python-modules/fairscale { }; + faiss = toPythonModule (pkgs.faiss.override { pythonSupport = true; pythonPackages = self;