feat(data-plane): enable RDMA transport for TransferQueue - #256
feat(data-plane): enable RDMA transport for TransferQueue#256overloadedHenry wants to merge 31 commits into
Conversation
|
@gongshaotian 此处按计划做了第一轮实现。 |
There was a problem hiding this comment.
Pull request overview
This PR wires up an RDMA-capable TransferQueue data plane in Relax by introducing intent-only CLI flags, probing RDMA/Mooncake capability across GPU nodes before tq.init, AND-reducing results into a job-unique effective config with graded fallback, and adding lifecycle helpers + tests/benchmarks/docs to validate failure paths and byte-exactness.
Changes:
- Add TransferQueue RDMA intent flags and a driver-side pre-
tq.initprobe/reduction flow to select MooncakeStore (RDMA/TCP) or fall back to SimpleStorage. - Introduce reusable helpers for TransferQueue controller reaping and Mooncake segment unmounting to prevent hangs/leaks.
- Add unit/integration tests, benchmarks, and operational docs covering degradation, retries, teardown ordering, and byte-exact round-trips.
Reviewed changes
Copilot reviewed 12 out of 12 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
relax/utils/arguments.py |
Adds CLI flags for backend selection and RDMA/GDR intent. |
relax/core/controller.py |
Adds backend resolution logic (probe → reduce → effective config) and uses lifecycle helpers on teardown. |
relax/utils/rdma_probe.py |
Implements node probe + cluster fan-out + AND-reduction + config validation. |
relax/utils/tq_config.py |
Centralizes TQ backend config building and segment-capacity precheck. |
relax/utils/tq_lifecycle.py |
Adds controller reaping and Mooncake segment unmount teardown helpers. |
tests/utils/test_rdma_probe.py |
CPU-only unit tests for validation, probing, reduction, and config building. |
tests/utils/test_tq_failure_paths.py |
Failure-path tests for controller reaping, retry/raise behavior, degradation, and Mooncake byte-exactness (skipped when unavailable). |
tests/utils/test_tq_dataplane_behavior.py |
Integration tests for SimpleStorage dataplane behavior contracts and byte-exactness. |
scripts/benchmarks/tq_rdma_bench.py |
Single-node benchmark comparing SimpleStorage vs Mooncake/TCP vs Mooncake/RDMA. |
scripts/benchmarks/tq_cross_node_bench.py |
Cross-node benchmark mirroring Relax’s persistent-actor usage and on-wire transport verification. |
scripts/benchmarks/cross_node_rdma_bench.py |
Raw MooncakeDistributedStore benchmark to isolate transport-layer TCP vs RDMA. |
docs/draft/transfer_queue_rdma.md |
Usage/ops guide including downgrade ladder, logs, and troubleshooting. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| # 2. SimpleStorage short-circuit (default, zero behavior change). | ||
| if backend == "simple" or mode == "off": | ||
| from relax.utils.tq_config import build_simple_storage_config | ||
|
|
||
| return build_simple_storage_config( | ||
| total_storage_size=total_storage_size, | ||
| num_data_storage_units=self.config.num_data_storage_units, | ||
| ) |
| manager = MagicMock() | ||
| if store_client is None: | ||
| del manager.storage_client # SimpleStorage manager has no storage_client | ||
| else: | ||
| manager.storage_client = store_client |
| def build_simple_storage_config(total_storage_size: int, num_data_storage_units: int) -> dict[str, Any]: | ||
| """Build the ``backend`` dict for SimpleStorage (current default | ||
| behavior).""" |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 20 out of 20 changed files in this pull request and generated no new comments.
Suppressed comments (3)
relax/core/controller.py:247
--tq-storage-backend=mooncake --tq-rdma-mode=offcurrently short-circuits to SimpleStorage (because ofif backend == "simple" or mode == "off"). This contradicts the CLI help text whereoffmeans "no RDMA" (i.e., MooncakeStore over TCP), and it also bypassesvalidate_mooncake_runtime_contract()for that configuration.
# 2. SimpleStorage short-circuit (default, zero behavior change).
if backend == "simple" or mode == "off":
from relax.utils.tq_config import build_simple_storage_config
return build_simple_storage_config(
relax/utils/tq_config.py:76
build_simple_storage_configis typed to requiretotal_storage_size: int, butscripts/benchmarks/tq_cross_node_bench.pycalls it withNoneto request unlimited capacity. The signature should reflect the actual accepted value to avoid type-checking drift.
def build_simple_storage_config(total_storage_size: int, num_data_storage_units: int) -> dict[str, Any]:
scripts/benchmarks/tq_rdma_bench.py:306
run_one()callsclose_tq_and_wait()before anytq.init(), andclose_tq_and_wait()callsray.get_actor(...). Without a priorray.init(), this benchmark will fail immediately with Ray not initialized (it only catchesValueError, not the runtime init error).
def main():
"""Run the benchmark across all requested payload/field/config
combinations."""
args = parse_args()
TransferQueue RDMA 性能补充：双节点真实多模态载荷结论摘要
测试环境
两端软件版本完全一致。测试通过 Ray 测试对象与配置配置
三种配置均通过 TransferQueue 完成 载荷 profile
真实 profile 的 fixture 由以下生产链路生成，不在测试中重新实现等价逻辑： fixture 共 12 个样本、约 210 MiB， 测试参数
正式实测结果下表为
结果解读
正确性与线连证明本次验收不是只采集吞吐，而是将正确性作为性能数据的准入条件：
最终结果为 36/36 测量点 byte-exact PASS，wire-proof 全部通过。 测试中发现并修复的 TCP 数据面问题双节点正式矩阵首次运行时，TCP 会话出现了返回成功但内容被截断的情况：
修复在 os.environ.setdefault("MC_STORE_MEMCPY", "0")该守卫在每个进程创建或附着 Mooncake client 前执行，默认关闭有问题的 memcpy 路径，同时尊重运维显式设置的环境变量。守卫对 RDMA 会话无影响，并增加了“默认置 0”和“显式值不覆盖”两条契约测试。 因此，守卫后的 C1 吞吐（0.78 稳定性与复现注意事项
复现命令在具备真实模型和数据集的环境中，先按 PYTHONPATH=. python -u scripts/benchmarks/tq_cross_node_bench.py \
--master <master-host>:50051 \
--nodeb-ip <node-b-ip> \
--device <rdma-device> \
--payload-profiles synthetic multimodal real-multimodal \
--payload-mib 256 1024 2048 4096 \
--repeats 5 \
--require-wire-proof建议 TCP、RDMA 和 SimpleStorage 分别在独立进程中运行，并在结果中标注 fixture 来源（ 实验结论在真实双节点环境下，TransferQueue 的 RDMA 数据面已经通过真实多模态载荷验收：三种配置、三类 profile、四个容量档位共 36 个测量点全部逐字节一致，RDMA 相对 TCP 的 |
|
@codex review following the repository AGENTS.md and skills/code-review/SKILL.md |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1df2b8d8ed
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
|
||
| def resolve_mooncake_master_address() -> str: | ||
| """Return the externally managed Mooncake master endpoint.""" | ||
| return os.environ.get("MC_MASTER_ADDRESS", "localhost:50051") |
There was a problem hiding this comment.
When MC_MASTER_ADDRESS is absent in a multi-node Mooncake run, every node interprets this fallback as its own loopback endpoint, so auto abandons Mooncake and off/required abort even if a shared master exists elsewhere. Reject the missing deployment configuration instead of embedding an endpoint; hardcoded endpoints are also explicitly prohibited by the repository rules.
AGENTS.md reference: AGENTS.md:L55-L57
Useful? React with 👍 / 👎.
| for dev in sorted(os.listdir(base)): | ||
| return _check_port_active(dev, port) |
There was a problem hiding this comment.
On a multi-HCA host where the lexicographically first device has an inactive port or no GID but a later device is usable, this loop returns the first failure immediately; the analogous GID loop does the same. Consequently auto unnecessarily degrades the whole job to TCP and required rejects a valid cluster. Select one device whose active port and usable GID both pass rather than assuming the first HCA and fixed port/GID configuration.
AGENTS.md reference: AGENTS.md:L59-L59
Useful? React with 👍 / 👎.
| # Conservative: 8 MiB per sample when multimodal is enabled (real range | ||
| # 7.4 MiB for a 400-token image to hundreds of MiB at max token budget). | ||
| per_sample_mb = 8 if getattr(args, "multimodal_keys", None) is not None else 0 | ||
| return rollout_batch * n_samples * per_sample_mb * 1024 * 1024 |
There was a problem hiding this comment.
For multimodal jobs near the configured token/image limits, the fixed 8 MiB estimate can be orders of magnitude below the real payload even though the comment acknowledges sizes of hundreds of MiB. For example, 32 samples at 77 MiB with max_staleness=1 pass this check as 512 MiB but require about 4.9 GiB, exceeding the fixed 4 GiB hard-pinned segment and failing puts after training has started. Derive a defensible upper bound from the processor limits or make the segment size configurable instead of treating this lower bound as capacity validation.
Useful? React with 👍 / 👎.
| self.data_system_client = attach_tq_client( | ||
| self.args.tq_config, | ||
| requested_gdr=getattr(self.args, "tq_use_gdr", False), | ||
| role="rollout_worker", | ||
| ) |
There was a problem hiding this comment.
When a normal shutdown, global restart, or in-place service restart destroys this worker, its newly attached Mooncake client is never detached: RolloutManager.dispose() only stops monitors and engines, and the analogous component/Megatron lifecycles also omit client cleanup. Since this change itself documents that an unclosed storage client leaves its segment registered until the master TTL, an immediate restart can encounter stale endpoints and Failed to open segment errors. Expose an attach-only detach operation and invoke it from every worker teardown hook.
Useful? React with 👍 / 👎.
| init_result = initialize_tq_with_fallback( | ||
| tq_config, | ||
| mode=getattr(self.config, "tq_rdma_mode", "off"), | ||
| fallback_conf=fallback_config, | ||
| ) | ||
| self._tq_owner = init_result.owner |
There was a problem hiding this comment.
If DCS creation or any later service registration raises after this owner is assigned, Controller() never returns, so train.main() never installs _ctrl, signal handlers, or the atexit cleanup. The owner actor can disappear while its healthy named TransferQueueController survives; the next launch then deliberately attaches with owner=None, making subsequent shutdown a no-op and leaving the global TQ state orphaned. Wrap the post-initialization sequence in exception cleanup that closes the newly created owner before re-raising.
Useful? React with 👍 / 👎.
RexFlux
left a comment
RexFlux
left a comment
There was a problem hiding this comment.
感谢补充真实 multimodal_train_inputs/list[dict] 路径和双节点验收数据，最新版本相比初版在真实性、逐字节校验和 benchmark 覆盖上完善了很多。
结合 Codex review 和补充检查，目前仍建议 Request changes，优先处理以下问题：
- Codex 已指出的 master 配置、多 HCA、容量估算、worker detach 和 Controller 构造失败清理问题；
- 当前已知会静默损坏数据的 MC_STORE_MEMCPY=1 仍可绕过 correctness guard，应在受影响版本上 fail closed；
- capability probe 只覆盖 GPU 节点，但实际 TQ owner 和 Serve endpoint 没有固定在这些节点上，探测结论不一定覆盖真实数据面；
- 普通 worker 的 tq.init attach 没有 timeout，也没有纳入 job-level auto fallback；
- 默认 SimpleStorage 仍会创建新的 owner Actor，与“zero behavior change/no extra resource”的描述不一致。
此外，PR 当前仍无法直接合入 main，存在以下 merge conflicts：
- relax/components/sft.py
- relax/core/controller.py
- relax/distributed/ray/rollout.py
PR 描述中的测试数量和部分 benchmark 说明仍是旧版本，也建议在 rebase 后一起更新。真实双节点结果已经在 Conversation 中补充得比较完整，但原始 CSV 目前只保存在运行节点的 /tmp 路径，外部 reviewer 无法访
问；建议上传或附加一份脱敏后的 CSV/日志作为验收附件。
考虑到当前 PR 已达到 25 个文件、约 5.7k 行新增，也建议重新确认首期范围：RDMA 接入主流程保留在本 PR，上游 TransferQueue correctness patch、运行时 monkey patch 和额外诊断工具尽量拆成独立 PR，降低
review 和后续维护成本。
| memcpy anyway, so this default is a no-op there; ``setdefault`` keeps an | ||
| explicit operator override (e.g. ``MC_STORE_MEMCPY=1``) possible. | ||
| """ | ||
| os.environ.setdefault("MC_STORE_MEMCPY", "0") |
There was a problem hiding this comment.
这里已经确认 mooncake 0.3.10 的 memcpy 路径会静默截断数据甚至触发 SIGSEGV，但 setdefault() 仍允许外部通过 MC_STORE_MEMCPY=1 绕过 correctness guard。对于当前明确 pin 且已确认存在数据损坏的版本，这里应
该 fail closed，而不是保留未经验证的运维覆盖能力。
建议当前版本强制设置为 0；如果检测到 MC_STORE_MEMCPY=1，则直接拒绝启动并给出明确错误。未来升级到已修复版本后，再根据版本判断是否允许开启。对应的 test_contract_respects_explicit_memcpy_override 也应该
改为断言 fail-fast。
| # --------------------------------------------------------------------------- | ||
|
|
||
|
|
||
| def _select_dataplane_node_ids(nodes: list[dict]) -> list[str]: |
There was a problem hiding this comment.
这里假设 TransferQueue data plane 只运行在 advertise GPU 的节点上，但当前 TQ client 是在 Actor/Rollout 等 Ray Serve replica 进程中初始化的，Serve deployment 本身并没有绑定到相应的 GPU placement group
或具体节点；新增的 _TransferQueueOwner 也是未设置 node affinity 的 0-CPU Actor。
因此可能出现 GPU 节点探测全部通过、driver 选择 RDMA，但实际 owner/producer/consumer 被调度到未探测的 CPU 节点，随后 tq.init/attach 失败。Codex 提到的多 HCA
问题只解决“单节点选错设备”，没有解决“探测的不是实际 endpoint”这一层。
建议先固定 owner 和必要 TQ client 的 placement，再探测这些实际节点；或者在所有必要 endpoint 完成有界 attach handshake 后，再确认 job-level effective config。
Supporting code：
- TQ client 在 Serve replica 初始化：relax/components/rollout.py:336 (https://github.com/redai-infra/Relax/blob/1df2b8d8ed4d43086520a61a0845999a1462f991/relax/components/rollout.py#L336-L338)
- Owner 未绑定节点：relax/utils/tq_lifecycle.py:320 (https://github.com/redai-infra/Relax/blob/1df2b8d8ed4d43086520a61a0845999a1462f991/relax/utils/tq_lifecycle.py#L320-L330)
| return status | ||
|
|
||
|
|
||
| def attach_tq_client(conf: Any, *, requested_gdr: bool, role: str) -> Any: |
There was a problem hiding this comment.
目前只有首次 _TransferQueueOwner 初始化通过 ray.get(..., timeout=...) 做了超时保护，但 Actor/Rollout/Critic/Megatron worker 都会从这里直接调用 tq.init()，没有 timeout，也没有被纳入
initialize_tq_with_fallback() 的事务。
如果某个实际 endpoint 上 Mooncake setup 卡住、controller 处于半初始化状态，或者该节点的 RDMA/master 条件与启动探测不同，这里仍可能无限等待或导致 Serve replica 启动失败；此前 owner 初始化成功后，auto
模式也不会再统一回退到 SimpleStorage。
建议给 attach 增加有界超时，并让必要 endpoint 的 attach 结果汇总到 driver；auto 模式下任一必要 endpoint 失败时，应统一清理 Mooncake 状态并收敛到同一个 fallback backend。
|
|
||
| # 2. SimpleStorage short-circuit (default, zero behavior change). | ||
| # ``mooncake + off`` is MooncakeStore/TCP, not SimpleStorage. | ||
| if backend == "simple": |
There was a problem hiding this comment.
这里的 short-circuit 只跳过了 RDMA probe，并没有保留原来的 SimpleStorage 初始化路径：_initialize_data_system() 后面仍然无条件调用 initialize_tq_with_fallback()，最终会创建额外的 _TransferQueueOwner
Actor，并把首次 tq.init 从 Controller 进程移动到该 Actor。
Mooncake 路径；如果确实希望同时改造默认路径，需要修正文档并补完整的默认路径回归测试。
| logger.debug(f"Failed to close TransferQueue notification socket: {error}") | ||
|
|
||
|
|
||
| def _install_store_guards(client_cls: type) -> None: |
There was a problem hiding this comment.
这里已经超出了普通 integration guard：Relax 在运行时替换 TransferQueue 的 init、_notify_and_wait 等私有实现，并直接依赖上游内部 ZMQ 协议。后续 TransferQueue 升级时，即使公开 API 没变，也可能因为
私有实现变化出现难以定位的问题。
更建议把这些 correctness fix 合入 TransferQueue 上游，然后 Relax 只升级 pin 并保留版本/能力校验。如果首期必须临时保留 monkey patch，建议拆成独立 PR，限制精确适用版本，并写清楚删除条件。
|
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
实际 Relax commit SHABenchmark 和 fully-async smoke 使用的是同一组 PR 功能代码，具体由以下两个 PR 合并而来：
两个节点在运行前均核对为：
资源监控口径本轮没有采集连续的平均 GPU utilization、peak VRAM 或 sender/receiver CPU utilization，因此不提供这些数值，也不从零散日志反推。 日志ERROR解释：
|
|
hi，考虑到目前PR的changes比较大，从维护性角度看，需要完成下代码瘦身才可合入。建议至少做到 1. 归类存放tq新增utilsrelax/utils下新增的大量tq专属文件，需要开一个submodule统一归类存放 2. 删除首期不必要的 GDR 支持GDR 不是题目要求，目前也没有完成实际生效验证。建议首期删除：
首期只交付 host-RDMA，后续有真实需求再单独实现 GDR。 3. 收窄配置面建议首期只保留： --tq-rdma-mode {off,auto,required} 语义可以收敛为：
Mooncake/TCP 继续作为 benchmark 的 C1 对照，不一定需要成为公开生产配置。这样可以删除 backend/mode 组合矩阵、TCP connection-pool 运维契约以及两级降级状态机。是否保留 Mooncake/TCP 中间回退，需 4. 用真实 attach 代替过度细化的 /sys 探测当前 rdma_probe.py 有 652 行，分别检查 HCA、port、GID、memlock、master 等，最后仍然要执行真实 attach handshake。 可以考虑：
真实 attach 已经是最终事实，重复维护一套启发式探测容易和底层实现漂移。 5. 按“单任务独占集群”简化生命周期当前 tq_lifecycle.py 有 734 行，其中包含：
但 RFC 已明确首期只支持单任务独占 Ray 集群。可以考虑收敛为：
要保留超时、半初始化清理和 Mooncake segment unmount，但可以删除多 job/并发 initializer 兼容逻辑。 6. 删除一次性验收工具建议最终只留一个精简版跨节点 benchmark：
测试也不应该简单删除覆盖，而是用参数化合并重复 case。建议从约 2,800 行收敛到 800～1,200 行，保留题目要求的连接、背压、超时、断连、重试、清理和逐字节一致性。 优先考虑以下删减：
|
了解，我会尽快完成修改。 |
|
还有一些不必要的代码封装可以调整下： relax/utils/tq_lifecycle.py:165 这里的 uses_mooncake() 只是对 _uses_mooncake() 的纯转发，没有增加新的语义。建议只保留一个公开的 uses_mooncake()，并让本模块内部也直接复用它，减少一层无收益封装。 relax/utils/rdma_probe.py:412 _alive_gpu_nodes() 目前只是对 ray.nodes() 和 _select_dataplane_node_ids() 的一行包装，主要作用只是方便测试 mock。建议直接在 probe_cluster_nodes() 中调用节点筛选逻辑，测试侧 mock ray.nodes()， relax/utils/rdma_probe.py:478 _probe_on_node() 只重新 import 并转发到 probe_node()，没有增加实际逻辑。这里可以直接使用 ray.remote(num_cpus=0.001)(probe_node) 创建远程函数，避免嵌套一个纯转发 wrapper。 这次代码瘦身建议同时检查类似的一次性 helper 和纯转发 wrapper：如果函数只服务一个调用点、没有形成独立语义或隔离必要的 Ray/线程边界，可以优先内联。需要保留的例外包括线程 watchdog、完整的远程 _attempt()、_handshake() 和 _bounded_tq_init() 中的 _run() 这种承载 fallback 复用、Ray 进程边界和超时异常传递的可以保留。 |
# ♻️ Refactor ## Limit the initial transport to host RDMA - Remove public GDR configuration and runtime status plumbing - Pin Mooncake client configuration to host memory transfers - Preserve compatibility checks for legacy GDR controllers --- # ✅ Tests ## Cover the host-RDMA contract - Remove deferred GDR behavior tests - Reject attaching to a controller configured for GDR --- # 📝 Documentation ## Clarify the initial transport scope - Document host RDMA as the supported first-phase path
# ♻️ Refactor ## Reduce public data-plane choices - Replace the backend and transport matrix with a three-state host RDMA mode - Route automatic capability failures directly to SimpleStorage - Keep Mooncake TCP confined to benchmark baselines --- # ✅ Tests ## Cover production backend decisions - Verify off, auto, and required mode behavior - Ensure the production resolver never selects Mooncake TCP - Validate the narrowed command-line interface on CPU-only environments --- # 📝 Documentation ## Describe the two-path production contract - Document direct fallback from host RDMA to SimpleStorage - Clarify master configuration behavior and benchmark-only TCP usage
- remove static RDMA probing and simplify backend configuration - enforce exclusive controller ownership and fail-closed cleanup - isolate handshake and train actor initialization failures - preserve byte-exact coverage with hardened payload checks
Remove redundant single-node and raw transport benchmarks along with the one-off fixture generator. Keep one C0/C1/C2 acceptance tool with selected-device wire proof, byte-exact NonTensorStack coverage, fail-closed teardown, and safe external fixture loading.
Parameterize straightforward backend and lifecycle cases while keeping complex timeout and cleanup paths explicit. Remove duplicate endpoint and fixture harnesses, retain byte-exact and process-isolation coverage, and align the acceptance documentation.
|
目前根据评审意见进行了深度修改，后面会立刻在 Relax 维护的 TransferQueue 中提 PR，并汇报多机通信情况。 |
Correct multimodal segment sizing, clean partial actor initialization, retain failed owner handles, and unmount storage on teardown. Enforce the Mooncake memcpy safety contract, sanitize topology logs, and align tests and deployment docs.
# 🐛 Bug Fix ## Make cross-node acceptance fail closed - Require exact counter selection and account for idle traffic - Persist byte-exact and wire-proof failures before aborting - Preserve validation failures when partition cleanup also fails - Record exact source and dependency provenance in result rows --- # ♻️ Refactor ## Consolidate correctness and lifecycle coverage - Reuse raw-byte payload helpers across acceptance and dataplane checks - Parameterize repeated configuration, lifecycle, and cleanup scenarios - Retain process-isolation, retry, backpressure, and multimodal coverage --- # 📝 Documentation ## Clarify acceptance requirements - Document counter scope, idle adjustment, TCP baseline, and result provenance
|
重测的时候会出现上游 TQ 的 Tensor shape 的问题。由于不属于本 PR 的范围，于是只在本地解决。解决后均可正常通过。 相关 issue： |
|
Another Question：上游 TQ 是否要传 Capability marker？然后本 PR 内的代码进行校验？ |
看了下最新的调整，目前PR瘦身是有效的，代码整体approve 关于nested tensor问题，这属于独立的 #286 数据消费兼容问题，不需要再并入 #256。你本地叠加修复完成重测是可以的，只需要在测试记录中注明实际使用了#256 head 和 #286 的本地修复，后续让 #286 独立推进即可。 关于 Capability marker，建议在 TransferQueue PR redai-studio/TransferQueue#5 中增加一个很小的 correctness contract version，用来表示 retry 结果校验、删除失败传播和 ACK fail-closed 已经具备。#256 中用这个 marker 替换当前的方法名和源码检查即可，不需要再扩展成 capability 矩阵，也不要重新引入之前删掉的探测或 patch 逻辑 Capability marker 可以尽量简单，例如TQ PR 5 在 transfer_queue/init.py 中增加： MOONCAKE_CORRECTNESS_CONTRACT_VERSION = 1 这里的 version=1 表示已经具备本次 PR #5 的三项正确性契约：
Relax #256 中只需要读取并比较一次： import transfer_queue as tq _REQUIRED_TQ_MOONCAKE_CONTRACT_VERSION = 1 actual = getattr(tq, "MOONCAKE_CORRECTNESS_CONTRACT_VERSION", 0) 建议用这个检查替换当前 retry method existence 和 inspect.getsource 检查，而不是在现有逻辑上继续叠加。这样 TQ 侧只增加一个常量，Relax 侧也是几行版本比较，不会抵消 这轮代码瘦身。 最后比较推荐的合入顺序是：
|
【Task.026】Enable RDMA transport for the TransferQueue data plane
Closes #217
这个 PR 解决什么
Relax 的数据面（rollout → train 的样本传输）走 TransferQueue 的 SimpleStorage，即 ZMQ over TCP。在 RoCE 机器上，多模态 payload 全程用 TCP 搬运、RDMA 网卡闲置：Qwen3-VL 的
pixel_values按默认 16384 token 预算单图就是 77 MB，32 图 batch 是 2.47 GB。TransferQueue 本身已带 MooncakeStore 后端，
protocol可取rdma，并支持 GDR staging、超大对象分块与重试。Relax 此前从未把它接出来：relax/core/controller.py把后端固定为 SimpleStorage。本 PR 做的是配置接入、能力探测与一致回退，不是传输层实现——transfer_queue/零改动。改了什么
一条默认关闭、不可用时安全回退的 RDMA 数据路径：
--tq-storage-backend、--tq-rdma-mode、--tq-rdma-device、--tq-use-gdr；mooncake + off表示 MooncakeStore/TCP（off只关 RDMA 传输，不改后端）tq.init之前探测每个存活且有 GPU 的节点，AND 归约出 job 级唯一的 effective config，所有组件读同一份auto统一收敛 SimpleStorage，off/required启动失败并列出节点tq.init，超时该 worker 立刻失败而不是无限挂起required直接失败，不静默降级hard_pin: true+ 容量预检：按 token 预算推导最坏情况 payload（文本seq_length × 32 B，多模态另加seq_length × 784 像素/token × 12 B，8k 序列约 77 MiB/样本）对比 segment 大小tq.init前回收半初始化的TransferQueueController（F10 防挂死）；拆除时卸载 Mooncake segment；所有 worker teardown 钩子（Serve 组件__del__、RolloutManager.dispose、Megatron actor）立即 detach 本地客户端，不把死端点泄漏到client_ttl过期MC_STORE_MEMCPY强制为 0 且 fail-closed（pin 的 mooncake 0.3.10 memcpy 路径已实证静默截断/SIGSEGV）；启动前做只读能力校验（重试 API 存在性、put-before-notify 顺序）默认参数（
simple）与接入前逐行为一致：首次tq.init仍在 Controller 进程内执行、不创建 owner actor、拆除仍是裸tq.close()；唯一新增是 F10 reaper，且它只在遗留 controller 确证半初始化（会导致挂死）时才动作。为什么这么设计
flag 只表达意图，不暴露 Mooncake 内部参数。 endpoint、buffer、segment、timeout、master 策略走内部默认与部署环境（见下方环境变量）。一对一暴露会让配置组合、文档和测试矩阵持续膨胀。
driver 决策一次，且必须在
tq.init之前；生效前再用真实 attach 验证。tq.init会 attach 到已存在的 controller 并忽略传入的 conf，各 worker 各自决策会发散。/sys探测覆盖不了实际调度位置，所以最终以每个存活节点的有界 attach 握手为准——验证的是真实 endpoint，不是能力启发式。AND 归约，退化结果不丢弃。 探测任务超时或崩溃会转成退化结果让归约器降级，而不是被过滤掉、导致高报集群能力。
首期单任务独占集群。 不做多 job 并发、端口租约、master 共享。清理只动本作业拥有的资源：零
pkill；master 完全不碰（auto_init: false）；健康的 controller 保持不动；attach 到他人 controller 的会话绝不替它做全局回退或拆除。GDR 明确定性为实验性。 可用性无法在启动时探测（探测进程无 CUDA context），真实判定在各 worker 的 TQ 客户端内部，不满足则回退 host RDMA 并打 WARNING；
--tq-rdma-mode=required只覆盖传输层。评审响应
MC_STORE_MEMCPY=1绕过（tq_correctness.py:179）MC_MASTER_ADDRESS必填，无 localhost 回退detach_tq_client()挂到全部 teardown 钩子；强杀场景退回 master TTLoff强制 SimpleStorage（controller.py:241 + Copilot）simple恢复 upstream 等价路径（见上）；mooncake+off= Mooncake/TCPRELAX_TQ_ATTACH_TIMEOUT_SECONDS）+ 全节点握手 + driver 汇总统一回退tq.init超时后 daemon thread 仍可能在复用的 Ray worker 中运行max_calls=1、max_retries=0，并补测试锁定一次性 worker isolationRELAX_TQ_GLOBAL_SEGMENT_SIZE_GB可调，容量校验与客户端配置读同一值环境变量
MC_MASTER_ADDRESShost:port，每个节点都要设置MC_STORE_MEMCPYMC_TCP_ENABLE_CONNECTION_POOL1，复用 TCP connection，避免长矩阵运行耗尽临时端口；不影响 RDMA pathRELAX_TQ_ATTACH_TIMEOUT_SECONDSRELAX_TQ_GLOBAL_SEGMENT_SIZE_GB改动清单（分组）
relax/utils/arguments.pyrelax/core/controller.py_resolve_tq_backend/ 默认路径短路 / attach 握手汇总 / 失败清理relax/utils/rdma_probe.pyrelax/utils/tq_config.pyrelax/utils/tq_correctness.pyrelax/utils/tq_lifecycle.pyrelax/components/*、relax/distributed/ray/rollout.py、relax/backends/megatron/actor.pytests/utils/test_rdma_probe.py、test_tq_failure_paths.py、test_tq_dataplane_behavior.pydocs/draft/transfer_queue_rdma.md、scripts/benchmarks/tq_cross_node_bench.py测试
cfe6153014a554e28c5e40a9002d054f44fee7b3；Lint、Pre-commit 和 Python 3.10/3.11/3.12 tests 全部通过MC_TCP_ENABLE_CONNECTION_POOL=1--fully-async、fully_async=True、hybrid=False58054a33834aadbcf76aacd6b1e32e25c030f2c9（package version0.1.10.dev0）；Mooncake package version0.3.10.post2