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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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 6f145fb2c245991b0daa756ff1bab6b289953876 Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Sun, 23 Aug 2026 12:28:31 +0800 Subject: [PATCH 19/30] refactor(tq): group transfer queue utilities --- relax/backends/megatron/actor.py | 2 +- relax/components/actor.py | 2 +- relax/components/actor_fwd.py | 2 +- relax/components/advantages.py | 2 +- relax/components/base.py | 2 +- relax/components/critic.py | 2 +- relax/components/rollout.py | 2 +- relax/components/sft.py | 2 +- relax/core/controller.py | 12 +-- relax/distributed/ray/rollout.py | 2 +- relax/utils/tq/__init__.py | 7 ++ relax/utils/{tq_config.py => tq/config.py} | 2 +- .../{tq_correctness.py => tq/correctness.py} | 0 .../{tq_lifecycle.py => tq/lifecycle.py} | 14 +-- scripts/benchmarks/tq_cross_node_bench.py | 2 +- scripts/benchmarks/tq_rdma_bench.py | 6 +- tests/utils/_tq_handshake_timeout_probe.py | 2 +- tests/utils/mm_payload_fixtures.py | 2 +- tests/utils/test_rdma_probe.py | 4 +- tests/utils/test_tq_dataplane_behavior.py | 2 +- tests/utils/test_tq_failure_paths.py | 8 +- tests/utils/tq/__init__.py | 3 + tests/utils/tq/_payload_assertions.py | 91 +++++++++++++++++++ tests/utils/tq/test_payload_assertions.py | 84 +++++++++++++++++ 24 files changed, 219 insertions(+), 38 deletions(-) create mode 100644 relax/utils/tq/__init__.py rename relax/utils/{tq_config.py => tq/config.py} (99%) rename relax/utils/{tq_correctness.py => tq/correctness.py} (100%) rename relax/utils/{tq_lifecycle.py => tq/lifecycle.py} (99%) create mode 100644 tests/utils/tq/__init__.py create mode 100644 tests/utils/tq/_payload_assertions.py create mode 100644 tests/utils/tq/test_payload_assertions.py diff --git a/relax/backends/megatron/actor.py b/relax/backends/megatron/actor.py index 51c897e66..e6fe7398c 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, detach_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 diff --git a/relax/components/actor.py b/relax/components/actor.py index d121e262e..3811e60ad 100644 --- a/relax/components/actor.py +++ b/relax/components/actor.py @@ -16,7 +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 +from relax.utils.tq.lifecycle import attach_tq_client app = FastAPI() diff --git a/relax/components/actor_fwd.py b/relax/components/actor_fwd.py index 9f259d9a0..1e1848610 100644 --- a/relax/components/actor_fwd.py +++ b/relax/components/actor_fwd.py @@ -11,7 +11,7 @@ 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 +from relax.utils.tq.lifecycle import attach_tq_client app = FastAPI() diff --git a/relax/components/advantages.py b/relax/components/advantages.py index 8cb8780c2..b28ac5ad4 100644 --- a/relax/components/advantages.py +++ b/relax/components/advantages.py @@ -16,7 +16,7 @@ apply_opd_to_advantages, consume_opd_advantage_data, ) -from relax.utils.tq_lifecycle import attach_tq_client +from relax.utils.tq.lifecycle import attach_tq_client from relax.utils.training.ppo_utils import ( compute_approx_kl, get_advantages_and_returns_batch, diff --git a/relax/components/base.py b/relax/components/base.py index c84088187..4b8b68aed 100644 --- a/relax/components/base.py +++ b/relax/components/base.py @@ -100,7 +100,7 @@ def __del__(self) -> 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 + from relax.utils.tq.lifecycle import detach_tq_client detach_tq_client(generation) self._tq_client_generation = None diff --git a/relax/components/critic.py b/relax/components/critic.py index b3e2240d5..6e367f1f2 100644 --- a/relax/components/critic.py +++ b/relax/components/critic.py @@ -15,7 +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 +from relax.utils.tq.lifecycle import attach_tq_client @serve.deployment( diff --git a/relax/components/rollout.py b/relax/components/rollout.py index 2da98f0df..4968161dd 100644 --- a/relax/components/rollout.py +++ b/relax/components/rollout.py @@ -19,7 +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 +from relax.utils.tq.lifecycle import attach_tq_client app = FastAPI() diff --git a/relax/components/sft.py b/relax/components/sft.py index 939b39dbf..0d3a98635 100644 --- a/relax/components/sft.py +++ b/relax/components/sft.py @@ -38,7 +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.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 diff --git a/relax/core/controller.py b/relax/core/controller.py index 00d63fdd3..9f511dc3a 100644 --- a/relax/core/controller.py +++ b/relax/core/controller.py @@ -49,13 +49,13 @@ ) 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 ( +from relax.utils.tq.config import ( build_backend_config, resolve_mooncake_master_address, resolve_tq_capacity_batch_size, validate_mooncake_runtime_contract, ) -from relax.utils.tq_lifecycle import ( +from relax.utils.tq.lifecycle import ( TqInitResult, close_tq_owner, initialize_tq_with_fallback, @@ -326,7 +326,7 @@ def _initialize_data_system(self): fallback_config = None if backend_config.get("storage_backend") == "MooncakeStore": - from relax.utils.tq_config import build_simple_storage_config + from relax.utils.tq.config import build_simple_storage_config fallback_config = OmegaConf.create( { @@ -414,7 +414,7 @@ def _resolve_tq_backend(self, total_storage_size: int) -> dict: # 2. SimpleStorage short-circuit (default, zero behavior change). # ``mooncake + off`` is MooncakeStore/TCP, not SimpleStorage. if backend == "simple": - from relax.utils.tq_config import build_simple_storage_config + from relax.utils.tq.config import build_simple_storage_config return build_simple_storage_config( total_storage_size=total_storage_size, @@ -434,7 +434,7 @@ def _resolve_tq_backend(self, total_storage_size: int) -> dict: 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 + from relax.utils.tq.config import build_simple_storage_config logger.warning( "[dataplane] Installed TransferQueue does not satisfy the Mooncake " @@ -512,7 +512,7 @@ def _close_data_system(self) -> None: Default in-process path keeps upstream's plain ``tq.close()``; owner-mediated runs delegate to - :func:`relax.utils.tq_lifecycle.close_tq_owner`. + :func:`relax.utils.tq.lifecycle.close_tq_owner`. """ if self._tq_legacy_init: self._tq_legacy_init = False diff --git a/relax/distributed/ray/rollout.py b/relax/distributed/ray/rollout.py index 5a4887a73..2e33cd810 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, detach_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, diff --git a/relax/utils/tq/__init__.py b/relax/utils/tq/__init__.py new file mode 100644 index 000000000..656f34359 --- /dev/null +++ b/relax/utils/tq/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""TransferQueue integration utilities. + +Import concrete submodules directly so importing :mod:`relax.utils.tq` stays +lightweight and does not eagerly load Ray, TransferQueue, or Mooncake. +""" diff --git a/relax/utils/tq_config.py b/relax/utils/tq/config.py similarity index 99% rename from relax/utils/tq_config.py rename to relax/utils/tq/config.py index 2cd646f34..83aeb5c63 100644 --- a/relax/utils/tq_config.py +++ b/relax/utils/tq/config.py @@ -20,7 +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 +from relax.utils.tq.correctness import ensure_mooncake_correctness_guards logger = get_logger(__name__) diff --git a/relax/utils/tq_correctness.py b/relax/utils/tq/correctness.py similarity index 100% rename from relax/utils/tq_correctness.py rename to relax/utils/tq/correctness.py diff --git a/relax/utils/tq_lifecycle.py b/relax/utils/tq/lifecycle.py similarity index 99% rename from relax/utils/tq_lifecycle.py rename to relax/utils/tq/lifecycle.py index 131777486..00a0315ae 100644 --- a/relax/utils/tq_lifecycle.py +++ b/relax/utils/tq/lifecycle.py @@ -157,24 +157,20 @@ def _backend_description(conf: Any) -> str: 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 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) + 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): + if not uses_mooncake(conf): return - from relax.utils.tq_config import validate_mooncake_runtime_contract + from relax.utils.tq.config import validate_mooncake_runtime_contract validate_mooncake_runtime_contract() @@ -503,7 +499,7 @@ def verify_cluster_attach(conf: Any, *, timeout: float | None = None) -> list[st @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 + 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() diff --git a/scripts/benchmarks/tq_cross_node_bench.py b/scripts/benchmarks/tq_cross_node_bench.py index 31b39ca97..88ddf16ca 100644 --- a/scripts/benchmarks/tq_cross_node_bench.py +++ b/scripts/benchmarks/tq_cross_node_bench.py @@ -410,7 +410,7 @@ 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 ( + from relax.utils.tq.config import ( build_mooncake_config, build_simple_storage_config, validate_mooncake_runtime_contract, diff --git a/scripts/benchmarks/tq_rdma_bench.py b/scripts/benchmarks/tq_rdma_bench.py index 1399054d7..6e4ac14ec 100644 --- a/scripts/benchmarks/tq_rdma_bench.py +++ b/scripts/benchmarks/tq_rdma_bench.py @@ -87,7 +87,7 @@ def parse_args() -> argparse.Namespace: 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 + Reuses :mod:`relax.utils.tq.config` builders so the benchmark cannot drift from the production config shape (single source of truth for keys/defaults). """ @@ -95,7 +95,7 @@ def build_tq_config(config_name: str, args: argparse.Namespace, num_storage_unit 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 cfg = CONFIG_MAP[config_name] sampler = GRPOGroupNSampler(n_samples_per_prompt=1) @@ -210,7 +210,7 @@ def run_one(config_name: str, payload: dict, args: argparse.Namespace) -> dict: import transfer_queue as tq if CONFIG_MAP[config_name]["backend"] == "MooncakeStore": - from relax.utils.tq_config import validate_mooncake_runtime_contract + from relax.utils.tq.config import validate_mooncake_runtime_contract validate_mooncake_runtime_contract() diff --git a/tests/utils/_tq_handshake_timeout_probe.py b/tests/utils/_tq_handshake_timeout_probe.py index 331b05b86..6a2fffafb 100644 --- a/tests/utils/_tq_handshake_timeout_probe.py +++ b/tests/utils/_tq_handshake_timeout_probe.py @@ -85,7 +85,7 @@ def main(probe_dir: Path) -> None: _temp_dir=str(probe_dir / "ray"), ) try: - from relax.utils import tq_lifecycle + from relax.utils.tq import lifecycle as tq_lifecycle conf = {"backend": {"storage_backend": "SimpleStorage"}, "controller": {}} diff --git a/tests/utils/mm_payload_fixtures.py b/tests/utils/mm_payload_fixtures.py index 690ab2fb0..be226e6f8 100644 --- a/tests/utils/mm_payload_fixtures.py +++ b/tests/utils/mm_payload_fixtures.py @@ -28,7 +28,7 @@ import torch -from relax.utils.payload_digest import diff_digests, leaf_digests +from tests.utils.tq._payload_assertions import diff_digests, leaf_digests _REPO_ROOT = Path(__file__).resolve().parents[2] diff --git a/tests/utils/test_rdma_probe.py b/tests/utils/test_rdma_probe.py index 1763b47f0..f673a47fb 100644 --- a/tests/utils/test_rdma_probe.py +++ b/tests/utils/test_rdma_probe.py @@ -30,7 +30,7 @@ reduce_results, validate_config, ) -from relax.utils.tq_config import ( +from relax.utils.tq.config import ( build_mooncake_config, build_simple_storage_config, estimate_payload_bytes, @@ -630,7 +630,7 @@ def test_segment_capacity_env_override_raises_the_ceiling(self, monkeypatch): 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 + 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"): diff --git a/tests/utils/test_tq_dataplane_behavior.py b/tests/utils/test_tq_dataplane_behavior.py index cb2509924..6dee42e57 100644 --- a/tests/utils/test_tq_dataplane_behavior.py +++ b/tests/utils/test_tq_dataplane_behavior.py @@ -332,9 +332,9 @@ 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 + from tests.utils.tq._payload_assertions import diff_digests, leaf_digests num_samples = 4 train_data, source = mm_train_data(num_samples) diff --git a/tests/utils/test_tq_failure_paths.py b/tests/utils/test_tq_failure_paths.py index 8cdd713a9..16db01c37 100644 --- a/tests/utils/test_tq_failure_paths.py +++ b/tests/utils/test_tq_failure_paths.py @@ -14,7 +14,7 @@ (``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``) +* 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. @@ -39,8 +39,8 @@ import pytest import torch -from relax.utils import tq_lifecycle from relax.utils.rdma_probe import ProbeResult, reduce_results +from relax.utils.tq import lifecycle as tq_lifecycle def _has_real_submodule(dotted: str) -> bool: @@ -140,8 +140,8 @@ def _mm_slow_path_worker(result_queue, protocol: str) -> 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 + from tests.utils.tq._payload_assertions import diff_digests, leaf_digests train_data, source = mm_train_data(4) samples = train_data["multimodal_train_inputs"] @@ -1068,7 +1068,7 @@ 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 + from relax.utils.tq.correctness import ensure_mooncake_correctness_guards ensure_mooncake_correctness_guards() return MooncakeStoreClient( diff --git a/tests/utils/tq/__init__.py b/tests/utils/tq/__init__.py new file mode 100644 index 000000000..24cc87577 --- /dev/null +++ b/tests/utils/tq/__init__.py @@ -0,0 +1,3 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""TransferQueue test helpers and focused tests.""" diff --git a/tests/utils/tq/_payload_assertions.py b/tests/utils/tq/_payload_assertions.py new file mode 100644 index 000000000..4ca3778e4 --- /dev/null +++ b/tests/utils/tq/_payload_assertions.py @@ -0,0 +1,91 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Raw-byte payload assertions shared by TransferQueue tests. + +These helpers intentionally live under ``tests``: production code does not need +generic payload traversal, while the data-plane contract must distinguish byte +identity from value equality for NaNs, signed zero, and nested tensors. +""" + +from __future__ import annotations + +import hashlib +from typing import Any + +import numpy as np +import torch + + +LeafDigest = tuple[str, str, str] + + +def _tensor_digest(value: torch.Tensor) -> LeafDigest: + 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: + raw = repr(value).encode("utf-8") + return (f"py.{type(value).__name__}", "", hashlib.sha256(raw).hexdigest()) + + +def _unwrap_non_tensor(value: Any) -> Any: + if type(value).__name__ == "NonTensorStack": + return value.tolist() + if type(value).__name__ == "NonTensorData": + return value.data + return value + + +def leaf_digests(payload: Any, prefix: str = "payload") -> dict[str, LeafDigest]: + """Map every supported leaf to ``(dtype, shape, raw-byte SHA-256)``.""" + 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): + 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 mismatch descriptions; an 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 diff --git a/tests/utils/tq/test_payload_assertions.py b/tests/utils/tq/test_payload_assertions.py new file mode 100644 index 000000000..b930c7d0a --- /dev/null +++ b/tests/utils/tq/test_payload_assertions.py @@ -0,0 +1,84 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""CPU-only tests for TransferQueue payload byte-identity assertions.""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +import pytest +import torch + +from tests.utils.tq._payload_assertions import diff_digests, leaf_digests + + +class NonTensorData: + """Minimal tensordict-compatible wrapper for an optional-dependency-free + test.""" + + def __init__(self, data: Any) -> None: + self.data = data + + +class NonTensorStack: + """Minimal tensordict-compatible wrapper for an optional-dependency-free + test.""" + + def __init__(self, data: Any) -> None: + self._data = data + + def tolist(self) -> Any: + return self._data + + +def test_leaf_digests_distinguishes_raw_bytes_from_value_equality(): + positive_zero = leaf_digests(torch.tensor([0.0], dtype=torch.float32)) + negative_zero = leaf_digests(torch.tensor([-0.0], dtype=torch.float32)) + + assert positive_zero["payload"][:2] == negative_zero["payload"][:2] + assert positive_zero["payload"][2] != negative_zero["payload"][2] + assert diff_digests(positive_zero, negative_zero) == [ + f"payload: sha256 mismatch (expected {positive_zero['payload'][2]}, got {negative_zero['payload'][2]})" + ] + + +def test_leaf_digests_preserves_dtype_shape_and_nested_tensor_rows(): + with pytest.warns(UserWarning, match="prototype stage"): + nested = torch.nested.nested_tensor( + [torch.tensor([1, 2], dtype=torch.int16), torch.tensor([3], dtype=torch.int16)] + ) + payload = { + "array": np.array([[1, 2]], dtype=np.uint16), + "nested": nested, + "tensor": torch.tensor([[1, 2]], dtype=torch.int32), + } + + digests = leaf_digests(payload) + + assert digests["payload.array"][:2] == ("np.uint16", "1x2") + assert digests["payload.nested[0]"][:2] == ("torch.int16", "2") + assert digests["payload.nested[1]"][:2] == ("torch.int16", "1") + assert digests["payload.tensor"][:2] == ("torch.int32", "1x2") + + +def test_leaf_digests_unwraps_non_tensor_containers(): + wrapped = NonTensorStack([NonTensorData({"text": "hello"}), NonTensorData(None)]) + + assert leaf_digests(wrapped) == leaf_digests([{"text": "hello"}, None]) + + +def test_leaf_digests_rejects_unknown_leaf_type(): + with pytest.raises(TypeError, match=r"Unsupported payload leaf at payload\.bad: object"): + leaf_digests({"bad": object()}) + + +def test_diff_digests_reports_missing_and_extra_leaves(): + expected = leaf_digests({"expected": 1}) + actual = leaf_digests({"actual": 1}) + + problems = diff_digests(expected, actual) + + assert len(problems) == 2 + assert problems[0].startswith("payload.actual: unexpected extra leaf") + assert problems[1].startswith("payload.expected: missing") From acee359dc84bec3970dee8dc047988263840c3c0 Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Sun, 23 Aug 2026 15:24:46 +0800 Subject: [PATCH 20/30] refactor(tq): remove deferred GDR support MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # ♻️ Refactor ## Limit the initial transport to host RDMA - Remove public GDR configuration and runtime status plumbing - Pin Mooncake client configuration to host memory transfers - Preserve compatibility checks for legacy GDR controllers --- # ✅ Tests ## Cover the host-RDMA contract - Remove deferred GDR behavior tests - Reject attaching to a controller configured for GDR --- # 📝 Documentation ## Clarify the initial transport scope - Document host RDMA as the supported first-phase path --- docs/draft/transfer_queue_rdma.md | 18 +++--- relax/backends/megatron/actor.py | 1 - relax/components/actor.py | 1 - relax/components/actor_fwd.py | 1 - relax/components/advantages.py | 1 - relax/components/critic.py | 1 - relax/components/rollout.py | 1 - relax/components/sft.py | 1 - relax/core/controller.py | 27 +-------- relax/distributed/ray/rollout.py | 1 - relax/utils/arguments.py | 17 +----- relax/utils/rdma_probe.py | 43 +------------- relax/utils/tq/config.py | 15 ++--- relax/utils/tq/lifecycle.py | 50 +++------------- scripts/benchmarks/tq_cross_node_bench.py | 2 +- scripts/benchmarks/tq_rdma_bench.py | 1 - tests/utils/test_rdma_probe.py | 68 ++++++---------------- tests/utils/test_tq_failure_paths.py | 70 +++++++---------------- 18 files changed, 67 insertions(+), 252 deletions(-) diff --git a/docs/draft/transfer_queue_rdma.md b/docs/draft/transfer_queue_rdma.md index 5f1400e2d..415128471 100644 --- a/docs/draft/transfer_queue_rdma.md +++ b/docs/draft/transfer_queue_rdma.md @@ -8,22 +8,19 @@ Relax 的数据面（rollout ↔ train 之间的样本传输）默认走 Transfe ## 配置入口 -只暴露四个表达使用意图的参数，Mooncake 底层参数（endpoint、buffer、segment、timeout、master 策略）不做 CLI，走内部默认与部署环境。 +只暴露三个表达使用意图的参数，Mooncake 底层参数（endpoint、buffer、segment、timeout、master 策略）不做 CLI，走内部默认与部署环境。 | 参数 | 取值 | 说明 | |---|---|---| | `--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` | 默认关 | **实验性**，见下文 | -`--tq-rdma-mode=required` 只覆盖**传输层**（MooncakeStore + RDMA 可用性与 segment 容量），不覆盖 GDR。 +`--tq-rdma-mode=required` 覆盖的是**传输层**：MooncakeStore + RDMA 可用性与 segment 容量。 -### GDR 为实验性 +### 首期只交付 host RDMA -GDR 的可用性无法由 driver 的启动探测代表：探测跑在独立 Ray task 中，该进程没有初始化 CUDA context；真正创建 staging buffer 的是每个 worker 的 TQ 客户端。因此首期不宣称 job 级 GDR 已验证，`required` 也不对 GDR 做 fail-fast。 - -每个 worker 附加 TQ 后都会记录两层信息：`requested=true` 表示用户请求开启；`status=host_rdma_fallback/enabled_unverified/inactive/unknown` 表示该 worker 的本地观察。即使本地 staging buffer 已创建也只记为 `enabled_unverified`，不等同于线上流量已经证明走 GDR。 +数据面 payload 先经过主机内存再通过 RDMA 跨节点传输。GDR（GPU Direct RDMA）不在本期范围内：它不是任务书要求，可用性也无法由 driver 的启动探测代表（探测进程没有初始化 CUDA context），因此 Relax 把 Mooncake 的 `use_gdr` 固定为 `false`，不提供开关。如后续出现真实需求，GDR 应由独立 PR 实现并单独验证。 ## 启动流程与降级 @@ -40,7 +37,6 @@ driver 在**第一次 `tq.init` 之前**完成探测并生成 job 级唯一的 e 降级阶梯： ``` -GDR → host RDMA （worker 运行时判定，记录 requested 与实际状态） RDMA → Mooncake/TCP （任一节点无 RDMA 能力，或指定设备缺失） Mooncake → SimpleStorage（任一节点 mooncake/master 不可用、运行时正确性契约不满足，或 segment 预检不足） ``` @@ -50,15 +46,15 @@ Mooncake → SimpleStorage（任一节点 mooncake/master 不可用、运行时 正常启动会打三段，排障时先看这三段： ``` -[dataplane] requested: backend=mooncake rdma_mode=auto device=mlx5_bond_0 gdr=False +[dataplane] requested: backend=mooncake rdma_mode=auto device=mlx5_bond_0 [dataplane] probe result: -[probe:] protocol=rdma device=mlx5_bond_0 gdr=True +[probe:] protocol=rdma device=mlx5_bond_0 [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_requested=false gdr_status=off +[dataplane] backend=MooncakeStore protocol=rdma device=mlx5_bond_0 ``` 第三段带 `fallback=...` 就说明发生了降级，原因直接写在里面（例如 `fallback=mooncake_unavailable:`）。 diff --git a/relax/backends/megatron/actor.py b/relax/backends/megatron/actor.py index e6fe7398c..936184f90 100644 --- a/relax/backends/megatron/actor.py +++ b/relax/backends/megatron/actor.py @@ -203,7 +203,6 @@ def _init( repatch(args) self.data_system_client = attach_tq_client( args.tq_config, - requested_gdr=getattr(args, "tq_use_gdr", False), role=role, lease_owner=self, ) diff --git a/relax/components/actor.py b/relax/components/actor.py index 3811e60ad..4aca7b8ec 100644 --- a/relax/components/actor.py +++ b/relax/components/actor.py @@ -73,7 +73,6 @@ def __init__( self.data_system_client = attach_tq_client( self.config.tq_config, - requested_gdr=getattr(self.config, "tq_use_gdr", False), role=self.role, lease_owner=self, ) diff --git a/relax/components/actor_fwd.py b/relax/components/actor_fwd.py index 1e1848610..7ba7ebb31 100644 --- a/relax/components/actor_fwd.py +++ b/relax/components/actor_fwd.py @@ -37,7 +37,6 @@ def __init__( 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, lease_owner=self, ) diff --git a/relax/components/advantages.py b/relax/components/advantages.py index b28ac5ad4..1919c2138 100644 --- a/relax/components/advantages.py +++ b/relax/components/advantages.py @@ -41,7 +41,6 @@ def __init__( self.data_system_client = attach_tq_client( self.config.tq_config, - requested_gdr=getattr(self.config, "tq_use_gdr", False), role="advantages", lease_owner=self, ) diff --git a/relax/components/critic.py b/relax/components/critic.py index 6e367f1f2..82e6c7575 100644 --- a/relax/components/critic.py +++ b/relax/components/critic.py @@ -42,7 +42,6 @@ def __init__( self.data_system_client = attach_tq_client( self.config.tq_config, - requested_gdr=getattr(self.config, "tq_use_gdr", False), role=self.role, lease_owner=self, ) diff --git a/relax/components/rollout.py b/relax/components/rollout.py index 4968161dd..952091ab1 100644 --- a/relax/components/rollout.py +++ b/relax/components/rollout.py @@ -335,7 +335,6 @@ def __init__( self.data_system_client = attach_tq_client( self.config.tq_config, - requested_gdr=getattr(self.config, "tq_use_gdr", False), role="rollout", lease_owner=self, ) diff --git a/relax/components/sft.py b/relax/components/sft.py index 0d3a98635..9cf732c2c 100644 --- a/relax/components/sft.py +++ b/relax/components/sft.py @@ -83,7 +83,6 @@ def __init__(self, healthy, pgs, num_gpus, config, role, runtime_env=None): # n self.data_system_client = attach_tq_client( self.config.tq_config, - requested_gdr=getattr(self.config, "tq_use_gdr", False), role=self.role, lease_owner=self, ) diff --git a/relax/core/controller.py b/relax/core/controller.py index 9f511dc3a..967900118 100644 --- a/relax/core/controller.py +++ b/relax/core/controller.py @@ -451,7 +451,6 @@ def _resolve_tq_backend(self, total_storage_size: int) -> dict: probe_results, requested_backend=backend, requested_device=device, - use_gdr=getattr(self.config, "tq_use_gdr", False), rdma_mode=mode, ) @@ -467,37 +466,15 @@ def _resolve_tq_backend(self, total_storage_size: int) -> dict: # 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. 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 - # 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. Log requested vs effective so the startup log alone explains the + # 7. 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}" - ) + logger.info(f"[dataplane] requested: backend={backend} rdma_mode={mode} device={device or 'auto'}") for result in probe_results: logger.info(f"[dataplane] probe result:\n{result.summary()}") if cap_error: diff --git a/relax/distributed/ray/rollout.py b/relax/distributed/ray/rollout.py index 2e33cd810..d3319e80b 100644 --- a/relax/distributed/ray/rollout.py +++ b/relax/distributed/ray/rollout.py @@ -815,7 +815,6 @@ def __init__(self, args, pg, data_source=None): self.data_system_client = attach_tq_client( self.args.tq_config, - requested_gdr=getattr(self.args, "tq_use_gdr", False), role="rollout_worker", lease_owner=self, ) diff --git a/relax/utils/arguments.py b/relax/utils/arguments.py index 7aaf6a3b7..c5c16f23c 100644 --- a/relax/utils/arguments.py +++ b/relax/utils/arguments.py @@ -269,8 +269,7 @@ def add_transfer_queue_arguments(parser): "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 " + "failure instead of degrading. Only effective with " "--tq-storage-backend mooncake." ), ) @@ -285,20 +284,6 @@ def add_transfer_queue_arguments(parser): "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 index f446cfc20..46eb9e8ce 100644 --- a/relax/utils/rdma_probe.py +++ b/relax/utils/rdma_probe.py @@ -60,7 +60,6 @@ class ProbeResult: checks: tuple[CheckResult, ...] effective_protocol: str | None # "rdma" | "tcp" | None effective_device: str - gdr_eligible: bool errors: tuple[str, ...] = () @property @@ -70,10 +69,7 @@ def ok(self) -> bool: 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}" - ) + header = f"[probe:{self.node}] protocol={self.effective_protocol} device={self.effective_device}" lines = [header] for c in self.checks: tag = "ok" if c.ok else "FAIL" @@ -88,18 +84,13 @@ class EffectiveConfig: 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_status = "unknown" if self.gdr else "off" dev = self.device or "auto" - base = ( - f"[dataplane] backend={self.backend} protocol={self.protocol} device={dev} " - f"gdr_requested={str(self.gdr).lower()} gdr_status={gdr_status}" - ) + base = f"[dataplane] backend={self.backend} protocol={self.protocol} device={dev}" if self.fallback_reason: return f"{base} fallback={self.fallback_reason}" return base @@ -369,20 +360,11 @@ def probe_node(device: str = "", master_address: str = "", *, probe_rdma: bool = 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), ) @@ -433,7 +415,6 @@ def _degenerate_result(node: str, error: str) -> ProbeResult: checks=tuple(), effective_protocol=None, effective_device="", - gdr_eligible=False, errors=(error,), ) @@ -516,7 +497,6 @@ def reduce_results( *, requested_backend: str, requested_device: str, - use_gdr: bool, fallback_backend: str = "SimpleStorage", rdma_mode: str = "auto", ) -> EffectiveConfig: @@ -530,8 +510,6 @@ def reduce_results( ``--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. rdma_mode @@ -545,7 +523,6 @@ def reduce_results( backend="SimpleStorage", protocol="tcp", device="", - gdr=False, fallback_reason="", ) @@ -554,7 +531,6 @@ def reduce_results( backend="SimpleStorage", protocol="tcp", device="", - gdr=False, fallback_reason="no probe results", ) @@ -574,7 +550,6 @@ def reduce_results( backend=fallback_backend, protocol="tcp", device="", - gdr=False, fallback_reason=reason, ) @@ -583,7 +558,6 @@ def reduce_results( backend="MooncakeStore", protocol="tcp", device="", - gdr=False, fallback_reason="", ) @@ -596,16 +570,12 @@ def reduce_results( 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, - # probe_node defines GDR eligibility as RDMA transport readiness; - # CUDA staging is deliberately decided and logged by each worker. - gdr=use_gdr, fallback_reason="", ) @@ -615,7 +585,6 @@ def reduce_results( backend="MooncakeStore", protocol="tcp", device=requested_device, - gdr=False, fallback_reason=f"rdma_unavailable:{','.join(rdma_failed)}", ) @@ -634,7 +603,6 @@ def validate_config(args: Any) -> list[str]: 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( @@ -642,11 +610,4 @@ def validate_config(args: Any) -> list[str]: "(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 index 83aeb5c63..0af035d55 100644 --- a/relax/utils/tq/config.py +++ b/relax/utils/tq/config.py @@ -3,9 +3,9 @@ """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. +(``--tq-storage-backend``, ``--tq-rdma-mode``, ``--tq-rdma-device``) 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 @@ -32,7 +32,6 @@ _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 @@ -163,9 +162,11 @@ def build_mooncake_config( "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, + # Host RDMA only in this phase: GDR is pinned off here rather than + # left to the upstream default so the payload always traverses host + # memory. A GDR path needs its own probe/verification and lands in + # a separate change (mooncake_client.py:75-88 for the client side). + "use_gdr": False, }, } return cfg diff --git a/relax/utils/tq/lifecycle.py b/relax/utils/tq/lifecycle.py index 00a0315ae..320163444 100644 --- a/relax/utils/tq/lifecycle.py +++ b/relax/utils/tq/lifecycle.py @@ -102,6 +102,13 @@ def _backend_signature(conf: Any) -> tuple[Any, ...]: _get_config_value(mooncake, "global_segment_size"), _get_config_value(mooncake, "local_buffer_size"), bool(_get_config_value(mooncake, "hard_pin", False)), + # Retained although Relax now always builds ``use_gdr=False``: this + # process may attach to a controller created by an older build that + # still enabled GDR, and upstream ``tq.init`` ignores the caller's + # conf when attaching (interface.py:130-135), so the worker would + # silently run the unverified GDR path. Delete together with the + # rest of the signature machinery once the exclusive-cluster + # simplification drops compatibility attach. bool(_get_config_value(mooncake, "use_gdr", False)), ) simple = _get_config_value(backend, "SimpleStorage", {}) @@ -293,42 +300,6 @@ def _close_local_tq_client() -> None: 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 _resolve_attach_timeout() -> float: """Attach deadline in seconds; override via ``RELAX_TQ_ATTACH_TIMEOUT_SECONDS``.""" @@ -407,13 +378,11 @@ def _run() -> None: 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. + """Attach a component process within a bounded deadline. The deadline covers both waiting for a served controller config and ``tq.init`` itself, because either phase can hang unboundedly (get_config @@ -435,7 +404,6 @@ def attach_tq_client( _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 @@ -501,7 +469,7 @@ def verify_cluster_attach(conf: Any, *, timeout: float | None = None) -> list[st 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") + attach_tq_client(handshake_conf, role="attach-handshake") detach_tq_client() refs: list[Any] = [] diff --git a/scripts/benchmarks/tq_cross_node_bench.py b/scripts/benchmarks/tq_cross_node_bench.py index 88ddf16ca..3d543272e 100644 --- a/scripts/benchmarks/tq_cross_node_bench.py +++ b/scripts/benchmarks/tq_cross_node_bench.py @@ -421,7 +421,7 @@ def build_conf(protocol: str, master: str, device: str, segment_gib: int): 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="") + eff = EffectiveConfig(backend="MooncakeStore", protocol=protocol, device=device, 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 6e4ac14ec..67cb1f44a 100644 --- a/scripts/benchmarks/tq_rdma_bench.py +++ b/scripts/benchmarks/tq_rdma_bench.py @@ -110,7 +110,6 @@ def build_tq_config(config_name: str, args: argparse.Namespace, num_storage_unit 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) diff --git a/tests/utils/test_rdma_probe.py b/tests/utils/test_rdma_probe.py index f673a47fb..0ec30d2b4 100644 --- a/tests/utils/test_rdma_probe.py +++ b/tests/utils/test_rdma_probe.py @@ -66,7 +66,6 @@ def _make_probe( checks=(CheckResult("mooncake_import", True),), effective_protocol=protocol, effective_device=device, - gdr_eligible=protocol == "rdma", ) @@ -75,7 +74,6 @@ def _make_args(**kwargs) -> argparse.Namespace: 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, @@ -101,16 +99,6 @@ def test_simple_backend_with_rdma_mode_rejected(self): 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) == [] @@ -119,8 +107,8 @@ 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) + def test_valid_mooncake_required(self): + args = _make_args(tq_storage_backend="mooncake", tq_rdma_mode="required") assert validate_config(args) == [] @@ -137,18 +125,15 @@ def test_simple_backend_short_circuits(self): [_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" @@ -159,7 +144,6 @@ def test_one_node_no_mooncake_falls_back(self): [_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 @@ -169,28 +153,16 @@ def test_one_node_no_rdma_degrades_to_tcp(self): [_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_request_is_forwarded_for_rdma(self): - eff = reduce_results( - [_make_probe(), _make_probe(node="node-B")], - requested_backend="mooncake", - requested_device="", - use_gdr=True, - ) - assert eff.gdr is True - assert eff.fallback_reason == "" - def test_requested_device_must_match_every_rdma_node(self): eff = reduce_results( [_make_probe(device="rdma0"), _make_probe(device="rdma1", node="node-B")], requested_backend="mooncake", requested_device="rdma0", - use_gdr=False, ) assert eff.protocol == "tcp" assert eff.fallback_reason == "device_mismatch:rdma0" @@ -200,7 +172,6 @@ def test_empty_results_falls_back(self): [], requested_backend="mooncake", requested_device="", - use_gdr=False, ) assert eff.backend == "SimpleStorage" @@ -209,7 +180,6 @@ def test_off_mode_keeps_mooncake_and_selects_tcp(self): [_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", "") @@ -220,7 +190,6 @@ def test_off_mode_reports_mooncake_unavailable(self): [_make_probe(protocol=None)], requested_backend="mooncake", requested_device="", - use_gdr=False, rdma_mode="off", ) assert eff.backend == "SimpleStorage" @@ -244,7 +213,6 @@ 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.""" @@ -273,7 +241,6 @@ 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): @@ -430,7 +397,6 @@ def test_reduce_treats_degenerate_as_no_mooncake(self): results, requested_backend="mooncake", requested_device="", - use_gdr=False, ) assert eff.backend == "SimpleStorage" assert "n1" in eff.fallback_reason @@ -441,14 +407,12 @@ def test_reduce_reports_master_unreachable_distinctly(self): 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" @@ -477,7 +441,7 @@ def test_simple_storage_config_allows_unlimited_capacity(self): 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="") + eff = EffectiveConfig(backend="MooncakeStore", protocol="rdma", device="rdma0", fallback_reason="") 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"] @@ -485,19 +449,21 @@ 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="") + eff = EffectiveConfig(backend="MooncakeStore", protocol="rdma", device="rdma0", fallback_reason="") cfg = build_mooncake_config(eff, master_address="master.example:50051") 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, master_address="master.example:50051") - assert cfg["MooncakeStore"]["use_gdr"] is True + def test_mooncake_config_pins_host_rdma(self): + """This phase ships host RDMA only: ``use_gdr`` is always False and no + GDR staging buffer is configured.""" + eff = EffectiveConfig(backend="MooncakeStore", protocol="rdma", device="", fallback_reason="") + mc = build_mooncake_config(eff, master_address="master.example:50051")["MooncakeStore"] + assert mc["use_gdr"] is False + assert "gdr_staging_buffer_mb" not in mc def test_master_address_is_required(self, monkeypatch): # A loopback default would point every node at itself in multi-node @@ -505,7 +471,7 @@ def test_master_address_is_required(self, monkeypatch): 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="") + eff = EffectiveConfig(backend="MooncakeStore", protocol="tcp", device="", fallback_reason="") with pytest.raises(RuntimeError, match="MC_MASTER_ADDRESS"): build_mooncake_config(eff) @@ -548,14 +514,14 @@ def test_contract_accepts_explicit_memcpy_disable(self, monkeypatch): 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="") + eff = EffectiveConfig(backend="MooncakeStore", protocol="rdma", device="", 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="") + eff = EffectiveConfig(backend="MooncakeStore", protocol="rdma", device="", fallback_reason="") err = validate_segment_capacity(args, eff) assert err is not None assert "insufficient" in err.lower() @@ -601,7 +567,7 @@ def test_dynamic_partial_rollout_capacity_rejects_oversampling_peak(self): use_dynamic_global_batch_size=True, over_sampling_batch_size=64, ) - eff = EffectiveConfig(backend="MooncakeStore", protocol="rdma", device="", gdr=False, fallback_reason="") + eff = EffectiveConfig(backend="MooncakeStore", protocol="rdma", device="", fallback_reason="") err = validate_segment_capacity(args, eff) assert err is not None assert "effective_batch=64" in err @@ -617,7 +583,7 @@ def test_segment_capacity_multimodal_staleness_no_longer_passes(self): 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="") + eff = EffectiveConfig(backend="MooncakeStore", protocol="rdma", device="", fallback_reason="") err = validate_segment_capacity(args, eff) assert err is not None and "RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB" in err @@ -625,7 +591,7 @@ 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="") + eff = EffectiveConfig(backend="MooncakeStore", protocol="rdma", device="", fallback_reason="") monkeypatch.setenv("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB", "8") assert validate_segment_capacity(args, eff) is None diff --git a/tests/utils/test_tq_failure_paths.py b/tests/utils/test_tq_failure_paths.py index 16db01c37..02724df1a 100644 --- a/tests/utils/test_tq_failure_paths.py +++ b/tests/utils/test_tq_failure_paths.py @@ -75,7 +75,6 @@ def _probe(node: str, protocol: str | None = "rdma", device: str = "rdma0") -> P checks=(), effective_protocol=protocol, effective_device=device if protocol else "", - gdr_eligible=protocol == "rdma", errors=() if protocol else ("mooncake not importable",), ) @@ -463,12 +462,11 @@ def test_stale_generation_does_not_close_successor(self, monkeypatch): 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 tq_lifecycle.attach_tq_client({}, role="old", lease_owner=old_owner) is client + assert tq_lifecycle.attach_tq_client({}, role="new", lease_owner=new_owner) is client assert new_owner._tq_client_generation > old_owner._tq_client_generation calls = [] @@ -501,43 +499,6 @@ def test_component_del_without_client_is_noop(self, monkeypatch): assert calls == [] -# --------------------------------------------------------------------------- -# 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 # --------------------------------------------------------------------------- @@ -549,7 +510,7 @@ def _conf(backend: str) -> dict: return {"controller": {}, "backend": {"storage_backend": backend}} @staticmethod - def _mooncake_conf(protocol: str) -> dict: + def _mooncake_conf(protocol: str, *, use_gdr: bool = False) -> dict: return { "controller": {}, "backend": { @@ -558,6 +519,7 @@ def _mooncake_conf(protocol: str) -> dict: "protocol": protocol, "master_server_address": "master.invalid:50051", "hard_pin": True, + "use_gdr": use_gdr, }, }, } @@ -626,6 +588,20 @@ def test_attach_accepts_matching_mooncake_config(self, monkeypatch): assert result.owns_controller is False assert calls["attempts"] == [] + def test_attach_rejects_legacy_gdr_controller(self, monkeypatch): + """A controller left behind by a GDR-enabled build is not compatible. + + Relax now always requests host RDMA (``use_gdr=False``), but upstream + ``tq.init`` ignores the caller's conf when attaching, so accepting such + a controller would silently run this worker on the unverified GDR path. + """ + requested = self._mooncake_conf("rdma") + stored = self._mooncake_conf("rdma", use_gdr=True) + 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="required") + assert calls["attempts"] == [] + def test_attach_rejects_different_polling_mode(self, monkeypatch): requested = self._conf("SimpleStorage") requested["controller"]["polling_mode"] = True @@ -1009,9 +985,7 @@ 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 - ) + eff = reduce_results([_probe("a"), _probe("b")], requested_backend="mooncake", requested_device="") assert (eff.backend, eff.protocol, eff.fallback_reason) == ("MooncakeStore", "rdma", "") def test_one_node_without_mooncake_degrades_whole_job(self): @@ -1020,7 +994,6 @@ def test_one_node_without_mooncake_degrades_whole_job(self): [_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 @@ -1032,9 +1005,7 @@ def test_crashed_probe_task_degrades_whole_job(self): 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 - ) + eff = reduce_results([_probe("a"), degenerate], requested_backend="mooncake", requested_device="") assert eff.backend == "SimpleStorage" def test_one_node_tcp_only_degrades_transport_not_backend(self): @@ -1042,7 +1013,6 @@ def test_one_node_tcp_only_degrades_transport_not_backend(self): [_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 From ea6ee65be4a5d425c9dc06350699098a07856cf4 Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Sun, 23 Aug 2026 16:43:29 +0800 Subject: [PATCH 21/30] refactor(tq): narrow host RDMA configuration MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # ♻️ Refactor ## Reduce public data-plane choices - Replace the backend and transport matrix with a three-state host RDMA mode - Route automatic capability failures directly to SimpleStorage - Keep Mooncake TCP confined to benchmark baselines --- # ✅ Tests ## Cover production backend decisions - Verify off, auto, and required mode behavior - Ensure the production resolver never selects Mooncake TCP - Validate the narrowed command-line interface on CPU-only environments --- # 📝 Documentation ## Describe the two-path production contract - Document direct fallback from host RDMA to SimpleStorage - Clarify master configuration behavior and benchmark-only TCP usage --- docs/draft/transfer_queue_rdma.md | 42 ++-- relax/core/controller.py | 85 ++++---- relax/utils/arguments.py | 22 +-- relax/utils/rdma_probe.py | 165 ++++++---------- relax/utils/tq/config.py | 2 +- tests/core/test_controller_tq_backend.py | 240 +++++++++++++++++++++++ tests/utils/test_rdma_probe.py | 113 ++++------- tests/utils/test_tq_failure_paths.py | 20 +- 8 files changed, 414 insertions(+), 275 deletions(-) create mode 100644 tests/core/test_controller_tq_backend.py diff --git a/docs/draft/transfer_queue_rdma.md b/docs/draft/transfer_queue_rdma.md index 415128471..8d017dbd5 100644 --- a/docs/draft/transfer_queue_rdma.md +++ b/docs/draft/transfer_queue_rdma.md @@ -8,14 +8,15 @@ Relax 的数据面（rollout ↔ train 之间的样本传输）默认走 Transfe ## 配置入口 -只暴露三个表达使用意图的参数，Mooncake 底层参数（endpoint、buffer、segment、timeout、master 策略）不做 CLI，走内部默认与部署环境。 +只暴露两个表达使用意图的参数，Mooncake 底层参数（endpoint、buffer、segment、timeout、master 策略）不做 CLI，走内部默认与部署环境。 | 参数 | 取值 | 说明 | |---|---|---| -| `--tq-storage-backend` | `simple`（默认）/ `mooncake` | `simple` 保留接入前的存储与 controller 所有权语义，并共享新增的有界 attach/失败清理 | -| `--tq-rdma-mode` | `off`（默认）/ `auto` / `required` | `off` 即使有硬件也不用 RDMA；`auto` 探测失败自动降级；`required` 探测失败直接报错退出 | +| `--tq-rdma-mode` | `off`（默认）/ `auto` / `required` | `off` 使用原有 SimpleStorage，保留接入前的存储与 controller 所有权语义；`auto` 尝试 host RDMA，不可用则回退 SimpleStorage；`required` 不可用则直接报错退出 | | `--tq-rdma-device` | 设备名，如 `mlx5_bond_0`；空为自动 | 多网卡机器上自动选择可能选错，跨节点时建议显式指定 | +生效路径只有两种：**MooncakeStore/host-RDMA** 和 **SimpleStorage**。Mooncake/TCP 不是公开的生产配置，只作为 benchmark 的 C1 对照存在，因此 `auto` 不经过 TCP 中间档，直接回退 SimpleStorage。 + `--tq-rdma-mode=required` 覆盖的是**传输层**：MooncakeStore + RDMA 可用性与 segment 容量。 ### 首期只交付 host RDMA @@ -24,21 +25,22 @@ Relax 的数据面（rollout ↔ train 之间的样本传输）默认走 Transfe ## 启动流程与降级 -driver 在**第一次 `tq.init` 之前**完成探测并生成 job 级唯一的 effective config，其余组件（actor / critic / rollout / sft / advantages / actor_fwd）都读同一份，不各自决策。 +driver 在**第一次 `tq.init` 之前**完成探测并生成 job 级唯一的 effective config，其余组件（actor / critic / rollout / sft / advantages / actor_fwd）都读同一份，不各自决策。`off` 直接短路到 SimpleStorage，不做任何探测。 -1. 校验参数组合（例如 `simple` + `rdma-mode` 会被拒绝） +1. 校验 `--tq-rdma-mode` 取值 2. `probe_cluster_nodes()` 通过 Ray 把探测任务绑定到每个**存活且有 GPU** 的节点，并额外探测 driver（driver 也会创建 Mooncake owner client）；各节点读取本机 `/sys`、mooncake 状态，并在 2 秒上限内检查外部 master 的 TCP 可达性；超时或崩溃的节点转为退化结果，不静默丢弃 -3. `reduce_results()` 做 AND 归约：整个作业只能跑在最低共同能力上 -4. Mooncake 生效前，driver 在**每个存活节点**（不限 GPU，因为 Serve replica 与 0-CPU actor 没有 placement 绑定）用真实配置各跑一次**有界 attach 握手**并立即 detach；`auto` 下任一节点失败则统一关闭 Mooncake 状态、全作业收敛到 SimpleStorage，`off`/`required` 下启动失败并列出失败节点 +3. `reduce_results()` 做 AND 归约：只有所有节点都具备 host RDMA 才选 MooncakeStore，否则整个作业回退 SimpleStorage +4. Mooncake 生效前，driver 在**每个存活节点**（不限 GPU，因为 Serve replica 与 0-CPU actor 没有 placement 绑定）用真实配置各跑一次**有界 attach 握手**并立即 detach；`auto` 下任一节点失败则统一关闭 Mooncake 状态、全作业收敛到 SimpleStorage，`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` 时必须把这份每节点启动资源足迹一并纳入容量规划。 +第 4 步不是轻量探针：每个节点都会创建真实 Mooncake client，并按配置请求挂载/注册完整 client segment（默认 `global_segment_size=4 GiB`，另有默认 1 GiB local buffer），完成后立即释放。具体物理 RSS、锁页与注册方式取决于 Mooncake 实现，但启动阶段会出现节点级瞬时内存/注册资源尖峰；CPU-only head 也在覆盖范围内。节点内存或 `memlock` 不足会表现为 attach 握手失败：`auto` 下整个作业统一回退 SimpleStorage，`required` 下启动失败。调大 `RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB` 时必须把这份每节点启动资源足迹一并纳入容量规划。 -降级阶梯： +回退只有一档： ``` -RDMA → Mooncake/TCP （任一节点无 RDMA 能力，或指定设备缺失） -Mooncake → SimpleStorage（任一节点 mooncake/master 不可用、运行时正确性契约不满足，或 segment 预检不足） +MooncakeStore/host-RDMA → SimpleStorage +（任一节点无 RDMA 能力或指定设备缺失、mooncake/master 不可用、 + 运行时正确性契约不满足，或 segment 容量预检不足） ``` ## 启动日志怎么读 @@ -46,7 +48,7 @@ Mooncake → SimpleStorage（任一节点 mooncake/master 不可用、运行时 正常启动会打三段，排障时先看这三段： ``` -[dataplane] requested: backend=mooncake rdma_mode=auto device=mlx5_bond_0 +[dataplane] requested: rdma_mode=auto device=mlx5_bond_0 [dataplane] probe result: [probe:] protocol=rdma device=mlx5_bond_0 [ok] mooncake_import: version=0.3.10.post2 @@ -69,9 +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`。该变量是必填项：未设置时启动直接失败，Relax 不会假定 loopback 端点（多节点作业里每个节点都把自己的 localhost 当 master 会导致误降级或误中止）。 - -Mooncake/TCP 长会话还应在 driver 与所有 Ray worker 中统一设置 `MC_TCP_ENABLE_CONNECTION_POOL=1`。未启用连接池时，大批量反复传输可能耗尽临时 TCP 端口并报 `Cannot assign requested address`；该变量必须通过作业运行时环境传播到所有节点，不能只在提交命令所在的 shell 中设置。RDMA transport 不依赖此选项。 +然后给**每个节点**的作业环境设置 `MC_MASTER_ADDRESS=:50051`。Relax 不会假定 loopback 端点（多节点作业里每个节点都把自己的 localhost 当 master 会导致误降级或误中止），因此未设置或格式非法时：`auto` 记 WARNING 并回退 SimpleStorage，`required` 启动失败。 启动前置条件：部署侧必须先启动 master，所有 GPU 节点和 driver 都能解析并连接 `MC_MASTER_ADDRESS`，防火墙允许 master RPC 端口；作业镜像中的 TQ 必须包含本文“正确性依赖”所列修复。Relax 不负责拉起、重启或终止 master。 @@ -79,9 +79,10 @@ Mooncake/TCP 长会话还应在 driver 与所有 Ray worker 中统一设置 `MC_ | 情形 | 表现 | 处理 | |---|---|---| +| **`MC_MASTER_ADDRESS` 未配置或格式非法** | `auto` 记 WARNING 后回退 SimpleStorage（不做探测）；`required` 启动失败 | 给每个节点的作业环境设置 `MC_MASTER_ADDRESS=:` | | **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` 启动失败并列出节点。握手使用一次性 Ray worker，超时的 `tq.init` watchdog thread 不会污染后续任务 | 检查失败节点到 master 的连通性、RDMA 状态、可用内存与 `memlock`；握手会瞬时创建完整 client segment，CPU-only head 也会执行 | +| **attach 握手在某节点失败/超时** | driver 汇总各节点结果：`auto` 关闭 Mooncake 状态并统一回退 SimpleStorage（日志 `attach_handshake_failed:*`）；`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` | @@ -179,12 +180,13 @@ 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 | +| `backend=SimpleStorage fallback=rdma_unavailable:`，但机器有 RDMA 卡 | 端口非 ACTIVE、GID 取不到、`memlock` 过低，或指定的 `--tq-rdma-device` 在部分节点不存在 | 看 `probe result` 里哪一项 FAIL；`memlock` 需要 unlimited | +| `backend=SimpleStorage fallback=...`，但预期跑 RDMA，且日志显示 master 未配置 | `auto` 下 `MC_MASTER_ADDRESS` 缺失或格式非法会记 WARNING 后回退 SimpleStorage（`required` 直接失败） | 给每个节点的作业环境设置 `MC_MASTER_ADDRESS=:` | | `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` 会被启动拒绝，unset 即可 | -| Mooncake/TCP 档在单机回环下原生 SIGSEGV | 与上一行同源（memcpy 路径），回环下表现为崩溃而非静默截断 | 同上 | +| Mooncake/TCP 档（仅 benchmark C1）get 数据尾部全零，但批量返回码全部成功（逐字节校验 FAIL） | mooncake 0.3.10 memcpy 快拷贝路径缺陷：TCP-only 环境被自动启用后，跨节点 get 会静默截断（坏行自 64 KiB 对齐偏移起全零） | 正确性守卫已强制 `MC_STORE_MEMCPY=0`（见“容量不足与正确性依赖”）；显式设 `1` 会被启动拒绝，unset 即可 | +| Mooncake/TCP 档（仅 benchmark C1）在单机回环下原生 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` 是否打出 | @@ -213,6 +215,6 @@ RDMA 生效时前者按 payload 增长、后者基本不动；反之则说明落 - 写侧（put）收益明显，读侧（get）收益有限：get 每次调用都会注册/注销 MR，且 key 粒度是 `样本 × 字段`，碎片化开销盖过了传输收益。MR 常驻注册与读路径零拷贝成型不在首期范围。 - 跨节点 RDMA vs TCP 的收益在多轮之间波动较大，验收结论应基于多轮分布而非单轮数据。 -- Mooncake/TCP（C1）在 0.3.10 上依赖 `MC_STORE_MEMCPY=0` 守卫保证字节正确，且守卫后 get 吞吐低于 SimpleStorage：TCP 档定位为 RDMA 不可用时的正确性兜底，不是性能选项。 +- Mooncake/TCP（C1）在 0.3.10 上依赖 `MC_STORE_MEMCPY=0` 守卫保证字节正确，且守卫后 get 吞吐低于 SimpleStorage。它只是 benchmark 的对照档，不是生产路径：RDMA 不可用时生产回退的是 SimpleStorage。 - 消费端节点在 get 过程中中途死亡的端到端行为需要双节点真机验证，未做成自动化测试。 - Mooncake 传输层自身的超时参数不由 Relax 控制。 diff --git a/relax/core/controller.py b/relax/core/controller.py index 967900118..ff6c05d3e 100644 --- a/relax/core/controller.py +++ b/relax/core/controller.py @@ -310,7 +310,7 @@ def _initialize_data_system(self): flags={"allow_objects": True}, ) - if getattr(self.config, "tq_storage_backend", "simple") == "simple": + if getattr(self.config, "tq_rdma_mode", "off") == "off": # 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 @@ -396,24 +396,21 @@ 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. - ``--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. + ``--tq-rdma-mode=off`` retains the previous SimpleStorage and ownership + semantics while sharing the bounded worker-attach and failure-cleanup + hardening. ``auto``/``required`` run the RDMA capability probe + *before* ``tq.init`` and emit the startup log line; ``auto`` falls back + to SimpleStorage, ``required`` fails fast. """ - # 1. Validate flag combinations (structural, before any probe). + # 1. Validate the requested mode (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). - # ``mooncake + off`` is MooncakeStore/TCP, not SimpleStorage. - if backend == "simple": + def _simple_storage() -> dict: from relax.utils.tq.config import build_simple_storage_config return build_simple_storage_config( @@ -421,7 +418,23 @@ def _resolve_tq_backend(self, total_storage_size: int) -> dict: num_data_storage_units=self.config.num_data_storage_units, ) - # 3. MooncakeStore path: probe → reduce → effective config. + def _fall_back_or_raise(reason: str, error: Exception) -> dict: + """Every host-RDMA precondition failure funnels through here. + + ``auto`` converges the whole job on SimpleStorage; ``required`` re- + raises so an operator who demanded RDMA never runs silently + downgraded. + """ + if mode != "auto": + raise RuntimeError(f"--tq-rdma-mode={mode} but {reason}: {error}") from error + logger.warning(f"[dataplane] {reason}; auto fallback to SimpleStorage: {error}") + return _simple_storage() + + # 2. SimpleStorage short-circuit (default, zero behavior change). + if mode == "off": + return _simple_storage() + + # 3. Host-RDMA 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. @@ -429,37 +442,25 @@ def _resolve_tq_backend(self, total_storage_size: int) -> dict: try: validate_mooncake_runtime_contract() except RuntimeError as e: - if mode != "auto": - raise RuntimeError( - 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 - - logger.warning( - "[dataplane] Installed TransferQueue does not satisfy the Mooncake " - f"correctness contract; auto fallback to SimpleStorage: {e}" + return _fall_back_or_raise( + "the installed TransferQueue does not satisfy the Mooncake correctness contract", 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, probe_rdma=mode != "off") - - effective = reduce_results( - probe_results, - requested_backend=backend, - requested_device=device, - rdma_mode=mode, - ) + try: + master_address = resolve_mooncake_master_address() + except RuntimeError as e: + # A missing/invalid MC_MASTER_ADDRESS makes MooncakeStore + # unreachable for the whole job, so it degrades like any other + # unmet precondition rather than aborting an ``auto`` run. + return _fall_back_or_raise("the Mooncake master endpoint is not configured", e) + probe_results = probe_cluster_nodes(device, master_address) + + effective = reduce_results(probe_results, requested_device=device) - # 4. Only auto mode may degrade. ``off`` explicitly requests - # Mooncake/TCP, while ``required`` explicitly requires RDMA. + # 4. Only auto mode may fall back; ``required`` demands host RDMA. if mode != "auto" and effective.fallback_reason: detail = "\n".join(r.summary() for r in probe_results) raise RuntimeError( - f"--tq-rdma-mode={mode} but the requested Mooncake path is unavailable: " + f"--tq-rdma-mode={mode} but host RDMA is unavailable: " f"{effective.fallback_reason}.\n" f"Probe details:\n{detail}" ) @@ -467,19 +468,19 @@ def _resolve_tq_backend(self, total_storage_size: int) -> dict: # 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) - # 6. Capacity fallback is also auto-only. Explicit Mooncake/TCP and - # required-RDMA requests fail instead of silently changing backend. + # 6. Capacity fallback is also auto-only; ``required`` fails 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. 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'}") + logger.info(f"[dataplane] requested: rdma_mode={mode} device={device or 'auto'}") 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)") + logger.info("[dataplane] backend=SimpleStorage fallback=segment_capacity_insufficient") else: logger.info(effective.log_line()) return backend_dict diff --git a/relax/utils/arguments.py b/relax/utils/arguments.py index c5c16f23c..5aaae88df 100644 --- a/relax/utils/arguments.py +++ b/relax/utils/arguments.py @@ -249,28 +249,16 @@ def add_transfer_queue_arguments(parser): # 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. Only effective with " - "--tq-storage-backend mooncake." + "TransferQueue data-plane transport. 'off' (default) uses " + "SimpleStorage/ZMQ and is equivalent to current behavior. " + "'auto' attempts MooncakeStore over host RDMA and falls back " + "to SimpleStorage when it is unavailable (with a WARNING). " + "'required' fails fast instead of falling back." ), ) parser.add_argument( diff --git a/relax/utils/rdma_probe.py b/relax/utils/rdma_probe.py index 46eb9e8ce..ce2a0554a 100644 --- a/relax/utils/rdma_probe.py +++ b/relax/utils/rdma_probe.py @@ -1,7 +1,6 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. -"""RDMA capability probe and graded-degradation state machine for -MooncakeStore. +"""RDMA capability probe and fallback decision for MooncakeStore. This module runs **before** ``tq.init`` to decide the job-level effective ``{backend, protocol, device}`` triple. The probe is intentionally side-effect @@ -9,6 +8,10 @@ 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. +Only two outcomes exist: MooncakeStore over host RDMA, or the original +SimpleStorage. Mooncake/TCP survives as a benchmark baseline only, so there is +no intermediate transport to degrade through. + 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 @@ -29,6 +32,9 @@ logger = get_logger(__name__) +# Accepted ``--tq-rdma-mode`` values; ``off`` keeps the SimpleStorage path. +TQ_RDMA_MODES = frozenset({"off", "auto", "required"}) + # --------------------------------------------------------------------------- # Data structures # --------------------------------------------------------------------------- @@ -87,10 +93,16 @@ class EffectiveConfig: fallback_reason: str # "" if no fallback occurred def log_line(self) -> str: - """Return the single-line startup log string for this effective - config.""" - dev = self.device or "auto" - base = f"[dataplane] backend={self.backend} protocol={self.protocol} device={dev}" + """Return the single-line startup log string for this effective config. + + SimpleStorage deliberately reports no protocol: naming one would imply + the data plane went through a Mooncake transport, and Mooncake/TCP is + not a production path. + """ + if self.backend != "MooncakeStore": + base = f"[dataplane] backend={self.backend}" + else: + base = f"[dataplane] backend={self.backend} protocol={self.protocol} device={self.device or 'auto'}" if self.fallback_reason: return f"{base} fallback={self.fallback_reason}" return base @@ -291,7 +303,7 @@ def _check_health_check() -> CheckResult: # pragma: no cover - retained for ad- # 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``. -def probe_node(device: str = "", master_address: str = "", *, probe_rdma: bool = True) -> ProbeResult: +def probe_node(device: str = "", master_address: str = "") -> ProbeResult: """Run all capability checks on the current node. Parameters @@ -299,20 +311,14 @@ def probe_node(device: str = "", master_address: str = "", *, probe_rdma: bool = device 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 - explicitly requested Mooncake/TCP backend never depends on RDMA - hardware. """ node = socket.gethostname() checks: list[CheckResult] = [] errors: list[str] = [] checks.append(_check_mooncake_import()) - if probe_rdma: - checks.append(_check_rdma_devices()) - checks.append(_check_memlock()) + 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 @@ -320,11 +326,8 @@ def probe_node(device: str = "", master_address: str = "", *, probe_rdma: bool = if master_address: checks.append(_check_master_reachable(master_address)) - # Device-dependent checks are irrelevant when RDMA is explicitly off. - selected_device = "" - if probe_rdma: - selected_device, device_checks = _select_usable_rdma_device(device) - checks.extend(device_checks) + 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) @@ -335,21 +338,22 @@ def probe_node(device: str = "", master_address: str = "", *, probe_rdma: bool = 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 probe_rdma else "" + 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 not probe_rdma: - effective_protocol = "tcp" elif rdma_dev_ok and port_ok and gid_ok and memlock_ok: effective_protocol = "rdma" # 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). + # Mooncake is usable but this node cannot do RDMA. ``tcp`` is recorded + # so the reduction can distinguish "no RDMA here" from "no Mooncake at + # all" in its fallback reason; Mooncake/TCP is not a production + # transport (it survives only as a benchmark baseline). effective_protocol = "tcp" if not rdma_dev_ok: errors.append("no RDMA device") @@ -424,7 +428,6 @@ 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. @@ -445,9 +448,7 @@ 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) if probe_rdma else probe_node(device, master_address, probe_rdma=False) - ) + driver_result = probe_node(device, master_address) if not node_ids: logger.debug("No alive GPU nodes discovered; probing driver node only.") return [driver_result] @@ -456,16 +457,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, should_probe_rdma: bool) -> ProbeResult: + def _probe_on_node(dev: str, master: str) -> ProbeResult: from relax.utils.rdma_probe import probe_node as _probe - return _probe(dev, master, probe_rdma=should_probe_rdma) + 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, master_address, probe_rdma) + ref = _probe_on_node.options(scheduling_strategy=strategy).remote(device, master_address) refs.append(ref) id_by_ref[ref] = node_id @@ -495,37 +496,22 @@ def _probe_on_node(dev: str, master: str, should_probe_rdma: bool) -> ProbeResul def reduce_results( results: list[ProbeResult], *, - requested_backend: str, requested_device: str, - fallback_backend: str = "SimpleStorage", - rdma_mode: str = "auto", ) -> EffectiveConfig: """AND-reduce per-node results into a single job-level effective config. + The only two outcomes are MooncakeStore over host RDMA and the original + SimpleStorage: Mooncake/TCP is a benchmark baseline, not a production + transport, so it is never selected here and there is no intermediate rung + to degrade through. + Parameters ---------- results One :class:`ProbeResult` per data-plane node. - requested_backend - ``--tq-storage-backend`` value (``"simple"`` or ``"mooncake"``). requested_device ``--tq-rdma-device`` 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": - return EffectiveConfig( - backend="SimpleStorage", - protocol="tcp", - device="", - fallback_reason="", - ) - if not results: return EffectiveConfig( backend="SimpleStorage", @@ -535,43 +521,15 @@ 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) - - 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="", - fallback_reason=reason, - ) - - if rdma_mode == "off": - return EffectiveConfig( - backend="MooncakeStore", - protocol="tcp", - device="", - fallback_reason="", - ) - - 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 for r in results) - if not device_ok: - return EffectiveConfig( - backend="MooncakeStore", - protocol="tcp", - device=requested_device, - fallback_reason=f"device_mismatch:{requested_device}", - ) + if all(r.effective_protocol == "rdma" for r in results): + # Device: every node must expose the explicitly requested device. + if requested_device and any(r.effective_device != requested_device for r in results): + return EffectiveConfig( + backend="SimpleStorage", + protocol="tcp", + device="", + fallback_reason=f"device_mismatch:{requested_device}", + ) return EffectiveConfig( backend="MooncakeStore", protocol="rdma", @@ -579,13 +537,21 @@ def reduce_results( fallback_reason="", ) - # Some nodes can't do RDMA → degrade to TCP (still MooncakeStore). - rdma_failed = [r.node for r in results if r.effective_protocol != "rdma"] + no_mooncake_nodes = [r.node for r in results if r.effective_protocol is None] + if no_mooncake_nodes: + 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(no_mooncake_nodes)}" + ) + else: + reason = f"rdma_unavailable:{','.join(r.node for r in results if r.effective_protocol != 'rdma')}" return EffectiveConfig( - backend="MooncakeStore", + backend="SimpleStorage", protocol="tcp", - device=requested_device, - fallback_reason=f"rdma_unavailable:{','.join(rdma_failed)}", + device="", + fallback_reason=reason, ) @@ -595,19 +561,16 @@ def reduce_results( def validate_config(args: Any) -> list[str]: - """Return a list of error messages for invalid flag combinations. + """Return a list of error messages for an invalid RDMA configuration. Called at startup *before* probing. An empty list means the config is structurally valid (semantic/runtime validity is checked by the probe). + ``argparse`` already constrains the mode for CLI runs; this also covers + configs restored from a checkpoint or built programmatically. """ errors: list[str] = [] - backend = getattr(args, "tq_storage_backend", "simple") mode = getattr(args, "tq_rdma_mode", "off") - 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 mode not in TQ_RDMA_MODES: + errors.append(f"--tq-rdma-mode={mode!r} must be one of {', '.join(sorted(TQ_RDMA_MODES))}.") return errors diff --git a/relax/utils/tq/config.py b/relax/utils/tq/config.py index 0af035d55..4b4e6442c 100644 --- a/relax/utils/tq/config.py +++ b/relax/utils/tq/config.py @@ -3,7 +3,7 @@ """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``) plus an +(``--tq-rdma-mode``, ``--tq-rdma-device``) plus an :class:`~relax.utils.rdma_probe.EffectiveConfig` into the OmegaConf dict that ``tq.init`` expects. diff --git a/tests/core/test_controller_tq_backend.py b/tests/core/test_controller_tq_backend.py new file mode 100644 index 000000000..a17818b3f --- /dev/null +++ b/tests/core/test_controller_tq_backend.py @@ -0,0 +1,240 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Controller._resolve_tq_backend: the production backend decision. + +``reduce_results`` is covered in ``tests/utils/test_rdma_probe.py``; these tests +pin the Controller branch that consumes it, because that is where ``off`` +short-circuits, where ``auto`` must converge on SimpleStorage for *every* unmet +host-RDMA precondition, and where ``required`` must fail fast instead. + +Mooncake/TCP is a benchmark baseline only, so no input may ever make this method +return a MooncakeStore config whose protocol is not ``rdma``. +""" + +from __future__ import annotations + +import argparse +from types import SimpleNamespace + +import pytest + +from relax.utils.rdma_probe import ProbeResult +from tests.core.test_controller_s3_model_cleanup import controller +from tests.utils.test_arguments_opd_teacher_colocate import ( + arguments_module as _arguments_module_fixture, +) + + +# ``relax.utils.arguments`` pulls in the full SGLang server args, which the +# GitHub CPU CI does not install (it ships only ``sglang-router``). Reuse the +# stub fixture so the CLI assertions build a real parser without that +# dependency. +arguments_module = _arguments_module_fixture + + +def _config(**overrides) -> SimpleNamespace: + """A config whose worst-case payload fits the default 4 GiB segment.""" + defaults = dict( + tq_rdma_mode="auto", + tq_rdma_device="", + num_data_storage_units=1, + n_samples_per_prompt=1, + rollout_batch_size=8, + seq_length=8192, + multimodal_keys=None, + max_staleness=0, + partial_rollout=False, + use_dynamic_global_batch_size=False, + ) + defaults.update(overrides) + return SimpleNamespace(**defaults) + + +def _probe(node: str, protocol: str | None = "rdma", device: str = "") -> ProbeResult: + return ProbeResult( + node=node, + checks=(), + effective_protocol=protocol, + effective_device=device, + errors=() if protocol else ("mooncake not importable",), + ) + + +class _Recorder: + """Records which host-RDMA preconditions the Controller actually + consulted.""" + + def __init__(self, monkeypatch, *, contract_error=None, master_error=None, probes=None): + self.calls: list[str] = [] + self._probes = probes if probes is not None else [_probe("node-A")] + # ``build_mooncake_config`` re-resolves the endpoint from the + # environment, so the success paths need a real value there too. + monkeypatch.setenv("MC_MASTER_ADDRESS", "master.invalid:50051") + + def fake_contract() -> None: + self.calls.append("contract") + if contract_error is not None: + raise contract_error + + def fake_master() -> str: + self.calls.append("master") + if master_error is not None: + raise master_error + return "master.invalid:50051" + + def fake_probe(device: str, master_address: str) -> list[ProbeResult]: + self.calls.append("probe") + return self._probes + + monkeypatch.setattr(controller, "validate_mooncake_runtime_contract", fake_contract) + monkeypatch.setattr(controller, "resolve_mooncake_master_address", fake_master) + monkeypatch.setattr(controller, "probe_cluster_nodes", fake_probe) + + +def _resolve(config) -> dict: + instance = controller.Controller.__new__(controller.Controller) + instance.config = config + return instance._resolve_tq_backend(total_storage_size=64) + + +def _assert_simple_storage(backend: dict) -> None: + assert backend["storage_backend"] == "SimpleStorage" + assert "MooncakeStore" not in backend + + +class TestOffMode: + """``off`` is the untouched SimpleStorage path: it probes nothing.""" + + def test_off_short_circuits_without_any_precondition_check(self, monkeypatch): + recorder = _Recorder(monkeypatch) + backend = _resolve(_config(tq_rdma_mode="off")) + _assert_simple_storage(backend) + assert recorder.calls == [] + + def test_missing_mode_attribute_defaults_to_off(self, monkeypatch): + recorder = _Recorder(monkeypatch) + config = _config() + del config.tq_rdma_mode + _assert_simple_storage(_resolve(config)) + assert recorder.calls == [] + + +class TestAutoFallsBackForEveryUnmetPrecondition: + """Gate A: in ``auto``, anything short of host RDMA yields + SimpleStorage.""" + + def test_contract_failure_falls_back_before_touching_master(self, monkeypatch): + recorder = _Recorder(monkeypatch, contract_error=RuntimeError("retry guard missing")) + _assert_simple_storage(_resolve(_config())) + assert recorder.calls == ["contract"] + + def test_missing_master_endpoint_falls_back_before_probing(self, monkeypatch): + recorder = _Recorder(monkeypatch, master_error=RuntimeError("MC_MASTER_ADDRESS required")) + _assert_simple_storage(_resolve(_config())) + assert recorder.calls == ["contract", "master"] + + def test_node_without_rdma_falls_back(self, monkeypatch): + _Recorder(monkeypatch, probes=[_probe("node-A"), _probe("node-B", protocol="tcp")]) + _assert_simple_storage(_resolve(_config())) + + def test_node_without_mooncake_falls_back(self, monkeypatch): + _Recorder(monkeypatch, probes=[_probe("node-A"), _probe("node-B", protocol=None)]) + _assert_simple_storage(_resolve(_config())) + + def test_device_mismatch_falls_back(self, monkeypatch): + _Recorder(monkeypatch, probes=[_probe("node-A", device="rdma0"), _probe("node-B", device="rdma1")]) + _assert_simple_storage(_resolve(_config(tq_rdma_device="rdma0"))) + + def test_insufficient_segment_capacity_falls_back(self, monkeypatch): + """Worst-case multimodal payload far exceeds the default segment.""" + _Recorder(monkeypatch) + config = _config(multimodal_keys=["pixel_values"], rollout_batch_size=64, n_samples_per_prompt=8) + _assert_simple_storage(_resolve(config)) + + def test_all_nodes_rdma_selects_mooncake(self, monkeypatch): + _Recorder(monkeypatch, probes=[_probe("node-A"), _probe("node-B")]) + backend = _resolve(_config()) + assert backend["storage_backend"] == "MooncakeStore" + assert backend["MooncakeStore"]["protocol"] == "rdma" + + +class TestRequiredFailsFast: + """``required`` must never silently downgrade the same failures.""" + + @pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"contract_error": RuntimeError("retry guard missing")}, "correctness contract"), + ({"master_error": RuntimeError("MC_MASTER_ADDRESS required")}, "master endpoint is not configured"), + ({"probes": [_probe("node-A"), _probe("node-B", protocol="tcp")]}, "host RDMA is unavailable"), + ({"probes": [_probe("node-A"), _probe("node-B", protocol=None)]}, "host RDMA is unavailable"), + ], + ) + def test_required_raises(self, monkeypatch, kwargs, match): + _Recorder(monkeypatch, **kwargs) + with pytest.raises(RuntimeError, match=match): + _resolve(_config(tq_rdma_mode="required")) + + def test_required_raises_on_insufficient_capacity(self, monkeypatch): + _Recorder(monkeypatch) + config = _config( + tq_rdma_mode="required", + multimodal_keys=["pixel_values"], + rollout_batch_size=64, + n_samples_per_prompt=8, + ) + with pytest.raises(RuntimeError, match="segment capacity insufficient"): + _resolve(config) + + def test_required_accepts_full_rdma_cluster(self, monkeypatch): + _Recorder(monkeypatch, probes=[_probe("node-A"), _probe("node-B")]) + backend = _resolve(_config(tq_rdma_mode="required")) + assert backend["MooncakeStore"]["protocol"] == "rdma" + + +class TestProductionNeverSelectsMooncakeTcp: + """Mooncake/TCP exists only as benchmark C1.""" + + @pytest.mark.parametrize("mode", ["off", "auto"]) + @pytest.mark.parametrize( + "probes", + [ + [_probe("node-A", protocol="tcp")], + [_probe("node-A"), _probe("node-B", protocol="tcp")], + [_probe("node-A", protocol=None)], + [], + ], + ) + def test_no_input_produces_a_tcp_mooncake_backend(self, monkeypatch, mode, probes): + _Recorder(monkeypatch, probes=probes) + backend = _resolve(_config(tq_rdma_mode=mode)) + if backend["storage_backend"] == "MooncakeStore": + assert backend["MooncakeStore"]["protocol"] == "rdma" + else: + _assert_simple_storage(backend) + + def test_invalid_mode_is_rejected(self, monkeypatch): + _Recorder(monkeypatch) + with pytest.raises(ValueError, match="--tq-rdma-mode"): + _resolve(_config(tq_rdma_mode="mooncake")) + + +def test_cli_exposes_only_mode_and_device(arguments_module): + """The narrowed CLI keeps exactly two TransferQueue RDMA flags.""" + arguments_module.RouterArgs = SimpleNamespace(add_cli_args=lambda parser, **_kwargs: parser) + parser = argparse.ArgumentParser() + arguments_module.get_slime_extra_args_provider()(parser) + + tq_flags = sorted( + option + for action in parser._actions + for option in action.option_strings + if option.startswith("--tq-") and "timeout" not in option + ) + assert tq_flags == ["--tq-rdma-device", "--tq-rdma-mode"] + + args = parser.parse_args([]) + assert args.tq_rdma_mode == "off" + assert args.tq_rdma_device == "" + assert not hasattr(args, "tq_storage_backend") + assert not hasattr(args, "tq_use_gdr") diff --git a/tests/utils/test_rdma_probe.py b/tests/utils/test_rdma_probe.py index 0ec30d2b4..ad8bc688a 100644 --- a/tests/utils/test_rdma_probe.py +++ b/tests/utils/test_rdma_probe.py @@ -71,7 +71,6 @@ def _make_probe( def _make_args(**kwargs) -> argparse.Namespace: defaults = dict( - tq_storage_backend="mooncake", tq_rdma_mode="auto", tq_rdma_device="", num_data_storage_units=1, @@ -91,25 +90,21 @@ def _make_args(**kwargs) -> argparse.Namespace: class TestValidateConfig: - """validate_config: structural flag-combination checks before any probe.""" + """validate_config: structural mode check 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_valid_simple_off(self): - args = _make_args(tq_storage_backend="simple", tq_rdma_mode="off") - assert validate_config(args) == [] + @pytest.mark.parametrize("mode", ["off", "auto", "required"]) + def test_accepts_every_supported_mode(self, mode): + assert validate_config(_make_args(tq_rdma_mode=mode)) == [] - def test_valid_mooncake_auto(self): - args = _make_args(tq_storage_backend="mooncake", tq_rdma_mode="auto") - assert validate_config(args) == [] + def test_rejects_unknown_mode(self): + """Guards configs restored from a checkpoint or built without + argparse.""" + errors = validate_config(_make_args(tq_rdma_mode="mooncake")) + assert len(errors) == 1 + assert "--tq-rdma-mode" in errors[0] - def test_valid_mooncake_required(self): - args = _make_args(tq_storage_backend="mooncake", tq_rdma_mode="required") - assert validate_config(args) == [] + def test_missing_attribute_defaults_to_off(self): + assert validate_config(argparse.Namespace()) == [] # --------------------------------------------------------------------------- @@ -118,21 +113,15 @@ def test_valid_mooncake_required(self): class TestReduceResults: - """reduce_results: per-node ProbeResult -> job-level EffectiveConfig (AND reduction).""" + """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", - ) - assert eff.backend == "SimpleStorage" - assert eff.protocol == "tcp" + The only two outcomes are MooncakeStore/RDMA and SimpleStorage; Mooncake/TCP + is a benchmark baseline and must never be selected as a production backend. + """ 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="", ) assert eff.backend == "MooncakeStore" @@ -142,58 +131,32 @@ def test_all_nodes_rdma(self): 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="", ) assert eff.backend == "SimpleStorage" + assert "mooncake_unavailable" in eff.fallback_reason assert "node-B" in eff.fallback_reason - def test_one_node_no_rdma_degrades_to_tcp(self): + def test_one_node_without_rdma_falls_back_to_simple_not_tcp(self): eff = reduce_results( [_make_probe(protocol="rdma"), _make_probe(protocol="tcp", node="node-B")], - requested_backend="mooncake", requested_device="", ) - assert eff.backend == "MooncakeStore" - assert eff.protocol == "tcp" - assert "node-B" in eff.fallback_reason + assert (eff.backend, eff.protocol, eff.device) == ("SimpleStorage", "tcp", "") + assert eff.fallback_reason == "rdma_unavailable:node-B" def test_requested_device_must_match_every_rdma_node(self): eff = reduce_results( [_make_probe(device="rdma0"), _make_probe(device="rdma1", node="node-B")], - requested_backend="mooncake", requested_device="rdma0", ) - assert eff.protocol == "tcp" + assert eff.backend == "SimpleStorage" assert eff.fallback_reason == "device_mismatch:rdma0" def test_empty_results_falls_back(self): - eff = reduce_results( - [], - requested_backend="mooncake", - requested_device="", - ) - 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", - 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="", - rdma_mode="off", - ) + eff = reduce_results([], requested_device="") assert eff.backend == "SimpleStorage" - assert "mooncake_unavailable" in eff.fallback_reason + assert eff.fallback_reason == "no probe results" # --------------------------------------------------------------------------- @@ -332,21 +295,19 @@ 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): + def test_master_unreachable_blocks_mooncake(self, monkeypatch): + """An unreachable master is not a transport degradation: without it the + job cannot run MooncakeStore at all.""" 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")), + lambda address: CheckResult("master_reachable", False, address), ) - 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"} + result = probe_node("", "master.example:50051") + assert result.effective_protocol is None + assert result.ok is False + assert "master unreachable" in result.errors # --------------------------------------------------------------------------- @@ -393,11 +354,7 @@ def test_reduce_treats_degenerate_as_no_mooncake(self): _make_probe(protocol="rdma", node="n0"), _degenerate_result("n1", "probe_task_failed:boom"), ] - eff = reduce_results( - results, - requested_backend="mooncake", - requested_device="", - ) + eff = reduce_results(results, requested_device="") assert eff.backend == "SimpleStorage" assert "n1" in eff.fallback_reason @@ -409,11 +366,7 @@ def test_reduce_reports_master_unreachable_distinctly(self): effective_device="", errors=("master unreachable",), ) - eff = reduce_results( - [_make_probe(protocol="rdma", node="n0"), unavailable], - requested_backend="mooncake", - requested_device="", - ) + eff = reduce_results([_make_probe(protocol="rdma", node="n0"), unavailable], requested_device="") assert eff.backend == "SimpleStorage" assert eff.fallback_reason == "master_unreachable:n1" diff --git a/tests/utils/test_tq_failure_paths.py b/tests/utils/test_tq_failure_paths.py index 02724df1a..a3b208d45 100644 --- a/tests/utils/test_tq_failure_paths.py +++ b/tests/utils/test_tq_failure_paths.py @@ -985,16 +985,12 @@ 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="") + eff = reduce_results([_probe("a"), _probe("b")], requested_device="") 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="", - ) + eff = reduce_results([_probe("a"), _probe("b", protocol=None)], requested_device="") assert eff.backend == "SimpleStorage" assert "mooncake_unavailable" in eff.fallback_reason and "b" in eff.fallback_reason @@ -1005,16 +1001,12 @@ def test_crashed_probe_task_degrades_whole_job(self): degenerate = _degenerate_result("b", "probe task raised") assert degenerate.effective_protocol is None - eff = reduce_results([_probe("a"), degenerate], requested_backend="mooncake", requested_device="") + eff = reduce_results([_probe("a"), degenerate], requested_device="") 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="", - ) - assert (eff.backend, eff.protocol) == ("MooncakeStore", "tcp") + def test_one_node_without_rdma_degrades_whole_job_to_simple(self): + eff = reduce_results([_probe("a"), _probe("b", protocol="tcp")], requested_device="") + assert (eff.backend, eff.protocol) == ("SimpleStorage", "tcp") assert eff.fallback_reason From 19a4b13b573ec6106757c2c4670576314370a0e6 Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Mon, 24 Aug 2026 01:32:11 +0800 Subject: [PATCH 22/30] refactor(tq): use isolated attach for RDMA capability - remove static RDMA probing and simplify backend configuration - enforce exclusive controller ownership and fail-closed cleanup - isolate handshake and train actor initialization failures - preserve byte-exact coverage with hardened payload checks --- docs/draft/transfer_queue_rdma.md | 103 ++-- relax/components/actor.py | 12 +- relax/components/actor_fwd.py | 2 +- relax/components/critic.py | 2 +- relax/core/controller.py | 165 +++-- relax/distributed/ray/actor_group.py | 85 +++ relax/distributed/ray/train_actor.py | 10 + relax/utils/payload_digest.py | 23 +- relax/utils/rdma_probe.py | 576 ------------------ relax/utils/tq/config.py | 182 ++++-- relax/utils/tq/correctness.py | 18 +- relax/utils/tq/lifecycle.py | 390 +++++++----- scripts/benchmarks/tq_cross_node_bench.py | 10 +- scripts/benchmarks/tq_rdma_bench.py | 10 +- tests/core/test_controller_tq_backend.py | 408 ++++++++++--- tests/utils/_tq_handshake_timeout_probe.py | 6 +- .../utils/_train_actor_init_cleanup_probe.py | 71 +++ tests/utils/test_rdma_probe.py | 559 ----------------- tests/utils/test_tq_failure_paths.py | 505 ++++++++++----- tests/utils/test_train_actor_init_cleanup.py | 140 +++++ tests/utils/tq/_payload_assertions.py | 20 +- tests/utils/tq/test_config.py | 382 ++++++++++++ tests/utils/tq/test_payload_assertions.py | 25 + 23 files changed, 1996 insertions(+), 1708 deletions(-) delete mode 100644 relax/utils/rdma_probe.py create mode 100644 tests/utils/_train_actor_init_cleanup_probe.py delete mode 100644 tests/utils/test_rdma_probe.py create mode 100644 tests/utils/test_train_actor_init_cleanup.py create mode 100644 tests/utils/tq/test_config.py diff --git a/docs/draft/transfer_queue_rdma.md b/docs/draft/transfer_queue_rdma.md index 8d017dbd5..23aaa1c87 100644 --- a/docs/draft/transfer_queue_rdma.md +++ b/docs/draft/transfer_queue_rdma.md @@ -4,7 +4,9 @@ Relax 的数据面（rollout ↔ train 之间的样本传输）默认走 TransferQueue 的 SimpleStorage/ZMQ。本特性把 TransferQueue 已有的 MooncakeStore 后端接出来，使数据面可以走 RDMA，并在能力不足时安全回退。 -首期只做配置接入、能力探测与一致回退，**不改变 payload 形状与数据分发语义**。默认参数仍使用 SimpleStorage 及原有 controller 所有权模型；同时所有 worker attach（包括 SimpleStorage）新增默认 60 秒 deadline，半初始化 controller 会被回收，Controller 构造失败时会关闭本进程已经完成的 legacy `tq.init`。 +首期只做配置接入、真实 attach 能力判定与一致回退，**不改变 payload 形状与数据分发语义**。默认参数仍使用 SimpleStorage 及原有 controller 所有权模型；同时所有 worker attach（包括 SimpleStorage）新增默认 60 秒 deadline，半初始化 controller 会被回收，Controller 构造失败时会关闭本进程已经完成的 legacy `tq.init`。 + +> 与 RFC #217 的差异：RFC 要求"在首次 `tq.init` 之前完成静态节点探测"（F10）来避免挂死。按维护者的瘦身要求，静态探测已删除，防挂死改由这些机制保证：`tq.init` 在隔离的 owner actor 中执行并带超时、每次初始化前回收半初始化的 controller、拆除时等待 controller 真正从 GCS 注销、handshake 使用一次性 Ray worker（`max_calls=1`/`max_retries=0`）使超时的 `tq.init` 线程随进程退出。 ## 配置入口 @@ -25,41 +27,58 @@ Relax 的数据面（rollout ↔ train 之间的样本传输）默认走 Transfe ## 启动流程与降级 -driver 在**第一次 `tq.init` 之前**完成探测并生成 job 级唯一的 effective config，其余组件（actor / critic / rollout / sft / advantages / actor_fwd）都读同一份，不各自决策。`off` 直接短路到 SimpleStorage，不做任何探测。 +driver 在**第一次 `tq.init` 之前**完成所有配置校验，并生成 job 级唯一的 backend config，其余组件（actor / critic / rollout / sft / advantages / actor_fwd）都读同一份，不各自决策。`off` 直接短路到 SimpleStorage，不做任何 Mooncake/RDMA 检查。 + +Relax **不做静态硬件探测**：不扫描 `/sys/class/infiniband`，不读 GID 表，不推断 `memlock`。这类启发式覆盖不到真实调度位置（TQ client 跑在 Ray Serve replica 与 0-CPU actor 里，没有 placement 绑定），而且容易与 Mooncake 底层实现漂移。运行时可用性以真实 attach/setup 的结果为准；线路是否确实经过 RDMA 仍由真机 wire-proof 验收确认。 -1. 校验 `--tq-rdma-mode` 取值 -2. `probe_cluster_nodes()` 通过 Ray 把探测任务绑定到每个**存活且有 GPU** 的节点，并额外探测 driver（driver 也会创建 Mooncake owner client）；各节点读取本机 `/sys`、mooncake 状态，并在 2 秒上限内检查外部 master 的 TCP 可达性；超时或崩溃的节点转为退化结果，不静默丢弃 -3. `reduce_results()` 做 AND 归约：只有所有节点都具备 host RDMA 才选 MooncakeStore，否则整个作业回退 SimpleStorage -4. Mooncake 生效前，driver 在**每个存活节点**（不限 GPU，因为 Serve replica 与 0-CPU actor 没有 placement 绑定）用真实配置各跑一次**有界 attach 握手**并立即 detach；`auto` 下任一节点失败则统一关闭 Mooncake 状态、全作业收敛到 SimpleStorage，`required` 下启动失败并列出失败节点 -5. `required` 模式下若发生任何回退，直接抛异常并打印每个节点的探测明细 +1. 校验 `--tq-rdma-mode` 取值（非法值始终报错，不回退） +2. 校验 TransferQueue/Mooncake 正确性契约 +3. 校验 `MC_MASTER_ADDRESS` 存在且是可用的 `host:port`（只查格式，不做 DNS 解析和连通性探测） +4. segment 容量预检：按 token 预算推导最坏情况 payload，对比 client segment 大小 +5. driver 在独立 owner actor 中执行 `tq.init`，随后在**每个存活节点**调度一次性 handshake worker：真实 attach/setup → 确认 storage manager 是 `MooncakeStorageManager` 且 client 配置请求 `protocol=rdma` → detach +6. 汇总各节点结果：零存活节点、节点发现/任务调度失败或任一节点 attach 失败都按失败处理；`auto` 只有在确认所有一次性 worker 已终止后，才会**先完整拆除 Mooncake 状态**（owner close + controller 退出等待 + segment unmount），再收敛到 SimpleStorage；`required` 下启动失败并列出失败摘要 -第 4 步不是轻量探针：每个节点都会创建真实 Mooncake client，并按配置请求挂载/注册完整 client segment（默认 `global_segment_size=4 GiB`，另有默认 1 GiB local buffer），完成后立即释放。具体物理 RSS、锁页与注册方式取决于 Mooncake 实现，但启动阶段会出现节点级瞬时内存/注册资源尖峰；CPU-only head 也在覆盖范围内。节点内存或 `memlock` 不足会表现为 attach 握手失败：`auto` 下整个作业统一回退 SimpleStorage，`required` 下启动失败。调大 `RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB` 时必须把这份每节点启动资源足迹一并纳入容量规划。 +第 5 步不是轻量探针：每个节点都会创建真实 Mooncake client，并按配置请求挂载/注册完整 client segment（默认 `global_segment_size=4 GiB`，另有默认 1 GiB local buffer），完成后立即释放。具体物理 RSS、锁页与注册方式取决于 Mooncake 实现，但启动阶段会出现节点级瞬时内存/注册资源尖峰；CPU-only head 也在覆盖范围内。节点内存或 `memlock` 不足会表现为 attach 握手失败：`auto` 下整个作业统一回退 SimpleStorage，`required` 下启动失败。调大 `RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB` 时必须把这份每节点启动资源足迹一并纳入容量规划。 回退只有一档： ``` MooncakeStore/host-RDMA → SimpleStorage -（任一节点无 RDMA 能力或指定设备缺失、mooncake/master 不可用、 - 运行时正确性契约不满足，或 segment 容量预检不足） +（正确性契约不满足、master 未配置或格式非法、segment 容量预检不足， + 或全节点真实 attach/setup 未通过） ``` +### 启动耗时 + +`auto` 在 RDMA 不可用的集群上会比 `off` 慢，因为它要真的走一遍完整流程才能得出结论：owner `tq.init`（上限 60 s）→ 全节点 handshake（每节点 attach 上限 60 s，driver 侧等待上限 90 s）→ 拆除（owner close 上限 30 s，外加等待 controller 从 GCS 注销）→ SimpleStorage 初始化。最坏情况下启动阶段可能达到**分钟级**。这是用启动时延换取"判定基于真实 endpoint"的取舍；确定不需要 RDMA 的作业应显式用 `--tq-rdma-mode=off`，它不触发上述任何步骤。 + ## 启动日志怎么读 -正常启动会打三段，排障时先看这三段： +正常启动会打这几段，排障时先看它们： ``` [dataplane] requested: rdma_mode=auto device=mlx5_bond_0 -[dataplane] probe result: -[probe:] protocol=rdma device=mlx5_bond_0 - [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 +[dataplane] backend=MooncakeStore protocol=rdma device=mlx5_bond_0 (pending handshake) +[dataplane] Mooncake attach handshake passed on all alive nodes. ``` -第三段带 `fallback=...` 就说明发生了降级，原因直接写在里面（例如 `fallback=mooncake_unavailable:`）。 +`(pending handshake)` 表示配置校验已通过但能力尚未确认；只有最后一行出现才说明每个存活节点都完成了真实 attach/setup，storage manager 是 MooncakeStore，且 Mooncake client 的配置请求为 `protocol=rdma`。 + +注意这条断言的强度：它核对的是 client 的**配置意图和 setup 结果**，不是 negotiated transport，也不能单独证明数据包没有经过 TCP。线路级证明（IB counter 增长、非 RDMA 数据接口 counter 不增长）由 benchmark 的 `--require-wire-proof` 提供，见下方验收章节。 + +发生回退时会看到： + +``` +[dataplane] <失败原因>; auto fallback to SimpleStorage: <错误详情> +``` + +或 handshake 阶段失败： + +``` +[dataplane] Mooncake attach handshake reported N failure(s) (<稳定的失败类型摘要>); closing Mooncake state and converging the whole job to SimpleStorage. +``` + +`required` 模式下这两种情况都是启动失败并抛出同样的原因，不会出现回退日志。 ## Mooncake master 生命周期 @@ -73,16 +92,16 @@ setsid mooncake_master -rpc_port=50051 -metrics_port=9004 > /var/log/mooncake_ma 然后给**每个节点**的作业环境设置 `MC_MASTER_ADDRESS=:50051`。Relax 不会假定 loopback 端点（多节点作业里每个节点都把自己的 localhost 当 master 会导致误降级或误中止），因此未设置或格式非法时：`auto` 记 WARNING 并回退 SimpleStorage，`required` 启动失败。 -启动前置条件：部署侧必须先启动 master，所有 GPU 节点和 driver 都能解析并连接 `MC_MASTER_ADDRESS`，防火墙允许 master RPC 端口；作业镜像中的 TQ 必须包含本文“正确性依赖”所列修复。Relax 不负责拉起、重启或终止 master。 +启动前置条件：部署侧必须先启动 master，所有存活 Ray 节点和 driver 都能解析并连接 `MC_MASTER_ADDRESS`，防火墙允许 master RPC 端口；作业镜像中的 TQ 必须包含本文“正确性依赖”所列修复。Relax 不负责拉起、重启或终止 master。 三种情形下的行为： | 情形 | 表现 | 处理 | |---|---|---| -| **`MC_MASTER_ADDRESS` 未配置或格式非法** | `auto` 记 WARNING 后回退 SimpleStorage（不做探测）；`required` 启动失败 | 给每个节点的作业环境设置 `MC_MASTER_ADDRESS=:` | -| **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:*`）；`required` 启动失败并列出节点。握手使用一次性 Ray worker，超时的 `tq.init` watchdog thread 不会污染后续任务 | 检查失败节点到 master 的连通性、RDMA 状态、可用内存与 `memlock`；握手会瞬时创建完整 client segment，CPU-only head 也会执行 | +| **`MC_MASTER_ADDRESS` 未配置或格式非法** | `auto` 记 WARNING 后回退 SimpleStorage（不初始化任何 Mooncake 状态）；`required` 启动失败 | 给每个节点的作业环境设置 `MC_MASTER_ADDRESS=:` | +| **master 不可达** | 没有单独的可达性探测：master 连不上会表现为 owner `tq.init` 或各节点 attach 握手失败，按下面两行处理 | 先确认 master 进程、DNS/路由和防火墙 | +| **`tq.init` 失败/超时** | 第一次初始化在独立 owner actor 中执行，driver 最多等待 60 秒。失败后回收该 actor 及其拥有的半初始化 controller；`auto` 只重试一次 SimpleStorage，`required` 清理后抛出不含 endpoint/PID/路径的稳定错误类型摘要 | 查看 `mooncake_init_failed:*`、master 日志和 owner 清理日志 | +| **attach 握手在某节点失败/超时/protocol 不符** | driver 汇总各节点结果：`auto` 先完整关闭 Mooncake 状态再统一回退 SimpleStorage（日志 `attach_handshake_failed:*`）；`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` | @@ -94,24 +113,28 @@ 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.init` 之前会检查已存在的 `TransferQueueController` 命名 actor：**只有取不到 config（半初始化）、明确超时或 actor 已死时才回收**；其他 GCS/control-plane 异常只做脱敏后中止，不据此杀 actor +- 健康的既有 controller 保持不动，但启动会明确失败，**不会 attach，也不会在退出时关闭它**；操作者应先停止前一作业并清理其 TQ 状态，确保 Ray 集群干净 +- 首次初始化在专用 owner actor 中执行，config 内保存随机 owner token；若初始化失败，只有 token 仍匹配时才回收该次初始化创建的全局 actor；若并发 initializer 抢先创建了健康 controller，本次启动只做本地 detach 后失败 - 全局 `tq.close()` 只能由 owner actor 调用；actor、critic、rollout 等附加 worker 只能做本地 detach +- 任一长生命周期 train actor 初始化失败时，driver 会 force-kill 整个 train actor group，并通过预先排队的终止 probe 确认 actor task 已进入终态后再传播失败，防止超时的原生初始化线程稍后修改可复用进程的 TQ 全局状态 - master 进程始终不被 Relax 触碰 ## 容量不足与正确性依赖 -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` 调整，容量校验与客户端配置读取同一个值。该上界是保守推导，不能替代运行时错误处理。 +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` 不进入 Mooncake 配置与容量检查。segment 大小可用 `RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB` 调整，容量校验与客户端配置读取同一个值。该上界是保守推导，不能替代运行时错误处理。 + +当前运行时依赖固定到 TransferQueue commit `58054a33834aadbcf76aacd6b1e32e25c030f2c9`。Relax 在 Mooncake 启动前能检查 retry method 是否存在，以及 `KVStorageManager.put_data` 的源码顺序；这些检查**不能证明**当前 pin 已具备完整的 fail-closed 语义。 -运行时依赖固定到 TransferQueue commit `58054a33834aadbcf76aacd6b1e32e25c030f2c9`，并在 Mooncake 启动前检查以下能力： +启用 Mooncake/RDMA 前，上游 TransferQueue PR 必须同时补齐： -- `batch_upsert_from` / `batch_get_into` 对每个 key 的返回码做有限次数重试，耗尽后抛异常，不能把失败当成功或无限重试； -- `KVStorageManager.put_data` 必须先等待 storage put 成功，之后才能更新 production-ready 状态；写入失败时不通知消费者； -- Relax 的契约测试用失败 store 验证“写失败、production 状态不更新”，并用隔离 master 的真机故障注入验证物理容量溢出在 30 秒内显式失败。 +- `batch_upsert_from` / `batch_get_into` 的**每次 retry** 都校验返回结果与请求 key 等长，短结果必须显式失败； +- `NOTIFY_DATA_UPDATE_ACK` 必须检查 positive ACK，controller 拒绝更新时 producer 不得按成功返回； +- 合入修复后更新 Relax 的明确 commit pin，并以稳定的 version/capability marker 校验，而不是仅凭 method name 或源码文本推断能力。 -因此，上游曾出现的“返回码未检查导致静默丢数据”不是已知限制，而是 Mooncake 启用的硬门槛：`auto` 在契约不满足时禁用 Mooncake 并回退，`required` 拒绝启动。Docker 镜像固定上述修复 commit，运行时检查用于防止环境被旧包覆盖。 +在上述上游 PR 与新 pin 落地前，这个正确性门槛仍是明确的合入 blocker；本 PR 不通过 monkey patch 改写 TransferQueue。Relax 侧继续保留失败 store 的“写失败、production 状态不更新”契约测试，以及隔离 master 的物理容量溢出故障注入。 -正确性守卫强制 `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 只保留只读的能力校验。 +正确性守卫强制 `MC_STORE_MEMCPY=0` 且 **fail-closed**：mooncake 0.3.10 在 TCP-only 环境会自动启用 memcpy 快拷贝路径，该路径存在已确认的静默截断缺陷（现象与处置见排障表）；RDMA 会话本就自动禁用 memcpy，不受影响。由于缺陷在当前 pin 上已实证，显式导出 `MC_STORE_MEMCPY=1` 会在启动时被直接拒绝，待 pin 升级到修复版本后再按版本重新放开。 真机容量故障注入会故意创建 64 MiB segment 并写入 96 MiB，仅允许在独立、可丢弃的 master 上运行： @@ -128,11 +151,12 @@ Mock/本机测试和真实双节点 RDMA 测试必须分别报告，前者不能 | 层级 | 验证内容 | 通过标准 | |---|---|---| -| CI/mock | 参数矩阵、节点 AND 归约、master 不可达、owner 超时/清理/token、auto/required、有限重试、写失败不发布状态 | `tests/utils/test_rdma_probe.py` 与 `tests/utils/test_tq_failure_paths.py` 全部通过；真机项允许明确 skip | +| CI/mock | 模式校验、master 端点格式、容量预检、owner 超时/清理/token、attach 握手的 manager/config 契约与拆除顺序、auto/required、有限重试、写失败不发布状态 | `tests/utils/tq/test_config.py`、`tests/core/test_controller_tq_backend.py` 与 `tests/utils/test_tq_failure_paths.py` 全部通过；真机项允许明确 skip | | 本机 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 | +| 真实回退（`auto` + 某节点 RDMA 不可用） | 静态探测删除后，`auto` 会真的创建 Mooncake owner 与 named controller，再在 handshake 阶段失败并拆除，因此这条路径必须实测 | 逐条确认：① 日志显示 Mooncake owner `tq.init` **成功**、随后 handshake 阶段失败（否则实际只测到 owner 初始化失败，覆盖不到拆除）；② 最终存在且仅存在一个预期的 SimpleStorage `TransferQueueController`；③ 该 controller 的 stored config 确认为 SimpleStorage；④ SimpleStorage 的 put/get 正常工作；⑤ 旧 Mooncake owner actor 与其 client 均已消失；⑥ Mooncake segment 已从 master 卸载——若测的是强杀/超时路径，需等过 `client_ttl`（默认 30 s）再检查 | 真实多模态 fixture 生成（需要本地数据集 parquet 与 Qwen-VL 模型目录；产物写入 `tests/fixtures/`，已 gitignore，不入库）： @@ -178,9 +202,10 @@ 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 的瞬时资源足迹 | -| `backend=SimpleStorage fallback=rdma_unavailable:`，但机器有 RDMA 卡 | 端口非 ACTIVE、GID 取不到、`memlock` 过低，或指定的 `--tq-rdma-device` 在部分节点不存在 | 看 `probe result` 里哪一项 FAIL；`memlock` 需要 unlimited | +| 启动日志 `the installed TransferQueue does not satisfy the Mooncake correctness contract; auto fallback to SimpleStorage` | 镜像里的 TransferQueue/Mooncake 缺少本文“正确性依赖”所列能力 | 核对该节点的 `transfer_queue` / `mooncake-transfer-engine` 版本；镜像是否一致 | +| 启动日志 `the Mooncake master endpoint is not configured` 或 `segment capacity insufficient` / `the segment-capacity configuration is unusable` | 分别是 `MC_MASTER_ADDRESS` 缺失/格式非法、最坏情况容量上界超过 client segment、以及 `RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB` 或 `seq_length` 本身不可用 | 容量类问题按日志核对参数；减少 batch / `n_samples_per_prompt` / `max_staleness`，或调大 `RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB` 并同步规划每节点 attach 的瞬时资源足迹 | +| `Mooncake attach handshake reported N failure(s)`，摘要为 `handshake task failed (...)` 或 `handshake did not return ...` | 该节点 native setup、Ray task 或 attach 超时；日志只保留稳定的错误类型，不回显 endpoint、PID、路径或底层异常正文 | 在失败节点上按下方“选卡核验”逐项确认，并检查 master 连通性、Ray worker 日志与 `memlock` | +| handshake 报错含 `attached storage manager is not MooncakeStorageManager` | 该进程 attach 到的不是预期 Mooncake controller（例如集群里残留了上一次作业的 SimpleStorage controller） | 确认集群干净：本期按单任务独占集群设计，不接管他人的 controller | | `backend=SimpleStorage fallback=...`，但预期跑 RDMA，且日志显示 master 未配置 | `auto` 下 `MC_MASTER_ADDRESS` 缺失或格式非法会记 WARNING 后回退 SimpleStorage（`required` 直接失败） | 给每个节点的作业环境设置 `MC_MASTER_ADDRESS=:` | | `setup failed with error code: -1` | master 不可达 | 检查 master 进程与 `MC_MASTER_ADDRESS` | | `Failed to open segment ... Connection refused` | 上一轮客户端异常退出，死 segment 仍在 master 注册 | 等 `client_ttl`（30 s）过期后重试 | @@ -193,7 +218,7 @@ PYTHONPATH=. python -u scripts/benchmarks/tq_cross_node_bench.py \ ### 选卡核验 -指定设备前先确认端口状态与 GID： +Relax 不再自动扫描这些内容——启动只判断真实 attach/setup 与配置契约是否通过，具体哪一项硬件条件不满足需要在失败节点上手工核验。指定设备前先确认端口状态与 GID： ```bash ls /sys/class/infiniband/ # 有哪些设备 diff --git a/relax/components/actor.py b/relax/components/actor.py index 4aca7b8ec..933fef015 100644 --- a/relax/components/actor.py +++ b/relax/components/actor.py @@ -77,13 +77,11 @@ def __init__( lease_owner=self, ) - self.steps = ray.get( - self.actor_model.async_init( - config, - role=self.role, - with_ref=config.kl_coef != 0 or config.use_kl_loss, - with_opd_teacher=self.config.opd_teacher_load, - ) + self.steps = self.actor_model.init_and_wait( + config, + role=self.role, + with_ref=config.kl_coef != 0 or config.use_kl_loss, + with_opd_teacher=self.config.opd_teacher_load, ) assert len(set(self.steps)) == 1 diff --git a/relax/components/actor_fwd.py b/relax/components/actor_fwd.py index 7ba7ebb31..bb303fc0b 100644 --- a/relax/components/actor_fwd.py +++ b/relax/components/actor_fwd.py @@ -41,7 +41,7 @@ def __init__( 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)) + self.actor_model.init_and_wait(config, role=self.role, with_ref=False) self.step = 0 async def run(self) -> None: diff --git a/relax/components/critic.py b/relax/components/critic.py index 82e6c7575..4f44f038b 100644 --- a/relax/components/critic.py +++ b/relax/components/critic.py @@ -50,7 +50,7 @@ def __init__( args=config, num_gpus=num_gpus, pg=pgs, role=self.role, runtime_env=runtime_env ) - ray.get(self.critic_model.async_init(config, role=self.role, with_ref=False)) + self.critic_model.init_and_wait(config, role=self.role, with_ref=False) self.step = getattr(self.config, "start_rollout_id", None) or 0 # Wired by controller in colocate PPO to gate wake_up on SGLang offload. self._rollout_barrier: Optional[RolloutOffloadBarrier] = None diff --git a/relax/core/controller.py b/relax/core/controller.py index ff6c05d3e..a3ff0ac7b 100644 --- a/relax/core/controller.py +++ b/relax/core/controller.py @@ -47,19 +47,21 @@ 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, resolve_mooncake_master_address, resolve_tq_capacity_batch_size, + validate_config, validate_mooncake_runtime_contract, ) from relax.utils.tq.lifecycle import ( + TqHandshakeIsolationError, TqInitResult, close_tq_owner, initialize_tq_with_fallback, reap_unusable_tq_controller, + safe_exception_kind, uses_mooncake, verify_cluster_attach, ) @@ -153,9 +155,11 @@ def __init__(self, config: Namespace, runtime_env: dict = None) -> None: # is launched (RL is resolved later in placement_group.py). resolve_sft_num_rollout(self.config) - # Initialize data management system - self._initialize_data_system() try: + # Include data-system initialization in the constructor cleanup + # transaction: a cluster-handshake orchestration error can happen + # after the owner actor has created global TQ state. + self._initialize_data_system() self.dcs, self.config.coordinator_url = create_dcs_deployment() self._metrics_service_enabled = getattr(config, "use_metrics_service", False) @@ -198,11 +202,13 @@ def __init__(self, config: Namespace, runtime_env: dict = None) -> None: # 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.") + logger.error("Controller construction failed; closing any initialized TQ state.") 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}") + logger.warning( + f"TQ owner cleanup during failed construction failed ({safe_exception_kind(cleanup_error)})." + ) raise def _cleanup_s3_model_weights_after_init(self) -> None: @@ -344,6 +350,11 @@ def _initialize_data_system(self): mode=getattr(self.config, "tq_rdma_mode", "off"), fallback_conf=fallback_config, ) + # The owner becomes this Controller's cleanup responsibility before + # the cluster handshake. Otherwise a driver-side scheduling error can + # escape while ``self._tq_owner`` is still None and orphan global TQ + # state during constructor failure. + self._tq_owner = init_result.owner if uses_mooncake(init_result.config): init_result = self._confirm_mooncake_attach(init_result, fallback_config) self._tq_owner = init_result.owner @@ -353,18 +364,46 @@ def _initialize_data_system(self): 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. + """Validate Mooncake attach/setup on every alive Ray node. + + There is no static ``/sys`` probe. TQ clients live in Ray Serve + replicas and 0-CPU actors with no placement binding, so they may land + on any alive node, and a worker-side ``tq.init`` has no timeout of its + own. Each alive node therefore performs a bounded attach with the + stored config, confirms a Mooncake manager whose client is configured + for RDMA, and detaches. This is configuration/setup evidence, not + negotiated-transport or wire proof. + Failures are aggregated here so ``auto`` converges the whole job on one + backend instead of failing a single replica mid-deployment. + + On failure the Mooncake state created for this attempt is torn down + *before* any SimpleStorage initialisation: a surviving half-initialised + controller would make the next ``tq.init`` poll forever (F10). """ - failures = verify_cluster_attach(init_result.config) + + def _close_owned_attempt() -> None: + if not init_result.owns_controller: + return + try: + close_tq_owner(init_result.owner) + finally: + self._tq_owner = None + + try: + failures = verify_cluster_attach(init_result.config) + except TqHandshakeIsolationError: + _close_owned_attempt() + raise RuntimeError( + "Mooncake attach handshake workers could not be confirmed stopped; global fallback is unsafe" + ) from None + except Exception as error: + # Defensive boundary: verify_cluster_attach normally converts + # driver failures into stable summaries. An unexpected exception + # still must not strand the owner created for this attempt. + _close_owned_attempt() + raise RuntimeError( + f"Mooncake attach handshake orchestration failed ({safe_exception_kind(error)})" + ) from None if not failures: logger.info("[dataplane] Mooncake attach handshake passed on all alive nodes.") return init_result @@ -372,38 +411,45 @@ def _confirm_mooncake_attach(self, init_result: TqInitResult, fallback_config) - 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) + # ``required`` must fail loudly, and a job that merely attached to a + # foreign controller must never tear it down or replace its backend + # unilaterally. + _close_owned_attempt() raise RuntimeError( - f"Mooncake attach handshake failed on {len(failures)} node(s) (--tq-rdma-mode={mode}): {detail}" + f"Mooncake attach handshake reported {len(failures)} failure(s) (--tq-rdma-mode={mode}): {detail}" ) logger.warning( - f"[dataplane] Mooncake attach handshake failed on {len(failures)} node(s) ({detail}); " + f"[dataplane] Mooncake attach handshake reported {len(failures)} failure(s) ({detail}); " "closing Mooncake state and converging the whole job to SimpleStorage." ) - close_tq_owner(init_result.owner) + # A cleanup failure leaves global TQ state unknown, so it must propagate + # instead of starting SimpleStorage on top of a possibly dirty cluster. + _close_owned_attempt() fallback_result = initialize_tq_with_fallback(fallback_config, mode="auto") + self._tq_owner = fallback_result.owner return TqInitResult( config=fallback_result.config, owner=fallback_result.owner, - fallback_reason=f"attach_handshake_failed:{len(failures)}_nodes", + fallback_reason=f"attach_handshake_failed:{len(failures)}_failures", ) def _resolve_tq_backend(self, total_storage_size: int) -> dict: """Resolve the TransferQueue ``backend`` config dict. - ``--tq-rdma-mode=off`` retains the previous SimpleStorage and ownership - semantics while sharing the bounded worker-attach and failure-cleanup - hardening. ``auto``/``required`` run the RDMA capability probe - *before* ``tq.init`` and emit the startup log line; ``auto`` falls back - to SimpleStorage, ``required`` fails fast. + This method only validates *configuration*: the requested mode, the + TransferQueue correctness contract, the master endpoint and the segment + capacity. It deliberately performs no hardware capability probe -- a + ``/sys`` scan cannot see which nodes the scheduler will actually use, + and host-RDMA capability is established afterwards by the real attach + handshake in :meth:`_confirm_mooncake_attach`. + + ``off`` retains the previous SimpleStorage and ownership semantics. For + ``auto``/``required``, any unmet precondition either falls back to + SimpleStorage (``auto``) or fails fast (``required``). """ - # 1. Validate the requested mode (structural, before any probe). - # getattr defaults keep old checkpoints / non-argparse configs safe. + # 1. Validate the requested mode. A malformed mode is a configuration + # error, never a reason to silently run on SimpleStorage. errors = validate_config(self.config) if errors: raise ValueError("Invalid TransferQueue RDMA configuration:\n " + "\n ".join(errors)) @@ -419,14 +465,14 @@ def _simple_storage() -> dict: ) def _fall_back_or_raise(reason: str, error: Exception) -> dict: - """Every host-RDMA precondition failure funnels through here. + """Every unmet host-RDMA precondition funnels through here. ``auto`` converges the whole job on SimpleStorage; ``required`` re- raises so an operator who demanded RDMA never runs silently downgraded. """ if mode != "auto": - raise RuntimeError(f"--tq-rdma-mode={mode} but {reason}: {error}") from error + raise RuntimeError(f"--tq-rdma-mode={mode} but {reason}: {error}") from None logger.warning(f"[dataplane] {reason}; auto fallback to SimpleStorage: {error}") return _simple_storage() @@ -434,10 +480,7 @@ def _fall_back_or_raise(reason: str, error: Exception) -> dict: if mode == "off": return _simple_storage() - # 3. Host-RDMA 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. + # 3. Configuration preconditions for the host-RDMA path. device = getattr(self.config, "tq_rdma_device", "") try: validate_mooncake_runtime_contract() @@ -452,37 +495,29 @@ def _fall_back_or_raise(reason: str, error: Exception) -> dict: # unreachable for the whole job, so it degrades like any other # unmet precondition rather than aborting an ``auto`` run. return _fall_back_or_raise("the Mooncake master endpoint is not configured", e) - probe_results = probe_cluster_nodes(device, master_address) - - effective = reduce_results(probe_results, requested_device=device) - # 4. Only auto mode may fall back; ``required`` demands host RDMA. - if mode != "auto" and effective.fallback_reason: - detail = "\n".join(r.summary() for r in probe_results) - raise RuntimeError( - f"--tq-rdma-mode={mode} but host RDMA is unavailable: " - f"{effective.fallback_reason}.\n" - f"Probe details:\n{detail}" + # 4. Build the backend dict. Capacity validation reports an + # insufficient segment via ``cap_error``, but it can also *raise* for + # an unusable capacity input (garbage + # ``RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB``, missing ``seq_length``); both + # are configuration failures and must degrade identically. + try: + backend_dict, cap_error = build_backend_config( + self.config, + device=device, + master_address=master_address, + total_storage_size=total_storage_size, ) + except RuntimeError as e: + return _fall_back_or_raise("the segment-capacity configuration is unusable", e) + if cap_error: + return _fall_back_or_raise("segment capacity insufficient", RuntimeError(cap_error)) - # 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) - - # 6. Capacity fallback is also auto-only; ``required`` fails 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. Log requested vs effective so the startup log alone explains the - # decision, plus one summary block per probed node. + # 5. One startup line stating what was requested and what will run. The + # effective transport is only confirmed once the cluster-wide attach + # handshake passes. logger.info(f"[dataplane] requested: rdma_mode={mode} device={device or 'auto'}") - 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 fallback=segment_capacity_insufficient") - else: - logger.info(effective.log_line()) + logger.info(f"[dataplane] backend=MooncakeStore protocol=rdma device={device or 'auto'} (pending handshake)") return backend_dict def _close_data_system(self) -> None: diff --git a/relax/distributed/ray/actor_group.py b/relax/distributed/ray/actor_group.py index 44ba269b9..178e05edc 100644 --- a/relax/distributed/ray/actor_group.py +++ b/relax/distributed/ray/actor_group.py @@ -121,6 +121,91 @@ def async_init(self, args, role, with_ref=False, with_opd_teacher=False): for actor in self._actor_handlers ] + def init_and_wait(self, args, role, with_ref=False, with_opd_teacher=False): + """Initialize every train actor, destroying the group on any failure. + + ``MegatronTrainRayActor.init`` attaches the process-global + TransferQueue client from a regular actor method. If its bounded + ``tq.init`` times out, the method raises while the native daemon thread + may still be running. A failed distributed group cannot be reused + safely, so all actor processes are force-killed and confirmed terminal + before the initialization error is propagated. + """ + refs = self.async_init(args, role, with_ref=with_ref, with_opd_teacher=with_opd_teacher) + pending = list(refs) + results: dict[Any, Any] = {} + try: + # ``ray.get(refs)`` may wait for every ref before surfacing one + # failure. A rank blocked in native initialization would then + # prevent cleanup forever, so consume whichever ref finishes first. + while pending: + ready, pending = ray.wait(pending, num_returns=1) + if not ready: + raise RuntimeError("Ray returned no completed train-actor initialization task") + ref = ready[0] + results[ref] = ray.get(ref) + return [results[ref] for ref in refs] + except Exception: + self._terminate_failed_init() + raise + + def _terminate_failed_init(self, timeout: float = 10.0) -> None: + """Force-kill a failed train group and confirm its actor tasks + ended.""" + probe_refs = [] + control_errors: list[str] = [] + + # Queue a task that can never return normally before issuing the force + # kill. Each ref must later fail with RayActorError, proving that Ray + # processed the actor death instead of merely accepting a kill request. + for actor in self._actor_handlers: + try: + probe_refs.append(actor.termination_probe.remote()) + except ray.exceptions.RayActorError: + # The actor was already dead before cleanup reached it. + continue + except Exception as error: + control_errors.append(f"termination probe submission ({type(error).__name__})") + + for actor in self._actor_handlers: + try: + ray.kill(actor, no_restart=True) + except ray.exceptions.RayActorError: + continue + except Exception as error: + control_errors.append(f"actor kill ({type(error).__name__})") + + pending = [] + ready = [] + if probe_refs: + try: + ready, pending = ray.wait( + probe_refs, + num_returns=len(probe_refs), + timeout=timeout, + ) + except Exception as error: + control_errors.append(f"termination wait ({type(error).__name__})") + + for ref in ready: + try: + ray.get(ref) + except ray.exceptions.RayActorError: + continue + except Exception as error: + control_errors.append(f"termination probe result ({type(error).__name__})") + else: + control_errors.append("termination probe returned normally") + + if control_errors or pending: + detail = ", ".join(control_errors) + if pending: + detail = f"{detail}, " if detail else "" + detail += f"{len(pending)} task(s) remained pending" + raise RuntimeError(f"Failed to confirm train actor cleanup after initialization error: {detail}") from None + + self._actor_handlers = [] + def async_train(self, rollout_id): """Do one rollout training.""" return [actor.train.remote(rollout_id) for actor in self._actor_handlers] diff --git a/relax/distributed/ray/train_actor.py b/relax/distributed/ray/train_actor.py index 9b497a93f..59b97f53f 100644 --- a/relax/distributed/ray/train_actor.py +++ b/relax/distributed/ray/train_actor.py @@ -3,6 +3,7 @@ import abc import os import random +import threading from datetime import timedelta import ray @@ -79,6 +80,15 @@ def init(self, args, role, with_ref=False, with_opd_teacher=False): numa_local_rank = Envs.RANK % args.num_gpus_per_node device_utils.set_numa_affinity(numa_local_rank) + def termination_probe(self) -> None: + """Never return normally; used to confirm force-killed actor exit. + + The driver queues this behind ``init`` before calling ``ray.kill``. Its + ObjectRef must become terminal with ``RayActorError``; a successful + return would mean the actor survived cleanup and is unsafe to reuse. + """ + threading.Event().wait() + def clear_memory(self): print_memory("before TrainRayActor.clear_memory") clear_memory() diff --git a/relax/utils/payload_digest.py b/relax/utils/payload_digest.py index cef6f4d86..efac03fac 100644 --- a/relax/utils/payload_digest.py +++ b/relax/utils/payload_digest.py @@ -18,6 +18,7 @@ from __future__ import annotations import hashlib +import struct from typing import Any import numpy as np @@ -50,7 +51,11 @@ def _scalar_digest(value: Any) -> LeafDigest: raw = value elif isinstance(value, str): raw = value.encode("utf-8") - else: # bool / int / float / None — repr is canonical for these types. + elif isinstance(value, float): + # repr(float("nan")) discards the payload bits. Pack the Python double + # directly so distinct NaNs and signed zero remain byte-distinguishable. + raw = struct.pack("!d", value) + else: # bool / int / None — repr is canonical for these types. raw = repr(value).encode("utf-8") return (f"py.{type(value).__name__}", "", hashlib.sha256(raw).hexdigest()) @@ -67,6 +72,15 @@ def _unwrap_non_tensor(value: Any) -> Any: return value +def _dict_child_path(prefix: str, key: Any) -> str: + """Render a string dict key without colliding with nested/list paths.""" + if not isinstance(key, str): + raise TypeError(f"Unsupported payload dict key at {prefix}: {type(key).__name__}") + if key.isidentifier(): + return f"{prefix}.{key}" + return f"{prefix}[{key!r}]" + + def leaf_digests(payload: Any, prefix: str = "payload") -> dict[str, LeafDigest]: """Map every leaf of *payload* to ``(dtype, shape, sha256)``. @@ -87,8 +101,11 @@ def leaf_digests(payload: Any, prefix: str = "payload") -> dict[str, LeafDigest] 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}")) + for key in payload: + if not isinstance(key, str): + raise TypeError(f"Unsupported payload dict key at {prefix}: {type(key).__name__}") + for key in sorted(payload): + digests.update(leaf_digests(payload[key], _dict_child_path(prefix, key))) elif isinstance(payload, (list, tuple)): for index, item in enumerate(payload): digests.update(leaf_digests(item, f"{prefix}[{index}]")) diff --git a/relax/utils/rdma_probe.py b/relax/utils/rdma_probe.py deleted file mode 100644 index ce2a0554a..000000000 --- a/relax/utils/rdma_probe.py +++ /dev/null @@ -1,576 +0,0 @@ -# Copyright (c) 2026 Relax Authors. All Rights Reserved. - -"""RDMA capability probe and fallback decision 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. - -Only two outcomes exist: MooncakeStore over host RDMA, or the original -SimpleStorage. Mooncake/TCP survives as a benchmark baseline only, so there is -no intermediate transport to degrade through. - -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__) - -# Accepted ``--tq-rdma-mode`` values; ``off`` keeps the SimpleStorage path. -TQ_RDMA_MODES = frozenset({"off", "auto", "required"}) - -# --------------------------------------------------------------------------- -# 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 - 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} device={self.effective_device}" - 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 - fallback_reason: str # "" if no fallback occurred - - def log_line(self) -> str: - """Return the single-line startup log string for this effective config. - - SimpleStorage deliberately reports no protocol: naming one would imply - the data plane went through a Mooncake transport, and Mooncake/TCP is - not a production path. - """ - if self.backend != "MooncakeStore": - base = f"[dataplane] backend={self.backend}" - else: - base = f"[dataplane] backend={self.backend} protocol={self.protocol} device={self.device or 'auto'}" - 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: - 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 _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("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: - 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 _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. 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 - - 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 because it depends on -# local Mooncake client initialization. External-master reachability is checked -# directly with a bounded TCP connect, then authoritatively by ``tq.init``. -def probe_node(device: str = "", master_address: str = "") -> ProbeResult: - """Run all capability checks on the current node. - - Parameters - ---------- - device - Explicit RDMA device name; empty = scan every HCA and select one - whose ACTIVE port and usable GID both pass. - """ - node = socket.gethostname() - checks: list[CheckResult] = [] - errors: list[str] = [] - - checks.append(_check_mooncake_import()) - 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 - # the job creates a global TransferQueue controller. - if master_address: - checks.append(_check_master_reachable(master_address)) - - 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) - 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) - 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" - # Report the jointly validated device (ACTIVE port + usable GID). - effective_device = selected_device or effective_device - else: - # Mooncake is usable but this node cannot do RDMA. ``tcp`` is recorded - # so the reduction can distinguish "no RDMA here" from "no Mooncake at - # all" in its fallback reason; Mooncake/TCP is not a production - # transport (it survives only as a benchmark baseline). - 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") - - return ProbeResult( - node=node, - checks=tuple(checks), - effective_protocol=effective_protocol, - effective_device=effective_device, - 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="", - errors=(error,), - ) - - -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 - 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. - - 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 [driver_result] - - import ray - from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy - - @ray.remote(num_cpus=0.001) - def _probe_on_node(dev: str, master: str) -> ProbeResult: - from relax.utils.rdma_probe import probe_node as _probe - - 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, 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] = [driver_result] - 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_device: str, -) -> EffectiveConfig: - """AND-reduce per-node results into a single job-level effective config. - - The only two outcomes are MooncakeStore over host RDMA and the original - SimpleStorage: Mooncake/TCP is a benchmark baseline, not a production - transport, so it is never selected here and there is no intermediate rung - to degrade through. - - Parameters - ---------- - results - One :class:`ProbeResult` per data-plane node. - requested_device - ``--tq-rdma-device`` value. - """ - if not results: - return EffectiveConfig( - backend="SimpleStorage", - protocol="tcp", - device="", - fallback_reason="no probe results", - ) - - # AND reduction: the job can only run at the lowest common capability. - if all(r.effective_protocol == "rdma" for r in results): - # Device: every node must expose the explicitly requested device. - if requested_device and any(r.effective_device != requested_device for r in results): - return EffectiveConfig( - backend="SimpleStorage", - protocol="tcp", - device="", - fallback_reason=f"device_mismatch:{requested_device}", - ) - return EffectiveConfig( - backend="MooncakeStore", - protocol="rdma", - device=requested_device, - fallback_reason="", - ) - - no_mooncake_nodes = [r.node for r in results if r.effective_protocol is None] - if no_mooncake_nodes: - 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(no_mooncake_nodes)}" - ) - else: - reason = f"rdma_unavailable:{','.join(r.node for r in results if r.effective_protocol != 'rdma')}" - return EffectiveConfig( - backend="SimpleStorage", - protocol="tcp", - device="", - fallback_reason=reason, - ) - - -# --------------------------------------------------------------------------- -# Validation -# --------------------------------------------------------------------------- - - -def validate_config(args: Any) -> list[str]: - """Return a list of error messages for an invalid RDMA configuration. - - Called at startup *before* probing. An empty list means the config is - structurally valid (semantic/runtime validity is checked by the probe). - ``argparse`` already constrains the mode for CLI runs; this also covers - configs restored from a checkpoint or built programmatically. - """ - errors: list[str] = [] - mode = getattr(args, "tq_rdma_mode", "off") - - if mode not in TQ_RDMA_MODES: - errors.append(f"--tq-rdma-mode={mode!r} must be one of {', '.join(sorted(TQ_RDMA_MODES))}.") - return errors diff --git a/relax/utils/tq/config.py b/relax/utils/tq/config.py index 4b4e6442c..c335d813b 100644 --- a/relax/utils/tq/config.py +++ b/relax/utils/tq/config.py @@ -3,9 +3,10 @@ """Build TransferQueue backend config dicts from Relax CLI args. This module is the single place that maps Relax-side *intent* flags -(``--tq-rdma-mode``, ``--tq-rdma-device``) plus an -:class:`~relax.utils.rdma_probe.EffectiveConfig` into the OmegaConf dict that -``tq.init`` expects. +(``--tq-rdma-mode``, ``--tq-rdma-device``) into the OmegaConf dict that +``tq.init`` expects, and the only place that validates the configuration +preconditions for MooncakeStore. Actual host-RDMA capability is established by +the real attach handshake in :mod:`relax.utils.tq.lifecycle`, not here. Mooncake internals (endpoint, buffer size, segment size, master address, timeout) are intentionally *not* exposed as CLI flags — they come from @@ -15,16 +16,19 @@ from __future__ import annotations import inspect +import math import os from typing import Any 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__) +# Accepted ``--tq-rdma-mode`` values; ``off`` keeps the SimpleStorage path. +TQ_RDMA_MODES = frozenset({"off", "auto", "required"}) + # --------------------------------------------------------------------------- # Defaults (kept here rather than in config.yaml so they are visible to # Relax contributors without reading the TQ package). @@ -35,13 +39,71 @@ _DEFAULT_METADATA_SERVER = "P2PHANDSHAKE" # config.yaml:42-43 +def validate_config(args: Any) -> list[str]: + """Return error messages for an invalid RDMA configuration. + + Called at startup before anything is initialised. An empty list means the + requested mode is structurally valid; whether host RDMA actually works is + decided later by the real attach handshake. ``argparse`` already + constrains the mode for CLI runs, so this also covers configs restored from + a checkpoint or built programmatically. + """ + errors: list[str] = [] + mode = getattr(args, "tq_rdma_mode", "off") + device = getattr(args, "tq_rdma_device", "") + + if not isinstance(mode, str) or mode not in TQ_RDMA_MODES: + errors.append(f"--tq-rdma-mode must be one of {', '.join(sorted(TQ_RDMA_MODES))}.") + if not isinstance(device, str): + errors.append("--tq-rdma-device must be a string.") + elif device and any(character.isspace() or not character.isprintable() for character in device): + errors.append("--tq-rdma-device must be empty or a single printable device name without whitespace.") + return errors + + +def _split_host_port(address: str) -> tuple[str, int]: + """Parse ``host:port`` (and bracketed IPv6) or raise ``ValueError``. + + Format-only: no DNS resolution and no connection attempt, so this stays a + configuration check rather than becoming another capability probe. + + Failure messages name the *kind* of defect and never echo ``address``: the + caller logs them, and a deployment endpoint is internal infrastructure + detail that must not leak into job logs. + """ + value = address.strip() + if value.startswith("["): + end = value.find("]") + if end < 0 or end + 2 > len(value) or value[end + 1] != ":": + raise ValueError("bracketed endpoint must be [host]:port") + host, port_text = value[1:end], value[end + 2 :] + else: + host, separator, port_text = value.rpartition(":") + if not separator: + raise ValueError("endpoint must be host:port") + if ":" in host: + # A bare IPv6 literal: rpartition would silently treat its last + # group as the port (``fe80::1`` -> host ``fe80:``, port 1). + raise ValueError("IPv6 endpoint must be bracketed as [host]:port") + if not host: + raise ValueError("endpoint has an empty host") + if any(character.isspace() or not character.isprintable() for character in host): + raise ValueError("endpoint host must not contain whitespace or control characters") + if not port_text.isdigit(): + raise ValueError("endpoint port must be a decimal number") + port = int(port_text) + if not 1 <= port <= 65535: + raise ValueError("endpoint port is outside 1-65535") + return host, port + + def resolve_mooncake_master_address() -> str: """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. + ``MC_MASTER_ADDRESS`` is required and must parse as ``host:port``. A + loopback default would make every node of a multi-node job treat its own + localhost as the master, so ``auto`` would degrade and ``required`` would + abort even when a shared master is healthy elsewhere. """ address = os.environ.get("MC_MASTER_ADDRESS", "").strip() if not address: @@ -50,6 +112,10 @@ def resolve_mooncake_master_address() -> str: "managed mooncake master on every node; Relax never assumes a loopback " "endpoint." ) + try: + _split_host_port(address) + except ValueError as error: + raise RuntimeError(f"MC_MASTER_ADDRESS is not a usable endpoint: {error}") from error return address @@ -67,26 +133,35 @@ def resolve_global_segment_size() -> int: 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) + except ValueError: + raise RuntimeError("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB must be a finite positive number of GiB") from None + if not math.isfinite(gib) or gib <= 0: + raise RuntimeError("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB must be a finite positive number of GiB") + size_bytes = int(gib * 1024**3) + if size_bytes <= 0: + raise RuntimeError("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB must resolve to at least one byte") + return size_bytes def validate_mooncake_runtime_contract() -> None: - """Install and validate the Mooncake loss-prevention contract. + """Validate the Relax-side portion of the Mooncake safety contract. - 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. + The environment guard and available read-only capability checks run before + every Mooncake client is created or attached. Retry result-length and + positive-ACK correctness must be supplied by the pinned upstream + TransferQueue revision; this function does not monkey-patch that package. """ ensure_mooncake_correctness_guards() from transfer_queue.storage.managers.base import KVStorageManager - put_source = inspect.getsource(KVStorageManager.put_data) + try: + put_source = inspect.getsource(KVStorageManager.put_data) + except (OSError, TypeError): + # No retrievable source (compiled/stripped install): the ordering + # contract cannot be proven, so fail like any other unmet gate instead + # of escaping this RuntimeError-only boundary. + raise RuntimeError("Cannot verify TransferQueue put/notify ordering because source is unavailable") from None 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: @@ -118,29 +193,38 @@ def build_simple_storage_config(total_storage_size: int | None, num_data_storage def build_mooncake_config( - effective: EffectiveConfig, *, - master_address: str | None = None, + master_address: str, + device: str = "", + protocol: str = "rdma", 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 required - ``MC_MASTER_ADDRESS`` env var (see - :func:`resolve_mooncake_master_address`). + External master server address, already validated by + :func:`resolve_mooncake_master_address`. It is passed in rather than + re-read from the environment so the value that was checked is the value + that gets used. + device + Explicit RDMA device name; empty lets Mooncake select one natively. + protocol + Transport. Production only ever builds ``rdma``; ``tcp`` exists solely + for the C1 benchmark baseline. global_segment_size 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() + if global_segment_size is None: + segment_size = resolve_global_segment_size() + elif isinstance(global_segment_size, bool) or not isinstance(global_segment_size, int) or global_segment_size <= 0: + raise ValueError("global_segment_size must be a positive integer number of bytes") + else: + segment_size = global_segment_size cfg: dict[str, Any] = { # Selects the manager inside ``tq.init`` (interface.py reads @@ -150,15 +234,15 @@ def build_mooncake_config( "storage_backend": "MooncakeStore", "MooncakeStore": { # Transport - "protocol": effective.protocol, # "rdma" or "tcp" - "device_name": effective.device, + "protocol": protocol, + "device_name": 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 resolve_global_segment_size(), + "global_segment_size": segment_size, "local_buffer_size": _DEFAULT_LOCAL_BUFFER_SIZE, # Do NOT silently evict produced-but-unconsumed data. "hard_pin": True, @@ -228,17 +312,14 @@ def estimate_payload_bytes(args: Any) -> int: return capacity_batch * n_samples * per_sample -def validate_segment_capacity(args: Any, effective: EffectiveConfig) -> str | None: +def validate_segment_capacity(args: Any) -> 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. + Only meaningful for MooncakeStore (SimpleStorage manages its own capacity + via ``total_storage_size``), so callers invoke it on the Mooncake path + only. The check is conservative: it compares 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) capacity_batch = resolve_tq_capacity_batch_size(args) payload = estimate_payload_bytes(args) @@ -263,23 +344,20 @@ def validate_segment_capacity(args: Any, effective: EffectiveConfig) -> str | No def build_backend_config( args: Any, - effective: EffectiveConfig, *, + device: str, + master_address: str, total_storage_size: int, ) -> tuple[dict[str, Any], str | None]: - """Return ``(backend_config_dict, error_or_none)``. + """Return ``(backend_config_dict, error_or_none)`` for the host-RDMA path. - On error, ``backend_config_dict`` is a safe SimpleStorage fallback and - ``error`` explains why MooncakeStore was rejected. + ``master_address`` must already be validated by + :func:`resolve_mooncake_master_address`; it is threaded through so the + checked value is the one the client receives. On a capacity error the + returned dict is a safe SimpleStorage fallback and ``error`` explains why + MooncakeStore was rejected -- the caller decides whether that is fatal. """ - 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) + cap_error = validate_segment_capacity(args) if cap_error: logger.error(cap_error) return build_simple_storage_config( @@ -287,4 +365,4 @@ def build_backend_config( num_data_storage_units=args.num_data_storage_units, ), cap_error - return build_mooncake_config(effective), None + return build_mooncake_config(master_address=master_address, device=device), None diff --git a/relax/utils/tq/correctness.py b/relax/utils/tq/correctness.py index 8a5e7c91c..501845e2f 100644 --- a/relax/utils/tq/correctness.py +++ b/relax/utils/tq/correctness.py @@ -2,14 +2,11 @@ """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. +Relax validates the capabilities and environment it can inspect without +modifying TransferQueue at runtime. Per-retry result-length validation and a +strict production-status ACK must be fixed in upstream TransferQueue and then +consumed through an updated, capability-marked pin; method-name checks alone do +not prove those semantics. The pinned mooncake 0.3.10 additionally corrupts TCP-protocol transfers through its auto-enabled memcpy fast path, so that path is force-disabled @@ -39,7 +36,7 @@ def _enforce_safe_memcpy() -> None: 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 " + "MC_STORE_MEMCPY explicitly enables an unsafe value: 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." ) @@ -49,8 +46,9 @@ def _enforce_safe_memcpy() -> None: def ensure_mooncake_correctness_guards() -> None: """Validate that the installed stack can run MooncakeStore safely. - Read-only: checks the memcpy environment contract and that the pinned + Enforces the memcpy environment contract and checks that the pinned TransferQueue ships the Mooncake retry APIs Relax's data plane relies on. + It does not modify TransferQueue code or objects at runtime. """ _enforce_safe_memcpy() try: diff --git a/relax/utils/tq/lifecycle.py b/relax/utils/tq/lifecycle.py index 320163444..bf4050634 100644 --- a/relax/utils/tq/lifecycle.py +++ b/relax/utils/tq/lifecycle.py @@ -21,6 +21,7 @@ from __future__ import annotations +import math import os import threading import time @@ -79,89 +80,38 @@ class TqConfigurationMismatch(RuntimeError): """Raised when an existing controller uses an incompatible job 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) +class TqControllerInspectionError(RuntimeError): + """Raised when Ray cannot prove the named controller's current state.""" -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)), - # Retained although Relax now always builds ``use_gdr=False``: this - # process may attach to a controller created by an older build that - # still enabled GDR, and upstream ``tq.init`` ignores the caller's - # conf when attaching (interface.py:130-135), so the worker would - # silently run the unverified GDR path. Delete together with the - # rest of the signature machinery once the exclusive-cluster - # simplification drops compatibility attach. - 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"), - ) +class TqControllerMissingConfig(RuntimeError): + """Raised when a reachable controller is provably half-initialised.""" + +class TqHandshakeIsolationError(RuntimeError): + """Raised when timed-out handshake workers cannot be confirmed stopped.""" -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. +def safe_exception_kind(error: BaseException) -> str: + """Return a log-safe exception category without rendering its payload. + + Ray exception strings can contain worker IPs, process IDs and remote + traceback paths. Lifecycle logs and job-level failure summaries therefore + record only the stable exception class. """ - 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")), - ) + ray_task_error_type = getattr(getattr(ray, "exceptions", None), "RayTaskError", None) + if isinstance(ray_task_error_type, type) and isinstance(error, ray_task_error_type): + return "RayTaskError" + name = type(error).__name__ + return name if name.isidentifier() else "Exception" -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 _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 uses_mooncake(conf: Any) -> bool: @@ -197,7 +147,7 @@ def kill_tq_controller_and_wait(timeout: float = 20.0) -> None: except ValueError: return # actor does not exist — nothing to kill or wait for. except Exception as e: - raise RuntimeError(f"Failed to kill TransferQueueController: {e}") from e + raise RuntimeError(f"Failed to kill TransferQueueController ({safe_exception_kind(e)})") from None deadline = time.time() + timeout while time.time() < deadline: @@ -205,32 +155,48 @@ def kill_tq_controller_and_wait(timeout: float = 20.0) -> None: ray.get_actor(CONTROLLER_NAME, namespace=CONTROLLER_NAMESPACE) except ValueError: return + except ray.exceptions.RayError as error: + raise RuntimeError( + f"Failed to confirm TransferQueueController cleanup ({safe_exception_kind(error)})" + ) from None time.sleep(0.4) raise TqCleanupTimeout(f"TransferQueueController still resolvable after {timeout}s") def reap_unusable_tq_controller(get_config_timeout: float = 10.0) -> bool: - """Kill the TransferQueueController only if it cannot serve a config. + """Require a clean exclusive cluster, reaping only unusable TQ state. - 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). + Returns ``True`` when a half-initialised or unresponsive controller was + reaped. A healthy controller is never attached or killed: the initial RDMA + release supports one Relax job per Ray cluster, so healthy existing state + means the cluster is not clean and startup fails explicitly. """ try: existing = ray.get_actor(CONTROLLER_NAME, namespace=CONTROLLER_NAMESPACE) except ValueError: return False # nothing there — nothing to reap. + except ray.exceptions.RayError as error: + raise RuntimeError(f"Failed to query TransferQueueController ({safe_exception_kind(error)})") from None 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.") + except (ray.exceptions.GetTimeoutError, ray.exceptions.RayActorError) as error: + # Only a proven timeout/dead actor is safe to classify as unusable. + # Other RayError subclasses can be transient GCS/control-plane failures; + # killing a healthy controller on those would violate the exclusive-job + # ownership boundary. + logger.warning( + f"[dataplane] Existing TransferQueueController is unusable ({safe_exception_kind(error)}); reaping it." + ) conf = None + except ray.exceptions.RayError as error: + raise RuntimeError(f"Failed to inspect TransferQueueController ({safe_exception_kind(error)})") from None if conf is not None: - logger.info("[dataplane] Existing TransferQueueController is healthy; tq.init will attach to it.") - return False + raise TqConfigurationMismatch( + "A healthy TransferQueueController already exists. The initial RDMA release requires an exclusive, " + "clean Ray cluster; stop the previous Relax job or remove its TQ state before retrying." + ) logger.warning("[dataplane] TransferQueueController has no stored config (half-initialised); reaping it.") kill_tq_controller_and_wait() @@ -243,6 +209,8 @@ def _controller_exists() -> bool: return True except ValueError: return False + except ray.exceptions.RayError as error: + raise RuntimeError(f"Failed to query TransferQueueController ({safe_exception_kind(error)})") from None def _set_owner_token(conf: Any, token: str) -> None: @@ -259,11 +227,49 @@ def _get_owner_token(conf: Any) -> str: return "" +def assert_mooncake_rdma_configured() -> None: + """Fail unless the attached client configuration requests Mooncake/RDMA. + + A successful :func:`attach_tq_client` only proves ``tq.init`` returned + without raising. This check also establishes that the stored controller + config produced a Mooncake manager and an RDMA-configured storage client: + ``tq.init`` ignores the caller's conf when attaching to an existing + controller (``interface.py:130-135``). ``storage_client.protocol`` is the + configured request, not a negotiated transport signal, so this is not + proof that bytes traversed an HCA or that a native transport fallback is + impossible. Wire-level proof remains the benchmark's counter check. + + Raises + ------ + RuntimeError + When the storage manager is not MooncakeStore, when no storage client + is present, or when the client's configured protocol is not ``rdma``. + """ + manager = tq.get_client().storage_manager + if type(manager).__name__ != "MooncakeStorageManager": + raise RuntimeError("attached storage manager is not MooncakeStorageManager") + + store_client = getattr(manager, "storage_client", None) + if store_client is None: + raise RuntimeError("MooncakeStorageManager exposes no storage_client after attach") + + protocol = getattr(store_client, "protocol", None) + if protocol != "rdma": + raise RuntimeError("MooncakeStore client is not configured for protocol=rdma") + + 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) + try: + controller = ray.get_actor(CONTROLLER_NAME, namespace=CONTROLLER_NAMESPACE) + conf = ray.get(controller.get_config.remote(), timeout=timeout) + except ValueError: + raise + except ray.exceptions.RayError as error: + raise TqControllerInspectionError( + f"Failed to read TransferQueueController config ({safe_exception_kind(error)})" + ) from None if conf is None: - raise RuntimeError("TransferQueueController returned no config after tq.init completed") + raise TqControllerMissingConfig("TransferQueueController returned no config after tq.init completed") return conf @@ -281,13 +287,13 @@ def _close_local_tq_client() -> None: 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}") + logger.warning(f"[dataplane] Failed to close attached MooncakeStore client ({safe_exception_kind(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}") + logger.warning(f"[dataplane] Failed to close attached TransferQueue client ({safe_exception_kind(e)}).") # TransferQueue has no public detach-only API. Reset only process-local # handles; never touch _TQ_STORAGE or the named controller actor. @@ -308,10 +314,10 @@ def _resolve_attach_timeout() -> float: 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") + except ValueError: + raise RuntimeError("RELAX_TQ_ATTACH_TIMEOUT_SECONDS must be a finite positive number of seconds") from None + if not math.isfinite(value) or value <= 0: + raise RuntimeError("RELAX_TQ_ATTACH_TIMEOUT_SECONDS must be a finite positive number of seconds") return value @@ -335,7 +341,7 @@ def _await_controller_config(deadline: float) -> None: 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}" + last_error = f"get_config failed ({safe_exception_kind(e)})" time.sleep(min(0.5, max(deadline - time.monotonic(), 0.0))) continue if conf is not None: @@ -445,7 +451,41 @@ def detach_tq_client(generation: int | None = None) -> None: 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")] + node_ids: list[str] = [] + for node in ray.nodes(): + if not node.get("Alive"): + continue + node_id = node.get("NodeID") + if not isinstance(node_id, str) or not node_id: + raise RuntimeError("an alive Ray node has no valid NodeID") + node_ids.append(node_id) + return node_ids + + +def _cancel_handshake_tasks(refs: list[Any], *, timeout: float = 10.0) -> bool: + """Force-cancel submitted one-shot workers and confirm every ref is ready. + + A timed-out worker may still own a daemon ``tq.init`` thread. The caller + may only tear down Mooncake and start SimpleStorage after Ray confirms the + force-cancelled task refs have reached a terminal state. + """ + cancellation_failed = False + for ref in refs: + try: + ray.cancel(ref, force=True) + except Exception as error: # pragma: no cover - Ray control-plane failure + cancellation_failed = True + logger.warning(f"[dataplane] Failed to force-cancel a TQ handshake worker ({safe_exception_kind(error)}).") + if cancellation_failed: + return False + try: + _ready, pending = ray.wait(refs, num_returns=len(refs), timeout=timeout) + except Exception as error: # pragma: no cover - Ray control-plane failure + logger.warning( + f"[dataplane] Could not confirm TQ handshake worker cancellation ({safe_exception_kind(error)})." + ) + return False + return not pending def verify_cluster_attach(conf: Any, *, timeout: float | None = None) -> list[str]: @@ -453,45 +493,81 @@ def verify_cluster_attach(conf: Any, *, timeout: float | None = None) -> list[st summaries. 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. + worker would perform (real stored config, real storage client), confirms a + Mooncake manager whose client is configured for ``protocol=rdma``, and + detaches immediately. This validates actual attach/setup and configuration + agreement rather than a ``/sys`` heuristic; it is not negotiated-transport + or wire proof. 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, max_calls=1) - def _handshake(handshake_conf: Any) -> None: - from relax.utils.tq.lifecycle import attach_tq_client, detach_tq_client + try: + node_ids = _alive_node_ids() + except Exception as error: + return [f"cluster: node discovery failed ({safe_exception_kind(error)})"] + if not node_ids: + return ["cluster: no alive Ray nodes discovered"] - attach_tq_client(handshake_conf, role="attach-handshake") - detach_tq_client() + expects_mooncake = uses_mooncake(conf) 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 + try: + from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy + + @ray.remote(num_cpus=0, max_retries=0, max_calls=1) + def _handshake(handshake_conf: Any, check_mooncake: bool, attach_timeout: float) -> None: + from relax.utils.tq.lifecycle import ( + assert_mooncake_rdma_configured, + attach_tq_client, + detach_tq_client, + ) + + attach_tq_client(handshake_conf, role="attach-handshake", timeout=attach_timeout) + try: + if check_mooncake: + assert_mooncake_rdma_configured() + finally: + # A failed assertion must not leak this node's registered + # segment; without the detach it lingers until client_ttl. + detach_tq_client() + + for node_id in node_ids: + strategy = NodeAffinitySchedulingStrategy(node_id=node_id, soft=False) + ref = _handshake.options(scheduling_strategy=strategy).remote(conf, expects_mooncake, timeout) + refs.append(ref) + id_by_ref[ref] = node_id + except Exception as error: + if refs and not _cancel_handshake_tasks(refs): + raise TqHandshakeIsolationError( + "submitted TQ handshake workers could not be confirmed stopped after scheduling failed" + ) from None + return [f"cluster: handshake scheduling failed ({safe_exception_kind(error)})"] # 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) + try: + ready, pending = ray.wait(refs, num_returns=len(refs), timeout=wait_bound) + except Exception as error: + if not _cancel_handshake_tasks(refs): + raise TqHandshakeIsolationError( + "submitted TQ handshake workers could not be confirmed stopped after wait failed" + ) from None + return [f"cluster: handshake wait failed ({safe_exception_kind(error)})"] 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}") + except Exception as error: + failures.append(f"node {id_by_ref[ref][:12]}: handshake task failed ({safe_exception_kind(error)})") + if pending and not _cancel_handshake_tasks(pending): + raise TqHandshakeIsolationError("timed-out TQ handshake workers could not be confirmed stopped") 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 @@ -525,7 +601,7 @@ def close_tq_and_unmount(*, is_owner: bool) -> None: 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}") + logger.warning(f"[dataplane] Failed to unmount MooncakeStore segment ({safe_exception_kind(e)}).") @ray.remote(num_cpus=0) @@ -554,7 +630,7 @@ 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}") + logger.debug(f"[dataplane] TQ owner actor already stopped ({safe_exception_kind(e)}).") def _cleanup_failed_owner(owner: Any, owner_token: str, *, timeout: float = 10.0) -> None: @@ -567,7 +643,7 @@ def _cleanup_failed_owner(owner: Any, owner_token: str, *, timeout: float = 10.0 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}") + logger.warning(f"[dataplane] TQ owner cleanup RPC failed; killing owner actor ({safe_exception_kind(e)}).") finally: _stop_owner_actor(owner) @@ -575,10 +651,22 @@ def _cleanup_failed_owner(owner: Any, owner_token: str, *, timeout: float = 10.0 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.") + except TqControllerMissingConfig as error: + logger.warning( + "[dataplane] Failed initializer left a half-initialised TQ controller " + f"({safe_exception_kind(error)}); reaping it." + ) kill_tq_controller_and_wait() return + except TqControllerInspectionError: + # The owner process has already been stopped, but a transient GCS error + # cannot prove that any visible controller belongs to this attempt. + # Fail closed instead of risking another job's healthy controller. + raise + except Exception as error: + raise TqControllerInspectionError( + f"Unexpected failure while inspecting TransferQueueController ({safe_exception_kind(error)})" + ) from None stored_token = _get_owner_token(stored_conf) if stored_token == owner_token: @@ -592,35 +680,40 @@ def _cleanup_failed_owner(owner: Any, owner_token: str, *, timeout: float = 10.0 def _start_owner(conf: Any, *, timeout: float) -> TqInitResult: - owner = _TransferQueueOwner.remote() owner_token = uuid.uuid4().hex + try: + owner = _TransferQueueOwner.remote() + except ray.exceptions.RayError as error: + raise RuntimeError(f"TransferQueue owner creation failed ({safe_exception_kind(error)})") from None try: stored_conf, owns_controller = ray.get(owner.initialize.remote(conf, owner_token), timeout=timeout) - except ray.exceptions.GetTimeoutError as e: + except ray.exceptions.GetTimeoutError: _cleanup_failed_owner(owner, owner_token) - raise TqInitializationTimeout(f"tq.init did not finish within {timeout:.0f}s") from e - except Exception: + raise TqInitializationTimeout(f"tq.init did not finish within {timeout:.0f}s") from None + except Exception as error: _cleanup_failed_owner(owner, owner_token) + if isinstance(error, ray.exceptions.RayError): + raise RuntimeError(f"TransferQueue owner initialization failed ({safe_exception_kind(error)})") from None 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(). - config_mismatch = _configuration_signature(stored_conf) != _configuration_signature(conf) + # A controller appeared after the exclusive-cluster pre-check. This owner + # only attached to it, so detach locally and fail instead of silently sharing + # global state with a concurrent initializer. try: ray.get(owner.detach.remote(), timeout=10.0) + except ray.exceptions.RayError as error: + raise RuntimeError( + f"TransferQueue concurrent-initializer detach failed ({safe_exception_kind(error)})" + ) from None finally: _stop_owner_actor(owner) - if config_mismatch: - raise TqConfigurationMismatch( - "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." - ) - return TqInitResult(config=stored_conf, owner=None) + raise TqConfigurationMismatch( + "A concurrent TransferQueue initializer created a healthy controller. " + "Detached without modifying it; the initial RDMA release requires an exclusive Ray cluster." + ) def close_tq_owner(owner: Any | None, *, timeout: float = 30.0) -> None: @@ -640,7 +733,7 @@ def close_tq_owner(owner: Any | None, *, timeout: float = 30.0) -> None: 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 + raise RuntimeError(f"TransferQueue owner cleanup failed ({safe_exception_kind(close_error)})") from None def initialize_tq_with_fallback( @@ -659,18 +752,6 @@ def initialize_tq_with_fallback( 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. - stored_conf = _get_stored_config() - 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." - ) - return TqInitResult(config=stored_conf, owner=None) return _start_owner(attempt_conf, timeout=timeout) try: @@ -681,7 +762,7 @@ def _attempt(attempt_conf: Any) -> TqInitResult: reason = f"mooncake_init_failed:{type(primary_error).__name__}" logger.warning( - f"[dataplane] Mooncake tq.init failed ({primary_error}); " + f"[dataplane] Mooncake tq.init failed ({safe_exception_kind(primary_error)}); " "cleaned partial state and retrying once with SimpleStorage." ) try: @@ -689,8 +770,9 @@ def _attempt(attempt_conf: Any) -> TqInitResult: except Exception as fallback_error: raise RuntimeError( "TransferQueue SimpleStorage fallback initialization failed after " - f"Mooncake initialization error: {primary_error}" - ) from fallback_error + f"Mooncake initialization error ({safe_exception_kind(primary_error)}); " + f"fallback error ({safe_exception_kind(fallback_error)})" + ) from None return TqInitResult( config=result.config, owner=result.owner, diff --git a/scripts/benchmarks/tq_cross_node_bench.py b/scripts/benchmarks/tq_cross_node_bench.py index 3d543272e..4bcec16cf 100644 --- a/scripts/benchmarks/tq_cross_node_bench.py +++ b/scripts/benchmarks/tq_cross_node_bench.py @@ -409,7 +409,6 @@ def build_conf(protocol: str, master: str, device: str, segment_gib: int): 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, @@ -421,8 +420,13 @@ def build_conf(protocol: str, master: str, device: str, segment_gib: int): 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, fallback_reason="") - backend = build_mooncake_config(eff, master_address=master, global_segment_size=segment_gib * 1024**3) + # ``protocol="tcp"`` is the C1 baseline; production only builds "rdma". + backend = build_mooncake_config( + master_address=master, + device=device, + protocol=protocol, + global_segment_size=segment_gib * 1024**3, + ) return OmegaConf.create( { "controller": {"sampler": GRPOGroupNSampler(n_samples_per_prompt=1), "polling_mode": True}, diff --git a/scripts/benchmarks/tq_rdma_bench.py b/scripts/benchmarks/tq_rdma_bench.py index 67cb1f44a..5978caaef 100644 --- a/scripts/benchmarks/tq_rdma_bench.py +++ b/scripts/benchmarks/tq_rdma_bench.py @@ -94,7 +94,6 @@ def build_tq_config(config_name: str, args: argparse.Namespace, num_storage_unit 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] @@ -106,13 +105,12 @@ def build_tq_config(config_name: str, args: argparse.Namespace, num_storage_unit ) else: master_addr = args.master_address or os.environ.get("MC_MASTER_ADDRESS", "localhost:50051") - eff = EffectiveConfig( - backend="MooncakeStore", - protocol=cfg["protocol"], + backend_dict = build_mooncake_config( + master_address=master_addr, device=args.device, - fallback_reason="", + protocol=cfg["protocol"], + global_segment_size=8 * 1024**3, ) - backend_dict = build_mooncake_config(eff, master_address=master_addr, global_segment_size=8 * 1024**3) return OmegaConf.create( { diff --git a/tests/core/test_controller_tq_backend.py b/tests/core/test_controller_tq_backend.py index a17818b3f..61f338d32 100644 --- a/tests/core/test_controller_tq_backend.py +++ b/tests/core/test_controller_tq_backend.py @@ -2,13 +2,17 @@ """Controller._resolve_tq_backend: the production backend decision. -``reduce_results`` is covered in ``tests/utils/test_rdma_probe.py``; these tests -pin the Controller branch that consumes it, because that is where ``off`` -short-circuits, where ``auto`` must converge on SimpleStorage for *every* unmet -host-RDMA precondition, and where ``required`` must fail fast instead. - -Mooncake/TCP is a benchmark baseline only, so no input may ever make this method -return a MooncakeStore config whose protocol is not ``rdma``. +After the static ``/sys`` probe was removed, this method makes a pure +*configuration* decision: which mode was requested, whether the TransferQueue +correctness contract holds, whether a master endpoint is configured, and whether +the segment can hold the worst-case payload. Host-RDMA capability itself is +established later by the real cluster-wide attach handshake +(``tests/utils/test_tq_failure_paths.py``). + +These tests pin the branch structure: ``off`` must check nothing, ``auto`` must +converge on SimpleStorage for every unmet precondition, ``required`` must fail +fast on the same ones, a malformed mode must always raise, and no input may ever +produce a MooncakeStore config whose protocol is not ``rdma``. """ from __future__ import annotations @@ -18,7 +22,7 @@ import pytest -from relax.utils.rdma_probe import ProbeResult +from relax.utils.tq.lifecycle import TqConfigurationMismatch from tests.core.test_controller_s3_model_cleanup import controller from tests.utils.test_arguments_opd_teacher_colocate import ( arguments_module as _arguments_module_fixture, @@ -31,6 +35,8 @@ # dependency. arguments_module = _arguments_module_fixture +_MASTER = "master.invalid:50051" + def _config(**overrides) -> SimpleNamespace: """A config whose worst-case payload fits the default 4 GiB segment.""" @@ -50,26 +56,16 @@ def _config(**overrides) -> SimpleNamespace: return SimpleNamespace(**defaults) -def _probe(node: str, protocol: str | None = "rdma", device: str = "") -> ProbeResult: - return ProbeResult( - node=node, - checks=(), - effective_protocol=protocol, - effective_device=device, - errors=() if protocol else ("mooncake not importable",), - ) - - class _Recorder: - """Records which host-RDMA preconditions the Controller actually - consulted.""" + """Records which configuration preconditions the Controller consulted. + + A recorded ``probe`` entry would mean a static capability probe crept back + in; the sequence assertions below are what keep that from happening + silently. + """ - def __init__(self, monkeypatch, *, contract_error=None, master_error=None, probes=None): + def __init__(self, monkeypatch, *, contract_error=None, master_error=None): self.calls: list[str] = [] - self._probes = probes if probes is not None else [_probe("node-A")] - # ``build_mooncake_config`` re-resolves the endpoint from the - # environment, so the success paths need a real value there too. - monkeypatch.setenv("MC_MASTER_ADDRESS", "master.invalid:50051") def fake_contract() -> None: self.calls.append("contract") @@ -80,15 +76,10 @@ def fake_master() -> str: self.calls.append("master") if master_error is not None: raise master_error - return "master.invalid:50051" - - def fake_probe(device: str, master_address: str) -> list[ProbeResult]: - self.calls.append("probe") - return self._probes + return _MASTER monkeypatch.setattr(controller, "validate_mooncake_runtime_contract", fake_contract) monkeypatch.setattr(controller, "resolve_mooncake_master_address", fake_master) - monkeypatch.setattr(controller, "probe_cluster_nodes", fake_probe) def _resolve(config) -> dict: @@ -103,12 +94,11 @@ def _assert_simple_storage(backend: dict) -> None: class TestOffMode: - """``off`` is the untouched SimpleStorage path: it probes nothing.""" + """``off`` is the untouched SimpleStorage path: it checks nothing.""" def test_off_short_circuits_without_any_precondition_check(self, monkeypatch): recorder = _Recorder(monkeypatch) - backend = _resolve(_config(tq_rdma_mode="off")) - _assert_simple_storage(backend) + _assert_simple_storage(_resolve(_config(tq_rdma_mode="off"))) assert recorder.calls == [] def test_missing_mode_attribute_defaults_to_off(self, monkeypatch): @@ -118,6 +108,69 @@ def test_missing_mode_attribute_defaults_to_off(self, monkeypatch): _assert_simple_storage(_resolve(config)) assert recorder.calls == [] + def test_healthy_existing_controller_aborts_before_legacy_init(self, monkeypatch): + """The default path must not attach to or later close another job. + + Upstream ``tq.init`` ignores the caller's SimpleStorage config when a + named controller already exists. The exclusive-cluster check must + therefore fail before setting the local ownership flag or calling it. + """ + config = _config( + tq_rdma_mode="off", + fully_async=False, + balance_data=False, + polling_mode=False, + ) + instance = controller.Controller.__new__(controller.Controller) + instance.config = config + instance._tq_owner = None + instance._tq_legacy_init = False + + monkeypatch.setattr(controller, "resolve_sft_algo_key", lambda _config: "grpo") + monkeypatch.setattr(controller, "resolve_tq_capacity_batch_size", lambda _config: 1) + monkeypatch.setattr(controller, "GRPOGroupNSampler", lambda **_kwargs: object()) + monkeypatch.setattr( + instance, + "_resolve_tq_backend", + lambda _total_storage_size: { + "storage_backend": "SimpleStorage", + "SimpleStorage": {"total_storage_size": 1, "num_data_storage_units": 1}, + }, + ) + monkeypatch.setattr( + controller, + "reap_unusable_tq_controller", + lambda: (_ for _ in ()).throw(TqConfigurationMismatch("exclusive cluster is not clean")), + ) + monkeypatch.setattr( + controller.tq, + "init", + lambda **_kwargs: pytest.fail("existing controller must be rejected before tq.init"), + ) + + with pytest.raises(TqConfigurationMismatch, match="exclusive cluster"): + instance._initialize_data_system() + + assert instance._tq_owner is None + assert instance._tq_legacy_init is False + + +class TestNoStaticCapabilityProbe: + """The ``/sys``/GID/memlock probe is gone and must not return.""" + + def test_resolver_never_probes_hardware(self, monkeypatch): + recorder = _Recorder(monkeypatch) + _resolve(_config()) + assert recorder.calls == ["contract", "master"] + + def test_probe_helpers_are_no_longer_importable(self): + with pytest.raises(ModuleNotFoundError): + __import__("relax.utils.rdma_probe") + + def test_controller_module_holds_no_probe_symbols(self): + for name in ("probe_cluster_nodes", "reduce_results", "probe_node", "EffectiveConfig"): + assert not hasattr(controller, name), f"{name} should be gone with the static probe" + class TestAutoFallsBackForEveryUnmetPrecondition: """Gate A: in ``auto``, anything short of host RDMA yields @@ -128,22 +181,17 @@ def test_contract_failure_falls_back_before_touching_master(self, monkeypatch): _assert_simple_storage(_resolve(_config())) assert recorder.calls == ["contract"] - def test_missing_master_endpoint_falls_back_before_probing(self, monkeypatch): + def test_missing_master_endpoint_falls_back(self, monkeypatch): recorder = _Recorder(monkeypatch, master_error=RuntimeError("MC_MASTER_ADDRESS required")) _assert_simple_storage(_resolve(_config())) assert recorder.calls == ["contract", "master"] - def test_node_without_rdma_falls_back(self, monkeypatch): - _Recorder(monkeypatch, probes=[_probe("node-A"), _probe("node-B", protocol="tcp")]) - _assert_simple_storage(_resolve(_config())) - - def test_node_without_mooncake_falls_back(self, monkeypatch): - _Recorder(monkeypatch, probes=[_probe("node-A"), _probe("node-B", protocol=None)]) + def test_malformed_master_endpoint_falls_back(self, monkeypatch): + recorder = _Recorder( + monkeypatch, master_error=RuntimeError("MC_MASTER_ADDRESS is not a usable endpoint: missing host") + ) _assert_simple_storage(_resolve(_config())) - - def test_device_mismatch_falls_back(self, monkeypatch): - _Recorder(monkeypatch, probes=[_probe("node-A", device="rdma0"), _probe("node-B", device="rdma1")]) - _assert_simple_storage(_resolve(_config(tq_rdma_device="rdma0"))) + assert recorder.calls == ["contract", "master"] def test_insufficient_segment_capacity_falls_back(self, monkeypatch): """Worst-case multimodal payload far exceeds the default segment.""" @@ -151,11 +199,44 @@ def test_insufficient_segment_capacity_falls_back(self, monkeypatch): config = _config(multimodal_keys=["pixel_values"], rollout_batch_size=64, n_samples_per_prompt=8) _assert_simple_storage(_resolve(config)) - def test_all_nodes_rdma_selects_mooncake(self, monkeypatch): - _Recorder(monkeypatch, probes=[_probe("node-A"), _probe("node-B")]) - backend = _resolve(_config()) + def test_unusable_segment_size_override_falls_back(self, monkeypatch): + """Capacity validation *raises* here rather than returning a reason. + + A garbage override is a configuration failure like any other and must + not abort an ``auto`` run. + """ + _Recorder(monkeypatch) + monkeypatch.setenv("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB", "four") + _assert_simple_storage(_resolve(_config())) + + @pytest.mark.parametrize("value", ["nan", "inf", "-inf"]) + def test_non_finite_segment_size_override_falls_back(self, monkeypatch, value): + _Recorder(monkeypatch) + monkeypatch.setenv("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB", value) + _assert_simple_storage(_resolve(_config())) + + def test_missing_seq_length_falls_back(self, monkeypatch): + """Without ``seq_length`` the payload bound cannot be derived.""" + _Recorder(monkeypatch) + _assert_simple_storage(_resolve(_config(seq_length=None))) + + def test_satisfied_preconditions_select_host_rdma(self, monkeypatch): + _Recorder(monkeypatch) + backend = _resolve(_config(tq_rdma_device="mlx5_0")) assert backend["storage_backend"] == "MooncakeStore" assert backend["MooncakeStore"]["protocol"] == "rdma" + assert backend["MooncakeStore"]["device_name"] == "mlx5_0" + + def test_validated_master_endpoint_reaches_the_client_config(self, monkeypatch): + """The checked endpoint must be the one handed to Mooncake. + + ``build_mooncake_config`` must not re-read ``MC_MASTER_ADDRESS``; a + divergent env value here would surface as the wrong address. + """ + _Recorder(monkeypatch) + monkeypatch.setenv("MC_MASTER_ADDRESS", "someone.else.invalid:9999") + backend = _resolve(_config()) + assert backend["MooncakeStore"]["master_server_address"] == _MASTER class TestRequiredFailsFast: @@ -166,8 +247,6 @@ class TestRequiredFailsFast: [ ({"contract_error": RuntimeError("retry guard missing")}, "correctness contract"), ({"master_error": RuntimeError("MC_MASTER_ADDRESS required")}, "master endpoint is not configured"), - ({"probes": [_probe("node-A"), _probe("node-B", protocol="tcp")]}, "host RDMA is unavailable"), - ({"probes": [_probe("node-A"), _probe("node-B", protocol=None)]}, "host RDMA is unavailable"), ], ) def test_required_raises(self, monkeypatch, kwargs, match): @@ -186,37 +265,62 @@ def test_required_raises_on_insufficient_capacity(self, monkeypatch): with pytest.raises(RuntimeError, match="segment capacity insufficient"): _resolve(config) - def test_required_accepts_full_rdma_cluster(self, monkeypatch): - _Recorder(monkeypatch, probes=[_probe("node-A"), _probe("node-B")]) + def test_required_raises_on_unusable_segment_size_override(self, monkeypatch): + _Recorder(monkeypatch) + monkeypatch.setenv("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB", "four") + with pytest.raises(RuntimeError, match="segment-capacity configuration is unusable"): + _resolve(_config(tq_rdma_mode="required")) + + def test_required_raises_on_missing_seq_length(self, monkeypatch): + _Recorder(monkeypatch) + with pytest.raises(RuntimeError, match="segment-capacity configuration is unusable"): + _resolve(_config(tq_rdma_mode="required", seq_length=None)) + + def test_required_accepts_satisfied_preconditions(self, monkeypatch): + _Recorder(monkeypatch) backend = _resolve(_config(tq_rdma_mode="required")) assert backend["MooncakeStore"]["protocol"] == "rdma" +class TestModeValidation: + """A malformed mode is a configuration error in every mode.""" + + def test_invalid_mode_always_raises_instead_of_falling_back(self, monkeypatch): + recorder = _Recorder(monkeypatch) + with pytest.raises(ValueError, match="--tq-rdma-mode"): + _resolve(_config(tq_rdma_mode="mooncake")) + assert recorder.calls == [] + + @pytest.mark.parametrize("device", [None, ["rdma0"], "rdma0\nforged", " "]) + def test_invalid_device_always_raises_before_preconditions(self, monkeypatch, device): + recorder = _Recorder(monkeypatch) + with pytest.raises(ValueError, match="--tq-rdma-device"): + _resolve(_config(tq_rdma_device=device)) + assert recorder.calls == [] + + class TestProductionNeverSelectsMooncakeTcp: """Mooncake/TCP exists only as benchmark C1.""" - @pytest.mark.parametrize("mode", ["off", "auto"]) - @pytest.mark.parametrize( - "probes", - [ - [_probe("node-A", protocol="tcp")], - [_probe("node-A"), _probe("node-B", protocol="tcp")], - [_probe("node-A", protocol=None)], - [], - ], - ) - def test_no_input_produces_a_tcp_mooncake_backend(self, monkeypatch, mode, probes): - _Recorder(monkeypatch, probes=probes) + @pytest.mark.parametrize("mode", ["off", "auto", "required"]) + def test_selected_backend_is_rdma_or_simple(self, monkeypatch, mode): + _Recorder(monkeypatch) backend = _resolve(_config(tq_rdma_mode=mode)) if backend["storage_backend"] == "MooncakeStore": assert backend["MooncakeStore"]["protocol"] == "rdma" else: _assert_simple_storage(backend) - def test_invalid_mode_is_rejected(self, monkeypatch): - _Recorder(monkeypatch) - with pytest.raises(ValueError, match="--tq-rdma-mode"): - _resolve(_config(tq_rdma_mode="mooncake")) + @pytest.mark.parametrize( + "kwargs", + [ + {"contract_error": RuntimeError("gate failed")}, + {"master_error": RuntimeError("no endpoint")}, + ], + ) + def test_auto_fallback_never_yields_mooncake(self, monkeypatch, kwargs): + _Recorder(monkeypatch, **kwargs) + _assert_simple_storage(_resolve(_config())) def test_cli_exposes_only_mode_and_device(arguments_module): @@ -226,10 +330,7 @@ def test_cli_exposes_only_mode_and_device(arguments_module): arguments_module.get_slime_extra_args_provider()(parser) tq_flags = sorted( - option - for action in parser._actions - for option in action.option_strings - if option.startswith("--tq-") and "timeout" not in option + option for action in parser._actions for option in action.option_strings if option.startswith("--tq-") ) assert tq_flags == ["--tq-rdma-device", "--tq-rdma-mode"] @@ -238,3 +339,166 @@ def test_cli_exposes_only_mode_and_device(arguments_module): assert args.tq_rdma_device == "" assert not hasattr(args, "tq_storage_backend") assert not hasattr(args, "tq_use_gdr") + + +# --------------------------------------------------------------------------- +# _confirm_mooncake_attach: the cleanup evidence chain +# --------------------------------------------------------------------------- + + +class _AttachRecorder: + """Orders the teardown steps taken after a failed attach handshake. + + The static probe used to reject unusable clusters before any Mooncake state + existed. Now the handshake fails *after* an owner and a named controller + were created, so the ordering recorded here is what keeps a half- + initialised controller from surviving into the next ``tq.init`` (F10 hang). + """ + + def __init__( + self, + monkeypatch, + *, + failures=None, + verify_error=None, + close_error=None, + fallback_owner="fallback-owner", + ): + self.events: list[str] = [] + self.closed: list[object] = [] + self.initialized: list[object] = [] + self._fallback_owner = fallback_owner + + def fake_verify(conf, **_kwargs): + self.events.append("handshake") + if verify_error is not None: + raise verify_error + return list(failures or []) + + def fake_close(owner, **_kwargs): + self.events.append("close") + self.closed.append(owner) + if close_error is not None: + raise close_error + + def fake_initialize(conf, **_kwargs): + self.events.append("init_simple") + self.initialized.append(conf) + return controller.TqInitResult(config=conf, owner=self._fallback_owner) + + monkeypatch.setattr(controller, "verify_cluster_attach", fake_verify) + monkeypatch.setattr(controller, "close_tq_owner", fake_close) + monkeypatch.setattr(controller, "initialize_tq_with_fallback", fake_initialize) + + +def _confirm(config, *, owner="mooncake-owner", owns_controller=True, fallback_config="simple-conf"): + instance = controller.Controller.__new__(controller.Controller) + instance.config = config + init_result = controller.TqInitResult(config="mooncake-conf", owner=owner if owns_controller else None) + return instance._confirm_mooncake_attach(init_result, fallback_config) + + +class TestAttachHandshakeCleanupChain: + def test_success_keeps_mooncake_and_touches_no_cleanup(self, monkeypatch): + recorder = _AttachRecorder(monkeypatch, failures=[]) + result = _confirm(_config()) + assert result.config == "mooncake-conf" + assert result.fallback_reason == "" + assert recorder.events == ["handshake"] + + def test_auto_closes_owner_before_initializing_simple_storage(self, monkeypatch): + recorder = _AttachRecorder(monkeypatch, failures=["node-B: attach timed out"]) + result = _confirm(_config()) + # Ordering is the point: SimpleStorage must not be initialised while a + # half-initialised Mooncake controller may still be registered. + assert recorder.events == ["handshake", "close", "init_simple"] + assert result.config == "simple-conf" + assert result.fallback_reason == "attach_handshake_failed:1_failures" + + def test_cleanup_receives_this_attempts_mooncake_owner(self, monkeypatch): + recorder = _AttachRecorder(monkeypatch, failures=["node-B: attach timed out"]) + _confirm(_config(), owner="the-mooncake-owner") + assert recorder.closed == ["the-mooncake-owner"] + + def test_cleanup_failure_aborts_instead_of_falling_back(self, monkeypatch): + """If teardown fails, global TQ state is unknown. + + Starting SimpleStorage on top of it could attach to a dirty controller, + so the cleanup error must propagate. + """ + recorder = _AttachRecorder( + monkeypatch, + failures=["node-B: attach timed out"], + close_error=RuntimeError("TransferQueue owner cleanup failed: close timed out"), + ) + with pytest.raises(RuntimeError, match="owner cleanup failed"): + _confirm(_config()) + assert recorder.events == ["handshake", "close"] + assert recorder.initialized == [] + + def test_required_closes_owner_then_raises(self, monkeypatch): + recorder = _AttachRecorder(monkeypatch, failures=["node-B: protocol=tcp"]) + with pytest.raises(RuntimeError, match="attach handshake reported"): + _confirm(_config(tq_rdma_mode="required")) + assert recorder.events == ["handshake", "close"] + assert recorder.initialized == [] + + def test_attached_session_never_tears_down_a_foreign_controller(self, monkeypatch): + """A job that only attached must not close state it does not own.""" + recorder = _AttachRecorder(monkeypatch, failures=["node-B: attach timed out"]) + with pytest.raises(RuntimeError, match="attach handshake reported"): + _confirm(_config(), owns_controller=False) + assert recorder.events == ["handshake"] + assert recorder.closed == [] + + def test_unexpected_driver_exception_closes_owner_and_is_sanitized(self, monkeypatch): + secret = "worker endpoint and traceback path must stay private" + recorder = _AttachRecorder(monkeypatch, verify_error=RuntimeError(secret)) + with pytest.raises(RuntimeError, match="orchestration failed") as excinfo: + _confirm(_config()) + assert recorder.events == ["handshake", "close"] + assert recorder.initialized == [] + assert secret not in str(excinfo.value) + + def test_unexpected_driver_exception_never_closes_foreign_owner(self, monkeypatch): + recorder = _AttachRecorder(monkeypatch, verify_error=RuntimeError("private detail")) + with pytest.raises(RuntimeError, match="orchestration failed"): + _confirm(_config(), owns_controller=False) + assert recorder.events == ["handshake"] + assert recorder.closed == [] + + def test_unconfirmed_worker_isolation_closes_owner_and_aborts(self, monkeypatch): + recorder = _AttachRecorder( + monkeypatch, + verify_error=controller.TqHandshakeIsolationError("private cancellation detail"), + ) + with pytest.raises(RuntimeError, match="could not be confirmed stopped") as excinfo: + _confirm(_config()) + assert recorder.events == ["handshake", "close"] + assert recorder.initialized == [] + assert "private cancellation detail" not in str(excinfo.value) + + def test_constructor_cleanup_boundary_includes_data_system_initialization(self, monkeypatch): + """An exception after owner creation must still invoke constructor + cleanup.""" + events: list[str] = [] + + monkeypatch.setattr(controller, "resolve_sft_num_rollout", lambda _config: None) + monkeypatch.setattr(controller, "HealthManager", lambda **_kwargs: object()) + + def fail_after_owner_created(instance): + instance._tq_owner = "mooncake-owner" + raise RuntimeError("driver-side handshake orchestration failed") + + def record_cleanup(instance): + events.append(instance._tq_owner) + instance._tq_owner = None + + monkeypatch.setattr(controller.Controller, "_initialize_data_system", fail_after_owner_created) + monkeypatch.setattr(controller.Controller, "_close_data_system", record_cleanup) + + config = SimpleNamespace(use_health_check=False, max_global_restart=3) + with pytest.raises(RuntimeError, match="handshake orchestration failed"): + controller.Controller(config) + + assert events == ["mooncake-owner"] diff --git a/tests/utils/_tq_handshake_timeout_probe.py b/tests/utils/_tq_handshake_timeout_probe.py index 6a2fffafb..b49c665ae 100644 --- a/tests/utils/_tq_handshake_timeout_probe.py +++ b/tests/utils/_tq_handshake_timeout_probe.py @@ -103,9 +103,11 @@ def get_config(self): ).remote(conf) assert ray.get(controller.get_config.remote()) == conf - failures = tq_lifecycle.verify_cluster_attach(conf, timeout=2.0) + failures = tq_lifecycle.verify_cluster_attach(conf, timeout=0.3) assert len(failures) == 1 - assert "did not finish" in failures[0] + # The public failure summary is deliberately scrubbed: the underlying + # RayTaskError contains worker addresses, PIDs and traceback paths. + assert failures[0].endswith("handshake task failed (RayTaskError)") 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)) diff --git a/tests/utils/_train_actor_init_cleanup_probe.py b/tests/utils/_train_actor_init_cleanup_probe.py new file mode 100644 index 000000000..88afa80dc --- /dev/null +++ b/tests/utils/_train_actor_init_cleanup_probe.py @@ -0,0 +1,71 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Subprocess probe for train-actor cleanup after one rank fails init.""" + +from __future__ import annotations + +import os +import sys +import threading +import time +from pathlib import Path + + +def main(probe_dir: Path) -> None: + os.environ.setdefault("RAY_ENABLE_UV_RUN_RUNTIME_ENV", "0") + os.environ.pop("RAY_ADDRESS", None) + + import ray + + from relax.distributed.ray.actor_group import RayTrainGroup + + probe_dir.mkdir(parents=True, exist_ok=True) + late_marker = probe_dir / "late-mutation" + + @ray.remote(max_restarts=0) + class _InitActor: + def __init__(self, fail: bool): + self.fail = fail + + def init(self, _args, _role, **_kwargs): + if not self.fail: + return "ready" + + def mutate_late(): + time.sleep(2.0) + late_marker.write_text("dirty", encoding="utf-8") + + threading.Thread(target=mutate_late, daemon=True).start() + raise RuntimeError("expected initialization failure") + + def termination_probe(self): + threading.Event().wait() + + ray.init( + address="local", + num_cpus=2, + include_dashboard=False, + logging_level="ERROR", + _temp_dir=str(probe_dir / "ray"), + ) + try: + actors = [_InitActor.remote(True), _InitActor.remote(False)] + group = object.__new__(RayTrainGroup) + group._actor_handlers = actors + + try: + group.init_and_wait(object(), "actor") + except ray.exceptions.RayTaskError: + pass + else: + raise AssertionError("rank initialization failure was not propagated") + + assert group._actor_handlers == [] + time.sleep(2.2) + assert not late_marker.exists(), "killed actor completed a delayed process-global mutation" + finally: + ray.shutdown() + + +if __name__ == "__main__": + main(Path(sys.argv[1])) diff --git a/tests/utils/test_rdma_probe.py b/tests/utils/test_rdma_probe.py deleted file mode 100644 index ad8bc688a..000000000 --- a/tests/utils/test_rdma_probe.py +++ /dev/null @@ -1,559 +0,0 @@ -# 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 -import importlib.util -import os -from unittest import mock - -import pytest - -import relax.utils.rdma_probe as rdma_probe -from relax.utils.rdma_probe import ( - CheckResult, - EffectiveConfig, - ProbeResult, - _check_master_reachable, - _degenerate_result, - _select_dataplane_node_ids, - _split_host_port, - 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, - resolve_mooncake_master_address, - resolve_tq_capacity_batch_size, - validate_mooncake_runtime_contract, - validate_segment_capacity, -) - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -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", - node: str = "node-A", -) -> ProbeResult: - return ProbeResult( - node=node, - checks=(CheckResult("mooncake_import", True),), - effective_protocol=protocol, - effective_device=device, - ) - - -def _make_args(**kwargs) -> argparse.Namespace: - defaults = dict( - tq_rdma_mode="auto", - tq_rdma_device="", - num_data_storage_units=1, - max_staleness=0, - n_samples_per_prompt=1, - rollout_batch_size=32, - multimodal_keys=None, - seq_length=8192, - ) - defaults.update(kwargs) - return argparse.Namespace(**defaults) - - -# --------------------------------------------------------------------------- -# validate_config -# --------------------------------------------------------------------------- - - -class TestValidateConfig: - """validate_config: structural mode check before any probe.""" - - @pytest.mark.parametrize("mode", ["off", "auto", "required"]) - def test_accepts_every_supported_mode(self, mode): - assert validate_config(_make_args(tq_rdma_mode=mode)) == [] - - def test_rejects_unknown_mode(self): - """Guards configs restored from a checkpoint or built without - argparse.""" - errors = validate_config(_make_args(tq_rdma_mode="mooncake")) - assert len(errors) == 1 - assert "--tq-rdma-mode" in errors[0] - - def test_missing_attribute_defaults_to_off(self): - assert validate_config(argparse.Namespace()) == [] - - -# --------------------------------------------------------------------------- -# reduce_results -# --------------------------------------------------------------------------- - - -class TestReduceResults: - """reduce_results: per-node ProbeResult -> job-level EffectiveConfig (AND reduction). - - The only two outcomes are MooncakeStore/RDMA and SimpleStorage; Mooncake/TCP - is a benchmark baseline and must never be selected as a production backend. - """ - - def test_all_nodes_rdma(self): - eff = reduce_results( - [_make_probe(protocol="rdma"), _make_probe(protocol="rdma", node="node-B")], - requested_device="", - ) - 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_device="", - ) - assert eff.backend == "SimpleStorage" - assert "mooncake_unavailable" in eff.fallback_reason - assert "node-B" in eff.fallback_reason - - def test_one_node_without_rdma_falls_back_to_simple_not_tcp(self): - eff = reduce_results( - [_make_probe(protocol="rdma"), _make_probe(protocol="tcp", node="node-B")], - requested_device="", - ) - assert (eff.backend, eff.protocol, eff.device) == ("SimpleStorage", "tcp", "") - assert eff.fallback_reason == "rdma_unavailable:node-B" - - def test_requested_device_must_match_every_rdma_node(self): - eff = reduce_results( - [_make_probe(device="rdma0"), _make_probe(device="rdma1", node="node-B")], - requested_device="rdma0", - ) - assert eff.backend == "SimpleStorage" - assert eff.fallback_reason == "device_mismatch:rdma0" - - def test_empty_results_falls_back(self): - eff = reduce_results([], requested_device="") - assert eff.backend == "SimpleStorage" - assert eff.fallback_reason == "no probe results" - - -# --------------------------------------------------------------------------- -# 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 - - 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", 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, - 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 - 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)) - 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 - - def test_master_unreachable_blocks_mooncake(self, monkeypatch): - """An unreachable master is not a transport degradation: without it the - job cannot run MooncakeStore at all.""" - 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.example:50051") - assert result.effective_protocol is None - assert result.ok is False - assert "master unreachable" in result.errors - - -# --------------------------------------------------------------------------- -# 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, master="": 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_device="") - 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="", - errors=("master unreachable",), - ) - eff = reduce_results([_make_probe(protocol="rdma", node="n0"), unavailable], requested_device="") - assert eff.backend == "SimpleStorage" - assert eff.fallback_reason == "master_unreachable:n1" - - -# --------------------------------------------------------------------------- -# 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_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.""" - eff = EffectiveConfig(backend="MooncakeStore", protocol="rdma", device="rdma0", fallback_reason="") - 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" - ) - - def test_mooncake_config_has_hard_pin_true(self): - eff = EffectiveConfig(backend="MooncakeStore", protocol="rdma", device="rdma0", fallback_reason="") - cfg = build_mooncake_config(eff, master_address="master.example:50051") - 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 - - def test_mooncake_config_pins_host_rdma(self): - """This phase ships host RDMA only: ``use_gdr`` is always False and no - GDR staging buffer is configured.""" - eff = EffectiveConfig(backend="MooncakeStore", protocol="rdma", device="", fallback_reason="") - mc = build_mooncake_config(eff, master_address="master.example:50051")["MooncakeStore"] - assert mc["use_gdr"] is False - assert "gdr_staging_buffer_mb" not in mc - - 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="", 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", - ) - 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" - - 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") - 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) - eff = EffectiveConfig(backend="MooncakeStore", protocol="rdma", device="", 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="", 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_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) == 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_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="", 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"): - 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="", 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="", 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() diff --git a/tests/utils/test_tq_failure_paths.py b/tests/utils/test_tq_failure_paths.py index a3b208d45..6cb9ce4bb 100644 --- a/tests/utils/test_tq_failure_paths.py +++ b/tests/utils/test_tq_failure_paths.py @@ -2,11 +2,11 @@ """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 +Covers the four gaps the maintainer review called out, which +``tests/utils/tq/test_config.py`` (pure config construction/validation) and ``test_tq_dataplane_behavior.py`` (real TQ on SimpleStorage) did not: -* timeout -- controller ``get_config`` timeout and probe-task timeout +* timeout -- controller ``get_config`` timeout and attach-handshake 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), including @@ -39,7 +39,6 @@ import pytest import torch -from relax.utils.rdma_probe import ProbeResult, reduce_results from relax.utils.tq import lifecycle as tq_lifecycle @@ -68,17 +67,6 @@ def _has_real_submodule(dotted: str) -> bool: # --------------------------------------------------------------------------- -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 "", - 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") @@ -272,6 +260,9 @@ 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() + # Exception handlers require real exception classes; a bare MagicMock + # here would make ``except ray.exceptions.RayError`` invalid at runtime. + fake.exceptions = tq_lifecycle.ray.exceptions if actor is None: fake.get_actor.side_effect = ValueError("actor not found") else: @@ -290,9 +281,10 @@ def test_no_controller_is_a_noop(self, monkeypatch): assert tq_lifecycle.reap_unusable_tq_controller() is False assert killed == [] - def test_healthy_controller_is_left_alone(self, monkeypatch): + def test_healthy_controller_fails_without_being_killed(self, monkeypatch): killed = self._fake_ray(monkeypatch, actor=MagicMock(), get_result={"backend": {}}) - assert tq_lifecycle.reap_unusable_tq_controller() is False + with pytest.raises(tq_lifecycle.TqConfigurationMismatch, match="exclusive, clean Ray cluster"): + tq_lifecycle.reap_unusable_tq_controller() assert killed == [] def test_half_initialised_controller_is_reaped(self, monkeypatch): @@ -303,15 +295,41 @@ def test_half_initialised_controller_is_reaped(self, monkeypatch): 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")) + error = tq_lifecycle.ray.exceptions.GetTimeoutError("private timeout detail") + killed = self._fake_ray(monkeypatch, actor=MagicMock(), get_raises=error) 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")) + error = tq_lifecycle.ray.exceptions.RayActorError(error_msg="private actor detail") + killed = self._fake_ray(monkeypatch, actor=MagicMock(), get_raises=error) assert tq_lifecycle.reap_unusable_tq_controller() is True assert killed == ["killed"] + def test_other_ray_error_fails_closed_without_killing_or_leaking_detail(self, monkeypatch): + private_detail = "private control-plane endpoint and traceback" + error = tq_lifecycle.ray.exceptions.RayError(private_detail) + killed = self._fake_ray(monkeypatch, actor=MagicMock(), get_raises=error) + + with pytest.raises(RuntimeError, match="Failed to inspect") as excinfo: + tq_lifecycle.reap_unusable_tq_controller() + + assert killed == [] + assert private_detail not in str(excinfo.value) + + def test_stored_config_ray_error_is_sanitized(self, monkeypatch): + private_detail = "private GCS endpoint and traceback" + monkeypatch.setattr( + tq_lifecycle.ray, + "get_actor", + lambda *_args, **_kwargs: (_ for _ in ()).throw(tq_lifecycle.ray.exceptions.RayError(private_detail)), + ) + + with pytest.raises(tq_lifecycle.TqControllerInspectionError) as excinfo: + tq_lifecycle._get_stored_config() + + assert private_detail not in str(excinfo.value) + # --------------------------------------------------------------------------- # Controller lifecycle: teardown unmounts the Mooncake segment @@ -372,6 +390,13 @@ def test_attach_timeout_env_rejects_garbage(self, monkeypatch): with pytest.raises(RuntimeError, match="RELAX_TQ_ATTACH_TIMEOUT_SECONDS"): tq_lifecycle._resolve_attach_timeout() + @pytest.mark.parametrize("value", ["nan", "inf", "-inf"]) + def test_attach_timeout_env_rejects_non_finite_values(self, monkeypatch, value): + monkeypatch.setenv("RELAX_TQ_ATTACH_TIMEOUT_SECONDS", value) + with pytest.raises(RuntimeError, match="finite positive") as excinfo: + tq_lifecycle._resolve_attach_timeout() + assert value not in str(excinfo.value) + 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 @@ -406,18 +431,266 @@ def no_actor(name, namespace=None): def test_cluster_attach_handshake_worker_is_one_shot(self, monkeypatch): remote_options = {} + class _Task: + def options(self, **_kwargs): + return self + + def remote(self, *_args): + return object() + def record_remote_options(**options): remote_options.update(options) - return lambda function: function + return lambda _function: _Task() 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: ([], [])) + monkeypatch.setattr(tq_lifecycle, "_alive_node_ids", lambda: ["a" * 56]) + monkeypatch.setattr(tq_lifecycle.ray, "wait", lambda refs, **_kwargs: (list(refs), [])) + monkeypatch.setattr(tq_lifecycle.ray, "get", lambda _ref: None) assert tq_lifecycle.verify_cluster_attach({}, timeout=0.1) == [] assert remote_options["max_calls"] == 1 assert remote_options["max_retries"] == 0 + def test_cluster_attach_with_no_alive_nodes_fails_closed_without_scheduling(self, monkeypatch): + monkeypatch.setattr(tq_lifecycle, "_alive_node_ids", lambda: []) + monkeypatch.setattr( + tq_lifecycle.ray, + "remote", + lambda **_kwargs: pytest.fail("zero-node validation must not create a remote worker"), + ) + monkeypatch.setattr( + tq_lifecycle.ray, + "wait", + lambda *_args, **_kwargs: pytest.fail("zero-node validation must not call ray.wait"), + ) + + assert tq_lifecycle.verify_cluster_attach({}, timeout=0.1) == ["cluster: no alive Ray nodes discovered"] + + def test_alive_node_without_node_id_fails_closed(self, monkeypatch): + monkeypatch.setattr(tq_lifecycle.ray, "nodes", lambda: [{"Alive": True}]) + failures = tq_lifecycle.verify_cluster_attach({}, timeout=0.1) + assert failures == ["cluster: node discovery failed (RuntimeError)"] + + @staticmethod + def _capture_handshake(monkeypatch): + """Return ``(captured_args, get_worker)`` for the real nested worker. + + ``verify_cluster_attach`` defines ``_handshake`` inline, so the only + way to exercise its body — and therefore its ``finally`` detach — is to + grab the function Ray's decorator receives. + """ + captured: list[tuple] = [] + worker: list = [] + + class _Task: + def options(self, **_kwargs): + return self + + def remote(self, *args): + captured.append(args) + return object() + + def fake_remote(**_options): + def decorate(function): + worker.append(function) + return _Task() + + return decorate + + monkeypatch.setattr(tq_lifecycle.ray, "remote", fake_remote) + # Ray validates node IDs as 28-byte hex, so use a well-formed one. + monkeypatch.setattr(tq_lifecycle, "_alive_node_ids", lambda: ["a" * 56]) + monkeypatch.setattr(tq_lifecycle.ray, "wait", lambda refs, **kwargs: (list(refs), [])) + monkeypatch.setattr(tq_lifecycle.ray, "get", lambda ref: None) + return captured, worker + + def test_handshake_checks_mooncake_config_only_for_mooncake(self, monkeypatch): + """The Mooncake flag controls the manager/config assertion. + + SimpleStorage has no Mooncake protocol configuration to verify. + """ + captured, _worker = self._capture_handshake(monkeypatch) + + mooncake_conf = {"backend": {"storage_backend": "MooncakeStore"}} + assert tq_lifecycle.verify_cluster_attach(mooncake_conf, timeout=0.1) == [] + assert captured[-1][1] is True + assert captured[-1][2] == 0.1 + + simple_conf = {"backend": {"storage_backend": "SimpleStorage"}} + assert tq_lifecycle.verify_cluster_attach(simple_conf, timeout=0.1) == [] + assert captured[-1][1] is False + + def _run_worker(self, monkeypatch, *, assert_error=None): + """Execute the real ``_handshake`` body and record its call order.""" + _captured, worker = self._capture_handshake(monkeypatch) + tq_lifecycle.verify_cluster_attach({"backend": {"storage_backend": "MooncakeStore"}}, timeout=0.1) + assert worker, "Ray decorator never received the handshake function" + + events: list[str] = [] + + def fake_attach(conf, *, role, timeout): + events.append(f"attach:{role}:{timeout}") + return object() + + def fake_assert(): + events.append("assert") + if assert_error is not None: + raise assert_error + + monkeypatch.setattr(tq_lifecycle, "attach_tq_client", fake_attach) + monkeypatch.setattr(tq_lifecycle, "assert_mooncake_rdma_configured", fake_assert) + monkeypatch.setattr(tq_lifecycle, "detach_tq_client", lambda: events.append("detach")) + return worker[0], events + + def test_handshake_worker_attaches_asserts_then_detaches(self, monkeypatch): + handshake, events = self._run_worker(monkeypatch) + handshake({"backend": {"storage_backend": "MooncakeStore"}}, True, 0.1) + assert events == ["attach:attach-handshake:0.1", "assert", "detach"] + + def test_handshake_worker_detaches_even_when_the_assertion_fails(self, monkeypatch): + """A rejected transport must not leave this node's segment registered. + + Without the ``finally`` the segment would linger until the master's + ``client_ttl`` expires and break fast restarts. + """ + handshake, events = self._run_worker(monkeypatch, assert_error=RuntimeError("protocol=tcp")) + with pytest.raises(RuntimeError, match="protocol=tcp"): + handshake({"backend": {"storage_backend": "MooncakeStore"}}, True, 0.1) + assert events == ["attach:attach-handshake:0.1", "assert", "detach"] + + def test_handshake_worker_skips_the_assertion_for_simple_storage(self, monkeypatch): + handshake, events = self._run_worker(monkeypatch) + handshake({"backend": {"storage_backend": "SimpleStorage"}}, False, 0.1) + assert events == ["attach:attach-handshake:0.1", "detach"] + + def test_ready_task_failure_is_sanitized(self, monkeypatch): + self._capture_handshake(monkeypatch) + secret = "worker endpoint and traceback path must stay private" + monkeypatch.setattr(tq_lifecycle.ray, "get", lambda _ref: (_ for _ in ()).throw(RuntimeError(secret))) + + failures = tq_lifecycle.verify_cluster_attach({}, timeout=0.1) + + assert failures[0].endswith("handshake task failed (RuntimeError)") + assert secret not in failures[0] + + def test_partial_scheduling_failure_cancels_submitted_workers(self, monkeypatch): + first_ref = object() + cancellations: list[tuple[object, bool]] = [] + + class _Task: + calls = 0 + + def options(self, **_kwargs): + return self + + def remote(self, *_args): + self.calls += 1 + if self.calls == 1: + return first_ref + raise RuntimeError("private scheduling detail") + + monkeypatch.setattr(tq_lifecycle.ray, "remote", lambda **_kwargs: lambda _function: _Task()) + monkeypatch.setattr(tq_lifecycle, "_alive_node_ids", lambda: ["a" * 56, "b" * 56]) + monkeypatch.setattr( + tq_lifecycle.ray, + "cancel", + lambda ref, force: cancellations.append((ref, force)), + ) + monkeypatch.setattr(tq_lifecycle.ray, "wait", lambda refs, **_kwargs: (list(refs), [])) + + failures = tq_lifecycle.verify_cluster_attach({}, timeout=0.1) + + assert failures == ["cluster: handshake scheduling failed (RuntimeError)"] + assert cancellations == [(first_ref, True)] + + def test_wait_failure_cancels_and_confirms_submitted_workers(self, monkeypatch): + captured, _worker = self._capture_handshake(monkeypatch) + wait_calls = 0 + cancellations: list[tuple[object, bool]] = [] + + def fake_wait(submitted, **_kwargs): + nonlocal wait_calls + wait_calls += 1 + if wait_calls == 1: + raise RuntimeError("private wait detail") + return list(submitted), [] + + monkeypatch.setattr(tq_lifecycle.ray, "wait", fake_wait) + monkeypatch.setattr( + tq_lifecycle.ray, + "cancel", + lambda ref, force: cancellations.append((ref, force)), + ) + assert captured == [] + + failures = tq_lifecycle.verify_cluster_attach({}, timeout=0.1) + + assert failures == ["cluster: handshake wait failed (RuntimeError)"] + assert len(cancellations) == 1 and cancellations[0][1] is True + + def test_unconfirmed_pending_worker_aborts_fallback_boundary(self, monkeypatch): + self._capture_handshake(monkeypatch) + monkeypatch.setattr(tq_lifecycle.ray, "wait", lambda refs, **_kwargs: ([], list(refs))) + monkeypatch.setattr(tq_lifecycle.ray, "cancel", lambda _ref, force: None) + + with pytest.raises(tq_lifecycle.TqHandshakeIsolationError, match="could not be confirmed stopped"): + tq_lifecycle.verify_cluster_attach({}, timeout=0.1) + + def test_cancel_failure_aborts_fallback_boundary_without_leaking_detail(self, monkeypatch): + self._capture_handshake(monkeypatch) + private_detail = "private worker address and traceback path" + monkeypatch.setattr(tq_lifecycle.ray, "wait", lambda refs, **_kwargs: ([], list(refs))) + monkeypatch.setattr( + tq_lifecycle.ray, + "cancel", + lambda _ref, force: (_ for _ in ()).throw(RuntimeError(private_detail)), + ) + + with pytest.raises(tq_lifecycle.TqHandshakeIsolationError) as excinfo: + tq_lifecycle.verify_cluster_attach({}, timeout=0.1) + assert private_detail not in str(excinfo.value) + + +class TestAssertMooncakeRdmaConfigured: + """The client field proves configured intent, not negotiated transport. + + ``tq.init`` ignores the caller's conf when attaching to an existing + controller, so the manager and configured protocol must still match the + job-level contract. Wire proof remains a benchmark responsibility. + """ + + class MooncakeStorageManager: + """Name matters: the production check compares + ``type(...).__name__``.""" + + def __init__(self, storage_client): + self.storage_client = storage_client + + def _patch_client(self, monkeypatch, manager): + fake = MagicMock() + fake.get_client.return_value = MagicMock(storage_manager=manager) + monkeypatch.setattr(tq_lifecycle, "tq", fake) + + def test_accepts_mooncake_rdma(self, monkeypatch): + self._patch_client(monkeypatch, self.MooncakeStorageManager(MagicMock(protocol="rdma"))) + tq_lifecycle.assert_mooncake_rdma_configured() + + def test_rejects_non_mooncake_manager(self, monkeypatch): + self._patch_client(monkeypatch, MagicMock()) + with pytest.raises(RuntimeError, match="not MooncakeStorageManager"): + tq_lifecycle.assert_mooncake_rdma_configured() + + def test_rejects_missing_storage_client(self, monkeypatch): + self._patch_client(monkeypatch, self.MooncakeStorageManager(None)) + with pytest.raises(RuntimeError, match="no storage_client"): + tq_lifecycle.assert_mooncake_rdma_configured() + + @pytest.mark.parametrize("protocol", ["tcp", None, ""]) + def test_rejects_non_rdma_protocol(self, monkeypatch, protocol): + self._patch_client(monkeypatch, self.MooncakeStorageManager(MagicMock(protocol=protocol))) + with pytest.raises(RuntimeError, match="not configured for protocol=rdma"): + tq_lifecycle.assert_mooncake_rdma_configured() + 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.""" @@ -525,13 +798,16 @@ def _mooncake_conf(protocol: str, *, use_gdr: bool = False) -> dict: } @staticmethod - def _patch_transaction(monkeypatch, *, existed: bool, init_effects: list[object], stored_conf=None): + def _patch_transaction(monkeypatch, *, init_effects: list[object], reap_error: BaseException | None = 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_reap(): + calls["reap"].append(True) + if reap_error is not None: + raise reap_error + + monkeypatch.setattr(tq_lifecycle, "reap_unusable_tq_controller", fake_reap) def fake_start(conf, *, timeout): calls["attempts"].append(conf) @@ -545,101 +821,24 @@ def fake_start(conf, *, timeout): 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"]) + calls = self._patch_transaction(monkeypatch, 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): + @pytest.mark.parametrize("mode", ["off", "auto", "required"]) + def test_healthy_existing_controller_fails_exclusive_without_starting_owner(self, monkeypatch, mode): 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"] == [] + mismatch = tq_lifecycle.TqConfigurationMismatch("exclusive cluster is not clean") + calls = self._patch_transaction(monkeypatch, init_effects=[], reap_error=mismatch) - 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"): + with pytest.raises(tq_lifecycle.TqConfigurationMismatch, match="exclusive cluster"): tq_lifecycle.initialize_tq_with_fallback( requested, - mode="auto", + mode=mode, 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_attach_rejects_legacy_gdr_controller(self, monkeypatch): - """A controller left behind by a GDR-enabled build is not compatible. - - Relax now always requests host RDMA (``use_gdr=False``), but upstream - ``tq.init`` ignores the caller's conf when attaching, so accepting such - a controller would silently run this worker on the unverified GDR path. - """ - requested = self._mooncake_conf("rdma") - stored = self._mooncake_conf("rdma", use_gdr=True) - 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="required") - 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): @@ -647,7 +846,6 @@ def test_auto_cleans_failed_mooncake_then_retries_simple_once(self, monkeypatch) 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) @@ -661,7 +859,6 @@ def test_required_cleans_failed_init_without_fallback(self, monkeypatch): fallback = self._conf("SimpleStorage") calls = self._patch_transaction( monkeypatch, - existed=False, init_effects=[RuntimeError("master unavailable")], ) with pytest.raises(RuntimeError, match="master unavailable"): @@ -673,7 +870,6 @@ def test_timeout_auto_retries_only_after_isolated_owner_cleanup(self, monkeypatc 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) @@ -681,16 +877,6 @@ 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 @@ -763,6 +949,45 @@ def test_failed_owner_cleanup_respects_controller_owner_token(self, monkeypatch, assert stopped == [owner] assert bool(killed) is should_kill + def test_failed_owner_cleanup_does_not_kill_when_controller_inspection_fails(self, monkeypatch): + 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: (_ for _ in ()).throw( + tq_lifecycle.TqControllerInspectionError("controller inspection failed (RayError)") + ), + ) + monkeypatch.setattr(tq_lifecycle, "kill_tq_controller_and_wait", lambda: killed.append(True)) + + with pytest.raises(tq_lifecycle.TqControllerInspectionError): + tq_lifecycle._cleanup_failed_owner(owner, "ours") + + assert stopped == [owner] + assert killed == [] + + def test_failed_owner_cleanup_reaps_proven_half_initialised_controller(self, monkeypatch): + owner = _FakeOwner() + killed: list[bool] = [] + monkeypatch.setattr(tq_lifecycle.ray, "get", lambda ref, timeout: None) + monkeypatch.setattr(tq_lifecycle, "_stop_owner_actor", lambda _handle: None) + monkeypatch.setattr( + tq_lifecycle, + "_get_stored_config", + lambda timeout: (_ for _ in ()).throw( + tq_lifecycle.TqControllerMissingConfig("controller returned no config") + ), + ) + monkeypatch.setattr(tq_lifecycle, "kill_tq_controller_and_wait", lambda: killed.append(True)) + + tq_lifecycle._cleanup_failed_owner(owner, "ours") + + assert killed == [True] + def test_owner_close_failure_still_reaps_global_controller(self, monkeypatch): owner = _FakeOwner() stopped: list[object] = [] @@ -976,40 +1201,6 @@ def test_real_mooncake_capacity_overflow_is_bounded_and_loud(): assert "batch_upsert_from failed" in detail or "capacity" in detail.lower(), detail -# --------------------------------------------------------------------------- -# 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_device="") - 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_device="") - 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_device="") - assert eff.backend == "SimpleStorage" - - def test_one_node_without_rdma_degrades_whole_job_to_simple(self): - eff = reduce_results([_probe("a"), _probe("b", protocol="tcp")], requested_device="") - assert (eff.backend, eff.protocol) == ("SimpleStorage", "tcp") - assert eff.fallback_reason - - # --------------------------------------------------------------------------- # Byte-exactness on MooncakeStore (was SimpleStorage-only) # --------------------------------------------------------------------------- diff --git a/tests/utils/test_train_actor_init_cleanup.py b/tests/utils/test_train_actor_init_cleanup.py new file mode 100644 index 000000000..df2aff1c4 --- /dev/null +++ b/tests/utils/test_train_actor_init_cleanup.py @@ -0,0 +1,140 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""CPU-only checks for fail-closed Ray train-actor initialization.""" + +from __future__ import annotations + +import os +import subprocess +import sys +import tempfile + +import pytest +import ray + +from relax.distributed.ray.actor_group import RayTrainGroup + + +class _RemoteMethod: + def __init__(self, result=None, error: BaseException | None = None): + self.result = result + self.error = error + + def remote(self, *args, **kwargs): + if self.error is not None: + raise self.error + return self.result + + +class _FakeActor: + def __init__(self, probe_ref=None, probe_error: BaseException | None = None): + self.termination_probe = _RemoteMethod(probe_ref, probe_error) + + +def _group(*actors) -> RayTrainGroup: + group = object.__new__(RayTrainGroup) + group._actor_handlers = list(actors) + return group + + +def test_successful_initialization_preserves_actor_group(monkeypatch): + actor = _FakeActor() + group = _group(actor) + refs = [object(), object()] + values = {refs[0]: "rank-0", refs[1]: "rank-1"} + monkeypatch.setattr(group, "async_init", lambda *args, **kwargs: refs) + monkeypatch.setattr(ray, "wait", lambda pending, num_returns: ([pending[0]], pending[1:])) + monkeypatch.setattr(ray, "get", lambda ref: values[ref]) + monkeypatch.setattr(ray, "kill", lambda *_args, **_kwargs: pytest.fail("success must not kill actors")) + + assert group.init_and_wait(object(), "actor") == ["rank-0", "rank-1"] + assert group._actor_handlers == [actor] + + +def test_first_rank_failure_kills_and_confirms_every_actor(monkeypatch): + actors = [_FakeActor("probe-0"), _FakeActor("probe-1")] + group = _group(*actors) + init_refs = ["init-0", "init-1"] + killed = [] + monkeypatch.setattr(group, "async_init", lambda *args, **kwargs: init_refs) + + def fake_wait(refs, *, num_returns, timeout=None): + if timeout is None: + # rank 1 fails before rank 0 finishes; cleanup must begin without + # waiting for rank 0. + return ["init-1"], ["init-0"] + return list(refs), [] + + def fake_get(ref): + if ref == "init-1": + raise RuntimeError("rank initialization failed") + if str(ref).startswith("probe-"): + raise ray.exceptions.RayActorError(error_msg="actor terminated") + pytest.fail(f"unexpected ray.get({ref!r})") + + monkeypatch.setattr(ray, "wait", fake_wait) + monkeypatch.setattr(ray, "get", fake_get) + monkeypatch.setattr(ray, "kill", lambda actor, no_restart: killed.append((actor, no_restart))) + + with pytest.raises(RuntimeError, match="rank initialization failed"): + group.init_and_wait(object(), "actor") + + assert killed == [(actors[0], True), (actors[1], True)] + assert group._actor_handlers == [] + + +def test_kill_failure_still_attempts_every_actor_and_fails_closed(monkeypatch): + actors = [_FakeActor("probe-0"), _FakeActor("probe-1")] + group = _group(*actors) + killed = [] + + def fake_kill(actor, *, no_restart): + killed.append(actor) + if actor is actors[0]: + raise RuntimeError("private control-plane detail") + + monkeypatch.setattr(ray, "kill", fake_kill) + monkeypatch.setattr(ray, "wait", lambda refs, **kwargs: (list(refs), [])) + monkeypatch.setattr( + ray, + "get", + lambda _ref: (_ for _ in ()).throw(ray.exceptions.RayActorError(error_msg="actor terminated")), + ) + + with pytest.raises(RuntimeError, match="Failed to confirm train actor cleanup") as excinfo: + group._terminate_failed_init(timeout=0.1) + + assert killed == actors + assert "private control-plane detail" not in str(excinfo.value) + assert group._actor_handlers == actors + + +def test_pending_termination_probe_keeps_group_and_fails_closed(monkeypatch): + actor = _FakeActor("probe") + group = _group(actor) + monkeypatch.setattr(ray, "kill", lambda *_args, **_kwargs: None) + monkeypatch.setattr(ray, "wait", lambda refs, **kwargs: ([], list(refs))) + + with pytest.raises(RuntimeError, match=r"1 task\(s\) remained pending"): + group._terminate_failed_init(timeout=0.1) + + assert group._actor_handlers == [actor] + + +def test_real_actor_failure_cannot_mutate_state_after_cleanup(): + """A killed actor process cannot complete a delayed daemon-thread write.""" + 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__), "..", "..")) + with tempfile.TemporaryDirectory(prefix="train-init-ray-") as probe_dir: + result = subprocess.run( + [sys.executable, "-m", "tests.utils._train_actor_init_cleanup_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}" diff --git a/tests/utils/tq/_payload_assertions.py b/tests/utils/tq/_payload_assertions.py index 4ca3778e4..7b797e6f4 100644 --- a/tests/utils/tq/_payload_assertions.py +++ b/tests/utils/tq/_payload_assertions.py @@ -10,6 +10,7 @@ from __future__ import annotations import hashlib +import struct from typing import Any import numpy as np @@ -37,6 +38,11 @@ def _scalar_digest(value: Any) -> LeafDigest: raw = value elif isinstance(value, str): raw = value.encode("utf-8") + elif isinstance(value, float): + # ``repr(float('nan'))`` discards the NaN payload bits. Pack the + # Python double directly so distinct NaNs and signed zero remain + # byte-distinguishable just like tensor/ndarray leaves. + raw = struct.pack("!d", value) else: raw = repr(value).encode("utf-8") return (f"py.{type(value).__name__}", "", hashlib.sha256(raw).hexdigest()) @@ -50,6 +56,15 @@ def _unwrap_non_tensor(value: Any) -> Any: return value +def _dict_child_path(prefix: str, key: Any) -> str: + """Render a dict key without colliding with nested/list paths.""" + if not isinstance(key, str): + raise TypeError(f"Unsupported payload dict key at {prefix}: {type(key).__name__}") + if key.isidentifier(): + return f"{prefix}.{key}" + return f"{prefix}[{key!r}]" + + def leaf_digests(payload: Any, prefix: str = "payload") -> dict[str, LeafDigest]: """Map every supported leaf to ``(dtype, shape, raw-byte SHA-256)``.""" payload = _unwrap_non_tensor(payload) @@ -63,8 +78,11 @@ def leaf_digests(payload: Any, prefix: str = "payload") -> dict[str, LeafDigest] elif isinstance(payload, np.ndarray): digests[prefix] = _ndarray_digest(payload) elif isinstance(payload, dict): + for key in payload: + if not isinstance(key, str): + raise TypeError(f"Unsupported payload dict key at {prefix}: {type(key).__name__}") for key in sorted(payload): - digests.update(leaf_digests(payload[key], f"{prefix}.{key}")) + digests.update(leaf_digests(payload[key], _dict_child_path(prefix, key))) elif isinstance(payload, (list, tuple)): for index, item in enumerate(payload): digests.update(leaf_digests(item, f"{prefix}[{index}]")) diff --git a/tests/utils/tq/test_config.py b/tests/utils/tq/test_config.py new file mode 100644 index 000000000..4af93423b --- /dev/null +++ b/tests/utils/tq/test_config.py @@ -0,0 +1,382 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Unit tests for TransferQueue backend config construction and validation. + +CPU-only: nothing here starts Ray, TransferQueue, or a Mooncake client. These +tests cover everything Relax decides *before* anything is initialised — the +requested mode, the config dicts handed to ``tq.init``, the master endpoint +format, and the segment-capacity pre-check. + +Actual host-RDMA capability is not testable here by design: it is established by +the real cluster-wide attach handshake, covered in +``tests/utils/test_tq_failure_paths.py``. +""" + +from __future__ import annotations + +import argparse +import importlib.util +import os + +import pytest + +from relax.utils.tq.config import ( + _split_host_port, + build_mooncake_config, + build_simple_storage_config, + estimate_payload_bytes, + resolve_mooncake_master_address, + resolve_tq_capacity_batch_size, + validate_config, + validate_mooncake_runtime_contract, + validate_segment_capacity, +) + + +_MASTER = "master.example:50051" + + +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_args(**kwargs) -> argparse.Namespace: + defaults = dict( + tq_rdma_mode="auto", + tq_rdma_device="", + num_data_storage_units=1, + max_staleness=0, + n_samples_per_prompt=1, + rollout_batch_size=32, + multimodal_keys=None, + seq_length=8192, + ) + defaults.update(kwargs) + return argparse.Namespace(**defaults) + + +class TestValidateConfig: + """Minimal mode/device checks left after the static probe was removed.""" + + @pytest.mark.parametrize("mode", ["off", "auto", "required"]) + def test_accepts_every_supported_mode(self, mode): + assert validate_config(_make_args(tq_rdma_mode=mode)) == [] + + @pytest.mark.parametrize("mode", ["mooncake", "auto\nprivate", None, ["auto"]]) + def test_rejects_unknown_mode_without_echoing_it(self, mode): + """Guards configs restored from a checkpoint or built without + argparse.""" + errors = validate_config(_make_args(tq_rdma_mode=mode)) + assert len(errors) == 1 + assert "--tq-rdma-mode" in errors[0] + assert repr(mode) not in errors[0] + + def test_missing_attribute_defaults_to_off(self): + assert validate_config(argparse.Namespace()) == [] + + @pytest.mark.parametrize("device", [None, ["rdma0"], "rdma0\nforged", " "]) + def test_rejects_non_string_or_whitespace_device(self, device): + errors = validate_config(_make_args(tq_rdma_device=device)) + assert len(errors) == 1 + assert "--tq-rdma-device" in errors[0] + assert repr(device) not in errors[0] + + @pytest.mark.parametrize("device", ["", "rdma0", "mlx5_0"]) + def test_accepts_empty_or_printable_device_name(self, device): + assert validate_config(_make_args(tq_rdma_device=device)) == [] + + +class TestMasterEndpoint: + """``MC_MASTER_ADDRESS`` is deployment configuration, validated by format + only. + + No DNS lookup and no connection attempt: reachability is proven later by + the real attach, and re-adding a network probe here would recreate the + capability heuristic this phase deleted. + """ + + def test_required_and_returned_from_env(self, monkeypatch): + monkeypatch.setenv("MC_MASTER_ADDRESS", _MASTER) + assert resolve_mooncake_master_address() == _MASTER + + def test_missing_is_rejected(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() + + @pytest.mark.parametrize( + "address", + [ + "master.example", # no port + "master.example:", # empty port + "master.example:abc", # non-numeric port + ":50051", # no host + "master.example:0", # port below range + "master.example:65536", # port above range + "fe80::1", # bare IPv6: rpartition would yield port=1 + "fe80::1:50051", # bare IPv6 with port + "[fe80::1]", # bracketed host, no port + "[fe80::1]50051", # bracketed host, missing colon + "master\nprivate:50051", # embedded control character + "master name:50051", # embedded whitespace + ], + ) + def test_malformed_endpoints_are_rejected(self, monkeypatch, address): + monkeypatch.setenv("MC_MASTER_ADDRESS", address) + with pytest.raises(RuntimeError, match="not a usable endpoint"): + resolve_mooncake_master_address() + + def test_rejection_never_echoes_the_endpoint(self, monkeypatch): + """The Controller logs this error verbatim. + + An internal hostname or IP is deployment detail that must not reach job + logs, so the message names the defect only. + """ + secret = "prod-master-07.internal.corp:0" + monkeypatch.setenv("MC_MASTER_ADDRESS", secret) + with pytest.raises(RuntimeError) as excinfo: + resolve_mooncake_master_address() + message = str(excinfo.value) + assert "prod-master-07" not in message + assert "internal.corp" not in message + assert "port is outside" in message + + def test_accepts_hostname_and_bracketed_ipv6(self): + assert _split_host_port(_MASTER) == ("master.example", 50051) + assert _split_host_port("[2001:db8::1]:50051") == ("2001:db8::1", 50051) + + +class TestBackendConfigDicts: + """The dicts handed to ``tq.init``.""" + + 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_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.""" + assert build_mooncake_config(master_address=_MASTER)["storage_backend"] == "MooncakeStore" + assert ( + build_simple_storage_config(total_storage_size=1, num_data_storage_units=1)["storage_backend"] + == "SimpleStorage" + ) + + def test_mooncake_defaults_to_host_rdma(self): + """Production never names a protocol: the default is the only one it + ships.""" + mc = build_mooncake_config(master_address=_MASTER, device="rdma0")["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["master_server_address"] == _MASTER + + def test_empty_device_is_left_to_mooncake(self): + mc = build_mooncake_config(master_address=_MASTER)["MooncakeStore"] + assert mc["device_name"] == "" + + def test_tcp_is_reachable_only_by_explicit_request(self): + """Mooncake/TCP survives as benchmark C1, never as a production + default.""" + mc = build_mooncake_config(master_address=_MASTER, protocol="tcp")["MooncakeStore"] + assert mc["protocol"] == "tcp" + + def test_mooncake_config_pins_host_rdma(self): + """This phase ships host RDMA only: ``use_gdr`` is always False and no + GDR staging buffer is configured.""" + mc = build_mooncake_config(master_address=_MASTER)["MooncakeStore"] + assert mc["use_gdr"] is False + assert "gdr_staging_buffer_mb" not in mc + + def test_master_address_is_not_re_read_from_env(self, monkeypatch): + """The validated endpoint must be the one the client receives.""" + monkeypatch.setenv("MC_MASTER_ADDRESS", "other.example:9999") + mc = build_mooncake_config(master_address=_MASTER)["MooncakeStore"] + assert mc["master_server_address"] == _MASTER + + @pytest.mark.parametrize("segment_size", [0, -1, True, 1.5]) + def test_explicit_segment_size_must_be_a_positive_integer(self, segment_size): + with pytest.raises(ValueError, match="positive integer"): + build_mooncake_config(master_address=_MASTER, global_segment_size=segment_size) + + +class TestCorrectnessContract: + """The Mooncake loss-prevention gate.""" + + @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() + + @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" + + 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") + with pytest.raises(RuntimeError, match="MC_STORE_MEMCPY"): + validate_mooncake_runtime_contract() + + def test_contract_rejection_does_not_echo_memcpy_override(self, monkeypatch): + private_detail = "1\nprivate deployment detail" + monkeypatch.setenv("MC_STORE_MEMCPY", private_detail) + with pytest.raises(RuntimeError, match="MC_STORE_MEMCPY") as excinfo: + validate_mooncake_runtime_contract() + assert private_detail not in str(excinfo.value) + + 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" + + @pytest.mark.skipif( + not _REAL_TQ_STORAGE, + reason="needs real TransferQueue storage submodules; CPU CI uses a single-file transfer_queue stub", + ) + def test_unreadable_source_becomes_a_runtime_error(self, monkeypatch): + """A compiled/stripped install cannot prove put/notify ordering. + + The failure must stay inside the RuntimeError boundary the Controller + catches; a bare OSError would abort an ``auto`` run that is supposed to + fall back. + """ + import inspect as inspect_module + + from relax.utils.tq import config as config_module + + def raise_oserror(_obj): + raise OSError("could not get source code") + + monkeypatch.setattr(config_module.inspect, "getsource", raise_oserror) + assert inspect_module is not None # the patch targets the module's own alias + with pytest.raises(RuntimeError, match="Cannot verify TransferQueue put/notify ordering"): + validate_mooncake_runtime_contract() + + +class TestSegmentCapacity: + """Configuration-level capacity pre-check (kept: not a hardware probe).""" + + def test_text_only_passes(self): + assert validate_segment_capacity(_make_args(multimodal_keys=None)) is None + + def test_multimodal_large_batch_fails(self): + args = _make_args( + multimodal_keys=["pixel_values"], rollout_batch_size=256, n_samples_per_prompt=8, max_staleness=1 + ) + err = validate_segment_capacity(args) + assert err is not None + assert "insufficient" in err.lower() + + def test_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 + ) + err = validate_segment_capacity(args) + assert err is not None and "RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB" in err + + def test_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 + ) + monkeypatch.setenv("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB", "8") + assert validate_segment_capacity(args) is None + + @pytest.mark.parametrize("value", ["four", "-1", "nan", "inf", "-inf"]) + def test_segment_size_env_override_rejects_unusable_values(self, monkeypatch, value): + from relax.utils.tq.config import resolve_global_segment_size + + monkeypatch.setenv("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB", value) + with pytest.raises(RuntimeError, match="finite positive") as excinfo: + resolve_global_segment_size() + assert value not in str(excinfo.value) + + def test_segment_size_env_override_cannot_round_down_to_zero_bytes(self, monkeypatch): + from relax.utils.tq.config import resolve_global_segment_size + + monkeypatch.setenv("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB", "1e-20") + with pytest.raises(RuntimeError, match="at least one byte"): + resolve_global_segment_size() + + +class TestPayloadEstimate: + """The token-budget bound behind the capacity check.""" + + def test_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. + assert estimate_payload_bytes(_make_args(multimodal_keys=None)) == 32 * 1 * 8192 * 32 + + def test_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_requires_seq_length(self): + with pytest.raises(RuntimeError, match="seq_length"): + estimate_payload_bytes(_make_args(seq_length=None)) + + 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, + ) + err = validate_segment_capacity(args) + assert err is not None + assert "effective_batch=64" in err diff --git a/tests/utils/tq/test_payload_assertions.py b/tests/utils/tq/test_payload_assertions.py index b930c7d0a..027b86fcf 100644 --- a/tests/utils/tq/test_payload_assertions.py +++ b/tests/utils/tq/test_payload_assertions.py @@ -4,12 +4,14 @@ from __future__ import annotations +import struct from typing import Any import numpy as np import pytest import torch +from relax.utils.payload_digest import leaf_digests as benchmark_leaf_digests from tests.utils.tq._payload_assertions import diff_digests, leaf_digests @@ -43,6 +45,29 @@ def test_leaf_digests_distinguishes_raw_bytes_from_value_equality(): ] +@pytest.mark.parametrize("digest_fn", [leaf_digests, benchmark_leaf_digests], ids=["test-helper", "benchmark-helper"]) +def test_scalar_nan_payload_bits_are_not_collapsed_by_repr(digest_fn): + first = struct.unpack("!d", bytes.fromhex("7ff8000000000001"))[0] + second = struct.unpack("!d", bytes.fromhex("7ff8000000000002"))[0] + + assert repr(first) == repr(second) == "nan" + assert digest_fn(first) != digest_fn(second) + + +@pytest.mark.parametrize("digest_fn", [leaf_digests, benchmark_leaf_digests], ids=["test-helper", "benchmark-helper"]) +def test_dict_paths_do_not_collapse_dotted_keys_into_nested_keys(digest_fn): + digests = digest_fn({"a.b": 1, "a": {"b": 2}}) + + assert len(digests) == 2 + assert set(digests) == {"payload['a.b']", "payload.a.b"} + + +@pytest.mark.parametrize("digest_fn", [leaf_digests, benchmark_leaf_digests], ids=["test-helper", "benchmark-helper"]) +def test_non_string_dict_keys_fail_loudly(digest_fn): + with pytest.raises(TypeError, match="Unsupported payload dict key at payload: int"): + digest_fn({1: "value"}) + + def test_leaf_digests_preserves_dtype_shape_and_nested_tensor_rows(): with pytest.warns(UserWarning, match="prototype stage"): nested = torch.nested.nested_tensor( From 7265caf94909c81df45f911ad01bb2ef8f4e83c3 Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Mon, 24 Aug 2026 02:06:33 +0800 Subject: [PATCH 23/30] refactor(tq): simplify exclusive-cluster lifecycle --- docs/draft/transfer_queue_rdma.md | 7 +- relax/backends/megatron/actor.py | 7 +- relax/components/actor.py | 1 - relax/components/actor_fwd.py | 1 - relax/components/advantages.py | 1 - relax/components/base.py | 7 +- relax/components/critic.py | 1 - relax/components/rollout.py | 1 - relax/components/sft.py | 1 - relax/core/controller.py | 10 +- relax/distributed/ray/rollout.py | 7 +- relax/utils/tq/lifecycle.py | 337 ++++++++---------- tests/core/test_controller_tq_backend.py | 24 +- tests/utils/_tq_handshake_timeout_probe.py | 14 + tests/utils/test_tq_failure_paths.py | 377 +++++++++++++-------- 15 files changed, 397 insertions(+), 399 deletions(-) diff --git a/docs/draft/transfer_queue_rdma.md b/docs/draft/transfer_queue_rdma.md index 23aaa1c87..3a4ade999 100644 --- a/docs/draft/transfer_queue_rdma.md +++ b/docs/draft/transfer_queue_rdma.md @@ -108,14 +108,15 @@ setsid mooncake_master -rpc_port=50051 -metrics_port=9004 > /var/log/mooncake_ma ## 资源所有权与安全清理 -首期按**单任务独占 Ray 集群**实现，**不承诺同一节点上多个 Relax job 并发**：多 job 并发、端口租约、master 共享机制都不在首期范围内。 +首期按**单任务独占 Ray 集群**实现：同一 Ray cluster 在初始化和运行期间只能有一个 Relax job，操作者不得并发启动第二个 job。多 job admission、端口租约和 master 共享机制都不在首期范围内；删除 token/signature 后，`reap → tq.init` 不再承担并发 initializer 仲裁。 清理只动本作业拥有的资源： - 不使用任何 `pkill` / `killall` - `tq.init` 之前会检查已存在的 `TransferQueueController` 命名 actor：**只有取不到 config（半初始化）、明确超时或 actor 已死时才回收**；其他 GCS/control-plane 异常只做脱敏后中止，不据此杀 actor - 健康的既有 controller 保持不动，但启动会明确失败，**不会 attach，也不会在退出时关闭它**；操作者应先停止前一作业并清理其 TQ 状态，确保 Ray 集群干净 -- 首次初始化在专用 owner actor 中执行，config 内保存随机 owner token；若初始化失败，只有 token 仍匹配时才回收该次初始化创建的全局 actor；若并发 initializer 抢先创建了健康 controller，本次启动只做本地 detach 后失败 +- 首次初始化在专用 owner actor 中执行；owner 调度成功后才进入 candidate 初始化。初始化失败时先 bounded best-effort close，再 force-kill owner，并通过预先排队且永不正常返回的 termination probe 确认 owner 进程已经终止；随后才清理并确认 named controller 注销。任一步无法确认都会中止，禁止进入 fallback +- 因为集群在每次初始化前已经通过 clean-cluster gate，owner 失败后出现的 controller 按本次尝试的残留处理。该判断依赖上面的“禁止并发启动第二个 Relax job”硬前提，不提供 concurrent initializer winner 兼容 - 全局 `tq.close()` 只能由 owner actor 调用；actor、critic、rollout 等附加 worker 只能做本地 detach - 任一长生命周期 train actor 初始化失败时，driver 会 force-kill 整个 train actor group，并通过预先排队的终止 probe 确认 actor task 已进入终态后再传播失败，防止超时的原生初始化线程稍后修改可复用进程的 TQ 全局状态 - master 进程始终不被 Relax 触碰 @@ -151,7 +152,7 @@ Mock/本机测试和真实双节点 RDMA 测试必须分别报告，前者不能 | 层级 | 验证内容 | 通过标准 | |---|---|---| -| CI/mock | 模式校验、master 端点格式、容量预检、owner 超时/清理/token、attach 握手的 manager/config 契约与拆除顺序、auto/required、有限重试、写失败不发布状态 | `tests/utils/tq/test_config.py`、`tests/core/test_controller_tq_backend.py` 与 `tests/utils/test_tq_failure_paths.py` 全部通过；真机项允许明确 skip | +| CI/mock | 模式校验、master 端点格式、容量预检、owner 超时与终止确认、controller 清理、attach 握手的 manager/config 契约与拆除顺序、auto/required、有限重试、写失败不发布状态 | `tests/utils/tq/test_config.py`、`tests/core/test_controller_tq_backend.py` 与 `tests/utils/test_tq_failure_paths.py` 全部通过；真机项允许明确 skip | | 本机 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 兜底）；交付报告须注明真实档在何处跑过 | diff --git a/relax/backends/megatron/actor.py b/relax/backends/megatron/actor.py index 936184f90..36cea659a 100644 --- a/relax/backends/megatron/actor.py +++ b/relax/backends/megatron/actor.py @@ -155,12 +155,10 @@ 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. - generation = getattr(self, "_tq_client_generation", None) - if getattr(self, "data_system_client", None) is None or generation is None: + if getattr(self, "data_system_client", None) is None: return try: - detach_tq_client(generation) - self._tq_client_generation = None + detach_tq_client() self.data_system_client = None except Exception: # destructor must never raise (interpreter shutdown) return @@ -204,7 +202,6 @@ def _init( self.data_system_client = attach_tq_client( args.tq_config, 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 933fef015..d5273e0ec 100644 --- a/relax/components/actor.py +++ b/relax/components/actor.py @@ -74,7 +74,6 @@ def __init__( self.data_system_client = attach_tq_client( self.config.tq_config, role=self.role, - lease_owner=self, ) self.steps = self.actor_model.init_and_wait( diff --git a/relax/components/actor_fwd.py b/relax/components/actor_fwd.py index bb303fc0b..831665616 100644 --- a/relax/components/actor_fwd.py +++ b/relax/components/actor_fwd.py @@ -38,7 +38,6 @@ def __init__( self.data_system_client = attach_tq_client( self.config.tq_config, role=self.role, - lease_owner=self, ) self.actor_model = allocate_train_group(args=config, num_gpus=num_gpus, pg=pgs, runtime_env=runtime_env) self.actor_model.init_and_wait(config, role=self.role, with_ref=False) diff --git a/relax/components/advantages.py b/relax/components/advantages.py index 1919c2138..a01c7c2e7 100644 --- a/relax/components/advantages.py +++ b/relax/components/advantages.py @@ -42,7 +42,6 @@ def __init__( self.data_system_client = attach_tq_client( self.config.tq_config, role="advantages", - lease_owner=self, ) self.step = 0 diff --git a/relax/components/base.py b/relax/components/base.py index 4b8b68aed..8f18e6720 100644 --- a/relax/components/base.py +++ b/relax/components/base.py @@ -89,21 +89,18 @@ 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. - generation = getattr(self, "_tq_client_generation", None) - if getattr(self, "data_system_client", None) is None or generation is None: + if getattr(self, "data_system_client", None) is None: return try: from relax.utils.tq.lifecycle import detach_tq_client - detach_tq_client(generation) - self._tq_client_generation = None + detach_tq_client() 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 4f44f038b..a8a45f221 100644 --- a/relax/components/critic.py +++ b/relax/components/critic.py @@ -43,7 +43,6 @@ def __init__( self.data_system_client = attach_tq_client( self.config.tq_config, 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 952091ab1..c20de9602 100644 --- a/relax/components/rollout.py +++ b/relax/components/rollout.py @@ -336,7 +336,6 @@ def __init__( self.data_system_client = attach_tq_client( self.config.tq_config, 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 9cf732c2c..d9a85ff77 100644 --- a/relax/components/sft.py +++ b/relax/components/sft.py @@ -84,7 +84,6 @@ def __init__(self, healthy, pgs, num_gpus, config, role, runtime_env=None): # n self.data_system_client = attach_tq_client( self.config.tq_config, role=self.role, - lease_owner=self, ) self._dataset: Any | None = None diff --git a/relax/core/controller.py b/relax/core/controller.py index a3ff0ac7b..4139db3a1 100644 --- a/relax/core/controller.py +++ b/relax/core/controller.py @@ -361,7 +361,7 @@ def _initialize_data_system(self): 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'}") + logger.info("[dataplane] controller ownership=owner") def _confirm_mooncake_attach(self, init_result: TqInitResult, fallback_config) -> TqInitResult: """Validate Mooncake attach/setup on every alive Ray node. @@ -382,8 +382,6 @@ def _confirm_mooncake_attach(self, init_result: TqInitResult, fallback_config) - """ def _close_owned_attempt() -> None: - if not init_result.owns_controller: - return try: close_tq_owner(init_result.owner) finally: @@ -410,10 +408,8 @@ def _close_owned_attempt() -> None: 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: - # ``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 mode != "auto" or fallback_config is None: + # ``required`` must fail loudly after cleaning this job's owner. _close_owned_attempt() raise RuntimeError( f"Mooncake attach handshake reported {len(failures)} failure(s) (--tq-rdma-mode={mode}): {detail}" diff --git a/relax/distributed/ray/rollout.py b/relax/distributed/ray/rollout.py index d3319e80b..c065d295b 100644 --- a/relax/distributed/ray/rollout.py +++ b/relax/distributed/ray/rollout.py @@ -816,7 +816,6 @@ def __init__(self, args, pg, data_source=None): self.data_system_client = attach_tq_client( self.args.tq_config, role="rollout_worker", - lease_owner=self, ) logger.info(f"import {self.args.rollout_function_path} as generate_rollout function.") @@ -922,10 +921,8 @@ 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. - generation = getattr(self, "_tq_client_generation", None) - if generation is not None: - detach_tq_client(generation) - self._tq_client_generation = None + if getattr(self, "data_system_client", None) is not None: + detach_tq_client() self.data_system_client = None def _shutdown_all_engines(self, timeout: float = 15.0): diff --git a/relax/utils/tq/lifecycle.py b/relax/utils/tq/lifecycle.py index bf4050634..2c0e31967 100644 --- a/relax/utils/tq/lifecycle.py +++ b/relax/utils/tq/lifecycle.py @@ -25,7 +25,6 @@ import os import threading import time -import uuid from dataclasses import dataclass from typing import Any @@ -39,53 +38,33 @@ CONTROLLER_NAME = "TransferQueueController" CONTROLLER_NAMESPACE = "transfer_queue" -OWNER_TOKEN_FIELD = "relax_owner_token" 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: - """Result of an owner-aware TransferQueue initialization transaction.""" + """Result of an exclusive-owner TransferQueue initialization.""" config: Any - owner: Any | None + owner: Any 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 TqInitializationError(RuntimeError): + """Raised after a failed ``tq.init`` candidate is fully cleaned.""" + + 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.""" - - -class TqConfigurationMismatch(RuntimeError): - """Raised when an existing controller uses an incompatible job config.""" - - -class TqControllerInspectionError(RuntimeError): - """Raised when Ray cannot prove the named controller's current state.""" - - -class TqControllerMissingConfig(RuntimeError): - """Raised when a reachable controller is provably half-initialised.""" + """Raised when an owner or controller cannot be confirmed gone.""" class TqHandshakeIsolationError(RuntimeError): @@ -193,7 +172,7 @@ def reap_unusable_tq_controller(get_config_timeout: float = 10.0) -> bool: raise RuntimeError(f"Failed to inspect TransferQueueController ({safe_exception_kind(error)})") from None if conf is not None: - raise TqConfigurationMismatch( + raise RuntimeError( "A healthy TransferQueueController already exists. The initial RDMA release requires an exclusive, " "clean Ray cluster; stop the previous Relax job or remove its TQ state before retrying." ) @@ -203,30 +182,6 @@ def reap_unusable_tq_controller(get_config_timeout: float = 10.0) -> bool: return True -def _controller_exists() -> bool: - try: - ray.get_actor(CONTROLLER_NAME, namespace=CONTROLLER_NAMESPACE) - return True - except ValueError: - return False - except ray.exceptions.RayError as error: - raise RuntimeError(f"Failed to query TransferQueueController ({safe_exception_kind(error)})") from None - - -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 assert_mooncake_rdma_configured() -> None: """Fail unless the attached client configuration requests Mooncake/RDMA. @@ -265,11 +220,9 @@ def _get_stored_config(timeout: float = 10.0) -> Any: except ValueError: raise except ray.exceptions.RayError as error: - raise TqControllerInspectionError( - f"Failed to read TransferQueueController config ({safe_exception_kind(error)})" - ) from None + raise RuntimeError(f"Failed to read TransferQueueController config ({safe_exception_kind(error)})") from None if conf is None: - raise TqControllerMissingConfig("TransferQueueController returned no config after tq.init completed") + raise RuntimeError("TransferQueueController returned no config after tq.init completed") return conf @@ -386,7 +339,6 @@ def attach_tq_client( *, role: str, timeout: float | None = None, - lease_owner: Any | None = None, ) -> Any: """Attach a component process within a bounded deadline. @@ -395,30 +347,20 @@ def attach_tq_client( 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. + The initial release assumes one component/replica per Ray actor process. + Components therefore own the one process-global TQ client they attach and + unconditionally detach it during teardown. """ - 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() - - _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 + 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) + return tq.get_client() -def detach_tq_client(generation: int | None = None) -> None: +def detach_tq_client() -> None: """Detach this worker's TQ client (attach-only inverse of :func:`attach_tq_client`). @@ -426,26 +368,12 @@ def detach_tq_client(generation: int | None = None) -> 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). 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. + controller or globally stored data). The initial release assumes one + component/replica per Ray actor process, so there is no cross-instance + generation lease. Force-killed workers still fall back to the master-side + TTL. """ - 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 + _close_local_tq_client() def _alive_node_ids() -> list[str]: @@ -572,22 +500,17 @@ def _handshake(handshake_conf: Any, check_mooncake: bool, attach_timeout: float) return failures -def close_tq_and_unmount(*, is_owner: bool) -> None: +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. - 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. + This function is reachable only inside the dedicated owner actor. Workers + call :func:`detach_tq_client` instead because upstream ``tq.close()`` kills + the named controller even from an attached process. """ - if not is_owner: - logger.info("[dataplane] Detaching local TQ client; global controller is owned by another process.") - detach_tq_client() - return - store_client = None try: store_client = getattr(tq.get_client().storage_manager, "storage_client", None) @@ -608,117 +531,117 @@ def close_tq_and_unmount(*, is_owner: bool) -> None: class _TransferQueueOwner: """Process boundary for first-time TQ initialization and global cleanup.""" - def __init__(self) -> None: - self._owns_controller = False + def ready(self) -> None: + """Scheduling barrier used before ``tq.init`` enters fallback scope.""" + + def termination_probe(self) -> None: + """Never return; a terminal ref proves this actor process was + killed.""" + threading.Event().wait() - def initialize(self, conf: Any, owner_token: str) -> tuple[Any, bool]: + def initialize(self, conf: Any) -> Any: _prepare_mooncake_runtime(conf) - _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 + return _get_stored_config() def close(self) -> None: - close_tq_and_unmount(is_owner=self._owns_controller) + close_tq_and_unmount() - def detach(self) -> None: - detach_tq_client() - -def _stop_owner_actor(owner: Any) -> None: +def _stop_owner_actor(owner: Any, *, timeout: float = 10.0) -> None: + """Force-kill ``owner`` and prove its process reached a terminal state.""" + probe_ref = None + control_error = "" 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 ({safe_exception_kind(e)}).") - + # Queue behind any blocked initialize/close call. It can only become + # ready after the actor dies because the method never returns. + probe_ref = owner.termination_probe.remote() + except ray.exceptions.RayActorError: + return # Already terminal before cleanup reached it. + except Exception as error: + control_error = f"termination probe submission ({safe_exception_kind(error)})" -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. + try: + ray.kill(owner, no_restart=True) + except ray.exceptions.RayActorError: + pass + except Exception as error: + kind = safe_exception_kind(error) + control_error = f"{control_error}, " if control_error else "" + control_error += f"actor kill ({kind})" - 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. - """ + pending: list[Any] = [] + if probe_ref is not None: + try: + ready, pending = ray.wait([probe_ref], num_returns=1, timeout=timeout) + except Exception as error: + ready = [] + control_error = f"{control_error}, " if control_error else "" + control_error += f"termination wait ({safe_exception_kind(error)})" + for ref in ready: + try: + ray.get(ref) + except ray.exceptions.RayActorError: + pass + except Exception as error: + control_error = f"{control_error}, " if control_error else "" + control_error += f"termination probe result ({safe_exception_kind(error)})" + else: + control_error = f"{control_error}, " if control_error else "" + control_error += "termination probe returned normally" + + if control_error or pending or probe_ref is None: + detail = control_error or "termination could not be observed" + if pending: + detail = f"{detail}, " if detail else "" + detail += "termination probe remained pending" + raise TqCleanupTimeout(f"Failed to confirm TransferQueue owner cleanup: {detail}") from None + + +def _cleanup_failed_owner(owner: Any, *, timeout: float = 10.0) -> None: + """Stop a failed initializer and remove its exclusive-cluster state.""" try: ray.get(owner.close.remote(), timeout=timeout) except Exception as e: logger.warning(f"[dataplane] TQ owner cleanup RPC failed; killing owner actor ({safe_exception_kind(e)}).") - finally: - _stop_owner_actor(owner) - - try: - stored_conf = _get_stored_config(timeout=timeout) - except ValueError: - return - except TqControllerMissingConfig as error: - logger.warning( - "[dataplane] Failed initializer left a half-initialised TQ controller " - f"({safe_exception_kind(error)}); reaping it." - ) - kill_tq_controller_and_wait() - return - except TqControllerInspectionError: - # The owner process has already been stopped, but a transient GCS error - # cannot prove that any visible controller belongs to this attempt. - # Fail closed instead of risking another job's healthy controller. - raise - except Exception as error: - raise TqControllerInspectionError( - f"Unexpected failure while inspecting TransferQueueController ({safe_exception_kind(error)})" - ) from None + _stop_owner_actor(owner, timeout=timeout) - 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." - ) + # Only remove global state after the actor is confirmed terminal. A + # pre-kill existence query races with a late controller creation by the + # blocked native initializer. The exclusive-cluster precondition proves + # any controller visible here belongs to this attempt. + kill_tq_controller_and_wait(timeout=timeout) -def _start_owner(conf: Any, *, timeout: float) -> TqInitResult: - owner_token = uuid.uuid4().hex +def _start_owner(*, timeout: float) -> Any: + """Create and schedule the isolated owner before candidate fallback.""" try: owner = _TransferQueueOwner.remote() - except ray.exceptions.RayError as error: + ray.get(owner.ready.remote(), timeout=timeout) + except Exception as error: + if "owner" in locals(): + _stop_owner_actor(owner, timeout=timeout) raise RuntimeError(f"TransferQueue owner creation failed ({safe_exception_kind(error)})") from None + return owner + + +def _initialize_owner(owner: Any, conf: Any, *, timeout: float) -> TqInitResult: + """Initialize one candidate and fully clean its owner on failure.""" try: - stored_conf, owns_controller = ray.get(owner.initialize.remote(conf, owner_token), timeout=timeout) + stored_conf = ray.get(owner.initialize.remote(conf), timeout=timeout) except ray.exceptions.GetTimeoutError: - _cleanup_failed_owner(owner, owner_token) + _cleanup_failed_owner(owner) raise TqInitializationTimeout(f"tq.init did not finish within {timeout:.0f}s") from None except Exception as error: - _cleanup_failed_owner(owner, owner_token) - if isinstance(error, ray.exceptions.RayError): - raise RuntimeError(f"TransferQueue owner initialization failed ({safe_exception_kind(error)})") from None - raise - - if owns_controller: - return TqInitResult(config=stored_conf, owner=owner) - - # A controller appeared after the exclusive-cluster pre-check. This owner - # only attached to it, so detach locally and fail instead of silently sharing - # global state with a concurrent initializer. - try: - ray.get(owner.detach.remote(), timeout=10.0) - except ray.exceptions.RayError as error: - raise RuntimeError( - f"TransferQueue concurrent-initializer detach failed ({safe_exception_kind(error)})" + _cleanup_failed_owner(owner) + raise TqInitializationError( + f"TransferQueue owner initialization failed ({safe_exception_kind(error)})" ) from None - finally: - _stop_owner_actor(owner) - raise TqConfigurationMismatch( - "A concurrent TransferQueue initializer created a healthy controller. " - "Detached without modifying it; the initial RDMA release requires an exclusive Ray cluster." - ) + return TqInitResult(config=stored_conf, owner=owner) 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.""" + """Ask the exclusive owner process to close global TQ state.""" if owner is None: return close_error: Exception | None = None @@ -726,12 +649,10 @@ def close_tq_owner(owner: Any | None, *, timeout: float = 30.0) -> None: 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() + _stop_owner_actor(owner, timeout=timeout) + # Ensure the controller has left GCS even when owner.close() failed before + # reaching upstream tq.close(). This is a no-op when it is already gone. + kill_tq_controller_and_wait(timeout=timeout) if close_error is not None: raise RuntimeError(f"TransferQueue owner cleanup failed ({safe_exception_kind(close_error)})") from None @@ -743,21 +664,25 @@ def initialize_tq_with_fallback( fallback_conf: Any | None = None, timeout: float = DEFAULT_TQ_INIT_TIMEOUT_SECONDS, ) -> TqInitResult: - """Initialize TQ atomically with owner tracking and one safe auto fallback. + """Initialize TQ with one exclusive owner 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() - return _start_owner(attempt_conf, timeout=timeout) + def _attempt(owner: Any, attempt_conf: Any) -> TqInitResult: + return _initialize_owner(owner, attempt_conf, timeout=timeout) + # Keep the exclusive-cluster gate outside the auto fallback boundary. A + # healthy pre-existing controller is not an RDMA candidate failure and + # must never be converted into a SimpleStorage attach attempt. + reap_unusable_tq_controller() + owner = _start_owner(timeout=timeout) try: - return _attempt(conf) - except Exception as primary_error: - if isinstance(primary_error, TqConfigurationMismatch) or mode != "auto" or fallback_conf is None: + return _attempt(owner, conf) + except (TqInitializationTimeout, TqInitializationError) as primary_error: + if mode != "auto" or fallback_conf is None: raise reason = f"mooncake_init_failed:{type(primary_error).__name__}" @@ -766,7 +691,11 @@ def _attempt(attempt_conf: Any) -> TqInitResult: "cleaned partial state and retrying once with SimpleStorage." ) try: - result = _attempt(fallback_conf) + # Failed-owner cleanup must leave a clean cluster. Re-check before + # starting the fallback so unknown residual state fails closed. + reap_unusable_tq_controller() + fallback_owner = _start_owner(timeout=timeout) + result = _attempt(fallback_owner, fallback_conf) except Exception as fallback_error: raise RuntimeError( "TransferQueue SimpleStorage fallback initialization failed after " diff --git a/tests/core/test_controller_tq_backend.py b/tests/core/test_controller_tq_backend.py index 61f338d32..4ebbf9036 100644 --- a/tests/core/test_controller_tq_backend.py +++ b/tests/core/test_controller_tq_backend.py @@ -22,7 +22,6 @@ import pytest -from relax.utils.tq.lifecycle import TqConfigurationMismatch from tests.core.test_controller_s3_model_cleanup import controller from tests.utils.test_arguments_opd_teacher_colocate import ( arguments_module as _arguments_module_fixture, @@ -140,7 +139,7 @@ def test_healthy_existing_controller_aborts_before_legacy_init(self, monkeypatch monkeypatch.setattr( controller, "reap_unusable_tq_controller", - lambda: (_ for _ in ()).throw(TqConfigurationMismatch("exclusive cluster is not clean")), + lambda: (_ for _ in ()).throw(RuntimeError("exclusive cluster is not clean")), ) monkeypatch.setattr( controller.tq, @@ -148,7 +147,7 @@ def test_healthy_existing_controller_aborts_before_legacy_init(self, monkeypatch lambda **_kwargs: pytest.fail("existing controller must be rejected before tq.init"), ) - with pytest.raises(TqConfigurationMismatch, match="exclusive cluster"): + with pytest.raises(RuntimeError, match="exclusive cluster"): instance._initialize_data_system() assert instance._tq_owner is None @@ -391,10 +390,10 @@ def fake_initialize(conf, **_kwargs): monkeypatch.setattr(controller, "initialize_tq_with_fallback", fake_initialize) -def _confirm(config, *, owner="mooncake-owner", owns_controller=True, fallback_config="simple-conf"): +def _confirm(config, *, owner="mooncake-owner", fallback_config="simple-conf"): instance = controller.Controller.__new__(controller.Controller) instance.config = config - init_result = controller.TqInitResult(config="mooncake-conf", owner=owner if owns_controller else None) + init_result = controller.TqInitResult(config="mooncake-conf", owner=owner) return instance._confirm_mooncake_attach(init_result, fallback_config) @@ -443,14 +442,6 @@ def test_required_closes_owner_then_raises(self, monkeypatch): assert recorder.events == ["handshake", "close"] assert recorder.initialized == [] - def test_attached_session_never_tears_down_a_foreign_controller(self, monkeypatch): - """A job that only attached must not close state it does not own.""" - recorder = _AttachRecorder(monkeypatch, failures=["node-B: attach timed out"]) - with pytest.raises(RuntimeError, match="attach handshake reported"): - _confirm(_config(), owns_controller=False) - assert recorder.events == ["handshake"] - assert recorder.closed == [] - def test_unexpected_driver_exception_closes_owner_and_is_sanitized(self, monkeypatch): secret = "worker endpoint and traceback path must stay private" recorder = _AttachRecorder(monkeypatch, verify_error=RuntimeError(secret)) @@ -460,13 +451,6 @@ def test_unexpected_driver_exception_closes_owner_and_is_sanitized(self, monkeyp assert recorder.initialized == [] assert secret not in str(excinfo.value) - def test_unexpected_driver_exception_never_closes_foreign_owner(self, monkeypatch): - recorder = _AttachRecorder(monkeypatch, verify_error=RuntimeError("private detail")) - with pytest.raises(RuntimeError, match="orchestration failed"): - _confirm(_config(), owns_controller=False) - assert recorder.events == ["handshake"] - assert recorder.closed == [] - def test_unconfirmed_worker_isolation_closes_owner_and_aborts(self, monkeypatch): recorder = _AttachRecorder( monkeypatch, diff --git a/tests/utils/_tq_handshake_timeout_probe.py b/tests/utils/_tq_handshake_timeout_probe.py index b49c665ae..380dc09ba 100644 --- a/tests/utils/_tq_handshake_timeout_probe.py +++ b/tests/utils/_tq_handshake_timeout_probe.py @@ -144,6 +144,20 @@ def _clean_worker_state() -> tuple[int, float, bool, bool]: assert (successor_pid, successor_create_time) != timed_out_identity assert successor_mutated is False assert late_marker_exists is False + + # The dedicated initializer uses the same fail-closed termination + # barrier. This exercises the actual RayActorError transition rather + # than relying only on mocked ray.kill/ray.wait behavior. + owner = tq_lifecycle._TransferQueueOwner.remote() + ray.get(owner.ready.remote()) + tq_lifecycle._stop_owner_actor(owner, timeout=10.0) + try: + after_stop_ref = owner.ready.remote() + ray.get(after_stop_ref) + except ray.exceptions.RayActorError: + pass + else: + raise AssertionError("owner accepted work after terminal cleanup was confirmed") finally: ray.shutdown() diff --git a/tests/utils/test_tq_failure_paths.py b/tests/utils/test_tq_failure_paths.py index 6cb9ce4bb..f2977029e 100644 --- a/tests/utils/test_tq_failure_paths.py +++ b/tests/utils/test_tq_failure_paths.py @@ -24,6 +24,7 @@ import asyncio import importlib.util +import inspect import multiprocessing import os import queue @@ -283,7 +284,7 @@ def test_no_controller_is_a_noop(self, monkeypatch): def test_healthy_controller_fails_without_being_killed(self, monkeypatch): killed = self._fake_ray(monkeypatch, actor=MagicMock(), get_result={"backend": {}}) - with pytest.raises(tq_lifecycle.TqConfigurationMismatch, match="exclusive, clean Ray cluster"): + with pytest.raises(RuntimeError, match="exclusive, clean Ray cluster"): tq_lifecycle.reap_unusable_tq_controller() assert killed == [] @@ -325,7 +326,7 @@ def test_stored_config_ray_error_is_sanitized(self, monkeypatch): lambda *_args, **_kwargs: (_ for _ in ()).throw(tq_lifecycle.ray.exceptions.RayError(private_detail)), ) - with pytest.raises(tq_lifecycle.TqControllerInspectionError) as excinfo: + with pytest.raises(RuntimeError) as excinfo: tq_lifecycle._get_stored_config() assert private_detail not in str(excinfo.value) @@ -353,29 +354,22 @@ 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(is_owner=True) + 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(is_owner=True) + 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(is_owner=True) # must not raise + tq_lifecycle.close_tq_and_unmount() # 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() - # --------------------------------------------------------------------------- # Bounded attach (worker-side tq.init used to hang forever) @@ -727,53 +721,50 @@ 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) - - old_owner = SimpleNamespace() - new_owner = SimpleNamespace() - assert tq_lifecycle.attach_tq_client({}, role="old", lease_owner=old_owner) is client - assert tq_lifecycle.attach_tq_client({}, 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 generation: calls.append(generation)) + monkeypatch.setattr(tq_lifecycle, "detach_tq_client", lambda: calls.append(True)) component = Base() component.data_system_client = object() - component._tq_client_generation = 7 component.__del__() - assert calls == [7] + assert calls == [True] 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 generation: calls.append(generation)) + monkeypatch.setattr(tq_lifecycle, "detach_tq_client", lambda: calls.append(True)) component = Base() component.__del__() assert calls == [] +class TestExclusiveLifecycleSurface: + """Phase 4 must not grow the removed shared-cluster machinery back.""" + + def test_token_generation_and_foreign_owner_symbols_are_absent(self): + removed = ( + "OWNER_TOKEN_FIELD", + "TqConfigurationMismatch", + "TqControllerInspectionError", + "TqControllerMissingConfig", + "_TQ_CLIENT_GENERATION", + "_CURRENT_TQ_CLIENT_GENERATION", + ) + for name in removed: + assert not hasattr(tq_lifecycle, name) + + def test_attach_and_teardown_signatures_express_single_process_owner(self): + assert "lease_owner" not in inspect.signature(tq_lifecycle.attach_tq_client).parameters + assert list(inspect.signature(tq_lifecycle.detach_tq_client).parameters) == [] + assert list(inspect.signature(tq_lifecycle.close_tq_and_unmount).parameters) == [] + + # --------------------------------------------------------------------------- -# Owner-aware initialization transaction +# Exclusive-owner initialization transaction # --------------------------------------------------------------------------- @@ -783,98 +774,155 @@ def _conf(backend: str) -> dict: return {"controller": {}, "backend": {"storage_backend": backend}} @staticmethod - def _mooncake_conf(protocol: str, *, use_gdr: bool = False) -> dict: - return { - "controller": {}, - "backend": { - "storage_backend": "MooncakeStore", - "MooncakeStore": { - "protocol": protocol, - "master_server_address": "master.invalid:50051", - "hard_pin": True, - "use_gdr": use_gdr, - }, - }, - } - - @staticmethod - def _patch_transaction(monkeypatch, *, init_effects: list[object], reap_error: BaseException | None = None): - calls: dict[str, list] = {"reap": [], "attempts": []} + def _patch_transaction( + monkeypatch, + *, + init_effects: list[object], + start_effects: list[object] | None = None, + reap_effects: list[BaseException | None] | None = None, + ) -> list[str]: + events: list[str] = [] effects = iter(init_effects) + owners = iter(start_effects or [f"owner-{index}" for index in range(len(init_effects))]) + reaps = iter(reap_effects or []) def fake_reap(): - calls["reap"].append(True) - if reap_error is not None: - raise reap_error + events.append("reap") + effect = next(reaps, None) + if effect is not None: + raise effect monkeypatch.setattr(tq_lifecycle, "reap_unusable_tq_controller", fake_reap) - def fake_start(conf, *, timeout): - calls["attempts"].append(conf) + def fake_start(*, timeout): + events.append("start") + effect = next(owners) + if isinstance(effect, BaseException): + raise effect + return effect + + def fake_initialize(owner, conf, *, timeout): + backend = conf["backend"]["storage_backend"] + events.append(f"init:{backend}") effect = next(effects) if isinstance(effect, BaseException): raise effect - return tq_lifecycle.TqInitResult(config=conf, owner=effect) + return tq_lifecycle.TqInitResult(config=conf, owner=owner) monkeypatch.setattr(tq_lifecycle, "_start_owner", fake_start) - return calls + monkeypatch.setattr(tq_lifecycle, "_initialize_owner", fake_initialize) + return events def test_simple_path_also_runs_pre_init_reaper_and_becomes_owner(self, monkeypatch): conf = self._conf("SimpleStorage") - calls = self._patch_transaction(monkeypatch, init_effects=["owner"]) + events = self._patch_transaction(monkeypatch, init_effects=[None]) result = tq_lifecycle.initialize_tq_with_fallback(conf, mode="off") - assert result.owns_controller is True - assert len(calls["reap"]) == 1 + assert result.owner == "owner-0" + assert events == ["reap", "start", "init:SimpleStorage"] @pytest.mark.parametrize("mode", ["off", "auto", "required"]) def test_healthy_existing_controller_fails_exclusive_without_starting_owner(self, monkeypatch, mode): requested = self._conf("SimpleStorage") - mismatch = tq_lifecycle.TqConfigurationMismatch("exclusive cluster is not clean") - calls = self._patch_transaction(monkeypatch, init_effects=[], reap_error=mismatch) + mismatch = RuntimeError("exclusive cluster is not clean") + events = self._patch_transaction(monkeypatch, init_effects=[], reap_effects=[mismatch]) - with pytest.raises(tq_lifecycle.TqConfigurationMismatch, match="exclusive cluster"): + with pytest.raises(RuntimeError, match="exclusive cluster"): tq_lifecycle.initialize_tq_with_fallback( requested, mode=mode, fallback_conf=self._conf("SimpleStorage"), ) - assert calls["attempts"] == [] + assert events == ["reap"] def test_auto_cleans_failed_mooncake_then_retries_simple_once(self, monkeypatch): primary = self._conf("MooncakeStore") fallback = self._conf("SimpleStorage") - calls = self._patch_transaction( + events = self._patch_transaction( monkeypatch, - init_effects=[RuntimeError("master unavailable"), "fallback-owner"], + init_effects=[tq_lifecycle.TqInitializationError("master unavailable"), None], ) 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 + assert result.fallback_reason == "mooncake_init_failed:TqInitializationError" + assert events == [ + "reap", + "start", + "init:MooncakeStore", + "reap", + "start", + "init:SimpleStorage", + ] def test_required_cleans_failed_init_without_fallback(self, monkeypatch): primary = self._conf("MooncakeStore") fallback = self._conf("SimpleStorage") - calls = self._patch_transaction( + events = self._patch_transaction( monkeypatch, - init_effects=[RuntimeError("master unavailable")], + init_effects=[tq_lifecycle.TqInitializationError("master unavailable")], ) - with pytest.raises(RuntimeError, match="master unavailable"): + with pytest.raises(tq_lifecycle.TqInitializationError, match="master unavailable"): tq_lifecycle.initialize_tq_with_fallback(primary, mode="required", fallback_conf=fallback) - assert len(calls["attempts"]) == 1 + assert events == ["reap", "start", "init:MooncakeStore"] def test_timeout_auto_retries_only_after_isolated_owner_cleanup(self, monkeypatch): primary = self._conf("MooncakeStore") fallback = self._conf("SimpleStorage") - calls = self._patch_transaction( + events = self._patch_transaction( monkeypatch, - init_effects=[tq_lifecycle.TqInitializationTimeout("timed out"), "fallback-owner"], + init_effects=[tq_lifecycle.TqInitializationTimeout("timed out"), None], ) 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 + assert events[-3:] == ["reap", "start", "init:SimpleStorage"] + + def test_cleanup_failure_aborts_auto_before_second_gate(self, monkeypatch): + events = self._patch_transaction( + monkeypatch, + init_effects=[tq_lifecycle.TqCleanupTimeout("owner still running")], + ) + + with pytest.raises(tq_lifecycle.TqCleanupTimeout, match="still running"): + tq_lifecycle.initialize_tq_with_fallback( + self._conf("MooncakeStore"), + mode="auto", + fallback_conf=self._conf("SimpleStorage"), + ) + + assert events == ["reap", "start", "init:MooncakeStore"] + + def test_second_gate_failure_aborts_before_fallback_owner(self, monkeypatch): + gate_error = RuntimeError("residual controller state is unknown") + events = self._patch_transaction( + monkeypatch, + init_effects=[tq_lifecycle.TqInitializationError("master unavailable")], + reap_effects=[None, gate_error], + ) + + with pytest.raises(RuntimeError, match="fallback initialization failed"): + tq_lifecycle.initialize_tq_with_fallback( + self._conf("MooncakeStore"), + mode="auto", + fallback_conf=self._conf("SimpleStorage"), + ) + + assert events == ["reap", "start", "init:MooncakeStore", "reap"] + + def test_owner_creation_failure_is_not_a_candidate_fallback(self, monkeypatch): + events = self._patch_transaction( + monkeypatch, + init_effects=[], + start_effects=[RuntimeError("owner scheduling failed")], + ) + + with pytest.raises(RuntimeError, match="owner scheduling failed"): + tq_lifecycle.initialize_tq_with_fallback( + self._conf("MooncakeStore"), + mode="auto", + fallback_conf=self._conf("SimpleStorage"), + ) + + assert events == ["reap", "start"] class _RemoteMethod: @@ -887,106 +935,148 @@ def remote(self, *args, **kwargs): class _FakeOwner: def __init__(self): + self.ready = _RemoteMethod("ready-ref") self.initialize = _RemoteMethod("initialize-ref") self.close = _RemoteMethod("close-ref") - self.detach = _RemoteMethod("detach-ref") + self.termination_probe = _RemoteMethod("termination-ref") class TestOwnerProcessBoundary: - def test_start_timeout_cleans_the_isolated_owner_before_raising(self, monkeypatch): + def test_start_scheduling_timeout_stops_owner_without_entering_init(self, monkeypatch): owner = _FakeOwner() - cleaned: list[tuple[object, str]] = [] + stopped: list[object] = [] monkeypatch.setattr(tq_lifecycle._TransferQueueOwner, "remote", lambda: owner) def timed_out(ref, *, timeout): - assert ref == "initialize-ref" + assert ref == "ready-ref" raise tq_lifecycle.ray.exceptions.GetTimeoutError("test timeout") monkeypatch.setattr(tq_lifecycle.ray, "get", timed_out) + monkeypatch.setattr(tq_lifecycle, "_stop_owner_actor", lambda handle, **_kwargs: stopped.append(handle)) + + with pytest.raises(RuntimeError, match="owner creation failed"): + tq_lifecycle._start_owner(timeout=0.1) + assert stopped == [owner] + + def test_start_success_returns_scheduled_owner(self, monkeypatch): + owner = _FakeOwner() + monkeypatch.setattr(tq_lifecycle._TransferQueueOwner, "remote", lambda: owner) + monkeypatch.setattr(tq_lifecycle.ray, "get", lambda ref, timeout: None) + + assert tq_lifecycle._start_owner(timeout=1) is owner + + def test_initialize_success_returns_the_exclusive_owner(self, monkeypatch): + owner = _FakeOwner() + stored = TestInitializeTqWithFallback._conf("SimpleStorage") + monkeypatch.setattr(tq_lifecycle.ray, "get", lambda ref, timeout: stored) + + result = tq_lifecycle._initialize_owner(owner, stored, timeout=1) + + assert result.config is stored + assert result.owner is owner + + def test_initialize_timeout_cleans_owner_before_candidate_failure(self, monkeypatch): + owner = _FakeOwner() + cleaned: list[object] = [] monkeypatch.setattr( - tq_lifecycle, - "_cleanup_failed_owner", - lambda handle, token: cleaned.append((handle, token)), + tq_lifecycle.ray, + "get", + lambda _ref, timeout: (_ for _ in ()).throw(tq_lifecycle.ray.exceptions.GetTimeoutError("timeout")), ) + monkeypatch.setattr(tq_lifecycle, "_cleanup_failed_owner", lambda handle: cleaned.append(handle)) with pytest.raises(tq_lifecycle.TqInitializationTimeout): - tq_lifecycle._start_owner({"controller": {}}, timeout=0.1) - assert cleaned[0][0] is owner - assert cleaned[0][1] + tq_lifecycle._initialize_owner(owner, {}, timeout=0.1) - def test_concurrent_initializer_with_different_config_detaches_and_fails(self, monkeypatch): + assert cleaned == [owner] + + def test_initialize_error_is_sanitized_only_after_cleanup(self, monkeypatch): owner = _FakeOwner() - stopped: list[object] = [] - requested = TestInitializeTqWithFallback._mooncake_conf("rdma") - stored = TestInitializeTqWithFallback._conf("SimpleStorage") - monkeypatch.setattr(tq_lifecycle._TransferQueueOwner, "remote", lambda: owner) + cleaned: list[object] = [] + private_detail = "private endpoint and traceback path" monkeypatch.setattr( tq_lifecycle.ray, "get", - lambda ref, timeout: (stored, False) if ref == "initialize-ref" else None, + lambda _ref, timeout: (_ for _ in ()).throw(RuntimeError(private_detail)), ) - monkeypatch.setattr(tq_lifecycle, "_stop_owner_actor", lambda handle: stopped.append(handle)) + monkeypatch.setattr(tq_lifecycle, "_cleanup_failed_owner", lambda handle: cleaned.append(handle)) - with pytest.raises(tq_lifecycle.TqConfigurationMismatch, match="concurrent TransferQueue initializer"): - tq_lifecycle._start_owner(requested, timeout=1) - assert stopped == [owner] + with pytest.raises(tq_lifecycle.TqInitializationError) as excinfo: + tq_lifecycle._initialize_owner(owner, {}, timeout=0.1) + + assert cleaned == [owner] + assert private_detail not in str(excinfo.value) - @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): + def test_stop_owner_requires_terminal_probe_result(self, monkeypatch): 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)) + killed: list[tuple[object, bool]] = [] + monkeypatch.setattr(tq_lifecycle.ray, "kill", lambda handle, no_restart: killed.append((handle, no_restart))) + monkeypatch.setattr(tq_lifecycle.ray, "wait", lambda refs, **_kwargs: (list(refs), [])) monkeypatch.setattr( - tq_lifecycle, - "_get_stored_config", - lambda timeout: {"controller": {tq_lifecycle.OWNER_TOKEN_FIELD: stored_token}}, + tq_lifecycle.ray, + "get", + lambda _ref: (_ for _ in ()).throw(tq_lifecycle.ray.exceptions.RayActorError(error_msg="dead")), ) - 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 + tq_lifecycle._stop_owner_actor(owner, timeout=0.1) - def test_failed_owner_cleanup_does_not_kill_when_controller_inspection_fails(self, monkeypatch): + assert killed == [(owner, True)] + + def test_stop_owner_fails_closed_when_probe_remains_pending(self, monkeypatch): 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.ray, "kill", lambda *_args, **_kwargs: None) + monkeypatch.setattr(tq_lifecycle.ray, "wait", lambda refs, **_kwargs: ([], list(refs))) + + with pytest.raises(tq_lifecycle.TqCleanupTimeout, match="remained pending"): + tq_lifecycle._stop_owner_actor(owner, timeout=0.1) + + def test_stop_owner_rejects_a_probe_that_returns_normally(self, monkeypatch): + owner = _FakeOwner() + monkeypatch.setattr(tq_lifecycle.ray, "kill", lambda *_args, **_kwargs: None) + monkeypatch.setattr(tq_lifecycle.ray, "wait", lambda refs, **_kwargs: (list(refs), [])) + monkeypatch.setattr(tq_lifecycle.ray, "get", lambda _ref: None) + + with pytest.raises(tq_lifecycle.TqCleanupTimeout, match="returned normally"): + tq_lifecycle._stop_owner_actor(owner, timeout=0.1) + + def test_failed_owner_cleanup_waits_for_owner_then_reaps_controller(self, monkeypatch): + owner = _FakeOwner() + events: list[str] = [] + monkeypatch.setattr(tq_lifecycle.ray, "get", lambda _ref, timeout: events.append("close")) monkeypatch.setattr( tq_lifecycle, - "_get_stored_config", - lambda timeout: (_ for _ in ()).throw( - tq_lifecycle.TqControllerInspectionError("controller inspection failed (RayError)") - ), + "_stop_owner_actor", + lambda handle, **_kwargs: events.append("owner-terminal"), + ) + monkeypatch.setattr( + tq_lifecycle, + "kill_tq_controller_and_wait", + lambda **_kwargs: events.append("controller-gone"), ) - monkeypatch.setattr(tq_lifecycle, "kill_tq_controller_and_wait", lambda: killed.append(True)) - with pytest.raises(tq_lifecycle.TqControllerInspectionError): - tq_lifecycle._cleanup_failed_owner(owner, "ours") + tq_lifecycle._cleanup_failed_owner(owner) - assert stopped == [owner] - assert killed == [] + assert events == ["close", "owner-terminal", "controller-gone"] - def test_failed_owner_cleanup_reaps_proven_half_initialised_controller(self, monkeypatch): + def test_unconfirmed_owner_aborts_before_controller_cleanup(self, monkeypatch): owner = _FakeOwner() - killed: list[bool] = [] - monkeypatch.setattr(tq_lifecycle.ray, "get", lambda ref, timeout: None) - monkeypatch.setattr(tq_lifecycle, "_stop_owner_actor", lambda _handle: None) + controller_cleanup: list[bool] = [] + monkeypatch.setattr(tq_lifecycle.ray, "get", lambda _ref, timeout: None) + monkeypatch.setattr( + tq_lifecycle, + "_stop_owner_actor", + lambda *_args, **_kwargs: (_ for _ in ()).throw(tq_lifecycle.TqCleanupTimeout("owner pending")), + ) monkeypatch.setattr( tq_lifecycle, - "_get_stored_config", - lambda timeout: (_ for _ in ()).throw( - tq_lifecycle.TqControllerMissingConfig("controller returned no config") - ), + "kill_tq_controller_and_wait", + lambda **_kwargs: controller_cleanup.append(True), ) - monkeypatch.setattr(tq_lifecycle, "kill_tq_controller_and_wait", lambda: killed.append(True)) - tq_lifecycle._cleanup_failed_owner(owner, "ours") + with pytest.raises(tq_lifecycle.TqCleanupTimeout, match="owner pending"): + tq_lifecycle._cleanup_failed_owner(owner) - assert killed == [True] + assert controller_cleanup == [] def test_owner_close_failure_still_reaps_global_controller(self, monkeypatch): owner = _FakeOwner() @@ -997,9 +1087,8 @@ def test_owner_close_failure_still_reaps_global_controller(self, monkeypatch): "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)) + monkeypatch.setattr(tq_lifecycle, "_stop_owner_actor", lambda handle, **_kwargs: stopped.append(handle)) + monkeypatch.setattr(tq_lifecycle, "kill_tq_controller_and_wait", lambda **_kwargs: killed.append(True)) with pytest.raises(RuntimeError, match="owner cleanup failed"): tq_lifecycle.close_tq_owner(owner) From 766a1d8615f80dc74afb3bb1d7b3393f2d6dc599 Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Mon, 24 Aug 2026 03:06:39 +0800 Subject: [PATCH 24/30] refactor(tq): consolidate cross-node acceptance benchmark Remove redundant single-node and raw transport benchmarks along with the one-off fixture generator. Keep one C0/C1/C2 acceptance tool with selected-device wire proof, byte-exact NonTensorStack coverage, fail-closed teardown, and safe external fixture loading. --- .gitignore | 3 +- docs/draft/transfer_queue_rdma.md | 43 +- relax/utils/payload_digest.py | 157 --- relax/utils/tq/config.py | 15 +- scripts/benchmarks/cross_node_rdma_bench.py | 226 ---- scripts/benchmarks/make_multimodal_fixture.py | 265 ----- scripts/benchmarks/tq_cross_node_bench.py | 1036 +++++++---------- scripts/benchmarks/tq_rdma_bench.py | 414 ------- tests/utils/mm_payload_fixtures.py | 47 +- tests/utils/test_mm_payload_fixtures.py | 71 ++ tests/utils/test_tq_benchmark_guards.py | 204 +++- tests/utils/test_tq_dataplane_behavior.py | 11 +- tests/utils/test_tq_failure_paths.py | 3 +- tests/utils/tq/test_config.py | 5 + tests/utils/tq/test_payload_assertions.py | 20 +- 15 files changed, 736 insertions(+), 1784 deletions(-) delete mode 100644 relax/utils/payload_digest.py delete mode 100644 scripts/benchmarks/cross_node_rdma_bench.py delete mode 100644 scripts/benchmarks/make_multimodal_fixture.py delete mode 100644 scripts/benchmarks/tq_rdma_bench.py create mode 100644 tests/utils/test_mm_payload_fixtures.py diff --git a/.gitignore b/.gitignore index 6c6d95fda..22481d5e3 100644 --- a/.gitignore +++ b/.gitignore @@ -55,6 +55,5 @@ tensorboard_log .github/copilot-instructions.md env.sh -# Machine-local multimodal acceptance fixtures (generated by -# scripts/benchmarks/make_multimodal_fixture.py; hundreds of MB, never commit) +# Machine-local multimodal acceptance artifacts (hundreds of MB, never commit) tests/fixtures/ diff --git a/docs/draft/transfer_queue_rdma.md b/docs/draft/transfer_queue_rdma.md index 3a4ade999..596dc88e5 100644 --- a/docs/draft/transfer_queue_rdma.md +++ b/docs/draft/transfer_queue_rdma.md @@ -64,7 +64,7 @@ MooncakeStore/host-RDMA → SimpleStorage `(pending handshake)` 表示配置校验已通过但能力尚未确认；只有最后一行出现才说明每个存活节点都完成了真实 attach/setup，storage manager 是 MooncakeStore，且 Mooncake client 的配置请求为 `protocol=rdma`。 -注意这条断言的强度：它核对的是 client 的**配置意图和 setup 结果**，不是 negotiated transport，也不能单独证明数据包没有经过 TCP。线路级证明（IB counter 增长、非 RDMA 数据接口 counter 不增长）由 benchmark 的 `--require-wire-proof` 提供，见下方验收章节。 +注意这条断言的强度：它核对的是 client 的**配置意图和 setup 结果**，不是 negotiated transport，也不能单独证明数据包没有经过 TCP。线路级证明由 benchmark 的强制 counter gate 提供：RDMA 档要求 IB receive counter 增长且大于 TCP counter，SimpleStorage/TCP 档要求 TCP receive counter 增长且不小于 IB counter。 发生回退时会看到： @@ -156,48 +156,33 @@ Mock/本机测试和真实双节点 RDMA 测试必须分别报告，前者不能 | 本机 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 | +| 真实双节点 | 同一 driver/consumer pair 下各 backend 的原生数据布局；SimpleStorage、Mooncake/TCP、Mooncake/RDMA；synthetic、production-shaped multimodal 两种 profile；256/1024/2048/4096 MiB；每档 warmup + 至少 5 轮 | 每次 get 的 dtype、shape 与 raw-byte SHA-256 全部 PASS；强制 counter gate 证明对应线路；逐协议、逐轮 CSV 留档并报告均值、median、stddev | | 真实回退（`auto` + 某节点 RDMA 不可用） | 静态探测删除后，`auto` 会真的创建 Mooncake owner 与 named controller，再在 handshake 阶段失败并拆除，因此这条路径必须实测 | 逐条确认：① 日志显示 Mooncake owner `tq.init` **成功**、随后 handshake 阶段失败（否则实际只测到 owner 初始化失败，覆盖不到拆除）；② 最终存在且仅存在一个预期的 SimpleStorage `TransferQueueController`；③ 该 controller 的 stored config 确认为 SimpleStorage；④ SimpleStorage 的 put/get 正常工作；⑤ 旧 Mooncake owner actor 与其 client 均已消失；⑥ Mooncake segment 已从 master 卸载——若测的是强杀/超时路径，需等过 `client_ttl`（默认 30 s）再检查 | -真实多模态 fixture 生成（需要本地数据集 parquet 与 Qwen-VL 模型目录；产物写入 `tests/fixtures/`，已 gitignore，不入库）： +真实模型/数据 fixture 及其生成流程属于外部 PR 验收附件，不再由生产仓库维护。可选的本机 dataplane 测试仍可通过 `RELAX_MM_FIXTURE` 指向已验证的外部 fixture；唯一保留的跨节点 benchmark 只接受显式的 `synthetic` 和 `multimodal` profile，不会在缺失 fixture 时替换 profile。 -```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 集群需由部署侧预先准备）： +双节点验收命令（master 与 Ray 集群需由部署侧预先准备；每个 protocol 必须启动一个全新的 Python 进程并使用独立 CSV）： ```bash PYTHONPATH=. python -u scripts/benchmarks/tq_cross_node_bench.py \ + --protocol rdma \ --master :50051 \ - --nodeb-ip \ + --consumer-node-id \ --device \ - --payload-profiles synthetic multimodal real-multimodal \ + --tcp-device \ + --payload-profiles synthetic multimodal \ --payload-mib 256 1024 2048 4096 \ - --repeats 5 --require-wire-proof \ - --csv tq_cross_node_acceptance.csv + --repeats 5 \ + --csv c2-rdma.csv ``` -`real-multimodal` profile 按目标档位循环平铺 fixture 样本，`multimodal_train_inputs` 列以 NonTensorStack 走存储层非张量路径（SimpleStorage pickle / Mooncake msgpack），字节校验用行多重集指纹（采样器可重排行序；张量行按 dtype+字节比较以兼容后端间标量行 `()` 与 `[1]` 的表示差异，dict 叶子仍全形状校验）。 +用同样方式分别以 `--protocol simple` 和 `--protocol tcp` 运行 C0/C1，但这两档必须删除 `--device`；`--tcp-device` 三档都必须显式指定，因为它选择的是 consumer 节点上的验收 counter。C2 只汇总 `--device` 指定 HCA 的所有端口，不能用其他 HCA 的后台流量充当证据。C1/C2 除要求目标线路流量不低于非目标线路外，还分别要求 TCP/RDMA 接收量至少达到 raw payload 的 20%/80%；C0 因 SimpleStorage unit 可能被放在 consumer 本地，不使用 payload-volume 门槛，只确认观测到 TCP 且没有被 RDMA 流量主导。 -若验收环境没有两个 RDMA 节点，交付结论必须写成“真机验收未执行”，不能用 mock 通过推导真机已经通过。 - -### 参考吞吐区间 +C1 仅是 benchmark 对照，脚本会在 producer 和 consumer 进程内固定启用 Mooncake TCP connection pool，不形成生产配置面。C0 的 SimpleStorage units 使用上游原生 placement，因此三档共享的是固定 driver/consumer pair，而不是完全相同的 storage topology。每轮 byte-exact 或 wire-proof 失败都会立即终止，失败轮仍会先写入并 flush 到 CSV。 -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 注册开销 +若验收环境没有两个 RDMA 节点，交付结论必须写成“真机验收未执行”，不能用 mock 通过推导真机已经通过。 -逐档明细、逐轮分布与原始 CSV 属于交付验收材料，随验收报告存档，不在本文档维护。 +完整矩阵结果、依赖版本、逐轮分布与原始 CSV 属于对应 commit 的外部 PR 验收材料，不在本文档维护易过期的历史性能数字。 ## 排障表 diff --git a/relax/utils/payload_digest.py b/relax/utils/payload_digest.py deleted file mode 100644 index efac03fac..000000000 --- a/relax/utils/payload_digest.py +++ /dev/null @@ -1,157 +0,0 @@ -# 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 -import struct -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") - elif isinstance(value, float): - # repr(float("nan")) discards the payload bits. Pack the Python double - # directly so distinct NaNs and signed zero remain byte-distinguishable. - raw = struct.pack("!d", value) - else: # bool / int / 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 _dict_child_path(prefix: str, key: Any) -> str: - """Render a string dict key without colliding with nested/list paths.""" - if not isinstance(key, str): - raise TypeError(f"Unsupported payload dict key at {prefix}: {type(key).__name__}") - if key.isidentifier(): - return f"{prefix}.{key}" - return f"{prefix}[{key!r}]" - - -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 payload: - if not isinstance(key, str): - raise TypeError(f"Unsupported payload dict key at {prefix}: {type(key).__name__}") - for key in sorted(payload): - digests.update(leaf_digests(payload[key], _dict_child_path(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/relax/utils/tq/config.py b/relax/utils/tq/config.py index c335d813b..864f414ed 100644 --- a/relax/utils/tq/config.py +++ b/relax/utils/tq/config.py @@ -204,10 +204,9 @@ def build_mooncake_config( Parameters ---------- master_address - External master server address, already validated by - :func:`resolve_mooncake_master_address`. It is passed in rather than - re-read from the environment so the value that was checked is the value - that gets used. + External master server address. It is validated here as well as by + :func:`resolve_mooncake_master_address` so direct benchmark callers + cannot bypass the format contract. device Explicit RDMA device name; empty lets Mooncake select one natively. protocol @@ -219,6 +218,14 @@ def build_mooncake_config( :func:`resolve_global_segment_size`). Benchmarks may pass a larger value (e.g. 8 GiB) to avoid staging-buffer pressure. """ + if not isinstance(master_address, str): + raise ValueError("master_address must be a host:port string") + try: + _split_host_port(master_address) + except ValueError as error: + raise ValueError(f"master_address is not a usable endpoint: {error}") from None + master_address = master_address.strip() + if global_segment_size is None: segment_size = resolve_global_segment_size() elif isinstance(global_segment_size, bool) or not isinstance(global_segment_size, int) or global_segment_size <= 0: diff --git a/scripts/benchmarks/cross_node_rdma_bench.py b/scripts/benchmarks/cross_node_rdma_bench.py deleted file mode 100644 index 1dfcf9755..000000000 --- a/scripts/benchmarks/cross_node_rdma_bench.py +++ /dev/null @@ -1,226 +0,0 @@ -#!/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/make_multimodal_fixture.py b/scripts/benchmarks/make_multimodal_fixture.py deleted file mode 100644 index 37e901789..000000000 --- a/scripts/benchmarks/make_multimodal_fixture.py +++ /dev/null @@ -1,265 +0,0 @@ -#!/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 4bcec16cf..b23cceac2 100644 --- a/scripts/benchmarks/tq_cross_node_bench.py +++ b/scripts/benchmarks/tq_cross_node_bench.py @@ -1,47 +1,13 @@ #!/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 - -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. - -Usage (node A driver; node B already in the Ray cluster): - - PYTHONPATH= python -u scripts/benchmarks/tq_cross_node_bench.py \\ - --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). -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 +"""Cross-node TransferQueue C0/C1/C2 acceptance benchmark. + +Invoke once per fresh process: ``simple`` is C0, benchmark-only Mooncake/TCP is +C1, and Mooncake/host-RDMA is C2. The driver produces while one persistent +consumer is hard-pinned to another Ray node. Dtype/shape/raw-byte digests and +receive-counter wire proof are mandatory; each measured round is flushed to CSV +before either gate can fail. """ from __future__ import annotations @@ -49,363 +15,275 @@ import argparse import csv import hashlib +import os import statistics +import struct import time +from collections.abc import Mapping +from pathlib import Path 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 = "" +PROTOCOL_LABELS = { + "simple": "C0 SimpleStorage", + "tcp": "C1 Mooncake/TCP", + "rdma": "C2 Mooncake/RDMA", +} +CSV_COLUMNS = "protocol profile payload_mib actual_mib run byte_exact wire_proven put_ms get_ms put_gbs get_gbs ib_mb tcp_mb".split() +_TCP_WIRE_MIN_PAYLOAD_RATIO = 0.20 +_RDMA_WIRE_MIN_PAYLOAD_RATIO = 0.80 -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( - "--payload-profiles", - nargs="+", - default=["synthetic", "multimodal"], - 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", - 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).", +def _is_safe_device_name(value: str) -> bool: + return ( + bool(value) + and value not in {".", ".."} + and "/" not in value + and "\\" not in value + and all(character.isprintable() and not character.isspace() for character in value) ) - 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() - -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. - """ +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Cross-node TQ byte-exact and wire-proof benchmark") + parser.add_argument("--protocol", required=True, choices=sorted(PROTOCOL_LABELS)) + parser.add_argument("--consumer-node-id", required=True, help="Alive Ray NodeID distinct from the driver node") + parser.add_argument("--master", default=os.environ.get("MC_MASTER_ADDRESS", ""), help="Mooncake master host:port") + parser.add_argument("--device", default="", help="RDMA device required by the C2 wire-proof counter") + parser.add_argument("--tcp-device", required=True, help="Network interface used for the TCP receive counter") + parser.add_argument( + "--payload-profiles", nargs="+", default=["synthetic", "multimodal"], choices=["synthetic", "multimodal"] + ) + parser.add_argument("--payload-mib", nargs="+", type=int, default=[256, 1024, 2048, 4096]) + parser.add_argument("--num-samples", type=int, default=256) + parser.add_argument("--repeats", type=int, default=5) + parser.add_argument("--segment-gib", type=int, default=16) + parser.add_argument("--csv", required=True, help="Per-round CSV output path") + args = parser.parse_args() + for name in ("num_samples", "repeats", "segment_gib"): + if getattr(args, name) <= 0: + parser.error(f"--{name.replace('_', '-')} must be positive") + if not args.payload_mib or any(size <= 0 for size in args.payload_mib): + parser.error("--payload-mib values must be positive") + if args.protocol != "simple" and not args.master: + parser.error("--master or MC_MASTER_ADDRESS is required for Mooncake protocols") + if not _is_safe_device_name(args.tcp_device): + parser.error("--tcp-device must be a single printable interface name without whitespace or path separators") + if args.protocol == "rdma": + if not _is_safe_device_name(args.device): + parser.error("--device is required for RDMA and must be a single printable device name") + elif args.device: + parser.error("--device is valid only with --protocol rdma") + return args + + +def _columns_for_budget(total_bytes: int, num_samples: int, dtype: Any) -> int: 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]) + return max(1, total_bytes // (num_samples * torch.tensor([], dtype=dtype).element_size())) -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. - """ +def make_synthetic_payload(num_samples: int, total_mib: int): + """Deterministic mixed-dtype payload with a fixed schema.""" 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) + generator = torch.Generator().manual_seed(20260824) + field_budget = total_mib * 1024**2 // 3 + fp32_cols = _columns_for_budget(field_budget, num_samples, torch.float32) + bf16_cols = _columns_for_budget(field_budget, num_samples, torch.bfloat16) + int64_cols = _columns_for_budget(field_budget, num_samples, torch.int64) 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, + "activations": torch.randn(num_samples, fp32_cols, dtype=torch.float32, generator=generator), + "logits": torch.randn(num_samples, bf16_cols, dtype=torch.bfloat16, generator=generator), + "tokens": torch.randint(0, 151_000, (num_samples, int64_cols), dtype=torch.int64, generator=generator), }, batch_size=[num_samples], ) -_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. - """ +def make_multimodal_payload(num_samples: int, total_mib: int): + """Deterministic production-shaped non-tensor vision-language payload.""" 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 payload bytes across all fields (tensor, NestedTensor, or - NonTensorStack columns).""" - from relax.utils.payload_digest import total_leaf_bytes + seq_len = min(4096, max(128, target_bytes // max(1, num_samples * 128 * 1024))) + fixed_bytes = num_samples * (seq_len * 8 + 8 + 4) + pixel_budget = max(num_samples * 1536 * 4, target_bytes - fixed_bytes) + base_patches = max(1, pixel_budget // (num_samples * 1536 * 4)) + generator = torch.Generator().manual_seed(20260824) + multimodal_rows = [] + for index in range(num_samples): + variation = (index % 3) - 1 + requested_patches = max(1, base_patches + variation * max(1, base_patches // 20)) + height = max(1, int(requested_patches**0.5)) + width = max(1, (requested_patches + height - 1) // height) + patches = height * width + pixel_values = torch.empty((patches, 1536), dtype=torch.float32) + pixel_values.uniform_(-1.0, 1.0, generator=generator) + multimodal_rows.append( + { + "pixel_values": pixel_values, + "image_grid_thw": torch.tensor([[1, height, width]], dtype=torch.int64), + } + ) + token_row = list(range(seq_len)) + payload = dict_to_tensordict( + { + "tokens": [token_row for _ in range(num_samples)], + "sample_id": list(range(num_samples)), + "rewards": [float(index) / max(1, num_samples - 1) for index in range(num_samples)], + "multimodal_train_inputs": multimodal_rows, + }, + batch_size=num_samples, + ) + if type(payload.get("multimodal_train_inputs")).__name__ != "NonTensorStack": + raise RuntimeError("multimodal benchmark payload did not produce a NonTensorStack column") + return payload - return sum(total_leaf_bytes(payload[k]) for k in payload.keys()) +def _unwrap_non_tensor(value: Any) -> Any: + if type(value).__name__ == "NonTensorStack": + return value.tolist() + if type(value).__name__ == "NonTensorData": + return value.data + return value -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. - """ +def _column_rows(value: Any) -> list[Any]: 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 + value = _unwrap_non_tensor(value) + if isinstance(value, torch.Tensor): + if value.is_nested: + return list(value.unbind()) + if value.ndim == 0: + return [value] + return list(value.unbind()) + if isinstance(value, (list, tuple)): + return [_unwrap_non_tensor(row) for row in value] + raise TypeError(f"unsupported benchmark payload column: {type(value).__name__}") -def _column_rows(column) -> list: - """Rows of a TensorDict column: jagged NestedTensor, NonTensorStack, dense - tensor, or plain list.""" +def _leaf_digest(value: Any) -> Any: + import numpy as np 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 + value = _unwrap_non_tensor(value) + if isinstance(value, torch.Tensor): + if value.is_nested: + return ("nested", tuple(_leaf_digest(row) for row in value.unbind())) + contiguous = value.detach().cpu().contiguous() + flat = contiguous.reshape(-1) + raw = flat.view(torch.uint8).numpy().tobytes() if flat.numel() else b"" + return (str(contiguous.dtype), tuple(contiguous.shape), hashlib.sha256(raw).hexdigest()) + if isinstance(value, np.ndarray): + contiguous = np.ascontiguousarray(value) + return ( + f"np.{contiguous.dtype}", + tuple(contiguous.shape), + hashlib.sha256(contiguous.tobytes()).hexdigest(), + ) + if isinstance(value, Mapping): + if any(not isinstance(key, str) for key in value): + raise TypeError("benchmark payload dictionaries require string keys") + return ("dict", tuple((key, _leaf_digest(value[key])) for key in sorted(value))) + if isinstance(value, (list, tuple)): + return (type(value).__name__, tuple(_leaf_digest(item) for item in value)) + if isinstance(value, bytes): + raw = value + elif isinstance(value, str): + raw = value.encode("utf-8") + elif isinstance(value, float): + raw = struct.pack("!d", value) + elif isinstance(value, (bool, int)) or value is None: + raw = repr(value).encode("utf-8") + else: + raise TypeError(f"unsupported benchmark payload leaf: {type(value).__name__}") + return (f"py.{type(value).__name__}", (), hashlib.sha256(raw).hexdigest()) - 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]] = {} +def field_byte_digests(payload, fields: list[str]) -> dict[str, tuple[Any, ...]]: + """Ordered row digests preserving dtype, shape and every raw byte.""" + digests: dict[str, tuple[Any, ...]] = {} 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 + digests[field] = tuple(_leaf_digest(row) for row in _column_rows(payload.get(field))) + return digests -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) - raise TimeoutError(f"Ray actor {name!r} is still registered after {timeout:.1f}s") +def _value_bytes(value: Any) -> int: + import numpy as np + import torch + value = _unwrap_non_tensor(value) + if isinstance(value, torch.Tensor): + if value.is_nested: + return sum(_value_bytes(row) for row in value.unbind()) + return value.numel() * value.element_size() + if isinstance(value, np.ndarray): + return value.nbytes + if isinstance(value, Mapping): + return sum(_value_bytes(item) for item in value.values()) + if isinstance(value, (list, tuple)): + return sum(_value_bytes(item) for item in value) + if isinstance(value, bytes): + return len(value) + if isinstance(value, str): + return len(value.encode("utf-8")) + if isinstance(value, float): + return 8 + if isinstance(value, (bool, int)) or value is None: + return len(repr(value).encode("utf-8")) + raise TypeError(f"unsupported benchmark payload leaf: {type(value).__name__}") -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 +def payload_bytes(payload) -> int: + return sum(_value_bytes(row) for field in payload.keys() for row in _column_rows(payload.get(field))) - 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) +def read_counters(tcp_device: str, rdma_device: str = "", sysfs_root: Path = Path("/sys/class")) -> dict[str, int]: + """Read receive-byte counters on a selected HCA and TCP interface.""" + counters: dict[str, int] = {} + ib_root = sysfs_root / "infiniband" + if ib_root.is_dir(): + pattern = ( + f"{rdma_device}/ports/*/counters/port_rcv_data" if rdma_device else "*/ports/*/counters/port_rcv_data" + ) + for counter_path in sorted(ib_root.glob(pattern)): + try: + counters[f"ib:{counter_path.parents[3].name}:{counter_path.parents[1].name}"] = ( + int(counter_path.read_text().strip()) * 4 + ) + except (OSError, ValueError): + continue + tcp_path = sysfs_root / "net" / tcp_device / "statistics" / "rx_bytes" + try: + counters[f"tcp:{tcp_device}"] = int(tcp_path.read_text().strip()) + except (OSError, ValueError): + pass + return counters - wait_actor_gone() +def _counter_delta(before: dict[str, int], after: dict[str, int], prefix: str) -> int: + return sum(after[key] - before.get(key, after[key]) for key in after if key.startswith(prefix)) -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 wire_is_proven(protocol: str, ib_bytes: int, tcp_bytes: int, payload_nbytes: int) -> bool: + if protocol == "rdma": + return ib_bytes >= _RDMA_WIRE_MIN_PAYLOAD_RATIO * payload_nbytes and ib_bytes >= tcp_bytes + if protocol == "tcp": + return tcp_bytes >= _TCP_WIRE_MIN_PAYLOAD_RATIO * payload_nbytes and tcp_bytes >= ib_bytes + # SimpleStorage may place some units on the consumer node, so C0 cannot + # require a payload-volume ratio. It still must not look like RDMA. + return tcp_bytes > 0 and tcp_bytes >= ib_bytes 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 @@ -416,11 +294,9 @@ def build_conf(protocol: str, master: str, device: str, segment_gib: int): ) 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() - # ``protocol="tcp"`` is the C1 baseline; production only builds "rdma". backend = build_mooncake_config( master_address=master, device=device, @@ -436,338 +312,210 @@ def build_conf(protocol: str, master: str, device: str, segment_gib: int): ) -# ---- 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. +def require_clean_cluster() -> None: + try: + ray.get_actor("TransferQueueController", namespace="transfer_queue") + except ValueError: + return + raise RuntimeError("TransferQueueController already exists; benchmark requires a clean exclusive Ray cluster") - 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 +def close_tq_unmount_and_wait() -> None: + 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() + from relax.utils.tq.lifecycle import kill_tq_controller_and_wait - def alive(self) -> bool: - return True + store_client = None + try: + store_client = getattr(tq.get_client().storage_manager, "storage_client", None) + except (AssertionError, AttributeError): + pass + try: + tq.close() + finally: + try: + if store_client is not None and hasattr(store_client, "close"): + store_client.close() + finally: + kill_tq_controller_and_wait() - 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, expected_digests, order_insensitive: bool = False): - """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. - - ``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() +@ray.remote(num_cpus=0.001, max_restarts=0) +class TQConsumer: + """Persistent, hard-pinned consumer mirroring a Relax component actor.""" + + def __init__(self, conf: Any, tcp_device: str, rdma_device: str, protocol: str): + counters = read_counters(tcp_device, rdma_device if protocol == "rdma" else "") + if f"tcp:{tcp_device}" not in counters: + raise RuntimeError("TCP receive counter is unavailable on the selected consumer node") + if protocol == "rdma" and not any(key.startswith("ib:") for key in counters): + raise RuntimeError("RDMA receive counter is unavailable on the selected consumer node") + if protocol == "tcp": + os.environ["MC_TCP_ENABLE_CONNECTION_POOL"] = "1" + from relax.utils.tq.lifecycle import attach_tq_client + + # This runs inside the pinned Ray worker. The bounded helper applies + # Mooncake correctness guards in this process before client creation. + self.client = attach_tq_client(conf, role="benchmark-consumer") + self.tcp_device = tcp_device + self.rdma_device = rdma_device if protocol == "rdma" else "" + + def describe(self) -> tuple[str, str]: + manager = self.client.storage_manager + storage_client = getattr(manager, "storage_client", None) + return type(manager).__name__, getattr(storage_client, "protocol", "") + + def fetch(self, fields: list[str], batch_size: int, partition: str, expected: dict) -> dict[str, Any]: + before = read_counters(self.tcp_device, self.rdma_device) + started = time.perf_counter() meta = self.client.get_meta( - data_fields=list(fields), + data_fields=fields, batch_size=batch_size, partition_id=partition, mode="fetch", - task_name="xfer", + task_name="tq-benchmark", ) - 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() - # Digesting occurs outside the timed interval. A mismatch is fatal: - # throughput from a corrupt or truncated transfer is never reported. - 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}") - 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, True - - -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 + received = self.client.get_data(meta) + get_ms = (time.perf_counter() - started) * 1000 + after = read_counters(self.tcp_device, self.rdma_device) + actual = field_byte_digests(received, fields) + if actual != expected: + mismatches = [field for field in fields if actual.get(field) != expected.get(field)] + raise AssertionError(f"byte-exact mismatch after TQ get: fields={mismatches}") + return { + "get_ms": get_ms, + "ib_bytes": _counter_delta(before, after, "ib:"), + "tcp_bytes": _counter_delta(before, after, "tcp:"), + } + def shutdown(self) -> None: + from relax.utils.tq.lifecycle import detach_tq_client -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 + detach_tq_client() - 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")) - 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][profile][total_mib][num_fields] = { - # "put_mean_gbs","get_mean_gbs","get_med_gbs","get_std_gbs","wire", "per_run":[...]} - results: dict[str, 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, +def _validate_attached_backend(protocol: str, manager: str, attached_protocol: str) -> None: + if protocol == "simple" and manager != "AsyncSimpleStorageManager": + raise RuntimeError(f"C0 attached unexpected storage manager {manager}") + if protocol != "simple" and (manager != "MooncakeStorageManager" or attached_protocol != protocol): + raise RuntimeError( + f"{PROTOCOL_LABELS[protocol]} attached manager={manager} protocol={attached_protocol or '-'}" ) - 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 profile in args.payload_profiles: - results[protocol].setdefault(profile, {}) - for total_mib in args.payload_mib: - results[protocol][profile].setdefault(total_mib, {}) - # 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, fixture_path) - nf = len(list(payload.keys())) +def _teardown_benchmark(consumer: Any, owner_attempted: bool) -> None: + """Unmount the consumer, clean owned global state, then stop local Ray.""" + try: + if consumer is not None: + try: + ray.get(consumer.shutdown.remote(), timeout=10) + finally: + ray.kill(consumer, no_restart=True) + finally: + try: + if owner_attempted: + close_tq_unmount_and_wait() + finally: + ray.shutdown() + + +def main() -> None: + args = parse_args() + if args.protocol == "tcp": + os.environ["MC_TCP_ENABLE_CONNECTION_POOL"] = "1" + import transfer_queue as tq + from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy + + with open(args.csv, "w", newline="") as csv_handle: + writer = csv.DictWriter(csv_handle, fieldnames=CSV_COLUMNS) + writer.writeheader() + csv_handle.flush() + ray.init(address="auto", ignore_reinit_error=True, logging_level="ERROR") + consumer = None + measured_runs = 0 + owner_attempted = False + try: + driver_node_id = ray.get_runtime_context().get_node_id() + alive_ids = {node["NodeID"] for node in ray.nodes() if node.get("Alive")} + if args.consumer_node_id not in alive_ids: + raise RuntimeError("--consumer-node-id is not an alive Ray node") + if args.consumer_node_id == driver_node_id: + raise RuntimeError("producer and consumer must run on different Ray nodes") + require_clean_cluster() + + conf = build_conf(args.protocol, args.master, args.device, args.segment_gib) + owner_attempted = True + tq.init(conf=conf) + producer = tq.get_client() + strategy = NodeAffinitySchedulingStrategy(node_id=args.consumer_node_id, soft=False) + consumer = TQConsumer.options(scheduling_strategy=strategy).remote( + conf, args.tcp_device, args.device, args.protocol + ) + manager, attached_protocol = ray.get(consumer.describe.remote()) + _validate_attached_backend(args.protocol, manager, attached_protocol) + print(f"[setup] {PROTOCOL_LABELS[args.protocol]} consumer=remote-node", flush=True) + + for profile in args.payload_profiles: + for requested_mib in args.payload_mib: + payload = ( + make_multimodal_payload(args.num_samples, requested_mib) + if profile == "multimodal" + else make_synthetic_payload(args.num_samples, requested_mib) + ) fields = sorted(payload.keys()) - if order_insensitive: - expected_digests = field_multiset_digests(payload, fields) - else: - expected_digests = field_byte_digests(payload, fields) + expected = 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. - 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, batch_rows, part, expected_digests, order_insensitive) - ) - 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, - ) + put_rates: list[float] = [] + get_rates: list[float] = [] + for run in range(args.repeats + 1): + partition = f"bench-{profile}-{requested_mib}-{run}" + started = time.perf_counter() + producer.put(payload, partition_id=partition) + put_ms = (time.perf_counter() - started) * 1000 + fetched = ray.get(consumer.fetch.remote(fields, payload.batch_size[0], partition, expected)) + producer.clear_partition(partition) + if run == 0: 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( + proven = wire_is_proven(args.protocol, fetched["ib_bytes"], fetched["tcp_bytes"], nbytes) + put_gbs = nbytes / put_ms / 1e6 + get_gbs = nbytes / fetched["get_ms"] / 1e6 + put_rates.append(put_gbs) + get_rates.append(get_gbs) + writer.writerow( { - "protocol": protocol, + "protocol": args.protocol, "profile": profile, - "payload_mib": total_mib, + "payload_mib": requested_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, + "run": run, + "byte_exact": True, + "wire_proven": 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), + "get_ms": round(fetched["get_ms"], 2), + "put_gbs": round(put_gbs, 3), + "get_gbs": round(get_gbs, 3), + "ib_mb": round(fetched["ib_bytes"] / 1e6, 1), + "tcp_mb": round(fetched["tcp_bytes"] / 1e6, 1), } ) - - 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 + csv_handle.flush() + measured_runs += 1 + if not proven: + raise RuntimeError( + f"wire proof failed for protocol={args.protocol} " + f"profile={profile} payload={requested_mib}MiB run={run}" + ) print( - 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 " - 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']}", + f"[{profile} {requested_mib}MiB] byte_exact=PASS wire=PASS " + f"put={statistics.mean(put_rates):.2f}GB/s " + f"get_mean={statistics.mean(get_rates):.2f}GB/s " + f"get_median={statistics.median(get_rates):.2f}GB/s " + f"get_std={statistics.pstdev(get_rates):.2f}GB/s", 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"{'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 = 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:<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 = [] - 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) - - 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", - "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) + print(f"[csv] wrote {measured_runs} measured rounds", flush=True) + finally: + _teardown_benchmark(consumer, owner_attempted) if __name__ == "__main__": diff --git a/scripts/benchmarks/tq_rdma_bench.py b/scripts/benchmarks/tq_rdma_bench.py deleted file mode 100644 index 5978caaef..000000000 --- a/scripts/benchmarks/tq_rdma_bench.py +++ /dev/null @@ -1,414 +0,0 @@ -#!/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.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") - backend_dict = build_mooncake_config( - master_address=master_addr, - device=args.device, - protocol=cfg["protocol"], - global_segment_size=8 * 1024**3, - ) - - return OmegaConf.create( - { - "controller": {"sampler": sampler, "polling_mode": True}, - "backend": backend_dict, - }, - flags={"allow_objects": True}, - ) - - -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 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 - 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 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 - - wait_actor_gone(timeout=timeout) - - -# --------------------------------------------------------------------------- # -# 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 - - 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() - 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 run_benchmark(args: argparse.Namespace) -> None: - """Run the benchmark across all requested payload/field/config - combinations.""" - 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}") - - 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/mm_payload_fixtures.py b/tests/utils/mm_payload_fixtures.py index be226e6f8..f77a110ea 100644 --- a/tests/utils/mm_payload_fixtures.py +++ b/tests/utils/mm_payload_fixtures.py @@ -4,8 +4,8 @@ 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). +* **real** — an external acceptance fixture containing real dataset images + processed through the production Qwen-VL 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. @@ -16,8 +16,8 @@ ``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. +Tests record the returned ``source`` string as a pytest/JUnit property so real- +payload acceptance is auditable in CI artifacts. """ from __future__ import annotations @@ -27,6 +27,7 @@ from typing import Any import torch +from torch.torch_version import TorchVersion from tests.utils.tq._payload_assertions import diff_digests, leaf_digests @@ -51,27 +52,36 @@ def load_real_fixture(max_samples: int | None = None) -> dict[str, Any] | None: 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) + try: + with torch.serialization.safe_globals([TorchVersion]): + bundle = torch.load(path, map_location="cpu", weights_only=True) + except Exception as error: + raise RuntimeError( + f"External multimodal fixture could not be safely loaded ({type(error).__name__})" + ) from None + try: + 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) + except Exception as error: + raise RuntimeError(f"External multimodal fixture has an invalid schema ({type(error).__name__})") from None 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]}" + f"External multimodal fixture failed its manifest ({len(problems)} leaf mismatches); " + "replace it with a verified acceptance artifact." ) 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} + return {"train_data": train_data} def synthetic_mm_train_data(num_samples: int, seed: int = 20260813) -> dict[str, list[Any]]: @@ -98,8 +108,7 @@ 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. + ``"synthetic"``; callers record it in their acceptance results. """ bundle = load_real_fixture(max_samples=num_samples) if bundle is not None and len(bundle["train_data"]["tokens"]) >= num_samples: diff --git a/tests/utils/test_mm_payload_fixtures.py b/tests/utils/test_mm_payload_fixtures.py new file mode 100644 index 000000000..cd1663e24 --- /dev/null +++ b/tests/utils/test_mm_payload_fixtures.py @@ -0,0 +1,71 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Safety contracts for optional external multimodal acceptance fixtures.""" + +import torch + +from tests.utils import mm_payload_fixtures +from tests.utils.tq._payload_assertions import leaf_digests + + +def _bundle() -> dict: + train_data = { + "tokens": [[1, 2, 3]], + "multimodal_train_inputs": [ + { + "pixel_values": torch.arange(12, dtype=torch.float32).reshape(2, 6), + "image_grid_thw": torch.tensor([[1, 1, 2]], dtype=torch.int64), + } + ], + } + manifest = {} + manifest.update(leaf_digests(train_data["multimodal_train_inputs"][0], "sample[0].multimodal_train_inputs")) + manifest.update(leaf_digests(torch.tensor(train_data["tokens"][0]), "sample[0].tokens")) + return {"meta": {"torch_version": torch.__version__}, "train_data": train_data, "manifest": manifest} + + +def test_external_fixture_uses_safe_weights_only_loader(monkeypatch, tmp_path): + path = tmp_path / "fixture.pt" + torch.save(_bundle(), path) + monkeypatch.setenv("RELAX_MM_FIXTURE", str(path)) + + loaded = mm_payload_fixtures.load_real_fixture() + + assert loaded is not None + assert set(loaded) == {"train_data"} + assert loaded["train_data"]["tokens"] == [[1, 2, 3]] + + +def test_external_fixture_load_failure_does_not_leak_path_or_exception(monkeypatch, tmp_path): + path = tmp_path / "private-host-and-model-path.pt" + path.write_bytes(b"not a torch fixture") + monkeypatch.setenv("RELAX_MM_FIXTURE", str(path)) + + try: + mm_payload_fixtures.load_real_fixture() + except RuntimeError as error: + message = str(error) + else: # pragma: no cover - corrupt input must never be accepted + raise AssertionError("corrupt fixture unexpectedly loaded") + + assert str(path) not in message + assert "private-host" not in message + assert "pickle" not in message.lower() + + +def test_external_fixture_manifest_failure_does_not_leak_path(monkeypatch, tmp_path): + path = tmp_path / "private-dataset-location.pt" + bundle = _bundle() + bundle["manifest"] = {} + torch.save(bundle, path) + monkeypatch.setenv("RELAX_MM_FIXTURE", str(path)) + + try: + mm_payload_fixtures.load_real_fixture() + except RuntimeError as error: + message = str(error) + else: # pragma: no cover - invalid manifest must never be accepted + raise AssertionError("fixture with an invalid manifest unexpectedly loaded") + + assert "failed its manifest" in message + assert str(path) not in message diff --git a/tests/utils/test_tq_benchmark_guards.py b/tests/utils/test_tq_benchmark_guards.py index 208cae58d..2fc166f1d 100644 --- a/tests/utils/test_tq_benchmark_guards.py +++ b/tests/utils/test_tq_benchmark_guards.py @@ -1,16 +1,206 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. -"""Regression tests for fail-closed TransferQueue benchmark teardown.""" +"""CPU contracts for the retained cross-node TransferQueue benchmark.""" + +import sys +import warnings +from types import SimpleNamespace import pytest +import torch + +from scripts.benchmarks import tq_cross_node_bench + + +def test_simple_config_builds_without_mooncake_runtime(monkeypatch): + import transfer_queue + + monkeypatch.setattr(transfer_queue, "GRPOGroupNSampler", lambda **_kwargs: object()) + conf = tq_cross_node_bench.build_conf("simple", master="", device="", segment_gib=1) + assert conf.backend.storage_backend == "SimpleStorage" + + +def test_digest_normalizes_dense_and_nested_rows_without_losing_shape(): + dense = {"field": torch.tensor([[1, 2], [3, 4]], dtype=torch.int16)} + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + nested = {"field": torch.nested.nested_tensor(list(dense["field"].unbind()))} + + expected = tq_cross_node_bench.field_byte_digests(dense, ["field"]) + assert tq_cross_node_bench.field_byte_digests(nested, ["field"]) == expected + assert expected["field"][0][:2] == ("torch.int16", (2,)) + + changed = {"field": dense["field"].clone()} + changed["field"][1, 1] += 1 + assert tq_cross_node_bench.field_byte_digests(changed, ["field"]) != expected + + same_bytes_different_shape = {"field": dense["field"].reshape(1, 2, 2)} + one_row = {"field": dense["field"].reshape(1, 4)} + assert tq_cross_node_bench.field_byte_digests(same_bytes_different_shape, ["field"]) != ( + tq_cross_node_bench.field_byte_digests(one_row, ["field"]) + ) + -from scripts.benchmarks import tq_cross_node_bench, tq_rdma_bench +def test_multimodal_profile_uses_production_non_tensor_stack_and_recursive_digest(): + payload = tq_cross_node_bench.make_multimodal_payload(num_samples=3, total_mib=1) + multimodal = payload.get("multimodal_train_inputs") + assert type(multimodal).__name__ == "NonTensorStack" + rows = multimodal.tolist() + assert all(set(row) == {"pixel_values", "image_grid_thw"} for row in rows) + assert len({tuple(row["pixel_values"].shape) for row in rows}) > 1 + + expected = tq_cross_node_bench.field_byte_digests(payload, ["multimodal_train_inputs"]) + rows[1]["pixel_values"][0, 0] += 1 + assert tq_cross_node_bench.field_byte_digests(payload, ["multimodal_train_inputs"]) != expected @pytest.mark.parametrize( - "wait_actor_gone", - [tq_cross_node_bench.wait_actor_gone, tq_rdma_bench.wait_actor_gone], + ("protocol", "ib_bytes", "tcp_bytes", "payload_bytes", "expected"), + [ + ("rdma", 800, 0, 1000, True), + ("rdma", 799, 0, 1000, False), + ("rdma", 900, 901, 1000, False), + ("tcp", 0, 200, 1000, True), + ("tcp", 0, 199, 1000, False), + ("tcp", 201, 200, 1000, False), + ("simple", 0, 1, 1000, True), + ("simple", 1, 0, 1000, False), + ("rdma", 1, 0, 4 * 1024**3, False), + ("tcp", 0, 1, 4 * 1024**3, False), + ], ) -def test_actor_wait_timeout_is_not_silently_ignored(wait_actor_gone): - with pytest.raises(TimeoutError, match="still registered"): - wait_actor_gone(timeout=0) +def test_wire_proof_is_protocol_and_volume_specific(protocol, ib_bytes, tcp_bytes, payload_bytes, expected): + assert tq_cross_node_bench.wire_is_proven(protocol, ib_bytes, tcp_bytes, payload_bytes) is expected + + +def test_read_counters_scopes_rdma_to_selected_hca(tmp_path): + for device, value in (("rdma0", 11), ("rdma1", 29)): + path = tmp_path / "infiniband" / device / "ports" / "1" / "counters" + path.mkdir(parents=True) + (path / "port_rcv_data").write_text(str(value)) + tcp = tmp_path / "net" / "eth0" / "statistics" + tcp.mkdir(parents=True) + (tcp / "rx_bytes").write_text("101") + + counters = tq_cross_node_bench.read_counters("eth0", rdma_device="rdma1", sysfs_root=tmp_path) + + assert counters == {"ib:rdma1:1": 29 * 4, "tcp:eth0": 101} + + +@pytest.mark.parametrize( + "argv", + [ + ["bench", "--protocol", "rdma", "--consumer-node-id", "node", "--tcp-device", "eth0", "--csv", "x"], + [ + "bench", + "--protocol", + "tcp", + "--consumer-node-id", + "node", + "--master", + "master.example:50051", + "--device", + "rdma0", + "--tcp-device", + "eth0", + "--csv", + "x", + ], + [ + "bench", + "--protocol", + "rdma", + "--consumer-node-id", + "node", + "--master", + "master.example:50051", + "--device", + "../rdma0", + "--tcp-device", + "eth0", + "--csv", + "x", + ], + ], +) +def test_cli_rejects_ambiguous_or_unsafe_counter_configuration(monkeypatch, argv): + monkeypatch.setattr(sys, "argv", argv) + with pytest.raises(SystemExit, match="2"): + tq_cross_node_bench.parse_args() + + +class _RemoteCall: + def __init__(self, result): + self.result = result + + def remote(self): + return self.result + + +def test_teardown_keeps_cleanup_order_when_consumer_shutdown_fails(monkeypatch): + events: list[str] = [] + consumer = SimpleNamespace(shutdown=_RemoteCall("shutdown-ref")) + monkeypatch.setattr( + tq_cross_node_bench.ray, + "get", + lambda _ref, timeout: (_ for _ in ()).throw(RuntimeError("consumer shutdown failed")), + ) + monkeypatch.setattr(tq_cross_node_bench.ray, "kill", lambda *_args, **_kwargs: events.append("consumer-killed")) + monkeypatch.setattr(tq_cross_node_bench, "close_tq_unmount_and_wait", lambda: events.append("owner-closed")) + monkeypatch.setattr(tq_cross_node_bench.ray, "shutdown", lambda: events.append("ray-shutdown")) + + with pytest.raises(RuntimeError, match="consumer shutdown failed"): + tq_cross_node_bench._teardown_benchmark(consumer, owner_attempted=True) + + assert events == ["consumer-killed", "owner-closed", "ray-shutdown"] + + +def test_teardown_dirty_cluster_does_not_close_unowned_controller(monkeypatch): + events: list[str] = [] + monkeypatch.setattr(tq_cross_node_bench, "close_tq_unmount_and_wait", lambda: events.append("owner-closed")) + monkeypatch.setattr(tq_cross_node_bench.ray, "shutdown", lambda: events.append("ray-shutdown")) + + tq_cross_node_bench._teardown_benchmark(None, owner_attempted=False) + + assert events == ["ray-shutdown"] + + +def test_clean_cluster_guard_never_kills_a_healthy_existing_controller(monkeypatch): + controller = object() + killed: list[object] = [] + monkeypatch.setattr(tq_cross_node_bench.ray, "get_actor", lambda *_args, **_kwargs: controller) + monkeypatch.setattr(tq_cross_node_bench.ray, "kill", lambda handle, **_kwargs: killed.append(handle)) + + with pytest.raises(RuntimeError, match="clean exclusive Ray cluster"): + tq_cross_node_bench.require_clean_cluster() + + assert killed == [] + + +def test_teardown_shutdowns_ray_when_owner_cleanup_fails(monkeypatch): + events: list[str] = [] + monkeypatch.setattr( + tq_cross_node_bench, + "close_tq_unmount_and_wait", + lambda: (_ for _ in ()).throw(RuntimeError("owner cleanup failed")), + ) + monkeypatch.setattr(tq_cross_node_bench.ray, "shutdown", lambda: events.append("ray-shutdown")) + + with pytest.raises(RuntimeError, match="owner cleanup failed"): + tq_cross_node_bench._teardown_benchmark(None, owner_attempted=True) + + assert events == ["ray-shutdown"] + + +def test_production_controller_cleanup_timeout_fails_closed(monkeypatch): + from relax.utils.tq import lifecycle + + controller = object() + killed: list[object] = [] + monkeypatch.setattr(lifecycle.ray, "get_actor", lambda *_args, **_kwargs: controller) + monkeypatch.setattr(lifecycle.ray, "kill", lambda handle: killed.append(handle)) + + with pytest.raises(lifecycle.TqCleanupTimeout, match="still resolvable"): + lifecycle.kill_tq_controller_and_wait(timeout=0) + + assert killed == [controller] diff --git a/tests/utils/test_tq_dataplane_behavior.py b/tests/utils/test_tq_dataplane_behavior.py index 6dee42e57..f2d696273 100644 --- a/tests/utils/test_tq_dataplane_behavior.py +++ b/tests/utils/test_tq_dataplane_behavior.py @@ -322,13 +322,13 @@ class TestRealMultimodalFullLink: *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). + Payload source is reported in every assertion: ``real`` (an external + fixture built from 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): + def test_multimodal_list_dict_full_link_byte_exact(self, tq_factory, record_property): """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).""" @@ -338,6 +338,7 @@ def test_multimodal_list_dict_full_link_byte_exact(self, tq_factory): num_samples = 4 train_data, source = mm_train_data(num_samples) + record_property("multimodal_payload_source", source) train_data = dict(train_data) train_data["sample_id"] = list(range(num_samples)) batch = dict_to_tensordict(train_data, batch_size=num_samples) diff --git a/tests/utils/test_tq_failure_paths.py b/tests/utils/test_tq_failure_paths.py index f2977029e..1c09a2e65 100644 --- a/tests/utils/test_tq_failure_paths.py +++ b/tests/utils/test_tq_failure_paths.py @@ -1339,7 +1339,7 @@ def test_multi_dtype_shape_roundtrip_is_byte_exact(self, protocol): 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): + def test_multimodal_list_dict_slow_path_roundtrip_is_byte_exact(self, protocol, record_property): """The container production actually ships: one dict per sample. ``multimodal_train_inputs`` reaches MooncakeStore as non-tensor values @@ -1355,4 +1355,5 @@ def test_multimodal_list_dict_slow_path_roundtrip_is_byte_exact(self, protocol): across protocols) and a wedged engine cannot hang the suite. """ status, detail = _run_isolated_roundtrip(_mm_slow_path_worker, protocol, timeout=240) + record_property("multimodal_roundtrip", detail) assert status == "ok", f"{status}: {detail}" diff --git a/tests/utils/tq/test_config.py b/tests/utils/tq/test_config.py index 4af93423b..1d3e353d1 100644 --- a/tests/utils/tq/test_config.py +++ b/tests/utils/tq/test_config.py @@ -210,6 +210,11 @@ def test_master_address_is_not_re_read_from_env(self, monkeypatch): mc = build_mooncake_config(master_address=_MASTER)["MooncakeStore"] assert mc["master_server_address"] == _MASTER + @pytest.mark.parametrize("master", ["not-an-endpoint", "host:0", "fe80::1", 50051]) + def test_builder_rejects_unvalidated_master_address(self, master): + with pytest.raises(ValueError, match="master_address"): + build_mooncake_config(master_address=master) + @pytest.mark.parametrize("segment_size", [0, -1, True, 1.5]) def test_explicit_segment_size_must_be_a_positive_integer(self, segment_size): with pytest.raises(ValueError, match="positive integer"): diff --git a/tests/utils/tq/test_payload_assertions.py b/tests/utils/tq/test_payload_assertions.py index 027b86fcf..37a9ebaf4 100644 --- a/tests/utils/tq/test_payload_assertions.py +++ b/tests/utils/tq/test_payload_assertions.py @@ -5,13 +5,13 @@ from __future__ import annotations import struct +import warnings from typing import Any import numpy as np import pytest import torch -from relax.utils.payload_digest import leaf_digests as benchmark_leaf_digests from tests.utils.tq._payload_assertions import diff_digests, leaf_digests @@ -45,31 +45,29 @@ def test_leaf_digests_distinguishes_raw_bytes_from_value_equality(): ] -@pytest.mark.parametrize("digest_fn", [leaf_digests, benchmark_leaf_digests], ids=["test-helper", "benchmark-helper"]) -def test_scalar_nan_payload_bits_are_not_collapsed_by_repr(digest_fn): +def test_scalar_nan_payload_bits_are_not_collapsed_by_repr(): first = struct.unpack("!d", bytes.fromhex("7ff8000000000001"))[0] second = struct.unpack("!d", bytes.fromhex("7ff8000000000002"))[0] assert repr(first) == repr(second) == "nan" - assert digest_fn(first) != digest_fn(second) + assert leaf_digests(first) != leaf_digests(second) -@pytest.mark.parametrize("digest_fn", [leaf_digests, benchmark_leaf_digests], ids=["test-helper", "benchmark-helper"]) -def test_dict_paths_do_not_collapse_dotted_keys_into_nested_keys(digest_fn): - digests = digest_fn({"a.b": 1, "a": {"b": 2}}) +def test_dict_paths_do_not_collapse_dotted_keys_into_nested_keys(): + digests = leaf_digests({"a.b": 1, "a": {"b": 2}}) assert len(digests) == 2 assert set(digests) == {"payload['a.b']", "payload.a.b"} -@pytest.mark.parametrize("digest_fn", [leaf_digests, benchmark_leaf_digests], ids=["test-helper", "benchmark-helper"]) -def test_non_string_dict_keys_fail_loudly(digest_fn): +def test_non_string_dict_keys_fail_loudly(): with pytest.raises(TypeError, match="Unsupported payload dict key at payload: int"): - digest_fn({1: "value"}) + leaf_digests({1: "value"}) def test_leaf_digests_preserves_dtype_shape_and_nested_tensor_rows(): - with pytest.warns(UserWarning, match="prototype stage"): + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) nested = torch.nested.nested_tensor( [torch.tensor([1, 2], dtype=torch.int16), torch.tensor([3], dtype=torch.int16)] ) From dc17f9a937e78928d3fd16408556cf4118602a84 Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Mon, 24 Aug 2026 04:25:20 +0800 Subject: [PATCH 25/30] test(tq): consolidate RDMA coverage Parameterize straightforward backend and lifecycle cases while keeping complex timeout and cleanup paths explicit. Remove duplicate endpoint and fixture harnesses, retain byte-exact and process-isolation coverage, and align the acceptance documentation. --- docs/draft/transfer_queue_rdma.md | 16 +- relax/distributed/ray/actor_group.py | 8 +- tests/core/test_controller_tq_backend.py | 215 ++++----- tests/utils/mm_payload_fixtures.py | 116 ----- tests/utils/test_mm_payload_fixtures.py | 71 --- tests/utils/test_tq_benchmark_guards.py | 94 ++-- tests/utils/test_tq_dataplane_behavior.py | 173 ++----- tests/utils/test_tq_failure_paths.py | 563 +++++----------------- tests/utils/tq/test_config.py | 187 ++++--- 9 files changed, 374 insertions(+), 1069 deletions(-) delete mode 100644 tests/utils/mm_payload_fixtures.py delete mode 100644 tests/utils/test_mm_payload_fixtures.py diff --git a/docs/draft/transfer_queue_rdma.md b/docs/draft/transfer_queue_rdma.md index 596dc88e5..1afe9825a 100644 --- a/docs/draft/transfer_queue_rdma.md +++ b/docs/draft/transfer_queue_rdma.md @@ -133,19 +133,10 @@ Mooncake 配置固定 `hard_pin=true`，不会为了腾空间静默驱逐已经 - `NOTIFY_DATA_UPDATE_ACK` 必须检查 positive ACK，controller 拒绝更新时 producer 不得按成功返回； - 合入修复后更新 Relax 的明确 commit pin，并以稳定的 version/capability marker 校验，而不是仅凭 method name 或源码文本推断能力。 -在上述上游 PR 与新 pin 落地前，这个正确性门槛仍是明确的合入 blocker；本 PR 不通过 monkey patch 改写 TransferQueue。Relax 侧继续保留失败 store 的“写失败、production 状态不更新”契约测试，以及隔离 master 的物理容量溢出故障注入。 +在上述上游 PR 与新 pin 落地前，这个正确性门槛仍是明确的合入 blocker；本 PR 不通过 monkey patch 改写 TransferQueue。Relax 侧继续保留容量预检、SimpleStorage backpressure，以及失败 store 的“写失败、production 状态不更新”契约测试。真实 Mooncake 的容量与传输验收统一由跨节点 benchmark 承担，不再在 pytest 中维护第二套 direct-client harness。 正确性守卫强制 `MC_STORE_MEMCPY=0` 且 **fail-closed**：mooncake 0.3.10 在 TCP-only 环境会自动启用 memcpy 快拷贝路径，该路径存在已确认的静默截断缺陷（现象与处置见排障表）；RDMA 会话本就自动禁用 memcpy，不受影响。由于缺陷在当前 pin 上已实证，显式导出 `MC_STORE_MEMCPY=1` 会在启动时被直接拒绝，待 pin 升级到修复版本后再按版本重新放开。 -真机容量故障注入会故意创建 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 测试必须分别报告，前者不能替代后者。 @@ -153,13 +144,12 @@ Mock/本机测试和真实双节点 RDMA 测试必须分别报告，前者不能 | 层级 | 验证内容 | 通过标准 | |---|---|---| | CI/mock | 模式校验、master 端点格式、容量预检、owner 超时与终止确认、controller 清理、attach 握手的 manager/config 契约与拆除顺序、auto/required、有限重试、写失败不发布状态 | `tests/utils/tq/test_config.py`、`tests/core/test_controller_tq_backend.py` 与 `tests/utils/test_tq_failure_paths.py` 全部通过；真机项允许明确 skip | -| 本机 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 | +| 本机 TQ | SimpleStorage 全链路 put/get、容量 backpressure、空读、清理、字节一致性；`multimodal_train_inputs` 以生产容器（`list[dict]` / NonTensorStack，存储层非张量路径）全链路逐叶子 SHA-256 一致 | `tests/utils/test_tq_dataplane_behavior.py` 通过（含 `TestMultimodalFullLink`） | | 真实多模态载荷 | 真实数据集图像走完整生产预处理链（`build_messages` → `apply_chat_template` → `process_vision_info` → HF processor → `remap_mm_train_inputs`）生成 fixture；上述两级多模态用例检测到 fixture 后自动升级为真实载荷档 | fixture 存在时以 `[real]` 档通过；无 fixture 环境回退 `[synthetic]`（生产同构状，CI 兜底）；交付报告须注明真实档在何处跑过 | | 真实双节点 | 同一 driver/consumer pair 下各 backend 的原生数据布局；SimpleStorage、Mooncake/TCP、Mooncake/RDMA；synthetic、production-shaped multimodal 两种 profile；256/1024/2048/4096 MiB；每档 warmup + 至少 5 轮 | 每次 get 的 dtype、shape 与 raw-byte SHA-256 全部 PASS；强制 counter gate 证明对应线路；逐协议、逐轮 CSV 留档并报告均值、median、stddev | | 真实回退（`auto` + 某节点 RDMA 不可用） | 静态探测删除后，`auto` 会真的创建 Mooncake owner 与 named controller，再在 handshake 阶段失败并拆除，因此这条路径必须实测 | 逐条确认：① 日志显示 Mooncake owner `tq.init` **成功**、随后 handshake 阶段失败（否则实际只测到 owner 初始化失败，覆盖不到拆除）；② 最终存在且仅存在一个预期的 SimpleStorage `TransferQueueController`；③ 该 controller 的 stored config 确认为 SimpleStorage；④ SimpleStorage 的 put/get 正常工作；⑤ 旧 Mooncake owner actor 与其 client 均已消失；⑥ Mooncake segment 已从 master 卸载——若测的是强杀/超时路径，需等过 `client_ttl`（默认 30 s）再检查 | -真实模型/数据 fixture 及其生成流程属于外部 PR 验收附件，不再由生产仓库维护。可选的本机 dataplane 测试仍可通过 `RELAX_MM_FIXTURE` 指向已验证的外部 fixture；唯一保留的跨节点 benchmark 只接受显式的 `synthetic` 和 `multimodal` profile，不会在缺失 fixture 时替换 profile。 +真实模型/数据 fixture 及其生成流程属于外部 PR 验收附件，不再由仓库测试代码加载或维护。本机 dataplane 测试使用确定性的 production-shaped synthetic payload；唯一保留的跨节点 benchmark 只接受显式的 `synthetic` 和 `multimodal` profile，不会在缺失外部 fixture 时替换验收口径。 双节点验收命令（master 与 Ray 集群需由部署侧预先准备；每个 protocol 必须启动一个全新的 Python 进程并使用独立 CSV）： diff --git a/relax/distributed/ray/actor_group.py b/relax/distributed/ray/actor_group.py index 178e05edc..544000c7c 100644 --- a/relax/distributed/ray/actor_group.py +++ b/relax/distributed/ray/actor_group.py @@ -121,7 +121,13 @@ def async_init(self, args, role, with_ref=False, with_opd_teacher=False): for actor in self._actor_handlers ] - def init_and_wait(self, args, role, with_ref=False, with_opd_teacher=False): + def init_and_wait( + self, + args: Any, + role: str, + with_ref: bool = False, + with_opd_teacher: bool = False, + ) -> list[Any]: """Initialize every train actor, destroying the group on any failure. ``MegatronTrainRayActor.init`` attaches the process-global diff --git a/tests/core/test_controller_tq_backend.py b/tests/core/test_controller_tq_backend.py index 4ebbf9036..df82dc282 100644 --- a/tests/core/test_controller_tq_backend.py +++ b/tests/core/test_controller_tq_backend.py @@ -92,6 +92,10 @@ def _assert_simple_storage(backend: dict) -> None: assert "MooncakeStore" not in backend +def _raise(error: BaseException) -> None: + raise error + + class TestOffMode: """``off`` is the untouched SimpleStorage path: it checks nothing.""" @@ -139,7 +143,7 @@ def test_healthy_existing_controller_aborts_before_legacy_init(self, monkeypatch monkeypatch.setattr( controller, "reap_unusable_tq_controller", - lambda: (_ for _ in ()).throw(RuntimeError("exclusive cluster is not clean")), + lambda: _raise(RuntimeError("exclusive cluster is not clean")), ) monkeypatch.setattr( controller.tq, @@ -162,62 +166,62 @@ def test_resolver_never_probes_hardware(self, monkeypatch): _resolve(_config()) assert recorder.calls == ["contract", "master"] - def test_probe_helpers_are_no_longer_importable(self): - with pytest.raises(ModuleNotFoundError): - __import__("relax.utils.rdma_probe") - - def test_controller_module_holds_no_probe_symbols(self): - for name in ("probe_cluster_nodes", "reduce_results", "probe_node", "EffectiveConfig"): - assert not hasattr(controller, name), f"{name} should be gone with the static probe" - - -class TestAutoFallsBackForEveryUnmetPrecondition: - """Gate A: in ``auto``, anything short of host RDMA yields - SimpleStorage.""" - - def test_contract_failure_falls_back_before_touching_master(self, monkeypatch): - recorder = _Recorder(monkeypatch, contract_error=RuntimeError("retry guard missing")) - _assert_simple_storage(_resolve(_config())) - assert recorder.calls == ["contract"] - - def test_missing_master_endpoint_falls_back(self, monkeypatch): - recorder = _Recorder(monkeypatch, master_error=RuntimeError("MC_MASTER_ADDRESS required")) - _assert_simple_storage(_resolve(_config())) - assert recorder.calls == ["contract", "master"] - - def test_malformed_master_endpoint_falls_back(self, monkeypatch): - recorder = _Recorder( - monkeypatch, master_error=RuntimeError("MC_MASTER_ADDRESS is not a usable endpoint: missing host") - ) - _assert_simple_storage(_resolve(_config())) - assert recorder.calls == ["contract", "master"] - - def test_insufficient_segment_capacity_falls_back(self, monkeypatch): - """Worst-case multimodal payload far exceeds the default segment.""" - _Recorder(monkeypatch) - config = _config(multimodal_keys=["pixel_values"], rollout_batch_size=64, n_samples_per_prompt=8) - _assert_simple_storage(_resolve(config)) - def test_unusable_segment_size_override_falls_back(self, monkeypatch): - """Capacity validation *raises* here rather than returning a reason. +class TestBackendPreconditions: + """The same unmet precondition degrades ``auto`` and aborts + ``required``.""" - A garbage override is a configuration failure like any other and must - not abort an ``auto`` run. - """ - _Recorder(monkeypatch) - monkeypatch.setenv("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB", "four") - _assert_simple_storage(_resolve(_config())) - - @pytest.mark.parametrize("value", ["nan", "inf", "-inf"]) - def test_non_finite_segment_size_override_falls_back(self, monkeypatch, value): - _Recorder(monkeypatch) - monkeypatch.setenv("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB", value) - _assert_simple_storage(_resolve(_config())) + @pytest.mark.parametrize("mode", ["auto", "required"]) + @pytest.mark.parametrize( + ("recorder_kwargs", "config_overrides", "segment_override", "match"), + [ + ({"contract_error": RuntimeError("retry guard missing")}, {}, None, "correctness contract"), + ({"master_error": RuntimeError("MC_MASTER_ADDRESS required")}, {}, None, "master endpoint"), + ({"master_error": RuntimeError("unusable master endpoint")}, {}, None, "master endpoint"), + ( + {}, + {"multimodal_keys": ["pixel_values"], "rollout_batch_size": 64, "n_samples_per_prompt": 8}, + None, + "segment capacity insufficient", + ), + ({}, {}, "four", "segment-capacity configuration is unusable"), + ({}, {}, "nan", "segment-capacity configuration is unusable"), + ({}, {}, "inf", "segment-capacity configuration is unusable"), + ({}, {}, "-inf", "segment-capacity configuration is unusable"), + ({}, {"seq_length": None}, None, "segment-capacity configuration is unusable"), + ], + ids=[ + "correctness-contract", + "missing-master", + "malformed-master", + "insufficient-capacity", + "invalid-segment-size", + "nan-segment-size", + "infinite-segment-size", + "negative-infinite-segment-size", + "missing-seq-length", + ], + ) + def test_unmet_precondition( + self, + monkeypatch, + mode, + recorder_kwargs, + config_overrides, + segment_override, + match, + ): + recorder = _Recorder(monkeypatch, **recorder_kwargs) + if segment_override is not None: + monkeypatch.setenv("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB", segment_override) - def test_missing_seq_length_falls_back(self, monkeypatch): - """Without ``seq_length`` the payload bound cannot be derived.""" - _Recorder(monkeypatch) - _assert_simple_storage(_resolve(_config(seq_length=None))) + config = _config(tq_rdma_mode=mode, **config_overrides) + if mode == "auto": + _assert_simple_storage(_resolve(config)) + else: + with pytest.raises(RuntimeError, match=match): + _resolve(config) + assert recorder.calls == (["contract"] if recorder_kwargs.get("contract_error") else ["contract", "master"]) def test_satisfied_preconditions_select_host_rdma(self, monkeypatch): _Recorder(monkeypatch) @@ -226,61 +230,20 @@ def test_satisfied_preconditions_select_host_rdma(self, monkeypatch): assert backend["MooncakeStore"]["protocol"] == "rdma" assert backend["MooncakeStore"]["device_name"] == "mlx5_0" - def test_validated_master_endpoint_reaches_the_client_config(self, monkeypatch): - """The checked endpoint must be the one handed to Mooncake. + def test_required_accepts_satisfied_preconditions(self, monkeypatch): + _Recorder(monkeypatch) + backend = _resolve(_config(tq_rdma_mode="required")) + assert backend["MooncakeStore"]["protocol"] == "rdma" - ``build_mooncake_config`` must not re-read ``MC_MASTER_ADDRESS``; a - divergent env value here would surface as the wrong address. - """ + def test_validated_master_endpoint_reaches_the_client_config(self, monkeypatch): + """The checked endpoint, rather than a later env value, reaches + Mooncake.""" _Recorder(monkeypatch) monkeypatch.setenv("MC_MASTER_ADDRESS", "someone.else.invalid:9999") backend = _resolve(_config()) assert backend["MooncakeStore"]["master_server_address"] == _MASTER -class TestRequiredFailsFast: - """``required`` must never silently downgrade the same failures.""" - - @pytest.mark.parametrize( - ("kwargs", "match"), - [ - ({"contract_error": RuntimeError("retry guard missing")}, "correctness contract"), - ({"master_error": RuntimeError("MC_MASTER_ADDRESS required")}, "master endpoint is not configured"), - ], - ) - def test_required_raises(self, monkeypatch, kwargs, match): - _Recorder(monkeypatch, **kwargs) - with pytest.raises(RuntimeError, match=match): - _resolve(_config(tq_rdma_mode="required")) - - def test_required_raises_on_insufficient_capacity(self, monkeypatch): - _Recorder(monkeypatch) - config = _config( - tq_rdma_mode="required", - multimodal_keys=["pixel_values"], - rollout_batch_size=64, - n_samples_per_prompt=8, - ) - with pytest.raises(RuntimeError, match="segment capacity insufficient"): - _resolve(config) - - def test_required_raises_on_unusable_segment_size_override(self, monkeypatch): - _Recorder(monkeypatch) - monkeypatch.setenv("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB", "four") - with pytest.raises(RuntimeError, match="segment-capacity configuration is unusable"): - _resolve(_config(tq_rdma_mode="required")) - - def test_required_raises_on_missing_seq_length(self, monkeypatch): - _Recorder(monkeypatch) - with pytest.raises(RuntimeError, match="segment-capacity configuration is unusable"): - _resolve(_config(tq_rdma_mode="required", seq_length=None)) - - def test_required_accepts_satisfied_preconditions(self, monkeypatch): - _Recorder(monkeypatch) - backend = _resolve(_config(tq_rdma_mode="required")) - assert backend["MooncakeStore"]["protocol"] == "rdma" - - class TestModeValidation: """A malformed mode is a configuration error in every mode.""" @@ -310,17 +273,6 @@ def test_selected_backend_is_rdma_or_simple(self, monkeypatch, mode): else: _assert_simple_storage(backend) - @pytest.mark.parametrize( - "kwargs", - [ - {"contract_error": RuntimeError("gate failed")}, - {"master_error": RuntimeError("no endpoint")}, - ], - ) - def test_auto_fallback_never_yields_mooncake(self, monkeypatch, kwargs): - _Recorder(monkeypatch, **kwargs) - _assert_simple_storage(_resolve(_config())) - def test_cli_exposes_only_mode_and_device(arguments_module): """The narrowed CLI keeps exactly two TransferQueue RDMA flags.""" @@ -406,29 +358,18 @@ def test_success_keeps_mooncake_and_touches_no_cleanup(self, monkeypatch): assert recorder.events == ["handshake"] def test_auto_closes_owner_before_initializing_simple_storage(self, monkeypatch): - recorder = _AttachRecorder(monkeypatch, failures=["node-B: attach timed out"]) - result = _confirm(_config()) - # Ordering is the point: SimpleStorage must not be initialised while a - # half-initialised Mooncake controller may still be registered. + recorder = _AttachRecorder(monkeypatch, failures=["node: attach timed out"]) + result = _confirm(_config(), owner="attempt-owner") assert recorder.events == ["handshake", "close", "init_simple"] + assert recorder.closed == ["attempt-owner"] assert result.config == "simple-conf" assert result.fallback_reason == "attach_handshake_failed:1_failures" - def test_cleanup_receives_this_attempts_mooncake_owner(self, monkeypatch): - recorder = _AttachRecorder(monkeypatch, failures=["node-B: attach timed out"]) - _confirm(_config(), owner="the-mooncake-owner") - assert recorder.closed == ["the-mooncake-owner"] - def test_cleanup_failure_aborts_instead_of_falling_back(self, monkeypatch): - """If teardown fails, global TQ state is unknown. - - Starting SimpleStorage on top of it could attach to a dirty controller, - so the cleanup error must propagate. - """ recorder = _AttachRecorder( monkeypatch, - failures=["node-B: attach timed out"], - close_error=RuntimeError("TransferQueue owner cleanup failed: close timed out"), + failures=["node: attach timed out"], + close_error=RuntimeError("TransferQueue owner cleanup failed"), ) with pytest.raises(RuntimeError, match="owner cleanup failed"): _confirm(_config()) @@ -436,7 +377,7 @@ def test_cleanup_failure_aborts_instead_of_falling_back(self, monkeypatch): assert recorder.initialized == [] def test_required_closes_owner_then_raises(self, monkeypatch): - recorder = _AttachRecorder(monkeypatch, failures=["node-B: protocol=tcp"]) + recorder = _AttachRecorder(monkeypatch, failures=["node: protocol=tcp"]) with pytest.raises(RuntimeError, match="attach handshake reported"): _confirm(_config(tq_rdma_mode="required")) assert recorder.events == ["handshake", "close"] @@ -486,3 +427,17 @@ def record_cleanup(instance): controller.Controller(config) assert events == ["mooncake-owner"] + + def test_repeated_data_system_cleanup_closes_the_owner_once(self, monkeypatch): + instance = controller.Controller.__new__(controller.Controller) + instance._tq_legacy_init = False + instance._tq_owner = "owner" + closed: list[object] = [] + + monkeypatch.setattr(controller, "close_tq_owner", lambda owner: closed.append(owner) if owner else None) + + instance._close_data_system() + instance._close_data_system() + + assert closed == ["owner"] + assert instance._tq_owner is None diff --git a/tests/utils/mm_payload_fixtures.py b/tests/utils/mm_payload_fixtures.py deleted file mode 100644 index f77a110ea..000000000 --- a/tests/utils/mm_payload_fixtures.py +++ /dev/null @@ -1,116 +0,0 @@ -# Copyright (c) 2026 Relax Authors. All Rights Reserved. - -"""Shared multimodal payload sources for TransferQueue byte-exact tests. - -Two tiers, selected automatically: - -* **real** — an external acceptance fixture containing real dataset images - processed through the production Qwen-VL 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 record the returned ``source`` string as a pytest/JUnit property so real- -payload acceptance is auditable in CI artifacts. -""" - -from __future__ import annotations - -import os -from pathlib import Path -from typing import Any - -import torch -from torch.torch_version import TorchVersion - -from tests.utils.tq._payload_assertions 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 - try: - with torch.serialization.safe_globals([TorchVersion]): - bundle = torch.load(path, map_location="cpu", weights_only=True) - except Exception as error: - raise RuntimeError( - f"External multimodal fixture could not be safely loaded ({type(error).__name__})" - ) from None - try: - 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) - except Exception as error: - raise RuntimeError(f"External multimodal fixture has an invalid schema ({type(error).__name__})") from None - if problems: - raise RuntimeError( - f"External multimodal fixture failed its manifest ({len(problems)} leaf mismatches); " - "replace it with a verified acceptance artifact." - ) - 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 {"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"``; callers record it in their acceptance results. - """ - 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_mm_payload_fixtures.py b/tests/utils/test_mm_payload_fixtures.py deleted file mode 100644 index cd1663e24..000000000 --- a/tests/utils/test_mm_payload_fixtures.py +++ /dev/null @@ -1,71 +0,0 @@ -# Copyright (c) 2026 Relax Authors. All Rights Reserved. - -"""Safety contracts for optional external multimodal acceptance fixtures.""" - -import torch - -from tests.utils import mm_payload_fixtures -from tests.utils.tq._payload_assertions import leaf_digests - - -def _bundle() -> dict: - train_data = { - "tokens": [[1, 2, 3]], - "multimodal_train_inputs": [ - { - "pixel_values": torch.arange(12, dtype=torch.float32).reshape(2, 6), - "image_grid_thw": torch.tensor([[1, 1, 2]], dtype=torch.int64), - } - ], - } - manifest = {} - manifest.update(leaf_digests(train_data["multimodal_train_inputs"][0], "sample[0].multimodal_train_inputs")) - manifest.update(leaf_digests(torch.tensor(train_data["tokens"][0]), "sample[0].tokens")) - return {"meta": {"torch_version": torch.__version__}, "train_data": train_data, "manifest": manifest} - - -def test_external_fixture_uses_safe_weights_only_loader(monkeypatch, tmp_path): - path = tmp_path / "fixture.pt" - torch.save(_bundle(), path) - monkeypatch.setenv("RELAX_MM_FIXTURE", str(path)) - - loaded = mm_payload_fixtures.load_real_fixture() - - assert loaded is not None - assert set(loaded) == {"train_data"} - assert loaded["train_data"]["tokens"] == [[1, 2, 3]] - - -def test_external_fixture_load_failure_does_not_leak_path_or_exception(monkeypatch, tmp_path): - path = tmp_path / "private-host-and-model-path.pt" - path.write_bytes(b"not a torch fixture") - monkeypatch.setenv("RELAX_MM_FIXTURE", str(path)) - - try: - mm_payload_fixtures.load_real_fixture() - except RuntimeError as error: - message = str(error) - else: # pragma: no cover - corrupt input must never be accepted - raise AssertionError("corrupt fixture unexpectedly loaded") - - assert str(path) not in message - assert "private-host" not in message - assert "pickle" not in message.lower() - - -def test_external_fixture_manifest_failure_does_not_leak_path(monkeypatch, tmp_path): - path = tmp_path / "private-dataset-location.pt" - bundle = _bundle() - bundle["manifest"] = {} - torch.save(bundle, path) - monkeypatch.setenv("RELAX_MM_FIXTURE", str(path)) - - try: - mm_payload_fixtures.load_real_fixture() - except RuntimeError as error: - message = str(error) - else: # pragma: no cover - invalid manifest must never be accepted - raise AssertionError("fixture with an invalid manifest unexpectedly loaded") - - assert "failed its manifest" in message - assert str(path) not in message diff --git a/tests/utils/test_tq_benchmark_guards.py b/tests/utils/test_tq_benchmark_guards.py index 2fc166f1d..f23f02fd6 100644 --- a/tests/utils/test_tq_benchmark_guards.py +++ b/tests/utils/test_tq_benchmark_guards.py @@ -12,6 +12,25 @@ from scripts.benchmarks import tq_cross_node_bench +def _argv(protocol: str, *extra: str) -> list[str]: + return [ + "bench", + "--protocol", + protocol, + "--consumer-node-id", + "node", + "--tcp-device", + "eth0", + "--csv", + "x", + *extra, + ] + + +def _raise(error: BaseException) -> None: + raise error + + def test_simple_config_builds_without_mooncake_runtime(monkeypatch): import transfer_queue @@ -90,37 +109,9 @@ def test_read_counters_scopes_rdma_to_selected_hca(tmp_path): @pytest.mark.parametrize( "argv", [ - ["bench", "--protocol", "rdma", "--consumer-node-id", "node", "--tcp-device", "eth0", "--csv", "x"], - [ - "bench", - "--protocol", - "tcp", - "--consumer-node-id", - "node", - "--master", - "master.example:50051", - "--device", - "rdma0", - "--tcp-device", - "eth0", - "--csv", - "x", - ], - [ - "bench", - "--protocol", - "rdma", - "--consumer-node-id", - "node", - "--master", - "master.example:50051", - "--device", - "../rdma0", - "--tcp-device", - "eth0", - "--csv", - "x", - ], + _argv("rdma"), + _argv("tcp", "--master", "master.example:50051", "--device", "rdma0"), + _argv("rdma", "--master", "master.example:50051", "--device", "../rdma0"), ], ) def test_cli_rejects_ambiguous_or_unsafe_counter_configuration(monkeypatch, argv): @@ -129,40 +120,33 @@ def test_cli_rejects_ambiguous_or_unsafe_counter_configuration(monkeypatch, argv tq_cross_node_bench.parse_args() -class _RemoteCall: - def __init__(self, result): - self.result = result - - def remote(self): - return self.result +@pytest.fixture +def teardown_events(monkeypatch): + events: list[str] = [] + monkeypatch.setattr(tq_cross_node_bench.ray, "kill", lambda *_args, **_kwargs: events.append("consumer-killed")) + monkeypatch.setattr(tq_cross_node_bench, "close_tq_unmount_and_wait", lambda: events.append("owner-closed")) + monkeypatch.setattr(tq_cross_node_bench.ray, "shutdown", lambda: events.append("ray-shutdown")) + return events -def test_teardown_keeps_cleanup_order_when_consumer_shutdown_fails(monkeypatch): - events: list[str] = [] - consumer = SimpleNamespace(shutdown=_RemoteCall("shutdown-ref")) +def test_teardown_keeps_cleanup_order_when_consumer_shutdown_fails(monkeypatch, teardown_events): + consumer = SimpleNamespace(shutdown=SimpleNamespace(remote=lambda: "shutdown-ref")) monkeypatch.setattr( tq_cross_node_bench.ray, "get", - lambda _ref, timeout: (_ for _ in ()).throw(RuntimeError("consumer shutdown failed")), + lambda _ref, timeout: _raise(RuntimeError("consumer shutdown failed")), ) - monkeypatch.setattr(tq_cross_node_bench.ray, "kill", lambda *_args, **_kwargs: events.append("consumer-killed")) - monkeypatch.setattr(tq_cross_node_bench, "close_tq_unmount_and_wait", lambda: events.append("owner-closed")) - monkeypatch.setattr(tq_cross_node_bench.ray, "shutdown", lambda: events.append("ray-shutdown")) with pytest.raises(RuntimeError, match="consumer shutdown failed"): tq_cross_node_bench._teardown_benchmark(consumer, owner_attempted=True) - assert events == ["consumer-killed", "owner-closed", "ray-shutdown"] + assert teardown_events == ["consumer-killed", "owner-closed", "ray-shutdown"] -def test_teardown_dirty_cluster_does_not_close_unowned_controller(monkeypatch): - events: list[str] = [] - monkeypatch.setattr(tq_cross_node_bench, "close_tq_unmount_and_wait", lambda: events.append("owner-closed")) - monkeypatch.setattr(tq_cross_node_bench.ray, "shutdown", lambda: events.append("ray-shutdown")) - +def test_teardown_dirty_cluster_does_not_close_unowned_controller(teardown_events): tq_cross_node_bench._teardown_benchmark(None, owner_attempted=False) - assert events == ["ray-shutdown"] + assert teardown_events == ["ray-shutdown"] def test_clean_cluster_guard_never_kills_a_healthy_existing_controller(monkeypatch): @@ -177,19 +161,17 @@ def test_clean_cluster_guard_never_kills_a_healthy_existing_controller(monkeypat assert killed == [] -def test_teardown_shutdowns_ray_when_owner_cleanup_fails(monkeypatch): - events: list[str] = [] +def test_teardown_shutdowns_ray_when_owner_cleanup_fails(monkeypatch, teardown_events): monkeypatch.setattr( tq_cross_node_bench, "close_tq_unmount_and_wait", - lambda: (_ for _ in ()).throw(RuntimeError("owner cleanup failed")), + lambda: _raise(RuntimeError("owner cleanup failed")), ) - monkeypatch.setattr(tq_cross_node_bench.ray, "shutdown", lambda: events.append("ray-shutdown")) with pytest.raises(RuntimeError, match="owner cleanup failed"): tq_cross_node_bench._teardown_benchmark(None, owner_attempted=True) - assert events == ["ray-shutdown"] + assert teardown_events == ["ray-shutdown"] def test_production_controller_cleanup_timeout_fails_closed(monkeypatch): diff --git a/tests/utils/test_tq_dataplane_behavior.py b/tests/utils/test_tq_dataplane_behavior.py index f2d696273..fad340c54 100644 --- a/tests/utils/test_tq_dataplane_behavior.py +++ b/tests/utils/test_tq_dataplane_behavior.py @@ -1,33 +1,10 @@ # 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). - 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. +"""Real local SimpleStorage data-plane contracts. + +Covers connection, dense/multimodal byte identity, backpressure, empty get, +repeat put, and clear/reinit. Cross-node transport and disconnect behavior +belong to the retained C0/C1/C2 benchmark. """ from __future__ import annotations @@ -40,20 +17,10 @@ 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. - """ + """Distinguish the real package from CI's single-file TQ stub.""" 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 @@ -72,11 +39,6 @@ def _has_real_submodule(dotted: str) -> bool: _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. @@ -106,24 +68,14 @@ def _force_kill_controller() -> None: 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. - """ + """Flatten a dense tensor or NestedTensor to its comparable storage.""" if type(t).__name__ == "NestedTensor": return t.values().reshape(-1) 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. - """ + """Extract one row from a dense, nested, or non-tensor column.""" if isinstance(column, torch.Tensor) and column.is_nested: return column.unbind()[row_position] return column[row_position] @@ -140,6 +92,24 @@ def _payload(n: int, fields: list[str], cols: int, dtype: str = "float32", seed: return TensorDict(data, batch_size=[n]) +def _multimodal_payload(num_samples: int) -> dict: + """Deterministic Qwen3-VL-shaped payload for the non-tensor TQ path.""" + grids = ((1, 58, 64), (1, 34, 64), (1, 64, 64), (1, 26, 40)) + generator = torch.Generator().manual_seed(20260813) + multimodal = [] + tokens = [] + for index in range(num_samples): + t, h, w = grids[index % len(grids)] + multimodal.append( + { + "pixel_values": torch.randn(t * h * w, 1536, generator=generator), + "image_grid_thw": torch.tensor([[t, h, w]], dtype=torch.int64), + } + ) + tokens.append(torch.randint(0, 151_000, (512 + 173 * index,), generator=generator).tolist()) + return {"tokens": tokens, "multimodal_train_inputs": multimodal} + + 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) @@ -147,11 +117,6 @@ def _round_trip(client, payload, partition: str, fields: list[str], n: int): return client.get_data(meta) -# --------------------------------------------------------------------------- -# Fixtures -# --------------------------------------------------------------------------- - - @pytest.fixture(scope="module") def _ray_cluster(): import ray @@ -166,12 +131,7 @@ def _ray_cluster(): @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. - """ + """Yield an F10-safe ``reinit(capacity, units) -> client`` factory.""" import transfer_queue as tq from omegaconf import OmegaConf from transfer_queue import GRPOGroupNSampler @@ -201,54 +161,32 @@ def _reinit(capacity: int = 1024, units: int = 1): 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.""" + @pytest.mark.parametrize( + ("fields", "columns", "samples", "seed"), + [ + (["a", "b"], 8, 4, 0), + (["img", "txt", "mask"], 16, 8, 42), + (["pixel_values"], 1176, 4, 7), + ], + ids=["connection", "multi-field", "multimodal-width"], + ) + def test_dense_round_trip_is_byte_exact(self, tq_factory, fields, columns, samples, seed): 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()}" + payload = _payload(n=samples, fields=fields, cols=columns, seed=seed) + got = _round_trip(client, payload, "dense", fields, samples) + assert set(fields) <= set(got.keys()) + for field in fields: + gv, av = _flat_values(got[field]), _flat_values(payload[field]) + assert 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) + assert torch.equal(gv, av), f"{field}: not byte-exact" def test_backpressure_raises_on_capacity_overflow(self, tq_factory): """A single put exceeding capacity raises rather than silently @@ -313,31 +251,16 @@ def test_cleanup_clear_partition_then_reinit_isolated(self, tq_factory): 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`` (an external - fixture built from actual dataset images through the production Qwen-VL - processor chain) or ``synthetic`` (production-structured fallback, - CI-safe). - """ +class TestMultimodalFullLink: + """Production NonTensorStack container survives the full link exactly.""" def test_multimodal_list_dict_full_link_byte_exact(self, tq_factory, record_property): - """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.utils import dict_to_tensordict - from tests.utils.mm_payload_fixtures import mm_train_data from tests.utils.tq._payload_assertions import diff_digests, leaf_digests num_samples = 4 - train_data, source = mm_train_data(num_samples) + train_data = _multimodal_payload(num_samples) + source = "synthetic" record_property("multimodal_payload_source", source) train_data = dict(train_data) train_data["sample_id"] = list(range(num_samples)) diff --git a/tests/utils/test_tq_failure_paths.py b/tests/utils/test_tq_failure_paths.py index 1c09a2e65..0732d3d66 100644 --- a/tests/utils/test_tq_failure_paths.py +++ b/tests/utils/test_tq_failure_paths.py @@ -8,33 +8,24 @@ * timeout -- controller ``get_config`` timeout and attach-handshake 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), 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 +* retry/disconnect -- a transient get recovers and a dead peer fails loudly * 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. +Real Mooncake TCP/RDMA byte-exact and wire-level checks live in the single +cross-node C0/C1/C2 benchmark. This module keeps CI-safe failure semantics and +uses real TransferQueue classes only where a stubbed store is sufficient. """ from __future__ import annotations import asyncio import importlib.util -import inspect -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 @@ -60,192 +51,10 @@ 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" -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -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 - - -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() - - -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. - """ - client = None - keys: list[str] = [] - put_meta: list[dict | None] | None = None - status = "error" - detail = "roundtrip did not run" - try: - from tests.utils.mm_payload_fixtures import mm_train_data - from tests.utils.tq._payload_assertions import diff_digests, leaf_digests - - train_data, source = mm_train_data(4) - samples = train_data["multimodal_train_inputs"] - client = TestMooncakeByteExact._client(protocol) - 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", - ) - 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, 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() +def _raise(error: BaseException) -> None: + raise error # --------------------------------------------------------------------------- @@ -294,15 +103,14 @@ def test_half_initialised_controller_is_reaped(self, monkeypatch): 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.""" - error = tq_lifecycle.ray.exceptions.GetTimeoutError("private timeout detail") - killed = self._fake_ray(monkeypatch, actor=MagicMock(), get_raises=error) - assert tq_lifecycle.reap_unusable_tq_controller() is True - assert killed == ["killed"] - - def test_dead_actor_is_reaped(self, monkeypatch): - error = tq_lifecycle.ray.exceptions.RayActorError(error_msg="private actor detail") + @pytest.mark.parametrize("error_kind", ["timeout", "dead-actor"]) + def test_unusable_controller_is_reaped(self, monkeypatch, error_kind): + """An unresponsive/dead controller must not turn tq.init into a + hang.""" + if error_kind == "timeout": + error = tq_lifecycle.ray.exceptions.GetTimeoutError("private timeout detail") + else: + error = tq_lifecycle.ray.exceptions.RayActorError(error_msg="private actor detail") killed = self._fake_ray(monkeypatch, actor=MagicMock(), get_raises=error) assert tq_lifecycle.reap_unusable_tq_controller() is True assert killed == ["killed"] @@ -323,7 +131,7 @@ def test_stored_config_ray_error_is_sanitized(self, monkeypatch): monkeypatch.setattr( tq_lifecycle.ray, "get_actor", - lambda *_args, **_kwargs: (_ for _ in ()).throw(tq_lifecycle.ray.exceptions.RayError(private_detail)), + lambda *_args, **_kwargs: _raise(tq_lifecycle.ray.exceptions.RayError(private_detail)), ) with pytest.raises(RuntimeError) as excinfo: @@ -379,13 +187,8 @@ def test_uninitialised_tq_does_not_raise(self, monkeypatch): 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() - - @pytest.mark.parametrize("value", ["nan", "inf", "-inf"]) - def test_attach_timeout_env_rejects_non_finite_values(self, monkeypatch, value): + @pytest.mark.parametrize("value", ["soon", "nan", "inf", "-inf"]) + def test_attach_timeout_env_rejects_unusable_values(self, monkeypatch, value): monkeypatch.setenv("RELAX_TQ_ATTACH_TIMEOUT_SECONDS", value) with pytest.raises(RuntimeError, match="finite positive") as excinfo: tq_lifecycle._resolve_attach_timeout() @@ -422,28 +225,38 @@ 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): + def test_cluster_attach_covers_all_nodes_with_hard_affinity_and_one_shot_workers(self, monkeypatch): remote_options = {} + scheduling_strategies = [] + submitted_refs = [] class _Task: - def options(self, **_kwargs): + def options(self, **options): + scheduling_strategies.append(options["scheduling_strategy"]) return self def remote(self, *_args): - return object() + ref = object() + submitted_refs.append(ref) + return ref def record_remote_options(**options): remote_options.update(options) return lambda _function: _Task() monkeypatch.setattr(tq_lifecycle.ray, "remote", record_remote_options) - monkeypatch.setattr(tq_lifecycle, "_alive_node_ids", lambda: ["a" * 56]) + node_ids = ["a" * 56, "b" * 56] + monkeypatch.setattr(tq_lifecycle, "_alive_node_ids", lambda: node_ids) monkeypatch.setattr(tq_lifecycle.ray, "wait", lambda refs, **_kwargs: (list(refs), [])) monkeypatch.setattr(tq_lifecycle.ray, "get", lambda _ref: None) assert tq_lifecycle.verify_cluster_attach({}, timeout=0.1) == [] assert remote_options["max_calls"] == 1 assert remote_options["max_retries"] == 0 + assert len(submitted_refs) == len(node_ids) + assert [(strategy.node_id, strategy.soft) for strategy in scheduling_strategies] == [ + (node_id, False) for node_id in node_ids + ] def test_cluster_attach_with_no_alive_nodes_fails_closed_without_scheduling(self, monkeypatch): monkeypatch.setattr(tq_lifecycle, "_alive_node_ids", lambda: []) @@ -536,31 +349,34 @@ def fake_assert(): monkeypatch.setattr(tq_lifecycle, "detach_tq_client", lambda: events.append("detach")) return worker[0], events - def test_handshake_worker_attaches_asserts_then_detaches(self, monkeypatch): - handshake, events = self._run_worker(monkeypatch) - handshake({"backend": {"storage_backend": "MooncakeStore"}}, True, 0.1) - assert events == ["attach:attach-handshake:0.1", "assert", "detach"] - - def test_handshake_worker_detaches_even_when_the_assertion_fails(self, monkeypatch): - """A rejected transport must not leave this node's segment registered. - - Without the ``finally`` the segment would linger until the master's - ``client_ttl`` expires and break fast restarts. - """ - handshake, events = self._run_worker(monkeypatch, assert_error=RuntimeError("protocol=tcp")) - with pytest.raises(RuntimeError, match="protocol=tcp"): - handshake({"backend": {"storage_backend": "MooncakeStore"}}, True, 0.1) - assert events == ["attach:attach-handshake:0.1", "assert", "detach"] - - def test_handshake_worker_skips_the_assertion_for_simple_storage(self, monkeypatch): - handshake, events = self._run_worker(monkeypatch) - handshake({"backend": {"storage_backend": "SimpleStorage"}}, False, 0.1) - assert events == ["attach:attach-handshake:0.1", "detach"] + @pytest.mark.parametrize( + ("backend", "verify_mooncake", "assert_error"), + [ + ("MooncakeStore", True, None), + ("MooncakeStore", True, RuntimeError("protocol=tcp")), + ("SimpleStorage", False, None), + ], + ids=["mooncake", "mooncake-rejected", "simple"], + ) + def test_handshake_worker_always_detaches(self, monkeypatch, backend, verify_mooncake, assert_error): + """A rejected transport must not leave this node's segment + registered.""" + handshake, events = self._run_worker(monkeypatch, assert_error=assert_error) + conf = {"backend": {"storage_backend": backend}} + if assert_error is None: + handshake(conf, verify_mooncake, 0.1) + else: + with pytest.raises(RuntimeError, match="protocol=tcp"): + handshake(conf, verify_mooncake, 0.1) + expected = ["attach:attach-handshake:0.1"] + if verify_mooncake: + expected.append("assert") + assert events == [*expected, "detach"] def test_ready_task_failure_is_sanitized(self, monkeypatch): self._capture_handshake(monkeypatch) secret = "worker endpoint and traceback path must stay private" - monkeypatch.setattr(tq_lifecycle.ray, "get", lambda _ref: (_ for _ in ()).throw(RuntimeError(secret))) + monkeypatch.setattr(tq_lifecycle.ray, "get", lambda _ref: _raise(RuntimeError(secret))) failures = tq_lifecycle.verify_cluster_attach({}, timeout=0.1) @@ -637,7 +453,7 @@ def test_cancel_failure_aborts_fallback_boundary_without_leaking_detail(self, mo monkeypatch.setattr( tq_lifecycle.ray, "cancel", - lambda _ref, force: (_ for _ in ()).throw(RuntimeError(private_detail)), + lambda _ref, force: _raise(RuntimeError(private_detail)), ) with pytest.raises(tq_lifecycle.TqHandshakeIsolationError) as excinfo: @@ -721,46 +537,19 @@ def test_detach_delegates_to_local_close(self, monkeypatch): 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] - assert component.data_system_client is None - - def test_component_del_without_client_is_noop(self, monkeypatch): + @pytest.mark.parametrize("has_client", [True, False], ids=["attached", "not-attached"]) + def test_component_del_detaches_only_an_attached_client(self, monkeypatch, has_client): from relax.components.base import Base calls = [] monkeypatch.setattr(tq_lifecycle, "detach_tq_client", lambda: calls.append(True)) component = Base() + if has_client: + component.data_system_client = object() component.__del__() - assert calls == [] - - -class TestExclusiveLifecycleSurface: - """Phase 4 must not grow the removed shared-cluster machinery back.""" - - def test_token_generation_and_foreign_owner_symbols_are_absent(self): - removed = ( - "OWNER_TOKEN_FIELD", - "TqConfigurationMismatch", - "TqControllerInspectionError", - "TqControllerMissingConfig", - "_TQ_CLIENT_GENERATION", - "_CURRENT_TQ_CLIENT_GENERATION", - ) - for name in removed: - assert not hasattr(tq_lifecycle, name) - - def test_attach_and_teardown_signatures_express_single_process_owner(self): - assert "lease_owner" not in inspect.signature(tq_lifecycle.attach_tq_client).parameters - assert list(inspect.signature(tq_lifecycle.detach_tq_client).parameters) == [] - assert list(inspect.signature(tq_lifecycle.close_tq_and_unmount).parameters) == [] + assert calls == ([True] if has_client else []) + if has_client: + assert component.data_system_client is None # --------------------------------------------------------------------------- @@ -835,16 +624,24 @@ def test_healthy_existing_controller_fails_exclusive_without_starting_owner(self assert events == ["reap"] - def test_auto_cleans_failed_mooncake_then_retries_simple_once(self, monkeypatch): + @pytest.mark.parametrize( + ("primary_error", "fallback_reason"), + [ + ( + tq_lifecycle.TqInitializationError("master unavailable"), + "mooncake_init_failed:TqInitializationError", + ), + (tq_lifecycle.TqInitializationTimeout("timed out"), "mooncake_init_failed:TqInitializationTimeout"), + ], + ids=["init-error", "init-timeout"], + ) + def test_auto_cleans_failed_mooncake_then_retries_simple_once(self, monkeypatch, primary_error, fallback_reason): primary = self._conf("MooncakeStore") fallback = self._conf("SimpleStorage") - events = self._patch_transaction( - monkeypatch, - init_effects=[tq_lifecycle.TqInitializationError("master unavailable"), None], - ) + events = self._patch_transaction(monkeypatch, init_effects=[primary_error, None]) 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:TqInitializationError" + assert result.fallback_reason == fallback_reason assert events == [ "reap", "start", @@ -865,17 +662,6 @@ def test_required_cleans_failed_init_without_fallback(self, monkeypatch): tq_lifecycle.initialize_tq_with_fallback(primary, mode="required", fallback_conf=fallback) assert events == ["reap", "start", "init:MooncakeStore"] - def test_timeout_auto_retries_only_after_isolated_owner_cleanup(self, monkeypatch): - primary = self._conf("MooncakeStore") - fallback = self._conf("SimpleStorage") - events = self._patch_transaction( - monkeypatch, - init_effects=[tq_lifecycle.TqInitializationTimeout("timed out"), None], - ) - result = tq_lifecycle.initialize_tq_with_fallback(primary, mode="auto", fallback_conf=fallback) - assert result.config["backend"]["storage_backend"] == "SimpleStorage" - assert events[-3:] == ["reap", "start", "init:SimpleStorage"] - def test_cleanup_failure_aborts_auto_before_second_gate(self, monkeypatch): events = self._patch_transaction( monkeypatch, @@ -975,37 +761,30 @@ def test_initialize_success_returns_the_exclusive_owner(self, monkeypatch): assert result.config is stored assert result.owner is owner - def test_initialize_timeout_cleans_owner_before_candidate_failure(self, monkeypatch): - owner = _FakeOwner() - cleaned: list[object] = [] - monkeypatch.setattr( - tq_lifecycle.ray, - "get", - lambda _ref, timeout: (_ for _ in ()).throw(tq_lifecycle.ray.exceptions.GetTimeoutError("timeout")), - ) - monkeypatch.setattr(tq_lifecycle, "_cleanup_failed_owner", lambda handle: cleaned.append(handle)) - - with pytest.raises(tq_lifecycle.TqInitializationTimeout): - tq_lifecycle._initialize_owner(owner, {}, timeout=0.1) - - assert cleaned == [owner] - - def test_initialize_error_is_sanitized_only_after_cleanup(self, monkeypatch): + @pytest.mark.parametrize("error_kind", ["timeout", "runtime"]) + def test_initialize_failure_cleans_owner_before_propagating(self, monkeypatch, error_kind): owner = _FakeOwner() cleaned: list[object] = [] private_detail = "private endpoint and traceback path" + if error_kind == "timeout": + source_error = tq_lifecycle.ray.exceptions.GetTimeoutError("timeout") + expected_error = tq_lifecycle.TqInitializationTimeout + else: + source_error = RuntimeError(private_detail) + expected_error = tq_lifecycle.TqInitializationError monkeypatch.setattr( tq_lifecycle.ray, "get", - lambda _ref, timeout: (_ for _ in ()).throw(RuntimeError(private_detail)), + lambda _ref, timeout: _raise(source_error), ) monkeypatch.setattr(tq_lifecycle, "_cleanup_failed_owner", lambda handle: cleaned.append(handle)) - with pytest.raises(tq_lifecycle.TqInitializationError) as excinfo: + with pytest.raises(expected_error) as excinfo: tq_lifecycle._initialize_owner(owner, {}, timeout=0.1) assert cleaned == [owner] - assert private_detail not in str(excinfo.value) + if error_kind == "runtime": + assert private_detail not in str(excinfo.value) def test_stop_owner_requires_terminal_probe_result(self, monkeypatch): owner = _FakeOwner() @@ -1015,28 +794,29 @@ def test_stop_owner_requires_terminal_probe_result(self, monkeypatch): monkeypatch.setattr( tq_lifecycle.ray, "get", - lambda _ref: (_ for _ in ()).throw(tq_lifecycle.ray.exceptions.RayActorError(error_msg="dead")), + lambda _ref: _raise(tq_lifecycle.ray.exceptions.RayActorError(error_msg="dead")), ) tq_lifecycle._stop_owner_actor(owner, timeout=0.1) assert killed == [(owner, True)] - def test_stop_owner_fails_closed_when_probe_remains_pending(self, monkeypatch): + @pytest.mark.parametrize( + ("probe_ready", "match"), + [(False, "remained pending"), (True, "returned normally")], + ids=["pending", "returned"], + ) + def test_stop_owner_requires_a_terminal_probe_failure(self, monkeypatch, probe_ready, match): owner = _FakeOwner() monkeypatch.setattr(tq_lifecycle.ray, "kill", lambda *_args, **_kwargs: None) - monkeypatch.setattr(tq_lifecycle.ray, "wait", lambda refs, **_kwargs: ([], list(refs))) - - with pytest.raises(tq_lifecycle.TqCleanupTimeout, match="remained pending"): - tq_lifecycle._stop_owner_actor(owner, timeout=0.1) - - def test_stop_owner_rejects_a_probe_that_returns_normally(self, monkeypatch): - owner = _FakeOwner() - monkeypatch.setattr(tq_lifecycle.ray, "kill", lambda *_args, **_kwargs: None) - monkeypatch.setattr(tq_lifecycle.ray, "wait", lambda refs, **_kwargs: (list(refs), [])) + monkeypatch.setattr( + tq_lifecycle.ray, + "wait", + lambda refs, **_kwargs: (list(refs), []) if probe_ready else ([], list(refs)), + ) monkeypatch.setattr(tq_lifecycle.ray, "get", lambda _ref: None) - with pytest.raises(tq_lifecycle.TqCleanupTimeout, match="returned normally"): + with pytest.raises(tq_lifecycle.TqCleanupTimeout, match=match): tq_lifecycle._stop_owner_actor(owner, timeout=0.1) def test_failed_owner_cleanup_waits_for_owner_then_reaps_controller(self, monkeypatch): @@ -1065,7 +845,7 @@ def test_unconfirmed_owner_aborts_before_controller_cleanup(self, monkeypatch): monkeypatch.setattr( tq_lifecycle, "_stop_owner_actor", - lambda *_args, **_kwargs: (_ for _ in ()).throw(tq_lifecycle.TqCleanupTimeout("owner pending")), + lambda *_args, **_kwargs: _raise(tq_lifecycle.TqCleanupTimeout("owner pending")), ) monkeypatch.setattr( tq_lifecycle, @@ -1085,7 +865,7 @@ def test_owner_close_failure_still_reaps_global_controller(self, monkeypatch): monkeypatch.setattr( tq_lifecycle.ray, "get", - lambda ref, timeout: (_ for _ in ()).throw(RuntimeError("close failed")), + lambda ref, timeout: _raise(RuntimeError("close failed")), ) monkeypatch.setattr(tq_lifecycle, "_stop_owner_actor", lambda handle, **_kwargs: stopped.append(handle)) monkeypatch.setattr(tq_lifecycle, "kill_tq_controller_and_wait", lambda **_kwargs: killed.append(True)) @@ -1102,15 +882,13 @@ def test_owner_close_failure_still_reaps_global_controller(self, monkeypatch): class _FlakyStore: - """Stub mooncake store: the first ``fail_times`` calls return error - codes.""" + """Stub store whose first ``fail_times`` reads return an error code.""" 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: @@ -1124,10 +902,6 @@ 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. @@ -1165,34 +939,15 @@ def run_in_executor(self, executor, fn, *args): 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.""" + """Behavior surface retained in Relax; retry internals belong upstream.""" - def test_get_retries_then_succeeds(self, monkeypatch): + def test_transient_get_failure_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]) + keys = ["0@f0", "1@f0"] + client._batch_get_into_with_retry(keys, [1, 2], [8, 8]) + assert store.get_calls == [keys, keys, keys] def test_disconnect_surfaces_instead_of_returning_garbage(self): """A dead peer must raise, never hand back a silently short buffer.""" @@ -1259,101 +1014,3 @@ async def notify(*args, **kwargs): 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 - - -# --------------------------------------------------------------------------- -# 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 - - from relax.utils.tq.correctness import ensure_mooncake_correctness_guards - - ensure_mooncake_correctness_guards() - 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): - """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, record_property): - """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. - """ - status, detail = _run_isolated_roundtrip(_mm_slow_path_worker, protocol, timeout=240) - record_property("multimodal_roundtrip", detail) - assert status == "ok", f"{status}: {detail}" diff --git a/tests/utils/tq/test_config.py b/tests/utils/tq/test_config.py index 1d3e353d1..f918218a3 100644 --- a/tests/utils/tq/test_config.py +++ b/tests/utils/tq/test_config.py @@ -157,50 +157,32 @@ def test_accepts_hostname_and_bracketed_ipv6(self): class TestBackendConfigDicts: """The dicts handed to ``tq.init``.""" - def test_simple_storage_config(self): - cfg = build_simple_storage_config(total_storage_size=1000, num_data_storage_units=2) + @pytest.mark.parametrize("total_storage_size", [1000, None], ids=["bounded", "unlimited"]) + def test_simple_storage_config(self, total_storage_size): + cfg = build_simple_storage_config(total_storage_size=total_storage_size, num_data_storage_units=2) assert cfg == { "storage_backend": "SimpleStorage", - "SimpleStorage": {"total_storage_size": 1000, "num_data_storage_units": 2}, + "SimpleStorage": {"total_storage_size": total_storage_size, "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.""" - assert build_mooncake_config(master_address=_MASTER)["storage_backend"] == "MooncakeStore" - assert ( - build_simple_storage_config(total_storage_size=1, num_data_storage_units=1)["storage_backend"] - == "SimpleStorage" - ) - - def test_mooncake_defaults_to_host_rdma(self): - """Production never names a protocol: the default is the only one it - ships.""" - mc = build_mooncake_config(master_address=_MASTER, device="rdma0")["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 + @pytest.mark.parametrize( + ("kwargs", "expected_protocol", "expected_device"), + [ + ({}, "rdma", ""), + ({"device": "rdma0"}, "rdma", "rdma0"), + ({"protocol": "tcp"}, "tcp", ""), + ], + ids=["production-default", "explicit-device", "benchmark-tcp"], + ) + def test_mooncake_config_contract(self, kwargs, expected_protocol, expected_device): + cfg = build_mooncake_config(master_address=_MASTER, **kwargs) + assert cfg["storage_backend"] == "MooncakeStore" + mc = cfg["MooncakeStore"] + assert mc["protocol"] == expected_protocol + assert mc["device_name"] == expected_device + assert mc["hard_pin"] is True + assert mc["auto_init"] is False assert mc["master_server_address"] == _MASTER - - def test_empty_device_is_left_to_mooncake(self): - mc = build_mooncake_config(master_address=_MASTER)["MooncakeStore"] - assert mc["device_name"] == "" - - def test_tcp_is_reachable_only_by_explicit_request(self): - """Mooncake/TCP survives as benchmark C1, never as a production - default.""" - mc = build_mooncake_config(master_address=_MASTER, protocol="tcp")["MooncakeStore"] - assert mc["protocol"] == "tcp" - - def test_mooncake_config_pins_host_rdma(self): - """This phase ships host RDMA only: ``use_gdr`` is always False and no - GDR staging buffer is configured.""" - mc = build_mooncake_config(master_address=_MASTER)["MooncakeStore"] assert mc["use_gdr"] is False assert "gdr_staging_buffer_mb" not in mc @@ -228,41 +210,30 @@ class TestCorrectnessContract: 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() - - @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) + @pytest.mark.parametrize("override", [None, "0"], ids=["default", "explicit-disable"]) + def test_contract_accepts_safe_memcpy_settings(self, monkeypatch, override): + if override is None: + monkeypatch.delenv("MC_STORE_MEMCPY", raising=False) + else: + monkeypatch.setenv("MC_STORE_MEMCPY", override) validate_mooncake_runtime_contract() assert os.environ["MC_STORE_MEMCPY"] == "0" - def test_contract_rejects_explicit_memcpy_enable(self, monkeypatch): + @pytest.mark.parametrize( + ("override", "private_marker"), + [("1", None), ("1\nprivate deployment detail", "private deployment detail")], + ids=["enable", "untrusted-value"], + ) + def test_contract_rejects_unsafe_memcpy_settings(self, monkeypatch, override, private_marker): # 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") - with pytest.raises(RuntimeError, match="MC_STORE_MEMCPY"): - validate_mooncake_runtime_contract() - - def test_contract_rejection_does_not_echo_memcpy_override(self, monkeypatch): - private_detail = "1\nprivate deployment detail" - monkeypatch.setenv("MC_STORE_MEMCPY", private_detail) + monkeypatch.setenv("MC_STORE_MEMCPY", override) with pytest.raises(RuntimeError, match="MC_STORE_MEMCPY") as excinfo: validate_mooncake_runtime_contract() - assert private_detail not in str(excinfo.value) - - 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" + if private_marker is not None: + assert private_marker not in str(excinfo.value) @pytest.mark.skipif( not _REAL_TQ_STORAGE, @@ -294,22 +265,45 @@ class TestSegmentCapacity: def test_text_only_passes(self): assert validate_segment_capacity(_make_args(multimodal_keys=None)) is None - def test_multimodal_large_batch_fails(self): - args = _make_args( - multimodal_keys=["pixel_values"], rollout_batch_size=256, n_samples_per_prompt=8, max_staleness=1 - ) + @pytest.mark.parametrize( + ("overrides", "message"), + [ + ( + { + "multimodal_keys": ["pixel_values"], + "rollout_batch_size": 256, + "n_samples_per_prompt": 8, + "max_staleness": 1, + }, + "insufficient", + ), + ( + { + "multimodal_keys": ["pixel_values"], + "rollout_batch_size": 32, + "n_samples_per_prompt": 1, + "max_staleness": 1, + }, + "RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB", + ), + ( + { + "multimodal_keys": ["pixel_values"], + "rollout_batch_size": 16, + "partial_rollout": True, + "use_dynamic_global_batch_size": True, + "over_sampling_batch_size": 64, + }, + "effective_batch=64", + ), + ], + ids=["large-batch", "staleness", "dynamic-oversampling"], + ) + def test_insufficient_capacity_is_rejected(self, overrides, message): + args = _make_args(**overrides) err = validate_segment_capacity(args) assert err is not None - assert "insufficient" in err.lower() - - def test_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 - ) - err = validate_segment_capacity(args) - assert err is not None and "RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB" in err + assert message.lower() in err.lower() def test_env_override_raises_the_ceiling(self, monkeypatch): args = _make_args( @@ -355,33 +349,18 @@ def test_requires_seq_length(self): with pytest.raises(RuntimeError, match="seq_length"): estimate_payload_bytes(_make_args(seq_length=None)) - 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): + @pytest.mark.parametrize( + ("partial_rollout", "expected_batch"), + [(True, 64), (False, 16)], + ids=["dynamic-partial", "nominal"], + ) + def test_capacity_batch_resolution(self, partial_rollout, expected_batch): args = _make_args( - multimodal_keys=["pixel_values"], rollout_batch_size=16, - partial_rollout=True, + partial_rollout=partial_rollout, use_dynamic_global_batch_size=True, over_sampling_batch_size=64, ) - err = validate_segment_capacity(args) - assert err is not None - assert "effective_batch=64" in err + assert resolve_tq_capacity_batch_size(args) == expected_batch + if partial_rollout: + assert estimate_payload_bytes(args) == expected_batch * 8192 * 32 From 14e73f2584b14bbaf22c53f39c2b8c0d9900da40 Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Mon, 24 Aug 2026 04:47:37 +0800 Subject: [PATCH 26/30] docs(tq): update host-RDMA documentation --- docs/draft/transfer_queue_rdma.md | 250 ++++++++++-------------------- 1 file changed, 85 insertions(+), 165 deletions(-) diff --git a/docs/draft/transfer_queue_rdma.md b/docs/draft/transfer_queue_rdma.md index 1afe9825a..20ca1e6f3 100644 --- a/docs/draft/transfer_queue_rdma.md +++ b/docs/draft/transfer_queue_rdma.md @@ -1,163 +1,90 @@ -# TransferQueue RDMA 数据面使用与运维指南 +# TransferQueue host-RDMA 使用与运维 -## 概述 +## 范围与配置 -Relax 的数据面（rollout ↔ train 之间的样本传输）默认走 TransferQueue 的 SimpleStorage/ZMQ。本特性把 TransferQueue 已有的 MooncakeStore 后端接出来，使数据面可以走 RDMA，并在能力不足时安全回退。 +Relax 默认使用 TransferQueue SimpleStorage。首期 RDMA 支持只接入 MooncakeStore/host-RDMA，不改变 payload 形状或数据分发语义，也不支持 GDR。生产路径只有 host-RDMA 和 SimpleStorage；Mooncake/TCP 仅作为跨节点 benchmark 的 C1 对照。 -首期只做配置接入、真实 attach 能力判定与一致回退，**不改变 payload 形状与数据分发语义**。默认参数仍使用 SimpleStorage 及原有 controller 所有权模型；同时所有 worker attach（包括 SimpleStorage）新增默认 60 秒 deadline，半初始化 controller 会被回收，Controller 构造失败时会关闭本进程已经完成的 legacy `tq.init`。 +| 参数 | 语义 | +|---|---| +| `--tq-rdma-mode=off` | 默认值；保持原有 SimpleStorage 路径，不执行 Mooncake 检查或 attach | +| `--tq-rdma-mode=auto` | 尝试 host-RDMA；任一检查或节点 attach 失败时，完成清理后统一回退 SimpleStorage | +| `--tq-rdma-mode=required` | 要求 host-RDMA；不可用时完成清理并终止启动，不允许降级 | +| `--tq-rdma-device=` | 指定 RDMA 设备；空值由 Mooncake 原生逻辑选择，多 HCA 环境建议显式指定 | -> 与 RFC #217 的差异：RFC 要求"在首次 `tq.init` 之前完成静态节点探测"（F10）来避免挂死。按维护者的瘦身要求，静态探测已删除，防挂死改由这些机制保证：`tq.init` 在隔离的 owner actor 中执行并带超时、每次初始化前回收半初始化的 controller、拆除时等待 controller 真正从 GCS 注销、handshake 使用一次性 Ray worker（`max_calls=1`/`max_retries=0`）使超时的 `tq.init` 线程随进程退出。 +Relax 固定 Mooncake `use_gdr=false`。GDR 若有实际需求，应由独立 PR 实现并做专项验证。 -## 配置入口 +## 部署前提 -只暴露两个表达使用意图的参数，Mooncake 底层参数（endpoint、buffer、segment、timeout、master 策略）不做 CLI，走内部默认与部署环境。 +首期只支持**单任务独占 Ray 集群**：同一个 Ray cluster 在初始化和运行期间只能有一个 Relax job，不支持 concurrent initializer、多 job admission 或复用其他作业的 TransferQueue controller。 -| 参数 | 取值 | 说明 | -|---|---|---| -| `--tq-rdma-mode` | `off`（默认）/ `auto` / `required` | `off` 使用原有 SimpleStorage，保留接入前的存储与 controller 所有权语义；`auto` 尝试 host RDMA，不可用则回退 SimpleStorage；`required` 不可用则直接报错退出 | -| `--tq-rdma-device` | 设备名，如 `mlx5_bond_0`；空为自动 | 多网卡机器上自动选择可能选错，跨节点时建议显式指定 | - -生效路径只有两种：**MooncakeStore/host-RDMA** 和 **SimpleStorage**。Mooncake/TCP 不是公开的生产配置，只作为 benchmark 的 C1 对照存在，因此 `auto` 不经过 TCP 中间档，直接回退 SimpleStorage。 - -`--tq-rdma-mode=required` 覆盖的是**传输层**：MooncakeStore + RDMA 可用性与 segment 容量。 - -### 首期只交付 host RDMA - -数据面 payload 先经过主机内存再通过 RDMA 跨节点传输。GDR（GPU Direct RDMA）不在本期范围内：它不是任务书要求，可用性也无法由 driver 的启动探测代表（探测进程没有初始化 CUDA context），因此 Relax 把 Mooncake 的 `use_gdr` 固定为 `false`，不提供开关。如后续出现真实需求，GDR 应由独立 PR 实现并单独验证。 - -## 启动流程与降级 - -driver 在**第一次 `tq.init` 之前**完成所有配置校验，并生成 job 级唯一的 backend config，其余组件（actor / critic / rollout / sft / advantages / actor_fwd）都读同一份，不各自决策。`off` 直接短路到 SimpleStorage，不做任何 Mooncake/RDMA 检查。 - -Relax **不做静态硬件探测**：不扫描 `/sys/class/infiniband`，不读 GID 表，不推断 `memlock`。这类启发式覆盖不到真实调度位置（TQ client 跑在 Ray Serve replica 与 0-CPU actor 里，没有 placement 绑定），而且容易与 Mooncake 底层实现漂移。运行时可用性以真实 attach/setup 的结果为准；线路是否确实经过 RDMA 仍由真机 wire-proof 验收确认。 - -1. 校验 `--tq-rdma-mode` 取值（非法值始终报错，不回退） -2. 校验 TransferQueue/Mooncake 正确性契约 -3. 校验 `MC_MASTER_ADDRESS` 存在且是可用的 `host:port`（只查格式，不做 DNS 解析和连通性探测） -4. segment 容量预检：按 token 预算推导最坏情况 payload，对比 client segment 大小 -5. driver 在独立 owner actor 中执行 `tq.init`，随后在**每个存活节点**调度一次性 handshake worker：真实 attach/setup → 确认 storage manager 是 `MooncakeStorageManager` 且 client 配置请求 `protocol=rdma` → detach -6. 汇总各节点结果：零存活节点、节点发现/任务调度失败或任一节点 attach 失败都按失败处理；`auto` 只有在确认所有一次性 worker 已终止后，才会**先完整拆除 Mooncake 状态**（owner close + controller 退出等待 + segment unmount），再收敛到 SimpleStorage；`required` 下启动失败并列出失败摘要 - -第 5 步不是轻量探针：每个节点都会创建真实 Mooncake client，并按配置请求挂载/注册完整 client segment（默认 `global_segment_size=4 GiB`，另有默认 1 GiB local buffer），完成后立即释放。具体物理 RSS、锁页与注册方式取决于 Mooncake 实现，但启动阶段会出现节点级瞬时内存/注册资源尖峰；CPU-only head 也在覆盖范围内。节点内存或 `memlock` 不足会表现为 attach 握手失败：`auto` 下整个作业统一回退 SimpleStorage，`required` 下启动失败。调大 `RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB` 时必须把这份每节点启动资源足迹一并纳入容量规划。 +Mooncake master 由部署环境管理，Relax 不启动、重启或停止它。所有节点必须设置同一个外部 endpoint： -回退只有一档： - -``` -MooncakeStore/host-RDMA → SimpleStorage -（正确性契约不满足、master 未配置或格式非法、segment 容量预检不足， - 或全节点真实 attach/setup 未通过） +```bash +export MC_MASTER_ADDRESS=master.example:50051 ``` -### 启动耗时 +Relax 只检查 `host:port` 格式；DNS、路由、防火墙和 master 健康状态由真实初始化与 attach 验证。 -`auto` 在 RDMA 不可用的集群上会比 `off` 慢，因为它要真的走一遍完整流程才能得出结论：owner `tq.init`（上限 60 s）→ 全节点 handshake（每节点 attach 上限 60 s，driver 侧等待上限 90 s）→ 拆除（owner close 上限 30 s，外加等待 controller 从 GCS 注销）→ SimpleStorage 初始化。最坏情况下启动阶段可能达到**分钟级**。这是用启动时延换取"判定基于真实 endpoint"的取舍；确定不需要 RDMA 的作业应显式用 `--tq-rdma-mode=off`，它不触发上述任何步骤。 +## 启动、回退与清理 -## 启动日志怎么读 +`off` 直接初始化 SimpleStorage。`auto` 和 `required` 按以下顺序执行： -正常启动会打这几段，排障时先看它们： - -``` -[dataplane] requested: rdma_mode=auto device=mlx5_bond_0 -[dataplane] backend=MooncakeStore protocol=rdma device=mlx5_bond_0 (pending handshake) -[dataplane] Mooncake attach handshake passed on all alive nodes. -``` +1. 校验 mode、TransferQueue/Mooncake correctness contract、master 格式和 segment 容量。 +2. 在独立 owner actor 中有界执行 Mooncake `tq.init`。 +3. 在每个 ALIVE Ray 节点运行一次性 worker（`max_calls=1`、`max_retries=0`），真实 attach 并确认 `MooncakeStorageManager` 和 `protocol=rdma`，随后 detach。 +4. 全部节点通过后启用 host-RDMA；任一失败则等待 worker 终止，并清理 owner、controller 和 segment。 +5. 清理可确认时，`auto` 初始化 SimpleStorage，`required` 报错；清理无法确认时两种模式都 fail closed。 -`(pending handshake)` 表示配置校验已通过但能力尚未确认；只有最后一行出现才说明每个存活节点都完成了真实 attach/setup，storage manager 是 MooncakeStore，且 Mooncake client 的配置请求为 `protocol=rdma`。 +Relax 不再维护 `/sys` 启发式能力探测。真实 attach 是运行时能力判据，但它只证明 manager、配置 protocol 和 setup 成功；线路是否真正传输 RDMA 数据必须由 benchmark counter 的 wire proof 证明。 -注意这条断言的强度：它核对的是 client 的**配置意图和 setup 结果**，不是 negotiated transport，也不能单独证明数据包没有经过 TCP。线路级证明由 benchmark 的强制 counter gate 提供：RDMA 档要求 IB receive counter 增长且大于 TCP counter，SimpleStorage/TCP 档要求 TCP receive counter 增长且不小于 IB counter。 +owner 初始化和 worker attach 默认各有 60 秒边界；worker attach 可通过 `RELAX_TQ_ATTACH_TIMEOUT_SECONDS` 调整。RDMA 不可用时，`auto` 需要完成真实初始化、全节点 attach 和清理，启动可能达到分钟级。已知无需 RDMA 的作业应使用 `off`。 -发生回退时会看到： +健康的既有 controller 不会被接管或关闭，启动会直接失败；半初始化、超时或已死亡的 controller 才允许回收。全局 `tq.close()` 只由 owner 调用，普通 worker 只 detach 本地 client，Mooncake master 始终不由 Relax 管理。 -``` -[dataplane] <失败原因>; auto fallback to SimpleStorage: <错误详情> -``` +## Segment 与容量 -或 handshake 阶段失败： +真实 handshake 会在每个 ALIVE 节点（包括 CPU-only head）创建 Mooncake client，瞬时挂载并注册完整 client segment，结束后立即 detach。默认配置为每 client 4 GiB global segment 和 1 GiB local buffer；实际 RSS、锁页内存及注册资源由 Mooncake 实现决定。内存或 `memlock` 不足会表现为 attach 失败。 -``` -[dataplane] Mooncake attach handshake reported N failure(s) (<稳定的失败类型摘要>); closing Mooncake state and converging the whole job to SimpleStorage. -``` - -`required` 模式下这两种情况都是启动失败并抛出同样的原因，不会出现回退日志。 - -## Mooncake master 生命周期 - -`auto_init` 固定为 `false`：**Relax 既不启动也不停止 master**，master 由部署环境管理。这样做是因为 TQ 的 `auto_init=true` 路径会执行 `pkill -f "[m]ooncake_master"`，在共享集群上会杀掉其他人的进程。 - -启动 master（部署侧，节点上执行一次）： +global segment 可按部署容量调整： ```bash -setsid mooncake_master -rpc_port=50051 -metrics_port=9004 > /var/log/mooncake_master.log 2>&1 < /dev/null & +export RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB=8 ``` -然后给**每个节点**的作业环境设置 `MC_MASTER_ADDRESS=:50051`。Relax 不会假定 loopback 端点（多节点作业里每个节点都把自己的 localhost 当 master 会导致误降级或误中止），因此未设置或格式非法时：`auto` 记 WARNING 并回退 SimpleStorage，`required` 启动失败。 - -启动前置条件：部署侧必须先启动 master，所有存活 Ray 节点和 driver 都能解析并连接 `MC_MASTER_ADDRESS`，防火墙允许 master RPC 端口；作业镜像中的 TQ 必须包含本文“正确性依赖”所列修复。Relax 不负责拉起、重启或终止 master。 - -三种情形下的行为： - -| 情形 | 表现 | 处理 | -|---|---|---| -| **`MC_MASTER_ADDRESS` 未配置或格式非法** | `auto` 记 WARNING 后回退 SimpleStorage（不初始化任何 Mooncake 状态）；`required` 启动失败 | 给每个节点的作业环境设置 `MC_MASTER_ADDRESS=:` | -| **master 不可达** | 没有单独的可达性探测：master 连不上会表现为 owner `tq.init` 或各节点 attach 握手失败，按下面两行处理 | 先确认 master 进程、DNS/路由和防火墙 | -| **`tq.init` 失败/超时** | 第一次初始化在独立 owner actor 中执行，driver 最多等待 60 秒。失败后回收该 actor 及其拥有的半初始化 controller；`auto` 只重试一次 SimpleStorage，`required` 清理后抛出不含 endpoint/PID/路径的稳定错误类型摘要 | 查看 `mooncake_init_failed:*`、master 日志和 owner 清理日志 | -| **attach 握手在某节点失败/超时/protocol 不符** | driver 汇总各节点结果：`auto` 先完整关闭 Mooncake 状态再统一回退 SimpleStorage（日志 `attach_handshake_failed:*`）；`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` | - -## 资源所有权与安全清理 - -首期按**单任务独占 Ray 集群**实现：同一 Ray cluster 在初始化和运行期间只能有一个 Relax job，操作者不得并发启动第二个 job。多 job admission、端口租约和 master 共享机制都不在首期范围内；删除 token/signature 后，`reap → tq.init` 不再承担并发 initializer 仲裁。 +启动前容量预检采用保守上界：文本按 `seq_length × 32 B`，多模态另加 `seq_length × 784 × 12 B`，再乘 rollout batch、`n_samples_per_prompt` 和 `max_staleness + 1`。容量不足时 `auto` 回退 SimpleStorage，`required` 失败。增大 segment 时必须同步规划每个节点的瞬时内存与 `memlock` 足迹。 -清理只动本作业拥有的资源： +异常退出时 client segment 可能继续注册到 master，直到 Mooncake `client_ttl`（默认 30 秒）到期。 -- 不使用任何 `pkill` / `killall` -- `tq.init` 之前会检查已存在的 `TransferQueueController` 命名 actor：**只有取不到 config（半初始化）、明确超时或 actor 已死时才回收**；其他 GCS/control-plane 异常只做脱敏后中止，不据此杀 actor -- 健康的既有 controller 保持不动，但启动会明确失败，**不会 attach，也不会在退出时关闭它**；操作者应先停止前一作业并清理其 TQ 状态，确保 Ray 集群干净 -- 首次初始化在专用 owner actor 中执行；owner 调度成功后才进入 candidate 初始化。初始化失败时先 bounded best-effort close，再 force-kill owner，并通过预先排队且永不正常返回的 termination probe 确认 owner 进程已经终止；随后才清理并确认 named controller 注销。任一步无法确认都会中止，禁止进入 fallback -- 因为集群在每次初始化前已经通过 clean-cluster gate，owner 失败后出现的 controller 按本次尝试的残留处理。该判断依赖上面的“禁止并发启动第二个 Relax job”硬前提，不提供 concurrent initializer winner 兼容 -- 全局 `tq.close()` 只能由 owner actor 调用；actor、critic、rollout 等附加 worker 只能做本地 detach -- 任一长生命周期 train actor 初始化失败时，driver 会 force-kill 整个 train actor group，并通过预先排队的终止 probe 确认 actor task 已进入终态后再传播失败，防止超时的原生初始化线程稍后修改可复用进程的 TQ 全局状态 -- master 进程始终不被 Relax 触碰 +## Correctness 与依赖 gate -## 容量不足与正确性依赖 +当前 Relax pin 为 TransferQueue `58054a33834aadbcf76aacd6b1e32e25c030f2c9`。现有检查只能确认 retry API 和部分源码顺序，不能证明完整 fail-closed 语义。启用 Mooncake/RDMA 前，上游 TransferQueue PR 必须提供并固定到明确版本或 capability marker，至少保证： -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` 不进入 Mooncake 配置与容量检查。segment 大小可用 `RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB` 调整，容量校验与客户端配置读取同一个值。该上界是保守推导，不能替代运行时错误处理。 +- `batch_upsert_from` 和 `batch_get_into` 的每次 retry 都校验返回结果与请求 key 等长； +- `NOTIFY_DATA_UPDATE_ACK` 必须验证 positive ACK，controller 拒绝时 producer 不能按成功结束。 -当前运行时依赖固定到 TransferQueue commit `58054a33834aadbcf76aacd6b1e32e25c030f2c9`。Relax 在 Mooncake 启动前能检查 retry method 是否存在，以及 `KVStorageManager.put_data` 的源码顺序；这些检查**不能证明**当前 pin 已具备完整的 fail-closed 语义。 +这些修复应在 TransferQueue 上游实现；Relax 不使用 monkey patch 替代依赖修复。 -启用 Mooncake/RDMA 前，上游 TransferQueue PR 必须同时补齐： +当前固定的 Mooncake 0.3.10 在 TCP memcpy 路径存在已确认的静默截断风险。Relax 强制 `MC_STORE_MEMCPY=0`，显式设置不安全值会拒绝启动；升级到已修复版本后再重新评估该 gate。 -- `batch_upsert_from` / `batch_get_into` 的**每次 retry** 都校验返回结果与请求 key 等长，短结果必须显式失败； -- `NOTIFY_DATA_UPDATE_ACK` 必须检查 positive ACK，controller 拒绝更新时 producer 不得按成功返回； -- 合入修复后更新 Relax 的明确 commit pin，并以稳定的 version/capability marker 校验，而不是仅凭 method name 或源码文本推断能力。 +每次验收必须记录 Relax commit SHA、TransferQueue commit 和 Mooncake 版本，不能只记录分支名。 -在上述上游 PR 与新 pin 落地前，这个正确性门槛仍是明确的合入 blocker；本 PR 不通过 monkey patch 改写 TransferQueue。Relax 侧继续保留容量预检、SimpleStorage backpressure，以及失败 store 的“写失败、production 状态不更新”契约测试。真实 Mooncake 的容量与传输验收统一由跨节点 benchmark 承担，不再在 pytest 中维护第二套 direct-client harness。 +## 跨节点验收 -正确性守卫强制 `MC_STORE_MEMCPY=0` 且 **fail-closed**：mooncake 0.3.10 在 TCP-only 环境会自动启用 memcpy 快拷贝路径，该路径存在已确认的静默截断缺陷（现象与处置见排障表）；RDMA 会话本就自动禁用 memcpy，不受影响。由于缺陷在当前 pin 上已实证，显式导出 `MC_STORE_MEMCPY=1` 会在启动时被直接拒绝，待 pin 升级到修复版本后再按版本重新放开。 +唯一保留的 benchmark 是 `scripts/benchmarks/tq_cross_node_bench.py`： -## 验收分层 - -Mock/本机测试和真实双节点 RDMA 测试必须分别报告，前者不能替代后者。 - -| 层级 | 验证内容 | 通过标准 | +| 档位 | 后端 | 用途 | |---|---|---| -| CI/mock | 模式校验、master 端点格式、容量预检、owner 超时与终止确认、controller 清理、attach 握手的 manager/config 契约与拆除顺序、auto/required、有限重试、写失败不发布状态 | `tests/utils/tq/test_config.py`、`tests/core/test_controller_tq_backend.py` 与 `tests/utils/test_tq_failure_paths.py` 全部通过；真机项允许明确 skip | -| 本机 TQ | SimpleStorage 全链路 put/get、容量 backpressure、空读、清理、字节一致性；`multimodal_train_inputs` 以生产容器（`list[dict]` / NonTensorStack，存储层非张量路径）全链路逐叶子 SHA-256 一致 | `tests/utils/test_tq_dataplane_behavior.py` 通过（含 `TestMultimodalFullLink`） | -| 真实多模态载荷 | 真实数据集图像走完整生产预处理链（`build_messages` → `apply_chat_template` → `process_vision_info` → HF processor → `remap_mm_train_inputs`）生成 fixture；上述两级多模态用例检测到 fixture 后自动升级为真实载荷档 | fixture 存在时以 `[real]` 档通过；无 fixture 环境回退 `[synthetic]`（生产同构状，CI 兜底）；交付报告须注明真实档在何处跑过 | -| 真实双节点 | 同一 driver/consumer pair 下各 backend 的原生数据布局；SimpleStorage、Mooncake/TCP、Mooncake/RDMA；synthetic、production-shaped multimodal 两种 profile；256/1024/2048/4096 MiB；每档 warmup + 至少 5 轮 | 每次 get 的 dtype、shape 与 raw-byte SHA-256 全部 PASS；强制 counter gate 证明对应线路；逐协议、逐轮 CSV 留档并报告均值、median、stddev | -| 真实回退（`auto` + 某节点 RDMA 不可用） | 静态探测删除后，`auto` 会真的创建 Mooncake owner 与 named controller，再在 handshake 阶段失败并拆除，因此这条路径必须实测 | 逐条确认：① 日志显示 Mooncake owner `tq.init` **成功**、随后 handshake 阶段失败（否则实际只测到 owner 初始化失败，覆盖不到拆除）；② 最终存在且仅存在一个预期的 SimpleStorage `TransferQueueController`；③ 该 controller 的 stored config 确认为 SimpleStorage；④ SimpleStorage 的 put/get 正常工作；⑤ 旧 Mooncake owner actor 与其 client 均已消失；⑥ Mooncake segment 已从 master 卸载——若测的是强杀/超时路径，需等过 `client_ttl`（默认 30 s）再检查 | - -真实模型/数据 fixture 及其生成流程属于外部 PR 验收附件，不再由仓库测试代码加载或维护。本机 dataplane 测试使用确定性的 production-shaped synthetic payload；唯一保留的跨节点 benchmark 只接受显式的 `synthetic` 和 `multimodal` profile，不会在缺失外部 fixture 时替换验收口径。 +| C0 | SimpleStorage | 默认路径基线 | +| C1 | Mooncake/TCP | benchmark 对照，不是生产配置或 fallback | +| C2 | Mooncake/host-RDMA | 生产 RDMA candidate | -双节点验收命令（master 与 Ray 集群需由部署侧预先准备；每个 protocol 必须启动一个全新的 Python 进程并使用独立 CSV）： +每个 protocol 必须使用全新 Python 进程和独立 CSV。示例仅展示 C2；C0/C1 使用 `--protocol simple` / `tcp`，并删除 `--device`： ```bash -PYTHONPATH=. python -u scripts/benchmarks/tq_cross_node_bench.py \ +python -u scripts/benchmarks/tq_cross_node_bench.py \ --protocol rdma \ - --master :50051 \ - --consumer-node-id \ + --master master.example:50051 \ + --consumer-node-id \ --device \ --tcp-device \ --payload-profiles synthetic multimodal \ @@ -166,56 +93,49 @@ PYTHONPATH=. python -u scripts/benchmarks/tq_cross_node_bench.py \ --csv c2-rdma.csv ``` -用同样方式分别以 `--protocol simple` 和 `--protocol tcp` 运行 C0/C1，但这两档必须删除 `--device`；`--tcp-device` 三档都必须显式指定，因为它选择的是 consumer 节点上的验收 counter。C2 只汇总 `--device` 指定 HCA 的所有端口，不能用其他 HCA 的后台流量充当证据。C1/C2 除要求目标线路流量不低于非目标线路外，还分别要求 TCP/RDMA 接收量至少达到 raw payload 的 20%/80%；C0 因 SimpleStorage unit 可能被放在 consumer 本地，不使用 payload-volume 门槛，只确认观测到 TCP 且没有被 RDMA 流量主导。 - -C1 仅是 benchmark 对照，脚本会在 producer 和 consumer 进程内固定启用 Mooncake TCP connection pool，不形成生产配置面。C0 的 SimpleStorage units 使用上游原生 placement，因此三档共享的是固定 driver/consumer pair，而不是完全相同的 storage topology。每轮 byte-exact 或 wire-proof 失败都会立即终止，失败轮仍会先写入并 flush 到 CSV。 - -若验收环境没有两个 RDMA 节点，交付结论必须写成“真机验收未执行”，不能用 mock 通过推导真机已经通过。 +所有档位都必须 byte-exact。C2 要求 IB receive counter 增长、至少覆盖 raw payload 的 80%，且不被 TCP 流量主导；C1 要求 TCP receive counter 至少覆盖 20%，且不被 RDMA 流量主导。C0 只要求观测到 TCP 且不被 RDMA 主导，因为 SimpleStorage unit 可能位于 consumer 本地。 -完整矩阵结果、依赖版本、逐轮分布与原始 CSV 属于对应 commit 的外部 PR 验收材料，不在本文档维护易过期的历史性能数字。 +真实多模态 fixture、原始 CSV、版本信息和性能分布属于 PR 验收附件，不在仓库文档维护生成教程或易过期的性能数字。没有双节点 RDMA 环境时必须明确记录“真机验收未执行”，不能用 mock 结果替代。 -## 排障表 +## 排障 -| 现象 | 可能原因 | 处理 | -|---|---|---| -| 启动日志 `the installed TransferQueue does not satisfy the Mooncake correctness contract; auto fallback to SimpleStorage` | 镜像里的 TransferQueue/Mooncake 缺少本文“正确性依赖”所列能力 | 核对该节点的 `transfer_queue` / `mooncake-transfer-engine` 版本；镜像是否一致 | -| 启动日志 `the Mooncake master endpoint is not configured` 或 `segment capacity insufficient` / `the segment-capacity configuration is unusable` | 分别是 `MC_MASTER_ADDRESS` 缺失/格式非法、最坏情况容量上界超过 client segment、以及 `RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB` 或 `seq_length` 本身不可用 | 容量类问题按日志核对参数；减少 batch / `n_samples_per_prompt` / `max_staleness`，或调大 `RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB` 并同步规划每节点 attach 的瞬时资源足迹 | -| `Mooncake attach handshake reported N failure(s)`，摘要为 `handshake task failed (...)` 或 `handshake did not return ...` | 该节点 native setup、Ray task 或 attach 超时；日志只保留稳定的错误类型，不回显 endpoint、PID、路径或底层异常正文 | 在失败节点上按下方“选卡核验”逐项确认，并检查 master 连通性、Ray worker 日志与 `memlock` | -| handshake 报错含 `attached storage manager is not MooncakeStorageManager` | 该进程 attach 到的不是预期 Mooncake controller（例如集群里残留了上一次作业的 SimpleStorage controller） | 确认集群干净：本期按单任务独占集群设计，不接管他人的 controller | -| `backend=SimpleStorage fallback=...`，但预期跑 RDMA，且日志显示 master 未配置 | `auto` 下 `MC_MASTER_ADDRESS` 缺失或格式非法会记 WARNING 后回退 SimpleStorage（`required` 直接失败） | 给每个节点的作业环境设置 `MC_MASTER_ADDRESS=:` | -| `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 档（仅 benchmark C1）get 数据尾部全零，但批量返回码全部成功（逐字节校验 FAIL） | mooncake 0.3.10 memcpy 快拷贝路径缺陷：TCP-only 环境被自动启用后，跨节点 get 会静默截断（坏行自 64 KiB 对齐偏移起全零） | 正确性守卫已强制 `MC_STORE_MEMCPY=0`（见“容量不足与正确性依赖”）；显式设 `1` 会被启动拒绝，unset 即可 | -| Mooncake/TCP 档（仅 benchmark C1）在单机回环下原生 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` 是否打出 | - -### 选卡核验 - -Relax 不再自动扫描这些内容——启动只判断真实 attach/setup 与配置契约是否通过，具体哪一项硬件条件不满足需要在失败节点上手工核验。指定设备前先确认端口状态与 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 -``` +| 现象 | 检查与处理 | +|---|---| +| correctness contract 不满足 | 核对各节点 TransferQueue/Mooncake 版本和 capability marker；不要绕过 gate | +| master 缺失、格式错误或不可达 | 检查所有节点的 `MC_MASTER_ADDRESS`、DNS、路由、防火墙和 master 服务 | +| segment capacity insufficient | 减少 batch、采样数或 staleness，或增大 `RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB` 并重新核算资源 | +| attach handshake 失败 | 在失败节点检查 HCA port、GID、内存、`memlock` 和到 master 的连接；CPU-only head 也必须满足 attach 条件 | +| manager/protocol 不符 | 清理前一个作业；首期要求单任务独占且不会接管健康的旧 controller | +| `Connection refused` 指向旧 segment | 等待 `client_ttl` 过期，再确认旧 client/segment 已从 master 清理 | +| C1 尾部全零或 SIGSEGV | 确认未显式启用 `MC_STORE_MEMCPY`；C1 仅用于 benchmark | +| 多 HCA 环境建连失败 | 显式设置 `--tq-rdma-device`，不要依赖自动选卡 | -判定数据面是否真的走了 RDMA（get 前后取差值）： +### 手工核验 HCA、GID、memlock、线路与 master ```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 +export TQ_RDMA_DEVICE= +export TQ_RDMA_PORT=1 +export TQ_TCP_DEVICE= + +cat "/sys/class/infiniband/${TQ_RDMA_DEVICE}/ports/${TQ_RDMA_PORT}/state" +cat "/sys/class/infiniband/${TQ_RDMA_DEVICE}/ports/${TQ_RDMA_PORT}/gids/0" +cat "/sys/class/infiniband/${TQ_RDMA_DEVICE}/ports/${TQ_RDMA_PORT}/rate" +ulimit -l + +# 在数据传输前后取差值；IB counter 单位为 4-byte words。 +cat "/sys/class/infiniband/${TQ_RDMA_DEVICE}/ports/${TQ_RDMA_PORT}/counters/port_rcv_data" +cat "/sys/class/net/${TQ_TCP_DEVICE}/statistics/rx_bytes" +# 使用通用占位 endpoint 检查 DNS 与 TCP。 +export TQ_MASTER_HOST=master.example TQ_MASTER_PORT=50051 +getent hosts "${TQ_MASTER_HOST}" +nc -vz "${TQ_MASTER_HOST}" "${TQ_MASTER_PORT}" ``` -RDMA 生效时前者按 payload 增长、后者基本不动；反之则说明落在 TCP。 +master 可达性可用部署环境已有的 DNS/TCP 工具检查；不要把真实 endpoint、hostname 或本地路径写入提交、公开日志和 PR 文档。 -## 已知限制 +## 已知边界 -- 写侧（put）收益明显，读侧（get）收益有限：get 每次调用都会注册/注销 MR，且 key 粒度是 `样本 × 字段`，碎片化开销盖过了传输收益。MR 常驻注册与读路径零拷贝成型不在首期范围。 -- 跨节点 RDMA vs TCP 的收益在多轮之间波动较大，验收结论应基于多轮分布而非单轮数据。 -- Mooncake/TCP（C1）在 0.3.10 上依赖 `MC_STORE_MEMCPY=0` 守卫保证字节正确，且守卫后 get 吞吐低于 SimpleStorage。它只是 benchmark 的对照档，不是生产路径：RDMA 不可用时生产回退的是 SimpleStorage。 -- 消费端节点在 get 过程中中途死亡的端到端行为需要双节点真机验证，未做成自动化测试。 -- Mooncake 传输层自身的超时参数不由 Relax 控制。 +- 首期不支持 GDR、多 job 或 concurrent initializer。 +- Mooncake 传输层自身的 timeout 不由 Relax 控制。 +- attach 成功不等于 wire proof；C2 合入前仍需真实双节点 byte-exact、wire-proof 和 fully-async smoke。 +- mock/CPU CI 不能替代真实 RDMA 验收，真实测试结果必须关联准确的代码和依赖 SHA。 From fc0a4e2f916b3dc90f611ff1f0444f68578d97c0 Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Mon, 24 Aug 2026 05:03:35 +0800 Subject: [PATCH 27/30] refactor(tq): inline worker detach cleanup --- relax/utils/tq/lifecycle.py | 63 +++++++++++++--------------- tests/utils/test_tq_failure_paths.py | 18 ++++++-- 2 files changed, 43 insertions(+), 38 deletions(-) diff --git a/relax/utils/tq/lifecycle.py b/relax/utils/tq/lifecycle.py index 2c0e31967..3d19d8388 100644 --- a/relax/utils/tq/lifecycle.py +++ b/relax/utils/tq/lifecycle.py @@ -226,39 +226,6 @@ def _get_stored_config(timeout: float = 10.0) -> Any: 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 ({safe_exception_kind(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 ({safe_exception_kind(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 _resolve_attach_timeout() -> float: """Attach deadline in seconds; override via ``RELAX_TQ_ATTACH_TIMEOUT_SECONDS``.""" @@ -373,7 +340,35 @@ def detach_tq_client() -> None: generation lease. Force-killed workers still fall back to the master-side TTL. """ - _close_local_tq_client() + 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 ({safe_exception_kind(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 ({safe_exception_kind(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 _alive_node_ids() -> list[str]: diff --git a/tests/utils/test_tq_failure_paths.py b/tests/utils/test_tq_failure_paths.py index 0732d3d66..e52ec426f 100644 --- a/tests/utils/test_tq_failure_paths.py +++ b/tests/utils/test_tq_failure_paths.py @@ -531,11 +531,21 @@ def test_cluster_attach_timeout_does_not_leave_process_global_state(self): 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)) + def test_detach_closes_clients_and_resets_process_handles(self, monkeypatch): + storage_client = MagicMock() + client = MagicMock(storage_manager=SimpleNamespace(storage_client=storage_client)) + fake_tq = MagicMock() + fake_tq.get_client.return_value = client + fake_interface = SimpleNamespace(_TQ_CLIENT=client, _TQ_CONTROLLER=object()) + monkeypatch.setattr(tq_lifecycle.tq, "interface", fake_interface, raising=False) + monkeypatch.setattr(tq_lifecycle, "tq", fake_tq) + tq_lifecycle.detach_tq_client() - assert calls == [True] + + storage_client.close.assert_called_once_with() + client.close.assert_called_once_with() + assert fake_interface._TQ_CLIENT is None + assert fake_interface._TQ_CONTROLLER is None @pytest.mark.parametrize("has_client", [True, False], ids=["attached", "not-attached"]) def test_component_del_detaches_only_an_attached_client(self, monkeypatch, has_client): From 241dc37e588453b41486543a0dcca4546e609604 Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Mon, 24 Aug 2026 13:58:17 +0800 Subject: [PATCH 28/30] fix(tq): harden RDMA lifecycle and capacity guards Correct multimodal segment sizing, clean partial actor initialization, retain failed owner handles, and unmount storage on teardown. Enforce the Mooncake memcpy safety contract, sanitize topology logs, and align tests and deployment docs. --- docs/draft/transfer_queue_rdma.md | 14 ++++----- relax/core/controller.py | 16 ++++++---- relax/distributed/ray/actor_group.py | 10 ++++-- relax/utils/arguments.py | 6 ++-- relax/utils/tq/config.py | 33 ++++++++++++-------- relax/utils/tq/correctness.py | 28 ++++++++--------- relax/utils/tq/lifecycle.py | 33 ++++++++++++-------- tests/core/test_controller_tq_backend.py | 19 +++++++++-- tests/utils/test_tq_failure_paths.py | 19 ++++++++++- tests/utils/test_train_actor_init_cleanup.py | 17 ++++++++++ tests/utils/tq/test_config.py | 10 +++--- 11 files changed, 137 insertions(+), 68 deletions(-) diff --git a/docs/draft/transfer_queue_rdma.md b/docs/draft/transfer_queue_rdma.md index 20ca1e6f3..6aca88b28 100644 --- a/docs/draft/transfer_queue_rdma.md +++ b/docs/draft/transfer_queue_rdma.md @@ -6,7 +6,7 @@ Relax 默认使用 TransferQueue SimpleStorage。首期 RDMA 支持只接入 Moo | 参数 | 语义 | |---|---| -| `--tq-rdma-mode=off` | 默认值；保持原有 SimpleStorage 路径，不执行 Mooncake 检查或 attach | +| `--tq-rdma-mode=off` | 默认值；保持 SimpleStorage 数据路径，不执行 Mooncake 检查或集群 handshake；worker 的 TQ attach 仍有 60 秒边界 | | `--tq-rdma-mode=auto` | 尝试 host-RDMA；任一检查或节点 attach 失败时，完成清理后统一回退 SimpleStorage | | `--tq-rdma-mode=required` | 要求 host-RDMA；不可用时完成清理并终止启动，不允许降级 | | `--tq-rdma-device=` | 指定 RDMA 设备；空值由 Mooncake 原生逻辑选择，多 HCA 环境建议显式指定 | @@ -17,17 +17,17 @@ Relax 固定 Mooncake `use_gdr=false`。GDR 若有实际需求，应由独立 PR 首期只支持**单任务独占 Ray 集群**：同一个 Ray cluster 在初始化和运行期间只能有一个 Relax job，不支持 concurrent initializer、多 job admission 或复用其他作业的 TransferQueue controller。 -Mooncake master 由部署环境管理，Relax 不启动、重启或停止它。所有节点必须设置同一个外部 endpoint： +Mooncake master 由部署环境管理，Relax 不启动、重启或停止它。driver 必须设置外部 endpoint： ```bash export MC_MASTER_ADDRESS=master.example:50051 ``` -Relax 只检查 `host:port` 格式；DNS、路由、防火墙和 master 健康状态由真实初始化与 attach 验证。 +driver 将该 endpoint 写入 job-level TQ config，owner 和 worker attach 复用已存储配置，因此 worker 节点不要求重复设置该环境变量；但所有节点都必须能访问同一个 endpoint。Relax 只检查 `host:port` 格式；DNS、路由、防火墙和 master 健康状态由真实初始化与 attach 验证。 ## 启动、回退与清理 -`off` 直接初始化 SimpleStorage。`auto` 和 `required` 按以下顺序执行： +`off` 直接初始化 SimpleStorage，不创建 owner actor 或执行 Mooncake handshake；各 worker 仍通过有界 helper attach 到 SimpleStorage，以避免旧版无界等待。启动时还会回收不可用的半初始化 controller，并因首期独占集群约束拒绝复用健康的既有 controller。`auto` 和 `required` 按以下顺序执行： 1. 校验 mode、TransferQueue/Mooncake correctness contract、master 格式和 segment 容量。 2. 在独立 owner actor 中有界执行 Mooncake `tq.init`。 @@ -51,7 +51,7 @@ global segment 可按部署容量调整： export RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB=8 ``` -启动前容量预检采用保守上界：文本按 `seq_length × 32 B`，多模态另加 `seq_length × 784 × 12 B`，再乘 rollout batch、`n_samples_per_prompt` 和 `max_staleness + 1`。容量不足时 `auto` 回退 SimpleStorage，`required` 失败。增大 segment 时必须同步规划每个节点的瞬时内存与 `memlock` 足迹。 +启动前容量预检采用保守上界：文本按 `seq_length × 32 B`；多模态按当前支持的最大 transported tensor layout（16×16 spatial patch、temporal patch 2、RGB、2×2 spatial merge、float32）另加 `seq_length × 24,576 B`。因此 8,192 token 的单样本多模态 tensor 上界约 192 MiB。最终再乘 rollout batch、`n_samples_per_prompt` 和 `max_staleness + 1`。容量不足时 `auto` 回退 SimpleStorage，`required` 失败。增大 segment 时必须同步规划每个节点的瞬时内存与 `memlock` 足迹。 异常退出时 client segment 可能继续注册到 master，直到 Mooncake `client_ttl`（默认 30 秒）到期。 @@ -64,7 +64,7 @@ export RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB=8 这些修复应在 TransferQueue 上游实现；Relax 不使用 monkey patch 替代依赖修复。 -当前固定的 Mooncake 0.3.10 在 TCP memcpy 路径存在已确认的静默截断风险。Relax 强制 `MC_STORE_MEMCPY=0`，显式设置不安全值会拒绝启动；升级到已修复版本后再重新评估该 gate。 +测试确认 `mooncake-transfer-engine==0.3.10.post2` 的 TCP memcpy 路径存在静默截断风险。Relax 当前会统一强制 `MC_STORE_MEMCPY=0`，显式设置不安全值会拒绝启动；只有在能够可靠识别已修复 build 后，才重新评估是否允许 memcpy 路径。 每次验收必须记录 Relax commit SHA、TransferQueue commit 和 Mooncake 版本，不能只记录分支名。 @@ -102,7 +102,7 @@ python -u scripts/benchmarks/tq_cross_node_bench.py \ | 现象 | 检查与处理 | |---|---| | correctness contract 不满足 | 核对各节点 TransferQueue/Mooncake 版本和 capability marker；不要绕过 gate | -| master 缺失、格式错误或不可达 | 检查所有节点的 `MC_MASTER_ADDRESS`、DNS、路由、防火墙和 master 服务 | +| master 缺失、格式错误或不可达 | 检查 driver 的 `MC_MASTER_ADDRESS`，并检查所有节点到该 endpoint 的 DNS、路由、防火墙和 master 服务 | | segment capacity insufficient | 减少 batch、采样数或 staleness，或增大 `RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB` 并重新核算资源 | | attach handshake 失败 | 在失败节点检查 HCA port、GID、内存、`memlock` 和到 master 的连接；CPU-only head 也必须满足 attach 条件 | | manager/protocol 不符 | 清理前一个作业；首期要求单任务独占且不会接管健康的旧 controller | diff --git a/relax/core/controller.py b/relax/core/controller.py index 4139db3a1..9c50be1a5 100644 --- a/relax/core/controller.py +++ b/relax/core/controller.py @@ -382,10 +382,11 @@ def _confirm_mooncake_attach(self, init_result: TqInitResult, fallback_config) - """ def _close_owned_attempt() -> None: - try: - close_tq_owner(init_result.owner) - finally: - self._tq_owner = None + # Retain the handle if teardown fails. The constructor's outer + # cleanup boundary can then retry instead of losing the only + # reference to possibly live process-global TQ state. + close_tq_owner(init_result.owner) + self._tq_owner = None try: failures = verify_cluster_attach(init_result.config) @@ -512,8 +513,11 @@ def _fall_back_or_raise(reason: str, error: Exception) -> dict: # 5. One startup line stating what was requested and what will run. The # effective transport is only confirmed once the cluster-wide attach # handshake passes. - logger.info(f"[dataplane] requested: rdma_mode={mode} device={device or 'auto'}") - logger.info(f"[dataplane] backend=MooncakeStore protocol=rdma device={device or 'auto'} (pending handshake)") + device_selection = "explicit" if device else "auto" + logger.info(f"[dataplane] requested: rdma_mode={mode} device_selection={device_selection}") + logger.info( + f"[dataplane] backend=MooncakeStore protocol=rdma device_selection={device_selection} (pending handshake)" + ) return backend_dict def _close_data_system(self) -> None: diff --git a/relax/distributed/ray/actor_group.py b/relax/distributed/ray/actor_group.py index 544000c7c..485d5e98e 100644 --- a/relax/distributed/ray/actor_group.py +++ b/relax/distributed/ray/actor_group.py @@ -137,10 +137,14 @@ def init_and_wait( safely, so all actor processes are force-killed and confirmed terminal before the initialization error is propagated. """ - refs = self.async_init(args, role, with_ref=with_ref, with_opd_teacher=with_opd_teacher) - pending = list(refs) - results: dict[Any, Any] = {} try: + # Submission itself can fail after one or more actor calls were + # already queued. Keep it inside the same cleanup boundary as + # asynchronous task failures so no partially initialized group is + # left reusable. + refs = self.async_init(args, role, with_ref=with_ref, with_opd_teacher=with_opd_teacher) + pending = list(refs) + results: dict[Any, Any] = {} # ``ray.get(refs)`` may wait for every ref before surfacing one # failure. A rank blocked in native initialization would then # prevent cleanup forever, so consume whichever ref finishes first. diff --git a/relax/utils/arguments.py b/relax/utils/arguments.py index 5aaae88df..dc07dc1a0 100644 --- a/relax/utils/arguments.py +++ b/relax/utils/arguments.py @@ -245,8 +245,8 @@ def add_transfer_queue_arguments(parser): 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, + # The default selects the existing SimpleStorage data path. 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( @@ -255,7 +255,7 @@ def add_transfer_queue_arguments(parser): default="off", help=( "TransferQueue data-plane transport. 'off' (default) uses " - "SimpleStorage/ZMQ and is equivalent to current behavior. " + "SimpleStorage/ZMQ with bounded worker attach and cleanup. " "'auto' attempts MooncakeStore over host RDMA and falls back " "to SimpleStorage when it is unavailable (with a WARNING). " "'required' fails fast instead of falling back." diff --git a/relax/utils/tq/config.py b/relax/utils/tq/config.py index 864f414ed..ce95c2163 100644 --- a/relax/utils/tq/config.py +++ b/relax/utils/tq/config.py @@ -109,8 +109,8 @@ def resolve_mooncake_master_address() -> str: 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." + "managed mooncake master in the driver environment; Relax never assumes " + "a loopback endpoint." ) try: _split_host_port(address) @@ -269,11 +269,18 @@ def build_mooncake_config( # 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 +# +# The largest image/video tensor layout among the currently supported +# multimodal processors is Qwen3-VL's 16x16 spatial patch, temporal patch 2, +# RGB input, spatial merge 2x2, stored as float32. One schedulable vision +# token can therefore retain four flattened patch rows: +# +# 16 * 16 * 2 * 3 * 4 rows/token * 4 bytes/value = 24,576 bytes/token +# +# Keep this as one explicit bound instead of loading model config during +# Controller startup. If a supported processor gains a wider transported +# feature row, this bound and its regression test must be updated together. +_MULTIMODAL_BYTES_PER_TOKEN = 16 * 16 * 2 * 3 * (2 * 2) * 4 # Text: token ids, logprobs, masks and rewards; 32 B/token rounds them up. _TEXT_BYTES_PER_TOKEN = 32 @@ -299,11 +306,11 @@ 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. + processor cannot emit more multimodal tokens than ``--seq-length`` allows, + so the transported tensor payload of one sample is bounded by + ``seq_length * _MULTIMODAL_BYTES_PER_TOKEN`` (192 MiB at + ``seq_length=8192``) rather than a fixed per-sample guess that can pass + configurations which later fail puts mid-training. """ n_samples = args.n_samples_per_prompt capacity_batch = resolve_tq_capacity_batch_size(args) @@ -315,7 +322,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 + per_sample += seq_length * _MULTIMODAL_BYTES_PER_TOKEN return capacity_batch * n_samples * per_sample diff --git a/relax/utils/tq/correctness.py b/relax/utils/tq/correctness.py index 501845e2f..a3fd54b83 100644 --- a/relax/utils/tq/correctness.py +++ b/relax/utils/tq/correctness.py @@ -8,9 +8,9 @@ consumed through an updated, capability-marked pin; method-name checks alone do not prove those semantics. -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`). +Mooncake 0.3.10.post2 was observed corrupting 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 @@ -21,24 +21,24 @@ 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 - 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 + mooncake 0.3.10.post2 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. 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. + that build, this guard fails closed: an explicit ``MC_STORE_MEMCPY=1`` is + rejected at startup instead of honoured. Re-gate this behavior only when a + reliable fixed-build capability is available. """ override = os.environ.get("MC_STORE_MEMCPY", "").strip() if override not in ("", "0"): raise RuntimeError( - "MC_STORE_MEMCPY explicitly enables an unsafe value: the pinned mooncake 0.3.10 " + "MC_STORE_MEMCPY explicitly enables an unsafe value: the mooncake 0.3.10.post2 " "memcpy fast path silently truncates TCP transfers and can SIGSEGV. " - "Unset MC_STORE_MEMCPY; Relax forces it to 0 on this version." + "Unset MC_STORE_MEMCPY; Relax forces it to 0 until a fixed-build capability is available." ) os.environ["MC_STORE_MEMCPY"] = "0" diff --git a/relax/utils/tq/lifecycle.py b/relax/utils/tq/lifecycle.py index 3d19d8388..3198cd94d 100644 --- a/relax/utils/tq/lifecycle.py +++ b/relax/utils/tq/lifecycle.py @@ -438,7 +438,7 @@ def verify_cluster_attach(conf: Any, *, timeout: float | None = None) -> list[st expects_mooncake = uses_mooncake(conf) refs: list[Any] = [] - id_by_ref: dict[Any, str] = {} + node_number_by_ref: dict[Any, int] = {} try: from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy @@ -459,11 +459,14 @@ def _handshake(handshake_conf: Any, check_mooncake: bool, attach_timeout: float) # segment; without the detach it lingers until client_ttl. detach_tq_client() - for node_id in node_ids: + for node_number, node_id in enumerate(node_ids, start=1): strategy = NodeAffinitySchedulingStrategy(node_id=node_id, soft=False) ref = _handshake.options(scheduling_strategy=strategy).remote(conf, expects_mooncake, timeout) refs.append(ref) - id_by_ref[ref] = node_id + # Keep the raw Ray NodeID only in the scheduling strategy. Failure + # summaries may reach public CI/PR logs, so identify nodes by a + # stable per-handshake ordinal instead. + node_number_by_ref[ref] = node_number except Exception as error: if refs and not _cancel_handshake_tasks(refs): raise TqHandshakeIsolationError( @@ -487,11 +490,11 @@ def _handshake(handshake_conf: Any, check_mooncake: bool, attach_timeout: float) try: ray.get(ref) except Exception as error: - failures.append(f"node {id_by_ref[ref][:12]}: handshake task failed ({safe_exception_kind(error)})") + failures.append(f"node#{node_number_by_ref[ref]}: handshake task failed ({safe_exception_kind(error)})") if pending and not _cancel_handshake_tasks(pending): raise TqHandshakeIsolationError("timed-out TQ handshake workers could not be confirmed stopped") for ref in pending: - failures.append(f"node {id_by_ref[ref][:12]}: handshake did not return within {wait_bound:.0f}s") + failures.append(f"node#{node_number_by_ref[ref]}: handshake did not return within {wait_bound:.0f}s") return failures @@ -512,14 +515,18 @@ def close_tq_and_unmount() -> 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 ({safe_exception_kind(e)}).") + try: + tq.close() + finally: + # ``tq.close`` may fail while removing queued data. Segment teardown + # is still mandatory; otherwise the failed owner leaves registered and + # locked memory behind until the Mooncake TTL expires. + 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 ({safe_exception_kind(e)}).") @ray.remote(num_cpus=0) diff --git a/tests/core/test_controller_tq_backend.py b/tests/core/test_controller_tq_backend.py index df82dc282..439d312ca 100644 --- a/tests/core/test_controller_tq_backend.py +++ b/tests/core/test_controller_tq_backend.py @@ -225,10 +225,16 @@ def test_unmet_precondition( def test_satisfied_preconditions_select_host_rdma(self, monkeypatch): _Recorder(monkeypatch) - backend = _resolve(_config(tq_rdma_device="mlx5_0")) + private_device = "private-device-name" + messages: list[str] = [] + monkeypatch.setattr(controller.logger, "info", lambda message: messages.append(message)) + + backend = _resolve(_config(tq_rdma_device=private_device)) assert backend["storage_backend"] == "MooncakeStore" assert backend["MooncakeStore"]["protocol"] == "rdma" - assert backend["MooncakeStore"]["device_name"] == "mlx5_0" + assert backend["MooncakeStore"]["device_name"] == private_device + assert any("device_selection=explicit" in message for message in messages) + assert all(private_device not in message for message in messages) def test_required_accepts_satisfied_preconditions(self, monkeypatch): _Recorder(monkeypatch) @@ -371,10 +377,17 @@ def test_cleanup_failure_aborts_instead_of_falling_back(self, monkeypatch): failures=["node: attach timed out"], close_error=RuntimeError("TransferQueue owner cleanup failed"), ) + owner = object() + instance = controller.Controller.__new__(controller.Controller) + instance.config = _config() + instance._tq_owner = owner + init_result = controller.TqInitResult(config="mooncake-conf", owner=owner) + with pytest.raises(RuntimeError, match="owner cleanup failed"): - _confirm(_config()) + instance._confirm_mooncake_attach(init_result, "simple-conf") assert recorder.events == ["handshake", "close"] assert recorder.initialized == [] + assert instance._tq_owner is owner def test_required_closes_owner_then_raises(self, monkeypatch): recorder = _AttachRecorder(monkeypatch, failures=["node: protocol=tcp"]) diff --git a/tests/utils/test_tq_failure_paths.py b/tests/utils/test_tq_failure_paths.py index e52ec426f..31ec62d49 100644 --- a/tests/utils/test_tq_failure_paths.py +++ b/tests/utils/test_tq_failure_paths.py @@ -166,6 +166,22 @@ def test_mooncake_segment_is_unmounted_after_close(self, monkeypatch): # Order matters: tq.close() still needs the store alive for remove_all(). assert calls == ["tq.close", "store.close"] + def test_segment_is_unmounted_even_when_tq_close_fails(self, monkeypatch): + store_client = MagicMock() + calls = self._fake_tq(monkeypatch, store_client=store_client) + + def close_failure(): + calls.append("tq.close") + raise RuntimeError("close failed") + + tq_lifecycle.tq.close.side_effect = close_failure + store_client.close.side_effect = lambda: calls.append("store.close") + + with pytest.raises(RuntimeError, match="close failed"): + tq_lifecycle.close_tq_and_unmount() + + 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() @@ -380,7 +396,8 @@ def test_ready_task_failure_is_sanitized(self, monkeypatch): failures = tq_lifecycle.verify_cluster_attach({}, timeout=0.1) - assert failures[0].endswith("handshake task failed (RuntimeError)") + assert failures == ["node#1: handshake task failed (RuntimeError)"] + assert "a" * 56 not in failures[0] assert secret not in failures[0] def test_partial_scheduling_failure_cancels_submitted_workers(self, monkeypatch): diff --git a/tests/utils/test_train_actor_init_cleanup.py b/tests/utils/test_train_actor_init_cleanup.py index df2aff1c4..a64d4f4ce 100644 --- a/tests/utils/test_train_actor_init_cleanup.py +++ b/tests/utils/test_train_actor_init_cleanup.py @@ -51,6 +51,23 @@ def test_successful_initialization_preserves_actor_group(monkeypatch): assert group._actor_handlers == [actor] +def test_synchronous_submission_failure_cleans_partially_initialized_group(monkeypatch): + actor = _FakeActor() + group = _group(actor) + cleanup_calls = [] + + def fail_submission(*_args, **_kwargs): + raise RuntimeError("actor initialization submission failed") + + monkeypatch.setattr(group, "async_init", fail_submission) + monkeypatch.setattr(group, "_terminate_failed_init", lambda: cleanup_calls.append(True)) + + with pytest.raises(RuntimeError, match="submission failed"): + group.init_and_wait(object(), "actor") + + assert cleanup_calls == [True] + + def test_first_rank_failure_kills_and_confirms_every_actor(monkeypatch): actors = [_FakeActor("probe-0"), _FakeActor("probe-1")] group = _group(*actors) diff --git a/tests/utils/tq/test_config.py b/tests/utils/tq/test_config.py index f918218a3..9e5ca36dc 100644 --- a/tests/utils/tq/test_config.py +++ b/tests/utils/tq/test_config.py @@ -309,7 +309,7 @@ def test_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 ) - monkeypatch.setenv("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB", "8") + monkeypatch.setenv("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB", "16") assert validate_segment_capacity(args) is None @pytest.mark.parametrize("value", ["four", "-1", "nan", "inf", "-inf"]) @@ -338,12 +338,12 @@ def test_text_only_is_small_but_nonzero(self): assert estimate_payload_bytes(_make_args(multimodal_keys=None)) == 32 * 1 * 8192 * 32 def test_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. + # Qwen3-VL transports four 1536-float32 patch rows per schedulable + # vision token, or 24,576 bytes/token: ~192 MiB at 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 + assert per_sample == 8192 * (32 + 24_576) + assert 192 * 1024**2 < per_sample < 193 * 1024**2 def test_requires_seq_length(self): with pytest.raises(RuntimeError, match="seq_length"): From dc2963357f7d54a92d2d1a822a4fc8954e532ae2 Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Tue, 25 Aug 2026 00:54:39 +0800 Subject: [PATCH 29/30] fix(tq): harden benchmark and consolidate tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 🐛 Bug Fix ## Make cross-node acceptance fail closed - Require exact counter selection and account for idle traffic - Persist byte-exact and wire-proof failures before aborting - Preserve validation failures when partition cleanup also fails - Record exact source and dependency provenance in result rows --- # ♻️ Refactor ## Consolidate correctness and lifecycle coverage - Reuse raw-byte payload helpers across acceptance and dataplane checks - Parameterize repeated configuration, lifecycle, and cleanup scenarios - Retain process-isolation, retry, backpressure, and multimodal coverage --- # 📝 Documentation ## Clarify acceptance requirements - Document counter scope, idle adjustment, TCP baseline, and result provenance --- docs/draft/transfer_queue_rdma.md | 9 +- relax/utils/tq/correctness.py | 138 +- scripts/benchmarks/tq_cross_node_bench.py | 617 ++++--- tests/core/test_controller_tq_backend.py | 611 +++---- tests/utils/_tq_handshake_timeout_probe.py | 134 +- .../utils/_train_actor_init_cleanup_probe.py | 45 +- tests/utils/test_tq_benchmark_guards.py | 281 ++-- tests/utils/test_tq_dataplane_behavior.py | 299 ++-- tests/utils/test_tq_failure_paths.py | 1416 ++++++----------- tests/utils/test_train_actor_init_cleanup.py | 144 +- tests/utils/tq/_payload_assertions.py | 109 -- tests/utils/tq/test_config.py | 513 +++--- tests/utils/tq/test_payload_assertions.py | 97 +- 13 files changed, 1844 insertions(+), 2569 deletions(-) delete mode 100644 tests/utils/tq/_payload_assertions.py diff --git a/docs/draft/transfer_queue_rdma.md b/docs/draft/transfer_queue_rdma.md index 6aca88b28..29410ab28 100644 --- a/docs/draft/transfer_queue_rdma.md +++ b/docs/draft/transfer_queue_rdma.md @@ -78,7 +78,7 @@ export RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB=8 | C1 | Mooncake/TCP | benchmark 对照，不是生产配置或 fallback | | C2 | Mooncake/host-RDMA | 生产 RDMA candidate | -每个 protocol 必须使用全新 Python 进程和独立 CSV。示例仅展示 C2；C0/C1 使用 `--protocol simple` / `tcp`，并删除 `--device`： +每个 protocol 必须使用全新 Python 进程和独立 CSV。示例仅展示 C2；C1 改用 `--protocol tcp` 并删除 `--device`/`--rdma-port`，C0 改用 `--protocol simple` 并额外删除 `--master`： ```bash python -u scripts/benchmarks/tq_cross_node_bench.py \ @@ -86,6 +86,7 @@ python -u scripts/benchmarks/tq_cross_node_bench.py \ --master master.example:50051 \ --consumer-node-id \ --device \ + --rdma-port 1 \ --tcp-device \ --payload-profiles synthetic multimodal \ --payload-mib 256 1024 2048 4096 \ @@ -93,7 +94,11 @@ python -u scripts/benchmarks/tq_cross_node_bench.py \ --csv c2-rdma.csv ``` -所有档位都必须 byte-exact。C2 要求 IB receive counter 增长、至少覆盖 raw payload 的 80%，且不被 TCP 流量主导；C1 要求 TCP receive counter 至少覆盖 20%，且不被 RDMA 流量主导。C0 只要求观测到 TCP 且不被 RDMA 主导，因为 SimpleStorage unit 可能位于 consumer 本地。 +所有档位都必须 byte-exact。C2 只读取明确指定的 HCA/port，要求一次完整 `put → get` round 的 idle-adjusted IB receive bytes 至少覆盖 raw payload 的 80%；完整 round 可以覆盖对象随机落在 producer、consumer 或 owner segment 的情况。C1 要求 idle-adjusted TCP receive bytes 至少覆盖 20%，C0 只要求观测到跨节点 TCP，因为 SimpleStorage unit 可能位于 consumer 本地。端口级 counter 不是 per-flow 指标，正式验收必须使用静默或独占的数据端口；CSV 同时记录 raw delta、紧邻 round 的 idle rate 和扣除后的 proof bytes，不能在共享端口有显著背景流量时宣称 wire proof。 + +C1 会在 driver 及其 Ray worker runtime 中设置 `MC_TCP_ENABLE_CONNECTION_POOL=1`，这是当前 Mooncake/TCP correctness baseline 的组成部分，必须随结果一并记录。benchmark 通过一次性 owner lifecycle 有界初始化 TQ，producer/consumer 均有界 attach；任一失败都非零退出，不执行 backend fallback。 + +CSV 使用 exclusive-create，不覆盖既有文件。每个 warmup/测量 round 都先落盘再执行 byte-exact/wire gate，失败行包含稳定的 `error_kind`；每轮无论成功失败都清理 partition。每行还记录 Relax SHA、安装的 TransferQueue VCS commit 和 Mooncake package version，且 Relax tracked worktree 不干净时拒绝作为正式验收运行。 真实多模态 fixture、原始 CSV、版本信息和性能分布属于 PR 验收附件，不在仓库文档维护生成教程或易过期的性能数字。没有双节点 RDMA 环境时必须明确记录“真机验收未执行”，不能用 mock 结果替代。 diff --git a/relax/utils/tq/correctness.py b/relax/utils/tq/correctness.py index a3fd54b83..2f1056818 100644 --- a/relax/utils/tq/correctness.py +++ b/relax/utils/tq/correctness.py @@ -1,6 +1,6 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. -"""Fail-closed capability validation for TransferQueue's Mooncake backend. +"""Mooncake safety guards and raw-byte payload correctness helpers. Relax validates the capabilities and environment it can inspect without modifying TransferQueue at runtime. Per-retry result-length validation and a @@ -15,7 +15,14 @@ from __future__ import annotations +import hashlib import os +import struct +from collections.abc import Mapping +from typing import Any + + +LeafDigest = tuple[str, str, str] def _enforce_safe_memcpy() -> None: @@ -60,3 +67,132 @@ 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)) + + +def _tensor_digest(value: Any) -> LeafDigest: + import torch + + 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: Any) -> LeafDigest: + import numpy as np + + 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") + elif isinstance(value, float): + raw = struct.pack("!d", value) + else: + raw = repr(value).encode("utf-8") + return f"py.{type(value).__name__}", "", hashlib.sha256(raw).hexdigest() + + +def _unwrap_non_tensor(value: Any) -> Any: + if type(value).__name__ == "NonTensorStack": + return value.tolist() + if type(value).__name__ == "NonTensorData": + return value.data + return value + + +def _dict_child_path(prefix: str, key: Any) -> str: + if not isinstance(key, str): + raise TypeError(f"Unsupported payload dict key at {prefix}: {type(key).__name__}") + return f"{prefix}.{key}" if key.isidentifier() else f"{prefix}[{key!r}]" + + +def leaf_digests(payload: Any, prefix: str = "payload") -> dict[str, LeafDigest]: + """Map every supported leaf to ``(dtype, shape, raw-byte SHA-256)``.""" + import numpy as np + import torch + + payload = _unwrap_non_tensor(payload) + digests: dict[str, LeafDigest] = {} + if isinstance(payload, torch.Tensor): + if payload.is_nested: + for index, row in enumerate(payload.unbind()): + digests[f"{prefix}[{index}]"] = _tensor_digest(row) + else: + digests[prefix] = _tensor_digest(payload) + elif isinstance(payload, np.ndarray): + digests[prefix] = _ndarray_digest(payload) + elif isinstance(payload, Mapping): + for key in payload: + if not isinstance(key, str): + raise TypeError(f"Unsupported payload dict key at {prefix}: {type(key).__name__}") + for key in sorted(payload): + digests.update(leaf_digests(payload[key], _dict_child_path(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 mismatch descriptions; an 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 payload_rows(value: Any) -> list[Any]: + """Return logical rows from a tensor or non-tensor TQ column.""" + import torch + + value = _unwrap_non_tensor(value) + if isinstance(value, torch.Tensor): + return [value] if value.ndim == 0 else list(value.unbind()) + if isinstance(value, (list, tuple)): + return [_unwrap_non_tensor(row) for row in value] + raise TypeError(f"Unsupported payload column: {type(value).__name__}") + + +def payload_nbytes(value: Any) -> int: + """Count raw bytes represented by supported TQ payload values.""" + import numpy as np + import torch + + value = _unwrap_non_tensor(value) + if isinstance(value, torch.Tensor): + if value.is_nested: + return sum(payload_nbytes(row) for row in value.unbind()) + return value.numel() * value.element_size() + if isinstance(value, np.ndarray): + return value.nbytes + if isinstance(value, Mapping): + return sum(payload_nbytes(item) for item in value.values()) + if isinstance(value, (list, tuple)): + return sum(payload_nbytes(item) for item in value) + if isinstance(value, bytes): + return len(value) + if isinstance(value, str): + return len(value.encode("utf-8")) + if isinstance(value, float): + return 8 + if isinstance(value, (bool, int)) or value is None: + return len(repr(value).encode("utf-8")) + raise TypeError(f"Unsupported payload leaf: {type(value).__name__}") diff --git a/scripts/benchmarks/tq_cross_node_bench.py b/scripts/benchmarks/tq_cross_node_bench.py index b23cceac2..4074ae7bd 100644 --- a/scripts/benchmarks/tq_cross_node_bench.py +++ b/scripts/benchmarks/tq_cross_node_bench.py @@ -14,26 +14,38 @@ import argparse import csv -import hashlib +import importlib.metadata as importlib_metadata +import json +import math import os import statistics -import struct +import subprocess import time -from collections.abc import Mapping from pathlib import Path -from typing import Any +from typing import Any, TextIO import ray +from relax.utils.tq.correctness import leaf_digests, payload_nbytes, payload_rows + PROTOCOL_LABELS = { "simple": "C0 SimpleStorage", "tcp": "C1 Mooncake/TCP", "rdma": "C2 Mooncake/RDMA", } -CSV_COLUMNS = "protocol profile payload_mib actual_mib run byte_exact wire_proven put_ms get_ms put_gbs get_gbs ib_mb tcp_mb".split() +CSV_COLUMNS = ( + "protocol profile payload_mib actual_mib run measured status error_kind byte_exact mismatch_fields wire_proven " + "put_ms get_ms round_ms put_gbs get_gbs ib_mb tcp_mb proof_ib_mb proof_tcp_mb idle_ib_mbps idle_tcp_mbps " + "relax_sha tq_commit mooncake_version" +).split() _TCP_WIRE_MIN_PAYLOAD_RATIO = 0.20 _RDMA_WIRE_MIN_PAYLOAD_RATIO = 0.80 +_IDLE_COUNTER_SAMPLE_SECONDS = 0.5 +_GLOB_MAGIC = frozenset("*?[") + +CounterMap = dict[str, int] +DigestMap = dict[str, tuple[Any, ...]] def _is_safe_device_name(value: str) -> bool: @@ -42,6 +54,7 @@ def _is_safe_device_name(value: str) -> bool: and value not in {".", ".."} and "/" not in value and "\\" not in value + and not any(character in _GLOB_MAGIC for character in value) and all(character.isprintable() and not character.isspace() for character in value) ) @@ -50,8 +63,13 @@ def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Cross-node TQ byte-exact and wire-proof benchmark") parser.add_argument("--protocol", required=True, choices=sorted(PROTOCOL_LABELS)) parser.add_argument("--consumer-node-id", required=True, help="Alive Ray NodeID distinct from the driver node") - parser.add_argument("--master", default=os.environ.get("MC_MASTER_ADDRESS", ""), help="Mooncake master host:port") + parser.add_argument( + "--master", + default=None, + help="Mooncake master host:port (defaults to MC_MASTER_ADDRESS for Mooncake protocols)", + ) parser.add_argument("--device", default="", help="RDMA device required by the C2 wire-proof counter") + parser.add_argument("--rdma-port", type=int, default=None, help="RDMA port required by the C2 wire-proof counter") parser.add_argument("--tcp-device", required=True, help="Network interface used for the TCP receive counter") parser.add_argument( "--payload-profiles", nargs="+", default=["synthetic", "multimodal"], choices=["synthetic", "multimodal"] @@ -60,22 +78,30 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--num-samples", type=int, default=256) parser.add_argument("--repeats", type=int, default=5) parser.add_argument("--segment-gib", type=int, default=16) - parser.add_argument("--csv", required=True, help="Per-round CSV output path") + parser.add_argument("--csv", required=True, type=Path, help="New per-round CSV output path") args = parser.parse_args() for name in ("num_samples", "repeats", "segment_gib"): if getattr(args, name) <= 0: parser.error(f"--{name.replace('_', '-')} must be positive") if not args.payload_mib or any(size <= 0 for size in args.payload_mib): parser.error("--payload-mib values must be positive") - if args.protocol != "simple" and not args.master: - parser.error("--master or MC_MASTER_ADDRESS is required for Mooncake protocols") + if args.protocol == "simple": + if args.master is not None: + parser.error("--master is valid only with Mooncake protocols") + args.master = "" + else: + args.master = args.master or os.environ.get("MC_MASTER_ADDRESS", "") + if not args.master: + parser.error("--master or MC_MASTER_ADDRESS is required for Mooncake protocols") if not _is_safe_device_name(args.tcp_device): parser.error("--tcp-device must be a single printable interface name without whitespace or path separators") if args.protocol == "rdma": if not _is_safe_device_name(args.device): parser.error("--device is required for RDMA and must be a single printable device name") - elif args.device: - parser.error("--device is valid only with --protocol rdma") + if args.rdma_port is None or args.rdma_port <= 0: + parser.error("--rdma-port is required for RDMA and must be positive") + elif args.device or args.rdma_port is not None: + parser.error("--device and --rdma-port are valid only with --protocol rdma") return args @@ -85,7 +111,7 @@ def _columns_for_budget(total_bytes: int, num_samples: int, dtype: Any) -> int: return max(1, total_bytes // (num_samples * torch.tensor([], dtype=dtype).element_size())) -def make_synthetic_payload(num_samples: int, total_mib: int): +def make_synthetic_payload(num_samples: int, total_mib: int) -> Any: """Deterministic mixed-dtype payload with a fixed schema.""" import torch from tensordict import TensorDict @@ -105,7 +131,7 @@ def make_synthetic_payload(num_samples: int, total_mib: int): ) -def make_multimodal_payload(num_samples: int, total_mib: int): +def make_multimodal_payload(num_samples: int, total_mib: int) -> Any: """Deterministic production-shaped non-tensor vision-language payload.""" import torch @@ -147,143 +173,82 @@ def make_multimodal_payload(num_samples: int, total_mib: int): return payload -def _unwrap_non_tensor(value: Any) -> Any: - if type(value).__name__ == "NonTensorStack": - return value.tolist() - if type(value).__name__ == "NonTensorData": - return value.data - return value - - -def _column_rows(value: Any) -> list[Any]: - import torch - - value = _unwrap_non_tensor(value) - if isinstance(value, torch.Tensor): - if value.is_nested: - return list(value.unbind()) - if value.ndim == 0: - return [value] - return list(value.unbind()) - if isinstance(value, (list, tuple)): - return [_unwrap_non_tensor(row) for row in value] - raise TypeError(f"unsupported benchmark payload column: {type(value).__name__}") - - -def _leaf_digest(value: Any) -> Any: - import numpy as np - import torch +def field_byte_digests(payload: Any, fields: list[str]) -> DigestMap: + """Ordered row digests preserving dtype, shape and every raw byte.""" + return {field: tuple(leaf_digests(row) for row in payload_rows(payload.get(field))) for field in fields} - value = _unwrap_non_tensor(value) - if isinstance(value, torch.Tensor): - if value.is_nested: - return ("nested", tuple(_leaf_digest(row) for row in value.unbind())) - contiguous = value.detach().cpu().contiguous() - flat = contiguous.reshape(-1) - raw = flat.view(torch.uint8).numpy().tobytes() if flat.numel() else b"" - return (str(contiguous.dtype), tuple(contiguous.shape), hashlib.sha256(raw).hexdigest()) - if isinstance(value, np.ndarray): - contiguous = np.ascontiguousarray(value) - return ( - f"np.{contiguous.dtype}", - tuple(contiguous.shape), - hashlib.sha256(contiguous.tobytes()).hexdigest(), - ) - if isinstance(value, Mapping): - if any(not isinstance(key, str) for key in value): - raise TypeError("benchmark payload dictionaries require string keys") - return ("dict", tuple((key, _leaf_digest(value[key])) for key in sorted(value))) - if isinstance(value, (list, tuple)): - return (type(value).__name__, tuple(_leaf_digest(item) for item in value)) - if isinstance(value, bytes): - raw = value - elif isinstance(value, str): - raw = value.encode("utf-8") - elif isinstance(value, float): - raw = struct.pack("!d", value) - elif isinstance(value, (bool, int)) or value is None: - raw = repr(value).encode("utf-8") - else: - raise TypeError(f"unsupported benchmark payload leaf: {type(value).__name__}") - return (f"py.{type(value).__name__}", (), hashlib.sha256(raw).hexdigest()) +def payload_bytes(payload: Any) -> int: + return sum(payload_nbytes(payload.get(field)) for field in payload.keys()) -def field_byte_digests(payload, fields: list[str]) -> dict[str, tuple[Any, ...]]: - """Ordered row digests preserving dtype, shape and every raw byte.""" - digests: dict[str, tuple[Any, ...]] = {} - for field in fields: - digests[field] = tuple(_leaf_digest(row) for row in _column_rows(payload.get(field))) - return digests +def _read_required_counter(path: Path, label: str, *, scale: int = 1) -> int: + try: + value = int(path.read_text().strip()) * scale + except (OSError, ValueError): + raise RuntimeError(f"{label} receive counter is unavailable or unreadable") from None + if value < 0: + raise RuntimeError(f"{label} receive counter returned a negative value") + return value -def _value_bytes(value: Any) -> int: - import numpy as np - import torch - value = _unwrap_non_tensor(value) - if isinstance(value, torch.Tensor): - if value.is_nested: - return sum(_value_bytes(row) for row in value.unbind()) - return value.numel() * value.element_size() - if isinstance(value, np.ndarray): - return value.nbytes - if isinstance(value, Mapping): - return sum(_value_bytes(item) for item in value.values()) - if isinstance(value, (list, tuple)): - return sum(_value_bytes(item) for item in value) - if isinstance(value, bytes): - return len(value) - if isinstance(value, str): - return len(value.encode("utf-8")) - if isinstance(value, float): - return 8 - if isinstance(value, (bool, int)) or value is None: - return len(repr(value).encode("utf-8")) - raise TypeError(f"unsupported benchmark payload leaf: {type(value).__name__}") - - -def payload_bytes(payload) -> int: - return sum(_value_bytes(row) for field in payload.keys() for row in _column_rows(payload.get(field))) - - -def read_counters(tcp_device: str, rdma_device: str = "", sysfs_root: Path = Path("/sys/class")) -> dict[str, int]: - """Read receive-byte counters on a selected HCA and TCP interface.""" - counters: dict[str, int] = {} - ib_root = sysfs_root / "infiniband" - if ib_root.is_dir(): - pattern = ( - f"{rdma_device}/ports/*/counters/port_rcv_data" if rdma_device else "*/ports/*/counters/port_rcv_data" - ) - for counter_path in sorted(ib_root.glob(pattern)): - try: - counters[f"ib:{counter_path.parents[3].name}:{counter_path.parents[1].name}"] = ( - int(counter_path.read_text().strip()) * 4 - ) - except (OSError, ValueError): - continue +def read_counters( + tcp_device: str, + rdma_device: str = "", + rdma_port: int | None = None, + sysfs_root: Path = Path("/sys/class"), +) -> CounterMap: + """Read exact receive-byte counters or fail closed. + + ``port_rcv_data`` is expressed in four-byte words. RDMA device and port + are an all-or-nothing pair so no glob or multi-port aggregation can enter + an acceptance result. + """ + if not _is_safe_device_name(tcp_device): + raise ValueError("TCP counter selection requires one safe interface name") + if bool(rdma_device) != (rdma_port is not None): + raise ValueError("RDMA counter selection requires both device and port") + if rdma_device and (not _is_safe_device_name(rdma_device) or rdma_port is None or rdma_port <= 0): + raise ValueError("RDMA counter selection requires one safe device and positive port") + + counters: CounterMap = {} + if rdma_device: + ib_path = sysfs_root / "infiniband" / rdma_device / "ports" / str(rdma_port) / "counters" / "port_rcv_data" + counters[f"ib:{rdma_device}:{rdma_port}"] = _read_required_counter(ib_path, "RDMA", scale=4) tcp_path = sysfs_root / "net" / tcp_device / "statistics" / "rx_bytes" - try: - counters[f"tcp:{tcp_device}"] = int(tcp_path.read_text().strip()) - except (OSError, ValueError): - pass + counters[f"tcp:{tcp_device}"] = _read_required_counter(tcp_path, "TCP") return counters -def _counter_delta(before: dict[str, int], after: dict[str, int], prefix: str) -> int: - return sum(after[key] - before.get(key, after[key]) for key in after if key.startswith(prefix)) +def _counter_delta(before: CounterMap, after: CounterMap, prefix: str) -> int: + before_keys = {key for key in before if key.startswith(prefix)} + after_keys = {key for key in after if key.startswith(prefix)} + if before_keys != after_keys: + raise RuntimeError(f"{prefix.rstrip(':').upper()} counter set changed during the measured round") + deltas = [after[key] - before[key] for key in sorted(before_keys)] + if any(delta < 0 for delta in deltas): + raise RuntimeError(f"{prefix.rstrip(':').upper()} counter reset or wrapped during the measured round") + return sum(deltas) + + +def _subtract_idle_noise(raw_delta: int, idle_delta: int, idle_seconds: float, round_seconds: float) -> int: + if idle_seconds <= 0 or round_seconds <= 0: + raise RuntimeError("counter sampling durations must be positive") + estimated_noise = math.ceil(idle_delta * round_seconds / idle_seconds) + return max(0, raw_delta - estimated_noise) def wire_is_proven(protocol: str, ib_bytes: int, tcp_bytes: int, payload_nbytes: int) -> bool: if protocol == "rdma": - return ib_bytes >= _RDMA_WIRE_MIN_PAYLOAD_RATIO * payload_nbytes and ib_bytes >= tcp_bytes + return ib_bytes >= _RDMA_WIRE_MIN_PAYLOAD_RATIO * payload_nbytes if protocol == "tcp": - return tcp_bytes >= _TCP_WIRE_MIN_PAYLOAD_RATIO * payload_nbytes and tcp_bytes >= ib_bytes + return tcp_bytes >= _TCP_WIRE_MIN_PAYLOAD_RATIO * payload_nbytes # SimpleStorage may place some units on the consumer node, so C0 cannot - # require a payload-volume ratio. It still must not look like RDMA. - return tcp_bytes > 0 and tcp_bytes >= ib_bytes + # require a payload-volume ratio. It must still produce cross-node TCP. + return tcp_bytes > 0 -def build_conf(protocol: str, master: str, device: str, segment_gib: int): +def build_conf(protocol: str, master: str, device: str, segment_gib: int) -> Any: from omegaconf import OmegaConf from transfer_queue import GRPOGroupNSampler @@ -312,44 +277,78 @@ def build_conf(protocol: str, master: str, device: str, segment_gib: int): ) -def require_clean_cluster() -> None: - try: - ray.get_actor("TransferQueueController", namespace="transfer_queue") - except ValueError: - return - raise RuntimeError("TransferQueueController already exists; benchmark requires a clean exclusive Ray cluster") - - -def close_tq_unmount_and_wait() -> None: - import transfer_queue as tq +def _require_commit_id(value: Any, label: str) -> str: + commit = value if isinstance(value, str) else "" + if len(commit) != 40 or any(character not in "0123456789abcdefABCDEF" for character in commit): + raise RuntimeError(f"{label} commit is unavailable or not a full SHA") + return commit.lower() - from relax.utils.tq.lifecycle import kill_tq_controller_and_wait - store_client = None +def _git_output(repo_root: Path, *args: str) -> str: try: - store_client = getattr(tq.get_client().storage_manager, "storage_client", None) - except (AssertionError, AttributeError): - pass + result = subprocess.run( + ["git", "-C", str(repo_root), *args], + capture_output=True, + text=True, + timeout=5, + check=False, + ) + except (OSError, subprocess.SubprocessError): + raise RuntimeError("Relax commit provenance is unavailable") from None + if result.returncode != 0: + raise RuntimeError("Relax commit provenance is unavailable") + return result.stdout.strip() + + +def collect_provenance(repo_root: Path | None = None) -> dict[str, str]: + """Collect exact source/dependency versions for every CSV row. + + Acceptance runs require a clean tracked Relax checkout. Untracked result + files are ignored so the benchmark can write artifacts below the checkout. + TransferQueue's PEP 610 metadata supplies the installed VCS commit. + """ + if repo_root is None: + repo_root = Path(__file__).resolve().parents[2] + if _git_output(repo_root, "status", "--porcelain", "--untracked-files=no"): + raise RuntimeError("benchmark acceptance requires a clean tracked Relax checkout") + revision = _git_output(repo_root, "rev-parse", "HEAD") + + distribution = importlib_metadata.distribution("transferqueue") + direct_url_text = distribution.read_text("direct_url.json") + if direct_url_text is None: + raise RuntimeError("installed TransferQueue has no VCS provenance metadata") try: - tq.close() - finally: - try: - if store_client is not None and hasattr(store_client, "close"): - store_client.close() - finally: - kill_tq_controller_and_wait() + direct_url = json.loads(direct_url_text) + tq_commit = direct_url["vcs_info"]["commit_id"] + except (KeyError, TypeError, json.JSONDecodeError): + raise RuntimeError("installed TransferQueue VCS provenance metadata is invalid") from None + try: + mooncake_version = importlib_metadata.version("mooncake-transfer-engine") + except importlib_metadata.PackageNotFoundError: + mooncake_version = "not-installed" + + return { + "relax_sha": _require_commit_id(revision, "Relax"), + "tq_commit": _require_commit_id(tq_commit, "TransferQueue"), + "mooncake_version": mooncake_version, + } @ray.remote(num_cpus=0.001, max_restarts=0) class TQConsumer: """Persistent, hard-pinned consumer mirroring a Relax component actor.""" - def __init__(self, conf: Any, tcp_device: str, rdma_device: str, protocol: str): - counters = read_counters(tcp_device, rdma_device if protocol == "rdma" else "") - if f"tcp:{tcp_device}" not in counters: - raise RuntimeError("TCP receive counter is unavailable on the selected consumer node") - if protocol == "rdma" and not any(key.startswith("ib:") for key in counters): - raise RuntimeError("RDMA receive counter is unavailable on the selected consumer node") + def __init__( + self, + conf: Any, + tcp_device: str, + rdma_device: str, + rdma_port: int | None, + protocol: str, + ) -> None: + selected_device = rdma_device if protocol == "rdma" else "" + selected_port = rdma_port if protocol == "rdma" else None + read_counters(tcp_device, selected_device, selected_port) if protocol == "tcp": os.environ["MC_TCP_ENABLE_CONNECTION_POOL"] = "1" from relax.utils.tq.lifecycle import attach_tq_client @@ -358,15 +357,43 @@ def __init__(self, conf: Any, tcp_device: str, rdma_device: str, protocol: str): # Mooncake correctness guards in this process before client creation. self.client = attach_tq_client(conf, role="benchmark-consumer") self.tcp_device = tcp_device - self.rdma_device = rdma_device if protocol == "rdma" else "" + self.rdma_device = selected_device + self.rdma_port = selected_port def describe(self) -> tuple[str, str]: manager = self.client.storage_manager storage_client = getattr(manager, "storage_client", None) return type(manager).__name__, getattr(storage_client, "protocol", "") - def fetch(self, fields: list[str], batch_size: int, partition: str, expected: dict) -> dict[str, Any]: - before = read_counters(self.tcp_device, self.rdma_device) + def counters(self) -> CounterMap: + return read_counters(self.tcp_device, self.rdma_device, self.rdma_port) + + def sample_idle_counters(self, seconds: float) -> dict[str, Any]: + if seconds <= 0: + raise ValueError("idle counter sample duration must be positive") + before = self.counters() + started = time.perf_counter() + time.sleep(seconds) + after = self.counters() + elapsed = time.perf_counter() - started + return { + "seconds": elapsed, + "ib_bytes": _counter_delta(before, after, "ib:"), + "tcp_bytes": _counter_delta(before, after, "tcp:"), + } + + def begin_round(self) -> tuple[CounterMap, float]: + return self.counters(), time.perf_counter() + + def fetch( + self, + fields: list[str], + batch_size: int, + partition: str, + expected: DigestMap, + before: CounterMap, + counter_started: float, + ) -> dict[str, Any]: started = time.perf_counter() meta = self.client.get_meta( data_fields=fields, @@ -377,15 +404,17 @@ def fetch(self, fields: list[str], batch_size: int, partition: str, expected: di ) received = self.client.get_data(meta) get_ms = (time.perf_counter() - started) * 1000 - after = read_counters(self.tcp_device, self.rdma_device) + after = self.counters() + round_seconds = time.perf_counter() - counter_started actual = field_byte_digests(received, fields) - if actual != expected: - mismatches = [field for field in fields if actual.get(field) != expected.get(field)] - raise AssertionError(f"byte-exact mismatch after TQ get: fields={mismatches}") + mismatches = [field for field in fields if actual.get(field) != expected.get(field)] return { "get_ms": get_ms, "ib_bytes": _counter_delta(before, after, "ib:"), "tcp_bytes": _counter_delta(before, after, "tcp:"), + "round_seconds": round_seconds, + "byte_exact": not mismatches, + "mismatch_fields": mismatches, } def shutdown(self) -> None: @@ -403,37 +432,184 @@ def _validate_attached_backend(protocol: str, manager: str, attached_protocol: s ) -def _teardown_benchmark(consumer: Any, owner_attempted: bool) -> None: - """Unmount the consumer, clean owned global state, then stop local Ray.""" +def _throughput_gbs(nbytes: int, milliseconds: float, label: str) -> float: + if milliseconds <= 0: + raise RuntimeError(f"{label} duration must be positive") + return nbytes / milliseconds / 1e6 + + +def _write_csv_record(writer: csv.DictWriter, csv_handle: TextIO, record: dict[str, Any]) -> None: + writer.writerow(record) + csv_handle.flush() + + +def _run_round( + *, + producer: Any, + consumer: Any, + payload: Any, + fields: list[str], + expected: DigestMap, + nbytes: int, + protocol: str, + profile: str, + requested_mib: int, + run: int, + writer: csv.DictWriter, + csv_handle: TextIO, + provenance: dict[str, str], +) -> dict[str, float]: + """Measure, persist and clean one warmup or acceptance round.""" + partition = f"bench-{profile}-{requested_mib}-{run}" + record: dict[str, Any] = {column: "" for column in CSV_COLUMNS} + record.update( + { + "protocol": protocol, + "profile": profile, + "payload_mib": requested_mib, + "actual_mib": round(nbytes / 1024**2, 2), + "run": run, + "measured": run > 0, + **provenance, + } + ) try: - if consumer is not None: - try: - ray.get(consumer.shutdown.remote(), timeout=10) - finally: - ray.kill(consumer, no_restart=True) - finally: + idle = ray.get(consumer.sample_idle_counters.remote(_IDLE_COUNTER_SAMPLE_SECONDS)) + before, counter_started = ray.get(consumer.begin_round.remote()) + started = time.perf_counter() + producer.put(payload, partition_id=partition) + put_ms = (time.perf_counter() - started) * 1000 + fetched = ray.get( + consumer.fetch.remote( + fields, + payload.batch_size[0], + partition, + expected, + before, + counter_started, + ) + ) + round_seconds = fetched["round_seconds"] + proof_ib_bytes = _subtract_idle_noise(fetched["ib_bytes"], idle["ib_bytes"], idle["seconds"], round_seconds) + proof_tcp_bytes = _subtract_idle_noise(fetched["tcp_bytes"], idle["tcp_bytes"], idle["seconds"], round_seconds) + proven = wire_is_proven(protocol, proof_ib_bytes, proof_tcp_bytes, nbytes) + put_gbs = _throughput_gbs(nbytes, put_ms, "put") + get_gbs = _throughput_gbs(nbytes, fetched["get_ms"], "get") + + if not fetched["byte_exact"]: + status, error_kind = "fail", "ByteExactMismatch" + gate_error: BaseException | None = AssertionError("byte-exact mismatch after TQ get") + elif not proven: + status, error_kind = "fail", "WireProofFailed" + gate_error = RuntimeError( + f"wire proof failed for protocol={protocol} profile={profile} payload={requested_mib}MiB run={run}" + ) + else: + status, error_kind = "pass", "" + gate_error = None + record.update( + { + "status": status, + "error_kind": error_kind, + "byte_exact": fetched["byte_exact"], + "mismatch_fields": ";".join(fetched["mismatch_fields"]), + "wire_proven": proven, + "put_ms": round(put_ms, 2), + "get_ms": round(fetched["get_ms"], 2), + "round_ms": round(round_seconds * 1000, 2), + "put_gbs": round(put_gbs, 3), + "get_gbs": round(get_gbs, 3), + "ib_mb": round(fetched["ib_bytes"] / 1e6, 1), + "tcp_mb": round(fetched["tcp_bytes"] / 1e6, 1), + "proof_ib_mb": round(proof_ib_bytes / 1e6, 1), + "proof_tcp_mb": round(proof_tcp_bytes / 1e6, 1), + "idle_ib_mbps": round(idle["ib_bytes"] / idle["seconds"] / 1e6, 3), + "idle_tcp_mbps": round(idle["tcp_bytes"] / idle["seconds"] / 1e6, 3), + } + ) + except BaseException as error: + record.update({"status": "error", "error_kind": type(error).__name__}) try: - if owner_attempted: - close_tq_unmount_and_wait() - finally: - ray.shutdown() + producer.clear_partition(partition) + except BaseException as cleanup_error: + _write_csv_record(writer, csv_handle, record) + raise error from cleanup_error + _write_csv_record(writer, csv_handle, record) + raise + + try: + producer.clear_partition(partition) + except BaseException as cleanup_error: + if gate_error is None: + record.update({"status": "error", "error_kind": type(cleanup_error).__name__}) + _write_csv_record(writer, csv_handle, record) + if gate_error is not None: + raise gate_error from cleanup_error + raise + _write_csv_record(writer, csv_handle, record) + if gate_error is not None: + raise gate_error + return {"put_gbs": put_gbs, "get_gbs": get_gbs} + + +def _teardown_benchmark(consumer: Any, producer_attached: bool, owner: Any) -> None: + """Attempt every cleanup step and propagate the first teardown error.""" + from relax.utils.tq.lifecycle import close_tq_owner, detach_tq_client + + first_error: BaseException | None = None + + def remember(error: BaseException) -> None: + nonlocal first_error + if first_error is None: + first_error = error + + if consumer is not None: + try: + ray.get(consumer.shutdown.remote(), timeout=10) + except BaseException as error: + remember(error) + try: + ray.kill(consumer, no_restart=True) + except BaseException as error: + remember(error) + if producer_attached: + try: + detach_tq_client() + except BaseException as error: + remember(error) + if owner is not None: + try: + close_tq_owner(owner) + except BaseException as error: + remember(error) + try: + ray.shutdown() + except BaseException as error: + remember(error) + if first_error is not None: + raise first_error def main() -> None: args = parse_args() + provenance = collect_provenance() + runtime_env: dict[str, Any] = {} if args.protocol == "tcp": os.environ["MC_TCP_ENABLE_CONNECTION_POOL"] = "1" - import transfer_queue as tq + runtime_env = {"env_vars": {"MC_TCP_ENABLE_CONNECTION_POOL": "1"}} from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy - with open(args.csv, "w", newline="") as csv_handle: + from relax.utils.tq.lifecycle import attach_tq_client, initialize_tq_with_fallback + + with args.csv.open("x", newline="") as csv_handle: writer = csv.DictWriter(csv_handle, fieldnames=CSV_COLUMNS) writer.writeheader() csv_handle.flush() - ray.init(address="auto", ignore_reinit_error=True, logging_level="ERROR") + ray.init(address="auto", logging_level="ERROR", runtime_env=runtime_env) consumer = None + owner = None + producer_attached = False measured_runs = 0 - owner_attempted = False try: driver_node_id = ray.get_runtime_context().get_node_id() alive_ids = {node["NodeID"] for node in ray.nodes() if node.get("Alive")} @@ -441,15 +617,19 @@ def main() -> None: raise RuntimeError("--consumer-node-id is not an alive Ray node") if args.consumer_node_id == driver_node_id: raise RuntimeError("producer and consumer must run on different Ray nodes") - require_clean_cluster() conf = build_conf(args.protocol, args.master, args.device, args.segment_gib) - owner_attempted = True - tq.init(conf=conf) - producer = tq.get_client() + init_result = initialize_tq_with_fallback(conf, mode="required") + owner = init_result.owner + producer = attach_tq_client(init_result.config, role="benchmark-producer") + producer_attached = True strategy = NodeAffinitySchedulingStrategy(node_id=args.consumer_node_id, soft=False) consumer = TQConsumer.options(scheduling_strategy=strategy).remote( - conf, args.tcp_device, args.device, args.protocol + init_result.config, + args.tcp_device, + args.device, + args.rdma_port, + args.protocol, ) manager, attached_protocol = ray.get(consumer.describe.remote()) _validate_attached_backend(args.protocol, manager, attached_protocol) @@ -468,43 +648,26 @@ def main() -> None: put_rates: list[float] = [] get_rates: list[float] = [] for run in range(args.repeats + 1): - partition = f"bench-{profile}-{requested_mib}-{run}" - started = time.perf_counter() - producer.put(payload, partition_id=partition) - put_ms = (time.perf_counter() - started) * 1000 - fetched = ray.get(consumer.fetch.remote(fields, payload.batch_size[0], partition, expected)) - producer.clear_partition(partition) + rates = _run_round( + producer=producer, + consumer=consumer, + payload=payload, + fields=fields, + expected=expected, + nbytes=nbytes, + protocol=args.protocol, + profile=profile, + requested_mib=requested_mib, + run=run, + writer=writer, + csv_handle=csv_handle, + provenance=provenance, + ) if run == 0: continue - proven = wire_is_proven(args.protocol, fetched["ib_bytes"], fetched["tcp_bytes"], nbytes) - put_gbs = nbytes / put_ms / 1e6 - get_gbs = nbytes / fetched["get_ms"] / 1e6 - put_rates.append(put_gbs) - get_rates.append(get_gbs) - writer.writerow( - { - "protocol": args.protocol, - "profile": profile, - "payload_mib": requested_mib, - "actual_mib": round(nbytes / 1024**2, 2), - "run": run, - "byte_exact": True, - "wire_proven": proven, - "put_ms": round(put_ms, 2), - "get_ms": round(fetched["get_ms"], 2), - "put_gbs": round(put_gbs, 3), - "get_gbs": round(get_gbs, 3), - "ib_mb": round(fetched["ib_bytes"] / 1e6, 1), - "tcp_mb": round(fetched["tcp_bytes"] / 1e6, 1), - } - ) - csv_handle.flush() + put_rates.append(rates["put_gbs"]) + get_rates.append(rates["get_gbs"]) measured_runs += 1 - if not proven: - raise RuntimeError( - f"wire proof failed for protocol={args.protocol} " - f"profile={profile} payload={requested_mib}MiB run={run}" - ) print( f"[{profile} {requested_mib}MiB] byte_exact=PASS wire=PASS " f"put={statistics.mean(put_rates):.2f}GB/s " @@ -514,8 +677,14 @@ def main() -> None: flush=True, ) print(f"[csv] wrote {measured_runs} measured rounds", flush=True) - finally: - _teardown_benchmark(consumer, owner_attempted) + except BaseException as primary_error: + try: + _teardown_benchmark(consumer, producer_attached, owner) + except BaseException as cleanup_error: + raise primary_error from cleanup_error + raise + else: + _teardown_benchmark(consumer, producer_attached, owner) if __name__ == "__main__": diff --git a/tests/core/test_controller_tq_backend.py b/tests/core/test_controller_tq_backend.py index 439d312ca..0b7d51ea7 100644 --- a/tests/core/test_controller_tq_backend.py +++ b/tests/core/test_controller_tq_backend.py @@ -1,456 +1,277 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. -"""Controller._resolve_tq_backend: the production backend decision. - -After the static ``/sys`` probe was removed, this method makes a pure -*configuration* decision: which mode was requested, whether the TransferQueue -correctness contract holds, whether a master endpoint is configured, and whether -the segment can hold the worst-case payload. Host-RDMA capability itself is -established later by the real cluster-wide attach handshake -(``tests/utils/test_tq_failure_paths.py``). - -These tests pin the branch structure: ``off`` must check nothing, ``auto`` must -converge on SimpleStorage for every unmet precondition, ``required`` must fail -fast on the same ones, a malformed mode must always raise, and no input may ever -produce a MooncakeStore config whose protocol is not ``rdma``. -""" +"""Controller decisions layered on top of the TQ config contracts.""" from __future__ import annotations import argparse +from contextlib import nullcontext from types import SimpleNamespace +from typing import Any import pytest from tests.core.test_controller_s3_model_cleanup import controller -from tests.utils.test_arguments_opd_teacher_colocate import ( - arguments_module as _arguments_module_fixture, -) +from tests.utils.test_arguments_opd_teacher_colocate import arguments_module as _arguments_module_fixture -# ``relax.utils.arguments`` pulls in the full SGLang server args, which the -# GitHub CPU CI does not install (it ships only ``sglang-router``). Reuse the -# stub fixture so the CLI assertions build a real parser without that -# dependency. arguments_module = _arguments_module_fixture - _MASTER = "master.invalid:50051" -def _config(**overrides) -> SimpleNamespace: - """A config whose worst-case payload fits the default 4 GiB segment.""" - defaults = dict( - tq_rdma_mode="auto", - tq_rdma_device="", - num_data_storage_units=1, - n_samples_per_prompt=1, - rollout_batch_size=8, - seq_length=8192, - multimodal_keys=None, - max_staleness=0, - partial_rollout=False, - use_dynamic_global_batch_size=False, - ) - defaults.update(overrides) - return SimpleNamespace(**defaults) - - -class _Recorder: - """Records which configuration preconditions the Controller consulted. - - A recorded ``probe`` entry would mean a static capability probe crept back - in; the sequence assertions below are what keep that from happening - silently. - """ - - def __init__(self, monkeypatch, *, contract_error=None, master_error=None): - self.calls: list[str] = [] - - def fake_contract() -> None: - self.calls.append("contract") - if contract_error is not None: - raise contract_error - - def fake_master() -> str: - self.calls.append("master") - if master_error is not None: - raise master_error - return _MASTER - - monkeypatch.setattr(controller, "validate_mooncake_runtime_contract", fake_contract) - monkeypatch.setattr(controller, "resolve_mooncake_master_address", fake_master) +def _config(**overrides: Any) -> SimpleNamespace: + values = { + "tq_rdma_mode": "auto", + "tq_rdma_device": "", + "num_data_storage_units": 1, + "n_samples_per_prompt": 1, + "rollout_batch_size": 8, + "seq_length": 8192, + "multimodal_keys": None, + "max_staleness": 0, + } + values.update(overrides) + return SimpleNamespace(**values) -def _resolve(config) -> dict: +def _resolve(config: SimpleNamespace) -> dict[str, Any]: instance = controller.Controller.__new__(controller.Controller) instance.config = config return instance._resolve_tq_backend(total_storage_size=64) -def _assert_simple_storage(backend: dict) -> None: - assert backend["storage_backend"] == "SimpleStorage" - assert "MooncakeStore" not in backend - - -def _raise(error: BaseException) -> None: - raise error - - -class TestOffMode: - """``off`` is the untouched SimpleStorage path: it checks nothing.""" +def _is_simple(backend: dict[str, Any]) -> bool: + return backend["storage_backend"] == "SimpleStorage" - def test_off_short_circuits_without_any_precondition_check(self, monkeypatch): - recorder = _Recorder(monkeypatch) - _assert_simple_storage(_resolve(_config(tq_rdma_mode="off"))) - assert recorder.calls == [] - def test_missing_mode_attribute_defaults_to_off(self, monkeypatch): - recorder = _Recorder(monkeypatch) - config = _config() - del config.tq_rdma_mode - _assert_simple_storage(_resolve(config)) - assert recorder.calls == [] - - def test_healthy_existing_controller_aborts_before_legacy_init(self, monkeypatch): - """The default path must not attach to or later close another job. +class _DecisionHarness: + def __init__(self, monkeypatch: pytest.MonkeyPatch, failure: str | None) -> None: + self.calls: list[str] = [] - Upstream ``tq.init`` ignores the caller's SimpleStorage config when a - named controller already exists. The exclusive-cluster check must - therefore fail before setting the local ownership flag or calling it. - """ - config = _config( - tq_rdma_mode="off", - fully_async=False, - balance_data=False, - polling_mode=False, - ) - instance = controller.Controller.__new__(controller.Controller) - instance.config = config - instance._tq_owner = None - instance._tq_legacy_init = False - - monkeypatch.setattr(controller, "resolve_sft_algo_key", lambda _config: "grpo") - monkeypatch.setattr(controller, "resolve_tq_capacity_batch_size", lambda _config: 1) - monkeypatch.setattr(controller, "GRPOGroupNSampler", lambda **_kwargs: object()) - monkeypatch.setattr( - instance, - "_resolve_tq_backend", - lambda _total_storage_size: { - "storage_backend": "SimpleStorage", - "SimpleStorage": {"total_storage_size": 1, "num_data_storage_units": 1}, - }, - ) - monkeypatch.setattr( - controller, - "reap_unusable_tq_controller", - lambda: _raise(RuntimeError("exclusive cluster is not clean")), - ) - monkeypatch.setattr( - controller.tq, - "init", - lambda **_kwargs: pytest.fail("existing controller must be rejected before tq.init"), - ) - - with pytest.raises(RuntimeError, match="exclusive cluster"): - instance._initialize_data_system() - - assert instance._tq_owner is None - assert instance._tq_legacy_init is False - - -class TestNoStaticCapabilityProbe: - """The ``/sys``/GID/memlock probe is gone and must not return.""" - - def test_resolver_never_probes_hardware(self, monkeypatch): - recorder = _Recorder(monkeypatch) - _resolve(_config()) - assert recorder.calls == ["contract", "master"] - - -class TestBackendPreconditions: - """The same unmet precondition degrades ``auto`` and aborts - ``required``.""" - - @pytest.mark.parametrize("mode", ["auto", "required"]) - @pytest.mark.parametrize( - ("recorder_kwargs", "config_overrides", "segment_override", "match"), - [ - ({"contract_error": RuntimeError("retry guard missing")}, {}, None, "correctness contract"), - ({"master_error": RuntimeError("MC_MASTER_ADDRESS required")}, {}, None, "master endpoint"), - ({"master_error": RuntimeError("unusable master endpoint")}, {}, None, "master endpoint"), - ( - {}, - {"multimodal_keys": ["pixel_values"], "rollout_batch_size": 64, "n_samples_per_prompt": 8}, - None, - "segment capacity insufficient", - ), - ({}, {}, "four", "segment-capacity configuration is unusable"), - ({}, {}, "nan", "segment-capacity configuration is unusable"), - ({}, {}, "inf", "segment-capacity configuration is unusable"), - ({}, {}, "-inf", "segment-capacity configuration is unusable"), - ({}, {"seq_length": None}, None, "segment-capacity configuration is unusable"), - ], - ids=[ - "correctness-contract", - "missing-master", - "malformed-master", - "insufficient-capacity", - "invalid-segment-size", - "nan-segment-size", - "infinite-segment-size", - "negative-infinite-segment-size", - "missing-seq-length", - ], - ) - def test_unmet_precondition( - self, - monkeypatch, - mode, - recorder_kwargs, - config_overrides, - segment_override, - match, - ): - recorder = _Recorder(monkeypatch, **recorder_kwargs) - if segment_override is not None: - monkeypatch.setenv("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB", segment_override) - - config = _config(tq_rdma_mode=mode, **config_overrides) - if mode == "auto": - _assert_simple_storage(_resolve(config)) - else: - with pytest.raises(RuntimeError, match=match): - _resolve(config) - assert recorder.calls == (["contract"] if recorder_kwargs.get("contract_error") else ["contract", "master"]) - - def test_satisfied_preconditions_select_host_rdma(self, monkeypatch): - _Recorder(monkeypatch) - private_device = "private-device-name" - messages: list[str] = [] - monkeypatch.setattr(controller.logger, "info", lambda message: messages.append(message)) - - backend = _resolve(_config(tq_rdma_device=private_device)) - assert backend["storage_backend"] == "MooncakeStore" - assert backend["MooncakeStore"]["protocol"] == "rdma" - assert backend["MooncakeStore"]["device_name"] == private_device - assert any("device_selection=explicit" in message for message in messages) - assert all(private_device not in message for message in messages) - - def test_required_accepts_satisfied_preconditions(self, monkeypatch): - _Recorder(monkeypatch) - backend = _resolve(_config(tq_rdma_mode="required")) - assert backend["MooncakeStore"]["protocol"] == "rdma" - - def test_validated_master_endpoint_reaches_the_client_config(self, monkeypatch): - """The checked endpoint, rather than a later env value, reaches - Mooncake.""" - _Recorder(monkeypatch) - monkeypatch.setenv("MC_MASTER_ADDRESS", "someone.else.invalid:9999") - backend = _resolve(_config()) - assert backend["MooncakeStore"]["master_server_address"] == _MASTER - - -class TestModeValidation: - """A malformed mode is a configuration error in every mode.""" - - def test_invalid_mode_always_raises_instead_of_falling_back(self, monkeypatch): - recorder = _Recorder(monkeypatch) - with pytest.raises(ValueError, match="--tq-rdma-mode"): - _resolve(_config(tq_rdma_mode="mooncake")) - assert recorder.calls == [] - - @pytest.mark.parametrize("device", [None, ["rdma0"], "rdma0\nforged", " "]) - def test_invalid_device_always_raises_before_preconditions(self, monkeypatch, device): - recorder = _Recorder(monkeypatch) - with pytest.raises(ValueError, match="--tq-rdma-device"): - _resolve(_config(tq_rdma_device=device)) - assert recorder.calls == [] - - -class TestProductionNeverSelectsMooncakeTcp: - """Mooncake/TCP exists only as benchmark C1.""" - - @pytest.mark.parametrize("mode", ["off", "auto", "required"]) - def test_selected_backend_is_rdma_or_simple(self, monkeypatch, mode): - _Recorder(monkeypatch) - backend = _resolve(_config(tq_rdma_mode=mode)) - if backend["storage_backend"] == "MooncakeStore": - assert backend["MooncakeStore"]["protocol"] == "rdma" - else: - _assert_simple_storage(backend) + def contract() -> None: + self.calls.append("contract") + if failure == "contract": + raise RuntimeError("contract unavailable") + def master() -> str: + self.calls.append("master") + if failure == "master": + raise RuntimeError("master unavailable") + return _MASTER -def test_cli_exposes_only_mode_and_device(arguments_module): - """The narrowed CLI keeps exactly two TransferQueue RDMA flags.""" + def backend(_args: Any, *, device: str, master_address: str, total_storage_size: int): + self.calls.append(f"build:{device}:{master_address}:{total_storage_size}") + if failure == "capacity-error": + return {"storage_backend": "SimpleStorage"}, "capacity insufficient" + if failure == "capacity-config": + raise RuntimeError("capacity configuration unusable") + return { + "storage_backend": "MooncakeStore", + "MooncakeStore": {"protocol": "rdma", "device_name": device, "master_server_address": master_address}, + }, None + + monkeypatch.setattr(controller, "validate_mooncake_runtime_contract", contract) + monkeypatch.setattr(controller, "resolve_mooncake_master_address", master) + monkeypatch.setattr(controller, "build_backend_config", backend) + + +@pytest.mark.parametrize( + ("mode", "failure", "expect_simple"), + [("off", None, True), ("auto", None, False), ("required", None, False)] + + [("auto", failure, True) for failure in ("contract", "master", "capacity-error", "capacity-config")] + + [("required", failure, False) for failure in ("contract", "master", "capacity-error", "capacity-config")], +) +def test_backend_decision_matrix( + monkeypatch: pytest.MonkeyPatch, mode: str, failure: str | None, expect_simple: bool +) -> None: + harness = _DecisionHarness(monkeypatch, failure) + context = pytest.raises(RuntimeError) if mode == "required" and failure else nullcontext() + with context: + backend = _resolve(_config(tq_rdma_mode=mode, tq_rdma_device="rdma0")) + assert _is_simple(backend) is expect_simple + if not expect_simple: + assert backend["MooncakeStore"] == { + "protocol": "rdma", + "device_name": "rdma0", + "master_server_address": _MASTER, + } + if mode == "off": + assert harness.calls == [] + elif failure == "contract": + assert harness.calls == ["contract"] + elif failure == "master": + assert harness.calls == ["contract", "master"] + else: + assert harness.calls == ["contract", "master", f"build:rdma0:{_MASTER}:64"] + + +@pytest.mark.parametrize("overrides", [{"tq_rdma_mode": "mooncake"}, {"tq_rdma_device": "bad\ndevice"}]) +def test_invalid_config_fails_before_preconditions(monkeypatch: pytest.MonkeyPatch, overrides: dict[str, Any]) -> None: + harness = _DecisionHarness(monkeypatch, None) + with pytest.raises(ValueError, match="TransferQueue RDMA configuration"): + _resolve(_config(**overrides)) + assert harness.calls == [] + + +def test_missing_mode_defaults_to_off(monkeypatch: pytest.MonkeyPatch) -> None: + harness = _DecisionHarness(monkeypatch, None) + config = _config() + del config.tq_rdma_mode + assert _is_simple(_resolve(config)) + assert harness.calls == [] + + +def test_cli_exposes_only_mode_and_device(arguments_module: Any) -> None: arguments_module.RouterArgs = SimpleNamespace(add_cli_args=lambda parser, **_kwargs: parser) parser = argparse.ArgumentParser() arguments_module.get_slime_extra_args_provider()(parser) - tq_flags = sorted( option for action in parser._actions for option in action.option_strings if option.startswith("--tq-") ) assert tq_flags == ["--tq-rdma-device", "--tq-rdma-mode"] - args = parser.parse_args([]) - assert args.tq_rdma_mode == "off" - assert args.tq_rdma_device == "" - assert not hasattr(args, "tq_storage_backend") - assert not hasattr(args, "tq_use_gdr") - - -# --------------------------------------------------------------------------- -# _confirm_mooncake_attach: the cleanup evidence chain -# --------------------------------------------------------------------------- + assert (args.tq_rdma_mode, args.tq_rdma_device) == ("off", "") -class _AttachRecorder: - """Orders the teardown steps taken after a failed attach handshake. +def test_off_mode_rejects_a_healthy_existing_controller_before_legacy_init( + monkeypatch: pytest.MonkeyPatch, +) -> None: + instance = controller.Controller.__new__(controller.Controller) + instance.config = _config( + tq_rdma_mode="off", + fully_async=False, + balance_data=False, + polling_mode=False, + ) + instance._tq_owner = None + instance._tq_legacy_init = False + monkeypatch.setattr(controller, "resolve_sft_algo_key", lambda _config: "grpo") + monkeypatch.setattr(controller, "resolve_tq_capacity_batch_size", lambda _config: 1) + monkeypatch.setattr(controller, "GRPOGroupNSampler", lambda **_kwargs: object()) + monkeypatch.setattr(instance, "_resolve_tq_backend", lambda _size: {"storage_backend": "SimpleStorage"}) + monkeypatch.setattr( + controller, + "reap_unusable_tq_controller", + lambda: (_ for _ in ()).throw(RuntimeError("exclusive cluster is not clean")), + ) + monkeypatch.setattr(controller.tq, "init", lambda **_kwargs: pytest.fail("tq.init must not run")) + with pytest.raises(RuntimeError, match="exclusive cluster"): + instance._initialize_data_system() + assert instance._tq_owner is None and instance._tq_legacy_init is False - The static probe used to reject unusable clusters before any Mooncake state - existed. Now the handshake fails *after* an owner and a named controller - were created, so the ordering recorded here is what keeps a half- - initialised controller from surviving into the next ``tq.init`` (F10 hang). - """ +class _AttachHarness: def __init__( self, - monkeypatch, + monkeypatch: pytest.MonkeyPatch, *, - failures=None, - verify_error=None, - close_error=None, - fallback_owner="fallback-owner", - ): + failures: list[str] | None = None, + verify_error: BaseException | None = None, + close_error: BaseException | None = None, + ) -> None: self.events: list[str] = [] - self.closed: list[object] = [] - self.initialized: list[object] = [] - self._fallback_owner = fallback_owner - def fake_verify(conf, **_kwargs): + def verify(_conf: Any, **_kwargs: Any) -> list[str]: self.events.append("handshake") - if verify_error is not None: + if verify_error: raise verify_error - return list(failures or []) + return failures or [] - def fake_close(owner, **_kwargs): + def close(_owner: Any, **_kwargs: Any) -> None: self.events.append("close") - self.closed.append(owner) - if close_error is not None: + if close_error: raise close_error - def fake_initialize(conf, **_kwargs): + def fallback(conf: Any, **_kwargs: Any) -> Any: self.events.append("init_simple") - self.initialized.append(conf) - return controller.TqInitResult(config=conf, owner=self._fallback_owner) + return controller.TqInitResult(config=conf, owner="fallback-owner") - monkeypatch.setattr(controller, "verify_cluster_attach", fake_verify) - monkeypatch.setattr(controller, "close_tq_owner", fake_close) - monkeypatch.setattr(controller, "initialize_tq_with_fallback", fake_initialize) + monkeypatch.setattr(controller, "verify_cluster_attach", verify) + monkeypatch.setattr(controller, "close_tq_owner", close) + monkeypatch.setattr(controller, "initialize_tq_with_fallback", fallback) -def _confirm(config, *, owner="mooncake-owner", fallback_config="simple-conf"): +def _confirm(mode: str = "auto") -> tuple[Any, Any]: instance = controller.Controller.__new__(controller.Controller) - instance.config = config - init_result = controller.TqInitResult(config="mooncake-conf", owner=owner) - return instance._confirm_mooncake_attach(init_result, fallback_config) - - -class TestAttachHandshakeCleanupChain: - def test_success_keeps_mooncake_and_touches_no_cleanup(self, monkeypatch): - recorder = _AttachRecorder(monkeypatch, failures=[]) - result = _confirm(_config()) - assert result.config == "mooncake-conf" - assert result.fallback_reason == "" - assert recorder.events == ["handshake"] - - def test_auto_closes_owner_before_initializing_simple_storage(self, monkeypatch): - recorder = _AttachRecorder(monkeypatch, failures=["node: attach timed out"]) - result = _confirm(_config(), owner="attempt-owner") - assert recorder.events == ["handshake", "close", "init_simple"] - assert recorder.closed == ["attempt-owner"] - assert result.config == "simple-conf" - assert result.fallback_reason == "attach_handshake_failed:1_failures" - - def test_cleanup_failure_aborts_instead_of_falling_back(self, monkeypatch): - recorder = _AttachRecorder( - monkeypatch, - failures=["node: attach timed out"], - close_error=RuntimeError("TransferQueue owner cleanup failed"), - ) - owner = object() - instance = controller.Controller.__new__(controller.Controller) - instance.config = _config() - instance._tq_owner = owner - init_result = controller.TqInitResult(config="mooncake-conf", owner=owner) - - with pytest.raises(RuntimeError, match="owner cleanup failed"): - instance._confirm_mooncake_attach(init_result, "simple-conf") - assert recorder.events == ["handshake", "close"] - assert recorder.initialized == [] - assert instance._tq_owner is owner - - def test_required_closes_owner_then_raises(self, monkeypatch): - recorder = _AttachRecorder(monkeypatch, failures=["node: protocol=tcp"]) - with pytest.raises(RuntimeError, match="attach handshake reported"): - _confirm(_config(tq_rdma_mode="required")) - assert recorder.events == ["handshake", "close"] - assert recorder.initialized == [] - - def test_unexpected_driver_exception_closes_owner_and_is_sanitized(self, monkeypatch): - secret = "worker endpoint and traceback path must stay private" - recorder = _AttachRecorder(monkeypatch, verify_error=RuntimeError(secret)) - with pytest.raises(RuntimeError, match="orchestration failed") as excinfo: - _confirm(_config()) - assert recorder.events == ["handshake", "close"] - assert recorder.initialized == [] - assert secret not in str(excinfo.value) - - def test_unconfirmed_worker_isolation_closes_owner_and_aborts(self, monkeypatch): - recorder = _AttachRecorder( - monkeypatch, - verify_error=controller.TqHandshakeIsolationError("private cancellation detail"), - ) - with pytest.raises(RuntimeError, match="could not be confirmed stopped") as excinfo: - _confirm(_config()) - assert recorder.events == ["handshake", "close"] - assert recorder.initialized == [] - assert "private cancellation detail" not in str(excinfo.value) - - def test_constructor_cleanup_boundary_includes_data_system_initialization(self, monkeypatch): - """An exception after owner creation must still invoke constructor - cleanup.""" - events: list[str] = [] - - monkeypatch.setattr(controller, "resolve_sft_num_rollout", lambda _config: None) - monkeypatch.setattr(controller, "HealthManager", lambda **_kwargs: object()) - - def fail_after_owner_created(instance): - instance._tq_owner = "mooncake-owner" - raise RuntimeError("driver-side handshake orchestration failed") - - def record_cleanup(instance): - events.append(instance._tq_owner) - instance._tq_owner = None - - monkeypatch.setattr(controller.Controller, "_initialize_data_system", fail_after_owner_created) - monkeypatch.setattr(controller.Controller, "_close_data_system", record_cleanup) - - config = SimpleNamespace(use_health_check=False, max_global_restart=3) - with pytest.raises(RuntimeError, match="handshake orchestration failed"): - controller.Controller(config) - - assert events == ["mooncake-owner"] - - def test_repeated_data_system_cleanup_closes_the_owner_once(self, monkeypatch): - instance = controller.Controller.__new__(controller.Controller) - instance._tq_legacy_init = False + instance.config = _config(tq_rdma_mode=mode) + instance._tq_owner = "mooncake-owner" + result = controller.TqInitResult(config="mooncake-conf", owner="mooncake-owner") + return instance, result + + +@pytest.mark.parametrize( + ("mode", "failures", "verify_error", "expected_events", "error_match"), + [ + ("auto", [], None, ["handshake"], None), + ("auto", ["timeout"], None, ["handshake", "close", "init_simple"], None), + ("required", ["protocol=tcp"], None, ["handshake", "close"], "reported"), + ("auto", None, RuntimeError("private detail"), ["handshake", "close"], "orchestration failed"), + ( + "auto", + None, + controller.TqHandshakeIsolationError("private detail"), + ["handshake", "close"], + "could not be confirmed stopped", + ), + ], +) +def test_attach_handshake_cleanup_matrix( + monkeypatch: pytest.MonkeyPatch, + mode: str, + failures: list[str] | None, + verify_error: BaseException | None, + expected_events: list[str], + error_match: str | None, +) -> None: + harness = _AttachHarness(monkeypatch, failures=failures, verify_error=verify_error) + instance, result = _confirm(mode) + context = pytest.raises(RuntimeError, match=error_match) if error_match else nullcontext() + with context: + resolved = instance._confirm_mooncake_attach(result, "simple-conf") + if failures: + assert resolved.config == "simple-conf" and resolved.owner == "fallback-owner" + else: + assert resolved is result + assert harness.events == expected_events + + +def test_attach_cleanup_failure_aborts_before_fallback(monkeypatch: pytest.MonkeyPatch) -> None: + harness = _AttachHarness( + monkeypatch, + failures=["timeout"], + close_error=RuntimeError("owner cleanup failed"), + ) + instance, result = _confirm() + with pytest.raises(RuntimeError, match="owner cleanup failed"): + instance._confirm_mooncake_attach(result, "simple-conf") + assert harness.events == ["handshake", "close"] and instance._tq_owner == "mooncake-owner" + + +def test_constructor_and_repeated_cleanup_preserve_owner_boundary(monkeypatch: pytest.MonkeyPatch) -> None: + events: list[Any] = [] + close_data_system = controller.Controller._close_data_system + monkeypatch.setattr(controller, "resolve_sft_num_rollout", lambda _config: None) + monkeypatch.setattr(controller, "HealthManager", lambda **_kwargs: object()) + + def fail_init(instance: Any) -> None: instance._tq_owner = "owner" - closed: list[object] = [] + raise RuntimeError("initialization failed") - monkeypatch.setattr(controller, "close_tq_owner", lambda owner: closed.append(owner) if owner else None) + def close(instance: Any) -> None: + events.append(instance._tq_owner) + instance._tq_owner = None - instance._close_data_system() - instance._close_data_system() + monkeypatch.setattr(controller.Controller, "_initialize_data_system", fail_init) + monkeypatch.setattr(controller.Controller, "_close_data_system", close) + with pytest.raises(RuntimeError, match="initialization failed"): + controller.Controller(SimpleNamespace(use_health_check=False, max_global_restart=3)) + assert events == ["owner"] - assert closed == ["owner"] - assert instance._tq_owner is None + monkeypatch.setattr(controller.Controller, "_close_data_system", close_data_system) + instance = controller.Controller.__new__(controller.Controller) + instance._tq_legacy_init = False + instance._tq_owner = "owner" + monkeypatch.setattr(controller, "close_tq_owner", lambda owner: events.append(owner) if owner else None) + instance._close_data_system() + instance._close_data_system() + assert events == ["owner", "owner"] and instance._tq_owner is None diff --git a/tests/utils/_tq_handshake_timeout_probe.py b/tests/utils/_tq_handshake_timeout_probe.py index 380dc09ba..e6ce61665 100644 --- a/tests/utils/_tq_handshake_timeout_probe.py +++ b/tests/utils/_tq_handshake_timeout_probe.py @@ -1,6 +1,7 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. -"""Subprocess probe for one-shot Ray worker cleanup after attach timeout.""" +"""Subprocess proof that a timed-out one-shot attach worker cannot mutate +later.""" from __future__ import annotations @@ -10,30 +11,27 @@ from pathlib import Path -def _process_is_running(pid: int, create_time: float) -> bool: +def _running(pid: int, created: float) -> bool: import psutil try: process = psutil.Process(pid) - return abs(process.create_time() - create_time) < 1e-3 and process.status() != psutil.STATUS_ZOMBIE + return abs(process.create_time() - created) < 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( +def _write_tq_stub(directory: Path) -> None: + directory.mkdir(parents=True) + (directory / "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( @@ -49,115 +47,71 @@ def init(conf=None): 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), - } - } + started = probe_dir / "started" + mutated = probe_dir / "mutated" + delay = 2.0 + _write_tq_stub(stub_dir) + pythonpath = os.pathsep.join(filter(None, (str(stub_dir), os.environ.get("PYTHONPATH", "")))) import ray - assert not ray.is_initialized() ray.init( address="local", num_cpus=1, include_dashboard=False, logging_level="ERROR", - runtime_env=runtime_env, + runtime_env={ + "env_vars": { + "PYTHONPATH": pythonpath, + "RELAX_TQ_ATTACH_TIMEOUT_SECONDS": "0.3", + "RELAX_TEST_TQ_STARTED_MARKER": str(started), + "RELAX_TEST_TQ_LATE_MARKER": str(mutated), + "RELAX_TEST_TQ_LATE_MUTATION_DELAY": str(delay), + } + }, _temp_dir=str(probe_dir / "ray"), ) try: - from relax.utils.tq import lifecycle as tq_lifecycle + from relax.utils.tq import 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 + class Controller: + def get_config(self) -> dict: + return conf - controller = _Controller.options( - name=tq_lifecycle.CONTROLLER_NAME, - namespace=tq_lifecycle.CONTROLLER_NAMESPACE, - ).remote(conf) + controller = Controller.options( + name=lifecycle.CONTROLLER_NAME, namespace=lifecycle.CONTROLLER_NAMESPACE + ).remote() assert ray.get(controller.get_config.remote()) == conf - - failures = tq_lifecycle.verify_cluster_attach(conf, timeout=0.3) - assert len(failures) == 1 - # The public failure summary is deliberately scrubbed: the underlying - # RayTaskError contains worker addresses, PIDs and traceback paths. - assert failures[0].endswith("handshake task failed (RayTaskError)") - 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: + failures = lifecycle.verify_cluster_attach(conf, timeout=0.3) + assert len(failures) == 1 and failures[0].endswith("handshake task failed (RayTaskError)") + pid_text, created_text = started.read_text(encoding="utf-8").split(",") + identity = int(pid_text), float(created_text) + + time.sleep(delay + 0.2) + assert not mutated.exists(), "timed-out tq.init mutated state after task failure" + deadline = time.monotonic() + 10 + while _running(*identity) and time.monotonic() < deadline: time.sleep(0.05) - assert not _process_is_running(*timed_out_identity), "one-shot handshake worker did not exit after timeout" + assert not _running(*identity), "one-shot handshake worker did not exit" @ray.remote(num_cpus=0, max_retries=0) - def _clean_worker_state() -> tuple[int, float, bool, bool]: - import os - from pathlib import Path - + def clean_successor() -> tuple[int, float, bool]: 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 - - # The dedicated initializer uses the same fail-closed termination - # barrier. This exercises the actual RayActorError transition rather - # than relying only on mocked ray.kill/ray.wait behavior. - owner = tq_lifecycle._TransferQueueOwner.remote() - ray.get(owner.ready.remote()) - tq_lifecycle._stop_owner_actor(owner, timeout=10.0) - try: - after_stop_ref = owner.ready.remote() - ray.get(after_stop_ref) - except ray.exceptions.RayActorError: - pass - else: - raise AssertionError("owner accepted work after terminal cleanup was confirmed") + return process.pid, process.create_time(), transfer_queue.MUTATED + + successor_pid, successor_created, successor_mutated = ray.get(clean_successor.remote()) + assert (successor_pid, successor_created) != identity + assert successor_mutated is False and not mutated.exists() finally: ray.shutdown() diff --git a/tests/utils/_train_actor_init_cleanup_probe.py b/tests/utils/_train_actor_init_cleanup_probe.py index 88afa80dc..355f8c241 100644 --- a/tests/utils/_train_actor_init_cleanup_probe.py +++ b/tests/utils/_train_actor_init_cleanup_probe.py @@ -1,6 +1,7 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. -"""Subprocess probe for train-actor cleanup after one rank fails init.""" +"""Subprocess proof that failed train-actor init cannot mutate after +cleanup.""" from __future__ import annotations @@ -11,58 +12,44 @@ from pathlib import Path -def main(probe_dir: Path) -> None: +def main(directory: Path) -> None: os.environ.setdefault("RAY_ENABLE_UV_RUN_RUNTIME_ENV", "0") os.environ.pop("RAY_ADDRESS", None) - import ray from relax.distributed.ray.actor_group import RayTrainGroup - probe_dir.mkdir(parents=True, exist_ok=True) - late_marker = probe_dir / "late-mutation" + directory.mkdir(parents=True, exist_ok=True) + marker = directory / "late-mutation" @ray.remote(max_restarts=0) - class _InitActor: - def __init__(self, fail: bool): + class InitActor: + def __init__(self, fail: bool) -> None: self.fail = fail - def init(self, _args, _role, **_kwargs): + def init(self, _args, _role, **_kwargs) -> str: if not self.fail: return "ready" - - def mutate_late(): - time.sleep(2.0) - late_marker.write_text("dirty", encoding="utf-8") - - threading.Thread(target=mutate_late, daemon=True).start() + threading.Thread(target=lambda: (time.sleep(2), marker.write_text("dirty")), daemon=True).start() raise RuntimeError("expected initialization failure") - def termination_probe(self): + def termination_probe(self) -> None: threading.Event().wait() ray.init( - address="local", - num_cpus=2, - include_dashboard=False, - logging_level="ERROR", - _temp_dir=str(probe_dir / "ray"), + address="local", num_cpus=2, include_dashboard=False, logging_level="ERROR", _temp_dir=str(directory / "ray") ) try: - actors = [_InitActor.remote(True), _InitActor.remote(False)] group = object.__new__(RayTrainGroup) - group._actor_handlers = actors - + group._actor_handlers = [InitActor.remote(True), InitActor.remote(False)] + with_error = False try: group.init_and_wait(object(), "actor") except ray.exceptions.RayTaskError: - pass - else: - raise AssertionError("rank initialization failure was not propagated") - - assert group._actor_handlers == [] + with_error = True + assert with_error and group._actor_handlers == [] time.sleep(2.2) - assert not late_marker.exists(), "killed actor completed a delayed process-global mutation" + assert not marker.exists() finally: ray.shutdown() diff --git a/tests/utils/test_tq_benchmark_guards.py b/tests/utils/test_tq_benchmark_guards.py index f23f02fd6..1a6e2c048 100644 --- a/tests/utils/test_tq_benchmark_guards.py +++ b/tests/utils/test_tq_benchmark_guards.py @@ -1,15 +1,19 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. -"""CPU contracts for the retained cross-node TransferQueue benchmark.""" +"""CPU contracts for the retained cross-node acceptance benchmark.""" +from __future__ import annotations + +import csv +import io import sys -import warnings +from contextlib import nullcontext from types import SimpleNamespace +from typing import Any import pytest -import torch -from scripts.benchmarks import tq_cross_node_bench +from scripts.benchmarks import tq_cross_node_bench as bench def _argv(protocol: str, *extra: str) -> list[str]: @@ -31,46 +35,20 @@ def _raise(error: BaseException) -> None: raise error -def test_simple_config_builds_without_mooncake_runtime(monkeypatch): +def test_simple_and_multimodal_profiles_use_the_expected_runtime_shapes(monkeypatch: pytest.MonkeyPatch) -> None: import transfer_queue monkeypatch.setattr(transfer_queue, "GRPOGroupNSampler", lambda **_kwargs: object()) - conf = tq_cross_node_bench.build_conf("simple", master="", device="", segment_gib=1) - assert conf.backend.storage_backend == "SimpleStorage" - - -def test_digest_normalizes_dense_and_nested_rows_without_losing_shape(): - dense = {"field": torch.tensor([[1, 2], [3, 4]], dtype=torch.int16)} - with warnings.catch_warnings(): - warnings.simplefilter("ignore", UserWarning) - nested = {"field": torch.nested.nested_tensor(list(dense["field"].unbind()))} - - expected = tq_cross_node_bench.field_byte_digests(dense, ["field"]) - assert tq_cross_node_bench.field_byte_digests(nested, ["field"]) == expected - assert expected["field"][0][:2] == ("torch.int16", (2,)) - - changed = {"field": dense["field"].clone()} - changed["field"][1, 1] += 1 - assert tq_cross_node_bench.field_byte_digests(changed, ["field"]) != expected - - same_bytes_different_shape = {"field": dense["field"].reshape(1, 2, 2)} - one_row = {"field": dense["field"].reshape(1, 4)} - assert tq_cross_node_bench.field_byte_digests(same_bytes_different_shape, ["field"]) != ( - tq_cross_node_bench.field_byte_digests(one_row, ["field"]) - ) - - -def test_multimodal_profile_uses_production_non_tensor_stack_and_recursive_digest(): - payload = tq_cross_node_bench.make_multimodal_payload(num_samples=3, total_mib=1) - multimodal = payload.get("multimodal_train_inputs") - assert type(multimodal).__name__ == "NonTensorStack" - rows = multimodal.tolist() + assert bench.build_conf("simple", master="", device="", segment_gib=1).backend.storage_backend == "SimpleStorage" + payload = bench.make_multimodal_payload(num_samples=3, total_mib=1) + column = payload.get("multimodal_train_inputs") + assert type(column).__name__ == "NonTensorStack" + rows = column.tolist() assert all(set(row) == {"pixel_values", "image_grid_thw"} for row in rows) assert len({tuple(row["pixel_values"].shape) for row in rows}) > 1 - - expected = tq_cross_node_bench.field_byte_digests(payload, ["multimodal_train_inputs"]) + digest = bench.field_byte_digests(payload, ["multimodal_train_inputs"]) rows[1]["pixel_values"][0, 0] += 1 - assert tq_cross_node_bench.field_byte_digests(payload, ["multimodal_train_inputs"]) != expected + assert bench.field_byte_digests(payload, ["multimodal_train_inputs"]) != digest @pytest.mark.parametrize( @@ -78,32 +56,44 @@ def test_multimodal_profile_uses_production_non_tensor_stack_and_recursive_diges [ ("rdma", 800, 0, 1000, True), ("rdma", 799, 0, 1000, False), - ("rdma", 900, 901, 1000, False), + ("rdma", 900, 901, 1000, True), ("tcp", 0, 200, 1000, True), ("tcp", 0, 199, 1000, False), - ("tcp", 201, 200, 1000, False), + ("tcp", 201, 200, 1000, True), ("simple", 0, 1, 1000, True), ("simple", 1, 0, 1000, False), ("rdma", 1, 0, 4 * 1024**3, False), ("tcp", 0, 1, 4 * 1024**3, False), ], ) -def test_wire_proof_is_protocol_and_volume_specific(protocol, ib_bytes, tcp_bytes, payload_bytes, expected): - assert tq_cross_node_bench.wire_is_proven(protocol, ib_bytes, tcp_bytes, payload_bytes) is expected +def test_wire_proof_matrix(protocol: str, ib_bytes: int, tcp_bytes: int, payload_bytes: int, expected: bool) -> None: + assert bench.wire_is_proven(protocol, ib_bytes, tcp_bytes, payload_bytes) is expected -def test_read_counters_scopes_rdma_to_selected_hca(tmp_path): - for device, value in (("rdma0", 11), ("rdma1", 29)): - path = tmp_path / "infiniband" / device / "ports" / "1" / "counters" - path.mkdir(parents=True) - (path / "port_rcv_data").write_text(str(value)) +def test_counter_scope_failure_modes_and_idle_subtraction(tmp_path) -> None: + for device, port, value in (("rdma0", 1, 11), ("rdma1", 1, 17), ("rdma1", 2, 29)): + directory = tmp_path / "infiniband" / device / "ports" / str(port) / "counters" + directory.mkdir(parents=True) + (directory / "port_rcv_data").write_text(str(value)) tcp = tmp_path / "net" / "eth0" / "statistics" tcp.mkdir(parents=True) (tcp / "rx_bytes").write_text("101") - - counters = tq_cross_node_bench.read_counters("eth0", rdma_device="rdma1", sysfs_root=tmp_path) - - assert counters == {"ib:rdma1:1": 29 * 4, "tcp:eth0": 101} + assert bench.read_counters("eth0", "rdma1", 2, sysfs_root=tmp_path) == { + "ib:rdma1:2": 29 * 4, + "tcp:eth0": 101, + } + with pytest.raises(ValueError, match="safe device"): + bench.read_counters("eth0", "*", 1, sysfs_root=tmp_path) + with pytest.raises(RuntimeError, match="TCP receive counter"): + bench.read_counters("missing", sysfs_root=tmp_path) + with pytest.raises(RuntimeError, match="counter set changed"): + bench._counter_delta({"tcp:x": 1}, {}, "tcp:") + with pytest.raises(RuntimeError, match="reset or wrapped"): + bench._counter_delta({"tcp:x": 2}, {"tcp:x": 1}, "tcp:") + assert bench._subtract_idle_noise(1_100, 100, 1.0, 1.0) == 1_000 + assert bench._subtract_idle_noise(50, 100, 1.0, 1.0) == 0 + with pytest.raises(RuntimeError, match="duration must be positive"): + bench._throughput_gbs(1_000, 0.0, "put") @pytest.mark.parametrize( @@ -111,78 +101,153 @@ def test_read_counters_scopes_rdma_to_selected_hca(tmp_path): [ _argv("rdma"), _argv("tcp", "--master", "master.example:50051", "--device", "rdma0"), - _argv("rdma", "--master", "master.example:50051", "--device", "../rdma0"), + *[ + _argv("rdma", "--master", "master.example:50051", "--device", device, "--rdma-port", "1") + for device in ("../rdma0", "*", "rdma*", "rdma?", "[ab]") + ], + _argv("tcp", "--master", "master.example:50051", "--rdma-port", "1"), + _argv("simple", "--master", "master.example:50051"), ], ) -def test_cli_rejects_ambiguous_or_unsafe_counter_configuration(monkeypatch, argv): +def test_cli_rejects_ambiguous_or_unsafe_counter_configuration( + monkeypatch: pytest.MonkeyPatch, argv: list[str] +) -> None: monkeypatch.setattr(sys, "argv", argv) with pytest.raises(SystemExit, match="2"): - tq_cross_node_bench.parse_args() + bench.parse_args() -@pytest.fixture -def teardown_events(monkeypatch): +@pytest.mark.parametrize("cleanup_fails", [False, True], ids=["cleanup-ok", "cleanup-fails"]) +@pytest.mark.parametrize( + ("byte_exact", "tcp_bytes", "error_type", "match", "error_kind"), + [ + (False, 1_000, AssertionError, "byte-exact", "ByteExactMismatch"), + (True, 0, RuntimeError, "wire proof failed", "WireProofFailed"), + ], + ids=["byte-exact", "wire-proof"], +) +def test_failed_gate_is_recorded_and_cleanup_does_not_mask_it( + monkeypatch: pytest.MonkeyPatch, + cleanup_fails: bool, + byte_exact: bool, + tcp_bytes: int, + error_type: type[BaseException], + match: str, + error_kind: str, +) -> None: events: list[str] = [] - monkeypatch.setattr(tq_cross_node_bench.ray, "kill", lambda *_args, **_kwargs: events.append("consumer-killed")) - monkeypatch.setattr(tq_cross_node_bench, "close_tq_unmount_and_wait", lambda: events.append("owner-closed")) - monkeypatch.setattr(tq_cross_node_bench.ray, "shutdown", lambda: events.append("ray-shutdown")) - return events - + monkeypatch.setattr(bench.ray, "get", lambda value: value) + consumer = SimpleNamespace( + sample_idle_counters=SimpleNamespace(remote=lambda _seconds: {"seconds": 1.0, "ib_bytes": 0, "tcp_bytes": 0}), + begin_round=SimpleNamespace(remote=lambda: ({"tcp:x": 0}, 1.0)), + fetch=SimpleNamespace( + remote=lambda *_args: { + "get_ms": 1.0, + "round_seconds": 1.0, + "ib_bytes": 0, + "tcp_bytes": tcp_bytes, + "byte_exact": byte_exact, + "mismatch_fields": [] if byte_exact else ["field"], + } + ), + ) -def test_teardown_keeps_cleanup_order_when_consumer_shutdown_fails(monkeypatch, teardown_events): - consumer = SimpleNamespace(shutdown=SimpleNamespace(remote=lambda: "shutdown-ref")) + class Producer: + def put(self, *_args: Any, **_kwargs: Any) -> None: + events.append("put") + + def clear_partition(self, _partition: str) -> None: + events.append("clear") + if cleanup_fails: + raise RuntimeError("partition cleanup failed") + + output = io.StringIO() + writer = csv.DictWriter(output, fieldnames=bench.CSV_COLUMNS) + writer.writeheader() + with pytest.raises(error_type, match=match) as excinfo: + bench._run_round( + producer=Producer(), + consumer=consumer, + payload=SimpleNamespace(batch_size=[1]), + fields=["field"], + expected={"field": ()}, + nbytes=1_000, + protocol="tcp", + profile="synthetic", + requested_mib=1, + run=1, + writer=writer, + csv_handle=output, + provenance={"relax_sha": "a" * 40, "tq_commit": "b" * 40, "mooncake_version": "test"}, + ) + row = next(csv.DictReader(io.StringIO(output.getvalue()))) + assert (row["status"], row["error_kind"], row["byte_exact"]) == ("fail", error_kind, str(byte_exact)) + assert row["mismatch_fields"] == ("" if byte_exact else "field") + assert isinstance(excinfo.value.__cause__, RuntimeError) is cleanup_fails + assert events == ["put", "clear"] + + +def test_provenance_records_versions_and_rejects_dirty_checkout(monkeypatch: pytest.MonkeyPatch, tmp_path) -> None: + results = iter([SimpleNamespace(returncode=0, stdout=""), SimpleNamespace(returncode=0, stdout="a" * 40 + "\n")]) + monkeypatch.setattr(bench.subprocess, "run", lambda *_args, **_kwargs: next(results)) + direct_url = '{"vcs_info":{"commit_id":"' + "b" * 40 + '"}}' monkeypatch.setattr( - tq_cross_node_bench.ray, - "get", - lambda _ref, timeout: _raise(RuntimeError("consumer shutdown failed")), + bench.importlib_metadata, + "distribution", + lambda _name: SimpleNamespace(read_text=lambda _filename: direct_url), ) - - with pytest.raises(RuntimeError, match="consumer shutdown failed"): - tq_cross_node_bench._teardown_benchmark(consumer, owner_attempted=True) - - assert teardown_events == ["consumer-killed", "owner-closed", "ray-shutdown"] - - -def test_teardown_dirty_cluster_does_not_close_unowned_controller(teardown_events): - tq_cross_node_bench._teardown_benchmark(None, owner_attempted=False) - - assert teardown_events == ["ray-shutdown"] - - -def test_clean_cluster_guard_never_kills_a_healthy_existing_controller(monkeypatch): - controller = object() - killed: list[object] = [] - monkeypatch.setattr(tq_cross_node_bench.ray, "get_actor", lambda *_args, **_kwargs: controller) - monkeypatch.setattr(tq_cross_node_bench.ray, "kill", lambda handle, **_kwargs: killed.append(handle)) - - with pytest.raises(RuntimeError, match="clean exclusive Ray cluster"): - tq_cross_node_bench.require_clean_cluster() - - assert killed == [] - - -def test_teardown_shutdowns_ray_when_owner_cleanup_fails(monkeypatch, teardown_events): + monkeypatch.setattr(bench.importlib_metadata, "version", lambda _name: "0.3.test") + assert bench.collect_provenance(tmp_path) == { + "relax_sha": "a" * 40, + "tq_commit": "b" * 40, + "mooncake_version": "0.3.test", + } monkeypatch.setattr( - tq_cross_node_bench, - "close_tq_unmount_and_wait", - lambda: _raise(RuntimeError("owner cleanup failed")), + bench.subprocess, + "run", + lambda *_args, **_kwargs: SimpleNamespace(returncode=0, stdout=" M tracked-file\n"), ) + with pytest.raises(RuntimeError, match="clean tracked Relax checkout"): + bench.collect_provenance(tmp_path) - with pytest.raises(RuntimeError, match="owner cleanup failed"): - tq_cross_node_bench._teardown_benchmark(None, owner_attempted=True) - - assert teardown_events == ["ray-shutdown"] - -def test_production_controller_cleanup_timeout_fails_closed(monkeypatch): +@pytest.fixture +def teardown_events(monkeypatch: pytest.MonkeyPatch) -> list[str]: from relax.utils.tq import lifecycle - controller = object() - killed: list[object] = [] - monkeypatch.setattr(lifecycle.ray, "get_actor", lambda *_args, **_kwargs: controller) - monkeypatch.setattr(lifecycle.ray, "kill", lambda handle: killed.append(handle)) + events: list[str] = [] + monkeypatch.setattr(bench.ray, "kill", lambda *_args, **_kwargs: events.append("consumer-killed")) + monkeypatch.setattr(lifecycle, "detach_tq_client", lambda: events.append("producer-detached")) + monkeypatch.setattr(lifecycle, "close_tq_owner", lambda _owner: events.append("owner-closed")) + monkeypatch.setattr(bench.ray, "shutdown", lambda: events.append("ray-shutdown")) + return events - with pytest.raises(lifecycle.TqCleanupTimeout, match="still resolvable"): - lifecycle.kill_tq_controller_and_wait(timeout=0) - assert killed == [controller] +@pytest.mark.parametrize("failure", ["consumer", "owner", "none"]) +def test_teardown_is_ordered_and_does_not_touch_unowned_state( + monkeypatch: pytest.MonkeyPatch, teardown_events: list[str], failure: str +) -> None: + from relax.utils.tq import lifecycle + + consumer, attached, owner = None, False, None + expected = ["ray-shutdown"] + error_match = None + if failure == "consumer": + consumer, attached, owner = SimpleNamespace(shutdown=SimpleNamespace(remote=lambda: "ref")), True, "owner" + monkeypatch.setattr(bench.ray, "get", lambda *_args, **_kwargs: _raise(RuntimeError("consumer failed"))) + + def fail_owner_cleanup(_owner: Any) -> None: + teardown_events.append("owner-closed") + raise RuntimeError("owner failed") + + monkeypatch.setattr(lifecycle, "close_tq_owner", fail_owner_cleanup) + expected = ["consumer-killed", "producer-detached", "owner-closed", "ray-shutdown"] + error_match = "consumer failed" + elif failure == "owner": + owner = "owner" + monkeypatch.setattr(lifecycle, "close_tq_owner", lambda _owner: _raise(RuntimeError("owner failed"))) + error_match = "owner failed" + context = pytest.raises(RuntimeError, match=error_match) if error_match else nullcontext() + with context: + bench._teardown_benchmark(consumer, producer_attached=attached, owner=owner) + assert teardown_events == expected diff --git a/tests/utils/test_tq_dataplane_behavior.py b/tests/utils/test_tq_dataplane_behavior.py index fad340c54..c2c20cffe 100644 --- a/tests/utils/test_tq_dataplane_behavior.py +++ b/tests/utils/test_tq_dataplane_behavior.py @@ -1,51 +1,35 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. -"""Real local SimpleStorage data-plane contracts. - -Covers connection, dense/multimodal byte identity, backpressure, empty get, -repeat put, and clear/reinit. Cross-node transport and disconnect behavior -belong to the retained C0/C1/C2 benchmark. -""" +"""Real local SimpleStorage connection and data-plane contracts.""" from __future__ import annotations import importlib.util import time +from typing import Any import pytest import torch +from relax.utils.tq.correctness import diff_digests, leaf_digests, payload_rows + def _has_real_submodule(dotted: str) -> bool: - """Distinguish the real package from CI's single-file TQ stub.""" try: return importlib.util.find_spec(dotted) is not None except (ImportError, ValueError, TypeError): 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)." - ), + not (_has_real_submodule("transfer_queue.storage") and importlib.util.find_spec("ray")), + reason="requires real TransferQueue, Ray, and a startable local CPU cluster", ) - _TQ_ACTOR = "TransferQueueController" _TQ_NS = "transfer_queue" 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 @@ -67,54 +51,41 @@ def _force_kill_controller() -> None: pass -def _flat_values(t): - """Flatten a dense tensor or NestedTensor to its comparable storage.""" - if type(t).__name__ == "NestedTensor": - return t.values().reshape(-1) - return t.reshape(-1) - - -def _row_value(column, row_position: int): - """Extract one row from a dense, nested, or non-tensor column.""" - 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).""" +def _payload(samples: int, fields: list[str], columns: int, seed: int = 0) -> Any: 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]) + generator = torch.Generator().manual_seed(seed) + return TensorDict( + {field: torch.randn(samples, columns, generator=generator) for field in fields}, + batch_size=[samples], + ) -def _multimodal_payload(num_samples: int) -> dict: - """Deterministic Qwen3-VL-shaped payload for the non-tensor TQ path.""" +def _multimodal_payload(samples: int) -> dict[str, Any]: grids = ((1, 58, 64), (1, 34, 64), (1, 64, 64), (1, 26, 40)) generator = torch.Generator().manual_seed(20260813) - multimodal = [] - tokens = [] - for index in range(num_samples): - t, h, w = grids[index % len(grids)] + multimodal, tokens = [], [] + for index in range(samples): + temporal, height, width = grids[index % len(grids)] multimodal.append( { - "pixel_values": torch.randn(t * h * w, 1536, generator=generator), - "image_grid_thw": torch.tensor([[t, h, w]], dtype=torch.int64), + "pixel_values": torch.randn(temporal * height * width, 1536, generator=generator), + "image_grid_thw": torch.tensor([[temporal, height, width]], dtype=torch.int64), } ) tokens.append(torch.randint(0, 151_000, (512 + 173 * index,), generator=generator).tolist()) return {"tokens": tokens, "multimodal_train_inputs": multimodal} -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) +def _get(client: Any, partition: str, fields: list[str], size: int) -> Any: + meta = client.get_meta( + data_fields=fields, + batch_size=size, + partition_id=partition, + mode="fetch", + task_name=partition, + ) + return meta, client.get_data(meta) @pytest.fixture(scope="module") @@ -123,169 +94,107 @@ def _ray_cluster(): ray.init(ignore_reinit_error=True, logging_level="ERROR") yield - try: - ray.shutdown() - except Exception: - pass + ray.shutdown() @pytest.fixture def tq_factory(_ray_cluster): - """Yield an F10-safe ``reinit(capacity, units) -> client`` factory.""" import transfer_queue as tq from omegaconf import OmegaConf from transfer_queue import GRPOGroupNSampler - def _reinit(capacity: int = 1024, units: int = 1): + def reinit(capacity: int = 1024): tq.close() if not _wait_controller_gone(): _force_kill_controller() - _wait_controller_gone() + assert _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, - } - }, + "controller": {"sampler": GRPOGroupNSampler(n_samples_per_prompt=1), "polling_mode": True}, + "backend": {"SimpleStorage": {"total_storage_size": capacity, "num_data_storage_units": 1}}, }, flags={"allow_objects": True}, ) tq.init(conf=conf) return tq.get_client() - yield _reinit - + yield reinit tq.close() _wait_controller_gone() _force_kill_controller() _wait_controller_gone() -class TestTqDataPlaneBehavior: - @pytest.mark.parametrize( - ("fields", "columns", "samples", "seed"), - [ - (["a", "b"], 8, 4, 0), - (["img", "txt", "mask"], 16, 8, 42), - (["pixel_values"], 1176, 4, 7), - ], - ids=["connection", "multi-field", "multimodal-width"], - ) - def test_dense_round_trip_is_byte_exact(self, tq_factory, fields, columns, samples, seed): - client = tq_factory() - payload = _payload(n=samples, fields=fields, cols=columns, seed=seed) - got = _round_trip(client, payload, "dense", fields, samples) - assert set(fields) <= set(got.keys()) - for field in fields: - gv, av = _flat_values(got[field]), _flat_values(payload[field]) - assert gv.numel() == av.numel() - assert gv.dtype == av.dtype - assert torch.equal(gv, av), f"{field}: not byte-exact" - - 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 - - -class TestMultimodalFullLink: - """Production NonTensorStack container survives the full link exactly.""" - - def test_multimodal_list_dict_full_link_byte_exact(self, tq_factory, record_property): - from relax.utils.utils import dict_to_tensordict - from tests.utils.tq._payload_assertions import diff_digests, leaf_digests - - num_samples = 4 - train_data = _multimodal_payload(num_samples) - source = "synthetic" - record_property("multimodal_payload_source", source) - 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" +@pytest.mark.parametrize( + ("fields", "columns", "samples", "seed"), + [(["a", "b"], 8, 4, 0), (["img", "txt", "mask"], 16, 8, 42), (["pixel_values"], 1176, 4, 7)], + ids=["connection", "multi-field", "multimodal-width"], +) +def test_dense_round_trip_is_byte_exact(tq_factory, fields: list[str], columns: int, samples: int, seed: int) -> None: + client = tq_factory() + payload = _payload(samples, fields, columns, seed) + client.put(payload, partition_id="dense") + _, received = _get(client, "dense", fields, samples) + assert set(fields) <= set(received.keys()) + for field in fields: + expected_rows = [leaf_digests(row) for row in payload_rows(payload[field])] + actual_rows = [leaf_digests(row) for row in payload_rows(received[field])] + assert len(actual_rows) == len(expected_rows) + assert all( + not diff_digests(expected, actual) for expected, actual in zip(expected_rows, actual_rows, strict=True) ) - 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]}" + +def test_backpressure_fails_without_publishing_data(tq_factory) -> None: + client = tq_factory(capacity=4) + with pytest.raises(RuntimeError, match="capacity"): + client.put(_payload(8, ["a"], 4), partition_id="backpressure") + meta, _ = _get(client, "backpressure", ["a"], 8) + assert getattr(meta, "size", None) == 0 + + +def test_empty_get_returns_without_hanging(tq_factory) -> None: + meta, data = _get(tq_factory(), "empty", ["a"], 4) + assert getattr(meta, "size", None) == 0 and list(data.keys()) == [] + + +def test_repeat_put_stays_bounded_and_uncorrupted(tq_factory) -> None: + client = tq_factory() + first, second = _payload(4, ["a"], 4, 1), _payload(4, ["a"], 4, 2) + client.put(first, partition_id="repeat") + client.put(second, partition_id="repeat") + meta, received = _get(client, "repeat", ["a"], 4) + candidates = [leaf_digests(row) for row in [*payload_rows(first["a"]), *payload_rows(second["a"])]] + assert getattr(meta, "size", None) == 4 + assert all(leaf_digests(row) in candidates for row in payload_rows(received["a"])) + + +def test_clear_partition_and_reinit_are_isolated(tq_factory) -> None: + client = tq_factory() + client.put(_payload(4, ["a"], 4), partition_id="cleanup") + assert getattr(_get(client, "cleanup", ["a"], 4)[0], "size", None) == 4 + client.clear_partition("cleanup") + assert getattr(_get(client, "cleanup", ["a"], 4)[0], "size", None) == 0 + assert getattr(_get(tq_factory(capacity=16), "cleanup", ["a"], 4)[0], "size", None) == 0 + + +def test_multimodal_non_tensor_full_link_is_byte_exact(tq_factory, record_property) -> None: + from relax.utils.utils import dict_to_tensordict + + samples = 4 + source = _multimodal_payload(samples) + expected_mm = [leaf_digests(row) for row in source["multimodal_train_inputs"]] + expected_tokens = [leaf_digests(torch.tensor(row, dtype=torch.int64)) for row in source["tokens"]] + batch = dict_to_tensordict({**source, "sample_id": list(range(samples))}, batch_size=samples) + assert type(batch.get("multimodal_train_inputs")).__name__ == "NonTensorStack" + record_property("multimodal_payload_source", "synthetic") + + fields = ["sample_id", "tokens", "multimodal_train_inputs"] + client = tq_factory() + client.put(batch, partition_id="multimodal") + _, received = _get(client, "multimodal", fields, samples) + sample_ids = [int(value) for value in payload_rows(received["sample_id"])] + assert sorted(sample_ids) == list(range(samples)) + for position, sample_id in enumerate(sample_ids): + assert not diff_digests(expected_mm[sample_id], leaf_digests(payload_rows(received[fields[2]])[position])) + assert not diff_digests(expected_tokens[sample_id], leaf_digests(payload_rows(received[fields[1]])[position])) diff --git a/tests/utils/test_tq_failure_paths.py b/tests/utils/test_tq_failure_paths.py index 31ec62d49..3ee84c1b1 100644 --- a/tests/utils/test_tq_failure_paths.py +++ b/tests/utils/test_tq_failure_paths.py @@ -1,21 +1,6 @@ # 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 -``tests/utils/tq/test_config.py`` (pure config construction/validation) and -``test_tq_dataplane_behavior.py`` (real TQ on SimpleStorage) did not: - -* timeout -- controller ``get_config`` timeout and attach-handshake timeout -* disconnect -- store errors surface instead of returning corrupt data -* retry/disconnect -- a transient get recovers and a dead peer fails loudly -* automatic degradation as pytest (was a manual two-node script) -* the controller reaper / teardown helpers (now in ``relax.utils.tq.lifecycle``) - -Real Mooncake TCP/RDMA byte-exact and wire-level checks live in the single -cross-node C0/C1/C2 benchmark. This module keeps CI-safe failure semantics and -uses real TransferQueue classes only where a stubbed store is sufficient. -""" +"""CI-safe timeout, retry, disconnect, fallback, and cleanup contracts.""" from __future__ import annotations @@ -25,7 +10,10 @@ import subprocess import sys import tempfile +import time +from contextlib import nullcontext from types import SimpleNamespace +from typing import Any from unittest.mock import AsyncMock, MagicMock import pytest @@ -35,18 +23,9 @@ 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 @@ -57,884 +36,506 @@ def _raise(error: BaseException) -> None: raise error -# --------------------------------------------------------------------------- -# 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() - # Exception handlers require real exception classes; a bare MagicMock - # here would make ``except ray.exceptions.RayError`` invalid at runtime. - fake.exceptions = tq_lifecycle.ray.exceptions - 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_fails_without_being_killed(self, monkeypatch): - killed = self._fake_ray(monkeypatch, actor=MagicMock(), get_result={"backend": {}}) - with pytest.raises(RuntimeError, match="exclusive, clean Ray cluster"): - tq_lifecycle.reap_unusable_tq_controller() - 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"] - - @pytest.mark.parametrize("error_kind", ["timeout", "dead-actor"]) - def test_unusable_controller_is_reaped(self, monkeypatch, error_kind): - """An unresponsive/dead controller must not turn tq.init into a - hang.""" - if error_kind == "timeout": - error = tq_lifecycle.ray.exceptions.GetTimeoutError("private timeout detail") - else: - error = tq_lifecycle.ray.exceptions.RayActorError(error_msg="private actor detail") - killed = self._fake_ray(monkeypatch, actor=MagicMock(), get_raises=error) - assert tq_lifecycle.reap_unusable_tq_controller() is True - assert killed == ["killed"] - - def test_other_ray_error_fails_closed_without_killing_or_leaking_detail(self, monkeypatch): - private_detail = "private control-plane endpoint and traceback" - error = tq_lifecycle.ray.exceptions.RayError(private_detail) - killed = self._fake_ray(monkeypatch, actor=MagicMock(), get_raises=error) - - with pytest.raises(RuntimeError, match="Failed to inspect") as excinfo: - tq_lifecycle.reap_unusable_tq_controller() - - assert killed == [] - assert private_detail not in str(excinfo.value) - - def test_stored_config_ray_error_is_sanitized(self, monkeypatch): - private_detail = "private GCS endpoint and traceback" - monkeypatch.setattr( - tq_lifecycle.ray, - "get_actor", - lambda *_args, **_kwargs: _raise(tq_lifecycle.ray.exceptions.RayError(private_detail)), - ) - - with pytest.raises(RuntimeError) as excinfo: - tq_lifecycle._get_stored_config() - - assert private_detail not in str(excinfo.value) - - -# --------------------------------------------------------------------------- -# 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 = 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) - 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_segment_is_unmounted_even_when_tq_close_fails(self, monkeypatch): - store_client = MagicMock() - calls = self._fake_tq(monkeypatch, store_client=store_client) - - def close_failure(): - calls.append("tq.close") - raise RuntimeError("close failed") - - tq_lifecycle.tq.close.side_effect = close_failure - store_client.close.side_effect = lambda: calls.append("store.close") - - with pytest.raises(RuntimeError, match="close failed"): - tq_lifecycle.close_tq_and_unmount() - - 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) +@pytest.mark.parametrize( + ("state", "expected", "raises"), + [ + ("absent", False, None), + ("healthy", None, RuntimeError), + ("half", True, None), + ("timeout", True, None), + ("dead", True, None), + ("unknown-error", None, RuntimeError), + ], +) +def test_reaper_state_matrix( + monkeypatch: pytest.MonkeyPatch, state: str, expected: bool | None, raises: type[BaseException] | None +) -> None: + actor = MagicMock() + killed: list[str] = [] + fake_ray = MagicMock() + fake_ray.exceptions = tq_lifecycle.ray.exceptions + if state == "absent": + fake_ray.get_actor.side_effect = ValueError("missing") + else: + fake_ray.get_actor.return_value = actor + if state == "healthy": + fake_ray.get.return_value = {"backend": {}} + elif state == "half": + fake_ray.get.return_value = None + elif state == "timeout": + fake_ray.get.side_effect = fake_ray.exceptions.GetTimeoutError("private") + elif state == "dead": + fake_ray.get.side_effect = fake_ray.exceptions.RayActorError(error_msg="private") + elif state == "unknown-error": + fake_ray.get.side_effect = fake_ray.exceptions.RayError("private endpoint") + monkeypatch.setattr(tq_lifecycle, "ray", fake_ray) + monkeypatch.setattr(tq_lifecycle, "kill_tq_controller_and_wait", lambda **_kwargs: killed.append("killed")) + context = pytest.raises(raises) if raises else nullcontext() + with context: + assert tq_lifecycle.reap_unusable_tq_controller() is expected + assert killed == (["killed"] if state in {"half", "timeout", "dead"} else []) + + +def test_controller_cleanup_timeout_fails_closed(monkeypatch: pytest.MonkeyPatch) -> None: + controller = object() + killed: list[Any] = [] + monkeypatch.setattr(tq_lifecycle.ray, "get_actor", lambda *_args, **_kwargs: controller) + monkeypatch.setattr(tq_lifecycle.ray, "kill", lambda handle: killed.append(handle)) + with pytest.raises(tq_lifecycle.TqCleanupTimeout, match="still resolvable"): + tq_lifecycle.kill_tq_controller_and_wait(timeout=0) + assert killed == [controller] + + +@pytest.mark.parametrize( + ("has_store", "close_error", "expected"), + [(True, None, ["tq", "store"]), (True, RuntimeError("close failed"), ["tq", "store"]), (False, None, ["tq"])], +) +def test_close_unmount_matrix( + monkeypatch: pytest.MonkeyPatch, + has_store: bool, + close_error: BaseException | None, + expected: list[str], +) -> None: + events: list[str] = [] + store = MagicMock() + manager = SimpleNamespace(storage_client=store) if has_store else SimpleNamespace() + fake_tq = MagicMock() + fake_tq.get_client.return_value = MagicMock(storage_manager=manager) + + def close() -> None: + events.append("tq") + if close_error: + raise close_error + + fake_tq.close.side_effect = close + store.close.side_effect = lambda: events.append("store") + monkeypatch.setattr(tq_lifecycle, "tq", fake_tq) + context = pytest.raises(RuntimeError, match="close failed") if close_error else nullcontext() + with context: 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() - - -# --------------------------------------------------------------------------- -# Bounded attach (worker-side tq.init used to hang forever) -# --------------------------------------------------------------------------- - - -class TestBoundedAttach: - """attach_tq_client: one deadline for get_config wait and tq.init.""" - - @pytest.mark.parametrize("value", ["soon", "nan", "inf", "-inf"]) - def test_attach_timeout_env_rejects_unusable_values(self, monkeypatch, value): - monkeypatch.setenv("RELAX_TQ_ATTACH_TIMEOUT_SECONDS", value) - with pytest.raises(RuntimeError, match="finite positive") as excinfo: - tq_lifecycle._resolve_attach_timeout() - assert value not in str(excinfo.value) - - 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) - - def test_cluster_attach_covers_all_nodes_with_hard_affinity_and_one_shot_workers(self, monkeypatch): - remote_options = {} - scheduling_strategies = [] - submitted_refs = [] + assert events == expected + + +@pytest.mark.parametrize("value", ["soon", "nan", "inf", "-inf"]) +def test_attach_timeout_rejects_unusable_env(monkeypatch: pytest.MonkeyPatch, value: str) -> None: + monkeypatch.setenv("RELAX_TQ_ATTACH_TIMEOUT_SECONDS", value) + with pytest.raises(RuntimeError, match="finite positive") as excinfo: + tq_lifecycle._resolve_attach_timeout() + assert value not in str(excinfo.value) + + +def test_bounded_init_timeout_error_and_env_override(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("RELAX_TQ_ATTACH_TIMEOUT_SECONDS", "12.5") + assert tq_lifecycle._resolve_attach_timeout() == 12.5 + monkeypatch.setattr(tq_lifecycle.tq, "init", lambda **_kwargs: time.sleep(5)) + with pytest.raises(tq_lifecycle.TqAttachTimeout, match="did not finish"): + tq_lifecycle._bounded_tq_init({}, time.monotonic() + 0.2, role="test") + monkeypatch.setattr(tq_lifecycle.tq, "init", lambda **_kwargs: _raise(ValueError("bad conf"))) + with pytest.raises(ValueError, match="bad conf"): + tq_lifecycle._bounded_tq_init({}, time.monotonic() + 5, role="test") + monkeypatch.setattr(tq_lifecycle.ray, "get_actor", lambda *_args, **_kwargs: _raise(ValueError("missing"))) + with pytest.raises(tq_lifecycle.TqAttachTimeout, match="attach timed out"): + tq_lifecycle._await_controller_config(time.monotonic() + 0.2) + + +class _HandshakeHarness: + def __init__(self, monkeypatch: pytest.MonkeyPatch, node_ids: list[str] | None = None) -> None: + self.args: list[tuple[Any, ...]] = [] + self.workers: list[Any] = [] + self.options: dict[str, Any] = {} + self.strategies: list[Any] = [] + + class Task: + def options(task_self, **options: Any) -> "Task": + self.strategies.append(options["scheduling_strategy"]) + return task_self + + def remote(task_self, *args: Any) -> object: + self.args.append(args) + return object() - class _Task: - def options(self, **options): - scheduling_strategies.append(options["scheduling_strategy"]) - return self + def remote(**options: Any): + self.options.update(options) - def remote(self, *_args): - ref = object() - submitted_refs.append(ref) - return ref + def decorate(function: Any) -> Task: + self.workers.append(function) + return Task() - def record_remote_options(**options): - remote_options.update(options) - return lambda _function: _Task() + return decorate - monkeypatch.setattr(tq_lifecycle.ray, "remote", record_remote_options) - node_ids = ["a" * 56, "b" * 56] - monkeypatch.setattr(tq_lifecycle, "_alive_node_ids", lambda: node_ids) + monkeypatch.setattr(tq_lifecycle.ray, "remote", remote) + monkeypatch.setattr(tq_lifecycle, "_alive_node_ids", lambda: node_ids or ["a" * 56]) monkeypatch.setattr(tq_lifecycle.ray, "wait", lambda refs, **_kwargs: (list(refs), [])) monkeypatch.setattr(tq_lifecycle.ray, "get", lambda _ref: None) - assert tq_lifecycle.verify_cluster_attach({}, timeout=0.1) == [] - assert remote_options["max_calls"] == 1 - assert remote_options["max_retries"] == 0 - assert len(submitted_refs) == len(node_ids) - assert [(strategy.node_id, strategy.soft) for strategy in scheduling_strategies] == [ - (node_id, False) for node_id in node_ids - ] - - def test_cluster_attach_with_no_alive_nodes_fails_closed_without_scheduling(self, monkeypatch): - monkeypatch.setattr(tq_lifecycle, "_alive_node_ids", lambda: []) - monkeypatch.setattr( - tq_lifecycle.ray, - "remote", - lambda **_kwargs: pytest.fail("zero-node validation must not create a remote worker"), - ) - monkeypatch.setattr( - tq_lifecycle.ray, - "wait", - lambda *_args, **_kwargs: pytest.fail("zero-node validation must not call ray.wait"), - ) - - assert tq_lifecycle.verify_cluster_attach({}, timeout=0.1) == ["cluster: no alive Ray nodes discovered"] - - def test_alive_node_without_node_id_fails_closed(self, monkeypatch): - monkeypatch.setattr(tq_lifecycle.ray, "nodes", lambda: [{"Alive": True}]) - failures = tq_lifecycle.verify_cluster_attach({}, timeout=0.1) - assert failures == ["cluster: node discovery failed (RuntimeError)"] - - @staticmethod - def _capture_handshake(monkeypatch): - """Return ``(captured_args, get_worker)`` for the real nested worker. - ``verify_cluster_attach`` defines ``_handshake`` inline, so the only - way to exercise its body — and therefore its ``finally`` detach — is to - grab the function Ray's decorator receives. - """ - captured: list[tuple] = [] - worker: list = [] +def test_cluster_attach_uses_every_node_hard_affinity_and_one_shot_workers(monkeypatch: pytest.MonkeyPatch) -> None: + nodes = ["a" * 56, "b" * 56] + harness = _HandshakeHarness(monkeypatch, nodes) + assert tq_lifecycle.verify_cluster_attach({}, timeout=0.1) == [] + assert (harness.options["max_calls"], harness.options["max_retries"]) == (1, 0) + assert [strategy.node_id for strategy in harness.strategies] == nodes + assert all(strategy.soft is False for strategy in harness.strategies) - class _Task: - def options(self, **_kwargs): - return self - def remote(self, *args): - captured.append(args) - return object() - - def fake_remote(**_options): - def decorate(function): - worker.append(function) - return _Task() - - return decorate - - monkeypatch.setattr(tq_lifecycle.ray, "remote", fake_remote) - # Ray validates node IDs as 28-byte hex, so use a well-formed one. - monkeypatch.setattr(tq_lifecycle, "_alive_node_ids", lambda: ["a" * 56]) - monkeypatch.setattr(tq_lifecycle.ray, "wait", lambda refs, **kwargs: (list(refs), [])) - monkeypatch.setattr(tq_lifecycle.ray, "get", lambda ref: None) - return captured, worker - - def test_handshake_checks_mooncake_config_only_for_mooncake(self, monkeypatch): - """The Mooncake flag controls the manager/config assertion. - - SimpleStorage has no Mooncake protocol configuration to verify. - """ - captured, _worker = self._capture_handshake(monkeypatch) - - mooncake_conf = {"backend": {"storage_backend": "MooncakeStore"}} - assert tq_lifecycle.verify_cluster_attach(mooncake_conf, timeout=0.1) == [] - assert captured[-1][1] is True - assert captured[-1][2] == 0.1 - - simple_conf = {"backend": {"storage_backend": "SimpleStorage"}} - assert tq_lifecycle.verify_cluster_attach(simple_conf, timeout=0.1) == [] - assert captured[-1][1] is False - - def _run_worker(self, monkeypatch, *, assert_error=None): - """Execute the real ``_handshake`` body and record its call order.""" - _captured, worker = self._capture_handshake(monkeypatch) - tq_lifecycle.verify_cluster_attach({"backend": {"storage_backend": "MooncakeStore"}}, timeout=0.1) - assert worker, "Ray decorator never received the handshake function" - - events: list[str] = [] - - def fake_attach(conf, *, role, timeout): - events.append(f"attach:{role}:{timeout}") - return object() - - def fake_assert(): - events.append("assert") - if assert_error is not None: - raise assert_error - - monkeypatch.setattr(tq_lifecycle, "attach_tq_client", fake_attach) - monkeypatch.setattr(tq_lifecycle, "assert_mooncake_rdma_configured", fake_assert) - monkeypatch.setattr(tq_lifecycle, "detach_tq_client", lambda: events.append("detach")) - return worker[0], events - - @pytest.mark.parametrize( - ("backend", "verify_mooncake", "assert_error"), - [ - ("MooncakeStore", True, None), - ("MooncakeStore", True, RuntimeError("protocol=tcp")), - ("SimpleStorage", False, None), - ], - ids=["mooncake", "mooncake-rejected", "simple"], +@pytest.mark.parametrize( + ("backend", "verify_mooncake", "assert_error"), + [ + ("MooncakeStore", True, None), + ("MooncakeStore", True, RuntimeError("protocol=tcp")), + ("SimpleStorage", False, None), + ], +) +def test_handshake_worker_verifies_backend_and_always_detaches( + monkeypatch: pytest.MonkeyPatch, + backend: str, + verify_mooncake: bool, + assert_error: BaseException | None, +) -> None: + harness = _HandshakeHarness(monkeypatch) + conf = {"backend": {"storage_backend": backend}} + assert tq_lifecycle.verify_cluster_attach(conf, timeout=0.1) == [] + assert harness.args[0][1] is verify_mooncake + events: list[str] = [] + monkeypatch.setattr(tq_lifecycle, "attach_tq_client", lambda *_args, **_kwargs: events.append("attach")) + monkeypatch.setattr( + tq_lifecycle, + "assert_mooncake_rdma_configured", + lambda: events.append("assert") or (_raise(assert_error) if assert_error else None), ) - def test_handshake_worker_always_detaches(self, monkeypatch, backend, verify_mooncake, assert_error): - """A rejected transport must not leave this node's segment - registered.""" - handshake, events = self._run_worker(monkeypatch, assert_error=assert_error) - conf = {"backend": {"storage_backend": backend}} - if assert_error is None: - handshake(conf, verify_mooncake, 0.1) - else: - with pytest.raises(RuntimeError, match="protocol=tcp"): - handshake(conf, verify_mooncake, 0.1) - expected = ["attach:attach-handshake:0.1"] - if verify_mooncake: - expected.append("assert") - assert events == [*expected, "detach"] - - def test_ready_task_failure_is_sanitized(self, monkeypatch): - self._capture_handshake(monkeypatch) - secret = "worker endpoint and traceback path must stay private" - monkeypatch.setattr(tq_lifecycle.ray, "get", lambda _ref: _raise(RuntimeError(secret))) - - failures = tq_lifecycle.verify_cluster_attach({}, timeout=0.1) - - assert failures == ["node#1: handshake task failed (RuntimeError)"] - assert "a" * 56 not in failures[0] - assert secret not in failures[0] - - def test_partial_scheduling_failure_cancels_submitted_workers(self, monkeypatch): - first_ref = object() - cancellations: list[tuple[object, bool]] = [] - - class _Task: - calls = 0 - - def options(self, **_kwargs): - return self - - def remote(self, *_args): - self.calls += 1 - if self.calls == 1: - return first_ref - raise RuntimeError("private scheduling detail") - - monkeypatch.setattr(tq_lifecycle.ray, "remote", lambda **_kwargs: lambda _function: _Task()) - monkeypatch.setattr(tq_lifecycle, "_alive_node_ids", lambda: ["a" * 56, "b" * 56]) - monkeypatch.setattr( - tq_lifecycle.ray, - "cancel", - lambda ref, force: cancellations.append((ref, force)), - ) - monkeypatch.setattr(tq_lifecycle.ray, "wait", lambda refs, **_kwargs: (list(refs), [])) - - failures = tq_lifecycle.verify_cluster_attach({}, timeout=0.1) - - assert failures == ["cluster: handshake scheduling failed (RuntimeError)"] - assert cancellations == [(first_ref, True)] - - def test_wait_failure_cancels_and_confirms_submitted_workers(self, monkeypatch): - captured, _worker = self._capture_handshake(monkeypatch) - wait_calls = 0 - cancellations: list[tuple[object, bool]] = [] - - def fake_wait(submitted, **_kwargs): - nonlocal wait_calls - wait_calls += 1 - if wait_calls == 1: - raise RuntimeError("private wait detail") - return list(submitted), [] - - monkeypatch.setattr(tq_lifecycle.ray, "wait", fake_wait) - monkeypatch.setattr( - tq_lifecycle.ray, - "cancel", - lambda ref, force: cancellations.append((ref, force)), - ) - assert captured == [] - - failures = tq_lifecycle.verify_cluster_attach({}, timeout=0.1) - - assert failures == ["cluster: handshake wait failed (RuntimeError)"] - assert len(cancellations) == 1 and cancellations[0][1] is True - - def test_unconfirmed_pending_worker_aborts_fallback_boundary(self, monkeypatch): - self._capture_handshake(monkeypatch) - monkeypatch.setattr(tq_lifecycle.ray, "wait", lambda refs, **_kwargs: ([], list(refs))) - monkeypatch.setattr(tq_lifecycle.ray, "cancel", lambda _ref, force: None) + monkeypatch.setattr(tq_lifecycle, "detach_tq_client", lambda: events.append("detach")) + context = pytest.raises(RuntimeError, match="protocol=tcp") if assert_error else nullcontext() + with context: + harness.workers[0](conf, verify_mooncake, 0.1) + assert events == ["attach", *(["assert"] if verify_mooncake else []), "detach"] + + +def test_handshake_scheduling_failure_cancels_submitted_workers(monkeypatch: pytest.MonkeyPatch) -> None: + first_ref = object() + cancelled: list[object] = [] + + class Task: + calls = 0 + + def options(self, **_kwargs: Any) -> "Task": + return self + + def remote(self, *_args: Any) -> object: + self.calls += 1 + if self.calls == 1: + return first_ref + raise RuntimeError("private scheduling detail") + + monkeypatch.setattr(tq_lifecycle.ray, "remote", lambda **_kwargs: lambda _function: Task()) + monkeypatch.setattr(tq_lifecycle, "_alive_node_ids", lambda: ["a" * 56, "b" * 56]) + monkeypatch.setattr(tq_lifecycle.ray, "cancel", lambda ref, force: cancelled.append(ref)) + monkeypatch.setattr(tq_lifecycle.ray, "wait", lambda refs, **_kwargs: (list(refs), [])) + assert tq_lifecycle.verify_cluster_attach({}, timeout=0.1) == [ + "cluster: handshake scheduling failed (RuntimeError)" + ] + assert cancelled == [first_ref] + + +@pytest.mark.parametrize("cancel_fails", [False, True]) +def test_unconfirmed_handshake_worker_aborts_fallback(monkeypatch: pytest.MonkeyPatch, cancel_fails: bool) -> None: + _HandshakeHarness(monkeypatch) + monkeypatch.setattr(tq_lifecycle.ray, "wait", lambda refs, **_kwargs: ([], list(refs))) + if cancel_fails: + monkeypatch.setattr(tq_lifecycle.ray, "cancel", lambda *_args, **_kwargs: _raise(RuntimeError("private"))) + else: + monkeypatch.setattr(tq_lifecycle.ray, "cancel", lambda *_args, **_kwargs: None) + with pytest.raises(tq_lifecycle.TqHandshakeIsolationError, match="could not be confirmed stopped"): + tq_lifecycle.verify_cluster_attach({}, timeout=0.1) + + +class MooncakeStorageManager: + def __init__(self, storage_client: Any) -> None: + self.storage_client = storage_client + + +@pytest.mark.parametrize( + ("manager", "match"), + [ + (MooncakeStorageManager(MagicMock(protocol="rdma")), None), + (MagicMock(), "not MooncakeStorageManager"), + (MooncakeStorageManager(None), "no storage_client"), + *[(MooncakeStorageManager(MagicMock(protocol=value)), "not configured") for value in ("tcp", None, "")], + ], +) +def test_configured_mooncake_manager_matrix(monkeypatch: pytest.MonkeyPatch, manager: Any, match: str | None) -> None: + fake_tq = MagicMock() + fake_tq.get_client.return_value = MagicMock(storage_manager=manager) + monkeypatch.setattr(tq_lifecycle, "tq", fake_tq) + context = pytest.raises(RuntimeError, match=match) if match else nullcontext() + with context: + tq_lifecycle.assert_mooncake_rdma_configured() - with pytest.raises(tq_lifecycle.TqHandshakeIsolationError, match="could not be confirmed stopped"): - tq_lifecycle.verify_cluster_attach({}, timeout=0.1) - def test_cancel_failure_aborts_fallback_boundary_without_leaking_detail(self, monkeypatch): - self._capture_handshake(monkeypatch) - private_detail = "private worker address and traceback path" - monkeypatch.setattr(tq_lifecycle.ray, "wait", lambda refs, **_kwargs: ([], list(refs))) - monkeypatch.setattr( - tq_lifecycle.ray, - "cancel", - lambda _ref, force: _raise(RuntimeError(private_detail)), +def test_attach_timeout_leaves_no_reusable_process_global_state() -> None: + 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__), "..", "..")) + 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, ) - - with pytest.raises(tq_lifecycle.TqHandshakeIsolationError) as excinfo: - tq_lifecycle.verify_cluster_attach({}, timeout=0.1) - assert private_detail not in str(excinfo.value) - - -class TestAssertMooncakeRdmaConfigured: - """The client field proves configured intent, not negotiated transport. - - ``tq.init`` ignores the caller's conf when attaching to an existing - controller, so the manager and configured protocol must still match the - job-level contract. Wire proof remains a benchmark responsibility. - """ - - class MooncakeStorageManager: - """Name matters: the production check compares - ``type(...).__name__``.""" - - def __init__(self, storage_client): - self.storage_client = storage_client - - def _patch_client(self, monkeypatch, manager): - fake = MagicMock() - fake.get_client.return_value = MagicMock(storage_manager=manager) - monkeypatch.setattr(tq_lifecycle, "tq", fake) - - def test_accepts_mooncake_rdma(self, monkeypatch): - self._patch_client(monkeypatch, self.MooncakeStorageManager(MagicMock(protocol="rdma"))) - tq_lifecycle.assert_mooncake_rdma_configured() - - def test_rejects_non_mooncake_manager(self, monkeypatch): - self._patch_client(monkeypatch, MagicMock()) - with pytest.raises(RuntimeError, match="not MooncakeStorageManager"): - tq_lifecycle.assert_mooncake_rdma_configured() - - def test_rejects_missing_storage_client(self, monkeypatch): - self._patch_client(monkeypatch, self.MooncakeStorageManager(None)) - with pytest.raises(RuntimeError, match="no storage_client"): - tq_lifecycle.assert_mooncake_rdma_configured() - - @pytest.mark.parametrize("protocol", ["tcp", None, ""]) - def test_rejects_non_rdma_protocol(self, monkeypatch, protocol): - self._patch_client(monkeypatch, self.MooncakeStorageManager(MagicMock(protocol=protocol))) - with pytest.raises(RuntimeError, match="not configured for protocol=rdma"): - tq_lifecycle.assert_mooncake_rdma_configured() - - 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) -# --------------------------------------------------------------------------- - - -class TestWorkerDetach: - """detach_tq_client and the teardown hooks that must invoke it.""" - - def test_detach_closes_clients_and_resets_process_handles(self, monkeypatch): - storage_client = MagicMock() - client = MagicMock(storage_manager=SimpleNamespace(storage_client=storage_client)) - fake_tq = MagicMock() - fake_tq.get_client.return_value = client - fake_interface = SimpleNamespace(_TQ_CLIENT=client, _TQ_CONTROLLER=object()) - monkeypatch.setattr(tq_lifecycle.tq, "interface", fake_interface, raising=False) - monkeypatch.setattr(tq_lifecycle, "tq", fake_tq) - - tq_lifecycle.detach_tq_client() - - storage_client.close.assert_called_once_with() - client.close.assert_called_once_with() - assert fake_interface._TQ_CLIENT is None - assert fake_interface._TQ_CONTROLLER is None - - @pytest.mark.parametrize("has_client", [True, False], ids=["attached", "not-attached"]) - def test_component_del_detaches_only_an_attached_client(self, monkeypatch, has_client): - from relax.components.base import Base - - calls = [] - monkeypatch.setattr(tq_lifecycle, "detach_tq_client", lambda: calls.append(True)) - component = Base() - if has_client: - component.data_system_client = object() - component.__del__() - assert calls == ([True] if has_client else []) - if has_client: - assert component.data_system_client is None - - -# --------------------------------------------------------------------------- -# Exclusive-owner initialization transaction -# --------------------------------------------------------------------------- - - -class TestInitializeTqWithFallback: - @staticmethod - def _conf(backend: str) -> dict: - return {"controller": {}, "backend": {"storage_backend": backend}} - - @staticmethod - def _patch_transaction( - monkeypatch, + assert result.returncode == 0, f"probe stdout:\n{result.stdout}\nprobe stderr:\n{result.stderr}" + + +def test_worker_detach_closes_clients_resets_handles_and_is_used_by_components( + monkeypatch: pytest.MonkeyPatch, +) -> None: + storage_client = MagicMock() + client = MagicMock(storage_manager=SimpleNamespace(storage_client=storage_client)) + fake_tq = MagicMock() + fake_tq.get_client.return_value = client + interface = SimpleNamespace(_TQ_CLIENT=client, _TQ_CONTROLLER=object()) + monkeypatch.setattr(tq_lifecycle.tq, "interface", interface, raising=False) + monkeypatch.setattr(tq_lifecycle, "tq", fake_tq) + tq_lifecycle.detach_tq_client() + storage_client.close.assert_called_once_with() + client.close.assert_called_once_with() + assert interface._TQ_CLIENT is None and interface._TQ_CONTROLLER is None + + from relax.components.base import Base + + calls: list[bool] = [] + monkeypatch.setattr(tq_lifecycle, "detach_tq_client", lambda: calls.append(True)) + component = Base() + component.data_system_client = object() + component.__del__() + assert calls == [True] and component.data_system_client is None + + +def _conf(backend: str) -> dict[str, Any]: + return {"controller": {}, "backend": {"storage_backend": backend}} + + +class _InitHarness: + def __init__( + self, + monkeypatch: pytest.MonkeyPatch, + init_effects: list[Any], *, - init_effects: list[object], - start_effects: list[object] | None = None, + start_effect: Any = None, reap_effects: list[BaseException | None] | None = None, - ) -> list[str]: - events: list[str] = [] + ) -> None: + self.events: list[str] = [] effects = iter(init_effects) - owners = iter(start_effects or [f"owner-{index}" for index in range(len(init_effects))]) reaps = iter(reap_effects or []) - def fake_reap(): - events.append("reap") + def reap() -> None: + self.events.append("reap") effect = next(reaps, None) - if effect is not None: + if effect: raise effect - monkeypatch.setattr(tq_lifecycle, "reap_unusable_tq_controller", fake_reap) + def start(*, timeout: float) -> str: + self.events.append("start") + if isinstance(start_effect, BaseException): + raise start_effect + return f"owner-{self.events.count('start')}" - def fake_start(*, timeout): - events.append("start") - effect = next(owners) - if isinstance(effect, BaseException): - raise effect - return effect - - def fake_initialize(owner, conf, *, timeout): + def initialize(owner: Any, conf: dict[str, Any], *, timeout: float) -> Any: backend = conf["backend"]["storage_backend"] - events.append(f"init:{backend}") + self.events.append(f"init:{backend}") effect = next(effects) if isinstance(effect, BaseException): raise effect return tq_lifecycle.TqInitResult(config=conf, owner=owner) - monkeypatch.setattr(tq_lifecycle, "_start_owner", fake_start) - monkeypatch.setattr(tq_lifecycle, "_initialize_owner", fake_initialize) - return events - - def test_simple_path_also_runs_pre_init_reaper_and_becomes_owner(self, monkeypatch): - conf = self._conf("SimpleStorage") - events = self._patch_transaction(monkeypatch, init_effects=[None]) - result = tq_lifecycle.initialize_tq_with_fallback(conf, mode="off") - assert result.owner == "owner-0" - assert events == ["reap", "start", "init:SimpleStorage"] - - @pytest.mark.parametrize("mode", ["off", "auto", "required"]) - def test_healthy_existing_controller_fails_exclusive_without_starting_owner(self, monkeypatch, mode): - requested = self._conf("SimpleStorage") - mismatch = RuntimeError("exclusive cluster is not clean") - events = self._patch_transaction(monkeypatch, init_effects=[], reap_effects=[mismatch]) - - with pytest.raises(RuntimeError, match="exclusive cluster"): - tq_lifecycle.initialize_tq_with_fallback( - requested, - mode=mode, - fallback_conf=self._conf("SimpleStorage"), - ) - - assert events == ["reap"] - - @pytest.mark.parametrize( - ("primary_error", "fallback_reason"), - [ - ( - tq_lifecycle.TqInitializationError("master unavailable"), - "mooncake_init_failed:TqInitializationError", - ), - (tq_lifecycle.TqInitializationTimeout("timed out"), "mooncake_init_failed:TqInitializationTimeout"), - ], - ids=["init-error", "init-timeout"], + monkeypatch.setattr(tq_lifecycle, "reap_unusable_tq_controller", reap) + monkeypatch.setattr(tq_lifecycle, "_start_owner", start) + monkeypatch.setattr(tq_lifecycle, "_initialize_owner", initialize) + + +@pytest.mark.parametrize( + ("mode", "init_effects", "start_effect", "reap_effects", "expected_events", "error"), + [ + ("off", [None], None, None, ["reap", "start", "init:SimpleStorage"], None), + ( + "auto", + [tq_lifecycle.TqInitializationError("failed"), None], + None, + None, + ["reap", "start", "init:MooncakeStore", "reap", "start", "init:SimpleStorage"], + None, + ), + ( + "required", + [tq_lifecycle.TqInitializationError("failed")], + None, + None, + ["reap", "start", "init:MooncakeStore"], + tq_lifecycle.TqInitializationError, + ), + ( + "auto", + [tq_lifecycle.TqCleanupTimeout("pending")], + None, + None, + ["reap", "start", "init:MooncakeStore"], + tq_lifecycle.TqCleanupTimeout, + ), + ( + "auto", + [tq_lifecycle.TqInitializationError("failed")], + None, + [None, RuntimeError("dirty")], + ["reap", "start", "init:MooncakeStore", "reap"], + RuntimeError, + ), + ("auto", [], RuntimeError("schedule failed"), None, ["reap", "start"], RuntimeError), + ], +) +def test_initialize_and_fallback_transaction_matrix( + monkeypatch: pytest.MonkeyPatch, + mode: str, + init_effects: list[Any], + start_effect: Any, + reap_effects: list[BaseException | None] | None, + expected_events: list[str], + error: type[BaseException] | None, +) -> None: + harness = _InitHarness( + monkeypatch, + init_effects, + start_effect=start_effect, + reap_effects=reap_effects, ) - def test_auto_cleans_failed_mooncake_then_retries_simple_once(self, monkeypatch, primary_error, fallback_reason): - primary = self._conf("MooncakeStore") - fallback = self._conf("SimpleStorage") - events = self._patch_transaction(monkeypatch, init_effects=[primary_error, None]) - result = tq_lifecycle.initialize_tq_with_fallback(primary, mode="auto", fallback_conf=fallback) - assert result.config["backend"]["storage_backend"] == "SimpleStorage" - assert result.fallback_reason == fallback_reason - assert events == [ - "reap", - "start", - "init:MooncakeStore", - "reap", - "start", - "init:SimpleStorage", - ] - - def test_required_cleans_failed_init_without_fallback(self, monkeypatch): - primary = self._conf("MooncakeStore") - fallback = self._conf("SimpleStorage") - events = self._patch_transaction( - monkeypatch, - init_effects=[tq_lifecycle.TqInitializationError("master unavailable")], - ) - with pytest.raises(tq_lifecycle.TqInitializationError, match="master unavailable"): - tq_lifecycle.initialize_tq_with_fallback(primary, mode="required", fallback_conf=fallback) - assert events == ["reap", "start", "init:MooncakeStore"] - - def test_cleanup_failure_aborts_auto_before_second_gate(self, monkeypatch): - events = self._patch_transaction( - monkeypatch, - init_effects=[tq_lifecycle.TqCleanupTimeout("owner still running")], - ) - - with pytest.raises(tq_lifecycle.TqCleanupTimeout, match="still running"): - tq_lifecycle.initialize_tq_with_fallback( - self._conf("MooncakeStore"), - mode="auto", - fallback_conf=self._conf("SimpleStorage"), - ) - - assert events == ["reap", "start", "init:MooncakeStore"] - - def test_second_gate_failure_aborts_before_fallback_owner(self, monkeypatch): - gate_error = RuntimeError("residual controller state is unknown") - events = self._patch_transaction( - monkeypatch, - init_effects=[tq_lifecycle.TqInitializationError("master unavailable")], - reap_effects=[None, gate_error], - ) - - with pytest.raises(RuntimeError, match="fallback initialization failed"): - tq_lifecycle.initialize_tq_with_fallback( - self._conf("MooncakeStore"), - mode="auto", - fallback_conf=self._conf("SimpleStorage"), - ) - - assert events == ["reap", "start", "init:MooncakeStore", "reap"] - - def test_owner_creation_failure_is_not_a_candidate_fallback(self, monkeypatch): - events = self._patch_transaction( - monkeypatch, - init_effects=[], - start_effects=[RuntimeError("owner scheduling failed")], - ) - - with pytest.raises(RuntimeError, match="owner scheduling failed"): - tq_lifecycle.initialize_tq_with_fallback( - self._conf("MooncakeStore"), - mode="auto", - fallback_conf=self._conf("SimpleStorage"), - ) - - assert events == ["reap", "start"] + primary = _conf("SimpleStorage" if mode == "off" else "MooncakeStore") + context = pytest.raises(error) if error else nullcontext() + with context: + result = tq_lifecycle.initialize_tq_with_fallback(primary, mode=mode, fallback_conf=_conf("SimpleStorage")) + assert result.owner is not None + if mode == "auto" and len(init_effects) == 2: + assert result.config["backend"]["storage_backend"] == "SimpleStorage" + assert harness.events == expected_events class _RemoteMethod: - def __init__(self, value): + def __init__(self, value: str) -> None: self.value = value - def remote(self, *args, **kwargs): + def remote(self, *_args: Any, **_kwargs: Any) -> str: return self.value class _FakeOwner: - def __init__(self): - self.ready = _RemoteMethod("ready-ref") - self.initialize = _RemoteMethod("initialize-ref") - self.close = _RemoteMethod("close-ref") - self.termination_probe = _RemoteMethod("termination-ref") - - -class TestOwnerProcessBoundary: - def test_start_scheduling_timeout_stops_owner_without_entering_init(self, monkeypatch): - owner = _FakeOwner() - stopped: list[object] = [] - monkeypatch.setattr(tq_lifecycle._TransferQueueOwner, "remote", lambda: owner) - - def timed_out(ref, *, timeout): - assert ref == "ready-ref" - raise tq_lifecycle.ray.exceptions.GetTimeoutError("test timeout") - - monkeypatch.setattr(tq_lifecycle.ray, "get", timed_out) - monkeypatch.setattr(tq_lifecycle, "_stop_owner_actor", lambda handle, **_kwargs: stopped.append(handle)) - - with pytest.raises(RuntimeError, match="owner creation failed"): - tq_lifecycle._start_owner(timeout=0.1) - assert stopped == [owner] - - def test_start_success_returns_scheduled_owner(self, monkeypatch): - owner = _FakeOwner() - monkeypatch.setattr(tq_lifecycle._TransferQueueOwner, "remote", lambda: owner) - monkeypatch.setattr(tq_lifecycle.ray, "get", lambda ref, timeout: None) - - assert tq_lifecycle._start_owner(timeout=1) is owner - - def test_initialize_success_returns_the_exclusive_owner(self, monkeypatch): - owner = _FakeOwner() - stored = TestInitializeTqWithFallback._conf("SimpleStorage") - monkeypatch.setattr(tq_lifecycle.ray, "get", lambda ref, timeout: stored) - - result = tq_lifecycle._initialize_owner(owner, stored, timeout=1) - - assert result.config is stored - assert result.owner is owner - - @pytest.mark.parametrize("error_kind", ["timeout", "runtime"]) - def test_initialize_failure_cleans_owner_before_propagating(self, monkeypatch, error_kind): - owner = _FakeOwner() - cleaned: list[object] = [] - private_detail = "private endpoint and traceback path" - if error_kind == "timeout": - source_error = tq_lifecycle.ray.exceptions.GetTimeoutError("timeout") - expected_error = tq_lifecycle.TqInitializationTimeout - else: - source_error = RuntimeError(private_detail) - expected_error = tq_lifecycle.TqInitializationError + def __init__(self) -> None: + self.ready = _RemoteMethod("ready") + self.initialize = _RemoteMethod("initialize") + self.close = _RemoteMethod("close") + self.termination_probe = _RemoteMethod("probe") + + +@pytest.mark.parametrize("outcome", ["success", "timeout"]) +def test_owner_start_boundary(monkeypatch: pytest.MonkeyPatch, outcome: str) -> None: + owner = _FakeOwner() + stopped: list[Any] = [] + monkeypatch.setattr(tq_lifecycle._TransferQueueOwner, "remote", lambda: owner) + if outcome == "timeout": monkeypatch.setattr( tq_lifecycle.ray, "get", - lambda _ref, timeout: _raise(source_error), + lambda *_args, **_kwargs: _raise(tq_lifecycle.ray.exceptions.GetTimeoutError("timeout")), ) - monkeypatch.setattr(tq_lifecycle, "_cleanup_failed_owner", lambda handle: cleaned.append(handle)) - - with pytest.raises(expected_error) as excinfo: + monkeypatch.setattr(tq_lifecycle, "_stop_owner_actor", lambda handle, **_kwargs: stopped.append(handle)) + with pytest.raises(RuntimeError, match="owner creation failed"): + tq_lifecycle._start_owner(timeout=0.1) + assert stopped == [owner] + else: + monkeypatch.setattr(tq_lifecycle.ray, "get", lambda *_args, **_kwargs: None) + assert tq_lifecycle._start_owner(timeout=0.1) is owner + + +def test_owner_initialize_success_timeout_and_exception_cleanup(monkeypatch: pytest.MonkeyPatch) -> None: + owner = _FakeOwner() + stored = _conf("SimpleStorage") + monkeypatch.setattr(tq_lifecycle.ray, "get", lambda *_args, **_kwargs: stored) + result = tq_lifecycle._initialize_owner(owner, stored, timeout=0.1) + assert result.config is stored and result.owner is owner + + cleaned: list[Any] = [] + monkeypatch.setattr(tq_lifecycle, "_cleanup_failed_owner", lambda handle: cleaned.append(handle)) + for source, expected in ( + (tq_lifecycle.ray.exceptions.GetTimeoutError("timeout"), tq_lifecycle.TqInitializationTimeout), + (RuntimeError("private detail"), tq_lifecycle.TqInitializationError), + ): + monkeypatch.setattr(tq_lifecycle.ray, "get", lambda *_args, error=source, **_kwargs: _raise(error)) + with pytest.raises(expected): tq_lifecycle._initialize_owner(owner, {}, timeout=0.1) + assert cleaned == [owner, owner] - assert cleaned == [owner] - if error_kind == "runtime": - assert private_detail not in str(excinfo.value) - def test_stop_owner_requires_terminal_probe_result(self, monkeypatch): - owner = _FakeOwner() - killed: list[tuple[object, bool]] = [] - monkeypatch.setattr(tq_lifecycle.ray, "kill", lambda handle, no_restart: killed.append((handle, no_restart))) - monkeypatch.setattr(tq_lifecycle.ray, "wait", lambda refs, **_kwargs: (list(refs), [])) - monkeypatch.setattr( - tq_lifecycle.ray, - "get", - lambda _ref: _raise(tq_lifecycle.ray.exceptions.RayActorError(error_msg="dead")), - ) - - tq_lifecycle._stop_owner_actor(owner, timeout=0.1) - - assert killed == [(owner, True)] - - @pytest.mark.parametrize( - ("probe_ready", "match"), - [(False, "remained pending"), (True, "returned normally")], - ids=["pending", "returned"], +@pytest.mark.parametrize( + ("ready", "ray_get", "error"), + [ + (True, lambda _ref: _raise(tq_lifecycle.ray.exceptions.RayActorError(error_msg="dead")), None), + (False, lambda _ref: None, tq_lifecycle.TqCleanupTimeout), + (True, lambda _ref: None, tq_lifecycle.TqCleanupTimeout), + ], +) +def test_owner_stop_requires_terminal_actor_failure( + monkeypatch: pytest.MonkeyPatch, + ready: bool, + ray_get: Any, + error: type[BaseException] | None, +) -> None: + owner = _FakeOwner() + monkeypatch.setattr(tq_lifecycle.ray, "kill", lambda *_args, **_kwargs: None) + monkeypatch.setattr( + tq_lifecycle.ray, "wait", lambda refs, **_kwargs: (list(refs), []) if ready else ([], list(refs)) ) - def test_stop_owner_requires_a_terminal_probe_failure(self, monkeypatch, probe_ready, match): - owner = _FakeOwner() - monkeypatch.setattr(tq_lifecycle.ray, "kill", lambda *_args, **_kwargs: None) - monkeypatch.setattr( - tq_lifecycle.ray, - "wait", - lambda refs, **_kwargs: (list(refs), []) if probe_ready else ([], list(refs)), - ) - monkeypatch.setattr(tq_lifecycle.ray, "get", lambda _ref: None) + monkeypatch.setattr(tq_lifecycle.ray, "get", ray_get) + context = pytest.raises(error) if error else nullcontext() + with context: + tq_lifecycle._stop_owner_actor(owner, timeout=0.1) - with pytest.raises(tq_lifecycle.TqCleanupTimeout, match=match): - tq_lifecycle._stop_owner_actor(owner, timeout=0.1) - def test_failed_owner_cleanup_waits_for_owner_then_reaps_controller(self, monkeypatch): - owner = _FakeOwner() - events: list[str] = [] - monkeypatch.setattr(tq_lifecycle.ray, "get", lambda _ref, timeout: events.append("close")) - monkeypatch.setattr( - tq_lifecycle, - "_stop_owner_actor", - lambda handle, **_kwargs: events.append("owner-terminal"), - ) - monkeypatch.setattr( - tq_lifecycle, - "kill_tq_controller_and_wait", - lambda **_kwargs: events.append("controller-gone"), - ) +def test_owner_cleanup_order_and_fail_closed_boundary(monkeypatch: pytest.MonkeyPatch) -> None: + owner = _FakeOwner() + events: list[str] = [] + monkeypatch.setattr(tq_lifecycle.ray, "get", lambda *_args, **_kwargs: events.append("close")) + monkeypatch.setattr(tq_lifecycle, "_stop_owner_actor", lambda *_args, **_kwargs: events.append("terminal")) + monkeypatch.setattr(tq_lifecycle, "kill_tq_controller_and_wait", lambda **_kwargs: events.append("reap")) + tq_lifecycle._cleanup_failed_owner(owner) + assert events == ["close", "terminal", "reap"] + monkeypatch.setattr( + tq_lifecycle, + "_stop_owner_actor", + lambda *_args, **_kwargs: _raise(tq_lifecycle.TqCleanupTimeout("pending")), + ) + with pytest.raises(tq_lifecycle.TqCleanupTimeout, match="pending"): tq_lifecycle._cleanup_failed_owner(owner) - - assert events == ["close", "owner-terminal", "controller-gone"] - - def test_unconfirmed_owner_aborts_before_controller_cleanup(self, monkeypatch): - owner = _FakeOwner() - controller_cleanup: list[bool] = [] - monkeypatch.setattr(tq_lifecycle.ray, "get", lambda _ref, timeout: None) - monkeypatch.setattr( - tq_lifecycle, - "_stop_owner_actor", - lambda *_args, **_kwargs: _raise(tq_lifecycle.TqCleanupTimeout("owner pending")), - ) - monkeypatch.setattr( - tq_lifecycle, - "kill_tq_controller_and_wait", - lambda **_kwargs: controller_cleanup.append(True), - ) - - with pytest.raises(tq_lifecycle.TqCleanupTimeout, match="owner pending"): - tq_lifecycle._cleanup_failed_owner(owner) - - assert controller_cleanup == [] - - 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: _raise(RuntimeError("close failed")), - ) - monkeypatch.setattr(tq_lifecycle, "_stop_owner_actor", lambda handle, **_kwargs: stopped.append(handle)) - monkeypatch.setattr(tq_lifecycle, "kill_tq_controller_and_wait", lambda **_kwargs: 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 -# --------------------------------------------------------------------------- + assert events == ["close", "terminal", "reap", "close"] class _FlakyStore: - """Stub store whose first ``fail_times`` reads return an error code.""" - - def __init__(self, fail_times: int, code: int = -800, raise_exc: Exception | None = None): + def __init__(self, fail_times: int, error: Exception | None = None) -> None: self.fail_times = fail_times - self.code = code - self.raise_exc = raise_exc - self.get_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.error = error + self.calls: list[list[str]] = [] + + def batch_get_into(self, keys: list[str], _ptrs: list[int], _sizes: list[int]) -> list[int]: + self.calls.append(list(keys)) + if self.error: + raise self.error + if self.fail_times: self.fail_times -= 1 - return [self.code] * len(keys) + return [-800] * 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 _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. - """ +def _client_with_store(store: _FlakyStore) -> Any: from transfer_queue.storage.clients.mooncake_client import MooncakeStoreClient client = object.__new__(MooncakeStoreClient) @@ -943,101 +544,52 @@ def _client_with_store(store) -> object: return client -class _InlineExecutorLoop: - """Run storage-manager sync calls inline for deterministic async tests. +@pytest.mark.skipif(not _REAL_MOONCAKE_CLIENT, reason="requires real TransferQueue Mooncake client") +@pytest.mark.parametrize("disconnect", [False, True], ids=["retry", "disconnect"]) +def test_retry_and_disconnect_surface(monkeypatch: pytest.MonkeyPatch, disconnect: bool) -> None: + monkeypatch.setattr("transfer_queue.storage.clients.mooncake_client.RETRY_DELAY_SECONDS", 0) + error = RuntimeError("Failed to open segment for endpoint='
'") if disconnect else None + store = _FlakyStore(0 if disconnect else 2, error) + client = _client_with_store(store) + context = pytest.raises(RuntimeError, match="Failed to open segment") if disconnect else nullcontext() + with context: + client._batch_get_into_with_retry(["0@f0", "1@f0"], [1, 2], [8, 8]) + assert len(store.calls) == (1 if disconnect else 3) - 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): +class _InlineExecutorLoop: + def run_in_executor(self, _executor: Any, function: Any, *args: Any) -> asyncio.Future[Any]: future = asyncio.get_running_loop().create_future() try: - future.set_result(fn(*args)) + future.set_result(function(*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", -) -class TestRetryAndDisconnect: - """Behavior surface retained in Relax; retry internals belong upstream.""" - - def test_transient_get_failure_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) - keys = ["0@f0", "1@f0"] - client._batch_get_into_with_retry(keys, [1, 2], [8, 8]) - assert store.get_calls == [keys, keys, keys] - - 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]) - - -@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 _REAL_MOONCAKE_CLIENT, reason="requires real TransferQueue Mooncake manager") +@pytest.mark.asyncio +@pytest.mark.parametrize("put_fails", [False, True], ids=["success", "failure"]) +async def test_production_status_follows_storage_success(monkeypatch: pytest.MonkeyPatch, put_fails: bool) -> None: + from tensordict import TensorDict + from transfer_queue.storage.managers.mooncake_manager import MooncakeStorageManager + + monkeypatch.setattr(asyncio, "get_event_loop", lambda: _InlineExecutorLoop()) + calls: list[str] = [] + storage = MagicMock() + if put_fails: + storage.put.side_effect = RuntimeError("capacity exhausted") + else: + storage.put.side_effect = lambda *_args: calls.append("put") + manager = object.__new__(MooncakeStorageManager) + manager.storage_client = storage + manager.notify_data_update = AsyncMock(side_effect=lambda *_args, **_kwargs: calls.append("notify")) + manager.controller_handshake_socket = None + manager.storage_manager_id = "contract-test" + manager.zmq_context = MagicMock() + meta = MagicMock(global_indexes=[7], partition_ids=["capacity"], _custom_backend_meta=[{}]) + meta.get_all_custom_meta.return_value = [{}] + context = pytest.raises(RuntimeError, match="capacity exhausted") if put_fails else nullcontext() + with context: + await manager.put_data(TensorDict({"pixel_values": torch.randn(1, 16)}, batch_size=[1]), meta) + assert calls == ([] if put_fails else ["put", "notify"]) diff --git a/tests/utils/test_train_actor_init_cleanup.py b/tests/utils/test_train_actor_init_cleanup.py index a64d4f4ce..2c1d98548 100644 --- a/tests/utils/test_train_actor_init_cleanup.py +++ b/tests/utils/test_train_actor_init_cleanup.py @@ -1,6 +1,6 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. -"""CPU-only checks for fail-closed Ray train-actor initialization.""" +"""Fail-closed Ray train-actor initialization contracts.""" from __future__ import annotations @@ -8,6 +8,8 @@ import subprocess import sys import tempfile +from types import SimpleNamespace +from typing import Any import pytest import ray @@ -15,139 +17,91 @@ from relax.distributed.ray.actor_group import RayTrainGroup -class _RemoteMethod: - def __init__(self, result=None, error: BaseException | None = None): - self.result = result - self.error = error - - def remote(self, *args, **kwargs): - if self.error is not None: - raise self.error - return self.result +def _group(*actors: Any) -> RayTrainGroup: + group = object.__new__(RayTrainGroup) + group._actor_handlers = list(actors) + return group -class _FakeActor: - def __init__(self, probe_ref=None, probe_error: BaseException | None = None): - self.termination_probe = _RemoteMethod(probe_ref, probe_error) +def _actor(probe: Any = None) -> Any: + return SimpleNamespace(termination_probe=SimpleNamespace(remote=lambda: probe)) -def _group(*actors) -> RayTrainGroup: - group = object.__new__(RayTrainGroup) - group._actor_handlers = list(actors) - return group +def _raise(error: BaseException) -> None: + raise error -def test_successful_initialization_preserves_actor_group(monkeypatch): - actor = _FakeActor() +def test_init_and_wait_preserves_success_and_cleans_submission_failure(monkeypatch: pytest.MonkeyPatch) -> None: + actor = _actor() group = _group(actor) refs = [object(), object()] - values = {refs[0]: "rank-0", refs[1]: "rank-1"} - monkeypatch.setattr(group, "async_init", lambda *args, **kwargs: refs) - monkeypatch.setattr(ray, "wait", lambda pending, num_returns: ([pending[0]], pending[1:])) + values = dict(zip(refs, ["rank-0", "rank-1"], strict=True)) + monkeypatch.setattr(group, "async_init", lambda *_args, **_kwargs: refs) + monkeypatch.setattr(ray, "wait", lambda pending, **_kwargs: ([pending[0]], pending[1:])) monkeypatch.setattr(ray, "get", lambda ref: values[ref]) - monkeypatch.setattr(ray, "kill", lambda *_args, **_kwargs: pytest.fail("success must not kill actors")) - assert group.init_and_wait(object(), "actor") == ["rank-0", "rank-1"] assert group._actor_handlers == [actor] - -def test_synchronous_submission_failure_cleans_partially_initialized_group(monkeypatch): - actor = _FakeActor() - group = _group(actor) - cleanup_calls = [] - - def fail_submission(*_args, **_kwargs): - raise RuntimeError("actor initialization submission failed") - - monkeypatch.setattr(group, "async_init", fail_submission) - monkeypatch.setattr(group, "_terminate_failed_init", lambda: cleanup_calls.append(True)) - + monkeypatch.setattr(group, "async_init", lambda *_args, **_kwargs: _raise(RuntimeError("submission failed"))) + monkeypatch.setattr(group, "_terminate_failed_init", lambda: group._actor_handlers.clear()) with pytest.raises(RuntimeError, match="submission failed"): group.init_and_wait(object(), "actor") - - assert cleanup_calls == [True] + assert group._actor_handlers == [] -def test_first_rank_failure_kills_and_confirms_every_actor(monkeypatch): - actors = [_FakeActor("probe-0"), _FakeActor("probe-1")] +def test_first_rank_failure_kills_and_confirms_every_actor(monkeypatch: pytest.MonkeyPatch) -> None: + actors = [_actor("probe-0"), _actor("probe-1")] group = _group(*actors) - init_refs = ["init-0", "init-1"] - killed = [] - monkeypatch.setattr(group, "async_init", lambda *args, **kwargs: init_refs) - - def fake_wait(refs, *, num_returns, timeout=None): - if timeout is None: - # rank 1 fails before rank 0 finishes; cleanup must begin without - # waiting for rank 0. - return ["init-1"], ["init-0"] - return list(refs), [] - - def fake_get(ref): + monkeypatch.setattr(group, "async_init", lambda *_args, **_kwargs: ["init-0", "init-1"]) + + def wait(refs: list[str], *, timeout: float | None = None, **_kwargs: Any): + return (list(refs), []) if timeout is not None else (["init-1"], ["init-0"]) + + def get(ref: str) -> None: if ref == "init-1": raise RuntimeError("rank initialization failed") - if str(ref).startswith("probe-"): - raise ray.exceptions.RayActorError(error_msg="actor terminated") - pytest.fail(f"unexpected ray.get({ref!r})") + raise ray.exceptions.RayActorError(error_msg="actor terminated") - monkeypatch.setattr(ray, "wait", fake_wait) - monkeypatch.setattr(ray, "get", fake_get) + killed: list[Any] = [] + monkeypatch.setattr(ray, "wait", wait) + monkeypatch.setattr(ray, "get", get) monkeypatch.setattr(ray, "kill", lambda actor, no_restart: killed.append((actor, no_restart))) - with pytest.raises(RuntimeError, match="rank initialization failed"): group.init_and_wait(object(), "actor") - assert killed == [(actors[0], True), (actors[1], True)] assert group._actor_handlers == [] -def test_kill_failure_still_attempts_every_actor_and_fails_closed(monkeypatch): - actors = [_FakeActor("probe-0"), _FakeActor("probe-1")] +@pytest.mark.parametrize("failure", ["kill", "pending"]) +def test_unconfirmed_cleanup_keeps_group_and_fails_closed(monkeypatch: pytest.MonkeyPatch, failure: str) -> None: + actors = [_actor("probe-0"), _actor("probe-1")] group = _group(*actors) - killed = [] - - def fake_kill(actor, *, no_restart): - killed.append(actor) - if actor is actors[0]: - raise RuntimeError("private control-plane detail") - - monkeypatch.setattr(ray, "kill", fake_kill) - monkeypatch.setattr(ray, "wait", lambda refs, **kwargs: (list(refs), [])) - monkeypatch.setattr( - ray, - "get", - lambda _ref: (_ for _ in ()).throw(ray.exceptions.RayActorError(error_msg="actor terminated")), - ) - + if failure == "kill": + monkeypatch.setattr(ray, "kill", lambda *_args, **_kwargs: _raise(RuntimeError("private detail"))) + monkeypatch.setattr(ray, "wait", lambda refs, **_kwargs: (list(refs), [])) + monkeypatch.setattr( + ray, + "get", + lambda _ref: _raise(ray.exceptions.RayActorError(error_msg="terminated")), + ) + else: + monkeypatch.setattr(ray, "kill", lambda *_args, **_kwargs: None) + monkeypatch.setattr(ray, "wait", lambda refs, **_kwargs: ([], list(refs))) with pytest.raises(RuntimeError, match="Failed to confirm train actor cleanup") as excinfo: group._terminate_failed_init(timeout=0.1) - - assert killed == actors - assert "private control-plane detail" not in str(excinfo.value) + assert "private detail" not in str(excinfo.value) assert group._actor_handlers == actors -def test_pending_termination_probe_keeps_group_and_fails_closed(monkeypatch): - actor = _FakeActor("probe") - group = _group(actor) - monkeypatch.setattr(ray, "kill", lambda *_args, **_kwargs: None) - monkeypatch.setattr(ray, "wait", lambda refs, **kwargs: ([], list(refs))) - - with pytest.raises(RuntimeError, match=r"1 task\(s\) remained pending"): - group._terminate_failed_init(timeout=0.1) - - assert group._actor_handlers == [actor] - - -def test_real_actor_failure_cannot_mutate_state_after_cleanup(): - """A killed actor process cannot complete a delayed daemon-thread write.""" +def test_real_actor_failure_cannot_mutate_state_after_cleanup() -> None: 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__), "..", "..")) + root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) with tempfile.TemporaryDirectory(prefix="train-init-ray-") as probe_dir: result = subprocess.run( [sys.executable, "-m", "tests.utils._train_actor_init_cleanup_probe", probe_dir], - cwd=repo_root, + cwd=root, env=env, capture_output=True, text=True, diff --git a/tests/utils/tq/_payload_assertions.py b/tests/utils/tq/_payload_assertions.py deleted file mode 100644 index 7b797e6f4..000000000 --- a/tests/utils/tq/_payload_assertions.py +++ /dev/null @@ -1,109 +0,0 @@ -# Copyright (c) 2026 Relax Authors. All Rights Reserved. - -"""Raw-byte payload assertions shared by TransferQueue tests. - -These helpers intentionally live under ``tests``: production code does not need -generic payload traversal, while the data-plane contract must distinguish byte -identity from value equality for NaNs, signed zero, and nested tensors. -""" - -from __future__ import annotations - -import hashlib -import struct -from typing import Any - -import numpy as np -import torch - - -LeafDigest = tuple[str, str, str] - - -def _tensor_digest(value: torch.Tensor) -> LeafDigest: - 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") - elif isinstance(value, float): - # ``repr(float('nan'))`` discards the NaN payload bits. Pack the - # Python double directly so distinct NaNs and signed zero remain - # byte-distinguishable just like tensor/ndarray leaves. - raw = struct.pack("!d", value) - else: - raw = repr(value).encode("utf-8") - return (f"py.{type(value).__name__}", "", hashlib.sha256(raw).hexdigest()) - - -def _unwrap_non_tensor(value: Any) -> Any: - if type(value).__name__ == "NonTensorStack": - return value.tolist() - if type(value).__name__ == "NonTensorData": - return value.data - return value - - -def _dict_child_path(prefix: str, key: Any) -> str: - """Render a dict key without colliding with nested/list paths.""" - if not isinstance(key, str): - raise TypeError(f"Unsupported payload dict key at {prefix}: {type(key).__name__}") - if key.isidentifier(): - return f"{prefix}.{key}" - return f"{prefix}[{key!r}]" - - -def leaf_digests(payload: Any, prefix: str = "payload") -> dict[str, LeafDigest]: - """Map every supported leaf to ``(dtype, shape, raw-byte SHA-256)``.""" - 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 payload: - if not isinstance(key, str): - raise TypeError(f"Unsupported payload dict key at {prefix}: {type(key).__name__}") - for key in sorted(payload): - digests.update(leaf_digests(payload[key], _dict_child_path(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 mismatch descriptions; an 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 diff --git a/tests/utils/tq/test_config.py b/tests/utils/tq/test_config.py index 9e5ca36dc..32c801951 100644 --- a/tests/utils/tq/test_config.py +++ b/tests/utils/tq/test_config.py @@ -1,22 +1,13 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. -"""Unit tests for TransferQueue backend config construction and validation. - -CPU-only: nothing here starts Ray, TransferQueue, or a Mooncake client. These -tests cover everything Relax decides *before* anything is initialised — the -requested mode, the config dicts handed to ``tq.init``, the master endpoint -format, and the segment-capacity pre-check. - -Actual host-RDMA capability is not testable here by design: it is established by -the real cluster-wide attach handshake, covered in -``tests/utils/test_tq_failure_paths.py``. -""" +"""CPU-only contracts for TransferQueue configuration construction.""" from __future__ import annotations import argparse import importlib.util import os +from typing import Any import pytest @@ -25,6 +16,7 @@ build_mooncake_config, build_simple_storage_config, estimate_payload_bytes, + resolve_global_segment_size, resolve_mooncake_master_address, resolve_tq_capacity_batch_size, validate_config, @@ -46,321 +38,190 @@ def _has_real_tq_storage() -> bool: _REAL_TQ_STORAGE = _has_real_tq_storage() -def _make_args(**kwargs) -> argparse.Namespace: - defaults = dict( - tq_rdma_mode="auto", - tq_rdma_device="", - num_data_storage_units=1, - max_staleness=0, - n_samples_per_prompt=1, - rollout_batch_size=32, - multimodal_keys=None, - seq_length=8192, - ) - defaults.update(kwargs) - return argparse.Namespace(**defaults) - - -class TestValidateConfig: - """Minimal mode/device checks left after the static probe was removed.""" - - @pytest.mark.parametrize("mode", ["off", "auto", "required"]) - def test_accepts_every_supported_mode(self, mode): - assert validate_config(_make_args(tq_rdma_mode=mode)) == [] - - @pytest.mark.parametrize("mode", ["mooncake", "auto\nprivate", None, ["auto"]]) - def test_rejects_unknown_mode_without_echoing_it(self, mode): - """Guards configs restored from a checkpoint or built without - argparse.""" - errors = validate_config(_make_args(tq_rdma_mode=mode)) - assert len(errors) == 1 - assert "--tq-rdma-mode" in errors[0] - assert repr(mode) not in errors[0] - - def test_missing_attribute_defaults_to_off(self): - assert validate_config(argparse.Namespace()) == [] - - @pytest.mark.parametrize("device", [None, ["rdma0"], "rdma0\nforged", " "]) - def test_rejects_non_string_or_whitespace_device(self, device): - errors = validate_config(_make_args(tq_rdma_device=device)) - assert len(errors) == 1 - assert "--tq-rdma-device" in errors[0] - assert repr(device) not in errors[0] - - @pytest.mark.parametrize("device", ["", "rdma0", "mlx5_0"]) - def test_accepts_empty_or_printable_device_name(self, device): - assert validate_config(_make_args(tq_rdma_device=device)) == [] - - -class TestMasterEndpoint: - """``MC_MASTER_ADDRESS`` is deployment configuration, validated by format - only. - - No DNS lookup and no connection attempt: reachability is proven later by - the real attach, and re-adding a network probe here would recreate the - capability heuristic this phase deleted. - """ - - def test_required_and_returned_from_env(self, monkeypatch): - monkeypatch.setenv("MC_MASTER_ADDRESS", _MASTER) - assert resolve_mooncake_master_address() == _MASTER - - def test_missing_is_rejected(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() - - @pytest.mark.parametrize( - "address", - [ - "master.example", # no port - "master.example:", # empty port - "master.example:abc", # non-numeric port - ":50051", # no host - "master.example:0", # port below range - "master.example:65536", # port above range - "fe80::1", # bare IPv6: rpartition would yield port=1 - "fe80::1:50051", # bare IPv6 with port - "[fe80::1]", # bracketed host, no port - "[fe80::1]50051", # bracketed host, missing colon - "master\nprivate:50051", # embedded control character - "master name:50051", # embedded whitespace - ], - ) - def test_malformed_endpoints_are_rejected(self, monkeypatch, address): - monkeypatch.setenv("MC_MASTER_ADDRESS", address) - with pytest.raises(RuntimeError, match="not a usable endpoint"): - resolve_mooncake_master_address() - - def test_rejection_never_echoes_the_endpoint(self, monkeypatch): - """The Controller logs this error verbatim. - - An internal hostname or IP is deployment detail that must not reach job - logs, so the message names the defect only. - """ - secret = "prod-master-07.internal.corp:0" - monkeypatch.setenv("MC_MASTER_ADDRESS", secret) - with pytest.raises(RuntimeError) as excinfo: - resolve_mooncake_master_address() - message = str(excinfo.value) - assert "prod-master-07" not in message - assert "internal.corp" not in message - assert "port is outside" in message - - def test_accepts_hostname_and_bracketed_ipv6(self): - assert _split_host_port(_MASTER) == ("master.example", 50051) - assert _split_host_port("[2001:db8::1]:50051") == ("2001:db8::1", 50051) - - -class TestBackendConfigDicts: - """The dicts handed to ``tq.init``.""" - - @pytest.mark.parametrize("total_storage_size", [1000, None], ids=["bounded", "unlimited"]) - def test_simple_storage_config(self, total_storage_size): - cfg = build_simple_storage_config(total_storage_size=total_storage_size, num_data_storage_units=2) - assert cfg == { - "storage_backend": "SimpleStorage", - "SimpleStorage": {"total_storage_size": total_storage_size, "num_data_storage_units": 2}, - } - - @pytest.mark.parametrize( - ("kwargs", "expected_protocol", "expected_device"), - [ - ({}, "rdma", ""), - ({"device": "rdma0"}, "rdma", "rdma0"), - ({"protocol": "tcp"}, "tcp", ""), - ], - ids=["production-default", "explicit-device", "benchmark-tcp"], - ) - def test_mooncake_config_contract(self, kwargs, expected_protocol, expected_device): - cfg = build_mooncake_config(master_address=_MASTER, **kwargs) - assert cfg["storage_backend"] == "MooncakeStore" - mc = cfg["MooncakeStore"] - assert mc["protocol"] == expected_protocol - assert mc["device_name"] == expected_device - assert mc["hard_pin"] is True - assert mc["auto_init"] is False - assert mc["master_server_address"] == _MASTER - assert mc["use_gdr"] is False - assert "gdr_staging_buffer_mb" not in mc - - def test_master_address_is_not_re_read_from_env(self, monkeypatch): - """The validated endpoint must be the one the client receives.""" - monkeypatch.setenv("MC_MASTER_ADDRESS", "other.example:9999") - mc = build_mooncake_config(master_address=_MASTER)["MooncakeStore"] - assert mc["master_server_address"] == _MASTER - - @pytest.mark.parametrize("master", ["not-an-endpoint", "host:0", "fe80::1", 50051]) - def test_builder_rejects_unvalidated_master_address(self, master): - with pytest.raises(ValueError, match="master_address"): - build_mooncake_config(master_address=master) - - @pytest.mark.parametrize("segment_size", [0, -1, True, 1.5]) - def test_explicit_segment_size_must_be_a_positive_integer(self, segment_size): - with pytest.raises(ValueError, match="positive integer"): - build_mooncake_config(master_address=_MASTER, global_segment_size=segment_size) - - -class TestCorrectnessContract: - """The Mooncake loss-prevention gate.""" - - @pytest.mark.skipif( - not _REAL_TQ_STORAGE, - reason="needs real TransferQueue storage submodules; CPU CI uses a single-file transfer_queue stub", - ) - @pytest.mark.parametrize("override", [None, "0"], ids=["default", "explicit-disable"]) - def test_contract_accepts_safe_memcpy_settings(self, monkeypatch, override): - if override is None: - monkeypatch.delenv("MC_STORE_MEMCPY", raising=False) - else: - monkeypatch.setenv("MC_STORE_MEMCPY", override) - validate_mooncake_runtime_contract() - assert os.environ["MC_STORE_MEMCPY"] == "0" +def _args(**overrides: Any) -> argparse.Namespace: + values = { + "tq_rdma_mode": "auto", + "tq_rdma_device": "", + "num_data_storage_units": 1, + "max_staleness": 0, + "n_samples_per_prompt": 1, + "rollout_batch_size": 32, + "multimodal_keys": None, + "seq_length": 8192, + } + values.update(overrides) + return argparse.Namespace(**values) + + +@pytest.mark.parametrize( + ("overrides", "expected_error"), + [({"tq_rdma_mode": mode}, None) for mode in ("off", "auto", "required")] + + [({"tq_rdma_device": device}, None) for device in ("", "rdma0", "mlx5_0")] + + [({"tq_rdma_mode": mode}, "--tq-rdma-mode") for mode in ("mooncake", "auto\nprivate", None, ["auto"])] + + [({"tq_rdma_device": device}, "--tq-rdma-device") for device in (None, ["rdma0"], "rdma0\nforged", " ")], +) +def test_validate_config_matrix(overrides: dict[str, Any], expected_error: str | None) -> None: + errors = validate_config(_args(**overrides)) + if expected_error is None: + assert errors == [] + else: + assert len(errors) == 1 and expected_error in errors[0] + assert repr(next(iter(overrides.values()))) not in errors[0] + + +def test_missing_config_attributes_default_to_off() -> None: + assert validate_config(argparse.Namespace()) == [] + + +@pytest.mark.parametrize( + "address", + [ + "master.example", + "master.example:", + "master.example:abc", + ":50051", + "master.example:0", + "master.example:65536", + "fe80::1", + "fe80::1:50051", + "[fe80::1]", + "[fe80::1]50051", + "master\nprivate:50051", + "master name:50051", + ], +) +def test_master_endpoint_rejects_malformed_values_without_echo(monkeypatch: pytest.MonkeyPatch, address: str) -> None: + monkeypatch.setenv("MC_MASTER_ADDRESS", address) + with pytest.raises(RuntimeError, match="not a usable endpoint") as excinfo: + resolve_mooncake_master_address() + assert address not in str(excinfo.value) + + +def test_master_endpoint_is_required_and_accepts_hostname_or_bracketed_ipv6( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("MC_MASTER_ADDRESS", raising=False) + with pytest.raises(RuntimeError, match="MC_MASTER_ADDRESS"): + resolve_mooncake_master_address() + monkeypatch.setenv("MC_MASTER_ADDRESS", _MASTER) + assert resolve_mooncake_master_address() == _MASTER + assert _split_host_port(_MASTER) == ("master.example", 50051) + assert _split_host_port("[2001:db8::1]:50051") == ("2001:db8::1", 50051) + + +@pytest.mark.parametrize("total_storage_size", [1000, None], ids=["bounded", "unlimited"]) +def test_simple_storage_config(total_storage_size: int | None) -> None: + assert build_simple_storage_config(total_storage_size, 2) == { + "storage_backend": "SimpleStorage", + "SimpleStorage": {"total_storage_size": total_storage_size, "num_data_storage_units": 2}, + } + + +@pytest.mark.parametrize( + ("kwargs", "protocol", "device"), + [({}, "rdma", ""), ({"device": "rdma0"}, "rdma", "rdma0"), ({"protocol": "tcp"}, "tcp", "")], + ids=["production-default", "explicit-device", "benchmark-tcp"], +) +def test_mooncake_config_contract(kwargs: dict[str, str], protocol: str, device: str) -> None: + config = build_mooncake_config(master_address=_MASTER, **kwargs)["MooncakeStore"] + assert (config["protocol"], config["device_name"]) == (protocol, device) + assert config["master_server_address"] == _MASTER + assert config["hard_pin"] is True and config["auto_init"] is False and config["use_gdr"] is False + assert "gdr_staging_buffer_mb" not in config + + +@pytest.mark.parametrize( + ("kwargs", "match"), + [({"master_address": value}, "master_address") for value in ("not-an-endpoint", "host:0", "fe80::1", 50051)] + + [ + ({"master_address": _MASTER, "global_segment_size": value}, "positive integer") for value in (0, -1, True, 1.5) + ], +) +def test_mooncake_builder_rejects_invalid_direct_inputs(kwargs: dict[str, Any], match: str) -> None: + with pytest.raises(ValueError, match=match): + build_mooncake_config(**kwargs) - @pytest.mark.parametrize( - ("override", "private_marker"), - [("1", None), ("1\nprivate deployment detail", "private deployment detail")], - ids=["enable", "untrusted-value"], - ) - def test_contract_rejects_unsafe_memcpy_settings(self, monkeypatch, override, private_marker): - # 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. + +@pytest.mark.parametrize("override", [None, "0"], ids=["default", "explicit-disable"]) +@pytest.mark.skipif(not _REAL_TQ_STORAGE, reason="requires real TransferQueue storage submodules") +def test_runtime_contract_accepts_safe_memcpy(monkeypatch: pytest.MonkeyPatch, override: str | None) -> None: + if override is None: + monkeypatch.delenv("MC_STORE_MEMCPY", raising=False) + else: monkeypatch.setenv("MC_STORE_MEMCPY", override) - with pytest.raises(RuntimeError, match="MC_STORE_MEMCPY") as excinfo: - validate_mooncake_runtime_contract() - if private_marker is not None: - assert private_marker not in str(excinfo.value) - - @pytest.mark.skipif( - not _REAL_TQ_STORAGE, - reason="needs real TransferQueue storage submodules; CPU CI uses a single-file transfer_queue stub", - ) - def test_unreadable_source_becomes_a_runtime_error(self, monkeypatch): - """A compiled/stripped install cannot prove put/notify ordering. - - The failure must stay inside the RuntimeError boundary the Controller - catches; a bare OSError would abort an ``auto`` run that is supposed to - fall back. - """ - import inspect as inspect_module - - from relax.utils.tq import config as config_module - - def raise_oserror(_obj): - raise OSError("could not get source code") - - monkeypatch.setattr(config_module.inspect, "getsource", raise_oserror) - assert inspect_module is not None # the patch targets the module's own alias - with pytest.raises(RuntimeError, match="Cannot verify TransferQueue put/notify ordering"): - validate_mooncake_runtime_contract() - - -class TestSegmentCapacity: - """Configuration-level capacity pre-check (kept: not a hardware probe).""" - - def test_text_only_passes(self): - assert validate_segment_capacity(_make_args(multimodal_keys=None)) is None - - @pytest.mark.parametrize( - ("overrides", "message"), - [ - ( - { - "multimodal_keys": ["pixel_values"], - "rollout_batch_size": 256, - "n_samples_per_prompt": 8, - "max_staleness": 1, - }, - "insufficient", - ), - ( - { - "multimodal_keys": ["pixel_values"], - "rollout_batch_size": 32, - "n_samples_per_prompt": 1, - "max_staleness": 1, - }, - "RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB", - ), - ( - { - "multimodal_keys": ["pixel_values"], - "rollout_batch_size": 16, - "partial_rollout": True, - "use_dynamic_global_batch_size": True, - "over_sampling_batch_size": 64, - }, - "effective_batch=64", - ), - ], - ids=["large-batch", "staleness", "dynamic-oversampling"], - ) - def test_insufficient_capacity_is_rejected(self, overrides, message): - args = _make_args(**overrides) - err = validate_segment_capacity(args) - assert err is not None - assert message.lower() in err.lower() - - def test_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 - ) - monkeypatch.setenv("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB", "16") - assert validate_segment_capacity(args) is None - - @pytest.mark.parametrize("value", ["four", "-1", "nan", "inf", "-inf"]) - def test_segment_size_env_override_rejects_unusable_values(self, monkeypatch, value): - from relax.utils.tq.config import resolve_global_segment_size - - monkeypatch.setenv("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB", value) - with pytest.raises(RuntimeError, match="finite positive") as excinfo: - resolve_global_segment_size() - assert value not in str(excinfo.value) - - def test_segment_size_env_override_cannot_round_down_to_zero_bytes(self, monkeypatch): - from relax.utils.tq.config import resolve_global_segment_size - - monkeypatch.setenv("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB", "1e-20") - with pytest.raises(RuntimeError, match="at least one byte"): - resolve_global_segment_size() - - -class TestPayloadEstimate: - """The token-budget bound behind the capacity check.""" - - def test_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. - assert estimate_payload_bytes(_make_args(multimodal_keys=None)) == 32 * 1 * 8192 * 32 - - def test_multimodal_is_token_budget_bound(self): - # Qwen3-VL transports four 1536-float32 patch rows per schedulable - # vision token, or 24,576 bytes/token: ~192 MiB at 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 + 24_576) - assert 192 * 1024**2 < per_sample < 193 * 1024**2 - - def test_requires_seq_length(self): - with pytest.raises(RuntimeError, match="seq_length"): - estimate_payload_bytes(_make_args(seq_length=None)) - - @pytest.mark.parametrize( - ("partial_rollout", "expected_batch"), - [(True, 64), (False, 16)], - ids=["dynamic-partial", "nominal"], + validate_mooncake_runtime_contract() + assert os.environ["MC_STORE_MEMCPY"] == "0" + + +@pytest.mark.parametrize( + ("override", "private_marker"), + [("1", None), ("1\nprivate deployment detail", "private deployment detail")], +) +def test_runtime_contract_rejects_unsafe_memcpy_without_echo( + monkeypatch: pytest.MonkeyPatch, override: str, private_marker: str | None +) -> None: + monkeypatch.setenv("MC_STORE_MEMCPY", override) + with pytest.raises(RuntimeError, match="MC_STORE_MEMCPY") as excinfo: + validate_mooncake_runtime_contract() + if private_marker is not None: + assert private_marker not in str(excinfo.value) + + +@pytest.mark.skipif(not _REAL_TQ_STORAGE, reason="requires real TransferQueue storage submodules") +def test_runtime_contract_fails_closed_when_source_is_unavailable(monkeypatch: pytest.MonkeyPatch) -> None: + from relax.utils.tq import config as config_module + + monkeypatch.setattr(config_module.inspect, "getsource", lambda _obj: (_ for _ in ()).throw(OSError())) + with pytest.raises(RuntimeError, match="Cannot verify TransferQueue put/notify ordering"): + validate_mooncake_runtime_contract() + + +@pytest.mark.parametrize( + ("overrides", "expected_fragment"), + [ + ({"multimodal_keys": None}, None), + ({"multimodal_keys": ["pixel_values"], "rollout_batch_size": 256, "n_samples_per_prompt": 8}, "insufficient"), + ({"multimodal_keys": ["pixel_values"], "rollout_batch_size": 32, "max_staleness": 1}, "staleness+1=2"), + ( + { + "multimodal_keys": ["pixel_values"], + "rollout_batch_size": 16, + "partial_rollout": True, + "use_dynamic_global_batch_size": True, + "over_sampling_batch_size": 64, + }, + "effective_batch=64", + ), + ], +) +def test_segment_capacity_matrix(overrides: dict[str, Any], expected_fragment: str | None) -> None: + error = validate_segment_capacity(_args(**overrides)) + if expected_fragment is None: + assert error is None + else: + assert error is not None and expected_fragment.lower() in error.lower() + + +@pytest.mark.parametrize("value", ["four", "-1", "nan", "inf", "-inf", "1e-20"]) +def test_segment_size_env_rejects_unusable_values(monkeypatch: pytest.MonkeyPatch, value: str) -> None: + monkeypatch.setenv("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB", value) + with pytest.raises(RuntimeError) as excinfo: + resolve_global_segment_size() + assert value not in str(excinfo.value) + + +def test_capacity_override_and_payload_bound(monkeypatch: pytest.MonkeyPatch) -> None: + multimodal = _args(multimodal_keys=["pixel_values"], rollout_batch_size=1) + assert estimate_payload_bytes(_args()) == 32 * 32 * 8192 + assert estimate_payload_bytes(multimodal) == 8192 * (32 + 24_576) + dynamic = _args( + rollout_batch_size=16, + partial_rollout=True, + use_dynamic_global_batch_size=True, + over_sampling_batch_size=64, ) - def test_capacity_batch_resolution(self, partial_rollout, expected_batch): - args = _make_args( - rollout_batch_size=16, - partial_rollout=partial_rollout, - use_dynamic_global_batch_size=True, - over_sampling_batch_size=64, - ) - assert resolve_tq_capacity_batch_size(args) == expected_batch - if partial_rollout: - assert estimate_payload_bytes(args) == expected_batch * 8192 * 32 + assert resolve_tq_capacity_batch_size(dynamic) == 64 + monkeypatch.setenv("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB", "16") + assert validate_segment_capacity(_args(multimodal_keys=["pixel_values"], max_staleness=1)) is None + with pytest.raises(RuntimeError, match="seq_length"): + estimate_payload_bytes(_args(seq_length=None)) diff --git a/tests/utils/tq/test_payload_assertions.py b/tests/utils/tq/test_payload_assertions.py index 37a9ebaf4..3040bde7f 100644 --- a/tests/utils/tq/test_payload_assertions.py +++ b/tests/utils/tq/test_payload_assertions.py @@ -1,6 +1,6 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. -"""CPU-only tests for TransferQueue payload byte-identity assertions.""" +"""Raw-byte identity contracts shared by dataplane and acceptance tests.""" from __future__ import annotations @@ -12,21 +12,15 @@ import pytest import torch -from tests.utils.tq._payload_assertions import diff_digests, leaf_digests +from relax.utils.tq.correctness import diff_digests, leaf_digests, payload_nbytes, payload_rows class NonTensorData: - """Minimal tensordict-compatible wrapper for an optional-dependency-free - test.""" - def __init__(self, data: Any) -> None: self.data = data class NonTensorStack: - """Minimal tensordict-compatible wrapper for an optional-dependency-free - test.""" - def __init__(self, data: Any) -> None: self._data = data @@ -34,74 +28,51 @@ def tolist(self) -> Any: return self._data -def test_leaf_digests_distinguishes_raw_bytes_from_value_equality(): - positive_zero = leaf_digests(torch.tensor([0.0], dtype=torch.float32)) - negative_zero = leaf_digests(torch.tensor([-0.0], dtype=torch.float32)) - - assert positive_zero["payload"][:2] == negative_zero["payload"][:2] - assert positive_zero["payload"][2] != negative_zero["payload"][2] - assert diff_digests(positive_zero, negative_zero) == [ - f"payload: sha256 mismatch (expected {positive_zero['payload'][2]}, got {negative_zero['payload'][2]})" - ] - - -def test_scalar_nan_payload_bits_are_not_collapsed_by_repr(): - first = struct.unpack("!d", bytes.fromhex("7ff8000000000001"))[0] - second = struct.unpack("!d", bytes.fromhex("7ff8000000000002"))[0] - - assert repr(first) == repr(second) == "nan" - assert leaf_digests(first) != leaf_digests(second) - - -def test_dict_paths_do_not_collapse_dotted_keys_into_nested_keys(): - digests = leaf_digests({"a.b": 1, "a": {"b": 2}}) +def test_digest_preserves_raw_bytes_dtype_shape_and_paths() -> None: + positive = leaf_digests(torch.tensor([0.0], dtype=torch.float32)) + negative = leaf_digests(torch.tensor([-0.0], dtype=torch.float32)) + assert positive["payload"][:2] == negative["payload"][:2] + assert positive["payload"][2] != negative["payload"][2] - assert len(digests) == 2 - assert set(digests) == {"payload['a.b']", "payload.a.b"} + first_nan = struct.unpack("!d", bytes.fromhex("7ff8000000000001"))[0] + second_nan = struct.unpack("!d", bytes.fromhex("7ff8000000000002"))[0] + assert repr(first_nan) == repr(second_nan) == "nan" + assert leaf_digests(first_nan) != leaf_digests(second_nan) + paths = leaf_digests({"a.b": 1, "a": {"b": 2}}) + assert set(paths) == {"payload['a.b']", "payload.a.b"} -def test_non_string_dict_keys_fail_loudly(): - with pytest.raises(TypeError, match="Unsupported payload dict key at payload: int"): - leaf_digests({1: "value"}) - -def test_leaf_digests_preserves_dtype_shape_and_nested_tensor_rows(): +def test_digest_supports_nested_numpy_and_non_tensor_containers() -> None: with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) nested = torch.nested.nested_tensor( [torch.tensor([1, 2], dtype=torch.int16), torch.tensor([3], dtype=torch.int16)] ) - payload = { - "array": np.array([[1, 2]], dtype=np.uint16), - "nested": nested, - "tensor": torch.tensor([[1, 2]], dtype=torch.int32), - } - - digests = leaf_digests(payload) - + digests = leaf_digests( + { + "array": np.array([[1, 2]], dtype=np.uint16), + "nested": nested, + "wrapped": NonTensorStack([NonTensorData({"text": "hello"}), NonTensorData(None)]), + } + ) assert digests["payload.array"][:2] == ("np.uint16", "1x2") assert digests["payload.nested[0]"][:2] == ("torch.int16", "2") assert digests["payload.nested[1]"][:2] == ("torch.int16", "1") - assert digests["payload.tensor"][:2] == ("torch.int32", "1x2") - - -def test_leaf_digests_unwraps_non_tensor_containers(): - wrapped = NonTensorStack([NonTensorData({"text": "hello"}), NonTensorData(None)]) - - assert leaf_digests(wrapped) == leaf_digests([{"text": "hello"}, None]) - - -def test_leaf_digests_rejects_unknown_leaf_type(): - with pytest.raises(TypeError, match=r"Unsupported payload leaf at payload\.bad: object"): - leaf_digests({"bad": object()}) + assert leaf_digests(NonTensorStack([NonTensorData(None)])) == leaf_digests([None]) -def test_diff_digests_reports_missing_and_extra_leaves(): - expected = leaf_digests({"expected": 1}) - actual = leaf_digests({"actual": 1}) +def test_rows_nbytes_and_diff_cover_non_tensor_payloads() -> None: + tensor = torch.tensor([[1, 2], [3, 4]], dtype=torch.int16) + wrapped = NonTensorStack([NonTensorData({"text": "é"}), NonTensorData(None)]) + assert all(torch.equal(got, want) for got, want in zip(payload_rows(tensor), tensor.unbind(), strict=True)) + assert payload_rows(wrapped) == [{"text": "é"}, None] + assert payload_nbytes({"tensor": tensor, "wrapped": wrapped}) == 8 + 2 + len(repr(None)) + problems = diff_digests(leaf_digests({"expected": 1}), leaf_digests({"actual": 1})) + assert len(problems) == 2 and "unexpected extra" in problems[0] and "missing" in problems[1] - problems = diff_digests(expected, actual) - assert len(problems) == 2 - assert problems[0].startswith("payload.actual: unexpected extra leaf") - assert problems[1].startswith("payload.expected: missing") +@pytest.mark.parametrize("payload", [{1: "value"}, {"bad": object()}], ids=["non-string-key", "unknown-leaf"]) +def test_digest_rejects_unsupported_payload(payload: dict[Any, Any]) -> None: + with pytest.raises(TypeError, match="Unsupported payload"): + leaf_digests(payload) From d977762496f119f21aef8d67385d1019b89ec4a4 Mon Sep 17 00:00:00 2001 From: overloadedHenry <115620759+overloadedHenry@users.noreply.github.com> Date: Thu, 27 Aug 2026 01:51:06 +0800 Subject: [PATCH 30/30] fix(tq): require versioned correctness contract --- docker/Dockerfile | 2 +- docker/Dockerfile.npu | 2 +- docs/draft/transfer_queue_rdma.md | 7 +++--- relax/core/controller.py | 2 +- relax/utils/arguments.py | 2 +- relax/utils/tq/config.py | 24 +++---------------- relax/utils/tq/correctness.py | 25 ++++++++++---------- tests/utils/tq/test_config.py | 38 +++++++++++++++---------------- 8 files changed, 43 insertions(+), 59 deletions(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index c473e89b0..3e8d7a4aa 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -116,7 +116,7 @@ COPY requirements.txt /tmp/requirements.txt RUN pip install --ignore-installed PyJWT && \ pip install -r /tmp/requirements.txt --no-cache-dir && \ pip install --no-cache-dir "compressed_tensors>=0.13.0" tensordict==0.10.0 pyvers==0.1.0 'nvidia-modelopt[hf]==0.44.0' --no-deps && \ - pip install "transferqueue @ git+https://github.com/redai-infra/TransferQueue.git@58054a33834aadbcf76aacd6b1e32e25c030f2c9" --no-deps + pip install "transferqueue @ git+https://github.com/redai-infra/TransferQueue.git@6c7a587292910af0827f027de99e005e1900310e" --no-deps # sgl-router: override the official wheel (pulled by requirements.txt above) with # slime's r3-capable fork. The official sglang-router drops the routed_experts diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu index 89ed93f0c..0d4055f5a 100644 --- a/docker/Dockerfile.npu +++ b/docker/Dockerfile.npu @@ -68,7 +68,7 @@ RUN pip install pyyaml && \ RUN cd /root && rm -rf /root/pytorch && \ pip install triton-ascend==3.2.0 && \ pip install tensordict==0.10.0 pyvers==0.1.0 --no-deps && \ - pip install "transferqueue @ git+https://github.com/redai-infra/TransferQueue.git@58054a33834aadbcf76aacd6b1e32e25c030f2c9" --no-deps + pip install "transferqueue @ git+https://github.com/redai-infra/TransferQueue.git@6c7a587292910af0827f027de99e005e1900310e" --no-deps # Clone Megatron-LM, MindSpeed, MindSpeed-Bridge, Megatron-Bridge and install RUN git clone https://gitcode.com/ascend/MindSpeed.git /root/MindSpeed && \ diff --git a/docs/draft/transfer_queue_rdma.md b/docs/draft/transfer_queue_rdma.md index 29410ab28..a63892b56 100644 --- a/docs/draft/transfer_queue_rdma.md +++ b/docs/draft/transfer_queue_rdma.md @@ -57,10 +57,11 @@ export RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB=8 ## Correctness 与依赖 gate -当前 Relax pin 为 TransferQueue `58054a33834aadbcf76aacd6b1e32e25c030f2c9`。现有检查只能确认 retry API 和部分源码顺序，不能证明完整 fail-closed 语义。启用 Mooncake/RDMA 前，上游 TransferQueue PR 必须提供并固定到明确版本或 capability marker，至少保证： +当前 Relax pin 为 TransferQueue `6c7a587292910af0827f027de99e005e1900310e`，并要求 `MOONCAKE_CORRECTNESS_CONTRACT_VERSION >= 1`。Contract version 1 保证： -- `batch_upsert_from` 和 `batch_get_into` 的每次 retry 都校验返回结果与请求 key 等长； -- `NOTIFY_DATA_UPDATE_ACK` 必须验证 positive ACK，controller 拒绝时 producer 不能按成功结束。 +- `batch_upsert_from` 和 `batch_get_into` 的每次 batch/retry 都校验返回结果与请求 key 等长； +- `batch_remove` 的非幂等失败向调用方传播； +- `NOTIFY_DATA_UPDATE_ACK` 验证 positive ACK，controller 拒绝时 producer 不会按成功结束。 这些修复应在 TransferQueue 上游实现；Relax 不使用 monkey patch 替代依赖修复。 diff --git a/relax/core/controller.py b/relax/core/controller.py index 9c50be1a5..585c87b22 100644 --- a/relax/core/controller.py +++ b/relax/core/controller.py @@ -21,7 +21,7 @@ raise ImportError( "transfer_queue is out of date (missing StreamingTokenBudgetSampler). Upgrade with:\n" ' pip install "transferqueue @ git+https://github.com/redai-infra/' - 'TransferQueue.git@58054a33834aadbcf76aacd6b1e32e25c030f2c9" --no-deps\n' + 'TransferQueue.git@6c7a587292910af0827f027de99e005e1900310e" --no-deps\n' "or use the latest image." ) from e diff --git a/relax/utils/arguments.py b/relax/utils/arguments.py index dc07dc1a0..36892e9f8 100644 --- a/relax/utils/arguments.py +++ b/relax/utils/arguments.py @@ -37,7 +37,7 @@ _MIN_TQ_VERSION = "0.1.10.dev0" _TQ_UPGRADE_CMD = ( 'pip install "transferqueue @ git+https://github.com/redai-infra/' - 'TransferQueue.git@58054a33834aadbcf76aacd6b1e32e25c030f2c9" --no-deps' + 'TransferQueue.git@6c7a587292910af0827f027de99e005e1900310e" --no-deps' ) diff --git a/relax/utils/tq/config.py b/relax/utils/tq/config.py index ce95c2163..6d93bdd70 100644 --- a/relax/utils/tq/config.py +++ b/relax/utils/tq/config.py @@ -15,7 +15,6 @@ from __future__ import annotations -import inspect import math import os from typing import Any @@ -146,29 +145,12 @@ def resolve_global_segment_size() -> int: def validate_mooncake_runtime_contract() -> None: """Validate the Relax-side portion of the Mooncake safety contract. - The environment guard and available read-only capability checks run before - every Mooncake client is created or attached. Retry result-length and - positive-ACK correctness must be supplied by the pinned upstream - TransferQueue revision; this function does not monkey-patch that package. + The environment guard and versioned correctness-contract check run before + every Mooncake client is created or attached. This function does not + monkey-patch TransferQueue. """ ensure_mooncake_correctness_guards() - from transfer_queue.storage.managers.base import KVStorageManager - - try: - put_source = inspect.getsource(KVStorageManager.put_data) - except (OSError, TypeError): - # No retrievable source (compiled/stripped install): the ordering - # contract cannot be proven, so fail like any other unmet gate instead - # of escaping this RuntimeError-only boundary. - raise RuntimeError("Cannot verify TransferQueue put/notify ordering because source is unavailable") from None - 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 diff --git a/relax/utils/tq/correctness.py b/relax/utils/tq/correctness.py index 2f1056818..3a37bb316 100644 --- a/relax/utils/tq/correctness.py +++ b/relax/utils/tq/correctness.py @@ -3,10 +3,9 @@ """Mooncake safety guards and raw-byte payload correctness helpers. Relax validates the capabilities and environment it can inspect without -modifying TransferQueue at runtime. Per-retry result-length validation and a -strict production-status ACK must be fixed in upstream TransferQueue and then -consumed through an updated, capability-marked pin; method-name checks alone do -not prove those semantics. +modifying TransferQueue at runtime. The pinned TransferQueue advertises a +versioned correctness contract for batch/retry result counts, remove failure +propagation, and fail-closed production-status ACKs. Mooncake 0.3.10.post2 was observed corrupting TCP-protocol transfers through its auto-enabled memcpy fast path, so that path is force-disabled here and an @@ -23,6 +22,7 @@ LeafDigest = tuple[str, str, str] +_REQUIRED_TQ_MOONCAKE_CONTRACT_VERSION = 1 def _enforce_safe_memcpy() -> None: @@ -54,19 +54,20 @@ def ensure_mooncake_correctness_guards() -> None: """Validate that the installed stack can run MooncakeStore safely. Enforces the memcpy environment contract and checks that the pinned - TransferQueue ships the Mooncake retry APIs Relax's data plane relies on. - It does not modify TransferQueue code or objects at runtime. + TransferQueue advertises the Mooncake correctness contract Relax's data + plane relies on. It does not modify TransferQueue at runtime. """ _enforce_safe_memcpy() try: - from transfer_queue.storage.clients.mooncake_client import MooncakeStoreClient + import transfer_queue as tq except ImportError as error: - raise RuntimeError("Installed TransferQueue has no MooncakeStore support") from error + raise RuntimeError( + "Installed TransferQueue does not satisfy the required Mooncake correctness contract" + ) 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)) + actual = getattr(tq, "MOONCAKE_CORRECTNESS_CONTRACT_VERSION", 0) + if isinstance(actual, bool) or not isinstance(actual, int) or actual < _REQUIRED_TQ_MOONCAKE_CONTRACT_VERSION: + raise RuntimeError("Installed TransferQueue does not satisfy the required Mooncake correctness contract") def _tensor_digest(value: Any) -> LeafDigest: diff --git a/tests/utils/tq/test_config.py b/tests/utils/tq/test_config.py index 32c801951..d14bba346 100644 --- a/tests/utils/tq/test_config.py +++ b/tests/utils/tq/test_config.py @@ -5,8 +5,9 @@ from __future__ import annotations import argparse -import importlib.util import os +import sys +from types import ModuleType from typing import Any import pytest @@ -28,16 +29,6 @@ _MASTER = "master.example:50051" -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 _args(**overrides: Any) -> argparse.Namespace: values = { "tq_rdma_mode": "auto", @@ -143,8 +134,13 @@ def test_mooncake_builder_rejects_invalid_direct_inputs(kwargs: dict[str, Any], @pytest.mark.parametrize("override", [None, "0"], ids=["default", "explicit-disable"]) -@pytest.mark.skipif(not _REAL_TQ_STORAGE, reason="requires real TransferQueue storage submodules") -def test_runtime_contract_accepts_safe_memcpy(monkeypatch: pytest.MonkeyPatch, override: str | None) -> None: +@pytest.mark.parametrize("contract_version", [1, 2], ids=["required", "forward-compatible"]) +def test_runtime_contract_accepts_safe_memcpy( + monkeypatch: pytest.MonkeyPatch, override: str | None, contract_version: int +) -> None: + tq_stub = ModuleType("transfer_queue") + tq_stub.MOONCAKE_CORRECTNESS_CONTRACT_VERSION = contract_version + monkeypatch.setitem(sys.modules, "transfer_queue", tq_stub) if override is None: monkeypatch.delenv("MC_STORE_MEMCPY", raising=False) else: @@ -167,12 +163,16 @@ def test_runtime_contract_rejects_unsafe_memcpy_without_echo( assert private_marker not in str(excinfo.value) -@pytest.mark.skipif(not _REAL_TQ_STORAGE, reason="requires real TransferQueue storage submodules") -def test_runtime_contract_fails_closed_when_source_is_unavailable(monkeypatch: pytest.MonkeyPatch) -> None: - from relax.utils.tq import config as config_module - - monkeypatch.setattr(config_module.inspect, "getsource", lambda _obj: (_ for _ in ()).throw(OSError())) - with pytest.raises(RuntimeError, match="Cannot verify TransferQueue put/notify ordering"): +@pytest.mark.parametrize("contract_version", [None, 0, "1", True], ids=["missing", "old", "string", "bool"]) +def test_runtime_contract_rejects_missing_or_invalid_marker( + monkeypatch: pytest.MonkeyPatch, contract_version: int | str | bool | None +) -> None: + monkeypatch.setenv("MC_STORE_MEMCPY", "0") + tq_stub = ModuleType("transfer_queue") + if contract_version is not None: + tq_stub.MOONCAKE_CORRECTNESS_CONTRACT_VERSION = contract_version + monkeypatch.setitem(sys.modules, "transfer_queue", tq_stub) + with pytest.raises(RuntimeError, match="required Mooncake correctness contract"): validate_mooncake_runtime_contract()