JupyterHub: Machine learning toolkits addition

master
David Crompton 2024-02-01 20:02:53 -05:00
parent fb5ebf1be0
commit 9fd4f2c67c
2 changed files with 31 additions and 11 deletions

View File

@ -1,5 +1,4 @@
{ config, lib, pkgs, ...}: { { config, lib, pkgs, ...}: {
config = {
# Make sure opengl is enabled # Make sure opengl is enabled
hardware.opengl = { hardware.opengl = {
enable = true; enable = true;
@ -21,7 +20,8 @@
# Use the open source version of the kernel module # Use the open source version of the kernel module
# Only available on driver 515.43.04+ # Only available on driver 515.43.04+
open = true; # Seems to cause some issues with Tensorflow/CUDA/PyTorch
open = false;
# Enable the nvidia settings menu # Enable the nvidia settings menu
nvidiaSettings = true; nvidiaSettings = true;
@ -32,15 +32,15 @@
# Allow Unfree, Nvidia, CUDA, and derived packages. # Allow Unfree, Nvidia, CUDA, and derived packages.
nixpkgs.config.allowUnfree = true; nixpkgs.config.allowUnfree = true;
nixpkgs.config.cudaSupport = true;
# Use Prebuilt Cached Versions of CUDA Packages # Use Prebuilt Cached Versions of CUDA Packages
nix.settings = { nix.settings = {
trusted-substituters = [ substituters = [
"https://cuda-maintainers.cachix.org" "https://cuda-maintainers.cachix.org"
]; ];
trusted-public-keys = [ trusted-public-keys = [
"cuda-maintainers.cachix.org-1:0dq3bujKpuEPMCX6U4WylrUDZ9JyUG0VpVZa7CNfq5E=" "cuda-maintainers.cachix.org-1:0dq3bujKpuEPMCX6U4WylrUDZ9JyUG0VpVZa7CNfq5E="
]; ];
}; };
};
} }

View File

@ -7,12 +7,14 @@
extraConfig = '' extraConfig = ''
c.SystemdSpawner.mem_limit = '2G' c.SystemdSpawner.mem_limit = '2G'
c.SystemdSpawner.cpu_limit = 1.0 c.SystemdSpawner.cpu_limit = 1.0
c.SystemdSpawner.extra_paths = ['${pkgs.buildEnv { c.SystemdSpawner.isolate_devices = False
name = "jupyterhub-system-spawner-env"; c.SystemdSpawner.extra_paths = [
paths = with pkgs; [ '${pkgs.buildEnv {
gcc name = "jupyterhub-system-spawner-env";
]; paths = with pkgs; [
}}/bin'] gcc
];
}}/bin']
''; '';
# Extra Paths used to add gcc et al for packages like Brian2 # Extra Paths used to add gcc et al for packages like Brian2
# https://github.com/jupyterhub/systemdspawner#extra_paths # https://github.com/jupyterhub/systemdspawner#extra_paths
@ -70,7 +72,25 @@
withMpi = false; withMpi = false;
}) })
brian2 brian2
]));
# Machine learning toolkits:
tensorflow
tensorboard
keras
edward
transformers
tflearn
torch
torchvision
torchsde
torchmetrics
torchio
torchdiffeq
botorch
lion-pytorch
])).override (args: { ignoreCollisions = true; });
# Odd collision between tensorboard of torch & tensorflow
# need to resolve later
in { in {
displayName = "Python 3 for Computational Neuroscience"; displayName = "Python 3 for Computational Neuroscience";
argv = [ argv = [