From 0fcf9014879b5288bb8a71c644e7815e24933b0f Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Tue, 11 Aug 2026 12:20:02 +0800 Subject: [PATCH 01/22] feat(data-plane): enable RDMA transport for TransferQueue 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%. --- docs/draft/transfer_queue_rdma.md | 130 +++++ relax/core/controller.py | 126 ++++- relax/utils/arguments.py | 55 +++ relax/utils/rdma_probe.py | 521 ++++++++++++++++++++ relax/utils/tq_config.py | 185 +++++++ relax/utils/tq_lifecycle.py | 114 +++++ scripts/benchmarks/cross_node_rdma_bench.py | 226 +++++++++ scripts/benchmarks/tq_cross_node_bench.py | 490 ++++++++++++++++++ scripts/benchmarks/tq_rdma_bench.py | 396 +++++++++++++++ tests/utils/test_rdma_probe.py | 338 +++++++++++++ tests/utils/test_tq_dataplane_behavior.py | 296 +++++++++++ tests/utils/test_tq_failure_paths.py | 374 ++++++++++++++ 12 files changed, 3244 insertions(+), 7 deletions(-) create mode 100644 docs/draft/transfer_queue_rdma.md create mode 100644 relax/utils/rdma_probe.py create mode 100644 relax/utils/tq_config.py create mode 100644 relax/utils/tq_lifecycle.py create mode 100644 scripts/benchmarks/cross_node_rdma_bench.py create mode 100644 scripts/benchmarks/tq_cross_node_bench.py create mode 100644 scripts/benchmarks/tq_rdma_bench.py create mode 100644 tests/utils/test_rdma_probe.py create mode 100644 tests/utils/test_tq_dataplane_behavior.py create mode 100644 tests/utils/test_tq_failure_paths.py diff --git a/docs/draft/transfer_queue_rdma.md b/docs/draft/transfer_queue_rdma.md new file mode 100644 index 000000000..a94895dcd --- /dev/null +++ b/docs/draft/transfer_queue_rdma.md @@ -0,0 +1,130 @@ +# TransferQueue RDMA 数据面使用与运维指南 + +## 概述 + +Relax 的数据面（rollout ↔ train 之间的样本传输）默认走 TransferQueue 的 SimpleStorage/ZMQ。本特性把 TransferQueue 已有的 MooncakeStore 后端接出来，使数据面可以走 RDMA，并在能力不足时安全回退。 + +首期只做配置接入、能力探测与一致回退，**不改变 payload 形状与数据分发语义**。默认参数下行为与接入前完全一致。 + +## 配置入口 + +只暴露四个表达使用意图的参数，Mooncake 底层参数（endpoint、buffer、segment、timeout、master 策略）不做 CLI，走内部默认与部署环境。 + +| 参数 | 取值 | 说明 | +|---|---|---| +| `--tq-storage-backend` | `simple`（默认）/ `mooncake` | `simple` 等价于接入前行为 | +| `--tq-rdma-mode` | `off`（默认）/ `auto` / `required` | `off` 即使有硬件也不用 RDMA；`auto` 探测失败自动降级；`required` 探测失败直接报错退出 | +| `--tq-rdma-device` | 设备名，如 `mlx5_bond_0`；空为自动 | 多网卡机器上自动选择可能选错，跨节点时建议显式指定 | +| `--tq-use-gdr` | 默认关 | **实验性**，见下文 | + +`--tq-rdma-mode=required` 只覆盖**传输层**（MooncakeStore + RDMA 可用性与 segment 容量），不覆盖 GDR。 + +### GDR 为实验性 + +GDR 的可用性无法在启动时探测：探测跑在独立的 Ray task 里，该进程没有初始化 CUDA context，`torch.cuda.is_initialized()` 恒为 False，若在此判定会让 GDR 永远不可达。真实判定发生在每个 worker 的 TQ 客户端内部（`mooncake_client.py`），没有 CUDA context 时**静默回退到 host RDMA**并打 WARNING。 + +因此首期：`--tq-use-gdr` 标记为实验性，`required` 不对 GDR 做 fail-fast。要把 GDR 纳入分级降级，需要补一条 worker → driver 的能力回报通道，属于后续工作。 + +## 启动流程与降级 + +driver 在**第一次 `tq.init` 之前**完成探测并生成 job 级唯一的 effective config，其余组件（actor / critic / rollout / sft / advantages / actor_fwd）都读同一份，不各自决策。 + +1. 校验参数组合（例如 `simple` + `rdma-mode` 会被拒绝） +2. `probe_cluster_nodes()` 通过 Ray 把探测任务绑定到每个**存活且有 GPU** 的节点，各自读本机 `/sys` 与 mooncake 状态；超时或崩溃的节点转为退化结果，不静默丢弃 +3. `reduce_results()` 做 AND 归约：整个作业只能跑在最低共同能力上 +4. `required` 模式下若发生任何回退，直接抛异常并打印每个节点的探测明细 + +降级阶梯： + +``` +GDR → host RDMA （worker 运行时判定，静默回退 + WARNING） +RDMA → Mooncake/TCP （任一节点无 RDMA 能力，或指定设备缺失） +Mooncake → SimpleStorage（任一节点 mooncake 不可导入，或 segment 容量不足） +``` + +## 启动日志怎么读 + +正常启动会打三段，排障时先看这三段： + +``` +[dataplane] requested: backend=mooncake rdma_mode=auto device=mlx5_bond_0 gdr=False +[dataplane] probe result: +[probe:] protocol=rdma device=mlx5_bond_0 gdr=True + [ok] mooncake_import: version=0.3.10.post2 + [ok] rdma_devices: mlx5_bond_0, ... + [ok] port_state: ACTIVE + [ok] gid: ... + [ok] memlock: unlimited +[dataplane] backend=MooncakeStore protocol=rdma device=mlx5_bond_0 gdr=off +``` + +第三段带 `fallback=...` 就说明发生了降级，原因直接写在里面（例如 `fallback=mooncake_unavailable:`）。 + +## Mooncake master 生命周期 + +`auto_init` 固定为 `false`：**Relax 既不启动也不停止 master**，master 由部署环境管理。这样做是因为 TQ 的 `auto_init=true` 路径会执行 `pkill -f "[m]ooncake_master"`，在共享集群上会杀掉其他人的进程。 + +启动 master（部署侧，节点上执行一次）： + +```bash +setsid mooncake_master -rpc_port=50051 -metrics_port=9004 > /var/log/mooncake_master.log 2>&1 < /dev/null & +``` + +然后给作业设置 `MC_MASTER_ADDRESS=:50051`。未设置时内部默认 `localhost:50051`，仅适用于单节点开发。 + +三种情形下的行为： + +| 情形 | 表现 | 处理 | +|---|---|---| +| **初始化失败**（master 不可达） | TQ 客户端 `setup` 返回 `-1`，抛 `Mooncake store setup failed with error code: -1`。`required` 模式下作业直接失败；`auto` 模式下这一步已过了探测，不会自动回退——探测只验证本机能力，不验证 master 连通性 | 先确认 master 进程与 `MC_MASTER_ADDRESS`，再重启作业 | +| **正常退出**（作业结束或全局重启） | Relax 在拆数据面时调用 `close_tq_and_unmount()`：先 `tq.close()`（它内部还要用 store 做 `remove_all()`），再显式 `storage_client.close()` 卸载 segment 并从 master 注销。master 本身不动 | 无需操作 | +| **异常退出**（worker 被 kill / OOM / 节点掉线） | Python 层不执行，segment 仍在 master 注册。master 要等 `client_ttl`（默认 30 s）才判定客户端过期，期间新作业的 put 会打到死端点并报 `Failed to open segment ... Connection refused` | 等 30 s 后重启，或部署侧调小 `-client_ttl` | + +## 资源所有权与安全清理 + +首期按**单任务独占 Ray 集群**实现，**不承诺同一节点上多个 Relax job 并发**：多 job 并发、端口租约、master 共享机制都不在首期范围内。 + +清理只动本作业拥有的资源： + +- 不使用任何 `pkill` / `killall` +- `tq.init` 之前会检查已存在的 `TransferQueueController` 命名 actor：**只有取不到 config（半初始化）或 actor 已死时才回收**，健康的 controller 保持不动并正常 attach +- master 进程始终不被 Relax 触碰 + +## 排障表 + +| 现象 | 可能原因 | 处理 | +|---|---|---| +| 启动日志 `backend=SimpleStorage fallback=mooncake_unavailable:` | 该节点上 `import mooncake` 失败 | 检查该节点的 `mooncake-transfer-engine` 安装；镜像是否一致 | +| `protocol=tcp fallback=...`，但机器有 RDMA 卡 | 端口非 ACTIVE、GID 取不到、`memlock` 过低，或指定的 `--tq-rdma-device` 在部分节点不存在 | 看 `probe result` 里哪一项 FAIL；`memlock` 需要 unlimited | +| `setup failed with error code: -1` | master 不可达 | 检查 master 进程与 `MC_MASTER_ADDRESS` | +| `Failed to open segment ... Connection refused` | 上一轮客户端异常退出，死 segment 仍在 master 注册 | 等 `client_ttl`（30 s）过期后重试 | +| `batch_get_into failed ... error codes [-800, ...]` | 会话内切换协议（0.3.10 上更敏感），或对端不可达 | 每个协议单独进程跑；确认对端存活 | +| 多网卡机器跨节点建连失败 | 自动选卡选到了不通的网卡 | 显式 `--tq-rdma-device`；必要时用 `MC_TCP_BIND_ADDRESS` 指定 TCP 侧绑定地址 | +| 训练卡在启动、无日志推进 | 半初始化的 controller（TQ 的 `_init_from_existing` 会无限轮询 config） | 本特性已加自动回收；若仍出现，确认 `[dataplane] ... reaping it` 是否打出 | + +### 选卡核验 + +指定设备前先确认端口状态与 GID： + +```bash +ls /sys/class/infiniband/ # 有哪些设备 +cat /sys/class/infiniband/mlx5_bond_0/ports/1/state # 需要 ACTIVE +cat /sys/class/infiniband/mlx5_bond_0/ports/1/rate +ulimit -l # 需要 unlimited +``` + +判定数据面是否真的走了 RDMA（get 前后取差值）： + +```bash +cat /sys/class/infiniband/mlx5_bond_0/ports/1/counters/port_rcv_data # RDMA，单位是 4 字节字 +cat /sys/class/net/bond0/statistics/rx_bytes # TCP +``` + +RDMA 生效时前者按 payload 增长、后者基本不动；反之则说明落在 TCP。 + +## 已知限制 + +- 写侧（put）收益明显，读侧（get）收益有限：get 每次调用都会注册/注销 MR，且 key 粒度是 `样本 × 字段`，碎片化开销盖过了传输收益。MR 常驻注册与读路径零拷贝成型不在首期范围。 +- 跨节点 RDMA vs TCP 的收益在多轮之间波动较大，验收结论应基于多轮分布而非单轮数据。 +- 消费端节点在 get 过程中中途死亡的端到端行为需要双节点真机验证，未做成自动化测试。 +- Mooncake 传输层自身的超时参数不由 Relax 控制。 diff --git a/relax/core/controller.py b/relax/core/controller.py index e59f5cfd7..f064184af 100644 --- a/relax/core/controller.py +++ b/relax/core/controller.py @@ -47,7 +47,10 @@ set_managed_opd_teacher_on_actor_service, shutdown_managed_opd_teacher, ) +from relax.utils.rdma_probe import probe_cluster_nodes, reduce_results, validate_config from relax.utils.s3_model_loader import cleanup_s3_model_weights_from_shm +from relax.utils.tq_config import build_backend_config +from relax.utils.tq_lifecycle import close_tq_and_unmount, kill_tq_controller_and_wait, reap_unusable_tq_controller from relax.utils.training.ppo_utils import validate_ppo_config from relax.utils.utils import compute_dp_size, recovery_load_path @@ -277,18 +280,127 @@ def _initialize_data_system(self): "sampler": sampler, "polling_mode": self.config.polling_mode, }, - "backend": { - "SimpleStorage": { - "total_storage_size": total_storage_size, - "num_data_storage_units": self.config.num_data_storage_units, - }, - }, + "backend": self._resolve_tq_backend(total_storage_size), }, flags={"allow_objects": True}, ) tq_config = tq.init(conf=tq_config) or tq_config self.config.tq_config = tq_config + def _resolve_tq_backend(self, total_storage_size: int) -> dict: + """Resolve the TransferQueue ``backend`` config dict. + + Default behavior (``--tq-storage-backend=simple``) is identical to the + previous hardcoded SimpleStorage path. When MooncakeStore is + requested, runs the RDMA capability probe *before* ``tq.init``, applies + graded degradation, and emits the startup log line. + """ + # 1. Validate flag combinations (structural, before any probe). + # getattr defaults keep old checkpoints / non-argparse configs safe. + errors = validate_config(self.config) + if errors: + raise ValueError("Invalid TransferQueue RDMA configuration:\n " + "\n ".join(errors)) + + backend = getattr(self.config, "tq_storage_backend", "simple") + mode = getattr(self.config, "tq_rdma_mode", "off") + + # 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, + ) + + # 3. MooncakeStore path: probe → reduce → effective config. + # The driver fans the probe out to every alive GPU node via Ray + # (probe_cluster_nodes), then AND-reduces to a single job-level + # effective config so all data-plane workers converge identically. + device = getattr(self.config, "tq_rdma_device", "") + probe_results = probe_cluster_nodes(device) + for r in probe_results: + logger.debug(r.summary()) + + effective = reduce_results( + probe_results, + requested_backend=backend, + requested_device=device, + use_gdr=getattr(self.config, "tq_use_gdr", False), + ) + + # 4. required mode: fail fast instead of silently degrading (probe level). + if mode == "required" and effective.fallback_reason: + detail = "\n".join(r.summary() for r in probe_results) + raise RuntimeError( + f"--tq-rdma-mode=required but RDMA probe failed: {effective.fallback_reason}.\n" + f"Probe details:\n{detail}" + ) + + # 5. Build backend dict (may fall back to SimpleStorage on capacity error). + backend_dict, cap_error = build_backend_config(self.config, effective, total_storage_size=total_storage_size) + actual_backend = "MooncakeStore" if "MooncakeStore" in backend_dict else "SimpleStorage" + + # 6. required mode: also fail fast on capacity-induced fallback. + if mode == "required" and cap_error: + raise RuntimeError(f"--tq-rdma-mode=required but segment capacity insufficient: {cap_error}") + + # 7. GDR is EXPERIMENTAL in this phase: --tq-rdma-mode=required only + # covers the *transport* (MooncakeStore + RDMA), not GDR. The probe + # cannot decide GDR eligibility -- it runs as a separate Ray task + # where torch.cuda.is_initialized() is always False -- so the real + # check happens per worker in the runtime client + # (mooncake_client.py:87), which silently falls back to host RDMA. + use_gdr = getattr(self.config, "tq_use_gdr", False) + if use_gdr and (effective.protocol != "rdma" or actual_backend != "MooncakeStore"): + logger.warning( + "[dataplane] --tq-use-gdr requested but the effective path is not RDMA " + f"(protocol={effective.protocol}, backend={actual_backend}); GDR inactive." + ) + elif use_gdr: + logger.warning( + "[dataplane] --tq-use-gdr is EXPERIMENTAL: eligibility is not probed, and " + "workers without an initialised CUDA context fall back to host RDMA silently. " + "--tq-rdma-mode=required does NOT fail fast on unavailable GDR." + ) + + # 8. F10 anti-hang: reap a half-initialised controller before tq.init, + # for every backend (a stale MooncakeStore controller used to survive + # because the guard only ran on the SimpleStorage fallback path). + # Only unusable controllers are killed -- a healthy one belongs to + # whoever created it and tq.init legitimately attaches to it. + self._reap_unusable_tq_controller() + + # 9. Log requested vs effective so the startup log alone explains the + # decision, plus one summary block per probed node. + logger.info( + f"[dataplane] requested: backend={backend} rdma_mode={mode} device={device or 'auto'} gdr={use_gdr}" + ) + for result in probe_results: + logger.info(f"[dataplane] probe result:\n{result.summary()}") + if cap_error: + logger.warning(f"[dataplane] MooncakeStore capacity fallback to SimpleStorage: {cap_error}") + logger.info("[dataplane] backend=SimpleStorage protocol=tcp (capacity fallback)") + else: + logger.info(effective.log_line()) + return backend_dict + + def _reap_unusable_tq_controller(self) -> None: + """Delegate to + :func:`relax.utils.tq_lifecycle.reap_unusable_tq_controller`.""" + reap_unusable_tq_controller() + + def _kill_stale_tq_controller(self) -> None: + """Delegate to + :func:`relax.utils.tq_lifecycle.kill_tq_controller_and_wait`.""" + kill_tq_controller_and_wait() + + @staticmethod + def _close_data_system() -> None: + """Delegate to + :func:`relax.utils.tq_lifecycle.close_tq_and_unmount`.""" + close_tq_and_unmount() + def _deploy_metrics_service(self): """Deploy the MetricsService as a lightweight Ray Serve deployment. @@ -1006,7 +1118,7 @@ def _global_restart(self) -> None: # --- 1.7 Tear down data system (storage units + controller) --- try: - tq.close() + self._close_data_system() except Exception as e: logger.warning(f"[Global Restart] Failed to tear down data system: {e}") diff --git a/relax/utils/arguments.py b/relax/utils/arguments.py index 3bf40c090..33f142491 100644 --- a/relax/utils/arguments.py +++ b/relax/utils/arguments.py @@ -244,6 +244,61 @@ def add_transfer_queue_arguments(parser): default=1, help="Fully async pipeline num of iters every global batch.", ) + # ── RDMA transport (MooncakeStore backend) ────────────────────────── + # Default values are equivalent to current main behavior (SimpleStorage, + # no RDMA). These expose *intent* only; Mooncake internals (endpoint, + # buffer, segment, timeout, master strategy) are handled via internal + # defaults / deployment environment, not CLI flags. + parser.add_argument( + "--tq-storage-backend", + choices=["simple", "mooncake"], + default="simple", + help=( + "TransferQueue storage backend. 'simple' (default) uses " + "SimpleStorage/ZMQ and is equivalent to current behavior. " + "'mooncake' uses MooncakeStore, which supports RDMA transport " + "via --tq-rdma-mode." + ), + ) + parser.add_argument( + "--tq-rdma-mode", + choices=["off", "auto", "required"], + default="off", + help=( + "RDMA transport mode for MooncakeStore backend. 'off' (default) " + "never uses RDMA even if hardware is available. 'auto' probes " + "RDMA capability at startup and degrades to TCP/SimpleStorage if " + "unavailable (with a WARNING). 'required' fails fast on probe " + "failure instead of degrading; it covers the transport only, not " + "GDR (see --tq-use-gdr). Only effective with " + "--tq-storage-backend mooncake." + ), + ) + parser.add_argument( + "--tq-rdma-device", + type=str, + default="", + help=( + "RDMA device name for MooncakeStore (e.g. mlx5_bond_0). Empty " + "(default) lets Mooncake auto-select. On multi-NIC hosts the " + "auto-selected device may fail cross-node; specify explicitly if " + "needed." + ), + ) + parser.add_argument( + "--tq-use-gdr", + action=argparse.BooleanOptionalAction, + default=False, + help=( + "EXPERIMENTAL: enable GPU Direct RDMA (GDR) staging for " + "MooncakeStore. Eligibility is NOT probed at startup (the probe " + "runs in a separate Ray task with no CUDA context), so each " + "worker decides at runtime: without an initialised CUDA context " + "it falls back to host RDMA with a WARNING, and " + "--tq-rdma-mode=required does not fail fast on that. Requires " + "RDMA protocol. Default off." + ), + ) return parser def add_cluster_arguments(parser): diff --git a/relax/utils/rdma_probe.py b/relax/utils/rdma_probe.py new file mode 100644 index 000000000..aa3e296ac --- /dev/null +++ b/relax/utils/rdma_probe.py @@ -0,0 +1,521 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""RDMA capability probe and graded-degradation state machine for +MooncakeStore. + +This module runs **before** ``tq.init`` to decide the job-level effective +``{backend, protocol, device}`` triple. The probe is intentionally side-effect +free: it only reads ``/sys``/``resource``/mooncake introspection and performs a +short handshake. The result is AND-reduced across all data-plane nodes by the +driver so that every worker converges on the *same* effective config. + +Key constraint (F10 in the RFC): probing must happen *before* ``tq.init``. +If the named actor ``TransferQueueController`` is created before the probe +succeeds, a subsequent ``tq.init`` retry via ``_init_from_existing`` will spin +in an unbounded ``while conf is None`` loop (``interface.py:109-118``) and hang +the job with no error message. +""" + +from __future__ import annotations + +import os +import resource +import socket +from dataclasses import dataclass +from typing import Any + +from relax.utils.logging_utils import get_logger + + +logger = get_logger(__name__) + +# --------------------------------------------------------------------------- +# Data structures +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class CheckResult: + """Outcome of a single capability check.""" + + name: str + ok: bool + detail: str = "" + + +@dataclass(frozen=True) +class ProbeResult: + """Aggregated RDMA capability report for a single node. + + ``effective_protocol`` is the highest transport this node can use after + graded degradation: + + * ``"rdma"`` – RDMA device ACTIVE + GID available + mooncake importable. + * ``"tcp"`` – mooncake importable but no usable RDMA device. + * ``None`` – mooncake not importable at all; must fall back to + SimpleStorage. + """ + + node: str + checks: tuple[CheckResult, ...] + effective_protocol: str | None # "rdma" | "tcp" | None + effective_device: str + gdr_eligible: bool + errors: tuple[str, ...] = () + + @property + def ok(self) -> bool: + """True if this node can run MooncakeStore (tcp or rdma).""" + return self.effective_protocol is not None + + def summary(self) -> str: + """Return a multi-line human-readable report of this node's probe.""" + header = ( + f"[probe:{self.node}] protocol={self.effective_protocol} " + f"device={self.effective_device} gdr={self.gdr_eligible}" + ) + lines = [header] + for c in self.checks: + tag = "ok" if c.ok else "FAIL" + lines.append(f" [{tag}] {c.name}: {c.detail}" if c.detail else f" [{tag}] {c.name}") + return "\n".join(lines) + + +@dataclass(frozen=True) +class EffectiveConfig: + """Job-level unique effective config after AND-reduction across nodes.""" + + backend: str # "MooncakeStore" or "SimpleStorage" + protocol: str # "rdma" or "tcp" + device: str + gdr: bool + fallback_reason: str # "" if no fallback occurred + + def log_line(self) -> str: + """Return the single-line startup log string for this effective + config.""" + gdr_str = "on" if self.gdr else "off" + dev = self.device or "auto" + base = f"[dataplane] backend={self.backend} protocol={self.protocol} device={dev} gdr={gdr_str}" + if self.fallback_reason: + return f"{base} fallback={self.fallback_reason}" + return base + + +# --------------------------------------------------------------------------- +# Individual checks (pure read, no side effects beyond mooncake import) +# --------------------------------------------------------------------------- + + +def _check_mooncake_import() -> CheckResult: + try: + import mooncake # noqa: F401 (import-time only) + + ver = getattr(mooncake, "__version__", "unknown") + return CheckResult("mooncake_import", True, f"version={ver}") + except Exception as e: # pragma: no cover - environment dependent + return CheckResult("mooncake_import", False, str(e)) + + +def _check_rdma_devices() -> CheckResult: + base = "/sys/class/infiniband" + if not os.path.isdir(base): + return CheckResult("rdma_devices", False, "no /sys/class/infiniband") + devs = sorted(os.listdir(base)) + if not devs: + return CheckResult("rdma_devices", False, "empty /sys/class/infiniband") + return CheckResult("rdma_devices", True, ",".join(devs)) + + +def _check_port_active(device: str = "", port: int = 1) -> CheckResult: + if device: + state_path = f"/sys/class/infiniband/{device}/ports/{port}/state" + try: + with open(state_path) as f: + state = f.read().strip() + ok = "ACTIVE" in state + return CheckResult(f"port_active:{device}/{port}", ok, state) + except FileNotFoundError: + return CheckResult(f"port_active:{device}/{port}", False, "state file missing") + except OSError as e: + return CheckResult(f"port_active:{device}/{port}", False, str(e)) + # No device specified: check the first available one. + base = "/sys/class/infiniband" + if not os.path.isdir(base): + return CheckResult("port_active", False, "no infiniband dir") + for dev in sorted(os.listdir(base)): + return _check_port_active(dev, port) + return CheckResult("port_active", False, "no devices") + + +def _check_gid_available(device: str = "", gid_index: int = 3) -> CheckResult: + if device: + gid_path = f"/sys/class/infiniband/{device}/ports/1/gids/{gid_index}" + try: + with open(gid_path) as f: + raw = f.read().strip() + ok = raw.replace(":", "") != "0" * 32 and bool(raw) + return CheckResult(f"gid:{device}/{gid_index}", ok, raw[:24] + "..." if len(raw) > 24 else raw) + except FileNotFoundError: + return CheckResult(f"gid:{device}/{gid_index}", False, "gid file missing") + except OSError as e: + return CheckResult(f"gid:{device}/{gid_index}", False, str(e)) + base = "/sys/class/infiniband" + if not os.path.isdir(base): + return CheckResult(f"gid:{gid_index}", False, "no infiniband dir") + for dev in sorted(os.listdir(base)): + return _check_gid_available(dev, gid_index) + return CheckResult(f"gid:{gid_index}", False, "no devices") + + +def _check_memlock() -> CheckResult: + try: + soft, hard = resource.getrlimit(resource.RLIMIT_MEMLOCK) + # ``soft`` of RLIM_INFINITY is typically -1 on Linux. + unlimited = soft in (-1, resource.RLIM_INFINITY) or soft > 2**30 + detail = f"soft={soft} hard={hard}" + return CheckResult("memlock", unlimited, detail) + except (ValueError, OSError) as e: + return CheckResult("memlock", False, str(e)) + + +def _check_health_check() -> CheckResult: # pragma: no cover - retained for ad-hoc use + """Call mooncake's native ``health_check()`` (NOT used by ``probe_node``). + + Returns 0=healthy, 1=not initialized/closed, 2=master unreachable. Kept as + a utility for post-init diagnostics; intentionally excluded from the pre- + init probe because the master is not running at probe time. + """ + try: + import mooncake + + hc = getattr(mooncake, "health_check", None) + if hc is None: + # Some builds expose it on the store module instead. + from mooncake import store # type: ignore + + hc = getattr(store, "health_check", None) + if hc is None: + return CheckResult("health_check", False, "health_check() not found in mooncake") + code = int(hc()) + ok = code == 0 + return CheckResult("health_check", ok, f"return_code={code}") + except Exception as e: + return CheckResult("health_check", False, str(e)) + + +# --------------------------------------------------------------------------- +# Per-node probe +# --------------------------------------------------------------------------- + +# NOTE: ``health_check()`` is intentionally NOT probed -- it queries the +# mooncake master, which is not yet running at probe time (auto_init=False), +# so it would always report failure and force a spurious SimpleStorage +# fallback. Master reachability is authoritatively checked by ``tq.init``. +_CHECK_FUNCS_NO_DEVICE = [ + _check_mooncake_import, + _check_rdma_devices, + _check_memlock, +] + + +def probe_node(device: str = "") -> ProbeResult: + """Run all capability checks on the current node. + + Parameters + ---------- + device + Explicit RDMA device name; empty = auto-detect first available. + """ + node = socket.gethostname() + checks: list[CheckResult] = [] + errors: list[str] = [] + + for fn in _CHECK_FUNCS_NO_DEVICE: + checks.append(fn()) + + # Device-dependent checks. + checks.append(_check_port_active(device)) + checks.append(_check_gid_available(device)) + + # Determine effective protocol via graded degradation. + mooncake_ok = any(c.name == "mooncake_import" and c.ok for c in checks) + rdma_dev_ok = any(c.name == "rdma_devices" and c.ok for c in checks) + port_ok = any(c.name.startswith("port_active") and c.ok for c in checks) + gid_ok = any(c.name.startswith("gid") and c.ok for c in checks) + memlock_ok = any(c.name == "memlock" and c.ok for c in checks) + + effective_protocol: str | None + effective_device = device + if not mooncake_ok: + effective_protocol = None + errors.append("mooncake not importable") + elif rdma_dev_ok and port_ok and gid_ok and memlock_ok: + effective_protocol = "rdma" + if not effective_device: + # Pick first available device for the summary. + base = "/sys/class/infiniband" + if os.path.isdir(base): + effective_device = sorted(os.listdir(base))[0] + else: + # mooncake usable but RDMA incomplete -> degrade to TCP (still MooncakeStore). + effective_protocol = "tcp" + if not rdma_dev_ok: + errors.append("no RDMA device") + elif not port_ok: + errors.append("HCA port not ACTIVE") + elif not gid_ok: + errors.append("GID unavailable") + elif not memlock_ok: + errors.append("memlock too low for RDMA MR registration") + + # GDR eligibility == RDMA transport available. The actual CUDA-context + # check (mooncake_client.py:87) runs in the *client* worker process at + # runtime, NOT in this probe task -- a separate Ray task always reports + # torch.cuda.is_initialized() == False, so probing it here would make GDR + # permanently unreachable. We assert transport capability only; the + # runtime client performs the CUDA check and warns/falls back if needed. + gdr_eligible = effective_protocol == "rdma" + + return ProbeResult( + node=node, + checks=tuple(checks), + effective_protocol=effective_protocol, + effective_device=effective_device, + gdr_eligible=gdr_eligible, + errors=tuple(errors), + ) + + +# --------------------------------------------------------------------------- +# Multi-node fan-out (driver -> every alive GPU node) +# --------------------------------------------------------------------------- + + +def _select_dataplane_node_ids(nodes: list[dict]) -> list[str]: + """Return node IDs of alive nodes that advertise GPU resources. + + The TransferQueue data plane runs on Actor + Rollout worker nodes, which + always advertise GPU resources. Head / CPU-only nodes are excluded so a + non-data-plane node cannot force a spurious RDMA degradation. + """ + ids: list[str] = [] + for n in nodes: + if not n.get("Alive"): + continue + resources = n.get("Resources") or {} + if resources.get("GPU", 0) >= 1: + ids.append(n["NodeID"]) + return ids + + +def _alive_gpu_nodes() -> list[str]: + """Discover alive GPU node IDs from the current Ray cluster. + + Thin seam around ``ray.nodes()``; kept separate so + :func:`probe_cluster_nodes` and its tests can stub discovery without + spinning up Ray. + """ + import ray + + return _select_dataplane_node_ids(ray.nodes()) + + +def _degenerate_result(node: str, error: str) -> ProbeResult: + """Build a :class:`ProbeResult` for a node whose probe failed or timed out. + + ``effective_protocol=None`` makes :func:`reduce_results` treat the node as + mooncake-unavailable (degrade toward TCP / SimpleStorage) instead of + silently dropping it, which would over-report cluster capability. + """ + return ProbeResult( + node=node, + checks=tuple(), + effective_protocol=None, + effective_device="", + gdr_eligible=False, + errors=(error,), + ) + + +def probe_cluster_nodes(device: str = "", *, timeout: float = 60.0) -> list[ProbeResult]: + """Probe every alive GPU-bearing node and return one result per node. + + The driver fans the probe out as a short-lived Ray remote task pinned to + each node via ``NodeAffinitySchedulingStrategy(soft=False)`` so that + ``probe_node`` reads *that node's* own ``/sys`` / mooncake state. The + caller then AND-reduces the returned list via :func:`reduce_results`, + producing a single job-level effective config that every worker converges + on — satisfying the requirement that the driver decide once for the whole + job rather than just probing its own node. + + A node whose probe task raises or exceeds ``timeout`` seconds is recorded + as a degenerate ``effective_protocol=None`` result so the reducer degrades + the whole job rather than silently omitting the node. + + Returns ``[probe_node(device)]`` when no GPU workers are discoverable + (single-node / local dev), preserving backward-compatible behavior. + """ + node_ids = _alive_gpu_nodes() + if not node_ids: + logger.debug("No alive GPU nodes discovered; probing driver node only.") + return [probe_node(device)] + + import ray + from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy + + @ray.remote(num_cpus=0.001) + def _probe_on_node(dev: str) -> ProbeResult: + from relax.utils.rdma_probe import probe_node as _probe + + return _probe(dev) + + refs: list[Any] = [] + id_by_ref: dict[Any, str] = {} + for node_id in node_ids: + strategy = NodeAffinitySchedulingStrategy(node_id=node_id, soft=False) + ref = _probe_on_node.options(scheduling_strategy=strategy).remote(device) + refs.append(ref) + id_by_ref[ref] = node_id + + ready, pending = ray.wait(refs, num_returns=len(refs), timeout=timeout) + results: list[ProbeResult] = [] + for ref in ready: + node_id = id_by_ref[ref] + try: + results.append(ray.get(ref)) + except Exception as e: # pragma: no cover - depends on remote task failure + logger.warning(f"[probe] node {node_id} probe task failed: {e}") + results.append(_degenerate_result(node_id, f"probe_task_failed:{e}")) + for ref in pending: + node_id = id_by_ref[ref] + ray.cancel(ref, force=True) + logger.warning(f"[probe] node {node_id} probe timed out after {timeout}s") + results.append(_degenerate_result(node_id, f"probe_timeout:{timeout}s")) + + return results + + +# --------------------------------------------------------------------------- +# Job-level AND reduction +# --------------------------------------------------------------------------- + + +def reduce_results( + results: list[ProbeResult], + *, + requested_backend: str, + requested_device: str, + use_gdr: bool, + fallback_backend: str = "SimpleStorage", +) -> EffectiveConfig: + """AND-reduce per-node results into a single job-level effective config. + + Parameters + ---------- + results + One :class:`ProbeResult` per data-plane node. + requested_backend + ``--tq-storage-backend`` value (``"simple"`` or ``"mooncake"``). + requested_device + ``--tq-rdma-device`` value. + use_gdr + ``--tq-use-gdr`` value. + fallback_backend + Backend to degrade to when probe fails in auto mode. + """ + # SimpleStorage short-circuits: no probing needed. + if requested_backend == "simple": + return EffectiveConfig( + backend="SimpleStorage", + protocol="tcp", + device="", + gdr=False, + fallback_reason="", + ) + + if not results: + return EffectiveConfig( + backend="SimpleStorage", + protocol="tcp", + device="", + gdr=False, + fallback_reason="no probe results", + ) + + # AND reduction: the job can only run at the lowest common capability. + any_no_mooncake = any(r.effective_protocol is None for r in results) + all_rdma = all(r.effective_protocol == "rdma" for r in results) + all_gdr = all(r.gdr_eligible for r in results) + + if any_no_mooncake: + failed_nodes = [r.node for r in results if r.effective_protocol is None] + return EffectiveConfig( + backend=fallback_backend, + protocol="tcp", + device="", + gdr=False, + fallback_reason=f"mooncake_unavailable:{','.join(failed_nodes)}", + ) + + if all_rdma: + # Device: if any node lacks the requested device, fall back to tcp. + if requested_device: + device_ok = all(r.effective_device == requested_device or not r.effective_device for r in results) + if not device_ok: + return EffectiveConfig( + backend="MooncakeStore", + protocol="tcp", + device=requested_device, + gdr=False, + fallback_reason=f"device_mismatch:{requested_device}", + ) + return EffectiveConfig( + backend="MooncakeStore", + protocol="rdma", + device=requested_device, + gdr=use_gdr and all_gdr, + fallback_reason="" if (not use_gdr or all_gdr) else "gdr_cuda_not_initialized", + ) + + # Some nodes can't do RDMA → degrade to TCP (still MooncakeStore). + rdma_failed = [r.node for r in results if r.effective_protocol != "rdma"] + return EffectiveConfig( + backend="MooncakeStore", + protocol="tcp", + device=requested_device, + gdr=False, + fallback_reason=f"rdma_unavailable:{','.join(rdma_failed)}", + ) + + +# --------------------------------------------------------------------------- +# Validation +# --------------------------------------------------------------------------- + + +def validate_config(args: Any) -> list[str]: + """Return a list of error messages for invalid flag combinations. + + Called at startup *before* probing. An empty list means the config is + structurally valid (semantic/runtime validity is checked by the probe). + """ + errors: list[str] = [] + backend = getattr(args, "tq_storage_backend", "simple") + mode = getattr(args, "tq_rdma_mode", "off") + use_gdr = getattr(args, "tq_use_gdr", False) + + if backend == "simple" and mode != "off": + errors.append( + f"--tq-rdma-mode={mode} is meaningless with --tq-storage-backend=simple " + "(RDMA only applies to MooncakeStore). Set --tq-rdma-mode=off or " + "--tq-storage-backend=mooncake." + ) + if backend == "simple" and use_gdr: + errors.append("--tq-use-gdr requires --tq-storage-backend=mooncake.") + if use_gdr and mode == "off": + errors.append( + "--tq-use-gdr is set but --tq-rdma-mode=off; GDR requires RDMA transport. " + "Set --tq-rdma-mode=auto or required." + ) + return errors diff --git a/relax/utils/tq_config.py b/relax/utils/tq_config.py new file mode 100644 index 000000000..9db8dbe6c --- /dev/null +++ b/relax/utils/tq_config.py @@ -0,0 +1,185 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Build TransferQueue backend config dicts from Relax CLI args. + +This module is the single place that maps Relax-side *intent* flags +(``--tq-storage-backend``, ``--tq-rdma-mode``, ``--tq-rdma-device``, +``--tq-use-gdr``) plus an :class:`~relax.utils.rdma_probe.EffectiveConfig` +into the OmegaConf dict that ``tq.init`` expects. + +Mooncake internals (endpoint, buffer size, segment size, master address, +timeout) are intentionally *not* exposed as CLI flags — they come from +internal defaults or the deployment environment, per maintainer guidance. +""" + +from __future__ import annotations + +import os +from typing import Any + +from relax.utils.logging_utils import get_logger +from relax.utils.rdma_probe import EffectiveConfig + + +logger = get_logger(__name__) + +# --------------------------------------------------------------------------- +# Defaults (kept here rather than in config.yaml so they are visible to +# Relax contributors without reading the TQ package). +# --------------------------------------------------------------------------- + +_DEFAULT_GLOBAL_SEGMENT_SIZE = 4 * 1024**3 # 4 GiB per client (config.yaml:52) +_DEFAULT_LOCAL_BUFFER_SIZE = 1 * 1024**3 # 1 GiB per client (config.yaml:54) +_DEFAULT_GDR_STAGING_MB = 1024 # config.yaml:103 +_DEFAULT_METADATA_SERVER = "P2PHANDSHAKE" # config.yaml:42-43 + + +# --------------------------------------------------------------------------- +# Config builders +# --------------------------------------------------------------------------- + + +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).""" + return { + # ``tq.init`` selects the manager from this key alone (TQ config.yaml:22 + # defaults it to SimpleStorage); a backend section without it is ignored. + "storage_backend": "SimpleStorage", + "SimpleStorage": { + "total_storage_size": total_storage_size, + "num_data_storage_units": num_data_storage_units, + }, + } + + +def build_mooncake_config( + effective: EffectiveConfig, + *, + master_address: str | None = None, + global_segment_size: int | None = None, +) -> dict[str, Any]: + """Build the ``backend`` dict for MooncakeStore. + + Parameters + ---------- + effective + The job-level :class:`EffectiveConfig` after probing. + master_address + External master server address. If ``None``, read from the + ``MC_MASTER_ADDRESS`` env var; if still unset, fall back to localhost + (single-node dev only — production must set the env var). + global_segment_size + Override the per-client segment size (default 4 GiB). Benchmarks may + pass a larger value (e.g. 8 GiB) to avoid staging-buffer pressure. + """ + if master_address is None: + master_address = os.environ.get("MC_MASTER_ADDRESS", "localhost:50051") + + cfg: dict[str, Any] = { + # Selects the manager inside ``tq.init`` (interface.py reads + # ``backend.storage_backend``, TQ config.yaml:22 defaults it to + # SimpleStorage). Without this key the MooncakeStore section below is + # parsed and then silently ignored -- the job still runs on ZMQ/TCP. + "storage_backend": "MooncakeStore", + "MooncakeStore": { + # Transport + "protocol": effective.protocol, # "rdma" or "tcp" + "device_name": effective.device, + # Master / metadata — externally managed, never auto-init. + "auto_init": False, + "master_server_address": master_address, + "metadata_server": _DEFAULT_METADATA_SERVER, + "local_hostname": "", # empty = auto-detect via Ray node IP + # Memory + "global_segment_size": global_segment_size or _DEFAULT_GLOBAL_SEGMENT_SIZE, + "local_buffer_size": _DEFAULT_LOCAL_BUFFER_SIZE, + # Do NOT silently evict produced-but-unconsumed data. + "hard_pin": True, + # GDR + "use_gdr": effective.gdr, + "gdr_staging_buffer_mb": _DEFAULT_GDR_STAGING_MB, + }, + } + return cfg + + +# --------------------------------------------------------------------------- +# Capacity validation +# --------------------------------------------------------------------------- + + +def estimate_payload_bytes(args: Any) -> int: + """Rough estimate of per-step multimodal payload size in bytes. + + Used only for segment-capacity pre-check. The real payload depends on + image resolution and patch count; this is a conservative lower bound based + on ``--multimodal-keys`` presence and ``n_samples_per_prompt``. + """ + n_samples = args.n_samples_per_prompt + rollout_batch = args.rollout_batch_size + # 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 + + +def validate_segment_capacity(args: Any, effective: EffectiveConfig) -> str | None: + """Return an error message if segment capacity is insufficient, else None. + + Only checked for MooncakeStore (SimpleStorage manages its own capacity via + ``total_storage_size``). The check is conservative: it uses the *per- + client* segment size (``global_segment_size``) against the in-flight upper + bound. + """ + if effective.backend != "MooncakeStore": + return None + + max_staleness = getattr(args, "max_staleness", 0) + payload = estimate_payload_bytes(args) + needed = payload * (max_staleness + 1) + available = _DEFAULT_GLOBAL_SEGMENT_SIZE + + if needed > available: + return ( + f"MooncakeStore segment capacity insufficient: estimated in-flight payload " + f"{needed / 1024**3:.1f} GiB (rollout_batch={args.rollout_batch_size} × " + f"n_samples={args.n_samples_per_prompt} × staleness+1={max_staleness + 1}) " + f"exceeds global_segment_size {available / 1024**3:.1f} GiB. " + f"Reduce batch size, increase global_segment_size, or reduce max_staleness." + ) + return None + + +# --------------------------------------------------------------------------- +# Top-level resolver +# --------------------------------------------------------------------------- + + +def build_backend_config( + args: Any, + effective: EffectiveConfig, + *, + total_storage_size: int, +) -> tuple[dict[str, Any], str | None]: + """Return ``(backend_config_dict, error_or_none)``. + + On error, ``backend_config_dict`` is a safe SimpleStorage fallback and + ``error`` explains why MooncakeStore was rejected. + """ + if effective.backend == "SimpleStorage": + return build_simple_storage_config( + total_storage_size=total_storage_size, + num_data_storage_units=args.num_data_storage_units, + ), None + + # MooncakeStore path. + cap_error = validate_segment_capacity(args, effective) + if cap_error: + logger.error(cap_error) + return build_simple_storage_config( + total_storage_size=total_storage_size, + num_data_storage_units=args.num_data_storage_units, + ), cap_error + + return build_mooncake_config(effective), None diff --git a/relax/utils/tq_lifecycle.py b/relax/utils/tq_lifecycle.py new file mode 100644 index 000000000..f6d11e7ee --- /dev/null +++ b/relax/utils/tq_lifecycle.py @@ -0,0 +1,114 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""TransferQueue controller/store lifecycle helpers. + +Extracted from :mod:`relax.core.controller` so the behavior can be unit-tested +without importing the whole Controller (which drags in Ray Serve and Megatron). +The Controller keeps thin wrappers that delegate here. + +Two problems these helpers exist for: + +* **F10 anti-hang** -- ``tq.init`` attaches to the existing named actor and polls + ``get_config`` forever while it returns ``None`` (TQ ``interface.py:109-118``), + so a controller left behind by a run that died between actor creation and + ``store_config`` turns the next start into a hang. +* **Segment leak on teardown** -- ``tq.close()`` only tears down ZMQ (TQ + ``storage/managers/base.py:378``); it never calls ``storage_client.close()``, + so a MooncakeStore segment stays mounted and registered in the master until + ``client_ttl`` (30 s) expires, and puts from the restarted job fail with + "Failed to open segment ... Connection refused". +""" + +from __future__ import annotations + +import time + +import ray +import transfer_queue as tq + +from relax.utils.logging_utils import get_logger + + +logger = get_logger(__name__) + +CONTROLLER_NAME = "TransferQueueController" +CONTROLLER_NAMESPACE = "transfer_queue" + + +def kill_tq_controller_and_wait(timeout: float = 20.0) -> None: + """Kill the TransferQueueController named actor, then wait for GCS + deregistration. + + Waiting matters because ``ray.kill`` is asynchronous: the handle stays + resolvable for a short window, and a re-init landing in that window + attaches to a dying actor and fails with ``ActorDiedError``. + """ + try: + stale = ray.get_actor(CONTROLLER_NAME, namespace=CONTROLLER_NAMESPACE) + ray.kill(stale) + logger.info("[dataplane] Killed TransferQueueController actor (F10 guard).") + except ValueError: + return # actor does not exist — nothing to kill or wait for. + except Exception as e: # pragma: no cover - best-effort cleanup + logger.warning(f"[dataplane] Failed to kill TransferQueueController: {e}") + return + + deadline = time.time() + timeout + while time.time() < deadline: + try: + ray.get_actor(CONTROLLER_NAME, namespace=CONTROLLER_NAMESPACE) + except ValueError: + return + time.sleep(0.4) + logger.warning(f"[dataplane] TransferQueueController still resolvable after {timeout}s; proceeding anyway.") + + +def reap_unusable_tq_controller(get_config_timeout: float = 10.0) -> bool: + """Kill the TransferQueueController only if it cannot serve a config. + + Returns ``True`` when a controller was reaped. A controller that *does* + return a config is left alone: it belongs to whoever created it, attaching + to it is the intended behavior, and this keeps the guard within "clean up + only what this job owns" (no broad pkill/killall). + """ + try: + existing = ray.get_actor(CONTROLLER_NAME, namespace=CONTROLLER_NAMESPACE) + except ValueError: + return False # nothing there — nothing to reap. + + try: + conf = ray.get(existing.get_config.remote(), timeout=get_config_timeout) + except Exception as e: # actor dead, unresponsive, or API missing + logger.warning(f"[dataplane] Existing TransferQueueController is unusable ({e}); reaping it.") + conf = None + + if conf is not None: + logger.info("[dataplane] Existing TransferQueueController is healthy; tq.init will attach to it.") + return False + + logger.warning("[dataplane] TransferQueueController has no stored config (half-initialised); reaping it.") + kill_tq_controller_and_wait() + return True + + +def close_tq_and_unmount() -> None: + """Close TransferQueue and unmount the MooncakeStore segment. + + Order matters: ``tq.close()`` still needs the store alive for its + ``remove_all()``, so the client handle is captured first and unmounted + after. SimpleStorage has no ``storage_client``, so this is a no-op there. + """ + store_client = None + try: + store_client = getattr(tq.get_client().storage_manager, "storage_client", None) + except (AssertionError, AttributeError): + pass # TQ not initialised in this process, or no KV client. + + tq.close() + + if store_client is not None and hasattr(store_client, "close"): + try: + store_client.close() + logger.info("[dataplane] Unmounted MooncakeStore segment on teardown.") + except Exception as e: # pragma: no cover - best-effort cleanup + logger.warning(f"[dataplane] Failed to unmount MooncakeStore segment: {e}") diff --git a/scripts/benchmarks/cross_node_rdma_bench.py b/scripts/benchmarks/cross_node_rdma_bench.py new file mode 100644 index 000000000..1dfcf9755 --- /dev/null +++ b/scripts/benchmarks/cross_node_rdma_bench.py @@ -0,0 +1,226 @@ +#!/usr/bin/env python +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Cross-node RDMA benchmark using raw MooncakeDistributedStore. + +This bypasses TransferQueue to measure the raw transport layer across two +nodes, isolating TCP vs RDMA net benefit. + +Setup: + Node A (holder): python scripts/benchmarks/cross_node_rdma_bench.py --role holder \ + --master 0.0.0.0:50051 --segment-gb 16 --device mlx5_bond_0 + Node B (client): python scripts/benchmarks/cross_node_rdma_bench.py --role client \ + --master :50051 --device mlx5_bond_0 \ + --protocol rdma --payload-mib 64 --repeats 5 +""" + +from __future__ import annotations + +import argparse +import os +import statistics +import time + +import torch + + +DEFAULT_DEVICE = "" + + +def parse_args() -> argparse.Namespace: + """Parse cross-node benchmark CLI arguments.""" + p = argparse.ArgumentParser(description="Cross-node RDMA benchmark") + p.add_argument( + "--role", + required=True, + choices=["holder", "client"], + help="holder=segment holder on node A, client=benchmark on node B", + ) + p.add_argument("--master", required=True, help="Master address host:port") + p.add_argument("--device", default=DEFAULT_DEVICE, help="RDMA device name") + p.add_argument("--protocol", default="rdma", choices=["tcp", "rdma"]) + p.add_argument("--segment-gb", type=int, default=16, help="Segment size in GiB (holder only)") + p.add_argument( + "--payload-mib", + nargs="+", + type=int, + default=[8, 64, 256], + help="Payload sizes in MiB (client only, --mode simple)", + ) + p.add_argument( + "--mode", + default="simple", + choices=["simple", "multimodal"], + help="simple=1D tensor, multimodal=32 samples x [patch,1176] variable-length", + ) + p.add_argument("--num-samples", type=int, default=32, help="Num samples (multimodal mode only)") + p.add_argument("--patch-min", type=int, default=1213, help="Min patch count per sample (multimodal)") + p.add_argument("--patch-max", type=int, default=2471, help="Max patch count per sample (multimodal)") + p.add_argument("--hidden", type=int, default=1176, help="Hidden dim per patch (multimodal)") + p.add_argument("--repeats", type=int, default=5, help="Repetitions (client only)") + p.add_argument("--warmup", type=int, default=2, help="Warmup rounds (client only)") + return p.parse_args() + + +def create_store(args, segment_size: int): + """Create a MooncakeDistributedStore.""" + from mooncake.store import MooncakeDistributedStore + + local_hostname = os.environ.get("MC_TCP_BIND_ADDRESS", "") + store = MooncakeDistributedStore() + store.setup( + local_hostname, # local_hostname + "P2PHANDSHAKE", # metadata_server + segment_size, # global_segment_size (0 on client = no local segment) + 1024 * 1024 * 1024, # local_buffer_size (1 GiB) + args.protocol, # protocol + args.device, # device_name + args.master, # master_server_address + ) + return store + + +def run_holder(args): + """Run as segment holder on node A — mounts a large segment and waits.""" + segment_size = args.segment_gb * 1024**3 + print( + f"[holder] Creating MooncakeDistributedStore: segment={args.segment_gb} GiB, " + f"protocol={args.protocol}, device={args.device}, master={args.master}" + ) + store = create_store(args, segment_size) + print("[holder] Segment mounted. Holding... (Ctrl+C to stop)") + print(f"[holder] Master: {args.master}") + try: + while True: + time.sleep(60) + except KeyboardInterrupt: + print("\n[holder] Shutting down...") + finally: + # Release the segment even on unexpected exit so the master does not + # leak a pinned-memory segment across benchmark runs. + try: + store.close() + except Exception as e: + print(f"[holder] store.close() failed: {e}") + + +def run_client(args): + """Run as client on node B — put/get payloads, measure throughput.""" + print( + f"[client] Connecting: protocol={args.protocol}, device={args.device}, " + f"master={args.master}, segment_size=0 (forces cross-node)" + ) + print(f"[client] mode={args.mode}, repeats={args.repeats}") + store = create_store(args, 0) # segment_size=0 → all data lands on holder + + if args.mode == "multimodal": + _run_multimodal(args, store) + else: + _run_simple(args, store) + + store.close() + print("[client] Done") + + +def _make_multimodal_payload(args): + """Create a realistic pixel_values payload: N samples x [patch, hidden] + float32. + + Patch counts are uniformly spread in [patch_min, patch_max] to match real + Qwen2-VL batches. All samples are packed into one contiguous 1D buffer + (what TQ serializes into for transport). + """ + import random + + rng = random.Random(42) # deterministic + patches = [rng.randint(args.patch_min, args.patch_max) for _ in range(args.num_samples)] + total_elements = sum(p * args.hidden for p in patches) + total_bytes = total_elements * 4 # float32 + data = torch.randn(total_elements, dtype=torch.float32) + return data, total_bytes, patches + + +def _run_simple(args, store): + for payload_mib in args.payload_mib: + payload_bytes = payload_mib * 1024 * 1024 + data = torch.randn(payload_bytes // 4, dtype=torch.float32) + key = f"bench_{payload_mib}mib" + _bench_one(store, key, data, payload_bytes, f"{payload_mib} MiB", args) + + +def _run_multimodal(args, store): + data, total_bytes, patches = _make_multimodal_payload(args) + total_mib = total_bytes / 1024 / 1024 + print( + f" [multimodal] {args.num_samples} samples, patches {min(patches)}-{max(patches)}, " + f"hidden={args.hidden}, total={total_mib:.1f} MiB ({total_bytes / 1e6:.1f} MB)" + ) + _bench_one(store, "bench_multimodal", data, total_bytes, f"{total_mib:.0f}M mm", args) + + # Also test with n_samples_per_prompt=8 group duplication (GRPO redundancy) + for mult in [4, 8]: + group_data = data.repeat(mult) + group_bytes = total_bytes * mult + group_mib = group_bytes / 1024 / 1024 + _bench_one(store, f"bench_mm_{mult}x", group_data, group_bytes, f"{group_mib:.0f}M mm×{mult}", args) + + +def _bench_one(store, key, data, payload_bytes, label, args): + """Put/get one payload, print timing.""" + # Warmup + for _ in range(args.warmup): + store.put_tensor(key, data) + _ = store.get_tensor(key) + + put_times = [] + get_times = [] + for i in range(args.repeats): + t0 = time.perf_counter() + store.put_tensor(key, data) + put_ms = (time.perf_counter() - t0) * 1000 + + t0 = time.perf_counter() + retrieved = store.get_tensor(key) + get_ms = (time.perf_counter() - t0) * 1000 + + # Correctness is a hard requirement, not an assert: ``python -O`` strips + # asserts, and a silent None skip would inflate throughput. Fail loudly. + if retrieved is None: + raise RuntimeError(f"get_tensor returned None for {key} (data lost)") + if not torch.equal(retrieved, data): + raise RuntimeError(f"Byte mismatch for {key}") + del retrieved + + put_gbs = payload_bytes / put_ms / 1e6 if put_ms > 0 else 0 + get_gbs = payload_bytes / get_ms / 1e6 if get_ms > 0 else 0 + put_times.append(put_ms) + get_times.append(get_ms) + print( + f" [{label:>12}] run {i + 1}/{args.repeats}: " + f"put={put_ms:>7.1f}ms ({put_gbs:.2f} GB/s) " + f"get={get_ms:>7.1f}ms ({get_gbs:.2f} GB/s)" + ) + + put_med = statistics.median(put_times) + get_med = statistics.median(get_times) + put_gbs = payload_bytes / put_med / 1e6 if put_med > 0 else 0 + get_gbs = payload_bytes / get_med / 1e6 if get_med > 0 else 0 + print( + f" [{label:>12}] MEDIAN: " + f"put={put_med:>7.1f}ms ({put_gbs:.2f} GB/s) " + f"get={get_med:>7.1f}ms ({get_gbs:.2f} GB/s)" + ) + print() + + +def main(): + """Dispatch to holder (node A) or client (node B) per ``--role``.""" + args = parse_args() + if args.role == "holder": + run_holder(args) + else: + run_client(args) + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmarks/tq_cross_node_bench.py b/scripts/benchmarks/tq_cross_node_bench.py new file mode 100644 index 000000000..a0c56e255 --- /dev/null +++ b/scripts/benchmarks/tq_cross_node_bench.py @@ -0,0 +1,490 @@ +#!/usr/bin/env python +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""TQ-layer cross-node benchmark matching REAL Relax usage. + +Real Relax: every component (actor/rollout/critic) is a PERSISTENT +``@serve.deployment`` Ray actor that calls ``tq.init`` once (attach) + +``tq.get_client``, then put/get for the whole job. This benchmark mirrors +that: a persistent consumer ACTOR on node B (declared via ``@ray.remote`` +class, scheduled with ``NodeAffinitySchedulingStrategy``) that attaches once +and fetches repeatedly -- NOT an ephemeral task. + +Three configs run in the SAME cross-node topology so they are directly +comparable: + + * C0 ``simple`` -- SimpleStorage / ZMQ / TCP (current default backend) + * C1 ``tcp`` -- MooncakeStore / mooncake / TCP + * C2 ``rdma`` -- MooncakeStore / mooncake / RDMA + +Every payload档's transport is proven on the wire by reading the IB +``port_rcv_data`` and bond0 ``rx_bytes`` counters around each get: RDMA档 must +show IB moving and bond0 flat; TCP档 the reverse. There is no "thought it was +RDMA but was TCP" ambiguity. + +Usage (node A driver; node B already in the Ray cluster): + + PYTHONPATH= python -u scripts/benchmarks/tq_cross_node_bench.py \\ + --payload-mib 256 1024 2048 4096 --repeats 5 --csv tq_cross_node_gib.csv + +On mooncake 0.3.10, switching protocols inside one driver session can make the +third protocol's ``batch_get_into`` return -800 (see task-26-dev-log §7.11). +If that happens, run one protocol per process and merge the CSVs: + + --protocols rdma --csv c2.csv + --protocols tcp --csv c1.csv + --protocols simple --csv c0.csv +""" + +from __future__ import annotations + +import argparse +import csv +import statistics +import time +from typing import Any + +import ray + + +# Placeholder defaults -- override --master / --nodeb-ip / --device with your own +# cluster's values before running. Do not commit real infrastructure IPs/devices. +DEFAULT_MASTER = ":50051" +DEFAULT_DEVICE = "" +DEFAULT_NODEB = "" + + +def parse_args() -> argparse.Namespace: + """Parse CLI arguments.""" + p = argparse.ArgumentParser(description="TQ-layer cross-node RDMA vs TCP benchmark") + p.add_argument("--master", default=DEFAULT_MASTER, help="mooncake master host:port (node A)") + p.add_argument("--device", default=DEFAULT_DEVICE, help="RDMA device name") + p.add_argument("--nodeb-ip", default=DEFAULT_NODEB, help="node B NodeManagerAddress") + p.add_argument( + "--payload-mib", + nargs="+", + type=int, + default=[256, 1024, 2048, 4096], + help="Total payload sizes per put in MiB (4096 == 4 GiB). Defaults span 256M..4G.", + ) + p.add_argument("--num-samples", type=int, default=256, help="Rows per put (rollout-batch-like)") + p.add_argument("--num-fields", nargs="+", type=int, default=[1], help="Tensor fields per put") + p.add_argument( + "--repeats", + type=int, + default=5, + help="Measured put/get rounds per档 (a separate warmup round is always added; the mean " + "of these repeats is the reported figure)", + ) + p.add_argument( + "--segment-gib", + type=int, + default=16, + help="MooncakeStore global_segment_size per client (GiB). Must exceed the largest payload.", + ) + p.add_argument( + "--protocols", + nargs="+", + default=["simple", "tcp", "rdma"], + choices=["simple", "tcp", "rdma"], + help="Subset of configs to run (use one per process to dodge the 0.3.10 -800 issue).", + ) + p.add_argument("--csv", default="", help="Optional path to write per-run rows + summary") + return p.parse_args() + + +def make_payload(num_samples: int, num_fields: int, total_mib: int): + """Build a TensorDict of num_samples rows x num_fields tensors totaling. + + ~total_mib. + """ + import torch + from tensordict import TensorDict + + dt = torch.float32 + elem = torch.tensor([], dtype=dt).element_size() + per_field = total_mib * 1024 * 1024 // max(1, num_fields) + cols = max(1, per_field // (elem * num_samples)) + g = torch.Generator().manual_seed(1234) + data = {f"field_{i}": torch.randn(num_samples, cols, dtype=dt, generator=g) for i in range(num_fields)} + return TensorDict(data, batch_size=[num_samples]) + + +def payload_bytes(payload) -> int: + """Total bytes across all tensor fields.""" + return sum(payload[k].nelement() * payload[k].element_size() for k in payload.keys()) + + +def wait_actor_gone(name: str = "TransferQueueController", timeout: float = 30.0) -> None: + """Wait for a named TQ actor to leave the GCS (F10-safe re-init).""" + deadline = time.time() + timeout + while time.time() < deadline: + try: + ray.get_actor(name, namespace="transfer_queue") + except ValueError: + return + time.sleep(0.4) + + +def close_tq_unmount_and_wait() -> None: + """Close TQ, unmount the Mooncake segment, then wait for controller + deregistration. + + ``tq.close()`` only tears down ZMQ (managers/base.py:378); the Mooncake + segment stays mounted and registered in the master, so the next config's + put hits a dead endpoint until client_ttl (30 s) expires. Unmount + explicitly after close (close itself still needs the store alive for + remove_all()). + """ + import transfer_queue as tq + + store_client = None + try: + store_client = getattr(tq.get_client().storage_manager, "storage_client", None) + except (AssertionError, AttributeError): + pass + + tq.close() + + if store_client is not None and hasattr(store_client, "close"): + try: + store_client.close() + except Exception as e: # pragma: no cover - best effort + print(f" [warn] store_client.close() failed: {e}", flush=True) + + wait_actor_gone() + + +def read_counters() -> dict[str, int]: + """IB port_rcv_data (all devices, bytes) + bond0 rx_bytes, for transport + proof.""" + import os + + out: dict[str, int] = {} + root = "/sys/class/infiniband" + try: + for dev in sorted(os.listdir(root)): + ports_dir = f"{root}/{dev}/ports" + if not os.path.isdir(ports_dir): + continue + for port in sorted(os.listdir(ports_dir)): + try: + with open(f"{ports_dir}/{port}/counters/port_rcv_data") as fh: + out[f"ib:{dev}:{port}"] = int(fh.read().strip()) * 4 + except OSError: + pass + except OSError: + pass + try: + with open("/sys/class/net/bond0/statistics/rx_bytes") as fh: + out["tcp:bond0"] = int(fh.read().strip()) + except OSError: + pass + return out + + +def build_conf(protocol: str, master: str, device: str, segment_gib: int): + """Build the tq.init OmegaConf. + + ``protocol="simple"`` selects the SimpleStorage/ZMQ baseline (C0) so the + three-config comparison (C0 / Mooncake-TCP / Mooncake-RDMA) runs in the + *same* cross-node topology; "tcp"/"rdma" select MooncakeStore. + """ + from omegaconf import OmegaConf + from transfer_queue import GRPOGroupNSampler + + from relax.utils.rdma_probe import EffectiveConfig + from relax.utils.tq_config import build_mooncake_config, build_simple_storage_config + + if protocol == "simple": + # total_storage_size=None == unlimited sample count (TQ config.yaml default). + backend = build_simple_storage_config(total_storage_size=None, num_data_storage_units=2) + else: + eff = EffectiveConfig(backend="MooncakeStore", protocol=protocol, device=device, gdr=False, fallback_reason="") + backend = build_mooncake_config(eff, master_address=master, global_segment_size=segment_gib * 1024**3) + return OmegaConf.create( + { + "controller": {"sampler": GRPOGroupNSampler(n_samples_per_prompt=1), "polling_mode": True}, + "backend": backend, + }, + flags={"allow_objects": True}, + ) + + +# ---- Persistent consumer actor (module-level, mirrors a Relax component) ---- + + +@ray.remote(num_cpus=0.001) +class TQConsumer: + """Persistent consumer on node B: attaches once, fetches many times. + + Mirrors a Relax component actor (actor.py / rollout.py): tq.init once in + __init__ (attach to the shared controller), tq.get_client, then repeated + get_meta/get_data for the job lifetime. + """ + + def __init__(self, master: str, device: str, protocol: str, segment_gib: int): + import transfer_queue as tq + + tq.init(conf=build_conf(protocol, master, device, segment_gib)) # attaches (conf ignored on attach) + self.client = tq.get_client() + + def alive(self) -> bool: + return True + + def describe(self) -> dict: + """Report the manager/client actually instantiated -- never trust the + conf alone.""" + mgr = self.client.storage_manager + inner = getattr(mgr, "storage_client", None) + return { + "manager": type(mgr).__name__, + "client": type(inner).__name__ if inner is not None else "-", + "protocol": getattr(inner, "protocol", "-"), + } + + def shutdown(self) -> None: + """Unmount the Mooncake segment before this actor is killed.""" + inner = getattr(self.client.storage_manager, "storage_client", None) + if inner is not None and hasattr(inner, "close"): + try: + inner.close() + except Exception: # pragma: no cover - best effort + pass + + def fetch(self, fields, batch_size: int, partition: str): + """One cross-node get; return (ms, ib_mb, tcp_mb, ib_tail_mb, + tcp_tail_mb). + + ``ib_tail`` is the IB delta in the 5 ms *after* ``get_data`` returns. + If ``get_data`` is synchronous it is ~0; if it returned before RDMA + finished (async), the tail keeps flowing and ``ib_tail`` > 0 -- a + definitive async-completion detector that does not need a costly full- + data touch. + """ + before = read_counters() + t0 = time.perf_counter() + meta = self.client.get_meta( + data_fields=list(fields), + batch_size=batch_size, + partition_id=partition, + mode="fetch", + task_name="xfer", + ) + got = self.client.get_data(meta) + ms = (time.perf_counter() - t0) * 1000 + after = read_counters() + time.sleep(0.005) # let any async RDMA tail register on the counters + settled = read_counters() + # Touch the data so the get is not optimized away and to sanity-check shape. + for f in fields: + v = got[f] + _ = v.values().reshape(-1) if type(v).__name__ == "NestedTensor" else v.reshape(-1) + ib = sum(after[k] - before.get(k, 0) for k in after if k.startswith("ib:")) / 1e6 + tcp = sum(after[k] - before.get(k, 0) for k in after if k.startswith("tcp:")) / 1e6 + ib_tail = sum(settled[k] - after.get(k, 0) for k in settled if k.startswith("ib:")) / 1e6 + tcp_tail = sum(settled[k] - after.get(k, 0) for k in settled if k.startswith("tcp:")) / 1e6 + return ms, ib, tcp, ib_tail, tcp_tail + + +def _mean(values: list[float]) -> float: + """Arithmetic mean (the reported statistic); 0 for an empty list.""" + return statistics.mean(values) if values else 0.0 + + +def _gbs(nbytes: int, ms: float) -> float: + """Convert a latency in ms to throughput in GB/s.""" + return nbytes / ms / 1e6 if ms > 0 else 0.0 + + +def main() -> None: + """Run the cross-node TQ benchmark for the requested protocols.""" + import transfer_queue as tq + from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy + + args = parse_args() + ray.init(ignore_reinit_error=True, address="auto", logging_level="ERROR") + + nodeb = next(n for n in ray.nodes() if n["NodeManagerAddress"] == args.nodeb_ip and n.get("Alive")) + nodeb_id = nodeb["NodeID"] + strat = NodeAffinitySchedulingStrategy(node_id=nodeb_id, soft=False) + print( + f"[setup] node A = driver | node B = {args.nodeb_ip} ({nodeb_id[:10]}) " + f"| master = {args.master} | device = {args.device} " + f"| segment = {args.segment_gib} GiB | repeats = {args.repeats} (mean)", + flush=True, + ) + + labels = {"simple": "C0 SimpleStorage", "tcp": "C1 Mooncake/TCP", "rdma": "C2 Mooncake/RDMA"} + # results[protocol][total_mib][num_fields] = { + # "put_mean_gbs","get_mean_gbs","get_med_gbs","get_std_gbs","wire", "per_run":[...]} + results: dict[str, dict[int, dict[int, dict[str, Any]]]] = {} + csv_rows: list[dict[str, Any]] = [] + + for protocol in args.protocols: + print(f"\n===== {labels[protocol]} =====", flush=True) + close_tq_unmount_and_wait() + tq.init(conf=build_conf(protocol, args.master, args.device, args.segment_gib)) + producer = tq.get_client() + prod_mgr = type(producer.storage_manager).__name__ + prod_client = getattr(producer.storage_manager, "storage_client", None) + print( + f" node A manager={prod_mgr} client={type(prod_client).__name__ if prod_client else '-'}", + flush=True, + ) + + consumer = TQConsumer.options(scheduling_strategy=strat).remote( + args.master, args.device, protocol, args.segment_gib + ) + ray.get(consumer.alive.remote()) # ensure attached before measuring + desc = ray.get(consumer.describe.remote()) + print( + f" node B manager={desc['manager']} client={desc['client']} protocol={desc['protocol']}", + flush=True, + ) + + results.setdefault(protocol, {}) + for total_mib in args.payload_mib: + results[protocol].setdefault(total_mib, {}) + for nf in args.num_fields: + payload = make_payload(args.num_samples, nf, total_mib) + fields = sorted(payload.keys()) + nbytes = payload_bytes(payload) + put_times: list[float] = [] + get_times: list[float] = [] + ibs: list[float] = [] + tcps: list[float] = [] + ib_tails: list[float] = [] # async-completion detector (~0 == sync get) + tcp_tails: list[float] = [] + # repeat 0 is a warm-up: first transfer pays RDMA endpoint handshake. + for r in range(args.repeats + 1): + part = f"xfer_{protocol}_{total_mib}_{nf}_{r}" + t0 = time.perf_counter() + producer.put(payload, partition_id=part) + put_ms = (time.perf_counter() - t0) * 1000 + get_ms, ib, tcp, ib_tail, tcp_tail = ray.get(consumer.fetch.remote(fields, args.num_samples, part)) + producer.clear_partition(part) + if r == 0: + print( + f" {total_mib}M f={nf} (warmup, not counted): " + f"put={put_ms:.0f}ms get={get_ms:.0f}ms " + f"ib_tail={ib_tail:.0f}MB", + flush=True, + ) + continue + put_times.append(put_ms) + get_times.append(get_ms) + ibs.append(ib) + tcps.append(tcp) + ib_tails.append(ib_tail) + tcp_tails.append(tcp_tail) + csv_rows.append( + { + "protocol": protocol, + "payload_mib": total_mib, + "num_fields": nf, + "run": r, + "put_ms": round(put_ms, 2), + "get_ms": round(get_ms, 2), + "put_gbs": round(_gbs(nbytes, put_ms), 3), + "get_gbs": round(_gbs(nbytes, get_ms), 3), + "ib_mb": round(ib, 1), + "tcp_mb": round(tcp, 1), + "ib_tail_mb": round(ib_tail, 1), + "tcp_tail_mb": round(tcp_tail, 1), + } + ) + + put_mean = _mean(put_times) + get_mean = _mean(get_times) + get_med = statistics.median(get_times) + # Std-dev of the per-run THROUGHPUT values (not of latency): converting a + # latency stddev via _gbs(nbytes, std_ms) would be a meaningless number. + get_gbs_runs = [_gbs(nbytes, ms) for ms in get_times] + get_std_gbs = statistics.pstdev(get_gbs_runs) if len(get_gbs_runs) > 1 else 0.0 + ib_med = statistics.median(ibs) + tcp_med = statistics.median(tcps) + ib_tail_med = statistics.median(ib_tails) if ib_tails else 0.0 + wire = "RDMA" if ib_med > tcp_med else "TCP" + rec = { + "put_mean_gbs": _gbs(nbytes, put_mean), + "get_mean_gbs": _gbs(nbytes, get_mean), + "get_med_gbs": _gbs(nbytes, get_med), + "get_std_gbs": get_std_gbs, + "wire": wire, + "ib_tail_med_mb": ib_tail_med, + "per_run_get_gbs": [round(g, 2) for g in get_gbs_runs], + } + results[protocol][total_mib][nf] = rec + print( + f" {str(total_mib) + 'M':<9} f={nf} put_mean={rec['put_mean_gbs']:6.2f} " + f"GB/s get_mean={rec['get_mean_gbs']:6.2f} (med {rec['get_med_gbs']:.2f}, " + f"std {rec['get_std_gbs']:.2f}) GB/s " + f"[wire: IB {ib_med:.0f}MB / bond0 {tcp_med:.0f}MB -> {wire}; " + f"tail {ib_tail_med:.0f}MB] " + f"runs={rec['per_run_get_gbs']}", + flush=True, + ) + + ray.get(consumer.shutdown.remote()) # unmount before kill, else the segment lingers + ray.kill(consumer) + close_tq_unmount_and_wait() + + # ---- Summary (mean-based, all requested protocols) ---- + print("\n===== SUMMARY: TQ-layer cross-node, same topology (get, MEAN of N runs) =====", flush=True) + header = ( + f"{'Payload':<9}{'f':<4}{'C0 mean':>9}{'C1 mean':>9}{'C2 mean':>9}" + f"{'C1/C0':>8}{'C2/C1':>8}{'C2 std':>8}{'>=20%':>7}{'wire C0/C1/C2':>18}" + ) + print(header, flush=True) + for total_mib in args.payload_mib: + for nf in args.num_fields: + c0 = results.get("simple", {}).get(total_mib, {}).get(nf) + c1 = results.get("tcp", {}).get(total_mib, {}).get(nf) + c2 = results.get("rdma", {}).get(total_mib, {}).get(nf) + if not (c0 and c1 and c2): + # A protocol was skipped (--protocols subset) -- print what we have. + parts = [] + for name, c in (("C0", c0), ("C1", c1), ("C2", c2)): + parts.append(f"{name}={c['get_mean_gbs']:.2f}" if c else f"{name}=-") + print(f"{str(total_mib) + 'M':<9}{nf:<4}" + " ".join(parts) + " (subset run)", flush=True) + continue + g0, g1, g2 = c0["get_mean_gbs"], c1["get_mean_gbs"], c2["get_mean_gbs"] + back_pct = (g1 - g0) / g0 * 100 if g0 > 0 else 0.0 + rdma_pct = (g2 - g1) / g1 * 100 if g1 > 0 else 0.0 + wire = f"{c0['wire']}/{c1['wire']}/{c2['wire']}" + print( + f"{str(total_mib) + 'M':<9}{nf:<4}{g0:>9.2f}{g1:>9.2f}{g2:>9.2f}" + f"{f'{back_pct:+.0f}%':>8}{f'{rdma_pct:+.0f}%':>8}{c2['get_std_gbs']:>8.2f}" + f"{('PASS' if rdma_pct >= 20 else 'no'):>7}{wire:>18}", + flush=True, + ) + print(" C1/C0 = MooncakeStore vs SimpleStorage (backend effect)", flush=True) + print(" C2/C1 = RDMA vs TCP on the same backend (transport effect, gate target, mean-based)", flush=True) + print(" std = population stddev of C2 get across the N runs (run-to-run variance)", flush=True) + + if args.csv: + cols = [ + "protocol", + "payload_mib", + "num_fields", + "run", + "put_ms", + "get_ms", + "put_gbs", + "get_gbs", + "ib_mb", + "tcp_mb", + "ib_tail_mb", + "tcp_tail_mb", + ] + with open(args.csv, "w", newline="") as fh: + w = csv.DictWriter(fh, fieldnames=cols) + w.writeheader() + for row in csv_rows: + w.writerow({k: row[k] for k in cols}) + print(f"\n[csv] wrote {len(csv_rows)} per-run rows to {args.csv}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmarks/tq_rdma_bench.py b/scripts/benchmarks/tq_rdma_bench.py new file mode 100644 index 000000000..6133dbae2 --- /dev/null +++ b/scripts/benchmarks/tq_rdma_bench.py @@ -0,0 +1,396 @@ +#!/usr/bin/env python +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""TransferQueue RDMA benchmark: SimpleStorage vs Mooncake/TCP vs +Mooncake/RDMA. + +Runs put/get on synthetic payloads and reports throughput for three +configurations, isolating the RDMA net benefit (C2 - C1). + +Usage (single-node loopback, no master needed for C0): + python scripts/benchmarks/tq_rdma_bench.py \ + --payload-mib 1 16 64 256 \ + --num-samples 256 \ + --num-fields 1 8 32 \ + --repeats 3 + +Usage (cross-node RDMA, needs external mooncake master): + MC_MASTER_ADDRESS=node-A:50051 \ + MC_TCP_BIND_ADDRESS= \ + python scripts/benchmarks/tq_rdma_bench.py \ + --configs C0 C1 C2 \ + --payload-mib 64 256 \ + --num-samples 64 \ + --repeats 3 \ + --device mlx5_bond_0 +""" + +from __future__ import annotations + +import argparse +import os +import statistics +import time + +import torch + + +# --------------------------------------------------------------------------- # +# Argument parsing +# --------------------------------------------------------------------------- # + + +def parse_args() -> argparse.Namespace: + """Parse benchmark CLI arguments.""" + p = argparse.ArgumentParser(description="TransferQueue RDMA benchmark") + p.add_argument( + "--configs", + nargs="+", + default=["C0", "C1", "C2"], + choices=["C0", "C1", "C2"], + help="C0=SimpleStorage, C1=Mooncake/TCP, C2=Mooncake/RDMA", + ) + p.add_argument( + "--payload-mib", + nargs="+", + type=int, + default=[16, 64, 256], + help="Payload sizes in MiB (total across all fields)", + ) + p.add_argument("--num-samples", type=int, default=256, help="Number of samples (rows)") + p.add_argument("--num-fields", nargs="+", type=int, default=[1, 8, 32], help="Number of fields per sample") + p.add_argument("--repeats", type=int, default=3, help="Repetitions per config (report median)") + p.add_argument("--warmup", type=int, default=1, help="Warmup rounds (not counted)") + p.add_argument("--device", type=str, default="", help="RDMA device name (e.g. mlx5_bond_0). Empty = auto.") + p.add_argument( + "--master-address", + type=str, + default=None, + help="Mooncake master address. Default: env MC_MASTER_ADDRESS or localhost:50051", + ) + p.add_argument("--dtype", type=str, default="float32", choices=["float32", "bfloat16", "float16"]) + p.add_argument("--output-csv", type=str, default=None, help="Write results to CSV file") + return p.parse_args() + + +# --------------------------------------------------------------------------- # +# Config builders +# --------------------------------------------------------------------------- # + +CONFIG_MAP = { + "C0": {"backend": "SimpleStorage", "protocol": "tcp"}, + "C1": {"backend": "MooncakeStore", "protocol": "tcp"}, + "C2": {"backend": "MooncakeStore", "protocol": "rdma"}, +} + + +def build_tq_config(config_name: str, args: argparse.Namespace, num_storage_units: int = 1): + """Build the tq.init config dict for a given benchmark config. + + Reuses :mod:`relax.utils.tq_config` builders so the benchmark cannot drift + from the production config shape (single source of truth for + keys/defaults). + """ + from omegaconf import OmegaConf + from transfer_queue import GRPOGroupNSampler + + from relax.utils.rdma_probe import EffectiveConfig + from relax.utils.tq_config import build_mooncake_config, build_simple_storage_config + + cfg = CONFIG_MAP[config_name] + sampler = GRPOGroupNSampler(n_samples_per_prompt=1) + + if cfg["backend"] == "SimpleStorage": + backend_dict = build_simple_storage_config( + total_storage_size=1024**3, num_data_storage_units=num_storage_units + ) + else: + master_addr = args.master_address or os.environ.get("MC_MASTER_ADDRESS", "localhost:50051") + eff = EffectiveConfig( + backend="MooncakeStore", + protocol=cfg["protocol"], + device=args.device, + gdr=False, + fallback_reason="", + ) + backend_dict = build_mooncake_config(eff, master_address=master_addr, global_segment_size=8 * 1024**3) + + return OmegaConf.create( + { + "controller": {"sampler": sampler, "polling_mode": True}, + "backend": backend_dict, + }, + flags={"allow_objects": True}, + ) + + +def close_tq_and_wait(timeout: float = 20.0) -> None: + """Close TQ, unmount the Mooncake segment, and wait for the controller to + leave the GCS. + + Required between configs: ``tq.init`` attaches to an existing controller + and ignores the new conf (interface.py:152), so without close+wait every + config after the first silently reuses the first config's backend. + + ``tq.close()`` tears down only the ZMQ layer (managers/base.py:378); it never + calls ``storage_client.close()``. With MooncakeStore that leaves the segment + mounted and still registered in the master, so the next config's put targets a + dead endpoint ("Failed to open segment ... Connection refused") until the + master's ``client_ttl`` (30 s) expires. Unmount explicitly instead. + """ + import time + + import ray + import transfer_queue as tq + + store_client = None + try: + store_client = getattr(tq.get_client().storage_manager, "storage_client", None) + except (AssertionError, AttributeError): + pass + + tq.close() # runs remove_all() through the store, so unmount has to come after + + if store_client is not None and hasattr(store_client, "close"): + store_client.close() # unmounts the segment and deregisters from the master + + deadline = time.time() + timeout + while time.time() < deadline: + try: + ray.get_actor("TransferQueueController", namespace="transfer_queue") + except ValueError: + return + time.sleep(0.4) + + +# --------------------------------------------------------------------------- # +# Payload generation +# --------------------------------------------------------------------------- # + + +def make_payload(num_samples: int, num_fields: int, total_mib: int, dtype: str): + """Create a synthetic TensorDict of ``num_fields`` tensors. + + Each tensor has shape (num_samples, N) with the requested dtype, where N is + chosen so total bytes ≈ total_mib * 1024^2. + + Returns a ``TensorDict`` with ``batch_size=[num_samples]`` (what TQ's + ``client.put`` expects). + """ + from tensordict import TensorDict + + dt = getattr(torch, dtype) + elem_size = torch.tensor([], dtype=dt).element_size() + total_bytes = total_mib * 1024 * 1024 + per_field_bytes = total_bytes // num_fields + cols = max(1, per_field_bytes // (elem_size * num_samples)) + + data = {} + for f in range(num_fields): + data[f"field_{f}"] = torch.randn(num_samples, cols, dtype=dt) + return TensorDict(data, batch_size=[num_samples]) + + +def payload_bytes(payload) -> int: + """Return total bytes across all tensor fields in ``payload``.""" + return sum(payload[key].nelement() * payload[key].element_size() for key in payload.keys()) + + +# --------------------------------------------------------------------------- # +# Benchmark core +# --------------------------------------------------------------------------- # + + +def run_one(config_name: str, payload: dict, args: argparse.Namespace) -> dict: + """Run put/get once and return timing.""" + import transfer_queue as tq + + # Close any prior controller and wait for GCS deregistration so this config + # gets a fresh backend (tq.init otherwise attaches to the existing one). + close_tq_and_wait() + tq_config = build_tq_config(config_name, args) + tq_config = tq.init(conf=tq_config) or tq_config + client = tq.get_client() + + nbytes = payload_bytes(payload) + + # Warmup + for _ in range(args.warmup): + client.put(payload, partition_id="bench") + client.clear_partition("bench") + + # Timed put + t0 = time.perf_counter() + client.put(payload, partition_id="bench") + put_ms = (time.perf_counter() - t0) * 1000 + + # Timed get: create a fetch meta, then get_data + field_names = sorted(payload.keys()) + t0 = time.perf_counter() + fetch_meta = client.get_meta( + data_fields=field_names, + batch_size=args.num_samples, + partition_id="bench", + mode="fetch", + task_name="bench", + ) + data = client.get_data(fetch_meta) + get_ms = (time.perf_counter() - t0) * 1000 + + # Correctness spot-check (best-effort: get_data may add non-tensor fields) + mismatches = [] + for k in field_names: + try: + if not torch.equal(data[k], payload[k]): + mismatches.append(k) + except Exception: + pass # non-tensor field, skip + if mismatches: + raise RuntimeError(f"Byte mismatch in fields: {mismatches}") + + client.clear_partition("bench") + + return { + "config": config_name, + "backend": CONFIG_MAP[config_name]["backend"], + "protocol": CONFIG_MAP[config_name]["protocol"], + "put_ms": put_ms, + "get_ms": get_ms, + "nbytes": nbytes, + "put_gbs": nbytes / put_ms / 1e6 if put_ms > 0 else 0, + "get_gbs": nbytes / get_ms / 1e6 if get_ms > 0 else 0, + } + + +def run_config(config_name: str, payload: dict, args: argparse.Namespace) -> dict: + """Run ``args.repeats`` times, return median.""" + results = [] + for i in range(args.repeats): + r = run_one(config_name, payload, args) + results.append(r) + print( + f" [{config_name}] run {i + 1}/{args.repeats}: put={r['put_ms']:.1f}ms " + f"({r['put_gbs']:.2f} GB/s) get={r['get_ms']:.1f}ms ({r['get_gbs']:.2f} GB/s)" + ) + + put_vals = sorted(r["put_ms"] for r in results) + get_vals = sorted(r["get_ms"] for r in results) + put_med = statistics.median(put_vals) + get_med = statistics.median(get_vals) + nbytes = results[0]["nbytes"] + return { + "config": config_name, + "backend": CONFIG_MAP[config_name]["backend"], + "protocol": CONFIG_MAP[config_name]["protocol"], + "put_ms_median": put_med, + "get_ms_median": get_med, + "put_ms_min": min(put_vals), + "put_ms_max": max(put_vals), + "get_ms_min": min(get_vals), + "get_ms_max": max(get_vals), + "nbytes": nbytes, + "put_gbs_median": nbytes / put_med / 1e6 if put_med > 0 else 0, + "get_gbs_median": nbytes / get_med / 1e6 if get_med > 0 else 0, + } + + +# --------------------------------------------------------------------------- # +# Main +# --------------------------------------------------------------------------- # + + +def main(): + """Run the benchmark across all requested payload/field/config + combinations.""" + args = parse_args() + + print("=" * 80) + print("TransferQueue RDMA Benchmark") + print(f" configs: {args.configs}") + print(f" payload sizes: {args.payload_mib} MiB") + print(f" samples: {args.num_samples}, fields: {args.num_fields}") + print(f" repeats: {args.repeats}, dtype: {args.dtype}") + print(f" device: {args.device or 'auto'}") + print("=" * 80) + + all_results = [] + + for total_mib in args.payload_mib: + for nf in args.num_fields: + payload = make_payload(args.num_samples, nf, total_mib, args.dtype) + actual_mib = payload_bytes(payload) / 1024 / 1024 + print(f"\n--- {actual_mib:.1f} MiB / {args.num_samples} samples / {nf} fields ---") + + for cfg in args.configs: + try: + result = run_config(cfg, payload, args) + all_results.append(result) + print( + f" [{cfg}] MEDIAN: put={result['put_ms_median']:.1f}ms " + f"({result['put_gbs_median']:.2f} GB/s) " + f"get={result['get_ms_median']:.1f}ms ({result['get_gbs_median']:.2f} GB/s)" + ) + except Exception as e: + print(f" [{cfg}] FAILED: {e}") + all_results.append( + { + "config": cfg, + "backend": CONFIG_MAP[cfg]["backend"], + "protocol": CONFIG_MAP[cfg]["protocol"], + "error": str(e), + "payload_mib": actual_mib, + "num_fields": nf, + } + ) + + # Summary table + print("\n" + "=" * 80) + print("SUMMARY (median throughput)") + print( + f"{'Config':<6} {'Backend':<16} {'Proto':<6} {'Payload':>8} {'Fields':>7} " + f"{'put ms':>8} {'put GB/s':>9} {'get ms':>8} {'get GB/s':>9}" + ) + print("-" * 80) + for r in all_results: + if "error" in r: + print(f"{r['config']:<6} {r['backend']:<16} {r['protocol']:<6} {'ERR':>8}") + continue + actual_mib = r["nbytes"] / 1024 / 1024 + print( + f"{r['config']:<6} {r['backend']:<16} {r['protocol']:<6} " + f"{actual_mib:>7.1f}M {'?':>7} " + f"{r['put_ms_median']:>8.1f} {r['put_gbs_median']:>9.2f} " + f"{r['get_ms_median']:>8.1f} {r['get_gbs_median']:>9.2f}" + ) + + # RDMA net benefit if both C1 and C2 present + c1_results = [r for r in all_results if r["config"] == "C1" and "error" not in r] + c2_results = [r for r in all_results if r["config"] == "C2" and "error" not in r] + if c1_results and c2_results: + print("\n--- RDMA net benefit (C2 - C1, same backend, protocol only) ---") + for c1, c2 in zip(c1_results, c2_results): + put_delta = c2["put_gbs_median"] - c1["put_gbs_median"] + get_delta = c2["get_gbs_median"] - c1["get_gbs_median"] + put_pct = (put_delta / c1["put_gbs_median"] * 100) if c1["put_gbs_median"] > 0 else 0 + get_pct = (get_delta / c1["get_gbs_median"] * 100) if c1["get_gbs_median"] > 0 else 0 + print(f" put: {put_delta:+.2f} GB/s ({put_pct:+.0f}%) get: {get_delta:+.2f} GB/s ({get_pct:+.0f}%)") + + # CSV output + if args.output_csv: + import csv + + with open(args.output_csv, "w", newline="") as f: + if all_results: + writer = csv.DictWriter(f, fieldnames=all_results[0].keys()) + writer.writeheader() + writer.writerows(all_results) + print(f"\nCSV written to {args.output_csv}") + + # Final cleanup: tear down the last config's controller so a subsequent + # benchmark run starts from a clean slate. + close_tq_and_wait() + print("\n[dataplane] benchmark complete") + + +if __name__ == "__main__": + main() diff --git a/tests/utils/test_rdma_probe.py b/tests/utils/test_rdma_probe.py new file mode 100644 index 000000000..b2cb0dc4f --- /dev/null +++ b/tests/utils/test_rdma_probe.py @@ -0,0 +1,338 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Unit tests for RDMA capability probe and config resolution. + +These tests are CPU-only and do NOT require a real TransferQueue or RDMA +hardware. They mock the filesystem and mooncake import to exercise every +branch of the probe, reduction, and config validation logic. +""" + +from __future__ import annotations + +import argparse +from unittest import mock + +import relax.utils.rdma_probe as rdma_probe +from relax.utils.rdma_probe import ( + CheckResult, + EffectiveConfig, + ProbeResult, + _degenerate_result, + _select_dataplane_node_ids, + probe_cluster_nodes, + probe_node, + reduce_results, + validate_config, +) +from relax.utils.tq_config import ( + build_mooncake_config, + build_simple_storage_config, + estimate_payload_bytes, + validate_segment_capacity, +) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_probe( + protocol: str | None = "rdma", + device: str = "rdma0", + gdr: bool = False, + node: str = "node-A", +) -> ProbeResult: + return ProbeResult( + node=node, + checks=(CheckResult("mooncake_import", True),), + effective_protocol=protocol, + effective_device=device, + gdr_eligible=gdr, + ) + + +def _make_args(**kwargs) -> argparse.Namespace: + defaults = dict( + tq_storage_backend="mooncake", + tq_rdma_mode="auto", + tq_rdma_device="", + tq_use_gdr=False, + num_data_storage_units=1, + max_staleness=0, + n_samples_per_prompt=1, + rollout_batch_size=32, + multimodal_keys=None, + ) + defaults.update(kwargs) + return argparse.Namespace(**defaults) + + +# --------------------------------------------------------------------------- +# validate_config +# --------------------------------------------------------------------------- + + +class TestValidateConfig: + """validate_config: structural flag-combination checks before any probe.""" + + def test_simple_backend_with_rdma_mode_rejected(self): + args = _make_args(tq_storage_backend="simple", tq_rdma_mode="auto") + errors = validate_config(args) + assert len(errors) == 1 + assert "simple" in errors[0] + + def test_simple_backend_with_gdr_rejected(self): + args = _make_args(tq_storage_backend="simple", tq_rdma_mode="off", tq_use_gdr=True) + errors = validate_config(args) + assert any("--tq-use-gdr" in e for e in errors) + + def test_gdr_without_rdma_rejected(self): + args = _make_args(tq_storage_backend="mooncake", tq_rdma_mode="off", tq_use_gdr=True) + errors = validate_config(args) + assert any("rdma-mode=off" in e for e in errors) + + def test_valid_simple_off(self): + args = _make_args(tq_storage_backend="simple", tq_rdma_mode="off") + assert validate_config(args) == [] + + def test_valid_mooncake_auto(self): + args = _make_args(tq_storage_backend="mooncake", tq_rdma_mode="auto") + assert validate_config(args) == [] + + def test_valid_mooncake_required_gdr(self): + args = _make_args(tq_storage_backend="mooncake", tq_rdma_mode="required", tq_use_gdr=True) + assert validate_config(args) == [] + + +# --------------------------------------------------------------------------- +# reduce_results +# --------------------------------------------------------------------------- + + +class TestReduceResults: + """reduce_results: per-node ProbeResult -> job-level EffectiveConfig (AND reduction).""" + + def test_simple_backend_short_circuits(self): + eff = reduce_results( + [_make_probe()], + requested_backend="simple", + requested_device="rdma0", + use_gdr=False, + ) + assert eff.backend == "SimpleStorage" + assert eff.protocol == "tcp" + assert eff.gdr is False + + def test_all_nodes_rdma(self): + eff = reduce_results( + [_make_probe(protocol="rdma"), _make_probe(protocol="rdma", node="node-B")], + requested_backend="mooncake", + requested_device="", + use_gdr=False, + ) + assert eff.backend == "MooncakeStore" + assert eff.protocol == "rdma" + assert eff.fallback_reason == "" + + def test_one_node_no_mooncake_falls_back(self): + eff = reduce_results( + [_make_probe(protocol="rdma"), _make_probe(protocol=None, node="node-B")], + requested_backend="mooncake", + requested_device="", + use_gdr=False, + ) + assert eff.backend == "SimpleStorage" + assert "node-B" in eff.fallback_reason + + def test_one_node_no_rdma_degrades_to_tcp(self): + eff = reduce_results( + [_make_probe(protocol="rdma"), _make_probe(protocol="tcp", node="node-B")], + requested_backend="mooncake", + requested_device="", + use_gdr=False, + ) + assert eff.backend == "MooncakeStore" + assert eff.protocol == "tcp" + assert "node-B" in eff.fallback_reason + + def test_gdr_eligible_only_when_all_nodes(self): + eff = reduce_results( + [_make_probe(gdr=True), _make_probe(gdr=False, node="node-B")], + requested_backend="mooncake", + requested_device="", + use_gdr=True, + ) + assert eff.gdr is False + assert "gdr_cuda_not_initialized" in eff.fallback_reason + + def test_gdr_eligible_all_nodes(self): + eff = reduce_results( + [_make_probe(gdr=True), _make_probe(gdr=True, node="node-B")], + requested_backend="mooncake", + requested_device="", + use_gdr=True, + ) + assert eff.gdr is True + + def test_empty_results_falls_back(self): + eff = reduce_results( + [], + requested_backend="mooncake", + requested_device="", + use_gdr=False, + ) + assert eff.backend == "SimpleStorage" + + +# --------------------------------------------------------------------------- +# probe_node (mocked filesystem) +# --------------------------------------------------------------------------- + + +class TestProbeNode: + """probe_node: per-node capability checks (mocked /sys and mooncake).""" + + def test_no_infiniband_dir_gives_tcp(self): + """When /sys/class/infiniband doesn't exist, protocol should be tcp (if + mooncake imports) or None (if not).""" + with mock.patch("os.path.isdir", return_value=False): + with mock.patch("relax.utils.rdma_probe._check_mooncake_import") as mi: + mi.return_value = CheckResult("mooncake_import", True, "version=0.3.10") + result = probe_node("") + # mooncake importable but no RDMA device → tcp + assert result.effective_protocol == "tcp" + + def test_active_rdma_device_gives_rdma(self): + """When all checks pass, protocol should be rdma.""" + + def fake_isdir(path): + return "infiniband" in path + + with ( + mock.patch("os.path.isdir", side_effect=fake_isdir), + mock.patch("os.listdir", return_value=["rdma0"]), + mock.patch("builtins.open", mock.mock_open(read_data="4: ACTIVE")), + mock.patch("relax.utils.rdma_probe._check_mooncake_import") as mi, + mock.patch("relax.utils.rdma_probe._check_health_check") as hc, + mock.patch("relax.utils.rdma_probe.resource.getrlimit", return_value=(-1, -1)), + ): + mi.return_value = CheckResult("mooncake_import", True, "ok") + hc.return_value = CheckResult("health_check", True, "return_code=0") + result = probe_node("") + assert result.effective_protocol == "rdma" + assert result.ok + + +# --------------------------------------------------------------------------- +# probe_cluster_nodes (multi-node fan-out) + helpers +# --------------------------------------------------------------------------- + + +class TestProbeClusterNodes: + """probe_cluster_nodes: multi-node fan-out helpers + degenerate-result handling.""" + + def test_select_nodes_filters_dead_and_cpu_only(self): + """Only alive nodes advertising GPU resources are data-plane nodes.""" + nodes = [ + {"NodeID": "n0", "Alive": True, "Resources": {"GPU": 8.0}}, + {"NodeID": "n1", "Alive": True, "Resources": {"GPU": 0}}, + {"NodeID": "n2", "Alive": False, "Resources": {"GPU": 8.0}}, + {"NodeID": "n3", "Alive": True, "Resources": {}}, + ] + assert _select_dataplane_node_ids(nodes) == ["n0"] + + def test_degenerate_result_is_no_mooncake(self): + """A failed/timed-out node reports effective_protocol=None so the AND- + reducer degrades instead of silently dropping the node.""" + r = _degenerate_result("node-X", "probe_timeout:60s") + assert r.node == "node-X" + assert r.effective_protocol is None + assert r.ok is False + assert "probe_timeout" in r.errors[0] + + def test_cluster_falls_back_to_local_when_no_gpu_nodes(self, monkeypatch): + """No alive GPU workers (single-node / local dev) -> probe driver only, + never touching Ray remote scheduling.""" + monkeypatch.setattr(rdma_probe, "_alive_gpu_nodes", lambda: []) + local = _make_probe(protocol="rdma", node="local-driver") + monkeypatch.setattr(rdma_probe, "probe_node", lambda dev: local) + results = probe_cluster_nodes("") + assert len(results) == 1 + assert results[0] is local + + def test_reduce_treats_degenerate_as_no_mooncake(self): + """A probe failure on one node forces job-level fallback (not a silent + drop that would over-report RDMA readiness).""" + results = [ + _make_probe(protocol="rdma", node="n0"), + _degenerate_result("n1", "probe_task_failed:boom"), + ] + eff = reduce_results( + results, + requested_backend="mooncake", + requested_device="", + use_gdr=False, + ) + assert eff.backend == "SimpleStorage" + assert "n1" in eff.fallback_reason + + +# --------------------------------------------------------------------------- +# tq_config builders +# --------------------------------------------------------------------------- + + +class TestTqConfigBuilder: + """tq_config builders: SimpleStorage/MooncakeStore dict + capacity + validation.""" + + def test_simple_storage_config(self): + cfg = build_simple_storage_config(total_storage_size=1000, num_data_storage_units=2) + assert cfg == { + "storage_backend": "SimpleStorage", + "SimpleStorage": {"total_storage_size": 1000, "num_data_storage_units": 2}, + } + + def test_storage_backend_key_selects_the_manager(self): + """``tq.init`` reads ``backend.storage_backend``; omitting it silently + keeps SimpleStorage.""" + eff = EffectiveConfig(backend="MooncakeStore", protocol="rdma", device="rdma0", gdr=False, fallback_reason="") + assert build_mooncake_config(eff)["storage_backend"] == "MooncakeStore" + assert ( + build_simple_storage_config(total_storage_size=1, num_data_storage_units=1)["storage_backend"] + == "SimpleStorage" + ) + + def test_mooncake_config_has_hard_pin_true(self): + eff = EffectiveConfig(backend="MooncakeStore", protocol="rdma", device="rdma0", gdr=False, fallback_reason="") + cfg = build_mooncake_config(eff) + mc = cfg["MooncakeStore"] + assert mc["protocol"] == "rdma" + assert mc["device_name"] == "rdma0" + assert mc["hard_pin"] is True # no silent eviction + assert mc["auto_init"] is False # external master + assert mc["use_gdr"] is False + + def test_mooncake_config_gdr_propagated(self): + eff = EffectiveConfig(backend="MooncakeStore", protocol="rdma", device="", gdr=True, fallback_reason="") + cfg = build_mooncake_config(eff) + assert cfg["MooncakeStore"]["use_gdr"] is True + + def test_segment_capacity_text_only_passes(self): + args = _make_args(multimodal_keys=None) + eff = EffectiveConfig(backend="MooncakeStore", protocol="rdma", device="", gdr=False, fallback_reason="") + assert validate_segment_capacity(args, eff) is None + + def test_segment_capacity_multimodal_large_batch_fails(self): + args = _make_args( + multimodal_keys=["pixel_values"], rollout_batch_size=256, n_samples_per_prompt=8, max_staleness=1 + ) + eff = EffectiveConfig(backend="MooncakeStore", protocol="rdma", device="", gdr=False, fallback_reason="") + err = validate_segment_capacity(args, eff) + assert err is not None + assert "insufficient" in err.lower() + + def test_estimate_payload_text_only_is_zero(self): + args = _make_args(multimodal_keys=None) + assert estimate_payload_bytes(args) == 0 diff --git a/tests/utils/test_tq_dataplane_behavior.py b/tests/utils/test_tq_dataplane_behavior.py new file mode 100644 index 000000000..fdfa5e7e5 --- /dev/null +++ b/tests/utils/test_tq_dataplane_behavior.py @@ -0,0 +1,296 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""TransferQueue data-plane behavior + byte-exact consistency tests. + +These are **integration tests**: they spin up a real TransferQueue over +SimpleStorage (ZMQ/TCP) inside a local Ray cluster and exercise the full +put -> get round-trip plus six storage-behavior contracts. They do NOT need +GPU or RDMA — SimpleStorage runs as Ray actors with in-process ZMQ servers, +so they are single-node and CI-runnable wherever a Ray cluster can start. + +Behaviors verified (mapped to the RFC's six required behaviors): + + 1. connection -- tq.init + get_client yield a usable client (put/get). + 2. byte-exact -- every put tensor returns byte-identical via NestedTensor + .values(), including a realistic multimodal column count. + 3. backpressure-- a single put exceeding capacity raises RuntimeError + ("Storage capacity exceeded") rather than silently dropping. + 4. empty-get -- get on an empty partition returns size==0 + empty TensorDict + without hanging (the consumer-side "no data yet" contract). + 5. retry -- re-putting the same partition overwrites; get returns latest. + 6. cleanup -- clear_partition empties data; close+reinit yields a fresh, + isolated controller (exercises the F10 anti-hang path). + +A true cross-node disconnect (consumer node death mid-get) is NOT covered here +-- it requires multi-node GPU hardware and is skipped per project rules. +""" + +from __future__ import annotations + +import importlib.util +import time + +import pytest +import torch + + +def _has_real_submodule(dotted: str) -> bool: + """True only if a REAL transfer_queue package is installed. + + CI installs a single-file ``transfer_queue`` stub whose ``__getattr__`` + returns a dummy for any attribute, so ``find_spec("transfer_queue")`` is + True on CI even though no real submodule exists. Probing a real submodule + (``transfer_queue.storage``) returns None for the stub, so these + integration tests skip on CPU CI and run only where real TransferQueue + + Ray are installed. + """ + try: + return importlib.util.find_spec(dotted) is not None + except (ImportError, ValueError, TypeError): + # CI's single-file transfer_queue stub returns a dummy for ``__path__``, + # so find_spec on a submodule raises TypeError instead of returning None. + return False + + +_TQ_OK = _has_real_submodule("transfer_queue.storage") +_RAY_OK = importlib.util.find_spec("ray") is not None + +pytestmark = pytest.mark.skipif( + not (_TQ_OK and _RAY_OK), + reason=( + "TransferQueue data-plane tests require the `transfer_queue` and `ray` " + "packages plus a startable local Ray cluster (SimpleStorage; no GPU/RDMA)." + ), +) + +_TQ_ACTOR = "TransferQueueController" +_TQ_NS = "transfer_queue" + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _wait_controller_gone(timeout: float = 20.0) -> bool: + """Poll until the named TQ controller is gone from the Ray GCS. + + Required between tests: after ``tq.close()`` kills the actor, its handle is + still resolvable for a short window. Re-init during that window would + attach to a *dead* controller (ActorDiedError) — the F10 race. + """ + import ray + + deadline = time.time() + timeout + while time.time() < deadline: + try: + ray.get_actor(_TQ_ACTOR, namespace=_TQ_NS) + except ValueError: + return True + time.sleep(0.4) + return False + + +def _force_kill_controller() -> None: + import ray + + try: + ray.kill(ray.get_actor(_TQ_ACTOR, namespace=_TQ_NS)) + except ValueError: + pass + + +def _flat_values(t): + """Flatten a dense tensor or a NestedTensor to a 1-D comparable view. + + TransferQueue returns per-sample data as ``NestedTensor`` (shape ``(N, + j0)``); ``.values()`` is the row-major concatenated storage, which is what + byte-exact comparison must use. Dense inputs flatten identically. + """ + if type(t).__name__ == "NestedTensor": + return t.values().reshape(-1) + return t.reshape(-1) + + +def _payload(n: int, fields: list[str], cols: int, dtype: str = "float32", seed: int = 0): + """Build a TensorDict of ``n`` samples with ``fields`` of shape (n, + cols).""" + from tensordict import TensorDict + + dt = getattr(torch, dtype) + g = torch.Generator().manual_seed(seed) + data = {f: torch.randn(n, cols, dtype=dt, generator=g) for f in fields} + return TensorDict(data, batch_size=[n]) + + +def _round_trip(client, payload, partition: str, fields: list[str], n: int): + """put -> get_data and return the retrieved TensorDict.""" + client.put(payload, partition_id=partition) + meta = client.get_meta(data_fields=fields, batch_size=n, partition_id=partition, mode="fetch", task_name=partition) + return client.get_data(meta) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def _ray_cluster(): + import ray + + ray.init(ignore_reinit_error=True, logging_level="ERROR") + yield + try: + ray.shutdown() + except Exception: + pass + + +@pytest.fixture +def tq_factory(_ray_cluster): + """Yield ``reinit(capacity, units=1) -> client``; tears down after each + test. + + Each call closes the prior TQ, waits for the controller to leave the GCS + (F10-safe), and starts a fresh controller with the requested capacity. + """ + import transfer_queue as tq + from omegaconf import OmegaConf + from transfer_queue import GRPOGroupNSampler + + def _reinit(capacity: int = 1024, units: int = 1): + tq.close() + if not _wait_controller_gone(): + _force_kill_controller() + _wait_controller_gone() + conf = OmegaConf.create( + { + "controller": { + "sampler": GRPOGroupNSampler(n_samples_per_prompt=1), + "polling_mode": True, + }, + "backend": { + "SimpleStorage": { + "total_storage_size": capacity, + "num_data_storage_units": units, + } + }, + }, + flags={"allow_objects": True}, + ) + tq.init(conf=conf) + return tq.get_client() + + yield _reinit + + # Final teardown: close + ensure controller gone so the next test/module + # starts clean. + tq.close() + _wait_controller_gone() + _force_kill_controller() + _wait_controller_gone() + + +# --------------------------------------------------------------------------- +# Behavior + byte-exact tests +# --------------------------------------------------------------------------- + + +class TestTqDataPlaneBehavior: + def test_connection_establishment(self, tq_factory): + """tq.init + get_client produce a client that can put and get.""" + client = tq_factory() + payload = _payload(n=4, fields=["a", "b"], cols=8) + got = _round_trip(client, payload, "conn", ["a", "b"], 4) + assert "a" in got.keys() and "b" in got.keys() + + def test_byte_exact_consistency(self, tq_factory): + """Every put tensor returns byte-identical (NestedTensor .values()).""" + client = tq_factory() + fields = ["img", "txt", "mask"] + payload = _payload(n=8, fields=fields, cols=16, seed=42) + got = _round_trip(client, payload, "be", fields, 8) + for k in fields: + gv, av = _flat_values(got[k]), _flat_values(payload[k]) + assert gv.numel() == av.numel(), f"{k}: numel {gv.numel()} != {av.numel()}" + assert gv.dtype == av.dtype + # Elementwise identity is the byte-exact contract. + assert torch.equal(gv, av), f"{k}: not byte-exact" + + def test_byte_exact_multimodal_column_count(self, tq_factory): + """Realistic multimodal hidden dim (1176, Qwen3-VL) survives round- + trip. + + Uses a bounded token count to keep the test fast; the point is that a + production column count round-trips byte-exactly through NestedTensor. + """ + client = tq_factory() + fields = ["pixel_values"] + payload = _payload(n=4, fields=fields, cols=1176, seed=7) + got = _round_trip(client, payload, "mm", fields, 4) + gv, av = _flat_values(got["pixel_values"]), _flat_values(payload["pixel_values"]) + assert gv.numel() == av.numel() + assert torch.equal(gv, av) + + def test_backpressure_raises_on_capacity_overflow(self, tq_factory): + """A single put exceeding capacity raises rather than silently + dropping.""" + client = tq_factory(capacity=4) + payload = _payload(n=8, fields=["a"], cols=4) # 8 samples > capacity 4 + with pytest.raises(RuntimeError, match="capacity"): + client.put(payload, partition_id="bp") + # Nothing was stored -> a subsequent get reports size 0 (no data). + meta = client.get_meta(data_fields=["a"], batch_size=8, partition_id="bp", mode="fetch", task_name="bp") + assert getattr(meta, "size", None) == 0 + + def test_empty_get_returns_zero_size_without_hanging(self, tq_factory): + """get on an empty partition returns size==0 + empty TensorDict (no + hang).""" + client = tq_factory() + meta = client.get_meta(data_fields=["a"], batch_size=4, partition_id="empty", mode="fetch", task_name="empty") + assert getattr(meta, "size", None) == 0 + data = client.get_data(meta) # empty TensorDict; must not KeyError on access + assert len(list(data.keys())) == 0 + + def test_repeat_put_same_partition_is_safe(self, tq_factory): + """Re-putting the same partition id is safe: no crash, no duplication, + no corruption. + + TQ's overwrite-vs-sample semantics are sampler-dependent, so we assert + the stable, observable contract -- the partition stays bounded at N and + every returned sample is byte-identical to a sample we actually put + (never garbage). + """ + client = tq_factory() + first = _payload(n=4, fields=["a"], cols=4, seed=1) + second = _payload(n=4, fields=["a"], cols=4, seed=2) + client.put(first, partition_id="rp") + client.put(second, partition_id="rp") # must not crash or hang + meta = client.get_meta(data_fields=["a"], batch_size=4, partition_id="rp", mode="fetch", task_name="rp") + assert getattr(meta, "size", None) == 4 # bounded; not duplicated to 8 + got = client.get_data(meta) + gv = _flat_values(got["a"]) + assert gv.numel() == 16 + # Every returned row matches some row we put (first or second); order may + # differ due to sampling, but no row may be corrupted. + got_rows = gv.reshape(4, 4) + candidates = torch.cat([first["a"], second["a"]], dim=0) # (8, 4) + for i in range(4): + row = got_rows[i] + assert torch.any(torch.all(candidates == row, dim=1)), f"row {i} matches no put sample (corrupted)" + + def test_cleanup_clear_partition_then_reinit_isolated(self, tq_factory): + """clear_partition empties data; a reinit yields a fresh controller.""" + client = tq_factory() + client.put(_payload(n=4, fields=["a"], cols=4), partition_id="cp") + meta = client.get_meta(data_fields=["a"], batch_size=4, partition_id="cp", mode="fetch", task_name="cp") + assert getattr(meta, "size", None) == 4 + client.clear_partition("cp") + meta2 = client.get_meta(data_fields=["a"], batch_size=4, partition_id="cp", mode="fetch", task_name="cp2") + assert getattr(meta2, "size", None) == 0 + + # Reinit with a different capacity -> fresh controller, old partition gone. + client2 = tq_factory(capacity=16) + meta3 = client2.get_meta(data_fields=["a"], batch_size=4, partition_id="cp", mode="fetch", task_name="cp3") + assert getattr(meta3, "size", None) == 0 diff --git a/tests/utils/test_tq_failure_paths.py b/tests/utils/test_tq_failure_paths.py new file mode 100644 index 000000000..8ff560c42 --- /dev/null +++ b/tests/utils/test_tq_failure_paths.py @@ -0,0 +1,374 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Failure-path tests for the TransferQueue dataplane enablement. + +Covers the four gaps the maintainer review called out, which the existing +``test_rdma_probe.py`` (pure config/probe logic) and +``test_tq_dataplane_behavior.py`` (real TQ on SimpleStorage) did not: + +* timeout -- controller ``get_config`` timeout and probe-task timeout +* disconnect -- store errors surface instead of returning corrupt data +* retry -- ``batch_get_into`` / ``batch_upsert_from`` retry-then-raise +* byte-exactness on **MooncakeStore** (SimpleStorage-only before) +* automatic degradation as pytest (was a manual two-node script) +* the controller reaper / teardown helpers (now in ``relax.utils.tq_lifecycle``) + +Everything except the MooncakeStore round-trip runs with stubs, so it is +CI-safe; the round-trip skips unless a reachable mooncake master is configured. +""" + +from __future__ import annotations + +import importlib.util +import os +import socket +from unittest.mock import MagicMock + +import pytest +import torch + +from relax.utils import tq_lifecycle +from relax.utils.rdma_probe import ProbeResult, reduce_results + + +def _has_real_submodule(dotted: str) -> bool: + """True only if a REAL transfer_queue submodule is importable. + + CI installs a single-file ``transfer_queue`` stub; + ``transfer_queue.storage`` does not exist there, so tests that touch the + real MooncakeStoreClient skip on CPU CI and run only where real + TransferQueue is installed. + """ + try: + return importlib.util.find_spec(dotted) is not None + except (ImportError, ValueError, TypeError): + # CI's single-file transfer_queue stub returns a dummy for ``__path__``, + # so find_spec on a submodule raises TypeError instead of returning None. + return False + + +_REAL_MOONCAKE_CLIENT = _has_real_submodule("transfer_queue.storage.clients.mooncake_client") + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _probe(node: str, protocol: str | None = "rdma", device: str = "rdma0") -> ProbeResult: + """Build a ProbeResult without running any real probe.""" + return ProbeResult( + node=node, + checks=(), + effective_protocol=protocol, + effective_device=device if protocol else "", + gdr_eligible=protocol == "rdma", + errors=() if protocol else ("mooncake not importable",), + ) + + +def _master_address() -> str: + """Mooncake master address for the round-trip test.""" + return os.environ.get("MC_MASTER_ADDRESS", "127.0.0.1:50051") + + +def _master_reachable(timeout: float = 1.0) -> bool: + """True if something accepts TCP connections on the master address.""" + host, _, port = _master_address().rpartition(":") + try: + with socket.create_connection((host, int(port)), timeout=timeout): + return True + except OSError: + return False + + +# --------------------------------------------------------------------------- +# Controller lifecycle: reaper (timeout / half-initialised / healthy) +# --------------------------------------------------------------------------- + + +class TestReapUnusableController: + """reap_unusable_tq_controller: only unusable controllers get killed.""" + + @staticmethod + def _fake_ray(monkeypatch, *, actor, get_result=None, get_raises=None): + """Stub the ray module used by tq_lifecycle; record kill calls.""" + killed: list = [] + fake = MagicMock() + if actor is None: + fake.get_actor.side_effect = ValueError("actor not found") + else: + fake.get_actor.return_value = actor + if get_raises is not None: + fake.get.side_effect = get_raises + else: + fake.get.return_value = get_result + fake.kill.side_effect = lambda handle: killed.append(handle) + monkeypatch.setattr(tq_lifecycle, "ray", fake) + monkeypatch.setattr(tq_lifecycle, "kill_tq_controller_and_wait", lambda *a, **k: killed.append("killed")) + return killed + + def test_no_controller_is_a_noop(self, monkeypatch): + killed = self._fake_ray(monkeypatch, actor=None) + assert tq_lifecycle.reap_unusable_tq_controller() is False + assert killed == [] + + def test_healthy_controller_is_left_alone(self, monkeypatch): + killed = self._fake_ray(monkeypatch, actor=MagicMock(), get_result={"backend": {}}) + assert tq_lifecycle.reap_unusable_tq_controller() is False + assert killed == [] + + def test_half_initialised_controller_is_reaped(self, monkeypatch): + """conf is None == actor created but store_config never ran (F10).""" + killed = self._fake_ray(monkeypatch, actor=MagicMock(), get_result=None) + assert tq_lifecycle.reap_unusable_tq_controller() is True + assert killed == ["killed"] + + def test_get_config_timeout_is_reaped(self, monkeypatch): + """An unresponsive controller must not turn tq.init into a hang.""" + killed = self._fake_ray(monkeypatch, actor=MagicMock(), get_raises=TimeoutError("get_config timed out")) + assert tq_lifecycle.reap_unusable_tq_controller() is True + assert killed == ["killed"] + + def test_dead_actor_is_reaped(self, monkeypatch): + killed = self._fake_ray(monkeypatch, actor=MagicMock(), get_raises=RuntimeError("ActorDiedError")) + assert tq_lifecycle.reap_unusable_tq_controller() is True + assert killed == ["killed"] + + +# --------------------------------------------------------------------------- +# Controller lifecycle: teardown unmounts the Mooncake segment +# --------------------------------------------------------------------------- + + +class TestCloseTqAndUnmount: + """close_tq_and_unmount: tq.close() first, then unmount the segment.""" + + @staticmethod + def _fake_tq(monkeypatch, store_client): + calls: list[str] = [] + fake = MagicMock() + manager = MagicMock() + if store_client is None: + del manager.storage_client # SimpleStorage manager has no storage_client + else: + manager.storage_client = store_client + fake.get_client.return_value = MagicMock(storage_manager=manager) + fake.close.side_effect = lambda: calls.append("tq.close") + monkeypatch.setattr(tq_lifecycle, "tq", fake) + return calls + + def test_mooncake_segment_is_unmounted_after_close(self, monkeypatch): + store_client = MagicMock() + calls = self._fake_tq(monkeypatch, store_client=store_client) + store_client.close.side_effect = lambda: calls.append("store.close") + tq_lifecycle.close_tq_and_unmount() + # Order matters: tq.close() still needs the store alive for remove_all(). + assert calls == ["tq.close", "store.close"] + + def test_simple_storage_teardown_is_noop_beyond_close(self, monkeypatch): + calls = self._fake_tq(monkeypatch, store_client=None) + tq_lifecycle.close_tq_and_unmount() + assert calls == ["tq.close"] + + def test_uninitialised_tq_does_not_raise(self, monkeypatch): + fake = MagicMock() + fake.get_client.side_effect = AssertionError("Please initialize the TransferQueue first") + monkeypatch.setattr(tq_lifecycle, "tq", fake) + tq_lifecycle.close_tq_and_unmount() # must not raise + fake.close.assert_called_once() + + +# --------------------------------------------------------------------------- +# Retry / disconnect on the MooncakeStore data path +# --------------------------------------------------------------------------- + + +class _FlakyStore: + """Stub mooncake store: the first ``fail_times`` calls return error + codes.""" + + def __init__(self, fail_times: int, code: int = -800, raise_exc: Exception | None = None): + self.fail_times = fail_times + self.code = code + self.raise_exc = raise_exc + self.get_calls: list[list[str]] = [] + self.put_calls: list[list[str]] = [] + + def _codes(self, keys): + if self.raise_exc is not None: + raise self.raise_exc + if self.fail_times > 0: + self.fail_times -= 1 + return [self.code] * len(keys) + return [0] * len(keys) + + def batch_get_into(self, keys, ptrs, sizes): + self.get_calls.append(list(keys)) + return self._codes(keys) + + def batch_upsert_from(self, keys, ptrs, sizes, config=None): + self.put_calls.append(list(keys)) + return self._codes(keys) + + +def _client_with_store(store) -> object: + """A MooncakeStoreClient with only ``_store``/``replica_config`` wired up. + + ``__init__`` is skipped on purpose: it would need a live mooncake master. + """ + from transfer_queue.storage.clients.mooncake_client import MooncakeStoreClient + + client = object.__new__(MooncakeStoreClient) + client._store = store + client.replica_config = None + return client + + +@pytest.mark.skipif( + not _REAL_MOONCAKE_CLIENT, + reason="needs a real transfer_queue (CI uses a single-file stub); run on a host with TransferQueue installed", +) +class TestRetryAndDisconnect: + """batch_get_into / batch_upsert_from: retry, then raise loudly.""" + + def test_get_retries_then_succeeds(self, monkeypatch): + monkeypatch.setattr("transfer_queue.storage.clients.mooncake_client.RETRY_DELAY_SECONDS", 0) + store = _FlakyStore(fail_times=2) + client = _client_with_store(store) + client._batch_get_into_with_retry(["0@f0", "1@f0"], [1, 2], [8, 8]) + assert len(store.get_calls) == 3 # initial + 2 retries + assert store.get_calls[-1] == ["0@f0", "1@f0"] + + def test_get_raises_after_max_retries(self, monkeypatch): + monkeypatch.setattr("transfer_queue.storage.clients.mooncake_client.RETRY_DELAY_SECONDS", 0) + client = _client_with_store(_FlakyStore(fail_times=99)) + with pytest.raises(RuntimeError, match="batch_get_into failed"): + client._batch_get_into_with_retry(["0@f0"], [1], [8]) + + def test_put_retries_then_succeeds(self, monkeypatch): + monkeypatch.setattr("transfer_queue.storage.clients.mooncake_client.RETRY_DELAY_SECONDS", 0) + store = _FlakyStore(fail_times=1, code=-1) + client = _client_with_store(store) + client._batch_upsert_with_retry(["0@f0"], [1], [8]) + assert len(store.put_calls) == 2 + + def test_put_raises_after_max_retries(self, monkeypatch): + monkeypatch.setattr("transfer_queue.storage.clients.mooncake_client.RETRY_DELAY_SECONDS", 0) + client = _client_with_store(_FlakyStore(fail_times=99, code=-1)) + with pytest.raises(RuntimeError, match="batch_upsert_from failed"): + client._batch_upsert_with_retry(["0@f0"], [1], [8]) + + def test_disconnect_surfaces_instead_of_returning_garbage(self): + """A dead peer must raise, never hand back a silently short buffer.""" + exc = RuntimeError("Failed to open segment for endpoint=':16675'") + client = _client_with_store(_FlakyStore(fail_times=0, raise_exc=exc)) + with pytest.raises(RuntimeError, match="Failed to open segment"): + client._batch_get_into_with_retry(["0@f0"], [1], [8]) + + +# --------------------------------------------------------------------------- +# Automatic degradation (was the manual two-node fault_inject_multinode.py) +# --------------------------------------------------------------------------- + + +class TestAutomaticDegradation: + """AND-reduction turns any node's failure into a job-level downgrade.""" + + def test_all_nodes_rdma_stays_rdma(self): + eff = reduce_results( + [_probe("a"), _probe("b")], requested_backend="mooncake", requested_device="", use_gdr=False + ) + assert (eff.backend, eff.protocol, eff.fallback_reason) == ("MooncakeStore", "rdma", "") + + def test_one_node_without_mooncake_degrades_whole_job(self): + """Mirrors the PYTHONPATH-poisoning case of the two-node script.""" + eff = reduce_results( + [_probe("a"), _probe("b", protocol=None)], + requested_backend="mooncake", + requested_device="", + use_gdr=False, + ) + assert eff.backend == "SimpleStorage" + assert "mooncake_unavailable" in eff.fallback_reason and "b" in eff.fallback_reason + + def test_crashed_probe_task_degrades_whole_job(self): + """A probe task that raises becomes a degenerate result, never + dropped.""" + from relax.utils.rdma_probe import _degenerate_result + + degenerate = _degenerate_result("b", "probe task raised") + assert degenerate.effective_protocol is None + eff = reduce_results( + [_probe("a"), degenerate], requested_backend="mooncake", requested_device="", use_gdr=False + ) + assert eff.backend == "SimpleStorage" + + def test_one_node_tcp_only_degrades_transport_not_backend(self): + eff = reduce_results( + [_probe("a"), _probe("b", protocol="tcp")], + requested_backend="mooncake", + requested_device="", + use_gdr=False, + ) + assert (eff.backend, eff.protocol) == ("MooncakeStore", "tcp") + assert eff.fallback_reason + + +# --------------------------------------------------------------------------- +# Byte-exactness on MooncakeStore (was SimpleStorage-only) +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif( + not (_master_reachable() and _REAL_MOONCAKE_CLIENT), + reason=( + "needs a reachable mooncake master and a real TransferQueue install " + "(CI uses a single-file transfer_queue stub and has no RDMA/mooncake " + "deployment), so the MooncakeStore round-trip is skipped" + ), +) +class TestMooncakeByteExact: + """Real MooncakeStoreClient put/get round-trip, byte-for-byte.""" + + @staticmethod + def _client(protocol: str): + from transfer_queue.storage.clients.mooncake_client import MooncakeStoreClient + + return MooncakeStoreClient( + { + "protocol": protocol, + "device_name": os.environ.get("MC_RDMA_DEVICE", ""), + "master_server_address": _master_address(), + "metadata_server": "P2PHANDSHAKE", + "local_hostname": "", + "global_segment_size": 2 * 1024**3, + "local_buffer_size": 512 * 1024**2, + "hard_pin": True, + "use_gdr": False, + } + ) + + @pytest.mark.parametrize("protocol", ["tcp", "rdma"]) + def test_multi_dtype_shape_roundtrip_is_byte_exact(self, protocol): + tensors = { + "f32_2d": torch.randn(64, 1176, dtype=torch.float32), + "bf16_2d": torch.randn(32, 512).to(torch.bfloat16), + "i64_1d": torch.arange(4096, dtype=torch.int64), + "noncontig": torch.randn(128, 256).t(), # transposed == non-contiguous + } + client = self._client(protocol) + try: + keys = [f"bx_{protocol}_{name}" for name in tensors] + values = list(tensors.values()) + client.put(keys, values) + got = client.get( + keys, + shapes=[tuple(v.shape) for v in values], + dtypes=[v.dtype for v in values], + ) + for name, want, have in zip(tensors, values, got, strict=True): + assert have is not None, f"{name} came back empty" + assert torch.equal(have, want.contiguous()), f"{name} is not byte-exact" + finally: + client.close() From 75424b0e9df04acc63a321750548de6ab6da548e Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Thu, 13 Aug 2026 16:06:18 +0800 Subject: [PATCH 02/22] fix(data-plane): harden TQ RDMA lifecycle - 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 --- docs/draft/transfer_queue_rdma.md | 68 ++++- relax/backends/megatron/actor.py | 7 +- relax/components/actor.py | 7 +- relax/components/actor_fwd.py | 7 +- relax/components/advantages.py | 7 +- relax/components/critic.py | 7 +- relax/components/rollout.py | 7 +- relax/components/sft.py | 7 +- relax/core/controller.py | 101 ++++-- relax/distributed/ray/rollout.py | 9 +- relax/utils/rdma_probe.py | 88 ++++-- relax/utils/tq_config.py | 36 ++- relax/utils/tq_lifecycle.py | 305 +++++++++++++++++- scripts/benchmarks/tq_cross_node_bench.py | 351 ++++++++++++++------- tests/utils/test_rdma_probe.py | 45 ++- tests/utils/test_tq_failure_paths.py | 357 +++++++++++++++++++++- 16 files changed, 1198 insertions(+), 211 deletions(-) diff --git a/docs/draft/transfer_queue_rdma.md b/docs/draft/transfer_queue_rdma.md index a94895dcd..8ea003b6a 100644 --- a/docs/draft/transfer_queue_rdma.md +++ b/docs/draft/transfer_queue_rdma.md @@ -21,25 +21,25 @@ Relax 的数据面（rollout ↔ train 之间的样本传输）默认走 Transfe ### GDR 为实验性 -GDR 的可用性无法在启动时探测：探测跑在独立的 Ray task 里，该进程没有初始化 CUDA context，`torch.cuda.is_initialized()` 恒为 False，若在此判定会让 GDR 永远不可达。真实判定发生在每个 worker 的 TQ 客户端内部（`mooncake_client.py`），没有 CUDA context 时**静默回退到 host RDMA**并打 WARNING。 +GDR 的可用性无法由 driver 的启动探测代表：探测跑在独立 Ray task 中，该进程没有初始化 CUDA context；真正创建 staging buffer 的是每个 worker 的 TQ 客户端。因此首期不宣称 job 级 GDR 已验证，`required` 也不对 GDR 做 fail-fast。 -因此首期：`--tq-use-gdr` 标记为实验性，`required` 不对 GDR 做 fail-fast。要把 GDR 纳入分级降级，需要补一条 worker → driver 的能力回报通道，属于后续工作。 +每个 worker 附加 TQ 后都会记录两层信息：`requested=true` 表示用户请求开启；`status=host_rdma_fallback/enabled_unverified/inactive/unknown` 表示该 worker 的本地观察。即使本地 staging buffer 已创建也只记为 `enabled_unverified`，不等同于线上流量已经证明走 GDR。 ## 启动流程与降级 driver 在**第一次 `tq.init` 之前**完成探测并生成 job 级唯一的 effective config，其余组件（actor / critic / rollout / sft / advantages / actor_fwd）都读同一份，不各自决策。 1. 校验参数组合（例如 `simple` + `rdma-mode` 会被拒绝） -2. `probe_cluster_nodes()` 通过 Ray 把探测任务绑定到每个**存活且有 GPU** 的节点，各自读本机 `/sys` 与 mooncake 状态；超时或崩溃的节点转为退化结果，不静默丢弃 +2. `probe_cluster_nodes()` 通过 Ray 把探测任务绑定到每个**存活且有 GPU** 的节点，并额外探测 driver（driver 也会创建 Mooncake owner client）；各节点读取本机 `/sys`、mooncake 状态，并在 2 秒上限内检查外部 master 的 TCP 可达性；超时或崩溃的节点转为退化结果，不静默丢弃 3. `reduce_results()` 做 AND 归约：整个作业只能跑在最低共同能力上 4. `required` 模式下若发生任何回退，直接抛异常并打印每个节点的探测明细 降级阶梯： ``` -GDR → host RDMA （worker 运行时判定，静默回退 + WARNING） +GDR → host RDMA （worker 运行时判定，记录 requested 与实际状态） RDMA → Mooncake/TCP （任一节点无 RDMA 能力，或指定设备缺失） -Mooncake → SimpleStorage（任一节点 mooncake 不可导入，或 segment 容量不足） +Mooncake → SimpleStorage（任一节点 mooncake/master 不可用、运行时正确性契约不满足，或 segment 预检不足） ``` ## 启动日志怎么读 @@ -55,7 +55,7 @@ Mooncake → SimpleStorage（任一节点 mooncake 不可导入，或 segment [ok] port_state: ACTIVE [ok] gid: ... [ok] memlock: unlimited -[dataplane] backend=MooncakeStore protocol=rdma device=mlx5_bond_0 gdr=off +[dataplane] backend=MooncakeStore protocol=rdma device=mlx5_bond_0 gdr_requested=false gdr_status=off ``` 第三段带 `fallback=...` 就说明发生了降级，原因直接写在里面（例如 `fallback=mooncake_unavailable:`）。 @@ -72,12 +72,15 @@ setsid mooncake_master -rpc_port=50051 -metrics_port=9004 > /var/log/mooncake_ma 然后给作业设置 `MC_MASTER_ADDRESS=:50051`。未设置时内部默认 `localhost:50051`，仅适用于单节点开发。 +启动前置条件：部署侧必须先启动 master，所有 GPU 节点和 driver 都能解析并连接 `MC_MASTER_ADDRESS`，防火墙允许 master RPC 端口；作业镜像中的 TQ 必须包含本文“正确性依赖”所列修复。Relax 不负责拉起、重启或终止 master。 + 三种情形下的行为： | 情形 | 表现 | 处理 | |---|---|---| -| **初始化失败**（master 不可达） | TQ 客户端 `setup` 返回 `-1`，抛 `Mooncake store setup failed with error code: -1`。`required` 模式下作业直接失败；`auto` 模式下这一步已过了探测，不会自动回退——探测只验证本机能力，不验证 master 连通性 | 先确认 master 进程与 `MC_MASTER_ADDRESS`，再重启作业 | -| **正常退出**（作业结束或全局重启） | Relax 在拆数据面时调用 `close_tq_and_unmount()`：先 `tq.close()`（它内部还要用 store 做 `remove_all()`），再显式 `storage_client.close()` 卸载 segment 并从 master 注销。master 本身不动 | 无需操作 | +| **master 在探测时不可达** | `auto` 统一降级到 SimpleStorage；`required` 启动失败并列出失败节点 | 先确认 master、DNS/路由和 `MC_MASTER_ADDRESS` | +| **master 探测通过、但 `tq.init` 时失败/超时** | 第一次初始化在独立 owner actor 中执行，driver 最多等待 60 秒。失败后回收该 actor 及其拥有的半初始化 controller；`auto` 只重试一次 SimpleStorage，`required` 清理后抛出原始错误 | 查看 `mooncake_init_failed:*`、master 日志和 owner 清理日志 | +| **正常退出**（作业结束或全局重启） | 只有 owner actor 调用全局 `tq.close()`，随后显式 `storage_client.close()` 卸载 segment；附加 worker 只关闭本地 client，不能删除全局数据或 controller。master 本身不动 | 无需操作 | | **异常退出**（worker 被 kill / OOM / 节点掉线） | Python 层不执行，segment 仍在 master 注册。master 要等 `client_ttl`（默认 30 s）才判定客户端过期，期间新作业的 put 会打到死端点并报 `Failed to open segment ... Connection refused` | 等 30 s 后重启，或部署侧调小 `-client_ttl` | ## 资源所有权与安全清理 @@ -88,8 +91,57 @@ setsid mooncake_master -rpc_port=50051 -metrics_port=9004 > /var/log/mooncake_ma - 不使用任何 `pkill` / `killall` - `tq.init` 之前会检查已存在的 `TransferQueueController` 命名 actor：**只有取不到 config（半初始化）或 actor 已死时才回收**，健康的 controller 保持不动并正常 attach +- 首次初始化在专用 owner actor 中执行，config 内保存随机 owner token；清理只有在 token 匹配或 controller 确认不可用时才回收全局 actor，避免初始化竞争中的附加者误杀所有者 +- 全局 `tq.close()` 只能由 owner actor 调用；actor、critic、rollout 等附加 worker 只能做本地 detach - master 进程始终不被 Relax 触碰 +## 容量不足与正确性依赖 + +Mooncake 配置固定 `hard_pin=true`，不会为了腾空间静默驱逐已经生产但尚未消费的数据。启动前还会按 rollout batch、采样数、staleness 和多模态 payload 的固定启发式估算检查 4 GiB client segment；`auto` 预检不足时回退 SimpleStorage，`required` 直接失败。该估算只是早期保护，不能替代运行时错误处理，因为真实图片 patch 数会变化。 + +运行时依赖固定到 TransferQueue commit `58054a33834aadbcf76aacd6b1e32e25c030f2c9`，并在 Mooncake 启动前检查以下能力： + +- `batch_upsert_from` / `batch_get_into` 对每个 key 的返回码做有限次数重试，耗尽后抛异常，不能把失败当成功或无限重试； +- `KVStorageManager.put_data` 必须先等待 storage put 成功，之后才能更新 production-ready 状态；写入失败时不通知消费者； +- Relax 的契约测试用失败 store 验证“写失败、production 状态不更新”，并用隔离 master 的真机故障注入验证物理容量溢出在 30 秒内显式失败。 + +因此，上游曾出现的“返回码未检查导致静默丢数据”不是已知限制，而是 Mooncake 启用的硬门槛：`auto` 在契约不满足时禁用 Mooncake 并回退，`required` 拒绝启动。Docker 镜像固定上述修复 commit，运行时检查用于防止环境被旧包覆盖。 + +真机容量故障注入会故意创建 64 MiB segment 并写入 96 MiB，仅允许在独立、可丢弃的 master 上运行： + +```bash +MC_MASTER_ADDRESS=:50051 \ +RELAX_RUN_REAL_MOONCAKE_CAPACITY_TEST=1 \ +pytest -q tests/utils/test_tq_failure_paths.py \ + -k real_mooncake_capacity_overflow +``` + +## 验收分层 + +Mock/本机测试和真实双节点 RDMA 测试必须分别报告，前者不能替代后者。 + +| 层级 | 验证内容 | 通过标准 | +|---|---|---| +| CI/mock | 参数矩阵、节点 AND 归约、master 不可达、owner 超时/清理/token、auto/required、有限重试、写失败不发布状态 | `tests/utils/test_rdma_probe.py` 与 `tests/utils/test_tq_failure_paths.py` 全部通过；真机项允许明确 skip | +| 本机 TQ | SimpleStorage 全链路 put/get、容量 backpressure、空读、清理、字节一致性 | `tests/utils/test_tq_dataplane_behavior.py` 通过 | +| 真实 Mooncake | TCP/RDMA direct-client 多模态字段与混合 dtype 逐字节一致 | `TestMooncakeByteExact` 的 TCP/RDMA 两档均通过，不得 skip | +| 真实双节点 | 同一拓扑的 SimpleStorage、Mooncake/TCP、Mooncake/RDMA；synthetic 与 production-shaped multimodal；256/1024/2048/4096 MiB；每档 warmup + 至少 5 轮 | 每次 get 的逐字段 SHA-256 全部 PASS；`--require-wire-proof` 证明 RDMA 档 IB counter 增长且 TCP 档网络 counter 增长；CSV 留档并报告均值、median、stddev | + +双节点验收命令（master 与 Ray 集群需由部署侧预先准备）： + +```bash +PYTHONPATH=. python -u scripts/benchmarks/tq_cross_node_bench.py \ + --master :50051 \ + --nodeb-ip \ + --device \ + --payload-profiles synthetic multimodal \ + --payload-mib 256 1024 2048 4096 \ + --repeats 5 --require-wire-proof \ + --csv tq_cross_node_acceptance.csv +``` + +若本次开发环境没有两个 RDMA 节点，交付结论必须写成“真机验收未执行”，不能用 mock 通过推导真机已经通过。 + ## 排障表 | 现象 | 可能原因 | 处理 | diff --git a/relax/backends/megatron/actor.py b/relax/backends/megatron/actor.py index 75143efa6..4ce5c3943 100644 --- a/relax/backends/megatron/actor.py +++ b/relax/backends/megatron/actor.py @@ -13,7 +13,6 @@ import requests import torch import torch.distributed as dist -import transfer_queue as tq from megatron.core import mpu @@ -61,6 +60,7 @@ from relax.utils.rotate_ckpt import rotate_ckpt from relax.utils.s3_model_loader import prepare_model_maybe_update_args from relax.utils.timer import Timer, inverse_timer, timer, with_defer +from relax.utils.tq_lifecycle import attach_tq_client from relax.utils.tracking_utils import init_tracking from relax.utils.training import train_dump_utils from relax.utils.training.data_fields import build_data_fields @@ -187,8 +187,9 @@ def _init( init(args) if repatch is not None: repatch(args) - tq.init(args.tq_config) - self.data_system_client = tq.get_client() + self.data_system_client = attach_tq_client( + args.tq_config, requested_gdr=getattr(args, "tq_use_gdr", False), role=role + ) if is_megatron_main_rank(): init_tracking(args, primary=False) diff --git a/relax/components/actor.py b/relax/components/actor.py index c915cf90f..1b199bd43 100644 --- a/relax/components/actor.py +++ b/relax/components/actor.py @@ -7,7 +7,6 @@ from typing import Any, Dict, Optional import ray -import transfer_queue as tq from fastapi import FastAPI from ray import serve @@ -17,6 +16,7 @@ from relax.engine.sft.runtime import is_sft_mode, sft_partition_id, sft_task_name from relax.utils.async_utils import run from relax.utils.opd.opd_utils import set_managed_opd_teacher_on_train_group +from relax.utils.tq_lifecycle import attach_tq_client app = FastAPI() @@ -71,8 +71,9 @@ def __init__( self.actor_model = allocate_train_group(args=config, num_gpus=num_gpus, pg=pgs, runtime_env=runtime_env) - tq.init(self.config.tq_config) - self.data_system_client = tq.get_client() + self.data_system_client = attach_tq_client( + self.config.tq_config, requested_gdr=getattr(self.config, "tq_use_gdr", False), role=self.role + ) self.steps = ray.get( self.actor_model.async_init( diff --git a/relax/components/actor_fwd.py b/relax/components/actor_fwd.py index 707b8add6..7cd068fcb 100644 --- a/relax/components/actor_fwd.py +++ b/relax/components/actor_fwd.py @@ -6,12 +6,12 @@ from typing import Any, Optional import ray -import transfer_queue as tq from fastapi import FastAPI from ray import serve from relax.components.base import Base from relax.distributed.ray.placement_group import allocate_train_group +from relax.utils.tq_lifecycle import attach_tq_client app = FastAPI() @@ -35,8 +35,9 @@ def __init__( self._run_thread = None self._done_event: Optional[asyncio.Event] = None self._thread_error: Optional[Exception] = None - tq.init(self.config.tq_config) - self.data_system_client = tq.get_client() + self.data_system_client = attach_tq_client( + self.config.tq_config, requested_gdr=getattr(self.config, "tq_use_gdr", False), role=self.role + ) self.actor_model = allocate_train_group(args=config, num_gpus=num_gpus, pg=pgs, runtime_env=runtime_env) ray.get(self.actor_model.async_init(config, role=self.role, with_ref=False)) self.step = 0 diff --git a/relax/components/advantages.py b/relax/components/advantages.py index 751db235a..e9f7b8521 100644 --- a/relax/components/advantages.py +++ b/relax/components/advantages.py @@ -6,7 +6,6 @@ from typing import Any, Dict import torch -import transfer_queue as tq from megatron.core import mpu from ray import serve from tensordict import TensorDict @@ -17,6 +16,7 @@ apply_opd_to_advantages, consume_opd_advantage_data, ) +from relax.utils.tq_lifecycle import attach_tq_client from relax.utils.training.ppo_utils import ( compute_approx_kl, get_advantages_and_returns_batch, @@ -39,8 +39,9 @@ def __init__( self._lock = threading.RLock() self.healthy = healthy - tq.init(self.config.tq_config) - self.data_system_client = tq.get_client() + self.data_system_client = attach_tq_client( + self.config.tq_config, requested_gdr=getattr(self.config, "tq_use_gdr", False), role="advantages" + ) self.step = 0 async def run(self) -> None: diff --git a/relax/components/critic.py b/relax/components/critic.py index ef01bbf5f..c0519e562 100644 --- a/relax/components/critic.py +++ b/relax/components/critic.py @@ -7,7 +7,6 @@ from typing import Any, Optional import ray -import transfer_queue as tq from ray import serve from ray.serve.schema import LoggingConfig @@ -16,6 +15,7 @@ from relax.distributed.ray.placement_group import allocate_train_group from relax.engine.sft.runtime import sft_partition_id from relax.utils.async_utils import run +from relax.utils.tq_lifecycle import attach_tq_client @serve.deployment( @@ -40,8 +40,9 @@ def __init__( self.healthy = healthy self.role = role - tq.init(self.config.tq_config) - self.data_system_client = tq.get_client() + self.data_system_client = attach_tq_client( + self.config.tq_config, requested_gdr=getattr(self.config, "tq_use_gdr", False), role=self.role + ) self.critic_model = allocate_train_group( args=config, num_gpus=num_gpus, pg=pgs, role=self.role, runtime_env=runtime_env diff --git a/relax/components/rollout.py b/relax/components/rollout.py index ebf065683..b4b3e7352 100644 --- a/relax/components/rollout.py +++ b/relax/components/rollout.py @@ -9,7 +9,6 @@ import httpx import ray -import transfer_queue as tq from fastapi import FastAPI, HTTPException, Request from fastapi.responses import StreamingResponse from pydantic import BaseModel, Field @@ -20,6 +19,7 @@ from relax.distributed.ray.placement_group import create_rollout_manager from relax.utils.env import Envs from relax.utils.http_utils import _wrap_ipv6 +from relax.utils.tq_lifecycle import attach_tq_client app = FastAPI() @@ -333,8 +333,9 @@ def __init__( self.config = config self.healthy = healthy - tq.init(self.config.tq_config) - self.data_system_client = tq.get_client() + self.data_system_client = attach_tq_client( + self.config.tq_config, requested_gdr=getattr(self.config, "tq_use_gdr", False), role="rollout" + ) self.rollout_manager, self.num_rollout_per_epoch = create_rollout_manager( config, pg, data_source=data_source, runtime_env=runtime_env ) diff --git a/relax/components/sft.py b/relax/components/sft.py index 15ca2b15c..96da5b48d 100644 --- a/relax/components/sft.py +++ b/relax/components/sft.py @@ -29,7 +29,6 @@ import random from typing import Any -import transfer_queue as tq from ray import serve from transformers import AutoConfig, AutoTokenizer @@ -39,6 +38,7 @@ from relax.utils.data.processor_pool import ProcessorPool from relax.utils.misc import load_function from relax.utils.s3_model_loader import prepare_model_maybe_update_args +from relax.utils.tq_lifecycle import attach_tq_client from relax.utils.training.eval_config import build_named_prompt_data_configs from relax.utils.utils import dict_to_tensordict @@ -81,8 +81,9 @@ def __init__(self, healthy, pgs, num_gpus, config, role, runtime_env=None): # n self.healthy = healthy self.step = getattr(config, "start_rollout_id", 0) - tq.init(self.config.tq_config) - self.data_system_client = tq.get_client() + self.data_system_client = attach_tq_client( + self.config.tq_config, requested_gdr=getattr(self.config, "tq_use_gdr", False), role=self.role + ) self._dataset: Any | None = None self._eval_dataset: Any | None = None diff --git a/relax/core/controller.py b/relax/core/controller.py index f064184af..21f489503 100644 --- a/relax/core/controller.py +++ b/relax/core/controller.py @@ -8,7 +8,6 @@ from typing import Any import ray -import transfer_queue as tq from omegaconf import OmegaConf from ray import serve from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy @@ -49,8 +48,12 @@ ) from relax.utils.rdma_probe import probe_cluster_nodes, reduce_results, validate_config from relax.utils.s3_model_loader import cleanup_s3_model_weights_from_shm -from relax.utils.tq_config import build_backend_config -from relax.utils.tq_lifecycle import close_tq_and_unmount, kill_tq_controller_and_wait, reap_unusable_tq_controller +from relax.utils.tq_config import ( + build_backend_config, + resolve_mooncake_master_address, + validate_mooncake_runtime_contract, +) +from relax.utils.tq_lifecycle import close_tq_owner, initialize_tq_with_fallback from relax.utils.training.ppo_utils import validate_ppo_config from relax.utils.utils import compute_dp_size, recovery_load_path @@ -134,6 +137,7 @@ def __init__(self, config: Namespace, runtime_env: dict = None) -> None: self._pending_task_refs_lock = threading.Lock() if not hasattr(self, "_global_restart_count"): self._global_restart_count = 0 + self._tq_owner = None # SFT: fill in num_rollout / num_rollout_per_epoch before any actor # is launched (RL is resolved later in placement_group.py). @@ -274,18 +278,44 @@ def _initialize_data_system(self): else: sampler = GRPOGroupNSampler(n_samples_per_prompt=self.config.n_samples_per_prompt) + controller_config = { + "sampler": sampler, + "polling_mode": self.config.polling_mode, + } + backend_config = self._resolve_tq_backend(total_storage_size) tq_config = OmegaConf.create( { - "controller": { - "sampler": sampler, - "polling_mode": self.config.polling_mode, - }, - "backend": self._resolve_tq_backend(total_storage_size), + "controller": controller_config, + "backend": backend_config, }, flags={"allow_objects": True}, ) - tq_config = tq.init(conf=tq_config) or tq_config - self.config.tq_config = tq_config + + fallback_config = None + if backend_config.get("storage_backend") == "MooncakeStore": + from relax.utils.tq_config import build_simple_storage_config + + fallback_config = OmegaConf.create( + { + "controller": controller_config, + "backend": build_simple_storage_config( + total_storage_size=total_storage_size, + num_data_storage_units=self.config.num_data_storage_units, + ), + }, + flags={"allow_objects": True}, + ) + + 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 + self.config.tq_config = init_result.config + if init_result.fallback_reason: + logger.warning(f"[dataplane] effective backend=SimpleStorage fallback={init_result.fallback_reason}") + logger.info(f"[dataplane] controller ownership={'owner' if init_result.owns_controller else 'attached'}") def _resolve_tq_backend(self, total_storage_size: int) -> dict: """Resolve the TransferQueue ``backend`` config dict. @@ -318,7 +348,26 @@ def _resolve_tq_backend(self, total_storage_size: int) -> dict: # (probe_cluster_nodes), then AND-reduces to a single job-level # effective config so all data-plane workers converge identically. device = getattr(self.config, "tq_rdma_device", "") - probe_results = probe_cluster_nodes(device) + try: + validate_mooncake_runtime_contract() + except RuntimeError as e: + if mode == "required": + raise RuntimeError( + "--tq-rdma-mode=required but the installed TransferQueue " + f"does not satisfy the Mooncake correctness contract: {e}" + ) from e + from relax.utils.tq_config import build_simple_storage_config + + logger.warning( + "[dataplane] Installed TransferQueue does not satisfy the Mooncake " + f"correctness contract; auto fallback to SimpleStorage: {e}" + ) + return build_simple_storage_config( + total_storage_size=total_storage_size, + num_data_storage_units=self.config.num_data_storage_units, + ) + master_address = resolve_mooncake_master_address() + probe_results = probe_cluster_nodes(device, master_address) for r in probe_results: logger.debug(r.summary()) @@ -364,14 +413,7 @@ def _resolve_tq_backend(self, total_storage_size: int) -> dict: "--tq-rdma-mode=required does NOT fail fast on unavailable GDR." ) - # 8. F10 anti-hang: reap a half-initialised controller before tq.init, - # for every backend (a stale MooncakeStore controller used to survive - # because the guard only ran on the SimpleStorage fallback path). - # Only unusable controllers are killed -- a healthy one belongs to - # whoever created it and tq.init legitimately attaches to it. - self._reap_unusable_tq_controller() - - # 9. Log requested vs effective so the startup log alone explains the + # 8. Log requested vs effective so the startup log alone explains the # decision, plus one summary block per probed node. logger.info( f"[dataplane] requested: backend={backend} rdma_mode={mode} device={device or 'auto'} gdr={use_gdr}" @@ -385,21 +427,11 @@ def _resolve_tq_backend(self, total_storage_size: int) -> dict: logger.info(effective.log_line()) return backend_dict - def _reap_unusable_tq_controller(self) -> None: - """Delegate to - :func:`relax.utils.tq_lifecycle.reap_unusable_tq_controller`.""" - reap_unusable_tq_controller() - - def _kill_stale_tq_controller(self) -> None: - """Delegate to - :func:`relax.utils.tq_lifecycle.kill_tq_controller_and_wait`.""" - kill_tq_controller_and_wait() - - @staticmethod - def _close_data_system() -> None: + def _close_data_system(self) -> None: """Delegate to :func:`relax.utils.tq_lifecycle.close_tq_and_unmount`.""" - close_tq_and_unmount() + close_tq_owner(self._tq_owner) + self._tq_owner = None def _deploy_metrics_service(self): """Deploy the MetricsService as a lightweight Ray Serve deployment. @@ -918,6 +950,11 @@ def shutdown(self) -> None: self._shutdown_agentic_rollout_services() + try: + self._close_data_system() + except Exception as e: + logger.warning(f"Failed to tear down data system during controller shutdown: {e}") + logger.info("Controller shutdown complete.") def add_serve(self, role: str) -> None: diff --git a/relax/distributed/ray/rollout.py b/relax/distributed/ray/rollout.py index ea59f4f63..d045fb95c 100644 --- a/relax/distributed/ray/rollout.py +++ b/relax/distributed/ray/rollout.py @@ -16,7 +16,6 @@ import numpy as np import ray -import transfer_queue as tq import yaml from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH, GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_WEIGHTS @@ -52,6 +51,7 @@ from relax.utils.opd.opd_utils import compute_mopd_metrics from relax.utils.reload_utils import ReloadableMixin from relax.utils.s3_model_loader import prepare_model_maybe_update_args +from relax.utils.tq_lifecycle import attach_tq_client from relax.utils.tracking_utils import init_tracking from relax.utils.training.train_dump_utils import ( save_debug_rollout_data, @@ -813,8 +813,11 @@ def __init__(self, args, pg, data_source=None): self.data_source = data_source - tq.init(self.args.tq_config) - self.data_system_client = tq.get_client() + self.data_system_client = attach_tq_client( + self.args.tq_config, + requested_gdr=getattr(self.args, "tq_use_gdr", False), + role="rollout_worker", + ) logger.info(f"import {self.args.rollout_function_path} as generate_rollout function.") logger.info(f"import {self.args.eval_function_path} as eval_generate_rollout function.") diff --git a/relax/utils/rdma_probe.py b/relax/utils/rdma_probe.py index aa3e296ac..d495cbd75 100644 --- a/relax/utils/rdma_probe.py +++ b/relax/utils/rdma_probe.py @@ -94,9 +94,12 @@ class EffectiveConfig: def log_line(self) -> str: """Return the single-line startup log string for this effective config.""" - gdr_str = "on" if self.gdr else "off" + gdr_status = "unknown" if self.gdr else "off" dev = self.device or "auto" - base = f"[dataplane] backend={self.backend} protocol={self.protocol} device={dev} gdr={gdr_str}" + base = ( + f"[dataplane] backend={self.backend} protocol={self.protocol} device={dev} " + f"gdr_requested={str(self.gdr).lower()} gdr_status={gdr_status}" + ) if self.fallback_reason: return f"{base} fallback={self.fallback_reason}" return base @@ -179,12 +182,39 @@ def _check_memlock() -> CheckResult: return CheckResult("memlock", False, str(e)) +def _split_host_port(address: str) -> tuple[str, int]: + """Parse ``host:port`` and bracketed IPv6 endpoints.""" + value = address.strip() + if value.startswith("["): + end = value.find("]") + if end < 0 or end + 2 > len(value) or value[end + 1] != ":": + raise ValueError(f"invalid bracketed endpoint: {address!r}") + return value[1:end], int(value[end + 2 :]) + host, separator, port = value.rpartition(":") + if not separator or not host or not port: + raise ValueError(f"expected host:port, got {address!r}") + return host, int(port) + + +def _check_master_reachable(address: str, timeout: float = 2.0) -> CheckResult: + """Verify that this node can establish a bounded TCP connection to + master.""" + try: + host, port = _split_host_port(address) + with socket.create_connection((host, port), timeout=timeout): + pass + return CheckResult("master_reachable", True, address) + except (OSError, ValueError) as e: + return CheckResult("master_reachable", False, f"{address}: {e}") + + def _check_health_check() -> CheckResult: # pragma: no cover - retained for ad-hoc use """Call mooncake's native ``health_check()`` (NOT used by ``probe_node``). Returns 0=healthy, 1=not initialized/closed, 2=master unreachable. Kept as - a utility for post-init diagnostics; intentionally excluded from the pre- - init probe because the master is not running at probe time. + a utility for post-init diagnostics. The pre-init path uses a bounded TCP + reachability check instead because the global health API is process-state + dependent before a local Mooncake client has been initialized. """ try: import mooncake @@ -208,10 +238,9 @@ def _check_health_check() -> CheckResult: # pragma: no cover - retained for ad- # Per-node probe # --------------------------------------------------------------------------- -# NOTE: ``health_check()`` is intentionally NOT probed -- it queries the -# mooncake master, which is not yet running at probe time (auto_init=False), -# so it would always report failure and force a spurious SimpleStorage -# fallback. Master reachability is authoritatively checked by ``tq.init``. +# NOTE: ``health_check()`` is intentionally NOT probed because it depends on +# local Mooncake client initialization. External-master reachability is checked +# directly with a bounded TCP connect, then authoritatively by ``tq.init``. _CHECK_FUNCS_NO_DEVICE = [ _check_mooncake_import, _check_rdma_devices, @@ -219,7 +248,7 @@ def _check_health_check() -> CheckResult: # pragma: no cover - retained for ad- ] -def probe_node(device: str = "") -> ProbeResult: +def probe_node(device: str = "", master_address: str = "") -> ProbeResult: """Run all capability checks on the current node. Parameters @@ -234,6 +263,12 @@ def probe_node(device: str = "") -> ProbeResult: for fn in _CHECK_FUNCS_NO_DEVICE: checks.append(fn()) + # Mooncake is externally managed by Relax deployments. When an endpoint is + # supplied, it must already be reachable from every data-plane node before + # the job creates a global TransferQueue controller. + if master_address: + checks.append(_check_master_reachable(master_address)) + # Device-dependent checks. checks.append(_check_port_active(device)) checks.append(_check_gid_available(device)) @@ -244,12 +279,16 @@ def probe_node(device: str = "") -> ProbeResult: port_ok = any(c.name.startswith("port_active") and c.ok for c in checks) gid_ok = any(c.name.startswith("gid") and c.ok for c in checks) memlock_ok = any(c.name == "memlock" and c.ok for c in checks) + master_ok = not master_address or any(c.name == "master_reachable" and c.ok for c in checks) effective_protocol: str | None effective_device = device if not mooncake_ok: effective_protocol = None errors.append("mooncake not importable") + elif not master_ok: + effective_protocol = None + errors.append("master unreachable") elif rdma_dev_ok and port_ok and gid_ok and memlock_ok: effective_protocol = "rdma" if not effective_device: @@ -338,7 +377,12 @@ def _degenerate_result(node: str, error: str) -> ProbeResult: ) -def probe_cluster_nodes(device: str = "", *, timeout: float = 60.0) -> list[ProbeResult]: +def probe_cluster_nodes( + device: str = "", + master_address: str = "", + *, + timeout: float = 60.0, +) -> list[ProbeResult]: """Probe every alive GPU-bearing node and return one result per node. The driver fans the probe out as a short-lived Ray remote task pinned to @@ -353,33 +397,35 @@ def probe_cluster_nodes(device: str = "", *, timeout: float = 60.0) -> list[Prob as a degenerate ``effective_protocol=None`` result so the reducer degrades the whole job rather than silently omitting the node. - Returns ``[probe_node(device)]`` when no GPU workers are discoverable - (single-node / local dev), preserving backward-compatible behavior. + The driver is always included because the first ``tq.init`` creates a local + Mooncake client there even when the Ray head is CPU-only. Returns only the + driver result when no GPU workers are discoverable (single-node/local dev). """ node_ids = _alive_gpu_nodes() + driver_result = probe_node(device, master_address) if not node_ids: logger.debug("No alive GPU nodes discovered; probing driver node only.") - return [probe_node(device)] + return [driver_result] import ray from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy @ray.remote(num_cpus=0.001) - def _probe_on_node(dev: str) -> ProbeResult: + def _probe_on_node(dev: str, master: str) -> ProbeResult: from relax.utils.rdma_probe import probe_node as _probe - return _probe(dev) + return _probe(dev, master) refs: list[Any] = [] id_by_ref: dict[Any, str] = {} for node_id in node_ids: strategy = NodeAffinitySchedulingStrategy(node_id=node_id, soft=False) - ref = _probe_on_node.options(scheduling_strategy=strategy).remote(device) + ref = _probe_on_node.options(scheduling_strategy=strategy).remote(device, master_address) refs.append(ref) id_by_ref[ref] = node_id ready, pending = ray.wait(refs, num_returns=len(refs), timeout=timeout) - results: list[ProbeResult] = [] + results: list[ProbeResult] = [driver_result] for ref in ready: node_id = id_by_ref[ref] try: @@ -450,12 +496,18 @@ def reduce_results( if any_no_mooncake: failed_nodes = [r.node for r in results if r.effective_protocol is None] + master_failed_nodes = [r.node for r in results if "master unreachable" in r.errors] + reason = ( + f"master_unreachable:{','.join(master_failed_nodes)}" + if master_failed_nodes + else f"mooncake_unavailable:{','.join(failed_nodes)}" + ) return EffectiveConfig( backend=fallback_backend, protocol="tcp", device="", gdr=False, - fallback_reason=f"mooncake_unavailable:{','.join(failed_nodes)}", + fallback_reason=reason, ) if all_rdma: diff --git a/relax/utils/tq_config.py b/relax/utils/tq_config.py index 9db8dbe6c..909a7fa88 100644 --- a/relax/utils/tq_config.py +++ b/relax/utils/tq_config.py @@ -14,6 +14,7 @@ from __future__ import annotations +import inspect import os from typing import Any @@ -34,6 +35,39 @@ _DEFAULT_METADATA_SERVER = "P2PHANDSHAKE" # config.yaml:42-43 +def resolve_mooncake_master_address() -> str: + """Return the externally managed Mooncake master endpoint.""" + return os.environ.get("MC_MASTER_ADDRESS", "localhost:50051") + + +def validate_mooncake_runtime_contract() -> None: + """Fail fast unless installed TransferQueue has the loss-prevention fixes. + + The RDMA integration depends on per-key put/get result validation and on + notifying production readiness only after storage succeeds. A version + number alone is insufficient for development builds, so validate the + concrete runtime capabilities before probing or creating a controller. + """ + try: + from transfer_queue.storage.clients.mooncake_client import MooncakeStoreClient + from transfer_queue.storage.managers.base import KVStorageManager + except ImportError as e: + raise RuntimeError("Installed TransferQueue has no MooncakeStore support") from e + + required_methods = ("_batch_upsert_with_retry", "_batch_get_into_with_retry") + missing = [name for name in required_methods if not callable(getattr(MooncakeStoreClient, name, None))] + if missing: + raise RuntimeError("Installed TransferQueue lacks required Mooncake failure handling: " + ", ".join(missing)) + + put_source = inspect.getsource(KVStorageManager.put_data) + storage_call = put_source.find("self.storage_client.put") + ready_notify = put_source.find("self.notify_data_update") + if storage_call < 0 or ready_notify < 0 or storage_call > ready_notify: + raise RuntimeError( + "Installed TransferQueue does not guarantee storage success before production-status notification" + ) + + # --------------------------------------------------------------------------- # Config builders # --------------------------------------------------------------------------- @@ -74,7 +108,7 @@ def build_mooncake_config( pass a larger value (e.g. 8 GiB) to avoid staging-buffer pressure. """ if master_address is None: - master_address = os.environ.get("MC_MASTER_ADDRESS", "localhost:50051") + master_address = resolve_mooncake_master_address() cfg: dict[str, Any] = { # Selects the manager inside ``tq.init`` (interface.py reads diff --git a/relax/utils/tq_lifecycle.py b/relax/utils/tq_lifecycle.py index f6d11e7ee..dcea6a79f 100644 --- a/relax/utils/tq_lifecycle.py +++ b/relax/utils/tq_lifecycle.py @@ -22,6 +22,9 @@ from __future__ import annotations import time +import uuid +from dataclasses import dataclass +from typing import Any import ray import transfer_queue as tq @@ -33,6 +36,29 @@ CONTROLLER_NAME = "TransferQueueController" CONTROLLER_NAMESPACE = "transfer_queue" +OWNER_TOKEN_FIELD = "relax_owner_token" +DEFAULT_TQ_INIT_TIMEOUT_SECONDS = 60.0 + + +@dataclass(frozen=True) +class TqInitResult: + """Result of an owner-aware TransferQueue initialization transaction.""" + + config: Any + owner: Any | None + fallback_reason: str = "" + + @property + def owns_controller(self) -> bool: + return self.owner is not None + + +class TqInitializationTimeout(TimeoutError): + """Raised when ``tq.init`` does not finish within the bounded timeout.""" + + +class TqCleanupTimeout(TimeoutError): + """Raised when a TQ controller cannot be confirmed gone after cleanup.""" def kill_tq_controller_and_wait(timeout: float = 20.0) -> None: @@ -49,9 +75,8 @@ def kill_tq_controller_and_wait(timeout: float = 20.0) -> None: logger.info("[dataplane] Killed TransferQueueController actor (F10 guard).") except ValueError: return # actor does not exist — nothing to kill or wait for. - except Exception as e: # pragma: no cover - best-effort cleanup - logger.warning(f"[dataplane] Failed to kill TransferQueueController: {e}") - return + except Exception as e: + raise RuntimeError(f"Failed to kill TransferQueueController: {e}") from e deadline = time.time() + timeout while time.time() < deadline: @@ -60,7 +85,7 @@ def kill_tq_controller_and_wait(timeout: float = 20.0) -> None: except ValueError: return time.sleep(0.4) - logger.warning(f"[dataplane] TransferQueueController still resolvable after {timeout}s; proceeding anyway.") + raise TqCleanupTimeout(f"TransferQueueController still resolvable after {timeout}s") def reap_unusable_tq_controller(get_config_timeout: float = 10.0) -> bool: @@ -91,13 +116,130 @@ def reap_unusable_tq_controller(get_config_timeout: float = 10.0) -> bool: return True -def close_tq_and_unmount() -> None: +def _controller_exists() -> bool: + try: + ray.get_actor(CONTROLLER_NAME, namespace=CONTROLLER_NAMESPACE) + return True + except ValueError: + return False + + +def _set_owner_token(conf: Any, token: str) -> None: + controller = conf.controller if hasattr(conf, "controller") else conf["controller"] + controller[OWNER_TOKEN_FIELD] = token + + +def _get_owner_token(conf: Any) -> str: + if conf is None: + return "" + controller = conf.controller if hasattr(conf, "controller") else conf.get("controller", {}) + if hasattr(controller, "get"): + return str(controller.get(OWNER_TOKEN_FIELD, "")) + return "" + + +def _get_stored_config(timeout: float = 10.0) -> Any: + controller = ray.get_actor(CONTROLLER_NAME, namespace=CONTROLLER_NAMESPACE) + conf = ray.get(controller.get_config.remote(), timeout=timeout) + if conf is None: + raise RuntimeError("TransferQueueController returned no config after tq.init completed") + return conf + + +def _close_local_tq_client() -> None: + """Detach this process without deleting global TQ data or controller.""" + client = None + store_client = None + try: + client = tq.get_client() + store_client = getattr(client.storage_manager, "storage_client", None) + except (AssertionError, AttributeError): + pass + + if store_client is not None and hasattr(store_client, "close"): + try: + store_client.close() + except Exception as e: # pragma: no cover - best-effort local cleanup + logger.warning(f"[dataplane] Failed to close attached MooncakeStore client: {e}") + + if client is not None and hasattr(client, "close"): + try: + client.close() + except Exception as e: # pragma: no cover - best-effort local cleanup + logger.warning(f"[dataplane] Failed to close attached TransferQueue client: {e}") + + # TransferQueue has no public detach-only API. Reset only process-local + # handles; never touch _TQ_STORAGE or the named controller actor. + try: + from transfer_queue import interface as tq_interface + + tq_interface._TQ_CLIENT = None + tq_interface._TQ_CONTROLLER = None + except (ImportError, AttributeError): # pragma: no cover - version dependent + pass + + +def log_tq_gdr_runtime_status(*, requested: bool, role: str) -> str: + """Log requested GDR intent separately from the local client's status. + + ``enabled_unverified`` means the client selected its GDR staging path, but + Relax has not proved that a transfer traversed GDR on the wire. This + avoids claiming job-wide GDR effectiveness from a driver-side capability + probe. + """ + if not requested: + return "not_requested" + + status = "unknown" + detail = "client introspection unavailable" + try: + manager = tq.get_client().storage_manager + store_client = getattr(manager, "storage_client", None) + if store_client is None or type(manager).__name__ != "MooncakeStorageManager": + status = "inactive" + detail = f"manager={type(manager).__name__}" + elif getattr(store_client, "protocol", "") != "rdma": + status = "inactive" + detail = f"protocol={getattr(store_client, 'protocol', 'unknown')}" + elif getattr(store_client, "_gdr_staging", None) is None: + status = "host_rdma_fallback" + detail = "GDR staging unavailable in this worker" + else: + status = "enabled_unverified" + detail = "local GDR path selected; wire effectiveness is unknown" + except Exception as e: # pragma: no cover - environment/version dependent + detail = str(e) + + log = logger.warning if status != "enabled_unverified" else logger.info + log(f"[dataplane:gdr] role={role} requested=true status={status} experimental=true detail={detail}") + return status + + +def attach_tq_client(conf: Any, *, requested_gdr: bool, role: str) -> Any: + """Attach a component process and report its local experimental GDR + state.""" + tq.init(conf=conf) + client = tq.get_client() + log_tq_gdr_runtime_status(requested=requested_gdr, role=role) + return client + + +def close_tq_and_unmount(*, is_owner: bool) -> None: """Close TransferQueue and unmount the MooncakeStore segment. Order matters: ``tq.close()`` still needs the store alive for its ``remove_all()``, so the client handle is captured first and unmounted after. SimpleStorage has no ``storage_client``, so this is a no-op there. + + Global ``tq.close()`` is owner-only because upstream kills the named + controller even in a process that merely attached to it. Non-owners only + detach their process-local client. """ + if not is_owner: + logger.info("[dataplane] Detaching local TQ client; global controller is owned by another process.") + _close_local_tq_client() + return + store_client = None try: store_client = getattr(tq.get_client().storage_manager, "storage_client", None) @@ -112,3 +254,156 @@ def close_tq_and_unmount() -> None: logger.info("[dataplane] Unmounted MooncakeStore segment on teardown.") except Exception as e: # pragma: no cover - best-effort cleanup logger.warning(f"[dataplane] Failed to unmount MooncakeStore segment: {e}") + + +@ray.remote(num_cpus=0) +class _TransferQueueOwner: + """Process boundary for first-time TQ initialization and global cleanup.""" + + def __init__(self) -> None: + self._owns_controller = False + + def initialize(self, conf: Any, owner_token: str) -> tuple[Any, bool]: + _set_owner_token(conf, owner_token) + tq.init(conf=conf) + stored_conf = _get_stored_config() + self._owns_controller = _get_owner_token(stored_conf) == owner_token + return stored_conf, self._owns_controller + + def close(self) -> None: + close_tq_and_unmount(is_owner=self._owns_controller) + + def detach(self) -> None: + _close_local_tq_client() + + +def _stop_owner_actor(owner: Any) -> None: + try: + ray.kill(owner) + except Exception as e: # pragma: no cover - actor may already be dead + logger.debug(f"[dataplane] TQ owner actor already stopped: {e}") + + +def _cleanup_failed_owner(owner: Any, owner_token: str, *, timeout: float = 10.0) -> None: + """Stop a failed initializer and remove only controller state it owns. + + A concurrent initializer may win the global named-actor race. In that case + this actor only attached, so neither its cleanup RPC nor this driver is + allowed to kill the winning controller. + """ + try: + ray.get(owner.close.remote(), timeout=timeout) + except Exception as e: + logger.warning(f"[dataplane] TQ owner cleanup RPC failed; killing owner actor: {e}") + finally: + _stop_owner_actor(owner) + + try: + stored_conf = _get_stored_config(timeout=timeout) + except ValueError: + return + except Exception as e: + logger.warning(f"[dataplane] Failed initializer left an unusable TQ controller ({e}); reaping it.") + kill_tq_controller_and_wait() + return + + stored_token = _get_owner_token(stored_conf) + if stored_token == owner_token: + logger.warning("[dataplane] Failed initializer left its TQ controller behind; reaping owned state.") + kill_tq_controller_and_wait() + else: + logger.info( + "[dataplane] Failed initializer had attached to a concurrently owned " + "TQ controller; leaving global state intact." + ) + + +def _start_owner(conf: Any, *, timeout: float) -> TqInitResult: + owner = _TransferQueueOwner.remote() + owner_token = uuid.uuid4().hex + try: + stored_conf, owns_controller = ray.get(owner.initialize.remote(conf, owner_token), timeout=timeout) + except ray.exceptions.GetTimeoutError as e: + _cleanup_failed_owner(owner, owner_token) + raise TqInitializationTimeout(f"tq.init did not finish within {timeout:.0f}s") from e + except Exception: + _cleanup_failed_owner(owner, owner_token) + raise + + if owns_controller: + return TqInitResult(config=stored_conf, owner=owner) + + # A concurrent initializer won the named-actor race. This process is only + # attached and must never retain an actor capable of global tq.close(). + try: + ray.get(owner.detach.remote(), timeout=10.0) + finally: + _stop_owner_actor(owner) + return TqInitResult(config=stored_conf, owner=None) + + +def close_tq_owner(owner: Any | None, *, timeout: float = 30.0) -> None: + """Ask the owner process to close global TQ state; attached sessions no- + op.""" + if owner is None: + return + close_error: Exception | None = None + try: + ray.get(owner.close.remote(), timeout=timeout) + except Exception as e: + close_error = e + finally: + _stop_owner_actor(owner) + # Do not proceed to a subsequent initialization until the actor name has + # actually left GCS. + if _controller_exists(): + kill_tq_controller_and_wait() + if close_error is not None: + raise RuntimeError(f"TransferQueue owner cleanup failed: {close_error}") from close_error + + +def initialize_tq_with_fallback( + conf: Any, + *, + mode: str, + fallback_conf: Any | None = None, + timeout: float = DEFAULT_TQ_INIT_TIMEOUT_SECONDS, +) -> TqInitResult: + """Initialize TQ atomically with owner tracking and one safe auto fallback. + + ``fallback_conf`` must be the equivalent SimpleStorage configuration. It + is used only for ``mode='auto'`` after a real Mooncake ``tq.init`` failure. + ``required`` always cleans up and re-raises the original error. + """ + + def _attempt(attempt_conf: Any) -> TqInitResult: + reap_unusable_tq_controller() + if _controller_exists(): + # Attach semantics: use the controller's actual config. Upstream + # tq.init(conf) returns the caller-provided config even when ignored. + return TqInitResult(config=_get_stored_config(), owner=None) + return _start_owner(attempt_conf, timeout=timeout) + + try: + return _attempt(conf) + except Exception as primary_error: + if mode != "auto" or fallback_conf is None: + raise + + reason = f"mooncake_init_failed:{type(primary_error).__name__}" + logger.warning( + f"[dataplane] Mooncake tq.init failed ({primary_error}); " + "cleaned partial state and retrying once with SimpleStorage." + ) + try: + result = _attempt(fallback_conf) + except Exception as fallback_error: + raise RuntimeError( + "TransferQueue SimpleStorage fallback initialization failed after " + f"Mooncake initialization error: {primary_error}" + ) from fallback_error + return TqInitResult( + config=result.config, + owner=result.owner, + fallback_reason=reason, + ) diff --git a/scripts/benchmarks/tq_cross_node_bench.py b/scripts/benchmarks/tq_cross_node_bench.py index a0c56e255..315e6d356 100644 --- a/scripts/benchmarks/tq_cross_node_bench.py +++ b/scripts/benchmarks/tq_cross_node_bench.py @@ -17,15 +17,19 @@ * C1 ``tcp`` -- MooncakeStore / mooncake / TCP * C2 ``rdma`` -- MooncakeStore / mooncake / RDMA -Every payload档's transport is proven on the wire by reading the IB -``port_rcv_data`` and bond0 ``rx_bytes`` counters around each get: RDMA档 must -show IB moving and bond0 flat; TCP档 the reverse. There is no "thought it was -RDMA but was TCP" ambiguity. +Both synthetic tensors and production-shaped multimodal payloads are tested; +every fetched field is compared byte-for-byte via a SHA-256 digest before a +throughput result is accepted. Every payload tier's transport is proven by +reading the IB ``port_rcv_data`` and bond0 ``rx_bytes`` counters around each +get: RDMA must show IB moving and bond0 flat; TCP the reverse. There is no +"thought it was RDMA but was TCP" ambiguity. Usage (node A driver; node B already in the Ray cluster): PYTHONPATH= python -u scripts/benchmarks/tq_cross_node_bench.py \\ - --payload-mib 256 1024 2048 4096 --repeats 5 --csv tq_cross_node_gib.csv + --payload-profiles synthetic multimodal \\ + --payload-mib 256 1024 2048 4096 --repeats 5 \\ + --require-wire-proof --csv tq_cross_node_gib.csv On mooncake 0.3.10, switching protocols inside one driver session can make the third protocol's ``batch_get_into`` return -800 (see task-26-dev-log §7.11). @@ -40,6 +44,7 @@ import argparse import csv +import hashlib import statistics import time from typing import Any @@ -69,6 +74,13 @@ def parse_args() -> argparse.Namespace: ) p.add_argument("--num-samples", type=int, default=256, help="Rows per put (rollout-batch-like)") p.add_argument("--num-fields", nargs="+", type=int, default=[1], help="Tensor fields per put") + p.add_argument( + "--payload-profiles", + nargs="+", + default=["synthetic", "multimodal"], + choices=["synthetic", "multimodal"], + help="Payload layouts to benchmark. Multimodal uses production field names, shapes, and dtypes.", + ) p.add_argument( "--repeats", type=int, @@ -90,6 +102,11 @@ def parse_args() -> argparse.Namespace: help="Subset of configs to run (use one per process to dodge the 0.3.10 -800 issue).", ) p.add_argument("--csv", default="", help="Optional path to write per-run rows + summary") + p.add_argument( + "--require-wire-proof", + action="store_true", + help="Fail unless counters prove RDMA traffic for rdma and TCP traffic for tcp/simple.", + ) return p.parse_args() @@ -110,11 +127,77 @@ def make_payload(num_samples: int, num_fields: int, total_mib: int): return TensorDict(data, batch_size=[num_samples]) +def make_multimodal_payload(num_samples: int, total_mib: int): + """Build a production-shaped vision-language rollout TensorDict. + + ``pixel_values`` follows the Qwen-VL hidden width (1176) and BF16 dtype; + token, mask, grid, reward, and sample-id fields exercise the mixed dtypes + present in real Relax batches. Patch count scales to the requested size, + so the same layout is validated at every benchmark tier. + """ + import torch + from tensordict import TensorDict + + target_bytes = total_mib * 1024 * 1024 + seq_len = min(4096, max(128, target_bytes // max(1, num_samples * 128 * 1024))) + fixed_bytes = num_samples * (seq_len * (8 + 8 + 8) + 3 * 8 + 8 + 4) + pixel_budget = max(num_samples * 1176 * 2, target_bytes - fixed_bytes) + patches = max(1, pixel_budget // (num_samples * 1176 * 2)) + + pixel_values = torch.arange(num_samples * patches * 1176, dtype=torch.int32) + pixel_values = (pixel_values.remainder(2048).to(torch.float32) / 128).to(torch.bfloat16) + pixel_values = pixel_values.reshape(num_samples, patches, 1176) + token_row = torch.arange(seq_len, dtype=torch.int64) + input_ids = token_row.repeat(num_samples, 1) + response_ids = (token_row + 100_000).repeat(num_samples, 1) + attention_mask = torch.ones((num_samples, seq_len), dtype=torch.int64) + image_grid_thw = torch.tensor([1, 1, patches], dtype=torch.int64).repeat(num_samples, 1, 1) + sample_id = torch.arange(num_samples, dtype=torch.int64).reshape(num_samples, 1) + rewards = torch.linspace(-1.0, 1.0, num_samples, dtype=torch.float32).reshape(num_samples, 1) + return TensorDict( + { + "pixel_values": pixel_values, + "image_grid_thw": image_grid_thw, + "input_ids": input_ids, + "response_ids": response_ids, + "attention_mask": attention_mask, + "sample_id": sample_id, + "rewards": rewards, + }, + batch_size=[num_samples], + ) + + +def make_profile_payload(profile: str, num_samples: int, num_fields: int, total_mib: int): + if profile == "multimodal": + return make_multimodal_payload(num_samples, total_mib) + return make_payload(num_samples, num_fields, total_mib) + + def payload_bytes(payload) -> int: """Total bytes across all tensor fields.""" return sum(payload[k].nelement() * payload[k].element_size() for k in payload.keys()) +def field_byte_digests(payload, fields: list[str]) -> dict[str, tuple[str, int, str]]: + """Return per-field byte digests normalized to TQ's row-major storage. + + TQ reconstructs dense input columns as jagged ``NestedTensor`` columns. + Comparing flattened values makes the digest representation-independent + while still checking every dtype bit and every payload byte. + """ + import torch + + out: dict[str, tuple[str, int, str]] = {} + for field in fields: + value = payload[field] + flat = value.values().reshape(-1) if type(value).__name__ == "NestedTensor" else value.reshape(-1) + flat = flat.detach().cpu().contiguous() + raw = flat.view(torch.uint8).numpy().tobytes() + out[field] = (str(flat.dtype), flat.numel(), hashlib.sha256(raw).hexdigest()) + return out + + def wait_actor_gone(name: str = "TransferQueueController", timeout: float = 30.0) -> None: """Wait for a named TQ actor to leave the GCS (F10-safe re-init).""" deadline = time.time() + timeout @@ -252,7 +335,7 @@ def shutdown(self) -> None: except Exception: # pragma: no cover - best effort pass - def fetch(self, fields, batch_size: int, partition: str): + def fetch(self, fields, batch_size: int, partition: str, expected_digests): """One cross-node get; return (ms, ib_mb, tcp_mb, ib_tail_mb, tcp_tail_mb). @@ -276,15 +359,17 @@ def fetch(self, fields, batch_size: int, partition: str): after = read_counters() time.sleep(0.005) # let any async RDMA tail register on the counters settled = read_counters() - # Touch the data so the get is not optimized away and to sanity-check shape. - for f in fields: - v = got[f] - _ = v.values().reshape(-1) if type(v).__name__ == "NestedTensor" else v.reshape(-1) + # Digesting occurs outside the timed interval. A mismatch is fatal: + # throughput from a corrupt or truncated transfer is never reported. + actual_digests = field_byte_digests(got, list(fields)) + if actual_digests != expected_digests: + mismatch = [field for field in fields if actual_digests.get(field) != expected_digests.get(field)] + raise AssertionError(f"byte-exact mismatch after TQ get: fields={mismatch}") ib = sum(after[k] - before.get(k, 0) for k in after if k.startswith("ib:")) / 1e6 tcp = sum(after[k] - before.get(k, 0) for k in after if k.startswith("tcp:")) / 1e6 ib_tail = sum(settled[k] - after.get(k, 0) for k in settled if k.startswith("ib:")) / 1e6 tcp_tail = sum(settled[k] - after.get(k, 0) for k in settled if k.startswith("tcp:")) / 1e6 - return ms, ib, tcp, ib_tail, tcp_tail + return ms, ib, tcp, ib_tail, tcp_tail, True def _mean(values: list[float]) -> float: @@ -316,9 +401,9 @@ def main() -> None: ) labels = {"simple": "C0 SimpleStorage", "tcp": "C1 Mooncake/TCP", "rdma": "C2 Mooncake/RDMA"} - # results[protocol][total_mib][num_fields] = { + # results[protocol][profile][total_mib][num_fields] = { # "put_mean_gbs","get_mean_gbs","get_med_gbs","get_std_gbs","wire", "per_run":[...]} - results: dict[str, dict[int, dict[int, dict[str, Any]]]] = {} + results: dict[str, dict[str, dict[int, dict[int, dict[str, Any]]]]] = {} csv_rows: list[dict[str, Any]] = [] for protocol in args.protocols: @@ -344,87 +429,115 @@ def main() -> None: ) results.setdefault(protocol, {}) - for total_mib in args.payload_mib: - results[protocol].setdefault(total_mib, {}) - for nf in args.num_fields: - payload = make_payload(args.num_samples, nf, total_mib) - fields = sorted(payload.keys()) - nbytes = payload_bytes(payload) - put_times: list[float] = [] - get_times: list[float] = [] - ibs: list[float] = [] - tcps: list[float] = [] - ib_tails: list[float] = [] # async-completion detector (~0 == sync get) - tcp_tails: list[float] = [] - # repeat 0 is a warm-up: first transfer pays RDMA endpoint handshake. - for r in range(args.repeats + 1): - part = f"xfer_{protocol}_{total_mib}_{nf}_{r}" - t0 = time.perf_counter() - producer.put(payload, partition_id=part) - put_ms = (time.perf_counter() - t0) * 1000 - get_ms, ib, tcp, ib_tail, tcp_tail = ray.get(consumer.fetch.remote(fields, args.num_samples, part)) - producer.clear_partition(part) - if r == 0: - print( - f" {total_mib}M f={nf} (warmup, not counted): " - f"put={put_ms:.0f}ms get={get_ms:.0f}ms " - f"ib_tail={ib_tail:.0f}MB", - flush=True, + for profile in args.payload_profiles: + results[protocol].setdefault(profile, {}) + for total_mib in args.payload_mib: + results[protocol][profile].setdefault(total_mib, {}) + # Multimodal has a fixed production schema; synthetic uses the + # requested field-count sweep. + field_counts = args.num_fields if profile == "synthetic" else [7] + for nf in field_counts: + payload = make_profile_payload(profile, args.num_samples, nf, total_mib) + nf = len(list(payload.keys())) + fields = sorted(payload.keys()) + expected_digests = field_byte_digests(payload, fields) + nbytes = payload_bytes(payload) + put_times: list[float] = [] + get_times: list[float] = [] + ibs: list[float] = [] + tcps: list[float] = [] + ib_tails: list[float] = [] # async-completion detector (~0 == sync get) + tcp_tails: list[float] = [] + # repeat 0 is a warm-up: first transfer pays RDMA endpoint handshake. + for r in range(args.repeats + 1): + part = f"xfer_{protocol}_{profile}_{total_mib}_{nf}_{r}" + t0 = time.perf_counter() + producer.put(payload, partition_id=part) + put_ms = (time.perf_counter() - t0) * 1000 + get_ms, ib, tcp, ib_tail, tcp_tail, byte_exact = ray.get( + consumer.fetch.remote(fields, args.num_samples, part, expected_digests) + ) + producer.clear_partition(part) + if r == 0: + print( + f" {profile} {total_mib}M f={nf} (warmup, not counted): " + f"put={put_ms:.0f}ms get={get_ms:.0f}ms " + f"ib_tail={ib_tail:.0f}MB byte_exact={byte_exact}", + flush=True, + ) + continue + put_times.append(put_ms) + get_times.append(get_ms) + ibs.append(ib) + tcps.append(tcp) + ib_tails.append(ib_tail) + tcp_tails.append(tcp_tail) + run_wire = "RDMA" if ib > tcp else "TCP" + run_wire_proven = (protocol == "rdma" and ib > 0 and ib > tcp) or ( + protocol != "rdma" and tcp > 0 and tcp >= ib + ) + csv_rows.append( + { + "protocol": protocol, + "profile": profile, + "payload_mib": total_mib, + "actual_mib": round(nbytes / 1024**2, 2), + "num_fields": nf, + "run": r, + "byte_exact": byte_exact, + "wire_observed": run_wire, + "wire_proven": run_wire_proven, + "put_ms": round(put_ms, 2), + "get_ms": round(get_ms, 2), + "put_gbs": round(_gbs(nbytes, put_ms), 3), + "get_gbs": round(_gbs(nbytes, get_ms), 3), + "ib_mb": round(ib, 1), + "tcp_mb": round(tcp, 1), + "ib_tail_mb": round(ib_tail, 1), + "tcp_tail_mb": round(tcp_tail, 1), + } ) - continue - put_times.append(put_ms) - get_times.append(get_ms) - ibs.append(ib) - tcps.append(tcp) - ib_tails.append(ib_tail) - tcp_tails.append(tcp_tail) - csv_rows.append( - { - "protocol": protocol, - "payload_mib": total_mib, - "num_fields": nf, - "run": r, - "put_ms": round(put_ms, 2), - "get_ms": round(get_ms, 2), - "put_gbs": round(_gbs(nbytes, put_ms), 3), - "get_gbs": round(_gbs(nbytes, get_ms), 3), - "ib_mb": round(ib, 1), - "tcp_mb": round(tcp, 1), - "ib_tail_mb": round(ib_tail, 1), - "tcp_tail_mb": round(tcp_tail, 1), - } - ) - put_mean = _mean(put_times) - get_mean = _mean(get_times) - get_med = statistics.median(get_times) - # Std-dev of the per-run THROUGHPUT values (not of latency): converting a - # latency stddev via _gbs(nbytes, std_ms) would be a meaningless number. - get_gbs_runs = [_gbs(nbytes, ms) for ms in get_times] - get_std_gbs = statistics.pstdev(get_gbs_runs) if len(get_gbs_runs) > 1 else 0.0 - ib_med = statistics.median(ibs) - tcp_med = statistics.median(tcps) - ib_tail_med = statistics.median(ib_tails) if ib_tails else 0.0 - wire = "RDMA" if ib_med > tcp_med else "TCP" - rec = { - "put_mean_gbs": _gbs(nbytes, put_mean), - "get_mean_gbs": _gbs(nbytes, get_mean), - "get_med_gbs": _gbs(nbytes, get_med), - "get_std_gbs": get_std_gbs, - "wire": wire, - "ib_tail_med_mb": ib_tail_med, - "per_run_get_gbs": [round(g, 2) for g in get_gbs_runs], - } - results[protocol][total_mib][nf] = rec - print( - f" {str(total_mib) + 'M':<9} f={nf} put_mean={rec['put_mean_gbs']:6.2f} " - f"GB/s get_mean={rec['get_mean_gbs']:6.2f} (med {rec['get_med_gbs']:.2f}, " - f"std {rec['get_std_gbs']:.2f}) GB/s " - f"[wire: IB {ib_med:.0f}MB / bond0 {tcp_med:.0f}MB -> {wire}; " - f"tail {ib_tail_med:.0f}MB] " - f"runs={rec['per_run_get_gbs']}", - flush=True, - ) + put_mean = _mean(put_times) + get_mean = _mean(get_times) + get_med = statistics.median(get_times) + # Std-dev of per-run throughput, not latency. + get_gbs_runs = [_gbs(nbytes, ms) for ms in get_times] + get_std_gbs = statistics.pstdev(get_gbs_runs) if len(get_gbs_runs) > 1 else 0.0 + ib_med = statistics.median(ibs) + tcp_med = statistics.median(tcps) + ib_tail_med = statistics.median(ib_tails) if ib_tails else 0.0 + wire = "RDMA" if ib_med > tcp_med else "TCP" + wire_proven = (protocol == "rdma" and ib_med > 0 and ib_med > tcp_med) or ( + protocol != "rdma" and tcp_med > 0 and tcp_med >= ib_med + ) + if args.require_wire_proof and not wire_proven: + raise RuntimeError( + f"wire proof failed for protocol={protocol} profile={profile} " + f"payload={total_mib}MiB: IB={ib_med:.1f}MB TCP={tcp_med:.1f}MB" + ) + rec = { + "put_mean_gbs": _gbs(nbytes, put_mean), + "get_mean_gbs": _gbs(nbytes, get_mean), + "get_med_gbs": _gbs(nbytes, get_med), + "get_std_gbs": get_std_gbs, + "wire": wire, + "wire_proven": wire_proven, + "byte_exact": True, + "ib_tail_med_mb": ib_tail_med, + "per_run_get_gbs": [round(g, 2) for g in get_gbs_runs], + } + results[protocol][profile][total_mib][nf] = rec + print( + f" {profile:<10} {str(total_mib) + 'M':<9} f={nf} " + f"put_mean={rec['put_mean_gbs']:6.2f} " + f"GB/s get_mean={rec['get_mean_gbs']:6.2f} (med {rec['get_med_gbs']:.2f}, " + f"std {rec['get_std_gbs']:.2f}) GB/s byte_exact=PASS " + f"[wire: IB {ib_med:.0f}MB / bond0 {tcp_med:.0f}MB -> {wire}; " + f"proof={'PASS' if wire_proven else 'UNKNOWN'}; tail {ib_tail_med:.0f}MB] " + f"runs={rec['per_run_get_gbs']}", + flush=True, + ) ray.get(consumer.shutdown.remote()) # unmount before kill, else the segment lingers ray.kill(consumer) @@ -433,32 +546,35 @@ def main() -> None: # ---- Summary (mean-based, all requested protocols) ---- print("\n===== SUMMARY: TQ-layer cross-node, same topology (get, MEAN of N runs) =====", flush=True) header = ( - f"{'Payload':<9}{'f':<4}{'C0 mean':>9}{'C1 mean':>9}{'C2 mean':>9}" + f"{'Profile':<11}{'Payload':<9}{'f':<4}{'C0 mean':>9}{'C1 mean':>9}{'C2 mean':>9}" f"{'C1/C0':>8}{'C2/C1':>8}{'C2 std':>8}{'>=20%':>7}{'wire C0/C1/C2':>18}" ) print(header, flush=True) - for total_mib in args.payload_mib: - for nf in args.num_fields: - c0 = results.get("simple", {}).get(total_mib, {}).get(nf) - c1 = results.get("tcp", {}).get(total_mib, {}).get(nf) - c2 = results.get("rdma", {}).get(total_mib, {}).get(nf) - if not (c0 and c1 and c2): - # A protocol was skipped (--protocols subset) -- print what we have. - parts = [] - for name, c in (("C0", c0), ("C1", c1), ("C2", c2)): - parts.append(f"{name}={c['get_mean_gbs']:.2f}" if c else f"{name}=-") - print(f"{str(total_mib) + 'M':<9}{nf:<4}" + " ".join(parts) + " (subset run)", flush=True) - continue - g0, g1, g2 = c0["get_mean_gbs"], c1["get_mean_gbs"], c2["get_mean_gbs"] - back_pct = (g1 - g0) / g0 * 100 if g0 > 0 else 0.0 - rdma_pct = (g2 - g1) / g1 * 100 if g1 > 0 else 0.0 - wire = f"{c0['wire']}/{c1['wire']}/{c2['wire']}" - print( - f"{str(total_mib) + 'M':<9}{nf:<4}{g0:>9.2f}{g1:>9.2f}{g2:>9.2f}" - f"{f'{back_pct:+.0f}%':>8}{f'{rdma_pct:+.0f}%':>8}{c2['get_std_gbs']:>8.2f}" - f"{('PASS' if rdma_pct >= 20 else 'no'):>7}{wire:>18}", - flush=True, - ) + for profile in args.payload_profiles: + field_counts = args.num_fields if profile == "synthetic" else [7] + for total_mib in args.payload_mib: + for nf in field_counts: + c0 = results.get("simple", {}).get(profile, {}).get(total_mib, {}).get(nf) + c1 = results.get("tcp", {}).get(profile, {}).get(total_mib, {}).get(nf) + c2 = results.get("rdma", {}).get(profile, {}).get(total_mib, {}).get(nf) + prefix = f"{profile:<11}{str(total_mib) + 'M':<9}{nf:<4}" + if not (c0 and c1 and c2): + # A protocol was skipped (--protocols subset) -- print what we have. + parts = [] + for name, c in (("C0", c0), ("C1", c1), ("C2", c2)): + parts.append(f"{name}={c['get_mean_gbs']:.2f}" if c else f"{name}=-") + print(prefix + " ".join(parts) + " (subset run)", flush=True) + continue + g0, g1, g2 = c0["get_mean_gbs"], c1["get_mean_gbs"], c2["get_mean_gbs"] + back_pct = (g1 - g0) / g0 * 100 if g0 > 0 else 0.0 + rdma_pct = (g2 - g1) / g1 * 100 if g1 > 0 else 0.0 + wire = f"{c0['wire']}/{c1['wire']}/{c2['wire']}" + print( + f"{prefix}{g0:>9.2f}{g1:>9.2f}{g2:>9.2f}" + f"{f'{back_pct:+.0f}%':>8}{f'{rdma_pct:+.0f}%':>8}{c2['get_std_gbs']:>8.2f}" + f"{('PASS' if rdma_pct >= 20 else 'no'):>7}{wire:>18}", + flush=True, + ) print(" C1/C0 = MooncakeStore vs SimpleStorage (backend effect)", flush=True) print(" C2/C1 = RDMA vs TCP on the same backend (transport effect, gate target, mean-based)", flush=True) print(" std = population stddev of C2 get across the N runs (run-to-run variance)", flush=True) @@ -466,9 +582,14 @@ def main() -> None: if args.csv: cols = [ "protocol", + "profile", "payload_mib", + "actual_mib", "num_fields", "run", + "byte_exact", + "wire_observed", + "wire_proven", "put_ms", "get_ms", "put_gbs", diff --git a/tests/utils/test_rdma_probe.py b/tests/utils/test_rdma_probe.py index b2cb0dc4f..24dc161f2 100644 --- a/tests/utils/test_rdma_probe.py +++ b/tests/utils/test_rdma_probe.py @@ -17,8 +17,10 @@ CheckResult, EffectiveConfig, ProbeResult, + _check_master_reachable, _degenerate_result, _select_dataplane_node_ids, + _split_host_port, probe_cluster_nodes, probe_node, reduce_results, @@ -28,6 +30,7 @@ build_mooncake_config, build_simple_storage_config, estimate_payload_bytes, + validate_mooncake_runtime_contract, validate_segment_capacity, ) @@ -223,6 +226,25 @@ def fake_isdir(path): assert result.effective_protocol == "rdma" assert result.ok + def test_unreachable_external_master_disables_mooncake(self, monkeypatch): + monkeypatch.setattr(rdma_probe, "_check_mooncake_import", lambda: CheckResult("mooncake_import", True)) + monkeypatch.setattr( + rdma_probe, + "_check_master_reachable", + lambda address: CheckResult("master_reachable", False, address), + ) + result = probe_node("", "master.invalid:50051") + assert result.effective_protocol is None + assert "master unreachable" in result.errors + + def test_master_endpoint_parser(self): + assert _split_host_port("master.example:50051") == ("master.example", 50051) + assert _split_host_port("[2001:db8::1]:50051") == ("2001:db8::1", 50051) + + def test_master_reachability_is_bounded_failure(self): + result = _check_master_reachable("127.0.0.1:1", timeout=0.01) + assert result.ok is False + # --------------------------------------------------------------------------- # probe_cluster_nodes (multi-node fan-out) + helpers @@ -256,7 +278,7 @@ def test_cluster_falls_back_to_local_when_no_gpu_nodes(self, monkeypatch): never touching Ray remote scheduling.""" monkeypatch.setattr(rdma_probe, "_alive_gpu_nodes", lambda: []) local = _make_probe(protocol="rdma", node="local-driver") - monkeypatch.setattr(rdma_probe, "probe_node", lambda dev: local) + monkeypatch.setattr(rdma_probe, "probe_node", lambda dev, master="": local) results = probe_cluster_nodes("") assert len(results) == 1 assert results[0] is local @@ -277,6 +299,24 @@ def test_reduce_treats_degenerate_as_no_mooncake(self): assert eff.backend == "SimpleStorage" assert "n1" in eff.fallback_reason + def test_reduce_reports_master_unreachable_distinctly(self): + unavailable = ProbeResult( + node="n1", + checks=(CheckResult("master_reachable", False),), + effective_protocol=None, + effective_device="", + gdr_eligible=False, + errors=("master unreachable",), + ) + eff = reduce_results( + [_make_probe(protocol="rdma", node="n0"), unavailable], + requested_backend="mooncake", + requested_device="", + use_gdr=False, + ) + assert eff.backend == "SimpleStorage" + assert eff.fallback_reason == "master_unreachable:n1" + # --------------------------------------------------------------------------- # tq_config builders @@ -319,6 +359,9 @@ def test_mooncake_config_gdr_propagated(self): cfg = build_mooncake_config(eff) assert cfg["MooncakeStore"]["use_gdr"] is True + def test_installed_tq_satisfies_loss_prevention_contract(self): + validate_mooncake_runtime_contract() + def test_segment_capacity_text_only_passes(self): args = _make_args(multimodal_keys=None) eff = EffectiveConfig(backend="MooncakeStore", protocol="rdma", device="", gdr=False, fallback_reason="") diff --git a/tests/utils/test_tq_failure_paths.py b/tests/utils/test_tq_failure_paths.py index 8ff560c42..b6f873732 100644 --- a/tests/utils/test_tq_failure_paths.py +++ b/tests/utils/test_tq_failure_paths.py @@ -19,10 +19,12 @@ from __future__ import annotations +import asyncio import importlib.util +import multiprocessing import os import socket -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import pytest import torch @@ -48,6 +50,7 @@ def _has_real_submodule(dotted: str) -> bool: _REAL_MOONCAKE_CLIENT = _has_real_submodule("transfer_queue.storage.clients.mooncake_client") +_RUN_REAL_CAPACITY = os.environ.get("RELAX_RUN_REAL_MOONCAKE_CAPACITY_TEST") == "1" # --------------------------------------------------------------------------- @@ -82,6 +85,35 @@ def _master_reachable(timeout: float = 1.0) -> bool: return False +def _real_capacity_worker(result_queue, segment_mib: int, payload_mib: int) -> None: + """Child-process target so a real Mooncake hang is externally bounded.""" + from transfer_queue.storage.clients.mooncake_client import MooncakeStoreClient + + client = None + try: + client = MooncakeStoreClient( + { + "protocol": "tcp", + "device_name": "", + "master_server_address": _master_address(), + "metadata_server": "P2PHANDSHAKE", + "local_hostname": "", + "global_segment_size": segment_mib * 1024**2, + "local_buffer_size": min(segment_mib, 64) * 1024**2, + "hard_pin": True, + "use_gdr": False, + } + ) + value = torch.arange(payload_mib * 1024**2, dtype=torch.uint8) + client.put(["real-capacity-overflow"], [value]) + result_queue.put(("unexpected_success", "put returned success")) + except BaseException as error: + result_queue.put(("error", f"{type(error).__name__}: {error}")) + finally: + if client is not None: + client.close() + + # --------------------------------------------------------------------------- # Controller lifecycle: reaper (timeout / half-initialised / healthy) # --------------------------------------------------------------------------- @@ -162,22 +194,224 @@ def test_mooncake_segment_is_unmounted_after_close(self, monkeypatch): store_client = MagicMock() calls = self._fake_tq(monkeypatch, store_client=store_client) store_client.close.side_effect = lambda: calls.append("store.close") - tq_lifecycle.close_tq_and_unmount() + tq_lifecycle.close_tq_and_unmount(is_owner=True) # Order matters: tq.close() still needs the store alive for remove_all(). assert calls == ["tq.close", "store.close"] def test_simple_storage_teardown_is_noop_beyond_close(self, monkeypatch): calls = self._fake_tq(monkeypatch, store_client=None) - tq_lifecycle.close_tq_and_unmount() + tq_lifecycle.close_tq_and_unmount(is_owner=True) assert calls == ["tq.close"] def test_uninitialised_tq_does_not_raise(self, monkeypatch): fake = MagicMock() fake.get_client.side_effect = AssertionError("Please initialize the TransferQueue first") monkeypatch.setattr(tq_lifecycle, "tq", fake) - tq_lifecycle.close_tq_and_unmount() # must not raise + tq_lifecycle.close_tq_and_unmount(is_owner=True) # must not raise fake.close.assert_called_once() + def test_attached_process_never_calls_global_close(self, monkeypatch): + store_client = MagicMock() + calls = self._fake_tq(monkeypatch, store_client=store_client) + tq_lifecycle.close_tq_and_unmount(is_owner=False) + assert calls == [] + store_client.close.assert_called_once() + + +# --------------------------------------------------------------------------- +# GDR requested-vs-runtime status +# --------------------------------------------------------------------------- + + +class MooncakeStorageManager: + def __init__(self, storage_client): + self.storage_client = storage_client + + +class TestGdrRuntimeStatus: + def test_not_requested_is_distinct(self): + assert tq_lifecycle.log_tq_gdr_runtime_status(requested=False, role="test") == "not_requested" + + def test_requested_on_simple_backend_is_inactive(self, monkeypatch): + fake = MagicMock() + fake.get_client.return_value = MagicMock(storage_manager=MagicMock()) + monkeypatch.setattr(tq_lifecycle, "tq", fake) + assert tq_lifecycle.log_tq_gdr_runtime_status(requested=True, role="test") == "inactive" + + def test_requested_without_worker_staging_reports_host_fallback(self, monkeypatch): + store = MagicMock(protocol="rdma") + store._gdr_staging = None + fake = MagicMock() + fake.get_client.return_value = MagicMock(storage_manager=MooncakeStorageManager(store)) + monkeypatch.setattr(tq_lifecycle, "tq", fake) + assert tq_lifecycle.log_tq_gdr_runtime_status(requested=True, role="test") == "host_rdma_fallback" + + def test_local_gdr_path_never_claims_verified_effectiveness(self, monkeypatch): + store = MagicMock(protocol="rdma") + store._gdr_staging = object() + fake = MagicMock() + fake.get_client.return_value = MagicMock(storage_manager=MooncakeStorageManager(store)) + monkeypatch.setattr(tq_lifecycle, "tq", fake) + assert tq_lifecycle.log_tq_gdr_runtime_status(requested=True, role="test") == "enabled_unverified" + + +# --------------------------------------------------------------------------- +# Owner-aware initialization transaction +# --------------------------------------------------------------------------- + + +class TestInitializeTqWithFallback: + @staticmethod + def _conf(backend: str) -> dict: + return {"controller": {}, "backend": {"storage_backend": backend}} + + @staticmethod + def _patch_transaction(monkeypatch, *, existed: bool, init_effects: list[object], stored_conf=None): + calls: dict[str, list] = {"reap": [], "attempts": []} + effects = iter(init_effects) + + monkeypatch.setattr(tq_lifecycle, "reap_unusable_tq_controller", lambda: calls["reap"].append(True)) + monkeypatch.setattr(tq_lifecycle, "_controller_exists", lambda: existed) + monkeypatch.setattr(tq_lifecycle, "_get_stored_config", lambda: stored_conf) + + def fake_start(conf, *, timeout): + calls["attempts"].append(conf) + effect = next(effects) + if isinstance(effect, BaseException): + raise effect + return tq_lifecycle.TqInitResult(config=conf, owner=effect) + + monkeypatch.setattr(tq_lifecycle, "_start_owner", fake_start) + return calls + + def test_simple_path_also_runs_pre_init_reaper_and_becomes_owner(self, monkeypatch): + conf = self._conf("SimpleStorage") + calls = self._patch_transaction(monkeypatch, existed=False, init_effects=["owner"]) + result = tq_lifecycle.initialize_tq_with_fallback(conf, mode="off") + assert result.owns_controller is True + assert len(calls["reap"]) == 1 + + def test_attach_is_not_owner(self, monkeypatch): + requested = self._conf("SimpleStorage") + stored = self._conf("SimpleStorage") + calls = self._patch_transaction(monkeypatch, existed=True, init_effects=[], stored_conf=stored) + result = tq_lifecycle.initialize_tq_with_fallback(requested, mode="off") + assert result.owns_controller is False + assert result.config is stored + assert calls["attempts"] == [] + + def test_auto_cleans_failed_mooncake_then_retries_simple_once(self, monkeypatch): + primary = self._conf("MooncakeStore") + fallback = self._conf("SimpleStorage") + calls = self._patch_transaction( + monkeypatch, + existed=False, + init_effects=[RuntimeError("master unavailable"), "fallback-owner"], + ) + result = tq_lifecycle.initialize_tq_with_fallback(primary, mode="auto", fallback_conf=fallback) + assert result.config["backend"]["storage_backend"] == "SimpleStorage" + assert result.fallback_reason == "mooncake_init_failed:RuntimeError" + assert len(calls["attempts"]) == 2 + assert len(calls["reap"]) == 2 + + def test_required_cleans_failed_init_without_fallback(self, monkeypatch): + primary = self._conf("MooncakeStore") + fallback = self._conf("SimpleStorage") + calls = self._patch_transaction( + monkeypatch, + existed=False, + init_effects=[RuntimeError("master unavailable")], + ) + with pytest.raises(RuntimeError, match="master unavailable"): + tq_lifecycle.initialize_tq_with_fallback(primary, mode="required", fallback_conf=fallback) + assert len(calls["attempts"]) == 1 + + def test_timeout_auto_retries_only_after_isolated_owner_cleanup(self, monkeypatch): + primary = self._conf("MooncakeStore") + fallback = self._conf("SimpleStorage") + calls = self._patch_transaction( + monkeypatch, + existed=False, + init_effects=[tq_lifecycle.TqInitializationTimeout("timed out"), "fallback-owner"], + ) + result = tq_lifecycle.initialize_tq_with_fallback(primary, mode="auto", fallback_conf=fallback) + assert result.config["backend"]["storage_backend"] == "SimpleStorage" + assert len(calls["attempts"]) == 2 + + +class _RemoteMethod: + def __init__(self, value): + self.value = value + + def remote(self, *args, **kwargs): + return self.value + + +class _FakeOwner: + def __init__(self): + self.initialize = _RemoteMethod("initialize-ref") + self.close = _RemoteMethod("close-ref") + self.detach = _RemoteMethod("detach-ref") + + +class TestOwnerProcessBoundary: + def test_start_timeout_cleans_the_isolated_owner_before_raising(self, monkeypatch): + owner = _FakeOwner() + cleaned: list[tuple[object, str]] = [] + monkeypatch.setattr(tq_lifecycle._TransferQueueOwner, "remote", lambda: owner) + + def timed_out(ref, *, timeout): + assert ref == "initialize-ref" + raise tq_lifecycle.ray.exceptions.GetTimeoutError("test timeout") + + monkeypatch.setattr(tq_lifecycle.ray, "get", timed_out) + monkeypatch.setattr( + tq_lifecycle, + "_cleanup_failed_owner", + lambda handle, token: cleaned.append((handle, token)), + ) + + with pytest.raises(tq_lifecycle.TqInitializationTimeout): + tq_lifecycle._start_owner({"controller": {}}, timeout=0.1) + assert cleaned[0][0] is owner + assert cleaned[0][1] + + @pytest.mark.parametrize("stored_token,should_kill", [("ours", True), ("theirs", False)]) + def test_failed_owner_cleanup_respects_controller_owner_token(self, monkeypatch, stored_token, should_kill): + owner = _FakeOwner() + stopped: list[object] = [] + killed: list[bool] = [] + monkeypatch.setattr(tq_lifecycle.ray, "get", lambda ref, timeout: None) + monkeypatch.setattr(tq_lifecycle, "_stop_owner_actor", lambda handle: stopped.append(handle)) + monkeypatch.setattr( + tq_lifecycle, + "_get_stored_config", + lambda timeout: {"controller": {tq_lifecycle.OWNER_TOKEN_FIELD: stored_token}}, + ) + monkeypatch.setattr(tq_lifecycle, "kill_tq_controller_and_wait", lambda: killed.append(True)) + + tq_lifecycle._cleanup_failed_owner(owner, "ours") + assert stopped == [owner] + assert bool(killed) is should_kill + + def test_owner_close_failure_still_reaps_global_controller(self, monkeypatch): + owner = _FakeOwner() + stopped: list[object] = [] + killed: list[bool] = [] + monkeypatch.setattr( + tq_lifecycle.ray, + "get", + lambda ref, timeout: (_ for _ in ()).throw(RuntimeError("close failed")), + ) + monkeypatch.setattr(tq_lifecycle, "_stop_owner_actor", lambda handle: stopped.append(handle)) + monkeypatch.setattr(tq_lifecycle, "_controller_exists", lambda: True) + monkeypatch.setattr(tq_lifecycle, "kill_tq_controller_and_wait", lambda: killed.append(True)) + + with pytest.raises(RuntimeError, match="owner cleanup failed"): + tq_lifecycle.close_tq_owner(owner) + assert stopped == [owner] + assert killed == [True] + # --------------------------------------------------------------------------- # Retry / disconnect on the MooncakeStore data path @@ -225,6 +459,24 @@ def _client_with_store(store) -> object: return client +class _InlineExecutorLoop: + """Run storage-manager sync calls inline for deterministic async tests. + + The production manager delegates KV operations to an executor. A real + default executor makes the failure-path test wait for worker shutdown on + some CI kernels, so these contract tests substitute an already-completed + Future without changing the storage-manager control flow. + """ + + def run_in_executor(self, executor, fn, *args): + future = asyncio.get_running_loop().create_future() + try: + future.set_result(fn(*args)) + except BaseException as error: + future.set_exception(error) + return future + + @pytest.mark.skipif( not _REAL_MOONCAKE_CLIENT, reason="needs a real transfer_queue (CI uses a single-file stub); run on a host with TransferQueue installed", @@ -267,6 +519,94 @@ def test_disconnect_surfaces_instead_of_returning_garbage(self): client._batch_get_into_with_retry(["0@f0"], [1], [8]) +@pytest.mark.skipif( + not _REAL_MOONCAKE_CLIENT, + reason="needs a real transfer_queue to verify the KV manager write/notify contract", +) +class TestMooncakeProductionStatusContract: + @staticmethod + def _manager(storage_client): + from transfer_queue.storage.managers.mooncake_manager import MooncakeStorageManager + + manager = object.__new__(MooncakeStorageManager) + manager.storage_client = storage_client + manager.notify_data_update = AsyncMock() + manager.controller_handshake_socket = None + manager.storage_manager_id = "capacity-contract-test" + manager.zmq_context = MagicMock() + return manager + + @staticmethod + def _data_and_meta(): + from tensordict import TensorDict + + data = TensorDict({"pixel_values": torch.randn(1, 16)}, batch_size=[1]) + meta = MagicMock() + meta.global_indexes = [7] + meta.partition_ids = ["capacity"] + meta._custom_backend_meta = [{}] + meta.get_all_custom_meta.return_value = [{}] + return data, meta + + @pytest.mark.asyncio + async def test_capacity_write_failure_never_notifies_production_ready(self, monkeypatch): + monkeypatch.setattr(asyncio, "get_event_loop", lambda: _InlineExecutorLoop()) + storage_client = MagicMock() + storage_client.put.side_effect = RuntimeError("batch_upsert_from failed: capacity exhausted") + manager = self._manager(storage_client) + data, meta = self._data_and_meta() + + with pytest.raises(RuntimeError, match="capacity exhausted"): + await manager.put_data(data, meta) + + manager.notify_data_update.assert_not_awaited() + + @pytest.mark.asyncio + async def test_production_ready_is_notified_only_after_storage_success(self, monkeypatch): + monkeypatch.setattr(asyncio, "get_event_loop", lambda: _InlineExecutorLoop()) + calls: list[str] = [] + storage_client = MagicMock() + storage_client.put.side_effect = lambda keys, values: calls.append("put") or [None] * len(keys) + manager = self._manager(storage_client) + + async def notify(*args, **kwargs): + calls.append("notify") + + manager.notify_data_update.side_effect = notify + data, meta = self._data_and_meta() + await manager.put_data(data, meta) + assert calls == ["put", "notify"] + + +@pytest.mark.skipif( + not (_RUN_REAL_CAPACITY and _master_reachable() and _REAL_MOONCAKE_CLIENT), + reason=( + "destructive real-capacity test is opt-in and needs an isolated reachable master; " + "set RELAX_RUN_REAL_MOONCAKE_CAPACITY_TEST=1 only on a disposable deployment" + ), +) +def test_real_mooncake_capacity_overflow_is_bounded_and_loud(): + """A physical segment overflow must fail, never hang or report success. + + Run this only against an isolated master: the deliberately tiny segment and + oversized put are fault injection, not a shared-cluster smoke test. + """ + context = multiprocessing.get_context("spawn") + result_queue = context.Queue() + process = context.Process(target=_real_capacity_worker, args=(result_queue, 64, 96)) + process.start() + process.join(timeout=30) + if process.is_alive(): + process.terminate() + process.join(timeout=5) + pytest.fail("Mooncake capacity-overflow put did not finish within 30 seconds") + + assert process.exitcode == 0 + status, detail = result_queue.get(timeout=2) + assert status == "error", detail + assert "batch_upsert_from failed" in detail or "capacity" in detail.lower(), detail + + # --------------------------------------------------------------------------- # Automatic degradation (was the manual two-node fault_inject_multinode.py) # --------------------------------------------------------------------------- @@ -352,9 +692,12 @@ def _client(protocol: str): @pytest.mark.parametrize("protocol", ["tcp", "rdma"]) def test_multi_dtype_shape_roundtrip_is_byte_exact(self, protocol): tensors = { - "f32_2d": torch.randn(64, 1176, dtype=torch.float32), - "bf16_2d": torch.randn(32, 512).to(torch.bfloat16), - "i64_1d": torch.arange(4096, dtype=torch.int64), + # Production multimodal field names, dimensions, and mixed dtypes. + "pixel_values": torch.randn(64, 1176, dtype=torch.float32).to(torch.bfloat16), + "image_grid_thw": torch.tensor([[1, 8, 8]], dtype=torch.int64), + "input_ids": torch.arange(4096, dtype=torch.int64), + "attention_mask": torch.ones(4096, dtype=torch.int64), + "rewards": torch.linspace(-1, 1, 64, dtype=torch.float32), "noncontig": torch.randn(128, 256).t(), # transposed == non-contiguous } client = self._client(protocol) From 2b14c07549bf899c5f4472011bc0a0a382dad47e Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Thu, 13 Aug 2026 20:32:46 +0800 Subject: [PATCH 03/22] fix(data-plane): enforce TQ correctness contracts 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. --- relax/core/controller.py | 25 +-- relax/utils/rdma_probe.py | 57 ++++--- relax/utils/tq_config.py | 31 ++-- relax/utils/tq_correctness.py | 179 ++++++++++++++++++++++ relax/utils/tq_lifecycle.py | 80 +++++++++- scripts/benchmarks/tq_cross_node_bench.py | 7 +- scripts/benchmarks/tq_rdma_bench.py | 29 +++- tests/utils/test_rdma_probe.py | 59 +++++++ tests/utils/test_tq_failure_paths.py | 175 ++++++++++++++++++++- 9 files changed, 583 insertions(+), 59 deletions(-) create mode 100644 relax/utils/tq_correctness.py diff --git a/relax/core/controller.py b/relax/core/controller.py index 21f489503..cf2934e42 100644 --- a/relax/core/controller.py +++ b/relax/core/controller.py @@ -335,7 +335,8 @@ def _resolve_tq_backend(self, total_storage_size: int) -> dict: mode = getattr(self.config, "tq_rdma_mode", "off") # 2. SimpleStorage short-circuit (default, zero behavior change). - if backend == "simple" or mode == "off": + # ``mooncake + off`` is MooncakeStore/TCP, not SimpleStorage. + if backend == "simple": from relax.utils.tq_config import build_simple_storage_config return build_simple_storage_config( @@ -351,9 +352,9 @@ def _resolve_tq_backend(self, total_storage_size: int) -> dict: try: validate_mooncake_runtime_contract() except RuntimeError as e: - if mode == "required": + if mode != "auto": raise RuntimeError( - "--tq-rdma-mode=required but the installed TransferQueue " + f"--tq-rdma-mode={mode} but the installed TransferQueue " f"does not satisfy the Mooncake correctness contract: {e}" ) from e from relax.utils.tq_config import build_simple_storage_config @@ -367,7 +368,7 @@ def _resolve_tq_backend(self, total_storage_size: int) -> dict: num_data_storage_units=self.config.num_data_storage_units, ) master_address = resolve_mooncake_master_address() - probe_results = probe_cluster_nodes(device, master_address) + probe_results = probe_cluster_nodes(device, master_address, probe_rdma=mode != "off") for r in probe_results: logger.debug(r.summary()) @@ -376,13 +377,16 @@ def _resolve_tq_backend(self, total_storage_size: int) -> dict: requested_backend=backend, requested_device=device, use_gdr=getattr(self.config, "tq_use_gdr", False), + rdma_mode=mode, ) - # 4. required mode: fail fast instead of silently degrading (probe level). - if mode == "required" and effective.fallback_reason: + # 4. Only auto mode may degrade. ``off`` explicitly requests + # Mooncake/TCP, while ``required`` explicitly requires RDMA. + if mode != "auto" and effective.fallback_reason: detail = "\n".join(r.summary() for r in probe_results) raise RuntimeError( - f"--tq-rdma-mode=required but RDMA probe failed: {effective.fallback_reason}.\n" + f"--tq-rdma-mode={mode} but the requested Mooncake path is unavailable: " + f"{effective.fallback_reason}.\n" f"Probe details:\n{detail}" ) @@ -390,9 +394,10 @@ def _resolve_tq_backend(self, total_storage_size: int) -> dict: backend_dict, cap_error = build_backend_config(self.config, effective, total_storage_size=total_storage_size) actual_backend = "MooncakeStore" if "MooncakeStore" in backend_dict else "SimpleStorage" - # 6. required mode: also fail fast on capacity-induced fallback. - if mode == "required" and cap_error: - raise RuntimeError(f"--tq-rdma-mode=required but segment capacity insufficient: {cap_error}") + # 6. Capacity fallback is also auto-only. Explicit Mooncake/TCP and + # required-RDMA requests fail instead of silently changing backend. + if mode != "auto" and cap_error: + raise RuntimeError(f"--tq-rdma-mode={mode} but segment capacity insufficient: {cap_error}") # 7. GDR is EXPERIMENTAL in this phase: --tq-rdma-mode=required only # covers the *transport* (MooncakeStore + RDMA), not GDR. The probe diff --git a/relax/utils/rdma_probe.py b/relax/utils/rdma_probe.py index d495cbd75..6a1e0c531 100644 --- a/relax/utils/rdma_probe.py +++ b/relax/utils/rdma_probe.py @@ -238,30 +238,31 @@ def _check_health_check() -> CheckResult: # pragma: no cover - retained for ad- # Per-node probe # --------------------------------------------------------------------------- + # NOTE: ``health_check()`` is intentionally NOT probed because it depends on # local Mooncake client initialization. External-master reachability is checked # directly with a bounded TCP connect, then authoritatively by ``tq.init``. -_CHECK_FUNCS_NO_DEVICE = [ - _check_mooncake_import, - _check_rdma_devices, - _check_memlock, -] - - -def probe_node(device: str = "", master_address: str = "") -> ProbeResult: +def probe_node(device: str = "", master_address: str = "", *, probe_rdma: bool = True) -> ProbeResult: """Run all capability checks on the current node. Parameters ---------- device Explicit RDMA device name; empty = auto-detect first available. + probe_rdma + When ``False``, validate only Mooncake importability and master + reachability, then select TCP. Used by ``--tq-rdma-mode=off`` so an + explicitly requested Mooncake/TCP backend never depends on RDMA + hardware. """ node = socket.gethostname() checks: list[CheckResult] = [] errors: list[str] = [] - for fn in _CHECK_FUNCS_NO_DEVICE: - checks.append(fn()) + checks.append(_check_mooncake_import()) + if probe_rdma: + checks.append(_check_rdma_devices()) + checks.append(_check_memlock()) # Mooncake is externally managed by Relax deployments. When an endpoint is # supplied, it must already be reachable from every data-plane node before @@ -269,9 +270,10 @@ def probe_node(device: str = "", master_address: str = "") -> ProbeResult: if master_address: checks.append(_check_master_reachable(master_address)) - # Device-dependent checks. - checks.append(_check_port_active(device)) - checks.append(_check_gid_available(device)) + # Device-dependent checks are irrelevant when RDMA is explicitly off. + if probe_rdma: + checks.append(_check_port_active(device)) + checks.append(_check_gid_available(device)) # Determine effective protocol via graded degradation. mooncake_ok = any(c.name == "mooncake_import" and c.ok for c in checks) @@ -282,13 +284,15 @@ def probe_node(device: str = "", master_address: str = "") -> ProbeResult: master_ok = not master_address or any(c.name == "master_reachable" and c.ok for c in checks) effective_protocol: str | None - effective_device = device + effective_device = device if probe_rdma else "" if not mooncake_ok: effective_protocol = None errors.append("mooncake not importable") elif not master_ok: effective_protocol = None errors.append("master unreachable") + elif not probe_rdma: + effective_protocol = "tcp" elif rdma_dev_ok and port_ok and gid_ok and memlock_ok: effective_protocol = "rdma" if not effective_device: @@ -382,6 +386,7 @@ def probe_cluster_nodes( master_address: str = "", *, timeout: float = 60.0, + probe_rdma: bool = True, ) -> list[ProbeResult]: """Probe every alive GPU-bearing node and return one result per node. @@ -402,7 +407,9 @@ def probe_cluster_nodes( driver result when no GPU workers are discoverable (single-node/local dev). """ node_ids = _alive_gpu_nodes() - driver_result = probe_node(device, master_address) + driver_result = ( + probe_node(device, master_address) if probe_rdma else probe_node(device, master_address, probe_rdma=False) + ) if not node_ids: logger.debug("No alive GPU nodes discovered; probing driver node only.") return [driver_result] @@ -411,16 +418,16 @@ def probe_cluster_nodes( from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy @ray.remote(num_cpus=0.001) - def _probe_on_node(dev: str, master: str) -> ProbeResult: + def _probe_on_node(dev: str, master: str, should_probe_rdma: bool) -> ProbeResult: from relax.utils.rdma_probe import probe_node as _probe - return _probe(dev, master) + return _probe(dev, master, probe_rdma=should_probe_rdma) refs: list[Any] = [] id_by_ref: dict[Any, str] = {} for node_id in node_ids: strategy = NodeAffinitySchedulingStrategy(node_id=node_id, soft=False) - ref = _probe_on_node.options(scheduling_strategy=strategy).remote(device, master_address) + ref = _probe_on_node.options(scheduling_strategy=strategy).remote(device, master_address, probe_rdma) refs.append(ref) id_by_ref[ref] = node_id @@ -454,6 +461,7 @@ def reduce_results( requested_device: str, use_gdr: bool, fallback_backend: str = "SimpleStorage", + rdma_mode: str = "auto", ) -> EffectiveConfig: """AND-reduce per-node results into a single job-level effective config. @@ -469,6 +477,10 @@ def reduce_results( ``--tq-use-gdr`` value. fallback_backend Backend to degrade to when probe fails in auto mode. + rdma_mode + ``off`` selects Mooncake/TCP after validating Mooncake and master + availability. ``auto`` and ``required`` reduce the probed RDMA + capability normally; the caller decides whether a fallback is fatal. """ # SimpleStorage short-circuits: no probing needed. if requested_backend == "simple": @@ -510,6 +522,15 @@ def reduce_results( fallback_reason=reason, ) + if rdma_mode == "off": + return EffectiveConfig( + backend="MooncakeStore", + protocol="tcp", + device="", + gdr=False, + fallback_reason="", + ) + if all_rdma: # Device: if any node lacks the requested device, fall back to tcp. if requested_device: diff --git a/relax/utils/tq_config.py b/relax/utils/tq_config.py index 909a7fa88..50193d7fc 100644 --- a/relax/utils/tq_config.py +++ b/relax/utils/tq_config.py @@ -20,6 +20,7 @@ from relax.utils.logging_utils import get_logger from relax.utils.rdma_probe import EffectiveConfig +from relax.utils.tq_correctness import ensure_mooncake_correctness_guards logger = get_logger(__name__) @@ -41,23 +42,16 @@ def resolve_mooncake_master_address() -> str: def validate_mooncake_runtime_contract() -> None: - """Fail fast unless installed TransferQueue has the loss-prevention fixes. + """Install and validate the Mooncake loss-prevention contract. - The RDMA integration depends on per-key put/get result validation and on - notifying production readiness only after storage succeeds. A version - number alone is insufficient for development builds, so validate the - concrete runtime capabilities before probing or creating a controller. + A version number alone is insufficient for development builds. Relax + therefore installs process-local guards that validate every batch response, + propagate removal failures, and require a positive production-status ACK. + Every process calls this before creating or attaching a Mooncake client. """ - try: - from transfer_queue.storage.clients.mooncake_client import MooncakeStoreClient - from transfer_queue.storage.managers.base import KVStorageManager - except ImportError as e: - raise RuntimeError("Installed TransferQueue has no MooncakeStore support") from e + ensure_mooncake_correctness_guards() - required_methods = ("_batch_upsert_with_retry", "_batch_get_into_with_retry") - missing = [name for name in required_methods if not callable(getattr(MooncakeStoreClient, name, None))] - if missing: - raise RuntimeError("Installed TransferQueue lacks required Mooncake failure handling: " + ", ".join(missing)) + from transfer_queue.storage.managers.base import KVStorageManager put_source = inspect.getsource(KVStorageManager.put_data) storage_call = put_source.find("self.storage_client.put") @@ -73,9 +67,12 @@ def validate_mooncake_runtime_contract() -> None: # --------------------------------------------------------------------------- -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).""" +def build_simple_storage_config(total_storage_size: int | None, num_data_storage_units: int) -> dict[str, Any]: + """Build the SimpleStorage backend config. + + ``total_storage_size=None`` preserves TransferQueue's unlimited-capacity + benchmark semantics; production Relax jobs pass a concrete sample count. + """ return { # ``tq.init`` selects the manager from this key alone (TQ config.yaml:22 # defaults it to SimpleStorage); a backend section without it is ignored. diff --git a/relax/utils/tq_correctness.py b/relax/utils/tq_correctness.py new file mode 100644 index 000000000..9b2b1d32d --- /dev/null +++ b/relax/utils/tq_correctness.py @@ -0,0 +1,179 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Fail-closed correctness guards for TransferQueue's Mooncake backend. + +The TransferQueue revision currently pinned by Relax validates the first +Mooncake batch result, but not every retry result, logs removal failures +without raising, and treats a missing/negative production-status ACK as a +successful notification. Those behaviours can turn an explicit storage or +controller failure into silent data loss. + +Keep the compatibility guards here, close to Relax's integration boundary, +until the equivalent checks are available in the pinned TransferQueue +revision. Installation is process-local and idempotent; every process that +creates or attaches a Mooncake client installs them before ``tq.init``. +""" + +from __future__ import annotations + +import asyncio +from functools import wraps +from typing import Any +from uuid import uuid4 + +from relax.utils.logging_utils import get_logger + + +logger = get_logger(__name__) + +_PATCH_MARKER = "_relax_mooncake_correctness_guards_v1" + + +def _validate_result_count(operation: str, keys: list[str], results: Any) -> None: + """Require one Mooncake result code for every requested key.""" + try: + actual = len(results) + except TypeError as error: + raise RuntimeError(f"{operation} returned a non-sized result, expected {len(keys)} codes") from error + if actual != len(keys): + raise RuntimeError(f"{operation} returned {actual} results, expected {len(keys)}") + + +class _StrictMooncakeStoreProxy: + """Validate every low-level batch response, including retry calls.""" + + def __init__(self, store: Any) -> None: + self._store = store + + def __getattr__(self, name: str) -> Any: + return getattr(self._store, name) + + def batch_upsert_from(self, keys: list[str], *args: Any, **kwargs: Any) -> Any: + results = self._store.batch_upsert_from(keys, *args, **kwargs) + _validate_result_count("batch_upsert_from", keys, results) + return results + + def batch_get_into(self, keys: list[str], *args: Any, **kwargs: Any) -> Any: + results = self._store.batch_get_into(keys, *args, **kwargs) + _validate_result_count("batch_get_into", keys, results) + return results + + def batch_remove(self, keys: list[str], *args: Any, **kwargs: Any) -> Any: + results = self._store.batch_remove(keys, *args, **kwargs) + _validate_result_count("batch_remove", keys, results) + failures = [(key, code) for key, code in zip(keys, results, strict=True) if code != 0] + if failures: + detail = ", ".join(f"{key}={code}" for key, code in failures) + raise RuntimeError(f"batch_remove failed: {detail}") + return results + + +async def _strict_notify_and_wait(self: Any, request_msg: list) -> None: + """Notify the controller and require a positive ACK within the deadline.""" + import zmq + + from transfer_queue.storage.managers import base as tq_base + from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, create_zmq_socket + + identity = f"{self.storage_manager_id}-notify-{uuid4().hex[:8]}".encode() + sock = create_zmq_socket( + ctx=self.zmq_context, + socket_type=zmq.DEALER, + ip=self.controller_info.ip, + identity=identity, + ) + sock.setsockopt(zmq.LINGER, 0) + sock.connect(self.controller_info.to_addr("request_handle_socket")) + + try: + await sock.send_multipart(request_msg) + loop = asyncio.get_running_loop() + deadline = loop.time() + tq_base.TQ_DATA_UPDATE_RESPONSE_TIMEOUT + + while True: + remaining = deadline - loop.time() + if remaining <= 0: + raise TimeoutError( + "Timed out waiting for TransferQueue production-status ACK " + f"after {tq_base.TQ_DATA_UPDATE_RESPONSE_TIMEOUT}s" + ) + try: + messages = await asyncio.wait_for( + sock.recv_multipart(copy=False), + timeout=min(tq_base.TQ_STORAGE_POLLER_TIMEOUT, remaining), + ) + except asyncio.TimeoutError: + continue + except Exception as error: + raise RuntimeError("Failed while waiting for TransferQueue production-status ACK") from error + + response = ZMQMessage.deserialize(messages) + if response.request_type != ZMQRequestType.NOTIFY_DATA_UPDATE_ACK: + continue + body = response.body if isinstance(response.body, dict) else {} + if body.get("success") is not True: + raise RuntimeError( + "TransferQueue controller rejected the production-status update " + f"for partition={body.get('partition_id', 'unknown')}" + ) + return + finally: + try: + if not sock.closed: + sock.close(linger=0) + except Exception as error: # pragma: no cover - best-effort socket cleanup + logger.debug(f"Failed to close TransferQueue notification socket: {error}") + + +def _install_store_guards(client_cls: type) -> None: + if getattr(client_cls, _PATCH_MARKER, False): + return + + original_init = client_cls.__init__ + + @wraps(original_init) + def guarded_init(self: Any, *args: Any, **kwargs: Any) -> None: + original_init(self, *args, **kwargs) + store = getattr(self, "_store", None) + if store is not None and not isinstance(store, _StrictMooncakeStoreProxy): + self._store = _StrictMooncakeStoreProxy(store) + + client_cls.__init__ = guarded_init + setattr(client_cls, _PATCH_MARKER, True) + + +def _install_notification_guards(manager_cls: type) -> None: + if getattr(manager_cls, _PATCH_MARKER, False): + return + + original_notify = manager_cls.notify_data_update + + @wraps(original_notify) + async def guarded_notify(self: Any, *args: Any, **kwargs: Any) -> None: + if not getattr(self, "controller_info", None): + raise RuntimeError("TransferQueue storage manager has no controller for production-status notification") + await original_notify(self, *args, **kwargs) + + manager_cls.notify_data_update = guarded_notify + manager_cls._notify_and_wait = _strict_notify_and_wait + setattr(manager_cls, _PATCH_MARKER, True) + + +def ensure_mooncake_correctness_guards() -> None: + """Install and validate all guards required for safe Mooncake operation.""" + try: + from transfer_queue.storage.clients.mooncake_client import MooncakeStoreClient + from transfer_queue.storage.managers.base import StorageManager + except ImportError as error: + raise RuntimeError("Installed TransferQueue has no MooncakeStore support") from error + + required_methods = ("_batch_upsert_with_retry", "_batch_get_into_with_retry") + missing = [name for name in required_methods if not callable(getattr(MooncakeStoreClient, name, None))] + if missing: + raise RuntimeError("Installed TransferQueue lacks required Mooncake retry APIs: " + ", ".join(missing)) + + _install_store_guards(MooncakeStoreClient) + _install_notification_guards(StorageManager) + + if not getattr(MooncakeStoreClient, _PATCH_MARKER, False) or not getattr(StorageManager, _PATCH_MARKER, False): + raise RuntimeError("Failed to install Mooncake silent-data-loss guards") diff --git a/relax/utils/tq_lifecycle.py b/relax/utils/tq_lifecycle.py index dcea6a79f..436d6cd43 100644 --- a/relax/utils/tq_lifecycle.py +++ b/relax/utils/tq_lifecycle.py @@ -61,6 +61,66 @@ class TqCleanupTimeout(TimeoutError): """Raised when a TQ controller cannot be confirmed gone after cleanup.""" +class TqConfigurationMismatch(RuntimeError): + """Raised when an existing controller uses a different backend config.""" + + +def _get_config_value(config: Any, key: str, default: Any = None) -> Any: + if config is None: + return default + if hasattr(config, "get"): + return config.get(key, default) + return getattr(config, key, default) + + +def _backend_signature(conf: Any) -> tuple[Any, ...]: + """Return the backend fields that must agree for a safe attach.""" + backend = _get_config_value(conf, "backend", {}) + storage_backend = _get_config_value(backend, "storage_backend", "SimpleStorage") + if storage_backend == "MooncakeStore": + mooncake = _get_config_value(backend, "MooncakeStore", {}) + return ( + storage_backend, + _get_config_value(mooncake, "protocol", "tcp"), + _get_config_value(mooncake, "device_name", "") or "", + _get_config_value(mooncake, "master_server_address", ""), + _get_config_value(mooncake, "metadata_server", ""), + _get_config_value(mooncake, "global_segment_size"), + _get_config_value(mooncake, "local_buffer_size"), + bool(_get_config_value(mooncake, "hard_pin", False)), + bool(_get_config_value(mooncake, "use_gdr", False)), + ) + simple = _get_config_value(backend, "SimpleStorage", {}) + return ( + "SimpleStorage", + _get_config_value(simple, "total_storage_size"), + _get_config_value(simple, "num_data_storage_units"), + ) + + +def _backend_description(conf: Any) -> str: + """Describe the backend without logging endpoints or host information.""" + backend = _get_config_value(conf, "backend", {}) + storage_backend = _get_config_value(backend, "storage_backend", "SimpleStorage") + if storage_backend != "MooncakeStore": + return "SimpleStorage" + mooncake = _get_config_value(backend, "MooncakeStore", {}) + return f"MooncakeStore/{_get_config_value(mooncake, 'protocol', 'tcp')}" + + +def _uses_mooncake(conf: Any) -> bool: + backend = _get_config_value(conf, "backend", {}) + return _get_config_value(backend, "storage_backend", "SimpleStorage") == "MooncakeStore" + + +def _prepare_mooncake_runtime(conf: Any) -> None: + if not _uses_mooncake(conf): + return + from relax.utils.tq_config import validate_mooncake_runtime_contract + + validate_mooncake_runtime_contract() + + def kill_tq_controller_and_wait(timeout: float = 20.0) -> None: """Kill the TransferQueueController named actor, then wait for GCS deregistration. @@ -218,6 +278,7 @@ def log_tq_gdr_runtime_status(*, requested: bool, role: str) -> str: def attach_tq_client(conf: Any, *, requested_gdr: bool, role: str) -> Any: """Attach a component process and report its local experimental GDR state.""" + _prepare_mooncake_runtime(conf) tq.init(conf=conf) client = tq.get_client() log_tq_gdr_runtime_status(requested=requested_gdr, role=role) @@ -264,6 +325,7 @@ def __init__(self) -> None: self._owns_controller = False def initialize(self, conf: Any, owner_token: str) -> tuple[Any, bool]: + _prepare_mooncake_runtime(conf) _set_owner_token(conf, owner_token) tq.init(conf=conf) stored_conf = _get_stored_config() @@ -335,10 +397,17 @@ def _start_owner(conf: Any, *, timeout: float) -> TqInitResult: # A concurrent initializer won the named-actor race. This process is only # attached and must never retain an actor capable of global tq.close(). + config_mismatch = _backend_signature(stored_conf) != _backend_signature(conf) try: ray.get(owner.detach.remote(), timeout=10.0) finally: _stop_owner_actor(owner) + if config_mismatch: + raise TqConfigurationMismatch( + "A concurrent TransferQueue initializer won with a different backend config " + f"(requested={_backend_description(conf)}, stored={_backend_description(stored_conf)}). " + "Detached without modifying the winning controller." + ) return TqInitResult(config=stored_conf, owner=None) @@ -381,13 +450,20 @@ def _attempt(attempt_conf: Any) -> TqInitResult: if _controller_exists(): # Attach semantics: use the controller's actual config. Upstream # tq.init(conf) returns the caller-provided config even when ignored. - return TqInitResult(config=_get_stored_config(), owner=None) + stored_conf = _get_stored_config() + if _backend_signature(stored_conf) != _backend_signature(attempt_conf): + raise TqConfigurationMismatch( + "Refusing to attach to an existing TransferQueueController with a different backend config " + f"(requested={_backend_description(attempt_conf)}, stored={_backend_description(stored_conf)}). " + "Only the owner may close the existing controller." + ) + return TqInitResult(config=stored_conf, owner=None) return _start_owner(attempt_conf, timeout=timeout) try: return _attempt(conf) except Exception as primary_error: - if mode != "auto" or fallback_conf is None: + if isinstance(primary_error, TqConfigurationMismatch) or mode != "auto" or fallback_conf is None: raise reason = f"mooncake_init_failed:{type(primary_error).__name__}" diff --git a/scripts/benchmarks/tq_cross_node_bench.py b/scripts/benchmarks/tq_cross_node_bench.py index 315e6d356..67193503a 100644 --- a/scripts/benchmarks/tq_cross_node_bench.py +++ b/scripts/benchmarks/tq_cross_node_bench.py @@ -277,12 +277,17 @@ def build_conf(protocol: str, master: str, device: str, segment_gib: int): from transfer_queue import GRPOGroupNSampler from relax.utils.rdma_probe import EffectiveConfig - from relax.utils.tq_config import build_mooncake_config, build_simple_storage_config + from relax.utils.tq_config import ( + build_mooncake_config, + build_simple_storage_config, + validate_mooncake_runtime_contract, + ) if protocol == "simple": # total_storage_size=None == unlimited sample count (TQ config.yaml default). backend = build_simple_storage_config(total_storage_size=None, num_data_storage_units=2) else: + validate_mooncake_runtime_contract() eff = EffectiveConfig(backend="MooncakeStore", protocol=protocol, device=device, gdr=False, fallback_reason="") backend = build_mooncake_config(eff, master_address=master, global_segment_size=segment_gib * 1024**3) return OmegaConf.create( diff --git a/scripts/benchmarks/tq_rdma_bench.py b/scripts/benchmarks/tq_rdma_bench.py index 6133dbae2..8be5243d0 100644 --- a/scripts/benchmarks/tq_rdma_bench.py +++ b/scripts/benchmarks/tq_rdma_bench.py @@ -205,6 +205,11 @@ def run_one(config_name: str, payload: dict, args: argparse.Namespace) -> dict: """Run put/get once and return timing.""" import transfer_queue as tq + if CONFIG_MAP[config_name]["backend"] == "MooncakeStore": + from relax.utils.tq_config import validate_mooncake_runtime_contract + + validate_mooncake_runtime_contract() + # Close any prior controller and wait for GCS deregistration so this config # gets a fresh backend (tq.init otherwise attaches to the existing one). close_tq_and_wait() @@ -299,11 +304,9 @@ def run_config(config_name: str, payload: dict, args: argparse.Namespace) -> dic # --------------------------------------------------------------------------- # -def main(): +def run_benchmark(args: argparse.Namespace) -> None: """Run the benchmark across all requested payload/field/config combinations.""" - args = parse_args() - print("=" * 80) print("TransferQueue RDMA Benchmark") print(f" configs: {args.configs}") @@ -386,11 +389,25 @@ def main(): writer.writerows(all_results) print(f"\nCSV written to {args.output_csv}") - # Final cleanup: tear down the last config's controller so a subsequent - # benchmark run starts from a clean slate. - close_tq_and_wait() print("\n[dataplane] benchmark complete") +def main() -> None: + """Initialize Ray before touching named actors and always clean up.""" + import ray + + args = parse_args() + ray.init(ignore_reinit_error=True) + try: + run_benchmark(args) + finally: + # Tear down the last config's controller so a subsequent benchmark run + # starts from a clean slate, then release the local Ray runtime. + try: + close_tq_and_wait() + finally: + ray.shutdown() + + if __name__ == "__main__": main() diff --git a/tests/utils/test_rdma_probe.py b/tests/utils/test_rdma_probe.py index 24dc161f2..57e36c3e3 100644 --- a/tests/utils/test_rdma_probe.py +++ b/tests/utils/test_rdma_probe.py @@ -10,8 +10,11 @@ from __future__ import annotations import argparse +import importlib.util from unittest import mock +import pytest + import relax.utils.rdma_probe as rdma_probe from relax.utils.rdma_probe import ( CheckResult, @@ -40,6 +43,16 @@ # --------------------------------------------------------------------------- +def _has_real_tq_storage() -> bool: + try: + return importlib.util.find_spec("transfer_queue.storage.clients.mooncake_client") is not None + except (ImportError, TypeError, ValueError): + return False + + +_REAL_TQ_STORAGE = _has_real_tq_storage() + + def _make_probe( protocol: str | None = "rdma", device: str = "rdma0", @@ -187,6 +200,28 @@ def test_empty_results_falls_back(self): ) assert eff.backend == "SimpleStorage" + def test_off_mode_keeps_mooncake_and_selects_tcp(self): + eff = reduce_results( + [_make_probe(protocol="rdma"), _make_probe(protocol="tcp", node="node-B")], + requested_backend="mooncake", + requested_device="rdma0", + use_gdr=False, + rdma_mode="off", + ) + assert (eff.backend, eff.protocol, eff.device) == ("MooncakeStore", "tcp", "") + assert eff.fallback_reason == "" + + def test_off_mode_reports_mooncake_unavailable(self): + eff = reduce_results( + [_make_probe(protocol=None)], + requested_backend="mooncake", + requested_device="", + use_gdr=False, + rdma_mode="off", + ) + assert eff.backend == "SimpleStorage" + assert "mooncake_unavailable" in eff.fallback_reason + # --------------------------------------------------------------------------- # probe_node (mocked filesystem) @@ -245,6 +280,22 @@ def test_master_reachability_is_bounded_failure(self): result = _check_master_reachable("127.0.0.1:1", timeout=0.01) assert result.ok is False + def test_off_mode_probe_does_not_touch_rdma_hardware(self, monkeypatch): + monkeypatch.setattr(rdma_probe, "_check_mooncake_import", lambda: CheckResult("mooncake_import", True)) + monkeypatch.setattr( + rdma_probe, + "_check_master_reachable", + lambda address: CheckResult("master_reachable", True, address), + ) + monkeypatch.setattr( + rdma_probe, + "_check_rdma_devices", + lambda: (_ for _ in ()).throw(AssertionError("RDMA probe must not run")), + ) + result = probe_node("", "master.example:50051", probe_rdma=False) + assert result.effective_protocol == "tcp" + assert {check.name for check in result.checks} == {"mooncake_import", "master_reachable"} + # --------------------------------------------------------------------------- # probe_cluster_nodes (multi-node fan-out) + helpers @@ -334,6 +385,10 @@ def test_simple_storage_config(self): "SimpleStorage": {"total_storage_size": 1000, "num_data_storage_units": 2}, } + def test_simple_storage_config_allows_unlimited_capacity(self): + cfg = build_simple_storage_config(total_storage_size=None, num_data_storage_units=2) + assert cfg["SimpleStorage"]["total_storage_size"] is None + def test_storage_backend_key_selects_the_manager(self): """``tq.init`` reads ``backend.storage_backend``; omitting it silently keeps SimpleStorage.""" @@ -359,6 +414,10 @@ def test_mooncake_config_gdr_propagated(self): cfg = build_mooncake_config(eff) assert cfg["MooncakeStore"]["use_gdr"] is True + @pytest.mark.skipif( + not _REAL_TQ_STORAGE, + reason="needs real TransferQueue storage submodules; CPU CI uses a single-file transfer_queue stub", + ) def test_installed_tq_satisfies_loss_prevention_contract(self): validate_mooncake_runtime_contract() diff --git a/tests/utils/test_tq_failure_paths.py b/tests/utils/test_tq_failure_paths.py index b6f873732..e04cd8111 100644 --- a/tests/utils/test_tq_failure_paths.py +++ b/tests/utils/test_tq_failure_paths.py @@ -24,6 +24,7 @@ import multiprocessing import os import socket +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock import pytest @@ -31,6 +32,7 @@ from relax.utils import tq_lifecycle from relax.utils.rdma_probe import ProbeResult, reduce_results +from relax.utils.tq_correctness import _strict_notify_and_wait, _StrictMooncakeStoreProxy def _has_real_submodule(dotted: str) -> bool: @@ -180,11 +182,7 @@ class TestCloseTqAndUnmount: def _fake_tq(monkeypatch, store_client): calls: list[str] = [] fake = MagicMock() - manager = MagicMock() - if store_client is None: - del manager.storage_client # SimpleStorage manager has no storage_client - else: - manager.storage_client = store_client + manager = SimpleNamespace() if store_client is None else SimpleNamespace(storage_client=store_client) fake.get_client.return_value = MagicMock(storage_manager=manager) fake.close.side_effect = lambda: calls.append("tq.close") monkeypatch.setattr(tq_lifecycle, "tq", fake) @@ -265,6 +263,20 @@ class TestInitializeTqWithFallback: def _conf(backend: str) -> dict: return {"controller": {}, "backend": {"storage_backend": backend}} + @staticmethod + def _mooncake_conf(protocol: str) -> dict: + return { + "controller": {}, + "backend": { + "storage_backend": "MooncakeStore", + "MooncakeStore": { + "protocol": protocol, + "master_server_address": "master.invalid:50051", + "hard_pin": True, + }, + }, + } + @staticmethod def _patch_transaction(monkeypatch, *, existed: bool, init_effects: list[object], stored_conf=None): calls: dict[str, list] = {"reap": [], "attempts": []} @@ -300,6 +312,35 @@ def test_attach_is_not_owner(self, monkeypatch): assert result.config is stored assert calls["attempts"] == [] + def test_attach_rejects_different_backend_without_closing_owner(self, monkeypatch): + requested = self._mooncake_conf("rdma") + stored = self._conf("SimpleStorage") + calls = self._patch_transaction(monkeypatch, existed=True, init_effects=[], stored_conf=stored) + with pytest.raises(tq_lifecycle.TqConfigurationMismatch, match="different backend config"): + tq_lifecycle.initialize_tq_with_fallback( + requested, + mode="auto", + fallback_conf=self._conf("SimpleStorage"), + ) + assert calls["attempts"] == [] + + def test_attach_rejects_different_mooncake_protocol(self, monkeypatch): + requested = self._mooncake_conf("rdma") + stored = self._mooncake_conf("tcp") + calls = self._patch_transaction(monkeypatch, existed=True, init_effects=[], stored_conf=stored) + with pytest.raises(tq_lifecycle.TqConfigurationMismatch, match="requested=MooncakeStore/rdma"): + tq_lifecycle.initialize_tq_with_fallback(requested, mode="required") + assert calls["attempts"] == [] + + def test_attach_accepts_matching_mooncake_config(self, monkeypatch): + requested = self._mooncake_conf("rdma") + stored = self._mooncake_conf("rdma") + calls = self._patch_transaction(monkeypatch, existed=True, init_effects=[], stored_conf=stored) + result = tq_lifecycle.initialize_tq_with_fallback(requested, mode="required") + assert result.config is stored + assert result.owns_controller is False + assert calls["attempts"] == [] + def test_auto_cleans_failed_mooncake_then_retries_simple_once(self, monkeypatch): primary = self._conf("MooncakeStore") fallback = self._conf("SimpleStorage") @@ -376,6 +417,23 @@ def timed_out(ref, *, timeout): assert cleaned[0][0] is owner assert cleaned[0][1] + def test_concurrent_initializer_with_different_config_detaches_and_fails(self, monkeypatch): + owner = _FakeOwner() + stopped: list[object] = [] + requested = TestInitializeTqWithFallback._mooncake_conf("rdma") + stored = TestInitializeTqWithFallback._conf("SimpleStorage") + monkeypatch.setattr(tq_lifecycle._TransferQueueOwner, "remote", lambda: owner) + monkeypatch.setattr( + tq_lifecycle.ray, + "get", + lambda ref, timeout: (stored, False) if ref == "initialize-ref" else None, + ) + monkeypatch.setattr(tq_lifecycle, "_stop_owner_actor", lambda handle: stopped.append(handle)) + + with pytest.raises(tq_lifecycle.TqConfigurationMismatch, match="concurrent TransferQueue initializer"): + tq_lifecycle._start_owner(requested, timeout=1) + assert stopped == [owner] + @pytest.mark.parametrize("stored_token,should_kill", [("ours", True), ("theirs", False)]) def test_failed_owner_cleanup_respects_controller_owner_token(self, monkeypatch, stored_token, should_kill): owner = _FakeOwner() @@ -459,6 +517,113 @@ def _client_with_store(store) -> object: return client +class _SequenceStore: + """Return a configured result sequence from low-level Mooncake calls.""" + + def __init__(self, results: list[list[int]]) -> None: + self.results = iter(results) + + def batch_upsert_from(self, keys, ptrs, sizes, config=None): + return next(self.results) + + def batch_get_into(self, keys, ptrs, sizes): + return next(self.results) + + def batch_remove(self, keys, force=True): + return next(self.results) + + +class _FakeNotifySocket: + def __init__(self) -> None: + self.closed = False + + def setsockopt(self, *args, **kwargs) -> None: + pass + + def connect(self, *args, **kwargs) -> None: + pass + + async def send_multipart(self, request) -> None: + pass + + async def recv_multipart(self, copy=False): + return [b"ack"] + + def close(self, linger=0) -> None: + self.closed = True + + +class TestMooncakeCorrectnessGuardPrimitives: + """Low-level response validation stays runnable on the CPU-only CI stub.""" + + def test_upsert_short_result_is_raised(self): + store = _StrictMooncakeStoreProxy(_SequenceStore([[0]])) + with pytest.raises(RuntimeError, match="returned 1 results, expected 2"): + store.batch_upsert_from(["k0", "k1"], [1, 2], [8, 8]) + + def test_remove_failure_is_raised(self): + store = _StrictMooncakeStoreProxy(_SequenceStore([[0, -704]])) + with pytest.raises(RuntimeError, match="batch_remove failed"): + store.batch_remove(["k0", "k1"], force=True) + + +@pytest.mark.skipif( + not _REAL_MOONCAKE_CLIENT, + reason="needs real TransferQueue storage submodules; CPU CI uses a single-file transfer_queue stub", +) +class TestMooncakeCorrectnessGuards: + """Integration with real TransferQueue internals; no GPU/master needed.""" + + def test_retry_short_result_is_never_treated_as_success(self, monkeypatch): + monkeypatch.setattr("transfer_queue.storage.clients.mooncake_client.RETRY_DELAY_SECONDS", 0) + store = _StrictMooncakeStoreProxy(_SequenceStore([[-1, -1], [0]])) + client = _client_with_store(store) + with pytest.raises(RuntimeError, match="returned 1 results, expected 2"): + client._batch_upsert_with_retry(["k0", "k1"], [1, 2], [8, 8]) + + @pytest.mark.asyncio + async def test_negative_production_status_ack_is_raised(self, monkeypatch): + from transfer_queue.utils import zmq_utils + + socket = _FakeNotifySocket() + monkeypatch.setattr(zmq_utils, "create_zmq_socket", lambda **kwargs: socket) + monkeypatch.setattr( + zmq_utils.ZMQMessage, + "deserialize", + staticmethod( + lambda messages: SimpleNamespace( + request_type=zmq_utils.ZMQRequestType.NOTIFY_DATA_UPDATE_ACK, + body={"success": False, "partition_id": "p0"}, + ) + ), + ) + manager = SimpleNamespace( + storage_manager_id="guard-test", + zmq_context=object(), + controller_info=SimpleNamespace(ip="redacted", to_addr=lambda name: "inproc://controller"), + ) + with pytest.raises(RuntimeError, match="rejected the production-status update"): + await _strict_notify_and_wait(manager, [b"request"]) + assert socket.closed is True + + @pytest.mark.asyncio + async def test_missing_production_status_ack_is_bounded(self, monkeypatch): + from transfer_queue.storage.managers import base as tq_base + from transfer_queue.utils import zmq_utils + + socket = _FakeNotifySocket() + monkeypatch.setattr(zmq_utils, "create_zmq_socket", lambda **kwargs: socket) + monkeypatch.setattr(tq_base, "TQ_DATA_UPDATE_RESPONSE_TIMEOUT", 0) + manager = SimpleNamespace( + storage_manager_id="guard-test", + zmq_context=object(), + controller_info=SimpleNamespace(ip="redacted", to_addr=lambda name: "inproc://controller"), + ) + with pytest.raises(TimeoutError, match="production-status ACK"): + await _strict_notify_and_wait(manager, [b"request"]) + assert socket.closed is True + + class _InlineExecutorLoop: """Run storage-manager sync calls inline for deterministic async tests. From 7e91be2aedf47c1c1ab8fac42264b12963544173 Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Fri, 14 Aug 2026 05:27:12 +0800 Subject: [PATCH 04/22] test(data-plane): real multimodal byte-exact tier MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # ✅ 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) --- .gitignore | 4 + docs/draft/transfer_queue_rdma.md | 25 +- relax/utils/payload_digest.py | 140 +++++++++ scripts/benchmarks/make_multimodal_fixture.py | 265 ++++++++++++++++++ scripts/benchmarks/tq_cross_node_bench.py | 193 +++++++++++-- tests/utils/mm_payload_fixtures.py | 107 +++++++ tests/utils/test_tq_dataplane_behavior.py | 71 +++++ tests/utils/test_tq_failure_paths.py | 81 +++++- 8 files changed, 861 insertions(+), 25 deletions(-) create mode 100644 relax/utils/payload_digest.py create mode 100644 scripts/benchmarks/make_multimodal_fixture.py create mode 100644 tests/utils/mm_payload_fixtures.py diff --git a/.gitignore b/.gitignore index e55036cb7..6c6d95fda 100644 --- a/.gitignore +++ b/.gitignore @@ -54,3 +54,7 @@ tensorboard_log # .github .github/copilot-instructions.md env.sh + +# Machine-local multimodal acceptance fixtures (generated by +# scripts/benchmarks/make_multimodal_fixture.py; hundreds of MB, never commit) +tests/fixtures/ diff --git a/docs/draft/transfer_queue_rdma.md b/docs/draft/transfer_queue_rdma.md index 8ea003b6a..4d7360d0e 100644 --- a/docs/draft/transfer_queue_rdma.md +++ b/docs/draft/transfer_queue_rdma.md @@ -123,9 +123,23 @@ Mock/本机测试和真实双节点 RDMA 测试必须分别报告，前者不能 | 层级 | 验证内容 | 通过标准 | |---|---|---| | CI/mock | 参数矩阵、节点 AND 归约、master 不可达、owner 超时/清理/token、auto/required、有限重试、写失败不发布状态 | `tests/utils/test_rdma_probe.py` 与 `tests/utils/test_tq_failure_paths.py` 全部通过；真机项允许明确 skip | -| 本机 TQ | SimpleStorage 全链路 put/get、容量 backpressure、空读、清理、字节一致性 | `tests/utils/test_tq_dataplane_behavior.py` 通过 | -| 真实 Mooncake | TCP/RDMA direct-client 多模态字段与混合 dtype 逐字节一致 | `TestMooncakeByteExact` 的 TCP/RDMA 两档均通过，不得 skip | -| 真实双节点 | 同一拓扑的 SimpleStorage、Mooncake/TCP、Mooncake/RDMA；synthetic 与 production-shaped multimodal；256/1024/2048/4096 MiB；每档 warmup + 至少 5 轮 | 每次 get 的逐字段 SHA-256 全部 PASS；`--require-wire-proof` 证明 RDMA 档 IB counter 增长且 TCP 档网络 counter 增长；CSV 留档并报告均值、median、stddev | +| 本机 TQ | SimpleStorage 全链路 put/get、容量 backpressure、空读、清理、字节一致性；`multimodal_train_inputs` 以生产容器（`list[dict]` / NonTensorStack，存储层非张量路径）全链路逐叶子 SHA-256 一致 | `tests/utils/test_tq_dataplane_behavior.py` 通过（含 `TestRealMultimodalFullLink`） | +| 真实 Mooncake | TCP/RDMA direct-client：混合 dtype 稠密张量逐字节一致；`list[dict]` 非张量 msgpack 慢路径逐叶子一致（每协议独立 spawn 子进程，规避 0.3.10 会话内协议切换问题） | `TestMooncakeByteExact` 的 TCP/RDMA 各两档均通过，不得 skip | +| 真实多模态载荷 | 真实数据集图像走完整生产预处理链（`build_messages` → `apply_chat_template` → `process_vision_info` → HF processor → `remap_mm_train_inputs`）生成 fixture；上述两级多模态用例检测到 fixture 后自动升级为真实载荷档 | fixture 存在时以 `[real]` 档通过；无 fixture 环境回退 `[synthetic]`（生产同构状，CI 兜底）；交付报告须注明真实档在何处跑过 | +| 真实双节点 | 同一拓扑的 SimpleStorage、Mooncake/TCP、Mooncake/RDMA；synthetic、production-shaped multimodal、real-multimodal 三种 profile；256/1024/2048/4096 MiB；每档 warmup + 至少 5 轮 | 每次 get 的逐字段 SHA-256 全部 PASS；`--require-wire-proof` 证明 RDMA 档 IB counter 增长且 TCP 档网络 counter 增长；CSV 留档并报告均值、median、stddev | + +真实多模态 fixture 生成（需要本地数据集 parquet 与 Qwen-VL 模型目录；产物写入 `tests/fixtures/`，已 gitignore，不入库）： + +```bash +PYTHONPATH=. python scripts/benchmarks/make_multimodal_fixture.py \ + --dataset /.parquet \ + --model / \ + --num-prompts 6 --n-samples-per-prompt 2 \ + --output tests/fixtures/tq_multimodal_fixture.pt \ + --manifest-json tests/fixtures/tq_multimodal_fixture.manifest.json +``` + +生成时自动做 processor 双跑字节一致性校验（以真实数据验证 F4 组共享假设）；fixture 载入时按叶子清单自校验，损坏即报错。测试与 bench 通过 `RELAX_MM_FIXTURE`（或默认路径 `tests/fixtures/tq_multimodal_fixture.pt`）发现 fixture。逐叶子指纹的规范实现在 `relax/utils/payload_digest.py`（哈希原始存储字节，对 NaN 也成立，严格强于 `torch.equal`）。 双节点验收命令（master 与 Ray 集群需由部署侧预先准备）： @@ -134,12 +148,14 @@ PYTHONPATH=. python -u scripts/benchmarks/tq_cross_node_bench.py \ --master :50051 \ --nodeb-ip \ --device \ - --payload-profiles synthetic multimodal \ + --payload-profiles synthetic multimodal real-multimodal \ --payload-mib 256 1024 2048 4096 \ --repeats 5 --require-wire-proof \ --csv tq_cross_node_acceptance.csv ``` +`real-multimodal` profile 按目标档位循环平铺 fixture 样本，`multimodal_train_inputs` 列以 NonTensorStack 走存储层非张量路径（SimpleStorage pickle / Mooncake msgpack），字节校验用行多重集指纹（采样器可重排行序；张量行按 dtype+字节比较以兼容后端间标量行 `()` 与 `[1]` 的表示差异，dict 叶子仍全形状校验）。 + 若本次开发环境没有两个 RDMA 节点，交付结论必须写成“真机验收未执行”，不能用 mock 通过推导真机已经通过。 ## 排障表 @@ -151,6 +167,7 @@ PYTHONPATH=. python -u scripts/benchmarks/tq_cross_node_bench.py \ | `setup failed with error code: -1` | master 不可达 | 检查 master 进程与 `MC_MASTER_ADDRESS` | | `Failed to open segment ... Connection refused` | 上一轮客户端异常退出，死 segment 仍在 master 注册 | 等 `client_ttl`（30 s）过期后重试 | | `batch_get_into failed ... error codes [-800, ...]` | 会话内切换协议（0.3.10 上更敏感），或对端不可达 | 每个协议单独进程跑；确认对端存活 | +| bench 的 Mooncake/TCP 档在单机回环下原生 SIGSEGV | mooncake 0.3.10 TCP transport 在单节点回环、64 MiB 级批量下崩溃（同机 direct-client 小批量正常，RDMA/SimpleStorage 同规模正常；响亮崩溃，非静默损坏） | C1 档在真实双节点拓扑上跑；单机开发环境用 simple/rdma 档验证 | | 多网卡机器跨节点建连失败 | 自动选卡选到了不通的网卡 | 显式 `--tq-rdma-device`；必要时用 `MC_TCP_BIND_ADDRESS` 指定 TCP 侧绑定地址 | | 训练卡在启动、无日志推进 | 半初始化的 controller（TQ 的 `_init_from_existing` 会无限轮询 config） | 本特性已加自动回收；若仍出现，确认 `[dataplane] ... reaping it` 是否打出 | diff --git a/relax/utils/payload_digest.py b/relax/utils/payload_digest.py new file mode 100644 index 000000000..cef6f4d86 --- /dev/null +++ b/relax/utils/payload_digest.py @@ -0,0 +1,140 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Leaf-level byte fingerprints for data-plane payloads. + +Byte-exact acceptance for the TransferQueue data plane needs one canonical +fingerprint definition shared by tests, benchmarks, and the multimodal fixture +generator. ``torch.equal`` is a *value* comparison (``NaN != NaN``, and +``-0.0 == 0.0``), so it cannot prove byte identity; these helpers hash the raw +storage bytes instead. + +The unit of comparison is a *leaf*: payloads are traversed recursively +(dict / list / tuple), tensors are normalized to contiguous row-major CPU +storage, and every leaf yields ``(dtype, shape, sha256)``. Comparing two +payloads reduces to comparing their digest maps, which also produces precise +"which leaf, which axis" mismatch reports for acceptance logs. +""" + +from __future__ import annotations + +import hashlib +from typing import Any + +import numpy as np +import torch + + +LeafDigest = tuple[str, str, str] # (dtype, shape, sha256 of raw bytes) + + +def _tensor_digest(value: torch.Tensor) -> LeafDigest: + """Digest a tensor's storage bytes, normalized to contiguous CPU layout. + + ``view(torch.uint8)`` reinterprets storage without conversion, so bfloat16 + and other numpy-unsupported dtypes hash losslessly. + """ + flat = value.detach().cpu().contiguous().reshape(-1) + raw = flat.view(torch.uint8).numpy().tobytes() if flat.numel() else b"" + shape = "x".join(str(dim) for dim in value.shape) + return (str(value.dtype), shape, hashlib.sha256(raw).hexdigest()) + + +def _ndarray_digest(value: np.ndarray) -> LeafDigest: + contiguous = np.ascontiguousarray(value) + shape = "x".join(str(dim) for dim in value.shape) + return (f"np.{contiguous.dtype}", shape, hashlib.sha256(contiguous.tobytes()).hexdigest()) + + +def _scalar_digest(value: Any) -> LeafDigest: + if isinstance(value, bytes): + raw = value + elif isinstance(value, str): + raw = value.encode("utf-8") + else: # bool / int / float / None — repr is canonical for these types. + raw = repr(value).encode("utf-8") + return (f"py.{type(value).__name__}", "", hashlib.sha256(raw).hexdigest()) + + +def _unwrap_non_tensor(value: Any) -> Any: + """Unwrap tensordict ``NonTensorData`` / ``NonTensorStack`` wrappers. + + TransferQueue returns non-tensor fields re-wrapped by tensordict; the + fingerprint must see the underlying Python object so that put-side and get- + side digests are comparable. + """ + if type(value).__name__ in ("NonTensorData", "NonTensorStack"): + return value.tolist() if type(value).__name__ == "NonTensorStack" else value.data + return value + + +def leaf_digests(payload: Any, prefix: str = "payload") -> dict[str, LeafDigest]: + """Map every leaf of *payload* to ``(dtype, shape, sha256)``. + + Supported nodes: dict (sorted keys), list/tuple, ``torch.Tensor`` + (including jagged ``NestedTensor``, digested per row so put-side lists and + get-side NestedTensors compare equal), ``np.ndarray``, and scalar leaves + (str/bytes/bool/int/float/None). Unknown node types raise ``TypeError`` so + no leaf is ever silently skipped. + """ + payload = _unwrap_non_tensor(payload) + digests: dict[str, LeafDigest] = {} + if isinstance(payload, torch.Tensor): + if payload.is_nested: + for row_index, row in enumerate(payload.unbind()): + digests[f"{prefix}[{row_index}]"] = _tensor_digest(row) + else: + digests[prefix] = _tensor_digest(payload) + elif isinstance(payload, np.ndarray): + digests[prefix] = _ndarray_digest(payload) + elif isinstance(payload, dict): + for key in sorted(payload.keys()): + digests.update(leaf_digests(payload[key], f"{prefix}.{key}")) + elif isinstance(payload, (list, tuple)): + for index, item in enumerate(payload): + digests.update(leaf_digests(item, f"{prefix}[{index}]")) + elif isinstance(payload, (str, bytes, bool, int, float)) or payload is None: + digests[prefix] = _scalar_digest(payload) + else: + raise TypeError(f"Unsupported payload leaf at {prefix}: {type(payload).__name__}") + return digests + + +def diff_digests(expected: dict[str, LeafDigest], actual: dict[str, LeafDigest]) -> list[str]: + """Return human-readable mismatch lines; empty list means byte-exact.""" + problems: list[str] = [] + for path in sorted(expected.keys() | actual.keys()): + want, have = expected.get(path), actual.get(path) + if want is None: + problems.append(f"{path}: unexpected extra leaf {have}") + elif have is None: + problems.append(f"{path}: missing (expected {want})") + elif want != have: + for axis, want_part, have_part in zip(("dtype", "shape", "sha256"), want, have, strict=True): + if want_part != have_part: + problems.append(f"{path}: {axis} mismatch (expected {want_part}, got {have_part})") + return problems + + +def total_leaf_bytes(payload: Any) -> int: + """Total payload bytes across all tensor/ndarray leaves (for effective- + bandwidth accounting). + + Scalar leaves count their encoded byte length; container overhead is + excluded because acceptance bandwidth is defined over payload bytes. + """ + payload = _unwrap_non_tensor(payload) + if isinstance(payload, torch.Tensor): + if payload.is_nested: + return sum(row.numel() * row.element_size() for row in payload.unbind()) + return payload.numel() * payload.element_size() + if isinstance(payload, np.ndarray): + return payload.nbytes + if isinstance(payload, dict): + return sum(total_leaf_bytes(value) for value in payload.values()) + if isinstance(payload, (list, tuple)): + return sum(total_leaf_bytes(item) for item in payload) + if isinstance(payload, bytes): + return len(payload) + if isinstance(payload, str): + return len(payload.encode("utf-8")) + return 0 diff --git a/scripts/benchmarks/make_multimodal_fixture.py b/scripts/benchmarks/make_multimodal_fixture.py new file mode 100644 index 000000000..37e901789 --- /dev/null +++ b/scripts/benchmarks/make_multimodal_fixture.py @@ -0,0 +1,265 @@ +#!/usr/bin/env python +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Generate a REAL multimodal payload fixture for byte-exact acceptance. + +The maintainer acceptance for the TransferQueue RDMA data plane requires +byte-exact consistency on *real* multimodal payloads, not synthetic +production-shaped tensors. "Real" means: actual dataset images pushed through +the exact production preprocessing chain, producing the exact structure the +Relax data plane ships — ``multimodal_train_inputs`` as ``list[dict]`` +(a tensordict ``NonTensorStack``, i.e. MooncakeStore's non-tensor slow path). + +This script replicates the production chain line-for-line by calling the same +functions rollout uses: + + parquet row {prompt, image} + -> relax.utils.data.data_utils.build_messages (placeholder -> content) + -> tokenizer.apply_chat_template (prompt string) + -> relax.utils.data.processing_utils.process_vision_info (bytes -> PIL, resize) + -> adapt_processor_kwargs -> HF processor -> strip input_ids/attention_mask + -> numpy->torch -> remap_mm_train_inputs (== sglang_rollout._run_processor) + -> GRPO group expansion: n_samples_per_prompt byte-identical copies per + prompt (production runs the processor once per sample; determinism is + verified below so per-sample clones are byte-equivalent) + +The fixture bundles the resulting ``train_data`` lists plus a leaf-level +SHA-256 manifest (see :mod:`relax.utils.payload_digest`). Consumers: + + * tests/utils/test_tq_dataplane_behavior.py (full tq.init/put/get link) + * tests/utils/test_tq_failure_paths.py (direct MooncakeStore client) + * scripts/benchmarks/tq_cross_node_bench.py (--payload-profiles real-multimodal) + +The .pt file is machine-local (hundreds of MB; NOT committed). Committable +provenance goes to ``--manifest-json``: generation args, dataset rows, patch +counts, and every leaf hash, so any regenerated fixture can be audited. + +Example (defaults target the acceptance dataset used in issue #217): + + PYTHONPATH=. python scripts/benchmarks/make_multimodal_fixture.py \\ + --dataset /path/to/.parquet \\ + --model /path/to/ \\ + --num-prompts 6 --n-samples-per-prompt 2 \\ + --output tests/fixtures/tq_multimodal_fixture.pt \\ + --manifest-json tests/fixtures/tq_multimodal_fixture.manifest.json +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import time +from typing import Any + +import numpy as np +import torch + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--dataset", required=True, help="Parquet file with 'prompt' and 'image' columns.") + parser.add_argument("--model", required=True, help="HF model dir providing the processor + tokenizer.") + parser.add_argument("--output", default="tests/fixtures/tq_multimodal_fixture.pt", help="Fixture .pt path.") + parser.add_argument("--manifest-json", default="", help="Optional committable provenance JSON path.") + parser.add_argument("--num-prompts", type=int, default=6, help="Distinct prompts (images) to process.") + parser.add_argument( + "--n-samples-per-prompt", + type=int, + default=2, + help="GRPO group size: byte-identical copies per prompt (production F4 semantics).", + ) + parser.add_argument("--row-offset", type=int, default=0, help="First dataset row to scan.") + parser.add_argument( + "--skip-determinism-check", + action="store_true", + help="Skip the double-run processor determinism verification (not recommended).", + ) + return parser.parse_args() + + +def _load_rows(dataset_path: str, num_prompts: int, row_offset: int) -> list[dict[str, Any]]: + """Sequentially collect rows that carry at least one image.""" + import pyarrow.parquet as pq + + table = pq.read_table(dataset_path, columns=["prompt", "image"]) + rows: list[dict[str, Any]] = [] + for index in range(row_offset, table.num_rows): + images = table["image"][index].as_py() + if not images: + continue + rows.append( + { + "row_index": index, + "prompt": table["prompt"][index].as_py(), + "image": images, + } + ) + if len(rows) >= num_prompts: + return rows + raise RuntimeError( + f"Only found {len(rows)} usable rows (needed {num_prompts}) in {dataset_path} from offset {row_offset}." + ) + + +def _run_production_processor(row: dict[str, Any], tokenizer: Any, processor: Any) -> tuple[list[int], dict[str, Any]]: + """Replicate ``sglang_rollout._run_image_processor``'s synchronous body. + + Every call below is the same production function the rollout worker uses; + nothing is re-implemented here. + """ + from relax.utils.data.data_utils import build_messages + from relax.utils.data.processing_utils import ( + adapt_processor_kwargs, + process_vision_info, + remap_mm_train_inputs, + ) + + messages = build_messages( + {"prompt": row["prompt"], "image": row["image"]}, + prompt_key="prompt", + system_prompt=None, + as_conversation=True, + multimodal_keys={"image": "image"}, + ) + prompt_str = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + multimodal_inputs = process_vision_info(messages, processor, use_audio_in_video=False) + + adapted = adapt_processor_kwargs( + processor, multimodal_inputs, {"use_audio_in_video": False, "return_mm_token_type_ids": False} + ) + processor_output = processor(text=prompt_str, **adapted) + prompt_ids = processor_output["input_ids"][0] + if isinstance(prompt_ids, torch.Tensor): + prompt_ids = prompt_ids.tolist() + train_inputs = { + key: (torch.from_numpy(value) if isinstance(value, np.ndarray) else value) + for key, value in processor_output.items() + if key not in ["input_ids", "attention_mask"] + } or None + train_inputs = remap_mm_train_inputs(processor, train_inputs) + if not isinstance(train_inputs, dict) or not train_inputs: + raise RuntimeError(f"Processor produced no multimodal train inputs for row {row['row_index']}.") + return list(prompt_ids), train_inputs + + +def _clone_train_inputs(train_inputs: dict[str, Any]) -> dict[str, Any]: + """Independent storage per sample, byte-identical content (F4 + semantics).""" + cloned: dict[str, Any] = {} + for key, value in train_inputs.items(): + cloned[key] = value.clone() if isinstance(value, torch.Tensor) else value + return cloned + + +def main() -> int: + from relax.utils.payload_digest import diff_digests, leaf_digests, total_leaf_bytes + + args = parse_args() + + from transformers import AutoProcessor, AutoTokenizer + + print(f"[fixture] loading processor/tokenizer from {args.model}") + processor = AutoProcessor.from_pretrained(args.model, trust_remote_code=True) + tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True) + + rows = _load_rows(args.dataset, args.num_prompts, args.row_offset) + print(f"[fixture] processing {len(rows)} dataset rows (offset {args.row_offset})") + + determinism: dict[str, Any] = {"checked": False, "byte_identical": None} + if not args.skip_determinism_check: + first_ids_a, first_mm_a = _run_production_processor(rows[0], tokenizer, processor) + first_ids_b, first_mm_b = _run_production_processor(rows[0], tokenizer, processor) + mismatches = diff_digests(leaf_digests(first_mm_a), leaf_digests(first_mm_b)) + determinism = {"checked": True, "byte_identical": not mismatches and first_ids_a == first_ids_b} + if mismatches: + # The GRPO group expansion below clones one processor run, which is + # only equivalent to production's N independent runs if the + # processor is byte-deterministic (RFC F4 caveat). Surface loudly. + print("[fixture] WARNING: processor is NOT byte-deterministic across runs:") + for line in mismatches[:8]: + print(f" {line}") + else: + print("[fixture] processor determinism verified (two runs byte-identical)") + + tokens: list[list[int]] = [] + multimodal_train_inputs: list[dict[str, Any]] = [] + group_index: list[int] = [] + prompt_meta: list[dict[str, Any]] = [] + for prompt_position, row in enumerate(rows): + prompt_ids, train_inputs = _run_production_processor(row, tokenizer, processor) + grid = train_inputs.get("image_grid_thw") + prompt_meta.append( + { + "row_index": row["row_index"], + "num_images": len(row["image"]), + "prompt_len": len(prompt_ids), + "mm_keys": sorted(train_inputs.keys()), + "mm_shapes": { + key: list(value.shape) for key, value in train_inputs.items() if isinstance(value, torch.Tensor) + }, + "image_grid_thw": grid.tolist() if isinstance(grid, torch.Tensor) else None, + } + ) + for _ in range(args.n_samples_per_prompt): + tokens.append(list(prompt_ids)) + multimodal_train_inputs.append(_clone_train_inputs(train_inputs)) + group_index.append(prompt_position) + print( + f"[fixture] row {row['row_index']}: prompt_len={len(prompt_ids)} " + f"mm={prompt_meta[-1]['mm_shapes']} x{args.n_samples_per_prompt} samples" + ) + + num_samples = len(tokens) + train_data = {"tokens": tokens, "multimodal_train_inputs": multimodal_train_inputs} + + manifest: dict[str, Any] = {} + for sample_index in range(num_samples): + manifest.update( + leaf_digests(multimodal_train_inputs[sample_index], f"sample[{sample_index}].multimodal_train_inputs") + ) + manifest.update( + leaf_digests(torch.tensor(tokens[sample_index], dtype=torch.int64), f"sample[{sample_index}].tokens") + ) + + payload_bytes = sum(total_leaf_bytes(sample) for sample in multimodal_train_inputs) + meta = { + "schema": 1, + "generated_at_unix": time.time(), + "model": os.path.abspath(args.model), + "processor_class": type(processor).__name__, + "image_processor_class": type(processor.image_processor).__name__, + "patch_size": getattr(processor.image_processor, "patch_size", None), + "dataset": os.path.abspath(args.dataset), + "row_offset": args.row_offset, + "num_prompts": len(rows), + "n_samples_per_prompt": args.n_samples_per_prompt, + "num_samples": num_samples, + "multimodal_payload_bytes": payload_bytes, + "processor_determinism": determinism, + "prompts": prompt_meta, + "versions": { + "torch": torch.__version__, + "transformers": __import__("transformers").__version__, + "tensordict": __import__("tensordict").__version__, + }, + } + + os.makedirs(os.path.dirname(os.path.abspath(args.output)), exist_ok=True) + torch.save({"meta": meta, "train_data": train_data, "manifest": manifest}, args.output) + print( + f"[fixture] wrote {args.output}: {num_samples} samples, " + f"{payload_bytes / 1024**2:.1f} MiB multimodal payload, {len(manifest)} manifest leaves" + ) + + if args.manifest_json: + os.makedirs(os.path.dirname(os.path.abspath(args.manifest_json)), exist_ok=True) + with open(args.manifest_json, "w") as handle: + json.dump({"meta": meta, "manifest": manifest}, handle, indent=1, sort_keys=True) + print(f"[fixture] wrote provenance manifest {args.manifest_json}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/benchmarks/tq_cross_node_bench.py b/scripts/benchmarks/tq_cross_node_bench.py index 67193503a..0c59e23c4 100644 --- a/scripts/benchmarks/tq_cross_node_bench.py +++ b/scripts/benchmarks/tq_cross_node_bench.py @@ -17,9 +17,13 @@ * C1 ``tcp`` -- MooncakeStore / mooncake / TCP * C2 ``rdma`` -- MooncakeStore / mooncake / RDMA -Both synthetic tensors and production-shaped multimodal payloads are tested; -every fetched field is compared byte-for-byte via a SHA-256 digest before a -throughput result is accepted. Every payload tier's transport is proven by +Three payload profiles: synthetic tensors, production-shaped multimodal +tensors, and ``real-multimodal`` — a replay of REAL images through the +production Qwen-VL processor (fixture from make_multimodal_fixture.py) with +``multimodal_train_inputs`` as a NonTensorStack column, i.e. the storage +backends' non-tensor slow path that real VL training exercises. Every fetched +field is compared byte-for-byte via a SHA-256 digest before a throughput +result is accepted. Every payload tier's transport is proven by reading the IB ``port_rcv_data`` and bond0 ``rx_bytes`` counters around each get: RDMA must show IB moving and bond0 flat; TCP the reverse. There is no "thought it was RDMA but was TCP" ambiguity. @@ -78,8 +82,16 @@ def parse_args() -> argparse.Namespace: "--payload-profiles", nargs="+", default=["synthetic", "multimodal"], - choices=["synthetic", "multimodal"], - help="Payload layouts to benchmark. Multimodal uses production field names, shapes, and dtypes.", + choices=["synthetic", "multimodal", "real-multimodal"], + help="Payload layouts to benchmark. 'multimodal' is production-shaped dense tensors; " + "'real-multimodal' replays a fixture from make_multimodal_fixture.py (real images through the " + "production processor, multimodal_train_inputs as list[dict] == storage non-tensor slow path).", + ) + p.add_argument( + "--fixture-path", + default="", + help="Fixture .pt for real-multimodal (default: $RELAX_MM_FIXTURE, else tests/fixtures/" + "tq_multimodal_fixture.pt). Generate with scripts/benchmarks/make_multimodal_fixture.py.", ) p.add_argument( "--repeats", @@ -168,15 +180,86 @@ def make_multimodal_payload(num_samples: int, total_mib: int): ) -def make_profile_payload(profile: str, num_samples: int, num_fields: int, total_mib: int): +_FIXTURE_CACHE: dict[str, Any] = {} + + +def resolve_fixture_path(cli_value: str) -> str: + """--fixture-path > $RELAX_MM_FIXTURE > repo default.""" + import os + + if cli_value: + return cli_value + if os.environ.get("RELAX_MM_FIXTURE"): + return os.environ["RELAX_MM_FIXTURE"] + repo_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + return os.path.join(repo_root, "tests", "fixtures", "tq_multimodal_fixture.pt") + + +def make_real_multimodal_payload(total_mib: int, fixture_path: str): + """Replay REAL processor outputs at the requested payload tier. + + Fixture samples (real dataset images through the production Qwen-VL chain) + are tiled cyclically until the multimodal payload reaches ``total_mib``. + The payload is assembled by the production ``dict_to_tensordict``, so + ``multimodal_train_inputs`` ships as a ``NonTensorStack`` column — + MooncakeStore's msgpack non-tensor slow path and SimpleStorage's pickled- + object path, exactly like a training job. Rows are shared references + (transport serializes each row anyway), so tiling does not multiply driver + RAM. + """ + import torch + + from relax.utils.payload_digest import total_leaf_bytes + from relax.utils.utils import dict_to_tensordict + + if fixture_path not in _FIXTURE_CACHE: + _FIXTURE_CACHE[fixture_path] = torch.load(fixture_path, map_location="cpu", weights_only=False) + bundle = _FIXTURE_CACHE[fixture_path] + base_tokens = bundle["train_data"]["tokens"] + base_mm = bundle["train_data"]["multimodal_train_inputs"] + + target_bytes = total_mib * 1024**2 + tokens: list[list[int]] = [] + multimodal: list[dict] = [] + accumulated = 0 + index = 0 + while accumulated < target_bytes or len(tokens) < 1: + source = index % len(base_mm) + tokens.append(base_tokens[source]) + multimodal.append(base_mm[source]) + accumulated += total_leaf_bytes(base_mm[source]) + index += 1 + num_samples = len(tokens) + train_data = { + "sample_id": list(range(num_samples)), + "tokens": tokens, + "multimodal_train_inputs": multimodal, + } + return dict_to_tensordict(train_data, batch_size=num_samples) + + +def make_profile_payload(profile: str, num_samples: int, num_fields: int, total_mib: int, fixture_path: str = ""): + if profile == "real-multimodal": + return make_real_multimodal_payload(total_mib, fixture_path) if profile == "multimodal": return make_multimodal_payload(num_samples, total_mib) return make_payload(num_samples, num_fields, total_mib) +def profile_field_counts(profile: str, num_fields: list[int]) -> list[int]: + """Field-count sweep per profile: synthetic sweeps --num-fields; the + multimodal profiles have fixed production schemas.""" + if profile == "synthetic": + return num_fields + return [7] if profile == "multimodal" else [3] + + def payload_bytes(payload) -> int: - """Total bytes across all tensor fields.""" - return sum(payload[k].nelement() * payload[k].element_size() for k in payload.keys()) + """Total payload bytes across all fields (tensor, NestedTensor, or + NonTensorStack columns).""" + from relax.utils.payload_digest import total_leaf_bytes + + return sum(total_leaf_bytes(payload[k]) for k in payload.keys()) def field_byte_digests(payload, fields: list[str]) -> dict[str, tuple[str, int, str]]: @@ -198,6 +281,55 @@ def field_byte_digests(payload, fields: list[str]) -> dict[str, tuple[str, int, return out +def _column_rows(column) -> list: + """Rows of a TensorDict column: jagged NestedTensor, NonTensorStack, dense + tensor, or plain list.""" + import torch + + if isinstance(column, torch.Tensor) and column.is_nested: + return list(column.unbind()) + if type(column).__name__ == "NonTensorStack": + return column.tolist() + return [column[i] for i in range(len(column))] + + +def field_multiset_digests(payload, fields: list[str]) -> dict[str, tuple[int, str]]: + """Order-insensitive per-field digest: hash of sorted per-row digests. + + Used for the real-multimodal profile, where columns include a + ``NonTensorStack`` of per-sample dicts: the sampler is free to reorder + rows, so the byte-exact contract is "the returned multiset of rows is + byte-identical to the put multiset". + + Dict rows (multimodal_train_inputs) are digested leaf-by-leaf with full + dtype+shape+bytes via relax.utils.payload_digest. Tensor rows are + digested as (dtype, numel, bytes) — same contract as ``_flat_values`` in + the dataplane tests — because backends legitimately differ in scalar-row + representation (SimpleStorage returns dense columns whose rows index as + 0-D; MooncakeStore reconstructs rows as shape ``[1]``; bytes and dtype + are identical). + """ + import torch + + from relax.utils.payload_digest import leaf_digests + + def _row_digest(row) -> str: + if isinstance(row, torch.Tensor) and not row.is_nested: + flat = row.detach().cpu().contiguous().reshape(-1) + raw = flat.view(torch.uint8).numpy().tobytes() if flat.numel() else b"" + token = f"{flat.dtype}|{flat.numel()}|{hashlib.sha256(raw).hexdigest()}" + else: + token = repr(sorted(leaf_digests(row).items())) + return hashlib.sha256(token.encode()).hexdigest() + + out: dict[str, tuple[int, str]] = {} + for field in fields: + rows = _column_rows(payload[field]) + row_hashes = sorted(_row_digest(row) for row in rows) + out[field] = (len(rows), hashlib.sha256("".join(row_hashes).encode()).hexdigest()) + return out + + def wait_actor_gone(name: str = "TransferQueueController", timeout: float = 30.0) -> None: """Wait for a named TQ actor to leave the GCS (F10-safe re-init).""" deadline = time.time() + timeout @@ -340,7 +472,7 @@ def shutdown(self) -> None: except Exception: # pragma: no cover - best effort pass - def fetch(self, fields, batch_size: int, partition: str, expected_digests): + def fetch(self, fields, batch_size: int, partition: str, expected_digests, order_insensitive: bool = False): """One cross-node get; return (ms, ib_mb, tcp_mb, ib_tail_mb, tcp_tail_mb). @@ -349,6 +481,10 @@ def fetch(self, fields, batch_size: int, partition: str, expected_digests): finished (async), the tail keeps flowing and ``ib_tail`` > 0 -- a definitive async-completion detector that does not need a costly full- data touch. + + ``order_insensitive`` selects the row-multiset digest (real-multimodal + profile: NonTensorStack columns, sampler may reorder rows); the + default column digest requires identical row order. """ before = read_counters() t0 = time.perf_counter() @@ -366,7 +502,10 @@ def fetch(self, fields, batch_size: int, partition: str, expected_digests): settled = read_counters() # Digesting occurs outside the timed interval. A mismatch is fatal: # throughput from a corrupt or truncated transfer is never reported. - actual_digests = field_byte_digests(got, list(fields)) + if order_insensitive: + actual_digests = field_multiset_digests(got, list(fields)) + else: + actual_digests = field_byte_digests(got, list(fields)) if actual_digests != expected_digests: mismatch = [field for field in fields if actual_digests.get(field) != expected_digests.get(field)] raise AssertionError(f"byte-exact mismatch after TQ get: fields={mismatch}") @@ -393,6 +532,15 @@ def main() -> None: from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy args = parse_args() + fixture_path = resolve_fixture_path(args.fixture_path) + if "real-multimodal" in args.payload_profiles: + import os + + if not os.path.isfile(fixture_path): + raise SystemExit( + f"real-multimodal profile needs a fixture at {fixture_path}; generate one with " + "scripts/benchmarks/make_multimodal_fixture.py (see its docstring), or pass --fixture-path." + ) ray.init(ignore_reinit_error=True, address="auto", logging_level="ERROR") nodeb = next(n for n in ray.nodes() if n["NodeManagerAddress"] == args.nodeb_ip and n.get("Alive")) @@ -438,14 +586,18 @@ def main() -> None: results[protocol].setdefault(profile, {}) for total_mib in args.payload_mib: results[protocol][profile].setdefault(total_mib, {}) - # Multimodal has a fixed production schema; synthetic uses the - # requested field-count sweep. - field_counts = args.num_fields if profile == "synthetic" else [7] + # Multimodal profiles have fixed production schemas; synthetic + # uses the requested field-count sweep. + field_counts = profile_field_counts(profile, args.num_fields) + order_insensitive = profile == "real-multimodal" for nf in field_counts: - payload = make_profile_payload(profile, args.num_samples, nf, total_mib) + payload = make_profile_payload(profile, args.num_samples, nf, total_mib, fixture_path) nf = len(list(payload.keys())) fields = sorted(payload.keys()) - expected_digests = field_byte_digests(payload, fields) + if order_insensitive: + expected_digests = field_multiset_digests(payload, fields) + else: + expected_digests = field_byte_digests(payload, fields) nbytes = payload_bytes(payload) put_times: list[float] = [] get_times: list[float] = [] @@ -454,13 +606,14 @@ def main() -> None: ib_tails: list[float] = [] # async-completion detector (~0 == sync get) tcp_tails: list[float] = [] # repeat 0 is a warm-up: first transfer pays RDMA endpoint handshake. + batch_rows = payload.batch_size[0] if payload.batch_size else args.num_samples for r in range(args.repeats + 1): part = f"xfer_{protocol}_{profile}_{total_mib}_{nf}_{r}" t0 = time.perf_counter() producer.put(payload, partition_id=part) put_ms = (time.perf_counter() - t0) * 1000 get_ms, ib, tcp, ib_tail, tcp_tail, byte_exact = ray.get( - consumer.fetch.remote(fields, args.num_samples, part, expected_digests) + consumer.fetch.remote(fields, batch_rows, part, expected_digests, order_insensitive) ) producer.clear_partition(part) if r == 0: @@ -534,7 +687,7 @@ def main() -> None: } results[protocol][profile][total_mib][nf] = rec print( - f" {profile:<10} {str(total_mib) + 'M':<9} f={nf} " + f" {profile:<16} {str(total_mib) + 'M':<9} f={nf} " f"put_mean={rec['put_mean_gbs']:6.2f} " f"GB/s get_mean={rec['get_mean_gbs']:6.2f} (med {rec['get_med_gbs']:.2f}, " f"std {rec['get_std_gbs']:.2f}) GB/s byte_exact=PASS " @@ -551,18 +704,18 @@ def main() -> None: # ---- Summary (mean-based, all requested protocols) ---- print("\n===== SUMMARY: TQ-layer cross-node, same topology (get, MEAN of N runs) =====", flush=True) header = ( - f"{'Profile':<11}{'Payload':<9}{'f':<4}{'C0 mean':>9}{'C1 mean':>9}{'C2 mean':>9}" + f"{'Profile':<16}{'Payload':<9}{'f':<4}{'C0 mean':>9}{'C1 mean':>9}{'C2 mean':>9}" f"{'C1/C0':>8}{'C2/C1':>8}{'C2 std':>8}{'>=20%':>7}{'wire C0/C1/C2':>18}" ) print(header, flush=True) for profile in args.payload_profiles: - field_counts = args.num_fields if profile == "synthetic" else [7] + field_counts = profile_field_counts(profile, args.num_fields) for total_mib in args.payload_mib: for nf in field_counts: c0 = results.get("simple", {}).get(profile, {}).get(total_mib, {}).get(nf) c1 = results.get("tcp", {}).get(profile, {}).get(total_mib, {}).get(nf) c2 = results.get("rdma", {}).get(profile, {}).get(total_mib, {}).get(nf) - prefix = f"{profile:<11}{str(total_mib) + 'M':<9}{nf:<4}" + prefix = f"{profile:<16}{str(total_mib) + 'M':<9}{nf:<4}" if not (c0 and c1 and c2): # A protocol was skipped (--protocols subset) -- print what we have. parts = [] diff --git a/tests/utils/mm_payload_fixtures.py b/tests/utils/mm_payload_fixtures.py new file mode 100644 index 000000000..690ab2fb0 --- /dev/null +++ b/tests/utils/mm_payload_fixtures.py @@ -0,0 +1,107 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Shared multimodal payload sources for TransferQueue byte-exact tests. + +Two tiers, selected automatically: + +* **real** — a fixture produced by ``scripts/benchmarks/make_multimodal_fixture.py`` + (real dataset images through the production Qwen-VL processing chain). + Found via ``$RELAX_MM_FIXTURE`` or ``tests/fixtures/tq_multimodal_fixture.pt``. + The fixture's leaf manifest is re-verified on load so a corrupted file can + never silently pass as ground truth. +* **synthetic** — production-*structured* fallback so CI (no dataset, no model + weights) still exercises the exact container shape the data plane ships: + ``multimodal_train_inputs`` as ``list[dict]`` with variable-length fp32 + ``pixel_values [patches, 1536]`` + int64 ``image_grid_thw [1, 3]`` and + ``t*h*w == patches`` (Qwen3-VL patch-16 geometry), which is MooncakeStore's + non-tensor msgpack slow path — NOT the dense-tensor fast path. + +Tests must report which tier ran (the returned ``source`` string) so real- +payload acceptance is auditable in CI logs. +""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any + +import torch + +from relax.utils.payload_digest import diff_digests, leaf_digests + + +_REPO_ROOT = Path(__file__).resolve().parents[2] +_DEFAULT_FIXTURE = _REPO_ROOT / "tests" / "fixtures" / "tq_multimodal_fixture.pt" + +# Grids observed on a real Qwen-VL processor (patch 16, spatial merge 2) over +# the acceptance dataset: mixed aspect ratios, ~0.6-3.7k patches/image. +_SYNTHETIC_GRIDS = ((1, 58, 64), (1, 34, 64), (1, 64, 64), (1, 26, 40), (1, 64, 58), (1, 40, 40)) +_QWEN3_VL_PATCH_DIM = 1536 # channels(3) * temporal(2) * patch(16)^2 + + +def fixture_path() -> Path: + """Fixture location: ``$RELAX_MM_FIXTURE`` wins, else the repo default.""" + override = os.environ.get("RELAX_MM_FIXTURE", "") + return Path(override) if override else _DEFAULT_FIXTURE + + +def load_real_fixture(max_samples: int | None = None) -> dict[str, Any] | None: + """Load + integrity-check the real fixture; ``None`` when unavailable.""" + path = fixture_path() + if not path.is_file(): + return None + bundle = torch.load(path, map_location="cpu", weights_only=False) + train_data = bundle["train_data"] + manifest = bundle["manifest"] + recomputed: dict[str, Any] = {} + for index, sample in enumerate(train_data["multimodal_train_inputs"]): + recomputed.update(leaf_digests(sample, f"sample[{index}].multimodal_train_inputs")) + recomputed.update( + leaf_digests(torch.tensor(train_data["tokens"][index], dtype=torch.int64), f"sample[{index}].tokens") + ) + problems = diff_digests(manifest, recomputed) + if problems: + raise RuntimeError( + f"Multimodal fixture {path} failed its own manifest ({len(problems)} leaf mismatches); " + f"regenerate it with scripts/benchmarks/make_multimodal_fixture.py. First: {problems[0]}" + ) + if max_samples is not None: + train_data = { + "tokens": train_data["tokens"][:max_samples], + "multimodal_train_inputs": train_data["multimodal_train_inputs"][:max_samples], + } + return {"meta": bundle["meta"], "train_data": train_data} + + +def synthetic_mm_train_data(num_samples: int, seed: int = 20260813) -> dict[str, list[Any]]: + """Production-structured synthetic ``train_data`` (see module + docstring).""" + generator = torch.Generator().manual_seed(seed) + tokens: list[list[int]] = [] + multimodal: list[dict[str, torch.Tensor]] = [] + for index in range(num_samples): + t, h, w = _SYNTHETIC_GRIDS[index % len(_SYNTHETIC_GRIDS)] + patches = t * h * w + multimodal.append( + { + "pixel_values": torch.randn(patches, _QWEN3_VL_PATCH_DIM, dtype=torch.float32, generator=generator), + "image_grid_thw": torch.tensor([[t, h, w]], dtype=torch.int64), + } + ) + prompt_len = 512 + 173 * index + tokens.append(torch.randint(0, 151_000, (prompt_len,), generator=generator).tolist()) + return {"tokens": tokens, "multimodal_train_inputs": multimodal} + + +def mm_train_data(num_samples: int) -> tuple[dict[str, list[Any]], str]: + """Real-fixture ``train_data`` when available, else synthetic. + + Returns ``(train_data, source)`` where source is ``"real"`` / + ``"synthetic"``; tests embed it in assertion ids so acceptance logs show + which tier actually ran. + """ + bundle = load_real_fixture(max_samples=num_samples) + if bundle is not None and len(bundle["train_data"]["tokens"]) >= num_samples: + return bundle["train_data"], "real" + return synthetic_mm_train_data(num_samples), "synthetic" diff --git a/tests/utils/test_tq_dataplane_behavior.py b/tests/utils/test_tq_dataplane_behavior.py index fdfa5e7e5..cb2509924 100644 --- a/tests/utils/test_tq_dataplane_behavior.py +++ b/tests/utils/test_tq_dataplane_behavior.py @@ -20,6 +20,11 @@ 5. retry -- re-putting the same partition overwrites; get returns latest. 6. cleanup -- clear_partition empties data; close+reinit yields a fresh, isolated controller (exercises the F10 anti-hang path). + 7. multimodal -- ``multimodal_train_inputs`` as ``list[dict]`` (the + production NonTensorStack container, storage's non-tensor + slow path) survives the full link byte-exactly; runs on the + REAL Qwen-VL fixture when present, production-structured + synthetic otherwise (see tests/utils/mm_payload_fixtures.py). A true cross-node disconnect (consumer node death mid-get) is NOT covered here -- it requires multi-node GPU hardware and is skipped per project rules. @@ -112,6 +117,18 @@ def _flat_values(t): return t.reshape(-1) +def _row_value(column, row_position: int): + """Extract one sample's value from a returned TensorDict column. + + Handles the three container types TQ can hand back: jagged ``NestedTensor`` + (variable-length tensor fields), ``NonTensorStack`` (list[dict] fields), + and plain dense tensors. + """ + if isinstance(column, torch.Tensor) and column.is_nested: + return column.unbind()[row_position] + return column[row_position] + + def _payload(n: int, fields: list[str], cols: int, dtype: str = "float32", seed: int = 0): """Build a TensorDict of ``n`` samples with ``fields`` of shape (n, cols).""" @@ -294,3 +311,57 @@ def test_cleanup_clear_partition_then_reinit_isolated(self, tq_factory): client2 = tq_factory(capacity=16) meta3 = client2.get_meta(data_fields=["a"], batch_size=4, partition_id="cp", mode="fetch", task_name="cp3") assert getattr(meta3, "size", None) == 0 + + +class TestRealMultimodalFullLink: + """Byte-exactness for the production multimodal container, full link. + + Production ships ``multimodal_train_inputs`` as one dict per sample + (``relax/utils/utils.py::dict_to_tensordict`` keeps the raw list -> + tensordict ``NonTensorStack``). That column takes the storage backends' + *non-tensor* path (SimpleStorage: pickled objects; MooncakeStore: msgpack + pack/unpack), which none of the dense-tensor tests above touch. + + Payload source is reported in every assertion: ``real`` (fixture from + ``scripts/benchmarks/make_multimodal_fixture.py`` — actual dataset images + through the production Qwen-VL processor chain) or ``synthetic`` + (production-structured fallback, CI-safe). + """ + + def test_multimodal_list_dict_full_link_byte_exact(self, tq_factory): + """Real assembly (dict_to_tensordict) -> tq put/get -> leaf-level + SHA-256 equality for every sample, aligned by sample_id (the sampler + may reorder rows).""" + from relax.utils.payload_digest import diff_digests, leaf_digests + from relax.utils.utils import dict_to_tensordict + from tests.utils.mm_payload_fixtures import mm_train_data + + num_samples = 4 + train_data, source = mm_train_data(num_samples) + train_data = dict(train_data) + train_data["sample_id"] = list(range(num_samples)) + batch = dict_to_tensordict(train_data, batch_size=num_samples) + assert type(batch.get("multimodal_train_inputs")).__name__ == "NonTensorStack", ( + "precondition: the multimodal column must be the production NonTensorStack container" + ) + + want_mm = [leaf_digests(sample) for sample in train_data["multimodal_train_inputs"]] + want_tokens = [leaf_digests(torch.tensor(row, dtype=torch.int64)) for row in train_data["tokens"]] + + client = tq_factory() + fields = ["sample_id", "tokens", "multimodal_train_inputs"] + got = _round_trip(client, batch, "mmreal", fields, num_samples) + + got_ids = [int(v) for v in _flat_values(got["sample_id"])] + assert sorted(got_ids) == list(range(num_samples)), f"[{source}] sample_id set mismatch: {got_ids}" + for row_position, sample_id in enumerate(got_ids): + mm_problems = diff_digests( + want_mm[sample_id], leaf_digests(_row_value(got["multimodal_train_inputs"], row_position)) + ) + assert not mm_problems, ( + f"[{source}] sample {sample_id} multimodal leaves not byte-exact: {mm_problems[:4]}" + ) + token_problems = diff_digests( + want_tokens[sample_id], leaf_digests(_row_value(got["tokens"], row_position)) + ) + assert not token_problems, f"[{source}] sample {sample_id} tokens not byte-exact: {token_problems[:4]}" diff --git a/tests/utils/test_tq_failure_paths.py b/tests/utils/test_tq_failure_paths.py index e04cd8111..86f196bd8 100644 --- a/tests/utils/test_tq_failure_paths.py +++ b/tests/utils/test_tq_failure_paths.py @@ -9,7 +9,10 @@ * timeout -- controller ``get_config`` timeout and probe-task timeout * disconnect -- store errors surface instead of returning corrupt data * retry -- ``batch_get_into`` / ``batch_upsert_from`` retry-then-raise -* byte-exactness on **MooncakeStore** (SimpleStorage-only before) +* byte-exactness on **MooncakeStore** (SimpleStorage-only before), including + the non-tensor msgpack slow path that production ``multimodal_train_inputs`` + (``list[dict]`` / NonTensorStack) actually takes — real Qwen-VL fixture when + available, production-structured synthetic otherwise * automatic degradation as pytest (was a manual two-node script) * the controller reaper / teardown helpers (now in ``relax.utils.tq_lifecycle``) @@ -116,6 +119,54 @@ def _real_capacity_worker(result_queue, segment_mib: int, payload_mib: int) -> N client.close() +def _mm_slow_path_worker(result_queue, protocol: str) -> None: + """Child-process target: multimodal list[dict] slow-path roundtrip. + + Runs in its own process (one mooncake session per protocol) because + mooncake 0.3.10 is unstable when one process re-creates clients across + protocols (see scripts/benchmarks/tq_cross_node_bench.py docstring); a + child also keeps a real engine hang externally bounded. Reports either + ("ok", ...), ("mismatch", ...) for data corruption, or ("error", ...) for + engine failures — the parent treats anything but "ok" as a hard failure. + """ + try: + from relax.utils.payload_digest import diff_digests, leaf_digests + from tests.utils.mm_payload_fixtures import mm_train_data + + train_data, source = mm_train_data(4) + samples = train_data["multimodal_train_inputs"] + client = TestMooncakeByteExact._client(protocol) + try: + keys = [f"mmslow_{protocol}_{index}@multimodal_train_inputs" for index in range(len(samples))] + put_meta = client.put(keys, samples) + if not all(isinstance(meta, dict) and meta.get("packed_size") for meta in put_meta): + result_queue.put( + ("error", f"[{source}] dict payloads did not take the non-tensor path, put meta: {put_meta}") + ) + return + got = client.get( + keys, + shapes=[[] for _ in keys], + dtypes=[None] * len(keys), + custom_backend_meta=put_meta, + ) + problems: list[str] = [] + for index, (want, have) in enumerate(zip(samples, got, strict=True)): + problems += [ + f"sample {index}: {line}" for line in diff_digests(leaf_digests(want), leaf_digests(have)) + ] + if problems: + result_queue.put(("mismatch", f"[{source}] {problems[:6]}")) + else: + leaves = sum(len(leaf_digests(sample)) for sample in samples) + packed = sum(meta["packed_size"] for meta in put_meta) + result_queue.put(("ok", f"[{source}] {len(samples)} samples, {leaves} leaves, {packed} packed bytes")) + finally: + client.close() + except BaseException as error: # pragma: no cover - transport/env failures + result_queue.put(("error", f"{type(error).__name__}: {error}")) + + # --------------------------------------------------------------------------- # Controller lifecycle: reaper (timeout / half-initialised / healthy) # --------------------------------------------------------------------------- @@ -880,3 +931,31 @@ def test_multi_dtype_shape_roundtrip_is_byte_exact(self, protocol): assert torch.equal(have, want.contiguous()), f"{name} is not byte-exact" finally: client.close() + + @pytest.mark.parametrize("protocol", ["tcp", "rdma"]) + def test_multimodal_list_dict_slow_path_roundtrip_is_byte_exact(self, protocol): + """The container production actually ships: one dict per sample. + + ``multimodal_train_inputs`` reaches MooncakeStore as non-tensor values + (tensordict NonTensorStack rows), which take the msgpack pack -> + registered-buffer memcpy slow path — a completely different code path + from the dense-tensor test above. Uses the REAL Qwen-VL fixture when + available (see tests/utils/mm_payload_fixtures.py), else + production-structured synthetic dicts; the payload source is part of + the reported result for acceptance auditing. + + Runs in a spawn child so each protocol gets a pristine mooncake + session (mooncake 0.3.10 misbehaves when one process cycles clients + across protocols) and a wedged engine cannot hang the suite. + """ + context = multiprocessing.get_context("spawn") + result_queue = context.Queue() + process = context.Process(target=_mm_slow_path_worker, args=(result_queue, protocol)) + process.start() + process.join(timeout=240) + if process.is_alive(): + process.terminate() + process.join(timeout=5) + pytest.fail(f"multimodal slow-path roundtrip ({protocol}) did not finish within 240 seconds") + status, detail = result_queue.get(timeout=2) + assert status == "ok", f"{status}: {detail}" From 0d8c26c3785f8e4a9134785d9c9cb07871ed2503 Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Fri, 14 Aug 2026 10:34:57 +0800 Subject: [PATCH 05/22] fix(data-plane): guard mooncake 0.3.10 TCP memcpy corruption MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 🐛 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/draft/transfer_queue_rdma.md | 30 +++++++++++++++++++++++++++++- relax/utils/tq_correctness.py | 25 +++++++++++++++++++++++-- tests/utils/test_rdma_probe.py | 21 +++++++++++++++++++++ 3 files changed, 73 insertions(+), 3 deletions(-) diff --git a/docs/draft/transfer_queue_rdma.md b/docs/draft/transfer_queue_rdma.md index 4d7360d0e..0a016ab24 100644 --- a/docs/draft/transfer_queue_rdma.md +++ b/docs/draft/transfer_queue_rdma.md @@ -158,6 +158,31 @@ PYTHONPATH=. python -u scripts/benchmarks/tq_cross_node_bench.py \ 若本次开发环境没有两个 RDMA 节点，交付结论必须写成“真机验收未执行”，不能用 mock 通过推导真机已经通过。 +### 双节点实测记录（2026-08，2 节点 × 8 GPU，mooncake 0.3.10.post2） + +同一拓扑下三配置全矩阵：每档 1 轮预热 + 5 轮测量，`--require-wire-proof`，全部 36 个测量点逐字节 SHA-256 PASS。C1 数值为 `MC_STORE_MEMCPY=0` 守卫生效后的诚实基线（守卫前的 TCP 读数混入了损坏路径，见排障表）。get 均值（GB/s）： + +| profile | 档位 | C0 Simple | C1 Mooncake/TCP | C2 Mooncake/RDMA | C2/C1 | +|---|---|---|---|---|---| +| synthetic | 256M | 2.26 | 0.93 | 2.12 | 2.3× | +| synthetic | 1G | 1.25 | 0.93 | 2.25 | 2.4× | +| synthetic | 2G | 1.32 | 0.91 | 2.36 | 2.6× | +| synthetic | 4G | 1.33 | 0.93 | 4.61 | 5.0× | +| multimodal | 256M | 1.73 | 0.78 | 1.61 | 2.1× | +| multimodal | 1G | 1.40 | 0.82 | 2.26 | 2.8× | +| multimodal | 2G | 1.49 | 0.86¹ | 2.51 | 2.9× | +| multimodal | 4G | 1.35 | 0.90¹ | 2.60 | 2.9× | +| real-multimodal | 256M | 1.93 | 1.10 | 2.83 | 2.6× | +| real-multimodal | 1G | 2.57 | 1.10 | 3.11 | 2.8× | +| real-multimodal | 2G | 3.21 | 1.03 | 2.91 | 2.8× | +| real-multimodal | 4G | 1.98 | 1.00 | 8.36 | 8.4× | + +- put 均值区间：C2 4.1–12.9 GB/s，C1 1.5–2.6 GB/s，C0 1.4–2.5 GB/s（写侧收益显著，与“已知限制”第一条一致）。 +- wire-proof：C2 各档对端 IB counter 增量 ≈ 载荷字节且 bond0 ≈ 0；C1/C0 反之，两类证明全部 PASS。 +- C2 在 4G 档吞吐跃升（synthetic 4.61、real-multimodal 8.36）：大批量摊薄了每 key 的 MR 注册开销。 +- 轮间波动：5 轮 std 绝大多数 ≤ 0.2 GB/s；C0 的 real-multimodal 256M/1G 档偶发波动（std 1.01/0.64）。 +- ¹ C1 multimodal 2G/4G 于重启 master 后的会话补测（长时间会话反复异常退出后 master 出现批量 `-800`，见排障表）；同会话其余档位与其他配置未受影响。 + ## 排障表 | 现象 | 可能原因 | 处理 | @@ -167,7 +192,9 @@ PYTHONPATH=. python -u scripts/benchmarks/tq_cross_node_bench.py \ | `setup failed with error code: -1` | master 不可达 | 检查 master 进程与 `MC_MASTER_ADDRESS` | | `Failed to open segment ... Connection refused` | 上一轮客户端异常退出，死 segment 仍在 master 注册 | 等 `client_ttl`（30 s）过期后重试 | | `batch_get_into failed ... error codes [-800, ...]` | 会话内切换协议（0.3.10 上更敏感），或对端不可达 | 每个协议单独进程跑；确认对端存活 | -| bench 的 Mooncake/TCP 档在单机回环下原生 SIGSEGV | mooncake 0.3.10 TCP transport 在单节点回环、64 MiB 级批量下崩溃（同机 direct-client 小批量正常，RDMA/SimpleStorage 同规模正常；响亮崩溃，非静默损坏） | C1 档在真实双节点拓扑上跑；单机开发环境用 simple/rdma 档验证 | +| TCP 档 get 回来的行“尾部全零”但所有批量码报成功（逐字节校验 FAIL） | mooncake 0.3.10 在 TCP-only 环境自动启用 `MC_STORE_MEMCPY` 快拷贝路径（`transfer_task.cpp` "auto-detected: TCP-only environment, memcpy enabled"），该路径跨节点静默截断：双节点取证探针中约半数全新会话的首次传输命中，坏行自 64 KiB 对齐偏移起全零，且不限于首传（金丝雀小传输不能完全预防）；`MC_STORE_MEMCPY=0` 下 12/12 会话干净 | 已修复：`ensure_mooncake_correctness_guards` 默认 `MC_STORE_MEMCPY=0`（RDMA 会话本就自动禁用 memcpy，无影响；运维可显式设 `MC_STORE_MEMCPY=1` 覆盖）。守卫后的 C1 吞吐是诚实基线，此前更高的 TCP 读数混入了未走网络的损坏路径 | +| bench 的 Mooncake/TCP 档在单机回环下原生 SIGSEGV | 与上一行同源：memcpy 快拷贝路径在单机回环、64 MiB 级批量下由静默损坏变为响亮崩溃 | 同上，`MC_STORE_MEMCPY=0` 守卫后已实测消除（回环 multimodal 256M 档干净跑完且逐字节 PASS） | +| 长时间反复起停会话后 `batch_upsert_from ... error codes [-800, ...]`，重试 3 次全败（响亮失败，非静默） | master 长期吸收异常退出的客户端后状态劣化（实测同一 master 连续两个全新会话在同一档位复现，master metrics 仍报 serving） | 重启 mooncake master 后同档位全过；长跑验收前先起新 master | | 多网卡机器跨节点建连失败 | 自动选卡选到了不通的网卡 | 显式 `--tq-rdma-device`；必要时用 `MC_TCP_BIND_ADDRESS` 指定 TCP 侧绑定地址 | | 训练卡在启动、无日志推进 | 半初始化的 controller（TQ 的 `_init_from_existing` 会无限轮询 config） | 本特性已加自动回收；若仍出现，确认 `[dataplane] ... reaping it` 是否打出 | @@ -195,5 +222,6 @@ RDMA 生效时前者按 payload 增长、后者基本不动；反之则说明落 - 写侧（put）收益明显，读侧（get）收益有限：get 每次调用都会注册/注销 MR，且 key 粒度是 `样本 × 字段`，碎片化开销盖过了传输收益。MR 常驻注册与读路径零拷贝成型不在首期范围。 - 跨节点 RDMA vs TCP 的收益在多轮之间波动较大，验收结论应基于多轮分布而非单轮数据。 +- Mooncake/TCP（C1）在 0.3.10 上必须依赖 `MC_STORE_MEMCPY=0` 守卫才能保证字节正确（见排障表）；守卫后 C1 get 吞吐低于 SimpleStorage（C0），TCP 档只作为 RDMA 不可用时的正确性兜底而非性能选项。 - 消费端节点在 get 过程中中途死亡的端到端行为需要双节点真机验证，未做成自动化测试。 - Mooncake 传输层自身的超时参数不由 Relax 控制。 diff --git a/relax/utils/tq_correctness.py b/relax/utils/tq_correctness.py index 9b2b1d32d..cbd0bccf1 100644 --- a/relax/utils/tq_correctness.py +++ b/relax/utils/tq_correctness.py @@ -6,7 +6,10 @@ Mooncake batch result, but not every retry result, logs removal failures without raising, and treats a missing/negative production-status ACK as a successful notification. Those behaviours can turn an explicit storage or -controller failure into silent data loss. +controller failure into silent data loss. The pinned mooncake 0.3.10 +additionally corrupts TCP-protocol transfers through its auto-enabled memcpy +fast path, so that path is disabled here by default (see +:func:`_enforce_safe_memcpy_default`). Keep the compatibility guards here, close to Relax's integration boundary, until the equivalent checks are available in the pinned TransferQueue @@ -17,6 +20,7 @@ from __future__ import annotations import asyncio +import os from functools import wraps from typing import Any from uuid import uuid4 @@ -71,7 +75,6 @@ def batch_remove(self, keys: list[str], *args: Any, **kwargs: Any) -> Any: async def _strict_notify_and_wait(self: Any, request_msg: list) -> None: """Notify the controller and require a positive ACK within the deadline.""" import zmq - from transfer_queue.storage.managers import base as tq_base from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, create_zmq_socket @@ -159,8 +162,26 @@ async def guarded_notify(self: Any, *args: Any, **kwargs: Any) -> None: setattr(manager_cls, _PATCH_MARKER, True) +def _enforce_safe_memcpy_default() -> None: + """Disable mooncake's memcpy fast path unless the operator overrides it. + + mooncake 0.3.10 auto-enables ``MC_STORE_MEMCPY`` in TCP-only environments + (``transfer_task.cpp`` "auto-detected: TCP-only environment, memcpy + enabled") and that path silently truncates cross-node gets: roughly half of + fresh-session first transfers returned rows whose tails were zero bytes + from a 64 KiB-aligned offset onward while every batch code reported success + (two-node forensic probes, 2026-08; 12/12 sessions clean with + ``MC_STORE_MEMCPY=0`` vs ~50% corrupt without). The same code path + SIGSEGVs on single-node loopback. RDMA-capable sessions auto-disable + 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") + + def ensure_mooncake_correctness_guards() -> None: """Install and validate all guards required for safe Mooncake operation.""" + _enforce_safe_memcpy_default() try: from transfer_queue.storage.clients.mooncake_client import MooncakeStoreClient from transfer_queue.storage.managers.base import StorageManager diff --git a/tests/utils/test_rdma_probe.py b/tests/utils/test_rdma_probe.py index 57e36c3e3..8f7f4f016 100644 --- a/tests/utils/test_rdma_probe.py +++ b/tests/utils/test_rdma_probe.py @@ -11,6 +11,7 @@ import argparse import importlib.util +import os from unittest import mock import pytest @@ -421,6 +422,26 @@ def test_mooncake_config_gdr_propagated(self): def test_installed_tq_satisfies_loss_prevention_contract(self): validate_mooncake_runtime_contract() + @pytest.mark.skipif( + not _REAL_TQ_STORAGE, + reason="needs real TransferQueue storage submodules; CPU CI uses a single-file transfer_queue stub", + ) + def test_contract_defaults_mooncake_memcpy_off(self, monkeypatch): + # mooncake 0.3.10 memcpy fast path silently truncates TCP transfers; + # the correctness guards must force it off when the operator is silent. + monkeypatch.delenv("MC_STORE_MEMCPY", raising=False) + validate_mooncake_runtime_contract() + assert os.environ["MC_STORE_MEMCPY"] == "0" + + @pytest.mark.skipif( + not _REAL_TQ_STORAGE, + reason="needs real TransferQueue storage submodules; CPU CI uses a single-file transfer_queue stub", + ) + def test_contract_respects_explicit_memcpy_override(self, monkeypatch): + monkeypatch.setenv("MC_STORE_MEMCPY", "1") + validate_mooncake_runtime_contract() + assert os.environ["MC_STORE_MEMCPY"] == "1" + def test_segment_capacity_text_only_passes(self): args = _make_args(multimodal_keys=None) eff = EffectiveConfig(backend="MooncakeStore", protocol="rdma", device="", gdr=False, fallback_reason="") From 4ab457820490a504b344c1c4157e9a5ff5469a5f Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Fri, 14 Aug 2026 10:45:27 +0800 Subject: [PATCH 06/22] docs: restructure transfer_queue_rdma.md as a usage guide MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 📝 docs - 将「双节点实测记录」日期化实验小节收敛为「参考吞吐区间」：只保留 get/put 量级区间与结论（供容量规划参考），逐档明细、逐轮分布与 原始 CSV 归入交付验收材料，不再在文档内维护 - 排障表三行（TCP 静默截断、回环 SIGSEGV、master 状态劣化）压缩为 「现象/原因/处理」一行式，剥离取证过程叙事（会话统计、探针细节） - `MC_STORE_MEMCPY=0` 守卫的行为说明移入「容量不足与正确性依赖」， 排障表引用之；补充 RDMA 会话不受影响与显式覆盖方式 - 「已知限制」与验收措辞去除开发过程口吻（"此前读数"、"本次开发 环境"），与 docs/draft 下其他使用指南的无时间性语态对齐 --- docs/draft/transfer_queue_rdma.md | 47 +++++++++++-------------------- 1 file changed, 17 insertions(+), 30 deletions(-) diff --git a/docs/draft/transfer_queue_rdma.md b/docs/draft/transfer_queue_rdma.md index 0a016ab24..2f94e4b19 100644 --- a/docs/draft/transfer_queue_rdma.md +++ b/docs/draft/transfer_queue_rdma.md @@ -107,6 +107,8 @@ Mooncake 配置固定 `hard_pin=true`，不会为了腾空间静默驱逐已经 因此，上游曾出现的“返回码未检查导致静默丢数据”不是已知限制，而是 Mooncake 启用的硬门槛：`auto` 在契约不满足时禁用 Mooncake 并回退，`required` 拒绝启动。Docker 镜像固定上述修复 commit，运行时检查用于防止环境被旧包覆盖。 +正确性守卫同时默认设置 `MC_STORE_MEMCPY=0`：mooncake 0.3.10 在 TCP-only 环境会自动启用 memcpy 快拷贝路径，该路径存在静默截断缺陷（现象与处置见排障表）；RDMA 会话本就自动禁用 memcpy，不受影响。运维确认环境安全后可在启动前显式导出 `MC_STORE_MEMCPY=1` 覆盖。 + 真机容量故障注入会故意创建 64 MiB segment 并写入 96 MiB，仅允许在独立、可丢弃的 master 上运行： ```bash @@ -156,32 +158,17 @@ PYTHONPATH=. python -u scripts/benchmarks/tq_cross_node_bench.py \ `real-multimodal` profile 按目标档位循环平铺 fixture 样本，`multimodal_train_inputs` 列以 NonTensorStack 走存储层非张量路径（SimpleStorage pickle / Mooncake msgpack），字节校验用行多重集指纹（采样器可重排行序；张量行按 dtype+字节比较以兼容后端间标量行 `()` 与 `[1]` 的表示差异，dict 叶子仍全形状校验）。 -若本次开发环境没有两个 RDMA 节点，交付结论必须写成“真机验收未执行”，不能用 mock 通过推导真机已经通过。 - -### 双节点实测记录（2026-08，2 节点 × 8 GPU，mooncake 0.3.10.post2） - -同一拓扑下三配置全矩阵：每档 1 轮预热 + 5 轮测量，`--require-wire-proof`，全部 36 个测量点逐字节 SHA-256 PASS。C1 数值为 `MC_STORE_MEMCPY=0` 守卫生效后的诚实基线（守卫前的 TCP 读数混入了损坏路径，见排障表）。get 均值（GB/s）： - -| profile | 档位 | C0 Simple | C1 Mooncake/TCP | C2 Mooncake/RDMA | C2/C1 | -|---|---|---|---|---|---| -| synthetic | 256M | 2.26 | 0.93 | 2.12 | 2.3× | -| synthetic | 1G | 1.25 | 0.93 | 2.25 | 2.4× | -| synthetic | 2G | 1.32 | 0.91 | 2.36 | 2.6× | -| synthetic | 4G | 1.33 | 0.93 | 4.61 | 5.0× | -| multimodal | 256M | 1.73 | 0.78 | 1.61 | 2.1× | -| multimodal | 1G | 1.40 | 0.82 | 2.26 | 2.8× | -| multimodal | 2G | 1.49 | 0.86¹ | 2.51 | 2.9× | -| multimodal | 4G | 1.35 | 0.90¹ | 2.60 | 2.9× | -| real-multimodal | 256M | 1.93 | 1.10 | 2.83 | 2.6× | -| real-multimodal | 1G | 2.57 | 1.10 | 3.11 | 2.8× | -| real-multimodal | 2G | 3.21 | 1.03 | 2.91 | 2.8× | -| real-multimodal | 4G | 1.98 | 1.00 | 8.36 | 8.4× | - -- put 均值区间：C2 4.1–12.9 GB/s，C1 1.5–2.6 GB/s，C0 1.4–2.5 GB/s（写侧收益显著，与“已知限制”第一条一致）。 -- wire-proof：C2 各档对端 IB counter 增量 ≈ 载荷字节且 bond0 ≈ 0；C1/C0 反之，两类证明全部 PASS。 -- C2 在 4G 档吞吐跃升（synthetic 4.61、real-multimodal 8.36）：大批量摊薄了每 key 的 MR 注册开销。 -- 轮间波动：5 轮 std 绝大多数 ≤ 0.2 GB/s；C0 的 real-multimodal 256M/1G 档偶发波动（std 1.01/0.64）。 -- ¹ C1 multimodal 2G/4G 于重启 master 后的会话补测（长时间会话反复异常退出后 master 出现批量 `-800`，见排障表）；同会话其余档位与其他配置未受影响。 +若验收环境没有两个 RDMA 节点，交付结论必须写成“真机验收未执行”，不能用 mock 通过推导真机已经通过。 + +### 参考吞吐区间 + +2 节点 × 8 GPU、mooncake 0.3.10.post2 上按上述命令完成过一次全矩阵验收（36/36 测量点逐字节 PASS、wire-proof 全部成立），量级供容量规划参考： + +- get 均值：C2 RDMA 2.1~4.6 GB/s（real-multimodal 最高 8.4），C1 TCP（守卫后）0.8~1.1 GB/s，C0 SimpleStorage 1.2~3.2 GB/s；C2 相对 C1 增益 2.1×~8.4×，全档位满足 ≥20% 门槛 +- put 均值：C2 4.1~12.9 GB/s，C1/C0 1.4~2.6 GB/s——写侧收益显著大于读侧，与“已知限制”第一条一致 +- 大档位（4G）下 C2 吞吐明显上扬：大批量摊薄了每 key 的 MR 注册开销 + +逐档明细、逐轮分布与原始 CSV 属于交付验收材料，随验收报告存档，不在本文档维护。 ## 排障表 @@ -192,9 +179,9 @@ PYTHONPATH=. python -u scripts/benchmarks/tq_cross_node_bench.py \ | `setup failed with error code: -1` | master 不可达 | 检查 master 进程与 `MC_MASTER_ADDRESS` | | `Failed to open segment ... Connection refused` | 上一轮客户端异常退出，死 segment 仍在 master 注册 | 等 `client_ttl`（30 s）过期后重试 | | `batch_get_into failed ... error codes [-800, ...]` | 会话内切换协议（0.3.10 上更敏感），或对端不可达 | 每个协议单独进程跑；确认对端存活 | -| TCP 档 get 回来的行“尾部全零”但所有批量码报成功（逐字节校验 FAIL） | mooncake 0.3.10 在 TCP-only 环境自动启用 `MC_STORE_MEMCPY` 快拷贝路径（`transfer_task.cpp` "auto-detected: TCP-only environment, memcpy enabled"），该路径跨节点静默截断：双节点取证探针中约半数全新会话的首次传输命中，坏行自 64 KiB 对齐偏移起全零，且不限于首传（金丝雀小传输不能完全预防）；`MC_STORE_MEMCPY=0` 下 12/12 会话干净 | 已修复：`ensure_mooncake_correctness_guards` 默认 `MC_STORE_MEMCPY=0`（RDMA 会话本就自动禁用 memcpy，无影响；运维可显式设 `MC_STORE_MEMCPY=1` 覆盖）。守卫后的 C1 吞吐是诚实基线，此前更高的 TCP 读数混入了未走网络的损坏路径 | -| bench 的 Mooncake/TCP 档在单机回环下原生 SIGSEGV | 与上一行同源：memcpy 快拷贝路径在单机回环、64 MiB 级批量下由静默损坏变为响亮崩溃 | 同上，`MC_STORE_MEMCPY=0` 守卫后已实测消除（回环 multimodal 256M 档干净跑完且逐字节 PASS） | -| 长时间反复起停会话后 `batch_upsert_from ... error codes [-800, ...]`，重试 3 次全败（响亮失败，非静默） | master 长期吸收异常退出的客户端后状态劣化（实测同一 master 连续两个全新会话在同一档位复现，master metrics 仍报 serving） | 重启 mooncake master 后同档位全过；长跑验收前先起新 master | +| Mooncake/TCP 档 get 数据尾部全零，但批量返回码全部成功（逐字节校验 FAIL） | mooncake 0.3.10 memcpy 快拷贝路径缺陷：TCP-only 环境被自动启用后，跨节点 get 会静默截断（坏行自 64 KiB 对齐偏移起全零） | 正确性守卫已默认 `MC_STORE_MEMCPY=0`（见“容量不足与正确性依赖”）；若被显式设为 `1`，改回 `0` | +| Mooncake/TCP 档在单机回环下原生 SIGSEGV | 与上一行同源（memcpy 路径），回环下表现为崩溃而非静默截断 | 同上 | +| 长时间反复起停会话后 `batch_upsert_from ... error codes [-800, ...]`，重试耗尽（响亮失败，非静默） | master 长期吸收异常退出的客户端后状态劣化，metrics 仍报 serving | 重启 mooncake master；长跑验收前先起新 master | | 多网卡机器跨节点建连失败 | 自动选卡选到了不通的网卡 | 显式 `--tq-rdma-device`；必要时用 `MC_TCP_BIND_ADDRESS` 指定 TCP 侧绑定地址 | | 训练卡在启动、无日志推进 | 半初始化的 controller（TQ 的 `_init_from_existing` 会无限轮询 config） | 本特性已加自动回收；若仍出现，确认 `[dataplane] ... reaping it` 是否打出 | @@ -222,6 +209,6 @@ RDMA 生效时前者按 payload 增长、后者基本不动；反之则说明落 - 写侧（put）收益明显，读侧（get）收益有限：get 每次调用都会注册/注销 MR，且 key 粒度是 `样本 × 字段`，碎片化开销盖过了传输收益。MR 常驻注册与读路径零拷贝成型不在首期范围。 - 跨节点 RDMA vs TCP 的收益在多轮之间波动较大，验收结论应基于多轮分布而非单轮数据。 -- Mooncake/TCP（C1）在 0.3.10 上必须依赖 `MC_STORE_MEMCPY=0` 守卫才能保证字节正确（见排障表）；守卫后 C1 get 吞吐低于 SimpleStorage（C0），TCP 档只作为 RDMA 不可用时的正确性兜底而非性能选项。 +- Mooncake/TCP（C1）在 0.3.10 上依赖 `MC_STORE_MEMCPY=0` 守卫保证字节正确，且守卫后 get 吞吐低于 SimpleStorage：TCP 档定位为 RDMA 不可用时的正确性兜底，不是性能选项。 - 消费端节点在 get 过程中中途死亡的端到端行为需要双节点真机验证，未做成自动化测试。 - Mooncake 传输层自身的超时参数不由 Relax 控制。 From 0057df6a7898d42f5ee7d359f39bdbbd67e49a85 Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Fri, 14 Aug 2026 21:01:54 +0800 Subject: [PATCH 07/22] fix(data-plane): harden env and HCA probe guards MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 🐛 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) --- relax/utils/rdma_probe.py | 143 +++++++++++++++++++++++---------- relax/utils/tq_config.py | 23 ++++-- relax/utils/tq_correctness.py | 26 +++--- tests/utils/test_rdma_probe.py | 122 +++++++++++++++++++++++++--- 4 files changed, 246 insertions(+), 68 deletions(-) diff --git a/relax/utils/rdma_probe.py b/relax/utils/rdma_probe.py index 6a1e0c531..ddd274fe5 100644 --- a/relax/utils/rdma_probe.py +++ b/relax/utils/rdma_probe.py @@ -130,45 +130,103 @@ def _check_rdma_devices() -> CheckResult: return CheckResult("rdma_devices", True, ",".join(devs)) -def _check_port_active(device: str = "", port: int = 1) -> CheckResult: - if device: - state_path = f"/sys/class/infiniband/{device}/ports/{port}/state" - try: - with open(state_path) as f: - state = f.read().strip() - ok = "ACTIVE" in state - return CheckResult(f"port_active:{device}/{port}", ok, state) - except FileNotFoundError: - return CheckResult(f"port_active:{device}/{port}", False, "state file missing") - except OSError as e: - return CheckResult(f"port_active:{device}/{port}", False, str(e)) - # No device specified: check the first available one. - base = "/sys/class/infiniband" - if not os.path.isdir(base): - return CheckResult("port_active", False, "no infiniband dir") - for dev in sorted(os.listdir(base)): - return _check_port_active(dev, port) - return CheckResult("port_active", False, "no devices") +def _check_port_active(device: str, port: int = 1) -> CheckResult: + state_path = f"/sys/class/infiniband/{device}/ports/{port}/state" + try: + with open(state_path) as f: + state = f.read().strip() + ok = "ACTIVE" in state + return CheckResult(f"port_active:{device}/{port}", ok, state) + except FileNotFoundError: + return CheckResult(f"port_active:{device}/{port}", False, "state file missing") + except OSError as e: + return CheckResult(f"port_active:{device}/{port}", False, str(e)) + + +def _check_gid_available(device: str, gid_index: int = 3, port: int = 1) -> CheckResult: + gid_path = f"/sys/class/infiniband/{device}/ports/{port}/gids/{gid_index}" + try: + with open(gid_path) as f: + raw = f.read().strip() + ok = raw.replace(":", "") != "0" * 32 and bool(raw) + return CheckResult(f"gid:{device}/{gid_index}", ok, raw[:24] + "..." if len(raw) > 24 else raw) + except FileNotFoundError: + return CheckResult(f"gid:{device}/{gid_index}", False, "gid file missing") + except OSError as e: + return CheckResult(f"gid:{device}/{gid_index}", False, str(e)) + + +def _list_numeric_entries(path: str) -> list[int]: + """Sorted numeric directory entries (port numbers / GID indices).""" + try: + return sorted(int(name) for name in os.listdir(path) if name.isdigit()) + except OSError: + return [] -def _check_gid_available(device: str = "", gid_index: int = 3) -> CheckResult: - if device: - gid_path = f"/sys/class/infiniband/{device}/ports/1/gids/{gid_index}" - try: - with open(gid_path) as f: - raw = f.read().strip() - ok = raw.replace(":", "") != "0" * 32 and bool(raw) - return CheckResult(f"gid:{device}/{gid_index}", ok, raw[:24] + "..." if len(raw) > 24 else raw) - except FileNotFoundError: - return CheckResult(f"gid:{device}/{gid_index}", False, "gid file missing") - except OSError as e: - return CheckResult(f"gid:{device}/{gid_index}", False, str(e)) +def _find_usable_gid(device: str, port: int) -> CheckResult: + """Return the first usable (non-zero) GID on ``device``/``port``. + + Index 3 (conventionally the RoCE v2 / IPv4-mapped entry) is preferred to + preserve the previous behaviour, then every other advertised index is + scanned so a host that populates a different index is not misreported as + GID-less. + """ + indices = _list_numeric_entries(f"/sys/class/infiniband/{device}/ports/{port}/gids") + if not indices: + return CheckResult(f"gid:{device}", False, f"no GID entries on port {port}") + ordered = ([3] if 3 in indices else []) + [i for i in indices if i != 3] + for gid_index in ordered: + check = _check_gid_available(device, gid_index, port=port) + if check.ok: + return check + return CheckResult(f"gid:{device}", False, f"no non-zero GID on port {port}") + + +def _select_usable_rdma_device(device: str = "") -> tuple[str, list[CheckResult]]: + """Pick one HCA whose ACTIVE port and usable GID both pass together. + + Scans every device under ``/sys/class/infiniband`` (or only ``device`` + when explicitly requested), every port of each device, and the GID table + of the first ACTIVE port — instead of assuming the lexicographically + first device with port 1 / GID index 3. A single down HCA on a + multi-HCA host therefore no longer degrades the whole node. + + Returns ``(selected_device, checks)``. ``selected_device`` is ``""`` + when no device qualifies; ``checks`` then summarises one failure per + inspected device for the probe report. + """ base = "/sys/class/infiniband" if not os.path.isdir(base): - return CheckResult(f"gid:{gid_index}", False, "no infiniband dir") - for dev in sorted(os.listdir(base)): - return _check_gid_available(dev, gid_index) - return CheckResult(f"gid:{gid_index}", False, "no devices") + return "", [CheckResult("port_active", False, "no infiniband dir")] + candidates = [device] if device else sorted(os.listdir(base)) + if not candidates: + return "", [CheckResult("port_active", False, "no devices")] + + failures: list[str] = [] + for dev in candidates: + ports = _list_numeric_entries(f"{base}/{dev}/ports") + if not ports: + failures.append(f"{dev}: no ports") + continue + port_check: CheckResult | None = None + active_port: int | None = None + for port in ports: + check = _check_port_active(dev, port) + if check.ok: + port_check, active_port = check, port + break + if port_check is None or active_port is None: + failures.append(f"{dev}: no ACTIVE port") + continue + gid_check = _find_usable_gid(dev, active_port) + if not gid_check.ok: + failures.append(f"{dev}/port{active_port}: {gid_check.detail}") + continue + return dev, [port_check, gid_check] + + detail = "; ".join(failures) + return "", [CheckResult("port_active", False, detail), CheckResult("gid", False, detail)] def _check_memlock() -> CheckResult: @@ -248,7 +306,8 @@ def probe_node(device: str = "", master_address: str = "", *, probe_rdma: bool = Parameters ---------- device - Explicit RDMA device name; empty = auto-detect first available. + Explicit RDMA device name; empty = scan every HCA and select one + whose ACTIVE port and usable GID both pass. probe_rdma When ``False``, validate only Mooncake importability and master reachability, then select TCP. Used by ``--tq-rdma-mode=off`` so an @@ -271,9 +330,10 @@ def probe_node(device: str = "", master_address: str = "", *, probe_rdma: bool = checks.append(_check_master_reachable(master_address)) # Device-dependent checks are irrelevant when RDMA is explicitly off. + selected_device = "" if probe_rdma: - checks.append(_check_port_active(device)) - checks.append(_check_gid_available(device)) + selected_device, device_checks = _select_usable_rdma_device(device) + checks.extend(device_checks) # Determine effective protocol via graded degradation. mooncake_ok = any(c.name == "mooncake_import" and c.ok for c in checks) @@ -295,11 +355,8 @@ def probe_node(device: str = "", master_address: str = "", *, probe_rdma: bool = effective_protocol = "tcp" elif rdma_dev_ok and port_ok and gid_ok and memlock_ok: effective_protocol = "rdma" - if not effective_device: - # Pick first available device for the summary. - base = "/sys/class/infiniband" - if os.path.isdir(base): - effective_device = sorted(os.listdir(base))[0] + # Report the jointly validated device (ACTIVE port + usable GID). + effective_device = selected_device or effective_device else: # mooncake usable but RDMA incomplete -> degrade to TCP (still MooncakeStore). effective_protocol = "tcp" diff --git a/relax/utils/tq_config.py b/relax/utils/tq_config.py index 50193d7fc..dc5cf7d77 100644 --- a/relax/utils/tq_config.py +++ b/relax/utils/tq_config.py @@ -37,8 +37,21 @@ def resolve_mooncake_master_address() -> str: - """Return the externally managed Mooncake master endpoint.""" - return os.environ.get("MC_MASTER_ADDRESS", "localhost:50051") + """Return the externally managed Mooncake master endpoint. + + ``MC_MASTER_ADDRESS`` is required. A loopback default would make every + node of a multi-node job treat its own localhost as the master, so the + reachability probe would degrade ``auto`` runs and abort ``off``/ + ``required`` runs even when a shared master is healthy elsewhere. + """ + address = os.environ.get("MC_MASTER_ADDRESS", "").strip() + if not address: + raise RuntimeError( + "MooncakeStore requires MC_MASTER_ADDRESS= of the externally " + "managed mooncake master on every node; Relax never assumes a loopback " + "endpoint." + ) + return address def validate_mooncake_runtime_contract() -> None: @@ -97,9 +110,9 @@ def build_mooncake_config( effective The job-level :class:`EffectiveConfig` after probing. master_address - External master server address. If ``None``, read from the - ``MC_MASTER_ADDRESS`` env var; if still unset, fall back to localhost - (single-node dev only — production must set the env var). + External master server address. If ``None``, read from the required + ``MC_MASTER_ADDRESS`` env var (see + :func:`resolve_mooncake_master_address`). global_segment_size Override the per-client segment size (default 4 GiB). Benchmarks may pass a larger value (e.g. 8 GiB) to avoid staging-buffer pressure. diff --git a/relax/utils/tq_correctness.py b/relax/utils/tq_correctness.py index cbd0bccf1..9d60d5bec 100644 --- a/relax/utils/tq_correctness.py +++ b/relax/utils/tq_correctness.py @@ -8,8 +8,8 @@ successful notification. Those behaviours can turn an explicit storage or controller failure into silent data loss. The pinned mooncake 0.3.10 additionally corrupts TCP-protocol transfers through its auto-enabled memcpy -fast path, so that path is disabled here by default (see -:func:`_enforce_safe_memcpy_default`). +fast path, so that path is force-disabled here and an explicit enable is +rejected (see :func:`_enforce_safe_memcpy`). Keep the compatibility guards here, close to Relax's integration boundary, until the equivalent checks are available in the pinned TransferQueue @@ -162,8 +162,8 @@ async def guarded_notify(self: Any, *args: Any, **kwargs: Any) -> None: setattr(manager_cls, _PATCH_MARKER, True) -def _enforce_safe_memcpy_default() -> None: - """Disable mooncake's memcpy fast path unless the operator overrides it. +def _enforce_safe_memcpy() -> None: + """Force-disable mooncake's memcpy fast path; reject attempts to enable it. mooncake 0.3.10 auto-enables ``MC_STORE_MEMCPY`` in TCP-only environments (``transfer_task.cpp`` "auto-detected: TCP-only environment, memcpy @@ -172,16 +172,24 @@ def _enforce_safe_memcpy_default() -> None: from a 64 KiB-aligned offset onward while every batch code reported success (two-node forensic probes, 2026-08; 12/12 sessions clean with ``MC_STORE_MEMCPY=0`` vs ~50% corrupt without). The same code path - SIGSEGVs on single-node loopback. RDMA-capable sessions auto-disable - memcpy anyway, so this default is a no-op there; ``setdefault`` keeps an - explicit operator override (e.g. ``MC_STORE_MEMCPY=1``) possible. + SIGSEGVs on single-node loopback. Because the corruption is confirmed on + the pinned mooncake build, this guard fails closed: an explicit + ``MC_STORE_MEMCPY=1`` is rejected at startup instead of honoured. Re-gate + on the mooncake version once the pin moves to a release with the fix. """ - os.environ.setdefault("MC_STORE_MEMCPY", "0") + override = os.environ.get("MC_STORE_MEMCPY", "").strip() + if override not in ("", "0"): + raise RuntimeError( + f"MC_STORE_MEMCPY={override!r} is rejected: the pinned mooncake 0.3.10 " + "memcpy fast path silently truncates TCP transfers and can SIGSEGV. " + "Unset MC_STORE_MEMCPY; Relax forces it to 0 on this version." + ) + os.environ["MC_STORE_MEMCPY"] = "0" def ensure_mooncake_correctness_guards() -> None: """Install and validate all guards required for safe Mooncake operation.""" - _enforce_safe_memcpy_default() + _enforce_safe_memcpy() try: from transfer_queue.storage.clients.mooncake_client import MooncakeStoreClient from transfer_queue.storage.managers.base import StorageManager diff --git a/tests/utils/test_rdma_probe.py b/tests/utils/test_rdma_probe.py index 8f7f4f016..25c53b285 100644 --- a/tests/utils/test_rdma_probe.py +++ b/tests/utils/test_rdma_probe.py @@ -34,6 +34,7 @@ build_mooncake_config, build_simple_storage_config, estimate_payload_bytes, + resolve_mooncake_master_address, validate_mooncake_runtime_contract, validate_segment_capacity, ) @@ -248,9 +249,16 @@ def test_active_rdma_device_gives_rdma(self): def fake_isdir(path): return "infiniband" in path + def fake_listdir(path): + if path.endswith("/ports"): + return ["1"] + if path.endswith("/gids"): + return ["3"] + return ["rdma0"] + with ( mock.patch("os.path.isdir", side_effect=fake_isdir), - mock.patch("os.listdir", return_value=["rdma0"]), + mock.patch("os.listdir", side_effect=fake_listdir), mock.patch("builtins.open", mock.mock_open(read_data="4: ACTIVE")), mock.patch("relax.utils.rdma_probe._check_mooncake_import") as mi, mock.patch("relax.utils.rdma_probe._check_health_check") as hc, @@ -261,6 +269,78 @@ def fake_isdir(path): result = probe_node("") assert result.effective_protocol == "rdma" assert result.ok + assert result.effective_device == "rdma0" + + @staticmethod + def _multi_hca_open(active_device: str): + """Path-aware ``open`` mock: only ``active_device`` has an ACTIVE port + and a non-zero GID at index 3.""" + + def fake_open(path, *args, **kwargs): + path = str(path) + if path.endswith("/state"): + state = "4: ACTIVE" if f"/{active_device}/" in path else "1: DOWN" + return mock.mock_open(read_data=state)(path) + if path.endswith("/gids/3"): + return mock.mock_open(read_data="0000:0000:0000:0000:0000:ffff:0a00:0001")(path) + if "/gids/" in path: + return mock.mock_open(read_data="0000:0000:0000:0000:0000:0000:0000:0000")(path) + raise FileNotFoundError(path) + + return fake_open + + def test_multi_hca_skips_down_first_device(self): + """A down lexicographically-first HCA must not degrade the node when a + later device has an ACTIVE port and a usable GID (review: probe every + usable HCA before degrading RDMA).""" + + def fake_isdir(path): + return "infiniband" in path + + def fake_listdir(path): + if path.endswith("/ports"): + return ["1"] + if path.endswith("/gids"): + return ["0", "3"] + return ["mlx5_0", "mlx5_1"] + + with ( + mock.patch("os.path.isdir", side_effect=fake_isdir), + mock.patch("os.listdir", side_effect=fake_listdir), + mock.patch("builtins.open", side_effect=self._multi_hca_open("mlx5_1")), + mock.patch("relax.utils.rdma_probe._check_mooncake_import") as mi, + mock.patch("relax.utils.rdma_probe.resource.getrlimit", return_value=(-1, -1)), + ): + mi.return_value = CheckResult("mooncake_import", True, "ok") + result = probe_node("") + assert result.effective_protocol == "rdma" + assert result.effective_device == "mlx5_1" + + def test_multi_hca_all_down_degrades_to_tcp(self): + """When no HCA has an ACTIVE port, the node degrades to + Mooncake/TCP.""" + + def fake_isdir(path): + return "infiniband" in path + + def fake_listdir(path): + if path.endswith("/ports"): + return ["1"] + if path.endswith("/gids"): + return ["3"] + return ["mlx5_0", "mlx5_1"] + + with ( + mock.patch("os.path.isdir", side_effect=fake_isdir), + mock.patch("os.listdir", side_effect=fake_listdir), + mock.patch("builtins.open", side_effect=self._multi_hca_open("none")), + mock.patch("relax.utils.rdma_probe._check_mooncake_import") as mi, + mock.patch("relax.utils.rdma_probe.resource.getrlimit", return_value=(-1, -1)), + ): + mi.return_value = CheckResult("mooncake_import", True, "ok") + result = probe_node("") + assert result.effective_protocol == "tcp" + assert "HCA port not ACTIVE" in result.errors def test_unreachable_external_master_disables_mooncake(self, monkeypatch): monkeypatch.setattr(rdma_probe, "_check_mooncake_import", lambda: CheckResult("mooncake_import", True)) @@ -394,7 +474,7 @@ def test_storage_backend_key_selects_the_manager(self): """``tq.init`` reads ``backend.storage_backend``; omitting it silently keeps SimpleStorage.""" eff = EffectiveConfig(backend="MooncakeStore", protocol="rdma", device="rdma0", gdr=False, fallback_reason="") - assert build_mooncake_config(eff)["storage_backend"] == "MooncakeStore" + assert build_mooncake_config(eff, master_address="master.example:50051")["storage_backend"] == "MooncakeStore" assert ( build_simple_storage_config(total_storage_size=1, num_data_storage_units=1)["storage_backend"] == "SimpleStorage" @@ -402,7 +482,7 @@ def test_storage_backend_key_selects_the_manager(self): def test_mooncake_config_has_hard_pin_true(self): eff = EffectiveConfig(backend="MooncakeStore", protocol="rdma", device="rdma0", gdr=False, fallback_reason="") - cfg = build_mooncake_config(eff) + cfg = build_mooncake_config(eff, master_address="master.example:50051") mc = cfg["MooncakeStore"] assert mc["protocol"] == "rdma" assert mc["device_name"] == "rdma0" @@ -412,9 +492,23 @@ def test_mooncake_config_has_hard_pin_true(self): def test_mooncake_config_gdr_propagated(self): eff = EffectiveConfig(backend="MooncakeStore", protocol="rdma", device="", gdr=True, fallback_reason="") - cfg = build_mooncake_config(eff) + cfg = build_mooncake_config(eff, master_address="master.example:50051") assert cfg["MooncakeStore"]["use_gdr"] is True + def test_master_address_is_required(self, monkeypatch): + # A loopback default would point every node at itself in multi-node + # runs; missing deployment configuration must be rejected instead. + monkeypatch.delenv("MC_MASTER_ADDRESS", raising=False) + with pytest.raises(RuntimeError, match="MC_MASTER_ADDRESS"): + resolve_mooncake_master_address() + eff = EffectiveConfig(backend="MooncakeStore", protocol="tcp", device="", gdr=False, fallback_reason="") + with pytest.raises(RuntimeError, match="MC_MASTER_ADDRESS"): + build_mooncake_config(eff) + + def test_master_address_from_env(self, monkeypatch): + monkeypatch.setenv("MC_MASTER_ADDRESS", "master.example:50051") + assert resolve_mooncake_master_address() == "master.example:50051" + @pytest.mark.skipif( not _REAL_TQ_STORAGE, reason="needs real TransferQueue storage submodules; CPU CI uses a single-file transfer_queue stub", @@ -433,14 +527,20 @@ def test_contract_defaults_mooncake_memcpy_off(self, monkeypatch): validate_mooncake_runtime_contract() assert os.environ["MC_STORE_MEMCPY"] == "0" - @pytest.mark.skipif( - not _REAL_TQ_STORAGE, - reason="needs real TransferQueue storage submodules; CPU CI uses a single-file transfer_queue stub", - ) - def test_contract_respects_explicit_memcpy_override(self, monkeypatch): + def test_contract_rejects_explicit_memcpy_enable(self, monkeypatch): + # mooncake 0.3.10's memcpy path is confirmed to corrupt data, so the + # guard fails closed instead of honouring an operator override. The + # rejection happens before any transfer_queue import, so this test + # runs on CPU CI too. monkeypatch.setenv("MC_STORE_MEMCPY", "1") - validate_mooncake_runtime_contract() - assert os.environ["MC_STORE_MEMCPY"] == "1" + with pytest.raises(RuntimeError, match="MC_STORE_MEMCPY"): + validate_mooncake_runtime_contract() + + def test_contract_accepts_explicit_memcpy_disable(self, monkeypatch): + monkeypatch.setenv("MC_STORE_MEMCPY", "0") + if _REAL_TQ_STORAGE: + validate_mooncake_runtime_contract() + assert os.environ["MC_STORE_MEMCPY"] == "0" def test_segment_capacity_text_only_passes(self): args = _make_args(multimodal_keys=None) From e209e1b080a6641e1ce623b6b0d53d8d4b939ea1 Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Fri, 14 Aug 2026 21:04:07 +0800 Subject: [PATCH 08/22] fix(data-plane): close TQ owner, detach clients MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 🐛 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/backends/megatron/actor.py | 13 ++++- relax/components/base.py | 14 ++++++ relax/core/controller.py | 75 ++++++++++++++++------------ relax/distributed/ray/rollout.py | 5 +- relax/utils/tq_lifecycle.py | 15 ++++++ tests/utils/test_tq_failure_paths.py | 35 +++++++++++++ 6 files changed, 124 insertions(+), 33 deletions(-) diff --git a/relax/backends/megatron/actor.py b/relax/backends/megatron/actor.py index 4ce5c3943..4012a4f7e 100644 --- a/relax/backends/megatron/actor.py +++ b/relax/backends/megatron/actor.py @@ -60,7 +60,7 @@ from relax.utils.rotate_ckpt import rotate_ckpt from relax.utils.s3_model_loader import prepare_model_maybe_update_args from relax.utils.timer import Timer, inverse_timer, timer, with_defer -from relax.utils.tq_lifecycle import attach_tq_client +from relax.utils.tq_lifecycle import attach_tq_client, detach_tq_client from relax.utils.tracking_utils import init_tracking from relax.utils.training import train_dump_utils from relax.utils.training.data_fields import build_data_fields @@ -151,6 +151,17 @@ def _per_step_rollout(self) -> bool: periodic predict steps; Megatron stays awake between.""" return not is_sft_mode(self.args) + def __del__(self) -> None: + # Best-effort detach on graceful teardown; ray.kill / fate-sharing + # kills skip destructors, in which case the Mooncake master TTL + # reclaims the segment. + if getattr(self, "data_system_client", None) is None: + return + try: + detach_tq_client() + except Exception: # destructor must never raise (interpreter shutdown) + return + def init( self, args: Namespace, diff --git a/relax/components/base.py b/relax/components/base.py index 7e1425361..e93b28a9e 100644 --- a/relax/components/base.py +++ b/relax/components/base.py @@ -90,6 +90,20 @@ def __init__(self) -> None: self._logger_instance = None self._lock = threading.Lock() + def __del__(self) -> None: + # Ray Serve calls the destructor on replica shutdown (normal stop, + # global restart, in-place restart). Components that attached a + # TransferQueue client must detach so a MooncakeStore segment + # deregisters before client_ttl instead of leaving a stale endpoint. + if getattr(self, "data_system_client", None) is None: + return + try: + from relax.utils.tq_lifecycle import detach_tq_client + + detach_tq_client() + except Exception: # destructor must never raise (interpreter shutdown) + return + @property def _logger(self): """Lazily create and cache a logger for this instance. diff --git a/relax/core/controller.py b/relax/core/controller.py index cf2934e42..cfb4be00f 100644 --- a/relax/core/controller.py +++ b/relax/core/controller.py @@ -145,42 +145,55 @@ def __init__(self, config: Namespace, runtime_env: dict = None) -> None: # Initialize data management system self._initialize_data_system() - self.dcs, self.config.coordinator_url = create_dcs_deployment() + try: + self.dcs, self.config.coordinator_url = create_dcs_deployment() - self._metrics_service_enabled = getattr(config, "use_metrics_service", False) - if self._metrics_service_enabled: - self._deploy_metrics_service() + self._metrics_service_enabled = getattr(config, "use_metrics_service", False) + if self._metrics_service_enabled: + self._deploy_metrics_service() - if self.config.use_agentic_rollout and not self.config.debug_train_only: - deploy_agentic_chat_api_services( - config=self.config, - runtime_env=self.runtime_env, - ) - self._autoscaler_config = None - try: - self.register_all_serve() - except Exception as e: - self._report_error_to_metrics_service(e) - raise + if self.config.use_agentic_rollout and not self.config.debug_train_only: + deploy_agentic_chat_api_services( + config=self.config, + runtime_env=self.runtime_env, + ) + self._autoscaler_config = None + try: + self.register_all_serve() + except Exception as e: + self._report_error_to_metrics_service(e) + raise - autoscaler_config_path = getattr(config, "autoscaler_config", None) - if autoscaler_config_path: - from relax.utils.autoscaler.config import AutoscalerConfig - from relax.utils.utils import get_serve_url + autoscaler_config_path = getattr(config, "autoscaler_config", None) + if autoscaler_config_path: + from relax.utils.autoscaler.config import AutoscalerConfig + from relax.utils.utils import get_serve_url - rollout_service_url = get_serve_url("/rollout") - self._autoscaler_config = AutoscalerConfig.from_yaml(autoscaler_config_path, rollout_service_url) - self._deploy_autoscaler_service() + rollout_service_url = get_serve_url("/rollout") + self._autoscaler_config = AutoscalerConfig.from_yaml(autoscaler_config_path, rollout_service_url) + self._deploy_autoscaler_service() - # Start health management with service restart callback - if self._health_check_enabled: - self._health_manager.start( - on_unhealthy=self._on_service_unhealthy, - on_fatal=self._on_service_fatal, - ) - logger.info("Global health check system enabled") - else: - logger.info("Global health check system disabled (use --use-health-check to enable)") + # Start health management with service restart callback + if self._health_check_enabled: + self._health_manager.start( + on_unhealthy=self._on_service_unhealthy, + on_fatal=self._on_service_fatal, + ) + logger.info("Global health check system enabled") + else: + logger.info("Global health check system disabled (use --use-health-check to enable)") + except Exception: + # Past this point a failed construction means Controller() never + # returns: train.main() cannot install its signal/atexit cleanup, + # so the TQ owner (and any TransferQueueController it owns) would + # be orphaned and the next launch would attach with owner=None, + # making its shutdown a no-op. Close what this job created. + logger.error("Controller construction failed after TQ initialization; closing TQ owner.") + try: + self._close_data_system() + except Exception as cleanup_error: # pragma: no cover - best effort + logger.warning(f"TQ owner cleanup during failed construction failed: {cleanup_error}") + raise def _cleanup_s3_model_weights_after_init(self) -> None: """Remove policy weight shards after every startup consumer is diff --git a/relax/distributed/ray/rollout.py b/relax/distributed/ray/rollout.py index d045fb95c..2b274b06a 100644 --- a/relax/distributed/ray/rollout.py +++ b/relax/distributed/ray/rollout.py @@ -51,7 +51,7 @@ from relax.utils.opd.opd_utils import compute_mopd_metrics from relax.utils.reload_utils import ReloadableMixin from relax.utils.s3_model_loader import prepare_model_maybe_update_args -from relax.utils.tq_lifecycle import attach_tq_client +from relax.utils.tq_lifecycle import attach_tq_client, detach_tq_client from relax.utils.tracking_utils import init_tracking from relax.utils.training.train_dump_utils import ( save_debug_rollout_data, @@ -920,6 +920,9 @@ def dispose(self): for monitor in self._health_monitors: monitor.stop() self._shutdown_all_engines() + # Deregister this worker's Mooncake segment before the actor dies so a + # fast restart does not hit stale endpoints until client_ttl expires. + detach_tq_client() def _shutdown_all_engines(self, timeout: float = 15.0): """Shut down all SGLang engine actors and their child processes. diff --git a/relax/utils/tq_lifecycle.py b/relax/utils/tq_lifecycle.py index 436d6cd43..66320b600 100644 --- a/relax/utils/tq_lifecycle.py +++ b/relax/utils/tq_lifecycle.py @@ -285,6 +285,21 @@ def attach_tq_client(conf: Any, *, requested_gdr: bool, role: str) -> Any: return client +def detach_tq_client() -> None: + """Detach this worker's TQ client (attach-only inverse of + :func:`attach_tq_client`). + + Closes the attached Mooncake storage client so its segment deregisters + from the master immediately instead of lingering until ``client_ttl`` + expires — a stale endpoint breaks fast restarts with "Failed to open + segment". Only process-local handles are touched (never the named + controller or globally stored data), so every worker teardown hook may + call this unconditionally; force-killed workers still fall back to the + master-side TTL. + """ + _close_local_tq_client() + + def close_tq_and_unmount(*, is_owner: bool) -> None: """Close TransferQueue and unmount the MooncakeStore segment. diff --git a/tests/utils/test_tq_failure_paths.py b/tests/utils/test_tq_failure_paths.py index 86f196bd8..8e076a9a4 100644 --- a/tests/utils/test_tq_failure_paths.py +++ b/tests/utils/test_tq_failure_paths.py @@ -267,6 +267,41 @@ def test_attached_process_never_calls_global_close(self, monkeypatch): store_client.close.assert_called_once() +# --------------------------------------------------------------------------- +# Worker detach (attach-only inverse used by every worker teardown hook) +# --------------------------------------------------------------------------- + + +class TestWorkerDetach: + """detach_tq_client and the teardown hooks that must invoke it.""" + + def test_detach_delegates_to_local_close(self, monkeypatch): + calls = [] + monkeypatch.setattr(tq_lifecycle, "_close_local_tq_client", lambda: calls.append(True)) + tq_lifecycle.detach_tq_client() + assert calls == [True] + + def test_component_del_detaches_attached_client(self, monkeypatch): + from relax.components.base import Base + + calls = [] + monkeypatch.setattr(tq_lifecycle, "detach_tq_client", lambda: calls.append(True)) + component = Base() + component.data_system_client = object() + component.__del__() + assert calls == [True] + component.data_system_client = None # keep GC-time __del__ a no-op + + def test_component_del_without_client_is_noop(self, monkeypatch): + from relax.components.base import Base + + calls = [] + monkeypatch.setattr(tq_lifecycle, "detach_tq_client", lambda: calls.append(True)) + component = Base() + component.__del__() + assert calls == [] + + # --------------------------------------------------------------------------- # GDR requested-vs-runtime status # --------------------------------------------------------------------------- From 0f579fde04e5b02fd8b9374a9b4639a64081a021 Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Fri, 14 Aug 2026 21:31:32 +0800 Subject: [PATCH 09/22] fix(data-plane): restore default path, bound attach MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 🐛 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 --- relax/core/controller.py | 80 ++++++++++++- relax/utils/tq_lifecycle.py | 161 ++++++++++++++++++++++++++- tests/utils/test_tq_failure_paths.py | 45 ++++++++ 3 files changed, 279 insertions(+), 7 deletions(-) diff --git a/relax/core/controller.py b/relax/core/controller.py index cfb4be00f..42e4b60ad 100644 --- a/relax/core/controller.py +++ b/relax/core/controller.py @@ -8,6 +8,7 @@ from typing import Any import ray +import transfer_queue as tq from omegaconf import OmegaConf from ray import serve from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy @@ -53,7 +54,14 @@ resolve_mooncake_master_address, validate_mooncake_runtime_contract, ) -from relax.utils.tq_lifecycle import close_tq_owner, initialize_tq_with_fallback +from relax.utils.tq_lifecycle import ( + TqInitResult, + close_tq_owner, + initialize_tq_with_fallback, + reap_unusable_tq_controller, + uses_mooncake, + verify_cluster_attach, +) from relax.utils.training.ppo_utils import validate_ppo_config from relax.utils.utils import compute_dp_size, recovery_load_path @@ -138,6 +146,7 @@ def __init__(self, config: Namespace, runtime_env: dict = None) -> None: if not hasattr(self, "_global_restart_count"): self._global_restart_count = 0 self._tq_owner = None + self._tq_legacy_init = False # SFT: fill in num_rollout / num_rollout_per_epoch before any actor # is launched (RL is resolved later in placement_group.py). @@ -304,6 +313,20 @@ def _initialize_data_system(self): flags={"allow_objects": True}, ) + if getattr(self.config, "tq_storage_backend", "simple") == "simple": + # Zero-behavior-change default path (review PR#256): identical to + # upstream, the first tq.init runs inside the Controller process + # and no _TransferQueueOwner actor is created. The only addition + # is the F10 reaper, which acts solely on a provably + # half-initialised leftover controller that would otherwise make + # this tq.init poll get_config forever. + reap_unusable_tq_controller() + self._tq_owner = None + self._tq_legacy_init = True + self.config.tq_config = tq.init(conf=tq_config) or tq_config + logger.info("[dataplane] backend=SimpleStorage (default in-process init, no owner actor)") + return + fallback_config = None if backend_config.get("storage_backend") == "MooncakeStore": from relax.utils.tq_config import build_simple_storage_config @@ -324,12 +347,55 @@ def _initialize_data_system(self): mode=getattr(self.config, "tq_rdma_mode", "off"), fallback_conf=fallback_config, ) + if uses_mooncake(init_result.config): + init_result = self._confirm_mooncake_attach(init_result, fallback_config) self._tq_owner = init_result.owner self.config.tq_config = init_result.config if init_result.fallback_reason: logger.warning(f"[dataplane] effective backend=SimpleStorage fallback={init_result.fallback_reason}") logger.info(f"[dataplane] controller ownership={'owner' if init_result.owns_controller else 'attached'}") + def _confirm_mooncake_attach(self, init_result: TqInitResult, fallback_config) -> TqInitResult: + """Bounded attach handshake on every alive node before the job-level + Mooncake config is confirmed. + + The RDMA probe cannot cover this: TQ clients live in Ray Serve replicas + and 0-CPU actors with no placement binding, so they may land on nodes + the GPU-only probe never saw, and a worker-side ``tq.init`` has no + timeout of its own. The handshake attaches (bounded) from every alive + node and reports failures back here, so ``auto`` degrades the whole job + to one backend instead of hanging or failing a single replica mid- + deployment. + """ + failures = verify_cluster_attach(init_result.config) + if not failures: + logger.info("[dataplane] Mooncake attach handshake passed on all alive nodes.") + return init_result + + detail = "; ".join(failures) + mode = getattr(self.config, "tq_rdma_mode", "off") + if mode != "auto" or not init_result.owns_controller or fallback_config is None: + # off/required must fail loudly, and a job that merely attached to + # a foreign controller must never tear it down or replace its + # backend unilaterally. + if init_result.owns_controller: + close_tq_owner(init_result.owner) + raise RuntimeError( + f"Mooncake attach handshake failed on {len(failures)} node(s) (--tq-rdma-mode={mode}): {detail}" + ) + + logger.warning( + f"[dataplane] Mooncake attach handshake failed on {len(failures)} node(s) ({detail}); " + "closing Mooncake state and converging the whole job to SimpleStorage." + ) + close_tq_owner(init_result.owner) + fallback_result = initialize_tq_with_fallback(fallback_config, mode="auto") + return TqInitResult( + config=fallback_result.config, + owner=fallback_result.owner, + fallback_reason=f"attach_handshake_failed:{len(failures)}_nodes", + ) + def _resolve_tq_backend(self, total_storage_size: int) -> dict: """Resolve the TransferQueue ``backend`` config dict. @@ -446,8 +512,16 @@ def _resolve_tq_backend(self, total_storage_size: int) -> dict: return backend_dict def _close_data_system(self) -> None: - """Delegate to - :func:`relax.utils.tq_lifecycle.close_tq_and_unmount`.""" + """Tear down the data system. + + Default in-process path keeps upstream's plain ``tq.close()``; + owner-mediated runs delegate to + :func:`relax.utils.tq_lifecycle.close_tq_owner`. + """ + if self._tq_legacy_init: + self._tq_legacy_init = False + tq.close() + return close_tq_owner(self._tq_owner) self._tq_owner = None diff --git a/relax/utils/tq_lifecycle.py b/relax/utils/tq_lifecycle.py index 66320b600..ed690e1e4 100644 --- a/relax/utils/tq_lifecycle.py +++ b/relax/utils/tq_lifecycle.py @@ -21,6 +21,8 @@ from __future__ import annotations +import os +import threading import time import uuid from dataclasses import dataclass @@ -38,6 +40,7 @@ CONTROLLER_NAMESPACE = "transfer_queue" OWNER_TOKEN_FIELD = "relax_owner_token" DEFAULT_TQ_INIT_TIMEOUT_SECONDS = 60.0 +DEFAULT_TQ_ATTACH_TIMEOUT_SECONDS = 60.0 @dataclass(frozen=True) @@ -57,6 +60,10 @@ class TqInitializationTimeout(TimeoutError): """Raised when ``tq.init`` does not finish within the bounded timeout.""" +class TqAttachTimeout(TimeoutError): + """Raised when a worker cannot attach to TransferQueue within the bound.""" + + class TqCleanupTimeout(TimeoutError): """Raised when a TQ controller cannot be confirmed gone after cleanup.""" @@ -113,6 +120,15 @@ def _uses_mooncake(conf: Any) -> bool: return _get_config_value(backend, "storage_backend", "SimpleStorage") == "MooncakeStore" +def uses_mooncake(conf: Any) -> bool: + """True when ``conf`` selects the MooncakeStore backend. + + Public so the Controller can decide whether the cluster-wide attach + handshake is required for the stored job-level config. + """ + return _uses_mooncake(conf) + + def _prepare_mooncake_runtime(conf: Any) -> None: if not _uses_mooncake(conf): return @@ -275,11 +291,96 @@ def log_tq_gdr_runtime_status(*, requested: bool, role: str) -> str: return status -def attach_tq_client(conf: Any, *, requested_gdr: bool, role: str) -> Any: - """Attach a component process and report its local experimental GDR - state.""" +def _resolve_attach_timeout() -> float: + """Attach deadline in seconds; override via + ``RELAX_TQ_ATTACH_TIMEOUT_SECONDS``.""" + raw = os.environ.get("RELAX_TQ_ATTACH_TIMEOUT_SECONDS", "").strip() + if not raw: + return DEFAULT_TQ_ATTACH_TIMEOUT_SECONDS + try: + value = float(raw) + except ValueError as error: + raise RuntimeError(f"RELAX_TQ_ATTACH_TIMEOUT_SECONDS={raw!r} must be a positive number of seconds") from error + if value <= 0: + raise RuntimeError(f"RELAX_TQ_ATTACH_TIMEOUT_SECONDS={raw!r} must be a positive number of seconds") + return value + + +def _await_controller_config(deadline: float) -> None: + """Bounded wait until the named controller serves a non-``None`` config. + + ``tq.init`` polls ``get_config`` forever while it returns ``None`` (the F10 + hang), so a worker refuses to enter that loop unless a config is provably + served before the deadline. + """ + last_error = "TransferQueueController named actor not found" + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TqAttachTimeout(f"TransferQueue attach timed out: {last_error}") + try: + controller = ray.get_actor(CONTROLLER_NAME, namespace=CONTROLLER_NAMESPACE) + except ValueError: + time.sleep(min(0.5, remaining)) + continue + try: + conf = ray.get(controller.get_config.remote(), timeout=max(min(remaining, 10.0), 0.1)) + except Exception as e: + last_error = f"get_config failed: {e}" + time.sleep(min(0.5, max(deadline - time.monotonic(), 0.0))) + continue + if conf is not None: + return + last_error = "controller exists but has stored no config yet" + time.sleep(min(0.5, max(deadline - time.monotonic(), 0.0))) + + +def _bounded_tq_init(conf: Any, deadline: float, *, role: str) -> None: + """Run ``tq.init`` under the remaining deadline. + + ``tq.init`` takes no timeout and its mooncake client setup blocks in native + code, so it runs on a daemon watchdog thread. On expiry the caller fails + fast with :class:`TqAttachTimeout`; the abandoned thread dies with the + failed worker process (Serve tears the replica down). + """ + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TqAttachTimeout(f"TransferQueue attach for role={role} timed out before tq.init") + error: list[BaseException] = [] + + def _run() -> None: + try: + tq.init(conf=conf) + except BaseException as e: # propagated to the attaching caller below + error.append(e) + + thread = threading.Thread(target=_run, name=f"tq-attach-{role}", daemon=True) + thread.start() + thread.join(remaining) + if thread.is_alive(): + raise TqAttachTimeout( + f"tq.init for role={role} did not finish within {remaining:.0f}s " + "(mooncake setup or controller poll hung); failing this worker fast." + ) + if error: + raise error[0] + + +def attach_tq_client(conf: Any, *, requested_gdr: bool, role: str, timeout: float | None = None) -> Any: + """Attach a component process within a bounded deadline and report its + local experimental GDR state. + + The deadline covers both waiting for a served controller config and + ``tq.init`` itself, because either phase can hang unboundedly (get_config + polling and mooncake endpoint setup respectively). ``None`` resolves the + deadline from ``RELAX_TQ_ATTACH_TIMEOUT_SECONDS`` (default 60 s). + """ + if timeout is None: + timeout = _resolve_attach_timeout() + deadline = time.monotonic() + timeout _prepare_mooncake_runtime(conf) - tq.init(conf=conf) + _await_controller_config(deadline) + _bounded_tq_init(conf, deadline, role=role) client = tq.get_client() log_tq_gdr_runtime_status(requested=requested_gdr, role=role) return client @@ -300,6 +401,58 @@ def detach_tq_client() -> None: _close_local_tq_client() +def _alive_node_ids() -> list[str]: + """Every alive node: TQ endpoints (Serve replicas and 0-CPU actors) carry + no placement binding, so any alive node may end up hosting one.""" + return [n["NodeID"] for n in ray.nodes() if n.get("Alive")] + + +def verify_cluster_attach(conf: Any, *, timeout: float | None = None) -> list[str]: + """Bounded attach handshake from every alive node; returns failure + summaries. + + Each task performs the same bounded :func:`attach_tq_client` a worker would + perform (real stored config, real storage client) and detaches immediately, + so it validates the *actual endpoints* instead of a ``/sys`` capability + heuristic on a node the scheduler may never use. An empty return means + every alive node attached within the deadline; the driver aggregates + failures and decides one job-level outcome. + """ + if timeout is None: + timeout = _resolve_attach_timeout() + from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy + + @ray.remote(num_cpus=0, max_retries=0) + def _handshake(handshake_conf: Any) -> None: + from relax.utils.tq_lifecycle import attach_tq_client, detach_tq_client + + attach_tq_client(handshake_conf, requested_gdr=False, role="attach-handshake") + detach_tq_client() + + refs: list[Any] = [] + id_by_ref: dict[Any, str] = {} + for node_id in _alive_node_ids(): + strategy = NodeAffinitySchedulingStrategy(node_id=node_id, soft=False) + ref = _handshake.options(scheduling_strategy=strategy).remote(conf) + refs.append(ref) + id_by_ref[ref] = node_id + + # Grace beyond the per-node attach deadline covers task scheduling and + # worker startup on a busy cluster. + wait_bound = timeout + 30.0 + ready, pending = ray.wait(refs, num_returns=len(refs), timeout=wait_bound) + failures: list[str] = [] + for ref in ready: + try: + ray.get(ref) + except Exception as e: + failures.append(f"node {id_by_ref[ref][:12]}: {e}") + for ref in pending: + ray.cancel(ref, force=True) + failures.append(f"node {id_by_ref[ref][:12]}: handshake did not return within {wait_bound:.0f}s") + return failures + + def close_tq_and_unmount(*, is_owner: bool) -> None: """Close TransferQueue and unmount the MooncakeStore segment. diff --git a/tests/utils/test_tq_failure_paths.py b/tests/utils/test_tq_failure_paths.py index 8e076a9a4..86b6c1b8e 100644 --- a/tests/utils/test_tq_failure_paths.py +++ b/tests/utils/test_tq_failure_paths.py @@ -267,6 +267,51 @@ def test_attached_process_never_calls_global_close(self, monkeypatch): store_client.close.assert_called_once() +# --------------------------------------------------------------------------- +# Bounded attach (worker-side tq.init used to hang forever) +# --------------------------------------------------------------------------- + + +class TestBoundedAttach: + """attach_tq_client: one deadline for get_config wait and tq.init.""" + + def test_attach_timeout_env_rejects_garbage(self, monkeypatch): + monkeypatch.setenv("RELAX_TQ_ATTACH_TIMEOUT_SECONDS", "soon") + with pytest.raises(RuntimeError, match="RELAX_TQ_ATTACH_TIMEOUT_SECONDS"): + tq_lifecycle._resolve_attach_timeout() + + def test_attach_timeout_env_override_is_used(self, monkeypatch): + monkeypatch.setenv("RELAX_TQ_ATTACH_TIMEOUT_SECONDS", "12.5") + assert tq_lifecycle._resolve_attach_timeout() == 12.5 + + def test_bounded_init_times_out_on_hung_tq_init(self, monkeypatch): + import time + + monkeypatch.setattr(tq_lifecycle.tq, "init", lambda conf: time.sleep(5)) + with pytest.raises(tq_lifecycle.TqAttachTimeout, match="did not finish"): + tq_lifecycle._bounded_tq_init({}, time.monotonic() + 0.2, role="test") + + def test_bounded_init_propagates_worker_error(self, monkeypatch): + import time + + def boom(conf): + raise ValueError("bad conf") + + monkeypatch.setattr(tq_lifecycle.tq, "init", boom) + with pytest.raises(ValueError, match="bad conf"): + tq_lifecycle._bounded_tq_init({}, time.monotonic() + 5.0, role="test") + + def test_await_controller_config_times_out_without_actor(self, monkeypatch): + import time + + def no_actor(name, namespace=None): + raise ValueError("actor not found") + + monkeypatch.setattr(tq_lifecycle.ray, "get_actor", no_actor) + with pytest.raises(tq_lifecycle.TqAttachTimeout, match="attach timed out"): + tq_lifecycle._await_controller_config(time.monotonic() + 0.3) + + # --------------------------------------------------------------------------- # Worker detach (attach-only inverse used by every worker teardown hook) # --------------------------------------------------------------------------- From 18c758dc00c82d8fc986a782f015158e43eb4986 Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Fri, 14 Aug 2026 21:32:20 +0800 Subject: [PATCH 10/22] fix(data-plane): size segment from token budget MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 🐛 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 --- relax/utils/tq_config.py | 72 +++++++++++++++++++++++++++------- tests/utils/test_rdma_probe.py | 48 ++++++++++++++++++++++- 2 files changed, 103 insertions(+), 17 deletions(-) diff --git a/relax/utils/tq_config.py b/relax/utils/tq_config.py index dc5cf7d77..224e78ef4 100644 --- a/relax/utils/tq_config.py +++ b/relax/utils/tq_config.py @@ -54,6 +54,27 @@ def resolve_mooncake_master_address() -> str: return address +def resolve_global_segment_size() -> int: + """Per-client Mooncake segment size in bytes. + + Defaults to 4 GiB (TQ config.yaml:52). Deployments whose worst-case in- + flight payload exceeds that (see :func:`estimate_payload_bytes`) set + ``RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB`` instead of editing code; capacity + validation and the client config read the same value so the check can never + pass a size the client does not actually mount. + """ + raw = os.environ.get("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB", "").strip() + if not raw: + return _DEFAULT_GLOBAL_SEGMENT_SIZE + try: + gib = float(raw) + except ValueError as error: + raise RuntimeError(f"RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB={raw!r} must be a positive number of GiB") from error + if gib <= 0: + raise RuntimeError(f"RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB={raw!r} must be a positive number of GiB") + return int(gib * 1024**3) + + def validate_mooncake_runtime_contract() -> None: """Install and validate the Mooncake loss-prevention contract. @@ -114,8 +135,10 @@ def build_mooncake_config( ``MC_MASTER_ADDRESS`` env var (see :func:`resolve_mooncake_master_address`). global_segment_size - Override the per-client segment size (default 4 GiB). Benchmarks may - pass a larger value (e.g. 8 GiB) to avoid staging-buffer pressure. + Override the per-client segment size. ``None`` resolves from + ``RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB`` (default 4 GiB, see + :func:`resolve_global_segment_size`). Benchmarks may pass a larger + value (e.g. 8 GiB) to avoid staging-buffer pressure. """ if master_address is None: master_address = resolve_mooncake_master_address() @@ -136,7 +159,7 @@ def build_mooncake_config( "metadata_server": _DEFAULT_METADATA_SERVER, "local_hostname": "", # empty = auto-detect via Ray node IP # Memory - "global_segment_size": global_segment_size or _DEFAULT_GLOBAL_SEGMENT_SIZE, + "global_segment_size": global_segment_size or resolve_global_segment_size(), "local_buffer_size": _DEFAULT_LOCAL_BUFFER_SIZE, # Do NOT silently evict produced-but-unconsumed data. "hard_pin": True, @@ -153,19 +176,38 @@ def build_mooncake_config( # --------------------------------------------------------------------------- -def estimate_payload_bytes(args: Any) -> int: - """Rough estimate of per-step multimodal payload size in bytes. +# Worst-case payload factors used by the segment-capacity pre-check. +# Vision: a ViT-style processor (Qwen-VL family: patch 14x14, spatial merge +# 2x2) maps one schedulable token to at most (14*2)^2 = 784 pixels, and +# ``pixel_values`` is float32 RGB, so vision bytes <= seq_length * 784 * 12. +_PIXELS_PER_VISION_TOKEN = 28 * 28 +_BYTES_PER_PIXEL_VALUE = 3 * 4 +# Text: token ids, logprobs, masks and rewards; 32 B/token rounds them up. +_TEXT_BYTES_PER_TOKEN = 32 - Used only for segment-capacity pre-check. The real payload depends on - image resolution and patch count; this is a conservative lower bound based - on ``--multimodal-keys`` presence and ``n_samples_per_prompt``. + +def estimate_payload_bytes(args: Any) -> int: + """Worst-case per-step payload upper bound in bytes. + + Derived from the token budget instead of a fixed per-sample constant: the + processor cannot emit more vision tokens than ``--seq-length`` allows, so + the pixel payload of one sample is bounded by + ``seq_length * _PIXELS_PER_VISION_TOKEN * _BYTES_PER_PIXEL_VALUE`` + (e.g. 77 MiB at seq_length=8192) rather than the old 8 MiB guess that + passed configurations which later failed puts mid-training. """ n_samples = args.n_samples_per_prompt rollout_batch = args.rollout_batch_size - # 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 + seq_length = int(getattr(args, "seq_length", 0) or 0) + if seq_length <= 0: + raise RuntimeError( + "MooncakeStore segment-capacity validation needs args.seq_length to bound " + "the per-sample payload; got a missing or non-positive value." + ) + per_sample = seq_length * _TEXT_BYTES_PER_TOKEN + if getattr(args, "multimodal_keys", None) is not None: + per_sample += seq_length * _PIXELS_PER_VISION_TOKEN * _BYTES_PER_PIXEL_VALUE + return rollout_batch * n_samples * per_sample def validate_segment_capacity(args: Any, effective: EffectiveConfig) -> str | None: @@ -182,15 +224,15 @@ def validate_segment_capacity(args: Any, effective: EffectiveConfig) -> str | No max_staleness = getattr(args, "max_staleness", 0) payload = estimate_payload_bytes(args) needed = payload * (max_staleness + 1) - available = _DEFAULT_GLOBAL_SEGMENT_SIZE + available = resolve_global_segment_size() if needed > available: return ( - f"MooncakeStore segment capacity insufficient: estimated in-flight payload " + f"MooncakeStore segment capacity insufficient: worst-case in-flight payload " f"{needed / 1024**3:.1f} GiB (rollout_batch={args.rollout_batch_size} × " f"n_samples={args.n_samples_per_prompt} × staleness+1={max_staleness + 1}) " f"exceeds global_segment_size {available / 1024**3:.1f} GiB. " - f"Reduce batch size, increase global_segment_size, or reduce max_staleness." + f"Reduce batch size / max_staleness, or raise RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB." ) return None diff --git a/tests/utils/test_rdma_probe.py b/tests/utils/test_rdma_probe.py index 25c53b285..5d80631d7 100644 --- a/tests/utils/test_rdma_probe.py +++ b/tests/utils/test_rdma_probe.py @@ -81,6 +81,7 @@ def _make_args(**kwargs) -> argparse.Namespace: n_samples_per_prompt=1, rollout_batch_size=32, multimodal_keys=None, + seq_length=8192, ) defaults.update(kwargs) return argparse.Namespace(**defaults) @@ -556,6 +557,49 @@ def test_segment_capacity_multimodal_large_batch_fails(self): assert err is not None assert "insufficient" in err.lower() - def test_estimate_payload_text_only_is_zero(self): + def test_estimate_payload_text_only_is_small_but_nonzero(self): + # Text payloads (ids/logprobs/masks) flow through the store too; the + # bound is seq_length * 32 B per sample. args = _make_args(multimodal_keys=None) - assert estimate_payload_bytes(args) == 0 + assert estimate_payload_bytes(args) == 32 * 1 * 8192 * 32 + + def test_estimate_payload_multimodal_is_token_budget_bound(self): + # One sample may not exceed seq_length vision tokens; at 784 pixels + # per token and 12 B per pixel that is ~77 MiB for seq_length=8192. + args = _make_args(multimodal_keys=["pixel_values"], rollout_batch_size=1, n_samples_per_prompt=1) + per_sample = estimate_payload_bytes(args) + assert per_sample == 8192 * (32 + 784 * 12) + assert 70 * 1024**2 < per_sample < 80 * 1024**2 + + def test_estimate_payload_requires_seq_length(self): + args = _make_args(seq_length=None) + with pytest.raises(RuntimeError, match="seq_length"): + estimate_payload_bytes(args) + + def test_segment_capacity_multimodal_staleness_no_longer_passes(self): + # Review (Codex P1): 32 in-flight samples x ~77 MiB x (staleness+1)=2 + # needs ~4.9 GiB and previously passed the 8 MiB/sample guess. + args = _make_args( + multimodal_keys=["pixel_values"], rollout_batch_size=32, n_samples_per_prompt=1, max_staleness=1 + ) + eff = EffectiveConfig(backend="MooncakeStore", protocol="rdma", device="", gdr=False, fallback_reason="") + err = validate_segment_capacity(args, eff) + assert err is not None and "RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB" in err + + def test_segment_capacity_env_override_raises_the_ceiling(self, monkeypatch): + args = _make_args( + multimodal_keys=["pixel_values"], rollout_batch_size=32, n_samples_per_prompt=1, max_staleness=1 + ) + eff = EffectiveConfig(backend="MooncakeStore", protocol="rdma", device="", gdr=False, fallback_reason="") + monkeypatch.setenv("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB", "8") + assert validate_segment_capacity(args, eff) is None + + def test_segment_size_env_override_rejects_garbage(self, monkeypatch): + from relax.utils.tq_config import resolve_global_segment_size + + monkeypatch.setenv("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB", "four") + with pytest.raises(RuntimeError, match="RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB"): + resolve_global_segment_size() + monkeypatch.setenv("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB", "-1") + with pytest.raises(RuntimeError, match="positive"): + resolve_global_segment_size() From 682400ab129a821b16c276c8470d2bcda35e22ac Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Fri, 14 Aug 2026 21:40:18 +0800 Subject: [PATCH 11/22] refactor(data-plane): split runtime patches out of PR MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # ♻️ 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 --- relax/utils/tq_correctness.py | 182 +++------------------------ tests/utils/test_tq_failure_paths.py | 108 ---------------- 2 files changed, 19 insertions(+), 271 deletions(-) diff --git a/relax/utils/tq_correctness.py b/relax/utils/tq_correctness.py index 9d60d5bec..8a5e7c91c 100644 --- a/relax/utils/tq_correctness.py +++ b/relax/utils/tq_correctness.py @@ -1,165 +1,24 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. -"""Fail-closed correctness guards for TransferQueue's Mooncake backend. - -The TransferQueue revision currently pinned by Relax validates the first -Mooncake batch result, but not every retry result, logs removal failures -without raising, and treats a missing/negative production-status ACK as a -successful notification. Those behaviours can turn an explicit storage or -controller failure into silent data loss. The pinned mooncake 0.3.10 -additionally corrupts TCP-protocol transfers through its auto-enabled memcpy -fast path, so that path is force-disabled here and an explicit enable is -rejected (see :func:`_enforce_safe_memcpy`). - -Keep the compatibility guards here, close to Relax's integration boundary, -until the equivalent checks are available in the pinned TransferQueue -revision. Installation is process-local and idempotent; every process that -creates or attaches a Mooncake client installs them before ``tq.init``. +"""Fail-closed capability validation for TransferQueue's Mooncake backend. + +Relax refuses to run MooncakeStore unless the installed TransferQueue and +mooncake expose the primitives that make silent data loss detectable. This +module validates capabilities and environment only; it never modifies +TransferQueue at runtime. The temporary runtime patches that harden the +remaining gaps of the pinned revision (per-retry result validation, raising +removal failures, and a strict production-status ACK) are maintained in a +separate version-gated PR so their exact applicability and removal condition +stay reviewable on their own. + +The pinned mooncake 0.3.10 additionally corrupts TCP-protocol transfers +through its auto-enabled memcpy fast path, so that path is force-disabled +here and an explicit enable is rejected (see :func:`_enforce_safe_memcpy`). """ from __future__ import annotations -import asyncio import os -from functools import wraps -from typing import Any -from uuid import uuid4 - -from relax.utils.logging_utils import get_logger - - -logger = get_logger(__name__) - -_PATCH_MARKER = "_relax_mooncake_correctness_guards_v1" - - -def _validate_result_count(operation: str, keys: list[str], results: Any) -> None: - """Require one Mooncake result code for every requested key.""" - try: - actual = len(results) - except TypeError as error: - raise RuntimeError(f"{operation} returned a non-sized result, expected {len(keys)} codes") from error - if actual != len(keys): - raise RuntimeError(f"{operation} returned {actual} results, expected {len(keys)}") - - -class _StrictMooncakeStoreProxy: - """Validate every low-level batch response, including retry calls.""" - - def __init__(self, store: Any) -> None: - self._store = store - - def __getattr__(self, name: str) -> Any: - return getattr(self._store, name) - - def batch_upsert_from(self, keys: list[str], *args: Any, **kwargs: Any) -> Any: - results = self._store.batch_upsert_from(keys, *args, **kwargs) - _validate_result_count("batch_upsert_from", keys, results) - return results - - def batch_get_into(self, keys: list[str], *args: Any, **kwargs: Any) -> Any: - results = self._store.batch_get_into(keys, *args, **kwargs) - _validate_result_count("batch_get_into", keys, results) - return results - - def batch_remove(self, keys: list[str], *args: Any, **kwargs: Any) -> Any: - results = self._store.batch_remove(keys, *args, **kwargs) - _validate_result_count("batch_remove", keys, results) - failures = [(key, code) for key, code in zip(keys, results, strict=True) if code != 0] - if failures: - detail = ", ".join(f"{key}={code}" for key, code in failures) - raise RuntimeError(f"batch_remove failed: {detail}") - return results - - -async def _strict_notify_and_wait(self: Any, request_msg: list) -> None: - """Notify the controller and require a positive ACK within the deadline.""" - import zmq - from transfer_queue.storage.managers import base as tq_base - from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, create_zmq_socket - - identity = f"{self.storage_manager_id}-notify-{uuid4().hex[:8]}".encode() - sock = create_zmq_socket( - ctx=self.zmq_context, - socket_type=zmq.DEALER, - ip=self.controller_info.ip, - identity=identity, - ) - sock.setsockopt(zmq.LINGER, 0) - sock.connect(self.controller_info.to_addr("request_handle_socket")) - - try: - await sock.send_multipart(request_msg) - loop = asyncio.get_running_loop() - deadline = loop.time() + tq_base.TQ_DATA_UPDATE_RESPONSE_TIMEOUT - - while True: - remaining = deadline - loop.time() - if remaining <= 0: - raise TimeoutError( - "Timed out waiting for TransferQueue production-status ACK " - f"after {tq_base.TQ_DATA_UPDATE_RESPONSE_TIMEOUT}s" - ) - try: - messages = await asyncio.wait_for( - sock.recv_multipart(copy=False), - timeout=min(tq_base.TQ_STORAGE_POLLER_TIMEOUT, remaining), - ) - except asyncio.TimeoutError: - continue - except Exception as error: - raise RuntimeError("Failed while waiting for TransferQueue production-status ACK") from error - - response = ZMQMessage.deserialize(messages) - if response.request_type != ZMQRequestType.NOTIFY_DATA_UPDATE_ACK: - continue - body = response.body if isinstance(response.body, dict) else {} - if body.get("success") is not True: - raise RuntimeError( - "TransferQueue controller rejected the production-status update " - f"for partition={body.get('partition_id', 'unknown')}" - ) - return - finally: - try: - if not sock.closed: - sock.close(linger=0) - except Exception as error: # pragma: no cover - best-effort socket cleanup - logger.debug(f"Failed to close TransferQueue notification socket: {error}") - - -def _install_store_guards(client_cls: type) -> None: - if getattr(client_cls, _PATCH_MARKER, False): - return - - original_init = client_cls.__init__ - - @wraps(original_init) - def guarded_init(self: Any, *args: Any, **kwargs: Any) -> None: - original_init(self, *args, **kwargs) - store = getattr(self, "_store", None) - if store is not None and not isinstance(store, _StrictMooncakeStoreProxy): - self._store = _StrictMooncakeStoreProxy(store) - - client_cls.__init__ = guarded_init - setattr(client_cls, _PATCH_MARKER, True) - - -def _install_notification_guards(manager_cls: type) -> None: - if getattr(manager_cls, _PATCH_MARKER, False): - return - - original_notify = manager_cls.notify_data_update - - @wraps(original_notify) - async def guarded_notify(self: Any, *args: Any, **kwargs: Any) -> None: - if not getattr(self, "controller_info", None): - raise RuntimeError("TransferQueue storage manager has no controller for production-status notification") - await original_notify(self, *args, **kwargs) - - manager_cls.notify_data_update = guarded_notify - manager_cls._notify_and_wait = _strict_notify_and_wait - setattr(manager_cls, _PATCH_MARKER, True) def _enforce_safe_memcpy() -> None: @@ -188,11 +47,14 @@ def _enforce_safe_memcpy() -> None: def ensure_mooncake_correctness_guards() -> None: - """Install and validate all guards required for safe Mooncake operation.""" + """Validate that the installed stack can run MooncakeStore safely. + + Read-only: checks the memcpy environment contract and that the pinned + TransferQueue ships the Mooncake retry APIs Relax's data plane relies on. + """ _enforce_safe_memcpy() try: from transfer_queue.storage.clients.mooncake_client import MooncakeStoreClient - from transfer_queue.storage.managers.base import StorageManager except ImportError as error: raise RuntimeError("Installed TransferQueue has no MooncakeStore support") from error @@ -200,9 +62,3 @@ def ensure_mooncake_correctness_guards() -> None: missing = [name for name in required_methods if not callable(getattr(MooncakeStoreClient, name, None))] if missing: raise RuntimeError("Installed TransferQueue lacks required Mooncake retry APIs: " + ", ".join(missing)) - - _install_store_guards(MooncakeStoreClient) - _install_notification_guards(StorageManager) - - if not getattr(MooncakeStoreClient, _PATCH_MARKER, False) or not getattr(StorageManager, _PATCH_MARKER, False): - raise RuntimeError("Failed to install Mooncake silent-data-loss guards") diff --git a/tests/utils/test_tq_failure_paths.py b/tests/utils/test_tq_failure_paths.py index 86b6c1b8e..6cc30d2dd 100644 --- a/tests/utils/test_tq_failure_paths.py +++ b/tests/utils/test_tq_failure_paths.py @@ -35,7 +35,6 @@ from relax.utils import tq_lifecycle from relax.utils.rdma_probe import ProbeResult, reduce_results -from relax.utils.tq_correctness import _strict_notify_and_wait, _StrictMooncakeStoreProxy def _has_real_submodule(dotted: str) -> bool: @@ -648,113 +647,6 @@ def _client_with_store(store) -> object: return client -class _SequenceStore: - """Return a configured result sequence from low-level Mooncake calls.""" - - def __init__(self, results: list[list[int]]) -> None: - self.results = iter(results) - - def batch_upsert_from(self, keys, ptrs, sizes, config=None): - return next(self.results) - - def batch_get_into(self, keys, ptrs, sizes): - return next(self.results) - - def batch_remove(self, keys, force=True): - return next(self.results) - - -class _FakeNotifySocket: - def __init__(self) -> None: - self.closed = False - - def setsockopt(self, *args, **kwargs) -> None: - pass - - def connect(self, *args, **kwargs) -> None: - pass - - async def send_multipart(self, request) -> None: - pass - - async def recv_multipart(self, copy=False): - return [b"ack"] - - def close(self, linger=0) -> None: - self.closed = True - - -class TestMooncakeCorrectnessGuardPrimitives: - """Low-level response validation stays runnable on the CPU-only CI stub.""" - - def test_upsert_short_result_is_raised(self): - store = _StrictMooncakeStoreProxy(_SequenceStore([[0]])) - with pytest.raises(RuntimeError, match="returned 1 results, expected 2"): - store.batch_upsert_from(["k0", "k1"], [1, 2], [8, 8]) - - def test_remove_failure_is_raised(self): - store = _StrictMooncakeStoreProxy(_SequenceStore([[0, -704]])) - with pytest.raises(RuntimeError, match="batch_remove failed"): - store.batch_remove(["k0", "k1"], force=True) - - -@pytest.mark.skipif( - not _REAL_MOONCAKE_CLIENT, - reason="needs real TransferQueue storage submodules; CPU CI uses a single-file transfer_queue stub", -) -class TestMooncakeCorrectnessGuards: - """Integration with real TransferQueue internals; no GPU/master needed.""" - - def test_retry_short_result_is_never_treated_as_success(self, monkeypatch): - monkeypatch.setattr("transfer_queue.storage.clients.mooncake_client.RETRY_DELAY_SECONDS", 0) - store = _StrictMooncakeStoreProxy(_SequenceStore([[-1, -1], [0]])) - client = _client_with_store(store) - with pytest.raises(RuntimeError, match="returned 1 results, expected 2"): - client._batch_upsert_with_retry(["k0", "k1"], [1, 2], [8, 8]) - - @pytest.mark.asyncio - async def test_negative_production_status_ack_is_raised(self, monkeypatch): - from transfer_queue.utils import zmq_utils - - socket = _FakeNotifySocket() - monkeypatch.setattr(zmq_utils, "create_zmq_socket", lambda **kwargs: socket) - monkeypatch.setattr( - zmq_utils.ZMQMessage, - "deserialize", - staticmethod( - lambda messages: SimpleNamespace( - request_type=zmq_utils.ZMQRequestType.NOTIFY_DATA_UPDATE_ACK, - body={"success": False, "partition_id": "p0"}, - ) - ), - ) - manager = SimpleNamespace( - storage_manager_id="guard-test", - zmq_context=object(), - controller_info=SimpleNamespace(ip="redacted", to_addr=lambda name: "inproc://controller"), - ) - with pytest.raises(RuntimeError, match="rejected the production-status update"): - await _strict_notify_and_wait(manager, [b"request"]) - assert socket.closed is True - - @pytest.mark.asyncio - async def test_missing_production_status_ack_is_bounded(self, monkeypatch): - from transfer_queue.storage.managers import base as tq_base - from transfer_queue.utils import zmq_utils - - socket = _FakeNotifySocket() - monkeypatch.setattr(zmq_utils, "create_zmq_socket", lambda **kwargs: socket) - monkeypatch.setattr(tq_base, "TQ_DATA_UPDATE_RESPONSE_TIMEOUT", 0) - manager = SimpleNamespace( - storage_manager_id="guard-test", - zmq_context=object(), - controller_info=SimpleNamespace(ip="redacted", to_addr=lambda name: "inproc://controller"), - ) - with pytest.raises(TimeoutError, match="production-status ACK"): - await _strict_notify_and_wait(manager, [b"request"]) - assert socket.closed is True - - class _InlineExecutorLoop: """Run storage-manager sync calls inline for deterministic async tests. From 830cada93e2f881be09da44df9eaf7a16cba992b Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Fri, 14 Aug 2026 21:44:26 +0800 Subject: [PATCH 12/22] docs: sync rdma guide with review-round changes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 📝 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) --- docs/draft/transfer_queue_rdma.md | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/docs/draft/transfer_queue_rdma.md b/docs/draft/transfer_queue_rdma.md index 2f94e4b19..fe3dc990b 100644 --- a/docs/draft/transfer_queue_rdma.md +++ b/docs/draft/transfer_queue_rdma.md @@ -32,7 +32,8 @@ driver 在**第一次 `tq.init` 之前**完成探测并生成 job 级唯一的 e 1. 校验参数组合（例如 `simple` + `rdma-mode` 会被拒绝） 2. `probe_cluster_nodes()` 通过 Ray 把探测任务绑定到每个**存活且有 GPU** 的节点，并额外探测 driver（driver 也会创建 Mooncake owner client）；各节点读取本机 `/sys`、mooncake 状态，并在 2 秒上限内检查外部 master 的 TCP 可达性；超时或崩溃的节点转为退化结果，不静默丢弃 3. `reduce_results()` 做 AND 归约：整个作业只能跑在最低共同能力上 -4. `required` 模式下若发生任何回退，直接抛异常并打印每个节点的探测明细 +4. Mooncake 生效前，driver 在**每个存活节点**（不限 GPU，因为 Serve replica 与 0-CPU actor 没有 placement 绑定）用真实配置各跑一次**有界 attach 握手**并立即 detach；`auto` 下任一节点失败则统一关闭 Mooncake 状态、全作业收敛到 SimpleStorage，`off`/`required` 下启动失败并列出失败节点 +5. `required` 模式下若发生任何回退，直接抛异常并打印每个节点的探测明细 降级阶梯： @@ -70,7 +71,7 @@ Mooncake → SimpleStorage（任一节点 mooncake/master 不可用、运行时 setsid mooncake_master -rpc_port=50051 -metrics_port=9004 > /var/log/mooncake_master.log 2>&1 < /dev/null & ``` -然后给作业设置 `MC_MASTER_ADDRESS=:50051`。未设置时内部默认 `localhost:50051`，仅适用于单节点开发。 +然后给**每个节点**的作业环境设置 `MC_MASTER_ADDRESS=:50051`。该变量是必填项：未设置时启动直接失败，Relax 不会假定 loopback 端点（多节点作业里每个节点都把自己的 localhost 当 master 会导致误降级或误中止）。 启动前置条件：部署侧必须先启动 master，所有 GPU 节点和 driver 都能解析并连接 `MC_MASTER_ADDRESS`，防火墙允许 master RPC 端口；作业镜像中的 TQ 必须包含本文“正确性依赖”所列修复。Relax 不负责拉起、重启或终止 master。 @@ -80,6 +81,8 @@ setsid mooncake_master -rpc_port=50051 -metrics_port=9004 > /var/log/mooncake_ma |---|---|---| | **master 在探测时不可达** | `auto` 统一降级到 SimpleStorage；`required` 启动失败并列出失败节点 | 先确认 master、DNS/路由和 `MC_MASTER_ADDRESS` | | **master 探测通过、但 `tq.init` 时失败/超时** | 第一次初始化在独立 owner actor 中执行，driver 最多等待 60 秒。失败后回收该 actor 及其拥有的半初始化 controller；`auto` 只重试一次 SimpleStorage，`required` 清理后抛出原始错误 | 查看 `mooncake_init_failed:*`、master 日志和 owner 清理日志 | +| **attach 握手在某节点失败/超时** | driver 汇总各节点结果：`auto` 关闭 Mooncake 状态并统一回退 SimpleStorage（日志 `attach_handshake_failed:*`）；`off`/`required` 启动失败并列出节点 | 检查失败节点到 master 的连通性与 RDMA 状态 | +| **worker attach 卡住**（controller 半初始化 / mooncake setup 挂起） | 每个 worker 的 attach 有统一 deadline（默认 60 s，`RELAX_TQ_ATTACH_TIMEOUT_SECONDS` 可调）：先有界等待 controller 提供配置，再在 watchdog 线程里跑 `tq.init`；超时该 worker 立刻失败而不是无限挂起 | 看 `TqAttachTimeout` 报错中的阶段描述 | | **正常退出**（作业结束或全局重启） | 只有 owner actor 调用全局 `tq.close()`，随后显式 `storage_client.close()` 卸载 segment；附加 worker 只关闭本地 client，不能删除全局数据或 controller。master 本身不动 | 无需操作 | | **异常退出**（worker 被 kill / OOM / 节点掉线） | Python 层不执行，segment 仍在 master 注册。master 要等 `client_ttl`（默认 30 s）才判定客户端过期，期间新作业的 put 会打到死端点并报 `Failed to open segment ... Connection refused` | 等 30 s 后重启，或部署侧调小 `-client_ttl` | @@ -97,7 +100,7 @@ setsid mooncake_master -rpc_port=50051 -metrics_port=9004 > /var/log/mooncake_ma ## 容量不足与正确性依赖 -Mooncake 配置固定 `hard_pin=true`，不会为了腾空间静默驱逐已经生产但尚未消费的数据。启动前还会按 rollout batch、采样数、staleness 和多模态 payload 的固定启发式估算检查 4 GiB client segment；`auto` 预检不足时回退 SimpleStorage，`required` 直接失败。该估算只是早期保护，不能替代运行时错误处理，因为真实图片 patch 数会变化。 +Mooncake 配置固定 `hard_pin=true`，不会为了腾空间静默驱逐已经生产但尚未消费的数据。启动前按 **token 预算推导最坏情况 payload** 检查 client segment（默认 4 GiB）：文本按 `seq_length × 32 B`，多模态样本另加 `seq_length × 784 像素/token × 12 B`（ViT patch 14、merge 2、float32 RGB，例如 8k 序列约 77 MiB/样本），再乘 rollout batch、采样数与 `staleness+1`；`auto` 预检不足时回退 SimpleStorage，`required`/`off` 直接失败。segment 大小可用 `RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB` 调整，容量校验与客户端配置读取同一个值。该上界是保守推导，不能替代运行时错误处理。 运行时依赖固定到 TransferQueue commit `58054a33834aadbcf76aacd6b1e32e25c030f2c9`，并在 Mooncake 启动前检查以下能力： @@ -107,7 +110,7 @@ Mooncake 配置固定 `hard_pin=true`，不会为了腾空间静默驱逐已经 因此，上游曾出现的“返回码未检查导致静默丢数据”不是已知限制，而是 Mooncake 启用的硬门槛：`auto` 在契约不满足时禁用 Mooncake 并回退，`required` 拒绝启动。Docker 镜像固定上述修复 commit，运行时检查用于防止环境被旧包覆盖。 -正确性守卫同时默认设置 `MC_STORE_MEMCPY=0`：mooncake 0.3.10 在 TCP-only 环境会自动启用 memcpy 快拷贝路径，该路径存在静默截断缺陷（现象与处置见排障表）；RDMA 会话本就自动禁用 memcpy，不受影响。运维确认环境安全后可在启动前显式导出 `MC_STORE_MEMCPY=1` 覆盖。 +正确性守卫强制 `MC_STORE_MEMCPY=0` 且 **fail-closed**：mooncake 0.3.10 在 TCP-only 环境会自动启用 memcpy 快拷贝路径，该路径存在已确认的静默截断缺陷（现象与处置见排障表）；RDMA 会话本就自动禁用 memcpy，不受影响。由于缺陷在当前 pin 上已实证，显式导出 `MC_STORE_MEMCPY=1` 会在启动时被直接拒绝，待 pin 升级到修复版本后再按版本重新放开。另外，针对 pin 版本剩余缺口的运行时补丁（逐次重试返回码校验、严格 production-status ACK）已拆分到独立的版本门控分支 `feat/tq-mooncake-loss-guards`（`relax/utils/tq_mooncake_patches.py`），本 PR 只保留只读的能力校验。 真机容量故障注入会故意创建 64 MiB segment 并写入 96 MiB，仅允许在独立、可丢弃的 master 上运行： @@ -179,7 +182,7 @@ PYTHONPATH=. python -u scripts/benchmarks/tq_cross_node_bench.py \ | `setup failed with error code: -1` | master 不可达 | 检查 master 进程与 `MC_MASTER_ADDRESS` | | `Failed to open segment ... Connection refused` | 上一轮客户端异常退出，死 segment 仍在 master 注册 | 等 `client_ttl`（30 s）过期后重试 | | `batch_get_into failed ... error codes [-800, ...]` | 会话内切换协议（0.3.10 上更敏感），或对端不可达 | 每个协议单独进程跑；确认对端存活 | -| Mooncake/TCP 档 get 数据尾部全零，但批量返回码全部成功（逐字节校验 FAIL） | mooncake 0.3.10 memcpy 快拷贝路径缺陷：TCP-only 环境被自动启用后，跨节点 get 会静默截断（坏行自 64 KiB 对齐偏移起全零） | 正确性守卫已默认 `MC_STORE_MEMCPY=0`（见“容量不足与正确性依赖”）；若被显式设为 `1`，改回 `0` | +| Mooncake/TCP 档 get 数据尾部全零，但批量返回码全部成功（逐字节校验 FAIL） | mooncake 0.3.10 memcpy 快拷贝路径缺陷：TCP-only 环境被自动启用后，跨节点 get 会静默截断（坏行自 64 KiB 对齐偏移起全零） | 正确性守卫已强制 `MC_STORE_MEMCPY=0`（见“容量不足与正确性依赖”）；显式设 `1` 会被启动拒绝，unset 即可 | | Mooncake/TCP 档在单机回环下原生 SIGSEGV | 与上一行同源（memcpy 路径），回环下表现为崩溃而非静默截断 | 同上 | | 长时间反复起停会话后 `batch_upsert_from ... error codes [-800, ...]`，重试耗尽（响亮失败，非静默） | master 长期吸收异常退出的客户端后状态劣化，metrics 仍报 serving | 重启 mooncake master；长跑验收前先起新 master | | 多网卡机器跨节点建连失败 | 自动选卡选到了不通的网卡 | 显式 `--tq-rdma-device`；必要时用 `MC_TCP_BIND_ADDRESS` 指定 TCP 侧绑定地址 | From b15ee18b4c8137817d7e27bf52d69af9b66a0b44 Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Fri, 14 Aug 2026 23:14:55 +0800 Subject: [PATCH 13/22] fix(tests): stub attach_tq_client now that SFT no longer imports tq 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. --- tests/components/test_sft.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/tests/components/test_sft.py b/tests/components/test_sft.py index 213a654e6..95dc9a74d 100644 --- a/tests/components/test_sft.py +++ b/tests/components/test_sft.py @@ -94,8 +94,7 @@ async def test_sft_step_pushes_one_batch_to_tq(monkeypatch): fake_client = MagicMock() fake_client.async_put = AsyncMock(return_value=None) - monkeypatch.setattr("relax.components.sft.tq.init", lambda *a, **kw: None, raising=False) - monkeypatch.setattr("relax.components.sft.tq.get_client", lambda: fake_client, raising=False) + monkeypatch.setattr("relax.components.sft.attach_tq_client", lambda *a, **kw: fake_client) args = _make_args(global_batch_size=4) SFTCls = SFT.func_or_class @@ -137,8 +136,7 @@ async def test_sft_step_rejects_empty_or_partial_batch(monkeypatch, returned_cou fake_client = MagicMock() fake_client.async_put = AsyncMock(return_value=None) - monkeypatch.setattr("relax.components.sft.tq.init", lambda *a, **kw: None, raising=False) - monkeypatch.setattr("relax.components.sft.tq.get_client", lambda: fake_client, raising=False) + monkeypatch.setattr("relax.components.sft.attach_tq_client", lambda *a, **kw: fake_client) args = _make_args(global_batch_size=4) SFTCls = SFT.func_or_class @@ -173,8 +171,7 @@ async def test_sft_eval_rejects_source_with_no_valid_samples(monkeypatch): _patch_pipeline_dependencies(monkeypatch) fake_client = MagicMock() fake_client.async_put = AsyncMock(return_value=None) - monkeypatch.setattr("relax.components.sft.tq.init", lambda *a, **kw: None, raising=False) - monkeypatch.setattr("relax.components.sft.tq.get_client", lambda: fake_client, raising=False) + monkeypatch.setattr("relax.components.sft.attach_tq_client", lambda *a, **kw: fake_client) args = _make_args(global_batch_size=4) args.eval_interval = 1 @@ -208,8 +205,7 @@ async def test_sft_loop_advances_step(monkeypatch): fake_client = MagicMock() fake_client.async_put = AsyncMock(return_value=None) fake_client.async_get_partition_list = AsyncMock(return_value=[]) - monkeypatch.setattr("relax.components.sft.tq.init", lambda *a, **kw: None, raising=False) - monkeypatch.setattr("relax.components.sft.tq.get_client", lambda: fake_client, raising=False) + monkeypatch.setattr("relax.components.sft.attach_tq_client", lambda *a, **kw: fake_client) args = _make_args(global_batch_size=2, num_rollout=3) SFTCls = SFT.func_or_class From 66ec9d533f7f40d54a1fe98e1b0389b940b1fe5e Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Sun, 16 Aug 2026 02:07:28 +0800 Subject: [PATCH 14/22] fix(data-plane): address RDMA review findings --- relax/core/controller.py | 7 +--- relax/utils/tq_config.py | 24 +++++++++-- relax/utils/tq_lifecycle.py | 45 ++++++++++++++++++-- scripts/benchmarks/tq_cross_node_bench.py | 1 + scripts/benchmarks/tq_rdma_bench.py | 28 +++++++------ tests/utils/test_rdma_probe.py | 33 +++++++++++++++ tests/utils/test_tq_benchmark_guards.py | 16 ++++++++ tests/utils/test_tq_failure_paths.py | 50 +++++++++++++++++++++++ 8 files changed, 180 insertions(+), 24 deletions(-) create mode 100644 tests/utils/test_tq_benchmark_guards.py diff --git a/relax/core/controller.py b/relax/core/controller.py index 42e4b60ad..523898ad7 100644 --- a/relax/core/controller.py +++ b/relax/core/controller.py @@ -51,6 +51,7 @@ from relax.utils.s3_model_loader import cleanup_s3_model_weights_from_shm from relax.utils.tq_config import ( build_backend_config, + resolve_tq_capacity_batch_size, resolve_mooncake_master_address, validate_mooncake_runtime_contract, ) @@ -271,11 +272,7 @@ def _cleanup_s3_model_weights_after_init(self) -> None: def _initialize_data_system(self): algo_key = resolve_sft_algo_key(self.config) - batch_size_for_capacity = ( - self.config.over_sampling_batch_size - if self.config.partial_rollout and self.config.use_dynamic_global_batch_size - else self.config.rollout_batch_size - ) + batch_size_for_capacity = resolve_tq_capacity_batch_size(self.config) total_storage_size = ( batch_size_for_capacity * (self.config.max_staleness + 1) * self.config.n_samples_per_prompt ) diff --git a/relax/utils/tq_config.py b/relax/utils/tq_config.py index 224e78ef4..2cd646f34 100644 --- a/relax/utils/tq_config.py +++ b/relax/utils/tq_config.py @@ -186,6 +186,23 @@ def build_mooncake_config( _TEXT_BYTES_PER_TOKEN = 32 +def resolve_tq_capacity_batch_size(args: Any) -> int: + """Return the largest batch that one TQ step may need to hold. + + Dynamic partial rollout schedules from the over-sampling pool, so its + capacity contract is ``over_sampling_batch_size`` rather than the smaller + nominal ``rollout_batch_size``. Keep this resolution shared with the + Controller's SimpleStorage sizing so both backends reserve for the same + number of samples. + """ + rollout_batch = getattr(args, "rollout_batch_size") + if getattr(args, "partial_rollout", False) and getattr(args, "use_dynamic_global_batch_size", False): + over_sampling_batch = getattr(args, "over_sampling_batch_size", None) + if over_sampling_batch is not None: + return int(over_sampling_batch) + return int(rollout_batch) + + def estimate_payload_bytes(args: Any) -> int: """Worst-case per-step payload upper bound in bytes. @@ -197,7 +214,7 @@ def estimate_payload_bytes(args: Any) -> int: passed configurations which later failed puts mid-training. """ n_samples = args.n_samples_per_prompt - rollout_batch = args.rollout_batch_size + capacity_batch = resolve_tq_capacity_batch_size(args) seq_length = int(getattr(args, "seq_length", 0) or 0) if seq_length <= 0: raise RuntimeError( @@ -207,7 +224,7 @@ def estimate_payload_bytes(args: Any) -> int: per_sample = seq_length * _TEXT_BYTES_PER_TOKEN if getattr(args, "multimodal_keys", None) is not None: per_sample += seq_length * _PIXELS_PER_VISION_TOKEN * _BYTES_PER_PIXEL_VALUE - return rollout_batch * n_samples * per_sample + return capacity_batch * n_samples * per_sample def validate_segment_capacity(args: Any, effective: EffectiveConfig) -> str | None: @@ -222,6 +239,7 @@ def validate_segment_capacity(args: Any, effective: EffectiveConfig) -> str | No return None max_staleness = getattr(args, "max_staleness", 0) + capacity_batch = resolve_tq_capacity_batch_size(args) payload = estimate_payload_bytes(args) needed = payload * (max_staleness + 1) available = resolve_global_segment_size() @@ -229,7 +247,7 @@ def validate_segment_capacity(args: Any, effective: EffectiveConfig) -> str | No if needed > available: return ( f"MooncakeStore segment capacity insufficient: worst-case in-flight payload " - f"{needed / 1024**3:.1f} GiB (rollout_batch={args.rollout_batch_size} × " + f"{needed / 1024**3:.1f} GiB (effective_batch={capacity_batch} × " f"n_samples={args.n_samples_per_prompt} × staleness+1={max_staleness + 1}) " f"exceeds global_segment_size {available / 1024**3:.1f} GiB. " f"Reduce batch size / max_staleness, or raise RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB." diff --git a/relax/utils/tq_lifecycle.py b/relax/utils/tq_lifecycle.py index ed690e1e4..95bd536be 100644 --- a/relax/utils/tq_lifecycle.py +++ b/relax/utils/tq_lifecycle.py @@ -69,7 +69,7 @@ class TqCleanupTimeout(TimeoutError): class TqConfigurationMismatch(RuntimeError): - """Raised when an existing controller uses a different backend config.""" + """Raised when an existing controller uses an incompatible job config.""" def _get_config_value(config: Any, key: str, default: Any = None) -> Any: @@ -105,6 +105,41 @@ def _backend_signature(conf: Any) -> tuple[Any, ...]: ) +def _sampler_signature(sampler: Any) -> tuple[Any, ...]: + """Return sampler identity and immutable construction-time parameters. + + TransferQueue samplers keep mutable scheduling state in underscore-prefixed + attributes. Those caches legitimately differ between processes and must + not prevent an attach; public attributes describe the sampling contract + that workers and the existing controller must agree on. + """ + if sampler is None: + return (None, ()) + if isinstance(sampler, str): + return ("string", sampler) + if isinstance(sampler, type): + return ("class", f"{sampler.__module__}.{sampler.__qualname__}") + + sampler_type = f"{type(sampler).__module__}.{type(sampler).__qualname__}" + try: + public_config = tuple( + sorted((name, value) for name, value in vars(sampler).items() if not name.startswith("_")) + ) + except TypeError: + public_config = () + return (sampler_type, public_config) + + +def _configuration_signature(conf: Any) -> tuple[Any, ...]: + """Return fields that must agree for workers to share a controller.""" + controller = _get_config_value(conf, "controller", {}) + return ( + _backend_signature(conf), + bool(_get_config_value(controller, "polling_mode", False)), + _sampler_signature(_get_config_value(controller, "sampler")), + ) + + def _backend_description(conf: Any) -> str: """Describe the backend without logging endpoints or host information.""" backend = _get_config_value(conf, "backend", {}) @@ -565,14 +600,15 @@ def _start_owner(conf: Any, *, timeout: float) -> TqInitResult: # A concurrent initializer won the named-actor race. This process is only # attached and must never retain an actor capable of global tq.close(). - config_mismatch = _backend_signature(stored_conf) != _backend_signature(conf) + config_mismatch = _configuration_signature(stored_conf) != _configuration_signature(conf) try: ray.get(owner.detach.remote(), timeout=10.0) finally: _stop_owner_actor(owner) if config_mismatch: raise TqConfigurationMismatch( - "A concurrent TransferQueue initializer won with a different backend config " + "A concurrent TransferQueue initializer won with a different backend config or controller sampling " + "contract " f"(requested={_backend_description(conf)}, stored={_backend_description(stored_conf)}). " "Detached without modifying the winning controller." ) @@ -619,9 +655,10 @@ def _attempt(attempt_conf: Any) -> TqInitResult: # Attach semantics: use the controller's actual config. Upstream # tq.init(conf) returns the caller-provided config even when ignored. stored_conf = _get_stored_config() - if _backend_signature(stored_conf) != _backend_signature(attempt_conf): + if _configuration_signature(stored_conf) != _configuration_signature(attempt_conf): raise TqConfigurationMismatch( "Refusing to attach to an existing TransferQueueController with a different backend config " + "or controller sampling contract " f"(requested={_backend_description(attempt_conf)}, stored={_backend_description(stored_conf)}). " "Only the owner may close the existing controller." ) diff --git a/scripts/benchmarks/tq_cross_node_bench.py b/scripts/benchmarks/tq_cross_node_bench.py index 0c59e23c4..31b39ca97 100644 --- a/scripts/benchmarks/tq_cross_node_bench.py +++ b/scripts/benchmarks/tq_cross_node_bench.py @@ -339,6 +339,7 @@ def wait_actor_gone(name: str = "TransferQueueController", timeout: float = 30.0 except ValueError: return time.sleep(0.4) + raise TimeoutError(f"Ray actor {name!r} is still registered after {timeout:.1f}s") def close_tq_unmount_and_wait() -> None: diff --git a/scripts/benchmarks/tq_rdma_bench.py b/scripts/benchmarks/tq_rdma_bench.py index 8be5243d0..1399054d7 100644 --- a/scripts/benchmarks/tq_rdma_bench.py +++ b/scripts/benchmarks/tq_rdma_bench.py @@ -124,9 +124,22 @@ def build_tq_config(config_name: str, args: argparse.Namespace, num_storage_unit ) +def wait_actor_gone(name: str = "TransferQueueController", timeout: float = 20.0) -> None: + """Wait for a named TQ actor to leave GCS, failing closed on timeout.""" + import ray + + deadline = time.time() + timeout + while time.time() < deadline: + try: + ray.get_actor(name, namespace="transfer_queue") + except ValueError: + return + time.sleep(0.4) + raise TimeoutError(f"Ray actor {name!r} is still registered after {timeout:.1f}s") + + def close_tq_and_wait(timeout: float = 20.0) -> None: - """Close TQ, unmount the Mooncake segment, and wait for the controller to - leave the GCS. + """Close TQ, unmount the Mooncake segment, and confirm controller exit. Required between configs: ``tq.init`` attaches to an existing controller and ignores the new conf (interface.py:152), so without close+wait every @@ -138,9 +151,6 @@ def close_tq_and_wait(timeout: float = 20.0) -> None: dead endpoint ("Failed to open segment ... Connection refused") until the master's ``client_ttl`` (30 s) expires. Unmount explicitly instead. """ - import time - - import ray import transfer_queue as tq store_client = None @@ -154,13 +164,7 @@ def close_tq_and_wait(timeout: float = 20.0) -> None: if store_client is not None and hasattr(store_client, "close"): store_client.close() # unmounts the segment and deregisters from the master - deadline = time.time() + timeout - while time.time() < deadline: - try: - ray.get_actor("TransferQueueController", namespace="transfer_queue") - except ValueError: - return - time.sleep(0.4) + wait_actor_gone(timeout=timeout) # --------------------------------------------------------------------------- # diff --git a/tests/utils/test_rdma_probe.py b/tests/utils/test_rdma_probe.py index 5d80631d7..4d4a16bd6 100644 --- a/tests/utils/test_rdma_probe.py +++ b/tests/utils/test_rdma_probe.py @@ -35,6 +35,7 @@ build_simple_storage_config, estimate_payload_bytes, resolve_mooncake_master_address, + resolve_tq_capacity_batch_size, validate_mooncake_runtime_contract, validate_segment_capacity, ) @@ -571,6 +572,38 @@ def test_estimate_payload_multimodal_is_token_budget_bound(self): assert per_sample == 8192 * (32 + 784 * 12) assert 70 * 1024**2 < per_sample < 80 * 1024**2 + def test_capacity_batch_uses_dynamic_partial_rollout_oversampling(self): + args = _make_args( + rollout_batch_size=16, + partial_rollout=True, + use_dynamic_global_batch_size=True, + over_sampling_batch_size=64, + ) + assert resolve_tq_capacity_batch_size(args) == 64 + assert estimate_payload_bytes(args) == 64 * 8192 * 32 + + def test_capacity_batch_uses_nominal_rollout_without_dynamic_partial_rollout(self): + args = _make_args( + rollout_batch_size=16, + partial_rollout=False, + use_dynamic_global_batch_size=True, + over_sampling_batch_size=64, + ) + assert resolve_tq_capacity_batch_size(args) == 16 + + def test_dynamic_partial_rollout_capacity_rejects_oversampling_peak(self): + args = _make_args( + multimodal_keys=["pixel_values"], + rollout_batch_size=16, + partial_rollout=True, + use_dynamic_global_batch_size=True, + over_sampling_batch_size=64, + ) + eff = EffectiveConfig(backend="MooncakeStore", protocol="rdma", device="", gdr=False, fallback_reason="") + err = validate_segment_capacity(args, eff) + assert err is not None + assert "effective_batch=64" in err + def test_estimate_payload_requires_seq_length(self): args = _make_args(seq_length=None) with pytest.raises(RuntimeError, match="seq_length"): diff --git a/tests/utils/test_tq_benchmark_guards.py b/tests/utils/test_tq_benchmark_guards.py new file mode 100644 index 000000000..208cae58d --- /dev/null +++ b/tests/utils/test_tq_benchmark_guards.py @@ -0,0 +1,16 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Regression tests for fail-closed TransferQueue benchmark teardown.""" + +import pytest + +from scripts.benchmarks import tq_cross_node_bench, tq_rdma_bench + + +@pytest.mark.parametrize( + "wait_actor_gone", + [tq_cross_node_bench.wait_actor_gone, tq_rdma_bench.wait_actor_gone], +) +def test_actor_wait_timeout_is_not_silently_ignored(wait_actor_gone): + with pytest.raises(TimeoutError, match="still registered"): + wait_actor_gone(timeout=0) diff --git a/tests/utils/test_tq_failure_paths.py b/tests/utils/test_tq_failure_paths.py index 6cc30d2dd..ac14de76f 100644 --- a/tests/utils/test_tq_failure_paths.py +++ b/tests/utils/test_tq_failure_paths.py @@ -471,6 +471,46 @@ def test_attach_accepts_matching_mooncake_config(self, monkeypatch): assert result.owns_controller is False assert calls["attempts"] == [] + def test_attach_rejects_different_polling_mode(self, monkeypatch): + requested = self._conf("SimpleStorage") + requested["controller"]["polling_mode"] = True + stored = self._conf("SimpleStorage") + stored["controller"]["polling_mode"] = False + calls = self._patch_transaction(monkeypatch, existed=True, init_effects=[], stored_conf=stored) + with pytest.raises(tq_lifecycle.TqConfigurationMismatch, match="controller sampling contract"): + tq_lifecycle.initialize_tq_with_fallback(requested, mode="off") + assert calls["attempts"] == [] + + def test_attach_rejects_different_sampler_type(self, monkeypatch): + requested = self._conf("SimpleStorage") + requested["controller"]["sampler"] = _SamplerA(n_samples_per_prompt=2) + stored = self._conf("SimpleStorage") + stored["controller"]["sampler"] = _SamplerB(n_samples_per_prompt=2) + calls = self._patch_transaction(monkeypatch, existed=True, init_effects=[], stored_conf=stored) + with pytest.raises(tq_lifecycle.TqConfigurationMismatch, match="controller sampling contract"): + tq_lifecycle.initialize_tq_with_fallback(requested, mode="off") + assert calls["attempts"] == [] + + def test_attach_rejects_different_sampler_public_config(self, monkeypatch): + requested = self._conf("SimpleStorage") + requested["controller"]["sampler"] = _SamplerA(n_samples_per_prompt=2) + stored = self._conf("SimpleStorage") + stored["controller"]["sampler"] = _SamplerA(n_samples_per_prompt=4) + calls = self._patch_transaction(monkeypatch, existed=True, init_effects=[], stored_conf=stored) + with pytest.raises(tq_lifecycle.TqConfigurationMismatch, match="controller sampling contract"): + tq_lifecycle.initialize_tq_with_fallback(requested, mode="off") + assert calls["attempts"] == [] + + def test_attach_ignores_sampler_private_runtime_state(self, monkeypatch): + requested = self._conf("SimpleStorage") + requested["controller"]["sampler"] = _SamplerA(n_samples_per_prompt=2, state={"request": 1}) + stored = self._conf("SimpleStorage") + stored["controller"]["sampler"] = _SamplerA(n_samples_per_prompt=2, state={"stored": 3}) + calls = self._patch_transaction(monkeypatch, existed=True, init_effects=[], stored_conf=stored) + result = tq_lifecycle.initialize_tq_with_fallback(requested, mode="off") + assert result.config is stored + assert calls["attempts"] == [] + def test_auto_cleans_failed_mooncake_then_retries_simple_once(self, monkeypatch): primary = self._conf("MooncakeStore") fallback = self._conf("SimpleStorage") @@ -510,6 +550,16 @@ def test_timeout_auto_retries_only_after_isolated_owner_cleanup(self, monkeypatc assert len(calls["attempts"]) == 2 +class _SamplerA: + def __init__(self, n_samples_per_prompt: int, state: dict | None = None): + self.n_samples_per_prompt = n_samples_per_prompt + self._states = state or {} + + +class _SamplerB(_SamplerA): + pass + + class _RemoteMethod: def __init__(self, value): self.value = value From 151b09735efaef73a62cada9ebb7bd5743113ffa Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Sun, 16 Aug 2026 02:23:51 +0800 Subject: [PATCH 15/22] style(controller): sort TQ config imports --- relax/core/controller.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/relax/core/controller.py b/relax/core/controller.py index 523898ad7..a48d3135e 100644 --- a/relax/core/controller.py +++ b/relax/core/controller.py @@ -51,8 +51,8 @@ from relax.utils.s3_model_loader import cleanup_s3_model_weights_from_shm from relax.utils.tq_config import ( build_backend_config, - resolve_tq_capacity_batch_size, resolve_mooncake_master_address, + resolve_tq_capacity_batch_size, validate_mooncake_runtime_contract, ) from relax.utils.tq_lifecycle import ( From dcbc85ba8ae835b1855b1c4870c76aba82ed8946 Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Sun, 16 Aug 2026 03:50:28 +0800 Subject: [PATCH 16/22] fix(data-plane): guard TQ client generations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 🐛 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 --- docs/draft/transfer_queue_rdma.md | 9 ++-- relax/backends/megatron/actor.py | 12 +++-- relax/components/actor.py | 5 +- relax/components/actor_fwd.py | 5 +- relax/components/advantages.py | 5 +- relax/components/base.py | 8 ++- relax/components/critic.py | 5 +- relax/components/rollout.py | 5 +- relax/components/sft.py | 5 +- relax/core/controller.py | 23 +++++---- relax/distributed/ray/rollout.py | 7 ++- relax/utils/rdma_probe.py | 9 ++-- relax/utils/tq_lifecycle.py | 74 ++++++++++++++++++++++------ tests/utils/test_rdma_probe.py | 24 ++++----- tests/utils/test_tq_failure_paths.py | 33 +++++++++++-- 15 files changed, 167 insertions(+), 62 deletions(-) diff --git a/docs/draft/transfer_queue_rdma.md b/docs/draft/transfer_queue_rdma.md index fe3dc990b..91622363c 100644 --- a/docs/draft/transfer_queue_rdma.md +++ b/docs/draft/transfer_queue_rdma.md @@ -4,7 +4,7 @@ Relax 的数据面（rollout ↔ train 之间的样本传输）默认走 TransferQueue 的 SimpleStorage/ZMQ。本特性把 TransferQueue 已有的 MooncakeStore 后端接出来，使数据面可以走 RDMA，并在能力不足时安全回退。 -首期只做配置接入、能力探测与一致回退，**不改变 payload 形状与数据分发语义**。默认参数下行为与接入前完全一致。 +首期只做配置接入、能力探测与一致回退，**不改变 payload 形状与数据分发语义**。默认参数仍使用 SimpleStorage 及原有 controller 所有权模型；同时所有 worker attach（包括 SimpleStorage）新增默认 60 秒 deadline，半初始化 controller 会被回收，Controller 构造失败时会关闭本进程已经完成的 legacy `tq.init`。 ## 配置入口 @@ -12,7 +12,7 @@ Relax 的数据面（rollout ↔ train 之间的样本传输）默认走 Transfe | 参数 | 取值 | 说明 | |---|---|---| -| `--tq-storage-backend` | `simple`（默认）/ `mooncake` | `simple` 等价于接入前行为 | +| `--tq-storage-backend` | `simple`（默认）/ `mooncake` | `simple` 保留接入前的存储与 controller 所有权语义，并共享新增的有界 attach/失败清理 | | `--tq-rdma-mode` | `off`（默认）/ `auto` / `required` | `off` 即使有硬件也不用 RDMA；`auto` 探测失败自动降级；`required` 探测失败直接报错退出 | | `--tq-rdma-device` | 设备名，如 `mlx5_bond_0`；空为自动 | 多网卡机器上自动选择可能选错，跨节点时建议显式指定 | | `--tq-use-gdr` | 默认关 | **实验性**，见下文 | @@ -35,6 +35,8 @@ driver 在**第一次 `tq.init` 之前**完成探测并生成 job 级唯一的 e 4. Mooncake 生效前，driver 在**每个存活节点**（不限 GPU，因为 Serve replica 与 0-CPU actor 没有 placement 绑定）用真实配置各跑一次**有界 attach 握手**并立即 detach；`auto` 下任一节点失败则统一关闭 Mooncake 状态、全作业收敛到 SimpleStorage，`off`/`required` 下启动失败并列出失败节点 5. `required` 模式下若发生任何回退，直接抛异常并打印每个节点的探测明细 +第 4 步不是轻量探针：每个节点都会创建真实 Mooncake client，并按配置请求挂载/注册完整 client segment（默认 `global_segment_size=4 GiB`，另有默认 1 GiB local buffer），完成后立即释放。具体物理 RSS、锁页与注册方式取决于 Mooncake 实现，但启动阶段会出现节点级瞬时内存/注册资源尖峰；CPU-only head 也在覆盖范围内。节点内存或 `memlock` 不足会表现为 attach 握手失败：`auto` 下整个作业统一回退 SimpleStorage，`off`/`required` 下启动失败。调大 `RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB` 时必须把这份每节点启动资源足迹一并纳入容量规划。 + 降级阶梯： ``` @@ -81,7 +83,7 @@ setsid mooncake_master -rpc_port=50051 -metrics_port=9004 > /var/log/mooncake_ma |---|---|---| | **master 在探测时不可达** | `auto` 统一降级到 SimpleStorage；`required` 启动失败并列出失败节点 | 先确认 master、DNS/路由和 `MC_MASTER_ADDRESS` | | **master 探测通过、但 `tq.init` 时失败/超时** | 第一次初始化在独立 owner actor 中执行，driver 最多等待 60 秒。失败后回收该 actor 及其拥有的半初始化 controller；`auto` 只重试一次 SimpleStorage，`required` 清理后抛出原始错误 | 查看 `mooncake_init_failed:*`、master 日志和 owner 清理日志 | -| **attach 握手在某节点失败/超时** | driver 汇总各节点结果：`auto` 关闭 Mooncake 状态并统一回退 SimpleStorage（日志 `attach_handshake_failed:*`）；`off`/`required` 启动失败并列出节点 | 检查失败节点到 master 的连通性与 RDMA 状态 | +| **attach 握手在某节点失败/超时** | driver 汇总各节点结果：`auto` 关闭 Mooncake 状态并统一回退 SimpleStorage（日志 `attach_handshake_failed:*`）；`off`/`required` 启动失败并列出节点 | 检查失败节点到 master 的连通性、RDMA 状态、可用内存与 `memlock`；握手会瞬时创建完整 client segment，CPU-only head 也会执行 | | **worker attach 卡住**（controller 半初始化 / mooncake setup 挂起） | 每个 worker 的 attach 有统一 deadline（默认 60 s，`RELAX_TQ_ATTACH_TIMEOUT_SECONDS` 可调）：先有界等待 controller 提供配置，再在 watchdog 线程里跑 `tq.init`；超时该 worker 立刻失败而不是无限挂起 | 看 `TqAttachTimeout` 报错中的阶段描述 | | **正常退出**（作业结束或全局重启） | 只有 owner actor 调用全局 `tq.close()`，随后显式 `storage_client.close()` 卸载 segment；附加 worker 只关闭本地 client，不能删除全局数据或 controller。master 本身不动 | 无需操作 | | **异常退出**（worker 被 kill / OOM / 节点掉线） | Python 层不执行，segment 仍在 master 注册。master 要等 `client_ttl`（默认 30 s）才判定客户端过期，期间新作业的 put 会打到死端点并报 `Failed to open segment ... Connection refused` | 等 30 s 后重启，或部署侧调小 `-client_ttl` | @@ -178,6 +180,7 @@ PYTHONPATH=. python -u scripts/benchmarks/tq_cross_node_bench.py \ | 现象 | 可能原因 | 处理 | |---|---|---| | 启动日志 `backend=SimpleStorage fallback=mooncake_unavailable:` | 该节点上 `import mooncake` 失败 | 检查该节点的 `mooncake-transfer-engine` 安装；镜像是否一致 | +| 启动日志 `MooncakeStore capacity fallback to SimpleStorage` | 保守的最坏情况容量上界超过 client segment；多模态按“每个 token 都可能是 vision token”估算，8k 序列约 77 MiB/样本，再乘 batch、采样数与 `staleness+1`，很容易超过默认 4 GiB | 按容量日志核对参数；减少 batch / `n_samples_per_prompt` / `max_staleness`，或调大 `RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB` 并同步规划每节点 attach 的瞬时资源足迹 | | `protocol=tcp fallback=...`，但机器有 RDMA 卡 | 端口非 ACTIVE、GID 取不到、`memlock` 过低，或指定的 `--tq-rdma-device` 在部分节点不存在 | 看 `probe result` 里哪一项 FAIL；`memlock` 需要 unlimited | | `setup failed with error code: -1` | master 不可达 | 检查 master 进程与 `MC_MASTER_ADDRESS` | | `Failed to open segment ... Connection refused` | 上一轮客户端异常退出，死 segment 仍在 master 注册 | 等 `client_ttl`（30 s）过期后重试 | diff --git a/relax/backends/megatron/actor.py b/relax/backends/megatron/actor.py index 4012a4f7e..51c897e66 100644 --- a/relax/backends/megatron/actor.py +++ b/relax/backends/megatron/actor.py @@ -155,10 +155,13 @@ def __del__(self) -> None: # Best-effort detach on graceful teardown; ray.kill / fate-sharing # kills skip destructors, in which case the Mooncake master TTL # reclaims the segment. - if getattr(self, "data_system_client", None) is None: + generation = getattr(self, "_tq_client_generation", None) + if getattr(self, "data_system_client", None) is None or generation is None: return try: - detach_tq_client() + detach_tq_client(generation) + self._tq_client_generation = None + self.data_system_client = None except Exception: # destructor must never raise (interpreter shutdown) return @@ -199,7 +202,10 @@ def _init( if repatch is not None: repatch(args) self.data_system_client = attach_tq_client( - args.tq_config, requested_gdr=getattr(args, "tq_use_gdr", False), role=role + args.tq_config, + requested_gdr=getattr(args, "tq_use_gdr", False), + role=role, + lease_owner=self, ) if is_megatron_main_rank(): init_tracking(args, primary=False) diff --git a/relax/components/actor.py b/relax/components/actor.py index 1b199bd43..d121e262e 100644 --- a/relax/components/actor.py +++ b/relax/components/actor.py @@ -72,7 +72,10 @@ def __init__( self.actor_model = allocate_train_group(args=config, num_gpus=num_gpus, pg=pgs, runtime_env=runtime_env) self.data_system_client = attach_tq_client( - self.config.tq_config, requested_gdr=getattr(self.config, "tq_use_gdr", False), role=self.role + self.config.tq_config, + requested_gdr=getattr(self.config, "tq_use_gdr", False), + role=self.role, + lease_owner=self, ) self.steps = ray.get( diff --git a/relax/components/actor_fwd.py b/relax/components/actor_fwd.py index 7cd068fcb..9f259d9a0 100644 --- a/relax/components/actor_fwd.py +++ b/relax/components/actor_fwd.py @@ -36,7 +36,10 @@ def __init__( self._done_event: Optional[asyncio.Event] = None self._thread_error: Optional[Exception] = None self.data_system_client = attach_tq_client( - self.config.tq_config, requested_gdr=getattr(self.config, "tq_use_gdr", False), role=self.role + self.config.tq_config, + requested_gdr=getattr(self.config, "tq_use_gdr", False), + role=self.role, + lease_owner=self, ) self.actor_model = allocate_train_group(args=config, num_gpus=num_gpus, pg=pgs, runtime_env=runtime_env) ray.get(self.actor_model.async_init(config, role=self.role, with_ref=False)) diff --git a/relax/components/advantages.py b/relax/components/advantages.py index e9f7b8521..7a2a7cc87 100644 --- a/relax/components/advantages.py +++ b/relax/components/advantages.py @@ -40,7 +40,10 @@ def __init__( self.healthy = healthy self.data_system_client = attach_tq_client( - self.config.tq_config, requested_gdr=getattr(self.config, "tq_use_gdr", False), role="advantages" + self.config.tq_config, + requested_gdr=getattr(self.config, "tq_use_gdr", False), + role="advantages", + lease_owner=self, ) self.step = 0 diff --git a/relax/components/base.py b/relax/components/base.py index e93b28a9e..c84088187 100644 --- a/relax/components/base.py +++ b/relax/components/base.py @@ -89,18 +89,22 @@ def __init__(self) -> None: self.step = 0 self._logger_instance = None self._lock = threading.Lock() + self._tq_client_generation: int | None = None def __del__(self) -> None: # Ray Serve calls the destructor on replica shutdown (normal stop, # global restart, in-place restart). Components that attached a # TransferQueue client must detach so a MooncakeStore segment # deregisters before client_ttl instead of leaving a stale endpoint. - if getattr(self, "data_system_client", None) is None: + generation = getattr(self, "_tq_client_generation", None) + if getattr(self, "data_system_client", None) is None or generation is None: return try: from relax.utils.tq_lifecycle import detach_tq_client - detach_tq_client() + detach_tq_client(generation) + self._tq_client_generation = None + self.data_system_client = None except Exception: # destructor must never raise (interpreter shutdown) return diff --git a/relax/components/critic.py b/relax/components/critic.py index c0519e562..b3e2240d5 100644 --- a/relax/components/critic.py +++ b/relax/components/critic.py @@ -41,7 +41,10 @@ def __init__( self.role = role self.data_system_client = attach_tq_client( - self.config.tq_config, requested_gdr=getattr(self.config, "tq_use_gdr", False), role=self.role + self.config.tq_config, + requested_gdr=getattr(self.config, "tq_use_gdr", False), + role=self.role, + lease_owner=self, ) self.critic_model = allocate_train_group( diff --git a/relax/components/rollout.py b/relax/components/rollout.py index b4b3e7352..2da98f0df 100644 --- a/relax/components/rollout.py +++ b/relax/components/rollout.py @@ -334,7 +334,10 @@ def __init__( self.healthy = healthy self.data_system_client = attach_tq_client( - self.config.tq_config, requested_gdr=getattr(self.config, "tq_use_gdr", False), role="rollout" + self.config.tq_config, + requested_gdr=getattr(self.config, "tq_use_gdr", False), + role="rollout", + lease_owner=self, ) self.rollout_manager, self.num_rollout_per_epoch = create_rollout_manager( config, pg, data_source=data_source, runtime_env=runtime_env diff --git a/relax/components/sft.py b/relax/components/sft.py index 96da5b48d..939b39dbf 100644 --- a/relax/components/sft.py +++ b/relax/components/sft.py @@ -82,7 +82,10 @@ def __init__(self, healthy, pgs, num_gpus, config, role, runtime_env=None): # n self.step = getattr(config, "start_rollout_id", 0) self.data_system_client = attach_tq_client( - self.config.tq_config, requested_gdr=getattr(self.config, "tq_use_gdr", False), role=self.role + self.config.tq_config, + requested_gdr=getattr(self.config, "tq_use_gdr", False), + role=self.role, + lease_owner=self, ) self._dataset: Any | None = None diff --git a/relax/core/controller.py b/relax/core/controller.py index a48d3135e..00d63fdd3 100644 --- a/relax/core/controller.py +++ b/relax/core/controller.py @@ -311,12 +311,12 @@ def _initialize_data_system(self): ) if getattr(self.config, "tq_storage_backend", "simple") == "simple": - # Zero-behavior-change default path (review PR#256): identical to - # upstream, the first tq.init runs inside the Controller process - # and no _TransferQueueOwner actor is created. The only addition - # is the F10 reaper, which acts solely on a provably - # half-initialised leftover controller that would otherwise make - # this tq.init poll get_config forever. + # Preserve the upstream SimpleStorage ownership model: the first + # tq.init runs inside Controller and no _TransferQueueOwner actor + # is created. Lifecycle hardening still applies: the F10 reaper + # removes a provably half-initialised controller, worker attaches + # use a 60-second default deadline, and constructor failure closes + # any legacy tq.init completed by this process. reap_unusable_tq_controller() self._tq_owner = None self._tq_legacy_init = True @@ -396,10 +396,11 @@ def _confirm_mooncake_attach(self, init_result: TqInitResult, fallback_config) - def _resolve_tq_backend(self, total_storage_size: int) -> dict: """Resolve the TransferQueue ``backend`` config dict. - Default behavior (``--tq-storage-backend=simple``) is identical to the - previous hardcoded SimpleStorage path. When MooncakeStore is - requested, runs the RDMA capability probe *before* ``tq.init``, applies - graded degradation, and emits the startup log line. + ``--tq-storage-backend=simple`` retains the previous storage and + ownership semantics while sharing the bounded worker-attach and + failure-cleanup hardening. When MooncakeStore is requested, this runs + the RDMA capability probe *before* ``tq.init``, applies graded + degradation, and emits the startup log line. """ # 1. Validate flag combinations (structural, before any probe). # getattr defaults keep old checkpoints / non-argparse configs safe. @@ -445,8 +446,6 @@ def _resolve_tq_backend(self, total_storage_size: int) -> dict: ) master_address = resolve_mooncake_master_address() probe_results = probe_cluster_nodes(device, master_address, probe_rdma=mode != "off") - for r in probe_results: - logger.debug(r.summary()) effective = reduce_results( probe_results, diff --git a/relax/distributed/ray/rollout.py b/relax/distributed/ray/rollout.py index 2b274b06a..95dddba72 100644 --- a/relax/distributed/ray/rollout.py +++ b/relax/distributed/ray/rollout.py @@ -817,6 +817,7 @@ def __init__(self, args, pg, data_source=None): self.args.tq_config, requested_gdr=getattr(self.args, "tq_use_gdr", False), role="rollout_worker", + lease_owner=self, ) logger.info(f"import {self.args.rollout_function_path} as generate_rollout function.") @@ -922,7 +923,11 @@ def dispose(self): self._shutdown_all_engines() # Deregister this worker's Mooncake segment before the actor dies so a # fast restart does not hit stale endpoints until client_ttl expires. - detach_tq_client() + generation = getattr(self, "_tq_client_generation", None) + if generation is not None: + detach_tq_client(generation) + self._tq_client_generation = None + self.data_system_client = None def _shutdown_all_engines(self, timeout: float = 15.0): """Shut down all SGLang engine actors and their child processes. diff --git a/relax/utils/rdma_probe.py b/relax/utils/rdma_probe.py index ddd274fe5..f446cfc20 100644 --- a/relax/utils/rdma_probe.py +++ b/relax/utils/rdma_probe.py @@ -561,7 +561,6 @@ def reduce_results( # AND reduction: the job can only run at the lowest common capability. any_no_mooncake = any(r.effective_protocol is None for r in results) all_rdma = all(r.effective_protocol == "rdma" for r in results) - all_gdr = all(r.gdr_eligible for r in results) if any_no_mooncake: failed_nodes = [r.node for r in results if r.effective_protocol is None] @@ -591,7 +590,7 @@ def reduce_results( if all_rdma: # Device: if any node lacks the requested device, fall back to tcp. if requested_device: - device_ok = all(r.effective_device == requested_device or not r.effective_device for r in results) + device_ok = all(r.effective_device == requested_device for r in results) if not device_ok: return EffectiveConfig( backend="MooncakeStore", @@ -604,8 +603,10 @@ def reduce_results( backend="MooncakeStore", protocol="rdma", device=requested_device, - gdr=use_gdr and all_gdr, - fallback_reason="" if (not use_gdr or all_gdr) else "gdr_cuda_not_initialized", + # probe_node defines GDR eligibility as RDMA transport readiness; + # CUDA staging is deliberately decided and logged by each worker. + gdr=use_gdr, + fallback_reason="", ) # Some nodes can't do RDMA → degrade to TCP (still MooncakeStore). diff --git a/relax/utils/tq_lifecycle.py b/relax/utils/tq_lifecycle.py index 95bd536be..419523769 100644 --- a/relax/utils/tq_lifecycle.py +++ b/relax/utils/tq_lifecycle.py @@ -42,6 +42,13 @@ DEFAULT_TQ_INIT_TIMEOUT_SECONDS = 60.0 DEFAULT_TQ_ATTACH_TIMEOUT_SECONDS = 60.0 +# TransferQueue stores its client in process-global module state. A generation +# identifies the most recent component that claimed that client so a delayed +# destructor from an in-place reload cannot close its successor's connection. +_TQ_CLIENT_LEASE_LOCK = threading.RLock() +_TQ_CLIENT_GENERATION = 0 +_CURRENT_TQ_CLIENT_GENERATION: int | None = None + @dataclass(frozen=True) class TqInitResult: @@ -401,7 +408,14 @@ def _run() -> None: raise error[0] -def attach_tq_client(conf: Any, *, requested_gdr: bool, role: str, timeout: float | None = None) -> Any: +def attach_tq_client( + conf: Any, + *, + requested_gdr: bool, + role: str, + timeout: float | None = None, + lease_owner: Any | None = None, +) -> Any: """Attach a component process within a bounded deadline and report its local experimental GDR state. @@ -409,19 +423,32 @@ def attach_tq_client(conf: Any, *, requested_gdr: bool, role: str, timeout: floa ``tq.init`` itself, because either phase can hang unboundedly (get_config polling and mooncake endpoint setup respectively). ``None`` resolves the deadline from ``RELAX_TQ_ATTACH_TIMEOUT_SECONDS`` (default 60 s). + + When ``lease_owner`` is provided, its private generation token is updated + after a successful attach. Teardown must pass that token to + :func:`detach_tq_client`; stale owners then leave a newer process-global + client untouched. """ - if timeout is None: - timeout = _resolve_attach_timeout() - deadline = time.monotonic() + timeout - _prepare_mooncake_runtime(conf) - _await_controller_config(deadline) - _bounded_tq_init(conf, deadline, role=role) - client = tq.get_client() - log_tq_gdr_runtime_status(requested=requested_gdr, role=role) + global _CURRENT_TQ_CLIENT_GENERATION, _TQ_CLIENT_GENERATION + + with _TQ_CLIENT_LEASE_LOCK: + if timeout is None: + timeout = _resolve_attach_timeout() + deadline = time.monotonic() + timeout + _prepare_mooncake_runtime(conf) + _await_controller_config(deadline) + _bounded_tq_init(conf, deadline, role=role) + client = tq.get_client() + log_tq_gdr_runtime_status(requested=requested_gdr, role=role) + + _TQ_CLIENT_GENERATION += 1 + _CURRENT_TQ_CLIENT_GENERATION = _TQ_CLIENT_GENERATION + if lease_owner is not None: + lease_owner._tq_client_generation = _CURRENT_TQ_CLIENT_GENERATION return client -def detach_tq_client() -> None: +def detach_tq_client(generation: int | None = None) -> None: """Detach this worker's TQ client (attach-only inverse of :func:`attach_tq_client`). @@ -429,11 +456,26 @@ def detach_tq_client() -> None: from the master immediately instead of lingering until ``client_ttl`` expires — a stale endpoint breaks fast restarts with "Failed to open segment". Only process-local handles are touched (never the named - controller or globally stored data), so every worker teardown hook may - call this unconditionally; force-killed workers still fall back to the - master-side TTL. + controller or globally stored data). A component teardown passes the + generation recorded by :func:`attach_tq_client`; a stale generation means + another component has since claimed the process-global client and is left + untouched. ``None`` remains an unconditional detach for short-lived + attach handshakes and owner cleanup. Force-killed workers still fall back + to the master-side TTL. """ - _close_local_tq_client() + global _CURRENT_TQ_CLIENT_GENERATION + + with _TQ_CLIENT_LEASE_LOCK: + if generation is not None and generation != _CURRENT_TQ_CLIENT_GENERATION: + logger.debug( + "[dataplane] Skipping stale TQ detach: " + f"owner_generation={generation} current_generation={_CURRENT_TQ_CLIENT_GENERATION}" + ) + return + try: + _close_local_tq_client() + finally: + _CURRENT_TQ_CLIENT_GENERATION = None def _alive_node_ids() -> list[str]: @@ -501,7 +543,7 @@ def close_tq_and_unmount(*, is_owner: bool) -> None: """ if not is_owner: logger.info("[dataplane] Detaching local TQ client; global controller is owned by another process.") - _close_local_tq_client() + detach_tq_client() return store_client = None @@ -539,7 +581,7 @@ def close(self) -> None: close_tq_and_unmount(is_owner=self._owns_controller) def detach(self) -> None: - _close_local_tq_client() + detach_tq_client() def _stop_owner_actor(owner: Any) -> None: diff --git a/tests/utils/test_rdma_probe.py b/tests/utils/test_rdma_probe.py index 4d4a16bd6..1763b47f0 100644 --- a/tests/utils/test_rdma_probe.py +++ b/tests/utils/test_rdma_probe.py @@ -59,7 +59,6 @@ def _has_real_tq_storage() -> bool: def _make_probe( protocol: str | None = "rdma", device: str = "rdma0", - gdr: bool = False, node: str = "node-A", ) -> ProbeResult: return ProbeResult( @@ -67,7 +66,7 @@ def _make_probe( checks=(CheckResult("mooncake_import", True),), effective_protocol=protocol, effective_device=device, - gdr_eligible=gdr, + gdr_eligible=protocol == "rdma", ) @@ -176,24 +175,25 @@ def test_one_node_no_rdma_degrades_to_tcp(self): assert eff.protocol == "tcp" assert "node-B" in eff.fallback_reason - def test_gdr_eligible_only_when_all_nodes(self): + def test_gdr_request_is_forwarded_for_rdma(self): eff = reduce_results( - [_make_probe(gdr=True), _make_probe(gdr=False, node="node-B")], + [_make_probe(), _make_probe(node="node-B")], requested_backend="mooncake", requested_device="", use_gdr=True, ) - assert eff.gdr is False - assert "gdr_cuda_not_initialized" in eff.fallback_reason + assert eff.gdr is True + assert eff.fallback_reason == "" - def test_gdr_eligible_all_nodes(self): + def test_requested_device_must_match_every_rdma_node(self): eff = reduce_results( - [_make_probe(gdr=True), _make_probe(gdr=True, node="node-B")], + [_make_probe(device="rdma0"), _make_probe(device="rdma1", node="node-B")], requested_backend="mooncake", - requested_device="", - use_gdr=True, + requested_device="rdma0", + use_gdr=False, ) - assert eff.gdr is True + assert eff.protocol == "tcp" + assert eff.fallback_reason == "device_mismatch:rdma0" def test_empty_results_falls_back(self): eff = reduce_results( @@ -244,6 +244,7 @@ def test_no_infiniband_dir_gives_tcp(self): result = probe_node("") # mooncake importable but no RDMA device → tcp assert result.effective_protocol == "tcp" + assert result.gdr_eligible is False def test_active_rdma_device_gives_rdma(self): """When all checks pass, protocol should be rdma.""" @@ -272,6 +273,7 @@ def fake_listdir(path): assert result.effective_protocol == "rdma" assert result.ok assert result.effective_device == "rdma0" + assert result.gdr_eligible is True @staticmethod def _multi_hca_open(active_device: str): diff --git a/tests/utils/test_tq_failure_paths.py b/tests/utils/test_tq_failure_paths.py index ac14de76f..de4b79791 100644 --- a/tests/utils/test_tq_failure_paths.py +++ b/tests/utils/test_tq_failure_paths.py @@ -325,22 +325,47 @@ def test_detach_delegates_to_local_close(self, monkeypatch): tq_lifecycle.detach_tq_client() assert calls == [True] + def test_stale_generation_does_not_close_successor(self, monkeypatch): + client = object() + monkeypatch.setattr(tq_lifecycle, "_TQ_CLIENT_GENERATION", 0) + monkeypatch.setattr(tq_lifecycle, "_CURRENT_TQ_CLIENT_GENERATION", None) + monkeypatch.setattr(tq_lifecycle, "_prepare_mooncake_runtime", lambda conf: None) + monkeypatch.setattr(tq_lifecycle, "_await_controller_config", lambda deadline: None) + monkeypatch.setattr(tq_lifecycle, "_bounded_tq_init", lambda conf, deadline, role: None) + monkeypatch.setattr(tq_lifecycle.tq, "get_client", lambda: client) + monkeypatch.setattr(tq_lifecycle, "log_tq_gdr_runtime_status", lambda **kwargs: "not_requested") + + old_owner = SimpleNamespace() + new_owner = SimpleNamespace() + assert tq_lifecycle.attach_tq_client({}, requested_gdr=False, role="old", lease_owner=old_owner) is client + assert tq_lifecycle.attach_tq_client({}, requested_gdr=False, role="new", lease_owner=new_owner) is client + assert new_owner._tq_client_generation > old_owner._tq_client_generation + + calls = [] + monkeypatch.setattr(tq_lifecycle, "_close_local_tq_client", lambda: calls.append(True)) + tq_lifecycle.detach_tq_client(old_owner._tq_client_generation) + assert calls == [] + tq_lifecycle.detach_tq_client(new_owner._tq_client_generation) + assert calls == [True] + def test_component_del_detaches_attached_client(self, monkeypatch): from relax.components.base import Base calls = [] - monkeypatch.setattr(tq_lifecycle, "detach_tq_client", lambda: calls.append(True)) + monkeypatch.setattr(tq_lifecycle, "detach_tq_client", lambda generation: calls.append(generation)) component = Base() component.data_system_client = object() + component._tq_client_generation = 7 component.__del__() - assert calls == [True] - component.data_system_client = None # keep GC-time __del__ a no-op + assert calls == [7] + assert component.data_system_client is None + assert component._tq_client_generation is None def test_component_del_without_client_is_noop(self, monkeypatch): from relax.components.base import Base calls = [] - monkeypatch.setattr(tq_lifecycle, "detach_tq_client", lambda: calls.append(True)) + monkeypatch.setattr(tq_lifecycle, "detach_tq_client", lambda generation: calls.append(generation)) component = Base() component.__del__() assert calls == [] From cfe6153014a554e28c5e40a9002d054f44fee7b3 Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Tue, 18 Aug 2026 01:15:59 +0800 Subject: [PATCH 17/22] fix(data-plane): isolate TQ attach workers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 🐛 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 --- docs/draft/transfer_queue_rdma.md | 4 +++- relax/utils/tq_lifecycle.py | 16 +++++++++------- tests/utils/test_tq_failure_paths.py | 15 +++++++++++++++ 3 files changed, 27 insertions(+), 8 deletions(-) diff --git a/docs/draft/transfer_queue_rdma.md b/docs/draft/transfer_queue_rdma.md index 91622363c..5f1400e2d 100644 --- a/docs/draft/transfer_queue_rdma.md +++ b/docs/draft/transfer_queue_rdma.md @@ -75,6 +75,8 @@ setsid mooncake_master -rpc_port=50051 -metrics_port=9004 > /var/log/mooncake_ma 然后给**每个节点**的作业环境设置 `MC_MASTER_ADDRESS=:50051`。该变量是必填项：未设置时启动直接失败，Relax 不会假定 loopback 端点（多节点作业里每个节点都把自己的 localhost 当 master 会导致误降级或误中止）。 +Mooncake/TCP 长会话还应在 driver 与所有 Ray worker 中统一设置 `MC_TCP_ENABLE_CONNECTION_POOL=1`。未启用连接池时，大批量反复传输可能耗尽临时 TCP 端口并报 `Cannot assign requested address`；该变量必须通过作业运行时环境传播到所有节点，不能只在提交命令所在的 shell 中设置。RDMA transport 不依赖此选项。 + 启动前置条件：部署侧必须先启动 master，所有 GPU 节点和 driver 都能解析并连接 `MC_MASTER_ADDRESS`，防火墙允许 master RPC 端口；作业镜像中的 TQ 必须包含本文“正确性依赖”所列修复。Relax 不负责拉起、重启或终止 master。 三种情形下的行为： @@ -83,7 +85,7 @@ setsid mooncake_master -rpc_port=50051 -metrics_port=9004 > /var/log/mooncake_ma |---|---|---| | **master 在探测时不可达** | `auto` 统一降级到 SimpleStorage；`required` 启动失败并列出失败节点 | 先确认 master、DNS/路由和 `MC_MASTER_ADDRESS` | | **master 探测通过、但 `tq.init` 时失败/超时** | 第一次初始化在独立 owner actor 中执行，driver 最多等待 60 秒。失败后回收该 actor 及其拥有的半初始化 controller；`auto` 只重试一次 SimpleStorage，`required` 清理后抛出原始错误 | 查看 `mooncake_init_failed:*`、master 日志和 owner 清理日志 | -| **attach 握手在某节点失败/超时** | driver 汇总各节点结果：`auto` 关闭 Mooncake 状态并统一回退 SimpleStorage（日志 `attach_handshake_failed:*`）；`off`/`required` 启动失败并列出节点 | 检查失败节点到 master 的连通性、RDMA 状态、可用内存与 `memlock`；握手会瞬时创建完整 client segment，CPU-only head 也会执行 | +| **attach 握手在某节点失败/超时** | driver 汇总各节点结果：`auto` 关闭 Mooncake 状态并统一回退 SimpleStorage（日志 `attach_handshake_failed:*`）；`off`/`required` 启动失败并列出节点。握手使用一次性 Ray worker，超时的 `tq.init` watchdog thread 不会污染后续任务 | 检查失败节点到 master 的连通性、RDMA 状态、可用内存与 `memlock`；握手会瞬时创建完整 client segment，CPU-only head 也会执行 | | **worker attach 卡住**（controller 半初始化 / mooncake setup 挂起） | 每个 worker 的 attach 有统一 deadline（默认 60 s，`RELAX_TQ_ATTACH_TIMEOUT_SECONDS` 可调）：先有界等待 controller 提供配置，再在 watchdog 线程里跑 `tq.init`；超时该 worker 立刻失败而不是无限挂起 | 看 `TqAttachTimeout` 报错中的阶段描述 | | **正常退出**（作业结束或全局重启） | 只有 owner actor 调用全局 `tq.close()`，随后显式 `storage_client.close()` 卸载 segment；附加 worker 只关闭本地 client，不能删除全局数据或 controller。master 本身不动 | 无需操作 | | **异常退出**（worker 被 kill / OOM / 节点掉线） | Python 层不执行，segment 仍在 master 注册。master 要等 `client_ttl`（默认 30 s）才判定客户端过期，期间新作业的 put 会打到死端点并报 `Failed to open segment ... Connection refused` | 等 30 s 后重启，或部署侧调小 `-client_ttl` | diff --git a/relax/utils/tq_lifecycle.py b/relax/utils/tq_lifecycle.py index 419523769..131777486 100644 --- a/relax/utils/tq_lifecycle.py +++ b/relax/utils/tq_lifecycle.py @@ -488,18 +488,20 @@ def verify_cluster_attach(conf: Any, *, timeout: float | None = None) -> list[st """Bounded attach handshake from every alive node; returns failure summaries. - Each task performs the same bounded :func:`attach_tq_client` a worker would - perform (real stored config, real storage client) and detaches immediately, - so it validates the *actual endpoints* instead of a ``/sys`` capability - heuristic on a node the scheduler may never use. An empty return means - every alive node attached within the deadline; the driver aggregates - failures and decides one job-level outcome. + Each one-shot task performs the same bounded :func:`attach_tq_client` a + worker would perform (real stored config, real storage client) and detaches + immediately, so it validates the *actual endpoints* instead of a ``/sys`` + capability heuristic on a node the scheduler may never use. The worker is + not reused because a timed-out ``tq.init`` daemon thread cannot be stopped + safely in-process. An empty return means every alive node attached within + the deadline; the driver aggregates failures and decides one job-level + outcome. """ if timeout is None: timeout = _resolve_attach_timeout() from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy - @ray.remote(num_cpus=0, max_retries=0) + @ray.remote(num_cpus=0, max_retries=0, max_calls=1) def _handshake(handshake_conf: Any) -> None: from relax.utils.tq_lifecycle import attach_tq_client, detach_tq_client diff --git a/tests/utils/test_tq_failure_paths.py b/tests/utils/test_tq_failure_paths.py index de4b79791..87ced0d6c 100644 --- a/tests/utils/test_tq_failure_paths.py +++ b/tests/utils/test_tq_failure_paths.py @@ -310,6 +310,21 @@ def no_actor(name, namespace=None): with pytest.raises(tq_lifecycle.TqAttachTimeout, match="attach timed out"): tq_lifecycle._await_controller_config(time.monotonic() + 0.3) + def test_cluster_attach_handshake_worker_is_one_shot(self, monkeypatch): + remote_options = {} + + def record_remote_options(**options): + remote_options.update(options) + return lambda function: function + + monkeypatch.setattr(tq_lifecycle.ray, "remote", record_remote_options) + monkeypatch.setattr(tq_lifecycle, "_alive_node_ids", lambda: []) + monkeypatch.setattr(tq_lifecycle.ray, "wait", lambda *args, **kwargs: ([], [])) + + assert tq_lifecycle.verify_cluster_attach({}, timeout=0.1) == [] + assert remote_options["max_calls"] == 1 + assert remote_options["max_retries"] == 0 + # --------------------------------------------------------------------------- # Worker detach (attach-only inverse used by every worker teardown hook) From 10ec8dec3693e3fc07749c0cabf660c12efe467d Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Tue, 18 Aug 2026 03:29:15 +0800 Subject: [PATCH 18/22] test(data-plane): harden TQ lifecycle tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # ✅ 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 --- tests/utils/_tq_handshake_timeout_probe.py | 150 ++++++++++++++ tests/utils/test_tq_failure_paths.py | 215 +++++++++++++++------ 2 files changed, 304 insertions(+), 61 deletions(-) create mode 100644 tests/utils/_tq_handshake_timeout_probe.py diff --git a/tests/utils/_tq_handshake_timeout_probe.py b/tests/utils/_tq_handshake_timeout_probe.py new file mode 100644 index 000000000..331b05b86 --- /dev/null +++ b/tests/utils/_tq_handshake_timeout_probe.py @@ -0,0 +1,150 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Subprocess probe for one-shot Ray worker cleanup after attach timeout.""" + +from __future__ import annotations + +import os +import sys +import time +from pathlib import Path + + +def _process_is_running(pid: int, create_time: float) -> bool: + import psutil + + try: + process = psutil.Process(pid) + return abs(process.create_time() - create_time) < 1e-3 and process.status() != psutil.STATUS_ZOMBIE + except psutil.NoSuchProcess: + return False + + +def _write_transfer_queue_stub(stub_dir: Path) -> None: + """Create a TQ stub whose delayed init visibly mutates module state.""" + stub_dir.mkdir(parents=True, exist_ok=True) + (stub_dir / "transfer_queue.py").write_text( + """\ +import os +import time +from pathlib import Path + +import psutil + +MUTATED = False + + +def init(conf=None): + process = psutil.Process() + Path(os.environ["RELAX_TEST_TQ_STARTED_MARKER"]).write_text( + f"{process.pid},{process.create_time()}", encoding="utf-8" + ) + time.sleep(float(os.environ["RELAX_TEST_TQ_LATE_MUTATION_DELAY"])) + global MUTATED + MUTATED = True + Path(os.environ["RELAX_TEST_TQ_LATE_MARKER"]).write_text("dirty", encoding="utf-8") +""", + encoding="utf-8", + ) + + +def main(probe_dir: Path) -> None: + # Ray reads this switch at import time. Keep the probe isolated from both + # uv parent-process discovery and any cluster selected by the caller. + os.environ.setdefault("RAY_ENABLE_UV_RUN_RUNTIME_ENV", "0") + os.environ.pop("RAY_ADDRESS", None) + + probe_dir.mkdir(parents=True, exist_ok=True) + stub_dir = probe_dir / "stub" + started_path = probe_dir / "timed-out-worker.started" + mutation_path = probe_dir / "late-global-mutation" + late_mutation_delay = 2.0 + _write_transfer_queue_stub(stub_dir) + + original_pythonpath = os.environ.get("PYTHONPATH", "") + worker_pythonpath = os.pathsep.join(path for path in (str(stub_dir), original_pythonpath) if path) + runtime_env = { + "env_vars": { + "PYTHONPATH": worker_pythonpath, + "RELAX_TQ_ATTACH_TIMEOUT_SECONDS": "0.3", + "RELAX_TEST_TQ_STARTED_MARKER": str(started_path), + "RELAX_TEST_TQ_LATE_MARKER": str(mutation_path), + "RELAX_TEST_TQ_LATE_MUTATION_DELAY": str(late_mutation_delay), + } + } + + import ray + + assert not ray.is_initialized() + ray.init( + address="local", + num_cpus=1, + include_dashboard=False, + logging_level="ERROR", + runtime_env=runtime_env, + _temp_dir=str(probe_dir / "ray"), + ) + try: + from relax.utils import tq_lifecycle + + conf = {"backend": {"storage_backend": "SimpleStorage"}, "controller": {}} + + @ray.remote(num_cpus=0) + class _Controller: + def __init__(self, config): + self.config = config + + def get_config(self): + return self.config + + controller = _Controller.options( + name=tq_lifecycle.CONTROLLER_NAME, + namespace=tq_lifecycle.CONTROLLER_NAMESPACE, + ).remote(conf) + assert ray.get(controller.get_config.remote()) == conf + + failures = tq_lifecycle.verify_cluster_attach(conf, timeout=2.0) + assert len(failures) == 1 + assert "did not finish" in failures[0] + assert started_path.exists(), "the production handshake never entered tq.init" + timed_out_pid_text, timed_out_create_time_text = started_path.read_text(encoding="utf-8").split(",") + timed_out_identity = (int(timed_out_pid_text), float(timed_out_create_time_text)) + + # Without max_calls=1, the reusable task worker survives and its daemon + # eventually dirties both the module global and this external marker. + time.sleep(late_mutation_delay + 0.2) + assert not mutation_path.exists(), "timed-out tq.init continued mutating state after task failure" + + exit_deadline = time.monotonic() + 10.0 + while _process_is_running(*timed_out_identity) and time.monotonic() < exit_deadline: + time.sleep(0.05) + assert not _process_is_running(*timed_out_identity), "one-shot handshake worker did not exit after timeout" + + @ray.remote(num_cpus=0, max_retries=0) + def _clean_worker_state() -> tuple[int, float, bool, bool]: + import os + from pathlib import Path + + import psutil + import transfer_queue + + process = psutil.Process() + return ( + os.getpid(), + process.create_time(), + transfer_queue.MUTATED, + Path(os.environ["RELAX_TEST_TQ_LATE_MARKER"]).exists(), + ) + + successor_pid, successor_create_time, successor_mutated, late_marker_exists = ray.get( + _clean_worker_state.remote() + ) + assert (successor_pid, successor_create_time) != timed_out_identity + assert successor_mutated is False + assert late_marker_exists is False + finally: + ray.shutdown() + + +if __name__ == "__main__": + main(Path(sys.argv[1])) diff --git a/tests/utils/test_tq_failure_paths.py b/tests/utils/test_tq_failure_paths.py index 87ced0d6c..8cdd713a9 100644 --- a/tests/utils/test_tq_failure_paths.py +++ b/tests/utils/test_tq_failure_paths.py @@ -26,8 +26,14 @@ import importlib.util import multiprocessing import os +import queue import socket +import subprocess +import sys +import tempfile +import uuid from types import SimpleNamespace +from typing import Any, Callable from unittest.mock import AsyncMock, MagicMock import pytest @@ -128,6 +134,11 @@ def _mm_slow_path_worker(result_queue, protocol: str) -> None: ("ok", ...), ("mismatch", ...) for data corruption, or ("error", ...) for engine failures — the parent treats anything but "ok" as a hard failure. """ + client = None + keys: list[str] = [] + put_meta: list[dict | None] | None = None + status = "error" + detail = "roundtrip did not run" try: from relax.utils.payload_digest import diff_digests, leaf_digests from tests.utils.mm_payload_fixtures import mm_train_data @@ -135,35 +146,118 @@ def _mm_slow_path_worker(result_queue, protocol: str) -> None: train_data, source = mm_train_data(4) samples = train_data["multimodal_train_inputs"] client = TestMooncakeByteExact._client(protocol) - try: - keys = [f"mmslow_{protocol}_{index}@multimodal_train_inputs" for index in range(len(samples))] - put_meta = client.put(keys, samples) - if not all(isinstance(meta, dict) and meta.get("packed_size") for meta in put_meta): - result_queue.put( - ("error", f"[{source}] dict payloads did not take the non-tensor path, put meta: {put_meta}") - ) - return - got = client.get( - keys, - shapes=[[] for _ in keys], - dtypes=[None] * len(keys), - custom_backend_meta=put_meta, + run_token = uuid.uuid4().hex + keys = [f"mmslow_{run_token}_{protocol}_{index}@multimodal_train_inputs" for index in range(len(samples))] + put_meta = client.put(keys, samples) + if not all(isinstance(meta, dict) and meta.get("packed_size") for meta in put_meta): + raise RuntimeError(f"[{source}] dict payloads did not take the non-tensor path, put meta: {put_meta}") + got = client.get( + keys, + shapes=[[] for _ in keys], + dtypes=[None] * len(keys), + custom_backend_meta=put_meta, + ) + problems: list[str] = [] + for index, (want, have) in enumerate(zip(samples, got, strict=True)): + problems += [f"sample {index}: {line}" for line in diff_digests(leaf_digests(want), leaf_digests(have))] + if problems: + status, detail = "mismatch", f"[{source}] {problems[:6]}" + else: + leaves = sum(len(leaf_digests(sample)) for sample in samples) + packed = sum(meta["packed_size"] for meta in put_meta) + status, detail = ( + "ok", + f"[{source}] {len(samples)} samples, {leaves} leaves, {packed} packed bytes", ) - problems: list[str] = [] - for index, (want, have) in enumerate(zip(samples, got, strict=True)): - problems += [ - f"sample {index}: {line}" for line in diff_digests(leaf_digests(want), leaf_digests(have)) - ] - if problems: - result_queue.put(("mismatch", f"[{source}] {problems[:6]}")) - else: - leaves = sum(len(leaf_digests(sample)) for sample in samples) - packed = sum(meta["packed_size"] for meta in put_meta) - result_queue.put(("ok", f"[{source}] {len(samples)} samples, {leaves} leaves, {packed} packed bytes")) - finally: - client.close() except BaseException as error: # pragma: no cover - transport/env failures - result_queue.put(("error", f"{type(error).__name__}: {error}")) + status, detail = "error", f"{type(error).__name__}: {error}" + finally: + if client is not None: + try: + if keys: + client.clear(keys, put_meta) + client.close() + except BaseException as error: # pragma: no cover - native cleanup failures + status, detail = "error", f"cleanup {type(error).__name__}: {error}" + result_queue.put((status, detail)) + + +def _dense_roundtrip_worker(result_queue, protocol: str) -> None: + """Child-process target: dense tensor roundtrip on one pristine session.""" + client = None + keys: list[str] = [] + status = "error" + detail = "roundtrip did not run" + try: + tensors = { + # Production multimodal field names, dimensions, and mixed dtypes. + "pixel_values": torch.randn(64, 1176, dtype=torch.float32).to(torch.bfloat16), + "image_grid_thw": torch.tensor([[1, 8, 8]], dtype=torch.int64), + "input_ids": torch.arange(4096, dtype=torch.int64), + "attention_mask": torch.ones(4096, dtype=torch.int64), + "rewards": torch.linspace(-1, 1, 64, dtype=torch.float32), + "noncontig": torch.randn(128, 256).t(), # transposed == non-contiguous + } + client = TestMooncakeByteExact._client(protocol) + run_token = uuid.uuid4().hex + keys = [f"bx_{run_token}_{protocol}_{name}" for name in tensors] + values = list(tensors.values()) + client.put(keys, values) + got = client.get( + keys, + shapes=[tuple(value.shape) for value in values], + dtypes=[value.dtype for value in values], + ) + problems = [ + name + for name, want, have in zip(tensors, values, got, strict=True) + if have is None or not torch.equal(have, want.contiguous()) + ] + if problems: + status, detail = "mismatch", f"non-byte-exact fields: {problems}" + else: + status, detail = "ok", f"{len(tensors)} dense tensors" + except BaseException as error: # pragma: no cover - transport/env failures + status, detail = "error", f"{type(error).__name__}: {error}" + finally: + if client is not None: + try: + if keys: + client.clear(keys) + client.close() + except BaseException as error: # pragma: no cover - native cleanup failures + status, detail = "error", f"cleanup {type(error).__name__}: {error}" + result_queue.put((status, detail)) + + +def _run_isolated_roundtrip(target: Callable[[Any, str], None], protocol: str, *, timeout: float) -> tuple[str, str]: + """Run one native Mooncake session with a hard process boundary.""" + context = multiprocessing.get_context("spawn") + result_queue = context.Queue() + process = context.Process(target=target, args=(result_queue, protocol)) + try: + process.start() + process.join(timeout=timeout) + timed_out = process.is_alive() + if timed_out: + process.terminate() + process.join(timeout=5) + if process.is_alive(): + process.kill() + process.join(timeout=5) + if process.is_alive(): + pytest.fail(f"{protocol} roundtrip process survived SIGKILL") + if timed_out: + pytest.fail(f"{protocol} roundtrip did not finish within {timeout:.0f} seconds") + if process.exitcode != 0: + pytest.fail(f"{protocol} roundtrip process exited with code {process.exitcode}") + try: + return result_queue.get(timeout=2) + except queue.Empty: + pytest.fail(f"{protocol} roundtrip process returned no result") + finally: + result_queue.close() + result_queue.join_thread() # --------------------------------------------------------------------------- @@ -325,6 +419,27 @@ def record_remote_options(**options): assert remote_options["max_calls"] == 1 assert remote_options["max_retries"] == 0 + def test_cluster_attach_timeout_does_not_leave_process_global_state(self): + """The one-shot worker dies before its abandoned tq.init can mutate + state.""" + env = os.environ.copy() + env["RAY_ENABLE_UV_RUN_RUNTIME_ENV"] = "0" + env.pop("RAY_ADDRESS", None) + repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) + # Ray places Unix-domain sockets below its temp directory. A pytest + # tmp_path can exceed Linux's 107-byte AF_UNIX path limit. + with tempfile.TemporaryDirectory(prefix="tq-ray-") as probe_dir: + result = subprocess.run( + [sys.executable, "-m", "tests.utils._tq_handshake_timeout_probe", probe_dir], + cwd=repo_root, + env=env, + capture_output=True, + text=True, + timeout=45, + check=False, + ) + assert result.returncode == 0, f"probe stdout:\n{result.stdout}\nprobe stderr:\n{result.stderr}" + # --------------------------------------------------------------------------- # Worker detach (attach-only inverse used by every worker teardown hook) @@ -953,6 +1068,9 @@ class TestMooncakeByteExact: def _client(protocol: str): from transfer_queue.storage.clients.mooncake_client import MooncakeStoreClient + from relax.utils.tq_correctness import ensure_mooncake_correctness_guards + + ensure_mooncake_correctness_guards() return MooncakeStoreClient( { "protocol": protocol, @@ -969,30 +1087,14 @@ def _client(protocol: str): @pytest.mark.parametrize("protocol", ["tcp", "rdma"]) def test_multi_dtype_shape_roundtrip_is_byte_exact(self, protocol): - tensors = { - # Production multimodal field names, dimensions, and mixed dtypes. - "pixel_values": torch.randn(64, 1176, dtype=torch.float32).to(torch.bfloat16), - "image_grid_thw": torch.tensor([[1, 8, 8]], dtype=torch.int64), - "input_ids": torch.arange(4096, dtype=torch.int64), - "attention_mask": torch.ones(4096, dtype=torch.int64), - "rewards": torch.linspace(-1, 1, 64, dtype=torch.float32), - "noncontig": torch.randn(128, 256).t(), # transposed == non-contiguous - } - client = self._client(protocol) - try: - keys = [f"bx_{protocol}_{name}" for name in tensors] - values = list(tensors.values()) - client.put(keys, values) - got = client.get( - keys, - shapes=[tuple(v.shape) for v in values], - dtypes=[v.dtype for v in values], - ) - for name, want, have in zip(tensors, values, got, strict=True): - assert have is not None, f"{name} came back empty" - assert torch.equal(have, want.contiguous()), f"{name} is not byte-exact" - finally: - client.close() + """Run each protocol in a pristine, externally bounded session. + + Mooncake 0.3.10 can wedge in ``setup()`` when one process closes a TCP + client and then creates an RDMA client. Process isolation also keeps a + native engine hang from blocking the pytest worker indefinitely. + """ + status, detail = _run_isolated_roundtrip(_dense_roundtrip_worker, protocol, timeout=120) + assert status == "ok", f"{status}: {detail}" @pytest.mark.parametrize("protocol", ["tcp", "rdma"]) def test_multimodal_list_dict_slow_path_roundtrip_is_byte_exact(self, protocol): @@ -1010,14 +1112,5 @@ def test_multimodal_list_dict_slow_path_roundtrip_is_byte_exact(self, protocol): session (mooncake 0.3.10 misbehaves when one process cycles clients across protocols) and a wedged engine cannot hang the suite. """ - context = multiprocessing.get_context("spawn") - result_queue = context.Queue() - process = context.Process(target=_mm_slow_path_worker, args=(result_queue, protocol)) - process.start() - process.join(timeout=240) - if process.is_alive(): - process.terminate() - process.join(timeout=5) - pytest.fail(f"multimodal slow-path roundtrip ({protocol}) did not finish within 240 seconds") - status, detail = result_queue.get(timeout=2) + status, detail = _run_isolated_roundtrip(_mm_slow_path_worker, protocol, timeout=240) assert status == "ok", f"{status}: {detail}" From 6245ee9a4a060bf69393b35f07c8319cad2c0efe Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Fri, 14 Aug 2026 23:15:20 +0800 Subject: [PATCH 19/22] feat(data-plane): version-gated mooncake loss guards MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # ⭐ 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 --- relax/utils/tq_correctness.py | 6 + relax/utils/tq_mooncake_patches.py | 199 ++++++++++++++++++++++++ tests/utils/test_tq_mooncake_patches.py | 178 +++++++++++++++++++++ 3 files changed, 383 insertions(+) create mode 100644 relax/utils/tq_mooncake_patches.py create mode 100644 tests/utils/test_tq_mooncake_patches.py diff --git a/relax/utils/tq_correctness.py b/relax/utils/tq_correctness.py index 8a5e7c91c..4c51ea126 100644 --- a/relax/utils/tq_correctness.py +++ b/relax/utils/tq_correctness.py @@ -62,3 +62,9 @@ def ensure_mooncake_correctness_guards() -> None: missing = [name for name in required_methods if not callable(getattr(MooncakeStoreClient, name, None))] if missing: raise RuntimeError("Installed TransferQueue lacks required Mooncake retry APIs: " + ", ".join(missing)) + + # Version-gated runtime patches for the pinned revision's remaining gaps; + # see relax/utils/tq_mooncake_patches.py for scope and removal condition. + from relax.utils.tq_mooncake_patches import install_mooncake_loss_guards + + install_mooncake_loss_guards() diff --git a/relax/utils/tq_mooncake_patches.py b/relax/utils/tq_mooncake_patches.py new file mode 100644 index 000000000..0e97d074d --- /dev/null +++ b/relax/utils/tq_mooncake_patches.py @@ -0,0 +1,199 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Version-gated runtime loss guards for the pinned TransferQueue revision. + +The TransferQueue revision currently pinned by Relax validates the first +Mooncake batch result but not every retry result, logs removal failures +without raising, and treats a missing or negative production-status ACK as a +successful notification. Those behaviours can turn an explicit storage or +controller failure into silent data loss. Until the equivalent checks land +upstream, this module patches the pinned revision at runtime; installation is +process-local and idempotent. + +These are monkey patches over *private* upstream internals +(``MooncakeStoreClient.__init__``, ``StorageManager._notify_and_wait``), so +they are gated on the exact pinned version: any other transfer_queue refuses +to start rather than running unvalidated patches +(see :func:`_require_pinned_transfer_queue`). + +Removal condition: delete this module and its single call site in +:func:`relax.utils.tq_correctness.ensure_mooncake_correctness_guards` once the +pinned TransferQueue itself validates every batch/retry result, raises on +removal failure, and requires a positive production-status ACK. +""" + +from __future__ import annotations + +import asyncio +from functools import wraps +from typing import Any +from uuid import uuid4 + +from relax.utils.logging_utils import get_logger + + +logger = get_logger(__name__) + +_PATCH_MARKER = "_relax_mooncake_correctness_guards_v1" + +# Exact pins these patches were written and forensically validated against. +_PATCHED_TQ_VERSIONS = ("0.1.10.dev0",) + + +def _require_pinned_transfer_queue() -> None: + """Refuse to patch any transfer_queue version the guards were not written + for.""" + import transfer_queue + + version = str(getattr(transfer_queue, "__version__", "unknown")) + if version not in _PATCHED_TQ_VERSIONS: + raise RuntimeError( + f"transfer_queue {version} is not covered by Relax's Mooncake loss guards " + f"(validated pins: {', '.join(_PATCHED_TQ_VERSIONS)}). These guards replace " + "private upstream internals (MooncakeStoreClient.__init__, " + "StorageManager._notify_and_wait); re-validate them against the new pin and " + "extend _PATCHED_TQ_VERSIONS, or delete relax/utils/tq_mooncake_patches.py " + "entirely if the fixes have landed upstream." + ) + + +def _validate_result_count(operation: str, keys: list[str], results: Any) -> None: + """Require one Mooncake result code for every requested key.""" + try: + actual = len(results) + except TypeError as error: + raise RuntimeError(f"{operation} returned a non-sized result, expected {len(keys)} codes") from error + if actual != len(keys): + raise RuntimeError(f"{operation} returned {actual} results, expected {len(keys)}") + + +class _StrictMooncakeStoreProxy: + """Validate every low-level batch response, including retry calls.""" + + def __init__(self, store: Any) -> None: + self._store = store + + def __getattr__(self, name: str) -> Any: + return getattr(self._store, name) + + def batch_upsert_from(self, keys: list[str], *args: Any, **kwargs: Any) -> Any: + results = self._store.batch_upsert_from(keys, *args, **kwargs) + _validate_result_count("batch_upsert_from", keys, results) + return results + + def batch_get_into(self, keys: list[str], *args: Any, **kwargs: Any) -> Any: + results = self._store.batch_get_into(keys, *args, **kwargs) + _validate_result_count("batch_get_into", keys, results) + return results + + def batch_remove(self, keys: list[str], *args: Any, **kwargs: Any) -> Any: + results = self._store.batch_remove(keys, *args, **kwargs) + _validate_result_count("batch_remove", keys, results) + failures = [(key, code) for key, code in zip(keys, results, strict=True) if code != 0] + if failures: + detail = ", ".join(f"{key}={code}" for key, code in failures) + raise RuntimeError(f"batch_remove failed: {detail}") + return results + + +async def _strict_notify_and_wait(self: Any, request_msg: list) -> None: + """Notify the controller and require a positive ACK within the deadline.""" + import zmq + from transfer_queue.storage.managers import base as tq_base + from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, create_zmq_socket + + identity = f"{self.storage_manager_id}-notify-{uuid4().hex[:8]}".encode() + sock = create_zmq_socket( + ctx=self.zmq_context, + socket_type=zmq.DEALER, + ip=self.controller_info.ip, + identity=identity, + ) + sock.setsockopt(zmq.LINGER, 0) + sock.connect(self.controller_info.to_addr("request_handle_socket")) + + try: + await sock.send_multipart(request_msg) + loop = asyncio.get_running_loop() + deadline = loop.time() + tq_base.TQ_DATA_UPDATE_RESPONSE_TIMEOUT + + while True: + remaining = deadline - loop.time() + if remaining <= 0: + raise TimeoutError( + "Timed out waiting for TransferQueue production-status ACK " + f"after {tq_base.TQ_DATA_UPDATE_RESPONSE_TIMEOUT}s" + ) + try: + messages = await asyncio.wait_for( + sock.recv_multipart(copy=False), + timeout=min(tq_base.TQ_STORAGE_POLLER_TIMEOUT, remaining), + ) + except asyncio.TimeoutError: + continue + except Exception as error: + raise RuntimeError("Failed while waiting for TransferQueue production-status ACK") from error + + response = ZMQMessage.deserialize(messages) + if response.request_type != ZMQRequestType.NOTIFY_DATA_UPDATE_ACK: + continue + body = response.body if isinstance(response.body, dict) else {} + if body.get("success") is not True: + raise RuntimeError( + "TransferQueue controller rejected the production-status update " + f"for partition={body.get('partition_id', 'unknown')}" + ) + return + finally: + try: + if not sock.closed: + sock.close(linger=0) + except Exception as error: # pragma: no cover - best-effort socket cleanup + logger.debug(f"Failed to close TransferQueue notification socket: {error}") + + +def _install_store_guards(client_cls: type) -> None: + if getattr(client_cls, _PATCH_MARKER, False): + return + + original_init = client_cls.__init__ + + @wraps(original_init) + def guarded_init(self: Any, *args: Any, **kwargs: Any) -> None: + original_init(self, *args, **kwargs) + store = getattr(self, "_store", None) + if store is not None and not isinstance(store, _StrictMooncakeStoreProxy): + self._store = _StrictMooncakeStoreProxy(store) + + client_cls.__init__ = guarded_init + setattr(client_cls, _PATCH_MARKER, True) + + +def _install_notification_guards(manager_cls: type) -> None: + if getattr(manager_cls, _PATCH_MARKER, False): + return + + original_notify = manager_cls.notify_data_update + + @wraps(original_notify) + async def guarded_notify(self: Any, *args: Any, **kwargs: Any) -> None: + if not getattr(self, "controller_info", None): + raise RuntimeError("TransferQueue storage manager has no controller for production-status notification") + await original_notify(self, *args, **kwargs) + + manager_cls.notify_data_update = guarded_notify + manager_cls._notify_and_wait = _strict_notify_and_wait + setattr(manager_cls, _PATCH_MARKER, True) + + +def install_mooncake_loss_guards() -> None: + """Install and verify all runtime guards (idempotent, process-local).""" + _require_pinned_transfer_queue() + from transfer_queue.storage.clients.mooncake_client import MooncakeStoreClient + from transfer_queue.storage.managers.base import StorageManager + + _install_store_guards(MooncakeStoreClient) + _install_notification_guards(StorageManager) + + if not getattr(MooncakeStoreClient, _PATCH_MARKER, False) or not getattr(StorageManager, _PATCH_MARKER, False): + raise RuntimeError("Failed to install Mooncake silent-data-loss guards") diff --git a/tests/utils/test_tq_mooncake_patches.py b/tests/utils/test_tq_mooncake_patches.py new file mode 100644 index 000000000..37faa40ac --- /dev/null +++ b/tests/utils/test_tq_mooncake_patches.py @@ -0,0 +1,178 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Tests for the version-gated Mooncake runtime loss guards. + +Moved out of ``test_tq_failure_paths.py`` together with the patches themselves +(``relax/utils/tq_mooncake_patches.py``); everything here is deleted with that +module once the fixes land upstream. +""" + +from __future__ import annotations + +import importlib.util +from types import SimpleNamespace + +import pytest + +from relax.utils.tq_mooncake_patches import ( + _require_pinned_transfer_queue, + _strict_notify_and_wait, + _StrictMooncakeStoreProxy, +) + + +def _has_real_submodule(dotted: str) -> bool: + """True only if a REAL transfer_queue submodule is importable. + + CI installs a single-file ``transfer_queue`` stub; + ``transfer_queue.storage`` does not exist there, so tests that touch the + real MooncakeStoreClient skip on CPU CI and run only where real + TransferQueue is installed. + """ + try: + return importlib.util.find_spec(dotted) is not None + except (ImportError, ValueError, TypeError): + # CI's single-file transfer_queue stub returns a dummy for ``__path__``, + # so find_spec on a submodule raises TypeError instead of returning None. + return False + + +_REAL_MOONCAKE_CLIENT = _has_real_submodule("transfer_queue.storage.clients.mooncake_client") + + +class _SequenceStore: + """Return a configured result sequence from low-level Mooncake calls.""" + + def __init__(self, results: list[list[int]]) -> None: + self.results = iter(results) + + def batch_upsert_from(self, keys, ptrs, sizes, config=None): + return next(self.results) + + def batch_get_into(self, keys, ptrs, sizes): + return next(self.results) + + def batch_remove(self, keys, force=True): + return next(self.results) + + +class _FakeNotifySocket: + def __init__(self) -> None: + self.closed = False + + def setsockopt(self, *args, **kwargs) -> None: + pass + + def connect(self, *args, **kwargs) -> None: + pass + + async def send_multipart(self, request) -> None: + pass + + async def recv_multipart(self, copy=False): + return [b"ack"] + + def close(self, linger=0) -> None: + self.closed = True + + +def _client_with_store(store) -> object: + """A MooncakeStoreClient with only ``_store``/``replica_config`` wired up. + + ``__init__`` is skipped on purpose: it would need a live mooncake master. + """ + from transfer_queue.storage.clients.mooncake_client import MooncakeStoreClient + + client = object.__new__(MooncakeStoreClient) + client._store = store + client.replica_config = None + return client + + +class TestVersionGate: + """Patches refuse to install on any transfer_queue they were not written + for.""" + + def test_unpinned_version_is_rejected(self, monkeypatch): + import transfer_queue + + monkeypatch.setattr(transfer_queue, "__version__", "9.9.9", raising=False) + with pytest.raises(RuntimeError, match="not covered by Relax"): + _require_pinned_transfer_queue() + + def test_pinned_version_is_accepted(self, monkeypatch): + import transfer_queue + + monkeypatch.setattr(transfer_queue, "__version__", "0.1.10.dev0", raising=False) + _require_pinned_transfer_queue() + + +class TestMooncakeCorrectnessGuardPrimitives: + """Low-level response validation stays runnable on the CPU-only CI stub.""" + + def test_upsert_short_result_is_raised(self): + store = _StrictMooncakeStoreProxy(_SequenceStore([[0]])) + with pytest.raises(RuntimeError, match="returned 1 results, expected 2"): + store.batch_upsert_from(["k0", "k1"], [1, 2], [8, 8]) + + def test_remove_failure_is_raised(self): + store = _StrictMooncakeStoreProxy(_SequenceStore([[0, -704]])) + with pytest.raises(RuntimeError, match="batch_remove failed"): + store.batch_remove(["k0", "k1"], force=True) + + +@pytest.mark.skipif( + not _REAL_MOONCAKE_CLIENT, + reason="needs real TransferQueue storage submodules; CPU CI uses a single-file transfer_queue stub", +) +class TestMooncakeCorrectnessGuards: + """Integration with real TransferQueue internals; no GPU/master needed.""" + + def test_retry_short_result_is_never_treated_as_success(self, monkeypatch): + monkeypatch.setattr("transfer_queue.storage.clients.mooncake_client.RETRY_DELAY_SECONDS", 0) + store = _StrictMooncakeStoreProxy(_SequenceStore([[-1, -1], [0]])) + client = _client_with_store(store) + with pytest.raises(RuntimeError, match="returned 1 results, expected 2"): + client._batch_upsert_with_retry(["k0", "k1"], [1, 2], [8, 8]) + + @pytest.mark.asyncio + async def test_negative_production_status_ack_is_raised(self, monkeypatch): + from transfer_queue.utils import zmq_utils + + socket = _FakeNotifySocket() + monkeypatch.setattr(zmq_utils, "create_zmq_socket", lambda **kwargs: socket) + monkeypatch.setattr( + zmq_utils.ZMQMessage, + "deserialize", + staticmethod( + lambda messages: SimpleNamespace( + request_type=zmq_utils.ZMQRequestType.NOTIFY_DATA_UPDATE_ACK, + body={"success": False, "partition_id": "p0"}, + ) + ), + ) + manager = SimpleNamespace( + storage_manager_id="guard-test", + zmq_context=object(), + controller_info=SimpleNamespace(ip="redacted", to_addr=lambda name: "inproc://controller"), + ) + with pytest.raises(RuntimeError, match="rejected the production-status update"): + await _strict_notify_and_wait(manager, [b"request"]) + assert socket.closed is True + + @pytest.mark.asyncio + async def test_missing_production_status_ack_is_bounded(self, monkeypatch): + from transfer_queue.storage.managers import base as tq_base + from transfer_queue.utils import zmq_utils + + socket = _FakeNotifySocket() + monkeypatch.setattr(zmq_utils, "create_zmq_socket", lambda **kwargs: socket) + monkeypatch.setattr(tq_base, "TQ_DATA_UPDATE_RESPONSE_TIMEOUT", 0) + manager = SimpleNamespace( + storage_manager_id="guard-test", + zmq_context=object(), + controller_info=SimpleNamespace(ip="redacted", to_addr=lambda name: "inproc://controller"), + ) + with pytest.raises(TimeoutError, match="production-status ACK"): + await _strict_notify_and_wait(manager, [b"request"]) + assert socket.closed is True From fdee456e3b18663bb651cec12c27849871cdc08b Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Sun, 16 Aug 2026 03:52:19 +0800 Subject: [PATCH 20/22] fix(data-plane): preserve Mooncake semantics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 🐛 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 --- relax/utils/tq_correctness.py | 17 +++-- relax/utils/tq_mooncake_patches.py | 31 ++++---- tests/utils/test_tq_mooncake_patches.py | 94 ++++++++++++++++++++++++- 3 files changed, 118 insertions(+), 24 deletions(-) diff --git a/relax/utils/tq_correctness.py b/relax/utils/tq_correctness.py index 4c51ea126..0ab4e56a6 100644 --- a/relax/utils/tq_correctness.py +++ b/relax/utils/tq_correctness.py @@ -4,12 +4,11 @@ Relax refuses to run MooncakeStore unless the installed TransferQueue and mooncake expose the primitives that make silent data loss detectable. This -module validates capabilities and environment only; it never modifies -TransferQueue at runtime. The temporary runtime patches that harden the -remaining gaps of the pinned revision (per-retry result validation, raising -removal failures, and a strict production-status ACK) are maintained in a -separate version-gated PR so their exact applicability and removal condition -stay reviewable on their own. +module validates capabilities and environment, then installs temporary, +version-gated runtime patches for the remaining gaps of the pinned revision +(per-retry result validation, raising non-idempotent removal failures, and a +strict production-status ACK). The exact applicability and removal condition +remain isolated in :mod:`relax.utils.tq_mooncake_patches`. The pinned mooncake 0.3.10 additionally corrupts TCP-protocol transfers through its auto-enabled memcpy fast path, so that path is force-disabled @@ -47,10 +46,10 @@ def _enforce_safe_memcpy() -> None: def ensure_mooncake_correctness_guards() -> None: - """Validate that the installed stack can run MooncakeStore safely. + """Validate the installed stack and install pinned Mooncake runtime guards. - Read-only: checks the memcpy environment contract and that the pinned - TransferQueue ships the Mooncake retry APIs Relax's data plane relies on. + Checks the memcpy environment contract and required retry APIs before + installing the exact-version patches for remaining upstream gaps. """ _enforce_safe_memcpy() try: diff --git a/relax/utils/tq_mooncake_patches.py b/relax/utils/tq_mooncake_patches.py index 0e97d074d..c37f07c02 100644 --- a/relax/utils/tq_mooncake_patches.py +++ b/relax/utils/tq_mooncake_patches.py @@ -19,7 +19,7 @@ Removal condition: delete this module and its single call site in :func:`relax.utils.tq_correctness.ensure_mooncake_correctness_guards` once the pinned TransferQueue itself validates every batch/retry result, raises on -removal failure, and requires a positive production-status ACK. +non-idempotent removal failure, and requires a positive production-status ACK. """ from __future__ import annotations @@ -35,6 +35,8 @@ logger = get_logger(__name__) _PATCH_MARKER = "_relax_mooncake_correctness_guards_v1" +# Mooncake's idempotent-delete result; the pinned upstream clear path accepts it. +_MOONCAKE_OBJECT_NOT_FOUND = -704 # Exact pins these patches were written and forensically validated against. _PATCHED_TQ_VERSIONS = ("0.1.10.dev0",) @@ -74,6 +76,8 @@ def __init__(self, store: Any) -> None: self._store = store def __getattr__(self, name: str) -> Any: + if name == "_store": + raise AttributeError(name) return getattr(self._store, name) def batch_upsert_from(self, keys: list[str], *args: Any, **kwargs: Any) -> Any: @@ -89,7 +93,9 @@ def batch_get_into(self, keys: list[str], *args: Any, **kwargs: Any) -> Any: def batch_remove(self, keys: list[str], *args: Any, **kwargs: Any) -> Any: results = self._store.batch_remove(keys, *args, **kwargs) _validate_result_count("batch_remove", keys, results) - failures = [(key, code) for key, code in zip(keys, results, strict=True) if code != 0] + failures = [ + (key, code) for key, code in zip(keys, results, strict=True) if code not in (0, _MOONCAKE_OBJECT_NOT_FOUND) + ] if failures: detail = ", ".join(f"{key}={code}" for key, code in failures) raise RuntimeError(f"batch_remove failed: {detail}") @@ -103,16 +109,17 @@ async def _strict_notify_and_wait(self: Any, request_msg: list) -> None: from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, create_zmq_socket identity = f"{self.storage_manager_id}-notify-{uuid4().hex[:8]}".encode() - sock = create_zmq_socket( - ctx=self.zmq_context, - socket_type=zmq.DEALER, - ip=self.controller_info.ip, - identity=identity, - ) - sock.setsockopt(zmq.LINGER, 0) - sock.connect(self.controller_info.to_addr("request_handle_socket")) - + sock = None try: + sock = create_zmq_socket( + ctx=self.zmq_context, + socket_type=zmq.DEALER, + ip=self.controller_info.ip, + identity=identity, + ) + sock.setsockopt(zmq.LINGER, 0) + sock.connect(self.controller_info.to_addr("request_handle_socket")) + await sock.send_multipart(request_msg) loop = asyncio.get_running_loop() deadline = loop.time() + tq_base.TQ_DATA_UPDATE_RESPONSE_TIMEOUT @@ -146,7 +153,7 @@ async def _strict_notify_and_wait(self: Any, request_msg: list) -> None: return finally: try: - if not sock.closed: + if sock is not None and not sock.closed: sock.close(linger=0) except Exception as error: # pragma: no cover - best-effort socket cleanup logger.debug(f"Failed to close TransferQueue notification socket: {error}") diff --git a/tests/utils/test_tq_mooncake_patches.py b/tests/utils/test_tq_mooncake_patches.py index 37faa40ac..7227a6994 100644 --- a/tests/utils/test_tq_mooncake_patches.py +++ b/tests/utils/test_tq_mooncake_patches.py @@ -15,6 +15,8 @@ import pytest from relax.utils.tq_mooncake_patches import ( + _install_notification_guards, + _install_store_guards, _require_pinned_transfer_queue, _strict_notify_and_wait, _StrictMooncakeStoreProxy, @@ -57,14 +59,16 @@ def batch_remove(self, keys, force=True): class _FakeNotifySocket: - def __init__(self) -> None: + def __init__(self, connect_error: Exception | None = None) -> None: self.closed = False + self.connect_error = connect_error def setsockopt(self, *args, **kwargs) -> None: pass def connect(self, *args, **kwargs) -> None: - pass + if self.connect_error is not None: + raise self.connect_error async def send_multipart(self, request) -> None: pass @@ -115,11 +119,56 @@ def test_upsert_short_result_is_raised(self): with pytest.raises(RuntimeError, match="returned 1 results, expected 2"): store.batch_upsert_from(["k0", "k1"], [1, 2], [8, 8]) - def test_remove_failure_is_raised(self): + def test_remove_object_not_found_is_allowed(self): store = _StrictMooncakeStoreProxy(_SequenceStore([[0, -704]])) + assert store.batch_remove(["k0", "k1"], force=True) == [0, -704] + + def test_remove_non_idempotent_failure_is_raised(self): + store = _StrictMooncakeStoreProxy(_SequenceStore([[0, -1]])) with pytest.raises(RuntimeError, match="batch_remove failed"): store.batch_remove(["k0", "k1"], force=True) + def test_missing_wrapped_store_raises_attribute_error(self): + store = object.__new__(_StrictMooncakeStoreProxy) + with pytest.raises(AttributeError): + store.close + + def test_store_guard_installation_is_idempotent(self): + raw_store = _SequenceStore([[0]]) + + class Client: + def __init__(self): + self._store = raw_store + + _install_store_guards(Client) + guarded_init = Client.__init__ + _install_store_guards(Client) + + client = Client() + assert Client.__init__ is guarded_init + assert isinstance(client._store, _StrictMooncakeStoreProxy) + assert client._store._store is raw_store + + +class TestNotificationGuardPrimitives: + @pytest.mark.asyncio + async def test_guarded_notify_rejects_missing_controller(self): + class Manager: + controller_info = None + + async def notify_data_update(self): + raise AssertionError("original notify must not run without a controller") + + async def _notify_and_wait(self, request_msg): + pass + + _install_notification_guards(Manager) + guarded_notify = Manager.notify_data_update + _install_notification_guards(Manager) + assert Manager.notify_data_update is guarded_notify + with pytest.raises(RuntimeError, match="has no controller"): + await Manager().notify_data_update() + @pytest.mark.skipif( not _REAL_MOONCAKE_CLIENT, @@ -160,6 +209,45 @@ async def test_negative_production_status_ack_is_raised(self, monkeypatch): await _strict_notify_and_wait(manager, [b"request"]) assert socket.closed is True + @pytest.mark.asyncio + async def test_positive_production_status_ack_returns(self, monkeypatch): + from transfer_queue.utils import zmq_utils + + socket = _FakeNotifySocket() + monkeypatch.setattr(zmq_utils, "create_zmq_socket", lambda **kwargs: socket) + monkeypatch.setattr( + zmq_utils.ZMQMessage, + "deserialize", + staticmethod( + lambda messages: SimpleNamespace( + request_type=zmq_utils.ZMQRequestType.NOTIFY_DATA_UPDATE_ACK, + body={"success": True, "partition_id": "p0"}, + ) + ), + ) + manager = SimpleNamespace( + storage_manager_id="guard-test", + zmq_context=object(), + controller_info=SimpleNamespace(ip="redacted", to_addr=lambda name: "inproc://controller"), + ) + await _strict_notify_and_wait(manager, [b"request"]) + assert socket.closed is True + + @pytest.mark.asyncio + async def test_connect_failure_closes_notification_socket(self, monkeypatch): + from transfer_queue.utils import zmq_utils + + socket = _FakeNotifySocket(connect_error=ConnectionError("controller unavailable")) + monkeypatch.setattr(zmq_utils, "create_zmq_socket", lambda **kwargs: socket) + manager = SimpleNamespace( + storage_manager_id="guard-test", + zmq_context=object(), + controller_info=SimpleNamespace(ip="redacted", to_addr=lambda name: "inproc://controller"), + ) + with pytest.raises(ConnectionError, match="controller unavailable"): + await _strict_notify_and_wait(manager, [b"request"]) + assert socket.closed is True + @pytest.mark.asyncio async def test_missing_production_status_ack_is_bounded(self, monkeypatch): from transfer_queue.storage.managers import base as tq_base From 6d984ce82a81af2fdd6e163ca05428ca3eb3edd0 Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Tue, 18 Aug 2026 03:53:47 +0800 Subject: [PATCH 21/22] test(data-plane): cover short get retry results --- tests/utils/test_tq_mooncake_patches.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tests/utils/test_tq_mooncake_patches.py b/tests/utils/test_tq_mooncake_patches.py index 7227a6994..9dd1ac8f6 100644 --- a/tests/utils/test_tq_mooncake_patches.py +++ b/tests/utils/test_tq_mooncake_patches.py @@ -119,6 +119,11 @@ def test_upsert_short_result_is_raised(self): with pytest.raises(RuntimeError, match="returned 1 results, expected 2"): store.batch_upsert_from(["k0", "k1"], [1, 2], [8, 8]) + def test_get_short_result_is_raised(self): + store = _StrictMooncakeStoreProxy(_SequenceStore([[0]])) + with pytest.raises(RuntimeError, match="returned 1 results, expected 2"): + store.batch_get_into(["k0", "k1"], [1, 2], [8, 8]) + def test_remove_object_not_found_is_allowed(self): store = _StrictMooncakeStoreProxy(_SequenceStore([[0, -704]])) assert store.batch_remove(["k0", "k1"], force=True) == [0, -704] @@ -184,6 +189,13 @@ def test_retry_short_result_is_never_treated_as_success(self, monkeypatch): with pytest.raises(RuntimeError, match="returned 1 results, expected 2"): client._batch_upsert_with_retry(["k0", "k1"], [1, 2], [8, 8]) + def test_get_retry_short_result_is_never_treated_as_success(self, monkeypatch): + monkeypatch.setattr("transfer_queue.storage.clients.mooncake_client.RETRY_DELAY_SECONDS", 0) + store = _StrictMooncakeStoreProxy(_SequenceStore([[-1, -1], [0]])) + client = _client_with_store(store) + with pytest.raises(RuntimeError, match="returned 1 results, expected 2"): + client._batch_get_into_with_retry(["k0", "k1"], [1, 2], [8, 8]) + @pytest.mark.asyncio async def test_negative_production_status_ack_is_raised(self, monkeypatch): from transfer_queue.utils import zmq_utils From c500eb4d10b434ba9decb3bfe26ec04a5433ea47 Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Tue, 18 Aug 2026 17:29:38 +0800 Subject: [PATCH 22/22] fix(data-plane): gate guards by TQ revision MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 🐛 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 --- relax/utils/tq_mooncake_patches.py | 75 ++++++++++++++--- tests/utils/test_tq_mooncake_patches.py | 106 +++++++++++++++++++++++- 2 files changed, 170 insertions(+), 11 deletions(-) diff --git a/relax/utils/tq_mooncake_patches.py b/relax/utils/tq_mooncake_patches.py index c37f07c02..a4d6b2abe 100644 --- a/relax/utils/tq_mooncake_patches.py +++ b/relax/utils/tq_mooncake_patches.py @@ -12,9 +12,9 @@ These are monkey patches over *private* upstream internals (``MooncakeStoreClient.__init__``, ``StorageManager._notify_and_wait``), so -they are gated on the exact pinned version: any other transfer_queue refuses -to start rather than running unvalidated patches -(see :func:`_require_pinned_transfer_queue`). +they are gated on the exact pinned package version and installed VCS revision: +any other build refuses to start rather than running unvalidated patches (see +:func:`_require_pinned_transfer_queue`). Removal condition: delete this module and its single call site in :func:`relax.utils.tq_correctness.ensure_mooncake_correctness_guards` once the @@ -25,7 +25,11 @@ from __future__ import annotations import asyncio +import json +import re from functools import wraps +from importlib import metadata +from pathlib import Path from typing import Any from uuid import uuid4 @@ -38,26 +42,77 @@ # Mooncake's idempotent-delete result; the pinned upstream clear path accepts it. _MOONCAKE_OBJECT_NOT_FOUND = -704 -# Exact pins these patches were written and forensically validated against. -_PATCHED_TQ_VERSIONS = ("0.1.10.dev0",) +# Exact builds these patches were written and forensically validated against. +_PATCHED_TQ_BUILDS = { + "0.1.10.dev0": frozenset({"58054a33834aadbcf76aacd6b1e32e25c030f2c9"}), +} + + +def _installed_transfer_queue_revision(transfer_queue_module: Any) -> str | None: + """Read the installed VCS revision from standard PEP 610 metadata.""" + try: + distribution = metadata.distribution("transferqueue") + if str(distribution.version) != str(getattr(transfer_queue_module, "__version__", "unknown")): + return None + + module_file = getattr(transfer_queue_module, "__file__", None) + if not module_file: + return None + module_path = Path(module_file).resolve() + package_root = Path(distribution.locate_file("transfer_queue")).resolve() + module_path.relative_to(package_root) + + direct_url = distribution.read_text("direct_url.json") + except (metadata.PackageNotFoundError, AttributeError, OSError, TypeError, UnicodeError, ValueError): + return None + if direct_url is None: + return None + + try: + provenance = json.loads(direct_url) + except (json.JSONDecodeError, TypeError): + return None + if not isinstance(provenance, dict): + return None + + vcs_info = provenance.get("vcs_info") + if not isinstance(vcs_info, dict) or vcs_info.get("vcs") != "git": + return None + commit_id = vcs_info.get("commit_id") + if not isinstance(commit_id, str) or not commit_id.strip(): + return None + revision = commit_id.strip().lower() + if re.fullmatch(r"[0-9a-f]{40}", revision) is None: + return None + return revision def _require_pinned_transfer_queue() -> None: - """Refuse to patch any transfer_queue version the guards were not written - for.""" + """Refuse to patch any TransferQueue source revision not validated here.""" import transfer_queue version = str(getattr(transfer_queue, "__version__", "unknown")) - if version not in _PATCHED_TQ_VERSIONS: + expected_revisions = _PATCHED_TQ_BUILDS.get(version) + if expected_revisions is None: raise RuntimeError( f"transfer_queue {version} is not covered by Relax's Mooncake loss guards " - f"(validated pins: {', '.join(_PATCHED_TQ_VERSIONS)}). These guards replace " + f"(validated versions: {', '.join(_PATCHED_TQ_BUILDS)}). These guards replace " "private upstream internals (MooncakeStoreClient.__init__, " "StorageManager._notify_and_wait); re-validate them against the new pin and " - "extend _PATCHED_TQ_VERSIONS, or delete relax/utils/tq_mooncake_patches.py " + "extend _PATCHED_TQ_BUILDS, or delete relax/utils/tq_mooncake_patches.py " "entirely if the fixes have landed upstream." ) + revision = _installed_transfer_queue_revision(transfer_queue) + if revision not in expected_revisions: + actual_revision = revision or "unknown" + raise RuntimeError( + f"transfer_queue {version} revision {actual_revision} is not covered by Relax's Mooncake loss guards " + f"(validated revisions: {', '.join(sorted(expected_revisions))}). The exact source revision is required " + "because these guards replace private upstream internals; install the validated pin or re-validate " + "the guards before extending _PATCHED_TQ_BUILDS." + ) + def _validate_result_count(operation: str, keys: list[str], results: Any) -> None: """Require one Mooncake result code for every requested key.""" diff --git a/tests/utils/test_tq_mooncake_patches.py b/tests/utils/test_tq_mooncake_patches.py index 9dd1ac8f6..1c9c36e69 100644 --- a/tests/utils/test_tq_mooncake_patches.py +++ b/tests/utils/test_tq_mooncake_patches.py @@ -17,6 +17,7 @@ from relax.utils.tq_mooncake_patches import ( _install_notification_guards, _install_store_guards, + _installed_transfer_queue_revision, _require_pinned_transfer_queue, _strict_notify_and_wait, _StrictMooncakeStoreProxy, @@ -97,6 +98,8 @@ class TestVersionGate: """Patches refuse to install on any transfer_queue they were not written for.""" + _PINNED_REVISION = "58054a33834aadbcf76aacd6b1e32e25c030f2c9" + def test_unpinned_version_is_rejected(self, monkeypatch): import transfer_queue @@ -104,12 +107,113 @@ def test_unpinned_version_is_rejected(self, monkeypatch): with pytest.raises(RuntimeError, match="not covered by Relax"): _require_pinned_transfer_queue() - def test_pinned_version_is_accepted(self, monkeypatch): + def test_pinned_version_and_revision_are_accepted(self, monkeypatch): import transfer_queue monkeypatch.setattr(transfer_queue, "__version__", "0.1.10.dev0", raising=False) + monkeypatch.setattr( + "relax.utils.tq_mooncake_patches._installed_transfer_queue_revision", + lambda module: self._PINNED_REVISION, + ) _require_pinned_transfer_queue() + def test_same_version_with_different_revision_is_rejected(self, monkeypatch): + import transfer_queue + + monkeypatch.setattr(transfer_queue, "__version__", "0.1.10.dev0", raising=False) + monkeypatch.setattr( + "relax.utils.tq_mooncake_patches._installed_transfer_queue_revision", + lambda module: "0" * 40, + ) + with pytest.raises(RuntimeError, match="revision .* is not covered"): + _require_pinned_transfer_queue() + + def test_missing_revision_metadata_is_rejected(self, monkeypatch): + import transfer_queue + + monkeypatch.setattr(transfer_queue, "__version__", "0.1.10.dev0", raising=False) + monkeypatch.setattr( + "relax.utils.tq_mooncake_patches._installed_transfer_queue_revision", + lambda module: None, + ) + with pytest.raises(RuntimeError, match="revision unknown is not covered"): + _require_pinned_transfer_queue() + + @pytest.mark.parametrize( + ("direct_url", "expected"), + [ + (None, None), + ("not-json", None), + ("{}", None), + ('{"vcs_info": {}}', None), + ('{"vcs_info": {"vcs": "hg", "commit_id": "58054a33834aadbcf76aacd6b1e32e25c030f2c9"}}', None), + ('{"vcs_info": {"vcs": "git", "commit_id": "abc123"}}', None), + ( + '{"vcs_info": {"vcs": "git", "commit_id": " 58054A33834AADBCF76AACD6B1E32E25C030F2C9 "}}', + _PINNED_REVISION, + ), + ], + ) + def test_revision_metadata_is_parsed_fail_closed(self, monkeypatch, tmp_path, direct_url, expected): + package_root = tmp_path / "installed" / "transfer_queue" + + class Distribution: + version = "0.1.10.dev0" + + def locate_file(self, filename): + assert filename == "transfer_queue" + return package_root + + def read_text(self, filename): + assert filename == "direct_url.json" + return direct_url + + monkeypatch.setattr( + "relax.utils.tq_mooncake_patches.metadata.distribution", + lambda name: Distribution(), + ) + module = SimpleNamespace( + __version__="0.1.10.dev0", + __file__=package_root / "__init__.py", + ) + assert _installed_transfer_queue_revision(module) == expected + + def test_distribution_version_mismatch_is_unverifiable(self, monkeypatch, tmp_path): + class Distribution: + version = "0.1.10.dev1" + + monkeypatch.setattr( + "relax.utils.tq_mooncake_patches.metadata.distribution", + lambda name: Distribution(), + ) + module = SimpleNamespace( + __version__="0.1.10.dev0", + __file__=tmp_path / "transfer_queue" / "__init__.py", + ) + assert _installed_transfer_queue_revision(module) is None + + def test_shadowed_module_is_unverifiable(self, monkeypatch, tmp_path): + package_root = tmp_path / "installed" / "transfer_queue" + + class Distribution: + version = "0.1.10.dev0" + + def locate_file(self, filename): + return package_root + + def read_text(self, filename): + return '{"vcs_info": {"vcs": "git", "commit_id": "58054a33834aadbcf76aacd6b1e32e25c030f2c9"}}' + + monkeypatch.setattr( + "relax.utils.tq_mooncake_patches.metadata.distribution", + lambda name: Distribution(), + ) + shadowed_module = SimpleNamespace( + __version__="0.1.10.dev0", + __file__=tmp_path / "shadow" / "transfer_queue.py", + ) + assert _installed_transfer_queue_revision(shadowed_module) is None + class TestMooncakeCorrectnessGuardPrimitives: """Low-level response validation stays runnable on the CPU-only CI stub."""