Introduce a cache system
Copying & patching all the DSOs is a time consuming process (~10s on a slow hard drive computer). We definitely don't want to go through it for each process start, we need to introduce a cache. For this cache, we go the concervative way. We're going to "resolve" a DSO name (ie. find the DSO absolute path) and sha256-hash each DSO. We're then going to compare the fingerprints to determine whether or not we need to nuke and rebuild the DSO cache. The cache state is persisted through a JSON file saved in the cache dir.
This commit is contained in:
parent
494bac3613
commit
375148c949
|
@ -19,12 +19,15 @@ pkgs.stdenvNoCC.mkDerivation {
|
||||||
patchShebangs $out/bin/nixglhost
|
patchShebangs $out/bin/nixglhost
|
||||||
'';
|
'';
|
||||||
|
|
||||||
postCheck = ''
|
doCheck = true;
|
||||||
black --check $out/bin/nixglhost
|
|
||||||
|
checkPhase = ''
|
||||||
|
black --check src/*.py
|
||||||
nixpkgs-fmt --check *.nix
|
nixpkgs-fmt --check *.nix
|
||||||
|
python src/nixglhost_wrapper_test.py
|
||||||
'';
|
'';
|
||||||
|
|
||||||
installPhase = ''
|
installPhase = ''
|
||||||
install -D -m0755 nixglhost-wrapper.py $out/bin/nixglhost
|
install -D -m0755 src/nixglhost_wrapper.py $out/bin/nixglhost
|
||||||
'';
|
'';
|
||||||
}
|
}
|
||||||
|
|
|
@ -1,16 +1,20 @@
|
||||||
#!/usr/bin/env python3
|
#!/usr/bin/env python3
|
||||||
|
|
||||||
import argparse
|
import argparse
|
||||||
|
import hashlib
|
||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
import re
|
import re
|
||||||
import shutil
|
import shutil
|
||||||
import subprocess
|
import subprocess
|
||||||
import sys
|
import sys
|
||||||
from typing import List, Dict
|
import time
|
||||||
|
from glob import glob
|
||||||
|
from typing import List, Literal, Dict, Tuple, TypedDict, TextIO, Optional
|
||||||
|
|
||||||
IN_NIX_STORE = False
|
IN_NIX_STORE = False
|
||||||
|
|
||||||
|
|
||||||
if IN_NIX_STORE:
|
if IN_NIX_STORE:
|
||||||
# The following paths are meant to be substituted by Nix at build
|
# The following paths are meant to be substituted by Nix at build
|
||||||
# time.
|
# time.
|
||||||
|
@ -19,6 +23,80 @@ else:
|
||||||
PATCHELF_PATH = "patchelf"
|
PATCHELF_PATH = "patchelf"
|
||||||
|
|
||||||
|
|
||||||
|
class ResolvedLib:
|
||||||
|
def __init__(self, name: str, fullpath: str, sha256: Optional[str] = None):
|
||||||
|
self.name: str = name
|
||||||
|
self.fullpath: str = fullpath
|
||||||
|
if sha256 is None:
|
||||||
|
h = hashlib.sha256()
|
||||||
|
with open(fullpath, "rb") as f:
|
||||||
|
h.update(f.read())
|
||||||
|
sha: str = h.hexdigest()
|
||||||
|
else:
|
||||||
|
sha = sha256
|
||||||
|
self.sha256: str = sha
|
||||||
|
|
||||||
|
def __repr__(self):
|
||||||
|
return f"ResolvedLib<{self.name}, {self.fullpath}, {self.sha256}>"
|
||||||
|
|
||||||
|
def to_dict(self) -> Dict:
|
||||||
|
return {"name": self.name, "fullpath": self.fullpath, "sha256": self.sha256}
|
||||||
|
|
||||||
|
def __eq__(self, o):
|
||||||
|
return (
|
||||||
|
self.name == o.name
|
||||||
|
and self.fullpath == o.fullpath
|
||||||
|
and self.sha256 == o.sha256
|
||||||
|
)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def from_dict(cls, d: Dict):
|
||||||
|
return ResolvedLib(d["name"], d["fullpath"], d["sha256"])
|
||||||
|
|
||||||
|
|
||||||
|
class HostDSOs:
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
glx: Dict[str, ResolvedLib],
|
||||||
|
cuda: Dict[str, ResolvedLib],
|
||||||
|
generic: Dict[str, ResolvedLib],
|
||||||
|
version: int = 1,
|
||||||
|
):
|
||||||
|
self.glx = glx
|
||||||
|
self.cuda = cuda
|
||||||
|
self.generic = generic
|
||||||
|
self.version = version
|
||||||
|
|
||||||
|
def __eq__(self, other):
|
||||||
|
return (
|
||||||
|
self.glx == other.glx
|
||||||
|
and self.cuda == other.cuda
|
||||||
|
and self.generic == other.generic
|
||||||
|
and self.version == other.version
|
||||||
|
)
|
||||||
|
|
||||||
|
def to_json(self) -> str:
|
||||||
|
return json.dumps(
|
||||||
|
{
|
||||||
|
"version": 1,
|
||||||
|
"glx": {k: v.to_dict() for k, v in self.glx.items()},
|
||||||
|
"cuda": {k: v.to_dict() for k, v in self.cuda.items()},
|
||||||
|
"generic": {k: v.to_dict() for k, v in self.generic.items()},
|
||||||
|
},
|
||||||
|
sort_keys=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def from_json(cls, o: str):
|
||||||
|
d: Dict = json.loads(o)
|
||||||
|
return HostDSOs(
|
||||||
|
version=d["version"],
|
||||||
|
glx={k: ResolvedLib.from_dict(v) for k, v in d["glx"].items()},
|
||||||
|
cuda={k: ResolvedLib.from_dict(v) for k, v in d["cuda"].items()},
|
||||||
|
generic={k: ResolvedLib.from_dict(v) for k, v in d["generic"].items()},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
# The following regexes list has been figured out by looking at the
|
# The following regexes list has been figured out by looking at the
|
||||||
# output of nix-build -A linuxPackages.nvidia_x11 before running
|
# output of nix-build -A linuxPackages.nvidia_x11 before running
|
||||||
# ls ./result/lib | grep -E ".so$".
|
# ls ./result/lib | grep -E ".so$".
|
||||||
|
@ -29,7 +107,6 @@ NVIDIA_DSO_PATTERNS = [
|
||||||
"libEGL_nvidia\.so.*$",
|
"libEGL_nvidia\.so.*$",
|
||||||
"libGLESv1_CM_nvidia\.so.*$",
|
"libGLESv1_CM_nvidia\.so.*$",
|
||||||
"libGLESv2_nvidia\.so.*$",
|
"libGLESv2_nvidia\.so.*$",
|
||||||
"libGLX_nvidia\.so.*$",
|
|
||||||
"libglxserver_nvidia\.so.*$",
|
"libglxserver_nvidia\.so.*$",
|
||||||
"libnvcuvid\.so.*$",
|
"libnvcuvid\.so.*$",
|
||||||
"libnvidia-allocator\.so.*$",
|
"libnvidia-allocator\.so.*$",
|
||||||
|
@ -82,12 +159,64 @@ CUDA_DSO_PATTERNS = ["libcudadebugger\.so.*$", "libcuda\.so.*$"]
|
||||||
GLX_DSO_PATTERNS = ["libGLX_nvidia\.so.*$"]
|
GLX_DSO_PATTERNS = ["libGLX_nvidia\.so.*$"]
|
||||||
|
|
||||||
|
|
||||||
def find_files(path: str, files_patterns: List[str]) -> List[str]:
|
def get_ld_paths() -> List[str]:
|
||||||
"""Scans the PATH directory looking for the files complying with
|
"""
|
||||||
the FILES_PATTERNS regexes list.
|
Vendored from https://github.com/albertz/system-tools/blob/master/bin/find-lib-in-path.py
|
||||||
|
|
||||||
Returns the list of the DSOs absolute paths."""
|
Find all the directories pointed by LD_LIBRARY_PATH and the ld cache."""
|
||||||
files = []
|
|
||||||
|
def parse_ld_conf_file(fn: str) -> List[str]:
|
||||||
|
paths = []
|
||||||
|
for l in open(fn).read().splitlines():
|
||||||
|
l = l.strip()
|
||||||
|
if not l:
|
||||||
|
continue
|
||||||
|
if l.startswith("#"):
|
||||||
|
continue
|
||||||
|
if l.startswith("include "):
|
||||||
|
dirglob = l[len("include ") :]
|
||||||
|
if dirglob[0] != "/":
|
||||||
|
dirglob = os.path.dirname(os.path.normpath(fn)) + "/" + dirglob
|
||||||
|
for sub_fn in glob(dirglob):
|
||||||
|
paths.extend(parse_ld_conf_file(sub_fn))
|
||||||
|
continue
|
||||||
|
paths.append(l)
|
||||||
|
return paths
|
||||||
|
|
||||||
|
LDPATH = os.getenv("LD_LIBRARY_PATH")
|
||||||
|
PREFIX = os.getenv("PREFIX") # Termux & etc.
|
||||||
|
paths = []
|
||||||
|
if LDPATH:
|
||||||
|
paths.extend(LDPATH.split(":"))
|
||||||
|
if os.path.exists("/etc/ld.so.conf"):
|
||||||
|
paths.extend(parse_ld_conf_file("/etc/ld.so.conf"))
|
||||||
|
else:
|
||||||
|
print('WARNING: file "/etc/ld.so.conf" not found.')
|
||||||
|
if PREFIX:
|
||||||
|
if os.path.exists(PREFIX + "/etc/ld.so.conf"):
|
||||||
|
paths.extend(parse_ld_conf_file(PREFIX + "/etc/ld.so.conf"))
|
||||||
|
else:
|
||||||
|
print('WARNING: file "' + PREFIX + '/etc/ld.so.conf" not found.')
|
||||||
|
paths.extend(
|
||||||
|
[
|
||||||
|
PREFIX + "/lib",
|
||||||
|
PREFIX + "/usr/lib",
|
||||||
|
PREFIX + "/lib64",
|
||||||
|
PREFIX + "/usr/lib64",
|
||||||
|
]
|
||||||
|
)
|
||||||
|
paths.extend(["/lib", "/usr/lib", "/lib64", "/usr/lib64"])
|
||||||
|
return [path for path in paths if os.path.isdir(path)]
|
||||||
|
|
||||||
|
|
||||||
|
def resolve_libraries(
|
||||||
|
paths: List[str], files_patterns: List[str]
|
||||||
|
) -> Dict[str, ResolvedLib]:
|
||||||
|
"""Scans the PATH directory looking for the files complying with
|
||||||
|
the FILES_PATTERNS regexes list. Each file matching the pattern will be found only once
|
||||||
|
|
||||||
|
Returns the list of the resolved DSOs."""
|
||||||
|
libraries: Dict[str, ResolvedLib] = {}
|
||||||
|
|
||||||
def is_dso_matching_pattern(filename):
|
def is_dso_matching_pattern(filename):
|
||||||
for pattern in files_patterns:
|
for pattern in files_patterns:
|
||||||
|
@ -95,15 +224,19 @@ def find_files(path: str, files_patterns: List[str]) -> List[str]:
|
||||||
return True
|
return True
|
||||||
return False
|
return False
|
||||||
|
|
||||||
for f in os.listdir(path):
|
for path in paths:
|
||||||
abs_file_path = os.path.abspath(os.path.join(path, f))
|
for fname in os.listdir(path):
|
||||||
if os.path.isfile(abs_file_path) and is_dso_matching_pattern(abs_file_path):
|
abs_file_path = os.path.abspath(os.path.join(path, fname))
|
||||||
files.append(abs_file_path)
|
if (
|
||||||
|
os.path.isfile(abs_file_path)
|
||||||
return files
|
and is_dso_matching_pattern(abs_file_path)
|
||||||
|
and (fname not in libraries)
|
||||||
|
):
|
||||||
|
libraries[fname] = ResolvedLib(fname, abs_file_path)
|
||||||
|
return libraries
|
||||||
|
|
||||||
|
|
||||||
def copy_and_patch_libs(dsos: List[str], libs_dir: str, rpath=None) -> None:
|
def copy_and_patch_libs(dsos: List[ResolvedLib], libs_dir: str, rpath=None) -> None:
|
||||||
"""Copies the graphic vendor DSOs to the cache directory before
|
"""Copies the graphic vendor DSOs to the cache directory before
|
||||||
patchelf-ing them.
|
patchelf-ing them.
|
||||||
|
|
||||||
|
@ -117,11 +250,11 @@ def copy_and_patch_libs(dsos: List[str], libs_dir: str, rpath=None) -> None:
|
||||||
runpath to point to the cache directory."""
|
runpath to point to the cache directory."""
|
||||||
rpath = rpath if (rpath is not None) else libs_dir
|
rpath = rpath if (rpath is not None) else libs_dir
|
||||||
for dso in dsos:
|
for dso in dsos:
|
||||||
basename = os.path.basename(dso)
|
basename = os.path.basename(dso.fullpath)
|
||||||
newpath = os.path.join(libs_dir, basename)
|
newpath = os.path.join(libs_dir, basename)
|
||||||
log_info(f"Copying {basename} to {newpath}")
|
log_info(f"Copying and patching {dso} to {newpath}")
|
||||||
shutil.copyfile(dso, newpath)
|
shutil.copyfile(dso.fullpath, newpath)
|
||||||
shutil.copymode(dso, newpath)
|
shutil.copymode(dso.fullpath, newpath)
|
||||||
patch_dso(newpath, rpath)
|
patch_dso(newpath, rpath)
|
||||||
|
|
||||||
|
|
||||||
|
@ -146,7 +279,9 @@ def patch_dso(dsoPath: str, rpath: str) -> None:
|
||||||
# some loosely connected parts together for no good reason.
|
# some loosely connected parts together for no good reason.
|
||||||
|
|
||||||
|
|
||||||
def generate_nvidia_egl_config_files(cache_dir: str, libs_dir: str) -> str:
|
def generate_nvidia_egl_config_files(
|
||||||
|
cache_dir: str, libs_dir: str, egl_conf_dir: str
|
||||||
|
) -> str:
|
||||||
"""Generates a set of JSON files describing the EGL exec
|
"""Generates a set of JSON files describing the EGL exec
|
||||||
envirnoment to libglvnd.
|
envirnoment to libglvnd.
|
||||||
|
|
||||||
|
@ -158,8 +293,6 @@ def generate_nvidia_egl_config_files(cache_dir: str, libs_dir: str) -> str:
|
||||||
{"file_format_version": "1.0.0", "ICD": {"library_path": dso}}
|
{"file_format_version": "1.0.0", "ICD": {"library_path": dso}}
|
||||||
)
|
)
|
||||||
|
|
||||||
egl_conf_dir = os.path.join(cache_dir, "egl-confs")
|
|
||||||
os.makedirs(egl_conf_dir, exist_ok=True)
|
|
||||||
dso_paths = [
|
dso_paths = [
|
||||||
("10_nvidia.json", f"{libs_dir}/libEGL_nvidia.so.0"),
|
("10_nvidia.json", f"{libs_dir}/libEGL_nvidia.so.0"),
|
||||||
("10_nvidia_wayland.json", f"{libs_dir}/libnvidia-egl-wayland.so.1"),
|
("10_nvidia_wayland.json", f"{libs_dir}/libnvidia-egl-wayland.so.1"),
|
||||||
|
@ -176,27 +309,32 @@ def generate_nvidia_egl_config_files(cache_dir: str, libs_dir: str) -> str:
|
||||||
return egl_conf_dir
|
return egl_conf_dir
|
||||||
|
|
||||||
|
|
||||||
def exec_binary(bin_path: str, args: List[str]) -> None:
|
def is_dso_cache_up_to_date(dsos: HostDSOs, cache_file_path: str) -> bool:
|
||||||
"""Replace the current python program with the program pointed by
|
"""Check whether or not we need to udate the host DSOs cache.
|
||||||
BIN_PATH.
|
|
||||||
|
|
||||||
Sets the relevant libGLvnd env variables."""
|
We keep what's in the cache through a JSON file stored at the root
|
||||||
log_info(f"Execv-ing {bin_path}")
|
of the cache_dir. We consider a DSO to be up to date if its name
|
||||||
log_info(f"Goodbye now.")
|
and its content sha256 are equivalent.
|
||||||
# The following two env variables are required by our patched libglvnd
|
"""
|
||||||
# implementation to figure out what kind of driver the host
|
log_info("Checking if the cache is up to date")
|
||||||
# machine is using.
|
if os.path.isfile(cache_file_path):
|
||||||
os.execv(bin_path, [bin_path] + args)
|
with open(cache_file_path, "r", encoding="utf8") as f:
|
||||||
|
try:
|
||||||
|
cached_dsos: HostDSOs = HostDSOs.from_json(f.read())
|
||||||
|
except:
|
||||||
|
return False
|
||||||
|
return dsos == cached_dsos
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
def nvidia_main(cache_dir: str, gl_vendor_path: str) -> Dict:
|
def nvidia_main(cache_dir: str, dso_vendor_paths: List[str]) -> Dict:
|
||||||
"""Prepares the environment necessary to run a opengl/cuda program
|
"""Prepares the environment necessary to run a opengl/cuda program
|
||||||
on a Nvidia graphics card. It is by definition really stateful.
|
on a Nvidia graphics card. It is by definition really stateful.
|
||||||
|
|
||||||
Roughly, we're going to:
|
Roughly, we're going to:
|
||||||
|
|
||||||
1. Setup the nvidia cache directory.
|
1. Setup the nvidia cache directory.
|
||||||
2. Find the nvidia DSOs in the GL_VENDOR_PATH.
|
2. Find the nvidia DSOs in the DSO_VENDOR_PATH.
|
||||||
3. Copy these DSOs to their appropriate cache directories.
|
3. Copy these DSOs to their appropriate cache directories.
|
||||||
4. Generate the EGL configuration files.
|
4. Generate the EGL configuration files.
|
||||||
5. Patchelf the runpath of what needs to be patched.
|
5. Patchelf the runpath of what needs to be patched.
|
||||||
|
@ -219,45 +357,47 @@ def nvidia_main(cache_dir: str, gl_vendor_path: str) -> Dict:
|
||||||
This function returns a dictionary containing the env variables
|
This function returns a dictionary containing the env variables
|
||||||
supposed to be added to the current process down the line."""
|
supposed to be added to the current process down the line."""
|
||||||
log_info("Nvidia routine begins")
|
log_info("Nvidia routine begins")
|
||||||
|
log_info("Setting up Nvidia cache directory")
|
||||||
cache_dir = os.path.join(cache_dir, "nvidia")
|
cache_dir = os.path.join(cache_dir, "nvidia")
|
||||||
libs_dir = os.path.join(cache_dir, "lib")
|
libs_dir = os.path.join(cache_dir, "lib")
|
||||||
cuda_dir = os.path.join(cache_dir, "cuda")
|
cuda_dir = os.path.join(cache_dir, "cuda")
|
||||||
glx_dir = os.path.join(cache_dir, "glx")
|
glx_dir = os.path.join(cache_dir, "glx")
|
||||||
|
egl_dir = os.path.join(cache_dir, "egl-confs")
|
||||||
|
cache_file_path = os.path.join(cache_dir, "cache.json")
|
||||||
log_info(f"Nvidia libs dir: {libs_dir}")
|
log_info(f"Nvidia libs dir: {libs_dir}")
|
||||||
log_info(f"Nvidia cuda dir: {libs_dir}")
|
log_info(f"Nvidia cuda dir: {libs_dir}")
|
||||||
os.makedirs(libs_dir, exist_ok=True)
|
os.makedirs(libs_dir, exist_ok=True)
|
||||||
os.makedirs(cuda_dir, exist_ok=True)
|
os.makedirs(cuda_dir, exist_ok=True)
|
||||||
os.makedirs(glx_dir, exist_ok=True)
|
os.makedirs(glx_dir, exist_ok=True)
|
||||||
log_info(f"Searching for the Nvidia OpenGL DSOs in {gl_vendor_path}")
|
os.makedirs(egl_dir, exist_ok=True)
|
||||||
# Nvidia OpenGL DSOs
|
# Find Host DSOS
|
||||||
opengl_dsos = find_files(gl_vendor_path, NVIDIA_DSO_PATTERNS)
|
log_info("Searching for the host DSOs")
|
||||||
log_info(f"Found the following DSOs:")
|
dsos: HostDSOs = HostDSOs(
|
||||||
for dso in opengl_dsos:
|
generic=resolve_libraries(dso_vendor_paths, NVIDIA_DSO_PATTERNS),
|
||||||
log_info(dso)
|
cuda=resolve_libraries(dso_vendor_paths, CUDA_DSO_PATTERNS),
|
||||||
log_info("Patching the DSOs.")
|
glx=resolve_libraries(dso_vendor_paths, GLX_DSO_PATTERNS),
|
||||||
copy_and_patch_libs(opengl_dsos, libs_dir)
|
)
|
||||||
# Nvidia Cuda DSOs
|
log_info("Caching and patching host DSOs")
|
||||||
log_info(f"Searching for the Nvidia Cuda DSOs in {gl_vendor_path}")
|
# Cache/Patch DSOs
|
||||||
cuda_dsos = find_files(gl_vendor_path, CUDA_DSO_PATTERNS)
|
if not is_dso_cache_up_to_date(dsos, cache_file_path):
|
||||||
log_info(f"Found the following DSOs:")
|
log_info("The cache is not up to date, regenerating it")
|
||||||
for dso in cuda_dsos:
|
shutil.rmtree(cache_dir)
|
||||||
log_info(dso)
|
os.makedirs(libs_dir, exist_ok=True)
|
||||||
log_info("Patching the DSOs.")
|
os.makedirs(cuda_dir, exist_ok=True)
|
||||||
copy_and_patch_libs(cuda_dsos, cuda_dir, libs_dir)
|
os.makedirs(glx_dir, exist_ok=True)
|
||||||
# GLX DSOs
|
os.makedirs(egl_dir, exist_ok=True)
|
||||||
log_info(f"Searching for the Nvidia GLX DSOs in {gl_vendor_path}")
|
copy_and_patch_libs(list(dsos.generic.values()), libs_dir, libs_dir)
|
||||||
glx_dsos = find_files(gl_vendor_path, GLX_DSO_PATTERNS)
|
copy_and_patch_libs(list(dsos.glx.values()), glx_dir, libs_dir)
|
||||||
log_info(f"Found the following DSOs:")
|
copy_and_patch_libs(list(dsos.cuda.values()), cuda_dir, libs_dir)
|
||||||
for dso in glx_dsos:
|
log_info("Setting up NVIDIA-specific execution env variables.")
|
||||||
log_info(dso)
|
with open(cache_file_path, "w", encoding="utf8") as f:
|
||||||
log_info("Patching the DSOs.")
|
f.write(dsos.to_json())
|
||||||
copy_and_patch_libs(glx_dsos, glx_dir, libs_dir)
|
else:
|
||||||
# Preparing the env
|
log_info("The cache is up to date.")
|
||||||
log_info("Setting NVIDIA-specific env variables.")
|
egl_config_files = generate_nvidia_egl_config_files(cache_dir, libs_dir, egl_dir)
|
||||||
new_env = {}
|
new_env = {}
|
||||||
log_info(f"__GLX_VENDOR_LIBRARY_NAME = nvidia")
|
log_info(f"__GLX_VENDOR_LIBRARY_NAME = nvidia")
|
||||||
new_env["__GLX_VENDOR_LIBRARY_NAME"] = "nvidia"
|
new_env["__GLX_VENDOR_LIBRARY_NAME"] = "nvidia"
|
||||||
egl_config_files = generate_nvidia_egl_config_files(cache_dir, libs_dir)
|
|
||||||
log_info(f"__EGL_VENDOR_LIBRARY_DIRS = {egl_config_files}")
|
log_info(f"__EGL_VENDOR_LIBRARY_DIRS = {egl_config_files}")
|
||||||
new_env["__EGL_VENDOR_LIBRARY_DIRS"] = egl_config_files
|
new_env["__EGL_VENDOR_LIBRARY_DIRS"] = egl_config_files
|
||||||
ld_library_path = os.environ.get("LD_LIBRARY_PATH", None)
|
ld_library_path = os.environ.get("LD_LIBRARY_PATH", None)
|
||||||
|
@ -272,15 +412,30 @@ def nvidia_main(cache_dir: str, gl_vendor_path: str) -> Dict:
|
||||||
return new_env
|
return new_env
|
||||||
|
|
||||||
|
|
||||||
|
def exec_binary(bin_path: str, args: List[str]) -> None:
|
||||||
|
"""Replace the current python program with the program pointed by
|
||||||
|
BIN_PATH.
|
||||||
|
|
||||||
|
Sets the relevant libGLvnd env variables."""
|
||||||
|
log_info(f"Execv-ing {bin_path}")
|
||||||
|
log_info(f"Goodbye now.")
|
||||||
|
# The following two env variables are required by our patched libglvnd
|
||||||
|
# implementation to figure out what kind of driver the host
|
||||||
|
# machine is using.
|
||||||
|
os.execv(bin_path, [bin_path] + args)
|
||||||
|
|
||||||
|
|
||||||
def main(args):
|
def main(args):
|
||||||
|
start_time = time.time()
|
||||||
home = os.path.expanduser("~")
|
home = os.path.expanduser("~")
|
||||||
xdg_cache_home = os.environ.get("XDG_CACHE_HOME", os.path.join(home, ".cache"))
|
xdg_cache_home = os.environ.get("XDG_CACHE_HOME", os.path.join(home, ".cache"))
|
||||||
cache_dir = os.path.join(xdg_cache_home, "nix-gl-host")
|
cache_dir = os.path.join(xdg_cache_home, "nix-gl-host")
|
||||||
log_info(f'Using "{cache_dir}" as cache dir.')
|
log_info(f'Using "{cache_dir}" as cache dir.')
|
||||||
os.makedirs(cache_dir, exist_ok=True)
|
os.makedirs(cache_dir, exist_ok=True)
|
||||||
log_info(f'Scanning "{args.GL_VENDOR_PATH}" for DSOs.')
|
host_dsos_paths: List[str] = get_ld_paths()
|
||||||
new_env = nvidia_main(cache_dir, args.GL_VENDOR_PATH)
|
new_env = nvidia_main(cache_dir, host_dsos_paths)
|
||||||
os.environ.update(new_env)
|
os.environ.update(new_env)
|
||||||
|
log_info(f"{time.time() - start_time} seconds elapsed since script start.")
|
||||||
exec_binary(args.NIX_BINARY, args.ARGS)
|
exec_binary(args.NIX_BINARY, args.ARGS)
|
||||||
return 0
|
return 0
|
||||||
|
|
||||||
|
@ -290,11 +445,6 @@ if __name__ == "__main__":
|
||||||
prog="nixglhost-wrapper",
|
prog="nixglhost-wrapper",
|
||||||
description="Wrapper used to massage the host GL drivers to work with your nix-built binary.",
|
description="Wrapper used to massage the host GL drivers to work with your nix-built binary.",
|
||||||
)
|
)
|
||||||
parser.add_argument(
|
|
||||||
"GL_VENDOR_PATH",
|
|
||||||
type=str,
|
|
||||||
help="a path pointing to the directory containing your GL driver shared libraries",
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"NIX_BINARY",
|
"NIX_BINARY",
|
||||||
type=str,
|
type=str,
|
|
@ -0,0 +1,48 @@
|
||||||
|
import unittest
|
||||||
|
|
||||||
|
from nixglhost_wrapper import HostDSOs, ResolvedLib
|
||||||
|
|
||||||
|
|
||||||
|
class TestCacheSerializer(unittest.TestCase):
|
||||||
|
def hostdso_json_golden_test(self):
|
||||||
|
hds = HostDSOs(
|
||||||
|
glx={
|
||||||
|
"dummyglx.so": ResolvedLib(
|
||||||
|
"dummyglx.so",
|
||||||
|
"/lib/dummyglx.so",
|
||||||
|
"031edd7d41651593c5fe5c006fa5752b37fddff7bc4e843aa6af0c950f4b9406",
|
||||||
|
)
|
||||||
|
},
|
||||||
|
cuda={
|
||||||
|
"dummycuda.so": ResolvedLib(
|
||||||
|
"dummycuda.so",
|
||||||
|
"/lib/dummycuda.so",
|
||||||
|
"031edd7d41651593c5fe5c006fa5752b37fddff7bc4e843aa6af0c950f4b9407",
|
||||||
|
)
|
||||||
|
},
|
||||||
|
generic={
|
||||||
|
"dummygeneric.so": ResolvedLib(
|
||||||
|
"dummygeneric.so",
|
||||||
|
"/lib/dummygeneric.so",
|
||||||
|
"031edd7d41651593c5fe5c006fa5752b37fddff7bc4e843aa6af0c950f4b9408",
|
||||||
|
)
|
||||||
|
},
|
||||||
|
)
|
||||||
|
json_hds = hds.to_json()
|
||||||
|
self.assertIsNotNone(json_hds)
|
||||||
|
golden_hds = HostDSOs.from_json(json_hds)
|
||||||
|
self.assertEqual(hds, golden_hds)
|
||||||
|
self.assertEqual(hds.to_json(), golden_hds.to_json())
|
||||||
|
|
||||||
|
def test_eq_commut_jsons(self):
|
||||||
|
"""Checks that object equality is not sensible to JSON keys commutations"""
|
||||||
|
hds_json = '{"version": 1, "glx": {"dummyglx.so": {"name": "dummyglx.so", "fullpath": "/lib/dummyglx.so", "sha256": "031edd7d41651593c5fe5c006fa5752b37fddff7bc4e843aa6af0c950f4b9406"}}, "cuda": {"dummycuda.so": {"name": "dummycuda.so", "fullpath": "/lib/dummycuda.so", "sha256": "031edd7d41651593c5fe5c006fa5752b37fddff7bc4e843aa6af0c950f4b9407"}, "dummycuda2.so": {"name": "dummycuda2.so", "fullpath": "/lib/dummycuda2.so", "sha256": "131edd7d41651593c5fe5c006fa5752b37fddff7bc4e843aa6af0c950f4b9407"}}, "generic": {"dummygeneric.so": {"name": "dummygeneric.so", "fullpath": "/lib/dummygeneric.so", "sha256": "031edd7d41651593c5fe5c006fa5752b37fddff7bc4e843aa6af0c950f4b9408"}}}'
|
||||||
|
commut_hds_json = '{"version": 1, "glx": {"dummyglx.so": {"name": "dummyglx.so", "fullpath": "/lib/dummyglx.so", "sha256": "031edd7d41651593c5fe5c006fa5752b37fddff7bc4e843aa6af0c950f4b9406"}}, "cuda": {"dummycuda2.so": {"name": "dummycuda2.so", "fullpath": "/lib/dummycuda2.so", "sha256": "131edd7d41651593c5fe5c006fa5752b37fddff7bc4e843aa6af0c950f4b9407"}, "dummycuda.so": {"name": "dummycuda.so", "fullpath": "/lib/dummycuda.so", "sha256": "031edd7d41651593c5fe5c006fa5752b37fddff7bc4e843aa6af0c950f4b9407"}}, "generic": {"dummygeneric.so": {"name": "dummygeneric.so", "fullpath": "/lib/dummygeneric.so", "sha256": "031edd7d41651593c5fe5c006fa5752b37fddff7bc4e843aa6af0c950f4b9408"}}}'
|
||||||
|
hds = HostDSOs.from_json(hds_json)
|
||||||
|
commut_hds = HostDSOs.from_json(commut_hds_json)
|
||||||
|
self.assertEqual(hds, commut_hds)
|
||||||
|
self.assertEqual(hds.to_json(), commut_hds.to_json())
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
Loading…
Reference in New Issue