Skip to content

Commit 4f3fa44

Browse files
committed
Add RL + batch-inference interleaving guide
Time-slices one verl fully-async RL training job with a stock vLLM batch-inference server on the same GPU: the trainer has absolute priority (cuda-checkpoint C/R), vLLM harvests the trainer's idle valleys via its native sleep mode driven through the snapshot agent's workload channel. The RL job manifest uses the timeslice-verl package (pkg/integrations/verl: a FullyAsyncTrainer subclass registered as trainer name "timeslice" via verl's fully-async lifecycle hooks + trainer registry, selected with async_training.trainer_name=timeslice; ray_pg_extra_resources PG pinning). The vLLM supervisor and load generator live under examples/ and are example-only: the supervisor demonstrates the polite-tenant pattern (waiter poll, readiness-gate drain before /sleep - required, vllm#28714 - workload-channel registration); production batch serving should sit behind a queue-based front-end with retries. Requires a snapshot-agent build with workload-channel default routing (config-less Snapshot/Restore resolves the job's registered workload channel) - see values-timeslice.yaml note.
1 parent eb5bbb3 commit 4f3fa44

4 files changed

Lines changed: 1477 additions & 0 deletions

File tree

Lines changed: 293 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,293 @@
1+
# Interleaving RL Training with Batch Inference on a Shared GPU
2+
3+
This guide time-slices **one verl fully-async RL training job** with a
4+
**stock vLLM batch-inference server** on the same GPU. The RL trainer has
5+
absolute priority; vLLM harvests the trainer's idle valleys (sample waits
6+
between training bursts — 40–70% of wall-clock for async RL workloads) to
7+
serve latency-tolerant batch traffic. Neither workload's code is modified:
8+
verl is integrated through its fully-async lifecycle hooks: the
9+
`timeslice-verl` package from this repo ships a FullyAsyncTrainer subclass
10+
registered under the trainer name `timeslice`
11+
(`async_training.trainer_name=timeslice`), and vLLM runs its normal
12+
OpenAI-compatible server wrapped by a ~100-line supervisor (shipped in
13+
`examples/shadow-vllm.yaml` — demo-quality, example-only).
14+
15+
What makes the batch side cheap to yield: vLLM's native **sleep mode** moves
16+
weights + KV cache to host RAM in ~1–2 s and back in ~100 ms — no process
17+
restart, no model reload. The RL trainer side uses `cuda-checkpoint`
18+
(~10–35 s per swap for a 1.5 B trainer), which is amortized against training
19+
bursts that run for minutes.
20+
21+
**Topology** (two 1-GPU nodes; H100-class):
22+
23+
```
24+
ROLLOUT NODE (dedicated) TRAINER NODE (SHARED, time-sliced)
25+
┌──────────────────────────┐ ┌──────────────────────────────────┐
26+
│ rl-rollout │ ray │ rl-head shadow-vllm │
27+
│ RL rollout engine (vLLM, │◀───────▶│ RL trainer ⇄ vLLM server │
28+
│ generates continuously) │ │ cuda-checkpoint sleep mode │
29+
└──────────────────────────┘ │ C/R (~1-2 s swaps)│
30+
└──────────────────────────────────┘
31+
lock owner over time on the trainer node's GPU:
32+
vLLM ████████░T██████████░T█████████░T████████ (T = training burst)
33+
▲ trainer queues → vLLM yields in ≤ poll(0.5s) + drain(~3s) + sleep(~1-2s)
34+
```
35+
36+
The RL job is a **2-node Ray cluster**: head pod `rl-head` (verl driver +
37+
FSDP trainer) on the shared trainer node, rollout pod `rl-rollout` (vLLM
38+
rollout engine) on a dedicated node. Placement is pinned with Ray custom
39+
resources (`trainer_node` / `rollout_node`) through verl's per-pool
40+
placement-group bundle resources (hydra override `ray_pg_extra_resources`).
41+
42+
## 1. Components
43+
44+
| Piece | What it is | Where it runs |
45+
|---|---|---|
46+
| TimeSlice Orchestrator | gRPC group-lock service | `timeslice-system` ns |
47+
| Snapshot Agent | DaemonSet; executes cuda-checkpoint C/R AND relays sleep/wake to registered apps (workload channel) | every GPU node |
48+
| RL job (`examples/rl-job.yaml`) | verl `fully_async_policy` code-RLVR training (Eurus-2 code split, rewards = live test execution), integrated via verl lifecycle hooks (trainer subclass registered as `timeslice`); head pod + rollout pod + headless Service | `default` ns |
49+
| Shadow vLLM (`examples/shadow-vllm.yaml`) | EXAMPLE-ONLY: stock `vllm serve` + supervisor: holds the lock only while nobody waits; registers sleep/wake callbacks with the agent | `default` ns, trainer node |
50+
| Load generator (`examples/load-generator.yaml`) | EXAMPLE-ONLY: continuous `/v1/completions` client + throughput log | `default` ns |
51+
52+
How a handoff works, end to end:
53+
54+
1. Trainer's sample batch becomes ready → the timeslice trainer's hook calls
55+
`acquire()`.
56+
2. Supervisor's 0.5 s poll sees `waiter_queue_depth > 0` → it **drains**
57+
(closes its readiness gate, so the pod leaves the Service endpoints and
58+
new batch requests fail fast; waits ~3 s for in-flight requests to
59+
finish — vLLM's `/sleep` does NOT drain the scheduler, and sleeping with
60+
a request mid-decode is a fatal CUDA error, vllm#28714) → calls
61+
`release()`.
62+
3. Orchestrator snapshots the vLLM job — the agent's per-job backend
63+
resolution (explicit config → live workload channel → pod annotation →
64+
default cuda) lands on its **registered workload channel** → the agent
65+
invokes the supervisor's callback → `POST /sleep` → HBM freed in ~1–2 s.
66+
4. Orchestrator restores the trainer (cuda-checkpoint, trainer node only) and
67+
grants the lock. Trainer runs its burst (minutes), then yields; the
68+
platform snapshots it and wakes vLLM the same way in reverse; the
69+
supervisor reopens its readiness gate after the wake.
70+
5. Batch requests sent while vLLM is draining or asleep fail fast at the
71+
Service (connection refused — the pod is NotReady); the demo client
72+
counts them (`errors_or_asleep`) and carries on. Production clients
73+
should queue and retry — see §8.
74+
75+
## 2. Prerequisites
76+
77+
- Kubernetes cluster (validated on GKE/COS) with **two free 1-GPU nodes**
78+
(H100-class recommended; reference shape GKE `a3-highgpu-1g`, 26 vCPU /
79+
234 Gi each): one shared trainer node (RL head pod + shadow vLLM pod) and
80+
one dedicated RL rollout node.
81+
- Cluster-admin; `kubectl`, `helm` v3, `git`, `envsubst` on your workstation.
82+
- Pods can reach Docker Hub, GitHub, HuggingFace. No HF token needed (models
83+
are fetched by HF id per node).
84+
- Host RAM on the trainer node: the RL head pod requests 80 Gi (110 Gi
85+
limit; trainer checkpoint ≈ full GPU allocation) + 60 Gi for the shadow
86+
vLLM's sleep offload.
87+
88+
**Version pins** (keep on first run): verl branch
89+
[`feat/fully-async-lifecycle-hooks`](https://github.com/aishukamal/verl/tree/feat/fully-async-lifecycle-hooks)
90+
of `github.com/aishukamal/verl` — upstream
91+
`983cb0f24443f87b3d161fad318445130a620b07` plus two feature commits
92+
(fully-async lifecycle hooks + trainer registry; per-pool PG bundle
93+
resources); temporary fork until the commits land upstream. Job image
94+
`verlai/verl:vllm020.dev2`; shadow vLLM image `vllm/vllm-openai:v0.9.2`;
95+
platform images: official `ghcr.io/llm-d-incubation/llm-d-rl-time-slicing/*`
96+
containing per-job backend resolution in the
97+
agent (config-less requests resolve: explicit config → live workload
98+
channel → `timeslice.io/backend` pod annotation → default cuda; the
99+
workload-channel step is required by this guide — see the release-pin
100+
note in §3);
101+
integration package `timeslice-verl` from this repo,
102+
`pkg/integrations/verl/` (ConfigMap install, see §5); trainer
103+
model `deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B`; batch model
104+
`Qwen/Qwen2.5-0.5B-Instruct`.
105+
106+
## 3. Step 1 — Install the platform
107+
108+
```bash
109+
git clone https://github.com/llm-d-incubation/llm-d-rl-time-slicing.git
110+
cd llm-d-rl-time-slicing
111+
helm dependency update ./deploy
112+
# If a previous `timeslice` release exists, `helm uninstall` it first —
113+
# helm install fails on an existing release name. NOTE: the chart owns the
114+
# timeslice-system namespace, so uninstall deletes it; wait for the
115+
# namespace to finish terminating (~30-60 s) or the reinstall fails with
116+
# "namespace is being terminated".
117+
helm install timeslice ./deploy -n timeslice-system --create-namespace
118+
119+
kubectl -n timeslice-system get pods -o wide # orchestrator + agents Running
120+
```
121+
122+
> **Release pin:** the chart defaults pull the official
123+
> `ghcr.io/llm-d-incubation/llm-d-rl-time-slicing/*` images at `latest`,
124+
> which include per-job backend
125+
> resolution in the snapshot agent (#159 — required here: the workload
126+
> channel is how the agent drives vLLM's sleep/wake instead of
127+
> cuda-checkpointing it). Once the first tagged release containing it is
128+
> cut, pin it with `--set timesliceorchestrator.image.tag=<version>
129+
> --set snapshot-agent.image.tag=<version>`.
130+
131+
No snapshot-device filter is needed in this topology: every node has one GPU,
132+
and only the labeled pods on the trainer node (rl-head, shadow-vllm) are
133+
snapshot targets — the rollout pod is unlabeled on another node.
134+
135+
## 4. Step 2 — Pick the nodes and label the trainer node
136+
137+
```bash
138+
export TRAINER_NODE=<the shared 1-GPU node (RL trainer + shadow vLLM)>
139+
export ROLLOUT_NODE=<the RL rollout's dedicated 1-GPU node>
140+
141+
kubectl label node "$TRAINER_NODE" group.timeslice.io/trainers=true --overwrite
142+
```
143+
144+
Only the trainer node gets the label (the rollout node must NOT carry it).
145+
Verify an agent runs on the trainer node and the orchestrator synced the
146+
group:
147+
148+
```bash
149+
AGENT_POD=$(kubectl -n timeslice-system get pods -l app.kubernetes.io/name=snapshot-agent \
150+
--field-selector spec.nodeName=$TRAINER_NODE -o jsonpath='{.items[0].metadata.name}')
151+
echo "$AGENT_POD"
152+
kubectl -n timeslice-system logs deploy/timeslice-timesliceorchestrator --tail=50 | grep -i trainers | tail -3
153+
```
154+
155+
## 5. Step 3 — Launch (order matters)
156+
157+
First, publish the `timeslice-verl` integration package into the cluster as
158+
ConfigMaps (the RL job pods install it from these — no GitHub fetch):
159+
160+
```bash
161+
# Package source: the repo you cloned in §3
162+
PKG=llm-d-rl-time-slicing/pkg/integrations/verl
163+
kubectl create configmap timeslice-verl-root --from-file=$PKG/pyproject.toml
164+
kubectl create configmap timeslice-verl-src --from-file=$PKG/timeslice_verl/
165+
```
166+
167+
ConfigMap keys cannot contain `/`, so the package tree is split into two maps
168+
(project root + module sources) and reassembled inside the pod; this also
169+
works in air-gapped clusters.
170+
171+
**verl source**: the RL pods clone the `feat/fully-async-lifecycle-hooks`
172+
branch of `github.com/aishukamal/verl` and install it with
173+
`pip install -e ".[gpu]"`. Once the commits land upstream, point the
174+
`VERL_REPO`/`VERL_REF` pod envs at mainline verl instead.
175+
176+
Then start the **shadow vLLM first** so it owns the GPU during the trainer's
177+
long CPU-side setup, then the RL job (both pods at once — the rollout pod's
178+
installs run in parallel with the head's), then the load:
179+
180+
```bash
181+
export RUN_SECONDS=5400 # RL training budget (~90 min)
182+
183+
envsubst '${TRAINER_NODE}' < examples/shadow-vllm.yaml | kubectl apply -f -
184+
# wait for "[supervisor] vLLM is up" (first model download ~2-4 min):
185+
kubectl logs shadow-vllm --tail=20
186+
187+
envsubst '${TRAINER_NODE} ${ROLLOUT_NODE} ${RUN_SECONDS}' < examples/rl-job.yaml | kubectl apply -f -
188+
kubectl apply -f examples/load-generator.yaml
189+
```
190+
191+
The RL job spends 20–45 min on setup (verl install on both pods, the head's
192+
wait for the rollout pod to join the ray cluster, model + dataset prep —
193+
**both pods prepare the dataset locally**: the `FullyAsyncRollouter` actor
194+
runs on the rollout pod and reads the parquet there; no shared volume, same
195+
seed ⇒ identical files) before it first requests the GPU — vLLM serves batch
196+
traffic the whole time.
197+
198+
## 6. Step 4 — Watch it work
199+
200+
```bash
201+
kubectl logs batch-load-generator --tail=6
202+
# [12:01:05] last 5s: completed=41 errors_or_asleep=0 <- vLLM holds GPU
203+
# [12:04:35] last 5s: completed=0 errors_or_asleep=12 <- trainer burst
204+
# [12:08:10] last 5s: completed=38 errors_or_asleep=1 <- vLLM back
205+
206+
kubectl logs shadow-vllm --tail=10
207+
# [supervisor] trainer is waiting - yielding GPU
208+
# [supervisor] drained in 2.82s
209+
# [supervisor] vLLM slept (HBM -> host RAM) in 0.97s
210+
# [supervisor] lock reacquired (waited 214380 ms, context_restored=True)
211+
# [supervisor] vLLM woke in 0.05s
212+
213+
# lock handoffs: the integration package logs every acquire/release
214+
# ([timeslice] ... ACQUIRE/RELEASE lines, forwarded into the head pod's log):
215+
kubectl logs rl-head --tail=2000 | grep -E "ACQUIRE|RELEASE" | tail -4
216+
# [timeslice] job=rl-trainer ACQUIRE group=trainers waited=8210ms context_restored=True
217+
# [timeslice] job=rl-trainer RELEASE group=trainers pending_waiters=1 snapshot_deferred=False
218+
```
219+
220+
Healthy steady state: the load generator alternates between full-throughput
221+
windows (trainer idle) and error windows a few minutes long (trainer burst);
222+
supervisor drain ≈ 3 s, sleep ≤ 3 s, wake ≤ 1 s; trainer `waited=` on
223+
ACQUIRE ≈ 5–11 s (poll latency + drain + vLLM sleep + the trainer's own
224+
cuda-checkpoint context restore, ~5 s once training state is large — the
225+
trainer never queues behind batch work).
226+
227+
Success criteria to check after ~3 training steps:
228+
229+
- Trainer step time within ~10% of a solo run (compare `train.log` timing to
230+
a run without the shadow vLLM, or to the reference: median 570 s/step).
231+
- Batch throughput > 0 in every trainer-idle window.
232+
- No cuda-checkpoint operations targeting the rollout node in any agent log
233+
(the rollout node must never show cuda-checkpoint activity).
234+
235+
## 7. Step 5 — Collect results
236+
237+
```bash
238+
mkdir -p results
239+
kubectl cp rl-head:/workspace/results/train.log results/train.log || true
240+
kubectl logs batch-load-generator --timestamps > results/batch_throughput.log
241+
kubectl logs shadow-vllm --timestamps > results/supervisor.log
242+
```
243+
244+
`train.log` (trainer bursts, with the `[timeslice]` ACQUIRE/RELEASE
245+
lock-handoff lines) + `batch_throughput.log` (harvested inference) together
246+
give the shared-GPU timeline: every second is either a training burst, batch
247+
serving, or a swap.
248+
249+
## 8. Troubleshooting
250+
251+
| Symptom | Cause | Fix |
252+
|---|---|---|
253+
| vLLM never sleeps at handoff; trainer restore then fails with OOM on the shared GPU | Workload registration didn't reach the agent (check supervisor log for `workload registered`) | Verify `TIMESLICE_AGENT_ADDR` resolves to the node IP :9001 and agent logs show the registration; the supervisor's local wake fallback covers restores but sleep MUST go through the channel |
254+
| `POST /sleep` returns 404 | vLLM started without dev endpoints | `VLLM_SERVER_DEV_MODE=1` and `--enable-sleep-mode` are both required (set in the manifest) |
255+
| vLLM dies seconds after a sleep (`CUDA error: an illegal memory access` in EngineCore), then the group faults on the failed wake | `/sleep` was called with requests in flight — vLLM does not drain the scheduler on sleep (vllm#28714, also affects ≥0.10.x; requests sent to a sleeping engine crash it too, vllm#15483) | Keep the supervisor's readiness-gate drain (in the manifest): gate closed + `running+waiting==0` BEFORE `release()`. Never call `/sleep` on a serving engine directly |
256+
| Trainer placement group lands on the rollout node (or rollout work on the trainer node) | PG pinning not active — `--resources` missing on a `ray start` line, or the `ray_pg_extra_resources` hydra override is missing/typo'd, or the verl build lacks the feature | Verify both `ray start` lines carry `--resources`, keep the override in `examples/rl-job.yaml`'s `run_head.sh` intact, and use the pinned verl branch (§2); check `ray status` custom resources |
257+
| cuda-checkpoint activity on the rollout node's agent, or rollout throughput collapses during handoffs | The rollout pod/node got labeled into the group | The rollout node must never show cuda-checkpoint activity: keep `timeslice.io/*` labels off rl-rollout and the group label off `$ROLLOUT_NODE` |
258+
| `rl-head` log stuck at `Phase 2b: wait for the rollout pod to join`, then `FATAL: rollout pod did not join` | rl-rollout Pending/crashed, or headless Service `rl-head` missing | `kubectl get pod rl-rollout; kubectl logs rl-rollout --tail=50`; apply the whole rendered `examples/rl-job.yaml` |
259+
| Load generator: 100% errors even when trainer is idle | vLLM crashed (supervisor exits, pod restarts) or Service selector mismatch | `kubectl logs shadow-vllm --previous`; check `kubectl get endpoints shadow-vllm` |
260+
| RL head or shadow vLLM pod Pending | Trainer node CPU/RAM too small for 8 CPU + 80 Gi (head) alongside 6 CPU + 60 Gi (vLLM) | Use an a3-highgpu-1g-class node (26 vCPU / 234 Gi), or shrink requests |
261+
262+
Demo-client caveat, worth stating to any customer: while vLLM is draining or
263+
asleep, new batch requests fail fast at the Service (the pod is NotReady).
264+
The demo load generator just counts those errors. A production batch
265+
front-end should be queue-based with retries — this composes naturally with
266+
[llm-d async processor](https://github.com/llm-d/llm-d-async) as the batch
267+
dispatcher, which is the planned productization path.
268+
269+
## 9. Adapting
270+
271+
- **Bigger batch model**: anything that fits in `VLLM_GPU_FRAC` of the shared
272+
GPU alongside zero trainer residency (they never co-reside). Sleep/wake
273+
time scales with weight size (~1 s per 15 GB to host RAM).
274+
- **Your RL workload**: swap model/dataset in `examples/rl-job.yaml` exactly as in the
275+
two-RL-jobs guide (§9 there); everything under "Required" comments stays.
276+
- **Yield latency vs. poll cost**: `WAITER_POLL_SECONDS` (default 0.5) is the
277+
worst-case extra wait the trainer sees.
278+
- **Trainer solo baseline** (for the ≤10% overhead check): run `examples/rl-job.yaml`
279+
alone — without the shadow vLLM the trainer acquires instantly every step
280+
and the platform defers snapshots (nothing contends), so it behaves like an
281+
unshared run.
282+
283+
## 10. Teardown
284+
285+
```bash
286+
kubectl delete pod shadow-vllm batch-load-generator --ignore-not-found
287+
kubectl delete pod rl-head rl-rollout --ignore-not-found
288+
kubectl delete service shadow-vllm rl-head --ignore-not-found
289+
kubectl delete configmap shadow-vllm-scripts rlbatch-trainer --ignore-not-found
290+
kubectl delete configmap timeslice-verl-root timeslice-verl-src --ignore-not-found
291+
helm uninstall timeslice -n timeslice-system
292+
kubectl label node "$TRAINER_NODE" group.timeslice.io/trainers- || true
293+
```
Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
# Continuous batch-inference client: hammers the shadow vLLM Service and
2+
# reports throughput every 5 s. Errors while the trainer holds the GPU are
3+
# EXPECTED (the server is asleep) — production clients should queue/retry;
4+
# this demo client just counts them.
5+
apiVersion: v1
6+
kind: Pod
7+
metadata:
8+
name: batch-load-generator
9+
namespace: default
10+
labels:
11+
app: batch-load-generator
12+
spec:
13+
restartPolicy: Always
14+
containers:
15+
- name: loadgen
16+
image: python:3.11-slim
17+
command: ["/bin/bash", "-c"]
18+
args:
19+
- |
20+
pip install --quiet requests
21+
python3 - <<'PY'
22+
import datetime, time, requests
23+
24+
URL = "http://shadow-vllm.default.svc.cluster.local:8000/v1/completions"
25+
PROMPT = "Explain, in three sentences, why sharing a GPU between RL training and batch inference raises utilization."
26+
MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
27+
28+
completed = errors = 0
29+
window = time.time()
30+
print("batch load generator started", flush=True)
31+
while True:
32+
try:
33+
r = requests.post(URL, json={"model": MODEL, "prompt": PROMPT,
34+
"max_tokens": 48}, timeout=10)
35+
if r.status_code == 200:
36+
completed += 1
37+
else:
38+
errors += 1
39+
except Exception:
40+
errors += 1
41+
time.sleep(0.5)
42+
if time.time() - window >= 5.0:
43+
ts = datetime.datetime.now().strftime("%H:%M:%S")
44+
print(f"[{ts}] last 5s: completed={completed} errors_or_asleep={errors}",
45+
flush=True)
46+
completed = errors = 0
47+
window = time.time()
48+
PY
49+
resources:
50+
requests:
51+
cpu: "1"
52+
memory: 1Gi

0 commit comments

Comments
 (0)