feat: support pinned Hugging Face model declarations (#4821)

Python Function declarations can now pin Hugging Face snapshots with
`models=[HuggingFaceModel(repo=..., revision=...)]`. The native Rust
registration path carries those declarations into `CREATE FUNCTION`, and
create/get/list results retain `image.models`. Functions without
declarations preserve their existing payload shape. The API reference
includes a model-loading example that initializes once per Function
instance.

Validation:
- The real SQL mock-service regression fails on the previous native
implementation and passes after the fix. It covers immediate/deferred
execution, requests with/without models, and create/get/list responses.
- Rust Function wire tests, strict Clippy, the public-type doctest, and
formatting.
- Rebuilt Python extension and targeted registration/golden tests.

The service-side snapshot preparation and cache are delivered separately
for ENT-3483.
This commit is contained in:
Xuanwo
2026-11-04 11:26:49 +08:00
committed by GitHub
parent aeab69e86a
commit f4271b9dcb
6 changed files with 202 additions and 18 deletions
+1
View File
@@ -41,6 +41,7 @@ from .functions import (
FunctionBinding as FunctionBinding,
FunctionRegistrationRequest as FunctionRegistrationRequest,
FunctionVersion as FunctionVersion,
HuggingFaceModel as HuggingFaceModel,
PythonRuntimeSpec as PythonRuntimeSpec,
RefreshColumnResult as RefreshColumnResult,
UdfDefinition as UdfDefinition,
+92
View File
@@ -318,6 +318,64 @@ class PythonEnvironmentSpec(_RemoteValue):
system_packages: tuple[str, ...] = ()
_HUGGING_FACE_MODEL_REPO = re.compile(
r"^[A-Za-z0-9][A-Za-z0-9._-]*/[A-Za-z0-9][A-Za-z0-9._-]*$"
)
_HUGGING_FACE_REVISION = re.compile(r"^[0-9a-f]{40}$")
class HuggingFaceModel(_RemoteValue):
"""One immutable Hugging Face model snapshot required by a Function.
``revision`` is a full 40-character lowercase commit SHA, not a branch or
tag. A Function may declare each repository once; separate versions of the
same repository use separate Function versions. The service prepares the
snapshot before initialization, so the model can load without a revision
argument or a Hugging Face Hub request::
@udf(
pip=["sentence-transformers==3.3.1"],
models=[HuggingFaceModel(
repo="sentence-transformers/all-MiniLM-L6-v2",
revision="1110a243fdf4706b3f48f1d95db1a4f5529b4d41",
)],
)
class Embed:
def __init__(self):
from sentence_transformers import SentenceTransformer
self.model = SentenceTransformer(
"sentence-transformers/all-MiniLM-L6-v2", device="cpu"
)
def __call__(self, text: str) -> list[float]:
return self.model.encode(text).tolist()
"""
repo: str = Field(
pattern=r"^[A-Za-z0-9][A-Za-z0-9._-]*/[A-Za-z0-9][A-Za-z0-9._-]*$"
)
revision: str
@field_validator("repo")
@classmethod
def _validate_repo(cls, value: str) -> str:
if not _HUGGING_FACE_MODEL_REPO.fullmatch(value) or any(
part in {".", ".."} for part in value.split("/")
):
raise ValueError("model repo must be a Hugging Face owner/name")
return value
@field_validator("revision", mode="before")
@classmethod
def _validate_revision(cls, value: Any) -> str:
if not isinstance(value, str) or not _HUGGING_FACE_REVISION.fullmatch(value):
raise ValueError(
"model revision must be a full 40-character lowercase commit SHA"
)
return value
class PythonRuntimeSpec(_RemoteValue):
"""Remote runtime definition with environment values.
@@ -392,6 +450,7 @@ class FunctionImage(_RemoteValue):
manifest_digest: str
descriptor: Mapping[str, Any]
source: bool
models: tuple[HuggingFaceModel, ...] = ()
def _validate_object_version(value: str) -> str:
@@ -571,6 +630,7 @@ class FunctionRegistrationRequest(_RemoteValue):
batch_size: Optional[int] = None
row_timeout_seconds: Optional[float] = None
secret_bindings: tuple[SecretBinding, ...] = ()
models: tuple[HuggingFaceModel, ...] = ()
@field_validator("batch_size", mode="before")
@classmethod
@@ -2025,6 +2085,7 @@ class UdfDefinition:
apt: tuple[str, ...] = (),
image: Optional[str] = None,
code: tuple[types.ModuleType, ...] = (),
models: tuple[HuggingFaceModel, ...] = (),
):
is_class = inspect.isclass(function)
if not is_class and (
@@ -2054,6 +2115,15 @@ class UdfDefinition:
raise ValueError(
f"Function image {image!r} must be host/repository@sha256:<digest>"
)
if any(not isinstance(model, HuggingFaceModel) for model in models):
raise TypeError("models must contain HuggingFaceModel values")
repos = [model.repo for model in models]
duplicates = sorted({repo for repo in repos if repos.count(repo) > 1})
if duplicates:
raise ValueError(
"a Function declares each Hugging Face repository once; "
f"duplicated: {duplicates!r}"
)
if conda:
environment_spec = PythonEnvironmentSpec(
kind="conda",
@@ -2075,6 +2145,16 @@ class UdfDefinition:
for key, value in environment.items()
):
raise TypeError("Function env keys and values must be strings")
reserved = sorted(
set(environment).intersection(
{"HF_HOME", "HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE"}
)
)
if models and reserved:
raise ValueError(
"model Functions reserve environment variables "
f"{reserved!r}; the worker sets them for the private snapshot"
)
if is_class:
if "<" in function.__qualname__ or "." in function.__qualname__:
raise ValueError(
@@ -2142,6 +2222,7 @@ class UdfDefinition:
runtime=runtime,
batch_size=batch_size,
row_timeout_seconds=row_timeout_seconds,
models=models,
)
# A class's own __dict__ holds its methods; copying it would shadow
# this wrapper's attributes.
@@ -2245,6 +2326,7 @@ def udf(
apt: tuple[str, ...] | list[str] = (),
image: Optional[str] = None,
code: Sequence[types.ModuleType] = (),
models: Sequence[HuggingFaceModel] = (),
) -> Callable[[Callable[..., Any]], UdfDefinition]: ...
@@ -2266,6 +2348,7 @@ def udf(
apt: tuple[str, ...] | list[str] = (),
image: Optional[str] = None,
code: Sequence[types.ModuleType] = (),
models: Sequence[HuggingFaceModel] = (),
):
"""Prepare a Python function or class for remote Function registration.
@@ -2407,6 +2490,13 @@ def udf(
package from ``code=`` so the copied Python package cannot shadow the
installed distribution. ``@udf`` warns when it skips files from a
package in ``code=``.
models : sequence of HuggingFaceModel, optional
Immutable Hugging Face repository snapshots prepared in the worker
node cache before user code starts. A Function may declare a
repository once. Different revisions of one repository use different
Function versions. Model Functions run offline with a private
``HF_HOME``; bind ``HF_TOKEN`` through Secrets when the repository
requires authentication.
Returns
-------
@@ -2516,6 +2606,7 @@ def udf(
apt=tuple(apt),
image=image,
code=tuple(code),
models=tuple(models),
)
if function is None:
@@ -2539,6 +2630,7 @@ __all__ = [
"FunctionImage",
"FunctionVersion",
"FunctionVersionRef",
"HuggingFaceModel",
"InputBinding",
"OutputMapping",
"PythonEnvironmentSpec",
@@ -26,6 +26,7 @@ import pytest
import lancedb
from lancedb import BlobReader
from lancedb.functions import (
HuggingFaceModel,
PythonRuntimeSpec,
SecretBinding,
SecretReference,
@@ -92,6 +93,38 @@ def test_scalar_udf_matches_shared_registration_golden_and_remains_callable():
}
def test_udf_declares_pinned_hugging_face_models():
model = HuggingFaceModel(
repo="sentence-transformers/all-MiniLM-L6-v2",
revision="1110a243fdf4706b3f48f1d95db1a4f5529b4d41",
)
@udf(pip=["sentence-transformers==3.3.1"], models=[model])
def embed_text(value: str) -> list[float]:
return [1.0]
request = json.loads(embed_text.registration_request.to_canonical_json())
assert request["models"] == [
{
"repo": "sentence-transformers/all-MiniLM-L6-v2",
"revision": "1110a243fdf4706b3f48f1d95db1a4f5529b4d41",
}
]
undeclared = json.loads(normalize_score.registration_request.to_canonical_json())
assert "models" not in undeclared
with pytest.raises(ValueError, match="once"):
udf(name="duplicate_model", models=[model, model])(lambda value: value)
with pytest.raises(ValueError, match="40-character"):
HuggingFaceModel(repo="owner/model", revision="main")
with pytest.raises(ValueError, match="reserve"):
udf(
name="reserved_model_env",
models=[model],
env={"HF_HOME": "/tmp/model"},
)(lambda value: value)
def test_udf_execution_controls_match_shared_registration_golden():
controlled = udf(
pip=["numpy>=2"],