diff --git a/docs/src/python/python.md b/docs/src/python/python.md index 3baa128ae..eb2c8d1db 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -465,6 +465,8 @@ with BlobReader(read_range, len(payload)) as reader: ::: lancedb.functions.FunctionRegistrationRequest +::: lancedb.functions.FunctionRetryPolicy + ::: lancedb.functions.FunctionArtifactRequest ::: lancedb.functions.FunctionArtifactContent diff --git a/docs/web/feature-engineering/refresh-column.mdx b/docs/web/feature-engineering/refresh-column.mdx index ea8b49b9b..24d050362 100644 --- a/docs/web/feature-engineering/refresh-column.mdx +++ b/docs/web/feature-engineering/refresh-column.mdx @@ -159,6 +159,34 @@ work sizes. For Function columns, the refresh defaults to failing rows and record them in the error store, subject to failure limits. See the [SQL refresh options](/sql/dml/refresh-column#options). +{/* +Draft for the Function retry policy release. + +For registered Function columns, pass `retry_policy` to +`refresh_column_async` to replace the Function's registered retry policy +for this refresh. For example, +`retry_policy=FunctionRetryPolicy(transient=True, max_attempts=3)` allows +three attempts, including the first call. Import `FunctionRetryPolicy` +from `lancedb.functions`. Use `FunctionRetryPolicy(max_attempts=1)` to +disable user-error retries. Omitting the override preserves the registered +policy. + +This requires a supporting LanceDB Enterprise server as well as an SDK +that exposes these options. An older Function image may need +re-registration before it accepts a refresh override. Only enable retries +when repeating a failed call's partial work is safe, including any external +requests made before the error. + +The attempt budget belongs to one failed Function call. With an explicit +retry policy and `error_policy="skip"`, exhaustion in a batch Function +writes nulls for the whole failed call and records its errors, subject to +failure limits. It does not split the batch to restart the retry budget. +Row Functions isolate failures per row. Infrastructure recovery, framework +deadlines and cancellation remain separate from this user-error policy. +See [Function failure handling](/management/function#failures) for the +retry and replay boundaries. +*/} + For a column filled by a registered Function, the Python and Rust job APIs can override the Function's GPU placement for this refresh: diff --git a/docs/web/management/function.mdx b/docs/web/management/function.mdx index 63a3f6446..4709264c7 100644 --- a/docs/web/management/function.mdx +++ b/docs/web/management/function.mdx @@ -401,9 +401,8 @@ full worked example. ## Failure handling {#failures} -By default, a Function that raises fails the refresh job. The Function -definition has no error-handling parameter; choose the row-error policy on -the refresh instead: +By default, a Function that raises fails the refresh job. Choose how the +refresh handles row errors with its error policy: ```sql SQL icon="database" @@ -424,6 +423,43 @@ but how much of the failed job's work was kept is not guaranteed, so do not count on the retry costing only the rows that did not complete. A refresh with `WHERE` recomputes every matching live row, including already-filled outputs. +{/* +Draft for the Function retry policy release. + +### Retry transient errors + +Native retries require an SDK that exposes `FunctionRetryPolicy` and a +LanceDB Enterprise server that supports Function retry policies. Upgrading +the SDK alone does not enable retries on an older server. If a refresh +override reports that the Function image does not support retry policies, +re-register the Function on a supporting server. + +Import `FunctionRetryPolicy` from `lancedb.functions` and pass a policy to +`@function` through `retry_policy`. For example, +`FunctionRetryPolicy(transient=True, max_attempts=3)` allows up to three +calls for matching transient failures. `max_attempts` includes the first +call; setting it to `1` disables user-error retries. A +[refresh override](/feature-engineering/refresh-column#resources) replaces +the registered policy for that refresh. + +Enable retries only when repeating a failed call's partial work is safe. +An external request may have succeeded before the Function raised, so this +policy does not guarantee exactly-once external effects. Functions declared +with external side effects reject retry policies. Calls that return +successfully are not repeated by this policy; infrastructure recovery and +later refreshes have separate replay behavior. + +Each failed callable invocation has one attempt budget. After exhaustion, +the refresh's `error_policy` applies. With an explicit retry policy, `skip` +on a batch Function writes nulls for every row in the failed call and records +their errors. It does not split that call into smaller retries with new +budgets. A row Function handles each row separately. Error records expose +`attempts` when the server records it; older records may omit the count. + +Worker recovery uses a separate retry budget. Framework deadlines and +cancellation stop execution without consuming further user-error retries. +*/} + ## What gets packaged {#packaging} diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index f1ddfda52..2208fd612 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -42,6 +42,7 @@ from .functions import ( FunctionBinding as FunctionBinding, FunctionDefinition as FunctionDefinition, FunctionRegistrationRequest as FunctionRegistrationRequest, + FunctionRetryPolicy as FunctionRetryPolicy, FunctionVersion as FunctionVersion, HuggingFaceModel as HuggingFaceModel, PythonRuntimeSpec as PythonRuntimeSpec, diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index f17856f3c..d1cee5091 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -617,6 +617,8 @@ class Table: batch_size: Optional[int] = None, gpu_class: Optional[str] = None, gpu_count: Optional[int] = None, + retry_policy: Optional[str] = None, + error_policy: Optional[str] = None, ) -> RefreshColumnResult: ... async def refresh_column_async( self, @@ -625,6 +627,8 @@ class Table: batch_size: Optional[int] = None, gpu_class: Optional[str] = None, gpu_count: Optional[int] = None, + retry_policy: Optional[str] = None, + error_policy: Optional[str] = None, ) -> Job: ... async def function_errors( self, @@ -1008,6 +1012,7 @@ class RefreshColumnResult: version: int class FunctionErrorRecord: + attempts: Optional[int] job_id: str fragment_id: int row_offset: Optional[int] @@ -1206,3 +1211,5 @@ class AuthzObject: ) -> AuthzObject: ... def __str__(self) -> str: ... def __repr__(self) -> str: ... + +def validate_function_retry_policy(json: str) -> None: ... diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index 4fa3340d5..4dad618ba 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -213,6 +213,45 @@ class _OpenRemoteValue(_RemoteValue): return set((self.__pydantic_extra__ or {}).keys()) +class FunctionRetryPolicy(_RemoteValue): + """Bound retries of selected errors inside one Enterprise Function call. + + ``max_attempts`` includes the first call. Matchers are alternatives; + ``transient=True`` selects connection/timeouts, HTTP 429/503 and the + provider codes ``SlowDown`` and ``ServiceUnavailable``. Exception names + may be qualified, and subclasses and wrapped causes also match. + + Backoff doubles with full jitter, bounded by ``max_backoff_seconds``. + Retry-After above that bound ends retries. A failed batch is retried as + one call: under ``skip``, exhaustion nulls every row in that call. + Completed calls are never repeated. Only opt in when repeating partial + work is safe; this does not make external side effects exactly once. + + >>> policy = FunctionRetryPolicy(transient=True, max_attempts=3) + >>> policy.max_attempts + 3 + >>> disabled = FunctionRetryPolicy(max_attempts=1) + """ + + model_config = ConfigDict(extra="forbid", frozen=True) + + max_attempts: Annotated[int, Field(strict=True)] = 3 + initial_backoff_seconds: Annotated[float, Field(strict=True)] = 0.2 + max_backoff_seconds: Annotated[float, Field(strict=True)] = 30.0 + exception_types: tuple[str, ...] = () + message_pattern: Optional[str] = None + error_codes: tuple[str, ...] = () + transient: Annotated[bool, Field(strict=True)] = False + honor_retry_after: Annotated[bool, Field(strict=True)] = True + + @model_validator(mode="after") + def _validate_policy(self): + from ._lancedb import validate_function_retry_policy + + validate_function_retry_policy(self.model_dump_json()) + return self + + class FunctionArtifact(_RemoteValue): """Content-addressed Python artifact identity.""" @@ -653,6 +692,7 @@ class FunctionRegistrationRequest(_RemoteValue): runtime: PythonRuntimeSpec batch_size: Optional[int] = None row_timeout_seconds: Optional[float] = None + retry_policy: Optional[FunctionRetryPolicy] = None secret_bindings: tuple[SecretBinding, ...] = () models: tuple[HuggingFaceModel, ...] = () @@ -2108,6 +2148,7 @@ class FunctionDefinition: python_version: Optional[str], batch_size: Optional[int] = None, row_timeout_seconds: Optional[float] = None, + retry_policy: Optional[FunctionRetryPolicy] = None, gpu: bool = False, gpu_class: Optional[str] = None, conda: tuple[str, ...] = (), @@ -2254,6 +2295,7 @@ class FunctionDefinition: runtime=runtime, batch_size=batch_size, row_timeout_seconds=row_timeout_seconds, + retry_policy=retry_policy, models=models, ) # A class's own __dict__ holds its methods; copying it would shadow @@ -2351,6 +2393,7 @@ def function( python_version: Optional[str] = None, batch_size: Optional[int] = None, row_timeout_seconds: Optional[float] = None, + retry_policy: Optional[FunctionRetryPolicy] = None, gpu: bool = False, gpu_class: Optional[str] = None, conda: tuple[str, ...] | list[str] = (), @@ -2373,6 +2416,7 @@ def function( python_version: Optional[str] = None, batch_size: Optional[int] = None, row_timeout_seconds: Optional[float] = None, + retry_policy: Optional[FunctionRetryPolicy] = None, gpu: bool = False, gpu_class: Optional[str] = None, conda: tuple[str, ...] | list[str] = (), @@ -2504,6 +2548,11 @@ def function( scalar Functions and one batch for vectorized Functions. This does not govern class initialization. When omitted, the service applies its default per-call budget. + retry_policy : FunctionRetryPolicy, optional + Retry selected failed calls inside the Function, with bounded attempts + and backoff. A refresh may override this default. Retrying a partial + call must be safe; completed calls are not repeated. For batch Functions, + skip after exhaustion applies to the whole failed call. gpu : bool, default False Whether every remote execution requires a GPU. The execution platform selects one compatible GPU for each worker. The requirement is part of @@ -2631,6 +2680,7 @@ def function( python_version=python_version, batch_size=batch_size, row_timeout_seconds=row_timeout_seconds, + retry_policy=retry_policy, gpu=gpu, gpu_class=gpu_class, conda=tuple(conda), @@ -2657,6 +2707,7 @@ __all__ = [ "FunctionOutput", "FunctionParameter", "FunctionRegistrationRequest", + "FunctionRetryPolicy", "FunctionResultField", "FunctionSignature", "FunctionImage", diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index db74c5375..c916fcb92 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -48,7 +48,11 @@ from lancedb.index import ( LabelList, ) from lancedb.job import Job -from lancedb.functions import FunctionApplication, RefreshColumnResult +from lancedb.functions import ( + FunctionApplication, + FunctionRetryPolicy, + RefreshColumnResult, +) from lancedb.remote.db import LOOP from lancedb.table import IndexConfigType, _maybe_add_fts_error_note import pyarrow as pa @@ -1137,6 +1141,8 @@ class RemoteTable(Table): batch_size: Optional[int] = None, gpu_class: Optional[str] = None, gpu_count: Optional[int] = None, + retry_policy: Optional[FunctionRetryPolicy] = None, + error_policy: Optional[Literal["fail", "skip"]] = None, ): return LOOP.run( self._table.refresh_column( @@ -1144,6 +1150,8 @@ class RemoteTable(Table): batch_size=batch_size, gpu_class=gpu_class, gpu_count=gpu_count, + retry_policy=retry_policy, + error_policy=error_policy, ) ) @@ -1154,6 +1162,8 @@ class RemoteTable(Table): batch_size: Optional[int] = None, gpu_class: Optional[str] = None, gpu_count: Optional[int] = None, + retry_policy: Optional[FunctionRetryPolicy] = None, + error_policy: Optional[Literal["fail", "skip"]] = None, ) -> Job[RefreshColumnResult]: return Job( LOOP.run( @@ -1162,6 +1172,8 @@ class RemoteTable(Table): batch_size=batch_size, gpu_class=gpu_class, gpu_count=gpu_count, + retry_policy=retry_policy, + error_policy=error_policy, ) ) ) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 59876a8c3..559552cee 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -82,6 +82,7 @@ from .index import ( ) from .expr import Expr from .functions import ( + FunctionRetryPolicy, FunctionApplication, RefreshColumnResult as RefreshColumnJobResult, ) @@ -3166,6 +3167,8 @@ class Table(ABC): batch_size: Optional[int] = None, gpu_class: Optional[str] = None, gpu_count: Optional[int] = None, + retry_policy: Optional[FunctionRetryPolicy] = None, + error_policy: Optional[Literal["fail", "skip"]] = None, ) -> "RefreshColumnResult": """ Fill the rows of a computed column that hold no value yet. @@ -3199,6 +3202,13 @@ class Table(ABC): Positive aggregate GPU budget for Function workers. None leaves the server's configured budget unchanged. LanceDB Enterprise only. + retry_policy: FunctionRetryPolicy, optional + Override retries of failed user-code calls. Only use when repeating + partial work is safe. Set max_attempts=1 to disable retries. + Enterprise Functions only; local refreshes reject this option. + error_policy: {"fail", "skip"}, optional + Action after failure or retry exhaustion. Skip records the error + and attempt count, subject to the server's failure limits. Returns ------- @@ -3215,6 +3225,8 @@ class Table(ABC): batch_size: Optional[int] = None, gpu_class: Optional[str] = None, gpu_count: Optional[int] = None, + retry_policy: Optional[FunctionRetryPolicy] = None, + error_policy: Optional[Literal["fail", "skip"]] = None, ) -> Job[RefreshColumnJobResult]: """ Like :meth:`refresh_column`, but returns a handle to the refresh job @@ -3247,6 +3259,13 @@ class Table(ABC): Positive aggregate GPU budget for Function workers. None leaves the server's configured budget unchanged. LanceDB Enterprise only. + retry_policy: FunctionRetryPolicy, optional + Override retries of failed user-code calls. Only use when repeating + partial work is safe. Set max_attempts=1 to disable retries. + Enterprise Functions only; local refreshes reject this option. + error_policy: {"fail", "skip"}, optional + Action after failure or retry exhaustion. Skip records the error + and attempt count, subject to the server's failure limits. Returns ------- @@ -5358,6 +5377,8 @@ class LanceTable(Table): batch_size: Optional[int] = None, gpu_class: Optional[str] = None, gpu_count: Optional[int] = None, + retry_policy: Optional[FunctionRetryPolicy] = None, + error_policy: Optional[Literal["fail", "skip"]] = None, ) -> "RefreshColumnResult": """Fill a computed column's unfilled rows and recompute those whose inputs changed. See @@ -5368,6 +5389,8 @@ class LanceTable(Table): batch_size=batch_size, gpu_class=gpu_class, gpu_count=gpu_count, + retry_policy=retry_policy, + error_policy=error_policy, ) ) @@ -5378,6 +5401,8 @@ class LanceTable(Table): batch_size: Optional[int] = None, gpu_class: Optional[str] = None, gpu_count: Optional[int] = None, + retry_policy: Optional[FunctionRetryPolicy] = None, + error_policy: Optional[Literal["fail", "skip"]] = None, ) -> Job[RefreshColumnJobResult]: """Fill a computed column's unfilled rows and recompute those whose inputs changed, returning a handle to the refresh job. See @@ -5390,6 +5415,8 @@ class LanceTable(Table): batch_size=batch_size, gpu_class=gpu_class, gpu_count=gpu_count, + retry_policy=retry_policy, + error_policy=error_policy, ) ) ) @@ -7745,6 +7772,8 @@ class AsyncTable: batch_size: Optional[int] = None, gpu_class: Optional[str] = None, gpu_count: Optional[int] = None, + retry_policy: Optional[FunctionRetryPolicy] = None, + error_policy: Optional[Literal["fail", "skip"]] = None, ) -> RefreshColumnResult: """ Fill the rows of a computed column that hold no value yet. @@ -7778,6 +7807,13 @@ class AsyncTable: Positive aggregate GPU budget for Function workers. None leaves the server's configured budget unchanged. LanceDB Enterprise only. + retry_policy: FunctionRetryPolicy, optional + Override retries of failed user-code calls. Only use when repeating + partial work is safe. Set max_attempts=1 to disable retries. + Enterprise Functions only; local refreshes reject this option. + error_policy: {"fail", "skip"}, optional + Action after failure or retry exhaustion. Skip records the error + and attempt count, subject to the server's failure limits. Returns ------- @@ -7789,6 +7825,10 @@ class AsyncTable: batch_size=batch_size, gpu_class=gpu_class, gpu_count=gpu_count, + retry_policy=None + if retry_policy is None + else FunctionRetryPolicy.model_validate(retry_policy).model_dump_json(), + error_policy=error_policy, ) async def function_errors( @@ -7887,6 +7927,8 @@ class AsyncTable: batch_size: Optional[int] = None, gpu_class: Optional[str] = None, gpu_count: Optional[int] = None, + retry_policy: Optional[FunctionRetryPolicy] = None, + error_policy: Optional[Literal["fail", "skip"]] = None, ) -> AsyncJob[RefreshColumnJobResult]: """ Like :meth:`refresh_column`, but returns a handle to the refresh job @@ -7919,6 +7961,13 @@ class AsyncTable: Positive aggregate GPU budget for Function workers. None leaves the server's configured budget unchanged. LanceDB Enterprise only. + retry_policy: FunctionRetryPolicy, optional + Override retries of failed user-code calls. Only use when repeating + partial work is safe. Set max_attempts=1 to disable retries. + Enterprise Functions only; local refreshes reject this option. + error_policy: {"fail", "skip"}, optional + Action after failure or retry exhaustion. Skip records the error + and attempt count, subject to the server's failure limits. Returns ------- @@ -7947,6 +7996,10 @@ class AsyncTable: batch_size=batch_size, gpu_class=gpu_class, gpu_count=gpu_count, + retry_policy=None + if retry_policy is None + else FunctionRetryPolicy.model_validate(retry_policy).model_dump_json(), + error_policy=error_policy, ), RefreshColumnJobResult.from_json, ) diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index 238d787d1..9b435b256 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -159,6 +159,34 @@ def test_function_execution_controls_apply_to_vectorized_functions(): assert request["row_timeout_seconds"] == 1.25 +def test_function_retry_policy_is_part_of_registration(): + policy = lancedb.FunctionRetryPolicy(transient=True, max_attempts=4) + definition = function(retry_policy=policy)(normalize_score.__wrapped__) + request = json.loads(definition.registration_request.to_canonical_json()) + assert request["retry_policy"] == {"max_attempts": 4, "transient": True} + assert lancedb.FunctionRetryPolicy(max_attempts=1).max_attempts == 1 + + +@pytest.mark.parametrize( + "options", + [ + {"max_attempts": 0, "transient": True}, + {"max_attempts": 101, "transient": True}, + {"max_attempts": True}, + {"max_attempts": 3}, + {"transient": True, "initial_backoff_seconds": -1}, + {"transient": True, "max_backoff_seconds": float("inf")}, + {"transient": True, "max_backoff_seconds": 0}, + {"message_pattern": "["}, + {"exception_types": [""]}, + {"transient": True, "typo": 1}, + ], +) +def test_function_retry_policy_rejects_invalid_options(options): + with pytest.raises(ValueError): + lancedb.FunctionRetryPolicy(**options) + + @pytest.mark.parametrize( ("options", "field"), [ diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index cd2261f04..994655ef9 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -1825,6 +1825,11 @@ def test_remote_create_index_async_returns_job(): {"batch_size": 1}, {"batch_size": 8}, {"gpu_class": "a10g", "gpu_count": 2}, + { + "retry_policy": {"transient": True, "max_attempts": 3}, + "error_policy": "skip", + }, + {"retry_policy": {"max_attempts": 1}, "error_policy": "fail"}, ], ) async def test_remote_refresh_async_returns_typed_terminal_result( @@ -1849,6 +1854,12 @@ async def test_remote_refresh_async_returns_typed_terminal_result( expected["gpu_class"] = refresh_kwargs["gpu_class"] if refresh_kwargs.get("gpu_count") is not None: expected["gpu_count"] = refresh_kwargs["gpu_count"] + if refresh_kwargs.get("retry_policy") is not None: + expected["retry_policy"] = lancedb.FunctionRetryPolicy.model_validate( + refresh_kwargs["retry_policy"] + ).model_dump(mode="json") + if refresh_kwargs.get("error_policy") is not None: + expected["error_policy"] = refresh_kwargs["error_policy"] assert json.loads(body) == expected request.send_response(202) request.send_header("Content-Type", "application/json") @@ -2010,10 +2021,12 @@ def test_remote_refresh_async_forwards_progress_from_submitted_handle(): ({"gpu_class": "a_b"}, "gpu_class"), ({"gpu_count": 0}, "gpu_count"), ({"gpu_count": -1}, "gpu_count"), + ({"retry_policy": {"transient": True, "max_attempts": 0}}, "retry_policy"), + ({"error_policy": "retry"}, "error_policy"), ], ) @pytest.mark.asyncio -async def test_remote_refresh_rejects_invalid_gpu_options_before_request( +async def test_remote_refresh_rejects_invalid_options_before_request( refresh_kwargs, message ): backfill_requests = [] @@ -2062,7 +2075,8 @@ async def test_remote_refresh_rejects_invalid_gpu_options_before_request( assert backfill_requests == [] -def test_remote_function_errors_lists_the_rows_a_refresh_skipped(): +@pytest.mark.parametrize("attempts", [None, 3]) +def test_remote_function_errors_lists_the_rows_a_refresh_skipped(attempts): listing = { "records": [ { @@ -2077,6 +2091,7 @@ def test_remote_function_errors_lists_the_rows_a_refresh_skipped(): "table_version": 11, "error_type": "ValueError", "error_message": "bad input 'x'", + "attempts": attempts, "traceback": ( "Traceback (most recent call last):\nValueError: bad input 'x'" ), @@ -2127,6 +2142,7 @@ def test_remote_function_errors_lists_the_rows_a_refresh_skipped(): assert errors.records[0].row_id == 42 assert errors.records[0].row_address == 12884901897 assert errors.records[0].function_version == "2" + assert errors.records[0].attempts == attempts assert errors.records[0].traceback == ( "Traceback (most recent call last):\nValueError: bad input 'x'" ) diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 88635c020..45ce9276e 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -5984,6 +5984,14 @@ def test_refresh_column_rejects_invalid_batch_size(tmp_path, method, batch_size) getattr(table, method)("x", batch_size=batch_size) +@pytest.mark.parametrize("method", ["refresh_column", "refresh_column_async"]) +def test_local_refresh_rejects_function_retry_policy(tmp_path, method): + table = lancedb.connect(tmp_path).create_table("rows", [{"x": 1}]) + policy = lancedb.FunctionRetryPolicy(transient=True) + with pytest.raises(NotImplementedError, match="Enterprise Function refresh"): + getattr(table, method)("x", retry_policy=policy) + + @pytest.mark.asyncio @pytest.mark.parametrize("method", ["refresh_column", "refresh_column_async"]) @pytest.mark.parametrize("batch_size", [0, -1]) diff --git a/python/src/lib.rs b/python/src/lib.rs index 0690d9622..27ff4cb01 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -145,6 +145,7 @@ pub fn _lancedb(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(table::tokenize, m)?)?; m.add_function(wrap_pyfunction!(permutation::async_permutation_builder, m)?)?; m.add_function(wrap_pyfunction!(util::validate_table_name, m)?)?; + m.add_function(wrap_pyfunction!(util::validate_function_retry_policy, m)?)?; m.add_function(wrap_pyfunction!(util::parse_field_path, m)?)?; m.add_function(wrap_pyfunction!(query::fts_query_to_json, m)?)?; m.add_function(wrap_pyfunction!(expr_col, m)?)?; diff --git a/python/src/table.rs b/python/src/table.rs index 44c8017c4..4c5a919d2 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -46,8 +46,22 @@ fn refresh_column_options( batch_size: Option, gpu_class: Option, gpu_count: Option, + retry_policy: Option, + error_policy: Option, ) -> PyResult { let mut options = RefreshColumnOptions::default(); + options.retry_policy = retry_policy + .map(|json| { + serde_json::from_str(&json) + .map_err(|error| PyValueError::new_err(format!("invalid retry_policy: {error}"))) + }) + .transpose()?; + options.error_policy = error_policy + .map(|value| { + serde_json::from_value(serde_json::Value::String(value)) + .map_err(|error| PyValueError::new_err(format!("invalid error_policy: {error}"))) + }) + .transpose()?; options.batch_size = batch_size .map(usize::try_from) .transpose() @@ -611,6 +625,7 @@ impl From for RefreshColumnResult { #[pyclass(module = "lancedb._lancedb", get_all, from_py_object)] #[derive(Clone, Debug)] pub struct FunctionErrorRecord { + pub attempts: Option, pub job_id: String, pub fragment_id: u64, pub row_offset: Option, @@ -653,6 +668,7 @@ impl FunctionErrorRecord { impl From for FunctionErrorRecord { fn from(record: lancedb::function::FunctionErrorRecord) -> Self { Self { + attempts: record.attempts, job_id: record.job_id, fragment_id: record.fragment_id, row_offset: record.row_offset, @@ -2126,15 +2142,18 @@ impl Table { }) } - #[pyo3(signature = (column, *, batch_size=None, gpu_class=None, gpu_count=None))] + #[pyo3(signature = (column, *, batch_size=None, gpu_class=None, gpu_count=None, retry_policy=None, error_policy=None))] pub fn refresh_column( self_: PyRef<'_, Self>, column: String, batch_size: Option, gpu_class: Option, gpu_count: Option, + retry_policy: Option, + error_policy: Option, ) -> PyResult> { - let options = refresh_column_options(batch_size, gpu_class, gpu_count)?; + let options = + refresh_column_options(batch_size, gpu_class, gpu_count, retry_policy, error_policy)?; let inner = self_.inner_ref()?.clone(); future_into_py(self_.py(), async move { let result = inner @@ -2145,15 +2164,18 @@ impl Table { }) } - #[pyo3(signature = (column, *, batch_size=None, gpu_class=None, gpu_count=None))] + #[pyo3(signature = (column, *, batch_size=None, gpu_class=None, gpu_count=None, retry_policy=None, error_policy=None))] pub fn refresh_column_async( self_: PyRef<'_, Self>, column: String, batch_size: Option, gpu_class: Option, gpu_count: Option, + retry_policy: Option, + error_policy: Option, ) -> PyResult> { - let options = refresh_column_options(batch_size, gpu_class, gpu_count)?; + let options = + refresh_column_options(batch_size, gpu_class, gpu_count, retry_policy, error_policy)?; let inner = self_.inner_ref()?.clone(); future_into_py(self_.py(), async move { let job = inner diff --git a/python/src/util.rs b/python/src/util.rs index 2145c0f14..a956066f1 100644 --- a/python/src/util.rs +++ b/python/src/util.rs @@ -70,3 +70,14 @@ pub fn parse_field_path(path: &str) -> PyResult> { /// A wrapper around a LanceDB type to allow it to be used in Python #[derive(Debug, Clone)] pub struct PyLanceDB(pub T); +/// Validate the Function retry wire contract using the same rules as Rust APIs. +#[pyo3::pyfunction] +pub fn validate_function_retry_policy(json: &str) -> pyo3::PyResult<()> { + let policy: lancedb::function::FunctionRetryPolicy = + serde_json::from_str(json).map_err(|error| { + pyo3::exceptions::PyValueError::new_err(format!("invalid retry_policy: {error}")) + })?; + policy + .validate() + .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string())) +} diff --git a/rust/lancedb/src/function.rs b/rust/lancedb/src/function.rs index 6063c34f3..41a683c54 100644 --- a/rust/lancedb/src/function.rs +++ b/rust/lancedb/src/function.rs @@ -16,6 +16,9 @@ use serde_json::Value; use crate::secrets::SecretBinding; use crate::{Error, Result}; +mod retry; +pub use retry::{FunctionErrorPolicy, FunctionRetryPolicy}; + /// Semantic Function type for a Blob v2 value. pub const FUNCTION_BLOB_V2_TYPE: &str = "blob_v2"; @@ -787,6 +790,9 @@ fn validate_row_timeout_value(value: Option) -> std::result::Result<(), &'s /// database already holds, which Sophon resolves inside the remote runtime. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct FunctionRegistrationRequest { + /// Default retries of failed user-code calls. Enterprise only. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub retry_policy: Option, pub name: String, pub artifact: FunctionArtifactRequest, pub signature: FunctionSignature, @@ -836,6 +842,9 @@ impl FunctionRegistrationRequest { }; validate_batch_size_value(self.batch_size).map_err(invalid_input)?; validate_row_timeout_value(self.row_timeout_seconds).map_err(invalid_input)?; + if let Some(policy) = &self.retry_policy { + policy.validate()?; + } Ok(()) } } @@ -1094,6 +1103,9 @@ impl FunctionErrorsRequest { /// read access to the table. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct FunctionErrorRecord { + /// Number of user-code calls for this failed invocation, when recorded. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub attempts: Option, /// The refresh job that recorded the error. pub job_id: String, /// The fragment holding the row. diff --git a/rust/lancedb/src/function/retry.rs b/rust/lancedb/src/function/retry.rs new file mode 100644 index 000000000..b52d59d2b --- /dev/null +++ b/rust/lancedb/src/function/retry.rs @@ -0,0 +1,160 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Client configuration for retries inside an Enterprise Function invocation. + +use serde::{Deserialize, Serialize}; + +use crate::{Error, Result}; + +/// Retry selected user-code errors without consuming a worker retry attempt. +/// +/// A retry repeats the entire failed call, so opt in only when repeating its +/// partial work is safe. Matchers are alternatives. With a batch Function and +/// `skip`, an exhausted call skips the whole batch rather than bisecting it. +/// +/// ``` +/// use lancedb::function::FunctionRetryPolicy; +/// let policy = FunctionRetryPolicy::transient(); +/// policy.validate().unwrap(); +/// ``` +#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)] +#[serde(default, deny_unknown_fields)] +#[non_exhaustive] +pub struct FunctionRetryPolicy { + /// Total calls, including the first; 1 disables retries, maximum 100. + pub max_attempts: u32, + /// Initial exponential-backoff ceiling, with full jitter; default 0.2 seconds. + pub initial_backoff_seconds: f64, + /// Maximum backoff, including Retry-After; default 30 seconds, maximum 3600. + pub max_backoff_seconds: f64, + /// Exception class names, preferably qualified; subclasses also match. + pub exception_types: Vec, + /// Regular expression matched against the exception class and message. + pub message_pattern: Option, + /// Structured provider error codes or HTTP status codes, represented as strings. + pub error_codes: Vec, + /// Match connection/timeouts, HTTP 429/503, SlowDown and ServiceUnavailable. + pub transient: bool, + /// Honor Retry-After within the configured maximum; default true. + pub honor_retry_after: bool, +} + +impl Default for FunctionRetryPolicy { + fn default() -> Self { + Self { + max_attempts: 3, + initial_backoff_seconds: 0.2, + max_backoff_seconds: 30.0, + exception_types: Vec::new(), + message_pattern: None, + error_codes: Vec::new(), + transient: false, + honor_retry_after: true, + } + } +} + +impl FunctionRetryPolicy { + /// Retry known transient errors up to three calls with bounded backoff. + pub fn transient() -> Self { + Self { + transient: true, + ..Self::default() + } + } + + /// Validate before registration or refresh submission. + pub fn validate(&self) -> Result<()> { + let invalid = |message: String| Error::InvalidInput { message }; + if !(1..=100).contains(&self.max_attempts) { + return Err(invalid(format!( + "retry_policy.max_attempts must be between 1 and 100, got {}", + self.max_attempts + ))); + } + for (name, value) in [ + ("initial_backoff_seconds", self.initial_backoff_seconds), + ("max_backoff_seconds", self.max_backoff_seconds), + ] { + if !value.is_finite() || !(0.0..=3600.0).contains(&value) { + return Err(invalid(format!( + "retry_policy.{name} must be finite and between 0 and 3600, got {value}" + ))); + } + } + if self.initial_backoff_seconds > self.max_backoff_seconds { + return Err(invalid( + "retry_policy.initial_backoff_seconds exceeds max_backoff_seconds".into(), + )); + } + if self + .exception_types + .iter() + .chain(&self.error_codes) + .any(String::is_empty) + { + return Err(invalid( + "retry_policy exception_types and error_codes must contain non-empty strings" + .into(), + )); + } + if let Some(pattern) = &self.message_pattern { + regex::Regex::new(pattern).map_err(|error| { + invalid(format!("invalid retry_policy.message_pattern: {error}")) + })?; + } + if self.max_attempts > 1 + && !self.transient + && self.message_pattern.is_none() + && self.exception_types.is_empty() + && self.error_codes.is_empty() + { + return Err(invalid( + "retry_policy requires at least one error matcher".into(), + )); + } + Ok(()) + } +} + +/// What a Function refresh does after a call fails or exhausts retries. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum FunctionErrorPolicy { + /// Fail the refresh job. + Fail, + /// Write null and record the error, subject to the server's failure limits. + Skip, +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + use serde_json::json; + + #[rstest] + #[case(json!({"max_attempts": 0, "transient": true}))] + #[case(json!({"max_attempts": 101, "transient": true}))] + #[case(json!({"max_attempts": 3}))] + #[case(json!({"initial_backoff_seconds": -1, "transient": true}))] + #[case(json!({"max_backoff_seconds": 0, "transient": true}))] + #[case(json!({"exception_types": [""]}))] + #[case(json!({"message_pattern": "["}))] + fn reject_invalid_policy(#[case] value: serde_json::Value) { + let policy: FunctionRetryPolicy = serde_json::from_value(value).unwrap(); + assert!(policy.validate().is_err()); + } + + #[test] + fn policy_defaults_and_disabled_policy() { + assert!(FunctionRetryPolicy::transient().validate().is_ok()); + let policy: FunctionRetryPolicy = + serde_json::from_value(json!({"max_attempts":1})).unwrap(); + assert!(policy.validate().is_ok()); + let mut invalid = policy; + invalid.initial_backoff_seconds = f64::NAN; + assert!(invalid.validate().is_err()); + } +} diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 6534f5519..5dee8eb1d 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -3977,6 +3977,12 @@ impl BaseTable for RemoteTable { if let Some(gpu_count) = options.gpu_count { body["gpu_count"] = serde_json::json!(gpu_count); } + if let Some(policy) = options.retry_policy { + body["retry_policy"] = serde_json::json!(policy); + } + if let Some(policy) = options.error_policy { + body["error_policy"] = serde_json::json!(policy); + } self.apply_branch_body(&mut body); let request = self .client @@ -10980,6 +10986,38 @@ mod tests { assert_eq!(job.id(), Some("j-gpu")); } + #[tokio::test] + async fn test_refresh_column_async_retry_policy() { + use crate::function::{FunctionErrorPolicy, FunctionRetryPolicy}; + use crate::table::RefreshColumnOptions; + let policy = FunctionRetryPolicy::transient(); + let expected_policy = serde_json::to_value(&policy).unwrap(); + let table = Table::new_with_handler("my_table", move |request| { + assert_eq!(request.url().path(), "/v1/table/my_table/backfill_column"); + assert_eq!( + request_body_json(&request), + serde_json::json!({ + "column":"quality", "retry_policy": expected_policy, "error_policy":"skip", + }) + ); + http::Response::builder() + .status(202) + .body(r#"{"job_id":"j-retry"}"#.to_string()) + .unwrap() + }); + let options = RefreshColumnOptions::default() + .retry_policy(policy) + .error_policy(FunctionErrorPolicy::Skip); + assert_eq!( + table + .refresh_column_async_with_options("quality", options) + .await + .unwrap() + .id(), + Some("j-retry") + ); + } + #[tokio::test] async fn test_refresh_column_rejects_invalid_gpu_options_before_request() { use crate::table::RefreshColumnOptions; @@ -11120,7 +11158,7 @@ mod tests { "column": "embedding", "function": "embed", "function_version": "2", "table_version": 11, "error_type": "ValueError", "error_message": "bad input 'x'", "traceback": "Traceback (most recent call last):\nValueError: bad input 'x'", - "created_at_millis": 1700000000000}], + "created_at_millis": 1700000000000, "attempts": 3}], "fragments": [{"job_id": "j-7", "fragment_id": 4, "rows_skipped": 500, "rows_recorded": 100}], "truncated": true}"#, @@ -11140,6 +11178,7 @@ mod tests { assert_eq!( errors.records, [FunctionErrorRecord { + attempts: Some(3), job_id: "j-7".into(), fragment_id: 3, row_offset: Some(9), diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 34846d33e..94a934c43 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -891,6 +891,12 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { options: RefreshColumnOptions, ) -> Result { options.validate()?; + if options.retry_policy.is_some() || options.error_policy.is_some() { + return Err(Error::NotSupported { + message: "Function retry and error policies require an Enterprise Function refresh" + .into(), + }); + } self.refresh_column(column).await } /// Fill a computed column's unfilled rows and recompute those whose @@ -910,6 +916,12 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { options: RefreshColumnOptions, ) -> Result> { options.validate()?; + if options.retry_policy.is_some() || options.error_policy.is_some() { + return Err(Error::NotSupported { + message: "Function retry and error policies require an Enterprise Function refresh" + .into(), + }); + } self.refresh_column_async(column).await } /// The per-row errors Function refreshes recorded on this table; see diff --git a/rust/lancedb/src/table/refresh.rs b/rust/lancedb/src/table/refresh.rs index 330cdfb6d..a9e69452d 100644 --- a/rust/lancedb/src/table/refresh.rs +++ b/rust/lancedb/src/table/refresh.rs @@ -59,6 +59,10 @@ use crate::{Error, Result}; #[derive(Debug, Clone, Default)] #[non_exhaustive] pub struct RefreshColumnOptions { + /// Override the Function's default user-code retry policy. Enterprise only. + pub retry_policy: Option, + /// Action after user-code failure or retry exhaustion. Enterprise only. + pub error_policy: Option, /// Maximum rows per Function call in LanceDB Enterprise. /// /// Must be greater than zero. `None` uses the server default. This has no @@ -77,6 +81,17 @@ pub struct RefreshColumnOptions { } impl RefreshColumnOptions { + /// Override retries of failed Function calls for this refresh. + pub fn retry_policy(mut self, policy: crate::function::FunctionRetryPolicy) -> Self { + self.retry_policy = Some(policy); + self + } + + /// Select fail or skip after the retry policy is exhausted. + pub fn error_policy(mut self, policy: crate::function::FunctionErrorPolicy) -> Self { + self.error_policy = Some(policy); + self + } /// Limit each Function call to at most `batch_size` rows. pub fn batch_size(mut self, batch_size: usize) -> Self { self.batch_size = Some(batch_size); @@ -97,6 +112,9 @@ impl RefreshColumnOptions { /// Validate client-side constraints before any request is sent. pub fn validate(&self) -> Result<()> { + if let Some(policy) = &self.retry_policy { + policy.validate()?; + } if self.batch_size == Some(0) { return Err(Error::InvalidInput { message: "batch_size must be greater than zero".into(),