Repository navigation
Conversation
ba2bd44 to
4adf8d0
Compare
| # Triton kernels persist on the shared volume, like vLLM's compile cache, so later | ||
| # replicas skip compiling them. Triton renames each kernel into place from a unique | ||
| # temp dir, so concurrent writers are safe. | ||
| env["TRITON_CACHE_DIR"] = str(Path(os.getenv("OPEN_RL_TMP_DIR", "/tmp/open-rl"), "cache", "triton")) |
There was a problem hiding this comment.
check if this is inline with the caching changes on the main @hodaaaaaaaaaa did as part of performance optimizations.
|
Question: pending records of a reaped or deleted job. Records leave the stream only when a result is saved (sampling_store.py:83). With the shutdown sentinel gone, delete_model and release_owner (scheduler_worker_manager.py:450) tell the pool nothing. For an owned pool, the stream key outlives the job in a Redis that never evicts. For a shared LoRA pool, the next job's dispatcher serves the dead job's leftovers first, then fails those older than an hour. The PR starts writing a job field on every record (api_server.py:1256) and adds job_of (sampling_store.py:38), but nothing reads either. |
|
minor: |
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.
4adf8d0 to
b64bb34
Compare
hodaaaaaaaaaa
left a comment
hodaaaaaaaaaa
left a comment
There was a problem hiding this comment.
I just realized these are getting merged into main. Is it possible to put the sampling pools work on a branch that we can merge into main all at once when it's ready? That way it'll be easier to roll it back if it does break things before our demo.
|
Yea I can consolidate the changes into one single branch, I'll also post numbers for the benchmarks to ensure no regressions before merge |
Switches sampling over to the pools, stacked on #357 (only the last commit is new).
asamplenow adds each request to its pool's Redis stream and returns, and the pool's dispatcher saves exactly one result per request, so a sampler that stops mid-batch loses nothing. The old queue sampler and its leftover patch test are deleted.A pool that one job owns runs in process at one replica, vLLM in the sampler's own process as on main, and takes one time-slicer turn per batch when it shares its trainer's GPU (
TimeSlicedDispatcher). With two or more replicas it getsvllm serveWorkloads behind a dispatcher pod with Envoy and the llm-d EPP. A shared LoRA pool is always routed, even at one replica, so it grows to the largestsampler_replicasany of its jobs asks for and a second job never has to replace an in-process sampler another job is using. On an FFT deployment, a session opened on a bare base-model name gets an FFT pool namedbase-<model>, apart from the base model's shared LoRA pool.Teardown also finds a routed pool's dispatcher Deployment and NetworkPolicy by an
openrl.io/ownerlabel, so the worker launcher's Role gainsliston both. Tested on GKE with the rest of the stack: a shared pool growing from one replica to two as jobs arrive, time-sliced and routed FFT, exclusive LoRA, and a replica killed mid-run.