GPUs

Run a function on a GPU

Name the GPU type in gpu=. The container starts on that GPU, runs the call and stops.

from klio_compute import function, Image

image = Image(python_version="python3.12", python_packages=["torch"])

@function(gpu="L40S", image=image)
def device_name() -> str:
    import torch
    return torch.cuda.get_device_name(0)

if __name__ == "__main__":
    print(device_name.remote())

Before a GPU workload starts (klio run, klio deploy, .remote(), klio model deploy), the command line prints what that GPU costs right now, once per GPU. klio run and klio deploy also take --gpu, for example --gpu L40S.

GPU types and prices

klio gpus
klio gpus --json

klio gpus lists every GPU type, its price per GPU-hour right now and whether it can start now. Prices follow the market and are checked every few minutes; a container keeps the price it started at for as long as it runs. The same list is on the pricing page.

GPU fallbacks

Give a list of GPU types and the workload uses whichever of them can start:

@function(gpu=["L40S", "A100-80"], image=image)
def generate(prompt: str) -> str:
    ...

If none of the GPUs you asked for can start right now, the workload stops before anything starts, and nothing is charged. With a list, the price of each GPU is printed.

Several GPUs in one container

A container can have 1, 2, 4 or 8 GPUs of one type, all on the same machine, so they can talk to each other directly (for example with NCCL or tensor parallelism).

from klio_compute import function

@function(gpu="L4", gpu_count=2)
def train():
    import torch
    return torch.cuda.device_count()  # 2

• Any other count stops before anything starts, and nothing is charged.
• Several GPUs need a paid workspace. On the free tier a container runs one GPU at a time.
• Each GPU count has its own price per GPU-hour, paid for each of the container's GPUs. klio gpus lists the counts that can start right now.
• With the Modal-compatible API, write gpu="L4:2".