Support DISABLE_KERNEL_MAPPING env var for completely disabling kernel mappings (#70)

* Disable kernel mappings with `DISABLE_KERNEL_MAPPING=1`

* Rename HF_KERNELS_CACHE to KERNELS_CACHE

But still recognize the old variant for compatibility.

* Add documentation for environment variables

Author: danieldk

README.md

@@ -47,6 +47,7 @@ the Hub.
 
 - [Using layers](docs/layers.md)
 - [Locking kernel versions](docs/locking.md)
+- [Environment variables](docs/env.md)
 - [Using kernels in a Docker container](docs/docker.md)
 - [Kernel requirements](docs/kernel-requirements.md)
 - [Writing kernels](https://github.com/huggingface/kernel-builder/blob/main/docs/writing-kernels.md) using [kernel-builder](https://github.com/huggingface/kernel-builder/)

docs/env.md

@@ -0,0 +1,10 @@
+# Environment variables
+
+## `KERNELS_CACHE`
+
+The directory to use as the local kernel cache. If not set, the cache
+of the `huggingface_hub` package is used.
+
+## `DISABLE_KERNEL_MAPPING`
+
+Disables kernel mappings for [`layers`](layers.md).

src/kernels/layer.py

@@ -1,4 +1,5 @@
 import inspect
+import os
 import warnings
 from contextvars import ContextVar
 from copy import deepcopy
@@ -10,6 +11,8 @@
 if TYPE_CHECKING:
     from torch import nn
 
+_DISABLE_KERNEL_MAPPING: bool = bool(int(os.environ.get("DISABLE_KERNEL_MAPPING", "0")))
+
 
 @dataclass(frozen=True)
 class Device:
@@ -131,6 +134,9 @@ def replace_kernel_forward_from_hub(cls, layer_name: str, *, use_fallback: bool
     cached_forward: Dict[LayerRepository, Callable] = {}
 
     def forward(self, x, *args, **kwargs):
+        if _DISABLE_KERNEL_MAPPING:
+            return fallback_forward(self, x, *args, **kwargs)
+
         kernel = _KERNEL_MAPPING.get().get(layer_name)
         if kernel is None:
             warnings.warn(

src/kernels/utils.py

@@ -4,6 +4,7 @@
 import importlib.metadata
 import inspect
 import json
+import logging
 import os
 import platform
 import sys
@@ -17,7 +18,20 @@
 
 from kernels.lockfile import KernelLock, VariantLock
 
-CACHE_DIR: Optional[str] = os.environ.get("HF_KERNELS_CACHE", None)
+
+def _get_cache_dir() -> Optional[str]:
+    """Returns the kernels cache directory."""
+    cache_dir = os.environ.get("HF_KERNELS_CACHE", None)
+    if cache_dir is not None:
+        logging.warning(
+            "HF_KERNELS_CACHE will be removed in the future, use KERNELS_CACHE instead"
+        )
+        return cache_dir
+
+    return os.environ.get("KERNELS_CACHE", None)
+
+
+CACHE_DIR: Optional[str] = _get_cache_dir()
 
 
 def build_variant() -> str: