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.

python/python/lancedb/__init__.py

@@ -19,6 +19,7 @@ from ._lancedb import FtsToken
 from ._lancedb import FunctionErrorFragment as FunctionErrorFragment
 from ._lancedb import FunctionErrorRecord as FunctionErrorRecord
 from ._lancedb import FunctionErrors as FunctionErrors
+from ._lancedb import GpuClass as GpuClass
 from ._lancedb import RTEWriteSpec
 from ._lancedb import tokenize as _tokenize
 from .common import URI, sanitize_uri

python/python/lancedb/_lancedb.pyi

@@ -112,6 +112,16 @@ class Session:
     @property
     def approx_num_items(self) -> int: ...
 
+class GpuClass:
+    name: str
+    resource_name: str
+    gpu_model: Optional[str]
+    gpu_memory_bytes: Optional[int]
+    max_gpu_count_per_worker: Optional[int]
+    cpu_millis: Optional[int]
+    memory_bytes: Optional[int]
+    ephemeral_storage_bytes: Optional[int]
+
 class Connection(object):
     uri: str
     async def is_open(self): ...
@@ -153,6 +163,7 @@ class Connection(object):
         limit: Optional[int],
     ) -> list[str]: ...  # Deprecated: Use list_tables instead
     async def open_job(self, job_id: str) -> Job: ...
+    async def list_gpu_classes(self) -> list[GpuClass]: ...
     async def create_function_async(
         self, request_json: str, namespace_path: Optional[List[str]] = None
     ) -> Job: ...
@@ -561,10 +572,20 @@ class Table:
         auto_refresh_interval_seconds: Optional[int] = None,
     ) -> AddColumnsResult: ...
     async def refresh_column(
-        self, column: str, *, batch_size: Optional[int] = None
+        self,
+        column: str,
+        *,
+        batch_size: Optional[int] = None,
+        gpu_class: Optional[str] = None,
+        gpu_count: Optional[int] = None,
     ) -> RefreshColumnResult: ...
     async def refresh_column_async(
-        self, column: str, *, batch_size: Optional[int] = None
+        self,
+        column: str,
+        *,
+        batch_size: Optional[int] = None,
+        gpu_class: Optional[str] = None,
+        gpu_count: Optional[int] = None,
     ) -> Job: ...
     async def function_errors(
         self,

python/python/lancedb/db.py

@@ -104,6 +104,7 @@ if TYPE_CHECKING:
     from .pydantic import LanceModel
 
     from ._lancedb import Connection as LanceDbConnection
+    from ._lancedb import GpuClass
     from ._lancedb import JobInfo
     from .common import DATA, URI
     from .embeddings import EmbeddingFunctionConfig
@@ -1348,6 +1349,21 @@ class DBConnection(EnforceOverrides):
         """
         raise NotImplementedError("open_job is not supported for this connection type")
 
+    def list_gpu_classes(self) -> List[GpuClass]:
+        """List the operator-configured Function GPU classes.
+
+        Each record names the class and reports its configured GPU model,
+        memory, per-worker GPU limit, and CPU, memory, and ephemeral-storage
+        limits. This describes configured hardware, not current free capacity.
+
+        LanceDB Distributed only; the principal needs System Usage permission.
+        Raises ``NotImplementedError`` when the server predates
+        ``GET /v1/gpu/classes``.
+        """
+        raise NotImplementedError(
+            "GPU class discovery is not supported for this connection type"
+        )
+
     def list_jobs(
         self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
     ) -> Listing[JobInfo]:
@@ -3450,6 +3466,12 @@ class AsyncConnection(object):
         """
         return AsyncJob(await self._inner.open_job(job_id))
 
+    async def list_gpu_classes(self) -> List[GpuClass]:
+        """List configured Function GPU classes. See
+        [DBConnection.list_gpu_classes][lancedb.db.DBConnection.list_gpu_classes].
+        """
+        return await self._inner.list_gpu_classes()
+
     async def create_function_async(
         self,
         definition: UdfDefinition,

python/python/lancedb/functions.py

@@ -61,6 +61,11 @@ from .secrets import EnvVarSecret
 _Int32 = conint(strict=True, ge=-(2**31), le=2**31 - 1)
 _UInt32 = conint(strict=True, ge=0, le=2**32 - 1)
 _UInt64 = conint(strict=True, ge=0, le=2**64 - 1)
+_GPU_CLASS_NAME = re.compile(r"^[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?$")
+_GPU_CLASS_ERROR = (
+    "gpu_class must be a lowercase name of at most 63 letters, digits, or "
+    "hyphens, and cannot start or end with a hyphen"
+)
 
 
 def _validate_gpu_wire_marker(value: Any) -> bool:
@@ -75,6 +80,12 @@ def _normalize_gpu_marker(value: bool) -> Optional[bool]:
     return True if value else None
 
 
+def _validate_gpu_class_name(value: Any) -> str:
+    if not isinstance(value, str) or not _GPU_CLASS_NAME.fullmatch(value):
+        raise ValueError(_GPU_CLASS_ERROR)
+    return value
+
+
 class _FrozenDict(dict):
     def _immutable(self, *args, **kwargs):
         raise TypeError("remote canonical values are immutable")
@@ -291,6 +302,8 @@ class PythonRuntimeSpec(_RemoteValue):
 
     V1 supports ``kind="python"``. Newer runtime kinds remain readable, while
     their unknown payload fields are intentionally not retained by the client.
+    ``kind="python_v2"`` requires ``gpu=True`` and may name an operator-defined
+    ``gpu_class``.
     """
 
     kind: str
@@ -298,6 +311,7 @@ class PythonRuntimeSpec(_RemoteValue):
     environment: Optional[PythonEnvironmentSpec] = None
     env: Optional[Mapping[str, str]] = None
     gpu: Optional[bool] = None
+    gpu_class: Optional[str] = None
 
     @model_validator(mode="before")
     @classmethod
@@ -315,6 +329,13 @@ class PythonRuntimeSpec(_RemoteValue):
             return None
         return _validate_gpu_wire_marker(value)
 
+    @field_validator("gpu_class", mode="before")
+    @classmethod
+    def _validate_gpu_class(cls, value):
+        if value is None:
+            return None
+        return _validate_gpu_class_name(value)
+
     @model_validator(mode="after")
     def _validate_runtime_kind(self):
         if self.kind == "python":
@@ -324,6 +345,10 @@ class PythonRuntimeSpec(_RemoteValue):
                 raise ValueError("python runtime requires environment")
             if self.gpu is not None:
                 raise ValueError("python runtime with gpu requires kind='python_v2'")
+            if self.gpu_class is not None:
+                raise ValueError(
+                    "python runtime with gpu_class requires kind='python_v2'"
+                )
         elif self.kind == "python_v2":
             if self.python_version is None:
                 raise ValueError("python_v2 runtime requires python_version")
@@ -336,6 +361,7 @@ class PythonRuntimeSpec(_RemoteValue):
             object.__setattr__(self, "environment", None)
             object.__setattr__(self, "env", None)
             object.__setattr__(self, "gpu", None)
+            object.__setattr__(self, "gpu_class", None)
         return self
 
 
@@ -373,6 +399,9 @@ class FunctionVersion(_RemoteValue):
     ``udf_def.registration_request.artifact.digest`` to check whether a local
     [UdfDefinition][lancedb.functions.UdfDefinition] has the same source.
     It is ``None`` when the server does not provide a source digest.
+
+    ``gpu_class`` is the operator-defined GPU pool fixed into this Function
+    version, or ``None`` when registration used the deployment default.
     """
 
     name: str
@@ -387,6 +416,7 @@ class FunctionVersion(_RemoteValue):
     created_at_millis: int
     metadata: Mapping[str, str]
     disabled: bool
+    gpu_class: Optional[str] = None
 
     def __call__(self, **arguments: Any) -> FunctionApplication:
         """Bind this exact version to named table columns.
@@ -1900,6 +1930,7 @@ class UdfDefinition:
         env: Mapping[str, str],
         python_version: Optional[str],
         gpu: bool = False,
+        gpu_class: Optional[str] = None,
         conda: tuple[str, ...] = (),
         conda_channels: tuple[str, ...] = (),
         apt: tuple[str, ...] = (),
@@ -1918,6 +1949,11 @@ class UdfDefinition:
             raise ValueError("a Function environment is pip or conda, not both")
         if conda_channels and not conda:
             raise ValueError("conda_channels requires conda packages")
+        gpu_marker = _normalize_gpu_marker(gpu)
+        if gpu_class is not None:
+            _validate_gpu_class_name(gpu_class)
+            if gpu_marker is None:
+                raise ValueError("gpu_class requires gpu=True")
         packages = tuple(sorted(set(conda if conda else pip)))
         system_packages = tuple(sorted(set(apt)))
         if any(
@@ -1988,7 +2024,6 @@ class UdfDefinition:
             kind = _PYTHON_CALLABLE_ARTIFACT
             source = entry_source.encode("utf-8")
         digest = f"sha256:{hashlib.sha256(source).hexdigest()}"
-        gpu_marker = _normalize_gpu_marker(gpu)
         runtime = PythonRuntimeSpec(
             kind="python_v2" if gpu_marker is not None else "python",
             python_version=python_version
@@ -1996,6 +2031,7 @@ class UdfDefinition:
             environment=environment_spec,
             env=environment,
             gpu=gpu_marker,
+            gpu_class=gpu_class,
         )
         self._function = function
         self._request = FunctionRegistrationRequest(
@@ -2110,6 +2146,7 @@ def udf(
     env: Optional[Mapping[str, str]] = None,
     python_version: Optional[str] = None,
     gpu: bool = False,
+    gpu_class: Optional[str] = None,
     conda: tuple[str, ...] | list[str] = (),
     conda_channels: tuple[str, ...] | list[str] = (),
     apt: tuple[str, ...] | list[str] = (),
@@ -2128,6 +2165,7 @@ def udf(
     env: Optional[Mapping[str, str]] = None,
     python_version: Optional[str] = None,
     gpu: bool = False,
+    gpu_class: Optional[str] = None,
     conda: tuple[str, ...] | list[str] = (),
     conda_channels: tuple[str, ...] | list[str] = (),
     apt: tuple[str, ...] | list[str] = (),
@@ -2251,6 +2289,12 @@ def udf(
         Whether every remote execution requires a GPU. The execution platform
         selects one compatible GPU for each worker. The requirement is part of
         the immutable Function version.
+    gpu_class : str, optional
+        Operator-defined GPU class for every execution of this Function, such
+        as ``"a10g"`` or ``"h100"``. Requires ``gpu=True`` and follows the
+        lowercase 63-character DNS-label rule. The class is part of the
+        immutable Function version; per-refresh aggregate GPU budgets are
+        selected through ``refresh_column`` instead of registration.
     code : sequence of module, optional
         Top-level modules or packages whose importable ``.py`` files ship
         with the Function and are importable by their own names.
@@ -2280,11 +2324,13 @@ def udf(
     ...     return value
     >>> batch_identity(pa.array([1, 2, 3])).to_pylist()
     [1, 2, 3]
-    >>> @udf(pip=["cupy-cuda12x"], gpu=True)
+    >>> @udf(pip=["cupy-cuda12x"], gpu=True, gpu_class="a10g")
     ... def gpu_score(value: int) -> int:
     ...     return value * 2
     >>> gpu_score.registration_request.runtime.gpu
     True
+    >>> gpu_score.registration_request.runtime.gpu_class
+    'a10g'
     >>> @udf
     ... class scale:
     ...     def __init__(self, factor: float, offset: float = 0.0):
@@ -2319,6 +2365,7 @@ def udf(
             env={} if env is None else env,
             python_version=python_version,
             gpu=gpu,
+            gpu_class=gpu_class,
             conda=tuple(conda),
             conda_channels=tuple(conda_channels),
             apt=tuple(apt),

python/python/lancedb/remote/db.py

@@ -49,7 +49,7 @@ from ..storage_integration import SecretReference, StorageIntegrationInfo
 from ..view import ViewDescription
 
 if TYPE_CHECKING:
-    from .._lancedb import JobInfo, Session
+    from .._lancedb import GpuClass, JobInfo, Session
 from ..embeddings import EmbeddingFunctionConfig
 from lance_namespace import (
     LanceNamespace,
@@ -958,6 +958,13 @@ class RemoteDBConnection(DBConnection):
         """
         return Job(LOOP.run(self._conn.open_job(job_id)))
 
+    @override
+    def list_gpu_classes(self) -> List[GpuClass]:
+        """List configured Function GPU classes. See
+        [DBConnection.list_gpu_classes][lancedb.db.DBConnection.list_gpu_classes].
+        """
+        return LOOP.run(self._conn.list_gpu_classes())
+
     @override
     def create_function_async(
         self,

python/python/lancedb/remote/table.py

@@ -1108,14 +1108,40 @@ class RemoteTable(Table):
             )
         )
 
-    def refresh_column(self, column: str, *, batch_size: Optional[int] = None):
-        return LOOP.run(self._table.refresh_column(column, batch_size=batch_size))
+    def refresh_column(
+        self,
+        column: str,
+        *,
+        batch_size: Optional[int] = None,
+        gpu_class: Optional[str] = None,
+        gpu_count: Optional[int] = None,
+    ):
+        return LOOP.run(
+            self._table.refresh_column(
+                column,
+                batch_size=batch_size,
+                gpu_class=gpu_class,
+                gpu_count=gpu_count,
+            )
+        )
 
     def refresh_column_async(
-        self, column: str, *, batch_size: Optional[int] = None
+        self,
+        column: str,
+        *,
+        batch_size: Optional[int] = None,
+        gpu_class: Optional[str] = None,
+        gpu_count: Optional[int] = None,
     ) -> Job[RefreshColumnResult]:
         return Job(
-            LOOP.run(self._table.refresh_column_async(column, batch_size=batch_size))
+            LOOP.run(
+                self._table.refresh_column_async(
+                    column,
+                    batch_size=batch_size,
+                    gpu_class=gpu_class,
+                    gpu_count=gpu_count,
+                )
+            )
         )
 
     def function_errors(

python/python/lancedb/table.py

@@ -3078,7 +3078,12 @@ class Table(ABC):
 
     @abstractmethod
     def refresh_column(
-        self, column: str, *, batch_size: Optional[int] = None
+        self,
+        column: str,
+        *,
+        batch_size: Optional[int] = None,
+        gpu_class: Optional[str] = None,
+        gpu_count: Optional[int] = None,
     ) -> "RefreshColumnResult":
         """
         Fill the rows of a computed column that hold no value yet.
@@ -3103,6 +3108,15 @@ class Table(ABC):
             Maximum rows per Function call on the remote server. Must be
             greater than zero. None uses the server default. Has no effect on
             local expression-backed columns.
+        gpu_class: str, optional
+            GPU class for Function workers. Must be a lowercase name of 1-63
+            letters, digits, or hyphens, without a leading or trailing hyphen.
+            None keeps the Function's registered class or the deployment
+            default. LanceDB Distributed only.
+        gpu_count: int, optional
+            Positive aggregate GPU budget for Function workers. None leaves
+            the server's configured budget unchanged. LanceDB Distributed
+            only.
 
         Returns
         -------
@@ -3113,7 +3127,12 @@ class Table(ABC):
 
     @abstractmethod
     def refresh_column_async(
-        self, column: str, *, batch_size: Optional[int] = None
+        self,
+        column: str,
+        *,
+        batch_size: Optional[int] = None,
+        gpu_class: Optional[str] = None,
+        gpu_count: Optional[int] = None,
     ) -> Job[RefreshColumnJobResult]:
         """
         Like :meth:`refresh_column`, but returns a handle to the refresh job
@@ -3137,6 +3156,15 @@ class Table(ABC):
             greater than zero. None uses the server default. Lower this for
             slow Functions to stay within the executor's per-call deadline.
             Has no effect on local expression-backed columns.
+        gpu_class: str, optional
+            GPU class for Function workers. Must be a lowercase name of 1-63
+            letters, digits, or hyphens, without a leading or trailing hyphen.
+            None keeps the Function's registered class or the deployment
+            default. LanceDB Distributed only.
+        gpu_count: int, optional
+            Positive aggregate GPU budget for Function workers. None leaves
+            the server's configured budget unchanged. LanceDB Distributed
+            only.
 
         Returns
         -------
@@ -5200,22 +5228,46 @@ class LanceTable(Table):
         )
 
     def refresh_column(
-        self, column: str, *, batch_size: Optional[int] = None
+        self,
+        column: str,
+        *,
+        batch_size: Optional[int] = None,
+        gpu_class: Optional[str] = None,
+        gpu_count: Optional[int] = None,
     ) -> "RefreshColumnResult":
         """Fill a computed column's unfilled rows and recompute those whose
         inputs changed. See
         [`AsyncTable.refresh_column`][lancedb.AsyncTable.refresh_column]."""
-        return LOOP.run(self._table.refresh_column(column, batch_size=batch_size))
+        return LOOP.run(
+            self._table.refresh_column(
+                column,
+                batch_size=batch_size,
+                gpu_class=gpu_class,
+                gpu_count=gpu_count,
+            )
+        )
 
     def refresh_column_async(
-        self, column: str, *, batch_size: Optional[int] = None
+        self,
+        column: str,
+        *,
+        batch_size: Optional[int] = None,
+        gpu_class: Optional[str] = None,
+        gpu_count: Optional[int] = None,
     ) -> Job[RefreshColumnJobResult]:
         """Fill a computed column's unfilled rows and recompute those whose
         inputs changed, returning a handle to the refresh job. See
         [`Table.refresh_column_async`][lancedb.table.Table.refresh_column_async].
         """
         return Job(
-            LOOP.run(self._table.refresh_column_async(column, batch_size=batch_size))
+            LOOP.run(
+                self._table.refresh_column_async(
+                    column,
+                    batch_size=batch_size,
+                    gpu_class=gpu_class,
+                    gpu_count=gpu_count,
+                )
+            )
         )
 
     def function_errors(
@@ -7475,7 +7527,12 @@ class AsyncTable:
             return await self._inner.add_columns(list(transforms.items()))
 
     async def refresh_column(
-        self, column: str, *, batch_size: Optional[int] = None
+        self,
+        column: str,
+        *,
+        batch_size: Optional[int] = None,
+        gpu_class: Optional[str] = None,
+        gpu_count: Optional[int] = None,
     ) -> RefreshColumnResult:
         """
         Fill the rows of a computed column that hold no value yet.
@@ -7500,13 +7557,27 @@ class AsyncTable:
             Maximum rows per Function call on the remote server. Must be
             greater than zero. None uses the server default. Has no effect on
             local expression-backed columns.
+        gpu_class: str, optional
+            GPU class for Function workers. Must be a lowercase name of 1-63
+            letters, digits, or hyphens, without a leading or trailing hyphen.
+            None keeps the Function's registered class or the deployment
+            default. LanceDB Distributed only.
+        gpu_count: int, optional
+            Positive aggregate GPU budget for Function workers. None leaves
+            the server's configured budget unchanged. LanceDB Distributed
+            only.
 
         Returns
         -------
         RefreshColumnResult
             The number of rows filled and the new version of the table.
         """
-        return await self._inner.refresh_column(column, batch_size=batch_size)
+        return await self._inner.refresh_column(
+            column,
+            batch_size=batch_size,
+            gpu_class=gpu_class,
+            gpu_count=gpu_count,
+        )
 
     async def function_errors(
         self,
@@ -7576,7 +7647,12 @@ class AsyncTable:
         )
 
     async def refresh_column_async(
-        self, column: str, *, batch_size: Optional[int] = None
+        self,
+        column: str,
+        *,
+        batch_size: Optional[int] = None,
+        gpu_class: Optional[str] = None,
+        gpu_count: Optional[int] = None,
     ) -> AsyncJob[RefreshColumnJobResult]:
         """
         Like :meth:`refresh_column`, but returns a handle to the refresh job
@@ -7600,6 +7676,15 @@ class AsyncTable:
             greater than zero. None uses the server default. Lower this for
             slow Functions to stay within the executor's per-call deadline.
             Has no effect on local expression-backed columns.
+        gpu_class: str, optional
+            GPU class for Function workers. Must be a lowercase name of 1-63
+            letters, digits, or hyphens, without a leading or trailing hyphen.
+            None keeps the Function's registered class or the deployment
+            default. LanceDB Distributed only.
+        gpu_count: int, optional
+            Positive aggregate GPU budget for Function workers. None leaves
+            the server's configured budget unchanged. LanceDB Distributed
+            only.
 
         Returns
         -------
@@ -7623,7 +7708,12 @@ class AsyncTable:
         'finished'
         """
         return _typed_job(
-            await self._inner.refresh_column_async(column, batch_size=batch_size),
+            await self._inner.refresh_column_async(
+                column,
+                batch_size=batch_size,
+                gpu_class=gpu_class,
+                gpu_count=gpu_count,
+            ),
             RefreshColumnJobResult.from_json,
         )
 

python/python/tests/test_first_class_function_slice1.py

@@ -167,6 +167,17 @@ def test_function_version_preserves_source_digest(source_digest):
         version.source_digest = "sha256:changed"
 
 
+def test_function_version_exposes_registered_gpu_class():
+    value = job_result("remote_function_job.json")
+    legacy = FunctionVersion.from_json(json.dumps(value))
+    assert legacy.gpu_class is None
+
+    value["gpu_class"] = "a10g"
+    version = FunctionVersion.from_json(json.dumps(value))
+    assert version.gpu_class == "a10g"
+    assert json.loads(version.to_canonical_json())["gpu_class"] == "a10g"
+
+
 def test_function_version_binds_named_columns_as_one_immutable_application():
     version = FunctionVersion.from_json(
         json.dumps(job_result("remote_function_job.json"))

python/python/tests/test_first_class_function_slice2.py

@@ -534,13 +534,14 @@ def test_udf_image_requires_digest_reference(image):
 
 
 def test_udf_gpu_marker_uses_gpu_runtime():
-    @udf(pip=["cupy-cuda12x"], gpu=True)
+    @udf(pip=["cupy-cuda12x"], gpu=True, gpu_class="a10g")
     def double_on_gpu(value: int) -> int:
         return value * 2
 
     request = json.loads(double_on_gpu.registration_request.to_canonical_json())
     assert request["runtime"]["kind"] == "python_v2"
     assert request["runtime"]["gpu"] is True
+    assert request["runtime"]["gpu_class"] == "a10g"
 
     @udf(pip=["pyarrow"])
     def cpu_function(value: int) -> int:
@@ -551,6 +552,7 @@ def test_udf_gpu_marker_uses_gpu_runtime():
     ]
     assert cpu_runtime["kind"] == "python"
     assert "gpu" not in cpu_runtime
+    assert "gpu_class" not in cpu_runtime
 
     def identity(value: int) -> int:
         return value
@@ -558,6 +560,11 @@ def test_udf_gpu_marker_uses_gpu_runtime():
     for invalid in [None, 0, 1, -1, 1.5, "", "true", "1", "H100"]:
         with pytest.raises(ValueError, match="gpu must be a boolean"):
             udf(name="invalid_gpu", gpu=invalid)(identity)
+    for invalid in ["H100", "a_b", "-a10g", "a10g-", "a" * 64]:
+        with pytest.raises(ValueError, match="gpu_class"):
+            udf(name="invalid_gpu_class", gpu=True, gpu_class=invalid)(identity)
+    with pytest.raises(ValueError, match="requires gpu=True"):
+        udf(name="gpu_class_without_gpu", gpu_class="a10g")(identity)
 
     base_runtime = {
         "kind": "python_v2",
@@ -566,14 +573,26 @@ def test_udf_gpu_marker_uses_gpu_runtime():
     }
     runtime = PythonRuntimeSpec.model_validate({**base_runtime, "gpu": True})
     assert runtime.gpu is True
+    runtime = PythonRuntimeSpec.model_validate(
+        {**base_runtime, "gpu": True, "gpu_class": "a10g"}
+    )
+    assert runtime.gpu_class == "a10g"
     for invalid in [False, 1, 0, "", "true", "1", "H100"]:
         with pytest.raises(ValueError, match="runtime.gpu must be true"):
             PythonRuntimeSpec.model_validate({**base_runtime, "gpu": invalid})
+    with pytest.raises(ValueError, match="gpu_class"):
+        PythonRuntimeSpec.model_validate(
+            {**base_runtime, "gpu": True, "gpu_class": "H100"}
+        )
 
 
 def test_unknown_runtime_discards_payload_before_known_field_validation():
     for payload in [
-        {"kind": "python_v3", "gpu": {"model": "H100"}},
+        {
+            "kind": "python_v3",
+            "gpu": {"model": "H100"},
+            "gpu_class": "a10g",
+        },
         {"kind": "python_v3", "resources": []},
         {
             "kind": "python_v3",

python/python/tests/test_package_metadata.py

@@ -31,6 +31,15 @@ def test_job_metadata_types_have_resolvable_modules(name):
     assert getattr(defining_module, public_type.__name__, None) is public_type
 
 
+def test_gpu_class_reexport_has_resolvable_module():
+    import lancedb
+
+    public_type = lancedb.GpuClass
+    defining_module = importlib.import_module(public_type.__module__)
+
+    assert getattr(defining_module, public_type.__name__, None) is public_type
+
+
 def test_pyo3_abi_matches_minimum_supported_python():
     project_dir = Path(__file__).parents[2]
     pyproject = (project_dir / "pyproject.toml").read_text()

python/python/tests/test_remote_db.py

@@ -1700,7 +1700,14 @@ def test_remote_create_index_async_returns_job():
 @pytest.mark.asyncio
 @pytest.mark.parametrize("async_table", [False, True])
 @pytest.mark.parametrize(
-    "refresh_kwargs", [{}, {"batch_size": None}, {"batch_size": 1}, {"batch_size": 8}]
+    "refresh_kwargs",
+    [
+        {},
+        {"batch_size": None},
+        {"batch_size": 1},
+        {"batch_size": 8},
+        {"gpu_class": "a10g", "gpu_count": 2},
+    ],
 )
 async def test_remote_refresh_async_returns_typed_terminal_result(
     async_table, refresh_kwargs
@@ -1720,6 +1727,10 @@ async def test_remote_refresh_async_returns_typed_terminal_result(
             expected = {"column": "derived"}
             if refresh_kwargs.get("batch_size") is not None:
                 expected["batch_size"] = refresh_kwargs["batch_size"]
+            if refresh_kwargs.get("gpu_class") is not None:
+                expected["gpu_class"] = refresh_kwargs["gpu_class"]
+            if refresh_kwargs.get("gpu_count") is not None:
+                expected["gpu_count"] = refresh_kwargs["gpu_count"]
             assert json.loads(body) == expected
             request.send_response(202)
             request.send_header("Content-Type", "application/json")
@@ -1788,6 +1799,66 @@ async def test_remote_refresh_async_returns_typed_terminal_result(
     assert result.version == 8
 
 
+@pytest.mark.parametrize(
+    "refresh_kwargs,message",
+    [
+        ({"gpu_class": "H100"}, "gpu_class"),
+        ({"gpu_class": "-a10g"}, "gpu_class"),
+        ({"gpu_class": "a_b"}, "gpu_class"),
+        ({"gpu_count": 0}, "gpu_count"),
+        ({"gpu_count": -1}, "gpu_count"),
+    ],
+)
+@pytest.mark.asyncio
+async def test_remote_refresh_rejects_invalid_gpu_options_before_request(
+    refresh_kwargs, message
+):
+    backfill_requests = []
+
+    def handler(request):
+        if request.path == "/v1/table/test/describe/":
+            request.send_response(200)
+            request.send_header("Content-Type", "application/json")
+            request.end_headers()
+            request.wfile.write(
+                json.dumps(
+                    {
+                        "version": 1,
+                        "schema": {
+                            "fields": [
+                                {
+                                    "name": "id",
+                                    "type": {"type": "int64"},
+                                    "nullable": False,
+                                }
+                            ]
+                        },
+                    }
+                ).encode()
+            )
+        elif request.path == "/v1/table/test/backfill_column":
+            backfill_requests.append(request.path)
+            request.send_response(202)
+            request.send_header("Content-Type", "application/json")
+            request.end_headers()
+            request.wfile.write(b'{"job_id": "unexpected"}')
+        else:
+            request.send_response(404)
+            request.end_headers()
+
+    with mock_lancedb_connection(handler) as db:
+        table = db.open_table("test")
+        with pytest.raises(ValueError, match=message):
+            table.refresh_column_async("derived", **refresh_kwargs)
+
+    async with mock_lancedb_connection_async(handler) as db:
+        table = await db.open_table("test")
+        with pytest.raises(ValueError, match=message):
+            await table.refresh_column_async("derived", **refresh_kwargs)
+
+    assert backfill_requests == []
+
+
 def test_remote_function_errors_lists_the_rows_a_refresh_skipped():
     listing = {
         "records": [
@@ -4914,6 +4985,83 @@ def test_remote_job_handle_reports_its_own_detail():
         }
 
 
+@pytest.mark.parametrize("is_async", [False, True])
+@pytest.mark.asyncio
+async def test_remote_list_gpu_classes(is_async):
+    requests = []
+
+    def handler(request):
+        requests.append((request.command, request.path))
+        request.send_response(200)
+        request.send_header("Content-Type", "application/json")
+        request.end_headers()
+        request.wfile.write(
+            json.dumps(
+                {
+                    "gpu_classes": [
+                        {
+                            "name": "default",
+                            "resource_name": "nvidia.com/gpu",
+                            "gpu_model": "NVIDIA A10G",
+                            "gpu_memory_bytes": 25769803776,
+                            "max_gpu_count_per_worker": 1,
+                            "cpu_millis": 4000,
+                            "memory_bytes": 25769803776,
+                            "ephemeral_storage_bytes": 1073741824,
+                        },
+                        {
+                            "name": "h100",
+                            "resource_name": "nvidia.com/gpu",
+                            "gpu_model": "NVIDIA H100",
+                            "gpu_memory_bytes": 85899345920,
+                            "max_gpu_count_per_worker": 8,
+                            "cpu_millis": None,
+                            "memory_bytes": None,
+                            "ephemeral_storage_bytes": None,
+                        },
+                    ]
+                }
+            ).encode()
+        )
+
+    if is_async:
+        async with mock_lancedb_connection_async(handler) as db:
+            classes = await db.list_gpu_classes()
+    else:
+        with mock_lancedb_connection(handler) as db:
+            classes = db.list_gpu_classes()
+
+    assert requests == [("GET", "/v1/gpu/classes")]
+    assert len(classes) == 2
+    assert isinstance(classes[0], lancedb.GpuClass)
+    assert classes[0].name == "default"
+    assert classes[0].gpu_model == "NVIDIA A10G"
+    assert classes[0].gpu_memory_bytes == 25769803776
+    assert classes[1].max_gpu_count_per_worker == 8
+    assert classes[1].cpu_millis is None
+
+
+@pytest.mark.parametrize("is_async", [False, True])
+@pytest.mark.asyncio
+async def test_remote_list_gpu_classes_older_server(is_async):
+    requests = []
+
+    def handler(request):
+        requests.append(request.path)
+        request.send_response(404)
+        request.end_headers()
+
+    with pytest.raises(NotImplementedError, match="/v1/gpu/classes"):
+        if is_async:
+            async with mock_lancedb_connection_async(handler) as db:
+                await db.list_gpu_classes()
+        else:
+            with mock_lancedb_connection(handler) as db:
+                db.list_gpu_classes()
+
+    assert requests == ["/v1/gpu/classes"]
+
+
 def _property_graph_handler(calls, respond):
     """A mock route handler recording ``(method, path, body)`` and answering
     with the ``(status, body)`` that ``respond(path, body)`` returns, as JSON."""