leavelet
9a91c73e09
Add batch_transfer_async_submit() (non-blocking submit -> batch_id) and wait_batch_transfers() (block-until-all, graceful) to MooncakeTransferEngine, exposing the mooncake async API for the upcoming per-layer overlapped transfer. Deliberately avoids batch_transfer_*_on_cuda: its CUDA host callback busy-waits on the stream and calls _exit(1) on transfer failure (verified in the mooncake source transfer_engine_py.cpp), which is unsafe for production (a decode-side abort mid-transfer would crash the prefill). The async submit + wait pair lets per-layer submits pipeline in the RDMA engine, then waits once on a background thread with graceful failure. Both bindings raise a clear upgrade error if the engine predates the API and degrade to -1 on transient errors. Mock-engine unit tested (8 cases). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
95 lines
3.1 KiB
Python
95 lines
3.1 KiB
Python
"""Unit tests for the async/non-blocking mooncake wrapper bindings (G1).
|
|
|
|
batch_transfer_async_submit() must do a non-blocking submit returning a batch_id;
|
|
wait_batch_transfers() must block-until-all and return 0/non-zero. Both must raise a
|
|
clear upgrade error if the engine lacks the method, and degrade to -1 on transient
|
|
engine errors. No GPU / real RDMA needed — the engine is mocked.
|
|
"""
|
|
|
|
import unittest
|
|
|
|
from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (
|
|
MooncakeTransferEngine,
|
|
)
|
|
|
|
|
|
class _GoodEngine:
|
|
def __init__(self, status_ret=0):
|
|
self.calls = []
|
|
self._status_ret = status_ret
|
|
|
|
def batch_transfer_async_write(self, sid, buffers, peers, lengths):
|
|
self.calls.append(("submit", sid, list(buffers), list(peers), list(lengths)))
|
|
return 7 # batch_id handle
|
|
|
|
def get_batch_transfer_status(self, ids):
|
|
self.calls.append(("wait", list(ids)))
|
|
return self._status_ret
|
|
|
|
|
|
class _NoAsyncEngine:
|
|
pass # lacks both methods
|
|
|
|
|
|
class _RaisingEngine:
|
|
def batch_transfer_async_write(self, *a):
|
|
raise RuntimeError("transient rdma error")
|
|
|
|
def get_batch_transfer_status(self, *a):
|
|
raise RuntimeError("transient rdma error")
|
|
|
|
|
|
def _wrapper(engine):
|
|
w = MooncakeTransferEngine.__new__(MooncakeTransferEngine) # skip real init/connect
|
|
w.engine = engine
|
|
return w
|
|
|
|
|
|
class TestAsyncSubmit(unittest.TestCase):
|
|
def test_submit_returns_batch_id_and_dispatches(self):
|
|
eng = _GoodEngine()
|
|
w = _wrapper(eng)
|
|
bid = w.batch_transfer_async_submit("sess", [1, 2], [10, 20], [64, 64])
|
|
self.assertEqual(bid, 7)
|
|
self.assertEqual(eng.calls[0], ("submit", "sess", [1, 2], [10, 20], [64, 64]))
|
|
|
|
def test_submit_missing_method_raises_upgrade(self):
|
|
w = _wrapper(_NoAsyncEngine())
|
|
with self.assertRaises(RuntimeError):
|
|
w.batch_transfer_async_submit("s", [1], [2], [3])
|
|
|
|
def test_submit_transient_error_returns_minus_one(self):
|
|
w = _wrapper(_RaisingEngine())
|
|
self.assertEqual(w.batch_transfer_async_submit("s", [1], [2], [3]), -1)
|
|
|
|
|
|
class TestWait(unittest.TestCase):
|
|
def test_wait_dispatches_and_returns_status(self):
|
|
eng = _GoodEngine(status_ret=0)
|
|
w = _wrapper(eng)
|
|
self.assertEqual(w.wait_batch_transfers([7, 8]), 0)
|
|
self.assertEqual(eng.calls[0], ("wait", [7, 8]))
|
|
|
|
def test_wait_empty_is_zero_without_engine_call(self):
|
|
eng = _GoodEngine()
|
|
w = _wrapper(eng)
|
|
self.assertEqual(w.wait_batch_transfers([]), 0)
|
|
self.assertEqual(eng.calls, [])
|
|
|
|
def test_wait_nonzero_status_propagates(self):
|
|
w = _wrapper(_GoodEngine(status_ret=-1))
|
|
self.assertEqual(w.wait_batch_transfers([7]), -1)
|
|
|
|
def test_wait_missing_method_raises_upgrade(self):
|
|
w = _wrapper(_NoAsyncEngine())
|
|
with self.assertRaises(RuntimeError):
|
|
w.wait_batch_transfers([7])
|
|
|
|
def test_wait_transient_error_returns_minus_one(self):
|
|
w = _wrapper(_RaisingEngine())
|
|
self.assertEqual(w.wait_batch_transfers([7]), -1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|