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:
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user