diff --git a/pkgs/development/python-modules/mmcv/default.nix b/pkgs/development/python-modules/mmcv/default.nix new file mode 100644 index 000000000000..40efd7589f21 --- /dev/null +++ b/pkgs/development/python-modules/mmcv/default.nix @@ -0,0 +1,122 @@ +{ lib +, buildPythonPackage +, fetchFromGitHub +, pytestCheckHook +, pythonOlder +, torch +, torchvision +, opencv4 +, yapf +, packaging +, pillow +, addict +, ninja +, which +, onnx +, onnxruntime +, scipy +, pyturbojpeg +, tifffile +, lmdb +, symlinkJoin +}: + +let + inherit (torch) cudaCapabilities cudaPackages cudaSupport; + inherit (cudaPackages) backendStdenv cudaVersion; + + cuda-common-redist = with cudaPackages; [ + cuda_cccl # + libcublas # cublas_v2.h + libcusolver # cusolverDn.h + libcusparse # cusparse.h + ]; + + cuda-native-redist = symlinkJoin { + name = "cuda-native-redist-${cudaVersion}"; + paths = with cudaPackages; [ + cuda_cudart # cuda_runtime.h + cuda_nvcc + ] ++ cuda-common-redist; + }; + + cuda-redist = symlinkJoin { + name = "cuda-redist-${cudaVersion}"; + paths = cuda-common-redist; + }; + +in +buildPythonPackage rec { + pname = "mmcv"; + version = "1.7.1"; + format = "setuptools"; + + disabled = pythonOlder "3.6"; + + src = fetchFromGitHub { + owner = "open-mmlab"; + repo = pname; + rev = "refs/tags/v${version}"; + hash = "sha256-b4MLBPNRCcPq1osUvqo71PCWVX7lOjAH+dXttd4ZapU"; + }; + + preConfigure = '' + export MMCV_WITH_OPS=1 + '' + lib.optionalString cudaSupport '' + export CC=${backendStdenv.cc}/bin/cc + export CXX=${backendStdenv.cc}/bin/c++ + export TORCH_CUDA_ARCH_LIST="${lib.concatStringsSep ";" cudaCapabilities}" + export FORCE_CUDA=1 + ''; + + postPatch = '' + substituteInPlace setup.py --replace "cpu_use = 4" "cpu_use = $NIX_BUILD_CORES" + ''; + + preCheck = '' + # remove the conflicting source directory + rm -rf mmcv + ''; + + # test_cnn test_ops really requires gpus to be useful. + # some of the tests take exceedingly long time. + # the rest of the tests are disabled due to sandbox env. + disabledTests = [ + "test_cnn" + "test_ops" + "test_fileclient" + "test_load_model_zoo" + "test_processing" + "test_checkpoint" + "test_hub" + "test_reader" + ]; + + nativeBuildInputs = [ ninja which ] + ++ lib.optionals cudaSupport [ cuda-native-redist ]; + + buildInputs = [ torch ] ++ lib.optionals cudaSupport [ cuda-redist ]; + + nativeCheckInputs = [ pytestCheckHook torchvision lmdb onnx onnxruntime scipy pyturbojpeg tifffile ]; + + propagatedBuildInputs = [ + torch + opencv4 + yapf + packaging + pillow + addict + ]; + + pythonImportsCheck = [ + "mmcv" + ]; + + meta = with lib; { + description = "A Foundational Library for Computer Vision Research"; + homepage = "https://github.com/open-mmlab/mmcv"; + changelog = "https://github.com/open-mmlab/mmcv/releases/tag/v${version}"; + license = with licenses; [ asl20 ]; + maintainers = with maintainers; [ rxiao ]; + }; +} diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 1085080f0cb6..c334f4f89ce3 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -6078,6 +6078,8 @@ self: super: with self; { enablePython = true; }); + mmcv = callPackage ../development/python-modules/mmcv { }; + mmh3 = callPackage ../development/python-modules/mmh3 { }; mmpython = callPackage ../development/python-modules/mmpython { };