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:
Xuanwo
2026-11-02 08:50:36 +08:00
committed by GitHub
parent 6a255926d7
commit 63aa06ef7d
24 changed files with 921 additions and 41 deletions
+65
View File
@@ -24,6 +24,7 @@ use lancedb::{
connection::NamespaceClientPushdownOperation,
database::namespace::LanceNamespaceDatabase,
database::{CreateTableMode, Database, ReadConsistency},
function::GpuClass as LanceGpuClass,
graph::PropertyGraphDefinition,
};
use pyo3::{
@@ -38,6 +39,62 @@ pub struct Connection {
inner: Option<LanceConnection>,
}
/// Configured hardware and worker limits of a Function GPU class.
#[pyclass(module = "lancedb._lancedb", get_all, from_py_object)]
#[derive(Clone, Debug)]
pub struct GpuClass {
pub name: String,
pub resource_name: String,
pub gpu_model: Option<String>,
pub gpu_memory_bytes: Option<u64>,
pub max_gpu_count_per_worker: Option<u32>,
pub cpu_millis: Option<u64>,
pub memory_bytes: Option<u64>,
pub ephemeral_storage_bytes: Option<u64>,
}
impl From<LanceGpuClass> for GpuClass {
fn from(class: LanceGpuClass) -> Self {
Self {
name: class.name,
resource_name: class.resource_name,
gpu_model: class.gpu_model,
gpu_memory_bytes: class.gpu_memory_bytes,
max_gpu_count_per_worker: class.max_gpu_count_per_worker,
cpu_millis: class.cpu_millis,
memory_bytes: class.memory_bytes,
ephemeral_storage_bytes: class.ephemeral_storage_bytes,
}
}
}
#[pymethods]
impl GpuClass {
fn __repr__(&self) -> String {
let gpu_model = self
.gpu_model
.as_ref()
.map_or_else(|| "None".to_string(), |value| format!("{value:?}"));
let optional = |value: Option<u64>| {
value.map_or_else(|| "None".to_string(), |value| value.to_string())
};
let max_gpu_count = self
.max_gpu_count_per_worker
.map_or_else(|| "None".to_string(), |value| value.to_string());
format!(
"GpuClass(name={:?}, resource_name={:?}, gpu_model={gpu_model}, \
gpu_memory_bytes={}, max_gpu_count_per_worker={max_gpu_count}, \
cpu_millis={}, memory_bytes={}, ephemeral_storage_bytes={})",
self.name,
self.resource_name,
optional(self.gpu_memory_bytes),
optional(self.cpu_millis),
optional(self.memory_bytes),
optional(self.ephemeral_storage_bytes),
)
}
}
impl Connection {
pub(crate) fn new(inner: LanceConnection) -> Self {
Self { inner: Some(inner) }
@@ -1202,6 +1259,14 @@ impl Connection {
))
}
pub fn list_gpu_classes(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
let classes = inner.list_gpu_classes().await.infer_error()?;
Ok(classes.into_iter().map(GpuClass::from).collect::<Vec<_>>())
})
}
pub fn cancel_job(self_: PyRef<'_, Self>, job_id: String) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
+1
View File
@@ -65,6 +65,7 @@ pub fn _lancedb(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
.write_style("LANCEDB_LOG_STYLE");
env_logger::init_from_env(env);
m.add_class::<Connection>()?;
m.add_class::<connection::GpuClass>()?;
m.add_class::<catalog::Catalog>()?;
m.add_class::<catalog::DatabaseNames>()?;
m.add_class::<listing::NameListing>()?;
+27 -5
View File
@@ -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