NixOS with cuda and micromamba

cuda-maintainers.nix

# cachix/cuda-maintainers.nix (create through `cachix use cuda-maintainers`)
{
  nix = {
    settings = {
      substituters = [
        "https://cuda-maintainers.cachix.org"
      ];
      trusted-public-keys = [
        "cuda-maintainers.cachix.org-1:0dq3bujKpuEPMCX6U4WylrUDZ9JyUG0VpVZa7CNfq5E="
      ];
    };
  };
}

mambarc.yml

channels:
  - conda-forge

ml-setup.nix

{ pkgs, ... }:
{
  # GPU
  nixpkgs.config.allowUnfree = true;
  services.xserver.videoDrivers = [ "nvidia" ];
  hardware.opengl = {
    # Also provides CUDA drivers
    enable = true;
    # Required to make Python packages find cuda without an FHS environment
    setLdLibraryPath = true;
  };
  hardware.nvidia = {
    nvidiaPersistenced = true;
  };

  # System-wide packages
  environment.systemPackages = with pkgs; [
    micromamba
  ];

  # Required to make mamba-managed Python run without an FHS environment
  programs.nix-ld.enable = true;

  # Mamba config, shell hooks and aliases
  environment.shellAliases = {
    # Shorter and easier to rememember
    conda = "micromamba";
  };
  environment.sessionVariables = {
    # Use conda-forge channel by default
    MAMBARC = "${./mambarc.yml}";
    # Personal preference, default is ~/.micromamba
    MAMBA_ROOT_PREFIX = "~/.local/share/mamba";
  };
  programs.bash = {
    interactiveShellInit = ''
      # Mamba does not create its root automatically
      mkdir -p $MAMBA_ROOT_PREFIX
      eval "$(micromamba shell hook -s bash)"
    '';
  };
  programs.zsh = {
    enable = true;
    interactiveShellInit = ''
      mkdir -p $MAMBA_ROOT_PREFIX
      eval "$(micromamba shell hook -s zsh)"
    '';
  };
}