This commit is contained in:
Félix Baylac Jacqué 2023-01-17 12:51:03 +01:00
parent 4f51b37c29
commit 4901b21072
No known key found for this signature in database
GPG Key ID: EFD315F31848DBA4
1 changed files with 23 additions and 13 deletions

View File

@ -6,9 +6,10 @@ import json
import os import os
import re import re
import shutil import shutil
import subprocess
import stat import stat
import subprocess
import sys import sys
import tempfile
import time import time
from glob import glob from glob import glob
from typing import List, Literal, Dict, Tuple, TypedDict, TextIO, Optional from typing import List, Literal, Dict, Tuple, TypedDict, TextIO, Optional
@ -412,7 +413,9 @@ def scan_dsos_from_dir(path: str) -> Optional[LibraryPath]:
return None return None
def cache_library_path(library_path: LibraryPath, cache_dir_root: str) -> str: def cache_library_path(
library_path: LibraryPath, temp_cache_dir_root: str, final_cache_dir_root: str
) -> str:
"""Generate a cache directory for the LIBRARY_PATH host directory. """Generate a cache directory for the LIBRARY_PATH host directory.
This cache directory is mirroring the host directory containing This cache directory is mirroring the host directory containing
@ -426,8 +429,9 @@ def cache_library_path(library_path: LibraryPath, cache_dir_root: str) -> str:
h.update(library_path.path.encode("utf8")) h.update(library_path.path.encode("utf8"))
path_hash: str = h.hexdigest() path_hash: str = h.hexdigest()
# Paths # Paths
cache_path_root: str = os.path.join(cache_dir_root, path_hash) cache_path_root: str = os.path.join(temp_cache_dir_root, path_hash)
lib_dir = os.path.join(cache_path_root, "lib") lib_dir = os.path.join(cache_path_root, "lib")
rpath_lib_dir = os.path.join(final_cache_dir_root, path_hash, "lib")
cuda_dir = os.path.join(cache_path_root, "cuda") cuda_dir = os.path.join(cache_path_root, "cuda")
egl_dir = os.path.join(cache_path_root, "egl") egl_dir = os.path.join(cache_path_root, "egl")
glx_dir = os.path.join(cache_path_root, "glx") glx_dir = os.path.join(cache_path_root, "glx")
@ -440,7 +444,7 @@ def cache_library_path(library_path: LibraryPath, cache_dir_root: str) -> str:
]: ]:
os.makedirs(d, exist_ok=True) os.makedirs(d, exist_ok=True)
if len(dsos) > 0: if len(dsos) > 0:
copy_and_patch_libs(dsos=dsos, dest_dir=d, rpath=lib_dir) copy_and_patch_libs(dsos=dsos, dest_dir=d, rpath=rpath_lib_dir)
else: else:
log_info(f"Did not find any DSO to put in {d}, skipping copy and patching.") log_info(f"Did not find any DSO to put in {d}, skipping copy and patching.")
return path_hash return path_hash
@ -534,15 +538,21 @@ def nvidia_main(
cache_content, cache_file_path cache_content, cache_file_path
) or not os.path.isfile(cached_ld_library_path): ) or not os.path.isfile(cached_ld_library_path):
log_info("The cache is not up to date, regenerating it") log_info("The cache is not up to date, regenerating it")
shutil.rmtree(cache_dir) with tempfile.TemporaryDirectory() as tmp_cache:
cache_paths: List[str] = [] new_cache = os.path.join(tmp_cache, "nix-gl-host")
for p in cache_content.paths: os.makedirs(new_cache)
log_info(f"Caching {p}") cache_paths: List[str] = []
cache_paths.append(cache_library_path(p, cache_dir)) for p in cache_content.paths:
cache_absolute_paths = [os.path.join(cache_dir, p) for p in cache_paths] log_info(f"Caching {p}")
nix_gl_ld_library_path = generate_cache_metadata( cache_paths.append(cache_library_path(p, new_cache, cache_dir))
cache_dir, cache_content, cache_absolute_paths # Pointing the LD_LIBRARY_PATH to the final destination
) # instead of the tmp dir.
cache_absolute_paths = [os.path.join(cache_dir, p) for p in cache_paths]
nix_gl_ld_library_path = generate_cache_metadata(
cache_dir, cache_content, cache_absolute_paths
)
log_info(f"Mv {new_cache} to {cache_dir}")
shutil.move(new_cache, os.path.split(cache_dir)[0])
else: else:
log_info("The cache is up to date, re-using it.") log_info("The cache is up to date, re-using it.")
with open(cached_ld_library_path, "r", encoding="utf8") as f: with open(cached_ld_library_path, "r", encoding="utf8") as f: