nemo_rl.utils.weight_transfer_zmq#

Transactional ZeroMQ value plane for remote sparse vLLM refit.

Module Contents#

Classes#

_RelayTransfer

ZmqSparseRefitClient

One-thread DEALER client with retry-safe payload identifiers.

ZmqSparseRefitServer

Bounded ROUTER relay that applies locally and fans out through a tree.

Functions#

_json_bytes

_configure_socket

stream_sparse_delta_payloads_via_zmq

Data#

_PROTOCOL

_DATA

_ACK

_NACK

API#

nemo_rl.utils.weight_transfer_zmq._PROTOCOL#

‘nemo-rl-sparse-zmq-v1’

nemo_rl.utils.weight_transfer_zmq._DATA#

b’DATA’

nemo_rl.utils.weight_transfer_zmq._ACK#

b’ACK’

nemo_rl.utils.weight_transfer_zmq._NACK#

b’NACK’

class nemo_rl.utils.weight_transfer_zmq._RelayTransfer#
checksums: dict[tuple[int, int], str]#

‘field(…)’

futures: list[concurrent.futures.Future[dict[str, Any]]]#

‘field(…)’

nemo_rl.utils.weight_transfer_zmq._json_bytes(value: collections.abc.Mapping[str, Any]) → bytes#
nemo_rl.utils.weight_transfer_zmq._configure_socket(socket: zmq.Socket, high_water_mark: int) → None#
class nemo_rl.utils.weight_transfer_zmq.ZmqSparseRefitClient(
address: str,
*,
timeout_s: float,
producer_id: int,
retries: int,
api_key: str | None = None,
)#

One-thread DEALER client with retry-safe payload identifiers.

Initialization

send_payload(
*,
transfer_id: str,
payload_id: int,
checksum: str,
verification_candidates: int,
body: bytes,
relay_root: str | None = None,
producer_id: int | None = None,
) → dict[str, Any]#
close() → None#
class nemo_rl.utils.weight_transfer_zmq.ZmqSparseRefitServer(
apply_payload: collections.abc.Callable[[bytes, collections.abc.Mapping[str, Any]], dict[str, Any]],
*,
bind_address: str,
api_key_env_var: str | None,
timeout_s: float,
tuning: nemo_rl.models.generation.vllm.config.VllmRefitTuningConfig,
)#

Bounded ROUTER relay that applies locally and fans out through a tree.

Initialization

configure_tree(
relay_addresses: collections.abc.Sequence[str],
*,
own_address: str,
) → None#
start() → str#
close() → None#
flush(
transfer_id: str,
expected_payloads: int = 0,
) → dict[str, Any]#

Wait for every staged fanout belonging to one transfer.

_fanout(
body: bytes,
metadata: collections.abc.Mapping[str, Any],
) → dict[str, Any]#
_forward(
body: bytes,
metadata: collections.abc.Mapping[str, Any],
address: str,
relay_root: str,
) → dict[str, Any]#
static _send_reply(
socket: zmq.Socket,
identity: bytes,
kind: bytes,
reply: collections.abc.Mapping[str, Any],
) → None#
_parse_data_message(
frames: list[bytes],
) → tuple[bytes, tuple[str, int, int], bytes, dict[str, Any]]#
_run() → None#
nemo_rl.utils.weight_transfer_zmq.stream_sparse_delta_payloads_via_zmq(
iterator: collections.abc.Iterable[nemo_rl.utils.weight_transfer_sparse_codec.NamedTensor],
*,
delta_tracker: nemo_rl.utils.weight_transfer_sparse_codec.DeltaCompressionTracker,
refit_targets: collections.abc.Sequence[str],
transfer_id: str,
api_key_env_var: str | None,
timeout_s: float,
shard_rank: int,
shard_count: int,
) → dict[str, int]#