API Reference

API reference for launching distributed GPU workloads from Databricks notebooks.

distributed

serverless_gpu.distributed(gpus, gpu_type=None, timeout=10800)

Decorator to launch a function on GPUs.

Parameters:

  • gpus (int) – Number of GPUs to use.

  • gpu_type (Optional[Union[GPUType, str]], optional) – The GPU type to use. If provided, must match the accelerator the notebook is connected to ("A10", "H100"). Defaults to None (automatically detected).

  • timeout (Optional[float], optional) – Wall-clock seconds to allow local execution to run before subprocesses are terminated and LocalExecutionTimeoutError is raised. Defaults to 10800 seconds (3 hours). Pass None to disable.

Example

from serverless_gpu import distributed

@distributed(gpus=8, gpu_type="H100", timeout=3600)  # 60 minutes
def train(lr, epochs):
    # import statements and function code here
    ...

# Run the function in parallel by calling `.distributed()`
train.distributed(lr=0.01, epochs=10)

The decorator returns a DistributedFunction object. Call .distributed() on it, with the same arguments you would pass to your original function, to launch execution on GPU.

GPU Types

Note

gpu_type is auto-detected from the notebook’s attached GPU. If you do specify it, you can directly pass "H100" or "A10" as a string.

class serverless_gpu.compute.GPUType(value)

Bases: Enum

Available GPU types for distributed computing on Databricks serverless GPU compute.

H100

NVIDIA H100 80GB GPU instances.

A10

NVIDIA A10 GPU instances.

classmethod names()

Get all GPU type names as lowercase strings.

Returns:

List of canonical user-facing GPU type names.

Return type:

List[str]