Skip to content

KubeRay Service

This service is the master's hands on Kubernetes. It runs on the master's host network — the one component that can talk to both the Kubernetes API and the Ray head — and it does two jobs:

  1. Onboards machines into the cluster (the join flow below), and
  2. Acts as the K8s/Ray proxy for everything else. The job_gateway runs inside a Docker bridge network and can't reach the K8s API or the Ray head, so it goes through this service instead.

It listens on port 8800. Its only local state is a small SQLite table of join keys (join_keys); the real cluster state lives in Kubernetes CRDs, so there's nothing to migrate to Postgres.

Onboarding a machine (the join flow)

The headline feature: turning a stranger's GPU box into a cluster worker over a secure tunnel, with one command.

  1. Admin mints a one-time-ish join key (POST /v1/nodes/join-keys) and gets back a ready-to-run command: sudo MASTER_IP=… bash join-node.sh <KEY>. Org keys are random tokens; provider keys are JWTs that double as the provider's auth token for the SDK/web UI. A VPN IP is reserved.
  2. The GPU owner runs join-node.sh, which does a two-step handshake:
  3. POST /v1/nodes/validate — checks the key and returns what the node needs to connect: the K3s token, the master's WireGuard public key, the node's assigned VPN IP, and the master's VPN IP. The master adds the WireGuard peer eagerly so the tunnel is live the moment the node brings up wg0.
  4. The node brings up its WireGuard tunnel + K3s agent, then calls POST /v1/nodes/register, which finalizes the peer and — for a GPU node — creates a Ray worker group (GPUs auto-detected via gpu_support). CPU-only nodes join without a worker group.
  5. Identity is pinned to the WireGuard public key. Re-running the script (after a reboot or a master switch) reuses the same VPN IP. A different machine trying to claim a hostname that's already taken is rejected (409) so it can't knock the original node off the cluster. Keys are reusable — one key can onboard many machines; expiry is the only gate.
  6. A major worker-vs-server version mismatch is rejected (426 → "run update.sh"); minor/patch differences are fine. Join/leave also maintains the WireGuard mesh (per-node pod CIDRs) and the Prometheus DCGM GPU-metrics scrape targets. POST /v1/nodes/disconnect tears all of that down for a node.

Managing the Ray cluster

CRUD for RayClusters and their worker groups (/v1/kuberay/clusters). The serve path reconfigures an endpoint's workers with atomic swap / replace-worker-groups so it can change hardware without dropping the endpoint. Worker-group add/remove retries on 409 Conflict, since concurrent CRD writes race.

The proxies (why they exist)

Because the gateway can't reach K8s or Ray directly, this service exposes thin pass-throughs — unauthenticated internal endpoints on the trusted host network:

  • /v1/crd — read/write the GridweaveJob and GridweaveEndpoint CRDs that hold the gateway's job/endpoint state (implemented via gridweave_controller.crd).
  • /v1/ray — Ray Serve application statuses and delete-by-name.

KubeTorch (distributed training)

/v1/kubetorch syncs a user's code to the worker pods and launches a distributed training job across them, then reports job status.

Endpoints

Group (prefix) Endpoints Auth
Clusters /v1/kuberay/clusters POST/GET `;GET/DELETE/{name};PATCH /{name}/scale;POST /{name}/worker-groups[-raw];GET /{name}/workers,/worker-groups,/worker-groups/{g}/logs;DELETE /{name}/worker-groups/{g}[/pods];POST /{name}/swap-worker-groups,/replace-worker-groups` admin
Nodes /v1/nodes POST/GET/DELETE /join-keys; POST /disconnect admin
POST /validate, /register join key
POST /verify-provider provider JWT
POST /report-version worker
CRD proxy /v1/crd /jobs + /endpoints CRUD (POST/GET/PATCH .../status/DELETE), GET /endpoints/by-name internal
Ray proxy /v1/ray GET /serve/statuses; DELETE /serve/applications/{app} internal
KubeTorch /v1/kubetorch POST /sync, /train; GET /health (admin); GET /jobs/{id} (user) admin/user
GET /healthz

Configuration

Variable Default Description
KUBERAY_PORT 8800 Service port
KUBERAY_DB /data/kuberay.db SQLite path (join keys only)
KUBERAY_NAMESPACE gridweave K8s namespace for RayClusters
ADMIN_TOKEN Admin bearer token
JWT_SECRET dev-secret-change-me Signs provider join-key JWTs; validates admin/provider JWTs
SERVER_VERSION 0.0.0 Master version, for the worker version gate
RAY_IMAGE rayproject/ray:2.54.0-py312 Ray head / CPU worker image
RAY_IMAGE_WORKER gridweave-worker:latest GPU worker image (built on each node by join-node.sh)
CLUSTER_NAME gridweave-cluster Default RayCluster name
MASTER_IP Master public IP (embedded in the join command)
MASTER_VPN_IP 10.100.0.1 Master's WireGuard address
RAY_DASHBOARD_URL / RAY_HEAD_ADDRESS derived from MASTER_VPN_IP Where the Ray proxy reaches Serve / the Ray client
CORS_ORIGINS * Allowed origins
RATE_LIMIT_PER_MINUTE 6000 Per-IP rate limit

Running & testing

# Runs on the master with host networking to reach K3s at localhost:6443
# (mounts /etc/rancher/k3s/k3s.yaml as kubeconfig):
docker compose -f docker/kuberay_service/docker-compose.yml up --build -d

# Unit tests (K8s + WireGuard + Ray are mocked):
pip install -e ".[kuberay,test]"
pytest tests/kuberay_service/ -v --ignore=tests/kuberay_service/test_docker.py

Coverage highlights: test_nodes / test_nodes_extra (join-key CRUD, validate/register, disconnect, verify-provider, version gate), test_wg_locking (concurrent WireGuard/VPN-IP mutation), test_worker_group_409_retry (retry on CRD conflict), test_worker_version (version annotation), test_cluster / test_routers_clusters (RayCluster CRUD), test_kubetorch* (code sync + training), test_ray (Serve proxy), test_dcgm_targets (Prometheus target lifecycle), test_docker (deployed). Interactive API docs at /docs.