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