feat: get_loaded_kernels() (#428)

* get_laoded_kernels

* Fix op_namespace

* Improve docstring

* Only for custom-op kernels

* Revert "Improve docstring"

This reverts commit b70f746.

* Improve doc on the right function

* Black

* Index by module_name + best-effort op namespace

* Only return values + get_kernel caching

* Docstring

* revision can't be None

* op namespace from .so name (works for both torch and tvm)

* module version + addressing comments

Author: cbensimon

kernels/src/kernels/__init__.py

@@ -23,6 +23,7 @@
 )
 from kernels.utils import (
     get_kernel,
+    get_loaded_kernels,
     get_local_kernel,
     get_locked_kernel,
     has_kernel,
@@ -45,6 +46,7 @@
     "LockedLayerRepository",
     "Mode",
     "get_kernel",
+    "get_loaded_kernels",
     "get_local_kernel",
     "get_locked_kernel",
     "has_kernel",

kernels/src/kernels/utils.py

@@ -12,6 +12,7 @@
 from importlib.metadata import Distribution
 from pathlib import Path
 from types import ModuleType
+from typing import NamedTuple
 
 from huggingface_hub import HfApi, constants
 
@@ -33,6 +34,26 @@
 KNOWN_BACKENDS = {"cpu", "cuda", "metal", "neuron", "rocm", "xpu", "npu"}
 
 
+class RepoInfos(NamedTuple):
+    repo_id: str
+    revision: str
+    backend: str | None
+
+
+class LoadedKernel(NamedTuple):
+    module: ModuleType
+    package_name: str
+    repo_infos: RepoInfos | None
+
+
+_loaded_kernels: dict[Path, LoadedKernel] = {}
+
+
+def get_loaded_kernels() -> list[LoadedKernel]:
+    """Returns a copy of the loaded kernels registry (see `kernels.utils.LoadedKernel` NamedTuple)."""
+    return list(_loaded_kernels.values())
+
+
 def _get_cache_dir() -> str | None:
     """Returns the kernels cache directory."""
     cache_dir = os.environ.get("HF_KERNELS_CACHE", None)
@@ -71,7 +92,12 @@ def _parse_local_kernel_overrides(local_kernels: str) -> dict[str, Path]:
 CACHE_DIR: str | None = _get_cache_dir()
 
 
-def _import_from_path(module_name: str, variant_path: Path) -> ModuleType:
+def _import_from_path(
+    module_name: str, variant_path: Path, _repo_infos: RepoInfos | None = None
+) -> ModuleType:
+    if (loaded_kernel := _loaded_kernels.get(variant_path)) is not None:
+        return loaded_kernel.module
+
     metadata = Metadata.load_from_variant(variant_path)
     validate_dependencies(module_name, metadata.python_depends, _backend())
 
@@ -83,6 +109,7 @@ def _import_from_path(module_name: str, variant_path: Path) -> ModuleType:
     # 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.
+    package_name = module_name
     path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path)).value)
     module_name = f"{module_name}_{path_hash}"
     spec = importlib.util.spec_from_file_location(module_name, file_path)
@@ -93,6 +120,12 @@ def _import_from_path(module_name: str, variant_path: Path) -> ModuleType:
         raise ImportError(f"Cannot load module {module_name} from spec")
     sys.modules[module_name] = module
     spec.loader.exec_module(module)  # type: ignore
+
+    _loaded_kernels[variant_path] = LoadedKernel(
+        module=module,
+        package_name=package_name,
+        repo_infos=_repo_infos,
+    )
     return module
 
 
@@ -279,10 +312,15 @@ def get_kernel(
         return get_local_kernel(override, package_name_from_repo_id(repo_id))
 
     revision = select_revision_or_version(repo_id, revision=revision, version=version)
+    repo_infos = RepoInfos(
+        repo_id=repo_id,
+        revision=revision,
+        backend=backend,
+    )
     package_name, variant_path = install_kernel(
         repo_id, revision=revision, backend=backend, user_agent=user_agent
     )
-    return _import_from_path(package_name, variant_path)
+    return _import_from_path(package_name, variant_path, _repo_infos=repo_infos)
 
 
 def get_local_kernel(