Skip to content

sampling: let FFT samplers time-slice with their trainer - #342

Closed
ShubyM wants to merge 5 commits into
openrl-dev:mainfrom
ShubyM:feat/sampling-gpu-leases
Closed

ShubyM wants to merge 5 commits into
openrl-dev:mainfrom
ShubyM:feat/sampling-gpu-leases

Conversation

@ShubyM

@ShubyM ShubyM commented Oct 19, 2026 •

Copy link
Copy Markdown
Collaborator

Stacked on #339; review the last commit.

This lets FFT trainers time-slice with their samplers. With #339 alone, every pool sampler gets its own GPU, so an FFT job needs two. With this PR, an FFT job's sampler shares its trainer's GPU through the node time-slicer.

  • The dispatcher holds the lease. It holds the replica's lease for one unit of work: wake, load weights, sample, sleep, release. A turn never ends in the middle of a weight load.
  • vLLM boots inside a lease. A small launcher starts vLLM inside a lease, then sleeps it. It kills vLLM if vLLM ignores SIGTERM during startup.
  • Faults reach the launcher. If parking the GPU fails, the slicer tells the launcher, which exits instead of hanging the trainer.
  • Only the holder releases. A grant can be released only by the connection that holds it.

Tested with unit tests (330) and on GKE with a 0.6B model: 0 of 423 samples overlapped, park and restore took 0.8 s each, and FFT ran at 13 s per step on one GPU. The fault path needs the time-slicer image from this PR.

ShubyM and others added 2 commits October 19, 2026 18:54
The dispatcher sends each pending request through the pool router, Envoy
in front of the llm-d EPP, and commits the first result. A result counts
only if the replica that answered still holds the expected weights and has
not paused to load others since the request was sent. When the pool moves
to a new revision, each replica loads the checkpoint or sparse delta,
resets its prefix cache and resumes. The dispatcher sleeps nowhere: it
waits on the stream for new requests, and wakes at least every second to
retry a failed send. On one machine, OPEN_RL_SAMPLER_ENDPOINTS lists the
replicas and no router is needed. FFT sparse deltas now record the save
they apply on top of, so a replica on any other version loads a
checkpoint instead.

A pool's only dispatcher reads each stream entry once and drops it when it
commits. A Redis or Kubernetes error in a dispatch pass is logged and
retried rather than ending the dispatcher, and commits wait for one of 64
Redis connections, since every in-flight send commits at once when replicas
come up. A request pending longer than OPEN_RL_SAMPLE_DEADLINE_SECONDS
(default one hour, measured from its stream id) fails with the pool and
revision in the error, so a pool that can never serve a revision does not
leave its clients waiting. The replica listing has a 10 s Kubernetes API
timeout.
- Recover replicas a dead dispatcher left paused or inside a weight
  update. A stop or kill during a load left the replica paused while it
  reported the new version, and a kill between start_weight_update and
  finish_weight_update left vLLM refusing every later update. The
  dispatcher resumes a replica left paused, closes an open update with a
  bare finish_weight_update and marks the replica unknown until a full
  load, and on shutdown gives loads in flight five seconds.
- Simplify the dispatcher: inline the deadline check into dispatch, drop
  the OPEN_RL_SAMPLE_DEADLINE_SECONDS knob and a concurrency default
  nothing used, chain raise_for_status, commit inside send's branches
  and flatten load_weights.
- Point an FFT delta at the save its optimizer step started from. The
  parent was the last sampler save, but a delta holds only the last
  optimizer step, so a save the pool never sampled broke the chain for
  every replica, and two optimizer steps between saves silently served
  wrong weights. The trainer lists every save a delta applies on as its
  bases, and writes a full checkpoint when no save holds the weights its
  step started from.
- Free a dispatch slot when its record fails at the deadline. A generate
  had no timeout and the deadline commit left the send running, so a
  replica that never answered kept its slots and starved healthy ones.
  The deadline commit cancels the send, and each send is bounded by the
  time its record has left.
- Keep a replica paused when its load is cut off mid-update. The
  cancelled load's finally resumed the replica inside its open update,
  so it served under its old version with weights that may be partly the
  new revision. Such a load now leaves the replica paused, and on first
  sight a dispatcher closes an open update on every replica, paused or
  not.
- Keep at most 32 idle connections in the dispatcher's HTTP client. It
  kept every connection a burst opened, hundreds of them, and httpx
  scans its whole pool whenever a request starts or ends, so a burst
  held the dispatcher's one loop at a full core. Its 5 s weight checks
  then timed out about 21 s after the send, and finished samples were
  generated again. Sends still have no connection limit.
- Probe a replica for an open weight update only on first sight and
  after one of the dispatcher's own loads breaks off. Each load ended by
  dropping the check, so the pass after every load sent a bare
  finish_weight_update, which vLLM refuses with a logged 500 traceback
  when no update is open.
@ShubyM
ShubyM force-pushed the feat/sampling-gpu-leases branch from c9c96c8 to 76fdbd4 Compare October 19, 2026 18:54
ShubyM and others added 3 commits October 19, 2026 19:57
Every sampler is now a sampling pool replica. On Kubernetes a pool is
vllm serve Workloads plus one CPU Deployment running the dispatcher, Envoy
and the llm-d EPP, with a NetworkPolicy and a readiness probe tied to the
expected weight version. On one machine the local worker manager starts
one vllm serve and a dispatcher that sends to it directly. asample and
create_sampling_session go through the pool accept path, LoRA saves are
frozen under their sampler reference, and checkpoints that pending
requests name are kept. Pool names are derived valid Kubernetes names, so
OPEN_RL_SAMPLER_POOLS is gone. The queue sampler, its Redis sampling queue
and the sampler_ready handshake are deleted.

A save under a frozen sampler name whose adapter is byte-identical to the
snapshot changes nothing, as when a script evaluates twice with no training
between; different weights under a used name still fail, since vLLM caches
adapters by name.

The dispatcher runs on the API-server image when OPEN_RL_DISPATCHER_IMAGE
names it, which the open-rl-images ConfigMap does; its pull from the 7 GB
worker image was the last thing a cold pool waited on. vLLM replicas keep
their torch.compile and Triton caches under OPEN_RL_TMP_DIR/cache on the
shared volume, so a new replica does not compile again.
- Call a pool sampler ready only once its pod is Ready. Its Workload
  reads Running from container start, through vLLM's load and through a
  crash loop, so the model view called it ready minutes before it could
  serve.
- Bind the worker-launcher Role to open-rl-sa, the account the API server
  runs as. open-rl-api-server is created nowhere, so a fresh install could
  not create Workloads, Deployments or NetworkPolicies.
- Refuse to start pool sampling without REDIS_URL, in every launcher
  mode. Before, the dispatcher died on the missing variable and asample
  failed on the in-memory store.
- Wait out NFS lookup caching before refusing a sample. The API server's
  NFS client could still hold a failed lookup of weights another node had
  just saved, so accept_pool_sample now lists the directory above the file
  and looks again for up to two seconds before the 400.
- Share only the Triton cache between vLLM replicas. vLLM writes its AOT
  compile artifact through a temp file named by PID alone, and replicas in
  different pods often run the same PID, so two replicas compiling at once
  could leave a mixed artifact.
- Drop pool manifest, router and RBAC settings that restate llm-d and
  Envoy defaults or that no code uses, and inline single-use helpers.
- Keep local pool files under OPEN_RL_TMP_DIR, default /tmp/open-rl. With
  the documented local setup the dispatcher put its revision file under
  /mnt/shared, died with PermissionError and left every local sample
  waiting, and vLLM got the same unwritable Triton cache.
- Fall back to the same /tmp/open-rl on the cluster. The readiness
  probes path took that default, but a pools dispatcher and every
  worker pod still fell back to /mnt/shared/open-rl, so an API server
  without OPEN_RL_TMP_DIR rendered a probe for a revision file no
  dispatcher wrote, and no replica ever turned Ready. The probe, the
  dispatcher and the pods now read one tmp_root(). The in-tree manifests
  all set OPEN_RL_TMP_DIR, so they never hit this.
- Create the LoRA snapshot directory before a LoRA pool starts. vllm
  serve exits when VLLM_LORA_RESOLVER_CACHE_DIR is missing, and only
  freeze_adapter created it, so a pool started before the first LoRA save
  crash-looped and base-model-only use never recovered.
- Pass trunk's 0.90 GPU memory fraction to every replica again, with the
  env still winning. Replicas ran at vLLM's 0.92, the setting tied to FFT
  delta applies running out of memory on an L4. The device-scaled batch
  limits are deleted instead, since the API server renders the args
  without knowing the device and vLLM 0.30 sizes them itself.
- Keep a local pool's dispatcher running beside its vLLM. A dead local
  dispatcher was never restarted, so its samples waited with nothing to
  fail them, and a vLLM restart left the old dispatcher reading the
  stream beside the new one.
- Hold pending prompts as array('I'). A prompt parsed into a list costs
  about 36 bytes a token, so the 2Gi dispatcher ran out near 40M pending
  prompt tokens and crash-looped re-reading the stream. A pending prompt
  now costs about 4 bytes a token once read.
- Call a pool sampler ready only while its dispatcher has a pod. The
  readiness probe compares vLLM's version with the revision file, which
  outlives the dispatcher, so a sampler read ready with its dispatcher
  dead while every sample waited.
…her lease

Every pool replica held a GPU of its own, so an FFT job that asked to
share one GPU between its trainer and sampler still took a second GPU for
its sampler. A pool sampler whose job shares a GPU is now non-exclusive, in
the scheduler manager and in local mode, and takes turns with its trainer
through the node time slicer, with no proxy in the request path. Exclusive
and LoRA pools keep a GPU of their own.

The worker manager marks such a replica with OPEN_RL_TIME_SLICED_SAMPLER=true
and runs it through server.sampler_launcher. The launcher registers the
replica workload, acquires, starts vllm serve --enable-sleep-mode inside that
turn, waits until it is healthy, sleeps it, releases, writes the marker the
startup probe waits for, and supervises vLLM for its lifetime, keeping the
registration open. vLLM ignores SIGTERM while its engine starts, so the
launcher follows a forwarded signal with SIGKILL after five seconds.

For each turn the pool dispatcher, which finds such replicas by the same
variable in the pod env, or in its own env on one machine, acquires the
workload through the slicer on the replica's node, wakes vLLM, runs any
weight load as a whole, sends pending requests to that replica directly,
sleeps it and releases. A turn ends when the replica goes idle, or after
OPEN_RL_SAMPLER_TURN_SECONDS with requests still pending, so the trainer
gets its turn.

The slicer changes that make this safe:
- A grant remembers the connection whose ACQUIRE holds it. When that
  connection closes and did not register the workload, as when a
  dispatcher dies mid-turn, the slicer releases the workload, which
  checkpoints the replica before the next waiter is granted.
- Only the connection holding a grant can release it. The launcher and the
  dispatcher share the workload name, so a turn that outlived a launcher
  restart could otherwise end the new launcher's boot turn and park vLLM
  mid-boot.
- When the park at the end of a turn, or for a dead dispatcher, fails, the
  slicer keeps the grant and sends the registering connection a FAULT line.
  The launcher then stops vLLM and exits, so the GPU frees for the trainer
  and kubelet restarts the replica, where before both waited forever with
  no error.

A turn is bounded when a call to its replica hangs. Sends still in flight
when a turn ends get TURN_GRACE_SECONDS (10) and are then cancelled, their
records staying pending for a later turn, and vLLM sleeps with mode=abort
before the lease is released. Calls that wake, prepare or sleep a replica
give up after 300 s, and /update_weights after 1800 s, on the router path
too.

Tests drive the real slicer over its socket, the real dispatcher with Redis
and the real launcher process against tests/vllm_stub.py: a turn never
splits a weight load, a dead dispatcher's lease is parked before the
trainer is granted, a failed park stops the launcher so the trainer goes
next, a stale release cannot end a restarted launcher's boot turn, and a
hung generate is cancelled so the trainer steps. The worker manager tests
check that an exclusive FFT job's workers are never time-sliced.
@ShubyM
ShubyM force-pushed the feat/sampling-gpu-leases branch from 76fdbd4 to 5e0953b Compare October 19, 2026 20:51
@ShubyM ShubyM changed the title sampling: let FFT samplers time-slice with their trainer in pools mode sampling: let FFT samplers time-slice with their trainer Oct 19, 2026
@droot

droot commented Oct 20, 2026

Copy link
Copy Markdown
Collaborator

@ShubyM

Looks like we have evolved the design significantly from what we initially reviewed. Is there a final design that i can take a look at ?

@ShubyM

ShubyM commented Oct 21, 2026

Copy link
Copy Markdown
Collaborator Author

Closing in favor of new design

@ShubyM ShubyM closed this Oct 21, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants