diff --git a/kernels/src/kernels/__init__.py b/kernels/src/kernels/__init__.py index 861afdc15..3dae0dc9a 100644 --- a/kernels/src/kernels/__init__.py +++ b/kernels/src/kernels/__init__.py @@ -7,6 +7,8 @@ from kernels._windows import _add_additional_dll_paths from kernels.benchmark import Benchmark from kernels.hf_hub import RepoInfo +from kernels.importer import LoadedKernel, get_loaded_kernels +from kernels.install import install_kernel from kernels.layer import ( CUDAProperties, Device, @@ -26,20 +28,17 @@ use_kernel_mapping, use_kernelized_func, ) -from kernels.utils import ( - LoadedKernel, +from kernels.load import ( get_kernel, - get_kernel_variants, - get_loaded_kernels, get_local_kernel, get_locked_kernel, has_kernel, - install_kernel, load_kernel, ) from kernels.variants import ( VariantAccepted, VariantRejected, + get_kernel_variants, ) _add_additional_dll_paths() diff --git a/kernels/src/kernels/cli/__init__.py b/kernels/src/kernels/cli/__init__.py index c8bb7cff4..55811d8b5 100644 --- a/kernels/src/kernels/cli/__init__.py +++ b/kernels/src/kernels/cli/__init__.py @@ -8,11 +8,11 @@ from kernels.cli.verify_signature import verify_signature from kernels.cli.versions import print_kernel_versions from kernels.compat import tomllib -from kernels.locking import KernelLock, get_kernel_locks -from kernels.utils import ( +from kernels.install import ( install_kernel, install_kernel_all_variants, ) +from kernels.locking import KernelLock, get_kernel_locks def main(): diff --git a/kernels/src/kernels/cli/benchmark.py b/kernels/src/kernels/cli/benchmark.py index 2aa82a728..5e601d582 100644 --- a/kernels/src/kernels/cli/benchmark.py +++ b/kernels/src/kernels/cli/benchmark.py @@ -20,9 +20,9 @@ hf_raise_for_status, ) +from kernels.backends import _backend from kernels.benchmark import Benchmark from kernels.hf_hub import _get_hf_api -from kernels.utils import _backend MISSING_DEPS: list[str] = [] diff --git a/kernels/src/kernels/cli/verify_signature.py b/kernels/src/kernels/cli/verify_signature.py index 052559d59..5afe2550d 100644 --- a/kernels/src/kernels/cli/verify_signature.py +++ b/kernels/src/kernels/cli/verify_signature.py @@ -7,7 +7,7 @@ from typing_extensions import assert_never from kernels._versions import select_revision_or_version -from kernels.utils import install_kernel, install_kernel_all_variants +from kernels.install import install_kernel, install_kernel_all_variants from kernels.variants import get_variants_local from kernels.verify import VerificationResult, verify_variant diff --git a/kernels/src/kernels/deps.py b/kernels/src/kernels/deps.py index 41188dc41..4723f418d 100644 --- a/kernels/src/kernels/deps.py +++ b/kernels/src/kernels/deps.py @@ -4,8 +4,9 @@ from pathlib import Path from huggingface_hub.dataclasses import strict +from kernels_data import Metadata -from kernels.backends import Backend +from kernels.backends import Backend, _backend @strict @@ -59,6 +60,11 @@ def from_dict(data: dict) -> "DependencyData": raise FileNotFoundError("Cannot load dependency data, is `kernels` correctly installed?") +def validate_variant_dependencies(variant_path: Path) -> None: + metadata = Metadata.read_from_file(variant_path / "metadata.json") + validate_dependencies(metadata.name.python_name, metadata.python_depends, _backend()) + + def validate_dependencies(kernel_module_name: str, dependencies: list[str], backend: Backend): """ Validate a list of dependencies to ensure they are installed. diff --git a/kernels/src/kernels/importer.py b/kernels/src/kernels/importer.py new file mode 100644 index 000000000..3df5807e8 --- /dev/null +++ b/kernels/src/kernels/importer.py @@ -0,0 +1,122 @@ +import importlib +import sys +import warnings +from dataclasses import dataclass +from pathlib import Path +from types import ModuleType + +from kernels_data import Metadata + +from kernels.hf_hub import RepoInfo + + +@dataclass(frozen=True) +class LoadedKernel: + """ + This dataclass provides information about a loaded kernel: + + - `metadata` (`Metadata`): kernel metadata. + - `module` (`ModuleType`): the imported kernel module. + - `repo_info` (`kernels.hf_hub.RepoInfo | None`): populated only for + kernels loaded via `get_kernel`. Loaders that work from a local path + (`get_local_kernel`) or a lockfile (`get_locked_kernel`, `load_kernel`) + leave this as `None`. + + The metadata includes the following properties that describe a kernel: + + - `id` (`str`): kernel identifier that is unique to the kernel version + backend. + - `name` (`str`): the name of the kernel. + - `version` (`int`): the version of the kernel. + - `license` (`str`): the license of the kernel. + - `upstream` (`str | None`): the original upstream repository of the kernel. + - `source` (`str | None`): the kernel-builder formatted source repository. + - `python_depends` (`list[str]`): required Python dependencies. + - `backend`: information about the kernel's backend. + """ + + metadata: Metadata + module: ModuleType + repo_info: RepoInfo | None + + +_loaded_kernels: dict[Path, LoadedKernel] = {} + + +def get_loaded_kernels() -> list[LoadedKernel]: + """ + Return a snapshot of every kernel that has been loaded into the current process. + + The returned list is a new list; mutating it does not affect the registry. + + Returns: + `list[LoadedKernel]`: One [`LoadedKernel`] per distinct kernel variant path + loaded in this process. + + Example: + ```python + from kernels import get_kernel, get_loaded_kernels + + get_kernel("kernels-community/activation", version=1) + for loaded in get_loaded_kernels(): + print(loaded.metadata.name, loaded.repo_info) + ``` + """ + return list(_loaded_kernels.values()) + + +def _warn_if_dirty(metadata: Metadata, variant_str: str) -> None: + """Warn when a kernel variant was built from a dirty git tree. + + A dirty build was produced from a working tree with uncommitted changes, + so its git SHA does not fully identify the sources it was built from and + the build may not be reproducible. + """ + provenance = metadata.provenance + if provenance is None or not provenance.dirty: + return + + dirty_sources = [] + if provenance.kernel is not None and provenance.kernel.dirty: + dirty_sources.append("kernel source") + builder_git = provenance.kernel_builder.git + if builder_git is not None and builder_git.dirty: + dirty_sources.append("kernel-builder") + + warnings.warn( + f"Kernel '{metadata.name}' variant '{variant_str}' was built from a dirty " + f"git tree ({', '.join(dirty_sources)} had uncommitted changes). Its " + "recorded git revision does not fully identify the sources it was built " + "from, so the build may not be reproducible.", + stacklevel=3, + ) + + +def _import_from_path(variant_path: Path, repo_info: RepoInfo | None = None) -> ModuleType: + if (loaded_kernel := _loaded_kernels.get(variant_path)) is not None: + return loaded_kernel.module + + metadata = Metadata.read_from_file(variant_path / "metadata.json") + _warn_if_dirty(metadata, variant_path.name) + module_name = metadata.name.python_name + + file_path = variant_path / "__init__.py" + if not file_path.exists(): + file_path = variant_path / module_name / "__init__.py" + if not file_path.exists(): + raise FileNotFoundError(f"No kernel module found at: `{variant_path}`") + + spec = importlib.util.spec_from_file_location(metadata.id, file_path) + if spec is None: + raise ImportError(f"Cannot load spec for {module_name} from {file_path}") + module = importlib.util.module_from_spec(spec) + if module is None: + raise ImportError(f"Cannot load module {module_name} from spec") + sys.modules[metadata.id] = module + spec.loader.exec_module(module) # type: ignore + + _loaded_kernels[variant_path] = LoadedKernel( + metadata=metadata, + module=module, + repo_info=repo_info, + ) + return module diff --git a/kernels/src/kernels/install.py b/kernels/src/kernels/install.py new file mode 100644 index 000000000..08d598b91 --- /dev/null +++ b/kernels/src/kernels/install.py @@ -0,0 +1,198 @@ +from pathlib import Path + +from huggingface_hub import HfApi +from huggingface_hub.errors import LocalEntryNotFoundError + +from kernels.deps import validate_variant_dependencies +from kernels.hf_hub import CACHE_DIR, _get_hf_api +from kernels.status import resolve_status +from kernels.variants import ( + Variant, + get_variants, + get_variants_local, + resolve_variant, + variants_trace_str, +) + +# Exclude pattern for bytecode. These are not included in kernel builds, +# but builds not done using kernel-builder might accidentally include +# bytecode. So these patterns are used to ensure that they are never +# downloaded. +_BYTECODE_IGNORE_PATTERNS = ["*.pyc", "**/__pycache__/**"] + + +def install_kernel( + repo_id: str, + *, + revision: str, + local_files_only: bool = False, + backend: str | None = None, + user_agent: str | dict | None = None, + validate_dependencies: bool = False, +) -> Path: + """ + Download a kernel for the current environment to the cache. + + The output path is validated against the hashes in `variant_locks` when provided. + + Args: + repo_id (`str`): + The Hub repository containing the kernel. + revision (`str`): + The specific revision (branch, tag, or commit) to download. + local_files_only (`bool`, *optional*, defaults to `False`): + Whether to only use local files and not download from the Hub. + backend (`str`, *optional*): + The backend to load the kernel for. Can only be `cpu` or the backend that Torch is compiled for. + The backend will be detected automatically if not provided. + user_agent (`Union[str, dict]`, *optional*): + The `user_agent` info to pass to `snapshot_download()` for internal telemetry. + validate_dependencies (`bool`, defaults to False): + When set to True, performs dependency validation. Useful for kernels that have + extra Python dependencies. + + Returns: + `Path`: The path to the variant directory. + """ + api = _get_hf_api(user_agent=user_agent) + if local_files_only: + # Same local-cache resolution path used by `load_kernel`, which is + # always offline. Sharing the helper avoids the network dependency + # that `get_variants` would otherwise introduce. + variant_path = _resolve_local_variant_path( + api, + repo_id, + revision=revision, + backend=backend, + ) + # For locally downloaded kernels, we run the validation after resolving the path + if validate_dependencies: + validate_variant_dependencies(variant_path) + return variant_path + + repo_id, revision = resolve_status(api, repo_id, revision) + variants = get_variants(api, repo_id=repo_id, revision=revision) + variant, trace = resolve_variant(variants, backend) + + if variant is None: + raise FileNotFoundError( + f"Cannot find a build variant for this system in {repo_id} (revision: {revision}):\n\n{variants_trace_str(trace)}" + ) + + # Validate Python dependencies before downloading the variant. + if validate_dependencies: + metadata_path = api.hf_hub_download( + repo_id, + repo_type="kernel", + filename=f"build/{variant.variant_str}/metadata.json", + cache_dir=CACHE_DIR, + revision=revision, + local_files_only=False, + ) + validate_variant_dependencies(Path(metadata_path).parent) + + allow_patterns = [f"build/{variant.variant_str}/*"] + ignore_patterns = _BYTECODE_IGNORE_PATTERNS + + repo_path = Path( + str( + api.snapshot_download( + repo_id, + repo_type="kernel", + allow_patterns=allow_patterns, + ignore_patterns=ignore_patterns, + cache_dir=CACHE_DIR, + revision=revision, + local_files_only=False, + ) + ) + ) + + try: + return _find_kernel_in_repo_path( + repo_path, + variant=variant, + ) + except FileNotFoundError: + raise FileNotFoundError(f"Cannot install kernel from repo {repo_id} (revision: {revision})") + + +def _resolve_local_variant_path( + api: HfApi, + repo_id: str, + *, + revision: str, + backend: str | None = None, +) -> Path: + """Resolve a kernel variant path from the local Hugging Face cache only. + + Used by `load_kernel` (which always operates on a pre-downloaded, locked + kernel) and by the offline branch of `install_kernel`. + """ + try: + local_repo_path = Path( + str( + api.snapshot_download( + repo_id, + repo_type="kernel", + ignore_patterns=_BYTECODE_IGNORE_PATTERNS, + cache_dir=CACHE_DIR, + revision=revision, + local_files_only=True, + ) + ) + ) + except LocalEntryNotFoundError as e: + raise FileNotFoundError( + f"Cannot find a local snapshot for {repo_id} (revision: {revision}). " + "When Hugging Face Hub is in offline mode the kernel must already " + "be present in the local cache." + ) from e + + variants = get_variants_local(local_repo_path / "build") + variant, status = resolve_variant(variants, backend) + if variant is None: + raise FileNotFoundError( + f"Cannot find a build variant for this system in {repo_id} (revision: {revision}):\n\n{variants_trace_str(status)}" + ) + + return _find_kernel_in_repo_path( + local_repo_path, + variant=variant, + ) + + +def _find_kernel_in_repo_path( + repo_path: Path, + *, + variant: Variant, +) -> Path: + variant_str = variant.variant_str + variant_path = repo_path / "build" / variant_str + if not variant_path.exists(): + raise FileNotFoundError(f"Variant path does not exist: `{variant_path}`") + + return variant_path + + +def install_kernel_all_variants( + repo_id: str, + *, + revision: str, +) -> Path: + api = _get_hf_api() + + repo_path = Path( + str( + api.snapshot_download( + repo_id, + repo_type="kernel", + allow_patterns="build/*", + ignore_patterns=_BYTECODE_IGNORE_PATTERNS, + cache_dir=CACHE_DIR, + revision=revision, + ) + ) + ) + + return repo_path / "build" diff --git a/kernels/src/kernels/layer/func.py b/kernels/src/kernels/layer/func.py index c4a3af8e3..70049bda5 100644 --- a/kernels/src/kernels/layer/func.py +++ b/kernels/src/kernels/layer/func.py @@ -5,14 +5,14 @@ from typing import TYPE_CHECKING, Protocol, Type from .._versions import select_revision_or_version +from ..load import ( + get_kernel, + get_local_kernel, +) from ..locking import ( get_caller_locked_kernel_revision, get_locked_kernel_revision, ) -from ..utils import ( - get_kernel, - get_local_kernel, -) from .layer import _create_func_module, use_kernel_forward_from_hub from .repos import RepositoryProtocol diff --git a/kernels/src/kernels/layer/layer.py b/kernels/src/kernels/layer/layer.py index 903ec14ea..3c822dd1d 100644 --- a/kernels/src/kernels/layer/layer.py +++ b/kernels/src/kernels/layer/layer.py @@ -10,14 +10,14 @@ from typing import TYPE_CHECKING, Callable, Protocol, Type from .._versions import select_revision_or_version +from ..load import ( + get_kernel, + get_local_kernel, +) from ..locking import ( get_caller_locked_kernel_revision, get_locked_kernel_revision, ) -from ..utils import ( - get_kernel, - get_local_kernel, -) from .device import Device from .globals import _DISABLE_KERNEL_MAPPING, _KERNEL_MAPPING from .mode import Mode diff --git a/kernels/src/kernels/load.py b/kernels/src/kernels/load.py new file mode 100644 index 000000000..9bc74f34a --- /dev/null +++ b/kernels/src/kernels/load.py @@ -0,0 +1,264 @@ +import functools +import os +from pathlib import Path +from types import ModuleType + +from huggingface_hub import constants + +from kernels._versions import select_revision_or_version +from kernels.deps import validate_variant_dependencies +from kernels.hf_hub import RepoInfo, _check_trust_remote_code, _get_hf_api +from kernels.importer import _import_from_path +from kernels.install import install_kernel +from kernels.locking import ( + get_caller_locked_kernel_revision, + get_locked_kernel_revision, +) +from kernels.variants import ( + get_variants, + get_variants_local, + resolve_variant, +) + + +def _get_local_kernel_overrides() -> dict[str, Path]: + """Returns list local overrides for kernels.""" + local_kerels = os.environ.get("LOCAL_KERNELS", None) + if local_kerels is None: + return dict() + return _parse_local_kernel_overrides(local_kerels) + + +@functools.lru_cache(maxsize=1) +def _parse_local_kernel_overrides(local_kernels: str) -> dict[str, Path]: + """Parse the LOCAL_KERNELS environment variable into a dictionary.""" + overrides = {} + for entry in local_kernels.split(":"): + if "=" not in entry: + raise ValueError( + f"Invalid LOCAL_KERNELS entry: {entry}. Expected format: repo_id_1=path_1:repo_id_2=path_2" + ) + repo_id, path = entry.split("=", 1) + overrides[repo_id] = Path(path) + + return overrides + + +def get_kernel( + repo_id: str, + revision: str | None = None, + version: int | None = None, + backend: str | None = None, + user_agent: str | dict | None = None, + trust_remote_code: bool | list[str] = False, +) -> ModuleType: + """ + Load a kernel from the kernel hub. + + This function downloads a kernel to the local Hugging Face Hub cache directory (if it was not downloaded before) + and then loads the kernel. + + Args: + repo_id (`str`): + The Hub repository containing the kernel. + revision (`str`, *optional*): + The specific revision (branch, tag, or commit) to download. Cannot be used together with `version`. + version (`int`, *optional*): + The kernel version to download. Cannot be used together with `revision`. + Either `version` or `revision` must be specified. + backend (`str`, *optional*): + The backend to load the kernel for. Can only be `cpu` or the backend that Torch is compiled for. + The backend will be detected automatically if not provided. + user_agent (`Union[str, dict]`, *optional*): + The `user_agent` info to pass to `snapshot_download()` for internal telemetry. + trust_remote_code (`bool | list[str]`, *optional*, defaults to `False`): + Whether to allow loading kernels from untrusted organisations. When ``False``, + only kernels from trusted organisations are allowed. When ``True``, all + repositories are allowed. A list of strings will be used to verify signing + identities in a future release; for now it emits a warning and falls + back to the default trust check. + + Returns: + `ModuleType`: The imported kernel module. + + Example: + ```python + import torch + from kernels import get_kernel + + activation = get_kernel("kernels-community/relu", version=1) + x = torch.randn(10, 20, device="cuda") + out = torch.empty_like(x) + result = activation.relu(out, x) + ``` + """ + override = _get_local_kernel_overrides().get(repo_id, None) + if override is not None: + return get_local_kernel(override) + + _check_trust_remote_code(repo_id, trust_remote_code) + + revision = select_revision_or_version(repo_id, revision=revision, version=version) + repo_info = RepoInfo( + repo_id=repo_id, + revision=revision, + ) + variant_path = install_kernel( + repo_id, + backend=backend, + revision=revision, + user_agent=user_agent, + validate_dependencies=True, + local_files_only=constants.HF_HUB_OFFLINE, + ) + return _import_from_path(variant_path, repo_info=repo_info) + + +def get_local_kernel( + repo_path: Path, + backend: str | None = None, +) -> ModuleType: + """ + Import a kernel from a local kernel repository path. + + Args: + repo_path (`Path`): + The local path to the kernel repository. + backend (`str`, *optional*): + The backend to load the kernel for. Can only be `cpu` or the backend that Torch is compiled for. + The backend will be detected automatically if not provided. + + Returns: + `ModuleType`: The imported kernel module. + """ + for base_path in [repo_path, repo_path / "build"]: + variants = get_variants_local(base_path) + variant, _ = resolve_variant(variants, backend) + + if variant is not None: + variant_path = base_path / variant.variant_str + validate_variant_dependencies(variant_path) + return _import_from_path(variant_path) + + # If we didn't find the package in the repo we may have a explicit + # package path. + variant_path = repo_path + if variant_path.exists(): + validate_variant_dependencies(variant_path) + return _import_from_path(variant_path) + + raise FileNotFoundError(f"Could not find kernel in {repo_path}") + + +def has_kernel( + repo_id: str, + revision: str | None = None, + version: int | None = None, + backend: str | None = None, +) -> bool: + """ + Check whether a kernel build exists for the current environment (Torch version and compute framework). + + Args: + repo_id (`str`): + The Hub repository containing the kernel. + revision (`str`, *optional*): + The specific revision (branch, tag, or commit) to download. Cannot be used together with `version`. + version (`int`, *optional*): + The kernel version to download. Cannot be used together with `revision`. + Either `version` or `revision` must be specified. + backend (`str`, *optional*): + The backend to load the kernel for. Can only be `cpu` or the backend that Torch is compiled for. + The backend will be detected automatically if not provided. + + Returns: + `bool`: `True` if a kernel is available for the current environment. + """ + revision = select_revision_or_version(repo_id, revision=revision, version=version) + + api = _get_hf_api() + variants = get_variants(api, repo_id=repo_id, revision=revision) + variant, _ = resolve_variant(variants, backend) + + if variant is None: + return False + + return api.file_exists( + repo_id, + repo_type="kernel", + revision=revision, + filename=f"build/{variant.variant_str}/metadata.json", + ) + + +def load_kernel( + repo_id: str, + *, + lockfile: Path | None, + backend: str | None = None, +) -> ModuleType: + """ + Get a pre-downloaded, locked kernel. + + If `lockfile` is not specified, the lockfile will be loaded from the caller's package metadata. + + Args: + repo_id (`str`): + The Hub repository containing the kernel. + lockfile (`Path`, *optional*): + Path to the lockfile. If not provided, the lockfile will be loaded from the caller's package metadata. + backend (`str`, *optional*): + The backend to load the kernel for. Can only be `cpu` or the backend that Torch is compiled for. + The backend will be detected automatically if not provided. + + Returns: + `ModuleType`: The imported kernel module. + """ + if lockfile is None: + locked_sha = get_caller_locked_kernel_revision(repo_id) + else: + with open(lockfile, "r") as f: + locked_sha = get_locked_kernel_revision(repo_id, f.read()) + + if locked_sha is None: + raise ValueError( + f"Kernel `{repo_id}` is not locked. Please lock it with `kernels lock ` and then reinstall the project." + ) + + try: + variant_path = install_kernel(repo_id, revision=locked_sha, backend=backend, local_files_only=True) + except FileNotFoundError as e: + raise FileNotFoundError( + f"Locked kernel `{repo_id}` was not downloaded or does not have an " + "applicable variant. Make sure it's downloaded locally via " + "`kernels download `." + ) from e + return _import_from_path(variant_path) + + +def get_locked_kernel(repo_id: str) -> ModuleType: + """ + Get a kernel using a lock file. + + Args: + repo_id (`str`): + The Hub repository containing the kernel. + local_files_only (`bool`, *optional*, defaults to `False`): + Whether to only use local files and not download from the Hub. + + Returns: + `ModuleType`: The imported kernel module. + """ + locked_sha = get_caller_locked_kernel_revision(repo_id) + + if locked_sha is None: + raise ValueError(f"Kernel `{repo_id}` is not locked") + + variant_path = install_kernel( + repo_id, + revision=locked_sha, + local_files_only=constants.HF_HUB_OFFLINE, + validate_dependencies=True, + ) + + return _import_from_path(variant_path) diff --git a/kernels/src/kernels/utils.py b/kernels/src/kernels/utils.py deleted file mode 100644 index e99174a6f..000000000 --- a/kernels/src/kernels/utils.py +++ /dev/null @@ -1,625 +0,0 @@ -import functools -import importlib -import os -import sys -import warnings -from dataclasses import dataclass -from pathlib import Path -from types import ModuleType - -from huggingface_hub import HfApi, constants -from huggingface_hub.errors import LocalEntryNotFoundError -from kernels_data import Metadata - -from kernels._versions import select_revision_or_version -from kernels.backends import _backend -from kernels.deps import validate_dependencies -from kernels.hf_hub import CACHE_DIR, RepoInfo, _check_trust_remote_code, _get_hf_api -from kernels.locking import ( - get_caller_locked_kernel_revision, - get_locked_kernel_revision, -) -from kernels.status import resolve_status -from kernels.variants import ( - Decision, - Variant, - get_variants, - get_variants_local, - resolve_variant, - resolve_variants, - variants_trace_str, -) - -KNOWN_BACKENDS = {"cpu", "cuda", "metal", "neuron", "rocm", "xpu", "npu"} - -# Exclude patter for bytecode. These are not included in kernel builds, -# but builds not done using kernel-builder might accidentally include -# bytcode. So these patterns are used to ensure that they are never -# downloaded. -_BYTECODE_IGNORE_PATTERNS = ["*.pyc", "**/__pycache__/**"] - - -@dataclass(frozen=True) -class LoadedKernel: - """ - This dataclass provides information about a loaded kernel: - - - `metadata` (`Metadata`): kernel metadata. - - `module` (`ModuleType`): the imported kernel module. - - `repo_info` (`kernels.utils.RepoInfo | None`): populated only for - kernels loaded via `get_kernel`. Loaders that work from a local path - (`get_local_kernel`) or a lockfile (`get_locked_kernel`, `load_kernel`) - leave this as `None`. - - The metadata includes the following properties that describe a kernel: - - - `id` (`str`): kernel identifier that is unique to the kernel version + backend. - - `name` (`str`): the name of the kernel. - - `version` (`int`): the version of the kernel. - - `license` (`str`): the license of the kernel. - - `upstream` (`str | None`): the original upstream repository of the kernel. - - `source` (`str | None`): the kernel-builder formatted source repository. - - `python_depends` (`list[str]`): required Python dependencies. - - `backend`: information about the kernel's backend. - """ - - metadata: Metadata - module: ModuleType - repo_info: RepoInfo | None - - -_loaded_kernels: dict[Path, LoadedKernel] = {} - - -def get_loaded_kernels() -> list[LoadedKernel]: - """ - Return a snapshot of every kernel that has been loaded into the current process. - - The returned list is a new list; mutating it does not affect the registry. - - Returns: - `list[LoadedKernel]`: One [`LoadedKernel`] per distinct kernel variant path - loaded in this process. - - Example: - ```python - from kernels import get_kernel, get_loaded_kernels - - get_kernel("kernels-community/activation", version=1) - for loaded in get_loaded_kernels(): - print(loaded.metadata.name, loaded.repo_info) - ``` - """ - return list(_loaded_kernels.values()) - - -def _get_local_kernel_overrides() -> dict[str, Path]: - """Returns list local overrides for kernels.""" - local_kerels = os.environ.get("LOCAL_KERNELS", None) - if local_kerels is None: - return dict() - return _parse_local_kernel_overrides(local_kerels) - - -@functools.lru_cache(maxsize=1) -def _parse_local_kernel_overrides(local_kernels: str) -> dict[str, Path]: - """Parse the LOCAL_KERNELS environment variable into a dictionary.""" - overrides = {} - for entry in local_kernels.split(":"): - if "=" not in entry: - raise ValueError( - f"Invalid LOCAL_KERNELS entry: {entry}. Expected format: repo_id_1=path_1:repo_id_2=path_2" - ) - repo_id, path = entry.split("=", 1) - overrides[repo_id] = Path(path) - - return overrides - - -def _validate_variant_dependencies(variant_path: Path) -> None: - metadata = Metadata.read_from_file(variant_path / "metadata.json") - validate_dependencies(metadata.name.python_name, metadata.python_depends, _backend()) - - -def _warn_if_dirty(metadata: Metadata, variant_str: str) -> None: - """Warn when a kernel variant was built from a dirty git tree. - - A dirty build was produced from a working tree with uncommitted changes, - so its git SHA does not fully identify the sources it was built from and - the build may not be reproducible. - """ - provenance = metadata.provenance - if provenance is None or not provenance.dirty: - return - - dirty_sources = [] - if provenance.kernel is not None and provenance.kernel.dirty: - dirty_sources.append("kernel source") - builder_git = provenance.kernel_builder.git - if builder_git is not None and builder_git.dirty: - dirty_sources.append("kernel-builder") - - warnings.warn( - f"Kernel '{metadata.name}' variant '{variant_str}' was built from a dirty " - f"git tree ({', '.join(dirty_sources)} had uncommitted changes). Its " - "recorded git revision does not fully identify the sources it was built " - "from, so the build may not be reproducible.", - stacklevel=3, - ) - - -def _import_from_path(variant_path: Path, repo_info: RepoInfo | None = None) -> ModuleType: - if (loaded_kernel := _loaded_kernels.get(variant_path)) is not None: - return loaded_kernel.module - - metadata = Metadata.read_from_file(variant_path / "metadata.json") - _warn_if_dirty(metadata, variant_path.name) - module_name = metadata.name.python_name - - file_path = variant_path / "__init__.py" - if not file_path.exists(): - file_path = variant_path / module_name / "__init__.py" - if not file_path.exists(): - raise FileNotFoundError(f"No kernel module found at: `{variant_path}`") - - spec = importlib.util.spec_from_file_location(metadata.id, file_path) - if spec is None: - raise ImportError(f"Cannot load spec for {module_name} from {file_path}") - module = importlib.util.module_from_spec(spec) - if module is None: - raise ImportError(f"Cannot load module {module_name} from spec") - sys.modules[metadata.id] = module - spec.loader.exec_module(module) # type: ignore - - _loaded_kernels[variant_path] = LoadedKernel( - metadata=metadata, - module=module, - repo_info=repo_info, - ) - return module - - -def install_kernel( - repo_id: str, - *, - revision: str, - local_files_only: bool = False, - backend: str | None = None, - user_agent: str | dict | None = None, - validate_dependencies: bool = False, -) -> Path: - """ - Download a kernel for the current environment to the cache. - - The output path is validated against the hashes in `variant_locks` when provided. - - Args: - repo_id (`str`): - The Hub repository containing the kernel. - revision (`str`): - The specific revision (branch, tag, or commit) to download. - local_files_only (`bool`, *optional*, defaults to `False`): - Whether to only use local files and not download from the Hub. - backend (`str`, *optional*): - The backend to load the kernel for. Can only be `cpu` or the backend that Torch is compiled for. - The backend will be detected automatically if not provided. - user_agent (`Union[str, dict]`, *optional*): - The `user_agent` info to pass to `snapshot_download()` for internal telemetry. - validate_dependencies (`bool`, defaults to False): - When set to True, performs dependency validation. Useful for kernels that have - extra Python dependencies. - - Returns: - `Path`: The path to the variant directory. - """ - api = _get_hf_api(user_agent=user_agent) - if local_files_only: - # Same local-cache resolution path used by `load_kernel`, which is - # always offline. Sharing the helper avoids the network dependency - # that `get_variants` would otherwise introduce. - variant_path = _resolve_local_variant_path( - api, - repo_id, - revision=revision, - backend=backend, - ) - # For locally downloaded kernels, we run the validation after resolving the path - if validate_dependencies: - _validate_variant_dependencies(variant_path) - return variant_path - - repo_id, revision = resolve_status(api, repo_id, revision) - variants = get_variants(api, repo_id=repo_id, revision=revision) - variant, trace = resolve_variant(variants, backend) - - if variant is None: - raise FileNotFoundError( - f"Cannot find a build variant for this system in {repo_id} (revision: {revision}):\n\n{variants_trace_str(trace)}" - ) - - # Validate Python dependencies before downloading the variant. - if validate_dependencies: - metadata_path = api.hf_hub_download( - repo_id, - repo_type="kernel", - filename=f"build/{variant.variant_str}/metadata.json", - cache_dir=CACHE_DIR, - revision=revision, - local_files_only=False, - ) - _validate_variant_dependencies(Path(metadata_path).parent) - - allow_patterns = [f"build/{variant.variant_str}/*"] - ignore_patterns = _BYTECODE_IGNORE_PATTERNS - - repo_path = Path( - str( - api.snapshot_download( - repo_id, - repo_type="kernel", - allow_patterns=allow_patterns, - ignore_patterns=ignore_patterns, - cache_dir=CACHE_DIR, - revision=revision, - local_files_only=False, - ) - ) - ) - - try: - return _find_kernel_in_repo_path( - repo_path, - variant=variant, - ) - except FileNotFoundError: - raise FileNotFoundError(f"Cannot install kernel from repo {repo_id} (revision: {revision})") - - -def _resolve_local_variant_path( - api: HfApi, - repo_id: str, - *, - revision: str, - backend: str | None = None, -) -> Path: - """Resolve a kernel variant path from the local Hugging Face cache only. - - Used by `load_kernel` (which always operates on a pre-downloaded, locked - kernel) and by the offline branch of `install_kernel`. - """ - try: - local_repo_path = Path( - str( - api.snapshot_download( - repo_id, - repo_type="kernel", - ignore_patterns=_BYTECODE_IGNORE_PATTERNS, - cache_dir=CACHE_DIR, - revision=revision, - local_files_only=True, - ) - ) - ) - except LocalEntryNotFoundError as e: - raise FileNotFoundError( - f"Cannot find a local snapshot for {repo_id} (revision: {revision}). " - "When Hugging Face Hub is in offline mode the kernel must already " - "be present in the local cache." - ) from e - - variants = get_variants_local(local_repo_path / "build") - variant, status = resolve_variant(variants, backend) - if variant is None: - raise FileNotFoundError( - f"Cannot find a build variant for this system in {repo_id} (revision: {revision}):\n\n{variants_trace_str(status)}" - ) - - return _find_kernel_in_repo_path( - local_repo_path, - variant=variant, - ) - - -def _find_kernel_in_repo_path( - repo_path: Path, - *, - variant: Variant, -) -> Path: - variant_str = variant.variant_str - variant_path = repo_path / "build" / variant_str - if not variant_path.exists(): - raise FileNotFoundError(f"Variant path does not exist: `{variant_path}`") - - return variant_path - - -def install_kernel_all_variants( - repo_id: str, - *, - revision: str, -) -> Path: - api = _get_hf_api() - - repo_path = Path( - str( - api.snapshot_download( - repo_id, - repo_type="kernel", - allow_patterns="build/*", - ignore_patterns=_BYTECODE_IGNORE_PATTERNS, - cache_dir=CACHE_DIR, - revision=revision, - ) - ) - ) - - return repo_path / "build" - - -def get_kernel( - repo_id: str, - revision: str | None = None, - version: int | None = None, - backend: str | None = None, - user_agent: str | dict | None = None, - trust_remote_code: bool | list[str] = False, -) -> ModuleType: - """ - Load a kernel from the kernel hub. - - This function downloads a kernel to the local Hugging Face Hub cache directory (if it was not downloaded before) - and then loads the kernel. - - Args: - repo_id (`str`): - The Hub repository containing the kernel. - revision (`str`, *optional*): - The specific revision (branch, tag, or commit) to download. Cannot be used together with `version`. - version (`int`, *optional*): - The kernel version to download. Cannot be used together with `revision`. - Either `version` or `revision` must be specified. - backend (`str`, *optional*): - The backend to load the kernel for. Can only be `cpu` or the backend that Torch is compiled for. - The backend will be detected automatically if not provided. - user_agent (`Union[str, dict]`, *optional*): - The `user_agent` info to pass to `snapshot_download()` for internal telemetry. - trust_remote_code (`bool | list[str]`, *optional*, defaults to `False`): - Whether to allow loading kernels from untrusted organisations. When ``False``, - only kernels from trusted organisations are allowed. When ``True``, all - repositories are allowed. A list of strings will be used to verify signing - identities in a future release; for now it emits a warning and falls - back to the default trust check. - - Returns: - `ModuleType`: The imported kernel module. - - Example: - ```python - import torch - from kernels import get_kernel - - activation = get_kernel("kernels-community/relu", version=1) - x = torch.randn(10, 20, device="cuda") - out = torch.empty_like(x) - result = activation.relu(out, x) - ``` - """ - override = _get_local_kernel_overrides().get(repo_id, None) - if override is not None: - return get_local_kernel(override) - - _check_trust_remote_code(repo_id, trust_remote_code) - - revision = select_revision_or_version(repo_id, revision=revision, version=version) - repo_info = RepoInfo( - repo_id=repo_id, - revision=revision, - ) - variant_path = install_kernel( - repo_id, - backend=backend, - revision=revision, - user_agent=user_agent, - validate_dependencies=True, - local_files_only=constants.HF_HUB_OFFLINE, - ) - return _import_from_path(variant_path, repo_info=repo_info) - - -def get_local_kernel( - repo_path: Path, - backend: str | None = None, -) -> ModuleType: - """ - Import a kernel from a local kernel repository path. - - Args: - repo_path (`Path`): - The local path to the kernel repository. - backend (`str`, *optional*): - The backend to load the kernel for. Can only be `cpu` or the backend that Torch is compiled for. - The backend will be detected automatically if not provided. - - Returns: - `ModuleType`: The imported kernel module. - """ - for base_path in [repo_path, repo_path / "build"]: - variants = get_variants_local(base_path) - variant, _ = resolve_variant(variants, backend) - - if variant is not None: - variant_path = base_path / variant.variant_str - _validate_variant_dependencies(variant_path) - return _import_from_path(variant_path) - - # If we didn't find the package in the repo we may have a explicit - # package path. - variant_path = repo_path - if variant_path.exists(): - _validate_variant_dependencies(variant_path) - return _import_from_path(variant_path) - - raise FileNotFoundError(f"Could not find kernel in {repo_path}") - - -def has_kernel( - repo_id: str, - revision: str | None = None, - version: int | None = None, - backend: str | None = None, -) -> bool: - """ - Check whether a kernel build exists for the current environment (Torch version and compute framework). - - Args: - repo_id (`str`): - The Hub repository containing the kernel. - revision (`str`, *optional*): - The specific revision (branch, tag, or commit) to download. Cannot be used together with `version`. - version (`int`, *optional*): - The kernel version to download. Cannot be used together with `revision`. - Either `version` or `revision` must be specified. - backend (`str`, *optional*): - The backend to load the kernel for. Can only be `cpu` or the backend that Torch is compiled for. - The backend will be detected automatically if not provided. - - Returns: - `bool`: `True` if a kernel is available for the current environment. - """ - revision = select_revision_or_version(repo_id, revision=revision, version=version) - - api = _get_hf_api() - variants = get_variants(api, repo_id=repo_id, revision=revision) - variant, _ = resolve_variant(variants, backend) - - if variant is None: - return False - - return api.file_exists( - repo_id, - repo_type="kernel", - revision=revision, - filename=f"build/{variant.variant_str}/metadata.json", - ) - - -def get_kernel_variants( - repo_id: str, - revision: str | None = None, - version: int | None = None, - backend: str | None = None, -) -> list[Decision]: - """ - Resolve all build variants of a kernel against the current environment. - - The decisions are sorted with compatible variants first, the most preferred - variant leading. - - Args: - repo_id (`str`): - The Hub repository containing the kernel. - revision (`str`, *optional*): - The specific revision (branch, tag, or commit) to inspect. Cannot be used together with `version`. - version (`int`, *optional*): - The kernel version to inspect. Cannot be used together with `revision`. - Either `version` or `revision` must be specified. - backend (`str`, *optional*): - The backend to resolve variants for. Can only be `cpu` or the backend that Torch is compiled for. - The backend will be detected automatically if not provided. - - Returns: - `list[Decision]`: One `VariantAccepted` or `VariantRejected` per build variant - in the repository, compatible variants first. - - Example: - ```python - from kernels import get_kernel_variants, VariantAccepted - - for decision in get_kernel_variants("kernels-community/activation", version=1): - name = decision.variant.variant_str - if isinstance(decision, VariantAccepted): - print(f"{name}: compatible") - else: - print(f"{name}: rejected ({decision.reason})") - - ``` - """ - revision = select_revision_or_version(repo_id, revision=revision, version=version) - - api = _get_hf_api() - variants = get_variants(api, repo_id=repo_id, revision=revision) - _, trace = resolve_variants(variants, backend) - return trace - - -def load_kernel( - repo_id: str, - *, - lockfile: Path | None, - backend: str | None = None, -) -> ModuleType: - """ - Get a pre-downloaded, locked kernel. - - If `lockfile` is not specified, the lockfile will be loaded from the caller's package metadata. - - Args: - repo_id (`str`): - The Hub repository containing the kernel. - lockfile (`Path`, *optional*): - Path to the lockfile. If not provided, the lockfile will be loaded from the caller's package metadata. - backend (`str`, *optional*): - The backend to load the kernel for. Can only be `cpu` or the backend that Torch is compiled for. - The backend will be detected automatically if not provided. - - Returns: - `ModuleType`: The imported kernel module. - """ - if lockfile is None: - locked_sha = get_caller_locked_kernel_revision(repo_id) - else: - with open(lockfile, "r") as f: - locked_sha = get_locked_kernel_revision(repo_id, f.read()) - - if locked_sha is None: - raise ValueError( - f"Kernel `{repo_id}` is not locked. Please lock it with `kernels lock ` and then reinstall the project." - ) - - try: - variant_path = install_kernel(repo_id, revision=locked_sha, backend=backend, local_files_only=True) - except FileNotFoundError as e: - raise FileNotFoundError( - f"Locked kernel `{repo_id}` was not downloaded or does not have an " - "applicable variant. Make sure it's downloaded locally via " - "`kernels download `." - ) from e - return _import_from_path(variant_path) - - -def get_locked_kernel(repo_id: str) -> ModuleType: - """ - Get a kernel using a lock file. - - Args: - repo_id (`str`): - The Hub repository containing the kernel. - local_files_only (`bool`, *optional*, defaults to `False`): - Whether to only use local files and not download from the Hub. - - Returns: - `ModuleType`: The imported kernel module. - """ - locked_sha = get_caller_locked_kernel_revision(repo_id) - - if locked_sha is None: - raise ValueError(f"Kernel `{repo_id}` is not locked") - - variant_path = install_kernel( - repo_id, - revision=locked_sha, - local_files_only=constants.HF_HUB_OFFLINE, - validate_dependencies=True, - ) - - return _import_from_path(variant_path) diff --git a/kernels/src/kernels/variants.py b/kernels/src/kernels/variants.py index 2e51b5e02..50ba0b1e6 100644 --- a/kernels/src/kernels/variants.py +++ b/kernels/src/kernels/variants.py @@ -13,6 +13,7 @@ from huggingface_hub.hf_api import RepoFolder from packaging.version import Version, parse +from kernels._versions import select_revision_or_version from kernels.backends import ( CANN, CUDA, @@ -554,6 +555,57 @@ def sort_key(vs: Decision) -> tuple[int, ...]: return sorted(variants, key=sort_key) +def get_kernel_variants( + repo_id: str, + revision: str | None = None, + version: int | None = None, + backend: str | None = None, +) -> list[Decision]: + """ + Resolve all build variants of a kernel against the current environment. + + The decisions are sorted with compatible variants first, the most preferred + variant leading. + + Args: + repo_id (`str`): + The Hub repository containing the kernel. + revision (`str`, *optional*): + The specific revision (branch, tag, or commit) to inspect. Cannot be used together with `version`. + version (`int`, *optional*): + The kernel version to inspect. Cannot be used together with `revision`. + Either `version` or `revision` must be specified. + backend (`str`, *optional*): + The backend to resolve variants for. Can only be `cpu` or the backend that Torch is compiled for. + The backend will be detected automatically if not provided. + + Returns: + `list[Decision]`: One `VariantAccepted` or `VariantRejected` per build variant + in the repository, compatible variants first. + + Example: + ```python + from kernels import get_kernel_variants, VariantAccepted + + for decision in get_kernel_variants("kernels-community/activation", version=1): + name = decision.variant.variant_str + if isinstance(decision, VariantAccepted): + print(f"{name}: compatible") + else: + print(f"{name}: rejected ({decision.reason})") + + ``` + """ + from kernels.hf_hub import _get_hf_api + + revision = select_revision_or_version(repo_id, revision=revision, version=version) + + api = _get_hf_api() + variants = get_variants(api, repo_id=repo_id, revision=revision) + _, trace = resolve_variants(variants, backend) + return trace + + def variants_trace_str(trace: list[Decision]) -> str: # Ensure that the list is sorted. sorted = _sort_variants(trace) diff --git a/kernels/tests/test_dirty_provenance.py b/kernels/tests/test_dirty_provenance.py index 436ec15aa..7a039c1cd 100644 --- a/kernels/tests/test_dirty_provenance.py +++ b/kernels/tests/test_dirty_provenance.py @@ -3,7 +3,7 @@ import pytest from kernels_data import Metadata -from kernels.utils import _import_from_path, _loaded_kernels, _warn_if_dirty +from kernels.importer import _import_from_path, _loaded_kernels, _warn_if_dirty def _write_variant(tmp_path, provenance): diff --git a/kernels/tests/test_layer.py b/kernels/tests/test_layer.py index 790e64458..699c4a119 100644 --- a/kernels/tests/test_layer.py +++ b/kernels/tests/test_layer.py @@ -19,13 +19,13 @@ use_kernel_forward_from_hub, use_kernel_mapping, ) +from kernels.install import ( + install_kernel, +) from kernels.layer.layer import ( _KERNEL_MAPPING, _validate_layer, ) -from kernels.utils import ( - install_kernel, -) @pytest.fixture diff --git a/kernels/tests/test_loaded_kernels.py b/kernels/tests/test_loaded_kernels.py index 5530e747f..e821b24c1 100644 --- a/kernels/tests/test_loaded_kernels.py +++ b/kernels/tests/test_loaded_kernels.py @@ -4,7 +4,7 @@ from kernels import get_kernel, get_loaded_kernels, get_local_kernel, install_kernel from kernels.hf_hub import RepoInfo -from kernels.utils import LoadedKernel, _loaded_kernels +from kernels.importer import LoadedKernel, _loaded_kernels _REPO_ID = "kernels-community/relu" _PACKAGE_NAME = "relu" diff --git a/kernels/tests/test_verify.py b/kernels/tests/test_verify.py index ce2be9339..c1f1b97b6 100644 --- a/kernels/tests/test_verify.py +++ b/kernels/tests/test_verify.py @@ -2,7 +2,7 @@ from sigstore.verify import policy from kernels import install_kernel -from kernels.utils import select_revision_or_version +from kernels._versions import select_revision_or_version from kernels.verify import VerificationResult, verify_variant TEST_POLICIES: list[policy.VerificationPolicy] = [