feat(weight_transfer): add trainer-inference weight sync with M2N and 2-D TP sharding#481
feat(weight_transfer): add trainer-inference weight sync with M2N and 2-D TP sharding#481ganeshku1 wants to merge 1 commit into
Conversation
… 2-D TP sharding - Protocol types: TrainerShard, TrainerTable, TrainerTensor, ResolvedRegion - 2-D tile sharding (TP x FSDP) with col_end=-1 for row-only backward compat - route_regions() pure-math router for RDMA descriptor generation - LocalPlanner and M2nPlanner with server coordination and fallback - PullRole and PushRole for bake/resolve/plan/execute lifecycle - NixlExecutor and M2nExecutor for RDMA transport - BakeRecorder / LazyWeight for op chain recording - MoEAdapter for stacked expert weight expansion - 58 pure-Python e2e tests; 17 real-model GPU tests (Qwen2.5-0.5B) Signed-off-by: Ganesh Kudleppanavar <ganeshku@nvidia.com>
WalkthroughThis PR adds a full trainer-inference weight synchronization system using NIXL RDMA. It introduces protocol types, lazy-weight baking/resolution, region routers, Local/Server/M2N planners, Pull/Push sync roles, transport executors, a new gRPC WeightSyncService (proto + Rust server implementation), client load-strategy integration, documentation, and extensive tests. ChangesTrainer-Inference Weight Sync
Estimated code review effort: 5 (Critical) | ~150 minutes Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Comment |
c41d4b0 to
9fb93e3
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 20
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
modelexpress_client/python/modelexpress/load_strategy/__init__.py (1)
64-76: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winGate
TrainerPullStrategybehind an explicit opt-in
is_available()only checks NIXL/RDMA/mx_client, butload()always waits up toMX_TRAINER_SYNC_TIMEOUTfor a TrainerTable.MX_WEIGHT_SYNC_SERVERonly changes the planner; it does not stop this strategy from being tried. In ordinary RDMA deployments that satisfy those checks, startup can stall for up to 300s before falling through toRdmaStrategy.Gate this on an explicit trainer-pull signal such as
MX_TRAINER_TABLE_KEYorMX_WEIGHT_SYNC_SERVERso non-trainer-pull loads skip it immediately.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/python/modelexpress/load_strategy/__init__.py` around lines 64 - 76, TrainerPullStrategy is being selected too broadly because load_strategy.__init__ always includes it when is_available() passes, causing non-trainer-pull deployments to wait in TrainerPullStrategy.load() for MX_TRAINER_SYNC_TIMEOUT. Update the strategy selection in the load strategy initialization so TrainerPullStrategy is only added when an explicit trainer-pull signal is present, such as MX_TRAINER_TABLE_KEY or MX_WEIGHT_SYNC_SERVER. Keep RdmaStrategy and the other fallbacks unchanged so ordinary RDMA startup skips TrainerPullStrategy immediately.
♻️ Duplicate comments (1)
modelexpress_client/python/modelexpress/weight_transfer/transport/nixl_executor.py (1)
94-104: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winSame handle-leak pattern as
M2nExecutor.execute.
agent.transfer(handle)at Line 102 can raise aftermake_prepped_xferalready allocated the handle, before it's appended tohandlesfor release — same issue flagged intransport/m2n_executor.py.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/python/modelexpress/weight_transfer/transport/nixl_executor.py` around lines 94 - 104, The transfer path in NixlExecutor has the same handle-leak risk as M2nExecutor.execute: if agent.transfer(handle) fails after agent.make_prepped_xfer() succeeds, the newly created handle is never recorded for cleanup. Update NixlExecutor.execute so the handle is protected immediately after creation, either by appending it to handles before calling transfer or by using a try/finally around agent.transfer(handle) that guarantees release tracking even on failure. Use the make_prepped_xfer, transfer, and handles logic in NixlExecutor.execute to locate the fix.
🧹 Nitpick comments (14)
modelexpress_client/python/pyproject.toml (1)
68-73: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
gpumarker defined but unused.
tests/gpu/test_e2e_primerl_real_model.pygates itself withpytest.mark.skipif(not _TWO_GPUS, ...)only — it never applies@pytest.mark.gpu. As-is, the marker can't be used for CI selection (pytest -m "not gpu"), andskipifalready causes a module-level skip only after import, whereas the marker would let CI deselect it upfront.♻️ Suggested fix
-_TWO_GPUS = torch.cuda.is_available() and torch.cuda.device_count() >= 2 -pytestmark = pytest.mark.skipif(not _TWO_GPUS, reason="requires 2 CUDA GPUs") +_TWO_GPUS = torch.cuda.is_available() and torch.cuda.device_count() >= 2 +pytestmark = [ + pytest.mark.gpu, + pytest.mark.skipif(not _TWO_GPUS, reason="requires 2 CUDA GPUs"), +]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/python/pyproject.toml` around lines 68 - 73, The gpu marker is declared in the pytest configuration but never applied to the GPU E2E test, so CI cannot select or exclude it by marker. Update tests/gpu/test_e2e_primerl_real_model.py to add `@pytest.mark.gpu` to the relevant test/class/module alongside the existing skipif gate, using the gpu marker name defined in pyproject.toml so pytest -m selection works as intended.modelexpress_client/python/tests/gpu/test_e2e_primerl_real_model.py (2)
22-26: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueHardcoded, environment-specific
libcudartpath.The first candidate path (
/usr/local/python-3.12.6/lib/python3.12/site-packages/nvidia/cuda_runtime/lib/libcudart.so.12) is tied to one specific Python/CUDA install layout. The function does fall back to a barelibcudart.so.12and skips gracefully if none load, so this is low risk, but consider dropping the overly-specific first path (or deriving it fromtorch.version.cuda/ctypes.util.find_library) to avoid stale references as environments change.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/python/tests/gpu/test_e2e_primerl_real_model.py` around lines 22 - 26, The hardcoded first entry in the _LIBCUDART_PATHS list is too environment-specific and should be removed or replaced with a dynamic lookup. Update the test setup around _LIBCUDART_PATHS in test_e2e_primerl_real_model.py to prefer a derived libcudart location using something like ctypes.util.find_library or values inferred from torch.version.cuda, while keeping the existing fallback and graceful-skip behavior if no library is found.
188-193: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winBlind
except Exception: continuesilently drops unresolvable regions without logging.If
resolve_chain_regionfails for a parameter, that parameter is silently excluded from the transfer plan with no diagnostic, making failures harder to root-cause (static analysis flags this as S112/BLE001).♻️ Suggested fix
+ import logging + logger = logging.getLogger(__name__) try: offset, rshape, stride = resolve_chain_region( copy.op_chain, torch.Size(copy.dst_shape), copy.dst_dtype ) - except Exception: + except Exception: + logger.warning("Failed to resolve region for %s", copy.src_name, exc_info=True) continue🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/python/tests/gpu/test_e2e_primerl_real_model.py` around lines 188 - 193, The blind exception handling around resolve_chain_region in the parameter copy loop silently skips parameters without any diagnostics. Update the try/except in test_e2e_primerl_real_model to catch the failure in a more explicit way and emit a warning or debug log that includes the parameter/op_chain context before continuing, so unresolved regions are still visible when the transfer plan is built.Source: Linters/SAST tools
modelexpress_client/python/tests/test_e2e_primerl.py (1)
193-199: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winTautological test — doesn't exercise
PullRole.sync_and_post_process.This test only calls a bare
MagicMock()and checks its own call count; it can't detect a regression in the actual integration path wherePullRole.sync_and_post_process()invokesself._adapter.post_pull_hook(model)aftersync().Consider building a minimal
PullRolewith mocked planner/executor/adapter and assertingadapter.post_pull_hook.call_countafter realsync_and_post_process()calls, so the test verifies actual wiring rather than mock bookkeeping.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/python/tests/test_e2e_primerl.py` around lines 193 - 199, The current test is tautological because it only increments a MagicMock and never exercises PullRole.sync_and_post_process(). Update test_post_pull_hook_called_per_sync to construct a minimal PullRole with mocked planner, executor, and adapter, then call sync_and_post_process() and assert adapter.post_pull_hook.call_count so the test verifies the real post-pull wiring after sync rather than mock self-counting.modelexpress_client/python/modelexpress/weight_transfer/protocol/ops.py (1)
11-27: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win
OpSpecis frozen but not actually hashable.
kwargs: dictmakes the auto-generated__hash__(fromfrozen=True) raiseTypeErroronhash(dict)if anyone ever hashes anOpSpec/OpChain(e.g., as a cache/memoization key). Freezing normally implies hashability; this silently violates that contract.🔧 Proposed fix
-from dataclasses import dataclass +from dataclasses import dataclass, field @@ name: str args: tuple - kwargs: dict + kwargs: dict = field(hash=False)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/python/modelexpress/weight_transfer/protocol/ops.py` around lines 11 - 27, `OpSpec` is marked frozen but still contains an unhashable `dict` in `kwargs`, which breaks the expected hashability of the dataclass. Update `OpSpec` in `ops.py` so its state is normalized into immutable, hashable forms in `__post_init__` and/or provide a custom `__hash__` that hashes a stable representation of `name`, `args`, and `kwargs`; keep `OpChain` compatibility in mind if it relies on hashing these specs.modelexpress_client/python/modelexpress/weight_transfer/protocol/router.py (1)
149-193: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winConsider asserting src/dst element-count parity in
_zip_src_dst.If
src_triplesanddst_runsever have mismatched total element counts (e.g., a bug upstream inresolve_copies), the loop silently stops once either iterator is exhausted, dropping the remainder without any error - a silent partial-copy in a weight-sync path is hard to detect and could produce a subtly wrong model.🛡️ Proposed defensive check
while cur_shard is not None and dst_rem > 0: ... return descriptors + # After the loop, both sides should be fully consumed. + # if cur_shard is not None or dst_rem > 0: + # raise ValueError("src/dst element run totals do not match")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/python/modelexpress/weight_transfer/protocol/router.py` around lines 149 - 193, Add a defensive parity check in _zip_src_dst so source and destination element totals must match before building RdmaDescriptor entries. Compute/validate the total counts for src_triples and dst_runs, and raise an explicit error if they differ instead of allowing the while loop to stop early and drop leftover elements. Keep the fix localized to _zip_src_dst in router.py so upstream bugs in resolve_copies fail fast with a clear message.modelexpress_client/python/modelexpress/weight_transfer/adapters/__init__.py (1)
7-7: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSort
__all__(Ruff RUF022).♻️ Suggested fix
-__all__ = ["ModelAdapter", "MoEAdapter"] +__all__ = ["MoEAdapter", "ModelAdapter"]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/python/modelexpress/weight_transfer/adapters/__init__.py` at line 7, The __all__ export list in the adapters package is not sorted and triggers Ruff RUF022; update the __all__ definition in the adapters __init__ module so the exported names are ordered consistently, keeping the existing symbols ModelAdapter and MoEAdapter but arranging them in sorted order.Source: Linters/SAST tools
modelexpress_client/python/modelexpress/weight_transfer/__init__.py (1)
65-91: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSort
__all__(Ruff RUF022).♻️ Suggested ordering
__all__ = [ - # protocol - "OpSpec", - "OpChain", - "TrainerShard", - "TrainerTensor", - "TrainerTable", - "InferenceShard", - "InferenceTable", - "ResolvedRegion", - "RdmaDescriptor", - "SyncMode", - "encode_trainer_table", - "decode_trainer_table", - "encode_inference_table", - "decode_inference_table", - # engine - "BakeRecorder", - "LazyWeight", - "RecordedCopy", - "WeightLoaderAdapter", - # roles - "PullRole", - "PushRole", - # adapters - "MoEAdapter", + "BakeRecorder", + "InferenceShard", + "InferenceTable", + "LazyWeight", + "MoEAdapter", + "OpChain", + "OpSpec", + "PullRole", + "PushRole", + "RdmaDescriptor", + "RecordedCopy", + "ResolvedRegion", + "SyncMode", + "TrainerShard", + "TrainerTable", + "TrainerTensor", + "WeightLoaderAdapter", + "decode_inference_table", + "decode_trainer_table", + "encode_inference_table", + "encode_trainer_table", ]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/python/modelexpress/weight_transfer/__init__.py` around lines 65 - 91, The __all__ list in weight_transfer/__init__.py is not sorted and triggers Ruff RUF022. Reorder the exported symbols in __all__ into a consistent sorted order, keeping the existing grouped sections if needed, and update the list near the module-level __all__ definition so the import export order is alphabetical and lint-compliant.Source: Linters/SAST tools
modelexpress_client/python/modelexpress/weight_transfer/adapters/moe.py (2)
24-27: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMutable class-level default dict is shared across instances.
_default_proj_mapis a mutabledictclass attribute;self._projection_map = projection_map or self._default_proj_mapaliases the shared class dict when no override is passed. Any in-place mutation ofself._projection_mapon one instance would leak to every other instance (and toMoEAdapter._default_proj_mapitself). Currently nothing mutates it, but it's a latent footgun flagged by Ruff (RUF012).🛡️ Suggested fix
- self._projection_map = projection_map or self._default_proj_map + self._projection_map = dict(projection_map) if projection_map else dict(self._default_proj_map)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/python/modelexpress/weight_transfer/adapters/moe.py` around lines 24 - 27, The shared mutable default mapping in MoEAdapter should not be reused directly across instances. Update the initialization path in the class that defines _default_proj_map so self._projection_map gets its own copy when no projection_map is provided, instead of aliasing the class attribute. Keep the existing symbols (_default_proj_map, __init__, and _projection_map) so the fix is localized and avoids cross-instance mutation.Source: Linters/SAST tools
43-67: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winNo validation that
num_expertsmatches the stacked tensor's leading dimension.
adapt_lazy_weightsblindly iteratesrange(self._num_experts)and builds a__getitem__op for each index, regardless oflazy.shape[0]. Ifnum_expertsis misconfigured relative to the actual stacked tensor (e.g., trainer checkpoint has a different expert count), this silently constructs out-of-range indexing ops that only fail much later when the op chain is replayed against the real tensor (apply_chain/resolve_chain_region), far from the root cause.🛡️ Suggested guard
prefix, proj_key = m.group(1), m.group(2) hf_proj = self._projection_map.get(proj_key) if hf_proj is None: yield hf_name, lazy continue + if lazy.shape[0] != self._num_experts: + raise ValueError( + f"{hf_name}: stacked dim {lazy.shape[0]} != configured num_experts {self._num_experts}" + ) import torch for j in range(self._num_experts):🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/python/modelexpress/weight_transfer/adapters/moe.py` around lines 43 - 67, The adapt_lazy_weights method in moe.py should validate that the configured num_experts matches the leading dimension of the stacked LazyWeight before yielding per-expert entries. Add a guard near the existing stacked-tensor handling in adapt_lazy_weights that compares self._num_experts with lazy.shape[0] (or the equivalent expert dimension) and raises a clear error if they differ. Keep the fix localized to adapt_lazy_weights and use the existing LazyWeight, self._stacked_pattern, and self._projection_map flow so invalid __getitem__ op_chain construction is prevented early.modelexpress_client/python/modelexpress/weight_transfer/protocol/test_weight_transfer.py (1)
59-59: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUnused unpacked variables (
offset,stride) flagged by Ruff.🧹 Suggested fix
- offset, shape, stride = resolve_chain_region( + offset, shape, _stride = resolve_chain_region( (), torch.Size([4, 8]), torch.float32 ) ... - offset, shape, stride = resolve_chain_region( + offset, shape, _stride = resolve_chain_region( chain, torch.Size([8, 16]), torch.bfloat16 ) ... - offset, shape, stride = resolve_chain_region( + _offset, shape, _stride = resolve_chain_region( chain, torch.Size([4, 8]), torch.float32 )Also applies to: 67-67, 76-76
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/python/modelexpress/weight_transfer/protocol/test_weight_transfer.py` at line 59, Ruff is flagging unused unpacked values in the weight transfer tests where resolve_chain_region is called and only shape is used. Update the affected test cases in test_weight_transfer.py to avoid binding offset and stride when they are not needed, either by unpacking only the required value or by assigning the extra values to placeholders, and apply the same cleanup to all resolve_chain_region usages in the referenced tests.Source: Linters/SAST tools
modelexpress_client/python/modelexpress/weight_transfer/engine/__init__.py (1)
7-7: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSort
__all__(Ruff RUF022).♻️ Suggested fix
-__all__ = ["WeightLoaderAdapter", "BakeRecorder", "LazyWeight", "RecordedCopy", "bake_model"] +__all__ = ["BakeRecorder", "LazyWeight", "RecordedCopy", "WeightLoaderAdapter", "bake_model"]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/python/modelexpress/weight_transfer/engine/__init__.py` at line 7, The __all__ export list in the engine package is unsorted and triggers Ruff RUF022. Update the __all__ declaration in the module that defines WeightLoaderAdapter, BakeRecorder, LazyWeight, RecordedCopy, and bake_model so the exported names are ordered consistently in alphabetical order.Source: Linters/SAST tools
modelexpress_client/python/modelexpress/weight_transfer/engine/lazy.py (1)
109-127: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicated
copy_/__copy__recording logic between__torch_dispatch__and__torch_function__.The exact same
RecordedCopyconstruction and guard conditions are repeated in both dispatch paths. Any future change toRecordedCopyfields or guard conditions must be kept in sync manually in two places.♻️ Suggested extraction
+def _maybe_record_copy(lazy: "LazyWeight", dst: Any) -> None: + if ( + _BAKE_STACK + and dst is not None + and isinstance(dst, torch.Tensor) + and not isinstance(dst, LazyWeight) + and dst.device.type != "meta" + ): + _BAKE_STACK[-1].record(RecordedCopy( + src_name=lazy._lazy_name, + op_chain=lazy._lazy_chain, + dst_addr=dst.data_ptr(), + dst_shape=tuple(dst.shape), + dst_stride=tuple(dst.stride()), + dst_dtype=dst.dtype, + dst_device_id=dst.device.index or 0, + ))Then call
_maybe_record_copy(lazy, dst)from bothcopy_branches.Also applies to: 163-181
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/python/modelexpress/weight_transfer/engine/lazy.py` around lines 109 - 127, The `copy_`/`__copy__` recording logic in `lazy.py` is duplicated across both dispatch paths, so extract the shared guard and `RecordedCopy` निर्माण into a helper such as `_maybe_record_copy(lazy, dst)` and call it from both `__torch_dispatch__` and `__torch_function__`; keep the helper responsible for the `_BAKE_STACK` checks, `LazyWeight`/meta-device filtering, and `RecordedCopy` field population so future changes only happen in one place.modelexpress_client/python/modelexpress/envs.py (1)
197-201: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
MX_TRAINER_SYNC_TIMEOUTis overloaded for two different timeout semantics.Per the downstream consumer, this single value drives both the
TrainerTablediscovery wait loop and thePullRole's RDMAsync_timeout: "sync_timeout=float(envs.MX_TRAINER_SYNC_TIMEOUT)". Tuning one use case (e.g., shortening table-discovery wait) inadvertently changes the RDMA transfer timeout too, which could cause legitimate large transfers to time out.Consider splitting into
MX_TRAINER_TABLE_TIMEOUTandMX_TRAINER_SYNC_TIMEOUTfor independent tuning.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/python/modelexpress/envs.py` around lines 197 - 201, `MX_TRAINER_SYNC_TIMEOUT` in the env mapping is being used for two different timeout behaviors, so split the shared setting into separate env vars for discovery and sync. Update `envs.py` to introduce distinct keys such as `MX_TRAINER_TABLE_TIMEOUT` for the trainer table wait loop and keep `MX_TRAINER_SYNC_TIMEOUT` for the RDMA transfer path, and adjust the downstream consumers in `TrainerTable` and `PullRole` so each reads the appropriate value instead of sharing one timeout.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In
`@modelexpress_client/python/modelexpress/load_strategy/trainer_pull_strategy.py`:
- Around line 226-231: The post-pull weight processing in post_pull_hook is
currently swallowed by a warning, so sync_and_post_process can succeed even when
process_weights_after_loading fails. Update post_pull_hook to surface that
failure to the caller by re-raising the exception or otherwise returning an
error signal, and ensure PullRole.sync_and_post_process handles that failure
path instead of treating the sync as successful; keep the logging in place, but
do not let the exception be silently ignored.
- Around line 190-200: The Redis fallback in the load path can block
indefinitely because `trainer_pull_strategy.py`’s Redis client creation and
`get` call do not enforce any timeout. Update the `_fetch_table` Redis branch to
pass explicit `socket_connect_timeout` and `socket_timeout` values when calling
`redis_lib.from_url`, and keep the failure path returning `None` so the
`MX_TRAINER_SYNC_TIMEOUT` retry loop is still effective. Use the existing
`redis_url`, `ctx.global_rank`, and `logger.debug` flow to locate and adjust the
Redis fallback logic.
- Around line 88-92: Preserve the original exception chain in the trainer-table
fetch failure path: in `TrainerPullStrategy`’s table-loading logic where
`_fetch_table(ctx)` is wrapped by `try/except Exception as e`, re-raise
`StrategyFailed` with `from e` so the original traceback is retained. Keep the
existing log message, and ensure the `StrategyFailed` raised from this block
carries the original exception context for easier debugging.
In `@modelexpress_client/python/modelexpress/weight_transfer/engine/lazy.py`:
- Around line 129-144: The exception handler in lazy weight transfer is
swallowing failures from func(*meta_args, **kwargs) and returning a plain meta
tensor with the original lazy.shape, which breaks op-chain tracking and can
silently drop weights. Update the fallback in LazyWeight.__torch_function__ /
the wrapped op path to avoid fabricating a torch.Tensor; instead re-raise the
exception or log the failure and re-raise, consistent with the safer handling
already used in __torch_function__ below. Keep the returned value as a
LazyWeight only on the successful path so subsequent ops and copy_() still
recognize the tracked weight.
- Around line 94-99: The docstring in the lazy tensor fallback is misleading
because it frames __torch_dispatch__ as a PyTorch 2.9-only requirement, which
conflicts with the package’s torch>=2.6.0 support. Update the documentation on
the __torch_dispatch__ / _make_wrapper_subclass fallback in lazy.py to describe
the behavior and why the ATen-level path exists, rather than tying it to a
specific torch version gate.
In `@modelexpress_client/python/modelexpress/weight_transfer/planner/resolver.py`:
- Around line 43-76: `region_elem_runs` is generating one run per element in
`_collect`, even when the innermost dimension is contiguous, which explodes
tuple creation for large tensors. Update the `_collect` logic in
`region_elem_runs` to fast-path the last dimension when `stride[dim] == 1` by
appending a single `(base, shape[dim])` run instead of per-element `(offset, 1)`
entries, while preserving the existing recursive behavior for non-contiguous
cases and the final merge pass.
In `@modelexpress_client/python/modelexpress/weight_transfer/planner/router.py`:
- Around line 1-234: This module duplicates the routing implementation from
protocol/router.py, so replace the local copy with a thin re-export of
route_regions from that module. Update planner.router to preserve the same
public API used by planner/local.py and tests, while removing the duplicated
helper implementations like _split_run_2d and _zip_src_dst. Keep the module
importable without changing call sites.
In `@modelexpress_client/python/modelexpress/weight_transfer/planner/server.py`:
- Around line 44-63: The server build path in ServerPlanner.build is using
payload field names that do not match the WeightSyncService proto, so the
request/response won’t interoperate correctly. Update the server-side request
construction in _build_via_server to use the BuildPlanRequest fields model_key
and regions, and read descriptors from the GetPlanResponse instead of
regions_payload/descriptors_payload. Keep the fallback behavior intact, but make
the server path align with WeightSyncService so it can return descriptors
successfully.
In
`@modelexpress_client/python/modelexpress/weight_transfer/protocol/__init__.py`:
- Around line 4-45: The package root for modelexpress.weight_transfer.protocol
is missing the M2N protocol exports, so importers cannot access the full API
from this module. Update the protocol package initializer to re-export
M2nDescriptor and the encode_m2n_descriptors / decode_m2n_descriptors functions
alongside the existing OpSpec, types, and serialization helpers, and add them to
__all__ so package-root imports expose the complete protocol surface.
In `@modelexpress_client/python/modelexpress/weight_transfer/roles/pull.py`:
- Around line 86-98: The PullRole refresh path is leaking NIXL remote-agent
registrations because initialize() rebuilds remote_agents and a new NixlExecutor
without releasing the previously registered agents, and teardown() never cleans
them up. Update PullRole.initialize(), refresh(), and/or teardown() to
explicitly unregister or reuse the existing table.agents entries via
_nixl_manager before creating a new NixlExecutor, so stale remote-agent state
does not accumulate across reshard cycles.
In `@modelexpress_client/python/modelexpress/weight_transfer/roles/push.py`:
- Around line 106-125: The missing-parameter branch in push.py is under-logged
compared with the size-mismatch branch, so update the shard handling in the push
flow to treat both as a meaningful sync problem. In the loop over
inference_table.shards inside the push logic, raise the log level for the
local_params miss from debug to warning (or equivalent) and add an aggregate
validation after the loop in the same method/class to detect when not all shards
were mapped. Use the existing identifiers inference_table.shards, local_params,
and the push method on the trainer role to locate and adjust the logging and
validation consistently.
- Around line 50-92: The push plan built in `PushRole.initialize()` caches RDMA
descriptors using parameter pointers that can change under FSDP2/`fully_shard`,
so `PushRole.sync()` may use stale addresses. Update `PushRole.sync()` to
rebuild the descriptors or refresh each descriptor’s `src_addr` from the current
model parameters before calling `NixlExecutor.execute()`, using the existing
`PushRole`, `_build_push_plan()`, and `_executor` flow as the main entry points.
In
`@modelexpress_client/python/modelexpress/weight_transfer/transport/m2n_executor.py`:
- Around line 44-130: `M2nExecutor.execute` duplicates nearly all of
`NixlExecutor.execute`, so shared transfer setup, polling, timeout, release, and
synchronize behavior can drift. Extract the common execute workflow into a
shared helper or base class in the transport layer, and have `M2nExecutor` and
`NixlExecutor` delegate to it while only supplying the path-specific pieces such
as grouping key, transfer operation, and descriptor-to-dlist mapping.
- Around line 93-104: The transfer handle can leak if `agent.transfer(handle)`
fails after `agent.make_prepped_xfer(...)` has already allocated it. Update
`m2n_executor` so the handle is registered for cleanup before calling
`agent.transfer(handle)`, or wrap the transfer in a `try`/`finally` that
guarantees release on failure; keep the fix localized around the `handles` list
and the `transfer(handle)` call site.
In `@modelexpress_client/python/tests/gpu/test_e2e_primerl_real_model.py`:
- Line 44: Remove the hardcoded developer checkout path from the module import
setup in test_e2e_primerl_real_model.py; the unconditional sys.path.insert(0,
...) should not alter import resolution for everyone. Update the test to rely on
the installed modelexpress package or a test-local path setup/fixture if needed,
and ensure any path manipulation happens only in the test environment and does
not take precedence over the real package.
- Around line 38-42: Gate the throughput assertions in _cuda_memcpy-based test
flow on peer-access availability, since the current cudaMemcpyDefault path
between GPU 0 and 1 can fall back to host staging without peer access enabled.
Update the test in test_e2e_primerl_real_model.py to explicitly check whether
peer access between the devices is supported and only run the >5 GB/s throughput
checks when that capability is present; otherwise skip or relax the benchmark
assertions while keeping _cuda_memcpy unchanged.
In `@modelexpress_server/src/weight_sync/router.rs`:
- Around line 125-127: The unpack_runs helper currently assumes every chunk has
two elements, so an odd-length payload can panic when indexing c[1]. Update
unpack_runs in router.rs to validate the input length before unpacking and
handle malformed data safely, returning an error or rejecting the request
instead of crashing; if this function is used by the weight sync request flow,
make sure the caller around unpack_runs propagates a Status for invalid client
payloads.
- Around line 16-25: `TrainerShard` and `split_run` in the Rust router only
handle row ranges, so column-sharded tensors are matched and addressed
incorrectly. Extend `TrainerShard` to carry the same column range metadata as
the Python router (`col_start`/`col_end`), then update `split_run` to match
shards using both row and column bounds when computing the target `agent_index`
and `src_addr`; if column-sharded layouts are not supported, add an explicit
validation/error path instead of falling back to the first row match.
In `@modelexpress_server/src/weight_sync/service.rs`:
- Around line 108-122: The invalidation in publish_trainer_table only removes
entries from plan_key_to_id, leaving orphaned CachedPlan values behind in
state.plans. Update publish_trainer_table to collect the plan IDs removed by the
retain call on plan_key_to_id, then use those IDs to delete the matching
CachedPlan entries from plans as part of the same invalidation path. Keep the
fix localized to publish_trainer_table and the state fields plan_key_to_id and
plans so both caches stay in sync.
- Around line 160-199: The build_plan flow in WeightSyncService has a
check-then-insert race between the initial state.read() lookup of plan_key_to_id
and the later state.write() insert, which can create duplicate plans for the
same req.plan_key. Fix it by rechecking the cache under the write lock or by
moving the lookup and insert into a single locked critical section in
build_plan, so only one plan_id is created and stored for a given plan_key. Use
the existing plan_key_to_id and plans fields, along with build_plan and
invalidate_plan, to keep the cached plan mapping consistent and idempotent.
---
Outside diff comments:
In `@modelexpress_client/python/modelexpress/load_strategy/__init__.py`:
- Around line 64-76: TrainerPullStrategy is being selected too broadly because
load_strategy.__init__ always includes it when is_available() passes, causing
non-trainer-pull deployments to wait in TrainerPullStrategy.load() for
MX_TRAINER_SYNC_TIMEOUT. Update the strategy selection in the load strategy
initialization so TrainerPullStrategy is only added when an explicit
trainer-pull signal is present, such as MX_TRAINER_TABLE_KEY or
MX_WEIGHT_SYNC_SERVER. Keep RdmaStrategy and the other fallbacks unchanged so
ordinary RDMA startup skips TrainerPullStrategy immediately.
---
Duplicate comments:
In
`@modelexpress_client/python/modelexpress/weight_transfer/transport/nixl_executor.py`:
- Around line 94-104: The transfer path in NixlExecutor has the same handle-leak
risk as M2nExecutor.execute: if agent.transfer(handle) fails after
agent.make_prepped_xfer() succeeds, the newly created handle is never recorded
for cleanup. Update NixlExecutor.execute so the handle is protected immediately
after creation, either by appending it to handles before calling transfer or by
using a try/finally around agent.transfer(handle) that guarantees release
tracking even on failure. Use the make_prepped_xfer, transfer, and handles logic
in NixlExecutor.execute to locate the fix.
---
Nitpick comments:
In `@modelexpress_client/python/modelexpress/envs.py`:
- Around line 197-201: `MX_TRAINER_SYNC_TIMEOUT` in the env mapping is being
used for two different timeout behaviors, so split the shared setting into
separate env vars for discovery and sync. Update `envs.py` to introduce distinct
keys such as `MX_TRAINER_TABLE_TIMEOUT` for the trainer table wait loop and keep
`MX_TRAINER_SYNC_TIMEOUT` for the RDMA transfer path, and adjust the downstream
consumers in `TrainerTable` and `PullRole` so each reads the appropriate value
instead of sharing one timeout.
In `@modelexpress_client/python/modelexpress/weight_transfer/__init__.py`:
- Around line 65-91: The __all__ list in weight_transfer/__init__.py is not
sorted and triggers Ruff RUF022. Reorder the exported symbols in __all__ into a
consistent sorted order, keeping the existing grouped sections if needed, and
update the list near the module-level __all__ definition so the import export
order is alphabetical and lint-compliant.
In
`@modelexpress_client/python/modelexpress/weight_transfer/adapters/__init__.py`:
- Line 7: The __all__ export list in the adapters package is not sorted and
triggers Ruff RUF022; update the __all__ definition in the adapters __init__
module so the exported names are ordered consistently, keeping the existing
symbols ModelAdapter and MoEAdapter but arranging them in sorted order.
In `@modelexpress_client/python/modelexpress/weight_transfer/adapters/moe.py`:
- Around line 24-27: The shared mutable default mapping in MoEAdapter should not
be reused directly across instances. Update the initialization path in the class
that defines _default_proj_map so self._projection_map gets its own copy when no
projection_map is provided, instead of aliasing the class attribute. Keep the
existing symbols (_default_proj_map, __init__, and _projection_map) so the fix
is localized and avoids cross-instance mutation.
- Around line 43-67: The adapt_lazy_weights method in moe.py should validate
that the configured num_experts matches the leading dimension of the stacked
LazyWeight before yielding per-expert entries. Add a guard near the existing
stacked-tensor handling in adapt_lazy_weights that compares self._num_experts
with lazy.shape[0] (or the equivalent expert dimension) and raises a clear error
if they differ. Keep the fix localized to adapt_lazy_weights and use the
existing LazyWeight, self._stacked_pattern, and self._projection_map flow so
invalid __getitem__ op_chain construction is prevented early.
In `@modelexpress_client/python/modelexpress/weight_transfer/engine/__init__.py`:
- Line 7: The __all__ export list in the engine package is unsorted and triggers
Ruff RUF022. Update the __all__ declaration in the module that defines
WeightLoaderAdapter, BakeRecorder, LazyWeight, RecordedCopy, and bake_model so
the exported names are ordered consistently in alphabetical order.
In `@modelexpress_client/python/modelexpress/weight_transfer/engine/lazy.py`:
- Around line 109-127: The `copy_`/`__copy__` recording logic in `lazy.py` is
duplicated across both dispatch paths, so extract the shared guard and
`RecordedCopy` निर्माण into a helper such as `_maybe_record_copy(lazy, dst)` and
call it from both `__torch_dispatch__` and `__torch_function__`; keep the helper
responsible for the `_BAKE_STACK` checks, `LazyWeight`/meta-device filtering,
and `RecordedCopy` field population so future changes only happen in one place.
In `@modelexpress_client/python/modelexpress/weight_transfer/protocol/ops.py`:
- Around line 11-27: `OpSpec` is marked frozen but still contains an unhashable
`dict` in `kwargs`, which breaks the expected hashability of the dataclass.
Update `OpSpec` in `ops.py` so its state is normalized into immutable, hashable
forms in `__post_init__` and/or provide a custom `__hash__` that hashes a stable
representation of `name`, `args`, and `kwargs`; keep `OpChain` compatibility in
mind if it relies on hashing these specs.
In `@modelexpress_client/python/modelexpress/weight_transfer/protocol/router.py`:
- Around line 149-193: Add a defensive parity check in _zip_src_dst so source
and destination element totals must match before building RdmaDescriptor
entries. Compute/validate the total counts for src_triples and dst_runs, and
raise an explicit error if they differ instead of allowing the while loop to
stop early and drop leftover elements. Keep the fix localized to _zip_src_dst in
router.py so upstream bugs in resolve_copies fail fast with a clear message.
In
`@modelexpress_client/python/modelexpress/weight_transfer/protocol/test_weight_transfer.py`:
- Line 59: Ruff is flagging unused unpacked values in the weight transfer tests
where resolve_chain_region is called and only shape is used. Update the affected
test cases in test_weight_transfer.py to avoid binding offset and stride when
they are not needed, either by unpacking only the required value or by assigning
the extra values to placeholders, and apply the same cleanup to all
resolve_chain_region usages in the referenced tests.
In `@modelexpress_client/python/pyproject.toml`:
- Around line 68-73: The gpu marker is declared in the pytest configuration but
never applied to the GPU E2E test, so CI cannot select or exclude it by marker.
Update tests/gpu/test_e2e_primerl_real_model.py to add `@pytest.mark.gpu` to the
relevant test/class/module alongside the existing skipif gate, using the gpu
marker name defined in pyproject.toml so pytest -m selection works as intended.
In `@modelexpress_client/python/tests/gpu/test_e2e_primerl_real_model.py`:
- Around line 22-26: The hardcoded first entry in the _LIBCUDART_PATHS list is
too environment-specific and should be removed or replaced with a dynamic
lookup. Update the test setup around _LIBCUDART_PATHS in
test_e2e_primerl_real_model.py to prefer a derived libcudart location using
something like ctypes.util.find_library or values inferred from
torch.version.cuda, while keeping the existing fallback and graceful-skip
behavior if no library is found.
- Around line 188-193: The blind exception handling around resolve_chain_region
in the parameter copy loop silently skips parameters without any diagnostics.
Update the try/except in test_e2e_primerl_real_model to catch the failure in a
more explicit way and emit a warning or debug log that includes the
parameter/op_chain context before continuing, so unresolved regions are still
visible when the transfer plan is built.
In `@modelexpress_client/python/tests/test_e2e_primerl.py`:
- Around line 193-199: The current test is tautological because it only
increments a MagicMock and never exercises PullRole.sync_and_post_process().
Update test_post_pull_hook_called_per_sync to construct a minimal PullRole with
mocked planner, executor, and adapter, then call sync_and_post_process() and
assert adapter.post_pull_hook.call_count so the test verifies the real post-pull
wiring after sync rather than mock self-counting.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: db68ddd6-8d69-4069-95f0-c681e5303985
📒 Files selected for processing (47)
docs/ARCHITECTURE.mdmodelexpress_client/python/modelexpress/envs.pymodelexpress_client/python/modelexpress/load_strategy/__init__.pymodelexpress_client/python/modelexpress/load_strategy/trainer_pull_strategy.pymodelexpress_client/python/modelexpress/weight_transfer/__init__.pymodelexpress_client/python/modelexpress/weight_transfer/adapters/__init__.pymodelexpress_client/python/modelexpress/weight_transfer/adapters/base.pymodelexpress_client/python/modelexpress/weight_transfer/adapters/moe.pymodelexpress_client/python/modelexpress/weight_transfer/engine/__init__.pymodelexpress_client/python/modelexpress/weight_transfer/engine/base.pymodelexpress_client/python/modelexpress/weight_transfer/engine/lazy.pymodelexpress_client/python/modelexpress/weight_transfer/planner/__init__.pymodelexpress_client/python/modelexpress/weight_transfer/planner/base.pymodelexpress_client/python/modelexpress/weight_transfer/planner/local.pymodelexpress_client/python/modelexpress/weight_transfer/planner/m2n_planner.pymodelexpress_client/python/modelexpress/weight_transfer/planner/resolver.pymodelexpress_client/python/modelexpress/weight_transfer/planner/router.pymodelexpress_client/python/modelexpress/weight_transfer/planner/server.pymodelexpress_client/python/modelexpress/weight_transfer/protocol/__init__.pymodelexpress_client/python/modelexpress/weight_transfer/protocol/ops.pymodelexpress_client/python/modelexpress/weight_transfer/protocol/router.pymodelexpress_client/python/modelexpress/weight_transfer/protocol/serialization.pymodelexpress_client/python/modelexpress/weight_transfer/protocol/test_weight_transfer.pymodelexpress_client/python/modelexpress/weight_transfer/protocol/types.pymodelexpress_client/python/modelexpress/weight_transfer/roles/__init__.pymodelexpress_client/python/modelexpress/weight_transfer/roles/base.pymodelexpress_client/python/modelexpress/weight_transfer/roles/pull.pymodelexpress_client/python/modelexpress/weight_transfer/roles/push.pymodelexpress_client/python/modelexpress/weight_transfer/transport/__init__.pymodelexpress_client/python/modelexpress/weight_transfer/transport/m2n_executor.pymodelexpress_client/python/modelexpress/weight_transfer/transport/nixl_executor.pymodelexpress_client/python/pyproject.tomlmodelexpress_client/python/tests/gpu/conftest.pymodelexpress_client/python/tests/gpu/test_e2e_primerl_real_model.pymodelexpress_client/python/tests/test_e2e_nccl_m2n.pymodelexpress_client/python/tests/test_e2e_primerl.pymodelexpress_client/python/tests/test_m2n.pymodelexpress_client/python/tests/test_rl_weight_sync.pymodelexpress_client/python/tests/test_weight_transfer.pymodelexpress_common/build.rsmodelexpress_common/proto/weight_sync.protomodelexpress_common/src/lib.rsmodelexpress_server/src/lib.rsmodelexpress_server/src/server.rsmodelexpress_server/src/weight_sync.rsmodelexpress_server/src/weight_sync/router.rsmodelexpress_server/src/weight_sync/service.rs
| try: | ||
| table = self._fetch_table(ctx) | ||
| except Exception as e: | ||
| logger.info("[Worker %d] TrainerTable not available: %s", ctx.global_rank, e) | ||
| raise StrategyFailed(f"TrainerTable not available: {e}", mutated=False) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Preserve exception chain when re-raising as StrategyFailed.
raise StrategyFailed(...) inside the except Exception as e: block drops the original traceback. Use from e for easier debugging of trainer-table fetch failures.
🔧 Proposed fix
except Exception as e:
logger.info("[Worker %d] TrainerTable not available: %s", ctx.global_rank, e)
- raise StrategyFailed(f"TrainerTable not available: {e}", mutated=False)
+ raise StrategyFailed(f"TrainerTable not available: {e}", mutated=False) from e📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| try: | |
| table = self._fetch_table(ctx) | |
| except Exception as e: | |
| logger.info("[Worker %d] TrainerTable not available: %s", ctx.global_rank, e) | |
| raise StrategyFailed(f"TrainerTable not available: {e}", mutated=False) | |
| try: | |
| table = self._fetch_table(ctx) | |
| except Exception as e: | |
| logger.info("[Worker %d] TrainerTable not available: %s", ctx.global_rank, e) | |
| raise StrategyFailed(f"TrainerTable not available: {e}", mutated=False) from e |
🧰 Tools
🪛 Ruff (0.15.20)
[warning] 90-90: Do not catch blind exception: Exception
(BLE001)
[warning] 92-92: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In
`@modelexpress_client/python/modelexpress/load_strategy/trainer_pull_strategy.py`
around lines 88 - 92, Preserve the original exception chain in the trainer-table
fetch failure path: in `TrainerPullStrategy`’s table-loading logic where
`_fetch_table(ctx)` is wrapped by `try/except Exception as e`, re-raise
`StrategyFailed` with `from e` so the original traceback is retained. Keep the
existing log message, and ensure the `StrategyFailed` raised from this block
carries the original exception context for easier debugging.
Source: Linters/SAST tools
| redis_url = envs.MX_REDIS_URL | ||
| try: | ||
| import redis as redis_lib | ||
| r = redis_lib.from_url(redis_url) | ||
| return r.get(key) | ||
| except ImportError: | ||
| logger.debug("[Worker %d] redis-py not installed", ctx.global_rank) | ||
| except Exception as e: | ||
| logger.debug("[Worker %d] Redis GET failed: %s", ctx.global_rank, e) | ||
|
|
||
| return None |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Redis fallback call has no timeout — can hang the load path indefinitely.
redis_lib.from_url(redis_url) and .get(key) are called with no socket_timeout/socket_connect_timeout. If Redis is unreachable but not actively refusing connections, this call can block forever, silently defeating the MX_TRAINER_SYNC_TIMEOUT-bounded retry loop in _fetch_table.
🔧 Proposed fix
try:
import redis as redis_lib
- r = redis_lib.from_url(redis_url)
+ r = redis_lib.from_url(
+ redis_url,
+ socket_connect_timeout=5.0,
+ socket_timeout=5.0,
+ )
return r.get(key)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| redis_url = envs.MX_REDIS_URL | |
| try: | |
| import redis as redis_lib | |
| r = redis_lib.from_url(redis_url) | |
| return r.get(key) | |
| except ImportError: | |
| logger.debug("[Worker %d] redis-py not installed", ctx.global_rank) | |
| except Exception as e: | |
| logger.debug("[Worker %d] Redis GET failed: %s", ctx.global_rank, e) | |
| return None | |
| redis_url = envs.MX_REDIS_URL | |
| try: | |
| import redis as redis_lib | |
| r = redis_lib.from_url( | |
| redis_url, | |
| socket_connect_timeout=5.0, | |
| socket_timeout=5.0, | |
| ) | |
| return r.get(key) | |
| except ImportError: | |
| logger.debug("[Worker %d] redis-py not installed", ctx.global_rank) | |
| except Exception as e: | |
| logger.debug("[Worker %d] Redis GET failed: %s", ctx.global_rank, e) | |
| return None |
🧰 Tools
🪛 Ruff (0.15.20)
[warning] 197-197: Do not catch blind exception: Exception
(BLE001)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In
`@modelexpress_client/python/modelexpress/load_strategy/trainer_pull_strategy.py`
around lines 190 - 200, The Redis fallback in the load path can block
indefinitely because `trainer_pull_strategy.py`’s Redis client creation and
`get` call do not enforce any timeout. Update the `_fetch_table` Redis branch to
pass explicit `socket_connect_timeout` and `socket_timeout` values when calling
`redis_lib.from_url`, and keep the failure path returning `None` so the
`MX_TRAINER_SYNC_TIMEOUT` retry loop is still effective. Use the existing
`redis_url`, `ctx.global_rank`, and `logger.debug` flow to locate and adjust the
Redis fallback logic.
| def post_pull_hook(self, model: Any) -> None: | ||
| if self._ctx.adapter is not None and hasattr(self._ctx.adapter, "process_weights_after_loading"): | ||
| try: | ||
| self._ctx.adapter.process_weights_after_loading(model) | ||
| except Exception as e: | ||
| logger.warning("post_pull_hook failed: %s", e) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Silent failure of post-pull weight post-processing.
process_weights_after_loading failures (e.g. FP8 repack) are caught and only logged as a warning; sync_and_post_process in PullRole will return as if sync succeeded even though weights may not have been correctly post-processed, risking silent incorrect inference results.
🔧 Proposed fix
def post_pull_hook(self, model: Any) -> None:
if self._ctx.adapter is not None and hasattr(self._ctx.adapter, "process_weights_after_loading"):
- try:
- self._ctx.adapter.process_weights_after_loading(model)
- except Exception as e:
- logger.warning("post_pull_hook failed: %s", e)
+ self._ctx.adapter.process_weights_after_loading(model)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def post_pull_hook(self, model: Any) -> None: | |
| if self._ctx.adapter is not None and hasattr(self._ctx.adapter, "process_weights_after_loading"): | |
| try: | |
| self._ctx.adapter.process_weights_after_loading(model) | |
| except Exception as e: | |
| logger.warning("post_pull_hook failed: %s", e) | |
| def post_pull_hook(self, model: Any) -> None: | |
| if self._ctx.adapter is not None and hasattr(self._ctx.adapter, "process_weights_after_loading"): | |
| self._ctx.adapter.process_weights_after_loading(model) |
🧰 Tools
🪛 Ruff (0.15.20)
[warning] 230-230: Do not catch blind exception: Exception
(BLE001)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In
`@modelexpress_client/python/modelexpress/load_strategy/trainer_pull_strategy.py`
around lines 226 - 231, The post-pull weight processing in post_pull_hook is
currently swallowed by a warning, so sync_and_post_process can succeed even when
process_weights_after_loading fails. Update post_pull_hook to surface that
failure to the caller by re-raising the exception or otherwise returning an
error signal, and ensure PullRole.sync_and_post_process handles that failure
path instead of treating the sync as successful; keep the logging in place, but
do not let the exception be silently ignored.
| """ATen-level fallback required by PyTorch ≥2.9 for _make_wrapper_subclass. | ||
|
|
||
| Primary interception lives in __torch_function__ (Python-level calls | ||
| like narrow/view/copy_). copy_ is also handled here so both paths | ||
| work; all other ops proxy through meta tensors for correct output shape. | ||
| """ |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
What PyTorch version introduced a requirement for __torch_dispatch__ alongside __torch_function__ when using _make_wrapper_subclass?
💡 Result:
There is no single PyTorch version that introduced a formal, mandatory requirement to use torch_dispatch alongside torch_function when using _make_wrapper_subclass. Rather, the use of torch_dispatch with _make_wrapper_subclass is a functional necessity for creating custom tensor subclasses that are intended to be performant, functional, or compatible with modern PyTorch features like torch.compile and functorch [1][2]. _make_wrapper_subclass, introduced around PyTorch 1.10 (PR #65340) [3], creates a tensor without immediate backing storage to support custom, "wrapper" subclasses that manage their own data or behavior [4][3]. Because these tensors lack standard storage, they are generally incompatible with standard dispatch paths. To ensure these subclasses function correctly across the PyTorch ecosystem—particularly for tracing and graph-based operations—implementers must provide a torch_dispatch override to define how operations are handled, as standard torch_function is often insufficient for deep integration into the PyTorch dispatcher [1][2]. In modern PyTorch (notably for compatibility with PT2/torch.compile), subclasses created with _make_wrapper_subclass are expected to implement not only torch_dispatch but also tensor_flatten and tensor_unflatten to be considered traceable [2]. While torch_function can still be used for high-level Python overrides, it is increasingly treated as a legacy or auxiliary mechanism compared to torch_dispatch for low-level subclass behavior [2][5][6]. Consequently, many developers explicitly disable torch_function in these subclasses (using torch._C._disabled_torch_function_impl) to force exclusive reliance on torch_dispatch [7][8].
Citations:
- 1: Subclasses with
__torch_function__()impls for composite ops don't work with dynamic shapes pytorch/pytorch#120654 - 2: https://github.com/pytorch/pytorch/blob/main/torch/utils/_python_dispatch.py
- 3: Add Tensor._make_wrapper_subclass pytorch/pytorch#65340
- 4: https://dev-discuss.pytorch.org/t/whats-the-difference-between-torch-tensor-make-subclass-and-torch-tensor-make-wrapper-subclass/1839
- 5: https://docs.pytorch.org/docs/stable/notes/extending.html
- 6: __torch_function__ mode pytorch/pytorch#75154
- 7:
_make_wrapper_subclass+__torch_dispatch__fails undertorch.compile+torch.vmapwith cross-device storage error pytorch/pytorch#176957 - 8: custom_op + TensorSubclass + torch.compile interaction : Better error message or doc pytorch/pytorch#175408
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== torch version references =="
rg -n --hidden --glob '!**/.git/**' 'torch[^0-9]*2\.[0-9]+\.[0-9]+' .
echo
echo "== candidate dependency files =="
git ls-files | rg '(^|/)(pyproject\.toml|requirements(\.txt)?|setup\.py|setup\.cfg|Pipfile|poetry\.lock|uv\.lock|conda\.ya?ml|environment\.ya?ml)$'
echo
echo "== target file excerpt =="
sed -n '70,120p' modelexpress_client/python/modelexpress/weight_transfer/engine/lazy.pyRepository: ai-dynamo/modelexpress
Length of output: 17131
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== repo root files mentioning torch version pin or constraint =="
rg -n --hidden --glob '!**/.git/**' '\btorch\b|PyTorch|pytorch' \
modelexpress_client/python \
pyproject.toml requirements.txt requirements*.txt setup.py setup.cfg Pipfile poetry.lock uv.lock \
2>/dev/null || trueRepository: ai-dynamo/modelexpress
Length of output: 50378
Clarify the torch version note
__torch_dispatch__ isn’t a PyTorch 2.9-only requirement for _make_wrapper_subclass; this package already declares torch>=2.6.0, so the docstring should describe the fallback behavior instead of a version gate.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@modelexpress_client/python/modelexpress/weight_transfer/engine/lazy.py`
around lines 94 - 99, The docstring in the lazy tensor fallback is misleading
because it frames __torch_dispatch__ as a PyTorch 2.9-only requirement, which
conflicts with the package’s torch>=2.6.0 support. Update the documentation on
the __torch_dispatch__ / _make_wrapper_subclass fallback in lazy.py to describe
the behavior and why the ATen-level path exists, rather than tying it to a
specific torch version gate.
| try: | ||
| meta_args = tuple( | ||
| torch.empty(a.shape, dtype=a.dtype, device="meta") | ||
| if isinstance(a, LazyWeight) else a | ||
| for a in args | ||
| ) | ||
| out = func(*meta_args, **kwargs) | ||
| if isinstance(out, torch.Tensor): | ||
| new_op = OpSpec(name=func_name, args=args[1:], kwargs=kwargs) | ||
| return LazyWeight( | ||
| lazy._lazy_name, out.shape, out.dtype, | ||
| lazy._lazy_chain + (new_op,), | ||
| ) | ||
| return out | ||
| except Exception: | ||
| return torch.empty(lazy.shape, dtype=lazy.dtype, device="meta") |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Broad exception swallow returns a wrong-shaped, untracked tensor instead of failing loudly.
When func(*meta_args, **kwargs) raises, the fallback returns torch.empty(lazy.shape, ...) — the tensor's original (pre-op) shape, not the actual result shape, and it is a plain torch.Tensor, not a LazyWeight. This silently breaks op-chain tracking: any subsequent op/copy_() on this value will no longer be recognized as a LazyWeight, so the corresponding weight is silently dropped from the bake pass instead of raising a visible error. Given this drives which RDMA regions get synchronized to inference workers, a silently dropped/mis-shaped weight is a serious, hard-to-diagnose correctness bug.
Consider re-raising (or logging + re-raising) instead of fabricating a fallback tensor, mirroring the safer except Exception: return super().__torch_function__(...) behavior used a few lines below in __torch_function__.
🧰 Tools
🪛 Ruff (0.15.20)
[warning] 140-140: Consider (*lazy._lazy_chain, new_op) instead of concatenation
Replace with (*lazy._lazy_chain, new_op)
(RUF005)
[warning] 143-143: Do not catch blind exception: Exception
(BLE001)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@modelexpress_client/python/modelexpress/weight_transfer/engine/lazy.py`
around lines 129 - 144, The exception handler in lazy weight transfer is
swallowing failures from func(*meta_args, **kwargs) and returning a plain meta
tensor with the original lazy.shape, which breaks op-chain tracking and can
silently drop weights. Update the fallback in LazyWeight.__torch_function__ /
the wrapped op path to avoid fabricating a torch.Tensor; instead re-raise the
exception or log the failure and re-raise, consistent with the safer handling
already used in __torch_function__ below. Keep the returned value as a
LazyWeight only on the successful path so subsequent ops and copy_() still
recognize the tracked weight.
| raise RuntimeError(f"cudaMemcpy failed rc={rc}") | ||
|
|
||
|
|
||
| sys.path.insert(0, os.path.expanduser("~/workspace/modelexpress/modelexpress_client/python")) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Remove developer-specific hardcoded path from sys.path.
sys.path.insert(0, os.path.expanduser("~/workspace/modelexpress/modelexpress_client/python")) hardcodes one author's local checkout location. It runs unconditionally at module import (even in CI runs where GPUs — and thus this test module's skip condition — mean it's still collected), and inserting at index 0 means it takes priority over any properly installed modelexpress package, risking tests silently running against a stale/different copy of the code on any machine where that path happens to exist.
🐛 Suggested fix
-sys.path.insert(0, os.path.expanduser("~/workspace/modelexpress/modelexpress_client/python"))
-
from modelexpress.weight_transfer.engine.lazy import BakeRecorder, LazyWeight📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| sys.path.insert(0, os.path.expanduser("~/workspace/modelexpress/modelexpress_client/python")) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@modelexpress_client/python/tests/gpu/test_e2e_primerl_real_model.py` at line
44, Remove the hardcoded developer checkout path from the module import setup in
test_e2e_primerl_real_model.py; the unconditional sys.path.insert(0, ...) should
not alter import resolution for everyone. Update the test to rely on the
installed modelexpress package or a test-local path setup/fixture if needed, and
ensure any path manipulation happens only in the test environment and does not
take precedence over the real package.
| /// Trainer shard descriptor, mirroring `protocol.types.TrainerShard`. | ||
| #[derive(Debug, Clone, Deserialize)] | ||
| pub struct TrainerShard { | ||
| pub agent_index: u32, | ||
| pub row_start: i64, | ||
| pub row_end: i64, | ||
| pub device_addr: u64, | ||
| pub row_bytes: i64, | ||
| pub device_id: i32, | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🔴 Critical | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm the Python router models column sharding that the Rust mirror lacks.
fd -t f 'router.py' modelexpress_client | while read -r f; do
echo "=== $f ==="
rg -nP -C2 'col_start|col_end|elems_per_row|row_start|def route_regions' "$f"
done
# Confirm TrainerShard field set on the Python side.
rg -nP -C1 'class TrainerShard|col_start|col_end' modelexpress_client/python/modelexpress/weight_transfer/protocol/types.pyRepository: ai-dynamo/modelexpress
Length of output: 7404
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the Rust router and related types/routing logic.
sed -n '1,260p' modelexpress_server/src/weight_sync/router.rs | cat -n
printf '\n=== related Rust protocol/type definitions ===\n'
rg -n -C 3 'TrainerShard|TrainerTensor|split_run|route_regions|row_start|row_end|device_addr|device_id' modelexpress_server/src -g '*.rs'Repository: ai-dynamo/modelexpress
Length of output: 33056
Support column-sharded tensors in the Rust router modelexpress_server/src/weight_sync/router.rs:16-25,130-162 — TrainerShard only tracks row_start/row_end, and split_run matches shards by row only. For 2-D/column-sharded tensors, that picks the first row-matching shard and computes the wrong agent_index/src_addr. Mirror the Python router’s col_start/col_end handling here, or reject non-row-only layouts explicitly.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@modelexpress_server/src/weight_sync/router.rs` around lines 16 - 25,
`TrainerShard` and `split_run` in the Rust router only handle row ranges, so
column-sharded tensors are matched and addressed incorrectly. Extend
`TrainerShard` to carry the same column range metadata as the Python router
(`col_start`/`col_end`), then update `split_run` to match shards using both row
and column bounds when computing the target `agent_index` and `src_addr`; if
column-sharded layouts are not supported, add an explicit validation/error path
instead of falling back to the first row match.
| fn unpack_runs(flat: &[i64]) -> Vec<(i64, i64)> { | ||
| flat.chunks(2).map(|c| (c[0], c[1])).collect() | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
unpack_runs panics on odd-length input.
chunks(2) yields a final 1-element slice when flat.len() is odd, and c[1] then indexes out of bounds, panicking the request handler. Element-run payloads should always be pairs, but this input arrives over the wire from a client, so a malformed request would crash the task rather than return a Status.
🛡️ Proposed guard
-fn unpack_runs(flat: &[i64]) -> Vec<(i64, i64)> {
- flat.chunks(2).map(|c| (c[0], c[1])).collect()
-}
+fn unpack_runs(flat: &[i64]) -> Vec<(i64, i64)> {
+ flat.chunks_exact(2).map(|c| (c[0], c[1])).collect()
+}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| fn unpack_runs(flat: &[i64]) -> Vec<(i64, i64)> { | |
| flat.chunks(2).map(|c| (c[0], c[1])).collect() | |
| } | |
| fn unpack_runs(flat: &[i64]) -> Vec<(i64, i64)> { | |
| flat.chunks_exact(2).map(|c| (c[0], c[1])).collect() | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@modelexpress_server/src/weight_sync/router.rs` around lines 125 - 127, The
unpack_runs helper currently assumes every chunk has two elements, so an
odd-length payload can panic when indexing c[1]. Update unpack_runs in router.rs
to validate the input length before unpacking and handle malformed data safely,
returning an error or rejecting the request instead of crashing; if this
function is used by the weight sync request flow, make sure the caller around
unpack_runs propagates a Status for invalid client payloads.
| async fn publish_trainer_table( | ||
| &self, | ||
| request: Request<PublishTrainerTableRequest>, | ||
| ) -> Result<Response<PublishTrainerTableResponse>, Status> { | ||
| let req = request.into_inner(); | ||
| let mut state = self.state.write().await; | ||
| state | ||
| .trainer_tables | ||
| .insert(req.model_key.clone(), req.table_payload); | ||
| // Invalidate any cached plans that used this model's table | ||
| state | ||
| .plan_key_to_id | ||
| .retain(|k, _| !k.starts_with(&req.model_key)); | ||
| Ok(Response::new(PublishTrainerTableResponse { ok: true })) | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Memory leak: invalidating a trainer table drops plan_key_to_id mappings but never removes the underlying plans.
retain prunes plan_key_to_id, yet the corresponding CachedPlan entries in state.plans are never removed, so they leak. Since publish_trainer_table is called "once per broadcast step", orphaned plans accumulate unbounded across training. Collect the removed plan_ids and drop them from plans too.
🔧 Proposed fix
- // Invalidate any cached plans that used this model's table
- state
- .plan_key_to_id
- .retain(|k, _| !k.starts_with(&req.model_key));
+ // Invalidate any cached plans that used this model's table
+ let stale_ids: Vec<String> = state
+ .plan_key_to_id
+ .iter()
+ .filter(|(k, _)| k.starts_with(&req.model_key))
+ .map(|(_, id)| id.clone())
+ .collect();
+ state
+ .plan_key_to_id
+ .retain(|k, _| !k.starts_with(&req.model_key));
+ for id in stale_ids {
+ state.plans.remove(&id);
+ }📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| async fn publish_trainer_table( | |
| &self, | |
| request: Request<PublishTrainerTableRequest>, | |
| ) -> Result<Response<PublishTrainerTableResponse>, Status> { | |
| let req = request.into_inner(); | |
| let mut state = self.state.write().await; | |
| state | |
| .trainer_tables | |
| .insert(req.model_key.clone(), req.table_payload); | |
| // Invalidate any cached plans that used this model's table | |
| state | |
| .plan_key_to_id | |
| .retain(|k, _| !k.starts_with(&req.model_key)); | |
| Ok(Response::new(PublishTrainerTableResponse { ok: true })) | |
| } | |
| async fn publish_trainer_table( | |
| &self, | |
| request: Request<PublishTrainerTableRequest>, | |
| ) -> Result<Response<PublishTrainerTableResponse>, Status> { | |
| let req = request.into_inner(); | |
| let mut state = self.state.write().await; | |
| state | |
| .trainer_tables | |
| .insert(req.model_key.clone(), req.table_payload); | |
| // Invalidate any cached plans that used this model's table | |
| let stale_ids: Vec<String> = state | |
| .plan_key_to_id | |
| .iter() | |
| .filter(|(k, _)| k.starts_with(&req.model_key)) | |
| .map(|(_, id)| id.clone()) | |
| .collect(); | |
| state | |
| .plan_key_to_id | |
| .retain(|k, _| !k.starts_with(&req.model_key)); | |
| for id in stale_ids { | |
| state.plans.remove(&id); | |
| } | |
| Ok(Response::new(PublishTrainerTableResponse { ok: true })) | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@modelexpress_server/src/weight_sync/service.rs` around lines 108 - 122, The
invalidation in publish_trainer_table only removes entries from plan_key_to_id,
leaving orphaned CachedPlan values behind in state.plans. Update
publish_trainer_table to collect the plan IDs removed by the retain call on
plan_key_to_id, then use those IDs to delete the matching CachedPlan entries
from plans as part of the same invalidation path. Keep the fix localized to
publish_trainer_table and the state fields plan_key_to_id and plans so both
caches stay in sync.
| // Reuse an existing plan for this plan_key if one was already built. | ||
| { | ||
| let state = self.state.read().await; | ||
| if let Some(plan_id) = state.plan_key_to_id.get(&req.plan_key) { | ||
| return Ok(Response::new(BuildPlanResponse { | ||
| plan_id: plan_id.clone(), | ||
| })); | ||
| } | ||
| } | ||
|
|
||
| // Fetch the TrainerTable and deserialize it. | ||
| let table_bytes = { | ||
| let state = self.state.read().await; | ||
| state | ||
| .trainer_tables | ||
| .get(&req.model_key) | ||
| .cloned() | ||
| .ok_or_else(|| { | ||
| Status::not_found(format!( | ||
| "TrainerTable not found for model_key {:?}", | ||
| req.model_key | ||
| )) | ||
| })? | ||
| }; | ||
|
|
||
| let table: TrainerTableJson = serde_json::from_slice(&table_bytes) | ||
| .map_err(|e| Status::internal(format!("Failed to decode TrainerTable: {e}")))?; | ||
|
|
||
| // Route the resolved regions in Rust (no torch). | ||
| let descriptors = route_regions(&req.regions, &table); | ||
|
|
||
| let plan_id = Uuid::new_v4().to_string(); | ||
| let mut state = self.state.write().await; | ||
| state | ||
| .plans | ||
| .insert(plan_id.clone(), CachedPlan { descriptors }); | ||
| state.plan_key_to_id.insert(req.plan_key, plan_id.clone()); | ||
|
|
||
| Ok(Response::new(BuildPlanResponse { plan_id })) | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the relevant file and locate the build_plan implementation and state types.
ast-grep outline modelexpress_server/src/weight_sync/service.rs --view expanded
printf '\n--- build_plan references ---\n'
rg -n "build_plan|plan_key_to_id|CachedPlan|trainer_tables|plans" modelexpress_server/src/weight_sync/service.rs
printf '\n--- surrounding lines ---\n'
sed -n '120,230p' modelexpress_server/src/weight_sync/service.rsRepository: ai-dynamo/modelexpress
Length of output: 7002
build_plan still has a race between the cache check and insert.
Two concurrent requests with the same plan_key can both miss the read-locked lookup, both route, and both insert different plan_ids; the later write overwrites plan_key_to_id, leaving the earlier CachedPlan unreachable from invalidate_plan. Recheck under the write lock before inserting, or keep the lookup/insert path under one lock, so the endpoint stays idempotent.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@modelexpress_server/src/weight_sync/service.rs` around lines 160 - 199, The
build_plan flow in WeightSyncService has a check-then-insert race between the
initial state.read() lookup of plan_key_to_id and the later state.write()
insert, which can create duplicate plans for the same req.plan_key. Fix it by
rechecking the cache under the write lock or by moving the lookup and insert
into a single locked critical section in build_plan, so only one plan_id is
created and stored for a given plan_key. Use the existing plan_key_to_id and
plans fields, along with build_plan and invalidate_plan, to keep the cached plan
mapping consistent and idempotent.
nicolasnoble
left a comment
nicolasnoble
left a comment
There was a problem hiding this comment.
Went through this properly. Inline comments are the code-level stuff, roughly blocker-first. Two I'd really want fixed before this merges: the M2N barrier has no timeout or reaper, so one inference worker that never registers hangs the whole collective and leaks the session, and the pull refit writes into live param memory with no rollback, so a mid-sync failure leaves the engine serving mixed old/new weights.
Few housekeeping things that don't need inline threads:
protocol/test_weight_transfer.pyis byte-identical totests/test_weight_transfer.pyand lives inside the package, soinclude=["modelexpress*"]ships it in the wheel and the suite double-counts it. Drop the copy, pintestpaths.tests/test_e2e_nccl_m2n.pyis named for NCCL e2e transport but it's fully mocked, no NCCL and no transport. Rename or drop, especially now that #497 has real ones.protocol/router.pyis a dead byte-identical copy ofplanner/router.py, and_shard_widthlooks unused.- The new
WeightLoaderAdapter+_CtxEngineAdapterduplicateEngineAdapterwithout adding a capability it lacks - could fold in. - The four new env vars (
MX_TRAINER_TABLE_KEY,MX_TRAINER_SYNC_TIMEOUT,MX_REDIS_URL,MX_WEIGHT_SYNC_SERVER) aren't indocs/DEPLOYMENT.md.
Happy to talk through any of it.
| // --------------------------------------------------------------------------- | ||
|
|
||
| fn unpack_runs(flat: &[i64]) -> Vec<(i64, i64)> { | ||
| flat.chunks(2).map(|c| (c[0], c[1])).collect() |
There was a problem hiding this comment.
chunks(2) yields a 1-element final chunk on odd-length input, so c[1] panics, and this comes straight off the wire. chunks_exact(2), or validate even length and return invalid_argument.
| Ok(Response::new(InvalidatePlanResponse { ok: true })) | ||
| } | ||
|
|
||
| async fn register_m2n_worker( |
There was a problem hiding this comment.
No timeout or reaper on this barrier. If a worker never registers, everyone else polls forever and the pending entry leaks. The p2p service has a reaper; this one doesn't. Can we stamp entries with a creation time and reap stale ones?
| }, | ||
| ); | ||
|
|
||
| if entry.registrations.len() < entry.total_workers as usize { |
There was a problem hiding this comment.
total_workers is an int32 off the wire; a negative value casts to a huge usize and the barrier can never fire. A late worker also gets a plan_id but no slice and polls forever. Reject total_workers <= 0 and unknown-rank late registrations.
| } | ||
|
|
||
| let count = src_rem.min(dst_rem); | ||
| let src_addr = shard.device_addr + (src_rel * element_size) as u64; |
There was a problem hiding this comment.
This address math runs on unvalidated wire ints, so it overflow-panics in debug and wraps in release into a bad descriptor. checked_add/checked_mul and reject on overflow.
| return result | ||
|
|
||
| def update_weights(self, ctx: LoadContext) -> None: | ||
| """Execute a weight sync using the pre-built static plan. |
There was a problem hiding this comment.
No rollback here. sync() writes into live param memory, so if it throws partway the model's left with mixed old/new weights and keeps serving. There's a rollback() on the strategy but nothing calls it on this path. Stage into scratch and commit atomically, or force a reinit on any sync failure?
|
|
||
| runs: list[tuple[int, int]] = [] | ||
|
|
||
| def _collect(base: int, dim: int) -> None: |
There was a problem hiding this comment.
Builds one tuple per element before merging; a contiguous lm_head/embedding is hundreds of millions of tuples in a Python loop. Worth a contiguous fast path so it doesn't OOM on real tensors.
| """Run model.load_weights() with LazyWeights and return recorded copies.""" | ||
| recorder = BakeRecorder() | ||
| with recorder: | ||
| if hasattr(model, "load_weights"): |
There was a problem hiding this comment.
When the model has no load_weights() this returns [] and the whole sync reports success having moved nothing, so the engine serves uninitialized weights. Should raise on a zero-byte plan.
| sum(d.nbytes for d in self._descriptors) / 1e9, | ||
| ) | ||
|
|
||
| remote_agents: dict[int, str] = {} |
There was a problem hiding this comment.
add_remote_agent runs every initialize() but teardown() never deregisters, so a trainer that reshards repeatedly accumulates registrations. Track and remove them? (Same on the push side, push.py ~L75.)
| logger = logging.getLogger("modelexpress.strategy_trainer_pull") | ||
|
|
||
|
|
||
| def _use_server_planner() -> bool: |
There was a problem hiding this comment.
bool(os.environ.get("MX_WEIGHT_SYNC_SERVER", "")) treats =0 and =false as enabled, since any non-empty string is truthy. Parse through the envs registry with real truthy semantics.
| } | ||
|
|
||
| /// Shared server state (plan cache + raw table blobs). | ||
| struct WeightSyncState { |
There was a problem hiding this comment.
These maps are insert-only with no TTL and superseded plans get orphaned, so it's a slow leak over a long run. Evict on retarget/invalidate.
2 participants
Summary
weight_transfersub-package for trainer → inference weight sync with two transfer modes: PULL (inference workers read from trainer) and PUSH (trainer writes to inference workers)[R, C/tp]tileTests
88 tests, 0 GPU required for CI (pure-Python stubs):
test_rl_weight_sync.pytest_e2e_nccl_m2n.pytest_e2e_primerl.pytests/gpu/test_e2e_primerl_real_model.pyReal-model GPU tests (
tests/gpu/, skipped in CI, run on devdesktop with 2× RTX 6000 Ada):Test plan
python -m pytest tests/test_rl_weight_sync.py tests/test_e2e_nccl_m2n.py tests/test_e2e_primerl.py— 58 tests, no GPUpython -m pytest tests/gpu/test_e2e_primerl_real_model.py— 17 tests, requires 2× CUDA GPUs🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Bug Fixes
Tests