diff --git a/.gitignore b/.gitignore index e55036cb7..22481d5e3 100644 --- a/.gitignore +++ b/.gitignore @@ -54,3 +54,6 @@ tensorboard_log # .github .github/copilot-instructions.md env.sh + +# Machine-local multimodal acceptance artifacts (hundreds of MB, never commit) +tests/fixtures/ diff --git a/docker/Dockerfile b/docker/Dockerfile index c473e89b0..3e8d7a4aa 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -116,7 +116,7 @@ COPY requirements.txt /tmp/requirements.txt RUN pip install --ignore-installed PyJWT && \ pip install -r /tmp/requirements.txt --no-cache-dir && \ pip install --no-cache-dir "compressed_tensors>=0.13.0" tensordict==0.10.0 pyvers==0.1.0 'nvidia-modelopt[hf]==0.44.0' --no-deps && \ - pip install "transferqueue @ git+https://github.com/redai-infra/TransferQueue.git@58054a33834aadbcf76aacd6b1e32e25c030f2c9" --no-deps + pip install "transferqueue @ git+https://github.com/redai-infra/TransferQueue.git@6c7a587292910af0827f027de99e005e1900310e" --no-deps # sgl-router: override the official wheel (pulled by requirements.txt above) with # slime's r3-capable fork. The official sglang-router drops the routed_experts diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu index 89ed93f0c..0d4055f5a 100644 --- a/docker/Dockerfile.npu +++ b/docker/Dockerfile.npu @@ -68,7 +68,7 @@ RUN pip install pyyaml && \ RUN cd /root && rm -rf /root/pytorch && \ pip install triton-ascend==3.2.0 && \ pip install tensordict==0.10.0 pyvers==0.1.0 --no-deps && \ - pip install "transferqueue @ git+https://github.com/redai-infra/TransferQueue.git@58054a33834aadbcf76aacd6b1e32e25c030f2c9" --no-deps + pip install "transferqueue @ git+https://github.com/redai-infra/TransferQueue.git@6c7a587292910af0827f027de99e005e1900310e" --no-deps # Clone Megatron-LM, MindSpeed, MindSpeed-Bridge, Megatron-Bridge and install RUN git clone https://gitcode.com/ascend/MindSpeed.git /root/MindSpeed && \ diff --git a/docs/draft/transfer_queue_rdma.md b/docs/draft/transfer_queue_rdma.md new file mode 100644 index 000000000..a63892b56 --- /dev/null +++ b/docs/draft/transfer_queue_rdma.md @@ -0,0 +1,147 @@ +# TransferQueue host-RDMA 使用与运维 + +## 范围与配置 + +Relax 默认使用 TransferQueue SimpleStorage。首期 RDMA 支持只接入 MooncakeStore/host-RDMA，不改变 payload 形状或数据分发语义，也不支持 GDR。生产路径只有 host-RDMA 和 SimpleStorage；Mooncake/TCP 仅作为跨节点 benchmark 的 C1 对照。 + +| 参数 | 语义 | +|---|---| +| `--tq-rdma-mode=off` | 默认值；保持 SimpleStorage 数据路径，不执行 Mooncake 检查或集群 handshake；worker 的 TQ attach 仍有 60 秒边界 | +| `--tq-rdma-mode=auto` | 尝试 host-RDMA；任一检查或节点 attach 失败时，完成清理后统一回退 SimpleStorage | +| `--tq-rdma-mode=required` | 要求 host-RDMA；不可用时完成清理并终止启动，不允许降级 | +| `--tq-rdma-device=` | 指定 RDMA 设备；空值由 Mooncake 原生逻辑选择，多 HCA 环境建议显式指定 | + +Relax 固定 Mooncake `use_gdr=false`。GDR 若有实际需求，应由独立 PR 实现并做专项验证。 + +## 部署前提 + +首期只支持**单任务独占 Ray 集群**：同一个 Ray cluster 在初始化和运行期间只能有一个 Relax job，不支持 concurrent initializer、多 job admission 或复用其他作业的 TransferQueue controller。 + +Mooncake master 由部署环境管理，Relax 不启动、重启或停止它。driver 必须设置外部 endpoint： + +```bash +export MC_MASTER_ADDRESS=master.example:50051 +``` + +driver 将该 endpoint 写入 job-level TQ config，owner 和 worker attach 复用已存储配置，因此 worker 节点不要求重复设置该环境变量；但所有节点都必须能访问同一个 endpoint。Relax 只检查 `host:port` 格式；DNS、路由、防火墙和 master 健康状态由真实初始化与 attach 验证。 + +## 启动、回退与清理 + +`off` 直接初始化 SimpleStorage，不创建 owner actor 或执行 Mooncake handshake；各 worker 仍通过有界 helper attach 到 SimpleStorage，以避免旧版无界等待。启动时还会回收不可用的半初始化 controller，并因首期独占集群约束拒绝复用健康的既有 controller。`auto` 和 `required` 按以下顺序执行： + +1. 校验 mode、TransferQueue/Mooncake correctness contract、master 格式和 segment 容量。 +2. 在独立 owner actor 中有界执行 Mooncake `tq.init`。 +3. 在每个 ALIVE Ray 节点运行一次性 worker（`max_calls=1`、`max_retries=0`），真实 attach 并确认 `MooncakeStorageManager` 和 `protocol=rdma`，随后 detach。 +4. 全部节点通过后启用 host-RDMA；任一失败则等待 worker 终止，并清理 owner、controller 和 segment。 +5. 清理可确认时，`auto` 初始化 SimpleStorage，`required` 报错；清理无法确认时两种模式都 fail closed。 + +Relax 不再维护 `/sys` 启发式能力探测。真实 attach 是运行时能力判据，但它只证明 manager、配置 protocol 和 setup 成功；线路是否真正传输 RDMA 数据必须由 benchmark counter 的 wire proof 证明。 + +owner 初始化和 worker attach 默认各有 60 秒边界；worker attach 可通过 `RELAX_TQ_ATTACH_TIMEOUT_SECONDS` 调整。RDMA 不可用时，`auto` 需要完成真实初始化、全节点 attach 和清理，启动可能达到分钟级。已知无需 RDMA 的作业应使用 `off`。 + +健康的既有 controller 不会被接管或关闭，启动会直接失败；半初始化、超时或已死亡的 controller 才允许回收。全局 `tq.close()` 只由 owner 调用，普通 worker 只 detach 本地 client，Mooncake master 始终不由 Relax 管理。 + +## Segment 与容量 + +真实 handshake 会在每个 ALIVE 节点（包括 CPU-only head）创建 Mooncake client，瞬时挂载并注册完整 client segment，结束后立即 detach。默认配置为每 client 4 GiB global segment 和 1 GiB local buffer；实际 RSS、锁页内存及注册资源由 Mooncake 实现决定。内存或 `memlock` 不足会表现为 attach 失败。 + +global segment 可按部署容量调整： + +```bash +export RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB=8 +``` + +启动前容量预检采用保守上界：文本按 `seq_length × 32 B`；多模态按当前支持的最大 transported tensor layout（16×16 spatial patch、temporal patch 2、RGB、2×2 spatial merge、float32）另加 `seq_length × 24,576 B`。因此 8,192 token 的单样本多模态 tensor 上界约 192 MiB。最终再乘 rollout batch、`n_samples_per_prompt` 和 `max_staleness + 1`。容量不足时 `auto` 回退 SimpleStorage，`required` 失败。增大 segment 时必须同步规划每个节点的瞬时内存与 `memlock` 足迹。 + +异常退出时 client segment 可能继续注册到 master，直到 Mooncake `client_ttl`（默认 30 秒）到期。 + +## Correctness 与依赖 gate + +当前 Relax pin 为 TransferQueue `6c7a587292910af0827f027de99e005e1900310e`，并要求 `MOONCAKE_CORRECTNESS_CONTRACT_VERSION >= 1`。Contract version 1 保证： + +- `batch_upsert_from` 和 `batch_get_into` 的每次 batch/retry 都校验返回结果与请求 key 等长； +- `batch_remove` 的非幂等失败向调用方传播； +- `NOTIFY_DATA_UPDATE_ACK` 验证 positive ACK，controller 拒绝时 producer 不会按成功结束。 + +这些修复应在 TransferQueue 上游实现；Relax 不使用 monkey patch 替代依赖修复。 + +测试确认 `mooncake-transfer-engine==0.3.10.post2` 的 TCP memcpy 路径存在静默截断风险。Relax 当前会统一强制 `MC_STORE_MEMCPY=0`，显式设置不安全值会拒绝启动；只有在能够可靠识别已修复 build 后，才重新评估是否允许 memcpy 路径。 + +每次验收必须记录 Relax commit SHA、TransferQueue commit 和 Mooncake 版本，不能只记录分支名。 + +## 跨节点验收 + +唯一保留的 benchmark 是 `scripts/benchmarks/tq_cross_node_bench.py`： + +| 档位 | 后端 | 用途 | +|---|---|---| +| C0 | SimpleStorage | 默认路径基线 | +| C1 | Mooncake/TCP | benchmark 对照，不是生产配置或 fallback | +| C2 | Mooncake/host-RDMA | 生产 RDMA candidate | + +每个 protocol 必须使用全新 Python 进程和独立 CSV。示例仅展示 C2；C1 改用 `--protocol tcp` 并删除 `--device`/`--rdma-port`，C0 改用 `--protocol simple` 并额外删除 `--master`： + +```bash +python -u scripts/benchmarks/tq_cross_node_bench.py \ + --protocol rdma \ + --master master.example:50051 \ + --consumer-node-id \ + --device \ + --rdma-port 1 \ + --tcp-device \ + --payload-profiles synthetic multimodal \ + --payload-mib 256 1024 2048 4096 \ + --repeats 5 \ + --csv c2-rdma.csv +``` + +所有档位都必须 byte-exact。C2 只读取明确指定的 HCA/port，要求一次完整 `put → get` round 的 idle-adjusted IB receive bytes 至少覆盖 raw payload 的 80%；完整 round 可以覆盖对象随机落在 producer、consumer 或 owner segment 的情况。C1 要求 idle-adjusted TCP receive bytes 至少覆盖 20%，C0 只要求观测到跨节点 TCP，因为 SimpleStorage unit 可能位于 consumer 本地。端口级 counter 不是 per-flow 指标，正式验收必须使用静默或独占的数据端口；CSV 同时记录 raw delta、紧邻 round 的 idle rate 和扣除后的 proof bytes，不能在共享端口有显著背景流量时宣称 wire proof。 + +C1 会在 driver 及其 Ray worker runtime 中设置 `MC_TCP_ENABLE_CONNECTION_POOL=1`，这是当前 Mooncake/TCP correctness baseline 的组成部分，必须随结果一并记录。benchmark 通过一次性 owner lifecycle 有界初始化 TQ，producer/consumer 均有界 attach；任一失败都非零退出，不执行 backend fallback。 + +CSV 使用 exclusive-create，不覆盖既有文件。每个 warmup/测量 round 都先落盘再执行 byte-exact/wire gate，失败行包含稳定的 `error_kind`；每轮无论成功失败都清理 partition。每行还记录 Relax SHA、安装的 TransferQueue VCS commit 和 Mooncake package version，且 Relax tracked worktree 不干净时拒绝作为正式验收运行。 + +真实多模态 fixture、原始 CSV、版本信息和性能分布属于 PR 验收附件，不在仓库文档维护生成教程或易过期的性能数字。没有双节点 RDMA 环境时必须明确记录“真机验收未执行”，不能用 mock 结果替代。 + +## 排障 + +| 现象 | 检查与处理 | +|---|---| +| correctness contract 不满足 | 核对各节点 TransferQueue/Mooncake 版本和 capability marker；不要绕过 gate | +| master 缺失、格式错误或不可达 | 检查 driver 的 `MC_MASTER_ADDRESS`，并检查所有节点到该 endpoint 的 DNS、路由、防火墙和 master 服务 | +| segment capacity insufficient | 减少 batch、采样数或 staleness，或增大 `RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB` 并重新核算资源 | +| attach handshake 失败 | 在失败节点检查 HCA port、GID、内存、`memlock` 和到 master 的连接；CPU-only head 也必须满足 attach 条件 | +| manager/protocol 不符 | 清理前一个作业；首期要求单任务独占且不会接管健康的旧 controller | +| `Connection refused` 指向旧 segment | 等待 `client_ttl` 过期，再确认旧 client/segment 已从 master 清理 | +| C1 尾部全零或 SIGSEGV | 确认未显式启用 `MC_STORE_MEMCPY`；C1 仅用于 benchmark | +| 多 HCA 环境建连失败 | 显式设置 `--tq-rdma-device`，不要依赖自动选卡 | + +### 手工核验 HCA、GID、memlock、线路与 master + +```bash +export TQ_RDMA_DEVICE= +export TQ_RDMA_PORT=1 +export TQ_TCP_DEVICE= + +cat "/sys/class/infiniband/${TQ_RDMA_DEVICE}/ports/${TQ_RDMA_PORT}/state" +cat "/sys/class/infiniband/${TQ_RDMA_DEVICE}/ports/${TQ_RDMA_PORT}/gids/0" +cat "/sys/class/infiniband/${TQ_RDMA_DEVICE}/ports/${TQ_RDMA_PORT}/rate" +ulimit -l + +# 在数据传输前后取差值；IB counter 单位为 4-byte words。 +cat "/sys/class/infiniband/${TQ_RDMA_DEVICE}/ports/${TQ_RDMA_PORT}/counters/port_rcv_data" +cat "/sys/class/net/${TQ_TCP_DEVICE}/statistics/rx_bytes" +# 使用通用占位 endpoint 检查 DNS 与 TCP。 +export TQ_MASTER_HOST=master.example TQ_MASTER_PORT=50051 +getent hosts "${TQ_MASTER_HOST}" +nc -vz "${TQ_MASTER_HOST}" "${TQ_MASTER_PORT}" +``` + +master 可达性可用部署环境已有的 DNS/TCP 工具检查；不要把真实 endpoint、hostname 或本地路径写入提交、公开日志和 PR 文档。 + +## 已知边界 + +- 首期不支持 GDR、多 job 或 concurrent initializer。 +- Mooncake 传输层自身的 timeout 不由 Relax 控制。 +- attach 成功不等于 wire proof；C2 合入前仍需真实双节点 byte-exact、wire-proof 和 fully-async smoke。 +- mock/CPU CI 不能替代真实 RDMA 验收，真实测试结果必须关联准确的代码和依赖 SHA。 diff --git a/relax/backends/megatron/actor.py b/relax/backends/megatron/actor.py index 75143efa6..36cea659a 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,18 @@ def _per_step_rollout(self) -> bool: periodic predict steps; Megatron stays awake between.""" return not is_sft_mode(self.args) + def __del__(self) -> None: + # Best-effort detach on graceful teardown; ray.kill / fate-sharing + # kills skip destructors, in which case the Mooncake master TTL + # reclaims the segment. + if getattr(self, "data_system_client", None) is None: + return + try: + detach_tq_client() + self.data_system_client = None + except Exception: # destructor must never raise (interpreter shutdown) + return + def init( self, args: Namespace, @@ -187,8 +199,10 @@ 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, + role=role, + ) if is_megatron_main_rank(): init_tracking(args, primary=False) diff --git a/relax/components/actor.py b/relax/components/actor.py index c915cf90f..d5273e0ec 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,16 +71,16 @@ 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, + role=self.role, + ) - self.steps = ray.get( - self.actor_model.async_init( - config, - role=self.role, - with_ref=config.kl_coef != 0 or config.use_kl_loss, - with_opd_teacher=self.config.opd_teacher_load, - ) + self.steps = self.actor_model.init_and_wait( + config, + role=self.role, + with_ref=config.kl_coef != 0 or config.use_kl_loss, + with_opd_teacher=self.config.opd_teacher_load, ) assert len(set(self.steps)) == 1 diff --git a/relax/components/actor_fwd.py b/relax/components/actor_fwd.py index 707b8add6..831665616 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,10 +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, + role=self.role, + ) self.actor_model = allocate_train_group(args=config, num_gpus=num_gpus, pg=pgs, runtime_env=runtime_env) - ray.get(self.actor_model.async_init(config, role=self.role, with_ref=False)) + self.actor_model.init_and_wait(config, role=self.role, with_ref=False) self.step = 0 async def run(self) -> None: diff --git a/relax/components/advantages.py b/relax/components/advantages.py index 3fd183fc4..a01c7c2e7 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,10 @@ 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, + role="advantages", + ) self.step = 0 async def run(self) -> None: diff --git a/relax/components/base.py b/relax/components/base.py index 7e1425361..8f18e6720 100644 --- a/relax/components/base.py +++ b/relax/components/base.py @@ -90,6 +90,21 @@ def __init__(self) -> None: self._logger_instance = None self._lock = threading.Lock() + def __del__(self) -> None: + # Ray Serve calls the destructor on replica shutdown (normal stop, + # global restart, in-place restart). Components that attached a + # TransferQueue client must detach so a MooncakeStore segment + # deregisters before client_ttl instead of leaving a stale endpoint. + if getattr(self, "data_system_client", None) is None: + return + try: + from relax.utils.tq.lifecycle import detach_tq_client + + detach_tq_client() + self.data_system_client = None + except Exception: # destructor must never raise (interpreter shutdown) + return + @property def _logger(self): """Lazily create and cache a logger for this instance. diff --git a/relax/components/critic.py b/relax/components/critic.py index ef01bbf5f..a8a45f221 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,14 +40,16 @@ 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, + role=self.role, + ) self.critic_model = allocate_train_group( args=config, num_gpus=num_gpus, pg=pgs, role=self.role, runtime_env=runtime_env ) - ray.get(self.critic_model.async_init(config, role=self.role, with_ref=False)) + self.critic_model.init_and_wait(config, role=self.role, with_ref=False) self.step = getattr(self.config, "start_rollout_id", None) or 0 # Wired by controller in colocate PPO to gate wake_up on SGLang offload. self._rollout_barrier: Optional[RolloutOffloadBarrier] = None diff --git a/relax/components/rollout.py b/relax/components/rollout.py index ebf065683..c20de9602 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,10 @@ 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, + role="rollout", + ) self.rollout_manager, self.num_rollout_per_epoch = create_rollout_manager( config, pg, data_source=data_source, runtime_env=runtime_env ) diff --git a/relax/components/sft.py b/relax/components/sft.py index 15ca2b15c..d9a85ff77 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,10 @@ 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, + role=self.role, + ) self._dataset: Any | None = None self._eval_dataset: Any | None = None diff --git a/relax/core/controller.py b/relax/core/controller.py index e59f5cfd7..585c87b22 100644 --- a/relax/core/controller.py +++ b/relax/core/controller.py @@ -21,7 +21,7 @@ raise ImportError( "transfer_queue is out of date (missing StreamingTokenBudgetSampler). Upgrade with:\n" ' pip install "transferqueue @ git+https://github.com/redai-infra/' - 'TransferQueue.git@58054a33834aadbcf76aacd6b1e32e25c030f2c9" --no-deps\n' + 'TransferQueue.git@6c7a587292910af0827f027de99e005e1900310e" --no-deps\n' "or use the latest image." ) from e @@ -48,6 +48,23 @@ shutdown_managed_opd_teacher, ) from relax.utils.s3_model_loader import cleanup_s3_model_weights_from_shm +from relax.utils.tq.config import ( + build_backend_config, + resolve_mooncake_master_address, + resolve_tq_capacity_batch_size, + validate_config, + validate_mooncake_runtime_contract, +) +from relax.utils.tq.lifecycle import ( + TqHandshakeIsolationError, + TqInitResult, + close_tq_owner, + initialize_tq_with_fallback, + reap_unusable_tq_controller, + safe_exception_kind, + uses_mooncake, + verify_cluster_attach, +) from relax.utils.training.ppo_utils import validate_ppo_config from relax.utils.utils import compute_dp_size, recovery_load_path @@ -131,49 +148,68 @@ 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). resolve_sft_num_rollout(self.config) - # Initialize data management system - self._initialize_data_system() - 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() - - 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 + # Include data-system initialization in the constructor cleanup + # transaction: a cluster-handshake orchestration error can happen + # after the owner actor has created global TQ state. + self._initialize_data_system() + self.dcs, self.config.coordinator_url = create_dcs_deployment() + + self._metrics_service_enabled = getattr(config, "use_metrics_service", False) + 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 - 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; closing any initialized TQ state.") + try: + self._close_data_system() + except Exception as cleanup_error: # pragma: no cover - best effort + logger.warning( + f"TQ owner cleanup during failed construction failed ({safe_exception_kind(cleanup_error)})." + ) + raise def _cleanup_s3_model_weights_after_init(self) -> None: """Remove policy weight shards after every startup consumer is @@ -242,11 +278,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 +303,236 @@ 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_rdma_mode", "off") == "off": + # Preserve the upstream SimpleStorage ownership model: the first + # tq.init runs inside Controller and no _TransferQueueOwner actor + # is created. Lifecycle hardening still applies: the F10 reaper + # 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, + ) + # The owner becomes this Controller's cleanup responsibility before + # the cluster handshake. Otherwise a driver-side scheduling error can + # escape while ``self._tq_owner`` is still None and orphan global TQ + # state during constructor failure. + self._tq_owner = init_result.owner + if uses_mooncake(init_result.config): + init_result = self._confirm_mooncake_attach(init_result, fallback_config) + self._tq_owner = init_result.owner + 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("[dataplane] controller ownership=owner") + + def _confirm_mooncake_attach(self, init_result: TqInitResult, fallback_config) -> TqInitResult: + """Validate Mooncake attach/setup on every alive Ray node. + + There is no static ``/sys`` probe. TQ clients live in Ray Serve + replicas and 0-CPU actors with no placement binding, so they may land + on any alive node, and a worker-side ``tq.init`` has no timeout of its + own. Each alive node therefore performs a bounded attach with the + stored config, confirms a Mooncake manager whose client is configured + for RDMA, and detaches. This is configuration/setup evidence, not + negotiated-transport or wire proof. + Failures are aggregated here so ``auto`` converges the whole job on one + backend instead of failing a single replica mid-deployment. + + On failure the Mooncake state created for this attempt is torn down + *before* any SimpleStorage initialisation: a surviving half-initialised + controller would make the next ``tq.init`` poll forever (F10). + """ + + def _close_owned_attempt() -> None: + # Retain the handle if teardown fails. The constructor's outer + # cleanup boundary can then retry instead of losing the only + # reference to possibly live process-global TQ state. + close_tq_owner(init_result.owner) + self._tq_owner = None + + try: + failures = verify_cluster_attach(init_result.config) + except TqHandshakeIsolationError: + _close_owned_attempt() + raise RuntimeError( + "Mooncake attach handshake workers could not be confirmed stopped; global fallback is unsafe" + ) from None + except Exception as error: + # Defensive boundary: verify_cluster_attach normally converts + # driver failures into stable summaries. An unexpected exception + # still must not strand the owner created for this attempt. + _close_owned_attempt() + raise RuntimeError( + f"Mooncake attach handshake orchestration failed ({safe_exception_kind(error)})" + ) from None + if not failures: + logger.info("[dataplane] Mooncake attach handshake passed on all alive nodes.") + return init_result + + detail = "; ".join(failures) + mode = getattr(self.config, "tq_rdma_mode", "off") + if mode != "auto" or fallback_config is None: + # ``required`` must fail loudly after cleaning this job's owner. + _close_owned_attempt() + raise RuntimeError( + f"Mooncake attach handshake reported {len(failures)} failure(s) (--tq-rdma-mode={mode}): {detail}" + ) + + logger.warning( + f"[dataplane] Mooncake attach handshake reported {len(failures)} failure(s) ({detail}); " + "closing Mooncake state and converging the whole job to SimpleStorage." + ) + # A cleanup failure leaves global TQ state unknown, so it must propagate + # instead of starting SimpleStorage on top of a possibly dirty cluster. + _close_owned_attempt() + fallback_result = initialize_tq_with_fallback(fallback_config, mode="auto") + self._tq_owner = fallback_result.owner + return TqInitResult( + config=fallback_result.config, + owner=fallback_result.owner, + fallback_reason=f"attach_handshake_failed:{len(failures)}_failures", + ) + + def _resolve_tq_backend(self, total_storage_size: int) -> dict: + """Resolve the TransferQueue ``backend`` config dict. + + This method only validates *configuration*: the requested mode, the + TransferQueue correctness contract, the master endpoint and the segment + capacity. It deliberately performs no hardware capability probe -- a + ``/sys`` scan cannot see which nodes the scheduler will actually use, + and host-RDMA capability is established afterwards by the real attach + handshake in :meth:`_confirm_mooncake_attach`. + + ``off`` retains the previous SimpleStorage and ownership semantics. For + ``auto``/``required``, any unmet precondition either falls back to + SimpleStorage (``auto``) or fails fast (``required``). + """ + # 1. Validate the requested mode. A malformed mode is a configuration + # error, never a reason to silently run on SimpleStorage. + errors = validate_config(self.config) + if errors: + raise ValueError("Invalid TransferQueue RDMA configuration:\n " + "\n ".join(errors)) + + mode = getattr(self.config, "tq_rdma_mode", "off") + + def _simple_storage() -> dict: + 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, + ) + + def _fall_back_or_raise(reason: str, error: Exception) -> dict: + """Every unmet host-RDMA precondition funnels through here. + + ``auto`` converges the whole job on SimpleStorage; ``required`` re- + raises so an operator who demanded RDMA never runs silently + downgraded. + """ + if mode != "auto": + raise RuntimeError(f"--tq-rdma-mode={mode} but {reason}: {error}") from None + logger.warning(f"[dataplane] {reason}; auto fallback to SimpleStorage: {error}") + return _simple_storage() + + # 2. SimpleStorage short-circuit (default, zero behavior change). + if mode == "off": + return _simple_storage() + + # 3. Configuration preconditions for the host-RDMA path. + device = getattr(self.config, "tq_rdma_device", "") + try: + validate_mooncake_runtime_contract() + except RuntimeError as e: + return _fall_back_or_raise( + "the installed TransferQueue does not satisfy the Mooncake correctness contract", e + ) + try: + master_address = resolve_mooncake_master_address() + except RuntimeError as e: + # A missing/invalid MC_MASTER_ADDRESS makes MooncakeStore + # unreachable for the whole job, so it degrades like any other + # unmet precondition rather than aborting an ``auto`` run. + return _fall_back_or_raise("the Mooncake master endpoint is not configured", e) + + # 4. Build the backend dict. Capacity validation reports an + # insufficient segment via ``cap_error``, but it can also *raise* for + # an unusable capacity input (garbage + # ``RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB``, missing ``seq_length``); both + # are configuration failures and must degrade identically. + try: + backend_dict, cap_error = build_backend_config( + self.config, + device=device, + master_address=master_address, + total_storage_size=total_storage_size, + ) + except RuntimeError as e: + return _fall_back_or_raise("the segment-capacity configuration is unusable", e) + if cap_error: + return _fall_back_or_raise("segment capacity insufficient", RuntimeError(cap_error)) + + # 5. One startup line stating what was requested and what will run. The + # effective transport is only confirmed once the cluster-wide attach + # handshake passes. + device_selection = "explicit" if device else "auto" + logger.info(f"[dataplane] requested: rdma_mode={mode} device_selection={device_selection}") + logger.info( + f"[dataplane] backend=MooncakeStore protocol=rdma device_selection={device_selection} (pending handshake)" + ) + return backend_dict + + def _close_data_system(self) -> None: + """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 +1051,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 +1256,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/actor_group.py b/relax/distributed/ray/actor_group.py index 44ba269b9..485d5e98e 100644 --- a/relax/distributed/ray/actor_group.py +++ b/relax/distributed/ray/actor_group.py @@ -121,6 +121,101 @@ def async_init(self, args, role, with_ref=False, with_opd_teacher=False): for actor in self._actor_handlers ] + def init_and_wait( + self, + args: Any, + role: str, + with_ref: bool = False, + with_opd_teacher: bool = False, + ) -> list[Any]: + """Initialize every train actor, destroying the group on any failure. + + ``MegatronTrainRayActor.init`` attaches the process-global + TransferQueue client from a regular actor method. If its bounded + ``tq.init`` times out, the method raises while the native daemon thread + may still be running. A failed distributed group cannot be reused + safely, so all actor processes are force-killed and confirmed terminal + before the initialization error is propagated. + """ + try: + # Submission itself can fail after one or more actor calls were + # already queued. Keep it inside the same cleanup boundary as + # asynchronous task failures so no partially initialized group is + # left reusable. + refs = self.async_init(args, role, with_ref=with_ref, with_opd_teacher=with_opd_teacher) + pending = list(refs) + results: dict[Any, Any] = {} + # ``ray.get(refs)`` may wait for every ref before surfacing one + # failure. A rank blocked in native initialization would then + # prevent cleanup forever, so consume whichever ref finishes first. + while pending: + ready, pending = ray.wait(pending, num_returns=1) + if not ready: + raise RuntimeError("Ray returned no completed train-actor initialization task") + ref = ready[0] + results[ref] = ray.get(ref) + return [results[ref] for ref in refs] + except Exception: + self._terminate_failed_init() + raise + + def _terminate_failed_init(self, timeout: float = 10.0) -> None: + """Force-kill a failed train group and confirm its actor tasks + ended.""" + probe_refs = [] + control_errors: list[str] = [] + + # Queue a task that can never return normally before issuing the force + # kill. Each ref must later fail with RayActorError, proving that Ray + # processed the actor death instead of merely accepting a kill request. + for actor in self._actor_handlers: + try: + probe_refs.append(actor.termination_probe.remote()) + except ray.exceptions.RayActorError: + # The actor was already dead before cleanup reached it. + continue + except Exception as error: + control_errors.append(f"termination probe submission ({type(error).__name__})") + + for actor in self._actor_handlers: + try: + ray.kill(actor, no_restart=True) + except ray.exceptions.RayActorError: + continue + except Exception as error: + control_errors.append(f"actor kill ({type(error).__name__})") + + pending = [] + ready = [] + if probe_refs: + try: + ready, pending = ray.wait( + probe_refs, + num_returns=len(probe_refs), + timeout=timeout, + ) + except Exception as error: + control_errors.append(f"termination wait ({type(error).__name__})") + + for ref in ready: + try: + ray.get(ref) + except ray.exceptions.RayActorError: + continue + except Exception as error: + control_errors.append(f"termination probe result ({type(error).__name__})") + else: + control_errors.append("termination probe returned normally") + + if control_errors or pending: + detail = ", ".join(control_errors) + if pending: + detail = f"{detail}, " if detail else "" + detail += f"{len(pending)} task(s) remained pending" + raise RuntimeError(f"Failed to confirm train actor cleanup after initialization error: {detail}") from None + + self._actor_handlers = [] + def async_train(self, rollout_id): """Do one rollout training.""" return [actor.train.remote(rollout_id) for actor in self._actor_handlers] diff --git a/relax/distributed/ray/rollout.py b/relax/distributed/ray/rollout.py index db8255b29..c065d295b 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,10 @@ 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, + role="rollout_worker", + ) logger.info(f"import {self.args.rollout_function_path} as generate_rollout function.") logger.info(f"import {self.args.eval_function_path} as eval_generate_rollout function.") @@ -917,6 +919,11 @@ 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. + if getattr(self, "data_system_client", None) is not None: + detach_tq_client() + self.data_system_client = None def _shutdown_all_engines(self, timeout: float = 15.0): """Shut down all SGLang engine actors and their child processes. diff --git a/relax/distributed/ray/train_actor.py b/relax/distributed/ray/train_actor.py index 9b497a93f..59b97f53f 100644 --- a/relax/distributed/ray/train_actor.py +++ b/relax/distributed/ray/train_actor.py @@ -3,6 +3,7 @@ import abc import os import random +import threading from datetime import timedelta import ray @@ -79,6 +80,15 @@ def init(self, args, role, with_ref=False, with_opd_teacher=False): numa_local_rank = Envs.RANK % args.num_gpus_per_node device_utils.set_numa_affinity(numa_local_rank) + def termination_probe(self) -> None: + """Never return normally; used to confirm force-killed actor exit. + + The driver queues this behind ``init`` before calling ``ray.kill``. Its + ObjectRef must become terminal with ``RayActorError``; a successful + return would mean the actor survived cleanup and is unsafe to reuse. + """ + threading.Event().wait() + def clear_memory(self): print_memory("before TrainRayActor.clear_memory") clear_memory() diff --git a/relax/utils/arguments.py b/relax/utils/arguments.py index f7d1e710e..36892e9f8 100644 --- a/relax/utils/arguments.py +++ b/relax/utils/arguments.py @@ -37,7 +37,7 @@ _MIN_TQ_VERSION = "0.1.10.dev0" _TQ_UPGRADE_CMD = ( 'pip install "transferqueue @ git+https://github.com/redai-infra/' - 'TransferQueue.git@58054a33834aadbcf76aacd6b1e32e25c030f2c9" --no-deps' + 'TransferQueue.git@6c7a587292910af0827f027de99e005e1900310e" --no-deps' ) @@ -244,6 +244,34 @@ def add_transfer_queue_arguments(parser): default=1, help="Fully async pipeline num of iters every global batch.", ) + # ── RDMA transport (MooncakeStore backend) ────────────────────────── + # The default selects the existing SimpleStorage data path. These + # expose *intent* only; Mooncake internals (endpoint, + # buffer, segment, timeout, master strategy) are handled via internal + # defaults / deployment environment, not CLI flags. + parser.add_argument( + "--tq-rdma-mode", + choices=["off", "auto", "required"], + default="off", + help=( + "TransferQueue data-plane transport. 'off' (default) uses " + "SimpleStorage/ZMQ with bounded worker attach and cleanup. " + "'auto' attempts MooncakeStore over host RDMA and falls back " + "to SimpleStorage when it is unavailable (with a WARNING). " + "'required' fails fast instead of falling back." + ), + ) + 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." + ), + ) return parser def add_cluster_arguments(parser): diff --git a/relax/utils/tq/__init__.py b/relax/utils/tq/__init__.py new file mode 100644 index 000000000..656f34359 --- /dev/null +++ b/relax/utils/tq/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""TransferQueue integration utilities. + +Import concrete submodules directly so importing :mod:`relax.utils.tq` stays +lightweight and does not eagerly load Ray, TransferQueue, or Mooncake. +""" diff --git a/relax/utils/tq/config.py b/relax/utils/tq/config.py new file mode 100644 index 000000000..6d93bdd70 --- /dev/null +++ b/relax/utils/tq/config.py @@ -0,0 +1,364 @@ +# 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-rdma-mode``, ``--tq-rdma-device``) into the OmegaConf dict that +``tq.init`` expects, and the only place that validates the configuration +preconditions for MooncakeStore. Actual host-RDMA capability is established by +the real attach handshake in :mod:`relax.utils.tq.lifecycle`, not here. + +Mooncake internals (endpoint, buffer size, segment size, master address, +timeout) are intentionally *not* exposed as CLI flags — they come from +internal defaults or the deployment environment, per maintainer guidance. +""" + +from __future__ import annotations + +import math +import os +from typing import Any + +from relax.utils.logging_utils import get_logger +from relax.utils.tq.correctness import ensure_mooncake_correctness_guards + + +logger = get_logger(__name__) + +# Accepted ``--tq-rdma-mode`` values; ``off`` keeps the SimpleStorage path. +TQ_RDMA_MODES = frozenset({"off", "auto", "required"}) + +# --------------------------------------------------------------------------- +# Defaults (kept here rather than in config.yaml so they are visible to +# Relax contributors without reading the TQ package). +# --------------------------------------------------------------------------- + +_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_METADATA_SERVER = "P2PHANDSHAKE" # config.yaml:42-43 + + +def validate_config(args: Any) -> list[str]: + """Return error messages for an invalid RDMA configuration. + + Called at startup before anything is initialised. An empty list means the + requested mode is structurally valid; whether host RDMA actually works is + decided later by the real attach handshake. ``argparse`` already + constrains the mode for CLI runs, so this also covers configs restored from + a checkpoint or built programmatically. + """ + errors: list[str] = [] + mode = getattr(args, "tq_rdma_mode", "off") + device = getattr(args, "tq_rdma_device", "") + + if not isinstance(mode, str) or mode not in TQ_RDMA_MODES: + errors.append(f"--tq-rdma-mode must be one of {', '.join(sorted(TQ_RDMA_MODES))}.") + if not isinstance(device, str): + errors.append("--tq-rdma-device must be a string.") + elif device and any(character.isspace() or not character.isprintable() for character in device): + errors.append("--tq-rdma-device must be empty or a single printable device name without whitespace.") + return errors + + +def _split_host_port(address: str) -> tuple[str, int]: + """Parse ``host:port`` (and bracketed IPv6) or raise ``ValueError``. + + Format-only: no DNS resolution and no connection attempt, so this stays a + configuration check rather than becoming another capability probe. + + Failure messages name the *kind* of defect and never echo ``address``: the + caller logs them, and a deployment endpoint is internal infrastructure + detail that must not leak into job logs. + """ + value = address.strip() + if value.startswith("["): + end = value.find("]") + if end < 0 or end + 2 > len(value) or value[end + 1] != ":": + raise ValueError("bracketed endpoint must be [host]:port") + host, port_text = value[1:end], value[end + 2 :] + else: + host, separator, port_text = value.rpartition(":") + if not separator: + raise ValueError("endpoint must be host:port") + if ":" in host: + # A bare IPv6 literal: rpartition would silently treat its last + # group as the port (``fe80::1`` -> host ``fe80:``, port 1). + raise ValueError("IPv6 endpoint must be bracketed as [host]:port") + if not host: + raise ValueError("endpoint has an empty host") + if any(character.isspace() or not character.isprintable() for character in host): + raise ValueError("endpoint host must not contain whitespace or control characters") + if not port_text.isdigit(): + raise ValueError("endpoint port must be a decimal number") + port = int(port_text) + if not 1 <= port <= 65535: + raise ValueError("endpoint port is outside 1-65535") + return host, port + + +def resolve_mooncake_master_address() -> str: + """Return the externally managed Mooncake master endpoint. + + ``MC_MASTER_ADDRESS`` is required and must parse as ``host:port``. A + loopback default would make every node of a multi-node job treat its own + localhost as the master, so ``auto`` would degrade and ``required`` would + abort even when a shared master is healthy elsewhere. + """ + address = os.environ.get("MC_MASTER_ADDRESS", "").strip() + if not address: + raise RuntimeError( + "MooncakeStore requires MC_MASTER_ADDRESS= of the externally " + "managed mooncake master in the driver environment; Relax never assumes " + "a loopback endpoint." + ) + try: + _split_host_port(address) + except ValueError as error: + raise RuntimeError(f"MC_MASTER_ADDRESS is not a usable endpoint: {error}") from error + return address + + +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: + raise RuntimeError("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB must be a finite positive number of GiB") from None + if not math.isfinite(gib) or gib <= 0: + raise RuntimeError("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB must be a finite positive number of GiB") + size_bytes = int(gib * 1024**3) + if size_bytes <= 0: + raise RuntimeError("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB must resolve to at least one byte") + return size_bytes + + +def validate_mooncake_runtime_contract() -> None: + """Validate the Relax-side portion of the Mooncake safety contract. + + The environment guard and versioned correctness-contract check run before + every Mooncake client is created or attached. This function does not + monkey-patch TransferQueue. + """ + ensure_mooncake_correctness_guards() + + +# --------------------------------------------------------------------------- +# 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( + *, + master_address: str, + device: str = "", + protocol: str = "rdma", + global_segment_size: int | None = None, +) -> dict[str, Any]: + """Build the ``backend`` dict for MooncakeStore. + + Parameters + ---------- + master_address + External master server address. It is validated here as well as by + :func:`resolve_mooncake_master_address` so direct benchmark callers + cannot bypass the format contract. + device + Explicit RDMA device name; empty lets Mooncake select one natively. + protocol + Transport. Production only ever builds ``rdma``; ``tcp`` exists solely + for the C1 benchmark baseline. + global_segment_size + Override the per-client segment size. ``None`` resolves from + ``RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB`` (default 4 GiB, see + :func:`resolve_global_segment_size`). Benchmarks may pass a larger + value (e.g. 8 GiB) to avoid staging-buffer pressure. + """ + if not isinstance(master_address, str): + raise ValueError("master_address must be a host:port string") + try: + _split_host_port(master_address) + except ValueError as error: + raise ValueError(f"master_address is not a usable endpoint: {error}") from None + master_address = master_address.strip() + + if global_segment_size is None: + segment_size = resolve_global_segment_size() + elif isinstance(global_segment_size, bool) or not isinstance(global_segment_size, int) or global_segment_size <= 0: + raise ValueError("global_segment_size must be a positive integer number of bytes") + else: + segment_size = global_segment_size + + cfg: dict[str, Any] = { + # Selects the manager inside ``tq.init`` (interface.py reads + # ``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": protocol, + "device_name": device, + # Master / metadata — externally managed, never auto-init. + "auto_init": False, + "master_server_address": master_address, + "metadata_server": _DEFAULT_METADATA_SERVER, + "local_hostname": "", # empty = auto-detect via Ray node IP + # Memory + "global_segment_size": segment_size, + "local_buffer_size": _DEFAULT_LOCAL_BUFFER_SIZE, + # Do NOT silently evict produced-but-unconsumed data. + "hard_pin": True, + # Host RDMA only in this phase: GDR is pinned off here rather than + # left to the upstream default so the payload always traverses host + # memory. A GDR path needs its own probe/verification and lands in + # a separate change (mooncake_client.py:75-88 for the client side). + "use_gdr": False, + }, + } + return cfg + + +# --------------------------------------------------------------------------- +# Capacity validation +# --------------------------------------------------------------------------- + + +# Worst-case payload factors used by the segment-capacity pre-check. +# +# The largest image/video tensor layout among the currently supported +# multimodal processors is Qwen3-VL's 16x16 spatial patch, temporal patch 2, +# RGB input, spatial merge 2x2, stored as float32. One schedulable vision +# token can therefore retain four flattened patch rows: +# +# 16 * 16 * 2 * 3 * 4 rows/token * 4 bytes/value = 24,576 bytes/token +# +# Keep this as one explicit bound instead of loading model config during +# Controller startup. If a supported processor gains a wider transported +# feature row, this bound and its regression test must be updated together. +_MULTIMODAL_BYTES_PER_TOKEN = 16 * 16 * 2 * 3 * (2 * 2) * 4 +# Text: token ids, logprobs, masks and rewards; 32 B/token rounds them up. +_TEXT_BYTES_PER_TOKEN = 32 + + +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 multimodal tokens than ``--seq-length`` allows, + so the transported tensor payload of one sample is bounded by + ``seq_length * _MULTIMODAL_BYTES_PER_TOKEN`` (192 MiB at + ``seq_length=8192``) rather than a fixed per-sample guess that can pass + configurations which later fail puts mid-training. + """ + n_samples = args.n_samples_per_prompt + capacity_batch = resolve_tq_capacity_batch_size(args) + 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 * _MULTIMODAL_BYTES_PER_TOKEN + return capacity_batch * n_samples * per_sample + + +def validate_segment_capacity(args: Any) -> str | None: + """Return an error message if segment capacity is insufficient, else None. + + Only meaningful for MooncakeStore (SimpleStorage manages its own capacity + via ``total_storage_size``), so callers invoke it on the Mooncake path + only. The check is conservative: it compares the *per-client* segment size + (``global_segment_size``) against the in-flight upper bound. + """ + 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, + *, + device: str, + master_address: str, + total_storage_size: int, +) -> tuple[dict[str, Any], str | None]: + """Return ``(backend_config_dict, error_or_none)`` for the host-RDMA path. + + ``master_address`` must already be validated by + :func:`resolve_mooncake_master_address`; it is threaded through so the + checked value is the one the client receives. On a capacity error the + returned dict is a safe SimpleStorage fallback and ``error`` explains why + MooncakeStore was rejected -- the caller decides whether that is fatal. + """ + cap_error = validate_segment_capacity(args) + 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(master_address=master_address, device=device), None diff --git a/relax/utils/tq/correctness.py b/relax/utils/tq/correctness.py new file mode 100644 index 000000000..3a37bb316 --- /dev/null +++ b/relax/utils/tq/correctness.py @@ -0,0 +1,199 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Mooncake safety guards and raw-byte payload correctness helpers. + +Relax validates the capabilities and environment it can inspect without +modifying TransferQueue at runtime. The pinned TransferQueue advertises a +versioned correctness contract for batch/retry result counts, remove failure +propagation, and fail-closed production-status ACKs. + +Mooncake 0.3.10.post2 was observed corrupting TCP-protocol transfers through +its auto-enabled memcpy fast path, so that path is force-disabled here and an +explicit enable is rejected (see :func:`_enforce_safe_memcpy`). +""" + +from __future__ import annotations + +import hashlib +import os +import struct +from collections.abc import Mapping +from typing import Any + + +LeafDigest = tuple[str, str, str] +_REQUIRED_TQ_MOONCAKE_CONTRACT_VERSION = 1 + + +def _enforce_safe_memcpy() -> None: + """Force-disable mooncake's memcpy fast path; reject attempts to enable it. + + mooncake 0.3.10.post2 auto-enables ``MC_STORE_MEMCPY`` in TCP-only + environments (``transfer_task.cpp`` "auto-detected: TCP-only environment, + memcpy enabled") and that path silently truncates cross-node gets: roughly + half of fresh-session first transfers returned rows whose tails were zero + bytes from a 64 KiB-aligned offset onward while every batch code reported + success (two-node forensic probes, 2026-08; 12/12 sessions clean with + ``MC_STORE_MEMCPY=0`` vs ~50% corrupt without). The same code path + SIGSEGVs on single-node loopback. Because the corruption is confirmed on + that build, this guard fails closed: an explicit ``MC_STORE_MEMCPY=1`` is + rejected at startup instead of honoured. Re-gate this behavior only when a + reliable fixed-build capability is available. + """ + override = os.environ.get("MC_STORE_MEMCPY", "").strip() + if override not in ("", "0"): + raise RuntimeError( + "MC_STORE_MEMCPY explicitly enables an unsafe value: the mooncake 0.3.10.post2 " + "memcpy fast path silently truncates TCP transfers and can SIGSEGV. " + "Unset MC_STORE_MEMCPY; Relax forces it to 0 until a fixed-build capability is available." + ) + os.environ["MC_STORE_MEMCPY"] = "0" + + +def ensure_mooncake_correctness_guards() -> None: + """Validate that the installed stack can run MooncakeStore safely. + + Enforces the memcpy environment contract and checks that the pinned + TransferQueue advertises the Mooncake correctness contract Relax's data + plane relies on. It does not modify TransferQueue at runtime. + """ + _enforce_safe_memcpy() + try: + import transfer_queue as tq + except ImportError as error: + raise RuntimeError( + "Installed TransferQueue does not satisfy the required Mooncake correctness contract" + ) from error + + actual = getattr(tq, "MOONCAKE_CORRECTNESS_CONTRACT_VERSION", 0) + if isinstance(actual, bool) or not isinstance(actual, int) or actual < _REQUIRED_TQ_MOONCAKE_CONTRACT_VERSION: + raise RuntimeError("Installed TransferQueue does not satisfy the required Mooncake correctness contract") + + +def _tensor_digest(value: Any) -> LeafDigest: + import torch + + flat = value.detach().cpu().contiguous().reshape(-1) + raw = flat.view(torch.uint8).numpy().tobytes() if flat.numel() else b"" + shape = "x".join(str(dim) for dim in value.shape) + return str(value.dtype), shape, hashlib.sha256(raw).hexdigest() + + +def _ndarray_digest(value: Any) -> LeafDigest: + import numpy as np + + contiguous = np.ascontiguousarray(value) + shape = "x".join(str(dim) for dim in value.shape) + return f"np.{contiguous.dtype}", shape, hashlib.sha256(contiguous.tobytes()).hexdigest() + + +def _scalar_digest(value: Any) -> LeafDigest: + if isinstance(value, bytes): + raw = value + elif isinstance(value, str): + raw = value.encode("utf-8") + elif isinstance(value, float): + raw = struct.pack("!d", value) + else: + raw = repr(value).encode("utf-8") + return f"py.{type(value).__name__}", "", hashlib.sha256(raw).hexdigest() + + +def _unwrap_non_tensor(value: Any) -> Any: + if type(value).__name__ == "NonTensorStack": + return value.tolist() + if type(value).__name__ == "NonTensorData": + return value.data + return value + + +def _dict_child_path(prefix: str, key: Any) -> str: + if not isinstance(key, str): + raise TypeError(f"Unsupported payload dict key at {prefix}: {type(key).__name__}") + return f"{prefix}.{key}" if key.isidentifier() else f"{prefix}[{key!r}]" + + +def leaf_digests(payload: Any, prefix: str = "payload") -> dict[str, LeafDigest]: + """Map every supported leaf to ``(dtype, shape, raw-byte SHA-256)``.""" + import numpy as np + import torch + + payload = _unwrap_non_tensor(payload) + digests: dict[str, LeafDigest] = {} + if isinstance(payload, torch.Tensor): + if payload.is_nested: + for index, row in enumerate(payload.unbind()): + digests[f"{prefix}[{index}]"] = _tensor_digest(row) + else: + digests[prefix] = _tensor_digest(payload) + elif isinstance(payload, np.ndarray): + digests[prefix] = _ndarray_digest(payload) + elif isinstance(payload, Mapping): + for key in payload: + if not isinstance(key, str): + raise TypeError(f"Unsupported payload dict key at {prefix}: {type(key).__name__}") + for key in sorted(payload): + digests.update(leaf_digests(payload[key], _dict_child_path(prefix, key))) + elif isinstance(payload, (list, tuple)): + for index, item in enumerate(payload): + digests.update(leaf_digests(item, f"{prefix}[{index}]")) + elif isinstance(payload, (str, bytes, bool, int, float)) or payload is None: + digests[prefix] = _scalar_digest(payload) + else: + raise TypeError(f"Unsupported payload leaf at {prefix}: {type(payload).__name__}") + return digests + + +def diff_digests(expected: dict[str, LeafDigest], actual: dict[str, LeafDigest]) -> list[str]: + """Return mismatch descriptions; an empty list means byte-exact.""" + problems: list[str] = [] + for path in sorted(expected.keys() | actual.keys()): + want, have = expected.get(path), actual.get(path) + if want is None: + problems.append(f"{path}: unexpected extra leaf {have}") + elif have is None: + problems.append(f"{path}: missing (expected {want})") + elif want != have: + for axis, want_part, have_part in zip(("dtype", "shape", "sha256"), want, have, strict=True): + if want_part != have_part: + problems.append(f"{path}: {axis} mismatch (expected {want_part}, got {have_part})") + return problems + + +def payload_rows(value: Any) -> list[Any]: + """Return logical rows from a tensor or non-tensor TQ column.""" + import torch + + value = _unwrap_non_tensor(value) + if isinstance(value, torch.Tensor): + return [value] if value.ndim == 0 else list(value.unbind()) + if isinstance(value, (list, tuple)): + return [_unwrap_non_tensor(row) for row in value] + raise TypeError(f"Unsupported payload column: {type(value).__name__}") + + +def payload_nbytes(value: Any) -> int: + """Count raw bytes represented by supported TQ payload values.""" + import numpy as np + import torch + + value = _unwrap_non_tensor(value) + if isinstance(value, torch.Tensor): + if value.is_nested: + return sum(payload_nbytes(row) for row in value.unbind()) + return value.numel() * value.element_size() + if isinstance(value, np.ndarray): + return value.nbytes + if isinstance(value, Mapping): + return sum(payload_nbytes(item) for item in value.values()) + if isinstance(value, (list, tuple)): + return sum(payload_nbytes(item) for item in value) + if isinstance(value, bytes): + return len(value) + if isinstance(value, str): + return len(value.encode("utf-8")) + if isinstance(value, float): + return 8 + if isinstance(value, (bool, int)) or value is None: + return len(repr(value).encode("utf-8")) + raise TypeError(f"Unsupported payload leaf: {type(value).__name__}") diff --git a/relax/utils/tq/lifecycle.py b/relax/utils/tq/lifecycle.py new file mode 100644 index 000000000..3198cd94d --- /dev/null +++ b/relax/utils/tq/lifecycle.py @@ -0,0 +1,711 @@ +# 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 math +import os +import threading +import time +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" +DEFAULT_TQ_INIT_TIMEOUT_SECONDS = 60.0 +DEFAULT_TQ_ATTACH_TIMEOUT_SECONDS = 60.0 + + +@dataclass(frozen=True) +class TqInitResult: + """Result of an exclusive-owner TransferQueue initialization.""" + + config: Any + owner: Any + fallback_reason: str = "" + + +class TqInitializationTimeout(TimeoutError): + """Raised when ``tq.init`` does not finish within the bounded timeout.""" + + +class TqInitializationError(RuntimeError): + """Raised after a failed ``tq.init`` candidate is fully cleaned.""" + + +class TqAttachTimeout(TimeoutError): + """Raised when a worker cannot attach to TransferQueue within the bound.""" + + +class TqCleanupTimeout(TimeoutError): + """Raised when an owner or controller cannot be confirmed gone.""" + + +class TqHandshakeIsolationError(RuntimeError): + """Raised when timed-out handshake workers cannot be confirmed stopped.""" + + +def safe_exception_kind(error: BaseException) -> str: + """Return a log-safe exception category without rendering its payload. + + Ray exception strings can contain worker IPs, process IDs and remote + traceback paths. Lifecycle logs and job-level failure summaries therefore + record only the stable exception class. + """ + ray_task_error_type = getattr(getattr(ray, "exceptions", None), "RayTaskError", None) + if isinstance(ray_task_error_type, type) and isinstance(error, ray_task_error_type): + return "RayTaskError" + name = type(error).__name__ + return name if name.isidentifier() else "Exception" + + +def _get_config_value(config: Any, key: str, default: Any = None) -> Any: + if config is None: + return default + if hasattr(config, "get"): + return config.get(key, default) + return getattr(config, key, default) + + +def uses_mooncake(conf: Any) -> bool: + """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. + """ + backend = _get_config_value(conf, "backend", {}) + return _get_config_value(backend, "storage_backend", "SimpleStorage") == "MooncakeStore" + + +def _prepare_mooncake_runtime(conf: Any) -> None: + if not uses_mooncake(conf): + return + from relax.utils.tq.config import validate_mooncake_runtime_contract + + validate_mooncake_runtime_contract() + + +def kill_tq_controller_and_wait(timeout: float = 20.0) -> None: + """Kill the TransferQueueController named actor, then wait for GCS + deregistration. + + 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 ({safe_exception_kind(e)})") from None + + deadline = time.time() + timeout + while time.time() < deadline: + try: + ray.get_actor(CONTROLLER_NAME, namespace=CONTROLLER_NAMESPACE) + except ValueError: + return + except ray.exceptions.RayError as error: + raise RuntimeError( + f"Failed to confirm TransferQueueController cleanup ({safe_exception_kind(error)})" + ) from None + time.sleep(0.4) + raise TqCleanupTimeout(f"TransferQueueController still resolvable after {timeout}s") + + +def reap_unusable_tq_controller(get_config_timeout: float = 10.0) -> bool: + """Require a clean exclusive cluster, reaping only unusable TQ state. + + Returns ``True`` when a half-initialised or unresponsive controller was + reaped. A healthy controller is never attached or killed: the initial RDMA + release supports one Relax job per Ray cluster, so healthy existing state + means the cluster is not clean and startup fails explicitly. + """ + try: + existing = ray.get_actor(CONTROLLER_NAME, namespace=CONTROLLER_NAMESPACE) + except ValueError: + return False # nothing there — nothing to reap. + except ray.exceptions.RayError as error: + raise RuntimeError(f"Failed to query TransferQueueController ({safe_exception_kind(error)})") from None + + try: + conf = ray.get(existing.get_config.remote(), timeout=get_config_timeout) + except (ray.exceptions.GetTimeoutError, ray.exceptions.RayActorError) as error: + # Only a proven timeout/dead actor is safe to classify as unusable. + # Other RayError subclasses can be transient GCS/control-plane failures; + # killing a healthy controller on those would violate the exclusive-job + # ownership boundary. + logger.warning( + f"[dataplane] Existing TransferQueueController is unusable ({safe_exception_kind(error)}); reaping it." + ) + conf = None + except ray.exceptions.RayError as error: + raise RuntimeError(f"Failed to inspect TransferQueueController ({safe_exception_kind(error)})") from None + + if conf is not None: + raise RuntimeError( + "A healthy TransferQueueController already exists. The initial RDMA release requires an exclusive, " + "clean Ray cluster; stop the previous Relax job or remove its TQ state before retrying." + ) + + logger.warning("[dataplane] TransferQueueController has no stored config (half-initialised); reaping it.") + kill_tq_controller_and_wait() + return True + + +def assert_mooncake_rdma_configured() -> None: + """Fail unless the attached client configuration requests Mooncake/RDMA. + + A successful :func:`attach_tq_client` only proves ``tq.init`` returned + without raising. This check also establishes that the stored controller + config produced a Mooncake manager and an RDMA-configured storage client: + ``tq.init`` ignores the caller's conf when attaching to an existing + controller (``interface.py:130-135``). ``storage_client.protocol`` is the + configured request, not a negotiated transport signal, so this is not + proof that bytes traversed an HCA or that a native transport fallback is + impossible. Wire-level proof remains the benchmark's counter check. + + Raises + ------ + RuntimeError + When the storage manager is not MooncakeStore, when no storage client + is present, or when the client's configured protocol is not ``rdma``. + """ + manager = tq.get_client().storage_manager + if type(manager).__name__ != "MooncakeStorageManager": + raise RuntimeError("attached storage manager is not MooncakeStorageManager") + + store_client = getattr(manager, "storage_client", None) + if store_client is None: + raise RuntimeError("MooncakeStorageManager exposes no storage_client after attach") + + protocol = getattr(store_client, "protocol", None) + if protocol != "rdma": + raise RuntimeError("MooncakeStore client is not configured for protocol=rdma") + + +def _get_stored_config(timeout: float = 10.0) -> Any: + try: + controller = ray.get_actor(CONTROLLER_NAME, namespace=CONTROLLER_NAMESPACE) + conf = ray.get(controller.get_config.remote(), timeout=timeout) + except ValueError: + raise + except ray.exceptions.RayError as error: + raise RuntimeError(f"Failed to read TransferQueueController config ({safe_exception_kind(error)})") from None + if conf is None: + raise RuntimeError("TransferQueueController returned no config after tq.init completed") + return conf + + +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: + raise RuntimeError("RELAX_TQ_ATTACH_TIMEOUT_SECONDS must be a finite positive number of seconds") from None + if not math.isfinite(value) or value <= 0: + raise RuntimeError("RELAX_TQ_ATTACH_TIMEOUT_SECONDS must be a finite positive number of seconds") + return value + + +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 ({safe_exception_kind(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, + *, + role: str, + timeout: float | None = None, +) -> Any: + """Attach a component process within a bounded deadline. + + The deadline covers both waiting for a served controller config and + ``tq.init`` itself, because either phase can hang unboundedly (get_config + polling and mooncake endpoint setup respectively). ``None`` resolves the + deadline from ``RELAX_TQ_ATTACH_TIMEOUT_SECONDS`` (default 60 s). + + The initial release assumes one component/replica per Ray actor process. + Components therefore own the one process-global TQ client they attach and + unconditionally detach it during teardown. + """ + if timeout is None: + timeout = _resolve_attach_timeout() + deadline = time.monotonic() + timeout + _prepare_mooncake_runtime(conf) + _await_controller_config(deadline) + _bounded_tq_init(conf, deadline, role=role) + return tq.get_client() + + +def detach_tq_client() -> 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). The initial release assumes one + component/replica per Ray actor process, so there is no cross-instance + generation lease. Force-killed workers still fall back to the master-side + TTL. + """ + client = None + store_client = None + try: + client = tq.get_client() + store_client = getattr(client.storage_manager, "storage_client", None) + except (AssertionError, AttributeError): + pass + + if store_client is not None and hasattr(store_client, "close"): + try: + store_client.close() + except Exception as e: # pragma: no cover - best-effort local cleanup + logger.warning(f"[dataplane] Failed to close attached MooncakeStore client ({safe_exception_kind(e)}).") + + if client is not None and hasattr(client, "close"): + try: + client.close() + except Exception as e: # pragma: no cover - best-effort local cleanup + logger.warning(f"[dataplane] Failed to close attached TransferQueue client ({safe_exception_kind(e)}).") + + # TransferQueue has no public detach-only API. Reset only process-local + # handles; never touch _TQ_STORAGE or the named controller actor. + try: + from transfer_queue import interface as tq_interface + + tq_interface._TQ_CLIENT = None + tq_interface._TQ_CONTROLLER = None + except (ImportError, AttributeError): # pragma: no cover - version dependent + pass + + +def _alive_node_ids() -> list[str]: + """Every alive node: TQ endpoints (Serve replicas and 0-CPU actors) carry + no placement binding, so any alive node may end up hosting one.""" + node_ids: list[str] = [] + for node in ray.nodes(): + if not node.get("Alive"): + continue + node_id = node.get("NodeID") + if not isinstance(node_id, str) or not node_id: + raise RuntimeError("an alive Ray node has no valid NodeID") + node_ids.append(node_id) + return node_ids + + +def _cancel_handshake_tasks(refs: list[Any], *, timeout: float = 10.0) -> bool: + """Force-cancel submitted one-shot workers and confirm every ref is ready. + + A timed-out worker may still own a daemon ``tq.init`` thread. The caller + may only tear down Mooncake and start SimpleStorage after Ray confirms the + force-cancelled task refs have reached a terminal state. + """ + cancellation_failed = False + for ref in refs: + try: + ray.cancel(ref, force=True) + except Exception as error: # pragma: no cover - Ray control-plane failure + cancellation_failed = True + logger.warning(f"[dataplane] Failed to force-cancel a TQ handshake worker ({safe_exception_kind(error)}).") + if cancellation_failed: + return False + try: + _ready, pending = ray.wait(refs, num_returns=len(refs), timeout=timeout) + except Exception as error: # pragma: no cover - Ray control-plane failure + logger.warning( + f"[dataplane] Could not confirm TQ handshake worker cancellation ({safe_exception_kind(error)})." + ) + return False + return not pending + + +def verify_cluster_attach(conf: Any, *, timeout: float | None = None) -> list[str]: + """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), confirms a + Mooncake manager whose client is configured for ``protocol=rdma``, and + detaches immediately. This validates actual attach/setup and configuration + agreement rather than a ``/sys`` heuristic; it is not negotiated-transport + or wire proof. The worker is not reused because a timed-out ``tq.init`` + daemon thread cannot be stopped safely in-process. An empty return means + every alive node attached within the deadline; the driver aggregates + failures and decides one job-level outcome. + """ + if timeout is None: + timeout = _resolve_attach_timeout() + + try: + node_ids = _alive_node_ids() + except Exception as error: + return [f"cluster: node discovery failed ({safe_exception_kind(error)})"] + if not node_ids: + return ["cluster: no alive Ray nodes discovered"] + + expects_mooncake = uses_mooncake(conf) + + refs: list[Any] = [] + node_number_by_ref: dict[Any, int] = {} + try: + from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy + + @ray.remote(num_cpus=0, max_retries=0, max_calls=1) + def _handshake(handshake_conf: Any, check_mooncake: bool, attach_timeout: float) -> None: + from relax.utils.tq.lifecycle import ( + assert_mooncake_rdma_configured, + attach_tq_client, + detach_tq_client, + ) + + attach_tq_client(handshake_conf, role="attach-handshake", timeout=attach_timeout) + try: + if check_mooncake: + assert_mooncake_rdma_configured() + finally: + # A failed assertion must not leak this node's registered + # segment; without the detach it lingers until client_ttl. + detach_tq_client() + + for node_number, node_id in enumerate(node_ids, start=1): + strategy = NodeAffinitySchedulingStrategy(node_id=node_id, soft=False) + ref = _handshake.options(scheduling_strategy=strategy).remote(conf, expects_mooncake, timeout) + refs.append(ref) + # Keep the raw Ray NodeID only in the scheduling strategy. Failure + # summaries may reach public CI/PR logs, so identify nodes by a + # stable per-handshake ordinal instead. + node_number_by_ref[ref] = node_number + except Exception as error: + if refs and not _cancel_handshake_tasks(refs): + raise TqHandshakeIsolationError( + "submitted TQ handshake workers could not be confirmed stopped after scheduling failed" + ) from None + return [f"cluster: handshake scheduling failed ({safe_exception_kind(error)})"] + + # Grace beyond the per-node attach deadline covers task scheduling and + # worker startup on a busy cluster. + wait_bound = timeout + 30.0 + try: + ready, pending = ray.wait(refs, num_returns=len(refs), timeout=wait_bound) + except Exception as error: + if not _cancel_handshake_tasks(refs): + raise TqHandshakeIsolationError( + "submitted TQ handshake workers could not be confirmed stopped after wait failed" + ) from None + return [f"cluster: handshake wait failed ({safe_exception_kind(error)})"] + failures: list[str] = [] + for ref in ready: + try: + ray.get(ref) + except Exception as error: + failures.append(f"node#{node_number_by_ref[ref]}: handshake task failed ({safe_exception_kind(error)})") + if pending and not _cancel_handshake_tasks(pending): + raise TqHandshakeIsolationError("timed-out TQ handshake workers could not be confirmed stopped") + for ref in pending: + failures.append(f"node#{node_number_by_ref[ref]}: handshake did not return within {wait_bound:.0f}s") + return failures + + +def close_tq_and_unmount() -> None: + """Close TransferQueue and unmount the MooncakeStore segment. + + Order matters: ``tq.close()`` still needs the store alive for its + ``remove_all()``, so the client handle is captured first and unmounted + after. SimpleStorage has no ``storage_client``, so this is a no-op there. + + This function is reachable only inside the dedicated owner actor. Workers + call :func:`detach_tq_client` instead because upstream ``tq.close()`` kills + the named controller even from an attached process. + """ + 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. + + try: + tq.close() + finally: + # ``tq.close`` may fail while removing queued data. Segment teardown + # is still mandatory; otherwise the failed owner leaves registered and + # locked memory behind until the Mooncake TTL expires. + if store_client is not None and hasattr(store_client, "close"): + try: + store_client.close() + logger.info("[dataplane] Unmounted MooncakeStore segment on teardown.") + except Exception as e: # pragma: no cover - best-effort cleanup + logger.warning(f"[dataplane] Failed to unmount MooncakeStore segment ({safe_exception_kind(e)}).") + + +@ray.remote(num_cpus=0) +class _TransferQueueOwner: + """Process boundary for first-time TQ initialization and global cleanup.""" + + def ready(self) -> None: + """Scheduling barrier used before ``tq.init`` enters fallback scope.""" + + def termination_probe(self) -> None: + """Never return; a terminal ref proves this actor process was + killed.""" + threading.Event().wait() + + def initialize(self, conf: Any) -> Any: + _prepare_mooncake_runtime(conf) + tq.init(conf=conf) + return _get_stored_config() + + def close(self) -> None: + close_tq_and_unmount() + + +def _stop_owner_actor(owner: Any, *, timeout: float = 10.0) -> None: + """Force-kill ``owner`` and prove its process reached a terminal state.""" + probe_ref = None + control_error = "" + try: + # Queue behind any blocked initialize/close call. It can only become + # ready after the actor dies because the method never returns. + probe_ref = owner.termination_probe.remote() + except ray.exceptions.RayActorError: + return # Already terminal before cleanup reached it. + except Exception as error: + control_error = f"termination probe submission ({safe_exception_kind(error)})" + + try: + ray.kill(owner, no_restart=True) + except ray.exceptions.RayActorError: + pass + except Exception as error: + kind = safe_exception_kind(error) + control_error = f"{control_error}, " if control_error else "" + control_error += f"actor kill ({kind})" + + pending: list[Any] = [] + if probe_ref is not None: + try: + ready, pending = ray.wait([probe_ref], num_returns=1, timeout=timeout) + except Exception as error: + ready = [] + control_error = f"{control_error}, " if control_error else "" + control_error += f"termination wait ({safe_exception_kind(error)})" + for ref in ready: + try: + ray.get(ref) + except ray.exceptions.RayActorError: + pass + except Exception as error: + control_error = f"{control_error}, " if control_error else "" + control_error += f"termination probe result ({safe_exception_kind(error)})" + else: + control_error = f"{control_error}, " if control_error else "" + control_error += "termination probe returned normally" + + if control_error or pending or probe_ref is None: + detail = control_error or "termination could not be observed" + if pending: + detail = f"{detail}, " if detail else "" + detail += "termination probe remained pending" + raise TqCleanupTimeout(f"Failed to confirm TransferQueue owner cleanup: {detail}") from None + + +def _cleanup_failed_owner(owner: Any, *, timeout: float = 10.0) -> None: + """Stop a failed initializer and remove its exclusive-cluster state.""" + try: + ray.get(owner.close.remote(), timeout=timeout) + except Exception as e: + logger.warning(f"[dataplane] TQ owner cleanup RPC failed; killing owner actor ({safe_exception_kind(e)}).") + _stop_owner_actor(owner, timeout=timeout) + + # Only remove global state after the actor is confirmed terminal. A + # pre-kill existence query races with a late controller creation by the + # blocked native initializer. The exclusive-cluster precondition proves + # any controller visible here belongs to this attempt. + kill_tq_controller_and_wait(timeout=timeout) + + +def _start_owner(*, timeout: float) -> Any: + """Create and schedule the isolated owner before candidate fallback.""" + try: + owner = _TransferQueueOwner.remote() + ray.get(owner.ready.remote(), timeout=timeout) + except Exception as error: + if "owner" in locals(): + _stop_owner_actor(owner, timeout=timeout) + raise RuntimeError(f"TransferQueue owner creation failed ({safe_exception_kind(error)})") from None + return owner + + +def _initialize_owner(owner: Any, conf: Any, *, timeout: float) -> TqInitResult: + """Initialize one candidate and fully clean its owner on failure.""" + try: + stored_conf = ray.get(owner.initialize.remote(conf), timeout=timeout) + except ray.exceptions.GetTimeoutError: + _cleanup_failed_owner(owner) + raise TqInitializationTimeout(f"tq.init did not finish within {timeout:.0f}s") from None + except Exception as error: + _cleanup_failed_owner(owner) + raise TqInitializationError( + f"TransferQueue owner initialization failed ({safe_exception_kind(error)})" + ) from None + return TqInitResult(config=stored_conf, owner=owner) + + +def close_tq_owner(owner: Any | None, *, timeout: float = 30.0) -> None: + """Ask the exclusive owner process to close global TQ state.""" + 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 + _stop_owner_actor(owner, timeout=timeout) + # Ensure the controller has left GCS even when owner.close() failed before + # reaching upstream tq.close(). This is a no-op when it is already gone. + kill_tq_controller_and_wait(timeout=timeout) + if close_error is not None: + raise RuntimeError(f"TransferQueue owner cleanup failed ({safe_exception_kind(close_error)})") from None + + +def initialize_tq_with_fallback( + conf: Any, + *, + mode: str, + fallback_conf: Any | None = None, + timeout: float = DEFAULT_TQ_INIT_TIMEOUT_SECONDS, +) -> TqInitResult: + """Initialize TQ with one exclusive owner and one safe auto fallback. + + ``fallback_conf`` must be the equivalent SimpleStorage configuration. It + is used only for ``mode='auto'`` after a real Mooncake ``tq.init`` failure. + ``required`` always cleans up and re-raises the original error. + """ + + def _attempt(owner: Any, attempt_conf: Any) -> TqInitResult: + return _initialize_owner(owner, attempt_conf, timeout=timeout) + + # Keep the exclusive-cluster gate outside the auto fallback boundary. A + # healthy pre-existing controller is not an RDMA candidate failure and + # must never be converted into a SimpleStorage attach attempt. + reap_unusable_tq_controller() + owner = _start_owner(timeout=timeout) + try: + return _attempt(owner, conf) + except (TqInitializationTimeout, TqInitializationError) as primary_error: + if mode != "auto" or fallback_conf is None: + raise + + reason = f"mooncake_init_failed:{type(primary_error).__name__}" + logger.warning( + f"[dataplane] Mooncake tq.init failed ({safe_exception_kind(primary_error)}); " + "cleaned partial state and retrying once with SimpleStorage." + ) + try: + # Failed-owner cleanup must leave a clean cluster. Re-check before + # starting the fallback so unknown residual state fails closed. + reap_unusable_tq_controller() + fallback_owner = _start_owner(timeout=timeout) + result = _attempt(fallback_owner, fallback_conf) + except Exception as fallback_error: + raise RuntimeError( + "TransferQueue SimpleStorage fallback initialization failed after " + f"Mooncake initialization error ({safe_exception_kind(primary_error)}); " + f"fallback error ({safe_exception_kind(fallback_error)})" + ) from None + return TqInitResult( + config=result.config, + owner=result.owner, + fallback_reason=reason, + ) diff --git a/scripts/benchmarks/tq_cross_node_bench.py b/scripts/benchmarks/tq_cross_node_bench.py new file mode 100644 index 000000000..4074ae7bd --- /dev/null +++ b/scripts/benchmarks/tq_cross_node_bench.py @@ -0,0 +1,691 @@ +#!/usr/bin/env python +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Cross-node TransferQueue C0/C1/C2 acceptance benchmark. + +Invoke once per fresh process: ``simple`` is C0, benchmark-only Mooncake/TCP is +C1, and Mooncake/host-RDMA is C2. The driver produces while one persistent +consumer is hard-pinned to another Ray node. Dtype/shape/raw-byte digests and +receive-counter wire proof are mandatory; each measured round is flushed to CSV +before either gate can fail. +""" + +from __future__ import annotations + +import argparse +import csv +import importlib.metadata as importlib_metadata +import json +import math +import os +import statistics +import subprocess +import time +from pathlib import Path +from typing import Any, TextIO + +import ray + +from relax.utils.tq.correctness import leaf_digests, payload_nbytes, payload_rows + + +PROTOCOL_LABELS = { + "simple": "C0 SimpleStorage", + "tcp": "C1 Mooncake/TCP", + "rdma": "C2 Mooncake/RDMA", +} +CSV_COLUMNS = ( + "protocol profile payload_mib actual_mib run measured status error_kind byte_exact mismatch_fields wire_proven " + "put_ms get_ms round_ms put_gbs get_gbs ib_mb tcp_mb proof_ib_mb proof_tcp_mb idle_ib_mbps idle_tcp_mbps " + "relax_sha tq_commit mooncake_version" +).split() +_TCP_WIRE_MIN_PAYLOAD_RATIO = 0.20 +_RDMA_WIRE_MIN_PAYLOAD_RATIO = 0.80 +_IDLE_COUNTER_SAMPLE_SECONDS = 0.5 +_GLOB_MAGIC = frozenset("*?[") + +CounterMap = dict[str, int] +DigestMap = dict[str, tuple[Any, ...]] + + +def _is_safe_device_name(value: str) -> bool: + return ( + bool(value) + and value not in {".", ".."} + and "/" not in value + and "\\" not in value + and not any(character in _GLOB_MAGIC for character in value) + and all(character.isprintable() and not character.isspace() for character in value) + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Cross-node TQ byte-exact and wire-proof benchmark") + parser.add_argument("--protocol", required=True, choices=sorted(PROTOCOL_LABELS)) + parser.add_argument("--consumer-node-id", required=True, help="Alive Ray NodeID distinct from the driver node") + parser.add_argument( + "--master", + default=None, + help="Mooncake master host:port (defaults to MC_MASTER_ADDRESS for Mooncake protocols)", + ) + parser.add_argument("--device", default="", help="RDMA device required by the C2 wire-proof counter") + parser.add_argument("--rdma-port", type=int, default=None, help="RDMA port required by the C2 wire-proof counter") + parser.add_argument("--tcp-device", required=True, help="Network interface used for the TCP receive counter") + parser.add_argument( + "--payload-profiles", nargs="+", default=["synthetic", "multimodal"], choices=["synthetic", "multimodal"] + ) + parser.add_argument("--payload-mib", nargs="+", type=int, default=[256, 1024, 2048, 4096]) + parser.add_argument("--num-samples", type=int, default=256) + parser.add_argument("--repeats", type=int, default=5) + parser.add_argument("--segment-gib", type=int, default=16) + parser.add_argument("--csv", required=True, type=Path, help="New per-round CSV output path") + args = parser.parse_args() + for name in ("num_samples", "repeats", "segment_gib"): + if getattr(args, name) <= 0: + parser.error(f"--{name.replace('_', '-')} must be positive") + if not args.payload_mib or any(size <= 0 for size in args.payload_mib): + parser.error("--payload-mib values must be positive") + if args.protocol == "simple": + if args.master is not None: + parser.error("--master is valid only with Mooncake protocols") + args.master = "" + else: + args.master = args.master or os.environ.get("MC_MASTER_ADDRESS", "") + if not args.master: + parser.error("--master or MC_MASTER_ADDRESS is required for Mooncake protocols") + if not _is_safe_device_name(args.tcp_device): + parser.error("--tcp-device must be a single printable interface name without whitespace or path separators") + if args.protocol == "rdma": + if not _is_safe_device_name(args.device): + parser.error("--device is required for RDMA and must be a single printable device name") + if args.rdma_port is None or args.rdma_port <= 0: + parser.error("--rdma-port is required for RDMA and must be positive") + elif args.device or args.rdma_port is not None: + parser.error("--device and --rdma-port are valid only with --protocol rdma") + return args + + +def _columns_for_budget(total_bytes: int, num_samples: int, dtype: Any) -> int: + import torch + + return max(1, total_bytes // (num_samples * torch.tensor([], dtype=dtype).element_size())) + + +def make_synthetic_payload(num_samples: int, total_mib: int) -> Any: + """Deterministic mixed-dtype payload with a fixed schema.""" + import torch + from tensordict import TensorDict + + generator = torch.Generator().manual_seed(20260824) + field_budget = total_mib * 1024**2 // 3 + fp32_cols = _columns_for_budget(field_budget, num_samples, torch.float32) + bf16_cols = _columns_for_budget(field_budget, num_samples, torch.bfloat16) + int64_cols = _columns_for_budget(field_budget, num_samples, torch.int64) + return TensorDict( + { + "activations": torch.randn(num_samples, fp32_cols, dtype=torch.float32, generator=generator), + "logits": torch.randn(num_samples, bf16_cols, dtype=torch.bfloat16, generator=generator), + "tokens": torch.randint(0, 151_000, (num_samples, int64_cols), dtype=torch.int64, generator=generator), + }, + batch_size=[num_samples], + ) + + +def make_multimodal_payload(num_samples: int, total_mib: int) -> Any: + """Deterministic production-shaped non-tensor vision-language payload.""" + import torch + + from relax.utils.utils import dict_to_tensordict + + target_bytes = total_mib * 1024**2 + seq_len = min(4096, max(128, target_bytes // max(1, num_samples * 128 * 1024))) + fixed_bytes = num_samples * (seq_len * 8 + 8 + 4) + pixel_budget = max(num_samples * 1536 * 4, target_bytes - fixed_bytes) + base_patches = max(1, pixel_budget // (num_samples * 1536 * 4)) + generator = torch.Generator().manual_seed(20260824) + multimodal_rows = [] + for index in range(num_samples): + variation = (index % 3) - 1 + requested_patches = max(1, base_patches + variation * max(1, base_patches // 20)) + height = max(1, int(requested_patches**0.5)) + width = max(1, (requested_patches + height - 1) // height) + patches = height * width + pixel_values = torch.empty((patches, 1536), dtype=torch.float32) + pixel_values.uniform_(-1.0, 1.0, generator=generator) + multimodal_rows.append( + { + "pixel_values": pixel_values, + "image_grid_thw": torch.tensor([[1, height, width]], dtype=torch.int64), + } + ) + token_row = list(range(seq_len)) + payload = dict_to_tensordict( + { + "tokens": [token_row for _ in range(num_samples)], + "sample_id": list(range(num_samples)), + "rewards": [float(index) / max(1, num_samples - 1) for index in range(num_samples)], + "multimodal_train_inputs": multimodal_rows, + }, + batch_size=num_samples, + ) + if type(payload.get("multimodal_train_inputs")).__name__ != "NonTensorStack": + raise RuntimeError("multimodal benchmark payload did not produce a NonTensorStack column") + return payload + + +def field_byte_digests(payload: Any, fields: list[str]) -> DigestMap: + """Ordered row digests preserving dtype, shape and every raw byte.""" + return {field: tuple(leaf_digests(row) for row in payload_rows(payload.get(field))) for field in fields} + + +def payload_bytes(payload: Any) -> int: + return sum(payload_nbytes(payload.get(field)) for field in payload.keys()) + + +def _read_required_counter(path: Path, label: str, *, scale: int = 1) -> int: + try: + value = int(path.read_text().strip()) * scale + except (OSError, ValueError): + raise RuntimeError(f"{label} receive counter is unavailable or unreadable") from None + if value < 0: + raise RuntimeError(f"{label} receive counter returned a negative value") + return value + + +def read_counters( + tcp_device: str, + rdma_device: str = "", + rdma_port: int | None = None, + sysfs_root: Path = Path("/sys/class"), +) -> CounterMap: + """Read exact receive-byte counters or fail closed. + + ``port_rcv_data`` is expressed in four-byte words. RDMA device and port + are an all-or-nothing pair so no glob or multi-port aggregation can enter + an acceptance result. + """ + if not _is_safe_device_name(tcp_device): + raise ValueError("TCP counter selection requires one safe interface name") + if bool(rdma_device) != (rdma_port is not None): + raise ValueError("RDMA counter selection requires both device and port") + if rdma_device and (not _is_safe_device_name(rdma_device) or rdma_port is None or rdma_port <= 0): + raise ValueError("RDMA counter selection requires one safe device and positive port") + + counters: CounterMap = {} + if rdma_device: + ib_path = sysfs_root / "infiniband" / rdma_device / "ports" / str(rdma_port) / "counters" / "port_rcv_data" + counters[f"ib:{rdma_device}:{rdma_port}"] = _read_required_counter(ib_path, "RDMA", scale=4) + tcp_path = sysfs_root / "net" / tcp_device / "statistics" / "rx_bytes" + counters[f"tcp:{tcp_device}"] = _read_required_counter(tcp_path, "TCP") + return counters + + +def _counter_delta(before: CounterMap, after: CounterMap, prefix: str) -> int: + before_keys = {key for key in before if key.startswith(prefix)} + after_keys = {key for key in after if key.startswith(prefix)} + if before_keys != after_keys: + raise RuntimeError(f"{prefix.rstrip(':').upper()} counter set changed during the measured round") + deltas = [after[key] - before[key] for key in sorted(before_keys)] + if any(delta < 0 for delta in deltas): + raise RuntimeError(f"{prefix.rstrip(':').upper()} counter reset or wrapped during the measured round") + return sum(deltas) + + +def _subtract_idle_noise(raw_delta: int, idle_delta: int, idle_seconds: float, round_seconds: float) -> int: + if idle_seconds <= 0 or round_seconds <= 0: + raise RuntimeError("counter sampling durations must be positive") + estimated_noise = math.ceil(idle_delta * round_seconds / idle_seconds) + return max(0, raw_delta - estimated_noise) + + +def wire_is_proven(protocol: str, ib_bytes: int, tcp_bytes: int, payload_nbytes: int) -> bool: + if protocol == "rdma": + return ib_bytes >= _RDMA_WIRE_MIN_PAYLOAD_RATIO * payload_nbytes + if protocol == "tcp": + return tcp_bytes >= _TCP_WIRE_MIN_PAYLOAD_RATIO * payload_nbytes + # SimpleStorage may place some units on the consumer node, so C0 cannot + # require a payload-volume ratio. It must still produce cross-node TCP. + return tcp_bytes > 0 + + +def build_conf(protocol: str, master: str, device: str, segment_gib: int) -> Any: + from omegaconf import OmegaConf + from transfer_queue import GRPOGroupNSampler + + from relax.utils.tq.config import ( + build_mooncake_config, + build_simple_storage_config, + validate_mooncake_runtime_contract, + ) + + if protocol == "simple": + backend = build_simple_storage_config(total_storage_size=None, num_data_storage_units=2) + else: + validate_mooncake_runtime_contract() + backend = build_mooncake_config( + master_address=master, + device=device, + protocol=protocol, + global_segment_size=segment_gib * 1024**3, + ) + return OmegaConf.create( + { + "controller": {"sampler": GRPOGroupNSampler(n_samples_per_prompt=1), "polling_mode": True}, + "backend": backend, + }, + flags={"allow_objects": True}, + ) + + +def _require_commit_id(value: Any, label: str) -> str: + commit = value if isinstance(value, str) else "" + if len(commit) != 40 or any(character not in "0123456789abcdefABCDEF" for character in commit): + raise RuntimeError(f"{label} commit is unavailable or not a full SHA") + return commit.lower() + + +def _git_output(repo_root: Path, *args: str) -> str: + try: + result = subprocess.run( + ["git", "-C", str(repo_root), *args], + capture_output=True, + text=True, + timeout=5, + check=False, + ) + except (OSError, subprocess.SubprocessError): + raise RuntimeError("Relax commit provenance is unavailable") from None + if result.returncode != 0: + raise RuntimeError("Relax commit provenance is unavailable") + return result.stdout.strip() + + +def collect_provenance(repo_root: Path | None = None) -> dict[str, str]: + """Collect exact source/dependency versions for every CSV row. + + Acceptance runs require a clean tracked Relax checkout. Untracked result + files are ignored so the benchmark can write artifacts below the checkout. + TransferQueue's PEP 610 metadata supplies the installed VCS commit. + """ + if repo_root is None: + repo_root = Path(__file__).resolve().parents[2] + if _git_output(repo_root, "status", "--porcelain", "--untracked-files=no"): + raise RuntimeError("benchmark acceptance requires a clean tracked Relax checkout") + revision = _git_output(repo_root, "rev-parse", "HEAD") + + distribution = importlib_metadata.distribution("transferqueue") + direct_url_text = distribution.read_text("direct_url.json") + if direct_url_text is None: + raise RuntimeError("installed TransferQueue has no VCS provenance metadata") + try: + direct_url = json.loads(direct_url_text) + tq_commit = direct_url["vcs_info"]["commit_id"] + except (KeyError, TypeError, json.JSONDecodeError): + raise RuntimeError("installed TransferQueue VCS provenance metadata is invalid") from None + try: + mooncake_version = importlib_metadata.version("mooncake-transfer-engine") + except importlib_metadata.PackageNotFoundError: + mooncake_version = "not-installed" + + return { + "relax_sha": _require_commit_id(revision, "Relax"), + "tq_commit": _require_commit_id(tq_commit, "TransferQueue"), + "mooncake_version": mooncake_version, + } + + +@ray.remote(num_cpus=0.001, max_restarts=0) +class TQConsumer: + """Persistent, hard-pinned consumer mirroring a Relax component actor.""" + + def __init__( + self, + conf: Any, + tcp_device: str, + rdma_device: str, + rdma_port: int | None, + protocol: str, + ) -> None: + selected_device = rdma_device if protocol == "rdma" else "" + selected_port = rdma_port if protocol == "rdma" else None + read_counters(tcp_device, selected_device, selected_port) + if protocol == "tcp": + os.environ["MC_TCP_ENABLE_CONNECTION_POOL"] = "1" + from relax.utils.tq.lifecycle import attach_tq_client + + # This runs inside the pinned Ray worker. The bounded helper applies + # Mooncake correctness guards in this process before client creation. + self.client = attach_tq_client(conf, role="benchmark-consumer") + self.tcp_device = tcp_device + self.rdma_device = selected_device + self.rdma_port = selected_port + + def describe(self) -> tuple[str, str]: + manager = self.client.storage_manager + storage_client = getattr(manager, "storage_client", None) + return type(manager).__name__, getattr(storage_client, "protocol", "") + + def counters(self) -> CounterMap: + return read_counters(self.tcp_device, self.rdma_device, self.rdma_port) + + def sample_idle_counters(self, seconds: float) -> dict[str, Any]: + if seconds <= 0: + raise ValueError("idle counter sample duration must be positive") + before = self.counters() + started = time.perf_counter() + time.sleep(seconds) + after = self.counters() + elapsed = time.perf_counter() - started + return { + "seconds": elapsed, + "ib_bytes": _counter_delta(before, after, "ib:"), + "tcp_bytes": _counter_delta(before, after, "tcp:"), + } + + def begin_round(self) -> tuple[CounterMap, float]: + return self.counters(), time.perf_counter() + + def fetch( + self, + fields: list[str], + batch_size: int, + partition: str, + expected: DigestMap, + before: CounterMap, + counter_started: float, + ) -> dict[str, Any]: + started = time.perf_counter() + meta = self.client.get_meta( + data_fields=fields, + batch_size=batch_size, + partition_id=partition, + mode="fetch", + task_name="tq-benchmark", + ) + received = self.client.get_data(meta) + get_ms = (time.perf_counter() - started) * 1000 + after = self.counters() + round_seconds = time.perf_counter() - counter_started + actual = field_byte_digests(received, fields) + mismatches = [field for field in fields if actual.get(field) != expected.get(field)] + return { + "get_ms": get_ms, + "ib_bytes": _counter_delta(before, after, "ib:"), + "tcp_bytes": _counter_delta(before, after, "tcp:"), + "round_seconds": round_seconds, + "byte_exact": not mismatches, + "mismatch_fields": mismatches, + } + + def shutdown(self) -> None: + from relax.utils.tq.lifecycle import detach_tq_client + + detach_tq_client() + + +def _validate_attached_backend(protocol: str, manager: str, attached_protocol: str) -> None: + if protocol == "simple" and manager != "AsyncSimpleStorageManager": + raise RuntimeError(f"C0 attached unexpected storage manager {manager}") + if protocol != "simple" and (manager != "MooncakeStorageManager" or attached_protocol != protocol): + raise RuntimeError( + f"{PROTOCOL_LABELS[protocol]} attached manager={manager} protocol={attached_protocol or '-'}" + ) + + +def _throughput_gbs(nbytes: int, milliseconds: float, label: str) -> float: + if milliseconds <= 0: + raise RuntimeError(f"{label} duration must be positive") + return nbytes / milliseconds / 1e6 + + +def _write_csv_record(writer: csv.DictWriter, csv_handle: TextIO, record: dict[str, Any]) -> None: + writer.writerow(record) + csv_handle.flush() + + +def _run_round( + *, + producer: Any, + consumer: Any, + payload: Any, + fields: list[str], + expected: DigestMap, + nbytes: int, + protocol: str, + profile: str, + requested_mib: int, + run: int, + writer: csv.DictWriter, + csv_handle: TextIO, + provenance: dict[str, str], +) -> dict[str, float]: + """Measure, persist and clean one warmup or acceptance round.""" + partition = f"bench-{profile}-{requested_mib}-{run}" + record: dict[str, Any] = {column: "" for column in CSV_COLUMNS} + record.update( + { + "protocol": protocol, + "profile": profile, + "payload_mib": requested_mib, + "actual_mib": round(nbytes / 1024**2, 2), + "run": run, + "measured": run > 0, + **provenance, + } + ) + try: + idle = ray.get(consumer.sample_idle_counters.remote(_IDLE_COUNTER_SAMPLE_SECONDS)) + before, counter_started = ray.get(consumer.begin_round.remote()) + started = time.perf_counter() + producer.put(payload, partition_id=partition) + put_ms = (time.perf_counter() - started) * 1000 + fetched = ray.get( + consumer.fetch.remote( + fields, + payload.batch_size[0], + partition, + expected, + before, + counter_started, + ) + ) + round_seconds = fetched["round_seconds"] + proof_ib_bytes = _subtract_idle_noise(fetched["ib_bytes"], idle["ib_bytes"], idle["seconds"], round_seconds) + proof_tcp_bytes = _subtract_idle_noise(fetched["tcp_bytes"], idle["tcp_bytes"], idle["seconds"], round_seconds) + proven = wire_is_proven(protocol, proof_ib_bytes, proof_tcp_bytes, nbytes) + put_gbs = _throughput_gbs(nbytes, put_ms, "put") + get_gbs = _throughput_gbs(nbytes, fetched["get_ms"], "get") + + if not fetched["byte_exact"]: + status, error_kind = "fail", "ByteExactMismatch" + gate_error: BaseException | None = AssertionError("byte-exact mismatch after TQ get") + elif not proven: + status, error_kind = "fail", "WireProofFailed" + gate_error = RuntimeError( + f"wire proof failed for protocol={protocol} profile={profile} payload={requested_mib}MiB run={run}" + ) + else: + status, error_kind = "pass", "" + gate_error = None + record.update( + { + "status": status, + "error_kind": error_kind, + "byte_exact": fetched["byte_exact"], + "mismatch_fields": ";".join(fetched["mismatch_fields"]), + "wire_proven": proven, + "put_ms": round(put_ms, 2), + "get_ms": round(fetched["get_ms"], 2), + "round_ms": round(round_seconds * 1000, 2), + "put_gbs": round(put_gbs, 3), + "get_gbs": round(get_gbs, 3), + "ib_mb": round(fetched["ib_bytes"] / 1e6, 1), + "tcp_mb": round(fetched["tcp_bytes"] / 1e6, 1), + "proof_ib_mb": round(proof_ib_bytes / 1e6, 1), + "proof_tcp_mb": round(proof_tcp_bytes / 1e6, 1), + "idle_ib_mbps": round(idle["ib_bytes"] / idle["seconds"] / 1e6, 3), + "idle_tcp_mbps": round(idle["tcp_bytes"] / idle["seconds"] / 1e6, 3), + } + ) + except BaseException as error: + record.update({"status": "error", "error_kind": type(error).__name__}) + try: + producer.clear_partition(partition) + except BaseException as cleanup_error: + _write_csv_record(writer, csv_handle, record) + raise error from cleanup_error + _write_csv_record(writer, csv_handle, record) + raise + + try: + producer.clear_partition(partition) + except BaseException as cleanup_error: + if gate_error is None: + record.update({"status": "error", "error_kind": type(cleanup_error).__name__}) + _write_csv_record(writer, csv_handle, record) + if gate_error is not None: + raise gate_error from cleanup_error + raise + _write_csv_record(writer, csv_handle, record) + if gate_error is not None: + raise gate_error + return {"put_gbs": put_gbs, "get_gbs": get_gbs} + + +def _teardown_benchmark(consumer: Any, producer_attached: bool, owner: Any) -> None: + """Attempt every cleanup step and propagate the first teardown error.""" + from relax.utils.tq.lifecycle import close_tq_owner, detach_tq_client + + first_error: BaseException | None = None + + def remember(error: BaseException) -> None: + nonlocal first_error + if first_error is None: + first_error = error + + if consumer is not None: + try: + ray.get(consumer.shutdown.remote(), timeout=10) + except BaseException as error: + remember(error) + try: + ray.kill(consumer, no_restart=True) + except BaseException as error: + remember(error) + if producer_attached: + try: + detach_tq_client() + except BaseException as error: + remember(error) + if owner is not None: + try: + close_tq_owner(owner) + except BaseException as error: + remember(error) + try: + ray.shutdown() + except BaseException as error: + remember(error) + if first_error is not None: + raise first_error + + +def main() -> None: + args = parse_args() + provenance = collect_provenance() + runtime_env: dict[str, Any] = {} + if args.protocol == "tcp": + os.environ["MC_TCP_ENABLE_CONNECTION_POOL"] = "1" + runtime_env = {"env_vars": {"MC_TCP_ENABLE_CONNECTION_POOL": "1"}} + from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy + + from relax.utils.tq.lifecycle import attach_tq_client, initialize_tq_with_fallback + + with args.csv.open("x", newline="") as csv_handle: + writer = csv.DictWriter(csv_handle, fieldnames=CSV_COLUMNS) + writer.writeheader() + csv_handle.flush() + ray.init(address="auto", logging_level="ERROR", runtime_env=runtime_env) + consumer = None + owner = None + producer_attached = False + measured_runs = 0 + try: + driver_node_id = ray.get_runtime_context().get_node_id() + alive_ids = {node["NodeID"] for node in ray.nodes() if node.get("Alive")} + if args.consumer_node_id not in alive_ids: + raise RuntimeError("--consumer-node-id is not an alive Ray node") + if args.consumer_node_id == driver_node_id: + raise RuntimeError("producer and consumer must run on different Ray nodes") + + conf = build_conf(args.protocol, args.master, args.device, args.segment_gib) + init_result = initialize_tq_with_fallback(conf, mode="required") + owner = init_result.owner + producer = attach_tq_client(init_result.config, role="benchmark-producer") + producer_attached = True + strategy = NodeAffinitySchedulingStrategy(node_id=args.consumer_node_id, soft=False) + consumer = TQConsumer.options(scheduling_strategy=strategy).remote( + init_result.config, + args.tcp_device, + args.device, + args.rdma_port, + args.protocol, + ) + manager, attached_protocol = ray.get(consumer.describe.remote()) + _validate_attached_backend(args.protocol, manager, attached_protocol) + print(f"[setup] {PROTOCOL_LABELS[args.protocol]} consumer=remote-node", flush=True) + + for profile in args.payload_profiles: + for requested_mib in args.payload_mib: + payload = ( + make_multimodal_payload(args.num_samples, requested_mib) + if profile == "multimodal" + else make_synthetic_payload(args.num_samples, requested_mib) + ) + fields = sorted(payload.keys()) + expected = field_byte_digests(payload, fields) + nbytes = payload_bytes(payload) + put_rates: list[float] = [] + get_rates: list[float] = [] + for run in range(args.repeats + 1): + rates = _run_round( + producer=producer, + consumer=consumer, + payload=payload, + fields=fields, + expected=expected, + nbytes=nbytes, + protocol=args.protocol, + profile=profile, + requested_mib=requested_mib, + run=run, + writer=writer, + csv_handle=csv_handle, + provenance=provenance, + ) + if run == 0: + continue + put_rates.append(rates["put_gbs"]) + get_rates.append(rates["get_gbs"]) + measured_runs += 1 + print( + f"[{profile} {requested_mib}MiB] byte_exact=PASS wire=PASS " + f"put={statistics.mean(put_rates):.2f}GB/s " + f"get_mean={statistics.mean(get_rates):.2f}GB/s " + f"get_median={statistics.median(get_rates):.2f}GB/s " + f"get_std={statistics.pstdev(get_rates):.2f}GB/s", + flush=True, + ) + print(f"[csv] wrote {measured_runs} measured rounds", flush=True) + except BaseException as primary_error: + try: + _teardown_benchmark(consumer, producer_attached, owner) + except BaseException as cleanup_error: + raise primary_error from cleanup_error + raise + else: + _teardown_benchmark(consumer, producer_attached, owner) + + +if __name__ == "__main__": + 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/core/test_controller_tq_backend.py b/tests/core/test_controller_tq_backend.py new file mode 100644 index 000000000..0b7d51ea7 --- /dev/null +++ b/tests/core/test_controller_tq_backend.py @@ -0,0 +1,277 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Controller decisions layered on top of the TQ config contracts.""" + +from __future__ import annotations + +import argparse +from contextlib import nullcontext +from types import SimpleNamespace +from typing import Any + +import pytest + +from tests.core.test_controller_s3_model_cleanup import controller +from tests.utils.test_arguments_opd_teacher_colocate import arguments_module as _arguments_module_fixture + + +arguments_module = _arguments_module_fixture +_MASTER = "master.invalid:50051" + + +def _config(**overrides: Any) -> SimpleNamespace: + values = { + "tq_rdma_mode": "auto", + "tq_rdma_device": "", + "num_data_storage_units": 1, + "n_samples_per_prompt": 1, + "rollout_batch_size": 8, + "seq_length": 8192, + "multimodal_keys": None, + "max_staleness": 0, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def _resolve(config: SimpleNamespace) -> dict[str, Any]: + instance = controller.Controller.__new__(controller.Controller) + instance.config = config + return instance._resolve_tq_backend(total_storage_size=64) + + +def _is_simple(backend: dict[str, Any]) -> bool: + return backend["storage_backend"] == "SimpleStorage" + + +class _DecisionHarness: + def __init__(self, monkeypatch: pytest.MonkeyPatch, failure: str | None) -> None: + self.calls: list[str] = [] + + def contract() -> None: + self.calls.append("contract") + if failure == "contract": + raise RuntimeError("contract unavailable") + + def master() -> str: + self.calls.append("master") + if failure == "master": + raise RuntimeError("master unavailable") + return _MASTER + + def backend(_args: Any, *, device: str, master_address: str, total_storage_size: int): + self.calls.append(f"build:{device}:{master_address}:{total_storage_size}") + if failure == "capacity-error": + return {"storage_backend": "SimpleStorage"}, "capacity insufficient" + if failure == "capacity-config": + raise RuntimeError("capacity configuration unusable") + return { + "storage_backend": "MooncakeStore", + "MooncakeStore": {"protocol": "rdma", "device_name": device, "master_server_address": master_address}, + }, None + + monkeypatch.setattr(controller, "validate_mooncake_runtime_contract", contract) + monkeypatch.setattr(controller, "resolve_mooncake_master_address", master) + monkeypatch.setattr(controller, "build_backend_config", backend) + + +@pytest.mark.parametrize( + ("mode", "failure", "expect_simple"), + [("off", None, True), ("auto", None, False), ("required", None, False)] + + [("auto", failure, True) for failure in ("contract", "master", "capacity-error", "capacity-config")] + + [("required", failure, False) for failure in ("contract", "master", "capacity-error", "capacity-config")], +) +def test_backend_decision_matrix( + monkeypatch: pytest.MonkeyPatch, mode: str, failure: str | None, expect_simple: bool +) -> None: + harness = _DecisionHarness(monkeypatch, failure) + context = pytest.raises(RuntimeError) if mode == "required" and failure else nullcontext() + with context: + backend = _resolve(_config(tq_rdma_mode=mode, tq_rdma_device="rdma0")) + assert _is_simple(backend) is expect_simple + if not expect_simple: + assert backend["MooncakeStore"] == { + "protocol": "rdma", + "device_name": "rdma0", + "master_server_address": _MASTER, + } + if mode == "off": + assert harness.calls == [] + elif failure == "contract": + assert harness.calls == ["contract"] + elif failure == "master": + assert harness.calls == ["contract", "master"] + else: + assert harness.calls == ["contract", "master", f"build:rdma0:{_MASTER}:64"] + + +@pytest.mark.parametrize("overrides", [{"tq_rdma_mode": "mooncake"}, {"tq_rdma_device": "bad\ndevice"}]) +def test_invalid_config_fails_before_preconditions(monkeypatch: pytest.MonkeyPatch, overrides: dict[str, Any]) -> None: + harness = _DecisionHarness(monkeypatch, None) + with pytest.raises(ValueError, match="TransferQueue RDMA configuration"): + _resolve(_config(**overrides)) + assert harness.calls == [] + + +def test_missing_mode_defaults_to_off(monkeypatch: pytest.MonkeyPatch) -> None: + harness = _DecisionHarness(monkeypatch, None) + config = _config() + del config.tq_rdma_mode + assert _is_simple(_resolve(config)) + assert harness.calls == [] + + +def test_cli_exposes_only_mode_and_device(arguments_module: Any) -> None: + arguments_module.RouterArgs = SimpleNamespace(add_cli_args=lambda parser, **_kwargs: parser) + parser = argparse.ArgumentParser() + arguments_module.get_slime_extra_args_provider()(parser) + tq_flags = sorted( + option for action in parser._actions for option in action.option_strings if option.startswith("--tq-") + ) + assert tq_flags == ["--tq-rdma-device", "--tq-rdma-mode"] + args = parser.parse_args([]) + assert (args.tq_rdma_mode, args.tq_rdma_device) == ("off", "") + + +def test_off_mode_rejects_a_healthy_existing_controller_before_legacy_init( + monkeypatch: pytest.MonkeyPatch, +) -> None: + instance = controller.Controller.__new__(controller.Controller) + instance.config = _config( + tq_rdma_mode="off", + fully_async=False, + balance_data=False, + polling_mode=False, + ) + instance._tq_owner = None + instance._tq_legacy_init = False + monkeypatch.setattr(controller, "resolve_sft_algo_key", lambda _config: "grpo") + monkeypatch.setattr(controller, "resolve_tq_capacity_batch_size", lambda _config: 1) + monkeypatch.setattr(controller, "GRPOGroupNSampler", lambda **_kwargs: object()) + monkeypatch.setattr(instance, "_resolve_tq_backend", lambda _size: {"storage_backend": "SimpleStorage"}) + monkeypatch.setattr( + controller, + "reap_unusable_tq_controller", + lambda: (_ for _ in ()).throw(RuntimeError("exclusive cluster is not clean")), + ) + monkeypatch.setattr(controller.tq, "init", lambda **_kwargs: pytest.fail("tq.init must not run")) + with pytest.raises(RuntimeError, match="exclusive cluster"): + instance._initialize_data_system() + assert instance._tq_owner is None and instance._tq_legacy_init is False + + +class _AttachHarness: + def __init__( + self, + monkeypatch: pytest.MonkeyPatch, + *, + failures: list[str] | None = None, + verify_error: BaseException | None = None, + close_error: BaseException | None = None, + ) -> None: + self.events: list[str] = [] + + def verify(_conf: Any, **_kwargs: Any) -> list[str]: + self.events.append("handshake") + if verify_error: + raise verify_error + return failures or [] + + def close(_owner: Any, **_kwargs: Any) -> None: + self.events.append("close") + if close_error: + raise close_error + + def fallback(conf: Any, **_kwargs: Any) -> Any: + self.events.append("init_simple") + return controller.TqInitResult(config=conf, owner="fallback-owner") + + monkeypatch.setattr(controller, "verify_cluster_attach", verify) + monkeypatch.setattr(controller, "close_tq_owner", close) + monkeypatch.setattr(controller, "initialize_tq_with_fallback", fallback) + + +def _confirm(mode: str = "auto") -> tuple[Any, Any]: + instance = controller.Controller.__new__(controller.Controller) + instance.config = _config(tq_rdma_mode=mode) + instance._tq_owner = "mooncake-owner" + result = controller.TqInitResult(config="mooncake-conf", owner="mooncake-owner") + return instance, result + + +@pytest.mark.parametrize( + ("mode", "failures", "verify_error", "expected_events", "error_match"), + [ + ("auto", [], None, ["handshake"], None), + ("auto", ["timeout"], None, ["handshake", "close", "init_simple"], None), + ("required", ["protocol=tcp"], None, ["handshake", "close"], "reported"), + ("auto", None, RuntimeError("private detail"), ["handshake", "close"], "orchestration failed"), + ( + "auto", + None, + controller.TqHandshakeIsolationError("private detail"), + ["handshake", "close"], + "could not be confirmed stopped", + ), + ], +) +def test_attach_handshake_cleanup_matrix( + monkeypatch: pytest.MonkeyPatch, + mode: str, + failures: list[str] | None, + verify_error: BaseException | None, + expected_events: list[str], + error_match: str | None, +) -> None: + harness = _AttachHarness(monkeypatch, failures=failures, verify_error=verify_error) + instance, result = _confirm(mode) + context = pytest.raises(RuntimeError, match=error_match) if error_match else nullcontext() + with context: + resolved = instance._confirm_mooncake_attach(result, "simple-conf") + if failures: + assert resolved.config == "simple-conf" and resolved.owner == "fallback-owner" + else: + assert resolved is result + assert harness.events == expected_events + + +def test_attach_cleanup_failure_aborts_before_fallback(monkeypatch: pytest.MonkeyPatch) -> None: + harness = _AttachHarness( + monkeypatch, + failures=["timeout"], + close_error=RuntimeError("owner cleanup failed"), + ) + instance, result = _confirm() + with pytest.raises(RuntimeError, match="owner cleanup failed"): + instance._confirm_mooncake_attach(result, "simple-conf") + assert harness.events == ["handshake", "close"] and instance._tq_owner == "mooncake-owner" + + +def test_constructor_and_repeated_cleanup_preserve_owner_boundary(monkeypatch: pytest.MonkeyPatch) -> None: + events: list[Any] = [] + close_data_system = controller.Controller._close_data_system + monkeypatch.setattr(controller, "resolve_sft_num_rollout", lambda _config: None) + monkeypatch.setattr(controller, "HealthManager", lambda **_kwargs: object()) + + def fail_init(instance: Any) -> None: + instance._tq_owner = "owner" + raise RuntimeError("initialization failed") + + def close(instance: Any) -> None: + events.append(instance._tq_owner) + instance._tq_owner = None + + monkeypatch.setattr(controller.Controller, "_initialize_data_system", fail_init) + monkeypatch.setattr(controller.Controller, "_close_data_system", close) + with pytest.raises(RuntimeError, match="initialization failed"): + controller.Controller(SimpleNamespace(use_health_check=False, max_global_restart=3)) + assert events == ["owner"] + + monkeypatch.setattr(controller.Controller, "_close_data_system", close_data_system) + instance = controller.Controller.__new__(controller.Controller) + instance._tq_legacy_init = False + instance._tq_owner = "owner" + monkeypatch.setattr(controller, "close_tq_owner", lambda owner: events.append(owner) if owner else None) + instance._close_data_system() + instance._close_data_system() + assert events == ["owner", "owner"] and instance._tq_owner is None diff --git a/tests/utils/_tq_handshake_timeout_probe.py b/tests/utils/_tq_handshake_timeout_probe.py new file mode 100644 index 000000000..e6ce61665 --- /dev/null +++ b/tests/utils/_tq_handshake_timeout_probe.py @@ -0,0 +1,120 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Subprocess proof that a timed-out one-shot attach worker cannot mutate +later.""" + +from __future__ import annotations + +import os +import sys +import time +from pathlib import Path + + +def _running(pid: int, created: float) -> bool: + import psutil + + try: + process = psutil.Process(pid) + return abs(process.create_time() - created) < 1e-3 and process.status() != psutil.STATUS_ZOMBIE + except psutil.NoSuchProcess: + return False + + +def _write_tq_stub(directory: Path) -> None: + directory.mkdir(parents=True) + (directory / "transfer_queue.py").write_text( + """\ +import os +import time +from pathlib import Path +import psutil + +MUTATED = False + +def init(conf=None): + process = psutil.Process() + Path(os.environ["RELAX_TEST_TQ_STARTED_MARKER"]).write_text( + 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: + 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 = probe_dir / "started" + mutated = probe_dir / "mutated" + delay = 2.0 + _write_tq_stub(stub_dir) + pythonpath = os.pathsep.join(filter(None, (str(stub_dir), os.environ.get("PYTHONPATH", "")))) + + import ray + + ray.init( + address="local", + num_cpus=1, + include_dashboard=False, + logging_level="ERROR", + runtime_env={ + "env_vars": { + "PYTHONPATH": pythonpath, + "RELAX_TQ_ATTACH_TIMEOUT_SECONDS": "0.3", + "RELAX_TEST_TQ_STARTED_MARKER": str(started), + "RELAX_TEST_TQ_LATE_MARKER": str(mutated), + "RELAX_TEST_TQ_LATE_MUTATION_DELAY": str(delay), + } + }, + _temp_dir=str(probe_dir / "ray"), + ) + try: + from relax.utils.tq import lifecycle + + conf = {"backend": {"storage_backend": "SimpleStorage"}, "controller": {}} + + @ray.remote(num_cpus=0) + class Controller: + def get_config(self) -> dict: + return conf + + controller = Controller.options( + name=lifecycle.CONTROLLER_NAME, namespace=lifecycle.CONTROLLER_NAMESPACE + ).remote() + assert ray.get(controller.get_config.remote()) == conf + failures = lifecycle.verify_cluster_attach(conf, timeout=0.3) + assert len(failures) == 1 and failures[0].endswith("handshake task failed (RayTaskError)") + pid_text, created_text = started.read_text(encoding="utf-8").split(",") + identity = int(pid_text), float(created_text) + + time.sleep(delay + 0.2) + assert not mutated.exists(), "timed-out tq.init mutated state after task failure" + deadline = time.monotonic() + 10 + while _running(*identity) and time.monotonic() < deadline: + time.sleep(0.05) + assert not _running(*identity), "one-shot handshake worker did not exit" + + @ray.remote(num_cpus=0, max_retries=0) + def clean_successor() -> tuple[int, float, bool]: + import psutil + import transfer_queue + + process = psutil.Process() + return process.pid, process.create_time(), transfer_queue.MUTATED + + successor_pid, successor_created, successor_mutated = ray.get(clean_successor.remote()) + assert (successor_pid, successor_created) != identity + assert successor_mutated is False and not mutated.exists() + finally: + ray.shutdown() + + +if __name__ == "__main__": + main(Path(sys.argv[1])) diff --git a/tests/utils/_train_actor_init_cleanup_probe.py b/tests/utils/_train_actor_init_cleanup_probe.py new file mode 100644 index 000000000..355f8c241 --- /dev/null +++ b/tests/utils/_train_actor_init_cleanup_probe.py @@ -0,0 +1,58 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Subprocess proof that failed train-actor init cannot mutate after +cleanup.""" + +from __future__ import annotations + +import os +import sys +import threading +import time +from pathlib import Path + + +def main(directory: Path) -> None: + os.environ.setdefault("RAY_ENABLE_UV_RUN_RUNTIME_ENV", "0") + os.environ.pop("RAY_ADDRESS", None) + import ray + + from relax.distributed.ray.actor_group import RayTrainGroup + + directory.mkdir(parents=True, exist_ok=True) + marker = directory / "late-mutation" + + @ray.remote(max_restarts=0) + class InitActor: + def __init__(self, fail: bool) -> None: + self.fail = fail + + def init(self, _args, _role, **_kwargs) -> str: + if not self.fail: + return "ready" + threading.Thread(target=lambda: (time.sleep(2), marker.write_text("dirty")), daemon=True).start() + raise RuntimeError("expected initialization failure") + + def termination_probe(self) -> None: + threading.Event().wait() + + ray.init( + address="local", num_cpus=2, include_dashboard=False, logging_level="ERROR", _temp_dir=str(directory / "ray") + ) + try: + group = object.__new__(RayTrainGroup) + group._actor_handlers = [InitActor.remote(True), InitActor.remote(False)] + with_error = False + try: + group.init_and_wait(object(), "actor") + except ray.exceptions.RayTaskError: + with_error = True + assert with_error and group._actor_handlers == [] + time.sleep(2.2) + assert not marker.exists() + finally: + ray.shutdown() + + +if __name__ == "__main__": + main(Path(sys.argv[1])) diff --git a/tests/utils/test_tq_benchmark_guards.py b/tests/utils/test_tq_benchmark_guards.py new file mode 100644 index 000000000..1a6e2c048 --- /dev/null +++ b/tests/utils/test_tq_benchmark_guards.py @@ -0,0 +1,253 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""CPU contracts for the retained cross-node acceptance benchmark.""" + +from __future__ import annotations + +import csv +import io +import sys +from contextlib import nullcontext +from types import SimpleNamespace +from typing import Any + +import pytest + +from scripts.benchmarks import tq_cross_node_bench as bench + + +def _argv(protocol: str, *extra: str) -> list[str]: + return [ + "bench", + "--protocol", + protocol, + "--consumer-node-id", + "node", + "--tcp-device", + "eth0", + "--csv", + "x", + *extra, + ] + + +def _raise(error: BaseException) -> None: + raise error + + +def test_simple_and_multimodal_profiles_use_the_expected_runtime_shapes(monkeypatch: pytest.MonkeyPatch) -> None: + import transfer_queue + + monkeypatch.setattr(transfer_queue, "GRPOGroupNSampler", lambda **_kwargs: object()) + assert bench.build_conf("simple", master="", device="", segment_gib=1).backend.storage_backend == "SimpleStorage" + payload = bench.make_multimodal_payload(num_samples=3, total_mib=1) + column = payload.get("multimodal_train_inputs") + assert type(column).__name__ == "NonTensorStack" + rows = column.tolist() + assert all(set(row) == {"pixel_values", "image_grid_thw"} for row in rows) + assert len({tuple(row["pixel_values"].shape) for row in rows}) > 1 + digest = bench.field_byte_digests(payload, ["multimodal_train_inputs"]) + rows[1]["pixel_values"][0, 0] += 1 + assert bench.field_byte_digests(payload, ["multimodal_train_inputs"]) != digest + + +@pytest.mark.parametrize( + ("protocol", "ib_bytes", "tcp_bytes", "payload_bytes", "expected"), + [ + ("rdma", 800, 0, 1000, True), + ("rdma", 799, 0, 1000, False), + ("rdma", 900, 901, 1000, True), + ("tcp", 0, 200, 1000, True), + ("tcp", 0, 199, 1000, False), + ("tcp", 201, 200, 1000, True), + ("simple", 0, 1, 1000, True), + ("simple", 1, 0, 1000, False), + ("rdma", 1, 0, 4 * 1024**3, False), + ("tcp", 0, 1, 4 * 1024**3, False), + ], +) +def test_wire_proof_matrix(protocol: str, ib_bytes: int, tcp_bytes: int, payload_bytes: int, expected: bool) -> None: + assert bench.wire_is_proven(protocol, ib_bytes, tcp_bytes, payload_bytes) is expected + + +def test_counter_scope_failure_modes_and_idle_subtraction(tmp_path) -> None: + for device, port, value in (("rdma0", 1, 11), ("rdma1", 1, 17), ("rdma1", 2, 29)): + directory = tmp_path / "infiniband" / device / "ports" / str(port) / "counters" + directory.mkdir(parents=True) + (directory / "port_rcv_data").write_text(str(value)) + tcp = tmp_path / "net" / "eth0" / "statistics" + tcp.mkdir(parents=True) + (tcp / "rx_bytes").write_text("101") + assert bench.read_counters("eth0", "rdma1", 2, sysfs_root=tmp_path) == { + "ib:rdma1:2": 29 * 4, + "tcp:eth0": 101, + } + with pytest.raises(ValueError, match="safe device"): + bench.read_counters("eth0", "*", 1, sysfs_root=tmp_path) + with pytest.raises(RuntimeError, match="TCP receive counter"): + bench.read_counters("missing", sysfs_root=tmp_path) + with pytest.raises(RuntimeError, match="counter set changed"): + bench._counter_delta({"tcp:x": 1}, {}, "tcp:") + with pytest.raises(RuntimeError, match="reset or wrapped"): + bench._counter_delta({"tcp:x": 2}, {"tcp:x": 1}, "tcp:") + assert bench._subtract_idle_noise(1_100, 100, 1.0, 1.0) == 1_000 + assert bench._subtract_idle_noise(50, 100, 1.0, 1.0) == 0 + with pytest.raises(RuntimeError, match="duration must be positive"): + bench._throughput_gbs(1_000, 0.0, "put") + + +@pytest.mark.parametrize( + "argv", + [ + _argv("rdma"), + _argv("tcp", "--master", "master.example:50051", "--device", "rdma0"), + *[ + _argv("rdma", "--master", "master.example:50051", "--device", device, "--rdma-port", "1") + for device in ("../rdma0", "*", "rdma*", "rdma?", "[ab]") + ], + _argv("tcp", "--master", "master.example:50051", "--rdma-port", "1"), + _argv("simple", "--master", "master.example:50051"), + ], +) +def test_cli_rejects_ambiguous_or_unsafe_counter_configuration( + monkeypatch: pytest.MonkeyPatch, argv: list[str] +) -> None: + monkeypatch.setattr(sys, "argv", argv) + with pytest.raises(SystemExit, match="2"): + bench.parse_args() + + +@pytest.mark.parametrize("cleanup_fails", [False, True], ids=["cleanup-ok", "cleanup-fails"]) +@pytest.mark.parametrize( + ("byte_exact", "tcp_bytes", "error_type", "match", "error_kind"), + [ + (False, 1_000, AssertionError, "byte-exact", "ByteExactMismatch"), + (True, 0, RuntimeError, "wire proof failed", "WireProofFailed"), + ], + ids=["byte-exact", "wire-proof"], +) +def test_failed_gate_is_recorded_and_cleanup_does_not_mask_it( + monkeypatch: pytest.MonkeyPatch, + cleanup_fails: bool, + byte_exact: bool, + tcp_bytes: int, + error_type: type[BaseException], + match: str, + error_kind: str, +) -> None: + events: list[str] = [] + monkeypatch.setattr(bench.ray, "get", lambda value: value) + consumer = SimpleNamespace( + sample_idle_counters=SimpleNamespace(remote=lambda _seconds: {"seconds": 1.0, "ib_bytes": 0, "tcp_bytes": 0}), + begin_round=SimpleNamespace(remote=lambda: ({"tcp:x": 0}, 1.0)), + fetch=SimpleNamespace( + remote=lambda *_args: { + "get_ms": 1.0, + "round_seconds": 1.0, + "ib_bytes": 0, + "tcp_bytes": tcp_bytes, + "byte_exact": byte_exact, + "mismatch_fields": [] if byte_exact else ["field"], + } + ), + ) + + class Producer: + def put(self, *_args: Any, **_kwargs: Any) -> None: + events.append("put") + + def clear_partition(self, _partition: str) -> None: + events.append("clear") + if cleanup_fails: + raise RuntimeError("partition cleanup failed") + + output = io.StringIO() + writer = csv.DictWriter(output, fieldnames=bench.CSV_COLUMNS) + writer.writeheader() + with pytest.raises(error_type, match=match) as excinfo: + bench._run_round( + producer=Producer(), + consumer=consumer, + payload=SimpleNamespace(batch_size=[1]), + fields=["field"], + expected={"field": ()}, + nbytes=1_000, + protocol="tcp", + profile="synthetic", + requested_mib=1, + run=1, + writer=writer, + csv_handle=output, + provenance={"relax_sha": "a" * 40, "tq_commit": "b" * 40, "mooncake_version": "test"}, + ) + row = next(csv.DictReader(io.StringIO(output.getvalue()))) + assert (row["status"], row["error_kind"], row["byte_exact"]) == ("fail", error_kind, str(byte_exact)) + assert row["mismatch_fields"] == ("" if byte_exact else "field") + assert isinstance(excinfo.value.__cause__, RuntimeError) is cleanup_fails + assert events == ["put", "clear"] + + +def test_provenance_records_versions_and_rejects_dirty_checkout(monkeypatch: pytest.MonkeyPatch, tmp_path) -> None: + results = iter([SimpleNamespace(returncode=0, stdout=""), SimpleNamespace(returncode=0, stdout="a" * 40 + "\n")]) + monkeypatch.setattr(bench.subprocess, "run", lambda *_args, **_kwargs: next(results)) + direct_url = '{"vcs_info":{"commit_id":"' + "b" * 40 + '"}}' + monkeypatch.setattr( + bench.importlib_metadata, + "distribution", + lambda _name: SimpleNamespace(read_text=lambda _filename: direct_url), + ) + monkeypatch.setattr(bench.importlib_metadata, "version", lambda _name: "0.3.test") + assert bench.collect_provenance(tmp_path) == { + "relax_sha": "a" * 40, + "tq_commit": "b" * 40, + "mooncake_version": "0.3.test", + } + monkeypatch.setattr( + bench.subprocess, + "run", + lambda *_args, **_kwargs: SimpleNamespace(returncode=0, stdout=" M tracked-file\n"), + ) + with pytest.raises(RuntimeError, match="clean tracked Relax checkout"): + bench.collect_provenance(tmp_path) + + +@pytest.fixture +def teardown_events(monkeypatch: pytest.MonkeyPatch) -> list[str]: + from relax.utils.tq import lifecycle + + events: list[str] = [] + monkeypatch.setattr(bench.ray, "kill", lambda *_args, **_kwargs: events.append("consumer-killed")) + monkeypatch.setattr(lifecycle, "detach_tq_client", lambda: events.append("producer-detached")) + monkeypatch.setattr(lifecycle, "close_tq_owner", lambda _owner: events.append("owner-closed")) + monkeypatch.setattr(bench.ray, "shutdown", lambda: events.append("ray-shutdown")) + return events + + +@pytest.mark.parametrize("failure", ["consumer", "owner", "none"]) +def test_teardown_is_ordered_and_does_not_touch_unowned_state( + monkeypatch: pytest.MonkeyPatch, teardown_events: list[str], failure: str +) -> None: + from relax.utils.tq import lifecycle + + consumer, attached, owner = None, False, None + expected = ["ray-shutdown"] + error_match = None + if failure == "consumer": + consumer, attached, owner = SimpleNamespace(shutdown=SimpleNamespace(remote=lambda: "ref")), True, "owner" + monkeypatch.setattr(bench.ray, "get", lambda *_args, **_kwargs: _raise(RuntimeError("consumer failed"))) + + def fail_owner_cleanup(_owner: Any) -> None: + teardown_events.append("owner-closed") + raise RuntimeError("owner failed") + + monkeypatch.setattr(lifecycle, "close_tq_owner", fail_owner_cleanup) + expected = ["consumer-killed", "producer-detached", "owner-closed", "ray-shutdown"] + error_match = "consumer failed" + elif failure == "owner": + owner = "owner" + monkeypatch.setattr(lifecycle, "close_tq_owner", lambda _owner: _raise(RuntimeError("owner failed"))) + error_match = "owner failed" + context = pytest.raises(RuntimeError, match=error_match) if error_match else nullcontext() + with context: + bench._teardown_benchmark(consumer, producer_attached=attached, owner=owner) + assert teardown_events == expected diff --git a/tests/utils/test_tq_dataplane_behavior.py b/tests/utils/test_tq_dataplane_behavior.py new file mode 100644 index 000000000..c2c20cffe --- /dev/null +++ b/tests/utils/test_tq_dataplane_behavior.py @@ -0,0 +1,200 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Real local SimpleStorage connection and data-plane contracts.""" + +from __future__ import annotations + +import importlib.util +import time +from typing import Any + +import pytest +import torch + +from relax.utils.tq.correctness import diff_digests, leaf_digests, payload_rows + + +def _has_real_submodule(dotted: str) -> bool: + try: + return importlib.util.find_spec(dotted) is not None + except (ImportError, ValueError, TypeError): + return False + + +pytestmark = pytest.mark.skipif( + not (_has_real_submodule("transfer_queue.storage") and importlib.util.find_spec("ray")), + reason="requires real TransferQueue, Ray, and a startable local CPU cluster", +) +_TQ_ACTOR = "TransferQueueController" +_TQ_NS = "transfer_queue" + + +def _wait_controller_gone(timeout: float = 20.0) -> bool: + 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 _payload(samples: int, fields: list[str], columns: int, seed: int = 0) -> Any: + from tensordict import TensorDict + + generator = torch.Generator().manual_seed(seed) + return TensorDict( + {field: torch.randn(samples, columns, generator=generator) for field in fields}, + batch_size=[samples], + ) + + +def _multimodal_payload(samples: int) -> dict[str, Any]: + grids = ((1, 58, 64), (1, 34, 64), (1, 64, 64), (1, 26, 40)) + generator = torch.Generator().manual_seed(20260813) + multimodal, tokens = [], [] + for index in range(samples): + temporal, height, width = grids[index % len(grids)] + multimodal.append( + { + "pixel_values": torch.randn(temporal * height * width, 1536, generator=generator), + "image_grid_thw": torch.tensor([[temporal, height, width]], dtype=torch.int64), + } + ) + tokens.append(torch.randint(0, 151_000, (512 + 173 * index,), generator=generator).tolist()) + return {"tokens": tokens, "multimodal_train_inputs": multimodal} + + +def _get(client: Any, partition: str, fields: list[str], size: int) -> Any: + meta = client.get_meta( + data_fields=fields, + batch_size=size, + partition_id=partition, + mode="fetch", + task_name=partition, + ) + return meta, client.get_data(meta) + + +@pytest.fixture(scope="module") +def _ray_cluster(): + import ray + + ray.init(ignore_reinit_error=True, logging_level="ERROR") + yield + ray.shutdown() + + +@pytest.fixture +def tq_factory(_ray_cluster): + import transfer_queue as tq + from omegaconf import OmegaConf + from transfer_queue import GRPOGroupNSampler + + def reinit(capacity: int = 1024): + tq.close() + if not _wait_controller_gone(): + _force_kill_controller() + assert _wait_controller_gone() + conf = OmegaConf.create( + { + "controller": {"sampler": GRPOGroupNSampler(n_samples_per_prompt=1), "polling_mode": True}, + "backend": {"SimpleStorage": {"total_storage_size": capacity, "num_data_storage_units": 1}}, + }, + flags={"allow_objects": True}, + ) + tq.init(conf=conf) + return tq.get_client() + + yield reinit + tq.close() + _wait_controller_gone() + _force_kill_controller() + _wait_controller_gone() + + +@pytest.mark.parametrize( + ("fields", "columns", "samples", "seed"), + [(["a", "b"], 8, 4, 0), (["img", "txt", "mask"], 16, 8, 42), (["pixel_values"], 1176, 4, 7)], + ids=["connection", "multi-field", "multimodal-width"], +) +def test_dense_round_trip_is_byte_exact(tq_factory, fields: list[str], columns: int, samples: int, seed: int) -> None: + client = tq_factory() + payload = _payload(samples, fields, columns, seed) + client.put(payload, partition_id="dense") + _, received = _get(client, "dense", fields, samples) + assert set(fields) <= set(received.keys()) + for field in fields: + expected_rows = [leaf_digests(row) for row in payload_rows(payload[field])] + actual_rows = [leaf_digests(row) for row in payload_rows(received[field])] + assert len(actual_rows) == len(expected_rows) + assert all( + not diff_digests(expected, actual) for expected, actual in zip(expected_rows, actual_rows, strict=True) + ) + + +def test_backpressure_fails_without_publishing_data(tq_factory) -> None: + client = tq_factory(capacity=4) + with pytest.raises(RuntimeError, match="capacity"): + client.put(_payload(8, ["a"], 4), partition_id="backpressure") + meta, _ = _get(client, "backpressure", ["a"], 8) + assert getattr(meta, "size", None) == 0 + + +def test_empty_get_returns_without_hanging(tq_factory) -> None: + meta, data = _get(tq_factory(), "empty", ["a"], 4) + assert getattr(meta, "size", None) == 0 and list(data.keys()) == [] + + +def test_repeat_put_stays_bounded_and_uncorrupted(tq_factory) -> None: + client = tq_factory() + first, second = _payload(4, ["a"], 4, 1), _payload(4, ["a"], 4, 2) + client.put(first, partition_id="repeat") + client.put(second, partition_id="repeat") + meta, received = _get(client, "repeat", ["a"], 4) + candidates = [leaf_digests(row) for row in [*payload_rows(first["a"]), *payload_rows(second["a"])]] + assert getattr(meta, "size", None) == 4 + assert all(leaf_digests(row) in candidates for row in payload_rows(received["a"])) + + +def test_clear_partition_and_reinit_are_isolated(tq_factory) -> None: + client = tq_factory() + client.put(_payload(4, ["a"], 4), partition_id="cleanup") + assert getattr(_get(client, "cleanup", ["a"], 4)[0], "size", None) == 4 + client.clear_partition("cleanup") + assert getattr(_get(client, "cleanup", ["a"], 4)[0], "size", None) == 0 + assert getattr(_get(tq_factory(capacity=16), "cleanup", ["a"], 4)[0], "size", None) == 0 + + +def test_multimodal_non_tensor_full_link_is_byte_exact(tq_factory, record_property) -> None: + from relax.utils.utils import dict_to_tensordict + + samples = 4 + source = _multimodal_payload(samples) + expected_mm = [leaf_digests(row) for row in source["multimodal_train_inputs"]] + expected_tokens = [leaf_digests(torch.tensor(row, dtype=torch.int64)) for row in source["tokens"]] + batch = dict_to_tensordict({**source, "sample_id": list(range(samples))}, batch_size=samples) + assert type(batch.get("multimodal_train_inputs")).__name__ == "NonTensorStack" + record_property("multimodal_payload_source", "synthetic") + + fields = ["sample_id", "tokens", "multimodal_train_inputs"] + client = tq_factory() + client.put(batch, partition_id="multimodal") + _, received = _get(client, "multimodal", fields, samples) + sample_ids = [int(value) for value in payload_rows(received["sample_id"])] + assert sorted(sample_ids) == list(range(samples)) + for position, sample_id in enumerate(sample_ids): + assert not diff_digests(expected_mm[sample_id], leaf_digests(payload_rows(received[fields[2]])[position])) + assert not diff_digests(expected_tokens[sample_id], leaf_digests(payload_rows(received[fields[1]])[position])) diff --git a/tests/utils/test_tq_failure_paths.py b/tests/utils/test_tq_failure_paths.py new file mode 100644 index 000000000..3ee84c1b1 --- /dev/null +++ b/tests/utils/test_tq_failure_paths.py @@ -0,0 +1,595 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""CI-safe timeout, retry, disconnect, fallback, and cleanup contracts.""" + +from __future__ import annotations + +import asyncio +import importlib.util +import os +import subprocess +import sys +import tempfile +import time +from contextlib import nullcontext +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest +import torch + +from relax.utils.tq import lifecycle as tq_lifecycle + + +def _has_real_submodule(dotted: str) -> bool: + try: + return importlib.util.find_spec(dotted) is not None + except (ImportError, ValueError, TypeError): + return False + + +_REAL_MOONCAKE_CLIENT = _has_real_submodule("transfer_queue.storage.clients.mooncake_client") + + +def _raise(error: BaseException) -> None: + raise error + + +@pytest.mark.parametrize( + ("state", "expected", "raises"), + [ + ("absent", False, None), + ("healthy", None, RuntimeError), + ("half", True, None), + ("timeout", True, None), + ("dead", True, None), + ("unknown-error", None, RuntimeError), + ], +) +def test_reaper_state_matrix( + monkeypatch: pytest.MonkeyPatch, state: str, expected: bool | None, raises: type[BaseException] | None +) -> None: + actor = MagicMock() + killed: list[str] = [] + fake_ray = MagicMock() + fake_ray.exceptions = tq_lifecycle.ray.exceptions + if state == "absent": + fake_ray.get_actor.side_effect = ValueError("missing") + else: + fake_ray.get_actor.return_value = actor + if state == "healthy": + fake_ray.get.return_value = {"backend": {}} + elif state == "half": + fake_ray.get.return_value = None + elif state == "timeout": + fake_ray.get.side_effect = fake_ray.exceptions.GetTimeoutError("private") + elif state == "dead": + fake_ray.get.side_effect = fake_ray.exceptions.RayActorError(error_msg="private") + elif state == "unknown-error": + fake_ray.get.side_effect = fake_ray.exceptions.RayError("private endpoint") + monkeypatch.setattr(tq_lifecycle, "ray", fake_ray) + monkeypatch.setattr(tq_lifecycle, "kill_tq_controller_and_wait", lambda **_kwargs: killed.append("killed")) + context = pytest.raises(raises) if raises else nullcontext() + with context: + assert tq_lifecycle.reap_unusable_tq_controller() is expected + assert killed == (["killed"] if state in {"half", "timeout", "dead"} else []) + + +def test_controller_cleanup_timeout_fails_closed(monkeypatch: pytest.MonkeyPatch) -> None: + controller = object() + killed: list[Any] = [] + monkeypatch.setattr(tq_lifecycle.ray, "get_actor", lambda *_args, **_kwargs: controller) + monkeypatch.setattr(tq_lifecycle.ray, "kill", lambda handle: killed.append(handle)) + with pytest.raises(tq_lifecycle.TqCleanupTimeout, match="still resolvable"): + tq_lifecycle.kill_tq_controller_and_wait(timeout=0) + assert killed == [controller] + + +@pytest.mark.parametrize( + ("has_store", "close_error", "expected"), + [(True, None, ["tq", "store"]), (True, RuntimeError("close failed"), ["tq", "store"]), (False, None, ["tq"])], +) +def test_close_unmount_matrix( + monkeypatch: pytest.MonkeyPatch, + has_store: bool, + close_error: BaseException | None, + expected: list[str], +) -> None: + events: list[str] = [] + store = MagicMock() + manager = SimpleNamespace(storage_client=store) if has_store else SimpleNamespace() + fake_tq = MagicMock() + fake_tq.get_client.return_value = MagicMock(storage_manager=manager) + + def close() -> None: + events.append("tq") + if close_error: + raise close_error + + fake_tq.close.side_effect = close + store.close.side_effect = lambda: events.append("store") + monkeypatch.setattr(tq_lifecycle, "tq", fake_tq) + context = pytest.raises(RuntimeError, match="close failed") if close_error else nullcontext() + with context: + tq_lifecycle.close_tq_and_unmount() + assert events == expected + + +@pytest.mark.parametrize("value", ["soon", "nan", "inf", "-inf"]) +def test_attach_timeout_rejects_unusable_env(monkeypatch: pytest.MonkeyPatch, value: str) -> None: + monkeypatch.setenv("RELAX_TQ_ATTACH_TIMEOUT_SECONDS", value) + with pytest.raises(RuntimeError, match="finite positive") as excinfo: + tq_lifecycle._resolve_attach_timeout() + assert value not in str(excinfo.value) + + +def test_bounded_init_timeout_error_and_env_override(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("RELAX_TQ_ATTACH_TIMEOUT_SECONDS", "12.5") + assert tq_lifecycle._resolve_attach_timeout() == 12.5 + monkeypatch.setattr(tq_lifecycle.tq, "init", lambda **_kwargs: time.sleep(5)) + with pytest.raises(tq_lifecycle.TqAttachTimeout, match="did not finish"): + tq_lifecycle._bounded_tq_init({}, time.monotonic() + 0.2, role="test") + monkeypatch.setattr(tq_lifecycle.tq, "init", lambda **_kwargs: _raise(ValueError("bad conf"))) + with pytest.raises(ValueError, match="bad conf"): + tq_lifecycle._bounded_tq_init({}, time.monotonic() + 5, role="test") + monkeypatch.setattr(tq_lifecycle.ray, "get_actor", lambda *_args, **_kwargs: _raise(ValueError("missing"))) + with pytest.raises(tq_lifecycle.TqAttachTimeout, match="attach timed out"): + tq_lifecycle._await_controller_config(time.monotonic() + 0.2) + + +class _HandshakeHarness: + def __init__(self, monkeypatch: pytest.MonkeyPatch, node_ids: list[str] | None = None) -> None: + self.args: list[tuple[Any, ...]] = [] + self.workers: list[Any] = [] + self.options: dict[str, Any] = {} + self.strategies: list[Any] = [] + + class Task: + def options(task_self, **options: Any) -> "Task": + self.strategies.append(options["scheduling_strategy"]) + return task_self + + def remote(task_self, *args: Any) -> object: + self.args.append(args) + return object() + + def remote(**options: Any): + self.options.update(options) + + def decorate(function: Any) -> Task: + self.workers.append(function) + return Task() + + return decorate + + monkeypatch.setattr(tq_lifecycle.ray, "remote", remote) + monkeypatch.setattr(tq_lifecycle, "_alive_node_ids", lambda: node_ids or ["a" * 56]) + monkeypatch.setattr(tq_lifecycle.ray, "wait", lambda refs, **_kwargs: (list(refs), [])) + monkeypatch.setattr(tq_lifecycle.ray, "get", lambda _ref: None) + + +def test_cluster_attach_uses_every_node_hard_affinity_and_one_shot_workers(monkeypatch: pytest.MonkeyPatch) -> None: + nodes = ["a" * 56, "b" * 56] + harness = _HandshakeHarness(monkeypatch, nodes) + assert tq_lifecycle.verify_cluster_attach({}, timeout=0.1) == [] + assert (harness.options["max_calls"], harness.options["max_retries"]) == (1, 0) + assert [strategy.node_id for strategy in harness.strategies] == nodes + assert all(strategy.soft is False for strategy in harness.strategies) + + +@pytest.mark.parametrize( + ("backend", "verify_mooncake", "assert_error"), + [ + ("MooncakeStore", True, None), + ("MooncakeStore", True, RuntimeError("protocol=tcp")), + ("SimpleStorage", False, None), + ], +) +def test_handshake_worker_verifies_backend_and_always_detaches( + monkeypatch: pytest.MonkeyPatch, + backend: str, + verify_mooncake: bool, + assert_error: BaseException | None, +) -> None: + harness = _HandshakeHarness(monkeypatch) + conf = {"backend": {"storage_backend": backend}} + assert tq_lifecycle.verify_cluster_attach(conf, timeout=0.1) == [] + assert harness.args[0][1] is verify_mooncake + events: list[str] = [] + monkeypatch.setattr(tq_lifecycle, "attach_tq_client", lambda *_args, **_kwargs: events.append("attach")) + monkeypatch.setattr( + tq_lifecycle, + "assert_mooncake_rdma_configured", + lambda: events.append("assert") or (_raise(assert_error) if assert_error else None), + ) + monkeypatch.setattr(tq_lifecycle, "detach_tq_client", lambda: events.append("detach")) + context = pytest.raises(RuntimeError, match="protocol=tcp") if assert_error else nullcontext() + with context: + harness.workers[0](conf, verify_mooncake, 0.1) + assert events == ["attach", *(["assert"] if verify_mooncake else []), "detach"] + + +def test_handshake_scheduling_failure_cancels_submitted_workers(monkeypatch: pytest.MonkeyPatch) -> None: + first_ref = object() + cancelled: list[object] = [] + + class Task: + calls = 0 + + def options(self, **_kwargs: Any) -> "Task": + return self + + def remote(self, *_args: Any) -> object: + self.calls += 1 + if self.calls == 1: + return first_ref + raise RuntimeError("private scheduling detail") + + monkeypatch.setattr(tq_lifecycle.ray, "remote", lambda **_kwargs: lambda _function: Task()) + monkeypatch.setattr(tq_lifecycle, "_alive_node_ids", lambda: ["a" * 56, "b" * 56]) + monkeypatch.setattr(tq_lifecycle.ray, "cancel", lambda ref, force: cancelled.append(ref)) + monkeypatch.setattr(tq_lifecycle.ray, "wait", lambda refs, **_kwargs: (list(refs), [])) + assert tq_lifecycle.verify_cluster_attach({}, timeout=0.1) == [ + "cluster: handshake scheduling failed (RuntimeError)" + ] + assert cancelled == [first_ref] + + +@pytest.mark.parametrize("cancel_fails", [False, True]) +def test_unconfirmed_handshake_worker_aborts_fallback(monkeypatch: pytest.MonkeyPatch, cancel_fails: bool) -> None: + _HandshakeHarness(monkeypatch) + monkeypatch.setattr(tq_lifecycle.ray, "wait", lambda refs, **_kwargs: ([], list(refs))) + if cancel_fails: + monkeypatch.setattr(tq_lifecycle.ray, "cancel", lambda *_args, **_kwargs: _raise(RuntimeError("private"))) + else: + monkeypatch.setattr(tq_lifecycle.ray, "cancel", lambda *_args, **_kwargs: None) + with pytest.raises(tq_lifecycle.TqHandshakeIsolationError, match="could not be confirmed stopped"): + tq_lifecycle.verify_cluster_attach({}, timeout=0.1) + + +class MooncakeStorageManager: + def __init__(self, storage_client: Any) -> None: + self.storage_client = storage_client + + +@pytest.mark.parametrize( + ("manager", "match"), + [ + (MooncakeStorageManager(MagicMock(protocol="rdma")), None), + (MagicMock(), "not MooncakeStorageManager"), + (MooncakeStorageManager(None), "no storage_client"), + *[(MooncakeStorageManager(MagicMock(protocol=value)), "not configured") for value in ("tcp", None, "")], + ], +) +def test_configured_mooncake_manager_matrix(monkeypatch: pytest.MonkeyPatch, manager: Any, match: str | None) -> None: + fake_tq = MagicMock() + fake_tq.get_client.return_value = MagicMock(storage_manager=manager) + monkeypatch.setattr(tq_lifecycle, "tq", fake_tq) + context = pytest.raises(RuntimeError, match=match) if match else nullcontext() + with context: + tq_lifecycle.assert_mooncake_rdma_configured() + + +def test_attach_timeout_leaves_no_reusable_process_global_state() -> None: + env = os.environ.copy() + env["RAY_ENABLE_UV_RUN_RUNTIME_ENV"] = "0" + env.pop("RAY_ADDRESS", None) + repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) + with tempfile.TemporaryDirectory(prefix="tq-ray-") as probe_dir: + result = subprocess.run( + [sys.executable, "-m", "tests.utils._tq_handshake_timeout_probe", probe_dir], + cwd=repo_root, + env=env, + capture_output=True, + text=True, + timeout=45, + check=False, + ) + assert result.returncode == 0, f"probe stdout:\n{result.stdout}\nprobe stderr:\n{result.stderr}" + + +def test_worker_detach_closes_clients_resets_handles_and_is_used_by_components( + monkeypatch: pytest.MonkeyPatch, +) -> None: + storage_client = MagicMock() + client = MagicMock(storage_manager=SimpleNamespace(storage_client=storage_client)) + fake_tq = MagicMock() + fake_tq.get_client.return_value = client + interface = SimpleNamespace(_TQ_CLIENT=client, _TQ_CONTROLLER=object()) + monkeypatch.setattr(tq_lifecycle.tq, "interface", interface, raising=False) + monkeypatch.setattr(tq_lifecycle, "tq", fake_tq) + tq_lifecycle.detach_tq_client() + storage_client.close.assert_called_once_with() + client.close.assert_called_once_with() + assert interface._TQ_CLIENT is None and interface._TQ_CONTROLLER is None + + from relax.components.base import Base + + calls: list[bool] = [] + monkeypatch.setattr(tq_lifecycle, "detach_tq_client", lambda: calls.append(True)) + component = Base() + component.data_system_client = object() + component.__del__() + assert calls == [True] and component.data_system_client is None + + +def _conf(backend: str) -> dict[str, Any]: + return {"controller": {}, "backend": {"storage_backend": backend}} + + +class _InitHarness: + def __init__( + self, + monkeypatch: pytest.MonkeyPatch, + init_effects: list[Any], + *, + start_effect: Any = None, + reap_effects: list[BaseException | None] | None = None, + ) -> None: + self.events: list[str] = [] + effects = iter(init_effects) + reaps = iter(reap_effects or []) + + def reap() -> None: + self.events.append("reap") + effect = next(reaps, None) + if effect: + raise effect + + def start(*, timeout: float) -> str: + self.events.append("start") + if isinstance(start_effect, BaseException): + raise start_effect + return f"owner-{self.events.count('start')}" + + def initialize(owner: Any, conf: dict[str, Any], *, timeout: float) -> Any: + backend = conf["backend"]["storage_backend"] + self.events.append(f"init:{backend}") + effect = next(effects) + if isinstance(effect, BaseException): + raise effect + return tq_lifecycle.TqInitResult(config=conf, owner=owner) + + monkeypatch.setattr(tq_lifecycle, "reap_unusable_tq_controller", reap) + monkeypatch.setattr(tq_lifecycle, "_start_owner", start) + monkeypatch.setattr(tq_lifecycle, "_initialize_owner", initialize) + + +@pytest.mark.parametrize( + ("mode", "init_effects", "start_effect", "reap_effects", "expected_events", "error"), + [ + ("off", [None], None, None, ["reap", "start", "init:SimpleStorage"], None), + ( + "auto", + [tq_lifecycle.TqInitializationError("failed"), None], + None, + None, + ["reap", "start", "init:MooncakeStore", "reap", "start", "init:SimpleStorage"], + None, + ), + ( + "required", + [tq_lifecycle.TqInitializationError("failed")], + None, + None, + ["reap", "start", "init:MooncakeStore"], + tq_lifecycle.TqInitializationError, + ), + ( + "auto", + [tq_lifecycle.TqCleanupTimeout("pending")], + None, + None, + ["reap", "start", "init:MooncakeStore"], + tq_lifecycle.TqCleanupTimeout, + ), + ( + "auto", + [tq_lifecycle.TqInitializationError("failed")], + None, + [None, RuntimeError("dirty")], + ["reap", "start", "init:MooncakeStore", "reap"], + RuntimeError, + ), + ("auto", [], RuntimeError("schedule failed"), None, ["reap", "start"], RuntimeError), + ], +) +def test_initialize_and_fallback_transaction_matrix( + monkeypatch: pytest.MonkeyPatch, + mode: str, + init_effects: list[Any], + start_effect: Any, + reap_effects: list[BaseException | None] | None, + expected_events: list[str], + error: type[BaseException] | None, +) -> None: + harness = _InitHarness( + monkeypatch, + init_effects, + start_effect=start_effect, + reap_effects=reap_effects, + ) + primary = _conf("SimpleStorage" if mode == "off" else "MooncakeStore") + context = pytest.raises(error) if error else nullcontext() + with context: + result = tq_lifecycle.initialize_tq_with_fallback(primary, mode=mode, fallback_conf=_conf("SimpleStorage")) + assert result.owner is not None + if mode == "auto" and len(init_effects) == 2: + assert result.config["backend"]["storage_backend"] == "SimpleStorage" + assert harness.events == expected_events + + +class _RemoteMethod: + def __init__(self, value: str) -> None: + self.value = value + + def remote(self, *_args: Any, **_kwargs: Any) -> str: + return self.value + + +class _FakeOwner: + def __init__(self) -> None: + self.ready = _RemoteMethod("ready") + self.initialize = _RemoteMethod("initialize") + self.close = _RemoteMethod("close") + self.termination_probe = _RemoteMethod("probe") + + +@pytest.mark.parametrize("outcome", ["success", "timeout"]) +def test_owner_start_boundary(monkeypatch: pytest.MonkeyPatch, outcome: str) -> None: + owner = _FakeOwner() + stopped: list[Any] = [] + monkeypatch.setattr(tq_lifecycle._TransferQueueOwner, "remote", lambda: owner) + if outcome == "timeout": + monkeypatch.setattr( + tq_lifecycle.ray, + "get", + lambda *_args, **_kwargs: _raise(tq_lifecycle.ray.exceptions.GetTimeoutError("timeout")), + ) + monkeypatch.setattr(tq_lifecycle, "_stop_owner_actor", lambda handle, **_kwargs: stopped.append(handle)) + with pytest.raises(RuntimeError, match="owner creation failed"): + tq_lifecycle._start_owner(timeout=0.1) + assert stopped == [owner] + else: + monkeypatch.setattr(tq_lifecycle.ray, "get", lambda *_args, **_kwargs: None) + assert tq_lifecycle._start_owner(timeout=0.1) is owner + + +def test_owner_initialize_success_timeout_and_exception_cleanup(monkeypatch: pytest.MonkeyPatch) -> None: + owner = _FakeOwner() + stored = _conf("SimpleStorage") + monkeypatch.setattr(tq_lifecycle.ray, "get", lambda *_args, **_kwargs: stored) + result = tq_lifecycle._initialize_owner(owner, stored, timeout=0.1) + assert result.config is stored and result.owner is owner + + cleaned: list[Any] = [] + monkeypatch.setattr(tq_lifecycle, "_cleanup_failed_owner", lambda handle: cleaned.append(handle)) + for source, expected in ( + (tq_lifecycle.ray.exceptions.GetTimeoutError("timeout"), tq_lifecycle.TqInitializationTimeout), + (RuntimeError("private detail"), tq_lifecycle.TqInitializationError), + ): + monkeypatch.setattr(tq_lifecycle.ray, "get", lambda *_args, error=source, **_kwargs: _raise(error)) + with pytest.raises(expected): + tq_lifecycle._initialize_owner(owner, {}, timeout=0.1) + assert cleaned == [owner, owner] + + +@pytest.mark.parametrize( + ("ready", "ray_get", "error"), + [ + (True, lambda _ref: _raise(tq_lifecycle.ray.exceptions.RayActorError(error_msg="dead")), None), + (False, lambda _ref: None, tq_lifecycle.TqCleanupTimeout), + (True, lambda _ref: None, tq_lifecycle.TqCleanupTimeout), + ], +) +def test_owner_stop_requires_terminal_actor_failure( + monkeypatch: pytest.MonkeyPatch, + ready: bool, + ray_get: Any, + error: type[BaseException] | None, +) -> None: + owner = _FakeOwner() + monkeypatch.setattr(tq_lifecycle.ray, "kill", lambda *_args, **_kwargs: None) + monkeypatch.setattr( + tq_lifecycle.ray, "wait", lambda refs, **_kwargs: (list(refs), []) if ready else ([], list(refs)) + ) + monkeypatch.setattr(tq_lifecycle.ray, "get", ray_get) + context = pytest.raises(error) if error else nullcontext() + with context: + tq_lifecycle._stop_owner_actor(owner, timeout=0.1) + + +def test_owner_cleanup_order_and_fail_closed_boundary(monkeypatch: pytest.MonkeyPatch) -> None: + owner = _FakeOwner() + events: list[str] = [] + monkeypatch.setattr(tq_lifecycle.ray, "get", lambda *_args, **_kwargs: events.append("close")) + monkeypatch.setattr(tq_lifecycle, "_stop_owner_actor", lambda *_args, **_kwargs: events.append("terminal")) + monkeypatch.setattr(tq_lifecycle, "kill_tq_controller_and_wait", lambda **_kwargs: events.append("reap")) + tq_lifecycle._cleanup_failed_owner(owner) + assert events == ["close", "terminal", "reap"] + + monkeypatch.setattr( + tq_lifecycle, + "_stop_owner_actor", + lambda *_args, **_kwargs: _raise(tq_lifecycle.TqCleanupTimeout("pending")), + ) + with pytest.raises(tq_lifecycle.TqCleanupTimeout, match="pending"): + tq_lifecycle._cleanup_failed_owner(owner) + assert events == ["close", "terminal", "reap", "close"] + + +class _FlakyStore: + def __init__(self, fail_times: int, error: Exception | None = None) -> None: + self.fail_times = fail_times + self.error = error + self.calls: list[list[str]] = [] + + def batch_get_into(self, keys: list[str], _ptrs: list[int], _sizes: list[int]) -> list[int]: + self.calls.append(list(keys)) + if self.error: + raise self.error + if self.fail_times: + self.fail_times -= 1 + return [-800] * len(keys) + return [0] * len(keys) + + +def _client_with_store(store: _FlakyStore) -> Any: + from transfer_queue.storage.clients.mooncake_client import MooncakeStoreClient + + client = object.__new__(MooncakeStoreClient) + client._store = store + client.replica_config = None + return client + + +@pytest.mark.skipif(not _REAL_MOONCAKE_CLIENT, reason="requires real TransferQueue Mooncake client") +@pytest.mark.parametrize("disconnect", [False, True], ids=["retry", "disconnect"]) +def test_retry_and_disconnect_surface(monkeypatch: pytest.MonkeyPatch, disconnect: bool) -> None: + monkeypatch.setattr("transfer_queue.storage.clients.mooncake_client.RETRY_DELAY_SECONDS", 0) + error = RuntimeError("Failed to open segment for endpoint='
'") if disconnect else None + store = _FlakyStore(0 if disconnect else 2, error) + client = _client_with_store(store) + context = pytest.raises(RuntimeError, match="Failed to open segment") if disconnect else nullcontext() + with context: + client._batch_get_into_with_retry(["0@f0", "1@f0"], [1, 2], [8, 8]) + assert len(store.calls) == (1 if disconnect else 3) + + +class _InlineExecutorLoop: + def run_in_executor(self, _executor: Any, function: Any, *args: Any) -> asyncio.Future[Any]: + future = asyncio.get_running_loop().create_future() + try: + future.set_result(function(*args)) + except BaseException as error: + future.set_exception(error) + return future + + +@pytest.mark.skipif(not _REAL_MOONCAKE_CLIENT, reason="requires real TransferQueue Mooncake manager") +@pytest.mark.asyncio +@pytest.mark.parametrize("put_fails", [False, True], ids=["success", "failure"]) +async def test_production_status_follows_storage_success(monkeypatch: pytest.MonkeyPatch, put_fails: bool) -> None: + from tensordict import TensorDict + from transfer_queue.storage.managers.mooncake_manager import MooncakeStorageManager + + monkeypatch.setattr(asyncio, "get_event_loop", lambda: _InlineExecutorLoop()) + calls: list[str] = [] + storage = MagicMock() + if put_fails: + storage.put.side_effect = RuntimeError("capacity exhausted") + else: + storage.put.side_effect = lambda *_args: calls.append("put") + manager = object.__new__(MooncakeStorageManager) + manager.storage_client = storage + manager.notify_data_update = AsyncMock(side_effect=lambda *_args, **_kwargs: calls.append("notify")) + manager.controller_handshake_socket = None + manager.storage_manager_id = "contract-test" + manager.zmq_context = MagicMock() + meta = MagicMock(global_indexes=[7], partition_ids=["capacity"], _custom_backend_meta=[{}]) + meta.get_all_custom_meta.return_value = [{}] + context = pytest.raises(RuntimeError, match="capacity exhausted") if put_fails else nullcontext() + with context: + await manager.put_data(TensorDict({"pixel_values": torch.randn(1, 16)}, batch_size=[1]), meta) + assert calls == ([] if put_fails else ["put", "notify"]) diff --git a/tests/utils/test_train_actor_init_cleanup.py b/tests/utils/test_train_actor_init_cleanup.py new file mode 100644 index 000000000..2c1d98548 --- /dev/null +++ b/tests/utils/test_train_actor_init_cleanup.py @@ -0,0 +1,111 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Fail-closed Ray train-actor initialization contracts.""" + +from __future__ import annotations + +import os +import subprocess +import sys +import tempfile +from types import SimpleNamespace +from typing import Any + +import pytest +import ray + +from relax.distributed.ray.actor_group import RayTrainGroup + + +def _group(*actors: Any) -> RayTrainGroup: + group = object.__new__(RayTrainGroup) + group._actor_handlers = list(actors) + return group + + +def _actor(probe: Any = None) -> Any: + return SimpleNamespace(termination_probe=SimpleNamespace(remote=lambda: probe)) + + +def _raise(error: BaseException) -> None: + raise error + + +def test_init_and_wait_preserves_success_and_cleans_submission_failure(monkeypatch: pytest.MonkeyPatch) -> None: + actor = _actor() + group = _group(actor) + refs = [object(), object()] + values = dict(zip(refs, ["rank-0", "rank-1"], strict=True)) + monkeypatch.setattr(group, "async_init", lambda *_args, **_kwargs: refs) + monkeypatch.setattr(ray, "wait", lambda pending, **_kwargs: ([pending[0]], pending[1:])) + monkeypatch.setattr(ray, "get", lambda ref: values[ref]) + assert group.init_and_wait(object(), "actor") == ["rank-0", "rank-1"] + assert group._actor_handlers == [actor] + + monkeypatch.setattr(group, "async_init", lambda *_args, **_kwargs: _raise(RuntimeError("submission failed"))) + monkeypatch.setattr(group, "_terminate_failed_init", lambda: group._actor_handlers.clear()) + with pytest.raises(RuntimeError, match="submission failed"): + group.init_and_wait(object(), "actor") + assert group._actor_handlers == [] + + +def test_first_rank_failure_kills_and_confirms_every_actor(monkeypatch: pytest.MonkeyPatch) -> None: + actors = [_actor("probe-0"), _actor("probe-1")] + group = _group(*actors) + monkeypatch.setattr(group, "async_init", lambda *_args, **_kwargs: ["init-0", "init-1"]) + + def wait(refs: list[str], *, timeout: float | None = None, **_kwargs: Any): + return (list(refs), []) if timeout is not None else (["init-1"], ["init-0"]) + + def get(ref: str) -> None: + if ref == "init-1": + raise RuntimeError("rank initialization failed") + raise ray.exceptions.RayActorError(error_msg="actor terminated") + + killed: list[Any] = [] + monkeypatch.setattr(ray, "wait", wait) + monkeypatch.setattr(ray, "get", get) + monkeypatch.setattr(ray, "kill", lambda actor, no_restart: killed.append((actor, no_restart))) + with pytest.raises(RuntimeError, match="rank initialization failed"): + group.init_and_wait(object(), "actor") + assert killed == [(actors[0], True), (actors[1], True)] + assert group._actor_handlers == [] + + +@pytest.mark.parametrize("failure", ["kill", "pending"]) +def test_unconfirmed_cleanup_keeps_group_and_fails_closed(monkeypatch: pytest.MonkeyPatch, failure: str) -> None: + actors = [_actor("probe-0"), _actor("probe-1")] + group = _group(*actors) + if failure == "kill": + monkeypatch.setattr(ray, "kill", lambda *_args, **_kwargs: _raise(RuntimeError("private detail"))) + monkeypatch.setattr(ray, "wait", lambda refs, **_kwargs: (list(refs), [])) + monkeypatch.setattr( + ray, + "get", + lambda _ref: _raise(ray.exceptions.RayActorError(error_msg="terminated")), + ) + else: + monkeypatch.setattr(ray, "kill", lambda *_args, **_kwargs: None) + monkeypatch.setattr(ray, "wait", lambda refs, **_kwargs: ([], list(refs))) + with pytest.raises(RuntimeError, match="Failed to confirm train actor cleanup") as excinfo: + group._terminate_failed_init(timeout=0.1) + assert "private detail" not in str(excinfo.value) + assert group._actor_handlers == actors + + +def test_real_actor_failure_cannot_mutate_state_after_cleanup() -> None: + env = os.environ.copy() + env["RAY_ENABLE_UV_RUN_RUNTIME_ENV"] = "0" + env.pop("RAY_ADDRESS", None) + root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) + with tempfile.TemporaryDirectory(prefix="train-init-ray-") as probe_dir: + result = subprocess.run( + [sys.executable, "-m", "tests.utils._train_actor_init_cleanup_probe", probe_dir], + cwd=root, + env=env, + capture_output=True, + text=True, + timeout=45, + check=False, + ) + assert result.returncode == 0, f"probe stdout:\n{result.stdout}\nprobe stderr:\n{result.stderr}" diff --git a/tests/utils/tq/__init__.py b/tests/utils/tq/__init__.py new file mode 100644 index 000000000..24cc87577 --- /dev/null +++ b/tests/utils/tq/__init__.py @@ -0,0 +1,3 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""TransferQueue test helpers and focused tests.""" diff --git a/tests/utils/tq/test_config.py b/tests/utils/tq/test_config.py new file mode 100644 index 000000000..d14bba346 --- /dev/null +++ b/tests/utils/tq/test_config.py @@ -0,0 +1,227 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""CPU-only contracts for TransferQueue configuration construction.""" + +from __future__ import annotations + +import argparse +import os +import sys +from types import ModuleType +from typing import Any + +import pytest + +from relax.utils.tq.config import ( + _split_host_port, + build_mooncake_config, + build_simple_storage_config, + estimate_payload_bytes, + resolve_global_segment_size, + resolve_mooncake_master_address, + resolve_tq_capacity_batch_size, + validate_config, + validate_mooncake_runtime_contract, + validate_segment_capacity, +) + + +_MASTER = "master.example:50051" + + +def _args(**overrides: Any) -> argparse.Namespace: + values = { + "tq_rdma_mode": "auto", + "tq_rdma_device": "", + "num_data_storage_units": 1, + "max_staleness": 0, + "n_samples_per_prompt": 1, + "rollout_batch_size": 32, + "multimodal_keys": None, + "seq_length": 8192, + } + values.update(overrides) + return argparse.Namespace(**values) + + +@pytest.mark.parametrize( + ("overrides", "expected_error"), + [({"tq_rdma_mode": mode}, None) for mode in ("off", "auto", "required")] + + [({"tq_rdma_device": device}, None) for device in ("", "rdma0", "mlx5_0")] + + [({"tq_rdma_mode": mode}, "--tq-rdma-mode") for mode in ("mooncake", "auto\nprivate", None, ["auto"])] + + [({"tq_rdma_device": device}, "--tq-rdma-device") for device in (None, ["rdma0"], "rdma0\nforged", " ")], +) +def test_validate_config_matrix(overrides: dict[str, Any], expected_error: str | None) -> None: + errors = validate_config(_args(**overrides)) + if expected_error is None: + assert errors == [] + else: + assert len(errors) == 1 and expected_error in errors[0] + assert repr(next(iter(overrides.values()))) not in errors[0] + + +def test_missing_config_attributes_default_to_off() -> None: + assert validate_config(argparse.Namespace()) == [] + + +@pytest.mark.parametrize( + "address", + [ + "master.example", + "master.example:", + "master.example:abc", + ":50051", + "master.example:0", + "master.example:65536", + "fe80::1", + "fe80::1:50051", + "[fe80::1]", + "[fe80::1]50051", + "master\nprivate:50051", + "master name:50051", + ], +) +def test_master_endpoint_rejects_malformed_values_without_echo(monkeypatch: pytest.MonkeyPatch, address: str) -> None: + monkeypatch.setenv("MC_MASTER_ADDRESS", address) + with pytest.raises(RuntimeError, match="not a usable endpoint") as excinfo: + resolve_mooncake_master_address() + assert address not in str(excinfo.value) + + +def test_master_endpoint_is_required_and_accepts_hostname_or_bracketed_ipv6( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("MC_MASTER_ADDRESS", raising=False) + with pytest.raises(RuntimeError, match="MC_MASTER_ADDRESS"): + resolve_mooncake_master_address() + monkeypatch.setenv("MC_MASTER_ADDRESS", _MASTER) + assert resolve_mooncake_master_address() == _MASTER + assert _split_host_port(_MASTER) == ("master.example", 50051) + assert _split_host_port("[2001:db8::1]:50051") == ("2001:db8::1", 50051) + + +@pytest.mark.parametrize("total_storage_size", [1000, None], ids=["bounded", "unlimited"]) +def test_simple_storage_config(total_storage_size: int | None) -> None: + assert build_simple_storage_config(total_storage_size, 2) == { + "storage_backend": "SimpleStorage", + "SimpleStorage": {"total_storage_size": total_storage_size, "num_data_storage_units": 2}, + } + + +@pytest.mark.parametrize( + ("kwargs", "protocol", "device"), + [({}, "rdma", ""), ({"device": "rdma0"}, "rdma", "rdma0"), ({"protocol": "tcp"}, "tcp", "")], + ids=["production-default", "explicit-device", "benchmark-tcp"], +) +def test_mooncake_config_contract(kwargs: dict[str, str], protocol: str, device: str) -> None: + config = build_mooncake_config(master_address=_MASTER, **kwargs)["MooncakeStore"] + assert (config["protocol"], config["device_name"]) == (protocol, device) + assert config["master_server_address"] == _MASTER + assert config["hard_pin"] is True and config["auto_init"] is False and config["use_gdr"] is False + assert "gdr_staging_buffer_mb" not in config + + +@pytest.mark.parametrize( + ("kwargs", "match"), + [({"master_address": value}, "master_address") for value in ("not-an-endpoint", "host:0", "fe80::1", 50051)] + + [ + ({"master_address": _MASTER, "global_segment_size": value}, "positive integer") for value in (0, -1, True, 1.5) + ], +) +def test_mooncake_builder_rejects_invalid_direct_inputs(kwargs: dict[str, Any], match: str) -> None: + with pytest.raises(ValueError, match=match): + build_mooncake_config(**kwargs) + + +@pytest.mark.parametrize("override", [None, "0"], ids=["default", "explicit-disable"]) +@pytest.mark.parametrize("contract_version", [1, 2], ids=["required", "forward-compatible"]) +def test_runtime_contract_accepts_safe_memcpy( + monkeypatch: pytest.MonkeyPatch, override: str | None, contract_version: int +) -> None: + tq_stub = ModuleType("transfer_queue") + tq_stub.MOONCAKE_CORRECTNESS_CONTRACT_VERSION = contract_version + monkeypatch.setitem(sys.modules, "transfer_queue", tq_stub) + if override is None: + monkeypatch.delenv("MC_STORE_MEMCPY", raising=False) + else: + monkeypatch.setenv("MC_STORE_MEMCPY", override) + validate_mooncake_runtime_contract() + assert os.environ["MC_STORE_MEMCPY"] == "0" + + +@pytest.mark.parametrize( + ("override", "private_marker"), + [("1", None), ("1\nprivate deployment detail", "private deployment detail")], +) +def test_runtime_contract_rejects_unsafe_memcpy_without_echo( + monkeypatch: pytest.MonkeyPatch, override: str, private_marker: str | None +) -> None: + monkeypatch.setenv("MC_STORE_MEMCPY", override) + with pytest.raises(RuntimeError, match="MC_STORE_MEMCPY") as excinfo: + validate_mooncake_runtime_contract() + if private_marker is not None: + assert private_marker not in str(excinfo.value) + + +@pytest.mark.parametrize("contract_version", [None, 0, "1", True], ids=["missing", "old", "string", "bool"]) +def test_runtime_contract_rejects_missing_or_invalid_marker( + monkeypatch: pytest.MonkeyPatch, contract_version: int | str | bool | None +) -> None: + monkeypatch.setenv("MC_STORE_MEMCPY", "0") + tq_stub = ModuleType("transfer_queue") + if contract_version is not None: + tq_stub.MOONCAKE_CORRECTNESS_CONTRACT_VERSION = contract_version + monkeypatch.setitem(sys.modules, "transfer_queue", tq_stub) + with pytest.raises(RuntimeError, match="required Mooncake correctness contract"): + validate_mooncake_runtime_contract() + + +@pytest.mark.parametrize( + ("overrides", "expected_fragment"), + [ + ({"multimodal_keys": None}, None), + ({"multimodal_keys": ["pixel_values"], "rollout_batch_size": 256, "n_samples_per_prompt": 8}, "insufficient"), + ({"multimodal_keys": ["pixel_values"], "rollout_batch_size": 32, "max_staleness": 1}, "staleness+1=2"), + ( + { + "multimodal_keys": ["pixel_values"], + "rollout_batch_size": 16, + "partial_rollout": True, + "use_dynamic_global_batch_size": True, + "over_sampling_batch_size": 64, + }, + "effective_batch=64", + ), + ], +) +def test_segment_capacity_matrix(overrides: dict[str, Any], expected_fragment: str | None) -> None: + error = validate_segment_capacity(_args(**overrides)) + if expected_fragment is None: + assert error is None + else: + assert error is not None and expected_fragment.lower() in error.lower() + + +@pytest.mark.parametrize("value", ["four", "-1", "nan", "inf", "-inf", "1e-20"]) +def test_segment_size_env_rejects_unusable_values(monkeypatch: pytest.MonkeyPatch, value: str) -> None: + monkeypatch.setenv("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB", value) + with pytest.raises(RuntimeError) as excinfo: + resolve_global_segment_size() + assert value not in str(excinfo.value) + + +def test_capacity_override_and_payload_bound(monkeypatch: pytest.MonkeyPatch) -> None: + multimodal = _args(multimodal_keys=["pixel_values"], rollout_batch_size=1) + assert estimate_payload_bytes(_args()) == 32 * 32 * 8192 + assert estimate_payload_bytes(multimodal) == 8192 * (32 + 24_576) + dynamic = _args( + rollout_batch_size=16, + partial_rollout=True, + use_dynamic_global_batch_size=True, + over_sampling_batch_size=64, + ) + assert resolve_tq_capacity_batch_size(dynamic) == 64 + monkeypatch.setenv("RELAX_TQ_GLOBAL_SEGMENT_SIZE_GB", "16") + assert validate_segment_capacity(_args(multimodal_keys=["pixel_values"], max_staleness=1)) is None + with pytest.raises(RuntimeError, match="seq_length"): + estimate_payload_bytes(_args(seq_length=None)) diff --git a/tests/utils/tq/test_payload_assertions.py b/tests/utils/tq/test_payload_assertions.py new file mode 100644 index 000000000..3040bde7f --- /dev/null +++ b/tests/utils/tq/test_payload_assertions.py @@ -0,0 +1,78 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Raw-byte identity contracts shared by dataplane and acceptance tests.""" + +from __future__ import annotations + +import struct +import warnings +from typing import Any + +import numpy as np +import pytest +import torch + +from relax.utils.tq.correctness import diff_digests, leaf_digests, payload_nbytes, payload_rows + + +class NonTensorData: + def __init__(self, data: Any) -> None: + self.data = data + + +class NonTensorStack: + def __init__(self, data: Any) -> None: + self._data = data + + def tolist(self) -> Any: + return self._data + + +def test_digest_preserves_raw_bytes_dtype_shape_and_paths() -> None: + positive = leaf_digests(torch.tensor([0.0], dtype=torch.float32)) + negative = leaf_digests(torch.tensor([-0.0], dtype=torch.float32)) + assert positive["payload"][:2] == negative["payload"][:2] + assert positive["payload"][2] != negative["payload"][2] + + first_nan = struct.unpack("!d", bytes.fromhex("7ff8000000000001"))[0] + second_nan = struct.unpack("!d", bytes.fromhex("7ff8000000000002"))[0] + assert repr(first_nan) == repr(second_nan) == "nan" + assert leaf_digests(first_nan) != leaf_digests(second_nan) + + paths = leaf_digests({"a.b": 1, "a": {"b": 2}}) + assert set(paths) == {"payload['a.b']", "payload.a.b"} + + +def test_digest_supports_nested_numpy_and_non_tensor_containers() -> None: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + nested = torch.nested.nested_tensor( + [torch.tensor([1, 2], dtype=torch.int16), torch.tensor([3], dtype=torch.int16)] + ) + digests = leaf_digests( + { + "array": np.array([[1, 2]], dtype=np.uint16), + "nested": nested, + "wrapped": NonTensorStack([NonTensorData({"text": "hello"}), NonTensorData(None)]), + } + ) + assert digests["payload.array"][:2] == ("np.uint16", "1x2") + assert digests["payload.nested[0]"][:2] == ("torch.int16", "2") + assert digests["payload.nested[1]"][:2] == ("torch.int16", "1") + assert leaf_digests(NonTensorStack([NonTensorData(None)])) == leaf_digests([None]) + + +def test_rows_nbytes_and_diff_cover_non_tensor_payloads() -> None: + tensor = torch.tensor([[1, 2], [3, 4]], dtype=torch.int16) + wrapped = NonTensorStack([NonTensorData({"text": "é"}), NonTensorData(None)]) + assert all(torch.equal(got, want) for got, want in zip(payload_rows(tensor), tensor.unbind(), strict=True)) + assert payload_rows(wrapped) == [{"text": "é"}, None] + assert payload_nbytes({"tensor": tensor, "wrapped": wrapped}) == 8 + 2 + len(repr(None)) + problems = diff_digests(leaf_digests({"expected": 1}), leaf_digests({"actual": 1})) + assert len(problems) == 2 and "unexpected extra" in problems[0] and "missing" in problems[1] + + +@pytest.mark.parametrize("payload", [{1: "value"}, {"bad": object()}], ids=["non-string-key", "unknown-leaf"]) +def test_digest_rejects_unsupported_payload(payload: dict[Any, Any]) -> None: + with pytest.raises(TypeError, match="Unsupported payload"): + leaf_digests(payload)