mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-13 22:26:53 +00:00
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:
@@ -1142,8 +1142,11 @@ LanceDB Enterprise it is the server's backfill job.
|
||||
`REFRESH COLUMN ... WHERE ... WITH (...)` form: `where` narrows the rows,
|
||||
`errorPolicy`, `maxFailedFraction`, `maxFailedRows`, `rowTimeoutSeconds`
|
||||
and `initTimeoutSeconds` govern failing rows. `batchSize`, `gpuClass`,
|
||||
`resourceBudget` and `maxExecutorProcessesPerWorker` size registered
|
||||
Function refreshes. A `resourceBudget` requires `cpuMillis` and
|
||||
`maxConcurrency`, `resourceBudget` and
|
||||
`maxExecutorProcessesPerWorker` size registered Function refreshes.
|
||||
`maxConcurrency` overrides the registered default for this logical
|
||||
execution only and is a positive integer; it is not an account-wide
|
||||
quota. A `resourceBudget` requires `cpuMillis` and
|
||||
`memoryBytes`; its `gpuCount` defaults to zero for a CPU-only Function.
|
||||
The budget applies only to this explicit refresh and does not configure
|
||||
a later automatic refresh.
|
||||
@@ -1167,6 +1170,7 @@ const job = await table.refreshColumnAsync("embedding", {
|
||||
cpuMillis: 8000,
|
||||
memoryBytes: 16 * 1024 ** 3,
|
||||
},
|
||||
maxConcurrency: 4,
|
||||
errorPolicy: "skip",
|
||||
maxFailedFraction: 0.01,
|
||||
rowTimeoutSeconds: 30,
|
||||
|
||||
@@ -64,6 +64,17 @@ Deadline for Function initialization, in seconds.
|
||||
|
||||
***
|
||||
|
||||
### maxConcurrency?
|
||||
|
||||
```ts
|
||||
optional maxConcurrency: number;
|
||||
```
|
||||
|
||||
Override the registered Function's maximum concurrent user-code calls
|
||||
for this logical execution.
|
||||
|
||||
***
|
||||
|
||||
### maxExecutorProcessesPerWorker?
|
||||
|
||||
```ts
|
||||
|
||||
@@ -60,6 +60,7 @@ from lancedb import function
|
||||
@function(
|
||||
pip=["numpy==2.2.0"],
|
||||
batch_size=32,
|
||||
max_concurrency=4,
|
||||
row_timeout_seconds=90,
|
||||
)
|
||||
def score_batch(value: pa.Array) -> pa.Array:
|
||||
@@ -72,12 +73,15 @@ once a function also does something with real fixed cost per call -- loading
|
||||
a model, opening a connection, making a network round trip.
|
||||
|
||||
`batch_size` caps the number of input rows in each execution batch; actual
|
||||
batches may be smaller. `row_timeout_seconds` bounds each Python call: one
|
||||
row for scalar Functions or the entire batch for vectorized Functions. It
|
||||
does not include the instance's `__init__`. Both settings are fixed into the
|
||||
registered version and apply to direct calls, computed columns, and
|
||||
materialized-view refreshes alike. Omit them to use adaptive batching and the
|
||||
service's default per-call deadline.
|
||||
batches may be smaller. `max_concurrency` limits user-code calls in flight for
|
||||
one Function within one logical execution; it does not limit calls across
|
||||
different jobs and is not an account-wide quota. `row_timeout_seconds` bounds
|
||||
each Python call: one row for scalar Functions or the entire batch for
|
||||
vectorized Functions. It does not include the instance's `__init__`. These
|
||||
settings are fixed into the registered version and apply to direct calls,
|
||||
computed columns, and materialized-view refreshes alike. A refresh can
|
||||
override `max_concurrency` for that execution. Omit it to preserve the
|
||||
existing behavior.
|
||||
|
||||
## Stateful Functions: load once, call many times {#stateful}
|
||||
|
||||
|
||||
@@ -225,6 +225,39 @@ Function declaration can separately set `auto_refresh_resource_budget` on
|
||||
refreshes. SQL-expression columns reject the automatic resource-budget option
|
||||
because they do not use Function workers.
|
||||
|
||||
To override the registered Function's call concurrency for this logical
|
||||
execution only:
|
||||
|
||||
<CodeGroup>
|
||||
```python Python icon="python"
|
||||
job = table.refresh_column_async("embedding", max_concurrency=4)
|
||||
job.wait()
|
||||
```
|
||||
|
||||
```rust Rust icon="rust"
|
||||
use lancedb::table::RefreshColumnOptions;
|
||||
|
||||
let _job = table
|
||||
.refresh_column_async_with_options(
|
||||
"embedding",
|
||||
RefreshColumnOptions::default().max_concurrency(4),
|
||||
)
|
||||
.await?;
|
||||
```
|
||||
|
||||
```typescript TypeScript icon="typescript"
|
||||
const job = await table.refreshColumnAsync("embedding", {
|
||||
maxConcurrency: 4,
|
||||
});
|
||||
await job.wait();
|
||||
```
|
||||
</CodeGroup>
|
||||
|
||||
`max_concurrency` caps user-code calls in flight for one Function within one
|
||||
logical execution. It is not an account-wide quota and does not resize or
|
||||
rewrite `resource_budget`. Omitting it uses the registration default; a
|
||||
refresh value overrides that default for this execution only.
|
||||
|
||||
`batch_size` caps rows in one Function call; LanceDB adapts actual call and
|
||||
work sizes. For Function columns, the refresh defaults to
|
||||
`error_policy = 'fail'`. Set `error_policy = 'skip'` to write null outputs for
|
||||
|
||||
@@ -682,6 +682,7 @@ describe("remote connection", () => {
|
||||
const job = await table.refreshColumnAsync("embedding", {
|
||||
where: "id % 2 = 0",
|
||||
batchSize: 16,
|
||||
maxConcurrency: 3,
|
||||
gpuClass: "a10g",
|
||||
gpuCount: 2,
|
||||
errorPolicy: "skip",
|
||||
@@ -699,6 +700,11 @@ describe("remote connection", () => {
|
||||
memoryBytes: 16 * 1024 ** 3,
|
||||
},
|
||||
});
|
||||
await expect(
|
||||
table.refreshColumnAsync("embedding", {
|
||||
maxConcurrency: true as unknown as number,
|
||||
}),
|
||||
).rejects.toThrow();
|
||||
// Invalid options reject before any request is sent.
|
||||
for (const [options, message] of [
|
||||
[{ maxFailedRows: -1 }, "maxFailedRows"],
|
||||
@@ -706,6 +712,12 @@ describe("remote connection", () => {
|
||||
[{ where: " " }, "where"],
|
||||
[{ maxFailedFraction: 2 }, "max_failed_fraction"],
|
||||
[{ rowTimeoutSeconds: 0 }, "row_timeout_seconds"],
|
||||
[{ maxConcurrency: 1.5 }, "maxConcurrency"],
|
||||
[{ maxConcurrency: 0 }, "maxConcurrency"],
|
||||
[{ maxConcurrency: -1 }, "maxConcurrency"],
|
||||
[{ maxConcurrency: 2 ** 32 }, "maxConcurrency"],
|
||||
[{ maxConcurrency: Number.NaN }, "maxConcurrency"],
|
||||
[{ maxConcurrency: Number.POSITIVE_INFINITY }, "maxConcurrency"],
|
||||
[{ numWorkers: 0 }, "num_workers"],
|
||||
[{ gpuCount: 0 }, "gpu_count"],
|
||||
[
|
||||
@@ -787,7 +799,8 @@ describe("remote connection", () => {
|
||||
expect(bodies).toEqual([
|
||||
JSON.parse(
|
||||
`{"column": "embedding", "where": "id % 2 = 0", "batch_size": 16,
|
||||
"gpu_class": "a10g", "gpu_count": 2, "error_policy": "skip",
|
||||
"max_concurrency": 3, "gpu_class": "a10g", "gpu_count": 2,
|
||||
"error_policy": "skip",
|
||||
"max_failed_fraction": 0.01, "max_failed_rows": 50,
|
||||
"row_timeout_seconds": 30, "init_timeout_seconds": 120,
|
||||
"max_executor_processes_per_worker": 2}`,
|
||||
|
||||
@@ -5686,6 +5686,9 @@ describe("computed columns", () => {
|
||||
await expect(
|
||||
table.refreshColumnAsync("x", { errorPolicy: "skip" }),
|
||||
).rejects.toThrow("error_policy requires a LanceDB Enterprise refresh");
|
||||
await expect(
|
||||
table.refreshColumnAsync("x", { maxConcurrency: 4 }),
|
||||
).rejects.toThrow("max_concurrency requires a LanceDB Enterprise refresh");
|
||||
await expect(
|
||||
table.refreshColumnAsync("x", {
|
||||
resourceBudget: { cpuMillis: 1_000, memoryBytes: 1_024 },
|
||||
|
||||
@@ -723,8 +723,11 @@ export abstract class Table {
|
||||
* `REFRESH COLUMN ... WHERE ... WITH (...)` form: `where` narrows the rows,
|
||||
* `errorPolicy`, `maxFailedFraction`, `maxFailedRows`, `rowTimeoutSeconds`
|
||||
* and `initTimeoutSeconds` govern failing rows. `batchSize`, `gpuClass`,
|
||||
* `resourceBudget` and `maxExecutorProcessesPerWorker` size registered
|
||||
* Function refreshes. A `resourceBudget` requires `cpuMillis` and
|
||||
* `maxConcurrency`, `resourceBudget` and
|
||||
* `maxExecutorProcessesPerWorker` size registered Function refreshes.
|
||||
* `maxConcurrency` overrides the registered default for this logical
|
||||
* execution only and is a positive integer; it is not an account-wide
|
||||
* quota. A `resourceBudget` requires `cpuMillis` and
|
||||
* `memoryBytes`; its `gpuCount` defaults to zero for a CPU-only Function.
|
||||
* The budget applies only to this explicit refresh and does not configure
|
||||
* a later automatic refresh.
|
||||
@@ -742,6 +745,7 @@ export abstract class Table {
|
||||
* cpuMillis: 8000,
|
||||
* memoryBytes: 16 * 1024 ** 3,
|
||||
* },
|
||||
* maxConcurrency: 4,
|
||||
* errorPolicy: "skip",
|
||||
* maxFailedFraction: 0.01,
|
||||
* rowTimeoutSeconds: 30,
|
||||
|
||||
@@ -1621,6 +1621,22 @@ fn resource_budget_gpu_count(value: Option<f64>) -> napi::Result<u32> {
|
||||
Ok(value as u32)
|
||||
}
|
||||
|
||||
fn positive_u32(value: Option<f64>, name: &str) -> napi::Result<Option<u32>> {
|
||||
value
|
||||
.map(|value| {
|
||||
if !value.is_finite()
|
||||
|| value.fract() != 0.0
|
||||
|| !(1.0..=u32::MAX as f64).contains(&value)
|
||||
{
|
||||
return Err(napi::Error::from_reason(format!(
|
||||
"{name} must be an integer between 1 and 4294967295"
|
||||
)));
|
||||
}
|
||||
Ok(value as u32)
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
|
||||
/// Options for `Table.refreshColumn` and `Table.refreshColumnAsync`. Every
|
||||
/// option is LanceDB Enterprise only; a LanceDB Embedded refresh ignores
|
||||
/// `batchSize`, `gpuClass` and `gpuCount` and rejects the rest, including
|
||||
@@ -1633,6 +1649,9 @@ pub struct RefreshColumnOptions {
|
||||
pub where_clause: Option<String>,
|
||||
/// Maximum rows per Function call; must be greater than zero.
|
||||
pub batch_size: Option<u32>,
|
||||
/// Override the registered Function's maximum concurrent user-code calls
|
||||
/// for this logical execution.
|
||||
pub max_concurrency: Option<f64>,
|
||||
/// GPU class for Function workers.
|
||||
pub gpu_class: Option<String>,
|
||||
/// Complete aggregate CPU and memory budget for a registered Function
|
||||
@@ -1666,6 +1685,7 @@ impl TryFrom<RefreshColumnOptions> for lancedb::table::RefreshColumnOptions {
|
||||
let mut options = Self::default();
|
||||
options.where_clause = value.where_clause;
|
||||
options.batch_size = value.batch_size.map(|size| size as usize);
|
||||
options.max_concurrency = positive_u32(value.max_concurrency, "maxConcurrency")?;
|
||||
options.gpu_class = value.gpu_class;
|
||||
options.resource_budget = value
|
||||
.resource_budget
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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"),
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
@@ -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,
|
||||
|
||||
@@ -611,6 +611,13 @@ pub struct FunctionVersion {
|
||||
deserialize_with = "deserialize_batch_size"
|
||||
)]
|
||||
batch_size: Option<u32>,
|
||||
/// Default maximum concurrent user-code calls per logical execution.
|
||||
#[serde(
|
||||
default,
|
||||
skip_serializing_if = "Option::is_none",
|
||||
deserialize_with = "deserialize_max_concurrency"
|
||||
)]
|
||||
max_concurrency: Option<u32>,
|
||||
/// Per-call deadline in seconds.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
row_timeout_seconds: Option<RowTimeoutSeconds>,
|
||||
@@ -681,6 +688,11 @@ impl FunctionVersion {
|
||||
self.batch_size
|
||||
}
|
||||
|
||||
/// Default maximum concurrent user-code calls, when declared.
|
||||
pub fn max_concurrency(&self) -> Option<u32> {
|
||||
self.max_concurrency
|
||||
}
|
||||
|
||||
/// Per-call deadline in seconds, when declared.
|
||||
pub fn row_timeout_seconds(&self) -> Option<f64> {
|
||||
self.row_timeout_seconds.map(RowTimeoutSeconds::get)
|
||||
@@ -769,6 +781,18 @@ where
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
fn deserialize_max_concurrency<'de, D>(
|
||||
deserializer: D,
|
||||
) -> std::result::Result<Option<u32>, D::Error>
|
||||
where
|
||||
D: Deserializer<'de>,
|
||||
{
|
||||
let value = Option::<u32>::deserialize(deserializer)
|
||||
.map_err(|error| de::Error::custom(format!("max_concurrency: {error}")))?;
|
||||
validate_max_concurrency_value(value).map_err(de::Error::custom)?;
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
fn validate_batch_size_value(value: Option<u32>) -> std::result::Result<(), &'static str> {
|
||||
if value == Some(0) {
|
||||
return Err("batch_size must be between 1 and u32::MAX");
|
||||
@@ -776,6 +800,13 @@ fn validate_batch_size_value(value: Option<u32>) -> std::result::Result<(), &'st
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_max_concurrency_value(value: Option<u32>) -> std::result::Result<(), &'static str> {
|
||||
if value == Some(0) {
|
||||
return Err("max_concurrency must be between 1 and u32::MAX");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_row_timeout_value(value: Option<f64>) -> std::result::Result<(), &'static str> {
|
||||
if value.is_some_and(|seconds| !seconds.is_finite() || seconds <= 0.0) {
|
||||
return Err("row_timeout_seconds must be positive and finite");
|
||||
@@ -804,6 +835,15 @@ pub struct FunctionRegistrationRequest {
|
||||
deserialize_with = "deserialize_batch_size"
|
||||
)]
|
||||
pub batch_size: Option<u32>,
|
||||
/// Default maximum concurrent user-code calls for each logical execution
|
||||
/// of this Function version. Omission leaves the current behavior
|
||||
/// unchanged.
|
||||
#[serde(
|
||||
default,
|
||||
skip_serializing_if = "Option::is_none",
|
||||
deserialize_with = "deserialize_max_concurrency"
|
||||
)]
|
||||
pub max_concurrency: Option<u32>,
|
||||
/// Deadline for each Python call, in seconds: one row for scalar Functions
|
||||
/// and one batch for vectorized Functions. It does not govern Function
|
||||
/// initialization.
|
||||
@@ -841,6 +881,7 @@ impl FunctionRegistrationRequest {
|
||||
message: message.to_string(),
|
||||
};
|
||||
validate_batch_size_value(self.batch_size).map_err(invalid_input)?;
|
||||
validate_max_concurrency_value(self.max_concurrency).map_err(invalid_input)?;
|
||||
validate_row_timeout_value(self.row_timeout_seconds).map_err(invalid_input)?;
|
||||
if let Some(policy) = &self.retry_policy {
|
||||
policy.validate()?;
|
||||
@@ -1318,6 +1359,15 @@ mod tests {
|
||||
assert_eq!(version.gpu_class(), Some("a10g"));
|
||||
let encoded: Value = serde_json::from_str(&version.to_canonical_json().unwrap()).unwrap();
|
||||
assert_eq!(encoded["gpu_class"], "a10g");
|
||||
|
||||
document["max_concurrency"] = serde_json::json!(4);
|
||||
let version = FunctionVersion::from_json(&document.to_string()).unwrap();
|
||||
assert_eq!(version.max_concurrency(), Some(4));
|
||||
let encoded: Value = serde_json::from_str(&version.to_canonical_json().unwrap()).unwrap();
|
||||
assert_eq!(encoded["max_concurrency"], 4);
|
||||
|
||||
document["max_concurrency"] = serde_json::json!(0);
|
||||
assert!(FunctionVersion::from_json(&document.to_string()).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -2959,6 +2959,7 @@ async fn function_creation_and_description_use_sql_query_handles() {
|
||||
for (defer, with_models) in [(false, false), (false, true), (true, false), (true, true)] {
|
||||
let mut request_json: serde_json::Value = serde_json::from_str(REQUEST).unwrap();
|
||||
request_json["batch_size"] = serde_json::json!(2);
|
||||
request_json["max_concurrency"] = serde_json::json!(4);
|
||||
request_json["row_timeout_seconds"] = serde_json::json!(0.75);
|
||||
request_json["secret_bindings"] = serde_json::json!([{
|
||||
"kind": "env",
|
||||
@@ -2967,6 +2968,7 @@ async fn function_creation_and_description_use_sql_query_handles() {
|
||||
}]);
|
||||
let mut version_json: serde_json::Value = serde_json::from_str(VERSION).unwrap();
|
||||
version_json["batch_size"] = serde_json::json!(12);
|
||||
version_json["max_concurrency"] = serde_json::json!(4);
|
||||
version_json["row_timeout_seconds"] = serde_json::json!(8.0);
|
||||
version_json["result_stability"] = serde_json::json!("input_and_context");
|
||||
if with_models {
|
||||
@@ -3010,6 +3012,7 @@ async fn function_creation_and_description_use_sql_query_handles() {
|
||||
assert_eq!(created.name(), "embed");
|
||||
assert_eq!(created.version(), "1");
|
||||
assert_eq!(created.batch_size(), Some(12));
|
||||
assert_eq!(created.max_concurrency(), Some(4));
|
||||
assert_eq!(created.row_timeout_seconds(), Some(8.0));
|
||||
assert_eq!(created.result_stability(), Some("input_and_context"));
|
||||
// Assert the actual SQL body, after the same native JSON decode used
|
||||
@@ -3029,6 +3032,7 @@ async fn function_creation_and_description_use_sql_query_handles() {
|
||||
serde_json::from_str(&created.to_canonical_json().unwrap()).unwrap();
|
||||
assert_eq!(created_json["image"].get("models"), expected_models);
|
||||
assert_eq!(created_json["batch_size"], 12);
|
||||
assert_eq!(created_json["max_concurrency"], 4);
|
||||
assert_eq!(created_json["row_timeout_seconds"], 8.0);
|
||||
assert_eq!(created_json["result_stability"], "input_and_context");
|
||||
assert_eq!(
|
||||
@@ -3041,12 +3045,14 @@ async fn function_creation_and_description_use_sql_query_handles() {
|
||||
.unwrap();
|
||||
assert_eq!(described.name(), "embed");
|
||||
assert_eq!(described.batch_size(), Some(12));
|
||||
assert_eq!(described.max_concurrency(), Some(4));
|
||||
assert_eq!(described.row_timeout_seconds(), Some(8.0));
|
||||
assert_eq!(described.result_stability(), Some("input_and_context"));
|
||||
let described_json: serde_json::Value =
|
||||
serde_json::from_str(&described.to_canonical_json().unwrap()).unwrap();
|
||||
assert_eq!(described_json["image"].get("models"), expected_models);
|
||||
assert_eq!(described_json["batch_size"], 12);
|
||||
assert_eq!(described_json["max_concurrency"], 4);
|
||||
assert_eq!(described_json["row_timeout_seconds"], 8.0);
|
||||
assert_eq!(described_json["result_stability"], "input_and_context");
|
||||
let latest = db
|
||||
@@ -3062,6 +3068,7 @@ async fn function_creation_and_description_use_sql_query_handles() {
|
||||
.unwrap();
|
||||
assert_eq!(listed.items.len(), 1);
|
||||
assert_eq!(listed.items[0].batch_size(), Some(12));
|
||||
assert_eq!(listed.items[0].max_concurrency(), Some(4));
|
||||
assert_eq!(listed.items[0].row_timeout_seconds(), Some(8.0));
|
||||
assert_eq!(
|
||||
listed.items[0].result_stability(),
|
||||
@@ -3071,6 +3078,7 @@ async fn function_creation_and_description_use_sql_query_handles() {
|
||||
serde_json::from_str(&listed.items[0].to_canonical_json().unwrap()).unwrap();
|
||||
assert_eq!(listed_json["image"].get("models"), expected_models);
|
||||
assert_eq!(listed_json["batch_size"], 12);
|
||||
assert_eq!(listed_json["max_concurrency"], 4);
|
||||
assert_eq!(listed_json["row_timeout_seconds"], 8.0);
|
||||
assert_eq!(listed_json["result_stability"], "input_and_context");
|
||||
let seen = statements.lock().unwrap();
|
||||
@@ -3106,6 +3114,13 @@ async fn function_creation_rejects_invalid_execution_controls_before_sql() {
|
||||
request.batch_size = Some(0);
|
||||
let error = db.create_function_async(request, &[]).await.unwrap_err();
|
||||
assert!(matches!(error, Error::InvalidInput { message } if message.contains("batch_size")));
|
||||
|
||||
let mut request = crate::function::FunctionRegistrationRequest::from_json(REQUEST).unwrap();
|
||||
request.max_concurrency = Some(0);
|
||||
let error = db.create_function_async(request, &[]).await.unwrap_err();
|
||||
assert!(
|
||||
matches!(error, Error::InvalidInput { message } if message.contains("max_concurrency"))
|
||||
);
|
||||
assert!(statements.lock().unwrap().is_empty());
|
||||
shutdown.send(()).unwrap();
|
||||
}
|
||||
|
||||
@@ -11131,6 +11131,7 @@ mod tests {
|
||||
serde_json::json!({
|
||||
"column": "quality",
|
||||
"where": "id % 2 = 0",
|
||||
"max_concurrency": 3,
|
||||
"publication_interval_rows": 60,
|
||||
"error_policy": "skip",
|
||||
"max_failed_fraction": 0.01,
|
||||
@@ -11148,6 +11149,7 @@ mod tests {
|
||||
});
|
||||
let options = RefreshColumnOptions::default()
|
||||
.where_clause("id % 2 = 0")
|
||||
.max_concurrency(3)
|
||||
.publication_interval_rows(60)
|
||||
.error_policy(FunctionErrorPolicy::Skip)
|
||||
.max_failed_fraction(0.01)
|
||||
@@ -11192,6 +11194,10 @@ mod tests {
|
||||
RefreshColumnOptions::default().num_workers(0),
|
||||
"num_workers",
|
||||
),
|
||||
(
|
||||
RefreshColumnOptions::default().max_concurrency(0),
|
||||
"max_concurrency",
|
||||
),
|
||||
(
|
||||
RefreshColumnOptions::default().max_executor_processes_per_worker(0),
|
||||
"max_executor_processes_per_worker",
|
||||
|
||||
@@ -149,6 +149,12 @@ pub struct RefreshColumnOptions {
|
||||
/// effect on expression-backed columns in LanceDB Embedded.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub batch_size: Option<usize>,
|
||||
/// Override the registered Function's maximum concurrent user-code calls
|
||||
/// for this logical execution. Must be greater than zero. `None` uses the
|
||||
/// registration default, or disabled when the Function declares none.
|
||||
/// This is not an account quota and does not change `resource_budget`.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub max_concurrency: Option<u32>,
|
||||
/// Opt-in row-count threshold for publishing completed Function outputs.
|
||||
///
|
||||
/// Must be greater than zero. `None` keeps whole-fragment publication.
|
||||
@@ -235,6 +241,13 @@ impl RefreshColumnOptions {
|
||||
self
|
||||
}
|
||||
|
||||
/// Override the registered Function's maximum concurrent user-code calls
|
||||
/// for this logical execution.
|
||||
pub fn max_concurrency(mut self, max_concurrency: u32) -> Self {
|
||||
self.max_concurrency = Some(max_concurrency);
|
||||
self
|
||||
}
|
||||
|
||||
/// Use `rows` as the opt-in publication threshold for completed Function
|
||||
/// outputs. The actual publication boundary can be larger than this value.
|
||||
pub fn publication_interval_rows(mut self, rows: u64) -> Self {
|
||||
@@ -278,6 +291,7 @@ impl RefreshColumnOptions {
|
||||
"max_executor_processes_per_worker",
|
||||
self.max_executor_processes_per_worker.is_some(),
|
||||
),
|
||||
("max_concurrency", self.max_concurrency.is_some()),
|
||||
(
|
||||
"publication_interval_rows",
|
||||
self.publication_interval_rows.is_some(),
|
||||
@@ -335,6 +349,11 @@ impl RefreshColumnOptions {
|
||||
message: "batch_size must be greater than zero".into(),
|
||||
});
|
||||
}
|
||||
if self.max_concurrency == Some(0) {
|
||||
return Err(Error::InvalidInput {
|
||||
message: "max_concurrency must be greater than zero".into(),
|
||||
});
|
||||
}
|
||||
if self.publication_interval_rows == Some(0) {
|
||||
return Err(Error::InvalidInput {
|
||||
message: "publication_interval_rows must be greater than zero".into(),
|
||||
|
||||
@@ -71,6 +71,7 @@ fn execution_controls_match_shared_canonical_golden() {
|
||||
|
||||
let value: Value = serde_json::from_str(&request.to_canonical_json().unwrap()).unwrap();
|
||||
assert_eq!(value["batch_size"], 2);
|
||||
assert_eq!(value["max_concurrency"], 4);
|
||||
assert_eq!(value["row_timeout_seconds"], 0.75);
|
||||
}
|
||||
|
||||
@@ -80,6 +81,10 @@ fn execution_controls_require_their_positive_ranges() {
|
||||
for (field, value) in [
|
||||
("batch_size", Value::from(0)),
|
||||
("batch_size", Value::from(4_294_967_296_u64)),
|
||||
("max_concurrency", Value::from(0)),
|
||||
("max_concurrency", Value::from(4_294_967_296_u64)),
|
||||
("max_concurrency", Value::from(-1)),
|
||||
("max_concurrency", Value::from(1.5)),
|
||||
("row_timeout_seconds", Value::from(0)),
|
||||
("row_timeout_seconds", Value::from(-0.25)),
|
||||
] {
|
||||
@@ -107,6 +112,11 @@ fn execution_controls_are_validated_before_canonical_serialization() {
|
||||
request.batch_size = Some(0);
|
||||
let error = request.to_canonical_json().unwrap_err();
|
||||
assert!(error.to_string().contains("batch_size"), "{error}");
|
||||
|
||||
let mut request = FunctionRegistrationRequest::from_json(&source).unwrap();
|
||||
request.max_concurrency = Some(0);
|
||||
let error = request.to_canonical_json().unwrap_err();
|
||||
assert!(error.to_string().contains("max_concurrency"), "{error}");
|
||||
}
|
||||
|
||||
/// The same shared golden as the Python suite builds from `@function(secrets=...)`
|
||||
|
||||
+1
-1
@@ -1 +1 @@
|
||||
{"artifact":{"adapter":{"kind":"scalar_to_arrow_batch","version":1},"content":{"data":"ZnJvbSBfX2Z1dHVyZV9fIGltcG9ydCBhbm5vdGF0aW9ucwoKZGVmIG5vcm1hbGl6ZV9zY29yZSh2YWx1ZTogZmxvYXQpIC0+IGZsb2F0OgogICAgcmV0dXJuIHZhbHVlIC8gMTAwLjAK","encoding":"base64"},"digest":"sha256:760784bdcef57b802f389b97804cc0b618aae39e86044733451bcd0b13089a7f","entrypoint":"normalize_score","kind":"python_callable"},"batch_size":2,"name":"normalize_score","row_timeout_seconds":0.75,"runtime":{"env":{"MODE":"test"},"environment":{"kind":"pip","packages":["numpy>=2"]},"kind":"python","python_version":"3.12"},"signature":{"inputs":[{"arrow_type":"float64","name":"value","nullable":false}],"output":{"arrow_type":"float64","kind":"scalar","nullable":false}}}
|
||||
{"artifact":{"adapter":{"kind":"scalar_to_arrow_batch","version":1},"content":{"data":"ZnJvbSBfX2Z1dHVyZV9fIGltcG9ydCBhbm5vdGF0aW9ucwoKZGVmIG5vcm1hbGl6ZV9zY29yZSh2YWx1ZTogZmxvYXQpIC0+IGZsb2F0OgogICAgcmV0dXJuIHZhbHVlIC8gMTAwLjAK","encoding":"base64"},"digest":"sha256:760784bdcef57b802f389b97804cc0b618aae39e86044733451bcd0b13089a7f","entrypoint":"normalize_score","kind":"python_callable"},"batch_size":2,"max_concurrency":4,"name":"normalize_score","row_timeout_seconds":0.75,"runtime":{"env":{"MODE":"test"},"environment":{"kind":"pip","packages":["numpy>=2"]},"kind":"python","python_version":"3.12"},"signature":{"inputs":[{"arrow_type":"float64","name":"value","nullable":false}],"output":{"arrow_type":"float64","kind":"scalar","nullable":false}}}
|
||||
|
||||
+1
@@ -1,6 +1,7 @@
|
||||
{
|
||||
"name": "normalize_score",
|
||||
"batch_size": 2,
|
||||
"max_concurrency": 4,
|
||||
"row_timeout_seconds": 0.75,
|
||||
"artifact": {
|
||||
"kind": "python_callable",
|
||||
|
||||
Reference in New Issue
Block a user