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

   1 { lib
   2 , buildPythonPackage
   3 , fetchFromGitHub
   4 , pytestCheckHook
   5 , pytest_xdist
   6 , pythonOlder
   7 , matplotlib
   8 , mock
   9 , pytorch
  10 , pynvml
  11 , scikitlearn
  12 , tqdm
  13 }:
  14
  15 buildPythonPackage rec {
  16   pname = "ignite";
  17   version = "0.4.2";
  18
  19   src = fetchFromGitHub {
  20     owner = "pytorch";
  21     repo = pname;
  22     rev = "v${version}";
  23     sha256 = "00vcmhnp14s54g386izgaxzrdr2nqv3pz9nvpyiwrn33zawr308z";
  24   };
  25
  26   checkInputs = [ pytestCheckHook matplotlib mock pytest_xdist ];
  27   propagatedBuildInputs = [ pytorch scikitlearn tqdm pynvml ];
  28
  29   # runs succesfully in 3.9, however, async isn't correctly closed so it will fail after test suite.
  30   doCheck = pythonOlder "3.9";
  31
  32   # Some packages are not in NixPkgs; other tests try to build distributed
  33   # models, which doesn't work in the sandbox.
  34   # avoid tests which need special packages
  35   pytestFlagsArray = [
  36     "--ignore=tests/ignite/contrib/handlers/test_trains_logger.py"
  37     "--ignore=tests/ignite/metrics/test_dill.py"
  38     "--ignore=tests/ignite/metrics/test_ssim.py"
  39     "tests/"
  40   ];
  41
  42   # disable tests which need specific packages
  43   disabledTests = [
  44     "idist"
  45     "tensorboard"
  46     "mlflow"
  47     "trains"
  48     "visdom"
  49     "test_setup_neptune"
  50     "test_output_handler" # needs mlflow
  51     "test_integration"
  52     "test_pbar" # slight output differences
  53     "test_write_results"
  54     "test_setup_plx"
  55   ];
  56
  57   meta = with lib; {
  58     description = "High-level training library for PyTorch";
  59     homepage = "https://pytorch.org/ignite";
  60     license = licenses.bsd3;
  61     maintainers = [ maintainers.bcdarwin ];
  62   };
  63 }