Compute $LD_LIBRARY_PATH from nix-gl-host

- This PR removes $CUDA_HOME/lib from $LD_LIBRARY_PATH because bitsandbytes can follow PyTorch's CUDA runtime since bitsandbytes-foundation/bitsandbytes#1041, and setting $LD_LIBRARY_PATH for `bitsandbytes` is not necessary any more
- Hard-coded `/run/opengl-driver/lib` are removed because nix-gl-host can print the driver path for both Ubuntu and NixOS since numtide/nix-gl-host#16
- `link-nvidia-drivers.nix` is removed because drivers are instead copied to a cache directory by nix-gl-host

examples/my-project/flake.nix

@@ -29,7 +29,7 @@
         inputs.nix-ml-ops.flakeModules.kubernetesJob
         inputs.nix-ml-ops.flakeModules.kubernetesService
         inputs.nix-ml-ops.flakeModules.devcontainer
-        inputs.nix-ml-ops.flakeModules.linkNvidiaDrivers
+        inputs.nix-ml-ops.flakeModules.cuda
         inputs.nix-ml-ops.flakeModules.volumeMountNfs
         inputs.nix-ml-ops.flakeModules.devcontainerGcpCliTools
         inputs.nix-ml-ops.flakeModules.gkeCredential

flake-modules/cuda.nix

@@ -2,24 +2,25 @@ topLevel@{ flake-parts-lib, inputs, ... }: {
   imports = [
     inputs.flake-parts.flakeModules.flakeModules
     ./common.nix
-    ./link-nvidia-drivers.nix
   ];
   flake.flakeModules.cuda = {
     imports = [
       topLevel.config.flake.flakeModules.common
-      topLevel.config.flake.flakeModules.linkNvidiaDrivers
     ];
     options.perSystem = flake-parts-lib.mkPerSystemOption ({ lib, pkgs, system, ... }: {
       config = lib.mkIf (system != "aarch64-darwin") {
         nixpkgs.config.allowUnfree = true;
         nixpkgs.config.cudaSupport = true;
 
         ml-ops.common = common: {
-          config.LD_LIBRARY_PATH = lib.mkMerge [
-            "/run/opengl-driver/lib"
-            # bitsandbytes need to search for CUDA libraries
-            "${common.config.environmentVariables.CUDA_HOME}/lib"
-          ];
+          config.devenvShellModule.enterShell = ''
+            export LD_LIBRARY_PATH="$(${
+              lib.escapeShellArgs [
+                "${inputs.nix-gl-host.defaultPackage.${system}}/bin/nixglhost"
+                "--print-ld-library-path"
+              ]
+            })":''${LD_LIBRARY_PATH:-}
+          '';
           config.devenvShellModule.packages = [
             common.config.cuda.home
           ];

flake.nix

@@ -50,6 +50,10 @@
       inputs.nixpkgs.follows = "nixpkgs";
     };
     nix-ld-rs.url = "github:nix-community/nix-ld-rs";
+    nix-gl-host = {
+      url = "github:Atry/nix-gl-host";
+      inputs.nixpkgs.follows = "nixpkgs";
+    };
   };
   outputs = inputs:
     let