Repository navigation
Conversation
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
force-pushed
the
feat/sampling-gpu-leases
branch
from
October 19, 2026 18:54
c9c96c8 to
76fdbd4
Compare
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
force-pushed
the
feat/sampling-gpu-leases
branch
from
October 19, 2026 20:51
76fdbd4 to
5e0953b
Compare
Collaborator
|
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 ? |
This was referenced Oct 21, 2026
Collaborator
Author
|
Closing in favor of new design |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.
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.