Merge pull request #330555 from SomeoneSerge/refactor/faiss

faiss: upkeep

pkgs/development/libraries/science/math/faiss/default.nix

@@ -1,26 +1,28 @@
-{ lib
-, config
-, fetchFromGitHub
-, stdenv
-, cmake
-, cudaPackages ? { }
-, cudaSupport ? config.cudaSupport
-, pythonSupport ? true
-, pythonPackages
-, llvmPackages
-, blas
-, swig
-, addOpenGLRunpath
-, optLevel ? let
-    optLevels =
-      lib.optionals stdenv.hostPlatform.avx2Support [ "avx2" ]
-      ++ lib.optionals stdenv.hostPlatform.sse4_1Support [ "sse4" ]
-      ++ [ "generic" ];
-  in
-  # Choose the maximum available optimization level
-  builtins.head optLevels
-, faiss # To run demos in the tests
-, runCommand
+{
+  lib,
+  config,
+  fetchFromGitHub,
+  stdenv,
+  cmake,
+  cudaPackages ? { },
+  cudaSupport ? config.cudaSupport,
+  pythonSupport ? true,
+  pythonPackages,
+  llvmPackages,
+  blas,
+  swig,
+  autoAddDriverRunpath,
+  optLevel ?
+    let
+      optLevels =
+        lib.optionals stdenv.hostPlatform.avx2Support [ "avx2" ]
+        ++ lib.optionals stdenv.hostPlatform.sse4_1Support [ "sse4" ]
+        ++ [ "generic" ];
+    in
+    # Choose the maximum available optimization level
+    builtins.head optLevels,
+  faiss, # To run demos in the tests
+  runCommand,
 }@inputs:
 
 let
@@ -44,7 +46,11 @@ in
 stdenv.mkDerivation {
   inherit pname version;
 
-  outputs = [ "out" "demos" ];
+  outputs = [
+    "out"
+    "demos"
+    "dist"
+  ];
 
   src = fetchFromGitHub {
     owner = "facebookresearch";
@@ -53,50 +59,46 @@ stdenv.mkDerivation {
     hash = "sha256-nS8nhkNGGb2oAJKfr/MIAZjAwMxBGbNd16/CkEtv67I=";
   };
 
+  # Remove the following substituteInPlace when updating
+  # to a release that contains change from PR
+  # https://github.com/facebookresearch/faiss/issues/3239
+  # that fixes building faiss with swig 4.2.x
   postPatch = ''
-    # Remove the following substituteInPlace when updating
-    # to a release that contains change from PR
-    # https://github.com/facebookresearch/faiss/issues/3239
-    # that fixes building faiss with swig 4.2.x
     substituteInPlace faiss/python/swigfaiss.swig \
       --replace-fail '#ifdef SWIGWORDSIZE64' '#if (__SIZEOF_LONG__ == 8)'
   '';
 
-  buildInputs = [
-    blas
-    swig
-  ] ++ lib.optionals pythonSupport [
-    pythonPackages.setuptools
-    pythonPackages.pip
-    pythonPackages.wheel
-  ] ++ lib.optionals stdenv.cc.isClang [
-    llvmPackages.openmp
-  ] ++ lib.optionals cudaSupport cudaComponents;
+  nativeBuildInputs =
+    [ cmake ]
+    ++ lib.optionals cudaSupport [
+      cudaPackages.cuda_nvcc
+      autoAddDriverRunpath
+    ]
+    ++ lib.optionals pythonSupport [
+      pythonPackages.python
+      pythonPackages.setuptools
+      pythonPackages.pip
+      pythonPackages.wheel
+    ];
 
-  propagatedBuildInputs = lib.optionals pythonSupport [
-    pythonPackages.numpy
-    pythonPackages.packaging
-  ];
-
-  nativeBuildInputs = [ cmake ] ++ lib.optionals cudaSupport [
-    cudaPackages.cuda_nvcc
-    addOpenGLRunpath
-  ] ++ lib.optionals pythonSupport [
-    pythonPackages.python
-  ];
-
-  passthru.extra-requires.all = [
-    pythonPackages.numpy
-  ];
-
-  cmakeFlags = [
-    "-DFAISS_ENABLE_GPU=${if cudaSupport then "ON" else "OFF"}"
-    "-DFAISS_ENABLE_PYTHON=${if pythonSupport then "ON" else "OFF"}"
-    "-DFAISS_OPT_LEVEL=${optLevel}"
-  ] ++ lib.optionals cudaSupport [
-    "-DCMAKE_CUDA_ARCHITECTURES=${flags.cmakeCudaArchitecturesString}"
-  ];
+  buildInputs =
+    [
+      blas
+      swig
+    ]
+    ++ lib.optionals pythonSupport [ pythonPackages.numpy ]
+    ++ lib.optionals stdenv.cc.isClang [ llvmPackages.openmp ]
+    ++ lib.optionals cudaSupport cudaComponents;
 
+  cmakeFlags =
+    [
+      (lib.cmakeBool "FAISS_ENABLE_GPU" cudaSupport)
+      (lib.cmakeBool "FAISS_ENABLE_PYTHON" pythonSupport)
+      (lib.cmakeFeature "FAISS_OPT_LEVEL" optLevel)
+    ]
+    ++ lib.optionals cudaSupport [
+      (lib.cmakeFeature "CMAKE_CUDA_ARCHITECTURES" flags.cmakeCudaArchitecturesString)
+    ];
 
   buildFlags =
     [ "faiss" ]
@@ -113,39 +115,29 @@ stdenv.mkDerivation {
      python -m pip wheel --verbose --no-index --no-deps --no-clean --no-build-isolation --wheel-dir dist .)
   '';
 
-  postInstall = ''
-    mkdir -p $demos/bin
-    if [[ "$buildInputs" == *demo_ivfpq_indexing* ]] ; then
-      cp ./demos/demo_ivfpq_indexing $demos/bin/
-    fi
-  '' + lib.optionalString pythonSupport ''
-    mkdir -p $out/${pythonPackages.python.sitePackages}
-    (cd faiss/python && python -m pip install dist/*.whl --no-index --no-warn-script-location --prefix="$out" --no-cache)
-  '';
-
-  postFixup = lib.optionalString (pythonSupport && cudaSupport) ''
-    addOpenGLRunpath $out/${pythonPackages.python.sitePackages}/faiss/*.so
-    addOpenGLRunpath $demos/bin/*
-  '';
-
-  # Need buildPythonPackage for this one
-  # pythonImportsCheck = [
-  #   "faiss"
-  # ];
+  postInstall =
+    ''
+      mkdir -p $demos/bin
+      if [[ "$buildInputs" == *demo_ivfpq_indexing* ]] ; then
+        cp ./demos/demo_ivfpq_indexing $demos/bin/
+      fi
+    ''
+    + lib.optionalString pythonSupport ''
+      mkdir "$dist"
+      cp faiss/python/dist/*.whl "$dist/"
+    '';
 
   passthru = {
     inherit cudaSupport cudaPackages pythonSupport;
 
     tests = {
-      runDemos = runCommand "${pname}-run-demos"
-        { buildInputs = [ faiss.demos ]; }
-        # There are more demos, we run just the one that documentation mentions
-        ''
-          demo_ivfpq_indexing && touch $out
-        '';
-    } // lib.optionalAttrs pythonSupport {
-      pytest = pythonPackages.callPackage ./tests.nix { };
-    };
+      runDemos =
+        runCommand "${pname}-run-demos" { buildInputs = [ faiss.demos ]; }
+          # There are more demos, we run just the one that documentation mentions
+          ''
+            demo_ivfpq_indexing && touch $out
+          '';
+    } // lib.optionalAttrs pythonSupport { pytest = pythonPackages.callPackage ./tests.nix { }; };
   };
 
   meta = with lib; {

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

@@ -0,0 +1,46 @@
+{
+  lib,
+  buildPythonPackage,
+  faiss-build,
+  numpy,
+  packaging,
+  setuptools,
+  pip,
+  wheel,
+}:
+
+buildPythonPackage {
+  inherit (faiss-build) pname version;
+  pyproject = true;
+
+  src = "${lib.getOutput "dist" faiss-build}";
+
+  postPatch = ''
+    mkdir dist
+    mv *.whl dist/
+  '';
+
+  build-system = [
+    setuptools
+    pip
+    wheel
+  ];
+
+  dependencies = [
+    numpy
+    packaging
+  ];
+
+  # E.g. cuda libraries; needed because reference scanning
+  # can't see inside the wheels
+  inherit (faiss-build) buildInputs;
+
+  dontBuild = true;
+
+  pythonImportsCheck = [ "faiss" ];
+
+  meta = lib.pipe (faiss-build.meta or { }) [
+    (lib.flip builtins.removeAttrs [ "mainProgram" ])
+    (m: m // { description = "Bindings for faiss, the similarity search library"; })
+  ];
+}

pkgs/top-level/python-packages.nix

@@ -4161,10 +4161,12 @@ self: super: with self; {
 
   fairseq = callPackage ../development/python-modules/fairseq { };
 
-  faiss = toPythonModule (pkgs.faiss.override {
-    pythonSupport = true;
-    pythonPackages = self;
-  });
+  faiss = callPackage ../development/python-modules/faiss {
+    faiss-build = pkgs.faiss.override {
+      pythonSupport = true;
+      pythonPackages = self;
+    };
+  };
 
   fake-useragent = callPackage ../development/python-modules/fake-useragent { };