test(disagg): set MC_GID_INDEX on RoCE hosts so mooncake KV transfer works (#30737)

This commit is contained in:
Alison Shao
2026-07-11 11:29:22 +08:00
committed by GitHub
parent 0663ebc783
commit 90688366d9
@@ -58,9 +58,12 @@ class PDDisaggregationServerBase(CustomTestCase):
cls._fail_fast_stop = None
# config transfer backend and rdma devices
cls._mc_gid_index_set = False
if is_in_ci():
cls.transfer_backend = ["--disaggregation-transfer-backend", "mooncake"]
cls.rdma_devices = ["--disaggregation-ib-device", get_rdma_devices_args()]
ib_devices = get_rdma_devices_args()
cls.rdma_devices = ["--disaggregation-ib-device", ib_devices]
cls._mc_gid_index_set = _maybe_set_roce_gid_index(ib_devices)
else:
cls.transfer_backend = [
"--disaggregation-transfer-backend",
@@ -183,6 +186,8 @@ class PDDisaggregationServerBase(CustomTestCase):
if cls._fail_fast_stop is not None:
cls._fail_fast_stop.set()
os.environ.pop("MC_TCP_ENABLE_CONNECTION_POOL")
if getattr(cls, "_mc_gid_index_set", False):
os.environ.pop("MC_GID_INDEX", None)
for process in [cls.process_lb, cls.process_decode, cls.process_prefill]:
if process:
try:
@@ -371,3 +376,83 @@ def get_rdma_devices_args():
# Deduplicate while preserving order
return ",".join(dict.fromkeys(rdma_devices))
_IB_SYSFS = "/sys/class/infiniband"
def _roce_v2_gid_index(device: str):
"""Return a RoCEv2 GID index for a device, preferring a global (routable)
GID over a link-local (fe80::) one, or None if the device has no RoCEv2 GID.
"""
port = os.path.join(_IB_SYSFS, device, "ports", "1")
types_dir = os.path.join(port, "gid_attrs", "types")
try:
indices = sorted(int(x) for x in os.listdir(types_dir) if x.isdigit())
except OSError:
return None
fallback = None
for i in indices:
try:
with open(os.path.join(types_dir, str(i))) as f:
if f.read().strip() != "RoCE v2":
continue
except OSError:
continue
if fallback is None:
fallback = i
try:
with open(os.path.join(port, "gids", str(i))) as f:
gid = f.read().strip()
except OSError:
gid = ""
# Prefer a global GID; link-local (fe80::) entries don't route between
# NICs on some fabrics.
if gid and not gid.lower().startswith("fe80"):
return i
return fallback
def _detect_roce_gid_index(devices):
"""Return a single RoCEv2 GID index shared by all `devices`, or None.
None when any device is InfiniBand (mooncake selects the GID automatically
there), when a device has no RoCEv2 GID, or when devices disagree on the
index — MC_GID_INDEX is a single global value, so a divergent set can't be
satisfied and is left to mooncake's own selection.
"""
picked = None
for device in [d.strip() for d in devices if d.strip()]:
try:
with open(os.path.join(_IB_SYSFS, device, "ports", "1", "link_layer")) as f:
if f.read().strip() != "Ethernet":
return None
except OSError:
return None
idx = _roce_v2_gid_index(device)
if idx is None:
return None
if picked is None:
picked = idx
elif picked != idx:
return None
return picked
def _maybe_set_roce_gid_index(ib_devices) -> bool:
"""Export MC_GID_INDEX for a RoCE fabric; return True if this call set it.
On RoCE-only hosts mooncake's automatic GID selection can come up empty
("GID is NULL, please check your GID index by specifying MC_GID_INDEX"),
leaving the KV-transfer RDMA endpoint with no GID so every prefill->decode
transfer fails and PD accuracy collapses to 0. InfiniBand hosts don't need
this (auto GID works), and a user-provided MC_GID_INDEX is left untouched.
"""
if not ib_devices or os.environ.get("MC_GID_INDEX"):
return False
gid_index = _detect_roce_gid_index(ib_devices.split(","))
if gid_index is None:
return False
os.environ["MC_GID_INDEX"] = str(gid_index)
logger.warning("RoCE fabric detected; set MC_GID_INDEX=%d for mooncake", gid_index)
return True