Last active
May 26, 2023 17:54
-
-
Save SomeoneSerge/4832997ab09e4e71301e5469eec3066a to your computer and use it in GitHub Desktop.
GPU tests in nixpkgs as per https://github.com/NixOS/nixpkgs/issues/225912#issuecomment-1562347597
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
{ pkgs, ... }: | |
{ | |
# ... | |
nix.settings.system-features = [ "expose-cuda" ]; | |
nix.settings.pre-build-hook = pkgs.writers.writePython3 "nix-pre-build.py" { } (builtins.readFile ./nix-pre-build-hook.py); | |
# ... | |
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
{ lib | |
, buildPythonPackage | |
, fetchPypi | |
, substituteAll | |
, pythonOlder | |
, addOpenGLRunpath | |
, pynvml | |
}: | |
buildPythonPackage rec { | |
pname = "pynvml"; | |
version = "11.5.0"; | |
disabled = pythonOlder "3.6"; | |
src = fetchPypi { | |
inherit pname version; | |
hash = "sha256-0CeyG5WxCIufwngRf59ht8Z/jjOnh+n4P3NfD3GsMtA="; | |
}; | |
patches = [ | |
(substituteAll { | |
src = ./0001-locate-libnvidia-ml.so.1-on-NixOS.patch; | |
inherit (addOpenGLRunpath) driverLink; | |
}) | |
]; | |
doCheck = false; | |
passthru.tests.testNvmlInit = pynvml.overridePythonAttrs (_: { | |
checkPhase = '' | |
python3 << \EOF | |
import pynvml | |
from pynvml.smi import nvidia_smi | |
try: | |
print("enter: nvmlInit") | |
print(f"{pynvml.nvmlInit()=}") | |
except Exception as e: | |
print(e) | |
raise | |
finally: | |
print("exit: nvmlInit") | |
EOF | |
''; | |
doCheck = true; | |
requiredSystemFeatures = [ "expose-cuda" ]; | |
}); | |
pythonImportsCheck = [ "pynvml" "pynvml.smi" ]; | |
meta = with lib; { | |
description = "Python bindings for the NVIDIA Management Library"; | |
homepage = "https://www.nvidia.com"; | |
license = licenses.bsd3; | |
maintainers = [ maintainers.bcdarwin ]; | |
}; | |
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import json | |
import subprocess | |
from argparse import ArgumentParser | |
from pathlib import Path | |
from typing import List, Tuple | |
parser = ArgumentParser("pre-build-hook") | |
parser.add_argument("derivation_path") | |
parser.add_argument("sandbox_path", nargs="?") | |
def symlink_parents(p: Path) -> List[Path]: | |
out = [] | |
while p.is_symlink() and p not in out: | |
p = p.readlink() | |
out.append(p) | |
return out | |
if __name__ == "__main__": | |
args = parser.parse_args() | |
drv_path = args.derivation_path | |
proc = subprocess.run( | |
[ | |
"nix", | |
"show-derivation", | |
drv_path, | |
], | |
capture_output=True, | |
) | |
drv = json.loads(proc.stdout) | |
assert drv_path in drv | |
paths: List[Path] = [] | |
for p in Path("/run/opengl-driver/lib").glob("lib*"): | |
if p.name.startswith("libcuda") or p.name.startswith("libnvidia"): | |
paths.append(p) | |
for p in Path("/dev").glob("video*"): | |
paths.append(p) | |
for p in Path("/dev").glob("nvidia*"): | |
paths.append(p) | |
for p in list(paths): | |
paths.extend(symlink_parents(p)) | |
paths = sorted(set(paths)) | |
binds: List[Tuple[str, str]] = [ | |
(p.as_posix(), p.as_posix()) for p in paths | |
] | |
features = drv[drv_path].get("env", {}).get("requiredSystemFeatures", []) | |
if "expose-cuda" in features: | |
print("extra-sandbox-paths") # command | |
for guest_path, host_path in binds: | |
print(f"{guest_path}={host_path}") # arguments, one per line | |
print() # terminated by an empty line |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment