Skip to content

Instantly share code, notes, and snippets.

@SomeoneSerge
Last active May 26, 2023 17:54
Show Gist options
  • Save SomeoneSerge/4832997ab09e4e71301e5469eec3066a to your computer and use it in GitHub Desktop.
Save SomeoneSerge/4832997ab09e4e71301e5469eec3066a to your computer and use it in GitHub Desktop.
{ 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);
# ...
}
{ 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 ];
};
}
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