diff --git a/docs/src/js/classes/Table.md b/docs/src/js/classes/Table.md index 1fad723fe..5ab1751d4 100644 --- a/docs/src/js/classes/Table.md +++ b/docs/src/js/classes/Table.md @@ -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, diff --git a/docs/src/js/globals.md b/docs/src/js/globals.md index dab1c61bc..50f55f937 100644 --- a/docs/src/js/globals.md +++ b/docs/src/js/globals.md @@ -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) diff --git a/docs/src/js/interfaces/FunctionWorkerResourceBudget.md b/docs/src/js/interfaces/FunctionWorkerResourceBudget.md new file mode 100644 index 000000000..0352bc833 --- /dev/null +++ b/docs/src/js/interfaces/FunctionWorkerResourceBudget.md @@ -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. diff --git a/docs/src/js/interfaces/RefreshColumnOptions.md b/docs/src/js/interfaces/RefreshColumnOptions.md index 840012040..151cbb72f 100644 --- a/docs/src/js/interfaces/RefreshColumnOptions.md +++ b/docs/src/js/interfaces/RefreshColumnOptions.md @@ -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 diff --git a/docs/web/feature-engineering/refresh-column.mdx b/docs/web/feature-engineering/refresh-column.mdx index 0dfe3b33d..043c712cb 100644 --- a/docs/web/feature-engineering/refresh-column.mdx +++ b/docs/web/feature-engineering/refresh-column.mdx @@ -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: + + +```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(); +``` + + +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( `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} diff --git a/nodejs/__test__/remote.test.ts b/nodejs/__test__/remote.test.ts index 242be6101..e26adc58d 100644 --- a/nodejs/__test__/remote.test.ts +++ b/nodejs/__test__/remote.test.ts @@ -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 + }}`, + ), ]); }); diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index c3888d6c0..84f831a7a 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -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 () => { diff --git a/nodejs/lancedb/index.ts b/nodejs/lancedb/index.ts index 9644b0060..6bd46c9c8 100644 --- a/nodejs/lancedb/index.ts +++ b/nodejs/lancedb/index.ts @@ -57,6 +57,7 @@ export { AddResult, AddColumnsResult, RefreshColumnOptions, + FunctionWorkerResourceBudget, RefreshColumnResult, FunctionErrors, FunctionErrorsOptions, diff --git a/nodejs/lancedb/table.ts b/nodejs/lancedb/table.ts index 5ef298c8c..d7aac1fa0 100644 --- a/nodejs/lancedb/table.ts +++ b/nodejs/lancedb/table.ts @@ -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, diff --git a/nodejs/src/table.rs b/nodejs/src/table.rs index ebbc4a060..a008c606b 100644 --- a/nodejs/src/table.rs +++ b/nodejs/src/table.rs @@ -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, } +/// 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, +} + +fn resource_budget_integer(value: f64, name: &str) -> napi::Result { + 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) -> napi::Result { + 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, /// GPU class for Function workers. pub gpu_class: Option, + /// Complete aggregate CPU and memory budget for a registered Function + /// refresh. `gpuCount` inside the budget defaults to zero. + pub resource_budget: Option, /// Positive aggregate GPU budget for Function workers. pub gpu_count: Option, /// `"fail"` (default) stops on a row error; `"skip"` writes null and @@ -1628,6 +1667,18 @@ impl TryFrom 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, diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index aec0f7fb1..81615ae73 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -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, diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index a963ab55e..f389f313f 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -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, diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 8f9a78e99..59a7384fd 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -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 diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 3a1d7f5db..109dbfbdb 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -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", diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 157870bb5..2422437c2 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -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"), ], diff --git a/python/src/table.rs b/python/src/table.rs index 23c5b69d7..56fc3bf5c 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -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, batch_size: Option, publication_interval_rows: Option, + resource_budget: Option, gpu_class: Option, gpu_count: Option, retry_policy: Option, @@ -70,11 +73,71 @@ fn positive_count>(value: Option, name: &str) -> PyResult, + name: &str, + allow_zero: bool, +) -> PyResult { + if value.is_none() || value.is_instance_of::() { + return Err(PyValueError::new_err(format!( + "resource_budget.{name} must be an integer" + ))); + } + let parsed = value + .extract::() + .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> { + let Some(budget) = budget else { + return Ok(None); + }; + for key in budget.keys() { + let key = key + .extract::() + .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 { 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, + pub fn refresh_column<'py>( + self_: PyRef<'py, Self>, column: String, r#where: Option, batch_size: Option, publication_interval_rows: Option, + resource_budget: Option>, gpu_class: Option, gpu_count: Option, retry_policy: Option, @@ -2228,11 +2293,12 @@ impl Table { init_timeout_seconds: Option, num_workers: Option, max_executor_processes_per_worker: Option, - ) -> PyResult> { + ) -> PyResult> { 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, batch_size: Option, publication_interval_rows: Option, + resource_budget: Option>, gpu_class: Option, gpu_count: Option, retry_policy: Option, @@ -2272,11 +2339,12 @@ impl Table { init_timeout_seconds: Option, num_workers: Option, max_executor_processes_per_worker: Option, - ) -> PyResult> { + ) -> PyResult> { 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, diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index a69611306..4f428809c 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -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, + #[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 { 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) diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 80b6bd7ad..bd9220754 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -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}; diff --git a/rust/lancedb/src/table/refresh.rs b/rust/lancedb/src/table/refresh.rs index a8be34b17..742d60f18 100644 --- a/rust/lancedb/src/table/refresh.rs +++ b/rust/lancedb/src/table/refresh.rs @@ -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, + /// 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, /// 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) -> 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)?; }