diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 0446bd9b..1fc5d1a3 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -397,9 +397,10 @@ identity or ordering key. `WeightVersionShard` remains the name of the per-worker manifest publication. Its identity is `(version_id, worker_id, source_slot_id)`: `source_slot_id` identifies the required, version-scoped source contribution it covers, and -`worker_id` identifies the publishing process. The trainer coordinator chooses -the opaque slots—for example, -`publisher:global-rank:12` for a selected Megatron publisher. Multiple +`worker_id` identifies the publishing process. The trainer engine adapter +derives the slot from its native topology; the Megatron adapter uses +`publisher:global-rank:12` for global rank 12. The orchestrator uses the same +adapter-defined convention when declaring the version's expected slots. Multiple publications may advertise the same source slot, including a replacement worker or a generator that becomes a peer source. Deployments configured with Kubernetes or the test-only memory backend do not expose `RefitService` yet. diff --git a/modelexpress_client/python/README.md b/modelexpress_client/python/README.md index f2b56270..39807206 100644 --- a/modelexpress_client/python/README.md +++ b/modelexpress_client/python/README.md @@ -18,7 +18,7 @@ pip install -e . # With test dependencies pip install -e ".[dev]" -# Additionally install the pinned protobuf code generator when changing p2p.proto +# Additionally install the pinned protobuf code generator when changing protobuf APIs pip install -e ".[codegen]" ``` @@ -99,6 +99,70 @@ deployment. ## Programmatic Usage +### RL trainer publication + +An RL framework creates a weight version through the external Refit API. Each +trainer actor then invokes its rank-local client to stage and publish one shard. +Worker registration, manifest serving, and internal shard CRUD remain hidden +behind the client. + +```python +from modelexpress_rl import ( + ModelExpressTrainerClient, + WeightVersionRef, + WeightVersionShardManifestService, + refit_pb2_grpc, +) + +manifest_service = WeightVersionShardManifestService(endpoint="trainer-0:9000") +refit_pb2_grpc.add_RefitWorkerServiceServicer_to_server( + manifest_service, + trainer_worker_grpc_server, +) + +trainer = ModelExpressTrainerClient.initialize( + manager=nixl_manager, + manifest_publisher=manifest_service, +) + +shard = trainer.stage_shard( + version=WeightVersionRef(version.uid), + tensors=megatron_tensor_specs, +) +shard.publish() +``` + +The deployment supplies `MODEL_NAME`, `MX_TRAINER_ENGINE`, +`MX_TRAINER_STAGING_MODE`, `MX_WEIGHT_PAYLOAD_FORMAT`, `MX_WORKER_HOST`, and the +normal ModelExpress server configuration. The Megatron adapter derives its +source slot from the engine's global distributed rank. The NIXL metadata +endpoint is derived from `MX_WORKER_HOST` and the supplied NIXL manager's listen +port. + +`worker_endpoint` is the trainer-side manifest service address advertised to +other workers. `server_url` selects the central ModelExpress control-plane +service and defaults to the normal ModelExpress server configuration. + +Initialization fixes the staging mode and payload format. `publish()` hides +manifest publication and the internal `CreateWeightVersionShard` RPC. The +current Megatron adapter exposes its already-registered live buffers through +`IN_PLACE`, so callers must keep those tensors immutable while the version is +published. Its `source_reuse_ready` fence raises `NotImplementedError` until +version retirement is wired to the adapter; it must not be interpreted as an +early reuse signal. The required lifecycle is synchronous: create and publish +the version, update every generator, retire the version, and only then resume +training or begin the next optimizer step. The adapter does not claim fully +asynchronous `COPY_TO_DEVICE` behavior until that staging implementation exists. + +Version creation and expected-source-slot declaration remain +framework-orchestrator responsibilities. Each trainer adapter derives its own +source slot from the engine's native topology; the orchestrator declares the +expected slots using the same adapter-defined convention. `initialize()` +selects the configured trainer engine and constructs its adapter internally; +Megatron is the first implementation. Megatron-specific APIs live under +`modelexpress_rl`; +`modelexpress.refit.reshard` remains the shared, engine-neutral transfer core. + ### MxClient `MxClient` is a lightweight gRPC client for communicating with the ModelExpress server: diff --git a/modelexpress_client/python/generate_proto.sh b/modelexpress_client/python/generate_proto.sh index 6b7cf14b..e8062def 100755 --- a/modelexpress_client/python/generate_proto.sh +++ b/modelexpress_client/python/generate_proto.sh @@ -7,32 +7,37 @@ set -euo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" PROTO_DIR="${SCRIPT_DIR}/../../modelexpress_common/proto" OUT_DIR="${SCRIPT_DIR}/modelexpress" +RL_OUT_DIR="${SCRIPT_DIR}/modelexpress_rl" YEAR="$(date +%Y)" SPDX_HEADER="# SPDX-FileCopyrightText: Copyright (c) 2025-${YEAR} NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 #" -PROTOS=(p2p model) - -for proto in "${PROTOS[@]}"; do - # Generate protobuf files +# Generate protobuf files. Keep the inference and RL surfaces in separate +# Python modules even though they are built from the same proto directory. +for package_proto in "${OUT_DIR}:p2p" "${OUT_DIR}:model" "${RL_OUT_DIR}:refit"; do + package_dir="${package_proto%%:*}" + proto="${package_proto##*:}" echo "Generating protobuf files from ${PROTO_DIR}/${proto}.proto..." python -m grpc_tools.protoc \ "-I${PROTO_DIR}" \ - "--python_out=${OUT_DIR}" \ - "--grpc_python_out=${OUT_DIR}" \ + "--python_out=${package_dir}" \ + "--grpc_python_out=${package_dir}" \ "${PROTO_DIR}/${proto}.proto" - # Fix relative import in grpc file + # Fix relative imports in gRPC files. + grpc_file="${package_dir}/${proto}_pb2_grpc.py" echo "Fixing imports in ${proto}_pb2_grpc.py..." tmp_file="$(mktemp)" - sed "s/^import ${proto}_pb2 as/from . import ${proto}_pb2 as/" \ - "${OUT_DIR}/${proto}_pb2_grpc.py" > "${tmp_file}" - mv "${tmp_file}" "${OUT_DIR}/${proto}_pb2_grpc.py" - - # Add SPDX header to generated files - for file in "${OUT_DIR}/${proto}_pb2.py" "${OUT_DIR}/${proto}_pb2_grpc.py"; do + sed \ + -e "s/^import ${proto}_pb2 as/from . import ${proto}_pb2 as/" \ + -e "s/^ + f' but the generated code/ + ' but the generated code/" \ + "${grpc_file}" > "${tmp_file}" + mv "${tmp_file}" "${grpc_file}" + + # Add SPDX headers to generated files. + for file in "${package_dir}/${proto}_pb2.py" "${package_dir}/${proto}_pb2_grpc.py"; do echo "Adding SPDX header to ${file}..." tmp_file=$(mktemp) echo "${SPDX_HEADER}" > "${tmp_file}" diff --git a/modelexpress_client/python/modelexpress/model_pb2_grpc.py b/modelexpress_client/python/modelexpress/model_pb2_grpc.py index 7e125833..d9e9d9de 100644 --- a/modelexpress_client/python/modelexpress/model_pb2_grpc.py +++ b/modelexpress_client/python/modelexpress/model_pb2_grpc.py @@ -21,7 +21,7 @@ if _version_not_supported: raise RuntimeError( f'The grpc package installed is at version {GRPC_VERSION},' - + f' but the generated code in model_pb2_grpc.py depends on' + + ' but the generated code in model_pb2_grpc.py depends on' + f' grpcio>={GRPC_GENERATED_VERSION}.' + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}' + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.' diff --git a/modelexpress_client/python/modelexpress/nixl_transfer.py b/modelexpress_client/python/modelexpress/nixl_transfer.py index 83c50b46..432fa8d2 100644 --- a/modelexpress_client/python/modelexpress/nixl_transfer.py +++ b/modelexpress_client/python/modelexpress/nixl_transfer.py @@ -167,6 +167,11 @@ def nixl_metadata(self) -> bytes: """Get NIXL metadata for this agent.""" return self._metadata + @property + def listen_port(self) -> int | None: + """Get the port serving this agent's NIXL metadata.""" + return self._listen_port + @property def tensor_descriptors(self) -> list[TensorDescriptor]: """Get tensor descriptors for registered tensors.""" diff --git a/modelexpress_client/python/modelexpress/p2p_pb2_grpc.py b/modelexpress_client/python/modelexpress/p2p_pb2_grpc.py index 5f8ee047..0687c49f 100644 --- a/modelexpress_client/python/modelexpress/p2p_pb2_grpc.py +++ b/modelexpress_client/python/modelexpress/p2p_pb2_grpc.py @@ -21,7 +21,7 @@ if _version_not_supported: raise RuntimeError( f'The grpc package installed is at version {GRPC_VERSION},' - + f' but the generated code in p2p_pb2_grpc.py depends on' + + ' but the generated code in p2p_pb2_grpc.py depends on' + f' grpcio>={GRPC_GENERATED_VERSION}.' + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}' + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.' diff --git a/modelexpress_client/python/modelexpress/refit/reshard/__init__.py b/modelexpress_client/python/modelexpress/refit/reshard/__init__.py index e9fccd13..c69df8fe 100644 --- a/modelexpress_client/python/modelexpress/refit/reshard/__init__.py +++ b/modelexpress_client/python/modelexpress/refit/reshard/__init__.py @@ -47,21 +47,6 @@ ) from modelexpress.refit.reshard.cuda_pool import classic_cuda_alloc from modelexpress.refit.reshard.receiver import ReshardReceiver -from modelexpress.refit.reshard.megatron import ( - MegatronTargetLayout, - MegatronTargetSpec, - lower_megatron_target, -) -from modelexpress.refit.reshard.megatron_receiver import MegatronReshardReceiver -from modelexpress.refit.reshard.megatron_aliases import ( - MegatronAliasInput, - build_hf_aliases, -) -from modelexpress.refit.reshard.megatron_publisher import ( - MegatronPublishedTensorSpec, - publish_megatron_reshard_view, - publish_registered_shard_table, -) from modelexpress.refit.reshard.rendezvous import ( MxReshardRendezvous, PublishedShard, @@ -77,11 +62,6 @@ "FullPullSource", "IncompleteRefit", "LazyWeight", - "MegatronAliasInput", - "MegatronPublishedTensorSpec", - "MegatronReshardReceiver", - "MegatronTargetLayout", - "MegatronTargetSpec", "MxReshardRendezvous", "NixlReshardTransport", "OpChain", @@ -97,19 +77,15 @@ "Transport", "TransferPlan", "UnsupportedReshard", - "build_hf_aliases", "capture_geometry", "classic_cuda_alloc", "execute_transfer", "gather_sources", "intersect", - "lower_megatron_target", "op_chain_to_box", "paired_runs", "plan_pull", "plan_transfer", - "publish_megatron_reshard_view", - "publish_registered_shard_table", "shard_region", "tensor_digest", "wrap_rendezvous_blob", diff --git a/modelexpress_client/python/modelexpress/refit/reshard/megatron_aliases.py b/modelexpress_client/python/modelexpress/refit/reshard/megatron_aliases.py index 18a3971b..66b74a2a 100644 --- a/modelexpress_client/python/modelexpress/refit/reshard/megatron_aliases.py +++ b/modelexpress_client/python/modelexpress/refit/reshard/megatron_aliases.py @@ -1,248 +1,11 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Expose native Megatron storage as HF-canonical reshard source shards.""" +"""Compatibility imports for Megatron source aliases.""" -from __future__ import annotations +from modelexpress_rl.train.engines.megatron.aliases import ( + MegatronAliasInput, + MegatronTensorSpec, + build_hf_aliases, +) -from dataclasses import dataclass, field -from typing import Any - -from modelexpress.refit.reshard.rendezvous import PublishedShard, PublishedTensor -from modelexpress.refit.reshard.verify import published_digest - - -@dataclass(frozen=True) -class MegatronAliasInput: - name: str - tensor: Any - role: str - hf_names: tuple[str, ...] - global_shape: tuple[int, ...] - placement_kind: str - shard_axis: int | None - local_shard_range: tuple[int, int] | None - extras: dict[str, str] = field(default_factory=dict) - - -def _source_rank_and_size(item: MegatronAliasInput, axis: int) -> tuple[int, int]: - local_extent = int(item.tensor.shape[axis]) - global_extent = int(item.global_shape[axis]) - if item.placement_kind != "SHARD": - return 0, 1 - if item.local_shard_range is None: - raise ValueError(f"{item.name}: SHARD has no local range") - lo, hi = (int(value) for value in item.local_shard_range) - # A range can pass every check below and still lie outside the tensor it - # claims part of: (16, 24) against a global extent of 16 has the right width, - # divides evenly, and yields source rank 2 of a 2-rank group. The alias that - # follows would then address bytes the full tensor does not have. - if not 0 <= lo < hi <= global_extent: - raise ValueError( - f"{item.name}: source shard range {(lo, hi)} is outside the global " - f"extent {global_extent} on axis {axis}" - ) - if hi - lo != local_extent or global_extent % local_extent: - raise ValueError(f"{item.name}: inconsistent source shard geometry") - if lo % local_extent: - raise ValueError(f"{item.name}: non-uniform source shard is unsupported") - return lo // local_extent, global_extent // local_extent - - -def _one_shard( - *, - name: str, - tensor: Any, - full_shape: tuple[int, ...], - agent_name: str, - shard_axis: int | None, - shard_range: tuple[int, int] | None, -) -> PublishedTensor: - local_shape = tuple(int(dim) for dim in tensor.shape) - offset = [0] * len(local_shape) - if shard_axis is not None: - if shard_range is None: - raise ValueError(f"{name}: shard axis has no range") - lo, hi = shard_range - if hi - lo != local_shape[shard_axis]: - raise ValueError(f"{name}: shard range does not match local shape") - offset[shard_axis] = lo - elif local_shape != full_shape: - raise ValueError(f"{name}: replicated shape mismatch") - return PublishedTensor( - name=name, - dtype=str(tensor.dtype), - elsize=int(tensor.element_size()), - full_shape=full_shape, - shards=[ - PublishedShard( - agent_name=agent_name, - device_id=int(tensor.device.index or 0), - addr=int(tensor.data_ptr()), - shard_offset=tuple(offset), - shape=local_shape, - digest=published_digest(tensor), - ) - ], - ) - - -GATE_THEN_UP = "gate_then_up" - - -def _build_gated_aliases( - item: MegatronAliasInput, agent_name: str -) -> list[PublishedTensor]: - if len(item.hf_names) != 2: - raise ValueError(f"{item.name}: gated tensor requires gate/up HF names") - # The halves are assigned to hf_names positionally, so the storage order is - # what decides which HF tensor each half becomes. Getting it wrong publishes - # the gate projection's bytes under the up projection's name, which no digest - # gate can see: both names receive the bytes their publisher advertised. The - # order is therefore required rather than assumed. - order = item.extras.get("gated_mlp_order") - if order != GATE_THEN_UP: - raise ValueError( - f"{item.name}: fused gate/up aliasing requires extras" - f"['gated_mlp_order'] == {GATE_THEN_UP!r}, got {order!r}" - ) - axis = int(item.shard_axis if item.shard_axis is not None else 0) - local_extent = int(item.tensor.shape[axis]) - if local_extent % 2: - raise ValueError(f"{item.name}: fused gate/up extent must be even") - half = local_extent // 2 - gate = item.tensor.narrow(axis, 0, half) - up = item.tensor.narrow(axis, half, half) - if not gate.is_contiguous() or not up.is_contiguous(): - raise ValueError( - f"{item.name}: fused gate/up aliases are not contiguous on axis {axis}" - ) - source_rank, source_size = _source_rank_and_size(item, axis) - full_shape = list(item.tensor.shape) - full_shape[axis] = half * source_size - shard_range = ( - (source_rank * half, (source_rank + 1) * half) if source_size > 1 else None - ) - return [ - _one_shard( - name=hf_name, - tensor=tensor, - full_shape=tuple(int(dim) for dim in full_shape), - agent_name=agent_name, - shard_axis=axis if shard_range is not None else None, - shard_range=shard_range, - ) - for hf_name, tensor in zip(item.hf_names, (gate, up), strict=True) - ] - - -def _build_qkv_aliases( - item: MegatronAliasInput, agent_name: str -) -> list[PublishedTensor]: - if len(item.hf_names) != 3 or item.tensor.ndim != 2: - raise ValueError(f"{item.name}: QKV aliasing requires 2D q/k/v weights") - head_dim = int(item.extras["head_dim"]) - q_heads_local = int(item.extras["num_heads_local"]) - kv_heads_local = int(item.extras["num_kv_heads_local"]) - if kv_heads_local < 1 or q_heads_local % kv_heads_local: - raise ValueError(f"{item.name}: invalid local Q/KV head geometry") - rows_per_group = (q_heads_local // kv_heads_local + 2) * head_dim - if rows_per_group * kv_heads_local != int(item.tensor.shape[0]): - raise ValueError(f"{item.name}: QKV rows disagree with head metadata") - source_rank, source_size = _source_rank_and_size(item, 0) - hidden = int(item.tensor.shape[1]) - q_heads_per_group = q_heads_local // kv_heads_local - q_shards = [] - k_shards = [] - v_shards = [] - for local_group in range(kv_heads_local): - group = item.tensor.narrow(0, local_group * rows_per_group, rows_per_group) - q_rows = q_heads_per_group * head_dim - q = group.narrow(0, 0, q_rows) - k = group.narrow(0, q_rows, head_dim) - v = group.narrow(0, q_rows + head_dim, head_dim) - global_group = source_rank * kv_heads_local + local_group - for tensor, shards, start in ( - (q, q_shards, global_group * q_rows), - (k, k_shards, global_group * head_dim), - (v, v_shards, global_group * head_dim), - ): - shards.append( - PublishedShard( - agent_name=agent_name, - device_id=int(tensor.device.index or 0), - addr=int(tensor.data_ptr()), - shard_offset=(start, 0), - shape=tuple(int(dim) for dim in tensor.shape), - # The narrow, not the fused parent: this is the box a receiver - # reads from ``addr``, so it is the box whose bytes must match. - digest=published_digest(tensor), - ) - ) - return [ - PublishedTensor( - name=name, - dtype=str(item.tensor.dtype), - elsize=int(item.tensor.element_size()), - full_shape=(rows, hidden), - shards=shards, - ) - for name, rows, shards in ( - (item.hf_names[0], q_heads_local * source_size * head_dim, q_shards), - (item.hf_names[1], kv_heads_local * source_size * head_dim, k_shards), - (item.hf_names[2], kv_heads_local * source_size * head_dim, v_shards), - ) - ] - - -def build_hf_aliases( - items: list[MegatronAliasInput], *, agent_name: str -) -> list[PublishedTensor]: - """Build zero-copy HF aliases whose addresses remain in registered storage.""" - - aliases = [] - for item in items: - # Every alias published below is a base address plus a shape, which tells - # a reader the bytes run contiguously from that address. A strided view - # satisfies neither, and nothing downstream can detect it: the read simply - # lands on whatever sits between the elements the view meant to select. - if not item.tensor.is_contiguous(): - raise ValueError( - f"{item.name}: aliasing publishes an address and a shape, which " - f"requires contiguous storage; got a non-contiguous tensor of " - f"shape {tuple(int(dim) for dim in item.tensor.shape)}" - ) - if item.role == "qkv_column": - aliases.extend(_build_qkv_aliases(item, agent_name)) - continue - if ( - item.role in {"gated_mlp_column", "expert_column"} - and len(item.hf_names) == 2 - ): - aliases.extend(_build_gated_aliases(item, agent_name)) - continue - if len(item.hf_names) != 1: - raise ValueError( - f"{item.name}: role {item.role!r} cannot map to " - f"{len(item.hf_names)} HF tensors" - ) - aliases.append( - _one_shard( - name=item.hf_names[0], - tensor=item.tensor, - full_shape=tuple(item.global_shape), - agent_name=agent_name, - shard_axis=( - int(item.shard_axis) if item.placement_kind == "SHARD" else None - ), - shard_range=( - tuple(item.local_shard_range) - if item.placement_kind == "SHARD" - and item.local_shard_range is not None - else None - ), - ) - ) - return aliases - - -__all__ = ["MegatronAliasInput", "build_hf_aliases"] +__all__ = ["MegatronAliasInput", "MegatronTensorSpec", "build_hf_aliases"] diff --git a/modelexpress_client/python/modelexpress/refit/reshard/megatron_publisher.py b/modelexpress_client/python/modelexpress/refit/reshard/megatron_publisher.py index 7ba7d21a..f92f2885 100644 --- a/modelexpress_client/python/modelexpress/refit/reshard/megatron_publisher.py +++ b/modelexpress_client/python/modelexpress/refit/reshard/megatron_publisher.py @@ -1,176 +1,19 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Publish registered native Megatron tensors through the reshard rendezvous.""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any - -from modelexpress.refit.reshard.rendezvous import ( - MxReshardRendezvous, - PublishedShard, - PublishedTensor, - wrap_rendezvous_blob, +"""Compatibility imports for Megatron source publication.""" + +from modelexpress_rl.train.engines.megatron.publisher import ( + MegatronPublishedTensorSpec, + MegatronReshardManifest, + build_megatron_reshard_manifest, + publish_megatron_reshard_view, + publish_registered_shard_table, ) -from modelexpress.refit.reshard.verify import published_digest - - -@dataclass(frozen=True) -class MegatronPublishedTensorSpec: - name: str - global_shape: tuple[int, ...] - shard_axis: int | None = None - local_shard_range: tuple[int, int] | None = None - - def __post_init__(self) -> None: - if not self.global_shape or any(int(dim) <= 0 for dim in self.global_shape): - raise ValueError(f"{self.name}: invalid global shape {self.global_shape}") - if (self.shard_axis is None) != (self.local_shard_range is None): - raise ValueError( - f"{self.name}: shard_axis and local_shard_range must be set together" - ) - - -def publish_megatron_reshard_view( - *, - manager: Any, - rendezvous: MxReshardRendezvous, - tensors: dict[str, Any], - specs: list[MegatronPublishedTensorSpec], - metadata_endpoint: str, -) -> str: - """Publish a reshard shard table over an existing NIXL registration. - - The framework's normal publisher owns tensor lifetime, registration, and - listen thread. This seam only describes those same stable addresses in the - reshard rendezvous format, avoiding a duplicate NIXL agent or second tensor - allocation. - """ - - if not metadata_endpoint or ":" not in metadata_endpoint: - raise ValueError( - "metadata_endpoint must be an explicit host:port reachable by receivers" - ) - by_name: dict[str, MegatronPublishedTensorSpec] = {} - for spec in specs: - # Last-writer-wins here would publish one spec's shard description under a - # name the other spec owns, and the missing/extra check below compares key - # sets, so it cannot see it. lower_megatron_target rejects the analogous - # duplicate staging_name for the same reason. - if spec.name in by_name: - raise ValueError(f"duplicate Megatron publish spec for {spec.name!r}") - by_name[spec.name] = spec - missing = sorted(set(by_name).difference(tensors)) - extra = sorted(set(tensors).difference(by_name)) - if missing or extra: - raise ValueError( - f"Megatron shard table/tensor mismatch: missing={missing[:10]} " - f"extra={extra[:10]}" - ) - - agent_name = str(manager.agent_name) - published = [] - for name in sorted(by_name): - tensor = tensors[name] - spec = by_name[name] - if not tensor.is_contiguous(): - raise ValueError(f"{name}: reshard publication requires contiguous storage") - local_shape = tuple(int(dim) for dim in tensor.shape) - if len(local_shape) != len(spec.global_shape): - raise ValueError( - f"{name}: local rank {len(local_shape)} != global rank " - f"{len(spec.global_shape)}" - ) - offset = [0] * len(local_shape) - if spec.shard_axis is not None: - axis = int(spec.shard_axis) - if not 0 <= axis < len(local_shape): - raise ValueError(f"{name}: invalid shard axis {axis}") - assert spec.local_shard_range is not None - lo, hi = (int(value) for value in spec.local_shard_range) - if hi - lo != local_shape[axis] or hi > int(spec.global_shape[axis]): - raise ValueError( - f"{name}: local shape {local_shape} disagrees with shard " - f"range {(lo, hi)} in global shape {spec.global_shape}" - ) - offset[axis] = lo - elif local_shape != tuple(spec.global_shape): - raise ValueError( - f"{name}: replicated local shape {local_shape} != " - f"global shape {spec.global_shape}" - ) - published.append( - PublishedTensor( - name=name, - dtype=str(tensor.dtype), - elsize=int(tensor.element_size()), - full_shape=tuple(spec.global_shape), - shards=[ - PublishedShard( - agent_name=agent_name, - device_id=int(tensor.device.index or 0), - addr=int(tensor.data_ptr()), - shard_offset=tuple(offset), - shape=local_shape, - digest=published_digest(tensor), - ) - ], - ) - ) - - return publish_registered_shard_table( - manager=manager, - rendezvous=rendezvous, - published=published, - metadata_endpoint=metadata_endpoint, - ) - - -def publish_registered_shard_table( - *, - manager: Any, - rendezvous: MxReshardRendezvous, - published: list[PublishedTensor], - metadata_endpoint: str, -) -> str: - """Publish a validated alias table over already-registered storage. - - The caller supplies the rendezvous and keeps it for the lifetime of the - publication: publishing starts the source's READY heartbeat thread, and only - the owner of the rendezvous can stop it and mark the source stale on shutdown. - Constructing one here would leave a heartbeat running with no handle to it. - """ - - if not metadata_endpoint or ":" not in metadata_endpoint: - raise ValueError( - "metadata_endpoint must be an explicit host:port reachable by receivers" - ) - if not published: - raise ValueError("published shard table must not be empty") - agent_name = str(manager.agent_name) - for tensor in published: - if not tensor.shards: - raise ValueError(f"{tensor.name}: no shards were published") - for shard in tensor.shards: - if shard.agent_name != agent_name: - raise ValueError( - f"{tensor.name}: shard agent {shard.agent_name!r} does not " - f"match manager agent {agent_name!r}" - ) - if shard.addr <= 0: - raise ValueError(f"{tensor.name}: shard has invalid address") - blob = wrap_rendezvous_blob( - manager.nixl_metadata, - agent_name, - metadata_endpoint, - published, - ) - return rendezvous.publish(blob) - __all__ = [ "MegatronPublishedTensorSpec", + "MegatronReshardManifest", + "build_megatron_reshard_manifest", "publish_megatron_reshard_view", "publish_registered_shard_table", ] diff --git a/modelexpress_client/python/modelexpress_rl/__init__.py b/modelexpress_client/python/modelexpress_rl/__init__.py new file mode 100644 index 00000000..6a7eeaad --- /dev/null +++ b/modelexpress_client/python/modelexpress_rl/__init__.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""ModelExpress clients and protobuf bindings for RL weight refit.""" + +from .client import ( + ModelExpressTrainerClient, + StagedWeightVersionShard, + WeightVersionRef, +) +from .train import ( + CompletionFence, + StagedWeightVersionShardData, + TrainerEngineAdapter, + TrainerStagingMode, + WeightPayloadFormat, + WeightVersionShardManifest, + WeightVersionShardManifestPublisher, + WeightVersionShardManifestService, +) + +__all__ = [ + "CompletionFence", + "ModelExpressTrainerClient", + "StagedWeightVersionShard", + "StagedWeightVersionShardData", + "TrainerEngineAdapter", + "TrainerStagingMode", + "WeightPayloadFormat", + "WeightVersionRef", + "WeightVersionShardManifest", + "WeightVersionShardManifestPublisher", + "WeightVersionShardManifestService", +] diff --git a/modelexpress_client/python/modelexpress_rl/client.py b/modelexpress_client/python/modelexpress_rl/client.py new file mode 100644 index 00000000..7740c3fe --- /dev/null +++ b/modelexpress_client/python/modelexpress_rl/client.py @@ -0,0 +1,312 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Framework-facing trainer lifecycle for ModelExpress RL refit.""" + +from __future__ import annotations + +import threading +import uuid +from dataclasses import dataclass +from typing import Any + +import grpc +from modelexpress import auth, envs +from modelexpress.client import _get_server_url + +from . import envs as rl_envs +from . import refit_pb2, refit_pb2_grpc +from .train.adapter import ( + CompletionFence, + NixlMetadataProvider, + StagedWeightVersionShardData, + TrainerEngineAdapter, + TrainerStagingMode, + WeightPayloadFormat, + WeightVersionShardManifestPublisher, +) + + +def _required(value: str, name: str) -> str: + if not value.strip(): + raise ValueError(f"{name} is required") + return value + + +def _trainer_adapter( + *, + manager: NixlMetadataProvider, + nixl_metadata_endpoint: str, +) -> TrainerEngineAdapter: + engine = rl_envs.MX_TRAINER_ENGINE + if engine != "MEGATRON": + raise ValueError(f"unsupported MX_TRAINER_ENGINE={engine!r}") + + from .train.engines.megatron import MegatronTrainerAdapter + + return MegatronTrainerAdapter( + manager=manager, + nixl_metadata_endpoint=nixl_metadata_endpoint, + ) + + +def _nixl_metadata_endpoint(manager: NixlMetadataProvider) -> str: + host = _required(envs.MX_WORKER_HOST, "MX_WORKER_HOST") + if manager.listen_port is None: + raise ValueError("NIXL manager must have a metadata listen port") + return f"{host}:{manager.listen_port}" + + +def _staging_mode(value: TrainerStagingMode | None) -> TrainerStagingMode: + try: + return value or TrainerStagingMode(rl_envs.MX_TRAINER_STAGING_MODE) + except ValueError as error: + raise ValueError( + f"invalid MX_TRAINER_STAGING_MODE={rl_envs.MX_TRAINER_STAGING_MODE!r}" + ) from error + + +def _payload_format(value: WeightPayloadFormat | None) -> WeightPayloadFormat: + try: + return value or WeightPayloadFormat(rl_envs.MX_WEIGHT_PAYLOAD_FORMAT) + except ValueError as error: + raise ValueError( + f"invalid MX_WEIGHT_PAYLOAD_FORMAT={rl_envs.MX_WEIGHT_PAYLOAD_FORMAT!r}" + ) from error + + +@dataclass(frozen=True) +class WeightVersionRef: + """Opaque reference to one global WeightVersion created by the orchestrator.""" + + version_id: str + + def __post_init__(self) -> None: + _required(self.version_id, "version.version_id") + + +class StagedWeightVersionShard: + """One immutable rank-local shard staged for a global weight version.""" + + def __init__( + self, + *, + client: ModelExpressTrainerClient, + version: WeightVersionRef, + staged: StagedWeightVersionShardData, + ) -> None: + self._client = client + self._version = version + self._staged = staged + self._publish_lock = threading.Lock() + self._published = False + + @property + def source_reuse_ready(self) -> CompletionFence: + """Fence after which the trainer may reuse or mutate its input tensors.""" + return self._staged.source_reuse_ready + + def publish(self) -> None: + """Publish this staged shard; repeated calls are idempotent.""" + with self._publish_lock: + if self._published: + return + self._client._publish_staged_shard( + version=self._version, + staged=self._staged, + ) + self._published = True + + +class ModelExpressTrainerClient: + """Rank-local capture, staging, and publication client for trainer actors.""" + + def __init__(self) -> None: + self._channel: grpc.Channel | None = None + self._stub: refit_pb2_grpc.RefitServiceStub | None = None + self._published_shards: dict[ + str, list[StagedWeightVersionShardData] + ] = {} + self._registration_stop = threading.Event() + self._registration_thread: threading.Thread | None = None + + @classmethod + def initialize( + cls, + *, + manager: NixlMetadataProvider, + manifest_publisher: WeightVersionShardManifestPublisher, + model_name: str | None = None, + staging_mode: TrainerStagingMode | None = None, + payload_format: WeightPayloadFormat | None = None, + worker_endpoint: str | None = None, + worker_id: str | None = None, + server_url: str | None = None, + registration_ttl_seconds: int | None = None, + rpc_timeout_seconds: float = 30.0, + ) -> ModelExpressTrainerClient: + """Initialize a trainer worker and connect it to the MX control plane. + + ``worker_endpoint`` is this trainer's peer-reachable manifest service. + ``server_url`` is the central ModelExpress ``RefitService`` address. + """ + model_name = _required(model_name or envs.MODEL_NAME or "", "model_name") + staging_mode = _staging_mode(staging_mode) + payload_format = _payload_format(payload_format) + nixl_metadata_endpoint = _nixl_metadata_endpoint(manager) + worker_endpoint = _required( + worker_endpoint + or ( + f"{envs.MX_WORKER_HOST}:{envs.MX_WORKER_GRPC_PORT}" + if envs.MX_WORKER_HOST + else "" + ), + "worker_endpoint", + ) + worker_id = _required(worker_id or uuid.uuid4().hex[:8], "worker_id") + if staging_mode is TrainerStagingMode.UNSPECIFIED: + raise ValueError("staging_mode must be specified") + if payload_format is WeightPayloadFormat.UNSPECIFIED: + raise ValueError("payload_format must be specified") + if registration_ttl_seconds is None: + registration_ttl_seconds = envs.MX_HEARTBEAT_INTERVAL_SECS * 3 + if registration_ttl_seconds <= 0: + raise ValueError("registration_ttl_seconds must be positive") + if rpc_timeout_seconds <= 0: + raise ValueError("rpc_timeout_seconds must be positive") + adapter = _trainer_adapter( + manager=manager, + nixl_metadata_endpoint=nixl_metadata_endpoint, + ) + if staging_mode not in adapter.supported_staging_modes: + raise ValueError( + f"adapter does not support staging mode {staging_mode.value}" + ) + if payload_format not in adapter.supported_payload_formats: + raise ValueError( + f"adapter does not support payload format {payload_format.value}" + ) + + client = cls() + client.model_name = model_name + client.staging_mode = staging_mode + client.payload_format = payload_format + client.worker_id = worker_id + client.worker_endpoint = worker_endpoint + client.server_url = _get_server_url(server_url) + client._adapter = adapter + client._manifest_publisher = manifest_publisher + client._registration_ttl_seconds = registration_ttl_seconds + client._rpc_timeout_seconds = rpc_timeout_seconds + client._register_worker() + client._registration_thread = threading.Thread( + target=client._renew_worker_registration, + name=f"modelexpress-refit-renew-{worker_id}", + daemon=True, + ) + client._registration_thread.start() + return client + + @property + def _service(self) -> refit_pb2_grpc.RefitServiceStub: + if self._channel is None: + self._channel = auth.with_auth(grpc.insecure_channel(self.server_url)) + self._stub = refit_pb2_grpc.RefitServiceStub(self._channel) + assert self._stub is not None + return self._stub + + def _register_worker(self) -> None: + self._service.RegisterWorker( + refit_pb2.RegisterWorkerRequest( + worker=refit_pb2.WorkerRegistration( + worker_id=self.worker_id, + role=refit_pb2.WORKER_ROLE_TRAINER, + model_name=self.model_name, + endpoint=self.worker_endpoint, + ), + ttl_seconds=self._registration_ttl_seconds, + ), + timeout=self._rpc_timeout_seconds, + ) + + def _renew_worker_registration(self) -> None: + interval_seconds = max(self._registration_ttl_seconds / 3, 0.1) + while not self._registration_stop.wait(interval_seconds): + try: + self._register_worker() + except grpc.RpcError: + # A later renewal retries after transient control-plane failure. + continue + + def stage_shard( + self, + *, + version: WeightVersionRef, + tensors: Any, + ) -> StagedWeightVersionShard: + """Capture one immutable rank-local shard for ``version``.""" + if not isinstance(version, WeightVersionRef): + raise TypeError("version must be a WeightVersionRef") + staged = self._adapter.stage_shard( + tensors=tensors, + staging_mode=self.staging_mode, + payload_format=self.payload_format, + ) + return StagedWeightVersionShard(client=self, version=version, staged=staged) + + def _publish_staged_shard( + self, + *, + version: WeightVersionRef, + staged: StagedWeightVersionShardData, + ) -> None: + source_slot_id = self._adapter.source_slot_id + staged.publish_ready.wait() + manifest_endpoint = self._manifest_publisher.publish_manifest( + version_id=version.version_id, + source_slot_id=source_slot_id, + manifest=staged.manifest, + ) + shard = refit_pb2.WeightVersionShard( + version_id=version.version_id, + source_slot_id=source_slot_id, + worker_id=self.worker_id, + tensor_count=staged.manifest.tensor_count, + total_bytes=staged.manifest.total_bytes, + manifest_digest=staged.manifest.digest, + manifest_endpoint=_required(manifest_endpoint, "manifest_endpoint"), + transport=staged.manifest.transport, + ) + self._service.CreateWeightVersionShard( + refit_pb2.CreateWeightVersionShardRequest(shard=shard), + timeout=self._rpc_timeout_seconds, + ) + # Keep the adapter-owned buffers alive while the published version can + # still be selected as a source. Eviction/release is a later lifecycle + # operation, not the staged handle's Python object lifetime. + self._published_shards.setdefault(version.version_id, []).append(staged) + + def close(self) -> None: + """Close the underlying gRPC channel.""" + if self._registration_thread is not None: + self._registration_stop.set() + self._registration_thread.join() + self._registration_thread = None + if self._channel is not None: + self._channel.close() + self._channel = None + self._stub = None + self._published_shards.clear() + + def __enter__(self) -> ModelExpressTrainerClient: + return self + + def __exit__(self, _exc_type, _exc_value, _traceback) -> None: + self.close() + + +__all__ = [ + "ModelExpressTrainerClient", + "StagedWeightVersionShard", + "WeightVersionRef", +] diff --git a/modelexpress_client/python/modelexpress_rl/envs.py b/modelexpress_client/python/modelexpress_rl/envs.py new file mode 100644 index 00000000..a6a179ec --- /dev/null +++ b/modelexpress_client/python/modelexpress_rl/envs.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""RL-specific deployment policy. + +Model identity, server connectivity, worker endpoints, NIXL ports, and +heartbeat timing use the shared :mod:`modelexpress.envs` configuration. +Rank-local identity and endpoints are derived from the initialized engine. +""" + +from __future__ import annotations + +import os +from collections.abc import Callable +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + MX_TRAINER_ENGINE: str + MX_TRAINER_STAGING_MODE: str + MX_WEIGHT_PAYLOAD_FORMAT: str + + +environment_variables: dict[str, Callable[[], Any]] = { + "MX_TRAINER_ENGINE": lambda: os.environ.get("MX_TRAINER_ENGINE", "MEGATRON") + .strip() + .upper(), + "MX_TRAINER_STAGING_MODE": lambda: os.environ.get( + "MX_TRAINER_STAGING_MODE", "IN_PLACE" + ) + .strip() + .upper(), + "MX_WEIGHT_PAYLOAD_FORMAT": lambda: os.environ.get( + "MX_WEIGHT_PAYLOAD_FORMAT", "FULL_TENSOR" + ) + .strip() + .upper(), +} + + +def __getattr__(name: str) -> Any: + if name in environment_variables: + return environment_variables[name]() + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def __dir__() -> list[str]: + return sorted(environment_variables) diff --git a/modelexpress_client/python/modelexpress_rl/inference/__init__.py b/modelexpress_client/python/modelexpress_rl/inference/__init__.py new file mode 100644 index 00000000..875b67d4 --- /dev/null +++ b/modelexpress_client/python/modelexpress_rl/inference/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Inference-side ModelExpress RL integrations.""" diff --git a/modelexpress_client/python/modelexpress_rl/inference/reshard/__init__.py b/modelexpress_client/python/modelexpress_rl/inference/reshard/__init__.py new file mode 100644 index 00000000..ec89865f --- /dev/null +++ b/modelexpress_client/python/modelexpress_rl/inference/reshard/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Inference-side reshard planning and compatibility receivers.""" diff --git a/modelexpress_client/python/modelexpress_rl/inference/reshard/megatron/__init__.py b/modelexpress_client/python/modelexpress_rl/inference/reshard/megatron/__init__.py new file mode 100644 index 00000000..41b4e5a6 --- /dev/null +++ b/modelexpress_client/python/modelexpress_rl/inference/reshard/megatron/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Inference-side planning for Megatron-native source tensors.""" + +from .layout import MegatronTargetLayout, MegatronTargetSpec, lower_megatron_target +from .receiver import MegatronReshardReceiver + +__all__ = [ + "MegatronReshardReceiver", + "MegatronTargetLayout", + "MegatronTargetSpec", + "lower_megatron_target", +] diff --git a/modelexpress_client/python/modelexpress/refit/reshard/megatron.py b/modelexpress_client/python/modelexpress_rl/inference/reshard/megatron/layout.py similarity index 92% rename from modelexpress_client/python/modelexpress/refit/reshard/megatron.py rename to modelexpress_client/python/modelexpress_rl/inference/reshard/megatron/layout.py index ec269dfe..20157d72 100644 --- a/modelexpress_client/python/modelexpress/refit/reshard/megatron.py +++ b/modelexpress_client/python/modelexpress_rl/inference/reshard/megatron/layout.py @@ -12,6 +12,7 @@ from __future__ import annotations from dataclasses import dataclass, field +from numbers import Integral from typing import Any from modelexpress.refit.reshard.slice_plan import _row_major_strides @@ -31,6 +32,10 @@ SUPPORTED_ROLES = COLUMN_ROLES | ROW_ROLES | {REPLICATED} +def _is_integral(value: object) -> bool: + return isinstance(value, Integral) and not isinstance(value, bool) + + @dataclass(frozen=True) class MegatronTargetSpec: """One native Megatron tensor requested by an inference TP rank.""" @@ -47,10 +52,10 @@ def __post_init__(self) -> None: if self.role not in SUPPORTED_ROLES: raise ValueError(f"unsupported Megatron role {self.role!r}") if not self.global_shape or any( - int(extent) <= 0 for extent in self.global_shape + not _is_integral(extent) or extent <= 0 for extent in self.global_shape ): raise ValueError( - f"{self.source_name}: global_shape must contain positive extents" + f"{self.source_name}: global_shape must contain positive integer extents" ) @@ -60,6 +65,10 @@ class MegatronTargetLayout: tp_rank: int def __post_init__(self) -> None: + if not _is_integral(self.tp_size): + raise ValueError("tp_size must be an integer") + if not _is_integral(self.tp_rank): + raise ValueError("tp_rank must be an integer") if self.tp_size < 1: raise ValueError("tp_size must be at least 1") if not 0 <= self.tp_rank < self.tp_size: diff --git a/modelexpress_client/python/modelexpress/refit/reshard/megatron_receiver.py b/modelexpress_client/python/modelexpress_rl/inference/reshard/megatron/receiver.py similarity index 98% rename from modelexpress_client/python/modelexpress/refit/reshard/megatron_receiver.py rename to modelexpress_client/python/modelexpress_rl/inference/reshard/megatron/receiver.py index 9927af3a..0e54df38 100644 --- a/modelexpress_client/python/modelexpress/refit/reshard/megatron_receiver.py +++ b/modelexpress_client/python/modelexpress_rl/inference/reshard/megatron/receiver.py @@ -7,12 +7,13 @@ from collections.abc import Callable from typing import Any -from modelexpress.refit.reshard.megatron import ( +from modelexpress.refit.reshard.receiver import ReshardReceiver + +from .layout import ( MegatronTargetLayout, MegatronTargetSpec, lower_megatron_target, ) -from modelexpress.refit.reshard.receiver import ReshardReceiver def _dtype_label(dtype: Any) -> str: diff --git a/modelexpress_client/python/modelexpress_rl/refit_pb2.py b/modelexpress_client/python/modelexpress_rl/refit_pb2.py new file mode 100644 index 00000000..b92e0cf1 --- /dev/null +++ b/modelexpress_client/python/modelexpress_rl/refit_pb2.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# NO CHECKED-IN PROTOBUF GENCODE +# source: refit.proto +# Protobuf Python Version: 5.27.2 +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import runtime_version as _runtime_version +from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import builder as _builder +_runtime_version.ValidateProtobufRuntimeVersion( + _runtime_version.Domain.PUBLIC, + 5, + 27, + 2, + '', + 'refit.proto' +) +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0brefit.proto\x12\x13model_express.refit\"\x98\x01\n\x12WorkerRegistration\x12\x11\n\tworker_id\x18\x01 \x01(\t\x12-\n\x04role\x18\x02 \x01(\x0e\x32\x1f.model_express.refit.WorkerRole\x12\x12\n\nmodel_name\x18\x03 \x01(\t\x12\x10\n\x08\x65ndpoint\x18\x04 \x01(\t\x12\x1a\n\x12\x65xpires_at_unix_ms\x18\x05 \x01(\x04\"e\n\x15RegisterWorkerRequest\x12\x37\n\x06worker\x18\x01 \x01(\x0b\x32\'.model_express.refit.WorkerRegistration\x12\x13\n\x0bttl_seconds\x18\x02 \x01(\r\"\xfa\x02\n\rWeightVersion\x12\x0b\n\x03uid\x18\x01 \x01(\t\x12\x12\n\nmodel_name\x18\x02 \x01(\t\x12\x1b\n\x0eversion_number\x18\x03 \x01(\x04H\x00\x88\x01\x01\x12\x17\n\x0fidempotency_key\x18\x04 \x01(\t\x12@\n\x0epayload_format\x18\x05 \x01(\x0e\x32(.model_express.refit.WeightPayloadFormat\x12\x1c\n\x0f\x62\x61se_version_id\x18\x06 \x01(\tH\x01\x88\x01\x01\x12\x1d\n\x15\x65xpected_source_slots\x18\x07 \x03(\t\x12\x18\n\x10layout_signature\x18\x08 \x01(\t\x12\x36\n\x05state\x18\t \x01(\x0e\x32\'.model_express.refit.WeightVersionState\x12\x1a\n\x12\x63reated_at_unix_ms\x18\n \x01(\x04\x42\x11\n\x0f_version_numberB\x12\n\x10_base_version_id\"\x8c\x02\n\x1a\x43reateWeightVersionRequest\x12\x12\n\nmodel_name\x18\x01 \x01(\t\x12\x1b\n\x0eversion_number\x18\x02 \x01(\x04H\x00\x88\x01\x01\x12\x17\n\x0fidempotency_key\x18\x03 \x01(\t\x12@\n\x0epayload_format\x18\x04 \x01(\x0e\x32(.model_express.refit.WeightPayloadFormat\x12\x1c\n\x0f\x62\x61se_version_id\x18\x05 \x01(\tH\x01\x88\x01\x01\x12\x1d\n\x15\x65xpected_source_slots\x18\x06 \x03(\tB\x11\n\x0f_version_numberB\x12\n\x10_base_version_id\"&\n\x17GetWeightVersionRequest\x12\x0b\n\x03uid\x18\x01 \x01(\t\")\n\x1a\x44\x65leteWeightVersionRequest\x12\x0b\n\x03uid\x18\x01 \x01(\t\"\xc5\x01\n\x12WeightVersionShard\x12\x12\n\nversion_id\x18\x01 \x01(\t\x12\x16\n\x0esource_slot_id\x18\x02 \x01(\t\x12\x11\n\tworker_id\x18\x03 \x01(\t\x12\x14\n\x0ctensor_count\x18\x04 \x01(\x04\x12\x13\n\x0btotal_bytes\x18\x05 \x01(\x04\x12\x17\n\x0fmanifest_digest\x18\x06 \x01(\t\x12\x19\n\x11manifest_endpoint\x18\x07 \x01(\t\x12\x11\n\ttransport\x18\x08 \x01(\t\"Y\n\x1f\x43reateWeightVersionShardRequest\x12\x36\n\x05shard\x18\x01 \x01(\x0b\x32\'.model_express.refit.WeightVersionShard\"\x8f\x01\n CreateWeightVersionShardResponse\x12\x36\n\x05shard\x18\x01 \x01(\x0b\x32\'.model_express.refit.WeightVersionShard\x12\x33\n\x07version\x18\x02 \x01(\x0b\x32\".model_express.refit.WeightVersion\"4\n\x1eListWeightVersionShardsRequest\x12\x12\n\nversion_id\x18\x01 \x01(\t\"Z\n\x1fListWeightVersionShardsResponse\x12\x37\n\x06shards\x18\x01 \x03(\x0b\x32\'.model_express.refit.WeightVersionShard\"R\n$GetWeightVersionShardManifestRequest\x12\x12\n\nversion_id\x18\x01 \x01(\t\x12\x16\n\x0esource_slot_id\x18\x02 \x01(\t\"R\n%GetWeightVersionShardManifestResponse\x12\x10\n\x08manifest\x18\x01 \x01(\x0c\x12\x17\n\x0fmanifest_digest\x18\x02 \x01(\t\"`\n\x1f\x44\x65leteWeightVersionShardRequest\x12\x12\n\nversion_id\x18\x01 \x01(\t\x12\x16\n\x0esource_slot_id\x18\x02 \x01(\t\x12\x11\n\tworker_id\x18\x03 \x01(\t\"3\n DeleteWeightVersionShardResponse\x12\x0f\n\x07\x64\x65leted\x18\x01 \x01(\x08\"c\n\x0cVersionLease\x12\x10\n\x08lease_id\x18\x01 \x01(\t\x12\x12\n\nversion_id\x18\x02 \x01(\t\x12\x11\n\tworker_id\x18\x03 \x01(\t\x12\x1a\n\x12\x65xpires_at_unix_ms\x18\x04 \x01(\x04\"Y\n\x1bRegisterVersionLeaseRequest\x12\x12\n\nversion_id\x18\x01 \x01(\t\x12\x11\n\tworker_id\x18\x02 \x01(\t\x12\x13\n\x0bttl_seconds\x18\x03 \x01(\r\"T\n\x19\x44\x65leteVersionLeaseRequest\x12\x12\n\nversion_id\x18\x01 \x01(\t\x12\x10\n\x08lease_id\x18\x02 \x01(\t\x12\x11\n\tworker_id\x18\x03 \x01(\t\"-\n\x1a\x44\x65leteVersionLeaseResponse\x12\x0f\n\x07\x64\x65leted\x18\x01 \x01(\x08*]\n\nWorkerRole\x12\x1b\n\x17WORKER_ROLE_UNSPECIFIED\x10\x00\x12\x17\n\x13WORKER_ROLE_TRAINER\x10\x01\x12\x19\n\x15WORKER_ROLE_GENERATOR\x10\x02*\x88\x01\n\x13WeightPayloadFormat\x12%\n!WEIGHT_PAYLOAD_FORMAT_UNSPECIFIED\x10\x00\x12%\n!WEIGHT_PAYLOAD_FORMAT_FULL_TENSOR\x10\x01\x12#\n\x1fWEIGHT_PAYLOAD_FORMAT_XOR_DELTA\x10\x02*\xa0\x01\n\x12WeightVersionState\x12$\n WEIGHT_VERSION_STATE_UNSPECIFIED\x10\x00\x12 \n\x1cWEIGHT_VERSION_STATE_STAGING\x10\x01\x12\x1e\n\x1aWEIGHT_VERSION_STATE_READY\x10\x02\x12\"\n\x1eWEIGHT_VERSION_STATE_RELEASING\x10\x03\x32\xb2\x08\n\x0cRefitService\x12j\n\x13\x43reateWeightVersion\x12/.model_express.refit.CreateWeightVersionRequest\x1a\".model_express.refit.WeightVersion\x12\x64\n\x10GetWeightVersion\x12,.model_express.refit.GetWeightVersionRequest\x1a\".model_express.refit.WeightVersion\x12j\n\x13\x44\x65leteWeightVersion\x12/.model_express.refit.DeleteWeightVersionRequest\x1a\".model_express.refit.WeightVersion\x12\x65\n\x0eRegisterWorker\x12*.model_express.refit.RegisterWorkerRequest\x1a\'.model_express.refit.WorkerRegistration\x12\x87\x01\n\x18\x43reateWeightVersionShard\x12\x34.model_express.refit.CreateWeightVersionShardRequest\x1a\x35.model_express.refit.CreateWeightVersionShardResponse\x12\x84\x01\n\x17ListWeightVersionShards\x12\x33.model_express.refit.ListWeightVersionShardsRequest\x1a\x34.model_express.refit.ListWeightVersionShardsResponse\x12\x87\x01\n\x18\x44\x65leteWeightVersionShard\x12\x34.model_express.refit.DeleteWeightVersionShardRequest\x1a\x35.model_express.refit.DeleteWeightVersionShardResponse\x12k\n\x14RegisterVersionLease\x12\x30.model_express.refit.RegisterVersionLeaseRequest\x1a!.model_express.refit.VersionLease\x12u\n\x12\x44\x65leteVersionLease\x12..model_express.refit.DeleteVersionLeaseRequest\x1a/.model_express.refit.DeleteVersionLeaseResponse2\xad\x01\n\x12RefitWorkerService\x12\x96\x01\n\x1dGetWeightVersionShardManifest\x12\x39.model_express.refit.GetWeightVersionShardManifestRequest\x1a:.model_express.refit.GetWeightVersionShardManifestResponseb\x06proto3') + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'refit_pb2', _globals) +if not _descriptor._USE_C_DESCRIPTORS: + DESCRIPTOR._loaded_options = None + _globals['_WORKERROLE']._serialized_start=2256 + _globals['_WORKERROLE']._serialized_end=2349 + _globals['_WEIGHTPAYLOADFORMAT']._serialized_start=2352 + _globals['_WEIGHTPAYLOADFORMAT']._serialized_end=2488 + _globals['_WEIGHTVERSIONSTATE']._serialized_start=2491 + _globals['_WEIGHTVERSIONSTATE']._serialized_end=2651 + _globals['_WORKERREGISTRATION']._serialized_start=37 + _globals['_WORKERREGISTRATION']._serialized_end=189 + _globals['_REGISTERWORKERREQUEST']._serialized_start=191 + _globals['_REGISTERWORKERREQUEST']._serialized_end=292 + _globals['_WEIGHTVERSION']._serialized_start=295 + _globals['_WEIGHTVERSION']._serialized_end=673 + _globals['_CREATEWEIGHTVERSIONREQUEST']._serialized_start=676 + _globals['_CREATEWEIGHTVERSIONREQUEST']._serialized_end=944 + _globals['_GETWEIGHTVERSIONREQUEST']._serialized_start=946 + _globals['_GETWEIGHTVERSIONREQUEST']._serialized_end=984 + _globals['_DELETEWEIGHTVERSIONREQUEST']._serialized_start=986 + _globals['_DELETEWEIGHTVERSIONREQUEST']._serialized_end=1027 + _globals['_WEIGHTVERSIONSHARD']._serialized_start=1030 + _globals['_WEIGHTVERSIONSHARD']._serialized_end=1227 + _globals['_CREATEWEIGHTVERSIONSHARDREQUEST']._serialized_start=1229 + _globals['_CREATEWEIGHTVERSIONSHARDREQUEST']._serialized_end=1318 + _globals['_CREATEWEIGHTVERSIONSHARDRESPONSE']._serialized_start=1321 + _globals['_CREATEWEIGHTVERSIONSHARDRESPONSE']._serialized_end=1464 + _globals['_LISTWEIGHTVERSIONSHARDSREQUEST']._serialized_start=1466 + _globals['_LISTWEIGHTVERSIONSHARDSREQUEST']._serialized_end=1518 + _globals['_LISTWEIGHTVERSIONSHARDSRESPONSE']._serialized_start=1520 + _globals['_LISTWEIGHTVERSIONSHARDSRESPONSE']._serialized_end=1610 + _globals['_GETWEIGHTVERSIONSHARDMANIFESTREQUEST']._serialized_start=1612 + _globals['_GETWEIGHTVERSIONSHARDMANIFESTREQUEST']._serialized_end=1694 + _globals['_GETWEIGHTVERSIONSHARDMANIFESTRESPONSE']._serialized_start=1696 + _globals['_GETWEIGHTVERSIONSHARDMANIFESTRESPONSE']._serialized_end=1778 + _globals['_DELETEWEIGHTVERSIONSHARDREQUEST']._serialized_start=1780 + _globals['_DELETEWEIGHTVERSIONSHARDREQUEST']._serialized_end=1876 + _globals['_DELETEWEIGHTVERSIONSHARDRESPONSE']._serialized_start=1878 + _globals['_DELETEWEIGHTVERSIONSHARDRESPONSE']._serialized_end=1929 + _globals['_VERSIONLEASE']._serialized_start=1931 + _globals['_VERSIONLEASE']._serialized_end=2030 + _globals['_REGISTERVERSIONLEASEREQUEST']._serialized_start=2032 + _globals['_REGISTERVERSIONLEASEREQUEST']._serialized_end=2121 + _globals['_DELETEVERSIONLEASEREQUEST']._serialized_start=2123 + _globals['_DELETEVERSIONLEASEREQUEST']._serialized_end=2207 + _globals['_DELETEVERSIONLEASERESPONSE']._serialized_start=2209 + _globals['_DELETEVERSIONLEASERESPONSE']._serialized_end=2254 + _globals['_REFITSERVICE']._serialized_start=2654 + _globals['_REFITSERVICE']._serialized_end=3728 + _globals['_REFITWORKERSERVICE']._serialized_start=3731 + _globals['_REFITWORKERSERVICE']._serialized_end=3904 +# @@protoc_insertion_point(module_scope) diff --git a/modelexpress_client/python/modelexpress_rl/refit_pb2_grpc.py b/modelexpress_client/python/modelexpress_rl/refit_pb2_grpc.py new file mode 100644 index 00000000..248eaa58 --- /dev/null +++ b/modelexpress_client/python/modelexpress_rl/refit_pb2_grpc.py @@ -0,0 +1,533 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! +"""Client and server classes corresponding to protobuf-defined services.""" +import grpc +import warnings + +from . import refit_pb2 as refit__pb2 + +GRPC_GENERATED_VERSION = '1.66.2' +GRPC_VERSION = grpc.__version__ +_version_not_supported = False + +try: + from grpc._utilities import first_version_is_lower + _version_not_supported = first_version_is_lower(GRPC_VERSION, GRPC_GENERATED_VERSION) +except ImportError: + _version_not_supported = True + +if _version_not_supported: + raise RuntimeError( + f'The grpc package installed is at version {GRPC_VERSION},' + + ' but the generated code in refit_pb2_grpc.py depends on' + + f' grpcio>={GRPC_GENERATED_VERSION}.' + + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}' + + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.' + ) + + +class RefitServiceStub(object): + """Control-plane metadata for RL weight publication. Weight bytes and full + manifests remain on trainer and generator workers. + """ + + def __init__(self, channel): + """Constructor. + + Args: + channel: A grpc.Channel. + """ + self.CreateWeightVersion = channel.unary_unary( + '/model_express.refit.RefitService/CreateWeightVersion', + request_serializer=refit__pb2.CreateWeightVersionRequest.SerializeToString, + response_deserializer=refit__pb2.WeightVersion.FromString, + _registered_method=True) + self.GetWeightVersion = channel.unary_unary( + '/model_express.refit.RefitService/GetWeightVersion', + request_serializer=refit__pb2.GetWeightVersionRequest.SerializeToString, + response_deserializer=refit__pb2.WeightVersion.FromString, + _registered_method=True) + self.DeleteWeightVersion = channel.unary_unary( + '/model_express.refit.RefitService/DeleteWeightVersion', + request_serializer=refit__pb2.DeleteWeightVersionRequest.SerializeToString, + response_deserializer=refit__pb2.WeightVersion.FromString, + _registered_method=True) + self.RegisterWorker = channel.unary_unary( + '/model_express.refit.RefitService/RegisterWorker', + request_serializer=refit__pb2.RegisterWorkerRequest.SerializeToString, + response_deserializer=refit__pb2.WorkerRegistration.FromString, + _registered_method=True) + self.CreateWeightVersionShard = channel.unary_unary( + '/model_express.refit.RefitService/CreateWeightVersionShard', + request_serializer=refit__pb2.CreateWeightVersionShardRequest.SerializeToString, + response_deserializer=refit__pb2.CreateWeightVersionShardResponse.FromString, + _registered_method=True) + self.ListWeightVersionShards = channel.unary_unary( + '/model_express.refit.RefitService/ListWeightVersionShards', + request_serializer=refit__pb2.ListWeightVersionShardsRequest.SerializeToString, + response_deserializer=refit__pb2.ListWeightVersionShardsResponse.FromString, + _registered_method=True) + self.DeleteWeightVersionShard = channel.unary_unary( + '/model_express.refit.RefitService/DeleteWeightVersionShard', + request_serializer=refit__pb2.DeleteWeightVersionShardRequest.SerializeToString, + response_deserializer=refit__pb2.DeleteWeightVersionShardResponse.FromString, + _registered_method=True) + self.RegisterVersionLease = channel.unary_unary( + '/model_express.refit.RefitService/RegisterVersionLease', + request_serializer=refit__pb2.RegisterVersionLeaseRequest.SerializeToString, + response_deserializer=refit__pb2.VersionLease.FromString, + _registered_method=True) + self.DeleteVersionLease = channel.unary_unary( + '/model_express.refit.RefitService/DeleteVersionLease', + request_serializer=refit__pb2.DeleteVersionLeaseRequest.SerializeToString, + response_deserializer=refit__pb2.DeleteVersionLeaseResponse.FromString, + _registered_method=True) + + +class RefitServiceServicer(object): + """Control-plane metadata for RL weight publication. Weight bytes and full + manifests remain on trainer and generator workers. + """ + + def CreateWeightVersion(self, request, context): + """External APIs: called by the RL framework orchestrator. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def GetWeightVersion(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def DeleteWeightVersion(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def RegisterWorker(self, request, context): + """Internal APIs: called by ModelExpress worker clients, not by the RL + framework orchestrator. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def CreateWeightVersionShard(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def ListWeightVersionShards(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def DeleteWeightVersionShard(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def RegisterVersionLease(self, request, context): + """A live lease protects every shard while a generator installs the version. + Re-registering the same lease owner renews its expiry. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def DeleteVersionLease(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + +def add_RefitServiceServicer_to_server(servicer, server): + rpc_method_handlers = { + 'CreateWeightVersion': grpc.unary_unary_rpc_method_handler( + servicer.CreateWeightVersion, + request_deserializer=refit__pb2.CreateWeightVersionRequest.FromString, + response_serializer=refit__pb2.WeightVersion.SerializeToString, + ), + 'GetWeightVersion': grpc.unary_unary_rpc_method_handler( + servicer.GetWeightVersion, + request_deserializer=refit__pb2.GetWeightVersionRequest.FromString, + response_serializer=refit__pb2.WeightVersion.SerializeToString, + ), + 'DeleteWeightVersion': grpc.unary_unary_rpc_method_handler( + servicer.DeleteWeightVersion, + request_deserializer=refit__pb2.DeleteWeightVersionRequest.FromString, + response_serializer=refit__pb2.WeightVersion.SerializeToString, + ), + 'RegisterWorker': grpc.unary_unary_rpc_method_handler( + servicer.RegisterWorker, + request_deserializer=refit__pb2.RegisterWorkerRequest.FromString, + response_serializer=refit__pb2.WorkerRegistration.SerializeToString, + ), + 'CreateWeightVersionShard': grpc.unary_unary_rpc_method_handler( + servicer.CreateWeightVersionShard, + request_deserializer=refit__pb2.CreateWeightVersionShardRequest.FromString, + response_serializer=refit__pb2.CreateWeightVersionShardResponse.SerializeToString, + ), + 'ListWeightVersionShards': grpc.unary_unary_rpc_method_handler( + servicer.ListWeightVersionShards, + request_deserializer=refit__pb2.ListWeightVersionShardsRequest.FromString, + response_serializer=refit__pb2.ListWeightVersionShardsResponse.SerializeToString, + ), + 'DeleteWeightVersionShard': grpc.unary_unary_rpc_method_handler( + servicer.DeleteWeightVersionShard, + request_deserializer=refit__pb2.DeleteWeightVersionShardRequest.FromString, + response_serializer=refit__pb2.DeleteWeightVersionShardResponse.SerializeToString, + ), + 'RegisterVersionLease': grpc.unary_unary_rpc_method_handler( + servicer.RegisterVersionLease, + request_deserializer=refit__pb2.RegisterVersionLeaseRequest.FromString, + response_serializer=refit__pb2.VersionLease.SerializeToString, + ), + 'DeleteVersionLease': grpc.unary_unary_rpc_method_handler( + servicer.DeleteVersionLease, + request_deserializer=refit__pb2.DeleteVersionLeaseRequest.FromString, + response_serializer=refit__pb2.DeleteVersionLeaseResponse.SerializeToString, + ), + } + generic_handler = grpc.method_handlers_generic_handler( + 'model_express.refit.RefitService', rpc_method_handlers) + server.add_generic_rpc_handlers((generic_handler,)) + server.add_registered_method_handlers('model_express.refit.RefitService', rpc_method_handlers) + + + # This class is part of an EXPERIMENTAL API. +class RefitService(object): + """Control-plane metadata for RL weight publication. Weight bytes and full + manifests remain on trainer and generator workers. + """ + + @staticmethod + def CreateWeightVersion(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/model_express.refit.RefitService/CreateWeightVersion', + refit__pb2.CreateWeightVersionRequest.SerializeToString, + refit__pb2.WeightVersion.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def GetWeightVersion(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/model_express.refit.RefitService/GetWeightVersion', + refit__pb2.GetWeightVersionRequest.SerializeToString, + refit__pb2.WeightVersion.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def DeleteWeightVersion(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/model_express.refit.RefitService/DeleteWeightVersion', + refit__pb2.DeleteWeightVersionRequest.SerializeToString, + refit__pb2.WeightVersion.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def RegisterWorker(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/model_express.refit.RefitService/RegisterWorker', + refit__pb2.RegisterWorkerRequest.SerializeToString, + refit__pb2.WorkerRegistration.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def CreateWeightVersionShard(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/model_express.refit.RefitService/CreateWeightVersionShard', + refit__pb2.CreateWeightVersionShardRequest.SerializeToString, + refit__pb2.CreateWeightVersionShardResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def ListWeightVersionShards(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/model_express.refit.RefitService/ListWeightVersionShards', + refit__pb2.ListWeightVersionShardsRequest.SerializeToString, + refit__pb2.ListWeightVersionShardsResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def DeleteWeightVersionShard(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/model_express.refit.RefitService/DeleteWeightVersionShard', + refit__pb2.DeleteWeightVersionShardRequest.SerializeToString, + refit__pb2.DeleteWeightVersionShardResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def RegisterVersionLease(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/model_express.refit.RefitService/RegisterVersionLease', + refit__pb2.RegisterVersionLeaseRequest.SerializeToString, + refit__pb2.VersionLease.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def DeleteVersionLease(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/model_express.refit.RefitService/DeleteVersionLease', + refit__pb2.DeleteVersionLeaseRequest.SerializeToString, + refit__pb2.DeleteVersionLeaseResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + +class RefitWorkerServiceStub(object): + """Internal worker-to-worker API. The generator fetches the small transfer + manifest here; tensor bytes remain on the advertised data-plane transport. + """ + + def __init__(self, channel): + """Constructor. + + Args: + channel: A grpc.Channel. + """ + self.GetWeightVersionShardManifest = channel.unary_unary( + '/model_express.refit.RefitWorkerService/GetWeightVersionShardManifest', + request_serializer=refit__pb2.GetWeightVersionShardManifestRequest.SerializeToString, + response_deserializer=refit__pb2.GetWeightVersionShardManifestResponse.FromString, + _registered_method=True) + + +class RefitWorkerServiceServicer(object): + """Internal worker-to-worker API. The generator fetches the small transfer + manifest here; tensor bytes remain on the advertised data-plane transport. + """ + + def GetWeightVersionShardManifest(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + +def add_RefitWorkerServiceServicer_to_server(servicer, server): + rpc_method_handlers = { + 'GetWeightVersionShardManifest': grpc.unary_unary_rpc_method_handler( + servicer.GetWeightVersionShardManifest, + request_deserializer=refit__pb2.GetWeightVersionShardManifestRequest.FromString, + response_serializer=refit__pb2.GetWeightVersionShardManifestResponse.SerializeToString, + ), + } + generic_handler = grpc.method_handlers_generic_handler( + 'model_express.refit.RefitWorkerService', rpc_method_handlers) + server.add_generic_rpc_handlers((generic_handler,)) + server.add_registered_method_handlers('model_express.refit.RefitWorkerService', rpc_method_handlers) + + + # This class is part of an EXPERIMENTAL API. +class RefitWorkerService(object): + """Internal worker-to-worker API. The generator fetches the small transfer + manifest here; tensor bytes remain on the advertised data-plane transport. + """ + + @staticmethod + def GetWeightVersionShardManifest(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/model_express.refit.RefitWorkerService/GetWeightVersionShardManifest', + refit__pb2.GetWeightVersionShardManifestRequest.SerializeToString, + refit__pb2.GetWeightVersionShardManifestResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) diff --git a/modelexpress_client/python/modelexpress_rl/train/__init__.py b/modelexpress_client/python/modelexpress_rl/train/__init__.py new file mode 100644 index 00000000..dc592f52 --- /dev/null +++ b/modelexpress_client/python/modelexpress_rl/train/__init__.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Trainer-side ModelExpress RL integrations.""" + +from .adapter import ( + CompletionFence, + NixlMetadataProvider, + StagedWeightVersionShardData, + TrainerEngineAdapter, + TrainerStagingMode, + WeightPayloadFormat, + WeightVersionShardManifest, + WeightVersionShardManifestPublisher, +) +from .manifest import WeightVersionShardManifestService + +__all__ = [ + "CompletionFence", + "NixlMetadataProvider", + "StagedWeightVersionShardData", + "TrainerEngineAdapter", + "TrainerStagingMode", + "WeightPayloadFormat", + "WeightVersionShardManifest", + "WeightVersionShardManifestPublisher", + "WeightVersionShardManifestService", +] diff --git a/modelexpress_client/python/modelexpress_rl/train/adapter.py b/modelexpress_client/python/modelexpress_rl/train/adapter.py new file mode 100644 index 00000000..dff225ca --- /dev/null +++ b/modelexpress_client/python/modelexpress_rl/train/adapter.py @@ -0,0 +1,155 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Trainer-engine boundary for ModelExpress RL refit publication.""" + +from __future__ import annotations + +import hashlib +from abc import ABC, abstractmethod +from collections.abc import Callable +from dataclasses import dataclass +from enum import Enum +from typing import Any, Protocol + + +class NixlMetadataProvider(Protocol): + """Narrow NIXL manager surface required to publish trainer manifests.""" + + @property + def agent_name(self) -> str: + """Return the local NIXL agent name.""" + ... + + @property + def nixl_metadata(self) -> bytes: + """Return serialized metadata for the local NIXL agent.""" + ... + + @property + def listen_port(self) -> int | None: + """Return the local NIXL metadata-listener port, when enabled.""" + ... + + +class TrainerStagingMode(str, Enum): + """How a trainer adapter preserves a version's immutable source bytes.""" + + UNSPECIFIED = "UNSPECIFIED" + COPY_TO_DEVICE = "COPY_TO_DEVICE" + COPY_TO_HOST = "COPY_TO_HOST" + WRITE_TO_STORAGE = "WRITE_TO_STORAGE" + IN_PLACE = "IN_PLACE" + + +class WeightPayloadFormat(str, Enum): + """Weight representation fixed for one initialized client.""" + + UNSPECIFIED = "UNSPECIFIED" + FULL_TENSOR = "FULL_TENSOR" + XOR_DELTA = "XOR_DELTA" + + +@dataclass(frozen=True) +class CompletionFence: + """Blocking completion fence for an adapter-owned asynchronous operation.""" + + _wait: Callable[[], None] + + def wait(self) -> None: + """Block until the operation represented by this fence completes.""" + self._wait() + + +@dataclass(frozen=True) +class WeightVersionShardManifest: + """Engine-neutral description of one trainer process's source buffers.""" + + data: bytes + tensor_count: int + total_bytes: int + transport: str + + def __post_init__(self) -> None: + if not self.data: + raise ValueError("manifest data must not be empty") + if self.tensor_count <= 0: + raise ValueError("tensor_count must be positive") + if self.total_bytes <= 0: + raise ValueError("total_bytes must be positive") + if not self.transport: + raise ValueError("transport must not be empty") + + @property + def digest(self) -> str: + """Return the SHA-256 digest advertised through RefitService.""" + return hashlib.sha256(self.data).hexdigest() + + +@dataclass(frozen=True) +class StagedWeightVersionShardData: + """Adapter-owned immutable buffers and their transfer manifest.""" + + manifest: WeightVersionShardManifest + publish_ready: CompletionFence + source_reuse_ready: CompletionFence + buffer_owner: object | None = None + + +class TrainerEngineAdapter(ABC): + """Engine-specific capture and staging boundary for trainer publication. + + ModelExpress owns worker registration, manifest serving, and control-plane + publication. An implementation captures engine tensors into immutable + source buffers and describes those buffers in a transfer manifest. + """ + + @property + @abstractmethod + def source_slot_id(self) -> str: + """Return this rank's required logical contribution identifier.""" + + @property + @abstractmethod + def supported_staging_modes(self) -> frozenset[TrainerStagingMode]: + """Return staging modes implemented by this engine adapter.""" + + @property + @abstractmethod + def supported_payload_formats(self) -> frozenset[WeightPayloadFormat]: + """Return payload formats implemented by this engine adapter.""" + + @abstractmethod + def stage_shard( + self, + *, + tensors: Any, + staging_mode: TrainerStagingMode, + payload_format: WeightPayloadFormat, + ) -> StagedWeightVersionShardData: + """Capture one immutable, rank-local version shard.""" + + +class WeightVersionShardManifestPublisher(Protocol): + """Worker endpoint that makes a manifest retrievable before advertisement.""" + + def publish_manifest( + self, + *, + version_id: str, + source_slot_id: str, + manifest: WeightVersionShardManifest, + ) -> str: + """Publish ``manifest`` and return its ready, worker-local endpoint.""" + + +__all__ = [ + "CompletionFence", + "NixlMetadataProvider", + "StagedWeightVersionShardData", + "TrainerEngineAdapter", + "TrainerStagingMode", + "WeightPayloadFormat", + "WeightVersionShardManifest", + "WeightVersionShardManifestPublisher", +] diff --git a/modelexpress_client/python/modelexpress_rl/train/engines/__init__.py b/modelexpress_client/python/modelexpress_rl/train/engines/__init__.py new file mode 100644 index 00000000..4d4bdda0 --- /dev/null +++ b/modelexpress_client/python/modelexpress_rl/train/engines/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Training-engine adapters for ModelExpress RL.""" diff --git a/modelexpress_client/python/modelexpress_rl/train/engines/megatron/__init__.py b/modelexpress_client/python/modelexpress_rl/train/engines/megatron/__init__.py new file mode 100644 index 00000000..043883d0 --- /dev/null +++ b/modelexpress_client/python/modelexpress_rl/train/engines/megatron/__init__.py @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Megatron integration for ModelExpress RL refit publication.""" + +from .adapter import MegatronTrainerAdapter +from .aliases import MegatronAliasInput, MegatronTensorSpec, build_hf_aliases +from .publisher import ( + MegatronPublishedTensorSpec, + MegatronReshardManifest, + build_megatron_reshard_manifest, + publish_megatron_reshard_view, + publish_registered_shard_table, +) + +__all__ = [ + "MegatronAliasInput", + "MegatronPublishedTensorSpec", + "MegatronReshardManifest", + "MegatronTensorSpec", + "MegatronTrainerAdapter", + "build_hf_aliases", + "build_megatron_reshard_manifest", + "publish_megatron_reshard_view", + "publish_registered_shard_table", +] diff --git a/modelexpress_client/python/modelexpress_rl/train/engines/megatron/adapter.py b/modelexpress_client/python/modelexpress_rl/train/engines/megatron/adapter.py new file mode 100644 index 00000000..38f3d7a8 --- /dev/null +++ b/modelexpress_client/python/modelexpress_rl/train/engines/megatron/adapter.py @@ -0,0 +1,119 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Megatron implementation of the trainer-engine adapter contract.""" + +from __future__ import annotations + +import math +from typing import Any + +import torch.distributed as dist + +from modelexpress_rl.train.adapter import ( + CompletionFence, + NixlMetadataProvider, + StagedWeightVersionShardData, + TrainerEngineAdapter, + TrainerStagingMode, + WeightPayloadFormat, + WeightVersionShardManifest, +) + +from .aliases import MegatronTensorSpec, build_hf_aliases +from .publisher import build_megatron_reshard_manifest + + +def _source_reuse_unsupported() -> None: + raise NotImplementedError( + "Megatron IN_PLACE requires the RL framework to retire the published " + "version before resuming training; version-retirement signaling is not wired" + ) + + +class MegatronTrainerAdapter(TrainerEngineAdapter): + """Expose existing Megatron/NIXL buffers through the trainer contract.""" + + def __init__( + self, + *, + manager: NixlMetadataProvider, + nixl_metadata_endpoint: str, + ) -> None: + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError("Megatron distributed process group is not initialized") + self._manager = manager + self._nixl_metadata_endpoint = nixl_metadata_endpoint + self._source_slot_id = f"publisher:global-rank:{dist.get_rank()}" + + @property + def source_slot_id(self) -> str: + """Return the logical trainer contribution represented by this rank.""" + return self._source_slot_id + + @property + def supported_staging_modes(self) -> frozenset[TrainerStagingMode]: + return frozenset({TrainerStagingMode.IN_PLACE}) + + @property + def supported_payload_formats(self) -> frozenset[WeightPayloadFormat]: + return frozenset({WeightPayloadFormat.FULL_TENSOR}) + + def stage_shard( + self, + *, + tensors: Any, + staging_mode: TrainerStagingMode, + payload_format: WeightPayloadFormat, + ) -> StagedWeightVersionShardData: + """Capture the current pre-registered Megatron source buffers.""" + if staging_mode not in self.supported_staging_modes: + raise NotImplementedError( + f"MegatronTrainerAdapter does not support {staging_mode.value} staging" + ) + if payload_format not in self.supported_payload_formats: + raise NotImplementedError( + f"MegatronTrainerAdapter does not support {payload_format.value} payloads" + ) + if not isinstance(tensors, list) or not all( + isinstance(item, MegatronTensorSpec) for item in tensors + ): + raise TypeError("tensors must be a list of MegatronTensorSpec") + published = build_hf_aliases( + tensors, + agent_name=str(self._manager.agent_name), + ) + manifest = build_megatron_reshard_manifest( + manager=self._manager, + published=published, + metadata_endpoint=self._nixl_metadata_endpoint, + ) + total_bytes = sum( + math.prod(shard.shape) * tensor.elsize + for tensor in manifest.tensors + for shard in tensor.shards + ) + + return StagedWeightVersionShardData( + manifest=WeightVersionShardManifest( + data=manifest.blob, + tensor_count=len(manifest.tensors), + total_bytes=total_bytes, + transport="NIXL", + ), + # IN_PLACE performs no asynchronous copy, so the shard is ready to + # publish as soon as its manifest has been built. + publish_ready=CompletionFence(lambda: None), + # The RL framework must not start the next optimizer step while the + # version is published: it waits for every generator update, retires + # the version, and only then resumes training. Retirement signaling + # is not wired to this rank-local fence yet, so fail rather than + # incorrectly reporting that the source is safe to mutate. + source_reuse_ready=CompletionFence(_source_reuse_unsupported), + # IN_PLACE borrows these live tensor objects until the version is + # retired; retaining them here makes that ownership explicit. + buffer_owner=tuple(tensors), + ) + + +__all__ = ["MegatronTrainerAdapter"] diff --git a/modelexpress_client/python/modelexpress_rl/train/engines/megatron/aliases.py b/modelexpress_client/python/modelexpress_rl/train/engines/megatron/aliases.py new file mode 100644 index 00000000..2ff688a0 --- /dev/null +++ b/modelexpress_client/python/modelexpress_rl/train/engines/megatron/aliases.py @@ -0,0 +1,288 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Expose native Megatron storage as HF-canonical RL refit source shards.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +from modelexpress.refit.reshard.rendezvous import PublishedShard, PublishedTensor +from modelexpress.refit.reshard.verify import published_digest + + +@dataclass(frozen=True) +class MegatronTensorSpec: + """One native Megatron tensor and its logical transfer representation.""" + + name: str + tensor: Any + role: str + hf_names: tuple[str, ...] + global_shape: tuple[int, ...] + placement_kind: str + shard_axis: int | None + local_shard_range: tuple[int, int] | None + extras: dict[str, str] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not self.name or not self.hf_names: + raise ValueError("name and hf_names are required") + if len(set(self.hf_names)) != len(self.hf_names): + raise ValueError(f"{self.name}: hf_names must be unique") + if not self.global_shape or any(int(dim) <= 0 for dim in self.global_shape): + raise ValueError(f"{self.name}: invalid global shape {self.global_shape}") + if self.placement_kind not in {"SHARD", "REPLICATE"}: + raise ValueError(f"{self.name}: unsupported placement_kind") + has_shard_layout = ( + self.shard_axis is not None and self.local_shard_range is not None + ) + if self.placement_kind == "SHARD" and not has_shard_layout: + raise ValueError( + f"{self.name}: SHARD requires shard_axis and local_shard_range" + ) + if self.placement_kind == "REPLICATE" and ( + self.shard_axis is not None or self.local_shard_range is not None + ): + raise ValueError(f"{self.name}: REPLICATE cannot set shard layout") + + +def _source_rank_and_size(item: MegatronTensorSpec, axis: int) -> tuple[int, int]: + local_extent = int(item.tensor.shape[axis]) + global_extent = int(item.global_shape[axis]) + if item.placement_kind != "SHARD": + return 0, 1 + if item.local_shard_range is None: + raise ValueError(f"{item.name}: SHARD has no local range") + lo, hi = (int(value) for value in item.local_shard_range) + # A range can pass every check below and still lie outside the tensor it + # claims part of: (16, 24) against a global extent of 16 has the right width, + # divides evenly, and yields source rank 2 of a 2-rank group. The alias that + # follows would then address bytes the full tensor does not have. + if not 0 <= lo < hi <= global_extent: + raise ValueError( + f"{item.name}: source shard range {(lo, hi)} is outside the global " + f"extent {global_extent} on axis {axis}" + ) + if hi - lo != local_extent or global_extent % local_extent: + raise ValueError(f"{item.name}: inconsistent source shard geometry") + if lo % local_extent: + raise ValueError(f"{item.name}: non-uniform source shard is unsupported") + return lo // local_extent, global_extent // local_extent + + +def _one_shard( + *, + name: str, + tensor: Any, + full_shape: tuple[int, ...], + agent_name: str, + shard_axis: int | None, + shard_range: tuple[int, int] | None, +) -> PublishedTensor: + local_shape = tuple(int(dim) for dim in tensor.shape) + offset = [0] * len(local_shape) + if shard_axis is not None: + if shard_range is None: + raise ValueError(f"{name}: shard axis has no range") + lo, hi = shard_range + if hi - lo != local_shape[shard_axis]: + raise ValueError(f"{name}: shard range does not match local shape") + offset[shard_axis] = lo + elif local_shape != full_shape: + raise ValueError(f"{name}: replicated shape mismatch") + return PublishedTensor( + name=name, + dtype=str(tensor.dtype), + elsize=int(tensor.element_size()), + full_shape=full_shape, + shards=[ + PublishedShard( + agent_name=agent_name, + device_id=int(tensor.device.index or 0), + addr=int(tensor.data_ptr()), + shard_offset=tuple(offset), + shape=local_shape, + digest=published_digest(tensor), + ) + ], + ) + + +GATE_THEN_UP = "gate_then_up" + + +def _build_gated_aliases( + item: MegatronTensorSpec, agent_name: str +) -> list[PublishedTensor]: + if len(item.hf_names) != 2: + raise ValueError(f"{item.name}: gated tensor requires gate/up HF names") + # The halves are assigned to hf_names positionally, so the storage order is + # what decides which HF tensor each half becomes. Getting it wrong publishes + # the gate projection's bytes under the up projection's name, which no digest + # gate can see: both names receive the bytes their publisher advertised. The + # order is therefore required rather than assumed. + order = item.extras.get("gated_mlp_order") + if order != GATE_THEN_UP: + raise ValueError( + f"{item.name}: fused gate/up aliasing requires extras" + f"['gated_mlp_order'] == {GATE_THEN_UP!r}, got {order!r}" + ) + axis = int(item.shard_axis if item.shard_axis is not None else 0) + local_extent = int(item.tensor.shape[axis]) + if local_extent % 2: + raise ValueError(f"{item.name}: fused gate/up extent must be even") + half = local_extent // 2 + gate = item.tensor.narrow(axis, 0, half) + up = item.tensor.narrow(axis, half, half) + if not gate.is_contiguous() or not up.is_contiguous(): + raise ValueError( + f"{item.name}: fused gate/up aliases are not contiguous on axis {axis}" + ) + source_rank, source_size = _source_rank_and_size(item, axis) + full_shape = list(item.tensor.shape) + full_shape[axis] = half * source_size + expected = [int(dim) for dim in item.global_shape] + if expected[axis] % 2: + raise ValueError(f"{item.name}: fused global extent must be even") + expected[axis] //= 2 + if expected != [int(dim) for dim in full_shape]: + raise ValueError( + f"{item.name}: derived gate/up shape {tuple(full_shape)} disagrees " + f"with declared global shape {item.global_shape}" + ) + shard_range = ( + (source_rank * half, (source_rank + 1) * half) if source_size > 1 else None + ) + return [ + _one_shard( + name=hf_name, + tensor=tensor, + full_shape=tuple(int(dim) for dim in full_shape), + agent_name=agent_name, + shard_axis=axis if shard_range is not None else None, + shard_range=shard_range, + ) + for hf_name, tensor in zip(item.hf_names, (gate, up), strict=True) + ] + + +def _build_qkv_aliases( + item: MegatronTensorSpec, agent_name: str +) -> list[PublishedTensor]: + if len(item.hf_names) != 3 or item.tensor.ndim != 2: + raise ValueError(f"{item.name}: QKV aliasing requires 2D q/k/v weights") + required = ("head_dim", "num_heads_local", "num_kv_heads_local") + missing = [key for key in required if key not in item.extras] + if missing: + raise ValueError(f"{item.name}: QKV aliasing requires extras {missing}") + head_dim = int(item.extras["head_dim"]) + q_heads_local = int(item.extras["num_heads_local"]) + kv_heads_local = int(item.extras["num_kv_heads_local"]) + if kv_heads_local < 1 or q_heads_local % kv_heads_local: + raise ValueError(f"{item.name}: invalid local Q/KV head geometry") + rows_per_group = (q_heads_local // kv_heads_local + 2) * head_dim + if rows_per_group * kv_heads_local != int(item.tensor.shape[0]): + raise ValueError(f"{item.name}: QKV rows disagree with head metadata") + source_rank, source_size = _source_rank_and_size(item, 0) + hidden = int(item.tensor.shape[1]) + q_heads_per_group = q_heads_local // kv_heads_local + q_shards = [] + k_shards = [] + v_shards = [] + for local_group in range(kv_heads_local): + group = item.tensor.narrow(0, local_group * rows_per_group, rows_per_group) + q_rows = q_heads_per_group * head_dim + q = group.narrow(0, 0, q_rows) + k = group.narrow(0, q_rows, head_dim) + v = group.narrow(0, q_rows + head_dim, head_dim) + global_group = source_rank * kv_heads_local + local_group + for tensor, shards, start in ( + (q, q_shards, global_group * q_rows), + (k, k_shards, global_group * head_dim), + (v, v_shards, global_group * head_dim), + ): + shards.append( + PublishedShard( + agent_name=agent_name, + device_id=int(tensor.device.index or 0), + addr=int(tensor.data_ptr()), + shard_offset=(start, 0), + shape=tuple(int(dim) for dim in tensor.shape), + # The narrow, not the fused parent: this is the box a receiver + # reads from ``addr``, so it is the box whose bytes must match. + digest=published_digest(tensor), + ) + ) + return [ + PublishedTensor( + name=name, + dtype=str(item.tensor.dtype), + elsize=int(item.tensor.element_size()), + full_shape=(rows, hidden), + shards=shards, + ) + for name, rows, shards in ( + (item.hf_names[0], q_heads_local * source_size * head_dim, q_shards), + (item.hf_names[1], kv_heads_local * source_size * head_dim, k_shards), + (item.hf_names[2], kv_heads_local * source_size * head_dim, v_shards), + ) + ] + + +def build_hf_aliases( + items: list[MegatronTensorSpec], *, agent_name: str +) -> list[PublishedTensor]: + """Build zero-copy HF aliases whose addresses remain in registered storage.""" + + aliases = [] + for item in items: + # Every alias published below is a base address plus a shape, which tells + # a reader the bytes run contiguously from that address. A strided view + # satisfies neither, and nothing downstream can detect it: the read simply + # lands on whatever sits between the elements the view meant to select. + if not item.tensor.is_contiguous(): + raise ValueError( + f"{item.name}: aliasing publishes an address and a shape, which " + f"requires contiguous storage; got a non-contiguous tensor of " + f"shape {tuple(int(dim) for dim in item.tensor.shape)}" + ) + if item.role == "qkv_column": + aliases.extend(_build_qkv_aliases(item, agent_name)) + continue + if ( + item.role in {"gated_mlp_column", "expert_column"} + and len(item.hf_names) == 2 + ): + aliases.extend(_build_gated_aliases(item, agent_name)) + continue + if len(item.hf_names) != 1: + raise ValueError( + f"{item.name}: role {item.role!r} cannot map to " + f"{len(item.hf_names)} HF tensors" + ) + aliases.append( + _one_shard( + name=item.hf_names[0], + tensor=item.tensor, + full_shape=tuple(item.global_shape), + agent_name=agent_name, + shard_axis=( + int(item.shard_axis) if item.placement_kind == "SHARD" else None + ), + shard_range=( + tuple(item.local_shard_range) + if item.placement_kind == "SHARD" + and item.local_shard_range is not None + else None + ), + ) + ) + return aliases + + +# Compatibility name for the existing reshard API. +MegatronAliasInput = MegatronTensorSpec + + +__all__ = ["MegatronAliasInput", "MegatronTensorSpec", "build_hf_aliases"] diff --git a/modelexpress_client/python/modelexpress_rl/train/engines/megatron/publisher.py b/modelexpress_client/python/modelexpress_rl/train/engines/megatron/publisher.py new file mode 100644 index 00000000..a6ab70c0 --- /dev/null +++ b/modelexpress_client/python/modelexpress_rl/train/engines/megatron/publisher.py @@ -0,0 +1,149 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Publish registered native Megatron tensors for RL refit.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from modelexpress.refit.reshard.rendezvous import ( + MxReshardRendezvous, + PublishedTensor, + wrap_rendezvous_blob, +) +from modelexpress_rl.train.adapter import NixlMetadataProvider + +from .aliases import MegatronTensorSpec, build_hf_aliases + + +@dataclass(frozen=True) +class MegatronReshardManifest: + """Serialized manifest and the tensor records used to build it.""" + + blob: bytes + tensors: tuple[PublishedTensor, ...] + + +@dataclass(frozen=True) +class MegatronPublishedTensorSpec: + """Legacy one-to-one native tensor publication descriptor.""" + + name: str + global_shape: tuple[int, ...] + shard_axis: int | None = None + local_shard_range: tuple[int, int] | None = None + + def __post_init__(self) -> None: + if not self.global_shape or any(int(dim) <= 0 for dim in self.global_shape): + raise ValueError(f"{self.name}: invalid global shape {self.global_shape}") + if (self.shard_axis is None) != (self.local_shard_range is None): + raise ValueError( + f"{self.name}: shard_axis and local_shard_range must be set together" + ) + + +def build_megatron_reshard_manifest( + *, + manager: NixlMetadataProvider, + published: list[PublishedTensor], + metadata_endpoint: str, +) -> MegatronReshardManifest: + """Describe already-registered Megatron tensors without publishing them.""" + + if not metadata_endpoint or ":" not in metadata_endpoint: + raise ValueError( + "metadata_endpoint must be an explicit host:port reachable by receivers" + ) + agent_name = str(manager.agent_name) + if not published: + raise ValueError("published tensor list must not be empty") + names = set() + for tensor in published: + if tensor.name in names: + raise ValueError(f"duplicate published tensor {tensor.name!r}") + names.add(tensor.name) + if not tensor.shards: + raise ValueError(f"{tensor.name}: no shards were published") + for shard in tensor.shards: + if shard.agent_name != agent_name: + raise ValueError( + f"{tensor.name}: shard agent {shard.agent_name!r} does not " + f"match manager agent {agent_name!r}" + ) + if shard.addr <= 0: + raise ValueError(f"{tensor.name}: shard has invalid address") + + blob = wrap_rendezvous_blob( + manager.nixl_metadata, + agent_name, + metadata_endpoint, + published, + ) + return MegatronReshardManifest(blob=blob, tensors=tuple(published)) + + +def publish_megatron_reshard_view( + *, + manager: NixlMetadataProvider, + rendezvous: MxReshardRendezvous, + tensors: dict[str, Any], + specs: list[MegatronPublishedTensorSpec], + metadata_endpoint: str, +) -> str: + """Publish registered Megatron tensors through the existing rendezvous.""" + by_name = {spec.name: spec for spec in specs} + if len(by_name) != len(specs): + raise ValueError("duplicate Megatron publish spec") + missing = sorted(set(by_name).difference(tensors)) + extra = sorted(set(tensors).difference(by_name)) + if missing or extra: + raise ValueError( + f"Megatron shard table/tensor mismatch: missing={missing[:10]} " + f"extra={extra[:10]}" + ) + items = [ + MegatronTensorSpec( + name=name, + tensor=tensors[name], + role="column", + hf_names=(name,), + global_shape=spec.global_shape, + placement_kind="SHARD" if spec.shard_axis is not None else "REPLICATE", + shard_axis=spec.shard_axis, + local_shard_range=spec.local_shard_range, + ) + for name, spec in sorted(by_name.items()) + ] + published = build_hf_aliases(items, agent_name=str(manager.agent_name)) + manifest = build_megatron_reshard_manifest( + manager=manager, + published=published, + metadata_endpoint=metadata_endpoint, + ) + return rendezvous.publish(manifest.blob) + + +def publish_registered_shard_table( + *, + manager: NixlMetadataProvider, + rendezvous: MxReshardRendezvous, + published: list[PublishedTensor], + metadata_endpoint: str, +) -> str: + """Publish validated aliases through the existing rendezvous.""" + manifest = build_megatron_reshard_manifest( + manager=manager, + published=published, + metadata_endpoint=metadata_endpoint, + ) + return rendezvous.publish(manifest.blob) + + +__all__ = [ + "MegatronPublishedTensorSpec", + "MegatronReshardManifest", + "build_megatron_reshard_manifest", + "publish_megatron_reshard_view", + "publish_registered_shard_table", +] diff --git a/modelexpress_client/python/modelexpress_rl/train/manifest.py b/modelexpress_client/python/modelexpress_rl/train/manifest.py new file mode 100644 index 00000000..babdea98 --- /dev/null +++ b/modelexpress_client/python/modelexpress_rl/train/manifest.py @@ -0,0 +1,66 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Worker-local serving for versioned trainer manifests.""" + +from __future__ import annotations + +import threading + +import grpc + +from .. import refit_pb2, refit_pb2_grpc +from .adapter import WeightVersionShardManifest + + +class WeightVersionShardManifestService(refit_pb2_grpc.RefitWorkerServiceServicer): + """Publish and serve immutable manifests from one trainer process. + + The records intentionally share the worker process lifetime. Durable + version and shard metadata remains in the central RefitService backend; + large tensor buffers and their manifest remain worker-local. + """ + + def __init__(self, *, endpoint: str) -> None: + if not endpoint.strip(): + raise ValueError("endpoint is required") + self.endpoint = endpoint + self._manifests: dict[tuple[str, str], WeightVersionShardManifest] = {} + self._lock = threading.Lock() + + def publish_manifest( + self, + *, + version_id: str, + source_slot_id: str, + manifest: WeightVersionShardManifest, + ) -> str: + """Make one immutable manifest retrievable before returning its endpoint.""" + if not version_id.strip(): + raise ValueError("version_id is required") + if not source_slot_id.strip(): + raise ValueError("source_slot_id is required") + key = (version_id, source_slot_id) + with self._lock: + existing = self._manifests.get(key) + if existing is not None and existing != manifest: + raise ValueError( + "a different manifest is already published for " + f"version_id={version_id!r}, source_slot_id={source_slot_id!r}" + ) + self._manifests[key] = manifest + return self.endpoint + + def GetWeightVersionShardManifest(self, request, context): + key = (request.version_id, request.source_slot_id) + with self._lock: + manifest = self._manifests.get(key) + if manifest is None: + context.abort(grpc.StatusCode.NOT_FOUND, "manifest was not found") + return refit_pb2.GetWeightVersionShardManifestResponse( + manifest=manifest.data, + manifest_digest=manifest.digest, + ) + + +__all__ = ["WeightVersionShardManifestService"] diff --git a/modelexpress_client/python/tests/test_refit_envs.py b/modelexpress_client/python/tests/test_refit_envs.py new file mode 100644 index 00000000..c41ca750 --- /dev/null +++ b/modelexpress_client/python/tests/test_refit_envs.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for ModelExpress RL-specific environment variables.""" + +import pytest +from modelexpress_rl import envs + + +def test_defaults_when_unset(monkeypatch): + for name in envs.environment_variables: + monkeypatch.delenv(name, raising=False) + + assert envs.MX_TRAINER_ENGINE == "MEGATRON" + assert envs.MX_TRAINER_STAGING_MODE == "IN_PLACE" + assert envs.MX_WEIGHT_PAYLOAD_FORMAT == "FULL_TENSOR" + + +def test_values_are_normalized_and_read_live(monkeypatch): + monkeypatch.setenv("MX_TRAINER_ENGINE", " megatron ") + monkeypatch.setenv("MX_TRAINER_STAGING_MODE", " copy_to_device ") + monkeypatch.setenv("MX_WEIGHT_PAYLOAD_FORMAT", " xor_delta ") + + assert envs.MX_TRAINER_ENGINE == "MEGATRON" + assert envs.MX_TRAINER_STAGING_MODE == "COPY_TO_DEVICE" + assert envs.MX_WEIGHT_PAYLOAD_FORMAT == "XOR_DELTA" + + +def test_unknown_attribute_raises(): + with pytest.raises(AttributeError): + _ = envs.NOT_A_REAL_ENV_VAR + + +def test_dir_lists_registered_names(): + assert set(envs.environment_variables).issubset(dir(envs)) diff --git a/modelexpress_client/python/tests/test_refit_megatron_adapter.py b/modelexpress_client/python/tests/test_refit_megatron_adapter.py new file mode 100644 index 00000000..5cd2d18c --- /dev/null +++ b/modelexpress_client/python/tests/test_refit_megatron_adapter.py @@ -0,0 +1,188 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import hashlib +from concurrent import futures +from types import SimpleNamespace + +import grpc +import pytest + +from modelexpress.refit.reshard.rendezvous import unwrap_rendezvous_blob +from modelexpress_rl import ( + ModelExpressTrainerClient, + TrainerEngineAdapter, + TrainerStagingMode, + WeightPayloadFormat, + WeightVersionRef, + WeightVersionShardManifestService, + refit_pb2, + refit_pb2_grpc, +) +from modelexpress_rl.train.engines.megatron import ( + MegatronTensorSpec, + MegatronTrainerAdapter, +) + + +class _Tensor: + shape = (8, 8) + dtype = "torch.bfloat16" + device = SimpleNamespace(index=0) + + def is_contiguous(self): + return True + + def element_size(self): + return 2 + + def data_ptr(self): + return 0x1234 + + +class _Manager: + agent_name = "trainer-r3" + nixl_metadata = b"agent-metadata" + listen_port = 19003 + + +class _RefitService(refit_pb2_grpc.RefitServiceServicer): + def __init__(self, events): + self.events = events + self.registration_ttl = None + self.shard = None + + def RegisterWorker(self, request, _context): + self.events.append("register-worker") + self.registration_ttl = request.ttl_seconds + return request.worker + + def CreateWeightVersionShard(self, request, _context): + self.events.append("publish-version-shard") + self.shard = request.shard + return refit_pb2.CreateWeightVersionShardResponse( + shard=request.shard, + version=refit_pb2.WeightVersion( + uid=request.shard.version_id, + state=refit_pb2.WEIGHT_VERSION_STATE_READY, + ), + ) + + +def test_megatron_adapter_requires_initialized_distributed_engine(monkeypatch): + monkeypatch.setattr( + "modelexpress_rl.train.engines.megatron.adapter.dist.is_initialized", + lambda: False, + ) + + with pytest.raises(RuntimeError, match="distributed process group"): + MegatronTrainerAdapter( + manager=_Manager(), + nixl_metadata_endpoint="10.0.0.3:19003", + ) + + +def test_megatron_adapter_uses_shared_trainer_publication_flow(monkeypatch): + monkeypatch.setattr( + "modelexpress_rl.train.engines.megatron.adapter.dist.is_initialized", + lambda: True, + ) + monkeypatch.setattr( + "modelexpress_rl.train.engines.megatron.adapter.dist.get_rank", + lambda: 3, + ) + monkeypatch.setenv("MX_WORKER_HOST", "10.0.0.3") + events = [] + refit_service = _RefitService(events) + server = grpc.server(futures.ThreadPoolExecutor(max_workers=2)) + refit_pb2_grpc.add_RefitServiceServicer_to_server(refit_service, server) + port = server.add_insecure_port("127.0.0.1:0") + manifest_service = WeightVersionShardManifestService(endpoint=f"127.0.0.1:{port}") + refit_pb2_grpc.add_RefitWorkerServiceServicer_to_server(manifest_service, server) + server.start() + adapter = MegatronTrainerAdapter( + manager=_Manager(), + nixl_metadata_endpoint="10.0.0.3:19003", + ) + tensors = [ + MegatronTensorSpec( + name="column", + tensor=_Tensor(), + role="column", + hf_names=("column",), + global_shape=(16, 8), + placement_kind="SHARD", + shard_axis=0, + local_shard_range=(8, 16), + ) + ] + + with pytest.raises(NotImplementedError, match="COPY_TO_DEVICE staging"): + adapter.stage_shard( + tensors=tensors, + staging_mode=TrainerStagingMode.COPY_TO_DEVICE, + payload_format=WeightPayloadFormat.FULL_TENSOR, + ) + with pytest.raises(NotImplementedError, match="XOR_DELTA payloads"): + adapter.stage_shard( + tensors=tensors, + staging_mode=TrainerStagingMode.IN_PLACE, + payload_format=WeightPayloadFormat.XOR_DELTA, + ) + + try: + refit_client = ModelExpressTrainerClient.initialize( + server_url=f"127.0.0.1:{port}", + manager=_Manager(), + model_name="model", + staging_mode=TrainerStagingMode.IN_PLACE, + payload_format=WeightPayloadFormat.FULL_TENSOR, + manifest_publisher=manifest_service, + worker_endpoint="trainer-3:9000", + worker_id="worker-3", + registration_ttl_seconds=60, + ) + staged = refit_client.stage_shard( + version=WeightVersionRef("version-a"), + tensors=tensors, + ) + with pytest.raises(NotImplementedError, match="version-retirement"): + staged.source_reuse_ready.wait() + staged.publish() + worker_stub = refit_pb2_grpc.RefitWorkerServiceStub( + grpc.insecure_channel(refit_service.shard.manifest_endpoint) + ) + fetched = worker_stub.GetWeightVersionShardManifest( + refit_pb2.GetWeightVersionShardManifestRequest( + version_id="version-a", + source_slot_id="publisher:global-rank:3", + ) + ) + finally: + if "refit_client" in locals(): + refit_client.close() + server.stop(grace=None).wait() + + assert isinstance(adapter, TrainerEngineAdapter) + assert isinstance(refit_client._adapter, MegatronTrainerAdapter) + assert events == [ + "register-worker", + "publish-version-shard", + ] + assert refit_service.registration_ttl == 60 + assert refit_service.shard.version_id == "version-a" + assert refit_service.shard.source_slot_id == "publisher:global-rank:3" + assert refit_service.shard.worker_id == "worker-3" + assert refit_service.shard.tensor_count == 1 + assert refit_service.shard.total_bytes == 128 + assert ( + refit_service.shard.manifest_digest + == hashlib.sha256(fetched.manifest).hexdigest() + ) + assert refit_service.shard.manifest_endpoint == f"127.0.0.1:{port}" + assert refit_service.shard.transport == "NIXL" + assert fetched.manifest_digest == refit_service.shard.manifest_digest + payload = unwrap_rendezvous_blob(fetched.manifest) + assert payload.metadata_endpoint == "10.0.0.3:19003" + assert payload.tensors[0].shards[0].addr == 0x1234 + assert payload.tensors[0].shards[0].shard_offset == (8, 0) diff --git a/modelexpress_client/python/tests/test_refit_trainer_client.py b/modelexpress_client/python/tests/test_refit_trainer_client.py new file mode 100644 index 00000000..17d53322 --- /dev/null +++ b/modelexpress_client/python/tests/test_refit_trainer_client.py @@ -0,0 +1,217 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import time +from concurrent import futures + +import grpc +import pytest + +import modelexpress_rl.client as client_module +from modelexpress_rl import ( + CompletionFence, + ModelExpressTrainerClient, + StagedWeightVersionShardData, + TrainerEngineAdapter, + TrainerStagingMode, + WeightPayloadFormat, + WeightVersionRef, + WeightVersionShardManifest, + WeightVersionShardManifestService, + refit_pb2, + refit_pb2_grpc, +) + + +class _RefitService(refit_pb2_grpc.RefitServiceServicer): + def __init__(self): + self.registrations = {} + self.registration_count = 0 + self.shards = [] + + def RegisterWorker(self, request, _context): + self.registration_count += 1 + worker = request.worker + worker.expires_at_unix_ms = 1234 + self.registrations[worker.worker_id] = worker + return worker + + def CreateWeightVersionShard(self, request, context): + shard = request.shard + if shard.worker_id not in self.registrations: + context.abort(grpc.StatusCode.FAILED_PRECONDITION, "worker not registered") + self.shards.append(shard) + return refit_pb2.CreateWeightVersionShardResponse( + shard=shard, + version=refit_pb2.WeightVersion( + uid=shard.version_id, + state=refit_pb2.WEIGHT_VERSION_STATE_READY, + ), + ) + + +class _Manager: + listen_port = 19000 + + +class _Adapter(TrainerEngineAdapter): + source_slot_id = "rank:0" + supported_staging_modes = frozenset({TrainerStagingMode.COPY_TO_DEVICE}) + supported_payload_formats = frozenset({WeightPayloadFormat.FULL_TENSOR}) + + def __init__(self): + self.calls = [] + + def stage_shard(self, *, tensors, staging_mode, payload_format): + self.calls.append((tensors, staging_mode, payload_format)) + + return StagedWeightVersionShardData( + manifest=WeightVersionShardManifest( + data=b"manifest", + tensor_count=2, + total_bytes=128, + transport="NIXL", + ), + publish_ready=CompletionFence(lambda: None), + source_reuse_ready=CompletionFence(lambda: None), + buffer_owner=tensors, + ) + + +def test_trainer_stages_then_publishes_one_rank_local_shard(monkeypatch): + service = _RefitService() + server = grpc.server(futures.ThreadPoolExecutor(max_workers=2)) + refit_pb2_grpc.add_RefitServiceServicer_to_server(service, server) + port = server.add_insecure_port("127.0.0.1:0") + manifest_service = WeightVersionShardManifestService(endpoint=f"127.0.0.1:{port}") + refit_pb2_grpc.add_RefitWorkerServiceServicer_to_server(manifest_service, server) + server.start() + adapter = _Adapter() + monkeypatch.setattr(client_module, "_trainer_adapter", lambda **_kwargs: adapter) + monkeypatch.setenv("MODEL_NAME", "test/model") + monkeypatch.setenv("MX_TRAINER_STAGING_MODE", "COPY_TO_DEVICE") + monkeypatch.setenv("MX_WEIGHT_PAYLOAD_FORMAT", "FULL_TENSOR") + monkeypatch.setenv("MX_WORKER_HOST", "127.0.0.1") + monkeypatch.setenv("MX_WORKER_GRPC_PORT", str(port)) + + try: + trainer = ModelExpressTrainerClient.initialize( + manager=_Manager(), + manifest_publisher=manifest_service, + worker_id="trainer-0", + server_url=f"127.0.0.1:{port}", + registration_ttl_seconds=1, + ) + deadline = time.monotonic() + 10.0 + while service.registration_count < 2 and time.monotonic() < deadline: + time.sleep(0.02) + assert service.registration_count >= 2 + shard = trainer.stage_shard( + version=WeightVersionRef("version-a"), + tensors="model", + ) + + assert service.shards == [] + shard.source_reuse_ready.wait() + shard.publish() + shard.publish() + + worker_stub = refit_pb2_grpc.RefitWorkerServiceStub( + grpc.insecure_channel(service.shards[0].manifest_endpoint) + ) + fetched = worker_stub.GetWeightVersionShardManifest( + refit_pb2.GetWeightVersionShardManifestRequest( + version_id="version-a", + source_slot_id="rank:0", + ) + ) + second = trainer.stage_shard( + version=WeightVersionRef("version-a"), + tensors="model-2", + ) + second.publish() + retained_owners = [ + staged.buffer_owner for staged in trainer._published_shards["version-a"] + ] + finally: + if "trainer" in locals(): + trainer.close() + server.stop(grace=None).wait() + + assert adapter.calls == [ + ( + "model", + TrainerStagingMode.COPY_TO_DEVICE, + WeightPayloadFormat.FULL_TENSOR, + ), + ( + "model-2", + TrainerStagingMode.COPY_TO_DEVICE, + WeightPayloadFormat.FULL_TENSOR, + ), + ] + assert retained_owners == ["model", "model-2"] + assert len(service.shards) == 2 + assert service.registrations["trainer-0"].role == refit_pb2.WORKER_ROLE_TRAINER + assert ( + service.registrations["trainer-0"].endpoint + == service.shards[0].manifest_endpoint + ) + assert service.shards[0].version_id == "version-a" + assert service.shards[0].source_slot_id == "rank:0" + assert service.shards[0].worker_id == "trainer-0" + assert service.shards[0].tensor_count == 2 + assert service.shards[0].total_bytes == 128 + assert fetched.manifest == b"manifest" + + +def test_trainer_initialization_rejects_unspecified_fixed_settings(monkeypatch): + monkeypatch.setenv("MX_WORKER_HOST", "trainer") + with pytest.raises(ValueError, match="staging_mode must be specified"): + ModelExpressTrainerClient.initialize( + model_name="test/model", + manager=_Manager(), + staging_mode=TrainerStagingMode.UNSPECIFIED, + payload_format=WeightPayloadFormat.FULL_TENSOR, + manifest_publisher=object(), + worker_endpoint="trainer:9000", + ) + + with pytest.raises(ValueError, match="payload_format must be specified"): + ModelExpressTrainerClient.initialize( + model_name="test/model", + manager=_Manager(), + staging_mode=TrainerStagingMode.COPY_TO_DEVICE, + payload_format=WeightPayloadFormat.UNSPECIFIED, + manifest_publisher=object(), + worker_endpoint="trainer:9000", + ) + + +def test_trainer_initialization_rejects_adapter_unsupported_mode(monkeypatch): + adapter = _Adapter() + monkeypatch.setattr(client_module, "_trainer_adapter", lambda **_kwargs: adapter) + monkeypatch.setenv("MX_WORKER_HOST", "trainer") + with pytest.raises(ValueError, match="does not support staging mode IN_PLACE"): + ModelExpressTrainerClient.initialize( + model_name="test/model", + manager=_Manager(), + staging_mode=TrainerStagingMode.IN_PLACE, + payload_format=WeightPayloadFormat.FULL_TENSOR, + manifest_publisher=object(), + worker_endpoint="trainer:9000", + ) + + +def test_trainer_initialization_rejects_unknown_configured_engine(monkeypatch): + monkeypatch.setenv("MX_TRAINER_ENGINE", "unknown") + monkeypatch.setenv("MX_WORKER_HOST", "trainer") + with pytest.raises(ValueError, match="unsupported MX_TRAINER_ENGINE='UNKNOWN'"): + ModelExpressTrainerClient.initialize( + model_name="test/model", + manager=_Manager(), + staging_mode=TrainerStagingMode.IN_PLACE, + payload_format=WeightPayloadFormat.FULL_TENSOR, + manifest_publisher=object(), + worker_endpoint="trainer:9000", + ) diff --git a/modelexpress_client/python/tests/test_reshard_megatron.py b/modelexpress_client/python/tests/test_reshard_megatron.py index 047b799c..5e649c1c 100644 --- a/modelexpress_client/python/tests/test_reshard_megatron.py +++ b/modelexpress_client/python/tests/test_reshard_megatron.py @@ -1,30 +1,31 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from importlib import import_module + import pytest import torch - -from modelexpress import p2p_pb2 -from modelexpress.refit.reshard.megatron import ( +from modelexpress.refit.reshard.rendezvous import ( + MxReshardRendezvous, + unwrap_rendezvous_blob, +) +from modelexpress.refit.reshard.slice_plan import Shard +from modelexpress.refit.reshard.transfer_plan import SourceInfo, plan_transfer +from modelexpress_rl.inference.reshard.megatron import ( + MegatronReshardReceiver, MegatronTargetLayout, MegatronTargetSpec, lower_megatron_target, ) -from modelexpress.refit.reshard.megatron_aliases import ( +from modelexpress_rl.train.engines.megatron import ( MegatronAliasInput, - build_hf_aliases, -) -from modelexpress.refit.reshard.megatron_publisher import ( MegatronPublishedTensorSpec, + MegatronTensorSpec, + build_hf_aliases, + build_megatron_reshard_manifest, publish_megatron_reshard_view, + publish_registered_shard_table, ) -from modelexpress.refit.reshard.megatron_receiver import MegatronReshardReceiver -from modelexpress.refit.reshard.rendezvous import ( - MxReshardRendezvous, - unwrap_rendezvous_blob, -) -from modelexpress.refit.reshard.slice_plan import Shard -from modelexpress.refit.reshard.transfer_plan import SourceInfo, plan_transfer def _bf16_sources() -> tuple[dict[str, SourceInfo], list[torch.Tensor]]: @@ -156,6 +157,42 @@ def test_non_divisible_target_geometry_fails_closed(): ) +@pytest.mark.parametrize("extent", [4.5, True]) +def test_target_spec_rejects_non_integral_global_shape(extent): + with pytest.raises(ValueError, match="positive integer extents"): + MegatronTargetSpec("column", "column", (extent, 8), torch.bfloat16) + + +@pytest.mark.parametrize( + ("tp_size", "tp_rank", "message"), + [ + (2.5, 0, "tp_size must be an integer"), + (True, 0, "tp_size must be an integer"), + (2, 0.5, "tp_rank must be an integer"), + (2, False, "tp_rank must be an integer"), + ], +) +def test_target_layout_rejects_non_integral_geometry(tp_size, tp_rank, message): + with pytest.raises(ValueError, match=message): + MegatronTargetLayout(tp_size=tp_size, tp_rank=tp_rank) + + +def test_legacy_megatron_publisher_imports_remain_available(): + legacy_aliases = import_module( + "modelexpress.refit.reshard.megatron_aliases" + ) + legacy_publisher = import_module( + "modelexpress.refit.reshard.megatron_publisher" + ) + + assert legacy_aliases.MegatronAliasInput is MegatronAliasInput + assert legacy_aliases.build_hf_aliases is build_hf_aliases + assert ( + legacy_publisher.publish_registered_shard_table + is publish_registered_shard_table + ) + + def test_receiver_seam_validates_manifest_and_invokes_installer(): installed = [] receiver = object.__new__(MegatronReshardReceiver) @@ -194,36 +231,66 @@ def test_receiver_seam_rejects_stale_manifest_geometry(): ) -class _PublishClient: - """Records the published worker record and accepts heartbeat status updates.""" +class _Manager: + agent_name = "trainer-r3" + nixl_metadata = b"agent-metadata" + +class _PublishClient: def __init__(self): self.worker = None - self.status_updates = [] def publish_metadata(self, _identity, worker, _worker_id): self.worker = worker return "source-id" - def update_status(self, **kwargs): - self.status_updates.append(kwargs) + def update_status(self, **_kwargs): return True -class _Manager: - agent_name = "trainer-r3" - nixl_metadata = b"agent-metadata" +def test_manifest_builder_reuses_registered_tensor_addresses(): + tensor = torch.zeros((8, 8), dtype=torch.bfloat16) + published = build_hf_aliases( + [ + MegatronTensorSpec( + name="column", + tensor=tensor, + role="column", + hf_names=("column",), + global_shape=(16, 8), + placement_kind="SHARD", + shard_axis=0, + local_shard_range=(8, 16), + ) + ], + agent_name="trainer-r3", + ) + manifest = build_megatron_reshard_manifest( + manager=_Manager(), + published=published, + metadata_endpoint="10.0.0.3:19003", + ) -def _trainer_rendezvous(client): - return MxReshardRendezvous( - client, role="trainer", rank=3, model_name="model", worker_id="worker-3" + payload = unwrap_rendezvous_blob(manifest.blob) + assert (payload.agent_metadata, payload.agent_name, payload.metadata_endpoint) == ( + b"agent-metadata", + "trainer-r3", + "10.0.0.3:19003", ) + assert payload.tensors[0].shards[0].addr == tensor.data_ptr() + assert payload.tensors[0].shards[0].shard_offset == (8, 0) -def test_publisher_seam_reuses_registered_tensor_addresses(): +def test_existing_reshard_publisher_remains_compatible(): client = _PublishClient() - rendezvous = _trainer_rendezvous(client) + rendezvous = MxReshardRendezvous( + client, + role="trainer", + rank=3, + model_name="model", + worker_id="worker-3", + ) tensor = torch.zeros((8, 8), dtype=torch.bfloat16) try: @@ -242,57 +309,35 @@ def test_publisher_seam_reuses_registered_tensor_addresses(): metadata_endpoint="10.0.0.3:19003", ) finally: - # The caller owns the rendezvous precisely so its heartbeat can be stopped. rendezvous.close() assert source_id == "source-id" - assert client.worker.status > 0 payload = unwrap_rendezvous_blob(client.worker.nixl_metadata) - assert (payload.agent_metadata, payload.agent_name, payload.metadata_endpoint) == ( - b"agent-metadata", - "trainer-r3", - "10.0.0.3:19003", - ) assert payload.tensors[0].shards[0].addr == tensor.data_ptr() - assert payload.tensors[0].shards[0].shard_offset == (8, 0) - -def test_publishing_leaves_the_heartbeat_with_its_owner(): - """Publishing starts the source's READY heartbeat. A rendezvous built inside the - seam would leave that thread running with no handle to stop it, so the source - would only be marked stale at interpreter exit.""" - client = _PublishClient() - rendezvous = _trainer_rendezvous(client) - publish_megatron_reshard_view( - manager=_Manager(), - rendezvous=rendezvous, - tensors={"column": torch.zeros((8, 8), dtype=torch.bfloat16)}, - specs=[MegatronPublishedTensorSpec(name="column", global_shape=(8, 8))], - metadata_endpoint="10.0.0.3:19003", +def test_manifest_builder_rejects_duplicate_tensor_names(): + tensor = torch.zeros((8, 8), dtype=torch.bfloat16) + published = build_hf_aliases( + [ + MegatronTensorSpec( + name="column", + tensor=tensor, + role="column", + hf_names=("column",), + global_shape=(8, 8), + placement_kind="REPLICATE", + shard_axis=None, + local_shard_range=None, + ) + ], + agent_name="trainer-r3", ) - rendezvous.close() - assert client.status_updates[-1]["status"] == p2p_pb2.SOURCE_STATUS_STALE - assert client.status_updates[-1]["worker_id"] == "worker-3" - - -def test_publisher_seam_rejects_duplicate_spec_names(): - """Last-writer-wins would publish one spec's shard description under a name the - other spec owns, and comparing key sets against the tensors cannot see it.""" - - class Client: - def publish_metadata(self, *_args, **_kwargs): - raise AssertionError("publication must not run") - - spec = MegatronPublishedTensorSpec(name="column", global_shape=(8, 8)) - - with pytest.raises(ValueError, match="duplicate Megatron publish spec"): - publish_megatron_reshard_view( + with pytest.raises(ValueError, match="duplicate published tensor"): + build_megatron_reshard_manifest( manager=_Manager(), - rendezvous=_trainer_rendezvous(Client()), - tensors={"column": torch.zeros((8, 8), dtype=torch.bfloat16)}, - specs=[spec, spec], + published=[published[0], published[0]], metadata_endpoint="10.0.0.3:19003", ) @@ -371,6 +416,30 @@ def test_a_missing_fused_gate_up_order_is_rejected(): ) +def test_gated_aliases_reject_inconsistent_declared_global_shape(): + fused = torch.arange(32, dtype=torch.bfloat16).reshape(8, 4) + + with pytest.raises( + ValueError, match=r"linear_fc1\.weight: derived gate/up shape" + ): + build_hf_aliases( + [ + MegatronAliasInput( + name="linear_fc1.weight", + tensor=fused, + role="gated_mlp_column", + hf_names=("gate_proj.weight", "up_proj.weight"), + global_shape=(16, 5), + placement_kind="SHARD", + shard_axis=0, + local_shard_range=(8, 16), + extras={"gated_mlp_order": "gate_then_up"}, + ) + ], + agent_name="trainer-tp1", + ) + + def test_qkv_aliases_expose_hf_head_ranges_without_copy(): qkv = torch.arange(48, dtype=torch.bfloat16).reshape(12, 4) @@ -406,6 +475,31 @@ def test_qkv_aliases_expose_hf_head_ranges_without_copy(): assert v.shards[0].addr == qkv[10:].data_ptr() +def test_qkv_aliases_report_missing_extras_with_tensor_name(): + qkv = torch.arange(48, dtype=torch.bfloat16).reshape(12, 4) + + with pytest.raises( + ValueError, + match=r"linear_qkv\.weight: QKV aliasing requires extras", + ): + build_hf_aliases( + [ + MegatronAliasInput( + name="linear_qkv.weight", + tensor=qkv, + role="qkv_column", + hf_names=("q_proj.weight", "k_proj.weight", "v_proj.weight"), + global_shape=(24, 4), + placement_kind="SHARD", + shard_axis=0, + local_shard_range=(12, 24), + extras={"head_dim": "2"}, + ) + ], + agent_name="trainer-tp1", + ) + + def test_a_shard_range_beyond_the_global_extent_is_rejected(): """A range of (16, 24) against a global extent of 16 clears every other check: it is the right width, it divides evenly, and it yields source rank 2 diff --git a/modelexpress_common/proto/refit.proto b/modelexpress_common/proto/refit.proto index 5c3844ae..6f511022 100644 --- a/modelexpress_common/proto/refit.proto +++ b/modelexpress_common/proto/refit.proto @@ -25,6 +25,12 @@ service RefitService { rpc DeleteVersionLease(DeleteVersionLeaseRequest) returns (DeleteVersionLeaseResponse); } +// Internal worker-to-worker API. The generator fetches the small transfer +// manifest here; tensor bytes remain on the advertised data-plane transport. +service RefitWorkerService { + rpc GetWeightVersionShardManifest(GetWeightVersionShardManifestRequest) returns (GetWeightVersionShardManifestResponse); +} + enum WorkerRole { WORKER_ROLE_UNSPECIFIED = 0; WORKER_ROLE_TRAINER = 1; @@ -129,6 +135,16 @@ message ListWeightVersionShardsResponse { repeated WeightVersionShard shards = 1; } +message GetWeightVersionShardManifestRequest { + string version_id = 1; + string source_slot_id = 2; +} + +message GetWeightVersionShardManifestResponse { + bytes manifest = 1; + string manifest_digest = 2; +} + message DeleteWeightVersionShardRequest { string version_id = 1; string source_slot_id = 2;