diff --git a/lib/torch-extension/arch.nix b/lib/torch-extension/arch.nix index 6e907029..0a4b4e02 100644 --- a/lib/torch-extension/arch.nix +++ b/lib/torch-extension/arch.nix @@ -13,6 +13,7 @@ cuda_nvcc, get-kernel-check, kernel-abi-check, + kernel-layout-check, ninja, python3, remove-bytecode-hook, @@ -79,7 +80,12 @@ in stdenv.mkDerivation (prevAttrs: { name = "${extensionName}-torch-ext"; - inherit doAbiCheck nvccThreads src; + inherit + doAbiCheck + extensionName + nvccThreads + src + ; # Generate build files. postPatch = '' @@ -123,10 +129,11 @@ stdenv.mkDerivation (prevAttrs: { ''; nativeBuildInputs = [ - kernel-abi-check cmake ninja build2cmake + kernel-abi-check + kernel-layout-check remove-bytecode-hook ] ++ lib.optionals doGetKernelCheck [ @@ -223,28 +230,31 @@ stdenv.mkDerivation (prevAttrs: { postInstall = '' ( cd .. - cp -r torch-ext/${extensionName} $out/ + cp -r torch-ext/${extensionName}/* $out/ ) - cp $out/_${extensionName}_*/* $out/${extensionName} - rm -rf $out/_${extensionName}_* + mv $out/_${extensionName}_*/* $out/ + rm -d $out/_${extensionName}_${rev} + + # Set up a compatibility module for older kernels versions, remove when + # the updated kernels has been around for a while. + mkdir $out/${extensionName} + cp ${./compat.py} $out/${extensionName}/__init__.py '' + (lib.optionalString (stripRPath && stdenv.hostPlatform.isLinux)) '' - find $out/${extensionName} -name '*.so' \ + find $out/ -name '*.so' \ -exec patchelf --set-rpath "" {} \; '' + (lib.optionalString (stripRPath && stdenv.hostPlatform.isDarwin)) '' - find $out/${extensionName} -name '*.so' \ + find $out/ -name '*.so' \ -exec rewrite-nix-paths-macho {} \; # Stub some rpath. - find $out/${extensionName} -name '*.so' \ + find $out/ -name '*.so' \ -exec install_name_tool -add_rpath "@loader_path/lib" {} \; ''; doInstallCheck = true; - getKernelCheck = extensionName; - # We need access to the host system on Darwin for the Metal compiler. __noChroot = metalSupport; diff --git a/lib/torch-extension/compat.py b/lib/torch-extension/compat.py new file mode 100644 index 00000000..03dbc1af --- /dev/null +++ b/lib/torch-extension/compat.py @@ -0,0 +1,26 @@ +import ctypes +import sys + +import importlib +from pathlib import Path +from types import ModuleType + +def _import_from_path(file_path: Path) -> ModuleType: + # We cannot use the module name as-is, after adding it to `sys.modules`, + # it would also be used for other imports. So, we make a module name that + # depends on the path for it to be unique using the hex-encoded hash of + # the path. + path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value) + module_name = path_hash + spec = importlib.util.spec_from_file_location(module_name, file_path) + if spec is None: + raise ImportError(f"Cannot load spec for {module_name} from {file_path}") + module = importlib.util.module_from_spec(spec) + if module is None: + raise ImportError(f"Cannot load module {module_name} from spec") + sys.modules[module_name] = module + spec.loader.exec_module(module) # type: ignore + return module + + +globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py"))) diff --git a/lib/torch-extension/no-arch.nix b/lib/torch-extension/no-arch.nix index c591c853..783a5ac4 100644 --- a/lib/torch-extension/no-arch.nix +++ b/lib/torch-extension/no-arch.nix @@ -4,6 +4,7 @@ build2cmake, get-kernel-check, + kernel-layout-check, remove-bytecode-hook, torch, }: @@ -23,7 +24,7 @@ stdenv.mkDerivation (prevAttrs: { name = "${extensionName}-torch-ext"; - inherit src; + inherit extensionName src; # Add Torch as a dependency, so that devshells for universal kernels # also get torch as a build input. @@ -31,6 +32,7 @@ stdenv.mkDerivation (prevAttrs: { nativeBuildInputs = [ build2cmake + kernel-layout-check remove-bytecode-hook ] ++ lib.optionals doGetKernelCheck [ @@ -48,10 +50,10 @@ stdenv.mkDerivation (prevAttrs: { installPhase = '' mkdir -p $out - cp -r torch-ext/${extensionName} $out/ + cp -r torch-ext/${extensionName}/* $out/ + mkdir $out/${extensionName} + cp ${./compat.py} $out/${extensionName}/__init__.py ''; doInstallCheck = true; - - getKernelCheck = extensionName; }) diff --git a/overlay.nix b/overlay.nix index 97976395..ccd54c75 100644 --- a/overlay.nix +++ b/overlay.nix @@ -9,6 +9,8 @@ final: prev: { kernel-abi-check = prev.callPackage ./pkgs/kernel-abi-check { }; + kernel-layout-check = prev.callPackage ./pkgs/kernel-layout-check { }; + rewrite-nix-paths-macho = prev.callPackage ./pkgs/rewrite-nix-paths-macho { }; remove-bytecode-hook = prev.callPackage ./pkgs/remove-bytecode-hook { }; diff --git a/pkgs/get-kernel-check/get-kernel-check-hook.sh b/pkgs/get-kernel-check/get-kernel-check-hook.sh index d59810b5..73915f7f 100755 --- a/pkgs/get-kernel-check/get-kernel-check-hook.sh +++ b/pkgs/get-kernel-check/get-kernel-check-hook.sh @@ -3,35 +3,39 @@ echo "Sourcing get-kernel-check-hook.sh" _getKernelCheckHook() { - if [ ! -z "${getKernelCheck}" ]; then - echo "Checking loading kernel with get_kernel" - echo "Check whether the kernel can be loaded with get-kernel: ${getKernelCheck}" - - # We strip the full library paths from the extension. Unfortunately, - # in a Nix environment, the library dependencies cannot be found - # anymore. So we have to add the Torch library directory to the - # dynamic linker path to get it to pick it up. - if [ $(uname -s) == "Darwin" ]; then - TORCH_DIR=$(python -c "from pathlib import Path; import torch; print(Path(torch.__file__).parent)") - export DYLD_LIBRARY_PATH="${TORCH_DIR}/lib:${DYLD_LIBRARY_PATH}" - fi - - TMPDIR=$(mktemp -d -t test.XXXXXX) || exit 1 - trap "rm -rf '$TMPDIR'" EXIT - - # Some kernels want to write stuff (especially when they use Triton). - HOME=$(mktemp -d -t test.XXXXXX) || exit 1 - trap "rm -rf '$HOME'" EXIT - - # Emulate the bundle layout that kernels expects. This even works - # for universal kernels, since kernels checks the non-universal - # path first. - BUILD_VARIANT=$(python -c "from kernels.utils import build_variant; print(build_variant())") - mkdir -p "${TMPDIR}/build" - ln -s "$out" "${TMPDIR}/build/${BUILD_VARIANT}" - - python -c "from pathlib import Path; import kernels; kernels.get_local_kernel(Path('${TMPDIR}'), '${getKernelCheck}')" + echo "Checking loading kernel with get_kernel" + + if [ -z ${extensionName+x} ]; then + echo "extensionName must be set in derivation" + exit 1 + fi + + echo "Check whether the kernel can be loaded with get-kernel: ${extensionName}" + + # We strip the full library paths from the extension. Unfortunately, + # in a Nix environment, the library dependencies cannot be found + # anymore. So we have to add the Torch library directory to the + # dynamic linker path to get it to pick it up. + if [ $(uname -s) == "Darwin" ]; then + TORCH_DIR=$(python -c "from pathlib import Path; import torch; print(Path(torch.__file__).parent)") + export DYLD_LIBRARY_PATH="${TORCH_DIR}/lib:${DYLD_LIBRARY_PATH}" fi + + TMPDIR=$(mktemp -d -t test.XXXXXX) || exit 1 + trap "rm -rf '$TMPDIR'" EXIT + + # Some kernels want to write stuff (especially when they use Triton). + HOME=$(mktemp -d -t test.XXXXXX) || exit 1 + trap "rm -rf '$HOME'" EXIT + + # Emulate the bundle layout that kernels expects. This even works + # for universal kernels, since kernels checks the non-universal + # path first. + BUILD_VARIANT=$(python -c "from kernels.utils import build_variant; print(build_variant())") + mkdir -p "${TMPDIR}/build" + ln -s "$out" "${TMPDIR}/build/${BUILD_VARIANT}" + + python -c "from pathlib import Path; import kernels; kernels.get_local_kernel(Path('${TMPDIR}'), '${extensionName}')" } postInstallCheckHooks+=(_getKernelCheckHook) diff --git a/pkgs/kernel-abi-check/kernel-abi-check-hook.sh b/pkgs/kernel-abi-check/kernel-abi-check-hook.sh index de7cf672..3bdcda6f 100755 --- a/pkgs/kernel-abi-check/kernel-abi-check-hook.sh +++ b/pkgs/kernel-abi-check/kernel-abi-check-hook.sh @@ -5,7 +5,7 @@ _checkAbiHook() { echo "Skipping ABI check" else echo "Checking of ABI compatibility" - find "$out/${extensionName}" -name '*.so' -print0 | \ + find "$out/" -name '*.so' -print0 | \ xargs -0 -n1 kernel-abi-check fi } diff --git a/pkgs/kernel-layout-check/default.nix b/pkgs/kernel-layout-check/default.nix new file mode 100644 index 00000000..259fa33c --- /dev/null +++ b/pkgs/kernel-layout-check/default.nix @@ -0,0 +1,5 @@ +{ makeSetupHook, python3 }: + +makeSetupHook { + name = "kernel-layout-check-hook"; +} ./kernel-layout-check-hook.sh diff --git a/pkgs/kernel-layout-check/kernel-layout-check-hook.sh b/pkgs/kernel-layout-check/kernel-layout-check-hook.sh new file mode 100755 index 00000000..1a72b553 --- /dev/null +++ b/pkgs/kernel-layout-check/kernel-layout-check-hook.sh @@ -0,0 +1,27 @@ +#!/bin/sh + +echo "Sourcing kernel-layout-check-hook.sh" + +kernelLayoutCheckHook() { + echo "Checking kernel layout" + + if [ -z ${extensionName+x} ]; then + echo "extensionName must be set in derivation" + exit 1 + fi + + if [ ! -f source/torch-ext/${extensionName}/__init__.py ]; then + echo "Python module at source/torch-ext/${extensionName} must contain __init__.py" + exit 1 + fi + + # TODO: remove once the old location is removed from kernels. + if [ -e source/torch-ext/${extensionName}/${extensionName} ]; then + echo "Python module at source/torch-ext/${extensionName} must not have ${extensionName} file or directory." + exit 1 + fi +} + +if [ -z "${dontCheckLayout-}" ]; then + postUnpackHooks+=(kernelLayoutCheckHook) +fi