feat(refit): add RL trainer client and Megatron adapter

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.

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:

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}"

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}.'

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."""

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}.'

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",