fancy nix environment that works with cuda 11+

flake.nix

{
  nixConfig.bash-prompt = ''\033[1;32m\[[nix-develop:\[\033[36m\]\w\[\033[32m\]]$\033[0m '';

  inputs = {
    # doesn't provide cuda_12_2
    # leaving this blank most likely loads whatever is on the system
    # nixpkgs.url = "github:nixos/nixpkgs/23.11-pre";
    whacked-setup = {
      url = "github:whacked/setup/6b73eed531bdf4fde58a8214a94d11c4e024082e";
      flake = false;
    };
    nixgl.url = "github:guibou/nixGL";
  };
  outputs = { self, nixpkgs, flake-utils, whacked-setup, nixgl, ... }:
    flake-utils.lib.eachDefaultSystem
    (system:
    let
      name = "refinery";
      # pkgs = nixpkgs.legacyPackages.${system};
      pkgs = import nixpkgs {
        inherit system;
        config.allowUnfree = true;
        config.cudaSupport = true;
        overlays = builtins.trace system (if system == "x86_64-linux" then [
          nixgl.overlay
        ] else []);
      };
      whacked-helpers = import (whacked-setup + /nix/flake-helpers.nix) { inherit pkgs; };

      # hacky way to do environment selection. don't know a better way now
      shouldUseGpu = (builtins.getEnv "ENABLE_CUDA") != "";
      # cudaPackages = pkgs.cudaPackages_12_2;
      cudaPackages = pkgs.cudaPackages_11_8;
      gpu-packages = [
        pkgs.nvtop
        pkgs.linuxPackages_6_1.nvidia_x11
        pkgs.cudaPackages_11.cudnn
        cudaPackages.nccl
        cudaPackages.cuda_cudart
        cudaPackages.cudatoolkit
        cudaPackages.cudatoolkit.lib
        pkgs.libGLU
        pkgs.libGL
        pkgs.xorg.libXi
        pkgs.xorg.libXmu
        pkgs.freeglut
        pkgs.xorg.libXext
        pkgs.xorg.libX11
        pkgs.xorg.libXv
        pkgs.xorg.libXrandr
        pkgs.zlib 
        pkgs.ncurses5
        pkgs.stdenv.cc
        pkgs.binutils
        pkgs.ffmpeg
      ];
      cpuShellEnvironment = ''
        pastel paint yellow --bold "to use cuda-enabled environment, export ENABLE_CUDA=1, and run"
        pastel paint cyan "    NIXPKGS_ALLOW_INSECURE=1 nix develop --impure"
      '';
      gpuShellEnvironment = ''
        export PATH=$PATH:${pkgs.nixgl.auto.nixGLNvidia}/bin
        export LD_LIBRARY_PATH="${pkgs.linuxPackages.nvidia_x11}/lib":${pkgs.stdenv.cc.cc.lib}/lib/:${pkgs.zlib}/lib:${pkgs.openssl_1_1.out}/lib:/run/opengl-driver/lib/
        export LD_LIBRARY_PATH=${cudaPackages.cudatoolkit}/lib:${cudaPackages.cudatoolkit.lib}/lib:$LD_LIBRARY_PATH
        export LD_LIBRARY_PATH=${cudaPackages.nccl}/lib:$LD_LIBRARY_PATH

        # hacks for more libs
        # WARN: linux+cuda only
        export LD_LIBRARY_PATH=${cudaPackages.cudatoolkit}/target-linux-x64:$LD_LIBRARY_PATH
        export LD_LIBRARY_PATH=${pkgs.cudaPackages_11.cudnn}/lib:$LD_LIBRARY_PATH
        export LD_LIBRARY_PATH=${pkgs.libGL}/lib:$LD_LIBRARY_PATH
        export LD_LIBRARY_PATH=${pkgs.glib.out}/lib:$LD_LIBRARY_PATH
       
        # doesn't seem to be used anywhere
        export CUDA_PATH=${cudaPackages.cudatoolkit}

        # sanity check for debugging
        ldd .venv/lib/python3.10/site-packages/torch/lib/libtorch_global_deps.so | grep -i 'not found'
        ldd .venv/lib/python3.10/site-packages/torch/_C.cpython-310-x86_64-linux-gnu.so | grep -i 'not found'
        pastel paint limegreen --bold "cuda environment enabled"
        alias python='nixGLNvidia-535.86.05 python'
      '';
    in {
      # it is not clear if whacked-helpers includes stuff that masks the
      # default c++ build environment, leading to greenlet build failure;
      # if so, try using the original `pkgs.mkShell` which doesn't bundle
      # stuff into the shell.
      # devShell = pkgs.mkShell {
      devShell = whacked-helpers.mkShell {
        flakeFile = __curPos.file;  # used to forward current file to echo-shortcuts
        includeScripts = [
          (whacked-setup + /bash/node_shortcuts.sh)
        ];
      } {
        buildInputs = [
          pkgs.jq
          pkgs.jsonnet

          pkgs.poetry
          pkgs.python3
          pkgs.bash-completion
          pkgs.check-jsonschema
          pkgs.unixtools.column
          pkgs.coreutils

          pkgs.openapi-generator-cli
          pkgs.minio-client
          pkgs.vault
          pkgs.pastel  # for colored messages
          pkgs.postgresql
          pkgs.glibcLocales
          pkgs.awscli
          pkgs.openssl
          pkgs.overmind
          pkgs.zlib
          pkgs.glib
          pkgs.imagemagick
        ] ++ (
          if shouldUseGpu then gpu-packages else []
        );  # join lists with ++

        # ref https://discourse.nixos.org/t/nix-shells-and-python-packages-with-c-extensions/26326
        # Include C++ headers for regular clang calls:
        NIX_CFLAGS_COMPILE = pkgs.lib.optionals pkgs.stdenv.isDarwin [
          "-I${pkgs.lib.getDev pkgs.libcxx}/include/c++/v1"
        ];

        nativeBuildInputs = [
          pkgs.which
          pkgs.git
          pkgs.openssh
        ] ++ pkgs.lib.optionals pkgs.stdenv.isDarwin [
          pkgs.darwin.apple_sdk.frameworks.IOKit
        ];

        shellHook = ''
          ${if pkgs.stdenv.isLinux then (
            ''export LD_LIBRARY_PATH=${pkgs.glibc};/lib:${pkgs.openssl.out}/lib:${pkgs.openssl}/lib:${pkgs.glib.out}/lib:${pkgs.libglvnd}/lib:${pkgs.stdenv.cc.cc.lib}/lib/:${pkgs.zlib}/lib''${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}''
          ) else (
            ""
          )}

          ${if shouldUseGpu then gpuShellEnvironment else cpuShellEnvironment}
          # apparently libssl.so.1.1 is used by paddle?
          export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${pkgs.openssl_1_1.out}/lib
          # for psycopg2
          export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${pkgs.postgresql.lib}/lib
          ### database
          source ${whacked-setup}/bash/postgresql_shortcuts.sh
          export LC_ALL=en_US.UTF-8
          echo-shortcuts ${whacked-setup}/bash/postgresql_shortcuts.sh
          
        '' + ''
          export WORKDIR=$PWD
          export POETRY_VIRTUALENVS_IN_PROJECT=true
          export VIRTUAL_ENV=$WORKDIR/.venv  # poetry-style, NOT pip-style
          if [ ! -e $VIRTUAL_ENV ]; then
            poetry install
          fi
          source $VIRTUAL_ENV/bin/activate
          source ${pkgs.bash-completion}/share/bash-completion/bash_completion
          source ${pkgs.poetry}/share/bash-completion/completions/poetry.bash
          source ${whacked-setup}/bash/nix_shortcuts.sh

          unset shellHook
        '';  # join strings with +
      };
    }
  );
}