Compare commits

..

4 Commits

3 changed files with 32 additions and 8 deletions

View File

@ -22,7 +22,7 @@ print(y)
## Docker Reference
build and run the reference [example/basic.py](example/basic.py) in a Docker container with the following commands:
build and run the reference [examples/basic.py](examples/basic.py) in a Docker container with the following commands:
```bash
docker build --platform linux/amd64 -t kernels-reference -f docker/Dockerfile.reference .

View File

@ -50,6 +50,8 @@ def download_kernels(args):
with open(args.project_dir / "hf-kernels.lock", "r") as f:
lock_json = json.load(f)
all_successful = True
for kernel_lock_json in lock_json:
kernel_lock = KernelLock.from_json(kernel_lock_json)
print(
@ -59,7 +61,14 @@ def download_kernels(args):
if args.all_variants:
install_kernel_all_variants(kernel_lock.repo_id, kernel_lock.sha)
else:
install_kernel(kernel_lock.repo_id, kernel_lock.sha)
try:
install_kernel(kernel_lock.repo_id, kernel_lock.sha)
except FileNotFoundError as e:
print(e, file=sys.stderr)
all_successful = False
if not all_successful:
sys.exit(1)
def lock_kernels(args):

View File

@ -4,11 +4,12 @@ import importlib.metadata
import inspect
import json
import os
from pathlib import Path
import platform
import sys
from importlib.metadata import Distribution
from types import ModuleType
from typing import List, Optional
from typing import List, Optional, Tuple
from huggingface_hub import hf_hub_download, snapshot_download
from packaging.version import parse
@ -22,6 +23,9 @@ CACHE_DIR: Optional[str] = os.environ.get("HF_KERNELS_CACHE", None)
def build_variant():
import torch
if torch.version.cuda is None:
raise AssertionError("This kernel requires CUDA to be installed. Torch was not compiled with CUDA enabled.")
torch_version = parse(torch.__version__)
cuda_version = parse(torch.version.cuda)
cxxabi = "cxx11" if torch.compiled_with_cxx11_abi() else "cxx98"
@ -45,10 +49,12 @@ def import_from_path(module_name: str, file_path):
return module
def install_kernel(repo_id: str, revision: str, local_files_only: bool = False):
package_name = get_metadata(repo_id, revision, local_files_only=local_files_only)[
"torch"
]["name"]
def install_kernel(
repo_id: str, revision: str, local_files_only: bool = False
) -> Tuple[str, str]:
"""Download a kernel for the current environment to the cache."""
package_name = repo_id.split('/')[-1]
package_name = package_name.replace('-', '_')
repo_path = snapshot_download(
repo_id,
allow_patterns=f"build/{build_variant()}/*",
@ -56,7 +62,16 @@ def install_kernel(repo_id: str, revision: str, local_files_only: bool = False):
revision=revision,
local_files_only=local_files_only,
)
return package_name, f"{repo_path}/build/{build_variant()}"
variant_path = f"{repo_path}/build/{build_variant()}"
module_init_path = f"{variant_path}/{package_name}/__init__.py"
if not os.path.exists(module_init_path):
raise FileNotFoundError(
f"Kernel `{repo_id}` at revision {revision} does not have build: {build_variant()}"
)
return package_name, variant_path
def install_kernel_all_variants(