nemo_rl.models.generation.dynamo.refit#

NCCL refit protocol for a fixed managed Dynamo worker fleet.

Module Contents#

Classes#

DynamoWorkerEndpoint

Serializable identity and admin endpoint for one Dynamo vLLM engine.

DynamoRefitChannel

Closed refit protocol shared by driver and serialized rollout copies.

Functions#

_post_worker_route

_update_worker_weights

API#

class nemo_rl.models.generation.dynamo.refit.DynamoWorkerEndpoint#

Serializable identity and admin endpoint for one Dynamo vLLM engine.

instance_id: str#

None

system_url: str#

None

classmethod from_metadata(
metadata: dict[str, Any],
) → nemo_rl.models.generation.dynamo.refit.DynamoWorkerEndpoint#
nemo_rl.models.generation.dynamo.refit._post_worker_route(
*,
system_url: str,
route: str,
payload: dict[str, Any],
timeout_s: float,
) → bool#
nemo_rl.models.generation.dynamo.refit._update_worker_weights(
*,
system_url: str,
update_info: dict[str, Any],
timeout_s: float,
) → bool#
class nemo_rl.models.generation.dynamo.refit.DynamoRefitChannel(
workers: collections.abc.Sequence[dict[str, Any] | nemo_rl.models.generation.dynamo.refit.DynamoWorkerEndpoint],
*,
engine_world_size: int,
control_timeout_s: float,
validate_workers: collections.abc.Callable[[list[dict[str, Any]]], list[dict[str, Any]]] | None = None,
)#

Closed refit protocol shared by driver and serialized rollout copies.

Initialization

client_copy() → nemo_rl.models.generation.dynamo.refit.DynamoRefitChannel#

Return a serializable endpoint-only channel for rollout actors.

_validated_workers() → tuple[nemo_rl.models.generation.dynamo.refit.DynamoWorkerEndpoint, ...]#
property inference_world_size: int#
property sender_spec: nemo_rl.models.generation.interfaces.CollectiveSenderSpec#

Return vLLM’s non-negotiated peer and packing contract.

prepare(state_dict_info: dict[str, Any] | None) → None#
init_collective(
ip: str,
port: int,
world_size: int,
*,
train_world_size: int,
) → list[ray.ObjectRef]#
update_weights() → list[ray.ObjectRef]#
flush_cache() → bool#

Drain, clear, and resume every worker using immutable endpoints.