mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-13 14:16:54 +00:00
fix: preserve FunctionVersion execution controls (#4856)
Function catalog responses include `batch_size`, `row_timeout_seconds`, and `result_stability`, but the Rust and Python `FunctionVersion` models silently discarded them. Preserve these optional execution controls through decoding, typed access, and canonical serialization; versions without the fields retain their existing representation. The regression exercises raw SQL transport responses and public Python sync/async get/list paths. Live Sophon catalog acceptance remains pending under ENT-2073 and the related FunctionVersion round-trip requirement in ENT-2176.
This commit is contained in:
@@ -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.
|
||||
|
||||
@@ -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"))
|
||||
|
||||
@@ -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}
|
||||
|
||||
Reference in New Issue
Block a user