pkgs/development/python-modules/pytorch-pfn-extras/default.nix

   1 {
   2   lib,
   3   stdenv,
   4   buildPythonPackage,
   5   fetchFromGitHub,
   6
   7   # build-system
   8   setuptools,
   9
  10   # dependencies
  11   numpy,
  12   packaging,
  13   torch,
  14   typing-extensions,
  15
  16   # tests
  17   onnx,
  18   pytestCheckHook,
  19   torchvision,
  20 }:
  21
  22 buildPythonPackage rec {
  23   pname = "pytorch-pfn-extras";
  24   version = "0.8.1";
  25   pyproject = true;
  26
  27   src = fetchFromGitHub {
  28     owner = "pfnet";
  29     repo = "pytorch-pfn-extras";
  30     rev = "refs/tags/v${version}";
  31     hash = "sha256-6KHVsUHN2KDKAaMdhBpZgTq0XILWUsHJPgeRD0m9m20=";
  32   };
  33
  34   build-system = [ setuptools ];
  35
  36   dependencies = [
  37     numpy
  38     packaging
  39     torch
  40     typing-extensions
  41   ];
  42
  43   nativeCheckInputs = [
  44     onnx
  45     pytestCheckHook
  46     torchvision
  47   ];
  48
  49   pytestFlagsArray = [
  50     # Requires CUDA access which is not possible in the nix environment.
  51     "-m 'not gpu and not mpi'"
  52     "-Wignore::DeprecationWarning"
  53   ];
  54
  55   pythonImportsCheck = [ "pytorch_pfn_extras" ];
  56
  57   disabledTests =
  58     [
  59       # AssertionError: assert 4 == 0
  60       # where 4 = <MagicMock id='140733587469184'>.call_count
  61       "test_lr_scheduler_wait_for_first_optimizer_step"
  62     ]
  63     ++ lib.optionals (stdenv.hostPlatform.isDarwin) [
  64       # torch.distributed is not available on darwin
  65       "test_create_distributed_evaluator"
  66       "test_distributed_evaluation"
  67       "test_distributed_evaluator_progress_bar"
  68     ];
  69
  70   disabledTestPaths =
  71     [
  72       # Requires optuna which is currently (2022-02-16) marked as broken.
  73       "tests/pytorch_pfn_extras_tests/test_config_types.py"
  74
  75       # requires onnxruntime which was removed because of poor maintainability
  76       # See https://github.com/NixOS/nixpkgs/pull/105951 https://github.com/NixOS/nixpkgs/pull/155058
  77       "tests/pytorch_pfn_extras_tests/onnx_tests/test_annotate.py"
  78       "tests/pytorch_pfn_extras_tests/onnx_tests/test_as_output.py"
  79       "tests/pytorch_pfn_extras_tests/onnx_tests/test_export.py"
  80       "tests/pytorch_pfn_extras_tests/onnx_tests/test_export_testcase.py"
  81       "tests/pytorch_pfn_extras_tests/onnx_tests/test_lax.py"
  82       "tests/pytorch_pfn_extras_tests/onnx_tests/test_load_model.py"
  83       "tests/pytorch_pfn_extras_tests/onnx_tests/test_torchvision.py"
  84       "tests/pytorch_pfn_extras_tests/onnx_tests/utils.py"
  85
  86       # RuntimeError: No Op registered for Gradient with domain_version of 9
  87       "tests/pytorch_pfn_extras_tests/onnx_tests/test_grad.py"
  88
  89       # torch._dynamo.exc.BackendCompilerFailed: backend='compiler_fn' raised:
  90       # AttributeError: module 'torch.fx.experimental.proxy_tensor' has no attribute 'maybe_disable_fake_tensor_mode'
  91       "tests/pytorch_pfn_extras_tests/dynamo_tests/test_compile.py"
  92     ]
  93     ++ lib.optionals (stdenv.hostPlatform.isDarwin) [
  94       # torch.distributed is not available on darwin
  95       "tests/pytorch_pfn_extras_tests/distributed_tests/test_distributed_validation_sampler.py"
  96       "tests/pytorch_pfn_extras_tests/nn_tests/parallel_tests/test_distributed.py"
  97       "tests/pytorch_pfn_extras_tests/profiler_tests/test_record.py"
  98       "tests/pytorch_pfn_extras_tests/profiler_tests/test_time_summary.py"
  99       "tests/pytorch_pfn_extras_tests/training_tests/extensions_tests/test_accumulate.py"
 100       "tests/pytorch_pfn_extras_tests/training_tests/extensions_tests/test_sharded_snapshot.py"
 101     ]
 102     ++ lib.optionals (stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isAarch64) [
 103       # RuntimeError: internal error
 104       # convolution (e.g. F.conv3d) causes runtime error
 105       "tests/pytorch_pfn_extras_tests/nn_tests/modules_tests/test_lazy_conv.py"
 106     ];
 107
 108   meta = {
 109     description = "Supplementary components to accelerate research and development in PyTorch";
 110     homepage = "https://github.com/pfnet/pytorch-pfn-extras";
 111     changelog = "https://github.com/pfnet/pytorch-pfn-extras/releases/tag/v${version}";
 112     license = lib.licenses.mit;
 113     maintainers = with lib.maintainers; [ samuela ];
 114   };
 115 }