mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-13 14:16:54 +00:00
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:
@@ -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 [
|
||||
{
|
||||
|
||||
Reference in New Issue
Block a user