diff --git a/docs/source/layers.md b/docs/source/layers.md index c19c7899..99587380 100644 --- a/docs/source/layers.md +++ b/docs/source/layers.md @@ -229,6 +229,24 @@ with use_kernel_mapping(kernel_layer_mapping): This ensures that the mapping is not active anymore outside the `with`-scope. +### Attributing Hub requests + +Libraries and applications can identify the Hub requests made while loading a +layer by passing optional user-agent metadata to [`~kernels.LayerRepository`]: + +```python +layer_repo = LayerRepository( + repo_id="kernels-community/activation", + layer_name="SiluAndMul", + version=1, + user_agent={"my-library": "1.0.0"}, +) +``` + +`user_agent` accepts either a string or a dictionary. If it is omitted, no +application-specific metadata is added. Setting `HF_HUB_DISABLE_TELEMETRY=1` +disables user-agent telemetry, including metadata supplied this way. + If the layer is stateless (it does not use member variables in its forward _or_ it was originally a function that was converted into a kernel layer with [`~kernels.use_kernel_func_from_hub`]), it can also be mapped to a kernel function: diff --git a/kernels/src/kernels/layer/layer.py b/kernels/src/kernels/layer/layer.py index 4d724365..8b8cac86 100644 --- a/kernels/src/kernels/layer/layer.py +++ b/kernels/src/kernels/layer/layer.py @@ -59,6 +59,8 @@ class LayerRepository: only kernels from trusted organisations are allowed. When `True`, all repositories are allowed. A list of repository IDs allows only those repositories in addition to repositories from trusted organisations. + user_agent (`Union[str, dict]`, *optional*): + Optional application metadata to include in the user-agent for Hub requests. Example: ```python @@ -81,6 +83,7 @@ def __init__( revision: str | None = None, version: int | None = None, trust_remote_code: bool | list[str] = False, + user_agent: str | dict | None = None, ): if revision is not None and version is not None: raise ValueError("Either a revision or a version must be specified, not both.") @@ -92,6 +95,7 @@ def __init__( self._trust_remote_code = ( trust_remote_code.copy() if isinstance(trust_remote_code, list) else trust_remote_code ) + self._user_agent = user_agent # We are going to resolve these lazily, since we do not want # to do a network request for every registered LayerRepository. @@ -112,6 +116,7 @@ def load(self) -> Type["nn.Module"]: self._repo_id, revision=self._resolve_revision(), trust_remote_code=self._trust_remote_code, + user_agent=self._user_agent, ) return _get_kernel_layer(self, kernel) diff --git a/kernels/tests/test_layer.py b/kernels/tests/test_layer.py index aa65de63..bac5dccd 100644 --- a/kernels/tests/test_layer.py +++ b/kernels/tests/test_layer.py @@ -1,6 +1,7 @@ import logging import sys from contextlib import nullcontext +from types import SimpleNamespace import pytest import torch @@ -673,6 +674,35 @@ def test_layer_repository_requires_version_or_revision(): LayerRepository(repo_id="kernels-test/silu-and-mul", layer_name="SiluAndMul") +@pytest.mark.parametrize("user_agent", [None, "transformers/5.0.0", {"transformers": "5.0.0"}]) +def test_layer_repository_forwards_user_agent(monkeypatch, user_agent): + calls = [] + + def get_kernel(repo_id, **kwargs): + calls.append((repo_id, kwargs)) + return SimpleNamespace(layers=SimpleNamespace(SiluAndMul=SiluAndMul)) + + monkeypatch.setattr("kernels.layer.layer.get_kernel", get_kernel) + repo = LayerRepository( + repo_id="kernels-test/silu-and-mul", + layer_name="SiluAndMul", + revision="main", + user_agent=user_agent, + ) + + assert repo.load() is SiluAndMul + assert calls == [ + ( + "kernels-test/silu-and-mul", + { + "revision": "main", + "trust_remote_code": False, + "user_agent": user_agent, + }, + ) + ] + + def test_layer_repository_with_trust_remote_code_allowlist_is_hashable(): allowlist = ["untrusted-org/allowed-kernel"] repo = LayerRepository(