Making a development shell for an AI/GPU Accelerated python project with nix flakes, poetry/poetry2nix, and pypi

Went through the pain of packaging a python project on Nixos. Here’s some issues I hit, and how I got lucky resolving them. I feel the most reliable way of doing this in the future is to use docker and just imperatively build.

Here’s how I got web drivers, AI dependencies, gpu dependencies, and an api dependency bundled together into an ephemeral shell for python development, on NixOS 23.11

Enable Flakes
Start with setting up poetry2nix
Get the template flake by running nix flake init --template github:nix-community/poetry2nix
in the flake.nix, sometimes changing projectDir = self to projectDir = ./. fixed some issues
in your terminal, run nix develop . to build the poetry app with python packages described in pyproject.toml
By default, just poetry and python latest should be installed. the dependencies for the project (which gets reflected in the pyproject.toml) are updated with poetry add, such as poetry add numpy selenium scikit-learn
Exit out of the ephemeral shell from nix develop ., and rerun to have poetry2nix rebuild and link the newly declared packages

Poetry2nix has worked pretty well for the more obscure python packages, but failed in others. For example, sentence-transformers would depend on maturin, which would fail to link setuptools. If poetry doesn’t work, you can try and get the package from nixpkgs, or specify sha256s from pypi.org

Here’s an example of what I added to my flake.nix to get gpu acceleration, sentence-transfomers, firefox drivers for selenium, and other packages poetry failed to setup:


<span style="color:#323232;">packages = [ pkgs.poetry pkgs.python311Packages.sentence-transformers pkgs.firefox 
</span><span style="color:#323232;">            pkgs.python311Packages.openai pkgs.python311Packages.yt-dlp pkgs.python311Packages.pyopencl
</span><span style="color:#323232;">];
</span>

was added to this flake.nix, as in,


<span style="color:#323232;">{
</span><span style="color:#323232;">  description = "Application packaged using poetry2nix";
</span><span style="color:#323232;">
</span><span style="color:#323232;">  inputs = {
</span><span style="color:#323232;">    flake-utils.url = "github:numtide/flake-utils";
</span><span style="color:#323232;">    nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
</span><span style="color:#323232;">    poetry2nix = {
</span><span style="color:#323232;">      url = "github:nix-community/poetry2nix";
</span><span style="color:#323232;">      inputs.nixpkgs.follows = "nixpkgs";
</span><span style="color:#323232;">    };
</span><span style="color:#323232;">  };
</span><span style="color:#323232;">  outputs = { self, nixpkgs, flake-utils, poetry2nix }:
</span><span style="color:#323232;">    flake-utils.lib.eachDefaultSystem (system:
</span><span style="color:#323232;">      let
</span><span style="color:#323232;">        # see https://github.com/nix-community/poetry2nix/tree/master#api for more functions and examples.
</span><span style="color:#323232;">        pkgs = nixpkgs.legacyPackages.${system};
</span><span style="color:#323232;">        inherit (poetry2nix.lib.mkPoetry2Nix { inherit pkgs; }) mkPoetryApplication;
</span><span style="color:#323232;">      in
</span><span style="color:#323232;">      {
</span><span style="color:#323232;">        packages = {
</span><span style="color:#323232;">          myapp = mkPoetryApplication {
</span><span style="color:#323232;">            projectDir = ./.;
</span><span style="color:#323232;">          };
</span><span style="color:#323232;">          default = self.packages.${system}.myapp;
</span><span style="color:#323232;">        };
</span><span style="color:#323232;">        devShells.default = pkgs.mkShell {
</span><span style="color:#323232;">          inputsFrom = [ self.packages.${system}.myapp ];
</span><span style="color:#323232;">          packages = [ pkgs.poetry pkgs.python311Packages.sentence-transformers pkgs.firefox 
</span><span style="color:#323232;">            pkgs.python311Packages.openai pkgs.python311Packages.yt-dlp pkgs.python311Packages.pyopencl
</span><span style="color:#323232;">          ];
</span><span style="color:#323232;">          nativeBuildInputs = [(
</span><span style="color:#323232;">            pkgs.python311Packages.buildPythonPackage rec {
</span><span style="color:#323232;">              pname = "serpapi";
</span><span style="color:#323232;">              version = "0.1.5";
</span><span style="color:#323232;">              src = pkgs.python311Packages.fetchPypi {
</span><span style="color:#323232;">                inherit pname version;
</span><span style="color:#323232;">                sha256 = "b9707ed54750fdd2f62dc3a17c6a3fb7fa421dc37902fd65b2263c0ac765a1a5";
</span><span style="color:#323232;">              };
</span><span style="color:#323232;">            }
</span><span style="color:#323232;">          )];
</span><span style="color:#323232;">        };
</span><span style="color:#323232;">      });
</span><span style="color:#323232;">}
</span>

There was one package (serpapi), which was not in nixpkgs, and poetry failed as well. Adding this to native build inputs got serpapi installed


<span style="color:#323232;">nativeBuildInputs = [(
</span><span style="color:#323232;">            pkgs.python311Packages.buildPythonPackage rec {
</span><span style="color:#323232;">              pname = "serpapi";
</span><span style="color:#323232;">              version = "0.1.5";
</span><span style="color:#323232;">              src = pkgs.python311Packages.fetchPypi {
</span><span style="color:#323232;">                inherit pname version;
</span><span style="color:#323232;">                sha256 = "b9707ed54750fdd2f62dc3a17c6a3fb7fa421dc37902fd65b2263c0ac765a1a5";
</span><span style="color:#323232;">              };
</span><span style="color:#323232;">            }
</span><span style="color:#323232;">)];
</span>

All in all, it works, and I have no doubt I’ve made a reproducible environment. What attracts me is I’ve never had an easier time setting up cuda/cudnn/tensorrt/… system drivers have been near effortless, and much faster to setup than on debian. Tools like sentence-transformers and torch default to packages which leverage the GPU.

What pushes me away, is I’ve had failures in each of the three methods for specifying package dependencies, even though one of the three eventually was the fix for integrating the dependencies into my shell. For now, I’ll stick with it, but it’s hard for me to suggest to a team we use this in development

Add comment