Python environment

.envrc

#!/bin/sh

use flake .#intel
source venv/bin/activate

flake.nix

{
    description = "rlan nix flake";

    inputs = {
        nixpkgs = {
            # url = "github:nixos/nixpkgs/nixos-unstable";
            url = "github:nixos/nixpkgs";
        };
    };

    outputs = { self, nixpkgs }:

        let
            import-nixpkgs = cudaSupport: import nixpkgs {
                system = "x86_64-linux";
                config = {
                    allowUnfree = true;
                    inherit cudaSupport;
                };
            };
            pkgs-with-cuda = import-nixpkgs true;
            pkgs-without-cuda = import-nixpkgs false;

            make-shell = pkgs: pkgs.mkShell {
                buildInputs = (
                    with pkgs; [
                        python3

                        # For building certain pip packages
                        cmake
                        cudatoolkit
                        swig
                    ]
                ) ++ (
                    with pkgs.python3Packages; [
                        pygame
                        tensorboard
                        pytorch-bin
                        torchaudio-bin
                    ]
                );

                shellHook =
                let
                    packageList = with pkgs; [
                        # for PyTorch
                        stdenv.cc.cc.lib

                        # for Numpy
                        zlib

                        # GL libraries (for gym environment rendering)
                        libGL
                        libGLU

                        # Provides libgthread (for pyqtgraph)
                        glib.out
                        xorg.libxcb
                    ];

                    libPath = nixpkgs.lib.strings.concatMapStringsSep ":" (p: "${p}/lib") packageList;
                in
                    ''
                        export LD_LIBRARY_PATH=${libPath}

                        export SDL_VIDEODRIVER=x11

                        export PYTHONPATH=$(pwd):$PYTHONPATH
                    '';
            };

        in {
            devShells.x86_64-linux = {
                default = make-shell pkgs-with-cuda;
                intel = make-shell pkgs-without-cuda;
            };
        };
}