fix: point get-kernel-check at the built variant

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.

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 {

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 ];

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 {

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),
 ]
 

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,