shell.nix: tmux, zsh, poetry, cuda, make(serve,env)

{ pkgs ? import  {
  # Enable the installation of unfree packages
  config = {
    allowUnfree = true;
  };
} }:

pkgs.mkShell {
  buildInputs = with pkgs; [
    gcc
    zsh
    tmux
    cudaPackages_12_2.cudatoolkit
  ];

  shellHook = ''

    # Set the shell =================================================
    export SHELL=${pkgs.zsh}/bin/zsh
    export TMPDIR=/var/tmp # to avoid python and pip oom errors: https://superuser.com/questions/1288308/why-am-i-getting-no-space-left-on-device-when-it-appears-that-theres-lots-of

    source $(poetry env info --path)/bin/activate # instead of poetry shell

    # Setup the environment =========================================

    ## gcc
    export LD_LIBRARY_PATH=${pkgs.gcc}/lib64:${pkgs.stdenv.cc.cc.lib}/lib:$LD_LIBRARY_PATH

    ## cuda
    export CUDA_PATH=${pkgs.cudaPackages_12_2.cudatoolkit}
    export LD_LIBRARY_PATH=${pkgs.cudaPackages_12_2.cudatoolkit}/lib64:$LD_LIBRARY_PATH
    export LD_LIBRARY_PATH=/run/opengl-driver/lib:$LD_LIBRARY_PATH

    # Initialize tmux ===============================================

    # Prepare the tmux project layout ==============================

    # Get the name of the current directory to use as the tmux session name
    SESSION_NAME=$(basename "$(pwd)")

    if [ -z "$TMUX" ]; then

      # Window: Code editing
      tmux new-session -d -s "$SESSION_NAME" -n 'code'
      tmux send-keys -t "$SESSION_NAME:code" 'nvim .' C-m

      # Window: Shell with the env
      tmux new-window -t "$SESSION_NAME" -n 'shell'
      tmux send-keys -t "$SESSION_NAME:shell" 'make env' C-m

      # Window: Start server
      tmux new-window -t "$SESSION_NAME" -n 'server'
      tmux send-keys -t "$SESSION_NAME:server" 'make serve' C-m

      # Attach to code editing
      tmux select-window -t "$SESSION_NAME:code"
      tmux attach-session -t "$SESSION_NAME"
    fi

    exec ${pkgs.zsh}/bin/zsh
  '';
}

Next step

poetry add torch torchvision torchaudio