diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 36cf36ddf..2d9de72c4 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -947,6 +947,7 @@ class FunctionErrorRecord: job_id: str fragment_id: int row_offset: Optional[int] + row_address: Optional[int] column: str function: str function_version: str diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index 1d80d0d32..721ba5ced 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -16,7 +16,7 @@ import builtins import base64 import functools import hashlib -import importlib +import importlib.util import inspect import symtable import json @@ -1464,6 +1464,28 @@ def _module_locations(module: types.ModuleType) -> tuple[str, ...]: return tuple(os.path.realpath(path) for path in getattr(module, "__path__", ())) +def _module_locations_for_name(module_name: str) -> tuple[str, ...]: + """Where a top-level module would be imported from, without importing it + when it is not already loaded.""" + import os + + top = module_name.partition(".")[0] + module = sys.modules.get(top) + if module is not None: + return _module_locations(module) + try: + spec = importlib.util.find_spec(top) + except (AttributeError, ImportError, ValueError): + return () + if spec is None: + return () + if spec.submodule_search_locations: + return tuple(os.path.realpath(path) for path in spec.submodule_search_locations) + if spec.origin and spec.origin not in {"built-in", "frozen"}: + return (os.path.realpath(spec.origin),) + return () + + class _SourcePackager: """Builds the entry module of a Function artifact. @@ -1480,6 +1502,7 @@ class _SourcePackager: self._definitions: list[str] = [] self._packaged: set[int] = set() self._names: dict[str, Any] = {} + self._warned_imports: set[str] = set() def package(self, target: Any, signature: FunctionSignature) -> str: entry = self._definition(target, entry=True, signature=signature) @@ -1512,6 +1535,7 @@ class _SourcePackager: raise ValueError( f"@udf cannot package dynamic namespace access: {dynamic!r}" ) + self._warn_unshipped_direct_imports(definition) self._packaged.add(id(target)) self._claim(target.__name__, ("definition", id(target))) if reader_name is not None: @@ -1652,13 +1676,14 @@ class _SourcePackager: or top in sys.stdlib_module_names or top in sys.builtin_module_names or _normalized_distribution(top) in self._requirements + or top in self._warned_imports ): return - module = sys.modules.get(top) - locations = () if module is None else _module_locations(module) + locations = _module_locations_for_name(top) roots = _installed_roots() if not locations or any(location.startswith(roots) for location in locations): return + self._warned_imports.add(top) warnings.warn( f"@udf source imports {module_name!r} from {locations[0]}, a local " "module the Function's environment may not be able to import; ship " @@ -1667,6 +1692,22 @@ class _SourcePackager: UserWarning, ) + def _warn_unshipped_direct_imports( + self, definition: Union[ast.FunctionDef, ast.ClassDef] + ) -> None: + """Warn for body imports that resolve to local, undeclared source. + + Imports from site-packages are deliberately not guessed about: the + declared distribution name need not match the imported module, and the + client cannot inspect the remote base image. + """ + for node in ast.walk(definition): + if isinstance(node, ast.Import): + for alias in node.names: + self._require_importable(alias.name) + elif isinstance(node, ast.ImportFrom) and node.level == 0 and node.module: + self._require_importable(node.module) + def _code_files(modules: Sequence[Any]) -> tuple[dict[str, str], frozenset[str]]: """The Python source files of the modules and packages in ``code=``, @@ -2147,9 +2188,15 @@ def udf( imported normally, so helpers shared by several Functions live in one place. - A reference to a module in a local source tree that is neither shipped - with ``code`` nor named by a declared package draws a warning at - registration, because the worker will most likely fail to import it. + A module imported by the packaged source -- including a direct ``import`` + inside a function or method body -- draws a warning at registration when + it resolves to local source outside ``code`` and no declared requirement + matches the imported module name. The client cannot reliably validate + imports provided only by the base image: ``pip``/``conda`` distribution + names need not match the modules they install, and the client does not + know the contents of a built or custom image. Declare every third-party + import in ``pip`` or ``conda`` even when it happens to be installed on the + registration host. Closures, lambdas, nested definitions, mutable module-level objects, code that reaches the module namespace another way (``globals()``/``eval``, ``sys.modules``, ``builtins``), and a diff --git a/python/python/lancedb/job.py b/python/python/lancedb/job.py index e61a188f2..5646a9e10 100644 --- a/python/python/lancedb/job.py +++ b/python/python/lancedb/job.py @@ -26,6 +26,28 @@ __all__ = [ ] +def _timeout_message( + job_id: Optional[str], timeout: timedelta, job_type: Optional[str] +) -> str: + seconds = f"{timeout.total_seconds():g}" + unit = "second" if seconds == "1" else "seconds" + if job_id is None: + return ( + f"operation still running after {seconds} {unit}; " + "it continues on the server" + ) + if job_type == "query": + return ( + f"query {job_id!r} still running after {seconds} {unit}; " + "it continues on the server; keep this handle to continue waiting" + ) + return ( + f"job {job_id!r} still running after {seconds} {unit}; " + "it continues on the server; " + f"reattach with db.open_job({job_id!r}).wait()" + ) + + class AsyncJob(Generic[T]): """A handle to an operation that may still be running. @@ -73,13 +95,24 @@ class AsyncJob(Generic[T]): Raises `JobFailedError` if the operation failed, `JobCancelledError` if it was cancelled, and `TimeoutError` if `timeout` elapses first. + The timeout error identifies the job and explains that it continues on + the server. Registry jobs include the ``db.open_job(id).wait()`` + reattach call; SQL query handles must be kept because their IDs cannot + be reopened through the registry jobs API. """ if self._inner is None: return cast(T, None) if timeout is None: result = await self._inner.wait() else: - result = await asyncio.wait_for(self._inner.wait(), timeout.total_seconds()) + try: + result = await asyncio.wait_for( + self._inner.wait(), timeout.total_seconds() + ) + except asyncio.TimeoutError as error: + raise TimeoutError( + _timeout_message(self.id, timeout, self.job_type) + ) from error if self._result_decoder is not None: return self._result_decoder(result) return cast(T, result) @@ -270,6 +303,10 @@ class Job(Generic[T]): Raises `JobFailedError` if the operation failed, `JobCancelledError` if it was cancelled, and `TimeoutError` if `timeout` elapses first. + The timeout error identifies the job and explains that it continues on + the server. Registry jobs include the ``db.open_job(id).wait()`` + reattach call; SQL query handles must be kept because their IDs cannot + be reopened through the registry jobs API. """ if self._inner is None: return cast(T, None) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 7d82409ef..e16a29df1 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -3191,6 +3191,30 @@ class Table(ABC): ``records``, ``fragments`` and ``truncated``, the last saying whether the listing stopped at its limit. + Each record identifies its row with ``fragment_id``, ``row_offset`` + and the physical ``row_address`` (``(fragment_id << 32) | + row_offset``). To retry the recorded rows, collect the non-null + addresses before the table is compacted or rewritten and run the + SQL refresh: + + .. code-block:: python + + errors = table.function_errors(column="embedding") + addresses = sorted( + record.row_address + for record in errors.records + if record.row_address is not None + ) + if addresses: + address_list = ", ".join(str(address) for address in addresses) + database.execute_query( + f"REFRESH COLUMN embedding ON {table.name()} " + f"WHERE _rowaddr IN ({address_list})" + ) + + Records whose per-row detail was capped have no address and appear + only in ``fragments``; they cannot be selected individually here. + Examples -------- >>> errors = table.function_errors(column="embedding") # doctest: +SKIP @@ -7515,6 +7539,30 @@ class AsyncTable: ``records``, ``fragments`` and ``truncated``, the last saying whether the listing stopped at its limit. + Each record identifies its row with ``fragment_id``, ``row_offset`` + and the physical ``row_address`` (``(fragment_id << 32) | + row_offset``). To retry the recorded rows, collect the non-null + addresses before the table is compacted or rewritten and run the + SQL refresh: + + .. code-block:: python + + errors = await table.function_errors(column="embedding") + addresses = sorted( + record.row_address + for record in errors.records + if record.row_address is not None + ) + if addresses: + address_list = ", ".join(str(address) for address in addresses) + await database.execute_query( + f"REFRESH COLUMN embedding ON {table.name()} " + f"WHERE _rowaddr IN ({address_list})" + ) + + Records whose per-row detail was capped have no address and appear + only in ``fragments``; they cannot be selected individually here. + Examples -------- >>> errors = await table.function_errors(column="embedding") # doctest: +SKIP diff --git a/python/python/tests/test_function_packaging.py b/python/python/tests/test_function_packaging.py index a61588777..169e54a36 100644 --- a/python/python/tests/test_function_packaging.py +++ b/python/python/tests/test_function_packaging.py @@ -388,6 +388,7 @@ def test_code_packages_ship_as_a_bundle_and_version_with_their_source( def define(): sys.modules.pop("fx_helpers_functions", None) + importlib.invalidate_caches() return importlib.import_module("fx_helpers_functions").limited first = define() @@ -454,6 +455,7 @@ def test_local_module_outside_code_warns_at_registration(tmp_path, helper_packag """ ) ) + importlib.invalidate_caches() return getattr(importlib.import_module(name), name) with pytest.warns(UserWarning, match=r"local module.*code=\[fx_local\]"): @@ -467,6 +469,44 @@ def test_local_module_outside_code_warns_at_registration(tmp_path, helper_packag assert "from fx_local import normalize" in artifact(installed).decode() +def test_local_module_imported_in_function_body_warns_at_registration( + tmp_path, helper_package +): + """A direct import inside the packaged definition is checked too.""" + helper_package( + "fx_body_local", + {"__init__.py": "def normalize(value):\n return value\n"}, + ) + sys.modules.pop("fx_body_local", None) + + def define(name, options): + (tmp_path / f"{name}.py").write_text( + textwrap.dedent( + f""" + from lancedb import udf + + + @udf({options}) + def {name}(value: int) -> int: + import fx_body_local + + return fx_body_local.normalize(value) + """ + ) + ) + importlib.invalidate_caches() + return getattr(importlib.import_module(name), name) + + with pytest.warns(UserWarning, match=r"local module.*code=\[fx_body_local\]"): + local = define("fx_body_uses_local", "") + assert "import fx_body_local" in artifact(local).decode() + + with warnings.catch_warnings(): + warnings.simplefilter("error") + installed = define("fx_body_uses_installed", "pip=['fx-body-local==1.0']") + assert "import fx_body_local" in artifact(installed).decode() + + MMLB_HELPERS = { "__init__.py": "", "rate.py": """ @@ -527,6 +567,7 @@ def test_shared_helpers_construct_state_on_every_path(tmp_path, helper_package): initialized through the same helper, so neither path can skip it.""" helper_package("fx_mmlb", MMLB_HELPERS) (tmp_path / "fx_mmlb_functions.py").write_text(MMLB_FUNCTIONS) + importlib.invalidate_caches() functions = importlib.import_module("fx_mmlb_functions") signatures = [ diff --git a/python/python/tests/test_job.py b/python/python/tests/test_job.py new file mode 100644 index 000000000..19182485f --- /dev/null +++ b/python/python/tests/test_job.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import asyncio +from datetime import timedelta + +import pytest + +from lancedb.job import AsyncJob, Job + + +class _PendingJob: + id = "job-timeout" + job_type = None + + async def wait(self): + await asyncio.sleep(60) + + +class _PendingQuery: + id = "query-timeout" + job_type = "query" + + async def wait(self): + await asyncio.sleep(60) + + +def _assert_timeout(error: TimeoutError) -> None: + message = str(error) + assert "job 'job-timeout' still running after 0.01 seconds" in message + assert "continues on the server" in message + assert "db.open_job('job-timeout').wait()" in message + + +@pytest.mark.asyncio +async def test_async_job_wait_timeout_explains_reattach(): + job = AsyncJob(_PendingJob()) + + with pytest.raises(TimeoutError) as exc_info: + await job.wait(timeout=timedelta(seconds=0.01)) + + _assert_timeout(exc_info.value) + + +def test_job_wait_timeout_explains_reattach(): + job = Job(AsyncJob(_PendingJob())) + + with pytest.raises(TimeoutError) as exc_info: + job.wait(timeout=timedelta(seconds=0.01)) + + _assert_timeout(exc_info.value) + + +@pytest.mark.asyncio +async def test_query_wait_timeout_requires_keeping_original_handle(): + job = AsyncJob(_PendingQuery()) + + with pytest.raises(TimeoutError) as exc_info: + await job.wait(timeout=timedelta(seconds=0.01)) + + message = str(exc_info.value) + assert "query 'query-timeout' still running after 0.01 seconds" in message + assert "continues on the server" in message + assert "keep this handle to continue waiting" in message + assert "db.open_job" not in message diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index ef48d8eff..b2350104b 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -1795,6 +1795,7 @@ def test_remote_function_errors_lists_the_rows_a_refresh_skipped(): "job_id": "j-7", "fragment_id": 3, "row_offset": 9, + "row_address": 12884901897, "column": "embedding", "function": "embed", "function_version": "2", @@ -1845,7 +1846,13 @@ def test_remote_function_errors_lists_the_rows_a_refresh_skipped(): assert errors.truncated is True assert [r.error_message for r in errors.records] == ["bad input 'x'"] assert errors.records[0].row_offset == 9 + assert errors.records[0].row_address == 12884901897 assert errors.records[0].function_version == "2" + assert repr(errors.records[0]) == ( + 'FunctionErrorRecord(job_id="j-7", fragment_id=3, row_offset=9, ' + 'row_address=12884901897, column="embedding", ' + 'error_type="ValueError", error_message="bad input \'x\'")' + ) assert (errors.fragments[0].rows_skipped, errors.fragments[0].rows_recorded) == ( 500, 100, @@ -4812,6 +4819,9 @@ def test_remote_connection_jobs_surface(): assert job.state == "failed" assert job.spec == {"column": "vec"} assert job.failure.message == "worker died" + assert repr(job.failure) == ( + 'JobFailureInfo(phase="execute", message="worker died", retryable=True)' + ) assert job.status() == "failed" with pytest.raises(JobFailedError, match="worker died"): job.wait(timeout=timedelta(seconds=5)) diff --git a/python/src/job.rs b/python/src/job.rs index 1740b3e00..0a01d9faf 100644 --- a/python/src/job.rs +++ b/python/src/job.rs @@ -200,10 +200,20 @@ pub struct JobFailureInfo { #[pymethods] impl JobFailureInfo { fn __repr__(&self) -> String { - format!( - "JobFailureInfo(phase={:?}, message={:?}, retryable={:?})", - self.phase, self.message, self.retryable - ) + let phase = self + .phase + .as_ref() + .map_or_else(|| "None".to_string(), |value| format!("{value:?}")); + let message = self + .message + .as_ref() + .map_or_else(|| "None".to_string(), |value| format!("{value:?}")); + let retryable = match self.retryable { + Some(true) => "True", + Some(false) => "False", + None => "None", + }; + format!("JobFailureInfo(phase={phase}, message={message}, retryable={retryable})") } } diff --git a/python/src/table.rs b/python/src/table.rs index 83e18ef72..d9e006a38 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -596,6 +596,7 @@ pub struct FunctionErrorRecord { pub job_id: String, pub fragment_id: u64, pub row_offset: Option, + pub row_address: Option, pub column: String, pub function: String, pub function_version: String, @@ -608,12 +609,19 @@ pub struct FunctionErrorRecord { #[pymethods] impl FunctionErrorRecord { pub fn __repr__(&self) -> String { + let row_offset = self + .row_offset + .map_or_else(|| "None".to_string(), |value| value.to_string()); + let row_address = self + .row_address + .map_or_else(|| "None".to_string(), |value| value.to_string()); format!( - "FunctionErrorRecord(job_id={:?}, fragment_id={}, row_offset={:?}, column={:?}, \ - error_type={:?}, error_message={:?})", + "FunctionErrorRecord(job_id={:?}, fragment_id={}, row_offset={}, row_address={}, \ + column={:?}, error_type={:?}, error_message={:?})", self.job_id, self.fragment_id, - self.row_offset, + row_offset, + row_address, self.column, self.error_type, self.error_message @@ -627,6 +635,7 @@ impl From for FunctionErrorRecord { job_id: record.job_id, fragment_id: record.fragment_id, row_offset: record.row_offset, + row_address: record.row_address, column: record.column, function: record.function, function_version: record.function_version, diff --git a/rust/lancedb/src/function.rs b/rust/lancedb/src/function.rs index be06be2fa..a9797254e 100644 --- a/rust/lancedb/src/function.rs +++ b/rust/lancedb/src/function.rs @@ -867,6 +867,12 @@ pub struct FunctionErrorRecord { /// detail was capped and only the fragment summary remains. #[serde(default)] pub row_offset: Option, + /// Portable physical row address, `(fragment_id << 32) | row_offset`. + /// + /// `None` when the record covers a fragment rather than one row, or when + /// an older server does not report row addresses. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub row_address: Option, /// The column being computed. pub column: String, /// The Function that failed. diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index fad790391..adc0c1bfc 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -10673,6 +10673,7 @@ mod tests { .status(200) .body( r#"{"records": [{"job_id": "j-7", "fragment_id": 3, "row_offset": 9, + "row_address": 12884901897, "column": "embedding", "function": "embed", "function_version": "2", "table_version": 11, "error_type": "ValueError", "error_message": "bad input 'x'", "created_at_millis": 1700000000000}], @@ -10698,6 +10699,7 @@ mod tests { job_id: "j-7".into(), fragment_id: 3, row_offset: Some(9), + row_address: Some(12_884_901_897), column: "embedding".into(), function: "embed".into(), function_version: "2".into(), @@ -10718,6 +10720,22 @@ mod tests { ); assert!(errors.truncated); + // Servers before row_address omit it; the record still decodes. + let legacy: FunctionErrorRecord = serde_json::from_value(serde_json::json!({ + "job_id": "j-6", + "fragment_id": 2, + "row_offset": null, + "column": "embedding", + "function": "embed", + "function_version": "1", + "table_version": 10, + "error_type": "ValueError", + "error_message": "fragment reached the record cap", + "created_at_millis": 1_700_000_000_000_i64, + })) + .unwrap(); + assert!(legacy.row_address.is_none()); + // No filter sends no filter, and an empty listing reads as such. let table = Table::new_with_handler("my_table", |request| { let body = request.body().unwrap().as_bytes().unwrap();