Setting Deep learning Python project with CUDA support using uv2nix

Guides
1
  • So this is a post about how I successfully migrated to NixOS and set it up for development with: Python, PyTorch, CUDA, etc.
  • Since for this scenario I find that documents, projects, and practices are scattered across the internet and many of them are not actually up-to-date, and for someone new to NixOS or Nix, setting things up can be challenging, I’m posting the way how I arrived at this conclusive solution in case someone else in the future might find this helpful.
  • Here are the hardware and software environments I’m working with:
  • I’m running NixOS, Nixpkgs unstable on an AMD 5600 CPU and NVIDIA RTX 3060 graphics card. I’m quite familiar with uv, which handles Python dependencies well, and since my collaborators might not all work on NixOS, I’ll stick with uv as I used to do. I think the community has come up with a really good solution: uv2nix, I find it really convenient to work on.
  • So this is my flake.nix:
  • {
  description = "WIP Analysis Project";

  nixConfig = {
    extra-substituters = [
      "https://nix-community.cachix.org"
    ];
    extra-trusted-public-keys = [
      "nix-community.cachix.org-1:0dq3bujKpuEPMCX6U4WylrUDZ9JyUG0VpVZa7CNfq5E="
    ];
  };

  inputs = {
    nixpkgs.url = "github:NixOS/nixpkgs/nixpkgs-unstable";
    flake-parts.url = "github:hercules-ci/flake-parts";

    pyproject-nix = {
      url = "github:pyproject-nix/pyproject.nix";
      inputs.nixpkgs.follows = "nixpkgs";
    };

    uv2nix = {
      url = "github:pyproject-nix/uv2nix";
      inputs.pyproject-nix.follows = "pyproject-nix";
      inputs.nixpkgs.follows = "nixpkgs";
    };

    pyproject-build-systems = {
      url = "github:pyproject-nix/build-system-pkgs";
      inputs.pyproject-nix.follows = "pyproject-nix";
      inputs.uv2nix.follows = "uv2nix";
      inputs.nixpkgs.follows = "nixpkgs";
    };
  };

  outputs =
    {
      self,
      nixpkgs,
      flake-parts,
      pyproject-nix,
      uv2nix,
      pyproject-build-systems,
      ...
    }@inputs:
    flake-parts.lib.mkFlake { inherit inputs; } {
      systems = [
        "x86_64-linux"
        "aarch64-linux"
      ];

      perSystem =
        {
          config,
          self',
          inputs',
          pkgs,
          system,
          ...
        }:
        let

          pkgs = import inputs.nixpkgs {
            inherit system;
            config = {
              allowUnfree = true;
              # cudaSupport = true;
            };
          };

          python = pkgs.python312;
          workspace = inputs.uv2nix.lib.workspace.loadWorkspace {
            workspaceRoot = ./backend;
          };

          overlay = workspace.mkPyprojectOverlay {
            sourcePreference = "wheel";
          };
          editableOverlay = workspace.mkEditablePyprojectOverlay {
            root = "./backend";
          };
          customOverlay = self: super: {
            torch = super.torch.overrideAttrs (old: {
              nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
                pkgs.cudaPackages.cudatoolkit
                pkgs.cudaPackages.cudnn
                pkgs.cudaPackages.libcusparse
                pkgs.cudaPackages.libcusparse_lt
                pkgs.cudaPackages.libcufile
                pkgs.cudaPackages.libnvshmem
                pkgs.cudaPackages.nccl
              ];
              autoPatchelfIgnoreMissingDeps = (old.autoPatchelfIgnoreMissingDeps or [ ]) ++ [ "libcuda.so.1" ];
            });
            nvidia-cufile-cu12 = super.nvidia-cufile-cu12.overrideAttrs (old: {
              nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
                pkgs.rdma-core
              ];
            });
            nvidia-nvshmem-cu12 = super.nvidia-nvshmem-cu12.overrideAttrs (old: {
              nativeBuildInputs = old.nativeBuildInputs ++ [
                pkgs.openmpi
                pkgs.pmix
                pkgs.ucx
                pkgs.libfabric
                pkgs.rdma-core
              ];
            });
            nvidia-cusparse-cu12 = super.nvidia-cusparse-cu12.overrideAttrs (old: {
              nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
                pkgs.cudaPackages.libnvjitlink
              ];
            });
            nvidia-cusolver-cu12 = super.nvidia-cusolver-cu12.overrideAttrs (old: {
              nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
                pkgs.cudaPackages.libnvjitlink
                pkgs.cudaPackages.libcusparse
                pkgs.cudaPackages.libcublas
              ];
            });

          };

          pythonSets =
            (pkgs.callPackage pyproject-nix.build.packages {
              inherit python;
            }).overrideScope
              (
                pkgs.lib.composeManyExtensions [
                  pyproject-build-systems.overlays.wheel
                  overlay
                  customOverlay
                ]
              );

        in
        {
          devShells.default =
            let
              pythonSet = pythonSets.overrideScope editableOverlay;
              virtualenv = pythonSet.mkVirtualEnv "wip-analysis-dev" workspace.deps.all;
            in
            pkgs.mkShell {
              packages = [
                virtualenv
                pkgs.uv
                pkgs.ruff
                pkgs.cudaPackages.cudatoolkit
                pkgs.cudaPackages.cudnn
                pkgs.cudaPackages.cuda_nvcc
              ];

              env = {
                UV_NO_SYNC = "1";
                UV_PYTHON = pythonSet.python.interpreter;
                UV_PYTHON_DOWNLOADS = "never";
              };

              shellHook = ''
                echo "WIP-ANALYSIS-DEV"
                export REPO_ROOT=$(git rev-parse --show-toplevel)
                unset PYTHONPATH
              '';
            };

          packages.default = pythonSets.mkVirtualEnv "wip-analysis-env" workspace.deps.default;

        };

    };
}
  • A few points I found really hard to overcome:
    • On consumer-grade hardware, you should definitely set up caching, since compiling from source will take hundreds of years and ultimately leave nothing for you.
    • Google Gemini or ChatGPT know very little about Nix or NixOS, Flake, or Python development on NixOS; I don’t trust them on general suggestions or solutions, But they really helped me find the Nix packages that provide the libraries needed for uv2nix to patch Python dependencies.
  • So this is my workflow: first, initiate a Python project with uv (nix shell nixpkgs#uv nixpkgs#python3 then uv init --app --package) and set up a flake.nix file in the project root then enter dev shell.
  • use uv add to add dependencies and uv lock to update uv.lock file, and if using direnv, direnv reload to re-evaluate the flake to make the new dependencies available in the environment.
  • a few more words about Using an Nvidia graphics card on NixOS, from my NixOS configuration:
  • { config, pkgs, ... }:

{
  imports = [
    ./hardware-configuration.nix
  ];

  networking.hostName = "Syy-PC";

  # for Nvidia GPU
  services.xserver.videoDrivers = [ "nvidia" ];
  hardware.graphics.enable = true;
  hardware.nvidia = {
    open = true;
    package = config.boot.kernelPackages.nvidiaPackages.stable;
    modesetting.enable = true;
  };

  # Use CUDA Binary Cache
  nix.settings = {
    substituters = [
      "https://nix-community.cachix.org"
    ];
    trusted-public-keys = [
      "nix-community.cachix.org-1:mB9FSh9qf2dCimDSUo8Zy7bkq5CX+/rkCWyvRCYg3Fs="
    ];
  };

  nixpkgs.config.cudaSupport = true;
  # nixpkgs.overlays = [
  #   (self: super:{
  #     btop = super.btop.override { cudaSupport = true; };
  #   })
  # ];

  # This value determines the NixOS release from which the default
  # settings for stateful data, like file locations and database versions
  # on your system were taken. It‘s perfectly fine and recommended to leave
  # this value at the release version of the first install of this system.
  # Before changing this value read the documentation for this option
  # (e.g. man configuration.nix or on https://nixos.org/nixos/options.html).
  system.stateVersion = "25.05"; # Did you read the comment?
}
  • Hope I didn’t get things wrong and this post helps someone out.
  • Thanks and references:
3 Likes