rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-11-13 14:16:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: expose Function refresh resource budgets in SDKs (#4951)

Browse Source 
## Summary

- Add the canonical Rust `FunctionWorkerResourceBudget` and
`RefreshColumnOptions::resource_budget()` for explicit Enterprise
Function refreshes.
- Forward the nested budget through Python sync/async, PyO3, Rust remote
requests, Node NAPI, and TypeScript.
- Require positive CPU and memory values; normalize an omitted nested
GPU count to numeric zero; reject the nested budget when combined with
the top-level `gpu_count` shorthand.
- Reject `resource_budget` on LanceDB Embedded refresh paths instead of
silently ignoring it.
- Document the Rust, Python, and TypeScript APIs, including that SQL
servers currently require all three budget fields.

The Rust, Python, and TypeScript additions use `cpu_millis` /
`cpuMillis`, `memory_bytes` / `memoryBytes`, and `gpu_count` /
`gpuCount`. The top-level GPU shorthand keeps its existing behavior and
remains mutually exclusive with the complete nested budget.

## Scope and compatibility

This change configures only the explicit refresh call. It does not
configure `add_columns` automatic refresh and does not introduce a
fleet-wide call-cap interface.

A LanceDB Enterprise server must expose the nested `resource_budget`
backfill contract. LanceDB Embedded continues to support its existing
expression refresh defaults and rejects the new complete budget
explicitly.

## Validation

Validated against the final signed product source:

- Rust: resource-budget remote serialization test `2 passed`; validation
test `1 passed`.
- Python: actual PyO3 build; remote-request matrix `22 passed`;
invalid-option matrix `20 passed`; local rejection `22 passed`; Ruff
check and format clean.
- Repository-standard Pyright entry point: `0 errors`, within the
repository's configured include set.
- Node: native build and Biome lint passed; the existing remote-request
test `1 passed`, including fractional, negative, overflow, NaN, and
infinity nested GPU cases; the existing local-rejection test `1 passed`.
- TypeDoc generation and generated-document cleanliness passed.
- The native-tested product source differs from this head only in two
semantically equivalent Node range checks required by Clippy. Current CI
compiles and lints the final blobs. No dependency or lockfile changes
are included.
This commit is contained in:
Xuanwo
2026-11-11 10:05:52 +08:00
committed by GitHub
parent d44f768920
commit 4d22c0ce2b
19 changed files with 623 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/src/js/classes/Table.md
+16 -5
View File
@@ -1141,11 +1141,18 @@ LanceDB Enterprise it is the server's backfill job.
Refresh options, mirroring the SQL
`REFRESH COLUMN ... WHERE ... WITH (...)` form: `where` narrows the rows,
`errorPolicy`, `maxFailedFraction`, `maxFailedRows`, `rowTimeoutSeconds`
and `initTimeoutSeconds` govern failing rows, and `batchSize`,
`gpuClass`, `gpuCount`, `numWorkers` and `maxExecutorProcessesPerWorker`
size the workers. All are LanceDB Enterprise only; a LanceDB Embedded
refresh rejects every option other than `batchSize`, `gpuClass` and
`gpuCount`.
and `initTimeoutSeconds` govern failing rows. `batchSize`, `gpuClass`,
`resourceBudget` and `maxExecutorProcessesPerWorker` size registered
Function refreshes. 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.
`gpuCount` at the top level is the shorthand that preserves the
deployment defaults for CPU and memory, and cannot be combined with
`resourceBudget`. `numWorkers` applies only to expression columns; use
`resourceBudget` for registered Functions. All are LanceDB Enterprise
only; a LanceDB Embedded refresh rejects every option other than
`batchSize`, `gpuClass` and `gpuCount`.
#### Returns
@@ -1156,6 +1163,10 @@ LanceDB Enterprise it is the server's backfill job.
```ts
const job = await table.refreshColumnAsync("embedding", {
where: "published_at > '2026-01-01'",
resourceBudget: {
cpuMillis: 8000,
memoryBytes: 16 * 1024 ** 3,
},
errorPolicy: "skip",
maxFailedFraction: 0.01,
rowTimeoutSeconds: 30,
docs/src/js/globals.md
+1
View File
@@ -99,6 +99,7 @@
- [FunctionErrorRecord](interfaces/FunctionErrorRecord.md)
- [FunctionErrors](interfaces/FunctionErrors.md)
- [FunctionErrorsOptions](interfaces/FunctionErrorsOptions.md)
- [FunctionWorkerResourceBudget](interfaces/FunctionWorkerResourceBudget.md)
- [HnswPqOptions](interfaces/HnswPqOptions.md)
- [HnswSqOptions](interfaces/HnswSqOptions.md)
- [IndexConfig](interfaces/IndexConfig.md)
docs/src/js/interfaces/FunctionWorkerResourceBudget.md
+39
View File
@@ -0,0 +1,39 @@
[**@lancedb/lancedb**](../README.md) • **Docs**
***
[@lancedb/lancedb](../globals.md) / FunctionWorkerResourceBudget
# Interface: FunctionWorkerResourceBudget
Aggregate worker resources for one Enterprise Function refresh.
## Properties
### cpuMillis
```ts
cpuMillis: number;
```
Aggregate CPU budget in millicores; 1000 is one CPU.
***
### gpuCount?
```ts
optional gpuCount: number;
```
Aggregate GPU budget. Omitted means zero, a CPU-only refresh.
***
### memoryBytes
```ts
memoryBytes: number;
```
Aggregate memory budget in bytes.
docs/src/js/interfaces/RefreshColumnOptions.md
+13 -1
View File
@@ -8,7 +8,8 @@
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.
`batchSize`, `gpuClass` and `gpuCount` and rejects the rest, including
`resourceBudget`.
## Properties
@@ -105,6 +106,17 @@ Worker count for an expression column; Function columns reject it.
***
### resourceBudget?
```ts
optional resourceBudget: FunctionWorkerResourceBudget;
```
Complete aggregate CPU and memory budget for a registered Function
refresh. `gpuCount` inside the budget defaults to zero.
***
### rowTimeoutSeconds?
```ts
docs/web/feature-engineering/refresh-column.mdx
+60 -4
View File
@@ -166,8 +166,62 @@ The resource values cap the Function workers together: this example allows
up to 8 CPUs, 16 GiB of memory, and 2 GPUs across the fleet. They are not
per-worker allocations or a guaranteed worker count. `num_workers` is only
accepted for SQL-expression columns; registered Function columns reject it.
When setting CPU or memory, supply all three budget fields, using
`gpu_count = 0` for a CPU-only Function.
The SDKs accept the same nested budget. CPU and memory are required; the GPU
count inside the budget defaults to zero, so a CPU-only Function can omit it:
<CodeGroup>
```python Python icon="python"
job = table.refresh_column_async(
"embedding",
resource_budget={
"cpu_millis": 8000,
"memory_bytes": 16 * 1024**3,
},
)
job.wait()
```
```rust Rust icon="rust"
use lancedb::table::{
FunctionWorkerResourceBudget,
RefreshColumnOptions,
};
let _job = table
.refresh_column_async_with_options(
"embedding",
RefreshColumnOptions::default().resource_budget(
FunctionWorkerResourceBudget::new(8000, 17_179_869_184),
),
)
.await?;
```
```typescript TypeScript icon="typescript"
const job = await table.refreshColumnAsync("embedding", {
resourceBudget: {
cpuMillis: 8000,
memoryBytes: 16 * 1024 ** 3,
},
});
await job.wait();
```
</CodeGroup>
On current SQL servers, the equivalent `REFRESH COLUMN ... WITH (...)` form
still takes all three fields when CPU or memory is set; use `gpu_count = 0`
for a CPU-only Function. The SDK normalizes an omitted nested GPU count to a
numeric zero before sending the request.
The existing top-level `gpu_count` / `gpuCount` option remains the shorthand
for changing only the aggregate GPU budget. It cannot be combined with the
complete resource budget. Omitting both preserves the server's existing
defaults.
These examples configure only the explicit refresh shown. They do not set the
budget for an automatic refresh scheduled after `add_columns`; that remains a
separate declaration and scheduler capability.
`batch_size` caps rows in one Function call; LanceDB adapts actual call and
work sizes. For Function columns, the refresh defaults to
@@ -248,7 +302,8 @@ async fn refresh(
</CodeGroup>
`gpu_count` is the aggregate GPU budget for the refresh's Function workers,
not a request for that many GPUs per worker. It must be positive.
not a request for that many GPUs per worker. It must be positive and preserves
the deployment's CPU and memory defaults.
`gpu_class` must name a class returned by `list_gpu_classes()`; class names
are lowercase DNS labels. Omit either option to use the Function's registered
class or the server's configured budget. `list_gpu_classes()` requires System
@@ -341,7 +396,8 @@ The SDK rejects a blank `where`, a `max_failed_fraction` outside 0 to 1, a
timeout that is not a positive number and a zero worker or process count
before sending the request; the server validates everything else. A LanceDB
Embedded refresh rejects every option other than `batch_size`, `gpu_class`
and `gpu_count`.
and `gpu_count`; in particular, it rejects `resource_budget` rather than
silently ignoring it.
## When a refresh fails {#failures}
nodejs/__test__/remote.test.ts
+83
View File
@@ -693,6 +693,12 @@ describe("remote connection", () => {
});
expect(job.id).toBe("j-9");
await table.refreshColumnAsync("embedding");
await table.refreshColumnAsync("embedding", {
resourceBudget: {
cpuMillis: 8_000,
memoryBytes: 16 * 1024 ** 3,
},
});
// Invalid options reject before any request is sent.
for (const [options, message] of [
[{ maxFailedRows: -1 }, "maxFailedRows"],
@@ -702,6 +708,75 @@ describe("remote connection", () => {
[{ rowTimeoutSeconds: 0 }, "row_timeout_seconds"],
[{ numWorkers: 0 }, "num_workers"],
[{ gpuCount: 0 }, "gpu_count"],
[
{
resourceBudget: {
cpuMillis: 1_000,
memoryBytes: 1_024,
},
gpuCount: 1,
},
"gpu_count cannot be combined with resource_budget",
],
[
{
resourceBudget: {
cpuMillis: 0,
memoryBytes: 1_024,
},
},
"resourceBudget.cpuMillis",
],
[
{
resourceBudget: {
cpuMillis: 1_000,
memoryBytes: 1_024,
gpuCount: 1.5,
},
},
"resourceBudget.gpuCount",
],
[
{
resourceBudget: {
cpuMillis: 1_000,
memoryBytes: 1_024,
gpuCount: -1,
},
},
"resourceBudget.gpuCount",
],
[
{
resourceBudget: {
cpuMillis: 1_000,
memoryBytes: 1_024,
gpuCount: 2 ** 32,
},
},
"resourceBudget.gpuCount",
],
[
{
resourceBudget: {
cpuMillis: 1_000,
memoryBytes: 1_024,
gpuCount: Number.NaN,
},
},
"resourceBudget.gpuCount",
],
[
{
resourceBudget: {
cpuMillis: 1_000,
memoryBytes: 1_024,
gpuCount: Number.POSITIVE_INFINITY,
},
},
"resourceBudget.gpuCount",
],
] as const) {
await expect(
table.refreshColumnAsync("embedding", options),
@@ -718,6 +793,14 @@ describe("remote connection", () => {
"max_executor_processes_per_worker": 2}`,
),
{ column: "embedding" },
JSON.parse(
`{"column": "embedding",
"resource_budget": {
"cpu_millis": 8000,
"memory_bytes": 17179869184,
"gpu_count": 0
}}`,
),
]);
});
nodejs/__test__/table.test.ts
+5
View File
@@ -5686,6 +5686,11 @@ describe("computed columns", () => {
await expect(
table.refreshColumnAsync("x", { errorPolicy: "skip" }),
).rejects.toThrow("error_policy requires a LanceDB Enterprise refresh");
await expect(
table.refreshColumnAsync("x", {
resourceBudget: { cpuMillis: 1_000, memoryBytes: 1_024 },
}),
).rejects.toThrow("resource_budget requires a LanceDB Enterprise refresh");
});
it("records Function errors only on remote tables", async () => {
nodejs/lancedb/index.ts
+1
View File
@@ -57,6 +57,7 @@ export {
AddResult,
AddColumnsResult,
RefreshColumnOptions,
FunctionWorkerResourceBudget,
RefreshColumnResult,
FunctionErrors,
FunctionErrorsOptions,
nodejs/lancedb/table.ts
+16 -5
View File
@@ -722,15 +722,26 @@ export abstract class Table {
* @param {RefreshColumnOptions} options Refresh options, mirroring the SQL
* `REFRESH COLUMN ... WHERE ... WITH (...)` form: `where` narrows the rows,
* `errorPolicy`, `maxFailedFraction`, `maxFailedRows`, `rowTimeoutSeconds`
* and `initTimeoutSeconds` govern failing rows, and `batchSize`,
* `gpuClass`, `gpuCount`, `numWorkers` and `maxExecutorProcessesPerWorker`
* size the workers. All are LanceDB Enterprise only; a LanceDB Embedded
* refresh rejects every option other than `batchSize`, `gpuClass` and
* `gpuCount`.
* and `initTimeoutSeconds` govern failing rows. `batchSize`, `gpuClass`,
* `resourceBudget` and `maxExecutorProcessesPerWorker` size registered
* Function refreshes. 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.
* `gpuCount` at the top level is the shorthand that preserves the
* deployment defaults for CPU and memory, and cannot be combined with
* `resourceBudget`. `numWorkers` applies only to expression columns; use
* `resourceBudget` for registered Functions. All are LanceDB Enterprise
* only; a LanceDB Embedded refresh rejects every option other than
* `batchSize`, `gpuClass` and `gpuCount`.
* @example
* ```ts
* const job = await table.refreshColumnAsync("embedding", {
* where: "published_at > '2026-01-01'",
* resourceBudget: {
* cpuMillis: 8000,
* memoryBytes: 16 * 1024 ** 3,
* },
* errorPolicy: "skip",
* maxFailedFraction: 0.01,
* rowTimeoutSeconds: 30,
nodejs/src/table.rs
+53 -2
View File
@@ -9,7 +9,8 @@ use lancedb::ipc::{ipc_file_to_batches, ipc_file_to_schema};
use lancedb::table::{
AddDataMode, ColumnAlteration as LanceColumnAlteration,
FieldMetadataUpdate as LanceFieldMetadataUpdate, FtsToken as LanceDbFtsToken,
NewColumnTransform, Ref, Table as LanceDbTable,
FunctionWorkerResourceBudget as LanceFunctionWorkerResourceBudget, NewColumnTransform, Ref,
Table as LanceDbTable,
};
use napi::Env;
use napi::bindgen_prelude::*;
@@ -1586,9 +1587,44 @@ pub struct FunctionErrorsOptions {
pub limit: Option<i64>,
}
/// Aggregate worker resources for one Enterprise Function refresh.
#[napi(object)]
#[derive(Clone, Debug)]
pub struct FunctionWorkerResourceBudget {
/// Aggregate CPU budget in millicores; 1000 is one CPU.
pub cpu_millis: f64,
/// Aggregate memory budget in bytes.
pub memory_bytes: f64,
/// Aggregate GPU budget. Omitted means zero, a CPU-only refresh.
pub gpu_count: Option<f64>,
}
fn resource_budget_integer(value: f64, name: &str) -> napi::Result<u64> {
const MAX_SAFE_INTEGER: f64 = 9_007_199_254_740_991.0;
if !value.is_finite() || value.fract() != 0.0 || !(1.0..=MAX_SAFE_INTEGER).contains(&value) {
return Err(napi::Error::from_reason(format!(
"{name} must be a positive safe integer"
)));
}
Ok(value as u64)
}
fn resource_budget_gpu_count(value: Option<f64>) -> napi::Result<u32> {
let Some(value) = value else {
return Ok(0);
};
if !value.is_finite() || value.fract() != 0.0 || !(0.0..=u32::MAX as f64).contains(&value) {
return Err(napi::Error::from_reason(
"resourceBudget.gpuCount must be an integer between 0 and 4294967295",
));
}
Ok(value as u32)
}
/// 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.
/// `batchSize`, `gpuClass` and `gpuCount` and rejects the rest, including
/// `resourceBudget`.
#[napi(object)]
#[derive(Clone, Debug, Default)]
pub struct RefreshColumnOptions {
@@ -1599,6 +1635,9 @@ pub struct RefreshColumnOptions {
pub batch_size: Option<u32>,
/// GPU class for Function workers.
pub gpu_class: Option<String>,
/// Complete aggregate CPU and memory budget for a registered Function
/// refresh. `gpuCount` inside the budget defaults to zero.
pub resource_budget: Option<FunctionWorkerResourceBudget>,
/// Positive aggregate GPU budget for Function workers.
pub gpu_count: Option<u32>,
/// `"fail"` (default) stops on a row error; `"skip"` writes null and
@@ -1628,6 +1667,18 @@ impl TryFrom<RefreshColumnOptions> for lancedb::table::RefreshColumnOptions {
options.where_clause = value.where_clause;
options.batch_size = value.batch_size.map(|size| size as usize);
options.gpu_class = value.gpu_class;
options.resource_budget = value
.resource_budget
.map(|budget| {
Ok::<_, napi::Error>(
LanceFunctionWorkerResourceBudget::new(
resource_budget_integer(budget.cpu_millis, "resourceBudget.cpuMillis")?,
resource_budget_integer(budget.memory_bytes, "resourceBudget.memoryBytes")?,
)
.with_gpu_count(resource_budget_gpu_count(budget.gpu_count)?),
)
})
.transpose()?;
options.gpu_count = value.gpu_count;
options.error_policy = match value.error_policy.as_deref() {
None => None,
python/python/lancedb/_lancedb.pyi
+2
View File
@@ -627,6 +627,7 @@ class Table:
where: Optional[str] = None,
batch_size: Optional[int] = None,
publication_interval_rows: Optional[int] = None,
resource_budget: Optional[dict[str, int]] = None,
gpu_class: Optional[str] = None,
gpu_count: Optional[int] = None,
retry_policy: Optional[str] = None,
@@ -645,6 +646,7 @@ class Table:
where: Optional[str] = None,
batch_size: Optional[int] = None,
publication_interval_rows: Optional[int] = None,
resource_budget: Optional[dict[str, int]] = None,
gpu_class: Optional[str] = None,
gpu_count: Optional[int] = None,
retry_policy: Optional[str] = None,
python/python/lancedb/remote/table.py
+4
View File
@@ -1144,6 +1144,7 @@ class RemoteTable(Table):
where: Optional[str] = None,
batch_size: Optional[int] = None,
publication_interval_rows: Optional[int] = None,
resource_budget: Optional[Dict[str, int]] = None,
gpu_class: Optional[str] = None,
gpu_count: Optional[int] = None,
retry_policy: Optional[FunctionRetryPolicy] = None,
@@ -1161,6 +1162,7 @@ class RemoteTable(Table):
where=where,
batch_size=batch_size,
publication_interval_rows=publication_interval_rows,
resource_budget=resource_budget,
gpu_class=gpu_class,
gpu_count=gpu_count,
retry_policy=retry_policy,
@@ -1181,6 +1183,7 @@ class RemoteTable(Table):
where: Optional[str] = None,
batch_size: Optional[int] = None,
publication_interval_rows: Optional[int] = None,
resource_budget: Optional[Dict[str, int]] = None,
gpu_class: Optional[str] = None,
gpu_count: Optional[int] = None,
retry_policy: Optional[FunctionRetryPolicy] = None,
@@ -1199,6 +1202,7 @@ class RemoteTable(Table):
where=where,
batch_size=batch_size,
publication_interval_rows=publication_interval_rows,
resource_budget=resource_budget,
gpu_class=gpu_class,
gpu_count=gpu_count,
retry_policy=retry_policy,
python/python/lancedb/table.py
+62 -16
View File
@@ -3187,6 +3187,7 @@ class Table(ABC):
where: Optional[str] = None,
batch_size: Optional[int] = None,
publication_interval_rows: Optional[int] = None,
resource_budget: Optional[Dict[str, int]] = None,
gpu_class: Optional[str] = None,
gpu_count: Optional[int] = None,
retry_policy: Optional[FunctionRetryPolicy] = None,
@@ -3231,15 +3232,23 @@ class Table(ABC):
actual publication boundary can be larger than this threshold
because it also depends on completed work and Function batching.
LanceDB Enterprise only; local refreshes reject this option.
resource_budget: dict[str, int], optional
Aggregate resources for one registered Function refresh. The
``cpu_millis`` and ``memory_bytes`` keys are required;
``gpu_count`` is optional and defaults to zero, which requests a
CPU-only budget. This budget cannot be combined with the top-level
``gpu_count`` shorthand. LanceDB Enterprise only; local refreshes
reject it.
gpu_class: str, optional
GPU class for Function workers. Must be a lowercase name of 1-63
letters, digits, or hyphens, without a leading or trailing hyphen.
None keeps the Function's registered class or the deployment
default. LanceDB Enterprise only.
gpu_count: int, optional
Positive aggregate GPU budget for Function workers. None leaves
the server's configured budget unchanged. LanceDB Enterprise
only.
Positive aggregate GPU budget for Function workers without changing
the deployment's CPU and memory defaults. None leaves the server's
configured budget unchanged. For a complete CPU and memory budget,
use ``resource_budget`` instead. LanceDB Enterprise only.
retry_policy: FunctionRetryPolicy, optional
Override retries of failed user-code calls. Only use when repeating
partial work is safe. Set max_attempts=1 to disable retries.
@@ -3259,7 +3268,8 @@ class Table(ABC):
Deadline for Function initialization. LanceDB Enterprise only.
num_workers: int, optional
Worker count for an expression column; registered Function
columns reject it. LanceDB Enterprise only.
columns reject it and use ``resource_budget`` instead. LanceDB
Enterprise only.
max_executor_processes_per_worker: int, optional
Cap on executor processes per Function worker. LanceDB
Enterprise only.
@@ -3279,6 +3289,7 @@ class Table(ABC):
where: Optional[str] = None,
batch_size: Optional[int] = None,
publication_interval_rows: Optional[int] = None,
resource_budget: Optional[Dict[str, int]] = None,
gpu_class: Optional[str] = None,
gpu_count: Optional[int] = None,
retry_policy: Optional[FunctionRetryPolicy] = None,
@@ -3322,15 +3333,23 @@ class Table(ABC):
actual publication boundary can be larger than this threshold
because it also depends on completed work and Function batching.
LanceDB Enterprise only; local refreshes reject this option.
resource_budget: dict[str, int], optional
Aggregate resources for one registered Function refresh. The
``cpu_millis`` and ``memory_bytes`` keys are required;
``gpu_count`` is optional and defaults to zero, which requests a
CPU-only budget. This budget cannot be combined with the top-level
``gpu_count`` shorthand. LanceDB Enterprise only; local refreshes
reject it.
gpu_class: str, optional
GPU class for Function workers. Must be a lowercase name of 1-63
letters, digits, or hyphens, without a leading or trailing hyphen.
None keeps the Function's registered class or the deployment
default. LanceDB Enterprise only.
gpu_count: int, optional
Positive aggregate GPU budget for Function workers. None leaves
the server's configured budget unchanged. LanceDB Enterprise
only.
Positive aggregate GPU budget for Function workers without changing
the deployment's CPU and memory defaults. None leaves the server's
configured budget unchanged. For a complete CPU and memory budget,
use ``resource_budget`` instead. LanceDB Enterprise only.
retry_policy: FunctionRetryPolicy, optional
Override retries of failed user-code calls. Only use when repeating
partial work is safe. Set max_attempts=1 to disable retries.
@@ -3350,7 +3369,8 @@ class Table(ABC):
Deadline for Function initialization. LanceDB Enterprise only.
num_workers: int, optional
Worker count for an expression column; registered Function
columns reject it. LanceDB Enterprise only.
columns reject it and use ``resource_budget`` instead. LanceDB
Enterprise only.
max_executor_processes_per_worker: int, optional
Cap on executor processes per Function worker. LanceDB
Enterprise only.
@@ -5465,6 +5485,7 @@ class LanceTable(Table):
where: Optional[str] = None,
batch_size: Optional[int] = None,
publication_interval_rows: Optional[int] = None,
resource_budget: Optional[Dict[str, int]] = None,
gpu_class: Optional[str] = None,
gpu_count: Optional[int] = None,
retry_policy: Optional[FunctionRetryPolicy] = None,
@@ -5485,6 +5506,7 @@ class LanceTable(Table):
where=where,
batch_size=batch_size,
publication_interval_rows=publication_interval_rows,
resource_budget=resource_budget,
gpu_class=gpu_class,
gpu_count=gpu_count,
retry_policy=retry_policy,
@@ -5505,6 +5527,7 @@ class LanceTable(Table):
where: Optional[str] = None,
batch_size: Optional[int] = None,
publication_interval_rows: Optional[int] = None,
resource_budget: Optional[Dict[str, int]] = None,
gpu_class: Optional[str] = None,
gpu_count: Optional[int] = None,
retry_policy: Optional[FunctionRetryPolicy] = None,
@@ -5527,6 +5550,7 @@ class LanceTable(Table):
where=where,
batch_size=batch_size,
publication_interval_rows=publication_interval_rows,
resource_budget=resource_budget,
gpu_class=gpu_class,
gpu_count=gpu_count,
retry_policy=retry_policy,
@@ -7895,6 +7919,7 @@ class AsyncTable:
where: Optional[str] = None,
batch_size: Optional[int] = None,
publication_interval_rows: Optional[int] = None,
resource_budget: Optional[Dict[str, int]] = None,
gpu_class: Optional[str] = None,
gpu_count: Optional[int] = None,
retry_policy: Optional[FunctionRetryPolicy] = None,
@@ -7939,15 +7964,23 @@ class AsyncTable:
actual publication boundary can be larger than this threshold
because it also depends on completed work and Function batching.
LanceDB Enterprise only; local refreshes reject this option.
resource_budget: dict[str, int], optional
Aggregate resources for one registered Function refresh. The
``cpu_millis`` and ``memory_bytes`` keys are required;
``gpu_count`` is optional and defaults to zero, which requests a
CPU-only budget. This budget cannot be combined with the top-level
``gpu_count`` shorthand. LanceDB Enterprise only; local refreshes
reject it.
gpu_class: str, optional
GPU class for Function workers. Must be a lowercase name of 1-63
letters, digits, or hyphens, without a leading or trailing hyphen.
None keeps the Function's registered class or the deployment
default. LanceDB Enterprise only.
gpu_count: int, optional
Positive aggregate GPU budget for Function workers. None leaves
the server's configured budget unchanged. LanceDB Enterprise
only.
Positive aggregate GPU budget for Function workers without changing
the deployment's CPU and memory defaults. None leaves the server's
configured budget unchanged. For a complete CPU and memory budget,
use ``resource_budget`` instead. LanceDB Enterprise only.
retry_policy: FunctionRetryPolicy, optional
Override retries of failed user-code calls. Only use when repeating
partial work is safe. Set max_attempts=1 to disable retries.
@@ -7967,7 +8000,8 @@ class AsyncTable:
Deadline for Function initialization. LanceDB Enterprise only.
num_workers: int, optional
Worker count for an expression column; registered Function
columns reject it. LanceDB Enterprise only.
columns reject it and use ``resource_budget`` instead. LanceDB
Enterprise only.
max_executor_processes_per_worker: int, optional
Cap on executor processes per Function worker. LanceDB
Enterprise only.
@@ -7982,6 +8016,7 @@ class AsyncTable:
where=where,
batch_size=batch_size,
publication_interval_rows=publication_interval_rows,
resource_budget=resource_budget,
gpu_class=gpu_class,
gpu_count=gpu_count,
retry_policy=None
@@ -8092,6 +8127,7 @@ class AsyncTable:
where: Optional[str] = None,
batch_size: Optional[int] = None,
publication_interval_rows: Optional[int] = None,
resource_budget: Optional[Dict[str, int]] = None,
gpu_class: Optional[str] = None,
gpu_count: Optional[int] = None,
retry_policy: Optional[FunctionRetryPolicy] = None,
@@ -8135,15 +8171,23 @@ class AsyncTable:
actual publication boundary can be larger than this threshold
because it also depends on completed work and Function batching.
LanceDB Enterprise only; local refreshes reject this option.
resource_budget: dict[str, int], optional
Aggregate resources for one registered Function refresh. The
``cpu_millis`` and ``memory_bytes`` keys are required;
``gpu_count`` is optional and defaults to zero, which requests a
CPU-only budget. This budget cannot be combined with the top-level
``gpu_count`` shorthand. LanceDB Enterprise only; local refreshes
reject it.
gpu_class: str, optional
GPU class for Function workers. Must be a lowercase name of 1-63
letters, digits, or hyphens, without a leading or trailing hyphen.
None keeps the Function's registered class or the deployment
default. LanceDB Enterprise only.
gpu_count: int, optional
Positive aggregate GPU budget for Function workers. None leaves
the server's configured budget unchanged. LanceDB Enterprise
only.
Positive aggregate GPU budget for Function workers without changing
the deployment's CPU and memory defaults. None leaves the server's
configured budget unchanged. For a complete CPU and memory budget,
use ``resource_budget`` instead. LanceDB Enterprise only.
retry_policy: FunctionRetryPolicy, optional
Override retries of failed user-code calls. Only use when repeating
partial work is safe. Set max_attempts=1 to disable retries.
@@ -8163,7 +8207,8 @@ class AsyncTable:
Deadline for Function initialization. LanceDB Enterprise only.
num_workers: int, optional
Worker count for an expression column; registered Function
columns reject it. LanceDB Enterprise only.
columns reject it and use ``resource_budget`` instead. LanceDB
Enterprise only.
max_executor_processes_per_worker: int, optional
Cap on executor processes per Function worker. LanceDB
Enterprise only.
@@ -8195,6 +8240,7 @@ class AsyncTable:
where=where,
batch_size=batch_size,
publication_interval_rows=publication_interval_rows,
resource_budget=resource_budget,
gpu_class=gpu_class,
gpu_count=gpu_count,
retry_policy=None
python/python/tests/test_remote_db.py
+47
View File
@@ -1864,6 +1864,14 @@ def test_remote_create_index_async_returns_job():
{"batch_size": 1},
{"batch_size": 8},
{"publication_interval_rows": 60},
{"resource_budget": {"cpu_millis": 2000, "memory_bytes": 4 * 1024**3}},
{
"resource_budget": {
"cpu_millis": 8000,
"memory_bytes": 16 * 1024**3,
"gpu_count": 2,
}
},
{"gpu_class": "a10g", "gpu_count": 2},
{
"retry_policy": {"transient": True, "max_attempts": 3},
@@ -1904,6 +1912,12 @@ async def test_remote_refresh_async_returns_typed_terminal_result(
expected["gpu_class"] = refresh_kwargs["gpu_class"]
if refresh_kwargs.get("gpu_count") is not None:
expected["gpu_count"] = refresh_kwargs["gpu_count"]
if refresh_kwargs.get("resource_budget") is not None:
budget = refresh_kwargs["resource_budget"]
expected["resource_budget"] = {
**budget,
"gpu_count": budget.get("gpu_count", 0),
}
if refresh_kwargs.get("retry_policy") is not None:
expected["retry_policy"] = lancedb.FunctionRetryPolicy.model_validate(
refresh_kwargs["retry_policy"]
@@ -2093,6 +2107,39 @@ def test_remote_refresh_async_forwards_progress_from_submitted_handle():
({"publication_interval_rows": 0}, "publication_interval_rows"),
({"publication_interval_rows": -1}, "publication_interval_rows"),
({"num_workers": 0}, "num_workers"),
(
{"resource_budget": {"cpu_millis": 1000}},
"resource_budget.memory_bytes",
),
(
{
"resource_budget": {
"cpu_millis": 0,
"memory_bytes": 1,
}
},
"resource_budget.cpu_millis",
),
(
{
"resource_budget": {
"cpu_millis": 1000,
"memory_bytes": 1,
},
"gpu_count": 1,
},
"gpu_count cannot be combined with resource_budget",
),
(
{
"resource_budget": {
"cpu_millis": 1000,
"memory_bytes": 1,
"unknown": 1,
}
},
"unsupported key 'unknown'",
),
(
{"max_executor_processes_per_worker": -2},
"max_executor_processes_per_worker",
python/python/tests/test_table.py
+4
View File
@@ -5996,6 +5996,10 @@ def test_refresh_column_rejects_invalid_batch_size(tmp_path, method, batch_size)
({"row_timeout_seconds": 30}, "row_timeout_seconds"),
({"init_timeout_seconds": 60}, "init_timeout_seconds"),
({"publication_interval_rows": 60}, "publication_interval_rows"),
(
{"resource_budget": {"cpu_millis": 1000, "memory_bytes": 1024}},
"resource_budget",
),
({"num_workers": 2}, "num_workers"),
({"max_executor_processes_per_worker": 2}, "max_executor_processes_per_worker"),
],
python/src/table.rs
+79 -11
View File
@@ -29,15 +29,17 @@ use lancedb::connection::LanceFileVersion;
use lancedb::index::scalar::FtsIndexBuilder;
use lancedb::table::{
AddDataMode, ColumnAlteration, CompactionMode, CompactionOptions, Duration,
FieldMetadataUpdate, FtsToken as LanceDbFtsToken, IndexRemapMode, NewColumnTransform, Ref,
RefreshColumnOptions, Table as LanceDbTable,
FieldMetadataUpdate, FtsToken as LanceDbFtsToken, FunctionWorkerResourceBudget, IndexRemapMode,
NewColumnTransform, Ref, RefreshColumnOptions, Table as LanceDbTable,
};
use lancedb::tokenize as lancedb_tokenize;
use pyo3::{
Bound, FromPyObject, Py, PyAny, PyRef, PyResult, Python,
exceptions::{PyRuntimeError, PyValueError},
pyclass, pyfunction, pymethods,
types::{IntoPyDict, PyAnyMethods, PyBytes, PyDict, PyDictMethods, PyList, PyListMethods},
types::{
IntoPyDict, PyAnyMethods, PyBool, PyBytes, PyDict, PyDictMethods, PyList, PyListMethods,
},
};
mod scannable;
@@ -47,6 +49,7 @@ struct RefreshColumnArgs {
where_clause: Option<String>,
batch_size: Option<i64>,
publication_interval_rows: Option<i64>,
resource_budget: Option<FunctionWorkerResourceBudget>,
gpu_class: Option<String>,
gpu_count: Option<i64>,
retry_policy: Option<String>,
@@ -70,11 +73,71 @@ fn positive_count<T: TryFrom<i64>>(value: Option<i64>, name: &str) -> PyResult<O
.transpose()
}
fn resource_budget_integer(
value: &Bound<'_, PyAny>,
name: &str,
allow_zero: bool,
) -> PyResult<u64> {
if value.is_none() || value.is_instance_of::<PyBool>() {
return Err(PyValueError::new_err(format!(
"resource_budget.{name} must be an integer"
)));
}
let parsed = value
.extract::<u64>()
.map_err(|_| PyValueError::new_err(format!("resource_budget.{name} must be an integer")))?;
if !allow_zero && parsed == 0 {
return Err(PyValueError::new_err(format!(
"resource_budget.{name} must be greater than zero"
)));
}
Ok(parsed)
}
fn parse_resource_budget(
budget: Option<&Bound<'_, PyDict>>,
) -> PyResult<Option<FunctionWorkerResourceBudget>> {
let Some(budget) = budget else {
return Ok(None);
};
for key in budget.keys() {
let key = key
.extract::<String>()
.map_err(|_| PyValueError::new_err("resource_budget keys must be strings"))?;
if !matches!(key.as_str(), "cpu_millis" | "memory_bytes" | "gpu_count") {
return Err(PyValueError::new_err(format!(
"resource_budget contains unsupported key '{key}'"
)));
}
}
let cpu_millis = budget
.get_item("cpu_millis")?
.ok_or_else(|| PyValueError::new_err("resource_budget.cpu_millis is required"))?;
let cpu_millis = resource_budget_integer(&cpu_millis, "cpu_millis", false)?;
let memory_bytes = budget
.get_item("memory_bytes")?
.ok_or_else(|| PyValueError::new_err("resource_budget.memory_bytes is required"))?;
let memory_bytes = resource_budget_integer(&memory_bytes, "memory_bytes", false)?;
let gpu_count = match budget.get_item("gpu_count")? {
None => 0,
Some(value) if value.is_none() => 0,
Some(value) => {
u32::try_from(resource_budget_integer(&value, "gpu_count", true)?).map_err(|_| {
PyValueError::new_err("resource_budget.gpu_count must be no greater than u32::MAX")
})?
}
};
Ok(Some(
FunctionWorkerResourceBudget::new(cpu_millis, memory_bytes).with_gpu_count(gpu_count),
))
}
fn refresh_column_options(args: RefreshColumnArgs) -> PyResult<RefreshColumnOptions> {
let RefreshColumnArgs {
where_clause,
batch_size,
publication_interval_rows,
resource_budget,
gpu_class,
gpu_count,
retry_policy,
@@ -118,6 +181,7 @@ fn refresh_column_options(args: RefreshColumnArgs) -> PyResult<RefreshColumnOpti
.map_err(|_| PyValueError::new_err("batch_size must be greater than zero"))?;
options.publication_interval_rows =
positive_count(publication_interval_rows, "publication_interval_rows")?;
options.resource_budget = resource_budget;
options.gpu_class = gpu_class;
options.gpu_count = gpu_count
.map(|count| {
@@ -2210,14 +2274,15 @@ impl Table {
})
}
#[pyo3(signature = (column, *, r#where=None, batch_size=None, publication_interval_rows=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, 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(
self_: PyRef<'_, Self>,
pub fn refresh_column<'py>(
self_: PyRef<'py, Self>,
column: String,
r#where: Option<String>,
batch_size: Option<i64>,
publication_interval_rows: Option<i64>,
resource_budget: Option<Bound<'py, PyDict>>,
gpu_class: Option<String>,
gpu_count: Option<i64>,
retry_policy: Option<String>,
@@ -2228,11 +2293,12 @@ impl Table {
init_timeout_seconds: Option<f64>,
num_workers: Option<i64>,
max_executor_processes_per_worker: Option<i64>,
) -> PyResult<Bound<'_, PyAny>> {
) -> PyResult<Bound<'py, PyAny>> {
let options = refresh_column_options(RefreshColumnArgs {
where_clause: r#where,
batch_size,
publication_interval_rows,
resource_budget: parse_resource_budget(resource_budget.as_ref())?,
gpu_class,
gpu_count,
retry_policy,
@@ -2254,14 +2320,15 @@ impl Table {
})
}
#[pyo3(signature = (column, *, r#where=None, batch_size=None, publication_interval_rows=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, 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(
self_: PyRef<'_, Self>,
pub fn refresh_column_async<'py>(
self_: PyRef<'py, Self>,
column: String,
r#where: Option<String>,
batch_size: Option<i64>,
publication_interval_rows: Option<i64>,
resource_budget: Option<Bound<'py, PyDict>>,
gpu_class: Option<String>,
gpu_count: Option<i64>,
retry_policy: Option<String>,
@@ -2272,11 +2339,12 @@ impl Table {
init_timeout_seconds: Option<f64>,
num_workers: Option<i64>,
max_executor_processes_per_worker: Option<i64>,
) -> PyResult<Bound<'_, PyAny>> {
) -> PyResult<Bound<'py, PyAny>> {
let options = refresh_column_options(RefreshColumnArgs {
where_clause: r#where,
batch_size,
publication_interval_rows,
resource_budget: parse_resource_budget(resource_budget.as_ref())?,
gpu_class,
gpu_count,
retry_policy,
rust/lancedb/src/remote/table.rs
+59 -1
View File
@@ -11028,6 +11028,48 @@ mod tests {
assert_eq!(job.id(), Some("j-gpu"));
}
#[rstest]
#[case(None, 0)]
#[case(Some(2), 2)]
#[tokio::test]
async fn test_refresh_column_async_resource_budget(
#[case] gpu_count: Option<u32>,
#[case] expected_gpu_count: u32,
) {
use crate::table::{FunctionWorkerResourceBudget, RefreshColumnOptions};
let table = Table::new_with_handler("my_table", move |request| {
assert_eq!(request.url().path(), "/v1/table/my_table/backfill_column");
assert_eq!(
request_body_json(&request),
serde_json::json!({
"column": "quality",
"resource_budget": {
"cpu_millis": 8_000,
"memory_bytes": 17_179_869_184_u64,
"gpu_count": expected_gpu_count,
},
})
);
http::Response::builder()
.status(202)
.body(r#"{"job_id":"j-budget"}"#)
.unwrap()
});
let mut budget = FunctionWorkerResourceBudget::new(8_000, 17_179_869_184);
if let Some(gpu_count) = gpu_count {
budget = budget.with_gpu_count(gpu_count);
}
let job = table
.refresh_column_async_with_options(
"quality",
RefreshColumnOptions::default().resource_budget(budget),
)
.await
.unwrap();
assert_eq!(job.id(), Some("j-budget"));
}
#[tokio::test]
async fn test_refresh_column_async_retry_policy() {
use crate::function::{FunctionErrorPolicy, FunctionRetryPolicy};
@@ -11155,7 +11197,7 @@ mod tests {
#[tokio::test]
async fn test_refresh_column_rejects_invalid_gpu_options_before_request() {
use crate::table::RefreshColumnOptions;
use crate::table::{FunctionWorkerResourceBudget, RefreshColumnOptions};
let table = Table::new_with_handler("my_table", |_| -> http::Response<String> {
panic!("invalid GPU options must not send a request")
@@ -11170,6 +11212,22 @@ mod tests {
"gpu_class",
),
(RefreshColumnOptions::default().gpu_count(0), "gpu_count"),
(
RefreshColumnOptions::default()
.resource_budget(FunctionWorkerResourceBudget::new(0, 1)),
"resource_budget.cpu_millis",
),
(
RefreshColumnOptions::default()
.resource_budget(FunctionWorkerResourceBudget::new(1, 0)),
"resource_budget.memory_bytes",
),
(
RefreshColumnOptions::default()
.resource_budget(FunctionWorkerResourceBudget::new(1, 1))
.gpu_count(1),
"gpu_count cannot be combined with resource_budget",
),
] {
let error = table
.refresh_column_async_with_options("quality", options)
rust/lancedb/src/table.rs
+16 -10
View File
@@ -114,7 +114,7 @@ pub use optimize::{
CompactionMode, CompactionOptions, CompactionOverrides, IndexRemapMode, OptimizeAction,
OptimizeStats,
};
pub use refresh::{RefreshColumnOptions, RefreshColumnResult};
pub use refresh::{FunctionWorkerResourceBudget, RefreshColumnOptions, RefreshColumnResult};
pub use rte_stats::{MemtableStats, RTEStats, SsTableStats, TableShardStats};
pub use schema_evolution::{
AddColumnsResult, AlterColumnsResult, DropColumnsResult, FieldMetadataUpdate,
@@ -2144,10 +2144,14 @@ impl Table {
/// [`RefreshColumnOptions::publication_interval_rows()`] sets an opt-in
/// row-count publication threshold for completed Function outputs; unset
/// keeps whole-fragment publication.
/// [`RefreshColumnOptions::resource_budget()`] sets the aggregate CPU,
/// memory, and GPU budget for a Function refresh.
/// [`RefreshColumnOptions::gpu_class()`] selects a GPU pool and
/// [`RefreshColumnOptions::gpu_count()`] sets the aggregate GPU budget.
/// [`RefreshColumnOptions::gpu_count()`] sets the aggregate GPU budget
/// without changing the deployment's CPU and memory defaults.
/// Batch-size and GPU options have no effect on expression-backed columns
/// in LanceDB Embedded; a publication threshold is rejected there.
/// in LanceDB Embedded; a resource budget and publication threshold are
/// rejected there.
/// Remote refreshes require
/// [`Table::refresh_column_async_with_options`].
///
@@ -2208,13 +2212,15 @@ impl Table {
/// [`RefreshColumnOptions::publication_interval_rows()`] to publish
/// completed Function outputs at a row-count threshold; unset keeps
/// whole-fragment publication. Set
/// [`RefreshColumnOptions::gpu_class()`] and
/// [`RefreshColumnOptions::gpu_count()`] to override the registered GPU
/// class or set an aggregate GPU budget for Function workers. Unset values
/// use the Function and server defaults; zero counts and malformed class
/// names are rejected before submitting a job. Embedded expression
/// refreshes ignore the batch-size and GPU options, and reject publication
/// thresholds and other Enterprise-only options.
/// [`RefreshColumnOptions::resource_budget()`] to pass a complete aggregate
/// CPU, memory, and GPU budget for a Function refresh. The standalone
/// [`RefreshColumnOptions::gpu_count()`] shorthand and
/// [`RefreshColumnOptions::gpu_class()`] placement remain available; a
/// resource budget cannot be combined with the gpu-count shorthand. Unset
/// values use the Function and server defaults; invalid counts and malformed
/// class names are rejected before submitting a job. Embedded expression
/// refreshes ignore the batch-size and GPU options, and reject resource
/// budgets, publication thresholds, and other Enterprise-only options.
///
/// ```
/// # use lancedb::{Table, table::RefreshColumnOptions};
rust/lancedb/src/table/refresh.rs
+63
View File
@@ -55,6 +55,38 @@ use crate::function::validate_gpu_class_name;
use crate::job::Job;
use crate::{Error, Result};
/// Aggregate worker resources for one Enterprise Function refresh.
///
/// CPU is in millicores (1000 = one core) and memory is in bytes. These
/// values cap the Function workers together; they are not per-worker
/// allocations. The default GPU budget is zero, which describes a CPU-only
/// Function. This option is rejected by LanceDB Embedded refreshes.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
#[non_exhaustive]
pub struct FunctionWorkerResourceBudget {
pub cpu_millis: u64,
pub memory_bytes: u64,
/// Aggregate GPU budget across the refresh. Zero is CPU-only.
pub gpu_count: u32,
}
impl FunctionWorkerResourceBudget {
/// Create a CPU-only aggregate budget.
pub fn new(cpu_millis: u64, memory_bytes: u64) -> Self {
Self {
cpu_millis,
memory_bytes,
gpu_count: 0,
}
}
/// Set the aggregate GPU budget.
pub fn with_gpu_count(mut self, gpu_count: u32) -> Self {
self.gpu_count = gpu_count;
self
}
}
/// Options for refreshing a computed column.
///
/// Serializes to the Enterprise backfill request body, one key per set
@@ -108,6 +140,10 @@ pub struct RefreshColumnOptions {
/// rejects this option.
#[serde(skip_serializing_if = "Option::is_none")]
pub publication_interval_rows: Option<u64>,
/// Complete aggregate CPU, memory, and GPU budget for a registered
/// Function refresh. Cannot be combined with `gpu_count`.
#[serde(skip_serializing_if = "Option::is_none")]
pub resource_budget: Option<FunctionWorkerResourceBudget>,
/// GPU class for Function workers in LanceDB Enterprise.
///
/// `None` keeps the Function's registered class, or the deployment default
@@ -189,6 +225,15 @@ impl RefreshColumnOptions {
self
}
/// Set the aggregate worker resource budget for a Function refresh.
///
/// Use [`FunctionWorkerResourceBudget::new`] for a CPU-only budget, or
/// [`FunctionWorkerResourceBudget::with_gpu_count`] to include GPUs.
pub fn resource_budget(mut self, budget: FunctionWorkerResourceBudget) -> Self {
self.resource_budget = Some(budget);
self
}
/// Request Function workers from `gpu_class`.
pub fn gpu_class(mut self, gpu_class: impl Into<String>) -> Self {
self.gpu_class = Some(gpu_class.into());
@@ -220,6 +265,7 @@ impl RefreshColumnOptions {
"publication_interval_rows",
self.publication_interval_rows.is_some(),
),
("resource_budget", self.resource_budget.is_some()),
]
.into_iter()
.find_map(|(name, set)| set.then_some(name))
@@ -277,6 +323,23 @@ impl RefreshColumnOptions {
message: "publication_interval_rows must be greater than zero".into(),
});
}
if self.resource_budget.is_some() && self.gpu_count.is_some() {
return Err(Error::InvalidInput {
message: "gpu_count cannot be combined with resource_budget".into(),
});
}
if let Some(budget) = &self.resource_budget {
if budget.cpu_millis == 0 {
return Err(Error::InvalidInput {
message: "resource_budget.cpu_millis must be greater than zero".into(),
});
}
if budget.memory_bytes == 0 {
return Err(Error::InvalidInput {
message: "resource_budget.memory_bytes must be greater than zero".into(),
});
}
}
if let Some(gpu_class) = &self.gpu_class {
validate_gpu_class_name(gpu_class)?;
}