From 1ca637bbe597a0a317b985304750e1a2ff0acca2 Mon Sep 17 00:00:00 2001 From: cbensimon Date: Tue, 7 Apr 2026 16:48:11 +0000 Subject: [PATCH 1/8] get_laoded_kernels --- kernels/src/kernels/__init__.py | 2 ++ kernels/src/kernels/utils.py | 43 +++++++++++++++++++++++++++++++-- 2 files changed, 43 insertions(+), 2 deletions(-) diff --git a/kernels/src/kernels/__init__.py b/kernels/src/kernels/__init__.py index 248f712f..17f7cef8 100644 --- a/kernels/src/kernels/__init__.py +++ b/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", diff --git a/kernels/src/kernels/utils.py b/kernels/src/kernels/utils.py index fbe4dbce..16d36694 100644 --- a/kernels/src/kernels/utils.py +++ b/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,30 @@ KNOWN_BACKENDS = {"cpu", "cuda", "metal", "neuron", "rocm", "xpu", "npu"} +class RepoInfos(NamedTuple): + repo_id: str + revision: str | None + version: int | str | None + backend: str | None + + +class LoadedKernel(NamedTuple): + op_namespace: str + variant_path: Path + package_name: str + module_name: str + repo_infos: RepoInfos | None + + +_loaded_kernels: dict[str, LoadedKernel] = {} + + +def get_loaded_kernels() -> dict[str, LoadedKernel]: + """Returns the `op_namespace -> LoadedKernel` mapping + """ + return _loaded_kernels.copy() + + def _get_cache_dir() -> str | None: """Returns the kernels cache directory.""" cache_dir = os.environ.get("HF_KERNELS_CACHE", None) @@ -71,7 +96,7 @@ 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: metadata = Metadata.load_from_variant(variant_path) validate_dependencies(module_name, metadata.python_depends, _backend()) @@ -83,6 +108,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 +119,13 @@ 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[module.op._namespace] = LoadedKernel( + op_namespace=module.op._namespace, + variant_path=variant_path, + package_name=package_name, + module_name=module_name, + repo_infos=_repo_infos, + ) return module @@ -282,7 +315,13 @@ def get_kernel( package_name, variant_path = install_kernel( repo_id, revision=revision, backend=backend, user_agent=user_agent ) - return _import_from_path(package_name, variant_path) + repo_infos = RepoInfos( + repo_id=repo_id, + revision=revision, + version=version, + backend=backend, + ) + return _import_from_path(package_name, variant_path, _repo_infos=repo_infos) def get_local_kernel( From 8d08c8d53ab5fc2f9fe0be99a0314f99440df14c Mon Sep 17 00:00:00 2001 From: cbensimon Date: Tue, 7 Apr 2026 17:06:50 +0000 Subject: [PATCH 2/8] Fix op_namespace --- kernels/src/kernels/utils.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/kernels/src/kernels/utils.py b/kernels/src/kernels/utils.py index 16d36694..2f85f8a1 100644 --- a/kernels/src/kernels/utils.py +++ b/kernels/src/kernels/utils.py @@ -119,8 +119,9 @@ def _import_from_path(module_name: str, variant_path: Path, _repo_infos: RepoInf raise ImportError(f"Cannot load module {module_name} from spec") sys.modules[module_name] = module spec.loader.exec_module(module) # type: ignore - _loaded_kernels[module.op._namespace] = LoadedKernel( - op_namespace=module.op._namespace, + op_namespace = sys.modules[f"{module_name}._ops"].ops.name + _loaded_kernels[op_namespace] = LoadedKernel( + op_namespace=op_namespace, variant_path=variant_path, package_name=package_name, module_name=module_name, From b70f7461da92a7796bd869498c63c627cd8e825a Mon Sep 17 00:00:00 2001 From: cbensimon Date: Tue, 7 Apr 2026 17:09:33 +0000 Subject: [PATCH 3/8] Improve docstring --- kernels/src/kernels/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/kernels/src/kernels/utils.py b/kernels/src/kernels/utils.py index 2f85f8a1..8018f4c8 100644 --- a/kernels/src/kernels/utils.py +++ b/kernels/src/kernels/utils.py @@ -71,7 +71,7 @@ def _get_cache_dir() -> str | None: def _get_local_kernel_overrides() -> dict[str, Path]: - """Returns list local overrides for kernels.""" + """Returns a copy of the loaded kernels registry (`op_namespace -> LoadedKernel` mapping).""" local_kerels = os.environ.get("LOCAL_KERNELS", None) if local_kerels is None: return dict() From c99761b1886f8aaa196e3f31dc25628e3f4aa402 Mon Sep 17 00:00:00 2001 From: cbensimon Date: Tue, 7 Apr 2026 17:36:37 +0000 Subject: [PATCH 4/8] Only for custom-op kernels --- kernels/src/kernels/utils.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/kernels/src/kernels/utils.py b/kernels/src/kernels/utils.py index 8018f4c8..0e3c536a 100644 --- a/kernels/src/kernels/utils.py +++ b/kernels/src/kernels/utils.py @@ -119,14 +119,16 @@ def _import_from_path(module_name: str, variant_path: Path, _repo_infos: RepoInf raise ImportError(f"Cannot load module {module_name} from spec") sys.modules[module_name] = module spec.loader.exec_module(module) # type: ignore - op_namespace = sys.modules[f"{module_name}._ops"].ops.name - _loaded_kernels[op_namespace] = LoadedKernel( - op_namespace=op_namespace, - variant_path=variant_path, - package_name=package_name, - module_name=module_name, - repo_infos=_repo_infos, - ) + # Only track kernels that register custom torch ops + if (ops := sys.modules.get(f"{module_name}._ops")) is not None: + op_namespace = ops.ops.name + _loaded_kernels[op_namespace] = LoadedKernel( + op_namespace=op_namespace, + variant_path=variant_path, + package_name=package_name, + module_name=module_name, + repo_infos=_repo_infos, + ) return module From 5e2513c4be68e1daaa4785d69cd92044ed67829b Mon Sep 17 00:00:00 2001 From: cbensimon Date: Tue, 7 Apr 2026 17:36:44 +0000 Subject: [PATCH 5/8] Revert "Improve docstring" This reverts commit b70f7461da92a7796bd869498c63c627cd8e825a. --- kernels/src/kernels/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/kernels/src/kernels/utils.py b/kernels/src/kernels/utils.py index 0e3c536a..deac6f0b 100644 --- a/kernels/src/kernels/utils.py +++ b/kernels/src/kernels/utils.py @@ -71,7 +71,7 @@ def _get_cache_dir() -> str | None: def _get_local_kernel_overrides() -> dict[str, Path]: - """Returns a copy of the loaded kernels registry (`op_namespace -> LoadedKernel` mapping).""" + """Returns list local overrides for kernels.""" local_kerels = os.environ.get("LOCAL_KERNELS", None) if local_kerels is None: return dict() From e10d8af2724a40ac3007856c2df37fc4b796bc8b Mon Sep 17 00:00:00 2001 From: cbensimon Date: Tue, 7 Apr 2026 17:37:25 +0000 Subject: [PATCH 6/8] Improve doc on the right function --- kernels/src/kernels/utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/kernels/src/kernels/utils.py b/kernels/src/kernels/utils.py index deac6f0b..02fedb5c 100644 --- a/kernels/src/kernels/utils.py +++ b/kernels/src/kernels/utils.py @@ -53,8 +53,7 @@ class LoadedKernel(NamedTuple): def get_loaded_kernels() -> dict[str, LoadedKernel]: - """Returns the `op_namespace -> LoadedKernel` mapping - """ + """Returns a copy of the loaded kernels registry (`op_namespace -> LoadedKernel` mapping).""" return _loaded_kernels.copy() From 0606f7068a16466853b3b422971c23fbbdabaab2 Mon Sep 17 00:00:00 2001 From: cbensimon Date: Tue, 7 Apr 2026 17:38:16 +0000 Subject: [PATCH 7/8] Black --- kernels/src/kernels/utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/kernels/src/kernels/utils.py b/kernels/src/kernels/utils.py index 02fedb5c..df59a2c2 100644 --- a/kernels/src/kernels/utils.py +++ b/kernels/src/kernels/utils.py @@ -95,7 +95,9 @@ 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, _repo_infos: RepoInfos | None = None) -> ModuleType: +def _import_from_path( + module_name: str, variant_path: Path, _repo_infos: RepoInfos | None = None +) -> ModuleType: metadata = Metadata.load_from_variant(variant_path) validate_dependencies(module_name, metadata.python_depends, _backend()) From 12d86e4d96cf5ea36307571e3fb6b352661aa39d Mon Sep 17 00:00:00 2001 From: cbensimon Date: Wed, 8 Apr 2026 10:01:59 +0000 Subject: [PATCH 8/8] Index by module_name + best-effort op namespace --- kernels/src/kernels/utils.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/kernels/src/kernels/utils.py b/kernels/src/kernels/utils.py index df59a2c2..11feaa56 100644 --- a/kernels/src/kernels/utils.py +++ b/kernels/src/kernels/utils.py @@ -42,10 +42,10 @@ class RepoInfos(NamedTuple): class LoadedKernel(NamedTuple): - op_namespace: str variant_path: Path package_name: str module_name: str + torch_namespace: str | None repo_infos: RepoInfos | None @@ -53,7 +53,7 @@ class LoadedKernel(NamedTuple): def get_loaded_kernels() -> dict[str, LoadedKernel]: - """Returns a copy of the loaded kernels registry (`op_namespace -> LoadedKernel` mapping).""" + """Returns a copy of the loaded kernels registry (`module_name -> LoadedKernel` mapping).""" return _loaded_kernels.copy() @@ -120,16 +120,16 @@ def _import_from_path( raise ImportError(f"Cannot load module {module_name} from spec") sys.modules[module_name] = module spec.loader.exec_module(module) # type: ignore - # Only track kernels that register custom torch ops + torch_namespace: str | None = None if (ops := sys.modules.get(f"{module_name}._ops")) is not None: - op_namespace = ops.ops.name - _loaded_kernels[op_namespace] = LoadedKernel( - op_namespace=op_namespace, - variant_path=variant_path, - package_name=package_name, - module_name=module_name, - repo_infos=_repo_infos, - ) + torch_namespace = getattr(ops.ops, "name", None) + _loaded_kernels[module_name] = LoadedKernel( + torch_namespace=torch_namespace, + variant_path=variant_path, + package_name=package_name, + module_name=module_name, + repo_infos=_repo_infos, + ) return module