如何使用 Nix 让 torch-geometric 工作？

Question

我正在尝试让 Python 包 torch-geometric 使用 Nix（我在 NixOS 上）工作。目前，我使用mach-nix 尝试设置 Python 环境。但是，困难在于一些依赖项应该从单独的文件服务器（不是 pypi）下载，即https://pytorch-geometric.com/whl/torch-1.8.0+cpu.html。我首先尝试设置一个包含单个 torch-geometric 依赖项的环境：torch-sparse。

目前我有以下shell.nix：

{ pkgs ? import <nixpkgs> {} }:

let
  mach-nix = import (builtins.fetchGit {
    url = "https://github.com/DavHau/mach-nix/";
    ref = "refs/tags/3.3.0";
  }) {
    python = "python38";
  };
  sparse = mach-nix.buildPythonPackage {
    pname = "torch_sparse";
    version = "0.6.9";
    requirements = ''
      torch
      scipy
      pytest
      pytest-cov
      pytest-runner
    '';
    src = builtins.fetchGit {
      url = "https://github.com/rusty1s/pytorch_sparse";
      ref = "refs/tags/0.6.9";
    };
  };
in mach-nix.mkPython {
  requirements = "torch-sparse";
  packagesExtra = [
    sparse
  ];
}

在运行 nix-shell 时失败并显示以下错误消息：

running build_ext
error: [Errno 2] No such file or directory: 'which'
builder for '/nix/store/fs9nrrd2a233xp5d6njy6639yjbxp4g0-python3.8-torch_sparse-0.6.9.drv' failed with exit code 1

我尝试将which 包添加到checkInputs 和buildInputs，但这并不能解决问题。显然，我尝试直接从其 GitHub 存储库构建包，因为我不确定如何在 mach-nix 中引用 wheel 包。我对 NixOS 环境比较陌生，坦率地说，我完全迷路了。

我应该如何安装 Python 包，例如 torch-sparse 或 torch-geometric？我是否使用了正确的工具？

Answer 1

我设法想出了一个有效的 Nix 表达式。我将把答案留在这里以供将来参考。使用nix-shell 运行以下表达式将创建一个带有torch-1.8.0 和torch-geometric-1.7.0 及其所需依赖项的shell。

{ pkgs ? import <nixpkgs> { } }:

let
  python = pkgs.python38;
  pytorch-180 = let
    pyVerNoDot = builtins.replaceStrings [ "." ] [ "" ] python.pythonVersion;
    unsupported = throw "Unsupported system";
    version = "1.8.0";
  in python.pkgs.buildPythonPackage {
    inherit version;

    pname = "pytorch";

    format = "wheel";

    src = pkgs.fetchurl {
      name = "torch-${version}-cp38-cp38-linux_x86_64.whl";
      url =
        "https://download.pytorch.org/whl/cu111/torch-${version}%2Bcu111-cp38-cp38-linux_x86_64.whl";
      hash = "sha256-4NYiAkYfGXm3orLT8Y5diepRMAg+WzJelncy2zJp+Ho=";
    };

    nativeBuildInputs = with pkgs; [ addOpenGLRunpath patchelf ];

    propagatedBuildInputs = with python.pkgs; [
      future
      numpy
      pyyaml
      requests
      typing-extensions
    ];

    postInstall = ''
      # ONNX conversion
      rm -rf $out/bin
    '';

    postFixup = let rpath = pkgs.lib.makeLibraryPath [ pkgs.stdenv.cc.cc.lib ];
    in ''
      find $out/${python.sitePackages}/torch/lib -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
        echo "setting rpath for $lib..."
        patchelf --set-rpath "${rpath}:$out/${python.sitePackages}/torch/lib" "$lib"
        addOpenGLRunpath "$lib"
      done
    '';

    pythonImportsCheck = [ "torch" ];

    meta = with pkgs.lib; {
      description =
        "Open source, prototype-to-production deep learning platform";
      homepage = "https://pytorch.org/";
      changelog = "https://github.com/pytorch/pytorch/releases/tag/v${version}";
      license = licenses.unfree; # Includes CUDA and Intel MKL.
      platforms = platforms.linux;
      maintainers = with maintainers; [ danieldk ];
    };
  };
  sparse = with python.pkgs;
    buildPythonPackage rec {
      pname = "torch_sparse";
      version = "0.6.9";

      src = pkgs.fetchurl {
        name = "${pname}-${version}-cp38-cp38-linux_x86_64.whl";
        url =
          "https://pytorch-geometric.com/whl/torch-1.8.0+cpu/${pname}-${version}-cp38-cp38-linux_x86_64.whl";
        hash = "sha256-6dmZNQ0FlwKdfESKhvv8PPwzgsJFWlP8tYXWu2JLiMk=";
      };

      format = "wheel";

      propagatedBuildInputs = [ pytorch-180 scipy ];
      # buildInputs = [ pybind11 ];
      # nativeBuildInputs = [ pytest-runner pkgs.which ];

      doCheck = false;

      postInstall = ''
        rm -rf $out/${python.sitePackages}/test
      '';
    };
  scatter = with python.pkgs;
    buildPythonPackage rec {
      pname = "torch_scatter";
      version = "2.0.7";

      src = pkgs.fetchurl {
        name = "${pname}-${version}-cp38-cp38-linux_x86_64.whl";
        url =
          "https://pytorch-geometric.com/whl/torch-1.8.0+cpu/${pname}-${version}-cp38-cp38-linux_x86_64.whl";
        hash = "sha256-MRoFretgyEpq+7aJZc0399Kd+f28Uhn5+CxW5ZIKwcg=";
      };

      format = "wheel";

      propagatedBuildInputs = [ pytorch-180 ];

      doCheck = false;

      postInstall = ''
        rm -rf $out/${python.sitePackages}/test
      '';
    };
  cluster = with python.pkgs;
    buildPythonPackage rec {
      pname = "torch_cluster";
      version = "1.5.9";

      src = pkgs.fetchurl {
        name = "${pname}-${version}-cp38-cp38-linux_x86_64.whl";
        url =
          "https://pytorch-geometric.com/whl/torch-1.8.0+cpu/${pname}-${version}-cp38-cp38-linux_x86_64.whl";
        hash = "sha256-E2nywtiZ7m7VA1J7AY7gAHYvyN9H3zl/W0/WsZLzwF8=";
      };

      format = "wheel";

      propagatedBuildInputs = [ pytorch-180 ];

      doCheck = false;

      postInstall = ''
        rm -rf $out/${python.sitePackages}/test
      '';
    };
  spline = with python.pkgs;
    buildPythonPackage rec {
      pname = "torch_spline_conv";
      version = "1.2.1";

      src = pkgs.fetchurl {
        name = "${pname}-${version}-cp38-cp38-linux_x86_64.whl";
        url =
          "https://pytorch-geometric.com/whl/torch-1.8.0+cpu/${pname}-${version}-cp38-cp38-linux_x86_64.whl";
        hash = "sha256-ghSzoxoqSccPAZzfcHJEPYySQ/KYqQ90mFsOdt1CjUw=";
      };

      format = "wheel";

      propagatedBuildInputs = [ pytorch-180 ];

      doCheck = false;

      postInstall = ''
        rm -rf $out/${python.sitePackages}/test
      '';
    };
  python-louvain = with python.pkgs;
    buildPythonPackage rec {
      pname = "python-louvain";
      version = "0.15";

      src = fetchPypi {
        inherit pname version;
        sha256 = "1sqp97fwh4asx0jr72x8hil8z8fcg2xq92jklmh2m599pvgnx19a";
      };

      propagatedBuildInputs = [ numpy networkx ];

      doCheck = false;
    };
  googledrivedownloader = with python.pkgs;
    buildPythonPackage rec {
      pname = "googledrivedownloader";
      version = "0.4";

      src = fetchPypi {
        inherit pname version;
        sha256 = "0172l1f8ys0913wcr16lzx87vsnapppih62qswmvzwrggcrw2d2b";
      };

      doCheck = false;
    };
  geometric = with python.pkgs;
    buildPythonPackage rec {
      pname = "torch_geometric";
      version = "1.7.0";

      src = fetchPypi {
        inherit pname version;
        sha256 = "1a7ym34ynhk5gb3yc5v4qkmkrkyjbv1fgisrsk0c9xay66w7nwz9";
      };

      propagatedBuildInputs = [
        pytorch-180
        numpy
        scipy
        tqdm
        networkx
        scikit-learn
        requests
        pandas
        rdflib
        jinja2
        numba
        ase
        h5py
        python-louvain
        googledrivedownloader
      ];
      nativeBuildInputs = [ pytest-runner ];

      doCheck = false;

      # postInstall = ''
      #   rm -rf $out/${python.sitePackages}/test
      # '';
    };
  python-with-pkgs = python.withPackages
    (ps: with ps; [ pytorch-180 scatter sparse cluster spline geometric ps ]);
in pkgs.mkShell { buildInputs = [ python-with-pkgs ]; }

Answer 2

使用 nix-shell 运行以下表达式将使用 Mach-nix 创建一个带有 torch-1.9.0 和 torch-geometric-1.7.2 及其所需依赖项的 shell。

我已经在带有 nix 的 Mac OS X 上对此进行了测试。

{ pkgs ? import <nixpkgs> {}}:
let
  mach-nix = import (builtins.fetchGit {
    url = "https://github.com/DavHau/mach-nix/";
    ref = "refs/tags/3.3.0";
  }) {
    # optionally bring your own nixpkgs
    #pkgs = import <nixpkgs> {};

    # optionally specify the python version
    python = "python39";

    # optionally update pypi data revision from https://github.com/DavHau/pypi-deps-db
    pypiDataRev = "9e07576970701f9618758f9b42d105a29b1cbf81";
    pypiDataSha256 = "1nsmcy3rbdkzfc4ymhnmcf6jp9hnb1ykzv4n6lvxl86l5bdcyz7f";
  };

  sparse = mach-nix.buildPythonPackage {
    pname = "torch_sparse";
    version = "0.6.10";
    doCheck = false;
    nativeBuildInputs = [ pkgs.which ];
    requirements = ''
      torch
      scipy
      pytest
      pytest-cov
      pytest-runner
    '';
    src = builtins.fetchGit {
      url = "https://github.com/rusty1s/pytorch_sparse";
      ref = "refs/tags/0.6.10";
    };
    postInstall = ''rm -rf $out/lib/python*/site-packages/test '';
  };

  scatter = mach-nix.buildPythonPackage {
    pname = "torch_scatter";
    version = "2.0.7";
    doCheck = false;
    nativeBuildInputs = [ pkgs.which ];
    requirements = ''
      torch
      scipy
      pytest
      pytest-cov
      pytest-runner
    '';
    src = builtins.fetchGit {
      url = "https://github.com/rusty1s/pytorch_scatter";
      ref = "refs/tags/2.0.7";
    };
    postInstall = ''rm -rf $out/lib/python*/site-packages/test '';
  };

in
mach-nix.mkPythonShell {  # replace with mkPythonShell if shell is wanted
  packagesExtra = [sparse scatter];
  requirements = 
    '' 
       torch >= 1.9.0
       torch-geometric == 1.7.2
     '';
  providers =
    {
        _default = "wheel,nixpkgs,sdist";
    };  
}