From 6ddb9a7f5e2119327ebfc8dbcc98cc4456ab4657 Mon Sep 17 00:00:00 2001 From: DavHau Date: Tue, 26 Mar 2024 12:49:09 +0700 Subject: [PATCH] override/python: fix devShells for nvidia packages --- modules/flake-parts/pythonEnv/flake-template.nix | 2 +- overrides/python/nvidia-cublas-cu12/default.nix | 5 +++++ .../python/nvidia-cublas-cu12/remove-nvidia-init.nix | 10 ++++++++++ overrides/python/nvidia-cuda-cupti-cu12/default.nix | 5 +++++ overrides/python/nvidia-cuda-nvrtc-cu12/default.nix | 5 +++++ overrides/python/nvidia-cuda-runtime-cu12/default.nix | 5 +++++ overrides/python/nvidia-cudnn-cu12/default.nix | 5 +++++ overrides/python/nvidia-cufft-cu12/default.nix | 5 +++++ overrides/python/nvidia-curand-cu12/default.nix | 5 +++++ overrides/python/nvidia-cusolver-cu12/default.nix | 5 +++++ overrides/python/nvidia-cusparse-cu12/default.nix | 5 +++++ overrides/python/nvidia-nccl-cu12/default.nix | 5 +++++ overrides/python/nvidia-nvjitlink-cu12/default.nix | 5 +++++ 13 files changed, 66 insertions(+), 1 deletion(-) create mode 100644 overrides/python/nvidia-cublas-cu12/default.nix create mode 100644 overrides/python/nvidia-cublas-cu12/remove-nvidia-init.nix create mode 100644 overrides/python/nvidia-cuda-cupti-cu12/default.nix create mode 100644 overrides/python/nvidia-cuda-nvrtc-cu12/default.nix create mode 100644 overrides/python/nvidia-cuda-runtime-cu12/default.nix create mode 100644 overrides/python/nvidia-cudnn-cu12/default.nix create mode 100644 overrides/python/nvidia-cufft-cu12/default.nix create mode 100644 overrides/python/nvidia-curand-cu12/default.nix create mode 100644 overrides/python/nvidia-cusolver-cu12/default.nix create mode 100644 overrides/python/nvidia-cusparse-cu12/default.nix create mode 100644 overrides/python/nvidia-nccl-cu12/default.nix create mode 100644 overrides/python/nvidia-nvjitlink-cu12/default.nix diff --git a/modules/flake-parts/pythonEnv/flake-template.nix b/modules/flake-parts/pythonEnv/flake-template.nix index 91045426..2515d1ba 100644 --- a/modules/flake-parts/pythonEnv/flake-template.nix +++ b/modules/flake-parts/pythonEnv/flake-template.nix @@ -2,7 +2,7 @@ description = "My flake with dream2nix packages"; inputs = { - dream2nix.url = "github:nix-community/dream2nix/pythonEnv"; + dream2nix.url = "github:nix-community/dream2nix"; nixpkgs.follows = "dream2nix/nixpkgs"; flake-parts.url = "github:hercules-ci/flake-parts"; }; diff --git a/overrides/python/nvidia-cublas-cu12/default.nix b/overrides/python/nvidia-cublas-cu12/default.nix new file mode 100644 index 00000000..eab27caf --- /dev/null +++ b/overrides/python/nvidia-cublas-cu12/default.nix @@ -0,0 +1,5 @@ +{ + imports = [ + ./remove-nvidia-init.nix + ]; +} diff --git a/overrides/python/nvidia-cublas-cu12/remove-nvidia-init.nix b/overrides/python/nvidia-cublas-cu12/remove-nvidia-init.nix new file mode 100644 index 00000000..791c4ddf --- /dev/null +++ b/overrides/python/nvidia-cublas-cu12/remove-nvidia-init.nix @@ -0,0 +1,10 @@ +{ + config, + lib, + ... +}: { + mkDerivation.postFixup = + # prevents conflicts in nixpkgs buildEnv for python + lib.mkIf (lib.hasSuffix ".whl" config.mkDerivation.src) + "rm $out/lib/*/site-packages/nvidia/__pycache__/__init__.*"; +} diff --git a/overrides/python/nvidia-cuda-cupti-cu12/default.nix b/overrides/python/nvidia-cuda-cupti-cu12/default.nix new file mode 100644 index 00000000..cfaf5a45 --- /dev/null +++ b/overrides/python/nvidia-cuda-cupti-cu12/default.nix @@ -0,0 +1,5 @@ +{ + imports = [ + ../nvidia-cublas-cu12/remove-nvidia-init.nix + ]; +} diff --git a/overrides/python/nvidia-cuda-nvrtc-cu12/default.nix b/overrides/python/nvidia-cuda-nvrtc-cu12/default.nix new file mode 100644 index 00000000..cfaf5a45 --- /dev/null +++ b/overrides/python/nvidia-cuda-nvrtc-cu12/default.nix @@ -0,0 +1,5 @@ +{ + imports = [ + ../nvidia-cublas-cu12/remove-nvidia-init.nix + ]; +} diff --git a/overrides/python/nvidia-cuda-runtime-cu12/default.nix b/overrides/python/nvidia-cuda-runtime-cu12/default.nix new file mode 100644 index 00000000..cfaf5a45 --- /dev/null +++ b/overrides/python/nvidia-cuda-runtime-cu12/default.nix @@ -0,0 +1,5 @@ +{ + imports = [ + ../nvidia-cublas-cu12/remove-nvidia-init.nix + ]; +} diff --git a/overrides/python/nvidia-cudnn-cu12/default.nix b/overrides/python/nvidia-cudnn-cu12/default.nix new file mode 100644 index 00000000..cfaf5a45 --- /dev/null +++ b/overrides/python/nvidia-cudnn-cu12/default.nix @@ -0,0 +1,5 @@ +{ + imports = [ + ../nvidia-cublas-cu12/remove-nvidia-init.nix + ]; +} diff --git a/overrides/python/nvidia-cufft-cu12/default.nix b/overrides/python/nvidia-cufft-cu12/default.nix new file mode 100644 index 00000000..cfaf5a45 --- /dev/null +++ b/overrides/python/nvidia-cufft-cu12/default.nix @@ -0,0 +1,5 @@ +{ + imports = [ + ../nvidia-cublas-cu12/remove-nvidia-init.nix + ]; +} diff --git a/overrides/python/nvidia-curand-cu12/default.nix b/overrides/python/nvidia-curand-cu12/default.nix new file mode 100644 index 00000000..cfaf5a45 --- /dev/null +++ b/overrides/python/nvidia-curand-cu12/default.nix @@ -0,0 +1,5 @@ +{ + imports = [ + ../nvidia-cublas-cu12/remove-nvidia-init.nix + ]; +} diff --git a/overrides/python/nvidia-cusolver-cu12/default.nix b/overrides/python/nvidia-cusolver-cu12/default.nix new file mode 100644 index 00000000..cfaf5a45 --- /dev/null +++ b/overrides/python/nvidia-cusolver-cu12/default.nix @@ -0,0 +1,5 @@ +{ + imports = [ + ../nvidia-cublas-cu12/remove-nvidia-init.nix + ]; +} diff --git a/overrides/python/nvidia-cusparse-cu12/default.nix b/overrides/python/nvidia-cusparse-cu12/default.nix new file mode 100644 index 00000000..cfaf5a45 --- /dev/null +++ b/overrides/python/nvidia-cusparse-cu12/default.nix @@ -0,0 +1,5 @@ +{ + imports = [ + ../nvidia-cublas-cu12/remove-nvidia-init.nix + ]; +} diff --git a/overrides/python/nvidia-nccl-cu12/default.nix b/overrides/python/nvidia-nccl-cu12/default.nix new file mode 100644 index 00000000..cfaf5a45 --- /dev/null +++ b/overrides/python/nvidia-nccl-cu12/default.nix @@ -0,0 +1,5 @@ +{ + imports = [ + ../nvidia-cublas-cu12/remove-nvidia-init.nix + ]; +} diff --git a/overrides/python/nvidia-nvjitlink-cu12/default.nix b/overrides/python/nvidia-nvjitlink-cu12/default.nix new file mode 100644 index 00000000..cfaf5a45 --- /dev/null +++ b/overrides/python/nvidia-nvjitlink-cu12/default.nix @@ -0,0 +1,5 @@ +{ + imports = [ + ../nvidia-cublas-cu12/remove-nvidia-init.nix + ]; +}