How to set up minimal environment with Pytorch and avoid long compile times

Help
1

Hi all,

I’m relatively new to Nix and NixOS and am trying to create a minimal shell.nix to create a shell environment in which I can run Python and Pytorch.

I followed Declarative shell environments with shell.nix — nix.dev documentation to create the following shell.nix:

let
    nixpkgs = fetchTarball "https://github.com/NixOS/nixpkgs/archive/8b8c811c7c2541c30382c5de7ed26be055569c60.tar.gz";
    pkgs = import nixpkgs { config = { allowUnfree = true; }; overlays = []; };
in

pkgs.mkShellNoCC {
    packages = with pkgs; [
        python314
        python314Packages.torch-bin
    ];
}

Next, I added cache.nixos-cuda.org to my system’s configuration.nix according to https://nixos.wiki/wiki/CUDA and did a rebuild:

nix.settings = {
    /* ... */;
    substituters = [ "https://cache.nixos-cuda.org" ];
    trusted-public-keys = [ "cache.nixos-cuda.org:74DUi4Ye579gUqzH4ziL9IyiJBlDpMRn9MBN8oNan9M=" ];
  };

Unfortunately, when I run nix-shell it launches a pretty long compilation step (20 minutes and still running…) that starts with this output:

these 2 derivations will be built:
  /nix/store/dviwdc36igadkkvaf1bmrvlzcxjfhpcj-cuda12.9-libnvshmem-3.6.5-0.drv
  /nix/store/7qdc8i55q7fd609nb9nz8n6n4ck5n5n7-python3.14-torch-2.11.0.drv

I do see a few lines that indicate that the cache might be working to some degree, for example `copying path ‘/nix/store/787q58vxkrmzdp5ddqvf0f06z2vw78rg-mailcap-2.1.54’ from ‘https://cache.nixos.org’…`.

I know that this this topic has been covered many times before, but most of solutions involve other nix features as well. What would be a minimal shell.nix and configuration.nix that I could use to get started with Python and PyTorch relatively quickly?

Thanks a lot for the help!

Update: The compilation actually failed. Here are the last lines:

[140/784] Building CUDA object src/CMakeFiles/nvshmem_host.dir/host/stream/coll/rdxn/reduce_sum.cu.o
ninja: build stopped: subcommand failed.
For full logs, run:
nix log /nix/store/dviwdc36igadkkvaf1bmrvlzcxjfhpcj-cuda12.9-libnvshmem-3.6.5-0.drv

error: Cannot build ‘/nix/store/7qdc8i55q7fd609nb9nz8n6n4ck5n5n7-python3.14-torch-2.11.0.drv’.
Reason: 1 dependency failed.
Output paths:
/nix/store/nxcp7bpybrlfxhhxfa0b2yiw2nbv87ca-python3.14-torch-2.11.0
/nix/store/rcl5lsabwqxclnqx6f5hz67nah977gn2-python3.14-torch-2.11.0-dist
2

I managed to get this working almost immediately after asking my question.

  1. I need to set cudaSupport = true; in the arguments to the evaluation of nixpkgs
  2. I need to use the packagetorch, not torch-bin

Here is my working minimal shell.nix:

let
    nixpkgs = fetchTarball "https://github.com/NixOS/nixpkgs/archive/8b8c811c7c2541c30382c5de7ed26be055569c60.tar.gz";
    pkgs = import nixpkgs { config = { cudaSupport = true; allowUnfree = true; }; overlays = []; };
in

pkgs.mkShellNoCC {
    packages = with pkgs; [
        (python314.withPackages (ps:
            with ps; [
              torch
              torchvision
            ]))
    ];
}

Testing the shell:

>>> import torch
>>> torch.cuda.is_available()
True
>>> torch.cuda.get_device_properties('cuda')
_CudaDeviceProperties(name='NVIDIA RTX A4500', ...)