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 new file mode 100644 index 000000000..5f1400e2d --- /dev/null +++ b/docs/draft/transfer_queue_rdma.md @@ -0,0 +1,222 @@ +# TransferQueue RDMA 数据面使用与运维指南 + +## 概述 + +Relax 的数据面（rollout ↔ train 之间的样本传输）默认走 TransferQueue 的 SimpleStorage/ZMQ。本特性把 TransferQueue 已有的 MooncakeStore 后端接出来，使数据面可以走 RDMA，并在能力不足时安全回退。 + +首期只做配置接入、能力探测与一致回退，**不改变 payload 形状与数据分发语义**。默认参数仍使用 SimpleStorage 及原有 controller 所有权模型；同时所有 worker attach（包括 SimpleStorage）新增默认 60 秒 deadline，半初始化 controller 会被回收，Controller 构造失败时会关闭本进程已经完成的 legacy `tq.init`。 + +## 配置入口 + +只暴露四个表达使用意图的参数，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。 + +### GDR 为实验性 + +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。 + +## 启动流程与降级 + +driver 在**第一次 `tq.init` 之前**完成探测并生成 job 级唯一的 effective config，其余组件（actor / critic / rollout / sft / advantages / actor_fwd）都读同一份，不各自决策。 + +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. 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` 时必须把这份每节点启动资源足迹一并纳入容量规划。 + +降级阶梯： + +``` +GDR → host RDMA （worker 运行时判定，记录 requested 与实际状态） +RDMA → Mooncake/TCP （任一节点无 RDMA 能力，或指定设备缺失） +Mooncake → SimpleStorage（任一节点 mooncake/master 不可用、运行时正确性契约不满足，或 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_requested=false gdr_status=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`。该变量是必填项：未设置时启动直接失败，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。 + +三种情形下的行为： + +| 情形 | 表现 | 处理 | +|---|---|---| +| **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 也会执行 | +| **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 集群**实现，**不承诺同一节点上多个 Relax job 并发**：多 job 并发、端口租约、master 共享机制都不在首期范围内。 + +清理只动本作业拥有的资源： + +- 不使用任何 `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`，不会为了腾空间静默驱逐已经生产但尚未消费的数据。启动前按 **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 启动前检查以下能力： + +- `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，运行时检查用于防止环境被旧包覆盖。 + +正确性守卫强制 `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 上运行： + +```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、空读、清理、字节一致性；`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 集群需由部署侧预先准备）： + +```bash +PYTHONPATH=. python -u scripts/benchmarks/tq_cross_node_bench.py \ + --master :50051 \ + --nodeb-ip \ + --device \ + --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 通过推导真机已经通过。 + +### 参考吞吐区间 + +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 属于交付验收材料，随验收报告存档，不在本文档维护。 + +## 排障表 + +| 现象 | 可能原因 | 处理 | +|---|---|---| +| 启动日志 `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）过期后重试 | +| `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 路径），回环下表现为崩溃而非静默截断 | 同上 | +| 长时间反复起停会话后 `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` 是否打出 | + +### 选卡核验 + +指定设备前先确认端口状态与 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 的收益在多轮之间波动较大，验收结论应基于多轮分布而非单轮数据。 +- Mooncake/TCP（C1）在 0.3.10 上依赖 `MC_STORE_MEMCPY=0` 守卫保证字节正确，且守卫后 get 吞吐低于 SimpleStorage：TCP 档定位为 RDMA 不可用时的正确性兜底，不是性能选项。 +- 消费端节点在 get 过程中中途死亡的端到端行为需要双节点真机验证，未做成自动化测试。 +- Mooncake 传输层自身的超时参数不由 Relax 控制。 diff --git a/relax/backends/megatron/actor.py b/relax/backends/megatron/actor.py index 75143efa6..51c897e66 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, 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,20 @@ 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. + 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(generation) + self._tq_client_generation = None + self.data_system_client = None + except Exception: # destructor must never raise (interpreter shutdown) + return + def init( self, args: Namespace, @@ -187,8 +201,12 @@ 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, + 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 c915cf90f..d121e262e 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,12 @@ 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, + lease_owner=self, + ) 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..9f259d9a0 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,12 @@ 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, + 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.step = 0 diff --git a/relax/components/advantages.py b/relax/components/advantages.py index 3fd183fc4..8cb8780c2 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,12 @@ 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", + lease_owner=self, + ) self.step = 0 async def run(self) -> None: diff --git a/relax/components/base.py b/relax/components/base.py index 7e1425361..c84088187 100644 --- a/relax/components/base.py +++ b/relax/components/base.py @@ -89,6 +89,24 @@ 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: + return + try: + from relax.utils.tq_lifecycle import 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 @property def _logger(self): diff --git a/relax/components/critic.py b/relax/components/critic.py index ef01bbf5f..b3e2240d5 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,12 @@ 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, + lease_owner=self, + ) 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..2da98f0df 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,12 @@ 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", + 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 15ca2b15c..939b39dbf 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,12 @@ 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, + lease_owner=self, + ) self._dataset: Any | None = None self._eval_dataset: Any | None = None diff --git a/relax/core/controller.py b/relax/core/controller.py index e59f5cfd7..00d63fdd3 100644 --- a/relax/core/controller.py +++ b/relax/core/controller.py @@ -47,7 +47,22 @@ 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_mooncake_runtime_contract, +) +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 @@ -131,6 +146,8 @@ 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 + 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). @@ -138,42 +155,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 @@ -242,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 ) @@ -271,23 +297,229 @@ 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": { - "SimpleStorage": { - "total_storage_size": total_storage_size, - "num_data_storage_units": self.config.num_data_storage_units, - }, - }, + "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 + + if getattr(self.config, "tq_storage_backend", "simple") == "simple": + # 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 + 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 + + 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, + ) + 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. + + ``--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. + 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": + 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", "") + 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 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, + use_gdr=getattr(self.config, "tq_use_gdr", False), + rdma_mode=mode, + ) + + # 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={mode} but the requested Mooncake path is unavailable: " + f"{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. 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 + # 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 _close_data_system(self) -> None: + """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 def _deploy_metrics_service(self): """Deploy the MetricsService as a lightweight Ray Serve deployment. @@ -806,6 +1038,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: @@ -1006,7 +1243,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/distributed/ray/rollout.py b/relax/distributed/ray/rollout.py index db8255b29..5a4887a73 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, detach_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,12 @@ 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", + lease_owner=self, + ) 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.") @@ -917,6 +921,13 @@ 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. + 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/arguments.py b/relax/utils/arguments.py index f7d1e710e..7aaf6a3b7 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/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/relax/utils/rdma_probe.py b/relax/utils/rdma_probe.py new file mode 100644 index 000000000..f446cfc20 --- /dev/null +++ b/relax/utils/rdma_probe.py @@ -0,0 +1,652 @@ +# 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_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}" + ) + 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 = "", *, probe_rdma: bool = True) -> 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. + 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()) + + # 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 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) + + # 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 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" + # 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" + 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 = "", + 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. + + 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 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] + + import ray + 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: + from relax.utils.rdma_probe import probe_node as _probe + + 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, probe_rdma) + 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_backend: str, + 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. + + 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. + 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="", + 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) + + 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=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: + 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, + 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="", + ) + + # 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..2cd646f34 --- /dev/null +++ b/relax/utils/tq_config.py @@ -0,0 +1,289 @@ +# 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 inspect +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__) + +# --------------------------------------------------------------------------- +# 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 + + +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. + """ + 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 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. + + 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. + """ + ensure_mooncake_correctness_guards() + + from transfer_queue.storage.managers.base import KVStorageManager + + 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 +# --------------------------------------------------------------------------- + + +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. + "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 required + ``MC_MASTER_ADDRESS`` env var (see + :func:`resolve_mooncake_master_address`). + 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() + + 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 resolve_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 +# --------------------------------------------------------------------------- + + +# 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 + + +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. + + 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 + capacity_batch = resolve_tq_capacity_batch_size(args) + 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 capacity_batch * n_samples * per_sample + + +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) + capacity_batch = resolve_tq_capacity_batch_size(args) + payload = estimate_payload_bytes(args) + needed = payload * (max_staleness + 1) + available = resolve_global_segment_size() + + if needed > available: + return ( + f"MooncakeStore segment capacity insufficient: worst-case in-flight payload " + 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." + ) + 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_correctness.py b/relax/utils/tq_correctness.py new file mode 100644 index 000000000..0ab4e56a6 --- /dev/null +++ b/relax/utils/tq_correctness.py @@ -0,0 +1,69 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""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, then installs temporary, +version-gated runtime patches for the remaining gaps of the pinned revision +(per-retry result validation, raising non-idempotent removal failures, and a +strict production-status ACK). The exact applicability and removal condition +remain isolated in :mod:`relax.utils.tq_mooncake_patches`. + +The pinned mooncake 0.3.10 additionally corrupts TCP-protocol transfers +through its auto-enabled memcpy fast path, so that path is force-disabled +here and an explicit enable is rejected (see :func:`_enforce_safe_memcpy`). +""" + +from __future__ import annotations + +import os + + +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 + ``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. + """ + 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: + """Validate the installed stack and install pinned Mooncake runtime guards. + + Checks the memcpy environment contract and required retry APIs before + installing the exact-version patches for remaining upstream gaps. + """ + _enforce_safe_memcpy() + try: + from transfer_queue.storage.clients.mooncake_client import MooncakeStoreClient + 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)) + + # Version-gated runtime patches for the pinned revision's remaining gaps; + # see relax/utils/tq_mooncake_patches.py for scope and removal condition. + from relax.utils.tq_mooncake_patches import install_mooncake_loss_guards + + install_mooncake_loss_guards() diff --git a/relax/utils/tq_lifecycle.py b/relax/utils/tq_lifecycle.py new file mode 100644 index 000000000..131777486 --- /dev/null +++ b/relax/utils/tq_lifecycle.py @@ -0,0 +1,734 @@ +# 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 os +import threading +import time +import uuid +from dataclasses import dataclass +from typing import Any + +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" +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.""" + + 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 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.""" + + +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 _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", {}) + 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 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 + 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. + + 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: + raise RuntimeError(f"Failed to kill TransferQueueController: {e}") from e + + deadline = time.time() + timeout + while time.time() < deadline: + try: + ray.get_actor(CONTROLLER_NAME, namespace=CONTROLLER_NAMESPACE) + except ValueError: + return + 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. + + 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 _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 _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, + lease_owner: Any | 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). + + 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. + """ + 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(generation: int | None = None) -> 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). 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. + """ + 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]: + """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 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, max_calls=1) + 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. + + 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.") + detach_tq_client() + return + + 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}") + + +@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]: + _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 + + def close(self) -> None: + close_tq_and_unmount(is_owner=self._owns_controller) + + def detach(self) -> None: + detach_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(). + 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 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) + + +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. + 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: + return _attempt(conf) + except Exception as primary_error: + if isinstance(primary_error, TqConfigurationMismatch) or 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/relax/utils/tq_mooncake_patches.py b/relax/utils/tq_mooncake_patches.py new file mode 100644 index 000000000..a4d6b2abe --- /dev/null +++ b/relax/utils/tq_mooncake_patches.py @@ -0,0 +1,261 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Version-gated runtime loss guards for the pinned TransferQueue revision. + +The TransferQueue revision currently pinned by Relax validates the first +Mooncake batch result but not every retry result, logs removal failures +without raising, and treats a missing or negative production-status ACK as a +successful notification. Those behaviours can turn an explicit storage or +controller failure into silent data loss. Until the equivalent checks land +upstream, this module patches the pinned revision at runtime; installation is +process-local and idempotent. + +These are monkey patches over *private* upstream internals +(``MooncakeStoreClient.__init__``, ``StorageManager._notify_and_wait``), so +they are gated on the exact pinned package version and installed VCS revision: +any other build refuses to start rather than running unvalidated patches (see +:func:`_require_pinned_transfer_queue`). + +Removal condition: delete this module and its single call site in +:func:`relax.utils.tq_correctness.ensure_mooncake_correctness_guards` once the +pinned TransferQueue itself validates every batch/retry result, raises on +non-idempotent removal failure, and requires a positive production-status ACK. +""" + +from __future__ import annotations + +import asyncio +import json +import re +from functools import wraps +from importlib import metadata +from pathlib import Path +from typing import Any +from uuid import uuid4 + +from relax.utils.logging_utils import get_logger + + +logger = get_logger(__name__) + +_PATCH_MARKER = "_relax_mooncake_correctness_guards_v1" +# Mooncake's idempotent-delete result; the pinned upstream clear path accepts it. +_MOONCAKE_OBJECT_NOT_FOUND = -704 + +# Exact builds these patches were written and forensically validated against. +_PATCHED_TQ_BUILDS = { + "0.1.10.dev0": frozenset({"58054a33834aadbcf76aacd6b1e32e25c030f2c9"}), +} + + +def _installed_transfer_queue_revision(transfer_queue_module: Any) -> str | None: + """Read the installed VCS revision from standard PEP 610 metadata.""" + try: + distribution = metadata.distribution("transferqueue") + if str(distribution.version) != str(getattr(transfer_queue_module, "__version__", "unknown")): + return None + + module_file = getattr(transfer_queue_module, "__file__", None) + if not module_file: + return None + module_path = Path(module_file).resolve() + package_root = Path(distribution.locate_file("transfer_queue")).resolve() + module_path.relative_to(package_root) + + direct_url = distribution.read_text("direct_url.json") + except (metadata.PackageNotFoundError, AttributeError, OSError, TypeError, UnicodeError, ValueError): + return None + if direct_url is None: + return None + + try: + provenance = json.loads(direct_url) + except (json.JSONDecodeError, TypeError): + return None + if not isinstance(provenance, dict): + return None + + vcs_info = provenance.get("vcs_info") + if not isinstance(vcs_info, dict) or vcs_info.get("vcs") != "git": + return None + commit_id = vcs_info.get("commit_id") + if not isinstance(commit_id, str) or not commit_id.strip(): + return None + revision = commit_id.strip().lower() + if re.fullmatch(r"[0-9a-f]{40}", revision) is None: + return None + return revision + + +def _require_pinned_transfer_queue() -> None: + """Refuse to patch any TransferQueue source revision not validated here.""" + import transfer_queue + + version = str(getattr(transfer_queue, "__version__", "unknown")) + expected_revisions = _PATCHED_TQ_BUILDS.get(version) + if expected_revisions is None: + raise RuntimeError( + f"transfer_queue {version} is not covered by Relax's Mooncake loss guards " + f"(validated versions: {', '.join(_PATCHED_TQ_BUILDS)}). These guards replace " + "private upstream internals (MooncakeStoreClient.__init__, " + "StorageManager._notify_and_wait); re-validate them against the new pin and " + "extend _PATCHED_TQ_BUILDS, or delete relax/utils/tq_mooncake_patches.py " + "entirely if the fixes have landed upstream." + ) + + revision = _installed_transfer_queue_revision(transfer_queue) + if revision not in expected_revisions: + actual_revision = revision or "unknown" + raise RuntimeError( + f"transfer_queue {version} revision {actual_revision} is not covered by Relax's Mooncake loss guards " + f"(validated revisions: {', '.join(sorted(expected_revisions))}). The exact source revision is required " + "because these guards replace private upstream internals; install the validated pin or re-validate " + "the guards before extending _PATCHED_TQ_BUILDS." + ) + + +def _validate_result_count(operation: str, keys: list[str], results: Any) -> None: + """Require one Mooncake result code for every requested key.""" + 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: + if name == "_store": + raise AttributeError(name) + 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 not in (0, _MOONCAKE_OBJECT_NOT_FOUND) + ] + if failures: + detail = ", ".join(f"{key}={code}" for key, code in failures) + raise RuntimeError(f"batch_remove failed: {detail}") + 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 = None + try: + sock = create_zmq_socket( + ctx=self.zmq_context, + socket_type=zmq.DEALER, + ip=self.controller_info.ip, + identity=identity, + ) + sock.setsockopt(zmq.LINGER, 0) + sock.connect(self.controller_info.to_addr("request_handle_socket")) + + await sock.send_multipart(request_msg) + loop = asyncio.get_running_loop() + deadline = loop.time() + tq_base.TQ_DATA_UPDATE_RESPONSE_TIMEOUT + + 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 sock is not None and not sock.closed: + sock.close(linger=0) + except Exception as error: # pragma: no cover - best-effort socket cleanup + logger.debug(f"Failed to close TransferQueue notification socket: {error}") + + +def _install_store_guards(client_cls: type) -> None: + if getattr(client_cls, _PATCH_MARKER, False): + return + + original_init = client_cls.__init__ + + @wraps(original_init) + def guarded_init(self: Any, *args: Any, **kwargs: Any) -> None: + original_init(self, *args, **kwargs) + store = getattr(self, "_store", None) + if store is not None and not isinstance(store, _StrictMooncakeStoreProxy): + self._store = _StrictMooncakeStoreProxy(store) + + client_cls.__init__ = guarded_init + setattr(client_cls, _PATCH_MARKER, True) + + +def _install_notification_guards(manager_cls: type) -> None: + if getattr(manager_cls, _PATCH_MARKER, False): + return + + original_notify = manager_cls.notify_data_update + + @wraps(original_notify) + async def guarded_notify(self: Any, *args: Any, **kwargs: Any) -> None: + if not getattr(self, "controller_info", None): + raise RuntimeError("TransferQueue storage manager has no controller for production-status notification") + await original_notify(self, *args, **kwargs) + + manager_cls.notify_data_update = guarded_notify + manager_cls._notify_and_wait = _strict_notify_and_wait + setattr(manager_cls, _PATCH_MARKER, True) + + +def install_mooncake_loss_guards() -> None: + """Install and verify all runtime guards (idempotent, process-local).""" + _require_pinned_transfer_queue() + from transfer_queue.storage.clients.mooncake_client import MooncakeStoreClient + from transfer_queue.storage.managers.base import StorageManager + + _install_store_guards(MooncakeStoreClient) + _install_notification_guards(StorageManager) + + if not getattr(MooncakeStoreClient, _PATCH_MARKER, False) or not getattr(StorageManager, _PATCH_MARKER, False): + raise RuntimeError("Failed to install Mooncake silent-data-loss guards") diff --git a/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/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 new file mode 100644 index 000000000..31b39ca97 --- /dev/null +++ b/scripts/benchmarks/tq_cross_node_bench.py @@ -0,0 +1,770 @@ +#!/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 +""" + +from __future__ import annotations + +import argparse +import csv +import hashlib +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( + "--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).", + ) + 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. + """ + 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 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], + ) + + +_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 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]]: + """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 _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 + 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_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, + 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( + { + "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, 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() + 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() + # 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 + + +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() + 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, + ) + + 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())) + fields = sorted(payload.keys()) + 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] = [] + 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, + ) + 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), + } + ) + + 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:<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']}", + 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) + + +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..1399054d7 --- /dev/null +++ b/scripts/benchmarks/tq_rdma_bench.py @@ -0,0 +1,417 @@ +#!/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 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/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 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/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_rdma_probe.py b/tests/utils/test_rdma_probe.py new file mode 100644 index 000000000..1763b47f0 --- /dev/null +++ b/tests/utils/test_rdma_probe.py @@ -0,0 +1,640 @@ +# 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, + gdr_eligible=protocol == "rdma", + ) + + +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, + seq_length=8192, + ) + 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_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" + + def test_empty_results_falls_back(self): + eff = reduce_results( + [], + requested_backend="mooncake", + requested_device="", + use_gdr=False, + ) + 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) +# --------------------------------------------------------------------------- + + +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" + assert result.gdr_eligible is False + + 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" + assert result.gdr_eligible is True + + @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_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 +# --------------------------------------------------------------------------- + + +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_backend="mooncake", + requested_device="", + use_gdr=False, + ) + 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 +# --------------------------------------------------------------------------- + + +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", gdr=False, 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", gdr=False, 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_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", + ) + 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="", 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_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="", 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"): + 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() 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_dataplane_behavior.py b/tests/utils/test_tq_dataplane_behavior.py new file mode 100644 index 000000000..cb2509924 --- /dev/null +++ b/tests/utils/test_tq_dataplane_behavior.py @@ -0,0 +1,367 @@ +# 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. +""" + +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 _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).""" + 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 + + +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 new file mode 100644 index 000000000..8cdd713a9 --- /dev/null +++ b/tests/utils/test_tq_failure_paths.py @@ -0,0 +1,1116 @@ +# 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), 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``) + +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 asyncio +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 +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") +_RUN_REAL_CAPACITY = os.environ.get("RELAX_RUN_REAL_MOONCAKE_CAPACITY_TEST") == "1" + + +# --------------------------------------------------------------------------- +# 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 + + +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 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) + 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() + + +# --------------------------------------------------------------------------- +# 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 = 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(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(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(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() + + +# --------------------------------------------------------------------------- +# 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) + + 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 + + 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_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_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 generation: calls.append(generation)) + component = Base() + component.data_system_client = object() + component._tq_client_generation = 7 + component.__del__() + 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 generation: calls.append(generation)) + component = Base() + component.__del__() + 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 +# --------------------------------------------------------------------------- + + +class TestInitializeTqWithFallback: + @staticmethod + 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": []} + 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_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_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") + 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 _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 + + 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] + + 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() + 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 +# --------------------------------------------------------------------------- + + +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 + + +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", +) +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]) + + +@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) +# --------------------------------------------------------------------------- + + +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 + + 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): + """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) + assert status == "ok", f"{status}: {detail}" diff --git a/tests/utils/test_tq_mooncake_patches.py b/tests/utils/test_tq_mooncake_patches.py new file mode 100644 index 000000000..1c9c36e69 --- /dev/null +++ b/tests/utils/test_tq_mooncake_patches.py @@ -0,0 +1,382 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Tests for the version-gated Mooncake runtime loss guards. + +Moved out of ``test_tq_failure_paths.py`` together with the patches themselves +(``relax/utils/tq_mooncake_patches.py``); everything here is deleted with that +module once the fixes land upstream. +""" + +from __future__ import annotations + +import importlib.util +from types import SimpleNamespace + +import pytest + +from relax.utils.tq_mooncake_patches import ( + _install_notification_guards, + _install_store_guards, + _installed_transfer_queue_revision, + _require_pinned_transfer_queue, + _strict_notify_and_wait, + _StrictMooncakeStoreProxy, +) + + +def _has_real_submodule(dotted: str) -> bool: + """True only if a REAL transfer_queue submodule is importable. + + CI installs a single-file ``transfer_queue`` stub; + ``transfer_queue.storage`` does not exist there, so tests that touch the + real MooncakeStoreClient skip on CPU CI and run only where real + TransferQueue is installed. + """ + try: + return importlib.util.find_spec(dotted) is not None + except (ImportError, ValueError, TypeError): + # CI's single-file transfer_queue stub returns a dummy for ``__path__``, + # so find_spec on a submodule raises TypeError instead of returning None. + return False + + +_REAL_MOONCAKE_CLIENT = _has_real_submodule("transfer_queue.storage.clients.mooncake_client") + + +class _SequenceStore: + """Return a configured result sequence from low-level Mooncake calls.""" + + def __init__(self, results: list[list[int]]) -> None: + self.results = iter(results) + + def batch_upsert_from(self, keys, ptrs, sizes, config=None): + return next(self.results) + + def batch_get_into(self, keys, ptrs, sizes): + return next(self.results) + + def batch_remove(self, keys, force=True): + return next(self.results) + + +class _FakeNotifySocket: + def __init__(self, connect_error: Exception | None = None) -> None: + self.closed = False + self.connect_error = connect_error + + def setsockopt(self, *args, **kwargs) -> None: + pass + + def connect(self, *args, **kwargs) -> None: + if self.connect_error is not None: + raise self.connect_error + + async def send_multipart(self, request) -> None: + pass + + async def recv_multipart(self, copy=False): + return [b"ack"] + + def close(self, linger=0) -> None: + self.closed = True + + +def _client_with_store(store) -> object: + """A MooncakeStoreClient with only ``_store``/``replica_config`` wired up. + + ``__init__`` is skipped on purpose: it would need a live mooncake master. + """ + from transfer_queue.storage.clients.mooncake_client import MooncakeStoreClient + + client = object.__new__(MooncakeStoreClient) + client._store = store + client.replica_config = None + return client + + +class TestVersionGate: + """Patches refuse to install on any transfer_queue they were not written + for.""" + + _PINNED_REVISION = "58054a33834aadbcf76aacd6b1e32e25c030f2c9" + + def test_unpinned_version_is_rejected(self, monkeypatch): + import transfer_queue + + monkeypatch.setattr(transfer_queue, "__version__", "9.9.9", raising=False) + with pytest.raises(RuntimeError, match="not covered by Relax"): + _require_pinned_transfer_queue() + + def test_pinned_version_and_revision_are_accepted(self, monkeypatch): + import transfer_queue + + monkeypatch.setattr(transfer_queue, "__version__", "0.1.10.dev0", raising=False) + monkeypatch.setattr( + "relax.utils.tq_mooncake_patches._installed_transfer_queue_revision", + lambda module: self._PINNED_REVISION, + ) + _require_pinned_transfer_queue() + + def test_same_version_with_different_revision_is_rejected(self, monkeypatch): + import transfer_queue + + monkeypatch.setattr(transfer_queue, "__version__", "0.1.10.dev0", raising=False) + monkeypatch.setattr( + "relax.utils.tq_mooncake_patches._installed_transfer_queue_revision", + lambda module: "0" * 40, + ) + with pytest.raises(RuntimeError, match="revision .* is not covered"): + _require_pinned_transfer_queue() + + def test_missing_revision_metadata_is_rejected(self, monkeypatch): + import transfer_queue + + monkeypatch.setattr(transfer_queue, "__version__", "0.1.10.dev0", raising=False) + monkeypatch.setattr( + "relax.utils.tq_mooncake_patches._installed_transfer_queue_revision", + lambda module: None, + ) + with pytest.raises(RuntimeError, match="revision unknown is not covered"): + _require_pinned_transfer_queue() + + @pytest.mark.parametrize( + ("direct_url", "expected"), + [ + (None, None), + ("not-json", None), + ("{}", None), + ('{"vcs_info": {}}', None), + ('{"vcs_info": {"vcs": "hg", "commit_id": "58054a33834aadbcf76aacd6b1e32e25c030f2c9"}}', None), + ('{"vcs_info": {"vcs": "git", "commit_id": "abc123"}}', None), + ( + '{"vcs_info": {"vcs": "git", "commit_id": " 58054A33834AADBCF76AACD6B1E32E25C030F2C9 "}}', + _PINNED_REVISION, + ), + ], + ) + def test_revision_metadata_is_parsed_fail_closed(self, monkeypatch, tmp_path, direct_url, expected): + package_root = tmp_path / "installed" / "transfer_queue" + + class Distribution: + version = "0.1.10.dev0" + + def locate_file(self, filename): + assert filename == "transfer_queue" + return package_root + + def read_text(self, filename): + assert filename == "direct_url.json" + return direct_url + + monkeypatch.setattr( + "relax.utils.tq_mooncake_patches.metadata.distribution", + lambda name: Distribution(), + ) + module = SimpleNamespace( + __version__="0.1.10.dev0", + __file__=package_root / "__init__.py", + ) + assert _installed_transfer_queue_revision(module) == expected + + def test_distribution_version_mismatch_is_unverifiable(self, monkeypatch, tmp_path): + class Distribution: + version = "0.1.10.dev1" + + monkeypatch.setattr( + "relax.utils.tq_mooncake_patches.metadata.distribution", + lambda name: Distribution(), + ) + module = SimpleNamespace( + __version__="0.1.10.dev0", + __file__=tmp_path / "transfer_queue" / "__init__.py", + ) + assert _installed_transfer_queue_revision(module) is None + + def test_shadowed_module_is_unverifiable(self, monkeypatch, tmp_path): + package_root = tmp_path / "installed" / "transfer_queue" + + class Distribution: + version = "0.1.10.dev0" + + def locate_file(self, filename): + return package_root + + def read_text(self, filename): + return '{"vcs_info": {"vcs": "git", "commit_id": "58054a33834aadbcf76aacd6b1e32e25c030f2c9"}}' + + monkeypatch.setattr( + "relax.utils.tq_mooncake_patches.metadata.distribution", + lambda name: Distribution(), + ) + shadowed_module = SimpleNamespace( + __version__="0.1.10.dev0", + __file__=tmp_path / "shadow" / "transfer_queue.py", + ) + assert _installed_transfer_queue_revision(shadowed_module) is None + + +class TestMooncakeCorrectnessGuardPrimitives: + """Low-level response validation stays runnable on the CPU-only CI stub.""" + + 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_get_short_result_is_raised(self): + store = _StrictMooncakeStoreProxy(_SequenceStore([[0]])) + with pytest.raises(RuntimeError, match="returned 1 results, expected 2"): + store.batch_get_into(["k0", "k1"], [1, 2], [8, 8]) + + def test_remove_object_not_found_is_allowed(self): + store = _StrictMooncakeStoreProxy(_SequenceStore([[0, -704]])) + assert store.batch_remove(["k0", "k1"], force=True) == [0, -704] + + def test_remove_non_idempotent_failure_is_raised(self): + store = _StrictMooncakeStoreProxy(_SequenceStore([[0, -1]])) + with pytest.raises(RuntimeError, match="batch_remove failed"): + store.batch_remove(["k0", "k1"], force=True) + + def test_missing_wrapped_store_raises_attribute_error(self): + store = object.__new__(_StrictMooncakeStoreProxy) + with pytest.raises(AttributeError): + store.close + + def test_store_guard_installation_is_idempotent(self): + raw_store = _SequenceStore([[0]]) + + class Client: + def __init__(self): + self._store = raw_store + + _install_store_guards(Client) + guarded_init = Client.__init__ + _install_store_guards(Client) + + client = Client() + assert Client.__init__ is guarded_init + assert isinstance(client._store, _StrictMooncakeStoreProxy) + assert client._store._store is raw_store + + +class TestNotificationGuardPrimitives: + @pytest.mark.asyncio + async def test_guarded_notify_rejects_missing_controller(self): + class Manager: + controller_info = None + + async def notify_data_update(self): + raise AssertionError("original notify must not run without a controller") + + async def _notify_and_wait(self, request_msg): + pass + + _install_notification_guards(Manager) + guarded_notify = Manager.notify_data_update + _install_notification_guards(Manager) + assert Manager.notify_data_update is guarded_notify + with pytest.raises(RuntimeError, match="has no controller"): + await Manager().notify_data_update() + + +@pytest.mark.skipif( + not _REAL_MOONCAKE_CLIENT, + 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]) + + def test_get_retry_short_result_is_never_treated_as_success(self, monkeypatch): + monkeypatch.setattr("transfer_queue.storage.clients.mooncake_client.RETRY_DELAY_SECONDS", 0) + store = _StrictMooncakeStoreProxy(_SequenceStore([[-1, -1], [0]])) + client = _client_with_store(store) + with pytest.raises(RuntimeError, match="returned 1 results, expected 2"): + client._batch_get_into_with_retry(["k0", "k1"], [1, 2], [8, 8]) + + @pytest.mark.asyncio + async def test_negative_production_status_ack_is_raised(self, monkeypatch): + from transfer_queue.utils import zmq_utils + + 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_positive_production_status_ack_returns(self, monkeypatch): + from transfer_queue.utils import zmq_utils + + socket = _FakeNotifySocket() + monkeypatch.setattr(zmq_utils, "create_zmq_socket", lambda **kwargs: socket) + monkeypatch.setattr( + zmq_utils.ZMQMessage, + "deserialize", + staticmethod( + lambda messages: SimpleNamespace( + request_type=zmq_utils.ZMQRequestType.NOTIFY_DATA_UPDATE_ACK, + body={"success": True, "partition_id": "p0"}, + ) + ), + ) + manager = SimpleNamespace( + storage_manager_id="guard-test", + zmq_context=object(), + controller_info=SimpleNamespace(ip="redacted", to_addr=lambda name: "inproc://controller"), + ) + await _strict_notify_and_wait(manager, [b"request"]) + assert socket.closed is True + + @pytest.mark.asyncio + async def test_connect_failure_closes_notification_socket(self, monkeypatch): + from transfer_queue.utils import zmq_utils + + socket = _FakeNotifySocket(connect_error=ConnectionError("controller unavailable")) + monkeypatch.setattr(zmq_utils, "create_zmq_socket", lambda **kwargs: socket) + manager = SimpleNamespace( + storage_manager_id="guard-test", + zmq_context=object(), + controller_info=SimpleNamespace(ip="redacted", to_addr=lambda name: "inproc://controller"), + ) + with pytest.raises(ConnectionError, match="controller unavailable"): + await _strict_notify_and_wait(manager, [b"request"]) + assert socket.closed is True + + @pytest.mark.asyncio + async def test_missing_production_status_ack_is_bounded(self, monkeypatch): + from transfer_queue.storage.managers import base as tq_base + from 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