feat: expose Function call concurrency limits (#4956)

## Summary

Adds an optional `max_concurrency` setting to Function registration and
explicit refresh:

- A registered Function version can define the default maximum number of
user-code calls in flight for one logical execution.
- A refresh can override that default for that execution only.
- The value must be a positive `u32` and is independent from
`resource_budget`; it is not an account-wide quota.
- Omitting the option preserves the previous request payload. Local
embedded refreshes reject this Enterprise-only option.

Rust, Python sync/async, PyO3, and the existing Node refresh surface
forward the option. Node has no Function registration surface today, so
no registration API is added there. The generated TypeDoc pages are
included for the new public option.

## Validation

The native-tested source covers Function version decoding, registration
validation and serialization, the Flight SQL registration ticket and
returned version, remote refresh forwarding, local rejection, and Node's
numeric boundary before conversion. Exact stage results are Rust `1+3`
plus four remote cases at `1` each; Python `14/1/49/24`; Node `57`
remote and `1` local rejection. Ruff, standard repository Pyright, and
Biome passed.

The tested Python extension was
`21aa256a970f4a16ea8253d68401d0fc573d946412e8752155dfcad63f30654d`; the
tested Node addon and generated native declaration were
`f6089f861981f96f8ce0f6dbef43de0907b9eeb4b63e84cdd5da14e00d6cbd30` and
`4c4ceb205824075d706897c726c215dcf839880c757f0d0582d24cec5f938416`. The
final 24-file diff preserves the 22 tested
product/test/fixture/source-doc blobs and adds only two generated
TypeDoc pages.
This commit is contained in:
Xuanwo
2026-11-11 14:32:26 +08:00
committed by GitHub
parent 8029428c71
commit cf8e6d0fd4
24 changed files with 376 additions and 18 deletions
+2
View File
@@ -631,6 +631,7 @@ class Table:
*,
where: Optional[str] = None,
batch_size: Optional[int] = None,
max_concurrency: Optional[int] = None,
publication_interval_rows: Optional[int] = None,
resource_budget: Optional[dict[str, int]] = None,
gpu_class: Optional[str] = None,
@@ -650,6 +651,7 @@ class Table:
*,
where: Optional[str] = None,
batch_size: Optional[int] = None,
max_concurrency: Optional[int] = None,
publication_interval_rows: Optional[int] = None,
resource_budget: Optional[dict[str, int]] = None,
gpu_class: Optional[str] = None,
+36 -3
View File
@@ -97,6 +97,16 @@ def _validate_batch_size(value: Any) -> Optional[int]:
return value
def _validate_max_concurrency(value: Any) -> Optional[int]:
if value is None:
return None
if isinstance(value, bool) or not isinstance(value, int):
raise ValueError("max_concurrency must be an integer")
if not 1 <= value <= 2**32 - 1:
raise ValueError("max_concurrency must be between 1 and 4294967295")
return value
def _validate_row_timeout(value: Any) -> Optional[float]:
if value is None:
return None
@@ -529,9 +539,10 @@ class FunctionVersion(_RemoteValue):
``gpu_class`` is the operator-defined GPU pool fixed into this Function
version, or ``None`` when registration used the deployment default.
``batch_size``, ``row_timeout_seconds``, and ``result_stability`` retain
the execution controls reported by the catalog. They are ``None`` for
versions registered before the corresponding declaration.
``batch_size``, ``max_concurrency``, ``row_timeout_seconds``, and
``result_stability`` retain the execution controls reported by the
catalog. They are ``None`` for versions registered before the
corresponding declaration.
"""
name: str
@@ -548,6 +559,7 @@ class FunctionVersion(_RemoteValue):
disabled: bool
gpu_class: Optional[str] = None
batch_size: Optional[int] = None
max_concurrency: Optional[int] = None
row_timeout_seconds: Optional[float] = None
result_stability: Optional[str] = None
@@ -556,6 +568,11 @@ class FunctionVersion(_RemoteValue):
def _validate_batch_size(cls, value):
return _validate_batch_size(value)
@field_validator("max_concurrency", mode="before")
@classmethod
def _validate_max_concurrency(cls, value):
return _validate_max_concurrency(value)
@field_validator("row_timeout_seconds", mode="before")
@classmethod
def _validate_row_timeout(cls, value):
@@ -691,6 +708,7 @@ class FunctionRegistrationRequest(_RemoteValue):
signature: FunctionSignature
runtime: PythonRuntimeSpec
batch_size: Optional[int] = None
max_concurrency: Optional[int] = None
row_timeout_seconds: Optional[float] = None
retry_policy: Optional[FunctionRetryPolicy] = None
secret_bindings: tuple[SecretBinding, ...] = ()
@@ -701,6 +719,11 @@ class FunctionRegistrationRequest(_RemoteValue):
def _validate_batch_size(cls, value):
return _validate_batch_size(value)
@field_validator("max_concurrency", mode="before")
@classmethod
def _validate_max_concurrency(cls, value):
return _validate_max_concurrency(value)
@field_validator("row_timeout_seconds", mode="before")
@classmethod
def _validate_row_timeout(cls, value):
@@ -2163,6 +2186,7 @@ class FunctionDefinition:
env: Mapping[str, str],
python_version: Optional[str],
batch_size: Optional[int] = None,
max_concurrency: Optional[int] = None,
row_timeout_seconds: Optional[float] = None,
retry_policy: Optional[FunctionRetryPolicy] = None,
gpu: bool = False,
@@ -2310,6 +2334,7 @@ class FunctionDefinition:
signature=signature,
runtime=runtime,
batch_size=batch_size,
max_concurrency=max_concurrency,
row_timeout_seconds=row_timeout_seconds,
retry_policy=retry_policy,
models=models,
@@ -2408,6 +2433,7 @@ def function(
env: Optional[Mapping[str, str]] = None,
python_version: Optional[str] = None,
batch_size: Optional[int] = None,
max_concurrency: Optional[int] = None,
row_timeout_seconds: Optional[float] = None,
retry_policy: Optional[FunctionRetryPolicy] = None,
gpu: bool = False,
@@ -2431,6 +2457,7 @@ def function(
env: Optional[Mapping[str, str]] = None,
python_version: Optional[str] = None,
batch_size: Optional[int] = None,
max_concurrency: Optional[int] = None,
row_timeout_seconds: Optional[float] = None,
retry_policy: Optional[FunctionRetryPolicy] = None,
gpu: bool = False,
@@ -2564,6 +2591,11 @@ def function(
Maximum number of input rows supplied in one execution batch. Must be
between 1 and 4294967295. Remote execution may use smaller batches.
When omitted, the service chooses the batch size adaptively.
max_concurrency : int, optional
Maximum user-code calls in flight for one Function within one logical
execution. Must be between 1 and 4294967295. A refresh can override
this default for that execution only. Omitting it preserves the
existing behavior. This is not an account-wide quota.
row_timeout_seconds : float, optional
Positive, finite deadline in seconds for each Python call: one row for
scalar Functions and one batch for vectorized Functions. This does not
@@ -2700,6 +2732,7 @@ def function(
env={} if env is None else env,
python_version=python_version,
batch_size=batch_size,
max_concurrency=max_concurrency,
row_timeout_seconds=row_timeout_seconds,
retry_policy=retry_policy,
gpu=gpu,
+4
View File
@@ -1145,6 +1145,7 @@ class RemoteTable(Table):
*,
where: Optional[str] = None,
batch_size: Optional[int] = None,
max_concurrency: Optional[int] = None,
publication_interval_rows: Optional[int] = None,
resource_budget: Optional[Dict[str, int]] = None,
gpu_class: Optional[str] = None,
@@ -1163,6 +1164,7 @@ class RemoteTable(Table):
column,
where=where,
batch_size=batch_size,
max_concurrency=max_concurrency,
publication_interval_rows=publication_interval_rows,
resource_budget=resource_budget,
gpu_class=gpu_class,
@@ -1184,6 +1186,7 @@ class RemoteTable(Table):
*,
where: Optional[str] = None,
batch_size: Optional[int] = None,
max_concurrency: Optional[int] = None,
publication_interval_rows: Optional[int] = None,
resource_budget: Optional[Dict[str, int]] = None,
gpu_class: Optional[str] = None,
@@ -1203,6 +1206,7 @@ class RemoteTable(Table):
column,
where=where,
batch_size=batch_size,
max_concurrency=max_concurrency,
publication_interval_rows=publication_interval_rows,
resource_budget=resource_budget,
gpu_class=gpu_class,
+34
View File
@@ -3193,6 +3193,7 @@ class Table(ABC):
*,
where: Optional[str] = None,
batch_size: Optional[int] = None,
max_concurrency: Optional[int] = None,
publication_interval_rows: Optional[int] = None,
resource_budget: Optional[Dict[str, int]] = None,
gpu_class: Optional[str] = None,
@@ -3233,6 +3234,12 @@ class Table(ABC):
Maximum rows per Function call on the remote server. Must be
greater than zero. None uses the server default. Has no effect on
local expression-backed columns.
max_concurrency: int, optional
Override the registered Function's maximum concurrent user-code
calls for this logical execution. Must be between 1 and
4294967295. None uses the registration default. This is not an
account-wide quota. LanceDB Enterprise only; local refreshes
reject this option.
publication_interval_rows: int, optional
Opt-in publication threshold for completed Function outputs. Must
be greater than zero. None keeps whole-fragment publication. The
@@ -3295,6 +3302,7 @@ class Table(ABC):
*,
where: Optional[str] = None,
batch_size: Optional[int] = None,
max_concurrency: Optional[int] = None,
publication_interval_rows: Optional[int] = None,
resource_budget: Optional[Dict[str, int]] = None,
gpu_class: Optional[str] = None,
@@ -3334,6 +3342,12 @@ class Table(ABC):
greater than zero. None uses the server default. Lower this for
slow Functions to stay within the executor's per-call deadline.
Has no effect on local expression-backed columns.
max_concurrency: int, optional
Override the registered Function's maximum concurrent user-code
calls for this logical execution. Must be between 1 and
4294967295. None uses the registration default. This is not an
account-wide quota. LanceDB Enterprise only; local refreshes
reject this option.
publication_interval_rows: int, optional
Opt-in publication threshold for completed Function outputs. Must
be greater than zero. None keeps whole-fragment publication. The
@@ -5493,6 +5507,7 @@ class LanceTable(Table):
*,
where: Optional[str] = None,
batch_size: Optional[int] = None,
max_concurrency: Optional[int] = None,
publication_interval_rows: Optional[int] = None,
resource_budget: Optional[Dict[str, int]] = None,
gpu_class: Optional[str] = None,
@@ -5514,6 +5529,7 @@ class LanceTable(Table):
column,
where=where,
batch_size=batch_size,
max_concurrency=max_concurrency,
publication_interval_rows=publication_interval_rows,
resource_budget=resource_budget,
gpu_class=gpu_class,
@@ -5535,6 +5551,7 @@ class LanceTable(Table):
*,
where: Optional[str] = None,
batch_size: Optional[int] = None,
max_concurrency: Optional[int] = None,
publication_interval_rows: Optional[int] = None,
resource_budget: Optional[Dict[str, int]] = None,
gpu_class: Optional[str] = None,
@@ -5558,6 +5575,7 @@ class LanceTable(Table):
column,
where=where,
batch_size=batch_size,
max_concurrency=max_concurrency,
publication_interval_rows=publication_interval_rows,
resource_budget=resource_budget,
gpu_class=gpu_class,
@@ -7940,6 +7958,7 @@ class AsyncTable:
*,
where: Optional[str] = None,
batch_size: Optional[int] = None,
max_concurrency: Optional[int] = None,
publication_interval_rows: Optional[int] = None,
resource_budget: Optional[Dict[str, int]] = None,
gpu_class: Optional[str] = None,
@@ -7980,6 +7999,12 @@ class AsyncTable:
Maximum rows per Function call on the remote server. Must be
greater than zero. None uses the server default. Has no effect on
local expression-backed columns.
max_concurrency: int, optional
Override the registered Function's maximum concurrent user-code
calls for this logical execution. Must be between 1 and
4294967295. None uses the registration default. This is not an
account-wide quota. LanceDB Enterprise only; local refreshes
reject this option.
publication_interval_rows: int, optional
Opt-in publication threshold for completed Function outputs. Must
be greater than zero. None keeps whole-fragment publication. The
@@ -8037,6 +8062,7 @@ class AsyncTable:
column,
where=where,
batch_size=batch_size,
max_concurrency=max_concurrency,
publication_interval_rows=publication_interval_rows,
resource_budget=resource_budget,
gpu_class=gpu_class,
@@ -8148,6 +8174,7 @@ class AsyncTable:
*,
where: Optional[str] = None,
batch_size: Optional[int] = None,
max_concurrency: Optional[int] = None,
publication_interval_rows: Optional[int] = None,
resource_budget: Optional[Dict[str, int]] = None,
gpu_class: Optional[str] = None,
@@ -8187,6 +8214,12 @@ class AsyncTable:
greater than zero. None uses the server default. Lower this for
slow Functions to stay within the executor's per-call deadline.
Has no effect on local expression-backed columns.
max_concurrency: int, optional
Override the registered Function's maximum concurrent user-code
calls for this logical execution. Must be between 1 and
4294967295. None uses the registration default. This is not an
account-wide quota. LanceDB Enterprise only; local refreshes
reject this option.
publication_interval_rows: int, optional
Opt-in publication threshold for completed Function outputs. Must
be greater than zero. None keeps whole-fragment publication. The
@@ -8261,6 +8294,7 @@ class AsyncTable:
column,
where=where,
batch_size=batch_size,
max_concurrency=max_concurrency,
publication_interval_rows=publication_interval_rows,
resource_budget=resource_budget,
gpu_class=gpu_class,
@@ -184,6 +184,7 @@ def test_function_version_preserves_execution_controls():
value = job_result("remote_function_job.json")
value.update(
batch_size=12,
max_concurrency=4,
row_timeout_seconds=8.0,
result_stability="input_and_context",
)
@@ -192,9 +193,14 @@ def test_function_version_preserves_execution_controls():
assert encoded == value
assert type(encoded["batch_size"]) is int
assert type(encoded["max_concurrency"]) is int
assert type(encoded["row_timeout_seconds"]) is float
assert type(encoded["result_stability"]) is str
for invalid in [True, 1.5, 0, -1, 2**32]:
with pytest.raises(ValueError, match="max_concurrency"):
FunctionVersion(**{**value, "max_concurrency": invalid})
def test_function_version_binds_named_columns_as_one_immutable_application():
version = FunctionVersion.from_json(
@@ -134,6 +134,7 @@ def test_function_execution_controls_match_shared_registration_golden():
env={"MODE": "test"},
python_version="3.12",
batch_size=2,
max_concurrency=4,
row_timeout_seconds=0.75,
)(normalize_score.__wrapped__)
assert (
@@ -193,6 +194,11 @@ def test_function_retry_policy_rejects_invalid_options(options):
({"batch_size": 0}, "batch_size"),
({"batch_size": 2**32}, "batch_size"),
({"batch_size": True}, "batch_size"),
({"max_concurrency": 0}, "max_concurrency"),
({"max_concurrency": -1}, "max_concurrency"),
({"max_concurrency": 2**32}, "max_concurrency"),
({"max_concurrency": True}, "max_concurrency"),
({"max_concurrency": 1.5}, "max_concurrency"),
({"row_timeout_seconds": 0}, "row_timeout_seconds"),
({"row_timeout_seconds": -0.25}, "row_timeout_seconds"),
({"row_timeout_seconds": float("inf")}, "row_timeout_seconds"),
+48 -1
View File
@@ -10,6 +10,7 @@ import io
import inspect
import json
import multiprocessing as mp
from pathlib import Path
import pickle
import re
import socket
@@ -36,7 +37,7 @@ from lancedb.graph import (
PropertyGraph,
PropertyGraphDescription,
)
from lancedb.functions import FunctionApplication
from lancedb.functions import FunctionApplication, function
from lancedb.query import (
AsyncQuery,
ColumnOrdering,
@@ -1864,6 +1865,7 @@ def test_remote_create_index_async_returns_job():
{"batch_size": None},
{"batch_size": 1},
{"batch_size": 8},
{"max_concurrency": 3},
{"publication_interval_rows": 60},
{"resource_budget": {"cpu_millis": 2000, "memory_bytes": 4 * 1024**3}},
{
@@ -1927,6 +1929,7 @@ async def test_remote_refresh_async_returns_typed_terminal_result(
expected["error_policy"] = refresh_kwargs["error_policy"]
for key in (
"where",
"max_concurrency",
"publication_interval_rows",
"max_failed_fraction",
"max_failed_rows",
@@ -2107,6 +2110,11 @@ def test_remote_refresh_async_forwards_progress_from_submitted_handle():
({"init_timeout_seconds": -5}, "init_timeout_seconds"),
({"publication_interval_rows": 0}, "publication_interval_rows"),
({"publication_interval_rows": -1}, "publication_interval_rows"),
({"max_concurrency": True}, "max_concurrency"),
({"max_concurrency": 1.5}, "max_concurrency"),
({"max_concurrency": 0}, "max_concurrency"),
({"max_concurrency": -1}, "max_concurrency"),
({"max_concurrency": 2**32}, "max_concurrency"),
({"num_workers": 0}, "num_workers"),
(
{"resource_budget": {"cpu_millis": 1000}},
@@ -6093,6 +6101,45 @@ def _mock_remote_table_with_sql(flight_server):
yield lancedb.connect(**kwargs).create_table("test", [{"id": 1}])
def test_remote_function_registration_transports_max_concurrency():
flight = pytest.importorskip("pyarrow.flight")
fixture = (
Path(__file__).parents[3]
/ "rust"
/ "lancedb"
/ "tests"
/ "fixtures"
/ "first_class_functions"
/ "v1"
/ "remote_function_version.canonical.json"
)
definition = json.loads(fixture.read_text())
definition["max_concurrency"] = 4
@function(name="normalize_score", max_concurrency=4)
def registered(value: float) -> float:
return value
class RegistrationServer(flight.FlightServerBase):
def __init__(self):
super().__init__("grpc://127.0.0.1:0")
self.tickets = []
def do_get(self, context, ticket):
self.tickets.append(ticket.ticket)
return flight.RecordBatchStream(
pa.table({"definition": [json.dumps(definition)]})
)
server = RegistrationServer()
with _mock_remote_sql_connection(server) as kwargs:
version = lancedb.connect(**kwargs).create_function(registered)
assert len(server.tickets) == 1
assert b'"max_concurrency":4' in server.tickets[0]
assert version.max_concurrency == 4
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize(
+1
View File
@@ -5995,6 +5995,7 @@ def test_refresh_column_rejects_invalid_batch_size(tmp_path, method, batch_size)
({"max_failed_rows": 10}, "max_failed_rows"),
({"row_timeout_seconds": 30}, "row_timeout_seconds"),
({"init_timeout_seconds": 60}, "init_timeout_seconds"),
({"max_concurrency": 4}, "max_concurrency"),
({"publication_interval_rows": 60}, "publication_interval_rows"),
(
{"resource_budget": {"cpu_millis": 1000, "memory_bytes": 1024}},
+34 -2
View File
@@ -48,6 +48,7 @@ mod scannable;
struct RefreshColumnArgs {
where_clause: Option<String>,
batch_size: Option<i64>,
max_concurrency: Option<u32>,
publication_interval_rows: Option<i64>,
resource_budget: Option<FunctionWorkerResourceBudget>,
gpu_class: Option<String>,
@@ -73,6 +74,31 @@ fn positive_count<T: TryFrom<i64>>(value: Option<i64>, name: &str) -> PyResult<O
.transpose()
}
fn positive_u32_integer(value: Option<&Bound<'_, PyAny>>, name: &str) -> PyResult<Option<u32>> {
let Some(value) = value else {
return Ok(None);
};
if value.is_none() || value.is_instance_of::<PyBool>() {
return Err(PyValueError::new_err(format!(
"{name} must be an integer between 1 and 4294967295"
)));
}
let parsed = value.extract::<i64>().map_err(|_| {
PyValueError::new_err(format!(
"{name} must be an integer between 1 and 4294967295"
))
})?;
u32::try_from(parsed)
.ok()
.filter(|parsed| *parsed > 0)
.map(Some)
.ok_or_else(|| {
PyValueError::new_err(format!(
"{name} must be an integer between 1 and 4294967295"
))
})
}
fn resource_budget_integer(
value: &Bound<'_, PyAny>,
name: &str,
@@ -136,6 +162,7 @@ fn refresh_column_options(args: RefreshColumnArgs) -> PyResult<RefreshColumnOpti
let RefreshColumnArgs {
where_clause,
batch_size,
max_concurrency,
publication_interval_rows,
resource_budget,
gpu_class,
@@ -179,6 +206,7 @@ fn refresh_column_options(args: RefreshColumnArgs) -> PyResult<RefreshColumnOpti
.map(usize::try_from)
.transpose()
.map_err(|_| PyValueError::new_err("batch_size must be greater than zero"))?;
options.max_concurrency = max_concurrency;
options.publication_interval_rows =
positive_count(publication_interval_rows, "publication_interval_rows")?;
options.resource_budget = resource_budget;
@@ -2282,13 +2310,14 @@ impl Table {
})
}
#[pyo3(signature = (column, *, r#where=None, batch_size=None, publication_interval_rows=None, resource_budget=None, gpu_class=None, gpu_count=None, retry_policy=None, error_policy=None, max_failed_fraction=None, max_failed_rows=None, row_timeout_seconds=None, init_timeout_seconds=None, num_workers=None, max_executor_processes_per_worker=None))]
#[pyo3(signature = (column, *, r#where=None, batch_size=None, max_concurrency=None, publication_interval_rows=None, resource_budget=None, gpu_class=None, gpu_count=None, retry_policy=None, error_policy=None, max_failed_fraction=None, max_failed_rows=None, row_timeout_seconds=None, init_timeout_seconds=None, num_workers=None, max_executor_processes_per_worker=None))]
#[allow(clippy::too_many_arguments)]
pub fn refresh_column<'py>(
self_: PyRef<'py, Self>,
column: String,
r#where: Option<String>,
batch_size: Option<i64>,
max_concurrency: Option<Bound<'py, PyAny>>,
publication_interval_rows: Option<i64>,
resource_budget: Option<Bound<'py, PyDict>>,
gpu_class: Option<String>,
@@ -2305,6 +2334,7 @@ impl Table {
let options = refresh_column_options(RefreshColumnArgs {
where_clause: r#where,
batch_size,
max_concurrency: positive_u32_integer(max_concurrency.as_ref(), "max_concurrency")?,
publication_interval_rows,
resource_budget: parse_resource_budget(resource_budget.as_ref())?,
gpu_class,
@@ -2328,13 +2358,14 @@ impl Table {
})
}
#[pyo3(signature = (column, *, r#where=None, batch_size=None, publication_interval_rows=None, resource_budget=None, gpu_class=None, gpu_count=None, retry_policy=None, error_policy=None, max_failed_fraction=None, max_failed_rows=None, row_timeout_seconds=None, init_timeout_seconds=None, num_workers=None, max_executor_processes_per_worker=None))]
#[pyo3(signature = (column, *, r#where=None, batch_size=None, max_concurrency=None, publication_interval_rows=None, resource_budget=None, gpu_class=None, gpu_count=None, retry_policy=None, error_policy=None, max_failed_fraction=None, max_failed_rows=None, row_timeout_seconds=None, init_timeout_seconds=None, num_workers=None, max_executor_processes_per_worker=None))]
#[allow(clippy::too_many_arguments)]
pub fn refresh_column_async<'py>(
self_: PyRef<'py, Self>,
column: String,
r#where: Option<String>,
batch_size: Option<i64>,
max_concurrency: Option<Bound<'py, PyAny>>,
publication_interval_rows: Option<i64>,
resource_budget: Option<Bound<'py, PyDict>>,
gpu_class: Option<String>,
@@ -2351,6 +2382,7 @@ impl Table {
let options = refresh_column_options(RefreshColumnArgs {
where_clause: r#where,
batch_size,
max_concurrency: positive_u32_integer(max_concurrency.as_ref(), "max_concurrency")?,
publication_interval_rows,
resource_budget: parse_resource_budget(resource_budget.as_ref())?,
gpu_class,