diff --git a/docs/src/python/python.md b/docs/src/python/python.md index d0ebbd00b..87eb8d37e 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -470,6 +470,8 @@ with BlobReader(read_range, len(payload)) as reader: ::: lancedb.functions.FunctionVersion +::: lancedb.functions.HuggingFaceModel + ::: lancedb.functions.PythonRuntimeSpec ::: lancedb.GpuClass diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index 67784b0e5..6ff0dd1ff 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -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, diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index ae6de3ab4..ba2bbeba3 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -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:" ) + 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", diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index 535df89b4..bb2e6f755 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -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"], diff --git a/rust/lancedb/src/function.rs b/rust/lancedb/src/function.rs index 1f6c4faaf..8933d0b37 100644 --- a/rust/lancedb/src/function.rs +++ b/rust/lancedb/src/function.rs @@ -497,12 +497,36 @@ impl Serialize for PythonRuntimeSpec { } } +/// An immutable Hugging Face model snapshot required by a Function. +/// +/// The service prepares this snapshot before user code runs. `revision` is a +/// full 40-character lowercase commit SHA, not a branch or tag. +/// +/// ``` +/// use lancedb::function::HuggingFaceModel; +/// +/// let model = HuggingFaceModel { +/// repo: "sentence-transformers/all-MiniLM-L6-v2".into(), +/// revision: "1110a243fdf4706b3f48f1d95db1a4f5529b4d41".into(), +/// }; +/// ``` +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct HuggingFaceModel { + /// Hugging Face repository in `owner/name` form. + pub repo: String, + /// Exact repository commit containing the model snapshot. + pub revision: String, +} + /// A complete OCI Function image. Its digest is independent of catalog names. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct FunctionImage { pub manifest_digest: String, pub descriptor: Value, pub source: bool, + /// Immutable snapshots declared by this Function version. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub models: Vec, } fn deserialize_object_version<'de, D: Deserializer<'de>>( @@ -729,6 +753,9 @@ pub struct FunctionRegistrationRequest { /// declared against this version, not here. #[serde(default, skip_serializing_if = "Vec::is_empty")] pub secret_bindings: Vec, + /// Model snapshots prepared by the service before user code runs. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub models: Vec, } impl FunctionRegistrationRequest { diff --git a/rust/lancedb/src/remote/sql_test.rs b/rust/lancedb/src/remote/sql_test.rs index b2973c679..3c9edb322 100644 --- a/rust/lancedb/src/remote/sql_test.rs +++ b/rust/lancedb/src/remote/sql_test.rs @@ -2298,7 +2298,7 @@ async fn function_creation_and_description_use_sql_query_handles() { const REQUEST: &str = include_str!( "../../tests/fixtures/first_class_functions/v1/remote_function_registration_request.json" ); - for defer in [false, true] { + for (defer, with_models) in [(false, false), (false, true), (true, false), (true, true)] { let mut request_json: serde_json::Value = serde_json::from_str(REQUEST).unwrap(); request_json["batch_size"] = serde_json::json!(2); request_json["row_timeout_seconds"] = serde_json::json!(0.75); @@ -2307,6 +2307,16 @@ async fn function_creation_and_description_use_sql_query_handles() { "variable": "HF_TOKEN", "secret_ref": {"name": "hf-prod"}, }]); + let mut version_json: serde_json::Value = serde_json::from_str(VERSION).unwrap(); + if with_models { + let models = serde_json::json!([{ + "repo": "sentence-transformers/all-MiniLM-L6-v2", + "revision": "1110a243fdf4706b3f48f1d95db1a4f5529b4d41", + }]); + request_json["models"] = models.clone(); + version_json["image"]["models"] = models; + } + let expected_models = version_json["image"].get("models"); let mut expected_request = request_json.clone(); expected_request.as_object_mut().unwrap().remove("name"); let request = @@ -2314,7 +2324,7 @@ async fn function_creation_and_description_use_sql_query_handles() { .unwrap(); let result = RecordBatch::try_from_iter(vec![( "definition", - Arc::new(StringArray::from(vec![VERSION])) as ArrayRef, + Arc::new(StringArray::from(vec![version_json.to_string()])) as ArrayRef, )]) .unwrap(); let service = TestSqlService { @@ -2335,27 +2345,46 @@ async fn function_creation_and_description_use_sql_query_handles() { .unwrap(); assert_eq!(job.job_type().as_deref(), Some("query")); let id = Uuid::parse_str(job.id().unwrap()).unwrap(); - assert_eq!(job.wait().await.unwrap().name(), "embed"); - assert_eq!(job.wait().await.unwrap().version(), "1"); + let created = job.wait().await.unwrap(); + assert_eq!(created.name(), "embed"); + assert_eq!(created.version(), "1"); + // Assert the actual SQL body, after the same native JSON decode used + // by the Python binding, rather than only the Python declaration. + { + let seen = statements.lock().unwrap(); + let encoded = seen[0] + .strip_prefix(r#"CREATE FUNCTION "normalize_score" FROM '"#) + .unwrap() + .strip_suffix('\'') + .unwrap(); + let sent: serde_json::Value = + serde_json::from_str(&encoded.replace("''", "'")).unwrap(); + assert_eq!(sent, expected_request); + } + let created_json: serde_json::Value = + serde_json::from_str(&created.to_canonical_json().unwrap()).unwrap(); + assert_eq!(created_json["image"].get("models"), expected_models); assert_eq!( client.describe(id).await.unwrap().status, QueryStatus::Finished ); - assert_eq!( - db.get_function("embed", "1", &["team".into()]) - .await - .unwrap() - .name(), - "embed" - ); - let seen = statements.lock().unwrap(); - let encoded = seen[0] - .strip_prefix("CREATE FUNCTION \"normalize_score\" FROM '") - .unwrap() - .strip_suffix('\'') + let described = db + .get_function("embed", "1", &["team".into()]) + .await .unwrap(); - let sent: serde_json::Value = serde_json::from_str(&encoded.replace("''", "'")).unwrap(); - assert_eq!(sent, expected_request); + assert_eq!(described.name(), "embed"); + let described_json: serde_json::Value = + serde_json::from_str(&described.to_canonical_json().unwrap()).unwrap(); + assert_eq!(described_json["image"].get("models"), expected_models); + let listed = db + .list_functions(&["team".into()], Default::default()) + .await + .unwrap(); + assert_eq!(listed.items.len(), 1); + let listed_json: serde_json::Value = + serde_json::from_str(&listed.items[0].to_canonical_json().unwrap()).unwrap(); + assert_eq!(listed_json["image"].get("models"), expected_models); + let seen = statements.lock().unwrap(); assert_eq!(seen[1], "DESCRIBE FUNCTION \"embed\" WITH (version = '1')"); shutdown.send(()).unwrap(); }