Repository navigation
Conversation
ShubyM
force-pushed
the
feat/sampling-pool-failures
branch
2 times, most recently
from
October 26, 2026 23:27
d9d9382 to
1604e82
Compare
Each model's sampling requests go to a pool stream, open_rl:sampling:<pool>, instead of a Redis list that a sampler emptied on read. asample adds a request and returns. The pool's dispatcher saves one result per request, so a sampler that stops loses nothing. A pool that one job owns runs server.vllm_sampler, vLLM in its own process, when it has one replica, for LoRA and FFT alike. It takes time-slicer turns only when it shares its trainer GPU. A pool with two or more replicas gets vllm serve Workloads behind a dispatcher pod with Envoy and the llm-d EPP. Local pools always have one replica. A pool that shared LoRA jobs use is always routed, even at one replica. Each job's ensure then only adds replicas, so the pool grows to the largest sampler_replicas any of its jobs asks for, and a second job never meets an in-process sampler it would have to replace. A sampler on its trainer's GPU is its pool's only replica, so create_model refuses sampler_replicas above 1 for an FFT job that shares its trainer's GPU, and tells the user to set openrl.exclusive=true. A routed pool's dispatcher Deployment and NetworkPolicy carry an openrl.io/owner label. release_owner deletes them by that label as well as through the pool's Workloads, so they do not outlive Workloads that went first. The worker launcher's Role gains list on deployments and networkpolicies for that lookup. The queue sampler and its vLLM patch are deleted, along with the store's sampler queue reads. The in-process sampler keeps sizing vLLM's batch for its GPU with sampler_batch_limits. sample_params and sample_result build the vLLM params and the saved result for both samplers. The in-process sampler raises RevisionUnloadable for a delta that does not apply on its weights, so those requests fail at once there too. It also fails a request vLLM refuses as a client error, such as a prompt past max_model_len or a token out of vocabulary, and a request that fails with EngineGenerateError, instead of sending it again until the deadline. vLLM loads a LoRA adapter inside the engine step, and a missing adapter raises LoRAAdapterNotFoundError there, which kills EngineCore. So the in-process sampler checks for the snapshot's adapter_config.json before it sends a record. On a miss it lists the snapshot directory once, since an NFS client keeps a failed lookup until it does, and if the file is still missing the record stays pending and is sent again. When the engine dies, every request in flight fails with EngineDeadError, and the sampler saves a failed result for each, as main's sampler did. A request that kills the engine would otherwise be served again by each restarted engine until the deadline. A request sent after the engine died never reached it, so it stays pending for the restarted worker. A sampler that shares its GPU with the trainer runs its batches through TimeSlicedDispatcher in vllm_sampler, one time-slicer turn per batch with main's holding_gpu. It loads weights through the same Dispatcher.ready as the routed dispatcher. The shared Dispatcher knows nothing about turns. A sampling session opened on a bare base-model name keeps main's rule, FFT when the deployment enables FFT and LoRA otherwise. Its runtime was the base model's name, which is also the name of the base model's shared LoRA runtime, so an FFT sampler for it took the LoRA pool's owner and read the LoRA pool's stream. TrainingModelMetadata.runtime now names such a runtime base-<model>. Cluster samplers keep Triton's kernel cache at <OPEN_RL_TMP_DIR>/vllm-cache/triton on the shared volume, beside vLLM's compile cache. Local samplers keep Triton's default, which already persists on one machine. The scheduler worker manager refuses to start without OPEN_RL_TMP_DIR, since its default of /tmp/open-rl is not on the shared volume. Each record names the job that sent it as "job", which job_of reads, so later changes can tell a pool's jobs apart.
A shared LoRA pool serves every job on its base model, and the dispatcher sent its records oldest first. One job's big batch then held every slot until it finished, and other jobs waited behind it. Dispatcher.to_send takes one record from each job, read from the record's "job" field, in turn, oldest first within a job, starting with the job served longest ago. A record accepted before this change has no job and groups under its model.
A request could wait forever in two ways. A pool whose Workloads or dispatcher were deleted while its jobs still had requests pending never came back, since only create_sampling_session and save_weights_for_sampler ensured it. And the deadline was enforced only by the pool's dispatcher, so with no dispatcher nothing failed the pool's requests. Two siblings of the session reaper, each a task of its own on the same 30 s interval, now tend the pool streams. The deadline sweep fails every entry older than SAMPLE_DEADLINE_SECONDS through SamplingStore.fail_older_than. The reconcile looks at each job with pending entries in each pool, under the owner lock the reaper takes. They run apart, so a reconcile stuck in an ensure never holds up the sweep. The reconcile handles a pass's jobs concurrently, since an ensure can wait minutes for an old pool to terminate, but at most RECONCILE_CONCURRENCY (4) at a time, so stuck ensures cannot take every thread the API server runs blocking calls on. A job whose pool still has a sampler Workload, or a local sampler process, is left alone, since a client may sample without a session. When the pool is gone, a job whose owner has a live session gets worker_manager.ensure(job, "sampler"), which creates only what is missing, so a deleted pool comes back within 30 s even when the client sends no new sample and only polls retrieve_future. A job with no live session has its pending entries failed instead, so a later pool on the same base model does not replay them. The lock keeps a reaper release from landing between the checks and the ensure, which would recreate a pool that nothing removes again. With no worker manager, standing worker deployments own the pools and nothing is reconciled. delete_model takes the same owner lock. Under it, delete_model marks the job completed, fails the job's pending pool requests, and releases its workers. The reconcile reads the job's status again under the lock, and fails a completed job's requests instead of bringing its pool back. set_result's XDEL leaves an empty stream behind, and Redis never evicts it, so each FFT or exclusive job left one key forever. SamplingStore.delete_if_empty now deletes a pool's stream when nothing is pending in it. The reaper calls it for each runtime release_owner returns, delete_model after it releases a job's workers, and the reconcile after it fails the requests of a job whose pool is gone. It checks XLEN and deletes in one transaction that watches the stream, so a request accepted at that moment stays. pool_name, split out of TrainingModelMetadata.pool, names the pool of a runtime. Stream entries carry the job and the revision as fields of their own. The deadline sweep, pending_jobs, fail_job and pinned_revisions read each stream 100 entries at a time, so a backlog of long prompts never sits in memory at once, and the last three read the job and revision fields without decoding prompts. An entry written before those fields existed decodes its record. A record from before records named their job counts under its model, and is skipped when that model is no job of the record's pool. Results are saved through set_result, which keeps only the first, so a dispatcher answering the same entry is harmless. Each dispatcher pass drops the records it holds whose result was saved elsewhere, found with one pipelined EXISTS on their result keys, so it does not keep sending them or let them hold the pool on their revision. OPEN_RL_SAMPLE_DEADLINE_SECONDS sets the deadline, default 3600. Worker pods and pool dispatchers get it from the API server's environment. The dispatcher_body docstring no longer claims two dispatchers never overlap. An eviction or a partition can still overlap them, and set_result's transaction keeps each request to one result. Out of scope: retrieve_future's "try again" answer still says nothing about the pool's state.
ShubyM
force-pushed
the
feat/sampling-pool-failures
branch
from
October 27, 2026 19:06
1604e82 to
9093a86
Compare
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.
Makes sure no sampling request waits forever, stacked on #381 (only the last commit is new). Before this, a pool whose Workloads or dispatcher were deleted never came back unless the client sampled again, and only the pool's own dispatcher enforced the deadline, so with no dispatcher nothing failed its requests.
Two tasks next to the session reaper now tend the pool streams every 30 s. A reconcile recreates a missing pool for a job with a live session, or fails the job's requests when it has none. A deadline sweep fails any request older than
OPEN_RL_SAMPLE_DEADLINE_SECONDS(default 3600).delete_modelfails the job's pending requests under the same owner lock, so the reconcile never brings a deleted job's pool back. Stream entries also carry the job and revision as their own fields, so these paths read them without decoding prompts.Tested on GKE with the rest of the stack: a pool deleted mid-run came back in about 2.5 minutes and its runs finished.