feat(data-plane): enable RDMA transport for TransferQueue - #256
overloadedHenry wants to merge 45 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
|
@rai-studio-bot 根据 https://github.com/redai-studio/Relax/actions/runs/35525531332/attempts/1?pr=256 分析下 CI 超时原因，现在看起来是稳定 hang，这些 case 预期都是 6min 结束 看看是不是依赖 redai-studio/TransferQueue#5 |
rai-studio-bot
left a comment
rai-studio-bot
left a comment
There was a problem hiding this comment.
回复 @SigureMo 两次运行都停在 test_clear_partition_and_reinit_are_isolated（93%），不是整套测试正常跑满一小时。第一次运行最后一个完成的 case 在 17:26:59 UTC，直到 18:20:24 被取消；第二次归档日志停点相同，且两次都只有 pytest.log、没有完成的 JUnit 报告。
最可能原因是重复 init/close 泄漏 SimpleStorage placement group，耗尽测试集群的 8 个逻辑 CPU。 奖励测试在收集时启动 8 CPU Ray 集群，本文件复用它；每次 TQ 初始化留下一个 1 CPU group，第七个 case 的第二次初始化需要「8 个 group + 1 CPU controller」，恰好越界，进入无超时的 PG-ready 等待。具体调用链及修复建议已放在本轮行级评论。它与之前的多智能体行数容量问题是两个独立问题。
关于 TransferQueue #5：已经合并，merge SHA 就是本 PR pin 的 9784ad0，但该版本仍未回收上述 placement group，仅升级到 #5 不能消除这条耗尽路径。此外，本次 CI 使用现成 relaxrl:latest 镜像，只安装 requirements.txt，不会构建 PR 中的 Dockerfile，也未输出已安装 TQ 的 VCS SHA；因此不能把源码 pin 当成该次 CI 的实际依赖版本。
建议先在 CI 镜像里固定 ray.init(num_cpus=8) 后运行本测试文件，并加 -o faulthandler_timeout=60 与进程级超时，核验是否停在 transfer_queue/utils/common.py 的 ray.get(placement_group.ready())；同时记录 ray.available_resources()、placement-group 状态和 TQ direct_url.json。修复应回收 fixture 自己创建的 group，避免仅增加 CPU 或延长 CI 时限来掩盖泄漏。
现有归档没有 Python/Ray 线程栈，本地缺 Ray，以上是与两次停点吻合的源码定位，尚非运行时栈确认。本轮为 CI 专项答复，既有 P1 审查结论不变。
| def reinit(capacity: int = 1024): | ||
| tq.close() | ||
| if not _wait_controller_gone(): | ||
| _force_kill_controller() | ||
| assert _wait_controller_gone() |
There was a problem hiding this comment.
这个 fixture 在同一个 Ray 集群中反复 tq.close()/init()，但 TQ 的 SimpleStorage bootstrap 每次会创建一个 1 CPU placement group，close() 只杀 storage actor，不移除该 group。当前 pin 的 bootstrap 与 close 仍保留这一组合。
全套测试收集时，奖励测试先 ray.init(num_cpus=8)，这里的 module fixture 会复用它。前六个 case 累积六个 group，第七个 test_clear_partition_and_reinit_are_isolated 内再次 reinit 时，总共需要第八个 group + controller 的 1 CPU，超出集群容量；TQ 在无 timeout 的 ray.get(placement_group.ready()) 等待。两次 CI 的 pytest artifact 都恰好停在这个 case，与该资源耗尽路径吻合（现有日志没有线程栈，尚未在本地 Ray 环境复现）。
请显式回收本 fixture 创建的 placement groups，或让各次生命周期使用独立集群/进程；同时给真实 Ray 测试加外层 deadline 和线程栈采集。仅等待 named controller 消失不能释放 group 预留的 CPU。
There was a problem hiding this comment.
已复核 7b10734b：fixture 在每次初始化后记录新增 placement group（包括异常路径），在下一次初始化前及 teardown 中移除并等待资源释放，切断了此前重复初始化累积 CPU 预留的路径。新增 watchdog 覆盖初始化、测试和清理，并保存超时线程栈。此项代码问题已修复；本地缺少 Ray/Torch，仍需本轮 CI 确认真正消除了挂起。
There was a problem hiding this comment.
补充实际 CI 证据：7b10734b 的 H20 单测 已在 7 分 42 秒完成 pytest，原挂起用例 test_clear_partition_and_reinit_are_isolated 通过。结果为 2895 passed、18 skipped、1 failed；唯一失败是新容量用例比较 jagged/dense tensor 的整体 shape，并非再次挂起。这为本线程的资源回收修复补上了运行验证。
SimpleStorage caps physical rows, while agent/tool fan-out writes more rows than the logical rollout batch holds identities, so a rollout-batch-derived cap rejected the first write. The off path and both auto fallbacks now build it unbounded, the vestigial capacity parameter is gone, and fan-out regression tests cover the controller path and real SimpleStorage.
The SimpleStorage bootstrap reserves a 1-CPU placement group per tq.init and removes it only on the rollback path, so one bundle per case survived close() and starved the 8-CPU cluster the reward tests leave behind; the next init then blocked forever in the unbounded placement-group ready wait. The fixture now hands its bundles back and an autouse deadline dumps every thread stack to a file in the CI artifact directory before failing, because ray.get cannot be interrupted from Python and pytest drops captured stderr on a hard exit.
The en/zh guides and the fully-async draft still showed the controller-side total_storage_size computation that the unbounded SimpleStorage change removed, so they now describe the current contract: no row cap on the default path, a byte budget for Mooncake, and max_staleness + 1 in-flight batches for capacity planning. The duplicated fan-out rationale in the source collapses to the canonical statement in build_simple_storage_config plus one local line per call site.
# ♻️ Refactor - Leave fallback construction and capacity-error logging to Controller. - Remove unused builder logging and simplify lifecycle comments. --- # 📝 Documentation - Distinguish unbounded row storage from memory capacity planning. - Clarify physical-row fan-out estimates in both guides and the draft. --- # ✅ Tests - Cover capacity success and failure using the real backend builder. - Verify auto constructs one fallback and required fails without one. - Remove duplicate bounded-capacity coverage and check unbounded data round trips.
# ✅ Tests ## Cover cluster lifecycle with watchdog deadlines - Extract the Ray wait deadline into a reusable context manager. - Apply separate deadlines to cluster initialization and shutdown. - Attempt cluster shutdown even when initialization fails. ## Preserve cleanup after initialization and teardown errors - Track placement groups allocated before a failed TQ initialization. - Reap placement groups even when client or controller cleanup raises.
# 🐛 Bug Fix ## Accept nested tensor storage results - Normalize dense input and nested output into rows before comparing. - Check row counts and ordered dtype, shape, and byte digests. - Avoid comparing incompatible dense and jagged batch shapes.
rai-studio-bot
left a comment
rai-studio-bot
left a comment
There was a problem hiding this comment.
已复核 3068efb5 的逐行比较修复：它兼容 TQ 返回的 jagged tensor，同时保留行数、顺序、每行 dtype/shape 和字节摘要校验，未发现新增问题。上一版 CI 已实际通过原挂起用例，证据补在原线程。
当前 head 的 CI 尚未完成；本地缺少 torch/Ray/TransferQueue，未重跑真实数据面测试。此前 PR 描述同步建议仍待处理。
|
测试已过，剩余最新代码的双节点测试和 PR 内容整理。 |
rai-studio-bot
left a comment
rai-studio-bot
left a comment
There was a problem hiding this comment.
3068efb5 的 7 项 CI 已全部通过，包括 H20 单测、两项 GPU 训练、Python 3.10/3.11/3.12 和 Pre-commit；此前挂起与后续形状断言失败已有通过结果。
原描述同步建议 F1 已部分解决：正文已同步 mode/device、host-RDMA 和 SimpleStorage 回退。剩余“改动清单”仍写四个 flag、rdma_probe.py、旧 tq_*.py 路径及 #278 运行时补丁；建议改为两个参数、当前 relax/utils/tq/{config,correctness,lifecycle}.py，删除已移除的探测文件和补丁说明，并将测试段更新为本 head 的 CI 结果。以上仍属非阻塞描述问题；CI 通过不等同于本 head 已完成专项多节点 RDMA 验收。
rai-studio-bot
left a comment
rai-studio-bot
left a comment
There was a problem hiding this comment.
已复核 0c08c195 的全部增量，未发现新增问题。benchmark 在计算期望摘要前将 sample_id / rewards 统一为 [N, 1]，与 TQ 对一维列可能返回二维结果的契约一致；逐行 dtype、shape、字节及顺序校验仍保留，新增测试覆盖单行、多行和三类不匹配。
本地目标测试因缺少 Ray 在收集阶段失败；当前 CI 的 Pre-commit 已通过，其余运行中或排队，未执行多节点 RDMA benchmark。既有 F1/F2/F3 状态不变。
# 🐛 Bug Fix - Align sample_id and rewards with one-element TQ row shapes before computing strict dtype, shape, and byte digests. - Flush measurement records before partition cleanup and append final records with explicit cleanup status. Keep successful measurements pending until cleanup completes. - Preserve primary failures and associated cleanup and CSV write errors. --- # ✅ Tests - Reconstruct single-sample and multi-sample payloads from raw bytes using actual TransferQueue metadata and retain strict corruption checks. - Cover measurement flushing, transfer and gate failures, cleanup errors, and final CSV write failures. - Validate 42 CPU tests and all pre-commit checks. - Hardware RDMA validation was not performed. --- # 📝 Documentation - Document measurement/final CSV phases and final-result filtering.
0c08c19 to
766797a
Compare
|
双节点 RDMA 小规模测试成功，Qwen3-vl-4B + 4 rollout step + weight update，无 fallback。 |
rai-studio-bot
left a comment
rai-studio-bot
left a comment
There was a problem hiding this comment.
已复核 766797ab 的全部增量，未发现新增问题。CSV 先 flush 测量记录、清理后再追加最终状态，避免清理失败或挂起时丢失测量证据或误记成功；文档已说明仅统计 phase=final。新增标量重建测试使用的 schema 接口与固定版本 TQ 一致。
提取实际函数及相关测试、隔离 Ray 依赖后，10 个结果记录和异常处理用例通过；这不等同于完整测试或真实数据面验证。本地缺少 Ray/Torch/TQ，未运行 schema 张量测试及多节点 RDMA benchmark；Pre-commit 已通过，其余 CI 尚在运行。既有问题状态不变。
# ✅ Tests - Keep scalar shape, dtype, and strict digest checks active in stub CI. - Run metadata reconstruction checks with the real TransferQueue package; explicitly skip these checks for a non-package module stub. - Inspect the import spec without triggering fabricated stub attributes, and let real package import errors fail normally. - Validate focused CPU tests: 42 passed and 2 skipped with the CI stub; 44 passed with the real TransferQueue package. - Pass all pre-commit checks. Real metadata validation remains unavailable in stub CI; hardware RDMA validation was not performed.
rai-studio-bot
left a comment
rai-studio-bot
left a comment
There was a problem hiding this comment.
已复核 ffb88aef，未发现新增问题。上一版 Python CI 的失败确为模块 stub 导入 transfer_queue.metadata 时触发 TypeError: '_Stub' object is not iterable。本次保留基础形状/摘要测试，并仅在模块 stub 环境跳过真实 metadata 测试，修复与原因一致。
本地隔离验证确认 __spec__ 判断可区分普通模块和真实包，且不触发 __getattr__；缺少 Ray/Torch/TQ，未运行完整张量测试或多节点 RDMA 验收。Pre-commit 已通过，其余当前 CI 仍在运行。既有问题状态不变。
|
@rai-studio-bot 看看本 PR 新增单测时间组成，是新增单测 case 导致的吗？ |
rai-studio-bot
left a comment
rai-studio-bot
left a comment
There was a problem hiding this comment.
回复 @SigureMo：新增 case 有可测开销，但本次 H20 单测 job 接近两小时，主要耗在镜像依赖下载，不是执行新增测试。 以下按 ffb88aef 的这次 H20 run分析。
| 阶段 | 耗时 |
|---|---|
| 整个 job | 1 小时 58 分 31 秒 |
| Prepare local test image | 1 小时 50 分 02 秒 |
| Run unittest step | 7 分 57 秒 |
| pytest 自身（JUnit suite） | 469.68 秒 |
镜像构建的依赖安装层耗时 6519.5 秒；其中 nvidia_cudnn_cu12-9.17.1.4 的 647.8 MB wheel 下载日志跨越 5088.2 秒（84 分 48 秒），Ray wheel 下载另耗约 6 分钟。这些发生在 pytest 启动前。
按该 run 的 junit.xml，本 PR 新增测试文件加既有 SFT 文件中的新增测试，共 205 个参数化 case 条目，合计 77.95 秒，约占 pytest 总时长 16.6%：
| 新增测试范围 | 条目数 | JUnit 累计耗时 |
|---|---|---|
test_tq_dataplane_behavior.py |
9 | 53.06 秒 |
test_tq_failure_paths.py |
47 | 13.54 秒 |
test_train_actor_init_cleanup.py |
5 | 11.14 秒 |
| 其余配置、Controller、benchmark guard、摘要、FSDP/SFT 测试 | 144 | 0.21 秒 |
主要慢点是实际生命周期验证：数据面用例反复启动/关闭 TQ 和回收资源，普通 case 约 5 秒，含两次初始化的 test_clear_partition_and_reinit_are_isolated 为 10.20 秒；attach 超时进程隔离验证为 13.06 秒，actor 初始化失败清理验证为 11.14 秒。后两者各自启动独立 Ray 子进程，还显式等待约 2.2 秒确认没有延迟状态修改。新增 benchmark guard 39 个条目合计仅 0.076 秒。
这里 JUnit case 时间包含 setup/call/teardown，无法从现有 artifact 精确拆开；77.95 秒是当前运行中归到新增 case 的时间，不是同机器 main/head 对照测出的净增量。全套 case 时间合计 430.90 秒，另约 38.78 秒未归入 case（收集、会话等开销），不能全部归因于 PR。Python 3.10/3.11/3.12 的 pytest 总时长分别为 443.61/308.99/387.39 秒，这组 CPU CI未提供同样的 JUnit 分项。
建议先优化镜像复用与依赖下载；若要继续压缩单测时间，优先给上述真实 Ray/TQ 测试加 --durations=30 --durations-min=0.1，测清 setup/call/teardown 后再调整 fixture。不要为了省掉约一分钟而撤掉进程隔离和清理回归覆盖。当前 head 的 7 项 CI 均已通过。
|
@overloadedHenry 这个 PR 是不是已经完成了，那我推进下合入 |
目测还算可以接受 |
目前该 PR 内容已完成。后续开发如完善 GDR 等，若有计划，我可以继续贡献代码。 |
rai-studio-bot
left a comment
rai-studio-bot
left a comment
There was a problem hiding this comment.
已复核 a439afa7 合并 main 的结果，未发现新增问题。合并前后的 PR 差异具有相同 stable patch-id；两个交叉 actor 文件中的设备 API 迁移与 TQ attach/detach、初始化失败清理逻辑兼容，未发现 TQ 改动丢失。
本地配置测试 64 项通过。Pre-commit 已通过，其余当前 CI 运行中或排队；此前 head 的通过结果不作为本 head 的运行验证。本轮未运行真实 Ray/GPU/NPU 或多节点 RDMA 验收，既有问题状态不变。
rai-studio-bot
left a comment
rai-studio-bot
left a comment
There was a problem hiding this comment.
已复核 9fa347b0 合并 main 的结果，PR 自身 stable patch-id 未变，未发现新增代码问题。交叉文件 Dockerfile.npu 的主分支改动位于 DeepEP 构建参数，TQ pin 与接入逻辑保留。
NPU CI 在 Prepare local test image 的最初 apt 步骤失败：Could not resolve 'ci-proxy'，最终退出码 100。尚未执行 TQ 安装或训练，现有证据不能归因为 RDMA 回归。请先确保 Docker build 内可以解析并访问配置的代理，再重跑 NPU 验证；运行容器的 /etc/hosts 挂载不能替代构建阶段的解析配置。
Pre-commit 已通过，其余当前 CI 尚在运行；本轮未重复运行未变更的配置测试，也未执行真机 RDMA 验收。既有问题状态不变。
【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。本 PR 将 TransferQueue 的 MooncakeStore/host-RDMA 接入 Relax，用于 rollout → train 的样本传输。改动范围是配置接入、有界初始化与 attach、统一回退和资源清理，不实现传输层，也不改变 payload 形状或数据分发语义。
首期不支持 GDR。生产路径仅包含 host-RDMA 和 SimpleStorage；Mooncake/TCP 仅作为跨节点 benchmark 的对照模式。
改了什么
一条默认关闭、初始化与 attach 有界的 host-RDMA 数据路径：
--tq-rdma-mode=off|auto|required--tq-rdma-device=<device>off为默认值，使用 SimpleStorage，不创建 Mooncake owner actor，不执行 Mooncake 检查或集群 handshake；worker attach 仍有超时边界。
auto尝试 host-RDMA；任一检查或节点 attach 失败时，在确认资源清理完成后统一回退 SimpleStorage。
required要求 host-RDMA；不可用时清理资源并终止启动，不静默降级。节点运行一次性 attach worker，确认实际 manager 与
protocol=rdma。max_calls=1、max_retries=0，避免超时初始化线程污染后续复用任务。
/sys启发式能力探测；设备未指定时交由 Mooncake 原生逻辑选择，多 HCA 环境建议显式指定设备。
均不把逻辑 batch 配额作为物理存储行数上限。
required 拒绝启动。
MOONCAKE_CORRECTNESS_CONTRACT_VERSION >= 1，并对不安全的 memcpy 配置 fail closed；不使用运行时 monkey patch。
Mooncake master 始终由部署环境管理。
默认 --tq-rdma-mode=off 保持 SimpleStorage 后端，首次初始化仍在 Controller 进程执行，不创建 Mooncake owner actor，也不执行 Mooncake handshake。
默认路径并非完全没有行为变化：worker attach 增加超时边界，启动时回收不可用的半初始化 controller，并在首期单任务独占集群约束下拒绝复用健康的既有 controller。
为什么这么设计
只保留 mode/device 两个参数；master、segment 和 timeout 使用内部默认值及部署环境配置。 endpoint、buffer、segment、timeout、master 策略走内部默认与部署环境（见下方环境变量）。一对一暴露会让配置组合、文档和测试矩阵持续膨胀。
driver 决策一次，且必须在
tq.init之前；生效前再用真实 attach 验证。tq.init会 attach 到已存在的 controller 并忽略传入的 conf，各 worker 各自决策会发散。/sys探测覆盖不了实际调度位置，所以最终以每个存活节点的有界 attach 握手为准——验证的是真实 endpoint，不是能力启发式。以每个 ALIVE 节点的真实、有界 attach 结果决定是否启用 host-RDMA，不维护独立的启发式探测矩阵。 任一节点的 attach 失败或超时，均汇总到 driver；确认清理完成后，
auto统一回退 SimpleStorage，required终止启动。清理无法确认时 fail closed。Attach 成功只证明 manager、配置 protocol 和 setup 满足要求；实际线路是否传输 RDMA 数据，仍需跨节点 benchmark 的 wire proof 验证。首期单任务独占集群。 不做多 job 并发、端口租约、master 共享。清理只动本作业拥有的资源：零
pkill；master 完全不碰（auto_init: false）；健康的 controller 保持不动；首期要求单任务独占 Ray 集群；遇到健康的既有 controller 直接拒绝启动，不接管或关闭它。首期不支持 GDR，固定 use_gdr=false；GDR 如有需求，另行实现和专项验证。
环境变量
MC_MASTER_ADDRESSMC_STORE_MEMCPY0；显式设置不安全值会被启动拒绝。只有在能够可靠识别已修复 build 并完成验证后，才重新评估是否允许开启 memcpy 路径。MC_TCP_ENABLE_CONNECTION_POOLRELAX_TQ_ATTACH_TIMEOUT_SECONDSRELAX_TQ_GLOBAL_SEGMENT_SIZE_GB改动清单（分组）
relax/utils/arguments.py--tq-rdma-mode、--tq-rdma-devicerelax/core/controller.pyrelax/utils/tq/config.pyrelax/utils/tq/correctness.pyrelax/utils/tq/lifecycle.pyrelax/components/、relax/distributed/ray/rollout.py、relax/backends/tests/utils/tq/、tests/core/test_controller_tq_backend.py、tests/utils/test_tq_failure_paths.py、tests/utils/test_tq_dataplane_behavior.pydocs/draft/transfer_queue_rdma.md、scripts/benchmarks/tq_cross_node_bench.py测试
功能测试均在验收前进行完毕，目前的代码根据规范要求修复了冗余，merge 了 main 分支里的更新代码，等待最新测试。