/test/registered/unit/spec/test_decoupled_spec_io.py

"""Unit tests for srt/speculative/decoupled_spec_io.

decoupled_spec_io is the schema-only IPC layer for decoupled speculative
decoding: protocol message dataclasses, the cross-process request id codec, and
the drafter-side reconciliation helpers. These tests drive the real logic (id
round-trip + parse errors, commit validation, segment coalescing / contiguity /
prefix extraction, and inbox routing) on CPU; there is no GPU or transport here.
"""

import unittest

from sglang.srt.speculative.decoupled_spec_io import (
    DraftClose,
    DraftControlBatch,
    DraftMeshMessage,
    DraftMeshMessageType,
    DraftReqKey,
    DraftSync,
    VerifierCommitSegment,
    VerifyCommit,
    build_draft_scheduler_rid,
    parse_draft_scheduler_rid,
)
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase

register_cpu_ci(est_time=6, suite="base-a-test-cpu")


def _commit(rid, *, pre, tokens, src_verifier_rank=0, drafter_rank=0) -> VerifyCommit:
    return VerifyCommit(
        request_id=rid,
        src_verifier_rank=src_verifier_rank,
        dst_drafter_rank=drafter_rank,
        pre_verify_committed_len=pre,
        committed_tokens=list(tokens),
    )


def _segment(
    rid, *, pre=0, drafter_rank=0, src_verifier_rank=0
) -> VerifierCommitSegment:
    return VerifierCommitSegment(
        draft_key=DraftReqKey(src_verifier_rank=src_verifier_rank, request_id=rid),
        dst_drafter_rank=drafter_rank,
        pre_verify_committed_len=pre,
    )


class TestDraftSchedulerRid(CustomTestCase):
    def test_round_trip(self):
        key = DraftReqKey(src_verifier_rank=3, request_id="req-1")
        rid = build_draft_scheduler_rid(key)
        self.assertEqual(rid, "draft:3:req-1")
        self.assertEqual(parse_draft_scheduler_rid(rid), key)

    def test_request_id_containing_colon_round_trips(self):
        # request_id may contain ':' — parse splits on the first ':' only.
        key = DraftReqKey(src_verifier_rank=1, request_id="a:b:c")
        self.assertEqual(parse_draft_scheduler_rid(build_draft_scheduler_rid(key)), key)

    def test_parse_invalid_rid_raises(self):
        for bad in ["no-prefix", "draft:", "draft:0:", "draft:notint:r"]:
            with self.subTest(rid=bad):
                with self.assertRaises(ValueError):
                    parse_draft_scheduler_rid(bad)


class TestVerifyCommitValidation(CustomTestCase):
    def test_empty_tokens_raises(self):
        with self.assertRaises(ValueError):
            _commit("r", pre=0, tokens=[]).validate_committed_tokens()

    def test_negative_pre_len_raises(self):
        with self.assertRaises(ValueError):
            _commit("r", pre=-1, tokens=[1]).validate_committed_tokens()

    def test_valid_commit_passes(self):
        # Should not raise.
        _commit("r", pre=0, tokens=[1, 2]).validate_committed_tokens()


class TestVerifierCommitSegment(CustomTestCase):
    def test_append_coalesces_contiguous_commits(self):
        seg = _segment("r", pre=0)
        seg.append_message(_commit("r", pre=0, tokens=[10, 11]))
        self.assertEqual(seg.committed_tokens, [10, 11])
        self.assertEqual(seg.end_committed_len, 2)
        seg.append_message(_commit("r", pre=2, tokens=[12]))
        self.assertEqual(seg.committed_tokens, [10, 11, 12])
        self.assertEqual(seg.end_committed_len, 3)

    def test_append_wrong_request_raises(self):
        seg = _segment("r", pre=0)
        with self.assertRaises(RuntimeError):
            seg.append_message(_commit("other", pre=0, tokens=[1]))

    def test_append_wrong_drafter_rank_raises(self):
        seg = _segment("r", pre=0, drafter_rank=0)
        with self.assertRaises(RuntimeError):
            seg.append_message(_commit("r", pre=0, tokens=[1], drafter_rank=7))

    def test_append_non_contiguous_raises(self):
        seg = _segment("r", pre=0)
        seg.append_message(_commit("r", pre=0, tokens=[10]))  # end -> 1
        with self.assertRaises(RuntimeError):
            seg.append_message(_commit("r", pre=5, tokens=[11]))  # gap

    def test_append_runs_message_validation(self):
        seg = _segment("r", pre=0)
        with self.assertRaises(ValueError):
            seg.append_message(_commit("r", pre=0, tokens=[]))  # empty -> validate

    def test_extract_prefix_splits_segment(self):
        seg = _segment("r", pre=0)
        seg.append_message(_commit("r", pre=0, tokens=[10, 11, 12, 13]))
        prefix = seg.extract_prefix(2)
        self.assertEqual(prefix.committed_tokens, [10, 11])
        self.assertEqual(prefix.pre_verify_committed_len, 0)
        # Remainder stays in the original segment, with pre advanced by 2.
        self.assertEqual(seg.committed_tokens, [12, 13])
        self.assertEqual(seg.pre_verify_committed_len, 2)
        self.assertEqual(seg.end_committed_len, 4)

    def test_extract_prefix_bounds(self):
        seg = _segment("r", pre=0)
        seg.append_message(_commit("r", pre=0, tokens=[10, 11]))
        with self.assertRaises(ValueError):
            seg.extract_prefix(0)
        with self.assertRaises(ValueError):
            seg.extract_prefix(3)  # exceeds segment length


class TestDraftControlInbox(CustomTestCase):
    def _inbox(self):
        from sglang.srt.speculative.decoupled_spec_io import DraftControlInbox

        return DraftControlInbox()

    def _sync(self, rid, drafter_rank=0):
        return DraftSync(
            request_id=rid, src_verifier_rank=0, dst_drafter_rank=drafter_rank
        )

    def _close(self, rid, drafter_rank=0):
        return DraftClose(
            request_id=rid,
            src_verifier_rank=0,
            dst_drafter_rank=drafter_rank,
            reason="x",
        )

    def test_add_control_batch_routes_each_message_type(self):
        inbox = self._inbox()
        inbox.add_control_batch_locked(
            DraftControlBatch(
                dst_drafter_rank=0,
                sync_messages=[self._sync("s")],
                verify_commit_messages=[_commit("c", pre=0, tokens=[1])],
                close_messages=[self._close("x")],
            )
        )
        self.assertEqual([m.request_id for m in inbox.sync_messages], ["s"])
        self.assertIn(DraftReqKey(0, "c"), inbox.verifier_commit_segments)
        self.assertIn(DraftReqKey(0, "x"), inbox.close_keys)

    def test_close_drops_pending_segment_and_sync(self):
        inbox = self._inbox()
        inbox.add_control_batch_locked(
            DraftControlBatch(
                dst_drafter_rank=0,
                sync_messages=[self._sync("r")],
                verify_commit_messages=[_commit("r", pre=0, tokens=[1])],
            )
        )
        inbox.add_close_key_locked(DraftReqKey(0, "r"))
        self.assertEqual(inbox.sync_messages, [])
        self.assertNotIn(DraftReqKey(0, "r"), inbox.verifier_commit_segments)
        self.assertIn(DraftReqKey(0, "r"), inbox.close_keys)

    def test_verify_commit_for_closed_key_is_ignored(self):
        inbox = self._inbox()
        inbox.add_close_key_locked(DraftReqKey(0, "r"))
        inbox.add_verify_commit_locked(_commit("r", pre=0, tokens=[1]))
        self.assertNotIn(DraftReqKey(0, "r"), inbox.verifier_commit_segments)

    def test_extract_ready_controls_full_consume(self):
        inbox = self._inbox()
        inbox.add_control_batch_locked(
            DraftControlBatch(
                dst_drafter_rank=0,
                sync_messages=[self._sync("s")],
                verify_commit_messages=[_commit("c", pre=0, tokens=[1, 2])],
                close_messages=[self._close("x")],
            )
        )
        ready = inbox.extract_ready_controls_locked(
            lambda seg: len(seg.committed_tokens)
        )
        self.assertEqual([m.request_id for m in ready.sync_messages], ["s"])
        self.assertEqual({k.request_id for k in ready.close_keys}, {"x"})
        self.assertEqual(len(ready.ready_commit_segments), 1)
        self.assertEqual(ready.ready_commit_segments[0].committed_tokens, [1, 2])
        # Fully consumed -> the segment is gone; inbox drained.
        self.assertTrue(inbox.is_empty())

    def test_extract_ready_controls_zero_consumable_keeps_segment(self):
        inbox = self._inbox()
        inbox.add_verify_commit_locked(_commit("c", pre=0, tokens=[1, 2]))
        ready = inbox.extract_ready_controls_locked(lambda seg: 0)
        self.assertEqual(ready.ready_commit_segments, [])
        # Segment is left buffered for a later step.
        self.assertIn(DraftReqKey(0, "c"), inbox.verifier_commit_segments)

    def test_extract_ready_controls_partial_consume_buffers_remainder(self):
        inbox = self._inbox()
        inbox.add_verify_commit_locked(_commit("c", pre=0, tokens=[1, 2, 3]))
        ready = inbox.extract_ready_controls_locked(lambda seg: 1)
        self.assertEqual(ready.ready_commit_segments[0].committed_tokens, [1])
        # Remainder [2, 3] stays buffered with pre advanced to 1.
        seg = inbox.verifier_commit_segments[DraftReqKey(0, "c")]
        self.assertEqual(seg.committed_tokens, [2, 3])
        self.assertEqual(seg.pre_verify_committed_len, 1)

    def test_close_in_same_batch_drops_sync_and_commit(self):
        # add_control_batch applies close first, so a same-key sync/commit in the
        # same batch is dropped/ignored: close wins within one batch.
        inbox = self._inbox()
        inbox.add_control_batch_locked(
            DraftControlBatch(
                dst_drafter_rank=0,
                sync_messages=[self._sync("r")],
                verify_commit_messages=[_commit("r", pre=0, tokens=[1])],
                close_messages=[self._close("r")],
            )
        )
        self.assertEqual(inbox.sync_messages, [])
        self.assertNotIn(DraftReqKey(0, "r"), inbox.verifier_commit_segments)
        self.assertEqual({k.request_id for k in inbox.close_keys}, {"r"})

    def test_two_commits_same_key_coalesce_in_inbox(self):
        # A second commit for an existing key appends to the buffered segment.
        inbox = self._inbox()
        inbox.add_verify_commit_locked(_commit("r", pre=0, tokens=[1, 2]))
        inbox.add_verify_commit_locked(_commit("r", pre=2, tokens=[3]))
        seg = inbox.verifier_commit_segments[DraftReqKey(0, "r")]
        self.assertEqual(seg.committed_tokens, [1, 2, 3])
        self.assertEqual(seg.end_committed_len, 3)


class TestDraftMeshMessageEnvelope(CustomTestCase):
    def test_from_control_batch_sets_discriminant_and_slot(self):
        batch = DraftControlBatch(dst_drafter_rank=0)
        msg = DraftMeshMessage.from_control_batch(batch)
        self.assertEqual(msg.message_type, DraftMeshMessageType.CONTROL_BATCH)
        self.assertIs(msg.control_batch, batch)
        self.assertIsNone(msg.tail_stream_output_batch)

    def test_from_tail_stream_output_batch_sets_discriminant_and_slot(self):
        from sglang.srt.speculative.decoupled_spec_io import DraftTailStreamOutputBatch

        batch = DraftTailStreamOutputBatch()
        msg = DraftMeshMessage.from_tail_stream_output_batch(batch)
        self.assertEqual(
            msg.message_type, DraftMeshMessageType.TAIL_STREAM_OUTPUT_BATCH
        )
        self.assertIs(msg.tail_stream_output_batch, batch)
        self.assertIsNone(msg.control_batch)


if __name__ == "__main__":
    unittest.main(verbosity=3)