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
+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