VLLMDoubleBufferWeightReceiver

class torchrl.weight_update.llm.VLLMDoubleBufferWeightReceiver(scheme: VLLMDoubleBufferSyncScheme, vllm_engine)

Receives weights in a vLLM worker using double-buffered storage.

This receiver reads weights from a shared directory and loads them into the vLLM engine using the engine’s load_weights interface.

Example

>>> receiver = scheme.create_receiver(vllm_engine)
>>>
>>> # Poll for new weights
>>> if receiver.poll_and_apply():
...     print("Weights updated!")

apply_weights(weights: TensorDict) → None

Apply weights to vLLM engine using RPC.

This method uses RPC to tell all vLLM workers to load weights from the shared storage directory. Similar to how AsyncVLLM._update_weights_with_nccl_broadcast_simple uses collective_rpc to coordinate workers.

Parameters:

weights – TensorDict with flattened keys containing weights.

poll_and_apply(timeout: float = 180.0) → bool

Poll for and apply weights from shared storage.

Parameters:

timeout – Not used for file-based transport (kept for API compatibility).

Returns:

True if weights were successfully read and applied, False otherwise.

register_model(model_ref: Any) → None

Register the model to apply weights to.

Parameters:

model_ref – Either a direct object reference or a string path like ‘policy’ or ‘env.value_net’.

register_worker_transport(pipe: Any) → None

Register this worker’s communication pipe.

Parameters:

pipe – The pipe connection for this worker.

set_context(context: Any) → None

Set the context object (inner_collector) for resolving references.

Parameters:

context – The inner collector instance in the worker process.