diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index ba2bbeba3..0257e4d7b 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -482,6 +482,10 @@ class FunctionVersion(_RemoteValue): ``gpu_class`` is the operator-defined GPU pool fixed into this Function version, or ``None`` when registration used the deployment default. + + ``batch_size``, ``row_timeout_seconds``, and ``result_stability`` retain + the execution controls reported by the catalog. They are ``None`` for + versions registered before the corresponding declaration. """ name: str @@ -497,6 +501,19 @@ class FunctionVersion(_RemoteValue): metadata: Mapping[str, str] disabled: bool gpu_class: Optional[str] = None + batch_size: Optional[int] = None + row_timeout_seconds: Optional[float] = None + result_stability: Optional[str] = None + + @field_validator("batch_size", mode="before") + @classmethod + def _validate_batch_size(cls, value): + return _validate_batch_size(value) + + @field_validator("row_timeout_seconds", mode="before") + @classmethod + def _validate_row_timeout(cls, value): + return _validate_row_timeout(value) def __call__(self, **arguments: Any) -> FunctionApplication: """Bind this exact version to named table columns. diff --git a/python/python/tests/test_first_class_function_slice1.py b/python/python/tests/test_first_class_function_slice1.py index cc4485133..980c11993 100644 --- a/python/python/tests/test_first_class_function_slice1.py +++ b/python/python/tests/test_first_class_function_slice1.py @@ -180,6 +180,22 @@ def test_function_version_exposes_registered_gpu_class(): assert json.loads(version.to_canonical_json())["gpu_class"] == "a10g" +def test_function_version_preserves_execution_controls(): + value = job_result("remote_function_job.json") + value.update( + batch_size=12, + row_timeout_seconds=8.0, + result_stability="input_and_context", + ) + version = FunctionVersion.from_json(json.dumps(value)) + encoded = json.loads(version.to_canonical_json()) + + assert encoded == value + assert type(encoded["batch_size"]) is int + assert type(encoded["row_timeout_seconds"]) is float + assert type(encoded["result_stability"]) is str + + def test_function_version_binds_named_columns_as_one_immutable_application(): version = FunctionVersion.from_json( json.dumps(job_result("remote_function_job.json")) diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index bb2e6f755..6df017cff 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -25,7 +25,9 @@ import pytest import lancedb from lancedb import BlobReader +from lancedb.db import AsyncConnection from lancedb.functions import ( + FunctionVersion, HuggingFaceModel, PythonRuntimeSpec, SecretBinding, @@ -37,6 +39,7 @@ from lancedb.functions import ( udf, ) from lancedb.secrets import EnvVarSecret +from lancedb.remote.db import RemoteDBConnection from utils import start_server_thread THRESHOLD = 20 @@ -1865,6 +1868,78 @@ def test_local_function_catalog_operations_are_not_supported(tmp_path): db.drop_function("normalize_score", version=FUNCTION_VERSION) +class _RawFunctionCatalog: + def __init__(self, definition, *, encode=False): + self.definition = json.dumps(definition) if encode else definition + + async def get_function(self, _name, version, namespace_path): + return self.definition + + def list_functions(self, *, namespace_path=None, page_token=None, page_limit=None): + return self + + def __aiter__(self): + return self + + async def __anext__(self): + if self.definition is None: + raise StopAsyncIteration + definition, self.definition = self.definition, None + return definition + + def num_page_results(self): + return int(self.definition is not None) + + def page_token(self): + return None + + +def _assert_execution_control_encoding(version): + encoded = json.loads(version.to_canonical_json()) + assert type(encoded["batch_size"]) is int + assert type(encoded["row_timeout_seconds"]) is float + assert type(encoded["result_stability"]) is str + return encoded + + +@pytest.mark.asyncio +async def test_async_catalog_path_preserves_execution_controls(): + definition = json.loads( + (FIXTURES / "remote_function_version.canonical.json").read_text() + ) + definition.update( + batch_size=12, + row_timeout_seconds=8.0, + result_stability="input_and_context", + ) + connection = AsyncConnection(_RawFunctionCatalog(definition, encode=True)) + + exact = await connection.get_function("embed", version="1") + listed = [item async for item in connection.list_functions(page_limit=32)] + assert len(listed) == 1 + for version in [exact, *listed]: + assert _assert_execution_control_encoding(version) == definition + + +def test_sync_catalog_path_preserves_execution_controls(): + definition = json.loads( + (FIXTURES / "remote_function_version.canonical.json").read_text() + ) + definition.update( + batch_size=12, + row_timeout_seconds=8.0, + result_stability="input_and_context", + ) + connection = object.__new__(RemoteDBConnection) + connection._conn = _RawFunctionCatalog(FunctionVersion(**definition)) + + exact = connection.get_function("embed", version="1") + listed = list(connection.list_functions(page_limit=32)) + assert len(listed) == 1 + for version in [exact, *listed]: + assert _assert_execution_control_encoding(version) == definition + + @contextlib.contextmanager def _mock_remote_function_catalog(): state = {"requests": [], "version": None} diff --git a/rust/lancedb/src/function.rs b/rust/lancedb/src/function.rs index ffa61f158..96eb4de10 100644 --- a/rust/lancedb/src/function.rs +++ b/rust/lancedb/src/function.rs @@ -541,6 +541,31 @@ fn deserialize_object_version<'de, D: Deserializer<'de>>( } } +/// A positive, finite per-call deadline reported by the catalog. +#[derive(Debug, Clone, Copy, PartialEq, Serialize)] +#[serde(transparent)] +struct RowTimeoutSeconds(f64); + +impl Eq for RowTimeoutSeconds {} + +impl RowTimeoutSeconds { + fn get(self) -> f64 { + self.0 + } +} + +impl<'de> Deserialize<'de> for RowTimeoutSeconds { + fn deserialize>(deserializer: D) -> std::result::Result { + let seconds = f64::deserialize(deserializer)?; + if !seconds.is_finite() || seconds <= 0.0 { + return Err(de::Error::custom( + "row_timeout_seconds must be positive and finite", + )); + } + Ok(Self(seconds)) + } +} + /// One immutable Function object revision and its executable artifact. /// /// The Function is named by its parts, `name` and `namespace_path`, the way a @@ -570,6 +595,19 @@ pub struct FunctionVersion { /// Operator-defined GPU class declared at registration. #[serde(default, skip_serializing_if = "Option::is_none")] gpu_class: Option, + /// Default maximum rows in one callable batch. + #[serde( + default, + skip_serializing_if = "Option::is_none", + deserialize_with = "deserialize_batch_size" + )] + batch_size: Option, + /// Per-call deadline in seconds. + #[serde(default, skip_serializing_if = "Option::is_none")] + row_timeout_seconds: Option, + /// Result-stability value declared at registration. + #[serde(default, skip_serializing_if = "Option::is_none")] + result_stability: Option, } impl FunctionVersion { @@ -629,6 +667,21 @@ impl FunctionVersion { self.gpu_class.as_deref() } + /// Default maximum rows in one callable batch, when declared. + pub fn batch_size(&self) -> Option { + self.batch_size + } + + /// Per-call deadline in seconds, when declared. + pub fn row_timeout_seconds(&self) -> Option { + self.row_timeout_seconds.map(RowTimeoutSeconds::get) + } + + /// Result-stability value declared at registration, when present. + pub fn result_stability(&self) -> Option<&str> { + self.result_stability.as_deref() + } + pub fn signature(&self) -> &FunctionSignature { &self.signature } diff --git a/rust/lancedb/src/remote/sql_test.rs b/rust/lancedb/src/remote/sql_test.rs index b25ab4ddb..cbca4e9f6 100644 --- a/rust/lancedb/src/remote/sql_test.rs +++ b/rust/lancedb/src/remote/sql_test.rs @@ -2775,6 +2775,9 @@ async fn function_creation_and_description_use_sql_query_handles() { "secret_ref": {"name": "hf-prod"}, }]); let mut version_json: serde_json::Value = serde_json::from_str(VERSION).unwrap(); + version_json["batch_size"] = serde_json::json!(12); + version_json["row_timeout_seconds"] = serde_json::json!(8.0); + version_json["result_stability"] = serde_json::json!("input_and_context"); if with_models { let models = serde_json::json!([{ "repo": "sentence-transformers/all-MiniLM-L6-v2", @@ -2815,6 +2818,9 @@ async fn function_creation_and_description_use_sql_query_handles() { let created = job.wait().await.unwrap(); assert_eq!(created.name(), "embed"); assert_eq!(created.version(), "1"); + assert_eq!(created.batch_size(), Some(12)); + assert_eq!(created.row_timeout_seconds(), Some(8.0)); + assert_eq!(created.result_stability(), Some("input_and_context")); // Assert the actual SQL body, after the same native JSON decode used // by the Python binding, rather than only the Python declaration. { @@ -2831,6 +2837,9 @@ async fn function_creation_and_description_use_sql_query_handles() { 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!(created_json["batch_size"], 12); + assert_eq!(created_json["row_timeout_seconds"], 8.0); + assert_eq!(created_json["result_stability"], "input_and_context"); assert_eq!( client.describe(id).await.unwrap().status, QueryStatus::Finished @@ -2840,17 +2849,32 @@ async fn function_creation_and_description_use_sql_query_handles() { .await .unwrap(); assert_eq!(described.name(), "embed"); + assert_eq!(described.batch_size(), Some(12)); + assert_eq!(described.row_timeout_seconds(), Some(8.0)); + assert_eq!(described.result_stability(), Some("input_and_context")); 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); + assert_eq!(described_json["batch_size"], 12); + assert_eq!(described_json["row_timeout_seconds"], 8.0); + assert_eq!(described_json["result_stability"], "input_and_context"); let listed = db .list_functions(&["team".into()], Default::default()) .await .unwrap(); assert_eq!(listed.items.len(), 1); + assert_eq!(listed.items[0].batch_size(), Some(12)); + assert_eq!(listed.items[0].row_timeout_seconds(), Some(8.0)); + assert_eq!( + listed.items[0].result_stability(), + Some("input_and_context") + ); 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); + assert_eq!(listed_json["batch_size"], 12); + assert_eq!(listed_json["row_timeout_seconds"], 8.0); + assert_eq!(listed_json["result_stability"], "input_and_context"); let seen = statements.lock().unwrap(); assert_eq!(seen[1], "DESCRIBE FUNCTION \"embed\" WITH (version = '1')"); shutdown.send(()).unwrap();