mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-13 14:16:54 +00:00
feat: expose Function refresh resource budgets in SDKs (#4951)
## 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
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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"),
|
||||
],
|
||||
|
||||
+79
-11
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user