From c774c7bffe0d5fbf3da5c61e9b8bc3a8d8b11cf7 Mon Sep 17 00:00:00 2001 From: Someone Serge Date: Sat, 20 Jul 2024 12:02:30 +0000 Subject: [PATCH 1/5] cudaPackages.writeGpuTestPython: allow a selector for `libraries` to accommodate different python versions --- .../cuda-modules/write-gpu-python-test.nix | 28 +++++++++++++++++-- 1 file changed, 26 insertions(+), 2 deletions(-) diff --git a/pkgs/development/cuda-modules/write-gpu-python-test.nix b/pkgs/development/cuda-modules/write-gpu-python-test.nix index 5f0d5c6b8fe6..10b10fa3f923 100644 --- a/pkgs/development/cuda-modules/write-gpu-python-test.nix +++ b/pkgs/development/cuda-modules/write-gpu-python-test.nix @@ -2,16 +2,40 @@ lib, writers, runCommand, + python3Packages, }: { feature ? "cuda", name ? feature, - libraries ? [ ], + libraries ? [ ], # [PythonPackage] | (PackageSet -> [PythonPackage]) }: + +let + inherit (builtins) isFunction all; + librariesFun = if isFunction libraries then libraries else (_: libraries); +in + +assert lib.assertMsg ( + isFunction libraries || all (python3Packages.hasPythonModule) libraries +) "writeGpuTestPython was passed `libraries` from the wrong python release"; + content: let - tester = writers.writePython3Bin "tester-${name}" { inherit libraries; } content; + interpreter = python3Packages.python.withPackages librariesFun; + tester = + runCommand "tester-${name}" + { + inherit content; + passAsFile = [ "content" ]; + } + '' + mkdir -p "$out"/bin + cat << EOF >"$out"/bin/"tester-${name}" + #!${lib.getExe interpreter} + EOF + cat "$contentPath" >>"$out"/bin/"tester-${name}" + ''; tester' = tester.overrideAttrs (oldAttrs: { passthru.gpuCheck = runCommand "test-${name}" From 2d5e573acf2af54df6a8d6bcc4dc5dce639b2d72 Mon Sep 17 00:00:00 2001 From: Someone Serge Date: Sat, 20 Jul 2024 12:03:36 +0000 Subject: [PATCH 2/5] cudaPackages.writeGpuTestPython: sync the attr and the filesystem paths --- .../{write-gpu-python-test.nix => write-gpu-test-python.nix} | 0 pkgs/top-level/cuda-packages.nix | 2 +- 2 files changed, 1 insertion(+), 1 deletion(-) rename pkgs/development/cuda-modules/{write-gpu-python-test.nix => write-gpu-test-python.nix} (100%) diff --git a/pkgs/development/cuda-modules/write-gpu-python-test.nix b/pkgs/development/cuda-modules/write-gpu-test-python.nix similarity index 100% rename from pkgs/development/cuda-modules/write-gpu-python-test.nix rename to pkgs/development/cuda-modules/write-gpu-test-python.nix diff --git a/pkgs/top-level/cuda-packages.nix b/pkgs/top-level/cuda-packages.nix index 5540b89f1b98..639fa70446be 100644 --- a/pkgs/top-level/cuda-packages.nix +++ b/pkgs/top-level/cuda-packages.nix @@ -81,7 +81,7 @@ let nccl = final.callPackage ../development/cuda-modules/nccl { }; nccl-tests = final.callPackage ../development/cuda-modules/nccl-tests { }; - writeGpuTestPython = final.callPackage ../development/cuda-modules/write-gpu-python-test.nix { }; + writeGpuTestPython = final.callPackage ../development/cuda-modules/write-gpu-test-python.nix { }; }); mkVersionedPackageName = From e863b5843a3223998adf3f634aa1e6e7af0db6d4 Mon Sep 17 00:00:00 2001 From: Someone Serge Date: Sat, 20 Jul 2024 13:05:15 +0000 Subject: [PATCH 3/5] cudaPackages.writeGpuTestPython: accept makeWrapperArgs --- .../cuda-modules/write-gpu-test-python.nix | 33 +++++++++++++------ 1 file changed, 23 insertions(+), 10 deletions(-) diff --git a/pkgs/development/cuda-modules/write-gpu-test-python.nix b/pkgs/development/cuda-modules/write-gpu-test-python.nix index 10b10fa3f923..23282e916929 100644 --- a/pkgs/development/cuda-modules/write-gpu-test-python.nix +++ b/pkgs/development/cuda-modules/write-gpu-test-python.nix @@ -1,14 +1,15 @@ { lib, - writers, runCommand, python3Packages, + makeWrapper, }: { feature ? "cuda", - name ? feature, + name ? if feature == null then "cpu" else feature, libraries ? [ ], # [PythonPackage] | (PackageSet -> [PythonPackage]) -}: + ... +}@args: let inherit (builtins) isFunction all; @@ -25,23 +26,35 @@ let interpreter = python3Packages.python.withPackages librariesFun; tester = runCommand "tester-${name}" - { - inherit content; - passAsFile = [ "content" ]; - } + ( + lib.removeAttrs args [ + "libraries" + "name" + ] + // { + inherit content; + nativeBuildInputs = args.nativeBuildInputs or [ ] ++ [ makeWrapper ]; + passAsFile = args.passAsFile or [ ] ++ [ "content" ]; + } + ) '' mkdir -p "$out"/bin - cat << EOF >"$out"/bin/"tester-${name}" + cat << EOF >"$out/bin/$name" #!${lib.getExe interpreter} EOF - cat "$contentPath" >>"$out"/bin/"tester-${name}" + cat "$contentPath" >>"$out/bin/$name" + chmod +x "$out/bin/$name" + + if [[ -n "''${makeWrapperArgs+''${makeWrapperArgs[@]}}" ]] ; then + wrapProgram "$out/bin/$name" ''${makeWrapperArgs[@]} + fi ''; tester' = tester.overrideAttrs (oldAttrs: { passthru.gpuCheck = runCommand "test-${name}" { nativeBuildInputs = [ tester' ]; - requiredSystemFeatures = [ feature ]; + requiredSystemFeatures = lib.optionals (feature != null) [ feature ]; } '' set -e From 3b76c33407a6b09d78b70717f43a2b32eda35902 Mon Sep 17 00:00:00 2001 From: Someone Serge Date: Sat, 20 Jul 2024 12:05:16 +0000 Subject: [PATCH 4/5] python311Packages.torch.tests.tester-*Available: unbreak for non-default python package sets --- .../python-modules/torch/mk-runtime-check.nix | 7 ++++--- pkgs/development/python-modules/torch/tests.nix | 10 +++------- 2 files changed, 7 insertions(+), 10 deletions(-) diff --git a/pkgs/development/python-modules/torch/mk-runtime-check.nix b/pkgs/development/python-modules/torch/mk-runtime-check.nix index 14560b06f87c..61180a19aaba 100644 --- a/pkgs/development/python-modules/torch/mk-runtime-check.nix +++ b/pkgs/development/python-modules/torch/mk-runtime-check.nix @@ -1,14 +1,15 @@ { cudaPackages, feature, - torch, + libraries, versionAttr, + pythonPackages, }: -cudaPackages.writeGpuTestPython +(cudaPackages.writeGpuTestPython.override { python3Packages = pythonPackages; }) { inherit feature; - libraries = [ torch ]; + inherit libraries; name = "${feature}Available"; } '' diff --git a/pkgs/development/python-modules/torch/tests.nix b/pkgs/development/python-modules/torch/tests.nix index 76b901cbcea9..92ddccccbaeb 100644 --- a/pkgs/development/python-modules/torch/tests.nix +++ b/pkgs/development/python-modules/torch/tests.nix @@ -1,8 +1,4 @@ -{ - callPackage, - torchWithCuda, - torchWithRocm, -}: +{ callPackage }: { # To perform the runtime check use either @@ -11,11 +7,11 @@ tester-cudaAvailable = callPackage ./mk-runtime-check.nix { feature = "cuda"; versionAttr = "cuda"; - torch = torchWithCuda; + libraries = ps: [ ps.torchWithCuda ]; }; tester-rocmAvailable = callPackage ./mk-runtime-check.nix { feature = "rocm"; versionAttr = "hip"; - torch = torchWithRocm; + libraries = ps: [ ps.torchWithRocm ]; }; } From 2cc58226b0205fae68ff9266aa6edef09df06f6d Mon Sep 17 00:00:00 2001 From: Someone Serge Date: Sat, 20 Jul 2024 12:05:52 +0000 Subject: [PATCH 5/5] python3Packages.torch.tests.*compile*: init --- .../torch/mk-torch-compile-check.nix | 38 +++++++++++++++++++ .../python-modules/torch/tests.nix | 16 +++++++- 2 files changed, 53 insertions(+), 1 deletion(-) create mode 100644 pkgs/development/python-modules/torch/mk-torch-compile-check.nix diff --git a/pkgs/development/python-modules/torch/mk-torch-compile-check.nix b/pkgs/development/python-modules/torch/mk-torch-compile-check.nix new file mode 100644 index 000000000000..268ed5297da9 --- /dev/null +++ b/pkgs/development/python-modules/torch/mk-torch-compile-check.nix @@ -0,0 +1,38 @@ +{ + cudaPackages, + feature ? null, + lib, + libraries, + name ? if feature == null then "torch-compile-cpu" else "torch-compile-${feature}", + pythonPackages, + stdenv, +}: +let + deviceStr = if feature == null then "" else '', device="cuda"''; +in +(cudaPackages.writeGpuTestPython.override { python3Packages = pythonPackages; }) + { + inherit name feature libraries; + makeWrapperArgs = [ + "--suffix" + "PATH" + ":" + "${lib.getBin stdenv.cc}/bin" + ]; + } + '' + import torch + + + @torch.compile + def opt_foo2(x, y): + a = torch.sin(x) + b = torch.cos(y) + return a + b + + + print( + opt_foo2( + torch.randn(10, 10${deviceStr}), + torch.randn(10, 10${deviceStr}))) + '' diff --git a/pkgs/development/python-modules/torch/tests.nix b/pkgs/development/python-modules/torch/tests.nix index 92ddccccbaeb..e3f2ca44ba5a 100644 --- a/pkgs/development/python-modules/torch/tests.nix +++ b/pkgs/development/python-modules/torch/tests.nix @@ -1,6 +1,6 @@ { callPackage }: -{ +rec { # To perform the runtime check use either # `nix run .#python3Packages.torch.tests.tester-cudaAvailable` (outside the sandbox), or # `nix build .#python3Packages.torch.tests.tester-cudaAvailable.gpuCheck` (in a relaxed sandbox) @@ -14,4 +14,18 @@ versionAttr = "hip"; libraries = ps: [ ps.torchWithRocm ]; }; + + compileCpu = tester-compileCpu.gpuCheck; + tester-compileCpu = callPackage ./mk-torch-compile-check.nix { + feature = null; + libraries = ps: [ ps.torch ]; + }; + tester-compileCuda = callPackage ./mk-torch-compile-check.nix { + feature = "cuda"; + libraries = ps: [ ps.torchWithCuda ]; + }; + tester-compileRocm = callPackage ./mk-torch-compile-check.nix { + feature = "rocm"; + libraries = ps: [ ps.torchWithRocm ]; + }; }