feat(python): warn on scalar GPU Functions (#4966)

GPU Functions that use a scalar callable currently invoke the callable
once per row. The SDK documents Arrow-vectorized callables, but silently
accepts a scalar GPU callable, which makes it easy to assume that the
runtime will batch model work automatically.

Warn during Function construction when a GPU callable does not use the
documented vectorized form. A callable is recognized as vectorized when
every input is annotated as `pyarrow.Array`, or when a single input is
annotated as `pyarrow.RecordBatch`. The warning explains the resulting
per-row invocation and shows how to opt into batch execution. CPU
Functions and the vectorized GPU forms do not emit it.

The warning is advisory and does not change registration serialization,
the artifact, the runtime adapter, or execution behavior. It is one part
of the ENT-3522 documentation-and-warning alternative; it does not claim
that a single A10G meets the original batch throughput requirement.

Validation on the designated EC2 host reused the compatible existing
native extension from `ent3563-sdk-freshness-9336cfb`; this does not
claim a new native package or GPU performance run. The final head has a
clean source tree and ran:

```bash
PYTHONPATH="$sdk/python/python" /home/ec2-user/.venvs/ent3563-0a839fa/bin/python -m pytest -q python/python/tests/test_first_class_function_slice2.py -k 'gpu_marker_uses_gpu_runtime or gpu_scalar_warning_names_the_vectorized_alternative'
/home/ec2-user/.venvs/ent3563-0a839fa/bin/ruff format --check python/python/lancedb/functions.py python/python/tests/test_first_class_function_slice2.py
/home/ec2-user/.venvs/ent3563-0a839fa/bin/ruff check python/python/lancedb/functions.py python/python/tests/test_first_class_function_slice2.py
```

The actual result was `2 passed, 123 deselected`; both Ruff commands
passed. The original 2,000-row single-A10G throughput requirement
remains a separate ENT-3522 acceptance gate.
This commit is contained in:
Xuanwo
2026-11-13 02:04:32 +08:00
committed by GitHub
parent e61e21035c
commit 8939ef1f01
2 changed files with 86 additions and 4 deletions
+26 -1
View File
@@ -1529,6 +1529,17 @@ def _infer_signature(
return FunctionSignature(inputs=tuple(inputs), output=output)
def _is_vectorized_callable(function: Callable[..., Any], *, method: bool) -> bool:
parameters = _callable_parameters(function)
if method:
parameters = parameters[1:]
annotations = get_type_hints(function, include_extras=True)
annotated = tuple(annotations.get(parameter.name) for parameter in parameters)
return (len(annotated) == 1 and annotated[0] is pa.RecordBatch) or (
bool(annotated) and all(annotation is pa.Array for annotation in annotated)
)
# Both spellings strip from packaged source; `udf` is the deprecated alias.
_DECORATOR_NAMES = frozenset({"function", "udf"})
@@ -2302,7 +2313,9 @@ class FunctionDefinition:
once per row. A callable whose parameters are all annotated as
``pyarrow.Array`` receives whole Arrow arrays and returns a
``pyarrow.Array`` instead. This authoring choice is independent of the
internal artifact adapter in the registration request.
internal artifact adapter in the registration request. GPU Functions should
use the Arrow form when the model supports batches: a scalar GPU callable is
invoked once per row and leaves the model's own batching unavailable.
"""
def __init__(
@@ -2424,6 +2437,18 @@ class FunctionDefinition:
signature = signature._copy(update={"initialization": initialization})
else:
signature = _infer_signature(function, input_schema, output_schema)
if gpu_marker is not None:
call = _class_call(function) if is_class else function
if not _is_vectorized_callable(call, method=is_class):
warnings.warn(
f"GPU Function {function_name!r} uses a scalar callable, so each "
"row invokes it separately and its model batch is not inferred. "
"For GPU throughput, annotate every input as pyarrow.Array or "
"take one pyarrow.RecordBatch, provide input_schema and "
"output_schema, and batch inside that callable.",
UserWarning,
stacklevel=3,
)
code_files, code_modules = _code_files(code)
entry_source = _SourcePackager(
code_modules, frozenset(_requirement_name(package) for package in packages)
@@ -18,6 +18,7 @@ import subprocess
import sys
import typing
import urllib.parse
import warnings
from typing import Optional
import pyarrow as pa
@@ -658,9 +659,11 @@ def test_function_image_requires_digest_reference(image):
def test_function_gpu_marker_uses_gpu_runtime():
@function(pip=["cupy-cuda12x"], gpu=True, gpu_class="a10g")
def double_on_gpu(value: int) -> int:
return value * 2
with pytest.warns(UserWarning, match="uses a scalar callable"):
@function(pip=["cupy-cuda12x"], gpu=True, gpu_class="a10g")
def double_on_gpu(value: int) -> int:
return value * 2
request = json.loads(double_on_gpu.registration_request.to_canonical_json())
assert request["runtime"]["kind"] == "python_v2"
@@ -710,6 +713,60 @@ def test_function_gpu_marker_uses_gpu_runtime():
)
class _GpuScalarFunction:
def __call__(self, value: str) -> int:
return len(value)
class _GpuArrayFunction:
def __call__(self, value: pa.Array) -> pa.Array:
return value
def test_gpu_scalar_warning_names_the_vectorized_alternative():
with pytest.warns(UserWarning, match="model batch is not inferred"):
@function(gpu=True)
def scalar_embedding(value: str) -> list[int]:
return [len(value)]
input_schema = pa.schema([pa.field("value", pa.string(), nullable=False)])
with pytest.warns(UserWarning, match="uses a scalar callable"):
class_scalar = function(_GpuScalarFunction, gpu=True)
with warnings.catch_warnings(record=True) as caught:
warnings.simplefilter("always")
@function(
gpu=True, input_schema=input_schema, output_schema=pa.list_(pa.int64())
)
def array_embedding(value: pa.Array) -> pa.Array:
return value
@function(
gpu=True, input_schema=input_schema, output_schema=pa.list_(pa.int64())
)
def record_batch_embedding(value: pa.RecordBatch) -> pa.Array:
return value.column(0)
@function
def cpu_scalar(value: str) -> int:
return len(value)
class_embedding = function(
_GpuArrayFunction,
gpu=True,
input_schema=input_schema,
output_schema=pa.list_(pa.int64()),
)
assert not caught
assert array_embedding.registration_request.batch_size is None
assert record_batch_embedding.registration_request.batch_size is None
assert cpu_scalar.registration_request.batch_size is None
assert class_scalar.registration_request.batch_size is None
assert class_embedding.registration_request.batch_size is None
def test_unknown_runtime_discards_payload_before_known_field_validation():
for payload in [
{