【Task.026】Version-gated Mooncake loss guards - #278
Open
overloadedHenry wants to merge 23 commits into
Open
Conversation
overloadedHenry
requested review from
Aurelius84,
NINGBENZHE and
Yangruipis
as code owners
August 14, 2026 13:57
Add a default-off, safely-degrading RDMA path for the rollout->train sample transfer, reusing TransferQueue's existing MooncakeStore backend. No transfer_queue/ or payload-shape changes; default flags (simple+off) short-circuit to the original SimpleStorage path, so existing jobs are unaffected. Code - 4 intent-only flags (--tq-storage-backend / --tq-rdma-mode / --tq-rdma-device / --tq-use-gdr); Mooncake internals (endpoint, buffer, segment, timeout, master) stay internal - driver probes every alive GPU node before tq.init and AND-reduces a single job-level effective config; graded fallback GDR -> host RDMA -> Mooncake/TCP -> SimpleStorage; required mode fails fast on probe failure and capacity shortfall - hard_pin=True + segment-capacity precheck so produced-but-unconsumed data is never silently evicted - reap half-initialised TransferQueueController before tq.init (F10 anti-hang, incl. get_config timeout) and unmount the Mooncake segment on teardown so dead endpoints don't leak past client_ttl - GDR marked EXPERIMENTAL: not probed (probe runs without a CUDA context); decided per worker at runtime with a fallback WARNING Tests (52 passed) - test_rdma_probe.py (26): config validation, AND-reduction, multi-node fan-out, capacity, storage_backend key selects the manager - test_tq_failure_paths.py (19): reaper/timeout, teardown order, retry, disconnect, auto-degradation, MooncakeStore byte-exact - test_tq_dataplane_behavior.py (7): real SimpleStorage connection, backpressure, empty-get, repeat-put, cleanup - CI-safe: tests needing real transfer_queue/mooncake skip on the CPU CI single-file stub via real-submodule detection Benchmark + docs - scripts/benchmarks/tq_cross_node_bench.py: C0/C1/C2 same-topology cross-node (256M-4.5G, 5-run mean, per-run wire verification + async-tail diagnostic) - docs/draft/transfer_queue_rdma.md: master lifecycle, resource ownership, log reading, troubleshooting, known limits Measured (2-node cluster, 5-run mean): cross-node get C2/C1 = +28%..+126% across 256M..4.5G; put +45%..+146%.
- Run first initialization in a dedicated Ray owner actor with a bounded timeout - Clean partial controllers with owner tokens and restrict global close to the owner - Fall back once in auto mode and fail fast in required mode - Probe the external master across all data-plane nodes before initialization - Require retry-and-raise storage operations and storage-before-ready notification - Report requested GDR intent separately from per-worker runtime status --- - Cover master failures, owner cleanup, capacity errors, and notification ordering - Add multimodal byte-exact validation and tiered cross-node benchmark output - Separate mock coverage from opt-in real-environment acceptance checks --- - Document external master prerequisites, fallback behavior, and ownership rules - Record capacity guarantees, upstream correctness requirements, and validation tiers
Preserve Mooncake TCP semantics when RDMA is off and reject unsafe controller attachments. Fail closed on incomplete Mooncake batch results, removal failures, and unsuccessful production-status notifications. Make RDMA benchmarks and correctness tests safe for CPU-only CI environments.
# ✅ Tests ## Cover the production multimodal container (list[dict] slow path) - relax/utils/payload_digest.py: canonical leaf-level SHA-256 fingerprints (contiguous-CPU-normalized storage bytes; NaN-safe, stricter than torch.equal; NestedTensor rows == list rows; NonTensorData/Stack unwrap) - tests/utils/mm_payload_fixtures.py: payload source shared by tests -- real fixture (auto-verified against its manifest) with production-structured synthetic fallback for CI; tier reported in every assertion - test_tq_dataplane_behavior.py: TestRealMultimodalFullLink -- full tq.init/put/get with multimodal_train_inputs as NonTensorStack via the production dict_to_tensordict, per-sample leaf digests aligned by sample_id - test_tq_failure_paths.py: TestMooncakeByteExact gains the msgpack non-tensor slow-path roundtrip (tcp/rdma), one spawn child per protocol to isolate the mooncake 0.3.10 in-session protocol-switch instability --- # ⭐ Feature ## Real-payload fixture generator + bench profile - scripts/benchmarks/make_multimodal_fixture.py: replays the exact rollout preprocessing chain (build_messages -> apply_chat_template -> process_vision_info -> HF processor -> remap_mm_train_inputs) on real dataset rows; double-run determinism check validates the F4 group-sharing assumption; emits leaf manifest + committable provenance JSON - tq_cross_node_bench.py: real-multimodal profile (fixture tiled to each payload tier, NonTensorStack column) with order-insensitive row-multiset digests; dtype+bytes row contract absorbs the scalar-row () vs [1] representation difference between SimpleStorage and MooncakeStore --- # 📝 Documentation ## Acceptance layering for real payloads - docs/draft/transfer_queue_rdma.md: fixture workflow, real vs synthetic tier reporting rules, real-multimodal bench command; troubleshooting row for the mooncake 0.3.10 TCP loopback SIGSEGV found by this tier - .gitignore: tests/fixtures/ (machine-local, hundreds of MB)
# 🐛 Bug Fix
## Silent TCP truncation traced to mooncake's memcpy fast path
- relax/utils/tq_correctness.py: correctness guards now default
MC_STORE_MEMCPY=0 (setdefault, operator can override). mooncake 0.3.10
auto-enables the memcpy fast path in TCP-only environments and that path
silently truncates cross-node gets: two-node forensic probes captured
rows zero-filled from 64 KiB-aligned offsets onward while every batch
code reported success (~50% of fresh-session first transfers; not
limited to the first transfer -- a canary transfer does not fully
prevent it; 12/12 sessions clean with memcpy off). The same path is
the single-node loopback SIGSEGV documented earlier; both symptoms are
gone with the guard (loopback multimodal re-run passes byte-exact).
RDMA sessions auto-disable memcpy, so the default is a no-op there.
---
# ✅ Tests
## Contract coverage for the new guard
- tests/utils/test_rdma_probe.py: validate_mooncake_runtime_contract now
must default MC_STORE_MEMCPY to "0" when unset and must respect an
explicit operator override ("1"); both skip on the CPU-CI transfer_queue
stub like the existing contract test
---
# 📝 Documentation
## Two-node acceptance record + updated troubleshooting
- docs/draft/transfer_queue_rdma.md: full 3x3x4-tier acceptance table
(36/36 byte-exact PASS, wire-proof PASS; C2/C1 get gain 2.1x-8.4x with
guarded-TCP as the honest C1 baseline); troubleshooting rows for the
memcpy silent truncation (fixed by guard), the loopback SIGSEGV (same
root cause, verified fixed), and master-aging batch_upsert -800 (fresh
master clears it); known-limitations note that C1 is a correctness
fallback, not a performance option
# 📝 docs - 将「双节点实测记录」日期化实验小节收敛为「参考吞吐区间」：只保留 get/put 量级区间与结论（供容量规划参考），逐档明细、逐轮分布与 原始 CSV 归入交付验收材料，不再在文档内维护 - 排障表三行（TCP 静默截断、回环 SIGSEGV、master 状态劣化）压缩为 「现象/原因/处理」一行式，剥离取证过程叙事（会话统计、探针细节） - `MC_STORE_MEMCPY=0` 守卫的行为说明移入「容量不足与正确性依赖」， 排障表引用之；补充 RDMA 会话不受影响与显式覆盖方式 - 「已知限制」与验收措辞去除开发过程口吻（"此前读数"、"本次开发 环境"），与 docs/draft 下其他使用指南的无时间性语态对齐
# 🐛 Bug Fix ## Fail closed on unsafe MC_STORE_MEMCPY (review: tq_correctness.py:179) - Reject startup when MC_STORE_MEMCPY=1 is set: the pinned mooncake 0.3.10 memcpy fast path silently truncates TCP transfers and can SIGSEGV; force the variable to 0 otherwise - Re-gate on the mooncake version once the pin moves past the fix ## Require an explicit Mooncake master endpoint (Codex P1) - resolve_mooncake_master_address() rejects a missing MC_MASTER_ADDRESS instead of assuming localhost:50051, which made every node of a multi-node job treat itself as the master ## Probe every usable HCA before degrading RDMA (Codex P1) - _select_usable_rdma_device() scans all devices, all ports, and the GID table of the first ACTIVE port; a node degrades only when no device passes both checks together - probe_node reports the jointly validated device instead of the lexicographically first one --- # ✅ Tests ## Cover the fail-closed and multi-HCA behaviours - test_contract_rejects_explicit_memcpy_enable asserts fail-fast - master-address tests for the required env contract - multi-HCA selection tests (down first device, all-down degradation)
# 🐛 Bug Fix ## Close the TQ owner if Controller construction fails (review) - Wrap the post-_initialize_data_system() construction sequence in exception cleanup: close the newly created TQ owner before re-raising so a failed Controller() cannot orphan a healthy named TransferQueueController whose next launch attaches with owner=None ## Detach worker Mooncake clients during teardown (Codex P1) - Expose detach_tq_client(), the attach-only inverse of attach_tq_client(); it deregisters the worker segment immediately instead of waiting for the master client_ttl - Base.__del__ detaches on Ray Serve replica shutdown (covers Actor, ActorFwd, Advantages, Critic, Rollout, SFT) - RolloutManager.dispose() and MegatronTrainRayActor.__del__ detach on worker teardown; force-kills still fall back to the master TTL --- # ✅ Tests ## Worker detach coverage - detach_tq_client delegates to the process-local close helper - Base.__del__ detaches only when a TQ client was attached
# 🐛 Bug Fix ## Restore the zero-change SimpleStorage default path (review) - --tq-storage-backend=simple runs upstream-identical init again: first tq.init in the Controller process, no _TransferQueueOwner actor, and plain tq.close() on teardown - Only addition is the F10 reaper, which acts solely on a provably half-initialised leftover controller that would otherwise hang init ## Bound worker attach and converge the job on handshake failure (review) - attach_tq_client now enforces one deadline over both hang sources: waiting for a served controller config (the F10 poll loop) and tq.init itself (native mooncake setup); override via RELAX_TQ_ATTACH_TIMEOUT_SECONDS, default 60 s - verify_cluster_attach runs a bounded attach handshake from every alive node -- Serve replicas and 0-CPU actors carry no placement binding, so the GPU-only probe cannot vouch for the real endpoints - The Controller aggregates handshake failures: auto closes Mooncake state and converges the whole job to SimpleStorage; off/required fail loudly; attached (foreign-owner) sessions never tear down or replace the winning controller --- # ✅ Tests ## Bounded-attach coverage - deadline expiry on a hung tq.init raises TqAttachTimeout - worker errors propagate unchanged; env override parsing validated - missing named controller times out instead of spinning forever
# 🐛 Bug Fix ## Derive worst-case payload from the token budget (Codex P1) - estimate_payload_bytes bounds one sample by seq_length: text at 32 B/token plus, for multimodal jobs, seq_length x 784 pixels/token x 12 B (ViT patch 14, merge 2, float32 RGB) -- ~77 MiB at 8k instead of the fixed 8 MiB guess that passed configs which later failed puts mid-training - Missing/non-positive seq_length fails fast instead of guessing ## Make the segment size configurable without code edits - resolve_global_segment_size() reads RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB (default 4 GiB); capacity validation and the client config share the value so the check can never pass a size the client does not mount - Capacity error message points at the env override --- # ✅ Tests ## Capacity derivation coverage - token-budget bound matches the review example (32 x ~77 MiB x 2 in-flight now exceeds the default 4 GiB and is rejected) - env override raises the ceiling and rejects garbage values - missing seq_length raises
# ♻️ Refactor ## Keep only read-only capability validation in tq_correctness (review) - ensure_mooncake_correctness_guards now validates the environment (memcpy fail-closed) and that the pinned TransferQueue ships the Mooncake retry APIs; it no longer replaces upstream private internals (__init__, _notify_and_wait) at runtime - The runtime loss guards move to a stacked, version-gated branch (feat/tq-mooncake-loss-guards) with an explicit removal condition, per maintainer guidance to keep monkey patches in their own PR - Patch-primitive tests move with them; retry/ACK-ordering contract tests stay because they assert upstream behaviour, not patches
# 📝 Documentation ## transfer_queue_rdma.md - MC_MASTER_ADDRESS is required on every node (no loopback default) - MC_STORE_MEMCPY=1 is rejected at startup (fail-closed); note the runtime patches split into feat/tq-mooncake-loss-guards - capacity pre-check derives worst-case payload from the token budget; RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB adjusts the segment size - startup flow gains the all-alive-node bounded attach handshake with unified auto fallback; lifecycle table gains handshake and worker-attach-timeout rows (RELAX_TQ_ATTACH_TIMEOUT_SECONDS)
overloadedHenry
force-pushed
the
feat/tq-mooncake-loss-guards
branch
from
August 14, 2026 14:05
59eb0e6 to
e328191
Compare
overloadedHenry
force-pushed
the
feat/tq-mooncake-loss-guards
branch
from
August 14, 2026 14:40
e328191 to
3d27f54
Compare
SFT switched from module-level `import transfer_queue as tq` to attach_tq_client, so monkeypatching relax.components.sft.tq.* fails path resolution in CI (ModuleNotFoundError; 'relax.components.sft' is not a package). Patch the attach helper the component actually uses.
overloadedHenry
changed the title
overloadedHenry
force-pushed
the
feat/tq-mooncake-loss-guards
branch
from
August 14, 2026 15:23
3d27f54 to
9e28778
Compare
# 🐛 Bug Fix ## Protect process-global client ownership - Record a generation lease for every worker attachment - Ignore stale destructor cleanup after an in-process replacement - Keep explicit teardown idempotent and generation-aware ## Simplify capability reduction - Remove unreachable GDR and device fallback branches - Avoid duplicate per-node capability logging --- # 📝 Documentation ## Clarify startup and capacity behavior - Document bounded attachment behavior on the default backend - Describe handshake resource pressure and capacity fallback symptoms --- # ✅ Tests ## Cover replacement and reduction behavior - Verify stale generations cannot close the current client - Verify RDMA transport and device selection invariants
overloadedHenry
force-pushed
the
feat/tq-mooncake-loss-guards
branch
from
August 15, 2026 20:06
fa8c6b1 to
964035c
Compare
# 🐛 Bug Fix ## Prevent attach timeout state from reaching reused workers - Limit each cluster attach handshake worker to one task invocation - Preserve zero retries so a failed handshake is not silently repeated --- # ✅ Tests ## Lock the one-shot Ray task contract - Assert cluster attach handshakes use max_calls=1 and max_retries=0 --- # 📝 Documentation ## Document lifecycle isolation and TCP connection pooling - Explain one-shot handshake cleanup after a timed-out initialization - Record the connection-pool requirement for long Mooncake TCP sessions
overloadedHenry
force-pushed
the
feat/tq-mooncake-loss-guards
branch
from
August 17, 2026 17:37
964035c to
c9404ea
Compare
# ✅ Tests ## Verify attach timeout isolation - Exercise the production cluster attach handshake in a one-shot Ray worker - Assert a timed-out initialization cannot mutate process-global state later ## Isolate native Mooncake roundtrips - Run TCP and RDMA sessions in separate bounded subprocesses - Use unique keys and require cleanup to finish before reporting success - Apply the production correctness guard to real backend tests
# ⭐ Feature ## Reintroduce the runtime loss guards behind an exact version gate - relax/utils/tq_mooncake_patches.py carries the patches split out of the RDMA enablement PR: per-retry batch result validation (_StrictMooncakeStoreProxy), raising removal failures, and a strict production-status ACK (_strict_notify_and_wait) - install refuses any transfer_queue other than the validated pin (0.1.10.dev0) so a pin bump forces re-validation instead of running unreviewed patches over private upstream internals - Removal condition documented in the module docstring: delete once the pinned TransferQueue ships the equivalent checks - ensure_mooncake_correctness_guards installs them after its read-only capability validation --- # ✅ Tests ## Moved and extended guard tests - patch-primitive and ACK tests moved from test_tq_failure_paths.py - new version-gate tests: unpinned transfer_queue is rejected
# 🐛 Bug Fix ## Preserve safe removal behavior - Accept the idempotent object-not-found result during batch removal - Continue raising non-idempotent removal failures ## Harden runtime guards - Prevent recursive proxy lookup when the wrapped store is absent - Close notification sockets when setup or connection fails - Keep guard installation idempotent --- # 📝 Documentation ## Align correctness guard descriptions - Describe the version-gated runtime modifications and removal conditions --- # ✅ Tests ## Cover successful and failure paths - Verify accepted removal results and rejected failures - Verify positive acknowledgements, missing controllers, and socket cleanup
overloadedHenry
force-pushed
the
feat/tq-mooncake-loss-guards
branch
from
August 17, 2026 19:56
c9404ea to
6d984ce
Compare
RexFlux
suggested changes
Aug 18, 2026
RexFlux
left a comment
RexFlux
left a comment
Collaborator
There was a problem hiding this comment.
本轮复核后，retry/get 短返回校验、remove 失败处理、positive/missing ACK，以及补丁安装幂等性本身都没有发现新的问题；使用固定的 TransferQueue 58054a3 运行新增测试为 15 passed。
目前只剩一项合入阻塞：版本门控只校验 0.1.10.dev0，不能唯一识别实际验证过的 58054a3 revision。请按 inline comment 补充精确 revision 或私有实现 fingerprint 校验，并增加同版本不同 revision 的拒
绝测试。
修复后，请先合入 #256，再将 #278 rebase 到最新 main，确认 Files changed 只剩 tq_correctness.py、tq_mooncake_patches.py 和 test_tq_mooncake_patches.py，重新跑完 CI。
另外，#278 中依赖真实 TransferQueue 的 6 个集成用例在 GitHub CPU CI 会被跳过，而现有双节点日志早于 #278 的 runtime patch。建议最终 rebase 后补一次小规模 Mooncake/TCP 或 RDMA roundtrip/fully-
async smoke；不需要重跑完整性能矩阵
# 🐛 Bug Fix ## Bind runtime guards to an exact TransferQueue build - Require the validated package version and full Git revision - Fail closed on missing or malformed provenance metadata - Reject mismatched distributions and shadowed modules --- # ✅ Tests ## Cover revision provenance failures - Reject same-version builds from an unvalidated revision - Cover invalid VCS metadata, version mismatches, and module shadowing - Keep flat-stub CPU CI independent of Mooncake and GPU packages
Contributor
Author
感谢耐心审核，已增加精确 revision校验。小规模 Mooncake/TCP 已跑完。稍后可提供 log. |
Contributor
Author
|
fully_async_4step_attempt4_sanitized.log 此处为最新的 RDMA fully async 小规模测试。 |
RexFlux
approved these changes
Aug 19, 2026
RexFlux
left a comment
RexFlux
left a comment
Collaborator
There was a problem hiding this comment.
复核了最新 head 2e2b99f。上一轮关于精确 revision 门控的问题已经关闭：当前实现同时校验 TransferQueue package version 和 direct_url.json 中的完整 Git commit，只接受已验证的
58054a33834aadbcf76aacd6b1e32e25c030f2c9；metadata 缺失、同版本不同 revision、格式错误和模块路径不一致都会 fail closed，并有对应测试覆盖。
本地使用精确 pin 的 TransferQueue 运行相关测试为 78 passed、5 skipped，最新 CI 全部通过。补充的 fully-async RDMA 日志也完成了 4 个 rollout 和 4 个训练 step，MooncakeStore/RDMA 生效且没有发生
fallback，最终任务正常退出。
因此 #278 代码层面可以 Approve。请仍按约定先合入 #256，再确认 #278 相对最新 main 只剩 tq_correctness.py、tq_mooncake_patches.py 和 test_tq_mooncake_patches.py 三个文件，重新跑 CI 后合入。
Collaborator
|
hi，本PR的修复内容，请将相关fix提交到red-infra下维护的TransferQueue仓库吧，就不在relax主仓使用monkey patch的方式合入了：https://github.com/redai-infra/TransferQueue |
Contributor
Author
好的。 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
【Task.026】Version-gated Mooncake loss guards（PR 2/2）
这个 PR 解决什么
pin 的 TransferQueue（0.1.10.dev0）在 MooncakeStore 路径上有三个可导致静默丢数据的缺口：
batch_remove失败只打日志不抛错；这些修复的正确归宿是 TransferQueue 上游。在上游合入并升 pin 之前，Relax 用运行时补丁兜底。按 #256 的评审意见（monkey patch 覆盖上游私有实现，应独立成 PR、限定精确版本、写明删除条件），从主 PR 拆出。
改了什么
新增
relax/utils/tq_mooncake_patches.py（安装过程进程内幂等）：_StrictMooncakeStoreProxy：包装 store，batch_upsert_from/batch_get_into/batch_remove的**每次调用（含重试）**逐 key 校验返回码数量与结果，remove 失败显式抛错_strict_notify_and_wait：production-status 通知必须在 deadline 内收到正 ACK，超时或被拒都显式抛错，不再静默当成功install_mooncake_loss_guards()：安装并验证补丁标记，失败即抛错_PATCHED_TQ_VERSIONS = ("0.1.10.dev0",)——任何其他 transfer_queue 版本直接拒绝启动，强制升 pin 时重新评估补丁适用性，而不是对新版本运行未验证的补丁接线只有一处：
tq_correctness.ensure_mooncake_correctness_guards在只读能力校验通过后调用安装。删除条件
当 pinned TransferQueue 自带以下能力时，删除本模块与其唯一调用点即可完成下线：
测试
tests/utils/test_tq_mooncake_patches.py：