tjh.dev
1{
2 lib,
3 runCommand,
4 python3Packages,
5 makeWrapper,
6 writableTmpDirAsHomeHook,
7}:
8{
9 feature ? "cuda",
10 name ? if feature == null then "cpu" else feature,
11 libraries ? [ ], # [PythonPackage] | (PackageSet -> [PythonPackage])
12 gpuCheckArgs ? { },
13 ...
14}@args:
15
16let
17 inherit (builtins) isFunction all;
18 librariesFun = if isFunction libraries then libraries else (_: libraries);
19in
20
21assert lib.assertMsg (
22 isFunction libraries || all (python3Packages.hasPythonModule) libraries
23) "writeGpuTestPython was passed `libraries` from the wrong python release";
24
25content:
26
27let
28 interpreter = python3Packages.python.withPackages librariesFun;
29 tester =
30 runCommand "tester-${name}"
31 (
32 lib.removeAttrs args [
33 "gpuCheckArgs"
34 "libraries"
35 "name"
36 ]
37 // {
38 inherit content;
39 nativeBuildInputs = args.nativeBuildInputs or [ ] ++ [ makeWrapper ];
40 passAsFile = args.passAsFile or [ ] ++ [ "content" ];
41 }
42 )
43 ''
44 mkdir -p "$out"/bin
45 cat << EOF >"$out/bin/$name"
46 #!${lib.getExe interpreter}
47 EOF
48 cat "$contentPath" >>"$out/bin/$name"
49 chmod +x "$out/bin/$name"
50
51 if [[ -n "''${makeWrapperArgs+''${makeWrapperArgs[@]}}" ]] ; then
52 wrapProgram "$out/bin/$name" ''${makeWrapperArgs[@]}
53 fi
54 '';
55 tester' = tester.overrideAttrs (oldAttrs: {
56 passthru.gpuCheck =
57 runCommand "test-${name}"
58 (
59 gpuCheckArgs
60 // {
61 nativeBuildInputs = [
62 tester'
63 ]
64 ++ gpuCheckArgs.nativeBuildInputs or [ ];
65
66 requiredSystemFeatures =
67 lib.optionals (feature != null) [ feature ] ++ gpuCheckArgs.requiredSystemFeatures or [ ];
68 }
69 )
70 ''
71 set -e
72 ${tester.meta.mainProgram or (lib.getName tester')}
73 touch $out
74 '';
75 });
76in
77tester'