diff --git a/nix-builder/lib/build.nix b/nix-builder/lib/build.nix index 4d2e450e..757a81e2 100644 --- a/nix-builder/lib/build.nix +++ b/nix-builder/lib/build.nix @@ -157,6 +157,7 @@ rec { kernelProvenance ; kernelName = kernelConfig.name; + variant = variants.kernelVariant kernelConfig; } else if kernelConfig.isTvmFfi then extension.mkTvmFfiExtension { @@ -176,6 +177,7 @@ rec { kernelName = kernelConfig.name; doAbiCheck = true; + variant = variants.kernelVariant kernelConfig; } else extension.mkTorchExtension { @@ -197,6 +199,7 @@ rec { kernelName = kernelConfig.name; doAbiCheck = true; + variant = variants.kernelVariant kernelConfig; }; # Build multiple Torch extensions. diff --git a/nix-builder/lib/extension/torch/arch.nix b/nix-builder/lib/extension/torch/arch.nix index e28eae3d..daa077ac 100644 --- a/nix-builder/lib/extension/torch/arch.nix +++ b/nix-builder/lib/extension/torch/arch.nix @@ -72,6 +72,9 @@ # The Torch stable ABI version to check for. torchStableAbiVersion ? null, + # The variant that the build is expected to produce. + variant, + # Revision to bake into the ops name. rev, @@ -236,7 +239,10 @@ stdenv.mkDerivation (prevAttrs: { ++ extraDeps; env = - lib.optionalAttrs cudaSupport { + { + inherit variant; + } + // lib.optionalAttrs cudaSupport { CUDAToolkit_ROOT = "${lib.getDev cudaPackages.cuda_nvcc}"; } // lib.optionalAttrs xpuSupport { diff --git a/nix-builder/lib/extension/torch/no-arch.nix b/nix-builder/lib/extension/torch/no-arch.nix index a1b1d36d..d24f64aa 100644 --- a/nix-builder/lib/extension/torch/no-arch.nix +++ b/nix-builder/lib/extension/torch/no-arch.nix @@ -46,6 +46,9 @@ backendPythonDeps, + # The variant that the build is expected to produce. + variant, + # Git provenance (`{ sha; dirty; }`, or `null`) of the kernel source, recorded # in the build metadata. kernelProvenance, @@ -93,6 +96,10 @@ stdenv.mkDerivation (prevAttrs: { framework = "torch"; + env = { + inherit variant; + }; + # Add Torch as a dependency, so that devshells for universal kernels # also get torch as a build input. buildInputs = [ torch ]; diff --git a/nix-builder/lib/extension/tvm-ffi/arch.nix b/nix-builder/lib/extension/tvm-ffi/arch.nix index bfdd0604..de5a58dd 100644 --- a/nix-builder/lib/extension/tvm-ffi/arch.nix +++ b/nix-builder/lib/extension/tvm-ffi/arch.nix @@ -69,6 +69,9 @@ # Wheter to strip rpath for non-nix use. stripRPath ? false, + # The variant that the build is expected to produce. + variant, + # Revision to bake into the ops name. rev, @@ -234,7 +237,10 @@ stdenv.mkDerivation (prevAttrs: { ++ extraDeps; env = - lib.optionalAttrs cudaSupport { + { + inherit variant; + } + // lib.optionalAttrs cudaSupport { CUDAToolkit_ROOT = "${lib.getDev cudaPackages.cuda_nvcc}"; } // lib.optionalAttrs xpuSupport { diff --git a/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.py b/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.py index f93db419..497041a3 100755 --- a/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.py +++ b/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.py @@ -18,9 +18,25 @@ with open(kernelDeps) as f: kernel_paths = KernelPaths.from_json(f.read()) +variant = os.getenv("variant") +if not variant: + raise ValueError("`variant` environment variable is not set by Nix derivation") + kernel = KernelDependency(repo_id=out, version=KernelVersion.Version(0)) + +# The build host may not expose the accelerator being targeted (e.g. Metal is +# undetectable inside the sandboxed macOS build), so point the resolver at the +# variant the derivation is expected to produce instead of relying on backend +# auto-detection. This also fails the check if the build produced a different +# variant than expected. +variant_path = Path(out) / variant +if not variant_path.is_dir(): + raise FileNotFoundError( + f"Expected build variant `{variant}` is not present in `{out}`" + ) + resolvers = [ - RepoPathsResolver(local_kernels={out: Path(out)}), + RepoPathsResolver(local_kernels={out: variant_path}), KernelPathsResolver(kernel_paths=kernel_paths), ] diff --git a/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.sh b/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.sh index 6b765987..f37c15b5 100755 --- a/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.sh +++ b/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.sh @@ -15,6 +15,11 @@ _getKernelCheckHook() { exit 1 fi + if [ -z ${variant+x} ]; then + echo "variant must be set in derivation" + exit 1 + fi + echo "Check whether the kernel can be loaded with get-kernel: ${moduleName}" # We strip the full library paths from the extension. Unfortunately,