API Reference

API reference for launching distributed GPU workloads and initializing Ray 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]

ray_init

serverless_gpu.ray_init(*args, **kwargs)

Initialize Ray with dashboard access and display the dashboard URL.

Parameters:

  • *args – Positional arguments passed to ray.init().

  • **kwargs – Keyword arguments passed to ray.init().

Returns:

The Ray context returned by ray.init(), with the dashboard URL configured for notebook access.

Raises:

ImportError – If Ray is not installed.

Example

from serverless_gpu import ray_init

ray_init()

# Use standard Ray APIs after initialization.
@ray.remote
def square(value):
    return value * value

ray.get(square.remote(4))