GPU Support¶
When a machine joins the cluster it has to tell the scheduler what GPUs it has — a box with two 80 GB A100s should be able to run a big model; a laptop with none should only get CPU work. That's this library's job: on each worker it looks at the hardware, figures out the GPUs, and hands Ray a label describing them. The gateway then schedules against those labels ("find me a node with 40 GB of free NVIDIA VRAM"). It's small — no service, no database — imported by the worker's startup command and also runnable from the CLI.
The cluster is deliberately mixed: NVIDIA, AMD, and Apple machines with wildly different VRAM sizes, all in one pool. So the library does two things — detect whatever's on the box, then describe it to Ray in a way the scheduler can filter on.
1. Detecting the GPUs¶
detect_gpus() asks each vendor's own tool, in order, and stops at the first that answers — returning one GPUInfo (vendor, index, name, VRAM) per GPU, or an empty list meaning "CPU-only node."
| Order | Vendor | How it looks |
|---|---|---|
| 1 | NVIDIA | nvidia-smi — the standard query for index, name, total memory |
| 2 | AMD | rocm-smi; if that's missing (common inside containers) but /dev/kfd exists, it falls back to asking PyTorch |
| 3 | Apple | Apple Silicon's unified memory via sysctl (~75% of RAM counted as usable VRAM) |
| — | CPU-only | nothing answered |
A node is single-vendor — the first brand that reports any GPUs wins.
2. Describing them to Ray¶
Ray schedules on "custom resources" — just named numbers a node advertises. build_ray_args() turns the detected GPUs into the ray start flags that publish them:
- One VRAM number per GPU —
VRAM_MB_GPU0,VRAM_MB_GPU1, … Giving each GPU its own entry (instead of one lumped total) is what lets the scheduler target a specific GPU and handle machines with different-sized GPUs side by side. A mixed RTX 4090 (24 GB) + A100 (80 GB) box advertises{VRAM_MB_GPU0: 24576, VRAM_MB_GPU1: 81920}. - One vendor flag —
GPU_VENDOR_NVIDIA(or_AMD/_APPLE), so a job can require a specific brand.
It also tells Ray the node's CPU and GPU counts. Apple is the exception: it advertises its vendor and VRAM so the hardware is visible, but registers --num-gpus=0 — MPS scheduling isn't wired up yet, so those machines take CPU work only.
Using it¶
python -m gpu_support detect # JSON: what was found + the resources it'd publish
python -m gpu_support ray-args # print the `ray start` command it would run
python -m gpu_support start # detect, then exec `ray start`
The same builder makes the head node's command when asked (build_ray_args(is_head=True) → --head, ports 6379 / dashboard 8265 / client 10001); workers instead get --address=<head> and --block.
| Variable | Default | Description |
|---|---|---|
RAY_HEAD_ADDRESS |
ray-head:6379 |
The head a worker dials into |
NUM_CPUS |
auto | Override the CPU count it registers |
The worker image¶
Every vendor builds from one Dockerfile (docker/ray-worker) on top of rayproject/ray:2.54.0-py312, selected by two build args: RAY_IMAGE (the Ray base — e.g. a CUDA variant for NVIDIA) and TORCH_INDEX (the matching PyTorch wheels — CUDA, ROCm, or CPU). AMD RDNA3 cards (RX 7800 XT, …) also need HSA_OVERRIDE_GFX_VERSION=11.0.0, already set in the AMD compose profile.
Testing¶
# Unit tests run anywhere — the hardware probes are mocked:
pytest tests/gpu_support/ -v --ignore=tests/gpu_support/test_docker.py
# The cluster test auto-detects real hardware via /v1/resources and skips
# any vendor that isn't present:
pytest tests/gpu_support/test_docker.py -v
What they cover: test_detection (each vendor probe + CPU fallback), test_registration (building the resources and ray start args), test_gateway_gpu (the gateway's VRAM ledger tracking each node's vendor), test_docker (a real cluster).