CUDA Cache for Nix Community

Disclaimer: the author of this note is not a member of the Nix Community infrastructure team

The official NixOS Hydra does not build unfree packages due to the associated legal and security concerns. This has made consuming CUDA-enabled Nixpkgs challenging for an end-user because of both the expensive mass rebuilds, triggered by importing nixpkgs with { config.cudaSupport = true; }, and the cost of building certain leaf derivations such as python3Packages.{torch,jaxlib,tensorflow} or blender. While setting up a private build and cache infrastructure may only be an item in a company’s quotation, for a non-commercial or individual consumer trying out the technology it is an upfront cost that creates a steep entry barrier. This hinders adoption in e.g. the academic environment. Nix Community (@nix-community) is an organization that offers a platform on which to bootstrap ecosystem projects, a flexible “community playground” with sustainable shared resources. For over three months Nix Community, in a joint effort with the Nixpkgs CUDA Team, has been building and caching CUDA-enabled packages in a Hydra jobset. The release definition is maintained upstream in Nixpkgs and the list can be updated in a normal pull request. The jobset currently follows the nixos-unstable-small and nixos-$release-small channels, you may inspect the terraform files for details. The infrastructure can be sponsored using the Nix Community’s OpenCollective^11 project.

In a NixOS config you can enable the binary cache (the “substituter”) using:

# Cf. `man nix.conf` to construct an equivalent config for non-NixOS runtimes
{
  nix.settings.substituters = [
    "https://nix-community.cachix.org"
  ];
  nix.settings.trusted-public-keys = [
    # Compare to the key published at https://nix-community.org/cache
    "nix-community.cachix.org-1:mB9FSh9qf2dCimDSUo8Zy7bkq5CX+/rkCWyvRCYg3Fs="
  ];
}

To enable CUDA support in packages use the NixOS configuration:

  nixpkgs.config.cudaSupport = true;
  nixpkgs.config.allowUnfreePredicate =
    p:
    builtins.all (
      license:
      license.free
      || builtins.elem license.shortName [
        "CUDA EULA"
        "cuDNN EULA"
        "cuTENSOR EULA"
        "NVidia OptiX EULA"
      ]
    ) (if builtins.isList p.meta.license then p.meta.license else [ p.meta.license ]);
}

When preparing a development environment use pkgs = import <nixpkgs> { config = { cudaSupport = true; allowUnfreePredicate = /* ... */; }; }. When using the stable CLI write

$ nix-build "<nixpkgs>" \
  --arg config '{ cudaSupport = true; allowUnfree = true; }' \
  -A blender

or, with the experimental nix-command:

$ nix build -f "<nixpkgs>" \
  --arg config '{ /* ... */ }' \
  blender
$ nix build -f "<nixpkgs>" \
  -I nixpkgs=flake:github:NixOS/nixpkgs/master \
  --arg config '{ /* ... */ }' \
  cudaPackages.saxpy

The inclusion of CUDA in Nix Community’s cachix allows to phase out the Hercules CI jobset that was used for testing by the Nixpkgs CUDA Team, freeing up the resource and hours for other projects.

The new jobset is only the beginning: our teams already consider the next steps we can take, as we secure more funding and gain more expertise, to make maintenance and consumption of coprocessor-accelerated software in Nixpkgs easier. Renting a GPU machine would let us systemically test CUDA, ROCm, OpenGL, Vulkan, and other runtimes. Providing a binary cache for NVidia Jetson devices is near on the radar, being essentially a matter of a modest refactoring on the Nixpkgs side. The current jobset only builds packages for maximum hardware compatibility, but depending on the resource and demand it would be possible to build package sets that target individual GPU architectures and optimize closure sizes. Building more branches at shorter intervals would help to further reduce the number of cache-misses.

Nixpkgs is a collective project vested in by many individuals and companies, driven each by their own motivation, with their own background and in their own circumstances. Ultimately, this is the ecosystem we base our own projects upon, and we seek to keep it sustainable. If you’re a user or a developer, you are encouraged to engage, help in maintenance and community building. If you own a company that leverages Nix and Nixpkgs, do consider committing to regular donations (such as via OpenCollective projects ^10 ^11) and researching whether individual maintainers accept sponsorships or offer consulting services.

11 Likes