JupyterHub: Machine learning toolkits addition

machines/kcnhub/nvidia.nix

@@ -1,5 +1,4 @@
 { config, lib, pkgs, ...}: {
-  config = {
   # Make sure opengl is enabled
   hardware.opengl = {
     enable = true;
@@ -21,7 +20,8 @@
 
     # Use the open source version of the kernel module
     # 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
     nvidiaSettings = true;
@@ -32,15 +32,15 @@
 
   # Allow Unfree, Nvidia, CUDA, and derived packages.
   nixpkgs.config.allowUnfree = true;
+  nixpkgs.config.cudaSupport = true;
 
   # Use Prebuilt Cached Versions of CUDA Packages
   nix.settings = {
-    trusted-substituters = [
+    substituters = [
       "https://cuda-maintainers.cachix.org"
     ];
     trusted-public-keys = [
       "cuda-maintainers.cachix.org-1:0dq3bujKpuEPMCX6U4WylrUDZ9JyUG0VpVZa7CNfq5E="
     ];
   };
-};
 }

machines/kcnhub/servers/jupyterhub.nix

@@ -7,7 +7,9 @@
     extraConfig = ''
       c.SystemdSpawner.mem_limit = '2G'
       c.SystemdSpawner.cpu_limit = 1.0
-      c.SystemdSpawner.extra_paths = ['${pkgs.buildEnv {
+      c.SystemdSpawner.isolate_devices = False
+      c.SystemdSpawner.extra_paths = [
+        '${pkgs.buildEnv {
           name = "jupyterhub-system-spawner-env";
           paths = with pkgs; [
             gcc
@@ -70,7 +72,25 @@
             withMpi = false;
           })
           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 {
         displayName = "Python 3 for Computational Neuroscience";
         argv = [