Unverified Commit d808a7b1 authored by Someone's avatar Someone Committed by GitHub
Browse files

Merge pull request #330555 from SomeoneSerge/refactor/faiss 

faiss: upkeep
parents 1a69f03f 8b2b3b6a

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

+83 −91
Original line number Diff line number Diff line 
{ lib

, config

, fetchFromGitHub

, stdenv

, cmake

, cudaPackages ? { }

, cudaSupport ? config.cudaSupport

, pythonSupport ? true

, pythonPackages

, llvmPackages

, blas

, swig

, addOpenGLRunpath

, optLevel ? let

{

  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

    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,51 +59,47 @@ stdenv.mkDerivation { 
    hash = "sha256-nS8nhkNGGb2oAJKfr/MIAZjAwMxBGbNd16/CkEtv67I=";

  };



  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

  postPatch = ''

    substituteInPlace faiss/python/swigfaiss.swig \

      --replace-fail '#ifdef SWIGWORDSIZE64' '#if (__SIZEOF_LONG__ == 8)'

  '';



  buildInputs = [

    blas

    swig

  ] ++ lib.optionals pythonSupport [

  nativeBuildInputs =

    [ cmake ]

    ++ lib.optionals cudaSupport [

      cudaPackages.cuda_nvcc

      autoAddDriverRunpath

    ]

    ++ lib.optionals pythonSupport [

      pythonPackages.python

      pythonPackages.setuptools

      pythonPackages.pip

      pythonPackages.wheel

  ] ++ lib.optionals stdenv.cc.isClang [

    llvmPackages.openmp

  ] ++ lib.optionals cudaSupport cudaComponents;



  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" ]

    # This is just a demo app used as a test.
@@ -113,39 +115,29 @@ stdenv.mkDerivation { 
     python -m pip wheel --verbose --no-index --no-deps --no-clean --no-build-isolation --wheel-dir dist .)

  '';



  postInstall = ''

  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/*

    ''

    + lib.optionalString pythonSupport ''

      mkdir "$dist"

      cp faiss/python/dist/*.whl "$dist/"

    '';



  # Need buildPythonPackage for this one

  # pythonImportsCheck = [

  #   "faiss"

  # ];



  passthru = {

    inherit cudaSupport cudaPackages pythonSupport;



    tests = {

      runDemos = runCommand "${pname}-run-demos"

        { buildInputs = [ faiss.demos ]; }

      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 { };

    };

    } // lib.optionalAttrs pythonSupport { pytest = pythonPackages.callPackage ./tests.nix { }; };

  };



  meta = with lib; {

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

0 → 100644
+46 −0
Original line number Diff line number Diff line 
{

  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

+6 −4
Original line number Diff line number Diff line
@@ -4161,10 +4161,12 @@ self: super: with self; { 














  fairseq = callPackage ../development/python-modules/fairseq { };



























  faiss = toPythonModule (pkgs.faiss.override {













  faiss = callPackage ../development/python-modules/faiss {













    faiss-build = pkgs.faiss.override {















      pythonSupport = true;















      pythonPackages = self;













  });













    };













  };





























  fake-useragent = callPackage ../development/python-modules/fake-useragent { };