rlightning.weights.weight_transfer_manager

class rlightning.weights.weight_transfer_manager.WeightTransferManager(buffer_strategy: str)

Bases: object

clear_cache()

record_gpu_memory(record_name: str)

recv_weights(weight_buffer, metadata_by_dtype)

Receive weights from the train policy, grouped by dtype. Delegates the buffering and application logic to the configured WeightBuffer. This method is called by policies that do NOT use the shared buffer actor.

recv_weights_local(weight_buffer, state_dict)

Receive weights directly for single process mode (non-remote).

send_weights(state_dict: Dict[str, Dict[str, MockModule('torch.Tensor')]], eval_policies: List[MockModule('ray.actor.ActorHandle')])

Broadcasts the state dict to the eval policies.

This function is called on the TrainPolicy under sync mode. It dispatches to the correct sending strategy based on the eval policies’ configuration.

Parameters:

• state_dict (Dict[str, Dict[str, torch.Tensor]]) – A dict of module state dicts.

• eval_policies (List[ActorHandle]) – A list of Actors which related to policies for evaluation

Raises:

ValueError – Invalid buffer strategy

send_weights_cpu(state_dict: Dict[str, Dict[str, MockModule('torch.Tensor')]], cpu_weight_buffer: MockModule('ray.actor.ActorHandle'))

Sends weights to the CPU weight buffer actor.