From 0b18dcdaec6c5a94717f34d98a9e02f3d8525e53 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Thu, 24 Sep 2026 18:02:25 +0800 Subject: [PATCH] fix(python): warn instead of failing on local imports in @udf source (#4318) Since 0.40.0b8 (#4254), `@udf` raises when the packaged source imports a module from a local source tree that is neither shipped with `code=` nor named by a declared pip/conda package. That turned registrations that used to succeed into errors: - Callers that rewrite the packaged source before registering it now fail at decoration. Sophon's datagen integration suite inlines its test-module helpers this way, and `test_registered_udf_source_stands_alone_on_a_worker` fails on 0.40.0b8. - The check cannot know which modules a distribution installs. A declared package whose distribution name differs from its module (for example an internal library installed editable) is reported as local, and there was no way to proceed. This keeps the diagnosis but makes it a `UserWarning`, so registration behaves as it did before #4254 and still tells the author, at registration time, why the worker is likely to fail. --- python/python/lancedb/functions.py | 25 +++++++++++++------ .../python/tests/test_function_packaging.py | 18 ++++++++----- 2 files changed, 29 insertions(+), 14 deletions(-) diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index 380625f57..5f26b590e 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -26,6 +26,7 @@ import re import sys import textwrap import types +import warnings from collections.abc import Mapping, Sequence from datetime import date, datetime from typing import ( @@ -1527,9 +1528,14 @@ class _SourcePackager: return f"{name} = {_literal_source(value)}" def _require_importable(self, module_name: str) -> None: - """Refuse an import the Function's environment has no way to satisfy: - a module that lives in a local source tree, is not shipped with - ``code=``, and names no declared package.""" + """Warn about an import the Function's environment likely cannot + satisfy: a module that lives in a local source tree, is not shipped + with ``code=``, and names no declared package. + + A warning rather than an error: a distribution name need not match + the module it installs, so a declared package can still provide the + module, and callers may rewrite the packaged source before + registering it.""" top = module_name.partition(".")[0] if ( top in self._code_modules @@ -1543,10 +1549,12 @@ class _SourcePackager: roots = _installed_roots() if not locations or any(location.startswith(roots) for location in locations): return - raise ValueError( + warnings.warn( f"@udf source imports {module_name!r} from {locations[0]}, a local " - "module the Function's environment cannot import; ship it with " - f"code=[{top}] or declare the package that provides it in pip/conda" + "module the Function's environment may not be able to import; ship " + f"it with code=[{top}] or declare the package that provides it in " + "pip/conda", + UserWarning, ) @@ -1991,8 +1999,9 @@ def udf( place. A reference to a module in a local source tree that is neither shipped - with ``code`` nor provided by a declared package is rejected at - registration. Closures, lambdas, nested definitions, mutable module-level + with ``code`` nor named by a declared package draws a warning at + registration, because the worker will most likely fail to import it. + Closures, lambdas, nested definitions, mutable module-level objects, code that reaches the module namespace another way (``globals()``/``eval``, ``sys.modules``, ``builtins``), and a non-standard ``__builtins__`` are rejected where they can be seen and diff --git a/python/python/tests/test_function_packaging.py b/python/python/tests/test_function_packaging.py index 239d07be1..6e1e1cf43 100644 --- a/python/python/tests/test_function_packaging.py +++ b/python/python/tests/test_function_packaging.py @@ -11,6 +11,7 @@ import json import subprocess import sys import textwrap +import warnings from pathlib import Path from typing import Annotated, Optional @@ -299,9 +300,11 @@ def test_code_rejects_submodules_duplicates_and_non_modules(helper_package): return value -def test_local_module_outside_code_is_rejected_at_registration( - tmp_path, helper_package -): +def test_local_module_outside_code_warns_at_registration(tmp_path, helper_package): + """The worker most likely cannot import a local module that is neither + shipped nor declared, but registration still succeeds: a declared + distribution may install it under another name, and callers may rewrite + the packaged source before registering it.""" helper_package( "fx_local", {"__init__.py": "def normalize(value):\n return value\n"} ) @@ -322,11 +325,14 @@ def test_local_module_outside_code_is_rejected_at_registration( ) return getattr(importlib.import_module(name), name) - with pytest.raises(ValueError, match=r"local module.*code=\[fx_local\]"): - define("fx_uses_local", "") + with pytest.warns(UserWarning, match=r"local module.*code=\[fx_local\]"): + local = define("fx_uses_local", "") + assert "from fx_local import normalize" in artifact(local).decode() # A declared package that provides the module is trusted to install it. - installed = define("fx_uses_installed", "pip=['fx-local==1.0']") + with warnings.catch_warnings(): + warnings.simplefilter("error") + installed = define("fx_uses_installed", "pip=['fx-local==1.0']") assert "from fx_local import normalize" in artifact(installed).decode()