mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-13 14:16:54 +00:00
feat(python): expose Function GPU classes and placement (#4791)
## Summary - Add sync/async `DBConnection.list_gpu_classes()` returning typed `GpuClass` records for `GET /v1/gpu/classes`. Records include the class name, Kubernetes resource, GPU model and memory, per-worker GPU maximum, and worker CPU/memory/ephemeral-storage limits. - Add `udf(gpu_class=...)` for immutable Function-level GPU placement. It requires `gpu=True`; Function versions expose their registered `gpu_class`. - Add `gpu_class` and `gpu_count` to `refresh_column` / `refresh_column_async` and Rust `RefreshColumnOptions`. Both are sent as top-level `backfill_column` fields and validated before any request. `gpu_count` is the aggregate GPU budget across Function workers. - Validate GPU class names with the server lowercase 63-character DNS-label rule. - Document the discovery-to-registration flow, configured hardware versus free capacity, System Usage permission, and per-refresh GPU placement in the API and web guides. Older servers without `/v1/gpu/classes` produce a clear `NotImplementedError` instead of a bare 404. Part of ENT-3470. The client follows the in-flight Sophon contract on `xuanwo/ent-3470-fix` for GPU class discovery and backfill placement options. ## Breaking change `PythonRuntimeSpec::PythonV2` adds `gpu_class`, and `RefreshColumnOptions` adds `gpu_class` and `gpu_count`. Rust callers constructing either value with all fields explicitly must supply the new fields. ## Validation At `7685c968`: Rust GPU tests (5 passed), Rust/PyO3 clippy, Python extension build, pyright, ruff, and Python Function/runtime, discovery, refresh placement, and refresh regressions (147 passed, 6 skipped). At `b22428fd`: Rust GPU tests (5 passed), rustfmt, ruff, pyright, Function slice 2 (85 passed), and remote GPU/refresh tests (20 passed). The final commit only updates documentation and Rust doc comments.
This commit is contained in:
+27
-5
@@ -42,12 +42,30 @@ use pyo3::{
|
||||
|
||||
mod scannable;
|
||||
|
||||
fn refresh_column_options(batch_size: Option<i64>) -> PyResult<RefreshColumnOptions> {
|
||||
fn refresh_column_options(
|
||||
batch_size: Option<i64>,
|
||||
gpu_class: Option<String>,
|
||||
gpu_count: Option<i64>,
|
||||
) -> PyResult<RefreshColumnOptions> {
|
||||
let mut options = RefreshColumnOptions::default();
|
||||
options.batch_size = batch_size
|
||||
.map(usize::try_from)
|
||||
.transpose()
|
||||
.map_err(|_| PyValueError::new_err("batch_size must be greater than zero"))?;
|
||||
options.gpu_class = gpu_class;
|
||||
options.gpu_count = gpu_count
|
||||
.map(|count| {
|
||||
u32::try_from(count)
|
||||
.ok()
|
||||
.filter(|count| *count > 0)
|
||||
.ok_or_else(|| {
|
||||
PyValueError::new_err(
|
||||
"gpu_count must be a positive integer no greater than u32::MAX",
|
||||
)
|
||||
})
|
||||
})
|
||||
.transpose()?;
|
||||
options.validate().infer_error()?;
|
||||
Ok(options)
|
||||
}
|
||||
|
||||
@@ -2080,13 +2098,15 @@ impl Table {
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (column, *, batch_size=None))]
|
||||
#[pyo3(signature = (column, *, batch_size=None, gpu_class=None, gpu_count=None))]
|
||||
pub fn refresh_column(
|
||||
self_: PyRef<'_, Self>,
|
||||
column: String,
|
||||
batch_size: Option<i64>,
|
||||
gpu_class: Option<String>,
|
||||
gpu_count: Option<i64>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let options = refresh_column_options(batch_size)?;
|
||||
let options = refresh_column_options(batch_size, gpu_class, gpu_count)?;
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let result = inner
|
||||
@@ -2097,13 +2117,15 @@ impl Table {
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (column, *, batch_size=None))]
|
||||
#[pyo3(signature = (column, *, batch_size=None, gpu_class=None, gpu_count=None))]
|
||||
pub fn refresh_column_async(
|
||||
self_: PyRef<'_, Self>,
|
||||
column: String,
|
||||
batch_size: Option<i64>,
|
||||
gpu_class: Option<String>,
|
||||
gpu_count: Option<i64>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let options = refresh_column_options(batch_size)?;
|
||||
let options = refresh_column_options(batch_size, gpu_class, gpu_count)?;
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let job = inner
|
||||
|
||||
Reference in New Issue
Block a user