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 {